mirror of
https://github.com/n8n-io/n8n.git
synced 2026-10-17 05:29:44 +02:00
refactor(core): Share preview agent execution across callers (no-changelog) (#39064)
This commit is contained in:
Robin Braumann
GitHub
parent
3d00ed98b0
commit
6f2bc13984
|
|
@ -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<Pick<AgentsService, 'findById' | 'findByProjectId' | 'findByProjectIdPaginated'>>();
|
||||
const agentExecutionOrchestratorService = mock<AgentExecutionOrchestratorService>();
|
||||
const agentsBuilderService = mock<AgentsBuilderService>();
|
||||
const agentTestRunService = mock<AgentTestRunService>();
|
||||
const agentChatAttachmentService = mock<AgentChatAttachmentService>();
|
||||
const agentExecutionService = mock<AgentExecutionService>();
|
||||
const agentValidationService = mock<AgentValidationService>();
|
||||
const backgroundJobService = mock<AgentBackgroundJobService>();
|
||||
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<N8NCheckpointStorage>(),
|
||||
);
|
||||
|
||||
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<void>((resolve) => {
|
||||
markStarted = resolve;
|
||||
});
|
||||
const runBlocked = new Promise<void>((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<void>((resolve) => {
|
||||
markStarted = resolve;
|
||||
});
|
||||
const runBlocked = new Promise<void>((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();
|
||||
|
|
|
|||
|
|
@ -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<AgentValidationService>(),
|
||||
service,
|
||||
checkpointStorage,
|
||||
);
|
||||
const finalization = createDeferredPromise<string>();
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -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<T>(items: T[]): AsyncIterable<T> {
|
|||
|
||||
async function collectEvents(chunks: StreamChunk[]): Promise<AgentSseEvent[]> {
|
||||
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<AgentSseEvent[]> {
|
||||
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',
|
||||
|
|
|
|||
|
|
@ -19,6 +19,13 @@ const agentId = 'agent-1';
|
|||
const projectId = 'project-1';
|
||||
const user = mock<User>({ id: 'user-1' });
|
||||
const credentialProvider = mock<CredentialProvider>();
|
||||
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);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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<User>({ id: 'user-1' }));
|
||||
agentTestRunService.resumeDraftRun.mockResolvedValue(result);
|
||||
const execution = createDeferredPromise<typeof result>();
|
||||
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',
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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<StreamChunk>,
|
||||
send: (e: AgentSseEvent) => void,
|
||||
): Promise<boolean> {
|
||||
const ctx: ChunkHandlerCtx = { send };
|
||||
let suspended = false;
|
||||
|
||||
for await (const chunk of chunks) {
|
||||
const result = emitChunkEvents(chunk, ctx);
|
||||
suspended ||= result.suspended;
|
||||
}
|
||||
return suspended;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<ExecuteForChatConfig, 'memory'>,
|
||||
DraftRunConsumptionOptions {
|
||||
sessionId: string;
|
||||
}
|
||||
|
||||
interface ResumeDraftRunInput {
|
||||
agentId: string;
|
||||
projectId: string;
|
||||
sessionId: string;
|
||||
runId: string;
|
||||
toolCallId: string;
|
||||
resumeData: unknown;
|
||||
export interface ResumePreparedDraftRunInput
|
||||
extends Omit<ResumeForChatConfig, 'integrationType' | 'usePublishedVersion'>,
|
||||
DraftRunConsumptionOptions {
|
||||
user: User;
|
||||
initialResponse?: string;
|
||||
}
|
||||
|
||||
type ExecuteDraftRunInput = PrepareDraftRunInput &
|
||||
Pick<ExecuteForChatConfig, 'message' | 'user' | 'source' | 'abortSignal'>;
|
||||
|
||||
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<AgentTestRunResult, { status: 'completed' | 'suspended' }>;
|
||||
export type AgentTestRunResult =
|
||||
| (PreparedDraftRunResult & { sessionId: string })
|
||||
| Exclude<PrepareDraftRunResult, { status: 'ready' }>;
|
||||
|
||||
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<StreamChunk> {
|
||||
return this.agentExecutionOrchestratorService.executeForChat({
|
||||
agentId,
|
||||
projectId,
|
||||
message,
|
||||
user,
|
||||
...execution
|
||||
}: ExecutePreparedDraftRunInput): Promise<PreparedDraftRunResult> {
|
||||
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<AgentTestRunResult> {
|
||||
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<AgentTestRunResult> {
|
||||
const existing = await this.agentExecutionService.findThreadById(input.sessionId);
|
||||
async resumePreparedDraftRun({
|
||||
initialResponse = '',
|
||||
errorMode,
|
||||
onChunk,
|
||||
onExecutionRecorded,
|
||||
...execution
|
||||
}: ResumePreparedDraftRunInput): Promise<PreparedDraftRunResult> {
|
||||
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<AgentTestRunResult> {
|
||||
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<AgentTestRunResult> {
|
||||
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<StreamChunk>,
|
||||
sessionId: string,
|
||||
initialResponse: string,
|
||||
getExecutionId: () => string | undefined,
|
||||
): Promise<CollectedDraftRunResult> {
|
||||
{ errorMode = 'throw', onChunk }: DraftRunConsumptionOptions,
|
||||
): Promise<PreparedDraftRunResult> {
|
||||
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
|
||||
|
|
|
|||
|
|
@ -84,100 +84,137 @@ export class AgentTurnExecutionService {
|
|||
) {}
|
||||
|
||||
async *execute(config: ExecuteTurnConfig): AsyncGenerator<StreamChunk> {
|
||||
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<StreamChunk>; 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<StreamChunk>,
|
||||
turn: AgentTurnRequest,
|
||||
config: ExecuteTurnConfig,
|
||||
recorder: ExecutionRecorder,
|
||||
): AsyncGenerator<StreamChunk, boolean> {
|
||||
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<void> {
|
||||
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,
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user