diff --git a/app/lib/workflows/__tests__/runAgentStep.test.ts b/app/lib/workflows/__tests__/runAgentStep.test.ts index a9faf935e..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,37 @@ vi.mock("workflow/api", () => ({ })), })); -// Captures the options runAgentStep passes to createUIMessageStream so -// tests can drive its onStepFinish / onFinish callbacks directly. -type CreateOpts = { - generateId?: () => string; - onStepFinish?: (e: { responseMessage: unknown }) => unknown; +type StreamOpts = { + messageMetadata?: unknown; + generateMessageId?: unknown; + originalMessages?: 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: [] }), }; } @@ -92,7 +87,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, @@ -100,7 +98,6 @@ const baseInput = { }, ], modelId: "anthropic/claude-haiku-4.5", - chatId: "chat-1", agentContext: { sandbox: { state: { type: "vercel" }, workingDirectory: "/sandbox/mono" }, }, @@ -108,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[] = []; @@ -201,69 +182,88 @@ 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" } } }); }); - it("the wired callback returns undefined for non-finish-step parts", async () => { - const captured: unknown[] = []; - vi.mocked(streamText).mockReturnValue(makeStreamResult({ metadataCalls: captured }) as never); + // 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 cb = captured[0] as (args: { part: { type: string } }) => unknown; - expect(cb({ part: { type: "text-delta" } })).toBeUndefined(); - expect(cb({ part: { type: "start" } })).toBeUndefined(); + const args = vi.mocked(streamText).mock.calls[0]?.[0] as { stopWhen?: unknown }; + expect(args.stopWhen).toBeUndefined(); }); - it("persists the assistant message on each step via onStepFinish", async () => { - vi.mocked(streamText).mockReturnValue(makeStreamResult() as never); + 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"), + response: Promise.resolve({ messages: [] }), + } 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(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"), + response: Promise.resolve({ messages: produced }), + } as never); + const { stream } = makeWritable(); + + const result = await runAgentStep({ ...baseInput, writable: stream } as never); - expect(persistAssistantMessage).toHaveBeenCalledWith("chat-1", msg); + expect(result.responseMessages).toEqual(produced); }); - it("persists the final assistant message via onFinish", async () => { - vi.mocked(streamText).mockReturnValue(makeStreamResult() as never); + it("the wired callback returns undefined for non-finish-step parts", async () => { + const captured: unknown[] = []; + vi.mocked(streamText).mockReturnValue(makeStreamResult({ metadataCalls: captured }) 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); + const cb = captured[0] as (args: { part: { type: string } }) => unknown; + expect(cb({ part: { type: "text-delta" } })).toBeUndefined(); + expect(cb({ part: { type: "start" } })).toBeUndefined(); }); it("forwards assistantMessageId into toUIMessageStream's generateMessageId (stable row id)", async () => { @@ -283,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(); @@ -313,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); @@ -366,6 +343,7 @@ describe("runAgentStep", () => { })(), ), finishReason: Promise.resolve("length"), + response: Promise.resolve({ messages: [] }), } as never); const { stream } = makeWritable(); @@ -375,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 be1adf5d0..bf5677eb3 100644 --- a/app/lib/workflows/__tests__/runAgentWorkflow.test.ts +++ b/app/lib/workflows/__tests__/runAgentWorkflow.test.ts @@ -23,6 +23,15 @@ 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/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", () => ({ handleChatCredits: vi.fn(), })); @@ -76,6 +85,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -98,6 +108,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -117,6 +128,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -129,6 +141,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -147,22 +160,28 @@ 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, + responseMessages: [], responseMessage: undefined, }); 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 () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -178,6 +197,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -218,6 +238,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: responseMessage as never, }); @@ -247,6 +268,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: responseMessage as never, }); @@ -266,6 +288,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -292,6 +315,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: responseMessageWithMetadata, }); @@ -308,6 +332,7 @@ describe("runAgentWorkflow", () => { responseMessage: responseMessageWithMetadata, finishReason: "stop", aborted: false, + responseMessages: [], sessionId: "session-1", sessionTitle: "test session", repoOwner: "recoupable", @@ -321,6 +346,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: responseMessageWithMetadata, }); @@ -334,6 +360,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: responseMessageWithMetadata, }); @@ -350,6 +377,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: false, + responseMessages: [], responseMessage: undefined, }); @@ -378,6 +406,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: true, + responseMessages: [], responseMessage: abortedResponseMessage as never, }); @@ -398,6 +427,7 @@ describe("runAgentWorkflow", () => { vi.mocked(runAgentStep).mockResolvedValue({ finishReason: "stop", aborted: true, + responseMessages: [], responseMessage: abortedResponseMessage as never, }); @@ -410,6 +440,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..d85c963cd --- /dev/null +++ b/app/lib/workflows/__tests__/runAgentWorkflowLoop.test.ts @@ -0,0 +1,221 @@ +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 { 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", () => ({ + 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); + }); + + // 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", + 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/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 730b4e40e..7d53f8d28 100644 --- a/app/lib/workflows/runAgentStep.ts +++ b/app/lib/workflows/runAgentStep.ts @@ -1,32 +1,36 @@ -import { - streamText, - convertToModelMessages, - createUIMessageStream, - 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"; -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"; -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 = { - 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. */ - chatId: string; /** * The JSON-serializable agent context that survives the durable * workflow input. `runAgentStep` widens it into a full `AgentContext` @@ -68,38 +72,49 @@ 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; + /** + * 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 +123,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 - // step — a stopped or crashed turn keeps the partial reply rather than - // dropping it. The stable assistantMessageId makes each upsert overwrite - // the same row. The final message is also captured so runAgentWorkflow can - // charge credits from its metadata. + // Drive the stream directly and write each part to the shared writable, + // mirroring upstream open-agents. There is deliberately no + // `createUIMessageStream` wrapper: its only draw was `onStepFinish` for + // in-step persistence, and with one model call per step that fires once + // anyway. The wrapper also has to be put in "persistence mode" separately + // from the inner stream, and missing that silently dropped every tool call + // from the transcript (chat#1918). Persistence now lives in the workflow + // body via `persistAssistantMessageStep`. let responseMessage: UIMessage | undefined; - const uiStream = createUIMessageStream({ - generateId: () => input.assistantMessageId, - 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, - }), - ); - }, - }); + let finishReason: string; + let responseMessages: ModelMessage[] = []; - // 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, - }); + 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(); + } + } - // Short-circuit on user-stop — `result.finishReason` rejects when streamText aborts. - let finishReason: string; - if (userAborted) { - finishReason = "stop"; - // Prevent the late-rejecting promise from becoming an unhandled rejection. + // `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(() => {}); - } else { - finishReason = await result.finishReason; - } + 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, - aborted: userAborted, + responseMessageCount: responseMessages.length, }); - return { finishReason, responseMessage, aborted: userAborted }; + return { finishReason, responseMessage, responseMessages, aborted: false }; } diff --git a/app/lib/workflows/runAgentWorkflow.ts b/app/lib/workflows/runAgentWorkflow.ts index b73353727..4cbab0700 100644 --- a/app/lib/workflows/runAgentWorkflow.ts +++ b/app/lib/workflows/runAgentWorkflow.ts @@ -3,6 +3,11 @@ import type { LanguageModelUsage, UIMessage, UIMessageChunk } from "ai"; import { closeChatStream } from "@/app/lib/workflows/closeChatStream"; import { generateAssistantMessageId } from "@/app/lib/workflows/generateAssistantMessageId"; 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"; import { handleChatCredits } from "@/lib/credits/handleChatCredits"; @@ -60,11 +65,16 @@ export type RunAgentWorkflowInput = { * client; this function writes UIMessage chunks into the workflow's writable * via `runAgentStep`. * - * Currently runs a SINGLE `runAgentStep` turn. Tool-call iteration (up to - * MAX_TOOL_STEPS) happens INSIDE `streamText` via `stopWhen` — so the - * single workflow turn covers the full "user → assistant → tool → tool - * result → assistant" cycle without our outer loop having to thread - * messages between iterations. + * Runs the agent loop in THIS body, one `runAgentStep` per LLM call, up to + * `CHAT_AGENT_MAX_ITERATIONS`. Each iteration is journaled, so a killed or + * retried run resumes at the last completed call instead of re-executing the + * whole turn from minute zero. + * + * This replaced a single step that wrapped the entire loop via + * `stopWhen: stepCountIs(111)`. That step ran 11-25 minutes, exceeded + * Vercel's 800 s function ceiling, and was killed and retried 4 times — five + * complete agent runs, five emails to the customer, then a failed workflow. + * See chat#1918. * * WDK constraints honored: * - All I/O (streamText, sandbox.exec, fetches) lives in `"use step"` functions. @@ -97,27 +107,84 @@ export async function runAgentWorkflow(input: RunAgentWorkflowInput): 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, + accountId: input.accountId, + artistId: input.artistId, + interactive: input.interactive, + agentContext: input.agentContext, + writable, + assistantMessageId, + }); + + if (result.responseMessage) { + pendingAssistantResponse = result.responseMessage; + // Persist per iteration so a long turn's transcript stays live rather + // than landing only at the end. The stable assistantMessageId makes + // each write overwrite the same row. + await persistAssistantMessageStep(input.chatId, pendingAssistantResponse); + } + 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 +200,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/__tests__/finalizeAbortedAssistantMessage.test.ts b/lib/chat/__tests__/finalizeAbortedAssistantMessage.test.ts deleted file mode 100644 index f3670d2c0..000000000 --- a/lib/chat/__tests__/finalizeAbortedAssistantMessage.test.ts +++ /dev/null @@ -1,39 +0,0 @@ -import { describe, it, expect, vi, beforeEach } from "vitest"; -import type { UIMessage } from "ai"; -import { finalizeAbortedAssistantMessage } from "@/lib/chat/finalizeAbortedAssistantMessage"; -import { closeOpenToolCalls } from "@/lib/chat/closeOpenToolCalls"; -import { persistAssistantMessage } from "@/lib/chat/persistAssistantMessage"; - -vi.mock("@/lib/chat/closeOpenToolCalls", () => ({ 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/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. --- 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; -}