refactor(ai): split the tool loop into small pieces (#61005)

packages/ai/README.md

@@ -59,8 +59,11 @@ it fails). Each step ends with its own `step-finish`; one `finish` ends the loop
 and the usage summed across steps. `until` is checked
 after each step that has tool calls, with every step's `LLMResponse`; returning `true` ends the loop without running
 them. Tools are an array, sent in the order given so the provider can reuse its cached prompt prefix. Each tool has a
-`name`, and two tools with the same name fail the request. `tool` takes a plain `async` execute; a throw or rejection
-goes back to the model as an error result, and `signal` aborts when the loop is interrupted.
+`name`, and two tools with the same name fail the request. Tool arguments that are not valid JSON go back to the model
+as an error result so it can retry. A step that hits the output token limit ends the loop without running its tool
+calls, which may be cut off. A forced `toolChoice` applies to the first step only; later steps use `auto`. `tool` takes
+a plain `async` execute; a throw or rejection goes back to the model as an error result, and `signal` aborts when the
+loop is interrupted.
 `input` is an Effect Schema that decodes and types the arguments, or a raw JSON Schema object, such as an MCP tool's,
 whose arguments arrive as `unknown` (pair it with an optional JSON Schema `outputSchema`). `tool` from
 `@opencode/ai/effect` takes the same options (including an optional `success` schema) with an Effect execute:

packages/ai/src/llm.ts

@@ -26,7 +26,6 @@ import {
   type CompactionCheckpointResponse,
   type CompactionResponse,
   type Finish,
-  type FinishReasonDetails,
   ToolChoice,
   ToolEntry,
   Usage,
@@ -37,7 +36,8 @@ import {
   withPartial,
 } from "./schema/index.js"
 import { tool, type ExecutableTool, type Tool, type ToolSchema } from "./tool.js"
-import { ToolRuntime } from "./tool-runtime.js"
+import { runTools } from "./tool-runtime.js"
+import { addUsage, nextRequest, pendingCalls } from "./loop.js"
 
 /** Input accepted by `ai.request`, `ai.generate`, `ai.stream` and `ai.compact`. An `LLMRequest` is also a valid input. */
 export type LLMRequestInput<Model extends LanguageModel = LanguageModel> = TextRequestFields<Model> & PromptOrMessages
@@ -117,18 +117,9 @@ export function stream(input: Input, options?: StreamOptions) {
   return Stream.unwrap(
     prepare(input).pipe(
       Effect.map(({ request, output }) =>
-        Stream.suspend(() => {
-          let state = LLMResponse.empty()
-          return LLMClient.stream(request, options).pipe(
-            Stream.mapEffect((event) => {
-              state = LLMResponse.reduce(state, event)
-              const response = event.type === "finish" ? LLMResponse.complete(state) : undefined
-              if (!response || event.type !== "finish") return Effect.succeed(event)
-              const finish = { ...event, usage: response.usage, message: response.message }
-              return output === undefined ? Effect.succeed(finish) : withObject(output, finish, state)
-            }),
-          )
-        }),
+        stepStream(request, undefined, options, (step) =>
+          Stream.fromEffect(finishStep(step, step.response.usage, output)),
+        ),
       ),
     ),
   )
@@ -142,8 +133,9 @@ export type LLMLoopInput<Model extends LanguageModel = LanguageModel> = Omit<Tex
   }
 
 /**
- * Stream model calls, running the model's tool calls between them, until a step has no client-executed tool calls or
- * `until` returns `true`. Each step's `step-start` and `step-finish` carry its step index and usage; one `finish` ends
+ * Stream model calls, running the model's tool calls between them, until a step has no client-executed tool calls,
+ * hits the output token limit, or `until` returns `true`. Tool calls whose input is not valid JSON get an error result
+ * so the model can retry. A forced `toolChoice` applies to the first step only. Each step's `step-start` and `step-finish` carry its step index and usage; one `finish` ends
  * the loop with the final step's reason and `message` and the usage summed across steps. With `schema`, the model ends
  * by calling the hidden answer tool, decoded into `object` on the `finish`.
  */
@@ -157,197 +149,104 @@ export function loop<const Model extends LanguageModel>(
 ): Stream.Stream<LLMEvent, AIError, Service>
 export function loop(input: LLMLoopInput & { readonly schema?: OutputSchema }, options?: StreamOptions) {
   const { tools, until, ...rest } = input
-  // A step that calls the answer tool is final, even if it also called other tools.
-  const clientCalls = (response: LLMResponse, output: Tool | undefined) =>
-    output && response.toolCalls.some((call) => call.name === OBJECT_TOOL)
-      ? []
-      : response.toolCalls.filter((call) => call.providerExecuted !== true)
-  const step = (
-    request: LLMRequest,
-    steps: ReadonlyArray<LLMResponse>,
-    usage: Usage | undefined,
-    output: Tool | undefined,
-  ): Stream.Stream<LLMEvent, AIError, Service> =>
-    Stream.suspend(() => {
-      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(() => (state = LLMResponse.reduce(state, event)))),
-        // Each call's `finish` ends only this step; the loop emits its own once it stops.
-        Stream.filter((event) => event.type !== "finish"),
-        Stream.mapError((error) => withPartial(error, LLMResponse.partial(state))),
-        Stream.concat(
-          Stream.unwrap(
-            Effect.gen(function* () {
-              const response = LLMResponse.complete(state)
-              const finish = state.events.findLast(LLMEvent.is.finish)
-              if (!response || !finish) return Stream.empty
-              const total = addUsage(usage, response.usage)
-              const calls = clientCalls(response, output)
-              const done = [...steps, response]
-              if (calls.length === 0 || until?.({ steps: done })) {
-                const event = { ...finish, usage: total, message: response.message }
-                return Stream.fromEffect(
-                  output === undefined || calls.length > 0 ? Effect.succeed(event) : withObject(output, event, state),
-                )
-              }
-              const results = yield* Effect.forEach(calls, (call) => ToolRuntime.dispatch(tools, call), {
-                concurrency: "unbounded",
-              })
-              const next = {
-                ...request,
-                messages: [
-                  ...request.messages,
-                  response.message,
-                  ...calls.map((call, i) =>
-                    Message.tool({
-                      id: call.id,
-                      name: call.name,
-                      namespace: call.namespace,
-                      result: results[i].result,
-                      providerMetadata: call.providerMetadata,
-                    }),
-                  ),
-                ],
-              }
-              return Stream.fromIterable(results.flatMap((result) => result.events)).pipe(
-                Stream.concat(step(next, done, total, output)),
-              )
-            }),
-          ),
-        ),
-      )
-    })
   return Stream.unwrap(
-    prepare({ ...rest, tools }).pipe(Effect.map(({ request, output }) => step(request, [], undefined, output))),
+    prepare({ ...rest, tools }).pipe(
+      Effect.map(({ request, output }) => driveLoop({ tools, until, output, options }, request, [], undefined)),
+    ),
   )
 }
 
