@strands-agents/sdk 1.14.0 → 1.16.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';
@@ -46,6 +47,7 @@ import { takeSnapshot as takeSnapshotInternal, loadSnapshot as loadSnapshotInter
 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';
+import { BackgroundTasks } from '../background-tasks/background-tasks.js';
 /**
  * Supported string presets for the `contextManager` parameter.
  */
@@ -142,6 +144,7 @@ export class Agent {
      * The model provider used by the agent for inference.
      */
     model;
+    _modelRouter;
     /**
      * The system prompt to pass to the model provider.
      */
@@ -166,6 +169,7 @@ export class Agent {
      * 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.
      */
@@ -184,6 +188,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;
@@ -208,6 +226,7 @@ export class Agent {
     _checkpointing;
     /** Direct tool caller — created via {@link ToolCaller.create} factory. */
     _toolCaller;
+    _backgroundTasks;
     /**
      * Creates an instance of the Agent.
      * @param config - The configuration for the agent.
@@ -231,11 +250,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) {
@@ -258,6 +285,7 @@ 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());
         }
@@ -284,10 +312,26 @@ export class Agent {
         // 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._backgroundTasks = config?.backgroundTasks
+            ? new BackgroundTasks(config.backgroundTasks === true ? {} : config.backgroundTasks, (tool, context, middlewareInterrupt) => this._toolExecutor.executeBackground({
+                agent: this,
+                middlewareRegistry: this._middlewareRegistry,
+                // Isolate mutable local trace state from overlapping agent work. OTel spans
+                // still use the global provider; local traces are not added to AgentResult.traces.
+                tracer: new Tracer(config?.traceAttributes),
+                meter: this._meter,
+                cancelSignal: context.cancelSignal,
+                toolInterrupt: context.interrupt,
+                middlewareInterrupt,
+                toolGuard: (selectedTool) => this._backgroundTasks.assertToolCanRun(selectedTool),
+            }, context.toolUse, tool, context.invocationState, (event) => this._invokeCallbacks(event)))
+            : undefined;
         this._pluginRegistry = new PluginRegistry([
+            ...(this._modelRouter ? [this._modelRouter] : []),
             this._conversationManager,
             ...retryStrategies,
             ...(config?.plugins ?? []),
+            ...(this._backgroundTasks ? [this._backgroundTasks] : []),
             ...(!hasAgentDelegation ? [new AgentDelegation()] : []),
             ...((config?.contextManager === 'auto' || config?.contextManager === 'agentic') && !hasOffloader
                 ? [
@@ -940,7 +984,11 @@ export class Agent {
      * ```
      */
     loadSnapshot(snapshot) {
+        if (this._initialized)
+            this._backgroundTasks?.assertCanLoadSnapshot();
         loadSnapshotInternal(this, snapshot);
+        if (this._initialized && 'state' in snapshot.data)
+            this._backgroundTasks?._loadAppState();
     }
     /**
      * Invokes hook callbacks and printer for a stream event.
@@ -1042,6 +1090,24 @@ export class Agent {
                     cycleId,
                     messages: this.messages,
                 });
+                // Closes cycle telemetry exactly once, however the cycle exits: return,
+                // continue, throw, or the consumer closing the public iterator (which
+                // resumes the generator at the suspended yield and runs the finally).
+                // Called before building an AgentResult so its metrics include this cycle.
+                let cycleClosed = false;
+                const closeCycle = (error) => {
+                    if (cycleClosed) {
+                        return;
+                    }
+                    cycleClosed = true;
+                    this._meter.endCycle(cycleStartTime);
+                    if (error) {
+                        this._tracer.endAgentLoopSpan(cycleSpan, { error });
+                    }
+                    else {
+                        this._tracer.endAgentLoopSpan(cycleSpan);
+                    }
+                };
                 try {
                     // Normalize input and append user messages on first invocation only
                     if (currentArgs !== undefined) {
@@ -1073,8 +1139,7 @@ export class Agent {
                             if (structuredOutputTool && structuredOutputChoice) {
                                 throw new StructuredOutputError('The model failed to invoke the structured output tool even after it was forced.');
                             }
-                            this._meter.endCycle(cycleStartTime);
-                            this._tracer.endAgentLoopSpan(cycleSpan);
+                            closeCycle();
                             // Schema set, model ignored the tool — drop the response and force the tool next cycle.
                             // Appending the plain-text turn here would leave the conversation ending on an
                             // assistant message, which providers like Bedrock reject as assistant prefill.
@@ -1105,8 +1170,7 @@ export class Agent {
                             const toolResultMessage = new Message({ role: 'user', content: cancelBlocks });
                             yield this._appendMessage(modelResult.message, invocationState);
                             yield this._appendMessage(toolResultMessage, invocationState);
-                            this._meter.endCycle(cycleStartTime);
-                            this._tracer.endAgentLoopSpan(cycleSpan);
+                            closeCycle();
                             result = new AgentResult({
                                 stopReason: 'cancelled',
                                 lastMessage: modelResult.message,
@@ -1126,8 +1190,7 @@ export class Agent {
                             const priorResumePosition = resumePosition;
                             resumePosition = undefined;
                             if (priorResumePosition !== 'afterModel') {
-                                this._meter.endCycle(cycleStartTime);
-                                this._tracer.endAgentLoopSpan(cycleSpan);
+                                closeCycle();
                                 result = new AgentResult({
                                     stopReason: 'checkpoint',
                                     lastMessage: modelResult.message,
@@ -1143,19 +1206,18 @@ export class Agent {
                     }
                     // Execute tools
                     const toolsResult = yield* this.executeTools(assistantMessage, invocationState, completedToolResults);
-                    // When the consumer breaks the stream (e.g. agent.cancel() + break),
-                    // yield* returns undefined because the inner generator was closed.
+                    // Reached when the consumer breaks the stream during tool execution.
+                    // _streamCore's drain loop calls .return(), which runs the finally in
+                    // executeTools and yields its AfterToolsEvent, and the drain's next()
+                    // then completes the closed generator, so this yield* evaluates to
+                    // undefined instead of unwinding.
                     if (!toolsResult) {
-                        this._meter.endCycle(cycleStartTime);
-                        this._tracer.endAgentLoopSpan(cycleSpan);
                         continue;
                     }
                     const toolResultMessage = toolsResult.message;
                     // Tools were skipped (not executed) — preserve pending state so the next resume
                     // can run them.
                     if (this.isCancelled && toolsResult.toolsSkipped && this._interruptState.pendingToolExecution) {
-                        this._meter.endCycle(cycleStartTime);
-                        this._tracer.endAgentLoopSpan(cycleSpan);
                         continue;
                     }
                     /**
@@ -1173,15 +1235,17 @@ export class Agent {
                     if (this._interruptState.activated) {
                         this._interruptState.deactivate();
                     }
-                    this._meter.endCycle(cycleStartTime);
-                    this._tracer.endAgentLoopSpan(cycleSpan);
-                    // Hook requested halt: exit without calling the model again
+                    closeCycle();
+                    // 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',
@@ -1224,10 +1288,12 @@ export class Agent {
                     }
                 }
                 catch (error) {
-                    this._meter.endCycle(cycleStartTime);
-                    this._tracer.endAgentLoopSpan(cycleSpan, { error: error });
+                    closeCycle(error);
                     throw error;
                 }
+                finally {
+                    closeCycle();
+                }
             }
         }
         catch (error) {
@@ -1451,23 +1517,23 @@ 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 }),
             });
@@ -1488,14 +1554,14 @@ export class Agent {
                 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' };
@@ -1510,9 +1576,11 @@ export class Agent {
                     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,
@@ -1520,7 +1588,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);
                 }
@@ -1531,45 +1598,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
@@ -1577,13 +1649,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: this.model,
+            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)),
@@ -1600,6 +1674,7 @@ 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) {
+            invokedModelRef.model = ctx.model;
             const modelId = ctx.model.modelId;
             const modelSpan = self._tracer.startModelInvokeSpan({
                 messages: ctx.messages,
@@ -1749,6 +1824,10 @@ export class Agent {
                     tracer: this._tracer,
                     meter: this._meter,
                     cancelSignal: this._abortSignal,
+                    ...(this._backgroundTasks && {
+                        backgroundTasks: this._backgroundTasks,
+                        backgroundTaskPassId: assistantMessage.trackingId,
+                    }),
                 }, {
                     toolUseBlocks,
                     toolResultBlocks,
@@ -1826,10 +1905,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
@@ -1846,7 +1925,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;