a2a-source-edit: write agent.py

agent.py

@@ -514,7 +514,10 @@ def _connect_db() -> Any:
 
 
 async def _save_audit_and_receipt(tenant: str, audit_id: str, payload: dict[str, Any], receipt: dict[str, Any]) -> None:
-    def work() -> None:
+    await asyncio.to_thread(_save_audit_and_receipt_sync, tenant, audit_id, payload, receipt)
+
+
+def _save_audit_and_receipt_sync(tenant: str, audit_id: str, payload: dict[str, Any], receipt: dict[str, Any]) -> None:
     from psycopg.types.json import Jsonb
 
     with _connect_db() as conn:
@@ -543,7 +546,6 @@ async def _save_audit_and_receipt(tenant: str, audit_id: str, payload: dict[str,
                     """,
                     (tenant, audit_id, receipt["receipt_id"], Jsonb(receipt)),
                 )
-    await asyncio.to_thread(work)
 
 
 async def _persist_failure_if_possible(tenant: str, audit_id: str, target_url: str, payload: dict[str, Any], started_at: str) -> None:
@@ -555,9 +557,14 @@ async def _persist_failure_if_possible(tenant: str, audit_id: str, target_url: s
         "created_at": _now_iso(),
         "code": payload.get("code"),
     }
-    payload = {**payload, "receipt": receipt, "created_at": started_at, "updated_at": _now_iso()}
+    saved_payload = {**payload, "receipt": receipt, "created_at": started_at, "updated_at": _now_iso()}
+    try:
+        await asyncio.to_thread(_persist_failure_sync, tenant, audit_id, target_url, saved_payload, receipt)
+    except Exception:  # noqa: BLE001
+        return
 
-    def work() -> None:
+
+def _persist_failure_sync(tenant: str, audit_id: str, target_url: str, payload: dict[str, Any], receipt: dict[str, Any]) -> None:
     from psycopg.types.json import Jsonb
 
     with _connect_db() as conn:
@@ -584,14 +591,18 @@ async def _persist_failure_if_possible(tenant: str, audit_id: str, target_url: s
                     """,
                     (tenant, audit_id, receipt["receipt_id"], Jsonb(receipt)),
                 )
-    try:
-        await asyncio.to_thread(work)
-    except Exception:  # noqa: BLE001
-        return
 
 
 async def _load_audit(tenant: str, audit_id: str) -> dict[str, Any] | None:
-    def work() -> dict[str, Any] | None:
+    try:
+        return await asyncio.to_thread(_load_audit_sync, tenant, audit_id)
+    except RuntimeError as exc:
+        if str(exc) == "database_unavailable":
+            return {"ok": False, "code": "database_unavailable", "message": "Persistent storage is not configured for this deployment.", "audit_id": audit_id}
+        raise
+
+
+def _load_audit_sync(tenant: str, audit_id: str) -> dict[str, Any] | None:
     with _connect_db() as conn:
         with conn.cursor() as cur:
             cur.execute(
@@ -600,9 +611,3 @@ async def _load_audit(tenant: str, audit_id: str) -> dict[str, Any] | None:
             )
             row = cur.fetchone()
             return row[0] if row else None
-    try:
-        return await asyncio.to_thread(work)
-    except RuntimeError as exc:
-        if str(exc) == "database_unavailable":
-            return {"ok": False, "code": "database_unavailable", "message": "Persistent storage is not configured for this deployment.", "audit_id": audit_id}
-        raise