Skip to content

Commit f44e8ee

Browse files
google-genai-botcopybara-github
authored andcommitted
fix(python): discard stale checkpoint token on non-resumable finish reason
When a streaming response terminates with a non-resumable `finish_reason` (e.g., `STOP`, `MAX_TOKENS`, `SAFETY`), clear any intermediate checkpoint `continuation_token` received earlier in the same hop so a trailing stream error does not resume from an obsolete checkpoint. PiperOrigin-RevId: 998072455
1 parent 411efe0 commit f44e8ee

2 files changed

Lines changed: 51 additions & 0 deletions

File tree

‎google/genai/models.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6173,6 +6173,13 @@ def _generate_content_stream_with_continuation(
61736173
hop_finish_reason = candidate.finish_reason
61746174
if candidate.continuation_token:
61756175
continuation_token = candidate.continuation_token
6176+
elif (
6177+
candidate.finish_reason
6178+
and not _extra_utils.is_resumable_finish_reason(
6179+
candidate.finish_reason
6180+
)
6181+
):
6182+
continuation_token = None
61766183
chunk = _extra_utils.clear_intermediate_finish_reason(
61776184
chunk, enable_continuation, continuation_token
61786185
)
@@ -8432,6 +8439,13 @@ async def _generate_content_stream_with_continuation(
84328439
hop_finish_reason = candidate.finish_reason
84338440
if candidate.continuation_token:
84348441
continuation_token = candidate.continuation_token
8442+
elif (
8443+
candidate.finish_reason
8444+
and not _extra_utils.is_resumable_finish_reason(
8445+
candidate.finish_reason
8446+
)
8447+
):
8448+
continuation_token = None
84358449
chunk = _extra_utils.clear_intermediate_finish_reason(
84368450
chunk, enable_continuation, continuation_token
84378451
)

‎google/genai/tests/models/test_gc_continuation_token.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2095,6 +2095,24 @@ def _failing_hop(chunks_before_error, exc=None):
20952095
next(stream)
20962096
assert mock_raw_stream.call_count == 2
20972097

2098+
# 5. Error after STOP finish_reason in same hop -> stale checkpoint token was cleared, so does not resume
2099+
with mock.patch.object(
2100+
models.Models,
2101+
'_generate_content_stream',
2102+
side_effect=[
2103+
_failing_hop([chunk_ckpt_35k, chunk_hop2_final], server_503),
2104+
],
2105+
) as mock_raw_stream:
2106+
stream = models_module.generate_content_stream(
2107+
model='gemini-2.5-pro',
2108+
contents='Write a long book',
2109+
)
2110+
assert next(stream).text == 'Checkpoint 35k. '
2111+
assert next(stream).text == 'Resumed and completed.'
2112+
with pytest.raises(errors.ServerError):
2113+
next(stream)
2114+
assert mock_raw_stream.call_count == 1
2115+
20982116

20992117
@pytest.mark.asyncio
21002118
async def test_async_models_generate_content_stream_recovers_from_midstream_checkpoint_error(
@@ -2250,3 +2268,22 @@ async def _make_async_stream(chunk_list, exc=None):
22502268
received.append(chunk.text)
22512269
assert received == ['Async Ckpt 35k. ']
22522270
assert mock_async_stream.call_count == 2
2271+
2272+
# 5. Error after STOP finish_reason in same hop -> stale checkpoint token was cleared, so does not resume
2273+
with mock.patch.object(
2274+
models.AsyncModels,
2275+
'_generate_content_stream',
2276+
new_callable=mock.AsyncMock,
2277+
side_effect=[
2278+
_make_async_stream([chunk_ckpt_35k, chunk_hop2_final], server_504),
2279+
],
2280+
) as mock_async_stream:
2281+
received = []
2282+
with pytest.raises(errors.ServerError):
2283+
async for chunk in await models_module.generate_content_stream(
2284+
model='gemini-2.5-pro',
2285+
contents='Write a long book async',
2286+
):
2287+
received.append(chunk.text)
2288+
assert received == ['Async Ckpt 35k. ', 'Async resumed and completed.']
2289+
assert mock_async_stream.call_count == 1

0 commit comments

Comments
 (0)