refactor(ai): replace loose Schema.Unknown and catch-all union checks in protocols (#57606)

packages/ai/src/protocols/cohere-chat.ts

@@ -4,7 +4,7 @@ import { Endpoint } from "../route/endpoint.js"
 import { Framing } from "../route/framing.js"
 import { Protocol } from "../route/protocol.js"
 import { LLMEvent, Usage, type FinishReasonDetails, type LLMRequest } from "../schema/index.js"
-import { ProviderShared } from "./shared.js"
+import { JsonObject, ProviderShared } from "./shared.js"
 import { Lifecycle } from "./utils/lifecycle.js"
 import { ToolStream } from "./utils/tool-stream.js"
 
@@ -49,7 +49,7 @@ const Body = Schema.Struct({
         function: Schema.Struct({
           name: Schema.String,
           description: Schema.optional(Schema.String),
-          parameters: Schema.Unknown,
+          parameters: JsonObject,
         }),
       }),
     ),

packages/ai/src/protocols/gemini.ts

@@ -248,8 +248,17 @@ const GeminiPromptFeedback = Schema.StructWithRest(
 )
 type GeminiPromptFeedback = Schema.Schema.Type<typeof GeminiPromptFeedback>
 
+const GeminiError = Schema.StructWithRest(
+  Schema.Struct({
+    code: lenient(Schema.Number),
+    message: lenient(Schema.String),
+    status: lenient(Schema.String),
+  }),
+  [Schema.Record(Schema.String, Schema.Unknown)],
+)
+
 const GeminiEvent = Schema.Struct({
-  error: Schema.optional(Schema.Unknown),
+  error: Schema.optional(GeminiError),
   candidates: optionalNull(Schema.Array(GeminiCandidate)),
   promptFeedback: optionalNull(GeminiPromptFeedback),
   usageMetadata: optionalNull(GeminiUsage),
@@ -637,27 +646,18 @@ const step = (state: ParserState, event: GeminiEvent | NextRequest) => {
       },
       [],
     ] as const)
