Middleware pipeline II · adjudication and cleanup after model output

Last stop, the request crossed before_agent, before_model, and wrap_model_call, and the model finally got a tidied request and spoke. This stop walks the other half: how model output is adjudicated, allowed into tool execution, and cleaned up at run end.

The inbound half was preparation — laying out state, messages, and tool schemas. The return half is adjudication and cleanup after model output — is this result trustworthy? Should we interrupt? Can the tool call run? What needs tidying when the run ends?

From last stop: before the call, middleware runs in registration order; after the return, it runs in reverse order. This stop covers that second half: after the model returns an AIMessage, later-registered middleware handles it earlier and gets the first chance to rewrite the result.

flowchart TD
MODEL["model invoked · returns AIMessage"] --> AM["after_model · takes over this answer in reverse"]
AM -->|"has tool_calls: tools still needed"| WTC["wrap_tool_call · wraps each tool execution"]
WTC --> TOOL["tool runs · produces ToolMessage"]
TOOL --> FEED["ToolMessage added back to history"]
FEED -.->|"fed back to the model, opens the next turn"| MODEL
AM -->|"no tool_calls: already the final answer"| AA["after_agent · run cleanup, end"]

after_model: taking over the answer in reverse

after_model runs after every model return. It’s the first middleware layer the answer meets on the way out — but “first” means by actual execution order, not registration order.

Before reading on, fix one distinction in mind:

one run    = one big "model ↔ tools" loop, which may go around several times
one turn   = one invocation of the model node (before_model → model → after_model together make one turn)

The key point: after_model is the tail of each turn, not the tail of the whole run. If this turn’s AIMessage carries tool_calls, the tools run and their results are fed back to the model, opening the next turn; only when some turn no longer carries tool_calls and gives a final answer does after_agent close out the whole run. So after_model runs on every turn, while after_agent runs once per run.

Go back to last stop’s lead-only list; of the ones that implement after_model, top to bottom their registration order is:

Todo → TokenUsage → Title → SubagentLimit → LoopDetection → SafetyFinishReason

Reversed, that is the runtime execution order:

SafetyFinishReason   runs first
LoopDetection
SubagentLimit
Title
TokenUsage
Todo                 runs last

SafetyFinishReasonMiddleware is deliberately registered late in ⌖ agent.py::build_middlewares , and the comment says so outright — so that in the reverse dispatch it runs first, clearing untrustworthy tool_calls ahead of loop detection and subagent accounting. The following sections use that execution order.

SafetyFinishReason comes from ⌖ safety_finish_reason_middleware.py::SafetyFinishReasonMiddleware . It is not a safety classifier; it doesn’t judge whether content is safe. It consumes the stop signal the provider already attached:

OpenAI-compatible:  finish_reason = content_filter
Anthropic:          stop_reason = refusal
Gemini:             finish_reason = SAFETY / BLOCKLIST / PROHIBITED_CONTENT / ...

The field names and values differ per provider, so a SafetyTerminationDetector adapter normalizes them into one internal SafetyTermination (recording which detector fired, and the field and value it matched) — downstream logic never has to care which provider it was.

The trouble starts when the model is safety-terminated yet still carries a half-formed tool_calls: LangGraph routes any AIMessage with non-empty tool_calls into the tools node. A truncated tool call may hold only half an argument — say a write_file whose content stops mid-stream. Executing it means writing a file with mangled arguments.

So in after_model, SafetyFinishReason has one core action: on “safety termination + non-empty tool_calls,” it clears tool_calls, appends a user-visible explanation, and lays down three observability signals at once — a safety_termination entry in this AIMessage’s additional_kwargs, a live SSE event to the frontend so the UI can reconcile its state, and a persistent middleware:safety_termination audit record for later debugging. It pairs cleanly with last stop’s LLMErrorHandling:

LLMErrorHandling:    never got a normal AIMessage (the provider raised)
SafetyFinishReason:  got a normal AIMessage, but its tool_calls aren't trustworthy

