from __future__ import annotations as _annotations

from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Iterable, Iterator
from contextlib import AbstractAsyncContextManager, aclosing
from copy import deepcopy
from dataclasses import dataclass, field, replace
from datetime import datetime
from decimal import Decimal
from types import TracebackType
from typing import TYPE_CHECKING, Any, Generic, cast, overload

import anyio
from pydantic import ValidationError
from typing_extensions import Self

from . import _utils, exceptions, messages as _messages, models
from ._genai_prices import best_effort_price
from ._output import (
    OutputDataT_inv,
    OutputSchema,
    OutputValidator,
    OutputValidatorFunc,
    TextOutputSchema,
    run_image_process_hooks,
    run_output_with_hooks,
)
from ._run_context import AgentDepsT, RunContext, dispatch_event_stream
from ._sync_stream import SyncStreamBridge
from .messages import AgentStreamEvent, ModelResponseStreamEvent
from .output import (
    OutputDataT,
    ToolOutput,
)
from .tool_manager import ToolManager
from .tools import DeferredToolRequests
from .usage import RunUsage, UsageLimits

if TYPE_CHECKING:
    from .capabilities.abstract import AbstractCapability
    from .run import AgentRunResult

__all__ = (
    'OutputDataT',
    'OutputDataT_inv',
    'ToolOutput',
    'OutputValidatorFunc',
    'StreamedRunResultSync',
)


