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:
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