Spaces:
Paused
Paused
| """SQLite store for the controller. | |
| One connection, WAL, guarded by a lock: request handlers and the reconcile thread share | |
| it. Times are stored as epoch seconds (REAL). Events are append-only JSONL files per | |
| agent under `<data>/events/<swarm>/<agent>.jsonl`; the store keeps a per-agent line count | |
| so `after=` cursors are line offsets. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import sqlite3 | |
| import threading | |
| from collections.abc import Callable, Iterable | |
| from pathlib import Path | |
| from typing import Any | |
| from flock.models import Event, SwarmSpec, TokenUsage | |
| SCHEMA = """ | |
| CREATE TABLE IF NOT EXISTS swarms ( | |
| name TEXT PRIMARY KEY, | |
| spec TEXT NOT NULL, | |
| applied_at REAL NOT NULL, | |
| removing INTEGER NOT NULL DEFAULT 0 | |
| ); | |
| CREATE TABLE IF NOT EXISTS types ( | |
| swarm TEXT NOT NULL, | |
| type TEXT NOT NULL, | |
| health TEXT NOT NULL DEFAULT 'ok', | |
| next_worker INTEGER NOT NULL DEFAULT 0, | |
| PRIMARY KEY (swarm, type) | |
| ); | |
| CREATE TABLE IF NOT EXISTS agents ( | |
| name TEXT PRIMARY KEY, | |
| swarm TEXT NOT NULL, | |
| type TEXT NOT NULL, | |
| slot INTEGER NOT NULL, | |
| worker TEXT, | |
| hb_state TEXT, | |
| turn_id TEXT, | |
| idle_since REAL, | |
| session_id TEXT, | |
| resumed INTEGER, | |
| tokens_base TEXT NOT NULL DEFAULT '{}', | |
| tokens_last TEXT NOT NULL DEFAULT '{}', | |
| cost REAL NOT NULL DEFAULT 0, | |
| harness_estimate_usd REAL, | |
| turns INTEGER NOT NULL DEFAULT 0, | |
| last_sync REAL, | |
| msgboard_principal TEXT, | |
| resume_pending INTEGER NOT NULL DEFAULT 0, | |
| created_at REAL NOT NULL, | |
| last_heartbeat REAL | |
| ); | |
| CREATE INDEX IF NOT EXISTS agents_by_swarm ON agents (swarm, type); | |
| CREATE INDEX IF NOT EXISTS agents_by_worker ON agents (worker); | |
| CREATE TABLE IF NOT EXISTS workers ( | |
| id TEXT PRIMARY KEY, | |
| swarm TEXT NOT NULL, | |
| type TEXT NOT NULL, | |
| idx INTEGER NOT NULL, | |
| job_id TEXT, | |
| incarnation TEXT, | |
| state TEXT NOT NULL DEFAULT 'launching', | |
| deadline REAL, | |
| launched_at REAL, | |
| last_heartbeat REAL, | |
| last_inspect REAL, | |
| draining INTEGER NOT NULL DEFAULT 0, | |
| crashes INTEGER NOT NULL DEFAULT 0, | |
| consecutive_crashes INTEGER NOT NULL DEFAULT 0, | |
| next_launch_at REAL, | |
| retire_at REAL, | |
| created_at REAL NOT NULL | |
| ); | |
| CREATE INDEX IF NOT EXISTS workers_by_swarm ON workers (swarm, type, idx); | |
| CREATE TABLE IF NOT EXISTS commands ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| worker TEXT NOT NULL, | |
| body TEXT NOT NULL | |
| ); | |
| CREATE INDEX IF NOT EXISTS commands_by_worker ON commands (worker, id); | |
| CREATE TABLE IF NOT EXISTS pending_prompts ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| agent TEXT NOT NULL, | |
| body TEXT NOT NULL, | |
| created_at REAL NOT NULL | |
| ); | |
| CREATE INDEX IF NOT EXISTS prompts_by_agent ON pending_prompts (agent, id); | |
| CREATE TABLE IF NOT EXISTS event_index ( | |
| agent TEXT PRIMARY KEY, | |
| swarm TEXT NOT NULL, | |
| lines INTEGER NOT NULL DEFAULT 0 | |
| ); | |
| CREATE TABLE IF NOT EXISTS launches (ts REAL NOT NULL); | |
| CREATE TABLE IF NOT EXISTS agent_tokens ( | |
| agent TEXT PRIMARY KEY, | |
| swarm TEXT NOT NULL, | |
| token_hash TEXT NOT NULL UNIQUE, | |
| nonce TEXT NOT NULL, | |
| created_at REAL NOT NULL | |
| ); | |
| CREATE TABLE IF NOT EXISTS spawn_keys ( | |
| parent TEXT NOT NULL, | |
| key TEXT NOT NULL, | |
| child TEXT NOT NULL, | |
| PRIMARY KEY (parent, key) | |
| ); | |
| """ | |
| Row = dict[str, Any] | |
| def _tokens(raw: str) -> dict[str, TokenUsage]: | |
| return {k: TokenUsage.model_validate(v) for k, v in json.loads(raw or "{}").items()} | |
| def _dump_tokens(tokens: dict[str, TokenUsage]) -> str: | |
| return json.dumps({k: v.model_dump() for k, v in tokens.items()}, sort_keys=True) | |
| def agent_tokens(row: Row) -> dict[str, TokenUsage]: | |
| """Total tokens for an agent: earlier incarnations (base) plus the current one (last).""" | |
| total = _tokens(row["tokens_base"]) | |
| for model, usage in _tokens(row["tokens_last"]).items(): | |
| total[model] = total.get(model, TokenUsage()) + usage | |
| return total | |
| class Store: | |
| def __init__(self, data_dir: str | Path): | |
| self.data_dir = Path(data_dir) | |
| self.data_dir.mkdir(parents=True, exist_ok=True) | |
| self._lock = threading.RLock() | |
| self._db = sqlite3.connect( | |
| self.data_dir / "controller.db", check_same_thread=False, isolation_level=None | |
| ) | |
| self._db.row_factory = sqlite3.Row | |
| self._db.execute("PRAGMA journal_mode=WAL") | |
| self._db.execute("PRAGMA synchronous=NORMAL") | |
| self._db.executescript(SCHEMA) | |
| self._migrate() | |
| # Called with ("commands", worker_id) / ("events", agent) / ("reconcile", "") after | |
| # writes; the API turns these into wake-ups for long-polls. Must be thread-safe. | |
| self.listener: Callable[[str, str], None] = lambda kind, key: None | |
| # HMAC key for per-agent tokens (controller/spawn.py), derived from the controller | |
| # token and never stored: the table holds only hashes. | |
| self.agent_token_key: bytes | None = None | |
| def _migrate(self) -> None: | |
| """Additive column changes for databases created by an older version.""" | |
| cols = {r["name"] for r in self._q("PRAGMA table_info(workers)")} | |
| if "incarnation" not in cols: | |
| self._x("ALTER TABLE workers ADD COLUMN incarnation TEXT") | |
| # Spawn (docs/spawn.md): lineage on agents, the never-reused spawn counter and the cost | |
| # of removed agents on swarms. | |
| cols = {r["name"] for r in self._q("PRAGMA table_info(agents)")} | |
| if "parent" not in cols: | |
| self._x("ALTER TABLE agents ADD COLUMN parent TEXT") | |
| self._x("ALTER TABLE agents ADD COLUMN depth INTEGER NOT NULL DEFAULT 0") | |
| self._x("CREATE INDEX IF NOT EXISTS agents_by_parent ON agents (parent)") | |
| if "removing" not in cols: # set while a cascade removal is in flight | |
| self._x("ALTER TABLE agents ADD COLUMN removing INTEGER NOT NULL DEFAULT 0") | |
| cols = {r["name"] for r in self._q("PRAGMA table_info(swarms)")} | |
| if "spawned" not in cols: | |
| self._x("ALTER TABLE swarms ADD COLUMN spawned INTEGER NOT NULL DEFAULT 0") | |
| self._x("ALTER TABLE swarms ADD COLUMN removed_cost REAL NOT NULL DEFAULT 0") | |
| # ----------------------------------------------------------------- low level | |
| def _q(self, sql: str, args: Iterable[Any] = ()) -> list[Row]: | |
| with self._lock: | |
| return [dict(r) for r in self._db.execute(sql, tuple(args)).fetchall()] | |
| def _one(self, sql: str, args: Iterable[Any] = ()) -> Row | None: | |
| rows = self._q(sql, args) | |
| return rows[0] if rows else None | |
| def _x(self, sql: str, args: Iterable[Any] = ()) -> int: | |
| with self._lock: | |
| cur = self._db.execute(sql, tuple(args)) | |
| return cur.lastrowid or 0 | |
| def _update(self, table: str, key: str, value: Any, fields: dict[str, Any]) -> None: | |
| if not fields: | |
| return | |
| cols = ", ".join(f"{k} = ?" for k in fields) # keys come from code, never from input | |
| self._x(f"UPDATE {table} SET {cols} WHERE {key} = ?", [*fields.values(), value]) | |
| def transaction(self): | |
| """Context manager for a multi-statement write.""" | |
| store = self | |
| class _Tx: | |
| def __enter__(self): | |
| store._lock.acquire() | |
| store._db.execute("BEGIN IMMEDIATE") | |
| def __exit__(self, exc_type, exc, tb): | |
| try: | |
| store._db.execute("ROLLBACK" if exc_type else "COMMIT") | |
| finally: | |
| store._lock.release() | |
| return _Tx() | |
| def close(self) -> None: | |
| with self._lock: | |
| self._db.close() | |
| # ----------------------------------------------------------------- swarms and types | |
| def put_swarm(self, spec: SwarmSpec, now: float) -> None: | |
| with self.transaction(): | |
| self._x( | |
| "INSERT INTO swarms (name, spec, applied_at, removing) VALUES (?, ?, ?, 0) " | |
| "ON CONFLICT(name) DO UPDATE SET spec = excluded.spec, " | |
| "applied_at = excluded.applied_at, removing = 0", | |
| [spec.swarm, spec.model_dump_json(), now], | |
| ) | |
| for type_name in spec.types: | |
| self._x( | |
| "INSERT INTO types (swarm, type) VALUES (?, ?) " | |
| "ON CONFLICT(swarm, type) DO UPDATE SET health = 'ok'", | |
| [spec.swarm, type_name], | |
| ) | |
| # Re-applying the spec clears crashlooping and its crash streaks. | |
| self._x("UPDATE workers SET consecutive_crashes = 0 WHERE swarm = ?", [spec.swarm]) | |
| self.listener("reconcile", "") | |
| def get_swarm(self, name: str) -> Row | None: | |
| row = self._one("SELECT * FROM swarms WHERE name = ?", [name]) | |
| if row: | |
| row["spec"] = SwarmSpec.model_validate_json(row["spec"]) | |
| return row | |
| def list_swarms(self) -> list[Row]: | |
| rows = self._q("SELECT * FROM swarms ORDER BY name") | |
| for row in rows: | |
| row["spec"] = SwarmSpec.model_validate_json(row["spec"]) | |
| return rows | |
| def set_removing(self, name: str) -> None: | |
| self._x("UPDATE swarms SET removing = 1 WHERE name = ?", [name]) | |
| self.listener("reconcile", "") | |
| def delete_swarm(self, name: str) -> None: | |
| with self.transaction(): | |
| self._x("DELETE FROM swarms WHERE name = ?", [name]) | |
| self._x("DELETE FROM types WHERE swarm = ?", [name]) | |
| def type_names(self, swarm: str) -> list[str]: | |
| """Types with any state in the store (spec'd or leftover from a removed type).""" | |
| rows = self._q( | |
| "SELECT type FROM types WHERE swarm = ? UNION SELECT type FROM agents WHERE swarm = ? " | |
| "UNION SELECT type FROM workers WHERE swarm = ?", | |
| [swarm, swarm, swarm], | |
| ) | |
| return sorted(r["type"] for r in rows) | |
| def type_health(self, swarm: str, type_name: str) -> str: | |
| row = self._one("SELECT health FROM types WHERE swarm = ? AND type = ?", [swarm, type_name]) | |
| return row["health"] if row else "ok" | |
| def set_type_health(self, swarm: str, type_name: str, health: str) -> None: | |
| self._x( | |
| "INSERT INTO types (swarm, type, health) VALUES (?, ?, ?) " | |
| "ON CONFLICT(swarm, type) DO UPDATE SET health = excluded.health", | |
| [swarm, type_name, health], | |
| ) | |
| def next_worker_index(self, swarm: str, type_name: str) -> int: | |
| with self.transaction(): | |
| self._x("INSERT OR IGNORE INTO types (swarm, type) VALUES (?, ?)", [swarm, type_name]) | |
| row = self._one( | |
| "SELECT next_worker FROM types WHERE swarm = ? AND type = ?", [swarm, type_name] | |
| ) | |
| idx = row["next_worker"] | |
| self._x( | |
| "UPDATE types SET next_worker = ? WHERE swarm = ? AND type = ?", | |
| [idx + 1, swarm, type_name], | |
| ) | |
| return idx | |
| def bump_worker_index(self, swarm: str, type_name: str, at_least: int) -> None: | |
| self._x("INSERT OR IGNORE INTO types (swarm, type) VALUES (?, ?)", [swarm, type_name]) | |
| self._x( | |
| "UPDATE types SET next_worker = MAX(next_worker, ?) WHERE swarm = ? AND type = ?", | |
| [at_least, swarm, type_name], | |
| ) | |
| def type_owner(self, type_name: str) -> str | None: | |
| """The swarm owning a type name (agent names are global, so type names must be).""" | |
| row = self._one("SELECT swarm FROM types WHERE type = ? LIMIT 1", [type_name]) | |
| return row["swarm"] if row else None | |
| # ----------------------------------------------------------------- agents | |
| def get_agent(self, name: str) -> Row | None: | |
| return self._one("SELECT * FROM agents WHERE name = ?", [name]) | |
| def list_agents( | |
| self, | |
| swarm: str | None = None, | |
| type_name: str | None = None, | |
| worker: str | None = None, | |
| parent: str | None = None, | |
| ) -> list[Row]: | |
| where, args = [], [] | |
| for col, val in ( | |
| ("swarm", swarm), ("type", type_name), ("worker", worker), ("parent", parent) | |
| ): | |
| if val is not None: | |
| where.append(f"{col} = ?") | |
| args.append(val) | |
| clause = f"WHERE {' AND '.join(where)}" if where else "" | |
| return self._q(f"SELECT * FROM agents {clause} ORDER BY swarm, type, slot", args) | |
| def insert_agent(self, name: str, swarm: str, type_name: str, slot: int, now: float) -> None: | |
| self._x( | |
| "INSERT OR IGNORE INTO agents (name, swarm, type, slot, created_at) " | |
| "VALUES (?, ?, ?, ?, ?)", | |
| [name, swarm, type_name, slot, now], | |
| ) | |
| def update_agent(self, name: str, **fields: Any) -> None: | |
| self._update("agents", "name", name, fields) | |
| def delete_agent(self, name: str) -> None: | |
| with self.transaction(): | |
| self.delete_agent_rows(name) | |
| def delete_agent_rows(self, name: str) -> None: | |
| """delete_agent's statements, for a caller already inside a transaction.""" | |
| # Removed agents keep counting against limits.budget_usd. | |
| self._x( | |
| "UPDATE swarms SET removed_cost = removed_cost + " | |
| "(SELECT cost FROM agents WHERE name = ?) " | |
| "WHERE name = (SELECT swarm FROM agents WHERE name = ?)", | |
| [name, name], | |
| ) | |
| self._x("DELETE FROM agents WHERE name = ?", [name]) | |
| self._x("DELETE FROM pending_prompts WHERE agent = ?", [name]) | |
| self._x("DELETE FROM agent_tokens WHERE agent = ?", [name]) # revoke | |
| self._x("DELETE FROM spawn_keys WHERE parent = ?", [name]) | |
| def mark_removing(self, name: str) -> None: | |
| """First phase of a cascade removal: refuse new children, revoke the token.""" | |
| self._x("UPDATE agents SET removing = 1 WHERE name = ?", [name]) | |
| self._x("DELETE FROM agent_tokens WHERE agent = ?", [name]) | |
| # ----------------------------------------------------------------- spawn | |
| def insert_child( | |
| self, name: str, swarm: str, type_name: str, n: int, parent: str, depth: int, now: float | |
| ) -> None: | |
| self._x( | |
| "INSERT INTO agents (name, swarm, type, slot, parent, depth, created_at) " | |
| "VALUES (?, ?, ?, ?, ?, ?, ?)", | |
| [name, swarm, type_name, n, parent, depth, now], | |
| ) | |
| def next_spawn_number(self, swarm: str) -> int: | |
| """The swarm's spawn counter, incremented; never reused while the swarm exists.""" | |
| with self._lock: | |
| self._x("UPDATE swarms SET spawned = spawned + 1 WHERE name = ?", [swarm]) | |
| return self._one("SELECT spawned FROM swarms WHERE name = ?", [swarm])["spawned"] | |
| def spawned_count(self, swarm: str) -> int: | |
| row = self._one("SELECT spawned FROM swarms WHERE name = ?", [swarm]) | |
| return row["spawned"] if row else 0 | |
| def children_by_parent(self, swarm: str) -> dict[str, list[str]]: | |
| out: dict[str, list[str]] = {} | |
| for r in self._q( | |
| "SELECT parent, name FROM agents WHERE swarm = ? AND parent IS NOT NULL ORDER BY name", | |
| [swarm], | |
| ): | |
| out.setdefault(r["parent"], []).append(r["name"]) | |
| return out | |
| def swarm_cost(self, swarm: str) -> float: | |
| """Cost of every agent the swarm has had: live rows plus removed ones.""" | |
| row = self._one( | |
| "SELECT s.removed_cost + COALESCE((SELECT SUM(cost) FROM agents WHERE swarm = ?), 0) " | |
| "AS c FROM swarms s WHERE s.name = ?", | |
| [swarm, swarm], | |
| ) | |
| return row["c"] if row else 0.0 | |
| def count_agents(self, swarm: str, type_name: str | None = None) -> int: | |
| sql, args = "SELECT COUNT(*) AS n FROM agents WHERE swarm = ?", [swarm] | |
| if type_name is not None: | |
| sql += " AND type = ?" | |
| args.append(type_name) | |
| return self._one(sql, args)["n"] | |
| def spawn_key_child(self, parent: str, key: str) -> str | None: | |
| row = self._one("SELECT child FROM spawn_keys WHERE parent = ? AND key = ?", [parent, key]) | |
| return row["child"] if row else None | |
| def put_spawn_key(self, parent: str, key: str, child: str) -> None: | |
| self._x("INSERT INTO spawn_keys (parent, key, child) VALUES (?, ?, ?)", [parent, key, child]) | |
| # ----------------------------------------------------------------- agent tokens | |
| def get_agent_token(self, agent: str) -> Row | None: | |
| return self._one("SELECT * FROM agent_tokens WHERE agent = ?", [agent]) | |
| def agent_for_token_hash(self, token_hash: str) -> str | None: | |
| row = self._one("SELECT agent FROM agent_tokens WHERE token_hash = ?", [token_hash]) | |
| return row["agent"] if row else None | |
| def put_agent_token( | |
| self, agent: str, swarm: str, token_hash: str, nonce: str, now: float | |
| ) -> None: | |
| self._x( | |
| "INSERT INTO agent_tokens (agent, swarm, token_hash, nonce, created_at) " | |
| "VALUES (?, ?, ?, ?, ?) ON CONFLICT(agent) DO UPDATE SET " | |
| "token_hash = excluded.token_hash, nonce = excluded.nonce", | |
| [agent, swarm, token_hash, nonce, now], | |
| ) | |
| def record_agent_heartbeat( | |
| self, name: str, fields: dict[str, Any], tokens: dict[str, TokenUsage], cost_fn | |
| ) -> None: | |
| """Store one AgentHeartbeat. Tokens are cumulative (the worker restores them from the | |
| bucket); if a model's counters go backwards the count restarted (a failed restore), so | |
| fold the old values into base. An empty map means "not known yet" (an agent still | |
| restoring on a relaunched worker) and leaves tokens and cost alone: folding there | |
| counted every restored agent twice (1k run).""" | |
| with self.transaction(): | |
| row = self.get_agent(name) | |
| if row is None: | |
| return | |
| if not tokens and row["tokens_last"] not in ("", "{}"): | |
| if fields.get("harness_estimate_usd") is None: | |
| fields.pop("harness_estimate_usd", None) | |
| self._update("agents", "name", name, fields) | |
| return | |
| base, last = _tokens(row["tokens_base"]), _tokens(row["tokens_last"]) | |
| for model, old in last.items(): | |
| new = tokens.get(model, TokenUsage()) | |
| if any(getattr(new, f) < getattr(old, f) for f in TokenUsage.model_fields): | |
| base[model] = base.get(model, TokenUsage()) + old | |
| fields["tokens_base"] = _dump_tokens(base) | |
| fields["tokens_last"] = _dump_tokens(tokens) | |
| total = dict(base) | |
| for model, usage in tokens.items(): | |
| total[model] = total.get(model, TokenUsage()) + usage | |
| fields["cost"] = cost_fn(total) | |
| self._update("agents", "name", name, fields) | |
| # ----------------------------------------------------------------- workers | |
| def get_worker(self, worker_id: str) -> Row | None: | |
| return self._one("SELECT * FROM workers WHERE id = ?", [worker_id]) | |
| def list_workers(self, swarm: str | None = None, type_name: str | None = None) -> list[Row]: | |
| where, args = [], [] | |
| for col, val in (("swarm", swarm), ("type", type_name)): | |
| if val is not None: | |
| where.append(f"{col} = ?") | |
| args.append(val) | |
| clause = f"WHERE {' AND '.join(where)}" if where else "" | |
| return self._q(f"SELECT * FROM workers {clause} ORDER BY swarm, type, idx", args) | |
| def insert_worker(self, worker_id: str, swarm: str, type_name: str, idx: int, now: float): | |
| self._x( | |
| "INSERT OR IGNORE INTO workers (id, swarm, type, idx, created_at) VALUES (?, ?, ?, ?, ?)", | |
| [worker_id, swarm, type_name, idx, now], | |
| ) | |
| def update_worker(self, worker_id: str, **fields: Any) -> None: | |
| self._update("workers", "id", worker_id, fields) | |
| def delete_worker(self, worker_id: str) -> None: | |
| with self.transaction(): | |
| self._x("DELETE FROM workers WHERE id = ?", [worker_id]) | |
| self._x("DELETE FROM commands WHERE worker = ?", [worker_id]) | |
| self._x("UPDATE agents SET worker = NULL WHERE worker = ?", [worker_id]) | |
| # ----------------------------------------------------------------- commands | |
| def enqueue_command(self, worker_id: str, body: str) -> None: | |
| self._x("INSERT INTO commands (worker, body) VALUES (?, ?)", [worker_id, body]) | |
| self.listener("commands", worker_id) | |
| def commands_for(self, worker_id: str) -> list[Row]: | |
| return self._q("SELECT id, body FROM commands WHERE worker = ? ORDER BY id", [worker_id]) | |
| def clear_commands(self, worker_id: str) -> None: | |
| self._x("DELETE FROM commands WHERE worker = ?", [worker_id]) | |
| def delete_commands(self, ids: list[int]) -> None: | |
| if ids: | |
| marks = ",".join("?" * len(ids)) | |
| self._x(f"DELETE FROM commands WHERE id IN ({marks})", ids) | |
| # ----------------------------------------------------------------- pending prompts | |
| def add_prompt(self, agent: str, body: str, now: float) -> None: | |
| self._x( | |
| "INSERT INTO pending_prompts (agent, body, created_at) VALUES (?, ?, ?)", | |
| [agent, body, now], | |
| ) | |
| row = self.get_agent(agent) | |
| if row and row["worker"]: | |
| self.listener("commands", row["worker"]) | |
| def prompts_for_worker(self, worker_id: str) -> list[Row]: | |
| return self._q( | |
| "SELECT p.id, p.body FROM pending_prompts p JOIN agents a ON a.name = p.agent " | |
| "WHERE a.worker = ? ORDER BY p.id", | |
| [worker_id], | |
| ) | |
| def delete_prompts(self, ids: list[int]) -> None: | |
| if ids: | |
| marks = ",".join("?" * len(ids)) | |
| self._x(f"DELETE FROM pending_prompts WHERE id IN ({marks})", ids) | |
| def pending_counts(self, swarm: str | None = None) -> dict[str, int]: | |
| sql = "SELECT p.agent, COUNT(*) AS n FROM pending_prompts p" | |
| args: list[Any] = [] | |
| if swarm is not None: | |
| sql += " JOIN agents a ON a.name = p.agent WHERE a.swarm = ?" | |
| args.append(swarm) | |
| return {r["agent"]: r["n"] for r in self._q(sql + " GROUP BY p.agent", args)} | |
| # ----------------------------------------------------------------- launch budget | |
| def record_launch(self, now: float) -> None: | |
| self._x("INSERT INTO launches (ts) VALUES (?)", [now]) | |
| def launches_since(self, since: float) -> int: | |
| self._x("DELETE FROM launches WHERE ts < ?", [since]) | |
| return self._one("SELECT COUNT(*) AS n FROM launches WHERE ts >= ?", [since])["n"] | |
| # ----------------------------------------------------------------- events | |
| def _events_path(self, swarm: str, agent: str) -> Path: | |
| return self.data_dir / "events" / swarm / f"{agent}.jsonl" | |
| def append_events(self, swarm: str, agent: str, events: list[Event]) -> int: | |
| """Append events to the agent's stream; returns the new line count.""" | |
| if not events: | |
| return self.event_count(agent) | |
| path = self._events_path(swarm, agent) | |
| with self._lock: | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| with path.open("a", encoding="utf-8") as f: | |
| for ev in events: | |
| f.write(ev.model_dump_json() + "\n") | |
| self._x( | |
| "INSERT INTO event_index (agent, swarm, lines) VALUES (?, ?, ?) " | |
| "ON CONFLICT(agent) DO UPDATE SET lines = lines + excluded.lines", | |
| [agent, swarm, len(events)], | |
| ) | |
| n = self.event_count(agent) | |
| self.listener("events", agent) | |
| return n | |
| def event_count(self, agent: str) -> int: | |
| row = self._one("SELECT lines FROM event_index WHERE agent = ?", [agent]) | |
| return row["lines"] if row else 0 | |
| def read_events(self, agent: str, after: int, limit: int = 1000) -> list[dict[str, Any]]: | |
| row = self._one("SELECT swarm, lines FROM event_index WHERE agent = ?", [agent]) | |
| if row is None or after >= row["lines"]: | |
| return [] | |
| out: list[dict[str, Any]] = [] | |
| with self._events_path(row["swarm"], agent).open(encoding="utf-8") as f: | |
| for i, line in enumerate(f): | |
| if i >= row["lines"] or len(out) >= limit: | |
| break | |
| if i >= after: | |
| out.append(json.loads(line)) | |
| return out | |