feat: add middleware system for wrapping agent stages

[zastrowm]

Port of strands-agents/sdk-typescript#1068 into the mono-repo.

Summary

Adds a middleware system that wraps agent stages using async generator handlers. Middleware controls flow (retry, cache, transform, short-circuit).

Public API: agent.addMiddleware(stage, handler) — returns a cleanup function.

Three built-in stages: InvokeModelStage, ExecuteToolStage, AgentStreamStage.

Motivation

Hooks let you observe operations and set flags, but they don't let you wrap them. If you want to do something both before and after a model call (timing it, adding a span, catching errors), hooks force you to manage state across two separate callbacks. With middleware you can wrap the entire invocation keeping state within your callback:

agent.addMiddleware(InvokeModelStage, async function* (context, next) {
  const start = Date.now()
  const result = yield* next(context)
  metrics.record(Date.now() - start)
  return result
})

Beyond the before/after pattern, middleware also makes several other use cases much more natural to express: caching, input transformation, short-circuiting, and error handling. All of these are awkward or impossible with hooks alone.

Public API Changes

import { Agent, InvokeModelStage, ExecuteToolStage, AgentStreamStage } from '@strands-agents/sdk'
import type { MiddlewareStage, MiddlewareHandler, MiddlewareNext } from '@strands-agents/sdk'

const agent = new Agent({ model, tools })

// Register middleware for any built-in stage
agent.addMiddleware(InvokeModelStage, async function* (context, next) {
  // pre-processing: inspect or transform context
  const modified = { ...context, messages: sanitize(context.messages) }
  // call next layer (or don't, to short-circuit)
  const result = yield* next(modified)
  // post-processing: inspect or transform result
  return result
})

Three stages ship with the SDK:

Stage	Wraps	Context fields
InvokeModelStage	Model call (between Before/AfterModelCallEvent)	messages, systemPrompt, toolSpecs, toolChoice, modelState
ExecuteToolStage	Single tool execution (between Before/AfterToolCallEvent)	tool, toolUse (name, id, input)
AgentStreamStage	Full agent.stream() output	args, options

Handlers are async generators and simple pass-through is return yield* next(context). Manual iteration of next() allows real-time event filtering or injection while not calling next at all short-circuits the operation.

Plugin Examples

class ToolResultCache implements Plugin {
  name = 'tool-result-cache'
  private readonly _cache = new Map<string, ToolResultBlock>()

  initAgent(agent: LocalAgent): void {
    const cache = this._cache
    agent.addMiddleware(ExecuteToolStage, async function* (context, next) {
      const key = `${context.toolUse.name}:${JSON.stringify(context.toolUse.input)}`
      const cached = cache.get(key)
      if (cached) return { result: new ToolResultBlock({ toolUseId: context.toolUse.toolUseId, status: cached.status, content: cached.content }) }
      const result = yield* next(context)
      cache.set(key, result.result)
      return result
    })
  }
}

class RetryOnThrottle implements Plugin {
  name = 'retry-on-throttle'
  initAgent(agent: LocalAgent): void {
    agent.addMiddleware(InvokeModelStage, async function* (context, next) {
      for (let attempt = 0; attempt < 3; attempt++) {
        try { return yield* next(context) }
        catch (e) {
          if (!(e as Error).message.includes('ThrottlingException') || attempt === 2) throw e
        }
      }
      throw new Error('exhausted retries')
    })
  }
}

class SystemPromptInjector implements Plugin {
  name = 'system-prompt-injector'
  constructor(private readonly _suffix: string) {}

  initAgent(agent: LocalAgent): void {
    const suffix = this._suffix
    agent.addMiddleware(InvokeModelStage.Input, async (context) => ({
      ...context,
      systemPrompt: `${context.systemPrompt ?? ''}\n\n${suffix}`.trim(),
    }))
  }
}

// Usage: inject safety guidelines into every model call
const agent = new Agent({
  model,
  systemPrompt: 'You are a helpful assistant.',
  plugins: [new SystemPromptInjector('Always cite your sources.')],
})

Middleware Interrupts

Middleware contexts expose interrupt() for human-in-the-loop gating:

agent.addMiddleware(ExecuteToolStage, async function* (context, next) {
  const { response } = context.interrupt<string>({ name: 'approve', reason: 'Confirm?' })
  if (response !== 'yes') return { result: new ToolResultBlock({ ... }) }
  return yield* next(context)
})

