Files
work-queue-api/app/routers/work.py
Steve W fbc88bb62b
Some checks failed
ci / build-test-push (push) Failing after 1m42s
feat: rebuild work queue api with fastapi and postgres
2026-04-11 19:24:52 +00:00

189 lines
7.3 KiB
Python

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)