@strands-agents/sdk 1.13.0 → 1.15.0

package/dist/src/agent/agent.js

@@ -7,8 +7,9 @@ import { isValidToolName } from '../tools/tool.js';
 import { cloneSystemPrompt, systemPromptFromData } from '../types/messages.js';
 import { normalizeError, ConcurrentInvocationError, StructuredOutputError } from '../errors.js';
 import { Model } from '../models/model.js';
+import { ModelRouter } from '../models/routing/router.js';
 import { ModelPlugin } from '../plugins/model-plugin.js';
-import { isModelStreamEvent } from '../models/streaming.js';
+import { totalPromptTokens, 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';
@@ -30,6 +31,7 @@ import { ConcurrentToolExecutor } from '../tools/executors/concurrent.js';
 import { SequentialToolExecutor } from '../tools/executors/sequential.js';
 import { AgentAsTool } from './agent-as-tool.js';
 import { ToolCaller } from './tool-caller.js';
+import { continuations } from './continuation.js';
 import { MemoryManager } from '../memory/memory-manager.js';
 import { SessionManager } from '../session/session-manager.js';
 import { Tracer } from '../telemetry/tracer.js';
@@ -44,8 +46,9 @@ 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';
+import { ContextManager } from '../context-manager/context-manager.js';
 /**
- * Supported values for the `contextManager` parameter.
+ * Supported string presets for the `contextManager` parameter.
  */
 export const CONTEXT_MANAGER_STRATEGIES = ['auto', 'agentic'];
 /** Benchmark-validated token threshold for offloading tool results. */
@@ -65,8 +68,16 @@ const CONTEXT_MANAGER_COMPRESSION_THRESHOLD = 0.85;
  * 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).
