chore: import upstream snapshot with attribution

python/samples/02-agents/harness/console/agent_runner.py

@@ -0,0 +1,427 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Agent runner orchestration for the harness console.
+
+This module provides the HarnessAgentRunner class, which orchestrates agent
+invocations with observer lifecycle management. It handles:
+- User input dispatch
+- Agent streaming with observer notifications
+- Follow-up action collection
+- Streaming state management
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from agent_framework import Agent, AgentSession, MessageInjectionMiddleware
+    from agent_framework import Message as FrameworkMessage
+
+    from .app_state import FollowUpAction
+    from .observers.base import ConsoleObserver
+    from .state_driver import IUXStateDriver
+
+
+class HarnessAgentRunner:
+    """Orchestrates agent invocations driven by user-input events from the UI.
+
+    The component invokes the runner's input handlers (run_turn) directly;
+    the runner mutates UI state through the supplied IUXStateDriver.
+
+    When the underlying agent has a ``MessageInjectionMiddleware`` wired in
+    (as ``create_harness_agent`` does by default), the runner supports message
+    injection: input submitted while a turn is streaming is enqueued via
+    ``on_streaming_input`` and drained into the ongoing run by the middleware.
+    """
+
+    def __init__(
+        self,
+        agent: Agent,
+        observers: list[ConsoleObserver],
+        state_driver: IUXStateDriver,
+        *,
+        max_context_window_tokens: int | None = None,
+        max_output_tokens: int | None = None,
+    ) -> None:
+        """Initialize the agent runner.
+
+        Args:
+            agent: The agent to orchestrate.
+            observers: List of console observers for lifecycle events.
+            state_driver: The UI state driver for observer updates.
+            max_context_window_tokens: Optional max context window size for usage display.
+            max_output_tokens: Optional max output tokens for usage display.
+        """
+        self._agent = agent
+        self._observers = observers
+        self._ux = state_driver
+        self._max_context_window_tokens = max_context_window_tokens
+        self._max_output_tokens = max_output_tokens
+        self._input_gate = asyncio.Semaphore(1)  # Single turn at a time
+
+        # Resolve the message-injection middleware (if any) so streaming-time
+        # input can be enqueued into the ongoing run. Absent => injection no-ops.
+        from agent_framework import MessageInjectionMiddleware
+
+        self._message_injector: MessageInjectionMiddleware | None = next(
+            (m for m in (agent.middleware or []) if isinstance(m, MessageInjectionMiddleware)),
+            None,
+        )
+        # Snapshot of pending injected messages, used to detect consumption
+        # during streaming. Safe as instance state because _input_gate
+        # serialises turns.
+        self._last_pending_messages: list[FrameworkMessage] = []
+
+    async def run_turn(
+        self,
+        user_input: str,
+        session: AgentSession | None = None,
+    ) -> None:
+        """Run a single agent turn with the given user input.
+
+        Echoes the input, then delegates to the agent loop.
+
+        Args:
+            user_input: The user's input text.
+            session: Optional agent session for conversation history.
+        """
+        async with self._input_gate:
+            self._ux.write_user_input_echo(user_input)
+
+            from agent_framework import Message
+
+            messages = [Message(role="user", contents=[user_input])]
+            await self._run_agent_loop(messages, session)
+
+    async def start_agent_turn(
+        self,
+        messages: list,
+        session: AgentSession | None = None,
+    ) -> None:
+        """Resume the agent loop with pre-built messages (from follow-up responses).
+
+        Called by the app after the user finishes answering follow-up questions.
+        If messages is empty, just completes the turn.
+
+        Args:
+            messages: List of Message objects to send to the agent.
+            session: Optional agent session.
+        """
+        async with self._input_gate:
+            if not messages:
+                self._complete_turn()
+                return
+            await self._run_agent_loop(messages, session)
+
+    def on_streaming_input(
+        self,
+        text: str,
+        session: AgentSession | None = None,
+    ) -> None:
+        """Handle user input submitted while an agent turn is streaming.
+
+        The text is enqueued via the ``MessageInjectionMiddleware`` so the agent
+        can pick it up on its next opportunity within the ongoing run. No-op if
+        the agent has no injection middleware or there is no active session.
+
+        Args:
+            text: The user's input text.
+            session: The active agent session.
+        """
+        if self._message_injector is None or session is None:
+            return
+
+        text = text.strip()
+        if not text:
+            return
+
+        from agent_framework import Message
+
+        self._message_injector.enqueue_messages(session, Message(role="user", contents=[text]))
+        pending = self._message_injector.get_pending_messages(session)
+        self._ux.set_queued_messages([m.text for m in pending])
+
+    def _sync_queued_message_display(self, session: AgentSession | None) -> None:
+        """Sync the queued-items display with the injector's pending messages.
+
+        Messages that have been consumed (drained by the middleware) since the
+        last sync are echoed to the output area as regular user-input entries.
+        No-op if there is no injection middleware or active session.
+
+        Args:
+            session: The active agent session.
+        """
+        if self._message_injector is None or session is None:
+            return
+
+        pending = self._message_injector.get_pending_messages(session)
+
+        # The injection middleware drains the whole queue at once, so a message
+        # is consumed when it is no longer present in the pending list. Compare
+        # by object identity (snapshots share the same Message objects until the
+        # queue is cleared) so consumed messages are echoed correctly even if a
+        # drain is followed by a new enqueue before the next sync.
+        current_ids = {id(m) for m in pending}
+        for msg in self._last_pending_messages:
+            if id(msg) not in current_ids:
+                self._ux.write_user_input_echo(msg.text or "")
+
+        self._last_pending_messages = pending
+        self._ux.set_queued_messages([m.text for m in pending])
+
+    async def _run_agent_loop(
+        self,
+        messages: list,
+        session: AgentSession | None,
+    ) -> None:
+        """Run the agent loop, re-invoking as needed for follow-up messages.
+
+        Loops while there are messages to send. After each stream:
+        - Collects follow-up actions from observers
+        - If questions exist → queue them and return (UI will collect answers)
+        - If only direct messages → loop with those messages
+        - If nothing → complete the turn
+
+        Args:
+            messages: Initial messages to send.
+            session: Optional agent session.
+        """
+        next_messages = messages
+
+        # Seed the pending-message snapshot so consumed injected messages can be
+        # detected and echoed during streaming.
+        self._last_pending_messages = (
+            self._message_injector.get_pending_messages(session)
+            if self._message_injector is not None and session is not None
+            else []
+        )
+
+        while next_messages:
+            # Configure run options
+            options = self._configure_run_options(session)
+
+            # Begin streaming
+            self._ux.begin_streaming()
+            self._ux.begin_streaming_output()
+            self._ux.set_show_spinner(True)
+
+            try:
+                await self._stream_response_messages(next_messages, session, options)
+            except Exception as ex:
+                self._ux.append_info_line(
+                    f"❌ Stream error: {ex.__class__.__name__}:\n{ex}",
+                    color="red",
+                )
+
+            # Final sync after streaming (echo any messages consumed on the last update).
+            self._sync_queued_message_display(session)
+
+            # Stop spinner and end streaming output
+            self._ux.set_show_spinner(False)
+
+            # Collect follow-up actions from observers
+            follow_up_actions = await self._collect_follow_up_actions(session)
+
+            # Separate direct messages from questions
+            has_follow_ups = len(follow_up_actions) > 0
+
+            # Write no-text warning if applicable
+            await self._ux.write_no_text_warning(has_follow_ups)
+
+            # Enqueue all follow-up actions
+            for action in follow_up_actions:
+                self._ux.enqueue_follow_up_action(action)
+
+            # Check if there are pending questions (UI needs user input)
+            if self._ux.has_pending_questions():
+                # Pause — the UI will collect answers and call start_agent_turn
+                return
+
+            # No questions — drain any accumulated direct messages and loop
+            drained = self._ux.take_follow_up_responses()
+            next_messages = drained if drained else None
+
+        self._complete_turn()
+
+    def _complete_turn(self) -> None:
+        """Complete the current turn (end streaming)."""
+        self._ux.end_streaming()
+
+    def _configure_run_options(
+        self,
+        session: AgentSession | None,
+    ) -> dict:
+        """Configure run options via observers.
+
+        Each observer can modify the options dict to influence agent behavior.
+
+        Args:
+            session: Optional agent session.
+
+        Returns:
+            Options dict for agent.run().
+        """
+        options = {}
+        for observer in self._observers:
+            observer.configure_run_options(options, self._agent, session)
+        return options
+
+    async def _stream_response(
+        self,
+        user_input: str,
+        session: AgentSession | None,
+        options: dict,
+    ) -> None:
+        """Stream agent response from a text input and dispatch to observers.
+
+        Args:
+            user_input: The user's input text.
+            session: Optional agent session.
+            options: Run options configured by observers.
+        """
+        # Stream response using agent.run(stream=True)
+        stream = self._agent.run(
+            user_input,
+            stream=True,
+            session=session,
+            options=options,
+        )
+
+        # Process each update chunk
+        async for update in stream:
+            await self._dispatch_update(update, session)
+
+        # Extract usage from the final response
+        self._extract_usage(stream)
+
+    async def _stream_response_messages(
+        self,
+        messages: list,
+        session: AgentSession | None,
+        options: dict,
+    ) -> None:
+        """Stream agent response from Message objects and dispatch to observers.
+
+        Args:
+            messages: List of Message objects to send.
+            session: Optional agent session.
+            options: Run options configured by observers.
+        """
+        stream = self._agent.run(
+            messages,
+            stream=True,
+            session=session,
+            options=options,
+        )
+
+        async for update in stream:
+            await self._dispatch_update(update, session)
+
+        self._extract_usage(stream)
+
+    def _extract_usage(self, stream) -> None:
+        """Extract token usage from a completed stream."""
+        try:
+            get_final = getattr(stream, "get_final_response", None)
+            if not get_final:
+                return
+
+            import inspect
+
+            if inspect.iscoroutinefunction(get_final):
+                return
+
+            final_response = get_final()
+            if final_response is None:
+                return
+
+            usage = getattr(final_response, "usage_details", None)
+            if not isinstance(usage, dict):
+                return
+
+            input_tokens = usage.get("input_token_count", 0) or 0
+            output_tokens = usage.get("output_token_count", 0) or 0
+            if input_tokens or output_tokens:
+                self._ux.set_usage_text(self._format_usage(input_tokens, output_tokens))
+        except (AttributeError, TypeError):
+            pass
+
+    async def _dispatch_update(
+        self,
+        update,  # AgentResponseUpdate
+        session: AgentSession | None,
+    ) -> None:
+        """Dispatch a single update to all observers.
+
+        Calls observer lifecycle methods in order:
+        1. on_response_update (once per update)
+        2. on_content (for each content item)
+        3. on_text (if text is present)
+
+        Args:
+            update: The agent response update.
+            session: Optional agent session.
+        """
+        # on_response_update
+        for observer in self._observers:
+            await observer.on_response_update(self._ux, update, self._agent, session)
+
+        # on_content for each content item
+        if hasattr(update, "contents") and update.contents:
+            for content in update.contents:
+                for observer in self._observers:
+                    await observer.on_content(self._ux, content, self._agent, session)
+
+        # on_text for text chunks
+        if hasattr(update, "text") and update.text:
+            for observer in self._observers:
+                await observer.on_text(self._ux, update.text, self._agent, session)
+
+        # Echo any injected messages consumed by the agent on this update.
+        self._sync_queued_message_display(session)
+
+    async def _collect_follow_up_actions(
+        self,
+        session: AgentSession | None,
+    ) -> list[FollowUpAction]:
+        """Collect follow-up actions from all observers.
+
+        Called after streaming completes to gather any follow-up questions
+        or messages from observers.
+
+        Args:
+            session: Optional agent session.
+
+        Returns:
+            List of follow-up actions from all observers.
+        """
+        actions: list[FollowUpAction] = []
+        for observer in self._observers:
+            observer_actions = await observer.on_stream_complete(self._ux, self._agent, session)
+            if observer_actions:
+                actions.extend(observer_actions)
+        return actions
+
+    def _format_usage(self, input_tokens: int, output_tokens: int) -> str:
+        """Format token counts matching C# harness style: 📊 Tokens — input: X | output: Y | total: Z."""
+        total_tokens = input_tokens + output_tokens
+
+        input_budget = None
+        if self._max_context_window_tokens and self._max_output_tokens:
+            input_budget = self._max_context_window_tokens - self._max_output_tokens
+
+        return (
+            f"📊 Tokens — input: {self._format_token_count(input_tokens, input_budget)}"
+            f" | output: {self._format_token_count(output_tokens, self._max_output_tokens)}"
+            f" | total: {self._format_token_count(total_tokens, self._max_context_window_tokens)}"
+        )
+
+    @staticmethod
+    def _format_token_count(count: int, budget: int | None) -> str:
+        """Format a token count, optionally showing budget percentage."""
+        if budget and budget > 0:
+            pct = count / budget * 100
+            return f"{count:,}/{budget:,} ({pct:.1f}%)"
+        return f"{count:,}"