From 302c948d788cd641db43438e582928e1a2a79dc5 Mon Sep 17 00:00:00 2001 From: Mark Daoust Date: Thu, 8 Oct 2026 01:19:46 -0700 Subject: [PATCH] feat(interactions): add automatic_continuation to interactions.create in the TypeScript SDK Foreground `interactions.create()` calls are capped per request by the server (32k output tokens unary, 64k streaming) and stop with `status='requires_continuation'` plus a `continuation_token`. `GeminiNextGenInteractions.create` now accepts `automatic_continuation` (default `true`) on `params` and chains continuation hops client-side so a plain `create()` returns the full decode. Request chaining: - Each hop after the first sends the `continuation_token` from the previous hop. With `store !== false` (default) it sets `previous_interaction_id` to the hop that issued the token and omits `input`; with `store: false` it resends the original `input`. - Unary foreground model calls default `generation_config.max_output_tokens` to 65536 when unset without mutating caller objects (streaming uses the server's 64k default). Explicit values are never overridden. - Stops on `completed`, `incomplete` (token preserved for manual resumption), `requires_action`, `failed`, `cancelled`, a missing token, or a repeated token. Skipped for `agent` / `agent_config` requests and `background: true`; `background: true` with `automatic_continuation: false` throws an error upfront. Unary (`stream: false`) merging: - Hops are merged into one `Interaction`: `steps` concatenated in hop order (leading `user_input` steps dropped on later hops), `usage` counts and per-modality breakdowns summed across hops with `usage.initial_input_tokens` taken from hop 1, `id` / `status` / `continuation_token` / `updated` from the final hop, and `previous_interaction_id` / `created` / `input` from hop 1. - `continuation_ids` lists every hop id in order, de-duplicated with any server-provided ids. `output_text` skips intervening thought steps across hop boundaries. Streaming (`stream: true`): - Events from all hops are yielded verbatim on a pull-on-demand `Stream`. Consumers should treat `interaction.completed` with `status='requires_continuation'` as intermediate. - Mid-hop transport drops or premature EOF resume from the latest `interaction.status_update` checkpoint token when no content was yielded since that checkpoint (up to 3 no-progress reconnects). SSE `error` events are terminal; HTTP status retries stay with `options.retries`. Timeouts and errors: - Numeric per-call `options.timeout_ms` (and alias `options.timeout`) bounds the whole `create()` call across hops via `CallDeadline`. Client-level `httpOptions.timeout`, `options.retries`, and `AbortSignal` apply per hop. - A failure or timeout after progress throws `ContinuationInterruptedError` (`cause` chained) carrying `partialInteraction` (unary), `continuationToken`, `interactionId`, and `continuationIds` so billed work is preserved for manual resumption. First-hop failures propagate unchanged. Opt-out: - `automatic_continuation: false` returns the single server response untouched (no budget injection, no merging) and emits a one-time `console.warn` if it stops at `requires_continuation` without a caller-supplied `continuation_token`. Tests: `test/unit/interactions_continuation_conformance_test.ts` (52 conformance cases) and `test/unit/interactions_continuation_test.ts` (whole-call timeout, wire hygiene, stream cancellation, and exports). PiperOrigin-RevId: 995657093 --- google/genai/_gaos/google_genai.py | 437 ++- google/genai/_gaos/interactions.py | 12 + google/genai/_interactions_continuation.py | 1162 +++++++ google/genai/errors.py | 45 + .../test_automatic_continuation.py | 2727 +++++++++++++++++ 5 files changed, 4286 insertions(+), 97 deletions(-) create mode 100644 google/genai/_interactions_continuation.py create mode 100644 google/genai/tests/interactions/test_automatic_continuation.py diff --git a/google/genai/_gaos/google_genai.py b/google/genai/_gaos/google_genai.py index 95c9fdf7c..74efdce22 100644 --- a/google/genai/_gaos/google_genai.py +++ b/google/genai/_gaos/google_genai.py @@ -26,6 +26,7 @@ from typing import TYPE_CHECKING, Any, Mapping, Optional, TypeVar, Union, cast +import functools import io import json import mimetypes @@ -33,6 +34,7 @@ import httpx +from .. import _interactions_continuation from ._hooks.google_genai_auth import ( GOOGLE_GENAI_API_REVISION as _GOOGLE_GENAI_API_REVISION, GoogleGenAISecurityProvider, @@ -41,6 +43,7 @@ from .agents import Agents as GeneratedAgents from .interactions import AsyncInteractions as GeneratedAsyncInteractions from .interactions import Interactions as GeneratedInteractions +from .lib import compat_errors from .lib.compat_errors import ( _AsyncRawResponseAccessorProxy, _RawResponseAccessorProxy, @@ -54,6 +57,7 @@ from .types import environments from .types import credentials from .types import interactions +from .types.basemodel import Unset from .types.security import Security from .utils import BackoffStrategy, RetryConfig, eventstreaming from .triggers import AsyncTriggers as GeneratedAsyncTriggers @@ -262,46 +266,173 @@ def create( extra_query: Optional[Mapping[str, Any]] = None, extra_body: Optional[Mapping[str, Any]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + automatic_continuation: Optional[bool] = None, **body: Any, ) -> Union[ interactions.Interaction, eventstreaming.Stream[interactions.InteractionSSEEvent], ]: - if request is not None: - if body: - raise TypeError(_REQUEST_AND_BODY_ERROR) - stream = _request_stream(request) - response = wrap_sdk_call( - super().create, - request=request, - api_version=api_version, - extra_headers=extra_headers, - extra_query=extra_query, - extra_body=extra_body, - timeout=timeout, - ) - if stream: - if is_stream(response): - return wrap_stream_errors(response) + """Creates a new interaction. + + :param request: Optional request object (`CreateInteractionRequest`). + :param api_version: API version for request routing. + :param extra_headers: Additional headers to set or replace on requests. + :param extra_query: Additional query parameters to append to requests. + :param extra_body: Additional JSON object fields to merge into request bodies. + :param timeout: Override the default request timeout configuration + for this method in seconds. When `automatic_continuation` is + enabled (the default) a numeric `timeout` bounds the whole call + (all continuation hops together), while the client-level + `http_options.timeout` and `http_options.retry_options` apply + to each hop; see the `automatic_continuation` notes below. + :param automatic_continuation: Defaults to `True`. When `True`, + automatically resumes decoding across multiple continuation hops + whenever a response or stream stops at the per-request decode + token limit (`status='requires_continuation'` with a + `continuation_token`) until the interaction completes, reaches + `incomplete` (e.g. `max_output_tokens`), requires action, or + fails. Pass `automatic_continuation=False` to disable. + - For `stream=True`, events from all continuation hops are + yielded sequentially on the returned `Stream`. Each + intermediate hop ends with an `interaction.completed` event + whose `event.interaction.status` is `'requires_continuation'`, + followed by `interaction.created` for the next hop; check + `event.interaction.status != 'requires_continuation'` before + treating an `interaction.completed` event as terminal. + `event.interaction.usage` on streamed events reflects only + that hop rather than cumulative usage across hops. + - For `stream=False`, each continuation hop is a separate + interaction on the server that only contains that hop's + output. The SDK merges the hops into one `Interaction`: + `steps` are concatenated in hop order, `usage` + token counts are summed across hops (with + `usage.initial_input_tokens` set to the first hop's input + tokens), and all other fields (`id`, `status`, ...) come from + the final hop. The final hop's `id` can be passed as + `previous_interaction_id` to continue the conversation. + `continuation_ids` lists the ids of every hop in order (the + last one is `id`); note that `interactions.get(id)` returns + only the final hop. `previous_interaction_id`, `created` and + `input` come from the first hop. + - With `store=True` (the default), each continuation hop sends + the hop that issued the `continuation_token` as + `previous_interaction_id` and omits `input` (the server + already has it). With `store=False`, the original `input` is + resent with the token. + - Client-side continuation is skipped for `background=True` and + agent interactions (`agent` / `agent_config`), whose loops are + managed by the server. + - Timeouts: a numeric per-call `timeout` is a wall-clock budget + for the whole `create()` call. Each hop is sent with the + time left, and once it runs out the call stops with + `ContinuationInterruptedError` (see Errors) so the work so + far can be resumed. The client-level `http_options.timeout` + (and an `httpx.Timeout` passed as `timeout`) is a per-request + setting and applies to each hop separately. + `http_options.retry_options` also applies per hop: transient + errors on a continuation hop retry that hop using the latest + `continuation_token` without restarting from the initial + request. + - Billing note: With `background=False`, each continuation hop + re-prefills the conversation and generated prefix, so input + tokens are billed on every hop (and summed in + `usage.total_input_tokens`). If an unrecoverable error occurs + after one or more hops have completed, those completed hops + are still billed. For server-managed continuation without + client-side hop loops, use `background=True` (with + `store=True`). + - Errors: a failure on the first request is raised as usual. + A failure after progress (a later hop, or a stream after a + resumable checkpoint) raises + `google.genai.errors.ContinuationInterruptedError` (with the + failure as `__cause__`) carrying `continuation_token`, + `interaction_id` and `continuation_ids`, plus + `partial_interaction` (the completed hops merged, for + `stream=False`), so the output already produced is not lost + and the call can be resumed manually. Running out of the + per-call `timeout` after progress is reported the same way + (with a `TimeoutError` as `__cause__`). + """ + prepared = _interactions_continuation.ContinuationRequest.from_create_args( + request=request, + body=body, + automatic_continuation=automatic_continuation, + normalize_body=_normalize_create_body, + unset_cls=Unset, + ) + call_kwargs = { + 'api_version': api_version, + 'extra_headers': extra_headers, + 'extra_query': extra_query, + 'extra_body': extra_body, + } + deadline = _interactions_continuation.CallDeadline( + timeout, enabled=prepared.enable_continuation + ) + send_once = functools.partial( + self._send_create_once, prepared, call_kwargs, deadline + ) + + response = send_once() + if prepared.stream: + if not is_stream(response): return response - return _add_output_properties_if_interaction(response) - stream = _optional_bool(body.get('stream'), default=False) - body = _normalize_create_body(body) - response = wrap_sdk_call( + wrapped_stream = wrap_stream_errors(response) + if not prepared.enable_continuation: + if prepared.warn_on_incomplete: + return _interactions_continuation.wrap_stream_incomplete_warning( + wrapped_stream + ) + return wrapped_stream + + return _interactions_continuation.wrap_stream_continuation( + wrapped_stream, + send_once, + is_stream_fn=is_stream, + wrap_errors_fn=wrap_stream_errors, + connection_error_type=compat_errors.APIConnectionError, + initial_continuation_token=prepared.caller_token, + deadline=deadline, + stateful=prepared.stateful, + ) + + if prepared.enable_continuation: + return _interactions_continuation.continue_unary_interaction( + response, + send_once, + initial_continuation_token=prepared.caller_token, + normalize_interaction=_normalize_interaction_shape, + finalize_interaction=_add_output_properties_if_interaction, + deadline=deadline, + stateful=prepared.stateful, + ) + result = _add_output_properties_if_interaction(response) + if prepared.warn_on_incomplete: + _interactions_continuation.warn_if_incomplete_unary(result) + return result + + def _send_create_once( + self, + prepared: _interactions_continuation.ContinuationRequest, + call_kwargs: Mapping[str, Any], + deadline: _interactions_continuation.CallDeadline, + token: Optional[str] = None, + token_owner_id: Optional[str] = None, + ) -> Any: + req, req_body = prepared.hop_args(token, token_owner_id) + call_kwargs = {**call_kwargs, 'timeout': deadline.hop_timeout()} + if req is not None: + return wrap_sdk_call( + super().create, + request=req, + **call_kwargs, + ) + return wrap_sdk_call( super().create, - api_version=api_version, - **cast(Any, body), - extra_headers=extra_headers, - extra_query=extra_query, - extra_body=extra_body, - timeout=timeout, + **cast(Any, req_body), + **call_kwargs, ) - if stream: - if is_stream(response): - return wrap_stream_errors(response) - return response - return _add_output_properties_if_interaction(response) if not TYPE_CHECKING: @@ -414,46 +545,173 @@ async def create( extra_query: Optional[Mapping[str, Any]] = None, extra_body: Optional[Mapping[str, Any]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + automatic_continuation: Optional[bool] = None, **body: Any, ) -> Union[ interactions.Interaction, eventstreaming.AsyncStream[interactions.InteractionSSEEvent], ]: - if request is not None: - if body: - raise TypeError(_REQUEST_AND_BODY_ERROR) - stream = _request_stream(request) - response = await async_wrap_sdk_call( - super().create, - request=request, - api_version=api_version, - extra_headers=extra_headers, - extra_query=extra_query, - extra_body=extra_body, - timeout=timeout, - ) - if stream: - if is_stream(response): - return wrap_async_stream_errors(response) + """Creates a new interaction asynchronously. + + :param request: Optional request object (`CreateInteractionRequest`). + :param api_version: API version for request routing. + :param extra_headers: Additional headers to set or replace on requests. + :param extra_query: Additional query parameters to append to requests. + :param extra_body: Additional JSON object fields to merge into request bodies. + :param timeout: Override the default request timeout configuration + for this method in seconds. When `automatic_continuation` is + enabled (the default) a numeric `timeout` bounds the whole call + (all continuation hops together), while the client-level + `http_options.timeout` and `http_options.retry_options` apply + to each hop; see the `automatic_continuation` notes below. + :param automatic_continuation: Defaults to `True`. When `True`, + automatically resumes decoding across multiple continuation hops + whenever a response or stream stops at the per-request decode + token limit (`status='requires_continuation'` with a + `continuation_token`) until the interaction completes, reaches + `incomplete` (e.g. `max_output_tokens`), requires action, or + fails. Pass `automatic_continuation=False` to disable. + - For `stream=True`, events from all continuation hops are + yielded sequentially on the returned `AsyncStream`. Each + intermediate hop ends with an `interaction.completed` event + whose `event.interaction.status` is `'requires_continuation'`, + followed by `interaction.created` for the next hop; check + `event.interaction.status != 'requires_continuation'` before + treating an `interaction.completed` event as terminal. + `event.interaction.usage` on streamed events reflects only + that hop rather than cumulative usage across hops. + - For `stream=False`, each continuation hop is a separate + interaction on the server that only contains that hop's + output. The SDK merges the hops into one `Interaction`: + `steps` are concatenated in hop order, `usage` + token counts are summed across hops (with + `usage.initial_input_tokens` set to the first hop's input + tokens), and all other fields (`id`, `status`, ...) come from + the final hop. The final hop's `id` can be passed as + `previous_interaction_id` to continue the conversation. + `continuation_ids` lists the ids of every hop in order (the + last one is `id`); note that `interactions.get(id)` returns + only the final hop. `previous_interaction_id`, `created` and + `input` come from the first hop. + - With `store=True` (the default), each continuation hop sends + the hop that issued the `continuation_token` as + `previous_interaction_id` and omits `input` (the server + already has it). With `store=False`, the original `input` is + resent with the token. + - Client-side continuation is skipped for `background=True` and + agent interactions (`agent` / `agent_config`), whose loops are + managed by the server. + - Timeouts: a numeric per-call `timeout` is a wall-clock budget + for the whole `create()` call. Each hop is sent with the + time left, and once it runs out the call stops with + `ContinuationInterruptedError` (see Errors) so the work so + far can be resumed. The client-level `http_options.timeout` + (and an `httpx.Timeout` passed as `timeout`) is a per-request + setting and applies to each hop separately. + `http_options.retry_options` also applies per hop: transient + errors on a continuation hop retry that hop using the latest + `continuation_token` without restarting from the initial + request. + - Billing note: With `background=False`, each continuation hop + re-prefills the conversation and generated prefix, so input + tokens are billed on every hop (and summed in + `usage.total_input_tokens`). If an unrecoverable error occurs + after one or more hops have completed, those completed hops + are still billed. For server-managed continuation without + client-side hop loops, use `background=True` (with + `store=True`). + - Errors: a failure on the first request is raised as usual. + A failure after progress (a later hop, or a stream after a + resumable checkpoint) raises + `google.genai.errors.ContinuationInterruptedError` (with the + failure as `__cause__`) carrying `continuation_token`, + `interaction_id` and `continuation_ids`, plus + `partial_interaction` (the completed hops merged, for + `stream=False`), so the output already produced is not lost + and the call can be resumed manually. Running out of the + per-call `timeout` after progress is reported the same way + (with a `TimeoutError` as `__cause__`). + """ + prepared = _interactions_continuation.ContinuationRequest.from_create_args( + request=request, + body=body, + automatic_continuation=automatic_continuation, + normalize_body=_normalize_create_body, + unset_cls=Unset, + ) + call_kwargs = { + 'api_version': api_version, + 'extra_headers': extra_headers, + 'extra_query': extra_query, + 'extra_body': extra_body, + } + deadline = _interactions_continuation.CallDeadline( + timeout, enabled=prepared.enable_continuation + ) + send_once = functools.partial( + self._send_create_once, prepared, call_kwargs, deadline + ) + + response = await send_once() + if prepared.stream: + if not is_stream(response): return response - return _add_output_properties_if_interaction(response) - stream = _optional_bool(body.get('stream'), default=False) - body = _normalize_create_body(body) - response = await async_wrap_sdk_call( + wrapped_stream = wrap_async_stream_errors(response) + if not prepared.enable_continuation: + if prepared.warn_on_incomplete: + return _interactions_continuation.wrap_async_stream_incomplete_warning( + wrapped_stream + ) + return wrapped_stream + + return _interactions_continuation.wrap_async_stream_continuation( + wrapped_stream, + send_once, + is_stream_fn=is_stream, + wrap_errors_fn=wrap_async_stream_errors, + connection_error_type=compat_errors.APIConnectionError, + initial_continuation_token=prepared.caller_token, + deadline=deadline, + stateful=prepared.stateful, + ) + + if prepared.enable_continuation: + return await _interactions_continuation.continue_async_unary_interaction( + response, + send_once, + initial_continuation_token=prepared.caller_token, + normalize_interaction=_normalize_interaction_shape, + finalize_interaction=_add_output_properties_if_interaction, + deadline=deadline, + stateful=prepared.stateful, + ) + result = _add_output_properties_if_interaction(response) + if prepared.warn_on_incomplete: + _interactions_continuation.warn_if_incomplete_unary(result) + return result + + async def _send_create_once( + self, + prepared: _interactions_continuation.ContinuationRequest, + call_kwargs: Mapping[str, Any], + deadline: _interactions_continuation.CallDeadline, + token: Optional[str] = None, + token_owner_id: Optional[str] = None, + ) -> Any: + req, req_body = prepared.hop_args(token, token_owner_id) + call_kwargs = {**call_kwargs, 'timeout': deadline.hop_timeout()} + if req is not None: + return await async_wrap_sdk_call( + super().create, + request=req, + **call_kwargs, + ) + return await async_wrap_sdk_call( super().create, - api_version=api_version, - **cast(Any, body), - extra_headers=extra_headers, - extra_query=extra_query, - extra_body=extra_body, - timeout=timeout, + **cast(Any, req_body), + **call_kwargs, ) - if stream: - if is_stream(response): - return wrap_async_stream_errors(response) - return response - return _add_output_properties_if_interaction(response) if not TYPE_CHECKING: @@ -1338,16 +1596,9 @@ def _add_output_properties_if_interaction(value: Any) -> Any: steps = _get_value(normalized, 'steps') updates = _output_properties_from_steps(steps) if not updates: - return value - - if hasattr(normalized, 'model_copy'): - return normalized.model_copy(update=updates) - if isinstance(normalized, dict): - return {**normalized, **updates} + return normalized - for name, output_value in updates.items(): - setattr(normalized, name, output_value) - return normalized + return _set_or_copy(normalized, updates) def _normalize_interaction_shape(value: Any) -> Optional[Any]: @@ -1398,13 +1649,15 @@ def _output_properties_from_steps(steps: list[Any]) -> dict[str, Any]: should_stop = False for item in reversed(content): - if _get_value(item, 'type') == 'text': - collecting = True - text = _get_value(item, 'text') - text_parts.append(text if isinstance(text, str) else '') - elif collecting: - should_stop = True - break + item_type = _get_value(item, 'type') + if item_type != 'text': + if collecting: + should_stop = True + break + continue + collecting = True + text = _get_value(item, 'text') + text_parts.append(text if isinstance(text, str) else '') if should_stop: break @@ -1447,7 +1700,7 @@ def _output_properties_from_steps(steps: list[Any]) -> dict[str, Any]: def _get_value(value: Any, name: str) -> Any: - if isinstance(value, dict): + if isinstance(value, Mapping): return value.get(name) return getattr(value, name, None) @@ -1464,12 +1717,6 @@ def _get_value(value: Any, name: str) -> Any: ) -_REQUEST_AND_BODY_ERROR = ( - 'create() accepts either request=... or individual body keyword arguments, ' - 'not both.' -) - - def _normalize_create_body(body: dict[str, Any]) -> dict[str, Any]: unknown = set(body) - _CREATE_BODY_KEYS if unknown: @@ -1546,15 +1793,11 @@ def _optional_str(value: Any) -> Optional[str]: return None -def _request_stream(request: Any) -> bool: - body = ( - request.get('body') - if isinstance(request, Mapping) - else getattr(request, 'body', None) - ) - stream = ( - body.get('stream') - if isinstance(body, Mapping) - else getattr(body, 'stream', None) - ) - return bool(_optional_bool(stream, default=False)) +def _set_or_copy(target: Any, updates: dict[str, Any]) -> Any: + if hasattr(target, 'model_copy'): + return target.model_copy(update=updates) + if isinstance(target, dict): + return {**target, **updates} + for name, value in updates.items(): + setattr(target, name, value) + return target diff --git a/google/genai/_gaos/interactions.py b/google/genai/_gaos/interactions.py index dbddded0a..f34d65f85 100644 --- a/google/genai/_gaos/interactions.py +++ b/google/genai/_gaos/interactions.py @@ -48,6 +48,7 @@ def create( extra_query: Optional[Mapping[str, Any]] = None, extra_body: Optional[Mapping[str, Any]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + automatic_continuation: Optional[bool] = None, ) -> Union[interactions.Interaction, Stream[interactions.InteractionSSEEvent]]: r"""Creates a new interaction. @@ -87,6 +88,7 @@ def create( extra_query: Optional[Mapping[str, Any]] = None, extra_body: Optional[Mapping[str, Any]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + automatic_continuation: Optional[bool] = None, ) -> interactions.Interaction: r"""Creates a new interaction. @@ -154,6 +156,7 @@ def create( extra_query: Optional[Mapping[str, Any]] = None, extra_body: Optional[Mapping[str, Any]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + automatic_continuation: Optional[bool] = None, ) -> Stream[interactions.InteractionSSEEvent]: r"""Creates a new interaction. @@ -220,6 +223,7 @@ def create( extra_query: Optional[Mapping[str, Any]] = None, extra_body: Optional[Mapping[str, Any]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + automatic_continuation: Optional[bool] = None, ) -> interactions.Interaction: r"""Creates a new interaction. @@ -285,6 +289,7 @@ def create( extra_query: Optional[Mapping[str, Any]] = None, extra_body: Optional[Mapping[str, Any]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + automatic_continuation: Optional[bool] = None, ) -> Stream[interactions.InteractionSSEEvent]: r"""Creates a new interaction. @@ -335,6 +340,7 @@ def create( extra_query: Optional[Mapping[str, Any]] = None, extra_body: Optional[Mapping[str, Any]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + automatic_continuation: Optional[bool] = None, **body_kwargs: Any, ) -> Union[interactions.Interaction, Stream[interactions.InteractionSSEEvent]]: r"""Creates a new interaction. @@ -1317,6 +1323,7 @@ async def create( extra_query: Optional[Mapping[str, Any]] = None, extra_body: Optional[Mapping[str, Any]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + automatic_continuation: Optional[bool] = None, ) -> Union[interactions.Interaction, AsyncStream[interactions.InteractionSSEEvent]]: r"""Creates a new interaction. @@ -1356,6 +1363,7 @@ async def create( extra_query: Optional[Mapping[str, Any]] = None, extra_body: Optional[Mapping[str, Any]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + automatic_continuation: Optional[bool] = None, ) -> interactions.Interaction: r"""Creates a new interaction. @@ -1423,6 +1431,7 @@ async def create( extra_query: Optional[Mapping[str, Any]] = None, extra_body: Optional[Mapping[str, Any]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + automatic_continuation: Optional[bool] = None, ) -> AsyncStream[interactions.InteractionSSEEvent]: r"""Creates a new interaction. @@ -1489,6 +1498,7 @@ async def create( extra_query: Optional[Mapping[str, Any]] = None, extra_body: Optional[Mapping[str, Any]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + automatic_continuation: Optional[bool] = None, ) -> interactions.Interaction: r"""Creates a new interaction. @@ -1554,6 +1564,7 @@ async def create( extra_query: Optional[Mapping[str, Any]] = None, extra_body: Optional[Mapping[str, Any]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + automatic_continuation: Optional[bool] = None, ) -> AsyncStream[interactions.InteractionSSEEvent]: r"""Creates a new interaction. @@ -1604,6 +1615,7 @@ async def create( extra_query: Optional[Mapping[str, Any]] = None, extra_body: Optional[Mapping[str, Any]] = None, timeout: Optional[Union[float, httpx.Timeout]] = None, + automatic_continuation: Optional[bool] = None, **body_kwargs: Any, ) -> Union[interactions.Interaction, AsyncStream[interactions.InteractionSSEEvent]]: r"""Creates a new interaction. diff --git a/google/genai/_interactions_continuation.py b/google/genai/_interactions_continuation.py new file mode 100644 index 000000000..c5e076301 --- /dev/null +++ b/google/genai/_interactions_continuation.py @@ -0,0 +1,1162 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Automatic continuation state and helpers for the Interactions API.""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Mapping +import dataclasses +import time +from typing import Any, Optional +import warnings + +from . import errors + +# Maximum consecutive reconnects after a mid-hop disconnect that made no +# progress (no new checkpoint token since the hop's start token). +MAX_STREAM_RECONNECTS_WITHOUT_PROGRESS = 3 + +_AUTO_CONTINUE_AND_BACKGROUND_ERROR = ( + 'automatic_continuation=False cannot be used with background=True.' +) + +_REQUEST_AND_BODY_ERROR = ( + 'create() accepts either request=... or individual body keyword arguments, ' + 'not both.' +) + +_INCOMPLETE_CONTINUATION_WARNING = ( + "Interaction stopped before completion (status='requires_continuation') " + 'with a continuation_token because the per-request decode token limit was ' + 'reached. To resume decoding automatically, pass ' + 'automatic_continuation=True (note: each continuation hop re-prefills and ' + 'bills input tokens), or use background=True (the most cost-efficient ' + 'option for server-managed continuation), or resume manually by passing ' + 'continuation_token to interactions.create() (with ' + "previous_interaction_id set to this interaction's id and no input, or " + 'with the original input when store=False).' +) + +_RESUMABLE_STREAM_HOP_STATUSES = frozenset({ + None, + 'in_progress', + 'requires_continuation', +}) + +# Events that carry no model output. A stream that drops after one of these (and +# nothing else) since the last checkpoint can be resumed without duplicating +# output. +_NON_CONTENT_STREAM_EVENTS = frozenset({ + 'interaction.created', + 'interaction.status_update', +}) + +# Usage list fields are merged per key (e.g. per modality), summing the value. +_USAGE_LIST_KEYS: dict[str, tuple[str, str]] = { + 'cached_tokens_by_modality': ('modality', 'tokens'), + 'input_tokens_by_modality': ('modality', 'tokens'), + 'output_tokens_by_modality': ('modality', 'tokens'), + 'tool_use_tokens_by_modality': ('modality', 'tokens'), + 'grounding_tool_count': ('type', 'count'), +} + +_incomplete_continuation_warned = False + + +class ContinuationTimeoutError(TimeoutError): + """The per-call `timeout` ran out before automatic continuation finished.""" + + +class CallDeadline: + """Wall-clock budget of one logical `interactions.create()` call. + + A numeric per-call `timeout` bounds the whole call rather than each + continuation hop (HTTP request): the deadline is fixed when the call starts, + every hop is sent with the time left, and no further hop starts once it has + passed. An `httpx.Timeout` (or `None`) is a per-request setting and is passed + to each hop unchanged, so the call is unbounded across hops. + """ + + def __init__( + self, + timeout: Any, + *, + enabled: bool = True, + clock: Optional[Callable[[], float]] = None, + ): + self._timeout = timeout + self._clock = clock or time.monotonic + self._deadline: Optional[float] = None + if ( + enabled + and isinstance(timeout, (int, float)) + and not isinstance(timeout, bool) + ): + self._deadline = self._clock() + float(timeout) + + @property + def timeout(self) -> Any: + """The per-call `timeout` this deadline was created from.""" + return self._timeout + + def remaining(self) -> Optional[float]: + """Seconds left before the deadline, or None if unbounded.""" + if self._deadline is None: + return None + return self._deadline - self._clock() + + def expired(self) -> bool: + remaining = self.remaining() + return remaining is not None and remaining <= 0 + + def hop_timeout(self) -> Any: + """The `timeout` to send with the next hop (the time left, if bounded).""" + remaining = self.remaining() + if remaining is None: + return self._timeout + return max(remaining, 0.0) + + def check(self) -> None: + """Raises `ContinuationTimeoutError` once the deadline has passed.""" + if self.expired(): + raise ContinuationTimeoutError( + f'interactions.create() timed out after {self._timeout} seconds' + ' before automatic continuation finished. Increase `timeout` (it' + ' bounds the whole call, across continuation hops) or resume from' + ' the continuation_token.' + ) + + +def _get_value(value: Any, name: str) -> Any: + if isinstance(value, Mapping): + return value.get(name) + return getattr(value, name, None) + + +def _set_or_copy(target: Any, updates: dict[str, Any]) -> Any: + if hasattr(target, 'model_copy'): + return target.model_copy(update=updates) + if isinstance(target, dict): + return {**target, **updates} + for name, value in updates.items(): + if hasattr(target, '__dict__') or hasattr(type(target), name): + setattr(target, name, value) + return target + + +def _optional_bool( + value: Any, default: Optional[bool] = None +) -> Optional[bool]: + if isinstance(value, bool): + return value + return default + + +def _non_empty_str(value: Any) -> Optional[str]: + return value if isinstance(value, str) and value else None + + +def _is_unset_or_none( + value: Any, unset_cls: Optional[type[Any]] = None +) -> bool: + if value is None: + return True + if unset_cls is not None and isinstance(value, unset_cls): + return True + return type(value).__name__ == 'Unset' + + +def _is_agent_body(body: Any, unset_cls: Optional[type[Any]] = None) -> bool: + """Whether the create body targets an agent (server-managed loop).""" + return not _is_unset_or_none( + _get_value(body, 'agent'), unset_cls + ) or not _is_unset_or_none(_get_value(body, 'agent_config'), unset_cls) + + +def _continuation_body( + body: Any, token: str, token_owner_id: Optional[str] +) -> Any: + """Returns the create body for a continuation hop resuming from `token`. + + With `store` enabled (the default) the server rebuilds the context from the + stored interaction that issued `token`, so the hop sends that interaction as + `previous_interaction_id` (the server checks the token belongs to it) and + omits `input`. With `store=False` nothing is stored, so the original `input` + is resent with the token instead. + + `token_owner_id` is None when the token did not come from an earlier hop + (e.g. the caller's own `continuation_token`); the original request is then + resent as-is with `token`. + """ + updates: dict[str, Any] = {'continuation_token': token} + stateful = bool(token_owner_id) and _get_value(body, 'store') is not False + if stateful: + updates['previous_interaction_id'] = token_owner_id + if isinstance(body, Mapping): + new_body = {**body, **updates} + if stateful: + new_body.pop('input', None) + return new_body + if hasattr(body, 'model_copy'): + if stateful: + updates['input'] = None + return body.model_copy(update=updates) + return body + + +def _with_request_continuation( + request: Any, token: str, token_owner_id: Optional[str] +) -> Any: + body = _get_value(request, 'body') + new_body = _continuation_body(body, token, token_owner_id) + if isinstance(request, Mapping): + return {**request, 'body': new_body} + if hasattr(request, 'model_copy'): + return request.model_copy(update={'body': new_body}) + return request + + +@dataclasses.dataclass(frozen=True) +class ContinuationRequest: + """Prepared `interactions.create()` request state and per-hop builder.""" + + request: Any + body: Mapping[str, Any] + stream: bool + caller_token: Optional[str] + enable_continuation: bool + warn_on_incomplete: bool = False + stateful: bool = True + + @classmethod + def from_create_args( + cls, + *, + request: Any, + body: dict[str, Any], + automatic_continuation: Optional[bool], + normalize_body: Callable[[dict[str, Any]], dict[str, Any]], + unset_cls: Optional[type[Any]] = None, + ) -> ContinuationRequest: + """Validates and normalizes `create()` arguments for continuation.""" + if request is not None: + if body: + raise TypeError(_REQUEST_AND_BODY_ERROR) + request_body = _get_value(request, 'body') + is_background = bool( + _optional_bool(_get_value(request_body, 'background'), default=False) + ) + stream = bool( + _optional_bool(_get_value(request_body, 'stream'), default=False) + ) + caller_token = _non_empty_str( + _get_value(request_body, 'continuation_token') + ) + if automatic_continuation is None: + req_auto = _get_value(request, 'automatic_continuation') + if req_auto is None: + req_auto = _get_value(request_body, 'automatic_continuation') + automatic_continuation = req_auto + normalized_body: Mapping[str, Any] = body + else: + is_background = bool( + _optional_bool(body.get('background'), default=False) + ) + stream = bool(_optional_bool(body.get('stream'), default=False)) + normalized_body = normalize_body(body) + caller_token = _non_empty_str(normalized_body.get('continuation_token')) + request_body = normalized_body + + if is_background and automatic_continuation is False: + raise ValueError(_AUTO_CONTINUE_AND_BACKGROUND_ERROR) + + auto_continue_enabled = bool( + _optional_bool(automatic_continuation, default=True) + ) + is_agent = _is_agent_body(request_body, unset_cls) + enable_continuation = ( + auto_continue_enabled and not is_background and not is_agent + ) + warn_on_incomplete = ( + not auto_continue_enabled + and not is_background + and not is_agent + and not caller_token + ) + stateful = _get_value(request_body, 'store') is not False + + return cls( + request=request, + body=normalized_body, + stream=stream, + caller_token=caller_token, + enable_continuation=enable_continuation, + warn_on_incomplete=warn_on_incomplete, + stateful=stateful, + ) + + def hop_args( + self, + token: Optional[str] = None, + token_owner_id: Optional[str] = None, + ) -> tuple[Any, Mapping[str, Any]]: + """Returns `(request, body)` for the initial hop or a continuation hop.""" + if self.request is not None: + req = ( + self.request + if token is None + else _with_request_continuation(self.request, token, token_owner_id) + ) + return req, self.body + req_body = ( + self.body + if token is None + else _continuation_body(self.body, token, token_owner_id) + ) + return None, req_body + + +def _extract_hop_continuation_state( + event: Any, + hop_status: Optional[str], + hop_continuation_token: Optional[str], +) -> tuple[Optional[str], Optional[str]]: + """Extracts `(status, continuation_token)` updates from one SSE event.""" + if _get_value(event, 'event_type') == 'error': + return 'error', None + + interaction = _get_value(event, 'interaction') + if interaction is not None: + status = _get_value(interaction, 'status') + if isinstance(status, str) and status: + hop_status = status + token = _non_empty_str(_get_value(interaction, 'continuation_token')) + if _get_value(event, 'event_type') == 'interaction.completed': + # The terminal snapshot is authoritative: a missing token clears the hop's + # seed token and any earlier mid-hop checkpoint token. + hop_continuation_token = token + elif token: + hop_continuation_token = token + + direct_status = _get_value(event, 'status') + if isinstance(direct_status, str) and direct_status: + hop_status = direct_status + direct_token = _get_value(event, 'continuation_token') + if isinstance(direct_token, str) and direct_token: + hop_continuation_token = direct_token + + return hop_status, hop_continuation_token + + +def _has_content_since_checkpoint( + event: Any, + content_since_checkpoint: bool, +) -> bool: + """Tracks whether output was yielded since the latest resume point.""" + event_type = _get_value(event, 'event_type') + if event_type == 'interaction.status_update': + if _non_empty_str(_get_value(event, 'continuation_token')): + return False + return content_since_checkpoint + if event_type == 'interaction.completed': + # The terminal snapshot's token (if any) is a fresh resume point. + return False + if event_type in _NON_CONTENT_STREAM_EVENTS: + return content_since_checkpoint + return True + + +def _event_interaction_id(event: Any) -> Optional[str]: + """The id of the interaction (hop) an SSE event belongs to, if present.""" + interaction_id = _non_empty_str(_get_value(event, 'interaction_id')) + if interaction_id: + return interaction_id + return _non_empty_str(_get_value(_get_value(event, 'interaction'), 'id')) + + +def _can_resume_stream_hop( + hop_status: Optional[str], + hop_continuation_token: Optional[str], +) -> bool: + return bool( + hop_continuation_token and hop_status in _RESUMABLE_STREAM_HOP_STATUSES + ) + + +def _next_no_progress_reconnects( + no_progress_reconnects: int, + hop_status: Optional[str], + hop_start_token: Optional[str], + hop_continuation_token: Optional[str], +) -> int: + """Updates the no-progress reconnect counter after a resumable hop ends.""" + if hop_status == 'requires_continuation': + # Clean hop boundary (per-request decode limit reached). + return 0 + if hop_continuation_token != hop_start_token: + # Dropped mid-hop after a new checkpoint token was received. + return 0 + # Dropped mid-hop before any new checkpoint: retry the hop's own token. + return no_progress_reconnects + 1 + + +@dataclasses.dataclass +class _StreamHop: + """What was observed during one hop (HTTP request) of a continuation stream.""" + + # The continuation_token and owner id this hop was requested with. + start_token: Optional[str] = None + start_token_owner_id: Optional[str] = None + # Seeded with `start_token` so a disconnect before the hop's first checkpoint + # retries from the hop start. + token: Optional[str] = None + # The interaction that issued `token`. With store=True the server only accepts + # a token together with the interaction that issued it as + # `previous_interaction_id`, so a retry of the seed token keeps the seed's + # owner even after this hop's own interaction was created. + token_owner_id: Optional[str] = None + # This hop's own interaction id, once known. + interaction_id: Optional[str] = None + # Latest status observed during the hop, or None if none arrived yet. + status: Optional[str] = None + # Whether the HTTP stream was opened (response headers received). + stream_opened: bool = False + # Whether output events were yielded since the latest checkpoint token. + content_since_checkpoint: bool = False + + +class StreamContinuation: + """State of a streaming `interactions.create()` call with continuation. + + A stream is made of one or more hops (HTTP requests). Per-hop state lives in a + `_StreamHop` record replaced wholesale by `start_hop()`, so per-hop fields + cannot leak across hops. Loop state (the next resume token + owner id, the + no-progress reconnect counter, and the pending transport error) persists on + `StreamContinuation` across hops. + """ + + def __init__( + self, + initial_continuation_token: Optional[str] = None, + *, + connection_error_type: type[BaseException] = Exception, + stateful: bool = True, + ): + self._next_token: Optional[str] = initial_continuation_token + # The caller's own token: a failure that leaves nothing newer to resume + # from is not "after progress" and is raised as-is. + self._initial_token: Optional[str] = initial_continuation_token + # The interaction that issued `_next_token`; None for the caller's own + # token (the original request is resent as-is). + self._next_token_owner_id: Optional[str] = None + # The ids of every hop's interaction, in order (for error reporting). + self._hop_ids: list[str] = [] + self._no_progress_reconnects = 0 + # The latest transport error not followed by progress; raised if the loop + # gives up so the failure is not hidden by later clean closes. + self._last_error: Optional[Exception] = None + self._connection_error_type = connection_error_type + self._stateful = stateful + self._hop = _StreamHop( + start_token=initial_continuation_token, + token=initial_continuation_token, + ) + + def start_hop( + self, *, stream_opened: bool = False + ) -> tuple[Optional[str], Optional[str]]: + """Starts a new hop and returns `(token, token_owner_id)` to request.""" + self._hop = _StreamHop( + start_token=self._next_token, + start_token_owner_id=self._next_token_owner_id, + token=self._next_token, + token_owner_id=self._next_token_owner_id, + stream_opened=stream_opened, + ) + return self._next_token, self._next_token_owner_id + + def mark_stream_opened(self) -> None: + """Marks that the current hop's HTTP stream was opened.""" + self._hop.stream_opened = True + + def observe(self, event: Any) -> Any: + """Records one streamed SSE event and returns it.""" + hop = self._hop + hop.stream_opened = True + event_interaction_id = _event_interaction_id(event) + if event_interaction_id: + hop.interaction_id = event_interaction_id + if event_interaction_id not in self._hop_ids: + self._hop_ids.append(event_interaction_id) + previous_token = hop.token + hop.status, hop.token = _extract_hop_continuation_state( + event, hop.status, hop.token + ) + if hop.token and hop.token != previous_token: + hop.token_owner_id = hop.interaction_id + hop.content_since_checkpoint = _has_content_since_checkpoint( + event, hop.content_since_checkpoint + ) + return event + + def _can_resume_hop(self) -> bool: + # Resuming from the checkpoint would re-generate (and duplicate) any output + # yielded since it. + hop = self._hop + return not hop.content_since_checkpoint and _can_resume_stream_hop( + hop.status, hop.token + ) + + def _fail(self, error: Exception) -> Exception: + """The error to raise when the stream gives up with `error`. + + If the server issued a resumable token during this call (a completed + `requires_continuation` hop, or a mid-hop checkpoint), the work so far is + resumable and `error` is wrapped in `errors.ContinuationInterruptedError` + carrying that token. Otherwise `error` is returned as-is. + """ + hop = self._hop + token = ( + hop.token if _can_resume_stream_hop(hop.status, hop.token) else None + ) + if not token or token == self._initial_token: + return error + interrupted = errors.ContinuationInterruptedError( + f'Automatic continuation of the stream was interrupted: {error}.' + ' The continuation_token can be used to resume (output yielded after' + ' a mid-hop checkpoint is re-generated).', + partial_interaction=None, + continuation_token=token, + interaction_id=hop.token_owner_id if self._stateful else None, + continuation_ids=list(self._hop_ids), + ) + interrupted.__cause__ = error + return interrupted + + def end_hop(self, error: Optional[Exception] = None) -> bool: + """Ends the current hop and returns True if another hop should start. + + Raises `error` (or the pending prior disconnect error) when the stream + cannot safely resume or exceeds `MAX_STREAM_RECONNECTS_WITHOUT_PROGRESS`, + wrapped in `errors.ContinuationInterruptedError` if a server-issued token + to resume from exists (see `_fail`). + """ + hop = self._hop + if isinstance(error, ContinuationTimeoutError): + raise self._fail(error) + if error is not None and not hop.stream_opened: + # Pre-stream HTTP status or validation errors (already retried by + # RetryConfig) must propagate immediately; only connection errors resume. + if not isinstance(error, self._connection_error_type): + raise self._fail(error) + + if not self._can_resume_hop(): + if error is not None: + raise self._fail(error) + return False + + if hop.status == 'requires_continuation' and hop.token == hop.start_token: + # The hop finished but handed back the token it was started with: + # resuming would loop forever without progress. + if error is not None: + raise self._fail(error) + return False + + self._no_progress_reconnects = _next_no_progress_reconnects( + self._no_progress_reconnects, + hop.status, + hop.start_token, + hop.token, + ) + if self._no_progress_reconnects == 0: + self._last_error = None + if error is not None: + self._last_error = error + if self._no_progress_reconnects > MAX_STREAM_RECONNECTS_WITHOUT_PROGRESS: + if self._last_error is not None: + raise self._fail(self._last_error) + return False + + self._next_token = hop.token + self._next_token_owner_id = hop.token_owner_id + return True + + +def wrap_stream_continuation( + stream: Any, + create_next_stream: Callable[[Optional[str], Optional[str]], Any], + *, + is_stream_fn: Callable[[Any], bool], + wrap_errors_fn: Callable[[Any], Any], + connection_error_type: type[BaseException], + initial_continuation_token: Optional[str] = None, + deadline: Optional[CallDeadline] = None, + stateful: bool = True, +) -> Any: + """Wraps a sync `Stream` to automatically continue across hops. + + `deadline` (if bounded) is checked before each event and before each new + hop; once it has passed the stream stops with a `ContinuationTimeoutError` + (wrapped in `errors.ContinuationInterruptedError` if resumable). + """ + current_gen = stream.generator + current_stream = stream + deadline = deadline or CallDeadline(None) + + def _continuation_gen() -> Generator[Any, None, None]: + nonlocal current_gen, current_stream + state = StreamContinuation( + initial_continuation_token=initial_continuation_token, + connection_error_type=connection_error_type, + stateful=stateful, + ) + is_first_hop = True + try: + while True: + token, token_owner_id = state.start_hop(stream_opened=is_first_hop) + error: Optional[Exception] = None + if is_first_hop: + is_first_hop = False + else: + if hasattr(current_stream, 'response'): + current_stream.response.close() + try: + deadline.check() + next_stream = create_next_stream(token, token_owner_id) + except Exception as e: # pylint: disable=broad-exception-caught + error = e + else: + if not is_stream_fn(next_stream): + break + state.mark_stream_opened() + next_stream = wrap_errors_fn(next_stream) + current_stream = next_stream + current_gen = next_stream.generator + if hasattr(stream, 'response') and hasattr(next_stream, 'response'): + stream.response = next_stream.response + + if error is None: + it = iter(current_gen) + while True: + try: + deadline.check() + event = next(it) + except StopIteration: + break + except Exception as e: # pylint: disable=broad-exception-caught + error = e + break + yield state.observe(event) + + if not state.end_hop(error): + break + finally: + if hasattr(current_gen, 'close'): + current_gen.close() + if current_stream is not stream and hasattr(current_stream, 'close'): + current_stream.close() + + continuation_gen = _continuation_gen() + stream.generator = continuation_gen + + orig_close = stream.close + + def _close() -> None: + try: + continuation_gen.close() + finally: + orig_close() + + stream.close = _close + return stream + + +def wrap_async_stream_continuation( + stream: Any, + create_next_stream: Callable[ + [Optional[str], Optional[str]], Awaitable[Any] + ], + *, + is_stream_fn: Callable[[Any], bool], + wrap_errors_fn: Callable[[Any], Any], + connection_error_type: type[BaseException], + initial_continuation_token: Optional[str] = None, + deadline: Optional[CallDeadline] = None, + stateful: bool = True, +) -> Any: + """Wraps an `AsyncStream` to automatically continue across hops.""" + current_gen = stream.generator + current_stream = stream + deadline = deadline or CallDeadline(None) + + async def _continuation_agen() -> AsyncGenerator[Any, None]: + nonlocal current_gen, current_stream + state = StreamContinuation( + initial_continuation_token=initial_continuation_token, + connection_error_type=connection_error_type, + stateful=stateful, + ) + is_first_hop = True + try: + while True: + token, token_owner_id = state.start_hop(stream_opened=is_first_hop) + error: Optional[Exception] = None + if is_first_hop: + is_first_hop = False + else: + if hasattr(current_stream, 'response'): + await current_stream.response.aclose() + try: + deadline.check() + next_stream = await create_next_stream(token, token_owner_id) + except Exception as e: # pylint: disable=broad-exception-caught + error = e + else: + if not is_stream_fn(next_stream): + break + state.mark_stream_opened() + next_stream = wrap_errors_fn(next_stream) + current_stream = next_stream + current_gen = next_stream.generator + if hasattr(stream, 'response') and hasattr(next_stream, 'response'): + stream.response = next_stream.response + + if error is None: + it = current_gen.__aiter__() + while True: + try: + deadline.check() + event = await it.__anext__() + except StopAsyncIteration: + break + except Exception as e: # pylint: disable=broad-exception-caught + error = e + break + yield state.observe(event) + + if not state.end_hop(error): + break + finally: + if hasattr(current_gen, 'aclose'): + await current_gen.aclose() + if current_stream is not stream and hasattr(current_stream, 'close'): + await current_stream.close() + + continuation_agen = _continuation_agen() + stream.generator = continuation_agen + + orig_close = stream.close + + async def _close() -> None: + try: + await continuation_agen.aclose() + finally: + await orig_close() + + stream.close = _close + return stream + + +def _warn_incomplete_continuation(*, stacklevel: int = 3) -> None: + global _incomplete_continuation_warned + if not _incomplete_continuation_warned: + _incomplete_continuation_warned = True + warnings.warn( + _INCOMPLETE_CONTINUATION_WARNING, + category=UserWarning, + stacklevel=stacklevel, + ) + + +def warn_if_incomplete_unary(interaction: Any) -> None: + """Emits a one-time warning when unary `create()` stops uncontinued.""" + if _should_continue_unary(interaction): + _warn_incomplete_continuation(stacklevel=4) + + +def wrap_stream_incomplete_warning(stream: Any) -> Any: + """Wraps a sync `Stream` to emit a one-time warning when left uncontinued.""" + current_gen = stream.generator + + def _warning_gen() -> Generator[Any, None, None]: + hop_status: Optional[str] = None + hop_continuation_token: Optional[str] = None + try: + for event in current_gen: + hop_status, hop_continuation_token = _extract_hop_continuation_state( + event, + hop_status, + hop_continuation_token, + ) + yield event + if hop_status == 'requires_continuation' and hop_continuation_token: + _warn_incomplete_continuation(stacklevel=3) + finally: + if hasattr(current_gen, 'close'): + current_gen.close() + + warning_gen = _warning_gen() + stream.generator = warning_gen + + orig_close = stream.close + + def _close() -> None: + try: + warning_gen.close() + finally: + orig_close() + + stream.close = _close + return stream + + +def wrap_async_stream_incomplete_warning(stream: Any) -> Any: + """Wraps an `AsyncStream` to emit a one-time warning when left uncontinued.""" + current_gen = stream.generator + + async def _warning_agen() -> AsyncGenerator[Any, None]: + hop_status: Optional[str] = None + hop_continuation_token: Optional[str] = None + try: + async for event in current_gen: + hop_status, hop_continuation_token = _extract_hop_continuation_state( + event, + hop_status, + hop_continuation_token, + ) + yield event + if hop_status == 'requires_continuation' and hop_continuation_token: + _warn_incomplete_continuation(stacklevel=3) + finally: + if hasattr(current_gen, 'aclose'): + await current_gen.aclose() + + warning_agen = _warning_agen() + stream.generator = warning_agen + + orig_close = stream.close + + async def _close() -> None: + try: + await warning_agen.aclose() + finally: + await orig_close() + + stream.close = _close + return stream + + +def _should_continue_unary(interaction: Any) -> Optional[str]: + if _get_value(interaction, 'status') != 'requires_continuation': + return None + token = _get_value(interaction, 'continuation_token') + if not isinstance(token, str) or not token: + return None + return token + + +def _next_unary_hop_token(hop: Any, sent_token: Optional[str]) -> Optional[str]: + """The token to continue from after `hop`, or None to stop.""" + token = _should_continue_unary(hop) + if token == sent_token: + # The hop handed back the token it was started with: continuing would loop + # forever without progress. + return None + return token + + +def _default_normalize_interaction(value: Any) -> Optional[Any]: + return value if isinstance(_get_value(value, 'steps'), list) else None + + +def _identity(value: Any) -> Any: + return value + + +def _interrupted_unary( + error: Exception, + hops: list[Any], + token: str, + *, + normalize_interaction: Callable[[Any], Optional[Any]], + finalize_interaction: Callable[[Any], Any], + stateful: bool = True, +) -> errors.ContinuationInterruptedError: + """Builds the error for a continuation hop that failed after `hops`.""" + partial = merge_unary_hops( + hops, + normalize_interaction=normalize_interaction, + finalize_interaction=finalize_interaction, + ) + interaction_id = ( + _non_empty_str(_get_value(hops[-1], 'id')) + if stateful and _get_value(hops[0], 'store') is not False + else None + ) + return errors.ContinuationInterruptedError( + f'Automatic continuation was interrupted after {len(hops)} completed' + f' hop(s): {error}. The partial interaction (status' + " 'requires_continuation') and its continuation_token can be used to" + ' resume.', + partial_interaction=partial, + continuation_token=token, + interaction_id=interaction_id, + continuation_ids=_merge_continuation_ids(hops), + ) + + +def continue_unary_interaction( + initial_response: Any, + create_next: Callable[[Optional[str], Optional[str]], Any], + *, + initial_continuation_token: Optional[str] = None, + normalize_interaction: Callable[ + [Any], Optional[Any] + ] = _default_normalize_interaction, + finalize_interaction: Callable[[Any], Any] = _identity, + deadline: Optional[CallDeadline] = None, + stateful: bool = True, +) -> Any: + """Continues a unary interaction across `requires_continuation` hops. + + A hop that fails after an earlier hop completed raises + `errors.ContinuationInterruptedError` (with the failure as `__cause__`) + carrying the hops merged so far and the token to resume from, so the output + already produced (and billed) is not lost. Running out of `deadline` before + a hop starts is such a failure (`ContinuationTimeoutError`). + """ + deadline = deadline or CallDeadline(None) + hops = [initial_response] + sent_token = initial_continuation_token + while True: + token = _next_unary_hop_token(hops[-1], sent_token) + if not token: + return merge_unary_hops( + hops, + normalize_interaction=normalize_interaction, + finalize_interaction=finalize_interaction, + ) + try: + deadline.check() + hop = create_next(token, _non_empty_str(_get_value(hops[-1], 'id'))) + except Exception as e: + raise _interrupted_unary( + e, + hops, + token, + normalize_interaction=normalize_interaction, + finalize_interaction=finalize_interaction, + stateful=stateful, + ) from e + hops.append(hop) + sent_token = token + + +async def continue_async_unary_interaction( + initial_response: Any, + create_next: Callable[[Optional[str], Optional[str]], Awaitable[Any]], + *, + initial_continuation_token: Optional[str] = None, + normalize_interaction: Callable[ + [Any], Optional[Any] + ] = _default_normalize_interaction, + finalize_interaction: Callable[[Any], Any] = _identity, + deadline: Optional[CallDeadline] = None, + stateful: bool = True, +) -> Any: + """Async `continue_unary_interaction`.""" + deadline = deadline or CallDeadline(None) + hops = [initial_response] + sent_token = initial_continuation_token + while True: + token = _next_unary_hop_token(hops[-1], sent_token) + if not token: + return merge_unary_hops( + hops, + normalize_interaction=normalize_interaction, + finalize_interaction=finalize_interaction, + ) + try: + deadline.check() + hop = await create_next( + token, _non_empty_str(_get_value(hops[-1], 'id')) + ) + except Exception as e: + raise _interrupted_unary( + e, + hops, + token, + normalize_interaction=normalize_interaction, + finalize_interaction=finalize_interaction, + stateful=stateful, + ) from e + hops.append(hop) + sent_token = token + + +def _strip_leading_user_input_steps(steps: list[Any]) -> list[Any]: + start = 0 + while start < len(steps) and _get_value(steps[start], 'type') == 'user_input': + start += 1 + return steps[start:] + + +def _merge_continuation_ids(hops: list[Any]) -> list[str]: + ids: list[str] = [] + for hop in hops: + hop_continuation_ids = _get_value(hop, 'continuation_ids') + if isinstance(hop_continuation_ids, list): + ids.extend( + cid for cid in hop_continuation_ids if isinstance(cid, str) and cid + ) + hop_id = _get_value(hop, 'id') + if isinstance(hop_id, str) and hop_id: + ids.append(hop_id) + return list(dict.fromkeys(ids)) + + +def merge_unary_hops( + hops: list[Any], + *, + normalize_interaction: Callable[ + [Any], Optional[Any] + ] = _default_normalize_interaction, + finalize_interaction: Callable[[Any], Any] = _identity, +) -> Any: + """Merges automatic continuation hops into a single interaction. + + Each hop is its own interaction on the server, and only contains the output + generated during that hop. The merged result is the final hop (so `id`, + `status`, etc. refer to it) with the steps of every hop concatenated, usage + summed across hops, and `continuation_ids` listing every hop's id. Fields + describing the turn's request (`previous_interaction_id`, `created`, `input`) + come from the first hop: later hops point back at the hop before them, not at + the caller's previous turn. + """ + final = hops[-1] + if normalize_interaction(final) is None: + return final + + steps: list[Any] = [] + for hop_index, hop in enumerate(hops): + hop_steps = _get_value(hop, 'steps') + if isinstance(hop_steps, list): + if hop_index > 0: + # Later hops continue the same turn; a repeated user_input step would + # duplicate the input and cut `output_text` (which stops at the latest + # user_input step) short. + hop_steps = _strip_leading_user_input_steps(hop_steps) + steps.extend(hop_steps) + + updates: dict[str, Any] = { + 'steps': steps, + 'continuation_ids': _merge_continuation_ids(hops), + } + if len(hops) > 1: + for name in ('previous_interaction_id', 'created', 'input'): + updates[name] = _get_value(hops[0], name) + usage = _merge_usages( + [u for u in (_get_value(hop, 'usage') for hop in hops) if u is not None], + first_hop_usage=_get_value(hops[0], 'usage'), + ) + if usage is not None: + updates['usage'] = usage + + return finalize_interaction(_set_or_copy(final, updates)) + + +def _usage_field_names(usage: Any) -> list[str]: + if isinstance(usage, Mapping): + return list(usage.keys()) + model_fields = getattr(type(usage), 'model_fields', None) + if model_fields is not None: + names = list(model_fields) + elif hasattr(usage, '__dict__'): + names = list(vars(usage)) + else: + names = [] + for cls in type(usage).__mro__: + slots = getattr(cls, '__slots__', ()) + if isinstance(slots, str): + slots = (slots,) + names.extend( + s + for s in slots + if s not in ('__dict__', '__weakref__') and s not in names + ) + extra = getattr(usage, '__pydantic_extra__', None) or {} + return names + [name for name in extra if name not in names] + + +def _is_int(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) + + +def _merge_usage_list( + lists: list[list[Any]], key_name: str, value_name: str +) -> list[Any]: + merged: dict[Any, Any] = {} + for items in lists: + for item in items: + key = _get_value(item, key_name) + value = _get_value(item, value_name) + if key not in merged: + merged[key] = item + continue + prev = _get_value(merged[key], value_name) + if _is_int(prev) and _is_int(value): + merged[key] = _set_or_copy(merged[key], {value_name: prev + value}) + return list(merged.values()) + + +_UNSET = object() + + +def _merge_usages( + usages: list[Any], *, first_hop_usage: Any = _UNSET +) -> Optional[Any]: + """Sums token usage across hops, keeping the first hop's input tokens.""" + if not usages: + return None + first = usages[0] if first_hop_usage is _UNSET else first_hop_usage + updates: dict[str, Any] = {} + names: list[str] = [] + for usage in usages: + names.extend(n for n in _usage_field_names(usage) if n not in names) + + for name in names: + if name == 'initial_input_tokens': + continue + values = [_get_value(u, name) for u in usages] + present = [v for v in values if v is not None] + if not present: + continue + if all(_is_int(v) for v in present): + updates[name] = sum(present) + elif name in _USAGE_LIST_KEYS and all(isinstance(v, list) for v in present): + updates[name] = _merge_usage_list(present, *_USAGE_LIST_KEYS[name]) + + initial_input_tokens = ( + _get_value(first, 'initial_input_tokens') if first is not None else None + ) + if initial_input_tokens is None and first is not None: + initial_input_tokens = _get_value(first, 'total_input_tokens') + if initial_input_tokens is not None or ( + _get_value(usages[-1], 'initial_input_tokens') is not None + ): + updates['initial_input_tokens'] = initial_input_tokens + + return _set_or_copy(usages[-1], updates) diff --git a/google/genai/errors.py b/google/genai/errors.py index d1f434b49..d96c3637d 100644 --- a/google/genai/errors.py +++ b/google/genai/errors.py @@ -334,4 +334,49 @@ class UnknownApiResponseError(ValueError): """Raised when the response from the API cannot be parsed as JSON.""" pass + +class ContinuationInterruptedError(Exception): + """Raised when Interactions automatic continuation fails after progress. + + `interactions.create()` with `automatic_continuation` enabled (the default) + may make several requests ("hops"), each one resuming the decode from the + previous hop's `continuation_token`. If a hop fails after the server already + produced (and billed) output in an earlier hop, or after a resumable + checkpoint was reached, this error is raised instead of the underlying + error so the work is not lost. The underlying error is `__cause__`. + + To resume, pass `continuation_token` to `interactions.create()`: with `store` + enabled (the default), also pass `interaction_id` as `previous_interaction_id` + and no `input`; with `store=False`, resend the original `input`. + + Attributes: + partial_interaction: For unary calls, the interaction merged from the hops + that completed (status `requires_continuation`, with the resumable + `continuation_token`), or None if no hop completed. Always None for + streams, whose events were already yielded. + continuation_token: The latest resumable `continuation_token`. For streams + this may be a mid-hop checkpoint: resuming from it re-generates any + output yielded after the checkpoint. + interaction_id: The id of the stored interaction that issued + `continuation_token` (the value to send as `previous_interaction_id` when + `store` is enabled), or None when `store=False` or unknown. + continuation_ids: The ids of every hop created during the call, in order. + """ + + def __init__( + self, + message: str, + *, + partial_interaction: Any = None, + continuation_token: Optional[str] = None, + interaction_id: Optional[str] = None, + continuation_ids: Optional[list[str]] = None, + ): + super().__init__(message) + self.partial_interaction = partial_interaction + self.continuation_token = continuation_token + self.interaction_id = interaction_id + self.continuation_ids = list(continuation_ids or []) + + ExperimentalWarning = _common.ExperimentalWarning diff --git a/google/genai/tests/interactions/test_automatic_continuation.py b/google/genai/tests/interactions/test_automatic_continuation.py new file mode 100644 index 000000000..ecd2846d3 --- /dev/null +++ b/google/genai/tests/interactions/test_automatic_continuation.py @@ -0,0 +1,2727 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for automatic_continuation in Interactions API.""" + +import json +import os +from typing import Any +from unittest import mock + +import httpx +import pytest + +from ... import _interactions_continuation +from ... import client as client_lib +from ... import errors +from ... import types +from ..._gaos import interactions as gaos_interactions +from ..._gaos.lib import compat_errors +from ..._gaos.models import createinteraction as create_interaction_models +from ..._gaos.types import interactions as interaction_types +from ..._gaos.utils import eventstreaming + +pytest_plugins = ("pytest_asyncio",) + +# Until `continuation_token` is promoted to the public generated schema, +# `run_tests.sh --private` rewrites the `_gaos` imports above to +# `_gaos_private`. Use `pytestmark` rather than +# `pytest.skip(..., allow_module_level=True)` so pytest collects and skips tests +# cleanly (exit code 0) instead of failing with exit code 5 (no tests collected) +# under Blaze / TAP when `continuation_token` is absent from public types. +pytestmark = pytest.mark.skipif( + "continuation_token" not in interaction_types.Interaction.model_fields, + reason="continuation_token is not yet in the generated Interaction schema.", +) + + +def _make_sync_stream( + events: list[Any], + exc: Exception | None = None, +) -> eventstreaming.Stream[Any]: + stream = eventstreaming.Stream.__new__(eventstreaming.Stream) + stream.response = mock.MagicMock(spec=httpx.Response) + stream.client_ref = None + stream._closed = False + + def _gen(): + try: + for ev in events: + yield ev + if exc is not None: + raise exc + finally: + stream.response.close() + + stream.generator = _gen() + return stream + + +def _make_async_stream( + events: list[Any], + exc: Exception | None = None, +) -> eventstreaming.AsyncStream[Any]: + stream = eventstreaming.AsyncStream.__new__(eventstreaming.AsyncStream) + stream.response = mock.MagicMock(spec=httpx.Response) + stream.response.aclose = mock.AsyncMock() + stream.client_ref = None + stream._closed = False + + async def _close(): + stream._closed = True + await stream.response.aclose() + + stream.close = _close + + async def _agen(): + try: + for ev in events: + yield ev + if exc is not None: + raise exc + finally: + await stream.response.aclose() + + stream.generator = _agen() + return stream + + +def _created_event(interaction_id: str = "int_123") -> dict[str, Any]: + return { + "event_type": "interaction.created", + "interaction": { + "id": interaction_id, + "status": "in_progress", + }, + } + + +def _text_delta_event(text: str, index: int = 0) -> dict[str, Any]: + return { + "event_type": "step.delta", + "index": index, + "delta": {"type": "text", "text": text}, + } + + +def _completed_event( + *, + status: str = "completed", + continuation_token: str | None = None, + output_tokens: int = 10, + thought_tokens: int = 0, + interaction_id: str = "int_123", +) -> dict[str, Any]: + interaction: dict[str, Any] = { + "id": interaction_id, + "status": status, + "usage": { + "total_output_tokens": output_tokens, + "total_thought_tokens": thought_tokens, + "total_tokens": output_tokens + thought_tokens, + }, + } + if continuation_token is not None: + interaction["continuation_token"] = continuation_token + return { + "event_type": "interaction.completed", + "interaction": interaction, + } + + +def test_stream_automatic_continuation_with_request_object(): + client = client_lib.Client(api_key="test-key") + hop1_events = [ + _created_event(), + _text_delta_event("Part 1 "), + _completed_event( + status="requires_continuation", + continuation_token="tok_req_1", + output_tokens=20, + ), + ] + hop2_events = [ + _created_event(), + _text_delta_event("Part 2"), + _completed_event( + status="completed", + output_tokens=20, + ), + ] + + req = create_interaction_models.CreateInteractionRequest( + body=interaction_types.CreateModelInteraction( + model="gemini-2.5-pro", + input="Hello", + stream=True, + ) + ) + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[ + _make_sync_stream(hop1_events), + _make_sync_stream(hop2_events), + ], + ) as mock_create: + stream = client.interactions.create( + request=req, + automatic_continuation=True, + ) + received = list(stream) + + assert received == hop1_events + hop2_events + assert mock_create.call_count == 2 + second_req = mock_create.call_args_list[1].kwargs["request"] + assert second_req.body.continuation_token == "tok_req_1" + # store defaults to on: the hop continues the stored hop that issued the + # token instead of resending the input. + assert second_req.body.previous_interaction_id == "int_123" + assert second_req.body.input is None + assert "input" not in second_req.body.model_dump() + assert req.body.continuation_token is None + assert req.body.input == "Hello" + + +def test_stream_automatic_continuation_close_on_second_hop(): + client = client_lib.Client(api_key="test-key") + hop1_events = [ + _created_event(), + _completed_event( + status="requires_continuation", + continuation_token="tok_hop1", + output_tokens=10, + ), + ] + hop2_events = [ + _created_event(), + _text_delta_event("chunk 1"), + _text_delta_event("chunk 2"), + _completed_event(status="completed", output_tokens=10), + ] + hop1_stream = _make_sync_stream(hop1_events) + hop2_stream = _make_sync_stream(hop2_events) + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[hop1_stream, hop2_stream], + ): + with client.interactions.create( + model="gemini-2.5-pro", + input="Hello", + stream=True, + automatic_continuation=True, + ) as stream: + # Consume hop 1 (2 events) + first event of hop 2 + _ = next(stream) + _ = next(stream) + _ = next(stream) + # Exit context manager early on hop 2 + + hop1_stream.response.close.assert_called() + hop2_stream.response.close.assert_called() + + +@pytest.mark.asyncio +async def test_async_stream_automatic_continuation_close_on_second_hop(): + client = client_lib.Client(api_key="test-key") + hop1_events = [ + _created_event(), + _completed_event( + status="requires_continuation", + continuation_token="tok_hop1", + output_tokens=10, + ), + ] + hop2_events = [ + _created_event(), + _text_delta_event("chunk 1"), + _text_delta_event("chunk 2"), + _completed_event(status="completed", output_tokens=10), + ] + hop1_stream = _make_async_stream(hop1_events) + hop2_stream = _make_async_stream(hop2_events) + + with mock.patch.object( + gaos_interactions.AsyncInteractions, + "create", + new=mock.AsyncMock(side_effect=[hop1_stream, hop2_stream]), + ): + async with await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Hello", + stream=True, + automatic_continuation=True, + ) as stream: + _ = await stream.__anext__() + _ = await stream.__anext__() + _ = await stream.__anext__() + + hop1_stream.response.aclose.assert_awaited() + hop2_stream.response.aclose.assert_awaited() + + +@pytest.mark.asyncio +async def test_async_stream_automatic_continuation_three_hops(): + client = client_lib.Client(api_key="test-key") + hop1_events = [ + _created_event(), + _text_delta_event("Hop 1 "), + _completed_event( + status="requires_continuation", + continuation_token="tok_1", + output_tokens=50, + ), + ] + hop2_events = [ + _created_event(), + _text_delta_event("Hop 2 "), + _completed_event( + status="requires_continuation", + continuation_token="tok_2", + output_tokens=60, + ), + ] + hop3_events = [ + _created_event(), + _text_delta_event("Hop 3"), + _completed_event( + status="completed", + output_tokens=30, + ), + ] + + with mock.patch.object( + gaos_interactions.AsyncInteractions, + "create", + new=mock.AsyncMock( + side_effect=[ + _make_async_stream(hop1_events), + _make_async_stream(hop2_events), + _make_async_stream(hop3_events), + ] + ), + ) as mock_create: + stream = await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Write a long story", + stream=True, + automatic_continuation=True, + ) + received = [ev async for ev in stream] + + assert received == hop1_events + hop2_events + hop3_events + assert mock_create.call_count == 3 + assert "continuation_token" not in mock_create.call_args_list[0].kwargs + assert mock_create.call_args_list[1].kwargs["continuation_token"] == "tok_1" + assert mock_create.call_args_list[2].kwargs["continuation_token"] == "tok_2" + for call in mock_create.call_args_list[1:]: + assert call.kwargs["previous_interaction_id"] == "int_123" + assert "input" not in call.kwargs + + +def test_automatic_continuation_false_with_background_raises_value_error(): + client = client_lib.Client(api_key="test-key") + with pytest.raises( + ValueError, + match="automatic_continuation=False cannot be used with background=True", + ): + client.interactions.create( + model="gemini-2.5-pro", + input="Hello", + background=True, + automatic_continuation=False, + ) + + with pytest.raises( + ValueError, + match="automatic_continuation=False cannot be used with background=True", + ): + client.interactions.create( + model="gemini-2.5-pro", + input="Hello", + background=True, + stream=True, + automatic_continuation=False, + ) + + with pytest.raises( + ValueError, + match="automatic_continuation=False cannot be used with background=True", + ): + client.interactions.create( + request={"body": {"model": "gemini-2.5-pro", "background": True}}, + automatic_continuation=False, + ) + + +@pytest.mark.asyncio +async def test_async_automatic_continuation_false_with_background_raises_value_error(): + client = client_lib.Client(api_key="test-key") + with pytest.raises( + ValueError, + match="automatic_continuation=False cannot be used with background=True", + ): + await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Hello", + background=True, + automatic_continuation=False, + ) + + with pytest.raises( + ValueError, + match="automatic_continuation=False cannot be used with background=True", + ): + await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Hello", + background=True, + stream=True, + automatic_continuation=False, + ) + + +def test_unary_automatic_continuation_with_store_false_merges_hops(): + client = client_lib.Client(api_key="test-key") + hop1 = _unary_hop( + interaction_id="int_sf_1", + status="requires_continuation", + text="Stateless 1 ", + continuation_token="tok_sf_1", + ) + hop2 = _unary_hop( + interaction_id="int_sf_2", + status="completed", + text="Stateless 2", + ) + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[hop1, hop2], + ) as mock_create: + res = client.interactions.create( + model="gemini-2.5-pro", + input="Hello", + store=False, + automatic_continuation=True, + ) + + assert mock_create.call_count == 2 + assert mock_create.call_args_list[1].kwargs["store"] is False + assert ( + mock_create.call_args_list[1].kwargs["continuation_token"] == "tok_sf_1" + ) + # Nothing is stored, so the input is resent without previous_interaction_id. + assert mock_create.call_args_list[1].kwargs["input"] == "Hello" + assert "previous_interaction_id" not in mock_create.call_args_list[1].kwargs + assert res.output_text == "Stateless 1 Stateless 2" + assert res.continuation_ids == ["int_sf_1", "int_sf_2"] + + +@pytest.mark.asyncio +async def test_async_unary_automatic_continuation_with_store_false_merges_hops(): + client = client_lib.Client(api_key="test-key") + hop1 = _unary_hop( + interaction_id="int_sf_1", + status="requires_continuation", + text="Stateless 1 ", + continuation_token="tok_sf_1", + ) + hop2 = _unary_hop( + interaction_id="int_sf_2", + status="completed", + text="Stateless 2", + ) + + with mock.patch.object( + gaos_interactions.AsyncInteractions, + "create", + new=mock.AsyncMock(side_effect=[hop1, hop2]), + ) as mock_create: + res = await client.aio.interactions.create( + request={"body": {"model": "gemini-2.5-pro", "store": False}}, + automatic_continuation=True, + ) + + assert mock_create.call_count == 2 + second_req = mock_create.call_args_list[1].kwargs["request"] + assert second_req["body"]["continuation_token"] == "tok_sf_1" + assert res.output_text == "Stateless 1 Stateless 2" + assert res.continuation_ids == ["int_sf_1", "int_sf_2"] + + +def test_output_text_skips_thought_steps_but_stops_at_tool_steps(): + client = client_lib.Client(api_key="test-key") + interaction = interaction_types.Interaction( + id="int_tool_thought", + status="completed", + steps=[ + interaction_types.ThoughtStep(type="thought", signature="sig_0"), + interaction_types.ModelOutputStep( + type="model_output", + content=[ + interaction_types.TextContent( + type="text", text="Pre-tool chatter. " + ) + ], + ), + interaction_types.FunctionCallStep( + type="function_call", + id="call_1", + name="search", + arguments={"q": "test"}, + ), + interaction_types.FunctionResultStep( + type="function_result", + call_id="call_1", + result="found", + ), + interaction_types.ThoughtStep(type="thought", signature="sig_1"), + interaction_types.ModelOutputStep( + type="model_output", + content=[ + interaction_types.TextContent( + type="text", text="Post-tool hop 1. " + ) + ], + ), + interaction_types.ThoughtStep(type="thought", signature="sig_2"), + interaction_types.ModelOutputStep( + type="model_output", + content=[ + interaction_types.TextContent( + type="text", text="Post-tool hop 2." + ) + ], + ), + ], + ) + with mock.patch.object( + gaos_interactions.Interactions, + "create", + return_value=interaction, + ): + res = client.interactions.create( + model="gemini-2.5-pro", + input="Search and summarize", + ) + assert res.output_text == "Post-tool hop 1. Post-tool hop 2." + + +def _unary_hop( + *, + status: str, + text: str, + interaction_id: str = "int_unary_123", + continuation_token: str | None = None, + input_tokens: int = 10, + output_tokens: int = 20, + thought_tokens: int = 0, + created: str = "2026-01-01T00:00:00Z", + updated: str = "2026-01-01T00:00:05Z", +) -> interaction_types.Interaction: + """Builds a single hop response. + + Each continuation hop is its own interaction on the server, and only contains + the output (and usage) generated during that hop. + """ + return interaction_types.Interaction( + id=interaction_id, + status=status, + continuation_token=continuation_token, + created=created, + updated=updated, + steps=[ + interaction_types.ModelOutputStep( + type="model_output", + content=[interaction_types.TextContent(type="text", text=text)], + ) + ], + usage=interaction_types.Usage( + total_input_tokens=input_tokens, + total_output_tokens=output_tokens, + total_thought_tokens=thought_tokens, + total_tokens=input_tokens + output_tokens + thought_tokens, + ), + ) + + +def test_unary_automatic_continuation_with_request_object(): + client = client_lib.Client(api_key="test-key") + hop1 = _unary_hop( + interaction_id="int_req_1", + status="requires_continuation", + text="Part A ", + continuation_token="tok_req_1", + input_tokens=10, + output_tokens=40, + thought_tokens=5, + ) + hop2 = _unary_hop( + interaction_id="int_req_2", + status="completed", + text="Part B", + continuation_token=None, + input_tokens=50, + output_tokens=30, + ) + + req = create_interaction_models.CreateInteractionRequest( + body=interaction_types.CreateModelInteraction( + model="gemini-2.5-pro", + input="Hello", + ) + ) + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[hop1, hop2], + ) as mock_create: + res = client.interactions.create( + request=req, + automatic_continuation=True, + ) + + assert mock_create.call_count == 2 + second_req = mock_create.call_args_list[1].kwargs["request"] + assert second_req.body.continuation_token == "tok_req_1" + # store defaults to on: the hop continues the stored hop that issued the + # token instead of resending the input. + assert second_req.body.previous_interaction_id == "int_req_1" + assert second_req.body.input is None + assert "input" not in second_req.body.model_dump() + assert req.body.continuation_token is None + assert req.body.input == "Hello" + assert res.status == "completed" + assert res.output_text == "Part A Part B" + # Mental model 2: the result is the final hop, with the hops merged. + assert res.id == "int_req_2" + assert res.continuation_token is None + assert res.continuation_ids == ["int_req_1", "int_req_2"] + assert [s.content[0].text for s in res.steps] == ["Part A ", "Part B"] + assert res.usage.total_input_tokens == 60 + assert res.usage.total_output_tokens == 70 + assert res.usage.total_thought_tokens == 5 + assert res.usage.total_tokens == 135 + assert res.usage.initial_input_tokens == 10 + + +_RAW_REQUIRES_CONTINUATION_BODY = { + "id": "int_raw_1", + "status": "requires_continuation", + "continuation_token": "dG9rX3Jhd18x", + "previous_interaction_id": "int_turn_0", + "steps": [ + {"type": "user_input", "content": [{"type": "text", "text": "Hello"}]}, + {"type": "model_output", "content": [{"type": "text", "text": "A"}]}, + ], + "usage": { + "total_input_tokens": 10, + "total_output_tokens": 5, + "total_tokens": 15, + }, +} + + +def _wire_client( + handler, async_handler=None +) -> tuple[client_lib.Client, list[dict[str, Any]]]: + """A client whose HTTP transport is `handler`; records request bodies.""" + bodies: list[dict[str, Any]] = [] + + def _sync(request: httpx.Request) -> httpx.Response: + bodies.append(json.loads(request.content.decode())) + return handler(request) + + async def _async(request: httpx.Request) -> httpx.Response: + bodies.append(json.loads(request.content.decode())) + return (async_handler or handler)(request) + + client = client_lib.Client( + api_key="test-key", + http_options=types.HttpOptions( + httpx_client=httpx.Client(transport=httpx.MockTransport(_sync)), + httpx_async_client=httpx.AsyncClient( + transport=httpx.MockTransport(_async) + ), + ), + ) + return client, bodies + + +def _raw_json_response(_: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "application/json"}, + text=json.dumps(_RAW_REQUIRES_CONTINUATION_BODY), + ) + + +def _assert_opt_out_passthrough(res: Any, bodies: list[dict[str, Any]]): + # Exactly one request, sent as the caller wrote it: no continuation fields + # added. + assert len(bodies) == 1 + assert "generation_config" not in bodies[0] + assert "continuation_token" not in bodies[0] + assert bodies[0]["previous_interaction_id"] == "int_turn_0" + assert bodies[0]["input"] == "Hello" + # The response is the server body field-for-field (opted-out callers drive + # continuation themselves, so continuation_token / status must survive). + serialized = res.model_dump(mode="json", by_alias=True, exclude_none=True) + for key, value in _RAW_REQUIRES_CONTINUATION_BODY.items(): + assert serialized.get(key) == value, key + assert "continuation_ids" not in serialized + assert res.continuation_ids is None + assert res.continuation_token == "dG9rX3Jhd18x" + assert res.status == "requires_continuation" + + +@pytest.mark.filterwarnings("ignore:Interaction stopped before completion") +def test_unary_automatic_continuation_disabled_with_request_object_is_passthrough(): + client, bodies = _wire_client(_raw_json_response) + req = create_interaction_models.CreateInteractionRequest( + body=interaction_types.CreateModelInteraction( + model="gemini-2.5-pro", + input="Hello", + previous_interaction_id="int_turn_0", + ) + ) + res = client.interactions.create(request=req, automatic_continuation=False) + _assert_opt_out_passthrough(res, bodies) + # The caller's request object is not mutated either. + assert req.body.generation_config is None + assert req.body.continuation_token is None + + +@pytest.mark.asyncio +@pytest.mark.filterwarnings("ignore:Interaction stopped before completion") +async def test_async_unary_automatic_continuation_disabled_with_request_object_is_passthrough(): + client, bodies = _wire_client(_raw_json_response) + req = create_interaction_models.CreateInteractionRequest( + body=interaction_types.CreateModelInteraction( + model="gemini-2.5-pro", + input="Hello", + previous_interaction_id="int_turn_0", + ) + ) + res = await client.aio.interactions.create( + request=req, automatic_continuation=False + ) + _assert_opt_out_passthrough(res, bodies) + assert req.body.generation_config is None + assert req.body.continuation_token is None + + +def test_unary_automatic_continuation_merges_usage_by_modality(): + client = client_lib.Client(api_key="test-key") + + def _hop(interaction_id, status, token, text_tokens, image_tokens=None): + by_modality = [ + interaction_types.ModalityTokens(modality="text", tokens=text_tokens) + ] + if image_tokens is not None: + by_modality.append( + interaction_types.ModalityTokens( + modality="image", tokens=image_tokens + ) + ) + return interaction_types.Interaction( + id=interaction_id, + status=status, + continuation_token=token, + steps=[], + usage=interaction_types.Usage( + total_output_tokens=sum(m.tokens for m in by_modality), + output_tokens_by_modality=by_modality, + ), + ) + + hop1 = _hop("int_m1", "requires_continuation", "tok_m1", 10, image_tokens=5) + hop2 = _hop("int_m2", "completed", None, 7) + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[hop1, hop2], + ): + res = client.interactions.create( + model="gemini-2.5-pro", + input="Hello", + automatic_continuation=True, + ) + + assert res.usage.total_output_tokens == 22 + assert { + m.modality: m.tokens for m in res.usage.output_tokens_by_modality + } == {"text": 17, "image": 5} + # The hop responses are not mutated. + assert hop1.usage.output_tokens_by_modality[0].tokens == 10 + + +@pytest.mark.asyncio +async def test_async_unary_automatic_continuation_three_hops(): + client = client_lib.Client(api_key="test-key") + hop1 = _unary_hop( + interaction_id="int_a1", + status="requires_continuation", + text="Async 1 ", + continuation_token="tok_1", + output_tokens=50, + ) + hop2 = _unary_hop( + interaction_id="int_a2", + status="requires_continuation", + text="Async 2 ", + continuation_token="tok_2", + output_tokens=60, + ) + hop3 = _unary_hop( + interaction_id="int_a3", + status="completed", + text="Async 3", + continuation_token=None, + output_tokens=30, + ) + + with mock.patch.object( + gaos_interactions.AsyncInteractions, + "create", + new=mock.AsyncMock(side_effect=[hop1, hop2, hop3]), + ) as mock_create: + res = await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Write a long story", + automatic_continuation=True, + ) + + assert mock_create.call_count == 3 + assert "continuation_token" not in mock_create.call_args_list[0].kwargs + assert mock_create.call_args_list[1].kwargs["continuation_token"] == "tok_1" + assert mock_create.call_args_list[2].kwargs["continuation_token"] == "tok_2" + assert res.status == "completed" + assert res.output_text == "Async 1 Async 2 Async 3" + assert res.id == "int_a3" + assert res.continuation_ids == ["int_a1", "int_a2", "int_a3"] + assert len(res.steps) == 3 + assert res.usage.total_output_tokens == 140 + assert res.usage.total_input_tokens == 30 + assert res.usage.initial_input_tokens == 10 + + +def test_unary_automatic_continuation_stops_when_token_missing_or_empty(): + client = client_lib.Client(api_key="test-key") + hop_no_token = _unary_hop( + status="requires_continuation", + text="Continuation required without token", + continuation_token=None, + ) + hop_empty_token = _unary_hop( + status="requires_continuation", + text="Continuation required with empty token", + continuation_token="", + ) + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[hop_no_token, hop_empty_token], + ) as mock_create: + res1 = client.interactions.create( + model="gemini-2.5-pro", + input="Test 1", + automatic_continuation=True, + ) + res2 = client.interactions.create( + model="gemini-2.5-pro", + input="Test 2", + automatic_continuation=True, + ) + + assert mock_create.call_count == 2 + assert res1.status == "requires_continuation" + assert res1.output_text == "Continuation required without token" + assert res2.status == "requires_continuation" + assert res2.output_text == "Continuation required with empty token" + assert res1.continuation_ids == ["int_unary_123"] + assert res1.usage.total_output_tokens == 20 + + +def test_unary_automatic_continuation_with_initial_user_continuation_token(): + client = client_lib.Client(api_key="test-key") + hop1 = _unary_hop( + interaction_id="int_resumed_1", + status="requires_continuation", + text="Resumed hop 1 ", + continuation_token="tok_from_hop1", + ) + hop2 = _unary_hop( + interaction_id="int_resumed_2", + status="completed", + text="Resumed hop 2", + continuation_token=None, + ) + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[hop1, hop2], + ) as mock_create: + res = client.interactions.create( + model="gemini-2.5-pro", + input="Resume", + continuation_token="user_initial_tok", + automatic_continuation=True, + ) + + assert mock_create.call_count == 2 + assert ( + mock_create.call_args_list[0].kwargs["continuation_token"] + == "user_initial_tok" + ) + assert ( + mock_create.call_args_list[1].kwargs["continuation_token"] + == "tok_from_hop1" + ) + assert res.status == "completed" + assert res.output_text == "Resumed hop 1 Resumed hop 2" + assert res.continuation_ids == ["int_resumed_1", "int_resumed_2"] + + +def test_stream_automatic_continuation_stops_on_requires_action_or_missing_token(): + client = client_lib.Client(api_key="test-key") + requires_action_events = [ + _created_event(), + _text_delta_event("Calling tool..."), + _completed_event( + status="requires_action", + continuation_token="tok_should_not_continue", + output_tokens=20, + ), + ] + missing_token_events = [ + _created_event(), + _text_delta_event("Partial..."), + _completed_event( + status="requires_continuation", + continuation_token=None, + output_tokens=20, + ), + ] + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[ + _make_sync_stream(requires_action_events), + _make_sync_stream(missing_token_events), + ], + ) as mock_create: + stream1 = client.interactions.create( + model="gemini-2.5-pro", + input="Tool call", + stream=True, + automatic_continuation=True, + ) + events1 = list(stream1) + + stream2 = client.interactions.create( + model="gemini-2.5-pro", + input="Missing token", + stream=True, + automatic_continuation=True, + ) + events2 = list(stream2) + + assert mock_create.call_count == 2 + assert events1 == requires_action_events + assert events2 == missing_token_events + + +def test_stream_automatic_continuation_terminal_event_without_token_clears_checkpoint(): + client = client_lib.Client(api_key="test-key") + hop1_events = [ + _created_event(), + _text_delta_event("Chunk 1 "), + { + "event_type": "interaction.status_update", + "status": "in_progress", + "continuation_token": "tok_checkpoint_1", + }, + _completed_event( + status="requires_continuation", + continuation_token=None, + output_tokens=40, + ), + ] + hop2_events = [ + _created_event(), + _text_delta_event("Chunk 2"), + _completed_event( + status="completed", + continuation_token=None, + output_tokens=20, + ), + ] + + # The terminal interaction.completed event is authoritative: without a + # token it clears the stale mid-hop checkpoint, so the stream stops (same + # as unary) instead of resuming and re-generating already-received content. + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[ + _make_sync_stream(hop1_events), + _make_sync_stream(hop2_events), + ], + ) as mock_create: + stream = client.interactions.create( + model="gemini-2.5-pro", + input="Checkpoint test", + stream=True, + automatic_continuation=True, + ) + received = list(stream) + + assert received == hop1_events + assert mock_create.call_count == 1 + + +def test_stream_automatic_continuation_resumes_on_dead_connection_with_last_token(): + client = client_lib.Client(api_key="test-key") + hop1_events_before_drop = [ + _created_event(), + _text_delta_event("Chunk 1 "), + { + "event_type": "interaction.status_update", + "status": "in_progress", + "continuation_token": "tok_checkpoint_1", + }, + _text_delta_event("Chunk 2 "), + { + "event_type": "interaction.status_update", + "status": "in_progress", + "continuation_token": "tok_checkpoint_2", + }, + ] + hop2_events = [ + _created_event(), + _text_delta_event("Chunk 3"), + _completed_event( + status="completed", + continuation_token=None, + output_tokens=30, + ), + ] + + # 1. Dead connection mid-stream (transport error or premature EOF) after + # checkpoint tokens -> reconnects and sends the last continuation_token. + for disconnect_exc in ( + httpx.RemoteProtocolError("peer closed connection"), + httpx.ReadTimeout("stream read timed out"), + None, + ): + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[ + _make_sync_stream(hop1_events_before_drop, exc=disconnect_exc), + _make_sync_stream(hop2_events), + ], + ) as mock_create: + stream = client.interactions.create( + model="gemini-2.5-pro", + input="Dead connection resume test", + stream=True, + automatic_continuation=True, + ) + received = list(stream) + + assert received == hop1_events_before_drop + hop2_events + assert mock_create.call_count == 2 + assert "continuation_token" not in mock_create.call_args_list[0].kwargs + assert ( + mock_create.call_args_list[1].kwargs["continuation_token"] + == "tok_checkpoint_2" + ) + + # 2. Dead connection before any continuation_token in the hop -> raises error. + events_without_token = [ + _created_event(), + _text_delta_event("Early chunk"), + ] + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[ + _make_sync_stream( + events_without_token, + exc=httpx.RemoteProtocolError("early disconnect"), + ), + ], + ) as mock_create: + stream = client.interactions.create( + model="gemini-2.5-pro", + input="Early dead connection test", + stream=True, + automatic_continuation=True, + ) + with pytest.raises(compat_errors.APIConnectionError): + _ = list(stream) + assert mock_create.call_count == 1 + + # 3. Dead connection when automatic_continuation=False -> raises even if + # checkpoint token was received. + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[ + _make_sync_stream( + hop1_events_before_drop, + exc=httpx.RemoteProtocolError( + "disconnect with auto_continue=False" + ), + ), + ], + ) as mock_create: + stream = client.interactions.create( + model="gemini-2.5-pro", + input="No auto-continue dead connection test", + stream=True, + automatic_continuation=False, + ) + with pytest.raises(compat_errors.APIConnectionError): + _ = list(stream) + assert mock_create.call_count == 1 + + # 4. Repeated failure on the resumed hop before a new token (or any content) + # is emitted -> retries the hop's own token up to the no-progress cap, then + # raises ContinuationInterruptedError (the checkpoint token is resumable) + # with the last disconnect as the cause. + max_retries = ( + _interactions_continuation.MAX_STREAM_RECONNECTS_WITHOUT_PROGRESS + ) + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[ + _make_sync_stream( + hop1_events_before_drop, + exc=httpx.RemoteProtocolError("hop 1 disconnect"), + ), + ] + + [ + _make_sync_stream( + [_created_event()], + exc=httpx.RemoteProtocolError("hop 2 disconnect before token"), + ) + for _ in range(max_retries + 1) + ], + ) as mock_create: + stream = client.interactions.create( + model="gemini-2.5-pro", + input="Consecutive dead connection test", + stream=True, + automatic_continuation=True, + ) + with pytest.raises(errors.ContinuationInterruptedError) as exc_info: + _ = list(stream) + assert isinstance(exc_info.value.__cause__, compat_errors.APIConnectionError) + assert exc_info.value.continuation_token == "tok_checkpoint_2" + assert exc_info.value.partial_interaction is None + assert mock_create.call_count == max_retries + 2 + assert [ + c.kwargs["continuation_token"] for c in mock_create.call_args_list[1:] + ] == ["tok_checkpoint_2"] * (max_retries + 1) + + # 5. Server SSE error event after a checkpoint token -> stops cleanly without + # reconnecting. + events_with_sse_error = hop1_events_before_drop + [ + { + "event_type": "error", + "error": {"code": "resource_exhausted", "message": "quota exceeded"}, + }, + ] + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[_make_sync_stream(events_with_sse_error)], + ) as mock_create: + stream = client.interactions.create( + model="gemini-2.5-pro", + input="SSE error event after checkpoint", + stream=True, + automatic_continuation=True, + ) + received = list(stream) + assert received == events_with_sse_error + assert mock_create.call_count == 1 + + +@pytest.mark.asyncio +async def test_async_stream_automatic_continuation_resumes_on_dead_connection_with_last_token(): + client = client_lib.Client(api_key="test-key") + hop1_events_before_drop = [ + _created_event(), + _text_delta_event("Async Chunk 1 "), + { + "event_type": "interaction.status_update", + "status": "in_progress", + "continuation_token": "async_tok_ckpt_1", + }, + _text_delta_event("Async Chunk 2 "), + { + "event_type": "interaction.status_update", + "status": "in_progress", + "continuation_token": "async_tok_ckpt_2", + }, + ] + hop2_events = [ + _created_event(), + _text_delta_event("Async Chunk 3"), + _completed_event( + status="completed", + continuation_token=None, + output_tokens=30, + ), + ] + + for disconnect_exc in ( + httpx.RemoteProtocolError("async peer closed connection"), + httpx.ReadTimeout("async stream read timed out"), + None, + ): + with mock.patch.object( + gaos_interactions.AsyncInteractions, + "create", + new=mock.AsyncMock( + side_effect=[ + _make_async_stream(hop1_events_before_drop, exc=disconnect_exc), + _make_async_stream(hop2_events), + ] + ), + ) as mock_create: + stream = await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Async dead connection resume test", + stream=True, + automatic_continuation=True, + ) + received = [ev async for ev in stream] + + assert received == hop1_events_before_drop + hop2_events + assert mock_create.call_count == 2 + assert "continuation_token" not in mock_create.call_args_list[0].kwargs + assert ( + mock_create.call_args_list[1].kwargs["continuation_token"] + == "async_tok_ckpt_2" + ) + + events_without_token = [ + _created_event(), + _text_delta_event("Early async chunk"), + ] + with mock.patch.object( + gaos_interactions.AsyncInteractions, + "create", + new=mock.AsyncMock( + side_effect=[ + _make_async_stream( + events_without_token, + exc=httpx.RemoteProtocolError("early async disconnect"), + ), + ] + ), + ) as mock_create: + stream = await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Early async dead connection test", + stream=True, + automatic_continuation=True, + ) + with pytest.raises(compat_errors.APIConnectionError): + _ = [ev async for ev in stream] + assert mock_create.call_count == 1 + + events_with_sse_error = hop1_events_before_drop + [ + { + "event_type": "error", + "error": {"code": "resource_exhausted", "message": "quota exceeded"}, + }, + ] + with mock.patch.object( + gaos_interactions.AsyncInteractions, + "create", + new=mock.AsyncMock( + side_effect=[_make_async_stream(events_with_sse_error)] + ), + ) as mock_create: + stream = await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Async SSE error event after checkpoint", + stream=True, + automatic_continuation=True, + ) + received = [ev async for ev in stream] + assert received == events_with_sse_error + assert mock_create.call_count == 1 + + +def _status_update_event(continuation_token: str) -> dict[str, Any]: + return { + "event_type": "interaction.status_update", + "status": "in_progress", + "continuation_token": continuation_token, + } + + +def test_stream_automatic_continuation_drop_before_checkpoint_retries_hop_token(): + client = client_lib.Client(api_key="test-key") + hop1_events = [ + _created_event(), + _text_delta_event("Part 1 "), + _completed_event( + status="requires_continuation", + continuation_token="tok_hop1", + ), + ] + # Hop 2 drops before its first checkpoint, and before yielding any output + # (output since the last checkpoint can't be replayed without duplicating + # it; see the conformance spec). + hop2_events_before_drop = [ + _created_event(), + ] + hop2_retry_events = [ + _created_event(), + _text_delta_event("Part 2"), + _completed_event(status="completed"), + ] + + for disconnect_exc in ( + httpx.RemoteProtocolError("peer closed connection"), + None, + ): + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[ + _make_sync_stream(hop1_events), + _make_sync_stream(hop2_events_before_drop, exc=disconnect_exc), + _make_sync_stream(hop2_retry_events), + ], + ) as mock_create: + stream = client.interactions.create( + model="gemini-2.5-pro", + input="Drop before checkpoint test", + stream=True, + automatic_continuation=True, + ) + received = list(stream) + + assert received == hop1_events + hop2_events_before_drop + hop2_retry_events + assert mock_create.call_count == 3 + assert "continuation_token" not in mock_create.call_args_list[0].kwargs + assert mock_create.call_args_list[1].kwargs["continuation_token"] == ( + "tok_hop1" + ) + assert mock_create.call_args_list[2].kwargs["continuation_token"] == ( + "tok_hop1" + ) + + +def test_stream_automatic_continuation_drop_before_checkpoint_retries_caller_token(): + client = client_lib.Client(api_key="test-key") + events_before_drop = [_created_event()] + retry_events = [ + _created_event(), + _text_delta_event("done"), + _completed_event(status="completed"), + ] + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[ + _make_sync_stream( + events_before_drop, + exc=httpx.RemoteProtocolError("peer closed connection"), + ), + _make_sync_stream(retry_events), + ], + ) as mock_create: + stream = client.interactions.create( + model="gemini-2.5-pro", + input="Caller token drop test", + stream=True, + continuation_token="tok_caller", + automatic_continuation=True, + ) + received = list(stream) + + assert received == events_before_drop + retry_events + assert mock_create.call_count == 2 + assert [ + c.kwargs["continuation_token"] for c in mock_create.call_args_list + ] == ["tok_caller", "tok_caller"] + + +def test_stream_automatic_continuation_checkpoint_progress_resets_reconnect_cap(): + client = client_lib.Client(api_key="test-key") + max_retries = ( + _interactions_continuation.MAX_STREAM_RECONNECTS_WITHOUT_PROGRESS + ) + num_drops = max_retries + 2 + dropping_streams = [ + _make_sync_stream( + [_created_event(), _status_update_event(f"tok_ckpt_{i}")], + exc=httpx.RemoteProtocolError(f"drop {i}"), + ) + for i in range(num_drops) + ] + final_events = [_created_event(), _completed_event(status="completed")] + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=dropping_streams + [_make_sync_stream(final_events)], + ) as mock_create: + stream = client.interactions.create( + model="gemini-2.5-pro", + input="Progress resets cap test", + stream=True, + automatic_continuation=True, + ) + received = list(stream) + + assert received[-len(final_events) :] == final_events + assert mock_create.call_count == num_drops + 1 + assert [ + c.kwargs["continuation_token"] for c in mock_create.call_args_list[1:] + ] == [f"tok_ckpt_{i}" for i in range(num_drops)] + + +@pytest.mark.asyncio +async def test_async_stream_automatic_continuation_drop_before_checkpoint_retries_hop_token(): + client = client_lib.Client(api_key="test-key") + hop1_events = [ + _created_event(), + _completed_event( + status="requires_continuation", + continuation_token="async_tok_hop1", + ), + ] + hop2_events_before_drop = [_created_event()] + hop2_retry_events = [ + _created_event(), + _text_delta_event("Part 2"), + _completed_event(status="completed"), + ] + + with mock.patch.object( + gaos_interactions.AsyncInteractions, + "create", + new=mock.AsyncMock( + side_effect=[ + _make_async_stream(hop1_events), + _make_async_stream( + hop2_events_before_drop, + exc=httpx.RemoteProtocolError("async peer closed"), + ), + _make_async_stream(hop2_retry_events), + ] + ), + ) as mock_create: + stream = await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Async drop before checkpoint test", + stream=True, + automatic_continuation=True, + ) + received = [ev async for ev in stream] + + assert received == hop1_events + hop2_events_before_drop + hop2_retry_events + assert mock_create.call_count == 3 + assert [ + c.kwargs["continuation_token"] for c in mock_create.call_args_list[1:] + ] == ["async_tok_hop1", "async_tok_hop1"] + + +@pytest.mark.asyncio +async def test_async_stream_automatic_continuation_reconnect_cap_raises(): + client = client_lib.Client(api_key="test-key") + max_retries = ( + _interactions_continuation.MAX_STREAM_RECONNECTS_WITHOUT_PROGRESS + ) + + with mock.patch.object( + gaos_interactions.AsyncInteractions, + "create", + new=mock.AsyncMock( + side_effect=[ + _make_async_stream( + [_created_event()], + exc=httpx.RemoteProtocolError(f"async drop {i}"), + ) + for i in range(max_retries + 1) + ] + ), + ) as mock_create: + stream = await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Async reconnect cap test", + stream=True, + continuation_token="async_tok_caller", + automatic_continuation=True, + ) + with pytest.raises(compat_errors.APIConnectionError): + _ = [ev async for ev in stream] + + assert mock_create.call_count == max_retries + 1 + assert [ + c.kwargs["continuation_token"] for c in mock_create.call_args_list + ] == ["async_tok_caller"] * (max_retries + 1) + + +def test_stream_continuation_state_machine_resets_hop_fields_and_tracks_owner(): + state = _interactions_continuation.StreamContinuation( + connection_error_type=compat_errors.APIConnectionError, + ) + + # Hop 1: initial hop emits content, then a checkpoint token, then drops before + # any post-checkpoint content. + assert state.start_hop(stream_opened=True) == (None, None) + state.observe(_created_event(interaction_id="hop_1")) + state.observe(_text_delta_event("hello")) + state.observe(_status_update_event("tok_ckpt_1")) + assert state.end_hop(httpx.RemoteProtocolError("drop 1")) is True + + # Hop 2: start_hop replaces the per-hop record cleanly, seeded with hop 1's + # checkpoint token and hop_1 as owner; pre-stream non-connection errors raise + # immediately, wrapped as ContinuationInterruptedError since hop 1's + # checkpoint token is resumable. + assert state.start_hop(stream_opened=False) == ("tok_ckpt_1", "hop_1") + with pytest.raises(errors.ContinuationInterruptedError) as exc_info: + state.end_hop(RuntimeError("429")) + assert isinstance(exc_info.value.__cause__, RuntimeError) + assert exc_info.value.continuation_token == "tok_ckpt_1" + assert exc_info.value.interaction_id == "hop_1" + assert exc_info.value.continuation_ids == ["hop_1"] + assert exc_info.value.partial_interaction is None + + # Pre-stream connection error retries up to MAX_STREAM_RECONNECTS_WITHOUT_PROGRESS + # while preserving ("tok_ckpt_1", "hop_1") as the resume target. + conn_err = compat_errors.APIConnectionError( + request=httpx.Request("POST", "https://example.com/interactions") + ) + state = _interactions_continuation.StreamContinuation( + connection_error_type=compat_errors.APIConnectionError, + ) + state.start_hop(stream_opened=True) + state.observe(_created_event(interaction_id="hop_1")) + state.observe(_status_update_event("tok_ckpt_1")) + assert state.end_hop() is True + + for _ in range( + _interactions_continuation.MAX_STREAM_RECONNECTS_WITHOUT_PROGRESS + ): + assert state.start_hop(stream_opened=False) == ("tok_ckpt_1", "hop_1") + assert state.end_hop(conn_err) is True + + assert state.start_hop(stream_opened=False) == ("tok_ckpt_1", "hop_1") + with pytest.raises(errors.ContinuationInterruptedError) as exc_info: + state.end_hop(conn_err) + assert exc_info.value.__cause__ is conn_err + assert exc_info.value.continuation_token == "tok_ckpt_1" + assert exc_info.value.interaction_id == "hop_1" + + # No server-issued token (only the caller's own): the error is raised as-is. + state = _interactions_continuation.StreamContinuation( + initial_continuation_token="tok_caller", + connection_error_type=compat_errors.APIConnectionError, + ) + for _ in range( + _interactions_continuation.MAX_STREAM_RECONNECTS_WITHOUT_PROGRESS + ): + assert state.start_hop(stream_opened=False) == ("tok_caller", None) + assert state.end_hop(conn_err) is True + assert state.start_hop(stream_opened=False) == ("tok_caller", None) + with pytest.raises(compat_errors.APIConnectionError): + state.end_hop(conn_err) + + +def test_stream_continuation_state_machine_stops_on_unchanged_token(): + state = _interactions_continuation.StreamContinuation( + initial_continuation_token="tok_same", + connection_error_type=compat_errors.APIConnectionError, + ) + assert state.start_hop(stream_opened=True) == ("tok_same", None) + state.observe(_created_event(interaction_id="hop_1")) + state.observe( + _completed_event( + interaction_id="hop_1", + status="requires_continuation", + continuation_token="tok_same", + ) + ) + assert state.end_hop() is False + + +def test_unary_automatic_continuation_transient_retry_on_hop2(): + client = client_lib.Client( + api_key="test-key", + http_options={ + "retry_options": { + "attempts": 2, + "initial_delay": 0.001, + "max_delay": 0.002, + } + }, + ) + + def _json_response( + status_code: int, payload: dict[str, Any] + ) -> httpx.Response: + req = httpx.Request( + "POST", + "https://generativelanguage.googleapis.com/v1beta/interactions", + ) + return httpx.Response( + status_code=status_code, + headers={"content-type": "application/json"}, + content=json.dumps(payload).encode("utf-8"), + request=req, + ) + + hop1_ok = _json_response( + 200, + { + "id": "int_retry_1", + "status": "requires_continuation", + "continuation_token": "tok_retry_1", + "steps": [{ + "type": "model_output", + "content": [{"type": "text", "text": "Hop 1 "}], + }], + "usage": {"total_output_tokens": 20}, + }, + ) + hop2_err_503 = _json_response( + 503, + {"error": {"message": "Service unavailable", "code": 503}}, + ) + hop2_ok = _json_response( + 200, + { + "id": "int_retry_2", + "status": "completed", + "steps": [{ + "type": "model_output", + "content": [{"type": "text", "text": "Hop 2"}], + }], + "usage": {"total_output_tokens": 40}, + }, + ) + + # Hop 1 succeeds on attempt 1. + # Hop 2 fails once with 503 then retries Hop 2 with tok_retry_1 without + # re-running Hop 1. + with mock.patch.object( + client._api_client._httpx_client, + "send", + side_effect=[hop1_ok, hop2_err_503, hop2_ok], + ) as mock_send: + res = client.interactions.create( + model="gemini-2.5-pro", + input="Retry test", + automatic_continuation=True, + ) + + assert mock_send.call_count == 3 + req0_body = json.loads(mock_send.call_args_list[0].args[0].content) + req1_body = json.loads(mock_send.call_args_list[1].args[0].content) + req2_body = json.loads(mock_send.call_args_list[2].args[0].content) + assert "continuation_token" not in req0_body + assert req1_body.get("continuation_token") == "tok_retry_1" + assert req2_body.get("continuation_token") == "tok_retry_1" + assert res.status == "completed" + assert res.output_text == "Hop 1 Hop 2" + assert res.id == "int_retry_2" + assert res.continuation_ids == ["int_retry_1", "int_retry_2"] + assert res.usage.total_output_tokens == 60 + + +def test_unary_later_hop_failure_raises_continuation_interrupted_error(): + """A hop-3 failure keeps hops 1-2 merged on ContinuationInterruptedError.""" + client = client_lib.Client(api_key="test-key") + hop1 = _unary_hop( + interaction_id="int_err_1", + status="requires_continuation", + text="Part 1 ", + continuation_token="tok_err_1", + ) + hop2 = _unary_hop( + interaction_id="int_err_2", + status="requires_continuation", + text="Part 2 ", + continuation_token="tok_err_2", + ) + cause = compat_errors.APIConnectionError( + request=httpx.Request("POST", "https://example.com/interactions") + ) + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[hop1, hop2, cause], + ) as mock_create: + with pytest.raises(errors.ContinuationInterruptedError) as exc_info: + client.interactions.create(model="gemini-2.5-pro", input="Essay") + + assert mock_create.call_count == 3 + err = exc_info.value + assert err.__cause__ is cause + assert err.continuation_token == "tok_err_2" + assert err.interaction_id == "int_err_2" + assert err.continuation_ids == ["int_err_1", "int_err_2"] + partial = err.partial_interaction + assert partial.id == "int_err_2" + assert partial.status == "requires_continuation" + assert partial.continuation_token == "tok_err_2" + assert partial.output_text == "Part 1 Part 2 " + assert partial.continuation_ids == ["int_err_1", "int_err_2"] + + # A first-hop failure has nothing to keep: the error is raised as-is. + with mock.patch.object( + gaos_interactions.Interactions, "create", side_effect=[cause] + ): + with pytest.raises(compat_errors.APIConnectionError): + client.interactions.create(model="gemini-2.5-pro", input="Essay") + + +@pytest.mark.asyncio +async def test_async_unary_later_hop_failure_raises_continuation_interrupted_error(): + client = client_lib.Client(api_key="test-key") + hop1 = _unary_hop( + interaction_id="int_aerr_1", + status="requires_continuation", + text="Part 1 ", + continuation_token="tok_aerr_1", + ) + cause = compat_errors.APIConnectionError( + request=httpx.Request("POST", "https://example.com/interactions") + ) + + with mock.patch.object( + gaos_interactions.AsyncInteractions, + "create", + new=mock.AsyncMock(side_effect=[hop1, cause]), + ): + with pytest.raises(errors.ContinuationInterruptedError) as exc_info: + await client.aio.interactions.create( + model="gemini-2.5-pro", input="Essay" + ) + + err = exc_info.value + assert err.__cause__ is cause + assert err.continuation_token == "tok_aerr_1" + assert err.interaction_id == "int_aerr_1" + assert err.continuation_ids == ["int_aerr_1"] + assert err.partial_interaction.output_text == "Part 1 " + assert err.partial_interaction.status == "requires_continuation" + + +def test_default_automatic_continuation_leaves_max_output_tokens_unset_unary(): + client = client_lib.Client(api_key="test-key") + hop1 = _unary_hop( + interaction_id="int_def_1", + status="requires_continuation", + text="Part 1 ", + continuation_token="tok_def_1", + ) + hop2 = _unary_hop( + interaction_id="int_def_2", + status="completed", + text="Part 2", + ) + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[hop1, hop2], + ) as mock_create: + res = client.interactions.create( + model="gemini-2.5-pro", + input="Write a long essay", + ) + + assert res.status == "completed" + assert res.output_text == "Part 1 Part 2" + assert res.continuation_ids == ["int_def_1", "int_def_2"] + assert mock_create.call_count == 2 + assert "generation_config" not in mock_create.call_args_list[0].kwargs + assert "generation_config" not in mock_create.call_args_list[1].kwargs + + +def test_generation_config_preserved_and_not_mutated_without_max_output_tokens(): + client = client_lib.Client(api_key="test-key") + hop = _unary_hop(interaction_id="int_cfg_1", status="completed", text="Done") + + caller_dict_cfg = {"seed": 42} + caller_model_cfg = interaction_types.GenerationConfig(seed=99) + caller_req = create_interaction_models.CreateInteractionRequest( + body=interaction_types.CreateModelInteraction( + model="gemini-2.5-pro", + input="Request object test", + generation_config=caller_model_cfg, + ) + ) + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[hop, hop, hop], + ) as mock_create: + client.interactions.create( + model="gemini-2.5-pro", + input="Dict config test", + generation_config=caller_dict_cfg, + ) + client.interactions.create( + model="gemini-2.5-pro", + input="Model config test", + generation_config=caller_model_cfg, + ) + client.interactions.create(request=caller_req) + + assert caller_dict_cfg == {"seed": 42} + assert caller_model_cfg.max_output_tokens is None + assert caller_req.body.generation_config.max_output_tokens is None + + assert mock_create.call_args_list[0].kwargs["generation_config"] == { + "seed": 42, + } + sent_model_cfg = mock_create.call_args_list[1].kwargs["generation_config"] + assert sent_model_cfg.seed == 99 + assert sent_model_cfg.max_output_tokens is None + sent_req = mock_create.call_args_list[2].kwargs["request"] + assert sent_req.body.generation_config.seed == 99 + assert sent_req.body.generation_config.max_output_tokens is None + + +@pytest.mark.asyncio +async def test_async_max_output_tokens_left_unset_unless_explicitly_provided(): + client = client_lib.Client(api_key="test-key") + hop = _unary_hop( + interaction_id="int_async_1", status="completed", text="Done" + ) + bg_hop = interaction_types.Interaction(id="int_bg_1", status="in_progress") + stream_events = [ + _created_event(interaction_id="int_s_1"), + _text_delta_event("Chunk"), + _completed_event(interaction_id="int_s_1", status="completed"), + ] + + with mock.patch.object( + gaos_interactions.AsyncInteractions, + "create", + new=mock.AsyncMock( + side_effect=[ + hop, + hop, + hop, + bg_hop, + hop, + _make_async_stream(stream_events), + ] + ), + ) as mock_create: + # 1. Default async unary -> leaves generation_config unset + await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Default async", + ) + # 2. Explicit generation_config.max_output_tokens -> preserved + await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Explicit budget", + generation_config={"max_output_tokens": 2048}, + ) + # 3. agent interaction -> unset + await client.aio.interactions.create( + agent="deep-research-pro-preview-12-2025", + input="Agent run", + ) + # 4. background=True -> unset + bg_res = await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Background", + background=True, + ) + # 5. automatic_continuation=False -> unset + await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Disabled", + automatic_continuation=False, + ) + # 6. stream=True -> unset + stream = await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Streaming", + stream=True, + ) + _ = [ev async for ev in stream] + + assert bg_res.status == "in_progress" + assert bg_res.continuation_ids is None + assert "generation_config" not in mock_create.call_args_list[0].kwargs + assert mock_create.call_args_list[1].kwargs["generation_config"] == { + "max_output_tokens": 2048 + } + assert "generation_config" not in mock_create.call_args_list[2].kwargs + assert "generation_config" not in mock_create.call_args_list[3].kwargs + assert "generation_config" not in mock_create.call_args_list[4].kwargs + assert "generation_config" not in mock_create.call_args_list[5].kwargs + + +def test_unary_incomplete_with_continuation_token_emits_warning_once( + monkeypatch, +): + monkeypatch.setattr( + _interactions_continuation, "_incomplete_continuation_warned", False + ) + client = client_lib.Client(api_key="test-key") + hop1 = _unary_hop( + status="requires_continuation", + text="Partial output", + continuation_token="tok_warn_1", + output_tokens=100, + ) + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[hop1, hop1], + ): + with pytest.warns( + UserWarning, + match="status='requires_continuation'.*automatic_continuation=True", + ): + res = client.interactions.create( + model="gemini-2.5-pro", + input="Write a long story", + automatic_continuation=False, + ) + + import warnings as std_warnings + + with std_warnings.catch_warnings(record=True) as recorded: + std_warnings.simplefilter("always") + res2 = client.interactions.create( + model="gemini-2.5-pro", + input="Write another long story", + automatic_continuation=False, + ) + + assert res.status == "requires_continuation" + assert res.continuation_token == "tok_warn_1" + assert res2.status == "requires_continuation" + incomplete_warnings = [ + w for w in recorded if "status='requires_continuation'" in str(w.message) + ] + assert not incomplete_warnings + + +def test_unary_incomplete_when_caller_passes_continuation_token_does_not_warn( + monkeypatch, recwarn +): + monkeypatch.setattr( + _interactions_continuation, "_incomplete_continuation_warned", False + ) + client = client_lib.Client(api_key="test-key") + hop2 = _unary_hop( + status="requires_continuation", + text="Manual hop 2 still incomplete", + continuation_token="tok_manual_2", + output_tokens=100, + ) + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[hop2], + ): + res = client.interactions.create( + model="gemini-2.5-pro", + input="Write a long story", + continuation_token="tok_manual_1", + automatic_continuation=False, + ) + + assert res.status == "requires_continuation" + incomplete_warnings = [ + w + for w in recwarn.list + if "status='requires_continuation'" in str(w.message) + ] + assert not incomplete_warnings + + +def test_stream_incomplete_with_continuation_token_emits_warning_once( + monkeypatch, +): + monkeypatch.setattr( + _interactions_continuation, "_incomplete_continuation_warned", False + ) + client = client_lib.Client(api_key="test-key") + hop1_events = [ + _created_event(), + _text_delta_event("Stream partial"), + _completed_event( + status="requires_continuation", + continuation_token="tok_stream_warn_1", + output_tokens=50, + ), + ] + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[ + _make_sync_stream(hop1_events), + _make_sync_stream(hop1_events), + ], + ): + stream = client.interactions.create( + model="gemini-2.5-pro", + input="Stream warning test", + stream=True, + automatic_continuation=False, + ) + with pytest.warns( + UserWarning, + match="status='requires_continuation'.*automatic_continuation=True", + ): + received = list(stream) + + import warnings as std_warnings + + stream2 = client.interactions.create( + model="gemini-2.5-pro", + input="Stream warning test 2", + stream=True, + automatic_continuation=False, + ) + with std_warnings.catch_warnings(record=True) as recorded: + std_warnings.simplefilter("always") + received2 = list(stream2) + + assert received == hop1_events + assert received2 == hop1_events + incomplete_warnings = [ + w for w in recorded if "status='requires_continuation'" in str(w.message) + ] + assert not incomplete_warnings + + +@pytest.mark.asyncio +async def test_async_stream_incomplete_with_continuation_token_emits_warning( + monkeypatch, +): + monkeypatch.setattr( + _interactions_continuation, "_incomplete_continuation_warned", False + ) + client = client_lib.Client(api_key="test-key") + hop1_events = [ + _created_event(), + _text_delta_event("Async stream partial"), + _completed_event( + status="requires_continuation", + continuation_token="tok_async_stream_warn_1", + output_tokens=50, + ), + ] + + with mock.patch.object( + gaos_interactions.AsyncInteractions, + "create", + new=mock.AsyncMock(side_effect=[_make_async_stream(hop1_events)]), + ): + stream = await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Async stream warning test", + stream=True, + automatic_continuation=False, + ) + with pytest.warns( + UserWarning, + match="status='requires_continuation'.*automatic_continuation=True", + ): + received = [ev async for ev in stream] + + assert received == hop1_events + + +class _FakeClock: + """A manual monotonic clock for deadline tests.""" + + def __init__(self, now: float = 1000.0): + self.now = now + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +def test_call_deadline_bounds_numeric_timeout_only(): + clock = _FakeClock() + deadline = _interactions_continuation.CallDeadline(10, clock=clock) + assert deadline.hop_timeout() == 10 + clock.advance(4) + assert deadline.remaining() == pytest.approx(6) + assert deadline.hop_timeout() == pytest.approx(6) + assert not deadline.expired() + deadline.check() + clock.advance(7) + assert deadline.expired() + assert deadline.hop_timeout() == 0 + with pytest.raises(_interactions_continuation.ContinuationTimeoutError): + deadline.check() + + # None, httpx.Timeout, and enabled=False pass timeout through unchanged and + # never expire (opt-out / background / agent calls stay pure passthroughs). + for per_hop in (None, httpx.Timeout(10.0)): + deadline = _interactions_continuation.CallDeadline(per_hop, clock=clock) + clock.advance(1_000_000) + assert deadline.remaining() is None + assert not deadline.expired() + assert deadline.hop_timeout() is per_hop + deadline.check() + disabled = _interactions_continuation.CallDeadline( + 60, enabled=False, clock=clock + ) + clock.advance(1_000_000) + assert disabled.remaining() is None + assert disabled.hop_timeout() == 60 + disabled.check() + + # ContinuationTimeoutError mid-hop right after a checkpoint (no content yet) + # fails immediately instead of triggering a reconnect attempt. + state = _interactions_continuation.StreamContinuation( + connection_error_type=compat_errors.APIConnectionError, + ) + state.start_hop(stream_opened=True) + state.observe(_created_event(interaction_id="hop_1")) + state.observe(_status_update_event("tok_ckpt_1")) + with pytest.raises(errors.ContinuationInterruptedError) as exc_info: + state.end_hop(_interactions_continuation.ContinuationTimeoutError("expired")) + assert exc_info.value.continuation_token == "tok_ckpt_1" + assert exc_info.value.interaction_id == "hop_1" + + +def test_unary_timeout_bounds_whole_call_and_passes_remaining_per_hop(): + client = client_lib.Client(api_key="test-key") + clock = _FakeClock() + hops = [ + _unary_hop( + interaction_id="int_t1", + status="requires_continuation", + text="A ", + continuation_token="tok_t1", + ), + _unary_hop( + interaction_id="int_t2", + status="requires_continuation", + text="B ", + continuation_token="tok_t2", + ), + _unary_hop(interaction_id="int_t3", status="completed", text="C"), + ] + + def _create(**kwargs): + del kwargs + clock.advance(10) # each hop takes 10s + return hops.pop(0) + + with mock.patch.object( + _interactions_continuation.time, "monotonic", clock + ), mock.patch.object( + gaos_interactions.Interactions, "create", side_effect=_create + ) as mock_create: + res = client.interactions.create( + model="gemini-2.5-pro", input="Hello", timeout=60 + ) + + assert res.status == "completed" + assert res.continuation_ids == ["int_t1", "int_t2", "int_t3"] + # Each hop is sent with the time left on the call, not the full timeout. + assert [c.kwargs["timeout"] for c in mock_create.call_args_list] == [ + pytest.approx(60), + pytest.approx(50), + pytest.approx(40), + ] + + +def test_unary_timeout_expiry_raises_interrupted_with_partial_result(): + client = client_lib.Client(api_key="test-key") + clock = _FakeClock() + hops = [ + _unary_hop( + interaction_id="int_t1", + status="requires_continuation", + text="A ", + continuation_token="tok_t1", + output_tokens=40, + ), + _unary_hop( + interaction_id="int_t2", + status="requires_continuation", + text="B ", + continuation_token="tok_t2", + output_tokens=30, + ), + ] + + def _create(**kwargs): + del kwargs + clock.advance(10) + return hops.pop(0) + + with mock.patch.object( + _interactions_continuation.time, "monotonic", clock + ), mock.patch.object( + gaos_interactions.Interactions, "create", side_effect=_create + ) as mock_create: + with pytest.raises(errors.ContinuationInterruptedError) as exc_info: + client.interactions.create( + model="gemini-2.5-pro", input="Hello", timeout=15 + ) + + # Hop 1 (10s) left 5s for hop 2, which used it up; hop 3 never started. + assert mock_create.call_count == 2 + assert mock_create.call_args_list[1].kwargs["timeout"] == pytest.approx(5) + err = exc_info.value + assert isinstance(err.__cause__, TimeoutError) + assert isinstance( + err.__cause__, _interactions_continuation.ContinuationTimeoutError + ) + assert err.continuation_token == "tok_t2" + assert err.interaction_id == "int_t2" + assert err.continuation_ids == ["int_t1", "int_t2"] + assert err.partial_interaction.status == "requires_continuation" + assert err.partial_interaction.output_text == "A B " + assert err.partial_interaction.usage.total_output_tokens == 70 + + +@pytest.mark.asyncio +async def test_async_unary_timeout_expiry_raises_interrupted(): + client = client_lib.Client(api_key="test-key") + clock = _FakeClock() + hops = [ + _unary_hop( + interaction_id="int_at1", + status="requires_continuation", + text="A ", + continuation_token="tok_at1", + ), + _unary_hop( + interaction_id="int_at2", + status="requires_continuation", + text="B ", + continuation_token="tok_at2", + ), + ] + + async def _create(**kwargs): + del kwargs + clock.advance(10) + return hops.pop(0) + + with mock.patch.object( + _interactions_continuation.time, "monotonic", clock + ), mock.patch.object( + gaos_interactions.AsyncInteractions, + "create", + new=mock.AsyncMock(side_effect=_create), + ) as mock_create: + with pytest.raises(errors.ContinuationInterruptedError) as exc_info: + await client.aio.interactions.create( + model="gemini-2.5-pro", input="Hello", timeout=15 + ) + + assert mock_create.call_count == 2 + assert [c.kwargs["timeout"] for c in mock_create.call_args_list] == [ + pytest.approx(15), + pytest.approx(5), + ] + assert isinstance(exc_info.value.__cause__, TimeoutError) + assert exc_info.value.continuation_token == "tok_at2" + assert exc_info.value.partial_interaction.output_text == "A B " + + +def test_unary_httpx_timeout_object_applies_per_hop(): + client = client_lib.Client(api_key="test-key") + clock = _FakeClock() + per_hop = httpx.Timeout(10.0) + hops = [ + _unary_hop( + interaction_id="int_h1", + status="requires_continuation", + text="A ", + continuation_token="tok_h1", + ), + _unary_hop(interaction_id="int_h2", status="completed", text="B"), + ] + + def _create(**kwargs): + del kwargs + clock.advance(1_000) + return hops.pop(0) + + with mock.patch.object( + _interactions_continuation.time, "monotonic", clock + ), mock.patch.object( + gaos_interactions.Interactions, "create", side_effect=_create + ) as mock_create: + res = client.interactions.create( + model="gemini-2.5-pro", input="Hello", timeout=per_hop + ) + + assert res.status == "completed" + assert [c.kwargs["timeout"] for c in mock_create.call_args_list] == [ + per_hop, + per_hop, + ] + + +def test_stream_timeout_expiry_between_hops_raises_interrupted(): + client = client_lib.Client(api_key="test-key") + clock = _FakeClock() + hop1_events = [ + _created_event(interaction_id="int_s1"), + _text_delta_event("Part 1 "), + _completed_event( + status="requires_continuation", + continuation_token="tok_s1", + interaction_id="int_s1", + ), + ] + + def _hop1(): + yield hop1_events[0] + yield hop1_events[1] + clock.advance(20) # hop 1 outlives the 15s budget + yield hop1_events[2] + + with mock.patch.object( + _interactions_continuation.time, "monotonic", clock + ), mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[_make_sync_stream(_hop1())], + ) as mock_create: + stream = client.interactions.create( + model="gemini-2.5-pro", input="Hello", stream=True, timeout=15 + ) + received = [] + with pytest.raises(errors.ContinuationInterruptedError) as exc_info: + for event in stream: + received.append(event) + + # Hop 1's events were all yielded; hop 2 never started. + assert received == hop1_events + assert mock_create.call_count == 1 + assert isinstance(exc_info.value.__cause__, TimeoutError) + assert exc_info.value.continuation_token == "tok_s1" + assert exc_info.value.interaction_id == "int_s1" + assert exc_info.value.partial_interaction is None + + +def test_stream_timeout_expiry_mid_hop_raises_interrupted_with_checkpoint(): + client = client_lib.Client(api_key="test-key") + clock = _FakeClock() + hop1_events = [ + _created_event(interaction_id="int_m1"), + _text_delta_event("Part 1 "), + _completed_event( + status="requires_continuation", + continuation_token="tok_m1", + interaction_id="int_m1", + ), + ] + + def _hop2_events(): + yield _created_event(interaction_id="int_m2") + yield _text_delta_event("Part 2a ") + clock.advance(100) # the deadline passes while waiting for the next event + yield _text_delta_event("Part 2b ") + yield _completed_event(status="completed", interaction_id="int_m2") + + with mock.patch.object( + _interactions_continuation.time, "monotonic", clock + ), mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[ + _make_sync_stream(hop1_events), + _make_sync_stream(_hop2_events()), + ], + ) as mock_create: + stream = client.interactions.create( + model="gemini-2.5-pro", input="Hello", stream=True, timeout=60 + ) + received = [] + with pytest.raises(errors.ContinuationInterruptedError) as exc_info: + for event in stream: + received.append(event) + + # The event that arrived as the deadline passed is still delivered; the + # stream stops before waiting for the next one. + assert [ev["event_type"] for ev in received] == [ + "interaction.created", + "step.delta", + "interaction.completed", + "interaction.created", + "step.delta", + "step.delta", + ] + assert mock_create.call_count == 2 + assert isinstance(exc_info.value.__cause__, TimeoutError) + # Hop 2 produced output after hop 1's token, so resuming from tok_m1 + # re-generates "Part 2a " (documented on ContinuationInterruptedError). + assert exc_info.value.continuation_token == "tok_m1" + assert exc_info.value.interaction_id == "int_m1" + assert exc_info.value.continuation_ids == ["int_m1", "int_m2"] + + +@pytest.mark.asyncio +async def test_async_stream_timeout_expiry_between_hops_raises_interrupted(): + client = client_lib.Client(api_key="test-key") + clock = _FakeClock() + hop1_events = [ + _created_event(interaction_id="int_as1"), + _text_delta_event("Part 1 "), + _completed_event( + status="requires_continuation", + continuation_token="tok_as1", + interaction_id="int_as1", + ), + ] + + def _hop1(): + yield hop1_events[0] + yield hop1_events[1] + clock.advance(20) + yield hop1_events[2] + + with mock.patch.object( + _interactions_continuation.time, "monotonic", clock + ), mock.patch.object( + gaos_interactions.AsyncInteractions, + "create", + new=mock.AsyncMock(side_effect=[_make_async_stream(_hop1())]), + ) as mock_create: + stream = await client.aio.interactions.create( + model="gemini-2.5-pro", input="Hello", stream=True, timeout=15 + ) + received = [] + with pytest.raises(errors.ContinuationInterruptedError) as exc_info: + async for event in stream: + received.append(event) + + assert received == hop1_events + assert mock_create.call_count == 1 + assert isinstance(exc_info.value.__cause__, TimeoutError) + assert exc_info.value.continuation_token == "tok_as1" + + +@pytest.mark.asyncio +async def test_background_and_agent_requires_continuation_do_not_warn( + monkeypatch, recwarn +): + monkeypatch.setattr( + _interactions_continuation, "_incomplete_continuation_warned", False + ) + client = client_lib.Client(api_key="test-key") + incomplete_hop = _unary_hop( + interaction_id="int_skip_warn", + status="requires_continuation", + text="Server-managed partial", + continuation_token="tok_skip_warn", + output_tokens=25, + ) + incomplete_events = [ + _created_event(interaction_id="int_skip_stream"), + _text_delta_event("Server-managed stream partial"), + _completed_event( + interaction_id="int_skip_stream", + status="requires_continuation", + continuation_token="tok_skip_stream", + output_tokens=25, + ), + ] + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[ + incomplete_hop, + incomplete_hop, + _make_sync_stream(incomplete_events), + ], + ), mock.patch.object( + gaos_interactions.AsyncInteractions, + "create", + new=mock.AsyncMock( + side_effect=[ + incomplete_hop, + _make_async_stream(incomplete_events), + ] + ), + ): + bg_res = client.interactions.create( + model="gemini-2.5-pro", + input="Background task", + background=True, + ) + agent_res = client.interactions.create( + agent="deep-research-pro-preview-12-2025", + input="Agent task", + ) + bg_stream = list( + client.interactions.create( + model="gemini-2.5-pro", + input="Background stream", + background=True, + stream=True, + ) + ) + async_agent_res = await client.aio.interactions.create( + agent="deep-research-pro-preview-12-2025", + input="Async agent task", + ) + async_agent_stream = [ + ev + async for ev in await client.aio.interactions.create( + agent="deep-research-pro-preview-12-2025", + input="Async agent stream", + stream=True, + ) + ] + + assert bg_res.status == "requires_continuation" + assert agent_res.status == "requires_continuation" + assert len(bg_stream) == 3 + assert async_agent_res.status == "requires_continuation" + assert len(async_agent_stream) == 3 + incomplete_warnings = [ + w + for w in recwarn.list + if "status='requires_continuation'" in str(w.message) + ] + assert not incomplete_warnings + + +@pytest.mark.asyncio +async def test_store_false_interrupted_sets_interaction_id_none(): + client = client_lib.Client(api_key="test-key") + hop1 = _unary_hop( + interaction_id="int_unstored_1", + status="requires_continuation", + text="Part 1 ", + continuation_token="tok_unstored_1", + output_tokens=10, + ) + hop1_events = [ + _created_event(interaction_id="int_unstored_s1"), + _text_delta_event("Stream Part 1 "), + _completed_event( + interaction_id="int_unstored_s1", + status="requires_continuation", + continuation_token="tok_unstored_s1", + ), + ] + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[hop1, RuntimeError("boom unary")], + ): + with pytest.raises(errors.ContinuationInterruptedError) as unary_exc: + client.interactions.create( + model="gemini-2.5-pro", + input="Hello", + store=False, + ) + assert unary_exc.value.continuation_token == "tok_unstored_1" + assert unary_exc.value.interaction_id is None + + with mock.patch.object( + gaos_interactions.Interactions, + "create", + side_effect=[ + _make_sync_stream(hop1_events), + RuntimeError("boom sync stream"), + ], + ): + stream = client.interactions.create( + model="gemini-2.5-pro", + input="Hello", + store=False, + stream=True, + ) + with pytest.raises(errors.ContinuationInterruptedError) as stream_exc: + list(stream) + assert stream_exc.value.continuation_token == "tok_unstored_s1" + assert stream_exc.value.interaction_id is None + + with mock.patch.object( + gaos_interactions.AsyncInteractions, + "create", + new=mock.AsyncMock( + side_effect=[ + _make_async_stream(hop1_events), + RuntimeError("boom async stream"), + ] + ), + ): + astream = await client.aio.interactions.create( + model="gemini-2.5-pro", + input="Hello", + store=False, + stream=True, + ) + with pytest.raises(errors.ContinuationInterruptedError) as astream_exc: + _ = [ev async for ev in astream] + assert astream_exc.value.continuation_token == "tok_unstored_s1" + assert astream_exc.value.interaction_id is None + + +def test_merge_usages_does_not_leak_later_hop_initial_input_tokens(): + hop1_missing_input = interaction_types.Interaction( + id="int_u1", + status="requires_continuation", + continuation_token="tok_u1", + steps=[], + usage=interaction_types.Usage(total_output_tokens=20), + ) + hop2 = interaction_types.Interaction( + id="int_u2", + status="completed", + steps=[], + usage=interaction_types.Usage( + initial_input_tokens=500, + total_input_tokens=500, + total_output_tokens=30, + ), + ) + merged = _interactions_continuation.merge_unary_hops( + [hop1_missing_input, hop2] + ) + assert merged.usage.initial_input_tokens is None + assert merged.usage.total_input_tokens == 500 + assert merged.usage.total_output_tokens == 50 + + hop1_no_usage = interaction_types.Interaction( + id="int_u1_none", + status="requires_continuation", + continuation_token="tok_u1", + steps=[], + usage=None, + ) + merged_no_usage = _interactions_continuation.merge_unary_hops( + [hop1_no_usage, hop2] + ) + assert merged_no_usage.usage.initial_input_tokens is None + assert merged_no_usage.usage.total_output_tokens == 30 + + +def test_merge_usages_supports_slotted_usage_objects(): + class _SlottedUsage: + __slots__ = ("total_input_tokens", "total_output_tokens") + + def __init__(self, total_input_tokens: int, total_output_tokens: int): + self.total_input_tokens = total_input_tokens + self.total_output_tokens = total_output_tokens + + class _SlottedUsageWithInitial(_SlottedUsage): + __slots__ = ("initial_input_tokens",) + + def __init__( + self, + total_input_tokens: int, + total_output_tokens: int, + initial_input_tokens: int | None = None, + ): + super().__init__(total_input_tokens, total_output_tokens) + self.initial_input_tokens = initial_input_tokens + + merged_basic = _interactions_continuation._merge_usages( + [_SlottedUsage(10, 20), _SlottedUsage(30, 40)] + ) + assert merged_basic.total_input_tokens == 40 + assert merged_basic.total_output_tokens == 60 + + merged_with_initial = _interactions_continuation._merge_usages([ + _SlottedUsageWithInitial(12, 20), + _SlottedUsageWithInitial(50, 40, initial_input_tokens=50), + ]) + assert merged_with_initial.total_input_tokens == 62 + assert merged_with_initial.total_output_tokens == 60 + assert merged_with_initial.initial_input_tokens == 12 + + +@pytest.mark.asyncio +async def test_async_stream_continuation_handles_secondary_stream_without_close(): + class _CloselessAsyncStream: + + def __init__(self, events: list[Any]): + async def _agen(): + for ev in events: + yield ev + + self.generator = _agen() + + hop1 = _make_async_stream([ + _created_event(interaction_id="int_c1"), + _completed_event( + interaction_id="int_c1", + status="requires_continuation", + continuation_token="tok_c1", + ), + ]) + hop2 = _CloselessAsyncStream([ + _created_event(interaction_id="int_c2"), + _text_delta_event("Done"), + _completed_event(interaction_id="int_c2", status="completed"), + ]) + + async def _create_next(token, owner_id): + assert (token, owner_id) == ("tok_c1", "int_c1") + return hop2 + + wrapped = _interactions_continuation.wrap_async_stream_continuation( + hop1, + _create_next, + is_stream_fn=lambda s: hasattr(s, "generator"), + wrap_errors_fn=lambda s: s, + connection_error_type=compat_errors.APIConnectionError, + ) + events = [ev async for ev in wrapped] + assert len(events) == 5