LoopDetection comes from ⌖ loop_detection_middleware.py::LoopDetectionMiddleware . First it has to decide what counts as a “loop,” and that has two layers:

Layer 1 · the same action over and over
  normalize each tool_call into a hash (tool name + key args)
  read_file buckets nearby line ranges together, so being off by a few lines
  doesn't count as two different calls
  content-editing tools like write_file and str_replace use the full args instead,
  so distinct edits aren't misread as repeats
  a batch of tool_calls is sorted before hashing, so the hash is the same even if
  the call order changes
  the same hash showing up enough times → hit

Layer 2 · the same kind of action too many times
  count by tool name, e.g. read_file called dozens of times
  even on different files each time, that's a kind of spinning → hit
  naturally high-frequency tools can raise or lower their threshold per tool name,
  to avoid false positives

Once it hits, the response splits into two levels of different severity. One middleware, two hooks, each holding half the brake: the soft warning lives in wrap_model_call (covered last stop — it has to wait for the tool results to come back so the reminder can be placed safely after the ToolMessage, without breaking the pairing protocol); this stop only looks at the other half, the hard stop in after_model.

The hard stop works like this: once the repetition reaches the hard limit, it rewrites the current AIMessage outright — clears the structured tool_calls, wipes the raw tool-call payload in additional_kwargs too (last stop’s DanglingToolCall noted that tool-call info can live in both places, so clearing only one isn’t enough), appends a [FORCED STOP] line, and changes finish_reason from tool_calls to stop. After that, the message is ordinary assistant text, and the graph won’t route into the tools node anymore.

SubagentLimit comes from ⌖ subagent_limit_middleware.py::SubagentLimitMiddleware . It only looks at the latest AIMessage (the one the model just emitted): count the task tool calls inside, and if there are more than max_concurrent, truncate the excess and swap it back under the same message id (same id → the LangGraph reducer replaces in place rather than appending). It leaves history unchanged, because a task call in history may already have its matching ToolMessage result, and rewriting it would break the pairing protocol. It sits in after_model, after the model has emitted tool_calls and before the tools run.

Prompting the model to launch fewer subagents is not enough. A resource limit cannot rely on the model obeying; this layer enforces the cap.

Title and TokenUsage are the “bypass” of this layer. Title ( ⌖ title_middleware.py::TitleMiddleware ) generates a thread title after the first complete exchange — reads messages, writes the title field, falls back to a local title on failure, never dragging the main run down. TokenUsage ( ⌖ token_usage_middleware.py::TokenUsageMiddleware ) stamps this step’s token usage, plus the actions classified from tool_calls (todo_start, subagent, search…), onto the AIMessage’s additional_kwargs for billing and analytics. Both write back with model_copy (make a new message, change only the fields that need changing, keep the original id). Why copy instead of mutating the original in place? Because that message may be shared by reference with the checkpoint, the streaming output, or later middleware, and an in-place edit would affect them; copying a fresh one that keeps the original id, then relying on “same id → the LangGraph reducer replaces in place,” supersedes the original rather than adding a new message.

TokenUsage also folds subagent usage back into the parent message. When the task tool dispatches a subagent, it caches that subagent’s usage keyed by tool_call_id; when TokenUsage sees the matching ToolMessage in after_model, it pops the cached value by that id and merges it into the dispatching AIMessage — pop is read-once-and-remove, so the same usage isn’t counted twice.

Todo in after_model does “premature-exit prevention”: if the model produces a clean final answer while todos are still incomplete, Todo queues a completion reminder and returns {"jump_to": "model"} to bounce control back to the model; the next wrap_model_call then injects that hidden reminder into the request. The reminder is deliberately not written into graph state — it’s a control prompt for the next model call, not part of the user-visible conversation. (Todo’s other two hooks were covered last stop; this is just the outbound piece.)

wrap_tool_call: the tool boundary

