import asyncio import os import time import unittest from datetime import datetime, timedelta, timezone from tempfile import NamedTemporaryFile from types import SimpleNamespace from unittest.mock import AsyncMock, patch from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from app.models import Base from app.services.ai_orchestration_run_service import AIOrchestrationRunService from app.services.signal_source_orchestration_service import SignalSourceOrchestrationService def _utc_now_naive() -> datetime: return datetime.now(timezone.utc).replace(tzinfo=None) class _FastModeSignalService(SignalSourceOrchestrationService): session_factory: async_sessionmaker[AsyncSession] | None = None async def _load_project(self, **_: object): return SimpleNamespace(id=3, title="项目A", name="项目A", description="desc", owner_id=1) async def _load_signal(self, **_: object): return SimpleNamespace( id=9, title="新增 GitHub 信号源", summary="监控 issue", signal_type="新需求", recommended_action="enter_only", content="新增 GitHub 信号源并自动拆任务", ) async def _plan_onboarding(self, **_: object): return ( { "should_onboard": True, "summary": "接入 GitHub 仓库信号源。", "source": { "source_kind": "github", "source_name": "核心仓库", "source_identifier": "github:openai/codex", "description": "监控 issue 与 PR", }, "worker_roles": ["pm_dispatcher", "tech_architect", "delivery_coach"], "openclaw_instruction": "创建接入任务。", }, "planned", { "used_fallback": False, "error_code": None, "source_field": "message.content", "provider": "glm", "model": "GLM-4.7", }, ) async def _create_or_reuse_registry(self, **_: object): return { "status": "created", "registry_id": 11, "source_kind": "github", "source_name": "核心仓库", "source_identifier": "github:openai/codex", "description": "监控 issue 与 PR", } async def _resolve_operator_user(self, **_: object): return SimpleNamespace(id=1, role="BOSS", name="Boss", username="boss") async def _dispatch_openclaw(self, **_: object): return { "status": "completed", "result": {"trace_id": "openclaw-123"}, } @classmethod async def _execute_workers_and_finalize(cls, **kwargs: object): await asyncio.sleep(0.2) assert cls.session_factory is not None async with cls.session_factory() as db: service = cls(db) trace_id = str(kwargs["trace_id"]) summary = str(kwargs["summary"]) planner_status = str(kwargs["planner_status"]) planner_diagnostics = kwargs["planner_diagnostics"] registry_result = dict(kwargs["registry_result"]) openclaw_result = dict(kwargs["openclaw_result"]) worker_role_keys = list(kwargs["worker_role_keys"]) worker_results = [ { "worker_role_key": role_key, "worker_role_name": role_key, "status": "completed", "summary": f"{role_key} done", } for role_key in worker_role_keys ] final_result = service._build_orchestration_result( trace_id=trace_id, response_mode="fast", planner_status=planner_status, planner_diagnostics=planner_diagnostics, summary=summary, registry_result=registry_result, source_payload=registry_result, openclaw_result=openclaw_result, worker_results=worker_results, pending_workers=[], status="completed", run_status="completed", ) await AIOrchestrationRunService(db).update_run( tenant_id=int(kwargs["tenant_id"]), trace_id=trace_id, status="completed", planner_status=planner_status, openclaw_status=str(openclaw_result.get("status") or ""), worker_total=len(worker_role_keys), worker_completed=len(worker_role_keys), worker_failed=0, summary=summary, planner_diagnostics=planner_diagnostics if isinstance(planner_diagnostics, dict) else {}, result_payload=final_result, completed_at=_utc_now_naive(), ) await db.commit() return final_result class _ConcurrentWorkerSignalService(SignalSourceOrchestrationService): @classmethod async def _execute_single_worker_role_isolated(cls, **kwargs: object): await asyncio.sleep(0.2) role = dict(kwargs["role"]) return { "worker_role_key": str(role["role_key"]), "worker_role_name": str(role["role_name"]), "status": "completed", "summary": f"{role['role_key']} done", } class AIOrchestrationRunServiceTest(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self) -> None: tmp = NamedTemporaryFile(prefix="pmflow-orchestration-test-", suffix=".db", delete=False) tmp.close() self.db_path = tmp.name self.engine = create_async_engine(f"sqlite+aiosqlite:///{self.db_path}", future=True) self.session_factory = async_sessionmaker(self.engine, expire_on_commit=False, class_=AsyncSession) async with self.engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) _FastModeSignalService.session_factory = self.session_factory _ConcurrentWorkerSignalService.session_factory = self.session_factory async def asyncTearDown(self) -> None: await self.engine.dispose() if getattr(self, "db_path", None): try: os.remove(self.db_path) except FileNotFoundError: pass async def test_fast_mode_returns_before_background_workers_finish_and_run_is_queryable(self) -> None: async with self.session_factory() as db: service = _FastModeSignalService(db) started = time.perf_counter() result = await service.orchestrate_from_signal( tenant_id=1, project_id=3, signal_id=9, initiating_text="新增一个 GitHub 信号源，监控核心仓库 issue。", requested_via="pmflow", operator_user_id=1, dry_run=False, response_mode="fast", ) elapsed = time.perf_counter() - started self.assertLess(elapsed, 0.15) self.assertEqual(result["run_status"], "running") self.assertEqual(len(result["pending_workers"]), 3) async with self.session_factory() as db: row = await AIOrchestrationRunService(db).get_run(tenant_id=1, trace_id=result["trace_id"]) self.assertIsNotNone(row) self.assertEqual(row.status, "running") await asyncio.sleep(0.3) async with self.session_factory() as db: row = await AIOrchestrationRunService(db).get_run(tenant_id=1, trace_id=result["trace_id"]) snapshot = AIOrchestrationRunService.serialize_run(row) self.assertEqual(snapshot["run_status"], "completed") self.assertEqual(snapshot["worker_completed"], 3) self.assertEqual(snapshot["worker_pending"], 0) async def test_complete_mode_worker_execution_is_concurrent(self) -> None: async with self.session_factory() as db: service = _ConcurrentWorkerSignalService(db) initial_result = service._build_orchestration_result( trace_id="trace-concurrent", response_mode="complete", planner_status="planned", planner_diagnostics={"used_fallback": False}, summary="接入 GitHub 仓库信号源。", registry_result={"registry_id": 11, "status": "created"}, source_payload={ "source_kind": "github", "source_name": "核心仓库", "source_identifier": "github:openai/codex", }, openclaw_result={"status": "completed", "result": {"trace_id": "openclaw-123"}}, worker_results=[], pending_workers=["pm_dispatcher", "tech_architect", "delivery_coach"], status="partial", run_status="running", ) await AIOrchestrationRunService(db).create_run( tenant_id=1, trace_id="trace-concurrent", run_type="signal_source_onboarding", project_id=3, signal_id=9, response_mode="complete", status="running", planner_status="planned", openclaw_status="completed", worker_total=3, worker_completed=0, worker_failed=0, summary="接入 GitHub 仓库信号源。", planner_diagnostics={"used_fallback": False}, result_payload=initial_result, created_by_user_id=1, ) await db.commit() fake_roles = { "roles": [ {"role_key": "pm_dispatcher", "role_name": "PM"}, {"role_key": "tech_architect", "role_name": "Arch"}, {"role_key": "delivery_coach", "role_name": "Coach"}, ] } with patch("app.services.signal_source_orchestration_service.prompt_profile_service.get_profiles", AsyncMock(return_value=fake_roles)): started = time.perf_counter() final_result = await _ConcurrentWorkerSignalService._execute_workers_and_finalize( tenant_id=1, trace_id="trace-concurrent", project_payload={"id": 3, "title": "项目A", "name": "项目A", "description": "desc", "owner_id": 1}, signal_payload={"id": 9, "title": "新增 GitHub 信号源", "summary": "监控 issue", "signal_type": "新需求", "recommended_action": "enter_only", "content": "content"}, registry_result={ "registry_id": 11, "status": "created", "source_kind": "github", "source_name": "核心仓库", "source_identifier": "github:openai/codex", "description": "监控 issue 与 PR", }, openclaw_result={"status": "completed", "result": {"trace_id": "openclaw-123"}}, worker_role_keys=["pm_dispatcher", "tech_architect", "delivery_coach"], operator_user_payload={"id": 1, "name": "Boss", "username": "boss", "role": "BOSS"}, planner_status="planned", planner_diagnostics={"used_fallback": False}, summary="接入 GitHub 仓库信号源。", signal_id=None, ) elapsed = time.perf_counter() - started self.assertLess(elapsed, 0.45) self.assertEqual(final_result["status"], "completed") self.assertEqual(len(final_result["worker_executions"]), 3) async def test_mark_stale_runs_as_interrupted(self) -> None: async with self.session_factory() as db: service = AIOrchestrationRunService(db) await service.create_run( tenant_id=1, trace_id="trace-stale", run_type="signal_source_onboarding", project_id=3, signal_id=9, response_mode="fast", status="running", planner_status="planned", openclaw_status="completed", worker_total=3, worker_completed=1, worker_failed=0, summary="stale run", planner_diagnostics={"used_fallback": False}, result_payload={"pending_workers": ["pm_dispatcher"], "run_status": "running"}, created_by_user_id=1, started_at=_utc_now_naive() - timedelta(minutes=11), ) await db.commit() updated = await service.mark_stale_runs_interrupted(older_than_minutes=10) await db.commit() row = await service.get_run(tenant_id=1, trace_id="trace-stale") self.assertEqual(updated, 1) self.assertEqual(row.status, "interrupted") if __name__ == "__main__": unittest.main()