"""Model-continuation primitives shared by the agent graph. A *continuation* happens when a model returns a `ModelResponse` with `state == 'suspended'` (Anthropic `pause_turn`, OpenAI background mode, …): the graph re-issues the request with the suspended response echoed back, and the provider resumes the same logical turn. This module owns the provider-agnostic glue for stitching those segments back into a single response/stream: - [`merge_responses`][pydantic_ai.models._continuation.merge_responses] folds a continuation response into the one it continues. - [`merge_mode`][pydantic_ai.models._continuation.merge_mode] reports whether a continuation *replaces* or *accumulates*, so the streamed composite can reindex parts consistently with the merge. - [`_ContinuationStreamedResponse`][pydantic_ai.models._continuation._ContinuationStreamedResponse] drives the streamed loop, presenting every segment as one continuous stream. This module is deliberately decoupled from `_agent_graph`: it imports only from `models`, `messages`, `usage`, `exceptions`, and the stdlib. Pluggable timing is injected as `sleep_func` so the loop stays free of `now_utc()`/RNG and replays deterministically under durable executors (e.g. Temporal). """ from __future__ import annotations import asyncio import time from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable from contextlib import AbstractContextManager, nullcontext, suppress from dataclasses import dataclass, field, replace from datetime import datetime, timezone from typing import Any, Literal from .. import _utils from .._run_context import RunContext from ..exceptions import UnexpectedModelBehavior from ..messages import ( FinalResultEvent, ModelMessage, ModelResponse, ModelResponseState, ModelResponseStreamEvent, PartDeltaEvent, PartEndEvent, PartStartEvent, ) from ..settings import ModelSettings from ..usage import RequestUsage from . import Model, StreamedResponse __all__ = [ 'MAX_BACKGROUND_POLLS', 'MAX_GENERATION_CONTINUATIONS', 'MergeMode', 'cancel_suspended_job', 'merge_mode', 'merge_responses', '_ContinuationStreamedResponse', ] # Framework-protocol markers a model may stamp on a continuation response's `metadata`, namespaced # under `__pydantic_ai__` so they never collide with a provider's own metadata. `FallbackModel` sets # `replace_previous_response` on the fresh response it produces after a rewind-and-restart to signal # "this response supersedes the suspended turn, do not accumulate onto it" — categorically different # from a same-`provider_response_id` background poll. The `FallbackModel` side must use these exact # keys. The marker is transient: `merge_responses` honors it as a replace, then pops it so it can't # persist into history and wrongly force a later legitimate `pause_turn` continuation to replace. _PYDANTIC_AI_METADATA_KEY = '__pydantic_ai__' _REPLACE_PREVIOUS_RESPONSE_KEY = 'replace_previous_response' MAX_GENERATION_CONTINUATIONS = 10 """Maximum number of *fresh-generation* continuation segments for a single model turn. Applies to every re-suspension that produces genuinely new generation: an *accumulate* (Anthropic `pause_turn`, appending new parts under a new `provider_response_id`), a model change, or a `FallbackModel` `replace_previous_response` directive. This is the guard against a model that never leaves the `'suspended'` state, endlessly emitting new segments; exceeding it raises [`UnexpectedModelBehavior`][pydantic_ai.exceptions.UnexpectedModelBehavior]. Only a *same-id* re-suspension — re-polling a single long-running job under the same `provider_response_id` (OpenAI background mode) — is bounded by the far more generous [`MAX_BACKGROUND_POLLS`][pydantic_ai.models._continuation.MAX_BACKGROUND_POLLS] instead, so a healthy background job that legitimately runs for minutes isn't killed after ~10 continuations. """ MAX_BACKGROUND_POLLS = 1000 """Backstop for *same-id* continuation polling of a single background job. A same-id re-suspension re-fetches one long-running job under the same `provider_response_id` (OpenAI background mode), so — unlike a fresh-generation re-suspension (see [`MAX_GENERATION_CONTINUATIONS`][pydantic_ai.models._continuation.MAX_GENERATION_CONTINUATIONS]) — it carries no risk of an unbounded model spawning endless new segments: it's the *same* job, and legitimate background jobs run for minutes (hundreds of polls at a ~2s interval). The real bounds here are the run's usage limits and explicit cancellation, which already work; this large ceiling is only a last-resort safety net against a provider stuck returning `'suspended'` for the same id forever (which usage limits wouldn't catch, since a pending poll adds no tokens). Exceeding it raises [`UnexpectedModelBehavior`][pydantic_ai.exceptions.UnexpectedModelBehavior]. """ MergeMode = Literal['replace-same-id', 'replace-new', 'accumulate'] """How a continuation response folds into the one it continues. - `'replace-same-id'`: both responses share a `provider_response_id` — a passive re-poll of one long-running job (OpenAI background `retrieve`, which returns the full response so far). Replaces wholesale, and — being the *same* job rather than fresh generation — is bounded by the generous [`MAX_BACKGROUND_POLLS`][pydantic_ai.models._continuation.MAX_BACKGROUND_POLLS] ceiling. - `'replace-new'`: the response supersedes the suspended turn with *fresh* generation — the model changed (accumulating parts from different models is always wrong) or a `FallbackModel` stamped the `replace_previous_response` marker after a rewind-and-restart. Replaces wholesale, but counts against the strict [`MAX_GENERATION_CONTINUATIONS`][pydantic_ai.models._continuation.MAX_GENERATION_CONTINUATIONS] ceiling, since a chain of fresh suspensions is the exact runaway that cap guards against. - `'accumulate'`: appends new parts onto the prior response (Anthropic `pause_turn`). Strict ceiling. """ # Deterministic fallback for `timestamp` before any segment has streamed. This is # never reached in practice (a segment is always in flight or finalized by the time # `timestamp` is read), but keeps the loop free of `now_utc()` for durable replay. _FALLBACK_TIMESTAMP = datetime.fromtimestamp(0, tz=timezone.utc) def _has_replace_marker(response: ModelResponse) -> bool: """Whether `response` carries the `FallbackModel` `replace_previous_response` directive.""" metadata = response.metadata if not _utils.is_str_dict(metadata): return False namespace = metadata.get(_PYDANTIC_AI_METADATA_KEY) return bool(_utils.is_str_dict(namespace) and namespace.get(_REPLACE_PREVIOUS_RESPONSE_KEY)) def _strip_replace_marker(metadata: Any) -> dict[str, Any] | None: """Return `metadata` without the transient `replace_previous_response` marker (other keys intact). Copies before mutating so the caller's dicts (including the shared `__pydantic_ai__` namespace, which also holds the `FallbackModel` continuation pin) aren't touched. """ if not _utils.is_str_dict(metadata): return metadata namespace = metadata.get(_PYDANTIC_AI_METADATA_KEY) if not (_utils.is_str_dict(namespace) and _REPLACE_PREVIOUS_RESPONSE_KEY in namespace): return metadata namespace = {k: v for k, v in namespace.items() if k != _REPLACE_PREVIOUS_RESPONSE_KEY} metadata = {**metadata} if namespace: metadata[_PYDANTIC_AI_METADATA_KEY] = namespace else: del metadata[_PYDANTIC_AI_METADATA_KEY] return metadata def merge_mode(existing: ModelResponse, new: ModelResponse) -> MergeMode: """Classify how `new` folds into `existing` — see [`MergeMode`][pydantic_ai.models._continuation.MergeMode]. The single decision path shared by [`merge_responses`][pydantic_ai.models._continuation.merge_responses], the continuation-count ceilings, and the streamed composite's part-index reindexing. """ # A `FallbackModel` rewind-and-restart marks its fresh response as superseding the suspended turn; # honor that first, since such a response may otherwise look like an accumulate. if _has_replace_marker(new): return 'replace-new' if existing.provider_response_id and existing.provider_response_id == new.provider_response_id: return 'replace-same-id' if existing.model_name and new.model_name and existing.model_name != new.model_name: return 'replace-new' return 'accumulate' def merge_responses(existing: ModelResponse, new: ModelResponse) -> ModelResponse: """Merge a continuation response into the one it continues. On any `'replace-*'` mode (same `provider_response_id`, a model change, or a `FallbackModel` `replace_previous_response` directive), replace the content with the new response. Fresh-generation replacements retain the prior response's billed usage; same-job polling replaces its cumulative usage snapshot. Otherwise accumulate parts and usage, and use other fields from the new response. Either way, `provider_details` and `metadata` accumulate across the turn's segments (latest-wins) so turn-scoped data a later segment omits isn't lost — see below. """ mode = merge_mode(existing, new) if mode == 'replace-same-id': merged = new elif mode == 'replace-new': merged = replace(new, usage=existing.usage + new.usage) else: # Same model, different response → accumulate parts and sum usage. # Preserve existing provider response IDs when continuation responses omit them # (e.g. resumed OpenAI streams that start after a sequence number). merged = replace( new, parts=[*existing.parts, *new.parts], usage=existing.usage + new.usage, provider_response_id=new.provider_response_id or existing.provider_response_id, ) # A turn's provider metadata accumulates across its segments. Turn-scoped identifiers — OpenAI # `background`/`conversation_id`, Anthropic `container_id` in `provider_details`, and the # `FallbackModel` continuation pin in `metadata` — are only stamped on segments whose payload # carries them, so a resumed or interrupted segment (e.g. a mid-flight cancel snapshot, whose # in-flight segment hasn't stamped the pin yet) can omit one an earlier segment set. Merge # latest-wins (`new` overrides) so they survive into the merged response; e.g. # `cancel_suspended_response` relies on OpenAI's `background` marker and the `FallbackModel` pin # to reach the server-side job. if existing.provider_details: merged = replace(merged, provider_details={**existing.provider_details, **(merged.provider_details or {})}) if existing.metadata: merged = replace(merged, metadata={**existing.metadata, **(merged.metadata or {})}) # Pop the transient `replace_previous_response` marker now that it's been honored above, so it # doesn't persist into history where it would wrongly force a later legitimate `pause_turn` # continuation to replace rather than accumulate. Other `__pydantic_ai__` keys (the continuation # pin) survive. stripped = _strip_replace_marker(merged.metadata) if stripped is not merged.metadata: merged = replace(merged, metadata=stripped) return merged async def cancel_suspended_job(model: Model, response: ModelResponse) -> None: """Best-effort teardown of a server-side suspended/background job that survives cancellation. When the trigger is a workflow/task cancellation (e.g. Temporal), awaiting the (activity-wrapped) cancel from inside an already-cancelled scope would raise `CancelledError` before the cancel runs, silently leaking the job. Shield the cancel so it completes before the cancellation propagates; Temporal's workflow loop respects `asyncio.shield`. Any error from the cancel itself is swallowed — a failing teardown must not replace the error (or cancellation) that aborted the run. """ job = asyncio.ensure_future(model.cancel_suspended_response(response)) try: await asyncio.shield(job) except asyncio.CancelledError: # Our scope was cancelled mid-teardown; let the shielded cancel finish before propagating. with suppress(Exception): await job raise except Exception: pass @dataclass class _ContinuationStreamedResponse(StreamedResponse): """A [`StreamedResponse`][pydantic_ai.models.StreamedResponse] that stitches continuation segments into one stream. Each segment is an ordinary `model.request_stream(...)` sub-stream. Their events are re-emitted as a single continuous stream, with part indices offset so parts from accumulated segments (Anthropic `pause_turn`) don't collide, while replaced segments (OpenAI background `retrieve`) keep reusing the same index space. `get()` returns the live merged snapshot at any point; `usage` sums/replaces in lockstep with the merge so the graph accounts for it exactly once. """ model: Model model_settings: ModelSettings | None base_messages: list[ModelMessage] run_context: RunContext[Any] | None max_generation_continuations: int sleep_func: Callable[[float], Awaitable[None]] check_usage: Callable[[RequestUsage], None] finalize_response: Callable[[ModelResponse], None] initial_suspended_response: ModelResponse | None = None # Ceiling for *replace*-style (single-job background poll) re-suspensions, kept separate from # `max_generation_continuations` (which bounds fresh-generation re-suspensions). See `MAX_BACKGROUND_POLLS`. max_background_polls: int = MAX_BACKGROUND_POLLS # Entered around each segment's `model.request_stream(...)`. The agent graph passes a factory # that re-attaches the ambient context (e.g. the OTel `chat` span opened by `wrap_model_request` # in a separate task) so span updates driven by `get_current_span()` land on the right span even # though segments are opened lazily in the consumer task. Opaque to this module (no OTel coupling). segment_context: Callable[[], AbstractContextManager[Any]] = nullcontext _merged_response: ModelResponse | None = field(default=None, init=False) _current_sub: StreamedResponse | None = field(default=None, init=False) _stopped: bool = field(default=False, init=False) # Set by `aclose()`: the consumer stopped iterating and the stream was torn down *without* a # `cancel()`/`close_stream()` (which would flip `_stopped`/`_cancelled` and cancel the server-side # job). Lets `get()` distinguish a deliberate detach — where a still-pending suspended job survives # server-side and the run should be recorded as resumable `'suspended'` — from a live, still-streaming # snapshot (`'incomplete'`). See `get()`. _detached: bool = field(default=False, init=False) # The inner segment-stitching generator, kept separately so `aclose()` can tear it down # directly: the outer cancel-guard's `async for … in iterator` does NOT forward `aclose()` # to it, and this generator owns each segment's `async with request_stream(...)`. _segment_iterator: AsyncGenerator[ModelResponseStreamEvent, None] | None = field(default=None, init=False) def __aiter__(self) -> AsyncIterator[ModelResponseStreamEvent]: """Stream every segment as one continuous event stream. This intentionally bypasses the base `StreamedResponse.__aiter__`'s `iterator_with_final_event` / `iterator_with_part_end` wrappers: each sub-stream is already wrapped by them, so it emits fully-formed `PartStart`/`PartDelta`/`PartEnd` and `FinalResultEvent`s. The composite only applies reindexing + final-result capture (inside `_get_event_iterator`) and the cancel-guard (reproducing the base `_finished`/`_cancelled` transitions) on top. One minor semantic gap: `PartEndEvent.next_part_kind` is `None` at each sub-stream boundary (a segment can't see the next segment's first part), whereas a single-segment stream would populate it. This is acceptable because parts never merge across segment boundaries. """ if self._event_iterator is None: self._segment_iterator = self._get_event_iterator() self._event_iterator = self._iterator_with_cancel_guard(self._segment_iterator) return self._event_iterator def get_stream_cancel_errors(self) -> tuple[type[BaseException], ...]: """Cancel-teardown errors to suppress, extended with the in-flight sub-stream's own. The cancel-guard tears the current segment down via its `close_stream()`, so the transport error it raises is whatever that sub's transport produces — httpx for most providers, but botocore (Bedrock) or grpc (xAI) for others, which each report their own types via an override. Consult the in-flight sub (`_current_sub` is still set at exception time) on top of the httpx default, so a non-httpx teardown error is suppressed into a clean `'interrupted'` stop rather than escaping to the consumer. """ errors = super().get_stream_cancel_errors() if self._current_sub is not None: errors = (*errors, *self._current_sub.get_stream_cancel_errors()) return errors async def _iterator_with_cancel_guard( self, iterator: AsyncIterator[ModelResponseStreamEvent] ) -> AsyncIterator[ModelResponseStreamEvent]: # Mirror `StreamedResponse.__aiter__`'s cancel-guard: suppress transport # errors caused by `cancel()` tearing down an in-flight sub-stream, and only # flip `_finished` on a natural `StopAsyncIteration` of a stream that wasn't # cancelled, so an early `break`/`aclose()`/in-flight error — or a `cancel()` # that still drains to completion — leaves `get()` reporting `'incomplete'`/ # `'interrupted'` rather than `'complete'`. try: async for event in iterator: if self._first_chunk_monotonic is None: # First event surfaced to the consumer: stamp the monotonic clock so # `time_to_first_chunk` works, mirroring the base cancel-guard this replaces. self._first_chunk_monotonic = time.perf_counter() yield event except self.get_stream_cancel_errors(): if not self.cancelled: raise else: if not self._cancelled: self._finished = True def _count_continuation( self, response: ModelResponse, last_mode: MergeMode | None, accumulate_count: int, replace_count: int ) -> tuple[int, int]: """Count a suspended re-issue against its ceiling (same-id poll vs everything else), raising if exceeded. Only a `'replace-same-id'` re-suspension — a passive re-poll of one long-running background job — gets the generous `max_background_polls` ceiling. A model-change or `FallbackModel`-directed replace is *fresh* generation, not the same job, so it counts against the strict `max_generation_continuations` cap alongside accumulate re-suspensions — otherwise a chain of fresh suspensions is the exact runaway the strict cap guards against. The first re-issue has no prior merge to classify (`last_mode is None`) so it counts as strict, harmless since both ceilings allow at least one. See `MAX_BACKGROUND_POLLS`. Returns the updated `(accumulate_count, replace_count)`. """ job_id = response.provider_response_id if last_mode == 'replace-same-id': replace_count += 1 if replace_count > self.max_background_polls: raise UnexpectedModelBehavior( f'Model response for job {job_id!r} remained suspended after polling the maximum of ' f'{self.max_background_polls} times' ) else: accumulate_count += 1 if accumulate_count > self.max_generation_continuations: raise UnexpectedModelBehavior( f'Model response {job_id!r} was suspended more than the maximum of ' f'{self.max_generation_continuations} times' ) return accumulate_count, replace_count async def _get_event_iterator(self) -> AsyncGenerator[ModelResponseStreamEvent, None]: # noqa: C901 # Two independent ceilings, distinguished by the generic `merge_mode` signal (the same one that # drives reindexing): every *fresh-generation* re-suspension (accumulate `pause_turn`, a model # change, or a `FallbackModel` replace directive) risks an unbounded model spawning new segments, # so it keeps the small `max_generation_continuations` cap; only a *same-id* re-suspension (OpenAI background # poll) re-fetches one long-running job under the same `provider_response_id`, so a healthy job # that legitimately runs for minutes must not be killed by the small cap — it gets the far more # generous `max_background_polls` backstop. Mirrors the non-streaming continuation loop in # `_agent_graph`. See `MAX_BACKGROUND_POLLS`. accumulate_count = 0 replace_count = 0 # Mode of the merge that produced the current suspended `response`, used to pick its ceiling. A # continuation chain is homogeneous in practice (a poll chain is all same-id, a `pause_turn` chain # all-accumulate), so the previous merge's mode reliably classifies the next re-issue. last_mode: MergeMode | None = None response = self.initial_suspended_response # Index at which the most recent segment's parts began in the stitched stream. A replaced # segment reuses this (same parts under the same id); an accumulated segment appends after # all prior parts. See `_segment_offset`. last_segment_offset = 0 try: while True: if self._cancelled or self._stopped: break if response is None: messages = self.base_messages elif response.state == 'suspended': accumulate_count, replace_count = self._count_continuation( response, last_mode, accumulate_count, replace_count ) if delay := self.model.continuation_delay(response): await self.sleep_func(delay) # A `cancel()`/`close_stream()` from another task during the inter-poll sleep # already tore down the server-side job; don't open the next sub-stream, which # for Anthropic `pause_turn` would actively resume generation and burn tokens. if self._cancelled or self._stopped: break messages = [*self.base_messages, response] else: break # While this sub is in flight, `_merged_response` holds the accumulator of # all prior segments (excluding the current sub) so `get()` can fold in the # live `sub.get()` snapshot without double-counting. self._merged_response = response # Resolved lazily on the first reindexable event, once `sub.provider_response_id` # is populated, so replace-vs-accumulate matches the eventual `merge_mode`. segment_offset: int | None = None with self.segment_context(): async with self.model.request_stream( messages, self.model_settings, self.model_request_parameters, self.run_context ) as sub: self._current_sub = sub async for event in sub: if isinstance(event, FinalResultEvent): self.final_result_event = event yield event continue if segment_offset is None: segment_offset = self._segment_offset(response, sub, last_segment_offset) yield self._reindex(event, segment_offset) last_segment_offset = segment_offset or 0 # Read `sub.get()` AFTER the `async with` exits so late-stamped metadata # (e.g. a `FallbackModel` continuation pin) is captured. sub_response = sub.get() if response is None: if sub_response.state == 'suspended': self.finalize_response(sub_response) merged = sub_response else: # Continuation segments are separately billed requests. Finalize their usage before # merging so additive costs preserve per-request pricing, including pricing tiers. self.finalize_response(response) self.finalize_response(sub_response) # Classify this transition (replace vs accumulate) so the next re-issue is counted # against the right ceiling. last_mode = merge_mode(response, sub_response) merged = merge_responses(response, sub_response) self._merged_response = merged self._current_sub = None self._usage = merged.usage self.check_usage(merged.usage) response = merged self._merged_response = response except GeneratorExit: # Deliberate `aclose()` detach: tear the connection down without cancelling the # server-side job (that stays on the `cancel()`/`close_stream()` path). Re-raise as-is. raise except BaseException: # A later segment failed (transport error, `check_usage` raising, or the max-continuations # raise) with a suspended job in hand. The non-streaming continuation loop cancels the # server-side job on exactly this class of failure; mirror it so streaming doesn't leak the # job (which history would otherwise record as unresumable and uncancellable). Skip when a # deliberate `cancel()`/`close_stream()` is already tearing things down — it cancels itself. if response is not None and response.state == 'suspended' and not (self._cancelled or self._stopped): await cancel_suspended_job(self.model, response) raise finally: # Finalize an interrupted segment only after its generator has unwound, as teardown can stamp # additional usage. This also handles `aclose()` racing a debounced consumer's prefetch task. if self._current_sub is not None: self.finalize_response(self._current_sub.get()) @staticmethod def _segment_offset(response: ModelResponse | None, sub: StreamedResponse, last_segment_offset: int) -> int: """Index at which the current segment's parts begin in the stitched stream. Shares [`merge_mode`][pydantic_ai.models._continuation.merge_mode]'s decision so reindexing matches the eventual merge: - `'accumulate'` appends after all prior parts (offset = number of prior parts). - `'replace-same-id'` (a background job re-polled under the same `provider_response_id`) re-emits the *same* parts in the *same* index space, so it reuses the replaced segment's offset. - `'replace-new'` (a model change, or a `FallbackModel` `replace_previous_response` directive) supersedes the whole prior response — `merge_responses` keeps only the new parts, indexed from 0 — so its events must start at offset 0 too, or the live event indices would drift past the final response's (e.g. after one or more accumulated segments). """ if response is None: return 0 mode = merge_mode(response, sub.get()) if mode == 'accumulate': return len(response.parts) if mode == 'replace-new': return 0 return last_segment_offset def _reindex(self, event: ModelResponseStreamEvent, offset: int) -> ModelResponseStreamEvent: if offset and isinstance(event, (PartStartEvent, PartDeltaEvent, PartEndEvent)): return replace(event, index=event.index + offset) return event def _snapshot(self) -> ModelResponse | None: """The merged response so far, folding in any in-flight sub-stream.""" merged = self._merged_response if (sub := self._current_sub) is not None: sub_response = sub.get() return sub_response if merged is None else merge_responses(merged, sub_response) return merged @property def usage(self) -> RequestUsage: """Live usage across all segments so far, including the in-flight sub-stream. The composite's `_usage` is only refreshed when a segment completes, so — unlike a plain segment, whose model updates `_usage` live during iteration — reading it mid segment would omit the in-flight sub's usage. Fold in the current sub's live snapshot so consumers (e.g. `AgentStream.usage`) see the running total at any point. """ snapshot = self._snapshot() return snapshot.usage if snapshot is not None else self._usage def get(self) -> ModelResponse: """Build the live merged [`ModelResponse`][pydantic_ai.messages.ModelResponse] across all segments so far. The composite normally resolves the whole `suspended → … → complete` chain, so mid-run it's `'complete'` once the loop exits, `'interrupted'` if cancelled, and `'incomplete'` while a segment is still in flight. The one case it *does* surface `'suspended'` is a **detach**: the consumer stopped iterating and the stream was torn down via `aclose()` — not `cancel()` — while the current/last segment is itself a still-pending suspended job. The server-side job survives (detach doesn't cancel it), so recording `'suspended'` makes the run resumable later, matching the non-streaming path where a persisted suspended response can be resumed. A real `cancel()` (`_cancelled`) also cancels the server-side job, so it stays `'interrupted'` and non-resumable. """ snapshot = self._snapshot() state: ModelResponseState if self._finished: state = 'complete' elif self._cancelled: # A real `cancel()` tore down the server-side job too, so this is not resumable. state = 'interrupted' elif self._detached and snapshot is not None and snapshot.state == 'suspended': # Detached with a still-pending suspended job in hand (and not cancelled): resumable. state = 'suspended' else: state = 'incomplete' if snapshot is None: return ModelResponse(parts=[], model_name=self.model_name, state=state) return replace(snapshot, state=state) async def close_stream(self) -> None: """Stop the continuation loop and cancel any server-side suspended/background job.""" self._stopped = True try: if self._current_sub is not None: await self._current_sub.close_stream() finally: # Cancel the server-side job even if tearing down the sub-stream connection raised, # so a failed connection teardown can't leave the background job running. Shielded so the # cancel completes even when the run is being torn down by a workflow/task cancellation. await cancel_suspended_job(self.model, self.get()) async def aclose(self) -> None: """Tear down the in-flight sub-stream (its HTTP connection) without cancelling any job. Each segment's `model.request_stream(...)` context manager lives inside the stitching async generator, so — unlike an ordinary `StreamedResponse` whose connection the agent graph closed via `async with request_stream(...)` — an early break from a consumer doesn't reliably propagate `aclose()` to it. Closing the generator here runs that context manager's teardown (mirroring the pre-stitching behavior), and is safe to call once the consumer has stopped iterating (including after a normal, fully-drained stream, where it's a no-op). Cancellation of a server-side job stays on the `close_stream()` path, driven by `AgentStream.cancel()`. Closes the inner segment generator directly: the outer cancel-guard's `async for … in iterator` does not forward `aclose()` to it, so closing the cancel-guard alone would leave each segment's `request_stream(...)` context manager (and its connection) open until garbage collection. """ # Record that the consumer stopped and the stream was torn down without cancelling. `get()` # reads this to report a still-pending suspended job as resumable `'suspended'` rather than a # live `'incomplete'` snapshot; `_finished`/`_cancelled` take precedence, so this is a no-op # after a fully-drained (`'complete'`) or cancelled (`'interrupted'`) stream. self._detached = True if self._segment_iterator is not None: try: await self._segment_iterator.aclose() except RuntimeError as exc: # A debounced consumer (`group_by_temporal`) can have a prefetch task parked mid-`__anext__` # inside the stitching generator, so `aclose()` raises `RuntimeError: aclose(): asynchronous # generator is already running`. That's exactly the case where a prefetch task exists, and its # own cancellation (when the consumer's debounce is torn down) unwinds the generator and runs # each segment's `request_stream(...)` teardown — so letting this pass still closes the # connection. The model adapters' `close_stream()` avoid this race by routing through # `PeekableAsyncStream.aclose()` (cancelling the in-flight pull first); the stitching generator # here is not `PeekableAsyncStream`-wrapped, so it still needs the suppression. if not _utils.is_async_generator_already_running(exc): raise @property def model_name(self) -> str: if self._current_sub is not None: return self._current_sub.model_name if self._merged_response is not None and self._merged_response.model_name: return self._merged_response.model_name return '' @property def provider_name(self) -> str | None: if self._current_sub is not None: return self._current_sub.provider_name return self._merged_response.provider_name if self._merged_response is not None else None @property def provider_url(self) -> str | None: if self._current_sub is not None: return self._current_sub.provider_url return self._merged_response.provider_url if self._merged_response is not None else None @property def timestamp(self) -> datetime: if self._current_sub is not None: return self._current_sub.timestamp return self._merged_response.timestamp if self._merged_response is not None else _FALLBACK_TIMESTAMP