From 6f2bc1398429b15b8d0efdce81164a5769a2dfb2 Mon Sep 17 00:00:00 2001 From: Robin Braumann <50590409+bjorger@users.noreply.github.com> Date: Fri, 18 Sep 2026 11:32:54 +0000 Subject: [PATCH] refactor(core): Share preview agent execution across callers (no-changelog) (#39064) --- .../__tests__/agent-chat.controller.test.ts | 296 +++++++++++------- ...ent-execution-orchestrator.service.test.ts | 56 ++++ .../agents/__tests__/agent-sse-stream.test.ts | 37 ++- .../__tests__/agent-test-run.service.test.ts | 137 +++++++- ...agent-workflow-tool-resume.service.test.ts | 10 +- .../modules/agents/agent-chat.controller.ts | 108 +++---- .../src/modules/agents/agent-sse-stream.ts | 105 +++---- .../modules/agents/agent-test-run.service.ts | 242 +++++++------- .../agents/agent-turn-execution.service.ts | 183 ++++++----- 9 files changed, 729 insertions(+), 445 deletions(-) diff --git a/packages/cli/src/modules/agents/__tests__/agent-chat.controller.test.ts b/packages/cli/src/modules/agents/__tests__/agent-chat.controller.test.ts index 7d76f4280a9..a235a1208de 100644 --- a/packages/cli/src/modules/agents/__tests__/agent-chat.controller.test.ts +++ b/packages/cli/src/modules/agents/__tests__/agent-chat.controller.test.ts @@ -1,4 +1,5 @@ import { EventEmitter } from 'node:events'; +import { createDeferredPromise } from '@n8n/utils/promise/deferred-promise'; import type { Mocked } from 'vitest'; import { mock } from 'vitest-mock-extended'; @@ -11,13 +12,15 @@ import type { AgentChatAttachmentService } from '../agent-chat-attachment.servic import { AgentChatController } from '../agent-chat.controller'; import type { AgentExecutionOrchestratorService } from '../agent-execution-orchestrator.service'; import type { AgentExecutionService } from '../agent-execution.service'; +import type { AgentValidationService } from '../agent-validation.service'; import type { AgentBackgroundJobService } from '../background/agent-background-job.service'; import type { AgentExecutionThread } from '../entities/agent-execution-thread.entity'; import type { FlushableResponse } from '../agent-sse-stream'; import type { AgentTestChatService } from '../agent-test-chat.service'; -import type { AgentTestRunService } from '../agent-test-run.service'; +import { AgentTestRunService } from '../agent-test-run.service'; import type { AgentsService } from '../agents.service'; import type { AgentsBuilderService } from '../builder/agents-builder.service'; +import type { N8NCheckpointStorage } from '../integrations/n8n-checkpoint-storage'; import { expectProjectScopedAgentRoutes, getRoutesByHandlerName, @@ -28,22 +31,17 @@ function makeController() { mock>(); const agentExecutionOrchestratorService = mock(); const agentsBuilderService = mock(); - const agentTestRunService = mock(); const agentChatAttachmentService = mock(); const agentExecutionService = mock(); + const agentValidationService = mock(); const backgroundJobService = mock(); - agentTestRunService.prepareDraftRun.mockResolvedValue({ - status: 'ready', - sessionId: 'thread-1', - }); - agentTestRunService.streamDraftRun.mockImplementation((config) => - agentExecutionOrchestratorService.executeForChat({ - ...config, - memory: { - threadId: config.sessionId, - resourceId: `draft-chat:${config.user.id}`, - }, - }), + agentExecutionService.findThreadById.mockResolvedValue(null); + agentValidationService.validateAgentIsRunnable.mockResolvedValue({ missing: [] }); + const agentTestRunService = new AgentTestRunService( + agentExecutionService, + agentValidationService, + agentExecutionOrchestratorService, + mock(), ); const controller = new AgentChatController( @@ -75,7 +73,10 @@ function makeController() { }; } -function makeSseResponse(writes: string[]): FlushableResponse { +function makeSseResponse( + writes: string[], + beforeWrite?: (chunk: string) => void, +): FlushableResponse { const emitter = new EventEmitter(); const res = Object.assign(emitter, { socket: { @@ -86,16 +87,19 @@ function makeSseResponse(writes: string[]): FlushableResponse { setHeader: vi.fn(), flushHeaders: vi.fn(), write: vi.fn((chunk: string) => { + beforeWrite?.(chunk); writes.push(String(chunk)); return true; }), flush: vi.fn(), - end: vi.fn(() => { - emitter.emit('finish'); - }), + end: vi.fn(), writableEnded: false, destroyed: false, }); + res.end.mockImplementation(() => { + res.writableEnded = true; + emitter.emit('finish'); + }); return res as unknown as FlushableResponse; } @@ -319,10 +323,37 @@ describe('AgentChatController chat message history', () => { }); describe('AgentChatController SSE done payload', () => { + const operations = [ + { + name: 'new chat', + method: 'executeForChat' as const, + start: async (controller: AgentChatController, res: FlushableResponse) => + await controller.chat( + { params: { projectId: 'project-1' }, user: { id: 'user-1' } } as never, + res, + 'agent-1', + { message: 'hi', sessionId: 'thread-1' } as never, + ), + done: { type: 'done', sessionId: 'thread-1', executionId: 'exec-99' }, + }, + { + name: 'resumed chat', + method: 'resumeForChat' as const, + start: async (controller: AgentChatController, res: FlushableResponse) => + await controller.chatResume( + { params: { projectId: 'project-1' }, user: { id: 'user-1' } } as never, + res, + 'agent-1', + { runId: 'run-1', toolCallId: 'tc-1', resumeData: { approved: true } } as never, + ), + done: { type: 'done', executionId: 'exec-99' }, + }, + ]; + it('does not start a chat after the response closes during preparation', async () => { - const { controller, agentTestRunService } = makeController(); + const { controller, agentTestRunService, agentExecutionOrchestratorService } = makeController(); let resolvePreparation = (_value: { status: 'ready'; sessionId: string }) => {}; - agentTestRunService.prepareDraftRun.mockReturnValue( + const prepareDraftRun = vi.spyOn(agentTestRunService, 'prepareDraftRun').mockReturnValue( new Promise((resolve) => { resolvePreparation = resolve; }), @@ -335,24 +366,118 @@ describe('AgentChatController SSE done payload', () => { 'agent-1', { message: 'hi' } as never, ); - await vi.waitFor(() => expect(agentTestRunService.prepareDraftRun).toHaveBeenCalled()); + await vi.waitFor(() => expect(prepareDraftRun).toHaveBeenCalled()); (res as unknown as EventEmitter).emit('close'); resolvePreparation({ status: 'ready', sessionId: 'thread-1' }); await request; - expect(agentTestRunService.streamDraftRun).not.toHaveBeenCalled(); + expect(agentExecutionOrchestratorService.executeForChat).not.toHaveBeenCalled(); }); - it('includes executionId on done when recorded', async () => { + it.each(operations)('sends done after the $name settles', async ({ method, start, done }) => { const { controller, agentExecutionOrchestratorService } = makeController(); - agentExecutionOrchestratorService.executeForChat.mockImplementation(async function* (config) { + const finalization = createDeferredPromise(); + const finalizationStarted = createDeferredPromise(); + let receivedSignal: AbortSignal | undefined; + agentExecutionOrchestratorService[method].mockImplementation(async function* (config) { + receivedSignal = config.abortSignal; + yield { type: 'finish', finishReason: 'stop' }; + finalizationStarted.resolve(); + await finalization.promise; config.onExecutionRecorded?.('exec-99'); - yield* []; }); const writes: string[] = []; const res = makeSseResponse(writes); + const request = start(controller, res); + await finalizationStarted.promise; + expect(res.end).not.toHaveBeenCalled(); + expect(writes.some((line) => line.includes('"done"'))).toBe(false); + finalization.resolve(); + await request; + + const events = writes + .filter((line) => line.startsWith('data: ')) + .map((line) => JSON.parse(line.slice(6).trim()) as { type: string }); + + expect(events).toContainEqual(done); + expect(res.end).toHaveBeenCalledOnce(); + res.emit('close'); + expect(receivedSignal?.aborted).toBe(false); + }); + + it.each(operations)('does not send done for a suspended $name', async ({ method, start }) => { + const { controller, agentExecutionOrchestratorService } = makeController(); + agentExecutionOrchestratorService[method].mockImplementation(async function* () { + yield { + type: 'tool-call-suspended', + runId: 'run-1', + toolCallId: 'tc-1', + toolName: 'ask_questions', + }; + }); + const writes: string[] = []; + + await start(controller, makeSseResponse(writes)); + + const eventTypes = writes + .filter((line) => line.startsWith('data: ')) + .map((line) => (JSON.parse(line.slice(6).trim()) as { type: string }).type); + expect(eventTypes).toContain('tool-call-suspended'); + expect(eventTypes).not.toContain('done'); + }); + + it.each(operations)( + 'aborts the $name when its SSE response closes early', + async ({ start, method }) => { + const { controller, agentExecutionOrchestratorService } = makeController(); + + let receivedSignal: AbortSignal | undefined; + let releaseRun = () => {}; + let markStarted = () => {}; + const runStarted = new Promise((resolve) => { + markStarted = resolve; + }); + const runBlocked = new Promise((resolve) => { + releaseRun = resolve; + }); + agentExecutionOrchestratorService[method].mockImplementation(async function* (config) { + receivedSignal = (config as { abortSignal?: AbortSignal }).abortSignal; + markStarted(); + await runBlocked; + yield* []; + }); + + const res = makeSseResponse([]); + const request = start(controller, res); + await runStarted; + (res as unknown as EventEmitter).emit('close'); + releaseRun(); + await request; + + expect(receivedSignal?.aborted).toBe(true); + }, + ); + + it('settles execution after SSE delivery fails', async () => { + const { controller, agentExecutionOrchestratorService } = makeController(); + const deliveryError = new Error('socket write failed'); + let receivedSignal: AbortSignal | undefined; + const lifecycle: string[] = []; + agentExecutionOrchestratorService.executeForChat.mockImplementation(async function* (config) { + receivedSignal = config.abortSignal; + try { + yield { type: 'text-delta', id: 'text-1', delta: 'first' }; + lifecycle.push('continued'); + yield { type: 'text-delta', id: 'text-1', delta: 'second' }; + } finally { + lifecycle.push('settled'); + } + }); + const res = makeSseResponse([], (chunk) => { + if (chunk.includes('"text-delta"')) throw deliveryError; + }); await controller.chat( { params: { projectId: 'project-1' }, user: { id: 'user-1' } } as never, @@ -361,94 +486,9 @@ describe('AgentChatController SSE done payload', () => { { message: 'hi', sessionId: 'thread-1' } as never, ); - const events = writes - .filter((line) => line.startsWith('data: ')) - .map((line) => JSON.parse(line.slice(6).trim()) as { type: string }); - - expect(events).toContainEqual({ - type: 'done', - sessionId: 'thread-1', - executionId: 'exec-99', - }); - }); - - it('includes executionId on resume done when recorded', async () => { - const { controller, agentExecutionOrchestratorService } = makeController(); - agentExecutionOrchestratorService.resumeForChat.mockImplementation(async function* (config) { - config.onExecutionRecorded?.('exec-resume-1'); - yield* []; - }); - - const writes: string[] = []; - const res = makeSseResponse(writes); - - await controller.chatResume( - { params: { projectId: 'project-1' }, user: { id: 'user-1' } } as never, - res, - 'agent-1', - { runId: 'run-1', toolCallId: 'tc-1', resumeData: { approved: true } } as never, - ); - - const events = writes - .filter((line) => line.startsWith('data: ')) - .map((line) => JSON.parse(line.slice(6).trim()) as { type: string }); - - expect(events).toContainEqual({ - type: 'done', - executionId: 'exec-resume-1', - }); - }); - - it.each([ - { - name: 'new chat', - start: async (controller: AgentChatController, res: FlushableResponse) => - await controller.chat( - { params: { projectId: 'project-1' }, user: { id: 'user-1' } } as never, - res, - 'agent-1', - { message: 'hi' } as never, - ), - method: 'executeForChat' as const, - }, - { - name: 'resumed chat', - start: async (controller: AgentChatController, res: FlushableResponse) => - await controller.chatResume( - { params: { projectId: 'project-1' }, user: { id: 'user-1' } } as never, - res, - 'agent-1', - { runId: 'run-1', toolCallId: 'tc-1', resumeData: { approved: true } } as never, - ), - method: 'resumeForChat' as const, - }, - ])('aborts the $name when its SSE response closes early', async ({ start, method }) => { - const { controller, agentExecutionOrchestratorService } = makeController(); - - let receivedSignal: AbortSignal | undefined; - let releaseRun = () => {}; - let markStarted = () => {}; - const runStarted = new Promise((resolve) => { - markStarted = resolve; - }); - const runBlocked = new Promise((resolve) => { - releaseRun = resolve; - }); - agentExecutionOrchestratorService[method].mockImplementation(async function* (config) { - receivedSignal = (config as { abortSignal?: AbortSignal }).abortSignal; - markStarted(); - await runBlocked; - yield* []; - }); - - const res = makeSseResponse([]); - const request = start(controller, res); - await runStarted; - (res as unknown as EventEmitter).emit('close'); - releaseRun(); - await request; - expect(receivedSignal?.aborted).toBe(true); + expect(lifecycle).toEqual(['continued', 'settled']); + expect(res.end).toHaveBeenCalledTimes(1); }); }); @@ -575,6 +615,38 @@ describe('AgentChatController attachment cleanup on failed turns', () => { expect(agentExecutionOrchestratorService.executeForChat).not.toHaveBeenCalled(); }); + it('deletes stored attachments when the response closes during upload', async () => { + const { controller, agentExecutionOrchestratorService, agentChatAttachmentService } = + makeController(); + let finishUpload = (_value: unknown) => {}; + agentChatAttachmentService.storeInbound.mockReturnValue( + new Promise((resolve) => { + finishUpload = resolve; + }), + ); + agentChatAttachmentService.deleteByIds.mockResolvedValue(undefined); + const { res } = makeCleanupSseResponse(); + const request = controller.chat( + { params: { projectId: 'project-1' }, user: { id: 'user-1' } } as never, + res, + 'agent-1', + { message: 'hi', attachments: [textAttachment('notes.txt')] } as never, + ); + await vi.waitFor(() => expect(agentChatAttachmentService.storeInbound).toHaveBeenCalled()); + + (res as unknown as EventEmitter).emit('close'); + finishUpload({ + id: 'att-1', + fileName: 'notes.txt', + mimeType: 'text/plain', + fileSizeBytes: 5, + }); + await request; + + expect(agentChatAttachmentService.deleteByIds).toHaveBeenCalledWith(['att-1']); + expect(agentExecutionOrchestratorService.executeForChat).not.toHaveBeenCalled(); + }); + it('rejects an empty attachment with a dedicated error message', async () => { const { controller, agentChatAttachmentService } = makeController(); const { res, events } = makeCleanupSseResponse(); diff --git a/packages/cli/src/modules/agents/__tests__/agent-execution-orchestrator.service.test.ts b/packages/cli/src/modules/agents/__tests__/agent-execution-orchestrator.service.test.ts index 118a3f37bf8..91eed8f21a5 100644 --- a/packages/cli/src/modules/agents/__tests__/agent-execution-orchestrator.service.test.ts +++ b/packages/cli/src/modules/agents/__tests__/agent-execution-orchestrator.service.test.ts @@ -26,7 +26,9 @@ import { AgentExecutionOrchestratorService } from '../agent-execution-orchestrat import type { AgentExecutionService } from '../agent-execution.service'; import type { AgentRunTracingService } from '../agent-run-tracing.service'; import type { AgentRuntimeCacheService } from '../agent-runtime-cache.service'; +import { AgentTestRunService } from '../agent-test-run.service'; import { AgentTurnExecutionService } from '../agent-turn-execution.service'; +import type { AgentValidationService } from '../agent-validation.service'; import type { Agent } from '../entities/agent.entity'; import type { AgentRepository } from '../repositories/agent.repository'; import { @@ -364,6 +366,60 @@ describe('AgentExecutionOrchestratorService', () => { ); }); + it('settles prepared execution after finalization and runtime release', async () => { + const { + service, + executionService, + checkpointStorage, + runtimeCacheService, + runtime, + onExecutionRecorded, + } = makeTurn(); + const testRunService = new AgentTestRunService( + executionService, + mock(), + service, + checkpointStorage, + ); + const finalization = createDeferredPromise(); + executionService.finalizeExecution.mockReturnValue(finalization.promise); + const onChunk = vi.fn(); + const input = { agentId, projectId, user, onChunk, onExecutionRecorded }; + let settled = false; + const execution = ( + operation === 'start' + ? testRunService.executePreparedDraftRun({ + ...input, + message: 'hello', + sessionId: 'thread-1', + }) + : testRunService.resumePreparedDraftRun({ + ...input, + runId: 'run-1', + toolCallId: 'tc-1', + resumeData: { approved: true }, + }) + ).then((result) => { + settled = true; + expect(runtimeCacheService.releaseRuntimeLease).toHaveBeenCalledWith(runtime.agent); + return result; + }); + await vi.waitFor(() => expect(executionService.finalizeExecution).toHaveBeenCalled()); + + expect(onChunk).toHaveBeenCalledWith({ type: 'finish', finishReason: 'stop' }); + expect(settled).toBe(false); + expect(onExecutionRecorded).not.toHaveBeenCalled(); + expect(runtimeCacheService.releaseRuntimeLease).not.toHaveBeenCalled(); + + finalization.resolve('finalized-execution'); + await expect(execution).resolves.toEqual({ + status: 'completed', + response: '', + executionId: 'finalized-execution', + }); + expect(onExecutionRecorded).toHaveBeenCalledWith('finalized-execution'); + }); + it('cancels an early-closed stream and finalizes before callbacks and lease release', async () => { const { stream, diff --git a/packages/cli/src/modules/agents/__tests__/agent-sse-stream.test.ts b/packages/cli/src/modules/agents/__tests__/agent-sse-stream.test.ts index 5e1ab7dac39..a7ac10cc148 100644 --- a/packages/cli/src/modules/agents/__tests__/agent-sse-stream.test.ts +++ b/packages/cli/src/modules/agents/__tests__/agent-sse-stream.test.ts @@ -3,7 +3,7 @@ import type { AgentSseEvent } from '@n8n/api-types'; import { LoggerProxy } from 'n8n-workflow'; import { EventEmitter } from 'node:events'; -import { initSseStream, pumpChunks, type FlushableResponse } from '../agent-sse-stream'; +import { emitChunkEvents, initSseStream, type FlushableResponse } from '../agent-sse-stream'; // --------------------------------------------------------------------------- // Helpers @@ -17,7 +17,9 @@ async function* toAsyncIterable(items: T[]): AsyncIterable { async function collectEvents(chunks: StreamChunk[]): Promise { const events: AgentSseEvent[] = []; - await pumpChunks(toAsyncIterable(chunks), (e) => events.push(e)); + for await (const chunk of toAsyncIterable(chunks)) { + emitChunkEvents(chunk, (event) => events.push(event)); + } return events; } @@ -32,6 +34,7 @@ function createResponse() { flushHeaders: vi.fn(), write: vi.fn(), flush: vi.fn(), + end: vi.fn(), socket, writableEnded: false, destroyed: false, @@ -43,7 +46,7 @@ function createResponse() { async function collectSerializedEvents(chunks: StreamChunk[]): Promise { const { res } = createResponse(); const { send } = initSseStream(res); - await pumpChunks(toAsyncIterable(chunks), send); + for await (const chunk of toAsyncIterable(chunks)) emitChunkEvents(chunk, send); const events = vi.mocked(res.write).mock.calls.flatMap(([payload]) => { if (typeof payload !== 'string' || !payload.startsWith('data: ')) return []; @@ -99,10 +102,25 @@ describe('agent-sse-stream — connection setup', () => { expect(res.write).not.toHaveBeenCalled(); }); + + it('closes completed delivery without cancelling execution', () => { + vi.useFakeTimers(); + const { res } = createResponse(); + const { close, abortSignal } = initSseStream(res); + + close(); + res.emit('close'); + vi.mocked(res.write).mockClear(); + vi.advanceTimersByTime(30_000); + + expect(abortSignal.aborted).toBe(false); + expect(res.end).toHaveBeenCalledOnce(); + expect(res.write).not.toHaveBeenCalled(); + }); }); // --------------------------------------------------------------------------- -// stringifyError — tested through pumpChunks / emitChunkEvents +// stringifyError — tested through emitChunkEvents // --------------------------------------------------------------------------- vi.mock('n8n-workflow', () => ({ @@ -111,7 +129,7 @@ vi.mock('n8n-workflow', () => ({ }, })); -describe('agent-sse-stream — stringifyError (via pumpChunks error chunk)', () => { +describe('agent-sse-stream — stringifyError (via error chunk)', () => { it('extracts .message from an Error instance', async () => { const events = await collectEvents([{ type: 'error', error: new Error('something broke') }]); expect(events).toEqual([{ type: 'error', message: 'something broke' }]); @@ -201,7 +219,7 @@ describe('agent-sse-stream — stream completion', () => { ]); }); - it('completes after the runtime stream closes even when a finish chunk is present', async () => { + it('leaves completion delivery to the caller when it receives a finish chunk', async () => { const events = await collectEvents([ { type: 'text-delta', id: 't-1', delta: 'hello' }, { type: 'text-end', id: 't-1' }, @@ -214,7 +232,7 @@ describe('agent-sse-stream — stream completion', () => { ]); }); - it('drains every suspension chunk before reporting that the run paused', async () => { + it('maps every suspension chunk', async () => { const chunks: StreamChunk[] = [ { type: 'tool-call-suspended', @@ -232,11 +250,8 @@ describe('agent-sse-stream — stream completion', () => { }, { type: 'finish', finishReason: 'other' }, ]; - const events: AgentSseEvent[] = []; + const events = await collectEvents(chunks); - const suspended = await pumpChunks(toAsyncIterable(chunks), (event) => events.push(event)); - - expect(suspended).toBe(true); expect(events).toEqual([ { type: 'tool-call-suspended', diff --git a/packages/cli/src/modules/agents/__tests__/agent-test-run.service.test.ts b/packages/cli/src/modules/agents/__tests__/agent-test-run.service.test.ts index 1dc61db4914..f3059a567de 100644 --- a/packages/cli/src/modules/agents/__tests__/agent-test-run.service.test.ts +++ b/packages/cli/src/modules/agents/__tests__/agent-test-run.service.test.ts @@ -19,6 +19,13 @@ const agentId = 'agent-1'; const projectId = 'project-1'; const user = mock({ id: 'user-1' }); const credentialProvider = mock(); +const preparedDraftRunInput = { + agentId, + projectId, + message: 'Hi', + user, + sessionId: 'session-1', +}; const approvalResumeSchema = (() => { const schema = zodToJsonSchema(APPROVAL_RESUME_SCHEMA); if (schema === null) throw new Error('Failed to generate approval resume schema'); @@ -77,6 +84,110 @@ function makeService() { } describe('AgentTestRunService', () => { + it('executes a prepared draft run without repeating preparation', async () => { + const { service, agentExecutionOrchestratorService, agentValidationService } = makeService(); + const executionError = new Error('streamed execution failed'); + const chunks: StreamChunk[] = [ + { type: 'text-delta', id: 'text-1', delta: 'Hello' }, + { type: 'error', error: executionError }, + { type: 'text-delta', id: 'text-1', delta: ' there' }, + ]; + agentExecutionOrchestratorService.executeForChat.mockImplementation(async function* (config) { + yield* chunks; + config.onExecutionRecorded?.('execution-1'); + }); + const onChunk = vi.fn(); + + await expect( + service.executePreparedDraftRun({ + ...preparedDraftRunInput, + attachments: [ + { id: 'attachment-1', fileName: 'notes.txt', mimeType: 'text/plain', sizeBytes: 5 }, + ], + errorMode: 'forward', + onChunk, + }), + ).resolves.toEqual({ + status: 'completed', + response: 'Hello there', + executionId: 'execution-1', + }); + expect(agentValidationService.validateAgentIsRunnable).not.toHaveBeenCalled(); + expect(onChunk).toHaveBeenCalledWith({ type: 'error', error: executionError }); + expect(onChunk).toHaveBeenCalledTimes(3); + expect(agentExecutionOrchestratorService.executeForChat).toHaveBeenCalledWith( + expect.objectContaining({ + memory: { threadId: 'session-1', resourceId: 'draft-chat:user-1' }, + attachments: [ + { id: 'attachment-1', fileName: 'notes.txt', mimeType: 'text/plain', sizeBytes: 5 }, + ], + }), + ); + }); + + it('resumes a prepared draft run with its selected memory scope', async () => { + const { service, agentExecutionOrchestratorService, agentExecutionService } = makeService(); + agentExecutionOrchestratorService.resumeForChat.mockImplementation(async function* (config) { + yield { type: 'text-delta', id: 'text-1', delta: 'continued' }; + config.onExecutionRecorded?.('execution-2'); + }); + + await expect( + service.resumePreparedDraftRun({ + agentId, + projectId, + runId: 'run-1', + toolCallId: 'tool-call-1', + resumeData: { approved: true }, + user, + expectedMemory: { threadId: 'session-1', resourceId: 'draft-chat:user-1' }, + initialResponse: 'Already ', + }), + ).resolves.toEqual({ + status: 'completed', + response: 'Already continued', + executionId: 'execution-2', + }); + expect(agentExecutionService.findThreadById).not.toHaveBeenCalled(); + expect(agentExecutionOrchestratorService.resumeForChat).toHaveBeenCalledWith( + expect.objectContaining({ + runId: 'run-1', + toolCallId: 'tool-call-1', + expectedMemory: { threadId: 'session-1', resourceId: 'draft-chat:user-1' }, + }), + ); + }); + + it('reports an observer failure after the execution stream settles', async () => { + const { service, agentExecutionOrchestratorService } = makeService(); + const abortController = new AbortController(); + const deliveryError = new Error('delivery failed'); + const produced: string[] = []; + agentExecutionOrchestratorService.executeForChat.mockImplementation(async function* () { + try { + yield { type: 'text-delta', id: 'text-1', delta: 'first' }; + produced.push('continued'); + yield { type: 'text-delta', id: 'text-1', delta: 'second' }; + } finally { + produced.push('settled'); + } + }); + const onChunk = vi.fn(() => { + throw deliveryError; + }); + + await expect( + service.executePreparedDraftRun({ + ...preparedDraftRunInput, + onChunk, + abortSignal: abortController.signal, + }), + ).rejects.toBe(deliveryError); + expect(onChunk).toHaveBeenCalledTimes(1); + expect(produced).toEqual(['continued', 'settled']); + expect(abortController.signal.aborted).toBe(false); + }); + it('runs a draft test and returns its response and execution identifiers', async () => { const { service, agentExecutionOrchestratorService } = makeService(); agentExecutionOrchestratorService.executeForChat.mockImplementation(async function* (config) { @@ -354,22 +465,24 @@ describe('AgentTestRunService', () => { expect(agentExecutionOrchestratorService.executeForChat).not.toHaveBeenCalled(); }); - it('propagates a streamed execution error instead of completing the draft run', async () => { + it.each(['headless', 'observed'])('rejects a streamed error for an %s run', async (caller) => { const { service, agentExecutionOrchestratorService } = makeService(); const executionError = new Error('streamed execution failed'); + let settled = false; agentExecutionOrchestratorService.executeForChat.mockImplementation(async function* () { - yield { type: 'error', error: executionError }; - yield { type: 'finish', finishReason: 'error' }; + try { + yield { type: 'error', error: executionError }; + yield { type: 'finish', finishReason: 'error' }; + } finally { + settled = true; + } }); - await expect( - service.executeDraftRun({ - agentId, - projectId, - message: 'Hi', - user, - credentialProvider, - }), - ).rejects.toBe(executionError); + const execution = + caller === 'headless' + ? service.executeDraftRun({ ...preparedDraftRunInput, credentialProvider }) + : service.executePreparedDraftRun({ ...preparedDraftRunInput, onChunk: vi.fn() }); + await expect(execution).rejects.toBe(executionError); + expect(settled).toBe(true); }); }); diff --git a/packages/cli/src/modules/agents/__tests__/agent-workflow-tool-resume.service.test.ts b/packages/cli/src/modules/agents/__tests__/agent-workflow-tool-resume.service.test.ts index 18b7ba80cfc..d87ebc52142 100644 --- a/packages/cli/src/modules/agents/__tests__/agent-workflow-tool-resume.service.test.ts +++ b/packages/cli/src/modules/agents/__tests__/agent-workflow-tool-resume.service.test.ts @@ -2,6 +2,7 @@ import { N8N_CHAT_INTEGRATION_TYPE } from '@n8n/api-types'; import type { Logger } from '@n8n/backend-common'; import type { User, UserRepository } from '@n8n/db'; import type { WorkflowExecuteAfterContext } from '@n8n/decorators'; +import { createDeferredPromise } from '@n8n/utils/promise/deferred-promise'; import type { InstanceSettings } from 'n8n-core'; import type { IRun, RelatedAgentRun } from 'n8n-workflow'; import { createRunExecutionData } from 'n8n-workflow'; @@ -393,9 +394,14 @@ describe('AgentWorkflowToolResumeService → preview chat', () => { ])('pushes the recorded execution after %s', async (_label, result) => { const { service, userRepository, agentTestRunService, broadcaster } = setup(); userRepository.findOneBy.mockResolvedValue(mock({ id: 'user-1' })); - agentTestRunService.resumeDraftRun.mockResolvedValue(result); + const execution = createDeferredPromise(); + agentTestRunService.resumeDraftRun.mockReturnValue(execution.promise); - await service.resume(previewRun, 'success'); + const resume = service.resume(previewRun, 'success'); + await vi.waitFor(() => expect(agentTestRunService.resumeDraftRun).toHaveBeenCalled()); + expect(broadcaster.notify).not.toHaveBeenCalled(); + execution.resolve(result); + await resume; expect(broadcaster.notify).toHaveBeenCalledWith({ projectId: 'project-1', diff --git a/packages/cli/src/modules/agents/agent-chat.controller.ts b/packages/cli/src/modules/agents/agent-chat.controller.ts index 03054700fac..f7bb923f941 100644 --- a/packages/cli/src/modules/agents/agent-chat.controller.ts +++ b/packages/cli/src/modules/agents/agent-chat.controller.ts @@ -6,7 +6,6 @@ import { AgentChatResumeDto, MAX_AGENT_CHAT_ATTACHMENT_SIZE_BYTES, MAX_AGENT_CHAT_ATTACHMENT_SIZE_MB, - N8N_CHAT_INTEGRATION_TYPE, ViewableMimeTypes, } from '@n8n/api-types'; import type { AuthenticatedRequest } from '@n8n/db'; @@ -29,7 +28,7 @@ import { import { AgentExecutionOrchestratorService } from './agent-execution-orchestrator.service'; import { AgentExecutionService, threadBelongsTo } from './agent-execution.service'; import { messagesToDto } from './agent-message-mapper'; -import { type FlushableResponse, initSseStream, pumpChunks } from './agent-sse-stream'; +import { type FlushableResponse, initSseStream } from './agent-sse-stream'; import { AgentTestChatService, chatThreadId } from './agent-test-chat.service'; import { AgentTestRunService } from './agent-test-run.service'; import { AgentsService } from './agents.service'; @@ -121,10 +120,7 @@ export class AgentChatController { req.user, ); - const { send } = initSseStream(res); - const abortController = new AbortController(); - const abortOnClose = () => abortController.abort(); - res.once('close', abortOnClose); + const { send, onChunk, abortSignal, close } = initSseStream(res); let executionId: string | undefined; let storedAttachments: StoredAttachmentRef[] | undefined; try { @@ -134,7 +130,7 @@ export class AgentChatController { sessionId, credentialProvider, }); - if (abortController.signal.aborted) return; + if (abortSignal.aborted) return; if (prepared.status === 'session_not_found') { send({ type: 'error', message: 'Session not found' }); return; @@ -157,24 +153,25 @@ export class AgentChatController { threadId, resourceId: draftChatMemoryResourceId(req.user.id), }); + abortSignal.throwIfAborted(); - const suspended = await pumpChunks( - this.agentTestRunService.streamDraftRun({ - agentId, - projectId, - message, - attachments: storedAttachments, - user: req.user, - sessionId: threadId, - previewChat: true, - onExecutionRecorded: (id) => { - executionId = id; - }, - abortSignal: abortController.signal, - }), - send, - ); - if (!suspended) { + const result = await this.agentTestRunService.executePreparedDraftRun({ + agentId, + projectId, + message, + attachments: storedAttachments, + user: req.user, + sessionId: threadId, + previewChat: true, + errorMode: 'forward', + onChunk, + onExecutionRecorded: (id) => { + executionId = id; + }, + abortSignal, + }); + executionId = result.executionId ?? executionId; + if (result.status === 'completed') { send({ type: 'done', sessionId: threadId, ...(executionId ? { executionId } : {}) }); } } catch (error) { @@ -186,13 +183,10 @@ export class AgentChatController { .deleteByIds(storedAttachments.map((ref) => ref.id)) .catch(() => {}); } - if (!abortController.signal.aborted) { - const errorMessage = error instanceof Error ? error.message : 'Chat failed'; - send({ type: 'error', message: errorMessage }); - } + const errorMessage = error instanceof Error ? error.message : 'Chat failed'; + send({ type: 'error', message: errorMessage }); } finally { - res.off('close', abortOnClose); - res.end(); + close(); } } @@ -206,42 +200,32 @@ export class AgentChatController { ) { const { projectId } = req.params; const { runId, toolCallId, resumeData } = payload; - const { send } = initSseStream(res); - - const abortController = new AbortController(); - const abortOnClose = () => abortController.abort(); - res.once('close', abortOnClose); + const { send, onChunk, abortSignal, close } = initSseStream(res); try { - let executionId: string | undefined; - const suspended = await pumpChunks( - this.agentExecutionOrchestratorService.resumeForChat({ - agentId, - projectId, - runId, - toolCallId, - resumeData, - user: req.user, - usePublishedVersion: false, - integrationType: N8N_CHAT_INTEGRATION_TYPE, - previewChat: true, - onExecutionRecorded: (id) => { - executionId = id; - }, - abortSignal: abortController.signal, - }), - send, - ); - if (!suspended) { - send({ type: 'done', ...(executionId ? { executionId } : {}) }); + abortSignal.throwIfAborted(); + const result = await this.agentTestRunService.resumePreparedDraftRun({ + agentId, + projectId, + runId, + toolCallId, + resumeData, + user: req.user, + previewChat: true, + errorMode: 'forward', + onChunk, + abortSignal, + }); + if (result.status === 'completed') { + send({ + type: 'done', + ...(result.executionId ? { executionId: result.executionId } : {}), + }); } } catch (error) { - if (!abortController.signal.aborted) { - const errorMessage = error instanceof Error ? error.message : 'Resume failed'; - send({ type: 'error', message: errorMessage }); - } + const errorMessage = error instanceof Error ? error.message : 'Resume failed'; + send({ type: 'error', message: errorMessage }); } finally { - res.off('close', abortOnClose); - res.end(); + close(); } } diff --git a/packages/cli/src/modules/agents/agent-sse-stream.ts b/packages/cli/src/modules/agents/agent-sse-stream.ts index 168058e2c88..12844ef7f12 100644 --- a/packages/cli/src/modules/agents/agent-sse-stream.ts +++ b/packages/cli/src/modules/agents/agent-sse-stream.ts @@ -13,13 +13,7 @@ export type FlushableResponse = Response & { flush?: () => void }; const SSE_HEARTBEAT_INTERVAL_MS = 30_000; -interface ChunkHandlerCtx { - send: (e: AgentSseEvent) => void; -} - -/** - * Set up SSE headers and return a typed `send(event)` helper. - */ +/** Set up preview delivery and request cancellation when the connection closes. */ export function initSseStream(res: FlushableResponse) { res.setHeader('Content-Type', 'text/event-stream; charset=UTF-8'); res.setHeader('Cache-Control', 'no-cache, no-transform'); @@ -43,12 +37,33 @@ export function initSseStream(res: FlushableResponse) { res.once('finish', stopHeartbeat); res.once('close', stopHeartbeat); + const abortController = new AbortController(); + const abortOnClose = () => abortController.abort(); + res.once('close', abortOnClose); + const close = () => { + res.off('close', abortOnClose); + stopHeartbeat(); + if (!res.writableEnded && !res.destroyed) res.end(); + }; + const send = (event: AgentSseEvent) => { + if (abortController.signal.aborted) return; res.write(`data: ${JSON.stringify(event)}\n\n`); res.flush?.(); }; - return { send }; + const onChunk = (chunk: StreamChunk) => { + if (abortController.signal.aborted) return; + try { + emitChunkEvents(chunk, send); + } catch (error) { + abortController.abort(); + close(); + throw error; + } + }; + + return { send, onChunk, abortSignal: abortController.signal, close }; } function toAgentSseMessage(message: AgentMessage): AgentSseMessage | undefined { @@ -111,10 +126,7 @@ function emitTextLikeChunk( } } -/** - * SSE-emit a tool-* chunk and fire any matching builder side-effect callback. - * Returns `{ suspended: true }` when the chunk was `tool-call-suspended`. - */ +/** Map tool chunks to SSE events. */ function emitToolChunk( chunk: Extract< StreamChunk, @@ -129,10 +141,8 @@ function emitToolChunk( | 'tool-call-suspended'; } >, - ctx: ChunkHandlerCtx, -): { suspended: boolean } { - const { send } = ctx; - + send: (e: AgentSseEvent) => void, +): void { switch (chunk.type) { case 'tool-input-start': send({ @@ -191,34 +201,30 @@ function emitToolChunk( input: chunk.suspendPayload, }; send({ type: 'tool-call-suspended', payload }); - return { suspended: true }; + break; } } - return { suspended: false }; } /** * Translate a single chunk into one or more SSE events. - * - * Returns `{ suspended: true }` when the chunk was a `tool-call-suspended` - * and `{ suspended: false }` for all other chunks. */ -function emitChunkEvents(chunk: StreamChunk, ctx: ChunkHandlerCtx): { suspended: boolean } { +export function emitChunkEvents(chunk: StreamChunk, send: (event: AgentSseEvent) => void): void { switch (chunk.type) { case 'start-step': - ctx.send({ type: 'start-step' }); - return { suspended: false }; + send({ type: 'start-step' }); + return; case 'finish-step': - ctx.send({ type: 'finish-step' }); - return { suspended: false }; + send({ type: 'finish-step' }); + return; case 'text-start': case 'text-delta': case 'text-end': case 'reasoning-start': case 'reasoning-delta': case 'reasoning-end': - emitTextLikeChunk(chunk, ctx.send); - return { suspended: false }; + emitTextLikeChunk(chunk, send); + return; case 'tool-input-start': case 'tool-input-delta': case 'tool-call': @@ -226,39 +232,40 @@ function emitChunkEvents(chunk: StreamChunk, ctx: ChunkHandlerCtx): { suspended: case 'tool-execution-end': case 'tool-result': case 'tool-call-suspended': - return emitToolChunk(chunk, ctx); + emitToolChunk(chunk, send); + return; case 'message': { const sseMessage = toAgentSseMessage(chunk.message); - if (sseMessage) ctx.send({ type: 'message', message: sseMessage }); - return { suspended: false }; + if (sseMessage) send({ type: 'message', message: sseMessage }); + return; } case 'subagent-chunk': { - if (chunk.parentToolCallId === undefined) return { suspended: false }; - ctx.send({ + if (chunk.parentToolCallId === undefined) return; + send({ type: 'subagent-chunk', parentToolCallId: chunk.parentToolCallId, taskPath: chunk.taskPath, chunk: chunk.chunk, }); - return { suspended: false }; + return; } case 'error': { const errMsg = stringifyError(chunk.error); - ctx.send({ type: 'error', message: errMsg }); - return { suspended: false }; + send({ type: 'error', message: errMsg }); + return; } case 'warning': { - ctx.send({ + send({ type: 'warning', message: chunk.message, ...(chunk.code !== undefined && { code: chunk.code }), ...(chunk.source !== undefined && { source: chunk.source }), ...(chunk.server !== undefined && { server: chunk.server }), }); - return { suspended: false }; + return; } default: - return { suspended: false }; + return; } } @@ -302,23 +309,3 @@ function stringifyError(error: unknown): string { } return 'Unknown error'; } - -/** - * Pump SDK stream chunks through a typed AgentSseEvent stream. - * - * Returns `true` when a suspension was emitted (the run paused), `false` - * otherwise. - */ -export async function pumpChunks( - chunks: AsyncIterable, - send: (e: AgentSseEvent) => void, -): Promise { - const ctx: ChunkHandlerCtx = { send }; - let suspended = false; - - for await (const chunk of chunks) { - const result = emitChunkEvents(chunk, ctx); - suspended ||= result.suspended; - } - return suspended; -} diff --git a/packages/cli/src/modules/agents/agent-test-run.service.ts b/packages/cli/src/modules/agents/agent-test-run.service.ts index 3be4be1db15..bd5c38a5f1c 100644 --- a/packages/cli/src/modules/agents/agent-test-run.service.ts +++ b/packages/cli/src/modules/agents/agent-test-run.service.ts @@ -17,8 +17,11 @@ import { Service } from '@n8n/di'; import { UserError } from 'n8n-workflow'; import { z } from 'zod'; -import type { StoredAttachmentRef } from './agent-chat-attachment.service'; -import { AgentExecutionOrchestratorService } from './agent-execution-orchestrator.service'; +import { + AgentExecutionOrchestratorService, + type ExecuteForChatConfig, + type ResumeForChatConfig, +} from './agent-execution-orchestrator.service'; import { AgentExecutionService, threadBelongsTo } from './agent-execution.service'; import { AgentValidationService } from './agent-validation.service'; import { N8NCheckpointStorage } from './integrations/n8n-checkpoint-storage'; @@ -36,40 +39,35 @@ export type PrepareDraftRunResult = | { status: 'session_not_found' } | { status: 'agent_misconfigured'; missing: string[] }; -interface StreamDraftRunInput { - agentId: string; - projectId: string; - message: string; - user: User; - sessionId: string; - attachments?: StoredAttachmentRef[]; - source?: string; - /** Set by the in-app preview chat only — see `ExecuteForChatConfig.previewChat`. */ - previewChat?: boolean; - onExecutionRecorded?: (executionId: string) => void; - abortSignal?: AbortSignal; +interface DraftRunConsumptionOptions { + errorMode?: 'throw' | 'forward'; + onChunk?: (chunk: StreamChunk) => void; } -interface ExecuteDraftRunInput extends PrepareDraftRunInput { - message: string; - user: User; - source?: string; - abortSignal?: AbortSignal; +export interface ExecutePreparedDraftRunInput + extends Omit, + DraftRunConsumptionOptions { + sessionId: string; } -interface ResumeDraftRunInput { - agentId: string; - projectId: string; - sessionId: string; - runId: string; - toolCallId: string; - resumeData: unknown; +export interface ResumePreparedDraftRunInput + extends Omit, + DraftRunConsumptionOptions { + user: User; + initialResponse?: string; +} + +type ExecuteDraftRunInput = PrepareDraftRunInput & + Pick; + +interface ResumeDraftRunInput + extends Omit< + ResumeForChatConfig, + 'expectedMemory' | 'integrationType' | 'usePublishedVersion' | 'onExecutionRecorded' + > { + sessionId: string; user: User; - source?: string; - /** Set by the in-app preview chat only — see `ExecuteForChatConfig.previewChat`. */ - previewChat?: boolean; response: string; - abortSignal?: AbortSignal; } export const agentTestRunContinuationSchema = z @@ -96,19 +94,14 @@ export interface AgentTestRunApproval extends ApprovalSuspendPayload { continuation: AgentTestRunContinuation; } -export type AgentTestRunResult = - | { status: 'completed'; response: string; sessionId: string; executionId?: string } - | { - status: 'suspended'; - response: string; - sessionId: string; - executionId?: string; - suspensions: AgentTestRunSuspension[]; - } - | { status: 'session_not_found' } - | { status: 'agent_misconfigured'; missing: string[] }; +export type PreparedDraftRunResult = { + response: string; + executionId?: string; +} & ({ status: 'completed' } | { status: 'suspended'; suspensions: AgentTestRunSuspension[] }); -type CollectedDraftRunResult = Extract; +export type AgentTestRunResult = + | (PreparedDraftRunResult & { sessionId: string }) + | Exclude; const expectedApprovalResumeJsonSchema = zodToJsonSchema(APPROVAL_RESUME_SCHEMA); @@ -185,80 +178,86 @@ export class AgentTestRunService { return { status: 'ready', sessionId: sessionId ?? randomUUID() }; } - streamDraftRun({ - agentId, - projectId, - message, - user, + async executePreparedDraftRun({ sessionId, - attachments, - source, - previewChat, + errorMode, + onChunk, onExecutionRecorded, - abortSignal, - }: StreamDraftRunInput): AsyncGenerator { - return this.agentExecutionOrchestratorService.executeForChat({ - agentId, - projectId, - message, - user, + ...execution + }: ExecutePreparedDraftRunInput): Promise { + let executionId: string | undefined; + const stream = this.agentExecutionOrchestratorService.executeForChat({ + ...execution, memory: { threadId: sessionId, - resourceId: draftChatMemoryResourceId(user.id), + resourceId: draftChatMemoryResourceId(execution.user.id), }, - attachments, - source, - previewChat, - onExecutionRecorded, - abortSignal, - }); - } - - async executeDraftRun(input: ExecuteDraftRunInput): Promise { - const prepared = await this.prepareDraftRun(input); - if (prepared.status !== 'ready') return prepared; - - let executionId: string | undefined; - const stream = this.streamDraftRun({ - ...input, - sessionId: prepared.sessionId, onExecutionRecorded: (id) => { executionId = id; + onExecutionRecorded?.(id); }, }); - return await this.collectDraftRun(stream, prepared.sessionId, '', () => executionId); + return await this.collectDraftRun(stream, '', () => executionId, { errorMode, onChunk }); } - async resumeDraftRun(input: ResumeDraftRunInput): Promise { - const existing = await this.agentExecutionService.findThreadById(input.sessionId); + async resumePreparedDraftRun({ + initialResponse = '', + errorMode, + onChunk, + onExecutionRecorded, + ...execution + }: ResumePreparedDraftRunInput): Promise { + let executionId: string | undefined; + const stream = this.agentExecutionOrchestratorService.resumeForChat({ + ...execution, + usePublishedVersion: false, + integrationType: N8N_CHAT_INTEGRATION_TYPE, + onExecutionRecorded: (id) => { + executionId = id; + onExecutionRecorded?.(id); + }, + }); + + return await this.collectDraftRun(stream, initialResponse, () => executionId, { + errorMode, + onChunk, + }); + } + + async executeDraftRun({ + credentialProvider, + ...input + }: ExecuteDraftRunInput): Promise { + const prepared = await this.prepareDraftRun({ ...input, credentialProvider }); + if (prepared.status !== 'ready') return prepared; + + const result = await this.executePreparedDraftRun({ + ...input, + sessionId: prepared.sessionId, + }); + return { ...result, sessionId: prepared.sessionId }; + } + + async resumeDraftRun({ + sessionId, + response, + ...input + }: ResumeDraftRunInput): Promise { + const existing = await this.agentExecutionService.findThreadById(sessionId); if (existing && !threadBelongsTo(existing, input.projectId, input.agentId)) { return { status: 'session_not_found' }; } - let executionId: string | undefined; - const stream = this.agentExecutionOrchestratorService.resumeForChat({ - agentId: input.agentId, - projectId: input.projectId, - runId: input.runId, - toolCallId: input.toolCallId, - resumeData: input.resumeData, - user: input.user, - usePublishedVersion: false, - integrationType: N8N_CHAT_INTEGRATION_TYPE, + const result = await this.resumePreparedDraftRun({ + ...input, expectedMemory: { - threadId: input.sessionId, + threadId: sessionId, resourceId: draftChatMemoryResourceId(input.user.id), }, - source: input.source, - previewChat: input.previewChat, - onExecutionRecorded: (id) => { - executionId = id; - }, - ...(input.abortSignal ? { abortSignal: input.abortSignal } : {}), + initialResponse: response, }); - - return await this.collectDraftRun(stream, input.sessionId, input.response, () => executionId); + return { ...result, sessionId }; } async resumeDraftApproval(input: { @@ -344,35 +343,50 @@ export class AgentTestRunService { private async collectDraftRun( stream: AsyncIterable, - sessionId: string, initialResponse: string, getExecutionId: () => string | undefined, - ): Promise { + { errorMode = 'throw', onChunk }: DraftRunConsumptionOptions, + ): Promise { let response = initialResponse; const suspensions: AgentTestRunSuspension[] = []; + let observerFailed = false; + let observerError: unknown; - for await (const chunk of stream) { - if (chunk.type === 'error') { - throw chunk.error; - } - if (chunk.type === 'text-delta') { - response += chunk.delta; - } else if (chunk.type === 'tool-call-suspended') { - suspensions.push({ - runId: chunk.runId, - toolCallId: chunk.toolCallId, - toolName: chunk.toolName, - ...(chunk.input !== undefined ? { input: chunk.input } : {}), - ...(chunk.suspendPayload !== undefined ? { suspendPayload: chunk.suspendPayload } : {}), - ...(chunk.resumeSchema !== undefined ? { resumeSchema: chunk.resumeSchema } : {}), - }); + try { + for await (const chunk of stream) { + if (!observerFailed && onChunk) { + try { + onChunk(chunk); + } catch (error) { + observerFailed = true; + observerError = error; + } + } + if (chunk.type === 'error' && errorMode === 'throw' && !observerFailed) { + throw chunk.error; + } + if (chunk.type === 'text-delta') { + response += chunk.delta; + } else if (chunk.type === 'tool-call-suspended') { + suspensions.push({ + runId: chunk.runId, + toolCallId: chunk.toolCallId, + toolName: chunk.toolName, + ...(chunk.input !== undefined ? { input: chunk.input } : {}), + ...(chunk.suspendPayload !== undefined ? { suspendPayload: chunk.suspendPayload } : {}), + ...(chunk.resumeSchema !== undefined ? { resumeSchema: chunk.resumeSchema } : {}), + }); + } } + } catch (error) { + if (observerFailed) throw observerError; + throw error; } + if (observerFailed) throw observerError; const executionId = getExecutionId(); const metadata = { response, - sessionId, ...(executionId ? { executionId } : {}), }; return suspensions.length > 0 diff --git a/packages/cli/src/modules/agents/agent-turn-execution.service.ts b/packages/cli/src/modules/agents/agent-turn-execution.service.ts index 7d2cd98e999..95d4451202d 100644 --- a/packages/cli/src/modules/agents/agent-turn-execution.service.ts +++ b/packages/cli/src/modules/agents/agent-turn-execution.service.ts @@ -84,100 +84,137 @@ export class AgentTurnExecutionService { ) {} async *execute(config: ExecuteTurnConfig): AsyncGenerator { - const { agentInstance, toolRegistry, backgroundJobSignal, onExecutionRecorded } = config; let executionId: string | undefined; let turn: AgentTurnRequest | undefined; let receivedFinish = false; const recorder = this.createRecorder( - toolRegistry, + config.toolRegistry, () => executionId, config.context, - backgroundJobSignal, + config.backgroundJobSignal, ); try { turn = await config.prepare(); - const result = - turn.type === 'start' - ? await agentInstance.stream(turn.input, turn.options) - : await agentInstance.resume('stream', turn.resumeData, turn.options); - if (turn.type === 'resume') { - recorder.recordHitlResponse(turn.options.toolCallId, turn.resumeData); - } - - // Keep SDK startup independent of execution-history storage. - executionId = await this.tryStartExecution( - { - ...turn.recording, - ...(backgroundJobSignal - ? { initialTimeline: structuredClone(recorder.getMessageRecord().timeline) } - : {}), - }, - recorder.startedAt, - turn.type === 'resume' - ? 'Failed to start resumed agent execution recording' - : 'Failed to start agent execution recording', - ); - const attributionTracker = createAttributionTracker(config.mcpServerAttributions); - for await (const value of streamAgentChunks(result.stream)) { - const chunk = config.includeHitlToolDetails - ? withApprovalToolDetails(value, toolRegistry) - : value; - recorder.record(chunk); - if (chunk.type === 'finish') receivedFinish = true; - if (turn.type === 'start') { - if (chunk.type === 'tool-call-suspended') { - this.logger.info('Chat: tool-call-suspended chunk received', { - agentId: config.context.agentId, - toolCallId: chunk.toolCallId, - toolName: chunk.toolName, - }); - } - if (chunk.type === 'finish' && chunk.finishReason === 'max-iterations') { - for (const chunk of getMaxIterationsChunks()) { - recorder.record(chunk); - yield chunk; - } - } - } - for (const attributionChunk of attributionTracker.observe(chunk)) { - recorder.record(attributionChunk); - yield attributionChunk; - } - yield chunk; - } + const startedTurn = await this.startTurn(turn, config, recorder); + executionId = startedTurn.executionId; + receivedFinish = yield* this.streamTurn(startedTurn.stream, turn, config, recorder); } catch (error) { recorder.record({ type: 'error', error }); recorder.record({ type: 'finish', finishReason: 'error' }); throw error; } finally { if (turn) { - const record = recorder.getMessageRecord(); - const cancelled = - turn.options.abortSignal?.aborted || - (!receivedFinish && !recorder.suspended && record.error === null); - await this.persistRecordedExecution({ - executionId, - onExecutionRecorded, - failureMessage: - turn.type === 'resume' - ? 'Failed to record resumed agent execution' - : 'Failed to record agent execution', - params: { - ...turn.recording, - record: cancelled ? { ...record, finishReason: 'cancelled', error: null } : record, - hitlStatus: recorder.suspended - ? 'suspended' - : turn.type === 'resume' - ? 'resumed' - : undefined, - }, - }); + await this.finalizeTurn(turn, config, recorder, executionId, receivedFinish); } await config.onSettled?.(recorder.suspended); } } + private async startTurn( + turn: AgentTurnRequest, + config: ExecuteTurnConfig, + recorder: ExecutionRecorder, + ): Promise<{ stream: ReadableStream; executionId: string | undefined }> { + const result = + turn.type === 'start' + ? await config.agentInstance.stream(turn.input, turn.options) + : await config.agentInstance.resume('stream', turn.resumeData, turn.options); + if (turn.type === 'resume') { + recorder.recordHitlResponse(turn.options.toolCallId, turn.resumeData); + } + + // Keep SDK startup independent of execution-history storage. + const executionId = await this.tryStartExecution( + { + ...turn.recording, + ...(config.backgroundJobSignal + ? { initialTimeline: structuredClone(recorder.getMessageRecord().timeline) } + : {}), + }, + recorder.startedAt, + turn.type === 'resume' + ? 'Failed to start resumed agent execution recording' + : 'Failed to start agent execution recording', + ); + + return { stream: result.stream, executionId }; + } + + private async *streamTurn( + stream: ReadableStream, + turn: AgentTurnRequest, + config: ExecuteTurnConfig, + recorder: ExecutionRecorder, + ): AsyncGenerator { + let receivedFinish = false; + const attributionTracker = createAttributionTracker(config.mcpServerAttributions); + + for await (const value of streamAgentChunks(stream)) { + const chunk = config.includeHitlToolDetails + ? withApprovalToolDetails(value, config.toolRegistry) + : value; + recorder.record(chunk); + if (chunk.type === 'finish') receivedFinish = true; + + if (turn.type === 'start') { + if (chunk.type === 'tool-call-suspended') { + this.logger.info('Chat: tool-call-suspended chunk received', { + agentId: config.context.agentId, + toolCallId: chunk.toolCallId, + toolName: chunk.toolName, + }); + } + if (chunk.type === 'finish' && chunk.finishReason === 'max-iterations') { + for (const maxIterationsChunk of getMaxIterationsChunks()) { + recorder.record(maxIterationsChunk); + yield maxIterationsChunk; + } + } + } + + for (const attributionChunk of attributionTracker.observe(chunk)) { + recorder.record(attributionChunk); + yield attributionChunk; + } + yield chunk; + } + + return receivedFinish; + } + + private async finalizeTurn( + turn: AgentTurnRequest, + config: ExecuteTurnConfig, + recorder: ExecutionRecorder, + executionId: string | undefined, + receivedFinish: boolean, + ): Promise { + const record = recorder.getMessageRecord(); + const cancelled = + turn.options.abortSignal?.aborted || + (!receivedFinish && !recorder.suspended && record.error === null); + const hitlStatus = recorder.suspended + ? 'suspended' + : turn.type === 'resume' + ? 'resumed' + : undefined; + + await this.persistRecordedExecution({ + executionId, + onExecutionRecorded: config.onExecutionRecorded, + failureMessage: + turn.type === 'resume' + ? 'Failed to record resumed agent execution' + : 'Failed to record agent execution', + params: { + ...turn.recording, + record: cancelled ? { ...record, finishReason: 'cancelled', error: null } : record, + hitlStatus, + }, + }); + } + createRecorder( toolRegistry?: ToolRegistry, getExecutionId: () => string | undefined = () => undefined,