@strands-agents/sdk 1.0.0-rc.4 → 1.0.0

package/dist/src/agent/agent.js

@@ -58,12 +58,14 @@ import {} from '../tools/tool.js';
 import { 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 { AgentPrinter, getDefaultAppender } from './printer.js';
 import { PluginRegistry } from '../plugins/registry.js';
 import { SlidingWindowConversationManager } from '../conversation-manager/sliding-window-conversation-manager.js';
+import { NullConversationManager } from '../conversation-manager/null-conversation-manager.js';
 import { ConversationManager } from '../conversation-manager/conversation-manager.js';
 import { HookRegistryImplementation } from '../hooks/registry.js';
 import { InitializedEvent, AfterInvocationEvent, AfterModelCallEvent, AfterToolCallEvent, AfterToolsEvent, BeforeInvocationEvent, BeforeModelCallEvent, BeforeToolCallEvent, BeforeToolsEvent, HookableEvent, MessageAddedEvent, ModelStreamUpdateEvent, ContentBlockEvent, ModelMessageEvent, ToolResultEvent, AgentResultEvent, ToolStreamUpdateEvent, } from '../hooks/events.js';
@@ -93,6 +95,12 @@ export class Agent {
      * State is not passed to the model during inference.
      */
     appState;
+    /**
+     * Runtime state for the model provider. Used by stateful models to persist
+     * provider-specific data (e.g., response IDs for conversation chaining)
+     * across invocations.
+     */
+    modelState;
     _conversationManager;
     /**
      * The model provider used by the agent for inference.
@@ -132,6 +140,8 @@ export class Agent {
     _tracer;
     /** Meter instance for accumulating loop metrics during invocation. */
     _meter;
+    /** Strategy for executing tool calls from a single assistant turn. */
+    _toolExecutor;
     /**
      * Creates an instance of the Agent.
      * @param config - The configuration for the agent.
@@ -140,7 +150,7 @@ export class Agent {
         // Initialize public fields
         this.messages = (config?.messages ?? []).map((msg) => (msg instanceof Message ? msg : Message.fromMessageData(msg)));
         this.appState = new StateStore(config?.appState);
-        this._conversationManager = config?.conversationManager ?? new SlidingWindowConversationManager({ windowSize: 40 });
+        this.modelState = new StateStore(config?.modelState);
         this.name = config?.name ?? DEFAULT_AGENT_NAME;
         this.id = config?.id ?? DEFAULT_AGENT_ID;
         if (config?.description !== undefined)
@@ -152,16 +162,30 @@ export class Agent {
         else {
             this.model = config?.model ?? new BedrockModel();
         }
+        // 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.');
+            }
+            this._conversationManager = new NullConversationManager();
+        }
+        else {
+            this._conversationManager =
+                config?.conversationManager ?? new SlidingWindowConversationManager({ windowSize: 40 });
+        }
         const { tools, mcpClients } = flattenTools(config?.tools ?? []);
         this._toolRegistry = new ToolRegistry(tools);
         this._mcpClients = mcpClients;
         // Initialize hooks registry
         this._hooksRegistry = new HookRegistryImplementation();
         // Initialize plugin registry with all plugins to be initialized during initialize()
+        // ModelPlugin is registered last so that on AfterInvocationEvent (which uses reverse
+        // callback ordering), it runs first — clearing messages before SessionManager saves.
         this._pluginRegistry = new PluginRegistry([
             this._conversationManager,
             ...(config?.plugins ?? []),
             ...(config?.sessionManager ? [config.sessionManager] : []),
+            new ModelPlugin(this.model),
         ]);
         if (config?.systemPrompt !== undefined) {
             this.systemPrompt = systemPromptFromData(config.systemPrompt);
@@ -177,6 +201,7 @@ export class Agent {
         this._tracer = new Tracer(config?.traceAttributes);
         // Initialize meter for local metrics accumulation
         this._meter = new Meter();
+        this._toolExecutor = config?.toolExecutor ?? 'concurrent';
         this._initialized = false;
     }
     /**
@@ -376,7 +401,7 @@ export class Agent {
                     yield await this._invokeCallbacks(result.value);
                     result = await streamGenerator.next();
                 }
-                yield await this._invokeCallbacks(new AgentResultEvent({ agent: this, result: result.value }));
+                yield await this._invokeCallbacks(new AgentResultEvent({ agent: this, result: result.value, invocationState: result.value.invocationState }));
                 return result.value;
             }
             catch (error) {
@@ -471,8 +496,26 @@ export class Agent {
         const structuredOutputSchema = options?.structuredOutputSchema ?? this._structuredOutputSchema;
         const structuredOutputTool = structuredOutputSchema ? new StructuredOutputTool(structuredOutputSchema) : undefined;
         let structuredOutputChoice;
-        // Emit event before the try block
-        yield new BeforeInvocationEvent({ agent: this });
+        // Resolve per-invocation state once. The same object is threaded through
+        // every lifecycle hook event, every tool context, and is surfaced on the
+        // AgentResult. Mutations by hooks/tools are visible across all recursive
+        // agent loop cycles within this invocation.
+        const invocationState = options?.invocationState ?? {};
+        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
@@ -510,11 +553,11 @@ export class Agent {
                     if (currentArgs !== undefined) {
                         const messagesToAppend = this._normalizeInput(currentArgs);
                         for (const message of messagesToAppend) {
-                            yield this._appendMessage(message);
+                            yield this._appendMessage(message, invocationState);
                         }
                         currentArgs = undefined;
                     }
-                    const modelResult = yield* this._invokeModel(structuredOutputChoice);
+                    const modelResult = yield* this._invokeModel(invocationState, structuredOutputChoice);
                     if (modelResult.stopReason !== 'toolUse') {
                         // If structured output is required, force it
                         if (structuredOutputTool) {
@@ -525,7 +568,7 @@ export class Agent {
                         }
                         this._meter.endCycle(cycleStartTime);
                         this._tracer.endAgentLoopSpan(cycleSpan);
-                        yield this._appendMessage(modelResult.message);
+                        yield this._appendMessage(modelResult.message, invocationState);
                         if (structuredOutputChoice) {
                             continue;
                         }
@@ -534,6 +577,7 @@ export class Agent {
                             lastMessage: modelResult.message,
                             traces: this._tracer.localTraces,
                             metrics: this._meter.metrics,
+                            invocationState,
                         });
                         return result;
                     }
@@ -546,8 +590,8 @@ export class Agent {
                             content: [new TextBlock('Tool execution cancelled')],
                         }));
                         const toolResultMessage = new Message({ role: 'user', content: cancelBlocks });
-                        yield this._appendMessage(modelResult.message);
-                        yield this._appendMessage(toolResultMessage);
+                        yield this._appendMessage(modelResult.message, invocationState);
+                        yield this._appendMessage(toolResultMessage, invocationState);
                         this._meter.endCycle(cycleStartTime);
                         this._tracer.endAgentLoopSpan(cycleSpan);
                         result = new AgentResult({
@@ -555,19 +599,20 @@ export class Agent {
                             lastMessage: modelResult.message,
                             traces: this._tracer.localTraces,
                             metrics: this._meter.metrics,
+                            invocationState,
                         });
                         return result;
                     }
                     // Execute tools
-                    const toolResultMessage = yield* this.executeTools(modelResult.message, this._toolRegistry);
+                    const toolResultMessage = yield* this.executeTools(modelResult.message, this._toolRegistry, invocationState);
                     /**
                      * Deferred append: both messages are added AFTER tool execution completes.
                      * This keeps agent.messages in a valid, reinvokable state at all times.
                      * If interrupted during tool execution, messages has no dangling toolUse
                      * without a matching toolResult, so the agent can be reinvoked cleanly.
                      */
-                    yield this._appendMessage(modelResult.message);
-                    yield this._appendMessage(toolResultMessage);
+                    yield this._appendMessage(modelResult.message, invocationState);
+                    yield this._appendMessage(toolResultMessage, invocationState);
                     this._meter.endCycle(cycleStartTime);
                     this._tracer.endAgentLoopSpan(cycleSpan);
                     // Structured output captured: exit
@@ -581,6 +626,7 @@ export class Agent {
                             traces: this._tracer.localTraces,
                             structuredOutput,
                             metrics: this._meter.metrics,
+                            invocationState,
                         });
                         return result;
                     }
