From 35a2cf341ef0b681b5b01eb7edfe6bfaa371329d Mon Sep 17 00:00:00 2001 From: Aiden Cline <63023139+rekram1-node@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:40:24 -0500 Subject: [PATCH] refactor(ai): replace loose Schema.Unknown and catch-all union checks in protocols (#57606) --- packages/ai/src/protocols/cohere-chat.ts | 4 +-- packages/ai/src/protocols/gemini.ts | 26 +++++++++---------- .../ai/src/protocols/google-interactions.ts | 25 +++++++++++++----- packages/ai/src/protocols/xai-images.ts | 7 +++-- 4 files changed, 37 insertions(+), 25 deletions(-) diff --git a/packages/ai/src/protocols/cohere-chat.ts b/packages/ai/src/protocols/cohere-chat.ts index 3b1756823c9..876d5e29766 100644 --- a/packages/ai/src/protocols/cohere-chat.ts +++ b/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, }), }), ), diff --git a/packages/ai/src/protocols/gemini.ts b/packages/ai/src/protocols/gemini.ts index 20d4fbfc2d4..c4af8a3b048 100644 --- a/packages/ai/src/protocols/gemini.ts +++ b/packages/ai/src/protocols/gemini.ts @@ -248,8 +248,17 @@ const GeminiPromptFeedback = Schema.StructWithRest( ) type GeminiPromptFeedback = Schema.Schema.Type +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, diff --git a/packages/ai/src/protocols/google-interactions.ts b/packages/ai/src/protocols/google-interactions.ts index 836177fdeb2..c6148646704 100644 --- a/packages/ai/src/protocols/google-interactions.ts +++ b/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) { diff --git a/packages/ai/src/protocols/xai-images.ts b/packages/ai/src/protocols/xai-images.ts index 438b93a701e..8011c0ba323 100644 --- a/packages/ai/src/protocols/xai-images.ts +++ b/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 } }, }) })