Merge pull request #50 from huggingface/fix/ctrl-c-and-streaming-rework

agent/main.py

@@ -10,6 +10,7 @@ import argparse
 import asyncio
 import json
 import os
+import signal
 import sys
 import time
 from dataclasses import dataclass
@@ -268,6 +269,8 @@ class _ThinkingShimmer:
         self._task = asyncio.ensure_future(self._animate())
 
     def stop(self):
+        if not self._running:
+            return  # no-op when never started (e.g. headless mode)
         self._running = False
         if self._task:
             self._task.cancel()
@@ -312,7 +315,10 @@ class _ThinkingShimmer:
 
 
 class _StreamBuffer:
-    """Accumulates streamed tokens, renders full markdown on finish."""
+    """Accumulates streamed tokens, renders markdown block-by-block as complete
+    blocks appear. A "block" is everything up to a paragraph break (\\n\\n).
+    Unclosed code fences (odd count of ```) hold back flushing until closed so
+    a code block is always rendered as one unit."""
 
     def __init__(self, console):
         self._console = console
@@ -321,10 +327,41 @@ class _StreamBuffer:
     def add_chunk(self, text: str):
         self._buffer += text
 
-    def finish(self):
-        """Render the accumulated text as markdown, then reset."""
+    def _pop_block(self) -> str | None:
+        """Extract the next complete block, or return None if nothing complete."""
+        if self._buffer.count("```") % 2 == 1:
+            return None  # inside an open code fence — wait for close
+        idx = self._buffer.find("\n\n")
+        if idx == -1:
+            return None
+        block = self._buffer[:idx]
+        self._buffer = self._buffer[idx + 2:]
+        return block
+
+    async def flush_ready(
+        self,
+        cancel_event: "asyncio.Event | None" = None,
+        instant: bool = False,
+    ):
+        """Render any complete blocks that have accumulated; leave the tail."""
+        while True:
+            if cancel_event is not None and cancel_event.is_set():
+                return
+            block = self._pop_block()
+            if block is None:
+                return
+            if block.strip():
+                await print_markdown(block, cancel_event=cancel_event, instant=instant)
+
+    async def finish(
+        self,
+        cancel_event: "asyncio.Event | None" = None,
+        instant: bool = False,
+    ):
+        """Flush complete blocks, then render whatever incomplete tail remains."""
+        await self.flush_ready(cancel_event=cancel_event, instant=instant)
         if self._buffer.strip():
-            print_markdown(self._buffer)
+            await print_markdown(self._buffer, cancel_event=cancel_event, instant=instant)
         self._buffer = ""
 
     def discard(self):
@@ -338,6 +375,7 @@ async def event_listener(
     ready_event: asyncio.Event,
     prompt_session: PromptSession,
     config=None,
+    session_holder=None,
 ) -> None:
     """Background task that listens for events and displays them"""
     submission_id = [1000]