-  if (ProviderShared.isRecord(event.error)) {
+  if (event.error) {
     const body = ProviderShared.encodeJson(event)
     return Effect.fail(
       new AIError({
         reason: classifyProviderFailure({
-          message:
-            typeof event.error.message === "string" && event.error.message.length > 0
-              ? event.error.message
-              : typeof event.error.status === "string" && event.error.status.length > 0
-                ? event.error.status
-                : "Gemini provider error",
-          status: typeof event.error.code === "number" ? event.error.code : undefined,
+          message: event.error.message || event.error.status || "Gemini provider error",
+          status: event.error.code,
           rawBody: body,
         }),
       }),
     )
   }
-  if ("error" in event)
-    return Effect.fail(
-      ProviderShared.eventError(state.route, `Invalid ${state.route} stream event`, ProviderShared.encodeJson(event)),
-    )
   const candidate = event.candidates?.[0]
   const nextState = {
     ...state,

packages/ai/src/protocols/google-interactions.ts

@@ -132,13 +132,22 @@ const OutputStep = Schema.Struct({
   content: Schema.optional(Schema.Array(Schema.Struct({ type: Schema.String, text: Schema.optional(Schema.String) }))),
 })
 type OutputStep = typeof OutputStep.Type
+const KNOWN_DELTA_TYPES = new Set(["text", "arguments_delta", "thought_signature", "thought_summary"])
+const UnknownDelta = Schema.Struct({
+  type: Schema.String.check(Schema.makeFilter((type) => !KNOWN_DELTA_TYPES.has(type))),
+}).pipe(
+  Schema.decodeTo(Schema.Struct({ type: Schema.Literal("unknown"), rawType: Schema.String }), {
+    decode: SchemaGetter.transform((delta) => ({ type: "unknown" as const, rawType: delta.type })),
+    encode: SchemaGetter.forbidden(() => "Unknown Google Interactions deltas are decode-only"),
+  }),
+)
 const Delta = Schema.Union([
   Schema.Struct({ type: Schema.Literal("text"), text: Schema.String }),
   Schema.Struct({ type: Schema.Literal("arguments_delta"), arguments: Schema.String }),
   Schema.Struct({ type: Schema.Literal("thought_signature"), signature: Schema.String }),
   Schema.Struct({ type: Schema.Literal("thought_summary"), content: Text }),
   // Unknown output modalities must fail explicitly rather than disappearing from a successful response.
-  Schema.Struct({ type: Schema.String }),
+  UnknownDelta,
 ])
 const Interaction = Schema.StructWithRest(
   Schema.Struct({
@@ -430,17 +439,17 @@ const onDelta = Effect.fnUntraced(function* (
   if (!step)
     return yield* ProviderShared.eventError(ADAPTER, "Interactions delta without step.start", encodeJson(delta))
   const events: LLMEvent[] = []
-  if (delta.type === "text" && "text" in delta && step.type === "model_output")
+  if (delta.type === "text" && step.type === "model_output")
     return [
       { ...state, lifecycle: Lifecycle.textDelta(state.lifecycle, events, String(index), delta.text) },
       events,
     ] satisfies StepResult
-  if (delta.type === "thought_summary" && "content" in delta && step.type === "thought")
+  if (delta.type === "thought_summary" && step.type === "thought")
     return [
       { ...state, lifecycle: Lifecycle.reasoningDelta(state.lifecycle, events, String(index), delta.content.text) },
       events,
     ] satisfies StepResult
-  if (delta.type === "thought_signature" && "signature" in delta) {
+  if (delta.type === "thought_signature") {
     const next = { ...step, signature: delta.signature }
     const tool = state.tools[index]
     return [
@@ -452,7 +461,7 @@ const onDelta = Effect.fnUntraced(function* (
       events,
     ] satisfies StepResult
   }
-  if (delta.type === "arguments_delta" && "arguments" in delta && step.type === "function_call") {
+  if (delta.type === "arguments_delta" && step.type === "function_call") {
     const result = ToolStream.appendExisting(
       ADAPTER,
       state.tools,
@@ -463,7 +472,11 @@ const onDelta = Effect.fnUntraced(function* (
     if (ToolStream.isError(result)) return yield* result
     return [{ ...state, tools: result.tools }, result.events] satisfies StepResult
   }
-  return yield* ProviderShared.eventError(ADAPTER, `Unsupported Interactions delta: ${delta.type}`, encodeJson(delta))
+  return yield* ProviderShared.eventError(
+    ADAPTER,
+    `Unsupported Interactions delta: ${delta.type === "unknown" ? delta.rawType : delta.type}`,
+    encodeJson(delta),
+  )
 })
 
 const onStop = Effect.fnUntraced(function* (state: ParserState, index: number) {

packages/ai/src/protocols/xai-images.ts

@@ -5,7 +5,7 @@ import { Media } from "../media.js"
 import { MediaProtocol } from "../route/media-protocol.js"
 import { MediaRoute } from "../route/media.js"
 import { mergeJsonRecords, type OpenString } from "../schema/index.js"
-import { ProviderShared, optionalNull } from "./shared.js"
+import { JsonObject, ProviderShared, lenient, optionalNull } from "./shared.js"
 import { MediaInput } from "./utils/media-input.js"
 
 const route = MediaProtocol.identity({ id: "xai-images", name: "xAI Images", provider: "xai" })
@@ -39,7 +39,7 @@ const XAIImageResponse = Schema.Struct({
       mime_type: optionalNull(Schema.String),
     }),
   ),
-  usage: Schema.optional(Schema.Unknown),
+  usage: lenient(JsonObject),
 })
 
 // ---------------------------------------------------------------------------
@@ -100,12 +100,11 @@ const decodeResponse = Effect.fn("XAIImages.decodeResponse")(function* (
     }),
   )
   if (images.length === 0) return yield* output.invalid(`${route.name} returned no images`)
-  const usage = ProviderShared.isRecord(decoded.usage) ? decoded.usage : undefined
   // xAI reports a USD cost (`cost_in_usd_ticks`) rather than tokens, seconds, or credits; the raw record stays in
   // provider metadata.
   return new ImageResponse({
     images,
-    providerMetadata: usage === undefined ? undefined : { xai: { usage } },
+    providerMetadata: decoded.usage === undefined ? undefined : { xai: { usage: decoded.usage } },
   })
 })