chore: initialize project repository

2026-09-09 21:55:37 +08:00
commit b1497fd2c6
167 changed files with 17690 additions and 0 deletions
@@ -0,0 +1,184 @@
+import asyncio
+import contextlib
+import logging
+from collections.abc import Awaitable, Callable
+from datetime import UTC, datetime, timedelta
+from typing import Any
+from uuid import uuid4
+
+from sqlalchemy import select, update
+
+from app.core.config import Settings, get_settings
+from app.core.contracts import AgentRequest, AgentResult, RequestContext
+from app.core.errors import AgentError, RecoverableAgentError, RunLeaseLostError
+from app.infrastructure.db import SessionFactory
+from app.model.audit import InteractionAudit
+from app.model.conversation import ConversationMessage
+from app.model.platform import AgentRun, DomainEventOutbox, RequestIdempotency
+from app.repository.agent_run_repository import AgentRunRepository
+from app.service.agent.bootstrap import get_agent_factory
+from app.service.agent.executor import AgentExecutor
+from app.service.agent.factory import AgentFactory
+from app.service.agent_persistence_service import AgentPersistenceService
+from app.service.identity_service import IdentityService
+from app.service.memory_service import MemoryService
+from app.worker.outbox_worker import OutboxWorker
+
+logger = logging.getLogger(__name__)
+
+
+class WorkerRuntime:
+    def __init__(
+        self, factory: AgentFactory | None = None, settings: Settings | None = None,
+        resolve_identity: Callable[[RequestContext], Awaitable[RequestContext]] | None = None,
+    ) -> None:
+        self.factory = factory if factory is not None else get_agent_factory()
+        self.settings = settings or get_settings()
+        self.resolve_identity = resolve_identity or IdentityService().resolve
+
+    async def dispatch_one(self, *, run_id: str | None = None) -> bool:
+        # Outbox acknowledges a durable SQL queue entry, not an in-memory task.
+        async with SessionFactory() as session:
+            async def dispatch(payload: dict[str, Any]) -> None:
+                run = await AgentRunRepository(session).get(str(payload["run_id"]))
+                if run is None:
+                    raise ValueError("run not found")
+
+            return await OutboxWorker(session, {"agent.run_requested": dispatch}).publish_one(
+                aggregate_id=run_id)
+
+    async def run_once(self) -> bool:
+        dispatched = await self.dispatch_one()
+        now = datetime.now(UTC).replace(tzinfo=None)
+        async with SessionFactory() as session:
+            run_id = await session.scalar(select(AgentRun.run_id).where(
+                AgentRun.status.in_(("queued", "running", "cancel_requested")),
+                (AgentRun.locked_until.is_(None) | (AgentRun.locked_until < now)),
+            ).order_by(AgentRun.created_at).limit(1))
+        if run_id is None:
+            return dispatched
+        return await self.execute(run_id) or dispatched
+
+    async def execute(self, run_id: str) -> bool:
+        # A new fencing token for each claim also fences restarts of the same process.
+        worker_id = str(uuid4())
+        async with SessionFactory() as session, session.begin():
+            run = await session.scalar(select(AgentRun).where(
+                AgentRun.run_id == run_id).with_for_update())
+            if run is None:
+                return False
+            if run.status == "cancel_requested":
+                run.status = "cancelled"
+                run.completed_at = datetime.now(UTC).replace(tzinfo=None)
+                run.locked_until = None
+                return True
+            claimed = await AgentRunRepository(session).claim(
+                run_id, worker_id, self.settings.worker_lease_seconds)
+            if not claimed:
+                return False
+        task = asyncio.create_task(self._execute_claimed(run_id, worker_id))
+        heartbeat = asyncio.create_task(self._heartbeat(run_id, worker_id, task))
+        try:
+            await task
+        except asyncio.CancelledError:
+            await self._failure(run_id, worker_id, "RUN_INTERRUPTED", retryable=True)
+            # Runtime cancellation is shutdown; a lost lease only cancels the child.
+            current = asyncio.current_task()
+            if current is not None and current.cancelling():
+                raise
+        except RunLeaseLostError:
+            await self._failure(run_id, worker_id, "RUN_LEASE_LOST", retryable=True)
+        except Exception as exc:
+            code = exc.code if isinstance(exc, AgentError) else "AGENT_INTERNAL_ERROR"
+            await self._failure(run_id, worker_id, code,
+                                retryable=isinstance(exc, RecoverableAgentError))
+            logger.warning("run failed run_id=%s error_type=%s", run_id, type(exc).__name__)
+        finally:
+            heartbeat.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await heartbeat
+        return True
+
+    async def _heartbeat(
+        self, run_id: str, worker_id: str, task: asyncio.Task[None]
+    ) -> None:
+        try:
+            while True:
+                await asyncio.sleep(self.settings.worker_lease_seconds / 3)
+                async with SessionFactory() as session, session.begin():
+                    renewed = await AgentRunRepository(session).renew(
+                        run_id, worker_id, self.settings.worker_lease_seconds)
+                if not renewed:
+                    task.cancel()
+                    return
+        except Exception:
+            task.cancel()
+            logger.warning("lease renewal failed run_id=%s", run_id)
+
+    async def _execute_claimed(self, run_id: str, worker_id: str) -> None:
+        async with SessionFactory() as session:
+            run = await AgentRunRepository(session).get(run_id)
+            if run is None:
+                raise ValueError("run not found")
+            message = await session.get(ConversationMessage, run.request_message_id)
+            idem = await session.get(RequestIdempotency, run.idempotency_id)
+            event = await session.scalar(select(DomainEventOutbox).where(
+                DomainEventOutbox.aggregate_id == run_id,
+                DomainEventOutbox.event_type == "agent.run_requested").limit(1))
+            if message is None or idem is None:
+                raise ValueError("run input missing")
+            request = AgentRequest(
+                agent_type=run.agent_type, message=message.content, session_id=run.session_id,
+                idempotency_key=idem.idempotency_key,
+                metadata=event.payload.get("metadata", {}) if event else {},
+            )
+            identity = RequestContext(user_id=str(run.user_id), trace_id=run.trace_id)
+        # Re-check account and permissions at execution time, including delayed jobs.
+        context = await self.resolve_identity(identity)
+        result: AgentResult | None = None
+        async for event_data in AgentExecutor(self.factory).execute(
+            request.agent_type, request, context, run_id
+        ):
+            if event_data.event_type == "done":
+                result = AgentResult.model_validate(event_data.payload["result"])
+        if result is None:
+            raise ValueError("Agent produced no terminal result")
+        async with SessionFactory() as session:
+            await AgentPersistenceService(session).complete_run(
+                run_id, result, worker_id=worker_id,
+                memory_extraction_requested=MemoryService.should_extract_memory(
+                    conversation_content=request.message, role="user"),
+            )
+
+    async def _failure(
+        self, run_id: str, worker_id: str, error_code: str, *, retryable: bool
+    ) -> None:
+        now = datetime.now(UTC).replace(tzinfo=None)
+        async with SessionFactory() as session, session.begin():
+            run = await session.scalar(select(AgentRun).where(
+                AgentRun.run_id == run_id).with_for_update())
+            if run is None or run.worker_id != worker_id:
+                return
+            if run.status not in {"running", "cancel_requested"}:
+                return
+            if run.status == "cancel_requested":
+                run.status = "cancelled"
+            elif retryable and run.attempt_count < self.settings.worker_retry_limit:
+                run.status = "queued"
+            else:
+                run.status = "failed"
+            run.error_code, run.updated_at = error_code, now
+            run.locked_until = (
+                now + timedelta(seconds=min(60, 2**run.attempt_count))
+                if run.status == "queued" else None
+            )
+            if run.status != "queued":
+                run.completed_at = now
+                await session.execute(update(RequestIdempotency).where(
+                    RequestIdempotency.id == run.idempotency_id
+                ).values(status="failed", error_code=error_code, updated_at=now))
+                session.add(InteractionAudit(
+                    actor_type="agent", actor_id=run.user_id, session_id=run.session_id,
+                    portal="api", action_type=f"agent.run_{run.status}",
+                    detail={"run_id": run_id, "error_code": error_code}, created_at=now,
+                ))