feat: extend operational investigation workflows

2026-10-30 17:30:19 +08:00
parent 63af998045
commit 05183aff93
82 changed files with 4644 additions and 330 deletions
@@ -19,6 +19,7 @@ class OperationalTaskRepository:
    def __init__(self, connection: AsyncConnection, project_id):
        self.connection = connection
        self.project_id = str(project_id)
+        self._agent_ownership_available: bool | None = None

    @asynccontextmanager
    async def transaction(self):
@@ -89,6 +90,34 @@ class OperationalTaskRepository:
        )
        return [str(row["job_id"]) for row in await cursor.fetchall()]

+    async def list_questions(self, task_id: str) -> list[dict]:
+        cursor = await self.connection.execute(
+            """SELECT document FROM workbench.operational_task_questions
+               WHERE project_id=%s AND task_id=%s
+               ORDER BY document->>'created_at' DESC, question_request_id DESC LIMIT 100""",
+            (self.project_id, task_id),
+        )
+        return [row["document"] for row in await cursor.fetchall()]
+
+    async def get_question(self, task_id: str, question_id: str) -> dict | None:
+        cursor = await self.connection.execute(
+            """SELECT document FROM workbench.operational_task_questions
+               WHERE project_id=%s AND task_id=%s AND question_request_id=%s""",
+            (self.project_id, task_id, question_id),
+        )
+        row = await cursor.fetchone()
+        return row["document"] if row else None
+
+    async def save_question(self, task_id: str, document: dict) -> None:
+        # All writers hold the owning task advisory lock.
+        await self.connection.execute(
+            """INSERT INTO workbench.operational_task_questions
+               (project_id,task_id,question_request_id,document) VALUES (%s,%s,%s,%s)
+               ON CONFLICT (project_id,question_request_id) DO UPDATE SET document=excluded.document
+               WHERE operational_task_questions.task_id=excluded.task_id""",
+            (self.project_id, task_id, document["question_request_id"], Jsonb(document)),
+        )
+
    async def get_task(self, task_id: str) -> dict | None:
        cursor = await self.connection.execute(
            """SELECT document FROM workbench.operational_tasks
@@ -101,7 +130,8 @@ class OperationalTaskRepository:
    async def get_active_run(self, task_id: str) -> dict | None:
        cursor = await self.connection.execute(
            """SELECT run_id,task_id,generation,execution_request_id,status,revision,
-                      agent_session_id,lease_expires_at,heartbeat_at,started_at,finished_at,
+                      agent_session_id,to_jsonb(operational_task_runs)->>'agent_instance_id' AS agent_instance_id,
+                      lease_expires_at,heartbeat_at,started_at,finished_at,
                      failure_code,failure_detail,dispatch_attempts,next_dispatch_at,
                      created_at,updated_at
               FROM workbench.operational_task_runs
@@ -116,7 +146,8 @@ class OperationalTaskRepository:
    async def get_run(self, run_id: str) -> dict | None:
        cursor = await self.connection.execute(
            """SELECT run_id,task_id,generation,execution_request_id,status,revision,
-                      agent_session_id,lease_expires_at,heartbeat_at,started_at,finished_at,
+                      agent_session_id,to_jsonb(operational_task_runs)->>'agent_instance_id' AS agent_instance_id,
+                      lease_expires_at,heartbeat_at,started_at,finished_at,
                      failure_code,failure_detail,dispatch_attempts,next_dispatch_at,
                      created_at,updated_at
               FROM workbench.operational_task_runs
@@ -143,13 +174,23 @@ class OperationalTaskRepository:
    async def require_execution_storage(self) -> None:
        """Reject claims before changing state if the inbox cannot be used."""
        cursor = await self.connection.execute(
-            "SELECT to_regclass('workbench.agent_messages') IS NOT NULL AS available"
+            """SELECT to_regclass('workbench.agent_messages') IS NOT NULL
+                      AND to_regclass('workbench.operational_task_questions') IS NOT NULL AS available,
+                      EXISTS (
+                        SELECT 1 FROM pg_attribute
+                        WHERE attrelid=to_regclass('workbench.operational_task_runs')
+                          AND attname='agent_instance_id' AND NOT attisdropped
+                      ) AS agent_ownership_available"""
        )
-        if not (await cursor.fetchone())["available"]:
+        readiness = await cursor.fetchone()
+        self._agent_ownership_available = bool(readiness.get("agent_ownership_available"))
+        if not readiness["available"] or not self._agent_ownership_available:
            raise DomainError(
                "DEPENDENCY_UNAVAILABLE",
-                "Agent message storage is not provisioned; apply "
-                "resources/sql/024_agent_messages.sql to this project database before starting runs",
+                "Agent execution storage is not provisioned; apply "
+                "resources/sql/024_agent_messages.sql, resources/sql/028_task_questions.sql, "
+                "and resources/sql/030_operational_run_agent_ownership.sql "
+                "to this project database before starting runs",
                status=503,
            )

@@ -190,14 +231,26 @@ class OperationalTaskRepository:
        )

    async def update_run(self, run: dict, expected_revision: int) -> None:
+        if self._agent_ownership_available is None:
+            cursor = await self.connection.execute(
+                """SELECT EXISTS (
+                     SELECT 1 FROM pg_attribute
+                     WHERE attrelid=to_regclass('workbench.operational_task_runs')
+                       AND attname='agent_instance_id' AND NOT attisdropped
+                   ) AS available"""
+            )
+            self._agent_ownership_available = bool((await cursor.fetchone())["available"])
+        owner_assignment = "agent_instance_id=%s," if self._agent_ownership_available else ""
+        owner_values = (run.get("agent_instance_id"),) if self._agent_ownership_available else ()
        cursor = await self.connection.execute(
-            """UPDATE workbench.operational_task_runs SET
-                 status=%s,revision=%s,agent_session_id=%s,lease_expires_at=%s,
+            f"""UPDATE workbench.operational_task_runs SET
+                 status=%s,revision=%s,agent_session_id=%s,{owner_assignment}lease_expires_at=%s,
                 heartbeat_at=%s,started_at=%s,finished_at=%s,failure_code=%s,
                 failure_detail=%s,dispatch_attempts=%s,next_dispatch_at=%s,updated_at=%s
               WHERE project_id=%s AND run_id=%s AND revision=%s""",
            (
                run["status"], run["revision"], run.get("agent_session_id"),
+                *owner_values,
                run.get("lease_expires_at"), run.get("heartbeat_at"),
                run.get("started_at"), run.get("finished_at"), run.get("failure_code"),
                run.get("failure_detail"), run.get("dispatch_attempts", 0),
@@ -214,7 +267,8 @@ class OperationalTaskRepository:
    ) -> list[dict]:
        result = await self.connection.execute(
            """SELECT run_id,task_id,generation,execution_request_id,status,revision,
-                      agent_session_id,lease_expires_at,heartbeat_at,started_at,finished_at,
+                      agent_session_id,to_jsonb(operational_task_runs)->>'agent_instance_id' AS agent_instance_id,
+                      lease_expires_at,heartbeat_at,started_at,finished_at,
                      failure_code,failure_detail,dispatch_attempts,next_dispatch_at,
                      created_at,updated_at
               FROM workbench.operational_task_runs
@@ -243,7 +297,8 @@ class OperationalTaskRepository:
    ) -> list[dict]:
        cursor = await self.connection.execute(
            """SELECT run_id,task_id,generation,execution_request_id,status,revision,
-                      agent_session_id,lease_expires_at,heartbeat_at,started_at,finished_at,
+                      agent_session_id,to_jsonb(operational_task_runs)->>'agent_instance_id' AS agent_instance_id,
+                      lease_expires_at,heartbeat_at,started_at,finished_at,
                      failure_code,failure_detail,dispatch_attempts,next_dispatch_at,
                      created_at,updated_at
               FROM workbench.operational_task_runs
@@ -300,6 +355,18 @@ class OperationalTaskRepository:
    async def current_event_cursor(self) -> int:
        return await WorkbenchEventRepository(self.connection, self.project_id).current_cursor()

+    async def list_presentation_events(self, task_id: str, after: int, through: int, limit: int) -> list[dict]:
+        cursor = await self.connection.execute(
+            """SELECT event_seq,occurred_at,metadata
+               FROM workbench.domain_event_log
+               WHERE project_id=%s AND topic='operational.presentation.requested'
+                 AND aggregate_type='operational_task_presentation' AND aggregate_id=%s
+                 AND event_seq>%s AND event_seq<=%s
+               ORDER BY event_seq LIMIT %s""",
+            (self.project_id, task_id, after, through, limit),
+        )
+        return [dict(row) for row in await cursor.fetchall()]
+
    async def get_task_by_create_request(self, request_id: str) -> dict | None:
        cursor = await self.connection.execute(
            """SELECT document FROM workbench.operational_tasks