@dataclass(kw_only=True)
class AgentStream(Generic[AgentDepsT, OutputDataT]):
    _raw_stream_response: models.StreamedResponse
    _output_schema: OutputSchema[OutputDataT]
    _model_request_parameters: models.ModelRequestParameters
    _output_validators: list[OutputValidator[AgentDepsT, OutputDataT]]
    _run_ctx: RunContext[AgentDepsT]
    _usage_limits: UsageLimits | None
    _tool_manager: ToolManager[AgentDepsT]
    _root_capability: AbstractCapability[AgentDepsT]
    _metadata_getter: Callable[[], dict[str, Any] | None] | None = field(default=None, repr=False)
    _event_stream_buffer_getter: Callable[[], list[AgentStreamEvent]] = field(default=list, repr=False)

    _events_iterator: AsyncIterator[AgentStreamEvent] | None = field(default=None, init=False)
    _initial_run_ctx_usage: RunUsage = field(init=False)
    _cached_output: OutputDataT | None = field(default=None, init=False)

    _anext_lock: anyio.Lock = field(default_factory=anyio.Lock, init=False)
    _pull_scopes: set[anyio.CancelScope] = field(default_factory=lambda: set[anyio.CancelScope](), init=False)

    def __post_init__(self):
        self._initial_run_ctx_usage = deepcopy(self._run_ctx.usage)

    async def stream_output(self, *, debounce_by: float | None = 0.1) -> AsyncIterator[OutputDataT]:
        """Asynchronously stream the (validated) agent outputs."""
        if self._cached_output is not None:
            yield deepcopy(self._cached_output)
            return

        last_response: _messages.ModelResponse | None = None
        async for response in self.stream_response(debounce_by=debounce_by):
            if self._raw_stream_response.final_result_event is None or (
                last_response and response.parts == last_response.parts
            ):
                continue
            last_response = response

            try:
                yield await self.validate_response_output(response, allow_partial=True)
            except (ValidationError, exceptions.ModelRetry):
                pass

        if self._raw_stream_response.final_result_event is not None:  # pragma: no branch
            response = self.response
            # Final validation with allow_partial=False (the default).
            # We always yield the final result even if the content matches the last partial yield, because:
            # 1. Output validators/functions receive partial_output=False only on this final call,
            #    and may behave differently based on that flag
            # 2. Users can rely on the last yielded item being the fully validated output
            self._cached_output = await self.validate_response_output(response)
            yield deepcopy(self._cached_output)

    async def stream_response(self, *, debounce_by: float | None = 0.1) -> AsyncIterator[_messages.ModelResponse]:
        """Asynchronously stream the (unvalidated) model responses for the agent.

        Yields `ModelResponse` snapshots — `state='incomplete'` while streaming is in flight,
        followed by one final `state='complete'` snapshot (or `'interrupted'` if `cancel()` was
        called). If the underlying response already has accumulated content when this is called,
        a pre-stream yield surfaces it before iteration begins.
        """
        msg = self.response
        if msg.state == 'incomplete':
            for part in msg.parts:
                if part.has_content():
                    yield msg
                    break

        async with _utils.group_by_temporal(self._model_response_events(), debounce_by) as group_iter:
            async for _items in group_iter:
                yield self.response  # state='incomplete' during streaming

        yield self.response  # final state='complete' (or 'interrupted')

    async def stream_text(self, *, delta: bool = False, debounce_by: float | None = 0.1) -> AsyncIterator[str]:
        """Stream the text result as an async iterable.

        !!! note
            [`TextOutput`][pydantic_ai.output.TextOutput] functions are not applied — use
            [`stream_output()`][pydantic_ai.result.AgentStream.stream_output] instead.
            Result validators will NOT be called on the text result if `delta=True`.

        Args:
            delta: if `True`, yield each chunk of text as it is received, if `False` (default), yield the full text
                up to the current point.
            debounce_by: by how much (if at all) to debounce/group the response chunks by. `None` means no debouncing.
                Debouncing is particularly important for long structured responses to reduce the overhead of
                performing validation as each token is received.
        """
        if not isinstance(self._output_schema, TextOutputSchema):
            raise exceptions.UserError('stream_text() can only be used with text responses')

        # Yield cached output for both delta and non-delta modes
        # This is expected that the subsequent calls to `stream_text()`
        # yield full not delta output even for `delta=True`
        if isinstance(self._cached_output, str):
            yield self._cached_output
            return

        if delta:
            async for text in self._stream_response_text(delta=True, debounce_by=debounce_by):
                yield text
        else:
            async for text in self._stream_response_text(delta=False, debounce_by=debounce_by):
                for validator in self._output_validators:
                    text = await validator.validate(text, replace(self._run_ctx, partial_output=True))
                yield text

    async def cancel(self) -> None:
        """Cancel local stream consumption and request provider shutdown.

        Whether this stops remote generation or closes the underlying transport depends on the provider SDK.

        This stops only the current model response; the run continues. To end the whole run,
        use [`AgentRun.cancel()`][pydantic_ai.run.AgentRun.cancel] or
        [`RunContext.cancel()`][pydantic_ai.tools.RunContext.cancel].
        """
        await self._raw_stream_response.cancel()

    async def drain(self) -> None:
        """Consume all remaining events from the stream, discarding them."""
        async for _ in self:
            pass

    @property
    def cancelled(self) -> bool:
        """Whether the stream has been cancelled via `cancel()`."""
        return self._raw_stream_response.cancelled

    @property
    def run_id(self) -> str:
        """The unique identifier for the agent run."""
        assert self._run_ctx.run_id is not None
        return self._run_ctx.run_id

    @property
    def conversation_id(self) -> str:
        """The unique identifier for the conversation this run belongs to."""
        assert self._run_ctx.conversation_id is not None
        return self._run_ctx.conversation_id

    @property
    def metadata(self) -> dict[str, Any] | None:
        """Metadata associated with this agent run, if configured."""
        if self._metadata_getter is not None:
            return self._metadata_getter()
        return self._run_ctx.metadata

    @property
    def response(self) -> _messages.ModelResponse:
        """Get the current state of the response."""
        return self._raw_stream_response.get()

    @property
    def usage(self) -> RunUsage:
        """Return the usage of the whole run.

        !!! note
            This won't return the full usage until the stream is finished.
        """
        # Mid-stream, `_raw_stream_response.usage` carries no cost yet (it's filled in when the response is
        # appended to history), so add a live best-effort estimate of this request's cost on top of the
        # earlier requests' cost. Once the cost has been filled in, `+` already accounts for it.
        usage = self._initial_run_ctx_usage + self._raw_stream_response.usage
        if self._raw_stream_response.usage.cost is None:
            price = best_effort_price(
                self._raw_stream_response.usage,
                model_name=self.response.model_name,
                provider_api_url=self.response.provider_url,
                provider_name=self.response.provider_name,
                genai_request_timestamp=self.response.timestamp,
            )
            if price is not None:
                usage.cost = (usage.cost or Decimal(0)) + price.total_price
        return usage

    @property
    def timestamp(self) -> datetime:
        """Get the timestamp of the response."""
        return self._raw_stream_response.timestamp

    async def get_output(self) -> OutputDataT:
        """Stream the whole response, validate the output and return it."""
        if self._cached_output is not None:
            return deepcopy(self._cached_output)

        # Iterate through any stream events
        async for _ in self:
            pass

        # Final validation with `allow_partial=False` (default)
        self._cached_output = await self.validate_response_output(self.response)
        return deepcopy(self._cached_output)

    def _settled_output(self) -> tuple[OutputDataT, str | None]:
        """The validated output and the name of the output tool that produced it, if any.

        Every path that consumes a stream to the end validates the response and caches the result,
        so by then `_cached_output` holds it. Cancelling does not: it marks the stream complete
        without a final response, and there is no output to settle on.
        """
        if self._cached_output is None and self.cancelled:
            raise exceptions.UserError(
                'The stream was cancelled before it produced an output, so this run has no settled '
                'result. The messages recorded up to the interruption are still available from '
                '`all_messages()`.'
            )
        final_result_event = self._raw_stream_response.final_result_event
        return (
            cast(OutputDataT, self._cached_output),
            final_result_event.tool_name if final_result_event is not None else None,
        )

    async def validate_response_output(
        self, message: _messages.ModelResponse, *, allow_partial: bool = False
    ) -> OutputDataT:
        """Validate a structured result message."""
        final_result_event = self._raw_stream_response.final_result_event
        if final_result_event is None:
            raise exceptions.UnexpectedModelBehavior('Invalid response, unable to find output')  # pragma: no cover

        output_tool_name = final_result_event.tool_name

        try:
            if self._output_schema.toolset and output_tool_name is not None:
                tool_call = next(
                    (part for part in message.tool_calls if part.tool_name == output_tool_name),
                    None,
                )
                if tool_call is None:
                    raise exceptions.UnexpectedModelBehavior(  # pragma: no cover
                        f'Invalid response, unable to find tool call for {output_tool_name!r}'
                    )
                return await self._tool_manager.handle_output_tool_call(
                    tool_call,
                    schema=self._output_schema,
                    allow_partial=allow_partial,
                    wrap_validation_errors=False,
                )
            elif deferred_tool_requests := _get_deferred_tool_requests(message.tool_calls, self._tool_manager):
                if not self._output_schema.allows_deferred_tools:
                    raise exceptions.UserError(
                        'A deferred tool call was present, but `DeferredToolRequests` is not among output types. To resolve this, add `DeferredToolRequests` to the list of output types for this agent.'
                    )
                return cast(OutputDataT, deferred_tool_requests)
            elif self._output_schema.allows_image and message.images:
                return await self._validate_image_output(message.images[0], allow_partial=allow_partial)
            elif text_processor := self._output_schema.text_processor:
                text = ''
                for part in message.parts:
                    if isinstance(part, _messages.TextPart):
                        text += part.content
                    elif isinstance(part, _messages.NativeToolCallPart):
                        # Text parts before a built-in tool call are essentially thoughts,
                        # not part of the final result output, so we reset the accumulated text
                        text = ''

                run_ctx = replace(self._run_ctx, partial_output=allow_partial)
                return await run_output_with_hooks(
                    text_processor,
                    text=text,
                    run_context=run_ctx,
                    capability=self._root_capability,
                    schema=self._output_schema,
                    allow_partial=allow_partial,
                    wrap_validation_errors=False,
                    output_validators=self._output_validators,
                )
            else:
                raise exceptions.UnexpectedModelBehavior(  # pragma: no cover
                    'Invalid response, unable to process text output'
                )
        except (ValidationError, exceptions.ModelRetry) as e:
            if not allow_partial:
                raise exceptions.UnexpectedModelBehavior(
                    'Output validation failed during streaming, and retries are not supported in `run_stream()`'
                ) from e
            raise

    async def _validate_image_output(self, image: _messages.BinaryImage, *, allow_partial: bool) -> OutputDataT:
        """Run process hooks (including output validators) for image output."""
        run_ctx = replace(self._run_ctx, partial_output=allow_partial)
        return cast(
            OutputDataT,
            await run_image_process_hooks(
                image,
                capability=self._root_capability,
                run_context=run_ctx,
                schema=self._output_schema,
                wrap_validation_errors=False,
                output_validators=self._output_validators,
            ),
        )

    async def _stream_response_text(
        self, *, delta: bool = False, debounce_by: float | None = 0.1
    ) -> AsyncIterator[str]:
        """Stream the response as an async iterable of text."""

        # Define a "merged" version of the iterator that will yield items that have already been retrieved
        # and items that we receive while streaming. We define a dedicated async iterator for this so we can
        # pass the combined stream to the group_by_temporal function within `_stream_text_deltas` below.
        async def _stream_text_deltas_ungrouped() -> AsyncIterator[tuple[str, int]]:
            # yields tuples of (text_content, part_index)
            # we don't currently make use of the part_index, but in principle this may be useful
            # so we retain it here for now to make possible future refactors simpler
            msg = self.response
            for i, part in enumerate(msg.parts):
                if isinstance(part, _messages.TextPart) and part.content:
                    yield part.content, i

            last_text_index: int | None = None
            async for event in self:
                if (
                    isinstance(event, _messages.PartStartEvent)
                    and isinstance(event.part, _messages.TextPart)
                    and event.part.content
                ):
                    last_text_index = event.index
                    yield event.part.content, event.index
                elif (
                    isinstance(event, _messages.PartDeltaEvent)
                    and isinstance(event.delta, _messages.TextPartDelta)
                    and event.delta.content_delta
                ):
                    last_text_index = event.index
                    yield event.delta.content_delta, event.index
                elif (
                    isinstance(event, _messages.PartStartEvent)
                    and isinstance(event.part, _messages.NativeToolCallPart)
                    and last_text_index is not None
                ):
                    # Text parts that are interrupted by a built-in tool call should not be joined together directly
                    yield '\n\n', event.index
                    last_text_index = None

        async def _stream_text_deltas() -> AsyncGenerator[str, None]:
            async with _utils.group_by_temporal(_stream_text_deltas_ungrouped(), debounce_by) as group_iter:
                async for items in group_iter:
                    # Note: we are currently just dropping the part index on the group here
                    yield ''.join([content for content, _ in items])

        async with aclosing(_stream_text_deltas()) as deltas_iter:
            if delta:
                async for text in deltas_iter:
                    yield text
            else:
                # a quick benchmark shows it's faster to build up a string with concat when we're
                # yielding at each step
                deltas: list[str] = []
                async for text in deltas_iter:
                    deltas.append(text)
                    yield ''.join(deltas)

    def __aiter__(self) -> AsyncIterator[AgentStreamEvent]:
        """Stream [`AgentStreamEvent`][pydantic_ai.messages.AgentStreamEvent]s, interleaving events emitted into the run's event buffer."""
        if self._events_iterator is None:
            # Token-limit checks run after every event and only look at token counts, so skip the per-event
            # cost calculation that the `usage` property does and pass the cheaper token-only usage.
            base_iter = _get_usage_checking_stream_response(
                self._raw_stream_response,
                self._usage_limits,
                lambda: self._initial_run_ctx_usage + self._raw_stream_response.usage,
            )
            # Wrap once, so a capability's `wrap_run_event_stream` sees each event exactly once no
            # matter how many times this stream is iterated (e.g. `stream_text()` then a drain).
            self._events_iterator = aiter(
                self._root_capability.wrap_run_event_stream(
                    self._run_ctx, stream=dispatch_event_stream(self._run_ctx, self._events_iter(base_iter))
                )
            )

        return self._pull_shared(self._events_iterator)

    async def aclose_events(self) -> None:
        """Close the event stream when a consumer walks away before exhausting it.

        The event iterator owns the capability chain, which can otherwise stay suspended with
        resources held, like a `ProcessEventStream` handler task parked on its receive stream.

        Any in-flight shared pull is cancelled and drained before the iterator is closed. The
        close is shielded because graph teardown can run inside an already-cancelled scope.

        The closed iterator is kept in place rather than discarded, so a later `__aiter__()` ends
        immediately instead of building a second chain (and a second handler) over a spent stream.
        """
        events_iterator = self._events_iterator
        if events_iterator is not None:
            for scope in self._pull_scopes:
                scope.cancel()
            with anyio.CancelScope(shield=True):
                async with self._anext_lock:
                    await _utils.aclose_if_supported(events_iterator)

    async def _pull_shared(self, events_iterator: AsyncIterator[AgentStreamEvent]) -> AsyncIterator[AgentStreamEvent]:
        # Serialize access to the shared iterator. An early break from stream_text() can leave a
        # pending `anext()` task in group_by_temporal while cleanup/drain starts iterating the same
        # stream.
        while True:
            async with self._anext_lock:
                event: AgentStreamEvent | None = None
                with anyio.CancelScope() as scope:
                    self._pull_scopes.add(scope)
                    try:
                        try:
                            event = await anext(events_iterator)
                        except StopAsyncIteration:
                            return
                    finally:
                        self._pull_scopes.discard(scope)
                if scope.cancel_called:
                    return
                assert event is not None
            yield event

    async def _model_response_events(self) -> AsyncIterator[ModelResponseStreamEvent]:
        """Iterate only the model response stream events, dropping events emitted into the run's event buffer."""
        async for event in self:
            if isinstance(
                event,
                _messages.PartStartEvent
                | _messages.PartDeltaEvent
                | _messages.PartEndEvent
                | _messages.FinalResultEvent,
            ):
                yield event

    async def _events_iter(self, base_iter: AsyncIterator[ModelResponseStreamEvent]) -> AsyncIterator[AgentStreamEvent]:
        while True:
            # Drain events emitted into the run's event buffer before each pull, so they interleave with the
            # model's own events. Events emitted while a pull is in flight surface on the next pull,
            # or through the response-handling node's stream once this stream is exhausted.
            while buffer := self._event_stream_buffer_getter():
                yield buffer.pop(0)

            try:
                event = await anext(base_iter)
            except StopAsyncIteration:
                return

            yield event