@@ -600,12 +646,13 @@ export class Agent {
                     role: 'assistant',
                     content: [new TextBlock('Cancelled by user')],
                 });
-                yield this._appendMessage(cancelMessage);
+                yield this._appendMessage(cancelMessage, invocationState);
                 result = new AgentResult({
                     stopReason: 'cancelled',
                     lastMessage: cancelMessage,
                     traces: this._tracer.localTraces,
                     metrics: this._meter.metrics,
+                    invocationState,
                 });
                 return result;
             }
@@ -621,7 +668,7 @@ export class Agent {
                     role: 'assistant',
                     content: [new TextBlock('Cancelled by user')],
                 });
-                yield this._appendMessage(cancelMessage);
+                yield this._appendMessage(cancelMessage, invocationState);
             }
             this._tracer.endAgentSpan(agentSpan, {
                 ...(caughtError && { error: caughtError }),
@@ -634,7 +681,7 @@ export class Agent {
                 this._toolRegistry.remove(STRUCTURED_OUTPUT_TOOL_NAME);
             }
             // Always emit final event
-            yield new AfterInvocationEvent({ agent: this });
+            yield new AfterInvocationEvent({ agent: this, invocationState });
         }
     }
     /**
@@ -716,9 +763,9 @@ export class Agent {
      * @param toolChoice - Optional tool choice to force specific tool usage
      * @returns Object containing the assistant message, stop reason, and optional redaction message
      */
