diff --git a/CHANGELOG.md b/CHANGELOG.md index 96877890..be975420 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,12 @@ size limit (#418, #423) - `getFilesToDelete` and `deleteFiles` take a `lastTouchedBefore` cutoff to clean up recent unreferenced files (#361, #437) +- Save every step the AI SDK runs, fail the generation when a step save fails, + and keep provider errors in `fullStream` with `saveStreamDeltas`. A caller's + `onStepEnd` now runs before its step is saved, its `onEnd` after the last one, + and a step cut off by `stopStream` fails instead of staying pending (#388, + #419, #428) +- Aborting a stream no longer fails a step that already finished, reversing #326 ## 0.7.9 diff --git a/src/component/messages.ts b/src/component/messages.ts index 3339b5ff..82d04128 100644 --- a/src/component/messages.ts +++ b/src/component/messages.ts @@ -46,6 +46,7 @@ import { finishHandler, releaseStreamFileOwnershipByIds, abortStreamsAtOrder, + abortById, } from "./streams.js"; import { partial } from "convex-helpers/validators"; import { @@ -208,6 +209,7 @@ const addMessagesArgs = { hideFromUserIdSearch: v.optional(v.boolean()), // If provided, finish this stream atomically with the message save. // This prevents UI flickering from separate mutations (issue #181). + // Failed messages abort it instead. finishStreamId: v.optional(v.id("streamingMessages")), }; export const addMessages = mutation({ @@ -426,7 +428,13 @@ async function addMessagesHandler( } // Atomically finish the stream if requested, preventing UI flickering // from separate mutations for message save and stream finish (issue #181). - if (finishStreamId) { + const failed = toReturn.find((m) => m.status === "failed"); + if (finishStreamId && failed) { + await abortById(ctx, { + streamId: finishStreamId, + reason: failed.error ?? "The generation failed", + }); + } else if (finishStreamId) { await finishHandler(ctx, { streamId: finishStreamId, }); diff --git a/src/component/streams.ts b/src/component/streams.ts index e5dcf7fe..aad76713 100644 --- a/src/component/streams.ts +++ b/src/component/streams.ts @@ -303,7 +303,7 @@ export const abort = mutation({ handler: abortById, }); -async function abortById( +export async function abortById( ctx: MutationCtx, args: { streamId: Id<"streamingMessages">; diff --git a/src/vercel/client/start.ts b/src/vercel/client/start.ts index e847c0b3..90da4cb7 100644 --- a/src/vercel/client/start.ts +++ b/src/vercel/client/start.ts @@ -34,6 +34,7 @@ import { wrapTools, type ToolCtx } from "./createTool.js"; import type { Agent } from "../index.js"; import { assert, omit } from "convex-helpers"; import { saveInputMessages } from "./saveInputMessages.js"; +import { errorToString } from "./utils.js"; import type { GenericActionCtx, GenericDataModel } from "convex/server"; export function resolveUsageModel( @@ -182,7 +183,16 @@ export async function startGeneration< createPendingMessage?: boolean, finishStreamId?: string, ) => Promise; - fail: (reason: string) => Promise; + fail: (reason: unknown) => Promise; + endStep: ( + step: StepResult, + error?: string, + ) => Promise; + beforeStep: (options: { + stepNumber: number; + responseMessages: ModelMessage[]; + }) => Promise; + end: (options?: { finishStreamId?: string }) => Promise; getSavedMessages: () => MessageDoc[]; }> { const userId = @@ -236,29 +246,61 @@ export async function startGeneration< assert(model, "model is required"); let activeModel: ModelOrMetadata = model; - // Both the caller's AbortSignal listener and AI SDK's onAbort can report - // the same cancellation. Share the one finalization instead of racing two - // mutations for the pending message. - let pendingMessageFailure: Promise | undefined; - const fail = (reason: string): Promise => { - if (!pendingMessageId) return Promise.resolve(); - if (!pendingMessageFailure) { - const messageId = pendingMessageId; - pendingMessageFailure = ctx - .runMutation(component.messages.finalizeMessage, { - messageId, - result: { status: "failed", error: reason }, - }) - .then(() => undefined); - } - return pendingMessageFailure; + let initialResponseMessages: ModelMessage[] = []; + let held: + | { + step: StepResult; + responseMessages: ModelMessage[]; + error?: string; + } + | undefined; + // The last step save. A failure waits for it, to fail the pending message + // it leaves, and reports it if it failed. + let saving: Promise = Promise.resolve(); + const saveHeld = (createPendingMessage: boolean, finishStreamId?: string) => { + const toSave = held; + held = undefined; + if (!toSave) return Promise.resolve(); + const saved = persist(toSave, createPendingMessage, { + finishStreamId, + error: toSave.error, + }).then((stoppedBy) => { + if (stoppedBy) throw new Error(stoppedBy); + }); + saving = saved; + return saved; + }; + // What ended the generation early, once something has. Both the caller's + // AbortSignal listener and AI SDK's onAbort can report the same + // cancellation, so there is one. A step that ends after it is only reported. + let failure: { error: unknown; persisted: Promise } | undefined; + let ending: Promise | undefined; + const fail = (reason: unknown): Promise => { + failure ??= { + error: typeof reason === "string" ? new Error(reason) : reason, + persisted: (async () => { + try { + await saving; + await saveHeld(false); + } finally { + if (pendingMessageId) { + await ctx.runMutation(component.messages.finalizeMessage, { + messageId: pendingMessageId, + result: { status: "failed", error: errorToString(reason) }, + }); + } + } + })(), + }; + return failure.persisted; }; if (args.abortSignal) { const abortSignal = args.abortSignal; + // The caller sees a failure to persist it when the generation ends. abortSignal.addEventListener( "abort", - async () => { - await fail(abortSignal.reason?.toString() ?? "abortSignal"); + () => { + fail(abortSignal.reason?.toString() ?? "abortSignal").catch(() => {}); }, { once: true }, ); @@ -301,8 +343,11 @@ export async function startGeneration< messages: context.messages, instructions: args.instructions ?? args.system, allowSystemInMessages, - stopWhen: - args.stopWhen ?? (opts.maxSteps ? isStepCount(opts.maxSteps) : undefined), + // A failure, or a failed step, also stops the AI SDK after this step. + stopWhen: [ + ...[args.stopWhen ?? isStepCount(opts.maxSteps ?? 1)].flat(), + () => failure !== undefined || held?.error !== undefined, + ], tools, ...(include ? { include } : {}), } as unknown as WithoutModel & { @@ -317,6 +362,119 @@ export async function startGeneration< // tool execution IDs, message IDs, etc.) and they must be unique. // The pending message is linked via the explicit `pendingMessageId` // parameter passed to addMessages in the save closure. + type ToSave = + | { + step: StepResult; + responseMessages?: ModelMessage[]; + } + | { object: GenerateObjectResult }; + const persist = async ( + toSave: ToSave, + createPendingMessage: boolean | undefined, + { finishStreamId, error }: { finishStreamId?: string; error?: string }, + ): Promise => { + if (threadId && saveMessages !== "none") { + let serialized; + if ("object" in toSave) { + serialized = await serializeObjectResult( + ctx, + component, + toSave.object, + activeModel, + ); + } else { + const newResponseMessages = + toSave.responseMessages ?? toSave.step.response.messages; + // Even an empty completed step needs a durable assistant result so + // the pending message can be finalized at the storage boundary. + const responseMessagesToSave: ModelMessage[] = + newResponseMessages.length > 0 + ? newResponseMessages + : [{ role: "assistant", content: [] }]; + serialized = await serializeResponseMessages( + ctx, + component, + toSave.step, + activeModel, + responseMessagesToSave, + ); + } + if (error !== undefined) { + for (const message of serialized.messages) { + message.status = "failed"; + message.error = error; + } + } + const embeddings = await embedMessages( + ctx, + { threadId, ...opts, userId }, + serialized.messages.map((m) => m.message), + ); + if (createPendingMessage) { + serialized.messages.push({ + message: { role: "assistant", content: [] }, + status: "pending", + }); + embeddings?.vectors.push(null); + } + const saved = await ctx.runMutation(component.messages.addMessages, { + userId, + threadId, + agentName: opts.agentName, + promptMessageId, + abandonIfPromptMissing: true, + pendingMessageId, + messages: serialized.messages, + embeddings, + failPendingSteps: false, + finishStreamId, + }); + const lastMessage = saved.messages.at(-1); + if (createPendingMessage && lastMessage?.status === "pending") { + pendingMessageId = lastMessage._id; + savedMessages.push(...saved.messages.slice(0, -1)); + } else { + pendingMessageId = undefined; + savedMessages.push(...saved.messages); + } + // A retry or an abort elsewhere failed what this save replaced. + if (lastMessage?.status === "failed" && error === undefined) { + return ( + lastMessage.error ?? + "Aborting - the pending message was marked as failed" + ); + } + // The prompt message was deleted, so the save was abandoned (#300). + if (!lastMessage && createPendingMessage) { + return "The prompt message was deleted"; + } + } + return undefined; + }; + const report = async (toSave: ToSave) => { + const output = "object" in toSave ? toSave.object : toSave.step; + if (opts.rawRequestResponseHandler) { + await opts.rawRequestResponseHandler(ctx, { + userId, + threadId, + agentName: opts.agentName, + request: output.request, + response: output.response, + }); + } + if (opts.usageHandler && output.usage) { + const usageModel = resolveUsageModel(toSave, activeModel); + await opts.usageHandler(ctx, { + userId, + threadId, + agentName: opts.agentName, + model: getModelName(usageModel), + provider: getProviderName(usageModel), + usage: output.usage, + providerMetadata: output.providerMetadata, + }); + } + }; return { args: aiArgs, order: order ?? 0, @@ -330,13 +488,37 @@ export async function startGeneration< } }, fail, + endStep: async (step, error) => { + held = { + step, + responseMessages: [ + ...initialResponseMessages, + ...step.response.messages, + ], + error, + }; + initialResponseMessages = []; + await report({ step }); + }, + beforeStep: async ({ stepNumber, responseMessages }) => { + if (stepNumber === 0) { + initialResponseMessages = [...responseMessages]; + return; + } + if (!failure) await saveHeld(true).catch(fail); + if (failure) { + await failure.persisted; + throw failure.error; + } + }, + // Throws if a save failed. Any other failure already reached the caller. + end: ({ finishStreamId } = {}) => + (ending ??= (async () => { + if (!failure) await saveHeld(false, finishStreamId).catch(fail); + await failure?.persisted; + })()), save: async ( - toSave: - | { - step: StepResult; - responseMessages?: ModelMessage[]; - } - | { object: GenerateObjectResult }, + toSave: ToSave, createPendingMessage?: boolean, /** * If provided, finish this stream atomically with the message save. @@ -344,97 +526,8 @@ export async function startGeneration< */ finishStreamId?: string, ) => { - if (threadId && saveMessages !== "none") { - let serialized; - if ("object" in toSave) { - serialized = await serializeObjectResult( - ctx, - component, - toSave.object, - activeModel, - ); - } else { - const newResponseMessages = - toSave.responseMessages ?? toSave.step.response.messages; - // Even an empty completed step needs a durable assistant result so - // the pending message can be finalized at the storage boundary. - const responseMessagesToSave: ModelMessage[] = - newResponseMessages.length > 0 - ? newResponseMessages - : [{ role: "assistant", content: [] }]; - serialized = await serializeResponseMessages( - ctx, - component, - toSave.step, - activeModel, - responseMessagesToSave, - ); - } - const embeddings = await embedMessages( - ctx, - { threadId, ...opts, userId }, - serialized.messages.map((m) => m.message), - ); - if (createPendingMessage) { - serialized.messages.push({ - message: { role: "assistant", content: [] }, - status: "pending", - }); - embeddings?.vectors.push(null); - } - const saved = await ctx.runMutation(component.messages.addMessages, { - userId, - threadId, - agentName: opts.agentName, - promptMessageId, - abandonIfPromptMissing: true, - pendingMessageId, - messages: serialized.messages, - embeddings, - failPendingSteps: false, - finishStreamId, - }); - const lastMessage = saved.messages.at(-1)!; - if (createPendingMessage) { - if (lastMessage.status === "failed") { - pendingMessageId = undefined; - savedMessages.push(...saved.messages); - await fail( - lastMessage.error ?? - "Aborting - the pending message was marked as failed", - ); - } else { - pendingMessageId = lastMessage._id; - pendingMessageFailure = undefined; - savedMessages.push(...saved.messages.slice(0, -1)); - } - } else { - pendingMessageId = undefined; - savedMessages.push(...saved.messages); - } - } - const output = "object" in toSave ? toSave.object : toSave.step; - if (opts.rawRequestResponseHandler) { - await opts.rawRequestResponseHandler(ctx, { - userId, - threadId, - agentName: opts.agentName, - request: output.request, - response: output.response, - }); - } - if (opts.usageHandler && output.usage) { - const usageModel = resolveUsageModel(toSave, activeModel); - await opts.usageHandler(ctx, { - userId, - threadId, - agentName: opts.agentName, - model: getModelName(usageModel), - provider: getProviderName(usageModel), - usage: output.usage, - providerMetadata: output.providerMetadata, - }); - } + await persist(toSave, createPendingMessage, { finishStreamId }); + await report(toSave); }, }; } diff --git a/src/vercel/client/steps.test.ts b/src/vercel/client/steps.test.ts new file mode 100644 index 00000000..c56551f4 --- /dev/null +++ b/src/vercel/client/steps.test.ts @@ -0,0 +1,937 @@ +import { describe, expect, test } from "vitest"; +import { v, type ObjectType } from "convex/values"; +import { isStepCount } from "ai"; +import { z } from "zod/v4"; +import { APICallError } from "@ai-sdk/provider"; +import type { LanguageModelV4StreamPart } from "@ai-sdk/provider"; +import { simulateReadableStream } from "ai"; +import type { ActionCtx } from "./types.js"; +import type { MessageDoc } from "../../validators.js"; +import { Agent, createThread, createTool } from "../index.js"; +import { app, components } from "./setup.test.js"; +import { mockModel } from "./mockModel.js"; +import { errorToString } from "./utils.js"; + +const USAGE = { + inputTokens: { total: 3, noCache: 3, cacheRead: 0, cacheWrite: 0 }, + outputTokens: { total: 10, text: 10, reasoning: 0 }, +}; +const CALL = { + type: "tool-call", + toolCallId: "tc-lookup", + toolName: "lookup", + input: "{}", +} as const; + +// The SDK continues past a tool call or a denied approval whatever the first +// step's finish reason says (#388). The first `rejections` requests are +// rejected before anything streams, retryably and without a backoff. +function model( + firstFinish: "tool-calls" | "stop", + failSecond: boolean, + stopAfterError = false, + rejections = 0, + failFirst = false, +) { + const first: LanguageModelV4StreamPart[] = [ + { type: "stream-start", warnings: [] }, + CALL, + ...(failFirst + ? [{ type: "error", error: "provider failed" } as const] + : []), + { + type: "finish", + finishReason: { unified: firstFinish, raw: undefined }, + usage: USAGE, + }, + ]; + const second: LanguageModelV4StreamPart[] = [ + { type: "stream-start", warnings: [] }, + { type: "text-start", id: "t" }, + { type: "text-delta", id: "t", delta: "Looked it up" }, + { type: "text-end", id: "t" }, + ...(failSecond + ? [{ type: "error", error: "provider failed" } as const] + : []), + { + type: "finish", + finishReason: { + unified: failSecond && !stopAfterError ? "error" : "stop", + raw: undefined, + }, + usage: USAGE, + }, + ]; + let calls = 0; + const call = () => { + if (calls++ < rejections) { + throw new APICallError({ + message: "provider rejected", + url: "mock", + requestBodyValues: {}, + isRetryable: true, + responseHeaders: { "retry-after-ms": "0" }, + }); + } + return calls - rejections - 1; + }; + const m = mockModel({ + doStream: async () => ({ + stream: simulateReadableStream({ chunks: call() ? second : first }), + }), + doGenerate: async () => { + if (call() === 0) { + return { + content: [CALL], + finishReason: { unified: firstFinish, raw: undefined }, + usage: USAGE, + warnings: [], + }; + } + if (failSecond) throw new Error("provider failed"); + return { + content: [{ type: "text", text: "Looked it up" }], + finishReason: { unified: "stop", raw: undefined }, + usage: USAGE, + warnings: [], + }; + }, + }); + return { model: m, calls: () => calls }; +} + +// The first attempt fails partway through, and the SDK's stream retry +// answers. +function retriedModel() { + let calls = 0; + const m = mockModel({ + doStream: async () => ({ + stream: simulateReadableStream({ + chunks: (calls++ + ? [ + { type: "stream-start", warnings: [] }, + { type: "text-start", id: "t" }, + { type: "text-delta", id: "t", delta: "Recovered" }, + { type: "text-end", id: "t" }, + { + type: "finish", + finishReason: { unified: "stop", raw: undefined }, + usage: USAGE, + }, + ] + : [ + { type: "stream-start", warnings: [] }, + { type: "text-start", id: "t" }, + { type: "text-delta", id: "t", delta: "Partial" }, + { type: "error", error: "provider failed" }, + ]) as LanguageModelV4StreamPart[], + }), + }), + }); + return { model: m, calls: () => calls }; +} + +type Row = { role?: string; parts: string[]; status: string; error?: string }; + +async function readThread(ctx: ActionCtx, threadId: string) { + const page = await ctx.runQuery( + components.agent.messages.listMessagesByThreadId, + { + threadId, + order: "asc", + paginationOpts: { cursor: null, numItems: 50 }, + }, + ); + return (page.page as MessageDoc[]).map( + (m): Row => ({ + role: m.message?.role, + parts: + typeof m.message?.content === "string" + ? ["text"] + : (m.message?.content ?? []).map((p) => p.type), + status: m.status, + ...(m.error !== undefined ? { error: m.error } : {}), + }), + ); +} + +const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); + +const runArgs = { + threadId: v.string(), + runner: v.union( + v.literal("generateText"), + v.literal("streamText"), + v.literal("returnImmediately"), + v.literal("noDeltas"), + ), + scenario: v.union( + v.literal("toolCalls"), + v.literal("stopWithToolCall"), + v.literal("deniedApproval"), + v.literal("toolError"), + v.literal("approvalRequest"), + ), + // Throws on the save that stores the assistant tool call, or the answer. + failSave: v.optional(v.union(v.literal("first"), v.literal("final"))), + failFirst: v.optional(v.boolean()), + failSecond: v.optional(v.boolean()), + stopAfterError: v.optional(v.boolean()), + // When a retry from the same prompt starts. + restartAt: v.optional( + v.union( + v.literal("duringFirstStep"), + v.literal("afterFirstStep"), + v.literal("beforeAnswer"), + ), + ), + deletePromptAt: v.optional( + v.union(v.literal("duringFirstStep"), v.literal("beforeAnswer")), + ), + abortBeforeStepEnd: v.optional(v.boolean()), + stopFromStepEnd: v.optional(v.boolean()), + stopOnUsage: v.optional(v.boolean()), + stopWhenThrows: v.optional(v.boolean()), + abortInStepEnd: v.optional(v.boolean()), + abortDuringSave: v.optional(v.boolean()), + awaitTextOnly: v.optional(v.boolean()), + // Reads fullStream and then text, as a caller would. + readStream: v.optional(v.boolean()), + retryFirstAttempt: v.optional(v.boolean()), + stopAtAnswer: v.optional(v.boolean()), + rejectRequests: v.optional(v.number()), + maxRetries: v.optional(v.number()), + throwAtFinish: v.optional(v.boolean()), + abortOnCreate: v.optional(v.boolean()), + slowFailure: v.optional(v.boolean()), + saveAll: v.optional(v.boolean()), + // Passes the deprecated onStepFinish and onFinish instead. + deprecatedCallbacks: v.optional(v.boolean()), +}; +type RunArgs = Omit, "threadId">; + +const run = app.action({ + args: runArgs, + handler: async (ctx, args) => { + const m = args.retryFirstAttempt + ? retriedModel() + : model( + args.scenario === "toolCalls" || args.scenario === "toolError" + ? "tool-calls" + : "stop", + args.failSecond ?? false, + args.stopAfterError, + args.rejectRequests, + args.failFirst, + ); + const firstMessageId = async () => { + const { page } = await ctx.runQuery( + components.agent.messages.listMessagesByThreadId, + { + threadId: args.threadId, + order: "asc", + paginationOpts: { cursor: null, numItems: 1 }, + }, + ); + return (page[0] as MessageDoc)._id; + }; + // What a retry from the same prompt does. + const restart = async () => { + await ctx.runMutation(components.agent.messages.addMessages, { + threadId: args.threadId, + promptMessageId: await firstMessageId(), + failPendingSteps: true, + messages: [ + { message: { role: "assistant", content: [] }, status: "pending" }, + ], + }); + }; + const deletePrompt = async () => { + await ctx.runMutation(components.agent.messages.deleteByIds, { + messageIds: [await firstMessageId()], + }); + }; + let usageCalls = 0; + let stepEnds = 0; + let endCalls = 0; + let savedAtEnd: Row[] | undefined; + let stopRequested = false; + const abort = new AbortController(); + let duringSecondStep: Row[] | undefined; + let streamsAtAnswerSave: string[] | undefined; + const agent = new Agent(components.agent, { + name: "steps", + languageModel: m.model, + tools: { + lookup: createTool({ + inputSchema: z.object({}), + needsApproval: args.scenario === "approvalRequest", + execute: async () => { + if (args.restartAt === "duringFirstStep") await restart(); + if (args.deletePromptAt === "duringFirstStep") await deletePrompt(); + if (args.scenario === "toolError") throw new Error("lookup failed"); + return { ok: true }; + }, + }), + }, + usageHandler: async () => { + usageCalls++; + }, + }); + const onStepEnd = async () => { + stepEnds++; + if (args.restartAt === "afterFirstStep" && stepEnds === 1) { + await restart(); + } + stopRequested = !!args.stopFromStepEnd; + if (args.abortInStepEnd) abort.abort("Stopped"); + }; + const onEnd = async () => { + endCalls++; + savedAtEnd = await readThread(ctx, args.threadId); + }; + const onErrors: string[] = []; + const callArgs = { + prompt: "Go.", + ...(args.readStream + ? { + onError: ({ error }: { error: unknown }) => { + onErrors.push(errorToString(error)); + }, + } + : {}), + stopWhen: [ + isStepCount(5), + () => { + if (args.stopWhenThrows) throw new Error("stopWhen failed"); + return stopRequested || (!!args.stopOnUsage && usageCalls > 0); + }, + ], + maxRetries: args.maxRetries ?? 0, + ...(args.retryFirstAttempt ? { streamRetries: 1 } : {}), + ...(args.throwAtFinish + ? { + experimental_transform: () => + new TransformStream({ + transform(part, controller) { + if (part.type === "finish") { + throw new Error("transform failed"); + } + controller.enqueue(part); + }, + }), + } + : {}), + ...(args.stopAtAnswer + ? { + experimental_transform: ({ + stopStream, + }: { + stopStream: () => void; + }) => + new TransformStream({ + transform(part, controller) { + if (part.type === "text-delta") stopStream(); + else controller.enqueue(part); + }, + }), + } + : {}), + ...(args.abortBeforeStepEnd + ? { + // The abort fails the generation before the SDK reaches the + // first step's onStepEnd. + onChunk: async ({ chunk }: { chunk: { type: string } }) => { + if (chunk.type !== "finish-step") return; + await sleep(50); + abort.abort("Stopped"); + await sleep(100); + }, + } + : {}), + abortSignal: abort.signal, + ...(args.scenario === "deniedApproval" + ? { toolApproval: () => ({ type: "denied" as const, reason: "No." }) } + : {}), + prepareStep: async ({ stepNumber }: { stepNumber: number }) => { + if (stepNumber === 1) { + duringSecondStep = await readThread(ctx, args.threadId); + if (args.restartAt === "beforeAnswer") await restart(); + if (args.deletePromptAt === "beforeAnswer") await deletePrompt(); + } + return undefined; + }, + ...(args.deprecatedCallbacks + ? { onStepFinish: onStepEnd, onFinish: onEnd } + : { onStepEnd, onEnd }), + }; + const saves = (args: unknown, part: string) => + (args as { messages?: MessageDoc[] }).messages?.some( + (m) => + m.message?.role === "assistant" && + Array.isArray(m.message.content) && + m.message.content.some((p) => p.type === part), + ) ?? false; + const intercepted: ActionCtx = { + ...ctx, + runMutation: (async (reference, mutationArgs) => { + if ( + args.failSave && + saves(mutationArgs, args.failSave === "first" ? "tool-call" : "text") + ) { + throw new Error(`${args.failSave} save failed`); + } + // Slower than the harness waits after the generation returns. + if ( + args.abortInStepEnd || + (args.awaitTextOnly && saves(mutationArgs, "text")) || + (args.slowFailure && "result" in (mutationArgs as object)) + ) { + await sleep(100); + } + const result = await ctx.runMutation(reference, mutationArgs); + const { pendingMessageId, format } = mutationArgs as { + pendingMessageId?: string; + format?: string; + }; + // Aborts once a step save has committed, before its reply arrives. + if (args.abortDuringSave && pendingMessageId) { + abort.abort("Stopped"); + await sleep(100); + } + // Someone else aborts the stream as soon as it is created, before + // the creation's reply arrives. + if (args.abortOnCreate && format !== undefined) { + await ctx.runMutation(components.agent.streams.abort, { + streamId: result as string, + reason: "Stopped elsewhere", + }); + await sleep(50); + } + if (saves(mutationArgs, "text")) { + const streams = await ctx.runQuery(components.agent.streams.list, { + threadId: args.threadId, + statuses: ["streaming", "finished", "aborted"], + }); + streamsAtAnswerSave = streams.map((s) => s.status); + } + return result; + }) as ActionCtx["runMutation"], + }; + const threadArgs = { threadId: args.threadId }; + const options = args.saveAll + ? { storageOptions: { saveMessages: "all" as const } } + : {}; + let error: string | null = null; + let seen: + | { errorParts: string[]; thrown: string | null; text: string } + | undefined; + try { + if (args.runner === "generateText") { + await agent.generateText(intercepted, threadArgs, callArgs, options); + } else { + const result = await agent.streamText( + intercepted, + threadArgs, + callArgs, + args.runner === "noDeltas" + ? options + : { + ...options, + saveStreamDeltas: { + throttleMs: 0, + returnImmediately: args.runner === "returnImmediately", + }, + }, + ); + if (args.readStream) { + seen = { errorParts: [], thrown: null, text: "" }; + try { + for await (const part of result.fullStream) { + if (part.type === "error") { + seen.errorParts.push(errorToString(part.error)); + } + } + } catch (e) { + seen.thrown = errorToString(e); + } + seen.text = await result.text.then( + (text) => `resolved: ${text}`, + (e) => `rejected: ${errorToString(e)}`, + ); + } else if (args.awaitTextOnly) await result.text; + else await result.consumeStream(); + } + } catch (e) { + error = errorToString(e); + } + // Let a returnImmediately generation's own cleanup land. + if (!args.awaitTextOnly) await sleep(50); + const streams = await ctx.runQuery(components.agent.streams.list, { + threadId: args.threadId, + statuses: ["streaming", "finished", "aborted"], + }); + return { + error, + usageCalls, + stepEnds, + endCalls, + modelCalls: m.calls(), + duringSecondStep: duringSecondStep ?? null, + saved: await readThread(ctx, args.threadId), + savedAtEnd: savedAtEnd ?? null, + streams: streams.map((s) => s.status), + streamsAtAnswerSave: streamsAtAnswerSave ?? null, + ...(seen ? { seen: { ...seen, onErrors } } : {}), + }; + }, +}); + +const { api, createTest } = app.defineModules({ steps: { run } }); + +async function go(args: RunArgs) { + const t = createTest(); + const threadId = await t.run((ctx) => + createThread(ctx, components.agent, {}), + ); + return t.action(api.steps.run, { threadId, ...args }); +} + +const USER: Row = { role: "user", parts: ["text"], status: "success" }; +const PENDING: Row = { role: "assistant", parts: [], status: "pending" }; +const EMPTY: Row = { role: "assistant", parts: [], status: "success" }; +const ANSWER: Row = { role: "assistant", parts: ["text"], status: "success" }; +const LOOKUP: Row[] = [ + { role: "assistant", parts: ["tool-call"], status: "success" }, + { role: "tool", parts: ["tool-result"], status: "success" }, +]; +const DENIED: Row[] = [ + { + role: "assistant", + parts: ["tool-call", "tool-approval-request"], + status: "success", + }, + { + role: "tool", + parts: ["tool-approval-response", "tool-result"], + status: "success", + }, +]; +const failed = (row: Row, error: string): Row => ({ + ...row, + status: "failed", + error, +}); + +const runners = [ + { runner: "generateText", streamed: false, awaited: true }, + { runner: "streamText", streamed: true, awaited: true }, + { runner: "returnImmediately", streamed: true, awaited: false }, + { runner: "noDeltas", streamed: false, awaited: false }, +] as const; + +describe.each(runners)("$runner", ({ runner, streamed, awaited }) => { + test.each([ + { scenario: "toolCalls", first: LOOKUP }, + { scenario: "stopWithToolCall", first: LOOKUP }, + { scenario: "deniedApproval", first: DENIED }, + { scenario: "toolError", first: LOOKUP }, + ] as const)( + "saves every step the SDK runs: $scenario (#388)", + async ({ scenario, first }) => { + const saved = [USER, ...first, ANSWER]; + expect(await go({ runner, scenario })).toEqual({ + error: null, + usageCalls: 2, + stepEnds: 2, + endCalls: 1, + modelCalls: 2, + duringSecondStep: [USER, ...first, PENDING], + saved, + savedAtEnd: saved, + streams: streamed ? ["finished"] : [], + streamsAtAnswerSave: streamed ? ["finished"] : [], + }); + }, + ); + + test("the deprecated callbacks run as onStepEnd and onEnd do", async () => { + const r = await go({ + runner, + scenario: "toolCalls", + deprecatedCallbacks: true, + }); + expect(r.stepEnds).toBe(2); + expect(r.endCalls).toBe(1); + expect(r.savedAtEnd).toEqual([USER, ...LOOKUP, ANSWER]); + }); + + test("saveMessages all saves every step", async () => { + const r = await go({ runner, scenario: "toolCalls", saveAll: true }); + expect(r.saved).toEqual([USER, ...LOOKUP, ANSWER]); + expect(r.streams).toEqual(streamed ? ["finished"] : []); + }); + + test("a tool approval request ends the generation without a pending message", async () => { + const r = await go({ runner, scenario: "approvalRequest" }); + expect(r.modelCalls).toBe(1); + expect(r.saved).toEqual([ + USER, + { + role: "assistant", + parts: ["tool-call", "tool-approval-request"], + status: "success", + }, + ]); + expect(r.streams).toEqual(streamed ? ["finished"] : []); + }); + + test("a failed step save stops the generation (#419)", async () => { + const r = await go({ runner, scenario: "toolCalls", failSave: "first" }); + expect(r.modelCalls).toBe(1); + // On main the save ran first, and its failure skipped the caller's + // onStepEnd. + expect(r.stepEnds).toBe(1); + expect(r.endCalls).toBe(0); + // A streamed step's parts are recovered into the failed message. + expect(r.saved).toEqual( + streamed + ? [USER, ...LOOKUP.map((row) => failed(row, "first save failed"))] + : [USER, failed(EMPTY, "first save failed")], + ); + if (awaited) expect(r.error).toBe("first save failed"); + expect(r.streams).toEqual(streamed ? ["aborted"] : []); + }); + + // An awaited generation rejects with it instead, as above. + test.skipIf(awaited)( + "a failed step save reaches the stream once, as an error part (#419)", + async () => { + const r = await go({ + runner, + scenario: "toolCalls", + failSave: "first", + readStream: true, + }); + // As after any error the SDK streams, the text is the finished steps'. + expect(r.seen).toEqual({ + errorParts: ["first save failed"], + onErrors: ["first save failed"], + thrown: null, + text: "resolved: ", + }); + }, + ); + + test("a failed final save fails the generation (#419)", async () => { + const r = await go({ runner, scenario: "toolCalls", failSave: "final" }); + expect(r.saved).toEqual([ + USER, + ...LOOKUP, + failed(streamed ? ANSWER : EMPTY, "final save failed"), + ]); + if (awaited) expect(r.error).toBe("final save failed"); + expect(r.endCalls).toBe(0); + expect(r.streams).toEqual(streamed ? ["aborted"] : []); + }); + + test("a provider error in the second step keeps the first and aborts the stream", async () => { + const r = await go({ runner, scenario: "toolCalls", failSecond: true }); + // streamText keeps the partial text, and its usage, from the errored step. + expect(r.saved).toEqual([ + USER, + ...LOOKUP, + failed(runner === "generateText" ? EMPTY : ANSWER, "provider failed"), + ]); + expect(r.usageCalls).toBe(runner === "generateText" ? 1 : 2); + expect(r.streams).toEqual(streamed ? ["aborted"] : []); + }); + + test.skipIf(runner === "generateText")( + "a provider error in a step the SDK would continue past stops it there", + async () => { + const r = await go({ runner, scenario: "toolCalls", failFirst: true }); + expect(r.modelCalls).toBe(1); + expect(r.saved).toEqual([ + USER, + ...LOOKUP.map((row) => failed(row, "provider failed")), + ]); + expect(r.error).toBe(null); + expect(r.streams).toEqual(streamed ? ["aborted"] : []); + }, + ); + + test.skipIf(runner === "generateText")( + "a provider error before a normal finish fails the answer", + async () => { + const r = await go({ + runner, + scenario: "toolCalls", + failSecond: true, + stopAfterError: true, + }); + expect(r.saved).toEqual([ + USER, + ...LOOKUP, + failed(ANSWER, "provider failed"), + ]); + expect(r.streams).toEqual(streamed ? ["aborted"] : []); + }, + ); + + test("a rejected request fails the generation", async () => { + const r = await go({ runner, scenario: "toolCalls", rejectRequests: 10 }); + expect(r.saved).toEqual([USER, failed(EMPTY, "provider rejected")]); + expect(r.streams).toEqual(streamed ? ["aborted"] : []); + if (runner === "generateText") expect(r.error).toBe("provider rejected"); + }); + + test("a rejected request is failed before the result rejects", async () => { + const r = await go({ + runner, + scenario: "toolCalls", + rejectRequests: 10, + awaitTextOnly: true, + slowFailure: true, + }); + expect(r.saved).toEqual([USER, failed(EMPTY, "provider rejected")]); + }); + + test("a request the SDK retries saves every step", async () => { + const r = await go({ + runner, + scenario: "toolCalls", + rejectRequests: 1, + maxRetries: 1, + }); + expect(r.modelCalls).toBe(3); + expect(r.saved).toEqual([USER, ...LOOKUP, ANSWER]); + expect(r.streams).toEqual(streamed ? ["finished"] : []); + }); + + test("stopWhen sees what onStepEnd set", async () => { + const r = await go({ + runner, + scenario: "toolCalls", + stopFromStepEnd: true, + }); + expect(r.modelCalls).toBe(1); + expect(r.saved).toEqual([USER, ...LOOKUP]); + }); + + test("stopWhen sees the usage of the step it follows", async () => { + const r = await go({ runner, scenario: "toolCalls", stopOnUsage: true }); + expect(r.modelCalls).toBe(1); + expect(r.saved).toEqual([USER, ...LOOKUP]); + }); + + test("a failure between steps keeps the step the SDK finished", async () => { + const r = await go({ runner, scenario: "toolCalls", stopWhenThrows: true }); + expect(r.modelCalls).toBe(1); + expect(r.saved).toEqual([USER, ...LOOKUP]); + if (awaited) expect(r.error).toBe("stopWhen failed"); + }); + + test.each([ + { failSave: undefined, saved: [USER, ...LOOKUP], error: null }, + { + failSave: "first", + saved: [USER, failed(EMPTY, "Stopped")], + error: "first save failed", + }, + ] as const)( + "an abort in onStepEnd saves the step first, or reports its failed save: $failSave", + async ({ failSave, saved, error }) => { + const r = await go({ + runner, + scenario: "toolCalls", + stopFromStepEnd: true, + abortInStepEnd: true, + failSave, + }); + expect(r.saved).toEqual(saved); + if (awaited) expect(r.error).toBe(error); + }, + ); + + test("an abort during a step save fails the next step", async () => { + const r = await go({ + runner, + scenario: "toolCalls", + abortDuringSave: true, + }); + expect(r.modelCalls).toBe(1); + expect(r.saved).toEqual([USER, ...LOOKUP, failed(EMPTY, "Stopped")]); + }); + + test("the answer is saved before the result resolves", async () => { + const r = await go({ runner, scenario: "toolCalls", awaitTextOnly: true }); + expect(r.saved).toEqual([USER, ...LOOKUP, ANSWER]); + }); + + test("the answer's save finishes its stream (#181)", async () => { + const r = await go({ runner, scenario: "toolCalls" }); + expect(r.streamsAtAnswerSave).toEqual(streamed ? ["finished"] : []); + }); + + test.skipIf(!streamed)( + "a stream aborted while it is being created ends aborted, with nothing pending", + async () => { + const r = await go({ + runner, + scenario: "toolCalls", + abortOnCreate: true, + }); + expect(r.saved).not.toContainEqual(PENDING); + expect(r.streams).toEqual(["aborted"]); + }, + ); + + test("the result rejects when the answer fails to save", async () => { + const r = await go({ + runner, + scenario: "toolCalls", + failSave: "final", + awaitTextOnly: true, + }); + expect(r.error).toBe("final save failed"); + }); + + test.skipIf(runner === "generateText")( + "a stream retry after a provider error saves the answer (#428)", + async () => { + expect( + await go({ runner, scenario: "toolCalls", retryFirstAttempt: true }), + ).toMatchObject({ + error: null, + usageCalls: 1, + stepEnds: 1, + modelCalls: 2, + duringSecondStep: null, + saved: [USER, ANSWER], + streams: streamed ? ["finished"] : [], + }); + }, + ); + + test.skipIf(runner === "generateText")( + "a transform that stops the stream fails the step it cut off before the result resolves", + async () => { + const r = await go({ + runner, + scenario: "toolCalls", + stopAtAnswer: true, + awaitTextOnly: true, + slowFailure: true, + }); + expect(r.modelCalls).toBe(2); + // A streamed step's parts are recovered into its message. + expect(r.saved).toEqual([ + USER, + ...LOOKUP, + failed( + streamed ? ANSWER : EMPTY, + "The generation stopped before its step ended", + ), + ]); + expect(r.streams).toEqual(streamed ? ["aborted"] : []); + // The caller stopped it, so nothing rejects. + expect(r.error).toBe(null); + }, + ); + + test.skipIf(runner === "generateText")( + "a transform that throws at the end keeps the answer", + async () => { + const r = await go({ + runner, + scenario: "toolCalls", + throwAtFinish: true, + }); + expect(r.saved).toEqual([USER, ...LOOKUP, ANSWER]); + expect(r.usageCalls).toBe(2); + if (awaited) expect(r.error).toBe("transform failed"); + }, + ); + + test("a retry from the same prompt stops the attempt it replaced", async () => { + const r = await go({ + runner, + scenario: "toolCalls", + restartAt: "duringFirstStep", + }); + expect(r.modelCalls).toBe(1); + }); + + test.skipIf(runner === "generateText")( + "an abort before onStepEnd saves the step once", + async () => { + const r = await go({ + runner, + scenario: "toolCalls", + abortBeforeStepEnd: true, + }); + expect( + r.saved.filter((row) => row.parts.includes("tool-result")), + ).toHaveLength(streamed ? 1 : 0); + expect(r.modelCalls).toBe(1); + }, + ); + + test("a retry during the answer fails the attempt it replaced", async () => { + const r = await go({ + runner, + scenario: "toolCalls", + restartAt: "beforeAnswer", + }); + expect(r.saved.filter((row) => row.parts.includes("text"))).toEqual([ + USER, + failed(ANSWER, "Restarting"), + ]); + if (awaited) expect(r.error).toBe("Restarting"); + expect(r.streams).toEqual(streamed ? ["aborted"] : []); + }); + + test("a retry between steps stops the attempt it replaced", async () => { + const r = await go({ + runner, + scenario: "toolCalls", + restartAt: "afterFirstStep", + }); + expect(r.modelCalls).toBe(1); + // The step saves after the retry, so it fails like one cut off by it. The + // retry's own pending message is never filled by this harness. + expect(r.saved).toEqual([ + USER, + failed(LOOKUP[0], "Restarting"), + PENDING, + failed(LOOKUP[1], "Restarting"), + failed(EMPTY, "Restarting"), + ]); + }); + + test("a prompt deleted during a step stops the generation (#300)", async () => { + const r = await go({ + runner, + scenario: "toolCalls", + deletePromptAt: "duringFirstStep", + }); + expect(r.modelCalls).toBe(1); + if (awaited) expect(r.error).toBe("The prompt message was deleted"); + }); + + test("a prompt deleted before the answer abandons its save (#300)", async () => { + const r = await go({ + runner, + scenario: "toolCalls", + deletePromptAt: "beforeAnswer", + }); + expect(r.modelCalls).toBe(2); + expect(r.error).toBe(null); + // A streamed generation's stream is aborted with its prompt. + expect(r.saved).toEqual([ + ...LOOKUP, + streamed ? failed(EMPTY, "async abort") : PENDING, + ]); + }); +}); diff --git a/src/vercel/client/streamText.test.ts b/src/vercel/client/streamText.test.ts index 83e90abc..dd922d1c 100644 --- a/src/vercel/client/streamText.test.ts +++ b/src/vercel/client/streamText.test.ts @@ -10,7 +10,7 @@ import { errorToString } from "./utils.js"; const FINAL_TEXT = "Hello from the model"; const PROVIDER_FAILURE_TEXT = "Mock provider failure"; -const CLEANUP_FAILURE_TEXT = "finalizeMessage rejected"; +const CLEANUP_FAILURE_TEXT = "saving the failure rejected"; function hasKeys( value: unknown, @@ -230,7 +230,11 @@ const streamTextCleanupFailure = app.action({ const failingCtx = { ...ctx, runMutation: (async (reference, args) => { - if (hasKeys(args, ["messageId", "result"])) { + if ( + hasKeys(args, ["messageId", "result"]) || + (hasKeys(args, ["failPendingSteps"]) && + JSON.stringify(args).includes('"status":"failed"')) + ) { throw new Error(CLEANUP_FAILURE_TEXT); } return ctx.runMutation(reference, args); @@ -259,6 +263,35 @@ const streamTextCleanupFailure = app.action({ }, }); +const streamTextFullStreamErrors = app.action({ + args: { + threadId: v.string(), + returnImmediately: v.boolean(), + saveMessages: v.boolean(), + }, + handler: async (ctx, { threadId, returnImmediately, saveMessages }) => { + const result = await failingAgent.streamText( + ctx, + { threadId }, + { prompt: "Test", onError: () => {} }, + { + saveStreamDeltas: { returnImmediately, throttleMs: 0 }, + ...(saveMessages ? {} : { storageOptions: { saveMessages: "none" } }), + }, + ); + const errors: string[] = []; + for await (const part of result.fullStream) { + if (part.type === "error") errors.push(errorToString(part.error)); + } + await new Promise((resolve) => setTimeout(resolve, 50)); + const streams = await ctx.runQuery(components.agent.streams.list, { + threadId, + statuses: ["streaming", "finished", "aborted"], + }); + return { errors, streams: streams.map((s) => s.status) }; + }, +}); + const streamTextRejectedRequest = app.action({ args: { threadId: v.string(), @@ -374,6 +407,7 @@ const { api, createTest } = app.defineModules({ streamTextNoStorage, streamTextNoStorageImmediate, streamTextCleanupFailure, + streamTextFullStreamErrors, streamTextRejectedRequest, streamTextAbortedMidStream, streamTextEmptyNoStorageAwaited, @@ -569,6 +603,29 @@ describe("streamText abort cleanup", () => { }); }); +describe("provider errors reach fullStream and abort the stream", () => { + test.each([ + { returnImmediately: false, saveMessages: true }, + { returnImmediately: true, saveMessages: true }, + { returnImmediately: false, saveMessages: false }, + { returnImmediately: true, saveMessages: false }, + ])( + "returnImmediately: $returnImmediately, saveMessages: $saveMessages", + async (args) => { + const t = createTest(); + const threadId = await t.run(async (ctx) => + createThread(ctx, components.agent, { userId: "u1" }), + ); + expect( + await t.action(api.streamText.streamTextFullStreamErrors, { + threadId, + ...args, + }), + ).toEqual({ errors: [PROVIDER_FAILURE_TEXT], streams: ["aborted"] }); + }, + ); +}); + describe("provider rejects the first request (issue #387)", () => { test.each([ ["awaited deltas", { chunking: "word" as const, throttleMs: 250 }], @@ -725,7 +782,7 @@ describe("saveStreamDeltas flushes buffered parts (issue #323)", () => { ).toBe(FINAL_TEXT); }); - test("an out of band abort fails the generation instead of saving it", async () => { + test("an out of band abort keeps the step that finished before it was seen", async () => { const t = createTest(); const threadId = await t.run(async (ctx) => createThread(ctx, components.agent, { userId: "u1" }), @@ -743,7 +800,8 @@ describe("saveStreamDeltas flushes buffered parts (issue #323)", () => { expect(streams).toHaveLength(1); expect(streams[0].status).toBe("aborted"); - // Nobody gets to save a successful message onto a row someone aborted. + // The generation learns of the abort only after its step finished. An + // abort stops what hasn't finished, and doesn't rewrite what has. const messages = await t.run(async (ctx) => agent.listMessages(ctx, { threadId, @@ -754,10 +812,10 @@ describe("saveStreamDeltas flushes buffered parts (issue #323)", () => { messages.page .filter((m) => m.message?.role === "assistant") .map((m) => m.status), - ).toEqual(["failed"]); + ).toEqual(["success"]); }); - test("the awaited path captures the stream-level finish chunk", async () => { + test("the awaited path keeps every part up to the last step end", async () => { const t = createTest(); const threadId = await t.run(async (ctx) => createThread(ctx, components.agent, { userId: "u1" }), @@ -790,12 +848,10 @@ describe("saveStreamDeltas flushes buffered parts (issue #323)", () => { const parts = deltas.flatMap((d) => d.parts); const types = parts.map((p) => p.type); - // Unlike the returnImmediately path, nothing stops accepting parts early - // here: consumeStream drains at EOF, so the trailing chunks the AI SDK - // emits after the last onStepEnd are persisted too. + // The answer's save finishes the stream before the stream-level "finish" + // chunk, as on the returnImmediately path. expect(types.at(0)).toBe("start"); - expect(types.at(-1)).toBe("finish"); - expect(types).toContain("finish-step"); + expect(types.at(-1)).toBe("finish-step"); expect( parts .filter((p) => p.type === "text-delta") diff --git a/src/vercel/client/streamText.ts b/src/vercel/client/streamText.ts index 083cd8cf..083ece4c 100644 --- a/src/vercel/client/streamText.ts +++ b/src/vercel/client/streamText.ts @@ -1,7 +1,7 @@ import type { - ModelMessage, - StepResult, StreamTextResult, + StreamTextTransform, + TextStreamPart, ToolSet, UIMessage as AIUIMessage, } from "ai"; @@ -25,7 +25,7 @@ import type { Output as AISDKOutput } from "ai"; import { startGeneration } from "./start.js"; import type { Agent } from "../index.js"; import { getModelName, getProviderName } from "../../shared.js"; -import { errorToString, willContinue } from "./utils.js"; +import { errorToString } from "./utils.js"; import { materializeUIMessageChunkFiles } from "../fileMaterialization.js"; export async function runStreamCleanup(cleanup: { @@ -109,32 +109,6 @@ export async function streamText< RUNTIME_CONTEXT >(ctx, component, streamTextArgs, options, "streamText"); - const steps: StepResult[] = []; - let firstStreamError: string | undefined; - let streamCleanupFailure: { error: unknown } | undefined; - let initialResponseMessages: ModelMessage[] = []; - let initialResponseMessagesSaved = false; - const responseMessagesForStep = ( - step: StepResult, - ) => [ - ...(initialResponseMessagesSaved ? [] : initialResponseMessages), - ...step.response.messages, - ]; - - // Track the final step for atomic save with stream finish (issue #181). - // Only used when streamText awaits stream consumption itself; the - // `returnImmediately` path saves inline instead (see onStepFinish below). - let pendingFinalStep: - | { - step: StepResult; - responseMessages: ModelMessage[]; - } - | undefined; - - // Whether streamText will await stream consumption before returning. - // When false (saveStreamDeltas.returnImmediately === true), we cannot - // defer the final-step save to a post-await block — the function has - // already returned by the time onStepFinish fires. See issue #265. const savesMessages = options?.storageOptions?.saveMessages !== "none"; const willAwaitStream = Boolean(threadId) && @@ -157,10 +131,6 @@ export async function streamText< materialize: (parts) => materializeUIMessageChunkFiles(ctx, component, parts), abortSignal: args.abortSignal, - // The message save finishes the stream row atomically (issue - // #181) — but only when there is a save. With saveMessages set to - // "none" nothing does, so the streamer keeps finish ownership. - finishHandledExternally: savesMessages, }, { threadId, @@ -176,31 +146,87 @@ export async function streamText< ) : undefined; + const failGeneration = (reason: string) => + runStreamCleanup({ + failCall: () => call.fail(reason), + failStreamer: async () => streamer?.fail(reason), + }); + const endGeneration = (options?: Parameters[0]) => + call.end(options).catch(async (error) => { + await streamer?.fail(errorToString(error)).catch(() => {}); + throw error; + }); + let stepError: string | undefined; + let finished = false; + let stopped = false; + let stepEnds = 0; + const transforms = mergeTransforms( + options?.saveStreamDeltas, + streamTextArgs.experimental_transform, + ).map( + (transform): StreamTextTransform => + (transformOptions) => + transform({ + ...transformOptions, + stopStream: () => { + stopped = true; + transformOptions.stopStream(); + }, + }), + ); + // The generation ends where its stream does, before the result resolves: + // at `finish`, after the last onStepEnd, the last step is saved and its + // stream finished before passing it on. A stream that ends without `finish` + // on an error part (one the SDK did not retry), or because a transform + // stopped it, cut its step off. It ends without either between retries. + const saveLastStep: StreamTextTransform = () => { + let errorPart: { error: unknown } | undefined; + let error: string | undefined; + return new TransformStream, TextStreamPart>({ + async transform(part, controller) { + if (part.type === "error") { + errorPart = part; + error = errorToString(part.error); + } + if (part.type === "finish-step") { + stepEnds++; + stepError = + error ?? + (part.finishReason === "error" + ? "The generation failed" + : undefined); + } + if (part.type === "finish") { + finished = true; + if (stepError !== undefined) await failGeneration(stepError); + else { + await endGeneration({ + finishStreamId: savesMessages + ? await streamer?.handOver(stepEnds) + : undefined, + }); + } + } + controller.enqueue(part); + }, + flush: async () => { + if (!finished && (error !== undefined || stopped)) { + await failGeneration( + error ?? "The generation stopped before its step ended", + ).catch((cleanupError) => { + // An error the stream has already carried isn't thrown again. + if (cleanupError !== errorPart?.error) throw cleanupError; + }); + } + }, + }); + }; + const result = streamTextAi({ ...args, abortSignal: streamer?.abortController.signal ?? args.abortSignal, - experimental_transform: mergeTransforms( - options?.saveStreamDeltas, - streamTextArgs.experimental_transform, - ), - onError: async (error) => { - console.error("onError", error); - const reason = (firstStreamError ??= errorToString(error.error)); - try { - await runStreamCleanup({ - failCall: () => call.fail(reason), - failStreamer: async () => streamer?.fail(reason), - }); - } catch (cleanupError) { - streamCleanupFailure ??= { error: cleanupError }; - console.error("Failed to clean up errored stream:", cleanupError); - } - return streamTextArgs.onError?.(error); - }, + experimental_transform: [...transforms, saveLastStep], onAbort: async (event) => { - const providerTriggeredAbort = - firstStreamError !== undefined && !args.abortSignal?.aborted; - if (providerTriggeredAbort) return; const reason = args.abortSignal?.reason ? errorToString(args.abortSignal.reason) : "streamText aborted"; @@ -211,70 +237,23 @@ export async function streamText< }); }, prepareStep: async (options) => { - if (options.stepNumber === 0) - initialResponseMessages = [...options.responseMessages]; + await call.beforeStep(options); const result = await streamTextArgs.prepareStep?.(options); if (result) { const model = result.model ?? options.model; call.updateModel(model); - // streamer?.updateMetadata({ - // model: getModelName(model), - // provider: getProviderName(model), - // providerOptions: options.messages.at(-1)?.providerOptions, - // }); return result; } return undefined; }, onStepEnd: async (step) => { - steps.push(step); - const createPendingMessage = await willContinue(steps, args.stopWhen); - if (!createPendingMessage && streamer) { - // Final step with streaming enabled. - if (willAwaitStream) { - // We're about to `await stream` below — defer the save so it - // happens atomically with stream finish (issue #181). Don't touch - // the streamer here: the stream-level `finish` chunk is emitted - // after this callback, so the row has to stay `streaming` and keep - // accepting parts until consumeStream reaches EOF, which is the only - // point where everything has actually been handed over. - pendingFinalStep = { - step, - responseMessages: responseMessagesForStep(step), - }; - } else { - // returnImmediately path: streamText is about to return without - // awaiting consumption, so the deferred-save block below won't - // see this step. Save inline now (issue #265). Nothing awaits the - // stream here, so this is the last moment we can drain deltas, see - // flushAndStopAccepting for the window that leaves. - await streamer.flushAndStopAccepting(); - const finishStreamId = await streamer.getOrCreateStreamId({ - ifAborted: "returnUndefined", - }); - if (finishStreamId) { - await call.save( - { step, responseMessages: responseMessagesForStep(step) }, - false, - finishStreamId, - ); - // The save finishes the row only when it stores messages. With - // saveMessages "none" nothing else will, so do it here. - if (!savesMessages) { - await streamer.finish(); - } - initialResponseMessagesSaved = true; - } - } - } else { - await call.save( - { step, responseMessages: responseMessagesForStep(step) }, - createPendingMessage, - ); - initialResponseMessagesSaved = true; - } + await call.endStep(step, stepError); return (streamTextArgs.onStepEnd ?? streamTextArgs.onStepFinish)?.(step); }, + onEnd: async (event) => { + await endGeneration(); + return (streamTextArgs.onEnd ?? streamTextArgs.onFinish)?.(event); + }, } as Parameters< typeof streamTextAi >[0]) as StreamTextResult; @@ -287,52 +266,37 @@ export async function streamText< : DEFAULT_STREAMING_OPTIONS.sendSources, }), ); + // The generation runs to its end whoever reads the stream. + let streamError: { error: unknown } | undefined; + const ended = Promise.resolve( + result.consumeStream({ + onError: (e) => { + streamError ??= { error: e }; + }, + }), + ).then(async () => { + if (streamError) await failGeneration(errorToString(streamError.error)); + await endGeneration(); + }); + void ended.catch(() => {}); if (willAwaitStream) { try { await stream; - await result.consumeStream(); } catch (e) { - // If the stream errored (e.g. onStepFinish threw), the DeltaStreamer's - // finish() was never called, leaving the streaming message stuck in - // "streaming" state. Clean it up by marking it as aborted. - try { - await streamer?.fail(errorToString(e)); - } catch (cleanupError) { - streamCleanupFailure ??= { error: cleanupError }; - } - if (pendingFinalStep) { - try { - await call.save(pendingFinalStep, false); - } catch (saveError) { - console.error("Failed to save deferred final step:", saveError); - } - } + await failGeneration(errorToString(e)).catch(() => {}); throw e; } + await ended; } else { // The caller consumes the provider stream, but this separate background - // consumer can still reject (e.g. if durable abort cleanup fails). - void stream?.catch((error) => { + // consumer can still reject (e.g. if durable abort cleanup fails, or the + // stream errored). + void stream?.catch(async (error) => { console.error("Failed to persist stream deltas:", error); + await call.fail(errorToString(error)).catch(() => {}); }); } - if (streamCleanupFailure) throw streamCleanupFailure.error; - - // If we deferred the final step save, do it now with atomic stream finish. - if (pendingFinalStep && streamer) { - const finishStreamId = await streamer.getOrCreateStreamId({ - ifAborted: "returnUndefined", - }); - if (finishStreamId) { - await call.save(pendingFinalStep, false, finishStreamId); - } - } else if (willAwaitStream && streamer) { - // No final step was deferred (e.g. the generation produced none), so no - // save will finish the stream. The streamer doesn't finish itself, so do - // it here rather than leaving the row to time out. - await streamer.finish(); - } const metadata: GenerationOutputMetadata = { promptMessageId, order, diff --git a/src/vercel/client/streaming.integration.test.ts b/src/vercel/client/streaming.integration.test.ts index 309747b8..105bfdda 100644 --- a/src/vercel/client/streaming.integration.test.ts +++ b/src/vercel/client/streaming.integration.test.ts @@ -835,7 +835,7 @@ describe("Fallback Behavior", () => { test("mergeTransforms adds smoothStream when streaming is enabled", () => { // No streaming options - returns existing transforms - expect(mergeTransforms(undefined, undefined)).toBeUndefined(); + expect(mergeTransforms(undefined, undefined)).toEqual([]); // Boolean true - adds smoothStream const transforms = mergeTransforms(true, undefined); diff --git a/src/vercel/client/streaming.ts b/src/vercel/client/streaming.ts index 38137134..6f74569b 100644 --- a/src/vercel/client/streaming.ts +++ b/src/vercel/client/streaming.ts @@ -181,20 +181,19 @@ export function mergeTransforms( | StreamTextTransform | Array> | undefined, -) { - if (!options) { - return existing; - } - const chunking = - typeof options === "boolean" - ? DEFAULT_STREAMING_OPTIONS.chunking - : options.chunking; +): Array> { const transforms = Array.isArray(existing) - ? existing + ? [...existing] : existing ? [existing] : []; - transforms.push(smoothStream({ delayInMs: null, chunking })); + if (options) { + const chunking = + typeof options === "boolean" + ? DEFAULT_STREAMING_OPTIONS.chunking + : options.chunking; + transforms.push(smoothStream({ delayInMs: null, chunking })); + } return transforms; } @@ -227,8 +226,7 @@ export class DeltaStreamer { public abortController: AbortController; /** * When true, external code finishes the stream row (atomically with the - * message save, issue #181) and `consumeStream` must not. Decided once, at - * construction. + * message save, issue #181) and `consumeStream` must not. */ #finishHandledExternally: boolean; /** @@ -237,6 +235,9 @@ export class DeltaStreamer { * non-streaming row. */ #stoppedAccepting: boolean = false; + #stepEnds = 0; + #sourceDone = false; + #progress: (() => void) | undefined; constructor( public readonly component: AgentComponent, @@ -274,6 +275,9 @@ export class DeltaStreamer { this.#finishHandledExternally = config.finishHandledExternally ?? false; this.#nextParts = []; this.abortController = new AbortController(); + this.abortController.signal.addEventListener("abort", () => + this.#progress?.(), + ); if (config.abortSignal) { const abortFromSignal = () => { void this.#abort("abortSignal").catch(() => { @@ -378,6 +382,10 @@ export class DeltaStreamer { try { for await (const chunk of stream) { await this.addParts([chunk]); + if ((chunk as { type?: string }).type === "finish-step") { + this.#stepEnds++; + this.#progress?.(); + } } } catch (error) { // A provider can throw while responding to an abort. Join the durable @@ -385,6 +393,9 @@ export class DeltaStreamer { // preserving the provider error for the caller. await this.#abort(errorToString(error)).catch(() => {}); throw error; + } finally { + this.#sourceDone = true; + this.#progress?.(); } // Abort cleanup owns the terminal component transition, so consumeStream // must wait for it instead of also trying to finish the stream. @@ -412,6 +423,25 @@ export class DeltaStreamer { } } + /** + * Waits until `stepEnds` step ends have come through the source, writes + * every part before them, and stops taking parts, for the caller to finish + * the stream with its message save (issue #181). Resolves undefined if the + * stream was aborted or never created. + */ + public async handOver(stepEnds: number): Promise { + const signal = this.abortController.signal; + while (this.#stepEnds < stepEnds && !this.#sourceDone && !signal.aborted) { + await new Promise((resolve) => (this.#progress = resolve)); + } + this.#progress = undefined; + this.#finishHandledExternally = true; + await this.flushAndStopAccepting(); + if (signal.aborted) return undefined; + await this.#creatingStreamIdPromise?.catch(() => {}); + return this.streamId; + } + /** * Drain everything currently buffered or in flight, so that after this * resolves no delta write is outstanding and #nextParts is empty. diff --git a/src/vercel/client/utils.test.ts b/src/vercel/client/utils.test.ts index 845a05fc..eb0fbcb5 100644 --- a/src/vercel/client/utils.test.ts +++ b/src/vercel/client/utils.test.ts @@ -1,9 +1,9 @@ import type { StepResult } from "ai"; import { describe, expect, test } from "vitest"; -import { hasSuccessfulToolCall, willContinue } from "./utils.js"; +import { hasSuccessfulToolCall } from "./utils.js"; -// Minimal StepResult builder — only the fields willContinue and -// hasSuccessfulToolCall actually read. Loosely typed on purpose so test +// Minimal StepResult builder — only the fields hasSuccessfulToolCall +// actually reads. Loosely typed on purpose so test // fixtures can be terse; cast at the boundary. type StepFixture = { finishReason?: string; @@ -22,10 +22,6 @@ function makeStep(partial: StepFixture): StepResult { } as unknown as StepResult; } -// Never stops, so `willContinue` returning true means it reached the end -// rather than taking the early `toolCalls > completed` bail. -const neverStop = () => false; - describe("hasSuccessfulToolCall", () => { test("returns true when last step has a tool-result for the named tool", () => { const step = makeStep({ @@ -64,40 +60,3 @@ describe("hasSuccessfulToolCall", () => { expect(hasSuccessfulToolCall("search")({ steps: [] })).toBe(false); }); }); - -describe("willContinue", () => { - test("does not stop when a tool-error fills in for a missing tool-result", async () => { - // Two tool calls; one returns a result, the other errors. - const step = makeStep({ - toolCalls: [ - { toolCallId: "1", toolName: "a" }, - { toolCallId: "2", toolName: "b" }, - ], - toolResults: [{ toolCallId: "1", toolName: "a" }], - content: [ - { type: "tool-result", toolName: "a" }, - { type: "tool-error", toolName: "b" }, - ], - }); - // A stop condition that never stops, so reaching the end returns true. - // With `undefined` the function returns false either way, which cannot - // distinguish "took the early bail" from "ran to the end". - expect(await willContinue([step], neverStop)).toBe(true); - }); - - test("stops when a tool call has neither a result nor an error yet", async () => { - const step = makeStep({ - toolCalls: [{ toolCallId: "1", toolName: "a" }], - toolResults: [], - content: [], - }); - // Same never-stopping condition as above, so a `false` here can only - // come from the early `toolCalls > completed` bail. - expect(await willContinue([step], neverStop)).toBe(false); - }); - - test("stops when finishReason is not tool-calls", async () => { - const step = makeStep({ finishReason: "stop" }); - expect(await willContinue([step], undefined)).toBe(false); - }); -}); diff --git a/src/vercel/client/utils.ts b/src/vercel/client/utils.ts index fc037286..e1907e19 100644 --- a/src/vercel/client/utils.ts +++ b/src/vercel/client/utils.ts @@ -1,5 +1,4 @@ -import type { Context } from "@ai-sdk/provider-utils"; -import type { StepResult, StopCondition, ToolSet } from "ai"; +import type { StopCondition } from "ai"; export { errorToString } from "../../errors.js"; @@ -20,35 +19,3 @@ export function hasSuccessfulToolCall(toolName: string): StopCondition { (p) => p.type === "tool-result" && p.toolName === toolName, ) ?? false; } - -export async function willContinue< - TOOLS extends ToolSet, - RUNTIME_CONTEXT extends Context, ->( - steps: StepResult[], - stopWhen: - | StopCondition - | Array> - | undefined, -): Promise { - const step = steps.at(-1)!; - // we aren't doing another round after a tool result - // TODO: whether to handle continuing after too much context used.. - if (step.finishReason !== "tool-calls") return false; - // Count both successful results and errors as completed outputs. - // Failed tool calls are represented as tool-error content parts, so only - // checking toolResults misses them. The fallback to step.toolResults.length - // is for callers whose steps lack content (mocks); the optional chain is - // defensive rather than load-bearing. - const completedOutputs = - step.content?.filter( - (p) => p.type === "tool-result" || p.type === "tool-error", - ).length ?? step.toolResults.length; - if (step.toolCalls.length > completedOutputs) return false; - if (Array.isArray(stopWhen)) { - return (await Promise.all(stopWhen.map(async (s) => s({ steps })))).every( - (stop) => !stop, - ); - } - return !!stopWhen && !(await stopWhen({ steps })); -} diff --git a/src/vercel/index.ts b/src/vercel/index.ts index e0e74115..6ed923e1 100644 --- a/src/vercel/index.ts +++ b/src/vercel/index.ts @@ -113,11 +113,7 @@ import type { WithoutModel, } from "./client/types.js"; import { streamText } from "./client/streamText.js"; -import { - errorToString, - hasSuccessfulToolCall, - willContinue, -} from "./client/utils.js"; +import { errorToString, hasSuccessfulToolCall } from "./client/utils.js"; export { isStepCount, stepCountIs } from "ai"; export { hasSuccessfulToolCall }; @@ -537,48 +533,54 @@ export class Agent< > & GenerationOutputMetadata > { - const { args, promptMessageId, order, ...call } = await this.start( + type Tools = TOOLS extends undefined ? AgentTools : TOOLS; + const { args, promptMessageId, order, ...call } = await startGeneration< + typeof generateTextArgs, + Tools, + CustomCtx, + RUNTIME_CONTEXT + >( ctx, - generateTextArgs, - { ...threadOpts, ...options }, + this.component, + { + ...generateTextArgs, + tools: (generateTextArgs.tools ?? this.options.tools) as Tools, + instructions: + generateTextArgs.instructions ?? + generateTextArgs.system ?? + this.options.instructions, + stopWhen: (generateTextArgs.stopWhen ?? this.options.stopWhen) as any, + }, + { + ...this.options, + ...threadOpts, + ...options, + agentName: this.options.name, + agentForToolCtx: this, + }, "generateText", ); - - type Tools = TOOLS extends undefined ? AgentTools : TOOLS; - const steps: StepResult[] = []; - let initialResponseMessages: ModelMessage[] = []; - let initialResponseMessagesSaved = false; try { const result = (await generateText({ ...args, prepareStep: async (options) => { - if (options.stepNumber === 0) { - initialResponseMessages = [...options.responseMessages]; - } + await call.beforeStep(options); const result = await generateTextArgs.prepareStep?.(options); call.updateModel(result?.model ?? options.model); return result; }, onStepEnd: async (step) => { - steps.push(step); - await call.save( - { - step, - responseMessages: [ - ...(initialResponseMessagesSaved - ? [] - : initialResponseMessages), - ...step.response.messages, - ], - }, - await willContinue(steps, args.stopWhen), - ); - initialResponseMessagesSaved = true; + await call.endStep(step); return ( generateTextArgs.onStepEnd ?? generateTextArgs.onStepFinish )?.(step); }, + onEnd: async (event) => { + await call.end(); + return (generateTextArgs.onEnd ?? generateTextArgs.onFinish)?.(event); + }, })) as GenerateTextResult; + await call.end(); const metadata: GenerationOutputMetadata = { promptMessageId, order,