Run Middleware

Run observers only watch a run. Middleware changes it: it sits around each model call and each tool call, so masking, guards, caching, argument rewriting, retries and logging live in one place instead of in every tool and every agent.

import { Agent, Cogitator, tripwire, type RunMiddleware } from '@cogitator-ai/core';

const timing: RunMiddleware = {
  name: 'timing',
  async wrapModelCall(call, next) {
    const started = Date.now();
    const response = await next(call);
    metrics.record(call.model, Date.now() - started);
    return response;
  },
};

const refundGuard: RunMiddleware = {
  name: 'refund-guard',
  beforeRun({ input }) {
    if (/refund/i.test(input)) tripwire('Refunds go to a person', { queue: 'billing' });
  },
};

const cog = new Cogitator({ middleware: [timing] });

const agent = new Agent({
  name: 'support',
  model: 'openai/gpt-6-luna',
  instructions: 'You help with orders.',
  middleware: [refundGuard],
});

await cog.run(agent, { input: 'Where is my order?', middleware: [auditFor(userId)] });

Where Middleware GoesCopy Anchor Link

The hooks of hooks run first, then the Cogitator's middleware, then the agent's, then the run's, then the built-in middleware closest to the model. before* and afterRun hooks run in that order, wrap* hooks nest with the first one outermost, and afterModel runs in reverse, so the middleware nearest the model sees its answer first.

Middleware is code, so the agent wire format does not carry it: the side that runs an agent gives its own.

HooksCopy Anchor Link

Every hook is optional. Each one gets the run's context as its last argument.

beforeRunCopy Anchor Link