-    async *_invokeModel(toolChoice) {
+    async *_invokeModel(invocationState, toolChoice) {
         const toolSpecs = this._toolRegistry.list().map((tool) => tool.toolSpec);
-        const streamOptions = { toolSpecs };
+        const streamOptions = { toolSpecs, modelState: this.modelState };
         if (this.systemPrompt !== undefined) {
             streamOptions.systemPrompt = this.systemPrompt;
         }
@@ -726,7 +773,37 @@ export class Agent {
         if (toolChoice) {
             streamOptions.toolChoice = toolChoice;
         }
-        yield new BeforeModelCallEvent({ agent: this, model: this.model });
+        // Estimate input tokens for the upcoming model call (non-fatal if estimation fails)
+        let projectedInputTokens;
+        try {
+            projectedInputTokens = await this._estimateInputTokens(streamOptions);
+        }
+        catch (e) {
+            logger.debug(`error=<${e}> | token estimation failed, proceeding without estimate`);
+        }
+        const beforeModelCallEvent = new BeforeModelCallEvent({
+            agent: this,
+            model: this.model,
+            invocationState,
+            ...(projectedInputTokens !== undefined && { projectedInputTokens }),
+        });
+        yield beforeModelCallEvent;
+        if (beforeModelCallEvent.cancel) {
+            const cancelText = typeof beforeModelCallEvent.cancel === 'string' ? beforeModelCallEvent.cancel : 'model call denied by hook';
+            const message = new Message({ role: 'assistant', content: [new TextBlock(cancelText)] });
+            const stopData = { message, stopReason: 'endTurn' };
+            const afterModelCallEvent = new AfterModelCallEvent({
+                agent: this,
+                model: this.model,
+                stopData,
+                invocationState,
+            });
+            yield afterModelCallEvent;
+            if (afterModelCallEvent.retry) {
+                return yield* this._invokeModel(invocationState, toolChoice);
+            }
+            return { message, stopReason: 'endTurn' };
+        }
         // Start model span within loop span context
         const modelId = this.model.modelId;
         const modelSpan = this._tracer.startModelInvokeSpan({
@@ -735,7 +812,7 @@ export class Agent {
             ...(this.systemPrompt !== undefined && { systemPrompt: this.systemPrompt }),
         });
         try {
-            const result = yield* this._streamFromModel(this.messages, streamOptions);
+            const result = yield* this._streamFromModel(this.messages, streamOptions, invocationState);
             // Accumulate token usage and model latency metrics
             this._meter.updateCycle(result.metadata);
             // End model span with usage
@@ -747,7 +824,12 @@ export class Agent {
                 ...(usage && { usage }),
                 ...(metrics && { metrics }),
             });
-            yield new ModelMessageEvent({ agent: this, message: result.message, stopReason: result.stopReason });
+            yield new ModelMessageEvent({
+                agent: this,
+                message: result.message,
+                stopReason: result.stopReason,
+                invocationState,
+            });
             // Handle user content redaction if guardrails blocked input
             if (result.redaction?.userMessage) {
                 this._redactLastMessage(result.redaction.userMessage);
@@ -757,10 +839,15 @@ export class Agent {
                 stopReason: result.stopReason,
                 ...(result.redaction && { redaction: result.redaction }),
             };
-            const afterModelCallEvent = new AfterModelCallEvent({ agent: this, model: this.model, stopData });
+            const afterModelCallEvent = new AfterModelCallEvent({
+                agent: this,
+                model: this.model,
+                stopData,
+                invocationState,
+            });
             yield afterModelCallEvent;
             if (afterModelCallEvent.retry) {
-                return yield* this._invokeModel(toolChoice);
+                return yield* this._invokeModel(invocationState, toolChoice);
             }
             return result;
         }
@@ -769,7 +856,12 @@ export class Agent {
             // End model span with error
             this._tracer.endModelInvokeSpan(modelSpan, { error: modelError });
             // Create error event
-            const errorEvent = new AfterModelCallEvent({ agent: this, model: this.model, error: modelError });
+            const errorEvent = new AfterModelCallEvent({
+                agent: this,
+                model: this.model,
+                error: modelError,
+                invocationState,
+            });
             // Yield error event - stream will invoke hooks
             yield errorEvent;
             // Let CancelledError propagate directly — no retry
@@ -779,7 +871,7 @@ export class Agent {
             }
             // After yielding, hooks have been invoked and may have set retry
             if (errorEvent.retry) {
-                return yield* this._invokeModel(toolChoice);
+                return yield* this._invokeModel(invocationState, toolChoice);
             }
             // Re-throw error
             throw error;
@@ -801,7 +893,7 @@ export class Agent {
      * @param streamOptions - Options for streaming
      * @returns StreamAggregatedResult containing message, stop reason, and optional redaction message
      */
-    async *_streamFromModel(messages, streamOptions) {
+    async *_streamFromModel(messages, streamOptions, invocationState) {
         const streamGenerator = this.model.streamAggregated(messages, streamOptions);
         let result = await streamGenerator.next();
         while (!result.done) {
@@ -809,11 +901,11 @@ export class Agent {
             const event = result.value;
             if (isModelStreamEvent(event)) {
                 // ModelStreamEvent: wrap in ModelStreamUpdateEvent
-                yield new ModelStreamUpdateEvent({ agent: this, event });
+                yield new ModelStreamUpdateEvent({ agent: this, event, invocationState });
             }
             else {
                 // ContentBlock: wrap in ContentBlockEvent
-                yield new ContentBlockEvent({ agent: this, contentBlock: event });
+                yield new ContentBlockEvent({ agent: this, contentBlock: event, invocationState });
             }
             result = await streamGenerator.next();
         }
@@ -821,61 +913,173 @@ export class Agent {
         return result.value;
     }
     /**
-     * Executes tools sequentially and streams all tool events.
+     * Emits `BeforeToolsEvent`, handles the pre-launch cancel paths, then
+     * delegates per-tool execution to the configured {@link ToolExecutorStrategy}.
+     * Always pairs `BeforeToolsEvent` with a terminal `AfterToolsEvent`, even on
+     * the invariant-violation throw path.
      *
      * @param assistantMessage - The assistant message containing tool use blocks
      * @param toolRegistry - Registry containing available tools
      * @returns User message containing tool results
      */
-    async *executeTools(assistantMessage, toolRegistry) {
-        const beforeToolsEvent = new BeforeToolsEvent({ agent: this, message: assistantMessage });
+    async *executeTools(assistantMessage, toolRegistry, invocationState) {
+        const beforeToolsEvent = new BeforeToolsEvent({ agent: this, message: assistantMessage, invocationState });
         yield beforeToolsEvent;
+        const toolUseBlocks = assistantMessage.content.filter((block) => block.type === 'toolUseBlock');
+        if (toolUseBlocks.length === 0) {
+            // Preserve BeforeToolsEvent/AfterToolsEvent bracket symmetry even on
+            // this invariant-violation branch.
+            yield new AfterToolsEvent({
+                agent: this,
+                message: new Message({ role: 'user', content: [] }),
+                invocationState,
+            });
+            throw new Error('Model indicated toolUse but no tool use blocks found in message');
+        }
+        // Pre-launch cancel paths are strategy-independent.
+        if (beforeToolsEvent.cancel) {
+            const message = typeof beforeToolsEvent.cancel === 'string' ? beforeToolsEvent.cancel : 'Tool cancelled by hook';
+            return yield* this._yieldCancelledToolResults(toolUseBlocks, message, invocationState);
+        }
+        if (this.isCancelled) {
+            return yield* this._yieldCancelledToolResults(toolUseBlocks, 'Tool execution cancelled', invocationState);
+        }
+        switch (this._toolExecutor) {
+            case 'sequential':
+                return yield* this._executeToolsSequential(toolUseBlocks, toolRegistry, invocationState);
+            case 'concurrent':
+                return yield* this._executeToolsConcurrent(toolUseBlocks, toolRegistry, invocationState);
+            default: {
+                const _exhaustive = this._toolExecutor;
+                throw new Error(`Unknown toolExecutor: ${_exhaustive}`);
+            }
+        }
+    }
+    /**
+     * Emits a `ToolResultEvent` for every block plus an `AfterToolsEvent`, and
+     * returns the resulting tool-result message. Used by the pre-launch cancel
+     * paths shared across executors.
+     */
+    async *_yieldCancelledToolResults(toolUseBlocks, message, invocationState) {
+        const cancelBlocks = this._cancelAllAsResults(toolUseBlocks, message);
+        for (const result of cancelBlocks) {
+            yield new ToolResultEvent({ agent: this, result, invocationState });
+        }
+        const toolResultMessage = new Message({ role: 'user', content: cancelBlocks });
+        yield new AfterToolsEvent({ agent: this, message: toolResultMessage, invocationState });
+        return toolResultMessage;
+    }
+    /**
+     * Executes tools one at a time, honoring `agent.cancelSignal` between
+     * iterations to short-circuit not-yet-started tools.
+     */
+    async *_executeToolsSequential(toolUseBlocks, toolRegistry, invocationState) {
         const toolResultBlocks = [];
         let toolResultMessage;
         try {
-            // Extract tool use blocks from assistant message
-            const toolUseBlocks = assistantMessage.content.filter((block) => block.type === 'toolUseBlock');
-            if (toolUseBlocks.length === 0) {
-                // No tool use blocks found even though stopReason is toolUse
-                throw new Error('Model indicated toolUse but no tool use blocks found in message');
-            }
-            // Cancel all tools if hook requested it
-            if (beforeToolsEvent.cancel) {
-                const cancelMessage = cancelToolMessage(beforeToolsEvent.cancel);
-                const cancelBlocks = toolUseBlocks.map((block) => new ToolResultBlock({
-                    toolUseId: block.toolUseId,
-                    status: 'error',
-                    content: [new TextBlock(cancelMessage)],
-                }));
-                for (const result of cancelBlocks) {
-                    yield new ToolResultEvent({ agent: this, result });
+            for (const toolUseBlock of toolUseBlocks) {
+                if (this.isCancelled) {
+                    const cancelBlock = new ToolResultBlock({
+                        toolUseId: toolUseBlock.toolUseId,
+                        status: 'error',
+                        content: [new TextBlock('Tool execution cancelled')],
+                    });
+                    toolResultBlocks.push(cancelBlock);
+                    yield new ToolResultEvent({ agent: this, result: cancelBlock, invocationState });
+                    continue;
                 }
-                toolResultBlocks.push(...cancelBlocks);
+                const toolResultBlock = yield* this.executeTool(toolUseBlock, toolRegistry, invocationState);
+                toolResultBlocks.push(toolResultBlock);
+                yield new ToolResultEvent({ agent: this, result: toolResultBlock, invocationState });
             }
-            else {
-                for (const toolUseBlock of toolUseBlocks) {
-                    if (this.isCancelled) {
-                        const cancelBlock = new ToolResultBlock({
-                            toolUseId: toolUseBlock.toolUseId,
-                            status: 'error',
-                            content: [new TextBlock('Tool execution cancelled')],
-                        });
-                        toolResultBlocks.push(cancelBlock);
-                        yield new ToolResultEvent({ agent: this, result: cancelBlock });
-                        continue;
-                    }
-                    const toolResultBlock = yield* this.executeTool(toolUseBlock, toolRegistry);
-                    toolResultBlocks.push(toolResultBlock);
-                    yield new ToolResultEvent({ agent: this, result: toolResultBlock });
+        }
+        finally {
+            toolResultMessage = new Message({ role: 'user', content: toolResultBlocks });
+            yield new AfterToolsEvent({ agent: this, message: toolResultMessage, invocationState });
+        }
+        return toolResultMessage;
+    }
+    /**
+     * Produces one error ToolResultBlock per tool use block, each carrying
+     * `message` as its error text. Shared by pre-launch cancel paths.
+     */
+    _cancelAllAsResults(toolUseBlocks, message) {
+        return toolUseBlocks.map((block) => new ToolResultBlock({
+            toolUseId: block.toolUseId,
+            status: 'error',
+            content: [new TextBlock(message)],
+        }));
+    }
+    /**
+     * Executes tools concurrently by merging N per-tool {@link executeTool}
+     * async generators via `Promise.race`. Per-tool event order is preserved
+     * (because each generator is iterated serially); cross-tool events may
+     * interleave at race resolution boundaries.
+     *
+     * Per-tool retry (`AfterToolCallEvent.retry`) is isolated — it lives inside
+     * `executeTool`'s own `while(true)` loop, so one tool retrying does not
+     * disturb its siblings.
+     */
+    async *_executeToolsConcurrent(toolUseBlocks, toolRegistry, invocationState) {
+        let toolResultMessage;
+        const gens = toolUseBlocks.map((block) => ({
+            block,
+            gen: this.executeTool(block, toolRegistry, invocationState),
+        }));
+        const step = (idx) => gens[idx].gen.next().then((res) => ({ idx, kind: 'next', res }), (error) => ({ idx, kind: 'throw', error }));
+        const pendingNext = new Map(gens.map((_, idx) => [idx, step(idx)]));
+        const resultsByToolUseId = new Map();
+        try {
+            while (pendingNext.size > 0) {
+                const winner = await Promise.race(pendingNext.values());
+                const { idx } = winner;
+                const block = gens[idx].block;
+                if (winner.kind === 'throw') {
+                    pendingNext.delete(idx);
+                    const err = normalizeError(winner.error);
+                    const result = new ToolResultBlock({
+                        toolUseId: block.toolUseId,
+                        status: 'error',
+                        content: [new TextBlock(err.message)],
+                        error: err,
+                    });
+                    resultsByToolUseId.set(block.toolUseId, result);
+                    yield new ToolResultEvent({ agent: this, result, invocationState });
+                    continue;
+                }
+                if (winner.res.done) {
+                    pendingNext.delete(idx);
+                    resultsByToolUseId.set(block.toolUseId, winner.res.value);
+                    yield new ToolResultEvent({ agent: this, result: winner.res.value, invocationState });
+                }
+                else {
+                    yield winner.res.value;
+                    pendingNext.set(idx, step(idx));
                 }
             }
         }
         finally {
-            toolResultMessage = new Message({
-                role: 'user',
-                content: toolResultBlocks,
-            });
-            yield new AfterToolsEvent({ agent: this, message: toolResultMessage });
+            // Close any generators still in-flight (e.g. consumer broke out of stream).
+            await Promise.allSettled(Array.from(pendingNext.keys(), (idx) => gens[idx].gen.return(undefined)));
+            // Build the result message from whatever completed, in source order.
+            // Missing entries get a fallback error block so the message always
+            // accounts for every toolUseBlock the model emitted.
+            const toolResultBlocks = [];
+            for (const block of toolUseBlocks) {
+                const result = resultsByToolUseId.get(block.toolUseId);
+                if (result) {
+                    toolResultBlocks.push(result);
+                }
+                else {
+                    toolResultBlocks.push(new ToolResultBlock({
+                        toolUseId: block.toolUseId,
+                        status: 'error',
+                        content: [new TextBlock('Tool execution interrupted')],
+                    }));
+                }
+            }
+            toolResultMessage = new Message({ role: 'user', content: toolResultBlocks });
+            yield new AfterToolsEvent({ agent: this, message: toolResultMessage, invocationState });
         }
         return toolResultMessage;
     }
@@ -889,7 +1093,7 @@ export class Agent {
      * @param toolRegistry - Registry containing available tools
      * @returns Tool result block
      */
-    async *executeTool(toolUseBlock, toolRegistry) {
+    async *executeTool(toolUseBlock, toolRegistry, invocationState) {
         const tool = toolRegistry.get(toolUseBlock.name);
         // Create toolUse object for hook events and telemetry
         const toolUse = {
@@ -899,11 +1103,11 @@ export class Agent {
         };
         // Retry loop for tool execution
         while (true) {
-            const beforeToolCallEvent = new BeforeToolCallEvent({ agent: this, toolUse, tool });
+            const beforeToolCallEvent = new BeforeToolCallEvent({ agent: this, toolUse, tool, invocationState });
             yield beforeToolCallEvent;
             // Cancel individual tool if hook requested it
             if (beforeToolCallEvent.cancel) {
-                const cancelMessage = cancelToolMessage(beforeToolCallEvent.cancel);
+                const cancelMessage = typeof beforeToolCallEvent.cancel === 'string' ? beforeToolCallEvent.cancel : 'Tool cancelled by hook';
                 const toolResult = new ToolResultBlock({
                     toolUseId: toolUseBlock.toolUseId,
                     status: 'error',
@@ -914,12 +1118,13 @@ export class Agent {
                     toolUse,
                     tool,
                     result: toolResult,
+                    invocationState,
                 });
                 yield afterToolCallEvent;
                 if (afterToolCallEvent.retry) {
                     continue;
                 }
-                return toolResult;
+                return afterToolCallEvent.result;
             }
             // Start tool span within loop span context
             const toolSpan = this._tracer.startToolCallSpan({
@@ -946,6 +1151,7 @@ export class Agent {
                         input: toolUseBlock.input,
                     },
                     agent: this,
+                    invocationState,
                 };
                 try {
                     // Manually iterate tool stream to wrap each ToolStreamEvent in ToolStreamUpdateEvent.
@@ -956,7 +1162,7 @@ export class Agent {
                     const toolGenerator = this._tracer.withSpanContext(toolSpan, () => tool.stream(toolContext));
                     let toolNext = await this._tracer.withSpanContext(toolSpan, () => toolGenerator.next());
                     while (!toolNext.done) {
-                        yield new ToolStreamUpdateEvent({ agent: this, event: toolNext.value });
+                        yield new ToolStreamUpdateEvent({ agent: this, event: toolNext.value, invocationState });
                         toolNext = await this._tracer.withSpanContext(toolSpan, () => toolGenerator.next());
                     }
                     const result = toolNext.value;
@@ -998,13 +1204,14 @@ export class Agent {
                 toolUse,
                 tool,
                 result: toolResult,
+                invocationState,
                 ...(error !== undefined && { error }),
             });
             yield afterToolCallEvent;
             if (afterToolCallEvent.retry) {
                 continue;
             }
-            return toolResult;
+            return afterToolCallEvent.result;
         }
     }
     /**
@@ -1052,25 +1259,58 @@ export class Agent {
             }
         }
     }
+    /**
+     * Estimate the input token count for the next model call.
+     *
+     * Uses the token counting strategy: reads inputTokens + outputTokens
+     * from the last assistant message's metadata as a known baseline, then estimates
+     * only new messages added after it. Falls back to full estimation when no metadata
+     * is available (cold start or first call).
+     *
+     * @param streamOptions - The stream options containing system prompt and tool specs
+     * @returns Estimated input token count
+     */
+    async _estimateInputTokens(streamOptions) {
+        // Find the last assistant message with usage metadata
+        let lastAssistantIdx = -1;
+        for (let i = this.messages.length - 1; i >= 0; i--) {
+            if (this.messages[i].role === 'assistant' && this.messages[i].metadata?.usage) {
+                lastAssistantIdx = i;
+                break;
+            }
+        }
+        let estimate;
+        if (lastAssistantIdx >= 0) {
+            const usage = this.messages[lastAssistantIdx].metadata.usage;
+            const knownBaseline = usage.inputTokens + usage.outputTokens;
+            const newMessages = this.messages.slice(lastAssistantIdx + 1);
+            if (newMessages.length === 0) {
+                estimate = knownBaseline;
+            }
+            else {
+                // System prompt and tool spec tokens are already included in the baseline from the prior model call
+                estimate = knownBaseline + (await this.model.countTokens(newMessages));
+            }
+        }
+        else {
+            estimate = await this.model.countTokens(this.messages, {
+                ...(streamOptions.systemPrompt !== undefined && { systemPrompt: streamOptions.systemPrompt }),
+                ...(streamOptions.toolSpecs !== undefined && { toolSpecs: streamOptions.toolSpecs }),
+            });
+        }
+        return estimate;
+    }
     /**
      * Appends a message to the conversation history and returns the event for yielding.
      *
      * @param message - The message to append
      * @returns MessageAddedEvent to be yielded
      */
-    _appendMessage(message) {
+    _appendMessage(message, invocationState) {
         this.messages.push(message);
-        return new MessageAddedEvent({ agent: this, message });
+        return new MessageAddedEvent({ agent: this, message, invocationState });
     }
 }
-/**
- * Returns the cancel message for a cancelled tool.
- * @param cancelTool - The cancel value (true or custom message)
- * @returns The cancel message string
- */
-function cancelToolMessage(cancelTool) {
-    return typeof cancelTool === 'string' ? cancelTool : 'tool cancelled by hook';
-}
 /**
  * Recursively flattens nested arrays of tools into a single flat array.
  * @param tools - Tools or nested arrays of tools