feat(ai): continue split Gemini responses (Gemini transport) (#55056)

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<typeof GeminiBody>
@@ -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<GeminiUsage>
   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<string, unknown> = 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<GeminiUsage>, 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<LLMEvent> => {
   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<LLMEvent> => {
         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<LLMEvent> => {
   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<string>
+  readonly input: TransportPrepareInput<GeminiBody>
+}
+
+const sse = HttpTransport.sseJson.with<GeminiBody>()
+
+/**
+ * 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<GeminiBody, Prepared, unknown> = {
+  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<object, string>()
+      const respond = (execution: TransportExecution<string>): Stream.Stream<unknown, AIError> =>
+        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"

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]

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<string | { readonly dropped: string } | ReadableStream<Uint8Array>>) => {
+    const bodies: Array<Record<string, unknown>> = []
+    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<unknown>) =>
+    values.map((value) => sseRaw(`data: ${JSON.stringify(value)}`)).join("")
+
+  const text = (value: string, extra: Record<string, unknown> = {}) => ({
+    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<void>()
+      let sent = false
+      const server = scripted([
+        new ReadableStream({
+          pull(controller) {
+            if (sent) return new Promise<void>(() => 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,
+            },
+          },
+        }),
+      )
+    }),
+  )
+})

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<string, unknown> }> = []
+      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({