From 8aae6c3b8fb9d444d169c762f9837854bf0518ac Mon Sep 17 00:00:00 2001 From: a2a-cloud Date: Sat, 18 Jul 2026 09:27:08 +0000 Subject: [PATCH] a2a-source-edit: write agent.py --- agent.py | 139 ++++++++++++++++++++++++++++--------------------------- 1 file changed, 72 insertions(+), 67 deletions(-) diff --git a/agent.py b/agent.py index 76f9cde..4ab6f23 100644 --- a/agent.py +++ b/agent.py @@ -514,36 +514,38 @@ 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: - from psycopg.types.json import Jsonb + await asyncio.to_thread(_save_audit_and_receipt_sync, tenant, audit_id, payload, receipt) - with _connect_db() as conn: - with conn.transaction(): - with conn.cursor() as cur: - cur.execute( - """ - INSERT INTO launch_audits (tenant_key, audit_id, target_url, ok, report, payload) - VALUES (%s, %s, %s, %s, %s, %s) - ON CONFLICT (tenant_key, audit_id) DO UPDATE SET - target_url = EXCLUDED.target_url, - ok = EXCLUDED.ok, - report = EXCLUDED.report, - payload = EXCLUDED.payload, - updated_at = NOW() - """, - (tenant, audit_id, payload["target_url"], True, Jsonb(payload["report"]), Jsonb(payload)), - ) - cur.execute( - """ - INSERT INTO launch_receipts (tenant_key, audit_id, receipt_id, payload) - VALUES (%s, %s, %s, %s) - ON CONFLICT (tenant_key, receipt_id) DO UPDATE SET - payload = EXCLUDED.payload, - created_at = NOW() - """, - (tenant, audit_id, receipt["receipt_id"], Jsonb(receipt)), - ) - await asyncio.to_thread(work) + +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: + with conn.transaction(): + with conn.cursor() as cur: + cur.execute( + """ + INSERT INTO launch_audits (tenant_key, audit_id, target_url, ok, report, payload) + VALUES (%s, %s, %s, %s, %s, %s) + ON CONFLICT (tenant_key, audit_id) DO UPDATE SET + target_url = EXCLUDED.target_url, + ok = EXCLUDED.ok, + report = EXCLUDED.report, + payload = EXCLUDED.payload, + updated_at = NOW() + """, + (tenant, audit_id, payload["target_url"], True, Jsonb(payload["report"]), Jsonb(payload)), + ) + cur.execute( + """ + INSERT INTO launch_receipts (tenant_key, audit_id, receipt_id, payload) + VALUES (%s, %s, %s, %s) + ON CONFLICT (tenant_key, receipt_id) DO UPDATE SET + payload = EXCLUDED.payload, + created_at = NOW() + """, + (tenant, audit_id, receipt["receipt_id"], Jsonb(receipt)), + ) async def _persist_failure_if_possible(tenant: str, audit_id: str, target_url: str, payload: dict[str, Any], started_at: str) -> None: @@ -555,54 +557,57 @@ 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()} - - def work() -> None: - from psycopg.types.json import Jsonb - - with _connect_db() as conn: - with conn.transaction(): - with conn.cursor() as cur: - cur.execute( - """ - INSERT INTO launch_audits (tenant_key, audit_id, target_url, ok, report, payload) - VALUES (%s, %s, %s, %s, %s, %s) - ON CONFLICT (tenant_key, audit_id) DO UPDATE SET - target_url = EXCLUDED.target_url, - ok = EXCLUDED.ok, - report = EXCLUDED.report, - payload = EXCLUDED.payload, - updated_at = NOW() - """, - (tenant, audit_id, _safe_url_echo(target_url), False, Jsonb({}), Jsonb(payload)), - ) - cur.execute( - """ - INSERT INTO launch_receipts (tenant_key, audit_id, receipt_id, payload) - VALUES (%s, %s, %s, %s) - ON CONFLICT (tenant_key, receipt_id) DO UPDATE SET payload = EXCLUDED.payload - """, - (tenant, audit_id, receipt["receipt_id"], Jsonb(receipt)), - ) + saved_payload = {**payload, "receipt": receipt, "created_at": started_at, "updated_at": _now_iso()} try: - await asyncio.to_thread(work) + await asyncio.to_thread(_persist_failure_sync, tenant, audit_id, target_url, saved_payload, receipt) except Exception: # noqa: BLE001 return -async def _load_audit(tenant: str, audit_id: str) -> dict[str, Any] | None: - def work() -> dict[str, Any] | None: - with _connect_db() as conn: +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: + with conn.transaction(): with conn.cursor() as cur: cur.execute( - "SELECT payload FROM launch_audits WHERE tenant_key = %s AND audit_id = %s", - (tenant, audit_id), + """ + INSERT INTO launch_audits (tenant_key, audit_id, target_url, ok, report, payload) + VALUES (%s, %s, %s, %s, %s, %s) + ON CONFLICT (tenant_key, audit_id) DO UPDATE SET + target_url = EXCLUDED.target_url, + ok = EXCLUDED.ok, + report = EXCLUDED.report, + payload = EXCLUDED.payload, + updated_at = NOW() + """, + (tenant, audit_id, _safe_url_echo(target_url), False, Jsonb({}), Jsonb(payload)), ) - row = cur.fetchone() - return row[0] if row else None + cur.execute( + """ + INSERT INTO launch_receipts (tenant_key, audit_id, receipt_id, payload) + VALUES (%s, %s, %s, %s) + ON CONFLICT (tenant_key, receipt_id) DO UPDATE SET payload = EXCLUDED.payload + """, + (tenant, audit_id, receipt["receipt_id"], Jsonb(receipt)), + ) + + +async def _load_audit(tenant: str, audit_id: str) -> dict[str, Any] | None: try: - return await asyncio.to_thread(work) + 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( + "SELECT payload FROM launch_audits WHERE tenant_key = %s AND audit_id = %s", + (tenant, audit_id), + ) + row = cur.fetchone() + return row[0] if row else None