Improve WebUI streaming and Python tooling

src/agent_runtime.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass, field, replace
 from datetime import datetime, timezone
 import json
 from pathlib import Path
-from typing import Any
+from typing import Any, Callable
 from uuid import uuid4
 
 from .account_runtime import AccountRuntime
@@ -68,6 +68,43 @@ from .tokenizer_runtime import describe_token_counter
 from .workflow_runtime import WorkflowRuntime
 from .worktree_runtime import WorktreeRuntime
 from .session_env_vars import clear_session_env_vars
+
+RuntimeEventSink = Callable[[dict[str, object]], None]
+
+
+class _RuntimeEventBuffer(list[dict[str, object]]):
+    def __init__(self, event_sink: RuntimeEventSink | None = None) -> None:
+        super().__init__()
+        self._event_sink = event_sink
+
+    def append(self, event: dict[str, object]) -> None:  # type: ignore[override]
+        super().append(event)
+        if self._event_sink is not None:
+            self._event_sink(dict(event))
+
+    def extend(self, events: Any) -> None:  # type: ignore[override]
+        for event in events:
+            self.append(event)
+
+
+def _append_final_text_stream_events(
+    stream_events: list[dict[str, object]],
+    text: str,
+) -> None:
+    if not text:
+        return
+    stream_events.append({'type': 'final_text_start'})
+    chunk_size = 96
+    for start in range(0, len(text), chunk_size):
+        stream_events.append(
+            {
+                'type': 'final_text_delta',
+                'delta': text[start : start + chunk_size],
+            }
+        )
+    stream_events.append({'type': 'final_text_end'})
+
+
 from .session_store import (
     StoredAgentSession,
     load_agent_session,
@@ -375,6 +412,7 @@ class LocalCodingAgent:
         session_id: str | None = None,
         *,
         runtime_context: str | None = None,
+        event_sink: RuntimeEventSink | None = None,
     ) -> AgentRunResult:
         self.managed_agent_id = None
         self.resume_source_session_id = None
@@ -389,6 +427,7 @@ class LocalCodingAgent:
             scratchpad_directory=scratchpad_directory,
             existing_file_history=(),
             runtime_context=runtime_context,
+            event_sink=event_sink,
         )
         self._accumulate_usage(result)
         self._finalize_managed_agent(result)
@@ -400,6 +439,7 @@ class LocalCodingAgent:
         stored_session: StoredAgentSession,
         *,
         runtime_context: str | None = None,
+        event_sink: RuntimeEventSink | None = None,
     ) -> AgentRunResult:
         self.managed_agent_id = None
         self.resume_source_session_id = stored_session.session_id
@@ -435,6 +475,7 @@ class LocalCodingAgent:
             scratchpad_directory=scratchpad_directory,
             existing_file_history=stored_session.file_history,
             runtime_context=runtime_context,
+            event_sink=event_sink,
         )
         self._accumulate_usage(result)
         self._finalize_managed_agent(result)
@@ -449,6 +490,7 @@ class LocalCodingAgent:
         scratchpad_directory: Path | None,
         existing_file_history: tuple[dict[str, object], ...],
         runtime_context: str | None = None,
+        event_sink: RuntimeEventSink | None = None,
     ) -> AgentRunResult:
         slash_result = preprocess_slash_command(self, prompt)
         if slash_result.handled and not slash_result.should_query:
@@ -495,6 +537,10 @@ class LocalCodingAgent:
         session.append_user(effective_prompt, model_content=model_prompt)
         self.last_session = session
         self.active_session_id = session_id
+        self.tool_context = replace(
+            self.tool_context,
+            scratchpad_directory=scratchpad_directory,
+        )
         tool_specs = [tool.to_openai_tool() for tool in self.tool_registry.values()]
         starting_usage = UsageStats()
         starting_cost_usd = 0.0
@@ -525,7 +571,7 @@ class LocalCodingAgent:
         total_usage = starting_usage
         total_cost_usd = starting_cost_usd
         file_history = list(existing_file_history)
-        stream_events: list[dict[str, object]] = []
+        stream_events: list[dict[str, object]] = _RuntimeEventBuffer(event_sink)
         assistant_response_segments: list[str] = []
         delegated_tasks = sum(
             1 for entry in file_history if entry.get('action') in ('delegate_agent', 'Agent')
@@ -715,8 +761,10 @@ class LocalCodingAgent:
                                 )
                                 last_content = ''.join(assistant_response_segments)
                                 continue
+                            final_output = ''.join(assistant_response_segments)
+                            _append_final_text_stream_events(stream_events, final_output)
                             result = AgentRunResult(
-                                final_output=''.join(assistant_response_segments),
+                                final_output=final_output,
                                 turns=turn_index,
                                 tool_calls=tool_calls,
                                 transcript=session.transcript(),
@@ -816,8 +864,10 @@ class LocalCodingAgent:
                     )
                     last_content = ''.join(assistant_response_segments)
                     continue
+                final_output = ''.join(assistant_response_segments)
+                _append_final_text_stream_events(stream_events, final_output)
                 result = AgentRunResult(
-                    final_output=''.join(assistant_response_segments),
+                    final_output=final_output,
                     turns=turn_index,
                     tool_calls=tool_calls,
                     transcript=session.transcript(),
@@ -896,6 +946,8 @@ class LocalCodingAgent:
                         'type': 'tool_start',
                         'tool_name': tool_call.name,
                         'tool_call_id': tool_call.id,
+                        'arguments': dict(tool_call.arguments),
+                        'assistant_content': turn.content,
                         'message_id': session.messages[tool_message_index].message_id,
                     }
                 )
@@ -1194,11 +1246,13 @@ class LocalCodingAgent:
                     self.last_run_result = result
                     return result
 
+        final_output = (
+            last_content
+            or 'Stopped: max turns reached before the model produced a final answer.'
+        )
+        _append_final_text_stream_events(stream_events, final_output)
         result = AgentRunResult(
-            final_output=(
-                last_content
-                or 'Stopped: max turns reached before the model produced a final answer.'
-            ),
+            final_output=final_output,
             turns=self.runtime_config.max_turns,
             tool_calls=tool_calls,
             transcript=session.transcript(),