Build agent pool manager with lifecycle and capability registration

Introduction

When you wire a supervisor to a fleet of specialised agents, the first thing that breaks is bookkeeping: which agent is alive, which one can answer this kind of task, and how many calls is it already juggling. Hardcode agent IDs into plans and every rolling deploy or crashed pod will strand work mid-flight; skip lifecycle tracking and a single dead pod can stall the entire orchestration loop. By the end of this lesson you'll be able to design an agent pool manager that registers capabilities on join, tracks lifecycle state through a fixed state machine, and routes tasks by capability rather than agent ID.

Key Terminology

• Capability: a named, versioned contract (e.g. rag.search@2.1.0) declaring an input/output schema, a max_concurrency, and a cost class. Plans pin capabilities, never agent IDs.
• Agent registration: the join-time handshake where a freshly booted agent posts its agent_id, endpoint, capability list, and heartbeat_interval_s to the manager so it can be indexed and routed work.
• Lifecycle state: the fixed set Spawning → Registering → Ready → Busy → Draining → Terminated that every agent walks through. The manager is the only writer, and every transition is appended to the orchestration log.
• Heartbeat: a periodic POST /heartbeat from agent to manager. Missing 3 × heartbeat_interval_s of pings flips the agent to Terminated via the watcher loop.
• Drain: a graceful shutdown signal (typically from a Kubernetes preStop hook) that stops new acquisitions while letting _inflight tasks finish before the pod exits.
• Capability index (_cap_index): the multimap capability_name → {agent_ids} that find_by_capability walks. Stale entries here are how dead pods get handed live work, so every _remove must prune it.

Concepts

Kubernetes-native scaling

Each agent type is a Deployment. The manager's job during scale-out is not to create pods directly — that breaks the reconciliation loop — but to patch the Deployment's replicas field. Kubernetes does the rest; the new pod boots, runs its registration handshake against the manager, and lands in the Ready set.

A typical Deployment skeleton (abridged):

• metadata.labels: app=search-agent, capability=rag.search, cap-version=2.1.0
• spec.template.spec.containers[0].env: MANAGER_URL, AGENT_ID (from downward API, the pod name), HEARTBEAT_INTERVAL_S
• readinessProbe: hits /ready which returns 200 only after the agent has successfully registered
• preStop lifecycle hook: calls the manager's /drain endpoint before SIGTERM

The readiness probe is the trick that keeps Kubernetes' service abstraction honest: a pod that has not registered yet is not in the load balancer rotation. Combined with the preStop drain, you get rolling deploys with zero in-flight task loss, because the manager already knows how to wait for _inflight to hit zero.

Operating discipline

• Treat the registry as ephemeral, not authoritative. If the manager restarts, agents must re-register on next heartbeat. Do not persist registrations; persist the lifecycle log instead.
• Cap concurrency at the agent, not at the manager. Each capability declares its own max_concurrency based on what its underlying model or sandbox can sustain.
• Emit every transition. Spawning, registering, ready, acquired, released, draining, lost, terminated — all of these go to the orchestration log with timestamp and reason. This is the only way to debug "why did my plan stall."
• Pin capability versions in plans, not agent IDs. A plan that names rag.search@>=2,<3 survives a rolling redeploy. A plan that names search-agent-7 does not.
• Backpressure at find_by_capability. When zero candidates are available, return empty and let the supervisor decide: queue, spill to a more expensive class, or scale the Deployment up.

Pitfalls

• Heartbeat too aggressive. A 1-second interval with a 3-miss rule will declare every agent dead during a GC pause. Start at 10s/30s and tune down only if you have evidence.
• Forgetting to remove agents from the capability index on terminate. Stale entries cause find_by_capability to return ghosts; acquire then races on a dead pod. Always remove from _cap_index inside _remove.
• Treating Busy as a hard lock. If max_concurrency > 1, an agent can be acquired again while busy. Track in-flight count, not a boolean.
• Scaling pods without checking pending registrations. A burst of find_by_capability misses can trigger a scale-up storm if you do not debounce. Coalesce scale requests over a short window (e.g. 5s).
• Silent capability drift. Two replicas of the same Deployment advertising different capability versions because the rollout is mid-flight is normal, but the manager must surface it. Log a warning when a capability name has more than one active version simultaneously, and let the supervisor's version pin do the routing.

With this manager in place, the rest of the orchestration platform — supervisor hierarchies, task queues, retry policies — has a stable substrate to build on. Tasks find capable agents, capable agents prove they are alive, and dead ones are reaped without human intervention.

Code Walkthrough

The agent pool as a first-class abstraction