-/** The result of `collect`: every step of a stream from `stream` or `loop`, folded into one value. */
-export interface Collected {
-  /** Text of the final step. */
-  readonly text: string
-  /** One response per model call, in order. */
-  readonly steps: ReadonlyArray<LLMResponse>
-  /** Each step's assistant message followed by its tool results, ready to append to the request's messages. */
-  readonly messages: ReadonlyArray<Message>
-  /** Usage from the terminal `finish`, summed across steps for a loop. */
-  readonly usage?: Usage
-  readonly finishReason: FinishReasonDetails
+export { collect, type Collected, type CollectResult } from "./loop.js"
+
+interface Step {
+  readonly response: LLMResponse
+  readonly finish: Finish
 }
 
-type FinishOf<Event> = Extract<Event, { readonly type: "finish" }>
-
-/** `Collected`, plus `object` when the stream was started with `schema`. */
-export type CollectResult<Event> = "object" extends keyof FinishOf<Event>
-  ? Collected & Pick<FinishOf<Event>, "object">
-  : Collected
-
-/** Run a stream from `stream` or `loop` to the end and fold its events into the final result. */
-export function collect<Event extends LLMEvent, E, R>(
-  events: Stream.Stream<Event, E, R>,
-): Effect.Effect<CollectResult<Event>, E | AIError, R>
-export function collect<E, R>(
-  events: Stream.Stream<LLMEvent, E, R>,
-): Effect.Effect<Collected & { readonly object?: unknown }, E | AIError, R> {
-  return Effect.suspend(() => {
-    const steps: LLMResponse[] = []
-    const messages: Message[] = []
-    let step = LLMResponse.empty()
-    // Between a `step-finish` and the next `step-start`, tool results come from the loop running tools.
-    let betweenSteps = false
+/**
+ * Stream one model call. Its `finish` is held back: `after` gets the completed step and streams what follows it.
+ * With `index`, the step's `step-start` and `step-finish` carry it.
+ */
+const stepStream = (
+  request: LLMRequest,
+  index: number | undefined,
+  options: StreamOptions | undefined,
+  after: (step: Step) => Stream.Stream<LLMEvent, AIError, Service>,
+): Stream.Stream<LLMEvent, AIError, Service> =>
+  Stream.suspend(() => {
+    let state = LLMResponse.empty()
     let finish: Finish | undefined
-    const close = (reason: FinishReasonDetails, usage: Usage | undefined) => {
-      if (step.events.length === 0) return
-      const response = new LLMResponse({
-        message: step.message,
-        events: step.events,
-        usage: usage ?? step.usage,
-        finishReason: reason,
-      })
-      steps.push(response)
-      messages.push(response.message)
-      step = LLMResponse.empty()
-    }
-    return Stream.runForEach(events, (event) =>
-      Effect.sync(() => {
-        if (event.type === "finish") {
-          close(event.reason, event.usage)
-          finish = event
-          return
-        }
-        if (event.type === "tool-result" && betweenSteps) {
-          messages.push(
-            Message.tool({
-              id: event.id,
-              name: event.name,
-              namespace: event.namespace,
-              result: event.result,
-              providerMetadata: event.providerMetadata,
-            }),
-          )
-          return
-        }
-        if (event.type === "step-start") betweenSteps = false
-        step = LLMResponse.reduce(step, event)
-        if (event.type !== "step-finish") return
-        close(event.reason, event.usage)
-        betweenSteps = true
-      }),
-    ).pipe(
-      Effect.flatMap(() => {
-        if (finish === undefined)
-          return Effect.fail(
-            new AIError({
-              reason: new InvalidProviderOutputError({
-                message: "The stream ended without a finish event.",
-                classification: "incomplete-stream",
-              }),
-            }),
-          )
-        const result = {
-          text: steps.at(-1)?.text ?? "",
-          steps,
-          messages,
-          usage: finish.usage,
-          finishReason: finish.reason,
-          ...("object" in finish ? { object: finish.object } : {}),
-        }
-        return Effect.succeed(result)
-      }),
+    return LLMClient.stream(request, options).pipe(
+      Stream.map((event) =>
+        index !== undefined && (event.type === "step-start" || event.type === "step-finish")
+          ? { ...event, index }
+          : event,
+      ),
+      Stream.tap((event) =>
+        Effect.sync(() => {
+          state = LLMResponse.reduce(state, event)
+          if (event.type === "finish") finish = event
+        }),
+      ),
+      Stream.filter((event) => event.type !== "finish"),
+      Stream.mapError((error) => withPartial(error, LLMResponse.partial(state))),
+      Stream.concat(
+        Stream.suspend(() => {
+          const response = LLMResponse.complete(state)
+          return response && finish ? after({ response, finish }) : Stream.empty
+        }),
+      ),
     )
   })
+
+/** The `finish` that ends a stream or loop: the final step's `message`, the given `usage`, and the decoded `object`. */
+const finishStep = (
+  step: Step,
+  usage: Usage | undefined,
+  output: Tool | undefined,
+): Effect.Effect<LLMEvent, AIError> => {
+  const event: Finish = { ...step.finish, usage, message: step.response.message }
+  if (output === undefined) return Effect.succeed(event)
+  return decodeObject(output, step.response).pipe(Effect.map((object) => ({ ...event, object })))
 }
 
-const addUsage = (total: Usage | undefined, step: Usage | undefined) => {
-  if (total === undefined || step === undefined) return total ?? step
-  const sum = (
-    key:
-      | "inputTokens"
-      | "outputTokens"
-      | "nonCachedInputTokens"
-      | "cacheReadInputTokens"
-      | "cacheWriteInputTokens"
-      | "reasoningTokens"
-      | "totalTokens"
-      | "cost",
-  ) => (total[key] === undefined && step[key] === undefined ? undefined : (total[key] ?? 0) + (step[key] ?? 0))
-  return new Usage({
-    // Context size describes the latest call, not a running total.
-    contextTokens: step.contextTokens ?? total.contextTokens,
-    inputTokens: sum("inputTokens"),
-    outputTokens: sum("outputTokens"),
-    nonCachedInputTokens: sum("nonCachedInputTokens"),
-    cacheReadInputTokens: sum("cacheReadInputTokens"),
-    cacheWriteInputTokens: sum("cacheWriteInputTokens"),
-    reasoningTokens: sum("reasoningTokens"),
-    totalTokens: sum("totalTokens"),
-    cost: sum("cost"),
-  })
+interface LoopContext {
+  readonly tools: ReadonlyArray<ExecutableTool>
+  readonly until: LLMLoopInput["until"]
+  readonly output: Tool | undefined
+  readonly options: StreamOptions | undefined
 }
 
+const driveLoop = (
+  context: LoopContext,
+  request: LLMRequest,
+  steps: ReadonlyArray<LLMResponse>,
+  usage: Usage | undefined,
+): Stream.Stream<LLMEvent, AIError, Service> =>
+  stepStream(request, steps.length, context.options, (step) => {
+    const total = addUsage(usage, step.response.usage)
+    const calls = pendingCalls(step.response, context.output && OBJECT_TOOL)
+    const done = [...steps, step.response]
+    if (calls.length === 0 || context.until?.({ steps: done }))
+      return Stream.fromEffect(finishStep(step, total, calls.length === 0 ? context.output : undefined))
+    return Stream.unwrap(
+      runTools(context.tools, calls).pipe(
+        Effect.map((results) =>
+          Stream.fromIterable(results.flatMap((result) => result.events)).pipe(
+            Stream.concat(
+              driveLoop(
+                context,
+                nextRequest(request, step.response, calls, results, context.output !== undefined),
+                done,
+                total,
+              ),
+            ),
+          ),
+        ),
+      ),
+    )
+  })
+
 const OBJECT_TOOL = "generate_object"
 
 type Input = LLMRequestInput & { readonly schema?: OutputSchema }
@@ -391,13 +290,6 @@ const decodeObject = (output: Tool, response: LLMResponse) => {
   return Effect.fail(invalid(`model did not call the forced \`${OBJECT_TOOL}\` tool`))
 }
 
-const withObject = (output: Tool, event: Finish, state: LLMResponse.State): Effect.Effect<LLMEvent, AIError> => {
-  const response = LLMResponse.complete(state)
-  return response === undefined
-    ? Effect.succeed(event)
-    : decodeObject(output, response).pipe(Effect.map((object) => ({ ...event, object })))
-}
-
 /** Normalize text request input into a plain `LLMRequest`, with `prompt` turned into the one user message. */
 export const request = <const Model extends LanguageModel>(input: LLMRequestInput<Model>): LLMRequest<Model> => {
   if ((input.prompt === undefined) === (input.messages === undefined))

packages/ai/src/loop.ts

@@ -0,0 +1,211 @@
+import { Effect, Stream } from "effect"
+import {
+  AIError,
+  InvalidProviderOutputError,
+  LLMResponse,
+  Message,
+  Usage,
+  LLMEvent,
+  ToolCallPart,
+  ToolChoice,
+  type Finish,
+  type FinishReasonDetails,
+  type LLMRequest,
+  type ToolInputError,
+} from "./schema/index.js"
+import type { DispatchResult, PendingCall } from "./tool-runtime.js"
+
+/**
+ * Client calls to answer before the next step. A call whose input was not valid JSON carries the error the model
+ * gets back, so it can retry. A step is final when it calls `answerTool`, even alongside other tools, or when it hits
+ * the output token limit, which may have cut its calls off.
+ */
+export const pendingCalls = (response: LLMResponse, answerTool?: string): ReadonlyArray<PendingCall> => {
+  if (response.finishReason.normalized === "length") return []
+  if (answerTool !== undefined && response.toolCalls.some((call) => call.name === answerTool)) return []
+  return [
+    ...response.toolCalls.filter((call) => call.providerExecuted !== true).map((call) => ({ call })),
+    ...response.events.filter(LLMEvent.is.toolInputError).map((event) => ({
+      call: invalidCall(event),
+      error: `Invalid tool arguments: the input is not valid JSON. Received: ${event.raw}`,
+    })),
+  ]
+}
+
+/** The call to replay for input that was not valid JSON. Providers require an object input, so it is empty. */
+const invalidCall = (event: ToolInputError) =>
+  ToolCallPart.make({ id: event.id, name: event.name, namespace: event.namespace, input: {} })
+
+/** Add the replayed calls for invalid input to the assistant message, so each tool result has its call. */
+const withInvalidCalls = (message: Message, calls: ReadonlyArray<ToolCallPart>) =>
+  calls.length === 0 ? message : new Message({ ...message, content: [...message.content, ...calls] })
+
+/** The step's assistant message followed by one tool message per call, in call order. */
+export const stepMessages = (
+  response: LLMResponse,
+  calls: ReadonlyArray<PendingCall>,
+  results: ReadonlyArray<DispatchResult>,
+): ReadonlyArray<Message> => [
+  withInvalidCalls(
+    response.message,
+    calls.flatMap((item) => (item.error === undefined ? [] : [item.call])),
+  ),
+  ...calls.map((item, i) =>
+    Message.tool({
+      id: item.call.id,
+      name: item.call.name,
+      namespace: item.call.namespace,
+      result: results[i].result,
+      providerMetadata: item.call.providerMetadata,
+    }),
+  ),
+]
+
+/**
+ * The request for the step after `response`, with its messages and tool results appended. A caller's forced
+ * `toolChoice` applies to the first step only; with `keepForcedChoice` (the answer tool's) it stays.
+ */
+export const nextRequest = (
+  request: LLMRequest,
+  response: LLMResponse,
+  calls: ReadonlyArray<PendingCall>,
+  results: ReadonlyArray<DispatchResult>,
+  keepForcedChoice: boolean,
+): LLMRequest => ({
+  ...request,
+  toolChoice:
+    !keepForcedChoice && (request.toolChoice?.type === "required" || request.toolChoice?.type === "tool")
+      ? new ToolChoice({ type: "auto", disableParallelToolUse: request.toolChoice.disableParallelToolUse })
+      : request.toolChoice,
+  messages: [...request.messages, ...stepMessages(response, calls, results)],
+})
+
+/** Sum usage across steps. */
+export const addUsage = (total: Usage | undefined, step: Usage | undefined) => {
+  if (total === undefined || step === undefined) return total ?? step
+  const sum = (
+    key:
+      | "inputTokens"
+      | "outputTokens"
+      | "nonCachedInputTokens"
+      | "cacheReadInputTokens"
+      | "cacheWriteInputTokens"
+      | "reasoningTokens"
+      | "totalTokens"
+      | "cost",
+  ) => (total[key] === undefined && step[key] === undefined ? undefined : (total[key] ?? 0) + (step[key] ?? 0))
+  return new Usage({
+    // Context size describes the latest call, not a running total.
+    contextTokens: step.contextTokens ?? total.contextTokens,
+    inputTokens: sum("inputTokens"),
+    outputTokens: sum("outputTokens"),
+    nonCachedInputTokens: sum("nonCachedInputTokens"),
+    cacheReadInputTokens: sum("cacheReadInputTokens"),
+    cacheWriteInputTokens: sum("cacheWriteInputTokens"),
+    reasoningTokens: sum("reasoningTokens"),
+    totalTokens: sum("totalTokens"),
+    cost: sum("cost"),
+  })
+}
+
+/** The result of `collect`: every step of a stream from `stream` or `loop`, folded into one value. */
+export interface Collected {
+  /** Text of the final step. */
+  readonly text: string
+  /** One response per model call, in order. */
+  readonly steps: ReadonlyArray<LLMResponse>
+  /** Each step's assistant message followed by its tool results, ready to append to the request's messages. */
+  readonly messages: ReadonlyArray<Message>
+  /** Usage from the terminal `finish`, summed across steps for a loop. */
+  readonly usage?: Usage
+  readonly finishReason: FinishReasonDetails
+}
+
+type FinishOf<Event> = Extract<Event, { readonly type: "finish" }>
+
+/** `Collected`, plus `object` when the stream was started with `schema`. */
+export type CollectResult<Event> = "object" extends keyof FinishOf<Event>
+  ? Collected & Pick<FinishOf<Event>, "object">
+  : Collected
+
+/** Run a stream from `stream` or `loop` to the end and fold its events into the final result. */
+export function collect<Event extends LLMEvent, E, R>(
+  events: Stream.Stream<Event, E, R>,
+): Effect.Effect<CollectResult<Event>, E | AIError, R>
+export function collect<E, R>(
+  events: Stream.Stream<LLMEvent, E, R>,
+): Effect.Effect<Collected & { readonly object?: unknown }, E | AIError, R> {
+  return Effect.suspend(() => {
+    const steps: LLMResponse[] = []
+    const messages: Message[] = []
+    let step = LLMResponse.empty()
+    // Between a `step-finish` and the next `step-start`, tool results come from the loop running tools.
+    let betweenSteps = false
+    // The last step's assistant message, and its calls with invalid input that the loop may answer.
+    let assistant = -1
+    let invalid: ReadonlyArray<ToolCallPart> = []
+    let finish: Finish | undefined
+    const close = (reason: FinishReasonDetails, usage: Usage | undefined) => {
+      if (step.events.length === 0) return
+      const response = new LLMResponse({
+        message: step.message,
+        events: step.events,
+        usage: usage ?? step.usage,
+        finishReason: reason,
+      })
+      steps.push(response)
+      assistant = messages.push(response.message) - 1
+      invalid = step.events.filter(LLMEvent.is.toolInputError).map(invalidCall)
+      step = LLMResponse.empty()
+    }
+    return Stream.runForEach(events, (event) =>
+      Effect.sync(() => {
+        if (event.type === "finish") {
+          close(event.reason, event.usage)
+          finish = event
+          return
+        }
+        if (event.type === "tool-result" && betweenSteps) {
+          const call = invalid.find((item) => item.id === event.id)
+          if (call) messages[assistant] = withInvalidCalls(messages[assistant], [call])
+          messages.push(
+            Message.tool({
+              id: event.id,
+              name: event.name,
+              namespace: event.namespace,
+              result: event.result,
+              providerMetadata: event.providerMetadata,
+            }),
+          )
+          return
+        }
+        if (event.type === "step-start") betweenSteps = false
+        step = LLMResponse.reduce(step, event)
+        if (event.type !== "step-finish") return
+        close(event.reason, event.usage)
+        betweenSteps = true
+      }),
+    ).pipe(
+      Effect.flatMap(() => {
+        if (finish === undefined)
+          return Effect.fail(
+            new AIError({
+              reason: new InvalidProviderOutputError({
+                message: "The stream ended without a finish event.",
+                classification: "incomplete-stream",
+              }),
+            }),
+          )
+        const result = {
+          text: steps.at(-1)?.text ?? "",
+          steps,
+          messages,
+          usage: finish.usage,
+          finishReason: finish.reason,
+          ...("object" in finish ? { object: finish.object } : {}),
+        }
+        return Effect.succeed(result)
+      }),
+    )
+  })
+}

packages/ai/src/tool-runtime.ts

@@ -41,6 +41,23 @@ export const dispatch = (tools: ReadonlyArray<Tool>, call: ToolCallPart): Effect
   )
 }
 
+/** A call the loop answers before the next step. With `error`, it is not run and the model gets the error back. */
+export interface PendingCall {
+  readonly call: ToolCallPart
+  readonly error?: string
+}
+
+/** Run a step's tool calls concurrently, with results in call order. */
+export const runTools = (tools: ReadonlyArray<Tool>, calls: ReadonlyArray<PendingCall>) =>
+  Effect.forEach(
+    calls,
+    (item) =>
+      item.error === undefined
+        ? dispatch(tools, item.call)
+        : Effect.succeed(result(item.call, { type: "error", value: item.error })),
+    { concurrency: "unbounded" },
+  )
+
 const decodeAndExecute = (tool: Tool, call: ToolCallPart): Effect.Effect<ToolSettlement, ToolFailure> =>
   tool._decode(call.input).pipe(
     Effect.mapError((error) => new ToolFailure({ message: `Invalid tool input: ${error.message}` })),

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

@@ -76,6 +76,18 @@ describe("AIError.partial", () => {
     }),
   )
 
+  it.effect("stream 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* ai.stream({ model, prompt: "Hello" }).pipe(Stream.runDrain, Effect.flip)
+
+      expect(error.reason).toBe(overloaded.reason)
+      expect(error.partial?.text).toBe("Partial answer")
+    }),
+  )
+
   it.effect("loop reports only the failed step as partial", () =>
     Effect.gen(function* () {
       const client = yield* TestLLM.Test

packages/ai/test/loop-steps.test.ts

@@ -0,0 +1,77 @@
+import { describe, expect, test } from "bun:test"
+import { LanguageModel, LLMEvent, LLMResponse, Message, Usage } from "../src/effect.js"
+import { addUsage, nextRequest, pendingCalls } from "../src/loop.js"
+import { OpenAIChat } from "../src/protocols/openai-chat.js"
+import { LLM } from "../src/llm.js"
+
+const model = LanguageModel.make({ id: "fictional-model", provider: "fixture", route: OpenAIChat.route })
+
+const step = (...events: ReadonlyArray<LLMEvent>) =>
+  LLMResponse.fromEvents([...events, LLMEvent.finish({ reason: { normalized: "tool-calls" } })])!
+const cutOff = (...events: ReadonlyArray<LLMEvent>) =>
+  LLMResponse.fromEvents([...events, LLMEvent.finish({ reason: { normalized: "length" } })])!
+
+const read = LLMEvent.toolCall({ id: "call-1", name: "read_file", input: { path: "a.txt" } })
+const hosted = LLMEvent.toolCall({ id: "call-2", name: "web_search", input: {}, providerExecuted: true })
+const invalid = LLMEvent.toolInputError({ id: "call-4", name: "read_file", raw: '{"path": "a' })
+const answer = LLMEvent.toolCall({ id: "call-3", name: "generate_object", input: { text: "hi" } })
+
+describe("pendingCalls", () => {
+  test("returns client calls and skips provider-executed ones", () => {
+    expect(pendingCalls(step(read, hosted)).map((item) => item.call.id)).toEqual(["call-1"])
+  })
+
+  test("answers input that is not valid JSON with an error and an empty replayed call", () => {
+    expect(pendingCalls(step(invalid))).toEqual([
+      {
+        call: expect.objectContaining({ type: "tool-call", id: "call-4", name: "read_file", input: {} }),
+        error: 'Invalid tool arguments: the input is not valid JSON. Received: {"path": "a',
+      },
+    ])
+  })
+
+  test("returns nothing when the step hit the output token limit", () => {
+    expect(pendingCalls(cutOff(read, invalid))).toEqual([])
+  })
+
+  test("returns nothing once the answer tool is called", () => {
+    expect(pendingCalls(step(read, answer), "generate_object")).toEqual([])
+    expect(pendingCalls(step(read, answer)).map((item) => item.call.id)).toEqual(["call-1", "call-3"])
+  })
+})
+
+describe("nextRequest", () => {
+  test("appends the assistant message and one tool message per call", () => {
+    const request = LLM.request({ model, prompt: "Read the file." })
+    const response = step(read)
+    const results = [{ result: { type: "text", value: "hello" } as const, events: [] }]
+    const next = nextRequest(request, response, [{ call: read }], results, false)
+
+    expect(next.messages.map((message) => message.role)).toEqual(["user", "assistant", "tool"])
+    expect(next.messages[1]).toBe(response.message)
+    expect(next.messages[2]).toEqual(
+      Message.tool({ id: "call-1", name: "read_file", result: { type: "text", value: "hello" } }),
+    )
+  })
+
+  test("resets a forced toolChoice unless it is kept, and adds the replayed call for invalid input", () => {
+    const request = LLM.request({ model, prompt: "Read the file.", toolChoice: "required" })
+    const response = step(invalid)
+    const calls = pendingCalls(response)
+    const results = [{ result: { type: "error", value: "bad" } as const, events: [] }]
+
+    expect(nextRequest(request, response, calls, results, false).toolChoice?.type).toBe("auto")
+    expect(nextRequest(request, response, calls, results, true).toolChoice?.type).toBe("required")
+    expect(nextRequest(request, response, calls, results, false).messages[1].content).toEqual([calls[0].call])
+  })
+})
+
+describe("addUsage", () => {
+  test("sums counts and keeps the latest context size", () => {
+    const total = addUsage(
+      new Usage({ inputTokens: 10, outputTokens: 2, contextTokens: 12 }),
+      new Usage({ inputTokens: 20, contextTokens: 34 }),
+    )
+    expect(total).toEqual(new Usage({ inputTokens: 30, outputTokens: 2, contextTokens: 34 }))
+  })
+})

packages/ai/test/loop.test.ts

@@ -167,6 +167,102 @@ describe("ai.loop", () => {
       expect(events.at(-1)?.type).toBe("finish")
     }),
   )
+  it.effect("sends unparseable tool input back to the model as an error result and continues", () =>
+    Effect.gen(function* () {
+      const client = yield* TestLLM.Test
+      yield* client.push(
+        TestLLM.toolCalls(LLMEvent.toolInputError({ id: "call-1", name: "read_file", raw: '{"path": "a.txt"' })),
+        TestLLM.toolCall("call-2", "read_file", { path: "a.txt" }),
+        TestLLM.text("The file says hello.", "text-1"),
+      )
+
+      const events = yield* collect()
+      const requests = yield* client.requests()
+      const error = {
+        type: "error",
+        value: 'Invalid tool arguments: the input is not valid JSON. Received: {"path": "a.txt"',
+      } as const
+
+      expect(requests).toHaveLength(3)
+      expect(requests[1].messages[1].content).toEqual([
+        expect.objectContaining({ type: "tool-call", id: "call-1", name: "read_file", input: {} }),
+      ])
+      expect(requests[1].messages[2].content).toEqual([
+        expect.objectContaining({ type: "tool-result", id: "call-1", result: error }),
+      ])
+      expect(events.filter(LLMEvent.is.toolResult).map((event) => event.result)).toEqual([
+        error,
+        { type: "text", value: "hello" },
+      ])
+    }),
+  )
+
+  it.effect("stops at the output token limit without running the step's tool calls", () =>
+    Effect.gen(function* () {
+      const client = yield* TestLLM.Test
+      yield* client.push(
+        TestLLM.complete(
+          { reason: { normalized: "length" } },
+          LLMEvent.toolCall({ id: "call-1", name: "read_file", input: { path: "a.txt" } }),
+          LLMEvent.toolInputError({ id: "call-2", name: "read_file", raw: '{"path": "b' }),
+        ),
+      )
+
+      const events = yield* collect()
+
+      expect(yield* client.requests()).toHaveLength(1)
+      expect(events.filter(LLMEvent.is.toolResult)).toHaveLength(0)
+      expect(events.at(-1)).toEqual(expect.objectContaining({ type: "finish", reason: { normalized: "length" } }))
+    }),
+  )
+
+  it.effect("forces the caller's toolChoice on the first step only", () =>
+    Effect.gen(function* () {
+      const client = yield* TestLLM.Test
+      yield* client.push(
+        TestLLM.toolCall("call-1", "read_file", { path: "a.txt" }),
+        TestLLM.text("The file says hello.", "text-1"),
+      )
+
+      yield* ai
+        .loop({
+          model,
+          prompt: "Read the file.",
+          tools: [read_file],
+          toolChoice: { type: "required", disableParallelToolUse: true },
+        })
+        .pipe(Stream.runDrain)
+      const requests = yield* client.requests()
+
+      expect(requests.map((request) => ({ ...request.toolChoice }))).toEqual([
+        { type: "required", disableParallelToolUse: true },
+        { type: "auto", disableParallelToolUse: true },
+      ])
+    }),
+  )
+
+  it.effect("keeps the answer tool forced on every step with schema", () =>
+    Effect.gen(function* () {
+      const client = yield* TestLLM.Test
+      yield* client.push(
+        TestLLM.toolCall("call-1", "read_file", { path: "a.txt" }),
+        TestLLM.toolCall("call-2", "generate_object", { text: "hello" }),
+      )
+
+      yield* ai
+        .loop({
+          model,
+          prompt: "Read the file.",
+          tools: [read_file],
+          toolChoice: "read_file",
+          schema: Schema.Struct({ text: Schema.String }),
+        })
+        .pipe(Stream.runDrain)
+
+      expect((yield* client.requests()).map((request) => request.toolChoice?.type)).toEqual(["required", "required"])
+    }),
+  )
+
   it.effect("fails without running tools when a step's stream fails", () =>
     Effect.gen(function* () {
       const client = yield* TestLLM.Test
@@ -287,6 +383,25 @@ describe("ai.collect", () => {
     }),
   )
 
+  it.effect("gives each error result for unparseable tool input its call, as the loop sent it", () =>
+    Effect.gen(function* () {
+      const client = yield* TestLLM.Test
+      yield* client.push(
+        TestLLM.toolCalls(LLMEvent.toolInputError({ id: "call-1", name: "read_file", raw: '{"path": "a.txt"' })),
+        TestLLM.text("Sorry.", "text-1"),
+      )
+
+      const result = yield* ai.collect(ai.loop({ model, prompt: "Read the file.", tools: [read_file] }))
+      const requests = yield* client.requests()
+
+      expect(result.messages.map((message) => message.role)).toEqual(["assistant", "tool", "assistant"])
+      expect(result.messages[0].content).toEqual([
+        expect.objectContaining({ type: "tool-call", id: "call-1", name: "read_file", input: {} }),
+      ])
+      expect(result.messages.slice(0, 2)).toEqual(requests[1].messages.slice(1))
+    }),
+  )
+
   it.effect("carries the decoded object when the stream has a schema", () =>
     Effect.gen(function* () {
       const client = yield* TestLLM.Test

packages/ai/test/recorded-scenarios.ts

@@ -1,5 +1,5 @@
 import { expect } from "bun:test"
-import { Effect, Schema } from "effect"
+import { Effect, Schema, Stream } from "effect"
 import {
   ai,
   type ContentPart,
@@ -11,10 +11,8 @@ import {
   LLMResponse,
   Media,
   Message,
-  toDefinitions,
   ToolChoice,
   ToolDefinition,
-  ToolRuntime,
 } from "../src/effect.js"
 import { LLMClient } from "../src/route.js"
 import { tool } from "../src/tool.js"
@@ -92,40 +90,7 @@ const restroomImage = () =>
   )
 
 export const runWeatherToolLoop = (request: LLMRequest) =>
-  Effect.gen(function* () {
-    const tools = [weatherRuntimeTool]
-    let next = { ...request, tools: toDefinitions(tools) }
-    const events: LLMEvent[] = []
-
-    for (let step = 0; step < 10; step++) {
-      const response = yield* LLMClient.generate(next)
-      events.push(...response.events.filter((event) => event.type !== "finish"))
-      const calls = response.events.filter(LLMEvent.is.toolCall).filter((call) => !call.providerExecuted)
-      if (calls.length === 0) {
-        const finish = response.events.find(LLMEvent.is.finish)
-        if (finish) events.push(finish)
-        return events
-      }
-
-      const dispatched = yield* Effect.forEach(calls, (call) =>
-        ToolRuntime.dispatch(tools, call).pipe(Effect.map((result) => [call, result] as const)),
-      )
-      events.push(...dispatched.flatMap(([, result]) => result.events))
-      next = {
-        ...next,
-        messages: [
-          ...next.messages,
-          Message.assistant(assistantContent(response.events)),
-          ...dispatched.map(([call, result]) => Message.tool({ id: call.id, name: call.name, result: result.result })),
-        ],
-      }
-    }
-
-    throw new Error("Weather tool loop exceeded 10 steps")
-  })
-
-const assistantContent = (events: ReadonlyArray<LLMEvent>) =>
-  events.reduce(LLMResponse.reduce, LLMResponse.empty()).message.content
+  ai.loop({ ...request, tools: [weatherRuntimeTool] }).pipe(Stream.runCollect)
 
 export const expectFinish = (events: ReadonlyArray<LLMEvent>, reason: FinishReason) =>
   expect(events.at(-1)).toMatchObject({ type: "finish", reason: { normalized: reason } })