diff --git a/packages/@n8n/agents/src/runtime/loop/agent-runtime.ts b/packages/@n8n/agents/src/runtime/loop/agent-runtime.ts index aeb7bd5910d..974b12d2d9b 100644 --- a/packages/@n8n/agents/src/runtime/loop/agent-runtime.ts +++ b/packages/@n8n/agents/src/runtime/loop/agent-runtime.ts @@ -6,8 +6,13 @@ import type { z } from 'zod'; import { incrementMessageCount, incrementTokenCountFromUsage } from './execution-counter'; import { GenerateSink } from './generate-sink'; import { hydrateFileParts } from './hydrate-file-parts'; -import type { ModelCallContext, RunOutputSink, RunServices } from './run-output-sink'; -import { RuntimeContextBuilder } from './runtime-context'; +import type { + ModelCallContext, + ModelTurnResult, + RunOutputSink, + RunServices, +} from './run-output-sink'; +import { RuntimeContextBuilder, type StaticLoopContext } from './runtime-context'; import { extractSettledToolCalls, formatMcpConnectionNote, @@ -75,6 +80,7 @@ import { BackgroundTaskTracker } from '../state/background-task-tracker'; import { AgentEventBus, type AgentAbortScope } from '../state/event-bus'; import { generateRunId, RunStateManager, StaleResumeError } from '../state/run-state'; import { startStreamSession } from '../streaming/stream-session'; +import type { StreamWriterGuard } from '../streaming/stream-writer-guard'; import { RuntimeTelemetry } from '../telemetry/runtime-telemetry'; import { DeferredToolManager } from '../tools/deferred-tool-manager'; import { fixToolCall } from '../tools/fix-tool-call'; @@ -178,6 +184,24 @@ interface LoopContext { pendingResume?: PendingResume; } +interface PreparedLoopContext extends LoopContext { + staticContext: StaticLoopContext; + instructionProviderOptions: ProviderOptions | undefined; + runTelemetry: BuiltTelemetry | undefined; + canDiscardRejectedInput: boolean; +} + +interface LoopState { + totalUsage: TokenUsage | undefined; + lastFinishReason: FinishReason; + structuredOutput: unknown; + maxIterations: number; + iterationCount: number; + reachedStopCondition: boolean; +} + +type ToolBatchSettlement = { suspended: false } | { suspended: true; result: T }; + /** * Core agent execution engine using the Vercel AI SDK directly. * @@ -391,43 +415,10 @@ export class AgentRuntime { options: ResumeOptions & ExecutionOptions, ): Promise { this.runId = options.runId; - const state = await this.runState.resume(this.runId); - if (!state) { - throw new StaleResumeError(`No suspended run found for runId: ${this.runId}`); - } - - const toolCall = state.pendingToolCalls[options.toolCallId]; - if (!toolCall) { - throw new StaleResumeError(`No tool call found for toolCallId: ${options.toolCallId}`); - } - if (options.hostMetadata !== undefined && !state.persistence) { - throw new Error('Cannot update host metadata without persistence'); - } - - const list = AgentMessageList.deserialize(state.messageList); - this.context.hydrateDeferredToolsFromList(list); - await hydrateFileParts(list.messages(), this.config.fileStore, { - threadId: state.persistence?.threadId, - }); - - const tool = this.context - .getCurrentTools(state.persistence) - .find((t) => t.name === toolCall.toolName); - if (!tool) throw new Error(`Tool ${toolCall.toolName} not found`); - - let resumeData: unknown = data; + const { state, list, resumeData } = await this.prepareToolResume(data, options); let abortScope: AgentAbortScope | undefined; let resumeClaimed = false; - const resumeSchema = toolCall.suspended ? toolCall.resumeSchema : tool.resumeSchema; - if (!isCancellation(resumeData) && resumeSchema) { - const parseResult = await parseWithSchema(resumeSchema, data, { stripUnknown: true }); - if (!parseResult.success) { - throw new Error(`Invalid resume payload: ${parseResult.error}`); - } - resumeData = parseResult.data; - } - try { // Merge persisted execution options with fresh caller options const { @@ -437,25 +428,7 @@ export class AgentRuntime { hostMetadata, ...callerExecOptions } = options; - const persisted = state.executionOptions ?? {}; - const persistedMaxIterations = persisted.maxIterations; - const callerMaxIterations = callerExecOptions.maxIterations; - if ( - callerMaxIterations !== undefined && - persistedMaxIterations !== undefined && - callerMaxIterations < persistedMaxIterations - ) { - throw new Error( - `Cannot decrease maxIterations when resuming a run. Expected >= ${persistedMaxIterations}, received ${callerMaxIterations}.`, - ); - } - - const mergedMaxIterations = callerMaxIterations ?? persistedMaxIterations; - const mergedExecOptions: ExecutionOptions & { iterationCount?: number } = { - ...callerExecOptions, - ...(mergedMaxIterations !== undefined ? { maxIterations: mergedMaxIterations } : {}), - ...(state.iterationCount !== undefined ? { iterationCount: state.iterationCount } : {}), - }; + const mergedExecOptions = this.mergeResumeExecutionOptions(state, callerExecOptions); const claimed = await this.runState.claimResume(this.runId, state); if (!claimed) { @@ -463,14 +436,7 @@ export class AgentRuntime { } resumeClaimed = true; const resumeOptions: RuntimeExecutionOptions = { - persistence: state.persistence - ? { - ...state.persistence, - ...(state.persistence.hostMetadata || hostMetadata - ? { hostMetadata: { ...state.persistence.hostMetadata, ...hostMetadata } } - : {}), - } - : undefined, + persistence: this.mergeResumePersistence(state.persistence, hostMetadata), ...mergedExecOptions, }; this.updateState({ persistence: resumeOptions.persistence }); @@ -485,68 +451,18 @@ export class AgentRuntime { resumeData, }; - await this.ensureModelCost(); + await this.prepareResumeMemory(list, state.persistence); - await this.memory.setListObservationLogMemory(list, state.persistence); - // The mask boundary is runtime-only state: re-derive it from the - // persisted cursor so a run that compacted mid-run before suspending - // does not resume with the full pre-compaction window. - await this.memory.applyObservationMask(list, state.persistence); - - if (method === 'generate') { - const sink = new GenerateSink(this.createRunServices()); - const rawResult = await this.telemetry.withRootSpan( - 'generate', - resumeOptions, - this.runId, - async () => - await this.runAgentLoop( - { - list, - options: resumeOptions, - abortScope: activeAbortScope, - pendingResume, - }, - sink, - ), - ); - try { - return this.finalizeGenerate(rawResult, list); - } finally { - abortScope.dispose(); - } - } - - return { - runId: this.runId, - stream: this.startStream({ - list, - options: resumeOptions, - abortScope: activeAbortScope, - pendingResume, - }), - getState: () => this.getState(), + const ctx: LoopContext = { + list, + options: resumeOptions, + abortScope: activeAbortScope, + pendingResume, }; + if (method === 'generate') return await this.generateResumedRun(ctx); + return this.createResumedStream(ctx); } catch (error) { - const isAbort = abortScope?.isAborted ?? false; - abortScope?.dispose(); - if (error instanceof StaleResumeError) throw error; - - this.updateState({ status: isAbort ? 'cancelled' : 'failed' }); - if (!isAbort) { - this.eventBus.emit({ type: AgentEvent.Error, message: String(error), error }); - } - if (resumeClaimed) await this.cleanupRun(); - if (method === 'generate') { - return { - runId: this.runId, - messages: [], - finishReason: 'error' as const, - error, - getState: () => this.getState(), - }; - } - return { runId: this.runId, stream: makeErrorStream(error), getState: () => this.getState() }; + return await this.handleResumeFailure(method, error, abortScope, resumeClaimed); } } @@ -564,6 +480,193 @@ export class AgentRuntime { options: { runId: string; contextNotes?: string[] } & ExecutionOptions, ): Promise { this.runId = options.runId; + const state = await this.loadStepCheckpoint(); + + const list = await this.restoreCheckpointMessages(state); + + let abortScope: AgentAbortScope | undefined; + try { + const { runId: _rid, contextNotes, ...callerExecOptions } = options; + const resumeOptions: RuntimeExecutionOptions = { + persistence: state.persistence, + ...this.mergeResumeExecutionOptions(state, callerExecOptions), + }; + + for (const note of contextNotes ?? []) { + list.addInput([{ role: 'user', content: [{ type: 'text', text: note }] }]); + } + + abortScope = this.eventBus.createAbortScope(resumeOptions.abortSignal); + const activeAbortScope = abortScope; + + await this.prepareResumeMemory(list, state.persistence); + + return this.createResumedStream({ + list, + options: resumeOptions, + abortScope: activeAbortScope, + }); + } catch (error) { + const isAbort = abortScope?.isAborted ?? false; + abortScope?.dispose(); + this.updateState({ status: isAbort ? 'cancelled' : 'failed' }); + if (!isAbort) { + this.eventBus.emit({ type: AgentEvent.Error, message: String(error), error }); + } + return { runId: this.runId, stream: makeErrorStream(error), getState: () => this.getState() }; + } + } + + // --- Private --- + + private async prepareToolResume(data: unknown, options: ResumeOptions & ExecutionOptions) { + const state = await this.runState.resume(this.runId); + if (!state) { + throw new StaleResumeError(`No suspended run found for runId: ${this.runId}`); + } + + const toolCall = state.pendingToolCalls[options.toolCallId]; + if (!toolCall) { + throw new StaleResumeError(`No tool call found for toolCallId: ${options.toolCallId}`); + } + if (options.hostMetadata !== undefined && !state.persistence) { + throw new Error('Cannot update host metadata without persistence'); + } + + const list = await this.restoreCheckpointMessages(state); + + const tool = this.context + .getCurrentTools(state.persistence) + .find((t) => t.name === toolCall.toolName); + if (!tool) throw new Error(`Tool ${toolCall.toolName} not found`); + + const resumeSchema = toolCall.suspended ? toolCall.resumeSchema : tool.resumeSchema; + const resumeData = await this.parseResumeData(data, resumeSchema); + return { state, list, resumeData }; + } + + private async handleResumeFailure( + method: 'generate' | 'stream', + error: unknown, + abortScope: AgentAbortScope | undefined, + resumeClaimed: boolean, + ): Promise { + const isAbort = abortScope?.isAborted ?? false; + abortScope?.dispose(); + if (error instanceof StaleResumeError) throw error; + + this.updateState({ status: isAbort ? 'cancelled' : 'failed' }); + if (!isAbort) { + this.eventBus.emit({ type: AgentEvent.Error, message: String(error), error }); + } + if (resumeClaimed) await this.cleanupRun(); + if (method === 'generate') { + return { + runId: this.runId, + messages: [], + finishReason: 'error' as const, + error, + getState: () => this.getState(), + }; + } + return { runId: this.runId, stream: makeErrorStream(error), getState: () => this.getState() }; + } + + private async restoreCheckpointMessages( + state: SerializableAgentState, + ): Promise { + const list = AgentMessageList.deserialize(state.messageList); + this.context.hydrateDeferredToolsFromList(list); + await hydrateFileParts(list.messages(), this.config.fileStore, { + threadId: state.persistence?.threadId, + }); + return list; + } + + private async parseResumeData( + data: unknown, + resumeSchema: BuiltTool['resumeSchema'], + ): Promise { + if (isCancellation(data) || !resumeSchema) return data; + const result = await parseWithSchema(resumeSchema, data, { stripUnknown: true }); + if (!result.success) throw new Error(`Invalid resume payload: ${result.error}`); + return result.data; + } + + private mergeResumeExecutionOptions( + state: SerializableAgentState, + callerExecOptions: ExecutionOptions, + ): ExecutionOptions & { iterationCount?: number } { + const persisted = state.executionOptions ?? {}; + const persistedMaxIterations = persisted.maxIterations; + const callerMaxIterations = callerExecOptions.maxIterations; + if ( + callerMaxIterations !== undefined && + persistedMaxIterations !== undefined && + callerMaxIterations < persistedMaxIterations + ) { + throw new Error( + `Cannot decrease maxIterations when resuming a run. Expected >= ${persistedMaxIterations}, received ${callerMaxIterations}.`, + ); + } + + const mergedMaxIterations = callerMaxIterations ?? persistedMaxIterations; + return { + ...callerExecOptions, + ...(mergedMaxIterations !== undefined ? { maxIterations: mergedMaxIterations } : {}), + ...(state.iterationCount !== undefined ? { iterationCount: state.iterationCount } : {}), + }; + } + + private mergeResumePersistence( + persistence: AgentPersistenceOptions | undefined, + hostMetadata: ResumeOptions['hostMetadata'], + ): AgentPersistenceOptions | undefined { + if (!persistence) return undefined; + const merged = { ...persistence }; + if (persistence.hostMetadata || hostMetadata) { + merged.hostMetadata = { ...persistence.hostMetadata, ...hostMetadata }; + } + return merged; + } + + private async prepareResumeMemory( + list: AgentMessageList, + persistence: AgentPersistenceOptions | undefined, + ): Promise { + await this.ensureModelCost(); + + await this.memory.setListObservationLogMemory(list, persistence); + // The mask boundary is runtime-only state: re-derive it from the + // persisted cursor so a run that compacted mid-run before suspending + // does not resume with the full pre-compaction window. + await this.memory.applyObservationMask(list, persistence); + } + + private async generateResumedRun(ctx: LoopContext): Promise { + const sink = new GenerateSink(this.createRunServices()); + const rawResult = await this.telemetry.withRootSpan( + 'generate', + ctx.options, + this.runId, + async () => await this.runAgentLoop(ctx, sink), + ); + try { + return this.finalizeGenerate(rawResult, ctx.list); + } finally { + ctx.abortScope.dispose(); + } + } + + private createResumedStream(ctx: LoopContext): StreamResult { + return { + runId: this.runId, + stream: this.startStream(ctx), + getState: () => this.getState(), + }; + } + + private async loadStepCheckpoint(): Promise { const state = await this.runState.loadForCrashResume(this.runId); if (!state) throw new Error(`No checkpoint found for runId: ${this.runId}`); if (state.status !== 'running') { @@ -579,69 +682,9 @@ export class AgentRuntime { `Checkpoint for runId ${this.runId} has pending tool calls — crashResume only accepts step checkpoints`, ); } - - const list = AgentMessageList.deserialize(state.messageList); - this.context.hydrateDeferredToolsFromList(list); - await hydrateFileParts(list.messages(), this.config.fileStore, { - threadId: state.persistence?.threadId, - }); - - let abortScope: AgentAbortScope | undefined; - try { - const { runId: _rid, contextNotes, ...callerExecOptions } = options; - const persisted = state.executionOptions ?? {}; - const persistedMaxIterations = persisted.maxIterations; - const callerMaxIterations = callerExecOptions.maxIterations; - if ( - callerMaxIterations !== undefined && - persistedMaxIterations !== undefined && - callerMaxIterations < persistedMaxIterations - ) { - throw new Error( - `Cannot decrease maxIterations when resuming a run. Expected >= ${persistedMaxIterations}, received ${callerMaxIterations}.`, - ); - } - const mergedMaxIterations = callerMaxIterations ?? persistedMaxIterations; - const resumeOptions: RuntimeExecutionOptions = { - persistence: state.persistence, - ...callerExecOptions, - ...(mergedMaxIterations !== undefined ? { maxIterations: mergedMaxIterations } : {}), - ...(state.iterationCount !== undefined ? { iterationCount: state.iterationCount } : {}), - }; - - for (const note of contextNotes ?? []) { - list.addInput([{ role: 'user', content: [{ type: 'text', text: note }] }]); - } - - abortScope = this.eventBus.createAbortScope(resumeOptions.abortSignal); - const activeAbortScope = abortScope; - - await this.ensureModelCost(); - await this.memory.setListObservationLogMemory(list, state.persistence); - await this.memory.applyObservationMask(list, state.persistence); - - return { - runId: this.runId, - stream: this.startStream({ - list, - options: resumeOptions, - abortScope: activeAbortScope, - }), - getState: () => this.getState(), - }; - } catch (error) { - const isAbort = abortScope?.isAborted ?? false; - abortScope?.dispose(); - this.updateState({ status: isAbort ? 'cancelled' : 'failed' }); - if (!isAbort) { - this.eventBus.emit({ type: AgentEvent.Error, message: String(error), error }); - } - return { runId: this.runId, stream: makeErrorStream(error), getState: () => this.getState() }; - } + return state; } - // --- Private --- - /** * Build an AgentMessageList for the current turn: * - Loads memory history (if configured) and adds as historySet @@ -793,36 +836,79 @@ export class AgentRuntime { * those throws into their terminal contract. */ private async runAgentLoop(ctx: LoopContext, sink: RunOutputSink): Promise { - const { list, options, abortScope, pendingResume } = ctx; + const { prepared, state } = await this.prepareLoop(ctx); + if (prepared.pendingResume) { + const settlement = await this.resumePendingTools( + prepared, + sink, + state, + prepared.pendingResume, + ); + if (settlement.suspended) return settlement.result; + } + + for (; state.iterationCount < state.maxIterations; state.iterationCount++) { + const settlement = await this.runLoopIteration(prepared, sink, state); + if (settlement.suspended) return settlement.result; + if (state.reachedStopCondition) break; + } + + if (!state.reachedStopCondition && state.iterationCount >= state.maxIterations) { + state.lastFinishReason = 'max-iterations'; + } + return await sink.finishComplete({ + list: ctx.list, + options: ctx.options, + finishReason: state.lastFinishReason, + usage: state.totalUsage, + structuredOutput: state.structuredOutput, + }); + } + + private async prepareLoop(ctx: LoopContext): Promise<{ + prepared: PreparedLoopContext; + state: LoopState; + }> { + const { list, options } = ctx; await this.activeSkills?.restore(list, options?.persistence); this.context.hydrateDeferredToolsFromList(list); - // Inject a model-facing note for any MCP servers that failed to connect - // during build(). The agent can mention the outage to the user when - // relevant; the note is system-message only and never persisted. + // This note reaches the model but is not stored in conversation history. list.mcpConnectionNote = formatMcpConnectionNote(this.config.mcpConnectionFailures ?? []); - let totalUsage: TokenUsage | undefined; - let lastFinishReason: FinishReason = 'stop'; - let structuredOutput: unknown; const runTelemetry = this.telemetry.resolve(options); - const staticLoopContext = this.context.buildStaticLoopContext({ + const staticContext = this.context.buildStaticLoopContext({ ...options, persistence: options?.persistence, }); - // Anthropic cache breakpoints are Anthropic-only and don't change across - // iterations, so compute once. Explicit .instructions() providerOptions win. + // Explicit instruction options take precedence over cache defaults. const instructionProviderOptions = mergeProviderOptions( buildInstructionPromptCacheOptions(this.config.promptCaching, this.modelIdString), this.config.instructionProviderOptions, ); - const maxIterations = options?.maxIterations ?? MAX_LOOP_ITERATIONS; - let iterationCount = options?.iterationCount ?? 0; - let reachedStopCondition = false; - const inputMessages = new Set(list.inputDelta()); - const inputIds = new Set([...inputMessages].map((message) => message.id)); + const state: LoopState = { + totalUsage: undefined, + lastFinishReason: 'stop', + structuredOutput: undefined, + maxIterations: options?.maxIterations ?? MAX_LOOP_ITERATIONS, + iterationCount: options?.iterationCount ?? 0, + reachedStopCondition: false, + }; + return { + prepared: { + ...ctx, + runTelemetry, + staticContext, + instructionProviderOptions, + canDiscardRejectedInput: this.canDiscardRejectedInput(ctx), + }, + state, + }; + } - // Can we discard an input in case of an error caused by its attachment - const canDiscardRejectedInput = + private canDiscardRejectedInput(ctx: LoopContext): boolean { + const inputMessages = new Set(ctx.list.inputDelta()); + const inputIds = new Set([...inputMessages].map((message) => message.id)); + return ( ctx.isFreshRun === true && [...inputMessages].some( (message) => @@ -831,237 +917,274 @@ export class AgentRuntime { Array.isArray(message.content) && message.content.some((part) => part.type === 'file' && part.data !== undefined), ) && - !list.messages().some((message) => !inputMessages.has(message) && inputIds.has(message.id)); + !ctx.list + .messages() + .some((message) => !inputMessages.has(message) && inputIds.has(message.id)) + ); + } - const buildToolBatchContext = (toolMap: Map): ToolBatchContext => ({ + private buildToolBatchContext( + ctx: PreparedLoopContext, + toolMap: Map, + ): ToolBatchContext { + return { toolMap, - list, + list: ctx.list, runId: this.runId, - persistence: options?.persistence, - telemetry: runTelemetry, - executionCounter: options?.executionCounter, - abortSignal: abortScope.signal, - isAborted: () => abortScope.isAborted, - }); - const finishToolBatch = async ( - batch: ToolCallBatchResult, - toolMap: Map, - nextIteration: number, - ) => { - const hasPending = Object.keys(batch.pending).length > 0; - let completed = false; - try { - this.assertNotAborted(abortScope); - await sink.emitToolBatch(batch); - this.assertNotAborted(abortScope); - if (!hasPending) { - completed = true; - return { suspended: false as const }; - } - - await this.persistSuspension( - batch.pending, - options, - list, - totalUsage, - maxIterations, - nextIteration, - ); - this.assertNotAborted(abortScope); - const result = await sink.finishSuspended({ - suspendRunId: this.runId, - list, - usage: totalUsage, - suspensions: batch.suspensions, - }); - this.assertNotAborted(abortScope); - completed = true; - return { suspended: true as const, result }; - } finally { - if (!completed && hasPending) { - await this.toolExecutor.cleanupPendingToolCalls( - batch.pending, - buildToolBatchContext(toolMap), - abortScope.isAborted ? 'Run aborted' : 'Parent run failed before suspension', - ); - try { - await this.runState.cancel(this.runId, this.getState()); - } catch { - // Preserve the failure that interrupted suspension finalization. - } - } - } + persistence: ctx.options?.persistence, + telemetry: ctx.runTelemetry, + executionCounter: ctx.options?.executionCounter, + abortSignal: ctx.abortScope.signal, + isAborted: () => ctx.abortScope.isAborted, }; + } - if (pendingResume) { - const pendingLoopContext = this.context.buildToolLoopContext( - staticLoopContext.aiProviderTools, - options?.persistence, - options?.executionCounter, - list, - ); - const batch = await this.toolExecutor.iteratePendingToolCallsConcurrent({ - ...buildToolBatchContext(pendingLoopContext.toolMap), - pendingResume, - }); - const finalized = await finishToolBatch(batch, pendingLoopContext.toolMap, iterationCount); - if (finalized.suspended) return finalized.result; - // The resumed batch is a clean boundary too: its tool results are new - // content no earlier boundary saw, so check before the next model call. - await this.memory.maybeObserveMidRun(list, options); - } - - for (; iterationCount < maxIterations; iterationCount++) { - this.assertNotAborted(abortScope); - - this.eventBus.emit({ type: AgentEvent.TurnStart }); - - for (const toolName of this.activeSkills?.toolDependencies() ?? []) { - this.deferredToolManager?.load(toolName); - } - - const { - toolMap, - aiTools, - hasTools, - effectiveInstructions, - volatileInstructions, - staticToolCacheName, - } = this.context.buildToolLoopContext( - staticLoopContext.aiProviderTools, - options?.persistence, - options?.executionCounter, - list, - ); - const hostVolatileInstructions = await this.resolveVolatileInstructions(options?.persistence); - const combinedVolatileInstructions = [volatileInstructions, hostVolatileInstructions] - .map((value) => value?.trim()) - .filter((value): value is string => Boolean(value)) - .join('\n\n'); - const { system, messages } = list.forLlm( - // Skill content changes only on activation. Keep it cached when memory compacts. - [effectiveInstructions, this.activeSkills?.instructions()] - .filter(Boolean) - .join('\n\n'), - instructionProviderOptions, - combinedVolatileInstructions || undefined, - supportsSplitSystemMessages(this.config.model), - ); - // Runtime breakpoints (conversation history, static tools) are per-call - // only — never persisted back to the message list or tool set. - const cached = applyRuntimeCacheBreakpoints({ - system, - messages: this.activeSkills?.modelMessages(messages, list) ?? messages, - aiTools, - promptCaching: this.config.promptCaching, - modelId: this.modelIdString, - staticToolCacheName, - }); - - const modelCallContext: ModelCallContext = { - model: staticLoopContext.model, - system, - messages: cached.messages, - abortSignal: abortScope.signal, - hasTools, - aiTools: cached.aiTools, - reasoning: staticLoopContext.reasoning, - providerOptions: staticLoopContext.providerOptions, - outputSpec: staticLoopContext.outputSpec, - maxOutputTokens: staticLoopContext.maxOutputTokens, - aiSdkOptions: this.buildAiSdkOptions(toolMap, options), - onInputRejected: - canDiscardRejectedInput && iterationCount === 0 - ? async () => { - if (abortScope.isAborted) return; - await this.memory.discardRejectedInput(list, options); - this.updateState({ messageList: list.serialize() }); - } - : undefined, - }; - let turn = await sink.callModel(modelCallContext); - - // Some providers occasionally return a `stop` turn with no output at - // all mid-task, which would silently end the run with work half-done. - // Retry the call a bounded number of times before accepting the empty - // turn; each discarded attempt still bills its usage. - for ( - let emptyRetry = 0; - emptyRetry < MAX_EMPTY_TURN_RETRIES && isEmptyModelTurn(turn); - emptyRetry++ - ) { - totalUsage = mergeUsage(totalUsage, turn.usage); - incrementTokenCountFromUsage(options?.executionCounter, turn.usage); - // Publish before the abort check so a cancel between the empty attempt - // and the retry still bills those tokens via getTerminalFinish(). - sink.reportUsage(totalUsage); - this.assertNotAborted(abortScope); - turn = await sink.callModel(modelCallContext); - } - - // Fold the just-finished turn's usage in before the abort check so a - // stop that lands right after the model call still bills its tokens. - totalUsage = mergeUsage(totalUsage, turn.usage); - incrementTokenCountFromUsage(options?.executionCounter, turn.usage); - sink.reportUsage(totalUsage); - - this.assertNotAborted(abortScope); - - lastFinishReason = turn.finishReason; - if (!isReasoningOnlyStop(turn)) list.addResponse(turn.newMessages); - // The turn is now in the list; drop any retained streamed text so a later - // abort's snapshot can't duplicate it (a stop before this point recovers it). - sink.onTurnFolded?.(); - - if (turn.aiFinishReason !== 'tool-calls') { - // A rejected/filtered request (e.g. a provider prompt safety block) - // surfaces as an output-less turn instead of an SDK error — throw so - // the failure reaches the caller rather than ending the run silently. - if (turn.errorReason) throw new Error(turn.errorReason.message); - structuredOutput = turn.structuredOutput; - this.emitTurnEnd(turn.newMessages, extractSettledToolCalls(turn.newMessages)); - reachedStopCondition = true; - break; - } - - const batch = await this.toolExecutor.iterateToolCallsConcurrent({ - ...buildToolBatchContext(toolMap), - toolCalls: turn.toolCalls, - }); - const finalized = await finishToolBatch(batch, toolMap, iterationCount + 1); - if (finalized.suspended) return finalized.result; - - // Emit TurnEnd after all tool calls in this iteration are processed - this.emitTurnEnd(turn.newMessages, extractSettledToolCalls(list.responseDelta())); - - // Clean loop boundary: all tool calls settled. Mid-run observation - // may compact the LLM window here before the next call. - await this.memory.maybeObserveMidRun(list, options); - - // Step boundary reached with nothing pending: durably checkpoint so a - // crash before the next model call loses only the in-flight step. - if (options?.stepCheckpoints) { - await this.persistStepCheckpoint( - list, - totalUsage, - options, - maxIterations, - iterationCount + 1, - ); - } - } - - if (!reachedStopCondition && iterationCount >= maxIterations) { - lastFinishReason = 'max-iterations'; - } - - return await sink.finishComplete({ - list, - options, - finishReason: lastFinishReason, - usage: totalUsage, - structuredOutput, + private async resumePendingTools( + ctx: PreparedLoopContext, + sink: RunOutputSink, + state: LoopState, + pendingResume: PendingResume, + ): Promise> { + const { toolMap } = this.context.buildToolLoopContext( + ctx.staticContext.aiProviderTools, + ctx.options?.persistence, + ctx.options?.executionCounter, + ctx.list, + ); + const batch = await this.toolExecutor.iteratePendingToolCallsConcurrent({ + ...this.buildToolBatchContext(ctx, toolMap), + pendingResume, }); + const settlement = await this.finishToolBatch( + ctx, + sink, + state, + batch, + toolMap, + state.iterationCount, + ); + if (settlement.suspended) return settlement; + // Resumed tool results form a new observation boundary before the next model call. + await this.memory.maybeObserveMidRun(ctx.list, ctx.options); + return settlement; + } + + private async runLoopIteration( + ctx: PreparedLoopContext, + sink: RunOutputSink, + state: LoopState, + ): Promise> { + this.assertNotAborted(ctx.abortScope); + this.eventBus.emit({ type: AgentEvent.TurnStart }); + const { toolMap, modelCallContext } = await this.prepareModelCall(ctx, state.iterationCount); + const turn = await this.callModelWithRetries(ctx, sink, state, modelCallContext); + this.assertNotAborted(ctx.abortScope); + + state.lastFinishReason = turn.finishReason; + if (!isReasoningOnlyStop(turn)) ctx.list.addResponse(turn.newMessages); + // Drop retained stream text only after the response is in the list. + sink.onTurnFolded?.(); + + if (turn.aiFinishReason !== 'tool-calls') { + if (turn.errorReason) throw new Error(turn.errorReason.message); + state.structuredOutput = turn.structuredOutput; + this.emitTurnEnd(turn.newMessages, extractSettledToolCalls(turn.newMessages)); + state.reachedStopCondition = true; + return { suspended: false }; + } + + const batch = await this.toolExecutor.iterateToolCallsConcurrent({ + ...this.buildToolBatchContext(ctx, toolMap), + toolCalls: turn.toolCalls, + }); + const settlement = await this.finishToolBatch( + ctx, + sink, + state, + batch, + toolMap, + state.iterationCount + 1, + ); + if (settlement.suspended) return settlement; + await this.completeToolTurn(ctx, state, turn); + return settlement; + } + + private async prepareModelCall(ctx: PreparedLoopContext, iterationCount: number) { + const { list, options, abortScope, staticContext } = ctx; + for (const toolName of this.activeSkills?.toolDependencies() ?? []) { + this.deferredToolManager?.load(toolName); + } + const tools = this.context.buildToolLoopContext( + staticContext.aiProviderTools, + options?.persistence, + options?.executionCounter, + list, + ); + const hostVolatileInstructions = await this.resolveVolatileInstructions(options?.persistence); + const { system, messages } = this.buildModelPrompt(ctx, tools, hostVolatileInstructions); + // Cache breakpoints apply to this call only. Do not change stored messages or tools. + const cached = applyRuntimeCacheBreakpoints({ + system, + messages: this.activeSkills?.modelMessages(messages, list) ?? messages, + aiTools: tools.aiTools, + promptCaching: this.config.promptCaching, + modelId: this.modelIdString, + staticToolCacheName: tools.staticToolCacheName, + }); + const modelCallContext: ModelCallContext = { + model: staticContext.model, + system, + messages: cached.messages, + abortSignal: abortScope.signal, + hasTools: tools.hasTools, + aiTools: cached.aiTools, + reasoning: staticContext.reasoning, + providerOptions: staticContext.providerOptions, + outputSpec: staticContext.outputSpec, + maxOutputTokens: staticContext.maxOutputTokens, + aiSdkOptions: this.buildAiSdkOptions(tools.toolMap, options), + onInputRejected: this.createInputRejectionHandler(ctx, iterationCount), + }; + return { toolMap: tools.toolMap, modelCallContext }; + } + + private buildModelPrompt( + ctx: PreparedLoopContext, + tools: ReturnType, + hostVolatileInstructions: string | undefined, + ) { + const combinedVolatileInstructions = [tools.volatileInstructions, hostVolatileInstructions] + .map((value) => value?.trim()) + .filter((value): value is string => Boolean(value)) + .join('\n\n'); + return ctx.list.forLlm( + // Skill content changes only on activation. Keep it cached when memory compacts. + [tools.effectiveInstructions, this.activeSkills?.instructions()] + .filter(Boolean) + .join('\n\n'), + ctx.instructionProviderOptions, + combinedVolatileInstructions || undefined, + supportsSplitSystemMessages(this.config.model), + ); + } + + private createInputRejectionHandler( + ctx: PreparedLoopContext, + iterationCount: number, + ): ModelCallContext['onInputRejected'] { + if (!ctx.canDiscardRejectedInput || iterationCount !== 0) return undefined; + return async () => { + if (ctx.abortScope.isAborted) return; + await this.memory.discardRejectedInput(ctx.list, ctx.options); + this.updateState({ messageList: ctx.list.serialize() }); + }; + } + + private async callModelWithRetries( + ctx: PreparedLoopContext, + sink: RunOutputSink, + state: LoopState, + modelCallContext: ModelCallContext, + ): Promise { + let turn = await sink.callModel(modelCallContext); + // Retry empty stop turns. Charge each attempt before checking for cancellation. + for (let retry = 0; retry < MAX_EMPTY_TURN_RETRIES && isEmptyModelTurn(turn); retry++) { + this.recordTurnUsage(ctx, sink, state, turn); + this.assertNotAborted(ctx.abortScope); + turn = await sink.callModel(modelCallContext); + } + this.recordTurnUsage(ctx, sink, state, turn); + return turn; + } + + private recordTurnUsage( + ctx: PreparedLoopContext, + sink: RunOutputSink, + state: LoopState, + turn: ModelTurnResult, + ): void { + state.totalUsage = mergeUsage(state.totalUsage, turn.usage); + incrementTokenCountFromUsage(ctx.options?.executionCounter, turn.usage); + sink.reportUsage(state.totalUsage); + } + + private async finishToolBatch( + ctx: PreparedLoopContext, + sink: RunOutputSink, + state: LoopState, + batch: ToolCallBatchResult, + toolMap: Map, + nextIteration: number, + ): Promise> { + const { list, options, abortScope } = ctx; + const hasPending = Object.keys(batch.pending).length > 0; + let completed = false; + try { + this.assertNotAborted(abortScope); + await sink.emitToolBatch(batch); + this.assertNotAborted(abortScope); + if (!hasPending) { + completed = true; + return { suspended: false }; + } + await this.persistSuspension( + batch.pending, + options, + list, + state.totalUsage, + state.maxIterations, + nextIteration, + ); + this.assertNotAborted(abortScope); + const result = await sink.finishSuspended({ + suspendRunId: this.runId, + list, + usage: state.totalUsage, + suspensions: batch.suspensions, + }); + this.assertNotAborted(abortScope); + completed = true; + return { suspended: true, result }; + } finally { + if (!completed && hasPending) await this.cleanupFailedSuspension(ctx, batch.pending, toolMap); + } + } + + private async cleanupFailedSuspension( + ctx: PreparedLoopContext, + pending: Record, + toolMap: Map, + ): Promise { + await this.toolExecutor.cleanupPendingToolCalls( + pending, + this.buildToolBatchContext(ctx, toolMap), + ctx.abortScope.isAborted ? 'Run aborted' : 'Parent run failed before suspension', + ); + try { + await this.runState.cancel(this.runId, this.getState()); + } catch { + // Preserve the failure that interrupted suspension finalization. + } + } + + private async completeToolTurn( + ctx: PreparedLoopContext, + state: LoopState, + turn: ModelTurnResult, + ): Promise { + this.emitTurnEnd(turn.newMessages, extractSettledToolCalls(ctx.list.responseDelta())); + // All tools have settled. Observe before the next call and its checkpoint. + await this.memory.maybeObserveMidRun(ctx.list, ctx.options); + if (ctx.options?.stepCheckpoints) { + await this.persistStepCheckpoint( + ctx.list, + state.totalUsage, + ctx.options, + state.maxIterations, + state.iterationCount + 1, + ); + } } private async resolveVolatileInstructions( @@ -1104,18 +1227,7 @@ export class AgentRuntime { withRootSpan: async (operation, options, runId, fn) => await this.telemetry.withRootSpan(operation, options, runId, fn), runLoop: async (guard) => { - // Surface MCP connection failures as non-fatal warnings before the - // first LLM step. Tools from these servers were skipped during - // build(); the run continues with the remaining tools. - for (const failure of this.config.mcpConnectionFailures ?? []) { - void guard.write({ - type: 'warning', - message: failure.error, - code: 'mcp_connection_failed', - source: 'mcp', - server: failure.server, - }); - } + this.writeMcpConnectionWarnings(guard); const resolvedList = ctx.list ?? (await this.initRun(ctx.input, ctx.options)); list = resolvedList; sink = new StreamSink(guard, this.createRunServices(), ctx.options); @@ -1150,6 +1262,21 @@ export class AgentRuntime { }); } + private writeMcpConnectionWarnings(guard: StreamWriterGuard): void { + // Surface MCP connection failures as non-fatal warnings before the + // first LLM step. Tools from these servers were skipped during + // build(); the run continues with the remaining tools. + for (const failure of this.config.mcpConnectionFailures ?? []) { + void guard.write({ + type: 'warning', + message: failure.error, + code: 'mcp_connection_failed', + source: 'mcp', + server: failure.server, + }); + } + } + /** * Persist a suspended run state and update the current state snapshot, and durably * save the turn-so-far to thread memory so a suspended turn that is later cancelled or @@ -1170,6 +1297,26 @@ export class AgentRuntime { const executionOptions: PersistedExecutionOptions | undefined = resolvedMaxIterations !== undefined ? { maxIterations: resolvedMaxIterations } : undefined; + this.markSuspendedToolCalls(list, pendingToolCalls); + + const state: SerializableAgentState = { + persistence: options?.persistence, + status: 'suspended', + messageList: list.serialize(), + pendingToolCalls, + usage: totalUsage, + executionOptions, + ...(resolvedIterationCount !== undefined ? { iterationCount: resolvedIterationCount } : {}), + }; + await this.runState.suspend(this.runId, state); + this.updateState({ status: 'suspended', pendingToolCalls, messageList: list.serialize() }); + await this.memory.persistTurnDelta(list, options); + } + + private markSuspendedToolCalls( + list: AgentMessageList, + pendingToolCalls: Record, + ): void { // Record what confirmation each suspended call showed the user, so an // abandoned suspension can be settled with that context on a later // history load instead of vanishing from the transcript. @@ -1184,19 +1331,6 @@ export class AgentRuntime { ...(typeof payload?.requestId === 'string' ? { requestId: payload.requestId } : {}), }); } - - const state: SerializableAgentState = { - persistence: options?.persistence, - status: 'suspended', - messageList: list.serialize(), - pendingToolCalls, - usage: totalUsage, - executionOptions, - ...(resolvedIterationCount !== undefined ? { iterationCount: resolvedIterationCount } : {}), - }; - await this.runState.suspend(this.runId, state); - this.updateState({ status: 'suspended', pendingToolCalls, messageList: list.serialize() }); - await this.memory.persistTurnDelta(list, options); } /** diff --git a/packages/@n8n/agents/src/runtime/loop/runtime-context.ts b/packages/@n8n/agents/src/runtime/loop/runtime-context.ts index c898fe0ed18..29494e097d6 100644 --- a/packages/@n8n/agents/src/runtime/loop/runtime-context.ts +++ b/packages/@n8n/agents/src/runtime/loop/runtime-context.ts @@ -39,6 +39,13 @@ function wrapBuiltInRules(fragments: string[]): string | undefined { return `\n${fragments.map((f) => `- ${f}`).join('\n')}\n`; } +function prependBuiltInRules(fragments: string[], instructions: string): string { + const rules = wrapBuiltInRules(fragments); + if (!rules) return instructions; + if (!instructions) return rules; + return `${rules}\n\n${instructions}`; +} + export interface StaticLoopContext { model: LanguageModel; aiProviderTools: ReturnType; @@ -68,7 +75,7 @@ export class RuntimeContextBuilder { buildStaticLoopContext( execOptions?: ExecutionOptions & { persistence?: AgentPersistenceOptions }, ): StaticLoopContext { - const { Output, jsonSchema } = loadAi(); + const ai = loadAi(); const aiProviderTools = toAiSdkProviderTools(this.config.providerTools); const model = createModel(this.config.model, this.config.modelFetch); const outputSchema = this.config.structuredOutput; @@ -78,16 +85,7 @@ export class RuntimeContextBuilder { isRawJsonSchemaOutput, ); - const outputSpec = outputSchema - ? Output.object({ - // Zod schemas pass through directly; a raw JSON Schema gets - // `additionalProperties: false` locked onto every object (required by Anthropic) - // and wrapped with the AI SDK's `jsonSchema()` helper. - schema: isZodSchema(outputSchema) - ? outputSchema - : jsonSchema(lockAdditionalProperties(outputSchema)), - }) - : undefined; + const outputSpec = this.buildOutputSpec(outputSchema, ai); return { model, @@ -99,6 +97,18 @@ export class RuntimeContextBuilder { }; } + private buildOutputSpec( + outputSchema: AgentRuntimeConfig['structuredOutput'], + { Output, jsonSchema }: ReturnType, + ): StaticLoopContext['outputSpec'] { + if (!outputSchema) return undefined; + // Anthropic requires additionalProperties: false on each raw schema object. + const schema = isZodSchema(outputSchema) + ? outputSchema + : jsonSchema(lockAdditionalProperties(outputSchema)); + return Output.object({ schema }); + } + /** Build the current local tool view; deferred loads can change this between iterations. */ buildToolLoopContext( aiProviderTools: ReturnType, @@ -260,6 +270,14 @@ export class RuntimeContextBuilder { instructions: string; volatileInstructions: string | undefined; } { + const { stableFragments, volatileFragments } = this.collectInstructionFragments(tools); + return { + instructions: prependBuiltInRules(stableFragments, this.config.instructions), + volatileInstructions: wrapBuiltInRules(volatileFragments), + }; + } + + private collectInstructionFragments(tools: BuiltTool[]) { const loadedToolNames = new Set( this.deferredToolManager?.getLoadedTools().map((t) => t.name) ?? [], ); @@ -272,9 +290,8 @@ export class RuntimeContextBuilder { ) { continue; } - (loadedToolNames.has(tool.name) ? volatileFragments : stableFragments).push( - tool.systemInstruction, - ); + const fragments = loadedToolNames.has(tool.name) ? volatileFragments : stableFragments; + fragments.push(tool.systemInstruction); } // Define the untrusted-data boundary ahead of the first wrapped result. @@ -288,18 +305,7 @@ export class RuntimeContextBuilder { target.unshift(UNTRUSTED_OUTPUT_DOCTRINE); } - const userInstructions = this.config.instructions; - const stableBlock = wrapBuiltInRules(stableFragments); - const instructions = stableBlock - ? userInstructions - ? `${stableBlock}\n\n${userInstructions}` - : stableBlock - : userInstructions; - - return { - instructions, - volatileInstructions: wrapBuiltInRules(volatileFragments), - }; + return { stableFragments, volatileFragments }; } /** Build the providerOptions object for provider-specific thinking config. */ diff --git a/packages/@n8n/agents/src/runtime/loop/stream-sink.ts b/packages/@n8n/agents/src/runtime/loop/stream-sink.ts index 043c72c9664..5ce5fd8bde1 100644 --- a/packages/@n8n/agents/src/runtime/loop/stream-sink.ts +++ b/packages/@n8n/agents/src/runtime/loop/stream-sink.ts @@ -1,4 +1,4 @@ -import type { StreamTextTransform, ToolSet } from 'ai'; +import type { StreamTextTransform, TextStreamPart, ToolSet } from 'ai'; import type { CompleteEmission, @@ -160,7 +160,7 @@ export class StreamSink implements RunOutputSink { return transforms.length > 0 ? { experimental_transform: transforms } : {}; } - async callModel(ctx: ModelCallContext): Promise { + private getStreamDeadlines(): { idleMs: number; firstOutputMs: number } { // Clamped to MAX_MODEL_STREAM_TIMEOUT_MS: above it setTimeout fires after // 1ms, turning an "effectively disable" override into instant stalls. const idleMs = Math.min( @@ -177,48 +177,55 @@ export class StreamSink implements RunOutputSink { ), MAX_MODEL_STREAM_TIMEOUT_MS, ); + return { idleMs, firstOutputMs }; + } + + async callModel(ctx: ModelCallContext): Promise { + const deadlines = this.getStreamDeadlines(); for (let attempt = 0; ; attempt++) { - // Opt-in: only build the reader when the host bills stopped runs, and only - // for providers with a reader. Re-created per attempt so a retried turn - // starts from a clean raw capture. - this.rawUsageReader = this.options?.recoverUsageOnAbort - ? createRawUsageReader(this.services.modelId) - : undefined; - // Some providers report failures (e.g. prompt safety blocks) only on the - // raw stream — no error, no content — so raw chunks are required to - // explain an otherwise silent empty response. - this.rawErrorReader = createRawErrorReader(this.services.modelId); + this.resetRawReaders(); // Per-attempt controller: a stall must be able to cancel this attempt's // fetch (releasing the socket the 1h network timeout would otherwise // hold) without touching the run-level signal. const turnAbort = new AbortController(); const attemptState = { streamedContent: false }; try { - return await this.streamModelTurn(ctx, turnAbort, { idleMs, firstOutputMs }, attemptState); + return await this.streamModelTurn(ctx, turnAbort, deadlines, attemptState); } catch (error) { if (isAttachmentValidationError(error)) { if (!attemptState.streamedContent) await ctx.onInputRejected?.(error); throw error; } - // A stalled or empty stream before any content is invisible to the user - // (and to the host's persistence) — re-issue the request instead of - // failing the run for what is usually a dead connection at request time. - // Once content has streamed, a retry would duplicate segments and orphan - // already-persisted tool-call facts, so the error surfaces instead. - // The abandoned attempt's prompt processing goes unbilled — accepted: - // it only has usage at all when the stream died between message_start - // and the first content chunk. - const retryable = - (error instanceof ModelStreamStallError || - loadAi().NoOutputGeneratedError.isInstance(error)) && - attempt < MAX_MODEL_STREAM_STALL_RETRIES && - !attemptState.streamedContent && - !ctx.abortSignal.aborted; - if (!retryable) throw error; + if (!this.canRetryStream(error, attempt, attemptState, ctx.abortSignal)) throw error; } } } + private resetRawReaders(): void { + // Each attempt needs fresh readers so a retry does not reuse captured usage. + this.rawUsageReader = this.options?.recoverUsageOnAbort + ? createRawUsageReader(this.services.modelId) + : undefined; + // Some providers report a blocked prompt only in raw chunks. + this.rawErrorReader = createRawErrorReader(this.services.modelId); + } + + private canRetryStream( + error: unknown, + attempt: number, + attemptState: { streamedContent: boolean }, + abortSignal: AbortSignal, + ): boolean { + // Retry only before content is visible or persisted. The abandoned attempt stays unbilled. + return ( + (error instanceof ModelStreamStallError || + loadAi().NoOutputGeneratedError.isInstance(error)) && + attempt < MAX_MODEL_STREAM_STALL_RETRIES && + !attemptState.streamedContent && + !abortSignal.aborted + ); + } + private async streamModelTurn( ctx: ModelCallContext, turnAbort: AbortController, @@ -227,6 +234,36 @@ export class StreamSink implements RunOutputSink { ): Promise { const { idleMs, firstOutputMs } = deadlines; const { NoOutputGeneratedError, streamText } = loadAi(); + const { result, activity } = this.startModelStream(ctx, turnAbort, idleMs, streamText); + + // Use the longer deadline before the first output to allow prompt processing. + const chunkStream = + idleMs > 0 + ? withChunkIdleTimeout( + result.stream, + () => (attemptState.streamedContent ? idleMs : firstOutputMs), + () => turnAbort.abort(), + () => activity.lastAt, + ) + : result.stream; + + for await (const chunk of chunkStream) { + if (chunk.type === 'error' && isAttachmentValidationError(chunk.error)) throw chunk.error; + // Result promises report this error. Do not forward it or count it as content before a retry. + if (chunk.type === 'error' && NoOutputGeneratedError.isInstance(chunk.error)) continue; + if (!this.captureStreamChunk(chunk, attemptState)) continue; + await this.writeStreamChunk(chunk); + } + + return await this.resolveModelTurn(result, ctx, idleMs, turnAbort); + } + + private startModelStream( + ctx: ModelCallContext, + turnAbort: AbortController, + idleMs: number, + streamText: ReturnType['streamText'], + ) { // Raw chunks serve two consumers: the usage/error readers parse them, and // the stall watchdog counts them as liveness — keepalive events the SDK // otherwise drops (e.g. Anthropic `ping`, a gateway's empty-delta frames) @@ -264,83 +301,58 @@ export class StreamSink implements RunOutputSink { : undefined, ), }); + return { result, activity }; + } - // A healthy streaming response emits chunks continuously (raw provider - // events included), so prolonged silence means the connection or provider - // is wedged — fail the turn instead of hanging until the network timeout. - // Pre-first-output silence gets the longer deadline (prompt processing on - // large cache-miss prompts sends nothing for minutes); once content has - // streamed, the tighter idle limit applies. - const chunkStream = - idleMs > 0 - ? withChunkIdleTimeout( - result.stream, - () => (attemptState.streamedContent ? idleMs : firstOutputMs), - () => turnAbort.abort(), - () => activity.lastAt, - ) - : result.stream; + private captureStreamChunk( + chunk: TextStreamPart, + attemptState: { streamedContent: boolean }, + ): boolean { + // Capture raw chunks here when a stream does not use the tap, such as a test stream. + if (chunk.type === 'raw') { + this.rawUsageReader?.capture(chunk.rawValue); + this.rawErrorReader?.capture(chunk.rawValue); + return false; + } + if (!STALL_RETRY_SAFE_CHUNK_TYPES.has(chunk.type)) attemptState.streamedContent = true; + // The runtime writes its own finish chunk after the loop completes. + if (chunk.type === 'finish') return false; + // Retain partial text so cancellation can save an unfinished response. + if (chunk.type === 'text-delta') this.partialText += chunk.text ?? ''; + return true; + } - // Consume the stream. When the AbortSignal fires mid-stream the AI SDK - // cancels the underlying fetch and the async iterator throws; the error - // propagates to the StreamSession which closes the consumer stream. - for await (const chunk of chunkStream) { - if (chunk.type === 'error' && isAttachmentValidationError(chunk.error)) throw chunk.error; - // The rejected result promises surface this error after the stream - // closes. Forwarding it here would mark a successful retry as failed. - // Runs before the `streamedContent` check below so a skipped chunk - // cannot make a stalled attempt look unretryable. - if (chunk.type === 'error' && NoOutputGeneratedError.isInstance(chunk.error)) continue; - // Raw chunks are normally consumed by the raw-chunk tap inside the SDK - // pipeline and never reach this loop; this branch is the fallback for - // streams the tap was not applied to (e.g. mocked streams in tests) and - // is why 'raw' stays in the retry-safe set. Capturing usage here also - // lets an aborted turn still be billed via getTerminalFinish. - if (chunk.type === 'raw') { - this.rawUsageReader?.capture(chunk.rawValue); - this.rawErrorReader?.capture(chunk.rawValue); - continue; - } - // Anything beyond transport bookkeeping counts as content: once seen, - // a stalled attempt is no longer silently retryable (see callModel). - if (!STALL_RETRY_SAFE_CHUNK_TYPES.has(chunk.type)) attemptState.streamedContent = true; - // Filter only the SDK's terminal `finish` chunk — the runtime emits its - // own consolidated `finish` after the loop completes. `start-step` / - // `finish-step` are passed through as LLM-iteration boundaries. - if (chunk.type === 'finish') continue; - - // Accumulate streamed text so an abort mid-response can persist it (the - // turn's `newMessages` are only built below, which the abort never reaches). - if (chunk.type === 'text-delta') this.partialText += chunk.text ?? ''; - - // Provider-executed tools (e.g. native web search) skip the local - // execution loop that emits tool-execution lifecycle events via the - // event bus. Stamp them here at chunk-arrival time so live chat and the - // persisted timeline both show a duration. A failed call arrives as a - // `tool-error` part (never a `tool-result`), so close its timing there. - if ((chunk.type === 'tool-result' || chunk.type === 'tool-error') && chunk.providerExecuted) { - await this.guard.write({ - type: 'tool-execution-end', - toolCallId: chunk.toolCallId, - toolName: chunk.toolName ?? '', - isError: chunk.type === 'tool-error', - endTime: Date.now(), - }); - } - - const converted = convertChunk(chunk); - if (converted) await this.guard.write(converted); - - if (chunk.type === 'tool-call' && chunk.providerExecuted) { - await this.guard.write({ - type: 'tool-execution-start', - toolCallId: chunk.toolCallId, - toolName: chunk.toolName ?? '', - startTime: Date.now(), - }); - } + private async writeStreamChunk(chunk: TextStreamPart): Promise { + // Provider tools skip the local executor. Record their timing when chunks arrive. + if ((chunk.type === 'tool-result' || chunk.type === 'tool-error') && chunk.providerExecuted) { + await this.guard.write({ + type: 'tool-execution-end', + toolCallId: chunk.toolCallId, + toolName: chunk.toolName ?? '', + isError: chunk.type === 'tool-error', + endTime: Date.now(), + }); } + const converted = convertChunk(chunk); + if (converted) await this.guard.write(converted); + + if (chunk.type === 'tool-call' && chunk.providerExecuted) { + await this.guard.write({ + type: 'tool-execution-start', + toolCallId: chunk.toolCallId, + toolName: chunk.toolName ?? '', + startTime: Date.now(), + }); + } + } + + private async resolveModelTurn( + result: ReturnType['result'], + ctx: ModelCallContext, + idleMs: number, + turnAbort: AbortController, + ): Promise { // The result promises settle as part of stream close on a healthy turn; // guard them with the same stall deadline so an SDK-internal promise that // never settles cannot hang the turn after the chunk loop ended. diff --git a/packages/@n8n/agents/src/runtime/tools/delegate-sub-agent-tool.ts b/packages/@n8n/agents/src/runtime/tools/delegate-sub-agent-tool.ts index 40ba669ee4a..b0ad041bde2 100644 --- a/packages/@n8n/agents/src/runtime/tools/delegate-sub-agent-tool.ts +++ b/packages/@n8n/agents/src/runtime/tools/delegate-sub-agent-tool.ts @@ -447,9 +447,25 @@ function resolveDelegateSubAgentSystemInstruction( * })); */ export function createDelegateSubAgentTool(options: CreateDelegateSubAgentToolOptions = {}) { - // Per-parent child path index for stable task paths (/root/name_0, /root/name_1, ...). + // Each parent has its own index so child task paths stay stable. const childPathIndexes = new Map(); const toolName = resolveDelegateSubAgentToolName(options.name); + const resolvedOptions = resolveDelegateSubAgentOptions(options, toolName); + const tool = buildDelegateSubAgentTool(toolName, resolvedOptions, childPathIndexes); + return withSdkOwnedBuiltInMetadata({ + ...tool, + ...createDelegateCancellationHandler(resolvedOptions, childPathIndexes), + metadata: { + ...tool.metadata, + [INLINE_DELEGATE_SUB_AGENT_TOOL_METADATA_KEY]: buildDelegateToolMetadata(resolvedOptions), + }, + }); +} + +function resolveDelegateSubAgentOptions( + options: CreateDelegateSubAgentToolOptions, + toolName: string, +): CreateDelegateSubAgentToolOptions { const resolvedOptions: CreateDelegateSubAgentToolOptions = { ...options, policy: resolveDelegateSubAgentPolicy(options.policy, toolName), @@ -462,6 +478,14 @@ export function createDelegateSubAgentTool(options: CreateDelegateSubAgentToolOp `${toolName} requires resumeSubAgent and cancelSubAgent to be configured together`, ); } + return resolvedOptions; +} + +function buildDelegateSubAgentTool( + toolName: string, + resolvedOptions: CreateDelegateSubAgentToolOptions, + childPathIndexes: Map, +) { const inlineProviderToolInstruction = resolvedOptions.resolveInlineSubAgentProviderTools ? "Provider-defined tools are loaded for the inline child's selected model provider." : 'Inline children do not inherit provider-defined tools.'; @@ -483,73 +507,77 @@ export function createDelegateSubAgentTool(options: CreateDelegateSubAgentToolOp ) => await handleDelegateSubAgent(input, ctx, resolvedOptions, childPathIndexes); const toModelOutput = (output: z.infer) => resolvedOptions.toModelOutput ? resolvedOptions.toModelOutput(output) : output; - const tool = resolvedOptions.resumeSubAgent - ? toolBuilder - .suspend(delegateSubAgentSuspendSchema) - .resume(delegateSubAgentResumeSchema) - .handler(handler) - .toModelOutput(toModelOutput) - .build() - : toolBuilder.handler(handler).toModelOutput(toModelOutput).build(); - return withSdkOwnedBuiltInMetadata({ - ...tool, - ...(resolvedOptions.cancelSubAgent !== undefined + if (resolvedOptions.resumeSubAgent) { + return toolBuilder + .suspend(delegateSubAgentSuspendSchema) + .resume(delegateSubAgentResumeSchema) + .handler(handler) + .toModelOutput(toModelOutput) + .build(); + } + return toolBuilder.handler(handler).toModelOutput(toModelOutput).build(); +} + +function createDelegateCancellationHandler( + options: CreateDelegateSubAgentToolOptions, + childPathIndexes: Map, +): Pick { + if (options.cancelSubAgent === undefined) return {}; + return { + onCancellation: async (rawInput: unknown, ctx: ToolCancellationContext) => { + const parsedInput = delegateSubAgentInputSchema.safeParse(rawInput); + if (!parsedInput.success) { + throw new Error('Delegated child input is missing or invalid'); + } + await cancelDelegatedSubAgent(parsedInput.data, ctx, options, childPathIndexes); + }, + }; +} + +function buildDelegateToolMetadata( + resolvedOptions: CreateDelegateSubAgentToolOptions, +): DelegateSubAgentToolMetadata { + return { + ...(resolvedOptions.name !== undefined ? { name: resolvedOptions.name } : {}), + ...(resolvedOptions.availableSubAgents !== undefined + ? { availableSubAgents: resolvedOptions.availableSubAgents } + : {}), + policy: resolvedOptions.policy, + ...(resolvedOptions.inlineSubAgentBlockedTools !== undefined + ? { inlineSubAgentBlockedTools: resolvedOptions.inlineSubAgentBlockedTools } + : {}), + ...(resolvedOptions.inlineSubAgentModelsByDifficulty !== undefined ? { - onCancellation: async (rawInput: unknown, ctx: ToolCancellationContext) => { - const parsedInput = delegateSubAgentInputSchema.safeParse(rawInput); - if (!parsedInput.success) { - throw new Error('Delegated child input is missing or invalid'); - } - await cancelDelegatedSubAgent(parsedInput.data, ctx, resolvedOptions, childPathIndexes); - }, + inlineSubAgentModelsByDifficulty: resolvedOptions.inlineSubAgentModelsByDifficulty, } : {}), - metadata: { - ...tool.metadata, - [INLINE_DELEGATE_SUB_AGENT_TOOL_METADATA_KEY]: { - ...(resolvedOptions.name !== undefined ? { name: resolvedOptions.name } : {}), - ...(resolvedOptions.availableSubAgents !== undefined - ? { availableSubAgents: resolvedOptions.availableSubAgents } - : {}), - policy: resolvedOptions.policy, - ...(resolvedOptions.inlineSubAgentBlockedTools !== undefined - ? { inlineSubAgentBlockedTools: resolvedOptions.inlineSubAgentBlockedTools } - : {}), - ...(resolvedOptions.inlineSubAgentModelsByDifficulty !== undefined - ? { - inlineSubAgentModelsByDifficulty: resolvedOptions.inlineSubAgentModelsByDifficulty, - } - : {}), - ...(resolvedOptions.resolveInlineSubAgentProviderTools !== undefined - ? { - resolveInlineSubAgentProviderTools: - resolvedOptions.resolveInlineSubAgentProviderTools, - } - : {}), - ...(resolvedOptions.runSubAgent !== undefined - ? { runSubAgent: resolvedOptions.runSubAgent } - : {}), - ...(resolvedOptions.resumeSubAgent !== undefined - ? { resumeSubAgent: resolvedOptions.resumeSubAgent } - : {}), - ...(resolvedOptions.shouldRetrySubAgentResumeError !== undefined - ? { shouldRetrySubAgentResumeError: resolvedOptions.shouldRetrySubAgentResumeError } - : {}), - ...(resolvedOptions.cancelSubAgent !== undefined - ? { cancelSubAgent: resolvedOptions.cancelSubAgent } - : {}), - ...(resolvedOptions.systemInstruction !== undefined - ? { systemInstruction: resolvedOptions.systemInstruction } - : {}), - ...(resolvedOptions.description !== undefined - ? { description: resolvedOptions.description } - : {}), - ...(resolvedOptions.toModelOutput !== undefined - ? { toModelOutput: resolvedOptions.toModelOutput } - : {}), - } satisfies DelegateSubAgentToolMetadata, - }, - }); + ...(resolvedOptions.resolveInlineSubAgentProviderTools !== undefined + ? { + resolveInlineSubAgentProviderTools: resolvedOptions.resolveInlineSubAgentProviderTools, + } + : {}), + ...(resolvedOptions.runSubAgent !== undefined + ? { runSubAgent: resolvedOptions.runSubAgent } + : {}), + ...(resolvedOptions.resumeSubAgent !== undefined + ? { resumeSubAgent: resolvedOptions.resumeSubAgent } + : {}), + ...(resolvedOptions.shouldRetrySubAgentResumeError !== undefined + ? { shouldRetrySubAgentResumeError: resolvedOptions.shouldRetrySubAgentResumeError } + : {}), + ...(resolvedOptions.cancelSubAgent !== undefined + ? { cancelSubAgent: resolvedOptions.cancelSubAgent } + : {}), + ...(resolvedOptions.systemInstruction !== undefined + ? { systemInstruction: resolvedOptions.systemInstruction } + : {}), + ...(resolvedOptions.description !== undefined + ? { description: resolvedOptions.description } + : {}), + ...(resolvedOptions.toModelOutput !== undefined + ? { toModelOutput: resolvedOptions.toModelOutput } + : {}), + }; } export function getInlineDelegateSubAgentToolOptions( @@ -632,26 +660,7 @@ async function handleDelegateSubAgent( request = restored.request; startedAt = Date.now(); emitSubAgentStarted(ctx, request, startedAt); - let output: DelegateSubAgentToolOutput; - try { - output = await options.resumeSubAgent( - { - ...request, - ...getChildCheckpointTarget(checkpoint), - resumeData: ctx.resumeData, - }, - createRunnerHelpers(ctx, request, options.name), - ); - } catch (error) { - if (ctx.abortSignal?.aborted || isAbortError(error)) throw error; - if (options.shouldRetrySubAgentResumeError?.(error) === false) throw error; - return await ctx.suspend(ctx.suspendPayload); - } - emitSubAgentCompleted(ctx, request, output, startedAt); - if (output.status === 'suspended') { - return await cascadeChildSuspension(ctx, request, output); - } - return output; + return await resumeDelegatedChild(ctx, restored, startedAt, options, options.resumeSubAgent); } const childPathIndexKey = getChildPathIndexKey(ctx); @@ -664,42 +673,90 @@ async function handleDelegateSubAgent( startedAt = Date.now(); emitSubAgentStarted(ctx, request, startedAt); - const toolName = options.name ?? DELEGATE_SUB_AGENT_TOOL_NAME; - if (!options.runSubAgent) { - throw new Error( - `${toolName} was registered without a runSubAgent callback, and no host runner was provided. Register it on an Agent (for inline delegation) or pass runSubAgent.`, - ); - } - const output = await options.runSubAgent( - request, + return await runDelegatedChild(ctx, request, startedAt, options); + } catch (error) { + return handleDelegateFailure(error, ctx, taskPath, request, startedAt); + } +} + +async function resumeDelegatedChild( + ctx: InterruptibleToolContext, + { request, checkpoint }: ReturnType, + startedAt: number, + options: CreateDelegateSubAgentToolOptions, + resumeSubAgent: DelegateSubAgentResumeRunner, +): Promise { + let output: DelegateSubAgentToolOutput; + try { + output = await resumeSubAgent( + { + ...request, + ...getChildCheckpointTarget(checkpoint), + resumeData: ctx.resumeData, + }, createRunnerHelpers(ctx, request, options.name), ); - emitSubAgentCompleted(ctx, request, output, startedAt); - if ( - output.status === 'suspended' && - options.resumeSubAgent !== undefined && - isInterruptibleToolContext(ctx) - ) { - return await cascadeChildSuspension(ctx, request, output); - } - return output; } catch (error) { - // When the parent has a signal it is the authority: `isAbortError` also - // matches by message text, and an unrelated child error must not be - // mistaken for a cancellation and kill the parent run. - const aborted = ctx.abortSignal ? ctx.abortSignal.aborted : isAbortError(error); - const output: DelegateSubAgentToolOutput = { - status: aborted ? 'cancelled' : 'failed', - ...(taskPath !== undefined ? { taskPath } : {}), - answer: '', - ...(aborted ? {} : { error: stringifyUnknown(error) }), - }; - if (request !== undefined && startedAt !== undefined) { - emitSubAgentCompleted(ctx, request, output, startedAt); - } - if (aborted) throw error; - return output; + if (ctx.abortSignal?.aborted || isAbortError(error)) throw error; + if (options.shouldRetrySubAgentResumeError?.(error) === false) throw error; + return await ctx.suspend(ctx.suspendPayload); } + emitSubAgentCompleted(ctx, request, output, startedAt); + if (output.status === 'suspended') { + return await cascadeChildSuspension(ctx, request, output); + } + return output; +} + +async function runDelegatedChild( + ctx: ToolContext | InterruptibleToolContext, + request: DelegateSubAgentRequest, + startedAt: number, + options: CreateDelegateSubAgentToolOptions, +): Promise { + const toolName = options.name ?? DELEGATE_SUB_AGENT_TOOL_NAME; + if (!options.runSubAgent) { + throw new Error( + `${toolName} was registered without a runSubAgent callback, and no host runner was provided. Register it on an Agent (for inline delegation) or pass runSubAgent.`, + ); + } + const output = await options.runSubAgent( + request, + createRunnerHelpers(ctx, request, options.name), + ); + emitSubAgentCompleted(ctx, request, output, startedAt); + if ( + output.status === 'suspended' && + options.resumeSubAgent !== undefined && + isInterruptibleToolContext(ctx) + ) { + return await cascadeChildSuspension(ctx, request, output); + } + return output; +} + +function handleDelegateFailure( + error: unknown, + ctx: ToolContext, + taskPath: SubAgentTaskPath | undefined, + request: DelegateSubAgentRequest | undefined, + startedAt: number | undefined, +): DelegateSubAgentToolOutput { + // When the parent has a signal it is the authority: `isAbortError` also + // matches by message text, and an unrelated child error must not be + // mistaken for a cancellation and kill the parent run. + const aborted = ctx.abortSignal ? ctx.abortSignal.aborted : isAbortError(error); + const output: DelegateSubAgentToolOutput = { + status: aborted ? 'cancelled' : 'failed', + ...(taskPath !== undefined ? { taskPath } : {}), + answer: '', + ...(aborted ? {} : { error: stringifyUnknown(error) }), + }; + if (request !== undefined && startedAt !== undefined) { + emitSubAgentCompleted(ctx, request, output, startedAt); + } + if (aborted) throw error; + return output; } function createDelegateSubAgentRequest( @@ -1001,12 +1058,7 @@ export function generateResultToDelegateSubAgentOutput( ...(result.structuredOutput !== undefined ? { structuredOutput: result.structuredOutput } : {}), ...(result.usage !== undefined ? { - usage: { - promptTokens: result.usage.promptTokens, - completionTokens: result.usage.completionTokens, - totalTokens: result.usage.totalTokens, - ...(result.usage.cost !== undefined ? { cost: result.usage.cost } : {}), - }, + usage: toDelegateUsage(result.usage), } : {}), ...(result.finishReason !== undefined ? { finishReason: result.finishReason } : {}), @@ -1017,6 +1069,15 @@ export function generateResultToDelegateSubAgentOutput( }; } +function toDelegateUsage(usage: TokenUsage): DelegateSubAgentToolOutput['usage'] { + return { + promptTokens: usage.promptTokens, + completionTokens: usage.completionTokens, + totalTokens: usage.totalTokens, + ...(usage.cost !== undefined ? { cost: usage.cost } : {}), + }; +} + /** Last non-empty assistant text across the run's messages. */ function lastText(messages: AgentMessage[]): string { const llmMessages = filterLlmMessages(messages); diff --git a/packages/@n8n/agents/src/runtime/tools/tool-adapter.ts b/packages/@n8n/agents/src/runtime/tools/tool-adapter.ts index 41fe09bac1f..ae58c6ff8dd 100644 --- a/packages/@n8n/agents/src/runtime/tools/tool-adapter.ts +++ b/packages/@n8n/agents/src/runtime/tools/tool-adapter.ts @@ -117,42 +117,14 @@ export async function executeTool( throw new Error(`No handler found for tool "${builtTool.name}"`); } - // Anchor programmatic skill activations to the calling tool's result by - // default, so the skill body rides on that result (kept out of `system`) - // instead of forcing a full prompt-prefix rewrite. An explicit anchor wins. - const baseLoadSkill = executionContext.loadSkill; - const loadSkill: ToolExecutionContext['loadSkill'] = baseLoadSkill - ? async (skillId, anchor) => - await baseLoadSkill(skillId, anchor ?? (toolCallId ? { toolCallId } : undefined)) - : undefined; - + const loadSkill = anchorSkillLoader(executionContext.loadSkill, toolCallId); if (builtTool.suspendSchema) { const isCancelled = isCancellation(resumeData); const ctx: InterruptibleToolContext = { - suspend: async (payload: unknown, options?: ToolSuspendOptions): Promise => { - const resolvedOptions: ToolSuspendOptions = { - continuation: executionContext.continuation, - ...options, - resumeSchema: options?.resumeSchema ?? executionContext.resumeSchema, - }; - await executionContext.onSuspend?.(payload, resolvedOptions); - return await Promise.resolve({ - [SUSPEND_BRAND]: true, - payload, - ...resolvedOptions, - } as never); - }, + suspend: createSuspendHandler(executionContext), resumeData: isCancelled ? undefined : resumeData, cancellation: isCancelled ? { message: resumeData.message } : undefined, - parentTelemetry, - toolCallId, - toolName: builtTool.name, - runId: executionContext.runId, - ...(loadSkill ? { loadSkill } : {}), - persistence: executionContext.persistence, - emitEvent: executionContext.emitEvent, - abortSignal: executionContext.abortSignal, - executionCounter: executionContext.executionCounter, + ...createToolContext(builtTool, executionContext, parentTelemetry, toolCallId, loadSkill), suspendPayload: executionContext.suspendPayload, continuation: executionContext.continuation, resumeSchema: executionContext.resumeSchema, @@ -160,7 +132,34 @@ export async function executeTool( return await builtTool.handler(args, ctx); } - const ctx: ToolContext = { + const ctx = createToolContext( + builtTool, + executionContext, + parentTelemetry, + toolCallId, + loadSkill, + ); + return await builtTool.handler(args, ctx); +} + +function anchorSkillLoader( + loadSkill: ToolExecutionContext['loadSkill'], + toolCallId: string | undefined, +): ToolExecutionContext['loadSkill'] { + if (!loadSkill) return undefined; + // Keep skill content on the tool result. An explicit anchor takes precedence. + return async (skillId, anchor) => + await loadSkill(skillId, anchor ?? (toolCallId ? { toolCallId } : undefined)); +} + +function createToolContext( + builtTool: BuiltTool, + executionContext: ToolExecutionContext, + parentTelemetry: BuiltTelemetry | undefined, + toolCallId: string | undefined, + loadSkill: ToolExecutionContext['loadSkill'], +): ToolContext { + return { parentTelemetry, toolCallId, toolName: builtTool.name, @@ -171,7 +170,24 @@ export async function executeTool( abortSignal: executionContext.abortSignal, executionCounter: executionContext.executionCounter, }; - return await builtTool.handler(args, ctx); +} + +function createSuspendHandler( + executionContext: ToolExecutionContext, +): InterruptibleToolContext['suspend'] { + return async (payload: unknown, options?: ToolSuspendOptions): Promise => { + const resolvedOptions: ToolSuspendOptions = { + continuation: executionContext.continuation, + ...options, + resumeSchema: options?.resumeSchema ?? executionContext.resumeSchema, + }; + await executionContext.onSuspend?.(payload, resolvedOptions); + return await Promise.resolve({ + [SUSPEND_BRAND]: true, + payload, + ...resolvedOptions, + } as never); + }; } /** diff --git a/packages/@n8n/agents/src/runtime/tools/tool-call-executor.ts b/packages/@n8n/agents/src/runtime/tools/tool-call-executor.ts index 95ea6bd6338..5d6c8b77659 100644 --- a/packages/@n8n/agents/src/runtime/tools/tool-call-executor.ts +++ b/packages/@n8n/agents/src/runtime/tools/tool-call-executor.ts @@ -122,13 +122,28 @@ export interface ToolCallBatchResult { pending: Record; } -interface RuntimeToolCall { +interface ToolCallInput { toolCallId: string; toolName: string; - input: JSONObject; + input: unknown; providerExecuted?: boolean; } +interface RuntimeToolCall extends ToolCallInput { + input: JSONObject; +} + +type ToolCallIdentity = Pick; +type SuspendedToolOutcome = Extract; + +interface InterruptedToolSuspension { + didSuspend: boolean; + abortObserved: boolean; + payload: unknown; + options?: ToolSuspendOptions; + cleanup?: Promise; +} + /** Shared input for the tool-call batch iterators. */ export interface ToolBatchContext { toolMap: Map; @@ -288,33 +303,54 @@ export class ToolCallExecutor { * even if one throws, then re-throws the first error. */ async iterateToolCallsConcurrent( - ctx: ToolBatchContext & { - toolCalls: Array<{ - toolCallId: string; - toolName: string; - input: unknown; - providerExecuted?: boolean; - }>; - }, + ctx: ToolBatchContext & { toolCalls: ToolCallInput[] }, ): Promise { - const { - toolCalls, - toolMap, - list, - runId, - telemetry: resolvedTelemetry, - executionCounter, - abortSignal, - } = ctx; const errors: ToolCallError[] = []; - const runtimeToolCalls: RuntimeToolCall[] = []; + const runtimeToolCalls = this.normalizeToolCalls(ctx.toolCalls, ctx, errors); + const executableCalls = runtimeToolCalls.filter((tc) => !tc.providerExecuted); + const providerExecutedCount = runtimeToolCalls.length - executableCalls.length; + for (let i = 0; i < providerExecutedCount; i++) { + incrementToolCallCount(ctx.executionCounter); + } + const executableCallsById = new Map(executableCalls.map((tc) => [tc.toolCallId, tc])); + const unexecutedIds = new Set(executableCalls.map((tc) => tc.toolCallId)); + const result: ToolCallBatchResult = { results: [], suspensions: [], errors, pending: {} }; + for (let batchStart = 0; batchStart < executableCalls.length; ) { + if (ctx.isAborted()) { + this.skipAbortedToolCalls(unexecutedIds, executableCallsById, ctx.list, result.results); + return await this.finalizeBatch(result, ctx); + } + const batch = this.takeNextToolCallBatch(executableCalls, batchStart, ctx.toolMap); + batchStart += batch.length; + const settled = await this.executeToolCallBatch(batch, ctx); + for (const tc of batch) unexecutedIds.delete(tc.toolCallId); + const hasSuspension = this.collectBatchResults(batch, settled, result, ctx); + + if (ctx.isAborted()) { + this.skipAbortedToolCalls(unexecutedIds, executableCallsById, ctx.list, result.results); + return await this.finalizeBatch(result, ctx); + } + if (hasSuspension) { + this.retainUnexecutedToolCalls(unexecutedIds, executableCallsById, result.pending); + break; + } + } + return await this.finalizeBatch(result, ctx); + } + + private normalizeToolCalls( + toolCalls: ToolCallInput[], + ctx: ToolBatchContext, + errors: ToolCallError[], + ): RuntimeToolCall[] { + const calls: RuntimeToolCall[] = []; for (const toolCall of toolCalls) { const normalizedInput = normalizeToolInputForModel(toolCall.input); if (!normalizedInput.ok) { const error = new Error(normalizedInput.error); - incrementToolCallCount(executionCounter); - list.setToolCallError(toolCall.toolCallId, error); + incrementToolCallCount(ctx.executionCounter); + ctx.list.setToolCallError(toolCall.toolCallId, error); errors.push({ toolCallId: toolCall.toolCallId, toolName: toolCall.toolName, @@ -323,177 +359,180 @@ export class ToolCallExecutor { }); continue; } - - runtimeToolCalls.push({ ...toolCall, input: normalizedInput.input }); + calls.push({ ...toolCall, input: normalizedInput.input }); } + return calls; + } - const executableCalls = runtimeToolCalls.filter((tc) => !tc.providerExecuted); - const providerExecutedCount = runtimeToolCalls.length - executableCalls.length; - for (let i = 0; i < providerExecutedCount; i++) { - incrementToolCallCount(executionCounter); + private async executeToolCallBatch( + batch: RuntimeToolCall[], + ctx: ToolBatchContext, + ): Promise>> { + return await Promise.allSettled( + batch.map( + async (tc) => + await this.processToolCall({ + toolCallId: tc.toolCallId, + toolName: tc.toolName, + input: tc.input, + toolMap: ctx.toolMap, + list: ctx.list, + runId: ctx.runId, + persistence: ctx.persistence, + resolvedTelemetry: ctx.telemetry, + executionCounter: ctx.executionCounter, + abortSignal: ctx.abortSignal, + countToolCall: true, + }), + ), + ); + } + + private collectBatchResults( + calls: RuntimeToolCall[], + settled: Array>, + batch: ToolCallBatchResult, + ctx: ToolBatchContext, + ): boolean { + let hasSuspension = false; + for (let i = 0; i < settled.length; i++) { + if (this.collectSettledToolCall(calls[i], settled[i], batch, ctx)) hasSuspension = true; } - const executableCallsById = new Map(executableCalls.map((tc) => [tc.toolCallId, tc])); - const unexecutedIds = new Set(executableCalls.map((tc) => tc.toolCallId)); - const results: ToolCallSuccess[] = []; - const suspensions: ToolCallSuspension[] = []; - const pending: Record = {}; + return hasSuspension; + } - for (let batchStart = 0; batchStart < executableCalls.length; ) { - if (ctx.isAborted()) { - this.deps.onCancelled(); - for (const id of unexecutedIds) { - const tc = executableCallsById.get(id)!; - const modelOutput = '[Skipped: run was aborted]'; - list.setToolCallResult(tc.toolCallId, modelOutput, { canceled: true }); - results.push({ - toolCallId: tc.toolCallId, - toolName: tc.toolName, - input: tc.input, - toolEntry: { - tool: tc.toolName, - input: tc.input, - output: modelOutput, - transformed: false, - canceled: true, - }, - modelOutput, - }); - } - return await this.finalizeBatch({ results, suspensions, errors, pending }, ctx); - } - - const batch = this.takeNextToolCallBatch(executableCalls, batchStart, toolMap); - batchStart += batch.length; - - const settledResults = await Promise.allSettled( - batch.map( - async (tc) => - await this.processToolCall({ - toolCallId: tc.toolCallId, - toolName: tc.toolName, - input: tc.input, - toolMap, - list, - runId, - persistence: ctx.persistence, - resolvedTelemetry, - executionCounter, - abortSignal, - countToolCall: true, - }), - ), - ); - - for (const tc of batch) { - unexecutedIds.delete(tc.toolCallId); - } - - let hasSuspension = false; - - for (let i = 0; i < settledResults.length; i++) { - const result = settledResults[i]; - const tc = batch[i]; - const toolInput = tc.input; - - if (result.status === 'rejected') { - list.setToolCallError(tc.toolCallId, result.reason); - errors.push({ - toolCallId: tc.toolCallId, - toolName: tc.toolName, - input: toolInput, - error: result.reason, - }); - } else if (result.value.outcome === 'suspended') { - hasSuspension = true; - suspensions.push({ - toolCallId: tc.toolCallId, - toolName: tc.toolName, - input: toolInput, - payload: result.value.payload, - resumeSchema: result.value.resumeSchema, - }); - pending[tc.toolCallId] = { - suspended: true, - toolCallId: tc.toolCallId, - toolName: tc.toolName, - input: toolInput, - suspendPayload: result.value.payload, - resumeSchema: result.value.resumeSchema, - ...(result.value.continuation !== undefined - ? { continuation: result.value.continuation } - : {}), - runId, - }; - } else if (result.value.outcome === 'success') { - results.push({ - toolCallId: tc.toolCallId, - toolName: tc.toolName, - input: toolInput, - toolEntry: result.value.toolEntry, - modelOutput: result.value.modelOutput, - customMessage: result.value.customMessage, - ...(result.value.mcpServerName !== undefined - ? { mcpServerName: result.value.mcpServerName } - : {}), - }); - } else if (result.value.outcome === 'cancelled') { - results.push({ - toolCallId: tc.toolCallId, - toolName: tc.toolName, - input: toolInput, - toolEntry: result.value.toolEntry, - modelOutput: result.value.modelOutput, - }); - } else if (result.value.outcome === 'error') { - errors.push({ - toolCallId: tc.toolCallId, - toolName: tc.toolName, - input: toolInput, - error: result.value.error, - }); - } else if (result.value.outcome === 'noop') { - // noop - } - } - - if (ctx.isAborted()) { - this.deps.onCancelled(); - for (const id of unexecutedIds) { - const tc = executableCallsById.get(id)!; - const modelOutput = '[Skipped: run was aborted]'; - list.setToolCallResult(tc.toolCallId, modelOutput, { canceled: true }); - results.push({ - toolCallId: tc.toolCallId, - toolName: tc.toolName, - input: tc.input, - toolEntry: { - tool: tc.toolName, - input: tc.input, - output: modelOutput, - transformed: false, - canceled: true, - }, - modelOutput, - }); - } - return await this.finalizeBatch({ results, suspensions, errors, pending }, ctx); - } - - if (hasSuspension) { - for (const id of unexecutedIds) { - const tc = executableCallsById.get(id)!; - pending[tc.toolCallId] = { - suspended: false, - toolCallId: tc.toolCallId, - toolName: tc.toolName, - input: tc.input, - }; - } + private collectSettledToolCall( + call: RuntimeToolCall, + result: PromiseSettledResult, + batch: ToolCallBatchResult, + ctx: ToolBatchContext, + ): boolean { + if (result.status === 'rejected') { + ctx.list.setToolCallError(call.toolCallId, result.reason); + batch.errors.push({ + toolCallId: call.toolCallId, + toolName: call.toolName, + input: call.input, + error: result.reason, + }); + return false; + } + const outcome = result.value; + switch (outcome.outcome) { + case 'suspended': + batch.suspensions.push(this.buildToolCallSuspension(call, outcome)); + batch.pending[call.toolCallId] = this.buildPendingToolCall(call, outcome, ctx.runId); + return true; + case 'success': + case 'cancelled': + this.appendCompletedToolCall(call, outcome, batch.results); + break; + case 'error': + batch.errors.push({ + toolCallId: call.toolCallId, + toolName: call.toolName, + input: call.input, + error: outcome.error, + }); break; - } } + return false; + } - return await this.finalizeBatch({ results, suspensions, errors, pending }, ctx); + private buildToolCallSuspension( + call: ToolCallIdentity, + outcome: SuspendedToolOutcome, + ): ToolCallSuspension { + return { + toolCallId: call.toolCallId, + toolName: call.toolName, + input: call.input, + payload: outcome.payload, + resumeSchema: outcome.resumeSchema, + }; + } + + private buildPendingToolCall( + call: ToolCallIdentity, + outcome: SuspendedToolOutcome, + runId: string, + ): PendingToolCall { + return { + suspended: true, + toolCallId: call.toolCallId, + toolName: call.toolName, + input: call.input, + suspendPayload: outcome.payload, + resumeSchema: outcome.resumeSchema, + ...(outcome.continuation !== undefined ? { continuation: outcome.continuation } : {}), + runId, + }; + } + + private appendCompletedToolCall( + call: ToolCallIdentity, + outcome: Extract, + results: ToolCallSuccess[], + ): void { + const result: ToolCallSuccess = { + toolCallId: call.toolCallId, + toolName: call.toolName, + input: call.input, + toolEntry: outcome.toolEntry, + modelOutput: outcome.modelOutput, + }; + if (outcome.outcome === 'success') { + result.customMessage = outcome.customMessage; + if (outcome.mcpServerName !== undefined) result.mcpServerName = outcome.mcpServerName; + } + results.push(result); + } + + private skipAbortedToolCalls( + ids: Set, + calls: Map, + list: AgentMessageList, + results: ToolCallSuccess[], + ): void { + this.deps.onCancelled(); + for (const id of ids) { + const call = calls.get(id)!; + const modelOutput = '[Skipped: run was aborted]'; + list.setToolCallResult(call.toolCallId, modelOutput, { canceled: true }); + results.push(this.buildSkippedToolResult(call, modelOutput)); + } + } + + private buildSkippedToolResult(call: ToolCallIdentity, modelOutput: string): ToolCallSuccess { + return { + toolCallId: call.toolCallId, + toolName: call.toolName, + input: call.input, + toolEntry: { + tool: call.toolName, + input: call.input, + output: modelOutput, + transformed: false, + canceled: true, + }, + modelOutput, + }; + } + + private retainUnexecutedToolCalls( + ids: Set, + calls: Map, + pending: Record, + ): void { + for (const id of ids) { + const call = calls.get(id)!; + pending[call.toolCallId] = { + suspended: false, + toolCallId: call.toolCallId, + toolName: call.toolName, + input: call.input, + }; + } } /** @@ -509,172 +548,112 @@ export class ToolCallExecutor { async iteratePendingToolCallsConcurrent( ctx: ToolBatchContext & { pendingResume: PendingResume }, ): Promise { - const { - pendingResume, - toolMap, - list, - runId, - persistence, - telemetry: resolvedTelemetry, - executionCounter, - abortSignal, - } = ctx; + const { pendingResume } = ctx; const resumedId = pendingResume.resumeToolCallId; const resumedEntry = pendingResume.pendingToolCalls[resumedId]; - if (!resumedEntry) { - throw new Error(`No pending tool call found for toolCallId: ${resumedId}`); - } - - const resumedToolName = resumedEntry.toolName; - const results: ToolCallSuccess[] = []; - const suspensions: ToolCallSuspension[] = []; - const errors: ToolCallError[] = []; - const pending: Record = {}; - - // 1. Execute the resumed tool - const processResult = await this.processToolCall( + if (!resumedEntry) throw new Error(`No pending tool call found for toolCallId: ${resumedId}`); + const batch: ToolCallBatchResult = { results: [], suspensions: [], errors: [], pending: {} }; + const outcome = await this.processToolCall( this.pendingToolCallParams(resumedEntry, ctx, pendingResume.resumeData), ); - if (processResult.outcome === 'suspended') { - pending[resumedId] = { - suspended: true, - toolCallId: resumedEntry.toolCallId, - toolName: resumedToolName, - input: resumedEntry.input, - suspendPayload: processResult.payload, - resumeSchema: processResult.resumeSchema, - ...(processResult.continuation !== undefined - ? { continuation: processResult.continuation } - : {}), - runId, - }; - suspensions.push({ - toolCallId: resumedId, - toolName: resumedToolName, - input: resumedEntry.input, - payload: processResult.payload, - resumeSchema: processResult.resumeSchema, - }); - } else if (processResult.outcome === 'success') { - results.push({ - toolCallId: resumedEntry.toolCallId, - toolName: resumedToolName, - input: resumedEntry.input, - toolEntry: processResult.toolEntry, - modelOutput: processResult.modelOutput, - customMessage: processResult.customMessage, - ...(processResult.mcpServerName !== undefined - ? { mcpServerName: processResult.mcpServerName } - : {}), - }); - } else if (processResult.outcome === 'cancelled') { - results.push({ - toolCallId: resumedEntry.toolCallId, - toolName: resumedToolName, - input: resumedEntry.input, - toolEntry: processResult.toolEntry, - modelOutput: processResult.modelOutput, - }); - list.addInput([ - { role: 'user', content: [{ type: 'text', text: processResult.userMessage }] }, - ]); - - for (const id of Object.keys(pendingResume.pendingToolCalls)) { - if (id !== resumedId) { - const siblingEntry = pendingResume.pendingToolCalls[id]; - const siblingTool = toolMap.get(siblingEntry.toolName); - const siblingParams = this.pendingToolCallParams( - siblingEntry, - ctx, - pendingResume.resumeData, - ); - if (siblingEntry.suspended && siblingTool?.onCancellation) { - try { - await this.runCancellationCleanup( - siblingParams, - siblingTool, - processResult.userMessage, - ); - } catch { - this.retainPendingToolCall(id, siblingEntry, pending, suspensions); - continue; - } - } - const modelOutput = '[Skipped: a sibling tool call was cancelled]'; - list.setToolCallResult(id, modelOutput, { - canceled: true, - }); - results.push({ - toolCallId: siblingEntry.toolCallId, - toolName: siblingEntry.toolName, - input: siblingEntry.input, - toolEntry: { - tool: siblingEntry.toolName, - input: siblingEntry.input, - output: modelOutput, - transformed: false, - canceled: true, - }, - modelOutput, - }); - } - } - - return await this.finalizeBatch({ results, suspensions, errors, pending }, ctx); - } else if (processResult.outcome === 'retryable-error') { + if (outcome.outcome === 'cancelled') { + this.appendCompletedToolCall(resumedEntry, outcome, batch.results); + ctx.list.addInput([{ role: 'user', content: [{ type: 'text', text: outcome.userMessage }] }]); + await this.cancelSiblingToolCalls(ctx, batch, outcome.userMessage); + return await this.finalizeBatch(batch, ctx); + } + if (outcome.outcome === 'retryable-error') { for (const [id, entry] of Object.entries(pendingResume.pendingToolCalls)) { - this.retainPendingToolCall(id, entry, pending, suspensions); + this.retainPendingToolCall(id, entry, batch.pending, batch.suspensions); } - return await this.finalizeBatch({ results, suspensions, errors, pending }, ctx); - } else if (processResult.outcome === 'error') { - errors.push({ - toolCallId: resumedEntry.toolCallId, - toolName: resumedToolName, - input: resumedEntry.input, - error: processResult.error, - }); - } else if (processResult.outcome === 'noop') { - // noop + return await this.finalizeBatch(batch, ctx); } - // 2. Process remaining items - const unexecuted: Array<{ toolCallId: string; toolName: string; input: unknown }> = []; + this.collectResumedToolCall(resumedId, resumedEntry, outcome, batch, ctx.runId); + await this.executeRemainingToolCalls(ctx, batch); + return await this.finalizeBatch(batch, ctx); + } - for (const [id, entry] of Object.entries(pendingResume.pendingToolCalls)) { - if (id === resumedId) continue; - - if (entry.suspended) { - this.retainPendingToolCall(id, entry, pending, suspensions); - } else { - // Unexecuted — collect for batch execution - unexecuted.push({ - toolCallId: id, + private collectResumedToolCall( + id: string, + entry: PendingToolCall, + outcome: ToolCallOutcome, + batch: ToolCallBatchResult, + runId: string, + ): void { + switch (outcome.outcome) { + case 'suspended': + batch.pending[id] = this.buildPendingToolCall(entry, outcome, runId); + batch.suspensions.push(this.buildToolCallSuspension({ ...entry, toolCallId: id }, outcome)); + break; + case 'success': + this.appendCompletedToolCall(entry, outcome, batch.results); + break; + case 'error': + batch.errors.push({ + toolCallId: entry.toolCallId, toolName: entry.toolName, input: entry.input, + error: outcome.error, }); + break; + } + } + + private async cancelSiblingToolCalls( + ctx: ToolBatchContext & { pendingResume: PendingResume }, + batch: ToolCallBatchResult, + userMessage: string, + ): Promise { + for (const id of Object.keys(ctx.pendingResume.pendingToolCalls)) { + if (id === ctx.pendingResume.resumeToolCallId) continue; + const entry = ctx.pendingResume.pendingToolCalls[id]; + const tool = ctx.toolMap.get(entry.toolName); + const params = this.pendingToolCallParams(entry, ctx, ctx.pendingResume.resumeData); + if (entry.suspended && tool?.onCancellation) { + try { + await this.runCancellationCleanup(params, tool, userMessage); + } catch { + this.retainPendingToolCall(id, entry, batch.pending, batch.suspensions); + continue; + } + } + const modelOutput = '[Skipped: a sibling tool call was cancelled]'; + ctx.list.setToolCallResult(id, modelOutput, { canceled: true }); + batch.results.push(this.buildSkippedToolResult(entry, modelOutput)); + } + } + + private async executeRemainingToolCalls( + ctx: ToolBatchContext & { pendingResume: PendingResume }, + result: ToolCallBatchResult, + ): Promise { + const unexecuted: ToolCallInput[] = []; + for (const [id, entry] of Object.entries(ctx.pendingResume.pendingToolCalls)) { + if (id === ctx.pendingResume.resumeToolCallId) continue; + if (entry.suspended) { + this.retainPendingToolCall(id, entry, result.pending, result.suspensions); + } else { + unexecuted.push({ toolCallId: id, toolName: entry.toolName, input: entry.input }); } } - - // Execute unexecuted tools via iterateToolCallsConcurrent - if (unexecuted.length > 0) { - const batch = await this.iterateToolCallsConcurrent({ - toolCalls: unexecuted, - toolMap, - list, - runId, - persistence, - telemetry: resolvedTelemetry, - executionCounter, - abortSignal, - isAborted: ctx.isAborted, - }); - results.push(...batch.results); - suspensions.push(...batch.suspensions); - errors.push(...batch.errors); - Object.assign(pending, batch.pending); - } - return await this.finalizeBatch({ results, suspensions, errors, pending }, ctx); + if (unexecuted.length === 0) return; + const batch = await this.iterateToolCallsConcurrent({ + toolCalls: unexecuted, + toolMap: ctx.toolMap, + list: ctx.list, + runId: ctx.runId, + persistence: ctx.persistence, + telemetry: ctx.telemetry, + executionCounter: ctx.executionCounter, + abortSignal: ctx.abortSignal, + isAborted: ctx.isAborted, + }); + result.results.push(...batch.results); + result.suspensions.push(...batch.suspensions); + result.errors.push(...batch.errors); + Object.assign(result.pending, batch.pending); } /** @@ -728,61 +707,83 @@ export class ToolCallExecutor { }); } + return await this.executeValidatedToolCall(params, builtTool, input); + } + + private async executeValidatedToolCall( + params: ProcessToolCallParams, + builtTool: BuiltTool, + input: JSONValue, + ): Promise { let toolResult: unknown; - let interruptedSuspendPayload: unknown; - let interruptedSuspendOptions: ToolSuspendOptions | undefined; - let didSuspend = false; - let abortObserved = false; - let suspensionCleanup: Promise | undefined; - const cleanupInterruptedSuspension = async () => { - suspensionCleanup ??= this.runCancellationCleanup( - { - ...params, - input, - suspendPayload: interruptedSuspendPayload, - continuation: interruptedSuspendOptions?.continuation, - resumeSchema: getToolResumeJsonSchema(builtTool, interruptedSuspendOptions?.resumeSchema), - }, - builtTool, - 'Run aborted', - ).catch(() => undefined); - await suspensionCleanup; + const suspension: InterruptedToolSuspension = { + didSuspend: false, + abortObserved: false, + payload: undefined, }; try { toolResult = await this.runToolHandler(params, builtTool, input, async (payload, options) => { - didSuspend = true; - interruptedSuspendPayload = payload; - interruptedSuspendOptions = options; - if (abortObserved || params.abortSignal?.aborted) { - await cleanupInterruptedSuspension(); + suspension.didSuspend = true; + suspension.payload = payload; + suspension.options = options; + if (suspension.abortObserved || params.abortSignal?.aborted) { + await this.cleanupInterruptedSuspension(params, builtTool, input, suspension); } }); } catch (error) { if (isAbortError(error) || params.abortSignal?.aborted) { - abortObserved = true; - if (didSuspend) { - await cleanupInterruptedSuspension(); - } else if (params.suspendPayload !== undefined || params.continuation !== undefined) { - try { - await this.runCancellationCleanup({ ...params, input }, builtTool, 'Run aborted'); - } catch { - // Parent shutdown must continue; persistent stores will prune stale checkpoints. - } - } + suspension.abortObserved = true; + await this.cleanupAbortedTool(params, builtTool, input, suspension); this.deps.onCancelled(); return this.buildCancelledOutcome(params, 'Run aborted'); } - return await this.toolError(params, error, builtTool); } if (isSuspendedToolResult(toolResult)) { return await this.buildSuspendedOutcome(params, builtTool, toolResult); } - return await this.buildSuccessOutcome(params, builtTool, input, toolResult); } + private async cleanupInterruptedSuspension( + params: ProcessToolCallParams, + builtTool: BuiltTool, + input: JSONValue, + suspension: InterruptedToolSuspension, + ): Promise { + // Both abort and late suspension can reach this path. Run cleanup only once. + suspension.cleanup ??= this.runCancellationCleanup( + { + ...params, + input, + suspendPayload: suspension.payload, + continuation: suspension.options?.continuation, + resumeSchema: getToolResumeJsonSchema(builtTool, suspension.options?.resumeSchema), + }, + builtTool, + 'Run aborted', + ).catch(() => undefined); + await suspension.cleanup; + } + + private async cleanupAbortedTool( + params: ProcessToolCallParams, + builtTool: BuiltTool, + input: JSONValue, + suspension: InterruptedToolSuspension, + ): Promise { + if (suspension.didSuspend) { + await this.cleanupInterruptedSuspension(params, builtTool, input, suspension); + } else if (params.suspendPayload !== undefined || params.continuation !== undefined) { + try { + await this.runCancellationCleanup({ ...params, input }, builtTool, 'Run aborted'); + } catch { + // Parent shutdown must continue; persistent stores will prune stale checkpoints. + } + } + } + private async runCancellationCleanup( params: ProcessToolCallParams, builtTool: BuiltTool, @@ -1095,24 +1096,12 @@ export class ToolCallExecutor { list.setToolCallResult(toolCallId, guardedResult.historyOutput); - let customMessage = await builtTool.toMessage?.(toolResult); - if (customMessage && builtTool.outputTrust === 'untrusted') { - customMessage = protectUntrustedToolMessage(customMessage, builtTool); - } - let guardedCustomMessage = customMessage - ? await guardToolMessageForModel(customMessage, this.deps.tokenCounter, storage) - : undefined; - // Stamp tool provenance so derived transcripts (e.g. the observation - // log observer) can keep this content inside untrusted-data boundaries. - if (guardedCustomMessage && 'role' in guardedCustomMessage) { - guardedCustomMessage = { - ...guardedCustomMessage, - origin: { kind: 'tool', toolName }, - }; - } - if (guardedCustomMessage) { - list.addResponse([guardedCustomMessage]); - } + const guardedCustomMessage = await this.appendCustomToolMessage( + params, + builtTool, + toolResult, + storage, + ); return { outcome: 'success', @@ -1128,6 +1117,33 @@ export class ToolCallExecutor { }; } + private async appendCustomToolMessage( + params: ProcessToolCallParams, + builtTool: BuiltTool, + toolResult: unknown, + storage: ToolResultGuardStorage | undefined, + ): Promise { + let customMessage = await builtTool.toMessage?.(toolResult); + if (customMessage && builtTool.outputTrust === 'untrusted') { + customMessage = protectUntrustedToolMessage(customMessage, builtTool); + } + let guardedCustomMessage = customMessage + ? await guardToolMessageForModel(customMessage, this.deps.tokenCounter, storage) + : undefined; + // Stamp tool provenance so derived transcripts (e.g. the observation + // log observer) can keep this content inside untrusted-data boundaries. + if (guardedCustomMessage && 'role' in guardedCustomMessage) { + guardedCustomMessage = { + ...guardedCustomMessage, + origin: { kind: 'tool', toolName: params.toolName }, + }; + } + if (guardedCustomMessage) { + params.list.addResponse([guardedCustomMessage]); + } + return guardedCustomMessage; + } + private getResultStorage(params: ProcessToolCallParams): ToolResultGuardStorage | undefined { const filesystem = this.deps.workspaceFilesystem; if (!filesystem) return undefined; diff --git a/packages/@n8n/agents/src/runtime/tools/tool-result-guard.ts b/packages/@n8n/agents/src/runtime/tools/tool-result-guard.ts index 5ba0dd55479..7386b7b6db0 100644 --- a/packages/@n8n/agents/src/runtime/tools/tool-result-guard.ts +++ b/packages/@n8n/agents/src/runtime/tools/tool-result-guard.ts @@ -197,39 +197,41 @@ export function sanitizeOffloadedToolResultsForMemory( return messages.map((message) => { if (!('content' in message)) return { ...message }; - const content = message.content.map((block): MessageContent => { - if (block.type === 'tool-call') { - if (block.state === 'resolved' && isOffloadedToolResult(block.output)) { - return { ...block, output: { ...EXPIRED_OFFLOADED_TOOL_RESULT } }; - } - if (block.state === 'resolved' && isContentToolResultOutput(block.output)) { - return { - ...block, - output: toJsonValue({ - type: 'content', - value: block.output.value.map((part) => - part.type === 'text' && isSerializedOffloadedToolResult(part.text) - ? { ...part, text: EXPIRED_OFFLOADED_TOOL_RESULT_JSON } - : part, - ), - }), - }; - } - if (block.state === 'rejected' && isSerializedOffloadedToolResult(block.error)) { - return { ...block, error: EXPIRED_OFFLOADED_TOOL_RESULT_JSON }; - } - } - - if (block.type === 'text' && isSerializedOffloadedToolResult(block.text)) { - return { ...block, text: EXPIRED_OFFLOADED_TOOL_RESULT_JSON }; - } - - return { ...block }; - }); + const content = message.content.map(sanitizeOffloadedContentBlock); return { ...message, content }; }); } +function sanitizeOffloadedContentBlock(block: MessageContent): MessageContent { + if (block.type === 'tool-call') { + if (block.state === 'resolved' && isOffloadedToolResult(block.output)) { + return { ...block, output: { ...EXPIRED_OFFLOADED_TOOL_RESULT } }; + } + if (block.state === 'resolved' && isContentToolResultOutput(block.output)) { + return { ...block, output: sanitizeOffloadedContentResult(block.output) }; + } + if (block.state === 'rejected' && isSerializedOffloadedToolResult(block.error)) { + return { ...block, error: EXPIRED_OFFLOADED_TOOL_RESULT_JSON }; + } + } + + if (block.type === 'text' && isSerializedOffloadedToolResult(block.text)) { + return { ...block, text: EXPIRED_OFFLOADED_TOOL_RESULT_JSON }; + } + + return { ...block }; +} + +function sanitizeOffloadedContentResult(output: ContentToolResultOutput): JSONValue { + return toJsonValue({ + type: 'content', + value: output.value.map((part) => { + if (part.type !== 'text' || !isSerializedOffloadedToolResult(part.text)) return part; + return { ...part, text: EXPIRED_OFFLOADED_TOOL_RESULT_JSON }; + }), + }); +} + async function tryOffloadResult( serialized: string, tokenCount: number,