diff --git a/packages/ai-orchestration/src/execute-agent.ts b/packages/ai-orchestration/src/execute-agent.ts index d7ddfd8dd..796a33ba8 100644 --- a/packages/ai-orchestration/src/execute-agent.ts +++ b/packages/ai-orchestration/src/execute-agent.ts @@ -33,8 +33,14 @@ export async function executeAgent(args: { }) if (isAgentStreamResult(result)) { + // `output` is already in flight, so observe it before draining can throw. + // Otherwise a failing stream leaves it rejected and unhandled, which + // terminates the worker under Node's default unhandled-rejection policy. + const output = Promise.resolve(result.output) + output.catch(() => undefined) + await drainAgentStream(definition.name, result.stream, args.emit) - return validateOutput(definition, await result.output) + return validateOutput(definition, await output) } if (isAsyncIterable(result)) { diff --git a/packages/ai-orchestration/tests/orchestration.test.ts b/packages/ai-orchestration/tests/orchestration.test.ts index 3bb276732..c556e4c1a 100644 --- a/packages/ai-orchestration/tests/orchestration.test.ts +++ b/packages/ai-orchestration/tests/orchestration.test.ts @@ -8,7 +8,9 @@ import { z } from 'zod' import { EventType } from '@tanstack/ai' import { AgentApprovalUnsupportedError, + AgentStreamError, agentMiddleware, + agentStream, createAIEventPublisher, defineAgent, toAIStream, @@ -215,6 +217,44 @@ describe('Workflow-backed agent orchestration', () => { expect(events.some((event) => event.type === 'RUN_FINISHED')).toBe(false) }) + it('fails the step without an unhandled rejection when an agentStream errors', async () => { + const unhandled: Array = [] + const onUnhandled = (reason: unknown) => unhandled.push(reason) + process.on('unhandledRejection', onUnhandled) + + const broken = defineAgent({ + name: 'broken', + run: () => + agentStream( + errorStream('model exploded'), + Promise.reject(new Error('output never settles')), + ), + }) + const workflow = createWorkflow({ id: 'broken' }) + .middleware([agentMiddleware()]) + .handler((ctx) => ctx.ai.agent('broken', broken, undefined)) + + const events = await collect( + runWorkflow({ + workflow, + runStore: inMemoryRunStore(), + input: {}, + }), + ) + // Let a rejection scheduled during the run reach the process handler. + await new Promise((resolve) => setTimeout(resolve, 0)) + process.off('unhandledRejection', onUnhandled) + + expect(unhandled).toEqual([]) + expect(events).toContainEqual( + expect.objectContaining({ + type: 'STEP_FAILED', + error: expect.objectContaining({ name: AgentStreamError.name }), + }), + ) + expect(events.some((event) => event.type === 'RUN_FINISHED')).toBe(false) + }) + it('adapts the Workflow publish hook without changing execution IDs', async () => { const published: Array<{ runId: string; chunk: StreamChunk }> = [] const publisher = createAIEventPublisher({ @@ -280,6 +320,14 @@ async function* streamText(text: string): AsyncIterable { } } +async function* errorStream(message: string): AsyncIterable { + yield { + type: EventType.RUN_ERROR, + message, + timestamp: Date.now(), + } +} + async function* approvalStream(): AsyncIterable { yield { type: 'CUSTOM',