Function run_query_as_agent_loop

src/query/agent_loop_compat.py:273–577

Drive the canonical query() loop and adapt to AgentLoopResult. Parameters mirror the existing ``run_agent_loop`` signature where practical, but the adapter takes ``initial_messages`` directly instead of a ``Conversation`` — the typed messages are the same shape query() consumes nati

(
    *,
    initial_messages: list[Message],
    provider: BaseProvider,
    tool_registry: ToolRegistry,
    tool_context: ToolContext,
    system_prompt: str | list[dict[str, Any]] = "You are a helpful assistant.",
    max_turns: int = 20,
    on_event: ToolEventHandler | None = None,
    on_text_chunk: TextChunkHandler | None = None,
    on_thinking_chunk: TextChunkHandler | None = None,
    on_message: Callable[[Message], None] | None = None,
    cancel_signal: AbortSignal | None = None,
    abort_controller: AbortController | None = None,
    extended_thinking: bool | None = None,
)

async def run_query_as_agent_loop(
    *,
    initial_messages: list[Message],
    provider: BaseProvider,
    tool_registry: ToolRegistry,
    tool_context: ToolContext,
    system_prompt: str | list[dict[str, Any]] = "You are a helpful assistant.",
    max_turns: int = 20,
    on_event: ToolEventHandler | None = None,
    on_text_chunk: TextChunkHandler | None = None,
    on_thinking_chunk: TextChunkHandler | None = None,
    on_message: Callable[[Message], None] | None = None,
    cancel_signal: AbortSignal | None = None,
    abort_controller: AbortController | None = None,
    extended_thinking: bool | None = None,
) -> AgentLoopRunResult:
    """Drive the canonical query() loop and adapt to AgentLoopResult.

    Parameters mirror the existing ``run_agent_loop`` signature where
    practical, but the adapter takes ``initial_messages`` directly
    instead of a ``Conversation`` — the typed messages are the same
    shape query() consumes natively.

    The ``on_event`` callback receives :class:`ToolEvent` instances
    for every tool_use observed in the model's responses and every
    tool_result yielded by the loop.

    ``on_text_chunk`` is forwarded into the QueryParams so the
    provider's streaming layer fires text chunks LIVE (per-delta).
    Callers MUST provide this if they need real-time text rendering
    (TUI live streaming, ESC-mid-stream cancel teardown).

    ``on_message`` is fired for EVERY :class:`Message` yielded by the
    loop (Anthropic-shape AssistantMessage with full content blocks
    including tool_use, UserMessage with tool_result blocks, etc.).
    Use this to persist the full conversation transcript faithfully —
    `response_text` alone loses tool_use/tool_result structure across
    multi-turn sessions.

    ``cancel_signal`` is bridged into the loop's abort_controller so
    user-initiated cancels (Ctrl+C, /exit) propagate cleanly. When
    not supplied, the function constructs its own AbortController.
    """
    # Critic C2 fix: do NOT mint a fresh controller when the caller
    # provided one. The provider's chat_stream_response listens on
    # ``QueryParams.abort_controller.signal`` to tear down HTTP streams
    # mid-flight on ESC. A fresh controller breaks that wiring — ESC
    # would flip the user's signal but the provider would never see it
    # because the per-message bridge below only fires when query()
    # yields a message, and a tool-use-only turn yields nothing during
    # the multi-second generation. Caller's controller IS the user's
    # signal source; reuse it.
    if abort_controller is None:
        if cancel_signal is not None:
            # We received only the signal, not its owning controller.
            # Mint a new controller and bridge cancellation into it
            # both pre- and per-iteration (legacy fallback path).
            abort_controller = AbortController()