"""AI 编排运行账本服务。""" from __future__ import annotations from datetime import datetime, timedelta, timezone from typing import Any, Optional from sqlalchemy import and_, select from sqlalchemy.ext.asyncio import AsyncSession from ...models import AIOrchestrationRun from ...utils.timezone import now_cn_naive def _utc_now_naive() -> datetime: return now_cn_naive() 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: 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 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.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), "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, "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, }