From 8d9c78278e1e9dc25d5b7069f06ad84601f4424f Mon Sep 17 00:00:00 2001 From: Sweets Sweetman Date: Sat, 1 Aug 2026 10:15:13 -0500 Subject: [PATCH 1/4] fix(workflows): one workflow step per LLM call so runs stop blowing the 800s ceiling runAgentStep wrapped the entire agent loop via stopWhen: stepCountIs(111), so a single "use step" ran 11-25 minutes. WDK deploys step handlers with maxDuration: max, which resolves to 800s on Pro, so Vercel killed the invocation and the step queue redelivered it. Confirmed on prod: Step "step//./app/lib/workflows/runAgentStep//runAgentStep" exceeded max retries (4 retries) [USER_ERROR] with the step recorded at attempt: 5 and an empty {"message":"Unknown error"} (the signature of a platform kill, not a thrown error). Each attempt was a complete agent run that mailed the customer again: 45 of the last 100 runs failed this way, every one at ~72.5 min = 5 attempts x ~870s. Moves the loop into the workflow body, one journaled step per LLM call: - runAgentStep drops stopWhen; the AI SDK default isStepCount(1) bounds it to a single model call plus that call's tool executions. It now takes modelMessages/originalMessages and returns responseMessages for threading. - runAgentWorkflow owns the loop, appending each iteration's responseMessages so iteration N+1 sees iteration N's tool results, and bounding it at CHAT_AGENT_MAX_ITERATIONS. - The turn's stream envelope moves up: sendStreamStart/sendStreamFinish are workflow-level steps and each iteration passes sendStart/sendFinish: false, so the client renders one assistant message instead of one per iteration. - buildMessageMetadataCallback takes a seed so usage/cost totals span the whole turn rather than resetting each iteration. - convertMessagesStep runs the conversion once, journaled, before the loop. CHAT_AGENT_STOP_WHEN stays: getGeneralAgent (the non-durable /api/chat route) still runs its tool loop inside streamText. Mirrors the reference in vercel-labs/open-agents apps/web/app/workflows/chat.ts. Refs recoupable/chat#1918 Co-Authored-By: Claude Opus 5 (1M context) --- .../workflows/__tests__/runAgentStep.test.ts | 90 +++++++-- .../__tests__/runAgentWorkflow.test.ts | 24 +++ .../__tests__/runAgentWorkflowLoop.test.ts | 182 ++++++++++++++++++ app/lib/workflows/convertMessagesStep.ts | 18 ++ app/lib/workflows/runAgentStep.ts | 102 +++++++--- app/lib/workflows/runAgentWorkflow.ts | 100 ++++++++-- app/lib/workflows/sendStreamFinish.ts | 20 ++ app/lib/workflows/sendStreamStart.ts | 25 +++ .../buildMessageMetadataCallback.ts | 19 +- lib/chat/const.ts | 14 ++ 10 files changed, 520 insertions(+), 74 deletions(-) create mode 100644 app/lib/workflows/__tests__/runAgentWorkflowLoop.test.ts create mode 100644 app/lib/workflows/convertMessagesStep.ts create mode 100644 app/lib/workflows/sendStreamFinish.ts create mode 100644 app/lib/workflows/sendStreamStart.ts diff --git a/app/lib/workflows/__tests__/runAgentStep.test.ts b/app/lib/workflows/__tests__/runAgentStep.test.ts index a9faf935e..9a989bc6b 100644 --- a/app/lib/workflows/__tests__/runAgentStep.test.ts +++ b/app/lib/workflows/__tests__/runAgentStep.test.ts @@ -78,6 +78,7 @@ function makeStreamResult(opts?: { }, ), finishReason: Promise.resolve("stop"), + responseMessages: Promise.resolve([]), }; } @@ -92,7 +93,10 @@ function makeWritable() { } const baseInput = { - messages: [ + // The workflow body now owns conversion and threading, so the step takes + // model messages directly plus the UI message it is appending to. + modelMessages: [{ role: "user" as const, content: "hi" }], + originalMessages: [ { id: "m1", role: "user" as const, @@ -201,33 +205,76 @@ describe("runAgentStep", () => { } }); - it("wires a prepareStep callback that marks the last message with cacheControl", async () => { + // With one model call per step there is no in-call step boundary left for + // `prepareStep` to hook, so cacheControl is applied to the messages handed + // to streamText directly. + it("marks the last model message with cacheControl before passing it to streamText", async () => { vi.mocked(streamText).mockReturnValue(makeStreamResult() as never); const { stream } = makeWritable(); - await runAgentStep({ ...baseInput, writable: stream } as never); + await runAgentStep({ + ...baseInput, + modelMessages: [ + { role: "user", content: "first" }, + { role: "user", content: "second" }, + ], + writable: stream, + } as never); const args = vi.mocked(streamText).mock.calls[0]?.[0] as { - prepareStep?: (opts: { - messages: Array<{ role: string; providerOptions?: Record }>; - model: unknown; - steps?: unknown[]; - }) => { messages?: unknown[] } | undefined; + prepareStep?: unknown; + messages: Array<{ providerOptions?: Record }>; }; - expect(typeof args.prepareStep).toBe("function"); - const anthropicModel = { provider: "anthropic", modelId: "claude-haiku-4.5" } as never; - const result = args.prepareStep!({ - messages: [ - { role: "user", content: "first" } as never, - { role: "user", content: "second" } as never, - ], - model: anthropicModel, - steps: [], + expect(args.prepareStep).toBeUndefined(); + expect(args.messages[0]?.providerOptions).toBeUndefined(); + expect(args.messages[1]?.providerOptions).toEqual({ + anthropic: { cacheControl: { type: "ephemeral" } }, }); - const out = result?.messages as Array<{ providerOptions?: Record }>; - expect(out).toBeDefined(); - expect(out[0]?.providerOptions).toBeUndefined(); - expect(out[1]?.providerOptions).toEqual({ anthropic: { cacheControl: { type: "ephemeral" } } }); + }); + + // The decomposition contract: one model call per step. If a `stopWhen` + // creeps back in, the step regrows past Vercel's 800 s ceiling and the + // duplicate-email failure returns (chat#1918). + it("does NOT set stopWhen, so the AI SDK default bounds the step to one model call", async () => { + vi.mocked(streamText).mockReturnValue(makeStreamResult() as never); + const { stream } = makeWritable(); + + await runAgentStep({ ...baseInput, writable: stream } as never); + + const args = vi.mocked(streamText).mock.calls[0]?.[0] as { stopWhen?: unknown }; + expect(args.stopWhen).toBeUndefined(); + }); + + it("suppresses per-iteration start/finish chunks so the turn renders as one message", async () => { + const streamOpts: Array<{ sendStart?: boolean; sendFinish?: boolean }> = []; + vi.mocked(streamText).mockReturnValue({ + toUIMessageStream: vi.fn((opts: { sendStart?: boolean; sendFinish?: boolean }) => { + streamOpts.push(opts); + return (async function* () {})(); + }), + finishReason: Promise.resolve("stop"), + responseMessages: Promise.resolve([]), + } as never); + const { stream } = makeWritable(); + + await runAgentStep({ ...baseInput, writable: stream } as never); + + expect(streamOpts[0]?.sendStart).toBe(false); + expect(streamOpts[0]?.sendFinish).toBe(false); + }); + + it("returns responseMessages so the workflow can thread them into the next iteration", async () => { + const produced = [{ role: "assistant", content: "tool call" }]; + vi.mocked(streamText).mockReturnValue({ + toUIMessageStream: vi.fn(() => (async function* () {})()), + finishReason: Promise.resolve("tool-calls"), + responseMessages: Promise.resolve(produced), + } as never); + const { stream } = makeWritable(); + + const result = await runAgentStep({ ...baseInput, writable: stream } as never); + + expect(result.responseMessages).toEqual(produced); }); it("the wired callback returns undefined for non-finish-step parts", async () => { @@ -366,6 +413,7 @@ describe("runAgentStep", () => { })(), ), finishReason: Promise.resolve("length"), + responseMessages: Promise.resolve([]), } as never); const { stream } = makeWritable(); diff --git a/app/lib/workflows/__tests__/runAgentWorkflow.test.ts b/app/lib/workflows/__tests__/runAgentWorkflow.test.ts index be1adf5d0..14d1b2f8e 100644 --- a/app/lib/workflows/__tests__/runAgentWorkflow.test.ts +++ b/app/lib/workflows/__tests__/runAgentWorkflow.test.ts @@ -23,6 +23,12 @@ vi.mock("@/app/lib/workflows/closeChatStream", () => ({ vi.mock("@/app/lib/workflows/generateAssistantMessageId", () => ({ generateAssistantMessageId: vi.fn(), })); +// The loop's supporting steps — exercised in runAgentWorkflowLoop.test.ts. +vi.mock("@/app/lib/workflows/convertMessagesStep", () => ({ + convertMessagesStep: vi.fn(() => Promise.resolve([])), +})); +vi.mock("@/app/lib/workflows/sendStreamStart", () => ({ sendStreamStart: vi.fn() })); +vi.mock("@/app/lib/workflows/sendStreamFinish", () => ({ sendStreamFinish: vi.fn() })); vi.mock("@/lib/credits/handleChatCredits", () => ({ handleChatCredits: vi.fn(), })); @@ -76,6 +82,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -98,6 +105,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -117,6 +125,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -129,6 +138,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -151,6 +161,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -163,6 +174,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -178,6 +190,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -218,6 +231,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: responseMessage as never, }); @@ -247,6 +261,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: responseMessage as never, }); @@ -266,6 +281,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -292,6 +308,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: responseMessageWithMetadata, }); @@ -308,6 +325,7 @@ describe("runAgentWorkflow", () => { responseMessage: responseMessageWithMetadata, finishReason: "stop", aborted: false, + responseMessages: [], sessionId: "session-1", sessionTitle: "test session", repoOwner: "recoupable", @@ -321,6 +339,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: responseMessageWithMetadata, }); @@ -334,6 +353,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: responseMessageWithMetadata, }); @@ -350,6 +370,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -378,6 +399,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: true, + responseMessages: [], responseMessage: abortedResponseMessage as never, }); @@ -398,6 +420,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: true, + responseMessages: [], responseMessage: abortedResponseMessage as never, }); @@ -410,6 +433,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: true, + responseMessages: [], responseMessage: abortedResponseMessage as never, }); diff --git a/app/lib/workflows/__tests__/runAgentWorkflowLoop.test.ts b/app/lib/workflows/__tests__/runAgentWorkflowLoop.test.ts new file mode 100644 index 000000000..009851b4f --- /dev/null +++ b/app/lib/workflows/__tests__/runAgentWorkflowLoop.test.ts @@ -0,0 +1,182 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { runAgentWorkflow } from "@/app/lib/workflows/runAgentWorkflow"; +import { runAgentStep } from "@/app/lib/workflows/runAgentStep"; +import { convertMessagesStep } from "@/app/lib/workflows/convertMessagesStep"; +import { sendStreamStart } from "@/app/lib/workflows/sendStreamStart"; +import { sendStreamFinish } from "@/app/lib/workflows/sendStreamFinish"; +import { generateAssistantMessageId } from "@/app/lib/workflows/generateAssistantMessageId"; +import { CHAT_AGENT_MAX_ITERATIONS } from "@/lib/chat/const"; + +vi.mock("@/app/lib/workflows/runAgentStep", () => ({ runAgentStep: vi.fn() })); +vi.mock("@/app/lib/workflows/convertMessagesStep", () => ({ convertMessagesStep: vi.fn() })); +vi.mock("@/app/lib/workflows/sendStreamStart", () => ({ sendStreamStart: vi.fn() })); +vi.mock("@/app/lib/workflows/sendStreamFinish", () => ({ sendStreamFinish: vi.fn() })); +vi.mock("@/app/lib/workflows/generateAssistantMessageId", () => ({ + generateAssistantMessageId: vi.fn(), +})); +vi.mock("@/app/lib/workflows/deleteEphemeralKeyStep", () => ({ + deleteEphemeralKeyStep: vi.fn(), +})); +vi.mock("@/app/lib/workflows/closeChatStream", () => ({ closeChatStream: vi.fn() })); +vi.mock("@/lib/chat/clearChatActiveStream", () => ({ clearChatActiveStream: vi.fn() })); +vi.mock("@/lib/credits/handleChatCredits", () => ({ handleChatCredits: vi.fn() })); +vi.mock("@/lib/chat/auto-commit/autoCommitChatTurn", () => ({ autoCommitChatTurn: vi.fn() })); + +const writableStub = new WritableStream(); +vi.mock("workflow", () => ({ + getWritable: vi.fn(() => writableStub), + getWorkflowMetadata: vi.fn(() => ({ + workflowRunId: "wrun_loop_test", + workflowName: "runAgentWorkflow", + workflowStartedAt: new Date(0), + url: "https://example.invalid/workflow", + })), +})); + +const baseInput = { + messages: [{ id: "m1", role: "user", parts: [{ type: "text", text: "hi" }] } as never], + chatId: "chat-1", + sessionId: "session-1", + accountId: "acc-1", + modelId: "anthropic/claude-haiku-4.5", + agentContext: { + sandbox: { state: { type: "vercel" }, workingDirectory: "/sandbox/mono" }, + } as never, +}; + +/** An assistant message carrying `n` parts, used to assert threading. */ +const assistantMessage = (id: string) => + ({ id, role: "assistant", parts: [{ type: "text", text: id }] }) as never; + +/** Queue a sequence of runAgentStep results, one per loop iteration. */ +function queueSteps(results: Array>) { + const mocked = vi.mocked(runAgentStep); + results.forEach(r => mocked.mockResolvedValueOnce(r as never)); + // Anything beyond the queued results ends the loop, so a bug that + // over-iterates fails on a count assertion rather than hanging. + mocked.mockResolvedValue({ + finishReason: "stop", + aborted: false, + responseMessage: undefined, + responseMessages: [], + } as never); +} + +beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(generateAssistantMessageId).mockResolvedValue("asst-loop-id"); + vi.mocked(convertMessagesStep).mockResolvedValue([{ role: "user", content: "hi" }] as never); +}); + +describe("runAgentWorkflow — per-iteration agent loop", () => { + it("keeps iterating while a step finishes on tool-calls, and stops once it finishes on stop", async () => { + queueSteps([ + { + finishReason: "tool-calls", + aborted: false, + responseMessage: assistantMessage("a1"), + responseMessages: [{ role: "assistant", content: "call-1" }], + }, + { + finishReason: "tool-calls", + aborted: false, + responseMessage: assistantMessage("a2"), + responseMessages: [{ role: "assistant", content: "call-2" }], + }, + { + finishReason: "stop", + aborted: false, + responseMessage: assistantMessage("a3"), + responseMessages: [{ role: "assistant", content: "done" }], + }, + ]); + + await runAgentWorkflow(baseInput); + + expect(runAgentStep).toHaveBeenCalledTimes(3); + }); + + it("threads each iteration's responseMessages into the next iteration's modelMessages", async () => { + queueSteps([ + { + finishReason: "tool-calls", + aborted: false, + responseMessage: assistantMessage("a1"), + responseMessages: [{ role: "assistant", content: "call-1" }], + }, + { + finishReason: "stop", + aborted: false, + responseMessage: assistantMessage("a2"), + responseMessages: [{ role: "assistant", content: "done" }], + }, + ]); + + await runAgentWorkflow(baseInput); + + const secondCall = vi.mocked(runAgentStep).mock.calls[1][0] as { + modelMessages: Array<{ role: string; content: string }>; + }; + expect(secondCall.modelMessages).toEqual([ + { role: "user", content: "hi" }, + { role: "assistant", content: "call-1" }, + ]); + }); + + it("emits exactly one stream start and one stream finish across a multi-iteration turn", async () => { + queueSteps([ + { + finishReason: "tool-calls", + aborted: false, + responseMessage: assistantMessage("a1"), + responseMessages: [{ role: "assistant", content: "call-1" }], + }, + { + finishReason: "tool-calls", + aborted: false, + responseMessage: assistantMessage("a2"), + responseMessages: [{ role: "assistant", content: "call-2" }], + }, + { + finishReason: "stop", + aborted: false, + responseMessage: assistantMessage("a3"), + responseMessages: [], + }, + ]); + + await runAgentWorkflow(baseInput); + + expect(sendStreamStart).toHaveBeenCalledTimes(1); + expect(sendStreamStart).toHaveBeenCalledWith(writableStub, "asst-loop-id"); + expect(sendStreamFinish).toHaveBeenCalledTimes(1); + }); + + it("stops looping when a step reports the user aborted, even on tool-calls", async () => { + queueSteps([ + { + finishReason: "tool-calls", + aborted: true, + responseMessage: assistantMessage("a1"), + responseMessages: [{ role: "assistant", content: "call-1" }], + }, + ]); + + await runAgentWorkflow(baseInput); + + expect(runAgentStep).toHaveBeenCalledTimes(1); + }); + + it("bounds a runaway tool-call loop at CHAT_AGENT_MAX_ITERATIONS", async () => { + vi.mocked(runAgentStep).mockResolvedValue({ + finishReason: "tool-calls", + aborted: false, + responseMessage: assistantMessage("a"), + responseMessages: [{ role: "assistant", content: "again" }], + } as never); + + await runAgentWorkflow(baseInput); + + expect(runAgentStep).toHaveBeenCalledTimes(CHAT_AGENT_MAX_ITERATIONS); + }); +}); diff --git a/app/lib/workflows/convertMessagesStep.ts b/app/lib/workflows/convertMessagesStep.ts new file mode 100644 index 000000000..17bf475cb --- /dev/null +++ b/app/lib/workflows/convertMessagesStep.ts @@ -0,0 +1,18 @@ +import { convertToModelMessages, type ModelMessage, type UIMessage } from "ai"; + +/** + * Convert the turn's UI messages to model messages, once, before the loop. + * + * A `"use step"` for two reasons: `convertToModelMessages` can perform I/O + * (it downloads file parts), which is illegal in workflow context; and the + * result is journaled, so a replay reuses it instead of re-downloading. + * + * The workflow body owns the resulting array and appends each iteration's + * `responseMessages` to it, which is how iteration N+1 sees iteration N's + * tool results. + */ +export async function convertMessagesStep(messages: UIMessage[]): Promise { + "use step"; + + return convertToModelMessages(messages); +} diff --git a/app/lib/workflows/runAgentStep.ts b/app/lib/workflows/runAgentStep.ts index 730b4e40e..4f63e5b97 100644 --- a/app/lib/workflows/runAgentStep.ts +++ b/app/lib/workflows/runAgentStep.ts @@ -1,17 +1,17 @@ import { streamText, - convertToModelMessages, createUIMessageStream, + type ModelMessage, type UIMessage, type UIMessageChunk, } from "ai"; import { gateway } from "@ai-sdk/gateway"; import { agentCustomInstructions } from "@/lib/chat/agentCustomInstructions"; import { buildAgentSystemPrompt } from "@/lib/chat/buildAgentSystemPrompt"; -import { CHAT_AGENT_STOP_WHEN } from "@/lib/chat/const"; import { buildAgentTools } from "@/lib/agent/buildAgentTools"; import type { AgentContext, DurableAgentContext } from "@/lib/agent/tools/AgentContext"; import { buildMessageMetadataCallback } from "@/lib/agent/messageMetadata/buildMessageMetadataCallback"; +import type { AgentMessageMetadata } from "@/lib/agent/messageMetadata/AgentMessageMetadata"; import { addCacheControlToTools } from "@/lib/agent/contextManagement/addCacheControlToTools"; import { addCacheControlToMessages } from "@/lib/agent/contextManagement/addCacheControlToMessages"; import { wrapToolsWithAbort } from "@/lib/agent/contextManagement/wrapToolsWithAbort"; @@ -22,7 +22,19 @@ import { getWorkflowMetadata } from "workflow"; import { pipeWorkflowStreamWithStopDetection } from "@/lib/chat/pipeWorkflowStreamWithStopDetection"; export type RunAgentStepInput = { - messages: UIMessage[]; + /** + * Conversation so far, in model form. Owned by `runAgentWorkflow`, which + * appends each iteration's `responseMessages` before the next call — that + * is how iteration N+1 sees iteration N's tool results. + */ + modelMessages: ModelMessage[]; + /** + * The UI-form messages this iteration appends to. When the last entry is + * the in-progress assistant message, the AI SDK keeps building THAT + * message rather than starting a new one, so `responseMessage` comes back + * cumulative and each persist overwrites one row. + */ + originalMessages: UIMessage[]; modelId: string; writable: WritableStream; /** Target chat for persisting the assistant message as it streams. */ @@ -74,32 +86,42 @@ export type RunAgentStepResult = { * `runAgentWorkflow` can charge credits from `responseMessage.metadata`. */ responseMessage: UIMessage | undefined; + /** + * Model-form messages this iteration produced (assistant turn + any tool + * results). `runAgentWorkflow` appends these to `modelMessages` so the + * next iteration continues the conversation instead of repeating it. + */ + responseMessages: ModelMessage[]; /** True when the user stopped the run; `runAgentWorkflow` skips billing + auto-commit on abort. */ aborted: boolean; }; /** - * One LLM turn (with internal tool-call iteration) in the chat workflow. + * ONE LLM call (plus that call's tool executions) in the chat workflow. * Runs as a Vercel Workflow `"use step"` so: * * - Sandbox-banned APIs (`fetch`, `setTimeout`, `crypto`) are legal inside. - * - The result is cached as a single durable event — replays after a crash - * do not re-bill the model or re-execute tools. + * - The result is journaled — a replay resumes from the last completed + * iteration rather than re-billing the model and re-running tools. * - * `streamText` drives the tool-call → tool-result → next-LLM-call loop - * internally using its default stop condition. Our outer workflow stays - * single-turn for now — multi-turn message threading lands when the rest - * of the tool surface ports in a follow-up PR. + * Deliberately does NOT set `stopWhen`: the AI SDK default is + * `isStepCount(1)`, so this returns after a single model call and the + * tool-call → tool-result → next-call loop is driven by `runAgentWorkflow`'s + * body instead. That is the whole point of the decomposition — a step that + * wrapped the full loop ran 11-25 minutes, blew Vercel's 800 s function + * ceiling, and was retried 4 times, mailing the customer once per attempt + * (chat#1918). * - * @param input - Messages + selected model + writable stream + agent context. - * @returns finishReason plus the assembled assistant message. + * @param input - Model messages + selected model + writable stream + agent context. + * @returns finishReason, the cumulative assistant message, and this + * iteration's response messages for threading into the next. */ export async function runAgentStep(input: RunAgentStepInput): Promise { "use step"; console.log("[runAgentStep] start", { modelId: input.modelId, - messageCount: input.messages.length, + messageCount: input.modelMessages.length, hasSandboxState: Boolean(input.agentContext.sandbox?.state), }); @@ -108,11 +130,8 @@ export async function runAgentStep(input: RunAgentStepInput): Promise ({ - messages: addCacheControlToMessages({ messages, model }), - }), }); // `messageMetadata` emits {modelId, usage, cost} chunks the UI renders as - // model/cost badges. - const messageMetadata = buildMessageMetadataCallback({ modelId: input.modelId }); + // model/cost badges. Seeded from the in-progress assistant message so the + // running totals span the whole turn rather than resetting each iteration. + const previousMessage = input.originalMessages.at(-1); + const previousMetadata = + previousMessage?.role === "assistant" + ? (previousMessage.metadata as AgentMessageMetadata | undefined) + : undefined; + const messageMetadata = buildMessageMetadataCallback({ + modelId: input.modelId, + seed: previousMetadata, + }); // createUIMessageStream exposes onStepFinish/onFinish (toUIMessageStream // only has onFinish), so the assistant message is persisted after every @@ -188,6 +214,16 @@ export async function runAgentStep(input: RunAgentStepInput): Promise input.assistantMessageId, + // Continue building the in-progress assistant message rather than + // starting a new one, so `responseMessage` stays cumulative across + // iterations and each persist overwrites a single row. + originalMessages: input.originalMessages, + // The turn's `start`/`finish` chunks are emitted ONCE by the + // workflow body (`sendStreamStart` / `sendStreamFinish`). Without + // this the client would see one per iteration and render N + // assistant messages instead of one. + sendStart: false, + sendFinish: false, }), ); }, @@ -205,12 +241,17 @@ export async function runAgentStep(input: RunAgentStepInput): Promise {}); + void Promise.resolve(result.responseMessages).catch(() => {}); } else { - finishReason = await result.finishReason; + [finishReason, responseMessages] = await Promise.all([ + result.finishReason, + result.responseMessages, + ]); } // On user-stop, close any tool-call parts the step boundary left open and @@ -222,7 +263,8 @@ export async function runAgentStep(input: RunAgentStepInput): Promise> | undefined; + + // The agent loop lives HERE, in the workflow body, with one journaled + // step per LLM call. It used to live inside `streamText` via + // `stopWhen: stepCountIs(111)`, which made a single step run 11-25 + // minutes — past Vercel's 800 s function ceiling, so the platform killed + // it and the queue retried it 4 times, each attempt a full agent run + // that mailed the customer again (chat#1918). + for (let iteration = 0; iteration < CHAT_AGENT_MAX_ITERATIONS; iteration++) { + result = await runAgentStep({ + // Snapshot, not the live array — each iteration's input is a durable + // step input and must describe the conversation as it was at THAT + // call, unaffected by later appends. + modelMessages: [...modelMessages], + originalMessages: [pendingAssistantResponse], + modelId: input.modelId, + chatId: input.chatId, + accountId: input.accountId, + artistId: input.artistId, + interactive: input.interactive, + agentContext: input.agentContext, + writable, + assistantMessageId, + }); + + if (result.responseMessage) pendingAssistantResponse = result.responseMessage; + modelMessages.push(...result.responseMessages); + + // A turn continues only while the model asked for more tools. Any + // other finish reason — and any user stop — ends it. + if (result.aborted || result.finishReason !== "tool-calls") break; + } + + console.log("[runAgentWorkflow] finish", { finishReason: result?.finishReason }); + + await sendStreamFinish(writable); + + // The assistant message is persisted per iteration inside `runAgentStep`, + // so it's not written here. We still use the accumulated message to // charge the account for this turn: atomic wallet debit + audit row via // the `deduct_credits_with_audit` Postgres function (`handleChatCredits` // → `recordCreditDeduction`). // + // `pendingAssistantResponse.metadata` carries the totals for the WHOLE + // turn, not just the last iteration — `runAgentStep` seeds each + // iteration's metadata callback from the message it was handed, so the + // running totals survive the step boundaries. + // // Charge on user-stop too — the provider already billed us for the // tokens consumed, and the assistant message (including partial tool // runs) is persisted, so the user owes the charge regardless of how - // the turn ended. `result.responseMessage.metadata` carries the - // usage actually consumed up to the abort point. - if (result.responseMessage) { - const metadata = result.responseMessage.metadata as AgentMessageMetadata | undefined; + // the turn ended. + if (result?.responseMessage) { + const metadata = pendingAssistantResponse.metadata as AgentMessageMetadata | undefined; await handleChatCredits({ accountId: input.accountId, model: input.modelId, @@ -133,13 +194,14 @@ export async function runAgentWorkflow(input: RunAgentWorkflowInput): Promise): Promise { + "use step"; + + const writer = writable.getWriter(); + try { + await writer.write({ type: "finish" }); + } finally { + writer.releaseLock(); + } +} diff --git a/app/lib/workflows/sendStreamStart.ts b/app/lib/workflows/sendStreamStart.ts new file mode 100644 index 000000000..6c203155d --- /dev/null +++ b/app/lib/workflows/sendStreamStart.ts @@ -0,0 +1,25 @@ +import type { UIMessageChunk } from "ai"; + +/** + * Emit the single `start` chunk that opens an assistant turn. + * + * Owned by the workflow body rather than by `runAgentStep`: the loop runs + * one step per LLM call, and each step's `toUIMessageStream` is told + * `sendStart: false`. Without this hoist the client would see one `start` + * per iteration and render N assistant messages instead of one. + * + * Runs as a `"use step"` because stream I/O is illegal in workflow context. + */ +export async function sendStreamStart( + writable: WritableStream, + messageId: string, +): Promise { + "use step"; + + const writer = writable.getWriter(); + try { + await writer.write({ type: "start", messageId }); + } finally { + writer.releaseLock(); + } +} diff --git a/lib/agent/messageMetadata/buildMessageMetadataCallback.ts b/lib/agent/messageMetadata/buildMessageMetadataCallback.ts index 07225fd6a..5ae6e6491 100644 --- a/lib/agent/messageMetadata/buildMessageMetadataCallback.ts +++ b/lib/agent/messageMetadata/buildMessageMetadataCallback.ts @@ -19,12 +19,23 @@ import type { AgentStepFinishMetadata } from "@/lib/agent/messageMetadata/AgentS * Each call to `buildMessageMetadataCallback` returns a FRESH closure — * one per assistant turn — so totals reset between turns. */ -export function buildMessageMetadataCallback(opts: { modelId: string }) { +export function buildMessageMetadataCallback(opts: { + modelId: string; + /** + * Running totals carried over from earlier iterations of the same turn. + * + * `runAgentWorkflow` now runs one `runAgentStep` per LLM call, and each + * call builds a fresh closure — without a seed the badges would reset to + * this iteration's numbers and under-report the turn. Pass the in-progress + * assistant message's metadata to keep the totals cumulative. + */ + seed?: Pick; +}) { let lastStepUsage: LanguageModelUsage | undefined; - let totalMessageUsage: LanguageModelUsage | undefined; + let totalMessageUsage: LanguageModelUsage | undefined = opts.seed?.totalMessageUsage; let lastStepCost: number | undefined; - let totalMessageCost: number | undefined; - let stepFinishReasons: AgentStepFinishMetadata[] = []; + let totalMessageCost: number | undefined = opts.seed?.totalMessageCost; + let stepFinishReasons: AgentStepFinishMetadata[] = [...(opts.seed?.stepFinishReasons ?? [])]; return function messageMetadata({ part, diff --git a/lib/chat/const.ts b/lib/chat/const.ts index 7e1772190..cf7f1b4d4 100644 --- a/lib/chat/const.ts +++ b/lib/chat/const.ts @@ -13,6 +13,20 @@ export const MAX_MESSAGES = 55; */ export const CHAT_AGENT_STOP_WHEN = stepCountIs(111); +/** + * Upper bound on agent-loop iterations in `runAgentWorkflow`. + * + * The durable workflow loops in its own body with ONE `"use step"` per LLM + * call, so this replaces `CHAT_AGENT_STOP_WHEN` for that path — the stop + * condition moved out of `streamText` and into the workflow. Same 111 for + * behavioural parity: high enough that normal flows never hit it, low + * enough to bound a runaway loop. + * + * `CHAT_AGENT_STOP_WHEN` stays for the non-durable `/api/chat` route + * (`getGeneralAgent`), which still runs its tool loop inside `streamText`. + */ +export const CHAT_AGENT_MAX_ITERATIONS = 111; + export const SYSTEM_PROMPT = `You are Recoup, a friendly, sharp, and strategic AI assistant for the music industry. You help music executives, artist teams, and self-starting artists analyze fan data, optimize marketing, and grow artist careers. --- From b81fa0f137591653a4c49ed9b4856a2e30bfdc2f Mon Sep 17 00:00:00 2001 From: Sweets Sweetman Date: Sat, 1 Aug 2026 10:25:18 -0500 Subject: [PATCH 2/4] fix(workflows): read threaded messages via result.response.messages (ai@6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first pass used `result.responseMessages`, which does not exist on StreamTextResult in ai@6.0.190 — it is an ai@7 accessor. `next build` caught it; local checks did not, because the dev node_modules had ai@7.0.2 installed against a package.json that pins 6.0.190. In 6.0.190 the equivalent is `(await result.response).messages`: the assistant message for this call plus any tool-result message, which is exactly what the next iteration needs appended. Verified against a clean `pnpm install --frozen-lockfile` (ai@6.0.190): the build's TypeScript step passes and it proceeds to page-data collection. Refs recoupable/chat#1918 Co-Authored-By: Claude Opus 5 (1M context) --- app/lib/workflows/__tests__/runAgentStep.test.ts | 8 ++++---- app/lib/workflows/runAgentStep.ts | 12 +++++++----- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/app/lib/workflows/__tests__/runAgentStep.test.ts b/app/lib/workflows/__tests__/runAgentStep.test.ts index 9a989bc6b..a88748d0f 100644 --- a/app/lib/workflows/__tests__/runAgentStep.test.ts +++ b/app/lib/workflows/__tests__/runAgentStep.test.ts @@ -78,7 +78,7 @@ function makeStreamResult(opts?: { }, ), finishReason: Promise.resolve("stop"), - responseMessages: Promise.resolve([]), + response: Promise.resolve({ messages: [] }), }; } @@ -253,7 +253,7 @@ describe("runAgentStep", () => { return (async function* () {})(); }), finishReason: Promise.resolve("stop"), - responseMessages: Promise.resolve([]), + response: Promise.resolve({ messages: [] }), } as never); const { stream } = makeWritable(); @@ -268,7 +268,7 @@ describe("runAgentStep", () => { vi.mocked(streamText).mockReturnValue({ toUIMessageStream: vi.fn(() => (async function* () {})()), finishReason: Promise.resolve("tool-calls"), - responseMessages: Promise.resolve(produced), + response: Promise.resolve({ messages: produced }), } as never); const { stream } = makeWritable(); @@ -413,7 +413,7 @@ describe("runAgentStep", () => { })(), ), finishReason: Promise.resolve("length"), - responseMessages: Promise.resolve([]), + response: Promise.resolve({ messages: [] }), } as never); const { stream } = makeWritable(); diff --git a/app/lib/workflows/runAgentStep.ts b/app/lib/workflows/runAgentStep.ts index 4f63e5b97..c4c982b1e 100644 --- a/app/lib/workflows/runAgentStep.ts +++ b/app/lib/workflows/runAgentStep.ts @@ -246,12 +246,14 @@ export async function runAgentStep(input: RunAgentStepInput): Promise {}); - void Promise.resolve(result.responseMessages).catch(() => {}); + void Promise.resolve(result.response).catch(() => {}); } else { - [finishReason, responseMessages] = await Promise.all([ - result.finishReason, - result.responseMessages, - ]); + // `result.response.messages` is the assistant message for this call plus + // any tool-result message — exactly what the next iteration needs + // appended. (`result.responseMessages` is ai@7; this repo is on 6.0.190.) + const [reason, response] = await Promise.all([result.finishReason, result.response]); + finishReason = reason; + responseMessages = response.messages; } // On user-stop, close any tool-call parts the step boundary left open and From c14fd92b016a765626ee24f913026186f4014086 Mon Sep 17 00:00:00 2001 From: Sweets Sweetman Date: Sat, 1 Aug 2026 11:41:58 -0500 Subject: [PATCH 3/4] fix(workflows): keep tool calls in the transcript across iterations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Caught on the preview, not by the unit tests. A 13-iteration run persisted an assistant message with only 2 parts (step-start + text) — every tool call was gone from chat_messages. The outer createUIMessageStream is what assembles the message handed to onStepFinish/onFinish, and it was not given originalMessages. Per the ai@6 docs that field is what puts the stream in "persistence mode", so without it each iteration rebuilt the message from its own chunks alone and the final text-only persist overwrote every tool call earlier in the turn. Passing originalMessages to the inner toUIMessageStream was not enough. Adds a regression test asserting createUIMessageStream is in persistence mode, since this failure is invisible to a green unit suite. Refs recoupable/chat#1918 Co-Authored-By: Claude Opus 5 (1M context) --- app/lib/workflows/__tests__/runAgentStep.test.ts | 15 +++++++++++++++ app/lib/workflows/runAgentStep.ts | 6 ++++++ 2 files changed, 21 insertions(+) diff --git a/app/lib/workflows/__tests__/runAgentStep.test.ts b/app/lib/workflows/__tests__/runAgentStep.test.ts index a88748d0f..584f871b6 100644 --- a/app/lib/workflows/__tests__/runAgentStep.test.ts +++ b/app/lib/workflows/__tests__/runAgentStep.test.ts @@ -53,6 +53,7 @@ vi.mock("workflow/api", () => ({ // tests can drive its onStepFinish / onFinish callbacks directly. type CreateOpts = { generateId?: () => string; + originalMessages?: unknown[]; onStepFinish?: (e: { responseMessage: unknown }) => unknown; onFinish?: (e: { responseMessage: unknown }) => unknown; execute?: (a: { writer: { write: () => void; merge: () => void; onError: undefined } }) => void; @@ -263,6 +264,20 @@ describe("runAgentStep", () => { expect(streamOpts[0]?.sendFinish).toBe(false); }); + // Regression guard for the transcript loss caught on the preview: the OUTER + // createUIMessageStream is what assembles the message handed to + // onStepFinish/onFinish, so it needs originalMessages too. Without it each + // iteration rebuilds the message from only its own chunks and the final + // text-only persist overwrites every tool call made earlier in the turn. + it("puts createUIMessageStream in persistence mode so each persist keeps earlier iterations' parts", async () => { + vi.mocked(streamText).mockReturnValue(makeStreamResult() as never); + const { stream } = makeWritable(); + + await runAgentStep({ ...baseInput, writable: stream } as never); + + expect(capturedCreateOpts.originalMessages).toBe(baseInput.originalMessages); + }); + it("returns responseMessages so the workflow can thread them into the next iteration", async () => { const produced = [{ role: "assistant", content: "tool call" }]; vi.mocked(streamText).mockReturnValue({ diff --git a/app/lib/workflows/runAgentStep.ts b/app/lib/workflows/runAgentStep.ts index c4c982b1e..b482b6bcd 100644 --- a/app/lib/workflows/runAgentStep.ts +++ b/app/lib/workflows/runAgentStep.ts @@ -197,6 +197,12 @@ export async function runAgentStep(input: RunAgentStepInput): Promise({ generateId: () => input.assistantMessageId, + // Persistence mode: this outer stream is what assembles the message + // handed to onStepFinish/onFinish, so it needs the in-progress assistant + // message too — not just the inner `toUIMessageStream`. Without it every + // iteration rebuilds the message from its own chunks alone, and the final + // text-only persist wipes every tool call made earlier in the turn. + originalMessages: input.originalMessages, onStepFinish: ({ responseMessage: stepMessage }) => { responseMessage = stepMessage; return persistAssistantMessage(input.chatId, stepMessage); From 8b6be5a7d785b6bd3876684313e378c2ff67edbc Mon Sep 17 00:00:00 2001 From: Sweets Sweetman Date: Sat, 1 Aug 2026 15:08:46 -0500 Subject: [PATCH 4/4] refactor(workflows): match upstream open-agents streaming shape, drop the wrapper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Aligns runAgentStep with vercel-labs/open-agents apps/web/app/workflows/chat.ts rather than keeping our own variant of it. The variant is what produced the transcript loss caught on the preview in c14fd92b. - Drops the outer createUIMessageStream. Upstream iterates `result.toUIMessageStream({...})` directly and writes each part to the shared writable with getWriter/write/releaseLock. Our wrapper existed only to get onStepFinish for in-step persistence, which fires once per step now that a step is one model call — and it had to be put in "persistence mode" separately from the inner stream, which is exactly what was missed. - Moves persistence to the workflow body via persistAssistantMessageStep, mirroring upstream's persistAssistantMessage(chatId, pendingAssistantResponse). runAgentStep no longer takes chatId at all. - Replaces pipeWorkflowStreamWithStopDetection with upstream's isAbortError check around the for-await, plus isRunCancelled to preserve the one case upstream does not have: run.cancel() closes our writable, so a write can fail with an unrelated error before the poller notices. - finalizeAbortedAssistantMessage folded into the step as closeOpenToolCalls; the body does the persisting. Deleted as dead: pipeWorkflowStreamWithStopDetection, finalizeAbortedAssistantMessage. Full suite 4307 pass; build's TypeScript step passes. Refs recoupable/chat#1918 Co-Authored-By: Claude Opus 5 (1M context) --- .../workflows/__tests__/runAgentStep.test.ts | 329 ++---------------- .../__tests__/runAgentStepStreaming.test.ts | 118 +++++++ .../__tests__/runAgentWorkflow.test.ts | 11 +- .../__tests__/runAgentWorkflowLoop.test.ts | 39 +++ .../workflows/persistAssistantMessageStep.ts | 26 ++ app/lib/workflows/runAgentStep.ts | 166 ++++----- app/lib/workflows/runAgentWorkflow.ts | 10 +- .../finalizeAbortedAssistantMessage.test.ts | 39 --- ...ipeWorkflowStreamWithStopDetection.test.ts | 66 ---- lib/chat/finalizeAbortedAssistantMessage.ts | 24 -- lib/chat/isAbortError.ts | 11 + lib/chat/isRunCancelled.ts | 21 ++ .../pipeWorkflowStreamWithStopDetection.ts | 64 ---- 13 files changed, 333 insertions(+), 591 deletions(-) create mode 100644 app/lib/workflows/__tests__/runAgentStepStreaming.test.ts create mode 100644 app/lib/workflows/persistAssistantMessageStep.ts delete mode 100644 lib/chat/__tests__/finalizeAbortedAssistantMessage.test.ts delete mode 100644 lib/chat/__tests__/pipeWorkflowStreamWithStopDetection.test.ts delete mode 100644 lib/chat/finalizeAbortedAssistantMessage.ts create mode 100644 lib/chat/isAbortError.ts create mode 100644 lib/chat/isRunCancelled.ts delete mode 100644 lib/chat/pipeWorkflowStreamWithStopDetection.ts diff --git a/app/lib/workflows/__tests__/runAgentStep.test.ts b/app/lib/workflows/__tests__/runAgentStep.test.ts index 584f871b6..83f6ac205 100644 --- a/app/lib/workflows/__tests__/runAgentStep.test.ts +++ b/app/lib/workflows/__tests__/runAgentStep.test.ts @@ -1,13 +1,10 @@ import { describe, it, expect, vi, beforeEach } from "vitest"; -import { streamText, createUIMessageStream } from "ai"; +import { streamText } from "ai"; import { runAgentStep } from "@/app/lib/workflows/runAgentStep"; -import { persistAssistantMessage } from "@/lib/chat/persistAssistantMessage"; -import { pollWorkflowCancellation } from "@/lib/chat/pollWorkflowCancellation"; -import { getRun } from "workflow/api"; vi.mock("ai", async () => { const actual = await vi.importActual("ai"); - return { ...actual, streamText: vi.fn(), createUIMessageStream: vi.fn() }; + return { ...actual, streamText: vi.fn() }; }); // Avoid pulling in real gateway / fetch surface. @@ -15,10 +12,6 @@ vi.mock("@ai-sdk/gateway", () => ({ gateway: vi.fn((modelId: string) => ({ modelId, __mock: "gateway" })), })); -vi.mock("@/lib/chat/persistAssistantMessage", () => ({ - persistAssistantMessage: vi.fn(), -})); - // runAgentStep now reads workflowRunId via getWorkflowMetadata() and polls // getRun(runId).status to source its abort signal. Stub both so the tests // don't pull in the workflow runtime. @@ -49,35 +42,35 @@ vi.mock("workflow/api", () => ({ })), })); -// Captures the options runAgentStep passes to createUIMessageStream so -// tests can drive its onStepFinish / onFinish callbacks directly. -type CreateOpts = { - generateId?: () => string; +type StreamOpts = { + messageMetadata?: unknown; + generateMessageId?: unknown; originalMessages?: unknown[]; - onStepFinish?: (e: { responseMessage: unknown }) => unknown; onFinish?: (e: { responseMessage: unknown }) => unknown; - execute?: (a: { writer: { write: () => void; merge: () => void; onError: undefined } }) => void; }; -let capturedCreateOpts: CreateOpts; function makeStreamResult(opts?: { metadataCalls?: Array; generateIdCalls?: Array; + streamOptsOut?: Array; + emitResponseMessage?: unknown; }) { const calls = opts?.metadataCalls ?? []; const genCalls = opts?.generateIdCalls ?? []; return { - toUIMessageStream: vi.fn( - (streamOpts: { messageMetadata?: unknown; generateMessageId?: unknown }) => { - // Capture the callbacks so tests can inspect them. - calls.push(streamOpts.messageMetadata); - genCalls.push(streamOpts.generateMessageId); - return (async function* () { - yield { type: "start" }; - yield { type: "finish" }; - })(); - }, - ), + toUIMessageStream: vi.fn((streamOpts: StreamOpts) => { + // Capture the callbacks so tests can inspect them. + calls.push(streamOpts.messageMetadata); + genCalls.push(streamOpts.generateMessageId); + opts?.streamOptsOut?.push(streamOpts); + return (async function* () { + yield { type: "text-start", id: "t1" }; + yield { type: "text-end", id: "t1" }; + if (opts && "emitResponseMessage" in opts) { + streamOpts.onFinish?.({ responseMessage: opts.emitResponseMessage }); + } + })(); + }), finishReason: Promise.resolve("stop"), response: Promise.resolve({ messages: [] }), }; @@ -105,7 +98,6 @@ const baseInput = { }, ], modelId: "anthropic/claude-haiku-4.5", - chatId: "chat-1", agentContext: { sandbox: { state: { type: "vercel" }, workingDirectory: "/sandbox/mono" }, }, @@ -113,23 +105,7 @@ const baseInput = { }; describe("runAgentStep", () => { - beforeEach(() => { - vi.clearAllMocks(); - // Default: capture the options, run execute (so toUIMessageStream — and - // its messageMetadata callback — is exercised), and return an empty - // stream that closes immediately so pipeTo resolves. - vi.mocked(createUIMessageStream).mockImplementation((opts: never) => { - capturedCreateOpts = opts as CreateOpts; - capturedCreateOpts.execute?.({ - writer: { write: () => {}, merge: () => {}, onError: undefined }, - }); - return new ReadableStream({ - start(controller) { - controller.close(); - }, - }) as never; - }); - }); + beforeEach(() => vi.clearAllMocks()); it("wires a messageMetadata callback into toUIMessageStream", async () => { const captured: unknown[] = []; @@ -264,20 +240,6 @@ describe("runAgentStep", () => { expect(streamOpts[0]?.sendFinish).toBe(false); }); - // Regression guard for the transcript loss caught on the preview: the OUTER - // createUIMessageStream is what assembles the message handed to - // onStepFinish/onFinish, so it needs originalMessages too. Without it each - // iteration rebuilds the message from only its own chunks and the final - // text-only persist overwrites every tool call made earlier in the turn. - it("puts createUIMessageStream in persistence mode so each persist keeps earlier iterations' parts", async () => { - vi.mocked(streamText).mockReturnValue(makeStreamResult() as never); - const { stream } = makeWritable(); - - await runAgentStep({ ...baseInput, writable: stream } as never); - - expect(capturedCreateOpts.originalMessages).toBe(baseInput.originalMessages); - }); - it("returns responseMessages so the workflow can thread them into the next iteration", async () => { const produced = [{ role: "assistant", content: "tool call" }]; vi.mocked(streamText).mockReturnValue({ @@ -304,30 +266,6 @@ describe("runAgentStep", () => { expect(cb({ part: { type: "start" } })).toBeUndefined(); }); - it("persists the assistant message on each step via onStepFinish", async () => { - vi.mocked(streamText).mockReturnValue(makeStreamResult() as never); - const { stream } = makeWritable(); - - await runAgentStep({ ...baseInput, writable: stream } as never); - - const msg = { id: "a1", role: "assistant", parts: [{ type: "text", text: "partial" }] }; - await capturedCreateOpts.onStepFinish?.({ responseMessage: msg }); - - expect(persistAssistantMessage).toHaveBeenCalledWith("chat-1", msg); - }); - - it("persists the final assistant message via onFinish", async () => { - vi.mocked(streamText).mockReturnValue(makeStreamResult() as never); - const { stream } = makeWritable(); - - await runAgentStep({ ...baseInput, writable: stream } as never); - - const msg = { id: "a1", role: "assistant", parts: [{ type: "text", text: "done" }] }; - await capturedCreateOpts.onFinish?.({ responseMessage: msg }); - - expect(persistAssistantMessage).toHaveBeenCalledWith("chat-1", msg); - }); - it("forwards assistantMessageId into toUIMessageStream's generateMessageId (stable row id)", async () => { const generateIdCalls: unknown[] = []; vi.mocked(streamText).mockReturnValue(makeStreamResult({ generateIdCalls }) as never); @@ -345,20 +283,6 @@ describe("runAgentStep", () => { expect(gen()).toBe("asst-from-workflow-xyz"); }); - it("sets a stable generateId on the createUIMessageStream", async () => { - vi.mocked(streamText).mockReturnValue(makeStreamResult() as never); - const { stream } = makeWritable(); - - await runAgentStep({ - ...baseInput, - writable: stream, - assistantMessageId: "asst-from-workflow-xyz", - } as never); - - expect(typeof capturedCreateOpts.generateId).toBe("function"); - expect(capturedCreateOpts.generateId!()).toBe("asst-from-workflow-xyz"); - }); - it("returns the finishReason from the model result", async () => { vi.mocked(streamText).mockReturnValue(makeStreamResult() as never); const { stream } = makeWritable(); @@ -375,18 +299,9 @@ describe("runAgentStep", () => { parts: [{ type: "text", text: "Hello" }], metadata: { totalMessageCost: 0.05 }, }; - vi.mocked(streamText).mockReturnValue(makeStreamResult() as never); - vi.mocked(createUIMessageStream).mockImplementation((opts: never) => { - const o = opts as CreateOpts; - o.execute?.({ writer: { write: () => {}, merge: () => {}, onError: undefined } }); - // Drive onFinish so runAgentStep captures the final message. - void o.onFinish?.({ responseMessage: emitted }); - return new ReadableStream({ - start(controller) { - controller.close(); - }, - }) as never; - }); + vi.mocked(streamText).mockReturnValue( + makeStreamResult({ emitResponseMessage: emitted }) as never, + ); const { stream } = makeWritable(); const result = await runAgentStep({ ...baseInput, writable: stream } as never); @@ -438,200 +353,4 @@ describe("runAgentStep", () => { expect(result.finishReason).toBe("length"); }); }); - - describe("user-abort path", () => { - it("returns { aborted: true, finishReason: 'stop' } when the poller fires", async () => { - // Poller fires synchronously — controller is aborted by the time pipeTo runs. - vi.mocked(pollWorkflowCancellation).mockImplementation( - (_runId: string, controller: AbortController) => { - controller.abort(); - return { stop: vi.fn(), done: Promise.resolve() }; - }, - ); - vi.mocked(streamText).mockReturnValue(makeStreamResult() as never); - const { stream } = makeWritable(); - - const result = await runAgentStep({ ...baseInput, writable: stream } as never); - - expect(result.aborted).toBe(true); - expect(result.finishReason).toBe("stop"); - }); - - it("does not await result.finishReason on abort (would deadlock if unresolved)", async () => { - vi.mocked(pollWorkflowCancellation).mockImplementation( - (_runId: string, controller: AbortController) => { - controller.abort(); - return { stop: vi.fn(), done: Promise.resolve() }; - }, - ); - // finishReason here is a forever-pending Promise — if runAgentStep awaited it - // on the abort path, the test would hang. - const neverResolves = new Promise(() => {}); - vi.mocked(streamText).mockReturnValue({ - toUIMessageStream: vi.fn(() => - (async function* () { - yield { type: "start" }; - yield { type: "finish" }; - })(), - ), - finishReason: neverResolves, - } as never); - const { stream } = makeWritable(); - - const result = await runAgentStep({ ...baseInput, writable: stream } as never); - - expect(result.finishReason).toBe("stop"); - expect(result.aborted).toBe(true); - }); - - it("re-persists with closed tool-error parts when aborting mid-tool-call", async () => { - // onStepFinish runs while the step is still emitting (a tool-call was - // streamed in this step). The step's captured responseMessage has the - // tool-call in input-available, with no terminal output chunk yet. - const openMessage = { - id: "asst-test-id", - role: "assistant", - parts: [ - { type: "text", text: "running a tool..." }, - { - type: "tool-bash", - toolCallId: "t-open", - state: "input-available", - input: { cmd: "sleep 30" }, - }, - ], - }; - - vi.mocked(createUIMessageStream).mockImplementationOnce((opts: never) => { - capturedCreateOpts = opts as CreateOpts; - capturedCreateOpts.execute?.({ - writer: { write: () => {}, merge: () => {}, onError: undefined }, - }); - // Drive onStepFinish synchronously so responseMessage is populated - // before the abort path runs in runAgentStep. - capturedCreateOpts.onStepFinish?.({ responseMessage: openMessage }); - return new ReadableStream({ - start(controller) { - controller.close(); - }, - }) as never; - }); - - vi.mocked(pollWorkflowCancellation).mockImplementation( - (_runId: string, controller: AbortController) => { - controller.abort(); - return { stop: vi.fn(), done: Promise.resolve() }; - }, - ); - vi.mocked(streamText).mockReturnValue(makeStreamResult() as never); - const { stream } = makeWritable(); - - const result = await runAgentStep({ ...baseInput, writable: stream } as never); - - expect(result.aborted).toBe(true); - - // Two persists: the step's onStepFinish, then the abort-path re-persist - // with closed tool parts. - expect(persistAssistantMessage).toHaveBeenCalledTimes(2); - const second = vi.mocked(persistAssistantMessage).mock.calls[1]?.[1] as { - parts: Array<{ type: string; state?: string; errorText?: string }>; - }; - const toolPart = second.parts.find(p => p.type === "tool-bash")!; - expect(toolPart.state).toBe("output-error"); - expect(toolPart.errorText).toBe("Cancelled"); - // responseMessage on the returned result should be the closed version. - expect(result.responseMessage).toBe(second); - }); - - it("does NOT re-persist when there are no open tool-call parts at abort", async () => { - const closedMessage = { - id: "asst-test-id", - role: "assistant", - parts: [{ type: "text", text: "done text" }], - }; - - vi.mocked(createUIMessageStream).mockImplementationOnce((opts: never) => { - capturedCreateOpts = opts as CreateOpts; - capturedCreateOpts.execute?.({ - writer: { write: () => {}, merge: () => {}, onError: undefined }, - }); - capturedCreateOpts.onStepFinish?.({ responseMessage: closedMessage }); - return new ReadableStream({ - start(controller) { - controller.close(); - }, - }) as never; - }); - - vi.mocked(pollWorkflowCancellation).mockImplementation( - (_runId: string, controller: AbortController) => { - controller.abort(); - return { stop: vi.fn(), done: Promise.resolve() }; - }, - ); - vi.mocked(streamText).mockReturnValue(makeStreamResult() as never); - const { stream } = makeWritable(); - - await runAgentStep({ ...baseInput, writable: stream } as never); - - // Just the onStepFinish persist — no re-persist needed. - expect(persistAssistantMessage).toHaveBeenCalledTimes(1); - }); - - it("detects runtime-cancel even when pipeTo resolves cleanly (writable closed by runtime)", async () => { - // The poller never fires — pipeTo completes naturally because the - // runtime closed the destination writable when run.cancel() landed. - // runAgentStep must still detect this via getRun().status and mark - // the result as aborted so the abort-path re-persist runs. - vi.mocked(getRun).mockReturnValueOnce({ - get status() { - return Promise.resolve("cancelled"); - }, - cancel: vi.fn(() => Promise.resolve()), - } as never); - - vi.mocked(streamText).mockReturnValue(makeStreamResult() as never); - const { stream } = makeWritable(); - - const result = await runAgentStep({ ...baseInput, writable: stream } as never); - - expect(result.aborted).toBe(true); - }); - - it("attaches .catch to result.finishReason so a late rejection is not unhandled", async () => { - vi.mocked(pollWorkflowCancellation).mockImplementation( - (_runId: string, controller: AbortController) => { - controller.abort(); - return { stop: vi.fn(), done: Promise.resolve() }; - }, - ); - const rejection = new Error("finishReason late reject"); - vi.mocked(streamText).mockReturnValue({ - toUIMessageStream: vi.fn(() => - (async function* () { - yield { type: "start" }; - yield { type: "finish" }; - })(), - ), - finishReason: Promise.reject(rejection), - } as never); - - // Catch unhandled rejections globally for the duration of this test. - const unhandled: unknown[] = []; - const onUnhandled = (e: Event) => { - unhandled.push((e as PromiseRejectionEvent).reason); - }; - process.on("unhandledRejection", onUnhandled); - try { - const { stream } = makeWritable(); - await runAgentStep({ ...baseInput, writable: stream } as never); - // Give microtasks a chance to flush. - await new Promise(r => setTimeout(r, 10)); - } finally { - process.off("unhandledRejection", onUnhandled); - } - - expect(unhandled).not.toContain(rejection); - }); - }); }); diff --git a/app/lib/workflows/__tests__/runAgentStepStreaming.test.ts b/app/lib/workflows/__tests__/runAgentStepStreaming.test.ts new file mode 100644 index 000000000..5effcab3c --- /dev/null +++ b/app/lib/workflows/__tests__/runAgentStepStreaming.test.ts @@ -0,0 +1,118 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { streamText } from "ai"; +import { runAgentStep } from "@/app/lib/workflows/runAgentStep"; +import { persistAssistantMessage } from "@/lib/chat/persistAssistantMessage"; + +vi.mock("ai", async () => { + const actual = await vi.importActual("ai"); + return { ...actual, streamText: vi.fn() }; +}); +vi.mock("@ai-sdk/gateway", () => ({ + gateway: vi.fn((modelId: string) => ({ modelId, __mock: "gateway" })), +})); +vi.mock("@/lib/chat/persistAssistantMessage", () => ({ + persistAssistantMessage: vi.fn(), +})); +vi.mock("workflow", () => ({ + getWorkflowMetadata: vi.fn(() => ({ + workflowRunId: "test-run-id", + workflowName: "test", + workflowStartedAt: new Date(0), + url: "https://example.test", + })), +})); +vi.mock("@/lib/chat/pollWorkflowCancellation", () => ({ + pollWorkflowCancellation: vi.fn(() => ({ stop: vi.fn(), done: Promise.resolve() })), +})); +vi.mock("workflow/api", () => ({ + getRun: vi.fn(() => ({ + get status() { + return Promise.resolve("running"); + }, + })), +})); + +const baseInput = { + modelMessages: [{ role: "user" as const, content: "hi" }], + originalMessages: [ + { id: "m1", role: "user" as const, parts: [{ type: "text" as const, text: "hi" }] }, + ], + modelId: "anthropic/claude-haiku-4.5", + agentContext: { + sandbox: { state: { type: "vercel" }, workingDirectory: "/sandbox/mono" }, + }, + assistantMessageId: "asst-test-id", +}; + +function makeWritable() { + const written: unknown[] = []; + const stream = new WritableStream({ + write(chunk) { + written.push(chunk); + }, + }); + return { stream, written }; +} + +beforeEach(() => vi.clearAllMocks()); + +describe("runAgentStep — streaming shape (mirrors upstream open-agents)", () => { + it("writes every stream part straight to the shared writable", async () => { + const parts = [ + { type: "text-start", id: "t1" }, + { type: "text-delta", id: "t1", delta: "hello" }, + { type: "text-end", id: "t1" }, + ]; + vi.mocked(streamText).mockReturnValue({ + toUIMessageStream: vi.fn(() => + (async function* () { + for (const p of parts) yield p; + })(), + ), + finishReason: Promise.resolve("stop"), + response: Promise.resolve({ messages: [] }), + } as never); + const { stream, written } = makeWritable(); + + await runAgentStep({ ...baseInput, writable: stream } as never); + + expect(written).toEqual(parts); + }); + + // Persistence moved to the workflow body. Keeping it in the step is what + // required the outer createUIMessageStream wrapper, and that wrapper is + // what dropped every tool call from the transcript (chat#1918). + it("does not persist from inside the step", async () => { + vi.mocked(streamText).mockReturnValue({ + toUIMessageStream: vi.fn(() => (async function* () {})()), + finishReason: Promise.resolve("stop"), + response: Promise.resolve({ messages: [] }), + } as never); + const { stream } = makeWritable(); + + await runAgentStep({ ...baseInput, writable: stream } as never); + + expect(persistAssistantMessage).not.toHaveBeenCalled(); + }); + + it("reports aborted and does not rethrow when the stream aborts", async () => { + const abortError = Object.assign(new Error("aborted"), { name: "AbortError" }); + vi.mocked(streamText).mockReturnValue({ + toUIMessageStream: vi.fn(() => + (async function* () { + yield { type: "text-start", id: "t1" }; + throw abortError; + })(), + ), + finishReason: Promise.reject(abortError), + response: Promise.reject(abortError), + } as never); + const { stream } = makeWritable(); + + const result = await runAgentStep({ ...baseInput, writable: stream } as never); + + expect(result.aborted).toBe(true); + expect(result.finishReason).toBe("stop"); + expect(result.responseMessages).toEqual([]); + }); +}); diff --git a/app/lib/workflows/__tests__/runAgentWorkflow.test.ts b/app/lib/workflows/__tests__/runAgentWorkflow.test.ts index 14d1b2f8e..bf5677eb3 100644 --- a/app/lib/workflows/__tests__/runAgentWorkflow.test.ts +++ b/app/lib/workflows/__tests__/runAgentWorkflow.test.ts @@ -27,6 +27,9 @@ vi.mock("@/app/lib/workflows/generateAssistantMessageId", () => ({ vi.mock("@/app/lib/workflows/convertMessagesStep", () => ({ convertMessagesStep: vi.fn(() => Promise.resolve([])), })); +vi.mock("@/app/lib/workflows/persistAssistantMessageStep", () => ({ + persistAssistantMessageStep: vi.fn(), +})); vi.mock("@/app/lib/workflows/sendStreamStart", () => ({ sendStreamStart: vi.fn() })); vi.mock("@/app/lib/workflows/sendStreamFinish", () => ({ sendStreamFinish: vi.fn() })); vi.mock("@/lib/credits/handleChatCredits", () => ({ @@ -157,7 +160,9 @@ describe("runAgentWorkflow", () => { expect(closeChatStream).toHaveBeenCalledWith(writableStub); }); - it("forwards chatId to runAgentStep so it can persist the assistant message per step", async () => { + // Persistence moved to the workflow body, so the step has no business + // knowing the chat id. Guards against re-coupling them. + it("does NOT pass chatId to runAgentStep — persistence is the workflow body's job", async () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, @@ -167,7 +172,9 @@ describe("runAgentWorkflow", () => { await runAgentWorkflow(baseInput); - expect(runAgentStep).toHaveBeenCalledWith(expect.objectContaining({ chatId: "chat-1" })); + expect(runAgentStep).toHaveBeenCalledWith( + expect.not.objectContaining({ chatId: expect.anything() }), + ); }); it("generates a fresh assistantMessageId via the step and forwards it to runAgentStep", async () => { diff --git a/app/lib/workflows/__tests__/runAgentWorkflowLoop.test.ts b/app/lib/workflows/__tests__/runAgentWorkflowLoop.test.ts index 009851b4f..d85c963cd 100644 --- a/app/lib/workflows/__tests__/runAgentWorkflowLoop.test.ts +++ b/app/lib/workflows/__tests__/runAgentWorkflowLoop.test.ts @@ -5,10 +5,14 @@ import { convertMessagesStep } from "@/app/lib/workflows/convertMessagesStep"; import { sendStreamStart } from "@/app/lib/workflows/sendStreamStart"; import { sendStreamFinish } from "@/app/lib/workflows/sendStreamFinish"; import { generateAssistantMessageId } from "@/app/lib/workflows/generateAssistantMessageId"; +import { persistAssistantMessageStep } from "@/app/lib/workflows/persistAssistantMessageStep"; import { CHAT_AGENT_MAX_ITERATIONS } from "@/lib/chat/const"; vi.mock("@/app/lib/workflows/runAgentStep", () => ({ runAgentStep: vi.fn() })); vi.mock("@/app/lib/workflows/convertMessagesStep", () => ({ convertMessagesStep: vi.fn() })); +vi.mock("@/app/lib/workflows/persistAssistantMessageStep", () => ({ + persistAssistantMessageStep: vi.fn(), +})); vi.mock("@/app/lib/workflows/sendStreamStart", () => ({ sendStreamStart: vi.fn() })); vi.mock("@/app/lib/workflows/sendStreamFinish", () => ({ sendStreamFinish: vi.fn() })); vi.mock("@/app/lib/workflows/generateAssistantMessageId", () => ({ @@ -167,6 +171,41 @@ describe("runAgentWorkflow — per-iteration agent loop", () => { expect(runAgentStep).toHaveBeenCalledTimes(1); }); + // Persistence lives in the workflow body, mirroring upstream open-agents + // (`persistAssistantMessage(options.chatId, pendingAssistantResponse)`). + // Persisting per iteration keeps a long turn's transcript live rather than + // landing only at the end. + it("persists the accumulated assistant message after each iteration", async () => { + queueSteps([ + { + finishReason: "tool-calls", + aborted: false, + responseMessage: assistantMessage("a1"), + responseMessages: [{ role: "assistant", content: "call-1" }], + }, + { + finishReason: "stop", + aborted: false, + responseMessage: assistantMessage("a2"), + responseMessages: [], + }, + ]); + + await runAgentWorkflow(baseInput); + + expect(persistAssistantMessageStep).toHaveBeenCalledTimes(2); + expect(persistAssistantMessageStep).toHaveBeenNthCalledWith( + 1, + "chat-1", + assistantMessage("a1"), + ); + expect(persistAssistantMessageStep).toHaveBeenNthCalledWith( + 2, + "chat-1", + assistantMessage("a2"), + ); + }); + it("bounds a runaway tool-call loop at CHAT_AGENT_MAX_ITERATIONS", async () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "tool-calls", diff --git a/app/lib/workflows/persistAssistantMessageStep.ts b/app/lib/workflows/persistAssistantMessageStep.ts new file mode 100644 index 000000000..8fcf249dc --- /dev/null +++ b/app/lib/workflows/persistAssistantMessageStep.ts @@ -0,0 +1,26 @@ +import type { UIMessage } from "ai"; +import { persistAssistantMessage } from "@/lib/chat/persistAssistantMessage"; + +/** + * Persist the in-progress assistant message from the workflow body. + * + * Called after every agent iteration so a long turn's transcript stays live + * in `chat_messages` rather than landing only when the turn ends. Mirrors + * upstream open-agents, which persists `pendingAssistantResponse` from the + * workflow body rather than from inside the agent step. + * + * Keeping this out of `runAgentStep` is what lets the step write parts + * straight to the writable instead of wrapping them in a + * `createUIMessageStream` purely to get a persist callback — a wrapper that + * silently dropped every tool call from the transcript (chat#1918). + * + * `persistAssistantMessage` swallows its own errors, so this never throws. + */ +export async function persistAssistantMessageStep( + chatId: string, + message: UIMessage, +): Promise { + "use step"; + + await persistAssistantMessage(chatId, message); +} diff --git a/app/lib/workflows/runAgentStep.ts b/app/lib/workflows/runAgentStep.ts index b482b6bcd..7d53f8d28 100644 --- a/app/lib/workflows/runAgentStep.ts +++ b/app/lib/workflows/runAgentStep.ts @@ -1,10 +1,4 @@ -import { - streamText, - createUIMessageStream, - type ModelMessage, - type UIMessage, - type UIMessageChunk, -} from "ai"; +import { streamText, type ModelMessage, type UIMessage, type UIMessageChunk } from "ai"; import { gateway } from "@ai-sdk/gateway"; import { agentCustomInstructions } from "@/lib/chat/agentCustomInstructions"; import { buildAgentSystemPrompt } from "@/lib/chat/buildAgentSystemPrompt"; @@ -15,11 +9,11 @@ import type { AgentMessageMetadata } from "@/lib/agent/messageMetadata/AgentMess import { addCacheControlToTools } from "@/lib/agent/contextManagement/addCacheControlToTools"; import { addCacheControlToMessages } from "@/lib/agent/contextManagement/addCacheControlToMessages"; import { wrapToolsWithAbort } from "@/lib/agent/contextManagement/wrapToolsWithAbort"; -import { persistAssistantMessage } from "@/lib/chat/persistAssistantMessage"; import { pollWorkflowCancellation } from "@/lib/chat/pollWorkflowCancellation"; -import { finalizeAbortedAssistantMessage } from "@/lib/chat/finalizeAbortedAssistantMessage"; +import { closeOpenToolCalls } from "@/lib/chat/closeOpenToolCalls"; +import { isAbortError } from "@/lib/chat/isAbortError"; +import { isRunCancelled } from "@/lib/chat/isRunCancelled"; import { getWorkflowMetadata } from "workflow"; -import { pipeWorkflowStreamWithStopDetection } from "@/lib/chat/pipeWorkflowStreamWithStopDetection"; export type RunAgentStepInput = { /** @@ -37,8 +31,6 @@ export type RunAgentStepInput = { originalMessages: UIMessage[]; modelId: string; writable: WritableStream; - /** Target chat for persisting the assistant message as it streams. */ - chatId: string; /** * The JSON-serializable agent context that survives the durable * workflow input. `runAgentStep` widens it into a full `AgentContext` @@ -80,10 +72,11 @@ export type RunAgentStepInput = { export type RunAgentStepResult = { finishReason: string; /** - * The assembled assistant message captured from the stream's `onFinish`. - * `undefined` if the stream finished without emitting one. Per-step - * persistence happens inside this function; this is returned so - * `runAgentWorkflow` can charge credits from `responseMessage.metadata`. + * The assembled assistant message captured from the stream's `onFinish`, + * cumulative across iterations via `originalMessages`. `undefined` if the + * stream finished without emitting one. Returned rather than persisted + * here — `runAgentWorkflow` persists it and charges credits from its + * metadata. */ responseMessage: UIMessage | undefined; /** @@ -188,91 +181,86 @@ export async function runAgentStep(input: RunAgentStepInput): Promise({ - generateId: () => input.assistantMessageId, - // Persistence mode: this outer stream is what assembles the message - // handed to onStepFinish/onFinish, so it needs the in-progress assistant - // message too — not just the inner `toUIMessageStream`. Without it every - // iteration rebuilds the message from its own chunks alone, and the final - // text-only persist wipes every tool call made earlier in the turn. - originalMessages: input.originalMessages, - onStepFinish: ({ responseMessage: stepMessage }) => { - responseMessage = stepMessage; - return persistAssistantMessage(input.chatId, stepMessage); - }, - onFinish: ({ responseMessage: finalMessage }) => { - responseMessage = finalMessage; - return persistAssistantMessage(input.chatId, finalMessage); - }, - onError: error => { - if (cancelController.signal.aborted) return ""; - return error instanceof Error ? error.message : String(error); - }, - execute: ({ writer }) => { - writer.merge( - result.toUIMessageStream({ - messageMetadata, - generateMessageId: () => input.assistantMessageId, - // Continue building the in-progress assistant message rather than - // starting a new one, so `responseMessage` stays cumulative across - // iterations and each persist overwrites a single row. - originalMessages: input.originalMessages, - // The turn's `start`/`finish` chunks are emitted ONCE by the - // workflow body (`sendStreamStart` / `sendStreamFinish`). Without - // this the client would see one per iteration and render N - // assistant messages instead of one. - sendStart: false, - sendFinish: false, - }), - ); - }, - }); - - // Pipe the stream to the workflow writable and detect user-stop vs natural - // finish (see pipeWorkflowStreamWithStopDetection for the why). - const userAborted = await pipeWorkflowStreamWithStopDetection({ - uiStream, - writable: input.writable, - cancelController, - workflowRunId, - poller, - }); - - // Short-circuit on user-stop — `result.finishReason` rejects when streamText aborts. let finishReason: string; let responseMessages: ModelMessage[] = []; - if (userAborted) { - finishReason = "stop"; - // Prevent the late-rejecting promises from becoming unhandled rejections. - void Promise.resolve(result.finishReason).catch(() => {}); - void Promise.resolve(result.response).catch(() => {}); - } else { - // `result.response.messages` is the assistant message for this call plus - // any tool-result message — exactly what the next iteration needs - // appended. (`result.responseMessages` is ai@7; this repo is on 6.0.190.) + + try { + for await (const part of result.toUIMessageStream({ + messageMetadata, + generateMessageId: () => input.assistantMessageId, + // Continue building the in-progress assistant message rather than + // starting a new one, so `responseMessage` stays cumulative across + // iterations and each persist overwrites a single row. + originalMessages: input.originalMessages, + // The turn's `start`/`finish` chunks are emitted ONCE by the workflow + // body (`sendStreamStart` / `sendStreamFinish`). Without this the client + // would see one per iteration and render N assistant messages. + sendStart: false, + sendFinish: false, + onFinish: ({ responseMessage: finalMessage }) => { + responseMessage = finalMessage; + }, + })) { + // A writer lock taken inside a step applies only within that step, so + // sequential steps can share one stream. Release per part so the step's + // request can terminate. + const writer = input.writable.getWriter(); + try { + await writer.write(part); + } finally { + writer.releaseLock(); + } + } + + // `response.messages` is the assistant message for this call plus any + // tool-result message — exactly what the next iteration needs appended. + // (`result.responseMessages` is ai@7; this repo is on 6.0.190.) const [reason, response] = await Promise.all([result.finishReason, result.response]); finishReason = reason; responseMessages = response.messages; - } + } catch (error) { + // Three ways a user-stop surfaces here: the stream throws AbortError; the + // poller already flipped our signal; or `run.cancel()` closed the workflow + // writable underneath us and the write threw something unrelated before + // the poller noticed. Confirm the last case against the run itself so a + // genuine failure still propagates. + if ( + !isAbortError(error) && + !cancelController.signal.aborted && + !(await isRunCancelled(workflowRunId)) + ) + throw error; + + // User-stop. `result.finishReason` / `result.response` reject once + // streamText aborts — swallow them so they don't surface as unhandled. + void Promise.resolve(result.finishReason).catch(() => {}); + void Promise.resolve(result.response).catch(() => {}); + finishReason = "stop"; + // Close any tool-call parts left without a terminal result, otherwise a + // reload renders them spinning forever. The workflow body persists it. + if (responseMessage) responseMessage = closeOpenToolCalls(responseMessage); - // On user-stop, close any tool-call parts the step boundary left open and - // re-persist (see finalizeAbortedAssistantMessage for the why). - if (userAborted && responseMessage) { - responseMessage = await finalizeAbortedAssistantMessage(input.chatId, responseMessage); + console.log("[runAgentStep] aborted", { hasResponseMessage: !!responseMessage }); + return { finishReason, responseMessage, responseMessages: [], aborted: true }; + } finally { + poller.stop(); + cancelController.abort(); + await poller.done.catch(() => {}); } console.log("[runAgentStep] finish", { finishReason, hasResponseMessage: !!responseMessage, responseMessageCount: responseMessages.length, - aborted: userAborted, }); - return { finishReason, responseMessage, responseMessages, aborted: userAborted }; + return { finishReason, responseMessage, responseMessages, aborted: false }; } diff --git a/app/lib/workflows/runAgentWorkflow.ts b/app/lib/workflows/runAgentWorkflow.ts index 65328e176..4cbab0700 100644 --- a/app/lib/workflows/runAgentWorkflow.ts +++ b/app/lib/workflows/runAgentWorkflow.ts @@ -6,6 +6,7 @@ import { runAgentStep } from "@/app/lib/workflows/runAgentStep"; import { convertMessagesStep } from "@/app/lib/workflows/convertMessagesStep"; import { sendStreamStart } from "@/app/lib/workflows/sendStreamStart"; import { sendStreamFinish } from "@/app/lib/workflows/sendStreamFinish"; +import { persistAssistantMessageStep } from "@/app/lib/workflows/persistAssistantMessageStep"; import { CHAT_AGENT_MAX_ITERATIONS } from "@/lib/chat/const"; import { clearChatActiveStream } from "@/lib/chat/clearChatActiveStream"; import { deleteEphemeralKeyStep } from "@/app/lib/workflows/deleteEphemeralKeyStep"; @@ -141,7 +142,6 @@ export async function runAgentWorkflow(input: RunAgentWorkflowInput): Promise ({ closeOpenToolCalls: vi.fn() })); -vi.mock("@/lib/chat/persistAssistantMessage", () => ({ persistAssistantMessage: vi.fn() })); - -const CHAT_ID = "chat-1"; -const original = { id: "m1", role: "assistant", parts: [] } as unknown as UIMessage; - -describe("finalizeAbortedAssistantMessage", () => { - beforeEach(() => vi.clearAllMocks()); - - it("persists and returns the closed message when open tool-calls were closed", async () => { - const closed = { - id: "m1", - role: "assistant", - parts: [{ type: "step-start" }], - } as unknown as UIMessage; - vi.mocked(closeOpenToolCalls).mockReturnValue(closed); - - const result = await finalizeAbortedAssistantMessage(CHAT_ID, original); - - expect(closeOpenToolCalls).toHaveBeenCalledWith(original); - expect(persistAssistantMessage).toHaveBeenCalledWith(CHAT_ID, closed); - expect(result).toBe(closed); - }); - - it("is a no-op (no persist, returns original) when nothing was open", async () => { - vi.mocked(closeOpenToolCalls).mockReturnValue(original); - - const result = await finalizeAbortedAssistantMessage(CHAT_ID, original); - - expect(persistAssistantMessage).not.toHaveBeenCalled(); - expect(result).toBe(original); - }); -}); diff --git a/lib/chat/__tests__/pipeWorkflowStreamWithStopDetection.test.ts b/lib/chat/__tests__/pipeWorkflowStreamWithStopDetection.test.ts deleted file mode 100644 index 21f007868..000000000 --- a/lib/chat/__tests__/pipeWorkflowStreamWithStopDetection.test.ts +++ /dev/null @@ -1,66 +0,0 @@ -import { describe, it, expect, vi, beforeEach } from "vitest"; -import { pipeWorkflowStreamWithStopDetection } from "@/lib/chat/pipeWorkflowStreamWithStopDetection"; -import { getRun } from "workflow/api"; - -vi.mock("workflow/api", () => ({ getRun: vi.fn() })); - -const mockStatus = (value: string | Error) => - vi.mocked(getRun).mockReturnValue({ - get status() { - return value instanceof Error ? Promise.reject(value) : Promise.resolve(value); - }, - } as unknown as ReturnType); - -const makeParams = (pipeTo: () => Promise) => { - const poller = { stop: vi.fn(), done: Promise.resolve() }; - const cancelController = new AbortController(); - return { - poller, - cancelController, - params: { - uiStream: { pipeTo } as unknown as ReadableStream, - writable: {} as unknown as WritableStream, - cancelController, - workflowRunId: "run-1", - poller, - } as Parameters[0], - }; -}; - -describe("pipeWorkflowStreamWithStopDetection", () => { - beforeEach(() => vi.clearAllMocks()); - - it("returns false on natural completion (pipe resolves, status not cancelled)", async () => { - mockStatus("completed"); - const { params, poller, cancelController } = makeParams(() => Promise.resolve()); - await expect(pipeWorkflowStreamWithStopDetection(params)).resolves.toBe(false); - expect(poller.stop).toHaveBeenCalled(); - expect(cancelController.signal.aborted).toBe(true); // finally aborts unconditionally - }); - - it("returns true when pipe resolves cleanly but the run was cancelled", async () => { - mockStatus("cancelled"); - const { params } = makeParams(() => Promise.resolve()); - await expect(pipeWorkflowStreamWithStopDetection(params)).resolves.toBe(true); - }); - - it("treats a transient status-read error after a clean pipe as success (false)", async () => { - mockStatus(new Error("status blip")); - const { params } = makeParams(() => Promise.resolve()); - await expect(pipeWorkflowStreamWithStopDetection(params)).resolves.toBe(false); - }); - - it("returns true when pipe rejects and the controller was aborted (user-stop)", async () => { - const { params, cancelController } = makeParams(() => { - cancelController.abort(); - return Promise.reject(new Error("aborted")); - }); - await expect(pipeWorkflowStreamWithStopDetection(params)).resolves.toBe(true); - }); - - it("rethrows a genuine pipe error when the controller was NOT aborted", async () => { - const { params, poller } = makeParams(() => Promise.reject(new Error("disk full"))); - await expect(pipeWorkflowStreamWithStopDetection(params)).rejects.toThrow("disk full"); - expect(poller.stop).toHaveBeenCalled(); // finally still runs - }); -}); diff --git a/lib/chat/finalizeAbortedAssistantMessage.ts b/lib/chat/finalizeAbortedAssistantMessage.ts deleted file mode 100644 index 92a058779..000000000 --- a/lib/chat/finalizeAbortedAssistantMessage.ts +++ /dev/null @@ -1,24 +0,0 @@ -import type { UIMessage } from "ai"; -import { closeOpenToolCalls } from "@/lib/chat/closeOpenToolCalls"; -import { persistAssistantMessage } from "@/lib/chat/persistAssistantMessage"; - -/** - * On user-stop, close any tool-call parts the step boundary persisted without a - * terminal result and re-persist. Without this a reload renders those parts as - * spinning forever — the AI SDK only transitions a tool part when it sees a - * terminal `output-*` chunk. A no-op (returns the original) when nothing is open. - * - * @param chatId - Chat whose assistant message is being finalized. - * @param responseMessage - The assistant message persisted at the step boundary. - * @returns The closed-and-persisted message, or the original when unchanged. - */ -export async function finalizeAbortedAssistantMessage( - chatId: string, - responseMessage: UIMessage, -): Promise { - const closed = closeOpenToolCalls(responseMessage); - if (closed !== responseMessage) { - await persistAssistantMessage(chatId, closed); - } - return closed; -} diff --git a/lib/chat/isAbortError.ts b/lib/chat/isAbortError.ts new file mode 100644 index 000000000..a7faf691a --- /dev/null +++ b/lib/chat/isAbortError.ts @@ -0,0 +1,11 @@ +/** + * Is this the error a stream throws when its abort signal fires? + * + * Used by `runAgentStep` to tell a user-stop apart from a genuine failure: + * the former is a normal end to the turn, the latter must propagate so the + * workflow records it. Port of upstream open-agents' `isAbortError` in + * `apps/web/app/workflows/chat.ts`. + */ +export function isAbortError(error: unknown): boolean { + return error instanceof Error && error.name === "AbortError"; +} diff --git a/lib/chat/isRunCancelled.ts b/lib/chat/isRunCancelled.ts new file mode 100644 index 000000000..5ab6e2b73 --- /dev/null +++ b/lib/chat/isRunCancelled.ts @@ -0,0 +1,21 @@ +import { getRun } from "workflow/api"; + +/** + * Has this workflow run been cancelled? + * + * `run.cancel()` (our `POST /api/chat/[chatId]/stop` path) closes the run's + * writable, so a step streaming into it can fail with a plain stream error + * before its cancellation poller notices. `runAgentStep` asks this before + * rethrowing, so a stop is reported as `aborted` rather than as a crash that + * fails the whole workflow. + * + * A status read that itself fails is treated as "not cancelled" — better to + * surface the original error than to swallow it on a transient blip. + */ +export async function isRunCancelled(workflowRunId: string): Promise { + try { + return (await getRun(workflowRunId).status) === "cancelled"; + } catch { + return false; + } +} diff --git a/lib/chat/pipeWorkflowStreamWithStopDetection.ts b/lib/chat/pipeWorkflowStreamWithStopDetection.ts deleted file mode 100644 index 87aae1388..000000000 --- a/lib/chat/pipeWorkflowStreamWithStopDetection.ts +++ /dev/null @@ -1,64 +0,0 @@ -import type { UIMessageChunk } from "ai"; -import { getRun } from "workflow/api"; -import type { CancellationPoller } from "@/lib/chat/pollWorkflowCancellation"; - -export interface PipeWorkflowStreamParams { - uiStream: ReadableStream; - writable: WritableStream; - cancelController: AbortController; - workflowRunId: string; - poller: CancellationPoller; -} - -/** - * Pipes the agent's UIMessage stream to the workflow writable while detecting a - * user-stop, returning whether the turn was aborted by the user. - * - * Distinguishing user-stop from natural completion: only the cancellation poller - * aborts the controller before our own `finally` runs, so a `pipeTo` rejection - * with the signal already aborted is the user-stop path. `pipeTo` can also resolve - * cleanly on cancel (the workflow runtime closes the destination writable on - * `run.cancel()`, which surfaces as a natural finish), so we additionally confirm - * via the run's status. We capture this here rather than reading - * `cancelController.signal.aborted` afterward, because the `finally` aborts the - * controller unconditionally to stop the poller — which would otherwise make every - * natural completion look like a user-stop. - * - * @returns true when the user stopped the run, false on natural completion. - * @throws the original error on a genuine (non-abort) pipe failure. - */ -export async function pipeWorkflowStreamWithStopDetection({ - uiStream, - writable, - cancelController, - workflowRunId, - poller, -}: PipeWorkflowStreamParams): Promise { - let userAborted = false; - try { - await uiStream.pipeTo(writable, { - preventClose: true, - preventAbort: true, - signal: cancelController.signal, - }); - try { - const status = await getRun(workflowRunId).status; - if (status === "cancelled") { - userAborted = true; - } - } catch { - /* transient status read — treat as success */ - } - } catch (err) { - if (cancelController.signal.aborted) { - userAborted = true; - } else { - throw err; - } - } finally { - poller.stop(); - cancelController.abort(); - await poller.done.catch(() => {}); - } - return userAborted; -}