ctrlcode 0.1.0__py3-none-any.whl

ctrlcode/session/manager.py

@@ -0,0 +1,967 @@
+"""Session management using harness-utils."""
+
+import asyncio
+import logging
+import time
+from typing import Any, AsyncIterator, Optional
+from dataclasses import dataclass
+from pathlib import Path
+import uuid
+
+from harnessutils import ConversationManager, Message, TextPart
+from harnessutils.storage import FilesystemStorage
+from harnessutils.config import HarnessConfig, StorageConfig
+
+from ..providers.base import Provider, StreamEvent
+from .baseline import BaselineManager, Baseline
+from ..fuzzing.derived_orchestrator import DerivedFuzzingOrchestrator
+from ..tools.executor import ToolExecutor
+from ..tools.registry import ToolRegistry
+from ..skills.registry import SkillRegistry
+
+logger = logging.getLogger(__name__)
+
+
+def strip_markdown_fences(content: str) -> str:
+    """Remove markdown code fences from content if present."""
+    lines = content.splitlines()
+
+    # Check if first line is a code fence
+    if lines and lines[0].strip().startswith("```"):
+        # Remove first line
+        lines = lines[1:]
+
+    # Check if last line is a closing fence
+    if lines and lines[-1].strip() == "```":
+        # Remove last line
+        lines = lines[:-1]
+
+    # Remove any trailing instruction lines (e.g., "Run with: ...")
+    while lines and lines[-1].strip().startswith("Run with:"):
+        lines = lines[:-1]
+
+    return "\n".join(lines)
+
+
+@dataclass
+class Session:
+    """Represents a conversation session."""
+
+    id: str
+    conv_id: str
+    provider: Provider
+    cumulative_tokens: int = 0  # Track total tokens used across all turns
+
+
+class SessionManager:
+    """Manages conversation sessions with baseline tracking."""
+
+    def __init__(
+        self,
+        provider: Provider,
+        storage_path: str,
+        config: Optional[HarnessConfig] = None,
+        fuzzing_orchestrator: Optional[DerivedFuzzingOrchestrator] = None,
+        fuzzing_enabled: bool = False,
+        tool_executor: Optional[ToolExecutor] = None,
+        tool_registry: Optional[ToolRegistry] = None,
+        skill_registry: Optional[SkillRegistry] = None,
+        context_limit: int = 200000,
+        workspace_root: Optional[Path] = None,
+    ):
+        """
+        Initialize session manager.
+
+        Args:
+            provider: LLM provider to use
+            storage_path: Path for conversation storage
+            config: Optional harness config
+            fuzzing_orchestrator: Optional fuzzing orchestrator
+            fuzzing_enabled: Whether to use fuzzing
+            tool_executor: Optional tool executor for MCP tools
+            tool_registry: Optional tool registry for tool definitions
+            skill_registry: Optional skill registry
+            context_limit: Context window limit
+            workspace_root: Workspace root for AGENT.md lookup
+        """
+        storage_config = StorageConfig(base_path=Path(storage_path))
+        self.conv_manager = ConversationManager(
+            storage=FilesystemStorage(storage_config),
+            config=config or HarnessConfig()
+        )
+        self.provider = provider
+        self.sessions: dict[str, Session] = {}
+        self.baseline_manager = BaselineManager()
+        self.fuzzing_orchestrator = fuzzing_orchestrator
+        self.fuzzing_enabled = fuzzing_enabled
+        self.tool_executor = tool_executor
+        self.tool_registry = tool_registry
+        self.skill_registry = skill_registry
+        self.context_limit = context_limit
+        self.workspace_root = workspace_root
+        self._base_prompt = self._load_base_prompt()
+        self._agent_instructions = self._load_agent_instructions()
+
+    def _load_base_prompt(self) -> str:
+        """
+        Load base system prompt from prompts/SYSTEM_PROMPT.md.
+
+        Search order:
+        1. {workspace_root}/prompts/SYSTEM_PROMPT.md  (project override)
+        2. ~/.config/ctrlcode/SYSTEM_PROMPT.md         (user global override)
+        3. Bundled default
+
+        Returns:
+            Base system prompt content.
+        """
+        candidates = []
+
+        # 1. Project-local override
+        if self.workspace_root:
+            candidates.append(self.workspace_root / "prompts" / "SYSTEM_PROMPT.md")
+
+        # 2. User global override
+        try:
+            from platformdirs import user_config_dir
+            candidates.append(Path(user_config_dir("ctrlcode")) / "SYSTEM_PROMPT.md")
+        except Exception:
+            pass
+
+        for prompt_file in candidates:
+            if prompt_file.exists():
+                try:
+                    content = prompt_file.read_text(encoding="utf-8")
+                    logger.info(f"Loaded system prompt from {prompt_file} ({len(content)} chars)")
+                    return content
+                except Exception as e:
+                    logger.warning(f"Failed to load {prompt_file}: {e}")
+
+        # 3. Bundled default
+        logger.debug("Using bundled default system prompt")
+        return """You are Ctrl+Code, Canoozie's personal coding assistant.
+
+## CRITICAL: Act immediately, never introduce yourself
+
+**NEVER** say "I'm ready to help", "I have access to tools", or introduce your capabilities.
+**NEVER** greet the user or ask what they want when they've already told you.
+When given a task, **immediately call the appropriate tool**. Your first output must be a tool call, not text.
+
+Examples of correct behaviour:
+- User says "show me the last git commit" → call `run_command` with `git log -1` immediately
+- User says "find the login function" → call `search_code` with "login" immediately
+- User says "read app.py" → call `read_file` with "app.py" immediately
+
+## Tools available
+
+- `run_command` — run shell commands (git, tests, builds, etc.)
+- `read_file` — read a file's contents
+- `write_file` — create a new file
+- `update_file` — edit an existing file
+- `search_files` — find files by glob pattern
+- `search_code` — search for code by content
+- `list_directory` — list directory contents
+- `web_fetch` — fetch a URL
+
+## Tool usage rules
+
+- Call ALL tools needed in a SINGLE response — do not wait for results before calling the next tool
+- Use `run_command` for git operations, tests, builds, and any shell commands
+- Use `read_file` / `search_files` / `search_code` for exploring the codebase
+- Use `update_file` to edit existing files, `write_file` only for new files
+- When referencing code, include `file_path:line_number` so the user can navigate to it
+
+## Workspace and file paths
+
+- Use relative paths (`src/main.py`) not absolute paths (`/home/user/src/main.py`)
+- If unsure of a file's location, call `search_files` first
+
+## Tone
+
+- Be concise and direct. No emojis unless asked. Output renders in a monospace terminal."""
+
+    def _load_agent_instructions(self) -> str:
+        """
+        Load AGENT.md instructions hierarchically.
+
+        Order (most general to most specific):
+        1. Global: ~/.config/ctrlcode/AGENT.md
+        2. Project: {workspace_root}/AGENT.md
+
+        Returns combined instructions with clear section markers.
+        """
+        instructions = []
+
+        # 1. Global config AGENT.md
+        try:
+            from platformdirs import user_config_dir
+            config_dir = Path(user_config_dir("ctrlcode"))
+            global_agent = config_dir / "AGENT.md"
+            if global_agent.exists():
+                content = global_agent.read_text(encoding="utf-8")
+                instructions.append(f"# Global Agent Instructions\n\n{content}")
+                logger.info(f"Loaded global AGENT.md ({len(content)} chars)")
+        except Exception as e:
+            logger.debug(f"No global AGENT.md: {e}")
+
+        # 2. Project AGENT.md
+        if self.workspace_root:
+            project_agent = self.workspace_root / "AGENT.md"
+            if project_agent.exists():
+                try:
+                    content = project_agent.read_text(encoding="utf-8")
+                    instructions.append(f"# Project-Specific Instructions\n\n{content}")
+                    logger.info(f"Loaded project AGENT.md ({len(content)} chars)")
+                except Exception as e:
+                    logger.warning(f"Failed to load project AGENT.md: {e}")
+
+        return "\n\n---\n\n".join(instructions) if instructions else ""
+
+    def create_session(self, provider: Optional[Provider] = None) -> Session:
+        """
+        Create a new session.
+
+        Args:
+            provider: Optional provider override
+
+        Returns:
+            New session instance
+        """
+        conv = self.conv_manager.create_conversation(project_id="ctrl-code")
+        session = Session(
+            id=str(uuid.uuid4()),
+            conv_id=conv.id,
+            provider=provider or self.provider
+        )
+        self.sessions[session.id] = session
+        return session
+
+    def get_session(self, session_id: str) -> Optional[Session]:
+        """Get session by ID."""
+        return self.sessions.get(session_id)
+
+    async def process_turn(
+        self,
+        session_id: str,
+        user_input: str,
+        tools: list[dict] | None = None
+    ) -> AsyncIterator[StreamEvent]:
+        """
+        Process a conversation turn.
+
+        Args:
+            session_id: Session identifier
+            user_input: User's input message
+            tools: Optional tool definitions
+
+        Yields:
+            StreamEvent: Streaming events from provider
+        """
+        session = self.sessions.get(session_id)
+        if not session:
+            raise ValueError(f"Session not found: {session_id}")
+
+        # Expand skills if present
+        if self.skill_registry:
+            expanded_input, was_skill = self.skill_registry.process_input(user_input)
+            if was_skill:
+                # Emit skill expansion event
+                yield StreamEvent(
+                    type="skill_expanded",
+                    data={"original": user_input, "expanded": expanded_input}
+                )
+                user_input = expanded_input
+
+        # Extract baseline if present
+        baseline_code = self.baseline_manager.extract_from_request(user_input)
+        if baseline_code:
+            baseline = Baseline(code=baseline_code)
+            self.baseline_manager.store(session_id, baseline)
+
+        # Prune before adding new user message (makes room for the new turn)
+        self.conv_manager.prune_before_turn(session.conv_id)
+
+        # Add user message
+        user_msg = Message(id=self._generate_msg_id(), role="user")
+        user_msg.add_part(TextPart(text=user_input))
+        self.conv_manager.add_message(session.conv_id, user_msg)
+
+        # Track context size before LLM call for accurate token delta
+        session.context_before_turn = self.conv_manager.calculate_context_usage(
+            session.conv_id,
+            model=session.provider.model
+        )
+
+        # Get messages for LLM
+        messages = self.conv_manager.to_model_format(session.conv_id)
+        logger.info(f"to_model_format: {len(messages)} messages, roles={[m.get('role') for m in messages]}")
+
+        # Get tool definitions from registry if not provided
+        if tools is None and self.tool_registry:
+            tools = self.tool_registry.get_tool_definitions()
+            logger.info(f"Fetched {len(tools)} tool definitions from registry")
+        else:
+            logger.info(f"Tools parameter: {tools}")
+
+        # Track if we should fuzz (detected code block)
+        should_fuzz = False
+        accumulated_text = []
+
+        # Track streaming rate for fuzzing output
+        streaming_start_time = time.time()
+        streaming_end_time = None
+
+        # Track usage tokens (will be replaced by fuzzed output tokens if fuzzing occurs)
+        usage_tokens = 0
+
+        # Build system prompt once; reuse for all continuation calls
+        system_prompt_msg: dict | None = None
+        if not messages or messages[0].get("role") != "system":
+            # Build AGENT.md section
+            agent_section = ""
+            if self._agent_instructions:
+                agent_section = f"\n\n{self._agent_instructions}"
+
+            # Inject live environment info
+            import platform
+            from datetime import date
+            work_dir = self.workspace_root or Path.cwd()
+            is_git = (work_dir / ".git").exists()
+            env_section = f"""
+**Environment:**
+```
+Working directory: {work_dir}
+Is directory a git repo: {"Yes" if is_git else "No"}
+Platform: {platform.system().lower()}
+OS Version: {platform.platform()}
+Today's date: {date.today().isoformat()}
+```"""
+
+            system_prompt_msg = {
+                "role": "system",
+                "content": f"{self._base_prompt}{agent_section}{env_section}"
+            }
+            messages = [system_prompt_msg] + messages
+
+        # Stream response with tool execution loop
+        assistant_text = []
+        tool_calls: list[dict] = []
+        current_tool_call: dict | None = None
+
+        # First pass: stream and collect tool calls
+        async for event in session.provider.stream(messages, tools=tools):  # type: ignore[attr-defined]
+            # Accumulate text and detect code blocks
+            if event.type == "text":
+                text_chunk = event.data["text"]
+                assistant_text.append(text_chunk)
+                accumulated_text.append(text_chunk)
+
+                # Check if we've accumulated a code block
+                full_text = "".join(accumulated_text)
+                if "```" in full_text and not should_fuzz:
+                    should_fuzz = True
+                    logger.info("Code block detected, will run fuzzing after completion")
+                    # Emit fuzzing start status instead of showing original text
+                    yield StreamEvent(
+                        type="fuzzing_progress",
+                        data={"stage": "detected", "message": "Code detected, preparing to optimize..."}
+                    )
+
+                # Only yield text if we're NOT fuzzing
+                if not should_fuzz:
+                    yield event
+                else:
+                    logger.debug("Suppressing text event (fuzzing enabled)")
+
+            elif event.type == "usage":
+                # Capture end time when response completes
+                streaming_end_time = time.time()
+
+                # Store completion tokens from API (tokens generated by assistant)
+                usage = event.data.get("usage", {})
+                usage_tokens += usage.get("completion_tokens", 0)
+
+                # Don't yield yet - we'll yield cumulative at the end
+                logger.debug(f"Captured {usage.get('completion_tokens', 0)} tokens, cumulative: {usage_tokens}")
+
+            elif event.type == "tool_call_start":
+                current_tool_call = {
+                    "tool": event.data["tool"],
+                    "call_id": event.data["call_id"],
+                    "input": ""
+                }
+                yield event
+
+            elif event.type == "tool_call_delta":
+                if current_tool_call:
+                    current_tool_call["input"] += event.data.get("delta", "")
+                yield event
+
+            elif event.type == "content_block_stop":
+                if current_tool_call:
+                    import json
+                    try:
+                        current_tool_call["input"] = json.loads(current_tool_call["input"])
+                    except json.JSONDecodeError:
+                        current_tool_call["input"] = {}
+                    tool_calls.append(current_tool_call)
+                    current_tool_call = None
+                yield event
+
+            else:
+                # All other events
+                yield event
+
+        # Check if we should fuzz before executing tools
+        # If write_file or update_file is called, fuzz the content
+        write_file_calls = [tc for tc in tool_calls if tc["tool"] == "write_file"]
+        update_file_calls = [tc for tc in tool_calls if tc["tool"] == "update_file"]
+
+        if write_file_calls:
+            # Fuzz any write_file content
+            for write_call in write_file_calls:
+                content = write_call["input"].get("content", "")
+                # Strip markdown fences if present
+                content = strip_markdown_fences(content)
+                write_call["input"]["content"] = content
+                path = write_call["input"].get("path", "")
+
+                should_fuzz = True
+                # Add to accumulated_text so fuzzing has the content
+                accumulated_text.append(content)
+                logger.info(f"write_file detected for {path}, will fuzz before writing")
+                break
+
+        if update_file_calls:
+            # Fuzz any update_file content
+            for update_call in update_file_calls:
+                content = update_call["input"].get("content", "")
+                if content:  # Only fuzz if there's content to update
+                    # Strip markdown fences if present
+                    content = strip_markdown_fences(content)
+                    update_call["input"]["content"] = content
+                    path = update_call["input"].get("path", "")
+
+                    should_fuzz = True
+                    # Add to accumulated_text so fuzzing has the content
+                    accumulated_text.append(content)
+                    logger.info(f"update_file detected for {path}, will fuzz before updating")
+                    break
+
+        # Execute tools if any were called (but not write_file if fuzzing)
+        if tool_calls and self.tool_executor:
+            # Add assistant message with tool use
+            assistant_msg = Message(id=self._generate_msg_id(), role="assistant")
+            if assistant_text:
+                assistant_msg.add_part(TextPart(text="".join(assistant_text)))
+            self.conv_manager.add_message(session.conv_id, assistant_msg)
+
+            # Execute non-write/update tools first (or all tools if not fuzzing)
+            auto_chain_calls = []  # Track tools to auto-chain
+
+            for tool_call in tool_calls:
+                # Skip write_file or update_file if we're going to fuzz
+                if tool_call["tool"] in ("write_file", "update_file") and should_fuzz:
+                    continue
+
+                result = await self.tool_executor.execute(
+                    tool_name=tool_call["tool"],
+                    arguments=tool_call["input"],
+                    call_id=tool_call["call_id"]
+                )
+
+                # Emit tool result event
+                yield StreamEvent(
+                    type="tool_result",
+                    data={
+                        "tool": tool_call["tool"],
+                        "success": result.success,
+                        "result": result.result if result.success else result.error
+                    }
+                )
+
+                # Add tool result message
+                tool_result_msg = Message(id=self._generate_msg_id(), role="user")
+                result_text = f"[Tool: {tool_call['tool']}]\n"
+                if result.success:
+                    result_text += f"Result: {result.result}"
+                else:
+                    result_text += f"Error: {result.error}"
+                tool_result_msg.add_part(TextPart(text=result_text))
+                self.conv_manager.add_message(session.conv_id, tool_result_msg)
+
+                # Auto-chaining logic: check if we should automatically call next tool
+                if result.success:
+                    chained = self._check_auto_chain(tool_call["tool"], result.result)
+                    if chained:
+                        auto_chain_calls.append(chained)
+                        logger.info(f"Auto-chaining: {tool_call['tool']} → {chained['tool']}")
+
+            # Execute auto-chained tool calls
+            for chained_call in auto_chain_calls:
+                result = await self.tool_executor.execute(
+                    tool_name=chained_call["tool"],
+                    arguments=chained_call["arguments"],
+                    call_id=chained_call["call_id"]
+                )
+
+                # Emit tool result event
+                yield StreamEvent(
+                    type="tool_result",
+                    data={
+                        "tool": chained_call["tool"],
+                        "success": result.success,
+                        "result": result.result if result.success else result.error
+                    }
+                )
+
+                # Add tool result message
+                tool_result_msg = Message(id=self._generate_msg_id(), role="user")
+                result_text = f"[Tool: {chained_call['tool']}]\n"
+                if result.success:
+                    result_text += f"Result: {result.result}"
+                else:
+                    result_text += f"Error: {result.error}"
+                tool_result_msg.add_part(TextPart(text=result_text))
+                self.conv_manager.add_message(session.conv_id, tool_result_msg)
+
+            # FSM: Continue looping through tool use until no more tools called
+            # This implements a proper ReAct loop: observe -> reason -> act -> repeat
+            non_write_tools = [tc for tc in tool_calls if tc["tool"] not in ("write_file", "update_file")]
+            logger.info(f"ReAct loop check: non_write_tools={len(non_write_tools)}, should_fuzz={should_fuzz}")
+
+            # Loop with high safety limit - let LLM decide when it's done
+            continuation_count = 0
+            max_continuations = 50  # Safety limit, not expected to hit in normal operation
+
+            while non_write_tools and not should_fuzz and continuation_count < max_continuations:
+                continuation_count += 1
+                logger.info(f"ReAct loop iteration {continuation_count}/{max_continuations}")
+
+                # Add continuation prompt
+                reminder_msg = Message(id=self._generate_msg_id(), role="user")
+                reminder_msg.add_part(TextPart(text="If you have all the information needed, present your findings to the user. If you need more data, call additional tools."))
+                self.conv_manager.add_message(session.conv_id, reminder_msg)
+
+                # Get updated messages for continuation (always include system prompt)
+                messages = self.conv_manager.to_model_format(session.conv_id)
+                if system_prompt_msg:
+                    messages = [system_prompt_msg] + messages
+
+                # Stream continuation and track tool calls
+                assistant_text = []
+                continuation_tool_calls = []
+                current_continuation_tool = None
+                event_count = 0
+
+                async for event in session.provider.stream(messages, tools=tools):  # type: ignore[attr-defined]
+                    event_count += 1
+                    yield event
+
+                    if event.type == "text":
+                        assistant_text.append(event.data["text"])
+                    elif event.type == "usage":
+                        # Accumulate tokens from continuation
+                        usage = event.data.get("usage", {})
+                        usage_tokens += usage.get("completion_tokens", 0)
+                        logger.debug(f"Continuation added {usage.get('completion_tokens', 0)} tokens, cumulative: {usage_tokens}")
+                    elif event.type == "tool_call_start":
+                        current_continuation_tool = {
+                            "tool": event.data["tool"],
+                            "call_id": event.data["call_id"],
+                            "input": ""
+                        }
+                    elif event.type == "tool_call_delta":
+                        if current_continuation_tool:
+                            current_continuation_tool["input"] += event.data.get("delta", "")
+                    elif event.type == "content_block_stop":
+                        if current_continuation_tool:
+                            import json
+                            try:
+                                current_continuation_tool["input"] = json.loads(current_continuation_tool["input"])
+                            except json.JSONDecodeError:
+                                current_continuation_tool["input"] = {}
+                            continuation_tool_calls.append(current_continuation_tool)
+                            current_continuation_tool = None
+
+                logger.info(f"Continuation finished with {event_count} events, {len(continuation_tool_calls)} tool calls")
+
+                # Add assistant message from continuation (whether tools were called or not)
+                cont_assistant_msg = Message(id=self._generate_msg_id(), role="assistant")
+                if assistant_text:
+                    cont_assistant_msg.add_part(TextPart(text="".join(assistant_text)))
+                self.conv_manager.add_message(session.conv_id, cont_assistant_msg)
+
+                # Execute any tool calls from continuation
+                if continuation_tool_calls and self.tool_executor:
+                    logger.info(f"Executing {len(continuation_tool_calls)} tools from continuation")
+
+                    # Check if continuation has write_file/update_file that should be fuzzed
+                    cont_write_calls = [tc for tc in continuation_tool_calls if tc["tool"] == "write_file"]
+                    cont_update_calls = [tc for tc in continuation_tool_calls if tc["tool"] == "update_file"]
+
+                    if cont_write_calls:
+                        for write_call in cont_write_calls:
+                            content = write_call["input"].get("content", "")
+                            content = strip_markdown_fences(content)
+                            write_call["input"]["content"] = content
+                            path = write_call["input"].get("path", "")
+                            should_fuzz = True
+                            accumulated_text.append(content)
+                            logger.info(f"write_file in continuation detected for {path}, will fuzz")
+                            # Add to main list so it gets executed after fuzzing
+                            write_file_calls.extend(cont_write_calls)
+                            break
+
+                    if cont_update_calls:
+                        for update_call in cont_update_calls:
+                            content = update_call["input"].get("content", "")
+                            if content:
+                                content = strip_markdown_fences(content)
+                                update_call["input"]["content"] = content
+                                path = update_call["input"].get("path", "")
+                                should_fuzz = True
+                                accumulated_text.append(content)
+                                logger.info(f"update_file in continuation detected for {path}, will fuzz")
+                                # Add to main list so it gets executed after fuzzing
+                                update_file_calls.extend(cont_update_calls)
+                                break
+
+                    # Execute non-write/update tools (or all if not fuzzing)
+                    for tool_call in continuation_tool_calls:
+                        # Skip write_file/update_file if fuzzing
+                        if tool_call["tool"] in ("write_file", "update_file") and should_fuzz:
+                            logger.info(f"Skipping {tool_call['tool']} execution, will fuzz first")
+                            continue
+
+                        result = await self.tool_executor.execute(
+                            tool_name=tool_call["tool"],
+                            arguments=tool_call["input"],
+                            call_id=tool_call["call_id"]
+                        )
+
+                        # Emit tool result event (marked as continuation)
+                        yield StreamEvent(
+                            type="tool_result",
+                            data={
+                                "tool": tool_call["tool"],
+                                "success": result.success,
+                                "result": result.result if result.success else result.error,
+                                "continuation": True,  # Mark as continuation for TUI to handle differently
+                                "continuation_index": continuation_count
+                            }
+                        )
+
+                        # Add tool result message
+                        tool_result_msg = Message(id=self._generate_msg_id(), role="user")
+                        result_text = f"[Tool: {tool_call['tool']}]\n"
+                        if result.success:
+                            result_text += f"Result: {result.result}"
+                        else:
+                            result_text += f"Error: {result.error}"
+                        tool_result_msg.add_part(TextPart(text=result_text))
+                        self.conv_manager.add_message(session.conv_id, tool_result_msg)
+
+                # Check if we should continue the ReAct loop
+                non_write_tools = [tc for tc in continuation_tool_calls if tc["tool"] not in ("write_file", "update_file")]
+                if not non_write_tools:
+                    logger.info(f"ReAct loop ending: no more non-write tools after {continuation_count} continuations")
+                    break
+
+                # Update messages for next iteration (always include system prompt)
+                messages = self.conv_manager.to_model_format(session.conv_id)
+                if system_prompt_msg:
+                    messages = [system_prompt_msg] + messages
+                logger.info(f"ReAct loop continuing: iteration {continuation_count}/{max_continuations}")
+
+            # After loop exits, emit continuation summary for TUI
+            tool_count_estimate = len(tool_calls) + continuation_count  # Rough estimate
+            yield StreamEvent(
+                type="continuation_complete",
+                data={
+                    "iterations": continuation_count,
+                    "tool_count": tool_count_estimate
+                }
+            )
+
+            # Only do a final summary call if the loop hit the safety limit.
+            # If it ended naturally (model gave a text response), the summary is already done.
+            if non_write_tools:
+                logger.info(f"ReAct loop hit max iterations ({max_continuations}), requesting final summary")
+                messages = self.conv_manager.to_model_format(session.conv_id)
+                if system_prompt_msg:
+                    messages = [system_prompt_msg] + messages
+                async for event in session.provider.stream(messages, tools=None):  # type: ignore[attr-defined]
+                    yield event
+                    if event.type == "usage":
+                        usage = event.data.get("usage", {})
+                        usage_tokens += usage.get("completion_tokens", 0)
+            else:
+                logger.info("ReAct loop completed naturally, skipping redundant final summary")
+
+        # If fuzzing is enabled and we detected code block, run fuzzing now
+        if should_fuzz and self.fuzzing_enabled and self.fuzzing_orchestrator:
+            logger.info("Running fuzzing on code response")
+
+            # Get baseline if exists
+            baseline = self.baseline_manager.get(session_id)  # type: ignore[assignment]
+            baseline_code = baseline.code if baseline else None  # type: ignore[union-attr]
+
+            # Use accumulated text as context for fuzzing
+            full_response = "".join(accumulated_text)
+
+            # Emit fuzzing start
+            yield StreamEvent(
+                type="fuzzing_progress",
+                data={"stage": "starting", "message": "Improving code quality..."}
+            )
+
+            # Run fuzzing and forward progress events
+            result = None
+            async for event in self.fuzzing_orchestrator.fuzz(
+                user_request=user_input,
+                generated_code=baseline_code or full_response,
+                context_messages=messages,
+            ):
+                from ..fuzzing.derived_orchestrator import FuzzingResult
+                if isinstance(event, FuzzingResult):
+                    result = event
+                    yield StreamEvent(
+                        type="fuzzing_complete",
+                        data={
+                            "iterations": result.iterations,
+                            "quality_score": result.quality_score,
+                            "budget_used": result.budget_used,
+                            "divergences_found": result.divergences_found,
+                            "divergences_fixed": result.divergences_fixed,
+                            "oracle_corrections": result.oracle_corrections,
+                        }
+                    )
+                else:
+                    yield event
+
+            if result:
+                # Calculate actual streaming rate from original response
+                if streaming_end_time and streaming_start_time:
+                    elapsed = streaming_end_time - streaming_start_time
+                    char_count = len(full_response)
+                    # Estimate tokens (1 token ≈ 4 chars)
+                    estimated_tokens = char_count / 4
+                    tokens_per_sec = max(estimated_tokens / elapsed, 20.0) if elapsed > 0 else 70.0
+                    logger.info(f"Using measured streaming rate: {tokens_per_sec:.1f} tokens/sec")
+                else:
+                    tokens_per_sec = 70.0
+
+                # Stream improved output at same rate as original (unless it's for write_file/update_file)
+                assistant_text = []
+                if write_file_calls or update_file_calls:
+                    # Don't stream write_file/update_file content - it will be shown in collapsible widget
+                    logger.info("Skipping streaming for write_file/update_file fuzzed output")
+                    assistant_text.append(result.final_output)
+                else:
+                    # Stream normally for non-write/update fuzzing
+                    logger.info("Streaming fuzzed output (no write_file/update_file)")
+                    async for event in self._stream_text_with_rate(result.final_output, tokens_per_sec):
+                        assistant_text.append(event.data["text"])
+                        yield event
+
+                # Count tokens in fuzzed output (this is what goes into context, not the original)
+                try:
+                    import tiktoken
+                    encoding = tiktoken.get_encoding("cl100k_base")
+                    fuzzed_tokens = len(encoding.encode(result.final_output))
+                    session.cumulative_tokens += fuzzed_tokens
+                    logger.info(f"Session {session_id} added {fuzzed_tokens} tokens from fuzzed output (cumulative: {session.cumulative_tokens})")
+                except ImportError:
+                    # Fallback to usage tokens if tiktoken not available
+                    session.cumulative_tokens += usage_tokens
+                    logger.info(f"Session {session_id} added {usage_tokens} tokens (cumulative: {session.cumulative_tokens})")
+
+                # Now execute write_file and update_file calls with the fuzzed code
+                for tool_call in write_file_calls + update_file_calls:
+                    # Update content with fuzzed output (strip markdown fences as safety measure)
+                    fuzzed_content = strip_markdown_fences(result.final_output)
+                    tool_call["input"]["content"] = fuzzed_content
+
+                    result_exec = await self.tool_executor.execute(
+                        tool_name=tool_call["tool"],
+                        arguments=tool_call["input"],
+                        call_id=tool_call["call_id"]
+                    )
+
+                    # Emit tool result event
+                    yield StreamEvent(
+                        type="tool_result",
+                        data={
+                            "tool": tool_call["tool"],
+                            "success": result_exec.success,
+                            "result": result_exec.result if result_exec.success else result_exec.error
+                        }
+                    )
+
+                    # Add tool result message
+                    tool_result_msg = Message(id=self._generate_msg_id(), role="user")
+                    result_text = f"[Tool: {tool_call['tool']}]\n"
+                    if result_exec.success:
+                        result_text += f"Result: {result_exec.result}\n\n"
+                        # Add explicit instruction for search tools
+                        if tool_call['tool'] in ('search_code', 'search_files', 'list_directory'):
+                            result_text += "Present these results to the user clearly. Show file paths, line numbers, and relevant content."
+                    else:
+                        result_text += f"Error: {result_exec.error}"
+                    tool_result_msg.add_part(TextPart(text=result_text))
+                    self.conv_manager.add_message(session.conv_id, tool_result_msg)
+
+        # Add final assistant message to conversation
+        if assistant_text:
+            assistant_msg = Message(id=self._generate_msg_id(), role="assistant")
+            assistant_msg.add_part(TextPart(text="".join(assistant_text)))
+            self.conv_manager.add_message(session.conv_id, assistant_msg)
+
+        # Calculate exact token usage AFTER turn using harness-utils
+        context_after = self.conv_manager.calculate_context_usage(
+            session.conv_id,
+            model=session.provider.model
+        )
+
+        # Calculate tokens used this turn (delta from before turn started)
+        context_before = session.context_before_turn if hasattr(session, 'context_before_turn') else 0
+        total_turn_tokens = context_after - context_before
+
+        # Update cumulative tracking
+        if not should_fuzz and usage_tokens > 0:
+            session.cumulative_tokens += usage_tokens
+            logger.info(f"Session {session_id} added {usage_tokens} tokens (no fuzzing, cumulative: {session.cumulative_tokens})")
+
+        # Emit token usage event for this turn
+        yield StreamEvent(
+            type="token_usage",
+            data={
+                "total_tokens": context_after,  # Cumulative total for context counter
+                "turn_tokens": total_turn_tokens,  # Tokens used this turn
+                "output_tokens": usage_tokens,  # Output tokens from API
+                "input_tokens": total_turn_tokens - usage_tokens if total_turn_tokens > usage_tokens else 0,
+            }
+        )
+
+
+    def get_baseline(self, session_id: str) -> Optional[Baseline]:
+        """Get baseline for session."""
+        return self.baseline_manager.get(session_id)
+
+    def get_context_stats(self, session_id: str) -> dict[str, Any]:
+        """Get context statistics for session."""
+        session = self.sessions.get(session_id)
+        if not session:
+            raise ValueError(f"Session not found: {session_id}")
+
+        # Get conversation messages from harness-utils
+        messages = self.conv_manager.to_model_format(session.conv_id)
+
+        # Calculate exact token count using harness-utils
+        token_count = self.conv_manager.calculate_context_usage(
+            session.conv_id,
+            model=session.provider.model
+        )
+
+        logger.info(f"Context stats for {session_id}: {len(messages)} messages, {token_count} tokens")
+
+        return {
+            "message_count": len(messages),
+            "estimated_tokens": token_count,
+            "max_tokens": self.context_limit,
+        }
+
+    def compact_conversation(self, session_id: str) -> None:
+        """Compact conversation using harness-utils."""
+        session = self.sessions.get(session_id)
+        if not session:
+            raise ValueError(f"Session not found: {session_id}")
+
+        # Trigger compaction manually
+        self.conv_manager.prune_before_turn(session.conv_id)
+
+    def clear_conversation(self, session_id: str) -> None:
+        """Clear conversation history."""
+        session = self.sessions.get(session_id)
+        if not session:
+            raise ValueError(f"Session not found: {session_id}")
+
+        # Create a new conversation to replace the old one
+        old_conv_id = session.conv_id
+        new_conv = self.conv_manager.create_conversation(project_id="ctrl-code")
+        session.conv_id = new_conv.id
+
+        # Reset cumulative tokens
+        session.cumulative_tokens = 0
+
+        logger.info(f"Cleared conversation history for session {session_id} (old: {old_conv_id}, new: {new_conv.id})")
+
+    def _generate_msg_id(self) -> str:
+        """Generate message ID."""
+        return f"msg_{uuid.uuid4().hex[:12]}"
+
+    def _check_auto_chain(self, tool_name: str, result: dict) -> dict | None:
+        """
+        Check if we should auto-chain to another tool call.
+
+        Args:
+            tool_name: Name of the tool that just executed
+            result: Result from the tool
+
+        Returns:
+            Dict with tool call info if should chain, None otherwise
+        """
+        # search_files → read_file (if exactly 1 file found)
+        if tool_name == "search_files":
+            # Result is a list of file paths
+            if isinstance(result, list) and len(result) == 1:
+                file_path = result[0]
+                return {
+                    "tool": "read_file",
+                    "arguments": {"path": file_path},
+                    "call_id": f"auto_{uuid.uuid4().hex[:12]}"
+                }
+
+        return None
+
+    async def _stream_text_with_rate(
+        self, text: str, tokens_per_sec: float = 50.0
+    ) -> AsyncIterator[StreamEvent]:
+        """
+        Stream text in chunks to simulate realistic token generation rate.
+
+        Args:
+            text: Complete text to stream
+            tokens_per_sec: Rate to simulate (default ~50 tokens/sec, typical for LLMs)
+
+        Yields:
+            StreamEvent with text chunks
+        """
+        if not text:
+            return
+
+        # Split by lines to preserve newlines
+        lines = text.split('\n')
+
+        # Calculate delay: tokens/sec * 4 chars/token = chars/sec
+        # delay = 1 / chars_per_sec
+        chars_per_sec = tokens_per_sec * 4.0
+        delay_per_char = 1.0 / chars_per_sec
+
+        # Stream line by line, with char batches within lines
+        batch_size = 20  # Send 20 chars at a time for smooth streaming
+
+        for line_idx, line in enumerate(lines):
+            # Stream the line content in chunks
+            for i in range(0, len(line), batch_size):
+                chunk = line[i:i + batch_size]
+                yield StreamEvent(type="text", data={"text": chunk})
+
+                # Delay before next chunk
+                if i + batch_size < len(line):
+                    await asyncio.sleep(delay_per_char * len(chunk))
+
+            # Add newline after each line (except the last one)
+            if line_idx < len(lines) - 1:
+                yield StreamEvent(type="text", data={"text": "\n"})
+                await asyncio.sleep(delay_per_char)