@strands-agents/sdk 1.4.0 → 1.6.0

package/dist/src/agent/agent.js

@@ -1,26 +1,33 @@
 import { AgentResult, } from '../types/agent.js';
 import { BedrockModel } from '../models/bedrock.js';
 import { contentBlockFromData, Message, TextBlock, ToolResultBlock, ToolUseBlock, } from '../types/messages.js';
+import { deepCopy } from '../types/json.js';
 import { McpClient } from '../mcp.js';
 import { isValidToolName } from '../tools/tool.js';
-import { systemPromptFromData } from '../types/messages.js';
+import { cloneSystemPrompt, systemPromptFromData } from '../types/messages.js';
 import { normalizeError, ConcurrentInvocationError, StructuredOutputError } from '../errors.js';
 import { Model } from '../models/model.js';
 import { ModelPlugin } from '../plugins/model-plugin.js';
 import { isModelStreamEvent } from '../models/streaming.js';
 import { ToolRegistry } from '../registry/tool-registry.js';
 import { StateStore } from '../state-store.js';
+import { serializeStateSerializable, loadStateSerializable } from '../types/serializable.js';
 import { AgentPrinter, getDefaultAppender } from './printer.js';
 import { InterventionRegistry } from '../interventions/registry.js';
 import { PluginRegistry } from '../plugins/registry.js';
 import { SlidingWindowConversationManager } from '../conversation-manager/sliding-window-conversation-manager.js';
+import { SummarizingConversationManager } from '../conversation-manager/summarizing-conversation-manager.js';
 import { NullConversationManager } from '../conversation-manager/null-conversation-manager.js';
 import { ConversationManager } from '../conversation-manager/conversation-manager.js';
+import { ContextOffloader } from '../vended-plugins/context-offloader/plugin.js';
+import { InMemoryStorage } from '../vended-plugins/context-offloader/storage.js';
 import { HookRegistryImplementation } from '../hooks/registry.js';
+import { MiddlewareRegistry, InvokeModelStage, ExecuteToolStage, AgentStreamStage } from '../middleware/index.js';
 import { InitializedEvent, AfterInvocationEvent, AfterModelCallEvent, AfterToolCallEvent, AfterToolsEvent, BeforeInvocationEvent, BeforeModelCallEvent, BeforeToolCallEvent, BeforeToolsEvent, HookableEvent, MessageAddedEvent, ModelStreamUpdateEvent, ContentBlockEvent, ModelMessageEvent, ToolResultEvent, AgentResultEvent, ToolStreamUpdateEvent, InterruptEvent, } from '../hooks/events.js';
 import { StructuredOutputTool, STRUCTURED_OUTPUT_TOOL_NAME } from '../tools/structured-output-tool.js';
 import { AgentAsTool } from './agent-as-tool.js';
 import { ToolCaller } from './tool-caller.js';
+import { MemoryManager } from '../memory/memory-manager.js';
 import { SessionManager } from '../session/session-manager.js';
 import { Tracer } from '../telemetry/tracer.js';
 import { Meter } from '../telemetry/meter.js';
@@ -28,13 +35,83 @@ import { logger } from '../logging/logger.js';
 import { CancelledError } from '../errors.js';
 import { DefaultModelRetryStrategy } from '../retry/default-model-retry-strategy.js';
 import { warnOnDuplicateRetryStrategyTypes } from '../retry/retry-strategy.js';
-import { InterruptError, InterruptState, interruptFromAgent } from '../interrupt.js';
+import { Interrupt, InterruptError, InterruptState, interruptFromAgent } from '../interrupt.js';
 import { isInterruptResponseContent } from '../types/interrupt.js';
 import { takeSnapshot as takeSnapshotInternal, loadSnapshot as loadSnapshotInternal } from './snapshot.js';