@dataclass(init=False)
class StreamedRunResult(Generic[AgentDepsT, OutputDataT]):
    """Result of a streamed run that returns structured data via a tool call."""

    _all_messages: list[_messages.ModelMessage]
    _new_message_index: int

    _stream_response: AgentStream[AgentDepsT, OutputDataT] | None = None
    _on_complete: Callable[[], Awaitable[None]] | None = None

    _run_result: AgentRunResult[OutputDataT] | None = None

    is_complete: bool = field(default=False, init=False)
    """Whether the stream has all been received.

    This is set to `True` when one of
    [`stream_output`][pydantic_ai.result.StreamedRunResult.stream_output],
    [`stream_text`][pydantic_ai.result.StreamedRunResult.stream_text],
    [`stream_response`][pydantic_ai.result.StreamedRunResult.stream_response] or
    [`get_output`][pydantic_ai.result.StreamedRunResult.get_output] completes.
    """

    @overload
    def __init__(
        self,
        all_messages: list[_messages.ModelMessage],
        new_message_index: int,
        stream_response: AgentStream[AgentDepsT, OutputDataT] | None,
        on_complete: Callable[[], Awaitable[None]] | None,
    ) -> None: ...

    @overload
    def __init__(
        self,
        all_messages: list[_messages.ModelMessage],
        new_message_index: int,
        *,
        run_result: AgentRunResult[OutputDataT],
    ) -> None: ...

    def __init__(
        self,
        all_messages: list[_messages.ModelMessage],
        new_message_index: int,
        stream_response: AgentStream[AgentDepsT, OutputDataT] | None = None,
        on_complete: Callable[[], Awaitable[None]] | None = None,
        run_result: AgentRunResult[OutputDataT] | None = None,
    ) -> None:
        self._all_messages = all_messages
        self._new_message_index = new_message_index

        self._stream_response = stream_response
        self._on_complete = on_complete
        self._run_result = run_result
        self._traceparent_value: str | None = None
        """Captured when the stream finishes, while the agent run span is still open.

        A settled result is for handing the run to code that outlives the stream, which is exactly
        when the span has closed and the ambient trace context is gone, so it cannot be read lazily.
        """

    @property
    def result(self) -> AgentRunResult[OutputDataT]:
        """This run as an [`AgentRunResult`][pydantic_ai.run.AgentRunResult], once the stream has finished.

        A `StreamedRunResult` reads its values off the stream that is producing them, so it lives only as
        long as that stream does. `AgentRunResult` is the settled form of the same run — the same output,
        messages, usage and IDs — and it has a
        [serialized shape](../message-history.md#storing-complete-run-results) you can store and load. Reach
        for this to hand the run to code that outlives the stream.

        Raises:
            UserError: If the stream hasn't finished, or was cancelled before producing an output,
                so the run has no settled output to settle on.
        """
        from ._agent_graph import GraphAgentState
        from .run import AgentRunResult

        if (run_result := self._run_result) is not None:
            return run_result
        elif self._stream_response is not None:
            if not self.is_complete:
                raise exceptions.UserError(
                    'The run is still streaming, so it has no settled result yet. Await `stream_output()`, '
                    '`stream_text()`, `stream_response()` or `get_output()` first.'
                )
            output, output_tool_name = self._stream_response._settled_output()  # pyright: ignore[reportPrivateUsage]
            return AgentRunResult(
                output=output,
                _output_tool_name=output_tool_name,
                _state=GraphAgentState(
                    message_history=self._all_messages,
                    usage=self.usage,
                    run_id=self.run_id,
                    conversation_id=self.conversation_id,
                    metadata=self.metadata,
                ),
                _new_message_index=self._new_message_index,
                _traceparent_value=self._traceparent_value,
            )
        else:
            raise ValueError('No stream response or run result provided')  # pragma: no cover

    def all_messages(self, *, output_tool_return_content: str | None = None) -> list[_messages.ModelMessage]:
        """Return the history of _messages.

        Args:
            output_tool_return_content: The return content of the tool call to set in the last message.
                This provides a convenient way to modify the content of the output tool call if you want to continue
                the conversation and want to set the response to the output tool call. If `None`, the last message will
                not be modified.

        Returns:
            List of messages.
        """
        # this is a method to be consistent with the other methods
        if output_tool_return_content is not None:
            raise NotImplementedError('Setting output tool return content is not supported for this result type.')
        return self._all_messages

    def all_messages_json(self, *, output_tool_return_content: str | None = None) -> bytes:
        """Return all messages from [`all_messages`][pydantic_ai.result.StreamedRunResult.all_messages] as JSON bytes.

        Args:
            output_tool_return_content: The return content of the tool call to set in the last message.
                This provides a convenient way to modify the content of the output tool call if you want to continue
                the conversation and want to set the response to the output tool call. If `None`, the last message will
                not be modified.

        Returns:
            JSON bytes representing the messages.
        """
        return _messages.ModelMessagesTypeAdapter.dump_json(
            self.all_messages(output_tool_return_content=output_tool_return_content)
        )

    def new_messages(self, *, output_tool_return_content: str | None = None) -> list[_messages.ModelMessage]:
        """Return the messages produced during this run.

        Messages provided via `message_history` and messages from older runs are excluded.

        Args:
            output_tool_return_content: The return content of the tool call to set in the last message.
                This provides a convenient way to modify the content of the output tool call if you want to continue
                the conversation and want to set the response to the output tool call. If `None`, the last message will
                not be modified.

        Returns:
            List of new messages.
        """
        return self.all_messages(output_tool_return_content=output_tool_return_content)[self._new_message_index :]

    def new_messages_json(self, *, output_tool_return_content: str | None = None) -> bytes:  # pragma: no cover
        """Return new messages from [`new_messages`][pydantic_ai.result.StreamedRunResult.new_messages] as JSON bytes.

        Args:
            output_tool_return_content: The return content of the tool call to set in the last message.
                This provides a convenient way to modify the content of the output tool call if you want to continue
                the conversation and want to set the response to the output tool call. If `None`, the last message will
                not be modified.

        Returns:
            JSON bytes representing the new messages.
        """
        return _messages.ModelMessagesTypeAdapter.dump_json(
            self.new_messages(output_tool_return_content=output_tool_return_content)
        )

    async def stream_output(self, *, debounce_by: float | None = 0.1) -> AsyncIterator[OutputDataT]:
        """Stream the output as an async iterable.

        The pydantic validator for structured data will be called in
        [partial mode](https://docs.pydantic.dev/dev/concepts/experimental/#partial-validation)
        on each iteration.

        Args:
            debounce_by: by how much (if at all) to debounce/group the output chunks by. `None` means no debouncing.
                Debouncing is particularly important for long structured outputs to reduce the overhead of
                performing validation as each token is received.

        Returns:
            An async iterable of the response data.
        """
        if self._run_result is not None:
            yield self._run_result.output
            await self._marked_completed()
        elif self._stream_response is not None:
            async for output in self._stream_response.stream_output(debounce_by=debounce_by):
                yield output
            await self._marked_completed(self.response)
        else:
            raise ValueError('No stream response or run result provided')  # pragma: no cover

    async def stream_text(self, *, delta: bool = False, debounce_by: float | None = 0.1) -> AsyncIterator[str]:
        """Stream the text result as an async iterable.

        !!! note
            [`TextOutput`][pydantic_ai.output.TextOutput] functions are not applied — use
            [`stream_output()`][pydantic_ai.result.StreamedRunResult.stream_output] instead.
            Result validators will NOT be called on the text result if `delta=True`.

        Args:
            delta: if `True`, yield each chunk of text as it is received, if `False` (default), yield the full text
                up to the current point.
            debounce_by: by how much (if at all) to debounce/group the response chunks by. `None` means no debouncing.
                Debouncing is particularly important for long structured responses to reduce the overhead of
                performing validation as each token is received.
        """
        if self._run_result is not None:  # pragma: no cover
            # We can't really get here, as `_run_result` is only set in `run_stream` when `CallToolsNode` produces `DeferredToolRequests` output
            # as a result of a tool function raising `CallDeferred` or `ApprovalRequired`.
            # That'll change if we ever support something like `raise EndRun(output: OutputT)` where `OutputT` could be `str`.
            if not isinstance(self._run_result.output, str):
                raise exceptions.UserError('stream_text() can only be used with text responses')
            yield self._run_result.output
            await self._marked_completed()
        elif self._stream_response is not None:
            async for text in self._stream_response.stream_text(delta=delta, debounce_by=debounce_by):
                yield text
            await self._marked_completed(self.response)
        else:
            raise ValueError('No stream response or run result provided')  # pragma: no cover

    async def stream_response(self, *, debounce_by: float | None = 0.1) -> AsyncIterator[_messages.ModelResponse]:
        """Stream the response as an async iterable of `ModelResponse` snapshots.

        Each yielded `ModelResponse` is the current state of the response: `response.state` is
        `'incomplete'` while streaming is in flight and `'complete'` (or `'interrupted'` if
        [`cancel()`][pydantic_ai.result.StreamedRunResult.cancel] was called) on the final yield.

        Args:
            debounce_by: by how much (if at all) to debounce/group the response chunks by. `None` means no debouncing.
                Debouncing is particularly important for long structured responses to reduce the overhead of
                performing validation as each token is received.

        Returns:
            An async iterable of `ModelResponse` snapshots.
        """
        if self._run_result is not None:
            yield self.response
            await self._marked_completed()
        elif self._stream_response is not None:
            last_msg: _messages.ModelResponse | None = None
            async for msg in self._stream_response.stream_response(debounce_by=debounce_by):
                yield msg
                last_msg = msg
            # `AgentStream.stream_response` always yields the final response, so `last_msg` is set.
            # Pass it to `_marked_completed` so `run_id` and `conversation_id` are stamped onto the
            # same instance the caller still holds a reference to in their iteration.
            assert last_msg is not None
            await self._marked_completed(last_msg)
        else:
            raise ValueError('No stream response or run result provided')  # pragma: no cover

    async def get_output(self) -> OutputDataT:
        """Stream the whole response, validate and return it."""
        if self._run_result is not None:
            output = self._run_result.output
            await self._marked_completed()
            return output
        elif self._stream_response is not None:
            output = await self._stream_response.get_output()
            await self._marked_completed(self.response)
            return output
        else:
            raise ValueError('No stream response or run result provided')  # pragma: no cover

    @property
    def response(self) -> _messages.ModelResponse:
        """Return the current state of the response."""
        if self._run_result is not None:
            return self._run_result.response
        elif self._stream_response is not None:
            return self._stream_response.response
        else:
            raise ValueError('No stream response or run result provided')  # pragma: no cover

    @property
    def metadata(self) -> dict[str, Any] | None:
        """Metadata associated with this agent run, if configured."""
        if self._run_result is not None:
            return self._run_result.metadata
        elif self._stream_response is not None:
            return self._stream_response.metadata
        else:
            return None

    @property
    def usage(self) -> RunUsage:
        """Return the usage of the whole run.

        !!! note
            This won't return the full usage until the stream is finished.
        """
        if self._run_result is not None:
            return self._run_result.usage
        elif self._stream_response is not None:
            return self._stream_response.usage
        else:
            raise ValueError('No stream response or run result provided')  # pragma: no cover

    @property
    def timestamp(self) -> datetime:
        """Get the timestamp of the response."""
        if self._run_result is not None:
            return self._run_result.timestamp
        elif self._stream_response is not None:
            return self._stream_response.timestamp
        else:
            raise ValueError('No stream response or run result provided')  # pragma: no cover

    @property
    def run_id(self) -> str:
        """The unique identifier for the agent run."""
        if self._run_result is not None:
            return self._run_result.run_id
        elif self._stream_response is not None:
            return self._stream_response.run_id
        else:
            raise ValueError('No stream response or run result provided')  # pragma: no cover

    @property
    def conversation_id(self) -> str:
        """The unique identifier for the conversation this run belongs to."""
        if self._run_result is not None:
            return self._run_result.conversation_id
        elif self._stream_response is not None:
            return self._stream_response.conversation_id
        else:
            raise ValueError('No stream response or run result provided')  # pragma: no cover

    async def validate_response_output(
        self, message: _messages.ModelResponse, *, allow_partial: bool = False
    ) -> OutputDataT:
        """Validate a structured result message."""
        if self._run_result is not None:
            return self._run_result.output
        elif self._stream_response is not None:
            return await self._stream_response.validate_response_output(message, allow_partial=allow_partial)
        else:
            raise ValueError('No stream response or run result provided')  # pragma: no cover

    def _record_response(self, message: _messages.ModelResponse) -> None:
        """Append a model response to the message history with the correct run and conversation IDs."""
        if self._stream_response:  # pragma: no branch
            message.run_id = self._stream_response.run_id
            message.conversation_id = self._stream_response.conversation_id
        self._all_messages.append(message)

    async def _marked_completed(self, message: _messages.ModelResponse | None = None) -> None:
        from ._instrumentation import current_otel_traceparent

        if self.is_complete:
            return
        self.is_complete = True
        self._traceparent_value = current_otel_traceparent()
        if message is not None:
            self._record_response(message)
        if self._on_complete is not None:
            await self._on_complete()

    async def cancel(self) -> None:
        """Cancel local stream consumption and request provider shutdown.

        Whether this stops remote generation or closes the underlying transport depends on the provider SDK.

        The interrupted response state is recorded in the message history so that
        `all_messages()` includes it.

        This stops only the current model response; the run continues. To end the whole run,
        use [`AgentRun.cancel()`][pydantic_ai.run.AgentRun.cancel] or
        [`RunContext.cancel()`][pydantic_ai.tools.RunContext.cancel].
        """
        if self._stream_response is not None:  # pragma: no branch
            await self._stream_response.cancel()
            # Record the interrupted response in _all_messages so all_messages()
            # includes it. is_complete guard prevents double-append if the stream
            # was already fully consumed before cancel was called.
            if not self.is_complete:
                self.is_complete = True
                self._record_response(self.response)

    @property
    def cancelled(self) -> bool:
        """Whether the stream has been cancelled via `cancel()`."""
        if self._stream_response is not None:
            return self._stream_response.cancelled
        # Only reachable via a `wrap_run` short-circuit, where there is no stream.
        return False  # pragma: no cover


