diff --git a/.changeset/fail-streamed-run-errors.md b/.changeset/fail-streamed-run-errors.md new file mode 100644 index 0000000000..7618a65cf7 --- /dev/null +++ b/.changeset/fail-streamed-run-errors.md @@ -0,0 +1,6 @@ +--- +'@tanstack/ai': patch +'@tanstack/ai-persistence': patch +--- + +Route adapter-emitted `RUN_ERROR` events through middleware `onError` hooks and preserve provider error codes in persisted run failures. diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index bbf368cf97..09d0731144 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -11,6 +11,7 @@ import { getGenericInterruptDefinitionRegistry, providePendingTurn, rehydrateInterruptRequest, + toRunErrorPayload, } from '@tanstack/ai/adapter-internals' import type { GenericInterruptRequest, @@ -1808,14 +1809,14 @@ async function failRun( error: unknown, usage?: TokenUsage, ): Promise { - // `RunRecord.error` is a structured `RunError`. Only `message` is filled in - // here: the middleware sees an opaque thrown value, and inventing a `code` - // from it would fabricate the stable classification consumers branch on. A - // provider-supplied code reaches the record through the adapter layer. + const runError = toRunErrorPayload(error) await runs?.update(runId, { status: 'failed', finishedAt: Date.now(), - error: { message: error instanceof Error ? error.message : String(error) }, + error: { + message: runError.message, + ...(runError.code !== undefined ? { code: runError.code } : {}), + }, ...(usage ? { usage } : {}), }) } diff --git a/packages/ai-persistence/tests/error-abort.test.ts b/packages/ai-persistence/tests/error-abort.test.ts index 7c7074aa46..b45f814aff 100644 --- a/packages/ai-persistence/tests/error-abort.test.ts +++ b/packages/ai-persistence/tests/error-abort.test.ts @@ -1,5 +1,11 @@ import { describe, expect, it, vi } from 'vitest' -import { EventType, chat, generateImage } from '@tanstack/ai' +import { + EventType, + chat, + defineChatMiddleware, + defineInterrupt, + generateImage, +} from '@tanstack/ai' import type { AnyTextAdapter, GenerationAbortInfo, @@ -52,6 +58,23 @@ const interruptFinished = (): StreamChunk => ({ }, }) +const booleanResponseSchema = { + '~standard': { + version: 1, + vendor: 'test', + validate(value: unknown) { + return typeof value === 'boolean' + ? { value } + : { issues: [{ message: 'response must be a boolean' }] } + }, + jsonSchema: { + input() { + return { type: 'boolean' } + }, + }, + }, +} as const + async function collect(stream: AsyncIterable) { const out: Array = [] for await (const c of stream) out.push(c) @@ -94,6 +117,70 @@ describe('chat persistence error/abort hooks', () => { expect(run?.error).toEqual({ message: 'provider exploded' }) }) + it('marks the run failed when the adapter emits RUN_ERROR', async () => { + const persistence = memoryPersistence() + const review = defineInterrupt({ + id: 'review', + responseSchema: booleanResponseSchema, + }) + const { adapter } = mockAdapter([ + [ + runStarted(), + { + type: EventType.RUN_ERROR, + message: 'provider failed', + code: 'provider_error', + timestamp: 1, + }, + ], + ]) + + const chunks = await collect( + chat({ + adapter, + interrupts: [review], + messages: [{ role: 'user', content: 'hi' }], + runId: 'r1', + threadId: 't1', + middleware: [ + defineChatMiddleware({ + onInterruptBoundary(ctx) { + if (ctx.phase !== 'afterModel') return + return { + interrupts: [ + review.interrupt({ + key: 'review', + reason: 'review', + message: 'Review the response', + }), + ], + } + }, + }), + withPersistence(persistence), + ], + }) as AsyncIterable, + ) + + expect(chunks).toContainEqual( + expect.objectContaining({ + type: EventType.RUN_ERROR, + message: 'provider failed', + code: 'provider_error', + }), + ) + expect(chunks).not.toContainEqual( + expect.objectContaining({ + type: EventType.RUN_FINISHED, + outcome: expect.objectContaining({ type: 'interrupt' }), + }), + ) + expect(await persistence.stores.runs!.get('r1')).toMatchObject({ + status: 'failed', + error: { message: 'provider failed', code: 'provider_error' }, + }) + }) + it('preserves known usage when structured-output finalization fails', async () => { const persistence = memoryPersistence() const usage = { diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 4ff000e55c..9bf2761aa0 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -1183,9 +1183,8 @@ class TextEngine< if (!skipAgentLoop) { do { - if (this.earlyTermination || this.isCancelled()) { - return - } + if (this.earlyTermination) break + if (this.isCancelled()) return this.logger.agentLoop(`iteration=${this.middlewareCtx.iteration}`, { iteration: this.middlewareCtx.iteration, @@ -1217,6 +1216,8 @@ class TextEngine< yield* this.streamModelResponse() + if (this.earlyTermination) break + if ( yield* this.emitBoundaryInterrupts( 'afterModel', @@ -1783,11 +1784,8 @@ class TextEngine< chunk: Extract, ): void { this.earlyTermination = true - if (this.finalStructuredOutput && this.finalizationError === null) { - const message = - chunk.message || - chunk.error?.message || - 'Run failed before structured output completed' + if (this.finalizationError === null) { + const message = chunk.message || chunk.error?.message || 'Run failed' this.finalizationError = { message, ...(chunk.code !== undefined diff --git a/packages/ai/tests/middleware.test.ts b/packages/ai/tests/middleware.test.ts index 300c5eca76..bebef54832 100644 --- a/packages/ai/tests/middleware.test.ts +++ b/packages/ai/tests/middleware.test.ts @@ -1,6 +1,10 @@ /* eslint-disable @typescript-eslint/require-await */ import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' import { chat } from '../src/activities/chat/index' +import { defineChatMiddleware } from '../src/activities/chat/middleware/define' +import { defineInterrupt } from '../src/interrupt-definition' +import { EventType } from '../src/types' import { collectChunks, createMockAdapter, @@ -94,6 +98,72 @@ describe('chat() middleware', () => { expect(onFinish).not.toHaveBeenCalled() }) + it('should call onError when an adapter emits RUN_ERROR', async () => { + const onError = vi.fn() + const onFinish = vi.fn() + const review = defineInterrupt({ + id: 'review', + responseSchema: z.object({ approved: z.boolean() }), + }) + + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + { + ...ev.runError('Provider failed'), + code: 'provider_error', + }, + ], + ], + }) + + const chunks = await collectChunks( + chat({ + adapter, + interrupts: [review], + messages: [{ role: 'user', content: 'Hi' }], + middleware: [ + defineChatMiddleware({ + onInterruptBoundary(ctx) { + if (ctx.phase !== 'afterModel') return + return { + interrupts: [ + review.interrupt({ + key: 'review', + reason: 'review', + message: 'Review the response', + }), + ], + } + }, + }), + { name: 'test', onError, onFinish }, + ], + }) as AsyncIterable, + ) + + expect(chunks).toContainEqual( + expect.objectContaining({ + type: EventType.RUN_ERROR, + message: 'Provider failed', + code: 'provider_error', + }), + ) + expect(chunks).not.toContainEqual( + expect.objectContaining({ + type: EventType.RUN_FINISHED, + outcome: expect.objectContaining({ type: 'interrupt' }), + }), + ) + expect(onError).toHaveBeenCalledOnce() + expect(onError.mock.calls[0]![1].error).toMatchObject({ + message: 'Provider failed', + code: 'provider_error', + }) + expect(onFinish).not.toHaveBeenCalled() + }) + it('should call exactly one terminal hook per run', async () => { const onStart = vi.fn() const onFinish = vi.fn() diff --git a/testing/e2e/src/lib/phase-capture.ts b/testing/e2e/src/lib/phase-capture.ts index 9fc78d432c..fa9219c74e 100644 --- a/testing/e2e/src/lib/phase-capture.ts +++ b/testing/e2e/src/lib/phase-capture.ts @@ -39,8 +39,10 @@ export interface PhaseCapture { * that only care about presence should use `.includes('structuredOutput')`. */ phases: Array - /** Count of `onFinish` invocations — must be exactly 1 per `chat()` call. */ + /** Count of `onFinish` invocations. */ onFinishCount: number + /** Count of `onError` invocations. */ + onErrorCount: number /** * Chunks that were yielded out of `chat()` to the SSE consumer. Captured * by teeing the iterable in `api.middleware-test.ts` after the middleware @@ -61,6 +63,7 @@ function bucketFor(captureId: string): PhaseCapture { bucket = { phases: [], onFinishCount: 0, + onErrorCount: 0, yieldedChunks: [], boundaries: [], resolutions: [], @@ -76,6 +79,7 @@ export function resetPhaseCapture(captureId: string): void { captures.set(captureId, { phases: [], onFinishCount: 0, + onErrorCount: 0, yieldedChunks: [], boundaries: [], resolutions: [], @@ -96,6 +100,10 @@ export function recordOnFinish(captureId: string): void { bucketFor(captureId).onFinishCount += 1 } +export function recordOnError(captureId: string): void { + bucketFor(captureId).onErrorCount += 1 +} + export function recordYieldedChunk( captureId: string, chunk: YieldedChunkSummary, diff --git a/testing/e2e/src/routes/api.middleware-test.ts b/testing/e2e/src/routes/api.middleware-test.ts index 031f38946d..41f14857f2 100644 --- a/testing/e2e/src/routes/api.middleware-test.ts +++ b/testing/e2e/src/routes/api.middleware-test.ts @@ -3,6 +3,7 @@ import { chat, chatParamsFromRequestBody, createCapability, + EventType, maxIterations, toServerSentEventsResponse, toolDefinition, @@ -18,7 +19,12 @@ import { } from '@/lib/memory-capture' import { SpanStatusCode } from '@opentelemetry/api' import { z } from 'zod' -import type { ChatMiddleware, StreamChunk, Tool } from '@tanstack/ai' +import type { + AnyTextAdapter, + ChatMiddleware, + StreamChunk, + Tool, +} from '@tanstack/ai' import type { AttributeValue, Attributes, @@ -38,6 +44,7 @@ import { recordGenericPolicy, recordGenericResolution, recordGenericToolExecution, + recordOnError, recordOnFinish, recordPhase, recordYieldedChunk, @@ -77,6 +84,54 @@ const weatherTool = toolDefinition({ JSON.stringify({ city: args.city, temperature: 72, condition: 'sunny' }), ) +function createRunErrorAdapter(): AnyTextAdapter { + return { + kind: 'text', + name: 'run-error-test', + model: 'run-error-test', + '~types': { + providerOptions: {}, + inputModalities: ['text'], + messageMetadataByModality: {}, + toolCapabilities: [], + toolCallMetadata: undefined, + systemPromptMetadata: undefined, + }, + async *chatStream(options): AsyncGenerator { + yield { + type: EventType.RUN_STARTED, + runId: options.runId ?? 'run-error-test', + threadId: options.threadId ?? 'run-error-test', + timestamp: Date.now(), + } + yield { + type: EventType.RUN_ERROR, + message: 'Provider failed', + code: 'provider_error', + timestamp: Date.now(), + } + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } +} + +const runErrorBoundaryMiddleware: ChatMiddleware = { + name: 'run-error-boundary', + onInterruptBoundary(ctx) { + if (ctx.phase !== 'afterModel') return + return { + interrupts: [ + reviewPlan.interrupt({ + key: 'run-error-review', + reason: 'review_required', + message: 'Review the response', + payload: { title: 'Run error review', boundary: ctx.phase }, + }), + ], + } + }, +} + const chunkTransformMiddleware: ChatMiddleware = { name: 'chunk-transform', onChunk(_ctx, chunk) { @@ -168,6 +223,9 @@ function createPhaseRecorderMiddleware(captureId: string): ChatMiddleware { onFinish() { recordOnFinish(captureId) }, + onError() { + recordOnError(captureId) + }, } } @@ -478,17 +536,21 @@ export const Route = createFileRoute('/api/middleware-test')({ typeof fp.model === 'string' ? fp.model : undefined try { - const adapterOptions = createTextAdapter( - provider as Parameters[0], - modelOverride, - aimockPort, - testId, - ) + const adapterOptions = + scenario === 'run-error' + ? { adapter: createRunErrorAdapter() } + : createTextAdapter( + provider as Parameters[0], + modelOverride, + aimockPort, + testId, + ) const middleware: Array = [] let genericLifecycleMiddleware: | ChatMiddleware - | undefined + | undefined = + scenario === 'run-error' ? runErrorBoundaryMiddleware : undefined const genericScenario = isGenericScenario(scenario) ? scenario : undefined diff --git a/testing/e2e/src/routes/middleware-test.tsx b/testing/e2e/src/routes/middleware-test.tsx index f9ae3899fe..c55a1cfada 100644 --- a/testing/e2e/src/routes/middleware-test.tsx +++ b/testing/e2e/src/routes/middleware-test.tsx @@ -227,6 +227,7 @@ function MiddlewareTestPage() { +