diff --git a/packages/ai/src/protocols/anthropic-messages.ts b/packages/ai/src/protocols/anthropic-messages.ts index 2196e8033a..14c6bc1b16 100644 --- a/packages/ai/src/protocols/anthropic-messages.ts +++ b/packages/ai/src/protocols/anthropic-messages.ts @@ -33,6 +33,18 @@ export const DEFAULT_BASE_URL = "https://api.anthropic.com/v1" export const PATH = "/messages" export const DEFAULT_MAX_TOKENS = 32_000 +const SSE_EVENTS = new Set([ + "message", + "message_start", + "message_delta", + "message_stop", + "content_block_start", + "content_block_delta", + "content_block_stop", + "error", +]) +export const framing = Framing.sseEvents(SSE_EVENTS) + export type ThinkingInput = | { readonly type: "adaptive" @@ -1039,7 +1051,7 @@ export const route = Route.make({ protocol, endpoint: Endpoint.path(PATH, { baseURL: DEFAULT_BASE_URL }), auth: Auth.none, - framing: Framing.sse, + framing, headers: () => ({ "anthropic-version": "2023-06-01" }), }) diff --git a/packages/ai/src/protocols/shared.ts b/packages/ai/src/protocols/shared.ts index 97d1e4b71d..12c7cf7031 100644 --- a/packages/ai/src/protocols/shared.ts +++ b/packages/ai/src/protocols/shared.ts @@ -197,19 +197,28 @@ export const errorText = (error: unknown) => { /** * `framing` step for Server-Sent Events. Decodes UTF-8, runs the SSE channel - * decoder, and drops empty / `[DONE]` keep-alive events so the protocol event - * schema sees one JSON string per element. The SSE channel emits a + * decoder, optionally filters named events, and drops empty / `[DONE]` + * keep-alive events so the protocol event schema sees one JSON string per + * element. The SSE channel emits a * `Retry` control event on its error channel; we drop it here (we don't * implement client-driven retries). Decoder failures become provider output * errors so the public error channel stays `AIError`. */ -export const sseFraming = (bytes: Stream.Stream): Stream.Stream => +export const sseFraming = ( + bytes: Stream.Stream, + events?: ReadonlySet, +): Stream.Stream => bytes.pipe( Stream.decodeText(), Stream.pipeThroughChannel(Sse.decode()), Stream.catchTag("Retry", () => Stream.empty), Stream.catchTag("SseError", (error) => Stream.fail(eventError("sse", error.message))), - Stream.filter((event) => event.data.length > 0 && event.data !== "[DONE]"), + Stream.filter( + (event) => + (events === undefined || events.has(event.event)) && + event.data.length > 0 && + (event.data !== "[DONE]" || (events !== undefined && event.event !== "message")), + ), Stream.map((event) => event.data), ) diff --git a/packages/ai/src/providers/google-vertex-messages.ts b/packages/ai/src/providers/google-vertex-messages.ts index bbfe2d241b..f21742b618 100644 --- a/packages/ai/src/providers/google-vertex-messages.ts +++ b/packages/ai/src/providers/google-vertex-messages.ts @@ -4,7 +4,6 @@ import { AnthropicMessages } from "../protocols/anthropic-messages.js" import { Auth } from "../route/auth.js" import { Route, type RouteDefaultsInput } from "../route/client.js" import { Endpoint } from "../route/endpoint.js" -import { Framing } from "../route/framing.js" import { Protocol } from "../route/protocol.js" import { ProviderID, type ModelID } from "../schema/index.js" import { GoogleVertexShared } from "./google-vertex-shared.js" @@ -57,7 +56,7 @@ const route = Route.make({ }), endpoint: Endpoint.path(({ request }) => `/${request.model.id}:streamRawPredict`), auth: Auth.none, - framing: Framing.sse, + framing: AnthropicMessages.framing, }) export const routes = [route] diff --git a/packages/ai/src/route/framing.ts b/packages/ai/src/route/framing.ts index 8c724aca73..c7b4968151 100644 --- a/packages/ai/src/route/framing.ts +++ b/packages/ai/src/route/framing.ts @@ -24,4 +24,10 @@ export interface Definition { /** Server-Sent Events framing. Used by every JSON-streaming HTTP provider. */ export const sse: Definition = { id: "sse", frame: ProviderShared.sseFraming } +/** SSE framing restricted to protocol-recognized event names. */ +export const sseEvents = (events: ReadonlySet): Definition => ({ + id: "sse", + frame: (bytes) => ProviderShared.sseFraming(bytes, events), +}) + export * as Framing from "./framing.js" diff --git a/packages/ai/test/lib/sse.ts b/packages/ai/test/lib/sse.ts index 80b275d296..ce0474efa8 100644 --- a/packages/ai/test/lib/sse.ts +++ b/packages/ai/test/lib/sse.ts @@ -10,6 +10,9 @@ export const sseEvents = (...chunks: ReadonlyArray): string => const formatChunk = (chunk: unknown) => `data: ${typeof chunk === "string" ? chunk : JSON.stringify(chunk)}\n\n` +export const sseNamedEvent = (event: string, data: unknown): string => + `event: ${event}\ndata: ${typeof data === "string" ? data : JSON.stringify(data)}` + /** * Build an SSE body from already-serialized strings (used when the chunk shape * itself is part of what's being tested, e.g. malformed chunks). diff --git a/packages/ai/test/provider/anthropic-messages.test.ts b/packages/ai/test/provider/anthropic-messages.test.ts index ae7f80e4ba..9a03840225 100644 --- a/packages/ai/test/provider/anthropic-messages.test.ts +++ b/packages/ai/test/provider/anthropic-messages.test.ts @@ -8,7 +8,7 @@ import * as AnthropicMessages from "../../src/protocols/anthropic-messages.js" import { continuationRequest, nativeAnthropicMessagesContinuation } from "../continuation-scenarios.js" import { it } from "../lib/effect.js" import { dynamicResponse, fixedResponse } from "../lib/http.js" -import { sseEvents } from "../lib/sse.js" +import { sseEvents, sseNamedEvent, sseRaw } from "../lib/sse.js" const model = AnthropicMessages.route .with({ endpoint: { baseURL: "https://api.anthropic.test/v1/" }, auth: Auth.header("x-api-key", "test") }) @@ -640,6 +640,59 @@ describe("Anthropic Messages route", () => { }), ) + it.effect("ignores unknown named SSE events", () => + Effect.gen(function* () { + const response = yield* LLMClient.generate(request).pipe( + Effect.provide( + fixedResponse( + sseRaw( + sseNamedEvent("message_start", { + type: "message_start", + message: { usage: { input_tokens: 5 } }, + }), + sseNamedEvent("proxy.stats", "not json"), + sseNamedEvent("content_block_start", { + type: "content_block_start", + index: 0, + content_block: { type: "text", text: "" }, + }), + sseNamedEvent("content_block_delta", { + type: "content_block_delta", + index: 0, + delta: { type: "text_delta", text: "Hello" }, + }), + sseNamedEvent("content_block_stop", { type: "content_block_stop", index: 0 }), + sseNamedEvent("message_delta", { + type: "message_delta", + delta: { stop_reason: "end_turn" }, + usage: { output_tokens: 1 }, + }), + sseNamedEvent("message_stop", { type: "message_stop" }), + sseNamedEvent("proxy.done", "still not json"), + ), + ), + ), + ) + + expect(response.message.content).toEqual([{ type: "text", text: "Hello" }]) + expect(response.finishReason).toEqual({ normalized: "stop", raw: "end_turn" }) + }), + ) + + it.effect("rejects malformed recognized SSE events", () => + Effect.gen(function* () { + const error = yield* LLMClient.generate(request).pipe( + Effect.provide(fixedResponse(sseRaw(sseNamedEvent("message_start", "[DONE]")))), + Effect.flip, + ) + + expect(error.reason).toMatchObject({ + _tag: "InvalidProviderOutput", + message: "Invalid anthropic/anthropic-messages stream event", + }) + }), + ) + it.effect("maps nullable input tokens and preserves unknown Anthropic usage fields", () => Effect.gen(function* () { const response = yield* LLMClient.generate(request).pipe(