effect 4.0.4 → 4.0.6

package/src/ai/Prompt.ts

@@ -465,9 +465,15 @@ export interface FilePartEncoded extends BasePartEncoded<"file", FilePartOptions
    */
   readonly fileName?: string | undefined
   /**
-   * File data as base64 string of data, a byte array, or a URL.
-   */
-  readonly data: string | Uint8Array | URL
+   * File data as a base64 string, a byte array, a URL, or the tagged JSON
+   * form that encoding produces for byte arrays and URLs.
+   */
+  readonly data:
+    | string
+    | Uint8Array
+    | URL
+    | { readonly _tag: "Uint8Array"; readonly base64: string }
+    | { readonly _tag: "URL"; readonly href: string }
 }
 
 /**
@@ -491,7 +497,18 @@ export const FilePart: Schema.Struct<{
   readonly type: Schema.Literal<"file">
   readonly mediaType: Schema.String
   readonly fileName: Schema.optional<Schema.String>
-  readonly data: Schema.Union<readonly [Schema.String, Schema.Uint8Array, Schema.URL]>
+  readonly data: Schema.Union<
+    readonly [
+      Schema.String,
+      Schema.decodeTo<
+        Schema.Uint8Array,
+        Schema.TaggedStruct<"Uint8Array", { readonly base64: Schema.Uint8ArrayFromBase64 }>
+      >,
+      Schema.decodeTo<Schema.URL, Schema.TaggedStruct<"URL", { readonly href: Schema.URLFromString }>>,
+      Schema.Uint8Array,
+      Schema.URL
+    ]
+  >
   readonly "~effect/ai/Prompt/Part": Schema.withDecodingDefaultKey<Schema.Literal<"~effect/ai/Prompt/Part">>
   readonly options: Schema.withDecodingDefault<
     Schema.$Record<
@@ -504,7 +521,33 @@ export const FilePart: Schema.Struct<{
   type: Schema.Literal("file"),
   mediaType: Schema.String,
   fileName: Schema.optional(Schema.String),
-  data: Schema.Union([Schema.String, Schema.Uint8Array, Schema.URL])
+  // The tagged members come first so that encoding produces JSON-safe data,
+  // while the raw members keep decoding byte arrays and URLs directly
+  data: Schema.Union([
+    Schema.String,
+    Schema.TaggedStruct("Uint8Array", { base64: Schema.Uint8ArrayFromBase64 }).pipe(Schema.decodeTo(
+      Schema.Uint8Array,
+      SchemaTransformation.transform({
+        decode: ({ base64 }) => base64,
+        encode: (base64) => (({
+          _tag: "Uint8Array",
+          base64
+        }) as const)
+      })
+    )),
+    Schema.TaggedStruct("URL", { href: Schema.URLFromString }).pipe(Schema.decodeTo(
+      Schema.URL,
+      SchemaTransformation.transform({
+        decode: ({ href }) => href,
+        encode: (href) => (({
+          _tag: "URL",
+          href
+        }) as const)
+      })
+    )),
+    Schema.Uint8Array,
+    Schema.URL
+  ])
 }).annotate({ identifier: "FilePart" })
 
 /**
@@ -1407,36 +1450,7 @@ export const UserMessage: Schema.Struct<{
         never,
         never
       >,
-      Schema.$Array<
-        Schema.Union<
-          readonly [
-            Schema.Struct<{
-              readonly type: Schema.Literal<"text">
-              readonly text: Schema.String
-              readonly "~effect/ai/Prompt/Part": Schema.withDecodingDefaultKey<Schema.Literal<"~effect/ai/Prompt/Part">>
-              readonly options: Schema.withDecodingDefault<
-                Schema.$Record<
-                  Schema.String,
-                  Schema.NullOr<Schema.Codec<Schema.Json>>
-                >
-              >
-            }>,
-            Schema.Struct<{
-              readonly type: Schema.Literal<"file">
-              readonly mediaType: Schema.String
-              readonly fileName: Schema.optional<Schema.String>
-              readonly data: Schema.Union<readonly [Schema.String, Schema.Uint8Array, Schema.URL]>
-              readonly "~effect/ai/Prompt/Part": Schema.withDecodingDefaultKey<Schema.Literal<"~effect/ai/Prompt/Part">>
-              readonly options: Schema.withDecodingDefault<
-                Schema.$Record<
-                  Schema.String,
-                  Schema.NullOr<Schema.Codec<Schema.Json>>
-                >
-              >
-            }>
-          ]
-        >
-      >
+      Schema.$Array<Schema.Union<readonly [typeof TextPart, typeof FilePart]>>
     ]
   >
   readonly "~effect/ai/Prompt/Message": Schema.withDecodingDefaultKey<Schema.Literal<"~effect/ai/Prompt/Message">>

package/src/ai/Response.ts

@@ -170,6 +170,21 @@ const toolPartSchemas = (tool: Tool.Any) => {
   return parts
 }
 
+type AnyToolParts = ToolCallPart<string, unknown> | ToolResultPart<string, unknown, unknown>
+
+const toolPartSchemasFor = (
+  toolkit: Toolkit.Any | Toolkit.WithHandler<any> | undefined
+): ReadonlyArray<Schema.Top> => {
+  if (toolkit === undefined) {
+    return [
+      makeToolCallPart(Schema.String, Schema.Unknown),
+      makeToolResultPart(Schema.String, Schema.Unknown, Schema.Unknown)
+    ]
+  }
+  const parts = Object.values(toolkit.tools as Record<string, Tool.Any>).map(toolPartSchemas)
+  return [...parts.map((part) => part.toolCall), ...parts.map((part) => part.toolResult)]
+}
+
 /**
  * Creates a Schema for all response parts based on a toolkit.
  *
@@ -178,6 +193,9 @@ const toolPartSchemas = (tool: Tool.Any) => {
  * Generates a schema that includes all possible response parts, with tool call
  * and tool result parts dynamically created based on the provided toolkit.
  *
+ * When no toolkit is provided, tool call and tool result parts accept any tool
+ * name, with parameters and results decoded as `unknown`.
+ *
  * **Example** (Building a response parts schema)
  *
  * ```ts import.meta.vitest
@@ -193,28 +211,95 @@ const toolPartSchemas = (tool: Tool.Any) => {
  *
  * const allPartsSchema = Response.AllParts(myToolkit)
  * Schema.isSchema(allPartsSchema) // => true
+ *
+ * const anyToolPartsSchema = Response.AllParts()
+ * Schema.isSchema(anyToolPartsSchema) // => true
  * ```
  *
  * @stability unstable
  * @category schemas
  * @since 4.0.0
  */