class StreamedRunResultSync(Generic[AgentDepsT, OutputDataT]):
    """Synchronous wrapper for [`StreamedRunResult`][pydantic_ai.result.StreamedRunResult] that only exposes sync methods.

    All of the run's async work happens on the caller's event loop. Context-manager and iterator
    lifecycles remain in stable tasks, so cancel scopes entered and exited by the agent graph never
    straddle tasks and OpenTelemetry spans stay correctly nested. The wrapper must be used and closed
    on the thread where it was created.

    This is a synchronous context manager; the underlying stream is cleaned up on exit:

    ```python
    from pydantic_ai import Agent

    agent = Agent('openai:gpt-5.2')

    def main():
        with agent.run_stream_sync('What is the capital of the UK?') as response:
            print(response.get_output())
            #> The capital of the UK is London.
    ```

    Using it without a `with` block also works for backwards compatibility. Garbage collection requests
    best-effort cleanup on the owner loop, but it cannot drive a stopped owner loop from another thread
    or while another loop is running. A `with` block should be used whenever deterministic cleanup matters.
    """

    _streamed_run_result: StreamedRunResult[AgentDepsT, OutputDataT]

    def __init__(self, run_stream_cm: AbstractAsyncContextManager[StreamedRunResult[AgentDepsT, OutputDataT]]) -> None:
        if isinstance(run_stream_cm, StreamedRunResult):
            # This wrapper used to take an already-entered `StreamedRunResult`, but it now needs the
            # `run_stream()` context manager so it can enter it in a stable task. Construct it via
            # `agent.run_stream_sync(...)` instead. TODO (v3): remove this check.
            raise TypeError(
                '`StreamedRunResultSync` now takes the `run_stream()` context manager rather than an '
                'already-entered `StreamedRunResult`; use `agent.run_stream_sync(...)` to construct it.'
            )
        self._bridge = SyncStreamBridge(run_stream_cm, async_alternative='`run_stream`')
        self._streamed_run_result = self._bridge.stream

    def __enter__(self) -> Self:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        self._bridge.shutdown((exc_type, exc_val, exc_tb))

    def all_messages(self, *, output_tool_return_content: str | None = None) -> list[_messages.ModelMessage]:
        """Return the history of messages.

        Args:
            output_tool_return_content: The return content of the tool call to set in the last message.
                This provides a convenient way to modify the content of the output tool call if you want to continue
                the conversation and want to set the response to the output tool call. If `None`, the last message will
                not be modified.

        Returns:
            List of messages.
        """
        return self._streamed_run_result.all_messages(output_tool_return_content=output_tool_return_content)

    def all_messages_json(self, *, output_tool_return_content: str | None = None) -> bytes:  # pragma: no cover
        """Return all messages from [`all_messages`][pydantic_ai.result.StreamedRunResultSync.all_messages] as JSON bytes.

        Args:
            output_tool_return_content: The return content of the tool call to set in the last message.
                This provides a convenient way to modify the content of the output tool call if you want to continue
                the conversation and want to set the response to the output tool call. If `None`, the last message will
                not be modified.

        Returns:
            JSON bytes representing the messages.
        """
        return self._streamed_run_result.all_messages_json(output_tool_return_content=output_tool_return_content)

    def new_messages(self, *, output_tool_return_content: str | None = None) -> list[_messages.ModelMessage]:
        """Return the messages produced during this run.

        Messages provided via `message_history` and messages from older runs are excluded.

        Args:
            output_tool_return_content: The return content of the tool call to set in the last message.
                This provides a convenient way to modify the content of the output tool call if you want to continue
                the conversation and want to set the response to the output tool call. If `None`, the last message will
                not be modified.

        Returns:
            List of new messages.
        """
        return self._streamed_run_result.new_messages(output_tool_return_content=output_tool_return_content)

    def new_messages_json(self, *, output_tool_return_content: str | None = None) -> bytes:  # pragma: no cover
        """Return new messages from [`new_messages`][pydantic_ai.result.StreamedRunResultSync.new_messages] as JSON bytes.

        Args:
            output_tool_return_content: The return content of the tool call to set in the last message.
                This provides a convenient way to modify the content of the output tool call if you want to continue
                the conversation and want to set the response to the output tool call. If `None`, the last message will
                not be modified.

        Returns:
            JSON bytes representing the new messages.
        """
        return self._streamed_run_result.new_messages_json(output_tool_return_content=output_tool_return_content)

    def stream_output(self, *, debounce_by: float | None = 0.1) -> Iterator[OutputDataT]:
        """Stream the output as an iterable.

        The pydantic validator for structured data will be called in
        [partial mode](https://docs.pydantic.dev/dev/concepts/experimental/#partial-validation)
        on each iteration.

        Args:
            debounce_by: by how much (if at all) to debounce/group the output chunks by. `None` means no debouncing.
                Debouncing is particularly important for long structured outputs to reduce the overhead of
                performing validation as each token is received.

        Returns:
            An iterable of the response data.
        """
        result = self._streamed_run_result
        return self._bridge.stream_sync(lambda: result.stream_output(debounce_by=debounce_by))

    def stream_text(self, *, delta: bool = False, debounce_by: float | None = 0.1) -> Iterator[str]:
        """Stream the text result as an iterable.

        !!! note
            [`TextOutput`][pydantic_ai.output.TextOutput] functions are not applied — use
            [`stream_output()`][pydantic_ai.result.StreamedRunResultSync.stream_output] instead.
            Result validators will NOT be called on the text result if `delta=True`.

        Args:
            delta: if `True`, yield each chunk of text as it is received, if `False` (default), yield the full text
                up to the current point.
            debounce_by: by how much (if at all) to debounce/group the response chunks by. `None` means no debouncing.
                Debouncing is particularly important for long structured responses to reduce the overhead of
                performing validation as each token is received.
        """
        result = self._streamed_run_result
        return self._bridge.stream_sync(lambda: result.stream_text(delta=delta, debounce_by=debounce_by))

    def stream_response(self, *, debounce_by: float | None = 0.1) -> Iterator[_messages.ModelResponse]:
        """Stream the response as an iterable of `ModelResponse` snapshots.

        Each yielded `ModelResponse` is the current state of the response: `response.state` is
        `'incomplete'` while streaming is in flight and `'complete'` on the final yield.

        Args:
            debounce_by: by how much (if at all) to debounce/group the response chunks by. `None` means no debouncing.
                Debouncing is particularly important for long structured responses to reduce the overhead of
                performing validation as each token is received.

        Returns:
            An iterable of `ModelResponse` snapshots.
        """
        result = self._streamed_run_result
        return self._bridge.stream_sync(lambda: result.stream_response(debounce_by=debounce_by))

    def get_output(self) -> OutputDataT:
        """Stream the whole response, validate and return it."""
        return self._bridge.call(self._streamed_run_result.get_output)

    @property
    def response(self) -> _messages.ModelResponse:
        """Return the current state of the response."""
        return self._streamed_run_result.response

    @property
    def usage(self) -> RunUsage:
        """Return the usage of the whole run.

        !!! note
            This won't return the full usage until the stream is finished.
        """
        return self._streamed_run_result.usage

    @property
    def timestamp(self) -> datetime:
        """Get the timestamp of the response."""
        return self._streamed_run_result.timestamp

    @property
    def run_id(self) -> str:
        """The unique identifier for the agent run."""
        return self._streamed_run_result.run_id

    @property
    def conversation_id(self) -> str:
        """The unique identifier for the conversation this run belongs to."""
        return self._streamed_run_result.conversation_id

    @property
    def metadata(self) -> dict[str, Any] | None:
        """Metadata associated with this agent run, if configured."""
        return self._streamed_run_result.metadata

    def validate_response_output(self, message: _messages.ModelResponse, *, allow_partial: bool = False) -> OutputDataT:
        """Validate a structured result message."""
        return self._bridge.call(
            lambda: self._streamed_run_result.validate_response_output(message, allow_partial=allow_partial),
        )

    @property
    def is_complete(self) -> bool:
        """Whether the stream has all been received.

        This is set to `True` when one of
        [`stream_output`][pydantic_ai.result.StreamedRunResultSync.stream_output],
        [`stream_text`][pydantic_ai.result.StreamedRunResultSync.stream_text],
        [`stream_response`][pydantic_ai.result.StreamedRunResultSync.stream_response] or
        [`get_output`][pydantic_ai.result.StreamedRunResultSync.get_output] completes.
        """
        return self._streamed_run_result.is_complete