Returns MiddlewareInterruptResult<T> (wrapper) — allows non-breaking additions (cached data, metadata) as the interrupt system evolves.

Phase Sub-Stages (Input / Wrap / Output)

Each stage exposes three phases with a fixed execution order in order to solve the middleware ordering problem without explicit priority numbers:

• Input: Modifies the input of the phase as a pure function - take in the original input, return the modified input
• Output: Modifies the output of the phase as a pure function - take in the original output, return the modified output
• Wrap: Wraps the entire invocation of the stage - this is more akin to traditional middleware from express or other frameworks

// Input: transform context before execution (plain async function)
agent.addMiddleware(InvokeModelStage.Input, async (context) => ({
  ...context,
  systemPrompt: injectToSystemPrompt(context),
}))

// Output: transform result after execution (plain async function)
agent.addMiddleware(InvokeModelStage.Output, async (result) => {
  log(`stopReason=${result.result.stopReason}`)
  return result
})

// Wrap: full async generator wrap (same as just providing `InvokeModelStage`)
agent.addMiddleware(InvokeModelStage.Wrap, async function* (context, next) {
  return yield* next(context)
})

Execution order is fixed: all Input → all Output → all Wrap → terminal. A retry plugin on .Wrap retries subsequent Wrap handlers and the terminal — Input/Output transforms are not re-executed. A response logger on .Output never needs to coordinate registration order with a system prompt injector on .Input.

Key Decisions

• First registered = outermost: follows existing middleware conventions (from Express/Koa)
• Phase ordering is fixed: Input/Output/Wrap reduces the need for explicit priority management (for now - we'll probably need to add it later)
• Hooks fire unconditionally before/after middleware: existing behavior is unchanged
• Result wrapper types (InvokeModelResult, etc.): allows future extension of middleware return values, including returning values "up the chain"
• readonly arrays on InvokeModelContext: enforce immutable-context pattern at the type level

Known Limitations

• Middleware cannot resume interrupts. AgentStreamStage middleware can observe a tool-level interrupt result (stopReason: 'interrupt') but cannot re-enter the stream with interrupt responses. Interrupt resolution runs in stream()'s outer loop, outside middleware. To programmatically resume interrupts within a single invocation, use AfterInvocationEvent.resume in a hook. A future enhancement may add a resume mechanism to AgentStreamResult or AgentStreamContext.

Checklist

• I have read the CONTRIBUTING document
• Tests prove the fix is effective / feature works
• No new warnings
• Documentation update (pending)

---

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[github-actions]

API Bar Raising: This PR introduces a new public primitive (addMiddleware, three built-in stages, phase sub-tokens, createStage) that customers are expected to frequently use. Per API_BAR_RAISING.md, substantial changes like this should have explicit API reviewer sign-off.

The PR description is excellent and provides all the information an API reviewer needs (use cases, examples, signatures, exports). Please ensure the PR has the needs-api-review or completed-api-review label before merge.

One API design question worth discussing with the reviewer: the execution order semantics of phases. The current model is "Input → Around → Output" which is intuitive for the simple case, but the comment in the code ("Compose layering: input → output → around") reveals that Output handlers are layered outside Around handlers. This means a retry plugin on .Around retries Input transforms but not Output transforms. Is this the desired behavior in all cases? It might be worth documenting this explicitly for users.

[github-actions]

Assessment: Comment

Well-designed middleware system with thoughtful API ergonomics and comprehensive test coverage. The PR description is exemplary for API documentation.

Review Categories

• API Surface: MiddlewareInputHandler/MiddlewareOutputHandler types should be exported from src/index.ts for plugin authors who need explicit typing. MiddlewareRegistry intentionally internal — worth confirming.
• Code Clarity: The compose-layering vs execution-order distinction in comments could be clearer for future maintainers.
• Performance: compose() re-sorts on every call; minor concern on hot paths but acceptable for v1.
• Test DRY: collect() helper duplicated across two test files — should use shared fixture.
• Process: This is a substantial new public primitive — ensure API bar-raising review is completed before merge.

The async generator approach for middleware is elegant and well-suited for streaming use cases. The phase sub-token design (Input/Around/Output) is a nice solution for the middleware ordering problem.