+import { defaultSandbox } from '../sandbox/default.js';
+import { summarizeContextTool, truncateContextTool, pinContextTool, createTokenUsageMiddleware, } from '../context-manager/modes/agentic/agentic-context.js';
+/**
+ * Supported values for the `contextManager` parameter.
+ */
+export const CONTEXT_MANAGER_STRATEGIES = ['auto', 'agentic'];
+/** Benchmark-validated token threshold for offloading tool results. */
+const CONTEXT_MANAGER_MAX_RESULT_TOKENS = 1_500;
+/** Higher offload threshold for agentic mode — the model manages its own context, so we preserve more inline. */
+const AGENTIC_CONTEXT_MANAGER_MAX_RESULT_TOKENS = 8_000;
+/** Benchmark-validated preview token count for offloaded results. */
+const CONTEXT_MANAGER_PREVIEW_TOKENS = 750;
+/** Benchmark-validated ratio of messages to summarize on overflow. */
+const CONTEXT_MANAGER_SUMMARY_RATIO = 0.3;
+/** Benchmark-validated context window ratio that triggers proactive compression. */
+const CONTEXT_MANAGER_COMPRESSION_THRESHOLD = 0.85;
+/**
+ * Resolve the contextManager facade into a concrete ConversationManager.
+ *
+ * When contextManager is undefined, falls back to the default SlidingWindowConversationManager.
+ * When "auto", uses SummarizingConversationManager with proactive compression.
+ * When "agentic", uses SummarizingConversationManager without proactive compression
+ * (the agent manages its context via tools; the context manager is only a reactive safety net).
+ */
+function resolveConversationManager(contextManager, conversationManager) {
+    if (contextManager === 'agentic') {
+        return (conversationManager ??
+            new SummarizingConversationManager({
+                summaryRatio: CONTEXT_MANAGER_SUMMARY_RATIO,
+            }));
+    }
+    if (contextManager === 'auto') {
+        return (conversationManager ??
+            new SummarizingConversationManager({
+                summaryRatio: CONTEXT_MANAGER_SUMMARY_RATIO,
+                proactiveCompression: { compressionThreshold: CONTEXT_MANAGER_COMPRESSION_THRESHOLD },
+            }));
+    }
+    if (contextManager !== undefined) {
+        throw new Error(`Unsupported contextManager value: "${contextManager}". Supported values: ${CONTEXT_MANAGER_STRATEGIES.map((s) => `"${s}"`).join(', ')}`);
+    }
+    return conversationManager ?? new SlidingWindowConversationManager({ windowSize: 40 });
+}
 /** Default name assigned to agents when none is provided. */
 const DEFAULT_AGENT_NAME = 'Strands Agent';
 /** Default identifier assigned to agents when none is provided. */
 const DEFAULT_AGENT_ID = 'agent';
+/**
+ * Creates a non-mutating interrupt function for middleware contexts.
+ * Reads existing responses from state (resume case) but never writes to it.
+ * On first run (no response), throws InterruptError with a locally-created Interrupt.
+ *
+ * @param interruptState - The agent's interrupt state (read-only access)
+ * @param idPrefix - Prefix for the interrupt ID (e.g., 'middleware:agentStream')
+ */
+function createMiddlewareInterrupt(interruptState, idPrefix) {
+    return (params) => {
+        const interruptId = `${idPrefix}:${params.name}`;
+        const existing = interruptState.interrupts[interruptId];
+        if (existing?.response !== undefined) {
+            return { response: existing.response };
+        }
+        if (params.response !== undefined) {
+            return { response: params.response };
+        }
+        const interrupt = new Interrupt({
+            id: interruptId,
+            name: params.name,
+            ...(params.reason !== undefined && { reason: params.reason }),
+            source: 'middleware',
+        });
+        throw new InterruptError(interrupt);
+    };
+}
 /**
  * Orchestrates the interaction between a model, a set of tools, and MCP clients.
  * The Agent is responsible for managing the lifecycle of tools and clients
@@ -81,7 +158,22 @@ export class Agent {
      * The session manager for saving and restoring agent sessions, if configured.
      */
     sessionManager;
+    /**
+     * The memory manager for cross-session memory retrieval and storage, if configured.
+     */
+    memoryManager;
+    _sandbox;
+    /**
+     * Execution environment for running commands, code, and file operations.
+     *
+     * @throws DefaultNotConfiguredError if no sandbox is configured for this
+     * environment (e.g. browsers, where no host default is registered).
+     */
+    get sandbox() {
+        return this._sandbox || defaultSandbox.get();
+    }
     _hooksRegistry;
+    _middlewareRegistry;
     _pluginRegistry;
     _interventionRegistry;
     _toolRegistry;