Think of the pool the way a thread pool works in a server runtime, but each "thread" is a specialised mini-service: a search-agent that knows how to issue web queries and summarise, a code-agent running a sandboxed interpreter, a math-agent wrapping a symbolic solver, a retrieval-agent bound to a particular vector index, an image-gen-agent calling a diffusion model. They share no logic — only a contract. That contract is what makes them poolable.

The orchestrator never names a specific agent by ID when planning a task. It names a capability and lets the manager pick. This indirection is what lets us scale, version, and fail over without rewriting plans.

Capability registration: the join handshake

When an agent process starts, the very first thing it does — before accepting any task — is register itself. The registration payload is a self-description: who I am, what I can do, what I cost, how much I can chew at once.

A minimal capability descriptor looks like this:

Code snippetpython
1from pydantic import BaseModel, Field
2from typing import Literal
3
4class Capability(BaseModel):
5    name: str                     # e.g. "rag.search", "code.python.exec"
6    version: str                  # semver, e.g. "2.1.0"
7    input_schema: dict            # JSON Schema for the task payload
8    output_schema: dict
9    max_concurrency: int = 1      # tasks this agent will run in parallel
10    cost_class: Literal["A", "B", "C"] = "B"   # A=cheap, C=expensive
11
12class AgentRegistration(BaseModel):
13    agent_id: str                 # unique, e.g. pod name
14    deployment: str               # k8s Deployment name
15    endpoint: str                 # gRPC/HTTP address
16    capabilities: list[Capability]
17    heartbeat_interval_s: int = 10
18    metadata: dict = Field(default_factory=dict)

The manager indexes registrations by (capability.name, capability.version). Because the same agent may advertise multiple capabilities — a code-agent that does both code.python.exec and code.python.lint — the index is a multimap, not a 1:1 mapping.

Why version the capability, not the agent

Agents are cattle; capabilities are contracts. If you bump the agent image from 1.4.7 to 1.4.8 for a bug fix, the supervisor does not need to know. But if you change the input schema for rag.search — say, you add a required tenant_id field — that is a breaking contract change and the capability version must move to 3.0.0. Plans pinned to >=2,<3 will route to the old fleet until they are migrated. This is the same discipline as API versioning, just inside the agent mesh.

The lifecycle state machine

Every agent moves through a fixed set of states. The manager is the only writer of these transitions, and every transition is appended to the orchestration log so that post-mortems and replays are possible.

Two transitions deserve a closer look. Drain is graceful: the agent stops accepting new work but finishes whatever it is currently doing, then exits. This is what runs during a rolling deploy. Heartbeat-lost-to-Terminated is the brutal path: if the watcher has not seen a ping within the agreed interval (typically 3 missed heartbeats), the manager declares the agent dead, removes it from the registry, and asks Kubernetes to confirm the pod is gone or to delete it.

The AgentPoolManager class

Here is the core of the manager. It is intentionally small: registration, lookup, acquire/release with concurrency tracking, and drain. Heartbeat handling and the K8s scaler are wired in around this skeleton.

Code snippetpython
1import asyncio
2import time
3from collections import defaultdict
4from packaging.version import Version, InvalidVersion
5
6class AgentPoolManager:
7    def __init__(self, k8s_client, log_sink):
8        self._agents: dict[str, AgentRegistration] = {}
9        self._state: dict[str, str] = {}                 # agent_id -> state
10        self._inflight: dict[str, int] = defaultdict(int)
11        self._last_seen: dict[str, float] = {}
12        self._cap_index: dict[str, set[str]] = defaultdict(set)  # cap_name -> agent_ids
13        self._lock = asyncio.Lock()
14        self._k8s = k8s_client
15        self._log = log_sink
16
17    async def register(self, reg: AgentRegistration) -> None:
18        async with self._lock:
19            self._agents[reg.agent_id] = reg
20            self._state[reg.agent_id] = "Ready"
21            self._last_seen[reg.agent_id] = time.time()
22            for cap in reg.capabilities:
23                self._cap_index[cap.name].add(reg.agent_id)
24            await self._log.emit("agent.registered", reg.agent_id,
25                                 caps=[c.name for c in reg.capabilities])
26
27    async def find_by_capability(self, name: str, version_spec: str = ">=0") -> list[str]:
28        async with self._lock:
29            candidates = []
30            for aid in self._cap_index.get(name, set()):
31                if self._state.get(aid) != "Ready":
32                    continue
33                reg = self._agents[aid]
34                cap = next(c for c in reg.capabilities if c.name == name)
35                if _matches(cap.version, version_spec):
36                    if self._inflight[aid] < cap.max_concurrency:
37                        candidates.append(aid)
38            # cheapest first, then least loaded
39            candidates.sort(key=lambda a: (
40                self._cost_rank(a, name), self._inflight[a]))
41            return candidates
42
43    async def acquire(self, agent_id: str) -> None:
44        async with self._lock:
45            reg = self._agents[agent_id]
46            cap_max = max(c.max_concurrency for c in reg.capabilities)
47            if self._inflight[agent_id] >= cap_max:
48                raise RuntimeError("agent saturated")
49            self._inflight[agent_id] += 1
50            if self._inflight[agent_id] == cap_max:
51                self._state[agent_id] = "Busy"
52            await self._log.emit("agent.acquired", agent_id)
53
54    async def release(self, agent_id: str) -> None:
55        async with self._lock:
56            self._inflight[agent_id] = max(0, self._inflight[agent_id] - 1)
57            if self._state.get(agent_id) == "Busy":
58                self._state[agent_id] = "Ready"
59            await self._log.emit("agent.released", agent_id)
60
61    async def drain(self, agent_id: str) -> None:
62        async with self._lock:
63            self._state[agent_id] = "Draining"
64            await self._log.emit("agent.draining", agent_id)
65        # wait for in-flight to settle, then scale the deployment down
66        while self._inflight[agent_id] > 0:
67            await asyncio.sleep(0.5)
68        await self._k8s.delete_pod(self._agents[agent_id].deployment, agent_id)
69        async with self._lock:
70            self._remove(agent_id)
71            await self._log.emit("agent.terminated", agent_id, reason="drain")
72
73    async def heartbeat_watcher(self) -> None:
74        # wakes every second; anything older than 3 * heartbeat_interval_s is dead
75        while True:
76            now = time.time()
77            async with self._lock:
78                dead = []
79                for aid, reg in self._agents.items():
80                    if self._state.get(aid) in ("Terminated", "Draining"):
81                        continue
82                    if now - self._last_seen[aid] > 3 * reg.heartbeat_interval_s:
83                        dead.append(aid)
84            for aid in dead:
85                await self._log.emit("agent.lost", aid)
86                await self._k8s.delete_pod(self._agents[aid].deployment, aid)
87                async with self._lock:
88                    self._remove(aid)
89            await asyncio.sleep(1.0)

The find_by_capability method is where the cost class actually pays off: when several agents qualify, the manager prefers the cheapest cost class first and the least-loaded agent within that class. That single line of policy is enough to keep expensive GPU-backed agents idle until the cheap ones are saturated.

Heartbeats and the watcher loop

The heartbeat_watcher method shown above runs as a background coroutine, waking every second to check _last_seen. Anything older than 3 * heartbeat_interval_s is presumed dead: the manager emits agent.lost, asks Kubernetes to delete the pod, and prunes it from the registry. The agent itself posts POST /heartbeat against the manager on its declared interval; the handler simply updates _last_seen[agent_id] = time.time(). Cheap, idempotent, and survives network blips up to the 3-miss threshold.

Do's and Don'ts

Do's

1. ✓Do version capabilities as contracts, not agent images — when rag.search gains a required tenant_id field in its input_schema, bump capability.version to 3.0.0 so plans pinned to >=2,<3 keep routing to the old fleet without a single plan rewrite; a pod image bump from 1.4.7 to 1.4.8 for a bug fix requires no version change at all.
2. ✓Do make AgentPoolManager the sole writer of lifecycle state transitions — every move through the state machine (Registering → Ready, Ready → Draining, heartbeat-lost → Terminated) must be appended to the orchestration log with timestamp and reason, because that append-only log is the only artifact that can answer why a plan stalled after a pod died mid-flight.
3. ✓Do gate acquire() on the per-agent _inflight counter rather than the Busy state label — an agent with max_concurrency > 1 stays in Busy while running concurrent tasks, so blocking new work the moment the label flips starves capacity that max_concurrency explicitly licensed; compare _inflight[agent_id] < cap.max_concurrency, not state != "Busy".

Don'ts

1. ✗Don't bind orchestration plans to specific agent IDs like search-agent-7 — use find_by_capability('rag.search', '>=2') so that pod restarts, rolling deploys, and Kubernetes-driven scale events do not strand in-flight work; the entire point of the capability registry indirection is that the supervisor never names the agent.
2. ✗Don't omit the _cap_index cleanup inside _remove — if a dead agent's ID lingers in the cap.name → set[agent_ids] multimap after termination, find_by_capability will hand the caller a pod that no longer exists, and the subsequent acquire() call silently fails or hangs with no obvious trace back to the stale index entry.
3. ✗Don't set heartbeat_interval_s to 1s with a 3-miss rule — a routine GC pause in an agent process can exceed 3 seconds and trigger a mass-eviction of healthy pods; the state machine's Ready → Terminated edge on heartbeat loss is intentionally blunt, so start at the lesson's reference point of 10s per beat (30s total window) and tighten only after measuring actual pause distributions.