diff --git a/packages/ai/src/protocols/gemini.ts b/packages/ai/src/protocols/gemini.ts index adc25fe5c87..5815602dcf4 100644 --- a/packages/ai/src/protocols/gemini.ts +++ b/packages/ai/src/protocols/gemini.ts @@ -1,10 +1,15 @@ -import { Effect, Option, Schema } from "effect" +import { Effect, Option, Schema, Stream } from "effect" import { Tool } from "@opencode/schema/tool" import { Route } from "../route/client.js" import { Auth } from "../route/auth.js" import { Endpoint } from "../route/endpoint.js" -import { Framing } from "../route/framing.js" import { Protocol } from "../route/protocol.js" +import { + HttpTransport, + type Transport, + type TransportExecution, + type TransportPrepareInput, +} from "../route/transport/index.js" import { AIError, LLMEvent, @@ -25,6 +30,7 @@ import { GeminiGenerateContent } from "./utils/gemini-generate-content.js" import { Lifecycle } from "./utils/lifecycle.js" const ADAPTER = "gemini" +const CONTINUATION = "CONTINUATION" // Google documents this sentinel for replaying Gemini 3 function calls after their original signature was lost. const SKIP_THOUGHT_SIGNATURE_VALIDATOR = "skip_thought_signature_validator" // Gemini 2.5 rejects a budget under the model's minimum: 512 on Flash-Lite, the highest, and 128 on Pro. @@ -212,6 +218,7 @@ const GeminiBodyFields = { tools: optionalArray(GeminiTool), toolConfig: Schema.optional(GeminiToolConfig), generationConfig: Schema.optional(GeminiGenerationConfig), + continuationToken: Schema.optional(Schema.String), } const GeminiBody = Schema.Struct(GeminiBodyFields) export type GeminiBody = Schema.Schema.Type @@ -254,7 +261,9 @@ interface ParserState { readonly finishReason?: string readonly hasToolCalls: boolean readonly promptFeedback?: GeminiPromptFeedback - readonly usage?: Usage + readonly usage?: GeminiUsage + /** Usage of the earlier requests of a response Gemini continued across requests. */ + readonly continuedUsage?: ReadonlyArray readonly lifecycle: Lifecycle.State readonly reasoningSignature?: string readonly textSignature?: string @@ -497,7 +506,11 @@ const fromRequest = Effect.fn("Gemini.fromRequest")(function* (request: LLMReque // `cachedContentTokenCount` subset. `candidatesTokenCount` is *exclusive* // of `thoughtsTokenCount` — visible-only, not a total — so we sum the two // to produce the inclusive `outputTokens` the rest of the contract expects. -const mapUsage = (usage: GeminiUsage | undefined, metadataKey: string) => { +const mapUsage = ( + usage: GeminiUsage | undefined, + metadataKey: string, + metadata: Record = usage ?? {}, +) => { if (!usage) return undefined // Explicit provider nulls decode as `null`; normalize to `undefined` so the // token arithmetic below treats them like absent counts. @@ -518,10 +531,35 @@ const mapUsage = (usage: GeminiUsage | undefined, metadataKey: string) => { cacheReadInputTokens: cached, reasoningTokens: thoughts, totalTokens: ProviderShared.totalTokens(promptTokens, outputTokens, usage.totalTokenCount ?? undefined), - providerMetadata: providerMetadata(metadataKey, usage), + providerMetadata: providerMetadata(metadataKey, metadata), }) } +// Each continued request re-bills the output so far as input. Reporting the first request's input with all output keeps +// input plus output equal to the response's final context size, under-reporting billed input; the summed billed counts +// remain in provider metadata. +const mapContinuedUsage = (requests: ReadonlyArray, metadataKey: string) => { + const sum = (field: keyof GeminiUsage) => + ProviderShared.sumTokens(...requests.map((usage) => usage[field] ?? undefined)) + return mapUsage( + { + promptTokenCount: requests[0].promptTokenCount, + cachedContentTokenCount: requests[0].cachedContentTokenCount, + candidatesTokenCount: sum("candidatesTokenCount"), + thoughtsTokenCount: sum("thoughtsTokenCount"), + }, + metadataKey, + { + promptTokenCount: sum("promptTokenCount"), + cachedContentTokenCount: sum("cachedContentTokenCount"), + candidatesTokenCount: sum("candidatesTokenCount"), + thoughtsTokenCount: sum("thoughtsTokenCount"), + totalTokenCount: sum("totalTokenCount"), + requests: requests.length, + }, + ) +} + const mapFinishReason = (finishReason: string | undefined, hasToolCalls: boolean): FinishReason => { if (finishReason === undefined) return hasToolCalls ? "tool-calls" : "unknown" if (finishReason === "STOP") return hasToolCalls ? "tool-calls" : "stop" @@ -545,7 +583,7 @@ const finish = (state: ParserState): ReadonlyArray => { const promptBlockReason = state.finishReason === undefined ? (state.promptFeedback?.blockReason ?? undefined) : undefined const finishReason = state.finishReason ?? promptBlockReason - if (finishReason === undefined && state.usage === undefined) return [] + if (finishReason === undefined && state.usage === undefined && state.continuedUsage === undefined) return [] const events: LLMEvent[] = [] let lifecycle = state.lifecycle @@ -573,7 +611,10 @@ const finish = (state: ParserState): ReadonlyArray => { promptBlockReason === undefined ? mapFinishReason(finishReason, state.hasToolCalls) : "content-filter", raw: finishReason, }, - usage: state.usage, + usage: + state.continuedUsage === undefined + ? mapUsage(state.usage, state.providerMetadataKey) + : mapContinuedUsage([...state.continuedUsage, state.usage ?? {}], state.providerMetadataKey), providerMetadata: state.promptFeedback === undefined ? undefined @@ -582,7 +623,19 @@ const finish = (state: ParserState): ReadonlyArray => { return events } -const step = (state: ParserState, event: GeminiEvent) => { +const step = (state: ParserState, event: GeminiEvent | NextRequest) => { + // Usage is cumulative within one request, so the next request of the same response counts from zero. + if (event instanceof NextRequest) + return Effect.succeed([ + { + ...state, + finishReason: undefined, + promptFeedback: undefined, + usage: undefined, + continuedUsage: [...(state.continuedUsage ?? []), state.usage ?? {}], + }, + [], + ] as const) if (ProviderShared.isRecord(event.error)) { const body = ProviderShared.encodeJson(event) return Effect.fail( @@ -604,14 +657,13 @@ const step = (state: ParserState, event: GeminiEvent) => { return Effect.fail( ProviderShared.eventError(state.route, `Invalid ${state.route} stream event`, ProviderShared.encodeJson(event)), ) + const candidate = event.candidates?.[0] const nextState = { ...state, promptFeedback: event.promptFeedback ?? state.promptFeedback, - usage: event.usageMetadata - ? (mapUsage(event.usageMetadata, state.providerMetadataKey) ?? state.usage) - : state.usage, + usage: event.usageMetadata ?? state.usage, + finishReason: candidate?.finishReason ?? state.finishReason, } - const candidate = event.candidates?.[0] if (candidate?.finishReason && mapFinishReason(candidate.finishReason, state.hasToolCalls) === "error") return Effect.fail( ProviderShared.eventError( @@ -620,11 +672,7 @@ const step = (state: ParserState, event: GeminiEvent) => { ProviderShared.encodeJson(event), ), ) - if (!candidate?.content) - return Effect.succeed([ - { ...nextState, finishReason: candidate?.finishReason ?? nextState.finishReason }, - [], - ] as const) + if (!candidate?.content) return Effect.succeed([nextState, []] as const) const events: LLMEvent[] = [] let hasToolCalls = nextState.hasToolCalls @@ -784,12 +832,104 @@ const step = (state: ParserState, event: GeminiEvent) => { nextReasoningId, nextTextId, seenCallIds, - finishReason: candidate.finishReason ?? nextState.finishReason, }, events, ] as const) } +// ============================================================================= +// Transport +// ============================================================================= +const GeminiCheckpoint = Schema.Struct({ + candidates: optionalNull( + Schema.Array( + Schema.Struct({ + finishReason: optionalNull(Schema.String), + continuationToken: optionalNull(Schema.String), + }), + ), + ), +}) +const decodeCheckpoint = Schema.decodeUnknownOption(GeminiCheckpoint) +/** Marks where the transport starts another request for the same response; provider JSON can never match it. */ +class NextRequest {} + +export interface Prepared { + readonly http: HttpTransport.HttpPrepared + readonly input: TransportPrepareInput +} + +const sse = HttpTransport.sseJson.with() + +/** + * SSE transport that continues responses Gemini splits across requests. Chunks carry `continuationToken` checkpoints, + * and a request that reaches its output cap finishes with `CONTINUATION`. The transport resends the original body with + * the latest token, which also resumes a request that ended or dropped before finishing. + */ +export const transport: Transport = { + id: sse.id, + prepare: (input) => sse.prepare(input).pipe(Effect.map((http) => ({ http, input }))), + execute: (prepared, request, runtime) => + Effect.gen(function* () { + const route = request.model.route.id + const label = `${request.model.provider}/${route}` + // Frames are parsed once here, so errors recover each frame's raw text for their body. + const raw = new WeakMap() + const respond = (execution: TransportExecution): Stream.Stream => + Stream.suspend(() => { + let token: string | undefined + let finishReason: string | undefined + let dropped: AIError | undefined + return execution.frames.pipe( + Stream.mapEffect((frame) => + ProviderShared.parseJson(route, frame, `Invalid ${label} stream event`).pipe( + Effect.map((event) => { + if (typeof event === "object" && event !== null) raw.set(event, frame) + const candidate = Option.getOrUndefined(decodeCheckpoint(event))?.candidates?.[0] + token = candidate?.continuationToken ?? (candidate?.finishReason ? undefined : token) + finishReason = candidate?.finishReason ?? finishReason + return event + }), + ), + ), + Stream.catchIf( + // A connection failure may resume; a timeout keeps its meaning and ends the response. + (error) => error.reason._tag === "Transport" && error.reason.code !== "Timeout", + (error) => { + dropped = error + return Stream.empty + }, + ), + Stream.concat( + Stream.suspend(() => { + if (token !== undefined && (finishReason === undefined || finishReason === CONTINUATION)) + return Stream.concat(Stream.succeed(new NextRequest()), next(token)) + if (dropped) return Stream.fail(dropped) + if (finishReason === CONTINUATION) + return Stream.fail(ProviderShared.eventError(label, "Gemini continued the response without a token")) + return Stream.empty + }), + ), + ) + }) + const next = (continuationToken: string) => + Stream.unwrap( + sse + .prepare({ ...prepared.input, body: { ...prepared.input.body, continuationToken } }) + .pipe( + Effect.flatMap((continued) => sse.execute(continued, request, runtime)), + Effect.map(respond), + ), + ) + const first = yield* sse.execute(prepared.http, request, runtime) + return { + frames: respond(first), + http: first.http, + body: (frame) => (typeof frame === "object" && frame !== null ? raw.get(frame) : undefined), + } + }), +} + // ============================================================================= // Protocol And Gemini Route // ============================================================================= @@ -806,8 +946,8 @@ export const protocol = Protocol.make({ // Gemini's schema rules are this API's default, including for tuned endpoints whose IDs do not name Gemini. sanitizer: "gemini", stream: { - event: Protocol.jsonEvent(GeminiEvent), - initial: (request) => ({ + event: Schema.Union([Schema.instanceOf(NextRequest), GeminiEvent]), + initial: (request): ParserState => ({ route: `${request.model.provider}/${request.model.route.id}`, providerMetadataKey: request.model.route.providerMetadataKey ?? String(request.model.provider), hasToolCalls: false, @@ -830,7 +970,7 @@ export const route = Route.make({ baseURL: DEFAULT_BASE_URL, }), auth: Auth.none, - framing: Framing.sse, + transport, }) export * as Gemini from "./gemini.js" diff --git a/packages/ai/src/providers/google-vertex.ts b/packages/ai/src/providers/google-vertex.ts index e5d5d53518d..085d93c575f 100644 --- a/packages/ai/src/providers/google-vertex.ts +++ b/packages/ai/src/providers/google-vertex.ts @@ -5,7 +5,6 @@ import { ProviderShared } from "../protocols/shared.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 { ProviderConfigurationError, ProviderID, type LLMRequest, type ModelID } from "../schema/index.js" import { GoogleVertexShared } from "./google-vertex-shared.js" @@ -69,7 +68,7 @@ const route = Route.make({ const serviceTier = request.providerOptions?.serviceTier return typeof serviceTier === "string" ? { "x-vertex-ai-llm-shared-request-type": serviceTier } : {} }, - framing: Framing.sse, + transport: Gemini.transport, }) export const routes = [route] diff --git a/packages/ai/test/provider/gemini.test.ts b/packages/ai/test/provider/gemini.test.ts index 26cd35b225e..032782822f7 100644 --- a/packages/ai/test/provider/gemini.test.ts +++ b/packages/ai/test/provider/gemini.test.ts @@ -1,11 +1,12 @@ import { describe, expect } from "bun:test" -import { Effect } from "effect" -import { LLM, AIError, LLMRequest, Message, ToolCallPart, ToolDefinition, Usage, Media } from "../../src/index.js" +import { Deferred, Effect, Fiber } from "effect" +import * as TestClock from "effect/testing/TestClock" +import { LLM, AIError, LLMEvent, LLMRequest, Message, ToolCallPart, ToolDefinition, Usage, Media } from "../../src/index.js" import { Auth, LLMClient } from "../../src/route.js" import { compileRequest } from "../../src/route/client.js" import * as Gemini from "../../src/protocols/gemini.js" import { it } from "../lib/effect.js" -import { fixedResponse } from "../lib/http.js" +import { dynamicResponse, fixedResponse } from "../lib/http.js" import { sseEvents, sseRaw } from "../lib/sse.js" const model = Gemini.route @@ -1832,3 +1833,316 @@ describe("Gemini route", () => { }), ) }) + +describe("Gemini response continuation", () => { + const continuedRequest = LLM.request({ + model: Gemini.route + .with({ + endpoint: { baseURL: "https://generativelanguage.test/v1beta/" }, + auth: Auth.header("x-goog-api-key", "test"), + }) + .model({ id: "gemini-test" }), + prompt: "Count.", + tools: [ToolDefinition.make({ name: "lookup", description: "Lookup data", inputSchema: { type: "object" } })], + }) + + // Serves one scripted SSE body per request, or a body whose connection drops after it, and records every sent body. + const scripted = (responses: ReadonlyArray>) => { + const bodies: Array> = [] + const layer = dynamicResponse((input) => + Effect.sync(() => { + bodies.push(JSON.parse(input.text)) + const response = responses[bodies.length - 1] + if (typeof response === "string" || response instanceof ReadableStream) + return input.respond(response, { headers: { "content-type": "text/event-stream" } }) + let sent = false + return input.respond( + new ReadableStream({ + pull(controller) { + if (sent) return controller.error(new Error("connection reset")) + sent = true + controller.enqueue(new TextEncoder().encode(response.dropped)) + }, + }), + { headers: { "content-type": "text/event-stream" } }, + ) + }), + ) + return { bodies, layer } + } + const chunks = (...values: ReadonlyArray) => + values.map((value) => sseRaw(`data: ${JSON.stringify(value)}`)).join("") + + const text = (value: string, extra: Record = {}) => ({ + candidates: [{ content: { role: "model", parts: [{ text: value }] }, ...extra }], + }) + const continuation = (token?: string) => ({ + candidates: [ + { content: { role: "model", parts: [{ text: "" }] }, finishReason: "CONTINUATION", continuationToken: token }, + ], + }) + + it.effect("stitches text split across three requests into one block", () => + Effect.gen(function* () { + const server = scripted([ + sseEvents( + { + candidates: [ + { content: { role: "model", parts: [{ text: "plan", thought: true, thoughtSignature: "sig_1" }] } }, + ], + }, + text("THIRTEEN THOUSAND TWO"), + continuation("token-1"), + ), + sseEvents( + text(" HUNDRED SEVEN"), + // The finishing chunk's token supersedes an earlier checkpoint's. + text("TY-TWO", { continuationToken: "checkpoint" }), + continuation("token-2"), + ), + sseEvents(text(" DONE", { finishReason: "STOP" })), + ]) + const response = yield* LLMClient.generate(continuedRequest).pipe(Effect.provide(server.layer)) + + expect(response.text).toBe("THIRTEEN THOUSAND TWO HUNDRED SEVENTY-TWO DONE") + expect(response.reasoning).toBe("plan") + expect(response.events.filter((event) => event.type !== "text-delta").map((event) => event.type)).toEqual([ + "step-start", + "reasoning-start", + "reasoning-delta", + "reasoning-end", + "text-start", + "text-end", + "step-finish", + "finish", + ]) + expect(response.events.find((event) => event.type === "reasoning-end")).toMatchObject({ + providerMetadata: { google: { thoughtSignature: "sig_1" } }, + }) + expect(response.finishReason).toEqual({ normalized: "stop", raw: "STOP" }) + expect(server.bodies).toHaveLength(3) + expect(server.bodies[0]).not.toHaveProperty("continuationToken") + expect(server.bodies[1]).toEqual({ ...server.bodies[0], continuationToken: "token-1" }) + expect(server.bodies[2]).toEqual({ ...server.bodies[0], continuationToken: "token-2" }) + }), + ) + + it.effect("finishes with a signed tool call from a later request", () => + Effect.gen(function* () { + const server = scripted([ + sseEvents(text("Looking"), continuation("token-1")), + sseEvents({ + candidates: [ + { + content: { + role: "model", + parts: [ + { text: " it up." }, + { functionCall: { id: "call_1", name: "lookup", args: {} }, thoughtSignature: "call_sig" }, + ], + }, + finishReason: "STOP", + }, + ], + }), + ]) + const response = yield* LLMClient.generate(continuedRequest).pipe(Effect.provide(server.layer)) + + expect(response.text).toBe("Looking it up.") + expect(response.events.filter((event) => event.type === "text-start")).toHaveLength(1) + expect(response.events.find((event) => event.type === "tool-call")).toMatchObject({ + id: "call_1", + providerMetadata: { google: { thoughtSignature: "call_sig" } }, + }) + expect(response.finishReason).toEqual({ normalized: "tool-calls", raw: "STOP" }) + }), + ) + + it.effect("forgets a block once the resumed request finishes", () => + Effect.gen(function* () { + const server = scripted([ + sseEvents(text("THIRTEEN", { continuationToken: "checkpoint-1" }), { promptFeedback: { blockReason: "OTHER" } }), + sseEvents(text(" THOUSAND", { finishReason: "STOP" })), + ]) + const response = yield* LLMClient.generate(continuedRequest).pipe(Effect.provide(server.layer)) + + expect(response.text).toBe("THIRTEEN THOUSAND") + expect(response.finishReason).toEqual({ normalized: "stop", raw: "STOP" }) + expect(response.events.filter(LLMEvent.is.finish).map((event) => event.providerMetadata)).toEqual([undefined]) + }), + ) + + it.effect("resumes a dropped stream from its latest checkpoint", () => + Effect.gen(function* () { + const server = scripted([ + { dropped: chunks(text("THIRTEEN"), text(" THOUSAND", { continuationToken: "checkpoint-1" })) }, + sseEvents(text(" TWO", { finishReason: "STOP" })), + ]) + const response = yield* LLMClient.generate(continuedRequest).pipe(Effect.provide(server.layer)) + + expect(response.text).toBe("THIRTEEN THOUSAND TWO") + expect(response.finishReason).toEqual({ normalized: "stop", raw: "STOP" }) + expect(server.bodies[1]).toEqual({ ...server.bodies[0], continuationToken: "checkpoint-1" }) + }), + ) + + it.effect("counts the usage of a dropped request once resumed", () => + Effect.gen(function* () { + const server = scripted([ + { + dropped: chunks( + { ...text("THIRTEEN"), usageMetadata: { promptTokenCount: 120, candidatesTokenCount: 5 } }, + { + ...text(" THOUSAND", { continuationToken: "checkpoint-1" }), + usageMetadata: { promptTokenCount: 142, candidatesTokenCount: 32768, totalTokenCount: 32910 }, + }, + ), + }, + sseEvents({ + ...text(" TWO", { finishReason: "STOP" }), + usageMetadata: { promptTokenCount: 32910, candidatesTokenCount: 100, totalTokenCount: 33010 }, + }), + ]) + const response = yield* LLMClient.generate(continuedRequest).pipe(Effect.provide(server.layer)) + + expect(response.usage).toEqual( + new Usage({ + inputTokens: 142, + outputTokens: 32868, + nonCachedInputTokens: 142, + totalTokens: 33010, + providerMetadata: { + google: { + promptTokenCount: 33052, + candidatesTokenCount: 32868, + totalTokenCount: 65920, + requests: 2, + }, + }, + }), + ) + }), + ) + + it.effect("continues a stream that ends at a checkpoint without finishing", () => + Effect.gen(function* () { + const server = scripted([ + sseEvents(text("THIRTEEN", { continuationToken: "checkpoint-1" })), + sseEvents(text(" THOUSAND", { finishReason: "STOP" })), + ]) + const response = yield* LLMClient.generate(continuedRequest).pipe(Effect.provide(server.layer)) + + expect(response.text).toBe("THIRTEEN THOUSAND") + expect(server.bodies[1]).toEqual({ ...server.bodies[0], continuationToken: "checkpoint-1" }) + }), + ) + + it.effect("fails a stream dropped before any checkpoint", () => + Effect.gen(function* () { + const server = scripted([{ dropped: chunks(text("Partial")) }]) + const error = yield* LLMClient.generate(continuedRequest).pipe(Effect.provide(server.layer), Effect.flip) + + expect(error.reason._tag).toBe("Transport") + expect(server.bodies).toHaveLength(1) + }), + ) + + it.effect("does not resume a response cut off by its own timeout", () => + Effect.gen(function* () { + const stalled = yield* Deferred.make() + let sent = false + const server = scripted([ + new ReadableStream({ + pull(controller) { + if (sent) return new Promise(() => Deferred.doneUnsafe(stalled, Effect.void)) + sent = true + controller.enqueue(new TextEncoder().encode(chunks(text("THIRTEEN", { continuationToken: "checkpoint-1" })))) + }, + }), + ]) + const fiber = yield* LLMClient.generate( + LLMRequest.update(continuedRequest, { http: { timeout: 1_000 } }), + ).pipe(Effect.provide(server.layer), Effect.flip, Effect.forkChild({ startImmediately: true })) + yield* Deferred.await(stalled) + yield* TestClock.adjust("1 second") + const error = yield* Fiber.join(fiber) + + expect(error.reason).toMatchObject({ _tag: "Transport", code: "Timeout" }) + expect(server.bodies).toHaveLength(1) + }), + ) + + it.effect("fails a continuation without a token", () => + Effect.gen(function* () { + const server = scripted([sseEvents(text("Partial"), continuation())]) + const error = yield* LLMClient.generate(continuedRequest).pipe(Effect.provide(server.layer), Effect.flip) + + expect(error).toMatchObject({ + reason: { _tag: "InvalidProviderOutput", http: { status: 200 } }, + message: "Gemini continued the response without a token", + }) + expect(server.bodies).toHaveLength(1) + }), + ) + + it.effect("does not continue a length finish", () => + Effect.gen(function* () { + const server = scripted([ + sseEvents(text("Partial", { finishReason: "MAX_TOKENS", continuationToken: "token-1" })), + ]) + const response = yield* LLMClient.generate(continuedRequest).pipe(Effect.provide(server.layer)) + + expect(response.finishReason).toEqual({ normalized: "length", raw: "MAX_TOKENS" }) + expect(server.bodies).toHaveLength(1) + }), + ) + + it.effect("reports the first request's input with all output", () => + Effect.gen(function* () { + const usage = [ + { promptTokenCount: 142, candidatesTokenCount: 32768, totalTokenCount: 32910 }, + { promptTokenCount: 32910, candidatesTokenCount: 32768, totalTokenCount: 65678 }, + { + promptTokenCount: 65678, + cachedContentTokenCount: 32911, + candidatesTokenCount: 32768, + totalTokenCount: 98446, + }, + { + promptTokenCount: 98446, + cachedContentTokenCount: 65679, + candidatesTokenCount: 7257, + totalTokenCount: 105703, + }, + ] + const server = scripted([ + sseEvents( + { ...text("a"), usageMetadata: { promptTokenCount: 120, candidatesTokenCount: 10, totalTokenCount: 130 } }, + { ...continuation("token-1"), usageMetadata: usage[0] }, + ), + sseEvents({ ...continuation("token-2"), usageMetadata: usage[1] }), + sseEvents({ ...continuation("token-3"), usageMetadata: usage[2] }), + sseEvents({ ...text("b", { finishReason: "STOP" }), usageMetadata: usage[3] }), + ]) + const response = yield* LLMClient.generate(continuedRequest).pipe(Effect.provide(server.layer)) + + expect(response.usage).toEqual( + new Usage({ + inputTokens: 142, + outputTokens: 105561, + nonCachedInputTokens: 142, + totalTokens: 105703, + providerMetadata: { + google: { + promptTokenCount: 197176, + cachedContentTokenCount: 98590, + candidatesTokenCount: 105561, + totalTokenCount: 302737, + requests: 4, + }, + }, + }), + ) + }), + ) +}) diff --git a/packages/ai/test/provider/google-vertex.test.ts b/packages/ai/test/provider/google-vertex.test.ts index 2f3da7b3a8a..227a545cdf6 100644 --- a/packages/ai/test/provider/google-vertex.test.ts +++ b/packages/ai/test/provider/google-vertex.test.ts @@ -539,6 +539,50 @@ describe("Google Vertex providers", () => { }), ) + it.effect("continues split Gemini responses with the Vertex request body", () => + Effect.gen(function* () { + const sent: Array<{ readonly authorization: string | null; readonly body: Record }> = [] + const responses = [ + sseEvents({ + candidates: [ + { + content: { role: "model", parts: [{ text: "Hel" }] }, + finishReason: "CONTINUATION", + continuationToken: "token-1", + }, + ], + }), + sseEvents({ candidates: [{ content: { role: "model", parts: [{ text: "lo." }] }, finishReason: "STOP" }] }), + ] + const response = yield* LLMClient.generate( + LLM.request({ + model: GoogleVertex.configure({ + accessToken: "vertex-token", + project: "vertex-project", + providerOptions: { labels: { component: "opencode" }, serviceTier: "flex" }, + }).model("gemini-test"), + prompt: "Say hello.", + }), + ).pipe( + Effect.provide( + dynamicResponse((input) => + Effect.sync(() => { + sent.push({ authorization: input.request.headers.authorization ?? null, body: JSON.parse(input.text) }) + return input.respond(responses[sent.length - 1], { headers: { "content-type": "text/event-stream" } }) + }), + ), + ), + ) + + expect(response.text).toBe("Hello.") + expect(response.events.filter((event) => event.type === "text-start")).toHaveLength(1) + expect(sent.map((request) => request.authorization)).toEqual(["Bearer vertex-token", "Bearer vertex-token"]) + expect(sent[0].body).toMatchObject({ labels: { component: "opencode" } }) + expect(sent[0].body).not.toHaveProperty("serviceTier") + expect(sent[1].body).toEqual({ ...sent[0].body, continuationToken: "token-1" }) + }), + ) + test("rejects tuned Gemini models in express mode", () => { expect(() => GoogleVertex.configure({ apiKey: "fixture" }).model("endpoints/1234567890")).toThrow( expect.objectContaining({