Skip to content

[langfun] Resume a Gemini generation that spent its output budget thinking - #782

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_1003633198
Open

copybara-service[bot] wants to merge 1 commit into
mainfrom
test_1003633198

Conversation

@copybara-service

@copybara-service copybara-service Bot commented Oct 21, 2026 •

Copy link
Copy Markdown

[langfun] Resume a Gemini generation that spent its output budget thinking

maxOutputTokens is shared between thinking and the answer. When thinking
consumes all of it, Vertex AI returns finishReason=MAX_TOKENS with an empty
text part -- and a valid continuationToken. _continuation_hop followed that
token only under finishReason=CONTINUATION, so the resumable offer was
discarded and the caller got GeminiEmptyGenerationError: billed for the
thinking, delivered nothing.

  • _continuation_hop follows a continuationToken under MAX_TOKENS as well
    as CONTINUATION.

  • A resume raises the output ceiling instead of re-asserting it.
    maxOutputTokens is a running total across continuation requests, so a
    MAX_TOKENS hop reports it exhausted and sending it again is rejected with
    max_output_tokens (N) has already been exhausted. Measured against the
    service, replaying one starved generation's token under four ceilings:
    unchanged -> 400; doubled, omitted, or above the model maximum -> 200.

  • The raise is bounded, so the number the caller pinned still bounds the call.
    A resume asks for what the generation has already burned plus the answer
    tokens still owed against the pinned cap -- never more -- so the answer
    cannot outgrow the ceiling even though the ceiling itself must rise. At most
    _MAX_PINNED_BUDGET_HOPS = 4 requests may do that, which for a generation
    that spends every hop thinking walks the ceiling 1x -> 2x -> 3x -> 4x and
    stops. The resume also stops as soon as the caller holds every answer token
    they asked for, or the hop returns answer content the service did not count.
    Dropping the ceiling instead was rejected: it hands the rest of the
    generation to the service default, and a 512-token request was measured
    answering with 61-75 KB of text. A request that pins no ceiling is resumed
    byte-identically, since there is nothing to extend.

  • The ceiling is always asserted, never dropped. Measured against the
    service, maxOutputTokens of 65,537, 131,072 and 1,000,000 -- all above
    the model's own 65,536-token output maximum -- are each accepted with HTTP
    200, so there is no value a resume has to omit. Dropping the key above some
    threshold would have handed the rest of the generation to the service
    default, which is the unbounded behaviour the ceiling exists to prevent.

  • A resume that fails for any other reason -- an expired token, a service
    refusal -- falls back to reporting what the generation itself did instead of
    the resume's own error. A failed CONTINUATION resume still propagates,
    since hiding it would turn a real failure into a silent truncation.

  • VertexAIGemini4Argon no longer pins max_tokens=65_536. It is a class
    default nobody asked for, and under it every starved generation on the model
    was unresumable. Unsetting it is not the fix on its own -- a caller pins the
    ceiling itself through ?max_tokens= in the model URI -- which is why the
    bounded extension above lives in the resume.

  • _check_answer_present stops calling every MAX_TOKENS candidate starved.
    candidatesTokenCount > 0 means the budget did reach the answer and the
    truncated hop was billed for tokens it did not return -- a different fault
    with a different remedy from thinking starvation, so the error now says
    which one occurred instead of always reporting the latter.

  • The non-retryable rule is NARROW: only a signature a re-send cannot change
    ends the caller's generation. A blocked prompt, an exhausted MAX_TOKENS
    budget, a STOP billed for neither answer nor thinking tokens, and a
    candidate carrying no content at all are reported as
    GeminiEmptyGenerationError. Everything else that renders no answer -- a
    tool-protocol finish reason such as MALFORMED_FUNCTION_CALL, a
    candidate-level content-policy flag such as SAFETY or RECITATION, a
    STOP over non-zero token counts, an absent or unrecognised finish
    reason, an unexplained empty candidates list -- is left on the retryable
    path, because that is the path that recovers it: those arrive
    mid-generation on requests that declared no tools at all, so there is
    nothing in the request to correct and the next sample succeeds. Reporting
    them instead converts a transient fault into a fatal one.
    _carries_unrendered_content still carves out candidates whose parts hold
    a functionCall or functionResponse: those carry the turn's payload in
    a non-text part, so an empty text rendering is correct output, not a
    starved one.

  • A candidate-level content-policy flag is NOT admitted, and the API
    contract is what settles it. SAFETY, RECITATION, BLOCKLIST,
    PROHIBITED_CONTENT, SPII, LANGUAGE and IMAGE_SAFETY are defined as
    flags raised on the response candidate content -- the tokens that happened
    to be sampled -- so the next sample is exactly what can clear them.
    promptFeedback.blockReason is the prompt-level counterpart: the API
    defines it on the prompt, returns no candidate with it, and names
    rephrasing the prompt as the remedy. Only that one is admitted. Admitting
    the candidate-level flags would end a live episode over an outcome a
    re-send can change, which is the same trade the closed rule lost 28 times
    out of 28, and no measured failure asks for it: across the 48 answer-less
    gemini-4-argon generations this CL was built from, every one reported
    MAX_TOKENS and none reported a content policy.

  • _merge_continuation_hops keeps a chain that rendered no answer labelled
    MAX_TOKENS instead of taking the last hop's reason. A hop that exhausts
    the budget followed by one that adds nothing and reports STOP merged into
    ('STOP', summed count) -- indistinguishable from a generation that ran
    and was lost, so it was re-sent max_attempts times, every re-send
    exhausting the same budget the same way. Closing that laundering where it
    happens is what lets the classifier above stay narrow. A chain that did
    render an answer keeps the last hop's reason, since it is the true one.

  • Each answer-less signature carries the remedy that fits it, instead of one
    starvation tail appended to all of them. Lowering the thinking budget is
    advice only for a budget that thinking actually consumed; it is wrong for a
    budget that reached the answer, and irrelevant when the model stopped
    having generated nothing at all.

  • GeminiEmptyGenerationError stays non-retryable for the signatures it
    admits. Every rescue that can work is spent inside the one request --
    following the continuationToken, then extending an explicit ceiling to 4x
    -- and each admitted signature reproduces on a re-send. Retrying would buy
    max_attempts full regenerations of the same outcome, each billed in full,
    so the message names that cost instead of incurring it.

  • _merge_continuation_hops keeps bare {'text': ''} parts rather than
    filtering them. Message.from_chunks writes no separator after an empty
    chunk and strips the fused text, so such a part is already inert in every
    position; the filter, and the fallback that restored what it had removed,
    were both dead weight. A test pins the property they relied on.