@dataclass(repr=False)
class FinalResult(Generic[OutputDataT]):
    """Marker class storing the final output of an agent run and associated metadata."""

    output: OutputDataT
    """The final result data."""

    tool_name: str | None = None
    """Name of the final output tool; `None` if the output came from unstructured text content."""

    tool_call_id: str | None = None
    """ID of the tool call that produced the final output; `None` if the output came from unstructured text content."""

    __repr__ = _utils.dataclasses_no_defaults_repr


def _get_usage_checking_stream_response(
    stream_response: models.StreamedResponse,
    limits: UsageLimits | None,
    get_usage: Callable[[], RunUsage],
) -> AsyncIterator[ModelResponseStreamEvent]:
    if limits is not None and limits.has_token_limits():

        async def _usage_checking_iterator():
            async for item in stream_response:
                limits.check_tokens(get_usage())
                limits.check_per_request_input_tokens(stream_response.usage.input_tokens)
                yield item

        return _usage_checking_iterator()
    else:
        return aiter(stream_response)


def _get_deferred_tool_requests(
    tool_calls: Iterable[_messages.ToolCallPart], tool_manager: ToolManager[AgentDepsT]
) -> DeferredToolRequests | None:
    """Get the deferred tool requests from the model response tool calls."""
    approvals: list[_messages.ToolCallPart] = []
    calls: list[_messages.ToolCallPart] = []

    for tool_call in tool_calls:
        tool_def = tool_manager.get_tool_def(tool_call.tool_name)
        if tool_def is not None:  # pragma: no branch
            if tool_def.kind == 'unapproved':
                approvals.append(tool_call)
            elif tool_def.kind == 'external':
                calls.append(tool_call)

    if not calls and not approvals:
        return None

    return DeferredToolRequests(calls=calls, approvals=approvals)