from __future__ import annotations import json from uuid import uuid4 from fastapi import APIRouter, HTTPException, Response from psycopg.errors import UniqueViolation from app.db import ( TERMINAL_STATUSES, complete_dispatch_log, create_dispatch_log, fetch_dispatch_log, fetch_project_or_404, fetch_work_or_404, get_conn, utcnow, validate_transition, ) from app.models import DispatchLog, WorkCreate, WorkItem, WorkUpdate router = APIRouter(prefix="/work", tags=["work"]) @router.post("", response_model=WorkItem, status_code=201) def create_work(payload: WorkCreate) -> WorkItem: if not payload.type.strip() or not payload.description.strip(): raise HTTPException(status_code=400, detail="type and description are required") now = utcnow() record = { "id": str(uuid4()), "project_id": str(payload.project_id) if payload.project_id else None, "type": payload.type.strip(), "description": payload.description.strip(), "payload": payload.payload, "priority": payload.priority, "status": "queued", "assigned_agent": payload.assigned_agent, "created_by": payload.created_by, "created_at": now, "updated_at": now, "completed_at": None, "outcome": None, "notes": None, "dispatch_log": [], } with get_conn() as conn: if record["project_id"]: fetch_project_or_404(conn, record["project_id"]) conn.execute( """ INSERT INTO work_items ( id, project_id, type, description, payload, priority, status, assigned_agent, created_by, created_at, updated_at, completed_at, outcome, notes ) VALUES (%s, %s, %s, %s, %s::jsonb, %s, %s, %s, %s, %s, %s, %s, %s, %s) """, ( record["id"], record["project_id"], record["type"], record["description"], json.dumps(record["payload"]) if record["payload"] is not None else None, record["priority"], record["status"], record["assigned_agent"], record["created_by"], record["created_at"], record["updated_at"], None, None, None, ), ) conn.commit() return WorkItem.model_validate(record) @router.get("", response_model=list[WorkItem]) def list_work(status: str | None = None, agent: str | None = None, project_id: str | None = None, since: str | None = None) -> list[WorkItem]: clauses: list[str] = [] params: list[object] = [] if status: clauses.append("status = %s") params.append(status) if agent: clauses.append("assigned_agent = %s") params.append(agent) if project_id: clauses.append("project_id = %s") params.append(project_id) if since: clauses.append("created_at > %s::timestamptz") params.append(since) query = """ SELECT id, project_id, type, description, payload, priority, status, assigned_agent, created_by, created_at, updated_at, completed_at, outcome, notes FROM work_items """ if clauses: query += " WHERE " + " AND ".join(clauses) query += " ORDER BY created_at ASC" if since else " ORDER BY priority ASC, created_at ASC" with get_conn() as conn: rows = conn.execute(query, tuple(params)).fetchall() return [WorkItem.model_validate({**row, "dispatch_log": []}) for row in rows] @router.get("/{work_id}", response_model=WorkItem) def get_work(work_id: str) -> WorkItem: with get_conn() as conn: row = fetch_work_or_404(conn, work_id) row["dispatch_log"] = fetch_dispatch_log(conn, work_id) return WorkItem.model_validate(row) @router.patch("/{work_id}", response_model=WorkItem) def update_work(work_id: str, payload: WorkUpdate) -> WorkItem: with get_conn() as conn: current = fetch_work_or_404(conn, work_id) if payload.assigned_agent is not None: current["assigned_agent"] = payload.assigned_agent if payload.notes is not None: current["notes"] = payload.notes requested_status = payload.status if requested_status is not None: validate_transition(current["status"], requested_status, current.get("assigned_agent"), payload.outcome, payload.notes) current["status"] = requested_status if requested_status in TERMINAL_STATUSES: current["completed_at"] = utcnow() current["outcome"] = payload.outcome else: current["completed_at"] = None current["outcome"] = None current["updated_at"] = utcnow() try: conn.execute( """ UPDATE work_items SET project_id = %s, type = %s, description = %s, payload = %s::jsonb, priority = %s, status = %s, assigned_agent = %s, created_by = %s, created_at = %s, updated_at = %s, completed_at = %s, outcome = %s, notes = %s WHERE id = %s """, ( current["project_id"], current["type"], current["description"], json.dumps(current["payload"]) if current["payload"] is not None else None, current["priority"], current["status"], current["assigned_agent"], current["created_by"], current["created_at"], current["updated_at"], current["completed_at"], current["outcome"], current["notes"], work_id, ), ) if requested_status == "dispatched": create_dispatch_log(conn, work_id, current["assigned_agent"], current["updated_at"]) if requested_status in TERMINAL_STATUSES: complete_dispatch_log(conn, work_id, current["completed_at"], current["outcome"]) conn.commit() except UniqueViolation: conn.rollback() raise HTTPException(status_code=409, detail="agent already has an in_progress work item") from None current["dispatch_log"] = fetch_dispatch_log(conn, work_id) return WorkItem.model_validate(current) @router.delete("/{work_id}", status_code=204) def delete_work(work_id: str) -> Response: with get_conn() as conn: current = fetch_work_or_404(conn, work_id) if current["status"] not in {"queued", "dispatched"}: raise HTTPException(status_code=400, detail="only queued or dispatched items can be cancelled") current["status"] = "cancelled" current["updated_at"] = utcnow() current["completed_at"] = utcnow() current["outcome"] = "cancelled" conn.execute( "UPDATE work_items SET status = %s, updated_at = %s, completed_at = %s, outcome = %s WHERE id = %s", (current["status"], current["updated_at"], current["completed_at"], current["outcome"], work_id), ) complete_dispatch_log(conn, work_id, current["completed_at"], current["outcome"]) conn.commit() return Response(status_code=204)