"""AI 编排运行账本服务。""" from __future__ import annotations import hashlib from datetime import datetime, timedelta from typing import Any, Optional from sqlalchemy import and_, or_, select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from ...core.database import AsyncSessionLocal from ...models import AIOrchestrationRun from ...utils.timezone import CN_TZ, now_cn_naive def _utc_now_naive() -> datetime: return now_cn_naive() def _normalize_naive_datetime(value: Optional[datetime]) -> Optional[datetime]: if value is None: return None if not isinstance(value, datetime): return None if value.tzinfo is None: return value return value.astimezone(CN_TZ).replace(tzinfo=None) INGRESS_RUN_TYPE = "ingress_feishu_message" INGRESS_RETRY_SCHEDULE_SECONDS: tuple[int, ...] = (10, 30, 90) INGRESS_RUNNING_STALE_MINUTES = 30 class AIOrchestrationRunService: def __init__(self, db: AsyncSession) -> None: self.db = db async def create_run( self, *, tenant_id: int, trace_id: str, run_type: str, project_id: Optional[int], signal_id: Optional[int], response_mode: str, status: str, planner_status: Optional[str], execution_status: Optional[str], worker_total: int, worker_completed: int, worker_failed: int, summary: str, planner_diagnostics: Optional[dict[str, Any]], result_payload: Optional[dict[str, Any]], created_by_user_id: Optional[int], started_at: Optional[datetime] = None, completed_at: Optional[datetime] = None, ) -> AIOrchestrationRun: row = AIOrchestrationRun( tenant_id=tenant_id, trace_id=trace_id, run_type=run_type, project_id=project_id, signal_id=signal_id, response_mode=response_mode, status=status, planner_status=planner_status, openclaw_status=execution_status, worker_total=max(0, int(worker_total)), worker_completed=max(0, int(worker_completed)), worker_failed=max(0, int(worker_failed)), summary=str(summary or "").strip()[:1000] or None, planner_diagnostics=planner_diagnostics if isinstance(planner_diagnostics, dict) else None, result_payload=result_payload if isinstance(result_payload, dict) else None, created_by_user_id=created_by_user_id, started_at=started_at or _utc_now_naive(), completed_at=completed_at, ) self.db.add(row) await self.db.flush() return row @staticmethod def build_ingress_trace_id( *, source_system: str, source_message_id: str, ) -> str: normalized_source = str(source_system or "").strip().lower()[:24] or "unknown" normalized_message_id = str(source_message_id or "").strip()[:96] if normalized_message_id: raw = f"{normalized_source}:{normalized_message_id}" else: digest = hashlib.sha1(f"{normalized_source}:{_utc_now_naive().isoformat()}".encode("utf-8"), usedforsecurity=False).hexdigest() raw = f"{normalized_source}:{digest[:40]}" return raw[:64] async def get_run_by_id(self, *, run_id: int) -> Optional[AIOrchestrationRun]: return ( await self.db.execute( select(AIOrchestrationRun).where(AIOrchestrationRun.id == int(run_id)) ) ).scalar_one_or_none() async def enqueue_ingress_run( self, *, tenant_id: int, source_system: str, source_message_id: str, sender: str, message_text: str, channel_ref: Optional[str], source_sender_open_id: Optional[str], chat_type: Optional[str], captured_at: Optional[datetime], project_hint_id: Optional[int], max_retries: int = 3, ) -> tuple[AIOrchestrationRun, bool]: trace_id = self.build_ingress_trace_id( source_system=source_system, source_message_id=source_message_id, ) existing = await self.get_run(tenant_id=tenant_id, trace_id=trace_id) if existing is not None: return existing, False result_payload = { "queue_mode": "async_ingress", "source_system": str(source_system or "").strip()[:32], "source_message_id": str(source_message_id or "").strip()[:128], "queued_at": _utc_now_naive().isoformat(), } row = AIOrchestrationRun( tenant_id=int(tenant_id), trace_id=trace_id, run_type=INGRESS_RUN_TYPE, project_id=int(project_hint_id) if project_hint_id is not None else None, signal_id=None, response_mode="fast", status="queued", planner_status=None, openclaw_status=None, worker_total=1, worker_completed=0, worker_failed=0, summary="消息已入异步处理队列。", planner_diagnostics=None, result_payload=result_payload, created_by_user_id=None, started_at=None, completed_at=None, source_system=str(source_system or "").strip()[:32] or None, source_message_id=str(source_message_id or "").strip()[:128] or None, sender=str(sender or "").strip()[:128] or "unknown", source_sender_open_id=str(source_sender_open_id or "").strip()[:128] or None, channel_ref=str(channel_ref or "").strip()[:128] or None, chat_type=str(chat_type or "").strip()[:32] or None, captured_at=_normalize_naive_datetime(captured_at), message_text=str(message_text or "").strip()[:8000] or "", project_hint_id=int(project_hint_id) if project_hint_id is not None else None, retry_count=0, max_retries=max(0, min(int(max_retries), 5)), next_retry_at=_utc_now_naive(), last_error=None, last_error_at=None, ack_sent_at=None, result_notified_at=None, work_item_id=None, ) self.db.add(row) try: await self.db.flush() except IntegrityError: await self.db.rollback() existing = await self.get_run(tenant_id=tenant_id, trace_id=trace_id) if existing is not None: return existing, False raise return row, True async def claim_next_ingress_run(self, *, now: Optional[datetime] = None) -> Optional[AIOrchestrationRun]: now_at = now or _utc_now_naive() stale_before = now_at - timedelta(minutes=INGRESS_RUNNING_STALE_MINUTES) stale_rows = ( await self.db.execute( select(AIOrchestrationRun) .where( and_( AIOrchestrationRun.run_type == INGRESS_RUN_TYPE, AIOrchestrationRun.status == "running", AIOrchestrationRun.updated_at <= stale_before, ) ) .with_for_update(skip_locked=True) ) ).scalars().all() for stale in stale_rows: if int(stale.retry_count or 0) < int(stale.max_retries or 0): delay = INGRESS_RETRY_SCHEDULE_SECONDS[min(int(stale.retry_count or 0), len(INGRESS_RETRY_SCHEDULE_SECONDS) - 1)] stale.retry_count = int(stale.retry_count or 0) + 1 stale.status = "retrying" stale.next_retry_at = now_at + timedelta(seconds=int(delay)) else: stale.status = "failed" stale.completed_at = now_at stale.next_retry_at = None stale.last_error = "stale_running_recovered" stale.last_error_at = now_at candidate_id = ( await self.db.execute( select(AIOrchestrationRun.id) .where( and_( AIOrchestrationRun.run_type == INGRESS_RUN_TYPE, or_( and_( AIOrchestrationRun.status.in_(("queued", "retrying")), or_( AIOrchestrationRun.next_retry_at.is_(None), AIOrchestrationRun.next_retry_at <= now_at, ), ), and_( AIOrchestrationRun.status == "failed", AIOrchestrationRun.last_error == "greenlet_spawn_error", AIOrchestrationRun.retry_count < AIOrchestrationRun.max_retries, or_( AIOrchestrationRun.next_retry_at.is_(None), AIOrchestrationRun.next_retry_at <= now_at, ), ), ), ) ) .order_by(AIOrchestrationRun.id.asc()) .with_for_update(skip_locked=True) .limit(1) ) ).scalar_one_or_none() if candidate_id is None: return None row = await self.get_run_by_id(run_id=int(candidate_id)) if row is None: return None if str(row.status or "") == "failed" and str(row.last_error or "") == "greenlet_spawn_error": row.retry_count = min( int(row.max_retries or 0), int(row.retry_count or 0) + 1, ) row.status = "running" row.openclaw_status = "running" row.started_at = row.started_at or now_at row.next_retry_at = None row.last_error = None row.last_error_at = None await self.db.flush() return row async def mark_ingress_ack_sent( self, *, tenant_id: int, trace_id: str, ) -> Optional[AIOrchestrationRun]: row = await self.get_run(tenant_id=tenant_id, trace_id=trace_id) if row is None: return None row.ack_sent_at = _utc_now_naive() await self.db.flush() return row async def mark_ingress_succeeded( self, *, run_id: int, result_payload: dict[str, Any], work_item_id: Optional[int], summary: Optional[str] = None, ) -> Optional[AIOrchestrationRun]: row = await self.get_run_by_id(run_id=run_id) if row is None: return None now_at = _utc_now_naive() row.status = "succeeded" row.openclaw_status = "succeeded" row.worker_total = max(1, int(row.worker_total or 1)) row.worker_completed = max(1, int(row.worker_completed or 0)) row.worker_failed = max(0, int(row.worker_failed or 0)) row.result_payload = result_payload if isinstance(result_payload, dict) else {} if work_item_id is not None: row.work_item_id = int(work_item_id) row.summary = (str(summary or "").strip()[:1000] or "异步处理完成。") row.completed_at = now_at row.next_retry_at = None row.last_error = None row.last_error_at = None await self.db.flush() return row async def mark_ingress_failed_or_retrying( self, *, run_id: int, error_message: str, ) -> tuple[Optional[AIOrchestrationRun], bool]: # Failure path must not reuse a possibly committed/invalid caller session. async with AsyncSessionLocal() as new_db: row = ( await new_db.execute( select(AIOrchestrationRun) .where(AIOrchestrationRun.id == int(run_id)) .with_for_update(skip_locked=True) ) ).scalar_one_or_none() if row is None: await new_db.commit() return None, False now_at = _utc_now_naive() normalized_error = str(error_message or "").strip()[:1000] or "unknown_error" retry_count = int(row.retry_count or 0) max_retries = max(0, int(row.max_retries or 0)) should_retry = retry_count < max_retries if should_retry: delay = INGRESS_RETRY_SCHEDULE_SECONDS[min(retry_count, len(INGRESS_RETRY_SCHEDULE_SECONDS) - 1)] row.retry_count = retry_count + 1 row.status = "retrying" row.openclaw_status = "retrying" row.next_retry_at = now_at + timedelta(seconds=int(delay)) row.completed_at = None row.worker_failed = max(0, int(row.worker_failed or 0)) else: row.status = "failed" row.openclaw_status = "failed" row.next_retry_at = None row.completed_at = now_at row.worker_failed = max(1, int(row.worker_failed or 0) + 1) row.last_error = normalized_error row.last_error_at = now_at row.summary = ( f"异步处理失败（已重试{int(row.retry_count or 0)}次）" if should_retry else f"异步处理失败（自动重试{max_retries}次后终止）" ) await new_db.flush() await new_db.commit() return row, (not should_retry) async def mark_ingress_result_notified( self, *, run_id: int, ) -> Optional[AIOrchestrationRun]: row = await self.get_run_by_id(run_id=run_id) if row is None: return None row.result_notified_at = _utc_now_naive() await self.db.flush() return row async def get_run(self, *, tenant_id: int, trace_id: str) -> Optional[AIOrchestrationRun]: return ( await self.db.execute( select(AIOrchestrationRun).where( and_( AIOrchestrationRun.tenant_id == int(tenant_id), AIOrchestrationRun.trace_id == str(trace_id), ) ) ) ).scalar_one_or_none() async def update_run( self, *, tenant_id: int, trace_id: str, status: Optional[str] = None, planner_status: Optional[str] = None, execution_status: Optional[str] = None, worker_total: Optional[int] = None, worker_completed: Optional[int] = None, worker_failed: Optional[int] = None, summary: Optional[str] = None, planner_diagnostics: Optional[dict[str, Any]] = None, result_payload: Optional[dict[str, Any]] = None, completed_at: Optional[datetime] = None, ) -> Optional[AIOrchestrationRun]: row = await self.get_run(tenant_id=tenant_id, trace_id=trace_id) if row is None: return None if status is not None: row.status = str(status or "").strip()[:20] if planner_status is not None: row.planner_status = str(planner_status or "").strip()[:20] or None if execution_status is not None: row.openclaw_status = str(execution_status or "").strip()[:20] or None if worker_total is not None: row.worker_total = max(0, int(worker_total)) if worker_completed is not None: row.worker_completed = max(0, int(worker_completed)) if worker_failed is not None: row.worker_failed = max(0, int(worker_failed)) if summary is not None: row.summary = str(summary or "").strip()[:1000] or None if planner_diagnostics is not None: row.planner_diagnostics = planner_diagnostics if result_payload is not None: row.result_payload = result_payload if completed_at is not None: row.completed_at = completed_at return row async def mark_stale_runs_interrupted( self, *, older_than_minutes: int = 10, ) -> int: threshold = _utc_now_naive() - timedelta(minutes=max(1, int(older_than_minutes))) rows = ( await self.db.execute( select(AIOrchestrationRun).where( and_( AIOrchestrationRun.status.in_(("queued", "running")), AIOrchestrationRun.run_type != INGRESS_RUN_TYPE, AIOrchestrationRun.started_at.is_not(None), AIOrchestrationRun.started_at < threshold, ) ) ) ).scalars().all() count = 0 for row in rows: payload = row.result_payload if isinstance(row.result_payload, dict) else {} payload["run_status"] = "interrupted" payload["pending_workers"] = payload.get("pending_workers") or [] row.status = "interrupted" row.completed_at = _utc_now_naive() row.result_payload = payload count += 1 return count @staticmethod def serialize_run(row: AIOrchestrationRun) -> dict[str, Any]: payload = row.result_payload if isinstance(row.result_payload, dict) else {} pending_workers = payload.get("pending_workers") if isinstance(payload.get("pending_workers"), list) else [] worker_total = int(row.worker_total or 0) worker_completed = int(row.worker_completed or 0) worker_failed = int(row.worker_failed or 0) worker_pending = max(0, worker_total - worker_completed - worker_failed) if not pending_workers and worker_pending > 0: pending_workers = [f"pending_{idx}" for idx in range(1, worker_pending + 1)] return { "trace_id": str(row.trace_id or ""), "run_type": str(row.run_type or ""), "project_id": int(row.project_id) if row.project_id is not None else None, "signal_id": int(row.signal_id) if row.signal_id is not None else None, "response_mode": str(row.response_mode or ""), "run_status": str(row.status or ""), "planner_status": str(row.planner_status or ""), "execution_status": str(row.openclaw_status or ""), "worker_total": worker_total, "worker_completed": worker_completed, "worker_failed": worker_failed, "worker_pending": worker_pending, "pending_workers": pending_workers, "summary": str(row.summary or payload.get("summary") or "").strip(), "planner_diagnostics": row.planner_diagnostics if isinstance(row.planner_diagnostics, dict) else None, "result_snapshot": payload, "source_system": str(row.source_system or ""), "source_message_id": str(row.source_message_id or ""), "sender": str(row.sender or ""), "channel_ref": str(row.channel_ref or ""), "project_hint_id": int(row.project_hint_id) if row.project_hint_id is not None else None, "retry_count": int(row.retry_count or 0), "max_retries": int(row.max_retries or 0), "next_retry_at": row.next_retry_at.isoformat() if row.next_retry_at else None, "last_error": str(row.last_error or ""), "work_item_id": int(row.work_item_id) if row.work_item_id is not None else None, "ack_sent_at": row.ack_sent_at.isoformat() if row.ack_sent_at else None, "result_notified_at": row.result_notified_at.isoformat() if row.result_notified_at else None, "created_by_user_id": int(row.created_by_user_id) if row.created_by_user_id is not None else None, "started_at": row.started_at.isoformat() if row.started_at else None, "completed_at": row.completed_at.isoformat() if row.completed_at else None, }