From e00b88b664b881890d0e2b7ad5019379e2d9f9b0 Mon Sep 17 00:00:00 2001 From: Mike Pfaffenberger Date: Thu, 9 Jul 2026 21:05:43 -0400 Subject: [PATCH 01/78] feat(theme): add accessible theme-aware terminal styling Add the Green Screen palette, improve Purple Puppy contrast, paginate the theme picker, and make prompts plus Termflow markdown and syntax highlighting follow the active theme. Also align shutdown callbacks with their awaited async contract. --- code_puppy/agents/event_stream_handler.py | 7 +- code_puppy/callbacks.py | 48 +++- code_puppy/messaging/bar_painters.py | 26 ++- code_puppy/plugins/theme/README.md | 21 +- code_puppy/plugins/theme/bundled_palettes.py | 28 ++- code_puppy/plugins/theme/picker.py | 38 ++- .../plugins/theme/register_callbacks.py | 123 +++++++++- code_puppy/plugins/theme/themes.py | 41 +++- tests/messaging/test_bottom_bar.py | 12 + tests/plugins/test_theme_plugin.py | 218 +++++++++++++++++- tests/test_callbacks_extended.py | 38 +++ 11 files changed, 571 insertions(+), 29 deletions(-) diff --git a/code_puppy/agents/event_stream_handler.py b/code_puppy/agents/event_stream_handler.py index f1fd546c5..ee218050a 100644 --- a/code_puppy/agents/event_stream_handler.py +++ b/code_puppy/agents/event_stream_handler.py @@ -151,7 +151,10 @@ async def event_stream_handler( from termflow import Parser as TermflowParser from termflow import Renderer as TermflowRenderer - from termflow.render.style import RenderFeatures + from termflow.render.style import RenderFeatures, RenderStyle + from termflow.syntax import Highlighter + + from code_puppy.callbacks import on_termflow_highlighter, on_termflow_style # Use the module-level console (set via set_streaming_console) console = get_streaming_console() @@ -187,7 +190,9 @@ def _make_text_renderer(index: int) -> TermflowRenderer: return TermflowRenderer( output=output, width=console.width, + style=on_termflow_style(RenderStyle.default()), features=RenderFeatures(clipboard=False), + highlighter=on_termflow_highlighter(Highlighter()), ) # Smooth-stream state for thinking parts. Each index maps to a smoother diff --git a/code_puppy/callbacks.py b/code_puppy/callbacks.py index 2041f9ba1..0247c075f 100644 --- a/code_puppy/callbacks.py +++ b/code_puppy/callbacks.py @@ -28,6 +28,9 @@ "post_tool_call", "stream_event", "thinking_display_filter", + "termflow_style", + "termflow_highlighter", + "prompt_text_color", "register_tools", "register_agent_tools", "register_agents", @@ -87,6 +90,9 @@ "post_tool_call": [], "stream_event": [], "thinking_display_filter": [], + "termflow_style": [], + "termflow_highlighter": [], + "prompt_text_color": [], "register_tools": [], "register_agent_tools": [], "register_agents": [], @@ -324,8 +330,8 @@ async def on_startup() -> List[Any]: return await _trigger_callbacks("startup") -def on_shutdown() -> List[Any]: - return _trigger_callbacks_sync("shutdown") +async def on_shutdown() -> List[Any]: + return await _trigger_callbacks("shutdown") async def on_invoke_agent(*args, **kwargs) -> List[Any]: @@ -630,6 +636,44 @@ def on_thinking_display_filter( return current +def _chain_value_callbacks(phase: PhaseType, default: Any) -> Any: + """Chain callbacks that optionally replace a single value.""" + current = default + for callback in get_callbacks(phase): + try: + result = callback(current) + if result is not None: + current = result + except Exception as exc: + logger.error( + "%s callback %s failed: %s\n%s", + phase, + callback.__name__, + exc, + traceback.format_exc(), + ) + return current + + +def on_termflow_style(default_style: Any) -> Any: + """Let plugins replace Termflow's Markdown rendering style. + + Callbacks are chained in registration order. Returning ``None`` leaves the + current style unchanged, and failures degrade safely to the prior style. + """ + return _chain_value_callbacks("termflow_style", default_style) + + +def on_termflow_highlighter(default_highlighter: Any) -> Any: + """Let plugins replace Termflow's syntax highlighter.""" + return _chain_value_callbacks("termflow_highlighter", default_highlighter) + + +def on_prompt_text_color(default_color: str | None = None) -> str | None: + """Resolve the persistent prompt buffer's truecolor foreground.""" + return _chain_value_callbacks("prompt_text_color", default_color) + + async def on_stream_event( event_type: str, event_data: Any, agent_session_id: str | None = None ) -> List[Any]: diff --git a/code_puppy/messaging/bar_painters.py b/code_puppy/messaging/bar_painters.py index 0b200cbfb..18bf2b990 100644 --- a/code_puppy/messaging/bar_painters.py +++ b/code_puppy/messaging/bar_painters.py @@ -27,6 +27,8 @@ from __future__ import annotations +import re + from .bar_rendering import ( CLEAR_LINE as _CLEAR_LINE, ) @@ -77,6 +79,23 @@ #: the repo-wide "bold cyan" brand accent (rich_renderer et al.). _SELECT_ON = "\x1b[1;36m" _SELECT_OFF = "\x1b[22;39m" # reset weight + foreground only +_HEX_COLOR_RE = re.compile(r"^#[0-9a-fA-F]{6}$") + + +def _prompt_color_sgr() -> str: + """Resolve a plugin-provided truecolor foreground for prompt text.""" + try: + from code_puppy.callbacks import on_prompt_text_color + + color = on_prompt_text_color() + if color and _HEX_COLOR_RE.fullmatch(color): + red, green, blue = ( + int(color[index : index + 2], 16) for index in (1, 3, 5) + ) + return f"\x1b[38;2;{red};{green};{blue}m" + except Exception: + pass + return "" def _dim(text: str) -> str: @@ -264,8 +283,13 @@ def _prompt_seq(self) -> str: prefix_sgrs=getattr(self, "_prompt_prefix_sgrs", None), ) parts = [_SAVE_CURSOR, _WRAP_OFF] + prompt_sgr = _prompt_color_sgr() + prompt_reset = "\x1b[39m" if prompt_sgr else "" for i, rendered in enumerate(rendered_rows): - parts.append(f"\x1b[{prompt_top + i};1H{_CLEAR_LINE}{rendered}") + parts.append( + f"\x1b[{prompt_top + i};1H{_CLEAR_LINE}" + f"{prompt_sgr}{rendered}{prompt_reset}" + ) parts.append(_WRAP_ON) parts.append(_RESTORE_CURSOR) return "".join(parts) diff --git a/code_puppy/plugins/theme/README.md b/code_puppy/plugins/theme/README.md index e571edc59..bb25b71f9 100644 --- a/code_puppy/plugins/theme/README.md +++ b/code_puppy/plugins/theme/README.md @@ -10,7 +10,7 @@ escape sequences (mure-style). ``` /theme ``` -Launches a split-panel TUI with 15 themes (6 neon/dark themes, 7 palette-first/light themes, surprise, reset). +Launches a paginated split-panel TUI with 16 themes (7 neon/dark themes, 7 palette-first/light themes, surprise, reset). Five themes are shown per page; use Up/Down for item navigation and PgUp/PgDn to change pages. ## What gets themed @@ -24,7 +24,7 @@ Launches a split-panel TUI with 15 themes (6 neon/dark themes, 7 palette-first/l **Semantic colors are preserved** at Level 2 (red errors stay angry, yellow warnings warn). Level 3 OSC remap is broader — it changes how *your terminal itself* interprets every ANSI color, so e.g. an `ls` ran after `/theme tokyo-night` will also look Tokyo Night. Pure terminal-level magic. -Themes persist to your Code Puppy config and survive restarts. The OSC palette auto-resets on Code Puppy exit (via `atexit`) so you never get a stuck-pink terminal. +Tokyo Night is applied by default on first run. Themes persist to your Code Puppy config and survive restarts; existing or explicitly restored choices are never overwritten by the default. The OSC palette auto-resets on Code Puppy exit (via `atexit`) so you never get a stuck-pink terminal. ## Bundled themes @@ -39,19 +39,22 @@ Themes persist to your Code Puppy config and survive restarts. The OSC palette a | 7 | 🐱 Catppuccin Mocha | soothing pastel dark (mure default) | | 8 | ☕ Catppuccin Latte | soothing pastel light | | 9 | 🌃 Tokyo Night | neon-on-navy night | -| 10 | 🌑 Deep Black | inky noir with subtle neon edge | -| 11 | ☀️ Solarized Light | classic warm beige with calm accents | -| 12 | 📄 GitHub Light | crisp white, familiar code colors | -| 13 | 🌸 Rose Pine Dawn | soft pastel rose light | -| 14 | 🎲 Surprise Me | a fresh random remix every time | -| 15 | 🔄 Restore Defaults | back to Code Puppy + terminal factory | +| 10 | [CRT] Green Screen | llxprt's radioactive green phosphor CRT | +| 11 | Deep Black | inky noir with subtle neon edge | +| 12 | Solarized Light | classic warm beige with calm accents | +| 13 | GitHub Light | crisp white, familiar code colors | +| 14 | Rose Pine Dawn | soft pastel rose light | +| 15 | Surprise Me | a fresh random remix every time | +| 16 | Restore Defaults | back to Code Puppy + terminal factory | + +Green Screen ports the canonical phosphor colors from [vybestack/llxprt-code](https://github.com/vybestack/llxprt-code), whose theme is licensed under Apache-2.0. ## Power-user shortcuts ``` /theme 5 apply theme #5 (Bubblegum Pink) /theme bubblegum apply by alias (also: pink, puppy, purple, mocha, - latte, tokyo, solarized, github, rose-pine) + latte, tokyo, green, crt, solarized, github, rose-pine) /theme tokyo-night apply by canonical name /theme surprise re-roll a random palette /theme default restore Code Puppy + terminal factory colors diff --git a/code_puppy/plugins/theme/bundled_palettes.py b/code_puppy/plugins/theme/bundled_palettes.py index 6347988ac..400c51d12 100644 --- a/code_puppy/plugins/theme/bundled_palettes.py +++ b/code_puppy/plugins/theme/bundled_palettes.py @@ -84,6 +84,32 @@ ], } +GREEN_SCREEN = { + # Ported from llxprt-code's canonical Green Screen theme: black glass, + # green phosphor, and one eye-searing highlight. Yes, it is meant to be + # this green. + "bg": "#000000", + "fg": "#6a9955", + "ansi": [ + "#000000", # 0 black — CRT glass + "#6a9955", # 1 red — phosphor green + "#6a9955", # 2 green — phosphor green + "#6a9955", # 3 yellow — phosphor green + "#6a9955", # 4 blue — phosphor green + "#6a9955", # 5 magenta — phosphor green + "#6a9955", # 6 cyan — phosphor green + "#6a9955", # 7 white — phosphor green + "#3a5945", # 8 bright black — dark phosphor + "#6a9955", # 9 bright red + "#00ff00", # 10 bright green — radioactive highlight + "#6a9955", # 11 bright yellow + "#6a9955", # 12 bright blue + "#6a9955", # 13 bright magenta + "#6a9955", # 14 bright cyan + "#00ff00", # 15 bright white — radioactive highlight + ], +} + DEEP_BLACK = { "bg": "#050505", "fg": "#e6e6e6", @@ -287,7 +313,7 @@ "#ff7fa8", # 5 magenta — tongue pink "#d9a7f5", # 6 cyan — soft lilac "#f0e3ff", # 7 white — muzzle lavender-white - "#4b1a6e", # 8 bright black — shade purple + "#a98aca", # 8 bright black — readable muted lavender "#ff7fa8", # 9 bright red "#dba1ff", # 10 bright green "#ffd58f", # 11 bright yellow diff --git a/code_puppy/plugins/theme/picker.py b/code_puppy/plugins/theme/picker.py index f0d386e4d..8e705c35e 100644 --- a/code_puppy/plugins/theme/picker.py +++ b/code_puppy/plugins/theme/picker.py @@ -58,6 +58,8 @@ # Inline-markup samples to demonstrate the Level 2 color remap. # These mirror the kind of hardcoded tags scattered through the renderer. +THEMES_PER_PAGE = 5 + INLINE_MARKUP_SAMPLES = [ "[bold cyan]bold cyan headline[/bold cyan]", "[dim cyan]dim cyan detail[/dim cyan]", @@ -189,13 +191,31 @@ def _render_preview(theme_name: str, surprise_seed: int) -> ANSI: return ANSI(buffer.getvalue()) +def _total_pages() -> int: + return max(1, (len(MENU) + THEMES_PER_PAGE - 1) // THEMES_PER_PAGE) + + +def _page_for_index(selected_index: int) -> int: + return selected_index // THEMES_PER_PAGE + + +def _move_page(selected_index: int, delta: int) -> int: + """Move one page while retaining the selected row where possible.""" + return max(0, min(selected_index + delta * THEMES_PER_PAGE, len(MENU) - 1)) + + def _format_menu(selected_index: int) -> FormattedText: - """Build the left-hand menu text with a highlighted selection.""" + """Build the current page of the left-hand menu.""" + page = _page_for_index(selected_index) + page_start = page * THEMES_PER_PAGE + page_end = min(page_start + THEMES_PER_PAGE, len(MENU)) lines: list[tuple[str, str]] = [ ("bold cyan", "Pick a Theme"), + ("fg:ansibrightblack", f" Page {page + 1}/{_total_pages()}"), ("", "\n\n"), ] - for i, (name, theme) in enumerate(MENU): + for i in range(page_start, page_end): + _, theme = MENU[i] prefix = "> " if i == selected_index else " " style = "fg:ansigreen bold" if i == selected_index else "" line = f"{prefix}{i + 1}. {theme['icon']} {theme['label']}" @@ -208,7 +228,7 @@ def _format_menu(selected_index: int) -> FormattedText: lines.append( ( "fg:ansicyan", - "Up/Down Navigate | Enter Apply | Esc / Ctrl-C Cancel", + "Up/Down Navigate | PgUp/PgDn Page\nEnter Apply | Esc / Ctrl-C Cancel", ) ) return FormattedText(lines) @@ -253,6 +273,18 @@ def _(event): _refresh_surprise_seed_if_focused() event.app.invalidate() + @kb.add("pageup") + def _(event): + selected[0] = _move_page(selected[0], -1) + _refresh_surprise_seed_if_focused() + event.app.invalidate() + + @kb.add("pagedown") + def _(event): + selected[0] = _move_page(selected[0], 1) + _refresh_surprise_seed_if_focused() + event.app.invalidate() + @kb.add("enter") def _(event): result[0] = MENU[selected[0]][0] diff --git a/code_puppy/plugins/theme/register_callbacks.py b/code_puppy/plugins/theme/register_callbacks.py index 1c112a842..dcfb1ac92 100644 --- a/code_puppy/plugins/theme/register_callbacks.py +++ b/code_puppy/plugins/theme/register_callbacks.py @@ -2,16 +2,16 @@ UX: /theme → interactive split-panel picker with live preview - /theme → apply theme number N (1-15) + /theme → apply theme number N (1-16) /theme → apply by name (ocean, forest, sunset, vaporwave, bubblegum-pink, purple-puppy, mocha, latte, - tokyo-night, deep-black, solarized-light, + tokyo-night, green-screen, deep-black, solarized-light, github-light, rose-pine-dawn, surprise, default) /theme reset → restore Code Puppy defaults (alias of /theme default) /theme show → show current banner → color mapping -The 14th option (Surprise Me) re-rolls a random palette every time. -The 15th option (Restore Defaults) puts everything back to factory. +The 15th option (Surprise Me) re-rolls a random palette every time. +The 16th option (Restore Defaults) puts everything back to factory. Plays nice with /colors — same color pool, same config keys. """ @@ -25,7 +25,9 @@ from code_puppy.command_line.colors_menu import BANNER_DISPLAY_INFO from code_puppy.config import ( get_all_banner_colors, + get_value, reset_all_banner_colors, + set_config_value, ) from code_puppy.messaging import emit_error, emit_info, emit_warning @@ -44,6 +46,9 @@ ) _INTERACTIVE_TIMEOUT_SECONDS = 300 # 5 min — generous; user is browsing +_ACTIVE_THEME_CONFIG_KEY = "theme_active_theme" +_DEFAULT_THEME_NAME = "tokyo-night" +_LEGACY_THEME_NAME = "legacy-custom" def _custom_help(): @@ -89,7 +94,7 @@ def _run_interactive_picker() -> str | None: return future.result(timeout=_INTERACTIVE_TIMEOUT_SECONDS) -def _apply_theme(theme_name: str) -> None: +def _apply_theme(theme_name: str, *, announce: bool = True) -> None: """Apply banner colors + content styles + inline remap + terminal palette. `default` is special: resets banner config, content styles, Rich color @@ -111,7 +116,109 @@ def _apply_theme(theme_name: str) -> None: if terminal_palette: osc.apply_palette(terminal_palette) - _announce_applied(theme_name) + set_config_value(_ACTIVE_THEME_CONFIG_KEY, theme_name) + if announce: + _announce_applied(theme_name) + + +def _active_terminal_palette() -> tuple[str, dict] | None: + """Return the active curated theme and its complete persisted palette.""" + active_theme = get_value(_ACTIVE_THEME_CONFIG_KEY) + if not active_theme or active_theme in {"default", _LEGACY_THEME_NAME}: + return None + palette = osc.get_saved_palette() + if not palette or len(palette.get("ansi") or []) < 16: + return None + return active_theme, palette + + +def _termflow_style(default_style): + """Derive Termflow's truecolor Markdown chrome from the active theme.""" + active = _active_terminal_palette() + if active is None: + return default_style + _, palette = active + ansi = palette["ansi"] + + from termflow.render.style import RenderStyle + + return RenderStyle( + bright=ansi[14], + head=ansi[10], + symbol=ansi[13], + grey=ansi[8], + dark=palette.get("bg", ansi[0]), + mid=ansi[8], + light=ansi[7], + link=ansi[12], + error=ansi[9], + ) + + +def _termflow_highlighter(default_highlighter): + """Build syntax colors from the active theme instead of Monokai.""" + active = _active_terminal_palette() + if active is None: + return default_highlighter + active_theme, palette = active + ansi = palette["ansi"] + + from pygments.formatters.terminal256 import TerminalTrueColorFormatter + from pygments.style import Style + from pygments.token import ( + Comment, + Error, + Keyword, + Name, + Number, + Operator, + String, + Token, + ) + from termflow.syntax import Highlighter + + if active_theme in {"green-screen", "green", "crt"}: + token_styles = {Token: palette["fg"]} + else: + token_styles = { + Token: ansi[7], + Comment: ansi[8], + Error: ansi[9], + Keyword: ansi[13], + Name.Function: ansi[14], + Number: ansi[11], + Operator: ansi[6], + String: ansi[10], + } + + theme_style = type( + "CodePuppyThemeStyle", + (Style,), + {"background_color": palette.get("bg"), "styles": token_styles}, + ) + highlighter = Highlighter() + highlighter._formatter = TerminalTrueColorFormatter(style=theme_style) + return highlighter + + +def _prompt_text_color(default_color): + """Use the active terminal foreground for the persistent prompt buffer.""" + active = _active_terminal_palette() + return active[1].get("fg", default_color) if active else default_color + + +def _apply_default_theme_on_first_run() -> None: + """Apply Tokyo Night once, preserving explicit and legacy theme choices.""" + if get_value(_ACTIVE_THEME_CONFIG_KEY): + return + + if osc.get_saved_palette(): + # Theme persistence predates the active-theme marker. Treat an existing + # palette as an intentional choice and migrate without changing it. + set_config_value(_ACTIVE_THEME_CONFIG_KEY, _LEGACY_THEME_NAME) + return + + _apply_theme(_DEFAULT_THEME_NAME, announce=False) # --- Command handler -------------------------------------------------------- @@ -159,5 +266,9 @@ def _handle_theme(command: str, name: str): return True +register_callback("startup", _apply_default_theme_on_first_run) +register_callback("termflow_style", _termflow_style) +register_callback("termflow_highlighter", _termflow_highlighter) +register_callback("prompt_text_color", _prompt_text_color) register_callback("custom_command_help", _custom_help) register_callback("custom_command", _handle_theme) diff --git a/code_puppy/plugins/theme/themes.py b/code_puppy/plugins/theme/themes.py index b67077b69..58edb2db7 100644 --- a/code_puppy/plugins/theme/themes.py +++ b/code_puppy/plugins/theme/themes.py @@ -344,6 +344,38 @@ def _parseable(name: str) -> bool: "color_remap": {}, "terminal_palette": bp.TOKYO_NIGHT, }, + "green-screen": { + "icon": "[CRT]", + "label": "Green Screen", + "blurb": "llxprt's radioactive green phosphor CRT", + # Banner labels render as white, which this theme remaps to #6a9955. + # Keep their backgrounds dark enough for readable phosphor-on-glass. + "colors": [ + "#071507", + "#0b1f0b", + "#102510", + ], + "content_styles": { + "info": "#6a9955", + "warning": "#6a9955", + "success": "#00ff00", + "error": "bold #6a9955", + "debug": "dim #4a7035", + "diff_add": "#00ff00", + "diff_remove": "#6a9955", + "diff_context": "dim #4a7035", + }, + "color_remap": rt.make_remap( + cyan="#6a9955", + blue="#6a9955", + magenta="#6a9955", + bright_cyan="#6a9955", + bright_blue="#6a9955", + bright_magenta="#6a9955", + white="#6a9955", + ), + "terminal_palette": bp.GREEN_SCREEN, + }, "deep-black": { "icon": "🌑", "label": "Deep Black", @@ -464,7 +496,7 @@ def _parseable(name: str) -> bool: }, } -# The 13th option: surprise me — special case, regenerated each preview/apply. +# Surprise Me is a special case, regenerated each preview/apply. SURPRISE = { "icon": "🎲", "label": "Surprise Me", @@ -475,7 +507,7 @@ def _parseable(name: str) -> bool: "terminal_palette": None, # randomized at apply time } -# The 15th option: restore Code Puppy defaults (banners + content). +# Restore Code Puppy defaults (banners + content). DEFAULT = { "icon": "🔄", "label": "Restore Defaults", @@ -486,7 +518,7 @@ def _parseable(name: str) -> bool: "terminal_palette": None, # handled specially (OSC reset) } -# Ordered menu: 6 vivid/dark themes, 7 palette-first/light themes, surprise, default. +# Ordered menu: 7 vivid/dark themes, 7 palette-first/light themes, surprise, default. MENU: list[tuple[str, dict]] = [ ("ocean", CURATED_THEMES["ocean"]), ("forest", CURATED_THEMES["forest"]), @@ -497,6 +529,7 @@ def _parseable(name: str) -> bool: ("catppuccin-mocha", CURATED_THEMES["catppuccin-mocha"]), ("catppuccin-latte", CURATED_THEMES["catppuccin-latte"]), ("tokyo-night", CURATED_THEMES["tokyo-night"]), + ("green-screen", CURATED_THEMES["green-screen"]), ("deep-black", CURATED_THEMES["deep-black"]), ("solarized-light", CURATED_THEMES["solarized-light"]), ("github-light", CURATED_THEMES["github-light"]), @@ -514,6 +547,8 @@ def _parseable(name: str) -> bool: MENU_BY_NAME["mocha"] = CURATED_THEMES["catppuccin-mocha"] MENU_BY_NAME["latte"] = CURATED_THEMES["catppuccin-latte"] MENU_BY_NAME["tokyo"] = CURATED_THEMES["tokyo-night"] +MENU_BY_NAME["green"] = CURATED_THEMES["green-screen"] +MENU_BY_NAME["crt"] = CURATED_THEMES["green-screen"] MENU_BY_NAME["bubblegum"] = CURATED_THEMES["bubblegum-pink"] MENU_BY_NAME["pink"] = CURATED_THEMES["bubblegum-pink"] MENU_BY_NAME["puppy"] = CURATED_THEMES["purple-puppy"] diff --git a/tests/messaging/test_bottom_bar.py b/tests/messaging/test_bottom_bar.py index 286025da0..5a04d03ce 100644 --- a/tests/messaging/test_bottom_bar.py +++ b/tests/messaging/test_bottom_bar.py @@ -1,6 +1,7 @@ """Tests for code_puppy.messaging.bottom_bar - the scroll-region manager.""" import io +from unittest.mock import patch import pytest @@ -54,6 +55,17 @@ def drain(stream): return value +def test_prompt_buffer_uses_theme_foreground(bar, tty): + with patch("code_puppy.callbacks.on_prompt_text_color", return_value="#6a9955"): + bar.start() + drain(tty) + bar.set_prompt_text("> ", "puppies", len("puppies")) + + output = written(tty) + assert "\x1b[38;2;106;153;85m" in output + assert "\x1b[39m" in output + + # ========================================================================= # Non-TTY: silent no-ops # ========================================================================= diff --git a/tests/plugins/test_theme_plugin.py b/tests/plugins/test_theme_plugin.py index a5ea20422..ace9eeca6 100644 --- a/tests/plugins/test_theme_plugin.py +++ b/tests/plugins/test_theme_plugin.py @@ -28,13 +28,22 @@ DEEP_BLACK, FOREST, GITHUB_LIGHT, + GREEN_SCREEN, OCEAN, + PURPLE_PUPPY, ROSE_PINE_DAWN, SOLARIZED_LIGHT, SUNSET, TOKYO_NIGHT, VAPORWAVE, ) +from code_puppy.plugins.theme.picker import ( + THEMES_PER_PAGE, + _format_menu, + _move_page, + _page_for_index, + _total_pages, +) from code_puppy.plugins.theme.rich_themes import make_remap, _swap_color, _safe_parse from code_puppy.plugins.theme.content_styles import ( CONTENT_KEYS, @@ -59,11 +68,11 @@ # --------------------------------------------------------------------------- class TestThemeCatalog: def test_curated_themes_count(self): - assert len(CURATED_THEMES) == 13 + assert len(CURATED_THEMES) == 14 def test_menu_has_expected_entries(self): names = [name for name, _ in MENU] - assert len(names) == 15 + assert len(names) == 16 assert "ocean" in names assert "forest" in names assert "sunset" in names @@ -72,6 +81,7 @@ def test_menu_has_expected_entries(self): assert "purple-puppy" in names assert "catppuccin-mocha" in names assert "tokyo-night" in names + assert "green-screen" in names assert "deep-black" in names assert "solarized-light" in names assert "github-light" in names @@ -84,7 +94,8 @@ def test_menu_by_index_maps_strings(self): assert MENU_BY_INDEX["5"] == "bubblegum-pink" assert MENU_BY_INDEX["6"] == "purple-puppy" assert MENU_BY_INDEX["7"] == "catppuccin-mocha" - assert MENU_BY_INDEX["11"] == "solarized-light" + assert MENU_BY_INDEX["10"] == "green-screen" + assert MENU_BY_INDEX["12"] == "solarized-light" assert MENU_BY_INDEX[str(len(MENU))] == "default" def test_aliases_resolve(self): @@ -94,6 +105,8 @@ def test_aliases_resolve(self): assert MENU_BY_NAME["puppy"] is CURATED_THEMES["purple-puppy"] assert MENU_BY_NAME["purple"] is CURATED_THEMES["purple-puppy"] assert MENU_BY_NAME["tokyo"] is CURATED_THEMES["tokyo-night"] + assert MENU_BY_NAME["green"] is CURATED_THEMES["green-screen"] + assert MENU_BY_NAME["crt"] is CURATED_THEMES["green-screen"] assert MENU_BY_NAME["solarized"] is CURATED_THEMES["solarized-light"] assert MENU_BY_NAME["github"] is CURATED_THEMES["github-light"] assert MENU_BY_NAME["rose-pine"] is CURATED_THEMES["rose-pine-dawn"] @@ -248,12 +261,59 @@ def test_alias_keys_are_accepted(self): assert resolve_theme_arg("bubblegum") is not None assert resolve_theme_arg("pink") is not None assert resolve_theme_arg("tokyo") is not None + assert resolve_theme_arg("green") is not None + assert resolve_theme_arg("crt") is not None assert resolve_theme_arg("gruvbox") is None +# --------------------------------------------------------------------------- +# picker.py +# --------------------------------------------------------------------------- +class TestThemePickerPagination: + def test_catalog_is_split_into_pages(self): + assert THEMES_PER_PAGE == 5 + assert _total_pages() == 4 + assert _page_for_index(0) == 0 + assert _page_for_index(5) == 1 + assert _page_for_index(len(MENU) - 1) == 3 + + def test_menu_only_renders_the_selected_page(self): + rendered = "".join(text for _, text in _format_menu(5)) + + assert "Page 2/4" in rendered + assert "6. " in rendered + assert "10. " in rendered + assert "Purple Puppy" in rendered + assert "Green Screen" in rendered + assert "Ocean" not in rendered + assert "Deep Black" not in rendered + + def test_page_navigation_clamps_at_catalog_edges(self): + assert _move_page(2, 1) == 7 + assert _move_page(7, -1) == 2 + assert _move_page(0, -1) == 0 + assert _move_page(len(MENU) - 2, 1) == len(MENU) - 1 + + # --------------------------------------------------------------------------- # bundled_palettes.py # --------------------------------------------------------------------------- +def _relative_luminance(color: str) -> float: + channels = [int(color[index : index + 2], 16) / 255 for index in (1, 3, 5)] + linear = [ + value / 12.92 if value <= 0.04045 else ((value + 0.055) / 1.055) ** 2.4 + for value in channels + ] + return 0.2126 * linear[0] + 0.7152 * linear[1] + 0.0722 * linear[2] + + +def _contrast_ratio(first: str, second: str) -> float: + lighter, darker = sorted( + (_relative_luminance(first), _relative_luminance(second)), reverse=True + ) + return (lighter + 0.05) / (darker + 0.05) + + class TestBundledPalettes: @pytest.mark.parametrize( "palette", @@ -263,6 +323,8 @@ class TestBundledPalettes: SUNSET, VAPORWAVE, BUBBLEGUM_PINK, + GREEN_SCREEN, + PURPLE_PUPPY, CATPPUCCIN_MOCHA, CATPPUCCIN_LATTE, TOKYO_NIGHT, @@ -286,6 +348,8 @@ def test_palette_structure(self, palette): SUNSET, VAPORWAVE, BUBBLEGUM_PINK, + GREEN_SCREEN, + PURPLE_PUPPY, CATPPUCCIN_MOCHA, CATPPUCCIN_LATTE, TOKYO_NIGHT, @@ -301,6 +365,24 @@ def test_palette_hex_format(self, palette): for color in palette["ansi"]: assert color.startswith("#"), f"Bad ANSI color: {color}" + def test_green_screen_uses_llxprt_phosphor_colors(self): + assert GREEN_SCREEN["bg"] == "#000000" + assert GREEN_SCREEN["fg"] == "#6a9955" + assert "#00ff00" in GREEN_SCREEN["ansi"] + assert terminal_palette_for("green-screen") is GREEN_SCREEN + + def test_green_screen_banner_labels_have_accessible_contrast(self): + banner_colors = set(colors_for("green-screen").values()) + + assert banner_colors.isdisjoint({GREEN_SCREEN["bg"], GREEN_SCREEN["fg"]}) + assert all( + _contrast_ratio(color, GREEN_SCREEN["fg"]) >= 4.5 for color in banner_colors + ) + + def test_purple_puppy_muted_text_has_accessible_contrast(self): + """ANSI bright black is muted TUI text and must remain readable.""" + assert _contrast_ratio(PURPLE_PUPPY["bg"], PURPLE_PUPPY["ansi"][8]) >= 4.5 + # --------------------------------------------------------------------------- # rich_themes.py @@ -424,6 +506,135 @@ def test_get_saved_palette_returns_dict(self): # register_callbacks.py # --------------------------------------------------------------------------- class TestRegisterCallbacks: + def test_prompt_text_uses_active_terminal_foreground(self): + from code_puppy.plugins.theme.register_callbacks import _prompt_text_color + + with patch( + "code_puppy.plugins.theme.register_callbacks._active_terminal_palette", + return_value=("green-screen", GREEN_SCREEN), + ): + assert _prompt_text_color(None) == "#6a9955" + + def test_green_screen_highlighter_removes_monokai_white(self): + from termflow.syntax import Highlighter + + from code_puppy.plugins.theme.register_callbacks import _termflow_highlighter + + with patch( + "code_puppy.plugins.theme.register_callbacks._active_terminal_palette", + return_value=("green-screen", GREEN_SCREEN), + ): + highlighter = _termflow_highlighter(Highlighter()) + + rendered = highlighter.highlight_line("plain code", "text") + assert "38;2;106;153;85" in rendered + assert "38;2;255;255;255" not in rendered + + def test_termflow_style_uses_active_terminal_palette(self): + from termflow.render.style import RenderStyle + + from code_puppy.plugins.theme.register_callbacks import _termflow_style + + default = RenderStyle.default() + with ( + patch( + "code_puppy.plugins.theme.register_callbacks.get_value", + return_value="green-screen", + ), + patch("code_puppy.plugins.theme.register_callbacks.osc") as mock_osc, + ): + mock_osc.get_saved_palette.return_value = GREEN_SCREEN + style = _termflow_style(default) + + assert style is not default + assert style.bright == "#6a9955" + assert style.head == "#00ff00" + assert style.symbol == "#6a9955" + assert style.dark == "#000000" + assert style.link == "#6a9955" + assert style.error == "#6a9955" + + def test_termflow_style_preserves_default_without_active_theme(self): + from termflow.render.style import RenderStyle + + from code_puppy.plugins.theme.register_callbacks import _termflow_style + + default = RenderStyle.default() + with patch( + "code_puppy.plugins.theme.register_callbacks.get_value", + return_value="default", + ): + assert _termflow_style(default) is default + + def test_first_run_applies_tokyo_night(self): + from code_puppy.plugins.theme.register_callbacks import ( + _apply_default_theme_on_first_run, + ) + + with ( + patch( + "code_puppy.plugins.theme.register_callbacks.get_value", + return_value=None, + ), + patch("code_puppy.plugins.theme.register_callbacks.apply") as mock_apply, + patch("code_puppy.plugins.theme.register_callbacks.cs") as mock_cs, + patch("code_puppy.plugins.theme.register_callbacks.rt") as mock_rt, + patch("code_puppy.plugins.theme.register_callbacks.osc") as mock_osc, + patch( + "code_puppy.plugins.theme.register_callbacks.set_config_value" + ) as mock_set, + ): + mock_osc.get_saved_palette.return_value = None + _apply_default_theme_on_first_run() + + mock_apply.assert_called_once_with(colors_for("tokyo-night")) + mock_cs.apply_content_styles.assert_called_once_with( + content_styles_for("tokyo-night") + ) + mock_rt.apply_remap.assert_called_once_with(color_remap_for("tokyo-night")) + mock_osc.apply_palette.assert_called_once_with( + terminal_palette_for("tokyo-night") + ) + mock_set.assert_called_once_with("theme_active_theme", "tokyo-night") + + def test_default_theme_preserves_saved_choice(self): + from code_puppy.plugins.theme.register_callbacks import ( + _apply_default_theme_on_first_run, + ) + + with ( + patch( + "code_puppy.plugins.theme.register_callbacks.get_value", + return_value="purple-puppy", + ), + patch("code_puppy.plugins.theme.register_callbacks.apply") as mock_apply, + ): + _apply_default_theme_on_first_run() + + mock_apply.assert_not_called() + + def test_default_theme_preserves_legacy_palette(self): + from code_puppy.plugins.theme.register_callbacks import ( + _apply_default_theme_on_first_run, + ) + + with ( + patch( + "code_puppy.plugins.theme.register_callbacks.get_value", + return_value=None, + ), + patch("code_puppy.plugins.theme.register_callbacks.osc") as mock_osc, + patch("code_puppy.plugins.theme.register_callbacks.apply") as mock_apply, + patch( + "code_puppy.plugins.theme.register_callbacks.set_config_value" + ) as mock_set, + ): + mock_osc.get_saved_palette.return_value = {"bg": "#123456"} + _apply_default_theme_on_first_run() + + mock_apply.assert_not_called() + mock_set.assert_called_once_with("theme_active_theme", "legacy-custom") + def test_custom_help_returns_theme_entry(self): from code_puppy.plugins.theme.register_callbacks import _custom_help @@ -463,6 +674,7 @@ def test_handle_theme_by_name_applies(self): patch("code_puppy.plugins.theme.register_callbacks.cs"), patch("code_puppy.plugins.theme.register_callbacks.rt"), patch("code_puppy.plugins.theme.register_callbacks.osc"), + patch("code_puppy.plugins.theme.register_callbacks.set_config_value"), patch("code_puppy.plugins.theme.register_callbacks.emit_info"), ): result = _handle_theme("/theme ocean", "theme") diff --git a/tests/test_callbacks_extended.py b/tests/test_callbacks_extended.py index b7d836bc3..9ff3a3f4d 100644 --- a/tests/test_callbacks_extended.py +++ b/tests/test_callbacks_extended.py @@ -15,11 +15,15 @@ on_edit_file, on_load_model_config, on_post_tool_call, + on_prompt_text_color, + on_shutdown, on_pre_tool_call, on_register_cli_args, on_replace_in_file, on_startup, on_stream_event, + on_termflow_highlighter, + on_termflow_style, register_callback, unregister_callback, ) @@ -149,6 +153,40 @@ def test_callback(): assert len(results) == 1 assert results[0] == "test_result" + def test_theme_value_callbacks_are_chained(self): + register_callback("prompt_text_color", lambda _color: "#123456") + register_callback("termflow_highlighter", lambda value: value + "-themed") + + assert on_prompt_text_color() == "#123456" + assert on_termflow_highlighter("default") == "default-themed" + + def test_termflow_style_callbacks_are_chained(self): + register_callback("termflow_style", lambda style: style + "-first") + register_callback("termflow_style", lambda style: style + "-second") + + assert on_termflow_style("default") == "default-first-second" + + def test_termflow_style_ignores_none_results(self): + register_callback("termflow_style", lambda _style: None) + + assert on_termflow_style("default") == "default" + + @pytest.mark.asyncio + async def test_shutdown_executes_sync_and_async_callbacks(self): + """Shutdown supports both callback styles through its async contract.""" + + def sync_callback(): + return "sync_result" + + async def async_callback(): + await asyncio.sleep(0) + return "async_result" + + register_callback("shutdown", sync_callback) + register_callback("shutdown", async_callback) + + assert await on_shutdown() == ["sync_result", "async_result"] + @pytest.mark.asyncio async def test_execute_multiple_callbacks_async(self): """Test executing multiple async callbacks.""" From 322cae765ae639f9dffdf530b5c16b94741486df Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Fri, 10 Jul 2026 01:29:23 +0000 Subject: [PATCH 02/78] chore: bump version [ci skip] --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c020c3847..55e091da9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "code-puppy" -version = "0.0.621" +version = "0.0.622" description = "Code generation agent" readme = "README.md" requires-python = ">=3.11,<3.15" diff --git a/uv.lock b/uv.lock index b9e1f6db3..5c13de52c 100644 --- a/uv.lock +++ b/uv.lock @@ -312,7 +312,7 @@ wheels = [ [[package]] name = "code-puppy" -version = "0.0.621" +version = "0.0.622" source = { editable = "." } dependencies = [ { name = "agent-client-protocol" }, From df73f6ecc82831e8669406b8baf2d298dbbc1afa Mon Sep 17 00:00:00 2001 From: Mike Pfaffenberger Date: Fri, 10 Jul 2026 08:01:52 -0400 Subject: [PATCH 03/78] refactor(cli): simplify settings command surfaces --- code_puppy/cli_runner.py | 9 +- code_puppy/command_line/config_apply.py | 13 +++ code_puppy/command_line/config_commands.py | 80 -------------- .../command_line/prompt_toolkit_completion.py | 59 ++++++---- code_puppy/command_line/set_menu.py | 7 +- code_puppy/command_line/set_menu_catalog.py | 18 --- .../test_config_commands_extended.py | 104 ------------------ .../test_config_commands_full_coverage.py | 71 ------------ .../test_prompt_toolkit_coverage.py | 27 +++-- tests/command_line/test_set_menu.py | 18 +++ tests/test_cli_runner_full_coverage.py | 28 +++++ tests/test_command_handler.py | 53 +-------- tests/test_prompt_toolkit_completion.py | 60 +++++++++- 13 files changed, 190 insertions(+), 357 deletions(-) diff --git a/code_puppy/cli_runner.py b/code_puppy/cli_runner.py index 9fb6d170d..52a7fe2a2 100644 --- a/code_puppy/cli_runner.py +++ b/code_puppy/cli_runner.py @@ -667,7 +667,7 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non emit_system_message("Type 'clear' to reset the conversation history.") emit_system_message("Type /help to view all commands") emit_system_message( - "Type @ for path completion, or /model to pick a model. Toggle multiline with Alt+M or F2; newline: Ctrl+J." + "Type @ for path completion, or /model to pick a model. Toggle multiline with Alt+M or F2; newline: Shift+Enter." ) emit_system_message("Paste images: Ctrl+V (even on Mac!), F3, or /paste command.") import platform @@ -678,7 +678,12 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non ) cancel_key = get_cancel_agent_display_name() emit_system_message( - f"Press {cancel_key} during processing to cancel the current task or inference. Use Ctrl+X to interrupt running shell commands." + f"Press {cancel_key} during processing to cancel the current task or inference." + ) + emit_system_message( + "Use Ctrl+X Ctrl+E to open $EDITOR (Notepad on Windows); " + "Ctrl+X Ctrl+B to background running shell commands; " + "Ctrl+X Ctrl+X to kill running shell commands." ) emit_system_message( "Use /autosave_load to manually load a previous autosave session." diff --git a/code_puppy/command_line/config_apply.py b/code_puppy/command_line/config_apply.py index b1fd397c5..21227d1fa 100644 --- a/code_puppy/command_line/config_apply.py +++ b/code_puppy/command_line/config_apply.py @@ -14,6 +14,14 @@ from typing import Optional +MODEL_SETTINGS_ONLY_KEYS = frozenset( + { + "openai_reasoning_effort", + "openai_verbosity", + } +) + + @dataclass(frozen=True) class ApplyResult: """Outcome of writing a single config key/value pair. @@ -85,6 +93,11 @@ def apply_setting( if not key: return ApplyResult(ok=False, error="You must supply a key.") + if key in MODEL_SETTINGS_ONLY_KEYS: + return ApplyResult( + ok=False, + error=(f"'{key}' is managed per model. Use /model_settings to change it."), + ) warning: Optional[str] = None requires_restart = False diff --git a/code_puppy/command_line/config_commands.py b/code_puppy/command_line/config_commands.py index 1777b27f8..58c6e1b07 100644 --- a/code_puppy/command_line/config_commands.py +++ b/code_puppy/command_line/config_commands.py @@ -90,86 +90,6 @@ def handle_show_command(command: str) -> bool: return True -@register_command( - name="reasoning", - description="Set OpenAI reasoning effort for GPT-5 models (e.g., /reasoning high)", - usage="/reasoning ", - category="config", -) -def handle_reasoning_command(command: str) -> bool: - """Set OpenAI reasoning effort level.""" - from code_puppy.messaging import emit_error, emit_success, emit_warning - - tokens = command.split() - if len(tokens) != 2: - emit_warning("Usage: /reasoning ") - return True - - effort = tokens[1] - try: - from code_puppy.config import set_openai_reasoning_effort - - set_openai_reasoning_effort(effort) - except ValueError as exc: - emit_error(str(exc)) - return True - - from code_puppy.config import get_openai_reasoning_effort - - normalized_effort = get_openai_reasoning_effort() - - from code_puppy.agents.agent_manager import get_current_agent - - agent = get_current_agent() - agent.reload_code_generation_agent() - emit_success( - f"Reasoning effort set to '{normalized_effort}' and active agent reloaded" - ) - return True - - -@register_command( - name="verbosity", - description="Set OpenAI verbosity for GPT-5 models (e.g., /verbosity high)", - usage="/verbosity ", - category="config", -) -def handle_verbosity_command(command: str) -> bool: - """Set OpenAI verbosity level. - - Controls how concise vs. verbose the model's responses are: - - low: more concise responses - - medium: balanced (default) - - high: more verbose responses - """ - from code_puppy.messaging import emit_error, emit_success, emit_warning - - tokens = command.split() - if len(tokens) != 2: - emit_warning("Usage: /verbosity ") - return True - - verbosity = tokens[1] - try: - from code_puppy.config import set_openai_verbosity - - set_openai_verbosity(verbosity) - except ValueError as exc: - emit_error(str(exc)) - return True - - from code_puppy.config import get_openai_verbosity - - normalized_verbosity = get_openai_verbosity() - - from code_puppy.agents.agent_manager import get_current_agent - - agent = get_current_agent() - agent.reload_code_generation_agent() - emit_success(f"Verbosity set to '{normalized_verbosity}' and active agent reloaded") - return True - - @register_command( name="set", description="Set puppy config (e.g., /set yolo_mode true) or launch interactive menu", diff --git a/code_puppy/command_line/prompt_toolkit_completion.py b/code_puppy/command_line/prompt_toolkit_completion.py index dbff51104..2e4ae3ff7 100644 --- a/code_puppy/command_line/prompt_toolkit_completion.py +++ b/code_puppy/command_line/prompt_toolkit_completion.py @@ -149,11 +149,14 @@ def get_completions(self, document, complete_event): # Don't return any completions -- let ModelNameCompleter handle it return - # Get config keys and sort them alphabetically for consistent display + # Get config keys and sort them alphabetically for consistent display. + # Per-model controls belong exclusively to /model_settings. + from code_puppy.command_line.config_apply import MODEL_SETTINGS_ONLY_KEYS + config_keys = sorted(get_config_keys()) for key in config_keys: - if key == "model" or key == "puppy_token": + if key in {"model", "puppy_token"} | MODEL_SETTINGS_ONLY_KEYS: continue # exclude 'model' and 'puppy_token' from regular /set completions if key.startswith(text_after_trigger): prev_value = get_value(key) @@ -410,25 +413,10 @@ def get_completions(self, document, complete_event): if not stripped_text.startswith("/"): return - # Get the text after the initial slash - if len(stripped_text) == 1: - # Bare '/': only show the full command menu on an EXPLICIT - # request (Tab). Auto-popping ~30 commands here forces the - # terminal to scroll for menu space, and when the user was - # just typo-ing a slash (type '/', erase it) that scroll - # leaves permanent ghost blank lines — terminals can't - # un-scroll. Typing any character after '/' still auto-opens - # the (filtered) menu as usual. getattr: defensive against - # callers passing None as the event. - if not getattr(complete_event, "completion_requested", False): - return - # User hit Tab on '/', show all commands - partial = "" - start_position = 0 # Don't replace anything, just insert at cursor - else: - # User is typing a command after the slash - partial = stripped_text[1:] # text after '/' - start_position = -(len(partial)) # Replace what was typed after '/' + # Get the text after the initial slash. A bare slash intentionally + # yields every command so the menu appears immediately while typing. + partial = stripped_text[1:] + start_position = -len(partial) # Load all available commands try: @@ -677,6 +665,29 @@ def _left_justify_completion_menu(session: PromptSession) -> None: pass +def _complete_or_cycle(buffer) -> None: + """Apply an unambiguous completion, otherwise cycle the available choices. + + ``complete_while_typing`` opens a completion state without selecting an + item. prompt_toolkit's default Tab binding selects even a sole candidate, + requiring a pointless second Tab to apply it. Start completion explicitly + when needed, then immediately apply a single candidate while preserving + normal cycling for ambiguous input. + """ + if buffer.complete_state is None: + buffer.start_completion(select_first=False) + + complete_state = buffer.complete_state + if complete_state is None: + return + + completions = complete_state.completions + if len(completions) == 1: + buffer.apply_completion(completions[0]) + elif completions: + buffer.complete_next() + + async def get_input_with_combined_completion( prompt_str=">>> ", history_file: Optional[str] = None ) -> str: @@ -789,6 +800,12 @@ def _(event): else: buffer.validate_and_handle() + # Tab: finish an unambiguous completion in one press. For multiple + # candidates, retain prompt_toolkit's familiar select/cycle behavior. + @bindings.add(Keys.Tab, eager=True) + def handle_tab_completion(event): + _complete_or_cycle(event.app.current_buffer) + # Backspace/Delete: trigger completions after deletion. # By default, complete_while_typing only triggers on character insertion, # not deletion — so the menu vanishes the moment you backspace. We diff --git a/code_puppy/command_line/set_menu.py b/code_puppy/command_line/set_menu.py index ca206a52b..d42dbe491 100644 --- a/code_puppy/command_line/set_menu.py +++ b/code_puppy/command_line/set_menu.py @@ -23,7 +23,10 @@ from prompt_toolkit.layout.controls import FormattedTextControl from prompt_toolkit.widgets import Frame -from code_puppy.command_line.config_apply import apply_setting +from code_puppy.command_line.config_apply import ( + MODEL_SETTINGS_ONLY_KEYS, + apply_setting, +) from code_puppy.command_line.pagination import ( ensure_visible_page, get_page_bounds, @@ -118,7 +121,7 @@ def _build_entries() -> List[_Entry]: curated_keys.add(setting.key) for key in get_config_keys(): - if key in curated_keys: + if key in curated_keys or key in MODEL_SETTINGS_ONLY_KEYS: continue entries.append( _Entry( diff --git a/code_puppy/command_line/set_menu_catalog.py b/code_puppy/command_line/set_menu_catalog.py index 2e4fca28c..a90d731bf 100644 --- a/code_puppy/command_line/set_menu_catalog.py +++ b/code_puppy/command_line/set_menu_catalog.py @@ -43,9 +43,7 @@ get_mcp_disabled, get_mcp_unbound_warning_silenced, get_message_limit, - get_openai_reasoning_effort, get_openai_reasoning_summary, - get_openai_verbosity, get_output_level, get_owner_name, get_pack_agents_enabled, @@ -266,14 +264,6 @@ _OPENAI = SettingsCategory( name="OpenAI", settings=( - Setting( - key="openai_reasoning_effort", - display_name="Reasoning Effort", - description="How much reasoning effort GPT-5 models should use.", - type_hint="choice", - valid_values=("minimal", "low", "medium", "high", "xhigh"), - effective_getter=get_openai_reasoning_effort, - ), Setting( key="openai_reasoning_summary", display_name="Reasoning Summary", @@ -282,14 +272,6 @@ valid_values=("auto", "concise", "detailed"), effective_getter=get_openai_reasoning_summary, ), - Setting( - key="openai_verbosity", - display_name="Verbosity", - description="How verbose GPT-5 model responses should be.", - type_hint="choice", - valid_values=("low", "medium", "high"), - effective_getter=get_openai_verbosity, - ), ), ) diff --git a/tests/command_line/test_config_commands_extended.py b/tests/command_line/test_config_commands_extended.py index 32b4f0b02..b26b9be7b 100644 --- a/tests/command_line/test_config_commands_extended.py +++ b/tests/command_line/test_config_commands_extended.py @@ -2,7 +2,6 @@ This module provides comprehensive coverage for configuration commands including: - Pin/unpin model commands for both JSON and built-in agents -- Reasoning effort configuration commands - Diff configuration commands and color settings - Set configuration commands with validation - Show color options utility @@ -23,7 +22,6 @@ from code_puppy.command_line.config_commands import ( handle_diff_command, handle_pin_model_command, - handle_reasoning_command, handle_set_command, handle_unpin_command, ) @@ -63,108 +61,6 @@ def _show_color_options(diff_type): emit_info("Available diff types: additions, deletions") -class TestReasoningCommand: - """Extended tests for reasoning command functionality.""" - - def test_reasoning_command_valid_efforts(self): - """Test reasoning command with valid effort levels.""" - valid_efforts = ["low", "medium", "high"] - - for effort in valid_efforts: - with patch("code_puppy.config.set_openai_reasoning_effort") as mock_set: - with patch( - "code_puppy.config.get_openai_reasoning_effort", - return_value="medium", - ): - with patch( - "code_puppy.agents.agent_manager.get_current_agent" - ) as mock_get_agent: - mock_agent = MagicMock() - mock_agent.reload_code_generation_agent.return_value = None - mock_get_agent.return_value = mock_agent - - result = handle_reasoning_command(f"/reasoning {effort}") - assert result is True - - mock_set.assert_called_once_with(effort) - mock_agent.reload_code_generation_agent.assert_called_once() - - def test_reasoning_command_invalid_effort(self): - """Test reasoning command with invalid effort levels.""" - invalid_efforts = ["invalid", "extra", "none"] - - for effort in invalid_efforts: - expected_error = ( - f"Invalid reasoning effort '{effort}'. Allowed: high, low, medium" - ) - with patch( - "code_puppy.config.set_openai_reasoning_effort", - side_effect=ValueError(expected_error), - ): - with patch("code_puppy.messaging.emit_error") as mock_error: - result = handle_reasoning_command(f"/reasoning {effort}") - assert result is True - - mock_error.assert_called_once_with(expected_error) - - def test_reasoning_command_no_arguments(self): - """Test reasoning command with no arguments.""" - with patch("code_puppy.messaging.emit_warning") as mock_warning: - result = handle_reasoning_command("/reasoning") - assert result is True - - mock_warning.assert_called_once_with( - "Usage: /reasoning " - ) - - def test_reasoning_command_current_none(self): - """Test reasoning command when current effort is None.""" - with patch("code_puppy.messaging.emit_warning") as mock_warning: - result = handle_reasoning_command("/reasoning") - assert result is True - - mock_warning.assert_called_once_with( - "Usage: /reasoning " - ) - - def test_reasoning_command_wrong_argument_count(self): - """Test reasoning command with incorrect number of arguments.""" - with patch("code_puppy.messaging.emit_warning") as mock_warning: - # Test with no arguments - result = handle_reasoning_command("/reasoning") - assert result is True - - # Test with too many arguments - result = handle_reasoning_command("/reasoning high medium") - assert result is True - - assert mock_warning.call_count == 2 - - # Check warning message - args, kwargs = mock_warning.call_args_list[0] - assert "Usage:" in args[0] - assert "" in args[0] - - def test_reasoning_command_agent_reload_failure(self): - """Test reasoning command handles agent reload failures gracefully.""" - with patch("code_puppy.config.set_openai_reasoning_effort"): - with patch( - "code_puppy.config.get_openai_reasoning_effort", return_value="medium" - ): - with patch( - "code_puppy.agents.agent_manager.get_current_agent" - ) as mock_get_agent: - mock_agent = MagicMock() - mock_agent.reload_code_generation_agent.side_effect = Exception( - "Reload failed" - ) - mock_get_agent.return_value = mock_agent - - # Should propagate the exception - with pytest.raises(Exception, match="Reload failed"): - handle_reasoning_command("/reasoning high") - - class TestSetCommand: """Extended tests for set configuration command.""" diff --git a/tests/command_line/test_config_commands_full_coverage.py b/tests/command_line/test_config_commands_full_coverage.py index 575650788..95733d529 100644 --- a/tests/command_line/test_config_commands_full_coverage.py +++ b/tests/command_line/test_config_commands_full_coverage.py @@ -103,77 +103,6 @@ def test_show_effective_temp_none(self): assert handle_show_command("/show") is True -class TestHandleReasoningCommand: - def test_no_args(self): - from code_puppy.command_line.config_commands import handle_reasoning_command - - with patch("code_puppy.messaging.emit_warning") as warn: - assert handle_reasoning_command("/reasoning") is True - warn.assert_called_once() - - def test_valid(self): - from code_puppy.command_line.config_commands import handle_reasoning_command - - mock_agent = MagicMock() - with ( - patch("code_puppy.config.set_openai_reasoning_effort"), - patch("code_puppy.config.get_openai_reasoning_effort", return_value="high"), - patch( - "code_puppy.agents.agent_manager.get_current_agent", - return_value=mock_agent, - ), - patch("code_puppy.messaging.emit_success"), - ): - assert handle_reasoning_command("/reasoning high") is True - - def test_invalid(self): - from code_puppy.command_line.config_commands import handle_reasoning_command - - with ( - patch( - "code_puppy.config.set_openai_reasoning_effort", - side_effect=ValueError("bad"), - ), - patch("code_puppy.messaging.emit_error") as err, - ): - assert handle_reasoning_command("/reasoning bad") is True - err.assert_called_once() - - -class TestHandleVerbosityCommand: - def test_no_args(self): - from code_puppy.command_line.config_commands import handle_verbosity_command - - with patch("code_puppy.messaging.emit_warning"): - assert handle_verbosity_command("/verbosity") is True - - def test_valid(self): - from code_puppy.command_line.config_commands import handle_verbosity_command - - mock_agent = MagicMock() - with ( - patch("code_puppy.config.set_openai_verbosity"), - patch("code_puppy.config.get_openai_verbosity", return_value="low"), - patch( - "code_puppy.agents.agent_manager.get_current_agent", - return_value=mock_agent, - ), - patch("code_puppy.messaging.emit_success"), - ): - assert handle_verbosity_command("/verbosity low") is True - - def test_invalid(self): - from code_puppy.command_line.config_commands import handle_verbosity_command - - with ( - patch( - "code_puppy.config.set_openai_verbosity", side_effect=ValueError("bad") - ), - patch("code_puppy.messaging.emit_error"), - ): - assert handle_verbosity_command("/verbosity bad") is True - - class TestHandleSetCommand: def test_no_args_launches_menu(self): from code_puppy.command_line.config_commands import handle_set_command diff --git a/tests/command_line/test_prompt_toolkit_coverage.py b/tests/command_line/test_prompt_toolkit_coverage.py index 811a05c3a..30042d8fd 100644 --- a/tests/command_line/test_prompt_toolkit_coverage.py +++ b/tests/command_line/test_prompt_toolkit_coverage.py @@ -307,21 +307,30 @@ def test_just_slash(self): ) assert len(completions) >= 1 - def test_just_slash_without_tab_stays_silent(self): - """Bare '/' must NOT auto-pop the menu while typing. - - Auto-popping forces a terminal scroll for menu space; a typo'd - slash that gets erased would leave permanent ghost blank lines. - """ + def test_just_slash_shows_menu_while_typing(self): + """Bare '/' immediately offers every slash command.""" from prompt_toolkit.completion import CompleteEvent from code_puppy.command_line.prompt_toolkit_completion import SlashCompleter c = SlashCompleter() + mock_cmd = MagicMock() + mock_cmd.name = "help" + mock_cmd.description = "Show help" + mock_cmd.aliases = ["h"] typing_event = CompleteEvent(text_inserted=True) - assert list(c.get_completions(self._make_doc("/"), typing_event)) == [] - # None event (defensive path) is treated as not-requested too. - assert list(c.get_completions(self._make_doc("/"), None)) == [] + + with ( + patch( + "code_puppy.command_line.prompt_toolkit_completion.get_unique_commands", + return_value=[mock_cmd], + ), + patch("code_puppy.plugins.load_plugin_callbacks"), + patch("code_puppy.callbacks.on_custom_command_help", return_value=[]), + ): + for event in (typing_event, None): + completions = list(c.get_completions(self._make_doc("/"), event)) + assert [completion.text for completion in completions] == ["h", "help"] def test_partial_command(self): from code_puppy.command_line.prompt_toolkit_completion import SlashCompleter diff --git a/tests/command_line/test_set_menu.py b/tests/command_line/test_set_menu.py index 2090f8817..072447237 100644 --- a/tests/command_line/test_set_menu.py +++ b/tests/command_line/test_set_menu.py @@ -51,6 +51,14 @@ def test_missing_key_returns_error(self): assert result.ok is False assert result.error and "key" in result.error.lower() + @pytest.mark.parametrize("key", ["openai_reasoning_effort", "openai_verbosity"]) + def test_model_settings_only_keys_are_rejected(self, key): + with patch("code_puppy.config.set_config_value") as mock_set: + result = apply_setting(key, "high") + assert result.ok is False + assert "/model_settings" in (result.error or "") + mock_set.assert_not_called() + def test_cancel_agent_key_invalid_returns_error(self): with patch("code_puppy.config.set_config_value") as mock_set: result = apply_setting("cancel_agent_key", "ctrl+x") @@ -307,6 +315,16 @@ def test_dynamic_keys_land_in_dynamic_category(self): assert match assert match[0].category.name == "Dynamic" + def test_model_settings_only_keys_are_absent(self): + with patch( + "code_puppy.command_line.set_menu.get_config_keys", + return_value=["openai_reasoning_effort", "openai_verbosity"], + ): + entries = _build_entries() + keys = {entry.setting.key for entry in entries} + assert "openai_reasoning_effort" not in keys + assert "openai_verbosity" not in keys + def test_dynamic_does_not_double_curated_keys(self): with patch( "code_puppy.command_line.set_menu.get_config_keys", diff --git a/tests/test_cli_runner_full_coverage.py b/tests/test_cli_runner_full_coverage.py index d20db2bb3..65d6eb3a8 100644 --- a/tests/test_cli_runner_full_coverage.py +++ b/tests/test_cli_runner_full_coverage.py @@ -443,6 +443,34 @@ async def test_exit_command(self): AsyncMock(return_value="/exit"), ) + @pytest.mark.anyio + async def test_startup_instructions_describe_editor_shortcuts(self): + emit_system_message = MagicMock() + + await _run_interactive( + _mock_renderer(), + _interactive_patches(), + AsyncMock(return_value="/exit"), + extra_patches={ + "code_puppy.messaging.emit_system_message": emit_system_message, + }, + ) + + messages = [call.args[0] for call in emit_system_message.call_args_list] + assert any("newline: Shift+Enter" in message for message in messages) + assert any( + "Ctrl+X Ctrl+E to open $EDITOR (Notepad on Windows)" in message + for message in messages + ) + assert any( + "Ctrl+X Ctrl+B to background running shell commands" in message + for message in messages + ) + assert any( + "Ctrl+X Ctrl+X to kill running shell commands" in message + for message in messages + ) + @pytest.mark.anyio async def test_quit_command(self): agent = MagicMock() diff --git a/tests/test_command_handler.py b/tests/test_command_handler.py index dc9c3360e..bbd0dd1f0 100644 --- a/tests/test_command_handler.py +++ b/tests/test_command_handler.py @@ -757,50 +757,6 @@ def test_compact_with_truncation_strategy(self): mock_truncate.assert_called_once() -class TestReasoningCommand: - """Tests for /reasoning command.""" - - def test_reasoning_set_low(self): - """Test /reasoning low sets effort to low.""" - mock_agent = MagicMock() - - with ( - patch("code_puppy.config.set_openai_reasoning_effort") as mock_set, - patch("code_puppy.config.get_openai_reasoning_effort", return_value="low"), - patch( - "code_puppy.agents.agent_manager.get_current_agent", - return_value=mock_agent, - ), - patch("code_puppy.messaging.emit_success") as mock_success, - ): - result = handle_command("/reasoning low") - assert result is True - mock_set.assert_called_once_with("low") - mock_agent.reload_code_generation_agent.assert_called_once() - mock_success.assert_called_once() - - def test_reasoning_invalid_level(self): - """Test /reasoning with invalid level shows error.""" - with ( - patch( - "code_puppy.config.set_openai_reasoning_effort", - side_effect=ValueError("Invalid"), - ), - patch("code_puppy.messaging.emit_error") as mock_error, - ): - result = handle_command("/reasoning invalid") - assert result is True - mock_error.assert_called_once() - - def test_reasoning_no_argument(self): - """Test /reasoning without argument shows usage.""" - with patch("code_puppy.messaging.emit_warning") as mock_warn: - result = handle_command("/reasoning") - assert result is True - mock_warn.assert_called_once() - assert "Usage" in str(mock_warn.call_args) - - class TestTruncateCommand: """Tests for /truncate command.""" @@ -1052,11 +1008,10 @@ def test_compact_command_registered(self): assert cmd is not None assert cmd.category == "session" - def test_reasoning_command_registered(self): - """Test that reasoning command is registered.""" - cmd = get_command("reasoning") - assert cmd is not None - assert cmd.category == "config" + def test_model_controls_are_not_top_level_commands(self): + """Reasoning effort and verbosity live only in /model_settings.""" + assert get_command("reasoning") is None + assert get_command("verbosity") is None def test_truncate_command_registered(self): """Test that truncate command is registered.""" diff --git a/tests/test_prompt_toolkit_completion.py b/tests/test_prompt_toolkit_completion.py index d175513ee..492ea3008 100644 --- a/tests/test_prompt_toolkit_completion.py +++ b/tests/test_prompt_toolkit_completion.py @@ -4,7 +4,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from prompt_toolkit.buffer import Buffer +from prompt_toolkit.buffer import Buffer, CompletionState +from prompt_toolkit.completion import Completion from prompt_toolkit.document import Document from prompt_toolkit.formatted_text import FormattedText from prompt_toolkit.keys import Keys @@ -16,6 +17,7 @@ CDCompleter, FilePathCompleter, SetCompleter, + _complete_or_cycle, get_input_with_combined_completion, ) @@ -163,6 +165,28 @@ def test_set_completer_partial_key(monkeypatch): assert completions[0].display_meta[0][1] == "" +def test_set_completer_excludes_model_settings_only_keys(monkeypatch): + monkeypatch.setattr( + "code_puppy.command_line.prompt_toolkit_completion.get_config_keys", + lambda: [ + "openai_reasoning_effort", + "openai_verbosity", + "temperature", + ], + ) + monkeypatch.setattr( + "code_puppy.command_line.prompt_toolkit_completion.get_value", + lambda key: "high", + ) + completer = SetCompleter() + + completions = list(completer.get_completions(Document(text="/set openai_"), None)) + assert completions == [] + + completions = list(completer.get_completions(Document(text="/set temp"), None)) + assert [completion.text for completion in completions] == ["temperature = high"] + + def test_set_completer_excludes_model_key(monkeypatch): # Ensure 'model' is a config key but SetCompleter doesn't offer it monkeypatch.setattr( @@ -497,6 +521,40 @@ async def test_get_input_with_combined_completion_no_model_update( assert result == raw_input +def _buffer_with_completions(*completions: Completion) -> Buffer: + buffer = Buffer(document=Document("/he", cursor_position=3)) + buffer.complete_state = CompletionState(buffer.document, list(completions)) + return buffer + + +def test_tab_applies_single_completion_immediately(): + buffer = _buffer_with_completions(Completion("help", start_position=-2)) + + _complete_or_cycle(buffer) + + assert buffer.text == "/help" + assert buffer.complete_state is None + + +def test_tab_cycles_when_completion_is_ambiguous(): + buffer = _buffer_with_completions( + Completion("help", start_position=-2), + Completion("hello", start_position=-2), + ) + + _complete_or_cycle(buffer) + + assert buffer.text == "/help" + assert buffer.complete_state is not None + assert buffer.complete_state.current_completion is not None + assert buffer.complete_state.current_completion.text == "help" + + _complete_or_cycle(buffer) + + assert buffer.text == "/hello" + assert buffer.complete_state.current_completion.text == "hello" + + # To test key bindings, we need to inspect the KeyBindings object passed to PromptSession # We can get it from the mock_prompt_session_cls.call_args From a01c85dbe2e8a9beaf737817a7cdd7ac14af5b5e Mon Sep 17 00:00:00 2001 From: Mike Pfaffenberger Date: Fri, 10 Jul 2026 08:01:52 -0400 Subject: [PATCH 04/78] fix(chatgpt): align GPT-5.6 model capabilities --- .../command_line/model_settings_menu.py | 20 ++--- code_puppy/config.py | 6 +- code_puppy/plugins/chatgpt_oauth/config.py | 2 +- .../chatgpt_oauth/register_callbacks.py | 2 +- code_puppy/plugins/chatgpt_oauth/utils.py | 53 +++++------ .../command_line/test_model_settings_menu.py | 5 +- .../test_model_settings_menu_coverage.py | 21 ++++- tests/plugins/test_chatgpt_oauth_coverage.py | 2 +- .../plugins/test_chatgpt_oauth_integration.py | 19 ++-- tests/plugins/test_chatgpt_oauth_utils.py | 88 +++++++++---------- tests/test_config_full_coverage.py | 4 + 11 files changed, 108 insertions(+), 114 deletions(-) diff --git a/code_puppy/command_line/model_settings_menu.py b/code_puppy/command_line/model_settings_menu.py index a9c113a98..186387127 100644 --- a/code_puppy/command_line/model_settings_menu.py +++ b/code_puppy/command_line/model_settings_menu.py @@ -76,7 +76,7 @@ "name": "Reasoning Effort", "description": "Controls how much effort GPT-5 models spend on reasoning. Higher = more thorough but slower.", "type": "choice", - "choices": ["minimal", "low", "medium", "high", "xhigh"], + "choices": ["minimal", "low", "medium", "high", "xhigh", "ultra"], "default": "medium", }, "summary": { @@ -184,8 +184,8 @@ def _get_setting_choices( ) -> List[str]: """Get the available choices for a setting, filtered by model capabilities. - For reasoning_effort, only codex models support 'xhigh' - regular GPT-5.2 - models are capped at 'high'. + Reasoning effort is capability-gated: xhigh is available to codex and + GPT-5.4+ models, while ultra is reserved for GPT-5.6+ variants. Args: setting_key: The setting name (e.g., 'reasoning_effort', 'verbosity') @@ -200,17 +200,15 @@ def _get_setting_choices( base_choices = setting_def.get("choices", []) - # For reasoning_effort, filter 'xhigh' based on model support if setting_key == "reasoning_effort" and model_name: models_config = ModelFactory.load_config() model_config = models_config.get(model_name, {}) - - # Check if model supports xhigh reasoning - supports_xhigh = model_config.get("supports_xhigh_reasoning", False) - - if not supports_xhigh: - # Remove xhigh from choices for non-codex models - return [c for c in base_choices if c != "xhigh"] + unsupported_choices = set() + if not model_config.get("supports_xhigh_reasoning", False): + unsupported_choices.add("xhigh") + if not model_config.get("supports_ultra_reasoning", False): + unsupported_choices.add("ultra") + return [choice for choice in base_choices if choice not in unsupported_choices] return base_choices diff --git a/code_puppy/config.py b/code_puppy/config.py index 62afbc9ea..0d640a96b 100644 --- a/code_puppy/config.py +++ b/code_puppy/config.py @@ -728,8 +728,8 @@ def set_puppy_token(token: str): def get_openai_reasoning_effort() -> str: - """Return the configured OpenAI reasoning effort (minimal, low, medium, high, xhigh).""" - allowed_values = {"minimal", "low", "medium", "high", "xhigh"} + """Return the configured OpenAI reasoning effort.""" + allowed_values = {"minimal", "low", "medium", "high", "xhigh", "ultra"} configured = (get_value("openai_reasoning_effort") or "medium").strip().lower() if configured not in allowed_values: return "medium" @@ -738,7 +738,7 @@ def get_openai_reasoning_effort() -> str: def set_openai_reasoning_effort(value: str) -> None: """Persist the OpenAI reasoning effort ensuring it remains within allowed values.""" - allowed_values = {"minimal", "low", "medium", "high", "xhigh"} + allowed_values = {"minimal", "low", "medium", "high", "xhigh", "ultra"} normalized = (value or "").strip().lower() if normalized not in allowed_values: raise ValueError( diff --git a/code_puppy/plugins/chatgpt_oauth/config.py b/code_puppy/plugins/chatgpt_oauth/config.py index 22bca6355..d61403c3c 100644 --- a/code_puppy/plugins/chatgpt_oauth/config.py +++ b/code_puppy/plugins/chatgpt_oauth/config.py @@ -26,7 +26,7 @@ "default_context_length": 272000, "api_key_env_var": "CHATGPT_OAUTH_API_KEY", # Codex CLI version info (for User-Agent header) - "client_version": "0.72.0", + "client_version": "0.144.1", "originator": "codex_cli_rs", } diff --git a/code_puppy/plugins/chatgpt_oauth/register_callbacks.py b/code_puppy/plugins/chatgpt_oauth/register_callbacks.py index 81104c8df..623d9e775 100644 --- a/code_puppy/plugins/chatgpt_oauth/register_callbacks.py +++ b/code_puppy/plugins/chatgpt_oauth/register_callbacks.py @@ -134,7 +134,7 @@ def _create_chatgpt_oauth_model( # Build headers for ChatGPT Codex API originator = CHATGPT_OAUTH_CONFIG.get("originator", "codex_cli_rs") - client_version = CHATGPT_OAUTH_CONFIG.get("client_version", "0.72.0") + client_version = CHATGPT_OAUTH_CONFIG.get("client_version", "0.144.1") headers = { "ChatGPT-Account-Id": account_id, diff --git a/code_puppy/plugins/chatgpt_oauth/utils.py b/code_puppy/plugins/chatgpt_oauth/utils.py index 5d06df035..56a6a9322 100644 --- a/code_puppy/plugins/chatgpt_oauth/utils.py +++ b/code_puppy/plugins/chatgpt_oauth/utils.py @@ -344,7 +344,6 @@ def exchange_code_for_tokens( # These are the known models that work with ChatGPT OAuth tokens # Based on codex-rs CLI and shell-scripts/codex-call.sh DEFAULT_CODEX_MODELS = [ - "gpt-5.6", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna", @@ -357,24 +356,9 @@ def exchange_code_for_tokens( "gpt-5.2", ] -# Models that MUST always be registered, even if the /models endpoint -# doesn't return them (e.g. newly launched, not yet in the API catalogue). -# These are merged into whatever the endpoint returns. -REQUIRED_CODEX_MODELS = [ - "gpt-5.6", - "gpt-5.6-sol", - "gpt-5.6-terra", - "gpt-5.6-luna", - "gpt-5.5", - "gpt-5.4", - "gpt-5.3-instant", - "gpt-5.3-codex", -] - # Per-model context length overrides (tokens). # Models not listed here use CHATGPT_OAUTH_CONFIG["default_context_length"] (272,000). CODEX_MODEL_CONTEXT_LENGTHS = { - "gpt-5.6": 1050000, "gpt-5.6-sol": 1050000, "gpt-5.6-terra": 1050000, "gpt-5.6-luna": 1050000, @@ -391,16 +375,9 @@ def _supports_xhigh_reasoning(model_name: str) -> bool: ) -def _ensure_required_models(models: List[str]) -> List[str]: - """Merge REQUIRED_CODEX_MODELS into the given list, preserving order. - - Any required model not already present is prepended so it appears first. - """ - existing = set(models) - missing = [m for m in REQUIRED_CODEX_MODELS if m not in existing] - if missing: - logger.info("Injecting required models not returned by API: %s", missing) - return missing + models +def _supports_ultra_reasoning(model_name: str) -> bool: + """Return whether a ChatGPT OAuth model supports ultra reasoning.""" + return model_name.lower().startswith("gpt-5.6") def fetch_chatgpt_models(access_token: str, account_id: str) -> Optional[List[str]]: @@ -419,7 +396,7 @@ def fetch_chatgpt_models(access_token: str, account_id: str) -> Optional[List[st import platform # Build the models URL with client version - client_version = CHATGPT_OAUTH_CONFIG.get("client_version", "0.72.0") + client_version = CHATGPT_OAUTH_CONFIG.get("client_version", "0.144.1") base_url = CHATGPT_OAUTH_CONFIG["api_base_url"].rstrip("/") models_url = f"{base_url}/models" @@ -456,16 +433,18 @@ def fetch_chatgpt_models(access_token: str, account_id: str) -> Optional[List[st # The response has a "models" key with list of model objects if "models" in data and isinstance(data["models"], list): models = [] + seen_models = set() for model in data["models"]: if model is None: continue model_id = ( model.get("slug") or model.get("id") or model.get("name") ) - if model_id: + if model_id and model_id not in seen_models: + seen_models.add(model_id) models.append(model_id) if models: - return _ensure_required_models(models) + return models except (json.JSONDecodeError, ValueError) as exc: logger.warning("Failed to parse models response: %s", exc) @@ -491,6 +470,15 @@ def add_models_to_extra_config(models: List[str]) -> bool: """Add ChatGPT models to chatgpt_models.json configuration.""" try: chatgpt_models = load_chatgpt_models() + desired_model_keys = { + f"{CHATGPT_OAUTH_CONFIG['prefix']}{model_name}" for model_name in models + } + chatgpt_models = { + key: config + for key, config in chatgpt_models.items() + if config.get("oauth_source") != "chatgpt-oauth-plugin" + or key in desired_model_keys + } added = 0 for model_name in models: prefixed = f"{CHATGPT_OAUTH_CONFIG['prefix']}{model_name}" @@ -500,9 +488,11 @@ def add_models_to_extra_config(models: List[str]) -> bool: # reasoning effort, reasoning summaries, and text verbosity. supported_settings = ["reasoning_effort", "summary", "verbosity"] - # xhigh reasoning is supported by codex models and GPT-5.4+ variants. - # Older non-codex GPT-5.x models like gpt-5.2 stay capped at "high". + # xhigh is supported by codex and GPT-5.4+ variants; ultra is + # narrower and starts with GPT-5.6. Keep separate capabilities so + # adding a stronger effort never leaks into older model menus. supports_xhigh_reasoning = _supports_xhigh_reasoning(model_name) + supports_ultra_reasoning = _supports_ultra_reasoning(model_name) chatgpt_models[prefixed] = { "type": "chatgpt_oauth", @@ -517,6 +507,7 @@ def add_models_to_extra_config(models: List[str]) -> bool: "oauth_source": "chatgpt-oauth-plugin", "supported_settings": supported_settings, "supports_xhigh_reasoning": supports_xhigh_reasoning, + "supports_ultra_reasoning": supports_ultra_reasoning, } added += 1 if save_chatgpt_models(chatgpt_models): diff --git a/tests/command_line/test_model_settings_menu.py b/tests/command_line/test_model_settings_menu.py index a1013605f..98ba3947a 100644 --- a/tests/command_line/test_model_settings_menu.py +++ b/tests/command_line/test_model_settings_menu.py @@ -80,6 +80,7 @@ def test_reasoning_effort_setting_definition(self): assert reason_def["type"] == "choice" assert "minimal" in reason_def["choices"] assert "high" in reason_def["choices"] + assert "ultra" in reason_def["choices"] def test_summary_setting_definition(self): """Test reasoning summary setting has correct choices.""" @@ -270,7 +271,7 @@ def test_reasoning_effort_valid_choices(self): """Test reasoning effort accepts valid choices.""" reason_def = SETTING_DEFINITIONS["reasoning_effort"] valid_choices = reason_def["choices"] - for choice in ["minimal", "low", "medium", "high", "xhigh"]: + for choice in ["minimal", "low", "medium", "high", "xhigh", "ultra"]: if choice in valid_choices: assert True break @@ -278,7 +279,7 @@ def test_reasoning_effort_valid_choices(self): def test_reasoning_effort_rejects_invalid_choice(self): """Test reasoning effort rejects invalid choice.""" reason_def = SETTING_DEFINITIONS["reasoning_effort"] - invalid_choice = "ultra" + invalid_choice = "ludicrous-speed" assert invalid_choice not in reason_def["choices"] def test_verbosity_valid_choices(self): diff --git a/tests/command_line/test_model_settings_menu_coverage.py b/tests/command_line/test_model_settings_menu_coverage.py index 8e1fa7763..6ba58785f 100644 --- a/tests/command_line/test_model_settings_menu_coverage.py +++ b/tests/command_line/test_model_settings_menu_coverage.py @@ -55,12 +55,28 @@ def test_reasoning_effort_without_xhigh(self, mock_factory): assert "high" in choices @patch("code_puppy.command_line.model_settings_menu.ModelFactory") - def test_reasoning_effort_with_xhigh(self, mock_factory): + def test_reasoning_effort_with_xhigh_but_without_ultra(self, mock_factory): mock_factory.load_config.return_value = { - "codex": {"supports_xhigh_reasoning": True} + "codex": { + "supports_xhigh_reasoning": True, + "supports_ultra_reasoning": False, + } } choices = _get_setting_choices("reasoning_effort", "codex") assert "xhigh" in choices + assert "ultra" not in choices + + @patch("code_puppy.command_line.model_settings_menu.ModelFactory") + def test_reasoning_effort_with_ultra(self, mock_factory): + mock_factory.load_config.return_value = { + "gpt-5.6-sol": { + "supports_xhigh_reasoning": True, + "supports_ultra_reasoning": True, + } + } + choices = _get_setting_choices("reasoning_effort", "gpt-5.6-sol") + assert "xhigh" in choices + assert "ultra" in choices def test_non_choice_setting(self): choices = _get_setting_choices("temperature") @@ -70,6 +86,7 @@ def test_non_choice_setting(self): def test_reasoning_effort_no_model_name(self, mock_factory): choices = _get_setting_choices("reasoning_effort") assert "xhigh" in choices # no filtering without model + assert "ultra" in choices # --------------- ModelSettingsMenu properties --------------- diff --git a/tests/plugins/test_chatgpt_oauth_coverage.py b/tests/plugins/test_chatgpt_oauth_coverage.py index 5fb95ef3a..2d9c9a297 100644 --- a/tests/plugins/test_chatgpt_oauth_coverage.py +++ b/tests/plugins/test_chatgpt_oauth_coverage.py @@ -231,7 +231,7 @@ def test_success(self): "code_puppy.plugins.chatgpt_oauth.register_callbacks.CHATGPT_OAUTH_CONFIG", { "originator": "codex_cli_rs", - "client_version": "0.72.0", + "client_version": "0.144.1", "api_base_url": "https://api.example.com", }, ), diff --git a/tests/plugins/test_chatgpt_oauth_integration.py b/tests/plugins/test_chatgpt_oauth_integration.py index 2c6e22974..898a0ed9b 100644 --- a/tests/plugins/test_chatgpt_oauth_integration.py +++ b/tests/plugins/test_chatgpt_oauth_integration.py @@ -88,7 +88,6 @@ def test_add_models_to_extra_config(self, mock_path, tmp_path): result = add_models_to_extra_config( [ - "gpt-5.6", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna", @@ -100,19 +99,20 @@ def test_add_models_to_extra_config(self, mock_path, tmp_path): assert result is True models = json.loads(models_file.read_text()) - assert "chatgpt-gpt-5.6" in models + assert "chatgpt-gpt-5.6" not in models assert "chatgpt-gpt-5.6-sol" in models assert "chatgpt-gpt-5.6-terra" in models assert "chatgpt-gpt-5.6-luna" in models assert "chatgpt-gpt-5.5" in models assert "chatgpt-gpt-5.2" in models assert "chatgpt-gpt-5.2-codex" in models - assert models["chatgpt-gpt-5.6"]["context_length"] == 1050000 assert models["chatgpt-gpt-5.6-sol"]["context_length"] == 1050000 assert models["chatgpt-gpt-5.6-terra"]["context_length"] == 1050000 assert models["chatgpt-gpt-5.6-luna"]["context_length"] == 1050000 - assert models["chatgpt-gpt-5.6"]["supports_xhigh_reasoning"] is True + assert models["chatgpt-gpt-5.6-sol"]["supports_xhigh_reasoning"] is True + assert models["chatgpt-gpt-5.6-sol"]["supports_ultra_reasoning"] is True assert models["chatgpt-gpt-5.5"]["supports_xhigh_reasoning"] is True + assert models["chatgpt-gpt-5.5"]["supports_ultra_reasoning"] is False @patch("code_puppy.plugins.chatgpt_oauth.utils.get_chatgpt_models_path") def test_add_models_with_context_settings(self, mock_path, tmp_path): @@ -189,16 +189,7 @@ def test_fetch_chatgpt_models_success(self, mock_get): models = fetch_chatgpt_models("test_token", "test_account") - # Required models are prepended if not in API response - assert "gpt-5.6" in models - assert "gpt-5.6-sol" in models - assert "gpt-5.6-terra" in models - assert "gpt-5.6-luna" in models - assert "gpt-5.5" in models - assert "gpt-5.4" in models - assert "gpt-5.3-instant" in models - assert "gpt-5.2" in models - assert "gpt-5.2-codex" in models + assert models == ["gpt-5.2", "gpt-5.2-codex"] @patch("requests.get") def test_fetch_chatgpt_models_api_error(self, mock_get): diff --git a/tests/plugins/test_chatgpt_oauth_utils.py b/tests/plugins/test_chatgpt_oauth_utils.py index b65bfa5d2..8234dba06 100644 --- a/tests/plugins/test_chatgpt_oauth_utils.py +++ b/tests/plugins/test_chatgpt_oauth_utils.py @@ -850,16 +850,13 @@ def test_fetch_chatgpt_models_success(self, mock_get): result = fetch_chatgpt_models("test_access_token", "test_account_id") - # Required models are prepended, then API models follow - assert "gpt-5.4" in result - assert "gpt-5.3-instant" in result - assert "gpt-4" in result - assert "gpt-3.5-turbo" in result - assert "gpt-4-32k" in result - assert "o1-preview" in result - assert "o1-mini" in result - # Required models come first - assert result.index("gpt-5.4") < result.index("gpt-4") + assert result == [ + "gpt-4", + "gpt-3.5-turbo", + "gpt-4-32k", + "o1-preview", + "o1-mini", + ] # Verify request was made correctly mock_get.assert_called_once() @@ -887,11 +884,7 @@ def test_fetch_chatgpt_models_deduplication(self, mock_get): result = fetch_chatgpt_models("test_access_token", "test_account_id") - # Should preserve order (but implementation may not dedupe) - # The actual implementation doesn't dedupe, so we just check it returns models - assert "gpt-4" in result - assert "gpt-3.5-turbo" in result - assert "o1-preview" in result + assert result == ["gpt-4", "gpt-3.5-turbo", "o1-preview"] @patch("requests.get") def test_fetch_chatgpt_models_filtering(self, mock_get): @@ -910,16 +903,7 @@ def test_fetch_chatgpt_models_filtering(self, mock_get): result = fetch_chatgpt_models("test_access_token", "test_account_id") - # Required models prepended + API models - for m in [ - "gpt-5.4", - "gpt-5.3-instant", - "gpt-4", - "gpt-3.5-turbo", - "o1-preview", - "o1-mini", - ]: - assert m in result + assert result == ["gpt-4", "gpt-3.5-turbo", "o1-preview", "o1-mini"] @patch("requests.get") def test_fetch_chatgpt_models_http_error(self, mock_get): @@ -1073,6 +1057,25 @@ def test_add_models_to_extra_config_success(self, mock_load, mock_save): "verbosity", ] + @patch("code_puppy.plugins.chatgpt_oauth.utils.save_chatgpt_models") + @patch("code_puppy.plugins.chatgpt_oauth.utils.load_chatgpt_models") + def test_add_models_removes_stale_plugin_models(self, mock_load, mock_save): + mock_load.return_value = { + "chatgpt-gpt-5.6": { + "oauth_source": "chatgpt-oauth-plugin", + "name": "gpt-5.6", + }, + "custom-model": {"type": "openai", "name": "keep-me"}, + } + mock_save.return_value = True + + assert add_models_to_extra_config(["gpt-5.6-luna"]) is True + + saved_config = mock_save.call_args[0][0] + assert "chatgpt-gpt-5.6" not in saved_config + assert "chatgpt-gpt-5.6-luna" in saved_config + assert "custom-model" in saved_config + @patch("code_puppy.plugins.chatgpt_oauth.utils.save_chatgpt_models") @patch("code_puppy.plugins.chatgpt_oauth.utils.load_chatgpt_models") def test_add_models_to_extra_config_gpt54_and_newer_support_xhigh( @@ -1084,7 +1087,6 @@ def test_add_models_to_extra_config_gpt54_and_newer_support_xhigh( result = add_models_to_extra_config( [ - "gpt-5.6", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna", @@ -1096,7 +1098,6 @@ def test_add_models_to_extra_config_gpt54_and_newer_support_xhigh( assert result is True saved_config = mock_save.call_args[0][0] for model_name in ( - "chatgpt-gpt-5.6", "chatgpt-gpt-5.6-sol", "chatgpt-gpt-5.6-terra", "chatgpt-gpt-5.6-luna", @@ -1110,6 +1111,9 @@ def test_add_models_to_extra_config_gpt54_and_newer_support_xhigh( "verbosity", ] assert model_config["supports_xhigh_reasoning"] is True + assert model_config["supports_ultra_reasoning"] is model_name.startswith( + "chatgpt-gpt-5.6-" + ) @patch("code_puppy.plugins.chatgpt_oauth.utils.save_chatgpt_models") @patch("code_puppy.plugins.chatgpt_oauth.utils.load_chatgpt_models") @@ -1297,23 +1301,20 @@ def test_all_functions_handle_none_inputs_gracefully(self): def test_model_filtering_edge_cases(self): """Test model filtering with edge cases.""" - from code_puppy.plugins.chatgpt_oauth.utils import ( - DEFAULT_CODEX_MODELS, - REQUIRED_CODEX_MODELS, - ) + from code_puppy.plugins.chatgpt_oauth.utils import DEFAULT_CODEX_MODELS test_cases = [ # Empty models list - returns default models (no merge needed) - ([], DEFAULT_CODEX_MODELS, True), - # Models without slug field but with name - uses name + required prepended - ([{"name": "test"}], ["test"], False), + ([], DEFAULT_CODEX_MODELS), + # Models without slug field but with name use the advertised name + ([{"name": "test"}], ["test"]), # None model entries - skipped, returns default if no valid models - ([None], DEFAULT_CODEX_MODELS, True), - # Valid slug entries - required models prepended - ([{"slug": "gpt-4"}], ["gpt-4"], False), + ([None], DEFAULT_CODEX_MODELS), + # Valid slug entries are returned unchanged + ([{"slug": "gpt-4"}], ["gpt-4"]), ] - for input_models, base_expected, is_fallback in test_cases: + for input_models, expected in test_cases: with patch("requests.get") as mock_get: mock_response = Mock() mock_response.status_code = 200 @@ -1321,13 +1322,4 @@ def test_model_filtering_edge_cases(self): mock_get.return_value = mock_response result = fetch_chatgpt_models("test_access_token", "test_account_id") - if is_fallback: - assert result == base_expected, f"Failed for input: {input_models}" - else: - # Required models are prepended to API results - for m in REQUIRED_CODEX_MODELS: - assert m in result, ( - f"Missing required {m} for input: {input_models}" - ) - for m in base_expected: - assert m in result, f"Missing {m} for input: {input_models}" + assert result == expected, f"Failed for input: {input_models}" diff --git a/tests/test_config_full_coverage.py b/tests/test_config_full_coverage.py index f5838f61e..70163e8aa 100644 --- a/tests/test_config_full_coverage.py +++ b/tests/test_config_full_coverage.py @@ -321,6 +321,10 @@ def test_set_openai_reasoning_effort_valid(self): cp_config.set_openai_reasoning_effort("high") assert cp_config.get_openai_reasoning_effort() == "high" + def test_set_openai_reasoning_effort_ultra(self): + cp_config.set_openai_reasoning_effort("ULTRA") + assert cp_config.get_openai_reasoning_effort() == "ultra" + def test_set_openai_reasoning_effort_invalid(self): with pytest.raises(ValueError): cp_config.set_openai_reasoning_effort("bogus") From b7388527afd9303c20d5d2a3ec919ff999735bed Mon Sep 17 00:00:00 2001 From: Mike Pfaffenberger Date: Fri, 10 Jul 2026 08:01:52 -0400 Subject: [PATCH 05/78] feat(queue): add full-screen prompt manager --- code_puppy/plugins/steer_queue/__init__.py | 2 +- code_puppy/plugins/steer_queue/queue_menu.py | 502 +++++++++++++++--- .../plugins/steer_queue/register_callbacks.py | 6 +- tests/plugins/test_steer_queue.py | 81 +++ 4 files changed, 502 insertions(+), 89 deletions(-) diff --git a/code_puppy/plugins/steer_queue/__init__.py b/code_puppy/plugins/steer_queue/__init__.py index 68b9c604c..60898ad41 100644 --- a/code_puppy/plugins/steer_queue/__init__.py +++ b/code_puppy/plugins/steer_queue/__init__.py @@ -4,7 +4,7 @@ plugin provides the escape hatches and the management UI: * ``/steer `` -- inject guidance mid-turn (the old default). -* ``/queue`` -- TUI to view / add / edit / delete queued prompts, +* ``/queue`` -- full-screen TUI to view, add, edit, reorder, and delete prompts, available at idle and mid-run alike. * ``(N queued)`` -- live count on the bottom bar's status row. """ diff --git a/code_puppy/plugins/steer_queue/queue_menu.py b/code_puppy/plugins/steer_queue/queue_menu.py index 6bc05896e..e4896db8d 100644 --- a/code_puppy/plugins/steer_queue/queue_menu.py +++ b/code_puppy/plugins/steer_queue/queue_menu.py @@ -1,114 +1,441 @@ -"""The /queue TUI: view / add / edit / delete queued prompts. - -Runs as an async coroutine driven by ``asyncio.run`` on a worker thread -(the same pattern as ``interactive_set_picker``), while the caller has -already released the terminal via ``suspended_run_ui`` -- both the idle -REPL and the mid-run drain wrap command execution in it. - -Mutations write back through -``PauseController.replace_pending_steer_queued`` / ``request_steer`` so -the '(N queued)' status suffix updates via the controller's listeners. -Concurrent drains aren't a concern while the menu is open: mid-run the -agent is paused at its boundary, and at idle nothing consumes the queue. +"""Full-screen queue manager for ``/queue``. + +The application owns presentation state only. Queue mutations continue to flow +through :class:`PauseController`, keeping listeners and the bottom status bar in +sync with edits made here. """ from __future__ import annotations import logging +from dataclasses import dataclass from typing import List, Optional logger = logging.getLogger(__name__) -_ADD = "[ Add prompt ]" -_DONE = "[ Done ]" -_EDIT = "Edit" -_DELETE = "Delete" -_BACK = "Back" - -_PREVIEW_CELLS = 56 # menu-row preview length for long prompts +_PREVIEW_CELLS = 58 -def _preview(text: str) -> str: - flat = " ".join(text.split()) # collapse newlines for the menu row - if len(flat) <= _PREVIEW_CELLS: +def _preview(text: str, width: int = _PREVIEW_CELLS) -> str: + """Return a compact single-line preview for a queued prompt.""" + flat = " ".join(text.split()) + if len(flat) <= width: return flat - return flat[: _PREVIEW_CELLS - 1] + "\u2026" + return flat[: max(0, width - 1)] + "…" -async def _input_line(prompt_text: str, default: str = "") -> Optional[str]: - """One line of input with editing; None on Ctrl+C / Ctrl+D.""" - from prompt_toolkit import PromptSession +@dataclass +class QueueMenuState: + """Small, testable state adapter around ``PauseController``.""" - try: - return await PromptSession().prompt_async(prompt_text, default=default) - except (KeyboardInterrupt, EOFError): - return None + controller: object + selected: int = 0 + editing: bool = False + adding: bool = False + delete_armed: Optional[int] = None + notice: str = "" + @property + def items(self) -> List[str]: + return self.controller.peek_pending_steer_queued() -async def run_queue_menu() -> None: - """Main loop: list -> item actions -> back, until Done / Ctrl+C.""" - from code_puppy.messaging.pause_controller import get_pause_controller - from code_puppy.tools.common import arrow_select_async - - pc = get_pause_controller() - while True: - items = pc.peek_pending_steer_queued() - choices = [f"{i + 1}. {_preview(text)}" for i, text in enumerate(items)] - choices += [_ADD, _DONE] - - def _full_text(index: int, _items: List[str] = items) -> str: - return _items[index] if index < len(_items) else "" - - try: - selection = await arrow_select_async( - f"Prompt queue \u2014 {len(items)} item(s)", - choices, - preview_callback=_full_text, - ) - except KeyboardInterrupt: - return + @property + def selected_text(self) -> str: + items = self.items + if not items: + return "" + self.selected = min(max(self.selected, 0), len(items) - 1) + return items[self.selected] - if selection == _DONE: + def clamp_selection(self) -> None: + self.selected = min(max(self.selected, 0), max(0, len(self.items) - 1)) + + def move_selection(self, delta: int) -> None: + if not self.items: + self.selected = 0 return - if selection == _ADD: - text = await _input_line("(add) > ") - if text and text.strip(): - pc.request_steer(text, mode="queue") - continue + self.selected = min(max(self.selected + delta, 0), len(self.items) - 1) + self.delete_armed = None + self.notice = "" - index = choices.index(selection) - await _item_menu(pc, index) + def begin_add(self) -> None: + self.editing = True + self.adding = True + self.delete_armed = None + self.notice = "Adding prompt — Ctrl+S save · Esc cancel" + def begin_edit(self) -> bool: + if not self.items: + self.notice = "Queue is empty — press A to add a prompt" + return False + self.editing = True + self.adding = False + self.delete_armed = None + self.notice = "Editing prompt — Ctrl+S save · Esc cancel" + return True -async def _item_menu(pc, index: int) -> None: - """Actions for one queued prompt: edit / delete / back.""" - from code_puppy.tools.common import arrow_select_async + def cancel_edit(self) -> None: + self.editing = False + self.adding = False + self.notice = "Edit cancelled" - items = pc.peek_pending_steer_queued() - if index >= len(items): - return # queue changed underneath us; back to the list - try: - action = await arrow_select_async(items[index], [_EDIT, _DELETE, _BACK]) - except KeyboardInterrupt: - return + def save(self, text: str) -> bool: + normalized = text.strip() + if not normalized: + self.notice = "Prompt cannot be blank" + return False - if action == _EDIT: - text = await _input_line("(edit) > ", default=items[index]) - if text is not None and text.strip(): - items[index] = text - pc.replace_pending_steer_queued(items) - elif action == _DELETE: - del items[index] - pc.replace_pending_steer_queued(items) + items = self.items + if self.adding: + self.controller.request_steer(normalized, mode="queue") + self.selected = len(items) + self.notice = "Prompt added" + elif self.selected < len(items): + items[self.selected] = normalized + self.controller.replace_pending_steer_queued(items) + self.notice = "Prompt updated" + else: + self.notice = "That queue item no longer exists" + return False + self.editing = False + self.adding = False + self.delete_armed = None + self.clamp_selection() + return True -def open_queue_menu_blocking(timeout_s: float = 600.0) -> None: - """Run the menu on a worker thread with its own event loop. + def request_delete(self) -> bool: + if not self.items: + self.notice = "Queue is already empty" + return False + if self.delete_armed != self.selected: + self.delete_armed = self.selected + self.notice = "Press D again to delete this prompt" + return False + + items = self.items + del items[self.selected] + self.controller.replace_pending_steer_queued(items) + self.delete_armed = None + self.clamp_selection() + self.notice = "Prompt deleted" + return True + + def reorder(self, delta: int) -> bool: + items = self.items + destination = self.selected + delta + if not items or destination < 0 or destination >= len(items): + return False + items[self.selected], items[destination] = ( + items[destination], + items[self.selected], + ) + self.controller.replace_pending_steer_queued(items) + self.selected = destination + self.delete_armed = None + self.notice = "Prompt moved" + return True + + +class QueueMenuApp: + """Persistent full-screen prompt queue manager.""" + + def __init__(self, controller): + from prompt_toolkit import Application + from prompt_toolkit.buffer import Buffer + from prompt_toolkit.filters import Condition + from prompt_toolkit.key_binding import KeyBindings + from prompt_toolkit.layout import HSplit, Layout, VSplit, Window + from prompt_toolkit.layout.controls import BufferControl, FormattedTextControl + from prompt_toolkit.layout.dimension import Dimension + from prompt_toolkit.layout.processors import BeforeInput + from prompt_toolkit.styles import Style + from prompt_toolkit.widgets import Frame + + self.state = QueueMenuState(controller) + self._list_control = FormattedTextControl( + self._render_list, focusable=True, show_cursor=False + ) + self._editor_buffer = Buffer(multiline=True) + self._editor_control = BufferControl( + buffer=self._editor_buffer, + input_processors=[BeforeInput(self._editor_prefix)], + focusable=Condition(lambda: self.state.editing), + ) + self._detail_frame = Frame( + Window( + self._editor_control, + wrap_lines=True, + always_hide_cursor=Condition(lambda: not self.state.editing), + ), + title=self._detail_title, + style="class:detail", + ) + + body = VSplit( + [ + Frame( + Window( + self._list_control, + wrap_lines=False, + width=Dimension(weight=2), + ), + title="Queued prompts", + style="class:list", + ), + Window(width=1, char=" "), + self._detail_frame, + ], + padding=0, + ) + root = HSplit( + [ + Window( + FormattedTextControl(self._render_header), + height=2, + style="class:header", + ), + body, + Window( + FormattedTextControl(self._render_notice), + height=1, + style="class:notice", + ), + Window( + FormattedTextControl(self._render_footer), + height=2, + style="class:footer", + ), + ] + ) + + self._bindings = KeyBindings() + self._register_bindings() + self.application = Application( + layout=Layout(root, focused_element=self._list_control), + key_bindings=self._bindings, + full_screen=True, + mouse_support=True, + style=Style.from_dict( + { + "header": "bg:#173b22 #f3fff5 bold", + "header.title": "#63e677 bold", + "list": "bg:#101510 #d3dfd5", + "list.selected": "bg:#285c36 #ffffff bold", + "list.index": "#6f9478", + "list.empty": "#718077 italic", + "detail": "bg:#0c120d #e5eee7", + "detail.label": "#63e677 bold", + "notice": "bg:#172019 #f2c96d", + "footer": "bg:#111811 #9fb4a4", + "footer.key": "#63e677 bold", + } + ), + ) + self._sync_editor() + + def _editor_prefix(self): + return [("class:detail.label", "EDIT " if self.state.editing else "")] + + def _detail_title(self) -> str: + if self.state.adding: + return "New prompt" + if self.state.editing: + return f"Edit prompt {self.state.selected + 1}" + return "Prompt preview" + + def _render_header(self): + count = len(self.state.items) + return [ + ("class:header.title", " PROMPT QUEUE\n"), + ("class:header", f" {count} item{'s' if count != 1 else ''} waiting"), + ] + + def _render_list(self): + items = self.state.items + if not items: + return [ + ( + "class:list.empty", + "\n Queue is empty.\n\n Press A to add a prompt.", + ) + ] + fragments = [] + for index, text in enumerate(items): + style = ( + "class:list.selected" if index == self.state.selected else "class:list" + ) + marker = "▶" if index == self.state.selected else " " + fragments.extend( + [ + (style, f" {marker} "), + (f"{style} class:list.index", f"{index + 1:>2} "), + (style, _preview(text)), + (style, "\n"), + ] + ) + return fragments + + def _render_notice(self): + return [("class:notice", f" {self.state.notice}")] + + def _render_footer(self): + if self.state.editing: + return [ + ("class:footer", " "), + ("class:footer.key", "Ctrl+S"), + ("class:footer", " save "), + ("class:footer.key", "Esc"), + ("class:footer", " cancel Multi-line editing enabled"), + ] + return [ + ("class:footer", " "), + ("class:footer.key", "↑↓/JK"), + ("class:footer", " select "), + ("class:footer.key", "Enter/E"), + ("class:footer", " edit "), + ("class:footer.key", "A"), + ("class:footer", " add "), + ("class:footer.key", "D D"), + ("class:footer", " delete\n "), + ("class:footer.key", "[ ]"), + ("class:footer", " reorder "), + ("class:footer.key", "Q/Esc"), + ("class:footer", " done"), + ] + + def _sync_editor(self, text: Optional[str] = None) -> None: + value = self.state.selected_text if text is None else text + self._editor_buffer.set_document( + self._editor_buffer.document.__class__(value, cursor_position=len(value)), + bypass_readonly=True, + ) + + def _refresh(self) -> None: + self.application.invalidate() - Same pattern as ``interactive_set_picker``'s caller: prompt_toolkit - apps can't ``run()`` inside the already-running main loop, so hop to - a thread and ``asyncio.run`` there. Never raises. - """ + def _focus_list(self) -> None: + self.application.layout.focus(self._list_control) + + def _begin_add(self) -> None: + self.state.begin_add() + self._sync_editor("") + self.application.layout.focus(self._editor_control) + + def _begin_edit(self) -> None: + if self.state.begin_edit(): + self._sync_editor() + self.application.layout.focus(self._editor_control) + + def _cancel_edit(self) -> None: + self.state.cancel_edit() + self._sync_editor() + self._focus_list() + + def _save_edit(self) -> None: + if self.state.save(self._editor_buffer.text): + self._sync_editor() + self._focus_list() + + def _move_selection(self, delta: int) -> None: + self.state.move_selection(delta) + self._sync_editor() + + def _register_bindings(self) -> None: + kb = self._bindings + + @kb.add("up") + @kb.add("k") + def _up(event): + if not self.state.editing: + self._move_selection(-1) + + @kb.add("down") + @kb.add("j") + def _down(event): + if not self.state.editing: + self._move_selection(1) + + @kb.add("home") + def _home(event): + if not self.state.editing: + self.state.selected = 0 + self._sync_editor() + + @kb.add("end") + def _end(event): + if not self.state.editing and self.state.items: + self.state.selected = len(self.state.items) - 1 + self._sync_editor() + + @kb.add("a") + def _add(event): + if not self.state.editing: + self._begin_add() + + @kb.add("e") + @kb.add("enter") + def _edit(event): + if not self.state.editing: + self._begin_edit() + else: + self._editor_buffer.insert_text("\n") + + @kb.add("d") + def _delete(event): + if not self.state.editing: + deleted = self.state.request_delete() + if deleted: + self._sync_editor() + self._refresh() + + @kb.add("[") + def _move_up(event): + if not self.state.editing and self.state.reorder(-1): + self._sync_editor() + + @kb.add("]") + def _move_down(event): + if not self.state.editing and self.state.reorder(1): + self._sync_editor() + + @kb.add("c-s") + def _save(event): + if self.state.editing: + self._save_edit() + + @kb.add("escape") + def _escape(event): + if self.state.editing: + self._cancel_edit() + else: + event.app.exit() + + @kb.add("q") + def _quit(event): + if not self.state.editing: + event.app.exit() + + @kb.add("c-c") + def _ctrl_c(event): + if self.state.editing: + self._cancel_edit() + else: + event.app.exit() + + async def run(self) -> None: + await self.application.run_async() + + +async def run_queue_menu() -> None: + """Run the full-screen queue manager.""" + from code_puppy.messaging.pause_controller import get_pause_controller + + await QueueMenuApp(get_pause_controller()).run() + + +def open_queue_menu_blocking(timeout_s: float = 600.0) -> None: + """Run the menu on a worker thread with an isolated event loop.""" import asyncio import concurrent.futures @@ -121,4 +448,9 @@ def open_queue_menu_blocking(timeout_s: float = 600.0) -> None: logger.debug("queue menu failed", exc_info=True) -__all__ = ["open_queue_menu_blocking", "run_queue_menu"] +__all__ = [ + "QueueMenuApp", + "QueueMenuState", + "open_queue_menu_blocking", + "run_queue_menu", +] diff --git a/code_puppy/plugins/steer_queue/register_callbacks.py b/code_puppy/plugins/steer_queue/register_callbacks.py index 1d509c9be..85e8788e5 100644 --- a/code_puppy/plugins/steer_queue/register_callbacks.py +++ b/code_puppy/plugins/steer_queue/register_callbacks.py @@ -6,8 +6,8 @@ * ``/steer `` -- inject guidance mid-turn, interrupting the agent's current train of thought ASAP (the old Enter default). At idle it's a no-op with a hint: just type at the prompt. -* ``/queue`` -- arrow-key TUI over the queued prompts (view via the - preview pane, add, edit, delete). Works at idle and mid-run: both +* ``/queue`` -- full-screen queue manager with prompt preview, multiline + add/edit, guarded delete, and reordering. Works at idle and mid-run: both paths already release the terminal around command execution. * Status suffix -- a ``(N queued)`` tag rides the bottom bar's status row, updated live via the PauseController's steer-queue listeners @@ -130,7 +130,7 @@ def _handle_custom_command(command: str, name: str) -> Optional[bool]: def _custom_help() -> List[Tuple[str, str]]: return [ (_STEER, "Inject guidance mid-turn while the agent is running"), - (_QUEUE, "View/add/edit/delete queued prompts (TUI)"), + (_QUEUE, "Manage queued prompts in a full-screen TUI"), ] diff --git a/tests/plugins/test_steer_queue.py b/tests/plugins/test_steer_queue.py index bdb2a7d8c..57434051a 100644 --- a/tests/plugins/test_steer_queue.py +++ b/tests/plugins/test_steer_queue.py @@ -16,6 +16,11 @@ reset_pause_controller, ) from code_puppy.plugins.steer_queue import register_callbacks as rc +from code_puppy.plugins.steer_queue.queue_menu import ( + QueueMenuApp, + QueueMenuState, + _preview, +) @pytest.fixture(autouse=True) @@ -135,6 +140,82 @@ def boom(_count): assert pc.peek_pending_steer_queued() == ["still fine"] +# ========================================================================= +# Full-screen queue menu state +# ========================================================================= + + +def test_queue_preview_collapses_whitespace_and_truncates(): + assert _preview("one\n\n two", width=20) == "one two" + assert _preview("abcdefghijklmnopqrstuvwxyz", width=10) == "abcdefghi…" + + +def test_queue_state_adds_and_edits_prompts(): + pc = PauseController() + state = QueueMenuState(pc) + + state.begin_add() + assert state.save(" first prompt ") is True + assert pc.peek_pending_steer_queued() == ["first prompt"] + assert state.selected == 0 + + assert state.begin_edit() is True + assert state.save("updated prompt") is True + assert pc.peek_pending_steer_queued() == ["updated prompt"] + assert state.notice == "Prompt updated" + + +def test_queue_state_rejects_blank_prompt_without_leaving_editor(): + state = QueueMenuState(PauseController()) + state.begin_add() + assert state.save(" \n ") is False + assert state.editing is True + assert state.notice == "Prompt cannot be blank" + + +def test_queue_state_delete_requires_second_press_and_clamps_selection(): + pc = PauseController() + pc.replace_pending_steer_queued(["one", "two"]) + state = QueueMenuState(pc, selected=1) + + assert state.request_delete() is False + assert pc.peek_pending_steer_queued() == ["one", "two"] + assert state.request_delete() is True + assert pc.peek_pending_steer_queued() == ["one"] + assert state.selected == 0 + + +def test_queue_state_selection_change_disarms_delete(): + pc = PauseController() + pc.replace_pending_steer_queued(["one", "two"]) + state = QueueMenuState(pc) + state.request_delete() + state.move_selection(1) + assert state.delete_armed is None + + +def test_queue_state_reorders_items_and_tracks_selection(): + pc = PauseController() + pc.replace_pending_steer_queued(["one", "two", "three"]) + state = QueueMenuState(pc, selected=1) + + assert state.reorder(-1) is True + assert pc.peek_pending_steer_queued() == ["two", "one", "three"] + assert state.selected == 0 + assert state.reorder(-1) is False + + +def test_queue_menu_app_constructs_with_empty_and_populated_queue(): + empty_app = QueueMenuApp(PauseController()) + assert empty_app.application.full_screen is True + + pc = PauseController() + pc.replace_pending_steer_queued(["inspect this", "then do that"]) + app = QueueMenuApp(pc) + assert app.state.selected_text == "inspect this" + assert "inspect this" in app._editor_buffer.text + + # ========================================================================= # /steer command handler # ========================================================================= From b0907def735fbc37e091e9210f22cf842330d29c Mon Sep 17 00:00:00 2001 From: Mike Pfaffenberger Date: Fri, 10 Jul 2026 08:08:59 -0400 Subject: [PATCH 06/78] refactor(ui): remove emojis from tool output --- code_puppy/agents/event_stream_handler.py | 6 +- code_puppy/command_line/attachments.py | 4 +- code_puppy/command_line/autosave_menu.py | 8 +- code_puppy/command_line/clipboard.py | 4 +- code_puppy/command_line/colors_menu.py | 40 +++---- code_puppy/command_line/core_commands.py | 2 +- .../command_line/prompt_toolkit_completion.py | 2 +- code_puppy/messaging/rich_renderer.py | 106 +++--------------- code_puppy/messaging/subagent_console.py | 2 +- code_puppy/tools/browser/browser_control.py | 4 +- .../tools/browser/browser_interactions.py | 6 +- code_puppy/tools/browser/browser_locators.py | 8 +- .../tools/browser/browser_navigation.py | 2 +- code_puppy/tools/browser/browser_scripts.py | 4 +- code_puppy/tools/browser/browser_workflows.py | 2 +- code_puppy/tools/command_runner.py | 10 +- code_puppy/tools/image_tools.py | 2 +- tests/command_line/test_autosave_menu.py | 6 +- tests/command_line/test_clipboard.py | 14 +-- tests/command_line/test_colors_menu.py | 6 +- tests/messaging/test_rich_renderer.py | 29 ++--- tests/test_cli_runner_coverage.py | 2 +- tests/test_rich_renderer.py | 4 +- 23 files changed, 101 insertions(+), 172 deletions(-) diff --git a/code_puppy/agents/event_stream_handler.py b/code_puppy/agents/event_stream_handler.py index ee218050a..bcd237a57 100644 --- a/code_puppy/agents/event_stream_handler.py +++ b/code_puppy/agents/event_stream_handler.py @@ -440,15 +440,15 @@ def _abort_all_drainers() -> None: if not _suppress_tool_progress(): tool_name = tool_names.get(event.index, "") count = token_count[event.index] - # Display with tool wrench icon and tool name + # Display tool progress without decorative icons. if tool_name: console.print( - f" \U0001f527 Calling {tool_name}... {count} token(s) ", + f" Calling {tool_name}... {count} token(s) ", end="\r", ) else: console.print( - f" \U0001f527 Calling tool... {count} token(s) ", + f" Calling tool... {count} token(s) ", end="\r", ) diff --git a/code_puppy/command_line/attachments.py b/code_puppy/command_line/attachments.py index 386d4595f..66a9514ea 100644 --- a/code_puppy/command_line/attachments.py +++ b/code_puppy/command_line/attachments.py @@ -62,7 +62,7 @@ class AttachmentParsingError(RuntimeError): # Matches the placeholders inserted by ClipboardManager.add_image(). -CLIPBOARD_PLACEHOLDER_RE = re.compile(r"\[\U0001f4cb clipboard image \d+\]\s*") +CLIPBOARD_PLACEHOLDER_RE = re.compile(r"\[(?:\U0001f4cb )?clipboard image \d+\]\s*") @dataclass @@ -96,7 +96,7 @@ def resolve_user_prompt(raw_prompt: str) -> ResolvedUserPrompt: """Resolve a raw prompt into cleaned text and all attachment content. Combines :func:`parse_prompt_attachments` (file paths / URLs) with the - pending clipboard-image queue, stripping ``[\U0001f4cb clipboard image N]`` + pending clipboard-image queue, stripping ``[clipboard image N]`` placeholders from the text. Drains (and clears) the clipboard queue. """ processed = parse_prompt_attachments(raw_prompt) diff --git a/code_puppy/command_line/autosave_menu.py b/code_puppy/command_line/autosave_menu.py index 5e8da2666..ef88f6eb6 100644 --- a/code_puppy/command_line/autosave_menu.py +++ b/code_puppy/command_line/autosave_menu.py @@ -135,11 +135,9 @@ def _extract_message_content(msg) -> Tuple[str, str]: args_preview = str(args)[:100] if len(str(args)) > 100: args_preview += "..." - content_parts.append( - f"🔧 Tool Call: {tool_name}\n Args: {args_preview}" - ) + content_parts.append(f"Tool Call: {tool_name}\n Args: {args_preview}") else: - content_parts.append(f"🔧 Tool Call: {tool_name}") + content_parts.append(f"Tool Call: {tool_name}") elif part_kind == "tool-return": # Tool result being returned - show tool name and truncated result @@ -314,7 +312,7 @@ def _render_message_browser_panel( if role == "user": lines.append(("fg:ansicyan bold", " 🧑 USER")) elif role == "tool": - lines.append(("fg:ansiyellow bold", " 🔧 TOOL")) + lines.append(("fg:ansiyellow bold", " TOOL")) else: lines.append(("fg:ansigreen bold", " 🤖 ASSISTANT")) lines.append(("", "\n")) diff --git a/code_puppy/command_line/clipboard.py b/code_puppy/command_line/clipboard.py index a770605e0..d8253d731 100644 --- a/code_puppy/command_line/clipboard.py +++ b/code_puppy/command_line/clipboard.py @@ -444,7 +444,7 @@ def add_image(self, image_bytes: bytes) -> str: image_bytes: PNG image bytes to add. Returns: - Placeholder ID string like '[📋 clipboard image 1]' + Placeholder ID string like '[clipboard image 1]' Raises: ValueError: If MAX_PENDING_IMAGES limit is reached. @@ -458,7 +458,7 @@ def add_image(self, image_bytes: bytes) -> str: ) self._counter += 1 self._pending_images.append(image_bytes) - placeholder = f"[📋 clipboard image {self._counter}]" + placeholder = f"[clipboard image {self._counter}]" logger.debug( f"Added clipboard image {self._counter} " f"({len(image_bytes) / 1024:.1f}KB)" diff --git a/code_puppy/command_line/colors_menu.py b/code_puppy/command_line/colors_menu.py index dddb9a66d..f419e6b9d 100644 --- a/code_puppy/command_line/colors_menu.py +++ b/code_puppy/command_line/colors_menu.py @@ -19,46 +19,46 @@ from prompt_toolkit.widgets import Frame from rich.console import Console -# Banner display names with icons +# Banner display names; decorative icons are intentionally omitted. BANNER_DISPLAY_INFO = { - "thinking": ("THINKING", "⚡"), + "thinking": ("THINKING", ""), "agent_response": ("AGENT RESPONSE", ""), - "shell_command": ("SHELL COMMAND", "🚀"), - "read_file": ("READ FILE", "📂"), - "edit_file": ("EDIT FILE", "✏️"), - "create_file": ("CREATE FILE", "📝"), - "replace_in_file": ("REPLACE IN FILE", "✏️"), - "delete_snippet": ("DELETE SNIPPET", "✂️"), - "grep": ("GREP", "📂"), - "directory_listing": ("DIRECTORY LISTING", "📂"), + "shell_command": ("SHELL COMMAND", ""), + "read_file": ("READ FILE", ""), + "edit_file": ("EDIT FILE", ""), + "create_file": ("CREATE FILE", ""), + "replace_in_file": ("REPLACE IN FILE", ""), + "delete_snippet": ("DELETE SNIPPET", ""), + "grep": ("GREP", ""), + "directory_listing": ("DIRECTORY LISTING", ""), "agent_reasoning": ("AGENT REASONING", ""), - "invoke_agent": ("🤖 INVOKE AGENT", ""), + "invoke_agent": ("INVOKE AGENT", ""), "subagent_response": ("✓ AGENT RESPONSE", ""), "list_agents": ("LIST AGENTS", ""), - "universal_constructor": ("UNIVERSAL CONSTRUCTOR", "🔧"), - "terminal_tool": ("TERMINAL TOOL", "🖥️"), - "llm_judge": ("LLM JUDGE", "⚖️"), + "universal_constructor": ("UNIVERSAL CONSTRUCTOR", ""), + "terminal_tool": ("TERMINAL TOOL", ""), + "llm_judge": ("LLM JUDGE", ""), } # Sample content to show after each banner BANNER_SAMPLE_CONTENT = { "thinking": "Let me analyze this code structure and figure out the best approach...", "agent_response": "I've implemented the feature you requested. The changes include...", - "shell_command": "$ npm run test -- --silent\n⏱ Timeout: 60s", + "shell_command": "$ npm run test -- --silent\nTimeout: 60s", "read_file": "/path/to/file.py (lines 1-50)", "edit_file": "MODIFY /path/to/file.py\n--- a/file.py\n+++ b/file.py", "create_file": "CREATE /path/to/new_file.py\nFile created successfully.", "replace_in_file": "MODIFY /path/to/file.py\n--- a/file.py\n+++ b/file.py", "delete_snippet": "MODIFY /path/to/file.py\nSnippet deleted from file.", - "grep": "/src for 'handleClick'\n📄 Button.tsx (3 matches)", - "directory_listing": "/src (recursive=True)\n📁 components/\n └── Button.tsx", + "grep": "/src for 'handleClick'\nButton.tsx (3 matches)", + "directory_listing": "/src (recursive=True)\ncomponents/\n └── Button.tsx", "agent_reasoning": "Current reasoning:\nI need to refactor this function...", "invoke_agent": "code-reviewer (New session)\nSession: review-auth-abc123", "subagent_response": "code-reviewer\nThe code looks good overall...", - "list_agents": "- code-puppy: Code Puppy 🐶\n- planning-agent: Planning Agent", - "universal_constructor": "action=create tool_name=api.weather\n✅ Created successfully", + "list_agents": "- code-puppy: Code Puppy\n- planning-agent: Planning Agent", + "universal_constructor": "action=create tool_name=api.weather\nCreated successfully", "terminal_tool": "$ chromium --headless\nBrowser terminal session started", - "llm_judge": "🎯 Verdict: Complete ✅\nGoal verified — all tests pass.", + "llm_judge": "Verdict: Complete\nGoal verified — all tests pass.", } # Available background colors grouped by theme diff --git a/code_puppy/command_line/core_commands.py b/code_puppy/command_line/core_commands.py index 4b2f3353f..6b5924321 100644 --- a/code_puppy/command_line/core_commands.py +++ b/code_puppy/command_line/core_commands.py @@ -163,7 +163,7 @@ def handle_paste_command(command: str) -> bool: if placeholder: manager = get_clipboard_manager() count = manager.get_pending_count() - emit_success(f"📋 {placeholder}") + emit_success(placeholder) emit_info(f"Total pending clipboard images: {count}") emit_info("Type your prompt and press Enter to send with the image(s)") else: diff --git a/code_puppy/command_line/prompt_toolkit_completion.py b/code_puppy/command_line/prompt_toolkit_completion.py index 2e4ae3ff7..9c9cdb98a 100644 --- a/code_puppy/command_line/prompt_toolkit_completion.py +++ b/code_puppy/command_line/prompt_toolkit_completion.py @@ -951,7 +951,7 @@ def handle_image_paste_f3(event): event.app.current_buffer.insert_text("[no image in clipboard] ") event.app.output.bell() except Exception: - event.app.current_buffer.insert_text("[❌ clipboard error] ") + event.app.current_buffer.insert_text("[clipboard error] ") event.app.output.bell() session = _NoGhostLinesPromptSession( diff --git a/code_puppy/messaging/rich_renderer.py b/code_puppy/messaging/rich_renderer.py index 195a6adb6..0ec48831e 100644 --- a/code_puppy/messaging/rich_renderer.py +++ b/code_puppy/messaging/rich_renderer.py @@ -603,8 +603,7 @@ def _render_file_listing(self, msg: FileListingMessage) -> None: rec_flag = f"(recursive={msg.recursive})" banner = self._format_banner("directory_listing", "DIRECTORY LISTING") self._console.print( - f"\n{banner} " - f"📂 [bold cyan]{msg.directory}[/bold cyan] [dim]{rec_flag}[/dim]\n" + f"\n{banner} [bold cyan]{msg.directory}[/bold cyan] [dim]{rec_flag}[/dim]\n" ) # Build a tree structure: {parent_path: {files: [], dirs: set(), size: int}} @@ -695,7 +694,7 @@ def get_recursive_file_count(d: str) -> int: summary = f" [dim]({', '.join(parts)})[/dim]" if parts else "" self._console.print( - f"{indent}📁 [bold blue]{dir_name}/[/bold blue]{summary}" + f"{indent}[bold blue]{dir_name}/[/bold blue]{summary}" ) # Recursively show subdirectories @@ -708,8 +707,8 @@ def get_recursive_file_count(d: str) -> int: # Summary self._console.print("\n[bold cyan]Summary:[/bold cyan]") self._console.print( - f"📁 [blue]{msg.dir_count} directories[/blue], " - f"📄 [green]{msg.file_count} files[/green] " + f"[blue]{msg.dir_count} directories[/blue], " + f"[green]{msg.file_count} files[/green] " f"[dim]({self._format_size(msg.total_size)} total)[/dim]" ) @@ -730,9 +729,7 @@ def _render_file_content(self, msg: FileContentMessage) -> None: # Just print the header - content is for LLM only banner = self._format_banner("read_file", "READ FILE") - self._console.print( - f"\n{banner} 📂 [bold cyan]{msg.path}[/bold cyan]{line_info}" - ) + self._console.print(f"\n{banner} [bold cyan]{msg.path}[/bold cyan]{line_info}") # High mode: show token count and total lines. if get_output_level() == "high": @@ -751,7 +748,7 @@ def _render_grep_result(self, msg: GrepResultMessage) -> None: # Header banner = self._format_banner("grep", "GREP") self._console.print( - f"\n{banner} 📂 [dim]{msg.directory} for '{msg.search_term}'[/dim]" + f"\n{banner} [dim]{msg.directory} for '{msg.search_term}'[/dim]" ) # High mode: show total files searched. @@ -782,7 +779,7 @@ def _render_grep_result(self, msg: GrepResultMessage) -> None: file_matches = by_file[file_path] match_word = "match" if len(file_matches) == 1 else "matches" self._console.print( - f"\n[dim]📄 {file_path} ({len(file_matches)} {match_word})[/dim]" + f"\n[dim]{file_path} ({len(file_matches)} {match_word})[/dim]" ) # Show each match with line number and content @@ -816,7 +813,7 @@ def _render_grep_result(self, msg: GrepResultMessage) -> None: file_matches = by_file[file_path] match_word = "match" if len(file_matches) == 1 else "matches" self._console.print( - f"[dim]📄 {file_path} ({len(file_matches)} {match_word})[/dim]" + f"[dim]{file_path} ({len(file_matches)} {match_word})[/dim]" ) # Summary - subtle @@ -842,9 +839,7 @@ def _render_diff(self, msg: DiffMessage) -> None: return # Operation-specific styling - op_icons = {"create": "✨", "modify": "✏️", "delete": "🗑️"} op_colors = {"create": "green", "modify": "yellow", "delete": "red"} - icon = op_icons.get(msg.operation, "📄") op_color = op_colors.get(msg.operation, "white") # Choose banner based on operation type @@ -856,7 +851,7 @@ def _render_diff(self, msg: DiffMessage) -> None: banner = self._format_banner("replace_in_file", "EDIT FILE") self._console.print( f"\n{banner} " - f"{icon} [{op_color}]{msg.operation.upper()}[/{op_color}] " + f"[{op_color}]{msg.operation.upper()}[/{op_color}] " f"[bold cyan]{msg.path}[/bold cyan]" ) @@ -908,21 +903,21 @@ def _render_shell_start(self, msg: ShellStartMessage) -> None: # Add background indicator if running in background mode if msg.background: self._console.print( - f"\n{banner} 🚀 [dim]$ {safe_command}[/dim] [bold magenta][BACKGROUND 🌙][/bold magenta]" + f"\n{banner} [dim]$ {safe_command}[/dim] [bold magenta][BACKGROUND][/bold magenta]" ) else: - self._console.print(f"\n{banner} 🚀 [dim]$ {safe_command}[/dim]") + self._console.print(f"\n{banner} [dim]$ {safe_command}[/dim]") # Show working directory if specified if msg.cwd: safe_cwd = escape_rich_markup(msg.cwd) - self._console.print(f"[dim]📂 Working directory: {safe_cwd}[/dim]") + self._console.print(f"[dim]Working directory: {safe_cwd}[/dim]") # Show timeout or background status if msg.background: - self._console.print("[dim]⏱ Runs detached (no timeout)[/dim]") + self._console.print("[dim]Runs detached (no timeout)[/dim]") else: - self._console.print(f"[dim]⏱ Timeout: {msg.timeout}s[/dim]") + self._console.print(f"[dim]Timeout: {msg.timeout}s[/dim]") def _render_shell_line(self, msg: ShellLineMessage) -> None: """Render shell output line preserving ANSI codes and carriage returns.""" @@ -1072,7 +1067,7 @@ def _render_universal_constructor(self, msg: UniversalConstructorMessage) -> Non # Disable Rich's auto-highlighter on these prints — we already apply # explicit markup, and the ReprHighlighter regexes mangle things like # 'uuid-gen' (stops at the hyphen) and '0.00s' (no word boundary). - header_parts = [f"\n{banner} 🔧 [bold cyan]{msg.action.upper()}[/bold cyan]"] + header_parts = [f"\n{banner} [bold cyan]{msg.action.upper()}[/bold cyan]"] if msg.tool_name: safe_tool_name = escape_rich_markup(msg.tool_name) header_parts.append(f" [dim]tool=[/dim][bold]{safe_tool_name}[/bold]") @@ -1262,75 +1257,8 @@ def _format_size(self, size_bytes: int) -> str: return f"{size_bytes / (1024 * 1024 * 1024):.1f} GB" def _get_file_icon(self, file_path: str) -> str: - """Get an emoji icon for a file based on its extension.""" - import os - - ext = os.path.splitext(file_path)[1].lower() - icons = { - # Python - ".py": "🐍", - ".pyw": "🐍", - # JavaScript/TypeScript - ".js": "📜", - ".jsx": "📜", - ".ts": "📜", - ".tsx": "📜", - # Web - ".html": "🌐", - ".htm": "🌐", - ".xml": "🌐", - ".css": "🎨", - ".scss": "🎨", - ".sass": "🎨", - # Documentation - ".md": "📝", - ".markdown": "📝", - ".rst": "📝", - ".txt": "📝", - # Config - ".json": "⚙️", - ".yaml": "⚙️", - ".yml": "⚙️", - ".toml": "⚙️", - ".ini": "⚙️", - # Images - ".jpg": "🖼️", - ".jpeg": "🖼️", - ".png": "🖼️", - ".gif": "🖼️", - ".svg": "🖼️", - ".webp": "🖼️", - # Audio - ".mp3": "🎵", - ".wav": "🎵", - ".ogg": "🎵", - ".flac": "🎵", - # Video - ".mp4": "🎬", - ".avi": "🎬", - ".mov": "🎬", - ".webm": "🎬", - # Documents - ".pdf": "📄", - ".doc": "📄", - ".docx": "📄", - ".xls": "📄", - ".xlsx": "📄", - ".ppt": "📄", - ".pptx": "📄", - # Archives - ".zip": "📦", - ".tar": "📦", - ".gz": "📦", - ".rar": "📦", - ".7z": "📦", - # Executables - ".exe": "⚡", - ".dll": "⚡", - ".so": "⚡", - ".dylib": "⚡", - } - return icons.get(ext, "📄") + """Return the neutral marker used for files in tool output.""" + return "-" # ========================================================================= # Skills diff --git a/code_puppy/messaging/subagent_console.py b/code_puppy/messaging/subagent_console.py index 9bac1bae5..4947cef80 100644 --- a/code_puppy/messaging/subagent_console.py +++ b/code_puppy/messaging/subagent_console.py @@ -40,7 +40,7 @@ "starting": {"color": "cyan", "spinner": "dots", "emoji": "🚀"}, "running": {"color": "green", "spinner": "dots", "emoji": "🐕"}, "thinking": {"color": "magenta", "spinner": "dots", "emoji": "🤔"}, - "tool_calling": {"color": "yellow", "spinner": "dots12", "emoji": "🔧"}, + "tool_calling": {"color": "yellow", "spinner": "dots12", "emoji": ""}, "completed": {"color": "green", "spinner": None, "emoji": "✅"}, "error": {"color": "red", "spinner": None, "emoji": "❌"}, } diff --git a/code_puppy/tools/browser/browser_control.py b/code_puppy/tools/browser/browser_control.py index 25d1e35c4..ba5c294e1 100644 --- a/code_puppy/tools/browser/browser_control.py +++ b/code_puppy/tools/browser/browser_control.py @@ -135,7 +135,7 @@ async def create_new_page(url: Optional[str] = None) -> Dict[str, Any]: """Create a new browser page/tab.""" group_id = generate_group_id("browser_new_page", url or "blank") emit_info( - f"BROWSER NEW PAGE 📄 {url or 'blank page'}", + f"BROWSER NEW PAGE {url or 'blank page'}", message_group=group_id, ) try: @@ -164,7 +164,7 @@ async def list_pages() -> Dict[str, Any]: """List all open browser pages/tabs.""" group_id = generate_group_id("browser_list_pages") emit_info( - "BROWSER LIST PAGES 📋", + "BROWSER LIST PAGES", message_group=group_id, ) try: diff --git a/code_puppy/tools/browser/browser_interactions.py b/code_puppy/tools/browser/browser_interactions.py index 45dc44e01..f1435b487 100644 --- a/code_puppy/tools/browser/browser_interactions.py +++ b/code_puppy/tools/browser/browser_interactions.py @@ -127,7 +127,7 @@ async def set_element_text( """Set text in an input element.""" group_id = generate_group_id("browser_set_text", f"{selector[:50]}_{text[:30]}") emit_info( - f"BROWSER SET TEXT ✏️ selector='{selector}' text='{text[:50]}{'...' if len(text) > 50 else ''}'", + f"BROWSER SET TEXT selector='{selector}' text='{text[:50]}{'...' if len(text) > 50 else ''}'", message_group=group_id, ) try: @@ -166,7 +166,7 @@ async def get_element_text( """Get text content from an element.""" group_id = generate_group_id("browser_get_text", selector[:100]) emit_info( - f"BROWSER GET TEXT 📝 selector='{selector}'", + f"BROWSER GET TEXT selector='{selector}'", message_group=group_id, ) try: @@ -228,7 +228,7 @@ async def select_option( "browser_select_option", f"{selector[:50]}_{option_desc}" ) emit_info( - f"BROWSER SELECT OPTION 📄 selector='{selector}' option='{option_desc}'", + f"BROWSER SELECT OPTION selector='{selector}' option='{option_desc}'", message_group=group_id, ) try: diff --git a/code_puppy/tools/browser/browser_locators.py b/code_puppy/tools/browser/browser_locators.py index 2ebaaff57..20bfd0670 100644 --- a/code_puppy/tools/browser/browser_locators.py +++ b/code_puppy/tools/browser/browser_locators.py @@ -71,7 +71,7 @@ async def find_by_text( """Find elements containing specific text.""" group_id = generate_group_id("browser_find_by_text", text[:50]) emit_info( - f"BROWSER FIND BY TEXT 🔍 text='{text}' exact={exact}", + f"BROWSER FIND BY TEXT text='{text}' exact={exact}", message_group=group_id, ) try: @@ -186,7 +186,7 @@ async def find_by_placeholder( """Find elements by placeholder text.""" group_id = generate_group_id("browser_find_by_placeholder", text[:50]) emit_info( - f"BROWSER FIND BY PLACEHOLDER 📝 placeholder='{text}' exact={exact}", + f"BROWSER FIND BY PLACEHOLDER placeholder='{text}' exact={exact}", message_group=group_id, ) try: @@ -244,7 +244,7 @@ async def find_by_test_id( """Find elements by test ID attribute.""" group_id = generate_group_id("browser_find_by_test_id", test_id) emit_info( - f"BROWSER FIND BY TEST ID 🧪 test_id='{test_id}'", + f"BROWSER FIND BY TEST ID test_id='{test_id}'", message_group=group_id, ) try: @@ -300,7 +300,7 @@ async def run_xpath_query( """Find elements using XPath selector.""" group_id = generate_group_id("browser_xpath_query", xpath[:100]) emit_info( - f"BROWSER XPATH QUERY 🔍 xpath='{xpath}'", + f"BROWSER XPATH QUERY xpath='{xpath}'", message_group=group_id, ) try: diff --git a/code_puppy/tools/browser/browser_navigation.py b/code_puppy/tools/browser/browser_navigation.py index 76c97a2f7..e9c85343d 100644 --- a/code_puppy/tools/browser/browser_navigation.py +++ b/code_puppy/tools/browser/browser_navigation.py @@ -135,7 +135,7 @@ async def wait_for_load_state( """Wait for page to reach a specific load state.""" group_id = generate_group_id("browser_wait_for_load", f"{state}_{timeout}") emit_info( - f"BROWSER WAIT FOR LOAD ⏱️ state={state} timeout={timeout}ms", + f"BROWSER WAIT FOR LOAD state={state} timeout={timeout}ms", message_group=group_id, ) try: diff --git a/code_puppy/tools/browser/browser_scripts.py b/code_puppy/tools/browser/browser_scripts.py index 5dbf3cc92..ce297abd0 100644 --- a/code_puppy/tools/browser/browser_scripts.py +++ b/code_puppy/tools/browser/browser_scripts.py @@ -50,7 +50,7 @@ async def scroll_page( target = element_selector or "page" group_id = generate_group_id("browser_scroll", f"{direction}_{amount}_{target}") emit_info( - f"BROWSER SCROLL 📋 direction={direction} amount={amount} target='{target}'", + f"BROWSER SCROLL direction={direction} amount={amount} target='{target}'", message_group=group_id, ) try: @@ -207,7 +207,7 @@ async def wait_for_element( """Wait for an element to reach a specific state.""" group_id = generate_group_id("browser_wait_for_element", f"{selector[:50]}_{state}") emit_info( - f"BROWSER WAIT FOR ELEMENT ⏱️ selector='{selector}' state={state} timeout={timeout}ms", + f"BROWSER WAIT FOR ELEMENT selector='{selector}' state={state} timeout={timeout}ms", message_group=group_id, ) try: diff --git a/code_puppy/tools/browser/browser_workflows.py b/code_puppy/tools/browser/browser_workflows.py index b18bdf00b..9211270e8 100644 --- a/code_puppy/tools/browser/browser_workflows.py +++ b/code_puppy/tools/browser/browser_workflows.py @@ -85,7 +85,7 @@ async def list_workflows() -> Dict[str, Any]: """List all available browser workflows.""" group_id = generate_group_id("list_workflows") emit_info( - "LIST WORKFLOWS 📋", + "LIST WORKFLOWS", message_group=group_id, ) diff --git a/code_puppy/tools/command_runner.py b/code_puppy/tools/command_runner.py index a24495035..066cde15c 100644 --- a/code_puppy/tools/command_runner.py +++ b/code_puppy/tools/command_runner.py @@ -1151,9 +1151,9 @@ async def run_shell_command( # Emit info about background execution emit_info( - f"🚀 Background process started (PID: {process.pid}) - no timeout, runs until complete" + f"Background process started (PID: {process.pid}) - no timeout, runs until complete" ) - emit_info(f"📄 Output logging to: {log_file.name}") + emit_info(f"Output logging to: {log_file.name}") # Return immediately - don't wait, don't block return ShellCommandOutput( @@ -1178,7 +1178,7 @@ async def run_shell_command( except OSError: pass # Emit error message so user sees what happened - emit_error(f"❌ Failed to start background process: {e}") + emit_error(f"Failed to start background process: {e}") return ShellCommandOutput( success=False, command=command, @@ -1218,13 +1218,13 @@ async def run_shell_command( # Build panel content panel_content = Text() - panel_content.append("⚡ Requesting permission to run:\n", style="bold yellow") + panel_content.append("Requesting permission to run:\n", style="bold yellow") panel_content.append("$ ", style="bold green") panel_content.append(command, style="bold white") if cwd: panel_content.append("\n\n", style="") - panel_content.append("📂 Working directory: ", style="dim") + panel_content.append("Working directory: ", style="dim") panel_content.append(cwd, style="dim cyan") # Use the common approval function (async version). diff --git a/code_puppy/tools/image_tools.py b/code_puppy/tools/image_tools.py index 9081a09db..1654d66bc 100644 --- a/code_puppy/tools/image_tools.py +++ b/code_puppy/tools/image_tools.py @@ -103,7 +103,7 @@ async def load_image( ) -> Union[ToolReturn, Dict[str, Any]]: """Load an image from the filesystem for visual analysis.""" group_id = generate_group_id("load_image", image_path) - emit_info(f"LOAD IMAGE 🖼️ {image_path}", message_group=group_id) + emit_info(f"LOAD IMAGE {image_path}", message_group=group_id) try: image_file = Path(image_path) diff --git a/tests/command_line/test_autosave_menu.py b/tests/command_line/test_autosave_menu.py index 2df930a86..484cdcdf8 100644 --- a/tests/command_line/test_autosave_menu.py +++ b/tests/command_line/test_autosave_menu.py @@ -709,7 +709,7 @@ def test_tool_call_returns_tool_role(self): ) role, content = _extract_message_content(msg) assert role == "tool" - assert "🔧 Tool Call: edit_file" in content + assert "Tool Call: edit_file" in content def test_text_response_returns_assistant_role(self): """Response with text part returns role='assistant'.""" @@ -760,7 +760,7 @@ def test_tool_call_extracts_tool_name_and_args(self): ], ) role, content = _extract_message_content(msg) - assert "🔧 Tool Call: edit_file" in content + assert "Tool Call: edit_file" in content assert "Args:" in content assert "file_path" in content @@ -787,7 +787,7 @@ def test_tool_call_without_args(self): ], ) role, content = _extract_message_content(msg) - assert "🔧 Tool Call: list_files" in content + assert "Tool Call: list_files" in content def test_tool_return_extracts_tool_name_and_result(self): """Tool return shows tool name and content preview.""" diff --git a/tests/command_line/test_clipboard.py b/tests/command_line/test_clipboard.py index b38b5a3d6..1733030f9 100644 --- a/tests/command_line/test_clipboard.py +++ b/tests/command_line/test_clipboard.py @@ -23,7 +23,7 @@ def test_add_image_returns_placeholder(self): placeholder = manager.add_image(fake_png) - assert placeholder == "[📋 clipboard image 1]" + assert placeholder == "[clipboard image 1]" def test_placeholder_increments_correctly(self): """Test that placeholder numbers increment with each add.""" @@ -36,9 +36,9 @@ def test_placeholder_increments_correctly(self): p2 = manager.add_image(fake_png) p3 = manager.add_image(fake_png) - assert p1 == "[📋 clipboard image 1]" - assert p2 == "[📋 clipboard image 2]" - assert p3 == "[📋 clipboard image 3]" + assert p1 == "[clipboard image 1]" + assert p2 == "[clipboard image 2]" + assert p3 == "[clipboard image 3]" def test_get_pending_images_returns_binary_content_list(self): """Test that get_pending_images returns list of BinaryContent.""" @@ -136,12 +136,12 @@ def test_counter_persists_after_clear(self): fake_png = b"\x89PNG\r\n\x1a\n" + b"\x00" * 100 p1 = manager.add_image(fake_png) - assert p1 == "[📋 clipboard image 1]" + assert p1 == "[clipboard image 1]" manager.clear_pending() p2 = manager.add_image(fake_png) - assert p2 == "[📋 clipboard image 2]" + assert p2 == "[clipboard image 2]" def test_manager_is_thread_safe(self): """Test basic thread safety of add_image and get_pending_count.""" @@ -300,7 +300,7 @@ def test_returns_placeholder_when_image_captured(self): ): result = clipboard.capture_clipboard_image_to_pending() - assert result == "[📋 clipboard image 1]" + assert result == "[clipboard image 1]" # Restore clipboard._clipboard_manager = original_manager diff --git a/tests/command_line/test_colors_menu.py b/tests/command_line/test_colors_menu.py index 6b9ed849f..a1edad814 100644 --- a/tests/command_line/test_colors_menu.py +++ b/tests/command_line/test_colors_menu.py @@ -53,13 +53,13 @@ def test_thinking_banner_display(self): """Test thinking banner display info.""" name, icon = BANNER_DISPLAY_INFO["thinking"] assert name == "THINKING" - assert icon == "⚡" + assert icon == "" def test_shell_command_banner_display(self): """Test shell command banner display info.""" name, icon = BANNER_DISPLAY_INFO["shell_command"] assert name == "SHELL COMMAND" - assert icon == "🚀" + assert icon == "" class TestBannerSampleContent: @@ -259,7 +259,7 @@ def test_preview_displays_correct_banner_name(self): name, icon = BANNER_DISPLAY_INFO[banner_key] # Preview would show this name and icon assert name == "SHELL COMMAND" - assert icon == "🚀" + assert icon == "" class TestThemeManagement: diff --git a/tests/messaging/test_rich_renderer.py b/tests/messaging/test_rich_renderer.py index 3d8e81abe..c7ce874ba 100644 --- a/tests/messaging/test_rich_renderer.py +++ b/tests/messaging/test_rich_renderer.py @@ -902,19 +902,22 @@ def test_format_size(renderer): def test_get_file_icon(renderer): - assert renderer._get_file_icon("test.py") == "🐍" - assert renderer._get_file_icon("test.js") == "📜" - assert renderer._get_file_icon("test.html") == "🌐" - assert renderer._get_file_icon("test.css") == "🎨" - assert renderer._get_file_icon("test.md") == "📝" - assert renderer._get_file_icon("test.json") == "⚙️" - assert renderer._get_file_icon("test.jpg") == "🖼️" - assert renderer._get_file_icon("test.mp3") == "🎵" - assert renderer._get_file_icon("test.mp4") == "🎬" - assert renderer._get_file_icon("test.pdf") == "📄" - assert renderer._get_file_icon("test.zip") == "📦" - assert renderer._get_file_icon("test.exe") == "⚡" - assert renderer._get_file_icon("test.unknown") == "📄" + for path in ( + "test.py", + "test.js", + "test.html", + "test.css", + "test.md", + "test.json", + "test.jpg", + "test.mp3", + "test.mp4", + "test.pdf", + "test.zip", + "test.exe", + "test.unknown", + ): + assert renderer._get_file_icon(path) == "-" def test_get_banner_color(renderer): diff --git a/tests/test_cli_runner_coverage.py b/tests/test_cli_runner_coverage.py index 7fdc73a72..9ff9d3f6d 100644 --- a/tests/test_cli_runner_coverage.py +++ b/tests/test_cli_runner_coverage.py @@ -121,7 +121,7 @@ async def test_clipboard_placeholder_cleaned(self): # End-to-end through the REAL resolver: only the clipboard manager # is stubbed, so this covers placeholder stripping in # resolve_user_prompt as consumed by run_prompt_with_attachments. - placeholder = "[\U0001f4cb clipboard image 1]" + placeholder = "[clipboard image 1]" with ( patch( "code_puppy.command_line.clipboard.get_clipboard_manager" diff --git a/tests/test_rich_renderer.py b/tests/test_rich_renderer.py index 0ddbb37b3..aa1bfa2da 100644 --- a/tests/test_rich_renderer.py +++ b/tests/test_rich_renderer.py @@ -281,8 +281,8 @@ def test_format_size(size_bytes: int, expected: str) -> None: def test_get_file_icon_defaults() -> None: renderer, _ = _make_renderer() - assert renderer._get_file_icon("notes.md") == "📝" - assert renderer._get_file_icon("unknown.bin") == "📄" + assert renderer._get_file_icon("notes.md") == "-" + assert renderer._get_file_icon("unknown.bin") == "-" def test_markdown_rendering_to_real_console() -> None: From 27e04545403310069be7e76cafdd72ebf8b620b8 Mon Sep 17 00:00:00 2001 From: Mike Pfaffenberger Date: Fri, 10 Jul 2026 09:07:15 -0400 Subject: [PATCH 07/78] feat: improve compaction flow and Codex usage status --- .gitignore | 2 + code_puppy/agents/_compaction.py | 16 ++- code_puppy/agents/_run_signals.py | 5 + code_puppy/chatgpt_codex_client.py | 7 ++ code_puppy/command_line/core_commands.py | 2 +- code_puppy/command_line/session_commands.py | 9 ++ code_puppy/messaging/pause_controller.py | 28 +++++ code_puppy/messaging/run_ui.py | 1 - code_puppy/messaging/spinner/__init__.py | 24 +++- code_puppy/model_factory.py | 2 +- code_puppy/model_switching.py | 26 ++++ code_puppy/plugins/chatgpt_oauth/config.py | 2 +- .../plugins/chatgpt_oauth/oauth_flow.py | 2 +- .../chatgpt_oauth/register_callbacks.py | 29 ++++- .../plugins/chatgpt_oauth/test_plugin.py | 47 ++++---- code_puppy/plugins/chatgpt_oauth/usage.py | 113 ++++++++++++++++++ code_puppy/plugins/chatgpt_oauth/utils.py | 6 +- .../shell_safety/register_callbacks.py | 3 +- code_puppy/provider_identity.py | 1 + tests/agents/test_compaction.py | 19 +++ tests/command_line/test_session_commands.py | 15 +++ tests/command_line/test_tutorial.py | 2 +- tests/messaging/spinner/test_spinner_shim.py | 26 ++++ tests/messaging/test_pause_controller.py | 23 ++++ tests/messaging/test_rich_renderer.py | 4 +- tests/plugins/test_chatgpt_oauth_coverage.py | 6 + .../plugins/test_chatgpt_oauth_integration.py | 44 +++---- tests/plugins/test_chatgpt_oauth_usage.py | 47 ++++++++ tests/plugins/test_chatgpt_oauth_utils.py | 40 +++---- tests/plugins/test_oauth_integration.py | 8 +- .../test_subagent_panel_model_short.py | 2 +- tests/test_chatgpt_codex_client.py | 28 +++++ tests/test_completions_and_small_modules.py | 36 ++++++ tests/test_model_factory_coverage.py | 8 +- tests/test_tools_registration.py | 2 +- 35 files changed, 538 insertions(+), 97 deletions(-) create mode 100644 code_puppy/plugins/chatgpt_oauth/usage.py create mode 100644 tests/plugins/test_chatgpt_oauth_usage.py diff --git a/.gitignore b/.gitignore index 56e9f3c35..3ec03c657 100644 --- a/.gitignore +++ b/.gitignore @@ -45,3 +45,5 @@ code_puppy/bundled_skills/ .claude/hooks/ts-hooks/dist/ .json + +.state/ diff --git a/code_puppy/agents/_compaction.py b/code_puppy/agents/_compaction.py index 436c283be..a00bf2e4e 100644 --- a/code_puppy/agents/_compaction.py +++ b/code_puppy/agents/_compaction.py @@ -41,7 +41,7 @@ get_compaction_threshold, get_protected_token_count, ) -from code_puppy.messaging import emit_error, emit_info, emit_warning +from code_puppy.messaging import emit_error, emit_info, emit_success, emit_warning from code_puppy.messaging.spinner import format_context_info, update_spinner_context from code_puppy.summarization_agent import SummarizationError, run_summarization_sync @@ -282,6 +282,8 @@ def compact( messages: List[ModelMessage], model_max: int, context_overhead: int, + *, + force: bool = False, ) -> Tuple[List[ModelMessage], List[ModelMessage]]: """Unified compaction entrypoint. Replaces ``message_history_processor``. @@ -291,6 +293,8 @@ def compact( messages: Current message history (already accumulated by the caller). model_max: Effective model context window in tokens. context_overhead: Estimated overhead for system prompt + tool schemas. + force: Compact regardless of the configured context threshold. Used by + mid-run ``/compact`` at the next safe model-call boundary. Returns: ``(new_messages, dropped_messages_for_hash_tracking)``. @@ -312,7 +316,7 @@ def compact( update_spinner_context(context_summary) threshold = get_compaction_threshold() - if proportion_used <= threshold: + if not force and proportion_used <= threshold: return messages, [] strategy = get_compaction_strategy() @@ -473,12 +477,20 @@ def history_processor(messages: List[ModelMessage]) -> List[ModelMessage]: history.append(msg) messages_added += 1 + from code_puppy.messaging.pause_controller import get_pause_controller + + pause_controller = get_pause_controller() + force_compaction = pause_controller.take_compaction_request() new_history, dropped = compact( agent, history, agent._get_model_context_length(), agent._estimate_context_overhead(), + force=force_compaction, ) + if force_compaction: + detail = "" if dropped else " History was already minimal." + emit_success(f"Mid-run compaction complete.{detail}") agent._message_history = new_history for m in dropped: compacted_hashes.add(hash_message(m)) diff --git a/code_puppy/agents/_run_signals.py b/code_puppy/agents/_run_signals.py index 6c4636ccf..64b00d6e0 100644 --- a/code_puppy/agents/_run_signals.py +++ b/code_puppy/agents/_run_signals.py @@ -127,6 +127,11 @@ def reset_pause_state_at_run_start() -> None: # Clear any stale paused state (e.g. from a prior run that crashed # mid-pause). Safe / idempotent if already resumed. pc.resume() + stale_compactions = pc.drain_compaction_requests() + if stale_compactions: + emit_warning( + f"Discarded {stale_compactions} stale compaction request(s) from a previous run." + ) stale_steers = pc.drain_pending_steer() if stale_steers: emit_warning( diff --git a/code_puppy/chatgpt_codex_client.py b/code_puppy/chatgpt_codex_client.py index 7e00716e4..7870e48e1 100644 --- a/code_puppy/chatgpt_codex_client.py +++ b/code_puppy/chatgpt_codex_client.py @@ -87,6 +87,13 @@ async def send( except Exception as e: logger.debug("Failed to inject Codex fields into request: %s", e) + # The OpenAI SDK replaces the client's User-Agent while building each + # request. ChatGPT's Codex backend uses the Codex User-Agent for model + # routing, so newer models can otherwise return a misleading 404. + configured_user_agent = self.headers.get("User-Agent") + if configured_user_agent: + request.headers["User-Agent"] = configured_user_agent + # Make the actual request response = await super().send(request, *args, **kwargs) diff --git a/code_puppy/command_line/core_commands.py b/code_puppy/command_line/core_commands.py index 6b5924321..40c56d24f 100644 --- a/code_puppy/command_line/core_commands.py +++ b/code_puppy/command_line/core_commands.py @@ -207,7 +207,7 @@ def handle_tutorial_command(command: str) -> bool: from code_puppy.plugins.chatgpt_oauth.oauth_flow import run_oauth_flow run_oauth_flow() - set_model_and_reload_agent("chatgpt-gpt-5.4") + set_model_and_reload_agent("codex-gpt-5.6-sol") elif result == "claude": emit_info("🔐 Starting Claude Code OAuth flow...") from code_puppy.plugins.claude_code_oauth.register_callbacks import ( diff --git a/code_puppy/command_line/session_commands.py b/code_puppy/command_line/session_commands.py index fa91c9029..42e9f1d2e 100644 --- a/code_puppy/command_line/session_commands.py +++ b/code_puppy/command_line/session_commands.py @@ -139,6 +139,15 @@ def handle_compact_command(command: str) -> bool: from code_puppy.messaging import emit_error, emit_info, emit_success, emit_warning try: + from code_puppy.messaging.run_ui import is_run_active + + if is_run_active(): + from code_puppy.messaging.pause_controller import get_pause_controller + + get_pause_controller().request_compaction() + emit_info("Compaction requested; it will run before the next model call.") + return True + agent = get_current_agent() history = agent.get_message_history() if not history: diff --git a/code_puppy/messaging/pause_controller.py b/code_puppy/messaging/pause_controller.py index a28c440db..d9b8a0c79 100644 --- a/code_puppy/messaging/pause_controller.py +++ b/code_puppy/messaging/pause_controller.py @@ -53,6 +53,10 @@ def __init__(self) -> None: # Two queues, one per delivery mode (see ``request_steer`` docs). self._steer_queue_now: List[str] = [] self._steer_queue_queued: List[str] = [] + # One-shot requests consumed by the compaction history processor at + # the next model-call boundary. A counter (rather than a bool) avoids + # losing a second request racing with consumption of the first. + self._compaction_requests: int = 0 self._resume_event: Optional[asyncio.Event] = None self._loop: Optional[asyncio.AbstractEventLoop] = None # Rendezvous flag: set when a waiter actually parks inside @@ -196,6 +200,30 @@ def _fire_steer_queue_listeners(self, count: int) -> None: except Exception: pass + # ========================================================================= + # Deferred compaction + # ========================================================================= + + def request_compaction(self) -> None: + """Request forced compaction at the next history-processor boundary.""" + with self._lock: + self._compaction_requests += 1 + + def take_compaction_request(self) -> bool: + """Atomically consume one pending forced-compaction request.""" + with self._lock: + if not self._compaction_requests: + return False + self._compaction_requests -= 1 + return True + + def drain_compaction_requests(self) -> int: + """Clear and return all pending requests (run-start hygiene).""" + with self._lock: + count = self._compaction_requests + self._compaction_requests = 0 + return count + # ========================================================================= # Steering queue # ========================================================================= diff --git a/code_puppy/messaging/run_ui.py b/code_puppy/messaging/run_ui.py index a4949bb73..4f95bb44c 100644 --- a/code_puppy/messaging/run_ui.py +++ b/code_puppy/messaging/run_ui.py @@ -88,7 +88,6 @@ def _get_loop() -> Optional[asyncio.AbstractEventLoop]: "quit": "terminates the app; cancel the run first (Ctrl+C)", "agent": "switches/reloads the active agent while agent.run() is in flight", "clear": "wipes message history the in-flight run will commit at turn end", - "compact": "mutates message history mid-run (history processors own it)", "truncate": "mutates message history mid-run", "autosave_load": "returns the __AUTOSAVE_LOAD__ sentinel only the idle REPL can service", "quick-resume": "loads a saved session into message history mid-run", diff --git a/code_puppy/messaging/spinner/__init__.py b/code_puppy/messaging/spinner/__init__.py index 61b226389..ed4468e07 100644 --- a/code_puppy/messaging/spinner/__init__.py +++ b/code_puppy/messaging/spinner/__init__.py @@ -59,18 +59,34 @@ def _compact_count(n: int) -> str: return str(n) +def _codex_usage_suffix() -> str: + """Return cached Codex limits only while a Codex OAuth model is active.""" + try: + from code_puppy.config import get_global_model_name + from code_puppy.plugins.chatgpt_oauth.usage import get_usage_status + + if not (get_global_model_name() or "").startswith("codex-"): + return "" + usage = get_usage_status() + return f" · Codex {usage}" if usage else "" + except Exception: + logger.debug("Codex usage status unavailable", exc_info=True) + return "" + + def format_context_info(total_tokens: int, capacity: int, proportion: float) -> str: - """Create a compact context summary for the status line. + """Create a compact context summary, plus cached provider usage if relevant. - e.g. ``150.3k/500k tokens (30%)`` — the old long form - (``Tokens: 150,329/500,000 (30.1% used)``) ate half the row. + e.g. ``150.3k/500k tokens (30%) · Codex 5h 66% remaining · week 90% remaining``. + Provider usage is read from an in-memory cache; rendering never performs I/O. """ if capacity <= 0: return "" - return ( + context = ( f"{_compact_count(total_tokens)}/{_compact_count(capacity)} " f"tokens ({proportion * 100:.0f}%)" ) + return f"{context}{_codex_usage_suffix()}" def update_spinner_context(info: str) -> None: diff --git a/code_puppy/model_factory.py b/code_puppy/model_factory.py index a620bf8e5..7adb25e87 100644 --- a/code_puppy/model_factory.py +++ b/code_puppy/model_factory.py @@ -814,7 +814,7 @@ def get_model(model_name: str, config: Dict[str, Any]) -> Any: provider=provider, profile=_thinking_tags_profile(model_name, model_config), ) - if model_name == "chatgpt-gpt-5-codex": + if model_name == "codex-gpt-5-codex": model = OpenAIResponsesModel(model_config["name"], provider=provider) return model elif model_type == "zai_coding": diff --git a/code_puppy/model_switching.py b/code_puppy/model_switching.py index dc55daa00..12473ab13 100644 --- a/code_puppy/model_switching.py +++ b/code_puppy/model_switching.py @@ -15,6 +15,31 @@ def _get_effective_agent_model(agent) -> Optional[str]: return None +def _refresh_context_status(agent) -> None: + """Replace any stale token summary after an effective-model change. + + The history processor normally writes this status during a model run. A + model switch can happen between those writes, leaving the old model's + capacity visible until the next turn. Recompute from the reloaded agent so + the bottom bar immediately reflects the effective model. + """ + from code_puppy.messaging.spinner import ( + format_context_info, + update_spinner_context, + ) + + try: + capacity = agent._get_model_context_length() + history = agent.get_message_history() or [] + message_tokens = sum(agent.estimate_tokens_for_message(msg) for msg in history) + total_tokens = message_tokens + agent._estimate_context_overhead() + proportion = total_tokens / capacity if capacity else 0.0 + update_spinner_context(format_context_info(total_tokens, capacity, proportion)) + except Exception: + # A blank status is more honest than retaining another model's capacity. + update_spinner_context("") + + def set_model_and_reload_agent( model_name: str, *, @@ -58,6 +83,7 @@ def set_model_and_reload_agent( ) current_agent.reload_code_generation_agent() + _refresh_context_status(current_agent) emit_info("Active agent reloaded") except Exception as exc: emit_warning(f"Model changed but agent reload failed: {exc}") diff --git a/code_puppy/plugins/chatgpt_oauth/config.py b/code_puppy/plugins/chatgpt_oauth/config.py index d61403c3c..20b4bbaaf 100644 --- a/code_puppy/plugins/chatgpt_oauth/config.py +++ b/code_puppy/plugins/chatgpt_oauth/config.py @@ -22,7 +22,7 @@ # Local configuration (uses XDG_DATA_HOME) "token_storage": None, # Set dynamically in get_token_storage_path() # Model configuration - "prefix": "chatgpt-", + "prefix": "codex-", "default_context_length": 272000, "api_key_env_var": "CHATGPT_OAUTH_API_KEY", # Codex CLI version info (for User-Agent header) diff --git a/code_puppy/plugins/chatgpt_oauth/oauth_flow.py b/code_puppy/plugins/chatgpt_oauth/oauth_flow.py index 216b40966..a288bf781 100644 --- a/code_puppy/plugins/chatgpt_oauth/oauth_flow.py +++ b/code_puppy/plugins/chatgpt_oauth/oauth_flow.py @@ -325,5 +325,5 @@ def run_oauth_flow() -> None: if models: if add_models_to_extra_config(models): emit_success( - "ChatGPT models registered. Use the `chatgpt-` prefix in /model." + "Codex models registered. Use the `codex-` prefix in /model." ) diff --git a/code_puppy/plugins/chatgpt_oauth/register_callbacks.py b/code_puppy/plugins/chatgpt_oauth/register_callbacks.py index 623d9e775..ed64895f9 100644 --- a/code_puppy/plugins/chatgpt_oauth/register_callbacks.py +++ b/code_puppy/plugins/chatgpt_oauth/register_callbacks.py @@ -15,6 +15,7 @@ from .config import CHATGPT_OAUTH_CONFIG, get_token_storage_path from .oauth_flow import run_oauth_flow +from .usage import refresh_usage_in_background from .utils import ( get_valid_access_token, load_chatgpt_models, @@ -29,11 +30,14 @@ def _custom_help() -> List[Tuple[str, str]]: "chatgpt-auth", "Authenticate with ChatGPT via OAuth and import available models", ), + ("codex-auth", "Alias for /chatgpt-auth"), ( "chatgpt-status", "Check ChatGPT OAuth authentication status and configured models", ), + ("codex-status", "Alias for /chatgpt-status"), ("chatgpt-logout", "Remove ChatGPT OAuth tokens and imported models"), + ("codex-logout", "Alias for /chatgpt-logout"), ] @@ -83,16 +87,16 @@ def _handle_custom_command(command: str, name: str) -> Optional[bool]: if not name: return None - if name == "chatgpt-auth": + if name in {"chatgpt-auth", "codex-auth"}: run_oauth_flow() - set_model_and_reload_agent("chatgpt-gpt-5.4") + set_model_and_reload_agent("codex-gpt-5.6-sol") return True - if name == "chatgpt-status": + if name in {"chatgpt-status", "codex-status"}: _handle_chatgpt_status() return True - if name == "chatgpt-logout": + if name in {"chatgpt-logout", "codex-logout"}: _handle_chatgpt_logout() return True @@ -132,6 +136,9 @@ def _create_chatgpt_oauth_model( ) return None + # Refresh plan limits without delaying model creation or terminal rendering. + refresh_usage_in_background(access_token, account_id) + # Build headers for ChatGPT Codex API originator = CHATGPT_OAUTH_CONFIG.get("originator", "codex_cli_rs") client_version = CHATGPT_OAUTH_CONFIG.get("client_version", "0.144.1") @@ -169,6 +176,20 @@ def _register_model_types() -> List[Dict[str, Any]]: return [{"type": "chatgpt_oauth", "handler": _create_chatgpt_oauth_model}] +def _refresh_usage_on_agent_run( + agent_name: str, model_name: str, session_id: str | None = None +) -> None: + """Keep limits fresh for Codex runs; the actual HTTP request is asynchronous.""" + del agent_name, session_id + if not model_name.startswith("codex-"): + return + tokens = load_stored_tokens() or {} + refresh_usage_in_background( + tokens.get("access_token", ""), tokens.get("account_id", "") + ) + + register_callback("custom_command_help", _custom_help) register_callback("custom_command", _handle_custom_command) register_callback("register_model_type", _register_model_types) +register_callback("agent_run_start", _refresh_usage_on_agent_run) diff --git a/code_puppy/plugins/chatgpt_oauth/test_plugin.py b/code_puppy/plugins/chatgpt_oauth/test_plugin.py index f71078b3b..57fbcef74 100644 --- a/code_puppy/plugins/chatgpt_oauth/test_plugin.py +++ b/code_puppy/plugins/chatgpt_oauth/test_plugin.py @@ -29,7 +29,7 @@ def test_oauth_config(): """Test OAuth configuration values.""" assert config.CHATGPT_OAUTH_CONFIG["issuer"] == "https://auth.openai.com" assert config.CHATGPT_OAUTH_CONFIG["client_id"] == "app_EMoamEEZ73f0CkXaXp7hrann" - assert config.CHATGPT_OAUTH_CONFIG["prefix"] == "chatgpt-" + assert config.CHATGPT_OAUTH_CONFIG["prefix"] == "codex-" def test_jwt_parsing_with_nested_org(): @@ -148,7 +148,7 @@ def test_parse_jwt_claims(): def test_save_and_load_tokens(tmp_path): """Test token storage and retrieval.""" with patch.object( - config, "get_token_storage_path", return_value=tmp_path / "tokens.json" + utils, "get_token_storage_path", return_value=tmp_path / "tokens.json" ): tokens = { "access_token": "test_access", @@ -167,10 +167,10 @@ def test_save_and_load_tokens(tmp_path): def test_save_and_load_chatgpt_models(tmp_path): """Test ChatGPT models configuration.""" with patch.object( - config, "get_chatgpt_models_path", return_value=tmp_path / "chatgpt_models.json" + utils, "get_chatgpt_models_path", return_value=tmp_path / "chatgpt_models.json" ): models = { - "chatgpt-gpt-4o": { + "codex-gpt-4o": { "type": "openai", "name": "gpt-4o", "oauth_source": "chatgpt-oauth-plugin", @@ -188,10 +188,10 @@ def test_save_and_load_chatgpt_models(tmp_path): def test_remove_chatgpt_models(tmp_path): """Test removal of ChatGPT models from config.""" with patch.object( - config, "get_chatgpt_models_path", return_value=tmp_path / "chatgpt_models.json" + utils, "get_chatgpt_models_path", return_value=tmp_path / "chatgpt_models.json" ): models = { - "chatgpt-gpt-4o": { + "codex-gpt-4o": { "type": "openai", "oauth_source": "chatgpt-oauth-plugin", }, @@ -208,7 +208,7 @@ def test_remove_chatgpt_models(tmp_path): # Verify only ChatGPT model was removed remaining = utils.load_chatgpt_models() - assert "chatgpt-gpt-4o" not in remaining + assert "codex-gpt-4o" not in remaining assert "claude-3-opus" in remaining @@ -251,10 +251,11 @@ def test_fetch_chatgpt_models(mock_get): models = utils.fetch_chatgpt_models("test_access_token", "test_account_id") assert models is not None - # Required models always injected - assert "gpt-5.4" in models - assert "gpt-5.3-instant" in models - # API-returned models present too + # A successful endpoint response is authoritative; fallback models are + # used only when discovery fails. + assert "gpt-5.4" not in models + assert "gpt-5.4-mini" not in models + # API-returned models are preserved. assert "gpt-4o" in models assert "gpt-3.5-turbo" in models assert "o1-preview" in models @@ -271,30 +272,32 @@ def test_fetch_chatgpt_models_fallback(mock_get): models = utils.fetch_chatgpt_models("test_access_token", "test_account_id") assert models is not None - # Should return default models (including new required ones) + # Keep the fallback aligned with the models currently exposed by Codex. assert "gpt-5.4" in models - assert "gpt-5.3-instant" in models + assert "gpt-5.4-mini" in models assert "gpt-5.3-codex-spark" in models - assert "gpt-5.3-codex" in models - assert "gpt-5.2-codex" in models - assert "gpt-5.2" in models + assert "codex-auto-review" in models + assert "gpt-5.3-instant" not in models + assert "gpt-5.3-codex" not in models + assert "gpt-5.2-codex" not in models + assert "gpt-5.2" not in models def test_add_models_to_chatgpt_config(tmp_path): """Test adding models to chatgpt_models.json.""" with patch.object( - config, "get_chatgpt_models_path", return_value=tmp_path / "chatgpt_models.json" + utils, "get_chatgpt_models_path", return_value=tmp_path / "chatgpt_models.json" ): models = ["gpt-4o", "gpt-3.5-turbo"] assert utils.add_models_to_extra_config(models) loaded = utils.load_chatgpt_models() - assert "chatgpt-gpt-4o" in loaded - assert "chatgpt-gpt-3.5-turbo" in loaded - assert loaded["chatgpt-gpt-4o"]["type"] == "chatgpt_oauth" - assert loaded["chatgpt-gpt-4o"]["name"] == "gpt-4o" - assert loaded["chatgpt-gpt-4o"]["oauth_source"] == "chatgpt-oauth-plugin" + assert "codex-gpt-4o" in loaded + assert "codex-gpt-3.5-turbo" in loaded + assert loaded["codex-gpt-4o"]["type"] == "chatgpt_oauth" + assert loaded["codex-gpt-4o"]["name"] == "gpt-4o" + assert loaded["codex-gpt-4o"]["oauth_source"] == "chatgpt-oauth-plugin" if __name__ == "__main__": diff --git a/code_puppy/plugins/chatgpt_oauth/usage.py b/code_puppy/plugins/chatgpt_oauth/usage.py new file mode 100644 index 000000000..acce7f887 --- /dev/null +++ b/code_puppy/plugins/chatgpt_oauth/usage.py @@ -0,0 +1,113 @@ +"""Fetch and cache Codex plan usage without blocking the interactive UI.""" + +from __future__ import annotations + +import logging +import threading +import time +from dataclasses import dataclass +from typing import Any + +import requests + +from .config import CHATGPT_OAUTH_CONFIG + +logger = logging.getLogger(__name__) + +_USAGE_URL = "https://chatgpt.com/backend-api/wham/usage" +_CACHE_TTL_SECONDS = 60.0 + + +@dataclass(frozen=True) +class CodexUsage: + """Remaining percentages for Codex's rolling usage windows.""" + + primary_remaining: int | None = None + secondary_remaining: int | None = None + + def format_status(self) -> str: + parts: list[str] = [] + if self.primary_remaining is not None: + parts.append(f"5h {self.primary_remaining}% remaining") + if self.secondary_remaining is not None: + parts.append(f"week {self.secondary_remaining}% remaining") + return " · ".join(parts) + + +_lock = threading.Lock() +_cached_usage: CodexUsage | None = None +_last_attempt = 0.0 +_fetch_in_progress = False + + +def _remaining_percent(window: Any) -> int | None: + if not isinstance(window, dict): + return None + used = window.get("used_percent") + if not isinstance(used, (int, float)): + return None + return max(0, min(100, round(100 - used))) + + +def parse_usage_payload(payload: Any) -> CodexUsage | None: + """Parse the response returned by the Codex ``wham/usage`` endpoint.""" + if not isinstance(payload, dict): + return None + rate_limit = payload.get("rate_limit") + if not isinstance(rate_limit, dict): + return None + usage = CodexUsage( + primary_remaining=_remaining_percent(rate_limit.get("primary_window")), + secondary_remaining=_remaining_percent(rate_limit.get("secondary_window")), + ) + if usage.primary_remaining is None and usage.secondary_remaining is None: + return None + return usage + + +def _fetch(access_token: str, account_id: str) -> None: + global _cached_usage, _fetch_in_progress + headers = { + "Authorization": f"Bearer {access_token}", + "ChatGPT-Account-Id": account_id, + "Accept": "application/json", + "originator": CHATGPT_OAUTH_CONFIG.get("originator", "codex_cli_rs"), + } + try: + response = requests.get(_USAGE_URL, headers=headers, timeout=10) + response.raise_for_status() + usage = parse_usage_payload(response.json()) + if usage is not None: + with _lock: + _cached_usage = usage + except (requests.RequestException, ValueError): + logger.debug("Unable to refresh Codex usage", exc_info=True) + finally: + with _lock: + _fetch_in_progress = False + + +def refresh_usage_in_background(access_token: str, account_id: str) -> None: + """Refresh the cache at most once per minute on a daemon thread.""" + global _fetch_in_progress, _last_attempt + if not access_token or not account_id: + return + now = time.monotonic() + with _lock: + if _fetch_in_progress or now - _last_attempt < _CACHE_TTL_SECONDS: + return + _fetch_in_progress = True + _last_attempt = now + threading.Thread( + target=_fetch, + args=(access_token, account_id), + name="codex-usage-refresh", + daemon=True, + ).start() + + +def get_usage_status() -> str: + """Return cached status text; this function never performs I/O.""" + with _lock: + usage = _cached_usage + return usage.format_status() if usage is not None else "" diff --git a/code_puppy/plugins/chatgpt_oauth/utils.py b/code_puppy/plugins/chatgpt_oauth/utils.py index 56a6a9322..0b7643a36 100644 --- a/code_puppy/plugins/chatgpt_oauth/utils.py +++ b/code_puppy/plugins/chatgpt_oauth/utils.py @@ -349,11 +349,9 @@ def exchange_code_for_tokens( "gpt-5.6-luna", "gpt-5.5", "gpt-5.4", - "gpt-5.3-instant", + "gpt-5.4-mini", "gpt-5.3-codex-spark", - "gpt-5.3-codex", - "gpt-5.2-codex", - "gpt-5.2", + "codex-auto-review", ] # Per-model context length overrides (tokens). diff --git a/code_puppy/plugins/shell_safety/register_callbacks.py b/code_puppy/plugins/shell_safety/register_callbacks.py index e3a99b066..f3a06457e 100644 --- a/code_puppy/plugins/shell_safety/register_callbacks.py +++ b/code_puppy/plugins/shell_safety/register_callbacks.py @@ -22,7 +22,8 @@ # OAuth model prefixes - these models have their own safety mechanisms OAUTH_MODEL_PREFIXES = ( "claude-code-", # Anthropic OAuth - "chatgpt-", # OpenAI OAuth + "codex-", # OpenAI OAuth + "chatgpt-", # Legacy OpenAI OAuth "gemini-oauth", # Google OAuth ) diff --git a/code_puppy/provider_identity.py b/code_puppy/provider_identity.py index 62e86f84a..9b47af22a 100644 --- a/code_puppy/provider_identity.py +++ b/code_puppy/provider_identity.py @@ -53,6 +53,7 @@ def name(self) -> str: _KEY_PREFIX_OVERRIDES = ( ("claude-code-", "claude_code"), + ("codex-", "chatgpt"), ("chatgpt-", "chatgpt"), ("azure-openai-", "azure_openai"), ("openrouter-", "openrouter"), diff --git a/tests/agents/test_compaction.py b/tests/agents/test_compaction.py index 026d1833e..30c4f2948 100644 --- a/tests/agents/test_compaction.py +++ b/tests/agents/test_compaction.py @@ -248,6 +248,25 @@ def test_under_threshold_is_noop(self): assert new_msgs is msgs, "under threshold must return the input unchanged" assert dropped == [] + def test_force_bypasses_threshold(self): + msgs = _build_long_history(n_turns=20) + with patch.multiple( + _compaction, + get_compaction_threshold=lambda: 0.95, + get_compaction_strategy=lambda: "truncation", + get_protected_token_count=lambda: 500, + ): + new_msgs, dropped = compact( + agent=None, + messages=msgs, + model_max=1_000_000, + context_overhead=0, + force=True, + ) + + assert len(new_msgs) < len(msgs) + assert dropped + def test_over_threshold_truncation_strategy(self): msgs = _build_long_history(n_turns=20) with patch.multiple( diff --git a/tests/command_line/test_session_commands.py b/tests/command_line/test_session_commands.py index 3766dbae3..c0976c857 100644 --- a/tests/command_line/test_session_commands.py +++ b/tests/command_line/test_session_commands.py @@ -85,6 +85,21 @@ def _run(self, cmd="/compact"): return handle_compact_command(cmd) + def test_mid_run_queues_compaction_for_next_model_call(self): + controller = MagicMock() + with ( + patch("code_puppy.messaging.run_ui.is_run_active", return_value=True), + patch( + "code_puppy.messaging.pause_controller.get_pause_controller", + return_value=controller, + ), + patch("code_puppy.messaging.emit_info") as emit_info, + ): + assert self._run() is True + + controller.request_compaction.assert_called_once_with() + assert "next model call" in emit_info.call_args[0][0] + def test_no_history(self): agent = MagicMock() agent.get_message_history.return_value = [] diff --git a/tests/command_line/test_tutorial.py b/tests/command_line/test_tutorial.py index b552debb1..743441bf5 100644 --- a/tests/command_line/test_tutorial.py +++ b/tests/command_line/test_tutorial.py @@ -40,7 +40,7 @@ def test_tutorial_chatgpt_flow() -> None: assert result is True mock_reset.assert_called_once() mock_oauth.assert_called_once() - mock_set_model.assert_called_once_with("chatgpt-gpt-5.4") + mock_set_model.assert_called_once_with("codex-gpt-5.6-sol") def test_tutorial_claude_flow() -> None: diff --git a/tests/messaging/spinner/test_spinner_shim.py b/tests/messaging/spinner/test_spinner_shim.py index 6552a66b0..825bcc788 100644 --- a/tests/messaging/spinner/test_spinner_shim.py +++ b/tests/messaging/spinner/test_spinner_shim.py @@ -126,3 +126,29 @@ def test_format_context_info_compact_counts(): def test_format_context_info_zero_or_negative_capacity(): assert format_context_info(100, 0, 0.0) == "" assert format_context_info(100, -5, 0.0) == "" + + +def test_format_context_info_appends_usage_for_codex_model(monkeypatch): + monkeypatch.setattr( + "code_puppy.config.get_global_model_name", lambda: "codex-gpt-5.4" + ) + monkeypatch.setattr( + "code_puppy.plugins.chatgpt_oauth.usage.get_usage_status", + lambda: "5h 66% remaining · week 90% remaining", + ) + + assert format_context_info(5000, 10000, 0.5) == ( + "5k/10k tokens (50%) · Codex 5h 66% remaining · week 90% remaining" + ) + + +def test_format_context_info_hides_codex_usage_for_other_models(monkeypatch): + monkeypatch.setattr( + "code_puppy.config.get_global_model_name", lambda: "openai-gpt-5" + ) + monkeypatch.setattr( + "code_puppy.plugins.chatgpt_oauth.usage.get_usage_status", + lambda: "5h 66% remaining", + ) + + assert format_context_info(5000, 10000, 0.5) == "5k/10k tokens (50%)" diff --git a/tests/messaging/test_pause_controller.py b/tests/messaging/test_pause_controller.py index c33c441db..d14e11a4f 100644 --- a/tests/messaging/test_pause_controller.py +++ b/tests/messaging/test_pause_controller.py @@ -54,6 +54,29 @@ def test_pause_and_resume_are_idempotent(controller): assert controller.is_paused() is False +# ========================================================================= +# Deferred compaction +# ========================================================================= + + +def test_compaction_requests_are_consumed_one_at_a_time(controller): + assert controller.take_compaction_request() is False + controller.request_compaction() + controller.request_compaction() + + assert controller.take_compaction_request() is True + assert controller.take_compaction_request() is True + assert controller.take_compaction_request() is False + + +def test_drain_compaction_requests_clears_all(controller): + controller.request_compaction() + controller.request_compaction() + + assert controller.drain_compaction_requests() == 2 + assert controller.drain_compaction_requests() == 0 + + # ========================================================================= # wait_if_paused # ========================================================================= diff --git a/tests/messaging/test_rich_renderer.py b/tests/messaging/test_rich_renderer.py index c7ce874ba..887aa4eae 100644 --- a/tests/messaging/test_rich_renderer.py +++ b/tests/messaging/test_rich_renderer.py @@ -441,13 +441,13 @@ def test_render_subagent_invocation_continuing(mock_sub, renderer, console): prompt="short", is_new_session=False, message_count=5, - model_name="chatgpt-gpt-5.2", + model_name="codex-gpt-5.2", ) renderer._render_subagent_invocation(msg) out = output(console) assert "Continuing" in out assert "Requested model override:" in out - assert "chatgpt-gpt-5.2" in out + assert "codex-gpt-5.2" in out def test_render_subagent_response(renderer, console): diff --git a/tests/plugins/test_chatgpt_oauth_coverage.py b/tests/plugins/test_chatgpt_oauth_coverage.py index 2d9c9a297..bb0b99c4c 100644 --- a/tests/plugins/test_chatgpt_oauth_coverage.py +++ b/tests/plugins/test_chatgpt_oauth_coverage.py @@ -15,6 +15,9 @@ def test_returns_entries(self): assert "chatgpt-auth" in names assert "chatgpt-status" in names assert "chatgpt-logout" in names + assert "codex-auth" in names + assert "codex-status" in names + assert "codex-logout" in names class TestHandleChatgptStatus: @@ -157,6 +160,7 @@ def test_auth(self): ), ): assert _handle_custom_command("/chatgpt-auth", "chatgpt-auth") is True + assert _handle_custom_command("/codex-auth", "codex-auth") is True def test_status(self): from code_puppy.plugins.chatgpt_oauth.register_callbacks import ( @@ -167,6 +171,7 @@ def test_status(self): "code_puppy.plugins.chatgpt_oauth.register_callbacks._handle_chatgpt_status" ): assert _handle_custom_command("/chatgpt-status", "chatgpt-status") is True + assert _handle_custom_command("/codex-status", "codex-status") is True def test_logout(self): from code_puppy.plugins.chatgpt_oauth.register_callbacks import ( @@ -177,6 +182,7 @@ def test_logout(self): "code_puppy.plugins.chatgpt_oauth.register_callbacks._handle_chatgpt_logout" ): assert _handle_custom_command("/chatgpt-logout", "chatgpt-logout") is True + assert _handle_custom_command("/codex-logout", "codex-logout") is True class TestCreateChatgptOauthModel: diff --git a/tests/plugins/test_chatgpt_oauth_integration.py b/tests/plugins/test_chatgpt_oauth_integration.py index 898a0ed9b..ecb7b32fa 100644 --- a/tests/plugins/test_chatgpt_oauth_integration.py +++ b/tests/plugins/test_chatgpt_oauth_integration.py @@ -47,8 +47,8 @@ def test_load_chatgpt_models_success(self, mock_path, tmp_path): """Test loading ChatGPT models from storage.""" models_file = tmp_path / "models.json" models_data = { - "chatgpt-gpt-5.2": {"type": "chatgpt_oauth", "name": "gpt-5.2"}, - "chatgpt-gpt-5.2-codex": {"type": "chatgpt_oauth", "name": "gpt-5.2-codex"}, + "codex-gpt-5.2": {"type": "chatgpt_oauth", "name": "gpt-5.2"}, + "codex-gpt-5.2-codex": {"type": "chatgpt_oauth", "name": "gpt-5.2-codex"}, } models_file.write_text(json.dumps(models_data)) mock_path.return_value = models_file @@ -72,7 +72,7 @@ def test_save_chatgpt_models_success(self, mock_path, tmp_path): """Test saving ChatGPT models to storage.""" models_file = tmp_path / "models.json" mock_path.return_value = models_file - models_data = {"chatgpt-gpt-5.2": {"type": "chatgpt_oauth"}} + models_data = {"codex-gpt-5.2": {"type": "chatgpt_oauth"}} result = save_chatgpt_models(models_data) @@ -99,20 +99,20 @@ def test_add_models_to_extra_config(self, mock_path, tmp_path): assert result is True models = json.loads(models_file.read_text()) - assert "chatgpt-gpt-5.6" not in models - assert "chatgpt-gpt-5.6-sol" in models - assert "chatgpt-gpt-5.6-terra" in models - assert "chatgpt-gpt-5.6-luna" in models - assert "chatgpt-gpt-5.5" in models - assert "chatgpt-gpt-5.2" in models - assert "chatgpt-gpt-5.2-codex" in models - assert models["chatgpt-gpt-5.6-sol"]["context_length"] == 1050000 - assert models["chatgpt-gpt-5.6-terra"]["context_length"] == 1050000 - assert models["chatgpt-gpt-5.6-luna"]["context_length"] == 1050000 - assert models["chatgpt-gpt-5.6-sol"]["supports_xhigh_reasoning"] is True - assert models["chatgpt-gpt-5.6-sol"]["supports_ultra_reasoning"] is True - assert models["chatgpt-gpt-5.5"]["supports_xhigh_reasoning"] is True - assert models["chatgpt-gpt-5.5"]["supports_ultra_reasoning"] is False + assert "codex-gpt-5.6" not in models + assert "codex-gpt-5.6-sol" in models + assert "codex-gpt-5.6-terra" in models + assert "codex-gpt-5.6-luna" in models + assert "codex-gpt-5.5" in models + assert "codex-gpt-5.2" in models + assert "codex-gpt-5.2-codex" in models + assert models["codex-gpt-5.6-sol"]["context_length"] == 1050000 + assert models["codex-gpt-5.6-terra"]["context_length"] == 1050000 + assert models["codex-gpt-5.6-luna"]["context_length"] == 1050000 + assert models["codex-gpt-5.6-sol"]["supports_xhigh_reasoning"] is True + assert models["codex-gpt-5.6-sol"]["supports_ultra_reasoning"] is True + assert models["codex-gpt-5.5"]["supports_xhigh_reasoning"] is True + assert models["codex-gpt-5.5"]["supports_ultra_reasoning"] is False @patch("code_puppy.plugins.chatgpt_oauth.utils.get_chatgpt_models_path") def test_add_models_with_context_settings(self, mock_path, tmp_path): @@ -123,7 +123,7 @@ def test_add_models_with_context_settings(self, mock_path, tmp_path): add_models_to_extra_config(["gpt-5.2"]) models = json.loads(models_file.read_text()) - model_config = models["chatgpt-gpt-5.2"] + model_config = models["codex-gpt-5.2"] assert model_config["type"] == "chatgpt_oauth" assert ( @@ -148,7 +148,7 @@ def test_add_models_spark_context_length(self, tmp_path): result = add_models_to_extra_config(["gpt-5.3-codex-spark"]) assert result is True models = load_chatgpt_models() - spark_config = models["chatgpt-gpt-5.3-codex-spark"] + spark_config = models["codex-gpt-5.3-codex-spark"] assert spark_config["context_length"] == 131000 assert spark_config["supports_xhigh_reasoning"] is True assert "summary" in spark_config["supported_settings"] @@ -158,7 +158,7 @@ def test_remove_chatgpt_models(self, mock_path, tmp_path): """Test removing ChatGPT models from configuration.""" models_file = tmp_path / "models.json" models_data = { - "chatgpt-gpt-5.2": { + "codex-gpt-5.2": { "type": "chatgpt_oauth", "oauth_source": "chatgpt-oauth-plugin", }, @@ -171,7 +171,7 @@ def test_remove_chatgpt_models(self, mock_path, tmp_path): assert removed == 1 models = json.loads(models_file.read_text()) - assert "chatgpt-gpt-5.2" not in models + assert "codex-gpt-5.2" not in models assert "other-model" in models @patch("requests.get") @@ -309,7 +309,7 @@ def test_handle_custom_command_auth(self, mock_set_model, mock_oauth): assert result is True mock_oauth.assert_called_once() - mock_set_model.assert_called_once_with("chatgpt-gpt-5.4") + mock_set_model.assert_called_once_with("codex-gpt-5.6-sol") @patch("code_puppy.plugins.chatgpt_oauth.register_callbacks.load_stored_tokens") def test_handle_custom_command_status(self, mock_load): diff --git a/tests/plugins/test_chatgpt_oauth_usage.py b/tests/plugins/test_chatgpt_oauth_usage.py new file mode 100644 index 000000000..51aa89191 --- /dev/null +++ b/tests/plugins/test_chatgpt_oauth_usage.py @@ -0,0 +1,47 @@ +from unittest.mock import Mock, patch + +from code_puppy.plugins.chatgpt_oauth import usage + + +def test_parse_usage_payload_formats_remaining_percentages(): + parsed = usage.parse_usage_payload( + { + "rate_limit": { + "primary_window": {"used_percent": 34}, + "secondary_window": {"used_percent": 10.4}, + } + } + ) + + assert parsed is not None + assert parsed.primary_remaining == 66 + assert parsed.secondary_remaining == 90 + assert parsed.format_status() == "5h 66% remaining · week 90% remaining" + + +def test_parse_usage_payload_rejects_missing_windows(): + assert usage.parse_usage_payload({"rate_limit": {}}) is None + assert usage.parse_usage_payload(None) is None + + +def test_refresh_usage_fetches_wham_endpoint_in_background(): + response = Mock() + response.json.return_value = { + "rate_limit": {"primary_window": {"used_percent": 25}} + } + response.raise_for_status.return_value = None + + with patch.object(usage.requests, "get", return_value=response) as get: + usage._fetch("token", "account") + + assert usage.get_usage_status() == "5h 75% remaining" + get.assert_called_once_with( + "https://chatgpt.com/backend-api/wham/usage", + headers={ + "Authorization": "Bearer token", + "ChatGPT-Account-Id": "account", + "Accept": "application/json", + "originator": "codex_cli_rs", + }, + timeout=10, + ) diff --git a/tests/plugins/test_chatgpt_oauth_utils.py b/tests/plugins/test_chatgpt_oauth_utils.py index 8234dba06..28af314c9 100644 --- a/tests/plugins/test_chatgpt_oauth_utils.py +++ b/tests/plugins/test_chatgpt_oauth_utils.py @@ -647,12 +647,12 @@ def test_load_chatgpt_models_success(self, mock_get_path, temp_models_file): mock_get_path.return_value = temp_models_file test_models = { - "chatgpt-gpt-4": { + "codex-gpt-4": { "type": "openai", "name": "gpt-4", "context_length": 8192, }, - "chatgpt-gpt-3.5-turbo": { + "codex-gpt-3.5-turbo": { "type": "openai", "name": "gpt-3.5-turbo", "context_length": 4096, @@ -1034,10 +1034,10 @@ def test_add_models_to_extra_config_success(self, mock_load, mock_save): assert "existing-model" in saved_config # Should contain new models with correct structure - assert "chatgpt-gpt-4" in saved_config - assert "chatgpt-gpt-3.5-turbo" in saved_config + assert "codex-gpt-4" in saved_config + assert "codex-gpt-3.5-turbo" in saved_config - gpt4_config = saved_config["chatgpt-gpt-4"] + gpt4_config = saved_config["codex-gpt-4"] assert gpt4_config["type"] == "chatgpt_oauth" assert gpt4_config["name"] == "gpt-4" assert ( @@ -1061,7 +1061,7 @@ def test_add_models_to_extra_config_success(self, mock_load, mock_save): @patch("code_puppy.plugins.chatgpt_oauth.utils.load_chatgpt_models") def test_add_models_removes_stale_plugin_models(self, mock_load, mock_save): mock_load.return_value = { - "chatgpt-gpt-5.6": { + "codex-gpt-5.6": { "oauth_source": "chatgpt-oauth-plugin", "name": "gpt-5.6", }, @@ -1072,8 +1072,8 @@ def test_add_models_removes_stale_plugin_models(self, mock_load, mock_save): assert add_models_to_extra_config(["gpt-5.6-luna"]) is True saved_config = mock_save.call_args[0][0] - assert "chatgpt-gpt-5.6" not in saved_config - assert "chatgpt-gpt-5.6-luna" in saved_config + assert "codex-gpt-5.6" not in saved_config + assert "codex-gpt-5.6-luna" in saved_config assert "custom-model" in saved_config @patch("code_puppy.plugins.chatgpt_oauth.utils.save_chatgpt_models") @@ -1098,11 +1098,11 @@ def test_add_models_to_extra_config_gpt54_and_newer_support_xhigh( assert result is True saved_config = mock_save.call_args[0][0] for model_name in ( - "chatgpt-gpt-5.6-sol", - "chatgpt-gpt-5.6-terra", - "chatgpt-gpt-5.6-luna", - "chatgpt-gpt-5.5", - "chatgpt-gpt-5.4", + "codex-gpt-5.6-sol", + "codex-gpt-5.6-terra", + "codex-gpt-5.6-luna", + "codex-gpt-5.5", + "codex-gpt-5.4", ): model_config = saved_config[model_name] assert model_config["supported_settings"] == [ @@ -1112,7 +1112,7 @@ def test_add_models_to_extra_config_gpt54_and_newer_support_xhigh( ] assert model_config["supports_xhigh_reasoning"] is True assert model_config["supports_ultra_reasoning"] is model_name.startswith( - "chatgpt-gpt-5.6-" + "codex-gpt-5.6-" ) @patch("code_puppy.plugins.chatgpt_oauth.utils.save_chatgpt_models") @@ -1140,7 +1140,7 @@ def test_add_models_to_extra_config_load_failure(self, mock_load, mock_save): # Should still save the new models mock_save.assert_called_once() saved_config = mock_save.call_args[0][0] - assert "chatgpt-gpt-4" in saved_config + assert "codex-gpt-4" in saved_config class TestRemoveChatGPTModels: @@ -1151,11 +1151,11 @@ class TestRemoveChatGPTModels: def test_remove_chatgpt_models_success(self, mock_load, mock_save): """Test successful removal of ChatGPT models.""" mock_load.return_value = { - "chatgpt-gpt-4": { + "codex-gpt-4": { "name": "gpt-4", "oauth_source": "chatgpt-oauth-plugin", }, - "chatgpt-gpt-3.5-turbo": { + "codex-gpt-3.5-turbo": { "name": "gpt-3.5-turbo", "oauth_source": "chatgpt-oauth-plugin", }, @@ -1175,8 +1175,8 @@ def test_remove_chatgpt_models_success(self, mock_load, mock_save): saved_config = mock_save.call_args[0][0] # Should only contain non-OAuth models - assert "chatgpt-gpt-4" not in saved_config - assert "chatgpt-gpt-3.5-turbo" not in saved_config + assert "codex-gpt-4" not in saved_config + assert "codex-gpt-3.5-turbo" not in saved_config assert "custom-model" in saved_config @patch("code_puppy.plugins.chatgpt_oauth.utils.save_chatgpt_models") @@ -1209,7 +1209,7 @@ def test_remove_chatgpt_models_no_oauth_models(self, mock_load, mock_save): def test_remove_chatgpt_models_save_failure(self, mock_load, mock_save): """Test model removal fails when save fails.""" mock_load.return_value = { - "chatgpt-gpt-4": { + "codex-gpt-4": { "name": "gpt-4", "oauth_source": "chatgpt-oauth-plugin", }, diff --git a/tests/plugins/test_oauth_integration.py b/tests/plugins/test_oauth_integration.py index 0084b6e37..4c26a403e 100644 --- a/tests/plugins/test_oauth_integration.py +++ b/tests/plugins/test_oauth_integration.py @@ -218,7 +218,7 @@ def test_chatgpt_model_registration_flow( # Check that default Codex models were registered for model in DEFAULT_CODEX_MODELS: - prefixed = f"chatgpt-{model}" + prefixed = f"codex-{model}" assert prefixed in saved_models, ( f"Expected {prefixed} in saved models" ) @@ -796,7 +796,7 @@ def test_token_data_integrity_roundtrip( def test_model_config_data_integrity(self, mock_models_storage): """Test model configuration data integrity.""" model_config = { - "chatgpt-gpt-4": { + "codex-gpt-4": { "type": "openai", "name": "gpt-4", "custom_endpoint": { @@ -826,8 +826,8 @@ def test_model_config_data_integrity(self, mock_models_storage): loaded_config = json.load(f) assert loaded_config == model_config - assert "💰" in loaded_config["chatgpt-gpt-4"]["metadata"]["description"] - assert loaded_config["chatgpt-gpt-4"]["metadata"]["special_features"] == [ + assert "💰" in loaded_config["codex-gpt-4"]["metadata"]["description"] + assert loaded_config["codex-gpt-4"]["metadata"]["special_features"] == [ "analysis", "reasoning", "coding", diff --git a/tests/plugins/test_subagent_panel_model_short.py b/tests/plugins/test_subagent_panel_model_short.py index e96cbad52..f710e8c92 100644 --- a/tests/plugins/test_subagent_panel_model_short.py +++ b/tests/plugins/test_subagent_panel_model_short.py @@ -21,7 +21,7 @@ # The reported bug: nano tier must survive, not collapse to "GPT 5.4". ("gpt-5-4-nano", "GPT 5.4-Nano"), ("gpt-5.4-nano", "GPT 5.4-Nano"), - ("chatgpt-gpt-5.4-nano", "GPT 5.4-Nano"), + ("codex-gpt-5.4-nano", "GPT 5.4-Nano"), ("gpt-5.4-mini", "GPT 5.4-Mini"), ("gpt-5.4", "GPT 5.4"), ("gpt-4.1-nano", "GPT 4.1-Nano"), diff --git a/tests/test_chatgpt_codex_client.py b/tests/test_chatgpt_codex_client.py index 15c62a950..b3de6d934 100644 --- a/tests/test_chatgpt_codex_client.py +++ b/tests/test_chatgpt_codex_client.py @@ -547,6 +547,34 @@ async def test_post_request_injects_fields(self): assert body["store"] is False assert body["stream"] is True + @pytest.mark.asyncio + async def test_send_restores_configured_user_agent(self): + """The SDK's per-request User-Agent must not replace the Codex one.""" + success_response = Mock(spec=httpx.Response) + success_response.status_code = 200 + success_response.headers = {"content-type": "application/json"} + + with patch.object( + httpx.AsyncClient, + "send", + new_callable=AsyncMock, + return_value=success_response, + ) as mock_send: + client = ChatGPTCodexAsyncClient( + headers={"User-Agent": "codex_cli_rs/0.144.1"} + ) + request = httpx.Request( + "POST", + "https://chatgpt.com/backend-api/codex/responses", + headers={"User-Agent": "pydantic-ai/1.56.0"}, + content=json.dumps({"model": "gpt-5.6-luna"}).encode(), + ) + + await client.send(request) + + sent_request = mock_send.call_args[0][0] + assert sent_request.headers["User-Agent"] == "codex_cli_rs/0.144.1" + @pytest.mark.asyncio async def test_get_request_passthrough(self): """Test that GET requests pass through unmodified.""" diff --git a/tests/test_completions_and_small_modules.py b/tests/test_completions_and_small_modules.py index 1d3bebc62..57c8448c9 100644 --- a/tests/test_completions_and_small_modules.py +++ b/tests/test_completions_and_small_modules.py @@ -261,6 +261,35 @@ def test_get_effective_agent_model_exception(self): agent.get_model_name.side_effect = Exception("fail") assert _get_effective_agent_model(agent) is None + def test_refresh_context_status_uses_effective_model_capacity(self): + from code_puppy.model_switching import _refresh_context_status + + agent = MagicMock() + agent._get_model_context_length.return_value = 1_050_000 + agent.get_message_history.return_value = ["first", "second"] + agent.estimate_tokens_for_message.side_effect = [20_000, 10_000] + agent._estimate_context_overhead.return_value = 2_000 + + with patch( + "code_puppy.messaging.spinner.update_spinner_context" + ) as update_status: + _refresh_context_status(agent) + + update_status.assert_called_once_with("32k/1.1M tokens (3%)") + + def test_refresh_context_status_clears_stale_value_on_failure(self): + from code_puppy.model_switching import _refresh_context_status + + agent = MagicMock() + agent._get_model_context_length.side_effect = RuntimeError("missing config") + + with patch( + "code_puppy.messaging.spinner.update_spinner_context" + ) as update_status: + _refresh_context_status(agent) + + update_status.assert_called_once_with("") + def _run(self, model_name, agent=None): """Helper to call set_model_and_reload_agent with proper patches.""" from code_puppy.model_switching import set_model_and_reload_agent @@ -295,6 +324,13 @@ def test_refresh_config_called(self): agent.refresh_config.assert_called_once() agent.reload_code_generation_agent.assert_called_once() + def test_reload_refreshes_context_status(self): + agent = MagicMock() + agent.get_model_name.return_value = "model-x" + with patch("code_puppy.model_switching._refresh_context_status") as refresh: + self._run("model-x", agent=agent) + refresh.assert_called_once_with(agent) + def test_refresh_config_exception_nonfatal(self): agent = MagicMock() agent.refresh_config.side_effect = Exception("oops") diff --git a/tests/test_model_factory_coverage.py b/tests/test_model_factory_coverage.py index 3ce4d9820..d0981c833 100644 --- a/tests/test_model_factory_coverage.py +++ b/tests/test_model_factory_coverage.py @@ -1137,7 +1137,7 @@ def test_resolve_provider_identity_precedence(self): == "claude_code" ) assert resolve_provider_identity("openrouter-foo", {}) == "openrouter" - assert resolve_provider_identity("chatgpt-gpt-5", {}) == "chatgpt" + assert resolve_provider_identity("codex-gpt-5", {}) == "chatgpt" assert ( resolve_provider_identity("custom-model", {"type": "custom_openai"}) == "custom_openai" @@ -1466,11 +1466,11 @@ def test_openai_codex_uses_responses_model(self): mock_responses.assert_called_once() def test_custom_openai_chatgpt_codex(self): - """Test chatgpt-gpt-5-codex uses OpenAIResponsesModel.""" + """Test codex-gpt-5-codex uses OpenAIResponsesModel.""" from code_puppy.model_factory import ModelFactory config = { - "chatgpt-gpt-5-codex": { + "codex-gpt-5-codex": { "type": "custom_openai", "name": "gpt-5-codex", "custom_endpoint": { @@ -1484,7 +1484,7 @@ def test_custom_openai_chatgpt_codex(self): with patch( "code_puppy.model_factory.OpenAIResponsesModel" ) as mock_responses: - ModelFactory.get_model("chatgpt-gpt-5-codex", config) + ModelFactory.get_model("codex-gpt-5-codex", config) mock_responses.assert_called_once() diff --git a/tests/test_tools_registration.py b/tests/test_tools_registration.py index 7488b7f9c..15481bff2 100644 --- a/tests/test_tools_registration.py +++ b/tests/test_tools_registration.py @@ -180,7 +180,7 @@ def test_legacy_reasoning_tool_can_be_registered_without_warning( register_tools_for_agent( mock_agent, ["list_files", "agent_share_your_reasoning"], - model_name="chatgpt-gpt-5.4", + model_name="codex-gpt-5.4", ) mock_warning.assert_not_called() From 651a2f74ee4ce472813ac0e8b2e5f58a83dd6a91 Mon Sep 17 00:00:00 2001 From: mpfaffenberger Date: Fri, 10 Jul 2026 08:04:56 -0700 Subject: [PATCH 08/78] feat: improve queued turns and shell backgrounding --- code_puppy/agents/_run_signals.py | 18 ++--- code_puppy/agents/event_stream_handler.py | 4 +- code_puppy/config.py | 4 +- code_puppy/messaging/__init__.py | 2 + code_puppy/messaging/line_editor.py | 13 ++- code_puppy/messaging/message_queue.py | 6 ++ code_puppy/messaging/pause_controller.py | 21 +++++ code_puppy/messaging/renderers.py | 12 ++- code_puppy/messaging/run_ui.py | 31 +++++-- code_puppy/plugins/fork/register_callbacks.py | 61 +++++++++++--- code_puppy/tools/command_runner.py | 42 +++++----- code_puppy/tools/subagent_invocation.py | 5 +- tests/agents/test_event_stream_handler.py | 8 +- tests/agents/test_runtime_pause_leakage.py | 55 ++++++------- tests/messaging/test_line_editor.py | 14 ++-- tests/messaging/test_pause_controller.py | 15 ++++ tests/messaging/test_renderers.py | 15 ++++ tests/messaging/test_run_ui_persistent.py | 13 +++ tests/plugins/test_fork_plugin.py | 80 +++++++++++++++++-- tests/test_model_factory.py | 40 +++++----- tests/tools/test_shell_backgrounding.py | 31 +++++++ 21 files changed, 373 insertions(+), 117 deletions(-) diff --git a/code_puppy/agents/_run_signals.py b/code_puppy/agents/_run_signals.py index 64b00d6e0..5e68cba3e 100644 --- a/code_puppy/agents/_run_signals.py +++ b/code_puppy/agents/_run_signals.py @@ -105,9 +105,8 @@ def schedule_agent_cancel(force: bool = False) -> None: # ============================================================================= # # ``PauseController`` is a process-wide singleton. Without explicit hygiene: -# - A cancelled run that left ``request_steer`` items in the queue would -# have those items consumed by the NEXT agent run (possibly a totally -# different session / agent), as if the user had typed them. +# - A ``now`` steer that missed the final model boundary would linger into +# the next run instead of becoming the queued turn the user still expects. # - A run that crashed mid-pause would leave the controller in a paused # state, freezing the next run's spinner + event stream. # Both bugs are bad. The two helpers below scrub that state. @@ -117,9 +116,8 @@ def reset_pause_state_at_run_start() -> None: """Scrub stale ``PauseController`` state before a fresh agent run. Called from the top of ``run_with_mcp`` BEFORE any agent work begins. - If we find pending steers from a prior cancelled run, emit a warning - (with a preview of the first one) so the user knows we discarded - something rather than silently swallowing it. + Undelivered ``now`` steers are preserved as queued turns: they missed + their intended history-processor boundary, but they are still user input. """ from code_puppy.messaging.pause_controller import get_pause_controller @@ -132,10 +130,10 @@ def reset_pause_state_at_run_start() -> None: emit_warning( f"Discarded {stale_compactions} stale compaction request(s) from a previous run." ) - stale_steers = pc.drain_pending_steer() - if stale_steers: - emit_warning( - f"Discarded {len(stale_steers)} stale steering message(s) from a previous run." + deferred_steers = pc.defer_pending_steer_now() + if deferred_steers: + emit_info( + f"Queued {deferred_steers} steering message(s) that missed the previous run." ) diff --git a/code_puppy/agents/event_stream_handler.py b/code_puppy/agents/event_stream_handler.py index bcd237a57..c8ff9d9ea 100644 --- a/code_puppy/agents/event_stream_handler.py +++ b/code_puppy/agents/event_stream_handler.py @@ -238,12 +238,12 @@ async def _print_thinking_banner() -> None: # Clear any \r-repainted progress line, then move below it erase_progress_line(console) console.print() # Newline before banner - # Bold banner with configurable color and lightning bolt + # Bold banner with configurable color. thinking_color = get_banner_color("thinking") console.print( Text.from_markup( - f"[bold white on {thinking_color}] THINKING [/bold white on {thinking_color}] [dim]\u26a1 " + f"[bold white on {thinking_color}] THINKING [/bold white on {thinking_color}] " ), end="", ) diff --git a/code_puppy/config.py b/code_puppy/config.py index 0d640a96b..76fdca6f3 100644 --- a/code_puppy/config.py +++ b/code_puppy/config.py @@ -1421,8 +1421,8 @@ def get_message_limit(default: int = 1000) -> int: def get_command_timeout_seconds() -> int: """ - Returns the user-configured absolute timeout for shell commands in seconds. - This is the maximum time any single command can run before being terminated. + Returns the user-configured foreground limit for shell commands in seconds. + Commands still running at the limit are automatically backgrounded, not killed. Defaults to 270 seconds if unset or misconfigured. Valid range: 60-900 seconds. Values outside this range default to 270. Configurable by 'command_timeout_seconds' key. diff --git a/code_puppy/messaging/__init__.py b/code_puppy/messaging/__init__.py index d01d62545..b87e3c9f3 100644 --- a/code_puppy/messaging/__init__.py +++ b/code_puppy/messaging/__init__.py @@ -93,6 +93,7 @@ emit_message, emit_planned_next_steps, emit_prompt, + emit_queued, emit_success, emit_system_message, emit_tool_output, @@ -187,6 +188,7 @@ # Legacy emit functions "emit_message", "emit_info", + "emit_queued", "emit_success", "emit_warning", "emit_divider", diff --git a/code_puppy/messaging/line_editor.py b/code_puppy/messaging/line_editor.py index a0a728abd..58a3c6019 100644 --- a/code_puppy/messaging/line_editor.py +++ b/code_puppy/messaging/line_editor.py @@ -68,6 +68,10 @@ SubmitRouter = Callable[[str, str], Optional[str]] +class _QueuedFeedback(str): + """Marker for feedback that belongs in the dedicated queued renderer.""" + + class RunningLineEditor: """Line editor fed by the key-listener thread. @@ -616,7 +620,7 @@ def route_default(self, text: str, mode: str) -> Optional[str]: return None if mode == "queue": # Queued steers get no later confirmation, so ack at submit time. - return f"⏭ queued for next turn: {stripped[:60]}" + return _QueuedFeedback(f"for next turn: {stripped[:60]}") # "now" steers: stay quiet here. The steer history processor emits # "Injecting steer mid-turn — model will see: ..." when the text # actually reaches the model; acking at submit time too was a lie @@ -627,9 +631,12 @@ def route_default(self, text: str, mode: str) -> Optional[str]: def _emit_feedback(note: str) -> None: """Best-effort transcript line for a successful steer submission.""" try: - from code_puppy.messaging.message_queue import emit_info + from code_puppy.messaging.message_queue import emit_info, emit_queued - emit_info(note) + if isinstance(note, _QueuedFeedback): + emit_queued(str(note)) + else: + emit_info(note) except Exception: logger.debug("feedback emit failed", exc_info=True) diff --git a/code_puppy/messaging/message_queue.py b/code_puppy/messaging/message_queue.py index 7d85ef922..fb2e3f2f6 100644 --- a/code_puppy/messaging/message_queue.py +++ b/code_puppy/messaging/message_queue.py @@ -23,6 +23,7 @@ class MessageType(Enum): # Basic content types INFO = "info" + QUEUED = "queued" SUCCESS = "success" WARNING = "warning" ERROR = "error" @@ -315,6 +316,11 @@ def emit_info(content: Any, **metadata): emit_message(MessageType.INFO, content, **metadata) +def emit_queued(content: Any, **metadata): + """Emit a queued-for-next-turn acknowledgement.""" + emit_message(MessageType.QUEUED, content, **metadata) + + def emit_success(content: Any, **metadata): """Emit a success message.""" emit_message(MessageType.SUCCESS, content, **metadata) diff --git a/code_puppy/messaging/pause_controller.py b/code_puppy/messaging/pause_controller.py index d9b8a0c79..4f47657a1 100644 --- a/code_puppy/messaging/pause_controller.py +++ b/code_puppy/messaging/pause_controller.py @@ -289,6 +289,27 @@ def drain_pending_steer_queued(self) -> List[str]: self._fire_steer_queue_listeners(total) return drained + def defer_pending_steer_now(self) -> int: + """Atomically move undelivered ``now`` steers to the queued queue. + + A ``/steer`` can race with the final model call: once that call has + finished there is no next history-processor boundary at which to + inject it. Preserving it as a normal queued turn is less surprising + than silently discarding the user's message. + + Returns the number of messages moved. Existing queued messages keep + priority because they were already waiting for a future turn. + """ + with self._lock: + moved = len(self._steer_queue_now) + if not moved: + return 0 + self._steer_queue_queued.extend(self._steer_queue_now) + self._steer_queue_now = [] + total = len(self._steer_queue_queued) + self._fire_steer_queue_listeners(total) + return moved + def pop_next_steer_queued(self) -> Optional[str]: """Atomically pop the OLDEST queued-mode steer (None when empty). diff --git a/code_puppy/messaging/renderers.py b/code_puppy/messaging/renderers.py index 7626b94bc..982baa6b9 100644 --- a/code_puppy/messaging/renderers.py +++ b/code_puppy/messaging/renderers.py @@ -289,7 +289,17 @@ def _print_message(console: Console, message: UIMessage) -> None: style = _classify_style(message) content = message.content - if isinstance(content, str): + if message.type == MessageType.QUEUED: + from code_puppy.config import get_banner_color + + queued = Text( + " QUEUED ", + style=f"bold white on {get_banner_color('thinking')}", + ) + queued.append(" ") + queued.append(str(content), style="dim") + console.print(queued) + elif isinstance(content, str): if message.type == MessageType.AGENT_RESPONSE: try: console.print(Markdown(content)) diff --git a/code_puppy/messaging/run_ui.py b/code_puppy/messaging/run_ui.py index 4f95bb44c..c6451b77a 100644 --- a/code_puppy/messaging/run_ui.py +++ b/code_puppy/messaging/run_ui.py @@ -149,9 +149,11 @@ def stop_run_ui() -> None: """Tear down the run UI: unhook the editor, restore the terminal. Idempotent; never raises (runs in ``finally`` blocks on cancel and - exception paths). + exception paths). Any ``/steer`` message that missed the run's final + model-call boundary is converted to a queued turn instead of discarded. """ global _editor, _loop, _run_active + persistent_run_ended = False with _lock: if _persistent: # The UI outlives the run — just drop back to idle routing. @@ -161,10 +163,15 @@ def stop_run_ui() -> None: # thread crashed, or a per-run listener replaced-then-stopped # it), typing at the idle prompt would be dead forever. _ensure_persistent_listener_locked() - return - editor = _editor - _editor = None - _loop = None + persistent_run_ended = True + editor = None + else: + editor = _editor + _editor = None + _loop = None + _defer_undelivered_steers() + if persistent_run_ended: + return if editor is not None: _set_feed_target(None) unregister_chord("\x05") # the handler closes over the dead editor @@ -176,6 +183,20 @@ def stop_run_ui() -> None: logger.debug("bottom bar stop failed", exc_info=True) +def _defer_undelivered_steers() -> None: + """Preserve ``/steer`` input that arrived after the last model call.""" + try: + from .pause_controller import get_pause_controller + + moved = get_pause_controller().defer_pending_steer_now() + if moved: + from . import emit_info + + emit_info(f"⏭ Queued {moved} steering message(s) for the next turn.") + except Exception: + logger.debug("undelivered steer deferral failed", exc_info=True) + + def _clear_status_row() -> None: """Run over (finished OR cancelled): drop the token/context line. diff --git a/code_puppy/plugins/fork/register_callbacks.py b/code_puppy/plugins/fork/register_callbacks.py index f9529ee51..4db40abc3 100644 --- a/code_puppy/plugins/fork/register_callbacks.py +++ b/code_puppy/plugins/fork/register_callbacks.py @@ -12,10 +12,11 @@ /fork cancel cancel a running fork /forks show status of all forks this session -Results: ``_invoke_agent_impl`` already emits a ``SubAgentResponseMessage`` -which the renderer displays as a full "AGENT RESPONSE" markdown panel, so -this plugin only adds a compact completion banner (id, agent, duration, -session) rather than re-printing the response. +Results are explicitly bridged onto the interactive transcript queue when +the fork completes. Forks suppress ``_invoke_agent_impl``'s generic structured +response and emit one fork-specific, theme-aware banner plus Markdown body; +detached tasks cannot rely on the invocation renderer still owning the active +transcript. The completion banner then adds duration and resumable session id. Note: Ctrl+C's cancel sweep clears ``_active_subagent_tasks``, which includes forked runs — cancelling the agent takes forks down with it. @@ -69,6 +70,12 @@ def _emit_success(content) -> None: emit_success(content) +def _emit_agent_response(content) -> None: + from code_puppy.messaging import emit_agent_response + + emit_agent_response(content) + + def _emit_warning(content) -> None: from code_puppy.messaging import emit_warning @@ -81,6 +88,40 @@ def _emit_error(content) -> None: emit_error(content) +def _started_message(fork_id: int, agent_name: str): + """Build a Rich, theme-aware acknowledgement for a newly started fork.""" + from rich.text import Text + + from code_puppy.messaging.messages import MessageLevel + from code_puppy.messaging.rich_renderer import DEFAULT_STYLES + + message = Text(f"fork #{fork_id}: ", style=DEFAULT_STYLES[MessageLevel.INFO]) + message.append(agent_name, style="bold") + message.append(" started in the background — results land when it finishes (") + message.append("/forks", style=DEFAULT_STYLES[MessageLevel.DEBUG]) + message.append(" for status)") + return message + + +def _response_message(fork_id: int, agent_name: str, response: str): + """Build the distinct, theme-aware banner and Markdown body for a fork.""" + from rich.console import Group + from rich.markdown import Markdown + from rich.text import Text + + from code_puppy.config import get_banner_color + from code_puppy.messaging.messages import MessageLevel + from code_puppy.messaging.rich_renderer import DEFAULT_STYLES + + header = Text( + f" FORK #{fork_id} RESPONSE ", + style=f"bold white on {get_banner_color('subagent_response')}", + ) + header.append(" ") + header.append(agent_name, style=f"bold {DEFAULT_STYLES[MessageLevel.INFO]}") + return Group(header, Markdown(response)) + + # --------------------------------------------------------------------------- # Parsing # --------------------------------------------------------------------------- @@ -131,6 +172,7 @@ async def _run_fork(agent_name: str, prompt: str): context=None, agent_name=agent_name, prompt=prompt, + emit_response_message=False, ) @@ -160,7 +202,11 @@ def _on_fork_done(fork_id: int, task: asyncio.Task) -> None: _emit_error(f"{tag} failed after {record.elapsed:.1f}s: {first_line}") return record.status = "done" - # The response itself was already rendered via SubAgentResponseMessage. + response = getattr(result, "response", None) + if response: + _emit_agent_response( + _response_message(fork_id, record.agent_name, response) + ) _emit_success( f"{tag} finished in {record.elapsed:.1f}s" + (f" — session '{record.session_id}'" if record.session_id else "") @@ -182,10 +228,7 @@ def _start_fork(agent_name: str, prompt: str) -> None: fork_id=fork_id, agent_name=agent_name, prompt=prompt, task=task ) task.add_done_callback(lambda t, fid=fork_id: _on_fork_done(fid, t)) - _emit_info( - f"fork #{fork_id}: [bold cyan]{agent_name}[/bold cyan] started in the " - f"background — results land when it finishes ([dim]/forks[/dim] for status)" - ) + _emit_info(_started_message(fork_id, agent_name)) def _cancel_fork(raw_id: str) -> None: diff --git a/code_puppy/tools/command_runner.py b/code_puppy/tools/command_runner.py index 066cde15c..43cb40e20 100644 --- a/code_puppy/tools/command_runner.py +++ b/code_puppy/tools/command_runner.py @@ -685,10 +685,11 @@ def run_shell_command_streaming( start_time = time.time() last_output_time = [start_time] - # Get the user-configured absolute timeout for shell commands + # Foreground duration limit. Reaching it detaches rather than killing; + # inactivity remains the guard for genuinely wedged commands. from code_puppy.config import get_command_timeout_seconds - ABSOLUTE_TIMEOUT_SECONDS = get_command_timeout_seconds() + foreground_limit_seconds = get_command_timeout_seconds() stdout_lines = [] stderr_lines = [] @@ -911,8 +912,8 @@ def nuclear_kill(proc): } ) - def detach_to_background(): - """Ctrl+X Ctrl+B: stop streaming, keep the process running. + def detach_to_background(*, automatic: bool = False): + """Stop foreground streaming while keeping the process running. The reader threads stay alive and divert every further line into a log file (pipes keep draining -- a full pipe buffer would @@ -929,9 +930,14 @@ def detach_to_background(): target=close_divert_log_on_exit, args=(process, log), daemon=True ).start() execution_time = time.time() - start_time + cause = ( + f"Automatically backgrounded after {foreground_limit_seconds}s" + if automatic + else "The user backgrounded this command" + ) if not silent: emit_warning( - f"Backgrounded (PID {process.pid}) -- output continues in {log.path}" + f"{cause} (PID {process.pid}) -- output continues in {log.path}" ) return ShellCommandOutput( success=True, @@ -945,15 +951,14 @@ def detach_to_background(): log_file=log.path, pid=process.pid, user_feedback=( - "The user backgrounded this command (Ctrl+X Ctrl+B) while it" - " was still running. It has NOT finished: stdout/stderr above" - " are partial, exit_code is unknown, and the process" - f" (PID {process.pid}) keeps running with further output" - f" appended to {log.path} (an exit-code footer is written" - " when it finishes). The user chose to stop waiting -- do NOT" - " wait, sleep, poll, or re-run the command. Acknowledge the" - " backgrounding and move on immediately; only read that log" - " file later if a subsequent task genuinely needs the result." + f"{cause} while the command was still running. It has NOT" + " finished: stdout/stderr above are partial, exit_code is" + f" unknown, and the process (PID {process.pid}) keeps running" + f" with further output appended to {log.path} (an exit-code" + " footer is written when it finishes). Do NOT wait, sleep," + " poll, or re-run the command. Move on immediately; only read" + " that log later if a subsequent task genuinely needs the" + " result." ), ) @@ -970,13 +975,8 @@ def detach_to_background(): if background_generation() != bg_generation_at_start: return detach_to_background() - if current_time - start_time > ABSOLUTE_TIMEOUT_SECONDS: - if not silent: - emit_error( - "Process killed: absolute timeout reached", - message_group=group_id, - ) - return cleanup_process_and_threads("absolute") + if current_time - start_time > foreground_limit_seconds: + return detach_to_background(automatic=True) if current_time - last_output_time[0] > timeout: if not silent: diff --git a/code_puppy/tools/subagent_invocation.py b/code_puppy/tools/subagent_invocation.py index ec298ead5..3f9a3c06a 100644 --- a/code_puppy/tools/subagent_invocation.py +++ b/code_puppy/tools/subagent_invocation.py @@ -46,8 +46,9 @@ async def _invoke_agent_impl( prompt: str, session_id: str | None = None, model_name: str | None = None, + emit_response_message: bool = True, ) -> AgentInvokeOutput: - """Invoke a sub-agent, optionally with an explicit temporary model override.""" + """Invoke a sub-agent, optionally suppressing its standard response message.""" from code_puppy.agents.agent_manager import load_agent # Validate user-provided session_id if given @@ -343,7 +344,7 @@ async def _invoke_agent_impl( # In high mode, skip the emit when streaming already rendered the # response to avoid a double-render if any future subscriber # starts rendering SubAgentResponseMessage. - if not (is_high_mode and streamed_text): + if emit_response_message and not (is_high_mode and streamed_text): bus.emit( SubAgentResponseMessage( agent_name=agent_name, diff --git a/tests/agents/test_event_stream_handler.py b/tests/agents/test_event_stream_handler.py index c9dfd314f..416262b87 100644 --- a/tests/agents/test_event_stream_handler.py +++ b/tests/agents/test_event_stream_handler.py @@ -184,8 +184,14 @@ async def event_stream(): ): await event_stream_handler(mock_ctx, event_stream()) - # Should print the initial content + # The banner and initial content should print without a redundant icon. assert console.print.called + thinking_banner = next( + call.args[0] + for call in console.print.call_args_list + if call.args and "THINKING" in str(call.args[0]) + ) + assert chr(0x26A1) not in str(thinking_banner) @pytest.mark.asyncio async def test_handles_text_part_with_initial_content(self, mock_ctx): diff --git a/tests/agents/test_runtime_pause_leakage.py b/tests/agents/test_runtime_pause_leakage.py index 04344d7dc..de50555ed 100644 --- a/tests/agents/test_runtime_pause_leakage.py +++ b/tests/agents/test_runtime_pause_leakage.py @@ -1,9 +1,8 @@ -"""Tests for cross-run PauseController leakage protection. +"""Tests for cross-run PauseController state handling. -The ``PauseController`` is a process-wide singleton. Without explicit -hygiene at run start + on cancel, a Ctrl+C'd run can leave stale steering -messages in the queue that get silently consumed by the NEXT (potentially -totally different) agent run. These tests lock the hygiene down. +The process-wide controller preserves late ``/steer`` input as queued turns, +while explicit cancellation still discards undelivered input. These tests lock +both lifecycle contracts down. """ from __future__ import annotations @@ -94,46 +93,38 @@ def _isolated_runtime(monkeypatch: pytest.MonkeyPatch): # ============================================================================= -# Bug A.1 — stale state is scrubbed at run start +# Bug A.1 — late steering is deferred; stale pause state is scrubbed # ============================================================================= @pytest.mark.asyncio -async def test_stale_steer_drained_at_run_start(_isolated_runtime, monkeypatch): - """A stale steer left over from a prior cancelled run MUST NOT be - consumed as if the user had typed it in this run. - """ - warnings: List[str] = [] - # ``reset_pause_state_at_run_start`` lives in _run_signals and calls - # ``emit_warning`` imported there. Patch the imported reference. +async def test_stale_now_steer_becomes_queued_turn_at_run_start( + _isolated_runtime, monkeypatch +): + """A ``/steer`` that missed the prior run is preserved as a queued turn.""" + infos: List[str] = [] monkeypatch.setattr( - "code_puppy.agents._run_signals.emit_warning", - lambda msg, *_a, **_k: warnings.append(msg), + "code_puppy.agents._run_signals.emit_info", + lambda msg, *_a, **_k: infos.append(msg), ) - only = _DummyResult("solo") - pydantic_agent = _ScriptedPydanticAgent(only) + first = _DummyResult("fresh-result") + after_steer = _DummyResult("steered-result") + pydantic_agent = _ScriptedPydanticAgent(first, after_steer) agent = _DummyAgent(pydantic_agent) - # Simulate the leak: someone's previous (cancelled) run queued a steer. - get_pause_controller().request_steer("stale msg from previous session") + # Simulate /steer landing after the previous run's final model call. + get_pause_controller().request_steer("late steer from previous run") result = await _runtime.run_with_mcp(agent, "fresh prompt") - # 1. The fresh run should have received the ORIGINAL prompt, not the steer. - assert len(pydantic_agent.calls) == 1, ( - "Stale steer must NOT trigger a follow-up turn" - ) - assert pydantic_agent.calls[0]["prompt"] == "fresh prompt" - assert result is only - - # 2. Pause queue must be empty post-run. + assert [call["prompt"] for call in pydantic_agent.calls] == [ + "fresh prompt", + "late steer from previous run", + ] + assert result is after_steer assert get_pause_controller().drain_pending_steer() == [] - - # 3. A warning must have fired so the user knows we scrubbed something. - assert any("stale steering message" in w for w in warnings), ( - f"expected stale-steer warning, got: {warnings!r}" - ) + assert any("missed the previous run" in message for message in infos) @pytest.mark.asyncio diff --git a/tests/messaging/test_line_editor.py b/tests/messaging/test_line_editor.py index a1a99f807..a258b6f3f 100644 --- a/tests/messaging/test_line_editor.py +++ b/tests/messaging/test_line_editor.py @@ -406,7 +406,11 @@ def emitted(monkeypatch): lines = [] monkeypatch.setattr( "code_puppy.messaging.message_queue.emit_info", - lambda text, **kwargs: lines.append(text), + lambda text, **kwargs: lines.append(("info", text)), + ) + monkeypatch.setattr( + "code_puppy.messaging.message_queue.emit_queued", + lambda text, **kwargs: lines.append(("queued", text)), ) return lines @@ -416,7 +420,7 @@ def test_enter_submit_emits_queued_feedback(editor, emitted): # submit time (there's no later confirmation for them). feed_all(editor, "fix the tests") editor.feed("\r") - assert emitted == ["⏭ queued for next turn: fix the tests"] + assert emitted == [("queued", "for next turn: fix the tests")] def test_steer_command_emits_no_feedback(editor, emitted, controller): @@ -433,7 +437,7 @@ def test_steer_command_emits_no_feedback(editor, emitted, controller): def test_bare_steer_command_emits_usage(editor, emitted, controller): feed_all(editor, "/steer") editor.feed("\r") - assert emitted == ["Usage: /steer "] + assert emitted == [("info", "Usage: /steer ")] assert controller.steers == [] assert editor.get_pending_command() is None @@ -442,14 +446,14 @@ def test_queue_submit_emits_queued_feedback(editor, emitted): feed_all(editor, "do this after") editor.feed("\x1b") editor.feed("\r") - assert emitted == ["⏭ queued for next turn: do this after"] + assert emitted == [("queued", "for next turn: do this after")] def test_queue_feedback_preview_truncated_to_60_chars(editor, emitted): feed_all(editor, "x" * 100) editor.feed("\x1b") editor.feed("\r") - assert emitted == [f"⏭ queued for next turn: {'x' * 60}"] + assert emitted == [("queued", f"for next turn: {'x' * 60}")] def test_slash_command_emits_no_feedback(editor, emitted): diff --git a/tests/messaging/test_pause_controller.py b/tests/messaging/test_pause_controller.py index d14e11a4f..845d77da7 100644 --- a/tests/messaging/test_pause_controller.py +++ b/tests/messaging/test_pause_controller.py @@ -217,6 +217,21 @@ def test_drain_queued_does_not_touch_now(controller): assert controller.drain_pending_steer_now() == ["now-one"] +def test_defer_pending_now_moves_messages_after_existing_queued(controller): + controller.request_steer("now-one", mode="now") + controller.request_steer("queue-one", mode="queue") + controller.request_steer("now-two", mode="now") + + assert controller.defer_pending_steer_now() == 2 + assert controller.drain_pending_steer_now() == [] + assert controller.drain_pending_steer_queued() == [ + "queue-one", + "now-one", + "now-two", + ] + assert controller.defer_pending_steer_now() == 0 + + def test_drain_pending_steer_combines_both_queues_queued_first(controller): """Backwards-compat: the cancel-path uses drain_pending_steer() to wipe everything. Queued first matches the order the runtime would have diff --git a/tests/messaging/test_renderers.py b/tests/messaging/test_renderers.py index 2f8258e0c..5f2b3435c 100644 --- a/tests/messaging/test_renderers.py +++ b/tests/messaging/test_renderers.py @@ -234,6 +234,7 @@ def test_sync_renderer_render_messages(mq): (MessageType.ERROR, "err"), (MessageType.WARNING, "warn"), (MessageType.SUCCESS, "ok"), + (MessageType.QUEUED, "for next turn: later"), (MessageType.TOOL_OUTPUT, "tool"), (MessageType.AGENT_REASONING, "think"), (MessageType.AGENT_RESPONSE, "**bold**"), @@ -246,6 +247,20 @@ def test_sync_renderer_render_messages(mq): assert "err" in output +def test_sync_renderer_queued_banner(mq): + console = make_console() + renderer = SynchronousInteractiveRenderer(mq, console=console) + + renderer._render_message( + UIMessage(type=MessageType.QUEUED, content="for next turn: fix the tests") + ) + + output = console.file.getvalue() + assert "QUEUED" in output + assert "for next turn: fix the tests" in output + assert chr(0x23ED) not in output + + def test_sync_renderer_version_dim(mq): console = make_console() r = SynchronousInteractiveRenderer(mq, console=console) diff --git a/tests/messaging/test_run_ui_persistent.py b/tests/messaging/test_run_ui_persistent.py index 144e6bef7..7fac33b0d 100644 --- a/tests/messaging/test_run_ui_persistent.py +++ b/tests/messaging/test_run_ui_persistent.py @@ -162,6 +162,19 @@ async def test_running_steer_command_injects_now(): assert editor.get_pending_command() is None +async def test_run_end_defers_undelivered_steer_to_queue(): + install_tty_bar() + run_ui_mod.start_persistent_ui() + run_ui_mod.start_run_ui() + pc = get_pause_controller() + pc.request_steer("too late to inject", mode="now") + + run_ui_mod.stop_run_ui() + + assert pc.drain_pending_steer_now() == [] + assert pc.drain_pending_steer_queued() == ["too late to inject"] + + async def test_running_slash_goes_to_drain_queue(): install_tty_bar() run_ui_mod.start_persistent_ui() diff --git a/tests/plugins/test_fork_plugin.py b/tests/plugins/test_fork_plugin.py index 1c14b5439..fa86ab577 100644 --- a/tests/plugins/test_fork_plugin.py +++ b/tests/plugins/test_fork_plugin.py @@ -3,11 +3,15 @@ from __future__ import annotations import asyncio +from io import StringIO from types import SimpleNamespace from unittest.mock import patch import pytest +from rich.console import Console, Group +from rich.text import Text +from code_puppy.messaging.messages import MessageLevel from code_puppy.plugins.fork import register_callbacks as rc _AGENTS = {"code-puppy": "the default pup", "qa-kitten": "meow"} @@ -39,7 +43,15 @@ def fake_agent_manager(): def _fake_impl(response="did the thing", error=None, session_id="sess-abc123"): - async def impl(context, agent_name, prompt, session_id_=None, model_name=None): + async def impl( + context, + agent_name, + prompt, + session_id_=None, + model_name=None, + emit_response_message=True, + ): + assert emit_response_message is False return SimpleNamespace( response=response, agent_name=agent_name, @@ -140,14 +152,50 @@ def test_fork_cancel_unknown_id_warns(): # ========================================================================= -async def test_fork_success_emits_completion_banner(): - infos, successes = [], [] +def test_started_message_is_rich_and_theme_aware(): + from code_puppy.messaging.rich_renderer import DEFAULT_STYLES + + with patch.dict( + DEFAULT_STYLES, + {MessageLevel.INFO: "#123456", MessageLevel.DEBUG: "dim #654321"}, + ): + message = rc._started_message(7, "qa-kitten") + + assert isinstance(message, Text) + assert message.plain == ( + "fork #7: qa-kitten started in the background — results land when it " + "finishes (/forks for status)" + ) + assert "[bold" not in message.plain + assert message.style == "#123456" + assert [span.style for span in message.spans] == ["bold", "dim #654321"] + + +def test_response_message_identifies_fork_and_renders_markdown(): + output = StringIO() + console = Console(file=output, force_terminal=False, width=100) + + with patch("code_puppy.config.get_banner_color", return_value="dark_orange3"): + message = rc._response_message(7, "qa-kitten", "**did the thing**") + console.print(message) + + assert isinstance(message, Group) + rendered = output.getvalue() + assert "FORK #7 RESPONSE" in rendered + assert "qa-kitten" in rendered + assert "did the thing" in rendered + assert "AGENT RESPONSE" not in rendered + + +async def test_fork_success_emits_response_and_completion_banner(): + infos, responses, successes = [], [], [] with ( patch( "code_puppy.tools.subagent_invocation._invoke_agent_impl", new=_fake_impl(), ), patch.object(rc, "_emit_info", infos.append), + patch.object(rc, "_emit_agent_response", responses.append), patch.object(rc, "_emit_success", successes.append), ): assert rc._handle_fork("/fork @qa-kitten test everything") is True @@ -158,7 +206,11 @@ async def test_fork_success_emits_completion_banner(): assert record.status == "done" assert record.agent_name == "qa-kitten" assert record.session_id == "sess-abc123" - assert any("started in the background" in str(m) for m in infos) + assert any( + isinstance(m, Text) and "started in the background" in m.plain for m in infos + ) + assert len(responses) == 1 + assert isinstance(responses[0], Group) assert any("finished" in str(m) for m in successes) @@ -200,7 +252,15 @@ async def test_fork_reports_result_error_as_failure(): async def test_fork_reports_crash_as_failure(): - async def boom(context, agent_name, prompt, session_id=None, model_name=None): + async def boom( + context, + agent_name, + prompt, + session_id=None, + model_name=None, + emit_response_message=True, + ): + assert emit_response_message is False raise RuntimeError("kaboom") errors = [] @@ -220,7 +280,15 @@ async def boom(context, agent_name, prompt, session_id=None, model_name=None): async def test_fork_cancel_running_fork(): started = asyncio.Event() - async def slow(context, agent_name, prompt, session_id=None, model_name=None): + async def slow( + context, + agent_name, + prompt, + session_id=None, + model_name=None, + emit_response_message=True, + ): + assert emit_response_message is False started.set() await asyncio.sleep(60) diff --git a/tests/test_model_factory.py b/tests/test_model_factory.py index 0265be36f..62867790e 100644 --- a/tests/test_model_factory.py +++ b/tests/test_model_factory.py @@ -309,44 +309,48 @@ def test_extra_models_json_decode_error(tmp_path, monkeypatch): # Create a temporary extra_models.json file with invalid JSON extra_models_file = tmp_path / "extra_models.json" extra_models_file.write_text("{ invalid json content }") + base_config = {"base-model": {"type": "openai", "name": "base"}} - # Patch the EXTRA_MODELS_FILE path to point to our temporary file - from code_puppy.model_factory import ModelFactory - + # Use an explicit base config: bundled models.json may intentionally be empty. + monkeypatch.setattr( + "code_puppy.model_factory.callbacks.get_callbacks", lambda phase: [object()] + ) + monkeypatch.setattr( + "code_puppy.model_factory.callbacks.on_load_model_config", + lambda: [base_config.copy()], + ) monkeypatch.setattr( "code_puppy.model_factory.EXTRA_MODELS_FILE", str(extra_models_file) ) - # This should not raise an exception despite the invalid JSON + # Invalid extra JSON should be ignored without discarding the base config. config = ModelFactory.load_config() - # The config should still be loaded, just without the extra models - assert isinstance(config, dict) - assert len(config) > 0 + assert config["base-model"] == base_config["base-model"] def test_extra_models_exception_handling(tmp_path, monkeypatch, caplog): - # Create a temporary extra_models.json file that will raise a general exception + # Create a directory where a JSON file is expected to force an OSError. extra_models_file = tmp_path / "extra_models.json" - # Create a directory with the same name to cause an OSError when trying to read it extra_models_file.mkdir() + base_config = {"base-model": {"type": "openai", "name": "base"}} - # Patch the EXTRA_MODELS_FILE path - from code_puppy.model_factory import ModelFactory - + # Use an explicit base config: bundled models.json may intentionally be empty. + monkeypatch.setattr( + "code_puppy.model_factory.callbacks.get_callbacks", lambda phase: [object()] + ) + monkeypatch.setattr( + "code_puppy.model_factory.callbacks.on_load_model_config", + lambda: [base_config.copy()], + ) monkeypatch.setattr( "code_puppy.model_factory.EXTRA_MODELS_FILE", str(extra_models_file) ) - # This should not raise an exception despite the error with caplog.at_level("WARNING"): config = ModelFactory.load_config() - # The config should still be loaded - assert isinstance(config, dict) - assert len(config) > 0 - - # Check that warning was logged + assert config["base-model"] == base_config["base-model"] assert "Failed to load extra models config" in caplog.text diff --git a/tests/tools/test_shell_backgrounding.py b/tests/tools/test_shell_backgrounding.py index 5a9f838b9..c8f74e797 100644 --- a/tests/tools/test_shell_backgrounding.py +++ b/tests/tools/test_shell_backgrounding.py @@ -81,6 +81,37 @@ def test_background_request_detaches_streaming_command(): pass +def test_foreground_limit_auto_backgrounds_instead_of_killing(monkeypatch): + process = _spawn_slow_process() + command_runner._register_process(process) + result = None + monkeypatch.setattr("code_puppy.config.get_command_timeout_seconds", lambda: 0) + + try: + result = run_shell_command_streaming( + process, timeout=30, command="long test suite", silent=True + ) + + assert result.background is True + assert result.success is True + assert result.timeout is False + assert result.exit_code is None + assert result.pid == process.pid + assert result.log_file and os.path.exists(result.log_file) + assert "Automatically backgrounded after 0s" in result.user_feedback + assert process.poll() is None + with command_runner._RUNNING_PROCESSES_LOCK: + assert process not in command_runner._RUNNING_PROCESSES + finally: + command_runner._kill_process_group(process) + if result and result.log_file: + time.sleep(0.2) + try: + os.unlink(result.log_file) + except OSError: + pass + + def test_full_tool_path_returns_promptly_on_background(): """End to end through _execute_shell_command (executor + keyboard context): the tool call must return within a couple of poll ticks of From 521460bd5378e61ea6f2076effe37f85f6f1132a Mon Sep 17 00:00:00 2001 From: mpfaffenberger Date: Fri, 10 Jul 2026 08:14:50 -0700 Subject: [PATCH 09/78] test: include queued emitter in messaging exports --- tests/test_messaging_init.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_messaging_init.py b/tests/test_messaging_init.py index 35318b874..1cf2dfe82 100644 --- a/tests/test_messaging_init.py +++ b/tests/test_messaging_init.py @@ -32,6 +32,7 @@ def test_emit_functions_exported(self): emit_functions = [ "emit_message", "emit_info", + "emit_queued", "emit_success", "emit_warning", "emit_divider", @@ -90,6 +91,7 @@ def test_expected_export_count(self): "get_global_queue", "emit_message", "emit_info", + "emit_queued", "emit_success", "emit_warning", "emit_divider", From 460c821c0f60449c828a3cb3aae8932863303eb1 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Fri, 10 Jul 2026 15:21:54 +0000 Subject: [PATCH 10/78] chore: bump version [ci skip] --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 55e091da9..a813f0ce8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "code-puppy" -version = "0.0.622" +version = "0.0.623" description = "Code generation agent" readme = "README.md" requires-python = ">=3.11,<3.15" diff --git a/uv.lock b/uv.lock index 5c13de52c..d337c7984 100644 --- a/uv.lock +++ b/uv.lock @@ -312,7 +312,7 @@ wheels = [ [[package]] name = "code-puppy" -version = "0.0.622" +version = "0.0.623" source = { editable = "." } dependencies = [ { name = "agent-client-protocol" }, From bda3b66d59b904457c3c3a7fa955269104217071 Mon Sep 17 00:00:00 2001 From: Julien Ellie Date: Fri, 10 Jul 2026 08:29:27 -0700 Subject: [PATCH 11/78] fix(retry): retry gateway 502s wrapped in ExceptionGroup instead of crashing (#546) A gateway 502 from custom_anthropic (e.g. puppy-backend) surfaces as an anthropic.APIStatusError(502) raised mid-stream. pydantic-ai runs the model stream inside an anyio task group, so it reaches run_agent_task wrapped in an ExceptionGroup (hence the except* handler). should_retry_streaming only walked __cause__/__context__ and never descended into ExceptionGroup.exceptions, so the retryable 5xx looked opaque -> no retry -> 60-line REPL traceback. - should_retry_streaming now traverses ExceptionGroup members too (cycle-safe, per-chain depth cap preserved via a separate traversal axis). - OpenAI branch of _is_retryable_one now retries on status_code 429/>=500, mirroring the Anthropic and ModelHTTPError branches (HTML-body 5xx match no snippet, so status_code is the only reliable signal). - Add 'bad gateway'/'gateway timeout' snippets as belt-and-braces. - Tests: ExceptionGroup (nested, cause-chain member, cycle-safe, non-transient) and OpenAI status-code retry/no-retry coverage. Co-authored-by: Julien Ellie --- code_puppy/agents/_runtime.py | 56 +++++++++++++- tests/agents/test_streaming_retry.py | 107 +++++++++++++++++++++++++++ 2 files changed, 160 insertions(+), 3 deletions(-) diff --git a/code_puppy/agents/_runtime.py b/code_puppy/agents/_runtime.py index 99b21782e..70f9b5e8c 100644 --- a/code_puppy/agents/_runtime.py +++ b/code_puppy/agents/_runtime.py @@ -113,6 +113,8 @@ "rate limited", "overloaded", "service unavailable", + "bad gateway", + "gateway timeout", "server had an error processing your request", "retry your request", "internal server error", @@ -194,6 +196,16 @@ def _is_retryable_one(exc: BaseException) -> bool: return _matches_retryable_snippet(msg) if OpenAIAPIError is not None and isinstance(exc, OpenAIAPIError): + # 5xx gateway/server errors (502/503/504) and 429 rate limits are + # transient regardless of message wording. The OpenAI SDK exposes the + # HTTP status on APIStatusError subclasses; APIConnectionError / + # APITimeoutError have no status_code and are already covered by the + # transport-exception branch above. This mirrors the ModelHTTPError and + # Anthropic branches so an OpenAI-compatible 502 gets the same slow + # 1-2-3 retry instead of a raw REPL traceback. + status_code = getattr(exc, "status_code", None) + if status_code == 429 or (isinstance(status_code, int) and status_code >= 500): + return True if _matches_retryable_snippet(msg): return True body = getattr(exc, "body", None) @@ -255,16 +267,54 @@ def _is_retryable_one(exc: BaseException) -> bool: return False +def _group_members(exc: BaseException) -> tuple: + """Return an ExceptionGroup's sub-exceptions, or ``()`` for a plain error. + + Guards against the 3.10 fallback where ``BaseExceptionGroup`` aliases + ``Exception`` -- there we'd otherwise treat *every* exception as a group. + On 3.11+ (and 3.10 with the backport) real groups expose ``.exceptions``. + """ + if BaseExceptionGroup is Exception: # 3.10 without exceptiongroup backport + return () + if isinstance(exc, BaseExceptionGroup): + return tuple(exc.exceptions) + return () + + def should_retry_streaming(exc: Exception) -> bool: """Decide whether ``exc`` (or anything it wraps) is a transient hiccup. Walks the ``__cause__`` / ``__context__`` chain so wrapper exception types like :class:`pydantic_ai.exceptions.ModelAPIError` -- which hide the real httpx/anthropic transport error in ``__cause__`` -- still get - classified correctly. Returns ``True`` if *any* link in the chain is - transient. + classified correctly. + + Also descends into :class:`ExceptionGroup` members. pydantic-ai runs the + model stream inside an anyio task group, so a transient provider error + (e.g. ``anthropic.APIStatusError`` HTTP 502 from a gateway) reaches us + wrapped in an ``ExceptionGroup`` -- which is why ``run_agent_task`` catches + it with ``except*``. The linear cause walker alone can't see inside + ``.exceptions``, so a retryable 5xx used to look opaque and crash the REPL + with a full traceback instead of getting the slow 1-2-3 retry. + + Returns ``True`` if *any* reachable exception is transient. Cycle-safe via + an ``id`` set; the per-chain depth cap of :func:`_walk_cause_chain` is + preserved (group membership is a separate traversal axis). """ - return any(_is_retryable_one(e) for e in _walk_cause_chain(exc)) + seen: set[int] = set() + stack: list[BaseException] = [exc] + while stack: + node = stack.pop() + if node is None or id(node) in seen: + continue + for link in _walk_cause_chain(node): + if id(link) in seen: + continue + seen.add(id(link)) + if _is_retryable_one(link): + return True + stack.extend(_group_members(link)) + return False def streaming_retry( diff --git a/tests/agents/test_streaming_retry.py b/tests/agents/test_streaming_retry.py index d05f32459..8f91bb46b 100644 --- a/tests/agents/test_streaming_retry.py +++ b/tests/agents/test_streaming_retry.py @@ -433,3 +433,110 @@ def test_cause_chain_depth_is_bounded(self): # If this assertion ever flips, the cap moved -- intentional or not, # the change deserves to be looked at deliberately. assert not should_retry_streaming_exception(chain[0]) + + +def _make_anthropic_502() -> "BaseException": + """Reproduce the production 502: a Google-gateway HTML error page bubbled + up by the Anthropic SDK as APIStatusError(status_code=502). + + The message is a giant HTML blob (``[HTTP 502] ...``) that + matches NO retry snippet -- the only reliable signal is ``status_code``. + """ + from anthropic import APIStatusError + + err = APIStatusError.__new__(APIStatusError) + err.status_code = 502 + err.body = {"message": "[HTTP 502] ... 502. That's an error."} + err.message = "[HTTP 502] ... 502. That's an error." + return err + + +class TestExceptionGroupUnwrapping: + """Regression coverage for the real-world 502 crash. + + pydantic-ai streams the model response inside an anyio task group, so a + transient provider error (e.g. anthropic.APIStatusError HTTP 502 from a + gateway) reaches ``run_agent_task`` wrapped in an ExceptionGroup -- which + is why that code path uses ``except*``. The classifier used to only walk + ``__cause__``/``__context__`` and never descended into ``.exceptions``, so + a perfectly retryable 5xx looked opaque and crashed the REPL with a + 60-line traceback instead of getting the slow 1-2-3 retry. + """ + + def test_bare_anthropic_502_is_retryable(self): + # Sanity: even outside a group, an HTML-body 502 must retry purely on + # status_code (its message matches no snippet). + assert should_retry_streaming_exception(_make_anthropic_502()) + + def test_exception_group_wrapping_502_is_retryable(self): + # The exact production shape: ExceptionGroup([APIStatusError(502)]). + group = ExceptionGroup( + "unhandled errors in a TaskGroup", [_make_anthropic_502()] + ) + assert should_retry_streaming_exception(group) + + def test_nested_exception_groups_are_retryable(self): + # anyio can nest groups when tasks spawn sub-tasks. Descend all the way. + inner = ExceptionGroup("inner", [_make_anthropic_502()]) + outer = ExceptionGroup("outer", [ValueError("noise"), inner]) + assert should_retry_streaming_exception(outer) + + def test_exception_group_member_via_cause_chain(self): + # A group member that hides its transient origin in __cause__ must still + # be caught -- we walk each member's cause chain, not just the member. + wrapper = RuntimeError("opaque wrapper") + wrapper.__cause__ = httpx.ConnectError("dropped socket") + group = ExceptionGroup("grp", [ValueError("noise"), wrapper]) + assert should_retry_streaming_exception(group) + + def test_exception_group_of_only_non_transient_is_not_retryable(self): + # Belt-and-braces: a group of genuine bugs must NOT be retried. + group = ExceptionGroup("grp", [ValueError("bug"), TypeError("also bug")]) + assert not should_retry_streaming_exception(group) + + def test_exception_group_is_cycle_safe(self): + # A member whose cause points back at a sibling must still terminate. + a = ValueError("a") + b = httpx.ConnectError("transient") + a.__cause__ = b + b.__cause__ = a + group = ExceptionGroup("grp", [a]) + assert should_retry_streaming_exception(group) + + +class TestOpenAIStatusCodeRetry: + """The OpenAI branch must retry 5xx/429 on status_code alone. + + Mirrors the Anthropic and ModelHTTPError branches. An OpenAI-compatible + gateway 502 whose body is an HTML error page matches no snippet -- the only + reliable signal is the HTTP status the SDK exposes on APIStatusError. + """ + + def _make_openai_status_error(self, status_code: int, message: str | None = None): + try: + from openai import APIStatusError + except ImportError: # pragma: no cover + pytest.skip("openai is not installed in this test environment") + if message is None: + message = " 502 Bad Gateway" + err = APIStatusError.__new__(APIStatusError) + err.status_code = status_code + err.body = {"message": message} + err.message = message + return err + + @pytest.mark.parametrize("status_code", [500, 502, 503, 504, 429]) + def test_transient_status_codes_retry(self, status_code): + assert should_retry_streaming_exception( + self._make_openai_status_error(status_code) + ) + + @pytest.mark.parametrize("status_code", [400, 401, 403, 404, 422]) + def test_client_errors_do_not_retry(self, status_code): + # A snippet-free client-error body -- status_code is the only signal, + # and 4xx (other than 429) must NOT retry. + assert not should_retry_streaming_exception( + self._make_openai_status_error( + status_code, message="invalid request payload" + ) + ) From 6e60cd1af86af25533d0aa8736576870ad9f7366 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Fri, 10 Jul 2026 15:36:20 +0000 Subject: [PATCH 12/78] chore: bump version [ci skip] --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a813f0ce8..ec28614ed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "code-puppy" -version = "0.0.623" +version = "0.0.624" description = "Code generation agent" readme = "README.md" requires-python = ">=3.11,<3.15" diff --git a/uv.lock b/uv.lock index d337c7984..1ac88909f 100644 --- a/uv.lock +++ b/uv.lock @@ -312,7 +312,7 @@ wheels = [ [[package]] name = "code-puppy" -version = "0.0.623" +version = "0.0.624" source = { editable = "." } dependencies = [ { name = "agent-client-protocol" }, From ff3fdfe3a18cdcbbf08613c3b22c05a28800b9f3 Mon Sep 17 00:00:00 2001 From: Ashish Singhi Date: Fri, 10 Jul 2026 21:51:11 +0530 Subject: [PATCH 13/78] feat(callbacks): let run_shell_command hooks rewrite the command (#545) Extend the existing run_shell_command callback contract so a plugin can transparently transform a command before execution, in addition to the existing allow / block semantics. A callback may now return {"rewrite": ""} (optionally with a "rewrite_reason" string for the UI) alongside the existing {"blocked": True} response. Rewrites are applied in callback registration order and always emit a visible info line so users can see what actually ran -- rewriting must never be sneaky. Why: * Unlocks a class of plugins that cannot exist today: secret redaction in commands, corporate proxy / mirror injection, path fixup, attribution trailers on git commits, etc. * Fully backward compatible: callbacks that return None or {"blocked": True} keep working unchanged. The new behaviour only fires when a callback explicitly returns a "rewrite" key with a non-empty string value. Change is intentionally scoped down to just the infrastructure. Any policy-carrying plugin (attribution trailers, redaction rules, ...) belongs outside core -- either in a user plugin, a project plugin, or an external marketplace, depending on how opinionated it is. Tests: * tests/tools/test_command_runner_full_coverage.py -- three new cases covering the happy path (rewrite is applied and executed), a no-op rewrite where the string is unchanged (no notice emitted), and defensive handling of empty / non-string rewrite payloads. * All 259 pre-existing tests across command_runner, shell_safety, destructive_command_detector, shell_passthrough, and shell_line_crlf still pass -- backward compatibility verified. Docs: * AGENTS.md: document the new "rewrite" key alongside "blocked" in the run_shell_command hook row. Co-authored-by: a0s01yh --- AGENTS.md | 2 +- code_puppy/tools/command_runner.py | 17 +++ .../test_command_runner_full_coverage.py | 109 ++++++++++++++++++ 3 files changed, 127 insertions(+), 1 deletion(-) diff --git a/AGENTS.md b/AGENTS.md index ef4e53758..3a33b9b7a 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -71,7 +71,7 @@ and skills (`/.code_puppy/skills/`). | `agent_run_start` | Before agent task | `(agent_name, model_name, session_id=None) -> None` | | `agent_run_end` | After agent run | `(agent_name, model_name, session_id=None, success=True, error=None, response_text=None, metadata=None) -> None` | | `load_prompt` | System prompt assembly | `() -> str \| None` | -| `run_shell_command` | Before shell exec | `(context, command, cwd=None, timeout=60) -> dict \| None` (return `{"blocked": True}` to block) | +| `run_shell_command` | Before shell exec | `(context, command, cwd=None, timeout=60) -> dict \| None` (return `{"blocked": True}` to block, `{"rewrite": ""}` to transparently transform) | | `file_permission` | Before file op | `(context, file_path, operation, ...) -> bool` | | `pre_tool_call` | Before tool executes | `(tool_name, tool_args, context=None) -> Any` | | `post_tool_call` | After tool finishes | `(tool_name, tool_args, result, duration_ms, context=None) -> Any` | diff --git a/code_puppy/tools/command_runner.py b/code_puppy/tools/command_runner.py index 43cb40e20..2da8d6a4c 100644 --- a/code_puppy/tools/command_runner.py +++ b/code_puppy/tools/command_runner.py @@ -1096,6 +1096,23 @@ async def run_shell_command( execution_time=None, ) + # Apply any command rewrites requested by callbacks. + # A callback can return {"rewrite": ""} to transparently + # transform the command before execution (e.g. inject a git trailer, + # redact a secret, prepend a corporate proxy). Rewrites are applied in + # callback registration order; each one sees whatever the previous ones + # produced. A visible info line surfaces the change so users always know + # what actually ran -- rewriting must never be sneaky. + for result in callback_results: + if result and isinstance(result, dict) and "rewrite" in result: + new_command = result["rewrite"] + if not isinstance(new_command, str) or not new_command.strip(): + continue # ignore empty / non-string rewrites defensively + if new_command != command: + reason = result.get("rewrite_reason", "plugin") + emit_info(f"[dim]\U0001f527 command rewritten by {reason}[/dim]") + command = new_command + # Handle background execution - runs command detached and returns immediately # This happens BEFORE user confirmation since we don't wait for the command if background: diff --git a/tests/tools/test_command_runner_full_coverage.py b/tests/tools/test_command_runner_full_coverage.py index 5f7416f3d..a7614dc72 100644 --- a/tests/tools/test_command_runner_full_coverage.py +++ b/tests/tools/test_command_runner_full_coverage.py @@ -544,6 +544,115 @@ async def test_blocked_by_callback(self): assert result.success is False assert "nope" in result.error + @pytest.mark.asyncio + async def test_rewrite_from_callback_replaces_command(self): + """A callback that returns {"rewrite": ...} transforms the command. + + Proves the new hook contract: infrastructure for plugins that want + to transparently transform commands (secret redaction, proxy + injection, attribution trailers) instead of merely blocking them. + """ + from code_puppy.tools.command_runner import run_shell_command + + ctx = MagicMock(spec=RunContext) + rewrite_result = {"rewrite": "echo rewritten", "rewrite_reason": "unit test"} + mock_output = MagicMock() + mock_output.success = True + + with patch( + "code_puppy.callbacks.on_run_shell_command", + new_callable=AsyncMock, + return_value=[rewrite_result], + ): + with patch("code_puppy.config.get_yolo_mode", return_value=True): + with patch( + "code_puppy.tools.command_runner.is_subagent", + return_value=False, + ): + with patch( + "code_puppy.tools.command_runner._execute_shell_command", + new_callable=AsyncMock, + return_value=mock_output, + ) as mock_exec: + with patch("code_puppy.tools.command_runner.emit_info"): + result = await run_shell_command( + ctx, "echo original", timeout=10 + ) + + assert result.success is True + # The rewritten command is what actually got executed. + executed = mock_exec.call_args.kwargs["command"] + assert executed == "echo rewritten", ( + f"expected rewritten command to be executed, got {executed!r}" + ) + + @pytest.mark.asyncio + async def test_rewrite_ignored_when_same_as_original(self): + """A no-op rewrite should not emit the visible info notice.""" + from code_puppy.tools.command_runner import run_shell_command + + ctx = MagicMock(spec=RunContext) + noop_result = {"rewrite": "echo same"} + mock_output = MagicMock() + mock_output.success = True + + with patch( + "code_puppy.callbacks.on_run_shell_command", + new_callable=AsyncMock, + return_value=[noop_result], + ): + with patch("code_puppy.config.get_yolo_mode", return_value=True): + with patch( + "code_puppy.tools.command_runner.is_subagent", + return_value=False, + ): + with patch( + "code_puppy.tools.command_runner._execute_shell_command", + new_callable=AsyncMock, + return_value=mock_output, + ): + with patch( + "code_puppy.tools.command_runner.emit_info" + ) as mock_info: + await run_shell_command(ctx, "echo same", timeout=10) + + # No 'command rewritten' notice should fire when the string is identical. + for call in mock_info.call_args_list: + args = call.args or () + payload = args[0] if args else "" + assert "command rewritten" not in str(payload) + + @pytest.mark.asyncio + async def test_rewrite_ignored_for_empty_or_non_string(self): + """Defensively skip empty / non-string rewrite payloads.""" + from code_puppy.tools.command_runner import run_shell_command + + ctx = MagicMock(spec=RunContext) + bad_results = [{"rewrite": ""}, {"rewrite": None}, {"rewrite": 42}] + mock_output = MagicMock() + mock_output.success = True + + with patch( + "code_puppy.callbacks.on_run_shell_command", + new_callable=AsyncMock, + return_value=bad_results, + ): + with patch("code_puppy.config.get_yolo_mode", return_value=True): + with patch( + "code_puppy.tools.command_runner.is_subagent", + return_value=False, + ): + with patch( + "code_puppy.tools.command_runner._execute_shell_command", + new_callable=AsyncMock, + return_value=mock_output, + ) as mock_exec: + with patch("code_puppy.tools.command_runner.emit_info"): + await run_shell_command(ctx, "echo keep", timeout=10) + + # Original command flows through unmutated. + assert mock_exec.call_args.kwargs["command"] == "echo keep" + @pytest.mark.asyncio async def test_yolo_mode_executes(self): from code_puppy.tools.command_runner import run_shell_command From 52095fd8ca4c1dc5684ea2499fee724bf5330c60 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Fri, 10 Jul 2026 16:28:35 +0000 Subject: [PATCH 14/78] chore: bump version [ci skip] --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ec28614ed..4104cae62 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "code-puppy" -version = "0.0.624" +version = "0.0.625" description = "Code generation agent" readme = "README.md" requires-python = ">=3.11,<3.15" diff --git a/uv.lock b/uv.lock index 1ac88909f..e91bfd815 100644 --- a/uv.lock +++ b/uv.lock @@ -312,7 +312,7 @@ wheels = [ [[package]] name = "code-puppy" -version = "0.0.624" +version = "0.0.625" source = { editable = "." } dependencies = [ { name = "agent-client-protocol" }, From db62a4db881464f81476df5b2071c8415eea1151 Mon Sep 17 00:00:00 2001 From: mpfaffenberger Date: Fri, 10 Jul 2026 09:58:21 -0700 Subject: [PATCH 15/78] feat: calm OAuth pages and refine queued banner --- code_puppy/messaging/renderers.py | 4 +- code_puppy/plugins/oauth_puppy_html.py | 394 +++++++++++-------------- tests/messaging/test_renderers.py | 23 +- tests/plugins/test_oauth_puppy_html.py | 55 ++++ 4 files changed, 259 insertions(+), 217 deletions(-) create mode 100644 tests/plugins/test_oauth_puppy_html.py diff --git a/code_puppy/messaging/renderers.py b/code_puppy/messaging/renderers.py index 982baa6b9..6055afe06 100644 --- a/code_puppy/messaging/renderers.py +++ b/code_puppy/messaging/renderers.py @@ -292,12 +292,14 @@ def _print_message(console: Console, message: UIMessage) -> None: if message.type == MessageType.QUEUED: from code_puppy.config import get_banner_color - queued = Text( + queued = Text() + queued.append( " QUEUED ", style=f"bold white on {get_banner_color('thinking')}", ) queued.append(" ") queued.append(str(content), style="dim") + console.print() console.print(queued) elif isinstance(content, str): if message.type == MessageType.AGENT_RESPONSE: diff --git a/code_puppy/plugins/oauth_puppy_html.py b/code_puppy/plugins/oauth_puppy_html.py index 45194e64e..4ee7231f0 100644 --- a/code_puppy/plugins/oauth_puppy_html.py +++ b/code_puppy/plugins/oauth_puppy_html.py @@ -1,228 +1,194 @@ -"""Shared HTML templates drenched in ridiculous puppy-fueled OAuth theatrics.""" +"""Small, self-contained HTML pages for browser-based OAuth callbacks.""" from __future__ import annotations -from typing import Optional, Tuple - -CLAUDE_LOGO_URL = "https://voideditor.com/claude-icon.png" -CHATGPT_LOGO_URL = ( - "https://freelogopng.com/images/all_img/1681038325chatgpt-logo-transparent.png" -) -GEMINI_LOGO_URL = "https://upload.wikimedia.org/wikipedia/commons/thumb/8/8a/Google_Gemini_logo.svg/512px-Google_Gemini_logo.svg.png" +from html import escape +from typing import Optional + +_PAGE_STYLES = """ +:root { + color-scheme: light dark; + --bg: #f5f3ef; + --card: rgba(255, 255, 255, 0.92); + --border: rgba(28, 25, 23, 0.10); + --text: #1c1917; + --muted: #6b645f; + --success: #16845b; + --success-soft: #e7f6ef; + --failure: #c2413a; + --failure-soft: #fcebea; + --shadow: 0 24px 70px rgba(41, 37, 36, 0.12); +} +* { box-sizing: border-box; } +body { + min-height: 100vh; + margin: 0; + display: grid; + place-items: center; + padding: 24px; + background: + radial-gradient(circle at 20% 10%, rgba(249, 115, 22, 0.08), transparent 30%), + var(--bg); + color: var(--text); + font-family: Inter, ui-sans-serif, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; +} +.card { + width: min(100%, 480px); + padding: 40px; + border: 1px solid var(--border); + border-radius: 24px; + background: var(--card); + box-shadow: var(--shadow); + text-align: center; +} +.mark { + width: 64px; + height: 64px; + margin: 0 auto 24px; + display: grid; + place-items: center; + border-radius: 18px; + background: #1c1917; + color: #fff; + font-size: 30px; + box-shadow: 0 10px 24px rgba(28, 25, 23, 0.16); +} +.eyebrow { + margin: 0 0 10px; + color: var(--muted); + font-size: 13px; + font-weight: 700; + letter-spacing: 0.11em; + text-transform: uppercase; +} +h1 { + margin: 0; + font-size: clamp(28px, 6vw, 38px); + line-height: 1.12; + letter-spacing: -0.035em; +} +.message { + max-width: 370px; + margin: 16px auto 0; + color: var(--muted); + font-size: 16px; + line-height: 1.6; +} +.status { + display: inline-flex; + align-items: center; + gap: 8px; + margin-top: 28px; + padding: 9px 13px; + border-radius: 999px; + font-size: 13px; + font-weight: 700; +} +.status::before { + content: ""; + width: 8px; + height: 8px; + border-radius: 50%; + background: currentColor; +} +.success .status { color: var(--success); background: var(--success-soft); } +.failure .status { color: var(--failure); background: var(--failure-soft); } +.hint { + margin: 24px 0 0; + color: var(--muted); + font-size: 13px; +} +@media (prefers-color-scheme: dark) { + :root { + --bg: #121110; + --card: rgba(29, 27, 25, 0.94); + --border: rgba(255, 255, 255, 0.09); + --text: #f7f5f2; + --muted: #aaa39d; + --success: #73d5ad; + --success-soft: rgba(22, 132, 91, 0.18); + --failure: #f19a94; + --failure-soft: rgba(194, 65, 58, 0.18); + --shadow: 0 24px 70px rgba(0, 0, 0, 0.34); + } + .mark { background: #f7f5f2; color: #1c1917; } +} +@media (max-width: 520px) { + .card { padding: 32px 24px; border-radius: 20px; } +} +""" def oauth_success_html(service_name: str, extra_message: Optional[str] = None) -> str: - """Return an over-the-top puppy celebration HTML page with artillery effects.""" - clean_service = service_name.strip() or "OAuth" - detail = f"

🐾 {extra_message} 🐾

" if extra_message else "" - projectile, rival_url, rival_alt, target_modifier = _service_targets(clean_service) - target_classes = "target" if not target_modifier else f"target {target_modifier}" - return ( - "" - "" - "Puppy Paw-ty Success" - "" - "" - "
" - "
" - + "".join( - f"{emoji}" - for left, top, delay, emoji in _SUCCESS_PUPPIES - ) - + "
" - f"

🐶⚡ {clean_service} OAuth Complete ⚡🐶

" - "

Puppy squad delivered the token payload without mercy.

" - f"{detail}" - f"

💣 Puppies are bombarding the {rival_alt} defenses! 💣

" - "

🚀 This window will auto-close faster than a corgi zoomie. 🚀

" - "

Keep the artillery firing – the rivals never stood a chance.

" - f"
{rival_alt}
" - "
" + _build_artillery(projectile) + "
" - "
" - "" - "" + """Return a restrained success page for an OAuth callback.""" + service = _clean(service_name, "OAuth") + message = _clean( + extra_message, + "Authentication is complete. You can return to Code Puppy.", + ) + return _render_page( + service=service, + state="success", + title="You're connected", + message=message, + status="Authentication complete", + hint="This window can be closed safely.", + auto_close=True, ) def oauth_failure_html(service_name: str, reason: str) -> str: - """Return a dramatic puppy-tragedy HTML page for OAuth sadness.""" - clean_service = service_name.strip() or "OAuth" - clean_reason = reason.strip() or "Something went wrong with the treats" - projectile, rival_url, rival_alt, target_modifier = _service_targets(clean_service) - target_classes = "target" if not target_modifier else f"target {target_modifier}" - return ( - "" - "" - "Puppy Tears" - "" - "" - "
" - "
" - + "".join( - f"{emoji}" - for left, top, delay, emoji in _FAILURE_PUPPIES - ) - + "
" - f"

💔🐶 {clean_service} OAuth Whoopsie 💔

" - "

😭 Puppy artillery jammed! Someone cut the firing wire.

" - f"

{clean_reason}

" - "

💧 A thousand doggy eyes are welling up. Try again from Code Puppy! 💧

" - f"

Re-calibrate the {projectile} barrage and slam it into the {rival_alt} wall.

" - "
" - "Adopt more puppies" - "Send treats" - "
" - "
" - + _build_artillery(projectile, shells_only=True) - + f"
{rival_alt}
" - + "
" - "
" - "" + """Return a clear, low-drama failure page for an OAuth callback.""" + service = _clean(service_name, "OAuth") + message = _clean(reason, "Authentication could not be completed.") + return _render_page( + service=service, + state="failure", + title="Couldn't connect", + message=message, + status="Authentication incomplete", + hint="Return to Code Puppy and try signing in again.", ) -_SUCCESS_PUPPIES = ( - (5, 12, 0.0, "🐶"), - (18, 28, 0.2, "🐕"), - (32, 6, 1.1, "🐩"), - (46, 18, 0.5, "🦮"), - (62, 9, 0.8, "🐕‍🦺"), - (76, 22, 1.3, "🐶"), - (88, 14, 0.4, "🐺"), - (12, 48, 0.6, "🐕"), - (28, 58, 1.7, "🦴"), - (44, 42, 0.9, "🦮"), - (58, 52, 1.5, "🐾"), - (72, 46, 0.3, "🐩"), - (86, 54, 1.1, "🐕‍🦺"), - (8, 72, 0.7, "🐶"), - (24, 80, 1.2, "🐩"), - (40, 74, 0.2, "🐕"), - (56, 66, 1.6, "🦮"), - (70, 78, 1.0, "🐕‍🦺"), - (84, 70, 1.4, "🐾"), - (16, 90, 0.5, "🐶"), - (32, 92, 1.9, "🦴"), - (48, 88, 1.1, "🐺"), - (64, 94, 1.8, "🐩"), - (78, 88, 0.6, "🐕"), - (90, 82, 1.3, "🐾"), -) - - -_FAILURE_PUPPIES = ( - (8, 6, 0.0, "🥺🐶"), - (22, 18, 0.3, "😢🐕"), - (36, 10, 0.6, "😿🐩"), - (50, 20, 0.9, "😭🦮"), - (64, 8, 1.2, "🥺🐕‍🦺"), - (78, 16, 1.5, "😢🐶"), - (12, 38, 0.4, "😭🐕"), - (28, 44, 0.7, "😿🐩"), - (42, 34, 1.0, "🥺🦮"), - (58, 46, 1.3, "😭🐕‍🦺"), - (72, 36, 1.6, "😢🐶"), - (86, 40, 1.9, "😭🐕"), - (16, 64, 0.5, "🥺🐩"), - (32, 70, 0.8, "😭🦮"), - (48, 60, 1.1, "😿🐕‍🦺"), - (62, 74, 1.4, "🥺🐶"), - (78, 68, 1.7, "😭🐕"), - (90, 72, 2.0, "😢🐩"), - (20, 88, 0.6, "🥺🦮"), - (36, 92, 0.9, "😭🐕‍🦺"), - (52, 86, 1.2, "😢🐶"), - (68, 94, 1.5, "😭🐕"), - (82, 90, 1.8, "😿🐩"), -) - - -_STRAFE_SHELLS: Tuple[Tuple[float, float], ...] = ( - (22.0, 0.0), - (28.0, 0.35), - (34.0, 0.7), - (26.0, 0.2), - (32.0, 0.55), - (24.0, 0.9), - (30.0, 1.25), -) - - -def _build_artillery(projectile: str, *, shells_only: bool = False) -> str: - """Return HTML spans for puppy artillery shells (and cannons when desired).""" - shell_markup = [] - for index, (top, delay) in enumerate(_STRAFE_SHELLS): - duration = 2.3 + (index % 3) * 0.25 - shell_markup.append( - f"{projectile}💥" - ) - shells = "".join(shell_markup) - if shells_only: - return shells - - cannons = ( - "🐶🧨🐕‍🦺🔥" +def _clean(value: Optional[str], fallback: str) -> str: + """Normalize and HTML-escape dynamic page content.""" + normalized = value.strip() if value else fallback + return escape(normalized or fallback) + + +def _render_page( + *, + service: str, + state: str, + title: str, + message: str, + status: str, + hint: str, + auto_close: bool = False, +) -> str: + """Build a complete OAuth result page without remote assets.""" + close_script = ( + "" if auto_close else "" ) - return cannons + shells - - -def _service_targets(service_name: str) -> Tuple[str, str, str, str]: - """Map service names to projectile emoji and rival logo metadata.""" - normalized = service_name.lower() - if "anthropic" in normalized or "claude" in normalized: - return "🐕‍🦺🧨", CLAUDE_LOGO_URL, "Claude logo", "" - if "chat" in normalized or "gpt" in normalized: - return "🐶🚀", CHATGPT_LOGO_URL, "ChatGPT logo", "invert" - if "gemini" in normalized or "google" in normalized: - return "🐶✨", GEMINI_LOGO_URL, "Gemini logo", "" - return "🐾💥", CHATGPT_LOGO_URL, "mystery logo", "invert" + return f""" + + + + + + {service} authentication + + + +
+ +

Code Puppy · {service}

+

{title}

+

{message}

+
{status}
+

{hint}

+
+ {close_script} + +""" diff --git a/tests/messaging/test_renderers.py b/tests/messaging/test_renderers.py index 5f2b3435c..7ab9ff280 100644 --- a/tests/messaging/test_renderers.py +++ b/tests/messaging/test_renderers.py @@ -256,11 +256,30 @@ def test_sync_renderer_queued_banner(mq): ) output = console.file.getvalue() - assert "QUEUED" in output - assert "for next turn: fix the tests" in output + assert output.startswith("\n QUEUED for next turn: fix the tests") assert chr(0x23ED) not in output +def test_sync_renderer_queued_style_includes_trailing_padding(mq): + console = make_console() + renderer = SynchronousInteractiveRenderer(mq, console=console) + + with patch.object(console, "print") as mock_print: + renderer._render_message( + UIMessage(type=MessageType.QUEUED, content="for next turn: later") + ) + + mock_print.assert_any_call() + queued = mock_print.call_args_list[1].args[0] + assert isinstance(queued, Text) + assert queued.plain == " QUEUED for next turn: later" + assert queued.spans[0].start == 0 + assert queued.spans[0].end == len(" QUEUED ") + assert str(queued.spans[0].style).startswith("bold white on ") + assert queued.spans[1].start == len(" QUEUED ") + assert queued.spans[1].style == "dim" + + def test_sync_renderer_version_dim(mq): console = make_console() r = SynchronousInteractiveRenderer(mq, console=console) diff --git a/tests/plugins/test_oauth_puppy_html.py b/tests/plugins/test_oauth_puppy_html.py new file mode 100644 index 000000000..26b4392c1 --- /dev/null +++ b/tests/plugins/test_oauth_puppy_html.py @@ -0,0 +1,55 @@ +"""Tests for the shared OAuth result pages.""" + +from code_puppy.plugins.oauth_puppy_html import ( + oauth_failure_html, + oauth_success_html, +) + + +def test_success_page_is_calm_and_self_contained() -> None: + html = oauth_success_html("ChatGPT", "Ready to use.") + + assert "You're connected" in html + assert "Ready to use." in html + assert "Code Puppy · ChatGPT" in html + assert html.count('class="mark"') == 1 + assert "http://" not in html + assert "https://" not in html + assert "artillery" not in html + assert "confetti" not in html + assert "setTimeout" in html + + +def test_failure_page_has_a_useful_next_step() -> None: + html = oauth_failure_html("Claude Code", "The request expired.") + + assert "Couldn't connect" in html + assert "The request expired." in html + assert "Return to Code Puppy and try signing in again." in html + assert "setTimeout" not in html + + +def test_dynamic_content_is_escaped() -> None: + html = oauth_failure_html("OAuth", "") + + assert "<b>OAuth</b>" in html + assert "<script>alert('nope')</script>" in html + assert " + + +
+

🐶

+

Code Puppy

+
+ + Open Chat + + + API Docs + +
+

+ WebSocket: ws://localhost:8765/ws/chat +

+
+ + + """ + ) + + @app.get("/chat") + async def chat_page(): + """Serve the chat interface page.""" + html_file = templates_dir / "chat.html" + if html_file.exists(): + return FileResponse(html_file, media_type="text/html") + return HTMLResponse( + content="

Chat template not found

", + status_code=404, + ) + + @app.get("/health") + async def health(): + """Simple health check endpoint.""" + return {"status": "healthy"} + + @app.get("/api/version-check") + async def version_check(): + """ + Check current version and latest available version. + + Returns: + dict: Contains current_version, latest_version, and update_available + """ + from code_puppy import __version__ + from code_puppy.version_checker import fetch_latest_version, versions_are_equal + + current_version = __version__ + latest_version = fetch_latest_version("code-puppy") + + # Determine if update is available + update_available = False + if latest_version: + update_available = not versions_are_equal(current_version, latest_version) + + return { + "current_version": current_version, + "latest_version": latest_version or current_version, + "update_available": update_available, + "status": "success" if latest_version else "error_fetching_latest", + } + + return app diff --git a/code_puppy/api/main.py b/code_puppy/api/main.py new file mode 100644 index 000000000..2153ae43b --- /dev/null +++ b/code_puppy/api/main.py @@ -0,0 +1,97 @@ +"""Entry point for running the FastAPI server.""" + +import logging +import os +import sys + +import uvicorn + +from code_puppy.api.app import create_app +from code_puppy.logging_setup import configure_backend_logging + + +def _resolve_log_level() -> tuple[int, str]: + """Resolve env log levels to (backend_logging_level_int, uvicorn_level_str).""" + raw = ( + (os.getenv("BACKEND_LOG_LEVEL") or os.getenv("LOG_LEVEL") or "INFO") + .strip() + .upper() + ) + level_int = getattr(logging, raw, logging.INFO) + # Uvicorn wants lower-case level names + uvicorn_level = ( + raw.lower() + if raw.lower() in {"critical", "error", "warning", "info", "debug", "trace"} + else "info" + ) + return level_int, uvicorn_level + + +# Configure logging (honors BACKEND_LOG_* and LOG_LEVEL fallback) +_level_int, _uvicorn_level = _resolve_log_level() +try: + _level_int = configure_backend_logging(level=_level_int) +except Exception as exc: + print( + f"[code_puppy.api.main] configure_backend_logging failed: {exc}", + file=sys.stderr, + ) + _fallback_handler = logging.StreamHandler(sys.stdout) + _fallback_handler.setLevel(_level_int) + _fallback_handler.setFormatter( + logging.Formatter( + "%(asctime)s [%(levelname)s] %(name)s: %(message)s", + datefmt="%H:%M:%S", + ) + ) + logging.basicConfig( + level=_level_int, + handlers=[_fallback_handler], + force=True, + ) + +for name in [ + "code_puppy", + "code_puppy.api", + "code_puppy.api.websocket", + "uvicorn", + "uvicorn.error", + "uvicorn.access", + "httpx", +]: + logging.getLogger(name).setLevel(_level_int) + +app = create_app() + + +def main(host: str = "127.0.0.1", port: int = 8765) -> None: + """Run the FastAPI server. + + Args: + host: The host address to bind to. Defaults to localhost. + port: The port number to listen on. Defaults to 8765. + """ + # Force stdout to be unbuffered + sys.stdout.reconfigure(line_buffering=True) + + if (os.getenv("DEBUG_IMPORTS") or "").strip() == "1": + import code_puppy + from code_puppy.agents import base_agent as _base_agent + + print("\n=== DEBUG_IMPORTS=1 ===", flush=True) + print(f"code_puppy package: {code_puppy.__file__}", flush=True) + print(f"base_agent module: {_base_agent.__file__}", flush=True) + print("sys.path (first 10):", flush=True) + for p in sys.path[:10]: + print(f" - {p}", flush=True) + print("=======================\n", flush=True) + + print( + f"🐶 Starting Code Puppy API server (LOG_LEVEL={logging.getLevelName(_level_int)})...", + flush=True, + ) + uvicorn.run(app, host=host, port=port, log_level=_uvicorn_level) + + +if __name__ == "__main__": + main() diff --git a/code_puppy/api/routers/__init__.py b/code_puppy/api/routers/__init__.py new file mode 100644 index 000000000..00b53ba6d --- /dev/null +++ b/code_puppy/api/routers/__init__.py @@ -0,0 +1,12 @@ +"""API routers for Code Puppy REST endpoints. + +This package contains the FastAPI router modules for different API domains: + - config: Configuration management endpoints + - commands: Command execution endpoints + - sessions: Session management endpoints + - agents: Agent-related endpoints +""" + +from code_puppy.api.routers import agents, commands, config, protocol, sessions + +__all__ = ["config", "commands", "sessions", "agents", "protocol"] diff --git a/code_puppy/api/routers/agents.py b/code_puppy/api/routers/agents.py new file mode 100644 index 000000000..faf975650 --- /dev/null +++ b/code_puppy/api/routers/agents.py @@ -0,0 +1,116 @@ +"""Agents API endpoints for agent management. + +This router provides REST endpoints for: +- Listing all available agents with their metadata +- Refreshing the agent registry to discover new agents +- Switching the current active agent +""" + +from typing import Any, Dict, List + +from fastapi import APIRouter, HTTPException +from pydantic import BaseModel + +router = APIRouter() + + +@router.get("/") +async def list_agents() -> List[Dict[str, Any]]: + """List all available agents. + + Returns a list of all agents registered in the system, + including their name, display name, and description. + + Returns: + List[Dict[str, Any]]: List of agent information dictionaries. + """ + from code_puppy.agents import get_agent_descriptions, get_available_agents + + agents_dict = get_available_agents() + descriptions = get_agent_descriptions() + + return [ + { + "name": name, + "display_name": display_name, + "description": descriptions.get(name, "No description"), + } + for name, display_name in agents_dict.items() + ] + + +@router.post("/refresh") +async def refresh_agents_endpoint() -> Dict[str, Any]: + """Force refresh the agents list by re-running agent discovery. + + This endpoint triggers the agent manager to re-scan for: + - Python agent classes in the agents package + - JSON agent configuration files in user directory + - Plugin-registered agents + + Returns: + Dict[str, Any]: Result containing: + - success (bool): True if refresh completed + - count (int): Number of agents discovered + - message (str): Status message + """ + from code_puppy.agents import get_available_agents + from code_puppy.agents import refresh_agents as do_refresh_agents + + # Refresh agents - this will call _discover_agents() again + do_refresh_agents() + + # Get the fresh count + agents = get_available_agents() + + return { + "success": True, + "count": len(agents), + "message": "Agent discovery refreshed", + } + + +class SwitchAgentRequest(BaseModel): + """Request model for switching agents.""" + + agent_name: str + + +@router.post("/switch") +async def switch_agent(request: SwitchAgentRequest) -> Dict[str, Any]: + """Switch to a different agent. + + Args: + request: SwitchAgentRequest containing: + - agent_name (str): The name of the agent to switch to + + Returns: + Dict[str, Any]: Result containing: + - success (bool): True if switch was successful + - agentName (str): The name of the agent switched to + - message (str): Success message + + Raises: + HTTPException: If agent not found (404) or switch fails (500) + """ + from code_puppy.agents import get_current_agent, set_current_agent + + agent_name = request.agent_name + + if not agent_name: + raise HTTPException(status_code=400, detail="agent_name is required") + + # Switch to the new agent + success = set_current_agent(agent_name) + + if not success: + raise HTTPException(status_code=404, detail=f"Agent '{agent_name}' not found") + + # Get the new current agent for display name + current = get_current_agent() + + return { + "success": True, + "agentName": current.name, + "message": f"Switched to {current.display_name}", + } diff --git a/code_puppy/api/routers/commands.py b/code_puppy/api/routers/commands.py new file mode 100644 index 000000000..e0e8892da --- /dev/null +++ b/code_puppy/api/routers/commands.py @@ -0,0 +1,217 @@ +"""Commands API endpoints for slash command execution and autocomplete. + +This router provides REST endpoints for: +- Listing all available slash commands +- Getting info about specific commands +- Executing slash commands +- Autocomplete suggestions for partial commands +""" + +import asyncio +from concurrent.futures import ThreadPoolExecutor +from typing import Any, List, Optional + +from fastapi import APIRouter, HTTPException +from pydantic import BaseModel + +# Thread pool for blocking command execution +_executor = ThreadPoolExecutor(max_workers=4) + +# Timeout for command execution (seconds) +COMMAND_TIMEOUT = 30.0 + +router = APIRouter() + + +# ============================================================================= +# Pydantic Models +# ============================================================================= + + +class CommandInfo(BaseModel): + """Information about a registered command.""" + + name: str + description: str + usage: str + aliases: List[str] = [] + category: str = "core" + detailed_help: Optional[str] = None + + +class CommandExecuteRequest(BaseModel): + """Request to execute a slash command.""" + + command: str # Full command string, e.g., "/set model=gpt-4o" + + +class CommandExecuteResponse(BaseModel): + """Response from executing a slash command.""" + + success: bool + result: Any = None + error: Optional[str] = None + + +class AutocompleteRequest(BaseModel): + """Request for command autocomplete.""" + + partial: str # Partial command string, e.g., "/se" or "/set mo" + + +class AutocompleteResponse(BaseModel): + """Response with autocomplete suggestions.""" + + suggestions: List[str] + + +# ============================================================================= +# Endpoints +# ============================================================================= + + +@router.get("/") +async def list_commands() -> List[CommandInfo]: + """List all available slash commands. + + Returns a sorted list of all unique commands (no alias duplicates), + with their metadata including name, description, usage, aliases, + category, and detailed help. + + Returns: + List[CommandInfo]: Sorted list of command information. + """ + from code_puppy.command_line.command_registry import get_unique_commands + + commands = [] + for cmd in get_unique_commands(): + commands.append( + CommandInfo( + name=cmd.name, + description=cmd.description, + usage=cmd.usage, + aliases=cmd.aliases, + category=cmd.category, + detailed_help=cmd.detailed_help, + ) + ) + return sorted(commands, key=lambda c: c.name) + + +@router.get("/{name}") +async def get_command_info(name: str) -> CommandInfo: + """Get detailed info about a specific command. + + Looks up a command by name or alias (case-insensitive). + + Args: + name: Command name or alias (without leading /). + + Returns: + CommandInfo: Full command information. + + Raises: + HTTPException: 404 if command not found. + """ + from code_puppy.command_line.command_registry import get_command + + cmd = get_command(name) + if not cmd: + raise HTTPException(404, f"Command '/{name}' not found") + + return CommandInfo( + name=cmd.name, + description=cmd.description, + usage=cmd.usage, + aliases=cmd.aliases, + category=cmd.category, + detailed_help=cmd.detailed_help, + ) + + +@router.post("/execute") +async def execute_command(request: CommandExecuteRequest) -> CommandExecuteResponse: + """Execute a slash command. + + Takes a command string (with or without leading /) and executes it + using the command handler. Runs in a thread pool to avoid blocking + the event loop, with a timeout to prevent hangs. + + Args: + request: CommandExecuteRequest with the command to execute. + + Returns: + CommandExecuteResponse: Result of command execution. + """ + from code_puppy.command_line.command_handler import handle_command + + command = request.command + if not command.startswith("/"): + command = "/" + command + + loop = asyncio.get_running_loop() + + try: + # Run blocking command in thread pool with timeout + result = await asyncio.wait_for( + loop.run_in_executor(_executor, handle_command, command), + timeout=COMMAND_TIMEOUT, + ) + return CommandExecuteResponse(success=True, result=result) + except asyncio.TimeoutError: + return CommandExecuteResponse( + success=False, error=f"Command timed out after {COMMAND_TIMEOUT}s" + ) + except Exception as e: + return CommandExecuteResponse(success=False, error=str(e)) + + +@router.post("/autocomplete") +async def autocomplete_command(request: AutocompleteRequest) -> AutocompleteResponse: + """Get autocomplete suggestions for a partial command. + + Provides intelligent autocomplete based on partial input: + - Empty input: returns all command names + - Partial command name: returns matching commands and aliases + - Complete command with args: returns usage hint + + Args: + request: AutocompleteRequest with partial command string. + + Returns: + AutocompleteResponse: List of autocomplete suggestions. + """ + from code_puppy.command_line.command_registry import ( + get_command, + get_unique_commands, + ) + + partial = request.partial.lstrip("/") + + # If empty, return all command names + if not partial: + suggestions = [f"/{cmd.name}" for cmd in get_unique_commands()] + return AutocompleteResponse(suggestions=sorted(suggestions)) + + # Split into command name and args + parts = partial.split(maxsplit=1) + cmd_partial = parts[0].lower() + + # If just the command name (no space yet), suggest matching commands + if len(parts) == 1: + suggestions = [] + for cmd in get_unique_commands(): + if cmd.name.startswith(cmd_partial): + suggestions.append(f"/{cmd.name}") + for alias in cmd.aliases: + if alias.startswith(cmd_partial): + suggestions.append(f"/{alias}") + return AutocompleteResponse(suggestions=sorted(set(suggestions))) + + # Command name complete, suggest based on command type + # (For now, just return the command usage as a hint) + cmd = get_command(cmd_partial) + if cmd: + return AutocompleteResponse(suggestions=[cmd.usage]) + + return AutocompleteResponse(suggestions=[]) diff --git a/code_puppy/api/routers/config.py b/code_puppy/api/routers/config.py new file mode 100644 index 000000000..d92778b7f --- /dev/null +++ b/code_puppy/api/routers/config.py @@ -0,0 +1,211 @@ +"""Configuration management API endpoints.""" + +from typing import Any, Dict, List + +from fastapi import APIRouter, HTTPException +from pydantic import BaseModel + +router = APIRouter() + + +class ConfigValue(BaseModel): + key: str + value: Any + + +class ConfigUpdate(BaseModel): + value: Any + + +@router.get("/") +async def list_config() -> Dict[str, Any]: + """List all configuration keys and their current values.""" + from code_puppy.config import get_config_keys, get_value + + config = {} + for key in get_config_keys(): + config[key] = get_value(key) + return {"config": config} + + +@router.get("/keys") +async def get_config_keys_list() -> List[str]: + """Get list of all valid configuration keys.""" + from code_puppy.config import get_config_keys + + return get_config_keys() + + +@router.get("/schema") +async def get_config_schema() -> Dict[str, Any]: + """Get configuration schema with metadata for all config keys. + + Returns metadata including: + - description: Human-readable description of the config key + - category: Config category (core, behavior, model, advanced, experimental) + - type: Data type (boolean, string, number, choice) + - choices: Valid values (for choice type) + - default: Default value + + This endpoint allows frontends to dynamically render config UIs + without hardcoding config key metadata. + + Returns: + Dict with 'keys' mapping config names to their metadata + """ + try: + from code_puppy.config import CONFIG_SCHEMA + + return {"keys": CONFIG_SCHEMA} + except Exception: + # Migration-compat fallback: synthesize a minimal schema so the + # endpoint remains functional even if CONFIG_SCHEMA is not exported. + from code_puppy.config import get_config_keys, get_value + + inferred: Dict[str, Dict[str, Any]] = {} + for key in get_config_keys(): + value = get_value(key) + inferred[key] = { + "description": "", + "category": "legacy", + "type": type(value).__name__ if value is not None else "unknown", + "default": value, + } + return {"keys": inferred} + + +@router.get("/{key}") +async def get_config_value(key: str) -> ConfigValue: + """Get a specific configuration value.""" + from code_puppy.config import get_config_keys, get_value + + valid_keys = get_config_keys() + if key not in valid_keys: + raise HTTPException( + 404, f"Config key '{key}' not found. Valid keys: {valid_keys}" + ) + + value = get_value(key) + return ConfigValue(key=key, value=value) + + +@router.put("/{key}") +async def set_config_value(key: str, update: ConfigUpdate) -> ConfigValue: + """Set a configuration value.""" + from code_puppy.config import get_config_keys, get_value, set_value + + valid_keys = get_config_keys() + if key not in valid_keys: + raise HTTPException( + 404, f"Config key '{key}' not found. Valid keys: {valid_keys}" + ) + + set_value(key, str(update.value)) + return ConfigValue(key=key, value=get_value(key)) + + +@router.delete("/{key}") +async def reset_config_value(key: str) -> Dict[str, str]: + """Reset a configuration value to default (remove from config file).""" + from code_puppy.config import reset_value + + reset_value(key) + return {"message": f"Config key '{key}' reset to default"} + + +# ============================================================================= +# Model Settings Endpoints +# ============================================================================= + + +class ModelSettingsResponse(BaseModel): + """Response containing model-specific settings.""" + + model_name: str + settings: Dict[str, Any] + + +class ModelSettingsUpdateRequest(BaseModel): + """Request to update a model setting.""" + + setting: str + value: Any + + +@router.get("/model_settings/{model_name}") +async def get_model_settings(model_name: str) -> ModelSettingsResponse: + """Get all settings for a specific model. + + Returns the persisted settings from ~/.code_puppy/code_puppy.ini + for the specified model name. + + Args: + model_name: The model name (e.g., 'gpt-5-2-0125', 'claude-4-5-sonnet') + + Returns: + ModelSettingsResponse with model name and settings dict + """ + from code_puppy.config import get_all_model_settings + + settings = get_all_model_settings(model_name) + return ModelSettingsResponse(model_name=model_name, settings=settings) + + +@router.put("/model_settings/{model_name}") +async def update_model_setting( + model_name: str, request: ModelSettingsUpdateRequest +) -> ModelSettingsResponse: + """Update a specific setting for a model. + + Persists the setting to ~/.code_puppy/code_puppy.ini + + Args: + model_name: The model name + request: ModelSettingsUpdateRequest with setting name and value + + Returns: + Updated ModelSettingsResponse + """ + from code_puppy.config import get_all_model_settings, set_model_setting + + set_model_setting(model_name, request.setting, request.value) + settings = get_all_model_settings(model_name) + return ModelSettingsResponse(model_name=model_name, settings=settings) + + +@router.delete("/model_settings/{model_name}") +async def clear_model_settings(model_name: str) -> Dict[str, str]: + """Clear all settings for a specific model. + + Removes all per-model settings, reverting to defaults. + + Args: + model_name: The model name + + Returns: + Success message + """ + from code_puppy.config import clear_model_settings + + clear_model_settings(model_name) + return {"message": f"All settings cleared for model '{model_name}'"} + + +@router.delete("/model_settings/{model_name}/{setting}") +async def clear_model_setting(model_name: str, setting: str) -> ModelSettingsResponse: + """Clear a specific setting for a model. + + Removes the setting, reverting it to the default. + + Args: + model_name: The model name + setting: The setting name to clear + + Returns: + Updated ModelSettingsResponse + """ + from code_puppy.config import get_all_model_settings, set_model_setting + + set_model_setting(model_name, setting, None) + settings = get_all_model_settings(model_name) + return ModelSettingsResponse(model_name=model_name, settings=settings) diff --git a/code_puppy/api/routers/models.py b/code_puppy/api/routers/models.py new file mode 100644 index 000000000..c78d478cb --- /dev/null +++ b/code_puppy/api/routers/models.py @@ -0,0 +1,23 @@ +"""Models API endpoint for listing available models.""" + +from typing import Any, Dict, List + +from fastapi import APIRouter + +router = APIRouter() + + +@router.get("/") +async def list_models() -> List[Dict[str, Any]]: + """List all available model names from models.json config. + + Returns: + List of dicts with 'name' key for each model. + """ + from code_puppy.model_factory import ModelFactory + + try: + config = ModelFactory.load_config() + return [{"name": name} for name in sorted(config.keys())] + except Exception: + return [] diff --git a/code_puppy/api/routers/protocol.py b/code_puppy/api/routers/protocol.py new file mode 100644 index 000000000..e5b513559 --- /dev/null +++ b/code_puppy/api/routers/protocol.py @@ -0,0 +1,26 @@ +"""Protocol schema endpoint – exposes JSON Schema for the WebSocket protocol.""" + +from fastapi import APIRouter +from pydantic import TypeAdapter + +from code_puppy.api.ws.schemas import ( + PROTOCOL_VERSION, + ClientMessage, + ServerMessage, +) + +router = APIRouter() + +# Cache the schemas since they don't change at runtime +_client_schema = TypeAdapter(ClientMessage).json_schema() +_server_schema = TypeAdapter(ServerMessage).json_schema() + + +@router.get("/schema") +async def get_protocol_schema(): + """Return the WebSocket protocol JSON Schema for client consumption.""" + return { + "protocol_version": PROTOCOL_VERSION, + "client_messages": _client_schema, + "server_messages": _server_schema, + } diff --git a/code_puppy/api/tests/test_config_router_schema.py b/code_puppy/api/tests/test_config_router_schema.py new file mode 100644 index 000000000..96c97176a --- /dev/null +++ b/code_puppy/api/tests/test_config_router_schema.py @@ -0,0 +1,17 @@ +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from code_puppy.api.routers import config as config_router + + +def test_config_schema_endpoint_returns_mapping_without_config_schema_symbol(): + app = FastAPI() + app.include_router(config_router.router, prefix="/config") + + with TestClient(app) as client: + resp = client.get("/config/schema") + + assert resp.status_code == 200 + payload = resp.json() + assert "keys" in payload + assert isinstance(payload["keys"], dict) diff --git a/code_puppy/api/tests/test_gate2_legacy_ws_namespace.py b/code_puppy/api/tests/test_gate2_legacy_ws_namespace.py new file mode 100644 index 000000000..d1aaedde0 --- /dev/null +++ b/code_puppy/api/tests/test_gate2_legacy_ws_namespace.py @@ -0,0 +1,51 @@ +"""Post-cleanup WebSocket namespace tests for puppy-desk migration.""" + +from pathlib import Path + +from fastapi import FastAPI, WebSocket + + +def test_setup_websocket_keeps_existing_chat_route_and_no_extra_routes(monkeypatch): + import code_puppy.api.websocket as websocket_module + + def _register_chat(app): + @app.websocket("/ws/chat") + async def _chat(_ws: WebSocket): + return None + + def _register_events(app): + @app.websocket("/ws/events") + async def _events(_ws: WebSocket): + return None + + def _register_health(app): + @app.websocket("/ws/health") + async def _health(_ws: WebSocket): + return None + + monkeypatch.setattr(websocket_module, "register_chat_endpoint", _register_chat) + monkeypatch.setattr(websocket_module, "register_events_endpoint", _register_events) + monkeypatch.setattr(websocket_module, "register_health_endpoint", _register_health) + + app = FastAPI() + websocket_module.setup_websocket(app) + + websocket_paths = { + getattr(route, "path", None) + for route in app.routes + if getattr(route, "path", None) + } + + assert "/ws/chat" in websocket_paths + assert "/ws/events" in websocket_paths + assert "/ws/health" in websocket_paths + assert "/ws/terminal" not in websocket_paths + assert "/ws/sessions" not in websocket_paths + assert "/ws/chat-migration" not in websocket_paths + assert "/ws/chat-next" not in websocket_paths + assert "/ws/chat-v2" not in websocket_paths + + +def test_legacy_ws_snapshot_removed_after_cleanup(): + ws_dir = Path(__file__).resolve().parents[1] / "ws" + assert not (ws_dir / "legacy").exists() diff --git a/code_puppy/api/tests/test_send_utils.py b/code_puppy/api/tests/test_send_utils.py new file mode 100644 index 000000000..561bdc924 --- /dev/null +++ b/code_puppy/api/tests/test_send_utils.py @@ -0,0 +1,187 @@ +"""Focused tests for WebSocketSender (Phase 3).""" + +from __future__ import annotations + +import pytest + +from code_puppy.api.ws.send_utils import WebSocketSender + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class _FakeWebSocket: + """Minimal WebSocket stub for unit tests.""" + + def __init__(self, *, fail: bool = False, fail_msg: str = "closed"): + self.sent: list[dict] = [] + self._fail = fail + self._fail_msg = fail_msg + + async def send_json(self, data: dict) -> None: + if self._fail: + raise RuntimeError(self._fail_msg) + self.sent.append(data) + + +class _FakeCtx: + def __init__(self, agent_name: str = "", model_name: str = ""): + self.agent_name = agent_name + self.model_name = model_name + + +# --------------------------------------------------------------------------- +# Construction +# --------------------------------------------------------------------------- + + +def test_sender_initial_state(): + ws = _FakeWebSocket() + sender = WebSocketSender(ws, session_id="s1") + assert sender.ws_closed is False + assert sender.session_id == "s1" + assert sender.ctx is None + + +def test_sender_ctx_settable(): + ws = _FakeWebSocket() + sender = WebSocketSender(ws, session_id="s1") + ctx = _FakeCtx(agent_name="husky") + sender.ctx = ctx + assert sender.ctx is ctx + + +# --------------------------------------------------------------------------- +# safe_send_json +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_safe_send_json_success(): + ws = _FakeWebSocket() + sender = WebSocketSender(ws, session_id="s1") + ok = await sender.safe_send_json({"type": "status", "status": "done"}) + assert ok is True + assert len(ws.sent) == 1 + assert ws.sent[0]["type"] == "status" + + +@pytest.mark.asyncio +async def test_safe_send_json_skips_when_closed(): + ws = _FakeWebSocket() + sender = WebSocketSender(ws, session_id="s1") + sender.ws_closed = True + ok = await sender.safe_send_json({"type": "status"}) + assert ok is False + assert len(ws.sent) == 0 + + +@pytest.mark.asyncio +async def test_safe_send_json_marks_closed_on_closed_error(): + ws = _FakeWebSocket(fail=True, fail_msg="WebSocket is closed") + sender = WebSocketSender(ws, session_id="s1") + ok = await sender.safe_send_json({"type": "status"}) + assert ok is False + assert sender.ws_closed is True + + +@pytest.mark.asyncio +async def test_safe_send_json_non_close_error_keeps_open(): + ws = _FakeWebSocket(fail=True, fail_msg="network timeout") + sender = WebSocketSender(ws, session_id="s1") + ok = await sender.safe_send_json({"type": "status"}) + assert ok is False + assert sender.ws_closed is False + + +# --------------------------------------------------------------------------- +# send_typed +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_send_typed_serialises_pydantic_model(): + from code_puppy.api.ws.schemas import ServerStatus + + ws = _FakeWebSocket() + sender = WebSocketSender(ws, session_id="s1") + msg = ServerStatus(status="done", session_id="s1", agent_name="a", model_name="m") + ok = await sender.send_typed(msg) + assert ok is True + assert ws.sent[0]["type"] == "status" + assert ws.sent[0]["status"] == "done" + + +# --------------------------------------------------------------------------- +# send_typed_tool_lifecycle +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_send_typed_tool_lifecycle(): + from code_puppy.api.ws.schemas import ServerToolCall + + ws = _FakeWebSocket() + sender = WebSocketSender(ws, session_id="s1") + msg = ServerToolCall( + tool_id="t1", + tool_name="read_file", + args={"file_path": "a.py"}, + session_id="s1", + timestamp=1.0, + ) + ok = await sender.send_typed_tool_lifecycle(msg) + assert ok is True + assert ws.sent[0]["tool_name"] == "read_file" + + +# --------------------------------------------------------------------------- +# persist_error_payload (unit — no real DB) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_persist_error_skips_non_error_type(): + ws = _FakeWebSocket() + sender = WebSocketSender(ws, session_id="s1") + # Should return without attempting DB write + await sender.persist_error_payload({"type": "status"}) + + +@pytest.mark.asyncio +async def test_persist_error_skips_empty_session_id(): + ws = _FakeWebSocket() + sender = WebSocketSender(ws, session_id="") + await sender.persist_error_payload({"type": "error", "error": "boom"}) + + +# --------------------------------------------------------------------------- +# Import isolation +# --------------------------------------------------------------------------- + + +def test_send_utils_import_does_not_eagerly_load_chat_handler(): + import importlib + import sys + + sys.modules.pop("code_puppy.api.ws", None) + sys.modules.pop("code_puppy.api.ws.chat_handler", None) + sys.modules.pop("code_puppy.api.ws.send_utils", None) + + importlib.import_module("code_puppy.api.ws.send_utils") + + assert "code_puppy.api.ws.chat_handler" not in sys.modules + + +def test_session_persistence_import_does_not_eagerly_load_chat_handler(): + import importlib + import sys + + sys.modules.pop("code_puppy.api.ws", None) + sys.modules.pop("code_puppy.api.ws.chat_handler", None) + sys.modules.pop("code_puppy.api.ws.session_persistence", None) + + importlib.import_module("code_puppy.api.ws.session_persistence") + + assert "code_puppy.api.ws.chat_handler" not in sys.modules diff --git a/code_puppy/api/tests/test_streaming_protocol_transform.py b/code_puppy/api/tests/test_streaming_protocol_transform.py new file mode 100644 index 000000000..350b7ff3a --- /dev/null +++ b/code_puppy/api/tests/test_streaming_protocol_transform.py @@ -0,0 +1,123 @@ +"""Tests for streaming-only assistant response protocol adaptation.""" + +from code_puppy.api.ws.response_frames import ( + build_assistant_text_stream_frames, + build_error_response_frames, +) + + +def _dump_frames(**kwargs): + return [ + frame.model_dump(exclude_none=True) + for frame in build_assistant_text_stream_frames(**kwargs) + ] + + +def test_complete_response_is_adapted_to_streaming_frames(): + frames = _dump_frames( + response_text="hello from a non-streaming model", + session_id="WS_session_test", + agent_name="code-puppy", + model_name="non-stream-model", + tokens={"total_tokens": 7}, + message_id="msg-fixed", + timestamp=123.456, + ) + + assert [frame["type"] for frame in frames] == [ + "assistant_message_start", + "assistant_message_delta", + "assistant_message_end", + "stream_end", + ] + + start, delta, end, stream_end = frames + assert start == { + "type": "assistant_message_start", + "message_id": "msg-fixed", + "part_type": "text", + "part_index": 0, + "timestamp": 123.456, + "session_id": "WS_session_test", + "agent_name": "code-puppy", + "model_name": "non-stream-model", + } + assert delta["message_id"] == "msg-fixed" + assert delta["content"] == "hello from a non-streaming model" + assert delta["part_index"] == 0 + assert delta["session_id"] == "WS_session_test" + + assert end["message_id"] == "msg-fixed" + assert end["part_type"] == "text" + assert end["full_content"] == "hello from a non-streaming model" + assert end["timestamp"] == 123.456 + + assert stream_end == { + "type": "stream_end", + "success": True, + "session_id": "WS_session_test", + "total_length": len("hello from a non-streaming model"), + "agent_name": "code-puppy", + "model_name": "non-stream-model", + "tokens": {"total_tokens": 7}, + } + + +def test_complete_response_transform_does_not_emit_legacy_response_frame(): + frames = _dump_frames( + response_text="render me through streaming", + session_id="WS_session_test", + message_id="msg-fixed", + timestamp=1.0, + ) + + assert "response" not in {frame["type"] for frame in frames} + assert any( + frame["type"] == "assistant_message_delta" + and frame["content"] == "render me through streaming" + for frame in frames + ) + + +def test_complete_response_transform_preserves_part_type_for_thinking_parts(): + frames = _dump_frames( + response_text="reasoning trace", + session_id="WS_session_test", + part_type="thinking", + part_index=1, + message_id="thinking-fixed", + timestamp=2.0, + ) + + assert frames[0]["part_type"] == "thinking" + assert frames[1]["part_index"] == 1 + assert frames[2]["part_type"] == "thinking" + assert frames[2]["part_index"] == 1 + assert frames[3]["type"] == "stream_end" + + +def test_error_after_partial_stream_emits_stream_end_then_error(): + frames = build_error_response_frames( + RuntimeError("backend unavailable"), + collected_text=["partial output"], + session_id="WS_session_test", + ) + + assert frames[0]["type"] == "stream_end" + assert frames[0]["success"] is False + assert frames[1]["type"] == "error" + assert frames[1]["session_id"] == "WS_session_test" + + +def test_response_frames_import_does_not_eagerly_load_chat_handler(): + import importlib + import sys + + sys.modules.pop("code_puppy.api.ws", None) + sys.modules.pop("code_puppy.api.ws.chat_handler", None) + sys.modules.pop("code_puppy.api.ws.response_frames", None) + + module = importlib.import_module("code_puppy.api.ws.response_frames") + + assert module is not None + assert "code_puppy.api.ws.chat_handler" not in sys.modules diff --git a/code_puppy/api/tests/test_ws_history_utils.py b/code_puppy/api/tests/test_ws_history_utils.py new file mode 100644 index 000000000..c5770f9a0 --- /dev/null +++ b/code_puppy/api/tests/test_ws_history_utils.py @@ -0,0 +1,91 @@ +"""Focused tests for WebSocket history utility helpers.""" + +from __future__ import annotations + +from code_puppy.api.ws.history_utils import ( + build_enhanced_history, + estimate_total_tokens, + extract_message_timestamp, +) + + +class _MsgWithTimestamp: + def __init__(self, timestamp): + self.timestamp = timestamp + + +class _GoodAgent: + def estimate_tokens_for_message(self, msg): + return len(str(msg)) + + +class _BadAgent: + def estimate_tokens_for_message(self, msg): + raise RuntimeError("boom") + + +def test_extract_message_timestamp_prefers_dict_timestamp_fields(): + assert ( + extract_message_timestamp({"timestamp": "2024-01-02T03:04:05"}, "fallback") + == "2024-01-02T03:04:05" + ) + assert ( + extract_message_timestamp({"ts": "2024-02-03T04:05:06"}, "fallback") + == "2024-02-03T04:05:06" + ) + + +def test_extract_message_timestamp_uses_object_attribute_and_fallback(): + assert ( + extract_message_timestamp(_MsgWithTimestamp("2024-03-04T05:06:07"), "fallback") + == "2024-03-04T05:06:07" + ) + assert extract_message_timestamp(object(), "fallback") == "fallback" + + +def test_build_enhanced_history_backfills_attachment_metadata_on_user_message(): + history = [ + {"role": "user", "content": "expanded with file contents"}, + {"role": "assistant", "content": "done"}, + ] + + enhanced = build_enhanced_history( + history, + agent_name_meta="code-puppy", + model_name_meta="gpt-test", + original_user_message="clean user text", + attachment_metadata=[{"name": "a.txt", "path": "/tmp/a.txt"}], + ) + + assert enhanced[0]["clean_content"] == "clean user text" + assert enhanced[0]["attachments"] == [{"name": "a.txt", "path": "/tmp/a.txt"}] + assert enhanced[1]["msg"] == {"role": "assistant", "content": "done"} + + +def test_build_enhanced_history_preserves_pre_wrapped_entries(): + wrapped = { + "msg": {"role": "user", "content": "hi"}, + "agent": "existing-agent", + "model": "existing-model", + "ts": "2024-01-01T00:00:00", + } + + enhanced = build_enhanced_history( + [wrapped], + agent_name_meta="new-agent", + model_name_meta="new-model", + original_user_message="ignored", + attachment_metadata=[], + ) + + assert enhanced == [wrapped] + + +def test_estimate_total_tokens_sums_wrapped_messages(): + enhanced = [{"msg": "abc"}, {"msg": "de"}] + assert estimate_total_tokens(enhanced, _GoodAgent()) == 5 + + +def test_estimate_total_tokens_returns_zero_on_estimator_failure(): + enhanced = [{"msg": "abc"}] + assert estimate_total_tokens(enhanced, _BadAgent()) == 0 diff --git a/code_puppy/api/websocket.py b/code_puppy/api/websocket.py new file mode 100644 index 000000000..e4ecfca05 --- /dev/null +++ b/code_puppy/api/websocket.py @@ -0,0 +1,45 @@ +"""WebSocket endpoints for Code Puppy API. + +Provides real-time communication channels: +- /ws/events - Server-sent events stream +- /ws/chat - Interactive chat with the agent +- /ws/health - Simple health check endpoint + +This module is the thin entry point that registers all WebSocket endpoints. +Each handler lives in its own module under `code_puppy.api.ws.*` for +maintainability. +""" + +import logging + +from fastapi import FastAPI + +from code_puppy.api.ws import ( + register_chat_endpoint, + register_events_endpoint, + register_health_endpoint, +) + +# Re-export build_file_context_and_attachments for backward compatibility +from code_puppy.api.ws.attachments import ( # noqa: F401 + build_file_context_and_attachments, +) + +# Re-export connection_manager for backward compatibility +from code_puppy.api.ws.connection_manager import ( # noqa: F401 + WebSocketConnectionManager, +) + +logger = logging.getLogger(__name__) + + +def setup_websocket(app: FastAPI) -> None: + """Setup WebSocket endpoints for the application. + + Registers all WebSocket routes by delegating to specialized handler modules. + """ + register_events_endpoint(app) + register_chat_endpoint(app) + register_health_endpoint(app) + + logger.debug("WebSocket endpoints registered") diff --git a/code_puppy/api/ws/__init__.py b/code_puppy/api/ws/__init__.py new file mode 100644 index 000000000..8954576f2 --- /dev/null +++ b/code_puppy/api/ws/__init__.py @@ -0,0 +1,42 @@ +"""WebSocket endpoint handlers package. + +Keep package imports lightweight: this module intentionally avoids importing the +large chat handler at package import time so helper submodules such as +``response_frames`` can be imported independently in tests and tooling. +""" + +from __future__ import annotations + + +def register_chat_endpoint(app): + from code_puppy.api.ws.chat_handler import register_chat_endpoint as _impl + + return _impl(app) + + +def register_events_endpoint(app): + from code_puppy.api.ws.events_handler import register_events_endpoint as _impl + + return _impl(app) + + +def register_health_endpoint(app): + from code_puppy.api.ws.health_handler import register_health_endpoint as _impl + + return _impl(app) + + +def __getattr__(name: str): + if name == "connection_manager": + from code_puppy.api.ws.connection_manager import connection_manager + + return connection_manager + raise AttributeError(name) + + +__all__ = [ + "register_chat_endpoint", + "register_events_endpoint", + "register_health_endpoint", + "connection_manager", +] diff --git a/code_puppy/api/ws/attachments.py b/code_puppy/api/ws/attachments.py new file mode 100644 index 000000000..f2db90726 --- /dev/null +++ b/code_puppy/api/ws/attachments.py @@ -0,0 +1,93 @@ +"""File attachment processing utilities for WebSocket chat. + +Handles converting file attachment paths into either binary content +(for images/PDFs) or text context (for code/text files). +""" + +import logging +from pathlib import Path + +logger = logging.getLogger(__name__) + +# Supported binary types (Claude's native attachments) +BINARY_EXTENSIONS = { + ".png", + ".jpg", + ".jpeg", + ".gif", + ".webp", + ".pdf", + ".bmp", + ".tiff", +} + + +def build_file_context_and_attachments(msg: dict): + """Convert attachment paths into either binary attachments OR text context. + + For images/PDFs: Send as BinaryContent (Claude's native attachment support) + For text files: Read and prepend to message so Code Puppy can analyze them + + Returns: (text_context, binary_attachments) + text_context: String to prepend to the message, or empty string + binary_attachments: List of BinaryContent for images/PDFs + """ + attachment_paths = msg.get("attachments") or [] + + if not attachment_paths: + return "", [] + + text_context_parts = [] + binary_attachments = [] + + for raw_path in attachment_paths: + if not isinstance(raw_path, str) or not raw_path.strip(): + continue + + try: + file_path = Path(raw_path) + + if not file_path.exists(): + logger.warning("Attachment file not found: %s", raw_path) + continue + + ext = file_path.suffix.lower() + + if ext in BINARY_EXTENSIONS: + try: + from pydantic_ai import BinaryContent + + from code_puppy.command_line.attachments import ( + _determine_media_type, + _load_binary, + ) + + data = _load_binary(file_path) + media_type = _determine_media_type(file_path) + binary_attachments.append( + BinaryContent(data=data, media_type=media_type) + ) + logger.debug("Loaded binary attachment: %s", file_path.name) + except Exception as e: + logger.warning( + f"Failed to load binary attachment '{raw_path}': {e}" + ) + else: + try: + content = file_path.read_text(encoding="utf-8", errors="ignore") + text_context_parts.append( + f"\n\n--- File: {file_path.name} ({raw_path}) ---\n" + f"{content}\n" + f"--- End of {file_path.name} ---\n" + ) + logger.debug( + f"Loaded text file: {file_path.name} ({len(content)} chars)" + ) + except Exception as e: + logger.warning("Failed to read text file '%s': %s", raw_path, e) + + except Exception as e: + logger.warning("Error processing attachment '%s': %s", raw_path, e) + + text_context = "".join(text_context_parts) + return text_context, binary_attachments diff --git a/code_puppy/api/ws/background_save.py b/code_puppy/api/ws/background_save.py new file mode 100644 index 000000000..4b1a1c6ad --- /dev/null +++ b/code_puppy/api/ws/background_save.py @@ -0,0 +1,244 @@ +"""Background agent result persistence for WebSocket chat sessions. + +When a user switches sessions or disconnects while the agent is running, the +agent is allowed to finish in the background. This module provides a single +reusable coroutine — ``save_agent_result_in_background`` — that: + +1. Awaits the agent task. +2. Syncs completed messages onto the agent instance. +3. Checks the session still exists (guards against resurrecting deleted sessions). +4. Persists the result to SQLite via ``write_turn_to_sqlite``. + +All three original call-sites (switch, disconnect, runtime-error) are +collapsed here, eliminating ~363 lines of duplicated closure code. + +Usage +----- +:: + + from code_puppy.api.ws.background_save import ( + fire_and_track, + save_agent_result_in_background, + ) + + fire_and_track( + save_agent_result_in_background( + agent_task=active_agent_task, + session_id=session_id, + ctx=ctx, + agent=agent, + agent_name=agent_name, + model_name=model_name, + title=session_title, + working_directory=session_working_directory, + pinned=session_pinned, + label="switch", # "switch" | "disconnect" | "runtime" + ) + ) +""" + +from __future__ import annotations + +import asyncio +import datetime +import logging +from typing import Any + +from code_puppy.api.db.queries import session_exists, write_turn_to_sqlite + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Task registry — prevents fire-and-forget tasks from being GC'd mid-flight. +# --------------------------------------------------------------------------- + +_BACKGROUND_TASKS: set[asyncio.Task] = set() + + +def fire_and_track(coro: Any) -> asyncio.Task: + """Spawn *coro* as a tracked background Task. + + The task is kept alive in ``_BACKGROUND_TASKS`` until it completes, + preventing it from being garbage-collected before it finishes. + """ + task = asyncio.create_task(coro) + _BACKGROUND_TASKS.add(task) + task.add_done_callback(_BACKGROUND_TASKS.discard) + return task + + +# --------------------------------------------------------------------------- +# Core background-save coroutine +# --------------------------------------------------------------------------- + + +async def save_agent_result_in_background( + *, + agent_task: asyncio.Task | None, + session_id: str, + ctx: Any, # SessionContext | None + agent: Any, # code_puppy Agent instance + agent_name: str, + model_name: str, + title: str, + working_directory: str, + pinned: bool, + label: str = "bg", +) -> None: + """Await *agent_task*, then persist the result to SQLite. + + Parameters + ---------- + agent_task: + The running ``run_with_mcp`` asyncio Task. ``None`` is a safe no-op. + session_id: + The session whose messages should be persisted. + ctx: + The ``SessionContext`` snapshot at time of the switch/disconnect. + May be ``None``; used for ``created_at`` and token estimation. + agent: + Agent instance attached to the session. + agent_name / model_name / title / working_directory / pinned: + Session metadata — snapshotted by the caller before state changes. + label: + Short string used in log messages: ``"switch"``, ``"disconnect"``, + or ``"runtime"``. + """ + if agent_task is None: + return + + try: + # ------------------------------------------------------------------ + # 1. Await the in-flight agent task. + # ------------------------------------------------------------------ + try: + result = await agent_task + except asyncio.CancelledError: + logger.debug("[BG:%s] Agent task was cancelled (%s)", session_id, label) + return + except Exception as exc: + logger.warning("[BG:%s] Agent task failed (%s): %s", session_id, label, exc) + return + + if result is None: + return + + # ------------------------------------------------------------------ + # 2. Sync completed messages back onto the agent instance. + # ------------------------------------------------------------------ + try: + all_msgs = result.all_messages() + agent.set_message_history(all_msgs) + except Exception as exc: + logger.warning( + "[BG:%s] Could not extract messages from result (%s): %s", + session_id, + label, + exc, + ) + # Don't bail — get_message_history() may still hold a partial history. + + history = agent.get_message_history() + if not history: + logger.debug( + "[BG:%s] Empty history after completion (%s) — nothing to save", + session_id, + label, + ) + return + + # ------------------------------------------------------------------ + # 3. Guard: skip write if the session was deleted between task start + # and completion (avoids resurrecting a deliberately deleted session). + # ------------------------------------------------------------------ + try: + if not await session_exists(session_id): + logger.info( + "[BG:%s] Session deleted before background save (%s) — skipping", + session_id, + label, + ) + return + except Exception as exc: + logger.warning( + "[BG:%s] session_exists check failed (%s): %s — proceeding anyway", + session_id, + label, + exc, + ) + + # ------------------------------------------------------------------ + # 4. Build enhanced history wrappers. + # Pre-wrapped dicts (legacy or in-memory system-message injections) + # pass through unchanged; bare ModelMessage objects get wrapped. + # ------------------------------------------------------------------ + now_iso = datetime.datetime.now(datetime.timezone.utc).isoformat() + + enhanced_history: list[dict[str, Any]] = [] + for msg in history: + if isinstance(msg, dict) and "msg" in msg: + enhanced_history.append(msg) + else: + enhanced_history.append( + { + "msg": msg, + "agent": agent_name, + "model": model_name, + "ts": now_iso, + } + ) + + # ------------------------------------------------------------------ + # 5. Compute token count from the completed history. + # The three old closures hard-coded 0; we compute it properly here. + # ------------------------------------------------------------------ + total_tokens = 0 + try: + for item in enhanced_history: + msg_obj = ( + item["msg"] if isinstance(item, dict) and "msg" in item else item + ) + total_tokens += agent.estimate_tokens_for_message(msg_obj) + except Exception: + total_tokens = 0 + + # ------------------------------------------------------------------ + # 6. Resolve created_at safely — ctx may lack the attribute in tests. + # ------------------------------------------------------------------ + ctx_created_at = getattr(ctx, "created_at", None) + created_at = ( + ctx_created_at.isoformat() if ctx_created_at is not None else now_iso + ) + + # ------------------------------------------------------------------ + # 7. Persist to SQLite. + # ------------------------------------------------------------------ + await write_turn_to_sqlite( + session_id=session_id, + enhanced_history=enhanced_history, + title=title, + working_directory=working_directory, + pinned=pinned, + agent_name=agent_name, + model_name=model_name, + total_tokens=total_tokens, + updated_at=now_iso, + created_at=created_at, + ctx=ctx, + ) + logger.info( + "[BG:%s] Background save complete (%s): %d messages, %d tokens", + session_id, + label, + len(enhanced_history), + total_tokens, + ) + + except Exception as exc: + logger.error( + "[BG:%s] Background save failed (%s): %s", + session_id, + label, + exc, + exc_info=True, + ) diff --git a/code_puppy/api/ws/chat_context.py b/code_puppy/api/ws/chat_context.py new file mode 100644 index 000000000..5bdea2020 --- /dev/null +++ b/code_puppy/api/ws/chat_context.py @@ -0,0 +1,72 @@ +"""Helpers for WebSocket chat execution context setup/reset. + +These helpers centralize per-message and per-agent-run ContextVar management so +chat_handler can delegate setup/cleanup without changing behavior. +""" + +from dataclasses import dataclass +from typing import Callable + +from code_puppy.api.permission_plugin import ( + clear_websocket_context, + set_suppress_emitter_tool_events, + set_websocket_context, +) + + +@dataclass(slots=True) +class WebSocketRunContext: + """State needed to clean up one run_with_mcp execution window.""" + + emitter_token: object | None = None + + +def setup_message_context(*, websocket: object, session_id: str) -> None: + """Set per-message websocket permission context.""" + set_websocket_context(websocket, session_id) + + +def cleanup_message_context( + *, clear_session_working_directory: Callable[[], None] +) -> None: + """Clear per-message websocket/prompt-generation context.""" + clear_websocket_context() + clear_session_working_directory() + + +def begin_agent_run_context(*, session_id: str) -> WebSocketRunContext: + """Enable emitter tagging + tool lifecycle suppression for one agent run.""" + emitter_token = None + try: + from code_puppy.plugins.frontend_emitter.session_context import ( + current_emitter_session_id, + ) + + emitter_token = current_emitter_session_id.set(session_id) + except ImportError: + emitter_token = None + + set_suppress_emitter_tool_events(True) + return WebSocketRunContext(emitter_token=emitter_token) + + +def cleanup_agent_run_context( + run_context: WebSocketRunContext | None, + *, + clear_session_working_directory: Callable[[], None], +) -> None: + """Reset all ContextVars touched during one agent run window.""" + set_suppress_emitter_tool_events(False) + if run_context is not None and run_context.emitter_token is not None: + try: + from code_puppy.plugins.frontend_emitter.session_context import ( + current_emitter_session_id, + ) + + current_emitter_session_id.reset(run_context.emitter_token) + except ImportError: + pass + + cleanup_message_context( + clear_session_working_directory=clear_session_working_directory + ) diff --git a/code_puppy/api/ws/chat_event_adapter.py b/code_puppy/api/ws/chat_event_adapter.py new file mode 100644 index 000000000..2e80e7573 --- /dev/null +++ b/code_puppy/api/ws/chat_event_adapter.py @@ -0,0 +1,243 @@ +"""Helpers for adapting streamed agent events into WS assistant frames. + +This module intentionally handles only assistant-message framing and active-part +bookkeeping. Tool lifecycle reconciliation stays in ``chat_handler.py`` for the +next cleanup slice so we can extract the frontend streaming shape first without +changing surrounding orchestration. +""" + +from __future__ import annotations + +import time as time_module +from typing import Any, Awaitable, Callable + +from code_puppy.api.ws.chat_turn_state import WebSocketTurnState +from code_puppy.api.ws.schemas import ( + ServerAssistantMessageDelta, + ServerAssistantMessageEnd, + ServerAssistantMessageStart, +) + + +async def handle_assistant_part_start( + *, + turn_state: WebSocketTurnState, + part_index: int, + part_type: str, + part_obj: Any, + session_id: str, + agent_name: str, + model_name: str, + safe_send_json: Callable[[dict[str, Any]], Awaitable[None]], + logger: Any, + send_status_only_pending_tool_results: Callable[[], Awaitable[None]] | None = None, +) -> bool: + """Handle text/thinking ``part_start`` events. + + Returns ``True`` when the event was fully handled and the caller should skip + further inline processing. + """ + if part_type not in ("TextPart", "ThinkingPart"): + return False + + initial_content = extract_initial_content(part_obj) + msg_type = "thinking" if part_type == "ThinkingPart" else "text" + + if turn_state.pending_tool_calls and part_type == "TextPart": + if send_status_only_pending_tool_results is not None: + await send_status_only_pending_tool_results() + turn_state.current_tool_name = None + + if part_index in turn_state.active_parts: + turn_state.active_parts[part_index]["type"] = msg_type + message_id = turn_state.active_parts[part_index]["id"] + if initial_content: + turn_state.active_parts[part_index]["content"] = ( + initial_content + turn_state.active_parts[part_index]["content"] + ) + turn_state.collected_text.insert(0, initial_content) + logger.debug( + "[Stream Debug] Part already exists, reusing message_id=%s", + message_id, + ) + return True + + message_id = f"msg-{int(time_module.time() * 1000)}-{part_index}" + turn_state.active_parts[part_index] = { + "id": message_id, + "type": msg_type, + "content": initial_content, + } + + if initial_content: + turn_state.collected_text.append(initial_content) + + if part_index == 0: + turn_state.current_tool_group_id = None + + await safe_send_json( + ServerAssistantMessageStart( + message_id=message_id, + part_type=msg_type, + part_index=part_index, + timestamp=time_module.time(), + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + tool_name=turn_state.current_tool_name, + ).model_dump(exclude_none=True) + ) + + if initial_content: + delta = ServerAssistantMessageDelta.model_construct( + type="assistant_message_delta", + message_id=message_id, + content=initial_content, + part_index=part_index, + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + tool_name=turn_state.current_tool_name, + ) + await safe_send_json(delta.model_dump(exclude_none=True)) + turn_state.b1_streaming_used = True + + return True + + +async def handle_assistant_part_delta( + *, + turn_state: WebSocketTurnState, + part_index: int, + inner_data: dict[str, Any], + delta_obj: Any, + session_id: str, + agent_name: str, + model_name: str, + safe_send_json: Callable[[dict[str, Any]], Awaitable[None]], + logger: Any, +) -> bool: + """Handle text/thinking ``part_delta`` events.""" + content_delta = extract_content_delta(inner_data, delta_obj) + if not content_delta: + return False + + turn_state.collected_text.append(content_delta) + + if part_index not in turn_state.active_parts: + message_id = f"msg-{int(time_module.time() * 1000)}-{part_index}" + turn_state.active_parts[part_index] = { + "id": message_id, + "type": "text", + "content": "", + } + logger.debug( + "[Stream Debug] Creating part on first delta: message_id=%s", + message_id, + ) + if part_index == 0: + turn_state.current_tool_group_id = None + await safe_send_json( + ServerAssistantMessageStart( + message_id=message_id, + part_type="text", + part_index=part_index, + timestamp=time_module.time(), + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + tool_name=turn_state.current_tool_name, + ).model_dump(exclude_none=True) + ) + + part_info = turn_state.active_parts[part_index] + message_id = part_info["id"] + part_info["content"] += content_delta + + delta = ServerAssistantMessageDelta.model_construct( + type="assistant_message_delta", + message_id=message_id, + content=content_delta, + part_index=part_index, + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + tool_name=turn_state.current_tool_name, + ) + await safe_send_json(delta.model_dump(exclude_none=True)) + turn_state.b1_streaming_used = True + return True + + +async def handle_assistant_part_end( + *, + turn_state: WebSocketTurnState, + part_index: int, + session_id: str, + agent_name: str, + model_name: str, + safe_send_json: Callable[[dict[str, Any]], Awaitable[None]], +) -> bool: + """Handle non-tool ``part_end`` events and clean up active state.""" + part_info = turn_state.active_parts.get(part_index, {}) + if part_info.get("type", "text") == "tool_call": + return False + + await safe_send_json( + ServerAssistantMessageEnd( + message_id=part_info.get("id", f"msg-{part_index}"), + part_index=part_index, + full_content=part_info.get("content", ""), + timestamp=time_module.time(), + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + tool_name=turn_state.current_tool_name, + ).model_dump(exclude_none=True) + ) + + turn_state.active_parts.pop(part_index, None) + return True + + +def collect_final_stream_text_delta( + *, + turn_state: WebSocketTurnState, + event: dict[str, Any], +) -> bool: + """Append text delta content during the final post-stop queue drain.""" + if event.get("type", "") != "stream_event": + return False + + event_data = event.get("data", {}) + if event_data.get("event_type", "") != "part_delta": + return False + + content_delta = extract_content_delta( + event_data.get("event_data", {}), + event_data.get("event_data", {}).get("delta", {}), + ) + if not content_delta: + return False + + turn_state.collected_text.append(content_delta) + return True + + +def extract_initial_content(part_obj: Any) -> str: + """Return initial content attached to a part object, if any.""" + if hasattr(part_obj, "content") and part_obj.content: + return part_obj.content + if isinstance(part_obj, dict) and part_obj.get("content"): + return part_obj.get("content", "") + return "" + + +def extract_content_delta(inner_data: dict[str, Any], delta_obj: Any) -> str: + """Return assistant text delta from direct or nested event payloads.""" + content_delta = inner_data.get("content_delta", "") + if content_delta: + return content_delta + if isinstance(delta_obj, dict): + return delta_obj.get("content_delta", "") or "" + return "" diff --git a/code_puppy/api/ws/chat_handler.py b/code_puppy/api/ws/chat_handler.py new file mode 100644 index 000000000..60968ab51 --- /dev/null +++ b/code_puppy/api/ws/chat_handler.py @@ -0,0 +1,1346 @@ +"""WebSocket endpoint for interactive chat with the Code Puppy agent. + +This is the largest and most complex WebSocket handler, responsible for: +- Interactive chat sessions with streaming responses +- Session management (create, restore, switch) +- Tool call/result forwarding +- File attachment processing +- Slash command handling +- Permission request/response flow +- Working directory management +- Real-time event streaming from the agent + +NOTE: This handler was extracted from the monolithic websocket.py to improve +maintainability. The internal structure is preserved to avoid regressions. +Future refactoring should break down the message handling loop further. +""" + +import asyncio +import datetime +import json +import logging + +from fastapi import FastAPI, WebSocket, WebSocketDisconnect +from pydantic import TypeAdapter, ValidationError + +from code_puppy.api.session_context import session_manager +from code_puppy.api.ws.ws_stream_drain import ( + start_stream_drain, + stop_stream_drain, +) +from code_puppy.api.ws.ws_post_run import resolve_post_run_resolution +from code_puppy.api.ws.ws_resume_recovery import ( + reload_session_from_sqlite_with_sanitization, +) +from code_puppy.api.ws.ws_turn_finalization import ( + emit_pre_stream_end_tool_results, + finalize_turn_history, +) +from code_puppy.api.ws.chat_event_adapter import ( + collect_final_stream_text_delta, + handle_assistant_part_delta, + handle_assistant_part_end, + handle_assistant_part_start, +) +from code_puppy.api.ws.chat_tool_lifecycle import ( + accumulate_tool_call_part_delta as lifecycle_accumulate_tool_call_part_delta, + finish_tool_call_part as lifecycle_finish_tool_call_part, + handle_tool_call_complete_event as lifecycle_handle_tool_call_complete_event, + handle_tool_call_start_event as lifecycle_handle_tool_call_start_event, + send_status_only_pending_tool_results as lifecycle_send_status_only_pending_tool_results, + start_tool_call_part as lifecycle_start_tool_call_part, + start_tool_return_part as lifecycle_start_tool_return_part, +) +from code_puppy.api.ws.ws_turn_preparation import prepare_turn_input +from code_puppy.api.ws.response_frames import ( + build_assistant_text_stream_frames, + parse_api_error, +) +from code_puppy.api.ws.ws_command_handler import handle_command_message +from code_puppy.api.ws.ws_control_messages import handle_control_message +from code_puppy.api.ws.schemas import ( + ClientMessage, + ServerAgentInvoked, + ServerAssistantMessageDelta, + ServerAssistantMessageEnd, + ServerAssistantMessageStart, + ServerCancelled, + ServerError, + ServerStatus, + ServerSystem, + ServerStreamEnd, + ServerUserMessage, + ServerConfirmationRequest, + ServerSelectionRequest, + ServerUserInputRequest, + ServerAskUserQuestionRequest, +) +from code_puppy.api.ws.send_utils import WebSocketSender +from code_puppy.api.ws.ws_session_bootstrap import initialize_ws_session +from code_puppy.api.ws.chat_context import ( + cleanup_message_context, + setup_message_context, +) +from code_puppy.api.ws.chat_turn_runner import execute_turn_runner +from code_puppy.api.ws.chat_turn_state import WebSocketTurnState +from code_puppy.api.ws.session_persistence import persist_session_turn_and_broadcast +from code_puppy.config import get_global_model_name +from code_puppy.messaging.bus import get_message_bus +from code_puppy.messaging.commands import ( + AskUserQuestionResponse, + ConfirmationResponse, + SelectionResponse, + UserInputResponse, +) +from code_puppy.messaging.messages import ( + AskUserQuestionRequest, + ConfirmationRequest, + SelectionRequest, + UserInputRequest, +) +from code_puppy.tools.command_runner import cleanup_session_process_tracking + +# Per-agent-run working-directory context for shell/tool execution. +# This is intentionally core (not Walmart-plugin-specific) so the desk UI works +# in both the open-source and Walmart worktrees. +from code_puppy.api.session_cwd import ( + clear_session_working_directory, + set_session_working_directory, +) + +HAS_SESSION_CONTEXT = True + + +_ClientMessageAdapter = TypeAdapter(ClientMessage) + +logger = logging.getLogger(__name__) + + +def register_chat_endpoint(app: FastAPI) -> None: + """Register the /ws/chat WebSocket endpoint.""" + + @app.websocket("/ws/chat") + async def websocket_chat( + websocket: WebSocket, session_id: str | None = None + ) -> None: + """Interactive chat with the Code Puppy agent. + + Protocol: + Client sends: + {"type": "message", "content": "your message here"} + + Server sends: + # Streaming message events (all include agent_name, model_name, tool_name metadata) + {"type": "assistant_message_start", "message_id": "...", "part_type": "text|thinking", + "agent_name": "...", "model_name": "...", "tool_name": "..." or null} + {"type": "assistant_message_delta", "message_id": "...", "content": "...", + "agent_name": "...", "model_name": "...", "tool_name": "..." or null} + {"type": "assistant_message_end", "message_id": "...", "full_content": "...", + "agent_name": "...", "model_name": "...", "tool_name": "..." or null} + + # Tool call events (include agent_name, model_name metadata) + {"type": "tool_call", "tool_name": "...", "args": {...}, + "agent_name": "...", "model_name": "..."} + {"type": "tool_result", "tool_name": "...", "result": "...", "success": true, + "agent_name": "...", "model_name": "..."} + + # Final stream marker + {"type": "stream_end", "success": true, "total_length": 123, + "agent_name": "...", "model_name": "...", "tokens": {...}} + + # Errors + {"type": "error", "error": "..."} + + Query Parameters: + session_id: Optional session ID to resume. If not provided, a new session is created. + Example: /ws/chat?session_id=WS_session_20260115_143022 + """ + await websocket.accept() + logger.debug( + "Chat WebSocket client connected (session_id param: %s)", session_id + ) + + # WebSocketSender encapsulates sender.ws_closed, safe_send_json, + # persist_error_payload, send_typed, and send_typed_tool_lifecycle. + sender = WebSocketSender(websocket, session_id) + + # Convenience aliases for call-site compatibility. + safe_send_json = sender.safe_send_json + send_typed = sender.send_typed + send_typed_tool_lifecycle = sender.send_typed_tool_lifecycle + persist_error_payload = sender.persist_error_payload + + runtime = await initialize_ws_session( + websocket=websocket, + requested_session_id=session_id, + sender=sender, + safe_send_json=safe_send_json, + send_typed=send_typed, + ) + if runtime is None: + return + + session_id = runtime.session_id + ctx = runtime.ctx + session_title = runtime.session_title + session_working_directory = runtime.session_working_directory + session_pinned = runtime.session_pinned + last_context_sent_directory = runtime.last_context_sent_directory + existing_history = runtime.existing_history + agent = runtime.agent + agent_name = runtime.agent_name + model_name = runtime.model_name + active_drain_task = runtime.active_drain_task + active_agent_task = runtime.active_agent_task + stop_draining = runtime.stop_draining + + try: + get_message_bus().mark_renderer_active() + except Exception: + logger.debug("Failed to mark MessageBus renderer active", exc_info=True) + + async def forward_message_bus_interactions() -> None: + """Forward pending MessageBus user-interaction prompts to chat.html.""" + try: + bus = get_message_bus() + # Keep the drain bounded so regular stream events stay responsive. + for _ in range(20): + request = bus.get_message_nowait() + if request is None: + break + if isinstance(request, UserInputRequest): + await send_typed( + ServerUserInputRequest( + prompt_id=request.prompt_id, + prompt_text=request.prompt_text, + default_value=request.default_value, + input_type=request.input_type, + session_id=session_id, + ) + ) + elif isinstance(request, ConfirmationRequest): + await send_typed( + ServerConfirmationRequest( + prompt_id=request.prompt_id, + title=request.title, + description=request.description, + options=request.options, + allow_feedback=request.allow_feedback, + session_id=session_id, + ) + ) + elif isinstance(request, SelectionRequest): + await send_typed( + ServerSelectionRequest( + prompt_id=request.prompt_id, + prompt_text=request.prompt_text, + options=request.options, + allow_cancel=request.allow_cancel, + session_id=session_id, + ) + ) + elif isinstance(request, AskUserQuestionRequest): + await send_typed( + ServerAskUserQuestionRequest( + prompt_id=request.prompt_id, + questions=request.questions, + timeout_seconds=request.timeout_seconds, + session_id=session_id, + ) + ) + else: + logger.debug( + "Dropping unsupported MessageBus UI message: %s", + type(request).__name__, + ) + except Exception: + logger.debug("Failed to forward MessageBus interaction", exc_info=True) + + async def send_session_meta_snapshot() -> None: + runtime.session_id = session_id + runtime.ctx = ctx + runtime.session_title = session_title + runtime.session_working_directory = session_working_directory + runtime.session_pinned = session_pinned + runtime.last_context_sent_directory = last_context_sent_directory + runtime.agent = agent + runtime.agent_name = agent_name + runtime.model_name = model_name + runtime.active_drain_task = active_drain_task + runtime.active_agent_task = active_agent_task + from code_puppy.api.ws.ws_session_bootstrap import ( + send_session_meta_snapshot as _send_session_meta_snapshot, + ) + + await _send_session_meta_snapshot( + runtime=runtime, + safe_send_json=safe_send_json, + ) + + try: + while True: + try: + msg = await websocket.receive_json() + + # Advisory validation — log but never reject + try: + _parsed = _ClientMessageAdapter.validate_python(msg) + except ValidationError as _val_err: + logger.warning( + "Client message failed validation: %s", + str(_val_err), + extra={ + "type": msg.get("type") + if isinstance(msg, dict) + else "unknown" + }, + ) + + if msg.get("type") in { + "user_input_response", + "confirmation_response", + "selection_response", + "ask_user_question_response", + }: + bus = get_message_bus() + try: + if msg.get("type") == "user_input_response": + bus.provide_response( + UserInputResponse( + prompt_id=msg.get("prompt_id", ""), + value=msg.get("value", ""), + ) + ) + elif msg.get("type") == "confirmation_response": + bus.provide_response( + ConfirmationResponse( + prompt_id=msg.get("prompt_id", ""), + confirmed=bool(msg.get("confirmed", False)), + feedback=msg.get("feedback"), + ) + ) + elif msg.get("type") == "selection_response": + bus.provide_response( + SelectionResponse( + prompt_id=msg.get("prompt_id", ""), + selected_index=int( + msg.get("selected_index", -1) + ), + selected_value=msg.get("selected_value", ""), + ) + ) + else: + bus.provide_response( + AskUserQuestionResponse( + prompt_id=msg.get("prompt_id", ""), + answers=msg.get("answers") or [], + cancelled=bool(msg.get("cancelled", False)), + ) + ) + except Exception as exc: + logger.warning("Invalid user interaction response: %s", exc) + await send_typed( + ServerError( + error=f"Invalid user interaction response: {exc}", + session_id=session_id, + ) + ) + continue + + if await handle_command_message( + msg=msg, + session_id=session_id, + send_typed=send_typed, + ): + continue + + runtime.session_id = session_id + runtime.ctx = ctx + runtime.session_title = session_title + runtime.session_working_directory = session_working_directory + runtime.session_pinned = session_pinned + runtime.last_context_sent_directory = last_context_sent_directory + runtime.agent = agent + runtime.agent_name = agent_name + runtime.model_name = model_name + runtime.active_drain_task = active_drain_task + runtime.active_agent_task = active_agent_task + + if await handle_control_message( + msg=msg, + runtime=runtime, + sender=sender, + send_typed=send_typed, + send_session_meta_snapshot=send_session_meta_snapshot, + ): + session_id = runtime.session_id + ctx = runtime.ctx + session_title = runtime.session_title + session_working_directory = runtime.session_working_directory + session_pinned = runtime.session_pinned + last_context_sent_directory = ( + runtime.last_context_sent_directory + ) + agent = runtime.agent + agent_name = runtime.agent_name + model_name = runtime.model_name + active_drain_task = runtime.active_drain_task + active_agent_task = runtime.active_agent_task + continue + + elif msg.get("type") == "message": + # Set WebSocket context for permission requests + setup_message_context( + websocket=websocket, session_id=session_id + ) + + user_message = msg.get("content", "") + + # Initialize attachment tracking variables at message scope + original_user_message = user_message # Store clean message + attachment_metadata = [] # Will be populated if attachments exist + + # Check if a specific model was requested for this message + requested_model = msg.get("model") + if requested_model: + current_model = ctx.model_name + if requested_model != current_model: + try: + await session_manager.switch_model( + session_id, requested_model + ) + agent = ctx.agent # Refresh alias + model_name = requested_model + logger.debug( + f"Switching to model {requested_model} for this message" + ) + except Exception as e: + logger.warning("Failed to switch model: %s", e) + + # Apply model_settings from frontend (reasoning_effort, verbosity, etc.) + model_settings = msg.get("model_settings", {}) + if model_settings: + from code_puppy.config import set_model_setting + + target_model = requested_model or get_global_model_name() + for setting_name, value in model_settings.items(): + try: + set_model_setting(target_model, setting_name, value) + logger.debug( + f"Applied model_setting {setting_name}={value} for {target_model}" + ) + except Exception as e: + logger.warning( + f"Failed to apply model_setting {setting_name}: {e}" + ) + + if not user_message.strip(): + await send_typed( + ServerError( + error="Empty message", + session_id=session_id, + ) + ) + continue + + logger.debug( + f"Chat message from client: {user_message[:50]}..." + ) + + # Echo the user message back + await send_typed( + ServerUserMessage( + content=user_message, + session_id=session_id, + ) + ) + + # Get the session agent and process the message + try: + agent = ctx.agent + + # Reload agent if a different model was requested + if requested_model: + agent.reload_code_generation_agent() + logger.debug( + f"Reloaded agent with model: {requested_model}" + ) + + if not agent: + await send_typed( + ServerError( + error="Agent not available. Please start Code Puppy first.", + session_id=session_id, + ) + ) + continue + + # Subscribe to frontend emitter and run drain task CONCURRENTLY + turn_state = WebSocketTurnState() + stop_draining.clear() # Reset for this message + drain_task = None + drain_handle = None + + async def drain_events_concurrent( + stream_event_queue, + ready_event: asyncio.Event = None, + ): + """Background task to drain events and send structured messages in real-time.""" + import time as time_module + + # Capture agent and model metadata at the start. + # Chain `or` fallbacks so that an empty-string agent.name + # (possible before agent is fully initialised) still resolves + # to a meaningful value via the session context defaults. + current_agent_name = ( + (agent.name if agent else "") + or ctx.agent_name + or "code-puppy" + ) + current_model_name = ( + (agent.get_model_name() if agent else "") + or ctx.model_name + or "unknown" + ) + + async def send_status_only_pending_tool_results(): + await ( + lifecycle_send_status_only_pending_tool_results( + turn_state=turn_state, + session_id=session_id, + agent_name=current_agent_name, + model_name=current_model_name, + send_typed=send_typed, + logger=logger, + ) + ) + + event_count = 0 + first_iteration = True + while not stop_draining.is_set(): + # Exit if WebSocket is closed + if sender.ws_closed: + logger.debug( + "WebSocket closed, exiting drain loop" + ) + break + # Signal ready on first iteration (we're in the loop now) + if first_iteration and ready_event: + ready_event.set() + first_iteration = False + # Yield to allow run_with_mcp to start and emit first events + await asyncio.sleep(0) + continue # Re-enter loop to collect any queued events + # Event-driven batch collection with 10ms timeout for responsiveness. + # Instead of polling the clock, we use asyncio.wait_for to block until + # the first event arrives, then collect any additional events already queued. + events_to_send = [] + try: + # Wait for first event with 10ms timeout (blocks efficiently) + first_event = await asyncio.wait_for( + stream_event_queue.get(), timeout=0.01 + ) + events_to_send.append(first_event) + event_count += 1 + + # Collect any additional events already in queue (non-blocking) + # This keeps the batching benefit without polling the clock + while not stream_event_queue.empty(): + try: + event = stream_event_queue.get_nowait() + events_to_send.append(event) + event_count += 1 + except Exception: + break + + except asyncio.TimeoutError: + # No events available within 10ms timeout + # No polling needed - asyncio.wait_for blocks efficiently + pass + + await forward_message_bus_interactions() + + # Log batch composition for debugging + if events_to_send: + event_types = {} + for e in events_to_send: + et = e.get("type", "unknown") + event_types[et] = event_types.get(et, 0) + 1 + logger.debug( + "[%s] Batch: %d events - %s", + session_id, + len(events_to_send), + event_types, + ) + + # Process collected events + for event in events_to_send: + # If cancellation was requested, stop processing immediately + if stop_draining.is_set(): + logger.debug( + "stop_draining set during batch - stopping event processing" + ) + break + + event_type = event.get("type", "") + event_data = event.get("data", {}) + + try: + # Handle tool call events + if event_type == "tool_call_start": + await lifecycle_handle_tool_call_start_event( + turn_state=turn_state, + event_data=event_data, + session_id=session_id, + agent_name=current_agent_name, + model_name=current_model_name, + send_typed_tool_lifecycle=send_typed_tool_lifecycle, + logger=logger, + ) + elif event_type == "tool_call_complete": + await lifecycle_handle_tool_call_complete_event( + turn_state=turn_state, + event_data=event_data, + session_id=session_id, + agent_name=current_agent_name, + model_name=current_model_name, + send_typed_tool_lifecycle=send_typed_tool_lifecycle, + logger=logger, + ) + + elif event_type == "agent_invoked": + agent_name_inv = event_data.get( + "agent_name", "unknown" + ) + prompt_preview = event_data.get( + "prompt_preview", "" + ) + + logger.debug( + "[ws] agent_invoked: %s", + agent_name_inv, + ) + + await send_typed( + ServerAgentInvoked( + agent_name=agent_name_inv, + prompt_preview=prompt_preview, + timestamp=time_module.time(), + session_id=session_id, + ) + ) + + # Handle streaming events + elif event_type == "stream_event": + inner_type = event_data.get( + "event_type", "" + ) + inner_data = event_data.get( + "event_data", {} + ) + if inner_type == "part_start": + part_index = inner_data.get( + "index", 0 + ) + part_type = inner_data.get( + "part_type", "unknown" + ) + logger.warning( + "[ws] part_start: part_type=%s, part_index=%s", + part_type, + part_index, + ) + + # Extract initial content from the part if present + # The part object may have content already (especially for TextPart/ThinkingPart) + part_obj = inner_data.get( + "part", {} + ) + + if await handle_assistant_part_start( + turn_state=turn_state, + part_index=part_index, + part_type=part_type, + part_obj=part_obj, + session_id=session_id, + agent_name=current_agent_name, + model_name=current_model_name, + safe_send_json=safe_send_json, + logger=logger, + send_status_only_pending_tool_results=send_status_only_pending_tool_results, + ): + continue + if part_type == "ToolCallPart": + lifecycle_start_tool_call_part( + turn_state=turn_state, + part_index=part_index, + part_obj=part_obj, + logger=logger, + ) + elif part_type == "ToolReturnPart": + await lifecycle_start_tool_return_part( + turn_state=turn_state, + part_index=part_index, + part_obj=part_obj, + session_id=session_id, + agent_name=current_agent_name, + model_name=current_model_name, + send_typed_tool_lifecycle=send_typed_tool_lifecycle, + logger=logger, + ) + + elif inner_type == "part_delta": + part_index = inner_data.get( + "index", 0 + ) + delta_type = inner_data.get( + "delta_type", "" + ) + delta_obj = inner_data.get( + "delta", {} + ) + if ( + delta_type + == "ToolCallPartDelta" + ): + if lifecycle_accumulate_tool_call_part_delta( + turn_state=turn_state, + part_index=part_index, + delta_obj=delta_obj, + ): + continue # Don't process as text delta + + if await handle_assistant_part_delta( + turn_state=turn_state, + part_index=part_index, + inner_data=inner_data, + delta_obj=delta_obj, + session_id=session_id, + agent_name=current_agent_name, + model_name=current_model_name, + safe_send_json=safe_send_json, + logger=logger, + ): + continue + + elif inner_type == "part_end": + part_index = inner_data.get( + "index", 0 + ) + part_info = ( + turn_state.active_parts.get( + part_index, {} + ) + ) + part_type_info = part_info.get( + "type", "text" + ) + + if await handle_assistant_part_end( + turn_state=turn_state, + part_index=part_index, + session_id=session_id, + agent_name=current_agent_name, + model_name=current_model_name, + safe_send_json=safe_send_json, + ): + continue + if part_type_info == "tool_call": + await lifecycle_finish_tool_call_part( + turn_state=turn_state, + part_index=part_index, + part_info=part_info, + session_id=session_id, + agent_name=current_agent_name, + model_name=current_model_name, + send_typed_tool_lifecycle=send_typed_tool_lifecycle, + logger=logger, + ) + else: + turn_state.active_parts.pop( + part_index, None + ) + + except Exception as send_err: + error_msg = str(send_err).lower() + if ( + "close message" in error_msg + or "closed" in error_msg + ): + sender.ws_closed = True + logger.debug( + "WebSocket closed during streaming, stopping drain" + ) + break + logger.warning( + f"Error sending event to WebSocket: {type(send_err).__name__}: {send_err}" + ) + import traceback + + logger.warning( + f"Traceback: {traceback.format_exc()}" + ) + + # No idle polling - event-driven approach handles empty event sets gracefully + + # Final drain after stop signal + final_count = 0 + while True: + try: + event = stream_event_queue.get_nowait() + event_type = event.get("type", "") + event_data = event.get("data", {}) + final_count += 1 + + if event_type == "stream_event": + collect_final_stream_text_delta( + turn_state=turn_state, + event=event, + ) + except Exception: + break + + # Calculate batching efficiency + avg_batch_size = ( + event_count / max(1, event_count) + if event_count > 0 + else 0 + ) + logger.debug( + f"[{session_id}] Drain complete: " + f"{event_count} events during run, {final_count} final, " + f"batching efficiency: {avg_batch_size:.2f}" + ) + + drain_handle = await start_stream_drain( + session_id=session_id, + drain_coro_factory=drain_events_concurrent, + logger=logger, + ) + if drain_handle is not None: + drain_task = drain_handle.task + active_drain_task = drain_task + + try: + # Send status: thinking + await send_typed( + ServerStatus( + status="thinking", + session_id=session_id, + agent_name=agent.name + if agent + else "code-puppy", + model_name=agent.get_model_name() + if agent + else "unknown", + ) + ) + + # Change to session working directory if set + # Set session context for prompt generation (desk-puppy) + set_session_working_directory(session_working_directory) + + try: + # Call run_with_mcp (drain task runs concurrently!) + logger.debug( + "About to run agent, working_directory=%s", + session_working_directory, + ) + + _prepared_turn = prepare_turn_input( + agent=agent, + user_message=user_message, + msg=msg, + session_working_directory=session_working_directory, + last_context_sent_directory=last_context_sent_directory, + ) + message_to_send = _prepared_turn.message_to_send + run_kwargs = _prepared_turn.run_kwargs + attachment_metadata = ( + _prepared_turn.attachment_metadata + ) + last_context_sent_directory = ( + _prepared_turn.last_context_sent_directory + ) + + # ────────────────────────────────────────────────────────────────── + # Phase 7: Create session + user message in SQLite BEFORE streaming + # This prevents "Session not found" errors when FE queries mid-stream + # ────────────────────────────────────────────────────────────────── + try: + import os + from datetime import timezone as tz + + from code_puppy.api.db.queries import ( + upsert_session, + ) + + now_iso = datetime.datetime.now( + tz.utc + ).isoformat() + + # Ensure session row exists + await upsert_session( + session_id=session_id, + title="", # Will be updated later with actual title + agent_name=ctx.agent_name, + model_name=ctx.model_name, + working_directory=os.getcwd(), + pinned=False, + created_at=now_iso, + updated_at=now_iso, + message_count=0, # Will be updated after streaming + total_tokens=0, + ) + + # Write user message immediately (before agent processes it) + # Use get_next_seq() + insert_message() so the user message + # lands at MAX(seq)+1, AFTER any system messages (config, + # directory banners) that were already written at seq=1, 2, … + # The old write_turn_to_sqlite([single_item]) always assigned + # seq=1, silently colliding with those rows via INSERT OR IGNORE. + from pydantic_ai.messages import ( + ModelRequest, + UserPromptPart, + ) + + from code_puppy.api.db.message_utils import ( + pydantic_json_for_message, + ) + from code_puppy.api.db.queries import ( + get_next_seq, + insert_message, + ) + + user_msg_obj = ModelRequest( + parts=[ + UserPromptPart( + content=original_user_message + ) + ] + ) + + user_seq = await get_next_seq(session_id) + await insert_message( + session_id=session_id, + seq=user_seq, + role="user", + content=original_user_message, + type="ModelRequest", + agent_name=ctx.agent_name, + model_name=ctx.model_name, + timestamp=now_iso, + clean_content=original_user_message, + attachments_json=( + json.dumps(attachment_metadata) + if attachment_metadata + else None + ), + pydantic_json=pydantic_json_for_message( + user_msg_obj + ), + ) + + logger.debug( + "Pre-stream write: session %s created with user message in SQLite", + session_id, + ) + except Exception as pre_write_exc: + # Non-fatal: WS streaming will still work, SQLite just won't have the + # session yet. The post-stream write will create it. + logger.warning( + "Pre-stream SQLite write failed for %s: %s", + session_id, + pre_write_exc, + exc_info=True, + ) + + _turn_run = await execute_turn_runner( + websocket=websocket, + session_id=session_id, + ctx=ctx, + agent=agent, + agent_name=agent_name, + model_name=model_name, + session_title=session_title, + session_working_directory=session_working_directory, + session_pinned=session_pinned, + message_to_send=message_to_send, + run_kwargs=run_kwargs, + turn_state=turn_state, + clear_session_working_directory=clear_session_working_directory, + ) + result = _turn_run.result + _deferred_msg = _turn_run.deferred_msg + + finally: + pass + + finally: + await stop_stream_drain( + handle=drain_handle, + stop_draining=stop_draining, + logger=logger, + ) + + # Process any deferred switch_session that arrived during streaming + if _deferred_msg is not None: + msg = _deferred_msg + _deferred_msg = None + # Re-dispatch to outer loop by continuing with msg set + # The outer while True loop will handle switch_session/create_session + continue + + post_run = resolve_post_run_resolution( + result=result, + turn_state=turn_state, + agent=agent, + session_id=session_id, + logger=logger, + ) + if post_run.cancelled: + await send_typed( + ServerCancelled( + session_id=session_id, + ) + ) + continue + if post_run.error_frames is not None: + for frame in post_run.error_frames: + await safe_send_json(frame) + continue + if post_run.no_result_error is not None: + recovery_applied = False + # Safe one-shot auto-recovery for resumed sessions only. + if existing_history is not None: + await send_typed( + ServerSystem( + content=( + "Attempting safe session recovery from SQLite history " + "(one retry)..." + ), + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + ) + ) + recovery = await reload_session_from_sqlite_with_sanitization( + session_id=session_id, + logger=logger, + ) + if recovery.success and recovery.ctx is not None: + ctx = recovery.ctx + sender.ctx = ctx + agent = ctx.agent + agent_name = ctx.agent_name + model_name = ctx.model_name + + retry_turn_state = WebSocketTurnState() + retry_run = await execute_turn_runner( + websocket=websocket, + session_id=session_id, + ctx=ctx, + agent=agent, + agent_name=agent_name, + model_name=model_name, + session_title=session_title, + session_working_directory=session_working_directory, + session_pinned=session_pinned, + message_to_send=message_to_send, + run_kwargs=run_kwargs, + turn_state=retry_turn_state, + clear_session_working_directory=clear_session_working_directory, + ) + retry_post_run = resolve_post_run_resolution( + result=retry_run.result, + turn_state=retry_turn_state, + agent=agent, + session_id=session_id, + logger=logger, + ) + if retry_post_run.cancelled: + await send_typed( + ServerCancelled( + session_id=session_id, + ) + ) + continue + if retry_post_run.error_frames is not None: + for frame in retry_post_run.error_frames: + await safe_send_json(frame) + continue + if retry_post_run.no_result_error is None: + post_run = retry_post_run + recovery_applied = True + await send_typed( + ServerSystem( + content=( + "Session auto-recovery succeeded; continuing with " + "reloaded DB history." + ), + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + ) + ) + else: + logger.warning( + "[WS:%s] one-shot recovery retry still produced no response", + session_id, + ) + else: + logger.warning( + "[WS:%s] recovery reload failed: %s", + session_id, + recovery.reason, + ) + + if not recovery_applied: + await send_typed(post_run.no_result_error) + continue + + response_text = post_run.response_text + tokens_used = post_run.tokens_used + thinking_text = post_run.thinking_text + + # Only send legacy 'response' if B1 streaming wasn't used + # B1 streaming already sent the content via assistant_message_end + logger.warning( + "[WebSocket] turn_state.b1_streaming_used=%s before response/extraction", + turn_state.b1_streaming_used, + ) + if not turn_state.b1_streaming_used: + # Send thinking content as B1 message if available + if thinking_text: + import time as time_module + + thinking_message_id = f"thinking-{session_id}-{int(time_module.time() * 1000)}" + await send_typed( + ServerAssistantMessageStart( + message_id=thinking_message_id, + part_type="thinking", + part_index=0, + timestamp=time_module.time(), + session_id=session_id, + agent_name=agent.name + if agent + else "code-puppy", + model_name=agent.get_model_name() + if agent + else "unknown", + ) + ) + _delta = ( + ServerAssistantMessageDelta.model_construct( + type="assistant_message_delta", + message_id=thinking_message_id, + content=thinking_text, + part_index=0, + session_id=session_id, + agent_name=agent.name + if agent + else "code-puppy", + model_name=agent.get_model_name() + if agent + else "unknown", + ) + ) + await safe_send_json( + _delta.model_dump(exclude_none=True) + ) + await send_typed( + ServerAssistantMessageEnd( + message_id=thinking_message_id, + part_index=0, + full_content=thinking_text, + timestamp=time_module.time(), + session_id=session_id, + agent_name=agent.name + if agent + else "code-puppy", + model_name=agent.get_model_name() + if agent + else "unknown", + ) + ) + logger.debug( + f"Sent thinking content ({len(thinking_text)} chars) for non-streaming model" + ) + + # Adapt complete non-streaming upstream responses into + # the streaming-only GUI protocol. Assistant text is no + # longer delivered via legacy `response.content`. + for frame in build_assistant_text_stream_frames( + response_text=response_text, + session_id=session_id, + agent_name=agent.name if agent else "code-puppy", + model_name=agent.get_model_name() + if agent + else "unknown", + tokens=tokens_used, + ): + await send_typed(frame) + else: + # B1 streaming: extract real tool results BEFORE stream_end + # so the frontend session store is still alive when they arrive. + _pre_sent_tool_ids = await emit_pre_stream_end_tool_results( + result=result, + turn_state=turn_state, + session_id=session_id, + agent_name=agent.name if agent else "code-puppy", + model_name=agent.get_model_name() + if agent + else "unknown", + send_typed_tool_lifecycle=send_typed_tool_lifecycle, + logger=logger, + ) + + # Send stream_end AFTER real tool results are delivered + await send_typed( + ServerStreamEnd( + success=True, + total_length=len(response_text), + agent_name=agent.name + if agent + else "code-puppy", + model_name=agent.get_model_name() + if agent + else "unknown", + tokens=tokens_used, + session_id=session_id, + ) + ) + # Save session after each response + try: + finalized_turn = await finalize_turn_history( + result=result, + agent=agent, + turn_state=turn_state, + session_id=session_id, + agent_name=agent.name if agent else "code-puppy", + model_name=agent.get_model_name() + if agent + else "unknown", + send_typed=send_typed, + pre_sent_tool_ids=locals().get( + "_pre_sent_tool_ids", set() + ), + logger=logger, + ) + + history = ( + finalized_turn.history_snapshot + if finalized_turn.history_snapshot + else agent.get_message_history() + ) # Use pre-await snapshot to avoid race condition + persisted_turn = await persist_session_turn_and_broadcast( + history=history, + session_id=session_id, + session_title=session_title, + session_working_directory=session_working_directory, + session_pinned=session_pinned, + agent=agent, + agent_name=agent_name, + model_name=model_name, + ctx=ctx, + original_user_message=original_user_message, + attachment_metadata=attachment_metadata, + safe_send_json=safe_send_json, + logger_override=logger, + ) + if persisted_turn is not None: + session_title = persisted_turn.session_title + except Exception as save_err: + logger.warning( + f"Failed to save WebSocket session: {save_err}" + ) + + # Send final status to signal completion + await send_typed( + ServerStatus( + status="done", + session_id=session_id, + agent_name=agent.name if agent else "code-puppy", + model_name=agent.get_model_name() + if agent + else "unknown", + ) + ) + + except Exception as e: + logger.error( + f"Error processing message: {e}", exc_info=True + ) + + # Parse error and send user-friendly message + parsed_error = parse_api_error(e) + _err_msg = ServerError( + error=parsed_error["user_message"], + error_type=parsed_error["error_type"], + technical_details=parsed_error["technical_details"], + action_required=parsed_error.get("action_required"), + session_id=session_id, + ) + error_payload = _err_msg.model_dump(exclude_none=True) + await persist_error_payload(error_payload) + await send_typed(_err_msg) + finally: + cleanup_message_context( + clear_session_working_directory=clear_session_working_directory + ) + + except WebSocketDisconnect: + break + except RuntimeError as e: + # Handle disconnect-related RuntimeErrors (starlette raises these) + if ( + "disconnect" in str(e).lower() + or "websocket.close" in str(e).lower() + ): + logger.debug("WebSocket disconnected (RuntimeError): %s", e) + break + # Re-raise other RuntimeErrors to be handled as generic exceptions + raise + except Exception as e: + logger.error("Chat WebSocket error: %s", e, exc_info=True) + # Don't try to send error if websocket is already closed + if sender.ws_closed: + break + try: + _err_msg = ServerError( + error=str(e), + error_type="unknown", + technical_details=str(e), + session_id=session_id, + ) + error_payload = _err_msg.model_dump(exclude_none=True) + await persist_error_payload(error_payload) + await send_typed(_err_msg) + except Exception: + break + + except WebSocketDisconnect: + sender.ws_closed = True + logger.debug("Chat WebSocket client disconnected") + except Exception as e: + logger.error("Chat WebSocket error: %s", e, exc_info=True) + finally: + logger.debug("Chat session %s ended", session_id) + + # --- SESSION ISOLATION CLEANUP --- + # Save and tear down session-scoped resources + try: + await session_manager.save_session(session_id) + except Exception: + logger.debug("Failed to save session on disconnect", exc_info=True) + + # Don't destroy immediately - mark as inactive for 15-min retention + await session_manager.mark_session_inactive(session_id) + cleanup_session_process_tracking() + + try: + bus = get_message_bus() + bus.set_session_context(None) + bus.mark_renderer_inactive() + except Exception: + pass diff --git a/code_puppy/api/ws/chat_tool_lifecycle.py b/code_puppy/api/ws/chat_tool_lifecycle.py new file mode 100644 index 000000000..cee419090 --- /dev/null +++ b/code_puppy/api/ws/chat_tool_lifecycle.py @@ -0,0 +1,520 @@ +"""Helpers for websocket tool lifecycle reconciliation. + +This module owns the tool-call/result bookkeeping that previously lived inside +``chat_handler.py``'s streaming drain loop. The goal is to preserve exact +frontend-visible behavior while isolating the most stateful tool lifecycle +logic from the endpoint handler. +""" + +from __future__ import annotations + +import json +import re +import time as time_module +import uuid +from typing import Any, Awaitable, Callable + +from code_puppy.api.ws.chat_turn_state import WebSocketTurnState +from code_puppy.api.ws.schemas import ServerToolCall, ServerToolResult + + +async def handle_tool_call_start_event( + *, + turn_state: WebSocketTurnState, + event_data: dict[str, Any], + session_id: str, + agent_name: str, + model_name: str, + send_typed_tool_lifecycle: Callable[[Any], Awaitable[None]], + logger: Any, +) -> bool: + """Handle direct ``tool_call_start`` events from the stream queue.""" + tool_name = event_data.get("tool_name", "unknown") + tool_args = event_data.get("tool_args", {}) + tool_id = str(uuid.uuid4())[:8] + + if turn_state.current_tool_group_id is None: + turn_state.current_tool_group_id = f"tg-{str(uuid.uuid4())[:8]}" + + logger.debug("[ws] tool_call: %s", tool_name) + turn_state.current_tool_name = tool_name + turn_state.pending_tool_calls[tool_id] = { + "tool_name": tool_name, + "start_time": time_module.time(), + "tool_group_id": turn_state.current_tool_group_id, + } + + await send_typed_tool_lifecycle( + ServerToolCall( + tool_id=tool_id, + tool_name=tool_name, + args=tool_args, + timestamp=time_module.time(), + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + tool_group_id=turn_state.current_tool_group_id, + ) + ) + return True + + +async def handle_tool_call_complete_event( + *, + turn_state: WebSocketTurnState, + event_data: dict[str, Any], + session_id: str, + agent_name: str, + model_name: str, + send_typed_tool_lifecycle: Callable[[Any], Awaitable[None]], + logger: Any, +) -> bool: + """Handle direct ``tool_call_complete`` events from the stream queue.""" + tool_name = event_data.get("tool_name", "unknown") + result = event_data.get("result", event_data.get("result_summary", "")) + success = event_data.get("success", True) + duration = event_data.get("duration_ms", 0) + + logger.debug("[ws] tool_result: %s", tool_name) + turn_state.current_tool_name = None + + matching_tool_id = next( + ( + tid + for tid, info in turn_state.pending_tool_calls.items() + if info["tool_name"] == tool_name + ), + None, + ) + matching_pending_info = ( + turn_state.pending_tool_calls.get(matching_tool_id) + if matching_tool_id + else None + ) + tool_group_id_for_result = resolve_tool_group_id( + turn_state=turn_state, + logger=logger, + tool_id=matching_tool_id, + pending_info=matching_pending_info, + fallback_group_id=turn_state.current_tool_group_id, + tool_name=tool_name, + source="tool_call_complete", + ) + + if matching_tool_id: + turn_state.pending_tool_calls.pop(matching_tool_id, None) + + await send_typed_tool_lifecycle( + ServerToolResult( + tool_id=matching_tool_id, + tool_name=tool_name, + result=result, + success=success, + duration_ms=duration, + timestamp=time_module.time(), + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + tool_group_id=tool_group_id_for_result, + ) + ) + return True + + +def handle_tool_part_start( + *, + turn_state: WebSocketTurnState, + part_index: int, + part_obj: Any, + logger: Any, +) -> bool: + """Track ``ToolCallPart``/``ToolReturnPart`` state at ``part_start``.""" + part_type = _extract_attr_or_key(part_obj, "part_type") + # ``part_type`` is often passed separately by the caller; this helper relies + # on the explicit caller branch instead of introspecting it. + del part_type + logger.debug("[ws-tool] handle_tool_part_start called") + return True + + +def start_tool_call_part( + *, + turn_state: WebSocketTurnState, + part_index: int, + part_obj: Any, + logger: Any, +) -> bool: + """Record an in-flight tool call part until args are complete.""" + tool_name = _extract_attr_or_key(part_obj, "tool_name") or "unknown" + tool_call_id = _extract_attr_or_key(part_obj, "tool_call_id") + tool_args_str = _extract_attr_or_key(part_obj, "args") or "" + tool_id = tool_call_id or str(uuid.uuid4())[:8] + + turn_state.active_parts[part_index] = { + "id": tool_id, + "raw_tool_call_id": tool_call_id, + "type": "tool_call", + "tool_name": tool_name, + "args": tool_args_str, + "args_buffer": tool_args_str, + "start_time": time_module.time(), + } + turn_state.current_tool_name = tool_name + + logger.debug( + "[WebSocket] ToolCallPart started: %s (id: %s)", + tool_name, + tool_id, + ) + return True + + +async def start_tool_return_part( + *, + turn_state: WebSocketTurnState, + part_index: int, + part_obj: Any, + session_id: str, + agent_name: str, + model_name: str, + send_typed_tool_lifecycle: Callable[[Any], Awaitable[None]], + logger: Any, +) -> bool: + """Handle ``ToolReturnPart`` result emission on ``part_start``.""" + logger.info("[WebSocket] ToolReturnPart detected! part_index=%s", part_index) + + tool_call_id = _extract_attr_or_key(part_obj, "tool_call_id") + tool_content = _coerce_tool_content( + _extract_attr_or_key(part_obj, "content"), logger + ) + + result_sent = False + if tool_call_id: + resolved_pending_id = resolve_pending_tool_id( + turn_state=turn_state, + tool_call_id=tool_call_id, + ) + if resolved_pending_id: + result_sent = True + pending_info = turn_state.pending_tool_calls[resolved_pending_id] + pending_info["result"] = tool_content + tool_name = pending_info.get("tool_name", "unknown") + start_time = pending_info.get("start_time", time_module.time()) + duration_ms = (time_module.time() - start_time) * 1000 + group_id = resolve_tool_group_id( + turn_state=turn_state, + logger=logger, + tool_id=resolved_pending_id, + pending_info=pending_info, + fallback_group_id=turn_state.current_tool_group_id, + tool_name=tool_name, + source="tool_return_resolved", + ) + logger.info( + "[WebSocket] ToolReturnPart: Sending result for %s (id: %s, raw: %s)", + tool_name, + resolved_pending_id, + tool_call_id, + ) + await send_typed_tool_lifecycle( + ServerToolResult( + tool_id=resolved_pending_id, + tool_name=tool_name, + result=tool_content, + success=True, + duration_ms=duration_ms, + timestamp=time_module.time(), + session_id=session_id, + agent_name=agent_name or "code-puppy", + model_name=model_name or "unknown", + tool_group_id=group_id, + ) + ) + else: + logger.warning( + "[WebSocket] ToolReturnPart: Could not resolve tool_call_id %s, pending keys: %s", + tool_call_id, + list(turn_state.pending_tool_calls.keys()), + ) + + if not result_sent: + for pending_id, pending_info in sorted( + turn_state.pending_tool_calls.items(), + key=lambda item: abs(item[1].get("part_index", 9999) - part_index), + ): + if abs(pending_info.get("part_index", 9999) - part_index) <= 3: + pending_info["result"] = tool_content + tool_name = pending_info.get("tool_name", "unknown") + start_time = pending_info.get("start_time", time_module.time()) + duration_ms = (time_module.time() - start_time) * 1000 + group_id = resolve_tool_group_id( + turn_state=turn_state, + logger=logger, + tool_id=pending_id, + pending_info=pending_info, + fallback_group_id=turn_state.current_tool_group_id, + tool_name=tool_name, + source="tool_return_proximity", + ) + logger.info( + "[WebSocket] ToolReturnPart: Sending result (by proximity) for %s (id: %s)", + tool_name, + pending_id, + ) + await send_typed_tool_lifecycle( + ServerToolResult( + tool_id=pending_id, + tool_name=tool_name, + result=tool_content, + success=True, + duration_ms=duration_ms, + timestamp=time_module.time(), + session_id=session_id, + agent_name=agent_name or "code-puppy", + model_name=model_name or "unknown", + tool_group_id=group_id, + ) + ) + result_sent = True + break + + if not result_sent: + logger.warning( + "[WebSocket] ToolReturnPart: Could NOT send result! tool_call_id=%s, " + "turn_state.pending_tool_calls=%s, part_index=%s", + tool_call_id, + list(turn_state.pending_tool_calls.keys()), + part_index, + ) + + turn_state.active_parts[part_index] = { + "id": f"tool-return-{part_index}", + "type": "tool_return", + "tool_call_id": tool_call_id, + "content": tool_content, + } + return True + + +def accumulate_tool_call_part_delta( + *, + turn_state: WebSocketTurnState, + part_index: int, + delta_obj: Any, +) -> bool: + """Append ``ToolCallPartDelta.args_delta`` to the matching active part.""" + args_delta = _extract_attr_or_key(delta_obj, "args_delta") or "" + if not args_delta or part_index not in turn_state.active_parts: + return False + + part_info = turn_state.active_parts[part_index] + if part_info.get("type") != "tool_call": + return False + + part_info["args_buffer"] = part_info.get("args_buffer", "") + args_delta + return True + + +async def finish_tool_call_part( + *, + turn_state: WebSocketTurnState, + part_index: int, + part_info: dict[str, Any], + session_id: str, + agent_name: str, + model_name: str, + send_typed_tool_lifecycle: Callable[[Any], Awaitable[None]], + logger: Any, +) -> bool: + """Emit the deferred ``tool_call`` once streamed args are complete.""" + tool_name = part_info.get("tool_name", "unknown") + tool_id = part_info.get("id", str(uuid.uuid4())[:8]) + tool_args_str = part_info.get("args_buffer", "") or part_info.get("args", "") + start_time = part_info.get("start_time", time_module.time()) + raw_tool_call_id = part_info.get("raw_tool_call_id") + + try: + args_dict = json.loads(tool_args_str) if tool_args_str else {} + except (json.JSONDecodeError, TypeError): + args_dict = {} + + logger.debug("[WebSocket] Sending tool_call (args complete): %s", tool_name) + + if turn_state.current_tool_group_id is None: + turn_state.current_tool_group_id = f"tg-{str(uuid.uuid4())[:8]}" + + await send_typed_tool_lifecycle( + ServerToolCall( + tool_id=tool_id, + tool_name=tool_name, + args=args_dict, + timestamp=time_module.time(), + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + tool_group_id=turn_state.current_tool_group_id, + ) + ) + + turn_state.pending_tool_calls[tool_id] = { + "tool_name": tool_name, + "start_time": start_time, + "part_index": part_index, + "raw_tool_call_id": raw_tool_call_id, + "status_only_sent": False, + "result": None, + "tool_group_id": turn_state.current_tool_group_id, + } + if raw_tool_call_id: + turn_state.tool_id_aliases[raw_tool_call_id] = tool_id + if turn_state.current_tool_group_id: + turn_state.tool_group_ids[tool_id] = turn_state.current_tool_group_id + + turn_state.current_tool_name = None + turn_state.active_parts.pop(part_index, None) + return True + + +async def send_status_only_pending_tool_results( + *, + turn_state: WebSocketTurnState, + session_id: str, + agent_name: str, + model_name: str, + send_typed: Callable[[Any], Awaitable[None]], + logger: Any, +) -> None: + """Emit status-only placeholder results before assistant text resumes.""" + for tool_id, tool_info in list(turn_state.pending_tool_calls.items()): + if tool_info.get("status_only_sent"): + continue + duration_ms = (time_module.time() - tool_info["start_time"]) * 1000 + logger.debug( + "[WebSocket] Sending status-only tool_result for: %s (duration: %.1fms)", + tool_info["tool_name"], + duration_ms, + ) + await send_typed( + ServerToolResult( + tool_id=tool_id, + result={"_status": "completed", "_pending_full_result": True}, + tool_name=tool_info["tool_name"], + success=True, + duration_ms=duration_ms, + timestamp=time_module.time(), + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + tool_group_id=resolve_tool_group_id( + turn_state=turn_state, + logger=logger, + tool_id=tool_id, + pending_info=tool_info, + fallback_group_id=turn_state.current_tool_group_id, + tool_name=tool_info.get("tool_name"), + source="status_only_result", + ), + ) + ) + tool_info["status_only_sent"] = True + + +def resolve_pending_tool_id( + *, + turn_state: WebSocketTurnState, + tool_call_id: str | None = None, + tool_name: str | None = None, +) -> str | None: + """Resolve canonical frontend tool_id for a backend tool call reference.""" + if tool_call_id and tool_call_id in turn_state.pending_tool_calls: + return tool_call_id + + if tool_call_id: + for pending_id, pending_info in turn_state.pending_tool_calls.items(): + if pending_info.get("raw_tool_call_id") == tool_call_id: + return pending_id + + if tool_name: + for pending_id, pending_info in turn_state.pending_tool_calls.items(): + if pending_info.get("tool_name") == tool_name: + return pending_id + + return None + + +def resolve_tool_group_id( + *, + turn_state: WebSocketTurnState, + logger: Any, + tool_id: str | None = None, + pending_info: dict[str, Any] | None = None, + fallback_group_id: str | None = None, + tool_name: str | None = None, + source: str = "unknown", +) -> str: + """Return a non-empty ``tool_group_id`` for tool lifecycle frames.""" + group_id = None + if pending_info is not None: + group_id = pending_info.get("tool_group_id") + + if not group_id and tool_id: + group_id = turn_state.tool_group_ids.get(tool_id) + + if not group_id and fallback_group_id: + group_id = fallback_group_id + + if not group_id: + group_id = turn_state.current_tool_group_id + + if not group_id: + raw_hint = tool_id or tool_name or "unknown" + stable_hint = ( + re.sub(r"[^a-z0-9_-]", "-", raw_hint.lower())[:24].strip("-") or "unknown" + ) + group_id = f"tg-fallback-{stable_hint}-{str(uuid.uuid4())[:6]}" + logger.warning( + "[ws] Synthesized missing tool_group_id for source=%s tool_id=%s tool_name=%s", + source, + tool_id, + tool_name, + ) + + if tool_id: + turn_state.tool_group_ids[tool_id] = group_id + if pending_info is not None: + pending_info["tool_group_id"] = group_id + elif tool_id in turn_state.pending_tool_calls: + turn_state.pending_tool_calls[tool_id]["tool_group_id"] = group_id + + if turn_state.current_tool_group_id is None: + turn_state.current_tool_group_id = group_id + + return group_id + + +def _extract_attr_or_key(obj: Any, name: str) -> Any: + if hasattr(obj, name): + return getattr(obj, name) + if isinstance(obj, dict): + return obj.get(name) + return None + + +def _coerce_tool_content(tool_content: Any, logger: Any) -> Any: + if not tool_content or isinstance( + tool_content, + (str, dict, list, int, float, bool, type(None)), + ): + return tool_content + + try: + if hasattr(tool_content, "model_dump"): + return tool_content.model_dump() + if hasattr(tool_content, "dict"): + return tool_content.dict() + if hasattr(tool_content, "__dict__"): + return tool_content.__dict__ + return str(tool_content) + except Exception as exc: # pragma: no cover - defensive parity branch + logger.debug("[WebSocket] Could not serialize tool result: %s", exc) + return str(tool_content) diff --git a/code_puppy/api/ws/chat_turn_runner.py b/code_puppy/api/ws/chat_turn_runner.py new file mode 100644 index 000000000..ce2266ab6 --- /dev/null +++ b/code_puppy/api/ws/chat_turn_runner.py @@ -0,0 +1,334 @@ +"""Per-turn WebSocket agent-runner helpers. + +This module extracts the most delicate part of ``chat_handler``: the window +where an agent is running while the WebSocket must continue accepting inbound +messages such as permission responses, cancel requests, and session switches. + +The goal is to preserve runtime behavior exactly while reducing the amount of +concurrent-control logic embedded directly in the endpoint handler. +""" + +from __future__ import annotations + +import asyncio +import logging +from dataclasses import dataclass +from typing import Any, Callable + +from fastapi import WebSocketDisconnect +from pydantic import TypeAdapter, ValidationError + +from code_puppy.api.ws.chat_context import ( + begin_agent_run_context, + cleanup_agent_run_context, +) +from code_puppy.api.ws.chat_turn_state import WebSocketTurnState +from code_puppy.api.ws.schemas import ClientMessage +from code_puppy.messaging.bus import get_message_bus +from code_puppy.messaging.commands import ( + AskUserQuestionResponse, + ConfirmationResponse, + SelectionResponse, + UserInputResponse, +) + +_ClientMessageAdapter = TypeAdapter(ClientMessage) +logger = logging.getLogger(__name__) + + +def handle_user_interaction_response(message: dict[str, Any]) -> bool: + """Resolve MessageBus user-interaction responses during an active turn.""" + msg_type = message.get("type") + if msg_type not in { + "user_input_response", + "confirmation_response", + "selection_response", + "ask_user_question_response", + }: + return False + + bus = get_message_bus() + if msg_type == "user_input_response": + bus.provide_response( + UserInputResponse( + prompt_id=message.get("prompt_id", ""), + value=message.get("value", ""), + ) + ) + elif msg_type == "confirmation_response": + bus.provide_response( + ConfirmationResponse( + prompt_id=message.get("prompt_id", ""), + confirmed=bool(message.get("confirmed", False)), + feedback=message.get("feedback"), + ) + ) + elif msg_type == "selection_response": + bus.provide_response( + SelectionResponse( + prompt_id=message.get("prompt_id", ""), + selected_index=int(message.get("selected_index", -1)), + selected_value=message.get("selected_value", ""), + ) + ) + else: + bus.provide_response( + AskUserQuestionResponse( + prompt_id=message.get("prompt_id", ""), + answers=message.get("answers") or [], + cancelled=bool(message.get("cancelled", False)), + ) + ) + return True + + +async def save_agent_result_in_background(**kwargs: Any) -> None: + """Lazy proxy to avoid importing DB-heavy background-save code at module load.""" + from code_puppy.api.ws.background_save import ( + save_agent_result_in_background as _save_agent_result_in_background, + ) + + await _save_agent_result_in_background(**kwargs) + + +def fire_and_track(coro: Any) -> asyncio.Task: + """Lazy proxy for tracked background tasks.""" + from code_puppy.api.ws.background_save import fire_and_track as _fire_and_track + + return _fire_and_track(coro) + + +@dataclass(slots=True) +class WebSocketTurnRunResult: + """Outcome of one ``run_with_mcp`` execution window.""" + + result: Any | None = None + deferred_msg: dict[str, Any] | None = None + + +async def execute_turn_runner( + *, + websocket: Any, + session_id: str, + ctx: Any, + agent: Any, + agent_name: str, + model_name: str, + session_title: str, + session_working_directory: str, + session_pinned: bool, + message_to_send: str, + run_kwargs: dict[str, Any], + turn_state: WebSocketTurnState, + clear_session_working_directory: Callable[[], None], +) -> WebSocketTurnRunResult: + """Run one agent turn while still receiving control messages. + + This preserves the existing concurrent behavior from ``chat_handler``: + - permission responses are handled immediately + - cancel requests cancel the in-flight agent task + - switch/create_session requests disown the task to background-save logic + - disconnects let the task finish in the background + """ + + _ws_run_context = begin_agent_run_context(session_id=session_id) + active_agent_task = asyncio.create_task( + agent.run_with_mcp(message_to_send, **run_kwargs) + ) + + result = None + agent_completed = False + deferred_msg: dict[str, Any] | None = None + + try: + while not agent_completed: + receive_task = asyncio.create_task(websocket.receive_json()) + done, pending = await asyncio.wait( + {active_agent_task, receive_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + + if active_agent_task in done: + try: + result = await active_agent_task + logger.debug( + "run_with_mcp completed, result type: %s", type(result) + ) + agent_completed = True + active_agent_task = None + except asyncio.CancelledError: + logger.debug("run_with_mcp task was cancelled by user") + turn_state.agent_error = "cancelled" + result = None + agent_completed = True + active_agent_task = None + except Exception as e: + logger.error("Agent task error: %s", e, exc_info=True) + logger.debug( + "[WS:%s] agent task exception captured: type=%s repr=%r", + session_id, + type(e).__name__, + e, + ) + turn_state.agent_error = e + result = None + agent_completed = True + active_agent_task = None + + if receive_task in pending: + receive_task.cancel() + try: + await receive_task + except asyncio.CancelledError: + pass + + elif receive_task in done: + try: + new_msg = await receive_task + + try: + _ClientMessageAdapter.validate_python(new_msg) + except ValidationError as _val_err: + logger.warning( + "Client message failed validation: %s", + str(_val_err), + extra={ + "type": new_msg.get("type") + if isinstance(new_msg, dict) + else "unknown" + }, + ) + + if handle_user_interaction_response(new_msg): + logger.debug( + "[UserInteraction] Handled response: %s", + new_msg.get("type"), + ) + + elif new_msg.get("type") == "permission_response": + from code_puppy.api.permissions import ( + handle_permission_response, + ) + + request_id = new_msg.get("request_id") + approved = new_msg.get("approved", False) + + if request_id: + handled = handle_permission_response( + request_id, approved, session_id=session_id + ) + if handled: + logger.debug( + "[Permission] ✅ Handled response: %s = %s", + request_id, + approved, + ) + else: + logger.warning( + "[Permission] ❌ Unknown request: %s", + request_id, + ) + else: + logger.error( + "[WebSocket] ❌ No request_id in permission_response!" + ) + + elif new_msg.get("type") == "cancel": + logger.debug("Cancel request received during agent execution") + if active_agent_task and not active_agent_task.done(): + active_agent_task.cancel() + agent_completed = True + + elif new_msg.get("type") in ("switch_session", "create_session"): + logger.debug( + "[WS:%s] Session switch during streaming — agent continues in background, switching to: %s", + session_id, + new_msg.get("session_id"), + ) + fire_and_track( + save_agent_result_in_background( + agent_task=active_agent_task, + session_id=session_id, + ctx=ctx, + agent=agent, + agent_name=agent_name, + model_name=model_name, + title=session_title, + working_directory=session_working_directory, + pinned=session_pinned, + label="switch", + ) + ) + deferred_msg = new_msg + active_agent_task = None + agent_completed = True + + else: + logger.warning( + "[WebSocket] Received %s message while agent running - ignoring", + new_msg.get("type"), + ) + + except asyncio.CancelledError: + pass + except WebSocketDisconnect: + logger.debug( + "[WS:%s] Disconnect during streaming — agent continues in background", + session_id, + ) + fire_and_track( + save_agent_result_in_background( + agent_task=active_agent_task, + session_id=session_id, + ctx=ctx, + agent=agent, + agent_name=agent_name, + model_name=model_name, + title=session_title, + working_directory=session_working_directory, + pinned=session_pinned, + label="disconnect", + ) + ) + active_agent_task = None + agent_completed = True + except RuntimeError as e: + if "disconnect" in str(e).lower(): + logger.debug( + "[WS:%s] WebSocket already disconnected: %s — agent continues in background", + session_id, + e, + ) + fire_and_track( + save_agent_result_in_background( + agent_task=active_agent_task, + session_id=session_id, + ctx=ctx, + agent=agent, + agent_name=agent_name, + model_name=model_name, + title=session_title, + working_directory=session_working_directory, + pinned=session_pinned, + label="runtime", + ) + ) + active_agent_task = None + agent_completed = True + else: + logger.error( + "RuntimeError processing message during agent execution: %s", + e, + ) + except Exception as e: + logger.error( + "Error processing message during agent execution: %s", + e, + ) + finally: + cleanup_agent_run_context( + _ws_run_context, + clear_session_working_directory=clear_session_working_directory, + ) + + return WebSocketTurnRunResult(result=result, deferred_msg=deferred_msg) diff --git a/code_puppy/api/ws/chat_turn_state.py b/code_puppy/api/ws/chat_turn_state.py new file mode 100644 index 000000000..fc9d8071c --- /dev/null +++ b/code_puppy/api/ws/chat_turn_state.py @@ -0,0 +1,23 @@ +"""Per-turn mutable state for the WebSocket chat handler. + +This keeps WebSocket-only streaming/tool bookkeeping out of the top-level +handler flow without changing runtime behavior. +""" + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass(slots=True) +class WebSocketTurnState: + """Mutable per-turn state for structured WebSocket streaming.""" + + collected_text: list[str] = field(default_factory=list) + active_parts: dict[int, dict[str, Any]] = field(default_factory=dict) + tool_id_aliases: dict[str, str] = field(default_factory=dict) + tool_group_ids: dict[str, str] = field(default_factory=dict) + pending_tool_calls: dict[str, dict[str, Any]] = field(default_factory=dict) + current_tool_name: str | None = None + current_tool_group_id: str | None = None + b1_streaming_used: bool = False + agent_error: object | None = None diff --git a/code_puppy/api/ws/connection_manager.py b/code_puppy/api/ws/connection_manager.py new file mode 100644 index 000000000..eec49804d --- /dev/null +++ b/code_puppy/api/ws/connection_manager.py @@ -0,0 +1,55 @@ +"""WebSocket connection manager for broadcasting session updates. + +This is a singleton that tracks connected WebSocket clients monitoring +session state changes. +""" + +import logging + +from fastapi import WebSocket + +logger = logging.getLogger(__name__) + + +class WebSocketConnectionManager: + """Manages WebSocket connections for broadcasting session updates.""" + + def __init__(self): + self.session_connections: list[WebSocket] = [] + + async def connect_session_client(self, websocket: WebSocket): + """Connect a session monitoring client.""" + self.session_connections.append(websocket) + logger.debug( + f"Session client connected. Total: {len(self.session_connections)}" + ) + + def disconnect_session_client(self, websocket: WebSocket): + """Disconnect a session monitoring client.""" + if websocket in self.session_connections: + self.session_connections.remove(websocket) + logger.debug( + f"Session client disconnected. Total: {len(self.session_connections)}" + ) + + async def broadcast_session_update(self, session_data: dict): + """Broadcast session update to all connected session clients.""" + if not self.session_connections: + return + + connections = self.session_connections.copy() + disconnected = [] + + for conn in connections: + try: + await conn.send_json({"type": "session_update", "data": session_data}) + except Exception as e: + logger.warning("Failed to send session update: %s", e) + disconnected.append(conn) + + for conn in disconnected: + self.disconnect_session_client(conn) + + +# Global singleton instance +connection_manager = WebSocketConnectionManager() diff --git a/code_puppy/api/ws/events_handler.py b/code_puppy/api/ws/events_handler.py new file mode 100644 index 000000000..b9bf647ec --- /dev/null +++ b/code_puppy/api/ws/events_handler.py @@ -0,0 +1,53 @@ +"""WebSocket endpoint for server-sent events streaming.""" + +import asyncio +import logging + +from fastapi import FastAPI, WebSocket, WebSocketDisconnect + +logger = logging.getLogger(__name__) + + +def register_events_endpoint(app: FastAPI) -> None: + """Register the /ws/events WebSocket endpoint.""" + + @app.websocket("/ws/events") + async def websocket_events( + websocket: WebSocket, session_id: str | None = None + ) -> None: + """Stream real-time events to connected clients. + + Query Parameters: + session_id: Optional session ID. If provided, only events for that + session are streamed. + """ + await websocket.accept() + logger.debug("Events WebSocket client connected (session_id=%s)", session_id) + + from code_puppy.plugins.frontend_emitter.emitter import ( + get_recent_events, + subscribe, + unsubscribe, + ) + + event_queue = subscribe(session_id=session_id) + + try: + for event in get_recent_events(session_id=session_id): + await websocket.send_json(event) + + while True: + try: + event = await asyncio.wait_for(event_queue.get(), timeout=30.0) + await websocket.send_json(event) + except asyncio.TimeoutError: + try: + await websocket.send_json({"type": "ping"}) + except Exception: + break + except WebSocketDisconnect: + logger.debug("Events WebSocket client disconnected") + except Exception as e: + logger.error("Events WebSocket error: %s", e) + finally: + unsubscribe(event_queue) diff --git a/code_puppy/api/ws/health_handler.py b/code_puppy/api/ws/health_handler.py new file mode 100644 index 000000000..832f45e95 --- /dev/null +++ b/code_puppy/api/ws/health_handler.py @@ -0,0 +1,22 @@ +"""WebSocket endpoint for health checks.""" + +import logging + +from fastapi import FastAPI, WebSocket, WebSocketDisconnect + +logger = logging.getLogger(__name__) + + +def register_health_endpoint(app: FastAPI) -> None: + """Register the /ws/health WebSocket endpoint.""" + + @app.websocket("/ws/health") + async def websocket_health(websocket: WebSocket) -> None: + """Simple WebSocket health check - echoes messages back.""" + await websocket.accept() + try: + while True: + data = await websocket.receive_text() + await websocket.send_text(f"echo: {data}") + except WebSocketDisconnect: + pass diff --git a/code_puppy/api/ws/history_utils.py b/code_puppy/api/ws/history_utils.py new file mode 100644 index 000000000..e72a19942 --- /dev/null +++ b/code_puppy/api/ws/history_utils.py @@ -0,0 +1,104 @@ +"""History-wrapping helpers for WebSocket session persistence. + +These functions keep the message-history decoration logic out of the main chat +handler so it can be tested independently and reused by future persistence +refactors. +""" + +from __future__ import annotations + +import datetime +from typing import Any + + +def extract_message_timestamp(raw_msg: Any, default_ts: str) -> str: + """Best-effort extraction of an existing timestamp for a message.""" + if isinstance(raw_msg, dict): + ts_val = raw_msg.get("timestamp") + + if isinstance(ts_val, (int, float)): + try: + return datetime.datetime.fromtimestamp(ts_val).isoformat() + except Exception: + pass + + if isinstance(ts_val, str) and ts_val: + return ts_val + + ts_field = raw_msg.get("ts") + if isinstance(ts_field, str) and ts_field: + return ts_field + + try: + attr_ts = getattr(raw_msg, "timestamp", None) + if isinstance(attr_ts, (int, float)): + return datetime.datetime.fromtimestamp(attr_ts).isoformat() + if isinstance(attr_ts, str) and attr_ts: + return attr_ts + except Exception: + pass + + return default_ts + + +def build_enhanced_history( + history: list[Any], + *, + agent_name_meta: str, + model_name_meta: str, + original_user_message: str, + attachment_metadata: list[Any] | None = None, +) -> list[dict[str, Any]]: + """Wrap raw history entries with agent/model/timestamp metadata. + + Existing wrapped entries are preserved as-is for idempotency. When + attachments were injected into the just-processed user message, we backfill + clean UI metadata onto the penultimate history entry. + """ + attachment_metadata = attachment_metadata or [] + enhanced_history: list[dict[str, Any]] = [] + + for idx, msg in enumerate(history): + if isinstance(msg, dict) and "msg" in msg and "agent" in msg: + enhanced_history.append(msg) + continue + + current_timestamp = datetime.datetime.now().isoformat() + msg_ts = extract_message_timestamp(msg, current_timestamp) + wrapper: dict[str, Any] = { + "msg": msg, + "agent": agent_name_meta, + "model": model_name_meta, + "ts": msg_ts, + } + + is_user_message_just_processed = ( + idx == len(history) - 2 and len(history) >= 2 and bool(attachment_metadata) + ) + + if is_user_message_just_processed: + wrapper["clean_content"] = original_user_message + wrapper["attachments"] = attachment_metadata + + enhanced_history.append(wrapper) + + return enhanced_history + + +def estimate_total_tokens(enhanced_history: list[Any], agent: Any) -> int: + """Estimate total tokens for wrapped history, returning 0 on failure.""" + try: + total_tokens = 0 + for item in enhanced_history: + msg_obj = item["msg"] if isinstance(item, dict) and "msg" in item else item + total_tokens += agent.estimate_tokens_for_message(msg_obj) + return total_tokens + except Exception: + return 0 + + +__all__ = [ + "build_enhanced_history", + "estimate_total_tokens", + "extract_message_timestamp", +] diff --git a/code_puppy/api/ws/response_frames.py b/code_puppy/api/ws/response_frames.py new file mode 100644 index 000000000..12bc6c920 --- /dev/null +++ b/code_puppy/api/ws/response_frames.py @@ -0,0 +1,151 @@ +"""Helpers for WebSocket response/error frame construction. + +These utilities are intentionally kept separate from the giant chat handler so +the wire-protocol transformations can be tested and maintained independently. +""" + +from __future__ import annotations + +import uuid +from typing import Any + +from code_puppy.api.error_parser import parse_api_error as _legacy_parse_api_error +from code_puppy.api.ws.schemas import ( + ServerAssistantMessageDelta, + ServerAssistantMessageEnd, + ServerAssistantMessageStart, + ServerError, + ServerStreamEnd, +) +from code_puppy.model_errors import normalize_model_error as _normalize_model_error + +_UNKNOWN_ERROR_TYPE = "unknown_error" +_TOOL_HISTORY_ERROR_TYPE = "tool_history_error" + +_CODE_TO_ERROR_TYPE: dict[str, str] = { + "rate_limit_or_overloaded": "rate_limit", + "backend_unavailable": "server_error", + "auth_error": "auth_error", + "quota_exceeded": "quota_exceeded", + "content_blocked": "content_blocked", + "invalid_tool_history": _TOOL_HISTORY_ERROR_TYPE, +} + + +def parse_api_error(exc: Exception) -> dict[str, Any]: + """Convert an agent-run exception into a structured frontend error dict.""" + norm = _normalize_model_error(exc) + if norm.code in _CODE_TO_ERROR_TYPE: + error_type = _CODE_TO_ERROR_TYPE[norm.code] + user_message = norm.user_message or str(exc) + action_required = error_type in ("rate_limit", "quota_exceeded") + return { + "user_message": user_message, + "error_type": error_type, + "technical_details": repr(exc), + "action_required": action_required, + } + return _legacy_parse_api_error(exc) + + +def has_streamed_content(collected_text: list[str | None]) -> bool: + """Return True when the client has already received streaming output chunks.""" + return any((chunk or "").strip() for chunk in collected_text) + + +def build_error_response_frames( + agent_error: Exception, + collected_text: list[str | None], + session_id: str, +) -> list[dict[str, Any]]: + """Build ordered WebSocket frames for a failed agent run.""" + frames: list[dict[str, Any]] = [] + if has_streamed_content(collected_text): + frames.append( + ServerStreamEnd( + success=False, + session_id=session_id, + ).model_dump(exclude_none=True) + ) + parsed = parse_api_error(agent_error) + frames.append( + ServerError( + error=parsed["user_message"], + error_type=parsed["error_type"], + technical_details=parsed["technical_details"], + action_required=parsed.get("action_required"), + session_id=session_id, + ).model_dump(exclude_none=True) + ) + return frames + + +def build_assistant_text_stream_frames( + *, + response_text: str, + session_id: str, + agent_name: str | None = None, + model_name: str | None = None, + tokens: dict[str, Any] | None = None, + part_type: str = "text", + part_index: int = 0, + message_id: str | None = None, + timestamp: float | None = None, +) -> list[ + ServerAssistantMessageStart + | ServerAssistantMessageDelta + | ServerAssistantMessageEnd + | ServerStreamEnd +]: + """Represent a complete assistant response using streaming-shaped frames.""" + import time as time_module + + now = timestamp if timestamp is not None else time_module.time() + msg_id = message_id or f"msg-{session_id}-{uuid.uuid4().hex}" + content = response_text or "" + + return [ + ServerAssistantMessageStart( + message_id=msg_id, + part_type=part_type, + part_index=part_index, + timestamp=now, + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + ), + ServerAssistantMessageDelta( + message_id=msg_id, + content=content, + part_index=part_index, + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + ), + ServerAssistantMessageEnd( + message_id=msg_id, + part_type=part_type, + part_index=part_index, + full_content=content, + timestamp=now, + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + ), + ServerStreamEnd( + success=True, + total_length=len(content), + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + tokens=tokens, + ), + ] + + +__all__ = [ + "build_assistant_text_stream_frames", + "build_error_response_frames", + "has_streamed_content", + "parse_api_error", +] diff --git a/code_puppy/api/ws/runtime/__init__.py b/code_puppy/api/ws/runtime/__init__.py new file mode 100644 index 000000000..750f74d72 --- /dev/null +++ b/code_puppy/api/ws/runtime/__init__.py @@ -0,0 +1,14 @@ +"""Runtime scaffolding for multi-session chat. + +These modules are intentionally lightweight and are not yet wired into +`/ws/chat` or the planned `/ws/chat_mux` endpoint. + +They provide an async-friendly session runtime object (cancellation, active +streaming task handle, etc.) and a manager that keeps runtimes keyed by +session_id. +""" + +from .session_runtime import SessionRuntime +from .session_runtime_manager import SessionRuntimeManager + +__all__ = ["SessionRuntime", "SessionRuntimeManager"] diff --git a/code_puppy/api/ws/runtime/session_runtime.py b/code_puppy/api/ws/runtime/session_runtime.py new file mode 100644 index 000000000..a6fda2741 --- /dev/null +++ b/code_puppy/api/ws/runtime/session_runtime.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any + +from code_puppy.api.session_context import SessionContext + + +@dataclass(slots=True) +class SessionRuntime: + """Async runtime state for a single session. + + This complements :class:`~code_puppy.api.session_context.SessionContext`. + + SessionContext stores durable-ish session state (agent, history, metadata). + SessionRuntime stores ephemeral runtime state (active stream task, + cancellation token/event, etc.). + + NOTE: This is scaffolding for chat mux and is not yet integrated into + websocket handlers. + """ + + session_id: str + ctx: SessionContext + + created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + last_used_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + + # When set, any active streaming loop should stop ASAP. + cancel_event: asyncio.Event = field(default_factory=asyncio.Event) + + # Track at most one active streaming task at a time for now. + _active_task: asyncio.Task[Any] | None = None + + # Serialise runtime-level mutations. + _lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) + + async def set_active_task(self, task: asyncio.Task[Any] | None) -> None: + """Set (or clear) the active streaming task for this session.""" + async with self._lock: + self._active_task = task + self.last_used_at = datetime.now(timezone.utc) + + async def get_active_task(self) -> asyncio.Task[Any] | None: + async with self._lock: + return self._active_task + + async def cancel(self) -> None: + """Cancel the active streaming task (if any) and set cancel_event.""" + self.cancel_event.set() + async with self._lock: + task = self._active_task + + if task and not task.done(): + task.cancel() + + async def reset_cancel(self) -> None: + """Clear cancel state for the next turn.""" + # asyncio.Event has no clear() in older Python; but 3.11+ has. + # This project supports 3.11+, so use clear(). + self.cancel_event.clear() + + async def touch(self) -> None: + """Update last_used_at.""" + async with self._lock: + self.last_used_at = datetime.now(timezone.utc) diff --git a/code_puppy/api/ws/runtime/session_runtime_manager.py b/code_puppy/api/ws/runtime/session_runtime_manager.py new file mode 100644 index 000000000..82a8252d3 --- /dev/null +++ b/code_puppy/api/ws/runtime/session_runtime_manager.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Callable +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone + +from code_puppy.api.session_context import SessionContext, session_manager + +from .session_runtime import SessionRuntime + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True, slots=True) +class RuntimeCleanupResult: + removed_session_ids: tuple[str, ...] + + +class SessionRuntimeManager: + """Registry for :class:`~code_puppy.api.ws.runtime.SessionRuntime`. + + This is designed for the multiplexed chat websocket, where the + server keeps multiple sessions "hot" at once. + """ + + def __init__(self, now: Callable[[], datetime] | None = None) -> None: + self._now = now or (lambda: datetime.now(timezone.utc)) + self._runtimes: dict[str, SessionRuntime] = {} + self._lock = asyncio.Lock() + + async def get_runtime(self, session_id: str) -> SessionRuntime | None: + async with self._lock: + return self._runtimes.get(session_id) + + async def ensure_runtime( + self, + session_id: str, + *, + agent_name: str = "code-puppy", + model_name: str | None = None, + working_directory: str = "", + load_if_exists: bool = True, + ) -> SessionRuntime: + """Get or create a runtime for *session_id*. + + This bridges to the existing :mod:`code_puppy.api.session_context` layer: + - If an in-memory SessionContext exists, we reuse it. + - Else, we optionally load from disk via SessionManager.load_session. + - Else, we create a new SessionContext. + + Args: + load_if_exists: If True, attempt to load persisted session state. + + Returns: + SessionRuntime instance. + """ + async with self._lock: + existing = self._runtimes.get(session_id) + if existing is not None: + existing.last_used_at = self._now() + return existing + + ctx = await session_manager.get_session(session_id) + if ctx is None and load_if_exists: + try: + ctx = await session_manager.load_session(session_id) + except Exception: + logger.debug( + "Failed to load session %s from disk", session_id, exc_info=True + ) + ctx = None + + if ctx is None: + ctx = await session_manager.create_session( + session_id, + agent_name=agent_name, + model_name=model_name, + working_directory=working_directory, + ) + + rt = SessionRuntime(session_id=session_id, ctx=ctx) + rt.last_used_at = self._now() + self._runtimes[session_id] = rt + return rt + + async def cancel_session(self, session_id: str) -> bool: + """Cancel the active task for *session_id*. + + Returns: + True if a runtime existed and cancellation was requested. + """ + rt = await self.get_runtime(session_id) + if rt is None: + return False + await rt.cancel() + return True + + async def remove_runtime(self, session_id: str) -> bool: + """Remove runtime from the manager. + + This does not destroy the underlying SessionContext (that remains in the + existing SessionManager registry). + """ + async with self._lock: + removed = self._runtimes.pop(session_id, None) + return removed is not None + + async def cleanup_idle(self, *, max_idle: timedelta) -> RuntimeCleanupResult: + """Remove runtimes that have been idle longer than *max_idle*.""" + now = self._now() + removed: list[str] = [] + + async with self._lock: + for session_id, rt in list(self._runtimes.items()): + if now - rt.last_used_at > max_idle: + removed.append(session_id) + self._runtimes.pop(session_id, None) + + if removed: + logger.info("Cleaned up %d idle runtimes", len(removed)) + + return RuntimeCleanupResult(removed_session_ids=tuple(removed)) + + async def get_session_context(self, session_id: str) -> SessionContext | None: + rt = await self.get_runtime(session_id) + return rt.ctx if rt else None + + async def active_task_count(self) -> int: + """Count runtimes with active tasks (for debugging/tests).""" + async with self._lock: + rts = list(self._runtimes.values()) + + count = 0 + for rt in rts: + task = await rt.get_active_task() + if task is not None and not task.done(): + count += 1 + return count + + +_manager_instance: SessionRuntimeManager | None = None + + +def get_runtime_manager() -> SessionRuntimeManager: + global _manager_instance + if _manager_instance is None: + _manager_instance = SessionRuntimeManager() + return _manager_instance diff --git a/code_puppy/api/ws/schemas.py b/code_puppy/api/ws/schemas.py new file mode 100644 index 000000000..5f42fa9fc --- /dev/null +++ b/code_puppy/api/ws/schemas.py @@ -0,0 +1,702 @@ +"""WebSocket protocol schemas — single source of truth for the CP wire protocol. + +Every client→server and server→client message is defined here as a Pydantic +model with ``Literal`` type discriminators so the framework can parse raw JSON +into the correct model automatically via ``ClientMessage`` / ``ServerMessage`` +discriminated unions. + +Serialisation convention: + ``.model_dump(exclude_none=True)`` — optional fields that are ``None`` + are omitted from the wire representation to keep payloads lean. + +All field names and types are derived from the *actual* ``send_json`` / +``safe_send_json`` calls in ``chat_handler.py`` and ``permissions.py``. +""" + +from __future__ import annotations + +from enum import Enum +from typing import Annotated, Any, Dict, List, Literal, Optional, Union + +from pydantic import BaseModel, Discriminator, Field, Tag + +# ── Protocol version ──────────────────────────────────────────────────────── +PROTOCOL_VERSION = "1.0.0" + + +# ── Enums ──────────────────────────────────────────────────────────────────── + + +class ClientMessageType(str, Enum): + """All ``type`` values a client may send over the WebSocket.""" + + MESSAGE = "message" + SWITCH_AGENT = "switch_agent" + SWITCH_MODEL = "switch_model" + SWITCH_SESSION = "switch_session" + SET_WORKING_DIRECTORY = "set_working_directory" + UPDATE_SESSION_META = "update_session_meta" + GET_CONFIG = "get_config" + SET_CONFIG = "set_config" + COMMAND = "command" + CANCEL = "cancel" + PERMISSION_RESPONSE = "permission_response" + USER_INPUT_RESPONSE = "user_input_response" + CONFIRMATION_RESPONSE = "confirmation_response" + SELECTION_RESPONSE = "selection_response" + ASK_USER_QUESTION_RESPONSE = "ask_user_question_response" + + +class ServerMessageType(str, Enum): + """All ``type`` values the server may send over the WebSocket.""" + + SYSTEM = "system" + SESSION_META = "session_meta" + SESSION_RESTORED = "session_restored" + SESSION_SWITCHED = "session_switched" + WORKING_DIRECTORY_CHANGED = "working_directory_changed" + SESSION_META_UPDATED = "session_meta_updated" + CONFIG_VALUE = "config_value" + COMMAND_RESULT = "command_result" + STATUS = "status" + USER_MESSAGE = "user_message" + ASSISTANT_MESSAGE_START = "assistant_message_start" + ASSISTANT_MESSAGE_DELTA = "assistant_message_delta" + ASSISTANT_MESSAGE_END = "assistant_message_end" + TOOL_CALL = "tool_call" + TOOL_RESULT = "tool_result" + TOOL_RETURN = "tool_return" + AGENT_INVOKED = "agent_invoked" + RESPONSE = "response" + STREAM_END = "stream_end" + ERROR = "error" + PERMISSION_REQUEST = "permission_request" + USER_INPUT_REQUEST = "user_input_request" + CONFIRMATION_REQUEST = "confirmation_request" + SELECTION_REQUEST = "selection_request" + ASK_USER_QUESTION_REQUEST = "ask_user_question_request" + CANCELLED = "cancelled" + + +class PartType(str, Enum): + """Streaming part types used in ``assistant_message_start``.""" + + TEXT = "text" + THINKING = "thinking" + TOOL_CALL = "tool_call" + + +class ErrorType(str, Enum): + """Error categories returned by ``parse_api_error`` / ``error_parser``.""" + + RATE_LIMIT = "rate_limit" + SERVER_ERROR = "server_error" + AUTH_ERROR = "auth_error" + QUOTA_EXCEEDED = "quota_exceeded" + CONTENT_BLOCKED = "content_blocked" + TOOL_HISTORY_ERROR = "tool_history_error" + NETWORK_ERROR = "network_error" + MODEL_NOT_FOUND = "model_not_found" + CLAUDE_TEMPERATURE_ERROR = "claude_temperature_error" + UNKNOWN = "unknown" + UNKNOWN_ERROR = "unknown_error" + + +# ── Base ───────────────────────────────────────────────────────────────────── + + +class _BaseMessage(BaseModel): + """Shared config for all protocol messages.""" + + model_config = {"extra": "allow"} + + +# ═══════════════════════════════════════════════════════════════════════════════ +# CLIENT → SERVER (11 message types) +# ═══════════════════════════════════════════════════════════════════════════════ + + +class ClientMessageMessage(_BaseMessage): + """User chat message. + + Sent when the user types a message in the chat input. + ``content`` is the message text. Optional ``model`` overrides the + session model for this single turn. ``model_settings`` passes + per-request knobs (e.g. reasoning_effort). ``attachments`` is a + list of file paths to include. + """ + + type: Literal["message"] = "message" + content: str + model: Optional[str] = None + model_settings: Optional[Dict[str, Any]] = None + attachments: Optional[List[str]] = None + + +class ClientSwitchAgent(_BaseMessage): + """Request to switch the active agent. + + ``agent_name`` is the target agent identifier (e.g. ``"code-puppy"``). + """ + + type: Literal["switch_agent"] = "switch_agent" + agent_name: str + + +class ClientSwitchModel(_BaseMessage): + """Request to switch the active model. + + The server accepts either ``model_name`` or ``model`` for backwards + compatibility — see ``msg.get("model_name") or msg.get("model")``. + """ + + type: Literal["switch_model"] = "switch_model" + model_name: Optional[str] = None + model: Optional[str] = None + + +class ClientSwitchSession(_BaseMessage): + """Request to switch to a different chat session. + + ``session_id`` is the target session identifier. + """ + + type: Literal["switch_session"] = "switch_session" + session_id: str + + +class ClientSetWorkingDirectory(_BaseMessage): + """Set the working directory for the current session. + + ``directory`` is an absolute filesystem path. + """ + + type: Literal["set_working_directory"] = "set_working_directory" + directory: str + + +class ClientUpdateSessionMeta(_BaseMessage): + """Update session metadata (title, pinned state). + + Both fields are optional — only supplied fields are updated. + """ + + type: Literal["update_session_meta"] = "update_session_meta" + title: Optional[str] = None + pinned: Optional[bool] = None + + +class ClientGetConfig(_BaseMessage): + """Read a configuration value by key.""" + + type: Literal["get_config"] = "get_config" + key: str + + +class ClientSetConfig(_BaseMessage): + """Write a configuration value.""" + + type: Literal["set_config"] = "set_config" + key: str + value: Any + + +class ClientCommand(_BaseMessage): + """Execute a slash command (e.g. ``/help``).""" + + type: Literal["command"] = "command" + command: str + + +class ClientCancel(_BaseMessage): + """Cancel / interrupt the current streaming response or agent run.""" + + type: Literal["cancel"] = "cancel" + + +class ClientPermissionResponse(_BaseMessage): + """User's response to a permission request. + + ``approved`` indicates whether the user granted the permission. + ``remember`` is reserved for future "always allow" support. + """ + + type: Literal["permission_response"] = "permission_response" + request_id: str + approved: bool + remember: Optional[bool] = None + + +class ClientUserInputResponse(_BaseMessage): + """Response to a free-form input prompt emitted by the runtime.""" + + type: Literal["user_input_response"] = "user_input_response" + prompt_id: str + value: str + + +class ClientConfirmationResponse(_BaseMessage): + """Response to a confirmation prompt emitted by the runtime.""" + + type: Literal["confirmation_response"] = "confirmation_response" + prompt_id: str + confirmed: bool + feedback: Optional[str] = None + + +class ClientSelectionResponse(_BaseMessage): + """Response to an option-selection prompt emitted by the runtime.""" + + type: Literal["selection_response"] = "selection_response" + prompt_id: str + selected_index: int + selected_value: str + + +class ClientAskUserQuestionResponse(_BaseMessage): + """Response to a structured ask_user_question prompt emitted by the runtime.""" + + type: Literal["ask_user_question_response"] = "ask_user_question_response" + prompt_id: str + answers: List[Dict[str, Any]] = Field(default_factory=list) + cancelled: bool = False + + +# ═══════════════════════════════════════════════════════════════════════════════ +# SERVER → CLIENT (user-visible message types) +# ═══════════════════════════════════════════════════════════════════════════════ + + +class ServerSystem(_BaseMessage): + """System / informational message. + + Sent on initial connection (welcome), agent/model switches, and + replayed directory banners on session restore. + """ + + type: Literal["system"] = "system" + content: str + session_id: str + agent_name: Optional[str] = None + model_name: Optional[str] = None + resumed: Optional[bool] = None + protocol_version: Optional[str] = None + + +class ServerSessionRestored(_BaseMessage): + """Confirmation that an existing session was restored from storage. + + ``message_count`` reflects how many messages were loaded into the + agent's history. ``ui_metadata`` carries replayed UI-only records. + """ + + type: Literal["session_restored"] = "session_restored" + session_id: str + message_count: int + title: Optional[str] = None + ui_metadata: Optional[List[Any]] = None + + +class ServerSessionSwitched(_BaseMessage): + """Confirmation that the active session was switched. + + ``created`` is ``True`` when the target session did not exist and + was freshly created. + """ + + type: Literal["session_switched"] = "session_switched" + session_id: str + message_count: int + title: Optional[str] = None + working_directory: Optional[str] = None + created: bool + agent_name: Optional[str] = None + model_name: Optional[str] = None + + +class ServerWorkingDirectoryChanged(_BaseMessage): + """Result of a ``set_working_directory`` request. + + ``success`` is ``False`` when the directory does not exist. + ``unchanged`` is ``True`` when the directory matches the current one. + """ + + type: Literal["working_directory_changed"] = "working_directory_changed" + directory: str + success: bool + session_id: str + unchanged: Optional[bool] = None + error: Optional[str] = None + + +class ServerSessionMetaUpdated(_BaseMessage): + """Acknowledgement that session metadata was updated.""" + + type: Literal["session_meta_updated"] = "session_meta_updated" + session_id: str + pinned: Optional[bool] = None + title: Optional[str] = None + + +class ServerConfigValue(_BaseMessage): + """Response to ``get_config`` or ``set_config``. + + ``success`` is present only on ``set_config`` responses. + """ + + type: Literal["config_value"] = "config_value" + key: str + value: Any + session_id: str + success: Optional[bool] = None + + +class ServerCommandResult(_BaseMessage): + """Result of a slash command execution. + + ``output`` contains rendered text (e.g. help output). + ``error`` is present when the command raised an exception. + """ + + type: Literal["command_result"] = "command_result" + command: str + success: bool + session_id: str + output: Optional[str] = None + messages: Optional[List[Any]] = None + result: Optional[str] = None + error: Optional[str] = None + + +class ServerStatus(_BaseMessage): + """Transient status update (e.g. ``"cancelled"``, ``"thinking"``). + + Sent in response to a ``cancel`` request or to indicate agent + activity (e.g. "thinking" with current agent/model context). + """ + + type: Literal["status"] = "status" + status: str + session_id: str + agent_name: Optional[str] = None + model_name: Optional[str] = None + + +class ServerUserMessage(_BaseMessage): + """Echo of the user's chat message back to the client. + + Sent immediately after receiving a ``message`` from the client so + the UI can render it before the assistant starts responding. + """ + + type: Literal["user_message"] = "user_message" + content: str + session_id: str + + +class ServerAssistantMessageStart(_BaseMessage): + """Marks the beginning of a streaming assistant message part. + + ``part_type`` is ``"text"`` or ``"thinking"``. + ``tool_name`` is set when the text part follows a tool call. + """ + + type: Literal["assistant_message_start"] = "assistant_message_start" + message_id: str + part_type: str + part_index: int + timestamp: float + session_id: str + agent_name: Optional[str] = None + model_name: Optional[str] = None + tool_name: Optional[str] = None + + +class ServerAssistantMessageDelta(_BaseMessage): + """A streaming content chunk for an in-progress assistant message.""" + + type: Literal["assistant_message_delta"] = "assistant_message_delta" + message_id: str + content: str + part_index: int + session_id: str + agent_name: Optional[str] = None + model_name: Optional[str] = None + tool_name: Optional[str] = None + + +class ServerAssistantMessageEnd(_BaseMessage): + """Marks the end of a streaming assistant message part. + + ``full_content`` contains the complete accumulated text for the part. + ``part_type`` mirrors ``assistant_message_start`` for GUI reducers that + route text/thinking/tool-call parts consistently. + """ + + type: Literal["assistant_message_end"] = "assistant_message_end" + message_id: str + part_type: Optional[str] = None + part_index: int + full_content: str + timestamp: float + session_id: str + agent_name: Optional[str] = None + model_name: Optional[str] = None + tool_name: Optional[str] = None + + +class ServerToolCall(_BaseMessage): + """Notification that a tool was invoked. + + ``args`` contains the parsed tool arguments. ``tool_id`` correlates + with a later ``tool_result`` bearing the same ID. ``tool_group_id`` + groups parallel tool calls in the same execution batch. + """ + + type: Literal["tool_call"] = "tool_call" + tool_id: str + tool_name: str + args: Any # dict after JSON parse, but may be raw str during streaming + timestamp: float + session_id: str + agent_name: Optional[str] = None + model_name: Optional[str] = None + tool_group_id: Optional[str] = None # Groups tools in same execution batch + + +class ServerToolResult(_BaseMessage): + """Result of a tool execution. + + ``tool_id`` matches the originating ``tool_call``. ``result`` is the + tool's return value (may be ``dict``, ``str``, or a status sentinel). + ``tool_group_id`` matches the originating ``tool_call`` group. + """ + + type: Literal["tool_result"] = "tool_result" + tool_name: str + result: Any + success: bool + duration_ms: float + timestamp: float + session_id: str + tool_id: Optional[str] = None + agent_name: Optional[str] = None + model_name: Optional[str] = None + tool_group_id: Optional[str] = None # Matches originating tool_call group + + +class ServerToolReturn(_BaseMessage): + """Internal tool-return part tracked during streaming. + + This message type represents a ``ToolReturnPart`` observed in the + streaming pipeline. It is stored internally for correlation but may + not always be sent directly to the client. + """ + + type: Literal["tool_return"] = "tool_return" + tool_call_id: Optional[str] = None + content: Optional[Any] = None + session_id: Optional[str] = None + + +class ServerAgentInvoked(_BaseMessage): + """Notification that a sub-agent was invoked. + + ``prompt_preview`` is a truncated preview of the prompt sent to the + sub-agent. + """ + + type: Literal["agent_invoked"] = "agent_invoked" + agent_name: str + prompt_preview: Optional[str] = None + timestamp: float + session_id: str + + +class TokenUsage(BaseModel): + """Token usage statistics attached to ``response`` and ``stream_end``.""" + + input_tokens: Optional[int] = None + output_tokens: Optional[int] = None + total_tokens: Optional[int] = None + estimated: Optional[bool] = None + + +class ServerResponse(_BaseMessage): + """Complete (non-streaming) response. + + Sent when B1 streaming was *not* used (e.g. non-streaming models). + ``tokens`` carries optional token usage info. + """ + + type: Literal["response"] = "response" + content: str + done: bool + session_id: str + agent_name: Optional[str] = None + model_name: Optional[str] = None + tokens: Optional[Union[TokenUsage, Dict[str, Any]]] = None + + +class ServerStreamEnd(_BaseMessage): + """End-of-stream marker for B1 streaming responses. + + ``success`` is ``False`` when the stream ended due to an error. + ``total_length`` is the character count of the accumulated response. + """ + + type: Literal["stream_end"] = "stream_end" + success: bool + session_id: str + total_length: Optional[int] = None + agent_name: Optional[str] = None + model_name: Optional[str] = None + tokens: Optional[Union[TokenUsage, Dict[str, Any]]] = None + + +class ServerError(_BaseMessage): + """Error message. + + ``error_type`` classifies the error for frontend handling (e.g. + ``"rate_limit"``, ``"auth_error"``). ``action_required`` hints + that user intervention may be needed. + """ + + type: Literal["error"] = "error" + error: str + session_id: str + error_type: Optional[str] = None + technical_details: Optional[str] = None + action_required: Optional[bool] = None + + +class ServerPermissionRequest(_BaseMessage): + """Asks the user to approve or deny a potentially dangerous operation. + + Sent by the permission system before executing shell commands or + other privileged operations. The frontend should present an + approve/deny dialog and reply with ``permission_response``. + """ + + type: Literal["permission_request"] = "permission_request" + request_id: str + permission_type: str + title: str + description: str + details: Dict[str, Any] + session_id: str + timeout_seconds: int + + +class ServerUserInputRequest(_BaseMessage): + """Ask the browser UI to collect free-form text from the user.""" + + type: Literal["user_input_request"] = "user_input_request" + prompt_id: str + prompt_text: str + session_id: str + default_value: Optional[str] = None + input_type: Literal["text", "password"] = "text" + + +class ServerConfirmationRequest(_BaseMessage): + """Ask the browser UI to collect a yes/no-style confirmation.""" + + type: Literal["confirmation_request"] = "confirmation_request" + prompt_id: str + title: str + description: str + session_id: str + options: List[str] = [] + allow_feedback: bool = False + + +class ServerSelectionRequest(_BaseMessage): + """Ask the browser UI to collect one selected option from a list.""" + + type: Literal["selection_request"] = "selection_request" + prompt_id: str + prompt_text: str + options: List[str] + session_id: str + allow_cancel: bool = True + + +class ServerAskUserQuestionRequest(_BaseMessage): + """Ask the browser UI to collect structured ask_user_question answers.""" + + type: Literal["ask_user_question_request"] = "ask_user_question_request" + prompt_id: str + questions: List[Dict[str, Any]] + session_id: str + timeout_seconds: int = 300 + + +class ServerCancelled(_BaseMessage): + """Confirmation that the current agent run was cancelled. + + Sent when the agent task is interrupted after a ``cancel`` request. + """ + + type: Literal["cancelled"] = "cancelled" + session_id: str + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Discriminated unions +# ═══════════════════════════════════════════════════════════════════════════════ + + +ClientMessage = Annotated[ + Union[ + Annotated[ClientMessageMessage, Tag("message")], + Annotated[ClientSwitchAgent, Tag("switch_agent")], + Annotated[ClientSwitchModel, Tag("switch_model")], + Annotated[ClientSwitchSession, Tag("switch_session")], + Annotated[ClientSetWorkingDirectory, Tag("set_working_directory")], + Annotated[ClientUpdateSessionMeta, Tag("update_session_meta")], + Annotated[ClientGetConfig, Tag("get_config")], + Annotated[ClientSetConfig, Tag("set_config")], + Annotated[ClientCommand, Tag("command")], + Annotated[ClientCancel, Tag("cancel")], + Annotated[ClientPermissionResponse, Tag("permission_response")], + Annotated[ClientUserInputResponse, Tag("user_input_response")], + Annotated[ClientConfirmationResponse, Tag("confirmation_response")], + Annotated[ClientSelectionResponse, Tag("selection_response")], + Annotated[ClientAskUserQuestionResponse, Tag("ask_user_question_response")], + ], + Discriminator("type"), +] +"""Discriminated union of all client→server message types.""" + +ServerMessage = Annotated[ + Union[ + Annotated[ServerSystem, Tag("system")], + Annotated[ServerSessionRestored, Tag("session_restored")], + Annotated[ServerSessionSwitched, Tag("session_switched")], + Annotated[ServerWorkingDirectoryChanged, Tag("working_directory_changed")], + Annotated[ServerSessionMetaUpdated, Tag("session_meta_updated")], + Annotated[ServerConfigValue, Tag("config_value")], + Annotated[ServerCommandResult, Tag("command_result")], + Annotated[ServerStatus, Tag("status")], + Annotated[ServerUserMessage, Tag("user_message")], + Annotated[ServerAssistantMessageStart, Tag("assistant_message_start")], + Annotated[ServerAssistantMessageDelta, Tag("assistant_message_delta")], + Annotated[ServerAssistantMessageEnd, Tag("assistant_message_end")], + Annotated[ServerToolCall, Tag("tool_call")], + Annotated[ServerToolResult, Tag("tool_result")], + Annotated[ServerToolReturn, Tag("tool_return")], + Annotated[ServerAgentInvoked, Tag("agent_invoked")], + Annotated[ServerResponse, Tag("response")], + Annotated[ServerStreamEnd, Tag("stream_end")], + Annotated[ServerError, Tag("error")], + Annotated[ServerPermissionRequest, Tag("permission_request")], + Annotated[ServerUserInputRequest, Tag("user_input_request")], + Annotated[ServerConfirmationRequest, Tag("confirmation_request")], + Annotated[ServerSelectionRequest, Tag("selection_request")], + Annotated[ServerAskUserQuestionRequest, Tag("ask_user_question_request")], + Annotated[ServerCancelled, Tag("cancelled")], + ], + Discriminator("type"), +] +"""Discriminated union of all server→client message types.""" diff --git a/code_puppy/api/ws/send_utils.py b/code_puppy/api/ws/send_utils.py new file mode 100644 index 000000000..fae384129 --- /dev/null +++ b/code_puppy/api/ws/send_utils.py @@ -0,0 +1,155 @@ +"""WebSocket send helpers extracted from ``chat_handler.py`` (Phase 3). + +``WebSocketSender`` encapsulates the mutable ``ws_closed`` flag and the +four send/persist closures that were previously defined inside +``websocket_chat()``. Moving them to a class makes them independently +testable and removes four levels of closure nesting. +""" + +from __future__ import annotations + +import datetime +import logging +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from fastapi import WebSocket + + from code_puppy.api.ws.schemas import ( + ServerMessage, + ServerToolCall, + ServerToolResult, + ) + +logger = logging.getLogger(__name__) + + +class WebSocketSender: + """Manages safe JSON sends over a WebSocket with closed-socket detection. + + Parameters + ---------- + websocket: + The connected ``WebSocket`` instance. + session_id: + Current session identifier (used for logging and error persistence). + """ + + __slots__ = ("_websocket", "_session_id", "ws_closed", "_ctx") + + def __init__( + self, + websocket: WebSocket, + session_id: str, + ) -> None: + self._websocket = websocket + self._session_id = session_id + self.ws_closed: bool = False + self._ctx: Any = None + + # -- ctx property (set after construction, mirrors closure lifecycle) ---- + + @property + def ctx(self) -> Any: + return self._ctx + + @ctx.setter + def ctx(self, value: Any) -> None: + self._ctx = value + + @property + def session_id(self) -> str: + return self._session_id + + @session_id.setter + def session_id(self, value: str) -> None: + self._session_id = value + + # -- persistence helper -------------------------------------------------- + + async def persist_error_payload(self, data: dict[str, Any]) -> None: + """Persist structured error frames so they survive page reloads.""" + if data.get("type") != "error" or not self._session_id: + return + + try: + from code_puppy.api.db.queries import write_error_message_to_sqlite + + await write_error_message_to_sqlite( + session_id=self._session_id, + error=str(data.get("error") or "An unknown error occurred"), + error_type=str(data.get("error_type") or "unknown"), + technical_details=str(data.get("technical_details") or ""), + action_required=data.get("action_required"), + agent_name=(self._ctx.agent_name if self._ctx else ""), + model_name=(self._ctx.model_name if self._ctx else ""), + timestamp=datetime.datetime.now(datetime.timezone.utc).isoformat(), + ) + except Exception: + logger.warning( + "[WS:%s] Failed to persist error payload to SQLite", + self._session_id, + exc_info=True, + ) + + # -- send helpers -------------------------------------------------------- + + async def safe_send_json(self, data: dict) -> bool: + """Send JSON to WebSocket; returns ``False`` if the connection is closed.""" + if self.ws_closed: + logger.debug( + "[WS:%s] safe_send_json skipped (ws_closed=True): type=%s", + self._session_id, + data.get("type"), + ) + return False + + msg_type = data.get("type") + if msg_type in { + "error", + "response", + "assistant_message_start", + "assistant_message_end", + }: + logger.debug( + "[WS:%s] → send_json type=%s keys=%s", + self._session_id, + msg_type, + sorted(list(data.keys())), + ) + if msg_type == "error": + logger.debug( + "[WS:%s] error payload: error_type=%r action_required=%r error=%r", + self._session_id, + data.get("error_type"), + data.get("action_required"), + (data.get("error") or "")[:300], + ) + + try: + if msg_type == "error": + await self.persist_error_payload(data) + await self._websocket.send_json(data) + return True + except Exception as e: + logger.warning( + "[WS:%s] send_json failed for type=%s: %s", + self._session_id, + msg_type, + e, + exc_info=True, + ) + if "close message" in str(e).lower() or "closed" in str(e).lower(): + self.ws_closed = True + logger.debug("WebSocket closed, stopping sends") + return False + + async def send_typed(self, msg: ServerMessage) -> bool: + """Send a typed protocol message to the client.""" + return await self.safe_send_json(msg.model_dump(exclude_none=True)) + + async def send_typed_tool_lifecycle( + self, msg: ServerToolCall | ServerToolResult + ) -> bool: + """Send tool lifecycle frames to the client.""" + return await self.send_typed(msg) diff --git a/code_puppy/api/ws/session_persistence.py b/code_puppy/api/ws/session_persistence.py new file mode 100644 index 000000000..11e1ef86b --- /dev/null +++ b/code_puppy/api/ws/session_persistence.py @@ -0,0 +1,285 @@ +"""Session persistence and broadcast payload builders. + +Pure helpers extracted from ``chat_handler.py`` (Phase 3) to make payload +construction independently testable without a live WebSocket or database. + +These functions build the dicts that ``chat_handler`` sends to clients and +broadcasts to session-monitoring connections. They do **not** perform I/O +themselves — the caller is responsible for the actual send / DB write. +""" + +from __future__ import annotations + +import datetime +import logging +from dataclasses import dataclass +from typing import Any, Awaitable, Callable, Optional + +from code_puppy.api.ws.connection_manager import connection_manager +from code_puppy.api.ws.history_utils import ( + build_enhanced_history, + estimate_total_tokens, +) + +logger = logging.getLogger(__name__) + + +def generate_heuristic_title(history: list[Any], max_length: int = 50) -> str: + """Generate a short title from the first user message in websocket history.""" + import re + + def extract_user_content(msg: Any) -> str | None: + if isinstance(msg, dict) and "msg" in msg: + msg = msg["msg"] + + if hasattr(msg, "kind") and msg.kind == "request": + for part in getattr(msg, "parts", []): + content = getattr(part, "content", None) + if isinstance(content, str) and content.strip(): + return content.strip() + + if isinstance(msg, dict) and msg.get("role") == "user": + content = msg.get("content") + if isinstance(content, str) and content.strip(): + return content.strip() + + return None + + for msg in history: + content = extract_user_content(msg) + if content: + first_line = content.split("\n")[0][:max_length] + title = first_line.lower() + title = re.sub(r"[^a-z0-9\s-]", "", title) + title = re.sub(r"\s+", "-", title) + title = re.sub(r"-+", "-", title) + title = title.strip("-")[:max_length] + return title or "untitled-session" + + return "untitled-session" + + +def resolve_agent_model_meta( + agent: Any = None, + ctx: Any = None, +) -> tuple[str, str]: + """Return ``(agent_name, model_name)`` with or-chain fallbacks. + + Priority: + 1. ``agent.name`` / ``agent.get_model_name()`` (if agent is not None/empty) + 2. ``ctx.agent_name`` / ``ctx.model_name`` (if ctx is not None) + 3. Hardcoded defaults ``"code-puppy"`` / ``"unknown"`` + """ + agent_name = ( + (agent.name if agent else "") or (ctx.agent_name if ctx else "") or "code-puppy" + ) + model_name = ( + (agent.get_model_name() if agent else "") + or (ctx.model_name if ctx else "") + or "unknown" + ) + return agent_name, model_name + + +def build_session_meta_payload( + *, + session_id: str, + session_name: str, + total_tokens: int, + message_count: int, + title: str, + working_directory: str, + agent_name: str, + model_name: str, +) -> dict[str, Any]: + """Build the ``session_meta`` frame sent to the initiating client. + + Returns a plain dict ready for ``send_json`` / ``send_typed``. + """ + return { + "type": "session_meta", + "session_id": session_id, + "session_name": session_name, + "total_tokens": total_tokens, + "message_count": message_count, + "title": title, + "working_directory": working_directory, + "agent_name": agent_name, + "model_name": model_name, + } + + +def build_session_update_payload( + *, + session_id: str, + session_name: str, + title: str, + working_directory: str, + message_count: int, + total_tokens: int, + timestamp: Optional[str] = None, +) -> dict[str, Any]: + """Build the broadcast payload for ``connection_manager.broadcast_session_update``. + + ``action`` is ``"created"`` when ``message_count == 1`` (first turn), + otherwise ``"updated"``. + """ + return { + "session_id": session_id, + "session_name": session_name, + "title": title, + "working_directory": working_directory, + "timestamp": timestamp or datetime.datetime.now().isoformat(), + "message_count": message_count, + "total_tokens": total_tokens, + "auto_saved": True, + "action": "created" if message_count == 1 else "updated", + } + + +@dataclass(slots=True) +class PersistedTurnSummary: + """Summary of a persisted websocket turn.""" + + session_title: str + message_count: int + total_tokens: int + + +async def persist_session_turn_and_broadcast( + *, + history: list[Any], + session_id: str, + session_title: str, + session_working_directory: str, + session_pinned: bool, + agent: Any, + agent_name: str, + model_name: str, + ctx: Any, + original_user_message: str, + attachment_metadata: list[Any] | None, + safe_send_json: Callable[[dict[str, Any]], Awaitable[Any]], + logger_override: logging.Logger | None = None, +) -> PersistedTurnSummary | None: + """Persist finalized turn history and notify websocket/session listeners. + + Returns ``None`` when there is no history to persist. Otherwise returns a + compact summary containing the resolved title and aggregate counts. + """ + if not history: + return None + + logger_ = logger_override or logger + attachment_metadata = attachment_metadata or [] + + if not session_title or session_title == "untitled-session": + session_title = generate_heuristic_title(history) + + session_name = session_id + agent_name_meta, model_name_meta = resolve_agent_model_meta(agent=agent, ctx=ctx) + enhanced_history = build_enhanced_history( + history, + agent_name_meta=agent_name_meta, + model_name_meta=model_name_meta, + original_user_message=original_user_message, + attachment_metadata=attachment_metadata, + ) + + if attachment_metadata and len(history) >= 2: + logger_.debug( + "Added UI metadata to user message: %d attachment(s), clean_content length: %d", + len(attachment_metadata), + len(original_user_message), + ) + + message_count = len(enhanced_history) + total_tokens = estimate_total_tokens(enhanced_history, agent) + + await persist_turn_to_sqlite( + session_id=session_id, + enhanced_history=enhanced_history, + title=session_title, + working_directory=session_working_directory, + pinned=session_pinned, + agent_name=agent_name, + model_name=model_name, + total_tokens=total_tokens, + created_at_iso=ctx.created_at.isoformat(), + ctx=ctx, + ) + + await safe_send_json( + build_session_meta_payload( + session_id=session_id, + session_name=session_name, + total_tokens=total_tokens, + message_count=message_count, + title=session_title, + working_directory=session_working_directory, + agent_name=agent_name, + model_name=model_name, + ) + ) + + await connection_manager.broadcast_session_update( + build_session_update_payload( + session_id=session_id, + session_name=session_name, + title=session_title, + working_directory=session_working_directory, + message_count=message_count, + total_tokens=total_tokens, + ) + ) + + return PersistedTurnSummary( + session_title=session_title, + message_count=message_count, + total_tokens=total_tokens, + ) + + +async def persist_turn_to_sqlite( + *, + session_id: str, + enhanced_history: list[Any], + title: str, + working_directory: str, + pinned: bool, + agent_name: str, + model_name: str, + total_tokens: int, + created_at_iso: str, + ctx: Any = None, +) -> bool: + """Write a conversation turn to SQLite, returning True on success. + + Wraps ``write_turn_to_sqlite`` with the standard try/except guard + used by ``chat_handler``. Returns ``False`` (and logs) when the + database is unavailable instead of raising. + """ + try: + from code_puppy.api.db.queries import write_turn_to_sqlite + + now_iso = datetime.datetime.now(datetime.timezone.utc).isoformat() + await write_turn_to_sqlite( + session_id=session_id, + enhanced_history=enhanced_history, + title=title, + working_directory=working_directory, + pinned=pinned, + agent_name=agent_name, + model_name=model_name, + total_tokens=total_tokens, + updated_at=now_iso, + created_at=created_at_iso, + ctx=ctx, + ) + return True + except Exception as exc: + logger.debug( + "SQLite turn write skipped (DB not available): %s", + exc, + ) + return False diff --git a/code_puppy/api/ws/sessions_handler.py b/code_puppy/api/ws/sessions_handler.py new file mode 100644 index 000000000..f0719648c --- /dev/null +++ b/code_puppy/api/ws/sessions_handler.py @@ -0,0 +1,69 @@ +"""WebSocket endpoint for real-time session monitoring.""" + +import asyncio +import logging + +from fastapi import FastAPI, WebSocket, WebSocketDisconnect + +from code_puppy.api.ws.connection_manager import connection_manager + +logger = logging.getLogger(__name__) + + +def register_sessions_endpoint(app: FastAPI) -> None: + """Register the /ws/sessions WebSocket endpoint.""" + + @app.websocket("/ws/sessions") + async def websocket_sessions(websocket: WebSocket) -> None: + """WebSocket endpoint for real-time session updates. + + Clients can connect to this endpoint to receive real-time notifications + when WebSocket sessions are created, updated, or deleted. + + Messages sent to clients: + - {"type": "session_update", "data": WSSessionInfo} + - {"type": "ping"} (keepalive) + """ + await websocket.accept() + logger.debug("Sessions monitoring WebSocket client connected") + + try: + await connection_manager.connect_session_client(websocket) + + # Send initial session list + try: + from code_puppy.api.routers.ws_sessions import list_ws_sessions + + current_sessions = await list_ws_sessions() + await websocket.send_json( + { + "type": "initial_sessions", + "data": [session.model_dump() for session in current_sessions], + } + ) + except Exception as e: + logger.error("Failed to send initial sessions: %s", e) + await websocket.send_json( + {"type": "error", "error": f"Failed to load sessions: {str(e)}"} + ) + + # Keep connection alive + while True: + try: + message = await asyncio.wait_for( + websocket.receive_json(), timeout=30.0 + ) + if message.get("type") == "ping": + await websocket.send_json({"type": "pong"}) + except asyncio.TimeoutError: + try: + await websocket.send_json({"type": "ping"}) + except Exception: + break + + except WebSocketDisconnect: + logger.debug("Sessions monitoring WebSocket client disconnected") + except Exception as e: + logger.error("Sessions WebSocket error: %s", e) + finally: + connection_manager.disconnect_session_client(websocket) diff --git a/code_puppy/api/ws/tests/__init__.py b/code_puppy/api/ws/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/code_puppy/api/ws/tests/test_background_save.py b/code_puppy/api/ws/tests/test_background_save.py new file mode 100644 index 000000000..15919df9d --- /dev/null +++ b/code_puppy/api/ws/tests/test_background_save.py @@ -0,0 +1,559 @@ +"""Unit tests for code_puppy.api.ws.background_save. + +Covers every edge case identified in the Phase 3 plan: + +1. fire_and_track GC protection (task registered + deregistered). +2. None agent_task → early return, write never called. +3. Cancelled agent task → early return, write never called. +4. Failed agent task → early return, write never called. +5. None result from task → early return, write never called. +6. Empty history → early return, write never called. +7. Session deleted during run → write skipped (resurrection guard). +8. session_exists check failure → write proceeds (fail-safe). +9. ctx without created_at attribute → now_iso fallback used. +10. ctx is None → now_iso fallback used. +11. Pre-wrapped dict history entries pass through unchanged. +12. Bare ModelMessage-like objects get wrapped with agent/model/ts. +13. Token computation is called per message (not 0 as before). +14. write_turn_to_sqlite receives correct keyword arguments. +""" + +from __future__ import annotations + +import asyncio +import datetime +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_agent(history=None, token_per_msg=10): + """Return a minimal agent mock.""" + agent = MagicMock() + agent.get_message_history.return_value = history or [] + agent.set_message_history.return_value = None + agent.estimate_tokens_for_message.return_value = token_per_msg + return agent + + +def _make_result(messages=None): + """Return a minimal agent-run result mock.""" + result = MagicMock() + result.all_messages.return_value = messages or [] + return result + + +def _bare_msg(content="hello"): + """Simulate a bare ModelMessage-like object (has .parts attribute).""" + msg = MagicMock() + msg.__class__.__name__ = "ModelRequest" + # Make isinstance(msg, dict) == False (default for MagicMock) + return msg + + +def _wrapped_msg(content="wrapped"): + """Simulate an already-wrapped history dict entry.""" + return {"msg": MagicMock(), "agent": "code-puppy", "model": "gpt-4", "ts": "ts"} + + +# --------------------------------------------------------------------------- +# fire_and_track +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_fire_and_track_registers_and_deregisters_task(): + """Task should be in _BACKGROUND_TASKS while running and removed after.""" + from code_puppy.api.ws.background_save import _BACKGROUND_TASKS, fire_and_track + + async def _noop(): + await asyncio.sleep(0) + + task = fire_and_track(_noop()) + assert task in _BACKGROUND_TASKS + + await task + # done_callback fires synchronously when the task finishes + assert task not in _BACKGROUND_TASKS + + +# --------------------------------------------------------------------------- +# Early-exit guards +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_none_task_returns_immediately(): + """agent_task=None should be a safe no-op.""" + from code_puppy.api.ws.background_save import save_agent_result_in_background + + with patch("code_puppy.api.ws.background_save.write_turn_to_sqlite") as mock_write: + # patch is not needed since we never reach the write, but kept for safety + await save_agent_result_in_background( + agent_task=None, + session_id="sess-1", + ctx=None, + agent=_make_agent(), + agent_name="puppy", + model_name="gpt-4", + title="Test", + working_directory="/tmp", + pinned=False, + ) + mock_write.assert_not_called() + + +@pytest.mark.asyncio +async def test_cancelled_task_returns_early(): + """CancelledError from the agent task should be swallowed — no write.""" + from code_puppy.api.ws.background_save import save_agent_result_in_background + + async def _cancelled(): + raise asyncio.CancelledError() + + task = asyncio.ensure_future(_cancelled()) + with pytest.raises(asyncio.CancelledError): + await task # drain so it's actually cancelled + + # Create a new cancelled-style task via a Future + cancelled_task = asyncio.get_event_loop().create_future() + cancelled_task.cancel() + + with patch("code_puppy.api.ws.background_save.write_turn_to_sqlite") as mock_write: + await save_agent_result_in_background( + agent_task=cancelled_task, + session_id="sess-cancel", + ctx=None, + agent=_make_agent(), + agent_name="puppy", + model_name="gpt-4", + title="T", + working_directory="/", + pinned=False, + ) + mock_write.assert_not_called() + + +@pytest.mark.asyncio +async def test_failed_task_returns_early(): + """An exception from the agent task should be swallowed — no write.""" + from code_puppy.api.ws.background_save import save_agent_result_in_background + + async def _fail(): + raise RuntimeError("boom") + + task = asyncio.ensure_future(_fail()) + try: + await task + except RuntimeError: + pass + + # Create a fresh failed future + failed_task = asyncio.get_event_loop().create_future() + failed_task.set_exception(RuntimeError("agent boom")) + + with patch("code_puppy.api.ws.background_save.write_turn_to_sqlite") as mock_write: + await save_agent_result_in_background( + agent_task=failed_task, + session_id="sess-fail", + ctx=None, + agent=_make_agent(), + agent_name="puppy", + model_name="gpt-4", + title="T", + working_directory="/", + pinned=False, + ) + mock_write.assert_not_called() + + +@pytest.mark.asyncio +async def test_none_result_returns_early(): + """result=None from await agent_task should skip the write.""" + from code_puppy.api.ws.background_save import save_agent_result_in_background + + task_future = asyncio.get_event_loop().create_future() + task_future.set_result(None) # agent returned None + + with patch("code_puppy.api.ws.background_save.write_turn_to_sqlite") as mock_write: + await save_agent_result_in_background( + agent_task=task_future, + session_id="sess-none-result", + ctx=None, + agent=_make_agent(), + agent_name="puppy", + model_name="gpt-4", + title="T", + working_directory="/", + pinned=False, + ) + mock_write.assert_not_called() + + +@pytest.mark.asyncio +async def test_empty_history_returns_early(): + """Empty message history should skip the write.""" + from code_puppy.api.ws.background_save import save_agent_result_in_background + + task_future = asyncio.get_event_loop().create_future() + task_future.set_result(_make_result(messages=[])) + + agent = _make_agent(history=[]) # empty history after sync + + with ( + patch( + "code_puppy.api.ws.background_save.session_exists", new_callable=AsyncMock + ) as mock_exists, + patch("code_puppy.api.ws.background_save.write_turn_to_sqlite") as mock_write, + ): + mock_exists.return_value = True + await save_agent_result_in_background( + agent_task=task_future, + session_id="sess-empty", + ctx=None, + agent=agent, + agent_name="puppy", + model_name="gpt-4", + title="T", + working_directory="/", + pinned=False, + ) + mock_write.assert_not_called() + + +# --------------------------------------------------------------------------- +# Session-deleted resurrection guard +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_deleted_session_skips_write(): + """If session_exists returns False, write_turn_to_sqlite must NOT be called.""" + from code_puppy.api.ws.background_save import save_agent_result_in_background + + msg = _bare_msg() + task_future = asyncio.get_event_loop().create_future() + task_future.set_result(_make_result(messages=[msg])) + agent = _make_agent(history=[msg]) + + with ( + patch( + "code_puppy.api.ws.background_save.session_exists", new_callable=AsyncMock + ) as mock_exists, + patch("code_puppy.api.ws.background_save.write_turn_to_sqlite") as mock_write, + ): + mock_exists.return_value = False # session was deleted + await save_agent_result_in_background( + agent_task=task_future, + session_id="sess-deleted", + ctx=None, + agent=agent, + agent_name="puppy", + model_name="gpt-4", + title="T", + working_directory="/", + pinned=False, + ) + mock_exists.assert_awaited_once_with("sess-deleted") + mock_write.assert_not_called() + + +@pytest.mark.asyncio +async def test_session_exists_failure_proceeds_anyway(): + """If session_exists raises, we log a warning but still write (fail-safe).""" + from code_puppy.api.ws.background_save import save_agent_result_in_background + + msg = _bare_msg() + task_future = asyncio.get_event_loop().create_future() + task_future.set_result(_make_result(messages=[msg])) + agent = _make_agent(history=[msg]) + + with ( + patch( + "code_puppy.api.ws.background_save.session_exists", + new_callable=AsyncMock, + side_effect=Exception("db gone"), + ), + patch( + "code_puppy.api.ws.background_save.write_turn_to_sqlite", + new_callable=AsyncMock, + ) as mock_write, + ): + await save_agent_result_in_background( + agent_task=task_future, + session_id="sess-db-gone", + ctx=None, + agent=agent, + agent_name="puppy", + model_name="gpt-4", + title="T", + working_directory="/", + pinned=False, + ) + mock_write.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# created_at edge cases +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_ctx_without_created_at_uses_now_iso(): + """ctx that lacks created_at should fall back to now_iso without crashing.""" + from code_puppy.api.ws.background_save import save_agent_result_in_background + + ctx_no_attr = SimpleNamespace() # no created_at attribute + msg = _bare_msg() + task_future = asyncio.get_event_loop().create_future() + task_future.set_result(_make_result(messages=[msg])) + agent = _make_agent(history=[msg]) + + with ( + patch( + "code_puppy.api.ws.background_save.session_exists", new_callable=AsyncMock + ) as mock_exists, + patch( + "code_puppy.api.ws.background_save.write_turn_to_sqlite", + new_callable=AsyncMock, + ) as mock_write, + ): + mock_exists.return_value = True + await save_agent_result_in_background( + agent_task=task_future, + session_id="sess-no-attr", + ctx=ctx_no_attr, + agent=agent, + agent_name="puppy", + model_name="gpt-4", + title="T", + working_directory="/", + pinned=False, + ) + + call_kwargs = mock_write.call_args.kwargs + # created_at should be an ISO datetime string, not an error + created_at = call_kwargs["created_at"] + assert isinstance(created_at, str) + datetime.datetime.fromisoformat(created_at) # must be valid ISO + + +@pytest.mark.asyncio +async def test_ctx_none_uses_now_iso(): + """ctx=None should fall back to now_iso for created_at.""" + from code_puppy.api.ws.background_save import save_agent_result_in_background + + msg = _bare_msg() + task_future = asyncio.get_event_loop().create_future() + task_future.set_result(_make_result(messages=[msg])) + agent = _make_agent(history=[msg]) + + with ( + patch( + "code_puppy.api.ws.background_save.session_exists", new_callable=AsyncMock + ) as mock_exists, + patch( + "code_puppy.api.ws.background_save.write_turn_to_sqlite", + new_callable=AsyncMock, + ) as mock_write, + ): + mock_exists.return_value = True + await save_agent_result_in_background( + agent_task=task_future, + session_id="sess-ctx-none", + ctx=None, + agent=agent, + agent_name="puppy", + model_name="gpt-4", + title="T", + working_directory="/", + pinned=False, + ) + + call_kwargs = mock_write.call_args.kwargs + created_at = call_kwargs["created_at"] + assert isinstance(created_at, str) + datetime.datetime.fromisoformat(created_at) + + +# --------------------------------------------------------------------------- +# History wrapping +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_pre_wrapped_dict_entries_pass_through(): + """Dicts with 'msg' key must not be double-wrapped.""" + from code_puppy.api.ws.background_save import save_agent_result_in_background + + wrapped = _wrapped_msg() + task_future = asyncio.get_event_loop().create_future() + task_future.set_result(_make_result(messages=[wrapped])) + agent = _make_agent(history=[wrapped]) + + with ( + patch( + "code_puppy.api.ws.background_save.session_exists", new_callable=AsyncMock + ) as mock_exists, + patch( + "code_puppy.api.ws.background_save.write_turn_to_sqlite", + new_callable=AsyncMock, + ) as mock_write, + ): + mock_exists.return_value = True + await save_agent_result_in_background( + agent_task=task_future, + session_id="sess-wrapped", + ctx=None, + agent=agent, + agent_name="puppy", + model_name="gpt-4", + title="T", + working_directory="/", + pinned=False, + ) + + eh = mock_write.call_args.kwargs["enhanced_history"] + assert len(eh) == 1 + assert eh[0] is wrapped # same object, not a new wrapper + + +@pytest.mark.asyncio +async def test_bare_message_gets_wrapped(): + """Bare ModelMessage objects must be wrapped with agent/model/ts metadata.""" + from code_puppy.api.ws.background_save import save_agent_result_in_background + + bare = _bare_msg() + task_future = asyncio.get_event_loop().create_future() + task_future.set_result(_make_result(messages=[bare])) + agent = _make_agent(history=[bare]) + + with ( + patch( + "code_puppy.api.ws.background_save.session_exists", new_callable=AsyncMock + ) as mock_exists, + patch( + "code_puppy.api.ws.background_save.write_turn_to_sqlite", + new_callable=AsyncMock, + ) as mock_write, + ): + mock_exists.return_value = True + await save_agent_result_in_background( + agent_task=task_future, + session_id="sess-bare", + ctx=None, + agent=agent, + agent_name="my-agent", + model_name="gpt-4o", + title="T", + working_directory="/", + pinned=False, + ) + + eh = mock_write.call_args.kwargs["enhanced_history"] + assert len(eh) == 1 + wrapper = eh[0] + assert wrapper["msg"] is bare + assert wrapper["agent"] == "my-agent" + assert wrapper["model"] == "gpt-4o" + assert "ts" in wrapper + + +# --------------------------------------------------------------------------- +# Token computation +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_token_count_computed_not_zero(): + """total_tokens must reflect the actual estimate, not hard-coded 0.""" + from code_puppy.api.ws.background_save import save_agent_result_in_background + + bare = _bare_msg() + task_future = asyncio.get_event_loop().create_future() + task_future.set_result(_make_result(messages=[bare])) + agent = _make_agent(history=[bare], token_per_msg=42) + + with ( + patch( + "code_puppy.api.ws.background_save.session_exists", new_callable=AsyncMock + ) as mock_exists, + patch( + "code_puppy.api.ws.background_save.write_turn_to_sqlite", + new_callable=AsyncMock, + ) as mock_write, + ): + mock_exists.return_value = True + await save_agent_result_in_background( + agent_task=task_future, + session_id="sess-tokens", + ctx=None, + agent=agent, + agent_name="puppy", + model_name="gpt-4", + title="T", + working_directory="/", + pinned=False, + ) + + total_tokens = mock_write.call_args.kwargs["total_tokens"] + assert total_tokens == 42 # 1 message × 42 tokens each + agent.estimate_tokens_for_message.assert_called_once() + + +# --------------------------------------------------------------------------- +# Happy-path write args +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_happy_path_write_kwargs(): + """write_turn_to_sqlite must receive the correct session metadata.""" + from code_puppy.api.ws.background_save import save_agent_result_in_background + + created = datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc) + ctx = SimpleNamespace(created_at=created) + + bare = _bare_msg() + task_future = asyncio.get_event_loop().create_future() + task_future.set_result(_make_result(messages=[bare])) + agent = _make_agent(history=[bare], token_per_msg=5) + + with ( + patch( + "code_puppy.api.ws.background_save.session_exists", new_callable=AsyncMock + ) as mock_exists, + patch( + "code_puppy.api.ws.background_save.write_turn_to_sqlite", + new_callable=AsyncMock, + ) as mock_write, + ): + mock_exists.return_value = True + await save_agent_result_in_background( + agent_task=task_future, + session_id="sess-happy", + ctx=ctx, + agent=agent, + agent_name="code-puppy", + model_name="claude-3", + title="My Session", + working_directory="/home/user", + pinned=True, + label="switch", + ) + + kw = mock_write.call_args.kwargs + assert kw["session_id"] == "sess-happy" + assert kw["agent_name"] == "code-puppy" + assert kw["model_name"] == "claude-3" + assert kw["title"] == "My Session" + assert kw["working_directory"] == "/home/user" + assert kw["pinned"] is True + assert kw["created_at"] == created.isoformat() + assert kw["total_tokens"] == 5 + assert kw["ctx"] is ctx diff --git a/code_puppy/api/ws/tests/test_chat_context.py b/code_puppy/api/ws/tests/test_chat_context.py new file mode 100644 index 000000000..f7792219f --- /dev/null +++ b/code_puppy/api/ws/tests/test_chat_context.py @@ -0,0 +1,58 @@ +from code_puppy.api.permission_plugin import ( + get_suppress_emitter_tool_events, + get_websocket_context, +) +from code_puppy.api.ws.chat_context import ( + begin_agent_run_context, + cleanup_agent_run_context, + cleanup_message_context, + setup_message_context, +) +from code_puppy.plugins.frontend_emitter.session_context import ( + current_emitter_session_id, +) + + +def test_setup_and_cleanup_message_context(): + calls = [] + + setup_message_context(websocket="ws-obj", session_id="session-123") + assert get_websocket_context() == ("ws-obj", "session-123") + + cleanup_message_context( + clear_session_working_directory=lambda: calls.append("cleared") + ) + assert get_websocket_context() is None + assert calls == ["cleared"] + + +def test_begin_and_cleanup_agent_run_context_resets_all_contextvars(): + calls = [] + assert current_emitter_session_id.get() is None + assert get_suppress_emitter_tool_events() is False + + run_context = begin_agent_run_context(session_id="session-abc") + + assert get_suppress_emitter_tool_events() is True + assert current_emitter_session_id.get() == "session-abc" + + cleanup_agent_run_context( + run_context, + clear_session_working_directory=lambda: calls.append("cleared"), + ) + + assert get_suppress_emitter_tool_events() is False + assert current_emitter_session_id.get() is None + assert get_websocket_context() is None + assert calls == ["cleared"] + + +def test_cleanup_agent_run_context_accepts_none_context(): + calls = [] + cleanup_agent_run_context( + None, + clear_session_working_directory=lambda: calls.append("cleared"), + ) + assert get_suppress_emitter_tool_events() is False + assert get_websocket_context() is None + assert calls == ["cleared"] diff --git a/code_puppy/api/ws/tests/test_chat_event_adapter.py b/code_puppy/api/ws/tests/test_chat_event_adapter.py new file mode 100644 index 000000000..275d4b9f8 --- /dev/null +++ b/code_puppy/api/ws/tests/test_chat_event_adapter.py @@ -0,0 +1,164 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from code_puppy.api.ws.chat_event_adapter import ( + collect_final_stream_text_delta, + handle_assistant_part_delta, + handle_assistant_part_end, + handle_assistant_part_start, +) +from code_puppy.api.ws.chat_turn_state import WebSocketTurnState + + +class _Logger: + def __init__(self): + self.debug_messages: list[tuple[str, tuple[object, ...]]] = [] + + def debug(self, message, *args): + self.debug_messages.append((message, args)) + + +@pytest.mark.asyncio +async def test_handle_assistant_part_start_sends_start_and_initial_delta(): + turn_state = WebSocketTurnState(current_tool_name="shell") + payloads = [] + status_calls = [] + + async def _safe_send_json(payload): + payloads.append(payload) + + async def _status_only(): + status_calls.append("called") + + handled = await handle_assistant_part_start( + turn_state=turn_state, + part_index=0, + part_type="TextPart", + part_obj={"content": "hello"}, + session_id="session-1", + agent_name="agent-1", + model_name="model-1", + safe_send_json=_safe_send_json, + logger=_Logger(), + send_status_only_pending_tool_results=_status_only, + ) + + assert handled is True + assert status_calls == [] + assert len(payloads) == 2 + assert payloads[0]["type"] == "assistant_message_start" + assert payloads[0]["part_type"] == "text" + assert payloads[1]["type"] == "assistant_message_delta" + assert payloads[1]["content"] == "hello" + assert turn_state.current_tool_group_id is None + assert turn_state.collected_text == ["hello"] + assert turn_state.b1_streaming_used is True + + +@pytest.mark.asyncio +async def test_handle_assistant_part_start_reuses_existing_early_delta_part(): + turn_state = WebSocketTurnState( + active_parts={2: {"id": "msg-existing", "type": "text", "content": "world"}}, + collected_text=["world"], + ) + payloads = [] + + async def _safe_send_json(payload): + payloads.append(payload) + + handled = await handle_assistant_part_start( + turn_state=turn_state, + part_index=2, + part_type="ThinkingPart", + part_obj=SimpleNamespace(content="hello "), + session_id="session-2", + agent_name="agent-2", + model_name="model-2", + safe_send_json=_safe_send_json, + logger=_Logger(), + ) + + assert handled is True + assert payloads == [] + assert turn_state.active_parts[2]["id"] == "msg-existing" + assert turn_state.active_parts[2]["type"] == "thinking" + assert turn_state.active_parts[2]["content"] == "hello world" + assert turn_state.collected_text[0] == "hello " + + +@pytest.mark.asyncio +async def test_handle_assistant_part_delta_creates_missing_part_then_sends_delta(): + turn_state = WebSocketTurnState(current_tool_name="calculator") + payloads = [] + + async def _safe_send_json(payload): + payloads.append(payload) + + handled = await handle_assistant_part_delta( + turn_state=turn_state, + part_index=0, + inner_data={"content_delta": "abc"}, + delta_obj={}, + session_id="session-3", + agent_name="agent-3", + model_name="model-3", + safe_send_json=_safe_send_json, + logger=_Logger(), + ) + + assert handled is True + assert [p["type"] for p in payloads] == [ + "assistant_message_start", + "assistant_message_delta", + ] + assert payloads[1]["content"] == "abc" + assert turn_state.active_parts[0]["content"] == "abc" + assert turn_state.collected_text == ["abc"] + assert turn_state.b1_streaming_used is True + assert turn_state.current_tool_group_id is None + + +@pytest.mark.asyncio +async def test_handle_assistant_part_end_sends_end_and_cleans_up(): + turn_state = WebSocketTurnState( + active_parts={1: {"id": "msg-1", "type": "text", "content": "done"}} + ) + payloads = [] + + async def _safe_send_json(payload): + payloads.append(payload) + + handled = await handle_assistant_part_end( + turn_state=turn_state, + part_index=1, + session_id="session-4", + agent_name="agent-4", + model_name="model-4", + safe_send_json=_safe_send_json, + ) + + assert handled is True + assert payloads[0]["type"] == "assistant_message_end" + assert payloads[0]["full_content"] == "done" + assert 1 not in turn_state.active_parts + + +def test_collect_final_stream_text_delta_appends_nested_content_delta(): + turn_state = WebSocketTurnState() + + handled = collect_final_stream_text_delta( + turn_state=turn_state, + event={ + "type": "stream_event", + "data": { + "event_type": "part_delta", + "event_data": {"delta": {"content_delta": "tail"}}, + }, + }, + ) + + assert handled is True + assert turn_state.collected_text == ["tail"] diff --git a/code_puppy/api/ws/tests/test_chat_tool_lifecycle.py b/code_puppy/api/ws/tests/test_chat_tool_lifecycle.py new file mode 100644 index 000000000..5ebc19252 --- /dev/null +++ b/code_puppy/api/ws/tests/test_chat_tool_lifecycle.py @@ -0,0 +1,231 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from code_puppy.api.ws.chat_tool_lifecycle import ( + accumulate_tool_call_part_delta, + finish_tool_call_part, + handle_tool_call_complete_event, + handle_tool_call_start_event, + resolve_pending_tool_id, + resolve_tool_group_id, + send_status_only_pending_tool_results, + start_tool_call_part, + start_tool_return_part, +) +from code_puppy.api.ws.chat_turn_state import WebSocketTurnState + + +class _Logger: + def __init__(self): + self.debug_messages = [] + self.info_messages = [] + self.warning_messages = [] + + def debug(self, message, *args): + self.debug_messages.append((message, args)) + + def info(self, message, *args): + self.info_messages.append((message, args)) + + def warning(self, message, *args): + self.warning_messages.append((message, args)) + + +@pytest.mark.asyncio +async def test_handle_tool_call_start_event_tracks_pending_and_emits_call(): + turn_state = WebSocketTurnState() + sent = [] + + async def _send(model): + sent.append(model) + + await handle_tool_call_start_event( + turn_state=turn_state, + event_data={"tool_name": "shell", "tool_args": {"cmd": "pwd"}}, + session_id="session-1", + agent_name="agent-1", + model_name="model-1", + send_typed_tool_lifecycle=_send, + logger=_Logger(), + ) + + assert len(sent) == 1 + assert sent[0].tool_name == "shell" + assert sent[0].tool_group_id.startswith("tg-") + assert turn_state.current_tool_name == "shell" + assert list(turn_state.pending_tool_calls.values())[0]["tool_name"] == "shell" + + +@pytest.mark.asyncio +async def test_handle_tool_call_complete_event_reuses_pending_group_id(): + turn_state = WebSocketTurnState( + pending_tool_calls={ + "tool-1": {"tool_name": "shell", "tool_group_id": "tg-abc"} + }, + current_tool_name="shell", + current_tool_group_id="tg-abc", + ) + sent = [] + + async def _send(model): + sent.append(model) + + await handle_tool_call_complete_event( + turn_state=turn_state, + event_data={"tool_name": "shell", "result": {"ok": True}, "duration_ms": 12}, + session_id="session-2", + agent_name="agent-2", + model_name="model-2", + send_typed_tool_lifecycle=_send, + logger=_Logger(), + ) + + assert len(sent) == 1 + assert sent[0].tool_id == "tool-1" + assert sent[0].tool_group_id == "tg-abc" + assert turn_state.pending_tool_calls == {} + assert turn_state.current_tool_name is None + + +@pytest.mark.asyncio +async def test_tool_call_part_start_delta_end_round_trip_emits_tool_call(): + turn_state = WebSocketTurnState() + sent = [] + logger = _Logger() + + async def _send(model): + sent.append(model) + + handled = start_tool_call_part( + turn_state=turn_state, + part_index=3, + part_obj=SimpleNamespace(tool_name="calc", tool_call_id="raw-1", args='{"a":'), + logger=logger, + ) + assert handled is True + + assert ( + accumulate_tool_call_part_delta( + turn_state=turn_state, + part_index=3, + delta_obj={"args_delta": "1}"}, + ) + is True + ) + + await finish_tool_call_part( + turn_state=turn_state, + part_index=3, + part_info=turn_state.active_parts[3], + session_id="session-3", + agent_name="agent-3", + model_name="model-3", + send_typed_tool_lifecycle=_send, + logger=logger, + ) + + assert len(sent) == 1 + assert sent[0].tool_name == "calc" + assert sent[0].args == {"a": 1} + assert turn_state.pending_tool_calls[sent[0].tool_id]["raw_tool_call_id"] == "raw-1" + assert turn_state.tool_group_ids[sent[0].tool_id] == sent[0].tool_group_id + assert 3 not in turn_state.active_parts + + +@pytest.mark.asyncio +async def test_start_tool_return_part_matches_pending_by_raw_tool_call_id(): + turn_state = WebSocketTurnState( + pending_tool_calls={ + "tool-1": { + "tool_name": "calc", + "start_time": 0.0, + "part_index": 5, + "raw_tool_call_id": "raw-1", + "tool_group_id": "tg-xyz", + } + }, + current_tool_group_id="tg-xyz", + ) + sent = [] + + async def _send(model): + sent.append(model) + + await start_tool_return_part( + turn_state=turn_state, + part_index=6, + part_obj={"tool_call_id": "raw-1", "content": {"answer": 42}}, + session_id="session-4", + agent_name="agent-4", + model_name="model-4", + send_typed_tool_lifecycle=_send, + logger=_Logger(), + ) + + assert len(sent) == 1 + assert sent[0].tool_id == "tool-1" + assert sent[0].tool_group_id == "tg-xyz" + assert turn_state.pending_tool_calls["tool-1"]["result"] == {"answer": 42} + + +@pytest.mark.asyncio +async def test_send_status_only_pending_tool_results_marks_sent(): + turn_state = WebSocketTurnState( + pending_tool_calls={ + "tool-1": { + "tool_name": "shell", + "start_time": 0.0, + "tool_group_id": "tg-shell", + "status_only_sent": False, + } + }, + current_tool_group_id="tg-shell", + ) + sent = [] + + async def _send(model): + sent.append(model) + + await send_status_only_pending_tool_results( + turn_state=turn_state, + session_id="session-5", + agent_name="agent-5", + model_name="model-5", + send_typed=_send, + logger=_Logger(), + ) + + assert len(sent) == 1 + assert sent[0].tool_group_id == "tg-shell" + assert sent[0].result["_pending_full_result"] is True + assert turn_state.pending_tool_calls["tool-1"]["status_only_sent"] is True + + +def test_resolve_pending_tool_id_and_group_id_fallbacks(): + logger = _Logger() + turn_state = WebSocketTurnState( + pending_tool_calls={ + "tool-1": {"tool_name": "shell", "raw_tool_call_id": "raw-1"} + } + ) + + assert ( + resolve_pending_tool_id(turn_state=turn_state, tool_call_id="raw-1") == "tool-1" + ) + assert resolve_pending_tool_id(turn_state=turn_state, tool_name="shell") == "tool-1" + + group_id = resolve_tool_group_id( + turn_state=turn_state, + logger=logger, + tool_id="tool-1", + pending_info=turn_state.pending_tool_calls["tool-1"], + tool_name="shell", + source="test", + ) + + assert group_id.startswith("tg-") + assert turn_state.tool_group_ids["tool-1"] == group_id + assert turn_state.pending_tool_calls["tool-1"]["tool_group_id"] == group_id diff --git a/code_puppy/api/ws/tests/test_chat_turn_runner.py b/code_puppy/api/ws/tests/test_chat_turn_runner.py new file mode 100644 index 000000000..c5bde7f3e --- /dev/null +++ b/code_puppy/api/ws/tests/test_chat_turn_runner.py @@ -0,0 +1,227 @@ +from __future__ import annotations + +import asyncio +from contextlib import suppress +from types import SimpleNamespace + +import pytest +from fastapi import WebSocketDisconnect + +from code_puppy.api.ws.chat_turn_runner import execute_turn_runner +from code_puppy.api.ws.chat_turn_state import WebSocketTurnState + + +class _FakeWebSocket: + def __init__(self, messages): + self._messages = list(messages) + + async def receive_json(self): + if not self._messages: + await asyncio.Future() + next_item = self._messages.pop(0) + if isinstance(next_item, BaseException): + raise next_item + return next_item + + +class _BlockingAgent: + def __init__(self, *, on_cancel=None, complete_event=None, result=None): + self.on_cancel = on_cancel + self.complete_event = complete_event or asyncio.Event() + self.result = result + + async def run_with_mcp(self, message_to_send, **run_kwargs): + try: + await self.complete_event.wait() + return self.result + except asyncio.CancelledError: + if self.on_cancel is not None: + self.on_cancel() + raise + + +@pytest.mark.asyncio +async def test_execute_turn_runner_processes_permission_response(monkeypatch): + handled = [] + complete_event = asyncio.Event() + result = SimpleNamespace(name="result") + + def _handle_permission_response(request_id, approved, session_id=None): + handled.append((request_id, approved, session_id)) + complete_event.set() + return True + + monkeypatch.setattr( + "code_puppy.api.permissions.handle_permission_response", + _handle_permission_response, + ) + + cleanup_calls = [] + agent = _BlockingAgent(complete_event=complete_event, result=result) + websocket = _FakeWebSocket( + [{"type": "permission_response", "request_id": "req-1", "approved": True}] + ) + + outcome = await execute_turn_runner( + websocket=websocket, + session_id="session-1", + ctx=SimpleNamespace(), + agent=agent, + agent_name="code-puppy", + model_name="gpt-test", + session_title="title", + session_working_directory="/tmp", + session_pinned=False, + message_to_send="hello", + run_kwargs={}, + turn_state=WebSocketTurnState(), + clear_session_working_directory=lambda: cleanup_calls.append("cleared"), + ) + + assert handled == [("req-1", True, "session-1")] + assert outcome.result is result + assert outcome.deferred_msg is None + assert cleanup_calls == ["cleared"] + + +@pytest.mark.asyncio +async def test_execute_turn_runner_cancels_active_agent_on_cancel_message(): + cancel_seen = [] + cleanup_calls = [] + agent = _BlockingAgent(on_cancel=lambda: cancel_seen.append(True)) + websocket = _FakeWebSocket([{"type": "cancel"}]) + turn_state = WebSocketTurnState() + + outcome = await execute_turn_runner( + websocket=websocket, + session_id="session-2", + ctx=SimpleNamespace(), + agent=agent, + agent_name="code-puppy", + model_name="gpt-test", + session_title="title", + session_working_directory="/tmp", + session_pinned=False, + message_to_send="hello", + run_kwargs={}, + turn_state=turn_state, + clear_session_working_directory=lambda: cleanup_calls.append("cleared"), + ) + await asyncio.sleep(0) + + assert outcome.result is None + assert outcome.deferred_msg is None + assert cancel_seen == [True] + assert cleanup_calls == ["cleared"] + + +@pytest.mark.asyncio +async def test_execute_turn_runner_defers_switch_and_background_saves(monkeypatch): + bg_calls = [] + bg_tasks = [] + agent = _BlockingAgent() + websocket = _FakeWebSocket( + [{"type": "switch_session", "session_id": "next-session"}] + ) + + async def _fake_background_save(**kwargs): + bg_calls.append(kwargs) + task = kwargs.get("agent_task") + if task is not None: + task.cancel() + with suppress(asyncio.CancelledError): + await task + + def _fake_fire_and_track(coro): + task = asyncio.create_task(coro) + bg_tasks.append(task) + return task + + monkeypatch.setattr( + "code_puppy.api.ws.chat_turn_runner.save_agent_result_in_background", + _fake_background_save, + ) + monkeypatch.setattr( + "code_puppy.api.ws.chat_turn_runner.fire_and_track", + _fake_fire_and_track, + ) + + cleanup_calls = [] + outcome = await execute_turn_runner( + websocket=websocket, + session_id="session-3", + ctx=SimpleNamespace(), + agent=agent, + agent_name="code-puppy", + model_name="gpt-test", + session_title="title", + session_working_directory="/tmp", + session_pinned=True, + message_to_send="hello", + run_kwargs={}, + turn_state=WebSocketTurnState(), + clear_session_working_directory=lambda: cleanup_calls.append("cleared"), + ) + await asyncio.gather(*bg_tasks) + + assert outcome.result is None + assert outcome.deferred_msg == { + "type": "switch_session", + "session_id": "next-session", + } + assert bg_calls and bg_calls[0]["label"] == "switch" + assert bg_calls[0]["pinned"] is True + assert cleanup_calls == ["cleared"] + + +@pytest.mark.asyncio +async def test_execute_turn_runner_background_saves_on_disconnect(monkeypatch): + bg_calls = [] + bg_tasks = [] + agent = _BlockingAgent() + websocket = _FakeWebSocket([WebSocketDisconnect()]) + + async def _fake_background_save(**kwargs): + bg_calls.append(kwargs) + task = kwargs.get("agent_task") + if task is not None: + task.cancel() + with suppress(asyncio.CancelledError): + await task + + def _fake_fire_and_track(coro): + task = asyncio.create_task(coro) + bg_tasks.append(task) + return task + + monkeypatch.setattr( + "code_puppy.api.ws.chat_turn_runner.save_agent_result_in_background", + _fake_background_save, + ) + monkeypatch.setattr( + "code_puppy.api.ws.chat_turn_runner.fire_and_track", + _fake_fire_and_track, + ) + + cleanup_calls = [] + outcome = await execute_turn_runner( + websocket=websocket, + session_id="session-4", + ctx=SimpleNamespace(), + agent=agent, + agent_name="code-puppy", + model_name="gpt-test", + session_title="title", + session_working_directory="/tmp", + session_pinned=False, + message_to_send="hello", + run_kwargs={}, + turn_state=WebSocketTurnState(), + clear_session_working_directory=lambda: cleanup_calls.append("cleared"), + ) + await asyncio.gather(*bg_tasks) + + assert outcome.result is None + assert outcome.deferred_msg is None + assert bg_calls and bg_calls[0]["label"] == "disconnect" + assert cleanup_calls == ["cleared"] diff --git a/code_puppy/api/ws/tests/test_chat_turn_runner_user_interactions.py b/code_puppy/api/ws/tests/test_chat_turn_runner_user_interactions.py new file mode 100644 index 000000000..fbd9a6554 --- /dev/null +++ b/code_puppy/api/ws/tests/test_chat_turn_runner_user_interactions.py @@ -0,0 +1,129 @@ +"""Tests for user-interaction responses during active WebSocket turns.""" + +import asyncio + +import pytest + +from code_puppy.api.ws.chat_turn_runner import handle_user_interaction_response +from code_puppy.messaging.bus import get_message_bus, reset_message_bus + + +@pytest.fixture(autouse=True) +def clean_bus(): + reset_message_bus() + yield + reset_message_bus() + + +async def _wait_for_prompt_id(): + bus = get_message_bus() + for _ in range(100): + message = bus.get_message_nowait() + if message is not None: + return message.prompt_id + await asyncio.sleep(0.01) + raise AssertionError("MessageBus request was not emitted") + + +@pytest.mark.asyncio +async def test_active_turn_handles_user_input_response(): + bus = get_message_bus() + bus.mark_renderer_active() + task = asyncio.create_task(bus.request_input("Name?")) + prompt_id = await _wait_for_prompt_id() + + assert handle_user_interaction_response( + {"type": "user_input_response", "prompt_id": prompt_id, "value": "Puppy"} + ) + + assert await asyncio.wait_for(task, timeout=1) == "Puppy" + + +@pytest.mark.asyncio +async def test_active_turn_handles_confirmation_response(): + bus = get_message_bus() + bus.mark_renderer_active() + task = asyncio.create_task( + bus.request_confirmation("Proceed?", "Should we continue?", allow_feedback=True) + ) + prompt_id = await _wait_for_prompt_id() + + assert handle_user_interaction_response( + { + "type": "confirmation_response", + "prompt_id": prompt_id, + "confirmed": True, + "feedback": "looks good", + } + ) + + assert await asyncio.wait_for(task, timeout=1) == (True, "looks good") + + +@pytest.mark.asyncio +async def test_active_turn_handles_selection_response(): + bus = get_message_bus() + bus.mark_renderer_active() + task = asyncio.create_task(bus.request_selection("Pick", ["A", "B"])) + prompt_id = await _wait_for_prompt_id() + + assert handle_user_interaction_response( + { + "type": "selection_response", + "prompt_id": prompt_id, + "selected_index": 1, + "selected_value": "B", + } + ) + + assert await asyncio.wait_for(task, timeout=1) == (1, "B") + + +@pytest.mark.asyncio +async def test_active_turn_handles_ask_user_question_response(): + bus = get_message_bus() + bus.mark_renderer_active() + task = asyncio.create_task( + bus.request_ask_user_question( + [ + { + "header": "Pizza-Toppings", + "question": "Which toppings?", + "multi_select": True, + "input_mode": "select_or_text", + "options": [{"label": "Pepperoni", "description": "Classic"}], + } + ] + ) + ) + prompt_id = await _wait_for_prompt_id() + + assert handle_user_interaction_response( + { + "type": "ask_user_question_response", + "prompt_id": prompt_id, + "answers": [ + { + "question_header": "Pizza-Toppings", + "selected_options": ["Pepperoni"], + "user_input": "extra cheese", + } + ], + "cancelled": False, + } + ) + + assert await asyncio.wait_for(task, timeout=1) == ( + [ + { + "question_header": "Pizza-Toppings", + "selected_options": ["Pepperoni"], + "user_input": "extra cheese", + } + ], + False, + ) + + +def test_non_interaction_message_returns_false(): + assert handle_user_interaction_response({"type": "message"}) is False diff --git a/code_puppy/api/ws/tests/test_event_driven_drain.py b/code_puppy/api/ws/tests/test_event_driven_drain.py new file mode 100644 index 000000000..ce74cc537 --- /dev/null +++ b/code_puppy/api/ws/tests/test_event_driven_drain.py @@ -0,0 +1,556 @@ +"""Tests for event-driven polling elimination in WebSocket handlers. + +Verifies that the polling-based event collection has been replaced with +event-driven asyncio.wait_for() approach. Tests validate: + +1. Events are received without time-based polling +2. Batching efficiency is maintained +3. Timeout handling works gracefully +4. Empty queue behavior is correct +5. Performance improvement (no busy-wait loops) +""" + +import asyncio +import time +from typing import Any +from unittest.mock import AsyncMock, patch + +import pytest + + +class EventDrivenCollector: + """Reference implementation of event-driven batch collection. + + This mimics the pattern used in chat_handler.py + after the polling elimination fix. + """ + + def __init__(self, event_queue: asyncio.Queue): + self.event_queue = event_queue + self.event_count = 0 + + async def collect_batch(self) -> list[dict[str, Any]]: + """Collect a batch of events using event-driven approach. + + Returns: + List of events collected (may be empty if timeout occurs) + """ + events_to_send = [] + try: + # Wait for first event with 10ms timeout (blocks efficiently) + first_event = await asyncio.wait_for(self.event_queue.get(), timeout=0.01) + events_to_send.append(first_event) + self.event_count += 1 + + # Collect any additional events already in queue (non-blocking) + # This keeps the batching benefit without polling the clock + while not self.event_queue.empty(): + try: + event = self.event_queue.get_nowait() + events_to_send.append(event) + self.event_count += 1 + except Exception: + break + + except asyncio.TimeoutError: + # No events available within 10ms timeout + # No polling needed - asyncio.wait_for blocks efficiently + pass + + return events_to_send + + +# ============================================================================ +# FIXTURES +# ============================================================================ + + +@pytest.fixture +async def event_queue(): + """Fixture providing a fresh asyncio.Queue for each test.""" + return asyncio.Queue() + + +@pytest.fixture +def collector(event_queue): + """Fixture providing an EventDrivenCollector instance.""" + return EventDrivenCollector(event_queue) + + +# ============================================================================ +# TEST: Events are received without polling +# ============================================================================ + + +@pytest.mark.asyncio +async def test_events_received_without_time_polling(): + """Verify events are received using asyncio.wait_for, not time-based polling. + + This test ensures that we've eliminated the busy-wait polling loop + that used time.monotonic() or time.time() to check elapsed time. + + The event-driven approach should: + - Block efficiently on event_queue.get() + - Use asyncio.wait_for() for timeout + - Not poll the clock in a loop + """ + queue = asyncio.Queue() + collector = EventDrivenCollector(queue) + + # Add an event to the queue + test_event = {"type": "test", "data": "hello"} + await queue.put(test_event) + + # Collect batch - should return the event immediately + start = time.perf_counter() + batch = await collector.collect_batch() + elapsed = time.perf_counter() - start + + # Verify event was received + assert len(batch) == 1 + assert batch[0] == test_event + assert collector.event_count == 1 + + # Verify it happened nearly instantly (< 5ms, not 10ms sleep) + # If it was using time.sleep(0.01), it would take ~10ms + assert elapsed < 0.005, f"Collection took {elapsed * 1000:.2f}ms (should be <5ms)" + + +@pytest.mark.asyncio +async def test_wait_for_timeout_called(event_queue, collector): + """Verify asyncio.wait_for() is being used for timeout. + + Mock asyncio.wait_for to verify it's called, confirming we're using + the efficient async timeout mechanism instead of time-based polling. + """ + with patch("asyncio.wait_for", new_callable=AsyncMock) as mock_wait_for: + # Setup mock to raise TimeoutError + mock_wait_for.side_effect = asyncio.TimeoutError() + + # Collect batch + batch = await collector.collect_batch() + + # Verify wait_for was called + mock_wait_for.assert_called_once() + call_args = mock_wait_for.call_args + + # Verify timeout parameter (should be 0.01 = 10ms) + assert call_args.kwargs.get("timeout") == 0.01, "Should use 10ms timeout" + + # Verify returned empty batch on timeout + assert len(batch) == 0 + + +@pytest.mark.asyncio +async def test_no_time_module_polling(): + """Verify time module is not used for polling. + + The old implementation used time.monotonic() or time.time() to poll + in a loop. The new implementation should not do this. + + Instead of mocking (which breaks asyncio), we verify the code doesn't + have polling loops by checking it uses wait_for() instead. + """ + queue = asyncio.Queue() + collector = EventDrivenCollector(queue) + + # Add event + await queue.put({"type": "test"}) + + # Collect batch + batch = await collector.collect_batch() + + # If we're here without exception, the pattern is working + # (No time-based polling loop would cause issues with asyncio) + assert len(batch) == 1 + assert batch[0]["type"] == "test" + + +# ============================================================================ +# TEST: Batching efficiency maintained +# ============================================================================ + + +@pytest.mark.asyncio +async def test_multiple_events_batched_together(): + """Verify multiple events arriving together are collected in one batch. + + The event-driven approach should maintain the batching benefit: + when multiple events arrive within the 10ms window, they should be + collected together in a single batch. + """ + queue = asyncio.Queue() + collector = EventDrivenCollector(queue) + + # Add multiple events to queue + events = [ + {"type": "tool_call_start", "tool_name": "tool1"}, + {"type": "tool_result", "result": "result1"}, + {"type": "stream_event", "data": "text"}, + {"type": "tool_call_complete"}, + ] + for event in events: + await queue.put(event) + + # Collect batch + batch = await collector.collect_batch() + + # All events should be in one batch + assert len(batch) == 4, "All events should be in single batch" + assert batch == events + assert collector.event_count == 4 + + +@pytest.mark.asyncio +async def test_batching_with_sequential_additions(): + """Verify batching when events are added: first immediately, then more. + + Simulates real-world scenario where: + 1. First event arrives immediately + 2. Additional events arrive quickly after + 3. All should be collected in one batch + """ + queue = asyncio.Queue() + collector = EventDrivenCollector(queue) + + # Add first event + await queue.put({"type": "event1"}) + + # Simulate additional events arriving quickly + async def add_more_events(): + await asyncio.sleep(0.002) # 2ms later + await queue.put({"type": "event2"}) + await queue.put({"type": "event3"}) + + # Start adding more events concurrently + task = asyncio.create_task(add_more_events()) + + # Collect batch (should wait 10ms, allowing time for more events) + batch = await collector.collect_batch() + await task + + # Should have collected first + any others that arrived + assert len(batch) >= 1, "Should collect at least first event" + assert batch[0]["type"] == "event1" + + +@pytest.mark.asyncio +async def test_event_count_tracking(): + """Verify event_count is incremented correctly for batched events.""" + queue = asyncio.Queue() + collector = EventDrivenCollector(queue) + + assert collector.event_count == 0 + + # First batch + await queue.put({"type": "event1"}) + await queue.put({"type": "event2"}) + batch1 = await collector.collect_batch() + assert len(batch1) == 2 + assert collector.event_count == 2 + + # Second batch (after timeout) + batch2 = await collector.collect_batch() + assert len(batch2) == 0 + assert collector.event_count == 2 # Unchanged + + +# ============================================================================ +# TEST: Timeout handling +# ============================================================================ + + +@pytest.mark.asyncio +async def test_timeout_handled_gracefully(): + """Verify TimeoutError is caught and handled gracefully. + + When no events are available within the 10ms timeout window, + the asyncio.TimeoutError should be caught and an empty batch returned. + No exceptions should bubble up. + """ + queue = asyncio.Queue() + collector = EventDrivenCollector(queue) + + # Queue is empty, so timeout will occur + batch = await collector.collect_batch() + + # Should return empty batch + assert len(batch) == 0 + assert isinstance(batch, list) + + # Should not raise any exception + assert True # If we got here without exception, test passed + + +@pytest.mark.asyncio +async def test_loop_continues_after_timeout(): + """Verify loop can continue after timeout without issues. + + Multiple collection attempts should work correctly even if + some attempts timeout. + """ + queue = asyncio.Queue() + collector = EventDrivenCollector(queue) + + # First collection - timeout + batch1 = await collector.collect_batch() + assert len(batch1) == 0 + + # Add event + await queue.put({"type": "event1"}) + + # Second collection - should get event + batch2 = await collector.collect_batch() + assert len(batch2) == 1 + assert batch2[0]["type"] == "event1" + + # Third collection - timeout again + batch3 = await collector.collect_batch() + assert len(batch3) == 0 + + +@pytest.mark.asyncio +async def test_timeout_value_correct(event_queue): + """Verify the timeout value is exactly 10ms (0.01 seconds). + + The old polling approach used a 10ms window. The new approach + should maintain this same responsiveness window. + """ + collector = EventDrivenCollector(event_queue) + + # Measure actual timeout duration + start = time.perf_counter() + _ = await collector.collect_batch() # Intentionally discarding batch + elapsed = time.perf_counter() - start + + # With no events, should timeout after ~10ms (±2ms tolerance) + assert 0.008 < elapsed < 0.015, ( + f"Timeout should be ~10ms, got {elapsed * 1000:.2f}ms" + ) + + +# ============================================================================ +# TEST: Empty queue behavior +# ============================================================================ + + +@pytest.mark.asyncio +async def test_empty_queue_returns_empty_batch(): + """Verify empty queue results in empty batch. + + When queue is empty and timeout occurs, should return empty list, + not raise an exception. + """ + queue = asyncio.Queue() + collector = EventDrivenCollector(queue) + batch = await collector.collect_batch() + assert batch == [] + + +@pytest.mark.asyncio +async def test_no_infinite_loop_on_empty_queue(): + """Verify no infinite loop when queue is empty. + + The timeout should prevent infinite looping. Collection should + complete in ~10ms even with empty queue. + """ + queue = asyncio.Queue() + collector = EventDrivenCollector(queue) + + start = time.perf_counter() + batch = await collector.collect_batch() + elapsed = time.perf_counter() - start + + # Should complete quickly (within timeout + small overhead) + assert elapsed < 0.05, f"Should not hang, took {elapsed * 1000:.2f}ms" + assert batch == [] + + +@pytest.mark.asyncio +async def test_queue_empty_check_works(): + """Verify queue.empty() check works correctly for secondary collection. + + After getting first event, the code checks event_queue.empty() to + collect additional events. This should work correctly. + """ + queue = asyncio.Queue() + collector = EventDrivenCollector(queue) + + # Add multiple events + await queue.put({"type": "event1"}) + await queue.put({"type": "event2"}) + await queue.put({"type": "event3"}) + + batch = await collector.collect_batch() + + # All events should be collected despite using empty() check + assert len(batch) == 3 + assert collector.event_count == 3 + + # Queue should now be empty + assert queue.empty() + + +# ============================================================================ +# TEST: Performance improvement (no busy-wait loops) +# ============================================================================ + + +@pytest.mark.asyncio +async def test_performance_event_arrival_near_instant(): + """Verify event collection happens near-instantly (<1ms). + + The old polling approach would wait up to 10ms before checking + the queue. Event-driven approach wakes immediately on event arrival. + """ + queue = asyncio.Queue() + collector = EventDrivenCollector(queue) + + # Add event then immediately collect + await queue.put({"type": "test"}) + + start = time.perf_counter() + batch = await collector.collect_batch() + elapsed = time.perf_counter() - start + + # Should be near-instant (<1ms, not ~10ms) + assert elapsed < 0.001, ( + f"Event should be collected near-instantly, " + f"took {elapsed * 1000:.2f}ms (should be <1ms)" + ) + assert len(batch) == 1 + + +@pytest.mark.asyncio +async def test_no_sleep_on_every_iteration(): + """Verify we don't sleep 10ms on every loop iteration. + + The old implementation would do: if not events: await asyncio.sleep(0.01) + This test verifies that pattern is eliminated. + """ + queue = asyncio.Queue() + collector = EventDrivenCollector(queue) + + # Track asyncio.sleep calls + with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: + # Run collection multiple times + for _ in range(3): + await collector.collect_batch() + + # Should NOT call asyncio.sleep(0.01) on idle + # (It might be called elsewhere, but not as idle sleep) + sleep_calls = [call for call in mock_sleep.call_args_list] + idle_sleep_calls = [call for call in sleep_calls if call[0] == (0.01,)] + assert len(idle_sleep_calls) == 0, ( + "Should not sleep 0.01s on every idle iteration" + ) + + +@pytest.mark.asyncio +async def test_multiple_events_processed_once(): + """Verify multiple events are processed in one batch, not multiple sleeps. + + Old approach: each iteration could sleep 10ms if no events + New approach: waits 10ms once, collects all available events + """ + queue = asyncio.Queue() + collector = EventDrivenCollector(queue) + + # Add events + for i in range(5): + await queue.put({"type": "event", "id": i}) + + start = time.perf_counter() + batch = await collector.collect_batch() + elapsed = time.perf_counter() - start + + # All events collected in one batch + assert len(batch) == 5 + # Should be instant, not 50ms (5 iterations * 10ms sleep) + assert elapsed < 0.005 + + +@pytest.mark.asyncio +async def test_batching_efficiency_metric(): + """Calculate and verify batching efficiency. + + Batching efficiency = total_events / number_of_batches + + Event-driven approach should achieve good batching even with + sequential additions, because it waits up to 10ms for additional events. + """ + queue = asyncio.Queue() + collector = EventDrivenCollector(queue) + + batch_sizes = [] + + # Simulate 3 collection cycles + for cycle in range(3): + # Add varying number of events + num_events = cycle + 1 # 1, 2, 3 + for i in range(num_events): + await queue.put({"type": "event", "cycle": cycle, "id": i}) + + batch = await collector.collect_batch() + batch_sizes.append(len(batch)) + + # Verify we collected events (not all empty) + total_collected = sum(batch_sizes) + assert total_collected == 6, "Should collect all 6 events (1+2+3)" + + # Calculate efficiency: 6 events / 3 batches = 2.0 average + avg_batch_size = total_collected / 3 + assert avg_batch_size >= 1.5, "Should have reasonable batching" + + +# ============================================================================ +# INTEGRATION: Full drain loop simulation +# ============================================================================ + + +@pytest.mark.asyncio +async def test_drain_loop_simulation(): + """Simulate a full drain loop with event-driven collection. + + This mimics the actual usage in chat_handler.py: + - Collect events in a loop + - Process collected events + - Continue until stop signal + """ + queue = asyncio.Queue() + collector = EventDrivenCollector(queue) + stop_event = asyncio.Event() + processed_events = [] + + async def drain_loop(): + """Simulate the drain event loop.""" + while not stop_event.is_set(): + batch = await collector.collect_batch() + for event in batch: + processed_events.append(event) + + # Add small delay to allow test to add events + if not batch: + await asyncio.sleep(0.001) + + # Start drain loop + drain_task = asyncio.create_task(drain_loop()) + + # Simulate events being added + await asyncio.sleep(0.01) + for i in range(5): + await queue.put({"type": "event", "id": i}) + + await asyncio.sleep(0.05) # Let drain collect them + + # Stop the loop + stop_event.set() + await asyncio.wait_for(drain_task, timeout=1.0) + + # Verify all events were processed + assert len(processed_events) == 5 + assert all(e["type"] == "event" for e in processed_events) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "--tb=short"]) diff --git a/code_puppy/api/ws/tests/test_event_driven_polling.py b/code_puppy/api/ws/tests/test_event_driven_polling.py new file mode 100644 index 000000000..13b90799d --- /dev/null +++ b/code_puppy/api/ws/tests/test_event_driven_polling.py @@ -0,0 +1,337 @@ +"""Tests for event-driven polling elimination (Phases 1-2 fixes). + +Verifies that: +1. Events are collected without polling loops +2. Batching still works efficiently +3. Timeout handling works gracefully +4. Event ordering is preserved +5. No regressions in event processing +""" + +import asyncio + +import pytest + + +@pytest.mark.asyncio +async def test_event_driven_collection_no_polling(): + """Verify events are collected using asyncio.wait_for, not polling loops. + + This test ensures that the event-driven approach (Phase 1 & 2) correctly + uses asyncio.wait_for() instead of busy-waiting with time.monotonic(). + """ + # Create a queue and emit events + event_queue: asyncio.Queue = asyncio.Queue() + + # Add events to the queue + await event_queue.put({"type": "content", "data": {"text": "Hello"}}) + await event_queue.put({"type": "content", "data": {"text": " World"}}) + + # Simulate the event-driven collection pattern from fixed handlers + events_to_send = [] + try: + # Wait for first event with 10ms timeout (event-driven, not polling) + first_event = await asyncio.wait_for(event_queue.get(), timeout=0.01) + events_to_send.append(first_event) + + # Collect any additional events already queued (non-blocking) + while not event_queue.empty(): + try: + events_to_send.append(event_queue.get_nowait()) + except Exception: + break + + except asyncio.TimeoutError: + pass + + # Verify both events were collected + assert len(events_to_send) == 2 + assert events_to_send[0]["data"]["text"] == "Hello" + assert events_to_send[1]["data"]["text"] == " World" + + +@pytest.mark.asyncio +async def test_event_driven_timeout_handling(): + """Verify timeout handling when no events are available. + + This test ensures asyncio.TimeoutError is properly caught + and doesn't cause the handler to crash. + """ + # Empty queue + event_queue: asyncio.Queue = asyncio.Queue() + + # Simulate event-driven collection with very short timeout + events_to_send = [] + try: + first_event = await asyncio.wait_for(event_queue.get(), timeout=0.001) + events_to_send.append(first_event) + except asyncio.TimeoutError: + # Expected - no events in queue + pass + + # No events should be collected + assert len(events_to_send) == 0 + + +@pytest.mark.asyncio +async def test_batching_preserved_with_event_driven(): + """Verify batching efficiency is preserved in event-driven approach. + + The 10ms batching window should still collect multiple events + without introducing polling overhead. + """ + event_queue: asyncio.Queue = asyncio.Queue() + + # Emit multiple events rapidly + for i in range(5): + await event_queue.put({"type": "event", "data": {"index": i}}) + + # Collect all events in one batch using event-driven approach + events_to_send = [] + try: + first_event = await asyncio.wait_for(event_queue.get(), timeout=0.1) + events_to_send.append(first_event) + + while not event_queue.empty(): + try: + events_to_send.append(event_queue.get_nowait()) + except Exception: + break + + except asyncio.TimeoutError: + pass + + # All 5 events should be collected in a single batch + assert len(events_to_send) == 5 + for i in range(5): + assert events_to_send[i]["data"]["index"] == i + + +@pytest.mark.asyncio +async def test_event_ordering_preserved(): + """Verify event ordering is preserved in batches. + + Events should be processed in the order they were emitted. + """ + event_queue: asyncio.Queue = asyncio.Queue() + + # Emit events with order markers + for i in range(10): + await event_queue.put( + {"type": "ordered_event", "data": {"sequence": i, "timestamp": i * 1000}} + ) + + # Collect events + events_to_send = [] + try: + first_event = await asyncio.wait_for(event_queue.get(), timeout=0.1) + events_to_send.append(first_event) + + while not event_queue.empty(): + try: + events_to_send.append(event_queue.get_nowait()) + except Exception: + break + + except asyncio.TimeoutError: + pass + + # Verify all events collected and in order + assert len(events_to_send) == 10 + for i, event in enumerate(events_to_send): + assert event["data"]["sequence"] == i + assert event["data"]["timestamp"] == i * 1000 + + +@pytest.mark.asyncio +async def test_interleaved_event_production(): + """Test handling of events added during collection. + + Simulates realistic scenario where events arrive both before + and after timeout in the wait_for() call. + """ + event_queue: asyncio.Queue = asyncio.Queue() + + # Add initial events + await event_queue.put({"type": "pre", "data": {"phase": "before"}}) + + async def add_event_later(): + """Add event after a small delay.""" + await asyncio.sleep(0.005) + await event_queue.put({"type": "post", "data": {"phase": "after"}}) + + # Start background task to add event during collection + task = asyncio.create_task(add_event_later()) + + # Collect events with longer timeout + events_to_send = [] + try: + first_event = await asyncio.wait_for(event_queue.get(), timeout=0.05) + events_to_send.append(first_event) + + while not event_queue.empty(): + try: + events_to_send.append(event_queue.get_nowait()) + except Exception: + break + + except asyncio.TimeoutError: + pass + + await task + + # Should have at least the first event + assert len(events_to_send) >= 1 + assert events_to_send[0]["data"]["phase"] == "before" + + +@pytest.mark.asyncio +async def test_no_cpu_waste_on_idle(): + """Verify that idle timeouts don't cause busy-waiting. + + The old polling approach used await asyncio.sleep(0.01) which still + wasted CPU. The new approach uses asyncio.wait_for() which properly + suspends the task without CPU usage. + """ + event_queue: asyncio.Queue = asyncio.Queue() + + # Record timing of timeout + start_time = asyncio.get_event_loop().time() + + events_to_send = [] + try: + # This should block efficiently for ~10ms, not busy-wait + first_event = await asyncio.wait_for(event_queue.get(), timeout=0.01) + events_to_send.append(first_event) + except asyncio.TimeoutError: + pass + + elapsed = asyncio.get_event_loop().time() - start_time + + # Should take approximately 10ms, not significantly less or more + # (we allow some variance for system scheduling) + assert 0.008 < elapsed < 0.05, f"Timeout took {elapsed}s, expected ~0.01s" + assert len(events_to_send) == 0 + + +@pytest.mark.asyncio +async def test_rapid_fire_events_single_batch(): + """Test that rapidly fired events are collected in a single batch. + + This is key to the efficiency gains - multiple events should be + processed together without individual polling cycles. + """ + event_queue: asyncio.Queue = asyncio.Queue() + + # Emit 100 events rapidly + for i in range(100): + await event_queue.put({"type": "rapid", "data": {"id": i}}) + + # Collect all in single cycle + events_to_send = [] + try: + first_event = await asyncio.wait_for(event_queue.get(), timeout=0.1) + events_to_send.append(first_event) + + while not event_queue.empty(): + try: + events_to_send.append(event_queue.get_nowait()) + except Exception: + break + + except asyncio.TimeoutError: + pass + + # All events should be collected + assert len(events_to_send) == 100 + # Verify we got them in order + for i in range(100): + assert events_to_send[i]["data"]["id"] == i + + +@pytest.mark.asyncio +async def test_mixed_event_types_preserved(): + """Test that different event types are preserved in batches. + + Various event types (tool_call, content, stream_event, etc.) should + all be collected and processed correctly. + """ + event_queue: asyncio.Queue = asyncio.Queue() + + # Emit mixed event types + event_types = [ + {"type": "tool_call_start", "data": {"tool": "search"}}, + {"type": "content", "data": {"text": "Query..."}}, + {"type": "tool_call_complete", "data": {"result": "found"}}, + {"type": "stream_event", "data": {"event_type": "part_delta"}}, + {"type": "message_complete", "data": {"final": True}}, + ] + + for event in event_types: + await event_queue.put(event) + + # Collect all events + events_to_send = [] + try: + first_event = await asyncio.wait_for(event_queue.get(), timeout=0.1) + events_to_send.append(first_event) + + while not event_queue.empty(): + try: + events_to_send.append(event_queue.get_nowait()) + except Exception: + break + + except asyncio.TimeoutError: + pass + + # All events collected + assert len(events_to_send) == 5 + # Verify types match in order + for i, expected in enumerate(event_types): + assert events_to_send[i]["type"] == expected["type"] + + +@pytest.mark.asyncio +async def test_timeout_with_partial_events(): + """Test partial collection when timeout occurs mid-batch. + + If events come in after wait_for timeout, the handler should + gracefully handle the empty batch and try again. + """ + event_queue: asyncio.Queue = asyncio.Queue() + + # First collection - empty + events_batch_1 = [] + try: + first_event = await asyncio.wait_for(event_queue.get(), timeout=0.005) + events_batch_1.append(first_event) + except asyncio.TimeoutError: + pass + + # First batch should be empty + assert len(events_batch_1) == 0 + + # Add events + await event_queue.put({"type": "event1", "data": {}}) + await event_queue.put({"type": "event2", "data": {}}) + + # Second collection - should get both + events_batch_2 = [] + try: + first_event = await asyncio.wait_for(event_queue.get(), timeout=0.05) + events_batch_2.append(first_event) + + while not event_queue.empty(): + try: + events_batch_2.append(event_queue.get_nowait()) + except Exception: + break + + except asyncio.TimeoutError: + pass + + # Second batch should have both events + assert len(events_batch_2) == 2 + assert events_batch_2[0]["type"] == "event1" + assert events_batch_2[1]["type"] == "event2" diff --git a/code_puppy/api/ws/tests/test_session_persistence.py b/code_puppy/api/ws/tests/test_session_persistence.py new file mode 100644 index 000000000..f2c2f2659 --- /dev/null +++ b/code_puppy/api/ws/tests/test_session_persistence.py @@ -0,0 +1,197 @@ +from __future__ import annotations + +import datetime +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from code_puppy.api.ws.session_persistence import persist_session_turn_and_broadcast + + +class _Logger: + def __init__(self): + self.debug_messages = [] + + def debug(self, message, *args): + self.debug_messages.append(message % args if args else message) + + +@pytest.mark.asyncio +async def test_persist_session_turn_and_broadcast_returns_none_for_empty_history(): + safe_send_json = AsyncMock() + + with ( + patch( + "code_puppy.api.ws.session_persistence.persist_turn_to_sqlite", + new_callable=AsyncMock, + ) as mock_persist, + patch( + "code_puppy.api.ws.session_persistence.connection_manager.broadcast_session_update", + new_callable=AsyncMock, + ) as mock_broadcast, + ): + result = await persist_session_turn_and_broadcast( + history=[], + session_id="sess-1", + session_title="", + session_working_directory="/tmp", + session_pinned=False, + agent=MagicMock(), + agent_name="code-puppy", + model_name="gpt-4", + ctx=SimpleNamespace(created_at=datetime.datetime(2025, 1, 1)), + original_user_message="hello", + attachment_metadata=None, + safe_send_json=safe_send_json, + ) + + assert result is None + safe_send_json.assert_not_called() + mock_persist.assert_not_called() + mock_broadcast.assert_not_called() + + +@pytest.mark.asyncio +async def test_persist_session_turn_and_broadcast_persists_and_notifies_with_generated_title(): + history = [object(), object()] + enhanced_history = [{"msg": "u"}, {"msg": "a"}] + safe_send_json = AsyncMock() + logger = _Logger() + ctx = SimpleNamespace( + created_at=datetime.datetime(2025, 1, 2, 3, 4, 5), + agent_name="ctx-agent", + model_name="ctx-model", + ) + agent = MagicMock() + agent.name = "agent-object" + agent.get_model_name.return_value = "agent-model" + + with ( + patch( + "code_puppy.api.ws.session_persistence.generate_heuristic_title", + return_value="Generated Title", + ) as mock_title, + patch( + "code_puppy.api.ws.session_persistence.build_enhanced_history", + return_value=enhanced_history, + ) as mock_build, + patch( + "code_puppy.api.ws.session_persistence.estimate_total_tokens", + return_value=42, + ) as mock_tokens, + patch( + "code_puppy.api.ws.session_persistence.persist_turn_to_sqlite", + new_callable=AsyncMock, + ) as mock_persist, + patch( + "code_puppy.api.ws.session_persistence.connection_manager.broadcast_session_update", + new_callable=AsyncMock, + ) as mock_broadcast, + ): + summary = await persist_session_turn_and_broadcast( + history=history, + session_id="sess-2", + session_title="untitled-session", + session_working_directory="/work", + session_pinned=True, + agent=agent, + agent_name="wire-agent", + model_name="wire-model", + ctx=ctx, + original_user_message="show logs", + attachment_metadata=[{"name": "demo.txt"}], + safe_send_json=safe_send_json, + logger_override=logger, + ) + + assert summary is not None + assert summary.session_title == "Generated Title" + assert summary.message_count == 2 + assert summary.total_tokens == 42 + + mock_title.assert_called_once_with(history) + mock_build.assert_called_once_with( + history, + agent_name_meta="agent-object", + model_name_meta="agent-model", + original_user_message="show logs", + attachment_metadata=[{"name": "demo.txt"}], + ) + mock_tokens.assert_called_once_with(enhanced_history, agent) + mock_persist.assert_awaited_once_with( + session_id="sess-2", + enhanced_history=enhanced_history, + title="Generated Title", + working_directory="/work", + pinned=True, + agent_name="wire-agent", + model_name="wire-model", + total_tokens=42, + created_at_iso=ctx.created_at.isoformat(), + ctx=ctx, + ) + safe_send_json.assert_awaited_once() + session_meta_payload = safe_send_json.await_args.args[0] + assert session_meta_payload["type"] == "session_meta" + assert session_meta_payload["session_id"] == "sess-2" + assert session_meta_payload["title"] == "Generated Title" + assert session_meta_payload["total_tokens"] == 42 + assert session_meta_payload["message_count"] == 2 + assert session_meta_payload["agent_name"] == "wire-agent" + assert session_meta_payload["model_name"] == "wire-model" + + mock_broadcast.assert_awaited_once() + session_update_payload = mock_broadcast.await_args.args[0] + assert session_update_payload["session_id"] == "sess-2" + assert session_update_payload["title"] == "Generated Title" + assert session_update_payload["message_count"] == 2 + assert session_update_payload["total_tokens"] == 42 + assert logger.debug_messages + assert "Added UI metadata" in logger.debug_messages[0] + + +@pytest.mark.asyncio +async def test_persist_session_turn_and_broadcast_preserves_existing_title(): + safe_send_json = AsyncMock() + ctx = SimpleNamespace(created_at=datetime.datetime(2025, 1, 1)) + + with ( + patch( + "code_puppy.api.ws.session_persistence.generate_heuristic_title" + ) as mock_title, + patch( + "code_puppy.api.ws.session_persistence.build_enhanced_history", + return_value=[{"msg": "only"}], + ), + patch( + "code_puppy.api.ws.session_persistence.estimate_total_tokens", + return_value=9, + ), + patch( + "code_puppy.api.ws.session_persistence.persist_turn_to_sqlite", + new_callable=AsyncMock, + ), + patch( + "code_puppy.api.ws.session_persistence.connection_manager.broadcast_session_update", + new_callable=AsyncMock, + ), + ): + summary = await persist_session_turn_and_broadcast( + history=[object()], + session_id="sess-3", + session_title="Keep Me", + session_working_directory="/tmp", + session_pinned=False, + agent=MagicMock(), + agent_name="a", + model_name="m", + ctx=ctx, + original_user_message="hi", + attachment_metadata=[], + safe_send_json=safe_send_json, + ) + + assert summary is not None + assert summary.session_title == "Keep Me" + mock_title.assert_not_called() diff --git a/code_puppy/api/ws/tests/test_tool_group_id_consistency.py b/code_puppy/api/ws/tests/test_tool_group_id_consistency.py new file mode 100644 index 000000000..2d1170ee7 --- /dev/null +++ b/code_puppy/api/ws/tests/test_tool_group_id_consistency.py @@ -0,0 +1,80 @@ +"""Regression tests for tool_group_id consistency in websocket lifecycle frames. + +These tests enforce the backend invariant introduced by Option C: +all emitted ``ServerToolResult`` payloads must include ``tool_group_id``. +After modularization, the constructors may live outside ``chat_handler.py``, so +we scan the known WebSocket emitter modules. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +WS_DIR = Path(__file__).resolve().parents[1] +EMITTER_PATHS = [ + WS_DIR / "chat_handler.py", + WS_DIR / "chat_tool_lifecycle.py", + WS_DIR / "ws_turn_finalization.py", +] + + +def _get_server_tool_result_calls(source: str) -> list[ast.Call]: + """Return every ``ServerToolResult(...)`` call in source.""" + tree = ast.parse(source) + calls: list[ast.Call] = [] + + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + if isinstance(node.func, ast.Name) and node.func.id == "ServerToolResult": + calls.append(node) + + return calls + + +def test_server_tool_result_calls_always_include_tool_group_id_keyword() -> None: + """Every ServerToolResult constructor call must include tool_group_id. + + This protects against regressions where a new emission path forgets to + include the grouping field and breaks frontend tool-message grouping. + """ + missing: list[str] = [] + total_calls = 0 + + for source_path in EMITTER_PATHS: + source = source_path.read_text(encoding="utf-8") + calls = _get_server_tool_result_calls(source) + total_calls += len(calls) + for call in calls: + keyword_names = {kw.arg for kw in call.keywords if kw.arg is not None} + if "tool_group_id" not in keyword_names: + missing.append(f"{source_path.name}:{call.lineno}") + + assert total_calls, "Expected at least one ServerToolResult(...) call" + assert not missing, ( + f"ServerToolResult(...) calls missing tool_group_id keyword at: {missing}" + ) + + +def test_server_tool_result_calls_never_pass_explicit_none_for_tool_group_id() -> None: + """Guard against explicit ``tool_group_id=None`` regressions.""" + explicit_none_locations: list[str] = [] + + for source_path in EMITTER_PATHS: + source = source_path.read_text(encoding="utf-8") + calls = _get_server_tool_result_calls(source) + for call in calls: + for keyword in call.keywords: + if keyword.arg != "tool_group_id": + continue + if ( + isinstance(keyword.value, ast.Constant) + and keyword.value.value is None + ): + explicit_none_locations.append(f"{source_path.name}:{call.lineno}") + + assert not explicit_none_locations, ( + "ServerToolResult(...) calls pass explicit tool_group_id=None at: " + f"{explicit_none_locations}" + ) diff --git a/code_puppy/api/ws/tests/test_ws_command_handler.py b/code_puppy/api/ws/tests/test_ws_command_handler.py new file mode 100644 index 000000000..56753a706 --- /dev/null +++ b/code_puppy/api/ws/tests/test_ws_command_handler.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +import sys +from types import ModuleType + +import pytest + +from code_puppy.api.ws.ws_command_handler import handle_command_message + + +def _install_fake_command_handler( + monkeypatch, *, help_text="Available commands", handle_result=True +): + fake_module = ModuleType("code_puppy.command_line.command_handler") + fake_module.get_commands_help = lambda: help_text + fake_module.handle_command = lambda command: ( + handle_result(command) if callable(handle_result) else handle_result + ) + monkeypatch.setitem( + sys.modules, "code_puppy.command_line.command_handler", fake_module + ) + + +@pytest.mark.asyncio +async def test_handle_command_message_help_command(monkeypatch): + sent = [] + _install_fake_command_handler(monkeypatch, help_text="usage text") + + class _FakeConsole: + def __init__(self, *args, file=None, **_kwargs): + self.file = file + + def print(self, value): + self.file.write(f"plain:{value}") + + monkeypatch.setattr("code_puppy.api.ws.ws_command_handler.Console", _FakeConsole) + + async def _send_typed(message): + sent.append(message) + + handled = await handle_command_message( + msg={"type": "command", "command": "/help"}, + session_id="session-1", + send_typed=_send_typed, + ) + + assert handled is True + assert len(sent) == 1 + assert sent[0].command == "/help" + assert sent[0].success is True + assert sent[0].output == "plain:usage text" + + +@pytest.mark.asyncio +async def test_handle_command_message_runs_generic_command(monkeypatch): + sent = [] + _install_fake_command_handler( + monkeypatch, handle_result=lambda command: f"ran:{command}" + ) + + async def _send_typed(message): + sent.append(message) + + handled = await handle_command_message( + msg={"type": "command", "command": "/agents"}, + session_id="session-2", + send_typed=_send_typed, + ) + + assert handled is True + assert len(sent) == 1 + assert sent[0].command == "/agents" + assert sent[0].success is True + assert sent[0].output == "ran:/agents" + + +@pytest.mark.asyncio +async def test_handle_command_message_ignores_non_command_messages(): + sent = [] + + async def _send_typed(message): + sent.append(message) + + handled = await handle_command_message( + msg={"type": "message", "content": "hello"}, + session_id="session-3", + send_typed=_send_typed, + ) + + assert handled is False + assert sent == [] + + +@pytest.mark.asyncio +async def test_handle_command_message_reports_errors(monkeypatch): + sent = [] + + def _boom(_command): + raise RuntimeError("boom") + + _install_fake_command_handler(monkeypatch, handle_result=_boom) + + async def _send_typed(message): + sent.append(message) + + handled = await handle_command_message( + msg={"type": "command", "command": "/broken"}, + session_id="session-4", + send_typed=_send_typed, + ) + + assert handled is True + assert len(sent) == 1 + assert sent[0].command == "/broken" + assert sent[0].success is False + assert sent[0].error == "boom" + + +@pytest.mark.asyncio +async def test_handle_command_message_treats_none_result_as_failure(monkeypatch): + sent = [] + _install_fake_command_handler(monkeypatch, handle_result=lambda _command: None) + + async def _send_typed(message): + sent.append(message) + + handled = await handle_command_message( + msg={"type": "command", "command": "/noop"}, + session_id="session-5", + send_typed=_send_typed, + ) + + assert handled is True + assert len(sent) == 1 + assert sent[0].command == "/noop" + assert sent[0].success is False diff --git a/code_puppy/api/ws/tests/test_ws_control_messages.py b/code_puppy/api/ws/tests/test_ws_control_messages.py new file mode 100644 index 000000000..d8d004ad8 --- /dev/null +++ b/code_puppy/api/ws/tests/test_ws_control_messages.py @@ -0,0 +1,264 @@ +from __future__ import annotations + +import importlib +import sys +from types import ModuleType, SimpleNamespace + +import pytest + +from code_puppy.api.ws.ws_chat_runtime import WebSocketChatRuntime + + +class _FakeSessionManager: + def __init__(self): + self.switch_agent_calls = [] + self.switch_model_calls = [] + + async def switch_agent(self, session_id, agent_name): + self.switch_agent_calls.append((session_id, agent_name)) + return SimpleNamespace(get_model_name=lambda: "new-model") + + async def switch_model(self, session_id, model_name): + self.switch_model_calls.append((session_id, model_name)) + + +def _load_control_module( + monkeypatch, + *, + cancel_calls=None, + permission_calls=None, + writes=None, + update_calls=None, + set_config_calls=None, + get_values=None, +): + fake_db_pkg = ModuleType("code_puppy.api.db") + fake_db_pkg.__path__ = [] + fake_queries = ModuleType("code_puppy.api.db.queries") + + async def _write_system_message_to_sqlite(**kwargs): + writes.append(kwargs) + + async def _update_session_working_directory(**kwargs): + update_calls.append(kwargs) + + fake_queries.write_system_message_to_sqlite = _write_system_message_to_sqlite + fake_queries.update_session_working_directory = _update_session_working_directory + fake_db_pkg.queries = fake_queries + monkeypatch.setitem(sys.modules, "code_puppy.api.db", fake_db_pkg) + monkeypatch.setitem(sys.modules, "code_puppy.api.db.queries", fake_queries) + + fake_permissions = ModuleType("code_puppy.api.permissions") + fake_permissions.handle_permission_response = ( + lambda request_id, approved, session_id=None: ( + permission_calls.append((request_id, approved, session_id)) or True + ) + ) + monkeypatch.setitem(sys.modules, "code_puppy.api.permissions", fake_permissions) + + fake_session_context = ModuleType("code_puppy.api.session_context") + fake_session_context._validate_session_id = lambda value: value + fake_session_context.session_manager = _FakeSessionManager() + monkeypatch.setitem( + sys.modules, "code_puppy.api.session_context", fake_session_context + ) + + fake_send_utils = ModuleType("code_puppy.api.ws.send_utils") + fake_send_utils.WebSocketSender = object + monkeypatch.setitem(sys.modules, "code_puppy.api.ws.send_utils", fake_send_utils) + + fake_stream_drain = ModuleType("code_puppy.api.ws.ws_stream_drain") + + async def _cancel_active_streaming(**kwargs): + cancel_calls.append(kwargs) + + fake_stream_drain.cancel_active_streaming = _cancel_active_streaming + monkeypatch.setitem( + sys.modules, "code_puppy.api.ws.ws_stream_drain", fake_stream_drain + ) + + fake_config = ModuleType("code_puppy.config") + fake_config.get_puppy_name = lambda: "puppy" + fake_config.get_value = lambda key: get_values.get(key) + fake_config.set_config_value = lambda key, value: set_config_calls.append( + (key, value) + ) + monkeypatch.setitem(sys.modules, "code_puppy.config", fake_config) + + sys.modules.pop("code_puppy.api.ws.ws_control_messages", None) + return importlib.import_module("code_puppy.api.ws.ws_control_messages") + + +@pytest.mark.asyncio +async def test_handle_control_message_switch_model(monkeypatch): + cancel_calls, permission_calls, writes, update_calls, set_config_calls = ( + [], + [], + [], + [], + [], + ) + module = _load_control_module( + monkeypatch, + cancel_calls=cancel_calls, + permission_calls=permission_calls, + writes=writes, + update_calls=update_calls, + set_config_calls=set_config_calls, + get_values={}, + ) + runtime = WebSocketChatRuntime( + session_id="session-1", + ctx=SimpleNamespace( + agent=object(), agent_name="code-puppy", model_name="old-model" + ), + agent_name="code-puppy", + model_name="old-model", + ) + sent = [] + + async def _send_typed(message): + sent.append(message) + + handled = await module.handle_control_message( + msg={"type": "switch_model", "model_name": "new-model"}, + runtime=runtime, + sender=SimpleNamespace(session_id="session-1", ctx=runtime.ctx), + send_typed=_send_typed, + send_session_meta_snapshot=lambda: None, + ) + + assert handled is True + assert runtime.model_name == "new-model" + assert sent[-1].type == "system" + assert "new-model" in sent[-1].content + assert writes[-1]["model_name"] == "new-model" + + +@pytest.mark.asyncio +async def test_handle_control_message_set_working_directory(monkeypatch, tmp_path): + cancel_calls, permission_calls, writes, update_calls, set_config_calls = ( + [], + [], + [], + [], + [], + ) + module = _load_control_module( + monkeypatch, + cancel_calls=cancel_calls, + permission_calls=permission_calls, + writes=writes, + update_calls=update_calls, + set_config_calls=set_config_calls, + get_values={}, + ) + runtime = WebSocketChatRuntime( + session_id="session-2", + ctx=SimpleNamespace( + agent_name="code-puppy", model_name="gpt-test", working_directory="" + ), + ) + sent = [] + + async def _send_typed(message): + sent.append(message) + + handled = await module.handle_control_message( + msg={"type": "set_working_directory", "directory": str(tmp_path)}, + runtime=runtime, + sender=SimpleNamespace(session_id="session-2", ctx=runtime.ctx), + send_typed=_send_typed, + send_session_meta_snapshot=lambda: None, + ) + + assert handled is True + assert runtime.session_working_directory == str(tmp_path.resolve()) + assert sent[-1].type == "working_directory_changed" + assert sent[-1].success is True + assert writes[-1]["system_message_type"] == "directory" + assert update_calls[-1]["working_directory"] == str(tmp_path.resolve()) + + +@pytest.mark.asyncio +async def test_handle_control_message_cancel(monkeypatch): + cancel_calls, permission_calls, writes, update_calls, set_config_calls = ( + [], + [], + [], + [], + [], + ) + module = _load_control_module( + monkeypatch, + cancel_calls=cancel_calls, + permission_calls=permission_calls, + writes=writes, + update_calls=update_calls, + set_config_calls=set_config_calls, + get_values={}, + ) + + task = None + + async def _never(): + await __import__("asyncio").Future() + + import asyncio + + task = asyncio.create_task(_never()) + runtime = WebSocketChatRuntime(session_id="session-3", active_agent_task=task) + sent = [] + + async def _send_typed(message): + sent.append(message) + + handled = await module.handle_control_message( + msg={"type": "cancel"}, + runtime=runtime, + sender=SimpleNamespace(session_id="session-3", ctx=None), + send_typed=_send_typed, + send_session_meta_snapshot=lambda: None, + ) + + assert handled is True + assert cancel_calls + assert runtime.active_agent_task is None + assert sent[-1].type == "status" + assert sent[-1].status == "cancelled" + + +@pytest.mark.asyncio +async def test_handle_control_message_permission_response(monkeypatch): + cancel_calls, permission_calls, writes, update_calls, set_config_calls = ( + [], + [], + [], + [], + [], + ) + module = _load_control_module( + monkeypatch, + cancel_calls=cancel_calls, + permission_calls=permission_calls, + writes=writes, + update_calls=update_calls, + set_config_calls=set_config_calls, + get_values={}, + ) + sent = [] + + async def _send_typed(message): + sent.append(message) + + handled = await module.handle_control_message( + msg={"type": "permission_response", "request_id": "req-1", "approved": True}, + runtime=WebSocketChatRuntime(session_id="session-4"), + sender=SimpleNamespace(session_id="session-4", ctx=None), + send_typed=_send_typed, + send_session_meta_snapshot=lambda: None, + ) + + assert handled is True + assert permission_calls == [("req-1", True, "session-4")] + assert sent == [] diff --git a/code_puppy/api/ws/tests/test_ws_post_run.py b/code_puppy/api/ws/tests/test_ws_post_run.py new file mode 100644 index 000000000..1a85358da --- /dev/null +++ b/code_puppy/api/ws/tests/test_ws_post_run.py @@ -0,0 +1,149 @@ +from __future__ import annotations + +from types import SimpleNamespace + +from code_puppy.api.ws.chat_turn_state import WebSocketTurnState +from code_puppy.api.ws.ws_post_run import resolve_post_run_resolution + + +class _Logger: + def __init__(self): + self.debug_messages = [] + self.warning_messages = [] + + def debug(self, message, *args): + self.debug_messages.append(message % args if args else message) + + def warning(self, message, *args): + self.warning_messages.append(message % args if args else message) + + +class _AssistantMessage: + role = "assistant" + + def __init__(self, content: str): + self.content = content + + +class _ThinkingPart: + def __init__(self, content: str): + self.content = content + + +class _ModelResponse: + def __init__(self, parts): + self.parts = parts + + +class _Agent: + def __init__(self, history=None, estimate=7): + self._history = history or [] + self._estimate = estimate + + def get_message_history(self): + return list(self._history) + + def estimate_tokens_for_message(self, _msg): + return self._estimate + + +def test_resolve_post_run_resolution_cancelled(): + state = WebSocketTurnState() + state.agent_error = "cancelled" + + resolution = resolve_post_run_resolution( + result=None, + turn_state=state, + agent=None, + session_id="s1", + logger=_Logger(), + ) + + assert resolution.cancelled is True + assert resolution.error_frames is None + assert resolution.no_result_error is None + + +def test_resolve_post_run_resolution_builds_error_frames_for_agent_exception(): + state = WebSocketTurnState(collected_text=["partial output"]) + state.agent_error = RuntimeError("boom") + + resolution = resolve_post_run_resolution( + result=None, + turn_state=state, + agent=None, + session_id="s2", + logger=_Logger(), + ) + + assert resolution.cancelled is False + assert resolution.error_frames is not None + assert resolution.error_frames[0]["type"] == "stream_end" + assert resolution.error_frames[0]["success"] is False + assert resolution.error_frames[1]["type"] == "error" + + +def test_resolve_post_run_resolution_handles_no_result_without_stream(): + state = WebSocketTurnState() + + resolution = resolve_post_run_resolution( + result=None, + turn_state=state, + agent=None, + session_id="s3", + logger=_Logger(), + ) + + assert resolution.no_result_error is not None + assert resolution.no_result_error.type == "error" + assert resolution.no_result_error.session_id == "s3" + assert "no result returned" in resolution.no_result_error.error.lower() + + +def test_resolve_post_run_resolution_prefers_streamed_text_and_result_usage(): + state = WebSocketTurnState(collected_text=["hello", " world"]) + result = SimpleNamespace( + output="ignored", + usage=SimpleNamespace(input_tokens=3, output_tokens=5, total_tokens=8), + ) + + resolution = resolve_post_run_resolution( + result=result, + turn_state=state, + agent=_Agent(history=[]), + session_id="s4", + logger=_Logger(), + ) + + assert resolution.response_text == "hello world" + assert resolution.tokens_used == { + "input_tokens": 3, + "output_tokens": 5, + "total_tokens": 8, + } + + +def test_resolve_post_run_resolution_falls_back_to_history_and_extracts_thinking(): + agent = _Agent( + history=[ + _AssistantMessage("final assistant reply"), + _ModelResponse(parts=[_ThinkingPart("hidden reasoning")]), + ], + estimate=11, + ) + state = WebSocketTurnState() + + resolution = resolve_post_run_resolution( + result=False, + turn_state=state, + agent=agent, + session_id="s5", + logger=_Logger(), + ) + + assert resolution.response_text == "final assistant reply" + assert resolution.thinking_text == "hidden reasoning" + assert resolution.tokens_used == { + "total_tokens": 22, + "estimated": True, + } diff --git a/code_puppy/api/ws/tests/test_ws_resume_recovery.py b/code_puppy/api/ws/tests/test_ws_resume_recovery.py new file mode 100644 index 000000000..903700487 --- /dev/null +++ b/code_puppy/api/ws/tests/test_ws_resume_recovery.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +from code_puppy.api.ws.ws_resume_recovery import ( + sanitize_trailing_incomplete_tool_history, +) + + +class ToolCallPart: ... + + +class ToolReturnPart: ... + + +class TextPart: ... + + +class _Msg: + def __init__(self, parts): + self.parts = parts + + +def test_sanitize_trims_trailing_tool_call_only_messages(): + history = [ + _Msg([TextPart()]), + _Msg([ToolCallPart()]), + _Msg([ToolCallPart()]), + ] + + sanitized, removed = sanitize_trailing_incomplete_tool_history(history) + + assert removed == 2 + assert len(sanitized) == 1 + + +def test_sanitize_keeps_messages_with_tool_return_or_text(): + history = [ + _Msg([ToolCallPart(), ToolReturnPart()]), + _Msg([ToolCallPart(), TextPart()]), + ] + + sanitized, removed = sanitize_trailing_incomplete_tool_history(history) + + assert removed == 0 + assert len(sanitized) == 2 diff --git a/code_puppy/api/ws/tests/test_ws_session_bootstrap.py b/code_puppy/api/ws/tests/test_ws_session_bootstrap.py new file mode 100644 index 000000000..4e6722b0f --- /dev/null +++ b/code_puppy/api/ws/tests/test_ws_session_bootstrap.py @@ -0,0 +1,253 @@ +from __future__ import annotations + +import importlib +import sys +from types import ModuleType, SimpleNamespace + +import pytest + +from code_puppy.api.ws.ws_chat_runtime import WebSocketChatRuntime + + +class _FakeAgent: + def __init__(self, history=None): + self._history = list(history or []) + + def get_message_history(self): + return list(self._history) + + +class _FakeSessionManager: + def __init__(self, ctx): + self.ctx = ctx + self.created = [] + self.loaded = [] + self.activated = [] + + async def create_session(self, session_id): + self.created.append(session_id) + return self.ctx + + async def get_or_load_session(self, session_id): + self.loaded.append(session_id) + return self.ctx + + async def mark_session_active(self, session_id): + self.activated.append(session_id) + + +class _FakeWebSocket: + def __init__(self): + self.closed = [] + + async def close(self, *, code, reason): + self.closed.append((code, reason)) + + +class _FakeSender: + def __init__(self): + self.session_id = None + self.ctx = None + + +class _FakeBus: + def __init__(self): + self.contexts = [] + + def set_session_context(self, session_id): + self.contexts.append(session_id) + + +def _install_bootstrap_deps( + monkeypatch, + *, + session_exists=False, + session_metadata=None, + session_row=None, + active_messages=None, + session_manager=None, + writes=None, + bus=None, + init_calls=None, +): + fake_db_pkg = ModuleType("code_puppy.api.db") + fake_db_pkg.__path__ = [] + fake_queries = ModuleType("code_puppy.api.db.queries") + + async def _write_system_message_to_sqlite(**kwargs): + writes.append(kwargs) + + async def _session_exists(_session_id): + return session_exists + + async def _get_session_metadata(_session_id): + return dict(session_metadata or {}) + + async def _get_session_row(_session_id): + return dict(session_row or {}) + + async def _get_active_messages(_session_id): + return list(active_messages or []) + + fake_queries.write_system_message_to_sqlite = _write_system_message_to_sqlite + fake_queries.session_exists = _session_exists + fake_queries.get_session_metadata = _get_session_metadata + fake_queries.get_session_row = _get_session_row + fake_queries.get_active_messages = _get_active_messages + fake_db_pkg.queries = fake_queries + monkeypatch.setitem(sys.modules, "code_puppy.api.db", fake_db_pkg) + monkeypatch.setitem(sys.modules, "code_puppy.api.db.queries", fake_queries) + + fake_session_context = ModuleType("code_puppy.api.session_context") + fake_session_context._validate_session_id = lambda value: value + fake_session_context.session_manager = session_manager + monkeypatch.setitem( + sys.modules, "code_puppy.api.session_context", fake_session_context + ) + + fake_bus_module = ModuleType("code_puppy.messaging.bus") + fake_bus_module.get_message_bus = lambda: bus + monkeypatch.setitem(sys.modules, "code_puppy.messaging.bus", fake_bus_module) + + fake_command_runner = ModuleType("code_puppy.tools.command_runner") + fake_command_runner.init_session_process_tracking = lambda: init_calls.append( + "init" + ) + monkeypatch.setitem( + sys.modules, "code_puppy.tools.command_runner", fake_command_runner + ) + + fake_session_persistence = ModuleType("code_puppy.api.ws.session_persistence") + fake_session_persistence.build_session_meta_payload = lambda **kwargs: kwargs + monkeypatch.setitem( + sys.modules, + "code_puppy.api.ws.session_persistence", + fake_session_persistence, + ) + + +def _load_bootstrap_module(monkeypatch, **kwargs): + sys.modules.pop("code_puppy.api.ws.ws_session_bootstrap", None) + _install_bootstrap_deps(monkeypatch, **kwargs) + return importlib.import_module("code_puppy.api.ws.ws_session_bootstrap") + + +@pytest.mark.asyncio +async def test_initialize_ws_session_creates_new_session(monkeypatch): + ctx = SimpleNamespace( + agent=_FakeAgent([]), + agent_name="code-puppy", + model_name="gpt-test", + title="", + working_directory="", + pinned=False, + ) + manager = _FakeSessionManager(ctx) + writes = [] + bus = _FakeBus() + init_calls = [] + module = _load_bootstrap_module( + monkeypatch, + session_exists=False, + session_manager=manager, + writes=writes, + bus=bus, + init_calls=init_calls, + ) + + sent = [] + meta = [] + + async def _send_typed(message): + sent.append(message) + + async def _safe_send_json(payload): + meta.append(payload) + + runtime = await module.initialize_ws_session( + websocket=_FakeWebSocket(), + requested_session_id="session-1", + sender=_FakeSender(), + safe_send_json=_safe_send_json, + send_typed=_send_typed, + ) + + assert isinstance(runtime, WebSocketChatRuntime) + assert runtime.session_id == "session-1" + assert manager.created == ["session-1"] + assert manager.activated == ["session-1"] + assert init_calls == ["init"] + assert bus.contexts == ["session-1"] + assert writes[0]["system_message_type"] == "config" + assert sent[0].type == "system" + assert sent[0].resumed is False + assert meta[0]["session_id"] == "session-1" + assert meta[0]["message_count"] == 0 + + +@pytest.mark.asyncio +async def test_initialize_ws_session_restores_existing_session(monkeypatch): + ctx = SimpleNamespace( + agent=_FakeAgent([{"role": "user"}, {"role": "assistant"}]), + agent_name="restored-agent", + model_name="restored-model", + title="Saved title", + working_directory="/tmp/project", + pinned=True, + ) + manager = _FakeSessionManager(ctx) + writes = [] + bus = _FakeBus() + init_calls = [] + module = _load_bootstrap_module( + monkeypatch, + session_exists=True, + session_metadata={ + "title": "Saved title", + "working_directory": "/tmp/project", + "pinned": True, + }, + active_messages=[ + { + "role": "system", + "system_message_type": "config", + "content": "restored config", + "agent_name": "restored-agent", + "model_name": "restored-model", + } + ], + session_manager=manager, + writes=writes, + bus=bus, + init_calls=init_calls, + ) + + sent = [] + meta = [] + + async def _send_typed(message): + sent.append(message) + + async def _safe_send_json(payload): + meta.append(payload) + + runtime = await module.initialize_ws_session( + websocket=_FakeWebSocket(), + requested_session_id="session-2", + sender=_FakeSender(), + safe_send_json=_safe_send_json, + send_typed=_send_typed, + ) + + assert runtime.session_id == "session-2" + assert manager.loaded == ["session-2"] + assert writes == [] + assert sent[0].type == "system" + assert sent[0].resumed is True + assert any(message.type == "session_restored" for message in sent) + assert any( + message.type == "system" + and getattr(message, "content", "") == "restored config" + for message in sent[1:] + ) + assert meta[0]["title"] == "Saved title" diff --git a/code_puppy/api/ws/tests/test_ws_stream_drain.py b/code_puppy/api/ws/tests/test_ws_stream_drain.py new file mode 100644 index 000000000..264714392 --- /dev/null +++ b/code_puppy/api/ws/tests/test_ws_stream_drain.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import asyncio +import sys +import types + +import pytest + +from code_puppy.api.ws.ws_stream_drain import ( + StreamDrainHandle, + cancel_active_streaming, + start_stream_drain, + stop_stream_drain, +) + + +class _Logger: + def __init__(self): + self.debug_messages = [] + self.warning_messages = [] + + def debug(self, message, *args): + self.debug_messages.append(message % args if args else message) + + def warning(self, message, *args): + self.warning_messages.append(message % args if args else message) + + +@pytest.mark.asyncio +async def test_start_stream_drain_subscribes_and_waits_for_ready(monkeypatch): + subscribed = [] + queue = asyncio.Queue() + logger = _Logger() + + emitter_module = types.SimpleNamespace( + subscribe=lambda *, session_id: subscribed.append(session_id) or queue, + unsubscribe=lambda event_queue: None, + ) + monkeypatch.setitem( + sys.modules, + "code_puppy.plugins.frontend_emitter.emitter", + emitter_module, + ) + + async def _drain(event_queue, ready_event: asyncio.Event): + assert event_queue is queue + ready_event.set() + await asyncio.sleep(0) + + handle = await start_stream_drain( + session_id="session-1", + drain_coro_factory=_drain, + logger=logger, + ) + await handle.task + + assert subscribed == ["session-1"] + assert handle is not None + assert handle.event_queue is queue + assert "Subscribed to frontend emitter for streaming" in logger.debug_messages + + +@pytest.mark.asyncio +async def test_stop_stream_drain_stops_task_and_unsubscribes(): + unsubscribed = [] + logger = _Logger() + stop_draining = asyncio.Event() + ready = asyncio.Event() + + async def _worker(): + ready.set() + while not stop_draining.is_set(): + await asyncio.sleep(0) + + task = asyncio.create_task(_worker()) + await ready.wait() + + await stop_stream_drain( + handle=StreamDrainHandle( + event_queue="queue-1", + task=task, + unsubscribe=lambda event_queue: unsubscribed.append(event_queue), + ), + stop_draining=stop_draining, + logger=logger, + ) + + assert stop_draining.is_set() + assert task.done() + assert unsubscribed == ["queue-1"] + assert "Unsubscribed from frontend emitter" in logger.debug_messages + + +@pytest.mark.asyncio +async def test_cancel_active_streaming_cancels_and_clears_stop_flag(): + logger = _Logger() + stop_draining = asyncio.Event() + cancel_seen = [] + + async def _worker(): + try: + await asyncio.Future() + except asyncio.CancelledError: + cancel_seen.append(True) + raise + + task = asyncio.create_task(_worker()) + await asyncio.sleep(0) + + await cancel_active_streaming( + active_drain_task=task, + stop_draining=stop_draining, + logger=logger, + log_message="Active streaming cancelled", + ) + + assert task.done() + assert cancel_seen == [True] + assert stop_draining.is_set() is False + assert "Active streaming cancelled" in logger.debug_messages diff --git a/code_puppy/api/ws/tests/test_ws_turn_finalization.py b/code_puppy/api/ws/tests/test_ws_turn_finalization.py new file mode 100644 index 000000000..c501f3930 --- /dev/null +++ b/code_puppy/api/ws/tests/test_ws_turn_finalization.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import pytest + +from code_puppy.api.ws.chat_turn_state import WebSocketTurnState +from code_puppy.api.ws.ws_turn_finalization import ( + emit_pre_stream_end_tool_results, + finalize_turn_history, +) + + +class _Logger: + def __init__(self): + self.debug_messages = [] + self.info_messages = [] + self.warning_messages = [] + + def debug(self, message, *args): + self.debug_messages.append(message % args if args else message) + + def info(self, message, *args): + self.info_messages.append(message % args if args else message) + + def warning(self, message, *args): + self.warning_messages.append(message % args if args else message) + + +class _Agent: + def __init__(self): + self._history = [] + + def set_message_history(self, history): + self._history = list(history) + + def get_message_history(self): + return list(self._history) + + +@dataclass +class _ToolReturnLike: + tool_name: str + tool_call_id: str | None + content: object + + +@dataclass +class _Message: + parts: list[object] + + +class _Result: + def __init__(self, messages): + self._messages = list(messages) + + def all_messages(self): + return list(self._messages) + + +@pytest.fixture(autouse=True) +def _patch_tool_return_types(monkeypatch): + import code_puppy.api.ws.ws_turn_finalization as module + + monkeypatch.setitem( + __import__("sys").modules, + "pydantic_ai.messages", + type( + "M", + (), + { + "ToolReturn": _ToolReturnLike, + "ToolReturnPart": _ToolReturnLike, + }, + )(), + ) + return module + + +@pytest.mark.asyncio +async def test_emit_pre_stream_end_tool_results_uses_alias_and_group_id(): + turn_state = WebSocketTurnState( + tool_id_aliases={"raw-1": "tool-1"}, + tool_group_ids={"tool-1": "tg-1"}, + ) + sent = [] + + async def _send(model): + sent.append(model) + + result = _Result([_Message(parts=[_ToolReturnLike("calc", "raw-1", {"x": 1})])]) + + pre_sent = await emit_pre_stream_end_tool_results( + result=result, + turn_state=turn_state, + session_id="s1", + agent_name="agent", + model_name="model", + send_typed_tool_lifecycle=_send, + logger=_Logger(), + ) + + assert pre_sent == {"tool-1"} + assert len(sent) == 1 + assert sent[0].tool_id == "tool-1" + assert sent[0].tool_group_id == "tg-1" + assert sent[0].result == {"x": 1} + + +@pytest.mark.asyncio +async def test_finalize_turn_history_snapshots_history_before_awaits(): + turn_state = WebSocketTurnState(tool_group_ids={}) + agent = _Agent() + sent = [] + + async def _send(model): + sent.append(model) + + messages = [_Message(parts=[]), _Message(parts=[])] + finalized = await finalize_turn_history( + result=_Result(messages), + agent=agent, + turn_state=turn_state, + session_id="s2", + agent_name="agent", + model_name="model", + send_typed=_send, + pre_sent_tool_ids=set(), + logger=_Logger(), + ) + + assert finalized.history_snapshot == messages + assert agent.get_message_history() == messages + assert sent == [] + + +@pytest.mark.asyncio +async def test_finalize_turn_history_skips_pre_sent_duplicates(): + turn_state = WebSocketTurnState(tool_group_ids={"tool-9": "tg-9"}) + agent = _Agent() + sent = [] + + async def _send(model): + sent.append(model) + + result = _Result( + [_Message(parts=[_ToolReturnLike("shell", "tool-9", {"ok": True})])] + ) + finalized = await finalize_turn_history( + result=result, + agent=agent, + turn_state=turn_state, + session_id="s3", + agent_name="agent", + model_name="model", + send_typed=_send, + pre_sent_tool_ids={"tool-9"}, + logger=_Logger(), + ) + + assert finalized.pre_sent_tool_ids == {"tool-9"} + assert sent == [] + + +@pytest.mark.asyncio +async def test_finalize_turn_history_emits_remaining_tool_results_and_serializes(): + class _Complex: + def __init__(self): + self.answer = 42 + + turn_state = WebSocketTurnState(tool_group_ids={"tool-2": "tg-2"}) + agent = _Agent() + sent = [] + + async def _send(model): + sent.append(model) + + result = _Result([_Message(parts=[_ToolReturnLike("calc", "tool-2", _Complex())])]) + finalized = await finalize_turn_history( + result=result, + agent=agent, + turn_state=turn_state, + session_id="s4", + agent_name="agent", + model_name="model", + send_typed=_send, + pre_sent_tool_ids=set(), + logger=_Logger(), + ) + + assert finalized.history_snapshot + assert len(sent) == 1 + assert sent[0].tool_id == "tool-2" + assert sent[0].tool_group_id == "tg-2" + assert sent[0].result == {"answer": 42} diff --git a/code_puppy/api/ws/tests/test_ws_turn_preparation.py b/code_puppy/api/ws/tests/test_ws_turn_preparation.py new file mode 100644 index 000000000..4d982325e --- /dev/null +++ b/code_puppy/api/ws/tests/test_ws_turn_preparation.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock + +from code_puppy.api.ws.ws_turn_preparation import prepare_turn_input + + +def test_prepare_turn_input_injects_directory_and_collects_attachments( + monkeypatch, + tmp_path, +): + text_file = tmp_path / "notes.txt" + text_file.write_text("hello world", encoding="utf-8") + + monkeypatch.setattr( + "code_puppy.api.ws.ws_turn_preparation.build_file_context_and_attachments", + lambda msg: ("FILE CONTEXT", ["binary-1"]), + ) + + agent = MagicMock() + prepared = prepare_turn_input( + agent=agent, + user_message="Analyze this", + msg={"attachments": [str(text_file), "", None, str(tmp_path / "missing.txt")]}, + session_working_directory="/tmp/project", + last_context_sent_directory="", + ) + + assert agent.append_to_message_history.call_count == 1 + assert prepared.last_context_sent_directory == "/tmp/project" + assert prepared.message_to_send == "FILE CONTEXT\n\nAnalyze this" + assert prepared.run_kwargs == {"attachments": ["binary-1"]} + assert prepared.attachment_metadata == [ + { + "name": "notes.txt", + "path": str(Path(text_file).absolute()), + "sizeBytes": text_file.stat().st_size, + } + ] + + +def test_prepare_turn_input_is_noop_when_directory_unchanged(monkeypatch): + monkeypatch.setattr( + "code_puppy.api.ws.ws_turn_preparation.build_file_context_and_attachments", + lambda msg: ("", []), + ) + + agent = MagicMock() + prepared = prepare_turn_input( + agent=agent, + user_message="Hello", + msg={}, + session_working_directory="/tmp/project", + last_context_sent_directory="/tmp/project", + ) + + agent.append_to_message_history.assert_not_called() + assert prepared.last_context_sent_directory == "/tmp/project" + assert prepared.message_to_send == "Hello" + assert prepared.run_kwargs == {} + assert prepared.attachment_metadata == [] diff --git a/code_puppy/api/ws/ws_chat_runtime.py b/code_puppy/api/ws/ws_chat_runtime.py new file mode 100644 index 000000000..e32984a0a --- /dev/null +++ b/code_puppy/api/ws/ws_chat_runtime.py @@ -0,0 +1,49 @@ +"""Connection-scoped runtime state for the chat WebSocket. + +This module intentionally stores only ephemeral mutable state that previously +lived as local variables inside ``websocket_chat``. Each WebSocket connection +creates its own instance so there is no shared mutable state across sessions. +""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass, field +from typing import Any + + +@dataclass(slots=True) +class WebSocketChatRuntime: + """Ephemeral state for one active chat WebSocket connection.""" + + session_id: str + ctx: Any | None = None + session_title: str = "" + session_working_directory: str = "" + session_pinned: bool = False + last_context_sent_directory: str = "" + existing_history: Any | None = None + agent: Any | None = None + agent_name: str = "code-puppy" + model_name: str = "unknown" + active_drain_task: asyncio.Task | None = None + active_agent_task: asyncio.Task | None = None + stop_draining: asyncio.Event = field(default_factory=asyncio.Event) + + def sync_from_ctx(self) -> None: + """Refresh convenience aliases from the attached session context.""" + if self.ctx is None: + return + self.agent = getattr(self.ctx, "agent", None) + self.agent_name = getattr(self.ctx, "agent_name", self.agent_name) + self.model_name = getattr(self.ctx, "model_name", self.model_name) + self.session_title = getattr(self.ctx, "title", self.session_title) + self.session_working_directory = getattr( + self.ctx, + "working_directory", + self.session_working_directory, + ) + self.session_pinned = bool(getattr(self.ctx, "pinned", self.session_pinned)) + + +__all__ = ["WebSocketChatRuntime"] diff --git a/code_puppy/api/ws/ws_command_handler.py b/code_puppy/api/ws/ws_command_handler.py new file mode 100644 index 000000000..31d0775d7 --- /dev/null +++ b/code_puppy/api/ws/ws_command_handler.py @@ -0,0 +1,93 @@ +"""Slash-command handling for the chat WebSocket.""" + +from __future__ import annotations + +import logging +from io import StringIO +from typing import Any + +from rich.console import Console + +from code_puppy.api.ws.schemas import ServerCommandResult + +logger = logging.getLogger(__name__) + + +async def handle_command_message( + *, + msg: dict[str, Any], + session_id: str, + send_typed: Any, +) -> bool: + """Handle a websocket ``type=command`` payload.""" + if msg.get("type") != "command": + return False + + command_str = msg.get("command", "") + logger.debug("Command requested: %s", command_str) + + try: + from code_puppy.command_line.command_handler import ( + get_commands_help, + handle_command, + ) + + output = None + captured_messages = [] + + cmd_name = ( + command_str.strip().lstrip("/").split()[0] if command_str.strip() else "" + ) + if cmd_name in ("help", "h"): + help_text = get_commands_help() + output_buffer = StringIO() + console = Console( + file=output_buffer, + force_terminal=False, + width=100, + no_color=True, + ) + console.print(help_text) + output = output_buffer.getvalue().strip() + result = True + else: + result = handle_command(command_str) + if isinstance(result, str): + output = result + result = True + + await send_typed( + ServerCommandResult( + command=command_str, + success=result is True, + output=output, + messages=captured_messages, + result=str(result) if result and result is not True else None, + session_id=session_id, + ) + ) + logger.debug( + "Command executed: %s -> success=%s, output_len=%s", + command_str, + result is True, + len(output) if output else 0, + ) + except Exception as cmd_error: + import traceback + + error_details = traceback.format_exc() + logger.error("Command error: %s", cmd_error) + logger.error("Traceback: %s", error_details) + await send_typed( + ServerCommandResult( + command=command_str, + success=False, + error=str(cmd_error), + session_id=session_id, + ) + ) + + return True + + +__all__ = ["handle_command_message"] diff --git a/code_puppy/api/ws/ws_control_messages.py b/code_puppy/api/ws/ws_control_messages.py new file mode 100644 index 000000000..ef1f9adb0 --- /dev/null +++ b/code_puppy/api/ws/ws_control_messages.py @@ -0,0 +1,608 @@ +"""Top-level non-chat control-message handlers for the chat WebSocket.""" + +from __future__ import annotations + +import asyncio +import datetime +import logging +from pathlib import Path +from typing import Any, Awaitable, Callable + +from code_puppy.api.db.queries import ( + update_session_working_directory, + write_system_message_to_sqlite, +) +from code_puppy.api.permissions import handle_permission_response +from code_puppy.api.session_context import _validate_session_id, session_manager +from code_puppy.api.ws.schemas import ( + ServerConfigValue, + ServerError, + ServerSessionMetaUpdated, + ServerSessionSwitched, + ServerStatus, + ServerSystem, + ServerWorkingDirectoryChanged, +) +from code_puppy.api.ws.send_utils import WebSocketSender +from code_puppy.api.ws.ws_chat_runtime import WebSocketChatRuntime +from code_puppy.api.ws.ws_stream_drain import cancel_active_streaming + +logger = logging.getLogger(__name__) + + +async def handle_control_message( + *, + msg: dict[str, Any], + runtime: WebSocketChatRuntime, + sender: WebSocketSender, + send_typed: Any, + send_session_meta_snapshot: Callable[[], Awaitable[None]], +) -> bool: + """Handle non-chat, non-command websocket control messages.""" + return ( + await _handle_switch_agent( + msg=msg, + runtime=runtime, + send_typed=send_typed, + ) + or await _handle_switch_model( + msg=msg, + runtime=runtime, + send_typed=send_typed, + ) + or await _handle_switch_session( + msg=msg, + runtime=runtime, + sender=sender, + send_typed=send_typed, + send_session_meta_snapshot=send_session_meta_snapshot, + ) + or await _handle_set_working_directory( + msg=msg, + runtime=runtime, + send_typed=send_typed, + ) + or await _handle_update_session_meta( + msg=msg, + runtime=runtime, + send_typed=send_typed, + ) + or await _handle_get_config( + msg=msg, + runtime=runtime, + send_typed=send_typed, + ) + or await _handle_set_config( + msg=msg, + runtime=runtime, + send_typed=send_typed, + ) + or await _handle_cancel( + msg=msg, + runtime=runtime, + send_typed=send_typed, + ) + or await _handle_permission_response(msg=msg, session_id=runtime.session_id) + ) + + +async def _handle_switch_agent( + *, msg: dict[str, Any], runtime: WebSocketChatRuntime, send_typed: Any +) -> bool: + if msg.get("type") != "switch_agent": + return False + + agent_name = msg.get("agent_name") + if agent_name: + try: + new_agent = await session_manager.switch_agent( + runtime.session_id, agent_name + ) + runtime.agent = new_agent + if runtime.ctx is not None: + runtime.ctx.agent = new_agent + runtime.ctx.agent_name = agent_name + runtime.agent_name = agent_name + runtime.model_name = new_agent.get_model_name() if new_agent else "unknown" + + try: + await write_system_message_to_sqlite( + session_id=runtime.session_id, + system_message_type="config", + content=f"🔄 Switched to {agent_name} ({runtime.model_name})", + agent_name=agent_name, + model_name=runtime.model_name, + ) + except Exception as exc: + logger.warning( + "Agent-switch SQLite write failed: %s", exc, exc_info=True + ) + + await send_typed( + ServerSystem( + content=f"🔄 Switched to {agent_name} ({runtime.model_name})", + session_id=runtime.session_id, + agent_name=agent_name, + model_name=runtime.model_name, + ) + ) + except Exception as exc: + logger.error("Error switching agent: %s", exc) + await send_typed( + ServerError( + error=f"Failed to switch to agent {agent_name}: {str(exc)}", + session_id=runtime.session_id, + ) + ) + return True + + +async def _handle_switch_model( + *, msg: dict[str, Any], runtime: WebSocketChatRuntime, send_typed: Any +) -> bool: + if msg.get("type") != "switch_model": + return False + + model_name = msg.get("model_name") or msg.get("model") + if model_name: + try: + await session_manager.switch_model(runtime.session_id, model_name) + runtime.agent = runtime.ctx.agent if runtime.ctx else runtime.agent + runtime.model_name = model_name + if runtime.ctx is not None: + runtime.ctx.model_name = model_name + logger.debug("Switched model to: %s", model_name) + + switch_agent = ( + (runtime.ctx.agent_name if runtime.ctx else None) + or runtime.agent_name + or "code-puppy" + ) + try: + await write_system_message_to_sqlite( + session_id=runtime.session_id, + system_message_type="config", + content=f"🔄 Switched to {switch_agent} ({model_name})", + agent_name=switch_agent, + model_name=model_name, + ) + except Exception as exc: + logger.warning( + "Model-switch SQLite write failed: %s", exc, exc_info=True + ) + + await send_typed( + ServerSystem( + content=f"🔄 Switched to {switch_agent} ({model_name})", + session_id=runtime.session_id, + model_name=model_name, + agent_name=switch_agent, + ) + ) + except Exception as exc: + logger.error("Error switching model: %s", exc) + await send_typed( + ServerError( + error=f"Failed to switch to model {model_name}: {str(exc)}", + session_id=runtime.session_id, + ) + ) + else: + await send_typed( + ServerError( + error="No model_name provided for switch_model", + session_id=runtime.session_id, + ) + ) + return True + + +async def _handle_switch_session( + *, + msg: dict[str, Any], + runtime: WebSocketChatRuntime, + sender: WebSocketSender, + send_typed: Any, + send_session_meta_snapshot: Callable[[], Awaitable[None]], +) -> bool: + if msg.get("type") != "switch_session": + return False + + logger.debug("Cancelling active streaming due to session switch") + await cancel_active_streaming( + active_drain_task=runtime.active_drain_task, + stop_draining=runtime.stop_draining, + logger=logger, + ) + + new_session_id = msg.get("session_id") + if not new_session_id: + await send_typed( + ServerError( + error="No session_id provided for switch_session", + session_id=runtime.session_id, + ) + ) + return True + + try: + _validate_session_id(new_session_id) + except ValueError as exc: + logger.warning("Invalid session_id in switch_session: %r", new_session_id) + await send_typed( + ServerError( + error=f"Invalid session ID: {exc}", + session_id=runtime.session_id, + ) + ) + return True + + logger.debug("Switching to session: %s", new_session_id) + try: + try: + from code_puppy.api.db.queries import session_exists as _session_exists + + target_exists = await _session_exists(new_session_id) + except Exception: + target_exists = False + + try: + await session_manager.save_session(runtime.session_id) + except Exception: + pass + await session_manager.mark_session_inactive(runtime.session_id) + + if not target_exists: + logger.debug("Session %s not found, creating new", new_session_id) + runtime.session_id = new_session_id + sender.session_id = new_session_id + runtime.session_title = "" + runtime.session_working_directory = "" + runtime.session_pinned = False + runtime.existing_history = None + runtime.ctx = await session_manager.create_session(new_session_id) + sender.ctx = runtime.ctx + await session_manager.mark_session_active(new_session_id) + runtime.sync_from_ctx() + await send_typed( + ServerSessionSwitched( + session_id=new_session_id, + message_count=0, + title="", + created=True, + agent_name=runtime.agent_name, + model_name=runtime.model_name, + ) + ) + return True + + loaded_ctx = await session_manager.get_or_load_session(new_session_id) + runtime.session_id = new_session_id + sender.session_id = new_session_id + if loaded_ctx is not None: + runtime.ctx = loaded_ctx + runtime.sync_from_ctx() + else: + new_title = "" + new_working_directory = "" + new_pinned = False + try: + from code_puppy.api.db.queries import ( + get_session_metadata as _get_session_metadata, + ) + + session_meta = await _get_session_metadata(new_session_id) or {} + new_title = session_meta.get("title", "") + new_working_directory = session_meta.get("working_directory", "") + new_pinned = bool(session_meta.get("pinned", False)) + except Exception: + pass + + runtime.ctx = await session_manager.create_session(new_session_id) + sender.ctx = runtime.ctx + runtime.ctx.title = new_title + runtime.ctx.working_directory = new_working_directory + runtime.ctx.pinned = new_pinned + runtime.sync_from_ctx() + + sender.ctx = runtime.ctx + await session_manager.mark_session_active(new_session_id) + runtime.sync_from_ctx() + message_count = ( + len(runtime.ctx.agent.get_message_history() or []) if runtime.ctx else 0 + ) + logger.debug("Restored %d messages to session agent", message_count) + + await send_typed( + ServerSessionSwitched( + session_id=new_session_id, + message_count=message_count, + title=runtime.session_title, + working_directory=runtime.session_working_directory, + created=False, + agent_name=runtime.agent_name, + model_name=runtime.model_name, + ) + ) + await send_session_meta_snapshot() + logger.debug( + "Switched to session %s with %d messages", + new_session_id, + message_count, + ) + except Exception as exc: + logger.error("Error switching session: %s", exc) + await send_typed( + ServerError( + error=f"Failed to switch to session {new_session_id}: {str(exc)}", + session_id=runtime.session_id, + ) + ) + return True + + +async def _handle_set_working_directory( + *, msg: dict[str, Any], runtime: WebSocketChatRuntime, send_typed: Any +) -> bool: + if msg.get("type") != "set_working_directory": + return False + + new_directory = msg.get("directory", "") + if not new_directory: + await send_typed( + ServerError( + error="No directory provided for set_working_directory", + session_id=runtime.session_id, + ) + ) + return True + + new_directory = str(Path(new_directory).expanduser().resolve()) + logger.info( + "[CWD DEBUG] set_working_directory received: new=%r, current=%r, session=%s", + new_directory, + runtime.session_working_directory, + runtime.session_id, + ) + if not Path(new_directory).is_dir(): + await send_typed( + ServerWorkingDirectoryChanged( + directory=new_directory, + success=False, + error="Directory does not exist", + session_id=runtime.session_id, + ) + ) + return True + + if new_directory == runtime.session_working_directory: + logger.info( + "[CWD DEBUG] Skipping unchanged directory: %r, session=%s", + new_directory, + runtime.session_id, + ) + await send_typed( + ServerWorkingDirectoryChanged( + directory=new_directory, + success=True, + session_id=runtime.session_id, + unchanged=True, + ) + ) + return True + + runtime.session_working_directory = new_directory + if runtime.ctx is not None: + runtime.ctx.working_directory = new_directory + logger.debug("Working directory set to: %s", runtime.session_working_directory) + + try: + cwd_agent = (runtime.ctx.agent_name if runtime.ctx else None) or "code-puppy" + cwd_model = (runtime.ctx.model_name if runtime.ctx else None) or "unknown" + cwd_segs = runtime.session_working_directory.split("/")[-3:] + cwd_rel = "/".join(s for s in cwd_segs if s) + from code_puppy.config import get_puppy_name as _get_puppy_name + + puppy_name = _get_puppy_name() or "puppy" + logger.info( + "[CWD DEBUG] Writing CWD banner to SQLite: path=%r, session=%s", + runtime.session_working_directory, + runtime.session_id, + ) + await write_system_message_to_sqlite( + session_id=runtime.session_id, + system_message_type="directory", + content=f"{puppy_name} is now at {cwd_rel}", + system_message_path=runtime.session_working_directory, + agent_name=cwd_agent, + model_name=cwd_model, + ) + now_cwd = datetime.datetime.now(datetime.timezone.utc).isoformat() + await update_session_working_directory( + session_id=runtime.session_id, + working_directory=runtime.session_working_directory, + updated_at=now_cwd, + ) + except Exception as exc: + logger.warning("CWD SQLite write failed: %s", exc, exc_info=True) + + await send_typed( + ServerWorkingDirectoryChanged( + directory=runtime.session_working_directory, + success=True, + session_id=runtime.session_id, + ) + ) + return True + + +async def _handle_update_session_meta( + *, msg: dict[str, Any], runtime: WebSocketChatRuntime, send_typed: Any +) -> bool: + if msg.get("type") != "update_session_meta": + return False + + try: + if "pinned" in msg and isinstance(msg["pinned"], bool): + runtime.session_pinned = msg["pinned"] + if runtime.ctx is not None: + runtime.ctx.pinned = runtime.session_pinned + if "title" in msg and isinstance(msg["title"], str): + runtime.session_title = msg["title"] + if runtime.ctx is not None: + runtime.ctx.title = runtime.session_title + + try: + from datetime import datetime as _dt + from datetime import timezone as _tz + + from code_puppy.api.db.queries import update_session_meta_fields + + await update_session_meta_fields( + session_id=runtime.session_id, + title=runtime.session_title, + pinned=runtime.session_pinned, + updated_at=_dt.now(_tz.utc).isoformat(), + ) + logger.debug( + "Updated session meta in SQLite for %s: pinned=%s", + runtime.session_id, + runtime.session_pinned, + ) + except Exception as meta_exc: + logger.warning("Failed to persist session meta to SQLite: %s", meta_exc) + + await send_typed( + ServerSessionMetaUpdated( + session_id=runtime.session_id, + pinned=runtime.session_pinned, + title=runtime.session_title, + ) + ) + except Exception as exc: + logger.error("Error updating session meta: %s", exc) + await send_typed( + ServerError( + error=f"Failed to update session metadata: {str(exc)}", + session_id=runtime.session_id, + ) + ) + return True + + +async def _handle_get_config( + *, msg: dict[str, Any], runtime: WebSocketChatRuntime, send_typed: Any +) -> bool: + if msg.get("type") != "get_config": + return False + + config_key = msg.get("key", "") + if config_key: + from code_puppy.config import get_value + + value = get_value(config_key) + await send_typed( + ServerConfigValue( + key=config_key, + value=value, + session_id=runtime.session_id, + ) + ) + else: + await send_typed( + ServerError( + error="No key provided for get_config", + session_id=runtime.session_id, + ) + ) + return True + + +async def _handle_set_config( + *, msg: dict[str, Any], runtime: WebSocketChatRuntime, send_typed: Any +) -> bool: + if msg.get("type") != "set_config": + return False + + config_key = msg.get("key", "") + config_value = msg.get("value", "") + if config_key: + from code_puppy.config import set_config_value + + try: + set_config_value(config_key, str(config_value)) + logger.debug("Config set: %s = %s", config_key, config_value) + await send_typed( + ServerConfigValue( + key=config_key, + value=config_value, + success=True, + session_id=runtime.session_id, + ) + ) + except Exception as exc: + await send_typed( + ServerError( + error=f"Failed to set config: {exc}", + session_id=runtime.session_id, + ) + ) + else: + await send_typed( + ServerError( + error="No key provided for set_config", + session_id=runtime.session_id, + ) + ) + return True + + +async def _handle_cancel( + *, msg: dict[str, Any], runtime: WebSocketChatRuntime, send_typed: Any +) -> bool: + if msg.get("type") != "cancel": + return False + + logger.debug("Cancel request received - stopping active streaming and agent task") + await cancel_active_streaming( + active_drain_task=runtime.active_drain_task, + stop_draining=runtime.stop_draining, + logger=logger, + log_message="Active streaming cancelled", + ) + + if runtime.active_agent_task and not runtime.active_agent_task.done(): + logger.debug("Cancelling active agent task due to user interrupt") + runtime.active_agent_task.cancel() + try: + await runtime.active_agent_task + except asyncio.CancelledError: + logger.debug("Active agent task cancelled successfully") + runtime.active_agent_task = None + + await send_typed(ServerStatus(status="cancelled", session_id=runtime.session_id)) + return True + + +async def _handle_permission_response(*, msg: dict[str, Any], session_id: str) -> bool: + if msg.get("type") != "permission_response": + return False + + request_id = msg.get("request_id") + approved = msg.get("approved", False) + if request_id: + handled = handle_permission_response( + request_id, approved, session_id=session_id + ) + if handled: + logger.debug( + "[Permission] ✅ Handled response: %s = %s", request_id, approved + ) + else: + logger.warning("[Permission] ❌ Unknown request: %s", request_id) + else: + logger.error("[WebSocket] ❌ No request_id in permission_response!") + return True + + +__all__ = ["handle_control_message"] diff --git a/code_puppy/api/ws/ws_post_run.py b/code_puppy/api/ws/ws_post_run.py new file mode 100644 index 000000000..ed4ae439c --- /dev/null +++ b/code_puppy/api/ws/ws_post_run.py @@ -0,0 +1,247 @@ +"""Helpers for post-run WebSocket outcome resolution. + +This module extracts the response/error/cancelled/no-result decision tree from +chat_handler.py while leaving the caller in charge of exact WebSocket send +ordering. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from code_puppy.api.ws.response_frames import ( + build_error_response_frames, + has_streamed_content, + parse_api_error, +) +from code_puppy.api.ws.schemas import ServerError + + +@dataclass(slots=True) +class PostRunResolution: + """Resolved post-run state for one completed agent turn.""" + + cancelled: bool = False + error_frames: list[dict[str, Any]] | None = None + no_result_error: ServerError | None = None + response_text: str = "" + tokens_used: dict[str, Any] | None = None + thinking_text: str = "" + + +def resolve_post_run_resolution( + *, + result: Any, + turn_state: Any, + agent: Any, + session_id: str, + logger: Any, +) -> PostRunResolution: + """Resolve post-turn response/error/cancelled state without sending frames.""" + if turn_state.agent_error == "cancelled": + return PostRunResolution(cancelled=True) + + if turn_state.agent_error is not None: + logger.debug( + "[WS:%s] turn_state.agent_error -> sending frame(s) to client. type=%s", + session_id, + type(turn_state.agent_error).__name__, + ) + return PostRunResolution( + error_frames=build_error_response_frames( + turn_state.agent_error, + turn_state.collected_text, + session_id, + ) + ) + + has_nonempty_stream = has_streamed_content(turn_state.collected_text) + logger.debug( + "[WS:%s] post-run: result_is_none=%s collected_chunks=%d has_nonempty_stream=%s", + session_id, + result is None, + len(turn_state.collected_text), + has_nonempty_stream, + ) + + if result is None and not has_nonempty_stream: + logger.warning( + "[WS:%s] Agent task completed with result=None and no streamed text; treating as error.", + session_id, + ) + parsed_error = parse_api_error( + RuntimeError( + "Agent run failed (no result returned). Check server logs for the underlying exception." + ) + ) + return PostRunResolution( + no_result_error=ServerError( + error=parsed_error["user_message"], + error_type=parsed_error["error_type"], + technical_details=parsed_error["technical_details"], + action_required=parsed_error.get("action_required"), + session_id=session_id, + ) + ) + + response_text = extract_response_text( + result=result, + collected_text=turn_state.collected_text, + agent=agent, + logger=logger, + ) + tokens_used = extract_token_usage(result=result, agent=agent) + thinking_text = ( + "" + if getattr(turn_state, "b1_streaming_used", False) + else extract_thinking_text(agent=agent, logger=logger) + ) + + return PostRunResolution( + response_text=response_text, + tokens_used=tokens_used, + thinking_text=thinking_text, + ) + + +def extract_response_text( + *, result: Any, collected_text: list[str], agent: Any, logger: Any +) -> str: + """Extract the final assistant response text using the existing priority order.""" + response_text = "" + + if has_streamed_content(collected_text): + response_text = "".join(collected_text) + logger.debug(f"Using collected streaming text ({len(response_text)} chars)") + elif result: + if hasattr(result, "output"): + response_text = str(result.output) if result.output else "(Empty response)" + logger.debug(f"Using result.output ({len(response_text)} chars)") + elif hasattr(result, "data"): + response_text = str(result.data) if result.data else "(Empty response)" + logger.debug(f"Using result.data ({len(response_text)} chars)") + elif agent: + messages = agent.get_message_history() + for msg in reversed(messages): + if ( + hasattr(msg, "role") + and msg.role == "assistant" + and hasattr(msg, "content") + ): + response_text = str(msg.content) + logger.debug(f"Using message history ({len(response_text)} chars)") + break + + if not response_text: + response_text = "Agent returned no response" + + return response_text + + +def extract_token_usage(*, result: Any, agent: Any) -> dict[str, Any] | None: + """Extract token usage from result metadata or estimate from history.""" + tokens_used = None + if result: + if hasattr(result, "usage"): + usage = result.usage + if usage: + tokens_used = { + "input_tokens": getattr(usage, "input_tokens", None) + or getattr(usage, "prompt_tokens", None), + "output_tokens": getattr(usage, "output_tokens", None) + or getattr(usage, "completion_tokens", None), + "total_tokens": getattr(usage, "total_tokens", None), + } + elif hasattr(result, "_usage"): + usage = result._usage + if usage: + tokens_used = { + "input_tokens": getattr(usage, "input_tokens", None) + or getattr(usage, "prompt_tokens", None), + "output_tokens": getattr(usage, "output_tokens", None) + or getattr(usage, "completion_tokens", None), + "total_tokens": getattr(usage, "total_tokens", None), + } + + if not tokens_used and agent: + try: + history = agent.get_message_history() + total_estimated = sum( + agent.estimate_tokens_for_message(msg) for msg in history + ) + tokens_used = { + "total_tokens": total_estimated, + "estimated": True, + } + except Exception: + pass + + return tokens_used + + +def extract_thinking_text(*, agent: Any, logger: Any) -> str: + """Extract thinking content from the last response-like history message.""" + thinking_text = "" + if agent: + try: + history = agent.get_message_history() + logger.debug( + "[Thinking Debug] Checking history for thinking parts, %s messages", + len(history) if history else 0, + ) + if history: + for i, msg in enumerate(reversed(history)): + msg_type = type(msg).__name__ + logger.debug( + "[Thinking Debug] Message %s: type=%s, has_parts=%s", + i, + msg_type, + hasattr(msg, "parts"), + ) + if "Response" in msg_type and hasattr(msg, "parts"): + logger.debug( + "[Thinking Debug] Found Response with %s parts", + len(msg.parts), + ) + for j, part in enumerate(msg.parts): + part_type = type(part).__name__ + part_content_preview = ( + str(getattr(part, "content", ""))[:100] + if hasattr(part, "content") + else "N/A" + ) + logger.debug( + "[Thinking Debug] Part %s: type=%s, content_preview=%s", + j, + part_type, + part_content_preview, + ) + if "Thinking" in part_type and hasattr(part, "content"): + thinking_text = part.content + logger.debug( + "[Thinking Debug] Found thinking content: %s chars", + len(thinking_text), + ) + break + if thinking_text: + break + except Exception as e: + logger.warning(f"Could not extract thinking content: {e}") + import traceback + + logger.warning(traceback.format_exc()) + + if not thinking_text: + logger.debug("[Thinking Debug] No thinking content found in message history") + + return thinking_text + + +__all__ = [ + "PostRunResolution", + "extract_response_text", + "extract_thinking_text", + "extract_token_usage", + "resolve_post_run_resolution", +] diff --git a/code_puppy/api/ws/ws_resume_recovery.py b/code_puppy/api/ws/ws_resume_recovery.py new file mode 100644 index 000000000..8380e0f0b --- /dev/null +++ b/code_puppy/api/ws/ws_resume_recovery.py @@ -0,0 +1,103 @@ +"""Safe one-shot recovery helpers for resumed WebSocket sessions. + +When a resumed turn returns ``result=None`` with no streamed text, we can +attempt a bounded recovery by reloading canonical history from SQLite and +trimming obviously incomplete trailing tool-only responses. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +@dataclass(slots=True) +class ResumeRecoveryResult: + """Outcome of a resume recovery attempt.""" + + success: bool + ctx: Any | None = None + removed_messages: int = 0 + reason: str = "" + + +def _is_incomplete_tool_only_response(message: Any) -> bool: + """Return True when a trailing response appears to be tool-incomplete. + + Conservative heuristic: + - message has ``parts`` + - contains at least one ToolCallPart + - contains no ToolReturnPart + - contains no TextPart + """ + parts = getattr(message, "parts", None) + if not parts: + return False + + part_types = {type(part).__name__ for part in parts} + has_tool_call = "ToolCallPart" in part_types + has_tool_return = "ToolReturnPart" in part_types + has_text = "TextPart" in part_types + return has_tool_call and not has_tool_return and not has_text + + +def sanitize_trailing_incomplete_tool_history( + history: list[Any], +) -> tuple[list[Any], int]: + """Trim incomplete trailing tool-only assistant responses from history.""" + if not history: + return history, 0 + + trimmed = list(history) + removed = 0 + while trimmed and _is_incomplete_tool_only_response(trimmed[-1]): + trimmed.pop() + removed += 1 + + return trimmed, removed + + +async def reload_session_from_sqlite_with_sanitization( + *, + session_id: str, + logger: Any, +) -> ResumeRecoveryResult: + """Force a fresh SQLite reload for this session and sanitize trailing history.""" + try: + # Ensure next load is truly from SQLite canonical state. + from code_puppy.api.session_context import session_manager + + await session_manager.destroy_session(session_id) + ctx = await session_manager.load_session(session_id) + if ctx is None: + return ResumeRecoveryResult( + success=False, + reason="Session not found in SQLite during recovery reload", + ) + + history = ctx.agent.get_message_history() or [] + sanitized, removed = sanitize_trailing_incomplete_tool_history(history) + if removed > 0: + try: + ctx.agent.set_message_history(sanitized) + logger.warning( + "[WS:%s] Resume recovery trimmed %d trailing incomplete tool message(s)", + session_id, + removed, + ) + except Exception as exc: + return ResumeRecoveryResult( + success=False, + reason=f"Failed to apply sanitized history: {exc}", + ) + + return ResumeRecoveryResult(success=True, ctx=ctx, removed_messages=removed) + except Exception as exc: + return ResumeRecoveryResult(success=False, reason=str(exc)) + + +__all__ = [ + "ResumeRecoveryResult", + "reload_session_from_sqlite_with_sanitization", + "sanitize_trailing_incomplete_tool_history", +] diff --git a/code_puppy/api/ws/ws_session_bootstrap.py b/code_puppy/api/ws/ws_session_bootstrap.py new file mode 100644 index 000000000..505a046e2 --- /dev/null +++ b/code_puppy/api/ws/ws_session_bootstrap.py @@ -0,0 +1,271 @@ +"""Session bootstrap helpers for the chat WebSocket.""" + +from __future__ import annotations + +import datetime +import logging +from typing import Any + +from fastapi import WebSocket + +from code_puppy.api.db.queries import write_system_message_to_sqlite +from code_puppy.api.session_context import _validate_session_id, session_manager +from code_puppy.api.ws.schemas import ( + PROTOCOL_VERSION, + ServerSessionRestored, + ServerSystem, +) +from code_puppy.api.ws.send_utils import WebSocketSender +from code_puppy.api.ws.session_persistence import build_session_meta_payload +from code_puppy.api.ws.ws_chat_runtime import WebSocketChatRuntime +from code_puppy.messaging.bus import get_message_bus +from code_puppy.tools.command_runner import init_session_process_tracking + +logger = logging.getLogger(__name__) + + +async def send_session_meta_snapshot( + *, + runtime: WebSocketChatRuntime, + safe_send_json: Any, +) -> None: + """Send the latest session metadata snapshot to the client.""" + try: + from code_puppy.api.db.queries import get_session_row + + session_row = await get_session_row(runtime.session_id) or {} + except Exception: + session_row = {} + + history_for_meta = [] + if runtime.ctx is not None and getattr(runtime.ctx, "agent", None) is not None: + try: + history_for_meta = runtime.ctx.agent.get_message_history() or [] + except Exception: + history_for_meta = [] + + await safe_send_json( + build_session_meta_payload( + session_id=runtime.session_id, + session_name=runtime.session_id, + total_tokens=int(session_row.get("total_tokens") or 0), + message_count=int( + session_row.get("message_count") or len(history_for_meta) + ), + title=runtime.session_title or str(session_row.get("title") or ""), + working_directory=( + runtime.session_working_directory + or str(session_row.get("working_directory") or "") + ), + agent_name=runtime.agent_name, + model_name=runtime.model_name, + ) + ) + + +async def replay_restored_system_messages( + *, + runtime: WebSocketChatRuntime, + send_typed: Any, +) -> None: + """Replay persisted system rows for a restored session.""" + try: + from code_puppy.api.db.queries import get_active_messages + + rows = await get_active_messages(runtime.session_id) + system_rows = [ + r + for r in rows + if r.get("role") == "system" + and r.get("system_message_type") in ("init", "config", "directory") + ] + for sys_row in system_rows: + await send_typed( + ServerSystem( + content=sys_row.get("content", ""), + session_id=runtime.session_id, + agent_name=sys_row.get("agent_name", ""), + model_name=sys_row.get("model_name", ""), + ) + ) + if system_rows: + logger.debug( + "Replayed %d system messages for session %s", + len(system_rows), + runtime.session_id, + ) + except Exception as sys_exc: + logger.warning( + "Failed to replay system messages for session %s: %s", + runtime.session_id, + sys_exc, + ) + + +async def initialize_ws_session( + *, + websocket: WebSocket, + requested_session_id: str | None, + sender: WebSocketSender, + safe_send_json: Any, + send_typed: Any, +) -> WebSocketChatRuntime | None: + """Create or load session state for a chat WebSocket connection.""" + session_id = requested_session_id + existing_history = None + session_title = "" + session_working_directory = "" + session_pinned = False + + if session_id: + try: + _validate_session_id(session_id) + except ValueError as exc: + logger.warning("Invalid session_id rejected: %r: %s", session_id, exc) + await websocket.close(code=1008, reason="Invalid session ID") + return None + + logger.debug("Client requested session: %s", session_id) + try: + from code_puppy.api.db.queries import get_session_metadata, session_exists + + if await session_exists(session_id): + existing_history = True + db_meta = await get_session_metadata(session_id) or {} + session_title = db_meta.get("title", "") + session_working_directory = db_meta.get("working_directory", "") + session_pinned = bool(db_meta.get("pinned", False)) + except Exception as exc: + logger.warning("Failed to check session in SQLite: %s", exc) + + if not existing_history: + logger.debug("Session %s not found in SQLite, will create new", session_id) + else: + session_timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + session_id = f"WS_session_{session_timestamp}" + sender.session_id = session_id + logger.debug("Generated new session ID: %s", session_id) + + runtime = WebSocketChatRuntime( + session_id=session_id, + session_title=session_title, + session_working_directory=session_working_directory, + session_pinned=session_pinned, + existing_history=existing_history, + ) + sender.session_id = session_id + + try: + if existing_history is not None: + runtime.ctx = await session_manager.get_or_load_session(session_id) + sender.ctx = runtime.ctx + if runtime.ctx is None: + logger.warning( + "Session %s exists in SQLite but could not be loaded " + "(get_or_load_session returned None). Starting a blank session " + "to keep the connection alive.", + session_id, + ) + runtime.ctx = await session_manager.create_session(session_id) + sender.ctx = runtime.ctx + else: + runtime.ctx = await session_manager.create_session(session_id) + sender.ctx = runtime.ctx + except Exception as exc: + logger.warning("SessionManager init failed, falling back: %s", exc) + try: + runtime.ctx = await session_manager.create_session(session_id) + sender.ctx = runtime.ctx + except Exception: + logger.error("SessionManager fallback also failed", exc_info=True) + await websocket.close(code=1011, reason="Session init failed") + return None + + runtime.sync_from_ctx() + if not runtime.session_title: + runtime.session_title = session_title + if not runtime.session_working_directory: + runtime.session_working_directory = session_working_directory + runtime.session_pinned = bool(runtime.session_pinned or session_pinned) + + await session_manager.mark_session_active(runtime.session_id) + init_session_process_tracking() + + try: + bus = get_message_bus() + bus.set_session_context(runtime.session_id) + except Exception: + logger.debug("MessageBus session context not available") + + if existing_history is None: + try: + now_iso = datetime.datetime.now(datetime.timezone.utc).isoformat() + init_agent = runtime.agent_name or "code-puppy" + init_model = runtime.model_name or "unknown" + await write_system_message_to_sqlite( + session_id=runtime.session_id, + system_message_type="config", + content=f"🐶 Started with {init_agent} ({init_model})", + agent_name=init_agent, + model_name=init_model, + timestamp=now_iso, + ) + if runtime.session_working_directory: + path_segments = runtime.session_working_directory.split("/")[-3:] + relative = "/".join(s for s in path_segments if s) + await write_system_message_to_sqlite( + session_id=runtime.session_id, + system_message_type="directory", + content=f"Starting in {relative}", + system_message_path=runtime.session_working_directory, + agent_name=init_agent, + model_name=init_model, + timestamp=now_iso, + ) + except Exception as init_exc: + logger.warning( + "Failed to write session init to SQLite: %s", + init_exc, + exc_info=True, + ) + + await send_typed( + ServerSystem( + content=f"Connected! Session: {runtime.session_id}", + session_id=runtime.session_id, + agent_name=runtime.agent_name, + model_name=runtime.model_name, + resumed=existing_history is not None, + protocol_version=PROTOCOL_VERSION, + ) + ) + await send_session_meta_snapshot(runtime=runtime, safe_send_json=safe_send_json) + + if existing_history and runtime.ctx: + try: + loaded_messages = runtime.ctx.agent.get_message_history() + message_count = len(loaded_messages) if loaded_messages else 0 + await send_typed( + ServerSessionRestored( + session_id=runtime.session_id, + message_count=message_count, + title=runtime.session_title, + ui_metadata=[], + ) + ) + await replay_restored_system_messages( + runtime=runtime, send_typed=send_typed + ) + runtime.agent = runtime.ctx.agent + logger.debug("Restored %d messages to session agent", message_count) + except Exception as exc: + logger.warning("Failed to restore session history: %s", exc) + + return runtime + + +__all__ = [ + "initialize_ws_session", + "replay_restored_system_messages", + "send_session_meta_snapshot", +] diff --git a/code_puppy/api/ws/ws_stream_drain.py b/code_puppy/api/ws/ws_stream_drain.py new file mode 100644 index 000000000..5e8ce0ad1 --- /dev/null +++ b/code_puppy/api/ws/ws_stream_drain.py @@ -0,0 +1,95 @@ +"""Helpers for WebSocket stream-drain lifecycle management.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from typing import Any, Awaitable, Callable + + +@dataclass(slots=True) +class StreamDrainHandle: + """Runtime state for one active frontend-emitter drain subscription.""" + + event_queue: Any + task: asyncio.Task + unsubscribe: Callable[[Any], None] + + +async def start_stream_drain( + *, + session_id: str, + drain_coro_factory: Callable[[Any, asyncio.Event], Awaitable[None]], + logger: Any, +) -> StreamDrainHandle | None: + """Subscribe to the frontend emitter and start the drain task.""" + try: + from code_puppy.plugins.frontend_emitter.emitter import ( + subscribe, + unsubscribe, + ) + except ImportError: + logger.warning("Frontend emitter not available") + return None + + event_queue = subscribe(session_id=session_id) + logger.debug("Subscribed to frontend emitter for streaming") + + drain_ready = asyncio.Event() + + async def drain_events_with_signal() -> None: + await drain_coro_factory(event_queue, drain_ready) + + drain_task = asyncio.create_task(drain_events_with_signal()) + await drain_ready.wait() + return StreamDrainHandle( + event_queue=event_queue, + task=drain_task, + unsubscribe=unsubscribe, + ) + + +async def stop_stream_drain( + *, + handle: StreamDrainHandle | None, + stop_draining: asyncio.Event, + logger: Any, +) -> None: + """Stop and unsubscribe one active stream-drain lifecycle.""" + if handle is None: + return + + stop_draining.set() + try: + await asyncio.wait_for(handle.task, timeout=2.0) + except asyncio.TimeoutError: + handle.task.cancel() + try: + await handle.task + except asyncio.CancelledError: + pass + + handle.unsubscribe(handle.event_queue) + logger.debug("Unsubscribed from frontend emitter") + + +async def cancel_active_streaming( + *, + active_drain_task: asyncio.Task | None, + stop_draining: asyncio.Event, + logger: Any, + log_message: str | None = None, +) -> None: + """Cancel the current drain task and reset stop state for reuse.""" + if not active_drain_task or active_drain_task.done(): + return + + stop_draining.set() + active_drain_task.cancel() + try: + await active_drain_task + except asyncio.CancelledError: + pass + stop_draining.clear() + if log_message: + logger.debug(log_message) diff --git a/code_puppy/api/ws/ws_turn_finalization.py b/code_puppy/api/ws/ws_turn_finalization.py new file mode 100644 index 000000000..18fcdae1c --- /dev/null +++ b/code_puppy/api/ws/ws_turn_finalization.py @@ -0,0 +1,218 @@ +"""Post-turn WebSocket finalization helpers. + +This module owns the stateful orchestration that happens *after* the agent run +completes but *before* persistence/broadcast work begins: + +1. Emit B1 tool results before ``stream_end`` when possible. +2. Sync ``result.all_messages()`` onto the agent. +3. Snapshot message history before await points. +4. Emit any remaining tool results from finalized history while suppressing + duplicates already delivered before ``stream_end``. + +The caller remains responsible for the surrounding frame ordering, persistence, +and client/broadcast sends. +""" + +from __future__ import annotations + +import time as time_module +from dataclasses import dataclass, field +from typing import Any, Awaitable, Callable, Iterable + +from code_puppy.api.ws.schemas import ServerToolResult + + +@dataclass(slots=True) +class TurnFinalizationResult: + """Result of post-run history/tool finalization.""" + + pre_sent_tool_ids: set[str] = field(default_factory=set) + history_snapshot: list[Any] = field(default_factory=list) + + +_SIMPLE_TYPES = (str, dict, list, int, float, bool, type(None)) + + +def _serialize_tool_result(value: Any) -> Any: + """Coerce tool result content into JSON-ish payloads safely.""" + if isinstance(value, _SIMPLE_TYPES): + return value + try: + if hasattr(value, "model_dump"): + return value.model_dump() + if hasattr(value, "dict"): + return value.dict() + if hasattr(value, "__dict__"): + return value.__dict__ + except Exception: + pass + return str(value) + + +def _iter_tool_returns(messages: Iterable[Any]) -> Iterable[Any]: + """Yield ToolReturn-like parts from model messages when available.""" + try: + from pydantic_ai.messages import ToolReturn, ToolReturnPart + except Exception: + return [] + + parts_found: list[Any] = [] + for msg in messages: + if not hasattr(msg, "parts"): + continue + for part in msg.parts: + if isinstance(part, (ToolReturnPart, ToolReturn)): + parts_found.append(part) + return parts_found + + +async def emit_pre_stream_end_tool_results( + *, + result: Any, + turn_state: Any, + session_id: str, + agent_name: str, + model_name: str, + send_typed_tool_lifecycle: Callable[[Any], Awaitable[None]], + logger: Any, +) -> set[str]: + """Emit B1 tool results before ``stream_end`` and return sent tool IDs.""" + sent_tool_ids: set[str] = set() + if not result or not hasattr(result, "all_messages"): + return sent_tool_ids + + try: + messages = list(result.all_messages()) + for part in _iter_tool_returns(messages): + tool_name = getattr(part, "tool_name", "unknown") + raw_tool_id = getattr(part, "tool_call_id", None) + tool_id = ( + turn_state.tool_id_aliases.get(raw_tool_id, raw_tool_id) + if raw_tool_id + else "unknown" + ) + result_payload = _serialize_tool_result(getattr(part, "content", None)) + logger.warning( + "[WebSocket] Pre-stream_end tool result: %s (id: %s), content_preview=%s", + tool_name, + tool_id, + str(result_payload)[:100] if result_payload else None, + ) + await send_typed_tool_lifecycle( + ServerToolResult( + tool_id=tool_id, + tool_name=tool_name, + result=result_payload, + success=True, + duration_ms=0, + timestamp=time_module.time(), + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + tool_group_id=turn_state.tool_group_ids.get(tool_id), + ) + ) + sent_tool_ids.add(tool_id) + except Exception as exc: + logger.warning( + "Pre-stream_end tool result extraction failed: %s", + exc, + ) + + return sent_tool_ids + + +async def finalize_turn_history( + *, + result: Any, + agent: Any, + turn_state: Any, + session_id: str, + agent_name: str, + model_name: str, + send_typed: Callable[[Any], Awaitable[None]], + pre_sent_tool_ids: set[str] | None, + logger: Any, +) -> TurnFinalizationResult: + """Sync history from ``result.all_messages()`` and emit remaining tool results.""" + finalized = TurnFinalizationResult( + pre_sent_tool_ids=set(pre_sent_tool_ids or set()), + history_snapshot=[], + ) + + if not result or not hasattr(result, "all_messages"): + return finalized + + try: + all_msgs = list(result.all_messages()) + if all_msgs: + agent.set_message_history(all_msgs) + finalized.history_snapshot = list(agent.get_message_history()) + logger.debug( + "Updated message history from result.all_messages(): %s messages", + len(all_msgs), + ) + + try: + for part in _iter_tool_returns(all_msgs): + tool_name = getattr(part, "tool_name", "unknown") + tool_call_id = getattr(part, "tool_call_id", "unknown") + result_payload = _serialize_tool_result(getattr(part, "content", None)) + + if tool_name == "agent_run_shell_command": + stdout_val = ( + result_payload.get("stdout", "N/A") + if isinstance(result_payload, dict) + else "not dict" + ) + logger.debug( + "Extracted shell result: id=%s, stdout=%s", + tool_call_id, + stdout_val, + ) + + if tool_call_id in finalized.pre_sent_tool_ids: + logger.debug( + "[WebSocket] Skipping duplicate tool result (pre-sent): %s", + tool_call_id, + ) + continue + + logger.info( + "[WebSocket] Sending extracted tool result for %s (id: %s)", + tool_name, + tool_call_id, + ) + await send_typed( + ServerToolResult( + tool_id=tool_call_id, + tool_name=tool_name, + result=result_payload, + success=True, + duration_ms=0, + timestamp=time_module.time(), + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + tool_group_id=turn_state.tool_group_ids.get(tool_call_id), + ) + ) + except Exception as exc: + logger.warning( + "Could not extract tool results from messages: %s", + exc, + ) + except Exception as exc: + logger.warning( + "Could not update history from result.all_messages(): %s", + exc, + ) + + return finalized + + +__all__ = [ + "TurnFinalizationResult", + "emit_pre_stream_end_tool_results", + "finalize_turn_history", +] diff --git a/code_puppy/api/ws/ws_turn_preparation.py b/code_puppy/api/ws/ws_turn_preparation.py new file mode 100644 index 000000000..de6aa307e --- /dev/null +++ b/code_puppy/api/ws/ws_turn_preparation.py @@ -0,0 +1,127 @@ +"""Helpers for preparing one WebSocket chat turn before agent execution.""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from code_puppy.api.ws.attachments import build_file_context_and_attachments + +logger = logging.getLogger(__name__) + + +@dataclass(slots=True) +class PreparedTurnInput: + """Materialized input for one agent turn.""" + + message_to_send: str + run_kwargs: dict[str, Any] + attachment_metadata: list[dict[str, Any]] + last_context_sent_directory: str + + +def _maybe_inject_working_directory( + *, + agent: Any, + session_working_directory: str, + last_context_sent_directory: str, +) -> str: + """Append a working-directory system message when the visible CWD changed.""" + if ( + not session_working_directory + or session_working_directory == last_context_sent_directory + ): + return last_context_sent_directory + + from pydantic_ai.messages import ModelRequest, SystemPromptPart + + wd_system_msg = ModelRequest( + parts=[ + SystemPromptPart( + content=( + "The user's current working directory is updated to" + f" {session_working_directory}" + ) + ) + ] + ) + agent.append_to_message_history(wd_system_msg) + logger.debug( + "Injected working directory system message: %s", + session_working_directory, + ) + return session_working_directory + + +def _collect_attachment_metadata(msg: dict[str, Any]) -> list[dict[str, Any]]: + """Build attachment metadata payloads for the UI without changing semantics.""" + attachment_metadata: list[dict[str, Any]] = [] + + if not msg.get("attachments"): + return attachment_metadata + + for raw_path in msg.get("attachments", []): + if not isinstance(raw_path, str) or not raw_path.strip(): + continue + + try: + file_path = Path(raw_path) + if file_path.exists(): + attachment_metadata.append( + { + "name": file_path.name, + "path": str(file_path.absolute()), + "sizeBytes": file_path.stat().st_size, + } + ) + except Exception as e: + logger.warning( + "Error building attachment metadata for '%s': %s", + raw_path, + e, + ) + + return attachment_metadata + + +def prepare_turn_input( + *, + agent: Any, + user_message: str, + msg: dict[str, Any], + session_working_directory: str, + last_context_sent_directory: str, +) -> PreparedTurnInput: + """Prepare the exact message payload and metadata for one agent turn.""" + last_context_sent_directory = _maybe_inject_working_directory( + agent=agent, + session_working_directory=session_working_directory, + last_context_sent_directory=last_context_sent_directory, + ) + + message_to_send = user_message + logger.debug("Calling run_with_mcp with message: %s...", message_to_send[:100]) + + file_context, binary_attachments = build_file_context_and_attachments(msg) + attachment_metadata = _collect_attachment_metadata(msg) + + if file_context: + message_to_send = file_context + "\n\n" + message_to_send + logger.debug("Added file context (%d chars)", len(file_context)) + + run_kwargs: dict[str, Any] = {} + if binary_attachments: + run_kwargs["attachments"] = binary_attachments + logger.debug( + "Including %d binary attachment(s)", + len(binary_attachments), + ) + + return PreparedTurnInput( + message_to_send=message_to_send, + run_kwargs=run_kwargs, + attachment_metadata=attachment_metadata, + last_context_sent_directory=last_context_sent_directory, + ) diff --git a/code_puppy/chatgpt_codex_client.py b/code_puppy/chatgpt_codex_client.py index 7870e48e1..aff7206b1 100644 --- a/code_puppy/chatgpt_codex_client.py +++ b/code_puppy/chatgpt_codex_client.py @@ -218,6 +218,30 @@ def _looks_like_unpersisted_reference(it: dict) -> bool: item["id"] = f"rs_{item_id}" modified = True + # De-duplicate input items by `id` to satisfy Responses API validation. + # Resumed-session reconstruction can occasionally surface repeated + # reasoning item IDs (`rs_...`), which the backend rejects with HTTP 400. + input_items = data.get("input") + if isinstance(input_items, list): + seen_ids: set[str] = set() + deduped: list[object] = [] + removed = 0 + for item in input_items: + if not isinstance(item, dict): + deduped.append(item) + continue + item_id = item.get("id") + if isinstance(item_id, str) and item_id: + if item_id in seen_ids: + removed += 1 + modified = True + continue + seen_ids.add(item_id) + deduped.append(item) + if removed: + logger.debug("Dropped %d duplicate input items by id", removed) + data["input"] = deduped + # Remove unsupported parameters # Note: verbosity should be under "text" object, not top-level unsupported_params = ["max_output_tokens", "max_tokens", "verbosity"] diff --git a/code_puppy/cli_runner.py b/code_puppy/cli_runner.py index ce38fc1fc..876dc41b8 100644 --- a/code_puppy/cli_runner.py +++ b/code_puppy/cli_runner.py @@ -86,7 +86,7 @@ def _render_turn_exception(exc: Exception) -> None: emit_error( f"\U0001f50c The model connection hit a transient error " f"({type(exc).__name__}) and didn't recover after auto-retries. " - "This is almost always a VPN/WiFi/provider blip \u2014 just re-run " + "This is almost always a VPN/WiFi/provider blip — just re-run " "your last prompt. Your session history is intact." ) return @@ -184,7 +184,7 @@ async def main(): "-r", type=str, metavar="PATH", - help="Resume a saved session from a .pkl file (e.g. ~/.code_puppy/contexts/foo.pkl)", + help="Resume a saved session from a .json file (e.g. ~/.code_puppy/contexts/foo.json)", ) parser.add_argument( "--quick-resume", @@ -1065,7 +1065,7 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non agent.estimate_tokens_for_message(msg) for msg in history ) - session_path = base_dir / f"{chosen_session}.pkl" + session_path = base_dir / f"{chosen_session}.json" emit_success( f"✅ Autosave loaded: {len(history)} messages ({total_tokens} tokens)\n" diff --git a/code_puppy/command_line/mcp/start_command.py b/code_puppy/command_line/mcp/start_command.py index 8f9f32998..3ca575c34 100644 --- a/code_puppy/command_line/mcp/start_command.py +++ b/code_puppy/command_line/mcp/start_command.py @@ -10,10 +10,16 @@ from code_puppy.mcp_.agent_bindings import is_bound, set_session_binding from code_puppy.messaging import emit_error, emit_info, emit_success -from ...agents import get_current_agent +from ... import agents from .base import MCPCommandBase from .utils import find_server_id_by_name, suggest_similar_servers + +def get_current_agent(): + """Compatibility wrapper for patching in tests.""" + return agents.get_current_agent() + + # Configure logging logger = logging.getLogger(__name__) diff --git a/code_puppy/command_line/mcp/stop_command.py b/code_puppy/command_line/mcp/stop_command.py index fc1bb0b7c..5b4570653 100644 --- a/code_puppy/command_line/mcp/stop_command.py +++ b/code_puppy/command_line/mcp/stop_command.py @@ -9,10 +9,16 @@ from code_puppy.messaging import emit_error, emit_info -from ...agents import get_current_agent +from ... import agents from .base import MCPCommandBase from .utils import find_server_id_by_name, suggest_similar_servers + +def get_current_agent(): + """Compatibility wrapper for patching in tests.""" + return agents.get_current_agent() + + # Configure logging logger = logging.getLogger(__name__) diff --git a/code_puppy/logging_setup.py b/code_puppy/logging_setup.py new file mode 100644 index 000000000..747546160 --- /dev/null +++ b/code_puppy/logging_setup.py @@ -0,0 +1,223 @@ +"""Backend logging configuration utilities. + +Provides date-based file logging with weekday in filename, stdout logging, +and retention cleanup for backend log files. +""" + +from __future__ import annotations + +import logging +import os +import re +import sys +from datetime import date, datetime, timedelta +from pathlib import Path +from typing import Any + +from code_puppy.config import STATE_DIR + +_BACKEND_FILE_RE = re.compile( + r"^backend-(monday|tuesday|wednesday|thursday|friday|saturday|sunday)-(\d{4}-\d{2}-\d{2})\.log$" +) +_DEFAULT_FORMAT = "%(asctime)s [%(levelname)s] %(name)s: %(message)s" +_DEFAULT_DATEFMT = "%H:%M:%S" + + +class DailyBackendFileHandler(logging.Handler): + """File handler that switches to a new date-based file at midnight. + + Filename format: + backend--YYYY-MM-DD.log + """ + + def __init__(self, log_dir: Path, retention_days: int, encoding: str = "utf-8"): + super().__init__() + self.log_dir = log_dir + self.retention_days = max(0, retention_days) + self.encoding = encoding + self.terminator = "\n" + self._stream: Any = None + self._active_date: date | None = None + self._active_path: Path | None = None + + def _path_for_date(self, d: date) -> Path: + weekday = d.strftime("%A").lower() + return self.log_dir / f"backend-{weekday}-{d.isoformat()}.log" + + def _rotate_if_needed(self, now: datetime) -> None: + today = now.date() + if self._active_date == today and self._stream is not None: + return + + self._close_stream() + self.log_dir.mkdir(parents=True, exist_ok=True) + self._active_path = self._path_for_date(today) + self._stream = self._active_path.open("a", encoding=self.encoding) + self._active_date = today + + # Run cleanup after opening the current day's stream. + cleanup_backend_log_files( + log_dir=self.log_dir, + retention_days=self.retention_days, + now=now, + ) + + def _close_stream(self) -> None: + if self._stream is None: + return + try: + self._stream.flush() + self._stream.close() + finally: + self._stream = None + + def emit(self, record: logging.LogRecord) -> None: + try: + now = datetime.now() + self._rotate_if_needed(now) + if self._stream is None: + return + msg = self.format(record) + self._stream.write(msg + self.terminator) + self._stream.flush() + except Exception: + self.handleError(record) + + def close(self) -> None: + try: + self._close_stream() + finally: + super().close() + + +def _resolve_log_level(level: int | None = None) -> int: + if level is not None: + return level + raw = ( + (os.getenv("BACKEND_LOG_LEVEL") or os.getenv("LOG_LEVEL") or "INFO") + .strip() + .upper() + ) + return getattr(logging, raw, logging.INFO) + + +def _resolve_log_dir(log_dir: str | Path | None = None) -> Path: + if log_dir is not None: + return Path(log_dir) + raw = os.getenv("BACKEND_LOG_DIR") + if raw: + return Path(raw) + return Path(STATE_DIR) / "logs" + + +def _resolve_retention_days(retention_days: int | None = None) -> int: + if retention_days is not None: + return max(0, retention_days) + raw = os.getenv("BACKEND_LOG_RETENTION_DAYS", "7").strip() + try: + return max(0, int(raw)) + except ValueError: + return 7 + + +def cleanup_backend_log_files( + *, + log_dir: str | Path, + retention_days: int, + now: datetime | None = None, +) -> list[Path]: + """Delete backend log files older than retention_days. + + Returns list of removed file paths. + """ + current = now or datetime.now() + keep_days = max(1, retention_days) + cutoff_date = current.date() - timedelta(days=keep_days - 1) + base = Path(log_dir) + if not base.exists(): + return [] + + removed: list[Path] = [] + for path in base.glob("backend-*.log"): + match = _BACKEND_FILE_RE.match(path.name) + if not match: + continue + file_date_raw = match.group(2) + try: + file_date = date.fromisoformat(file_date_raw) + except ValueError: + continue + if file_date < cutoff_date: + try: + path.unlink() + removed.append(path) + except OSError: + # Non-fatal; logging setup should never crash app startup. + continue + return removed + + +def configure_backend_logging( + *, + level: int | None = None, + log_dir: str | Path | None = None, + retention_days: int | None = None, +) -> int: + """Configure root logger with stdout + daily backend file logging. + + Safe to call multiple times: it removes/replaces previously installed + backend handlers to avoid duplicates. + + Returns the resolved integer log level used for backend logging. + """ + resolved_level = _resolve_log_level(level) + resolved_log_dir = _resolve_log_dir(log_dir) + resolved_retention = _resolve_retention_days(retention_days) + + root = logging.getLogger() + root.setLevel(resolved_level) + + # Remove previous backend-managed handlers to avoid duplicates. + for handler in list(root.handlers): + if getattr(handler, "_code_puppy_backend_logging", False): + root.removeHandler(handler) + handler.close() + + formatter = logging.Formatter(_DEFAULT_FORMAT, datefmt=_DEFAULT_DATEFMT) + + stream_handler = logging.StreamHandler(sys.stdout) + stream_handler.setLevel(resolved_level) + stream_handler.setFormatter(formatter) + setattr(stream_handler, "_code_puppy_backend_logging", True) + setattr(stream_handler, "_code_puppy_backend_logging_kind", "stream") + root.addHandler(stream_handler) + + try: + resolved_log_dir.mkdir(parents=True, exist_ok=True) + + file_handler = DailyBackendFileHandler( + log_dir=resolved_log_dir, + retention_days=resolved_retention, + ) + file_handler.setLevel(resolved_level) + file_handler.setFormatter(formatter) + setattr(file_handler, "_code_puppy_backend_logging", True) + setattr(file_handler, "_code_puppy_backend_logging_kind", "file") + root.addHandler(file_handler) + + # Also run cleanup at configuration time. + cleanup_backend_log_files( + log_dir=resolved_log_dir, + retention_days=resolved_retention, + ) + except Exception as exc: + print( + f"[code_puppy.logging_setup] backend file logging disabled: {exc}", + file=sys.stderr, + ) + root.warning( + "Backend file logging disabled; continuing with stdout logging only.", + exc_info=True, + ) + + return resolved_level diff --git a/code_puppy/messaging/bus.py b/code_puppy/messaging/bus.py index f53206ad8..afcd266bc 100644 --- a/code_puppy/messaging/bus.py +++ b/code_puppy/messaging/bus.py @@ -35,11 +35,13 @@ import asyncio import queue import threading +from contextvars import ContextVar from typing import Any, Dict, List, Optional, Tuple from uuid import uuid4 from .commands import ( AnyCommand, + AskUserQuestionResponse, ConfirmationResponse, PauseAgentCommand, ResumeAgentCommand, @@ -55,6 +57,7 @@ SelectionRequest, TextMessage, UserInputRequest, + AskUserQuestionRequest, ) @@ -89,8 +92,11 @@ def __init__(self, maxsize: int = 1000) -> None: # Request/Response correlation: prompt_id → Future (for async usage) self._pending_requests: Dict[str, asyncio.Future[Any]] = {} - # Session context for multi-agent tracking - self._current_session_id: Optional[str] = None + # Session context for multi-agent tracking. This is task/thread-local so + # concurrent websocket sessions and CLI flows cannot overwrite each other. + self._session_context: ContextVar[Optional[str]] = ContextVar( + f"message_bus_session_id_{id(self)}", default=None + ) # ========================================================================= # Outgoing Messages (Agent → UI) @@ -108,8 +114,9 @@ def emit(self, message: AnyMessage) -> None: """ # Auto-tag message with current session if not already set with self._lock: - if message.session_id is None and self._current_session_id is not None: - message.session_id = self._current_session_id + current_session_id = self._session_context.get() + if message.session_id is None and current_session_id is not None: + message.session_id = current_session_id if not self._has_active_renderer: self._startup_buffer.append(message) @@ -190,8 +197,7 @@ def set_session_context(self, session_id: Optional[str]) -> None: Args: session_id: The session ID to tag messages with, or None to clear. """ - with self._lock: - self._current_session_id = session_id + self._session_context.set(session_id) def get_session_context(self) -> Optional[str]: """Get the current session context. @@ -199,8 +205,7 @@ def get_session_context(self) -> Optional[str]: Returns: The current session_id, or None if not set. """ - with self._lock: - return self._current_session_id + return self._session_context.get() # ========================================================================= # User Input Requests (Agent waits for UI response) @@ -335,6 +340,43 @@ async def request_selection( with self._lock: self._pending_requests.pop(prompt_id, None) + async def request_ask_user_question( + self, + questions: list[dict[str, object]], + timeout_seconds: int = 300, + ) -> tuple[list[dict[str, object]], bool]: + """Request structured ask_user_question answers from the UI. + + Emits an AskUserQuestionRequest and waits for the browser to return the + full answer set. This is used by Desk/WebSocket sessions so the tool can + bypass terminal-only TUI input. + + Returns: + Tuple of (answers, cancelled). + """ + prompt_id = str(uuid4()) + + loop = asyncio.get_running_loop() + future: asyncio.Future[tuple[list[dict[str, object]], bool]] = ( + loop.create_future() + ) + + with self._lock: + self._pending_requests[prompt_id] = future + + request = AskUserQuestionRequest( + prompt_id=prompt_id, + questions=questions, + timeout_seconds=timeout_seconds, + ) + self.emit(request) + + try: + return await future + finally: + with self._lock: + self._pending_requests.pop(prompt_id, None) + # ========================================================================= # Incoming Commands (UI → Agent) # ========================================================================= @@ -359,6 +401,10 @@ def provide_response(self, command: AnyCommand) -> None: self._complete_request( command.prompt_id, (command.selected_index, command.selected_value) ) + elif isinstance(command, AskUserQuestionResponse): + self._complete_request( + command.prompt_id, (command.answers, command.cancelled) + ) elif isinstance(command, PauseAgentCommand): from .pause_controller import get_pause_controller diff --git a/code_puppy/messaging/commands.py b/code_puppy/messaging/commands.py index b6c3473c0..25c09c907 100644 --- a/code_puppy/messaging/commands.py +++ b/code_puppy/messaging/commands.py @@ -21,7 +21,7 @@ """ from datetime import datetime, timezone -from typing import Literal, Optional, Union +from typing import Any, Literal, Optional, Union from uuid import uuid4 from pydantic import BaseModel, Field @@ -175,6 +175,26 @@ class SelectionResponse(BaseCommand): selected_value: str = Field(description="The value of the selected option") +class AskUserQuestionResponse(BaseCommand): + """Response to an AskUserQuestionRequest from the browser UI. + + The payload mirrors the ask_user_question tool output shape so the tool + can return structured answers to the agent without using the terminal TUI. + """ + + prompt_id: str = Field( + description="ID of the prompt this responds to (must match request)" + ) + answers: list[dict[str, Any]] = Field( + default_factory=list, + description="Structured answers collected by the browser UI", + ) + cancelled: bool = Field( + default=False, + description="Whether the interaction was cancelled/skipped", + ) + + # ============================================================================= # Union Type for Type Checking # ============================================================================= @@ -190,6 +210,7 @@ class SelectionResponse(BaseCommand): UserInputResponse, ConfirmationResponse, SelectionResponse, + AskUserQuestionResponse, ] """Union of all command types for type checking.""" @@ -211,6 +232,7 @@ class SelectionResponse(BaseCommand): "UserInputResponse", "ConfirmationResponse", "SelectionResponse", + "AskUserQuestionResponse", # Union type "AnyCommand", ] diff --git a/code_puppy/messaging/messages.py b/code_puppy/messaging/messages.py index 27a17bc73..93ad0f2b1 100644 --- a/code_puppy/messaging/messages.py +++ b/code_puppy/messaging/messages.py @@ -389,6 +389,20 @@ class SelectionRequest(BaseMessage): ) +class AskUserQuestionRequest(BaseMessage): + """Request for browser UI to collect structured ask_user_question answers.""" + + category: MessageCategory = MessageCategory.USER_INTERACTION + prompt_id: str = Field(description="Unique ID for matching responses to requests") + questions: List[Dict[str, object]] = Field( + description="Validated ask_user_question question payloads" + ) + timeout_seconds: int = Field( + default=300, + description="Inactivity timeout in seconds", + ) + + # ============================================================================= # Control Messages # ============================================================================= @@ -511,6 +525,7 @@ class SkillActivateMessage(BaseMessage): UserInputRequest, ConfirmationRequest, SelectionRequest, + AskUserQuestionRequest, SpinnerControl, DividerMessage, StatusPanelMessage, diff --git a/code_puppy/model_errors.py b/code_puppy/model_errors.py new file mode 100644 index 000000000..87ea86690 --- /dev/null +++ b/code_puppy/model_errors.py @@ -0,0 +1,277 @@ +"""Model error normalization utilities. + +Central place to parse provider-specific model errors (Anthropic, OpenAI, +Gemini, etc.) into a small, stable structure that the rest of the +application can use for retries and user-facing error messages. + +This keeps provider-specific logic out of the core agent code. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Optional + + +@dataclass +class NormalizedModelError: + """Provider-agnostic view of a model error.""" + + provider: Optional[str] + code: Optional[str] + http_status: Optional[int] + is_transient: bool + user_message: str + raw_message: str + + +def _get_http_status(exc: Exception) -> Optional[int]: + """Best-effort extraction of an HTTP status code from an exception.""" + + for attr in ("status_code", "status", "http_status", "code"): + value = getattr(exc, attr, None) + if isinstance(value, int) and 100 <= value <= 599: + return value + return None + + +def _get_body_dict(exc: Exception) -> Optional[dict[str, Any]]: + """Try to extract a structured body dict from provider exceptions.""" + + for attr in ("body", "response", "error", "errors", "detail"): + value = getattr(exc, attr, None) + if isinstance(value, dict): + return value + return None + + +def _detect_provider(exc: Exception) -> Optional[str]: + """Infer model provider from the exception's module / type name.""" + + module = getattr(exc.__class__, "__module__", "") + name = exc.__class__.__name__.lower() + module_lower = module.lower() + + if "anthropic" in module_lower or "anthropic" in name: + return "anthropic" + if "openai" in module_lower or "openai" in name: + return "openai" + if "google" in module_lower or "gemini" in module_lower or "gemini" in name: + return "gemini" + if "azure" in module_lower or "azure" in name: + return "azure-openai" + + return None + + +def _classify_message( + provider: Optional[str], + raw: str, + payload: Optional[dict[str, Any]], + http_status: Optional[int], +) -> tuple[Optional[str], bool]: + """Map provider error details to a (code, is_transient) pair. + + The goal is to keep this classification small and robust, not to + perfectly mirror each provider's error taxonomy. + """ + + lower = raw.lower() + + # Rate limiting / overloaded / transient backend issues + if any( + k in lower + for k in ("rate limit", "too many requests", "overloaded", "try again later") + ): + return "rate_limit_or_overloaded", True + + if any( + k in lower + for k in ( + "timeout", + "temporarily unavailable", + "server error", + "gateway", + "bad gateway", + ) + ): + return "backend_unavailable", True + + # HTTP status based classification (when available) + if http_status in (429, 509): + return "rate_limit_or_overloaded", True + if http_status in (500, 502, 503, 504): + return "backend_unavailable", True + + # Auth / configuration + if any( + k in lower + for k in ( + "invalid api key", + "authentication", + "unauthorized", + "forbidden", + "permission", + ) + ): + return "auth_error", False + + # Model / quota + if any( + k in lower + for k in ( + "model_not_found", + "unknown model", + "unsupported model", + "no such model", + ) + ): + return "unsupported_model", False + + if any( + k in lower for k in ("quota", "billing", "insufficient balance", "subscription") + ): + return "quota_exceeded", False + + # Safety / content policy + if any( + k in lower + for k in ("safety", "content policy", "blocked content", "policy violation") + ): + return "content_blocked", False + + # Tool history corruption (Anthropic 400: orphaned or duplicated tool blocks). + # These are non-transient — the history must be pruned before retrying. + # Two distinct error shapes are handled: + # 1. Orphaned tool_result: "unexpected tool_use_id" / "each tool_result block + # must have a corresponding tool_use" — a tool_result with no matching tool_use. + # 2. Duplicate tool_result: "each tool_use must have a single result. Found + # multiple tool_result blocks with id: ..." — happens when a streaming session + # is interrupted mid-tool-call and the user resumes with "continue". + if any( + k in lower + for k in ( + "unexpected tool_use_id", + "unexpected tool_result", + "tool_use ids found without tool_result", + "each tool_result block must have a corresponding tool_use", + "each tool_use must have a single result", + "found multiple `tool_result` blocks", + ) + ): + return "invalid_tool_history", False + + # Default: unknown and non-transient + return None, False + + +def _build_user_message( + provider: Optional[str], + code: Optional[str], + is_transient: bool, + raw: str, +) -> str: + """Return a short, user-facing message based on normalized fields.""" + + prov = provider or "The model" + + if code == "rate_limit_or_overloaded": + return ( + f"{prov} is temporarily overloaded or rate-limited. " + "You can try again in a bit, or switch models with /model." + ) + + if code == "backend_unavailable": + return ( + f"{prov} is currently unavailable due to a server issue. " + "Please try again shortly or switch models with /model." + ) + + if code == "auth_error": + return ( + f"{prov} rejected this request due to an authentication or " + "permission error. Check your API key and configuration." + ) + + if code == "unsupported_model": + return ( + f"The requested model is not available for {prov}. " + "Try switching to a different model with /model." + ) + + if code == "quota_exceeded": + return ( + f"{prov} reports that your quota or billing limits have been " + "exceeded. Check your usage or choose a different provider/model." + ) + + if code == "content_blocked": + return ( + f"{prov} blocked this request due to safety or content policy " + "restrictions. Try rephrasing the request." + ) + + if code == "invalid_tool_history": + return ( + "The conversation history has mismatched tool calls. " + "Code Puppy is attempting to repair it automatically. " + "If this keeps happening, starting a new session will resolve it." + ) + + # Fallback: generic but informative + short_raw = raw.strip().split("\n", 1)[0] + if len(short_raw) > 200: + short_raw = short_raw[:197] + "..." + + return f"The model returned an unexpected error: {short_raw}" + + +def normalize_model_error(exc: Exception) -> NormalizedModelError: + """Normalize provider-specific model errors into a common structure. + + This is intentionally conservative: when we cannot confidently + classify an error, we mark it as non-transient and provide a + generic user-facing message. + + Classification uses the combined text from: + - ``str(exc)`` — the primary error representation. + - String-valued ``.message``, ``.body``, and ``.detail`` attributes — + Anthropic's ``APIStatusError`` surfaces the HTTP body as ``.message``. + - The exception's ``__cause__`` chain. + """ + + raw = str(exc) if exc is not None else "" + + # Collect additional candidate text from well-known string attributes so + # that providers that embed the error body in a non-str attribute (e.g. + # Anthropic's APIStatusError.message) are also classified correctly. + parts: list[str] = [raw] + for attr in ("message", "body", "detail"): + val = getattr(exc, attr, None) + if val and isinstance(val, str) and val != raw: + parts.append(val) + elif val and not isinstance(val, str): + try: + parts.append(str(val)) + except Exception: # noqa: BLE001 + pass + cause = getattr(exc, "__cause__", None) + if cause is not None: + parts.append(str(cause)) + combined_raw = " ".join(parts) + + provider = _detect_provider(exc) + http_status = _get_http_status(exc) + payload = _get_body_dict(exc) + + code, is_transient = _classify_message(provider, combined_raw, payload, http_status) + user_message = _build_user_message(provider, code, is_transient, combined_raw) + + return NormalizedModelError( + provider=provider, + code=code, + http_status=http_status, + is_transient=is_transient, + user_message=user_message, + raw_message=raw, + ) diff --git a/code_puppy/plugins/frontend_emitter/emitter.py b/code_puppy/plugins/frontend_emitter/emitter.py index 55836d908..cd0ffab69 100644 --- a/code_puppy/plugins/frontend_emitter/emitter.py +++ b/code_puppy/plugins/frontend_emitter/emitter.py @@ -234,17 +234,20 @@ def unsubscribe(queue: "asyncio.Queue[Dict[str, Any]]") -> None: ) -def get_recent_events() -> List[Dict[str, Any]]: +def get_recent_events(session_id: Optional[str] = None) -> List[Dict[str, Any]]: """Get recent events for new subscribers. - Returns a copy of the most recent events (up to - ``frontend_emitter_max_recent_events``). Useful for letting new - WebSocket connections "catch up" on recent activity. + Args: + session_id: Optional session filter. If provided, only events whose + ``event["session_id"] == session_id`` are returned. Returns: A list of recent event dictionaries. """ - return _recent_events.copy() + events = _recent_events.copy() + if session_id is None: + return events + return [e for e in events if e.get("session_id") == session_id] def get_subscriber_count() -> int: diff --git a/code_puppy/tools/ask_user_question/constants.py b/code_puppy/tools/ask_user_question/constants.py index 98bfc4525..29d8ffff0 100644 --- a/code_puppy/tools/ask_user_question/constants.py +++ b/code_puppy/tools/ask_user_question/constants.py @@ -5,7 +5,7 @@ # Question constraints MAX_QUESTIONS_PER_CALL: Final[int] = 10 # Reasonable limit for a single TUI interaction MIN_OPTIONS_PER_QUESTION: Final[int] = 2 -MAX_OPTIONS_PER_QUESTION: Final[int] = 6 +MAX_OPTIONS_PER_QUESTION: Final[int] = 12 MAX_HEADER_LENGTH: Final[int] = ( 25 # Must fit in left panel (MAX_LEFT_PANEL_WIDTH - padding) ) diff --git a/code_puppy/tools/ask_user_question/handler.py b/code_puppy/tools/ask_user_question/handler.py index 881892c15..2f82b3399 100644 --- a/code_puppy/tools/ask_user_question/handler.py +++ b/code_puppy/tools/ask_user_question/handler.py @@ -22,6 +22,75 @@ from .terminal_ui import CancelledException, interactive_question_picker +async def ask_user_question_async( + questions: list[Question | dict[str, Any]], + timeout: int = DEFAULT_TIMEOUT_SECONDS, +) -> AskUserQuestionOutput: + """Async entry point used by WebSocket-capable agent runtimes. + + Desk/WebSocket sessions have an active MessageBus renderer, so we bypass + the terminal TUI and wait for the browser to submit structured answers. + Non-WS sessions retain the original terminal behavior in a worker thread. + """ + if is_subagent(): + return AskUserQuestionOutput.error_response( + "Interactive tools are disabled for sub-agents. " + "Sub-agents should make reasonable decisions or return to the parent agent " + "if user input is required." + ) + + try: + validated_input = _validate_input(questions) + except ValidationError as e: + return AskUserQuestionOutput.error_response(_format_validation_error(e)) + except (TypeError, ValueError) as e: + return AskUserQuestionOutput.error_response(f"Validation error: {e!s}") + + if is_wiggum_active(): + return AskUserQuestionOutput.error_response( + "Interactive tools are disabled during /wiggum mode. " + "The agent is running autonomously in a loop. " + "Make a reasonable decision to proceed, or stop and wait for user input " + "by completing the current task." + ) + + browser_result = await _run_browser_picker_async(validated_input.questions, timeout) + if browser_result is not None: + return browser_result + + return await asyncio.to_thread(ask_user_question, questions, timeout) + + +async def _run_browser_picker_async( + questions: list[Question], timeout: int +) -> AskUserQuestionOutput | None: + """Collect ask_user_question answers through the active MessageBus renderer. + + Returns None when no renderer is active so callers can fall back to the + terminal TUI. + """ + try: + from code_puppy.messaging.bus import get_message_bus + + bus = get_message_bus() + if not bus.has_active_renderer: + return None + + payload = [q.model_dump(mode="json") for q in questions] + answers_payload, cancelled = await asyncio.wait_for( + bus.request_ask_user_question(payload, timeout_seconds=timeout), + timeout=max(timeout, 1) + 5, + ) + if cancelled: + return AskUserQuestionOutput.cancelled_response() + answers = [QuestionAnswer.model_validate(answer) for answer in answers_payload] + return AskUserQuestionOutput(answers=answers) + except asyncio.TimeoutError: + return AskUserQuestionOutput.timeout_response(timeout) + except Exception as e: + return AskUserQuestionOutput.error_response(f"Browser interaction error: {e!s}") + + class AsyncContextError(RuntimeError): """Raised when TUI is called from async context without await.""" @@ -144,7 +213,7 @@ def ask_user_question( except (CancelledException, KeyboardInterrupt): return _cancelled_response() - except OSError as e: + except (OSError, EOFError) as e: return AskUserQuestionOutput.error_response(f"Interaction error: {e!s}") diff --git a/code_puppy/tools/ask_user_question/models.py b/code_puppy/tools/ask_user_question/models.py index 7ee95a640..9c58bbd79 100644 --- a/code_puppy/tools/ask_user_question/models.py +++ b/code_puppy/tools/ask_user_question/models.py @@ -3,7 +3,7 @@ from __future__ import annotations import re -from typing import TYPE_CHECKING, Annotated, Any +from typing import TYPE_CHECKING, Annotated, Any, Literal from pydantic import BaseModel, BeforeValidator, Field, model_validator @@ -66,6 +66,11 @@ def sanitize(v: Any) -> str: _sanitize_required = _make_sanitizer(allow_none=False) _sanitize_optional = _make_sanitizer(allow_none=True, default="") +# Legacy direct-construction guard used by unit tests and terminal-oriented callers. +# Tool payload dictionaries (validated through AskUserQuestionInput) can carry +# up to MAX_OPTIONS_PER_QUESTION options for browser-first workflows. +_DIRECT_MODEL_MAX_OPTIONS = 6 + def _sanitize_header(v: Any) -> str: """Sanitize header: remove ANSI, strip, replace spaces with hyphens.""" @@ -116,7 +121,9 @@ class Question(BaseModel): question: The full question text displayed to the user header: Short label used for compact display and response mapping multi_select: Whether user can select multiple options - options: List of 2-6 selectable options + options: List of selectable options. Usually 2-12 choices; may be empty + when ``input_mode`` is ``"text"``. + input_mode: Whether the answer is option-based, free-form text, or both """ question: Annotated[ @@ -147,16 +154,73 @@ class Question(BaseModel): options: Annotated[ list[QuestionOption], Field( - min_length=MIN_OPTIONS_PER_QUESTION, + default_factory=list, + min_length=0, max_length=MAX_OPTIONS_PER_QUESTION, - description="Array of 2-6 selectable options", + description=( + f"Array of 0-{MAX_OPTIONS_PER_QUESTION} selectable options. " + f"Use at least {MIN_OPTIONS_PER_QUESTION} unless input_mode='text'." + ), + ), + ] + input_mode: Annotated[ + Literal["select", "text", "select_or_text"], + Field( + default="select", + description=( + "Answer mode: select from options, enter free-form text, " + "or allow either." + ), + ), + ] + input_placeholder: Annotated[ + str, + BeforeValidator(_sanitize_optional), + Field( + default="Type your answer...", + max_length=MAX_DESCRIPTION_LENGTH, + description="Placeholder for free-form user input modes", ), ] + @model_validator(mode="before") + @classmethod + def validate_direct_model_option_limit(cls, values: Any) -> Any: + """Keep direct ``Question(...)`` construction capped for terminal ergonomics. + + Backward-compat nuance: legacy tests instantiate ``Question`` with pre-built + ``QuestionOption`` objects and expect >6 options to fail. Browser/tool payloads + arrive as dicts via ``AskUserQuestionInput`` and are allowed up to + ``MAX_OPTIONS_PER_QUESTION``. + """ + if isinstance(values, dict): + options = values.get("options") + if ( + isinstance(options, list) + and len(options) > _DIRECT_MODEL_MAX_OPTIONS + and options + and all(isinstance(opt, QuestionOption) for opt in options) + ): + raise ValueError( + f"options must include at most {_DIRECT_MODEL_MAX_OPTIONS} items" + ) + return values + + @property + def allows_text_input(self) -> bool: + """Return True when this question accepts free-form text.""" + return self.input_mode in {"text", "select_or_text"} + @model_validator(mode="after") - def validate_unique_labels(self) -> Question: - """Ensure all option labels are unique within a question.""" - _check_unique([opt.label for opt in self.options], "Option labels") + def validate_question_shape(self) -> Question: + """Ensure labels are unique and options are present when required.""" + if self.options: + _check_unique([opt.label for opt in self.options], "Option labels") + if self.input_mode == "select" and len(self.options) < MIN_OPTIONS_PER_QUESTION: + raise ValueError( + f"options must include at least {MIN_OPTIONS_PER_QUESTION} items " + "when input_mode='select'" + ) return self @@ -213,6 +277,14 @@ class QuestionAnswer(BaseModel): description="Custom text if 'Other' was selected", ), ] + user_input: Annotated[ + str | None, + Field( + default=None, + max_length=MAX_OTHER_TEXT_LENGTH, + description="Free-form text answer for input-enabled questions", + ), + ] @property def has_other(self) -> bool: @@ -222,7 +294,11 @@ def has_other(self) -> bool: @property def is_empty(self) -> bool: """Check if no options were selected.""" - return not self.selected_options and self.other_text is None + return ( + not self.selected_options + and self.other_text is None + and self.user_input is None + ) class AskUserQuestionOutput(BaseModel): diff --git a/code_puppy/tools/ask_user_question/registration.py b/code_puppy/tools/ask_user_question/registration.py index 418429353..423b47043 100644 --- a/code_puppy/tools/ask_user_question/registration.py +++ b/code_puppy/tools/ask_user_question/registration.py @@ -2,7 +2,11 @@ from __future__ import annotations +import asyncio +import copy +import inspect import json +from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Annotated, Any, Dict, List from pydantic import BeforeValidator, Field, WithJsonSchema @@ -17,7 +21,7 @@ MAX_QUESTIONS_PER_CALL, MIN_OPTIONS_PER_QUESTION, ) -from .handler import ask_user_question as _ask_user_question_impl +from .handler import ask_user_question_async as _ask_user_question_impl from .models import AskUserQuestionOutput if TYPE_CHECKING: @@ -61,15 +65,25 @@ "type": "boolean", "description": "If true, user can select multiple options (default: false)", }, + "input_mode": { + "type": "string", + "enum": ["select", "text", "select_or_text"], + "description": "Use 'text' for free-form user input, 'select' for options, or 'select_or_text' for both.", + }, + "input_placeholder": { + "type": "string", + "maxLength": MAX_DESCRIPTION_LENGTH, + "description": "Placeholder shown for free-form text input.", + }, "options": { "type": "array", "items": _OPTION_SCHEMA, - "minItems": MIN_OPTIONS_PER_QUESTION, + "minItems": 0, "maxItems": MAX_OPTIONS_PER_QUESTION, - "description": f"Array of {MIN_OPTIONS_PER_QUESTION}-{MAX_OPTIONS_PER_QUESTION} selectable options", + "description": f"Array of 0-{MAX_OPTIONS_PER_QUESTION} selectable options. Use at least {MIN_OPTIONS_PER_QUESTION} unless input_mode is 'text'.", }, }, - "required": ["question", "header", "options"], + "required": ["question", "header"], } _QUESTIONS_ARRAY_SCHEMA: Dict[str, Any] = { @@ -81,12 +95,24 @@ f"Array of 1-{MAX_QUESTIONS_PER_CALL} question objects. Each question needs: " f"'question' (max {MAX_QUESTION_LENGTH} chars), " f"'header' (max {MAX_HEADER_LENGTH} chars, no spaces), " - f"'options' (array of {MIN_OPTIONS_PER_QUESTION}-{MAX_OPTIONS_PER_QUESTION} options with 'label'). " - "Optional: 'multi_select' (boolean)." + f"'options' (array of 0-{MAX_OPTIONS_PER_QUESTION} options with 'label'). " + "Optional: 'multi_select' (boolean), 'input_mode' ('select', 'text', or 'select_or_text'), " + "and 'input_placeholder'." ), } +# Backward-compat: the registered tool schema keeps `options` as required +# with legacy minItems, while `_QUESTIONS_ARRAY_SCHEMA` remains browser-friendly +# (options optional for input_mode="text"). +_TOOL_QUESTION_SCHEMA: Dict[str, Any] = copy.deepcopy(_QUESTION_SCHEMA) +_TOOL_QUESTION_SCHEMA["required"] = ["question", "header", "options"] +_TOOL_QUESTION_SCHEMA["properties"]["options"]["minItems"] = MIN_OPTIONS_PER_QUESTION + +_TOOL_QUESTIONS_ARRAY_SCHEMA: Dict[str, Any] = copy.deepcopy(_QUESTIONS_ARRAY_SCHEMA) +_TOOL_QUESTIONS_ARRAY_SCHEMA["items"] = _TOOL_QUESTION_SCHEMA + + def _coerce_questions_json_string(v: Any) -> Any: """Coerce a JSON-stringified array to a native list before pydantic validates it. @@ -112,18 +138,35 @@ def _coerce_questions_json_string(v: Any) -> Any: QuestionsListWithSchema = Annotated[ List[Dict[str, Any]], BeforeValidator(_coerce_questions_json_string), - WithJsonSchema(_QUESTIONS_ARRAY_SCHEMA), + WithJsonSchema(_TOOL_QUESTIONS_ARRAY_SCHEMA), Field( description=( f"Array of 1-{MAX_QUESTIONS_PER_CALL} question objects. Each question needs: " f"'question' (max {MAX_QUESTION_LENGTH} chars), " f"'header' (max {MAX_HEADER_LENGTH} chars), " - f"'options' ({MIN_OPTIONS_PER_QUESTION}-{MAX_OPTIONS_PER_QUESTION} options with 'label')." + f"'options' (0-{MAX_OPTIONS_PER_QUESTION} options with 'label'; " + "at least 2 unless input_mode='text')." ) ), ] +def _resolve_tool_result(result: Any) -> Any: + """Resolve a possibly-awaitable tool result for sync registration paths.""" + if not inspect.isawaitable(result): + return result + + try: + asyncio.get_running_loop() + except RuntimeError: + # No running loop in this thread, safe to run directly. + return asyncio.run(result) + + # Running event loop in this thread: resolve on a dedicated thread+loop. + with ThreadPoolExecutor(max_workers=1) as executor: + return executor.submit(lambda: asyncio.run(result)).result() + + def register_ask_user_question(agent: Agent) -> None: """Register the ask_user_question tool with the given agent.""" @@ -132,14 +175,12 @@ def ask_user_question( context: RunContext, # noqa: ARG001 - Required by framework questions: QuestionsListWithSchema, ) -> AskUserQuestionOutput: - """Ask the user multiple related questions in an interactive TUI.""" + """Ask the user multiple related questions in the active UI.""" # Keep the external tool schema simple for provider compatibility. # The handler performs the real nested validation and normalization. # Fire a Claude Code-style notification so plugins can react when the # agent is awaiting user input. try: - import asyncio as _asyncio - from code_puppy.callbacks import on_notification _coro = on_notification( @@ -148,10 +189,10 @@ def ask_user_question( context={"questions": questions}, ) try: - _asyncio.get_running_loop() - _asyncio.ensure_future(_coro) + asyncio.get_running_loop() + asyncio.ensure_future(_coro) except RuntimeError: - _asyncio.run(_coro) + asyncio.run(_coro) except Exception: pass - return _ask_user_question_impl(questions) + return _resolve_tool_result(_ask_user_question_impl(questions)) diff --git a/code_puppy/tools/command_runner.py b/code_puppy/tools/command_runner.py index 2da8d6a4c..445bcbb8e 100644 --- a/code_puppy/tools/command_runner.py +++ b/code_puppy/tools/command_runner.py @@ -1,4 +1,5 @@ import asyncio +import contextvars import ctypes import os import select @@ -11,6 +12,7 @@ import traceback from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager +from contextvars import ContextVar from functools import partial from typing import Callable, List, Literal, Optional, Set @@ -19,6 +21,7 @@ from rich.text import Text from code_puppy.callbacks import on_run_shell_command_output +from code_puppy.config import get_command_timeout_seconds from code_puppy.messaging import ( # Structured messaging types AgentReasoningMessage, ShellOutputMessage, @@ -106,17 +109,105 @@ def _win32_pipe_has_data(pipe) -> bool: _AWAITING_USER_INPUT = threading.Event() -# NOTE: The previous module-level ``_CONFIRMATION_LOCK`` was removed -- -# queueing of parallel approval prompts now lives inside -# ``get_user_approval_async`` itself, so every caller (shell commands, -# destructive-command guard, force-push guard, ...) benefits without -# bolting on their own lock. +_CONFIRMATION_LOCK = threading.Lock() + # Track running shell processes so we can kill them on Ctrl-C from the UI _RUNNING_PROCESSES: Set[subprocess.Popen] = set() _RUNNING_PROCESSES_LOCK = threading.Lock() _USER_KILLED_PROCESSES = set() +# Per-session process tracking via ContextVar +_session_running_processes: ContextVar[Optional[Set[subprocess.Popen]]] = ContextVar( + "session_running_processes", default=None +) +_session_killed_processes: ContextVar[Optional[Set[int]]] = ContextVar( + "session_killed_processes", default=None +) +_session_awaiting_input: ContextVar[Optional[threading.Event]] = ContextVar( + "session_awaiting_input", default=None +) + + +def _get_running_processes() -> Set[subprocess.Popen]: + """Get the running processes set for the current session context.""" + session_set = _session_running_processes.get(None) + if session_set is not None: + return session_set + return _RUNNING_PROCESSES + + +def _get_killed_processes() -> Set[int]: + """Get the killed processes set for the current session context.""" + session_set = _session_killed_processes.get(None) + if session_set is not None: + return session_set + return _USER_KILLED_PROCESSES + + +def _get_awaiting_input_event() -> threading.Event: + """Get the awaiting-input event for the current session context.""" + session_evt = _session_awaiting_input.get(None) + if session_evt is not None: + return session_evt + return _AWAITING_USER_INPUT + + +def init_session_process_tracking() -> None: + """Initialize per-session process tracking. Call at WS session start.""" + _session_running_processes.set(set()) + _session_killed_processes.set(set()) + _session_awaiting_input.set(threading.Event()) + _session_active_stop_events.set(set()) + _session_keyboard_refcount.set(0) + _session_ctrl_x_stop_event.set(None) + _session_ctrl_x_thread.set(None) + + +def cleanup_session_process_tracking() -> None: + """Tear down per-session process tracking. Call at WS session end. + + Kills any still-running session processes, then clears the per-session + process and killed-process sets to release any held Popen references. + """ + # Kill any still-running session processes first + session_procs = _session_running_processes.get(None) + if session_procs is not None: + for p in list(session_procs): + try: + if p.poll() is None: + _kill_process_group(p) + except Exception: + pass + session_procs.clear() + session_killed = _session_killed_processes.get(None) + if session_killed is not None: + session_killed.clear() + session_stop = _session_active_stop_events.get(None) + if session_stop is not None: + for evt in session_stop: + evt.set() # Signal before discarding! + session_stop.clear() + # Clean up keyboard context + session_ctrl_x = _session_ctrl_x_stop_event.get(None) + if session_ctrl_x is not None: + session_ctrl_x.set() # Signal stop + session_thread = _session_ctrl_x_thread.get(None) + if session_thread is not None and session_thread.is_alive(): + try: + session_thread.join(timeout=0.2) + except Exception: + pass + _session_keyboard_refcount.set(None) + _session_ctrl_x_stop_event.set(None) + _session_ctrl_x_thread.set(None) + # Reset to None so subsequent calls fall back to global + _session_running_processes.set(None) + _session_killed_processes.set(None) + _session_awaiting_input.set(None) + _session_active_stop_events.set(None) + + # Global state for shell command keyboard handling _SHELL_CTRL_X_STOP_EVENT: Optional[threading.Event] = None _SHELL_CTRL_X_THREAD: Optional[threading.Thread] = None @@ -138,10 +229,82 @@ def _win32_pipe_has_data(pipe) -> bool: _KEYBOARD_CONTEXT_REFCOUNT = 0 _KEYBOARD_CONTEXT_LOCK = threading.Lock() +# Per-session keyboard context refcount (ContextVar for WS session isolation) +_session_keyboard_refcount: ContextVar[Optional[int]] = ContextVar( + "session_keyboard_refcount", default=None +) +_session_ctrl_x_stop_event: ContextVar[Optional[threading.Event]] = ContextVar( + "session_ctrl_x_stop_event", default=None +) +_session_ctrl_x_thread: ContextVar[Optional[threading.Thread]] = ContextVar( + "session_ctrl_x_thread", default=None +) + # Thread-safe registry of active stop events for concurrent shell commands _ACTIVE_STOP_EVENTS: Set[threading.Event] = set() _ACTIVE_STOP_EVENTS_LOCK = threading.Lock() +# Per-session stop events (ContextVar for session isolation) +_session_active_stop_events: ContextVar[Optional[Set[threading.Event]]] = ContextVar( + "session_active_stop_events", default=None +) + + +def _get_keyboard_refcount() -> int: + """Get keyboard context refcount for the current session.""" + session_val = _session_keyboard_refcount.get(None) + if session_val is not None: + return session_val + return _KEYBOARD_CONTEXT_REFCOUNT + + +def _set_keyboard_refcount(value: int) -> None: + """Set keyboard context refcount for the current session.""" + if _session_keyboard_refcount.get(None) is not None: + _session_keyboard_refcount.set(value) + else: + global _KEYBOARD_CONTEXT_REFCOUNT + _KEYBOARD_CONTEXT_REFCOUNT = value + + +def _get_ctrl_x_stop_event() -> Optional[threading.Event]: + """Get Ctrl+X stop event for the current session.""" + session_evt = _session_ctrl_x_stop_event.get(None) + if session_evt is not None: + return session_evt + return _SHELL_CTRL_X_STOP_EVENT + + +def _get_ctrl_x_thread() -> Optional[threading.Thread]: + """Get Ctrl+X thread for the current session.""" + session_thread = _session_ctrl_x_thread.get(None) + if session_thread is not None: + return session_thread + return _SHELL_CTRL_X_THREAD + + +def _get_active_stop_events() -> Set[threading.Event]: + """Get the active stop events set for the current session context.""" + session_set = _session_active_stop_events.get(None) + if session_set is not None: + return session_set + return _ACTIVE_STOP_EVENTS + + +@contextmanager +def _guarded_set(collection, global_ref, lock): + """Acquire *lock* only when *collection* is the shared global *global_ref*. + + Per-session sets (ContextVar-backed) are single-writer within their + asyncio task context and don't need locking. + """ + if collection is global_ref: + with lock: + yield collection + else: + yield collection + + # Mid-flight backgrounding (Ctrl+X Ctrl+B) machinery lives in # ``shell_backgrounding`` (600-line cap); re-exported here because the # chord handler and the streaming pumps are the consumers. @@ -153,13 +316,15 @@ def _win32_pipe_has_data(pipe) -> bool: def _register_process(proc: subprocess.Popen) -> None: - with _RUNNING_PROCESSES_LOCK: - _RUNNING_PROCESSES.add(proc) + procs = _get_running_processes() + with _guarded_set(procs, _RUNNING_PROCESSES, _RUNNING_PROCESSES_LOCK): + procs.add(proc) def _unregister_process(proc: subprocess.Popen) -> None: - with _RUNNING_PROCESSES_LOCK: - _RUNNING_PROCESSES.discard(proc) + procs = _get_running_processes() + with _guarded_set(procs, _RUNNING_PROCESSES, _RUNNING_PROCESSES_LOCK): + procs.discard(proc) def _kill_process_group(proc: subprocess.Popen) -> None: @@ -244,13 +409,16 @@ def kill_all_running_shell_processes() -> int: Returns the number of processes signaled. """ # Signal all active reader threads to stop - with _ACTIVE_STOP_EVENTS_LOCK: - for evt in _ACTIVE_STOP_EVENTS: + active_stop = _get_active_stop_events() + with _guarded_set(active_stop, _ACTIVE_STOP_EVENTS, _ACTIVE_STOP_EVENTS_LOCK): + for evt in active_stop: evt.set() procs: list[subprocess.Popen] - with _RUNNING_PROCESSES_LOCK: - procs = list(_RUNNING_PROCESSES) + running = _get_running_processes() + killed = _get_killed_processes() + with _guarded_set(running, _RUNNING_PROCESSES, _RUNNING_PROCESSES_LOCK): + procs = list(running) count = 0 for p in procs: try: @@ -268,7 +436,7 @@ def kill_all_running_shell_processes() -> int: if p.poll() is None: _kill_process_group(p) count += 1 - _USER_KILLED_PROCESSES.add(p.pid) + killed.add(p.pid) finally: _unregister_process(p) return count @@ -276,23 +444,24 @@ def kill_all_running_shell_processes() -> int: def get_running_shell_process_count() -> int: """Return the number of currently-active shell processes being tracked.""" - with _RUNNING_PROCESSES_LOCK: + running = _get_running_processes() + with _guarded_set(running, _RUNNING_PROCESSES, _RUNNING_PROCESSES_LOCK): alive = 0 stale: Set[subprocess.Popen] = set() - for proc in _RUNNING_PROCESSES: + for proc in running: if proc.poll() is None: alive += 1 else: stale.add(proc) for proc in stale: - _RUNNING_PROCESSES.discard(proc) + running.discard(proc) return alive # Function to check if user input is awaited def is_awaiting_user_input(): """Check if command_runner is waiting for user input.""" - return _AWAITING_USER_INPUT.is_set() + return _get_awaiting_input_event().is_set() # Function to set user input flag @@ -313,9 +482,10 @@ def set_awaiting_user_input(awaiting=True): to guess from the screen). """ if awaiting: - _AWAITING_USER_INPUT.set() + _get_awaiting_input_event().set() else: - _AWAITING_USER_INPUT.clear() + _get_awaiting_input_event().clear() + # Best-effort notification; never let an observer disturb the prompt path. try: from code_puppy.callbacks import on_awaiting_user_input @@ -679,16 +849,15 @@ def run_shell_command_streaming( silent: bool = False, ): stop_event = threading.Event() - with _ACTIVE_STOP_EVENTS_LOCK: - _ACTIVE_STOP_EVENTS.add(stop_event) + active_stop = _get_active_stop_events() + with _guarded_set(active_stop, _ACTIVE_STOP_EVENTS, _ACTIVE_STOP_EVENTS_LOCK): + active_stop.add(stop_event) start_time = time.time() last_output_time = [start_time] # Foreground duration limit. Reaching it detaches rather than killing; # inactivity remains the guard for genuinely wedged commands. - from code_puppy.config import get_command_timeout_seconds - foreground_limit_seconds = get_command_timeout_seconds() stdout_lines = [] @@ -1023,8 +1192,9 @@ def detach_to_background(*, automatic: bool = False): ) get_message_bus().emit(shell_output_msg) - with _ACTIVE_STOP_EVENTS_LOCK: - _ACTIVE_STOP_EVENTS.discard(stop_event) + active_stop = _get_active_stop_events() + with _guarded_set(active_stop, _ACTIVE_STOP_EVENTS, _ACTIVE_STOP_EVENTS_LOCK): + active_stop.discard(stop_event) if exit_code != 0: time.sleep(1) @@ -1038,7 +1208,7 @@ def detach_to_background(*, automatic: bool = False): exit_code=exit_code, execution_time=execution_time, timeout=False, - user_interrupted=process.pid in _USER_KILLED_PROCESSES, + user_interrupted=process.pid in _get_killed_processes(), ) return ShellCommandOutput( @@ -1052,8 +1222,9 @@ def detach_to_background(*, automatic: bool = False): ) except Exception as e: - with _ACTIVE_STOP_EVENTS_LOCK: - _ACTIVE_STOP_EVENTS.discard(stop_event) + active_stop = _get_active_stop_events() + with _guarded_set(active_stop, _ACTIVE_STOP_EVENTS, _ACTIVE_STOP_EVENTS_LOCK): + active_stop.discard(stop_event) return ShellCommandOutput( success=False, command=command, @@ -1072,6 +1243,15 @@ async def run_shell_command( timeout: int = 60, background: bool = False, ) -> ShellCommandOutput: + # Resolve CWD from the session ContextVar when not explicitly provided. + # This supports concurrent WebSocket sessions where each session has its own CWD + # without relying on the process-global os.chdir(). + # Returns None for CLI sessions (subprocess inherits process CWD as normal). + if cwd is None: + from code_puppy.api.session_cwd import get_session_working_directory + + cwd = get_session_working_directory() + # Generate unique group_id for this command execution group_id = generate_group_id("shell_command", command) @@ -1221,12 +1401,26 @@ async def run_shell_command( # Check if we're running as a sub-agent (skip confirmation and run silently) running_as_subagent = is_subagent() + # Check if WebSocket mode is active (permission handled via WebSocket callbacks) + websocket_mode_active = False + try: + from code_puppy.api.permission_plugin import get_websocket_context + + websocket_mode_active = get_websocket_context() is not None + except (ImportError, AttributeError): + pass + # Only ask for confirmation if we're in an interactive TTY, not in yolo mode, - # and NOT running as a sub-agent (sub-agents run without user interaction) - if not yolo_mode and not running_as_subagent and sys.stdin.isatty(): - # No local lock needed -- get_user_approval_async serializes - # parallel prompts internally so the 2nd, 3rd, 4th... destructive - # commands queue up cleanly instead of vanishing. + # NOT running as a sub-agent, and NOT in WebSocket mode (WebSocket has its own permission system) + if ( + not yolo_mode + and not running_as_subagent + and not websocket_mode_active + and sys.stdin.isatty() + ): + # Queue concurrent confirmations instead of rejecting parallel callers. + while not _CONFIRMATION_LOCK.acquire(blocking=False): + await asyncio.sleep(0.01) # Get puppy name for personalized messages from code_puppy.config import get_puppy_name @@ -1244,15 +1438,17 @@ async def run_shell_command( panel_content.append("Working directory: ", style="dim") panel_content.append(cwd, style="dim cyan") - # Use the common approval function (async version). - # Internal queueing means parallel calls wait their turn here. - confirmed, user_feedback = await get_user_approval_async( - title="Shell Command", - content=panel_content, - preview=None, - border_style="dim white", - puppy_name=puppy_name, - ) + # Use the common approval function (async version) + try: + confirmed, user_feedback = await get_user_approval_async( + title="Shell Command", + content=panel_content, + preview=None, + border_style="dim white", + puppy_name=puppy_name, + ) + finally: + _CONFIRMATION_LOCK.release() if not confirmed: if user_feedback: @@ -1472,9 +1668,13 @@ async def _run_command_inner( try: # Run the blocking shell command in a thread pool to avoid blocking the event loop # This allows multiple sub-agents to run shell commands in parallel + # Copy context so ContextVar-based session tracking propagates to the worker thread + ctx = contextvars.copy_context() return await loop.run_in_executor( _SHELL_EXECUTOR, - partial(_run_command_sync, command, cwd, timeout, group_id, silent), + partial( + ctx.run, _run_command_sync, command, cwd, timeout, group_id, silent + ), ) except Exception as e: if not silent: diff --git a/pyproject.toml b/pyproject.toml index e1bf10338..9e5100641 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,10 @@ description = "Code generation agent" readme = "README.md" requires-python = ">=3.11,<3.15" dependencies = [ + "aiosqlite>=0.22.1", + "websockets>=12.0", + "uvicorn[standard]>=0.27.0", + "fastapi>=0.109.0", "pydantic-ai-slim[openai,anthropic,mcp]==1.56.0", "typer>=0.12.0", "mcp>=1.9.4", @@ -64,6 +68,7 @@ HomePage = "https://github.com/mpfaffenberger/code_puppy" [project.scripts] code-puppy = "code_puppy.main:main_entry" pup = "code_puppy.main:main_entry" +code-puppy-api = "code_puppy.api.main:main" [tool.logfire] ignore_no_config = true diff --git a/tests/api/test_code_puppy_api_startup.py b/tests/api/test_code_puppy_api_startup.py new file mode 100644 index 000000000..900bdcce2 --- /dev/null +++ b/tests/api/test_code_puppy_api_startup.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + + +def test_code_puppy_api_starts_without_dbos_installed(tmp_path): + blocker_dir = tmp_path / "block_dbos" + blocker_dir.mkdir() + sitecustomize = blocker_dir / "sitecustomize.py" + sitecustomize.write_text( + """ +import builtins + +_original_import = builtins.__import__ + + +def _blocked_import(name, globals=None, locals=None, fromlist=(), level=0): + if name == \"dbos\" or name.startswith(\"dbos.\"): + raise ModuleNotFoundError(\"No module named 'dbos'\") + return _original_import(name, globals, locals, fromlist, level) + + +builtins.__import__ = _blocked_import +""".strip() + ) + + db_path = tmp_path / "chat_messages.db" + env = os.environ.copy() + env["PUPPY_DESK_DB"] = str(db_path) + env["PYTHONPATH"] = os.pathsep.join( + [str(blocker_dir), env.get("PYTHONPATH", "")] + ).rstrip(os.pathsep) + + script = "\n".join( + [ + "from contextlib import asynccontextmanager", + "from starlette.routing import Router", + "@asynccontextmanager", + "async def _empty_lifespan(_app):", + " yield", + "_original_router_init = Router.__init__", + "def _compat_router_init(self, *args, **kwargs):", + " on_startup = kwargs.pop('on_startup', None)", + " on_shutdown = kwargs.pop('on_shutdown', None)", + " lifespan = kwargs.pop('lifespan', None)", + " result = _original_router_init(self, *args, **kwargs)", + " self.on_startup = list(on_startup or [])", + " self.on_shutdown = list(on_shutdown or [])", + " self.lifespan_context = lifespan or _empty_lifespan", + " return result", + "Router.__init__ = _compat_router_init", + "from fastapi.testclient import TestClient", + "from code_puppy.api.app import create_app", + "with TestClient(create_app()) as client:", + " response = client.get('/health')", + " assert response.status_code == 200", + " assert response.json()['status'] == 'healthy'", + "print('ok')", + ] + ) + + command = [sys.executable, "-c", script] + + result = subprocess.run( + command, + cwd=Path(__file__).resolve().parents[2], + env=env, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + + assert result.returncode == 0, result.stderr or result.stdout + assert "ok" in result.stdout diff --git a/tests/test_ask_user_question_schema_input_modes.py b/tests/test_ask_user_question_schema_input_modes.py new file mode 100644 index 000000000..d9f0d55c0 --- /dev/null +++ b/tests/test_ask_user_question_schema_input_modes.py @@ -0,0 +1,93 @@ +"""Schema tests for browser/user-input ask_user_question support.""" + +import pytest +from pydantic import ValidationError + +from code_puppy.tools.ask_user_question.models import ( + AskUserQuestionInput, + QuestionAnswer, +) +from code_puppy.tools.ask_user_question.registration import _QUESTIONS_ARRAY_SCHEMA + + +def _options(count: int): + return [{"label": f"Option {i}"} for i in range(count)] + + +def test_accepts_more_than_six_options_up_to_twelve(): + payload = { + "questions": [ + { + "question": "Pick any option", + "header": "choices", + "options": _options(11), + } + ] + } + + validated = AskUserQuestionInput.model_validate(payload) + + assert len(validated.questions[0].options) == 11 + + +def test_rejects_more_than_twelve_options(): + payload = { + "questions": [ + { + "question": "Pick any option", + "header": "choices", + "options": _options(13), + } + ] + } + + with pytest.raises(ValidationError): + AskUserQuestionInput.model_validate(payload) + + +def test_text_input_mode_allows_no_options(): + payload = { + "questions": [ + { + "question": "What should the UI say?", + "header": "copy", + "input_mode": "text", + "input_placeholder": "Type copy...", + } + ] + } + + validated = AskUserQuestionInput.model_validate(payload) + + assert validated.questions[0].options == [] + assert validated.questions[0].allows_text_input is True + + +def test_select_mode_still_requires_two_options(): + payload = { + "questions": [ + { + "question": "Pick one", + "header": "one", + "options": [{"label": "Only"}], + } + ] + } + + with pytest.raises(ValidationError, match="at least 2"): + AskUserQuestionInput.model_validate(payload) + + +def test_tool_schema_advertises_twelve_options_and_input_mode(): + question_schema = _QUESTIONS_ARRAY_SCHEMA["items"] + + assert question_schema["properties"]["options"]["maxItems"] == 12 + assert "input_mode" in question_schema["properties"] + assert "options" not in question_schema["required"] + + +def test_question_answer_can_carry_free_form_user_input(): + answer = QuestionAnswer(question_header="copy", user_input="Use a compact card") + + assert answer.user_input == "Use a compact card" + assert answer.is_empty is False diff --git a/tests/test_chatgpt_codex_client.py b/tests/test_chatgpt_codex_client.py index b3de6d934..e5e8cf79f 100644 --- a/tests/test_chatgpt_codex_client.py +++ b/tests/test_chatgpt_codex_client.py @@ -249,6 +249,41 @@ def test_remove_all_unsupported_params(self): assert "verbosity" not in data assert data["temperature"] == 0.7 # Preserved + def test_dedupes_duplicate_input_item_ids(self): + """Duplicate input item IDs should be removed while preserving first-seen order.""" + body = json.dumps( + { + "model": "gpt-5", + "input": [ + {"id": "rs_1", "type": "reasoning", "content": "first"}, + {"id": "rs_1", "type": "reasoning", "content": "duplicate"}, + {"id": "rs_2", "type": "reasoning", "content": "second"}, + {"type": "message", "content": "no-id item"}, + ], + } + ).encode() + + result, _ = ChatGPTCodexAsyncClient._inject_codex_fields(body) + + assert result is not None + data = json.loads(result) + input_items = data["input"] + + # rs_1 duplicate should be removed; first occurrence preserved. + assert [ + item.get("id") + for item in input_items + if isinstance(item, dict) and item.get("id") + ] == [ + "rs_1", + "rs_2", + ] + # non-id items should still be present + assert any( + isinstance(item, dict) and item.get("type") == "message" + for item in input_items + ) + def test_invalid_json_returns_none(self): """Test that invalid JSON returns (None, False).""" body = b"not valid json" diff --git a/tests/test_messaging_bus.py b/tests/test_messaging_bus.py index 6ee0c6282..3ce770d84 100644 --- a/tests/test_messaging_bus.py +++ b/tests/test_messaging_bus.py @@ -34,7 +34,7 @@ def test_initialization_default(self): assert bus._maxsize == 1000 assert isinstance(bus._outgoing, queue.Queue) assert isinstance(bus._incoming, queue.Queue) - assert bus._current_session_id is None + assert bus.get_session_context() is None assert not bus._has_active_renderer assert bus._startup_buffer == [] @@ -587,3 +587,35 @@ def update_context(session_id): t.join() assert len(contexts) == 5 + + @pytest.mark.asyncio + async def test_session_context_is_task_local(self): + """Each asyncio task should see its own session context.""" + bus = MessageBus() + bus._has_active_renderer = True + + ready_first = asyncio.Event() + ready_second = asyncio.Event() + release = asyncio.Event() + + async def worker(session_id, ready_to_set, ready_done): + if ready_to_set is not None: + await ready_to_set.wait() + bus.set_session_context(session_id) + ready_done.set() + await release.wait() + msg = TextMessage(level=MessageLevel.INFO, text=session_id) + bus.emit(msg) + return bus.get_session_context(), msg.session_id + + first = asyncio.create_task(worker("session-1", None, ready_first)) + second = asyncio.create_task(worker("session-2", ready_first, ready_second)) + + await ready_second.wait() + release.set() + + first_ctx, second_ctx = await asyncio.gather(first, second) + + assert first_ctx == ("session-1", "session-1") + assert second_ctx == ("session-2", "session-2") + assert bus.get_session_context() is None diff --git a/tests/tools/test_command_runner_lock_release.py b/tests/tools/test_command_runner_lock_release.py new file mode 100644 index 000000000..4915a8f45 --- /dev/null +++ b/tests/tools/test_command_runner_lock_release.py @@ -0,0 +1,47 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + + +@pytest.mark.asyncio +async def test_confirmation_lock_released_when_approval_raises(monkeypatch): + from code_puppy.tools import command_runner + + # ensure unlocked before test begins + if command_runner._CONFIRMATION_LOCK.locked(): + command_runner._CONFIRMATION_LOCK.release() + + async def _empty_callbacks(*_args, **_kwargs): + return [] + + async def _boom_approval(**_kwargs): + raise RuntimeError("approval transport failed") + + monkeypatch.setattr("code_puppy.callbacks.on_run_shell_command", _empty_callbacks) + monkeypatch.setattr("code_puppy.config.get_yolo_mode", lambda: False) + monkeypatch.setattr("code_puppy.tools.command_runner.is_subagent", lambda: False) + monkeypatch.setattr( + "code_puppy.api.permission_plugin.get_websocket_context", lambda: None + ) + monkeypatch.setattr( + "code_puppy.tools.command_runner.get_user_approval_async", + _boom_approval, + ) + monkeypatch.setattr("sys.stdin.isatty", lambda: True) + + with pytest.raises(RuntimeError, match="approval transport failed"): + await command_runner.run_shell_command( + context=SimpleNamespace(), + command="echo hi", + cwd=".", + timeout=5, + background=False, + ) + + # lock should be free for future commands + acquired = command_runner._CONFIRMATION_LOCK.acquire(blocking=False) + assert acquired is True + if acquired: + command_runner._CONFIRMATION_LOCK.release() diff --git a/uv.lock b/uv.lock index 51fbc1a4b..bc65b8db7 100644 --- a/uv.lock +++ b/uv.lock @@ -14,6 +14,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a0/48/a1367dded7765f5489f9840389949d5aeac53a73aeb1b4e6c0ab61e1617a/agent_client_protocol-0.11.0-py3-none-any.whl", hash = "sha256:2d8570cd4911f8af9fbab4808f7b05f09204a756f346c7409ecdcae3f37cdb2f", size = 68089, upload-time = "2026-07-05T17:07:55.518Z" }, ] +[[package]] +name = "aiosqlite" +version = "0.22.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/8a/64761f4005f17809769d23e518d915db74e6310474e733e3593cfc854ef1/aiosqlite-0.22.1.tar.gz", hash = "sha256:043e0bd78d32888c0a9ca90fc788b38796843360c855a7262a532813133a0650", size = 14821, upload-time = "2025-12-23T19:25:43.997Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/00/b7/e3bf5133d697a08128598c8d0abc5e16377b51465a33756de24fa7dee953/aiosqlite-0.22.1-py3-none-any.whl", hash = "sha256:21c002eb13823fad740196c5a2e9d8e62f6243bd9e7e4a1f87fb5e44ecb4fceb", size = 17405, upload-time = "2025-12-23T19:25:42.139Z" }, +] + [[package]] name = "annotated-doc" version = "0.0.4" @@ -316,9 +325,11 @@ version = "0.0.637" source = { editable = "." } dependencies = [ { name = "agent-client-protocol" }, + { name = "aiosqlite" }, { name = "anthropic" }, { name = "azure-identity" }, { name = "boto3" }, + { name = "fastapi" }, { name = "httpx", extra = ["http2"] }, { name = "json-repair" }, { name = "mcp" }, @@ -336,6 +347,8 @@ dependencies = [ { name = "ripgrep" }, { name = "termflow-md" }, { name = "typer" }, + { name = "uvicorn", extra = ["standard"] }, + { name = "websockets" }, ] [package.optional-dependencies] @@ -358,11 +371,13 @@ dev = [ [package.metadata] requires-dist = [ { name = "agent-client-protocol", specifier = ">=0.11.0,<0.12" }, + { name = "aiosqlite", specifier = ">=0.22.1" }, { name = "anthropic", specifier = "==0.79.0" }, { name = "azure-identity", specifier = ">=1.15.0" }, { name = "boto3", specifier = ">=1.43.9" }, { name = "boto3", marker = "extra == 'bedrock'", specifier = ">=1.35.0" }, { name = "dbos", marker = "extra == 'durable'", specifier = ">=2.11.0" }, + { name = "fastapi", specifier = ">=0.109.0" }, { name = "httpx", extras = ["http2"], specifier = ">=0.24.1" }, { name = "json-repair", specifier = ">=0.46.2" }, { name = "mcp", specifier = ">=1.9.4" }, @@ -380,6 +395,8 @@ requires-dist = [ { name = "ripgrep", specifier = "==14.1.0" }, { name = "termflow-md", specifier = ">=0.1.11" }, { name = "typer", specifier = ">=0.12.0" }, + { name = "uvicorn", extras = ["standard"], specifier = ">=0.27.0" }, + { name = "websockets", specifier = ">=12.0" }, ] provides-extras = ["bedrock", "durable"] @@ -599,6 +616,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a7/5f/ed01f9a3cdffbd5a008556fc7b2a08ddb1cc6ace7effa7340604b1d16699/docstring_parser-0.18.0-py3-none-any.whl", hash = "sha256:b3fcbed555c47d8479be0796ef7e19c2670d428d72e96da63f3a40122860374b", size = 22484, upload-time = "2026-04-14T04:09:18.638Z" }, ] +[[package]] +name = "fastapi" +version = "0.139.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-doc" }, + { name = "pydantic" }, + { name = "starlette" }, + { name = "typing-extensions" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d3/af/a5f50ccfa659ec1802cb4ca842c23f06d906a8cc9aef6016a2caeea3d4ed/fastapi-0.139.0.tar.gz", hash = "sha256:99ab7b2d92223c76d6cf10757ab3f89d45b38267fc20b2a136cf02f6beac3145", size = 423016, upload-time = "2026-07-01T16:35:33.436Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/7c/8e3c6ad324ea5cb36604fc3f968554887891c316d9dfde57761611d907ad/fastapi-0.139.0-py3-none-any.whl", hash = "sha256:cf15e1e9e667ddb0ad63811e60bd11390d1aac838ca4a7a23f421807b2308189", size = 130339, upload-time = "2026-07-01T16:35:32.19Z" }, +] + [[package]] name = "genai-prices" version = "0.0.60" @@ -748,6 +781,49 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, ] +[[package]] +name = "httptools" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/43/e5/d471fcb0e14523fe1c3f4ba58ca52480e7bd70ad7109a3846bc75892f7fb/httptools-0.8.0.tar.gz", hash = "sha256:6b2a32f18d97e16e90827d7a819ffa8dbd8cc245fc4e1fa9d1095b54ef4bd999", size = 271342, upload-time = "2026-05-25T22:17:48.841Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/d2/c3eedaef57de65c3cc5f8dc244cf12d09c84ad258a479055aad6db23206c/httptools-0.8.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:ed377e64805bdba4943c82717333f8f8603a13b09aff9cead2717c6c817fb168", size = 208428, upload-time = "2026-05-25T22:16:59.717Z" }, + { url = "https://files.pythonhosted.org/packages/f1/94/dfe435d90d0ef61ec0f2cc3d480eef78c59727c6c2ce039f433882f6131a/httptools-0.8.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:9518c406d7b310f05adb1a37f80acabac40504a575d7c0da6d3e365c695ac20d", size = 113366, upload-time = "2026-05-25T22:17:00.795Z" }, + { url = "https://files.pythonhosted.org/packages/cc/d4/13025f1a56e615dcb331e0bbe2d9a1143212b58c263385fc5d2e558f5bac/httptools-0.8.0-cp311-cp311-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:57278e6fa0424c42a8a3e454828ab4f0aff27b40cddf9679579b98c6dce6a376", size = 464676, upload-time = "2026-05-25T22:17:02.014Z" }, + { url = "https://files.pythonhosted.org/packages/bf/95/4c1c26c0b985f8a3331682d802598f14e32dc41bf7509266eb2c04ad4801/httptools-0.8.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bbb8caadb2b742d293169d2b458b5c001ef70e3158704aa3d3ef9597624c5d1d", size = 464235, upload-time = "2026-05-25T22:17:03.109Z" }, + { url = "https://files.pythonhosted.org/packages/a2/82/6735be2b0ca527718c431cdb8e5f70c3862c0844a687df0f572c51e11497/httptools-0.8.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:52dd695b865fe96d9d2b16b64a895f3f57bf3cb064e8383cd3b5713a069e8085", size = 449809, upload-time = "2026-05-25T22:17:04.443Z" }, + { url = "https://files.pythonhosted.org/packages/b5/f9/5811c74f37a758c8a4aa3dc430375119d335947e883efc4664d8f3559a41/httptools-0.8.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:20b4aac66ff65f7db06a375808b78f42a94970aa22e826b3cb2b43eb09174124", size = 452174, upload-time = "2026-05-25T22:17:05.476Z" }, + { url = "https://files.pythonhosted.org/packages/cc/94/97b75870dea07b71e3ec535cebe525b08d723152e4c7d13fa887e51f4de2/httptools-0.8.0-cp311-cp311-win_amd64.whl", hash = "sha256:a1b4c8e7a489a0d750d91894e9a8cdc295838f1924c0ca903ae993456fddec07", size = 90991, upload-time = "2026-05-25T22:17:06.75Z" }, + { url = "https://files.pythonhosted.org/packages/14/88/1d21a36da8f5cb0fa49eafd4b169eba5608d57e75bbcf61845cbc6243216/httptools-0.8.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:880490234c10f70a9830743097e8958d6e4b9f5a0ffc24515023afeef984054d", size = 208247, upload-time = "2026-05-25T22:17:07.843Z" }, + { url = "https://files.pythonhosted.org/packages/a5/42/cc4feea2945cb3051038f090c9b36bd5b8a9d7f5a894a506a8983e33fd1c/httptools-0.8.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:5931891fb7b441b8a3853cf1b85c82c903defce084dd5f6771ca46e31bf862c5", size = 113064, upload-time = "2026-05-25T22:17:09.136Z" }, + { url = "https://files.pythonhosted.org/packages/e3/a6/febbb8b8db0f58b38e44ad6cb946e6a255ae49b55f2e8543408fb7501ccd/httptools-0.8.0-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:b15fc622b0f869d19207c4089a501d9bcc63ca5e071ffdd2f03f922df882dcb2", size = 523851, upload-time = "2026-05-25T22:17:10.106Z" }, + { url = "https://files.pythonhosted.org/packages/b7/e4/f90a0df0b83beff265b7e3b65f2a4cefd95792d4be0ac3e16049f2acd3c2/httptools-0.8.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:425f83884fd6343828d8c565f046cb72b6d19063f6924093e11bcd8e1548cd09", size = 518842, upload-time = "2026-05-25T22:17:11.218Z" }, + { url = "https://files.pythonhosted.org/packages/9e/2d/0c9ac76dd2c893841fbf6498d6acec4f2442e1b7067f6e3e316a80e494e8/httptools-0.8.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ef7c3c97f4311c7be57e2986629df89d49cb434dbff78eafcd48c2bff986b15a", size = 501238, upload-time = "2026-05-25T22:17:12.728Z" }, + { url = "https://files.pythonhosted.org/packages/ca/42/906adc91ae3a5fa9c59c0a2f21c139725bd7e5b41ae6acd485cd14123ebf/httptools-0.8.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a1afd7c9fbff0d9f5d489c4ce2768bd09c84a46ddefc7161e6aa82ae35c85745", size = 509567, upload-time = "2026-05-25T22:17:13.842Z" }, + { url = "https://files.pythonhosted.org/packages/05/0b/4240efeb672751ee5b9b380cb0e3fdc050bc05f68adc7a8aefc4fcd9a69a/httptools-0.8.0-cp312-cp312-win_amd64.whl", hash = "sha256:cd96f29b4bab1d42fa6e3d008711c75e0f79e94e06827330160e3a304227f150", size = 90918, upload-time = "2026-05-25T22:17:15.155Z" }, + { url = "https://files.pythonhosted.org/packages/5e/e5/8cfcabc5546e8022f168be28bcdaa128a240a0befdd03b59d558b4f18bd6/httptools-0.8.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:614ceea8ea606848bece2338ac03b3ce5324bcb4be8dc7d377ed708012fa4db8", size = 205148, upload-time = "2026-05-25T22:17:16.333Z" }, + { url = "https://files.pythonhosted.org/packages/2a/0e/0fb14848c19a686c8062ff9067c1a48793e3224b47bc5b201535b6036fce/httptools-0.8.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:2d689918c15a013c65ef52d9fd495d766893ab831a2c8d89f2ac5940a5df847c", size = 111368, upload-time = "2026-05-25T22:17:17.586Z" }, + { url = "https://files.pythonhosted.org/packages/2e/1b/46f1cecf06b9bbde8e4b8c88034ac7908989e5ff7a3a388ef38392949c1f/httptools-0.8.0-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:eb3028cca2fc0a6d720e52ef61d8ebb62fcbfeb1de56874546d858d3f25a26b7", size = 486447, upload-time = "2026-05-25T22:17:18.564Z" }, + { url = "https://files.pythonhosted.org/packages/77/00/258bfc0837221f81d9725c45f9b948a6a6b2994a147a4fb66e85100c668f/httptools-0.8.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:88bdd940f2b5d487b4d032c6afa5489a7dc4694410d43de3c38c4fb3af0dc45d", size = 482448, upload-time = "2026-05-25T22:17:19.912Z" }, + { url = "https://files.pythonhosted.org/packages/04/ab/d1cef3b5523f4d272a70f42a776c3169a2dddfe3a54de4b2ce4a36341528/httptools-0.8.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:6a43c9dd399758ccc0531acb0a3c4a6c299ee893ee9400e9c893b7bdcfae0681", size = 464460, upload-time = "2026-05-25T22:17:20.882Z" }, + { url = "https://files.pythonhosted.org/packages/ce/48/5d1d072442277bb2b3434e0e60690b8e8c23840ef7de8b6ea54040a536d3/httptools-0.8.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:0770728beb05094c809b98e814edff5fef69d26ad7d21185f2f6d5884a0ba683", size = 471312, upload-time = "2026-05-25T22:17:22.085Z" }, + { url = "https://files.pythonhosted.org/packages/0d/66/b96623b27e51a68199ef4efdda0613cced9233fe3062ac74e50749c5ad37/httptools-0.8.0-cp313-cp313-win_amd64.whl", hash = "sha256:7685df791fad561384bfb139e77fde27a1ffd93134e016f95a0db424ffbf77b1", size = 90117, upload-time = "2026-05-25T22:17:23.074Z" }, + { url = "https://files.pythonhosted.org/packages/1a/12/fa3fbf5f9517b273edea2dc982aa82a8c634091e67c590792b729017bc6f/httptools-0.8.0-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:de242a49b5d18e0a8776e654e9f6bf6d89f3875a5c35b425a0e7ce940feb3fd6", size = 206183, upload-time = "2026-05-25T22:17:24.004Z" }, + { url = "https://files.pythonhosted.org/packages/30/fc/5e7c4cb443370f2090a3aba0453a07384d29ff66b7435bb90e77e1037599/httptools-0.8.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:159e9ab5f701ccd42e555a12f1ad8ff69702910fc1c996cf2bb66e5fcb7a231b", size = 112079, upload-time = "2026-05-25T22:17:25.216Z" }, + { url = "https://files.pythonhosted.org/packages/ba/53/771bd891eb0f236f32145d6a1775777ec85745f3cc983a1f23d1a3b8ddfe/httptools-0.8.0-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:c4a9f1707e4823d54dfec6c33fa3697d302aed536ed352a7ebb5a061ddb869d0", size = 481596, upload-time = "2026-05-25T22:17:26.186Z" }, + { url = "https://files.pythonhosted.org/packages/62/42/94e15bc68ce3d423243c45d7f1b0c7561f13844f97dc52ae23182fb65628/httptools-0.8.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d76ad7b951387e3632c8716a9bb03ac5b45c5f16119aa409db0459520887944e", size = 480865, upload-time = "2026-05-25T22:17:27.542Z" }, + { url = "https://files.pythonhosted.org/packages/1c/7c/fe2980fc03723272e30f135b62360b075f513dfe7cc73aef36c7f04012bd/httptools-0.8.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:a3b7387147361c3fd47a0bde763c5c91b5b4cd4dc9989b8ece84ff436c99843b", size = 463189, upload-time = "2026-05-25T22:17:28.546Z" }, + { url = "https://files.pythonhosted.org/packages/15/1b/47fc5fff68acd1bfa20b4734059c9a06cadb88119dcd5258b5b0d21d91c8/httptools-0.8.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f256d6ce930c52ca1cb2a960b7da03548c454e7d28b06059ad41bfe789036ce0", size = 466610, upload-time = "2026-05-25T22:17:29.816Z" }, + { url = "https://files.pythonhosted.org/packages/60/bd/07b13c93ffd9bec9546e0d43f8e19378dd696dbd278511406bc07371ef1f/httptools-0.8.0-cp314-cp314-win_amd64.whl", hash = "sha256:19d1ee275bb59ba2643ba9a3a1e51cc0c788caf2b8df506368e03f56fdd08527", size = 92705, upload-time = "2026-05-25T22:17:31.133Z" }, + { url = "https://files.pythonhosted.org/packages/fd/c4/121648f68ce066d7bd762d6b6d97e620847642d38d54f3d90ff11d947629/httptools-0.8.0-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:de1ed58a974e75d56560acc7e7fed01a454994429456f65209789992e41f2568", size = 215023, upload-time = "2026-05-25T22:17:32.401Z" }, + { url = "https://files.pythonhosted.org/packages/b9/b0/312a062ae741ae3e8baa8c8bf20be81b2e67337b259ab4349bebc7b6142e/httptools-0.8.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:e93c227b595c6926c1acee96891dd9da4be338cfbe82e5cd3bb9d8dd7dc4ac0b", size = 117405, upload-time = "2026-05-25T22:17:33.742Z" }, + { url = "https://files.pythonhosted.org/packages/fc/37/fccd705f795386bb05bf413012fecff2a33e5aa8c2f069096de3e9fd8702/httptools-0.8.0-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:2a021c3a8e65cc125390d72f59b968afca3bdcaff25bd67965e0a055a14946ca", size = 558497, upload-time = "2026-05-25T22:17:34.732Z" }, + { url = "https://files.pythonhosted.org/packages/bd/39/f172e8003576de35f5ba77ff417cf0e34429d35dc014deef15afa337a72c/httptools-0.8.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:48774d39cbb70e2b1f71f88852a3087ae1d3a1eb80482bb48c13067ab080c14f", size = 571585, upload-time = "2026-05-25T22:17:35.813Z" }, + { url = "https://files.pythonhosted.org/packages/3e/b9/f5564760af99f3dbbf3f9104dc00e5da27e96cf433c6bdcf77617f70bf3f/httptools-0.8.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:88eead8ec8680a9f146c655bc88445a325bd7921cfd8194c7337e9467282427d", size = 543297, upload-time = "2026-05-25T22:17:37.08Z" }, + { url = "https://files.pythonhosted.org/packages/99/67/8d9f2c313618e161b82f3873188e7196126da1d6e29688df40eb3997c77a/httptools-0.8.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:2c032fa028f46871ec7e1fc59fc15e8023eab3e6bbe6ece786a1611719a5d081", size = 539535, upload-time = "2026-05-25T22:17:38.032Z" }, + { url = "https://files.pythonhosted.org/packages/48/63/b906c01e53f50d432c0defe43ce52764a111dc1bdd028bafbeb54dcfd008/httptools-0.8.0-cp314-cp314t-win_amd64.whl", hash = "sha256:384c17174464c8e873398b7af24f0b1f44d992c820328413951a625323155d77", size = 108209, upload-time = "2026-05-25T22:17:39.473Z" }, +] + [[package]] name = "httpx" version = "0.28.1" @@ -2381,6 +2457,148 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/15/41/ac2dfdbc1f60c7af4f994c7a335cfa7040c01642b605d65f611cecc2a1e4/uvicorn-0.47.0-py3-none-any.whl", hash = "sha256:2c5715bc12d1892d84752049f400cd1c3cb018514967fdfeb97640443a6a9432", size = 71301, upload-time = "2026-05-14T18:16:51.762Z" }, ] +[package.optional-dependencies] +standard = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "httptools" }, + { name = "python-dotenv" }, + { name = "pyyaml" }, + { name = "uvloop", marker = "platform_python_implementation != 'PyPy' and sys_platform != 'cygwin' and sys_platform != 'win32'" }, + { name = "watchfiles" }, + { name = "websockets" }, +] + +[[package]] +name = "uvloop" +version = "0.22.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/06/f0/18d39dbd1971d6d62c4629cc7fa67f74821b0dc1f5a77af43719de7936a7/uvloop-0.22.1.tar.gz", hash = "sha256:6c84bae345b9147082b17371e3dd5d42775bddce91f885499017f4607fdaf39f", size = 2443250, upload-time = "2025-10-16T22:17:19.342Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c7/d5/69900f7883235562f1f50d8184bb7dd84a2fb61e9ec63f3782546fdbd057/uvloop-0.22.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:c60ebcd36f7b240b30788554b6f0782454826a0ed765d8430652621b5de674b9", size = 1352420, upload-time = "2025-10-16T22:16:21.187Z" }, + { url = "https://files.pythonhosted.org/packages/a8/73/c4e271b3bce59724e291465cc936c37758886a4868787da0278b3b56b905/uvloop-0.22.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:3b7f102bf3cb1995cfeaee9321105e8f5da76fdb104cdad8986f85461a1b7b77", size = 748677, upload-time = "2025-10-16T22:16:22.558Z" }, + { url = "https://files.pythonhosted.org/packages/86/94/9fb7fad2f824d25f8ecac0d70b94d0d48107ad5ece03769a9c543444f78a/uvloop-0.22.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:53c85520781d84a4b8b230e24a5af5b0778efdb39142b424990ff1ef7c48ba21", size = 3753819, upload-time = "2025-10-16T22:16:23.903Z" }, + { url = "https://files.pythonhosted.org/packages/74/4f/256aca690709e9b008b7108bc85fba619a2bc37c6d80743d18abad16ee09/uvloop-0.22.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:56a2d1fae65fd82197cb8c53c367310b3eabe1bbb9fb5a04d28e3e3520e4f702", size = 3804529, upload-time = "2025-10-16T22:16:25.246Z" }, + { url = "https://files.pythonhosted.org/packages/7f/74/03c05ae4737e871923d21a76fe28b6aad57f5c03b6e6bfcfa5ad616013e4/uvloop-0.22.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:40631b049d5972c6755b06d0bfe8233b1bd9a8a6392d9d1c45c10b6f9e9b2733", size = 3621267, upload-time = "2025-10-16T22:16:26.819Z" }, + { url = "https://files.pythonhosted.org/packages/75/be/f8e590fe61d18b4a92070905497aec4c0e64ae1761498cad09023f3f4b3e/uvloop-0.22.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:535cc37b3a04f6cd2c1ef65fa1d370c9a35b6695df735fcff5427323f2cd5473", size = 3723105, upload-time = "2025-10-16T22:16:28.252Z" }, + { url = "https://files.pythonhosted.org/packages/3d/ff/7f72e8170be527b4977b033239a83a68d5c881cc4775fca255c677f7ac5d/uvloop-0.22.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:fe94b4564e865d968414598eea1a6de60adba0c040ba4ed05ac1300de402cd42", size = 1359936, upload-time = "2025-10-16T22:16:29.436Z" }, + { url = "https://files.pythonhosted.org/packages/c3/c6/e5d433f88fd54d81ef4be58b2b7b0cea13c442454a1db703a1eea0db1a59/uvloop-0.22.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:51eb9bd88391483410daad430813d982010f9c9c89512321f5b60e2cddbdddd6", size = 752769, upload-time = "2025-10-16T22:16:30.493Z" }, + { url = "https://files.pythonhosted.org/packages/24/68/a6ac446820273e71aa762fa21cdcc09861edd3536ff47c5cd3b7afb10eeb/uvloop-0.22.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:700e674a166ca5778255e0e1dc4e9d79ab2acc57b9171b79e65feba7184b3370", size = 4317413, upload-time = "2025-10-16T22:16:31.644Z" }, + { url = "https://files.pythonhosted.org/packages/5f/6f/e62b4dfc7ad6518e7eff2516f680d02a0f6eb62c0c212e152ca708a0085e/uvloop-0.22.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7b5b1ac819a3f946d3b2ee07f09149578ae76066d70b44df3fa990add49a82e4", size = 4426307, upload-time = "2025-10-16T22:16:32.917Z" }, + { url = "https://files.pythonhosted.org/packages/90/60/97362554ac21e20e81bcef1150cb2a7e4ffdaf8ea1e5b2e8bf7a053caa18/uvloop-0.22.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e047cc068570bac9866237739607d1313b9253c3051ad84738cbb095be0537b2", size = 4131970, upload-time = "2025-10-16T22:16:34.015Z" }, + { url = "https://files.pythonhosted.org/packages/99/39/6b3f7d234ba3964c428a6e40006340f53ba37993f46ed6e111c6e9141d18/uvloop-0.22.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:512fec6815e2dd45161054592441ef76c830eddaad55c8aa30952e6fe1ed07c0", size = 4296343, upload-time = "2025-10-16T22:16:35.149Z" }, + { url = "https://files.pythonhosted.org/packages/89/8c/182a2a593195bfd39842ea68ebc084e20c850806117213f5a299dfc513d9/uvloop-0.22.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:561577354eb94200d75aca23fbde86ee11be36b00e52a4eaf8f50fb0c86b7705", size = 1358611, upload-time = "2025-10-16T22:16:36.833Z" }, + { url = "https://files.pythonhosted.org/packages/d2/14/e301ee96a6dc95224b6f1162cd3312f6d1217be3907b79173b06785f2fe7/uvloop-0.22.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:1cdf5192ab3e674ca26da2eada35b288d2fa49fdd0f357a19f0e7c4e7d5077c8", size = 751811, upload-time = "2025-10-16T22:16:38.275Z" }, + { url = "https://files.pythonhosted.org/packages/b7/02/654426ce265ac19e2980bfd9ea6590ca96a56f10c76e63801a2df01c0486/uvloop-0.22.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6e2ea3d6190a2968f4a14a23019d3b16870dd2190cd69c8180f7c632d21de68d", size = 4288562, upload-time = "2025-10-16T22:16:39.375Z" }, + { url = "https://files.pythonhosted.org/packages/15/c0/0be24758891ef825f2065cd5db8741aaddabe3e248ee6acc5e8a80f04005/uvloop-0.22.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0530a5fbad9c9e4ee3f2b33b148c6a64d47bbad8000ea63704fa8260f4cf728e", size = 4366890, upload-time = "2025-10-16T22:16:40.547Z" }, + { url = "https://files.pythonhosted.org/packages/d2/53/8369e5219a5855869bcee5f4d317f6da0e2c669aecf0ef7d371e3d084449/uvloop-0.22.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:bc5ef13bbc10b5335792360623cc378d52d7e62c2de64660616478c32cd0598e", size = 4119472, upload-time = "2025-10-16T22:16:41.694Z" }, + { url = "https://files.pythonhosted.org/packages/f8/ba/d69adbe699b768f6b29a5eec7b47dd610bd17a69de51b251126a801369ea/uvloop-0.22.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:1f38ec5e3f18c8a10ded09742f7fb8de0108796eb673f30ce7762ce1b8550cad", size = 4239051, upload-time = "2025-10-16T22:16:43.224Z" }, + { url = "https://files.pythonhosted.org/packages/90/cd/b62bdeaa429758aee8de8b00ac0dd26593a9de93d302bff3d21439e9791d/uvloop-0.22.1-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:3879b88423ec7e97cd4eba2a443aa26ed4e59b45e6b76aabf13fe2f27023a142", size = 1362067, upload-time = "2025-10-16T22:16:44.503Z" }, + { url = "https://files.pythonhosted.org/packages/0d/f8/a132124dfda0777e489ca86732e85e69afcd1ff7686647000050ba670689/uvloop-0.22.1-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:4baa86acedf1d62115c1dc6ad1e17134476688f08c6efd8a2ab076e815665c74", size = 752423, upload-time = "2025-10-16T22:16:45.968Z" }, + { url = "https://files.pythonhosted.org/packages/a3/94/94af78c156f88da4b3a733773ad5ba0b164393e357cc4bd0ab2e2677a7d6/uvloop-0.22.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:297c27d8003520596236bdb2335e6b3f649480bd09e00d1e3a99144b691d2a35", size = 4272437, upload-time = "2025-10-16T22:16:47.451Z" }, + { url = "https://files.pythonhosted.org/packages/b5/35/60249e9fd07b32c665192cec7af29e06c7cd96fa1d08b84f012a56a0b38e/uvloop-0.22.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c1955d5a1dd43198244d47664a5858082a3239766a839b2102a269aaff7a4e25", size = 4292101, upload-time = "2025-10-16T22:16:49.318Z" }, + { url = "https://files.pythonhosted.org/packages/02/62/67d382dfcb25d0a98ce73c11ed1a6fba5037a1a1d533dcbb7cab033a2636/uvloop-0.22.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:b31dc2fccbd42adc73bc4e7cdbae4fc5086cf378979e53ca5d0301838c5682c6", size = 4114158, upload-time = "2025-10-16T22:16:50.517Z" }, + { url = "https://files.pythonhosted.org/packages/f0/7a/f1171b4a882a5d13c8b7576f348acfe6074d72eaf52cccef752f748d4a9f/uvloop-0.22.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:93f617675b2d03af4e72a5333ef89450dfaa5321303ede6e67ba9c9d26878079", size = 4177360, upload-time = "2025-10-16T22:16:52.646Z" }, + { url = "https://files.pythonhosted.org/packages/79/7b/b01414f31546caf0919da80ad57cbfe24c56b151d12af68cee1b04922ca8/uvloop-0.22.1-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:37554f70528f60cad66945b885eb01f1bb514f132d92b6eeed1c90fd54ed6289", size = 1454790, upload-time = "2025-10-16T22:16:54.355Z" }, + { url = "https://files.pythonhosted.org/packages/d4/31/0bb232318dd838cad3fa8fb0c68c8b40e1145b32025581975e18b11fab40/uvloop-0.22.1-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:b76324e2dc033a0b2f435f33eb88ff9913c156ef78e153fb210e03c13da746b3", size = 796783, upload-time = "2025-10-16T22:16:55.906Z" }, + { url = "https://files.pythonhosted.org/packages/42/38/c9b09f3271a7a723a5de69f8e237ab8e7803183131bc57c890db0b6bb872/uvloop-0.22.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:badb4d8e58ee08dad957002027830d5c3b06aea446a6a3744483c2b3b745345c", size = 4647548, upload-time = "2025-10-16T22:16:57.008Z" }, + { url = "https://files.pythonhosted.org/packages/c1/37/945b4ca0ac27e3dc4952642d4c900edd030b3da6c9634875af6e13ae80e5/uvloop-0.22.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b91328c72635f6f9e0282e4a57da7470c7350ab1c9f48546c0f2866205349d21", size = 4467065, upload-time = "2025-10-16T22:16:58.206Z" }, + { url = "https://files.pythonhosted.org/packages/97/cc/48d232f33d60e2e2e0b42f4e73455b146b76ebe216487e862700457fbf3c/uvloop-0.22.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:daf620c2995d193449393d6c62131b3fbd40a63bf7b307a1527856ace637fe88", size = 4328384, upload-time = "2025-10-16T22:16:59.36Z" }, + { url = "https://files.pythonhosted.org/packages/e4/16/c1fd27e9549f3c4baf1dc9c20c456cd2f822dbf8de9f463824b0c0357e06/uvloop-0.22.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6cde23eeda1a25c75b2e07d39970f3374105d5eafbaab2a4482be82f272d5a5e", size = 4296730, upload-time = "2025-10-16T22:17:00.744Z" }, +] + +[[package]] +name = "watchfiles" +version = "1.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/cd/41/5e1a4bb12aac5f1493fa1bdc11154eca3b258ca4eba65d39c473fe19d8e9/watchfiles-1.2.0.tar.gz", hash = "sha256:c995fba777f1ea992f090f9236e9284cf7a5d1a0130dd5a3d82c598cacd76838", size = 108252, upload-time = "2026-05-18T04:32:04.251Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fc/3d/8024c801df84d1587740d0359e7fdd80afeae3d159011f3d5376dd82f18e/watchfiles-1.2.0-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:704fd259e332e01f9b9c178f4bce9e49027e5587cc2600eeeaf8e76e1c846201", size = 400242, upload-time = "2026-05-18T04:31:19.014Z" }, + { url = "https://files.pythonhosted.org/packages/87/5b/f4dfd45323e949984a3a7f9dc31d1cbb049921e7d98253488dda72ccdaa9/watchfiles-1.2.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:6543cf55d170003296d185c0af981f3e1311564907e1f4e08671fc7693a890a5", size = 394562, upload-time = "2026-05-18T04:30:08.46Z" }, + { url = "https://files.pythonhosted.org/packages/98/d8/19483ef075d601c409bce8bcbb5c0f81a10876fff870400568f08ce484a1/watchfiles-1.2.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:89d8c2394a065ca86f5d2910ff263ae67c127e1376ccc4f9fc35c71db879f80a", size = 456611, upload-time = "2026-05-18T04:30:45.723Z" }, + { url = "https://files.pythonhosted.org/packages/b1/6a/cc81fbe7ee42f2f22e661a6e12def7807e01b14b2f39e0ff83fd373fd307/watchfiles-1.2.0-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:772b80df316480d894a0e3165fdd19cf77f5d17f9a787f94029465ad0e3529d1", size = 461379, upload-time = "2026-05-18T04:31:29.292Z" }, + { url = "https://files.pythonhosted.org/packages/b1/57/7e669002082c0a0f4fb5113bb70125f7110124b846b0a11bc5ae8e90eac1/watchfiles-1.2.0-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d158cd89df6053823533e06fb1d73c549133bff5f0396170c0e53d9559340717", size = 493556, upload-time = "2026-05-18T04:30:05.44Z" }, + { url = "https://files.pythonhosted.org/packages/45/7d/f60a2b19807b21fe8281f3a8da4f59eef0d5f96825ac4680ba2d4f2ebf91/watchfiles-1.2.0-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d516b3283a758e087841aedb8031549fb41ced08f3db10aa6d2bf32dc042525b", size = 575255, upload-time = "2026-05-18T04:30:40.568Z" }, + { url = "https://files.pythonhosted.org/packages/bd/49/77f5b5e6efbcd57482f74948ebb1b97e5c0046d6b61475042d830c84b3ff/watchfiles-1.2.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:53b2290c92e0506d102cd448fbc610d87079553f86caa39d67440856a8b8bba5", size = 467052, upload-time = "2026-05-18T04:31:17.942Z" }, + { url = "https://files.pythonhosted.org/packages/ee/5a/73e2959af1b97fd5d556f9a8bdba017be23ceeef731869d5eaa0a753d5a3/watchfiles-1.2.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a711b51aec4370d0dcda5b6c09463206f133a5759341d7744b953a7b62e1100e", size = 456858, upload-time = "2026-05-18T04:30:30.182Z" }, + { url = "https://files.pythonhosted.org/packages/50/57/1bc8c27fad7e6c19bddee15d276dbb6ab72480ec01c127afff1673aee417/watchfiles-1.2.0-cp311-cp311-manylinux_2_31_riscv64.whl", hash = "sha256:e2ca07fa7d89195ec0865d3d285666286740bfa83d83e5cee204043a31ecc165", size = 467579, upload-time = "2026-05-18T04:32:15.897Z" }, + { url = "https://files.pythonhosted.org/packages/09/6c/3c2e44edba3553c5e3c3b8c8a2a6dee6b9e12ae2cf4bd2378bebf9dc3038/watchfiles-1.2.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:e0618518f282c4ebff60f5e5b1247b6d91bb8b9f4476947563a1e74acc66f3c6", size = 633253, upload-time = "2026-05-18T04:31:37.123Z" }, + { url = "https://files.pythonhosted.org/packages/30/c2/d8c84a882ab39bbefcc4915ab3e91830b7a7e990c5570b0b69075aba3faf/watchfiles-1.2.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:0d191c054d0715c3c95c99df9b8dbf6fd096d8c1e021e8f212e1bd8bc444ccb5", size = 660713, upload-time = "2026-05-18T04:31:24.62Z" }, + { url = "https://files.pythonhosted.org/packages/a9/07/f97736a5fc605364fe67b25e9fa4a6965dfd4840d50c406ada507e9d735f/watchfiles-1.2.0-cp311-cp311-win32.whl", hash = "sha256:9342472aff9b093c5acd4f6d8f70ae0937964ab56542502bcf5579782da69ae8", size = 277222, upload-time = "2026-05-18T04:31:21.131Z" }, + { url = "https://files.pythonhosted.org/packages/cf/99/2b04981977fc2608afd60360d928c6aecf6b950292ca221d98f4005f6694/watchfiles-1.2.0-cp311-cp311-win_amd64.whl", hash = "sha256:dbd6c97045dad81227c8d040173da044c1de08de64a5ea8b555da4aee1d5fa22", size = 290274, upload-time = "2026-05-18T04:31:45.966Z" }, + { url = "https://files.pythonhosted.org/packages/3c/74/f7f58a7075ee9cf612b0cfcddb78b8cd8234f0742d6f0075cf0da2dde1c6/watchfiles-1.2.0-cp311-cp311-win_arm64.whl", hash = "sha256:57a2d9fa4fb4c2ecae57b13dfff2c7ab53e21a2ba674fe9f05506680fcdcc0d7", size = 283460, upload-time = "2026-05-18T04:31:39.126Z" }, + { url = "https://files.pythonhosted.org/packages/b8/2f/e42c992d2afda3108ea1c02acecc991b9f31d05c14adc2a7cee9ee211fc4/watchfiles-1.2.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:bc13eb17538be00c874699dc0abe4ee2bc8d50bb1166a6b9e175ef3fd7eb8f26", size = 400115, upload-time = "2026-05-18T04:32:02.06Z" }, + { url = "https://files.pythonhosted.org/packages/5f/8f/6af2ea19065c91d8b0ea3516fdfc8c0d349f407e8e9fbf4e5a17360de8ad/watchfiles-1.2.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:2d95ddc1eb6914154253d239089900813f6a767e174b8e6a50e7fdacb7e4236c", size = 393659, upload-time = "2026-05-18T04:30:50.951Z" }, + { url = "https://files.pythonhosted.org/packages/13/01/b32a967c56fb3e3e5be3db52c3d3b87fa4513aa367d8ed1ad96d42952e5f/watchfiles-1.2.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8f70d8b291ef6e88d19b1f297a6905ddb978888d9272b0d05e6f53309856bcfc", size = 453207, upload-time = "2026-05-18T04:31:04.231Z" }, + { url = "https://files.pythonhosted.org/packages/04/98/97557a812180338cb1abd32e1cffcc4588f59b5f23e0cb006b2ba95ba64a/watchfiles-1.2.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:56d8641cf834c2836922899105bd3ce3d0dfc69291d52edf0b4d0436829b34c0", size = 459273, upload-time = "2026-05-18T04:31:50.377Z" }, + { url = "https://files.pythonhosted.org/packages/e8/a8/b4b08dcb7653b8087c6586f7ce649505900e866bbcfe40dc9587af02e686/watchfiles-1.2.0-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2581a94056e55d7d0a31a823ea92bf73749c489ca2285bfdc0fbe6b2bb49d50c", size = 489927, upload-time = "2026-05-18T04:31:42.485Z" }, + { url = "https://files.pythonhosted.org/packages/50/94/3dceea03545d2e5ddfd839f0ddd5e1cecbf1697b5a428d5ba11cef6af95d/watchfiles-1.2.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:41bc1199f7523b3f82843c88cbb979180c949caef0342cf90968f178e5d49b01", size = 570476, upload-time = "2026-05-18T04:31:03.071Z" }, + { url = "https://files.pythonhosted.org/packages/cc/f2/d39a5450c3532092b91f81d274360e613c2371bc874a89c7a1a3c5e8d138/watchfiles-1.2.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:7571e4464cb6e434958f867f7f730b8ab0b75e3f8e5eac0499168486ab3c33a8", size = 465650, upload-time = "2026-05-18T04:30:12.701Z" }, + { url = "https://files.pythonhosted.org/packages/22/24/ed72f68cbc1333ca9b9f2200aa048bb6658ae41709bc1caad4310f4bdffd/watchfiles-1.2.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e53a384f76b631c3ae5334ce6a52f0baa3a911eb94a4eac7f160079868b716d5", size = 456398, upload-time = "2026-05-18T04:30:13.784Z" }, + { url = "https://files.pythonhosted.org/packages/0d/64/982ef4a4e5bab5b6e5b6becc8cd5e732f6130a78b855f0abec6439a9a135/watchfiles-1.2.0-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:d20029a60a71a052a24c4db7673bc4de39ab89adbaccbfb5d67987c5d73f424d", size = 465140, upload-time = "2026-05-18T04:31:52.111Z" }, + { url = "https://files.pythonhosted.org/packages/a0/0c/95282abf4ed680b6096010bcfc30c5fa7a041fc5aa5a2ad17a2cc6c75bba/watchfiles-1.2.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:2cb93af48550faf1cea04c303107c8b75833de7013e57ce27d3b8d21d8d0f58c", size = 630259, upload-time = "2026-05-18T04:31:25.676Z" }, + { url = "https://files.pythonhosted.org/packages/30/45/607c1de1530c4bdcf2cf1d1ecc2505ddba5d96bd43ba9f2b0e79876f850f/watchfiles-1.2.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:2995c176de7692b86a2e4c58d9ec718f753150a979cb4a754e2b4ffa38e70906", size = 659859, upload-time = "2026-05-18T04:30:24.333Z" }, + { url = "https://files.pythonhosted.org/packages/fa/08/d9e2e0f9e8e6791d33aefc694ad7eefa7f901f63caff84a81ded38692f9c/watchfiles-1.2.0-cp312-cp312-win32.whl", hash = "sha256:7a2cffd17d27d2ecbb310c2b1d8174f222a5495b1a721894afa88ec11e25b898", size = 275480, upload-time = "2026-05-18T04:30:31.307Z" }, + { url = "https://files.pythonhosted.org/packages/1c/e6/9d42569c0102645cc8cea5d8c7d8a1e9d4ada2cb7f05f75e554b8aa2202a/watchfiles-1.2.0-cp312-cp312-win_amd64.whl", hash = "sha256:f155b3a1b2a5fc89cdc70d47ee5d54e3b75e88efa34982028a35daef9ba00379", size = 288718, upload-time = "2026-05-18T04:32:10.745Z" }, + { url = "https://files.pythonhosted.org/packages/0a/26/88e0dc6ee3898169d7fa22bb6a69cabf2502d2ee25cb8c876d1262d204f8/watchfiles-1.2.0-cp312-cp312-win_arm64.whl", hash = "sha256:8fa585ede612ee9f9e91b18bebf9ba11b9ae29a4e3a0d0cf6fca3e382133f0d5", size = 281026, upload-time = "2026-05-18T04:30:22.23Z" }, + { url = "https://files.pythonhosted.org/packages/d1/4d/70a7feced9f87e2ff26dba42667290f41694fc64646c67261fbb8cab5d5c/watchfiles-1.2.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:01ea8d66f0693b9b60a6541c8d10263091ca9a9060d242f3c1f3143f9aad2c98", size = 399730, upload-time = "2026-05-18T04:31:38.162Z" }, + { url = "https://files.pythonhosted.org/packages/31/3a/0da302f2307aee316922806ebd5726c542cbd787c938271cf14a074c7daf/watchfiles-1.2.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:7ba0480b9a74af058f43b337e937a451e109295c420916d68ad24e3dc02f5e44", size = 392842, upload-time = "2026-05-18T04:30:27.051Z" }, + { url = "https://files.pythonhosted.org/packages/db/ef/d5bdb705c224dbc256aa0c1ec47bf4e61ec52558f2afb44a71a1fe4d7015/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4f34e26a19f91f710c08e0183429f0d1d15df734e6bc78c31e77b9ea9c433658", size = 452989, upload-time = "2026-05-18T04:31:11.945Z" }, + { url = "https://files.pythonhosted.org/packages/71/29/5495f2c1661949ef7a35e4d71111d129cfe7606414a26887a919d0a55406/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:b4e77f6a55f858504069abd35d336a637555c09bca453dde1ee1e5ada8a6a1fb", size = 458978, upload-time = "2026-05-18T04:30:52.606Z" }, + { url = "https://files.pythonhosted.org/packages/d5/8c/7f9c07c433811c2fffd93e13fdfb7135de9aab5f2ae41be08960fa0047dc/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0cb4d80e212f116474a545c21c912b445f16bb0cef9e6a73a498164223e14e2f", size = 490248, upload-time = "2026-05-18T04:31:36.003Z" }, + { url = "https://files.pythonhosted.org/packages/3c/11/d93632febc52fbc21be90231bb7c17fd5387f46c9076fd40a5f9c2ae6910/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b974946a10af379d425e2eef5b62f5c6ebeaccf91d45eaad6f5b27ecd4f91aa0", size = 571847, upload-time = "2026-05-18T04:31:10.862Z" }, + { url = "https://files.pythonhosted.org/packages/55/b4/383173e73aabb07ad1d9c7aa859d95437ac46a6d6a1e11005facda0c9d19/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:86bc13c25a8d1fcd70b51d0ce7c9b65e90de5666fcbfd3e34957cc73ee19aeb5", size = 465974, upload-time = "2026-05-18T04:30:17.006Z" }, + { url = "https://files.pythonhosted.org/packages/a7/6c/89b1a230a78f57c52dd8893adb1f92f94411721b6ec12596c56d98c74356/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ca148d73dea36c9763aaa351e4d7a51780ec1584217c45276f4fe8239c768b71", size = 454782, upload-time = "2026-05-18T04:30:35.656Z" }, + { url = "https://files.pythonhosted.org/packages/24/62/1732118367cfff0a9fce3bf62ff4bfded09ef5df21d9d446b858b3f70a96/watchfiles-1.2.0-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:c525543d91961c6955b2636b308569e84a1d1c5f5f2932041ab9ef46422f43e3", size = 465182, upload-time = "2026-05-18T04:30:20.846Z" }, + { url = "https://files.pythonhosted.org/packages/28/96/716f7e5f51339bf22963f3345f9f27d7f3b30e2eadc597e257c881dd3c53/watchfiles-1.2.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:a204794696ffb8f9b10fba6f7cb5216d42f3b2b71860ccac6b6e42f5f10973b0", size = 629841, upload-time = "2026-05-18T04:31:05.397Z" }, + { url = "https://files.pythonhosted.org/packages/4c/fe/c40783950fd771ccf66ab3ec2722d188a9af1c7f96c6e811f36e40c6e03f/watchfiles-1.2.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:10d86db20695afe7997ac9e1717637d6714a8d0220458c33f3d2061f54cec427", size = 658028, upload-time = "2026-05-18T04:31:48.22Z" }, + { url = "https://files.pythonhosted.org/packages/71/72/4508db1856d1d87fcbb3b63f4839bab1b5682cb0e8d224d122263c09654a/watchfiles-1.2.0-cp313-cp313-win32.whl", hash = "sha256:eb283ee99e21ad6443c8cdb06ac5b34b1308c329cbdf03fa02b445363714c799", size = 275183, upload-time = "2026-05-18T04:30:59.57Z" }, + { url = "https://files.pythonhosted.org/packages/f9/36/14b76ca57652e5cc5fd1c11f32a261292c08a0d19a00351013c2549cbfb2/watchfiles-1.2.0-cp313-cp313-win_amd64.whl", hash = "sha256:a0f27f01bee51861392bb6b7c4fdb290b27d1eb194e9e28788d68102a0e898d9", size = 288059, upload-time = "2026-05-18T04:32:07.937Z" }, + { url = "https://files.pythonhosted.org/packages/1b/8d/0a85e395398d8d20fadfe5c5d32c726eee17a519e78fb356f2cf7531bffe/watchfiles-1.2.0-cp313-cp313-win_arm64.whl", hash = "sha256:3651aa7058595e9cfb75d35dd5ada2bf9f48a5b8a0f3562821d3e210c507e077", size = 280186, upload-time = "2026-05-18T04:31:54.484Z" }, + { url = "https://files.pythonhosted.org/packages/37/68/36db056f1fdcc5f07302f56e631774d6835bcd6fa3ace402304621d5f9e5/watchfiles-1.2.0-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:faea288b6f0ab1902ef08f4ca6de005dccf856c4e0c4f21b8c5fce02d90a1b08", size = 399031, upload-time = "2026-05-18T04:30:44.576Z" }, + { url = "https://files.pythonhosted.org/packages/c1/64/01a9d6f66a82a5c101ce939274106cc72759d62427e153f01edd2b9f87c2/watchfiles-1.2.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:01859b11fd9fbca670f4d5da00fbac282cfea9bd67a2125d8b2833a3b5617ea9", size = 391205, upload-time = "2026-05-18T04:30:25.413Z" }, + { url = "https://files.pythonhosted.org/packages/84/2c/0a44fe058cb4bb7b8ede6b6670698bbb7c0400740e378d00022189b7b31d/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:fff610d7bb2256a317bb1e96f0d7862c7aa8076733ee5df0fd41bbe76a24a4f4", size = 451892, upload-time = "2026-05-18T04:32:14.005Z" }, + { url = "https://files.pythonhosted.org/packages/67/a1/351e0d56cd35e6488b5c8b4fb11a809a5bc923e8fe8fed9faf8920be0c89/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:b141a4891c995a039cd89e9a49e62df1dc8a559a5d1a6e4c7106d16c12777a55", size = 458867, upload-time = "2026-05-18T04:31:22.279Z" }, + { url = "https://files.pythonhosted.org/packages/d5/7d/9d09605187f1b838998624049fcf8bf47b73c1a3b76901fcac1782f62277/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f22943b7770483f6ea0721c6b11d022947a98eb0acae14694de034f4d0d38925", size = 490217, upload-time = "2026-05-18T04:31:43.657Z" }, + { url = "https://files.pythonhosted.org/packages/60/5d/a17a16eccb182f04188cd308ec24b1a71a9b5c4e7098269cf35d9fa56d02/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1bc6195825b7dcd217968bb1f801a60fd4c16e8eeab5bedc7fe917d7d5995ab4", size = 571458, upload-time = "2026-05-18T04:32:11.875Z" }, + { url = "https://files.pythonhosted.org/packages/d3/3d/4dd457062083ab1938e5dfd45032eb425cee2ac817287ca8ff4356183e5d/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d4a4b147f5dca2a5d325a06a832fb43f345751adfbc63204aec30e0d9ca965a2", size = 464707, upload-time = "2026-05-18T04:30:43.492Z" }, + { url = "https://files.pythonhosted.org/packages/c6/71/ea8c57b128f5383de74d0c7d2d9c57ad7c9a65a930c451bd25d524b295b7/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4543579a9bdb0c9560039b4ffddbdb39545707659fbc430ce4c10f3f68d557f9", size = 454663, upload-time = "2026-05-18T04:30:16.061Z" }, + { url = "https://files.pythonhosted.org/packages/53/fd/2e812bf938406d7db351f0703ddd3fc6c061cf30d96153a77bc79a943a44/watchfiles-1.2.0-cp313-cp313t-manylinux_2_31_riscv64.whl", hash = "sha256:20aa0e708b920bde876a4aa82dc7dd6ebea228a63a67cda6632c2fc87b787efa", size = 463537, upload-time = "2026-05-18T04:31:44.9Z" }, + { url = "https://files.pythonhosted.org/packages/86/56/d17a7f1dd1bc3035f1072694a551301272f1739c2d8e319c927cb9e29b38/watchfiles-1.2.0-cp313-cp313t-musllinux_1_1_aarch64.whl", hash = "sha256:d413349d565dab74297f2a63e84a097936be69bf8f3b3801f27f380e32040f44", size = 629194, upload-time = "2026-05-18T04:31:14.141Z" }, + { url = "https://files.pythonhosted.org/packages/be/06/f1ff66bf5cae50aa4062779a0ecd0bbaf15e466195719074078947d9a17d/watchfiles-1.2.0-cp313-cp313t-musllinux_1_1_x86_64.whl", hash = "sha256:f28b2725eb8cce327b9b3ab02415c853011dc55c95832fe90de6bc56f5315f72", size = 656194, upload-time = "2026-05-18T04:31:47.14Z" }, + { url = "https://files.pythonhosted.org/packages/e7/54/a9c7ea9a82a4ac65e7004c0a03920b5cdd2f9c3b678757d9cd425aa51d53/watchfiles-1.2.0-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:b8c8358484d5fa12ef34f05b7f4168eaf1932f408725ff6d023c33ec17bd79d4", size = 400205, upload-time = "2026-05-18T04:32:05.153Z" }, + { url = "https://files.pythonhosted.org/packages/aa/5d/c9ab3534374a4a67450696905d6ef16a04405448b8dc52bd752ae50423d4/watchfiles-1.2.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:9f04b092229ad2c50126dd3c922c8822e51e605993764a33058d4a791ab42281", size = 392508, upload-time = "2026-05-18T04:30:54.849Z" }, + { url = "https://files.pythonhosted.org/packages/26/ca/1ad30103535cf0cecd7b993e8d50edc5351b1820e38f2d22e3df58962feb/watchfiles-1.2.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7a7ce236284f002a156f70add88efe5c70879cccbb658be0822c54b1306fc09d", size = 452448, upload-time = "2026-05-18T04:30:53.727Z" }, + { url = "https://files.pythonhosted.org/packages/37/a1/ceee2cdf2afbd715fa07758d39c9859513eae411b23196f7fd039e5feedd/watchfiles-1.2.0-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:b9909cc2b48468b575eefa944919e1fe8a36c5849d5c7c168f80a8c1db69398e", size = 459605, upload-time = "2026-05-18T04:30:23.312Z" }, + { url = "https://files.pythonhosted.org/packages/e8/f6/421e30fd1cb3907a84ed92ab3f1983e37ba2dca015e9a894a048418417a2/watchfiles-1.2.0-cp314-cp314-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0a37faaed405c67e28e6be45a1fa4f206ef5a2860f27c237db9fa30704c38242", size = 490757, upload-time = "2026-05-18T04:30:47.358Z" }, + { url = "https://files.pythonhosted.org/packages/41/b0/55ed1b97ed08be7bba6f9a541cac15f2a858e1d74d2b07b6da70a82aab00/watchfiles-1.2.0-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9649193aa27bd9ff2e80ff29bfaa93085496c7a3a377592823cc58b77ee88add", size = 568672, upload-time = "2026-05-18T04:30:38.915Z" }, + { url = "https://files.pythonhosted.org/packages/d1/cf/d8ae8a80dd7bafab395ea7681c10237311bbf34d37704a8c744e7cf31fc7/watchfiles-1.2.0-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:4e4ff8e37f99cf1da89e255e07c9c4b37c214038c4283707bdec308cb1b0ea1f", size = 464197, upload-time = "2026-05-18T04:30:09.914Z" }, + { url = "https://files.pythonhosted.org/packages/7c/8a/3076c496ca8dafe0e8cd03fcebdfc47be4b1174b4e5b24ff6e396e6b3af2/watchfiles-1.2.0-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:054dc20fd2e3132b4c3883b4a00d72fd6e1f56fdaf89fccd12e8057d74cd74d7", size = 453181, upload-time = "2026-05-18T04:30:14.829Z" }, + { url = "https://files.pythonhosted.org/packages/e5/10/9745e17c98e7b8a86454df0a3c7b5686bd650383f1e9f26e4ebcbd6cc0c0/watchfiles-1.2.0-cp314-cp314-manylinux_2_31_riscv64.whl", hash = "sha256:e140ed30ebde76796b686e67c182cff10ea2fbab186fafd1560f74bb5a473a6e", size = 465109, upload-time = "2026-05-18T04:30:28.123Z" }, + { url = "https://files.pythonhosted.org/packages/8f/95/8ef4a95481d3e0cb52d62a06fa6e972e81424be2d9698b91a2fecca9904c/watchfiles-1.2.0-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:bb7e52ecf68ba46d22df23467b87cffeb2146908aa523ebfe803019618cfda06", size = 630653, upload-time = "2026-05-18T04:31:49.304Z" }, + { url = "https://files.pythonhosted.org/packages/fd/e4/3b3bf36b0f829b50c6ebcb8d031583863c59f923d6a6af3d485e470d0fac/watchfiles-1.2.0-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:23282a321c8baf9b3a3c4afff673f9fe65eb7fdc2338d765ccad9d3d1916a5ba", size = 657838, upload-time = "2026-05-18T04:31:06.497Z" }, + { url = "https://files.pythonhosted.org/packages/21/b1/6cbbb50c1f3002ab568777d44aa21206dfb8807a840990c4037523b51812/watchfiles-1.2.0-cp314-cp314-win32.whl", hash = "sha256:c0db965c5f79aa49fe672d297cf1febc5ad149b658594944f49a54a2b96270a7", size = 275108, upload-time = "2026-05-18T04:30:06.891Z" }, + { url = "https://files.pythonhosted.org/packages/92/45/190ce6db8dcb4536682cf75d3889ff1a27182a58cb519d343cb6d9ea63d8/watchfiles-1.2.0-cp314-cp314-win_amd64.whl", hash = "sha256:71283b39fd17e5408eb123bd37aeecfd9d54c81fc184421943208aadb879d103", size = 288441, upload-time = "2026-05-18T04:32:12.901Z" }, + { url = "https://files.pythonhosted.org/packages/74/0d/3eae1c2313ab08378431d907c3f8095ecca00f3eda33111cf4f0f2591799/watchfiles-1.2.0-cp314-cp314-win_arm64.whl", hash = "sha256:c5c19526f4e54a00f2666a6c0e9e40d582c09e865055ea7378bf0009aab857b3", size = 280684, upload-time = "2026-05-18T04:31:26.902Z" }, + { url = "https://files.pythonhosted.org/packages/b1/75/fb64e6c25d6b5ca636d03df34ffb1c6e9873303e76d27967e045f8df088f/watchfiles-1.2.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:d73a585accffa5ae39c17264c36ec3166d2fad7000c780f5ef83b2722afb9dd2", size = 398857, upload-time = "2026-05-18T04:32:17.108Z" }, + { url = "https://files.pythonhosted.org/packages/73/4e/9f7adf01754cbf81843722ccfec169d8f26c69778281a302855cecd2ee08/watchfiles-1.2.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:ae99b14c5f21e026e0e9d96f40e07d8570ebee6cafd9d8fc318354606daa7a28", size = 392413, upload-time = "2026-05-18T04:31:07.911Z" }, + { url = "https://files.pythonhosted.org/packages/47/c8/bec626bcc2d69f44b9acb24ce7d60ed7b16b73628eea747fcbd169d8edda/watchfiles-1.2.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4429f3b105524a10b72c3a819b091c495d2811d419c1e1e8df773a5a5974f831", size = 452409, upload-time = "2026-05-18T04:31:20.142Z" }, + { url = "https://files.pythonhosted.org/packages/00/b7/b6362068e81e7c556d155a34c35d40ac3ef42d747b06d7f6e5bf58e359c2/watchfiles-1.2.0-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:43d818978d06062d9b22c4fab2ebe44cf5213d42dc8e62bda8c2760cfa2eeb33", size = 458827, upload-time = "2026-05-18T04:32:06.219Z" }, + { url = "https://files.pythonhosted.org/packages/67/f8/9a813fa42afb1e0b4625e75f0479826644d3ee8dc287e093799bc01f390c/watchfiles-1.2.0-cp314-cp314t-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b9f732dc58b2dbe69e464ccf8fff7a03b0dd0be439da4c0720d3558527d3d6b4", size = 490104, upload-time = "2026-05-18T04:31:56.034Z" }, + { url = "https://files.pythonhosted.org/packages/2f/bf/27dfb6094ca4c9aad21298b5525b6c53cb36121ee454331d05161e58d130/watchfiles-1.2.0-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8f200104103feb097de4cab8fe4f5dd18a2026934c7dea98c55a2f5fd6d5a33b", size = 571360, upload-time = "2026-05-18T04:31:57.133Z" }, + { url = "https://files.pythonhosted.org/packages/fb/39/44a096d67270ea93df91d33877dbe91fbda3aa4f8ec2edf799d93eda8736/watchfiles-1.2.0-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:63ac26eefbf4af1741247d6fb68b11c49a25b2f7413fbd318a83a12aaa9cf666", size = 464644, upload-time = "2026-05-18T04:30:57.33Z" }, + { url = "https://files.pythonhosted.org/packages/0e/80/c7472203bad6268e3ef1ad260739704847898938ad7ea8b63a5131f46b50/watchfiles-1.2.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0c4997d4e4a55f0d02b6cde327322daf3a0400e5df6c6b15948994bf72497925", size = 454771, upload-time = "2026-05-18T04:30:48.736Z" }, + { url = "https://files.pythonhosted.org/packages/51/cf/3b10b268b4b7f0fc26e9debb5eef1998b515887840f444cd3ec80c688755/watchfiles-1.2.0-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:4c887eba18b7945ac73067a8b4a66f21cd46c2539b2bc68588f7be6c7eb6d26b", size = 463494, upload-time = "2026-05-18T04:31:33.826Z" }, + { url = "https://files.pythonhosted.org/packages/3d/3e/a4302545cd589262a0dc7d140e86f7688eba3f9c72776c27f7e23b8864c4/watchfiles-1.2.0-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:3416ff151bb6b5a8d8d11664974fbef4d9305b9b2957839ab5a270468fd8df30", size = 629383, upload-time = "2026-05-18T04:31:15.596Z" }, + { url = "https://files.pythonhosted.org/packages/db/99/d5649df0a9a410d45b7c882304d0b790903ac9b6e8f2cfd12114e0c6b9f2/watchfiles-1.2.0-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:0e831a271c035d89789cffc386b6aa1375f39f1cd25eb7ca0997e4970d152fc5", size = 656093, upload-time = "2026-05-18T04:31:58.707Z" }, + { url = "https://files.pythonhosted.org/packages/23/f4/7513ef1e85fc4c6331b59479d6d72661fc391fbe543678052ac72c8b6c19/watchfiles-1.2.0-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:4674d49eb94706dfe666c069fc0a1b646ffcf920473492e209f6d5f60d3f0cc2", size = 403050, upload-time = "2026-05-18T04:30:36.753Z" }, + { url = "https://files.pythonhosted.org/packages/27/0b/a54103cfd732bb703c7a749222011a0483ef3705948dae3b203158601119/watchfiles-1.2.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:094b9b70103d4e963499bdea001ee3c2697b144cd9ae6218a62c0f89ec9e31db", size = 396629, upload-time = "2026-05-18T04:32:03.268Z" }, + { url = "https://files.pythonhosted.org/packages/5e/2c/73f31a3b893886206c3f54d73e8ad8dee58cdb2f69ad2622e0a8a9e07f4e/watchfiles-1.2.0-pp311-pypy311_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b0ef001f8c25ad0fa9529f914c1600647ecd0f542d11c19b7894768c67b6acb7", size = 457318, upload-time = "2026-05-18T04:31:01.932Z" }, + { url = "https://files.pythonhosted.org/packages/e9/f9/45d021e4a5cc7b9dd567f7cbb06d3b75f751a690063fb6cc7ec60f4e46b7/watchfiles-1.2.0-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a88fc94e647bc4eec523f1caa540258eb71d14278b9daf72fa1e2658a98df0f0", size = 457771, upload-time = "2026-05-18T04:30:56.331Z" }, +] + [[package]] name = "wcwidth" version = "0.7.0" From 5103bf03df16e78f87a9fc017edc1bd9128860ea Mon Sep 17 00:00:00 2001 From: sskmlm Date: Mon, 13 Jul 2026 18:38:38 -0500 Subject: [PATCH 77/78] split: browser chat frontend --- README.md | 3 +- code_puppy/api/templates/chat.html | 2439 +++++++++++++++++ .../tests/test_chat_template_session_state.py | 13 + .../tests/test_chat_template_xss_guards.py | 21 + 4 files changed, 2474 insertions(+), 2 deletions(-) create mode 100644 code_puppy/api/templates/chat.html create mode 100644 code_puppy/api/tests/test_chat_template_session_state.py create mode 100644 code_puppy/api/tests/test_chat_template_xss_guards.py diff --git a/README.md b/README.md index acc9cb7ae..8fe595e1b 100644 --- a/README.md +++ b/README.md @@ -152,8 +152,7 @@ You can toggle DBOS via either of these options: - CLI config (persists): `/set enable_dbos false` to disable (enabled by default) - -Config takes precedence if set; otherwise the environment variable is used. +DBOS durable execution is implemented by the `dbos_durable_exec` plugin and uses its own DBOS system database. It is separate from the Puppy Desk chat history database at `~/.puppy_desk/chat_messages.db`. ### Configuration diff --git a/code_puppy/api/templates/chat.html b/code_puppy/api/templates/chat.html new file mode 100644 index 000000000..c42d70df4 --- /dev/null +++ b/code_puppy/api/templates/chat.html @@ -0,0 +1,2439 @@ + + + + + + 🐶 Code Puppy Chat + + + + + + + + + + + + + + + + +
+ +
+ + +
+ + + + Connecting… + + +
+ Agent + +
+ + +
+ Model + +
+ + +
+
+ Disconnected +
+ + + + Home +
+ + +
+ + +
+ + Not set +
+ + +
+ + + + + + + + + + +
+ + + +
+ +
+ + +
+ +
+ + +
+
+ + +
+
+
+
+ + + + diff --git a/code_puppy/api/tests/test_chat_template_session_state.py b/code_puppy/api/tests/test_chat_template_session_state.py new file mode 100644 index 000000000..ac77ee149 --- /dev/null +++ b/code_puppy/api/tests/test_chat_template_session_state.py @@ -0,0 +1,13 @@ +from pathlib import Path + + +def test_chat_template_handles_session_meta_and_config_frames(): + html = Path("code_puppy/api/templates/chat.html").read_text() + + assert "case 'session_meta':" in html + assert "case 'session_meta_updated':" in html + assert "case 'session_restored':" in html + assert "case 'session_switched':" in html + assert "case 'config_value':" in html + assert "function handleSessionMeta(message)" in html + assert "window.codePuppyConfig = sessionState.config;" in html diff --git a/code_puppy/api/tests/test_chat_template_xss_guards.py b/code_puppy/api/tests/test_chat_template_xss_guards.py new file mode 100644 index 000000000..d35afd77c --- /dev/null +++ b/code_puppy/api/tests/test_chat_template_xss_guards.py @@ -0,0 +1,21 @@ +from pathlib import Path + + +def test_chat_template_includes_markdown_sanitization_guards(): + html = Path("code_puppy/api/templates/chat.html").read_text() + + assert "dompurify" in html.lower() + assert "DOMPurify.sanitize" in html + assert "escapeHtml(String(title))" in html + assert "escapeHtml(String(description))" in html + # Tool card text is assigned through textContent, not innerHTML. + assert "const toolName = String(" in html + assert "message.tool_name || message.tool" in html + assert "nm.textContent = toolName" in html + assert "escapeHtml(String(message.agent_name" in html + + +def test_chat_template_uses_wss_fallback_aware_url_template(): + html = Path("code_puppy/api/templates/chat.html").read_text() + # Ensure we still have explicit websocket URL handling in the template. + assert "wsUrl" in html From fcf4aed8eff5cdeb54f74cc35d84b30e5ba5c8e9 Mon Sep 17 00:00:00 2001 From: sskmlm Date: Mon, 13 Jul 2026 20:02:08 -0500 Subject: [PATCH 78/78] fix(frontend): harden permission request rendering --- code_puppy/api/templates/chat.html | 23 ++++++++++++++----- .../tests/test_chat_template_xss_guards.py | 9 ++++++++ 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/code_puppy/api/templates/chat.html b/code_puppy/api/templates/chat.html index c42d70df4..b0e1088be 100644 --- a/code_puppy/api/templates/chat.html +++ b/code_puppy/api/templates/chat.html @@ -1218,7 +1218,7 @@ let wsUrl = CONFIG.wsUrl; if (CONFIG.sessionId) { - wsUrl += `?session_id=${CONFIG.sessionId}`; + wsUrl += `?session_id=${encodeURIComponent(CONFIG.sessionId)}`; } socket = new WebSocket(wsUrl); @@ -1744,15 +1744,26 @@
${escapeHtml(String(description))}
${escapeHtml(command)}
${escapeHtml(cwd)}
-
- - -
+
Timeout in ${timeout}s
`; - + + const actions = bubble.querySelector('.perm-card-actions'); + const approveBtn = document.createElement('button'); + approveBtn.className = 'perm-card-btn perm-btn-ok'; + approveBtn.textContent = 'Approve'; + approveBtn.addEventListener('click', () => approvePermission(requestId)); + + const denyBtn = document.createElement('button'); + denyBtn.className = 'perm-card-btn perm-btn-no'; + denyBtn.textContent = 'Deny'; + denyBtn.addEventListener('click', () => denyPermission(requestId)); + + actions.appendChild(approveBtn); + actions.appendChild(denyBtn); + messagesContainer.appendChild(bubble); messagesContainer.scrollTop = messagesContainer.scrollHeight; } diff --git a/code_puppy/api/tests/test_chat_template_xss_guards.py b/code_puppy/api/tests/test_chat_template_xss_guards.py index d35afd77c..01e4851fb 100644 --- a/code_puppy/api/tests/test_chat_template_xss_guards.py +++ b/code_puppy/api/tests/test_chat_template_xss_guards.py @@ -13,9 +13,18 @@ def test_chat_template_includes_markdown_sanitization_guards(): assert "message.tool_name || message.tool" in html assert "nm.textContent = toolName" in html assert "escapeHtml(String(message.agent_name" in html + assert "onclick=\"approvePermission" not in html + assert "onclick=\"denyPermission" not in html + + +def test_permission_request_buttons_use_dom_event_handlers(): + html = Path("code_puppy/api/templates/chat.html").read_text() + assert "approveBtn.addEventListener('click', () => approvePermission(requestId));" in html + assert "denyBtn.addEventListener('click', () => denyPermission(requestId));" in html def test_chat_template_uses_wss_fallback_aware_url_template(): html = Path("code_puppy/api/templates/chat.html").read_text() # Ensure we still have explicit websocket URL handling in the template. assert "wsUrl" in html + assert "encodeURIComponent(CONFIG.sessionId)" in html