Skip to content

agent-loop.md

Latest commit

 

History

History
219 lines (178 loc) · 11.2 KB

agent-loop.md

File metadata and controls

219 lines (178 loc) · 11.2 KB
Raw

Agent runtime — state machine, budgets, resumability (Step 3.6)

Public users wanting to use the agent surface should start at reference/agent.md and guides/agent-quickstart.md. This document is for contributors: it explains why the runtime is shaped the way it is, and the invariants a reviewer should hold it to.

Why a new package

The Step 2.11 AgentLoopV0 spike proved iterative retrieval works but was a fixed for loop with heuristic stop conditions living inside rag-retrieval (see spikes/agent-loop-v0-gaps.md). Step 3.6 graduates it into a first-class runtime with three properties the spike lacked: a pluggable decision policy, first-class budget enforcement, and per-step resumability.

rag-agent is a new package depending only on rag-core. It owns orchestration behaviour; the wire types live in rag_core.agent_types so the gateway surfaces, SDKs, and the eval harness share one definition without a dependency cycle. The dependency graph stays acyclic: gateway → agent → core; retrieval and policy enter through structural seams, never imports.

The state machine

Each iteration of the loop is one step that walks one of two paths:

            ┌──────────── planning ────────────┐
            │                                   │
        CallTool                             Finish
            │                                   │
         acting                            finalizing → done
            │
        observing ──► (next step) planning
  • planning — the Controller inspects the AgentSnapshot + tool specs and returns an AgentAction (CallTool or Finish). This is where the LLM-as-controller call happens in production.
  • acting — the chosen tool is invoked.
  • observing — the tool result is folded into the accumulated chunk set (deduped by chunk_id) and the transcript.
  • finalizing — the controller produced an answer (or a ceiling forced one); egress governance + citation assembly run here; the run terminates done.
  • failed — a non-recoverable orchestration fault (unknown tool, a tool raising a non-RetrievalError) terminates the run.

Why an explicit state machine over the spike's for loop: each transition is a distinct, observable event (phase_changed, tool_started, …), the loop can be suspended and resumed at any step boundary, and stop conditions are checked at one well-defined point rather than scattered through the body.

Stop conditions — checked at the top of each step

A budget exhausted by step N must be observed before step N+1 does any work, so the checks live at the top of the loop body, in order:

  1. Any of the four Budget dimensions has hit zero → the matching StopReason (budget_tokens / budget_dollars / budget_wall / budget_iter).
  2. step_index >= max_steps (the config ceiling) → StopReason.max_iterations.

Either forces a clean finalize. The controller's own Finish yields StopReason.final_answer. The loop asserts a stop reason is always set before finalizing — exiting with stop_reason is None is an invariant violation that raises AgentLoopError.

Budget model

Budgets are four independent nullable dimensions (max_tokens, max_dollars, max_wall_ms, max_iter); None means uncapped on that dimension. Each step the loop spends against the working Budget:

  • wall_ms is measured from a monotonic clock.
  • tokens / dollars use the AgentConfig flat per-step estimates when the controller/tools don't return measured costs. This is deliberate: a token-capped run against cost-free noop backends must still terminate. A production LLMController can return measured usage to replace the estimate.

The terminal AgentRunResult.budget_consumed (BudgetSpend) is the delta between the budget as it stood when this drive began and where it ended — so on a resume it attributes spend to the resuming invocation, not the whole run.

Two ceilings, on purpose

max_steps (config) and budget_iter (per-request) are distinct. max_steps is defence-in-depth — always present, even when a caller sets no budget, so a misbehaving controller can't spin forever. budget_iter is an operator's explicit per-request iteration cap. Whichever bites first wins and sets a distinct StopReason, so an operator can tell why a run stopped.

Resumability

The whole run state is one frozen, serialisable AgentSnapshot: identity, phase, step_index, steps, accumulated chunk_refs + seen_chunk_ids, final_answer, status, and the remaining budget. After every step (when checkpoint_each_step) the loop hands the fresh snapshot to the CheckpointStore; resume(ctx, run_id) / AgentRequest.resume_run_id reload the latest snapshot and continue the state machine — resume is literally "load, then keep stepping."

Two design details worth calling out:

  • Last-write-wins per run_id. A snapshot already encodes the full state (step history lives in steps), so the store keeps only the latest row. This keeps the SPI tiny — save / load / delete — and makes resume a single load.
  • Budget stored as raw caps, not a nested Budget. Budget rejects zero caps at construction (invariant G-03), but a spent budget legitimately reaches zero — and a JSON round-trip on resume would re-run that validator and reject the snapshot. AgentSnapshot stores the four raw nullable caps and rebuilds via Budget.model_construct (validator-bypassed) in budget().