@copybara-service copybara-service Bot changed the title [langfun] Resume a Gemini generation that spent its decode budget thinking [langfun] Resume a Gemini generation that spent its output budget thinking Oct 22, 2026
@copybara-service
copybara-service Bot force-pushed the test_1003633198 branch 5 times, most recently from d8adae3 to da63c49 Compare October 24, 2026 11:15
…nking

`maxOutputTokens` is shared between thinking and the answer. When thinking
consumes all of it, Vertex AI returns `finishReason=MAX_TOKENS` with an empty
text part -- and a valid `continuationToken`. `_continuation_hop` followed that
token only under `finishReason=CONTINUATION`, so the resumable offer was
discarded and the caller got `GeminiEmptyGenerationError`: billed for the
thinking, delivered nothing.

* `_continuation_hop` follows a `continuationToken` under `MAX_TOKENS` as well
  as `CONTINUATION`.

* A resume raises the output ceiling instead of re-asserting it.
  `maxOutputTokens` is a running total across continuation requests, so a
  `MAX_TOKENS` hop reports it exhausted and sending it again is rejected with
  `max_output_tokens (N) has already been exhausted`. Measured against the
  service, replaying one starved generation's token under four ceilings:
  unchanged -> 400; doubled, omitted, or above the model maximum -> 200.

* The raise is bounded, so the number the caller pinned still bounds the call.
  A resume asks for what the generation has already burned plus the answer
  tokens still owed against the pinned cap -- never more -- so the answer
  cannot outgrow the ceiling even though the ceiling itself must rise. At most
  `_MAX_PINNED_BUDGET_HOPS = 4` requests may do that, which for a generation
  that spends every hop thinking walks the ceiling `1x -> 2x -> 3x -> 4x` and
  stops. The resume also stops as soon as the caller holds every answer token
  they asked for, or the hop returns answer content the service did not count.
  Dropping the ceiling instead was rejected: it hands the rest of the
  generation to the service default, and a 512-token request was measured
  answering with 61-75 KB of text. A request that pins no ceiling is resumed
  byte-identically, since there is nothing to extend.

* The ceiling is always asserted, never dropped. Measured against the
  service, `maxOutputTokens` of 65,537, 131,072 and 1,000,000 -- all above
  the model's own 65,536-token output maximum -- are each accepted with HTTP
  200, so there is no value a resume has to omit. Dropping the key above some
  threshold would have handed the rest of the generation to the service
  default, which is the unbounded behaviour the ceiling exists to prevent.