If, after after_model, this turn still calls tools (the AIMessage has tool_calls), control enters the tools node. Every tool execution passes through wrap_tool_call — the same call-chain wrapping pattern as wrap_model_call, except this time it wraps the tool, not the model.

The handler semantics are the same: handler(request) means “continue the call chain and eventually run the tool.” Before calling handler, a middleware can stop the tool from running (returning a ToolMessage short-circuits it); after handler returns, it can post-process the tool’s result.

Outer to inner, the middleware on this layer are:

ToolOutputBudget   (outermost)
Sandbox
Guardrail?
SandboxAudit
ToolErrorHandling
DeferredToolFilter
Clarification      (innermost)

Two groups make it clearer: those that decide before handler whether the tool runs at all, and those that handle the result after.

Clarification comes from ⌖ clarification_middleware.py::ClarificationMiddleware . The behavior of ask_clarification lives in the middleware. When the model calls it, wrap_tool_call intercepts, builds a ToolMessage, and returns:

Command(update={"messages": [tool_message]}, goto=END)

That turns a “question” into a run-level pause: the graph goes straight to END and waits for the user’s answer. The frontend keys on message.name == "ask_clarification" to render it as a clarification card rather than an ordinary tool result.

Here’s an easily-misread boundary: wrap_tool_call sees only one tool call at a time. If the model emits [a normal tool, ask_clarification] in the same AIMessage, the goto=END that Clarification returns does not guarantee the sibling tool is cancelled — in the current LangGraph version, multiple tool calls from one message run together (async uses asyncio.gather). “Clarify first, and only clarify” is the prompt asking ask_clarification to come alone, not a runtime guarantee.

So once that turn’s tools have run, END vs. back-to-the-model is decided by a rule on create_agent’s “tools → model” edge: it only goes to END when every client-side tool that ran this turn is return_direct. ask_clarification is return_direct=True, so when it’s the only call, the whole run stops cleanly and waits for the user. But mix in one ordinary tool and “all return_direct” no longer holds, so the graph carries both tools’ results back to the model — and clarification did not stop the sibling tool path. There’s a subtler constraint too: ask_clarification’s question and options are generated before that turn’s tools run, so it can’t ask anything that depends on the sibling tool’s return value — those results can’t matter until the next turn back at the model.

DeferredToolFilter’s tool side comes from ⌖ deferred_tool_filter_middleware.py::DeferredToolFilterMiddleware . Last stop covered its model side: hiding unpromoted schemas from request.tools. But hiding a schema isn’t enough — the model may still call a hidden tool from history, a prompt leak, or a guess. So on the wrap_tool_call side it adds enforcement: if a hidden tool is called, it returns ToolMessage(status="error") and doesn’t execute. The two sides together are fail-closed: the model side controls what’s visible, the tool side controls what can run. (Last stop already abstracted this into a ToolVisibilityPolicy.)

SandboxAudit comes from ⌖ sandbox_audit_middleware.py::SandboxAuditMiddleware . It only handles tool calls named bash: read args["command"], validate the input first, then grade the risk with regex / shlex, and pick one of three:

pass:   everything else — run as usual
warn:   matches a medium-risk pattern (pip install, apt install, chmod 777,
        sudo/su, PATH=…) — run as usual, but append a warning to the result
block:  empty / too long / contains a null byte, or matches a high-risk pattern
        (rm -rf /, dd if=, mkfs, curl | bash, /dev/tcp, fork bomb…)
        — don't call handler, return an error, the command never runs

The key safety property is one line: block returns before handler. Because execution lives inside the handler chain, a blocked command never reaches it. This isn’t a full sandbox; it’s a bash command policy gate, and every bash call — whatever the verdict — leaves an audit log carrying the timestamp, thread_id, command, and verdict (pass / warn / block).