@@ -346,6 +384,12 @@ async def event_listener(
     shimmer = _ThinkingShimmer(console)
     stream_buf = _StreamBuffer(console)
 
+    def _cancel_event():
+        """Return the session's cancellation Event so print_markdown can abort
+        its typewriter loop mid-stream when Ctrl+C fires."""
+        s = session_holder[0] if session_holder else None
+        return s._cancelled if s is not None else None
+
     while True:
         try:
             event = await event_queue.get()
@@ -358,14 +402,19 @@ async def event_listener(
                 shimmer.stop()
                 content = event.data.get("content", "") if event.data else ""
                 if content:
-                    print_markdown(content)
+                    await print_markdown(content, cancel_event=_cancel_event())
             elif event.event_type == "assistant_chunk":
                 content = event.data.get("content", "") if event.data else ""
                 if content:
                     stream_buf.add_chunk(content)
+                    # Flush any complete markdown blocks progressively so the
+                    # user sees paragraphs appear as they're produced, not just
+                    # at the end of the whole response.
+                    shimmer.stop()
+                    await stream_buf.flush_ready(cancel_event=_cancel_event())
             elif event.event_type == "assistant_stream_end":
                 shimmer.stop()
-                stream_buf.finish()
+                await stream_buf.finish(cancel_event=_cancel_event())
             elif event.event_type == "tool_call":
                 shimmer.stop()
                 stream_buf.discard()
@@ -660,10 +709,33 @@ async def event_listener(
                             if gated is not None:
                                 print(f"Gated: {gated}")
 
-                    # Get user decision for this item
-                    response = await prompt_session.prompt_async(
-                        f"Approve item {i}? (y=yes, yolo=approve all, n=no, or provide feedback): "
-                    )
+                    # Get user decision for this item. Ctrl+C / EOF here is
+                    # treated as "reject remaining" (matches Codex's modal
+                    # priority and Forgecode's approval-cancel path). Without
+                    # this, KeyboardInterrupt kills the event listener and
+                    # the main loop deadlocks waiting for turn_complete.
+                    try:
+                        response = await prompt_session.prompt_async(
+                            f"Approve item {i}? (y=yes, yolo=approve all, n=no, or provide feedback): "
+                        )
+                    except (KeyboardInterrupt, EOFError):
+                        get_console().print("[dim]Approval cancelled — rejecting remaining items[/dim]")
+                        approvals.append(
+                            {
+                                "tool_call_id": tool_call_id,
+                                "approved": False,
+                                "feedback": "User cancelled approval",
+                            }
+                        )
+                        for remaining in tools_data[i:]:
+                            approvals.append(
+                                {
+                                    "tool_call_id": remaining.get("tool_call_id", ""),
+                                    "approved": False,
+                                    "feedback": None,
+                                }
+                            )
+                        break
 
                     response = response.strip().lower()
 
@@ -910,44 +982,94 @@ async def main():
             ready_event,
             prompt_session,
             config,
+            session_holder=session_holder,
         )
     )
 
     await ready_event.wait()
 
     submission_id = [0]
-    last_interrupt_time = 0.0
-    agent_busy = False  # True only while the agent is processing a submission
+    # Mirrors codex-rs/tui/src/bottom_pane/mod.rs:137
+    # (`QUIT_SHORTCUT_TIMEOUT = Duration::from_secs(1)`). Two Ctrl+C presses
+    # within this window quit; a single press cancels the in-flight turn.
+    CTRL_C_QUIT_WINDOW = 1.0
+    # Hint string matches codex-rs/tui/src/bottom_pane/footer.rs:746
+    # (`" again to quit"` prefixed with the key binding, rendered dim).
+    CTRL_C_HINT = "[dim]ctrl + c again to quit[/dim]"
+    interrupt_state = {"last": 0.0, "exit": False}
+
+    loop = asyncio.get_running_loop()
+
+    def _on_sigint() -> None:
+        """SIGINT handler — fires while the agent is generating (terminal is
+        in cooked mode between prompts). Mirrors Codex's `on_ctrl_c` in
+        codex-rs/tui/src/chatwidget.rs: first press cancels active work and
+        arms the quit hint; second press within the window quits."""
+        now = time.monotonic()
+        session = session_holder[0]
+
+        if now - interrupt_state["last"] < CTRL_C_QUIT_WINDOW:
+            interrupt_state["exit"] = True
+            if session:
+                session.cancel()
+            # Wake the main loop out of turn_complete_event.wait()
+            turn_complete_event.set()
+            return
+
+        interrupt_state["last"] = now
+        if session and not session.is_cancelled:
+            session.cancel()
+        get_console().print(f"\n{CTRL_C_HINT}")
+
+    def _install_sigint() -> bool:
+        try:
+            loop.add_signal_handler(signal.SIGINT, _on_sigint)
+            return True
+        except (NotImplementedError, RuntimeError):
+            return False  # Windows or non-main thread
+
+    # prompt_toolkit's prompt_async installs its own SIGINT handler and, on
+    # exit, calls loop.remove_signal_handler(SIGINT) — which wipes ours too.
+    # So we re-arm at the top of every loop iteration, right before the busy
+    # wait. Without this, Ctrl+C during agent streaming after the first turn
+    # falls through to the default handler and the terminal just echoes ^C.
+    sigint_available = _install_sigint()
 
     try:
         while True:
-            # Wait for previous turn to complete, with interrupt support
+            if sigint_available:
+                _install_sigint()
+
             try:
                 await turn_complete_event.wait()
             except asyncio.CancelledError:
                 break
             turn_complete_event.clear()
-            agent_busy = False
 
-            # Get user input
+            if interrupt_state["exit"]:
+                break
+
+            # Get user input. prompt_toolkit puts the terminal in raw mode and
+            # installs its own SIGINT handling; ^C arrives as \x03 and surfaces
+            # as KeyboardInterrupt here. On return, prompt_toolkit removes the
+            # loop's SIGINT handler — we re-arm at the top of the next iter.
             try:
                 user_input = await get_user_input(prompt_session)
             except EOFError:
                 break
             except KeyboardInterrupt:
                 now = time.monotonic()
-                if now - last_interrupt_time < 3.0:
+                if now - interrupt_state["last"] < CTRL_C_QUIT_WINDOW:
                     break
-                last_interrupt_time = now
-                # If agent is actually working, cancel it
-                session = session_holder[0]
-                if agent_busy and session:
-                    session.cancel()
-                else:
-                    get_console().print("[dim]Ctrl+C again to exit[/dim]")
-                    turn_complete_event.set()
+                interrupt_state["last"] = now
+                get_console().print(CTRL_C_HINT)
+                turn_complete_event.set()
                 continue
 
+            # A successful read ends the double-press window — an unrelated
+            # Ctrl+C during the next turn should start a fresh arming.
+            interrupt_state["last"] = 0.0
+
             # Check for exit commands
             if user_input.strip().lower() in ["exit", "quit", "/quit", "/exit"]:
                 break
@@ -967,7 +1089,6 @@ async def main():
                     turn_complete_event.set()
                     continue
                 else:
-                    agent_busy = True
                     await submission_queue.put(sub)
                     continue
 
@@ -979,11 +1100,16 @@ async def main():
                     op_type=OpType.USER_INPUT, data={"text": user_input}
                 ),
             )
-            agent_busy = True
             await submission_queue.put(submission)
 
     except KeyboardInterrupt:
         pass
+    finally:
+        if sigint_available:
+            try:
+                loop.remove_signal_handler(signal.SIGINT)
+            except (NotImplementedError, RuntimeError):
+                pass
 
     # Shutdown
     shutdown_submission = Submission(
@@ -1071,13 +1197,17 @@ async def headless_main(
     )
     await submission_queue.put(submission)
 
-    # Process events until turn completes
+    # Process events until turn completes. Headless mode is for scripts /
+    # log capture: no shimmer animation, no typewriter, no live-redrawing
+    # research overlay. Output is plain, append-only text.
     console = _create_rich_console()
-    shimmer = _ThinkingShimmer(console)
     stream_buf = _StreamBuffer(console)
     _hl_last_tool = [None]
     _hl_sub_id = [1]
-    shimmer.start()
+    # Research sub-agent tool calls are buffered and dumped once the sub-agent
+    # finishes, instead of streaming via the live redrawing SubAgentDisplay.
+    _hl_research_calls: list[str] = []
+    _hl_in_research = [False]
 
     while True:
         event = await event_queue.get()
@@ -1086,16 +1216,14 @@ async def headless_main(
             content = event.data.get("content", "") if event.data else ""
             if content:
                 stream_buf.add_chunk(content)
+                await stream_buf.flush_ready(instant=True)
         elif event.event_type == "assistant_stream_end":
-            shimmer.stop()
-            stream_buf.finish()
+            await stream_buf.finish(instant=True)
         elif event.event_type == "assistant_message":
-            shimmer.stop()
             content = event.data.get("content", "") if event.data else ""
             if content:
-                print_markdown(content)
+                await print_markdown(content, instant=True)
         elif event.event_type == "tool_call":
-            shimmer.stop()
             stream_buf.discard()
             tool_name = event.data.get("tool", "") if event.data else ""
             arguments = event.data.get("arguments", {}) if event.data else {}
@@ -1109,11 +1237,33 @@ async def headless_main(
             success = event.data.get("success", False) if event.data else False
             if _hl_last_tool[0] == "plan_tool" and output:
                 print_tool_output(output, success, truncate=False)
-            shimmer.start()
         elif event.event_type == "tool_log":
             tool = event.data.get("tool", "") if event.data else ""
             log = event.data.get("log", "") if event.data else ""
-            if log:
+            if not log:
+                pass
+            elif tool == "research":
+                # Buffer research sub-agent activity; on completion, dump a
+                # single static block that mirrors the live overlay's styling
+                # without its line-erasing redraws (unfit for non-TTY output).
+                if log == "Starting research sub-agent...":
+                    _hl_in_research[0] = True
+                    _hl_research_calls.clear()
+                elif log == "Research complete.":
+                    _hl_in_research[0] = False
+                    f = get_console().file
+                    f.write("  \033[38;2;255;200;80m▸ research\033[0m\n")
+                    for call in _hl_research_calls:
+                        f.write(f"    \033[2m{call}\033[0m\n")
+                    f.flush()
+                    _hl_research_calls.clear()
+                elif log.startswith("tokens:") or log.startswith("tools:"):
+                    pass  # stats updates — only useful for the live display
+                elif _hl_in_research[0]:
+                    _hl_research_calls.append(log)
+                else:
+                    print_tool_log(tool, log)
+            else:
                 print_tool_log(tool, log)
         elif event.event_type == "approval_required":
             # Auto-approve everything in headless mode (safety net if yolo_mode
@@ -1140,13 +1290,11 @@ async def headless_main(
             new_tokens = event.data.get("new_tokens", 0) if event.data else 0
             print_compacted(old_tokens, new_tokens)
         elif event.event_type == "error":
-            shimmer.stop()
             stream_buf.discard()
             error = event.data.get("error", "Unknown error") if event.data else "Unknown error"
             print_error(error)
             break
         elif event.event_type in ("turn_complete", "interrupted"):
-            shimmer.stop()
             stream_buf.discard()
             history_size = event.data.get("history_size", "?") if event.data else "?"
             print(f"\n--- Agent {event.event_type} (history_size={history_size}) ---", file=sys.stderr)

agent/utils/terminal_display.py

@@ -3,10 +3,22 @@ Terminal display utilities — rich-powered CLI formatting.
 """
 
 from rich.console import Console
-from rich.markdown import Markdown
+from rich.markdown import Heading, Markdown
 from rich.panel import Panel
 from rich.theme import Theme
 
+
+class _LeftHeading(Heading):
+    """Rich's default Markdown renders h1/h2 centered via Align.center.
+    Yield the styled text directly so headings stay left-aligned."""
+
+    def __rich_console__(self, console, options):
+        self.text.justify = "left"
+        yield self.text
+
+
+Markdown.elements["heading_open"] = _LeftHeading
+
 _THEME = Theme({
     "tool.name": "bold rgb(255,200,80)",
     "tool.args": "dim",
@@ -235,8 +247,13 @@ def print_tool_log(tool: str, log: str) -> None:
 
 # ── Messages ───────────────────────────────────────────────────────────
 
-def print_markdown(text: str) -> None:
-    import io, time, random
+async def print_markdown(
+    text: str,
+    cancel_event: "asyncio.Event | None" = None,
+    instant: bool = False,
+) -> None:
+    import asyncio
+    import io, random
     from rich.padding import Padding
 
     _console.print()
@@ -260,21 +277,36 @@ def print_markdown(text: str) -> None:
     lines = rendered.split("\n")
     rendered = "\n".join(line.rstrip() for line in lines)
 
-    # CRT typewriter effect — fast, with occasional glitch
-    rng = random.Random(42)
     f = _console.file
+
+    # Headless / non-interactive: dump the rendered markdown in one write.
+    if instant:
+        f.write(rendered)
+        f.write("\n")
+        f.flush()
+        return
+
+    # CRT typewriter effect — async so the event loop can service signal
+    # handlers (Ctrl+C during streaming) between characters. If cancelled
+    # mid-type, stop cleanly: write an ANSI reset so half-open color state
+    # doesn't bleed onto the "interrupted" line, and return.
+    rng = random.Random(42)
+    cancelled = False
     for ch in rendered:
+        if cancel_event is not None and cancel_event.is_set():
+            cancelled = True
+            break
         f.write(ch)
         f.flush()
         if ch == "\n":
-            time.sleep(0.002)
+            await asyncio.sleep(0.002)
         elif ch == " ":
-            time.sleep(0.002)
+            await asyncio.sleep(0.002)
         elif rng.random() < 0.03:
-            time.sleep(0.015)
+            await asyncio.sleep(0.015)
         else:
-            time.sleep(0.004)
-    f.write("\n")
+            await asyncio.sleep(0.004)
+    f.write("\033[0m\n" if cancelled else "\n")
     f.flush()