Spaces:
Paused
Paused
| """Spawned agents (docs/spawn.md): admission, removal with cascade, per-agent tokens and | |
| what an agent token may do. | |
| A spawned agent is an ordinary agent row with `parent` set and no replica slot (its `slot` | |
| holds the swarm's spawn counter value instead). Reconcile places it on its type's workers like | |
| any other row; scaling `replicas` never touches it. | |
| Per-agent tokens are HMAC-SHA256(key, "<agent>\\0<nonce>") under a key derived from the | |
| controller token. The table keeps the nonce and a SHA-256 of the token, never the token, and | |
| the controller can re-derive the same value for every later assign of the agent (a relaunched | |
| worker, a re-sent assign). Removing the agent deletes the row, which revokes the token; | |
| rotating FLOCK_TOKEN invalidates every agent token until the next assign re-issues it. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import hmac | |
| import logging | |
| import secrets | |
| import uuid | |
| from collections.abc import Mapping | |
| from datetime import UTC, datetime | |
| from flock.controller.store import Row, Store | |
| from flock.models import ( | |
| Event, | |
| PromptCommand, | |
| SpawnAck, | |
| SpawnRequest, | |
| SwarmSpec, | |
| TypeHealth, | |
| UnassignCommand, | |
| child_depth, | |
| child_name, | |
| max_workers, | |
| root_depth, | |
| ) | |
| log = logging.getLogger("flock.spawn") | |
| REASONS = ( | |
| "type_not_allowed", | |
| "max_children", | |
| "depth", | |
| "workers", | |
| "alive", | |
| "spawned_total", | |
| "budget", | |
| "type_crashlooping", | |
| "parent_removing", | |
| ) | |
| MAX_LINEAGE = 64 # parent-chain walk bound; depth is finite by construction, this is a guard | |
| class Refused(Exception): | |
| """A spawn the policy does not admit: 409 with `reason` (one of REASONS).""" | |
| def __init__(self, reason: str, message: str) -> None: | |
| super().__init__(message) | |
| self.reason, self.message = reason, message | |
| def _ts(now: float) -> datetime: | |
| return datetime.fromtimestamp(now, UTC) | |
| # --------------------------------------------------------------------------- tokens | |
| def set_token_key(store: Store, controller_token: str) -> None: | |
| store.agent_token_key = ( | |
| hashlib.sha256(b"flock-agent-token\0" + controller_token.encode()).digest() | |
| if controller_token | |
| else None | |
| ) | |
| def _mint(key: bytes, agent: str, nonce: str) -> str: | |
| return "flock_at_" + hmac.new(key, f"{agent}\0{nonce}".encode(), hashlib.sha256).hexdigest() | |
| def token_hash(token: str) -> str: | |
| return hashlib.sha256(token.encode()).hexdigest() | |
| def agent_token(store: Store, agent: str, swarm: str, now: float) -> str | None: | |
| """The agent's token: issued on first call, the same value on every later call. None for | |
| an agent that does not exist or is being removed (its token stays revoked).""" | |
| key = store.agent_token_key | |
| if key is None: | |
| return None | |
| a = store.get_agent(agent) | |
| if a is None or a["removing"]: | |
| return None | |
| row = store.get_agent_token(agent) | |
| nonce = row["nonce"] if row else secrets.token_hex(16) | |
| token = _mint(key, agent, nonce) | |
| if row is None or row["token_hash"] != token_hash(token): # new, or FLOCK_TOKEN rotated | |
| store.put_agent_token(agent, swarm, token_hash(token), nonce, now) | |
| return token | |
| def agent_for_token(store: Store, token: str) -> Row | None: | |
| """The live agent a presented token belongs to, or None (also while it is being removed).""" | |
| key = store.agent_token_key | |
| if key is None or not token.startswith("flock_at_"): | |
| return None | |
| name = store.agent_for_token_hash(token_hash(token)) | |
| row = store.get_agent_token(name) if name else None | |
| if row is None or not hmac.compare_digest(_mint(key, name, row["nonce"]), token): | |
| return None # minted under an older controller token | |
| a = store.get_agent(name) | |
| return None if a is None or a["removing"] else a | |
| # --------------------------------------------------------------------------- lineage | |
| def agent_depth(spec: SwarmSpec | None, a: Row) -> int: | |
| """Generations `a` may still spawn below itself. A replica follows the current spec; a | |
| spawned agent keeps the depth it was born with.""" | |
| if a["parent"] is not None: | |
| return a["depth"] | |
| return root_depth(spec, a["type"]) if spec else 0 | |
| def is_descendant(store: Store, name: str | None, ancestor: str) -> bool: | |
| a = store.get_agent(name) if name else None | |
| for _ in range(MAX_LINEAGE): | |
| if a is None or a["parent"] is None: | |
| return False | |
| if a["parent"] == ancestor: | |
| return True | |
| a = store.get_agent(a["parent"]) | |
| return False | |
| def agent_may( | |
| store: Store, | |
| agent: Row, | |
| method: str, | |
| route: str, | |
| params: Mapping[str, str], | |
| query: Mapping[str, str], | |
| ) -> bool: | |
| """What an agent token allows: itself (status, events, whoami), its subtree (status, | |
| events, spawn, prompt, remove; it cannot remove itself) and prompting its parent.""" | |
| me, target = agent["name"], params.get("name") | |
| def below(t: str | None) -> bool: | |
| return is_descendant(store, t, me) | |
| match method, route: | |
| case "GET", "/whoami": | |
| return True | |
| case "GET", "/agents/{name}" | "/agents/{name}/events": | |
| return target == me or below(target) | |
| case "POST", "/agents/{name}/spawn": | |
| return target == me or below(target) | |
| case "POST", "/agents/{name}/prompt": | |
| return (target is not None and target == agent["parent"]) or below(target) | |
| case "DELETE", "/agents/{name}": | |
| return below(target) | |
| case "GET", "/swarms/{swarm}/agents": | |
| parent = query.get("parent") | |
| return params.get("swarm") == agent["swarm"] and (parent == me or below(parent)) | |
| return False | |
| # --------------------------------------------------------------------------- spawn | |
| def spawn(store: Store, parent_name: str, req: SpawnRequest, now: float) -> tuple[SpawnAck, bool]: | |
| """Admit and create one child of `parent_name`. Returns (ack, created). Raises KeyError for | |
| an unknown parent and Refused for a policy refusal. A repeated `key` returns the child it | |
| created first, whatever the caps say now.""" | |
| with store.transaction(): | |
| parent = store.get_agent(parent_name) | |
| if parent is None: | |
| raise KeyError(parent_name) | |
| if parent["removing"]: | |
| raise Refused("parent_removing", f"{parent_name} is being removed") | |
| if req.key is not None and (child := store.spawn_key_child(parent_name, req.key)): | |
| row = store.get_agent(child) | |
| depth = row["depth"] if row else 0 | |
| ack = SpawnAck(name=child, created=False, parent=parent_name, type=req.type, depth=depth) | |
| return ack, False | |
| sw = store.get_swarm(parent["swarm"]) | |
| if sw is None or sw["removing"]: | |
| raise Refused("parent_removing", f"swarm {parent['swarm']!r} is being removed") | |
| spec: SwarmSpec = sw["spec"] | |
| _admit(store, spec, parent, req.type) | |
| swarm, depth = parent["swarm"], child_depth(spec, agent_depth(spec, parent), req.type) | |
| while True: # a replica of an oddly named type could hold the name: skip it | |
| n = store.next_spawn_number(swarm) | |
| name = child_name(parent_name, req.type, n) | |
| if store.get_agent(name) is None: | |
| break | |
| store.insert_child(name, swarm, req.type, n, parent_name, depth, now) | |
| if req.key is not None: | |
| store.put_spawn_key(parent_name, req.key, name) | |
| agent_token(store, name, swarm, now) # issued at creation | |
| if req.prompt: | |
| cmd = PromptCommand(agent=name, prompt_id=uuid.uuid4().hex, text=req.prompt) | |
| store.add_prompt(name, cmd.model_dump_json(), now) | |
| ev = Event(type="spawned", ts=_ts(now), agent=parent_name, | |
| data={"child": name, "type": req.type, "depth": depth, "key": req.key}) | |
| store.append_events(swarm, parent_name, [ev]) | |
| store.listener("reconcile", "") | |
| return SpawnAck(name=name, created=True, parent=parent_name, type=req.type, depth=depth), True | |
| def _admit(store: Store, spec: SwarmSpec, parent: Row, type_name: str) -> None: | |
| rule = spec.spawn.get(parent["type"]) | |
| if rule is None or type_name not in rule.allowed or type_name not in spec.types: | |
| raise Refused("type_not_allowed", f"type {parent['type']!r} may not spawn {type_name!r}") | |
| if agent_depth(spec, parent) < 1: | |
| raise Refused("depth", f"{parent['name']} may not spawn (depth 0)") | |
| kids = store.list_agents(parent=parent["name"]) | |
| cap = rule.cap_for(type_name) | |
| if cap is not None and sum(k["type"] == type_name for k in kids) >= cap: | |
| raise Refused("max_children", f"{parent['name']} already has {cap} {type_name} children") | |
| if rule.max is not None and len(kids) >= rule.max: | |
| raise Refused("max_children", f"{parent['name']} already has {rule.max} children") | |
| if store.type_health(spec.swarm, type_name) == TypeHealth.crashlooping: | |
| raise Refused("type_crashlooping", f"type {type_name!r} is crashlooping") | |
| lim = spec.limits | |
| if lim.alive is not None and store.count_agents(spec.swarm) >= lim.alive: | |
| raise Refused("alive", f"swarm has {lim.alive} agents (limits.alive)") | |
| if lim.spawned_total is not None and store.spawned_count(spec.swarm) >= lim.spawned_total: | |
| raise Refused("spawned_total", f"swarm spawned {lim.spawned_total} agents") | |
| if lim.budget_usd is not None and store.swarm_cost(spec.swarm) >= lim.budget_usd: | |
| raise Refused("budget", f"swarm cost reached ${lim.budget_usd:g} (limits.budget_usd)") | |
| ts = spec.types[type_name] | |
| room = max_workers(ts) * max(1, ts.agents_per_worker) | |
| if store.count_agents(spec.swarm, type_name) >= room: | |
| raise Refused( | |
| "workers", | |
| f"type {type_name!r} is full: {max_workers(ts)} workers x {ts.agents_per_worker}", | |
| ) | |
| def worker_cap_conflicts(store: Store, spec: SwarmSpec) -> list[str]: | |
| """Types whose declared worker cap is below the workers their agents occupy now. Agents | |
| never move between live workers, so such an edit could not be honoured: apply refuses it.""" | |
| out = [] | |
| for name, ts in sorted(spec.types.items()): | |
| if ts.workers is None: | |
| continue | |
| in_use = len({a["worker"] for a in store.list_agents(spec.swarm, name) if a["worker"]}) | |
| if in_use > max_workers(ts): | |
| out.append( | |
| f"type {name!r}: workers.initial + workers.max_extra = {max_workers(ts)} is below " | |
| f"the {in_use} workers its agents occupy; remove agents first" | |
| ) | |
| return out | |
| # --------------------------------------------------------------------------- removal | |
| def _subtree(store: Store, name: str) -> list[Row]: | |
| """`name` and its descendants, children before parents.""" | |
| out: list[Row] = [] | |
| for child in store.list_agents(parent=name): | |
| out += _subtree(store, child["name"]) | |
| a = store.get_agent(name) | |
| return out + ([a] if a else []) | |
| def remove_subtree(store: Store, name: str, reason: str, now: float) -> list[Row]: | |
| """Remove `name` and everything it spawned, children first. Each placed agent gets an | |
| unassign (final sync; bucket state is kept), each token is revoked, and each spawned | |
| agent's parent gets `child_removed`. Returns the removed rows. | |
| Two transactions, each serialized with spawn admission (same store transaction): the | |
| first enumerates the subtree, marks every row `removing` and revokes every token, so no | |
| child can be added below it from then on (a spawn there gets 409 parent_removing); the | |
| second deletes the rows. A crash between them leaves marked rows that reconcile removes.""" | |
| with store.transaction(): | |
| rows = _subtree(store, name) | |
| for r in rows: | |
| store.mark_removing(r["name"]) | |
| if not rows: | |
| return [] | |
| events: list[Event] = [] | |
| with store.transaction(): | |
| for r in rows: | |
| w = store.get_worker(r["worker"]) if r["worker"] else None | |
| if w and w["job_id"]: | |
| store.enqueue_command(w["id"], UnassignCommand(agent=r["name"]).model_dump_json()) | |
| store.delete_agent_rows(r["name"]) | |
| if r["parent"] and store.get_agent(r["parent"]) is not None: | |
| why = reason if r["name"] == name else "parent_removed" | |
| events.append(Event(type="child_removed", ts=_ts(now), agent=r["parent"], | |
| data={"child": r["name"], "type": r["type"], "reason": why})) | |
| for ev in events: | |
| store.append_events(rows[0]["swarm"], ev.agent, [ev]) | |
| return rows | |