* A resume that fails for any other reason -- an expired token, a service
  refusal -- falls back to reporting what the generation itself did instead of
  the resume's own error. A failed `CONTINUATION` resume still propagates,
  since hiding it would turn a real failure into a silent truncation.

* `VertexAIGemini4Argon` no longer pins `max_tokens=65_536`. It is a class
  default nobody asked for, and under it every starved generation on the model
  was unresumable. Unsetting it is not the fix on its own -- a caller pins the
  ceiling itself through `?max_tokens=` in the model URI -- which is why the
  bounded extension above lives in the resume.

* `_check_answer_present` stops calling every `MAX_TOKENS` candidate starved.
  `candidatesTokenCount > 0` means the budget did reach the answer and the
  truncated hop was billed for tokens it did not return -- a different fault
  with a different remedy from thinking starvation, so the error now says
  which one occurred instead of always reporting the latter.

* The non-retryable rule is NARROW: only a signature a re-send cannot change
  ends the caller's generation. A blocked prompt, an exhausted `MAX_TOKENS`
  budget, a `STOP` billed for neither answer nor thinking tokens, and a
  candidate carrying no `content` at all are reported as
  `GeminiEmptyGenerationError`. Everything else that renders no answer -- a
  tool-protocol finish reason such as `MALFORMED_FUNCTION_CALL`, a
  candidate-level content-policy flag such as `SAFETY` or `RECITATION`, a
  `STOP` over non-zero token counts, an absent or unrecognised finish
  reason, an unexplained empty `candidates` list -- is left on the retryable
  path, because that is the path that recovers it: those arrive
  mid-generation on requests that declared no tools at all, so there is
  nothing in the request to correct and the next sample succeeds. Reporting
  them instead converts a transient fault into a fatal one.
  `_carries_unrendered_content` still carves out candidates whose parts hold
  a `functionCall` or `functionResponse`: those carry the turn's payload in
  a non-text part, so an empty text rendering is correct output, not a
  starved one.

* A candidate-level content-policy flag is NOT admitted, and the API
  contract is what settles it. `SAFETY`, `RECITATION`, `BLOCKLIST`,
  `PROHIBITED_CONTENT`, `SPII`, `LANGUAGE` and `IMAGE_SAFETY` are defined as
  flags raised on the response candidate content -- the tokens that happened
  to be sampled -- so the next sample is exactly what can clear them.
  `promptFeedback.blockReason` is the prompt-level counterpart: the API
  defines it on the prompt, returns no candidate with it, and names
  rephrasing the prompt as the remedy. Only that one is admitted. Admitting
  the candidate-level flags would end a live episode over an outcome a
  re-send can change, which is the same trade the closed rule lost 28 times
  out of 28, and no measured failure asks for it: across the 48 answer-less
  `gemini-4-argon` generations this CL was built from, every one reported
  `MAX_TOKENS` and none reported a content policy.

* `_merge_continuation_hops` keeps a chain that rendered no answer labelled
  `MAX_TOKENS` instead of taking the last hop's reason. A hop that exhausts
  the budget followed by one that adds nothing and reports `STOP` merged into
  `('STOP', summed count)` -- indistinguishable from a generation that ran
  and was lost, so it was re-sent `max_attempts` times, every re-send
  exhausting the same budget the same way. Closing that laundering where it
  happens is what lets the classifier above stay narrow. A chain that did
  render an answer keeps the last hop's reason, since it is the true one.

* Each answer-less signature carries the remedy that fits it, instead of one
  starvation tail appended to all of them. Lowering the thinking budget is
  advice only for a budget that thinking actually consumed; it is wrong for a
  budget that reached the answer, and irrelevant when the model stopped
  having generated nothing at all.

* `GeminiEmptyGenerationError` stays non-retryable for the signatures it
  admits. Every rescue that can work is spent inside the one request --
  following the `continuationToken`, then extending an explicit ceiling to 4x
  -- and each admitted signature reproduces on a re-send. Retrying would buy
  `max_attempts` full regenerations of the same outcome, each billed in full,
  so the message names that cost instead of incurring it.

* `_merge_continuation_hops` keeps bare `{'text': ''}` parts rather than
  filtering them. `Message.from_chunks` writes no separator after an empty
  chunk and strips the fused text, so such a part is already inert in every
  position; the filter, and the fallback that restored what it had removed,
  were both dead weight. A test pins the property they relied on.

PiperOrigin-RevId: 1003633198
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant