feat: extend operational investigation workflows
This commit is contained in:
@@ -3,6 +3,14 @@ from uuid import uuid4
|
||||
|
||||
from fastapi import HTTPException
|
||||
from psycopg.types.json import Jsonb
|
||||
from app.mcp.errors import CalculationCapacityExceeded
|
||||
|
||||
|
||||
# Explicit metadata projection prevents status polling from decoding large JSON results.
|
||||
JOB_COLUMNS = "job_id,project_id,owner_user_id,request_key,tool,parameters,binding,timeout_seconds,status,error,lease_generation,lease_owner,lease_expires_at,created_at,updated_at,started_at,finished_at,queue_expires_at,result_expires_at"
|
||||
|
||||
|
||||
PUBLIC_JOB_COLUMNS = "job_id,status,tool,error,created_at,updated_at,started_at,finished_at,queue_expires_at,result_expires_at"
|
||||
|
||||
|
||||
class CalculationJobRepository:
|
||||
@@ -10,30 +18,65 @@ class CalculationJobRepository:
|
||||
self.connection = connection
|
||||
self.project_id = str(project_id)
|
||||
|
||||
async def enqueue(self, *, owner_user_id, request_key, tool, parameters, binding, timeout_seconds):
|
||||
cursor = await self.connection.execute(
|
||||
"""INSERT INTO workbench.calculation_jobs
|
||||
(job_id,project_id,owner_user_id,request_key,tool,parameters,binding,timeout_seconds)
|
||||
VALUES (%s,%s,%s,%s,%s,%s,%s,%s)
|
||||
ON CONFLICT(project_id,owner_user_id,request_key) DO NOTHING RETURNING *""",
|
||||
(uuid4(), self.project_id, owner_user_id, request_key, tool,
|
||||
Jsonb(parameters), Jsonb(binding), timeout_seconds),
|
||||
async def maintain(self):
|
||||
await self.connection.execute(
|
||||
"""UPDATE workbench.calculation_jobs SET status='timed_out',
|
||||
error='{"code":"QUEUE_TIMEOUT"}'::jsonb,finished_at=clock_timestamp(),updated_at=clock_timestamp()
|
||||
WHERE project_id=%s AND status='queued' AND queue_expires_at<=clock_timestamp()""",
|
||||
(self.project_id,),
|
||||
)
|
||||
row = await cursor.fetchone()
|
||||
if row:
|
||||
return dict(row)
|
||||
await self.connection.execute(
|
||||
"""UPDATE workbench.calculation_jobs SET result=NULL
|
||||
WHERE project_id=%s AND result IS NOT NULL AND result_expires_at<=clock_timestamp()""",
|
||||
(self.project_id,),
|
||||
)
|
||||
|
||||
async def enqueue(self, *, owner_user_id, request_key, tool, parameters, binding, timeout_seconds,
|
||||
queue_wait_seconds=3600, max_active_jobs=100):
|
||||
async with self.connection.transaction():
|
||||
return await self._enqueue(owner_user_id=owner_user_id, request_key=request_key,
|
||||
tool=tool, parameters=parameters, binding=binding, timeout_seconds=timeout_seconds,
|
||||
queue_wait_seconds=queue_wait_seconds, max_active_jobs=max_active_jobs)
|
||||
|
||||
async def _enqueue(self, *, owner_user_id, request_key, tool, parameters, binding,
|
||||
timeout_seconds, queue_wait_seconds, max_active_jobs):
|
||||
# Serialize admission within the caller's transaction, across every API process.
|
||||
await self.connection.execute(
|
||||
"SELECT pg_advisory_xact_lock(hashtextextended(%s, 1748531781))", (self.project_id,))
|
||||
await self.maintain()
|
||||
cursor = await self.connection.execute(
|
||||
"SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND owner_user_id=%s AND request_key=%s",
|
||||
f"SELECT {JOB_COLUMNS} FROM workbench.calculation_jobs WHERE project_id=%s AND owner_user_id=%s AND request_key=%s",
|
||||
(self.project_id, owner_user_id, request_key),
|
||||
)
|
||||
row = dict(await cursor.fetchone())
|
||||
if row['tool'] != tool or row['parameters'] != parameters or row['binding'] != binding:
|
||||
raise HTTPException(409, 'Calculation request_key was used with different inputs')
|
||||
return row
|
||||
|
||||
async def get(self, job_id, owner_user_id=None):
|
||||
existing = await cursor.fetchone()
|
||||
if existing:
|
||||
row = dict(existing)
|
||||
if (row['tool'] != tool
|
||||
or {key: value for key, value in row['parameters'].items() if key != 'timeout_seconds'}
|
||||
!= {key: value for key, value in parameters.items() if key != 'timeout_seconds'}
|
||||
or row['binding'] != binding):
|
||||
raise HTTPException(409, 'Calculation request_key was used with different inputs')
|
||||
return row
|
||||
cursor = await self.connection.execute(
|
||||
"""SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND job_id=%s
|
||||
"""SELECT count(*) AS active FROM workbench.calculation_jobs WHERE project_id=%s
|
||||
AND status IN ('queued','running','cancel_requested')""", (self.project_id,))
|
||||
if (await cursor.fetchone())['active'] >= max_active_jobs:
|
||||
raise CalculationCapacityExceeded()
|
||||
cursor = await self.connection.execute(
|
||||
f"""INSERT INTO workbench.calculation_jobs
|
||||
(job_id,project_id,owner_user_id,request_key,tool,parameters,binding,timeout_seconds,queue_expires_at)
|
||||
VALUES (%s,%s,%s,%s,%s,%s,%s,%s,clock_timestamp()+%s*interval '1 second')
|
||||
RETURNING {JOB_COLUMNS}""",
|
||||
(uuid4(), self.project_id, owner_user_id, request_key, tool,
|
||||
Jsonb(parameters), Jsonb(binding), timeout_seconds, queue_wait_seconds),
|
||||
)
|
||||
return dict(await cursor.fetchone())
|
||||
|
||||
async def get(self, job_id, owner_user_id=None, *, include_result=False):
|
||||
await self.maintain()
|
||||
columns = JOB_COLUMNS + (",result" if include_result else "")
|
||||
cursor = await self.connection.execute(
|
||||
f"""SELECT {columns} FROM workbench.calculation_jobs WHERE project_id=%s AND job_id=%s
|
||||
AND (%s::uuid IS NULL OR owner_user_id=%s)""",
|
||||
(self.project_id, job_id, owner_user_id, owner_user_id),
|
||||
)
|
||||
@@ -41,8 +84,9 @@ class CalculationJobRepository:
|
||||
return dict(row) if row else None
|
||||
|
||||
async def list(self, *, owner_user_id, binding, limit, offset, status=None):
|
||||
await self.maintain()
|
||||
cursor = await self.connection.execute(
|
||||
"""SELECT * FROM workbench.calculation_jobs WHERE project_id=%s AND owner_user_id=%s
|
||||
f"""SELECT {PUBLIC_JOB_COLUMNS} FROM workbench.calculation_jobs WHERE project_id=%s AND owner_user_id=%s
|
||||
AND binding @> %s AND (%s::text IS NULL OR status=%s) ORDER BY created_at DESC,job_id LIMIT %s OFFSET %s""",
|
||||
(self.project_id, owner_user_id, Jsonb(binding), status, status, limit, offset),
|
||||
)
|
||||
@@ -51,28 +95,30 @@ class CalculationJobRepository:
|
||||
async def cancel(self, job_id, owner_user_id):
|
||||
await self.connection.execute(
|
||||
"""UPDATE workbench.calculation_jobs SET status=CASE WHEN status='queued' THEN 'cancelled'
|
||||
ELSE 'cancel_requested' END,updated_at=clock_timestamp()
|
||||
ELSE 'cancel_requested' END,finished_at=CASE WHEN status='queued' THEN clock_timestamp() ELSE finished_at END,updated_at=clock_timestamp()
|
||||
WHERE project_id=%s AND job_id=%s AND owner_user_id=%s AND status IN ('queued','running')""",
|
||||
(self.project_id, job_id, owner_user_id),
|
||||
)
|
||||
return await self.get(job_id, owner_user_id)
|
||||
|
||||
async def recover(self):
|
||||
await self.maintain()
|
||||
# A crashed write is never automatically replayed.
|
||||
await self.connection.execute(
|
||||
"""UPDATE workbench.calculation_jobs SET status='interrupted',
|
||||
error='{"code":"WORKER_INTERRUPTED"}'::jsonb,updated_at=clock_timestamp()
|
||||
error='{"code":"WORKER_INTERRUPTED"}'::jsonb,finished_at=clock_timestamp(),updated_at=clock_timestamp()
|
||||
WHERE project_id=%s AND status IN ('running','cancel_requested')
|
||||
AND lease_expires_at < clock_timestamp()""", (self.project_id,),
|
||||
)
|
||||
|
||||
async def claim(self, owner, lease_seconds):
|
||||
await self.maintain()
|
||||
cursor = await self.connection.execute(
|
||||
"""UPDATE workbench.calculation_jobs j SET status='running',lease_owner=%s,
|
||||
"""UPDATE workbench.calculation_jobs j SET status='running',started_at=clock_timestamp(),lease_owner=%s,
|
||||
lease_generation=j.lease_generation+1,
|
||||
lease_expires_at=clock_timestamp()+%s*interval '1 second',updated_at=clock_timestamp()
|
||||
WHERE j.job_id=(SELECT job_id FROM workbench.calculation_jobs
|
||||
WHERE project_id=%s AND status='queued' ORDER BY created_at,job_id
|
||||
WHERE project_id=%s AND status='queued' AND queue_expires_at>clock_timestamp() ORDER BY created_at,job_id
|
||||
FOR UPDATE SKIP LOCKED LIMIT 1) RETURNING j.*""",
|
||||
(owner, lease_seconds, self.project_id),
|
||||
)
|
||||
@@ -90,14 +136,15 @@ class CalculationJobRepository:
|
||||
result = await cursor.fetchone()
|
||||
return result['status'] if result else None
|
||||
|
||||
async def finish(self, row, owner, status, *, result=None, error=None):
|
||||
async def finish(self, row, owner, status, *, result=None, error=None, result_retention_seconds=604800):
|
||||
cursor = await self.connection.execute(
|
||||
"""UPDATE workbench.calculation_jobs SET status=%s,result=%s,error=%s,updated_at=clock_timestamp()
|
||||
"""UPDATE workbench.calculation_jobs SET status=%s,result=%s,error=%s,finished_at=clock_timestamp(),
|
||||
result_expires_at=clock_timestamp()+%s*interval '1 second',updated_at=clock_timestamp()
|
||||
WHERE project_id=%s AND job_id=%s AND lease_owner=%s AND lease_generation=%s
|
||||
AND lease_expires_at>clock_timestamp() AND status IN ('running','cancel_requested')
|
||||
AND (status!='cancel_requested' OR %s IN ('cancelled','interrupted')
|
||||
OR (%s='failed' AND %s='CLEANUP_FAILED')) RETURNING job_id""",
|
||||
(status, Jsonb(result) if result is not None else None, Jsonb(error) if error else None,
|
||||
(status, Jsonb(result) if result is not None else None, Jsonb(error) if error else None, result_retention_seconds,
|
||||
self.project_id, row['job_id'], owner, row['lease_generation'], status, status, (error or {}).get('code')),
|
||||
)
|
||||
return await cursor.fetchone() is not None
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -337,6 +337,7 @@ async def task_context_page(
|
||||
"evidence": ("operational_task_evidence", "evidence_id"),
|
||||
"actions": ("operational_task_actions", "action_id"),
|
||||
"access_requests": ("operational_task_access_requests", "access_request_id"),
|
||||
"questions": ("operational_task_questions", "question_request_id"),
|
||||
}
|
||||
documents = {"task": "operational_tasks", "report": "operational_task_reports", "plan": "operational_task_plans"}
|
||||
if section in {"access_requests", "progress"}:
|
||||
|
||||
Reference in New Issue
Block a user