@@ -116,6 +208,13 @@ export class Agent {
         if (config?.description !== undefined)
             this.description = config.description;
         this.sessionManager = config?.sessionManager;
+        this.memoryManager =
+            config?.memoryManager instanceof MemoryManager
+                ? config.memoryManager
+                : config?.memoryManager
+                    ? new MemoryManager(config.memoryManager)
+                    : undefined;
+        this._sandbox = config?.sandbox;
         if (typeof config?.model === 'string') {
             this.model = new BedrockModel({ modelId: config.model });
         }
@@ -124,21 +223,28 @@ export class Agent {
         }
         // Validate and assign conversation manager
         if (this.model.stateful) {
-            if (config?.conversationManager) {
-                throw new Error('Cannot use a conversationManager with a stateful model. The model manages conversation state server-side.');
+            if (config?.conversationManager || config?.contextManager) {
+                throw new Error('contextManager and conversationManager cannot be used with a stateful model. The model manages conversation state server-side.');
             }
             this._conversationManager = new NullConversationManager();
         }
         else {
-            this._conversationManager =
-                config?.conversationManager ?? new SlidingWindowConversationManager({ windowSize: 40 });
+            this._conversationManager = resolveConversationManager(config?.contextManager, config?.conversationManager);
         }
         const { tools, mcpClients } = flattenTools(config?.tools ?? []);
+        if (config?.contextManager === 'agentic') {
+            tools.push(summarizeContextTool, truncateContextTool, pinContextTool);
+        }
         this._toolRegistry = new ToolRegistry(tools);
         this._mcpClients = mcpClients;
         // Initialize hooks registry
         this._hooksRegistry = new HookRegistryImplementation();
         this._interventionRegistry = new InterventionRegistry(config?.interventions ?? [], this._hooksRegistry);
+        // Initialize middleware registry
+        this._middlewareRegistry = new MiddlewareRegistry();
+        if (config?.contextManager === 'agentic') {
+            this._middlewareRegistry.addInput(InvokeModelStage.Input, createTokenUsageMiddleware(this.model));
+        }
         // `undefined` (omitted) → install the default; `null`/`[]` → explicit opt-out.
         const retryStrategies = config?.retryStrategy === null
             ? []
@@ -156,10 +262,23 @@ export class Agent {
         // - Retry-strategy ordering is not load-bearing for correctness: `DefaultModelRetryStrategy`
         //   guards on `event.retry`, so a user hook that already set it short-circuits
         //   the strategy regardless of registration order.
+        const hasOffloader = (config?.plugins ?? []).some((p) => p.name === 'strands:context-offloader');
         this._pluginRegistry = new PluginRegistry([
             this._conversationManager,
             ...retryStrategies,
             ...(config?.plugins ?? []),
+            ...((config?.contextManager === 'auto' || config?.contextManager === 'agentic') && !hasOffloader
+                ? [
+                    new ContextOffloader({
+                        storage: new InMemoryStorage(),
+                        maxResultTokens: config?.contextManager === 'agentic'
+                            ? AGENTIC_CONTEXT_MANAGER_MAX_RESULT_TOKENS
+                            : CONTEXT_MANAGER_MAX_RESULT_TOKENS,
+                        previewTokens: CONTEXT_MANAGER_PREVIEW_TOKENS,
+                    }),
+                ]
+                : []),
+            ...(this.memoryManager ? [this.memoryManager] : []),
             ...(config?.sessionManager ? [config.sessionManager] : []),
             new ModelPlugin(this.model),
         ]);
@@ -208,6 +327,33 @@ export class Agent {
     addHook(eventType, callback, options) {
         return this._hooksRegistry.addCallback(eventType, callback, options);
     }
+    addMiddleware(stageOrPhase, handler) {
+        if ('_phase' in stageOrPhase) {
+            const phase = stageOrPhase;
+            const stage = phase._stage;
+            switch (phase._phase) {
+                case 'input': {
+                    const adapted = this._middlewareRegistry.addInput(stageOrPhase, handler);
+                    return () => this._middlewareRegistry.remove(stage, adapted);
+                }
+                case 'output': {
+                    const adapted = this._middlewareRegistry.addOutput(stageOrPhase, handler);
+                    return () => this._middlewareRegistry.remove(stage, adapted);
+                }
+                case 'wrap': {
+                    const wrapHandler = handler;
+                    this._middlewareRegistry.add(stage, wrapHandler);
+                    return () => this._middlewareRegistry.remove(stage, wrapHandler);
+                }
+                default:
+                    throw new Error(`Unknown middleware phase: ${phase._phase}`);
+            }
+        }
+        const stage = stageOrPhase;
+        const wrapHandler = handler;
+        this._middlewareRegistry.add(stage, wrapHandler);
+        return () => this._middlewareRegistry.remove(stage, wrapHandler);
+    }
     async initialize() {
         if (this._initialized) {
             return;
@@ -221,6 +367,18 @@ export class Agent {
                 this._toolRegistry.addOrReplace(newTools);
             };
         }));
+        // Register tools vended by the sandbox. The host default vends nothing. A tool
+        // is skipped if the user already registered one with that name.
+        if (this._sandbox) {
+            for (const sandboxTool of this._sandbox.getTools()) {
+                if (this._toolRegistry.get(sandboxTool.name)) {
+                    logger.debug(`tool_name=<${sandboxTool.name}> | sandbox-vended tool skipped, user already registered a tool with this name`);
+                }
+                else {
+                    this._toolRegistry.add(sandboxTool);
+                }
+            }
+        }
         await this._pluginRegistry.initialize(this);
         for (const handler of this._interventionRegistry.handlers) {
             const observer = handler;
@@ -454,90 +612,183 @@ export class Agent {
         this.acquireLock();
         try {
             await this.initialize();
+            // Thread the resolved invocationState so all layers share the same reference.
+            const invocationState = options?.invocationState ?? {};
+            const resolvedOptions = options?.invocationState ? options : { ...options, invocationState };
             let currentArgs = args;
-            // Outer loop: re-enters _stream when a hook sets AfterInvocationEvent.resume.
-            // One invocation lock spans the whole resume chain.
             while (true) {
-                // Fresh AbortController per invocation iteration, composed with any external signal.
+                // Fresh AbortController per iteration, composed with any external signal.
                 this._abortController = new AbortController();
-                this._abortSignal = options?.cancelSignal
-                    ? AbortSignal.any([this._abortController.signal, options.cancelSignal])
+                this._abortSignal = resolvedOptions?.cancelSignal
+                    ? AbortSignal.any([this._abortController.signal, resolvedOptions.cancelSignal])
                     : this._abortController.signal;
-                const streamGenerator = this._stream(currentArgs, options);
+                // Process interrupt responses before middleware runs so context.interrupt() can find them
+                const interruptResponses = this._extractInterruptResponses(currentArgs);
+                if (interruptResponses.length > 0) {
+                    this._interruptState.resume(interruptResponses);
+                }
+                // Hooks fire outside middleware — always, even on short-circuit.
+                const beforeInvocationEvent = new BeforeInvocationEvent({ agent: this, invocationState });
+                yield await this._invokeCallbacks(beforeInvocationEvent);
+                if (beforeInvocationEvent.cancel) {
+                    const cancelText = typeof beforeInvocationEvent.cancel === 'string'
+                        ? beforeInvocationEvent.cancel
+                        : 'invocation denied by hook';
+                    const message = new Message({ role: 'assistant', content: [new TextBlock(cancelText)] });
+                    yield this._appendMessage(message, invocationState);
+                    const afterEvent = new AfterInvocationEvent({ agent: this, invocationState });
+                    await this._invokeCallbacks(afterEvent);
+                    yield afterEvent;
+                    return new AgentResult({
+                        stopReason: 'endTurn',
+                        lastMessage: message,
+                        traces: this._tracer.localTraces,
+                        metrics: this._meter.metrics,
+                        invocationState,
+                    });
+                }
+                let result;
                 let caughtError;
-                let lastAfterInvocation;
-                let iterationResult;
+                const afterInvocationEvent = new AfterInvocationEvent({ agent: this, invocationState });
                 try {
-                    iterationResult = await streamGenerator.next();
-                    while (!iterationResult.done) {
-                        try {
-                            const processed = await this._invokeCallbacks(iterationResult.value);
-                            if (processed instanceof AfterInvocationEvent) {
-                                lastAfterInvocation = processed;
-                            }
-                            yield processed;
-                            iterationResult = await streamGenerator.next();
-                        }
-                        catch (error) {
-                            // Throw interrupt errors back into _stream so executeTools can store the
-                            // assistant message as pending execution state for resume.
-                            if (error instanceof InterruptError) {
-                                iterationResult = await streamGenerator.throw(error);
-                            }
-                            else {
-                                throw error;
-                            }
-                        }
-                    }
-                    // Suppress AgentResultEvent for resumed iterations — only the final
-                    // invocation in a resume chain reports an agent result.
-                    if (lastAfterInvocation?.resume === undefined) {
-                        yield await this._invokeCallbacks(new AgentResultEvent({
-                            agent: this,
-                            result: iterationResult.value,
-                            invocationState: iterationResult.value.invocationState,
-                        }));
-                    }
+                    result = yield* this._streamWithMiddleware(currentArgs, resolvedOptions, invocationState);
                 }
                 catch (error) {
                     caughtError = error;
-                    throw error;
                 }
                 finally {
-                    // Drain _stream() so cleanup hooks and printer still fire.
-                    // Yield only on error (consumer may still be iterating); on a consumer
-                    // break, yielding would suspend the generator and leak the lock.
-                    let drainResult = await streamGenerator.return(undefined);
-                    while (!drainResult.done) {
-                        try {
-                            if (caughtError) {
-                                yield await this._invokeCallbacks(drainResult.value);
-                            }
-                            else {
-                                await this._invokeCallbacks(drainResult.value);
-                            }
-                        }
-                        catch (error) {
-                            logger.warn(`event_type=<${drainResult.value.type}>, error=<${error}> | error invoking callbacks during cleanup`);
-                        }
-                        drainResult = await streamGenerator.next();
-                    }
-                    // Reset controller and signal for next iteration / invocation
-                    this._abortController = new AbortController();
-                    this._abortSignal = this._abortController.signal;
+                    // AfterInvocationEvent always fires — even on error or consumer break. Outside middleware.
+                    // Invoke hooks (so .resume can be set) but don't yield in finally (yields in finally
+                    // suspend the generator on consumer break instead of completing cleanup).
+                    await this._invokeCallbacks(afterInvocationEvent);
+                }
+                // Yield outside finally — in JS, a `yield` inside `finally` suspends the generator
+                // mid-cleanup when the consumer breaks, preventing subsequent cleanup code from running.
+                // This line is only reached on normal completion or caught error, never on consumer break.
+                yield afterInvocationEvent;
+                // Re-throw after hooks have fired
+                if (caughtError) {
+                    throw caughtError;
                 }
                 // Resume only on a clean invocation — errors propagate above.
-                if (lastAfterInvocation?.resume !== undefined) {
-                    currentArgs = lastAfterInvocation.resume;
+                if (afterInvocationEvent.resume !== undefined) {
+                    currentArgs = afterInvocationEvent.resume;
                     continue;
                 }
-                return iterationResult.value;
+                // Only emit AgentResultEvent on the final iteration (not on resumed ones).
+                yield await this._invokeCallbacks(new AgentResultEvent({
+                    agent: this,
+                    result: result,
+                    invocationState,
+                }));
+                return result;
             }
         }
         finally {
             this._isInvoking = false;
         }
     }
+    /**
+     * Invokes the AgentStreamStage middleware chain.
+     * Hooks fire outside this method (in stream()'s resume loop).
+     */
+    async *_streamWithMiddleware(args, options, invocationState) {
+        const context = {
+            agent: this,
+            args,
+            ...(options !== undefined && { options }),
+            interrupt: createMiddlewareInterrupt(this._interruptState, 'middleware:agentStream'),
+        };
+        // async function* doesn't bind lexical `this`; capture for the terminal callback.
+        // eslint-disable-next-line @typescript-eslint/no-this-alias
+        const self = this;
+        try {
+            const { result } = yield* this._middlewareRegistry.invoke(AgentStreamStage, context, async function* (ctx) {
+                const result = yield* self._streamCore(ctx.args, ctx.options);
+                return { result };
+            });
+            return result;
+        }
+        catch (error) {
+            if (error instanceof InterruptError) {
+                for (const interrupt of error.interrupts) {
+                    this._interruptState.registerInterrupt(interrupt);
+                }
+                this._interruptState.activate();
+                for (const interrupt of error.interrupts) {
+                    yield new InterruptEvent({ agent: this, interrupt, invocationState });
+                }
+                return new AgentResult({
+                    stopReason: 'interrupt',
+                    lastMessage: this.messages.length > 0
+                        ? this.messages[this.messages.length - 1]
+                        : new Message({ role: 'assistant', content: [new TextBlock('Interrupted')] }),
+                    traces: this._tracer.localTraces,
+                    metrics: this._meter.metrics,
+                    interrupts: this._interruptState.getUnansweredInterrupts(),
+                    invocationState,
+                });
+            }
+            throw error;
+        }
+    }
+    /**
+     * Single-pass stream through _stream() with event processing.
+     * No resume loop, no lifecycle events — those are handled by stream()'s resume loop.
+     */
+    async *_streamCore(args, options) {
+        const streamGenerator = this._stream(args, options);
+        let caughtError;
+        let iterationResult;
+        try {
+            iterationResult = await streamGenerator.next();
+            while (!iterationResult.done) {
+                try {
+                    const processed = await this._invokeCallbacks(iterationResult.value);
+                    yield processed;
+                    iterationResult = await streamGenerator.next();
+                }
+                catch (error) {
+                    // Throw interrupt errors back into _stream so executeTools can store the
+                    // assistant message as pending execution state for resume.
+                    if (error instanceof InterruptError) {
+                        iterationResult = await streamGenerator.throw(error);
+                    }
+                    else {
+                        throw error;
+                    }
+                }
+            }
+        }
+        catch (error) {
+            caughtError = error;
+            throw error;
+        }
+        finally {
+            // Drain _stream() so cleanup hooks and printer still fire.
+            // Yield only on error (consumer may still be iterating); on a consumer
+            // break, yielding would suspend the generator and leak the lock.
+            let drainResult = await streamGenerator.return(undefined);
+            while (!drainResult.done) {
+                try {
+                    if (caughtError) {
+                        yield await this._invokeCallbacks(drainResult.value);
+                    }
+                    else {
+                        await this._invokeCallbacks(drainResult.value);
+                    }
+                }
+                catch (error) {
+                    logger.warn(`event_type=<${drainResult.value.type}>, error=<${error}> | error invoking callbacks during cleanup`);
+                }
+                drainResult = await streamGenerator.next();
+            }
+            // Reset controller and signal for next iteration / invocation
+            this._abortController = new AbortController();
+            this._abortSignal = this._abortController.signal;
+        }
+        return iterationResult.value;
+    }
     /**
      * Returns a {@link Tool} that wraps this agent, allowing it to be used
      * as a tool by another agent.
@@ -661,30 +912,14 @@ export class Agent {
         // AgentResult. Mutations by hooks/tools are visible across all recursive
         // agent loop cycles within this invocation.
         const invocationState = options?.invocationState ?? {};
-        // Handle interrupt responses if present in input
+        // Interrupt responses are already consumed in stream() before middleware
+        // runs (so middleware-level interrupt() can find them). Re-extract here
+        // to gate the "non-interrupt input while interrupted" check below.
         const interruptResponses = this._extractInterruptResponses(args);
-        if (interruptResponses.length > 0) {
-            this._interruptState.resume(interruptResponses);
-        }
         // Reject non-interrupt input while in interrupted state
         if (this._interruptState.activated && interruptResponses.length === 0) {
             throw new TypeError('Agent is in an interrupted state. Resume by invoking with interruptResponse content blocks.');
         }
-        const beforeInvocationEvent = new BeforeInvocationEvent({ agent: this, invocationState });
-        yield beforeInvocationEvent;
-        if (beforeInvocationEvent.cancel) {
-            const cancelText = typeof beforeInvocationEvent.cancel === 'string' ? beforeInvocationEvent.cancel : 'invocation denied by hook';
-            const message = new Message({ role: 'assistant', content: [new TextBlock(cancelText)] });
-            yield this._appendMessage(message, invocationState);
-            yield new AfterInvocationEvent({ agent: this, invocationState });
-            return new AgentResult({
-                stopReason: 'endTurn',
-                lastMessage: message,
-                traces: this._tracer.localTraces,
-                metrics: this._meter.metrics,
-                invocationState,
-            });
-        }
         // Normalize input to get the user messages for telemetry
         const inputMessages = this._normalizeInput(args);
         // Start agent trace span
@@ -884,6 +1119,11 @@ export class Agent {
                 return result;
             }
             if (error instanceof InterruptError) {
+                // Handles interrupts from tools/hooks that propagated up through the agent loop.
+                // AgentStreamStage middleware interrupts are caught separately in _streamWithMiddleware().
+                for (const interrupt of error.interrupts) {
+                    this._interruptState.registerInterrupt(interrupt);
+                }
                 // Fan out one event per interrupt. Each event exposes `interrupt.source` so
                 // consumers can filter by origin (tool callback vs hook callback) without
                 // subscribing to separate event types.
@@ -917,8 +1157,6 @@ export class Agent {
             if (structuredOutputTool) {
                 this._toolRegistry.remove(STRUCTURED_OUTPUT_TOOL_NAME);
             }
-            // Always emit final event
-            yield new AfterInvocationEvent({ agent: this, invocationState });
         }
     }
     /**
@@ -1095,26 +1333,10 @@ export class Agent {
                 }
                 return { message, stopReason: 'endTurn' };
             }
-            // Start model span within loop span context
-            const modelId = this.model.modelId;
-            const modelSpan = this._tracer.startModelInvokeSpan({
-                messages: this.messages,
-                ...(modelId && { modelId }),
-                ...(this.systemPrompt !== undefined && { systemPrompt: this.systemPrompt }),
-            });
             try {
-                const result = yield* this._streamFromModel(this.messages, streamOptions, invocationState);
+                const result = yield* this._invokeModelWithMiddleware(invocationState, toolChoice, projectedInputTokens);
                 // Accumulate token usage and model latency metrics
                 this._meter.updateCycle(result.metadata);
-                // End model span with usage
-                const usage = result.metadata?.usage;
-                const metrics = result.metadata?.metrics;
-                this._tracer.endModelInvokeSpan(modelSpan, {
-                    output: result.message,
-                    stopReason: result.stopReason,
-                    ...(usage && { usage }),
-                    ...(metrics && { metrics }),
-                });
                 yield new ModelMessageEvent({
                     agent: this,
                     message: result.message,
@@ -1146,8 +1368,6 @@ export class Agent {
             }
             catch (error) {
                 const modelError = normalizeError(error);
-                // End model span with error
-                this._tracer.endModelInvokeSpan(modelSpan, { error: modelError });
                 // Create error event
                 const errorEvent = new AfterModelCallEvent({
                     agent: this,
@@ -1173,6 +1393,80 @@ export class Agent {
             }
         }
     }
+    /**
+     * Invokes the model through the InvokeModelStage middleware chain.
+     * Builds an InvokeModelContext from current agent state and composes the
+     * middleware chain with a terminal function that calls _streamFromModel
+     * using context fields directly (not re-derived from the agent).
+     *
+     * @param invocationState - Per-invocation state shared across hooks and tools
+     * @param toolChoice - Optional tool choice to force specific tool usage
+     * @returns StreamAggregatedResult from the model (or middleware short-circuit)
+     */
+    async *_invokeModelWithMiddleware(invocationState, toolChoice, projectedInputTokens) {
+        const context = {
+            agent: this,
+            messages: this.messages.map((msg) => msg.clone()),
+            ...(this.systemPrompt !== undefined && { systemPrompt: cloneSystemPrompt(this.systemPrompt) }),
+            toolSpecs: deepCopy(this._toolRegistry.list().map((tool) => tool.toolSpec)),
+            ...(toolChoice !== undefined && { toolChoice: deepCopy(toolChoice) }),
+            invocationState,
+            ...(projectedInputTokens !== undefined && { projectedInputTokens }),
+        };
+        // Snapshot model state before middleware runs so concurrent mutations don't leak in.
+        // The writeback happens after the entire middleware chain completes, so middleware
+        // cannot affect modelState at any point (before or after next()).
+        const modelStateSnapshot = this.modelState.getAll();
+        let tempModelState;
+        // async function* doesn't bind lexical `this`; capture for the terminal callback.
+        // eslint-disable-next-line @typescript-eslint/no-this-alias
+        const self = this;
+        const middlewareResult = yield* this._middlewareRegistry.invoke(InvokeModelStage, context, async function* (ctx) {
+            const modelId = self.model.modelId;
+            const modelSpan = self._tracer.startModelInvokeSpan({
+                messages: ctx.messages,
+                ...(modelId && { modelId }),
+                ...(ctx.systemPrompt !== undefined && { systemPrompt: ctx.systemPrompt }),
+            });
+            try {
+                // Wrap the snapshot into a StateStore for the model provider, which expects
+                // get/set methods.
+                tempModelState = new StateStore(modelStateSnapshot);
+                const streamOptions = {
+                    toolSpecs: ctx.toolSpecs,
+                    modelState: tempModelState,
+                    ...(ctx.systemPrompt !== undefined && { systemPrompt: ctx.systemPrompt }),
+                    ...(ctx.toolChoice && { toolChoice: ctx.toolChoice }),
+                };
+                const gen = self._streamFromModel(ctx.messages, streamOptions, ctx.invocationState);
+                let iterResult = await gen.next();
+                while (!iterResult.done) {
+                    yield iterResult.value;
+                    iterResult = await gen.next();
+                }
+                const usage = iterResult.value.metadata?.usage;
+                const metrics = iterResult.value.metadata?.metrics;
+                self._tracer.endModelInvokeSpan(modelSpan, {
+                    output: iterResult.value.message,
+                    stopReason: iterResult.value.stopReason,
+                    ...(usage && { usage }),
+                    ...(metrics && { metrics }),
+                });
+                return { result: iterResult.value };
+            }
+            catch (error) {
+                self._tracer.endModelInvokeSpan(modelSpan, { error: normalizeError(error) });
+                throw error;
+            }
+        });
+        // Sync model state after the entire middleware chain has completed, so no
+        // middleware mutation to agent.modelState (before or after next()) takes effect.
+        // Intentionally skipped on error — partial provider writes should not persist.
+        if (tempModelState) {
+            loadStateSerializable(this.modelState, serializeStateSerializable(tempModelState));
+        }
+        return middlewareResult.result;
+    }
     /**
      * Streams events from the model and dispatches appropriate events for each.
      *
@@ -1529,86 +1823,9 @@ export class Agent {
                 }
                 return afterToolCallEvent.result;
             }
-            // Start tool span within loop span context
-            const toolSpan = this._tracer.startToolCallSpan({
-                tool: toolUse,
-            });
-            // Track tool execution time for metrics
-            const toolStartTime = Date.now();
-            let toolResult;
-            let error;
-            if (!effectiveTool) {
-                // Tool not found
-                toolResult = new ToolResultBlock({
-                    toolUseId: toolUse.toolUseId,
-                    status: 'error',
-                    content: [new TextBlock(`Tool '${toolUse.name}' not found in registry`)],
-                });
-            }
-            else {
-                // Execute tool within the tool span context
-                const toolContext = {
-                    toolUse: {
-                        name: toolUse.name,
-                        toolUseId: toolUse.toolUseId,
-                        input: toolUse.input,
-                    },
-                    agent: this,
-                    invocationState,
-                    interrupt: (params) => {
-                        return interruptFromAgent(this, `tool:${toolUseBlock.toolUseId}:${params.name}`, params, 'tool');
-                    },
-                };
-                try {
-                    // Manually iterate tool stream to wrap each ToolStreamEvent in ToolStreamUpdateEvent.
-                    // This keeps the tool authoring interface unchanged — tools construct ToolStreamEvent
-                    // without knowledge of agents or hooks, and we wrap at the boundary.
-                    // Tool execution is ran within the tool span's context so that
-                    // downstream calls (e.g., MCP clients) can propagate trace context
-                    const toolGenerator = this._tracer.withSpanContext(toolSpan, () => effectiveTool.stream(toolContext));
-                    let toolNext = await this._tracer.withSpanContext(toolSpan, () => toolGenerator.next());
-                    while (!toolNext.done) {
-                        yield new ToolStreamUpdateEvent({ agent: this, event: toolNext.value, invocationState });
-                        toolNext = await this._tracer.withSpanContext(toolSpan, () => toolGenerator.next());
-                    }
-                    const result = toolNext.value;
-                    if (!result) {
-                        // Tool didn't return a result
-                        toolResult = new ToolResultBlock({
-                            toolUseId: toolUse.toolUseId,
-                            status: 'error',
-                            content: [new TextBlock(`Tool '${toolUse.name}' did not return a result`)],
-                        });
-                    }
-                    else {
-                        toolResult = result;
-                        error = result.error;
-                    }
-                }
-                catch (e) {
-                    // Re-throw InterruptError to allow interrupt handling
-                    if (e instanceof InterruptError) {
-                        throw e;
-                    }
-                    // Tool execution failed with error
-                    error = normalizeError(e);
-                    toolResult = new ToolResultBlock({
-                        toolUseId: toolUse.toolUseId,
-                        status: 'error',
-                        content: [new TextBlock(error.message)],
-                        error,
-                    });
-                }
-            }
-            // End tool span with the raw tool result — telemetry reflects what the
-            // tool actually returned, independent of AfterToolCallEvent mutations.
-            this._tracer.endToolCallSpan(toolSpan, { toolResult, ...(error && { error }) });
-            // End tool metrics tracking
-            this._meter.endToolCall({
-                tool: toolUse,
-                duration: Date.now() - toolStartTime,
-                success: toolResult.status === 'success',
-            });
+            // Execute tool core logic through middleware chain
+            const toolResult = yield* this._executeToolWithMiddleware(effectiveTool, toolUse, invocationState);
+            const error = toolResult.error;
             // Single point for AfterToolCallEvent
             const afterToolCallEvent = new AfterToolCallEvent({
                 agent: this,
@@ -1627,6 +1844,106 @@ export class Agent {
             return afterToolCallEvent.result;
         }
     }
+    async *_executeToolWithMiddleware(tool, toolUse, invocationState) {
+        const context = {
+            agent: this,
+            tool,
+            toolUse: deepCopy(toolUse),
+            invocationState,
+            interrupt: createMiddlewareInterrupt(this._interruptState, `middleware:executeTool:${toolUse.toolUseId}`),
+        };
+        // async function* doesn't bind lexical `this`; capture for the terminal callback.
+        // eslint-disable-next-line @typescript-eslint/no-this-alias
+        const self = this;
+        const middlewareResult = yield* this._middlewareRegistry.invoke(ExecuteToolStage, context, async function* (ctx) {
+            const result = yield* self._executeToolCore(ctx.tool, ctx.toolUse, ctx.invocationState);
+            return { result };
+        });
+        return middlewareResult.result;
+    }
+    async *_executeToolCore(effectiveTool, toolUse, invocationState) {
+        // Start tool span within loop span context
+        const toolSpan = this._tracer.startToolCallSpan({
+            tool: toolUse,
+        });
+        // Track tool execution time for metrics
+        const toolStartTime = Date.now();
+        let toolResult;
+        let error;
+        if (!effectiveTool) {
+            // Tool not found
+            toolResult = new ToolResultBlock({
+                toolUseId: toolUse.toolUseId,
+                status: 'error',
+                content: [new TextBlock(`Tool '${toolUse.name}' not found in registry`)],
+            });
+        }
+        else {
+            // Execute tool within the tool span context
+            const toolContext = {
+                toolUse: {
+                    name: toolUse.name,
+                    toolUseId: toolUse.toolUseId,
+                    input: toolUse.input,
+                },
+                agent: this,
+                invocationState,
+                interrupt: (params) => {
+                    return interruptFromAgent(this, `tool:${toolUse.toolUseId}:${params.name}`, params, 'tool');
+                },
+            };
+            try {
+                // Manually iterate tool stream to wrap each ToolStreamEvent in ToolStreamUpdateEvent.
+                // This keeps the tool authoring interface unchanged — tools construct ToolStreamEvent
+                // without knowledge of agents or hooks, and we wrap at the boundary.
+                // Tool execution is ran within the tool span's context so that
+                // downstream calls (e.g., MCP clients) can propagate trace context
+                const toolGenerator = this._tracer.withSpanContext(toolSpan, () => effectiveTool.stream(toolContext));
+                let toolNext = await this._tracer.withSpanContext(toolSpan, () => toolGenerator.next());
+                while (!toolNext.done) {
+                    yield new ToolStreamUpdateEvent({ agent: this, event: toolNext.value, invocationState });
+                    toolNext = await this._tracer.withSpanContext(toolSpan, () => toolGenerator.next());
+                }
+                const result = toolNext.value;
+                if (!result) {
+                    // Tool didn't return a result
+                    toolResult = new ToolResultBlock({
+                        toolUseId: toolUse.toolUseId,
+                        status: 'error',
+                        content: [new TextBlock(`Tool '${toolUse.name}' did not return a result`)],
+                    });
+                }
+                else {
+                    toolResult = result;
+                    error = result.error;
+                }
+            }
+            catch (e) {
+                // Re-throw InterruptError to allow interrupt handling
+                if (e instanceof InterruptError) {
+                    throw e;
+                }
+                // Tool execution failed with error
+                error = normalizeError(e);
+                toolResult = new ToolResultBlock({
+                    toolUseId: toolUse.toolUseId,
+                    status: 'error',
+                    content: [new TextBlock(error.message)],
+                    error,
+                });
+            }
+        }
+        // End tool span with the raw tool result — telemetry reflects what the
+        // tool actually returned, independent of AfterToolCallEvent mutations.
+        this._tracer.endToolCallSpan(toolSpan, { toolResult, ...(error && { error }) });
+        // End tool metrics tracking
+        this._meter.endToolCall({
+            tool: toolUse,
+            duration: Date.now() - toolStartTime,
+            success: toolResult.status === 'success',
+        });
+        return toolResult;
+    }
     /**
      * Redacts the last message in the conversation history.
      * Called when guardrails block user input and redaction is enabled.