a2a-source-edit: write agent.py

This commit is contained in:
a2a-cloud
2026-07-18 09:27:08 +00:00
parent d2a22e762d
commit 8aae6c3b8f

139
agent.py
View File

@@ -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: 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)
from psycopg.types.json import Jsonb
with _connect_db() as conn:
with conn.transaction(): def _save_audit_and_receipt_sync(tenant: str, audit_id: str, payload: dict[str, Any], receipt: dict[str, Any]) -> None:
with conn.cursor() as cur: from psycopg.types.json import Jsonb
cur.execute(
""" with _connect_db() as conn:
INSERT INTO launch_audits (tenant_key, audit_id, target_url, ok, report, payload) with conn.transaction():
VALUES (%s, %s, %s, %s, %s, %s) with conn.cursor() as cur:
ON CONFLICT (tenant_key, audit_id) DO UPDATE SET cur.execute(
target_url = EXCLUDED.target_url, """
ok = EXCLUDED.ok, INSERT INTO launch_audits (tenant_key, audit_id, target_url, ok, report, payload)
report = EXCLUDED.report, VALUES (%s, %s, %s, %s, %s, %s)
payload = EXCLUDED.payload, ON CONFLICT (tenant_key, audit_id) DO UPDATE SET
updated_at = NOW() target_url = EXCLUDED.target_url,
""", ok = EXCLUDED.ok,
(tenant, audit_id, payload["target_url"], True, Jsonb(payload["report"]), Jsonb(payload)), report = EXCLUDED.report,
) payload = EXCLUDED.payload,
cur.execute( updated_at = NOW()
""" """,
INSERT INTO launch_receipts (tenant_key, audit_id, receipt_id, payload) (tenant, audit_id, payload["target_url"], True, Jsonb(payload["report"]), Jsonb(payload)),
VALUES (%s, %s, %s, %s) )
ON CONFLICT (tenant_key, receipt_id) DO UPDATE SET cur.execute(
payload = EXCLUDED.payload, """
created_at = NOW() INSERT INTO launch_receipts (tenant_key, audit_id, receipt_id, payload)
""", VALUES (%s, %s, %s, %s)
(tenant, audit_id, receipt["receipt_id"], Jsonb(receipt)), ON CONFLICT (tenant_key, receipt_id) DO UPDATE SET
) payload = EXCLUDED.payload,
await asyncio.to_thread(work) 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: 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(), "created_at": _now_iso(),
"code": payload.get("code"), "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()}
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)),
)
try: 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 except Exception: # noqa: BLE001
return return
async def _load_audit(tenant: str, audit_id: str) -> dict[str, Any] | None: def _persist_failure_sync(tenant: str, audit_id: str, target_url: str, payload: dict[str, Any], receipt: dict[str, Any]) -> None:
def work() -> dict[str, Any] | None: from psycopg.types.json import Jsonb
with _connect_db() as conn:
with _connect_db() as conn:
with conn.transaction():
with conn.cursor() as cur: with conn.cursor() as cur:
cur.execute( 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() cur.execute(
return row[0] if row else None """
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: try:
return await asyncio.to_thread(work) return await asyncio.to_thread(_load_audit_sync, tenant, audit_id)
except RuntimeError as exc: except RuntimeError as exc:
if str(exc) == "database_unavailable": if str(exc) == "database_unavailable":
return {"ok": False, "code": "database_unavailable", "message": "Persistent storage is not configured for this deployment.", "audit_id": audit_id} return {"ok": False, "code": "database_unavailable", "message": "Persistent storage is not configured for this deployment.", "audit_id": audit_id}
raise 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