+ * When a ContextManager instance, uses NullConversationManager — the ContextManager owns
+ * overflow recovery via apply().
  */
 function resolveConversationManager(contextManager, conversationManager) {
+    if (contextManager === false) {
+        return conversationManager ?? new NullConversationManager();
+    }
+    if (contextManager instanceof ContextManager) {
+        return new NullConversationManager();
+    }
     if (contextManager === 'agentic') {
         return (conversationManager ??
             new SummarizingConversationManager({
@@ -132,6 +143,7 @@ export class Agent {
      * The model provider used by the agent for inference.
      */
     model;
+    _modelRouter;
     /**
      * The system prompt to pass to the model provider.
      */
@@ -148,10 +160,15 @@ export class Agent {
      * Optional description of what the agent does.
      */
     description;
+    /**
+     * The context manager for strategy-driven offloading, if configured.
+     */
+    contextManager;
     /**
      * The session manager for saving and restoring agent sessions, if configured.
      */
     sessionManager;
+    _sessionId;
     /**
      * The memory manager for cross-session memory retrieval and storage, if configured.
      */
@@ -170,6 +187,20 @@ export class Agent {
     get sandbox() {
         return this._sandbox || defaultSandbox.get();
     }
+    /**
+     * A stable, unique identifier for the current conversation session.
+     *
+     * If a SessionManager is attached, delegates to its sessionId.
+     * Otherwise, lazily generates and caches a random 8-character hex string.
+     */
+    get sessionId() {
+        if (this.sessionManager)
+            return this.sessionManager.sessionId;
+        if (!this._sessionId) {
+            this._sessionId = globalThis.crypto.randomUUID().slice(0, 8);
+        }
+        return this._sessionId;
+    }
     _hooksRegistry;
     _middlewareRegistry;
     _pluginRegistry;
@@ -207,6 +238,7 @@ export class Agent {
         this.id = config?.id ?? DEFAULT_AGENT_ID;
         if (config?.description !== undefined)
             this.description = config.description;
+        this.contextManager = config?.contextManager instanceof ContextManager ? config.contextManager : undefined;
         this.sessionManager = config?.sessionManager;
         this.storage = config?.storage;
         this.memoryManager =
@@ -216,11 +248,19 @@ export class Agent {
                     ? new MemoryManager(config.memoryManager)
                     : undefined;
         this._sandbox = config?.sandbox;
-        if (typeof config?.model === 'string') {
-            this.model = new BedrockModel({ modelId: config.model });
+        const configuredModel = config?.model;
+        if (typeof configuredModel === 'string') {
+            this.model = new BedrockModel({ modelId: configuredModel });
+        }
+        else if (configuredModel instanceof ModelRouter) {
+            this._modelRouter = configuredModel;
+            this.model = configuredModel.defaultModel;
         }
         else {
-            this.model = config?.model ?? new BedrockModel();
+            this.model = configuredModel ?? new BedrockModel();
+        }
+        if (config?.plugins?.some((plugin) => plugin instanceof ModelRouter)) {
+            throw new Error('ModelRouter must be passed through Agent({ model }), not plugins');
         }
         // Validate and assign conversation manager
         if (this.model.stateful) {
@@ -243,8 +283,9 @@ export class Agent {
         this._interventionRegistry = new InterventionRegistry(config?.interventions ?? [], this._hooksRegistry);
         // Initialize middleware registry
         this._middlewareRegistry = new MiddlewareRegistry();
+        this._modelRouter?.attachToAgent(this);
         if (config?.contextManager === 'agentic') {
-            this._middlewareRegistry.addInput(InvokeModelStage.Input, createTokenUsageMiddleware(this.model));
+            this._middlewareRegistry.addInput(InvokeModelStage.Input, createTokenUsageMiddleware());
         }
         // `undefined` (omitted) → install the default; `null`/`[]` → explicit opt-out.
         const retryStrategies = config?.retryStrategy === null
@@ -268,7 +309,9 @@ export class Agent {
         // when a delegate tool is added (construction, plugin getTools, MCP, runtime).
         // The plugin is a no-op when no delegation tools fire.
         const hasAgentDelegation = (config?.plugins ?? []).some((p) => p.name === 'strands:agent-delegation');
+        const contextManagerPlugin = config?.contextManager instanceof ContextManager ? config.contextManager : undefined;
         this._pluginRegistry = new PluginRegistry([
+            ...(this._modelRouter ? [this._modelRouter] : []),
             this._conversationManager,
             ...retryStrategies,
             ...(config?.plugins ?? []),
@@ -284,6 +327,7 @@ export class Agent {
                 ]
                 : []),
             ...(this.memoryManager ? [this.memoryManager] : []),
+            ...(contextManagerPlugin ? [contextManagerPlugin] : []),
             ...(config?.sessionManager ? [config.sessionManager] : []),
             new ModelPlugin(this.model),
         ]);
@@ -533,7 +577,7 @@ export class Agent {
     /**
      * The cancellation signal for the current invocation.
      *
-     * Tools can pass this to cancellable operations (e.g., `fetch(url, { signal: agent.cancelSignal })`).
+     * SDK-managed tool contexts receive this as `context.cancelSignal` for cancellable operations.
      * Hooks can check `event.agent.cancelSignal.aborted` to detect cancellation.
      */
     get cancelSignal() {
@@ -549,7 +593,7 @@ export class Agent {
      * - At the top of each agent loop cycle
      *
      * If a tool is already executing, it will run to completion unless
-     * the tool checks {@link LocalAgent.cancelSignal | cancelSignal} internally.
+     * the tool checks `context.cancelSignal` internally.
      *
      * Hook callbacks can check `event.agent.cancelSignal.aborted` to detect
      * cancellation and adjust their behavior accordingly.
@@ -637,6 +681,7 @@ export class Agent {
      */
     async *stream(args, options) {
         this.acquireLock();
+        let continuationEvent;
         try {
             await this.initialize();
             // Thread the resolved invocationState so all layers share the same reference.
@@ -664,6 +709,8 @@ export class Agent {
                     const message = new Message({ role: 'assistant', content: [new TextBlock(cancelText)] });
                     yield this._appendMessage(message, invocationState);
                     const afterEvent = new AfterInvocationEvent({ agent: this, invocationState });
+                    await continuations.abandon(continuationEvent, new Error('Continuation was not incorporated into agent history'));
+                    continuationEvent = afterEvent;
                     await this._invokeCallbacks(afterEvent);
                     yield afterEvent;
                     return new AgentResult({
@@ -678,7 +725,7 @@ export class Agent {
                 let caughtError;
                 const afterInvocationEvent = new AfterInvocationEvent({ agent: this, invocationState });
                 try {
-                    result = yield* this._streamWithMiddleware(currentArgs, resolvedOptions, invocationState);
+                    result = yield* this._streamWithMiddleware(currentArgs, resolvedOptions, invocationState, continuationEvent);
                 }
                 catch (error) {
                     caughtError = error;
@@ -687,6 +734,8 @@ export class Agent {
                     // 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 continuations.abandon(continuationEvent, new Error('Continuation was not incorporated into agent history'));
+                    continuationEvent = afterInvocationEvent;
                     await this._invokeCallbacks(afterInvocationEvent);
                 }
                 // Yield outside finally — in JS, a `yield` inside `finally` suspends the generator
@@ -697,9 +746,15 @@ export class Agent {
                 if (caughtError) {
                     throw caughtError;
                 }
-                // Resume only on a clean invocation — errors propagate above.
-                if (afterInvocationEvent.resume !== undefined) {
-                    currentArgs = afterInvocationEvent.resume;
+                const stopReason = result.stopReason;
+                const allowsContinuation = stopReason === 'endTurn' || stopReason === 'stopSequence';
+                if (!allowsContinuation && stopReason !== 'interrupt') {
+                    await continuations.abandon(afterInvocationEvent, new Error(`Continuation abandoned after ${stopReason}`));
+                }
+                const hasContinuation = (await continuations.prepare(afterInvocationEvent, (continuationArgs) => this._normalizeInput(continuationArgs), stopReason)) !== undefined;
+                continuationEvent = hasContinuation ? afterInvocationEvent : undefined;
+                if (hasContinuation || afterInvocationEvent.resume !== undefined) {
+                    currentArgs = afterInvocationEvent.resume ?? [];
                     continue;
                 }
                 // Only emit AgentResultEvent on the final iteration (not on resumed ones).
@@ -712,6 +767,7 @@ export class Agent {
             }
         }
         finally {
+            await continuations.abandon(continuationEvent, new Error('Agent stream closed before continuation input was incorporated into agent history'));
             this._isInvoking = false;
         }
     }
@@ -719,7 +775,7 @@ export class Agent {
      * Invokes the AgentStreamStage middleware chain.
      * Hooks fire outside this method (in stream()'s resume loop).
      */
-    async *_streamWithMiddleware(args, options, invocationState) {
+    async *_streamWithMiddleware(args, options, invocationState, continuationEvent) {
         // Snapshot so a gate that re-reads its response after next() still resolves even if a tool cycle called deactivate().
         const interruptsSnapshot = { ...this._interruptState.interrupts };
         const context = {
@@ -733,7 +789,8 @@ export class Agent {
         const self = this;
         try {
             const { result } = yield* this._middlewareRegistry.invoke(AgentStreamStage, context, async function* (ctx) {
-                const result = yield* self._streamCore(ctx.args, ctx.options);
+                const streamArgs = continuations.combine(continuationEvent, ctx.args, (continuationArgs) => self._normalizeInput(continuationArgs));
+                const result = yield* self._streamCore(streamArgs, ctx.options, streamArgs === ctx.args ? undefined : continuationEvent);
                 return { result };
             });
             if (this._interruptState.activated &&
@@ -770,8 +827,8 @@ export class Agent {
      * 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);
+    async *_streamCore(args, options, continuationEvent) {
+        const streamGenerator = this._stream(args, options, continuationEvent);
         let caughtError;
         let iterationResult;
         try {
@@ -933,7 +990,7 @@ export class Agent {
      * @param options - Optional per-invocation options
      * @returns Async generator that yields AgentStreamEvent objects and returns AgentResult
      */
-    async *_stream(args, options) {
+    async *_stream(args, options, continuationEvent) {
         let currentArgs = args;
         let result;
         this._validateLimits(options);
@@ -1016,8 +1073,13 @@ export class Agent {
                     // Normalize input and append user messages on first invocation only
                     if (currentArgs !== undefined) {
                         const messagesToAppend = this._normalizeInput(currentArgs);
-                        for (const message of messagesToAppend) {
-                            yield this._appendMessage(message, invocationState);
+                        if (continuationEvent) {
+                            yield* await this._appendContinuationMessages(messagesToAppend, continuationEvent, invocationState);
+                        }
+                        else {
+                            for (const message of messagesToAppend) {
+                                yield this._appendMessage(message, invocationState);
+                            }
                         }
                         currentArgs = undefined;
                     }
@@ -1140,13 +1202,16 @@ export class Agent {
                     }
                     this._meter.endCycle(cycleStartTime);
                     this._tracer.endAgentLoopSpan(cycleSpan);
-                    // Hook requested halt: exit without calling the model again
+                    // Hook requested halt with content: exit without calling the model again.
                     const { afterToolsEvent } = toolsResult;
-                    if (afterToolsEvent.endTurn) {
-                        const endTurnText = typeof afterToolsEvent.endTurn === 'string'
-                            ? afterToolsEvent.endTurn
-                            : 'Turn ended early by hook after tool execution';
-                        const lastMessage = new Message({ role: 'assistant', content: [new TextBlock(endTurnText)] });
+                    const endTurnValue = afterToolsEvent.endTurn;
+                    if (endTurnValue === true || (endTurnValue !== false && endTurnValue.length > 0)) {
+                        const endTurnContent = Array.isArray(endTurnValue)
+                            ? [...endTurnValue]
+                            : [
+                                new TextBlock(typeof endTurnValue === 'string' ? endTurnValue : 'Turn ended early by hook after tool execution'),
+                            ];
+                        const lastMessage = new Message({ role: 'assistant', content: endTurnContent });
                         yield this._appendMessage(lastMessage, invocationState);
                         result = new AgentResult({
                             stopReason: 'endTurn',
@@ -1203,7 +1268,9 @@ export class Agent {
                     role: 'assistant',
                     content: [new TextBlock('Cancelled by user')],
                 });
-                yield this._appendMessage(cancelMessage, invocationState);
+                if (this._hasOpenUserTurn()) {
+                    yield this._appendMessage(cancelMessage, invocationState);
+                }
                 result = new AgentResult({
                     stopReason: 'cancelled',
                     lastMessage: cancelMessage,
@@ -1233,14 +1300,16 @@ export class Agent {
         }
         finally {
             // If cancelled but the catch block was bypassed (generator terminated
-            // via .return() when the consumer breaks out of for-await), append an
-            // assistant message so the agent can be reinvoked with a new user prompt.
+            // via .return() when the consumer breaks out of for-await), close an
+            // existing user turn so the agent can be reinvoked.
             if (!caughtError && !result && this.isCancelled) {
                 const cancelMessage = new Message({
                     role: 'assistant',
                     content: [new TextBlock('Cancelled by user')],
                 });
-                yield this._appendMessage(cancelMessage, invocationState);
+                if (this._hasOpenUserTurn()) {
+                    yield this._appendMessage(cancelMessage, invocationState);
+                }
             }
             this._tracer.endAgentSpan(agentSpan, {
                 ...(caughtError && { error: caughtError }),
@@ -1412,48 +1481,70 @@ export class Agent {
         if (this.systemPrompt !== undefined) {
             streamOptions.systemPrompt = this.systemPrompt;
         }
-        // Add tool choice if provided
         if (toolChoice) {
             streamOptions.toolChoice = toolChoice;
         }
         let attemptCount = 1;
         while (true) {
-            // Estimate input tokens for the upcoming model call (non-fatal if estimation fails)
+            const selectedModel = this._modelForAttempt(invocationState);
             let projectedInputTokens;
             try {
+                // Context management continues to size against the agent's default model.
                 projectedInputTokens = await this._estimateInputTokens(streamOptions);
             }
-            catch (e) {
-                logger.debug(`error=<${e}> | token estimation failed, proceeding without estimate`);
+            catch (error) {
+                logger.debug(`error=<${error}> | token estimation failed, proceeding without estimate`);
             }
             const beforeModelCallEvent = new BeforeModelCallEvent({
                 agent: this,
-                model: this.model,
+                model: selectedModel,
                 invocationState,
                 ...(projectedInputTokens !== undefined && { projectedInputTokens }),
             });
-            yield beforeModelCallEvent;
+            let modelContinuation;
+            try {
+                yield beforeModelCallEvent;
+                modelContinuation = await continuations.prepare(beforeModelCallEvent, (continuationArgs) => this._normalizeInput(continuationArgs));
+            }
+            finally {
+                if (modelContinuation === undefined) {
+                    await continuations.abandon(beforeModelCallEvent, new Error('Agent stream closed before continuation input was incorporated into agent history'));
+                }
+            }
             if (beforeModelCallEvent.cancel) {
+                await continuations.abandon(beforeModelCallEvent, new Error('Continuation abandoned by BeforeModelCallEvent'));
                 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,
+                    model: selectedModel,
                     attemptCount,
                     stopData,
                     invocationState,
                 });
                 yield afterModelCallEvent;
                 if (afterModelCallEvent.retry) {
-                    attemptCount += 1;
+                    attemptCount = this._nextAttemptCount(selectedModel, invocationState, attemptCount);
                     continue;
                 }
                 return { message, stopReason: 'endTurn' };
             }
+            yield* await this._appendContinuationMessages(modelContinuation ?? [], beforeModelCallEvent, invocationState);
+            if (modelContinuation !== undefined) {
+                try {
+                    projectedInputTokens = await this._estimateInputTokens(streamOptions);
+                }
+                catch (error) {
+                    projectedInputTokens = undefined;
+                    logger.debug(`error=<${error}> | token estimation failed after continuation input, proceeding without estimate`);
+                }
+            }
+            const invokedModelRef = {};
             try {
-                const result = yield* this._invokeModelWithMiddleware(invocationState, toolChoice, projectedInputTokens);
-                // Accumulate token usage and model latency metrics
+                const result = yield* this._invokeModelWithMiddleware(invocationState, selectedModel, invokedModelRef, toolChoice, projectedInputTokens);
+                const routedModel = this._modelRouter?.getRoutedModel(this, invocationState) ?? selectedModel;
+                const model = invokedModelRef.model ?? selectedModel;
                 this._meter.updateCycle(result.metadata);
                 yield new ModelMessageEvent({
                     agent: this,
@@ -1461,7 +1552,6 @@ export class Agent {
                     stopReason: result.stopReason,
                     invocationState,
                 });
-                // Handle user content redaction if guardrails blocked input
                 if (result.redaction?.userMessage) {
                     this._redactLastMessage(result.redaction.userMessage);
                 }
@@ -1472,45 +1562,50 @@ export class Agent {
                 };
                 const afterModelCallEvent = new AfterModelCallEvent({
                     agent: this,
-                    model: this.model,
+                    model,
                     attemptCount,
                     stopData,
                     invocationState,
                 });
                 yield afterModelCallEvent;
                 if (afterModelCallEvent.retry) {
-                    attemptCount += 1;
+                    attemptCount = this._nextAttemptCount(routedModel, invocationState, attemptCount);
                     continue;
                 }
                 return result;
             }
             catch (error) {
+                const routedModel = this._modelRouter?.getRoutedModel(this, invocationState) ?? selectedModel;
+                // A failure before the terminal ran (e.g. in input middleware) is attributed to the routed model.
+                const failedModel = invokedModelRef.model ?? routedModel;
                 const modelError = normalizeError(error);
-                // Create error event
                 const errorEvent = new AfterModelCallEvent({
                     agent: this,
-                    model: this.model,
+                    model: failedModel,
                     attemptCount,
                     error: modelError,
                     invocationState,
                 });
-                // Yield error event - stream will invoke hooks
                 yield errorEvent;
-                // Let CancelledError propagate directly — no retry
-                // (we emit the AfterModelCall because we already emitted Before and we guarentee the pair)
+                // Preserve the Before/After event pair, but never retry cancellation.
                 if (error instanceof CancelledError) {
                     throw error;
                 }
-                // After yielding, hooks have been invoked and may have set retry
                 if (errorEvent.retry) {
-                    attemptCount += 1;
+                    attemptCount = this._nextAttemptCount(routedModel, invocationState, attemptCount);
                     continue;
                 }
-                // Re-throw error
                 throw error;
             }
         }
     }
+    _modelForAttempt(invocationState) {
+        return this._modelRouter?.getRoutedModel(this, invocationState) ?? this.model;
+    }
+    _nextAttemptCount(model, invocationState, attemptCount) {
+        const nextModel = this._modelRouter?.getRoutedModel(this, invocationState);
+        return nextModel !== undefined && nextModel !== model ? 1 : attemptCount + 1;
+    }
     /**
      * Invokes the model through the InvokeModelStage middleware chain.
      * Builds an InvokeModelContext from current agent state and composes the
@@ -1518,12 +1613,15 @@ export class Agent {
      * using context fields directly (not re-derived from the agent).
      *
      * @param invocationState - Per-invocation state shared across hooks and tools
+     * @param selectedModel - Model the chain starts with; middleware may replace it
+     * @param invokedModelRef - Slot filled with the model the terminal actually invoked
      * @param toolChoice - Optional tool choice to force specific tool usage
      * @returns StreamAggregatedResult from the model (or middleware short-circuit)
      */
-    async *_invokeModelWithMiddleware(invocationState, toolChoice, projectedInputTokens) {
+    async *_invokeModelWithMiddleware(invocationState, selectedModel, invokedModelRef, toolChoice, projectedInputTokens) {
         const context = {
             agent: this,
+            model: selectedModel,
             messages: this.messages.map((msg) => msg.clone()),
             ...(this.systemPrompt !== undefined && { systemPrompt: cloneSystemPrompt(this.systemPrompt) }),
             toolSpecs: deepCopy(this._toolRegistry.list().map((tool) => tool.toolSpec)),
@@ -1540,7 +1638,8 @@ export class Agent {
         // 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;
+            invokedModelRef.model = ctx.model;
+            const modelId = ctx.model.modelId;
             const modelSpan = self._tracer.startModelInvokeSpan({
                 messages: ctx.messages,
                 ...(modelId && { modelId }),
@@ -1556,8 +1655,10 @@ export class Agent {
                     modelState: tempModelState,
                     ...(ctx.systemPrompt !== undefined && { systemPrompt: ctx.systemPrompt }),
                     ...(ctx.toolChoice && { toolChoice: ctx.toolChoice }),
+                    // Omitted when zero, so an ordinary call's options are unchanged.
+                    ...(ctx.dynamicTrailingBlocks ? { dynamicTrailingBlocks: ctx.dynamicTrailingBlocks } : {}),
                 };
-                const gen = self._streamFromModel(ctx.messages, streamOptions, ctx.invocationState);
+                const gen = self._streamFromModel(ctx.model, ctx.messages, streamOptions, ctx.invocationState);
                 let iterResult = await gen.next();
                 while (!iterResult.done) {
                     yield iterResult.value;
@@ -1598,13 +1699,14 @@ export class Agent {
      * These are separate event classes because they represent different granularities
      * (partial deltas vs finished blocks). Both are yielded in the stream and hookable.
      *
+     * @param model - Model selected for this call
      * @param messages - Messages to send to the model
      * @param streamOptions - Options for streaming
      * @returns StreamAggregatedResult containing message, stop reason, and optional redaction message
      */
-    async *_streamFromModel(messages, streamOptions, invocationState) {
+    async *_streamFromModel(model, messages, streamOptions, invocationState) {
         messages = normalizeToolUseNames(messages);
-        const streamGenerator = this.model.streamAggregated(messages, streamOptions);
+        const streamGenerator = model.streamAggregated(messages, streamOptions);
         try {
             let result = await streamGenerator.next();
             while (!result.done) {
@@ -1685,6 +1787,7 @@ export class Agent {
                     middlewareRegistry: this._middlewareRegistry,
                     tracer: this._tracer,
                     meter: this._meter,
+                    cancelSignal: this._abortSignal,
                 }, {
                     toolUseBlocks,
                     toolResultBlocks,
@@ -1762,10 +1865,10 @@ 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).
+     * Uses the token counting strategy: reads the total prompt the model processed (including cached
+     * tokens) plus 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
@@ -1782,7 +1885,7 @@ export class Agent {
         let estimate;
         if (lastAssistantIdx >= 0) {
             const usage = this.messages[lastAssistantIdx].metadata.usage;
-            const knownBaseline = usage.inputTokens + usage.outputTokens;
+            const knownBaseline = totalPromptTokens(usage) + usage.outputTokens;
             const newMessages = this.messages.slice(lastAssistantIdx + 1);
             if (newMessages.length === 0) {
                 estimate = knownBaseline;
@@ -1821,6 +1924,36 @@ export class Agent {
         this.messages.push(message);
         return new MessageAddedEvent({ agent: this, message, invocationState });
     }
+    async _appendContinuationMessages(messages, continuationEvent, invocationState) {
+        const events = [];
+        for (const message of messages) {
+            const lastMessage = this.messages.at(-1);
+            if (lastMessage?.role !== message.role) {
+                events.push(this._appendMessage(message, invocationState));
+                continue;
+            }
+            const appendedMessage = new Message({
+                role: message.role,
+                content: [...lastMessage.content, ...message.content],
+                trackingId: lastMessage.trackingId,
+                ...(lastMessage.metadata !== undefined && { metadata: lastMessage.metadata }),
+            });
+            this.messages[this.messages.length - 1] = appendedMessage;
+            const messageEvent = new MessageAddedEvent({ agent: this, message: appendedMessage, invocationState });
+            if (events.at(-1)?.message === lastMessage) {
+                events[events.length - 1] = messageEvent;
+            }
+            else {
+                events.push(messageEvent);
+            }
+        }
+        await continuations.markAppended(continuationEvent);
+        return events;
+    }
+    _hasOpenUserTurn() {
+        const lastMessage = this.messages[this.messages.length - 1];
+        return lastMessage?.role === 'user';
+    }
 }
 const INVALID_TOOL_NAME_PLACEHOLDER = 'INVALID_TOOL_NAME';
 /**