diff --git a/api/agent_ops.py b/api/agent_ops.py index 7634b13..2a35a50 100644 --- a/api/agent_ops.py +++ b/api/agent_ops.py @@ -23,8 +23,8 @@ from __future__ import annotations import asyncio import os import random -import time import uuid +from collections import deque from datetime import datetime, timezone from typing import Any @@ -118,14 +118,34 @@ def _tag(run_id: str) -> str: return f"{AGENT_TAG} {run_id}" +# In-memory pools of agent-created primary keys per source. All UPDATE/DELETE +# operate by PK on these (instant), so we never scan the large unindexed +# `notes` column. Bounded by _POOL_CAP; orphaned tagged rows after a restart +# are harmless (clearly marked) and simply not re-tracked. +_pools: dict[str, deque] = { + "postgres": deque(maxlen=_POOL_CAP), + "mysql": deque(maxlen=_POOL_CAP), + "mongodb": deque(maxlen=_POOL_CAP), +} + + +def _resolve_op(op: str, pool: deque) -> str: + """UPDATE/DELETE require an existing agent row; otherwise fall back to INSERT.""" + if op in ("update", "delete") and not pool: + return "insert" + if op == "insert" and len(pool) >= _POOL_CAP: + return "delete" # keep the agent footprint bounded + return op + + def _pg_dml(op: str) -> str: + pool = _pools["postgres"] + op = _resolve_op(op, pool) conn = _pg_conn() try: conn.autocommit = True with conn.cursor() as cur: - cur.execute("SELECT count(*) FROM public.sales_orders WHERE notes LIKE %s", (AGENT_TAG + "%",)) - pool = cur.fetchone()[0] - if op == "insert" or (op == "delete" and pool == 0) or (op == "update" and pool == 0): + if op == "insert": rid = uuid.uuid4().hex[:8] cur.execute( """INSERT INTO public.sales_orders @@ -139,38 +159,29 @@ def _pg_dml(op: str) -> str: ), ) oid = cur.fetchone()[0] + pool.append(oid) return f"INSERT sales_orders order_id={oid} ({_tag(rid)})" if op == "update": + oid = random.choice(list(pool)) cur.execute( - """UPDATE public.sales_orders SET order_status=%s, amount=round(amount*%s,2) - WHERE order_id IN ( - SELECT order_id FROM public.sales_orders WHERE notes LIKE %s - ORDER BY order_id DESC LIMIT 1) - RETURNING order_id""", - (random.choice(ORDER_STATUS), round(random.uniform(0.9, 1.2), 2), AGENT_TAG + "%"), + "UPDATE public.sales_orders SET order_status=%s, amount=round(amount*%s,2) WHERE order_id=%s", + (random.choice(ORDER_STATUS), round(random.uniform(0.9, 1.2), 2), oid), ) - row = cur.fetchone() - return f"UPDATE sales_orders order_id={row[0]}" if row else "UPDATE sales_orders (no agent rows)" - # delete - cur.execute( - """DELETE FROM public.sales_orders WHERE order_id IN ( - SELECT order_id FROM public.sales_orders WHERE notes LIKE %s - ORDER BY order_id ASC LIMIT 1) RETURNING order_id""", - (AGENT_TAG + "%",), - ) - row = cur.fetchone() - return f"DELETE sales_orders order_id={row[0]}" if row else "DELETE sales_orders (no agent rows)" + return f"UPDATE sales_orders order_id={oid}" + oid = pool.popleft() + cur.execute("DELETE FROM public.sales_orders WHERE order_id=%s", (oid,)) + return f"DELETE sales_orders order_id={oid}" finally: conn.close() def _mysql_dml(op: str) -> str: + pool = _pools["mysql"] + op = _resolve_op(op, pool) conn = _mysql_conn() try: with conn.cursor() as cur: - cur.execute("SELECT count(*) FROM hr.employee_events WHERE notes LIKE %s", (AGENT_TAG + "%",)) - pool = cur.fetchone()[0] - if op == "insert" or (op in ("update", "delete") and pool == 0): + if op == "insert": rid = uuid.uuid4().hex[:8] cur.execute( """INSERT INTO hr.employee_events @@ -182,28 +193,29 @@ def _mysql_dml(op: str) -> str: round(random.uniform(-5000, 15000), 2), datetime.now(timezone.utc), _tag(rid), ), ) - return f"INSERT employee_events event_id={cur.lastrowid} ({_tag(rid)})" + eid = cur.lastrowid + pool.append(eid) + return f"INSERT employee_events event_id={eid} ({_tag(rid)})" if op == "update": + eid = random.choice(list(pool)) cur.execute( - """UPDATE hr.employee_events SET salary_change=%s, event_type=%s - WHERE notes LIKE %s ORDER BY event_id DESC LIMIT 1""", - (round(random.uniform(-5000, 15000), 2), random.choice(EVENT_TYPES), AGENT_TAG + "%"), + "UPDATE hr.employee_events SET salary_change=%s, event_type=%s WHERE event_id=%s", + (round(random.uniform(-5000, 15000), 2), random.choice(EVENT_TYPES), eid), ) - return f"UPDATE employee_events ({cur.rowcount} row)" - cur.execute( - "DELETE FROM hr.employee_events WHERE notes LIKE %s ORDER BY event_id ASC LIMIT 1", - (AGENT_TAG + "%",), - ) - return f"DELETE employee_events ({cur.rowcount} row)" + return f"UPDATE employee_events event_id={eid}" + eid = pool.popleft() + cur.execute("DELETE FROM hr.employee_events WHERE event_id=%s", (eid,)) + return f"DELETE employee_events event_id={eid}" finally: conn.close() def _mongo_dml(op: str) -> str: + pool = _pools["mongodb"] + op = _resolve_op(op, pool) client, coll = _mongo_coll() try: - pool = coll.count_documents({"atc_agent": True}, limit=1) - if op == "insert" or (op in ("update", "delete") and pool == 0): + if op == "insert": rid = uuid.uuid4().hex[:8] doc = { "event_id": str(uuid.uuid4()), @@ -217,21 +229,18 @@ def _mongo_dml(op: str) -> str: "agent_run": rid, } res = coll.insert_one(doc) + pool.append(res.inserted_id) return f"INSERT events _id={res.inserted_id} (agent_run={rid})" if op == "update": - doc = coll.find_one({"atc_agent": True}, sort=[("_id", -1)]) - if not doc: - return "UPDATE events (no agent docs)" + oid = random.choice(list(pool)) coll.update_one( - {"_id": doc["_id"]}, + {"_id": oid}, {"$set": {"type": random.choice(MONGO_TYPES), "amount": round(random.uniform(10, 50000), 4)}}, ) - return f"UPDATE events _id={doc['_id']}" - doc = coll.find_one({"atc_agent": True}, sort=[("_id", 1)]) - if not doc: - return "DELETE events (no agent docs)" - coll.delete_one({"_id": doc["_id"]}) - return f"DELETE events _id={doc['_id']}" + return f"UPDATE events _id={oid}" + oid = pool.popleft() + coll.delete_one({"_id": oid}) + return f"DELETE events _id={oid}" finally: client.close() @@ -285,7 +294,8 @@ async def agent_dml_loop() -> None: # ── Endpoints ──────────────────────────────────────────────────────────────── @router.get("/status") async def status() -> JSONResponse: - return JSONResponse({"ok": True, **_state}) + pools = {k: len(v) for k, v in _pools.items()} + return JSONResponse({"ok": True, "pools": pools, **_state}) @router.post("/toggle")