InMemoryCheckpointStore keys by (tenant_id, run_id) so a shared instance can't leak one tenant's run state to another even if run ids collide. It is not durable across restarts; production swaps a Postgres/Redis-backed store with no loop change.

Controller seam

The loop never decides what to do next — that is the Controller's job (decide(ctx, *, snapshot, tools) -> CallTool | Finish). This split keeps the loop pure orchestration and makes it trivially testable: ScriptedController replays a fixed action queue, so every phase sequence, budget edge, and checkpoint path is exercised deterministically with no model. HeuristicController is the credential-free default (retrieve once, then finalize); LLMController is production (JSON-action prompt, tolerant parse degrading to Finish).

Tool seam + retrieval governance

Tools are the only way the loop reaches the outside world. RetrieveTool wraps a structural Retriever Protocol — the same kind of seam the spike used for query understanding — so rag-agent never imports rag-retrieval. The gateway supplies the concrete _GatewayRetriever, which runs query understanding then the policed corpus/retrieval router. Because that path's HybridRetriever is the read_chunk PDP call site, retrieval governance is enforced identically to /v1/query without the loop knowing anything about policy.

A backend RetrievalError is recoverable at the agent layer: RetrieveTool returns AgentToolResult(ok=False, error=…) so the controller can retry, pick a different tool, or finalize from what it has — the loop degrades to "answer from what I gathered" rather than aborting the whole run on one bad fetch. Only a non-recoverable fault (unknown tool, a tool raising something other than RetrievalError) raises AgentLoopError and trips failed.

Egress governance is the gateway's job

The loop emits the raw final_answer. The final answer is text leaving the platform, so each surface consults PolicyEngine.egress_text over it at the boundary — once, when the answer_delta frame is reached, caching the decision:

  • deny → withhold the text; the denial notice goes on the wire, citations are cleared on the terminal result.
  • transform → substitute the redacted answer everywhere.
  • allow → pass through.

Running the check before the answer_delta frame is emitted guarantees denied text never reaches the wire. REST (apps/gateway/src/rag_gateway/agent.py) and gRPC (grpc_service.py Converse) share the exact same helpers (_apply_egress_policy, _redact_result, _failed_event) so the two surfaces behave identically.

Surfaces serialise the same stream

AgentEvent is a single flat model rather than a discriminated union of one class per kind. This keeps both the SSE JSON encoding and the gRPC ConverseStreamEvent mapping mechanical and keeps proto-compat parity field-for-field (enforced by tests/contract/test_grpc_proto_compat.py): the proto message mirrors the Pydantic model exactly, so the two surfaces can never silently drift. REST closes the stream with data: [DONE]; gRPC ends the server stream naturally. A resume-of-unknown-run failure is surfaced in-band as a run_failed frame on both surfaces, not as a transport-level abort.

Invariants (reviewer checklist)

  • rag-agent imports only rag-core — no rag-retrieval, rag-policy, rag-gateway, or backend imports.
  • The loop makes no governed SPI call directly; retrieval governance is in the wired Retriever, egress governance is at the gateway boundary. loop.py is not on the tests/policy/coverage.py allowlist; controller.py is (it calls llm.complete).
  • Stop conditions are checked at the top of the step; the loop never exits without a StopReason.
  • max_steps (config) and budget_iter (request) remain distinct, each with its own StopReason.
  • Every step persists a snapshot when checkpoint_each_step; resume is a single load + continue.
  • Budget is stored as raw caps in the snapshot (validator-bypassed on rebuild) so a spent (zero) budget round-trips on resume.
  • A backend RetrievalError is recoverable (ok=False); only non-recoverable faults raise AgentLoopError → failed.
  • AgentEvent stays a flat model and field-for-field aligned with ConverseStreamEvent (the proto-compat test must pass).
  • Denied egress text never reaches the wire (check runs before answer_delta); REST + gRPC reuse the same egress helpers.

Deferred

  • Durable checkpoint store. Only InMemoryCheckpointStore ships; the Postgres/Redis-backed store lands with the durable-runtime work.
  • Measured token/dollar accounting. The default uses flat per-step estimates; wiring real usage from LLMController / tools is a follow-up.
  • More tools. Only retrieve ships; web fetch, calculators, and sub-agents are future steps behind the same Tool seam.
  • MCP / OpenAI agent-mode. Step 3.6's surfaces are REST SSE + gRPC Converse only (ADR-0011 named Converse as the agent-loop home). Exposing the runtime through the MCP tool surface or an OpenAI-assistant dialect is out of scope.

See also