beforeRun({ input, messages }, context) sees the input of a new run and the messages its first turn starts from: the system prompt, the history and the input, with what memory recalled for the turn and the brief in a <memory_context> block in front of it (the run's context and the insights of reflection are added after beforeRun). The input is not saved to memory and no model turn starts before it returns, so it is the place for input guards. Audio input is already transcribed into input by then. Return { context } and the text goes after the input as a system message, saved to the thread with it, so a later run reads it too.

startRun(event, context) runs before, once per run: event.source is startup for a run on a thread without history and resume for one with history or a resume() (event.resumed). Its { context } goes right after the system prompt and is not saved: each run gets its own.

Tool calls before the permissionsCopy Anchor Link

beforeToolCall({ toolCall, tool, resumed }, context) runs for each call right before it starts (after the calls before it in the step ran, or as each is dispatched when calls run side by side) and before the permissions decide it, and again for a call that waited in a paused run once resume() continues it. Return a ToolCallVerdict:

updatedArguments replaces the call's arguments: they are checked against the tool's schema and the permissions decide them, not the ones the model sent. allow skips only what the mode or the tool's own approval setting asks (in dontAsk and auto mode too, without the classifier), never a deny or ask rule or an ask that guards a deny rule. With several middleware, deny wins over defer, over ask, over allow, and the last updatedArguments counts. Decisions a hook makes are logged with decidedBy: 'hook'.

decideApproval({ request, tool }, context) answers a call that asks, before onApproval and the guardrails' onToolApproval: return a ToolApprovalDecision ({ approved: true }, with updatedArguments when the permissions are set up, or { approved: false, reason }) and nobody is asked. Return nothing to leave it to them.

afterToolBatch({ calls }, context) runs once every call of a step has its result, before the next model call. beforeStop({ output, continuations }, context) runs when the agent would stop: { continueWith } keeps it going with that message, saved as a user message. continuations counts how often in a row that happened since the agent's last tool call, continued whether it happened since the last user message, and the hooks are not asked once the run used every turn it may take.

beforeModelCopy Anchor Link

beforeModel({ messages, model }, context) runs before each turn, model as the agent names it (before any fallback). Return { messages } and the run goes on with them, in this turn and every later one. The built-in compaction works this way.

const reminder: RunMiddleware = {
  name: 'reminder',
  beforeModel({ messages }, context) {
    if (context.iteration !== 5) return;
    return { messages: [...messages, { role: 'user', content: 'Wrap up soon.' }] };
  },
};

wrapModelCallCopy Anchor Link

wrapModelCall(call, next, context) wraps one call to one model. call.request is the ChatRequest the backend gets, call.model the provider/model it goes to, and call.stream where streamed text goes when the run streams. Hand next the call, or a changed copy of it, once or more, or return a ChatResponse of your own without calling it. When a turn falls back to one of the agent's fallbackModels, the hook runs again for that model. Every call that reaches a backend counts toward the run's usage, cost and budgets and gets its llm.chat span, an answer the hook gives without calling the model counts nothing. In a streamed run such an answer reaches onToken in one piece.

const cache: RunMiddleware = {
  name: 'cache',
  async wrapModelCall(call, next) {
    const key = JSON.stringify(call.request.messages);
    const hit = answers.get(key);
    if (hit) return hit;
    const response = await next(call);
    answers.set(key, response);
    return response;
  },
};

To see or rewrite streamed text, hand next a call with a stream sink of your own that passes text on to call.stream.

afterModelCopy Anchor Link

afterModel({ response, messages, model }, context) sees each answer before the run uses it. Return { content } to replace its text, { toolCalls } to replace the calls it makes (an empty list turns it into the final answer). Text a streamed run already sent to onToken cannot be taken back: set holdsStream: true on the middleware and the run holds each turn's text until the turn is done and the hooks saw it, then sends the final text in one piece. The same goes for a wrapModelCall hook that retries a call or blocks its answer: without holdsStream the text of every call it sends streams as it arrives.

A run's own middleware sits inside the Cogitator's and the agent's, so its afterModel sees an answer before theirs do. To see each turn as the run keeps it, after every layer, use onModelTurn in the run options instead.

const noEmails: RunMiddleware = {
  name: 'no-emails',
  holdsStream: true,
  afterModel: ({ response }) => ({
    content: response.content.replace(/[\w.+-]+@[\w-]+\.[\w.]+/g, '[email removed]'),
  }),
};

wrapToolCallCopy Anchor Link

wrapToolCall(call, next, context) wraps each tool call that passed approval, also the calls invokeTool() makes for the Cogitator's middleware. call.toolCall is the call as the model made it, call.args the arguments as the tool's schema parsed them, call.approved whether a person approved it. Change the arguments by handing next a toolCall with new ones: they are checked against the schema and the tool's requiresApproval again, and a call they do not fit, or that would need an approval for them, answers the model with the reason. An approval never carries over to new arguments: the layers inside see approved: false. Each layer and the tool get their own copy of plain data (objects, arrays, maps, sets), so an edit in place changes only that copy. Other values (a Buffer, a URL, an instance of a class) are passed as they are: an edit inside them is seen through their own fields or toJSON, but it reaches the caller's object, and a change a class keeps in private fields is not seen. The tool, the call's id and name and approved cannot be changed, and onToolCall, RunResult.toolCalls and spans keep the call as the model made it. Return a ToolResult of your own to answer for the tool.

const normalizeIds: RunMiddleware = {
  name: 'normalize-ids',
  wrapToolCall(call, next) {
    const id = String((call.args as { id: string }).id)
      .trim()
      .toUpperCase();
    return next({ ...call, toolCall: { ...call.toolCall, arguments: { id } } });
  },
};

afterRunCopy Anchor Link

afterRun(event, context) runs once a run ended, before run() returns or throws: event.status is completed (after the run saved what it said to memory), paused (with the paused result) or failed (with event.error, an abort among them). A completed or paused run has its result, a failed one has none. model is the model the run ended on, as provider/model, after any fallback, and memoryAgentId the id the agent's long-term memory is kept under (memoryAgentId(agent): its configured id, else its name). It is the place for work that keeps memory in shape for the next run, which the caller waits for: the built-in extraction of memory.extraction and compaction of memory.compaction work this way, for completed runs only. The run is done by then, so an error a hook throws is logged as a warning and the next afterRun hook still runs, the caller keeps the result or the error. Skip slow work when context.signal is aborted, the caller has stopped waiting.

const archive: RunMiddleware = {
  name: 'archive',
  async afterRun(event, context) {
    if (event.status !== 'completed') return;
    await runs.insert({
      threadId: context.threadId,
      model: event.model,
      output: event.result.output,
    });
  },
};

The Run ContextCopy Anchor Link

addContext and stop are unset outside a run (invokeTool()), logEvent when there is no thread to log to. A run a middleware stopped completes with result.stopped: the reason and the name of the middleware that asked.

TripwiresCopy Anchor Link

A guard that finds something the run must not go on with calls tripwire(reason, info?). The run stops where it is and fails with GUARDRAIL_TRIPPED (HTTP 400): the reason, info and the name of the middleware in details. Unlike a tool error, which the model reads and works around, a tripwire ends the run.

import { isTripwire } from '@cogitator-ai/core';

try {
  await cog.run(agent, { input });
} catch (error) {
  if (isTripwire(error)) {
    console.log(error.details.middleware, error.details.reason, error.details.info);
  }
}

A tripwire raised inside a tool, such as a nested run of an agent used as a tool, stops the outer run too instead of reaching its model as a tool error. Any other error a hook throws fails the run as it is, and a hook still running when the run is cancelled or times out is no longer waited for. Outside a run, invokeTool() throws an error its wrapToolCall middleware throws, a tripwire among them, while a failing tool still comes back in the result's error. A hook inside wrapModelCall that throws a provider error (LLM_UNAVAILABLE and the like) moves the turn to the next fallback model, as the model's own failure would.

Built-in MiddlewareCopy Anchor Link

The runtime's own features are middleware too, added from the config closest to the model, in this order:

The hooks of hooks become one more, cogitator:hooks, the outermost layer of every run, so a hook sees a tool call as the model made it and its result as the model reads it. The compactions of context and memory.compaction ask beforeCompaction first and tell afterCompaction after.

Because PII masking is the innermost layer, it also masks whatever your middleware adds to a request, while your middleware sees the real values. The guardrails check the model's answer and a tool's result before your hooks see them: what your afterModel or wrapToolCall hooks change afterwards is not checked again. piiMiddleware(config), promptInjectionMiddleware(detector), guardrailsMiddleware(constitutionalAI), clearToolOutputsMiddleware(options) and compactionMiddleware(contextManager) build them, to give a feature to one agent only. memoryCompactionMiddleware(config, deps) builds the compaction after a run, given the memory, a summarizer per model and run, and the context window of a model, and memoryExtractionMiddleware(config, deps) the extraction of memory items after a run, given the memory and a summarizer per model and run. The extraction's afterRun runs before the compaction's, so the items of a thread are taken before it is compacted. The example below gives one agent PII masking:

import { piiMiddleware } from '@cogitator-ai/core';

const intake = new Agent({
  name: 'intake',
  instructions: 'Take the customer details.',
  middleware: [piiMiddleware({ detect: ['email', 'phone'] })],
});

See Security and Constitutional AI for what the built-in features check.