feat(ai): keep partial output on errors (#60477)

packages/ai/src/index.ts

@@ -316,5 +316,6 @@ export const tool = <
 /** Default client over `RequestExecutor.fetchLayer` for scripts; the runtime builds its layer on first use. */
 export const ai = make()
 
+export { AIError } from "./schema/index.js"
 export * as AI from "./index.js"
 export type { ExecutableTool, Tool, ToolOptions, Tools } from "./tool.js"

packages/ai/src/llm.ts

@@ -17,6 +17,7 @@ import {
   type ContentPart,
   type LanguageModelProviderOptions,
   type ToolEntryInput,
+  withPartial,
 } from "./schema/index.js"
 import { tool, toDefinitions, type ExecutableTool, type Tool, type ToolSchema } from "./tool.js"
 import { ToolRuntime } from "./tool-runtime.js"
@@ -74,18 +75,19 @@ export function loop(input: LoopInput, options?: StreamOptions) {
   const { tools, until, ...rest } = input
   const step = (request: LLMRequest, steps: ReadonlyArray<LLMResponse>): Stream.Stream<LLMEvent, AIError, Service> =>
     Stream.suspend(() => {
-      const events: LLMEvent[] = []
+      let state = LLMResponse.empty()
       return LLMClient.stream(request, options).pipe(
         Stream.map((event) =>
           event.type === "step-start" || event.type === "step-finish" ? { ...event, index: steps.length } : event,
         ),
-        Stream.tap((event) => Effect.sync(() => events.push(event))),
+        Stream.tap((event) => Effect.sync(() => (state = LLMResponse.reduce(state, event)))),
+        Stream.mapError((error) => withPartial(error, LLMResponse.partial(state))),
         Stream.concat(
           Stream.unwrap(
             Effect.gen(function* () {
-              const response = LLMResponse.fromEvents(events)
+              const response = LLMResponse.complete(state)
               const calls = response?.toolCalls.filter((call) => call.providerExecuted !== true) ?? []
-              if (!response || calls.length === 0 || events.some(LLMEvent.is.providerError)) return Stream.empty
+              if (!response || calls.length === 0 || state.events.some(LLMEvent.is.providerError)) return Stream.empty
               const done = [...steps, response]
               if (until?.({ steps: done })) return Stream.empty
               const results = yield* Effect.forEach(calls, (call) => ToolRuntime.dispatch(tools, call), {
@@ -193,21 +195,22 @@ const runGenerateObject = Effect.fn("LLM.generateObject")(function* (options: Ge
   })
   const response = yield* LLMClient.generate(generateRequest)
   const call = response.toolCalls.find((event) => event.name === GENERATE_OBJECT_TOOL_NAME)
-  if (!call)
-    return yield* new AIError({
-      reason: new InvalidProviderOutputError({
-        message: `generateObject: model did not call the forced \`${GENERATE_OBJECT_TOOL_NAME}\` tool`,
-      }),
-    })
+  const invalidOutput = (reason: InvalidProviderOutputError) => new AIError({ reason, partial: response })
+  if (!call) {
+    const message =
+      response.finishReason.normalized === "length"
+        ? "generateObject: output hit the token limit before the forced tool call completed"
+        : `generateObject: model did not call the forced \`${GENERATE_OBJECT_TOOL_NAME}\` tool`
+    return yield* invalidOutput(new InvalidProviderOutputError({ message }))
+  }
   const object = yield* tool._decode(call.input).pipe(
-    Effect.mapError(
-      (error) =>
-        new AIError({
-          reason: new InvalidProviderOutputError({
-            message: `generateObject: tool input failed schema decode: ${error.message}`,
-            cause: error,
-          }),
+    Effect.mapError((error) =>
+      invalidOutput(
+        new InvalidProviderOutputError({
+          message: `generateObject: tool input failed schema decode: ${error.message}`,
+          cause: error,
         }),
+      ),
     ),
   )
   return new GenerateObjectResponse(object, response)

packages/ai/src/route/client.ts

@@ -30,6 +30,7 @@ import {
   mergeGenerationOptions,
   mergeHttpOptions,
   mergeProviderOptions,
+  withPartial,
 } from "../schema/index.js"
 
 export interface RouteBody<Body> {
@@ -629,12 +630,19 @@ const streamRequestWith = (runtime: TransportRuntime) => (request: LLMRequest, o
     }),
   )
 
-const generateWith = (stream: Interface["stream"]) =>
+export const generateWith = (stream: Interface["stream"]) =>
   Effect.fn("LLM.generate")(function* (request: LLMRequest, options?: StreamOptions) {
-    const state = yield* stream(request, options).pipe(Stream.runFold(LLMResponse.empty, LLMResponse.reduce))
+    let state = LLMResponse.empty()
+    yield* stream(request, options).pipe(
+      Stream.runForEach((event) => Effect.sync(() => (state = LLMResponse.reduce(state, event)))),
+      Effect.mapError((error) => withPartial(error, LLMResponse.partial(state))),
+    )
     const response = LLMResponse.complete(state)
     if (response) return response
-    return yield* incompleteStreamError(`${request.model.provider}/${request.model.route.id}`)
+    return yield* withPartial(
+      incompleteStreamError(`${request.model.provider}/${request.model.route.id}`),
+      LLMResponse.partial(state),
+    )
   })
 
 export function stream(request: LLMRequest, options?: StreamOptions): Stream.Stream<LLMEvent, AIError, Service> {

packages/ai/src/schema/errors.ts

@@ -2,6 +2,7 @@ import { Schema } from "effect"
 import { NonNegativeInt } from "@opencode/schema"
 import { Tool } from "@opencode/schema/tool"
 import { ModelID, ProviderID, RouteID } from "./ids.js"
+import type { LLMResponse } from "./events.js"
 
 export const ProviderFailureClassification = Schema.Literals(["context-overflow", "payload-too-large"])
 export type ProviderFailureClassification = typeof ProviderFailureClassification.Type
@@ -163,12 +164,23 @@ export class AIError extends Schema.TaggedError<AIError>()("AI.Error", {
   reason: AIErrorReason,
 }) {
   override readonly cause = this.reason
+  /** Output that arrived before `generate`, `generateObject` or a `loop` step failed. Not in the schema, so encoding omits it. */
+  readonly partial?: LLMResponse
+
+  constructor(props: Pick<AIError, "reason" | "partial">) {
+    super({ reason: props.reason })
+    this.partial = props.partial
+  }
 
   override get message(): string {
     return this.reason.message
   }
 }
 
+/** Return `error` with `partial` attached, or unchanged when there is no partial output. */
+export const withPartial = (error: AIError, partial: LLMResponse | undefined) =>
+  partial === undefined ? error : new AIError({ reason: error.reason, partial })
+
 /**
  * Failure type for tool execute handlers. Handlers must map their internal
  * errors to this shape; the runtime catches `ToolFailure`s and surfaces them

packages/ai/src/schema/events.ts

@@ -756,6 +756,12 @@ export namespace LLMResponse {
   /** Convenience reducer for callers that already have a collected event list. */
   export const fromEvents = (events: ReadonlyArray<LLMEvent>) => complete(events.reduce(reduce, empty()))
 
+  /** Output that arrived before a failure, or `undefined` if only `step-start` did; an unfinished response reports `error`. */
+  export const partial = (state: State) =>
+    state.events.every(LLMEvent.is.stepStart)
+      ? undefined
+      : complete({ ...state, finishReason: state.finishReason ?? { normalized: "error" } })
+
   /** Concatenate assistant text from a response or collected event list. */
   export const text = (response: Output) => responseText(response.events)
 

packages/ai/src/testing.ts

@@ -8,6 +8,7 @@ export { TestVideo } from "./testing/video.js"
 
 import {
   LLMClient,
+  generateWith,
   type CompactionRequest,
   type CheckpointRequest,
   type EndpointCompactOptions,
@@ -15,7 +16,6 @@ import {
 } from "./route/client.js"
 import {
   LLMEvent,
-  LLMResponse,
   CompactionResponse,
   CompactionCheckpointResponse,
   type FinishReasonDetails,
@@ -131,15 +131,7 @@ const make = (options: LayerOptions) =>
       ...script.controls,
       compact,
       stream,
-      generate: (request) =>
-        stream(request).pipe(
-          Stream.runFold(LLMResponse.empty, LLMResponse.reduce),
-          Effect.flatMap((state) => {
-            const response = LLMResponse.complete(state)
-            if (response) return Effect.succeed(response)
-            return Effect.die("TestLLM response ended without a terminal finish event")
-          }),
-        ),
+      generate: generateWith(stream),
     })
   })
 

packages/ai/test/error-partial.test.ts

@@ -0,0 +1,103 @@
+import { describe, expect, test } from "bun:test"
+import { Effect, Schema, Stream } from "effect"
+import { LanguageModel, LLM, LLMEvent, ProviderInternalError, tool } from "../src/effect.js"
+import { AI, AIError } from "../src/index.js"
+import { OpenAIChat } from "../src/protocols/openai-chat.js"
+import { TestLLM } from "../src/testing.js"
+import { testEffect } from "./lib/effect.js"
+
+const model = LanguageModel.make({ id: "fictional-model", provider: "fixture", route: OpenAIChat.route })
+const it = testEffect(TestLLM.testLayer())
+const overloaded = new AIError({ reason: new ProviderInternalError({ message: "Overloaded" }) })
+
+const textThenFail = TestLLM.failAfter(
+  overloaded,
+  LLMEvent.stepStart({ index: 0 }),
+  LLMEvent.textStart({ id: "text-0" }),
+  LLMEvent.textDelta({ id: "text-0", text: "Partial answer" }),
+  LLMEvent.stepFinish({
+    index: 0,
+    reason: { normalized: "stop" },
+    usage: { inputTokens: 5, nonCachedInputTokens: 5, outputTokens: 2 },
+  }),
+)
+
+describe("AIError.partial", () => {
+  it.effect("generate keeps the output that arrived before a mid-stream failure", () =>
+    Effect.gen(function* () {
+      const client = yield* TestLLM.Test
+      yield* client.push(textThenFail)
+
+      const error = yield* LLM.generate({ model, prompt: "Hello" }).pipe(Effect.flip)
+
+      expect(error.reason).toBe(overloaded.reason)
+      expect(error.partial?.text).toBe("Partial answer")
+      expect(error.partial?.message.content).toEqual([{ type: "text", text: "Partial answer" }])
+      expect(error.partial?.usage?.outputTokens).toBe(2)
+      expect(error.partial?.finishReason).toEqual({ normalized: "error" })
+    }),
+  )
+
+  it.effect("generate keeps the same error when the failure comes before any output", () =>
+    Effect.gen(function* () {
+      const client = yield* TestLLM.Test
+      yield* client.push(TestLLM.failAfter(overloaded, LLMEvent.stepStart({ index: 0 })))
+
+      expect(yield* LLM.generate({ model, prompt: "Hello" }).pipe(Effect.flip)).toBe(overloaded)
+    }),
+  )
+
+  it.effect("generateObject keeps the raw output when the token limit cuts off the forced tool call", () =>
+    Effect.gen(function* () {
+      const client = yield* TestLLM.Test
+      yield* client.push(
+        TestLLM.complete(
+          { reason: { normalized: "length" } },
+          LLMEvent.textStart({ id: "text-0" }),
+          LLMEvent.textDelta({ id: "text-0", text: "Thinking about the answer" }),
+        ),
+      )
+
+      const error = yield* LLM.generateObject({
+        model,
+        prompt: "Answer",
+        schema: Schema.Struct({ answer: Schema.String }),
+      }).pipe(Effect.flip)
+
+      expect(error.message).toContain("token limit")
+      expect(error.partial?.text).toBe("Thinking about the answer")
+      expect(error.partial?.finishReason.normalized).toBe("length")
+    }),
+  )
+
+  it.effect("loop reports only the failed step as partial", () =>
+    Effect.gen(function* () {
+      const client = yield* TestLLM.Test
+      yield* client.push(TestLLM.toolCall("call-1", "echo", { value: "hi" }))
+      yield* client.push(textThenFail)
+
+      const echo = tool({
+        description: "Echo a value.",
+        parameters: Schema.Struct({ value: Schema.String }),
+        success: Schema.String,
+        execute: ({ value }) => Effect.succeed(value),
+      })
+      const error = yield* LLM.loop({ model, prompt: "Echo hi.", tools: { echo } }).pipe(Stream.runDrain, Effect.flip)
+
+      expect(error.partial?.text).toBe("Partial answer")
+      expect(error.partial?.toolCalls).toEqual([])
+    }),
+  )
+})
+
+describe("AIError.partial from the promise client", () => {
+  test("rejects with the root-exported AIError carrying partial", async () => {
+    const ai = AI.make({ clients: TestLLM.testLayer({ fallback: textThenFail }) })
+
+    const error = await ai.llm.generate({ model, prompt: "Hello" }).catch((error: unknown) => error)
+
+    expect(error).toBeInstanceOf(AIError)
+    expect(error instanceof AIError && error.partial?.text).toBe("Partial answer")
+    await ai.dispose()
+  })
+})

packages/ai/test/testing.test.ts

@@ -337,7 +337,9 @@ describe("TestLLM first-class client", () => {
         ),
       ).toBe(failure)
       expect(observed).toEqual(events)
-      expect(yield* llm.generate(request).pipe(Effect.flip)).toBe(failure)
+      const generated = yield* llm.generate(request).pipe(Effect.flip)
+      expect(generated.reason).toBe(failure.reason)
+      expect(generated.partial?.text).toBe("partial")
 
       const defect = new Error("Broken fixture responder")
       yield* llm.serve(() => {