-export const AllParts = <T extends Toolkit.Any | Toolkit.WithHandler<any>>(
-  toolkit: T
-): Schema.Codec<
-  AllParts<T extends Toolkit.Any ? Toolkit.Tools<T> : Toolkit.WithHandlerTools<T>>,
-  AllPartsEncoded,
-  Tool.ResultDecodingServices<Toolkit.Tools<T>[keyof Toolkit.Tools<T>]>,
-  Tool.ResultEncodingServices<Toolkit.Tools<T>[keyof Toolkit.Tools<T>]>
-> => {
-  const toolCalls: Array<Schema.Top> = []
-  const toolResults: Array<Schema.Top> = []
-  for (const tool of Object.values(toolkit.tools as Record<string, Tool.Any>)) {
-    const parts = toolPartSchemas(tool)
-    toolCalls.push(parts.toolCall)
-    toolResults.push(parts.toolResult)
-  }
-  return Schema.Union([
+export const AllParts: {
+  /**
+   * Creates a Schema for all response parts based on a toolkit.
+   *
+   * **Details**
+   *
+   * Generates a schema that includes all possible response parts, with tool call
+   * and tool result parts dynamically created based on the provided toolkit.
+   *
+   * When no toolkit is provided, tool call and tool result parts accept any tool
+   * name, with parameters and results decoded as `unknown`.
+   *
+   * **Example** (Building a response parts schema)
+   *
+   * ```ts import.meta.vitest
+   * import { Schema } from "effect"
+   * import { Response, Tool, Toolkit } from "effect/ai"
+   *
+   * const myToolkit = Toolkit.make(
+   *   Tool.make("GetWeather", {
+   *     parameters: Schema.Struct({ city: Schema.String }),
+   *     success: Schema.Struct({ temperature: Schema.Number })
+   *   })
+   * )
+   *
+   * const allPartsSchema = Response.AllParts(myToolkit)
+   * Schema.isSchema(allPartsSchema) // => true
+   *
+   * const anyToolPartsSchema = Response.AllParts()
+   * Schema.isSchema(anyToolPartsSchema) // => true
+   * ```
+   *
+   * @stability unstable
+   * @category schemas
+   * @since 4.0.0
+   */
+  <T extends Toolkit.Any | Toolkit.WithHandler<any>>(toolkit: T): Schema.Codec<
+    AllParts<T extends Toolkit.Any ? Toolkit.Tools<T> : Toolkit.WithHandlerTools<T>>,
+    AllPartsEncoded,
+    Tool.ResultDecodingServices<Toolkit.Tools<T>[keyof Toolkit.Tools<T>]>,
+    Tool.ResultEncodingServices<Toolkit.Tools<T>[keyof Toolkit.Tools<T>]>
+  >
+  /**
+   * Creates a Schema for all response parts based on a toolkit.
+   *
+   * **Details**
+   *
+   * Generates a schema that includes all possible response parts, with tool call
+   * and tool result parts dynamically created based on the provided toolkit.
+   *
+   * When no toolkit is provided, tool call and tool result parts accept any tool
+   * name, with parameters and results decoded as `unknown`.
+   *
+   * **Example** (Building a response parts schema)
+   *
+   * ```ts import.meta.vitest
+   * import { Schema } from "effect"
+   * import { Response, Tool, Toolkit } from "effect/ai"
+   *
+   * const myToolkit = Toolkit.make(
+   *   Tool.make("GetWeather", {
+   *     parameters: Schema.Struct({ city: Schema.String }),
+   *     success: Schema.Struct({ temperature: Schema.Number })
+   *   })
+   * )
+   *
+   * const allPartsSchema = Response.AllParts(myToolkit)
+   * Schema.isSchema(allPartsSchema) // => true
+   *
+   * const anyToolPartsSchema = Response.AllParts()
+   * Schema.isSchema(anyToolPartsSchema) // => true
+   * ```
+   *
+   * @stability unstable
+   * @category schemas
+   * @since 4.0.0
+   */
+  (): Schema.Codec<AllParts<{}> | AnyToolParts, AllPartsEncoded>
+} = (toolkit?: Toolkit.Any | Toolkit.WithHandler<any>) =>
+  Schema.Union([
     TextPart,
     TextStartPart,
     TextDeltaPart,
@@ -233,10 +318,8 @@ export const AllParts = <T extends Toolkit.Any | Toolkit.WithHandler<any>>(
     ResponseMetadataPart,
     FinishPart,
     ErrorPart,
-    ...toolCalls,
-    ...toolResults
+    ...toolPartSchemasFor(toolkit)
   ]) as any
-}
 
 // =============================================================================
 // Parts
@@ -289,26 +372,49 @@ export type PartEncoded =
 /**
  * Creates a Schema for non-streaming response parts based on a toolkit.
  *
+ * **Details**
+ *
+ * When no toolkit is provided, tool call and tool result parts accept any tool
+ * name, with parameters and results decoded as `unknown`.
+ *
  * @stability unstable
  * @category schemas
  * @since 4.0.0
  */
-export const Part = <T extends Toolkit.Any | Toolkit.WithHandler<any>>(
-  toolkit: T
-): Schema.Codec<
-  Part<T extends Toolkit.Any ? Toolkit.Tools<T> : Toolkit.WithHandlerTools<T>>,
-  PartEncoded,
-  Tool.ResultDecodingServices<Toolkit.Tools<T>[keyof Toolkit.Tools<T>]>,
-  Tool.ResultEncodingServices<Toolkit.Tools<T>[keyof Toolkit.Tools<T>]>
-> => {
-  const toolCalls: Array<Schema.Top> = []
-  const toolResults: Array<Schema.Top> = []
-  for (const tool of Object.values(toolkit.tools as Record<string, Tool.Any>)) {
-    const parts = toolPartSchemas(tool)
-    toolCalls.push(parts.toolCall)
-    toolResults.push(parts.toolResult)
-  }
-  return Schema.Union([
+export const Part: {
+  /**
+   * Creates a Schema for non-streaming response parts based on a toolkit.
+   *
+   * **Details**
+   *
+   * When no toolkit is provided, tool call and tool result parts accept any tool
+   * name, with parameters and results decoded as `unknown`.
+   *
+   * @stability unstable
+   * @category schemas
+   * @since 4.0.0
+   */
+  <T extends Toolkit.Any | Toolkit.WithHandler<any>>(toolkit: T): Schema.Codec<
+    Part<T extends Toolkit.Any ? Toolkit.Tools<T> : Toolkit.WithHandlerTools<T>>,
+    PartEncoded,
+    Tool.ResultDecodingServices<Toolkit.Tools<T>[keyof Toolkit.Tools<T>]>,
+    Tool.ResultEncodingServices<Toolkit.Tools<T>[keyof Toolkit.Tools<T>]>
+  >
+  /**
+   * Creates a Schema for non-streaming response parts based on a toolkit.
+   *
+   * **Details**
+   *
+   * When no toolkit is provided, tool call and tool result parts accept any tool
+   * name, with parameters and results decoded as `unknown`.
+   *
+   * @stability unstable
+   * @category schemas
+   * @since 4.0.0
+   */
+  (): Schema.Codec<Part<{}> | AnyToolParts, PartEncoded>
+} = (toolkit?: Toolkit.Any | Toolkit.WithHandler<any>) =>
+  Schema.Union([
     TextPart,
     ReasoningPart,
     ToolApprovalRequestPart,
@@ -317,10 +423,8 @@ export const Part = <T extends Toolkit.Any | Toolkit.WithHandler<any>>(
     UrlSourcePart,
     ResponseMetadataPart,
     FinishPart,
-    ...toolCalls,
-    ...toolResults
+    ...toolPartSchemasFor(toolkit)
   ]) as any
-}
 
 // =============================================================================
 // Stream Parts
@@ -386,26 +490,49 @@ export type StreamPartEncoded =
 /**
  * Creates a Schema for streaming response parts based on a toolkit.
  *
+ * **Details**
+ *
+ * When no toolkit is provided, tool call and tool result parts accept any tool
+ * name, with parameters and results decoded as `unknown`.
+ *
  * @stability unstable
  * @category schemas
  * @since 4.0.0
  */
-export const StreamPart = <T extends Toolkit.Any | Toolkit.WithHandler<any>>(
-  toolkit: T
-): Schema.Codec<
-  StreamPart<T extends Toolkit.Any ? Toolkit.Tools<T> : Toolkit.WithHandlerTools<T>>,
-  StreamPartEncoded,
-  Tool.ResultDecodingServices<Toolkit.Tools<T>[keyof Toolkit.Tools<T>]>,
-  Tool.ResultEncodingServices<Toolkit.Tools<T>[keyof Toolkit.Tools<T>]>
-> => {
-  const toolCalls: Array<Schema.Top> = []
-  const toolResults: Array<Schema.Top> = []
-  for (const tool of Object.values(toolkit.tools as Record<string, Tool.Any>)) {
-    const parts = toolPartSchemas(tool)
-    toolCalls.push(parts.toolCall)
-    toolResults.push(parts.toolResult)
-  }
-  return Schema.Union([
+export const StreamPart: {
+  /**
+   * Creates a Schema for streaming response parts based on a toolkit.
+   *
+   * **Details**
+   *
+   * When no toolkit is provided, tool call and tool result parts accept any tool
+   * name, with parameters and results decoded as `unknown`.
+   *
+   * @stability unstable
+   * @category schemas
+   * @since 4.0.0
+   */
+  <T extends Toolkit.Any | Toolkit.WithHandler<any>>(toolkit: T): Schema.Codec<
+    StreamPart<T extends Toolkit.Any ? Toolkit.Tools<T> : Toolkit.WithHandlerTools<T>>,
+    StreamPartEncoded,
+    Tool.ResultDecodingServices<Toolkit.Tools<T>[keyof Toolkit.Tools<T>]>,
+    Tool.ResultEncodingServices<Toolkit.Tools<T>[keyof Toolkit.Tools<T>]>
+  >
+  /**
+   * Creates a Schema for streaming response parts based on a toolkit.
+   *
+   * **Details**
+   *
+   * When no toolkit is provided, tool call and tool result parts accept any tool
+   * name, with parameters and results decoded as `unknown`.
+   *
+   * @stability unstable
+   * @category schemas
+   * @since 4.0.0
+   */
+  (): Schema.Codec<StreamPart<{}> | AnyToolParts, StreamPartEncoded>
+} = (toolkit?: Toolkit.Any | Toolkit.WithHandler<any>) =>
+  Schema.Union([
     TextStartPart,
     TextDeltaPart,
     TextEndPart,
@@ -422,10 +549,8 @@ export const StreamPart = <T extends Toolkit.Any | Toolkit.WithHandler<any>>(
     ResponseMetadataPart,
     FinishPart,
     ErrorPart,
-    ...toolCalls,
-    ...toolResults
+    ...toolPartSchemasFor(toolkit)
   ]) as any
-}
 
 // =============================================================================
 // utility types
@@ -1522,15 +1647,17 @@ export const ToolCallPart: <const Name extends string, Params extends Schema.Con
 > = <const Name extends string, Params extends Schema.Constraint>(
   name: Name,
   params: Params
-) =>
+) => makeToolCallPart(Schema.Literal(name), params) as any
+
+const makeToolCallPart = (name: Schema.Top, params: Schema.Constraint) =>
   Schema.Struct({
     ...BasePart.fields,
     type: Schema.Literal("tool-call"),
     id: Schema.String,
-    name: Schema.Literal(name),
+    name,
     params,
     providerExecuted: Schema.Boolean.pipe(Schema.withDecodingDefaultKey(Effect.succeed(false)))
-  }).annotate({ identifier: "ToolCallPart" }) as any
+  }).annotate({ identifier: "ToolCallPart" })
 
 /**
  * Constructs a new tool call part.
@@ -1739,14 +1866,19 @@ export const ToolResultPart = <
   ToolResultPartEncoded,
   Success["DecodingServices"] | Failure["DecodingServices"],
   Success["EncodingServices"] | Failure["EncodingServices"]
-> => {
+> =>
+  makeToolResultPart(Schema.Literal(name), success, failure).annotate({
+    identifier: `ToolResultPart(${name})`
+  }) as any
+
+const makeToolResultPart = (name: Schema.Top, success: Schema.Constraint, failure: Schema.Constraint) => {
   const makeSchema = (result: Schema.Constraint, isFailure: boolean) => {
     const encodedResult = Schema.toEncoded(result)
     const Common = {
       id: Schema.String,
       type: Schema.Literal("tool-result"),
       isFailure: Schema.Literal(isFailure),
-      name: Schema.Literal(name)
+      name
     }
     const Decoded = Schema.Struct({
       ...Common,
@@ -1779,9 +1911,7 @@ export const ToolResultPart = <
       })
     ))
   }
-  return Schema.Union([makeSchema(success, false), makeSchema(failure, true)]).annotate({
-    identifier: `ToolResultPart(${name})`
-  })
+  return Schema.Union([makeSchema(success, false), makeSchema(failure, true)])
 }
 
 /**

package/src/ai/Tool.ts

@@ -115,7 +115,9 @@ export type DynamicTypeId = "~effect/ai/Tool/Dynamic"
  * execution will be returned in the error channel of the calling effect.
  *
  * If set to `"return"`, errors that occur during tool call handler execution
- * will be captured and returned as part of the tool call result.
+ * will be captured and returned as part of the tool call result. Like
+ * `Effect.catch`, only the first typed failure is returned, and a cause
+ * containing a defect is raised as a defect.
  *
  * @stability unstable
  * @category models
@@ -246,7 +248,9 @@ export interface Tool<
    * effect.
    *
    * If set to `"return"`, errors that occur during tool call handler execution
-   * will be captured and returned as part of the tool call result.
+   * will be captured and returned as part of the tool call result. Like
+   * `Effect.catch`, only the first typed failure is returned, and a cause
+   * containing a defect is raised as a defect.
    */
   readonly failureMode: FailureMode
 
@@ -1295,7 +1299,9 @@ export const make = <
    * execution will be returned in the error channel of the calling effect.
    *
    * If set to `"return"`, errors that occur during tool call handler execution
-   * will be captured and returned as part of the tool call result.
+   * will be captured and returned as part of the tool call result. Like
+   * `Effect.catch`, only the first typed failure is returned, and a cause
+   * containing a defect is raised as a defect.
    */
   readonly failureMode?: Mode
   /**
@@ -1578,7 +1584,9 @@ export const providerDefined = <
          * execution will be returned in the error channel of the calling effect.
          *
          * If set to `"return"`, errors that occur during tool call handler execution
-         * will be captured and returned as part of the tool call result.
+         * will be captured and returned as part of the tool call result. Like
+         * `Effect.catch`, only the first typed failure is returned, and a cause
+         * containing a defect is raised as a defect.
          */
         readonly failureMode?: Mode | undefined
       }

package/src/ai/Toolkit.ts

@@ -125,7 +125,8 @@ export interface HandlerContext<Tool extends Tool.Any> {
    * **Details**
    *
    * Preliminary results are streamed to the caller before the handler completes,
-   * enabling real-time progress updates for lengthy operations.
+   * enabling real-time progress updates for lengthy operations. Results are
+   * buffered without limit unless the caller provides `PreliminaryBuffer`.
    */
   readonly preliminary: (result: Tool.Success<Tool>) => Effect.Effect<void>
 }
@@ -296,9 +297,60 @@ export const FailureOrigin = Context.Reference<Tool.FailureOrigin>("effect/ai/To
   defaultValue: () => "result"
 })
 
+/**
+ * Limits how many preliminary results `Toolkit.handle` buffers per tool call.
+ *
+ * **Details**
+ *
+ * Defaults to `"unbounded"`. Provide it as a positive integer around `handle`,
+ * or around a program that calls it, and `HandlerContext.preliminary` suspends
+ * once that many results are unconsumed. Any other value is a defect.
+ *
+ * The bound only applies to the stream returned by `handle`. Consume that
+ * stream while the handler runs; waiting for the handler before reading it
+ * deadlocks once the buffer fills. `LanguageModel` moves results into its own
+ * buffers, so the bound does not limit memory end to end there.
+ *
+ * **Example** (Bounding preliminary results)
+ *
+ * ```ts import.meta.vitest
+ * import { Effect, Schema, Stream } from "effect"
+ * import { Tool, Toolkit } from "effect/ai"
+ *
+ * const toolkit = Toolkit.make(Tool.make("Progress", { success: Schema.Number }))
+ *
+ * const program = Effect.gen(function*() {
+ *   const handlers = yield* toolkit
+ *   const stream = yield* handlers.handle("Progress", {}).pipe(
+ *     Effect.provideService(Toolkit.PreliminaryBuffer, 16)
+ *   )
+ *   const results = yield* Stream.runCollect(stream)
+ *   return results.length
+ * }).pipe(
+ *   Effect.provide(toolkit.toLayer({
+ *     Progress: Effect.fnUntraced(function*(_, ctx) {
+ *       for (let i = 0; i < 64; i++) yield* ctx.preliminary(i)
+ *       return 64
+ *     })
+ *   }))
+ * )
+ *
+ * await Effect.runPromise(program) // => 65
+ * ```
+ *
+ * @stability unstable
+ * @category services
+ * @since 4.1.0
+ */
+export const PreliminaryBuffer = Context.Reference<number | "unbounded">("effect/ai/Toolkit/PreliminaryBuffer", {
+  defaultValue: () => "unbounded"
+})
+
 const failureCause = <E>(error: E, origin: Tool.FailureOrigin): Cause.Cause<E> =>
   Cause.annotate(Cause.fail(error), Context.make(FailureOrigin, origin))
 
+const failureOrigin = (error: unknown): Tool.FailureOrigin => Schema.isSchemaError(error) ? "result" : "handler"
+
 const Proto = {
   ...Effectable.Prototype({
     label: "Toolkit",
@@ -402,11 +454,15 @@ const Proto = {
         }
         const decodedParams = decodedParamsResult.success
 
+        const capacity = yield* PreliminaryBuffer
+        if (capacity !== "unbounded" && !(Number.isInteger(capacity) && capacity > 0)) {
+          return yield* Effect.die(`Tool "${name}": PreliminaryBuffer must be a positive integer, got ${capacity}`)
+        }
         const queue = yield* Queue.make<{
           readonly result: any
           readonly isFailure: boolean
           readonly preliminary: boolean
-        }, Cause.Done>()
+        }, Cause.Done>(capacity === "unbounded" ? undefined : { capacity })
         const context: HandlerContext<any> = {
           toolCallId,
           preliminary: (result) =>
@@ -448,18 +504,30 @@ const Proto = {
         }
 
         return Stream.fromQueue(queue).pipe(
-          Stream.catch((error) => {
-            const normalizedError = normalizeError(error)
-            const failureOrigin: Tool.FailureOrigin = Schema.isSchemaError(error) ? "result" : "handler"
-            return tool.failureMode === "error"
-              ? Stream.failCause(failureCause(normalizedError, failureOrigin))
-              : Stream.succeed({ result: normalizedError, isFailure: true, failureOrigin, preliminary: false })
-          }),
+          tool.failureMode === "error"
+            ? Stream.catchCause((cause) =>
+              Stream.failCause(Cause.map(
+                Cause.fromReasons(cause.reasons.map((reason) =>
+                  Cause.isFailReason(reason)
+                    ? reason.annotate(Context.make(FailureOrigin, failureOrigin(reason.error)), { overwrite: true })
+                    : reason
+                )),
+                normalizeError
+              ))
+            )
+            : Stream.catch((error) =>
+              Stream.succeed({
+                result: normalizeError(error),
+                isFailure: true,
+                failureOrigin: failureOrigin(error),
+                preliminary: false
+              })
+            ),
           Stream.mapEffect(Effect.fnUntraced(function*(output) {
             const encodedResult = yield* encodeResult(output.result, output.isFailure)
             return { ...output, encodedResult }
           })),
-          Stream.onEnd(Fiber.interrupt(fiber))
+          Stream.ensuring(Fiber.interrupt(fiber))
         ) satisfies Stream.Stream<Tool.HandlerResult<any>, any>
       })
 

package/src/ai/index.ts

@@ -53,6 +53,18 @@ export * as IdGenerator from "./IdGenerator.ts"
  */
 export * as LanguageModel from "./LanguageModel.ts"
 
+/**
+ * @stability experimental
+ * @since 4.0.6
+ */
+export * as McpClient from "./McpClient.ts"
+
+/**
+ * @stability experimental
+ * @since 4.0.6
+ */
+export * as McpClientError from "./McpClientError.ts"
+
 /**
  * @stability unstable
  * @since 4.0.0