Guardrail (optional) comes from ⌖ middleware.py::GuardrailMiddleware , installed only when a provider is configured. It hands the tool call (name, arguments, agent identity) to a GuardrailProvider for an allow / deny decision: deny skips handler and returns an error. It has a different role from static allowed_tools — allowed_tools decides “which tools the model can see/use in principle,” Guardrail decides “this concrete call, with these arguments, right now.” Typical enterprise use: allow read_file but deny path=/secrets/*, allow send_email but deny external recipients. Today it’s still a reserved extension point; the value depends on plugging in an enterprise policy provider.

Now the layers that tidy the result after handler.

ToolErrorHandling comes from ⌖ tool_error_handling_middleware.py::ToolErrorHandlingMiddleware . It wraps tool execution and turns ordinary exceptions into protocol-valid results:

try:
    result = handler(request)
except GraphBubbleUp:
    raise               # control-flow signal, re-raised as-is
except Exception as exc:
    return ToolMessage(status="error", ...)

Why bother appending a status="error" ToolMessage at all, instead of letting the exception propagate? Two reasons. One is protocol: the provider requires every AIMessage(tool_calls) to have a matching ToolMessage. The other is resume hygiene — if a checkpoint holds a tool_call with no result, on resume you can’t tell whether it never ran, failed, or lost its result; an explicit error result makes the state self-consistent, and the next model turn can recover from it.

The GraphBubbleUp branch matters: it’s LangGraph’s interrupt / pause / resume signal, not a tool failure. Swallow it into an error and you’ve disguised control flow as a business failure. It also normalizes task (the subagent-entry tool) results, stamping a structured subagent_status (completed / failed / timed_out) onto the ToolMessage so the frontend doesn’t have to parse strings like “Task Succeeded.”

Sandbox’s wrap_tool_call side comes from ⌖ middleware.py::SandboxMiddleware . Last stop noted sandbox defaults to lazy: before_agent usually doesn’t acquire; acquire happens on the first sandbox-tool call. The catch is that mutating runtime.state["sandbox"] inside a tool is local to that invocation and won’t flow back into graph state on its own. Sandbox’s wrap_tool_call detects “a new sandbox_id appeared during this call” and wraps it into a Command committed back to graph state:

Command(update={"sandbox": {"sandbox_id": sandbox_id}, "messages": [tool_message]})

That’s how checkpoint / resume can see the sandbox identity.

ToolOutputBudget’s tool side comes from ⌖ tool_output_budget_middleware.py::ToolOutputBudgetMiddleware . Last stop covered its model side (swapping oversized historical ToolMessages for a preview). On the wrap_tool_call side it does source control. To be precise about its job: it governs the size of ToolMessage.content, catching the case where a tool stuffs a big log, HTML page, or grep result straight into ToolMessage.content. When it spots one, it externalizes the full content to a file and leaves only a head/tail preview plus the path in the message.

Where that file lands depends on whether this call has a sandbox, and whether the sandbox mounts the thread directory:

no sandbox                              → write to the host's outputs dir
sandbox that mounts the thread dir      → also the host's outputs (the sandbox sees it
                                          through the same virtual path)
sandbox without the thread dir mounted  → write straight into the sandbox via
                                          sandbox.write_file

(It reads the capability flag the provider reports, rather than guessing from the provider’s name.) Note it also handles ToolMessages inside Command.update, so a large output can’t slip past the budget through the Command path.

after_agent: closing out the run

When this turn no longer has tool_calls (the model gave a final answer), the graph heads for the end. after_agent runs once as the whole run wraps up, and does two main things.

Memory comes from ⌖ memory_middleware.py::MemoryMiddleware . It writes no ThreadState field directly; instead it pushes this turn’s conversation (user input, final answer, correction / reinforcement signals) into a global ⌖ queue.py::MemoryUpdateQueue . That queue is debounced, not a plain FIFO: for the same (thread, user, agent) within the debounce window it keeps only the latest snapshot and merges the yes/no flags like “correction / reinforcement.” So a session firing several runs in quick succession triggers just one memory update. Internally it uses two locks: one guards creation of the global singleton queue, the other guards a single queue’s enqueue / timer / processing state; the actual memory update, model calls, and storage writes all happen outside the locks, so they don’t stall the run’s close-out.

Sandbox release also runs in after_agent: sandbox calls provider.release(sandbox_id). But “release” doesn’t mean “destroy” — LocalSandboxProvider’s release is a no-op (local per-thread sandboxes are reused), while other providers may close a client, return a container to a pool, or destroy resources. The lifecycle semantics live in the provider, not in this middleware; the details wait for the sandbox stop ahead.

On top of that, Todo and LoopDetection also do cleanup in after_agent: clearing this run’s leftover reminder / warning bookkeeping. They’re long-lived middleware instances, so they have to wipe their own cross-run residue — symmetric with last stop’s before_agent cleanup. But the symmetry is offset — the two ends don’t clear the same data:

after_agent deletes this run’s own unused reminder. That “reminder” is the loop-detection soft warning from earlier: on spotting the model spinning, it queues a “you’re repeating, wrap up” line, meaning to inject it into the request at the next model call. But if this turn gives a final answer outright and the run ends on the spot, that line never gets its “next model call” and shouldn’t stay in the queue; after_agent deletes it. before_agent instead deletes another run’s leftover — because a run may crash or be cancelled and never reach after_agent, a new run clears the stale entries others left in the shared bookkeeping before it starts. One cleans up after itself, the other backstops for others, so both are needed — it’s not the same thing done twice.

Easy-to-misread points

• To make a middleware run first on the way out, register it last. after_model dispatches in reverse — the later registered, the earlier it runs — which is exactly how Safety ends up first. Register it earlier and it runs last instead.
• Before rewriting a message on this layer, make sure the message id is unchanged. SubagentLimit, TokenUsage, SafetyFinishReason, and LoopDetection all rely on “reuse the original id → the reducer replaces in place.” Once the id changes, the reducer treats it as a new message and appends, duplicating tool_calls and breaking the pairing protocol.
• goto=END doesn’t guarantee cancelling sibling tools; whether it terminates depends on return_direct. It stops at END only when every client-side tool this turn is return_direct; ask_clarification mixed with an ordinary tool flows back to the model. “Only clarify” is a prompt convention, not a runtime guarantee.
• A soft warning can neither be injected immediately nor counts as already intercepting. The matching ToolMessage doesn’t exist yet, so injecting now gets rejected by the provider; it must be queued and injected at the next wrap_model_call. And it only appends a reminder — the tool still runs. The hard stop is what prevents the graph from entering the tools node. Conflate the two and you’ll mistake “loop detected” for “already terminated.”
• SafetyFinishReason doesn’t judge content safety. It only consumes the stop signal the provider already gave (content_filter / refusal / SAFETY…); relying on it to intercept harmful content is a misuse.
• Middleware that wraps tools must let GraphBubbleUp through. It’s LangGraph’s interrupt / resume signal, not a tool failure; if an except swallows it too, interrupt / resume break. When catching tool exceptions, always except GraphBubbleUp: raise first.
• ToolOutputBudget governs the size of ToolMessage.content, not disk-file size. Writing a big file with a short returned result stays outside its scope; it governs big logs / HTML / grep output placed directly into content (ToolMessages carried inside Command.update are handled too).
• release ≠ destroy. after_agent calls provider.release(sandbox_id), but LocalSandboxProvider’s release is a no-op and the local sandbox is kept for reuse. Don’t assume the sandbox is gone just because the run ended.

Put the two stops together and the full middleware picture of a run is clear: before the model call is preparation; after the model returns is adjudication. The piece touched again and again but never opened up is the sandbox — both the tool’s execution environment and the thing carrying a whole acquire / persist / release lifecycle. The next stop goes inside it: what local, container, and Kubernetes sandboxes isolate, and what they do not isolate.