effect 4.0.0-beta.104 → 4.0.0-beta.106

package/src/internal/schema/toArbitrary.ts

@@ -34,20 +34,6 @@ type LazyOption<T> = (
   recursionStack: RecursionStack
 ) => FastCheck.Arbitrary<T> | undefined
 
-export interface MutableReport {
-  readonly warnings: Array<Schema.Annotations.ToArbitrary.Warning>
-}
-
-/** @internal */
-export function makeReport(): MutableReport {
-  return { warnings: [] }
-}
-
-/** @internal */
-export function toReport(report: MutableReport): Schema.Annotations.ToArbitrary.Report {
-  return { warnings: report.warnings.slice() }
-}
-
 function arbitraryError(what: string) {
   return new Error(`Unable to derive an arbitrary for ${what}`)
 }
@@ -394,64 +380,6 @@ function finiteNumberContext(ctx: Context): Context {
   }
 }
 
-function reportChecks(report: MutableReport, checks: SchemaAST.Checks | undefined, path: ReadonlyArray<PropertyKey>) {
-  function visit(check: SchemaAST.Check<any>, covered: boolean) {
-    const arbitrary = check.annotations?.arbitrary
-    const nextCovered = covered || arbitrary?.constraint !== undefined || arbitrary?.candidate !== undefined
-    if (check._tag !== "Filter") {
-      for (const child of check.checks) {
-        visit(child, nextCovered)
-      }
-    } else if (!nextCovered) {
-      const description = check.annotations?.representation?.id ?? check.annotations?.identifier ??
-        check.annotations?.expected
-      report.warnings.push({ _tag: "OpaqueFilter", path, ...(description === undefined ? {} : { description }) })
-    }
-  }
-  checks?.forEach((check) => visit(check, false))
-}
-
-/** @internal */
-export function collectReport(ast: SchemaAST.AST, report: MutableReport) {
-  const stack = new WeakSet<SchemaAST.AST>()
-  function visit(ast: SchemaAST.AST, path: ReadonlyArray<PropertyKey>) {
-    if (stack.has(ast)) {
-      return
-    }
-    stack.add(ast)
-    reportChecks(report, ast.checks, path)
-    switch (ast._tag) {
-      case "Declaration":
-        ast.typeParameters.forEach((tp) => visit(tp, path))
-        break
-      case "Arrays": {
-        for (const [i, type] of [...ast.elements, ...ast.rest].entries()) {
-          visit(type, [...path, i])
-        }
-        break
-      }
-      case "Objects":
-        ast.propertySignatures.forEach((ps) => visit(ps.type, [...path, ps.name]))
-        ast.indexSignatures.forEach((is) => {
-          visit(is.parameter, path)
-          visit(is.type, path)
-        })
-        break
-      case "Union":
-        ast.types.forEach((type) => visit(type, path))
-        break
-      case "TemplateLiteral":
-        ast.parts.forEach((part, i) => visit(SchemaAST.toEncoded(part), [...path, i]))
-        break
-      case "Suspend":
-        visit(ast.thunk(), path)
-        break
-    }
-    stack.delete(ast)
-  }
-  visit(ast, [])
-}
-
 function applyCandidates(
   fc: typeof FastCheck,
   ctx: Context,

package/src/internal/schema/toJsonSchemaDocument.ts

@@ -1,7 +1,8 @@
 import * as Arr from "../../Array.ts"
 import * as Equal from "../../Equal.ts"
-import { escapeToken } from "../../JsonPointer.ts"
+import { escapeToken, unescapeToken } from "../../JsonPointer.ts"
 import type * as JsonSchema from "../../JsonSchema.ts"
+import { rewriteRefs } from "../../JsonSchema.ts"
 import * as RegEx from "../../RegExp.ts"
 import type * as Schema from "../../Schema.ts"
 import * as SchemaAST from "../../SchemaAST.ts"
@@ -21,6 +22,11 @@ const jsonSchemaAnnotationExcludedKeys = new Set([
   ...InternalAnnotations.jsonSchemaAnnotationKeys
 ])
 
+/** @internal */
+export const toRepresentationOptions = {
+  isAnonymousReferenceAllowed: (ast: SchemaAST.AST): boolean => !SchemaAST.isDeclaration(ast)
+}
+
 function collectJsonSchemaAnnotations(
   annotations: Schema.Annotations.Annotations | undefined,
   options: Schema.ToJsonSchemaOptions | undefined
@@ -142,22 +148,26 @@ function compileJsonSchema(
   references: SchemaRepresentation.References,
   options: Schema.ToJsonSchemaOptions | undefined
 ): JsonSchema.MultiDocument<"draft-2020-12"> {
-  const definitions: Record<string, JsonSchema.JsonSchema> = {}
   // null = compiling, string = canonical key, object = compiled schema
   const definitionStates = new Map<string, JsonSchema.JsonSchema | string | null>()
   const compiledRepresentations = new WeakMap<SchemaRepresentation.Representation, JsonSchema.JsonSchema>()
   const fallbackDefinitions = new Map<string, Array<string>>()
+  let hasAliases = false
   const referenceKeys = Object.keys(references)
   for (const key of referenceKeys) {
     compileDefinition(key, ["references", key])
   }
+  const schemas = Arr.map(
+    representations,
+    (representation, index) => finalizeJsonSchema(recur(representation, rootPaths[index]))
+  )
+  const definitions: Record<string, JsonSchema.JsonSchema> = {}
   for (const key of referenceKeys) {
     const compiled = definitionStates.get(key)!
     if (typeof compiled !== "string") {
-      InternalRecord.assignProperty(definitions, key, compiled)
+      InternalRecord.assignProperty(definitions, key, finalizeJsonSchema(compiled))
     }
   }
-  const schemas = Arr.map(representations, (representation, index) => recur(representation, rootPaths[index]))
   return { dialect: "draft-2020-12", schemas, definitions }
 
   function compileDefinition(key: string, path: Path): string {
@@ -179,6 +189,7 @@ function compileJsonSchema(
         if (candidates === undefined) fallbackDefinitions.set(fallback, [key])
         else candidates.push(key)
       } else {
+        hasAliases = true
         definitionStates.set(key, match)
         return match
       }
@@ -187,6 +198,15 @@ function compileJsonSchema(
     return key
   }
 
+  function finalizeJsonSchema(schema: JsonSchema.JsonSchema): JsonSchema.JsonSchema {
+    if (!hasAliases) return schema
+    return rewriteRefs(schema, ($ref) =>
+      $ref.replace(/^#\/\$defs\/([^/]*)/, (match, token) => {
+        const canonical = definitionStates.get(unescapeToken(token))
+        return typeof canonical === "string" ? `#/$defs/${escapeToken(canonical)}` : match
+      }))
+  }
+
   function getIdentifierFallback(
     representation: SchemaRepresentation.Representation
   ): string | undefined {

package/src/internal/schema/toRepresentation.ts

@@ -7,17 +7,24 @@ import * as InternalAnnotations from "./annotations.ts"
 
 /** @internal */
 export function toRepresentation(
-  ast: SchemaAST.AST
+  ast: SchemaAST.AST,
+  options?: Options
 ): SchemaRepresentation.Document {
-  const { references, representations } = toRepresentations([ast])
+  const { references, representations } = toRepresentations([ast], options)
   return { representation: representations[0], references }
 }
 
 /** @internal */
 export function toRepresentations(
-  asts: readonly [SchemaAST.AST, ...Array<SchemaAST.AST>]
+  asts: readonly [SchemaAST.AST, ...Array<SchemaAST.AST>],
+  options?: Options
 ): SchemaRepresentation.MultiDocument {
-  return fromASTs(asts)
+  return fromASTs(asts, options)
+}
+
+/** @internal */
+export interface Options {
+  readonly isAnonymousReferenceAllowed?: ((ast: SchemaAST.AST) => boolean) | undefined
 }
 
 type CheckRepresentationAnnotation = SchemaRepresentation.CheckRepresentationAnnotation<
@@ -28,11 +35,49 @@ function annotationsField<A>(annotations: A | undefined): { readonly annotations
   return annotations === undefined ? undefined : { annotations }
 }
 
-// Preserve repeated structural nodes as references without adding noise for leaf nodes.
-function isShareable(ast: SchemaAST.AST): boolean {
-  return SchemaAST.isArrays(ast) ||
-    SchemaAST.isObjects(ast) ||
-    (SchemaAST.isUnion(ast) && ast.types.some(isShareable))
+function hasShareableStructure(
+  ast: SchemaAST.AST,
+  isAnonymousReferenceAllowed: Options["isAnonymousReferenceAllowed"]
+): boolean {
+  if (isAnonymousReferenceAllowed?.(ast) === false) return false
+  switch (ast._tag) {
+    case "Arrays":
+    case "Objects":
+    case "Suspend":
+      return true
+    case "Declaration":
+      return true
+    case "Union":
+      return ast.types.some((ast) => hasShareableStructure(ast, isAnonymousReferenceAllowed))
+    default:
+      return false
+  }
+}
+
+function isWorthReferencing(bodyCost: number, occurrences: number): boolean {
+  return occurrences * bodyCost > bodyCost + occurrences + 1
+}
+
+function isAnonymousReferenceEligible(
+  ast: SchemaAST.AST,
+  occurrences: number,
+  isAnonymousReferenceAllowed: Options["isAnonymousReferenceAllowed"]
+): boolean {
+  if (isAnonymousReferenceAllowed?.(ast) === false) return false
+  if (hasShareableStructure(ast, isAnonymousReferenceAllowed)) return true
+  switch (ast._tag) {
+    case "Union":
+      return isWorthReferencing(ast.types.length + 1, occurrences)
+    case "Enum":
+      return isWorthReferencing(ast.enums.length + 1, occurrences)
+    case "TemplateLiteral":
+      return isWorthReferencing(ast.parts.length + 1, occurrences)
+    case "Literal":
+      return typeof ast.literal === "string" &&
+        isWorthReferencing(ast.literal.length / 32 + 1, occurrences)
+    default:
+      return false
+  }
 }
 
 interface ReferenceIdentifier {
@@ -54,17 +99,15 @@ function resolveReferenceIdentifier(
 }
 
 function fromASTs(
-  asts: readonly [SchemaAST.AST, ...Array<SchemaAST.AST>]
+  asts: readonly [SchemaAST.AST, ...Array<SchemaAST.AST>],
+  options: Options | undefined
 ): SchemaRepresentation.MultiDocument {
   const references: Record<string, SchemaRepresentation.Representation> = {}
   const anonymousReferences = new Map<SchemaAST.AST, string>()
   const referenceOwners = new Map<string, SchemaAST.AST>()
-  const valueIds = new Map<unknown, number>()
-  const canonicalByKey = new Map<string, SchemaAST.AST>()
-  let nextValueId = 0
   const buildingReferences = new Set<string>()
   const visiting = new Set<SchemaAST.AST>()
-  const visited = new Set<SchemaAST.AST>()
+  const occurrences = new Map<SchemaAST.AST, number>()
   const shared = new Set<SchemaAST.AST>()
 
   for (const ast of asts) visit(ast)
@@ -84,35 +127,6 @@ function fromASTs(
     return candidate
   }
 
-  function getValueId(value: unknown): number {
-    if (typeof value === "number" && globalThis.Number.isNaN(value)) {
-      return nextValueId++
-    }
-    const found = valueIds.get(value)
-    if (found !== undefined) return found
-    const id = nextValueId++
-    valueIds.set(value, id)
-    return id
-  }
-
-  function getIdentityKey(ast: SchemaAST.AST): string {
-    let identity = ast._tag
-    for (const [key, value] of Object.entries(ast)) {
-      if (key !== "_tag" && key !== "context") identity += `:${getValueId(value)}`
-    }
-    return identity
-  }
-
-  function getCanonicalAST(ast: SchemaAST.AST): SchemaAST.AST {
-    const key = getIdentityKey(ast)
-    const canonical = canonicalByKey.get(key)
-    if (canonical === undefined) {
-      canonicalByKey.set(key, ast)
-      return ast
-    }
-    return canonical
-  }
-
   function annotateReference(
     ast: SchemaAST.AST,
     referenceIdentifier: ReferenceIdentifier,
@@ -143,12 +157,16 @@ function fromASTs(
 
   function visit(input: SchemaAST.AST): void {
     const ast = SchemaAST.getLastEncoding(input)
-    const owner = getCanonicalAST(ast)
-    if (visited.has(owner)) {
-      if (isShareable(ast)) shared.add(owner)
+    const owner = SchemaAST.getContextOwner(ast)
+    const count = (occurrences.get(owner) ?? 0) + 1
+    occurrences.set(owner, count)
+    if (count > 1) {
+      if (
+        !shared.has(owner) &&
+        isAnonymousReferenceEligible(owner, count, options?.isAnonymousReferenceAllowed)
+      ) shared.add(owner)
       return
     }
-    visited.add(owner)
     visitChecks(ast.checks)
     switch (ast._tag) {
       case "Declaration":
@@ -178,7 +196,7 @@ function fromASTs(
 
   function recur(input: SchemaAST.AST): SchemaRepresentation.Representation {
     const ast = SchemaAST.getLastEncoding(input)
-    const owner = getCanonicalAST(ast)
+    const owner = SchemaAST.getContextOwner(ast)
     const referenceIdentifier = resolveReferenceIdentifier(input, ast)
     if (referenceIdentifier !== undefined) {
       const reference = getReference(referenceIdentifier.identifier, owner)

package/src/testing/TestSchema.ts

@@ -17,7 +17,7 @@ import * as Record from "../Record.ts"
 import * as Result from "../Result.ts"
 import * as Schema from "../Schema.ts"
 import * as SchemaAST from "../SchemaAST.ts"
-import type * as SchemaIssue from "../SchemaIssue.ts"
+import * as SchemaIssue from "../SchemaIssue.ts"
 import * as SchemaParser from "../SchemaParser.ts"
 import * as FastCheck from "../testing/FastCheck.ts"
 
@@ -123,7 +123,7 @@ export class Asserts<S extends Schema.Constraint> {
     async function succeed(input: S["~type.make.in"], expected?: S["Type"]) {
       const r = await Effect.runPromise(
         makeEffect(input, options).pipe(
-          Effect.mapErrorEager((issue) => issue.toString()),
+          Effect.mapErrorEager(SchemaIssue.defaultFormatter),
           Effect.result
         )
       )
@@ -135,7 +135,7 @@ export class Asserts<S extends Schema.Constraint> {
       async fail(input: unknown, message: string) {
         const r = await Effect.runPromise(
           makeEffect(input, options).pipe(
-            Effect.mapErrorEager((issue) => issue.toString()),
+            Effect.mapErrorEager(SchemaIssue.defaultFormatter),
             Effect.result
           )
         )
@@ -172,13 +172,13 @@ export class Asserts<S extends Schema.Constraint> {
   }) {
     const decodeUnknownEffect = SchemaParser.decodeUnknownEffect(this.schema)
     const encodeEffect = SchemaParser.encodeEffect(this.schema)
-    const arbitrary = Schema.toArbitrary(this.schema)
+    const arbitrary = Schema.toArbitrary(this.schema)(FastCheck)
     return FastCheck.assert(
       FastCheck.asyncProperty(arbitrary, async (t) => {
         const r = await Effect.runPromise(
           encodeEffect(t).pipe(
             Effect.flatMapEager((e) => decodeUnknownEffect(e)),
-            Effect.mapErrorEager((issue) => issue.toString()),
+            Effect.mapErrorEager(SchemaIssue.defaultFormatter),
             Effect.result
           )
         )
@@ -280,7 +280,7 @@ export class Asserts<S extends Schema.Constraint> {
       }) {
         const params = options?.params
         const is = Schema.is(schema)
-        const arb = Schema.toArbitrary(schema)
+        const arb = Schema.toArbitrary(schema)(FastCheck)
         FastCheck.assert(FastCheck.property(arb, (a) => is(a)), { numRuns: 20, ...params })
       }
     }
@@ -367,7 +367,7 @@ export class Decoding<S extends Schema.Constraint> {
   ) {
     const r = await Effect.runPromise(
       this.decodeUnknownEffect(input, this.options?.parseOptions).pipe(
-        Effect.mapErrorEager((issue) => issue.toString()),
+        Effect.mapErrorEager(SchemaIssue.defaultFormatter),
         Effect.result
       )
     )
@@ -401,7 +401,7 @@ export class Decoding<S extends Schema.Constraint> {
   ) {
     const r = await Effect.runPromise(
       this.decodeUnknownEffect(input, this.options?.parseOptions).pipe(
-        Effect.mapErrorEager((issue) => issue.toString()),
+        Effect.mapErrorEager(SchemaIssue.defaultFormatter),
         Effect.result
       )
     )
@@ -507,7 +507,7 @@ export class Encoding<S extends Schema.Constraint> {
   ) {
     const r = await Effect.runPromise(
       this.encodeUnknownEffect(input, this.options?.parseOptions).pipe(
-        Effect.mapErrorEager((issue) => issue.toString()),
+        Effect.mapErrorEager(SchemaIssue.defaultFormatter),
         Effect.result
       )
     )
@@ -541,7 +541,7 @@ export class Encoding<S extends Schema.Constraint> {
   ) {
     const r = await Effect.runPromise(
       this.encodeUnknownEffect(input, this.options?.parseOptions).pipe(
-        Effect.mapErrorEager((issue) => issue.toString()),
+        Effect.mapErrorEager(SchemaIssue.defaultFormatter),
         Effect.result
       )
     )

package/src/unstable/ai/McpServer.ts

@@ -489,7 +489,7 @@ const runWithProtocolState = Effect.fnUntraced(function*(options: {
   const isHttp = Option.isSome(yield* Effect.serviceOption(HttpRouter.HttpRouter))
   const sessions = protocolState.sessions
   const clientProtocols = new Map<number, McpProtocol.ProtocolAdapter>()
-  const cancelledRequests = new Map<number, Set<string>>()
+  const activeRequests = new Map<number, Map<string, boolean>>()
   const handlers = yield* Layer.build(layerHandlers(options, {
     sessions,
     protocolRegistry
@@ -577,11 +577,13 @@ const runWithProtocolState = Effect.fnUntraced(function*(options: {
     ...protocol,
     send: (clientId, response) => {
       if (response._tag === "Exit") {
-        const requests = cancelledRequests.get(clientId)
-        if (requests?.delete(requestKey(response.requestId)) === true) {
-          if (requests.size === 0) {
-            cancelledRequests.delete(clientId)
-          }
+        const requests = activeRequests.get(clientId)
+        const key = requestKey(response.requestId)
+        const cancelled = requests?.get(key)
+        if (requests !== undefined && requests.delete(key) && requests.size === 0) {
+          activeRequests.delete(clientId)
+        }
+        if (cancelled === true) {
           return Effect.void
         }
       }
@@ -670,9 +672,12 @@ const runWithProtocolState = Effect.fnUntraced(function*(options: {
               if (request.tag === "notifications/cancelled") {
                 return decodeCancelledNotification(request.payload).pipe(
                   Effect.flatMap(({ requestId }) => {
-                    const requests = cancelledRequests.get(clientId) ?? new Set<string>()
-                    requests.add(requestKey(requestId))
-                    cancelledRequests.set(clientId, requests)
+                    const key = requestKey(requestId)
+                    const requests = activeRequests.get(clientId)
+                    if (requests?.has(key) !== true) {
+                      return Effect.void
+                    }
+                    requests.set(key, true)
                     return f(clientId, {
                       _tag: "Interrupt",
                       requestId: RpcMessage.RequestId(requestId)
@@ -732,7 +737,14 @@ const runWithProtocolState = Effect.fnUntraced(function*(options: {
             }
             return selectedProtocol.payloadCodecs(rpc).decode(request.payload).pipe(
               Effect.matchEffect({
-                onSuccess: () => f(clientId, routedRequest),
+                onSuccess: () => {
+                  if (request.isNotification !== true) {
+                    const requests = activeRequests.get(clientId) ?? new Map<string, boolean>()
+                    requests.set(requestKey(request.id), false)
+                    activeRequests.set(clientId, requests)
+                  }
+                  return f(clientId, routedRequest)
+                },
                 onFailure: () =>
                   request.isNotification
                     ? Effect.void
@@ -752,7 +764,7 @@ const runWithProtocolState = Effect.fnUntraced(function*(options: {
           case "Interrupt":
             return f(clientId, request)
           case "Eof":
-            cancelledRequests.delete(clientId)
+            activeRequests.delete(clientId)
             return f(clientId, request)
           case "Pong":
           case "Exit":

package/src/unstable/ai/Prompt.ts

@@ -1840,20 +1840,30 @@ export const Prompt: Schema.Codec<Prompt, PromptEncoded> = Schema.Struct({
   Schema.decodeTo(
     $Prompt,
     SchemaTransformation.transformOrFail({
-      decode: (input) =>
+      decode: (input, options) =>
         Effect.mapBothEager(
           SchemaParser.decodeEffect(Schema.Array(Message))(input.content),
           {
             onSuccess: makePrompt,
-            onFailure: () => new SchemaIssue.InvalidValue({ message: "Invalid Prompt messages" })
+            onFailure: () =>
+              new SchemaIssue.InvalidValue(
+                { message: "Invalid Prompt messages" },
+                input.content,
+                options
+              )
           }
         ),
-      encode: (prompt) =>
+      encode: (prompt, options) =>
         Effect.mapBothEager(
           SchemaParser.encodeEffect(Schema.Array(Message))(prompt.content),
           {
             onSuccess: (messages) => ({ content: messages }),
-            onFailure: () => new SchemaIssue.InvalidValue({ message: "Invalid Prompt messages" })
+            onFailure: () =>
+              new SchemaIssue.InvalidValue(
+                { message: "Invalid Prompt messages" },
+                prompt.content,
+                options
+              )
           }
         )
     })
@@ -1994,11 +2004,23 @@ export const make = (input: RawInput): Prompt => {
  */
 export const fromMessages = (messages: ReadonlyArray<Message>): Prompt => makePrompt(messages)
 
+const mergeOptions = (left: ProviderOptions, right: ProviderOptions): ProviderOptions => {
+  const result: Record<string, ProviderOptions[string]> = { ...left }
+  for (const [provider, metadata] of Object.entries(right)) {
+    const previous = result[provider]
+    result[provider] = Predicate.isObject(previous) && Predicate.isObject(metadata)
+      ? Object.assign({}, previous, metadata)
+      : metadata
+  }
+  return result
+}
+
 /**
  * Creates a `Prompt` from response parts by folding completed text and
- * reasoning streams into assistant parts, placing tool calls and approval
- * requests in an assistant message, and placing non-preliminary tool results
- * in a tool message using their encoded results.
+ * reasoning streams into assistant parts, preserving provider metadata as
+ * prompt options, placing tool calls and approval requests in an assistant
+ * message, and placing non-preliminary tool results in a tool message using
+ * their encoded results.
  *
  * **Example** (Creating prompts from response parts)
  *
@@ -2042,55 +2064,63 @@ export const fromResponseParts = (parts: ReadonlyArray<Response.AnyPart>): Promp
   const assistantParts: Array<AssistantMessagePart> = []
   const toolParts: Array<ToolMessagePart> = []
 
-  const activeTextDeltas = new Map<string, { text: string }>()
-  const activeReasoningDeltas = new Map<string, { text: string }>()
+  const activeTextDeltas = new Map<string, { text: string; options: ProviderOptions }>()
+  const activeReasoningDeltas = new Map<string, { text: string; options: ProviderOptions }>()
 
   for (const part of parts) {
     switch (part.type) {
       // Text Parts
       case "text": {
-        assistantParts.push(makePart("text", { text: part.text }))
+        assistantParts.push(makePart("text", { text: part.text, options: part.metadata }))
         break
       }
 
       // Text Parts (streaming)
       case "text-start": {
-        activeTextDeltas.set(part.id, { text: "" })
+        activeTextDeltas.set(part.id, { text: "", options: part.metadata })
         break
       }
       case "text-delta": {
         if (activeTextDeltas.has(part.id)) {
-          activeTextDeltas.get(part.id)!.text += part.delta
+          const active = activeTextDeltas.get(part.id)!
+          active.text += part.delta
+          active.options = mergeOptions(active.options, part.metadata)
         }
         break
       }
       case "text-end": {
         if (activeTextDeltas.has(part.id)) {
-          assistantParts.push(makePart("text", activeTextDeltas.get(part.id)!))
+          const active = activeTextDeltas.get(part.id)!
+          active.options = mergeOptions(active.options, part.metadata)
+          assistantParts.push(makePart("text", active))
         }
         break
       }
 
       // Reasoning Parts
       case "reasoning": {
-        assistantParts.push(makePart("reasoning", { text: part.text }))
+        assistantParts.push(makePart("reasoning", { text: part.text, options: part.metadata }))
         break
       }
 
       // Reasoning Parts (streaming)
       case "reasoning-start": {
-        activeReasoningDeltas.set(part.id, { text: "" })
+        activeReasoningDeltas.set(part.id, { text: "", options: part.metadata })
         break
       }
       case "reasoning-delta": {
         if (activeReasoningDeltas.has(part.id)) {
-          activeReasoningDeltas.get(part.id)!.text += part.delta
+          const active = activeReasoningDeltas.get(part.id)!
+          active.text += part.delta
+          active.options = mergeOptions(active.options, part.metadata)
         }
         break
       }
       case "reasoning-end": {
         if (activeReasoningDeltas.has(part.id)) {
-          assistantParts.push(makePart("reasoning", activeReasoningDeltas.get(part.id)!))
+          const active = activeReasoningDeltas.get(part.id)!
+          active.options = mergeOptions(active.options, part.metadata)
+          assistantParts.push(makePart("reasoning", active))
         }
         break
       }
@@ -2101,7 +2131,8 @@ export const fromResponseParts = (parts: ReadonlyArray<Response.AnyPart>): Promp
           id: part.id,
           name: part.name,
           params: part.params,
-          providerExecuted: part.providerExecuted ?? false
+          providerExecuted: part.providerExecuted ?? false,
+          options: part.metadata
         }))
         break
       }
@@ -2115,7 +2146,8 @@ export const fromResponseParts = (parts: ReadonlyArray<Response.AnyPart>): Promp
             name: part.name,
             isFailure: part.isFailure,
             result: part.encodedResult,
-            providerExecuted: part.providerExecuted ?? false
+            providerExecuted: part.providerExecuted ?? false,
+            options: part.metadata
           }))
         }
         break

package/src/unstable/cluster/Entity.ts

@@ -34,7 +34,12 @@ import * as RpcClient from "../rpc/RpcClient.ts"
 import * as RpcGroup from "../rpc/RpcGroup.ts"
 import * as RpcSchema from "../rpc/RpcSchema.ts"
 import * as RpcServer from "../rpc/RpcServer.ts"
-import type { AlreadyProcessingMessage, MailboxFull, PersistenceError } from "./ClusterError.ts"
+import type {
+  AlreadyProcessingMessage,
+  EntityNotAssignedToRunner,
+  MailboxFull,
+  PersistenceError
+} from "./ClusterError.ts"
 import { Persisted, ShardGroup, Uninterruptible } from "./ClusterSchema.ts"
 import { EntityAddress } from "./EntityAddress.ts"
 import type { EntityId } from "./EntityId.ts"
@@ -117,7 +122,7 @@ export interface Entity<
       entityId: string
     ) => RpcClient.RpcClient.From<
       Rpcs,
-      MailboxFull | AlreadyProcessingMessage | PersistenceError
+      MailboxFull | AlreadyProcessingMessage | PersistenceError | EntityNotAssignedToRunner
     >,
     never,
     Sharding