diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index efcc3a998..3456a92ed 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -54,7 +54,10 @@ jobs: } } JSON - echo "Wrote $HOME/.code_puppy/extra_models.json" + # Tests isolate user config, and several fallback contracts require a + # non-empty bundled catalog. + cp "$HOME/.code_puppy/extra_models.json" code_puppy/models.json + echo "Wrote CI model catalog and test-only bundled fallback" - name: Debug environment variables env: diff --git a/.github/workflows/ocr-review.yml b/.github/workflows/ocr-review.yml new file mode 100644 index 000000000..8623822f2 --- /dev/null +++ b/.github/workflows/ocr-review.yml @@ -0,0 +1,77 @@ +name: OpenCodeReview + +on: + issue_comment: + types: [created] + +# Only an authorized /open-code-review comment shares the PR's concurrency +# group. Unrelated comments must not cancel a review already in progress. +concurrency: + group: >- + ${{ + ( + github.event.issue.pull_request + && github.event.comment.user.type != 'Bot' + && ( + github.event.comment.author_association == 'MEMBER' + || github.event.comment.author_association == 'OWNER' + || github.event.comment.author_association == 'COLLABORATOR' + ) + && startsWith(github.event.comment.body, '/open-code-review') + ) + && format('ocr-{0}', github.event.issue.number) + || format('ocr-noop-{0}', github.run_id) + }} + cancel-in-progress: true + +permissions: + contents: read + pull-requests: write + +jobs: + review: + name: Review pull request + if: | + github.event.issue.pull_request + && github.event.comment.user.type != 'Bot' + && ( + github.event.comment.author_association == 'MEMBER' + || github.event.comment.author_association == 'OWNER' + || github.event.comment.author_association == 'COLLABORATOR' + ) + && startsWith(github.event.comment.body, '/open-code-review') + runs-on: ubuntu-latest + timeout-minutes: 30 + + steps: + - name: Resolve pull request + id: pr + uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7 + with: + script: | + const { data: pullRequest } = await github.rest.pulls.get({ + owner: context.repo.owner, + repo: context.repo.repo, + pull_number: context.issue.number, + }); + core.setOutput('base_ref', pullRequest.base.ref); + core.setOutput('head_sha', pullRequest.head.sha); + core.setOutput('title', pullRequest.title); + + - name: Run OpenCodeReview + uses: alibaba/open-code-review@c3da87885086ff41ed8f2f77b65c7a3c471e01aa + with: + llm_url: ${{ secrets.OCR_LLM_URL }} + llm_auth_token: ${{ secrets.OCR_LLM_AUTH_TOKEN }} + llm_model: ${{ vars.OCR_LLM_MODEL }} + llm_use_anthropic: ${{ vars.OCR_LLM_USE_ANTHROPIC }} + llm_extra_body: '{"chat_template_kwargs":{"thinking_mode":"disabled"}}' + llm_timeout: "900" + base_ref: ${{ steps.pr.outputs.base_ref }} + head_sha: ${{ steps.pr.outputs.head_sha }} + background: ${{ steps.pr.outputs.title }} + ocr_version: "1.7.7" + review_concurrency: "1" + sticky_summary: "true" + incremental: "true" + upload_artifacts: "true" diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 07b100d2a..129a95bd8 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -62,7 +62,11 @@ jobs: } } JSON - echo "Wrote $HOME/.code_puppy/extra_models.json" + # Tests isolate user config, and several fallback contracts require a + # non-empty bundled catalog. Keep this test-only copy out of releases: + # the publish job starts from a fresh checkout. + cp "$HOME/.code_puppy/extra_models.json" code_puppy/models.json + echo "Wrote CI model catalog and test-only bundled fallback" - name: Debug environment variables env: diff --git a/.gitignore b/.gitignore index 56e9f3c35..e1a999037 100644 --- a/.gitignore +++ b/.gitignore @@ -45,3 +45,6 @@ code_puppy/bundled_skills/ .claude/hooks/ts-hooks/dist/ .json + +.state/ +.maestro/* 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/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/agents/_builder.py b/code_puppy/agents/_builder.py index 6e11a2445..a1958bf2b 100644 --- a/code_puppy/agents/_builder.py +++ b/code_puppy/agents/_builder.py @@ -370,12 +370,16 @@ def load_model_with_fallback( continue try: model = ModelFactory.get_model(candidate, models_config) - emit_info( - f"Using fallback model: {candidate}", message_group=message_group - ) - return model, candidate except ValueError: continue + if model is None: + # Missing credentials/provider reachability can make a model + # "configured" but unavailable at runtime. Keep searching for + # a *real* fallback instead of returning a None model that only + # explodes later in pydantic-ai run(). + continue + emit_info(f"Using fallback model: {candidate}", message_group=message_group) + return model, candidate friendly = ( "No valid model could be loaded. Update the model configuration or " 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/_key_listeners.py b/code_puppy/agents/_key_listeners.py index 7db1cf33b..1598bf5bc 100644 --- a/code_puppy/agents/_key_listeners.py +++ b/code_puppy/agents/_key_listeners.py @@ -885,6 +885,15 @@ def _enter_cbreak() -> None: # every control char is delivered verbatim, exactly like the # raw mode (tty.setraw) the classic prompt_toolkit path used. # + # setcbreak ALSO leaves software flow control enabled. With + # IXON set, an accidental Ctrl+S tells the tty to STOP output; + # the next editor repaint then blocks forever in stdout.flush() + # ON THE KEY-LISTENER THREAD. Input is alive but can't dispatch + # another key (including Ctrl+Q) because it is trapped writing + # the prompt — indistinguishable from a bricked terminal. This + # exact stack was captured live on 2026-07-11. Clear IXON/IXOFF + # while we own stdin; original attrs restore on suspend/exit. + # # Pure-keybinding Ctrl+C: disable the tty's INTR character so # ^C is delivered as a raw \x03 byte instead of becoming a # SIGINT — the POSIX mirror of Windows' session-wide @@ -896,8 +905,11 @@ def _enter_cbreak() -> None: # semantics return whenever we release stdin. try: attrs = termios.tcgetattr(fd) - attrs[0] &= ~termios.ICRNL # iflag - attrs[3] &= ~termios.IEXTEN # lflag + attrs[0] &= ~termios.ICRNL # preserve Ctrl+J vs Enter + attrs[0] &= ~termios.IXON # Ctrl+S must reach the editor + if hasattr(termios, "IXOFF"): + attrs[0] &= ~termios.IXOFF + attrs[3] &= ~termios.IEXTEN # deliver Ctrl+V/Ctrl+O raw try: vdisable = os.fpathconf(fd, "PC_VDISABLE") except (OSError, ValueError, AttributeError): diff --git a/code_puppy/agents/_run_signals.py b/code_puppy/agents/_run_signals.py index 6c4636ccf..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 @@ -127,10 +125,15 @@ 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_steers = pc.drain_pending_steer() - if stale_steers: + stale_compactions = pc.drain_compaction_requests() + if stale_compactions: emit_warning( - f"Discarded {len(stale_steers)} stale steering message(s) from a previous run." + f"Discarded {stale_compactions} stale compaction request(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/_runtime.py b/code_puppy/agents/_runtime.py index 99b21782e..770fd661c 100644 --- a/code_puppy/agents/_runtime.py +++ b/code_puppy/agents/_runtime.py @@ -18,6 +18,7 @@ from __future__ import annotations import asyncio +import re import signal import threading import uuid @@ -113,6 +114,8 @@ "rate limited", "overloaded", "service unavailable", + "bad gateway", + "gateway timeout", "server had an error processing your request", "retry your request", "internal server error", @@ -156,6 +159,34 @@ def _matches_retryable_snippet(msg: str) -> bool: return "stream" in msg and "ended" in msg +# Matches a ``[HTTP 502]`` style marker some gateways embed in the message. +_EMBEDDED_HTTP_STATUS_RE = re.compile(r"\[HTTP\s+(\d{3})\]", re.IGNORECASE) + + +def _is_transient_status(status_code: object) -> bool: + """True for HTTP statuses worth a silent retry: 429 or any 5xx.""" + return status_code == 429 or (isinstance(status_code, int) and status_code >= 500) + + +def _embedded_http_status(text: str) -> Optional[int]: + """Dig a real upstream HTTP status out of a ``[HTTP NNN]`` message marker. + + A gateway (e.g. the puppy-backend ``custom_anthropic`` proxy) can deliver an + upstream failure as an *in-band SSE ``error`` event* over a connection that + itself returned HTTP 200. The Anthropic SDK then builds an ``APIStatusError`` + whose ``.status_code`` is 200 -- the stream's status, not the failure's -- + and whose ``body.error.type`` is a generic ``"internal_error"``. The only + faithful trace of the real failure is a ``[HTTP 502]`` prefix baked into the + message/body text. Without recovering it, the classifier sees a 200 with no + matching snippet and refuses to retry, so a transient gateway 5xx crashes + the REPL instead of getting the slow spaced-out retry. + + Returns the embedded status as an int, or ``None`` if no marker is present. + """ + match = _EMBEDDED_HTTP_STATUS_RE.search(text or "") + return int(match.group(1)) if match else None + + def _walk_cause_chain( exc: BaseException, max_depth: int = 5 ) -> "Iterator[BaseException]": @@ -190,10 +221,26 @@ def _is_retryable_one(exc: BaseException) -> bool: msg = str(exc) + # Provider-agnostic: an in-band SSE 5xx surfaces with a 200 status_code and + # only a "[HTTP 5xx]" marker in the text (see _embedded_http_status). A 4xx + # marker (e.g. [HTTP 400]) is a genuine client error and must NOT retry. + if _is_transient_status(_embedded_http_status(msg)): + return True + if isinstance(exc, UnexpectedModelBehavior): 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 + # spaced-out 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) @@ -227,10 +274,18 @@ def _is_retryable_one(exc: BaseException) -> bool: if isinstance(err, dict): if _matches_retryable_snippet(str(err.get("message", ""))): return True + # A gateway 5xx wrapped in an SSE error event over a 200 stream + # lands here with status_code=200 and type "internal_error"; its + # real status survives only as a "[HTTP 5xx]" marker in the msg. + if _is_transient_status( + _embedded_http_status(str(err.get("message", ""))) + ): + return True if str(err.get("type", "")).lower() in { "api_error", "server_error", "internal_server_error", + "internal_error", }: return True @@ -255,24 +310,71 @@ 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 spaced-out 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( max_attempts: int = 3, - delays: Sequence[float] = (1, 2, 4), + delays: Sequence[float] = (5, 30), ) -> Callable[[Callable[[], Any]], Callable[[], Any]]: """Wrap a no-arg async callable with streaming-retry semantics. + Budget: 3 attempts spaced ``(5, 30)`` seconds apart. With N attempts only + the first N-1 delays fire, so this waits ~5s before retry #2 and ~30s before + retry #3 -- a ~35s total window. That spacing is deliberate: gateway 5xx + outages (the Google "502 ... try again in 30 seconds" page the puppy-backend + proxy surfaces) routinely outlast a tight 1-2-4s window, so a fast burst of + retries just exhausts the budget before the upstream recovers. The quick + first retry still catches instantaneous SSE blips; the long second gap rides + out the advertised ~30s outage before giving up. + Every retry (and the final exhaustion, if it happens) is logged with full detail to the on-disk error log via ``log_error``. Users only see a short UI banner, but SRE / power-users grepping ``~/.code_puppy/logs/errors.log`` diff --git a/code_puppy/agents/base_agent.py b/code_puppy/agents/base_agent.py index ff433ccee..4576819d1 100644 --- a/code_puppy/agents/base_agent.py +++ b/code_puppy/agents/base_agent.py @@ -73,6 +73,7 @@ def __init__(self) -> None: self._code_generation_agent: Any = None self._last_model_name: Optional[str] = None self._runtime_model_name_override: Optional[str] = None + self._session_model_name: Optional[str] = None self._puppy_rules: Optional[str] = None self._mcp_servers: List[Any] = [] self.cur_model: Optional[pydantic_ai.models.Model] = None @@ -143,6 +144,8 @@ def get_model_name(self) -> Optional[str]: override = self.get_runtime_model_name_override() if override: return override + if self._session_model_name: + return self._session_model_name pinned = get_agent_pinned_model(self.name) return pinned if pinned else get_global_model_name() @@ -191,6 +194,27 @@ def clear_message_history(self) -> None: def append_to_message_history(self, message: Any) -> None: self._message_history.append(message) + # ---- Session model + compaction compatibility helpers ---------------- + def set_session_model(self, model_name: Optional[str]) -> None: + """Set a per-session model override for this agent instance.""" + self._session_model_name = model_name or None + + def get_session_model(self) -> Optional[str]: + """Return the per-session model override, if any.""" + return self._session_model_name + + def reset_session_model(self) -> None: + """Clear the per-session model override for this agent instance.""" + self._session_model_name = None + + def get_compacted_message_hashes(self) -> Set[int]: + """Expose compacted-message hashes for session state transfers.""" + return set(self._compacted_message_hashes) + + def add_compacted_message_hash(self, message_hash: int) -> None: + """Track a compacted-message hash for this agent instance.""" + self._compacted_message_hashes.add(message_hash) + # ---- Token / context helpers ------------------------------------------ def estimate_tokens_for_message(self, message: Any) -> int: return estimate_tokens_for_message(message, self.get_model_name()) diff --git a/code_puppy/agents/event_stream_handler.py b/code_puppy/agents/event_stream_handler.py index f1fd546c5..54524633b 100644 --- a/code_puppy/agents/event_stream_handler.py +++ b/code_puppy/agents/event_stream_handler.py @@ -151,7 +151,14 @@ 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_prompt_text_color, + on_termflow_highlighter, + on_termflow_style, + ) # Use the module-level console (set via set_streaming_console) console = get_streaming_console() @@ -175,6 +182,21 @@ async def event_stream_handler( # Optional smooth (typewriter) writers wrapping the console for text parts. termflow_writers: dict[int, SmoothTermflowWriter] = {} + class _ThemedBoldWriter: + """Keep terminal bold from brightening default text to profile white.""" + + def __init__(self, target, color: str) -> None: + self._target = target + self._color_sgr = f"\x1b[38;2;{int(color[1:3], 16)};{int(color[3:5], 16)};{int(color[5:7], 16)}m" + + def write(self, text): + return self._target.write( + text.replace("\x1b[1m", f"\x1b[1m{self._color_sgr}") + ) + + def flush(self): + return self._target.flush() + def _make_text_renderer(index: int) -> TermflowRenderer: """Build a termflow renderer, optionally typed out smoothly.""" writer = make_smooth_termflow_writer(console.file) @@ -184,10 +206,15 @@ def _make_text_renderer(index: int) -> TermflowRenderer: output = writer else: output = console.file + prompt_color = on_prompt_text_color() + if prompt_color and len(prompt_color) == 7 and prompt_color.startswith("#"): + output = _ThemedBoldWriter(output, prompt_color) 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 @@ -233,12 +260,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="", ) @@ -435,15 +462,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/api/__init__.py b/code_puppy/api/__init__.py new file mode 100644 index 000000000..9e503eeb1 --- /dev/null +++ b/code_puppy/api/__init__.py @@ -0,0 +1,13 @@ +"""Code Puppy REST API module. + +This module provides a FastAPI-based REST API for Code Puppy configuration, +sessions, commands, and real-time WebSocket communication. + +Exports: + create_app: Factory function to create the FastAPI application + main: Entry point to run the server +""" + +from code_puppy.api.app import create_app + +__all__ = ["create_app"] diff --git a/code_puppy/api/app.py b/code_puppy/api/app.py new file mode 100644 index 000000000..ed59679ed --- /dev/null +++ b/code_puppy/api/app.py @@ -0,0 +1,249 @@ +"""FastAPI application factory for Code Puppy API.""" + +import asyncio +import logging +from contextlib import asynccontextmanager +from pathlib import Path +from typing import AsyncGenerator + +from fastapi import FastAPI, Request +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import FileResponse, HTMLResponse, JSONResponse +from starlette.middleware.base import BaseHTTPMiddleware + + +logger = logging.getLogger(__name__) + +# Default request timeout (seconds) - fail fast! +REQUEST_TIMEOUT = 30.0 + + +class TimeoutMiddleware(BaseHTTPMiddleware): + """Middleware to enforce request timeouts and prevent hanging requests.""" + + def __init__(self, app, timeout: float = REQUEST_TIMEOUT): + super().__init__(app) + self.timeout = timeout + + async def dispatch(self, request: Request, call_next): + # Skip timeout for WebSocket upgrades and streaming endpoints + if request.headers.get( + "upgrade", "" + ).lower() == "websocket" or request.url.path.startswith("/ws/"): + return await call_next(request) + + try: + return await asyncio.wait_for( + call_next(request), + timeout=self.timeout, + ) + except asyncio.TimeoutError: + return JSONResponse( + status_code=504, + content={ + "detail": f"Request timed out after {self.timeout}s", + "error": "timeout", + }, + ) + + +@asynccontextmanager +async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: + """Lifespan context manager for startup and shutdown events. + + Handles graceful cleanup of resources when the server shuts down. + """ + # Startup + logger.info("🐶 Code Puppy API starting up...") + + # Load plugin callbacks (including frontend_emitter for real-time streaming) + from code_puppy import plugins + + result = plugins.load_plugin_callbacks() + logger.info( + f"✓ Loaded plugins: builtin={result.get('builtin', [])}, user={result.get('user', [])}, external={result.get('external', [])}" + ) + + # Initialise shared SQLite database + try: + from code_puppy.api.db.connection import init_db + + await init_db() + logger.info("✓ aiosqlite DB initialised") + except Exception as _db_exc: + logger.error("SQLite DB init failed (continuing without DB): %s", _db_exc) + + yield + # Shutdown: clean up all the things! + logger.info("🐶 Code Puppy API shutting down, cleaning up...") + + # 1. Shutdown session cache thread pool executor + try: + from code_puppy.api.session_cache import shutdown_executor + + await shutdown_executor() + logger.info("✓ Session cache executor shut down") + except Exception as e: + logger.error("Error shutting down session cache executor: %s", e) + + # 2. Shutdown ws_sessions thread pool executor + try: + from code_puppy.api.routers import ws_sessions + + ws_sessions._executor.shutdown(wait=False) + logger.info("✓ WS sessions executor shut down") + except Exception as e: + logger.error("Error shutting down ws_sessions executor: %s", e) + + # 4. Remove PID file so /api status knows we're gone + try: + from code_puppy.config import STATE_DIR + + pid_file = Path(STATE_DIR) / "api_server.pid" + if pid_file.exists(): + pid_file.unlink() + logger.info("✓ PID file removed") + except Exception as e: + logger.error("Error removing PID file: %s", e) + + # 6. Close SQLite database + try: + from code_puppy.api.db.connection import close_db + + await close_db() + logger.info("✓ aiosqlite DB closed") + except Exception as e: + logger.error("Error closing SQLite DB: %s", e) + + +def create_app() -> FastAPI: + """Create and configure the FastAPI application.""" + app = FastAPI( + lifespan=lifespan, + title="Code Puppy API", + description="REST API and Interactive Terminal for Code Puppy", + version="1.0.0", + docs_url="/docs", + redoc_url="/redoc", + ) + + # Timeout middleware - added first so it wraps everything + app.add_middleware(TimeoutMiddleware, timeout=REQUEST_TIMEOUT) + + # CORS middleware for frontend access + app.add_middleware( + CORSMiddleware, + allow_origins=["*"], # Local/trusted + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], + ) + + # Include routers + from code_puppy.api.routers import ( + agents, + commands, + config, + sessions, + ws_sessions, + ) + + app.include_router(config.router, prefix="/api/config", tags=["config"]) + app.include_router(commands.router, prefix="/api/commands", tags=["commands"]) + app.include_router(sessions.router, prefix="/api/sessions", tags=["sessions"]) + app.include_router( + ws_sessions.router, prefix="/api/ws-sessions", tags=["ws-sessions"] + ) + app.include_router(agents.router, prefix="/api/agents", tags=["agents"]) + + from code_puppy.api.routers import models + + app.include_router(models.router, prefix="/api/models", tags=["models"]) + + from code_puppy.api.routers import protocol + + app.include_router(protocol.router, prefix="/api/protocol", tags=["protocol"]) + + # WebSocket endpoints + from code_puppy.api.websocket import setup_websocket + + setup_websocket(app) + + # Templates directory + templates_dir = Path(__file__).parent / "templates" + + @app.get("/") + async def root(): + """Landing page with links to terminal and docs.""" + return HTMLResponse( + content=""" + + + + Code Puppy 🐶 + + + +
+

🐶

+

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/compaction_tracker.py b/code_puppy/api/compaction_tracker.py new file mode 100644 index 000000000..8145b39c9 --- /dev/null +++ b/code_puppy/api/compaction_tracker.py @@ -0,0 +1,130 @@ +"""Compaction tracking — writes compaction metadata to SQLite.""" + +import logging +from datetime import datetime, timezone +from typing import Any, List + +logger = logging.getLogger(__name__) + + +def track_compaction_event( + session_id: str, + messages_before: List[Any], + messages_after: List[Any], + strategy: str, + agent: Any | None = None, +) -> None: + """ + Track a compaction event in SQLite. + + Called directly from message_history_processor after compaction completes. + + Args: + session_id: Session identifier + messages_before: Message history before compaction + messages_after: Message history after compaction + strategy: 'summarization' or 'truncation' + agent: Agent instance (optional, used for token estimation) + """ + # Only proceed if compaction actually happened (message count decreased) + if len(messages_after) >= len(messages_before): + logger.debug( + "Compaction tracker: no size reduction (%d -> %d), skipping", + len(messages_before), + len(messages_after), + ) + return + + try: + from code_puppy.api.db.connection import get_db, get_write_lock + + # Estimate tokens (use agent's method if available, else default) + def fallback_estimate(m: Any) -> int: + return 100 + + if agent and hasattr(agent, "estimate_tokens_for_message"): + estimate_tokens = agent.estimate_tokens_for_message + else: + estimate_tokens = fallback_estimate + + # Calculate compaction metadata + compacted_count = len(messages_before) - len(messages_after) + + # Estimate tokens for source messages (messages that were compacted) + # Assume first message is system prompt, compaction starts at index 1 + source_messages = messages_before[1 : compacted_count + 1] + source_tokens = sum(estimate_tokens(m) for m in source_messages) + + # Extract summary text from new messages (compare object IDs) + before_ids = {id(m) for m in messages_before} + new_messages = [m for m in messages_after if id(m) not in before_ids] + + summary_text = "(no summary found)" + summary_tokens = 0 + if new_messages: + summary_msg = new_messages[0] + text_parts = [] + for part in getattr(summary_msg, "parts", []) or []: + if hasattr(part, "content"): + text_parts.append(str(part.content)) + summary_text = " ".join(text_parts) if text_parts else "(empty summary)" + summary_tokens = estimate_tokens(summary_msg) + + # Assume compaction removes from the start (after system message at seq 1) + source_start = 2 # First message after system prompt + source_end = source_start + compacted_count - 1 + + # Create compaction_log row + created_at = datetime.now(timezone.utc).isoformat() + + with get_write_lock(): + db = get_db() + cursor = db.execute( + """ + INSERT INTO compaction_log + (session_id, summary_text, source_start, source_end, + source_count, source_tokens, summary_tokens, strategy, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + session_id, + summary_text[:1000], # Truncate very long summaries + source_start, + source_end, + compacted_count, + source_tokens, + summary_tokens, + strategy, + created_at, + ), + ) + compaction_log_id = cursor.lastrowid + + # Mark compacted messages in the messages table + rows_updated = db.execute( + """ + UPDATE messages + SET compacted = 1, compaction_log_id = ? + WHERE session_id = ? AND seq >= ? AND seq <= ? AND compacted = 0 + """, + (compaction_log_id, session_id, source_start, source_end), + ).rowcount + + db.commit() + + logger.info( + "Compaction tracked: session=%s log_id=%d seq=%d-%d count=%d marked=%d strategy=%s", + session_id, + compaction_log_id, + source_start, + source_end, + compacted_count, + rows_updated, + strategy, + ) + + except Exception as e: + # Non-fatal — compaction still worked, we just didn't track it + logger.warning( + "Compaction tracker failed for %s: %s", session_id, e, exc_info=True + ) diff --git a/code_puppy/api/db/__init__.py b/code_puppy/api/db/__init__.py new file mode 100644 index 000000000..83cffcb97 --- /dev/null +++ b/code_puppy/api/db/__init__.py @@ -0,0 +1,47 @@ +"""SQLite persistence layer for the Code Puppy WS API. + +Stores session history in ~/.puppy_desk/chat_messages.db — the same +database that the Desk Puppy Electron FE reads from. + +All database operations are async coroutines backed by aiosqlite. +""" + +from code_puppy.api.db.connection import close_db, get_db, init_db +from code_puppy.api.db.queries import ( + get_active_messages, + get_next_seq, + get_session_metadata, + get_session_row, + insert_compaction_log, + insert_message, + insert_tool_call, + mark_messages_compacted, + session_exists, + soft_delete_session, + update_session_stats, + update_session_working_directory, + upsert_session, + write_system_message_to_sqlite, + write_turn_to_sqlite, +) + +__all__ = [ + "init_db", + "close_db", + "get_db", + "get_session_row", + "get_session_metadata", + "session_exists", + "upsert_session", + "soft_delete_session", + "update_session_stats", + "update_session_working_directory", + "write_system_message_to_sqlite", + "get_next_seq", + "insert_message", + "insert_tool_call", + "get_active_messages", + "mark_messages_compacted", + "insert_compaction_log", + "write_turn_to_sqlite", +] diff --git a/code_puppy/api/db/backfill.py b/code_puppy/api/db/backfill.py new file mode 100644 index 000000000..142c0f471 --- /dev/null +++ b/code_puppy/api/db/backfill.py @@ -0,0 +1,381 @@ +"""Backfill helpers for the messages table. + +Currently contains one migration helper: + + backfill_pydantic_json(db) + Reconstructs pydantic_json from the plain `content` column for any + message row where pydantic_json IS NULL. Used by the v4→v5 schema + migration in connection.py. + +Design notes +------------ +* Only 'user' and 'assistant' roles are processed — 'system' rows have no + pydantic_json and the load path already skips them. +* The reconstruction is lossy for assistant messages that originally contained + tool calls or extended thinking, because that structure was never stored in + the `content` column. The backfill preserves the visible text so sessions + resume with *something* rather than blank history. +* The function is idempotent: it filters on ``pydantic_json IS NULL`` so + running it twice is safe. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import aiosqlite + +logger = logging.getLogger(__name__) + + +def _build_message(role: str, content: str, model_name: str) -> Any: + """Build a pydantic-ai ModelMessage from a plain content string. + + Args: + role: 'user' or 'assistant'. + content: Plain text content extracted from the messages row. + model_name: Model name stored on the row (used for ModelResponse). + + Returns: + A ModelRequest or ModelResponse instance ready for serialisation. + + Raises: + ValueError: If *role* is not 'user' or 'assistant'. + """ + from pydantic_ai.messages import ( + ModelRequest, + ModelResponse, + TextPart, + UserPromptPart, + ) + + if role == "user": + return ModelRequest(parts=[UserPromptPart(content=content)]) + if role == "assistant": + return ModelResponse( + parts=[TextPart(content=content)], + model_name=model_name or "unknown", + ) + raise ValueError(f"Unsupported role for backfill: {role!r}") + + +def _serialize(msg: Any) -> str: + """Serialise a single ModelMessage to the pydantic_json wire format. + + Stores as a single-element JSON array — identical to what + pydantic_json_for_message() produces during normal chat saves. + """ + from pydantic_ai.messages import ModelMessagesTypeAdapter + + return ModelMessagesTypeAdapter.dump_json([msg]).decode("utf-8") + + +async def backfill_pydantic_json( + db: "aiosqlite.Connection", + *, + batch_size: int = 200, +) -> tuple[int, int]: + """Backfill pydantic_json for message rows where it is currently NULL. + + Reads all eligible rows in one query, groups UPDATE statements into + batches of *batch_size* to keep transactions small, and commits after + each batch so a mid-run crash leaves the DB consistent. + + Args: + db: Open aiosqlite connection (must have row_factory set). + batch_size: Number of rows updated per transaction. + + Returns: + ``(updated, skipped)`` — row counts for rows successfully backfilled + and rows skipped due to serialisation errors. + """ + cursor = await db.execute( + """ + SELECT id, session_id, seq, role, content, model_name + FROM messages + WHERE pydantic_json IS NULL + AND compacted = 0 + AND role IN ('user', 'assistant') + ORDER BY session_id, seq + """ + ) + rows = await cursor.fetchall() + + if not rows: + logger.info("backfill_pydantic_json: no NULL rows found — nothing to do") + return 0, 0 + + logger.info("backfill_pydantic_json: %d rows to backfill", len(rows)) + + updated = 0 + skipped = 0 + batch: list[tuple[str, int]] = [] # (pydantic_json, id) + + for row in rows: + row_id: int = row["id"] + role: str = row["role"] + content: str = row["content"] or "" + model_name: str = row["model_name"] or "unknown" + session_id: str = row["session_id"] + seq: int = row["seq"] + + try: + msg = _build_message(role, content, model_name) + pj = _serialize(msg) + batch.append((pj, row_id)) + except Exception as exc: + logger.warning( + "backfill_pydantic_json: skipping id=%s (session=%s seq=%s role=%s): %s", + row_id, + session_id, + seq, + role, + exc, + ) + skipped += 1 + + if len(batch) >= batch_size: + await _flush_batch(db, batch) + updated += len(batch) + batch = [] + + if batch: + await _flush_batch(db, batch) + updated += len(batch) + + logger.info( + "backfill_pydantic_json: done — %d updated, %d skipped", + updated, + skipped, + ) + return updated, skipped + + +async def _flush_batch( + db: "aiosqlite.Connection", + batch: list[tuple[str, int]], +) -> None: + """Write one batch of (pydantic_json, id) pairs and commit.""" + await db.executemany( + "UPDATE messages SET pydantic_json = ? WHERE id = ?", + batch, + ) + await db.commit() + + +# --------------------------------------------------------------------------- +# Tool-call backfill (added to fix the early-continue bug in +# write_turn_to_sqlite that prevented tool_calls rows from ever being written) +# --------------------------------------------------------------------------- + + +async def backfill_tool_calls( + db: "aiosqlite.Connection", + *, + session_id: str | None = None, + batch_size: int = 100, +) -> tuple[int, int]: + """Backfill tool_call rows from pydantic_json for sessions saved without them. + + Sessions written before the fix for the early-continue bug in + ``write_turn_to_sqlite`` have no rows in ``tool_calls`` even though the + pydantic_json column contains full ToolCallPart / ToolReturnPart data. + + This function replays each session's message history in seq order, + matches ToolCallPart entries in assistant messages with ToolReturnPart + entries in the following user messages, and inserts the reconstructed + rows into ``tool_calls``. Uses ``INSERT OR IGNORE`` — fully idempotent. + + Args: + db: Open aiosqlite connection with ``row_factory`` set to + ``aiosqlite.Row``. + session_id: If provided, only process that one session; otherwise + process every session that has pydantic_json rows. + batch_size: Number of ``tool_calls`` rows inserted per commit. + + Returns: + ``(inserted, skipped)`` row counts. + """ + import json + import time + import uuid + + from pydantic_ai.messages import ModelMessagesTypeAdapter + + where_clause = "WHERE m.pydantic_json IS NOT NULL AND m.compacted = 0" + params: tuple = () + if session_id: + where_clause += " AND m.session_id = ?" + params = (session_id,) + + cursor = await db.execute( + f""" + SELECT m.id, m.session_id, m.seq, m.role, + m.agent_name, m.model_name, m.timestamp, m.pydantic_json + FROM messages m + {where_clause} + ORDER BY m.session_id, m.seq + """, + params, + ) + rows = await cursor.fetchall() + + if not rows: + logger.info("backfill_tool_calls: no eligible messages found — nothing to do") + return 0, 0 + + logger.info( + "backfill_tool_calls: processing %d message rows%s", + len(rows), + f" for session={session_id}" if session_id else "", + ) + + inserted = 0 + skipped = 0 + batch: list[tuple] = [] + + # Per-session state: pending tool calls keyed by tool_call_id + current_session: str | None = None + # tool_call_id → {name, args, parent_seq, agent, model, ts} + pending: dict[str, dict] = {} + + for row in rows: + row_session = row["session_id"] + if row_session != current_session: + # Reset pending state when we enter a new session + current_session = row_session + pending = {} + + try: + parsed = ModelMessagesTypeAdapter.validate_json(row["pydantic_json"]) + except Exception as exc: + logger.warning( + "backfill_tool_calls: skip seq=%s session=%s parse error: %s", + row["seq"], + row["session_id"], + exc, + ) + skipped += 1 + continue + + if not parsed: + continue + + msg = parsed[0] + role: str = row["role"] + seq: int = row["seq"] + agent_name: str = row["agent_name"] or "" + model_name: str = row["model_name"] or "" + + # Collect ToolCallParts from assistant (ModelResponse) messages + if role == "assistant": + for part in getattr(msg, "parts", []): + if type(part).__name__ != "ToolCallPart": + continue + tc_id = getattr(part, "tool_call_id", "") or str(uuid.uuid4()) + args = getattr(part, "args", {}) + if hasattr(args, "model_dump"): + args = args.model_dump() + elif hasattr(args, "__dict__"): + args = dict(args.__dict__) + pending[tc_id] = { + "name": getattr(part, "tool_name", "unknown"), + "args": args, + "parent_seq": seq, + "agent": agent_name, + "model": model_name, + "ts": row["timestamp"], + } + + # Match ToolReturnParts from user (ModelRequest) messages + elif role == "user": + for part in getattr(msg, "parts", []): + if type(part).__name__ != "ToolReturnPart": + continue + tc_id = getattr(part, "tool_call_id", "") + if not tc_id or tc_id not in pending: + continue + + tc = pending.pop(tc_id) + content = getattr(part, "content", "") + try: + result_val = ( + json.loads(content) if isinstance(content, str) else content + ) + except Exception: + result_val = content + + try: + args_json_str = json.dumps(tc["args"]) + except Exception: + args_json_str = str(tc["args"]) + try: + result_json_str = json.dumps(result_val) + except Exception: + result_json_str = str(result_val) + + try: + from datetime import datetime as _dt + + tool_ts = _dt.fromisoformat( + str(tc["ts"]).replace("Z", "+00:00") + ).timestamp() + except Exception: + tool_ts = time.time() + + batch.append( + ( + tc_id, # id + row_session, # session_id + tc["parent_seq"], # parent_message_seq + tc[ + "parent_seq" + ], # seq (best approximation without original seq) + tc["name"], # tool_name + args_json_str, # args_json + result_json_str, # result_json + "success", # status + None, # duration_ms + None, # error_text + tc["agent"], # agent_name + tc["model"], # model_name + tool_ts, # timestamp + ) + ) + + if len(batch) >= batch_size: + n = await _flush_tool_call_batch(db, batch) + inserted += n + batch = [] + + if batch: + n = await _flush_tool_call_batch(db, batch) + inserted += n + + logger.info( + "backfill_tool_calls: done — %d inserted, %d skipped", + inserted, + skipped, + ) + return inserted, skipped + + +async def _flush_tool_call_batch( + db: "aiosqlite.Connection", + batch: list[tuple], +) -> int: + """INSERT OR IGNORE a batch of tool_call rows and commit.""" + await db.executemany( + """ + INSERT OR IGNORE INTO tool_calls + (id, session_id, parent_message_seq, seq, tool_name, + args_json, result_json, status, duration_ms, error_text, + agent_name, model_name, timestamp) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + batch, + ) + await db.commit() + return len(batch) diff --git a/code_puppy/api/db/connection.py b/code_puppy/api/db/connection.py new file mode 100644 index 000000000..d60f9e2b2 --- /dev/null +++ b/code_puppy/api/db/connection.py @@ -0,0 +1,426 @@ +"""SQLite connection singleton using aiosqlite for non-blocking async access. + +The Desk Puppy Electron FE owns the schema (creates it at user_version=2). +This module opens the same database from the Python BE, runs any needed +migrations if the FE hasn't run yet, and exposes get_db() / init_db() / +close_db() for use by queries.py and the seeder. + +Schema versions: + 0 = no schema yet (fresh DB) + 1 = added token_count, compacted, pydantic_json, compaction_log_id to messages + 2 = added deleted_at to sessions (soft delete) + 3 = rebuilt messages table with UNIQUE(session_id, seq) + dedup + 4 = backfill blank agent_name / model_name in messages + tool_calls from session defaults + +Concurrency: + aiosqlite serialises all DB operations through its own internal background + thread — no threading.Lock required. All callers simply await the relevant + coroutine; the event loop is never blocked. +""" + +from __future__ import annotations + +import logging +import os +from pathlib import Path +from typing import Optional + +import aiosqlite + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Path resolution +# --------------------------------------------------------------------------- + +_PUPPY_DESK_DB_ENV = "PUPPY_DESK_DB" + + +def get_db_path() -> Path: + """Return path to the shared SQLite database. + + Override with PUPPY_DESK_DB env var for testing. + """ + env = os.environ.get(_PUPPY_DESK_DB_ENV) + if env: + return Path(env) + return Path.home() / ".puppy_desk" / "chat_messages.db" + + +# --------------------------------------------------------------------------- +# Full schema (v4) — created on first BE run if FE hasn't seeded it yet +# --------------------------------------------------------------------------- + +_SCHEMA_SQL = """ +CREATE TABLE IF NOT EXISTS sessions ( + session_id TEXT PRIMARY KEY, + title TEXT DEFAULT '', + agent_name TEXT DEFAULT 'code-puppy', + model_name TEXT DEFAULT '', + working_directory TEXT DEFAULT '', + pinned INTEGER DEFAULT 0, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + message_count INTEGER DEFAULT 0, + total_tokens INTEGER DEFAULT 0, + deleted_at TEXT +); +CREATE INDEX IF NOT EXISTS idx_sessions_updated ON sessions(updated_at DESC); +CREATE INDEX IF NOT EXISTS idx_sessions_deleted ON sessions(deleted_at); + +CREATE TABLE IF NOT EXISTS compaction_log ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + summary_text TEXT NOT NULL, + source_start INTEGER NOT NULL, + source_end INTEGER NOT NULL, + source_count INTEGER NOT NULL, + source_tokens INTEGER NOT NULL, + summary_tokens INTEGER NOT NULL, + strategy TEXT DEFAULT 'summarization', + created_at TEXT NOT NULL, + FOREIGN KEY (session_id) REFERENCES sessions(session_id) ON DELETE CASCADE +); +CREATE INDEX IF NOT EXISTS idx_compaction_session ON compaction_log(session_id); + +CREATE TABLE IF NOT EXISTS messages ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + seq INTEGER NOT NULL, + role TEXT NOT NULL, + content TEXT DEFAULT '', + type TEXT DEFAULT '', + agent_name TEXT DEFAULT '', + model_name TEXT DEFAULT '', + timestamp TEXT NOT NULL, + thinking TEXT, + attachments_json TEXT, + clean_content TEXT, + system_message_type TEXT, + system_message_path TEXT, + token_count INTEGER DEFAULT 0, + compacted INTEGER DEFAULT 0, + pydantic_json TEXT, + compaction_log_id INTEGER REFERENCES compaction_log(id) ON DELETE SET NULL, + FOREIGN KEY (session_id) REFERENCES sessions(session_id) ON DELETE CASCADE, + UNIQUE(session_id, seq) +); +CREATE INDEX IF NOT EXISTS idx_messages_session_seq ON messages(session_id, seq); +CREATE INDEX IF NOT EXISTS idx_messages_active ON messages(session_id, compacted); + +CREATE TABLE IF NOT EXISTS tool_calls ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + parent_message_seq INTEGER, + seq INTEGER NOT NULL, + tool_name TEXT NOT NULL, + args_json TEXT, + result_json TEXT, + status TEXT DEFAULT 'running', + duration_ms INTEGER, + error_text TEXT, + agent_name TEXT DEFAULT '', + model_name TEXT DEFAULT '', + timestamp REAL NOT NULL, + FOREIGN KEY (session_id) REFERENCES sessions(session_id) ON DELETE CASCADE +); +CREATE INDEX IF NOT EXISTS idx_tool_calls_session_seq ON tool_calls(session_id, seq); +CREATE INDEX IF NOT EXISTS idx_tool_calls_parent ON tool_calls(parent_message_seq); +""" + +# Migrations applied to existing databases that are below SCHEMA_VERSION +_MIGRATION_V0_TO_V1: list[str] = [ + "ALTER TABLE messages ADD COLUMN token_count INTEGER DEFAULT 0", + "ALTER TABLE messages ADD COLUMN compacted INTEGER DEFAULT 0", + "ALTER TABLE messages ADD COLUMN pydantic_json TEXT", + "ALTER TABLE messages ADD COLUMN compaction_log_id INTEGER REFERENCES compaction_log(id) ON DELETE SET NULL", +] + +_MIGRATION_V1_TO_V2: list[str] = [ + "ALTER TABLE sessions ADD COLUMN deleted_at TEXT", +] + +# v3: Rebuild messages table to add UNIQUE(session_id, seq). +# SQLite cannot add UNIQUE constraints via ALTER TABLE, so we must rebuild. +_MIGRATION_V2_TO_V3_SQL = """ +PRAGMA foreign_keys = OFF; +BEGIN; + +CREATE TABLE messages_v3 ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + seq INTEGER NOT NULL, + role TEXT NOT NULL, + content TEXT DEFAULT '', + type TEXT DEFAULT '', + agent_name TEXT DEFAULT '', + model_name TEXT DEFAULT '', + timestamp TEXT NOT NULL, + thinking TEXT, + attachments_json TEXT, + clean_content TEXT, + system_message_type TEXT, + system_message_path TEXT, + token_count INTEGER DEFAULT 0, + compacted INTEGER DEFAULT 0, + pydantic_json TEXT, + compaction_log_id INTEGER REFERENCES compaction_log(id) ON DELETE SET NULL, + FOREIGN KEY (session_id) REFERENCES sessions(session_id) ON DELETE CASCADE, + UNIQUE(session_id, seq) +); + +INSERT INTO messages_v3 +SELECT id, session_id, seq, role, content, type, agent_name, model_name, + timestamp, thinking, attachments_json, clean_content, + system_message_type, system_message_path, + token_count, compacted, pydantic_json, compaction_log_id +FROM messages AS m1 +WHERE id = ( + SELECT id FROM messages AS m2 + WHERE m2.session_id = m1.session_id AND m2.seq = m1.seq + ORDER BY + (m2.clean_content IS NOT NULL) DESC, + (m2.pydantic_json IS NOT NULL) DESC, + m2.id DESC + LIMIT 1 +); + +DROP TABLE messages; +ALTER TABLE messages_v3 RENAME TO messages; + +CREATE INDEX IF NOT EXISTS idx_messages_session_seq ON messages(session_id, seq); +CREATE INDEX IF NOT EXISTS idx_messages_active ON messages(session_id, compacted); + +COMMIT; +PRAGMA foreign_keys = ON; +""".strip() + +SCHEMA_VERSION = 5 + +# --------------------------------------------------------------------------- +# Async singleton (aiosqlite) +# --------------------------------------------------------------------------- + +_aconn: Optional[aiosqlite.Connection] = None + + +async def init_db() -> None: + """Open the database and ensure schema is at SCHEMA_VERSION. + + Safe to call multiple times — subsequent calls are no-ops. + """ + global _aconn + + if _aconn is not None: + return # already initialised + + db_path = get_db_path() + db_path.parent.mkdir(parents=True, exist_ok=True) + + logger.info("Opening aiosqlite DB at %s", db_path) + _aconn = await aiosqlite.connect(str(db_path)) + _aconn.row_factory = aiosqlite.Row + + # Performance + safety pragmas + await _aconn.execute("PRAGMA journal_mode = WAL") + await _aconn.execute("PRAGMA foreign_keys = ON") + await _aconn.execute("PRAGMA busy_timeout = 5000") + + cursor = await _aconn.execute("PRAGMA user_version") + row = await cursor.fetchone() + current_version: int = row[0] + + if current_version == 0: + # Fresh database — create everything at the current schema version + await _aconn.executescript(_SCHEMA_SQL) + await _aconn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") + await _aconn.commit() + logger.info( + "\u2713 aiosqlite DB created (schema v%d) at %s", SCHEMA_VERSION, db_path + ) + else: + # Existing database — run incremental migrations + await _run_alters_and_migrate(current_version) + + +async def _run_alters_and_migrate(current_version: int) -> None: + """Apply all pending schema migrations in order.""" + assert _aconn is not None + + if current_version < 1: + await _run_alters(_MIGRATION_V0_TO_V1) + await _aconn.execute("PRAGMA user_version = 1") + await _aconn.commit() + logger.info("\u2713 aiosqlite DB migrated v0 \u2192 v1") + + if current_version < 2: + await _run_alters(_MIGRATION_V1_TO_V2) + # Also create compaction_log and its index if missing (v0 DBs) + await _aconn.execute(""" + CREATE TABLE IF NOT EXISTS compaction_log ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + summary_text TEXT NOT NULL, + source_start INTEGER NOT NULL, + source_end INTEGER NOT NULL, + source_count INTEGER NOT NULL, + source_tokens INTEGER NOT NULL, + summary_tokens INTEGER NOT NULL, + strategy TEXT DEFAULT 'summarization', + created_at TEXT NOT NULL, + FOREIGN KEY (session_id) REFERENCES sessions(session_id) ON DELETE CASCADE + ) + """) + await _aconn.execute( + "CREATE INDEX IF NOT EXISTS idx_compaction_session ON compaction_log(session_id)" + ) + await _aconn.execute( + "CREATE INDEX IF NOT EXISTS idx_sessions_deleted ON sessions(deleted_at)" + ) + await _aconn.execute( + "CREATE INDEX IF NOT EXISTS idx_messages_active ON messages(session_id, compacted)" + ) + await _aconn.execute("PRAGMA user_version = 2") + await _aconn.commit() + logger.info("\u2713 aiosqlite DB migrated to v2") + + if current_version < 3: + logger.info( + "Running v2\u21923 migration: rebuilding messages table with UNIQUE(session_id, seq)..." + ) + await _aconn.executescript(_MIGRATION_V2_TO_V3_SQL) + await _aconn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") + await _aconn.commit() + cursor = await _aconn.execute("SELECT COUNT(*) FROM messages") + row = await cursor.fetchone() + count = row[0] + logger.info( + "\u2713 aiosqlite DB migrated to v%d (%d messages after dedup)", + SCHEMA_VERSION, + count, + ) + + if current_version < 4: + logger.info( + "Running v3\u2192v4 migration: backfilling blank agent_name/model_name \u2026" + ) + await _aconn.executescript(""" + -- messages: blank agent_name → session.agent_name + UPDATE messages + SET agent_name = ( + SELECT s.agent_name + FROM sessions s + WHERE s.session_id = messages.session_id + ) + WHERE (agent_name IS NULL OR agent_name = '') + AND ( + SELECT COALESCE(s.agent_name, '') + FROM sessions s + WHERE s.session_id = messages.session_id + ) != ''; + + -- messages: blank model_name → session.model_name + UPDATE messages + SET model_name = ( + SELECT s.model_name + FROM sessions s + WHERE s.session_id = messages.session_id + ) + WHERE (model_name IS NULL OR model_name = '') + AND ( + SELECT COALESCE(s.model_name, '') + FROM sessions s + WHERE s.session_id = messages.session_id + ) != ''; + + -- tool_calls: blank agent_name → session.agent_name + UPDATE tool_calls + SET agent_name = ( + SELECT s.agent_name + FROM sessions s + WHERE s.session_id = tool_calls.session_id + ) + WHERE (agent_name IS NULL OR agent_name = '') + AND ( + SELECT COALESCE(s.agent_name, '') + FROM sessions s + WHERE s.session_id = tool_calls.session_id + ) != ''; + + -- tool_calls: blank model_name → session.model_name + UPDATE tool_calls + SET model_name = ( + SELECT s.model_name + FROM sessions s + WHERE s.session_id = tool_calls.session_id + ) + WHERE (model_name IS NULL OR model_name = '') + AND ( + SELECT COALESCE(s.model_name, '') + FROM sessions s + WHERE s.session_id = tool_calls.session_id + ) != ''; + """) + await _aconn.execute("PRAGMA user_version = 4") + await _aconn.commit() + cursor = await _aconn.execute( + "SELECT COUNT(*) FROM messages WHERE agent_name != '' AND model_name != ''" + ) + row = await cursor.fetchone() + msg_count = row[0] + logger.info( + "\u2713 aiosqlite DB migrated to v4 (%d messages with agent+model)", + msg_count, + ) + + if current_version < 5: + logger.info( + "Running v4\u2192v5 migration: backfilling pydantic_json for NULL rows \u2026" + ) + from code_puppy.api.db.backfill import backfill_pydantic_json + + updated, skipped = await backfill_pydantic_json(_aconn) + await _aconn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") + await _aconn.commit() + logger.info( + "\u2713 aiosqlite DB migrated to v5 " + "(%d pydantic_json rows backfilled, %d skipped)", + updated, + skipped, + ) + + logger.info("\u2713 aiosqlite DB ready (schema v%d)", SCHEMA_VERSION) + + +async def _run_alters(statements: list[str]) -> None: + """Run ALTER TABLE statements, skipping 'duplicate column' errors gracefully.""" + assert _aconn is not None + for sql in statements: + try: + await _aconn.execute(sql) + except Exception as exc: + if "duplicate column" in str(exc).lower(): + logger.debug("Skipping already-applied migration: %s", sql) + else: + raise + + +def get_db() -> aiosqlite.Connection: + """Return the open aiosqlite connection. + + Raises RuntimeError if init_db() has not been awaited. + """ + if _aconn is None: + raise RuntimeError( + "aiosqlite DB not initialised — call `await init_db()` first" + ) + return _aconn + + +async def close_db() -> None: + """Close the database connection. Called during app shutdown.""" + global _aconn + if _aconn is not None: + await _aconn.close() + _aconn = None + logger.info("aiosqlite DB closed") diff --git a/code_puppy/api/db/message_utils.py b/code_puppy/api/db/message_utils.py new file mode 100644 index 000000000..8c38a1851 --- /dev/null +++ b/code_puppy/api/db/message_utils.py @@ -0,0 +1,219 @@ +"""Shared helpers for extracting display fields from pydantic-ai ModelMessages. + +Used by the SQLite persistence helpers and session_context.py. +Must not import from seeder.py or session_context.py to avoid circular imports. +""" + +from __future__ import annotations + +import logging +import warnings +from typing import Any, Optional + +logger = logging.getLogger(__name__) + +_PART_TYPE_TEXT = {"TextPart", "UserPromptPart", "SystemPromptPart"} +_PART_TYPE_THINKING = {"ThinkingPart"} + + +def get_role(msg: Any) -> str: + """Map a ModelMessage to a DB role string.""" + kind = getattr(msg, "kind", None) + if kind == "request": + return "user" + if kind == "response": + return "assistant" + return "system" + + +def extract_content(msg: Any) -> str: + """Extract human-readable text from a ModelMessage's parts.""" + parts = getattr(msg, "parts", []) + texts = [] + for part in parts: + if type(part).__name__ in _PART_TYPE_TEXT: + c = getattr(part, "content", "") + if isinstance(c, str) and c.strip(): + texts.append(c) + return "\n".join(texts) + + +def extract_thinking(msg: Any) -> Optional[str]: + """Extract extended thinking text from a ModelResponse's parts.""" + parts = getattr(msg, "parts", []) + texts = [ + getattr(p, "content", "") + for p in parts + if type(p).__name__ in _PART_TYPE_THINKING + and isinstance(getattr(p, "content", ""), str) + ] + joined = "\n".join(t for t in texts if t) + return joined or None + + +def get_message_timestamp(msg: Any) -> Optional[str]: + """Return the message's own timestamp as an ISO string, or None.""" + ts = getattr(msg, "timestamp", None) + if ts is None: + return None + try: + return ts.isoformat() + except Exception: + return str(ts) + + +def pydantic_json_for_message(msg: Any) -> Optional[str]: + """Serialise a single ModelMessage to a JSON string via ModelMessagesTypeAdapter. + + Stores as a single-element list: `[msg_json]`. + Deserialise with: `ModelMessagesTypeAdapter.validate_json(s)[0]` + + Returns None if serialisation fails for unsupported historical message shapes. + """ + try: + from pydantic_ai.messages import ModelMessagesTypeAdapter + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return ModelMessagesTypeAdapter.dump_json([msg]).decode("utf-8") + except Exception as exc: + logger.debug("pydantic_json serialisation failed: %s", exc) + return None + + +_PART_TYPE_TOOL_CALL = {"ToolCallPart"} +# Extended set of tool return part type names for runtime compatibility +# Different versions of pydantic-ai and streaming may use different class names +_PART_TYPE_TOOL_RETURN = {"ToolReturnPart", "ToolReturn", "RetryPromptPart"} + + +def extract_tool_calls(msg: Any) -> list[dict[str, Any]]: + """Extract ToolCallPart entries from a ModelResponse. + + Returns list of dicts: {id, name, args_dict} + """ + import uuid + + parts = getattr(msg, "parts", []) + result = [] + for part in parts: + if type(part).__name__ not in _PART_TYPE_TOOL_CALL: + continue + args = getattr(part, "args", {}) + if hasattr(args, "model_dump"): + args = args.model_dump() + elif hasattr(args, "__dict__"): + args = dict(args.__dict__) + result.append( + { + "id": getattr(part, "tool_call_id", "") or str(uuid.uuid4()), + "name": getattr(part, "tool_name", "unknown"), + "args": args, + } + ) + return result + + +def _is_tool_return_like(part: Any) -> bool: + """Check if a part is a tool return using duck-typing. + + Handles various runtime representations: + - ToolReturnPart (standard pydantic-ai) + - ToolReturn (streaming/runtime variant) + - Any object with tool_call_id + (content or return_value) + """ + class_name = type(part).__name__ + + # Fast path: known class names + if class_name in _PART_TYPE_TOOL_RETURN: + return True + + # Duck-typing fallback: must have tool_call_id and some result content + has_tool_call_id = hasattr(part, "tool_call_id") and getattr( + part, "tool_call_id", None + ) + has_content = hasattr(part, "content") or hasattr(part, "return_value") + + return bool(has_tool_call_id and has_content) + + +def _extract_tool_return_data(part: Any) -> dict[str, Any]: + """Extract tool return data from a part using multiple field strategies. + + Handles different field naming conventions: + - content (ToolReturnPart) + - return_value (some runtime variants) + - result (potential future variants) + """ + import json + + # Get the result content - try multiple field names + raw_content = None + for field in ("content", "return_value", "result"): + if hasattr(part, field): + raw_content = getattr(part, field, None) + if raw_content is not None: + break + + # Parse JSON if string, otherwise use as-is + if raw_content is None: + result_val = None + elif isinstance(raw_content, str): + try: + result_val = json.loads(raw_content) + except (json.JSONDecodeError, TypeError): + result_val = raw_content + else: + result_val = raw_content + + # Get tool name - try multiple field names + tool_name = "unknown" + for field in ("tool_name", "name"): + if hasattr(part, field): + name = getattr(part, field, None) + if name: + tool_name = name + break + + return { + "id": getattr(part, "tool_call_id", "") or "", + "name": tool_name, + "result": result_val, + } + + +def extract_tool_returns(msg: Any) -> list[dict[str, Any]]: + """Extract tool return entries from a ModelRequest. + + Returns list of dicts: {id, name, result} + + Handles multiple runtime representations: + - ToolReturnPart (standard pydantic-ai) + - ToolReturn (streaming variant) + - Duck-typed objects with tool_call_id + content/return_value + + This extended compatibility ensures tool results are persisted to SQLite + regardless of which pydantic-ai version or streaming mode is in use. + """ + parts = getattr(msg, "parts", []) + result = [] + + for part in parts: + if not _is_tool_return_like(part): + continue + + try: + data = _extract_tool_return_data(part) + if data["id"]: # Only add if we have a valid tool_call_id + result.append(data) + else: + logger.debug( + "Skipping tool return with empty tool_call_id: %s", + type(part).__name__, + ) + except Exception as exc: + logger.warning( + "Failed to extract tool return from %s: %s", type(part).__name__, exc + ) + + return result diff --git a/code_puppy/api/db/queries.py b/code_puppy/api/db/queries.py new file mode 100644 index 000000000..1a26ad0b1 --- /dev/null +++ b/code_puppy/api/db/queries.py @@ -0,0 +1,1154 @@ +"""Typed async query helpers for the shared chat_messages.db. + +All functions are coroutines. aiosqlite serialises writes internally via its +own background thread — no threading.Lock is needed here. +""" + +from __future__ import annotations + +import json +import logging +from typing import Any, Optional + +from code_puppy.api.db.connection import get_db +from code_puppy.api.db.sql_loader import load_sql + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Sessions +# --------------------------------------------------------------------------- + + +async def session_exists(session_id: str) -> bool: + """Return True if a session row already exists (regardless of deleted_at).""" + db = get_db() + cursor = await db.execute( + "SELECT 1 FROM sessions WHERE session_id = ? LIMIT 1", (session_id,) + ) + row = await cursor.fetchone() + return row is not None + + +async def get_session_row(session_id: str) -> Optional[dict]: + """Return the full session row as a plain dict, or None if not found. + + Non-fatal: returns None on any exception (DB unavailable, etc.). + """ + try: + db = get_db() + cursor = await db.execute( + "SELECT * FROM sessions WHERE session_id = ?", (session_id,) + ) + row = await cursor.fetchone() + if row is None: + return None + return dict(row) + except Exception: + return None + + +async def get_session_metadata(session_id: str) -> Optional[dict]: + """Return a subset of session columns useful for the chat handler. + + Returns keys: session_id, title, agent_name, model_name, + working_directory, pinned, created_at — or None if not found. + Non-fatal: returns None on any exception. + """ + try: + db = get_db() + cursor = await db.execute( + """ + SELECT session_id, title, agent_name, model_name, + working_directory, pinned, created_at + FROM sessions + WHERE session_id = ? + """, + (session_id,), + ) + row = await cursor.fetchone() + if row is None: + return None + return dict(row) + except Exception: + return None + + +async def upsert_session( + *, + session_id: str, + title: str = "", + agent_name: str = "code-puppy", + model_name: str = "", + working_directory: str = "", + pinned: bool = False, + created_at: str, + updated_at: str, + message_count: int = 0, + total_tokens: int = 0, + deleted_at: Optional[str] = None, +) -> None: + """Insert or replace a session row.""" + db = get_db() + try: + await db.execute( + """ + INSERT INTO sessions + (session_id, title, agent_name, model_name, working_directory, + pinned, created_at, updated_at, message_count, total_tokens, deleted_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(session_id) DO UPDATE SET + title = excluded.title, + agent_name = excluded.agent_name, + model_name = excluded.model_name, + working_directory = excluded.working_directory, + pinned = excluded.pinned, + updated_at = excluded.updated_at, + message_count = excluded.message_count, + total_tokens = excluded.total_tokens, + deleted_at = COALESCE(sessions.deleted_at, excluded.deleted_at) + """, + ( + session_id, + title, + agent_name, + model_name, + working_directory, + 1 if pinned else 0, + created_at, + updated_at, + message_count, + total_tokens, + deleted_at, + ), + ) + await db.commit() + except Exception: + await db.rollback() + raise + + +async def update_session_stats( + session_id: str, + *, + message_count: int, + total_tokens: int, + updated_at: str, +) -> None: + """Bump message_count, total_tokens, and updated_at after a save.""" + db = get_db() + try: + await db.execute( + """ + UPDATE sessions + SET message_count = ?, total_tokens = ?, updated_at = ? + WHERE session_id = ? + """, + (message_count, total_tokens, updated_at, session_id), + ) + await db.commit() + except Exception: + await db.rollback() + raise + + +async def update_session_working_directory( + session_id: str, + working_directory: str, + updated_at: str, +) -> None: + """Update only the working_directory and updated_at columns for a session. + + Narrow UPDATE — does not touch title, agent, model, pinned, counts, etc. + Called after a successful set_working_directory WS message. + """ + db = get_db() + await db.execute( + "UPDATE sessions SET working_directory = ?, updated_at = ? WHERE session_id = ?", + (working_directory, updated_at, session_id), + ) + await db.commit() + + +async def update_session_meta_fields( + session_id: str, + *, + title: str, + pinned: bool, + updated_at: str, +) -> None: + """Update only title, pinned, and updated_at for an existing session. + + Narrow UPDATE — does NOT touch message_count, total_tokens, agent_name, + model_name, or working_directory. Called by the update_session_meta WS + message handler (user renames or pins/unpins a session). + """ + db = get_db() + await db.execute( + "UPDATE sessions SET title = ?, pinned = ?, updated_at = ? WHERE session_id = ?", + (title, 1 if pinned else 0, updated_at, session_id), + ) + await db.commit() + + +async def soft_delete_session(session_id: str, deleted_at: str) -> None: + """Set deleted_at on a session (soft delete). Data is preserved.""" + db = get_db() + await db.execute( + "UPDATE sessions SET deleted_at = ? WHERE session_id = ?", + (deleted_at, session_id), + ) + await db.commit() + + +# --------------------------------------------------------------------------- +# Messages +# --------------------------------------------------------------------------- + + +async def get_next_seq(session_id: str) -> int: + """Return the next available seq number for a session (1-based).""" + db = get_db() + cursor = await db.execute( + "SELECT COALESCE(MAX(seq), 0) FROM messages WHERE session_id = ?", + (session_id,), + ) + row = await cursor.fetchone() + return (row[0] or 0) + 1 + + +async def _insert_immediate_message_and_sync_session( + *, + session_id: str, + role: str, + content: str, + type: str, + agent_name: str = "", + model_name: str = "", + timestamp: str, + thinking: Optional[str] = None, + attachments_json: Optional[str] = None, + clean_content: Optional[str] = None, + system_message_type: Optional[str] = None, + system_message_path: Optional[str] = None, + token_count: int = 0, + increment_message_count: bool = True, +) -> int: + """Insert an immediate WS message and sync session metadata atomically. + + This path is used for system/error frames that must persist immediately, + outside the normal end-of-turn batch save. The session row is created if + needed, the next seq is allocated atomically, the message row is inserted, + and session metadata is updated in the same transaction so + message_count/updated_at cannot drift. + + increment_message_count controls whether the sessions.message_count should + track this row. Persisted system banners/config rows should not affect the + user-visible chat turn count, while persisted error rows should. + """ + db = get_db() + try: + await db.execute("BEGIN IMMEDIATE") + cursor = await db.execute( + "SELECT 1 FROM sessions WHERE session_id = ?", (session_id,) + ) + existing = await cursor.fetchone() + if not existing: + await db.execute( + """ + INSERT INTO sessions + (session_id, title, agent_name, model_name, working_directory, + pinned, created_at, updated_at, message_count, total_tokens, deleted_at) + VALUES (?, '', ?, ?, '', 0, ?, ?, 0, 0, NULL) + """, + ( + session_id, + agent_name or "code-puppy", + model_name or "", + timestamp, + timestamp, + ), + ) + + cursor = await db.execute( + "SELECT COALESCE(MAX(seq), 0) + 1 AS next_seq FROM messages WHERE session_id = ?", + (session_id,), + ) + seq_row = await cursor.fetchone() + seq = int(seq_row["next_seq"] if seq_row is not None else 1) + + await db.execute( + """ + INSERT INTO messages + (session_id, seq, role, content, type, agent_name, model_name, + timestamp, thinking, attachments_json, clean_content, + system_message_type, system_message_path, + token_count, compacted, pydantic_json, compaction_log_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, NULL, NULL) + """, + ( + session_id, + seq, + role, + content, + type, + agent_name, + model_name, + timestamp, + thinking, + attachments_json, + clean_content, + system_message_type, + system_message_path, + token_count, + ), + ) + await db.execute( + """ + UPDATE sessions + SET message_count = COALESCE(message_count, 0) + ?, + total_tokens = COALESCE(total_tokens, 0) + ?, + updated_at = ? + WHERE session_id = ? + """, + ( + 1 if increment_message_count else 0, + token_count, + timestamp, + session_id, + ), + ) + await db.commit() + return seq + except Exception: + await db.rollback() + raise + + +async def write_system_message_to_sqlite( + *, + session_id: str, + system_message_type: str, + content: str, + system_message_path: str = "", + agent_name: str = "", + model_name: str = "", + timestamp: Optional[str] = None, +) -> None: + """Write a single system-event row to SQLite immediately. + + Used by the WS chat handler to persist agent/model switches, CWD changes, + and new-session initialisation messages so the FE cold-load path sees them + without relying on ephemeral WS-only state. + + For directory banners, deduplicates to avoid multiple banners for the same path. + """ + import datetime as _dt + + ts = timestamp or _dt.datetime.now(_dt.timezone.utc).isoformat() + + # Dedup: for directory system messages, check if one already exists with same path + if system_message_type == "directory" and system_message_path: + try: + db = get_db() + cursor = await db.execute( + """ + SELECT 1 FROM messages + WHERE session_id = ? + AND system_message_type = 'directory' + AND system_message_path = ? + LIMIT 1 + """, + (session_id, system_message_path), + ) + existing = await cursor.fetchone() + if existing: + import logging + + logging.getLogger(__name__).debug( + "Skipping duplicate directory banner: session=%s path=%s", + session_id, + system_message_path, + ) + return # Already have a banner for this path + except Exception as e: + import logging + + logging.getLogger(__name__).warning("Directory dedup check failed: %s", e) + # Continue with insert on error + + try: + await _insert_immediate_message_and_sync_session( + session_id=session_id, + role="system", + content=content, + type="system", + agent_name=agent_name, + model_name=model_name, + timestamp=ts, + system_message_type=system_message_type, + system_message_path=system_message_path, + increment_message_count=False, + ) + except Exception as exc: + import logging + + logging.getLogger(__name__).warning( + "write_system_message_to_sqlite failed (session=%s type=%s): %s", + session_id, + system_message_type, + exc, + exc_info=True, + ) + + +async def write_error_message_to_sqlite( + *, + session_id: str, + error: str, + error_type: str = "unknown", + technical_details: str = "", + action_required: Optional[str] = None, + agent_name: str = "", + model_name: str = "", + timestamp: Optional[str] = None, +) -> None: + """Write a structured error row to SQLite immediately. + + Error rows are stored as regular message rows with type='error' so they + survive session reloads. The structured payload is serialized into + attachments_json for the read path to reconstruct for the frontend. + """ + import datetime as _dt + + ts = timestamp or _dt.datetime.now(_dt.timezone.utc).isoformat() + payload = { + "error": error, + "error_type": error_type, + "technical_details": technical_details, + "action_required": action_required, + "session_id": session_id, + } + + try: + await _insert_immediate_message_and_sync_session( + session_id=session_id, + role="system", + content=error, + type="error", + agent_name=agent_name, + model_name=model_name, + timestamp=ts, + clean_content=error, + attachments_json=json.dumps(payload), + ) + except Exception as exc: + logging.getLogger(__name__).warning( + "write_error_message_to_sqlite failed (session=%s error_type=%s): %s", + session_id, + error_type, + exc, + exc_info=True, + ) + + +async def insert_message( + *, + session_id: str, + seq: int, + role: str, + content: str = "", + type: str = "", + agent_name: str = "", + model_name: str = "", + timestamp: str, + thinking: Optional[str] = None, + attachments_json: Optional[str] = None, + clean_content: Optional[str] = None, + system_message_type: Optional[str] = None, + system_message_path: Optional[str] = None, + token_count: int = 0, + compacted: int = 0, + pydantic_json: Optional[str] = None, + compaction_log_id: Optional[int] = None, +) -> None: + """Insert a single message row.""" + # Guard: skip empty visible content for user/assistant roles + # unless attachments are present (non-text payloads are still valid). + if ( + role in ("user", "assistant") + and not content.strip() + and not attachments_json + and not pydantic_json + ): + logger.debug( + "insert_message: skipping empty %s message for session=%s seq=%d", + role, + session_id, + seq, + ) + return + + db = get_db() + await db.execute( + """ + INSERT OR IGNORE INTO messages + (session_id, seq, role, content, type, agent_name, model_name, + timestamp, thinking, attachments_json, clean_content, + system_message_type, system_message_path, + token_count, compacted, pydantic_json, compaction_log_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + session_id, + seq, + role, + content, + type, + agent_name, + model_name, + timestamp, + thinking, + attachments_json, + clean_content, + system_message_type, + system_message_path, + token_count, + compacted, + pydantic_json, + compaction_log_id, + ), + ) + await db.commit() + + +async def insert_messages_batch(rows: list[dict[str, Any]]) -> None: + """Insert multiple message rows in a single transaction. + + Each dict in *rows* must have all the same keys as insert_message kwargs. + Uses OR IGNORE so duplicate (session_id, seq) pairs are skipped silently. + Rolls back on any failure so later immediate error/system writes can still + open a fresh transaction. + """ + if not rows: + return + + # Guard: filter out empty visible content for user/assistant roles, + # unless attachments are present (non-text payloads are still valid). + original_count = len(rows) + rows = [ + r + for r in rows + if not ( + r.get("role") in ("user", "assistant") + and not r.get("content", "").strip() + and not r.get("attachments_json") + and not r.get("pydantic_json") + ) + ] + dropped = original_count - len(rows) + if dropped > 0: + logger.info( + "insert_messages_batch: filtered %d empty user/assistant rows", + dropped, + ) + if not rows: + return + + db = get_db() + params = [ + ( + r["session_id"], + r["seq"], + r["role"], + r.get("content", ""), + r.get("type", ""), + r.get("agent_name", ""), + r.get("model_name", ""), + r["timestamp"], + r.get("thinking"), + r.get("attachments_json"), + r.get("clean_content"), + r.get("system_message_type"), + r.get("system_message_path"), + r.get("token_count", 0), + r.get("compacted", 0), + r.get("pydantic_json"), + r.get("compaction_log_id"), + ) + for r in rows + ] + try: + await db.executemany( + """ + INSERT OR IGNORE INTO messages + (session_id, seq, role, content, type, agent_name, model_name, + timestamp, thinking, attachments_json, clean_content, + system_message_type, system_message_path, + token_count, compacted, pydantic_json, compaction_log_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + params, + ) + await db.commit() + except Exception: + await db.rollback() + raise + + +async def get_active_messages(session_id: str) -> list[dict[str, Any]]: + """Return all non-compacted message rows for a session, ordered by seq. + + This is what the agent uses to reconstruct its context via pydantic_json. + """ + db = get_db() + cursor = await db.execute( + """ + SELECT session_id, seq, role, content, type, agent_name, model_name, + timestamp, thinking, attachments_json, clean_content, + system_message_type, system_message_path, + token_count, compacted, pydantic_json, compaction_log_id + FROM messages + WHERE session_id = ? AND compacted = 0 + ORDER BY seq + """, + (session_id,), + ) + rows = await cursor.fetchall() + return [dict(r) for r in rows] + + +# --------------------------------------------------------------------------- +# Tool calls +# --------------------------------------------------------------------------- + + +async def insert_tool_call( + *, + id: str, + session_id: str, + parent_message_seq: Optional[int], + seq: int, + tool_name: str, + args_json: Optional[str] = None, + result_json: Optional[str] = None, + status: str = "success", + duration_ms: Optional[int] = None, + error_text: Optional[str] = None, + agent_name: str = "", + model_name: str = "", + timestamp: float, +) -> None: + db = get_db() + await db.execute( + """ + INSERT OR IGNORE INTO tool_calls + (id, session_id, parent_message_seq, seq, tool_name, + args_json, result_json, status, duration_ms, error_text, + agent_name, model_name, timestamp) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + id, + session_id, + parent_message_seq, + seq, + tool_name, + args_json, + result_json, + status, + duration_ms, + error_text, + agent_name, + model_name, + timestamp, + ), + ) + await db.commit() + + +async def insert_tool_calls_batch(rows: list[dict[str, Any]]) -> None: + """Insert multiple tool call rows in a single transaction.""" + if not rows: + return + db = get_db() + params = [ + ( + r["id"], + r["session_id"], + r.get("parent_message_seq"), + r["seq"], + r["tool_name"], + r.get("args_json"), + r.get("result_json"), + r.get("status", "success"), + r.get("duration_ms"), + r.get("error_text"), + r.get("agent_name", ""), + r.get("model_name", ""), + r["timestamp"], + ) + for r in rows + ] + await db.executemany( + """ + INSERT OR IGNORE INTO tool_calls + (id, session_id, parent_message_seq, seq, tool_name, + args_json, result_json, status, duration_ms, error_text, + agent_name, model_name, timestamp) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + params, + ) + await db.commit() + + +# --------------------------------------------------------------------------- +# Compaction +# --------------------------------------------------------------------------- + + +async def mark_messages_compacted( + session_id: str, start_seq: int, end_seq: int +) -> None: + """Soft-delete a range of messages by setting compacted=1.""" + db = get_db() + await db.execute( + """ + UPDATE messages SET compacted = 1 + WHERE session_id = ? AND seq BETWEEN ? AND ? AND compacted = 0 + """, + (session_id, start_seq, end_seq), + ) + await db.commit() + + +async def insert_compaction_log( + *, + session_id: str, + summary_text: str, + source_start: int, + source_end: int, + source_count: int, + source_tokens: int, + summary_tokens: int, + strategy: str = "summarization", + created_at: str, +) -> int: + """Insert a compaction log entry and return its id.""" + db = get_db() + cursor = await db.execute( + """ + INSERT INTO compaction_log + (session_id, summary_text, source_start, source_end, + source_count, source_tokens, summary_tokens, strategy, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + session_id, + summary_text, + source_start, + source_end, + source_count, + source_tokens, + summary_tokens, + strategy, + created_at, + ), + ) + await db.commit() + return cursor.lastrowid # type: ignore[return-value] + + +# --------------------------------------------------------------------------- +# Per-turn chat write +# --------------------------------------------------------------------------- + + +async def write_turn_to_sqlite( + *, + session_id: str, + enhanced_history: list[Any], + title: str = "", + working_directory: str = "", + pinned: bool = False, + agent_name: str = "code-puppy", + model_name: str = "", + total_tokens: int = 0, + updated_at: str, + created_at: str, + ctx: Any, +) -> None: + """Write an entire enhanced_history to SQLite after a chat turn. + + Called from chat_handler.py after save_session() succeeds. + Uses INSERT OR IGNORE throughout — fully idempotent. + + Message seq: For non-user messages, we query the current max seq from the DB + and increment from there. User messages are skipped (already written pre-stream + via insert_message() with get_next_seq()). + Tool call seq: max_message_seq + tc_global_idx + 1 (after all messages). + Tool calls carry parent_message_seq so the FE doesn't need to infer it. + + Args: + enhanced_history: List of dicts (already-wrapped format from chat_handler): + {'msg': ModelMessage, 'agent': str, 'model': str, 'ts': str} + Possibly also: {'clean_content': str, 'attachments': [...]} + ctx: SessionContext — used for estimate_tokens_for_message(). + Pass None to use a simple len(content)//4 fallback. + """ + import re + import time + from datetime import datetime + + from code_puppy.api.db.message_utils import ( + extract_content, + extract_thinking, + extract_tool_calls, + extract_tool_returns, + get_message_timestamp, + get_role, + pydantic_json_for_message, + ) + + # Upsert session first + try: + await upsert_session( + session_id=session_id, + title=title, + agent_name=agent_name, + model_name=model_name, + working_directory=working_directory, + pinned=pinned, + created_at=created_at, + updated_at=updated_at, + message_count=len(enhanced_history), + total_tokens=total_tokens, + deleted_at=None, + ) + except Exception as exc: + logger.warning("write_turn_to_sqlite: upsert_session failed: %s", exc) + return + + # ------------------------------------------------------------------ # + # Build message rows and collect pending tool calls # + # ------------------------------------------------------------------ # + message_rows: list[dict[str, Any]] = [] + # pending_tool_calls: tool_call_id → {name, args, parent_seq, agent, model, ts} + pending_tool_calls: dict[str, dict[str, Any]] = {} + tool_call_rows_pending: list[dict[str, Any]] = [] + + _ctx_re = re.compile(r"^\[Session Context:[^\]]*\]\n\n?") + + # Get the current max seq from the DB so we can assign correct seq numbers + # to non-user messages. User messages are already written pre-stream with + # get_next_seq(), so we must not collide with them. + current_max_seq = await get_next_seq(session_id) - 1 # get_next_seq returns max+1 + + for idx, item in enumerate(enhanced_history): + # Unwrap the enhanced wrapper + if isinstance(item, dict) and "msg" in item: + actual_msg = item["msg"] + wrapper = item + else: + actual_msg = item + wrapper = {} + + if not hasattr(actual_msg, "parts"): + continue # skip system dict entries + + role = get_role(actual_msg) + ts = wrapper.get("ts") or get_message_timestamp(actual_msg) or updated_at + content = extract_content(actual_msg) + + if role in {"user", "assistant"} and not content.strip(): + diag_tool_calls_present = bool(extract_tool_calls(actual_msg)) + diag_tool_returns_present = bool(extract_tool_returns(actual_msg)) + logger.debug( + "write_turn_to_sqlite empty visible content: " + "session=%s role=%s history_idx=%d has_tool_calls=%s has_tool_returns=%s", + session_id, + role, + idx, + diag_tool_calls_present, + diag_tool_returns_present, + ) + + # Extract tool returns from user (ModelRequest) messages BEFORE skipping. + # BUG FIX: previously this block was placed AFTER the early `continue` + # below, making it dead code — tool_call_rows_pending was never populated + # and insert_tool_calls_batch() always received an empty list. + if role == "user": + for tr in extract_tool_returns(actual_msg): + tid = tr["id"] + if tid in pending_tool_calls: + tc = pending_tool_calls.pop(tid) + try: + args_json_str = json.dumps(tc["args"]) + except Exception: + args_json_str = str(tc["args"]) + # Serialize result - handle Pydantic models and complex objects + result_val = tr["result"] + try: + # Try Pydantic model_dump first + if hasattr(result_val, "model_dump"): + result_json_str = json.dumps(result_val.model_dump()) + elif hasattr(result_val, "dict"): + # Older Pydantic v1 style + result_json_str = json.dumps(result_val.dict()) + elif hasattr(result_val, "__dict__"): + # Generic object - convert to dict + result_json_str = json.dumps(vars(result_val)) + else: + result_json_str = json.dumps(result_val) + except Exception: + # Last resort - try to make it JSON-serializable + try: + result_json_str = json.dumps(str(result_val)) + except Exception: + result_json_str = json.dumps({"raw": str(result_val)}) + + try: + tool_ts = datetime.fromisoformat( + str(tc["ts"]).replace("Z", "+00:00") + ).timestamp() + except Exception: + tool_ts = time.time() + + tool_call_rows_pending.append( + { + "id": tid, + "session_id": session_id, + "parent_message_seq": tc["parent_seq"], + "tool_name": tc["name"], + "args_json": args_json_str, + "result_json": result_json_str, + "status": "success", + "agent_name": tc["agent"], + "model_name": tc["model"], + "timestamp": tool_ts, + } + ) + # Skip writing user message rows — already pre-written by + # chat_handler.py via insert_message()/get_next_seq(). Writing + # them again here would create duplicates. + continue + + # Assign seq for non-user messages: increment from current max + current_max_seq += 1 + seq = current_max_seq + + thinking = extract_thinking(actual_msg) + pj = pydantic_json_for_message(actual_msg) + + msg_agent = wrapper.get("agent") or agent_name or "code-puppy" + msg_model = wrapper.get("model") or model_name or "unknown" + + # clean_content: prefer wrapper value, then strip session context for user + clean_content: str | None = wrapper.get("clean_content") or None + if clean_content is None and role == "user" and content: + stripped = _ctx_re.sub("", content).lstrip() + if stripped != content: + clean_content = stripped + + # attachments_json: from wrapper + raw_attachments = wrapper.get("attachments") + attachments_json: str | None = ( + json.dumps(raw_attachments) if raw_attachments else None + ) + + # token count + try: + token_count = ( + ctx.agent.estimate_tokens_for_message(actual_msg) + if ctx + else max(1, len(content) // 4) + ) + except Exception: + token_count = max(1, len(content) // 4) + + assistant_tool_calls = ( + extract_tool_calls(actual_msg) if role == "assistant" else [] + ) + + # Preserve non-text assistant payloads (tool-call-only responses). + # This keeps the row from being filtered by insert_messages_batch. + if role == "assistant" and assistant_tool_calls and not attachments_json: + attachments_json = json.dumps({"tool_calls": assistant_tool_calls}) + + # Guard: skip only truly empty assistant messages + # (no text AND no tool calls). + if role == "assistant" and not content.strip() and not assistant_tool_calls: + logger.debug( + "write_turn_to_sqlite: skipping empty assistant message " + "session=%s history_idx=%d (no text, no tool calls)", + session_id, + idx, + ) + continue # truly empty - skip + + message_rows.append( + { + "session_id": session_id, + "seq": seq, + "role": role, + "content": content, + "type": type(actual_msg).__name__, + "agent_name": msg_agent, + "model_name": msg_model, + "timestamp": ts, + "thinking": thinking, + "attachments_json": attachments_json, + "clean_content": clean_content, + "token_count": token_count, + "pydantic_json": pj, + } + ) + + # Collect tool calls from assistant messages + if role == "assistant": + for tc in assistant_tool_calls: + pending_tool_calls[tc["id"]] = { + "name": tc["name"], + "args": tc["args"], + "parent_seq": seq, + "agent": msg_agent, + "model": msg_model, + "ts": ts, + } + + # Assign seqs to tool calls: current_max_seq + 1, current_max_seq + 2, ... + # (current_max_seq was already incremented for each message we added) + tool_call_rows: list[dict[str, Any]] = [] + for tc_idx, tc_row in enumerate(tool_call_rows_pending): + tc_row["seq"] = current_max_seq + tc_idx + 1 + tool_call_rows.append(tc_row) + + # Write everything + try: + await insert_messages_batch(message_rows) + except Exception as exc: + logger.warning("write_turn_to_sqlite: insert_messages_batch failed: %s", exc) + + if tool_call_rows: + try: + await insert_tool_calls_batch(tool_call_rows) + except Exception as exc: + logger.warning( + "write_turn_to_sqlite: insert_tool_calls_batch failed: %s", exc + ) + + # Update session stats + computed_tokens = sum(r.get("token_count", 0) for r in message_rows) + final_tokens = total_tokens if total_tokens > 0 else computed_tokens + try: + await update_session_stats( + session_id, + message_count=len(message_rows), + total_tokens=final_tokens, + updated_at=updated_at, + ) + except Exception as exc: + logger.warning("write_turn_to_sqlite: update_session_stats failed: %s", exc) + + # Log with diagnostic counts for forensic debugging + unmatched_calls = len(pending_tool_calls) + logger.info( + "write_turn_to_sqlite: session=%s messages=%d tool_calls=%d pending_unmatched=%d", + session_id, + len(message_rows), + len(tool_call_rows), + unmatched_calls, + ) + + # Warn if tool calls were detected but not matched to returns + if unmatched_calls > 0: + logger.warning( + "write_turn_to_sqlite: %d tool calls not matched to returns for session=%s: %s", + unmatched_calls, + session_id, + list(pending_tool_calls.keys()), + ) + + +# --------------------------------------------------------------------------- +# Interleaved session history query (Full Parity Model) +# +# Returns messages + tool_calls in a single ordered result set with a +# row_type discriminator, matching the desktop GUI's SESSION_LOAD_SQL. +# This enables the browser API to return the same data shape as the desktop. +# --------------------------------------------------------------------------- + +# SQL query for interleaved messages + tool_calls, ordered by seq +_SESSION_HISTORY_PARITY_SQL = load_sql("session_history_parity.sql") +_SESSION_HISTORY_PARITY_NO_COMPACTED_SQL = load_sql( + "session_history_parity_no_compacted.sql" +) + + +async def get_session_history_parity( + session_id: str, + *, + include_compacted: bool = True, +) -> list[dict[str, Any]]: + """Return interleaved messages + tool_calls for a session. + + This is the Full Parity Model query — returns the same row shape as the + desktop GUI's SESSION_LOAD_SQL, enabling consistent behavior across: + - Desktop GUI (Electron) + - Browser chat UI + - API consumers + + Args: + session_id: The session to load + include_compacted: If False, filter out compacted message rows + (tool_call rows are never compacted) + + Returns: + List of dicts with keys: + - row_type: 'message' | 'tool_call' + - row_id: str (message id or tool_call id) + - seq: int (ordering key) + - role: str + - content, type, agent_name, model_name, timestamp, ... + - tool_name, args_json, result_json, status, ... (for tool_call rows) + - parent_message_seq: int | None (for tool_call rows) + """ + db = get_db() + + if include_compacted: + # Use the standard parity SQL + cursor = await db.execute( + _SESSION_HISTORY_PARITY_SQL, + (session_id, session_id), + ) + else: + cursor = await db.execute( + _SESSION_HISTORY_PARITY_NO_COMPACTED_SQL, + (session_id, session_id), + ) + + rows = await cursor.fetchall() + return [dict(r) for r in rows] + + +async def get_session_tool_calls(session_id: str) -> list[dict[str, Any]]: + """Return all tool_call rows for a session, ordered by seq. + + Useful for diagnostics and validation. + """ + db = get_db() + cursor = await db.execute( + """ + SELECT id, session_id, parent_message_seq, seq, tool_name, + args_json, result_json, status, duration_ms, error_text, + agent_name, model_name, timestamp + FROM tool_calls + WHERE session_id = ? + ORDER BY seq + """, + (session_id,), + ) + rows = await cursor.fetchall() + return [dict(r) for r in rows] diff --git a/code_puppy/api/db/seeder.py b/code_puppy/api/db/seeder.py new file mode 100644 index 000000000..d2d31ffde --- /dev/null +++ b/code_puppy/api/db/seeder.py @@ -0,0 +1,17 @@ +"""Retired legacy session seeder. + +The application now treats SQLite + JSON session files as the only supported +runtime storage format. The old pickle-to-SQLite migration path has been +removed; this module remains only as a compatibility import target. +""" + +from __future__ import annotations + +import logging + +logger = logging.getLogger(__name__) + + +async def seed_from_pkl_dirs() -> None: + """No-op compatibility shim for the retired pickle migration path.""" + logger.debug("Legacy session seeder is retired; nothing to migrate") diff --git a/code_puppy/api/db/sql/session_history_parity.sql b/code_puppy/api/db/sql/session_history_parity.sql new file mode 100644 index 000000000..84c85295f --- /dev/null +++ b/code_puppy/api/db/sql/session_history_parity.sql @@ -0,0 +1,59 @@ +SELECT + 'message' AS row_type, + CAST(m.id AS TEXT) AS row_id, + m.seq, + m.role, + m.content, + m.type, + m.agent_name, + m.model_name, + m.timestamp, + m.thinking, + m.attachments_json, + m.clean_content, + m.system_message_type, + m.system_message_path, + m.token_count, + m.compacted, + m.compaction_log_id, + NULL AS tool_name, + NULL AS args_json, + NULL AS result_json, + NULL AS status, + NULL AS duration_ms, + NULL AS error_text, + NULL AS parent_message_seq +FROM messages m +WHERE m.session_id = ? + +UNION ALL + +SELECT + 'tool_call' AS row_type, + tc.id AS row_id, + tc.seq, + 'tool' AS role, + NULL AS content, + NULL AS type, + tc.agent_name, + tc.model_name, + CAST(tc.timestamp AS TEXT) AS timestamp, + NULL AS thinking, + NULL AS attachments_json, + NULL AS clean_content, + NULL AS system_message_type, + NULL AS system_message_path, + 0 AS token_count, + 0 AS compacted, + NULL AS compaction_log_id, + tc.tool_name, + tc.args_json, + tc.result_json, + tc.status, + tc.duration_ms, + tc.error_text, + tc.parent_message_seq +FROM tool_calls tc +WHERE tc.session_id = ? + +ORDER BY seq; diff --git a/code_puppy/api/db/sql/session_history_parity_no_compacted.sql b/code_puppy/api/db/sql/session_history_parity_no_compacted.sql new file mode 100644 index 000000000..fad44d5f5 --- /dev/null +++ b/code_puppy/api/db/sql/session_history_parity_no_compacted.sql @@ -0,0 +1,59 @@ +SELECT + 'message' AS row_type, + CAST(m.id AS TEXT) AS row_id, + m.seq, + m.role, + m.content, + m.type, + m.agent_name, + m.model_name, + m.timestamp, + m.thinking, + m.attachments_json, + m.clean_content, + m.system_message_type, + m.system_message_path, + m.token_count, + m.compacted, + m.compaction_log_id, + NULL AS tool_name, + NULL AS args_json, + NULL AS result_json, + NULL AS status, + NULL AS duration_ms, + NULL AS error_text, + NULL AS parent_message_seq +FROM messages m +WHERE m.session_id = ? AND m.compacted = 0 + +UNION ALL + +SELECT + 'tool_call' AS row_type, + tc.id AS row_id, + tc.seq, + 'tool' AS role, + NULL AS content, + NULL AS type, + tc.agent_name, + tc.model_name, + CAST(tc.timestamp AS TEXT) AS timestamp, + NULL AS thinking, + NULL AS attachments_json, + NULL AS clean_content, + NULL AS system_message_type, + NULL AS system_message_path, + 0 AS token_count, + 0 AS compacted, + NULL AS compaction_log_id, + tc.tool_name, + tc.args_json, + tc.result_json, + tc.status, + tc.duration_ms, + tc.error_text, + tc.parent_message_seq +FROM tool_calls tc +WHERE tc.session_id = ? + +ORDER BY seq; diff --git a/code_puppy/api/db/sql_loader.py b/code_puppy/api/db/sql_loader.py new file mode 100644 index 000000000..47ac4a1a7 --- /dev/null +++ b/code_puppy/api/db/sql_loader.py @@ -0,0 +1,22 @@ +"""Helpers for loading external SQL query files. + +Keeps complex SQL out of Python source while preserving parameterized execution. +""" + +from __future__ import annotations + +from functools import lru_cache +from pathlib import Path + +_SQL_DIR = Path(__file__).with_name("sql") + + +@lru_cache(maxsize=64) +def load_sql(name: str) -> str: + """Load an SQL file from ``code_puppy/api/db/sql``. + + Args: + name: Filename like ``session_history_parity.sql``. + """ + sql_path = _SQL_DIR / name + return sql_path.read_text(encoding="utf-8") diff --git a/code_puppy/api/error_parser.py b/code_puppy/api/error_parser.py new file mode 100644 index 000000000..ffd60fafe --- /dev/null +++ b/code_puppy/api/error_parser.py @@ -0,0 +1,239 @@ +"""Error parsing utilities for API errors. + +Provides human-readable error messages and actionable guidance for common API errors. +""" + +import logging +import re +from typing import Any, Dict + +logger = logging.getLogger(__name__) + + +def parse_api_error(error: Exception) -> Dict[str, Any]: + """Parse API errors into user-friendly messages with actionable guidance. + + Debugging: this function logs the chosen classification at DEBUG level + to help trace WS error propagation to the UI. + + Args: + error: The exception that was raised + + Returns: + Dictionary with: + - error_type: Category of error (quota_exceeded, rate_limit, auth_error, network_error, unknown) + - user_message: Human-readable message for the user + - technical_details: Technical error information + - action_required: Specific action the user should take (optional) + - original_error: String representation of the original error + """ + error_str = str(error) + error_type_name = type(error).__name__ + + logger.debug( + "parse_api_error: incoming type=%s msg=%r", + error_type_name, + error_str[:400], + ) + + # Check for quota exceeded errors (HTTP 429 or explicit quota messages) + if _is_quota_exceeded_error(error_str): + parsed = { + "error_type": "quota_exceeded", + "user_message": "Your API quota has been exceeded for the current model. Please switch to a different model and click continue to retry your request.", + "technical_details": _extract_quota_details(error_str), + "action_required": "switch_model", + "original_error": error_str, + } + logger.debug("parse_api_error: classified as %s", parsed["error_type"]) + return parsed + + # Check for rate limit errors + if _is_rate_limit_error(error_str): + parsed = { + "error_type": "rate_limit", + "user_message": "You're making requests too quickly. Please wait a moment and try again, or switch to a different model.", + "technical_details": error_str, + "action_required": "wait_or_switch_model", + "original_error": error_str, + } + logger.debug("parse_api_error: classified as %s", parsed["error_type"]) + return parsed + + # Check for authentication/permission errors + if _is_auth_error(error_str): + return { + "error_type": "auth_error", + "user_message": "Authentication failed. Please check your credentials or API key configuration.", + "technical_details": error_str, + "action_required": "check_credentials", + "original_error": error_str, + } + + # Check for network/timeout errors + if _is_network_error(error_str, error_type_name): + return { + "error_type": "network_error", + "user_message": "Network connection failed. Please check your internet connection and try again.", + "technical_details": error_str, + "action_required": "retry", + "original_error": error_str, + } + + # Check for model not found errors + if _is_model_not_found_error(error_str): + return { + "error_type": "model_not_found", + "user_message": "The requested model is not available. Please switch to a different model.", + "technical_details": error_str, + "action_required": "switch_model", + "original_error": error_str, + } + + # Check for Claude extended thinking temperature errors + if _is_claude_temperature_thinking_error(error_str): + parsed = { + "error_type": "claude_temperature_error", + "user_message": "Claude's extended thinking mode requires temperature to be set to 1.0. Please adjust your model settings or disable extended thinking.", + "technical_details": error_str, + "action_required": "adjust_temperature", + "original_error": error_str, + } + logger.debug("parse_api_error: classified as %s", parsed["error_type"]) + return parsed + + # Generic error fallback + parsed = { + "error_type": "unknown", + "user_message": f"An unexpected error occurred: {error_str}", + "technical_details": error_str, + "action_required": None, + "original_error": error_str, + } + logger.debug("parse_api_error: classified as %s", parsed["error_type"]) + return parsed + + +def _is_quota_exceeded_error(error_str: str) -> bool: + """Check if error is a quota exceeded error.""" + quota_patterns = [ + r"quota\s+exceeded", + r"HTTP\s+429", + r"aiplatform\.googleapis\.com/online_prediction_requests_per_base_model", + r"resource_exhausted", + r"too\s+many\s+requests.*quota", + ] + + for pattern in quota_patterns: + if re.search(pattern, error_str, re.IGNORECASE): + return True + return False + + +def _is_rate_limit_error(error_str: str) -> bool: + """Check if error is a rate limit error.""" + rate_limit_patterns = [ + r"rate\s+limit", + r"too\s+many\s+requests", + r"throttle", + r"requests.*per.*second", + r"requests.*per.*minute", + ] + + for pattern in rate_limit_patterns: + if re.search(pattern, error_str, re.IGNORECASE): + # Don't double-match quota errors + if not _is_quota_exceeded_error(error_str): + return True + return False + + +def _is_auth_error(error_str: str) -> bool: + """Check if error is an authentication error.""" + auth_patterns = [ + r"authentication\s+failed", + r"unauthorized", + r"HTTP\s+401", + r"HTTP\s+403", + r"permission\s+denied", + r"invalid.*api.*key", + r"invalid.*credentials", + r"access\s+denied", + ] + + for pattern in auth_patterns: + if re.search(pattern, error_str, re.IGNORECASE): + return True + return False + + +def _is_network_error(error_str: str, error_type_name: str) -> bool: + """Check if error is a network/connectivity error.""" + # Check error type names + network_error_types = [ + "TimeoutError", + "ConnectionError", + "ConnectTimeout", + "ReadTimeout", + "HTTPError", + ] + + if any(err_type in error_type_name for err_type in network_error_types): + return True + + # Check error message patterns + network_patterns = [ + r"timeout", + r"connection\s+refused", + r"connection\s+reset", + r"connection\s+failed", + r"network\s+error", + r"unable\s+to\s+connect", + r"host\s+unreachable", + ] + + for pattern in network_patterns: + if re.search(pattern, error_str, re.IGNORECASE): + return True + return False + + +def _is_model_not_found_error(error_str: str) -> bool: + """Check if error is a model not found error.""" + # Use simple patterns that match anywhere in the string + if re.search(r"model.*not\s+found", error_str, re.IGNORECASE): + return True + if re.search(r"invalid\s+model", error_str, re.IGNORECASE): + return True + if re.search(r"unknown\s+model", error_str, re.IGNORECASE): + return True + if re.search(r"model.*does\s+not\s+exist", error_str, re.IGNORECASE): + return True + if re.search(r"HTTP\s+404.*model", error_str, re.IGNORECASE): + return True + + return False + + +def _is_claude_temperature_thinking_error(error_str: str) -> bool: + """Check if error is Claude's temperature/thinking configuration error.""" + return ( + "temperature" in error_str.lower() + and "thinking" in error_str.lower() + and ("may only be set to 1" in error_str or "must be 1" in error_str) + ) + + +def _extract_quota_details(error_str: str) -> str: + """Extract specific quota details from error message.""" + # Try to extract the specific quota metric if present + match = re.search(r"quota\s+exceeded\s+for\s+([\w\./]+)", error_str, re.IGNORECASE) + if match: + return f"Quota exceeded for: {match.group(1)}" + + # Try to extract model name if present + match = re.search(r"base\s+model[:\s]+([\w\-]+)", error_str, re.IGNORECASE) + if match: + return f"Quota exceeded for base model: {match.group(1)}" + + return error_str 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/permission_plugin.py b/code_puppy/api/permission_plugin.py new file mode 100644 index 000000000..91a2e9f3d --- /dev/null +++ b/code_puppy/api/permission_plugin.py @@ -0,0 +1,222 @@ +""" +Permission Plugin for Tool Call Execution + +This plugin registers callbacks that intercept tool executions +and request user permission via the WebSocket API. + +It handles: +- All tool calls via pre_tool_call callback +- Shell commands via run_shell_command callback (for backward compatibility) + +Permission is automatically granted in yolo mode. +For WebSocket mode (CLI web interface), permission is requested via UI. +For terminal mode, permission is bypassed to maintain backward compatibility. +""" + +import logging +from contextvars import ContextVar +from typing import Any, Dict, Optional, Tuple + +logger = logging.getLogger(__name__) + +# Task-scoped WebSocket context — each asyncio task (i.e. each WS connection) +# gets its own isolated value via Python's ContextVar mechanism. +_ws_context: ContextVar[Optional[Tuple[Any, str]]] = ContextVar( + "ws_permission_context", default=None +) + +# Task-scoped flag used by chat websocket streaming to suppress duplicate +# tool_call/tool_result events from frontend_emitter pre/post_tool_call hooks. +_suppress_emitter_tool_events: ContextVar[bool] = ContextVar( + "ws_suppress_emitter_tool_events", default=False +) + + +def set_websocket_context(websocket: Any, session_id: str) -> None: + """Set the current WebSocket context for permission requests (task-scoped).""" + _ws_context.set((websocket, session_id)) + logger.debug("[Permission] WebSocket context set for session %s", session_id) + + +def get_websocket_context() -> Optional[Tuple[Any, str]]: + """Get the current WebSocket context. Returns ``(websocket, session_id)`` or ``None``.""" + return _ws_context.get() + + +def clear_websocket_context() -> None: + """Clear the WebSocket context for the current task.""" + _ws_context.set(None) + logger.debug("[Permission] WebSocket context cleared") + + +def set_suppress_emitter_tool_events(suppress: bool) -> None: + """Enable/disable suppression of frontend_emitter tool lifecycle events.""" + _suppress_emitter_tool_events.set(bool(suppress)) + + +def get_suppress_emitter_tool_events() -> bool: + """Return True when frontend_emitter tool_call_start/complete should be skipped.""" + return _suppress_emitter_tool_events.get() + + +async def pre_tool_call_permission( + tool_name: str, tool_args: dict, context: Any = None +) -> Optional[Dict[str, Any]]: + """ + Permission callback for all tool executions via pre_tool_call. + + Args: + tool_name: Name of the tool being called + tool_args: Arguments for the tool + context: Execution context + + Returns: + None to allow the tool, or a dict with error info to block it + """ + # Check if we have a WebSocket context (CLI web mode) + ctx = get_websocket_context() + if ctx is None: + # Terminal mode or no WebSocket - allow tool + logger.debug("[Permission] No WebSocket context, allowing tool: %s", tool_name) + return None + + ws, sid = ctx + + # WebSocket mode - request permission + logger.info("[Permission] Requesting permission for tool: %s", tool_name) + + try: + from code_puppy.api.permissions import request_permission + + # Build human-readable description using formatters + from code_puppy.api.tool_formatters import format_tool_call + + try: + formatted_args = format_tool_call(tool_name, tool_args) + description = formatted_args + except Exception as e: + # Fallback to raw args if formatting fails + logger.warning("[Permission] Failed to format tool args: %s", e) + description = f"Tool: {tool_name}" + if tool_args: + arg_preview = str(tool_args)[:200] + if len(str(tool_args)) > 200: + arg_preview += "..." + description = f"{description}\nArguments: {arg_preview}" + + approved = await request_permission( + websocket=ws, + session_id=sid, + request_type="tool_call", + title=f"Execute Tool: {tool_name}", + description=description, + details={ + "tool_name": tool_name, + "tool_args": tool_args, + }, + timeout=300, # 5 minute timeout + ) + + if not approved: + logger.info("[Permission] ❌ Tool DENIED: %s", tool_name) + # Return error dict to block the tool execution + return { + "error": "Permission denied by user", + "blocked": True, + "tool_name": tool_name, + } + + logger.info("[Permission] Tool APPROVED: %s", tool_name) + return None # Allow the tool + + except Exception as e: + logger.error("[Permission] Error requesting permission: %s", e) + return { + "error": "Permission system error", + "blocked": True, + "tool_name": tool_name, + } + + +async def shell_command_permission( + context: Any, command: str, cwd: str, timeout: int +) -> Optional[Dict[str, Any]]: + """ + Permission callback for shell command execution (legacy/backward compatibility). + + This is kept for backward compatibility but the main permission handling + now happens via pre_tool_call_permission. + + Args: + context: The execution context + command: The shell command to execute + cwd: Working directory + timeout: Command timeout + + Returns: + None to allow the command, or a dict with 'blocked': True to deny + """ + # Check if we have a WebSocket context (CLI web mode) + ctx = get_websocket_context() + if ctx is None: + # Terminal mode or no WebSocket - allow command + logger.debug("[Permission] No WebSocket context, allowing command: %s", command) + return None + + ws, sid = ctx + + # WebSocket mode - request permission + logger.info("[Permission] Requesting permission for shell command: %s", command) + + try: + from code_puppy.api.permissions import request_permission + + approved = await request_permission( + websocket=ws, + session_id=sid, + request_type="shell_command", + title="Execute Shell Command", + description=f"Run: {command}", + details={ + "command": command, + "cwd": cwd or "current directory", + "timeout": timeout, + }, + timeout=300, # 5 minute timeout + ) + + if not approved: + logger.info("[Permission] Command DENIED: %s", command) + return { + "blocked": True, + "error": "Permission denied by user", + "reasoning": f"User denied execution of command: {command}", + } + + logger.info("[Permission] Command APPROVED: %s", command) + return None # Allow the command + + except Exception as e: + logger.error("[Permission] Error requesting permission for command: %s", e) + return { + "blocked": True, + "error": "Permission system error", + "reasoning": f"Permission request failed for command: {command}", + } + + +def register_permission_callbacks(): + """Register the permission callbacks with code-puppy.""" + from code_puppy.callbacks import register_callback + + # Register for all tool calls + register_callback("pre_tool_call", pre_tool_call_permission) + logger.info("[Permission Plugin] Registered pre_tool_call permission callback") + + # Keep shell command callback for backward compatibility + register_callback("run_shell_command", shell_command_permission) + logger.info("[Permission Plugin] Registered shell command permission callback") + + +# Auto-register when imported +register_permission_callbacks() diff --git a/code_puppy/api/permissions.py b/code_puppy/api/permissions.py new file mode 100644 index 000000000..7059f6a5e --- /dev/null +++ b/code_puppy/api/permissions.py @@ -0,0 +1,144 @@ +""" +Permission System for Tool Call Execution. + +WebSocket-side permission requests are tracked per request_id and bound to the +originating session_id to prevent cross-session response confusion. +""" + +import asyncio +import logging +import uuid +from dataclasses import dataclass +from typing import Any, Dict, Optional + +from fastapi import WebSocket + +from code_puppy.api.ws.schemas import ServerPermissionRequest + +logger = logging.getLogger(__name__) + + +@dataclass(slots=True) +class PendingPermissionRequest: + future: asyncio.Future + session_id: str + + +# Global dictionary to track pending permission requests +permission_futures: Dict[str, PendingPermissionRequest] = {} + + +async def request_permission( + websocket: Optional[WebSocket], + session_id: str, + request_type: str, + title: str, + description: str, + details: Dict[str, Any], + timeout: int = 300, +) -> bool: + """Request permission from user via WebSocket.""" + # Check if yolo mode is enabled (auto-approve everything). + # Use the canonical config helper instead of reading CONFIG_FILE directly: + # the config backend is not guaranteed to be JSON, supports values like + # yes/on/1, and defaults YOLO to enabled when unset. + try: + from code_puppy.config import get_yolo_mode + + if get_yolo_mode(): + logger.info( + "[Permission] YOLO mode enabled, auto-approving %s", + request_type, + ) + return True + except Exception as exc: + logger.warning( + "[Permission] Could not read YOLO mode; falling back to UI prompt for %s: %s", + request_type, + exc, + ) + + if websocket is None: + logger.warning("[Permission] No WebSocket connection, denying %s", request_type) + return False + + request_id = str(uuid.uuid4()) + future: asyncio.Future = asyncio.get_running_loop().create_future() + permission_futures[request_id] = PendingPermissionRequest( + future=future, + session_id=session_id, + ) + + try: + await websocket.send_json( + ServerPermissionRequest( + request_id=request_id, + permission_type=request_type, + title=title, + description=description, + details=details, + session_id=session_id, + timeout_seconds=timeout, + ).model_dump(exclude_none=True) + ) + logger.info("[Permission] Sent request %s for %s", request_id, request_type) + except Exception as e: + permission_futures.pop(request_id, None) + logger.error("[Permission] Failed to send request: %s", e) + return False + + try: + approved = await asyncio.wait_for(future, timeout=timeout) + logger.info( + "[Permission] Got response for %s: %s", + request_id, + "APPROVED" if approved else "DENIED", + ) + return bool(approved) + except asyncio.TimeoutError: + logger.warning( + "[Permission] Request %s timed out after %ss", request_id, timeout + ) + return False + except Exception as e: + logger.error("[Permission] Error waiting for permission: %s", e) + return False + finally: + permission_futures.pop(request_id, None) + + +def handle_permission_response( + request_id: str, + approved: bool, + *, + session_id: Optional[str] = None, +) -> bool: + """Handle a permission response from the client. + + If ``session_id`` is provided, response handling is restricted to matching + pending requests from the same session. + """ + pending = permission_futures.get(request_id) + if pending is None: + logger.warning( + "[Permission] Received response for unknown/expired request: %s", + request_id, + ) + return False + + if session_id is not None and pending.session_id != session_id: + logger.warning( + "[Permission] Session mismatch for request %s (expected %s, got %s)", + request_id, + pending.session_id, + session_id, + ) + return False + + if pending.future.done(): + logger.warning("[Permission] Request already resolved: %s", request_id) + return False + + pending.future.set_result(bool(approved)) + logger.info("[Permission] Handled response for %s: %s", request_id, approved) + return True 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/routers/sessions.py b/code_puppy/api/routers/sessions.py new file mode 100644 index 000000000..c1404ad79 --- /dev/null +++ b/code_puppy/api/routers/sessions.py @@ -0,0 +1,187 @@ +"""Sessions API endpoints for retrieving subagent session data.""" + +import asyncio +import json +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Any, Dict, List, Optional + +from fastapi import APIRouter, HTTPException +from pydantic import BaseModel + +from code_puppy.session_storage import load_session + +_executor = ThreadPoolExecutor(max_workers=2) +FILE_IO_TIMEOUT = 10.0 + +router = APIRouter() + + +class SessionInfo(BaseModel): + """Session metadata information.""" + + session_id: str + agent_name: Optional[str] = None + initial_prompt: Optional[str] = None + created_at: Optional[str] = None + last_updated: Optional[str] = None + message_count: int = 0 + + +class MessageContent(BaseModel): + """Message content with role and optional timestamp.""" + + role: str + content: Any + timestamp: Optional[str] = None + + +class SessionDetail(SessionInfo): + """Session info with full message history.""" + + messages: List[Dict[str, Any]] = [] + + +def _get_sessions_dir() -> Path: + from code_puppy.config import DATA_DIR + + return Path(DATA_DIR) / "subagent_sessions" + + +def _serialize_message(msg: Any) -> Dict[str, Any]: + def _serialize_obj(obj: Any) -> Any: + if hasattr(obj, "model_dump"): + return obj.model_dump(mode="json") + if isinstance(obj, dict): + return {k: _serialize_obj(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_serialize_obj(item) for item in obj] + if isinstance(obj, (str, int, float, bool, type(None))): + return obj + if hasattr(obj, "__dict__"): + return {k: _serialize_obj(v) for k, v in obj.__dict__.items()} + return str(obj) + + if isinstance(msg, dict) and "msg" in msg: + actual_msg = msg["msg"] + return { + "msg": _serialize_obj(actual_msg), + "agent": msg.get("agent"), + "model": msg.get("model"), + "ts": msg.get("ts"), + } + + result = _serialize_obj(msg) + if not isinstance(result, dict): + return {"content": str(result)} + return result + + +def _load_json_sync(file_path: Path) -> dict: + with open(file_path, "r", encoding="utf-8") as f: + return json.load(f) + + +def _load_session_history_sync(file_path: Path) -> Any: + return load_session(file_path.stem, file_path.parent) + + +@router.get("/") +async def list_sessions() -> List[SessionInfo]: + sessions_dir = _get_sessions_dir() + if not sessions_dir.exists(): + return [] + + loop = asyncio.get_running_loop() + sessions: list[SessionInfo] = [] + + for txt_file in sessions_dir.glob("*.txt"): + session_id = txt_file.stem + try: + metadata = await asyncio.wait_for( + loop.run_in_executor(_executor, _load_json_sync, txt_file), + timeout=FILE_IO_TIMEOUT, + ) + sessions.append( + SessionInfo( + session_id=session_id, + agent_name=metadata.get("agent_name"), + initial_prompt=metadata.get("initial_prompt"), + created_at=metadata.get("created_at"), + last_updated=metadata.get("last_updated"), + message_count=metadata.get("message_count", 0), + ) + ) + except asyncio.TimeoutError: + sessions.append(SessionInfo(session_id=session_id)) + except Exception: + sessions.append(SessionInfo(session_id=session_id)) + + return sessions + + +@router.get("/{session_id}") +async def get_session(session_id: str) -> SessionInfo: + sessions_dir = _get_sessions_dir() + txt_file = sessions_dir / f"{session_id}.txt" + + if not txt_file.exists(): + raise HTTPException(404, f"Session '{session_id}' not found") + + loop = asyncio.get_running_loop() + try: + metadata = await asyncio.wait_for( + loop.run_in_executor(_executor, _load_json_sync, txt_file), + timeout=FILE_IO_TIMEOUT, + ) + except asyncio.TimeoutError: + raise HTTPException(504, f"Timeout reading session '{session_id}'") from None + + return SessionInfo( + session_id=session_id, + agent_name=metadata.get("agent_name"), + initial_prompt=metadata.get("initial_prompt"), + created_at=metadata.get("created_at"), + last_updated=metadata.get("last_updated"), + message_count=metadata.get("message_count", 0), + ) + + +@router.get("/{session_id}/messages") +async def get_session_messages(session_id: str) -> List[Dict[str, Any]]: + sessions_dir = _get_sessions_dir() + session_file = sessions_dir / f"{session_id}.pkl" + + if not session_file.exists(): + raise HTTPException(404, f"Session '{session_id}' messages not found") + + loop = asyncio.get_running_loop() + try: + messages = await asyncio.wait_for( + loop.run_in_executor(_executor, _load_session_history_sync, session_file), + timeout=FILE_IO_TIMEOUT, + ) + return [_serialize_message(msg) for msg in messages] + except asyncio.TimeoutError: + raise HTTPException( + 504, f"Timeout loading session '{session_id}' messages" + ) from None + except Exception as e: + raise HTTPException(500, f"Error loading session messages: {e}") from e + + +@router.delete("/{session_id}") +async def delete_session(session_id: str) -> Dict[str, str]: + sessions_dir = _get_sessions_dir() + txt_file = sessions_dir / f"{session_id}.txt" + session_file = sessions_dir / f"{session_id}.pkl" + + if not txt_file.exists() and not session_file.exists(): + raise HTTPException(404, f"Session '{session_id}' not found") + + if txt_file.exists(): + txt_file.unlink() + if session_file.exists(): + session_file.unlink() + + return {"message": f"Session '{session_id}' deleted"} diff --git a/code_puppy/api/routers/ws_sessions.py b/code_puppy/api/routers/ws_sessions.py new file mode 100644 index 000000000..dff3a5c84 --- /dev/null +++ b/code_puppy/api/routers/ws_sessions.py @@ -0,0 +1,290 @@ +"""WebSocket session management router. + +Handles CRUD operations for chat sessions. +Now exclusively uses SQLite as the source of truth. +""" + +import json +import logging +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, List, Optional + +from fastapi import APIRouter, HTTPException, Query + +from code_puppy.api.db.queries import ( + get_active_messages, + get_session_history_parity, + get_session_metadata, + get_session_tool_calls, + soft_delete_session, + update_session_meta_fields, +) +from code_puppy.config import get_ws_sessions_dir + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +def _get_ws_sessions_dir() -> Path: + """Get the WebSocket sessions directory.""" + return get_ws_sessions_dir() + + +def _validate_session_name(session_name: str, ws_dir: Path) -> str: + """Validate session name prevents path traversal.""" + if not session_name or ".." in session_name or "/" in session_name: + raise HTTPException(400, "Invalid session name") + return session_name + + +@router.get("/{session_name}/messages") +async def get_ws_session_messages( + session_name: str, + include_tool_calls: Optional[bool] = Query( + default=None, + description="Include tool_call rows interleaved with messages (Full Parity Model). " + "When true, returns unified row format with row_type discriminator.", + ), + include_compacted: bool = Query( + default=True, + description="Include compacted (summarized) messages. Only applies when include_tool_calls=true.", + ), +) -> List[Dict[str, Any]]: + """Get the full message history for a WebSocket session from SQLite. + + Args: + session_name: The session name / session_id + include_tool_calls: Feature flag for Full Parity Model. + - None/False: Legacy behavior, returns only message rows + - True: Returns interleaved message + tool_call rows with row_type discriminator + include_compacted: When using parity mode, whether to include compacted messages + + Returns: + List of serialized message dictionaries. + + When include_tool_calls=true, each row has: + - row_type: 'message' | 'tool_call' + - row_id: unique identifier + - seq: ordering key + - Plus type-specific fields (content for messages, result_json for tool_calls, etc.) + """ + _validate_session_name(session_name, _get_ws_sessions_dir()) + + try: + # ────────────────────────────────────────────────────────────────── + # Full Parity Model: interleaved messages + tool_calls + # ────────────────────────────────────────────────────────────────── + if include_tool_calls: + rows = await get_session_history_parity( + session_name, + include_compacted=include_compacted, + ) + + if rows: + logger.debug( + "get_ws_session_messages (parity): session=%s total_rows=%d " + "messages=%d tool_calls=%d", + session_name, + len(rows), + sum(1 for r in rows if r.get("row_type") == "message"), + sum(1 for r in rows if r.get("row_type") == "tool_call"), + ) + return rows + + # Check if session exists but has no data + meta = await get_session_metadata(session_name) + if meta: + return [] + + raise HTTPException(404, f"WebSocket session '{session_name}' not found") + + # ────────────────────────────────────────────────────────────────── + # Legacy behavior: message rows only (backward compatible) + # ────────────────────────────────────────────────────────────────── + rows = await get_active_messages(session_name) + serialized_rows: List[Dict[str, Any]] = [] + + for row in rows: + base_row = { + "role": row["role"], + "content": row["content"], + "type": row["type"], + "agent_name": row["agent_name"], + "model_name": row["model_name"], + "timestamp": row["timestamp"], + "thinking": row["thinking"], + "clean_content": row["clean_content"], + "seq": row["seq"], + } + + if row.get("pydantic_json"): + serialized_rows.append(base_row) + continue + + if row.get("type") == "error": + payload: Dict[str, Any] = {} + raw_payload = row.get("attachments_json") + if raw_payload: + try: + payload = json.loads(raw_payload) + except (TypeError, ValueError): + logger.warning( + "get_ws_session_messages: invalid error payload JSON for %s seq=%s", + session_name, + row.get("seq"), + ) + + serialized_rows.append( + { + **base_row, + "type": "error", + "error": payload.get("error", row["content"]), + "error_type": payload.get("error_type", "unknown"), + "technical_details": payload.get("technical_details", ""), + "action_required": payload.get("action_required"), + "session_id": payload.get("session_id", session_name), + } + ) + continue + + if row.get("type") == "system": + serialized_rows.append( + { + **base_row, + "system_message_type": row.get("system_message_type"), + "system_message_path": row.get("system_message_path"), + } + ) + continue + + serialized_rows.append(base_row) + + if serialized_rows: + return serialized_rows + + # If no messages found, check if session exists at all + meta = await get_session_metadata(session_name) + if meta: + # Session exists but has no messages yet - return empty list + return [] + + raise HTTPException(404, f"WebSocket session '{session_name}' not found") + + except HTTPException: + raise + except Exception as e: + logger.error( + "get_ws_session_messages: SQLite error for %s: %s", session_name, e + ) + raise HTTPException(503, "Service unavailable: database error") + + +@router.get("/{session_name}/tool-calls") +async def get_ws_session_tool_calls(session_name: str) -> List[Dict[str, Any]]: + """Get all tool calls for a session. + + This is a diagnostic/debugging endpoint that returns just the tool_calls + table rows for a session, useful for verifying persistence. + + Args: + session_name: The session name / session_id + + Returns: + List of tool_call rows with id, tool_name, args_json, result_json, status, etc. + """ + _validate_session_name(session_name, _get_ws_sessions_dir()) + + try: + rows = await get_session_tool_calls(session_name) + + if rows: + logger.debug( + "get_ws_session_tool_calls: session=%s count=%d", + session_name, + len(rows), + ) + return rows + + # Check if session exists + meta = await get_session_metadata(session_name) + if meta: + return [] # Session exists but has no tool calls + + raise HTTPException(404, f"WebSocket session '{session_name}' not found") + + except HTTPException: + raise + except Exception as e: + logger.error( + "get_ws_session_tool_calls: SQLite error for %s: %s", session_name, e + ) + raise HTTPException(503, "Service unavailable: database error") + + +@router.delete("/{session_name}") +async def delete_ws_session(session_name: str) -> Dict[str, str]: + """Soft-delete a WebSocket session in SQLite.""" + _validate_session_name(session_name, _get_ws_sessions_dir()) + + # Soft delete in SQLite + try: + await soft_delete_session(session_name, datetime.now(timezone.utc).isoformat()) + # We don't check if it existed, soft_delete is idempotent-ish + # But we might want to know if we actually updated anything. + # For now, we assume success if no error. + except Exception as e: + logger.error("delete_ws_session: SQLite error for %s: %s", session_name, e) + raise HTTPException(503, "Service unavailable: database error") + + return {"message": f"WebSocket session '{session_name}' deleted"} + + +@router.patch("/{session_name}") +async def update_ws_session( + session_name: str, updates: Dict[str, Any] +) -> Dict[str, Any]: + """Update WebSocket session metadata in SQLite. + + Supports updating: title, pinned. + """ + _validate_session_name(session_name, _get_ws_sessions_dir()) + + # Get current metadata + current_meta = await get_session_metadata(session_name) + if not current_meta: + raise HTTPException(404, f"WebSocket session '{session_name}' not found") + + # Prepare updates + new_title = current_meta["title"] + new_pinned = bool(current_meta["pinned"]) + updated_at = datetime.now(timezone.utc).isoformat() + + if "title" in updates or "name" in updates: + val = updates.get("title") or updates.get("name") + if val and isinstance(val, str) and val.strip(): + new_title = val.strip() + + if "pinned" in updates: + if isinstance(updates["pinned"], bool): + new_pinned = updates["pinned"] + + # Write to SQLite + try: + await update_session_meta_fields( + session_name, + title=new_title, + pinned=new_pinned, + updated_at=updated_at, + ) + except Exception as e: + logger.error("update_ws_session: SQLite error for %s: %s", session_name, e) + raise HTTPException(503, "Service unavailable: database error") + + return { + "session_id": session_name, + "title": new_title, + "pinned": new_pinned, + "updated_at": updated_at, + } diff --git a/code_puppy/api/session_cache.py b/code_puppy/api/session_cache.py new file mode 100644 index 000000000..bbcd4e858 --- /dev/null +++ b/code_puppy/api/session_cache.py @@ -0,0 +1,393 @@ +"""Server-side session caching for fast session loading. + +This module provides an LRU cache for deserialized session data, +dramatically reducing load times for frequently accessed sessions. +""" + +import asyncio +import logging +import time +from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +logger = logging.getLogger(__name__) + +# Configuration +MAX_CACHED_SESSIONS = 50 # Maximum number of sessions to cache +CACHE_TTL_SECONDS = 3600 # 1 hour TTL +MAX_WORKERS = 4 + +# Thread pool for async file I/O +_executor = ThreadPoolExecutor(max_workers=MAX_WORKERS) + + +class SessionCacheEntry: + """A cached session entry with metadata.""" + + __slots__ = [ + "messages", + "json_messages", + "metadata", + "loaded_at", + "last_accessed", + "file_mtime", + ] + + def __init__( + self, + messages: List[Any], + json_messages: List[Dict[str, Any]], + metadata: Dict[str, Any], + file_mtime: float, + ): + self.messages = messages + self.json_messages = json_messages # Pre-serialized for fast API response + self.metadata = metadata + self.loaded_at = time.time() + self.last_accessed = time.time() + self.file_mtime = file_mtime + + @property + def age_seconds(self) -> float: + return time.time() - self.loaded_at + + @property + def is_expired(self) -> bool: + return self.age_seconds > CACHE_TTL_SECONDS + + def touch(self): + """Update last accessed time.""" + self.last_accessed = time.time() + + +class SessionCache: + """LRU cache for session data with async support. + + Features: + - LRU eviction when cache is full + - TTL-based expiration + - File modification tracking for cache invalidation + - Pre-serialized JSON for fast API responses + - Thread-safe async operations + """ + + def __init__(self, max_size: int = MAX_CACHED_SESSIONS): + self._cache: OrderedDict[str, SessionCacheEntry] = OrderedDict() + self._max_size = max_size + self._lock = asyncio.Lock() + self._stats = {"hits": 0, "misses": 0, "evictions": 0, "expirations": 0} + + async def get( + self, session_id: str, session_file_path: Path + ) -> Optional[Tuple[List[Dict[str, Any]], Dict[str, Any]]]: + """Get cached session messages (pre-serialized JSON). + + Returns: + Tuple of (json_messages, metadata) if cached and valid, None otherwise + """ + async with self._lock: + entry = self._cache.get(session_id) + + if entry is None: + self._stats["misses"] += 1 + return None + + # Check TTL + if entry.is_expired: + self._stats["expirations"] += 1 + del self._cache[session_id] + logger.debug("[Cache] Session %s expired", session_id) + return None + + # Check if file was modified + try: + current_mtime = session_file_path.stat().st_mtime + if current_mtime > entry.file_mtime: + self._stats["expirations"] += 1 + del self._cache[session_id] + logger.debug( + f"[Cache] Session {session_id} file modified, invalidating cache" + ) + return None + except FileNotFoundError: + del self._cache[session_id] + return None + + # Cache hit! + self._stats["hits"] += 1 + entry.touch() + + # Move to end (most recently used) + self._cache.move_to_end(session_id) + + logger.debug( + f"[Cache] HIT for session {session_id} (age: {entry.age_seconds:.1f}s)" + ) + return entry.json_messages, entry.metadata + + async def put( + self, + session_id: str, + session_file_path: Path, + messages: List[Any], + json_messages: List[Dict[str, Any]], + metadata: Dict[str, Any], + ): + """Cache session data. + + Args: + session_id: Unique session identifier + session_file_path: Path to the session JSON file (for mtime tracking) + messages: Raw message objects + json_messages: Pre-serialized JSON messages + metadata: Session metadata + """ + async with self._lock: + # Evict if at capacity + while len(self._cache) >= self._max_size: + # Remove oldest (first) item + oldest_key = next(iter(self._cache)) + del self._cache[oldest_key] + self._stats["evictions"] += 1 + logger.debug("[Cache] Evicted session %s", oldest_key) + + # Get file mtime + try: + file_mtime = session_file_path.stat().st_mtime + except FileNotFoundError: + file_mtime = time.time() + + # Create entry + entry = SessionCacheEntry( + messages=messages, + json_messages=json_messages, + metadata=metadata, + file_mtime=file_mtime, + ) + + self._cache[session_id] = entry + logger.debug( + f"[Cache] Cached session {session_id} ({len(messages)} messages)" + ) + + async def invalidate(self, session_id: str): + """Remove a session from cache.""" + async with self._lock: + if session_id in self._cache: + del self._cache[session_id] + logger.debug("[Cache] Invalidated session %s", session_id) + + async def clear(self): + """Clear all cached sessions.""" + async with self._lock: + self._cache.clear() + logger.info("[Cache] All sessions cleared") + + @property + def stats(self) -> Dict[str, Any]: + """Get cache statistics.""" + hit_rate = 0.0 + total = self._stats["hits"] + self._stats["misses"] + if total > 0: + hit_rate = self._stats["hits"] / total * 100 + + return { + **self._stats, + "size": len(self._cache), + "max_size": self._max_size, + "hit_rate_percent": round(hit_rate, 1), + } + + +# Global cache instance +_session_cache: Optional[SessionCache] = None + + +def get_session_cache() -> SessionCache: + """Get the global session cache instance.""" + global _session_cache + if _session_cache is None: + _session_cache = SessionCache() + return _session_cache + + +def _load_session_sync(session_file_path: Path) -> List[Any]: + """Synchronous session loading (for executor).""" + from code_puppy.session_storage import load_session + + return load_session(session_file_path.stem, session_file_path.parent) + + +def _serialize_message(msg: Any) -> Dict[str, Any]: + """Serialize a pydantic-ai message to JSON-safe dict. + + Handles both wrapped and unwrapped message formats: + - Wrapped (WS sessions): {'msg': , 'agent': ..., 'model': ..., 'ts': ...} + - Unwrapped: or + + Preserves all message parts including ThinkingPart, ToolCallPart, TextPart, etc. + """ + + def _serialize_obj(obj: Any) -> Any: + """Recursively serialize an object to JSON-safe types.""" + # Pydantic v2 models with model_dump + if hasattr(obj, "model_dump"): + return obj.model_dump(mode="json") + + # Handle dicts recursively + if isinstance(obj, dict): + return {k: _serialize_obj(v) for k, v in obj.items()} + + # Handle lists/tuples recursively + if isinstance(obj, (list, tuple)): + return [_serialize_obj(item) for item in obj] + + # JSON-safe primitives + if isinstance(obj, (str, int, float, bool, type(None))): + return obj + + # pydantic-ai message objects (ModelRequest, ModelResponse, Part subclasses) + # They have __dict__ but are not Pydantic models + if hasattr(obj, "__dict__"): + result = {} + for k, v in obj.__dict__.items(): + result[k] = _serialize_obj(v) + return result + + # Fallback: convert to string + return str(obj) + + # Handle wrapped message format (used in WS sessions) + # {'msg': , 'agent': ..., 'model': ..., 'ts': ...} + if isinstance(msg, dict) and "msg" in msg: + actual_msg = msg["msg"] + return { + "msg": _serialize_obj(actual_msg), + "agent": msg.get("agent"), + "model": msg.get("model"), + "ts": msg.get("ts"), + } + + # Handle unwrapped messages + return _serialize_obj(msg) + + +def _serialize_messages_sync(messages: List[Any]) -> List[Dict[str, Any]]: + """Synchronous batch serialization (for executor). + + Serializes all messages to JSON-safe dicts. Runs in thread pool + to avoid blocking the event loop for large sessions. + """ + return [_serialize_message(msg) for msg in messages] + + +async def load_session_cached( + session_id: str, session_file_path: Path, timeout: float = 10.0 +) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: + """Load session messages with caching. + + This is the main entry point for loading session data. + It checks the cache first, and only loads from disk on cache miss. + + Args: + session_id: Unique session identifier + session_file_path: Path to the session JSON file + timeout: Timeout for file operations + + Returns: + Tuple of (json_messages, metadata) + + Raises: + FileNotFoundError: If session JSON file does not exist + asyncio.TimeoutError: If load times out + """ + cache = get_session_cache() + + # Try cache first + cached = await cache.get(session_id, session_file_path) + if cached is not None: + return cached + + # Cache miss - load from disk + if not session_file_path.exists(): + raise FileNotFoundError(f"Session file not found: {session_file_path}") + + loop = asyncio.get_running_loop() + + # Load session file in thread pool + start_time = time.time() + messages = await asyncio.wait_for( + loop.run_in_executor(_executor, _load_session_sync, session_file_path), + timeout=timeout, + ) + load_time = time.time() - start_time + + # Serialize to JSON in thread pool to avoid blocking event loop + json_messages = await loop.run_in_executor( + _executor, _serialize_messages_sync, messages + ) + + # Build metadata + metadata = { + "message_count": len(messages), + "load_time_ms": round(load_time * 1000, 1), + } + + # Cache for next time + await cache.put(session_id, session_file_path, messages, json_messages, metadata) + + logger.info( + f"[Cache] Loaded session {session_id} from disk in {load_time * 1000:.1f}ms" + ) + + return json_messages, metadata + + +async def invalidate_session_cache(session_id: str): + """Invalidate a specific session in the cache. + + Call this when a session is modified or deleted. + """ + cache = get_session_cache() + await cache.invalidate(session_id) + + +async def get_cache_stats() -> Dict[str, Any]: + """Get cache statistics for monitoring.""" + cache = get_session_cache() + return cache.stats + + +async def shutdown_executor() -> None: + """Shutdown the thread pool executor gracefully. + + This function ensures all pending tasks are completed before shutdown + and prevents resource leaks. Should be called during application shutdown. + + Thread-safe and idempotent - safe to call multiple times. + """ + global _executor + + if _executor is None: + logger.debug("[Executor] Already shut down or never initialized") + return + + try: + logger.info("[Executor] Shutting down thread pool executor...") + + # Shutdown gracefully - wait for current tasks to complete + # but don't cancel futures (they might be critical) + _executor.shutdown(wait=True, cancel_futures=False) + + logger.info("[Executor] Thread pool executor shutdown complete") + + # Clear the reference + _executor = None + + except Exception as e: + logger.error("[Executor] Error during shutdown: %s", e, exc_info=True) + # Still clear the reference even if shutdown had issues + _executor = None + raise diff --git a/code_puppy/api/session_context.py b/code_puppy/api/session_context.py new file mode 100644 index 000000000..f341f735e --- /dev/null +++ b/code_puppy/api/session_context.py @@ -0,0 +1,675 @@ +"""Multi-session isolation layer for Code Puppy. + +Provides per-session state management so that each browser tab / WebSocket +connection gets its own agent instance, model selection, working directory, +and message history — without touching any global singletons. + +Usage: + from code_puppy.api.session_context import session_manager + + ctx = session_manager.create_session("my-session-123") + ctx.agent.run_sync("Hello!") # uses the session’s own agent +""" + +from __future__ import annotations + +import asyncio +import logging +import re +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any, Dict, List, Optional, Set + +from code_puppy.agents.agent_manager import get_available_agents, load_agent +from code_puppy.agents.base_agent import BaseAgent +from code_puppy.api.db.message_utils import ( + extract_content, + extract_thinking, + get_message_timestamp, + get_role, + pydantic_json_for_message, +) +from code_puppy.config import get_global_model_name + +logger = logging.getLogger(__name__) + + +def _apply_session_model(agent: Any, model_name: Optional[str]) -> None: + """Best-effort session-model setter for migration compatibility. + + Some migrated/legacy agent classes may not yet implement + ``set_session_model``. In that case we set the internal override field + directly to avoid failing model switches at runtime. + """ + setter = getattr(agent, "set_session_model", None) + if callable(setter): + setter(model_name) + return + + # Last-resort compatibility for older agent implementations + try: + setattr(agent, "_session_model_name", model_name or None) + except Exception: + pass + + +def _reload_agent_if_supported(agent: Any) -> None: + reloader = getattr(agent, "reload_code_generation_agent", None) + if callable(reloader): + reloader() + + +# --------------------------------------------------------------------------- +# Validation helpers +# --------------------------------------------------------------------------- + +_SAFE_SESSION_ID_RE = re.compile(r"^[\w.\-]{1,256}$", re.ASCII) + + +def _validate_session_id(session_id: str) -> None: + """Reject session IDs that could cause path-traversal or other mischief.""" + if not session_id or not _SAFE_SESSION_ID_RE.match(session_id): + raise ValueError( + f"Invalid session_id: {session_id!r}. " + "Must be 1–256 chars of [a-zA-Z0-9_.-]." + ) + # Extra guard: no `..` anywhere + if ".." in session_id: + raise ValueError(f"Path traversal detected in session_id: {session_id!r}") + + +def _validate_agent_name(agent_name: str) -> None: + """Ensure the agent actually exists in the registry.""" + available = get_available_agents() + if agent_name not in available: + raise ValueError( + f"Unknown agent {agent_name!r}. Available: {', '.join(sorted(available))}" + ) + + +@dataclass +class SessionContext: + """All mutable state scoped to a single user session. + + Every field lives *only* on this object — nothing is written to global + singletons like ``_SESSION_MODEL`` or ``_CURRENT_AGENT``. + """ + + session_id: str + agent: BaseAgent + agent_name: str + model_name: str + working_directory: str + title: str = "" + pinned: bool = False + created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + websocket: Any | None = None + + # Per-agent-name history cache (survives agent switches within session) + _saved_histories: Dict[str, List[Any]] = field(default_factory=dict) + # Per-agent-name compacted-message hashes + _compacted_hashes: Dict[str, Set[str]] = field(default_factory=dict) + # Serialises mutations (switch_agent, switch_model, save) on this session. + # Stored as Optional so it can be created lazily inside a running event loop. + _op_lock: Optional[asyncio.Lock] = field(default=None, repr=False, compare=False) + + @property + def op_lock(self) -> asyncio.Lock: + """Lazy asyncio.Lock — created on first access inside the event loop.""" + if self._op_lock is None: + self._op_lock = asyncio.Lock() + return self._op_lock + + async def switch_agent_inline(self, new_agent_name: str) -> "BaseAgent": + """Hot-swap the agent on this context, preserving/restoring history. + + Unlike SessionManager.switch_agent(), operates directly on self without + requiring a session_id lookup. Used by MuxChatHandler which manages + the runtime reference directly. + + Returns: + The newly loaded BaseAgent. + """ + from code_puppy.agent_loader import load_agent # local to avoid circular import + + old_name = self.agent_name + + async with self.op_lock: + # Persist outgoing agent state + if hasattr(self.agent, "get_message_history"): + self._saved_histories[old_name] = self.agent.get_message_history() + if hasattr(self.agent, "get_compacted_message_hashes"): + self._compacted_hashes[old_name] = ( + self.agent.get_compacted_message_hashes() + ) + + # Load new agent + new_agent = load_agent(new_agent_name) + + # Restore history if we've visited this agent before + if new_agent_name in self._saved_histories and hasattr( + new_agent, "set_message_history" + ): + new_agent.set_message_history(self._saved_histories[new_agent_name]) + if new_agent_name in self._compacted_hashes and hasattr( + new_agent, "add_compacted_message_hash" + ): + for h in self._compacted_hashes[new_agent_name]: + new_agent.add_compacted_message_hash(h) + + # Carry session model across + _apply_session_model(new_agent, self.model_name) + _reload_agent_if_supported(new_agent) + + self.agent = new_agent + self.agent_name = new_agent_name + + return new_agent + + +# --------------------------------------------------------------------------- +# SessionManager +# --------------------------------------------------------------------------- + + +class SessionManager: + """Thread-safe registry of active ``SessionContext`` instances.""" + + _SESSION_RETENTION_SECONDS = 15 * 60 # 15 minutes + + def __init__(self) -> None: + self._sessions: Dict[str, SessionContext] = {} + self._lock = asyncio.Lock() + self._inactive_since: Dict[ + str, datetime + ] = {} # Track when session became inactive + self._cleanup_task: Optional[asyncio.Task] = None # Background cleanup task + + # -- creation / lookup / teardown ------------------------------------ + + async def create_session( + self, + session_id: str, + agent_name: str = "code-puppy", + model_name: Optional[str] = None, + working_directory: str = "", + ) -> SessionContext: + """Spin up a brand-new isolated session. + + Args: + session_id: Unique identifier (validated for path safety). + agent_name: Which agent to load (default ``code-puppy``). + model_name: Session-local model override. ``None`` → global default. + working_directory: Session-local CWD (never calls ``os.chdir``). + + Returns: + The freshly-created ``SessionContext``. + + Raises: + ValueError: On invalid *session_id* or unknown *agent_name*. + """ + _validate_session_id(session_id) + _validate_agent_name(agent_name) + + agent = load_agent(agent_name) + resolved_model = model_name or get_global_model_name() + + if model_name is not None: + _apply_session_model(agent, model_name) + + ctx = SessionContext( + session_id=session_id, + agent=agent, + agent_name=agent_name, + model_name=resolved_model, + working_directory=working_directory, + ) + + async with self._lock: + self._sessions[session_id] = ctx + + logger.info( + "Session created: id=%s agent=%s model=%s", + session_id, + agent_name, + resolved_model, + ) + return ctx + + async def get_session(self, session_id: str) -> Optional[SessionContext]: + """Look up an active session. Returns ``None`` if not found.""" + async with self._lock: + return self._sessions.get(session_id) + + async def destroy_session(self, session_id: str) -> None: + """Remove a session from the registry.""" + async with self._lock: + removed = self._sessions.pop(session_id, None) + if removed: + logger.info("Session destroyed: id=%s", session_id) + else: + logger.warning("destroy_session called for unknown id=%s", session_id) + + # -- agent switching ------------------------------------------------- + + async def switch_agent( + self, + session_id: str, + new_agent_name: str, + ) -> BaseAgent: + """Hot-swap the agent inside an existing session. + + * Saves current agent’s history & compacted hashes. + * Loads a fresh agent for *new_agent_name*. + * Restores any previously-saved history if switching *back*. + * Carries over the session model. + + Returns: + The newly-loaded ``BaseAgent``. + + Raises: + KeyError: If *session_id* is unknown. + ValueError: If *new_agent_name* is invalid. + """ + ctx = await self._require_session(session_id) + _validate_agent_name(new_agent_name) + + async with ctx.op_lock: # Serialise all mutations on this session + old_name = ctx.agent_name + + # Persist outgoing agent's state + ctx._saved_histories[old_name] = ctx.agent.get_message_history() + ctx._compacted_hashes[old_name] = ctx.agent.get_compacted_message_hashes() + + # Build the replacement agent + new_agent = load_agent(new_agent_name) + + # Restore history if we've visited this agent before + if new_agent_name in ctx._saved_histories: + new_agent.set_message_history(ctx._saved_histories[new_agent_name]) + if new_agent_name in ctx._compacted_hashes: + for h in ctx._compacted_hashes[new_agent_name]: + new_agent.add_compacted_message_hash(h) + + # Carry session model across agents + _apply_session_model(new_agent, ctx.model_name) + _reload_agent_if_supported(new_agent) + + ctx.agent = new_agent + ctx.agent_name = new_agent_name + + logger.info( + "Session %s: switched agent %s → %s", + session_id, + old_name, + new_agent_name, + ) + return new_agent + + # -- model switching ------------------------------------------------- + + async def switch_model(self, session_id: str, new_model: str) -> None: + """Change the model for a session without touching global config. + + Raises: + KeyError: If *session_id* is unknown. + """ + ctx = await self._require_session(session_id) + async with ctx.op_lock: + ctx.model_name = new_model + _apply_session_model(ctx.agent, new_model) + _reload_agent_if_supported(ctx.agent) + logger.info("Session %s: model → %s", session_id, new_model) + + # -- persistence ----------------------------------------------------- + + async def _write_to_sqlite( + self, + session_id: str, + ctx: "SessionContext", + history: List[Any], + ) -> None: + """Write session metadata + any new messages to SQLite. + + Called from save_session() as the sole durable write path. + All errors are caught — a DB failure never breaks the WS session. + + Uses INSERT OR IGNORE on (session_id, seq) so re-running after a + partial write is fully idempotent. + """ + try: + from code_puppy.api.db.queries import ( + insert_messages_batch, + update_session_stats, + upsert_session, + ) + except Exception as exc: + logger.debug("SQLite not available, skipping DB write: %s", exc) + return + + from datetime import datetime + from datetime import timezone as _tz + + updated_at = datetime.now(_tz.utc).isoformat() + + try: + # Upsert session row (creates it if this is the first save) + await upsert_session( + session_id=session_id, + title=ctx.title, + agent_name=ctx.agent_name, + model_name=ctx.model_name, + working_directory=ctx.working_directory, + pinned=ctx.pinned, + created_at=ctx.created_at.isoformat(), + updated_at=updated_at, + message_count=len(history), + total_tokens=0, # updated below + deleted_at=None, + ) + except Exception as exc: + logger.warning("Failed to upsert session %s in SQLite: %s", session_id, exc) + return + + # Build message rows — seq is 1-based index into the full history. + # INSERT OR IGNORE on UNIQUE(session_id, seq) makes this idempotent. + message_rows = [] + total_tokens = 0 + + for idx, msg in enumerate(history): + # Skip system dict entries that aren't proper ModelMessages + if isinstance(msg, dict) and not hasattr(msg, "parts"): + continue + + seq = idx + 1 + ts = get_message_timestamp(msg) or updated_at + content = extract_content(msg) + thinking = extract_thinking(msg) + pj = pydantic_json_for_message(msg) + + try: + token_count = ctx.agent.estimate_tokens_for_message(msg) + except Exception: + token_count = max(1, len(content) // 4) + + total_tokens += token_count + + message_rows.append( + { + "session_id": session_id, + "seq": seq, + "role": get_role(msg), + "content": content, + "type": type(msg).__name__, + "agent_name": ctx.agent_name, + "model_name": ctx.model_name, + "timestamp": ts, + "thinking": thinking, + "token_count": token_count, + "pydantic_json": pj, + } + ) + + try: + await insert_messages_batch(message_rows) + except Exception as exc: + logger.warning( + "Failed to insert messages for %s in SQLite: %s", session_id, exc + ) + return + + # Update the token count now that we've summed it + try: + await update_session_stats( + session_id, + message_count=len(history), + total_tokens=total_tokens, + updated_at=updated_at, + ) + except Exception as exc: + logger.warning( + "Failed to update session stats for %s in SQLite: %s", session_id, exc + ) + + logger.debug( + "Session %s written to SQLite (%d messages, %d tokens)", + session_id, + len(history), + total_tokens, + ) + + async def save_session(self, session_id: str) -> None: + """Persist session history and metadata to SQLite (single source of truth). + + Previously wrote session files under {WS_SESSION_DIR} — that legacy file write + has been removed. SQLite at ~/.puppy_desk/chat_messages.db is now the + only durable store. + + aiosqlite serialises all DB I/O through its own background thread so + this coroutine yields to the event loop on every await — no + thread-pool needed. + """ + ctx = await self._require_session(session_id) + _validate_session_id(session_id) + + async with ctx.op_lock: + history: List[Any] = ctx.agent.get_message_history() + + await self._write_to_sqlite(session_id, ctx, history) + logger.debug( + "Session %s saved to SQLite (%d messages)", session_id, len(history) + ) + + async def _load_from_sqlite( + self, session_id: str + ) -> Optional[tuple[List[Any], dict]]: + """Try to load message history + metadata from SQLite. + + Returns (messages, meta_dict) on success, None if SQLite is unavailable + or the session has no pydantic_json rows (seeded without BE). + """ + try: + from code_puppy.api.db.queries import get_active_messages, session_exists + except Exception: + return None + + try: + if not await session_exists(session_id): + return None + + rows = await get_active_messages(session_id) + # Filter to only rows with pydantic_json (seeder may leave some NULL). + # Sessions with no parseable messages still resume — just with empty history. + rows_with_json = [r for r in rows if r.get("pydantic_json")] + + from pydantic_ai.messages import ModelMessagesTypeAdapter + + messages: List[Any] = [] + for row in rows_with_json: + try: + parsed = ModelMessagesTypeAdapter.validate_json( + row["pydantic_json"] + ) + if parsed: + messages.append(parsed[0]) + except Exception as exc: + logger.warning( + "Failed to deserialise message seq=%s for session %s: %s", + row.get("seq"), + session_id, + exc, + ) + + # Build a metadata dict from the sessions table + try: + from code_puppy.api.db.connection import get_db + + db = get_db() + cursor = await db.execute( + "SELECT * FROM sessions WHERE session_id = ?", (session_id,) + ) + row_s = await cursor.fetchone() + meta: dict = dict(row_s) if row_s else {} + except Exception: + meta = {} + + logger.info( + "Session %s loaded from SQLite (%d messages)", + session_id, + len(messages), + ) + return messages, meta + + except Exception as exc: + logger.warning( + "SQLite load failed for session %s: %s", + session_id, + exc, + ) + return None + + async def load_session(self, session_id: str) -> Optional[SessionContext]: + """Restore a session from disk, creating a fresh agent. + + Returns ``None`` if no persisted data exists or loading fails. + """ + _validate_session_id(session_id) + + # ---- try SQLite first (sessions saved by BE write path) ---------- + sqlite_result = await self._load_from_sqlite(session_id) + if sqlite_result is not None: + messages, db_meta = sqlite_result + + agent_name = db_meta.get("agent_name", "code-puppy") + model_name = db_meta.get("model_name", get_global_model_name()) + title = db_meta.get("title", "") + pinned = bool(db_meta.get("pinned", False)) + working_directory = db_meta.get("working_directory", "") + created_at_raw = db_meta.get("created_at", "") + try: + created_at = ( + datetime.fromisoformat(created_at_raw) + if created_at_raw + else datetime.now(timezone.utc) + ) + except Exception: + created_at = datetime.now(timezone.utc) + + try: + agent = load_agent(agent_name) + except ValueError: + logger.warning( + "Agent %s unavailable, falling back to code-puppy", agent_name + ) + agent_name = "code-puppy" + agent = load_agent(agent_name) + + agent.set_message_history(messages) + _apply_session_model(agent, model_name) + + ctx = SessionContext( + session_id=session_id, + agent=agent, + agent_name=agent_name, + model_name=model_name, + working_directory=working_directory, + title=title, + pinned=pinned, + created_at=created_at, + ) + + async with self._lock: + self._sessions[session_id] = ctx + + return ctx + # SQLite is the single source of truth — no file fallback. + # If _load_from_sqlite() returned None the session does not exist in the DB. + logger.debug("Session %s not found in SQLite — treating as new", session_id) + return None + + # -- session retention (15-min inactive cleanup) ------------------------- + + async def mark_session_inactive(self, session_id: str) -> None: + """Mark a session as inactive (WS disconnected). Starts retention timer.""" + async with self._lock: + if session_id in self._sessions: + self._inactive_since[session_id] = datetime.now(timezone.utc) + logger.debug( + "Session %s marked inactive, will be cleaned up in 15 min", + session_id, + ) + # Ensure cleanup task is running + self._ensure_cleanup_task() + + async def mark_session_active(self, session_id: str) -> None: + """Mark a session as active (WS connected). Cancels pending cleanup.""" + async with self._lock: + self._inactive_since.pop(session_id, None) + logger.debug("Session %s marked active", session_id) + + def _ensure_cleanup_task(self) -> None: + """Start the background cleanup task if not already running.""" + if self._cleanup_task is None or self._cleanup_task.done(): + self._cleanup_task = asyncio.create_task(self._cleanup_inactive_sessions()) + + async def _cleanup_inactive_sessions(self) -> None: + """Background task that cleans up sessions inactive for > 15 minutes.""" + while True: + await asyncio.sleep(60) # Check every minute + now = datetime.now(timezone.utc) + to_cleanup: List[str] = [] + + async with self._lock: + for session_id, inactive_time in list(self._inactive_since.items()): + elapsed = (now - inactive_time).total_seconds() + if elapsed > self._SESSION_RETENTION_SECONDS: + to_cleanup.append(session_id) + + for session_id in to_cleanup: + logger.info( + "Cleaning up inactive session %s (inactive > 15 min)", session_id + ) + await self.destroy_session(session_id) + async with self._lock: + self._inactive_since.pop(session_id, None) + + # Stop task if no more inactive sessions to monitor + async with self._lock: + if not self._inactive_since: + break + + async def get_or_load_session(self, session_id: str) -> Optional[SessionContext]: + """Get an existing in-memory session OR load from SQLite. + + This is the preferred method for session access - it checks + in-memory first (including inactive sessions pending cleanup), + then falls back to SQLite load. + """ + # First check if session exists in memory (active or inactive) + async with self._lock: + ctx = self._sessions.get(session_id) + if ctx is not None: + # Session exists in memory - mark it active and return + self._inactive_since.pop(session_id, None) + logger.debug("Reusing in-memory session %s", session_id) + return ctx + + # Not in memory - try to load from SQLite + return await self.load_session(session_id) + + # -- internal helpers ------------------------------------------------ + + async def _require_session(self, session_id: str) -> SessionContext: + """Return the session or raise ``KeyError``.""" + async with self._lock: + ctx = self._sessions.get(session_id) + if ctx is None: + raise KeyError(f"No active session with id={session_id!r}") + return ctx + + +# --------------------------------------------------------------------------- +# Module-level singleton +# --------------------------------------------------------------------------- + +session_manager = SessionManager() diff --git a/code_puppy/api/session_cwd.py b/code_puppy/api/session_cwd.py new file mode 100644 index 000000000..270b4188a --- /dev/null +++ b/code_puppy/api/session_cwd.py @@ -0,0 +1,39 @@ +"""Per-agent-run working-directory ContextVar helpers. + +The WebSocket chat UI stores a session working directory on the runtime/session +metadata. Shell command execution needs access to that value without calling +``os.chdir()``, because multiple browser sessions may run concurrently in one +server process. +""" + +from __future__ import annotations + +from contextvars import ContextVar +from typing import Optional + +_session_working_directory: ContextVar[Optional[str]] = ContextVar( + "code_puppy_session_working_directory", + default=None, +) + + +def set_session_working_directory(directory: str | None) -> object: + """Set the current agent-run working directory and return reset token.""" + value = directory or None + return _session_working_directory.set(value) + + +def get_session_working_directory() -> str | None: + """Return the current agent-run working directory, if one is set.""" + return _session_working_directory.get(None) + + +def clear_session_working_directory(token: object | None = None) -> None: + """Clear or reset the current agent-run working directory.""" + if token is not None: + try: + _session_working_directory.reset(token) # type: ignore[arg-type] + return + except Exception: + pass + _session_working_directory.set(None) diff --git a/code_puppy/api/templates/chat.html b/code_puppy/api/templates/chat.html new file mode 100644 index 000000000..b0e1088be --- /dev/null +++ b/code_puppy/api/templates/chat.html @@ -0,0 +1,2450 @@ + + + + + + 🐶 Code Puppy Chat + + + + + + + + + + + + + + + + +
+ +
+ + +
+ + + + Connecting… + + +
+ Agent + +
+ + +
+ Model + +
+ + +
+
+ Disconnected +
+ + + + Home +
+ + +
+ + +
+ + Not set +
+ + +
+ + + + + + + + + + +
+ + + +
+ +
+ + +
+ +
+ + +
+
+ + +
+
+
+
+ + + + diff --git a/code_puppy/api/tests/test_backfill_pydantic_json.py b/code_puppy/api/tests/test_backfill_pydantic_json.py new file mode 100644 index 000000000..db026dc03 --- /dev/null +++ b/code_puppy/api/tests/test_backfill_pydantic_json.py @@ -0,0 +1,349 @@ +"""Unit tests for code_puppy.api.db.backfill.backfill_pydantic_json. + +All tests spin up a real in-memory aiosqlite database with the minimal schema +needed — no mocking of the DB layer, no patching of pydantic-ai. The goal +is to exercise the actual serialisation round-trip so we catch pydantic-ai +API changes as early as possible. + +Test inventory +-------------- + 1. Empty table → (0, 0) returned, no crash + 2. All rows already filled → (0, 0), nothing overwritten + 3. User message backfilled → pydantic_json set, deserialises to ModelRequest + 4. Assistant message backf. → pydantic_json set, deserialises to ModelResponse + 5. System rows skipped → pydantic_json remains NULL + 6. Compacted rows skipped → pydantic_json remains NULL + 7. Mixed NULL + non-NULL → only NULL rows updated + 8. Empty content string → backfilled with empty UserPromptPart (no crash) + 9. Idempotent (run twice) → second run updates 0 rows +10. Batch splitting → batch_size=2 with 5 rows all updated correctly +11. Serialised JSON is valid → round-trips through ModelMessagesTypeAdapter +12. _build_message bad role → raises ValueError +""" + +from __future__ import annotations + +import pytest + +# --------------------------------------------------------------------------- +# In-memory DB fixture +# --------------------------------------------------------------------------- + +_CREATE_MESSAGES = """ +CREATE TABLE messages ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + seq INTEGER NOT NULL, + role TEXT NOT NULL, + content TEXT DEFAULT '', + model_name TEXT DEFAULT '', + compacted INTEGER DEFAULT 0, + pydantic_json TEXT +); +""" + + +async def _make_db(): + """Return an open in-memory aiosqlite connection with the messages table.""" + import aiosqlite + + db = await aiosqlite.connect(":memory:") + db.row_factory = aiosqlite.Row + await db.execute(_CREATE_MESSAGES) + await db.commit() + return db + + +async def _insert( + db, + *, + session_id: str = "s1", + seq: int = 1, + role: str = "user", + content: str = "hello", + model_name: str = "gpt-4o", + compacted: int = 0, + pydantic_json: str | None = None, +) -> int: + """Insert one row and return its id.""" + cursor = await db.execute( + """ + INSERT INTO messages (session_id, seq, role, content, model_name, compacted, pydantic_json) + VALUES (?, ?, ?, ?, ?, ?, ?) + """, + (session_id, seq, role, content, model_name, compacted, pydantic_json), + ) + await db.commit() + return cursor.lastrowid + + +async def _fetch_pj(db, row_id: int) -> str | None: + """Fetch pydantic_json for a specific row id.""" + cursor = await db.execute( + "SELECT pydantic_json FROM messages WHERE id = ?", (row_id,) + ) + row = await cursor.fetchone() + return row["pydantic_json"] if row else None + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_empty_table_returns_zero_counts(): + """No rows → (0, 0) with no crash.""" + from code_puppy.api.db.backfill import backfill_pydantic_json + + db = await _make_db() + updated, skipped = await backfill_pydantic_json(db) + await db.close() + + assert updated == 0 + assert skipped == 0 + + +@pytest.mark.asyncio +async def test_all_rows_already_filled_returns_zero_counts(): + """Rows with existing pydantic_json are not touched.""" + from code_puppy.api.db.backfill import backfill_pydantic_json + + db = await _make_db() + existing_pj = '{"already": "there"}' + await _insert(db, pydantic_json=existing_pj) + + updated, skipped = await backfill_pydantic_json(db) + await db.close() + + assert updated == 0 + assert skipped == 0 + + +@pytest.mark.asyncio +async def test_user_message_backfilled_to_model_request(): + """A user row with NULL pydantic_json is backfilled with a ModelRequest.""" + from pydantic_ai.messages import ModelMessagesTypeAdapter, ModelRequest + + from code_puppy.api.db.backfill import backfill_pydantic_json + + db = await _make_db() + row_id = await _insert(db, role="user", content="hello world", pydantic_json=None) + + updated, skipped = await backfill_pydantic_json(db) + + pj = await _fetch_pj(db, row_id) + await db.close() + + assert updated == 1 + assert skipped == 0 + assert pj is not None + + # Round-trip check + msgs = ModelMessagesTypeAdapter.validate_json(pj) + assert len(msgs) == 1 + assert isinstance(msgs[0], ModelRequest) + # Content preserved + content_parts = [p.content for p in msgs[0].parts if hasattr(p, "content")] + assert "hello world" in content_parts + + +@pytest.mark.asyncio +async def test_assistant_message_backfilled_to_model_response(): + """An assistant row with NULL pydantic_json is backfilled with a ModelResponse.""" + from pydantic_ai.messages import ModelMessagesTypeAdapter, ModelResponse + + from code_puppy.api.db.backfill import backfill_pydantic_json + + db = await _make_db() + row_id = await _insert( + db, + role="assistant", + content="here is my answer", + model_name="claude-3", + pydantic_json=None, + ) + + updated, skipped = await backfill_pydantic_json(db) + + pj = await _fetch_pj(db, row_id) + await db.close() + + assert updated == 1 + assert skipped == 0 + assert pj is not None + + msgs = ModelMessagesTypeAdapter.validate_json(pj) + assert len(msgs) == 1 + assert isinstance(msgs[0], ModelResponse) + content_parts = [p.content for p in msgs[0].parts if hasattr(p, "content")] + assert "here is my answer" in content_parts + + +@pytest.mark.asyncio +async def test_system_rows_are_skipped(): + """Rows with role='system' are excluded from the backfill query.""" + from code_puppy.api.db.backfill import backfill_pydantic_json + + db = await _make_db() + row_id = await _insert( + db, role="system", content="you are helpful", pydantic_json=None + ) + + updated, skipped = await backfill_pydantic_json(db) + pj = await _fetch_pj(db, row_id) + await db.close() + + assert updated == 0 + assert skipped == 0 + assert pj is None, "system rows must stay NULL" + + +@pytest.mark.asyncio +async def test_compacted_rows_are_skipped(): + """Rows with compacted=1 are excluded from the backfill query.""" + from code_puppy.api.db.backfill import backfill_pydantic_json + + db = await _make_db() + row_id = await _insert( + db, role="user", content="old message", compacted=1, pydantic_json=None + ) + + updated, skipped = await backfill_pydantic_json(db) + pj = await _fetch_pj(db, row_id) + await db.close() + + assert updated == 0 + assert skipped == 0 + assert pj is None, "compacted rows must stay NULL" + + +@pytest.mark.asyncio +async def test_mixed_null_and_filled_rows(): + """Only NULL rows are updated; already-filled rows are untouched.""" + from code_puppy.api.db.backfill import backfill_pydantic_json + + db = await _make_db() + existing_pj = '{"sentinel": true}' + + id_null = await _insert(db, seq=1, pydantic_json=None) + id_filled = await _insert(db, seq=2, pydantic_json=existing_pj) + + updated, skipped = await backfill_pydantic_json(db) + + pj_null = await _fetch_pj(db, id_null) + pj_filled = await _fetch_pj(db, id_filled) + await db.close() + + assert updated == 1 + assert pj_null is not None, "NULL row should have been filled" + assert pj_filled == existing_pj, "Already-filled row must not be overwritten" + + +@pytest.mark.asyncio +async def test_empty_content_backfilled_without_crash(): + """Empty string content produces a valid pydantic_json (no crash).""" + from pydantic_ai.messages import ModelMessagesTypeAdapter + + from code_puppy.api.db.backfill import backfill_pydantic_json + + db = await _make_db() + row_id = await _insert(db, role="user", content="", pydantic_json=None) + + updated, skipped = await backfill_pydantic_json(db) + pj = await _fetch_pj(db, row_id) + await db.close() + + assert updated == 1 + assert skipped == 0 + assert pj is not None + # Must be parseable + msgs = ModelMessagesTypeAdapter.validate_json(pj) + assert len(msgs) == 1 + + +@pytest.mark.asyncio +async def test_idempotent_second_run_updates_nothing(): + """Running the backfill twice produces updated=0 on the second run.""" + from code_puppy.api.db.backfill import backfill_pydantic_json + + db = await _make_db() + await _insert(db, role="user", content="test", pydantic_json=None) + + updated1, _ = await backfill_pydantic_json(db) + updated2, _ = await backfill_pydantic_json(db) + await db.close() + + assert updated1 == 1 + assert updated2 == 0, "Second run must be a no-op" + + +@pytest.mark.asyncio +async def test_batch_splitting_all_rows_updated(): + """batch_size=2 with 5 rows — all rows must be updated correctly.""" + from pydantic_ai.messages import ModelMessagesTypeAdapter + + from code_puppy.api.db.backfill import backfill_pydantic_json + + db = await _make_db() + row_ids = [] + for i in range(5): + rid = await _insert( + db, + seq=i + 1, + role="user", + content=f"message {i}", + pydantic_json=None, + ) + row_ids.append(rid) + + updated, skipped = await backfill_pydantic_json(db, batch_size=2) + + for rid in row_ids: + pj = await _fetch_pj(db, rid) + assert pj is not None, f"Row {rid} should have been backfilled" + msgs = ModelMessagesTypeAdapter.validate_json(pj) + assert len(msgs) == 1 + + await db.close() + assert updated == 5 + assert skipped == 0 + + +@pytest.mark.asyncio +async def test_serialised_json_round_trips_correctly(): + """The generated pydantic_json must survive a full serialise → deserialise cycle.""" + from pydantic_ai.messages import ( + ModelMessagesTypeAdapter, + ModelRequest, + UserPromptPart, + ) + + from code_puppy.api.db.backfill import backfill_pydantic_json + + original_content = "does this survive the round trip? 🐶" + db = await _make_db() + row_id = await _insert( + db, role="user", content=original_content, pydantic_json=None + ) + + await backfill_pydantic_json(db) + pj = await _fetch_pj(db, row_id) + await db.close() + + msgs = ModelMessagesTypeAdapter.validate_json(pj) + assert isinstance(msgs[0], ModelRequest) + prompt_parts = [p for p in msgs[0].parts if isinstance(p, UserPromptPart)] + assert len(prompt_parts) == 1 + assert prompt_parts[0].content == original_content + + +def test_build_message_bad_role_raises(): + """_build_message raises ValueError for unsupported roles.""" + from code_puppy.api.db.backfill import _build_message + + with pytest.raises(ValueError, match="Unsupported role"): + _build_message("system", "content", "model") + + with pytest.raises(ValueError, match="Unsupported role"): + _build_message("tool", "content", "model") 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..01e4851fb --- /dev/null +++ b/code_puppy/api/tests/test_chat_template_xss_guards.py @@ -0,0 +1,30 @@ +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 + 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 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_db_sql_loader.py b/code_puppy/api/tests/test_db_sql_loader.py new file mode 100644 index 000000000..d60a3c912 --- /dev/null +++ b/code_puppy/api/tests/test_db_sql_loader.py @@ -0,0 +1,22 @@ +from pathlib import Path + + +def test_external_session_history_sql_files_exist_and_contain_expected_clauses(): + sql_dir = Path("code_puppy/api/db/sql") + parity_sql = (sql_dir / "session_history_parity.sql").read_text(encoding="utf-8") + no_compacted_sql = (sql_dir / "session_history_parity_no_compacted.sql").read_text( + encoding="utf-8" + ) + + assert "FROM messages m" in parity_sql + assert "FROM tool_calls tc" in parity_sql + assert "ORDER BY seq" in parity_sql + + assert "m.compacted = 0" in no_compacted_sql + assert "ORDER BY seq" in no_compacted_sql + + +def test_queries_module_uses_sql_loader_for_complex_session_query(): + queries_src = Path("code_puppy/api/db/queries.py").read_text(encoding="utf-8") + assert 'load_sql("session_history_parity.sql")' in queries_src + assert "session_history_parity_no_compacted.sql" in queries_src 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_permission_plugin_fail_closed.py b/code_puppy/api/tests/test_permission_plugin_fail_closed.py new file mode 100644 index 000000000..9ff60b81e --- /dev/null +++ b/code_puppy/api/tests/test_permission_plugin_fail_closed.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +import pytest + +from code_puppy.api.permission_plugin import ( + clear_websocket_context, + pre_tool_call_permission, + set_websocket_context, + shell_command_permission, +) + + +@pytest.mark.asyncio +async def test_pre_tool_call_permission_fails_closed_on_request_error(monkeypatch): + set_websocket_context(websocket=object(), session_id="session-1") + try: + + async def _boom(**_kwargs): + raise RuntimeError("permission transport down") + + monkeypatch.setattr("code_puppy.api.permissions.request_permission", _boom) + + result = await pre_tool_call_permission( + "agent_run_shell_command", {"command": "ls"} + ) + assert isinstance(result, dict) + assert result["blocked"] is True + assert result["error"] == "Permission system error" + finally: + clear_websocket_context() + + +@pytest.mark.asyncio +async def test_shell_command_permission_fails_closed_on_request_error(monkeypatch): + set_websocket_context(websocket=object(), session_id="session-1") + try: + + async def _boom(**_kwargs): + raise RuntimeError("permission transport down") + + monkeypatch.setattr("code_puppy.api.permissions.request_permission", _boom) + + result = await shell_command_permission( + context=None, + command="ls", + cwd=".", + timeout=30, + ) + assert isinstance(result, dict) + assert result["blocked"] is True + assert result["error"] == "Permission system error" + finally: + clear_websocket_context() 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_session_load.py b/code_puppy/api/tests/test_session_load.py new file mode 100644 index 000000000..597cd3620 --- /dev/null +++ b/code_puppy/api/tests/test_session_load.py @@ -0,0 +1,575 @@ +"""Unit tests for SessionManager._load_from_sqlite and .load_session. + +Phase 4 coverage from the session-management bug plan: + +_load_from_sqlite: + 1. Session not in DB → returns None + 2. SQLite import unavailable → returns None + 3. get_active_messages raises → returns None (outer except) + 4. All rows NULL pydantic_json → returns ([], meta) – graceful degradation + 5. Mixed NULL + valid rows → only valid rows parsed + 6. Valid rows → messages parsed correctly + 7. Malformed pydantic_json → row skipped + warning, rest succeed + 8. Metadata populated from sessions table + 9. No session row in sessions table → meta = {} +10. Sessions table query fails → meta = {} + +load_session: +11. Invalid session_id → ValueError before touching DB +12. _load_from_sqlite → None → load_session returns None +13. Normal happy path → history set, ctx in _sessions +14. Unknown agent_name in DB → falls back to code-puppy +15. Valid created_at string → parsed as datetime correctly +16. Malformed created_at → datetime.now() fallback +17. Empty created_at → datetime.now() fallback +18. Session registered in _sessions after load +""" + +from __future__ import annotations + +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_pydantic_json(content: str = "hello from user") -> str: + """Produce a valid pydantic_json string using real pydantic-ai types.""" + from pydantic_ai.messages import ( + ModelMessagesTypeAdapter, + ModelRequest, + UserPromptPart, + ) + + msg = ModelRequest(parts=[UserPromptPart(content=content)]) + return ModelMessagesTypeAdapter.dump_json([msg]).decode("utf-8") + + +def _make_response_pydantic_json(content: str = "hello from assistant") -> str: + """Produce a valid pydantic_json string for a ModelResponse.""" + from pydantic_ai.messages import ModelMessagesTypeAdapter, ModelResponse, TextPart + + msg = ModelResponse(parts=[TextPart(content=content)], model_name="test-model") + return ModelMessagesTypeAdapter.dump_json([msg]).decode("utf-8") + + +def _make_row( + seq: int = 1, + role: str = "user", + pydantic_json: str | None = None, +) -> dict: + """Build a minimal message row dict as returned by get_active_messages.""" + return { + "seq": seq, + "role": role, + "content": "some content", + "pydantic_json": pydantic_json, + } + + +def _make_session_row( + session_id: str = "test-session", + agent_name: str = "code-puppy", + model_name: str = "gpt-4o", + title: str = "Test Session", + pinned: bool = False, + working_directory: str = "/tmp", + created_at: str = "2026-01-01T00:00:00+00:00", +) -> dict: + """Build a minimal sessions table row dict.""" + return { + "session_id": session_id, + "agent_name": agent_name, + "model_name": model_name, + "title": title, + "pinned": 1 if pinned else 0, + "working_directory": working_directory, + "created_at": created_at, + } + + +def _make_db_mock(session_row: dict | None) -> MagicMock: + """Return a db mock whose cursor.fetchone() returns *session_row*.""" + cursor = AsyncMock() + cursor.fetchone.return_value = session_row + db = MagicMock() + db.execute = AsyncMock(return_value=cursor) + return db + + +def _make_fresh_manager(): + """Return a brand-new SessionManager instance (not the module singleton).""" + from code_puppy.api.session_context import SessionManager + + return SessionManager() + + +def _make_agent_mock(history: list | None = None) -> MagicMock: + """Return a minimal agent mock.""" + agent = MagicMock() + agent.get_message_history.return_value = history or [] + agent.set_message_history.return_value = None + agent.set_session_model.return_value = None + return agent + + +# --------------------------------------------------------------------------- +# _load_from_sqlite tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_load_from_sqlite_session_not_in_db(): + """session_exists returns False → _load_from_sqlite returns None.""" + mgr = _make_fresh_manager() + db_mock = _make_db_mock(None) + + with ( + patch( + "code_puppy.api.db.queries.session_exists", + new=AsyncMock(return_value=False), + ), + patch( + "code_puppy.api.db.queries.get_active_messages", + new=AsyncMock(return_value=[]), + ), + patch("code_puppy.api.db.connection.get_db", return_value=db_mock), + ): + result = await mgr._load_from_sqlite("test-session") + + assert result is None + + +@pytest.mark.asyncio +async def test_load_from_sqlite_import_error_returns_none(): + """If the DB queries module is unavailable, _load_from_sqlite returns None.""" + mgr = _make_fresh_manager() + + # Raising ImportError inside the try block that wraps the lazy import + # is the canonical way to simulate "SQLite not available". + with patch( + "code_puppy.api.db.queries.session_exists", + side_effect=ImportError("aiosqlite not installed"), + ): + result = await mgr._load_from_sqlite("test-session") + + assert result is None + + +@pytest.mark.asyncio +async def test_load_from_sqlite_get_active_messages_raises_returns_none(): + """An exception from get_active_messages is caught → None.""" + mgr = _make_fresh_manager() + + with ( + patch( + "code_puppy.api.db.queries.session_exists", new=AsyncMock(return_value=True) + ), + patch( + "code_puppy.api.db.queries.get_active_messages", + new=AsyncMock(side_effect=RuntimeError("db locked")), + ), + ): + result = await mgr._load_from_sqlite("test-session") + + assert result is None + + +@pytest.mark.asyncio +async def test_load_from_sqlite_all_null_pydantic_json_returns_empty_history(): + """All rows have pydantic_json=None → empty message list, but NOT None. + + This is graceful degradation: the session exists but has no serialisable + messages (e.g. seeded without the BE write path). The agent starts with + empty history rather than crashing. + """ + mgr = _make_fresh_manager() + rows = [ + _make_row(seq=1, role="user", pydantic_json=None), + _make_row(seq=2, role="assistant", pydantic_json=None), + ] + db_mock = _make_db_mock(_make_session_row()) + + with ( + patch( + "code_puppy.api.db.queries.session_exists", new=AsyncMock(return_value=True) + ), + patch( + "code_puppy.api.db.queries.get_active_messages", + new=AsyncMock(return_value=rows), + ), + patch("code_puppy.api.db.connection.get_db", return_value=db_mock), + ): + result = await mgr._load_from_sqlite("test-session") + + assert result is not None, "Should return (messages, meta) not None" + messages, meta = result + assert messages == [], ( + "Empty history expected when all rows have NULL pydantic_json" + ) + + +@pytest.mark.asyncio +async def test_load_from_sqlite_mixed_null_and_valid_rows(): + """NULL pydantic_json rows are skipped; valid rows are parsed.""" + mgr = _make_fresh_manager() + valid_json = _make_pydantic_json("hello") + rows = [ + _make_row(seq=1, role="user", pydantic_json=None), # NULL – skip + _make_row(seq=2, role="user", pydantic_json=valid_json), # valid + _make_row(seq=3, role="user", pydantic_json=None), # NULL – skip + ] + db_mock = _make_db_mock(_make_session_row()) + + with ( + patch( + "code_puppy.api.db.queries.session_exists", new=AsyncMock(return_value=True) + ), + patch( + "code_puppy.api.db.queries.get_active_messages", + new=AsyncMock(return_value=rows), + ), + patch("code_puppy.api.db.connection.get_db", return_value=db_mock), + ): + result = await mgr._load_from_sqlite("test-session") + + assert result is not None + messages, _ = result + assert len(messages) == 1, "Only the valid row should be parsed" + + +@pytest.mark.asyncio +async def test_load_from_sqlite_valid_rows_parsed_correctly(): + """Valid pydantic_json rows deserialise into proper ModelMessage objects.""" + from pydantic_ai.messages import ModelRequest + + mgr = _make_fresh_manager() + user_json = _make_pydantic_json("user turn") + assistant_json = _make_response_pydantic_json("assistant reply") + rows = [ + _make_row(seq=1, role="user", pydantic_json=user_json), + _make_row(seq=2, role="assistant", pydantic_json=assistant_json), + ] + db_mock = _make_db_mock(_make_session_row()) + + with ( + patch( + "code_puppy.api.db.queries.session_exists", new=AsyncMock(return_value=True) + ), + patch( + "code_puppy.api.db.queries.get_active_messages", + new=AsyncMock(return_value=rows), + ), + patch("code_puppy.api.db.connection.get_db", return_value=db_mock), + ): + result = await mgr._load_from_sqlite("test-session") + + assert result is not None + messages, _ = result + assert len(messages) == 2 + # First message must be a ModelRequest (user turn) + assert isinstance(messages[0], ModelRequest) + + +@pytest.mark.asyncio +async def test_load_from_sqlite_malformed_pydantic_json_skipped(caplog): + """A row with corrupt JSON is skipped with a warning; valid rows still load.""" + import logging + + mgr = _make_fresh_manager() + valid_json = _make_pydantic_json("good message") + rows = [ + _make_row(seq=1, role="user", pydantic_json="{this is not valid json!!!"), + _make_row(seq=2, role="user", pydantic_json=valid_json), + ] + db_mock = _make_db_mock(_make_session_row()) + + with ( + patch( + "code_puppy.api.db.queries.session_exists", new=AsyncMock(return_value=True) + ), + patch( + "code_puppy.api.db.queries.get_active_messages", + new=AsyncMock(return_value=rows), + ), + patch("code_puppy.api.db.connection.get_db", return_value=db_mock), + caplog.at_level(logging.WARNING, logger="code_puppy.api.session_context"), + ): + result = await mgr._load_from_sqlite("test-session") + + assert result is not None + messages, _ = result + # Only the valid row should come through + assert len(messages) == 1 + # A warning must have been emitted for the bad row + assert any("Failed to deserialise" in r.message for r in caplog.records) + + +@pytest.mark.asyncio +async def test_load_from_sqlite_metadata_populated(): + """Meta dict is populated from the sessions table row.""" + mgr = _make_fresh_manager() + session_row = _make_session_row( + agent_name="planning-agent", + model_name="claude-3", + title="My Session", + pinned=True, + working_directory="/home/user/project", + ) + db_mock = _make_db_mock(session_row) + + with ( + patch( + "code_puppy.api.db.queries.session_exists", new=AsyncMock(return_value=True) + ), + patch( + "code_puppy.api.db.queries.get_active_messages", + new=AsyncMock(return_value=[]), + ), + patch("code_puppy.api.db.connection.get_db", return_value=db_mock), + ): + result = await mgr._load_from_sqlite("test-session") + + assert result is not None + _, meta = result + assert meta["agent_name"] == "planning-agent" + assert meta["model_name"] == "claude-3" + assert meta["title"] == "My Session" + assert meta["working_directory"] == "/home/user/project" + + +@pytest.mark.asyncio +async def test_load_from_sqlite_no_session_row_gives_empty_meta(): + """If the sessions table has no row, meta is an empty dict.""" + mgr = _make_fresh_manager() + db_mock = _make_db_mock(None) # cursor.fetchone() → None + + with ( + patch( + "code_puppy.api.db.queries.session_exists", new=AsyncMock(return_value=True) + ), + patch( + "code_puppy.api.db.queries.get_active_messages", + new=AsyncMock(return_value=[]), + ), + patch("code_puppy.api.db.connection.get_db", return_value=db_mock), + ): + result = await mgr._load_from_sqlite("test-session") + + assert result is not None + _, meta = result + assert meta == {} + + +@pytest.mark.asyncio +async def test_load_from_sqlite_sessions_table_query_fails_gives_empty_meta(): + """If the sessions table query raises, meta falls back to {} without crashing.""" + mgr = _make_fresh_manager() + + # get_db() itself raises on first call (for the metadata query only) + failing_db = MagicMock() + failing_db.execute = AsyncMock(side_effect=RuntimeError("table locked")) + + with ( + patch( + "code_puppy.api.db.queries.session_exists", new=AsyncMock(return_value=True) + ), + patch( + "code_puppy.api.db.queries.get_active_messages", + new=AsyncMock(return_value=[]), + ), + patch("code_puppy.api.db.connection.get_db", return_value=failing_db), + ): + result = await mgr._load_from_sqlite("test-session") + + assert result is not None + _, meta = result + assert meta == {} + + +# --------------------------------------------------------------------------- +# load_session tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_load_session_invalid_session_id_raises(): + """load_session raises ValueError before touching the DB for bad IDs.""" + mgr = _make_fresh_manager() + with pytest.raises(ValueError, match="Invalid session_id"): + await mgr.load_session("../etc/passwd") + + +@pytest.mark.asyncio +async def test_load_session_returns_none_when_sqlite_has_no_data(): + """If _load_from_sqlite returns None, load_session returns None.""" + mgr = _make_fresh_manager() + mgr._load_from_sqlite = AsyncMock(return_value=None) + + result = await mgr.load_session("valid-session-id") + + assert result is None + + +@pytest.mark.asyncio +async def test_load_session_happy_path_sets_history_and_registers_context(): + """Normal path: history loaded from SQLite is set on the agent, and the + context is registered in the manager's _sessions dict. + """ + from pydantic_ai.messages import ModelRequest, UserPromptPart + + mgr = _make_fresh_manager() + real_msg = ModelRequest(parts=[UserPromptPart(content="hello")]) + db_meta = _make_session_row( + agent_name="code-puppy", + model_name="gpt-4o", + title="My Chat", + created_at="2026-01-15T12:00:00+00:00", + ) + mgr._load_from_sqlite = AsyncMock(return_value=([real_msg], db_meta)) + agent_mock = _make_agent_mock() + + with ( + patch("code_puppy.api.session_context.load_agent", return_value=agent_mock), + patch( + "code_puppy.api.session_context.get_global_model_name", + return_value="gpt-4o", + ), + ): + ctx = await mgr.load_session("my-chat-session") + + assert ctx is not None + assert ctx.session_id == "my-chat-session" + assert ctx.agent is agent_mock + assert ctx.agent_name == "code-puppy" + assert ctx.model_name == "gpt-4o" + assert ctx.title == "My Chat" + # set_message_history must have been called with the loaded messages + agent_mock.set_message_history.assert_called_once_with([real_msg]) + + +@pytest.mark.asyncio +async def test_load_session_registers_context_in_sessions_dict(): + """The loaded SessionContext must end up in SessionManager._sessions.""" + mgr = _make_fresh_manager() + db_meta = _make_session_row() + mgr._load_from_sqlite = AsyncMock(return_value=([], db_meta)) + agent_mock = _make_agent_mock() + + with ( + patch("code_puppy.api.session_context.load_agent", return_value=agent_mock), + patch( + "code_puppy.api.session_context.get_global_model_name", + return_value="gpt-4o", + ), + ): + ctx = await mgr.load_session("registered-session") + + assert "registered-session" in mgr._sessions + assert mgr._sessions["registered-session"] is ctx + + +@pytest.mark.asyncio +async def test_load_session_unknown_agent_falls_back_to_code_puppy(): + """If the agent stored in DB is no longer available, code-puppy is used.""" + mgr = _make_fresh_manager() + db_meta = _make_session_row(agent_name="deleted-custom-agent") + mgr._load_from_sqlite = AsyncMock(return_value=([], db_meta)) + fallback_agent = _make_agent_mock() + + def _load_agent_side_effect(name): + if name == "deleted-custom-agent": + raise ValueError(f"Unknown agent {name!r}") + return fallback_agent + + with ( + patch( + "code_puppy.api.session_context.load_agent", + side_effect=_load_agent_side_effect, + ), + patch( + "code_puppy.api.session_context.get_global_model_name", + return_value="gpt-4o", + ), + ): + ctx = await mgr.load_session("fallback-session") + + assert ctx is not None + assert ctx.agent_name == "code-puppy" + assert ctx.agent is fallback_agent + + +@pytest.mark.asyncio +async def test_load_session_valid_created_at_parsed(): + """A valid ISO created_at string in DB meta is parsed into a datetime.""" + mgr = _make_fresh_manager() + iso = "2026-03-15T09:30:00+00:00" + db_meta = _make_session_row(created_at=iso) + mgr._load_from_sqlite = AsyncMock(return_value=([], db_meta)) + + with ( + patch( + "code_puppy.api.session_context.load_agent", return_value=_make_agent_mock() + ), + patch( + "code_puppy.api.session_context.get_global_model_name", + return_value="gpt-4o", + ), + ): + ctx = await mgr.load_session("dated-session") + + assert ctx is not None + expected = datetime.fromisoformat(iso) + assert ctx.created_at == expected + + +@pytest.mark.asyncio +async def test_load_session_malformed_created_at_falls_back_to_now(): + """A malformed created_at string results in datetime.now() being used.""" + mgr = _make_fresh_manager() + db_meta = _make_session_row(created_at="not-a-date") + mgr._load_from_sqlite = AsyncMock(return_value=([], db_meta)) + before = datetime.now(timezone.utc) + + with ( + patch( + "code_puppy.api.session_context.load_agent", return_value=_make_agent_mock() + ), + patch( + "code_puppy.api.session_context.get_global_model_name", + return_value="gpt-4o", + ), + ): + ctx = await mgr.load_session("bad-date-session") + + after = datetime.now(timezone.utc) + assert ctx is not None + assert before <= ctx.created_at <= after, "created_at should be close to now()" + + +@pytest.mark.asyncio +async def test_load_session_empty_created_at_falls_back_to_now(): + """An empty created_at string also results in datetime.now() being used.""" + mgr = _make_fresh_manager() + db_meta = _make_session_row(created_at="") + mgr._load_from_sqlite = AsyncMock(return_value=([], db_meta)) + before = datetime.now(timezone.utc) + + with ( + patch( + "code_puppy.api.session_context.load_agent", return_value=_make_agent_mock() + ), + patch( + "code_puppy.api.session_context.get_global_model_name", + return_value="gpt-4o", + ), + ): + ctx = await mgr.load_session("empty-date-session") + + after = datetime.now(timezone.utc) + assert ctx is not None + assert before <= ctx.created_at <= after diff --git a/code_puppy/api/tests/test_session_model_compat.py b/code_puppy/api/tests/test_session_model_compat.py new file mode 100644 index 000000000..1309badd6 --- /dev/null +++ b/code_puppy/api/tests/test_session_model_compat.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +from datetime import datetime, timezone + +import pytest + +from code_puppy.api.session_context import ( + SessionContext, + SessionManager, + _apply_session_model, +) + + +class LegacyAgentWithoutSetter: + def __init__(self) -> None: + self._session_model_name = None + self.reload_count = 0 + + def reload_code_generation_agent(self) -> None: + self.reload_count += 1 + + +class ModernAgentWithSetter(LegacyAgentWithoutSetter): + def __init__(self) -> None: + super().__init__() + self.via_method = None + + def set_session_model(self, model_name): + self.via_method = model_name + + +@pytest.mark.parametrize( + "agent_cls,expected_attr,expected_via_method", + [ + (LegacyAgentWithoutSetter, "gpt-5.1", None), + (ModernAgentWithSetter, None, "gpt-5.1"), + ], +) +def test_apply_session_model_compat(agent_cls, expected_attr, expected_via_method): + agent = agent_cls() + _apply_session_model(agent, "gpt-5.1") + + if expected_attr is not None: + assert getattr(agent, "_session_model_name") == expected_attr + if expected_via_method is not None: + assert getattr(agent, "via_method") == expected_via_method + + +@pytest.mark.asyncio +async def test_switch_model_supports_legacy_agent_without_setter(): + mgr = SessionManager() + agent = LegacyAgentWithoutSetter() + ctx = SessionContext( + session_id="legacy-session", + agent=agent, + agent_name="code-puppy", + model_name="synthetic-GLM-5.1", + working_directory="", + created_at=datetime.now(timezone.utc), + ) + + mgr._sessions[ctx.session_id] = ctx + + await mgr.switch_model(ctx.session_id, "gpt-5.1") + + assert ctx.model_name == "gpt-5.1" + assert agent._session_model_name == "gpt-5.1" + assert agent.reload_count == 1 diff --git a/code_puppy/api/tests/test_session_persistence.py b/code_puppy/api/tests/test_session_persistence.py new file mode 100644 index 000000000..2a5b828b4 --- /dev/null +++ b/code_puppy/api/tests/test_session_persistence.py @@ -0,0 +1,172 @@ +"""Focused tests for session persistence payload builders (Phase 3).""" + +from __future__ import annotations + +from code_puppy.api.ws.session_persistence import ( + build_session_meta_payload, + build_session_update_payload, + resolve_agent_model_meta, +) + +# --------------------------------------------------------------------------- +# resolve_agent_model_meta +# --------------------------------------------------------------------------- + + +class _FakeAgent: + def __init__(self, name: str = "", model_name: str = ""): + self.name = name + self._model_name = model_name + + def get_model_name(self) -> str: + return self._model_name + + +class _FakeCtx: + def __init__(self, agent_name: str = "", model_name: str = ""): + self.agent_name = agent_name + self.model_name = model_name + + +def test_resolve_uses_agent_when_available(): + agent = _FakeAgent(name="husky", model_name="gpt-5") + name, model = resolve_agent_model_meta(agent=agent) + assert name == "husky" + assert model == "gpt-5" + + +def test_resolve_falls_back_to_ctx(): + agent = _FakeAgent(name="", model_name="") + ctx = _FakeCtx(agent_name="shepherd", model_name="claude-4") + name, model = resolve_agent_model_meta(agent=agent, ctx=ctx) + assert name == "shepherd" + assert model == "claude-4" + + +def test_resolve_falls_back_to_defaults(): + name, model = resolve_agent_model_meta() + assert name == "code-puppy" + assert model == "unknown" + + +def test_resolve_prefers_agent_over_ctx(): + agent = _FakeAgent(name="terrier", model_name="gpt-5") + ctx = _FakeCtx(agent_name="shepherd", model_name="claude-4") + name, model = resolve_agent_model_meta(agent=agent, ctx=ctx) + assert name == "terrier" + assert model == "gpt-5" + + +def test_resolve_agent_none_uses_ctx(): + ctx = _FakeCtx(agent_name="watchdog", model_name="o3") + name, model = resolve_agent_model_meta(agent=None, ctx=ctx) + assert name == "watchdog" + assert model == "o3" + + +# --------------------------------------------------------------------------- +# build_session_meta_payload +# --------------------------------------------------------------------------- + + +def test_session_meta_payload_contains_all_fields(): + payload = build_session_meta_payload( + session_id="WS_session_001", + session_name="WS_session_001", + total_tokens=42, + message_count=3, + title="test chat", + working_directory="/tmp", + agent_name="code-puppy", + model_name="gpt-5", + ) + assert payload["type"] == "session_meta" + assert payload["session_id"] == "WS_session_001" + assert payload["total_tokens"] == 42 + assert payload["message_count"] == 3 + assert payload["title"] == "test chat" + assert payload["working_directory"] == "/tmp" + assert payload["agent_name"] == "code-puppy" + assert payload["model_name"] == "gpt-5" + + +def test_session_meta_payload_has_no_extra_keys(): + payload = build_session_meta_payload( + session_id="s1", + session_name="s1", + total_tokens=0, + message_count=0, + title="", + working_directory="", + agent_name="", + model_name="", + ) + expected_keys = { + "type", + "session_id", + "session_name", + "total_tokens", + "message_count", + "title", + "working_directory", + "agent_name", + "model_name", + } + assert set(payload.keys()) == expected_keys + + +# --------------------------------------------------------------------------- +# build_session_update_payload +# --------------------------------------------------------------------------- + + +def test_session_update_action_created_for_first_message(): + payload = build_session_update_payload( + session_id="s1", + session_name="s1", + title="new chat", + working_directory="/home", + message_count=1, + total_tokens=10, + timestamp="2026-01-01T00:00:00", + ) + assert payload["action"] == "created" + assert payload["auto_saved"] is True + assert payload["timestamp"] == "2026-01-01T00:00:00" + + +def test_session_update_action_updated_for_subsequent_messages(): + payload = build_session_update_payload( + session_id="s1", + session_name="s1", + title="chat", + working_directory="", + message_count=5, + total_tokens=100, + ) + assert payload["action"] == "updated" + # timestamp auto-generated when not provided + assert "timestamp" in payload and payload["timestamp"] + + +def test_session_update_payload_has_expected_keys(): + payload = build_session_update_payload( + session_id="s1", + session_name="s1", + title="", + working_directory="", + message_count=2, + total_tokens=0, + ) + expected_keys = { + "session_id", + "session_name", + "title", + "working_directory", + "timestamp", + "message_count", + "total_tokens", + "auto_saved", + "action", + } + assert set(payload.keys()) == expected_keys 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/tool_formatters.py b/code_puppy/api/tool_formatters.py new file mode 100644 index 000000000..dc571818f --- /dev/null +++ b/code_puppy/api/tool_formatters.py @@ -0,0 +1,138 @@ +""" +Tool Call Argument Formatters + +Provides human-readable formatting for different tool types. +Used by permission system to display tool details in a user-friendly way. +""" + +from typing import Any, Dict + + +def format_shell_command(args: Dict[str, Any]) -> str: + """Format shell command arguments for display.""" + command = args.get("command", "") + cwd = args.get("cwd") + timeout = args.get("timeout", 60) + background = args.get("background", False) + + lines = [] + lines.append(f"Command: {command}") + if cwd: + lines.append(f"Directory: {cwd}") + lines.append(f"Timeout: {timeout}s") + if background: + lines.append("Mode: Background") + + return "\n".join(lines) + + +def format_file_operation(tool_name: str, args: Dict[str, Any]) -> str: + """Format file operation arguments for display.""" + file_path = args.get("file_path", "") + + lines = [] + if "read_file" in tool_name.lower(): + lines.append(f"Read file: {file_path}") + if "start_line" in args: + lines.append( + f"Lines: {args['start_line']}-{args['start_line'] + args.get('num_lines', 0)}" + ) + elif "edit_file" in tool_name.lower() or "write_file" in tool_name.lower(): + lines.append(f"Edit file: {file_path}") + content_preview = str(args.get("new_content", ""))[:100] + if len(content_preview) == 100: + content_preview += "..." + lines.append(f"Content preview: {content_preview}") + elif "delete_file" in tool_name.lower(): + lines.append(f"⚠️ DELETE file: {file_path}") + else: + lines.append(f"File: {file_path}") + + return "\n".join(lines) + + +def format_agent_invocation(args: Dict[str, Any]) -> str: + """Format agent invocation arguments for display.""" + agent_name = args.get("agent_name", "") + prompt = args.get("prompt", "") + session_id = args.get("session_id") + + lines = [] + lines.append(f"Agent: {agent_name}") + + # Truncate long prompts + prompt_preview = prompt[:200] + if len(prompt) > 200: + prompt_preview += "..." + lines.append(f"Prompt: {prompt_preview}") + + if session_id: + lines.append(f"Session: {session_id}") + + return "\n".join(lines) + + +def format_grep(args: Dict[str, Any]) -> str: + """Format grep arguments for display.""" + search_string = args.get("search_string", "") + directory = args.get("directory", ".") + + return f"Search: {search_string}\nDirectory: {directory}" + + +def format_list_files(args: Dict[str, Any]) -> str: + """Format list_files arguments for display.""" + directory = args.get("directory", ".") + recursive = args.get("recursive", True) + + mode = "Recursive" if recursive else "Non-recursive" + return f"List files: {directory}\nMode: {mode}" + + +def format_tool_call(tool_name: str, args: Dict[str, Any]) -> str: + """ + Format tool call arguments in a human-readable way. + + Args: + tool_name: Name of the tool being called + args: Tool arguments dict + + Returns: + Human-readable formatted string + """ + # Shell commands + if "shell" in tool_name.lower() or "command" in tool_name.lower(): + return format_shell_command(args) + + # File operations + if any( + word in tool_name.lower() + for word in ["file", "read", "write", "edit", "delete"] + ): + return format_file_operation(tool_name, args) + + # Agent invocations + if "agent" in tool_name.lower() or "invoke" in tool_name.lower(): + return format_agent_invocation(args) + + # Grep + if "grep" in tool_name.lower() or "search" in tool_name.lower(): + return format_grep(args) + + # List files + if "list" in tool_name.lower(): + return format_list_files(args) + + # Generic fallback - show all args + lines = [] + for key, value in args.items(): + value_str = str(value) + if len(value_str) > 100: + value_str = value_str[:100] + "..." + lines.append(f"{key}: {value_str}") + + return "\n".join(lines) if lines else "No arguments" + + +# Export main function +__all__ = ["format_tool_call"] 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/callbacks.py b/code_puppy/callbacks.py index 2041f9ba1..7dc972056 100644 --- a/code_puppy/callbacks.py +++ b/code_puppy/callbacks.py @@ -28,6 +28,10 @@ "post_tool_call", "stream_event", "thinking_display_filter", + "termflow_style", + "prompt_toolkit_style", + "termflow_highlighter", + "prompt_text_color", "register_tools", "register_agent_tools", "register_agents", @@ -87,6 +91,10 @@ "post_tool_call": [], "stream_event": [], "thinking_display_filter": [], + "termflow_style": [], + "prompt_toolkit_style": [], + "termflow_highlighter": [], + "prompt_text_color": [], "register_tools": [], "register_agent_tools": [], "register_agents": [], @@ -324,8 +332,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 +638,49 @@ 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_prompt_toolkit_style(default_style: Any = None) -> Any: + """Let plugins replace a prompt_toolkit Application style.""" + return _chain_value_callbacks("prompt_toolkit_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/chatgpt_codex_client.py b/code_puppy/chatgpt_codex_client.py index 7e00716e4..aff7206b1 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) @@ -211,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 9fb6d170d..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", @@ -593,6 +593,17 @@ def _persistent_prompt_parts() -> tuple: return ">>> ", [] +def _prompt_echo_text(task: str): + """Build a transcript echo using the active prompt foreground.""" + from rich.text import Text + + from code_puppy.callbacks import on_prompt_text_color + + prompt_color = on_prompt_text_color() + style = f"bold {prompt_color}" if prompt_color else "bold" + return Text(f"\n> {task}", style=style) + + def _interactive_sigint_guard(_sig, _frame): """Baseline SIGINT handler for the interactive REPL. @@ -667,7 +678,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 +689,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." @@ -841,8 +857,6 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non try: if persistent_prompt: - from rich.text import Text as _EchoText - from code_puppy.messaging.run_ui import ( set_idle_prompt_prefix, wait_for_idle_submission, @@ -862,7 +876,7 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non queued_task = _get_pc().pop_next_steer_queued() if queued_task is not None: task = queued_task - emit_info(_EchoText(f"\n> {task}", style="bold")) + emit_info(_prompt_echo_text(task)) emit_info("⏭ running queued prompt") else: # Raises EOFError on Ctrl+D-with-empty-buffer, which the @@ -874,7 +888,7 @@ async def interactive_mode(message_renderer, initial_command: str = None) -> Non # repeating the whole prompt chrome doubled every # line's noise. Text() (not markup) so bracket-y input # renders as-is. - emit_info(_EchoText(f"\n> {task}", style="bold")) + emit_info(_prompt_echo_text(task)) else: # Use prompt_toolkit for enhanced input with path completion try: @@ -1051,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/add_model_menu.py b/code_puppy/command_line/add_model_menu.py index 4f6ea6283..5c28bc887 100644 --- a/code_puppy/command_line/add_model_menu.py +++ b/code_puppy/command_line/add_model_menu.py @@ -35,6 +35,7 @@ save_credential, ) from code_puppy.tools.command_runner import set_awaiting_user_input +from code_puppy.callbacks import on_prompt_toolkit_style PAGE_SIZE = 15 # Items per page @@ -343,22 +344,24 @@ def _render_provider_list(self) -> List: """Render the provider list panel.""" lines = [] - lines.append(("", " Providers")) + lines.append(("class:tui.header", " Providers")) lines.append(("", "\n\n")) if not self.providers: - lines.append(("fg:yellow", " No providers available.")) + lines.append(("class:tui.warning", " No providers available.")) lines.append(("", "\n\n")) self._render_navigation_hints(lines) return lines filter_label = getattr(self, "provider_filter", "") or "type to filter" - lines.append(("fg:ansibrightblack", f" Filter: {filter_label}")) + lines.append(("class:tui.muted", f" Filter: {filter_label}")) lines.append(("", "\n\n")) filtered_providers = self._filtered_providers() if not filtered_providers: - lines.append(("fg:yellow", " No providers match the current filter.")) + lines.append( + ("class:tui.warning", " No providers match the current filter.") + ) lines.append(("", "\n\n")) self._render_navigation_hints(lines) return lines @@ -381,19 +384,18 @@ def _render_provider_list(self) -> List: suffix = " ⚠️" if is_unsupported else "" label = f"{prefix}{provider.name} ({provider.model_count} models){suffix}" - # Use dimmed color for unsupported providers - if is_unsupported: - lines.append(("fg:ansibrightblack dim", label)) - elif is_selected: - lines.append(("fg:ansibrightblack", label)) + if is_selected: + lines.append(("class:tui.selected", label)) + elif is_unsupported: + lines.append(("class:tui.muted", label)) else: - lines.append(("fg:ansibrightblack", label)) + lines.append(("class:tui.body", label)) lines.append(("", "\n")) lines.append(("", "\n")) lines.append( - ("fg:ansibrightblack", f" Page {self.current_page + 1}/{total_pages}") + ("class:tui.muted", f" Page {self.current_page + 1}/{total_pages}") ) lines.append(("", "\n")) @@ -405,21 +407,21 @@ def _render_model_list(self) -> List: lines = [] if not self.current_provider: - lines.append(("fg:yellow", " No provider selected.")) + lines.append(("class:tui.warning", " No provider selected.")) lines.append(("", "\n\n")) self._render_navigation_hints(lines) return lines - lines.append(("", f" {self.current_provider.name} Models")) + lines.append(("class:tui.header", f" {self.current_provider.name} Models")) lines.append(("", "\n")) filter_label = getattr(self, "model_filter", "") or "type to filter" - lines.append(("fg:ansibrightblack", f" Filter: {filter_label}")) + lines.append(("class:tui.muted", f" Filter: {filter_label}")) lines.append(("", "\n\n")) filtered_models = self._filtered_models() custom_visible = self._should_show_custom_model() if not filtered_models and not custom_visible: - lines.append(("fg:yellow", " No models match the current filter.")) + lines.append(("class:tui.warning", " No models match the current filter.")) lines.append(("", "\n\n")) self._render_navigation_hints(lines) return lines @@ -435,9 +437,9 @@ def _render_model_list(self) -> List: if custom_visible and i == len(filtered_models): is_selected = i == self.selected_model_idx if is_selected: - lines.append(("fg:ansicyan bold", " > ✨ Custom model...")) + lines.append(("class:tui.selected", " > ✨ Custom model...")) else: - lines.append(("fg:ansicyan", " ✨ Custom model...")) + lines.append(("class:tui.body", " ✨ Custom model...")) lines.append(("", "\n")) continue @@ -456,15 +458,15 @@ def _render_model_list(self) -> List: icon_str = " ".join(icons) + " " if icons else "" if is_selected: - lines.append(("fg:ansibrightblack", f" > {icon_str}{model.name}")) + lines.append(("class:tui.selected", f" > {icon_str}{model.name}")) else: - lines.append(("fg:ansibrightblack", f" {icon_str}{model.name}")) + lines.append(("class:tui.body", f" {icon_str}{model.name}")) lines.append(("", "\n")) lines.append(("", "\n")) lines.append( - ("fg:ansibrightblack", f" Page {self.current_page + 1}/{total_pages}") + ("class:tui.muted", f" Page {self.current_page + 1}/{total_pages}") ) lines.append(("", "\n")) @@ -474,61 +476,63 @@ def _render_model_list(self) -> List: def _render_navigation_hints(self, lines: List): """Render navigation hints at the bottom of the list panel.""" lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " ↑/↓ ")) - lines.append(("", "Navigate ")) - lines.append(("fg:ansibrightblack", "←/→ ")) - lines.append(("", "Page\n")) - lines.append(("fg:ansibrightblack", " Type ")) - lines.append(("", "Filter list\n")) - lines.append(("fg:ansibrightblack", " Backspace ")) - lines.append(("", "Delete filter char\n")) - lines.append(("fg:ansibrightblack", " Ctrl+U ")) - lines.append(("", "Clear filter\n")) + lines.append(("class:tui.help-key", " ↑/↓ ")) + lines.append(("class:tui.help", "Navigate ")) + lines.append(("class:tui.help-key", "←/→ ")) + lines.append(("class:tui.help", "Page\n")) + lines.append(("class:tui.help-key", " Type ")) + lines.append(("class:tui.help", "Filter list\n")) + lines.append(("class:tui.help-key", " Backspace ")) + lines.append(("class:tui.help", "Delete filter char\n")) + lines.append(("class:tui.help-key", " Ctrl+U ")) + lines.append(("class:tui.help", "Clear filter\n")) if self.view_mode == "providers": - lines.append(("fg:green", " Enter ")) - lines.append(("", "Select\n")) - lines.append(("fg:cyan", " Ctrl+E ")) - lines.append(("", "Edit credentials\n")) + lines.append(("class:tui.help-key", " Enter ")) + lines.append(("class:tui.help", "Select\n")) + lines.append(("class:tui.help-key", " Ctrl+E ")) + lines.append(("class:tui.help", "Edit credentials\n")) else: - lines.append(("fg:green", " Enter ")) - lines.append(("", "Add Model\n")) - lines.append(("fg:ansibrightblack", " Esc/Back ")) - lines.append(("", "Back\n")) - lines.append(("fg:ansibrightred", " Ctrl+C ")) - lines.append(("", "Cancel")) + lines.append(("class:tui.help-key", " Enter ")) + lines.append(("class:tui.help", "Add Model\n")) + lines.append(("class:tui.help-key", " Esc/Back ")) + lines.append(("class:tui.help", "Back\n")) + lines.append(("class:tui.help-key", " Ctrl+C ")) + lines.append(("class:tui.help", "Cancel")) def _render_model_details(self) -> List: """Render the model details panel.""" lines = [] - lines.append(("dim cyan", " MODEL DETAILS")) + lines.append(("class:tui.title", " MODEL DETAILS")) lines.append(("", "\n\n")) if self.view_mode == "providers": provider = self._get_current_provider() if not provider: - lines.append(("fg:yellow", " No provider selected.")) + lines.append(("class:tui.warning", " No provider selected.")) return lines - lines.append(("bold", f" {provider.name}")) + lines.append(("class:tui.label", f" {provider.name}")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", f" ID: {provider.id}")) + lines.append(("class:tui.body", f" ID: {provider.id}")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", f" Models: {provider.model_count}")) + lines.append(("class:tui.body", f" Models: {provider.model_count}")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", f" API: {provider.api}")) + lines.append(("class:tui.body", f" API: {provider.api}")) lines.append(("", "\n")) # Show unsupported warning if applicable if provider.id in UNSUPPORTED_PROVIDERS: lines.append(("", "\n")) - lines.append(("fg:ansired bold", " ⚠️ UNSUPPORTED PROVIDER")) + lines.append(("class:tui.error", " ⚠️ UNSUPPORTED PROVIDER")) lines.append(("", "\n")) - lines.append(("fg:ansired", f" {UNSUPPORTED_PROVIDERS[provider.id]}")) + lines.append( + ("class:tui.error", f" {UNSUPPORTED_PROVIDERS[provider.id]}") + ) lines.append(("", "\n")) lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", " Models from this provider cannot be added.", ) ) @@ -536,27 +540,27 @@ def _render_model_details(self) -> List: if provider.env: lines.append(("", "\n")) - lines.append(("bold", " Credentials:")) + lines.append(("class:tui.label", " Credentials:")) lines.append(("", "\n")) for env_var in provider.env: status = credential_display(env_var) hint = credential_hint(env_var) color = ( - "fg:ansigreen" + "class:tui.success" if is_credential_set(env_var) - else "fg:ansiyellow" + else "class:tui.warning" ) lines.append((color, f" • {env_var}: {status}")) lines.append(("", "\n")) if hint: - lines.append(("fg:ansibrightblack", f" {hint}")) + lines.append(("class:tui.muted", f" {hint}")) lines.append(("", "\n")) if provider.doc: lines.append(("", "\n")) - lines.append(("bold", " Documentation:")) + lines.append(("class:tui.label", " Documentation:")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", f" {provider.doc}")) + lines.append(("class:tui.body", f" {provider.doc}")) lines.append(("", "\n")) else: # models view @@ -564,67 +568,69 @@ def _render_model_details(self) -> List: provider = self.current_provider if not provider: - lines.append(("fg:yellow", " No model selected.")) + lines.append(("class:tui.warning", " No model selected.")) return lines # Handle custom model option if self._is_custom_model_selected(): - lines.append(("bold", " ✨ Custom Model")) + lines.append(("class:tui.label", " ✨ Custom Model")) lines.append(("", "\n\n")) - lines.append(("fg:ansicyan", " Add a model not listed in models.dev")) + lines.append( + ("class:tui.body", " Add a model not listed in models.dev") + ) lines.append(("", "\n\n")) - lines.append(("bold", " How it works:")) + lines.append(("class:tui.label", " How it works:")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " 1. Press Enter to select")) + lines.append(("class:tui.muted", " 1. Press Enter to select")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " 2. Enter the model ID/name")) + lines.append(("class:tui.muted", " 2. Enter the model ID/name")) lines.append(("", "\n")) lines.append( - ("fg:ansibrightblack", f" 3. Uses {provider.name}'s API endpoint") + ("class:tui.muted", f" 3. Uses {provider.name}'s API endpoint") ) lines.append(("", "\n\n")) - lines.append(("bold", " Use cases:")) + lines.append(("class:tui.label", " Use cases:")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " • Newly released models")) + lines.append(("class:tui.muted", " • Newly released models")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " • Fine-tuned models")) + lines.append(("class:tui.muted", " • Fine-tuned models")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " • Preview/beta models")) + lines.append(("class:tui.muted", " • Preview/beta models")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " • Custom deployments")) + lines.append(("class:tui.muted", " • Custom deployments")) lines.append(("", "\n\n")) if provider.env: - lines.append(("bold", " Required credentials:")) + lines.append(("class:tui.label", " Required credentials:")) lines.append(("", "\n")) for env_var in provider.env: - lines.append(("fg:ansibrightblack", f" • {env_var}")) + lines.append(("class:tui.muted", f" • {env_var}")) lines.append(("", "\n")) return lines if not model: - lines.append(("fg:yellow", " No model selected.")) + lines.append(("class:tui.warning", " No model selected.")) return lines - lines.append(("bold", f" {provider.name} - {model.name}")) + lines.append(("class:tui.label", f" {provider.name} - {model.name}")) lines.append(("", "\n\n")) # BIG WARNING for models without tool calling if not model.tool_call: - lines.append(("fg:ansiyellow bold", " ⚠️ NO TOOL CALLING SUPPORT")) + lines.append(("class:tui.warning", " ⚠️ NO TOOL CALLING SUPPORT")) lines.append(("", "\n")) lines.append( - ("fg:ansiyellow", " This model cannot use tools (file ops,") + ("class:tui.warning", " This model cannot use tools (file ops,") ) lines.append(("", "\n")) lines.append( - ("fg:ansiyellow", " shell commands, etc). It will be very") + ("class:tui.warning", " shell commands, etc). It will be very") ) lines.append(("", "\n")) - lines.append(("fg:ansiyellow", " limited for coding tasks!")) + lines.append(("class:tui.warning", " limited for coding tasks!")) lines.append(("", "\n\n")) # Capabilities - lines.append(("bold", " Capabilities:")) + lines.append(("class:tui.label", " Capabilities:")) lines.append(("", "\n")) capabilities = [ @@ -638,21 +644,21 @@ def _render_model_details(self) -> List: for cap_name, has_cap in capabilities: if has_cap: - lines.append(("fg:green", f" ✓ {cap_name}")) + lines.append(("class:tui.success", f" ✓ {cap_name}")) else: - lines.append(("fg:ansibrightblack", f" ✗ {cap_name}")) + lines.append(("class:tui.muted", f" ✗ {cap_name}")) lines.append(("", "\n")) # Pricing lines.append(("", "\n")) - lines.append(("bold", " Pricing:")) + lines.append(("class:tui.label", " Pricing:")) lines.append(("", "\n")) if model.cost_input is not None or model.cost_output is not None: if model.cost_input is not None: lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", f" Input: ${model.cost_input:.6f}/token", ) ) @@ -660,7 +666,7 @@ def _render_model_details(self) -> List: if model.cost_output is not None: lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", f" Output: ${model.cost_output:.6f}/token", ) ) @@ -668,24 +674,24 @@ def _render_model_details(self) -> List: if model.cost_cache_read is not None: lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", f" Cache Read: ${model.cost_cache_read:.6f}/token", ) ) lines.append(("", "\n")) else: - lines.append(("fg:ansibrightblack", " Pricing not available")) + lines.append(("class:tui.muted", " Pricing not available")) lines.append(("", "\n")) # Limits lines.append(("", "\n")) - lines.append(("bold", " Limits:")) + lines.append(("class:tui.label", " Limits:")) lines.append(("", "\n")) if model.context_length > 0: lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", f" Context: {model.context_length:,} tokens", ) ) @@ -693,7 +699,7 @@ def _render_model_details(self) -> List: if model.max_output > 0: lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", f" Max Output: {model.max_output:,} tokens", ) ) @@ -702,13 +708,13 @@ def _render_model_details(self) -> List: # Modalities if model.input_modalities or model.output_modalities: lines.append(("", "\n")) - lines.append(("bold", " Modalities:")) + lines.append(("class:tui.label", " Modalities:")) lines.append(("", "\n")) if model.input_modalities: lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", f" Input: {', '.join(model.input_modalities)}", ) ) @@ -716,7 +722,7 @@ def _render_model_details(self) -> List: if model.output_modalities: lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", f" Output: {', '.join(model.output_modalities)}", ) ) @@ -724,29 +730,23 @@ def _render_model_details(self) -> List: # Metadata lines.append(("", "\n")) - lines.append(("bold", " Metadata:")) + lines.append(("class:tui.label", " Metadata:")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", f" Model ID: {model.model_id}")) + lines.append(("class:tui.muted", f" Model ID: {model.model_id}")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", f" Full ID: {model.full_id}")) + lines.append(("class:tui.muted", f" Full ID: {model.full_id}")) lines.append(("", "\n")) if model.knowledge: - lines.append( - ("fg:ansibrightblack", f" Knowledge: {model.knowledge}") - ) + lines.append(("class:tui.muted", f" Knowledge: {model.knowledge}")) lines.append(("", "\n")) if model.release_date: - lines.append( - ("fg:ansibrightblack", f" Released: {model.release_date}") - ) + lines.append(("class:tui.muted", f" Released: {model.release_date}")) lines.append(("", "\n")) - lines.append( - ("fg:ansibrightblack", f" Open Weights: {model.open_weights}") - ) + lines.append(("class:tui.muted", f" Open Weights: {model.open_weights}")) lines.append(("", "\n")) return lines @@ -1307,6 +1307,7 @@ def _(event): key_bindings=kb, full_screen=False, mouse_support=False, + style=on_prompt_toolkit_style(), ) # Run application in a background thread to avoid event loop conflicts diff --git a/code_puppy/command_line/agent_menu.py b/code_puppy/command_line/agent_menu.py index 420d80fc1..d5f074032 100644 --- a/code_puppy/command_line/agent_menu.py +++ b/code_puppy/command_line/agent_menu.py @@ -42,6 +42,7 @@ ) from code_puppy.messaging import emit_info, emit_success, emit_warning from code_puppy.tools.command_runner import set_awaiting_user_input +from code_puppy.callbacks import on_prompt_toolkit_style PAGE_SIZE = 10 # Agents per page @@ -315,12 +316,12 @@ def _render_menu_panel( total_pages = get_total_pages(len(entries), PAGE_SIZE) start_idx, end_idx = get_page_bounds(page, len(entries), PAGE_SIZE) - lines.append(("bold", "Agents")) - lines.append(("fg:ansibrightblack", f" (Page {page + 1}/{total_pages})")) + lines.append(("class:tui.header", "Agents")) + lines.append(("class:tui.muted", f" (Page {page + 1}/{total_pages})")) lines.append(("", "\n\n")) if not entries: - lines.append(("fg:yellow", " No agents found.")) + lines.append(("class:tui.warning", " No agents found.")) lines.append(("", "\n\n")) else: # Show agents for current page @@ -335,40 +336,38 @@ def _render_menu_panel( # Build the line if is_selected: - lines.append(("fg:ansigreen", "▶ ")) - lines.append(("fg:ansigreen bold", safe_display_name)) + lines.append(("class:tui.selected", f"▶ {safe_display_name}")) else: - lines.append(("", " ")) - lines.append(("", safe_display_name)) + lines.append(("class:tui.body", f" {safe_display_name}")) if pinned_model: safe_pinned_model = _sanitize_display_text(pinned_model) - lines.append(("fg:ansiyellow", f" → {safe_pinned_model}")) + lines.append(("class:tui.label", f" → {safe_pinned_model}")) # Add current marker if is_current: - lines.append(("fg:ansicyan", " ← current")) + lines.append(("class:tui.success", " ← current")) lines.append(("", "\n")) # Navigation hints lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " ↑↓ ")) - lines.append(("", "Navigate\n")) - lines.append(("fg:ansibrightblack", " ←→ ")) - lines.append(("", "Page\n")) - lines.append(("fg:green", " Enter ")) - lines.append(("", "Select\n")) - lines.append(("fg:ansibrightblack", " P ")) - lines.append(("", "Pin model\n")) - lines.append(("fg:ansibrightblack", " B ")) - lines.append(("", "Bind MCP servers\n")) - lines.append(("fg:ansibrightblack", " C ")) - lines.append(("", "Clone\n")) - lines.append(("fg:ansibrightblack", " D ")) - lines.append(("", "Delete clone\n")) - lines.append(("fg:ansibrightred", " Ctrl+C ")) - lines.append(("", "Cancel")) + lines.append(("class:tui.help-key", " ↑↓ ")) + lines.append(("class:tui.help", "Navigate\n")) + lines.append(("class:tui.help-key", " ←→ ")) + lines.append(("class:tui.help", "Page\n")) + lines.append(("class:tui.help-key", " Enter ")) + lines.append(("class:tui.help", "Select\n")) + lines.append(("class:tui.help-key", " P ")) + lines.append(("class:tui.help", "Pin model\n")) + lines.append(("class:tui.help-key", " B ")) + lines.append(("class:tui.help", "Bind MCP servers\n")) + lines.append(("class:tui.help-key", " C ")) + lines.append(("class:tui.help", "Clone\n")) + lines.append(("class:tui.help-key", " D ")) + lines.append(("class:tui.help", "Delete clone\n")) + lines.append(("class:tui.help-key", " Ctrl+C ")) + lines.append(("class:tui.help", "Cancel")) return lines @@ -388,11 +387,11 @@ def _render_preview_panel( """ lines = [] - lines.append(("dim cyan", " AGENT DETAILS")) + lines.append(("class:tui.title", " AGENT DETAILS")) lines.append(("", "\n\n")) if not entry: - lines.append(("fg:yellow", " No agent selected.")) + lines.append(("class:tui.warning", " No agent selected.")) lines.append(("", "\n")) return lines @@ -405,22 +404,22 @@ def _render_preview_panel( safe_description = _sanitize_display_text(description) # Agent name (identifier) - lines.append(("bold", "Name: ")) + lines.append(("class:tui.label", "Name: ")) lines.append(("", name)) lines.append(("", "\n\n")) # Display name - lines.append(("bold", "Display Name: ")) - lines.append(("fg:ansicyan", safe_display_name)) + lines.append(("class:tui.label", "Display Name: ")) + lines.append(("class:tui.body", safe_display_name)) lines.append(("", "\n\n")) # Pinned model - lines.append(("bold", "Pinned Model: ")) + lines.append(("class:tui.label", "Pinned Model: ")) if pinned_model: safe_pinned_model = _sanitize_display_text(pinned_model) - lines.append(("fg:ansiyellow", safe_pinned_model)) + lines.append(("class:tui.body", safe_pinned_model)) else: - lines.append(("fg:ansibrightblack", "default")) + lines.append(("class:tui.muted", "default")) lines.append(("", "\n\n")) # MCP bindings summary @@ -428,19 +427,19 @@ def _render_preview_panel( bound = get_bound_servers(name) except Exception: bound = {} - lines.append(("bold", "MCP Servers: ")) + lines.append(("class:tui.label", "MCP Servers: ")) if bound: auto_count = sum(1 for opts in bound.values() if opts.get("auto_start")) summary = f"{len(bound)} bound" if auto_count: summary += f" ({auto_count} auto-start)" - lines.append(("fg:ansigreen", summary)) + lines.append(("class:tui.success", summary)) else: - lines.append(("fg:ansibrightblack", "none bound (strict opt-in)")) + lines.append(("class:tui.muted", "none bound (strict opt-in)")) lines.append(("", "\n\n")) # Description - lines.append(("bold", "Description:")) + lines.append(("class:tui.label", "Description:")) lines.append(("", "\n")) # Wrap description to fit panel @@ -451,7 +450,7 @@ def _render_preview_panel( current_line = "" for word in words: if len(current_line) + len(word) + 1 > 55: - lines.append(("fg:ansibrightblack", current_line)) + lines.append(("class:tui.body", current_line)) lines.append(("", "\n")) current_line = word else: @@ -460,17 +459,17 @@ def _render_preview_panel( else: current_line += " " + word if current_line.strip(): - lines.append(("fg:ansibrightblack", current_line)) + lines.append(("class:tui.body", current_line)) lines.append(("", "\n")) lines.append(("", "\n")) # Current status - lines.append(("bold", " Status: ")) + lines.append(("class:tui.label", " Status: ")) if is_current: - lines.append(("fg:ansigreen bold", "✓ Currently Active")) + lines.append(("class:tui.success", "✓ Currently Active")) else: - lines.append(("fg:ansibrightblack", "Not active")) + lines.append(("class:tui.muted", "Not active")) lines.append(("", "\n")) return lines @@ -639,6 +638,7 @@ def _(event): key_bindings=kb, full_screen=False, mouse_support=False, + style=on_prompt_toolkit_style(), ) set_awaiting_user_input(True) 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..65000e316 100644 --- a/code_puppy/command_line/autosave_menu.py +++ b/code_puppy/command_line/autosave_menu.py @@ -31,6 +31,7 @@ get_page_for_index, get_total_pages, ) +from code_puppy.callbacks import on_prompt_toolkit_style from code_puppy.config import AUTOSAVE_DIR from code_puppy.session_storage import list_sessions, load_session from code_puppy.tools.command_runner import set_awaiting_user_input @@ -135,11 +136,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 @@ -192,27 +191,27 @@ def _render_menu_panel( lines.append(status_line) elif in_search_mode: lines.append(("", "\n")) - lines.append(("fg:ansiyellow", f" Searching: '{search_buffer}'")) + lines.append(("class:tui.input.focused", f" Searching: '{search_buffer}'")) elif search_text: lines.append(("", "\n")) - lines.append(("fg:ansiyellow", f" Filter: '{search_text}'")) + lines.append(("class:tui.warning", f" Filter: '{search_text}'")) lines.append(("", "\n\n")) if not entries: if search_text or in_search_mode: - lines.append(("fg:yellow", " No sessions match your search.")) + lines.append(("class:tui.warning", " No sessions match your search.")) else: - lines.append(("fg:yellow", " No autosave sessions found.")) + lines.append(("class:tui.warning", " No autosave sessions found.")) lines.append(("", "\n\n")) # Navigation hints (always show) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " ↑/↓ ")) + lines.append(("class:tui.muted", " ↑/↓ ")) lines.append(("", "Navigate\n")) - lines.append(("fg:ansibrightblack", " ←/→ ")) + lines.append(("class:tui.muted", " ←/→ ")) lines.append(("", "Page\n")) - lines.append(("fg:green", " Enter ")) + lines.append(("class:tui.help-key", " Enter ")) lines.append(("", "Load\n")) - lines.append(("fg:ansibrightred", " Ctrl+C ")) + lines.append(("class:tui.help-key", " Ctrl+C ")) lines.append(("", "Cancel")) return lines @@ -238,31 +237,31 @@ def _render_menu_panel( # Highlight selected item if is_selected: - lines.append(("fg:ansibrightblack", f" > {label}")) + lines.append(("class:tui.selected", f" > {label}")) else: - lines.append(("fg:ansibrightblack", f" {label}")) + lines.append(("class:tui.muted", f" {label}")) lines.append(("", "\n")) # Navigation hints - change based on browse mode lines.append(("", "\n")) if browse_mode: - lines.append(("fg:ansicyan", " ↑/↓ ")) + lines.append(("class:tui.help-key", " ↑/↓ ")) lines.append(("", "Browse msgs\n")) - lines.append(("fg:ansiyellow", " Esc ")) + lines.append(("class:tui.help-key", " Esc ")) lines.append(("", "Exit browser\n")) else: - lines.append(("fg:ansibrightblack", " ↑/↓ ")) + lines.append(("class:tui.muted", " ↑/↓ ")) lines.append(("", "Navigate\n")) - lines.append(("fg:ansibrightblack", " ←/→ ")) + lines.append(("class:tui.muted", " ←/→ ")) lines.append(("", "Page\n")) - lines.append(("fg:ansicyan", " e ")) + lines.append(("class:tui.help-key", " e ")) lines.append(("", "Browse msgs\n")) - lines.append(("fg:ansibrightblack", " / ")) + lines.append(("class:tui.help-key", " / ")) lines.append(("", "Search content\n")) - lines.append(("fg:green", " Enter ")) + lines.append(("class:tui.help-key", " Enter ")) lines.append(("", "Load\n")) - lines.append(("fg:ansibrightred", " Ctrl+C ")) + lines.append(("class:tui.help-key", " Ctrl+C ")) lines.append(("", "Cancel")) return lines @@ -282,12 +281,12 @@ def _render_message_browser_panel( """ lines = [] - lines.append(("fg:ansicyan bold", " MESSAGE BROWSER")) + lines.append(("class:tui.header", " MESSAGE BROWSER")) lines.append(("", "\n\n")) total_messages = len(history) if total_messages == 0: - lines.append(("fg:yellow", " No messages in this session.")) + lines.append(("class:tui.warning", " No messages in this session.")) lines.append(("", "\n")) return lines @@ -302,25 +301,25 @@ def _render_message_browser_panel( role, content = _extract_message_content(msg) # Session info - lines.append(("fg:ansibrightblack", f" Session: {session_name}")) + lines.append(("class:tui.muted", f" Session: {session_name}")) lines.append(("", "\n")) # Message position indicator display_num = message_idx + 1 # 1-based for display - lines.append(("bold", f" Message {display_num} of {total_messages}")) + lines.append(("class:tui.label", f" Message {display_num} of {total_messages}")) lines.append(("", "\n\n")) # Role indicator with icon and color if role == "user": - lines.append(("fg:ansicyan bold", " 🧑 USER")) + lines.append(("class:tui.title", " \U0001f9d1 USER")) elif role == "tool": - lines.append(("fg:ansiyellow bold", " 🔧 TOOL")) + lines.append(("class:tui.warning", " TOOL")) else: - lines.append(("fg:ansigreen bold", " 🤖 ASSISTANT")) + lines.append(("class:tui.success", " \U0001f916 ASSISTANT")) lines.append(("", "\n")) # Separator line - lines.append(("fg:ansibrightblack", " " + "─" * 40)) + lines.append(("class:tui.muted", " " + "─" * 40)) lines.append(("", "\n")) # Render content - use markdown for user/assistant, plain text for tool @@ -329,7 +328,7 @@ def _render_message_browser_panel( # Tool messages are already formatted, don't pass through markdown # Use yellow color for tool output rendered = content - text_color = "fg:ansiyellow" + text_color = "class:tui.warning" else: # User and assistant messages should be rendered as markdown # Rich will handle the styling via ANSI codes @@ -354,12 +353,12 @@ def _render_message_browser_panel( lines.append(("", "\n")) except Exception as e: - lines.append(("fg:red", f" Error rendering message: {e}")) + lines.append(("class:tui.error", f" Error rendering message: {e}")) lines.append(("", "\n")) # Navigation hint at bottom lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " ↑ older ↓ newer Esc exit")) + lines.append(("class:tui.help", " ↑ older ↓ newer Esc exit")) lines.append(("", "\n")) return lines @@ -369,18 +368,18 @@ def _render_preview_panel(base_dir: Path, entry: Optional[Tuple[str, dict]]) -> """Render the right preview panel with message content using rich markdown.""" lines = [] - lines.append(("dim cyan", " PREVIEW")) + lines.append(("class:tui.title", " PREVIEW")) lines.append(("", "\n\n")) if not entry: - lines.append(("fg:yellow", " No session selected.")) + lines.append(("class:tui.warning", " No session selected.")) lines.append(("", "\n")) return lines session_name, metadata = entry # Show metadata - lines.append(("bold", " Session: ")) + lines.append(("class:tui.label", " Session: ")) lines.append(("", session_name)) lines.append(("", "\n")) @@ -390,18 +389,16 @@ def _render_preview_panel(base_dir: Path, entry: Optional[Tuple[str, dict]]) -> time_str = dt.strftime("%Y-%m-%d %H:%M:%S") except Exception: time_str = timestamp - lines.append(("fg:ansibrightblack", f" Saved: {time_str}")) + lines.append(("class:tui.muted", f" Saved: {time_str}")) lines.append(("", "\n")) msg_count = metadata.get("message_count", 0) tokens = metadata.get("total_tokens", 0) - lines.append( - ("fg:ansibrightblack", f" Messages: {msg_count} • Tokens: {tokens:,}") - ) + lines.append(("class:tui.muted", f" Messages: {msg_count} • Tokens: {tokens:,}")) lines.append(("", "\n\n")) - lines.append(("bold", " Last Message:")) - lines.append(("fg:ansibrightblack", " (press 'e' to browse full history)")) + lines.append(("class:tui.label", " Last Message:")) + lines.append(("class:tui.muted", " (press 'e' to browse full history)")) lines.append(("", "\n")) # Try to load and preview the last message @@ -426,11 +423,11 @@ def _render_preview_panel(base_dir: Path, entry: Optional[Tuple[str, dict]]) -> for line in message_lines: # Rich already rendered the markdown, just display it dimmed - lines.append(("fg:ansibrightblack", f" {line}")) + lines.append(("class:tui.muted", f" {line}")) lines.append(("", "\n")) except Exception as e: - lines.append(("fg:red", f" Error loading preview: {e}")) + lines.append(("class:tui.error", f" Error loading preview: {e}")) lines.append(("", "\n")) return lines @@ -457,10 +454,10 @@ def display_resumed_history( Configurable via: /set resume_message_count=50 """ from rich.console import Console - from rich.markdown import Markdown from rich.rule import Rule from code_puppy.config import get_banner_color, get_resume_message_count + from code_puppy.tools.display import render_markdown if not history: return @@ -523,9 +520,8 @@ def display_resumed_history( # Use the exact same banner format as normal AGENT RESPONSE banner = f"[bold white on {response_color}] AGENT RESPONSE [/bold white on {response_color}]" console.print(f"\n{banner}") - # Render content as markdown (same as normal chat) - md = Markdown(content) - console.print(md) + # Resume uses the same Termflow pipeline as live agent output. + render_markdown(content, console) console.print() # Blank line between messages @@ -621,13 +617,13 @@ def _compute_status_line() -> Optional[Tuple[str, str]]: other search activity to crowd out. """ if is_filtering[0]: - return ("fg:ansicyan bold", " Filtering...") + return ("class:tui.input.focused", " Filtering...") if in_search_mode[0] or search_text[0]: return None # Let the renderer show the search/filter line. cached = content_index.count() if 0 < cached < total_to_index: return ( - "fg:ansibrightblack", + "class:tui.muted", f" Indexing {cached}/{total_to_index}...", ) return None @@ -853,6 +849,7 @@ def _(event): layout = Layout(root_container) app = Application( + style=on_prompt_toolkit_style(), layout=layout, key_bindings=kb, full_screen=False, 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..2920c607e 100644 --- a/code_puppy/command_line/colors_menu.py +++ b/code_puppy/command_line/colors_menu.py @@ -18,47 +18,48 @@ from prompt_toolkit.layout.controls import FormattedTextControl from prompt_toolkit.widgets import Frame from rich.console import Console +from code_puppy.callbacks import on_prompt_toolkit_style -# 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 @@ -269,21 +270,21 @@ def get_left_panel_text(): """Generate the selector menu text.""" try: lines = [] - lines.append(("bold cyan", title)) + lines.append(("class:tui.header", title)) lines.append(("", "\n\n")) if not choices: - lines.append(("fg:ansiyellow", "No choices available")) + lines.append(("class:tui.warning", "No choices available")) lines.append(("", "\n")) else: for i, choice in enumerate(choices): # Skip separator lines for selection highlighting if "───" in choice: - lines.append(("fg:ansigray", f" {choice}")) + lines.append(("class:tui.muted", f" {choice}")) lines.append(("", "\n")) elif i == selected_index[0]: - lines.append(("fg:ansigreen", "▶ ")) - lines.append(("fg:ansigreen bold", choice)) + lines.append(("class:tui.selected", "▶ ")) + lines.append(("class:tui.selected", choice)) lines.append(("", "\n")) else: lines.append(("", " ")) @@ -292,11 +293,11 @@ def get_left_panel_text(): lines.append(("", "\n")) lines.append( - ("fg:ansicyan", "↑↓ Navigate │ Enter Select │ Ctrl-C Cancel") + ("class:tui.help-key", "↑↓ Navigate │ Enter Select │ Ctrl-C Cancel") ) return FormattedText(lines) except Exception as e: - return FormattedText([("fg:ansired", f"Error: {e}")]) + return FormattedText([("class:tui.error", f"Error: {e}")]) def get_right_panel_text(): """Generate the preview panel text.""" @@ -304,7 +305,7 @@ def get_right_panel_text(): preview = get_preview() return preview except Exception as e: - return FormattedText([("fg:ansired", f"Preview error: {e}")]) + return FormattedText([("class:tui.error", f"Preview error: {e}")]) kb = KeyBindings() @@ -370,6 +371,7 @@ def cancel(event): full_screen=False, mouse_support=False, color_depth="DEPTH_24_BIT", + style=on_prompt_toolkit_style(), ) sys.stdout.flush() diff --git a/code_puppy/command_line/config_apply.py b/code_puppy/command_line/config_apply.py index b1fd397c5..b3c6ff84a 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. @@ -50,10 +58,14 @@ def invalidate_post_write_caches(key: str) -> None: Called from both the slash-set path (:func:`apply_setting`) and the menu reset path so the two stay in lock-step. - Today only the ``model`` key has a cache that needs this treatment; - add new keys here as they grow caches. + Runtime CLI overrides are also cleared when their persisted setting is + explicitly changed, so an in-session ``/set`` always wins. """ - if key == "model": + if key == "yolo_mode": + from code_puppy.config import set_cli_yolo_override + + set_cli_yolo_override(None) + elif key == "model": from code_puppy.config import clear_model_cache, reset_session_model reset_session_model() @@ -85,6 +97,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 @@ -108,7 +125,12 @@ def apply_setting( warning = _restart_notice("DBOS configuration") requires_restart = True - set_config_value(key, normalized_value) + if key == "yolo_mode" and normalized_value.strip().lower() == "config": + # ``config`` only drops the process-local CLI override. The persisted + # value remains the source of truth once that override is gone. + normalized_value = "" + else: + set_config_value(key, normalized_value) invalidate_post_write_caches(key) reload_error: Optional[str] = None diff --git a/code_puppy/command_line/config_commands.py b/code_puppy/command_line/config_commands.py index 1777b27f8..ad2b925ae 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", @@ -215,7 +135,10 @@ def handle_set_command(command: str) -> bool: if is_sensitive_key(key) else result.value_after ) - emit_success(f'Set {key} = "{display}" in puppy.cfg!') + if key == "yolo_mode" and (value or "").strip().lower() == "config": + emit_success("Using YOLO mode from puppy.cfg; configuration unchanged.") + else: + emit_success(f'Set {key} = "{display}" in puppy.cfg!') # Restart notices (warning) and the reload-success/failure signal # are independent: a restart-required key like ``enable_dbos`` # should still report whether the live agent reload happened. The diff --git a/code_puppy/command_line/core_commands.py b/code_puppy/command_line/core_commands.py index 4b2f3353f..40c56d24f 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: @@ -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/diff_menu.py b/code_puppy/command_line/diff_menu.py index 9b2ee8c1c..65d360c15 100644 --- a/code_puppy/command_line/diff_menu.py +++ b/code_puppy/command_line/diff_menu.py @@ -16,6 +16,7 @@ from prompt_toolkit.layout.controls import FormattedTextControl from prompt_toolkit.widgets import Frame from rich.console import Console +from code_puppy.callbacks import on_prompt_toolkit_style # Sample code snippets for each language LANGUAGE_SAMPLES = { @@ -489,17 +490,17 @@ def get_left_panel_text(): """Generate the selector menu text.""" try: lines = [] - lines.append(("bold cyan", title)) + lines.append(("class:tui.header", title)) lines.append(("", "\n\n")) if not choices: - lines.append(("fg:ansiyellow", "No choices available")) + lines.append(("class:tui.warning", "No choices available")) lines.append(("", "\n")) else: for i, choice in enumerate(choices): if i == selected_index[0]: - lines.append(("fg:ansigreen", "▶ ")) - lines.append(("fg:ansigreen bold", choice)) + lines.append(("class:tui.selected", "▶ ")) + lines.append(("class:tui.selected", choice)) else: lines.append(("", " ")) lines.append(("", choice)) @@ -511,15 +512,18 @@ def get_left_panel_text(): if config is not None: current_lang = config.get_current_language() lang_hint = f"Language: {current_lang.upper()} (←→ to change)" - lines.append(("fg:ansiyellow", lang_hint)) + lines.append(("class:tui.warning", lang_hint)) lines.append(("", "\n")) lines.append( - ("fg:ansicyan", "↑↓ Navigate │ Enter Confirm │ Ctrl-C Cancel") + ( + "class:tui.help-key", + "↑↓ Navigate │ Enter Confirm │ Ctrl-C Cancel", + ) ) return FormattedText(lines) except Exception as e: - return FormattedText([("fg:ansired", f"Error: {e}")]) + return FormattedText([("class:tui.error", f"Error: {e}")]) def get_right_panel_text(): """Generate the preview panel text.""" @@ -528,7 +532,7 @@ def get_right_panel_text(): # get_preview() now returns ANSI, which is already FormattedText-compatible return preview except Exception as e: - return FormattedText([("fg:ansired", f"Preview error: {e}")]) + return FormattedText([("class:tui.error", f"Preview error: {e}")]) kb = KeyBindings() @@ -598,6 +602,7 @@ def cancel(event): full_screen=False, # Don't use full_screen to avoid buffer issues mouse_support=False, color_depth="DEPTH_24_BIT", # Enable truecolor support + style=on_prompt_toolkit_style(), ) sys.stdout.flush() @@ -759,53 +764,35 @@ def _get_preview_text_for_prompt_toolkit(config: DiffConfiguration) -> ANSI: header_text = "\n".join(header_parts) - # Temporarily override config to use current preview settings - from code_puppy.config import ( - get_diff_addition_color, - get_diff_deletion_color, - set_diff_addition_color, - set_diff_deletion_color, + # Pass preview colors directly. A preview should not scribble in puppy.cfg. + formatted_diff = format_diff_with_colors( + sample_diff, + addition_color=config.current_add_color, + deletion_color=config.current_del_color, ) - # Save original values - original_add_color = get_diff_addition_color() - original_del_color = get_diff_deletion_color() - - try: - # Temporarily set config to preview values - set_diff_addition_color(config.current_add_color) - set_diff_deletion_color(config.current_del_color) - - # Get the formatted diff (either Rich Text or Rich markup string) - formatted_diff = format_diff_with_colors(sample_diff) - - # Render everything with Rich Console to get ANSI output with proper color support - buffer = io.StringIO() - console = Console( - file=buffer, - force_terminal=True, - width=90, - legacy_windows=False, - color_system="truecolor", - no_color=False, - force_interactive=True, # Force interactive mode for better color support - ) - - # Print header - console.print(header_text, end="\n") + # Render everything with Rich Console to get ANSI output with proper color support + buffer = io.StringIO() + console = Console( + file=buffer, + force_terminal=True, + width=90, + legacy_windows=False, + color_system="truecolor", + no_color=False, + force_interactive=True, # Force interactive mode for better color support + ) - # Print diff (handles both Text objects and markup strings) - console.print(formatted_diff, end="\n\n") + # Print header + console.print(header_text, end="\n") - # Print footer - console.print("[bold]═" * 50 + "[/bold]", end="") + # Print diff (handles both Text objects and markup strings) + console.print(formatted_diff, end="\n\n") - ansi_output = buffer.getvalue() + # Print footer + console.print("[bold]═" * 50 + "[/bold]", end="") - finally: - # Restore original config values - set_diff_addition_color(original_add_color) - set_diff_deletion_color(original_del_color) + ansi_output = buffer.getvalue() # Wrap in ANSI() so prompt_toolkit can render it return ANSI(ansi_output) diff --git a/code_puppy/command_line/file_path_completion.py b/code_puppy/command_line/file_path_completion.py index e4c118e19..7b2c135ee 100644 --- a/code_puppy/command_line/file_path_completion.py +++ b/code_puppy/command_line/file_path_completion.py @@ -199,6 +199,10 @@ def get_completions( text = document.text cursor_position = document.cursor_position text_before_cursor = text[:cursor_position] + # ``/fork @...`` reserves its first argument for an agent name. Let + # Fork's AgentCompleter own that slot instead of mixing in project files. + if text_before_cursor.lstrip().startswith("/fork @"): + return if self.symbol not in text_before_cursor: return symbol_pos = text_before_cursor.rfind(self.symbol) diff --git a/code_puppy/command_line/judges_menu.py b/code_puppy/command_line/judges_menu.py index 31d84f08e..6b0e7211f 100644 --- a/code_puppy/command_line/judges_menu.py +++ b/code_puppy/command_line/judges_menu.py @@ -52,6 +52,7 @@ validate_name, ) from code_puppy.tools.command_runner import set_awaiting_user_input +from code_puppy.callbacks import on_prompt_toolkit_style PAGE_SIZE = 10 @@ -137,11 +138,11 @@ def _render_model_list( lines: list = [] if not models: - lines.append(("fg:yellow", " No models available.")) + lines.append(("class:tui.warning", " No models available.")) lines.append(("", "\n")) lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", " Configure models first — see /model in the main CLI.", ) ) @@ -152,12 +153,12 @@ def _render_model_list( # Header: (Page x/y, focused indicator) if focused: - lines.append(("fg:ansigreen bold", "▼ ")) + lines.append(("class:tui.selected", "▼ ")) else: - lines.append(("fg:ansibrightblack", " ")) + lines.append(("class:tui.muted", " ")) lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", f"Page {page + 1}/{max(total_pages, 1)} " f"(↑↓ to move, ←→ / PgUp PgDn to page)\n", ) @@ -167,11 +168,11 @@ def _render_model_list( is_sel = i == selected_idx name = _sanitize(models[i]) if is_sel and focused: - lines.append(("fg:ansigreen bold", " ▶ ")) - lines.append(("fg:ansigreen bold", name)) + lines.append(("class:tui.selected", " ▶ ")) + lines.append(("class:tui.selected", name)) elif is_sel: - lines.append(("fg:ansiyellow", " · ")) - lines.append(("fg:ansiyellow", name)) + lines.append(("class:tui.warning", " · ")) + lines.append(("class:tui.warning", name)) else: lines.append(("", " ")) lines.append(("", name)) @@ -294,22 +295,22 @@ def refresh() -> None: focused=is_model_focused(), ) if status_line[0]: - status_control.text = [("fg:ansired", status_line[0])] + status_control.text = [("class:tui.error", status_line[0])] else: # Hint at the bottom of the form: show current model selection. current = current_model() or "(no models available)" status_control.text = [ - ("fg:ansibrightblack", "Selected model: "), - ("fg:ansiyellow", _sanitize(current)), + ("class:tui.muted", "Selected model: "), + ("class:tui.warning", _sanitize(current)), ] help_control.text = [ - ("fg:ansibrightblack", " Tab "), + ("class:tui.help-key", " Tab "), ("", "next field "), - ("fg:ansigreen", " ↑↓ "), + ("class:tui.help-key", " ↑↓ "), ("", "select model "), - ("fg:ansigreen", " Ctrl+S "), + ("class:tui.help-key", " Ctrl+S "), ("", "save "), - ("fg:ansibrightred", " Esc/Ctrl+C "), + ("class:tui.error", " Esc/Ctrl+C "), ("", "cancel"), ] @@ -451,6 +452,7 @@ def _(event): key_bindings=kb, full_screen=False, mouse_support=False, + style=on_prompt_toolkit_style(), ) layout.focus(name_area) @@ -476,80 +478,80 @@ def _render_menu( total_pages = get_total_pages(len(judges), PAGE_SIZE) start, end = get_page_bounds(page, len(judges), PAGE_SIZE) - lines.append(("bold", "Goal Judges")) - lines.append(("fg:ansibrightblack", f" (Page {page + 1}/{max(total_pages, 1)})")) + lines.append(("class:tui.header", "Goal Judges")) + lines.append(("class:tui.muted", f" (Page {page + 1}/{max(total_pages, 1)})")) lines.append(("", "\n\n")) if not judges: - lines.append(("fg:yellow", " No judges configured.")) + lines.append(("class:tui.warning", " No judges configured.")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " Press ")) - lines.append(("fg:ansigreen bold", "N")) - lines.append(("fg:ansibrightblack", " to add one.")) + lines.append(("class:tui.muted", " Press ")) + lines.append(("class:tui.help-key", "N")) + lines.append(("class:tui.muted", " to add one.")) lines.append(("", "\n\n")) else: for i in range(start, end): judge = judges[i] is_selected = i == selected_idx marker = "▶ " if is_selected else " " - row_style = "fg:ansigreen bold" if is_selected else "" + row_style = "class:tui.selected" if is_selected else "" enabled_glyph = "[on] " if judge.enabled else "[off]" - enabled_style = "fg:ansigreen" if judge.enabled else "fg:ansibrightblack" + enabled_style = "class:tui.success" if judge.enabled else "class:tui.muted" - lines.append((row_style or "fg:ansigreen", marker)) + lines.append((row_style or "class:tui.success", marker)) lines.append((enabled_style, enabled_glyph + " ")) lines.append((row_style, _sanitize(judge.name))) - lines.append(("fg:ansibrightblack", " ")) - lines.append(("fg:ansiyellow", _sanitize(judge.model))) + lines.append(("class:tui.muted", " ")) + lines.append(("class:tui.warning", _sanitize(judge.model))) lines.append(("", "\n")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " ↑↓ ")) + lines.append(("class:tui.help-key", " ↑↓ ")) lines.append(("", "Navigate\n")) - lines.append(("fg:ansibrightblack", " ←→ ")) + lines.append(("class:tui.help-key", " ←→ ")) lines.append(("", "Page\n")) - lines.append(("fg:ansigreen", " N ")) + lines.append(("class:tui.help-key", " N ")) lines.append(("", "New judge\n")) - lines.append(("fg:ansigreen", " Enter ")) + lines.append(("class:tui.help-key", " Enter ")) lines.append(("", "Edit (or E)\n")) - lines.append(("fg:ansibrightblack", " T ")) + lines.append(("class:tui.help-key", " T ")) lines.append(("", "Toggle enabled\n")) - lines.append(("fg:ansibrightred", " D ")) + lines.append(("class:tui.error", " D ")) lines.append(("", "Delete\n")) - lines.append(("fg:ansibrightblack", " Esc ")) + lines.append(("class:tui.help-key", " Esc ")) lines.append(("", "Close (or Ctrl+C)")) return lines def _render_preview(judge: Optional[JudgeConfig]) -> list: lines = [] - lines.append(("dim cyan", " JUDGE DETAILS")) + lines.append(("class:tui.title", " JUDGE DETAILS")) lines.append(("", "\n\n")) if judge is None: - lines.append(("fg:yellow", " No judge selected.")) + lines.append(("class:tui.warning", " No judge selected.")) lines.append(("", "\n")) return lines - lines.append(("bold", "Name: ")) + lines.append(("class:tui.label", "Name: ")) lines.append(("", _sanitize(judge.name))) lines.append(("", "\n\n")) - lines.append(("bold", "Model: ")) - lines.append(("fg:ansiyellow", _sanitize(judge.model))) + lines.append(("class:tui.label", "Model: ")) + lines.append(("class:tui.warning", _sanitize(judge.model))) lines.append(("", "\n\n")) - lines.append(("bold", "Enabled: ")) + lines.append(("class:tui.label", "Enabled: ")) if judge.enabled: - lines.append(("fg:ansigreen", "yes")) + lines.append(("class:tui.success", "yes")) else: - lines.append(("fg:ansibrightblack", "no")) + lines.append(("class:tui.muted", "no")) lines.append(("", "\n\n")) - lines.append(("bold", "Prompt:")) + lines.append(("class:tui.label", "Prompt:")) lines.append(("", "\n")) for wrapped in _wrap(judge.prompt or "", width=58): - lines.append(("fg:ansibrightblack", wrapped or " ")) + lines.append(("class:tui.muted", wrapped or " ")) lines.append(("", "\n")) return lines @@ -742,7 +744,11 @@ def _(event): layout = Layout(root) app = Application( - layout=layout, key_bindings=kb, full_screen=False, mouse_support=False + layout=layout, + key_bindings=kb, + full_screen=False, + mouse_support=False, + style=on_prompt_toolkit_style(), ) set_awaiting_user_input(True) diff --git a/code_puppy/command_line/load_context_completion.py b/code_puppy/command_line/load_context_completion.py index 5b8157a6b..62b5ff45b 100644 --- a/code_puppy/command_line/load_context_completion.py +++ b/code_puppy/command_line/load_context_completion.py @@ -38,8 +38,13 @@ def get_completions(self, document, complete_event): try: contexts_dir = Path(CONFIG_DIR) / "contexts" if contexts_dir.exists(): - for pkl_file in contexts_dir.glob("*.pkl"): - session_name = pkl_file.stem # removes .pkl extension + # Support both legacy .pkl and newer .json session artifacts. + session_names = set() + for pattern in ("*.pkl", "*.json"): + for session_file in contexts_dir.glob(pattern): + session_names.add(session_file.stem) + + for session_name in sorted(session_names): if session_name.startswith(session_filter): yield Completion( session_name, diff --git a/code_puppy/command_line/mcp/custom_server_form.py b/code_puppy/command_line/mcp/custom_server_form.py index f73699419..3a5ecf957 100644 --- a/code_puppy/command_line/mcp/custom_server_form.py +++ b/code_puppy/command_line/mcp/custom_server_form.py @@ -27,6 +27,7 @@ from code_puppy.messaging import emit_info, emit_success from code_puppy.tools.command_runner import set_awaiting_user_input +from code_puppy.callbacks import on_prompt_toolkit_style # Example configurations for each server type CUSTOM_SERVER_EXAMPLES = { @@ -129,28 +130,32 @@ def _render_form(self) -> List: lines = [] title = " ✏️ EDIT MCP SERVER" if self.edit_mode else " ➕ ADD CUSTOM MCP SERVER" - lines.append(("bold cyan", title)) + lines.append(("class:tui.header", title)) lines.append(("", "\n\n")) # Server Name field - now in separate frame below - name_style = "fg:ansibrightcyan bold" if self.focused_field == 0 else "bold" + name_style = ( + "class:tui.input.focused" if self.focused_field == 0 else "class:tui.label" + ) lines.append((name_style, " 1. Server Name:")) lines.append(("", "\n")) if self.focused_field == 0: - lines.append(("fg:ansibrightgreen", " ▶ Type in the box below")) + lines.append(("class:tui.input.focused", " ▶ Type in the box below")) else: name_display = self.server_name if self.server_name else "(not set)" - lines.append(("fg:ansibrightblack", f" {name_display}")) + lines.append(("class:tui.muted", f" {name_display}")) # Show name validation hint inline name_error = self._validate_server_name(self.server_name) if name_error and self.server_name: # Only show if there's input lines.append(("", "\n")) - lines.append(("fg:ansiyellow", f" ⚠ {name_error}")) + lines.append(("class:tui.warning", f" \u26a0 {name_error}")) lines.append(("", "\n\n")) # Server Type field - type_style = "fg:ansibrightcyan bold" if self.focused_field == 1 else "bold" + type_style = ( + "class:tui.input.focused" if self.focused_field == 1 else "class:tui.label" + ) lines.append((type_style, " 2. Server Type:")) lines.append(("", "\n")) @@ -165,62 +170,66 @@ def _render_form(self) -> List: icon = type_icons.get(server_type, "") if self.focused_field == 1 and is_selected: - lines.append(("fg:ansibrightgreen", " ▶ ")) + lines.append(("class:tui.selected", " ▶ ")) elif is_selected: - lines.append(("fg:ansigreen", " ✓ ")) + lines.append(("class:tui.success", " \u2713 ")) else: lines.append(("", " ")) if is_selected: - lines.append(("fg:ansibrightcyan bold", f"{icon} {server_type}")) + lines.append(("class:tui.selected", f"{icon} {server_type}")) else: - lines.append(("fg:ansibrightblack", f"{icon} {server_type}")) + lines.append(("class:tui.muted", f"{icon} {server_type}")) lines.append(("", "\n")) lines.append(("", "\n")) # JSON Configuration field - json_style = "fg:ansibrightcyan bold" if self.focused_field == 2 else "bold" + json_style = ( + "class:tui.input.focused" if self.focused_field == 2 else "class:tui.label" + ) lines.append((json_style, " 3. JSON Configuration:")) lines.append(("", "\n")) if self.focused_field == 2: - lines.append(("fg:ansibrightgreen", " ▶ Editing in box below")) + lines.append(("class:tui.input.focused", " ▶ Editing in box below")) else: - lines.append(("fg:ansibrightblack", " (Tab to edit)")) + lines.append(("class:tui.muted", " (Tab to edit)")) lines.append(("", "\n\n")) # Validation status if self.validation_error: - lines.append(("fg:ansired bold", f" ❌ {self.validation_error}")) + lines.append(("class:tui.error", f" \u274c {self.validation_error}")) else: - lines.append(("fg:ansigreen", " ✓ Valid JSON")) + lines.append(("class:tui.success", " \u2713 Valid JSON")) lines.append(("", "\n\n")) # Navigation hints - lines.append(("fg:ansibrightblack", " Tab ")) + lines.append(("class:tui.help-key", " Tab ")) lines.append(("", "Next field ")) - lines.append(("fg:ansibrightblack", "Shift+Tab ")) + lines.append(("class:tui.help-key", "Shift+Tab ")) lines.append(("", "Prev\n")) if self.focused_field == 1: - lines.append(("fg:ansibrightblack", " ↑/↓ ")) + lines.append(("class:tui.help-key", " ↑/↓ ")) lines.append(("", "Change type\n")) - lines.append(("fg:green bold", " Ctrl+S ")) + lines.append(("class:tui.help-key", " Ctrl+S ")) lines.append(("", "Save & Install\n")) - lines.append(("fg:ansired", " Ctrl+C/Esc ")) + lines.append(("class:tui.help-key", " Ctrl+C/Esc ")) lines.append(("", "Cancel")) # Status message bar - shows feedback for user actions if self.status_message: lines.append(("", "\n\n")) - lines.append(("bold", " ─" * 20)) + lines.append(("class:tui.border", " ─" * 20)) lines.append(("", "\n")) if self.status_is_error: - lines.append(("fg:ansired bold", f" ⚠️ {self.status_message}")) + lines.append( + ("class:tui.error", f" \u26a0\ufe0f {self.status_message}") + ) else: - lines.append(("fg:ansigreen bold", f" ✓ {self.status_message}")) + lines.append(("class:tui.success", f" \u2713 {self.status_message}")) return lines @@ -230,62 +239,62 @@ def _render_preview(self) -> List: current_type = self._get_current_type() - lines.append(("bold cyan", " 📝 HELP & PREVIEW")) + lines.append(("class:tui.header", " \U0001f4dd HELP & PREVIEW")) lines.append(("", "\n\n")) # Type description - lines.append(("bold", f" {current_type.upper()} Server")) + lines.append(("class:tui.label", f" {current_type.upper()} Server")) lines.append(("", "\n")) desc = SERVER_TYPE_DESCRIPTIONS.get(current_type, "") - lines.append(("fg:ansibrightblack", f" {desc}")) + lines.append(("class:tui.muted", f" {desc}")) lines.append(("", "\n\n")) # Required fields - lines.append(("bold", " Required Fields:")) + lines.append(("class:tui.label", " Required Fields:")) lines.append(("", "\n")) if current_type == "stdio": - lines.append(("fg:ansicyan", ' • "command"')) - lines.append(("fg:ansibrightblack", " - executable to run")) + lines.append(("class:tui.title", ' • "command"')) + lines.append(("class:tui.muted", " - executable to run")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " Optional:")) + lines.append(("class:tui.muted", " Optional:")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", ' • "args" - command arguments')) + lines.append(("class:tui.muted", ' • "args" - command arguments')) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", ' • "env" - environment variables')) + lines.append(("class:tui.muted", ' • "env" - environment variables')) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", ' • "timeout" - seconds')) + lines.append(("class:tui.muted", ' • "timeout" - seconds')) lines.append(("", "\n")) else: # http or sse - lines.append(("fg:ansicyan", ' • "url"')) - lines.append(("fg:ansibrightblack", " - server endpoint")) + lines.append(("class:tui.title", ' • "url"')) + lines.append(("class:tui.muted", " - server endpoint")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " Optional:")) + lines.append(("class:tui.muted", " Optional:")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", ' • "headers" - HTTP headers')) + lines.append(("class:tui.muted", ' • "headers" - HTTP headers')) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", ' • "timeout" - seconds')) + lines.append(("class:tui.muted", ' • "timeout" - seconds')) lines.append(("", "\n")) lines.append(("", "\n")) # Example - lines.append(("bold", " Example:")) + lines.append(("class:tui.label", " Example:")) lines.append(("", "\n")) example = CUSTOM_SERVER_EXAMPLES.get(current_type, "{}") for line in example.split("\n"): - lines.append(("fg:ansibrightblack", f" {line}")) + lines.append(("class:tui.muted", f" {line}")) lines.append(("", "\n")) lines.append(("", "\n")) # Tips - lines.append(("bold", " 💡 Tips:")) + lines.append(("class:tui.label", " 💡 Tips:")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " • Use $ENV_VAR for secrets")) + lines.append(("class:tui.muted", " • Use $ENV_VAR for secrets")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " • Ctrl+N loads example")) + lines.append(("class:tui.muted", " • Ctrl+N loads example")) lines.append(("", "\n")) return lines @@ -474,6 +483,7 @@ def run(self) -> bool: wrap_lines=False, focusable=True, height=1, + style="class:tui.input", ) # Create JSON text area with syntax highlighting @@ -485,6 +495,7 @@ def run(self) -> bool: focusable=True, height=Dimension(min=8, max=15), lexer=PygmentsLexer(JsonLexer), + style="class:tui.input", ) # Layout with form on left, preview on right @@ -616,6 +627,7 @@ def _(event): key_bindings=kb, full_screen=False, mouse_support=False, + style=on_prompt_toolkit_style(), ) set_awaiting_user_input(True) diff --git a/code_puppy/command_line/mcp/install_menu.py b/code_puppy/command_line/mcp/install_menu.py index 279f01526..663cd4aef 100644 --- a/code_puppy/command_line/mcp/install_menu.py +++ b/code_puppy/command_line/mcp/install_menu.py @@ -24,6 +24,7 @@ prompt_for_server_config, ) from .custom_server_form import run_custom_server_form +from code_puppy.callbacks import on_prompt_toolkit_style logger = logging.getLogger(__name__) @@ -99,22 +100,22 @@ def _get_current_server(self): def _get_category_icon(self, category: str) -> str: """Get an icon for a category.""" if category == CUSTOM_SERVER_CATEGORY: - return "➕" + return "[+]" icons = { - "Code": "💻", - "Storage": "💾", - "Database": "🗄️", - "Documentation": "📝", - "DevOps": "🔧", - "Monitoring": "📊", - "Package Management": "📦", - "Communication": "💬", - "AI": "🤖", - "Search": "🔍", - "Development": "🛠️", - "Cloud": "☁️", + "Code": "[C]", + "Storage": "[S]", + "Database": "[D]", + "Documentation": "[D]", + "DevOps": "[O]", + "Monitoring": "[M]", + "Package Management": "[P]", + "Communication": "[C]", + "AI": "[A]", + "Search": "[S]", + "Development": "[D]", + "Cloud": "[C]", } - return icons.get(category, "📁") + return icons.get(category, "[ ]") def _is_custom_server_selected(self) -> bool: """Check if the custom server category is selected.""" @@ -129,11 +130,11 @@ def _render_category_list(self) -> List: """Render the category list panel.""" lines = [] - lines.append(("bold cyan", " 📂 CATEGORIES")) + lines.append(("class:tui.header", " CATEGORIES")) lines.append(("", "\n\n")) if not self.categories: - lines.append(("fg:yellow", " No categories available.")) + lines.append(("class:tui.warning", " No categories available.")) lines.append(("", "\n\n")) self._render_navigation_hints(lines) return lines @@ -154,9 +155,9 @@ def _render_category_list(self) -> List: if category == CUSTOM_SERVER_CATEGORY: label = f"{prefix}{icon} Custom Server (JSON)" if is_selected: - lines.append(("fg:ansibrightgreen bold", label)) + lines.append(("class:tui.selected", label)) else: - lines.append(("fg:ansigreen", label)) + lines.append(("class:tui.success", label)) else: # Count servers in category server_count = ( @@ -164,16 +165,16 @@ def _render_category_list(self) -> List: ) label = f"{prefix}{icon} {category} ({server_count})" if is_selected: - lines.append(("fg:ansibrightcyan bold", label)) + lines.append(("class:tui.selected", label)) else: - lines.append(("fg:ansibrightblack", label)) + lines.append(("class:tui.muted", label)) lines.append(("", "\n")) lines.append(("", "\n")) if total_pages > 1: lines.append( - ("fg:ansibrightblack", f" Page {self.current_page + 1}/{total_pages}") + ("class:tui.muted", f" Page {self.current_page + 1}/{total_pages}") ) lines.append(("", "\n")) @@ -185,17 +186,17 @@ def _render_server_list(self) -> List: lines = [] if not self.current_category: - lines.append(("fg:yellow", " No category selected.")) + lines.append(("class:tui.warning", " No category selected.")) lines.append(("", "\n\n")) self._render_navigation_hints(lines) return lines icon = self._get_category_icon(self.current_category) - lines.append(("bold cyan", f" {icon} {self.current_category.upper()}")) + lines.append(("class:tui.header", f" {icon} {self.current_category.upper()}")) lines.append(("", "\n\n")) if not self.current_servers: - lines.append(("fg:yellow", " No servers in this category.")) + lines.append(("class:tui.warning", " No servers in this category.")) lines.append(("", "\n\n")) self._render_navigation_hints(lines) return lines @@ -222,16 +223,16 @@ def _render_server_list(self) -> List: label = f"{prefix}{icon_str}{server.display_name}" if is_selected: - lines.append(("fg:ansibrightcyan bold", label)) + lines.append(("class:tui.selected", label)) else: - lines.append(("fg:ansibrightblack", label)) + lines.append(("class:tui.muted", label)) lines.append(("", "\n")) lines.append(("", "\n")) if total_pages > 1: lines.append( - ("fg:ansibrightblack", f" Page {self.current_page + 1}/{total_pages}") + ("class:tui.muted", f" Page {self.current_page + 1}/{total_pages}") ) lines.append(("", "\n")) @@ -241,32 +242,32 @@ def _render_server_list(self) -> List: def _render_navigation_hints(self, lines: List): """Render navigation hints at the bottom of the list panel.""" lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " ↑/↓ ")) + lines.append(("class:tui.help-key", " ↑/↓ ")) lines.append(("", "Navigate ")) - lines.append(("fg:ansibrightblack", "←/→ ")) + lines.append(("class:tui.help-key", "←/→ ")) lines.append(("", "Page\n")) if self.view_mode == "categories": - lines.append(("fg:green", " Enter ")) + lines.append(("class:tui.help-key", " Enter ")) lines.append(("", "Browse Servers\n")) else: - lines.append(("fg:green", " Enter ")) + lines.append(("class:tui.help-key", " Enter ")) lines.append(("", "Install Server\n")) - lines.append(("fg:ansibrightblack", " Esc/Back ")) + lines.append(("class:tui.help-key", " Esc/Back ")) lines.append(("", "Back\n")) - lines.append(("fg:ansired", " Ctrl+C ")) + lines.append(("class:tui.help-key", " Ctrl+C ")) lines.append(("", "Cancel")) def _render_details(self) -> List: """Render the details panel.""" lines = [] - lines.append(("bold cyan", " 📋 DETAILS")) + lines.append(("class:tui.header", " DETAILS")) lines.append(("", "\n\n")) if self.view_mode == "categories": category = self._get_current_category() if not category: - lines.append(("fg:yellow", " No category selected.")) + lines.append(("class:tui.warning", " No category selected.")) return lines # Special handling for custom server category @@ -274,27 +275,27 @@ def _render_details(self) -> List: return self._render_custom_server_details() icon = self._get_category_icon(category) - lines.append(("bold", f" {icon} {category}")) + lines.append(("class:tui.label", f" {icon} {category}")) lines.append(("", "\n\n")) # Show servers in this category servers = self.catalog.get_by_category(category) if self.catalog else [] - lines.append(("fg:ansibrightblack", f" {len(servers)} servers available")) + lines.append(("class:tui.muted", f" {len(servers)} servers available")) lines.append(("", "\n\n")) # Show popular servers in this category popular = [s for s in servers if s.popular] if popular: - lines.append(("bold", " ⭐ Popular:")) + lines.append(("class:tui.label", " * Popular:")) lines.append(("", "\n")) for server in popular[:5]: - lines.append(("fg:ansibrightblack", f" • {server.display_name}")) + lines.append(("class:tui.muted", f" • {server.display_name}")) lines.append(("", "\n")) else: # servers view server = self._get_current_server() if not server: - lines.append(("fg:yellow", " No server selected.")) + lines.append(("class:tui.warning", " No server selected.")) return lines # Server name with indicators @@ -302,19 +303,19 @@ def _render_details(self) -> List: if server.verified: indicators.append("✓ Verified") if server.popular: - indicators.append("⭐ Popular") + indicators.append("Popular") - lines.append(("bold", f" {server.display_name}")) + lines.append(("class:tui.label", f" {server.display_name}")) lines.append(("", "\n")) if indicators: - lines.append(("fg:green", f" {' | '.join(indicators)}")) + lines.append(("class:tui.success", f" {' | '.join(indicators)}")) lines.append(("", "\n")) lines.append(("", "\n")) # Description - lines.append(("bold", " Description:")) + lines.append(("class:tui.label", " Description:")) lines.append(("", "\n")) # Wrap description desc = server.description or "No description available" @@ -323,31 +324,31 @@ def _render_details(self) -> List: line = " " for word in words: if len(line) + len(word) > 50: - lines.append(("fg:ansibrightblack", line)) + lines.append(("class:tui.muted", line)) lines.append(("", "\n")) line = " " + word + " " else: line += word + " " if line.strip(): - lines.append(("fg:ansibrightblack", line)) + lines.append(("class:tui.muted", line)) lines.append(("", "\n")) lines.append(("", "\n")) # Type - lines.append(("bold", " Type:")) + lines.append(("class:tui.label", " Type:")) lines.append(("", "\n")) - type_icons = {"stdio": "📟", "http": "🌐", "sse": "📡"} - type_icon = type_icons.get(server.type, "❓") - lines.append(("fg:ansibrightblack", f" {type_icon} {server.type}")) + type_icons = {"stdio": "[stdio]", "http": "[http]", "sse": "[sse]"} + type_icon = type_icons.get(server.type, "[?]") + lines.append(("class:tui.muted", f" {type_icon} {server.type}")) lines.append(("", "\n\n")) # Tags if server.tags: - lines.append(("bold", " Tags:")) + lines.append(("class:tui.label", " Tags:")) lines.append(("", "\n")) tag_line = " " + ", ".join(server.tags[:6]) - lines.append(("fg:ansicyan", tag_line)) + lines.append(("class:tui.body", tag_line)) lines.append(("", "\n\n")) # Requirements @@ -356,22 +357,22 @@ def _render_details(self) -> List: # Environment variables env_vars = server.get_environment_vars() if env_vars: - lines.append(("bold", " 🔑 Environment Variables:")) + lines.append(("class:tui.label", " Environment Variables:")) lines.append(("", "\n")) for var in env_vars: # Check if already set is_set = os.environ.get(var) if is_set: - lines.append(("fg:green", f" ✓ {var}")) + lines.append(("class:tui.success", f" \u2713 {var}")) else: - lines.append(("fg:yellow", f" ○ {var}")) + lines.append(("class:tui.warning", f" ○ {var}")) lines.append(("", "\n")) lines.append(("", "\n")) # Command line args cmd_args = server.get_command_line_args() if cmd_args: - lines.append(("bold", " ⚙️ Configuration:")) + lines.append(("class:tui.label", " Configuration:")) lines.append(("", "\n")) for arg in cmd_args: name = arg.get("name", "unknown") @@ -380,7 +381,7 @@ def _render_details(self) -> List: marker = "*" if required else "?" default_str = f" [{default}]" if default else "" lines.append( - ("fg:ansibrightblack", f" {marker} {name}{default_str}") + ("class:tui.muted", f" {marker} {name}{default_str}") ) lines.append(("", "\n")) lines.append(("", "\n")) @@ -388,16 +389,16 @@ def _render_details(self) -> List: # Required tools required_tools = requirements.required_tools if required_tools: - lines.append(("bold", " 🛠️ Required Tools:")) + lines.append(("class:tui.label", " Required Tools:")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", f" {', '.join(required_tools)}")) + lines.append(("class:tui.muted", f" {', '.join(required_tools)}")) lines.append(("", "\n\n")) # Example usage if server.example_usage: - lines.append(("bold", " 💡 Example:")) + lines.append(("class:tui.label", " Example:")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", f" {server.example_usage}")) + lines.append(("class:tui.muted", f" {server.example_usage}")) lines.append(("", "\n")) return lines @@ -406,44 +407,44 @@ def _render_custom_server_details(self) -> List: """Render details for the custom server option.""" lines = [] - lines.append(("bold cyan", " 📋 DETAILS")) + lines.append(("class:tui.header", " 📋 DETAILS")) lines.append(("", "\n\n")) - lines.append(("bold green", " ➕ Add Custom MCP Server")) + lines.append(("class:tui.success", " \u2795 Add Custom MCP Server")) lines.append(("", "\n\n")) - lines.append(("fg:ansibrightblack", " Add your own MCP server by providing")) + lines.append(("class:tui.muted", " Add your own MCP server by providing")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " a JSON configuration.")) + lines.append(("class:tui.muted", " a JSON configuration.")) lines.append(("", "\n\n")) - lines.append(("bold", " 📟 Supported Types:")) + lines.append(("class:tui.label", " Supported Types:")) lines.append(("", "\n\n")) - lines.append(("fg:ansicyan bold", " 1. stdio")) + lines.append(("class:tui.title", " 1. stdio")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " Runs a local command (npx, python,")) + lines.append(("class:tui.muted", " Runs a local command (npx, python,")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " uvx, etc.) and communicates via")) + lines.append(("class:tui.muted", " uvx, etc.) and communicates via")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " stdin/stdout.")) + lines.append(("class:tui.muted", " stdin/stdout.")) lines.append(("", "\n\n")) - lines.append(("fg:ansicyan bold", " 2. http")) + lines.append(("class:tui.title", " 2. http")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " Connects to an HTTP endpoint that")) + lines.append(("class:tui.muted", " Connects to an HTTP endpoint that")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " implements the MCP protocol.")) + lines.append(("class:tui.muted", " implements the MCP protocol.")) lines.append(("", "\n\n")) - lines.append(("fg:ansicyan bold", " 3. sse")) + lines.append(("class:tui.title", " 3. sse")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " Connects via Server-Sent Events")) + lines.append(("class:tui.muted", " Connects via Server-Sent Events")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " for real-time streaming.")) + lines.append(("class:tui.muted", " for real-time streaming.")) lines.append(("", "\n\n")) - lines.append(("bold", " 💡 Press Enter to configure")) + lines.append(("class:tui.help", " \U0001f4a1 Press Enter to configure")) lines.append(("", "\n")) return lines @@ -639,6 +640,7 @@ def _(event): # pragma: no cover key_bindings=kb, full_screen=False, mouse_support=False, + style=on_prompt_toolkit_style(), ) set_awaiting_user_input(True) 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/command_line/mcp_binding_menu.py b/code_puppy/command_line/mcp_binding_menu.py index 0908554f3..b931ea2fe 100644 --- a/code_puppy/command_line/mcp_binding_menu.py +++ b/code_puppy/command_line/mcp_binding_menu.py @@ -34,6 +34,7 @@ ) from code_puppy.messaging import emit_info, emit_warning from code_puppy.tools.command_runner import set_awaiting_user_input +from code_puppy.callbacks import on_prompt_toolkit_style def _list_servers() -> List[Tuple[str, str, str]]: @@ -59,13 +60,13 @@ def _render_menu( """Format the left binding panel.""" bindings = get_bound_servers(agent_name) lines: List = [] - lines.append(("bold", "MCP bindings for agent: ")) - lines.append(("fg:ansicyan bold", agent_name)) + lines.append(("class:tui.label", "MCP bindings for agent: ")) + lines.append(("class:tui.title", agent_name)) lines.append(("", "\n\n")) if not servers: - lines.append(("fg:yellow", " No MCP servers installed yet.\n")) - lines.append(("fg:ansibrightblack", " Run /mcp install to add some.\n")) + lines.append(("class:tui.warning", " No MCP servers installed yet.\n")) + lines.append(("class:tui.help", " Run /mcp install to add some.\n")) else: for i, (name, _type, _state) in enumerate(servers): bound = name in bindings @@ -73,25 +74,25 @@ def _render_menu( checkbox = "[x]" if bound else "[ ]" auto_marker = " ⚡auto" if auto else "" prefix = "▶ " if i == selected_idx else " " - style_prefix = "fg:ansigreen bold" if i == selected_idx else "" - style_box = "fg:ansigreen" if bound else "fg:ansibrightblack" + style_prefix = "class:tui.selected" if i == selected_idx else "" + style_box = "class:tui.success" if bound else "class:tui.muted" lines.append((style_prefix, prefix)) lines.append((style_box, f"{checkbox} ")) lines.append((style_prefix or "", name)) if auto_marker: - lines.append(("fg:ansiyellow", auto_marker)) + lines.append(("class:tui.warning", auto_marker)) lines.append(("", "\n")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " ↑↓ ")) + lines.append(("class:tui.help-key", " ↑↓ ")) lines.append(("", "Navigate\n")) - lines.append(("fg:green", " Space ")) + lines.append(("class:tui.help-key", " Space ")) lines.append(("", "Toggle bind\n")) - lines.append(("fg:ansiyellow", " A ")) + lines.append(("class:tui.help-key", " A ")) lines.append(("", "Toggle auto-start\n")) - lines.append(("fg:ansicyan", " Enter / Q ")) + lines.append(("class:tui.help-key", " Enter / Q ")) lines.append(("", "Done\n")) - lines.append(("fg:ansibrightred", " Ctrl+C ")) + lines.append(("class:tui.help-key", " Ctrl+C ")) lines.append(("", "Cancel")) return lines @@ -103,10 +104,10 @@ def _render_preview( ) -> List: """Format the right detail panel.""" lines: List = [] - lines.append(("dim cyan", " SERVER DETAILS")) + lines.append(("class:tui.title", " SERVER DETAILS")) lines.append(("", "\n\n")) if not servers or not (0 <= selected_idx < len(servers)): - lines.append(("fg:ansibrightblack", " Nothing to preview.\n")) + lines.append(("class:tui.muted", " Nothing to preview.\n")) return lines name, type_, state = servers[selected_idx] @@ -114,25 +115,25 @@ def _render_preview( bound = name in bindings auto = bool(bindings.get(name, {}).get("auto_start")) - lines.append(("bold", "Name: ")) - lines.append(("fg:ansicyan", name)) + lines.append(("class:tui.label", "Name: ")) + lines.append(("class:tui.body", name)) lines.append(("", "\n\n")) - lines.append(("bold", "Type: ")) + lines.append(("class:tui.label", "Type: ")) lines.append(("", type_)) lines.append(("", "\n\n")) - lines.append(("bold", "State: ")) + lines.append(("class:tui.label", "State: ")) lines.append( - ("fg:ansigreen" if state == "running" else "fg:ansibrightblack", state) + ("class:tui.success" if state == "running" else "class:tui.muted", state) ) lines.append(("", "\n\n")) - lines.append(("bold", "Bound: ")) + lines.append(("class:tui.label", "Bound: ")) lines.append( - ("fg:ansigreen" if bound else "fg:ansibrightblack", "yes" if bound else "no"), + ("class:tui.success" if bound else "class:tui.muted", "yes" if bound else "no"), ) lines.append(("", "\n\n")) - lines.append(("bold", "Auto-start: ")) + lines.append(("class:tui.label", "Auto-start: ")) lines.append( - ("fg:ansiyellow" if auto else "fg:ansibrightblack", "yes" if auto else "no"), + ("class:tui.warning" if auto else "class:tui.muted", "yes" if auto else "no"), ) lines.append(("", "\n")) return lines @@ -211,6 +212,7 @@ def _(event): key_bindings=kb, full_screen=False, mouse_support=False, + style=on_prompt_toolkit_style(), ) set_awaiting_user_input(True) @@ -256,31 +258,31 @@ async def prompt_bind_after_install(server_name: str) -> None: def render() -> List: lines: List = [] - lines.append(("bold", "Bind '")) - lines.append(("fg:ansicyan bold", server_name)) - lines.append(("bold", "' to which agents?")) + lines.append(("class:tui.label", "Bind '")) + lines.append(("class:tui.title", server_name)) + lines.append(("class:tui.label", "' to which agents?")) lines.append(("", "\n\n")) for i, agent in enumerate(agents): bound = is_bound(agent, server_name) auto = get_auto_start(agent, server_name) if bound else False checkbox = "[x]" if bound else "[ ]" prefix = "▶ " if i == selected_idx[0] else " " - style_prefix = "fg:ansigreen bold" if i == selected_idx[0] else "" - style_box = "fg:ansigreen" if bound else "fg:ansibrightblack" + style_prefix = "class:tui.selected" if i == selected_idx[0] else "" + style_box = "class:tui.success" if bound else "class:tui.muted" lines.append((style_prefix, prefix)) lines.append((style_box, f"{checkbox} ")) lines.append((style_prefix or "", agent)) if auto: - lines.append(("fg:ansiyellow", " ⚡auto")) + lines.append(("class:tui.warning", " \u26a1auto")) lines.append(("", "\n")) lines.append(("", "\n")) - lines.append(("fg:green", " Space ")) + lines.append(("class:tui.help-key", " Space ")) lines.append(("", "Toggle bind ")) - lines.append(("fg:ansiyellow", "A ")) + lines.append(("class:tui.help-key", "A ")) lines.append(("", "Toggle auto-start\n")) - lines.append(("fg:ansicyan", " Enter / Q ")) + lines.append(("class:tui.help-key", " Enter / Q ")) lines.append(("", "Done ")) - lines.append(("fg:ansibrightred", "Ctrl+C ")) + lines.append(("class:tui.help-key", "Ctrl+C ")) lines.append(("", "Skip")) return lines @@ -325,7 +327,11 @@ def _(event): window = Window(content=menu_control, wrap_lines=False) frame = Frame(window, title=f"Bind {server_name}") app = Application( - layout=Layout(frame), key_bindings=kb, full_screen=False, mouse_support=False + layout=Layout(frame), + key_bindings=kb, + full_screen=False, + mouse_support=False, + style=on_prompt_toolkit_style(), ) set_awaiting_user_input(True) diff --git a/code_puppy/command_line/model_picker_completion.py b/code_puppy/command_line/model_picker_completion.py index 1cc4c63c7..d59107a24 100644 --- a/code_puppy/command_line/model_picker_completion.py +++ b/code_puppy/command_line/model_picker_completion.py @@ -27,6 +27,7 @@ save_credential, ) from code_puppy.command_line.utils import safe_input +from code_puppy.callbacks import on_prompt_toolkit_style logger = logging.getLogger(__name__) @@ -340,12 +341,12 @@ def _page_down(self) -> None: self.selected_index = self.page_start def _render(self): - lines = [("bold cyan", " 🤖 Select Active Model")] + lines = [("class:tui.header", " 🤖 Select Active Model")] filter_label = self.filter_text or "type to filter" - lines.append(("fg:ansibrightblack", f"\n Filter: {filter_label}")) + lines.append(("class:tui.muted", f"\n Filter: {filter_label}")) if self.total_pages > 1: lines.append( - ("fg:ansibrightblack", f" (Page {self.page + 1}/{self.total_pages})") + ("class:tui.muted", f" (Page {self.page + 1}/{self.total_pages})") ) lines.append(("", "\n")) @@ -355,19 +356,19 @@ def _render(self): if self.filter_text else "No models available." ) - lines.append(("fg:ansiyellow", f"\n {empty_message}\n")) - lines.append(("fg:ansibrightblack", " Type ")) - lines.append(("", "Adjust filter\n")) - lines.append(("fg:ansibrightblack", " Backspace ")) - lines.append(("", "Delete filter char\n")) + lines.append(("class:tui.warning", f"\n {empty_message}\n")) + lines.append(("class:tui.help-key", " Type ")) + lines.append(("class:tui.help", "Adjust filter\n")) + lines.append(("class:tui.help-key", " Backspace ")) + lines.append(("class:tui.help", "Delete filter char\n")) if self.filter_text: - lines.append(("fg:ansibrightblack", " Ctrl+U ")) - lines.append(("", "Clear filter\n")) - lines.append(("fg:ansiyellow", " Esc ")) - lines.append(("", "Exit\n")) + lines.append(("class:tui.help-key", " Ctrl+U ")) + lines.append(("class:tui.help", "Clear filter\n")) + lines.append(("class:tui.help-key", " Esc ")) + lines.append(("class:tui.help", "Exit\n")) return lines - lines.append(("fg:ansibrightblack", f"\n Current: {self.current_model}\n\n")) + lines.append(("class:tui.muted", f"\n Current: {self.current_model}\n\n")) for offset, model_name in enumerate(self.models_on_page): absolute_index = self.page_start + offset @@ -375,30 +376,30 @@ def _render(self): is_current = model_name == self.current_model prefix = " › " if is_selected else " " - style = "fg:ansiwhite bold" if is_selected else "fg:ansibrightblack" + style = "class:tui.selected" if is_selected else "class:tui.body" lines.append((style, f"{prefix}{model_name}")) if is_current: - lines.append(("fg:ansigreen", " (active)")) + lines.append(("class:tui.success", " (active)")) lines.append(("", "\n")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " ↑/↓ ")) - lines.append(("", "Navigate\n")) + lines.append(("class:tui.help-key", " ↑/↓ ")) + lines.append(("class:tui.help", "Navigate\n")) if self.total_pages > 1: - lines.append(("fg:ansibrightblack", " PgUp/PgDn ")) - lines.append(("", "Change page\n")) - lines.append(("fg:ansibrightblack", " Type ")) - lines.append(("", "Filter models\n")) - lines.append(("fg:ansibrightblack", " Backspace ")) - lines.append(("", "Delete filter char\n")) - lines.append(("fg:ansibrightblack", " Ctrl+U ")) - lines.append(("", "Clear filter\n")) - lines.append(("fg:ansigreen", " Enter ")) - lines.append(("", "Select model\n")) - lines.append(("fg:cyan", " Ctrl+E ")) - lines.append(("", "Edit credentials\n")) - lines.append(("fg:ansiyellow", " Esc ")) - lines.append(("", "Cancel\n")) + lines.append(("class:tui.help-key", " PgUp/PgDn ")) + lines.append(("class:tui.help", "Change page\n")) + lines.append(("class:tui.help-key", " Type ")) + lines.append(("class:tui.help", "Filter models\n")) + lines.append(("class:tui.help-key", " Backspace ")) + lines.append(("class:tui.help", "Delete filter char\n")) + lines.append(("class:tui.help-key", " Ctrl+U ")) + lines.append(("class:tui.help", "Clear filter\n")) + lines.append(("class:tui.help-key", " Enter ")) + lines.append(("class:tui.help", "Select model\n")) + lines.append(("class:tui.help-key", " Ctrl+E ")) + lines.append(("class:tui.help", "Edit credentials\n")) + lines.append(("class:tui.help-key", " Esc ")) + lines.append(("class:tui.help", "Cancel\n")) return lines def _edit_credentials_for_model(self, model_name: str) -> None: @@ -530,6 +531,7 @@ def _(event): layout=Layout(Window(content=control, wrap_lines=True)), key_bindings=kb, full_screen=False, + style=on_prompt_toolkit_style(), ) await app.run_async() @@ -601,6 +603,7 @@ async def get_input_with_model_completion( completer=ModelNameCompleter(trigger), history=history, complete_while_typing=True, + style=on_prompt_toolkit_style(), ) text = await session.prompt_async(prompt_str) possibly_stripped = update_model_in_input(text) diff --git a/code_puppy/command_line/model_settings_menu.py b/code_puppy/command_line/model_settings_menu.py index a9c113a98..7adae6c30 100644 --- a/code_puppy/command_line/model_settings_menu.py +++ b/code_puppy/command_line/model_settings_menu.py @@ -35,6 +35,7 @@ from code_puppy.messaging import emit_info from code_puppy.model_factory import ModelFactory from code_puppy.tools.command_runner import set_awaiting_user_input +from code_puppy.callbacks import on_prompt_toolkit_style # Pagination config MODELS_PER_PAGE = 15 @@ -76,7 +77,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 +185,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 +201,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 @@ -357,18 +356,18 @@ def _render_main_list(self) -> List: if self.view_mode == "models": # Header with page indicator - lines.append(("bold cyan", " 🐕 Select a Model to Configure")) + lines.append(("class:tui.header", " 🐕 Select a Model to Configure")) if self.total_pages > 1: lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", f" (Page {self.page + 1}/{self.total_pages})", ) ) lines.append(("", "\n\n")) if not self.all_models: - lines.append(("fg:ansiyellow", " No models available.")) + lines.append(("class:tui.warning", " No models available.")) lines.append(("", "\n\n")) self._add_model_nav_hints(lines) return lines @@ -384,7 +383,7 @@ def _render_main_list(self) -> List: is_current = model_name == self.current_model_name prefix = " › " if is_selected else " " - style = "fg:ansiwhite bold" if is_selected else "fg:ansibrightblack" + style = "class:tui.selected" if is_selected else "class:tui.body" # Check if model has any custom settings model_settings = get_all_model_settings(model_name) @@ -394,26 +393,26 @@ def _render_main_list(self) -> List: # Show indicators if is_current: - lines.append(("fg:ansigreen", " (active)")) + lines.append(("class:tui.success", " (active)")) if has_settings: - lines.append(("fg:ansicyan", " ⚙")) + lines.append(("class:tui.body", " ⚙")) lines.append(("", "\n")) if is_selected: description = get_model_description(models_config, model_name) - lines.append(("fg:ansiyellow italic", f" {description}\n")) + lines.append(("class:tui.body", f" {description}\n")) lines.append(("", "\n")) self._add_model_nav_hints(lines) else: # Settings view - lines.append(("bold cyan", f" ⚙ Settings for {self.selected_model}")) + lines.append(("class:tui.header", f" ⚙ Settings for {self.selected_model}")) lines.append(("", "\n\n")) if not self.supported_settings: lines.append( - ("fg:ansiyellow", " No configurable settings for this model.") + ("class:tui.warning", " No configurable settings for this model.") ) lines.append(("", "\n\n")) self._add_settings_nav_hints(lines) @@ -428,18 +427,18 @@ def _render_main_list(self) -> List: if is_selected and self.editing_mode: display_value = self._format_value(setting_key, self.edit_value) prefix = " ✏️ " - style = "fg:ansigreen bold" + style = "class:tui.success" else: display_value = self._format_value(setting_key, current_value) prefix = " › " if is_selected else " " - style = "fg:ansiwhite" if is_selected else "fg:ansibrightblack" + style = "class:tui.selected" if is_selected else "class:tui.body" # Setting name and value lines.append((style, f"{prefix}{setting_def['name']}: ")) if current_value is not None or (is_selected and self.editing_mode): - lines.append(("fg:ansicyan", display_value)) + lines.append(("class:tui.body", display_value)) else: - lines.append(("fg:ansibrightblack dim", display_value)) + lines.append(("class:tui.muted", display_value)) lines.append(("", "\n")) lines.append(("", "\n")) @@ -450,89 +449,89 @@ def _render_main_list(self) -> List: def _add_model_nav_hints(self, lines: List): """Add navigation hints for model list view.""" lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " ↑/↓ ")) - lines.append(("", "Navigate models\n")) + lines.append(("class:tui.help-key", " ↑/↓ ")) + lines.append(("class:tui.help", "Navigate models\n")) if self.total_pages > 1: - lines.append(("fg:ansibrightblack", " PgUp/PgDn ")) - lines.append(("", "Change page\n")) - lines.append(("fg:ansigreen", " Enter ")) - lines.append(("", "Configure model\n")) - lines.append(("fg:ansiyellow", " Esc ")) - lines.append(("", "Exit\n")) + lines.append(("class:tui.help-key", " PgUp/PgDn ")) + lines.append(("class:tui.help", "Change page\n")) + lines.append(("class:tui.help-key", " Enter ")) + lines.append(("class:tui.help", "Configure model\n")) + lines.append(("class:tui.help-key", " Esc ")) + lines.append(("class:tui.help", "Exit\n")) def _add_settings_nav_hints(self, lines: List): """Add navigation hints for settings view.""" lines.append(("", "\n")) if self.editing_mode: - lines.append(("fg:ansibrightblack", " ←/→ ")) - lines.append(("", "Adjust value\n")) - lines.append(("fg:ansigreen", " Enter ")) - lines.append(("", "Save\n")) - lines.append(("fg:ansiyellow", " Esc ")) - lines.append(("", "Cancel edit\n")) - lines.append(("fg:ansired", " d ")) - lines.append(("", "Reset to default\n")) + lines.append(("class:tui.help-key", " ←/→ ")) + lines.append(("class:tui.help", "Adjust value\n")) + lines.append(("class:tui.help-key", " Enter ")) + lines.append(("class:tui.help", "Save\n")) + lines.append(("class:tui.help-key", " Esc ")) + lines.append(("class:tui.help", "Cancel edit\n")) + lines.append(("class:tui.help-key", " d ")) + lines.append(("class:tui.help", "Reset to default\n")) else: - lines.append(("fg:ansibrightblack", " ↑/↓ ")) - lines.append(("", "Navigate settings\n")) - lines.append(("fg:ansigreen", " Enter ")) - lines.append(("", "Edit setting\n")) - lines.append(("fg:ansired", " d ")) - lines.append(("", "Reset to default\n")) - lines.append(("fg:ansiyellow", " Esc ")) - lines.append(("", "Back to models\n")) + lines.append(("class:tui.help-key", " ↑/↓ ")) + lines.append(("class:tui.help", "Navigate settings\n")) + lines.append(("class:tui.help-key", " Enter ")) + lines.append(("class:tui.help", "Edit setting\n")) + lines.append(("class:tui.help-key", " d ")) + lines.append(("class:tui.help", "Reset to default\n")) + lines.append(("class:tui.help-key", " Esc ")) + lines.append(("class:tui.help", "Back to models\n")) def _render_details_panel(self) -> List: """Render the details/help panel.""" lines = [] if self.view_mode == "models": - lines.append(("bold cyan", " Model Info")) + lines.append(("class:tui.title", " Model Info")) lines.append(("", "\n\n")) if not self.all_models: - lines.append(("fg:ansibrightblack", " No models available.")) + lines.append(("class:tui.muted", " No models available.")) return lines model_name = self.all_models[self.model_index] is_current = model_name == self.current_model_name - lines.append(("bold", f" {model_name}")) + lines.append(("class:tui.label", f" {model_name}")) lines.append(("", "\n\n")) if is_current: - lines.append(("fg:ansigreen", " ✓ Currently active model")) + lines.append(("class:tui.success", " ✓ Currently active model")) lines.append(("", "\n\n")) # Show current settings for this model model_settings = _get_model_display_settings(model_name) if model_settings: - lines.append(("bold", " Effective Settings:")) + lines.append(("class:tui.label", " Effective Settings:")) lines.append(("", "\n")) for setting_key, value in model_settings.items(): setting_def = SETTING_DEFINITIONS.get(setting_key, {}) name = setting_def.get("name", setting_key) display = self._format_value(setting_key, value) - lines.append(("fg:ansicyan", f" {name}: {display}")) + lines.append(("class:tui.body", f" {name}: {display}")) lines.append(("", "\n")) else: - lines.append(("fg:ansibrightblack", " Using all default settings")) + lines.append(("class:tui.muted", " Using all default settings")) lines.append(("", "\n")) # Show supported settings supported = self._get_supported_settings(model_name) lines.append(("", "\n")) - lines.append(("bold", " Configurable Settings:")) + lines.append(("class:tui.label", " Configurable Settings:")) lines.append(("", "\n")) if supported: for s in supported: setting_def = SETTING_DEFINITIONS.get(s, {}) name = setting_def.get("name", s) - lines.append(("fg:ansibrightblack", f" • {name}")) + lines.append(("class:tui.muted", f" • {name}")) lines.append(("", "\n")) else: - lines.append(("fg:ansibrightblack dim", " None")) + lines.append(("class:tui.muted", " None")) lines.append(("", "\n")) # Show pagination info at the bottom of details @@ -540,7 +539,7 @@ def _render_details_panel(self) -> List: lines.append(("", "\n")) lines.append( ( - "fg:ansibrightblack dim", + "class:tui.muted", f" Model {self.model_index + 1} of {len(self.all_models)}", ) ) @@ -548,12 +547,12 @@ def _render_details_panel(self) -> List: else: # Settings detail view - lines.append(("bold cyan", " Setting Details")) + lines.append(("class:tui.title", " Setting Details")) lines.append(("", "\n\n")) if not self.supported_settings: lines.append( - ("fg:ansibrightblack", " This model doesn't expose any settings.") + ("class:tui.muted", " This model doesn't expose any settings.") ) return lines @@ -562,84 +561,82 @@ def _render_details_panel(self) -> List: current_value = self._get_current_value(setting_key) # Setting name - lines.append(("bold", f" {setting_def['name']}")) + lines.append(("class:tui.label", f" {setting_def['name']}")) lines.append(("", "\n")) # Show if this is a global setting if setting_key in ("reasoning_effort", "verbosity"): lines.append( ( - "fg:ansiyellow", + "class:tui.warning", " ⚠ Global setting (applies to all GPT-5 models)", ) ) lines.append(("", "\n\n")) # Description - lines.append(("fg:ansibrightblack", f" {setting_def['description']}")) + lines.append(("class:tui.muted", f" {setting_def['description']}")) lines.append(("", "\n\n")) # Range/choices info if setting_def.get("type") == "choice": - lines.append(("bold", " Options:")) + lines.append(("class:tui.label", " Options:")) lines.append(("", "\n")) # Get filtered choices based on model capabilities choices = _get_setting_choices(setting_key, self.selected_model) lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", f" {' | '.join(choices)}", ) ) elif setting_def.get("type") == "boolean": - lines.append(("bold", " Options:")) + lines.append(("class:tui.label", " Options:")) lines.append(("", "\n")) lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", " Enabled | Disabled", ) ) else: - lines.append(("bold", " Range:")) + lines.append(("class:tui.label", " Range:")) lines.append(("", "\n")) lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", f" Min: {setting_def['min']} Max: {setting_def['max']} Step: {setting_def['step']}", ) ) lines.append(("", "\n\n")) # Current value - lines.append(("bold", " Current Value:")) + lines.append(("class:tui.label", " Current Value:")) lines.append(("", "\n")) if current_value is not None: lines.append( ( - "fg:ansicyan", + "class:tui.body", f" {self._format_value(setting_key, current_value)}", ) ) else: - lines.append(("fg:ansibrightblack dim", " (using model default)")) + lines.append(("class:tui.muted", " (using model default)")) lines.append(("", "\n\n")) # Editing hint if self.editing_mode: - lines.append(("fg:ansigreen bold", " ✏️ EDITING MODE")) + lines.append(("class:tui.success", " ✏️ EDITING MODE")) lines.append(("", "\n")) if self.edit_value is not None: lines.append( ( - "fg:ansicyan", + "class:tui.body", f" New value: {self._format_value(setting_key, self.edit_value)}", ) ) else: - lines.append( - ("fg:ansibrightblack", " New value: (model default)") - ) + lines.append(("class:tui.muted", " New value: (model default)")) lines.append(("", "\n")) return lines @@ -934,6 +931,7 @@ def _(event): key_bindings=kb, full_screen=False, mouse_support=False, + style=on_prompt_toolkit_style(), ) set_awaiting_user_input(True) diff --git a/code_puppy/command_line/onboarding_slides.py b/code_puppy/command_line/onboarding_slides.py index 535abe51d..9f39184fe 100644 --- a/code_puppy/command_line/onboarding_slides.py +++ b/code_puppy/command_line/onboarding_slides.py @@ -1,15 +1,13 @@ -"""Slide content for the onboarding wizard. +"""Semantic slide content for the onboarding wizard. -🐶 Lean, mean, ADHD-friendly slides. 5 slides max! + Lean, mean, ADHD-friendly slides. Five slides max, because this is +onboarding rather than a hostage situation. """ from typing import List, Tuple -# ============================================================================ -# Slide Data Constants -# ============================================================================ +SlideContent = list[tuple[str, str]] -# Model subscription options MODEL_OPTIONS: List[Tuple[str, str, str]] = [ ("chatgpt", "ChatGPT Plus/Pro/Max", "OAuth login - no API key needed"), ("claude", "Claude Code Pro/Max", "OAuth login - no API key needed"), @@ -19,161 +17,172 @@ ] -# ============================================================================ -# Navigation Footer (shown on ALL slides) -# ============================================================================ +def _add(content: SlideContent, style: str, text: str) -> None: + content.append((style, text)) -def get_nav_footer() -> str: - """Navigation hints shown at bottom of every slide.""" - return ( - "\n[dim]─────────────────────────────────────[/dim]\n" - "[green]→/l[/green] Next " - "[green]←/h[/green] Back " - "[green]↑↓/jk[/green] Options " - "[green]Enter[/green] Select " - "[yellow]ESC[/yellow] Skip" - ) - - -# ============================================================================ -# Gradient Banner -# ============================================================================ +def get_nav_footer() -> SlideContent: + """Return semantic navigation hints shown on every slide.""" + return [ + ("class:tui.muted", "\n─────────────────────────────────────\n"), + ("class:tui.help-key", "→/l"), + ("class:tui.help", " Next "), + ("class:tui.help-key", "←/h"), + ("class:tui.help", " Back "), + ("class:tui.help-key", "↑↓/jk"), + ("class:tui.help", " Options "), + ("class:tui.help-key", "Enter"), + ("class:tui.help", " Select "), + ("class:tui.warning", "ESC"), + ("class:tui.help", " Skip"), + ] -def get_gradient_banner() -> str: - """Generate the gradient CODE PUPPY banner.""" +def get_gradient_banner() -> SlideContent: + """Generate the CODE PUPPY banner using the semantic header style.""" try: import pyfiglet - lines = pyfiglet.figlet_format("CODE PUPPY", font="ansi_shadow").split("\n") - colors = ["bright_blue", "bright_cyan", "bright_green"] - result = [] - for i, line in enumerate(lines): - if line.strip(): - color = colors[min(i // 2, len(colors) - 1)] - result.append(f"[{color}]{line}[/{color}]") - return "\n".join(result) + banner = pyfiglet.figlet_format("CODE PUPPY", font="ansi_shadow").rstrip() except ImportError: - return "[bold bright_cyan]═══ CODE PUPPY 🐶 ═══[/bold bright_cyan]" + banner = "═══ CODE PUPPY ═══" + return [("class:tui.header", banner)] -# ============================================================================ -# Slide Content (5 slides total) -# ============================================================================ - - -def slide_welcome() -> str: - """Slide 1: Welcome - quick intro.""" +def slide_welcome() -> SlideContent: + """Slide 1: welcome and quick intro.""" content = get_gradient_banner() - content += "\n\n" - content += "[bold white]Welcome! 🐶[/bold white]\n\n" - content += "[cyan]Quick setup:[/cyan]\n" - content += " 1. Pick your model provider\n" - content += " 2. Optional: MCP servers\n" - content += " 3. Learn when to use which agent\n" - content += " 4. Start coding!\n\n" - content += "[dim]Takes ~1 minute. Let's go![/dim]" - content += get_nav_footer() + content.extend( + [ + ("class:tui.header", "\n\nWelcome! \n\n"), + ("class:tui.title", "Quick setup:\n"), + ("class:tui.body", " 1. Pick your model provider\n"), + ("class:tui.body", " 2. Optional: MCP servers\n"), + ("class:tui.body", " 3. Learn when to use which agent\n"), + ("class:tui.body", " 4. Start coding!\n\n"), + ("class:tui.muted", "Takes ~1 minute. Let's go!"), + ] + ) + content.extend(get_nav_footer()) return content -def slide_models(selected_option: int, options: List[Tuple[str, str]]) -> str: - """Slide 2: Model selection.""" - content = "[bold cyan]📦 Pick Your Models[/bold cyan]\n\n" - content += "[white]How do you want to access LLMs?[/white]\n\n" - +def slide_models(selected_option: int, options: List[Tuple[str, str]]) -> SlideContent: + """Slide 2: model selection.""" + content: SlideContent = [ + ("class:tui.header", " Pick Your Models\n\n"), + ("class:tui.body", "How do you want to access LLMs?\n\n"), + ] for i, (_, label) in enumerate(options): - if i == selected_option: - content += f"[bold green]▶ {label}[/bold green]\n" - else: - content += f"[dim] {label}[/dim]\n" - - content += "\n" + style = "class:tui.selected" if i == selected_option else "class:tui.muted" + marker = "▶ " if i == selected_option else " " + _add(content, style, f"{marker}{label}\n") + _add(content, "", "\n") - # Context based on selection opt = options[selected_option][0] if options else None if opt == "chatgpt": - content += "[yellow]💡 ChatGPT OAuth[/yellow]\n" - content += " Uses your existing subscription\n" - content += " GPT-5.2, GPT-5.2-codex\n" + _add(content, "class:tui.warning", " ChatGPT OAuth\n") + _add( + content, + "class:tui.body", + " Uses your existing subscription\n GPT-5.2, GPT-5.2-codex\n", + ) elif opt == "claude": - content += "[yellow]💡 Claude OAuth[/yellow]\n" - content += " Uses your existing subscription\n" - content += " Opus/Sonnet/Haiku 4.5\n" + _add(content, "class:tui.warning", " Claude OAuth\n") + _add( + content, + "class:tui.body", + " Uses your existing subscription\n Opus/Sonnet/Haiku 4.5\n", + ) elif opt == "api_keys": - content += "[yellow]💡 API Keys[/yellow]\n" - content += " [cyan]/set OPENAI_API_KEY=sk-...[/cyan]\n" - content += " [cyan]/add_model[/cyan] to browse 1500+ models\n" + _add(content, "class:tui.warning", " API Keys\n") + _add( + content, "class:tui.help-key", " /set OPENAI_API_KEY=sk-...\n /add_model" + ) + _add(content, "class:tui.body", " to browse 1500+ models\n") elif opt == "openrouter": - content += "[yellow]💡 OpenRouter[/yellow]\n" - content += " One API key, all providers\n" - content += " [cyan]/set OPENROUTER_API_KEY=...[/cyan]\n" + _add(content, "class:tui.warning", " OpenRouter\n") + _add(content, "class:tui.body", " One API key, all providers\n") + _add(content, "class:tui.help-key", " /set OPENROUTER_API_KEY=...\n") else: - content += "[dim]No worries! Use /set or /add_model later[/dim]\n" + _add(content, "class:tui.muted", "No worries! Use /set or /add_model later\n") - content += get_nav_footer() + content.extend(get_nav_footer()) return content -def slide_mcp() -> str: - """Slide 3: MCP servers (optional power-ups).""" - content = "[bold cyan]🔌 MCP Servers (Optional)[/bold cyan]\n\n" - content += "[white]Supercharge with external tools![/white]\n\n" - content += "[green]Commands:[/green]\n" - content += " [cyan]/mcp install[/cyan] Browse catalog\n" - content += " [cyan]/mcp list[/cyan] See your servers\n\n" - content += "[yellow]🌟 Popular picks:[/yellow]\n" - content += " • GitHub integration\n" - content += " • Postgres/databases\n" - content += " • Slack, Linear, etc.\n\n" - content += "[dim]Skip this if you just want to code![/dim]" - content += get_nav_footer() +def slide_mcp() -> SlideContent: + """Slide 3: optional MCP server power-ups.""" + content: SlideContent = [ + ("class:tui.header", " MCP Servers (Optional)\n\n"), + ("class:tui.body", "Supercharge with external tools!\n\n"), + ("class:tui.title", "Commands:\n"), + ("class:tui.help-key", " /mcp install"), + ("class:tui.body", " Browse catalog\n"), + ("class:tui.help-key", " /mcp list"), + ("class:tui.body", " See your servers\n\n"), + ("class:tui.warning", " Popular picks:\n"), + ( + "class:tui.body", + " • GitHub integration\n • Postgres/databases\n • Slack, Linear, etc.\n\n", + ), + ("class:tui.muted", "Skip this if you just want to code!"), + ] + content.extend(get_nav_footer()) return content -def slide_use_cases() -> str: - """Slide 4: When to use which agent - THE IMPORTANT ONE.""" - content = "[bold cyan]🎯 When to Use What[/bold cyan]\n\n" - - content += "[bold yellow]🐶 Code Puppy (default)[/bold yellow]\n" - content += " [green]USE FOR:[/green] Direct coding tasks\n" - content += " • Fix this bug\n" - content += " • Add a feature to this file\n" - content += " • Refactor this function\n" - content += " • Write tests for X\n\n" - - content += "[bold yellow]📋 Planning Agent[/bold yellow]\n" - content += " [green]USE FOR:[/green] Complex multi-step projects\n" - content += " • Build me a REST API with auth\n" - content += " • Create a CLI tool from scratch\n" - content += " • Refactor entire codebase\n" - content += " • Multi-file architectural changes\n\n" - - content += "[cyan]Switch: /agent planning-agent[/cyan]\n" - content += "[dim]Planning breaks big tasks into steps,[/dim]\n" - content += "[dim]then delegates to specialists.[/dim]" - content += get_nav_footer() +def slide_use_cases() -> SlideContent: + """Slide 4: when to use each agent.""" + content: SlideContent = [ + ("class:tui.header", " When to Use What\n\n"), + ("class:tui.warning", " Code Puppy (default)\n"), + ("class:tui.success", " USE FOR:"), + ( + "class:tui.body", + " Direct coding tasks\n • Fix this bug\n • Add a feature to this file\n • Refactor this function\n • Write tests for X\n\n", + ), + ("class:tui.warning", " Planning Agent\n"), + ("class:tui.success", " USE FOR:"), + ( + "class:tui.body", + " Complex multi-step projects\n • Build me a REST API with auth\n • Create a CLI tool from scratch\n • Refactor entire codebase\n • Multi-file architectural changes\n\n", + ), + ("class:tui.help-key", "Switch: /agent planning-agent\n"), + ( + "class:tui.muted", + "Planning breaks big tasks into steps,\nthen delegates to specialists.", + ), + ] + content.extend(get_nav_footer()) return content -def slide_done(trigger_oauth: str | None) -> str: - """Slide 5: You're ready!""" - content = "[bold green]🎉 Ready to Roll![/bold green]\n\n" - content += "[bold cyan]Essential commands:[/bold cyan]\n" - content += " [cyan]/model[/cyan] Switch models\n" - content += " [cyan]/agent[/cyan] Switch agents\n" - content += " [cyan]/help[/cyan] All commands\n\n" - - content += "[bold yellow]Pro tips:[/bold yellow]\n" - content += " • Be specific in prompts\n" - content += " • Use Planning Agent for big tasks\n" - content += " • @ for file path completion\n\n" - +def slide_done(trigger_oauth: str | None) -> SlideContent: + """Slide 5: ready to roll.""" + content: SlideContent = [ + ("class:tui.success", " Ready to Roll!\n\n"), + ("class:tui.header", "Essential commands:\n"), + ("class:tui.help-key", " /model"), + ("class:tui.body", " Switch models\n"), + ("class:tui.help-key", " /agent"), + ("class:tui.body", " Switch agents\n"), + ("class:tui.help-key", " /help"), + ("class:tui.body", " All commands\n\n"), + ("class:tui.warning", "Pro tips:\n"), + ( + "class:tui.body", + " • Be specific in prompts\n • Use Planning Agent for big tasks\n • @ for file path completion\n\n", + ), + ] if trigger_oauth: - content += f"[bold cyan]→ {trigger_oauth.title()} OAuth next![/bold cyan]\n\n" - - content += "[dim]Re-run anytime: [/dim][cyan]/tutorial[/cyan]\n" - content += "\n[bold yellow]Press Enter to start coding! 🐶[/bold yellow]" - content += get_nav_footer() + _add(content, "class:tui.header", f"→ {trigger_oauth.title()} OAuth next!\n\n") + content.extend( + [ + ("class:tui.muted", "Re-run anytime: "), + ("class:tui.help-key", "/tutorial\n\n"), + ("class:tui.warning", "Press Enter to start coding! "), + ] + ) + content.extend(get_nav_footer()) return content diff --git a/code_puppy/command_line/onboarding_wizard.py b/code_puppy/command_line/onboarding_wizard.py index b00fc9b21..208ff2f18 100644 --- a/code_puppy/command_line/onboarding_wizard.py +++ b/code_puppy/command_line/onboarding_wizard.py @@ -13,29 +13,28 @@ """ import asyncio -import io import os import sys from typing import List, Optional, Tuple from prompt_toolkit import Application -from prompt_toolkit.formatted_text import ANSI +from prompt_toolkit.formatted_text import FormattedText from prompt_toolkit.key_binding import KeyBindings from prompt_toolkit.layout import Layout, Window from prompt_toolkit.layout.controls import FormattedTextControl from prompt_toolkit.widgets import Frame -from rich.console import Console - from code_puppy.config import CONFIG_DIR from .onboarding_slides import ( MODEL_OPTIONS, + SlideContent, slide_done, slide_mcp, slide_models, slide_use_cases, slide_welcome, ) +from code_puppy.callbacks import on_prompt_toolkit_style # ============================================================================ # State Tracking @@ -109,7 +108,7 @@ def get_progress_indicator(self) -> str: "●" if i == self.current_slide else "○" for i in range(self.TOTAL_SLIDES) ) - def get_slide_content(self) -> str: + def get_slide_content(self) -> SlideContent: """Get content for current slide.""" if self.current_slide == 0: return slide_welcome() @@ -175,30 +174,18 @@ def prev_option(self) -> None: # ============================================================================ -def _get_slide_panel_content(wizard: OnboardingWizard) -> ANSI: - """Generate slide content for display.""" - buffer = io.StringIO() - console = Console( - file=buffer, - force_terminal=True, - width=80, - legacy_windows=False, - color_system="truecolor", - no_color=False, - force_interactive=True, - ) - - # Progress indicator +def _get_slide_panel_content(wizard: OnboardingWizard) -> FormattedText: + """Generate semantically styled slide content for display.""" progress = wizard.get_progress_indicator() - console.print(f"[dim]{progress}[/dim]") - console.print( - f"[dim]Slide {wizard.current_slide + 1} of {wizard.TOTAL_SLIDES}[/dim]\n" - ) - - # Slide content (includes nav footer) - console.print(wizard.get_slide_content()) - - return ANSI(buffer.getvalue()) + content: SlideContent = [ + ("class:tui.muted", f"{progress}\n"), + ( + "class:tui.muted", + f"Slide {wizard.current_slide + 1} of {wizard.TOTAL_SLIDES}\n\n", + ), + ] + content.extend(wizard.get_slide_content()) + return FormattedText(content) # ============================================================================ @@ -300,6 +287,7 @@ def cancel_wizard(event): full_screen=False, mouse_support=False, color_depth="DEPTH_24_BIT", + style=on_prompt_toolkit_style(), ) sys.stdout.write("\033[2J\033[H") diff --git a/code_puppy/command_line/prompt_toolkit_completion.py b/code_puppy/command_line/prompt_toolkit_completion.py index dbff51104..30eba995e 100644 --- a/code_puppy/command_line/prompt_toolkit_completion.py +++ b/code_puppy/command_line/prompt_toolkit_completion.py @@ -42,6 +42,7 @@ from code_puppy.command_line.pin_command_completion import PinCompleter, UnpinCompleter from code_puppy.command_line.skills_completion import SkillsCompleter from code_puppy.command_line.utils import list_directory +from code_puppy.callbacks import on_prompt_text_color, on_prompt_toolkit_style from code_puppy.config import ( COMMAND_HISTORY_FILE, get_config_keys, @@ -142,18 +143,21 @@ def get_completions(self, document, complete_event): # Extract the input after /set and space (up to cursor) trigger_end = actual_trigger_pos + len(self.trigger) + 1 # +1 for the space text_after_trigger = text_before_cursor[trigger_end:cursor_position].lstrip() - start_position = -(len(text_after_trigger)) + start_position = -len(text_after_trigger) # --- SPECIAL HANDLING FOR 'model' KEY --- if text_after_trigger == "model": # 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) @@ -170,7 +174,7 @@ def get_completions(self, document, complete_event): class AttachmentPlaceholderProcessor(Processor): """Display friendly placeholders for recognised attachments.""" - _PLACEHOLDER_STYLE = "class:attachment-placeholder" + _PLACEHOLDER_STYLE = "class:attachment-placeholder class:tui.title" # Skip expensive path detection for very long input (likely pasted content) _MAX_TEXT_LENGTH_FOR_REALTIME = 500 @@ -351,8 +355,9 @@ class AgentCompleter(Completer): Usage: /agent """ - def __init__(self, trigger: str = "/agent"): + def __init__(self, trigger: str = "/agent", prefix: str = ""): self.trigger = trigger + self.prefix = prefix def get_completions(self, document, complete_event): cursor_position = document.cursor_position @@ -367,7 +372,11 @@ def get_completions(self, document, complete_event): trigger_pos = text_before_cursor.find(self.trigger) trigger_end = trigger_pos + len(self.trigger) + 1 # +1 for the space text_after_trigger = text_before_cursor[trigger_end:cursor_position].lstrip() - start_position = -(len(text_after_trigger)) + if self.prefix: + if not text_after_trigger.startswith(self.prefix): + return + text_after_trigger = text_after_trigger[len(self.prefix) :] + start_position = -len(text_after_trigger) # Load all available agent names try: @@ -410,25 +419,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: @@ -564,7 +558,6 @@ def _normalize_emoji_spacing(text: str) -> str: # palette remap (Level 3) restyles the prompt to the chosen theme. PROMPT_STYLES = { "puppy": "bold ansimagenta", - "owner": "bold ansiwhite", "agent": "bold ansiblue", "model": "bold ansicyan", "cwd": "bold ansigreen", @@ -609,12 +602,15 @@ def get_prompt_with_active_model(base: str = ">>> "): cwd_display = cwd return FormattedText( [ - ("class:puppy", f"{puppy}"), + ("class:puppy class:tui.header", f"{puppy}"), ("", " "), - ("class:agent", f"[{_normalize_emoji_spacing(agent_display)}] "), - ("class:model", model_display + " "), - ("class:cwd", "(" + str(cwd_display) + ") "), - ("class:arrow", str(base)), + ( + "class:agent class:tui.label", + f"[{_normalize_emoji_spacing(agent_display)}] ", + ), + ("class:model class:tui.title", model_display + " "), + ("class:cwd class:tui.muted", "(" + str(cwd_display) + ") "), + ("class:arrow class:tui.help-key", str(base)), ] ) @@ -677,6 +673,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: @@ -699,6 +718,7 @@ async def get_input_with_combined_completion( AgentCompleter(trigger="/a"), AgentCompleter(trigger="/switch-agent"), AgentCompleter(trigger="/sa"), + AgentCompleter(trigger="/fork", prefix="@"), MCPCompleter(trigger="/mcp"), SkillsCompleter(trigger="/skills"), OllamaSetupCompleter(), @@ -789,6 +809,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 @@ -934,7 +960,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( @@ -949,36 +975,29 @@ def handle_image_paste_f3(event): # NOTE: the style field must be a str — `None` crashes `to_formatted_text`. if isinstance(prompt_str, str): prompt_str = FormattedText([("", prompt_str)]) - style = Style.from_dict( + prompt_text_color = on_prompt_text_color() + default_input_style = f"fg:{prompt_text_color}" if prompt_text_color else "" + local_style = Style.from_dict( { - # Keys must AVOID the 'class:' prefix – that prefix is used only when - # tagging tokens in `FormattedText`. See prompt_toolkit docs. - **PROMPT_STYLES, - "attachment-placeholder": "italic ansicyan", - # ── Completion menu (pi-inspired minimal look) ──────────────── - # Drop the chunky default highlight bar. Use terminal-default bg - # so the menu blends, then `ansibrightblack` for a soft "dim" - # look that stays readable on both light and dark themes (since - # ANSI bright-black gets remapped to a mid-grey by most terms). - "completion-menu": "bg:default fg:ansibrightblack", - "completion-menu.completion": "bg:default fg:ansibrightblack", - # Selection highlight: use cyan to match the puppy/model banner - # colors (green is already the `cwd` color, and the bright-green - # bold variant was way too loud — Mike's eyeballs filed a complaint). - # `noreverse` is critical: prompt_toolkit's default for this class - # is `fg:#888888 bg:#ffffff reverse`, which would otherwise swap - # our fg/bg and paint a chunky cyan block behind the left column. - "completion-menu.completion.current": "noreverse bg:default fg:ansibrightcyan bold", - # `display_meta` is the right-hand path column — same dim treatment, - # italicized to differentiate it from the primary column. - "completion-menu.meta.completion": "bg:default fg:ansibrightblack italic", - "completion-menu.meta.completion.current": "bg:default fg:ansicyan italic", - # Scrollbar — keep it subtle so it doesn't reintroduce a bar. - "completion-menu.multi-column-meta": "bg:default", - "scrollbar.background": "bg:default", - "scrollbar.button": "bg:ansibrightblack", + # Keep the prompt useful without the theme plugin. With the plugin + # active, its semantic root supplies the palette underneath these + # structural rules while an explicit prompt-text override still wins. + "": default_input_style, + "attachment-placeholder": "italic", + # Suppress prompt_toolkit's fixed white/grey/reverse completion + # presentation. Colors now inherit from the semantic theme root; + # only hierarchy and emphasis belong to this local component. + "completion-menu": "noreverse", + "completion-menu.completion": "noreverse", + "completion-menu.completion.current": "noreverse bold underline", + "completion-menu.meta.completion": "noreverse italic", + "completion-menu.meta.completion.current": "noreverse italic bold", + "completion-menu.multi-column-meta": "noreverse", + "scrollbar.background": "noreverse", + "scrollbar.button": "noreverse bold", } ) + style = on_prompt_toolkit_style(local_style) text = await session.prompt_async(prompt_str, style=style) # NOTE: We used to call update_model_in_input(text) here to handle /model and /m # commands at the prompt level, but that prevented the command handler from running 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/command_line/set_menu.py b/code_puppy/command_line/set_menu.py index ca206a52b..4a34536a4 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, @@ -46,6 +49,7 @@ reset_value, ) from code_puppy.tools.command_runner import set_awaiting_user_input +from code_puppy.callbacks import on_prompt_toolkit_style PAGE_SIZE = 12 @@ -118,7 +122,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( @@ -195,7 +199,11 @@ async def _prompt_for_value( prompt = ( f"New value for '{setting.key}' (current: {current_val or '(not set)'}): " ) - session = PromptSession(prompt, is_password=setting.sensitive) + session = PromptSession( + prompt, + is_password=setting.sensitive, + style=on_prompt_toolkit_style(), + ) try: new_val = await session.prompt_async() except KeyboardInterrupt: @@ -340,6 +348,7 @@ def update_display() -> None: key_bindings=kb, full_screen=False, mouse_support=False, + style=on_prompt_toolkit_style(), ) set_awaiting_user_input(True) 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/code_puppy/command_line/set_menu_render.py b/code_puppy/command_line/set_menu_render.py index 68adbf4ba..a68d5fb68 100644 --- a/code_puppy/command_line/set_menu_render.py +++ b/code_puppy/command_line/set_menu_render.py @@ -15,7 +15,7 @@ KeyValueLine = Tuple[str, str] _DEFAULT_PREFIX = "(Default) " -_DEFAULT_STYLE = "fg:ansibrightblack italic" +_DEFAULT_STYLE = "class:tui.muted" def truncate(text: str, max_len: int = 30) -> str: @@ -50,13 +50,13 @@ def wrap(text: str, width: int = 55) -> List[str]: _KEY_HELP = ( - ("up/down", "Navigate", "fg:ansibrightblack"), - ("left/right", "Page", "fg:ansibrightblack"), - ("Enter", "Edit value", "fg:green"), - ("/", "Search", "fg:ansibrightblack"), - ("r", "Reset to default", "fg:ansibrightblack"), - ("Esc", "Save & Exit", "fg:ansicyan"), - ("Ctrl+C", "Cancel (discard)", "fg:ansired"), + ("up/down", "Navigate", "class:tui.help-key"), + ("left/right", "Page", "class:tui.help-key"), + ("Enter", "Edit value", "class:tui.help-key"), + ("/", "Search", "class:tui.help-key"), + ("r", "Reset to default", "class:tui.help-key"), + ("Esc", "Save & Exit", "class:tui.help-key"), + ("Ctrl+C", "Cancel (discard)", "class:tui.error"), ) @@ -77,12 +77,12 @@ def render_left_panel( total_pages = total_pages_fn(len(entries), page_size) start_idx, end_idx = page_bounds(page, len(entries), page_size) - lines.append(("bold cyan", " Puppy Config Settings")) - lines.append(("fg:ansibrightblack", f" (Page {page + 1}/{max(total_pages, 1)})")) + lines.append(("class:tui.header", " Puppy Config Settings")) + lines.append(("class:tui.muted", f" (Page {page + 1}/{max(total_pages, 1)})")) if in_search_mode: - lines.append(("fg:ansiyellow", f" Searching: '{search_buffer}'")) + lines.append(("class:tui.warning", f" Searching: '{search_buffer}'")) elif search_text: - lines.append(("fg:ansiyellow", f" Filter: '{search_text}'")) + lines.append(("class:tui.warning", f" Filter: '{search_text}'")) lines.append(("", "\n\n")) current_category = "" @@ -91,7 +91,7 @@ def render_left_panel( if entry.category.name != current_category: if current_category: lines.append(("", "\n")) - lines.append(("bold fg:ansiblue", f" {entry.category.name}")) + lines.append(("class:tui.title", f" {entry.category.name}")) lines.append(("", "\n")) current_category = entry.category.name @@ -99,17 +99,17 @@ def render_left_panel( val_display = truncate(value_for_display(entry.setting)) from_default = is_default_value(entry.setting) if is_selected: - lines.append(("fg:ansigreen bold", f"> {entry.setting.display_name}")) - lines.append(("fg:ansigreen", " = ")) + lines.append(("class:tui.selected", f"> {entry.setting.display_name}")) + lines.append(("class:tui.selected", " = ")) if from_default: lines.append((_DEFAULT_STYLE, _DEFAULT_PREFIX)) - lines.append(("fg:ansigreen", val_display)) + lines.append(("class:tui.selected", val_display)) else: lines.append(("", f" {entry.setting.display_name}")) - lines.append(("fg:ansibrightblack", " = ")) + lines.append(("class:tui.muted", " = ")) if from_default: lines.append((_DEFAULT_STYLE, _DEFAULT_PREFIX)) - lines.append(("fg:ansibrightblack", val_display)) + lines.append(("class:tui.muted", val_display)) lines.append(("", "\n")) lines.append(("", "\n")) @@ -136,11 +136,11 @@ def _type_display(setting: Setting, valid_values_str: str) -> str: def render_right_panel(entry: Optional[object]) -> List[KeyValueLine]: """Render the right details panel for the currently-selected entry.""" lines: List[KeyValueLine] = [ - ("bold cyan", " Setting Details"), + ("class:tui.header", " Setting Details"), ("", "\n\n"), ] if entry is None: - lines.append(("fg:ansiyellow", " No setting selected.")) + lines.append(("class:tui.warning", " No setting selected.")) lines.append(("", "\n")) return lines @@ -148,47 +148,47 @@ def render_right_panel(entry: Optional[object]) -> List[KeyValueLine]: valid_values_str = ", ".join(setting.valid_values) for label, value, style in ( - ("Key: ", setting.key, "fg:ansicyan"), - ("Name: ", setting.display_name, "fg:ansigreen"), - ("Category: ", entry.category.name, "fg:ansiblue"), + ("Key: ", setting.key, "class:tui.help-key"), + ("Name: ", setting.display_name, "class:tui.success"), + ("Category: ", entry.category.name, "class:tui.title"), ): - lines.append(("bold", label)) + lines.append(("class:tui.label", label)) lines.append((style, value)) lines.append(("", "\n\n")) - lines.append(("bold", "Type: ")) - lines.append(("fg:ansiyellow", _type_display(setting, valid_values_str))) + lines.append(("class:tui.label", "Type: ")) + lines.append(("class:tui.warning", _type_display(setting, valid_values_str))) lines.append(("", "\n\n")) - lines.append(("bold", "Current Value: ")) + lines.append(("class:tui.label", "Current Value: ")) current = display_value(setting) from_default = is_default_value(setting) if current: if from_default: lines.append((_DEFAULT_STYLE, _DEFAULT_PREFIX)) - lines.append(("fg:ansigreen", current)) + lines.append(("class:tui.success", current)) else: - lines.append(("fg:ansibrightblack", "(not set)")) + lines.append(("class:tui.muted", "(not set)")) if setting.requires_restart: - lines.append(("fg:ansiyellow", " (restart required)")) + lines.append(("class:tui.warning", " (restart required)")) lines.append(("", "\n\n")) - lines.append(("bold", "Description:")) + lines.append(("class:tui.label", "Description:")) lines.append(("", "\n")) for wrapped in wrap(setting.description): - lines.append(("fg:ansibrightblack", " " + wrapped)) + lines.append(("class:tui.muted", " " + wrapped)) lines.append(("", "\n")) lines.append(("", "\n")) if setting.type_hint == "choice" and setting.valid_values: - lines.append(("bold", "Valid Values:")) + lines.append(("class:tui.label", "Valid Values:")) lines.append(("", "\n")) for val in setting.valid_values: if val == current: - lines.append(("fg:ansigreen", f" {val} (current)")) + lines.append(("class:tui.success", f" {val} (current)")) else: - lines.append(("fg:ansibrightblack", f" {val}")) + lines.append(("class:tui.muted", f" {val}")) lines.append(("", "\n")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " Tip: Press Enter to edit this setting.")) + lines.append(("class:tui.muted", " Tip: Press Enter to edit this setting.")) return lines diff --git a/code_puppy/command_line/uc_menu.py b/code_puppy/command_line/uc_menu.py index 228c522b9..f5c5657d4 100644 --- a/code_puppy/command_line/uc_menu.py +++ b/code_puppy/command_line/uc_menu.py @@ -27,6 +27,7 @@ from code_puppy.plugins.universal_constructor.models import UCToolInfo from code_puppy.plugins.universal_constructor.registry import get_registry from code_puppy.tools.command_runner import set_awaiting_user_input +from code_puppy.callbacks import on_prompt_toolkit_style PAGE_SIZE = 10 # Tools per page SOURCE_PAGE_SIZE = 30 # Lines of source per page @@ -208,13 +209,13 @@ def _render_menu_panel( total_pages = get_total_pages(len(tools), PAGE_SIZE) start_idx, end_idx = get_page_bounds(page, len(tools), PAGE_SIZE) - lines.append(("bold", "UC Tools")) - lines.append(("fg:ansibrightblack", f" (Page {page + 1}/{total_pages})")) + lines.append(("class:tui.header", "UC Tools")) + lines.append(("class:tui.muted", f" (Page {page + 1}/{total_pages})")) lines.append(("", "\n\n")) if not tools: - lines.append(("fg:yellow", " No UC tools found.\n")) - lines.append(("fg:ansibrightblack", " Ask the LLM to create one!\n")) + lines.append(("class:tui.warning", " No UC tools found.\n")) + lines.append(("class:tui.muted", " Ask the LLM to create one!\n")) lines.append(("", "\n")) else: for i in range(start_idx, end_idx): @@ -225,37 +226,37 @@ def _render_menu_panel( # Selection indicator if is_selected: - lines.append(("fg:ansigreen", "> ")) - lines.append(("fg:ansigreen bold", safe_name)) + lines.append(("class:tui.selected", "> ")) + lines.append(("class:tui.selected", safe_name)) else: lines.append(("", " ")) lines.append(("", safe_name)) # Status indicator if tool.meta.enabled: - lines.append(("fg:ansigreen", " [on]")) + lines.append(("class:tui.success", " [on]")) else: - lines.append(("fg:ansired", " [off]")) + lines.append(("class:tui.error", " [off]")) # Namespace tag if present if tool.meta.namespace: - lines.append(("fg:ansiblue", f" ({tool.meta.namespace})")) + lines.append(("class:tui.title", f" ({tool.meta.namespace})")) lines.append(("", "\n")) # Navigation hints lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " [up]/[down] ")) + lines.append(("class:tui.help-key", " [up]/[down] ")) lines.append(("", "Navigate\n")) - lines.append(("fg:ansibrightblack", " [left]/[right] ")) + lines.append(("class:tui.help-key", " [left]/[right] ")) lines.append(("", "Page\n")) - lines.append(("fg:green", " Enter ")) + lines.append(("class:tui.help-key", " Enter ")) lines.append(("", "View source\n")) - lines.append(("fg:ansiyellow", " E ")) + lines.append(("class:tui.help-key", " E ")) lines.append(("", "Toggle enabled\n")) - lines.append(("fg:ansired", " D ")) + lines.append(("class:tui.error", " D ")) lines.append(("", "Delete tool\n")) - lines.append(("fg:ansibrightblack", " Esc ")) + lines.append(("class:tui.help-key", " Esc ")) lines.append(("", "Exit")) return lines @@ -272,85 +273,85 @@ def _render_preview_panel(tool: Optional[UCToolInfo]) -> List: """ lines = [] - lines.append(("dim cyan", " TOOL DETAILS")) + lines.append(("class:tui.title", " TOOL DETAILS")) lines.append(("", "\n\n")) if not tool: - lines.append(("fg:yellow", " No tool selected.\n")) - lines.append(("fg:ansibrightblack", " Create some with the LLM!\n")) + lines.append(("class:tui.warning", " No tool selected.\n")) + lines.append(("class:tui.muted", " Create some with the LLM!\n")) return lines safe_name = _sanitize_display_text(tool.meta.name) safe_desc = _sanitize_display_text(tool.meta.description) # Tool name - lines.append(("bold", "Name: ")) - lines.append(("fg:ansicyan", safe_name)) + lines.append(("class:tui.label", "Name: ")) + lines.append(("class:tui.help-key", safe_name)) lines.append(("", "\n\n")) # Full name (with namespace) if tool.meta.namespace: - lines.append(("bold", "Full Name: ")) + lines.append(("class:tui.label", "Full Name: ")) lines.append(("", tool.full_name)) lines.append(("", "\n\n")) # Status - lines.append(("bold", "Status: ")) + lines.append(("class:tui.label", "Status: ")) if tool.meta.enabled: - lines.append(("fg:ansigreen bold", "ENABLED")) + lines.append(("class:tui.success", "ENABLED")) else: - lines.append(("fg:ansired bold", "DISABLED")) + lines.append(("class:tui.error", "DISABLED")) lines.append(("", "\n\n")) # Version - lines.append(("bold", "Version: ")) + lines.append(("class:tui.label", "Version: ")) lines.append(("", tool.meta.version)) lines.append(("", "\n\n")) # Author (if present) if tool.meta.author: - lines.append(("bold", "Author: ")) + lines.append(("class:tui.label", "Author: ")) lines.append(("", tool.meta.author)) lines.append(("", "\n\n")) # Signature - lines.append(("bold", "Signature: ")) - lines.append(("fg:ansiyellow", tool.signature)) + lines.append(("class:tui.label", "Signature: ")) + lines.append(("class:tui.warning", tool.signature)) lines.append(("", "\n\n")) # Description (word-wrapped) - lines.append(("bold", "Description:")) + lines.append(("class:tui.label", "Description:")) lines.append(("", "\n")) words = safe_desc.split() current_line = "" for word in words: if len(current_line) + len(word) + 1 > 50: - lines.append(("fg:ansibrightblack", f" {current_line}")) + lines.append(("class:tui.muted", f" {current_line}")) lines.append(("", "\n")) current_line = word else: current_line = word if not current_line else current_line + " " + word if current_line: - lines.append(("fg:ansibrightblack", f" {current_line}")) + lines.append(("class:tui.muted", f" {current_line}")) lines.append(("", "\n")) lines.append(("", "\n")) # Docstring preview (if available) if tool.docstring: - lines.append(("bold", "Docstring:")) + lines.append(("class:tui.label", "Docstring:")) lines.append(("", "\n")) doc_preview = tool.docstring[:150] if len(tool.docstring) > 150: doc_preview += "..." - lines.append(("fg:ansibrightblack", f" {doc_preview}")) + lines.append(("class:tui.muted", f" {doc_preview}")) lines.append(("", "\n\n")) # Source path - lines.append(("bold", "Source:")) + lines.append(("class:tui.label", "Source:")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", f" {tool.source_path}")) + lines.append(("class:tui.muted", f" {tool.source_path}")) lines.append(("", "\n")) return lines @@ -376,19 +377,19 @@ def _render_source_panel( lines = [] # Header - lines.append(("bold cyan", f" SOURCE: {tool.full_name}")) + lines.append(("class:tui.header", f" SOURCE: {tool.full_name}")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", f" {tool.source_path}")) + lines.append(("class:tui.muted", f" {tool.source_path}")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", "─" * 70)) + lines.append(("class:tui.muted", "─" * 70)) lines.append(("", "\n")) if error: - lines.append(("fg:ansired", f" Error: {error}\n")) + lines.append(("class:tui.error", f" Error: {error}\n")) return lines if not source_lines: - lines.append(("fg:yellow", " (empty file)\n")) + lines.append(("class:tui.warning", " (empty file)\n")) return lines # Calculate visible range @@ -405,7 +406,7 @@ def _render_source_panel( line_content = source_lines[i] # Line number - lines.append(("fg:ansibrightblack", f" {line_num:>{line_num_width}} │ ")) + lines.append(("class:tui.muted", f" {line_num:>{line_num_width}} │ ")) # Basic syntax highlighting highlighted = _highlight_python_line(line_content) @@ -413,7 +414,7 @@ def _render_source_panel( lines.append(("", "\n")) # Footer with scroll info - lines.append(("fg:ansibrightblack", "─" * 70)) + lines.append(("class:tui.muted", "─" * 70)) lines.append(("", "\n")) # Scroll position indicator @@ -421,21 +422,21 @@ def _render_source_panel( total_pages = (total_lines + SOURCE_PAGE_SIZE - 1) // SOURCE_PAGE_SIZE lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", f" Lines {scroll_offset + 1}-{end_offset} of {total_lines}", ) ) - lines.append(("fg:ansibrightblack", f" (Page {current_page}/{total_pages})")) + lines.append(("class:tui.muted", f" (Page {current_page}/{total_pages})")) lines.append(("", "\n\n")) # Navigation hints for source view - lines.append(("fg:ansibrightblack", " [up]/[down] ")) + lines.append(("class:tui.muted", " [up]/[down] ")) lines.append(("", "Scroll\n")) - lines.append(("fg:ansibrightblack", " [PgUp]/[PgDn] ")) + lines.append(("class:tui.muted", " [PgUp]/[PgDn] ")) lines.append(("", "Page\n")) - lines.append(("fg:ansiyellow", " Esc/Q ")) + lines.append(("class:tui.help-key", " Esc/Q ")) lines.append(("", "Back to list\n")) - lines.append(("fg:ansibrightred", " Ctrl+C ")) + lines.append(("class:tui.error", " Ctrl+C ")) lines.append(("", "Exit")) return lines @@ -494,7 +495,7 @@ def _highlight_python_line(line: str) -> List[Tuple[str, str]]: # Check for comments if line.lstrip().startswith("#"): - result.append(("fg:ansibrightblack italic", line)) + result.append(("class:tui.muted italic", line)) return result # Check for strings (simplified) @@ -815,6 +816,7 @@ def _source_exit(event): key_bindings=list_kb, full_screen=False, mouse_support=False, + style=on_prompt_toolkit_style(), ) else: # Source view @@ -825,6 +827,7 @@ def _source_exit(event): key_bindings=source_kb, full_screen=False, mouse_support=False, + style=on_prompt_toolkit_style(), ) await app.run_async() diff --git a/code_puppy/config.py b/code_puppy/config.py index 62afbc9ea..df9932239 100644 --- a/code_puppy/config.py +++ b/code_puppy/config.py @@ -53,18 +53,27 @@ def _get_xdg_dir(env_var: str, fallback: str) -> str: SKILLS_DIR = os.path.join(DATA_DIR, "skills") CONTEXTS_DIR = os.path.join(DATA_DIR, "contexts") -# OAuth plugin model files (XDG_DATA_HOME) -GEMINI_MODELS_FILE = os.path.join(DATA_DIR, "gemini_models.json") -CHATGPT_MODELS_FILE = os.path.join(DATA_DIR, "chatgpt_models.json") -CLAUDE_MODELS_FILE = os.path.join(DATA_DIR, "claude_models.json") -COPILOT_MODELS_FILE = os.path.join(DATA_DIR, "copilot_models.json") - # Cache files (XDG_CACHE_HOME) AUTOSAVE_DIR = os.path.join(CACHE_DIR, "autosaves") +WS_SESSION_DIR = os.path.join(CACHE_DIR, "ws_sessions") + + +def get_ws_sessions_dir() -> pathlib.Path: + """Return the ws_sessions directory path and ensure it exists.""" + p = pathlib.Path(WS_SESSION_DIR) + p.mkdir(parents=True, exist_ok=True, mode=0o700) + return p + # State files (XDG_STATE_HOME) COMMAND_HISTORY_FILE = os.path.join(STATE_DIR, "command_history.txt") +# OAuth plugin model files (XDG_DATA_HOME) +GEMINI_MODELS_FILE = os.path.join(DATA_DIR, "gemini_models.json") +CHATGPT_MODELS_FILE = os.path.join(DATA_DIR, "chatgpt_models.json") +CLAUDE_MODELS_FILE = os.path.join(DATA_DIR, "claude_models.json") +COPILOT_MODELS_FILE = os.path.join(DATA_DIR, "copilot_models.json") + def get_subagent_verbose() -> bool: """Return True if sub-agent verbose output is enabled (default False). @@ -728,8 +737,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 +747,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( @@ -1193,18 +1202,29 @@ def initialize_command_history_file(): ) -def get_yolo_mode(): - """ - Checks puppy.cfg for 'yolo_mode' (case-insensitive in value only). - Defaults to True if not set. - Allowed values for ON: 1, '1', 'true', 'yes', 'on' (all case-insensitive for value). - """ +_cli_yolo_override: Optional[bool] = None + + +def set_cli_yolo_override(value: Optional[bool]) -> None: + """Set a process-local YOLO value supplied by the CLI.""" + global _cli_yolo_override + _cli_yolo_override = value + + +def get_cli_yolo_override() -> Optional[bool]: + """Return the process-local CLI override, if one was supplied.""" + return _cli_yolo_override + + +def get_yolo_mode() -> bool: + """Return effective YOLO mode using CLI > persisted config precedence.""" + if _cli_yolo_override is not None: + return _cli_yolo_override + true_vals = {"1", "true", "yes", "on"} cfg_val = get_value("yolo_mode") if cfg_val is not None: - if str(cfg_val).strip().lower() in true_vals: - return True - return False + return str(cfg_val).strip().lower() in true_vals return True @@ -1421,8 +1441,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. @@ -1600,9 +1620,55 @@ def set_diff_highlight_style(style: str): pass -# Defaults for diff highlight colors — single source of truth. +# Diff colors use these only when no curated terminal palette is active. _DEFAULT_DIFF_ADDITION_HEX = "#0b1f0b" # darker green _DEFAULT_DIFF_DELETION_HEX = "#390e1a" # wine +_THEME_PALETTE_CONFIG_KEY = "osc_palette_json" + + +def _blend_hex(background: str, accent: str, accent_weight: float) -> str: + """Blend an accent into a background, returning a subtle highlight.""" + background_rgb = tuple( + int(background[index : index + 2], 16) for index in (1, 3, 5) + ) + accent_rgb = tuple(int(accent[index : index + 2], 16) for index in (1, 3, 5)) + channels = ( + round(base * (1 - accent_weight) + highlight * accent_weight) + for base, highlight in zip(background_rgb, accent_rgb) + ) + return "#" + "".join(f"{channel:02x}" for channel in channels) + + +def _theme_diff_defaults() -> tuple[str, str]: + """Derive quiet add/remove backgrounds from the active terminal theme. + + ANSI slots 2 and 1 are the theme's semantic green and red. Blending them + into the terminal background keeps highlights legible on both dark and + light themes instead of dropping a dark green rectangle onto everything. + """ + raw_palette = get_value(_THEME_PALETTE_CONFIG_KEY) + if not raw_palette: + return _DEFAULT_DIFF_ADDITION_HEX, _DEFAULT_DIFF_DELETION_HEX + + try: + palette = json.loads(raw_palette) + background = _coerce_to_hex(palette.get("bg"), "") + if not background: + raise ValueError("theme has no valid background") + ansi = palette.get("ansi") or [] + addition = _coerce_to_hex(ansi[2] if len(ansi) > 2 else "#2ea043", "#2ea043") + deletion = _coerce_to_hex(ansi[1] if len(ansi) > 1 else "#cf222e", "#cf222e") + red, green, blue = ( + int(background[index : index + 2], 16) for index in (1, 3, 5) + ) + luminance = (0.2126 * red + 0.7152 * green + 0.0722 * blue) / 255 + accent_weight = 0.14 if luminance > 0.5 else 0.20 + return ( + _blend_hex(background, addition, accent_weight), + _blend_hex(background, deletion, accent_weight), + ) + except (AttributeError, TypeError, ValueError, json.JSONDecodeError): + return _DEFAULT_DIFF_ADDITION_HEX, _DEFAULT_DIFF_DELETION_HEX def _coerce_to_hex(value: Optional[str], fallback: str) -> str: @@ -1640,12 +1706,12 @@ def _coerce_to_hex(value: Optional[str], fallback: str) -> str: def get_diff_addition_color() -> str: """Get the base color for diff additions, always as a valid '#RRGGBB' hex. - Falls back to the default darker green if the configured value is missing - or unparseable. + An explicit ``/diff`` choice wins. When unset, the color is derived from + the active theme's background and semantic green. """ - return _coerce_to_hex( - get_value("highlight_addition_color"), _DEFAULT_DIFF_ADDITION_HEX - ) + configured = get_value("highlight_addition_color") + theme_default, _ = _theme_diff_defaults() + return _coerce_to_hex(configured, theme_default) def set_diff_addition_color(color: str): @@ -1664,12 +1730,12 @@ def set_diff_addition_color(color: str): def get_diff_deletion_color() -> str: """Get the base color for diff deletions, always as a valid '#RRGGBB' hex. - Falls back to the default wine if the configured value is missing or - unparseable. + An explicit ``/diff`` choice wins. When unset, the color is derived from + the active theme's background and semantic red. """ - return _coerce_to_hex( - get_value("highlight_deletion_color"), _DEFAULT_DIFF_DELETION_HEX - ) + configured = get_value("highlight_deletion_color") + _, theme_default = _theme_diff_defaults() + return _coerce_to_hex(configured, theme_default) def set_diff_deletion_color(color: str): 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/__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/bar_painters.py b/code_puppy/messaging/bar_painters.py index 0b200cbfb..48a10401b 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,18 @@ 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}") + if prompt_sgr: + # Prefix styling uses full SGR resets between color runs. Restore + # the theme foreground after each reset so user text doesn't + # fall back to terminal-default white on light themes. + rendered = rendered.replace("\x1b[0m", f"\x1b[0m{prompt_sgr}") + 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/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/editor_completion.py b/code_puppy/messaging/editor_completion.py index 69f874cb9..48984f290 100644 --- a/code_puppy/messaging/editor_completion.py +++ b/code_puppy/messaging/editor_completion.py @@ -72,6 +72,7 @@ def build_completer(): AgentCompleter(trigger="/a"), AgentCompleter(trigger="/switch-agent"), AgentCompleter(trigger="/sa"), + AgentCompleter(trigger="/fork", prefix="@"), MCPCompleter(trigger="/mcp"), SkillsCompleter(trigger="/skills"), OllamaSetupCompleter(), 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/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/messaging/pause_controller.py b/code_puppy/messaging/pause_controller.py index a28c440db..4f47657a1 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 # ========================================================================= @@ -261,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..6055afe06 100644 --- a/code_puppy/messaging/renderers.py +++ b/code_puppy/messaging/renderers.py @@ -289,7 +289,19 @@ 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.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: try: console.print(Markdown(content)) diff --git a/code_puppy/messaging/rich_renderer.py b/code_puppy/messaging/rich_renderer.py index 195a6adb6..c71af3564 100644 --- a/code_puppy/messaging/rich_renderer.py +++ b/code_puppy/messaging/rich_renderer.py @@ -150,7 +150,12 @@ def __init__( import threading self._bus = bus - self._console = console or Console() + self._console = console or Console(highlight=False) + # ReprHighlighter is useful in a Python REPL, but this is a themed + # message bus: it unpredictably recolors numbers, parentheses, paths, + # and shell/grep output. Explicit markup and syntax renderables still + # carry their own styles, so disable only the console's implicit pass. + self._console.highlighter = None self._styles = styles or DEFAULT_STYLES.copy() self._running = False self._thread: Optional[threading.Thread] = None @@ -563,7 +568,9 @@ def _render_text(self, msg: TextMessage) -> None: prefix = self._get_level_prefix(msg.level) # Escape Rich markup to prevent crashes from malformed tags safe_text = escape_rich_markup(msg.text) - self._console.print(f"{prefix}{safe_text}", style=style) + # Rich's default repr highlighter restyles numbers and parentheses, + # leaking terminal-profile white through otherwise themed info lines. + self._console.print(f"{prefix}{safe_text}", style=style, highlight=False) def _get_level_prefix(self, level: MessageLevel) -> str: """Get a prefix icon for the message level.""" @@ -603,8 +610,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 +701,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 +714,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 +736,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 +755,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 +786,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 +820,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 +846,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 +858,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 +910,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.""" @@ -1000,7 +1002,12 @@ def _render_agent_response(self, msg: AgentResponseMessage) -> None: # Content (markdown or plain) if msg.is_markdown: md = Markdown(msg.content) - self._console.print(md) + # Give colorless Markdown styles (bold/headings) an explicit base + # foreground. Some terminals ignore OSC 10, making Rich resets leak + # the terminal profile's white into otherwise themed responses. + from code_puppy.callbacks import on_prompt_text_color + + self._console.print(md, style=on_prompt_text_color()) else: self._console.print(msg.content) @@ -1072,7 +1079,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 +1269,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/run_ui.py b/code_puppy/messaging/run_ui.py index a4949bb73..c6451b77a 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", @@ -150,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. @@ -162,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 @@ -177,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/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/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/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/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/agent_skills/metadata.py b/code_puppy/plugins/agent_skills/metadata.py index 405bf80b8..87cf049e5 100644 --- a/code_puppy/plugins/agent_skills/metadata.py +++ b/code_puppy/plugins/agent_skills/metadata.py @@ -111,7 +111,7 @@ def parse_skill_metadata(skill_path: Path) -> Optional[SkillMetadata]: SkillMetadata if successful, None if parsing fails. """ if not skill_path.exists(): - logger.warning(f"Skill path does not exist: {skill_path}") + logger.debug(f"Skill path does not exist: {skill_path}") return None skill_md_path = skill_path / "SKILL.md" @@ -173,12 +173,12 @@ def load_full_skill_content(skill_path: Path) -> Optional[str]: Full file content as string, or None if not found. """ if not skill_path.exists(): - logger.warning(f"Skill path does not exist: {skill_path}") + logger.debug(f"Skill path does not exist: {skill_path}") return None skill_md_path = skill_path / "SKILL.md" if not skill_md_path.exists(): - logger.warning(f"SKILL.md not found in skill directory: {skill_path}") + logger.debug(f"SKILL.md not found in skill directory: {skill_path}") return None try: @@ -200,7 +200,7 @@ def get_skill_resources(skill_path: Path) -> List[Path]: List of paths to resource files (excluding SKILL.md). """ if not skill_path.exists(): - logger.warning(f"Skill path does not exist: {skill_path}") + logger.debug(f"Skill path does not exist: {skill_path}") return [] if not skill_path.is_dir(): diff --git a/code_puppy/plugins/agent_skills/skills_install_menu.py b/code_puppy/plugins/agent_skills/skills_install_menu.py index 10d381c6c..d6a51278e 100644 --- a/code_puppy/plugins/agent_skills/skills_install_menu.py +++ b/code_puppy/plugins/agent_skills/skills_install_menu.py @@ -36,6 +36,7 @@ from code_puppy.plugins.agent_skills.installer import InstallResult from code_puppy.plugins.agent_skills.skill_catalog import SkillCatalogEntry, catalog from code_puppy.tools.command_runner import set_awaiting_user_input +from code_puppy.callbacks import on_prompt_toolkit_style logger = logging.getLogger(__name__) @@ -165,21 +166,21 @@ def _render_navigation_hints(self, lines: List) -> None: """Render keyboard shortcut hints at the bottom.""" lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " ↑/↓ ")) + lines.append(("class:tui.help-key", " ↑/↓ ")) lines.append(("", "Navigate ")) - lines.append(("fg:ansibrightblack", "←/→ ")) + lines.append(("class:tui.help-key", "←/→ ")) lines.append(("", "Page\n")) if self.view_mode == "categories": - lines.append(("fg:ansigreen", " Enter ")) + lines.append(("class:tui.help-key", " Enter ")) lines.append(("", "Browse Skills\n")) else: - lines.append(("fg:ansigreen", " Enter ")) + lines.append(("class:tui.success", " Enter ")) lines.append(("", "Install Skill\n")) - lines.append(("fg:ansibrightblack", " Esc/Back ")) + lines.append(("class:tui.help-key", " Esc/Back ")) lines.append(("", "Back\n")) - lines.append(("fg:ansired", " Ctrl+C ")) + lines.append(("class:tui.help-key", " Ctrl+C ")) lines.append(("", "Cancel")) def _render_category_list(self) -> List: @@ -187,15 +188,15 @@ def _render_category_list(self) -> List: lines = [] - lines.append(("bold cyan", " 📂 CATEGORIES")) + lines.append(("class:tui.title", " 📂 CATEGORIES")) lines.append(("", "\n\n")) if not self.categories: - lines.append(("fg:ansiyellow", " No remote categories available.")) + lines.append(("class:tui.warning", " No remote categories available.")) lines.append(("", "\n")) lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", " (Remote catalog unavailable or empty)\n", ) ) @@ -223,15 +224,15 @@ def _render_category_list(self) -> List: label = f"{prefix}{icon} {category} ({count})" if is_selected: - lines.append(("fg:ansibrightcyan bold", label)) + lines.append(("class:tui.selected", label)) else: - lines.append(("fg:ansibrightblack", label)) + lines.append(("class:tui.muted", label)) lines.append(("", "\n")) lines.append(("", "\n")) if total_pages > 1: lines.append( - ("fg:ansibrightblack", f" Page {self.current_page + 1}/{total_pages}") + ("class:tui.muted", f" Page {self.current_page + 1}/{total_pages}") ) lines.append(("", "\n")) @@ -244,17 +245,17 @@ def _render_skill_list(self) -> List: lines = [] if not self.current_category: - lines.append(("fg:ansiyellow", " No category selected.")) + lines.append(("class:tui.warning", " No category selected.")) lines.append(("", "\n\n")) self._render_navigation_hints(lines) return lines icon = self._get_category_icon(self.current_category) - lines.append(("bold cyan", f" {icon} {self.current_category.upper()}")) + lines.append(("class:tui.title", f" {icon} {self.current_category.upper()}")) lines.append(("", "\n\n")) if not self.current_skills: - lines.append(("fg:ansiyellow", " No skills in this category.")) + lines.append(("class:tui.warning", " No skills in this category.")) lines.append(("", "\n\n")) self._render_navigation_hints(lines) return lines @@ -270,13 +271,13 @@ def _render_skill_list(self) -> List: installed = is_skill_installed(entry.id) status_icon = "✓" if installed else "○" - status_style = "fg:ansigreen" if installed else "fg:ansibrightblack" + status_style = "class:tui.success" if installed else "class:tui.muted" prefix = " > " if is_selected else " " label = f"{prefix}{status_icon} {entry.display_name}" if is_selected: - lines.append(("fg:ansibrightcyan bold", label)) + lines.append(("class:tui.selected", label)) else: lines.append((status_style, label)) @@ -285,7 +286,7 @@ def _render_skill_list(self) -> List: lines.append(("", "\n")) if total_pages > 1: lines.append( - ("fg:ansibrightblack", f" Page {self.current_page + 1}/{total_pages}") + ("class:tui.muted", f" Page {self.current_page + 1}/{total_pages}") ) lines.append(("", "\n")) @@ -297,13 +298,13 @@ def _render_details(self) -> List: lines = [] - lines.append(("bold cyan", " 📋 DETAILS")) + lines.append(("class:tui.title", " 📋 DETAILS")) lines.append(("", "\n\n")) if self.view_mode == "categories": category = self._get_current_category() if not category: - lines.append(("fg:ansiyellow", " No category selected.")) + lines.append(("class:tui.warning", " No category selected.")) return lines icon = self._get_category_icon(category) @@ -316,7 +317,7 @@ def _render_details(self) -> List: except Exception: skills = [] - lines.append(("fg:ansibrightblack", f" {len(skills)} skills available")) + lines.append(("class:tui.muted", f" {len(skills)} skills available")) lines.append(("", "\n\n")) # Show a preview of the first few skills @@ -324,19 +325,19 @@ def _render_details(self) -> List: lines.append(("bold", " Preview:")) lines.append(("", "\n")) for entry in skills[:6]: - lines.append(("fg:ansibrightblack", f" • {entry.display_name}")) + lines.append(("class:tui.muted", f" • {entry.display_name}")) lines.append(("", "\n")) return lines entry = self._get_current_skill() if not entry: - lines.append(("fg:ansiyellow", " No skill selected.")) + lines.append(("class:tui.warning", " No skill selected.")) return lines installed = is_skill_installed(entry.id) installed_text = "Installed" if installed else "Not installed" - installed_style = "fg:ansigreen" if installed else "fg:ansiyellow" + installed_style = "class:tui.success" if installed else "class:tui.warning" lines.append(("bold", f" {entry.display_name}")) lines.append(("", "\n")) @@ -345,57 +346,59 @@ def _render_details(self) -> List: lines.append(("bold", " ID:")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", f" {entry.id}")) + lines.append(("class:tui.muted", f" {entry.id}")) lines.append(("", "\n\n")) lines.append(("bold", " Description:")) lines.append(("", "\n")) desc = entry.description or "No description available" for line in _wrap_text(desc, 56): - lines.append(("fg:ansibrightblack", f" {line}")) + lines.append(("class:tui.muted", f" {line}")) lines.append(("", "\n")) lines.append(("", "\n")) lines.append(("bold", " Category:")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", f" {entry.category}")) + lines.append(("class:tui.muted", f" {entry.category}")) lines.append(("", "\n\n")) lines.append(("bold", " Tags:")) lines.append(("", "\n")) tags = entry.tags or [] - lines.append(("fg:ansicyan", f" {', '.join(tags) if tags else '(none)'}")) + lines.append( + ("class:tui.header", f" {', '.join(tags) if tags else '(none)'}") + ) lines.append(("", "\n\n")) lines.append(("bold", " Contents:")) lines.append(("", "\n")) lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", f" scripts: {'yes' if entry.has_scripts else 'no'}", ) ) lines.append(("", "\n")) lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", f" references: {'yes' if entry.has_references else 'no'}", ) ) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", f" files: {entry.file_count}")) + lines.append(("class:tui.muted", f" files: {entry.file_count}")) lines.append(("", "\n\n")) lines.append(("bold", " Download:")) lines.append(("", "\n")) lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", f" size: {_format_bytes(entry.zip_size_bytes)}", ) ) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", f" url: {entry.download_url}")) + lines.append(("class:tui.muted", f" url: {entry.download_url}")) lines.append(("", "\n")) return lines @@ -587,6 +590,7 @@ def _(event): key_bindings=kb, full_screen=False, mouse_support=False, + style=on_prompt_toolkit_style(), ) set_awaiting_user_input(True) diff --git a/code_puppy/plugins/agent_skills/skills_menu.py b/code_puppy/plugins/agent_skills/skills_menu.py index fe86f6a15..c855da35b 100644 --- a/code_puppy/plugins/agent_skills/skills_menu.py +++ b/code_puppy/plugins/agent_skills/skills_menu.py @@ -43,6 +43,7 @@ parse_skill_metadata, ) from code_puppy.tools.command_runner import set_awaiting_user_input +from code_puppy.callbacks import on_prompt_toolkit_style PAGE_SIZE = 15 # Items per page @@ -119,19 +120,19 @@ def _render_skill_list(self) -> List: lines = [] # Header with status - status_color = "fg:ansigreen" if self.skills_enabled else "fg:ansired" + status_style = "class:tui.success" if self.skills_enabled else "class:tui.error" status_text = "ENABLED" if self.skills_enabled else "DISABLED" - lines.append((status_color, f" Skills: {status_text}")) + lines.append((status_style, f" Skills: {status_text}")) lines.append(("", "\n\n")) if not self.skills: - lines.append(("fg:ansiyellow", " No skills found.")) + lines.append(("class:tui.warning", " No skills found.")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " Create skills in:")) + lines.append(("class:tui.muted", " Create skills in:")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " ~/.code_puppy/skills/")) + lines.append(("class:tui.muted", " ~/.code_puppy/skills/")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " ./skills/")) + lines.append(("class:tui.muted", " ./skills/")) lines.append(("", "\n\n")) self._render_navigation_hints(lines) return lines @@ -150,7 +151,7 @@ def _render_skill_list(self) -> List: # Status icon status_icon = "✗" if is_disabled else "✓" - status_style = "fg:ansired" if is_disabled else "fg:ansigreen" + status_style = "class:tui.error" if is_disabled else "class:tui.success" # Get skill name from metadata if available metadata = self._get_skill_metadata(skill) @@ -160,20 +161,20 @@ def _render_skill_list(self) -> List: prefix = " > " if is_selected else " " if is_selected: - lines.append(("bold", prefix)) - lines.append((status_style + " bold", status_icon)) - lines.append(("bold", f" {display_name}")) + lines.append(("class:tui.selected", prefix)) + lines.append(("class:tui.selected", status_icon)) + lines.append(("class:tui.selected", f" {display_name}")) else: lines.append(("", prefix)) lines.append((status_style, status_icon)) - lines.append(("fg:ansibrightblack", f" {display_name}")) + lines.append(("class:tui.muted", f" {display_name}")) lines.append(("", "\n")) # Pagination info lines.append(("", "\n")) lines.append( - ("fg:ansibrightblack", f" Page {self.current_page + 1}/{total_pages}") + ("class:tui.muted", f" Page {self.current_page + 1}/{total_pages}") ) lines.append(("", "\n")) @@ -183,39 +184,39 @@ def _render_skill_list(self) -> List: def _render_navigation_hints(self, lines: List) -> None: """Render navigation hints at the bottom.""" lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " ↑/↓ or j/k ")) + lines.append(("class:tui.help-key", " ↑/↓ or j/k ")) lines.append(("", "Navigate ")) - lines.append(("fg:ansibrightblack", "←/→ ")) + lines.append(("class:tui.help-key", "←/→ ")) lines.append(("", "Page\n")) - lines.append(("fg:ansigreen", " Enter ")) + lines.append(("class:tui.help-key", " Enter ")) lines.append(("", "Toggle ")) - lines.append(("fg:ansicyan", " t ")) + lines.append(("class:tui.help-key", " t ")) lines.append(("", "Toggle System\n")) - lines.append(("fg:ansimagenta", " Ctrl+A ")) + lines.append(("class:tui.help-key", " Ctrl+A ")) lines.append(("", "Add Dir ")) - lines.append(("fg:ansiyellow", " Ctrl+D ")) + lines.append(("class:tui.help-key", " Ctrl+D ")) lines.append(("", "Show Dirs\n")) - lines.append(("fg:ansimagenta", " i ")) + lines.append(("class:tui.help-key", " i ")) lines.append(("", "Install from catalog\n")) - lines.append(("fg:ansiyellow", " r ")) + lines.append(("class:tui.help-key", " r ")) lines.append(("", "Refresh ")) - lines.append(("fg:ansired", " q ")) + lines.append(("class:tui.help-key", " q ")) lines.append(("", "Exit")) def _render_skill_details(self) -> List: """Render the skill details panel.""" lines = [] - lines.append(("dim cyan", " SKILL DETAILS")) + lines.append(("class:tui.title dim", " SKILL DETAILS")) lines.append(("", "\n\n")) skill = self._get_current_skill() if not skill: - lines.append(("fg:ansiyellow", " No skill selected.")) + lines.append(("class:tui.warning", " No skill selected.")) lines.append(("", "\n\n")) - lines.append(("fg:ansibrightblack", " Select a skill from the list")) + lines.append(("class:tui.muted", " Select a skill from the list")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " to view its details.")) + lines.append(("class:tui.muted", " to view its details.")) return lines metadata = self._get_skill_metadata(skill) @@ -223,7 +224,7 @@ def _render_skill_details(self) -> List: # Status status_text = "Disabled" if is_disabled else "Enabled" - status_style = "fg:ansired bold" if is_disabled else "fg:ansigreen bold" + status_style = "class:tui.error" if is_disabled else "class:tui.success" lines.append(("bold", " Status: ")) lines.append((status_style, status_text)) lines.append(("", "\n\n")) @@ -241,7 +242,7 @@ def _render_skill_details(self) -> List: desc = metadata.description wrapped = self._wrap_text(desc, 50) for line in wrapped: - lines.append(("fg:ansibrightblack", f" {line}")) + lines.append(("class:tui.muted", f" {line}")) lines.append(("", "\n")) lines.append(("", "\n")) @@ -250,7 +251,7 @@ def _render_skill_details(self) -> List: lines.append(("bold", " Tags:")) lines.append(("", "\n")) tags_str = ", ".join(metadata.tags) - lines.append(("fg:ansicyan", f" {tags_str}")) + lines.append(("class:tui.header", f" {tags_str}")) lines.append(("", "\n\n")) # Resources @@ -260,11 +261,11 @@ def _render_skill_details(self) -> List: lines.append(("", "\n")) for resource in resources[:5]: # Show first 5 resource_name = getattr(resource, "name", str(resource)) - lines.append(("fg:ansiyellow", f" • {resource_name}")) + lines.append(("class:tui.warning", f" • {resource_name}")) lines.append(("", "\n")) if len(resources) > 5: lines.append( - ("fg:ansibrightblack", f" ... and {len(resources) - 5} more") + ("class:tui.muted", f" ... and {len(resources) - 5} more") ) lines.append(("", "\n")) lines.append(("", "\n")) @@ -273,13 +274,11 @@ def _render_skill_details(self) -> List: # No metadata available lines.append(("bold", f" {skill.name}")) lines.append(("", "\n\n")) - lines.append(("fg:ansiyellow", " No metadata available")) + lines.append(("class:tui.warning", " No metadata available")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " Add a SKILL.md with frontmatter to")) + lines.append(("class:tui.muted", " Add a SKILL.md with frontmatter to")) lines.append(("", "\n")) - lines.append( - ("fg:ansibrightblack", " define name, description, and tags.") - ) + lines.append(("class:tui.muted", " define name, description, and tags.")) lines.append(("", "\n\n")) # Path @@ -288,7 +287,7 @@ def _render_skill_details(self) -> List: path_str = str(skill.path) if len(path_str) > 45: path_str = "..." + path_str[-42:] - lines.append(("fg:ansibrightblack", f" {path_str}")) + lines.append(("class:tui.muted", f" {path_str}")) lines.append(("", "\n")) return lines @@ -454,6 +453,7 @@ def _(event): key_bindings=kb, full_screen=False, mouse_support=False, + style=on_prompt_toolkit_style(), ) set_awaiting_user_input(True) diff --git a/code_puppy/plugins/chatgpt_oauth/config.py b/code_puppy/plugins/chatgpt_oauth/config.py index 22bca6355..20b4bbaaf 100644 --- a/code_puppy/plugins/chatgpt_oauth/config.py +++ b/code_puppy/plugins/chatgpt_oauth/config.py @@ -22,11 +22,11 @@ # 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) - "client_version": "0.72.0", + "client_version": "0.144.1", "originator": "codex_cli_rs", } 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 81104c8df..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,9 +136,12 @@ 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.72.0") + client_version = CHATGPT_OAUTH_CONFIG.get("client_version", "0.144.1") headers = { "ChatGPT-Account-Id": account_id, @@ -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 5d06df035..0b7643a36 100644 --- a/code_puppy/plugins/chatgpt_oauth/utils.py +++ b/code_puppy/plugins/chatgpt_oauth/utils.py @@ -344,37 +344,19 @@ 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", "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", -] - -# 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", + "codex-auto-review", ] # 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 +373,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 +394,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 +431,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 +468,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 +486,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 +505,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/code_puppy/plugins/fork/register_callbacks.py b/code_puppy/plugins/fork/register_callbacks.py index f9529ee51..621e10850 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(Text(""), header, Markdown(response)) + + # --------------------------------------------------------------------------- # Parsing # --------------------------------------------------------------------------- @@ -123,15 +164,34 @@ def _resolve_agent_name(requested: Optional[str]) -> Optional[str]: # Fork lifecycle # --------------------------------------------------------------------------- async def _run_fork(agent_name: str, prompt: str): - """Run the sub-agent. Reuses the invoke_agent implementation wholesale.""" + """Run the sub-agent and publish its tool-equivalent completion signal.""" + from code_puppy.callbacks import on_post_tool_call from code_puppy.tools.subagent_invocation import _invoke_agent_impl - # ``context`` is unused by the implementation; forks have no RunContext. - return await _invoke_agent_impl( - context=None, - agent_name=agent_name, - prompt=prompt, - ) + started_at = time.monotonic() + result = None + try: + # ``context`` is unused by the implementation; forks have no RunContext. + result = await _invoke_agent_impl( + context=None, + agent_name=agent_name, + prompt=prompt, + emit_response_message=False, + ) + return result + finally: + # /fork deliberately calls the shared implementation directly rather than + # entering through pydantic-ai's tool wrapper. Publish the same completion + # lifecycle event so observers (notably the sub-agent panel) can retire + # the row instead of displaying "writing response" forever. + if result is not None: + await on_post_tool_call( + "invoke_agent", + {"agent_name": agent_name, "prompt": prompt}, + result, + (time.monotonic() - started_at) * 1000, + {"detached_fork": True}, + ) def _on_fork_done(fork_id: int, task: asyncio.Task) -> None: @@ -160,7 +220,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 +246,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/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/plugins/hook_manager/hooks_menu.py b/code_puppy/plugins/hook_manager/hooks_menu.py index 7c87cd4b1..72696b7cc 100644 --- a/code_puppy/plugins/hook_manager/hooks_menu.py +++ b/code_puppy/plugins/hook_manager/hooks_menu.py @@ -30,21 +30,10 @@ save_hooks_config, toggle_hook_enabled, ) +from code_puppy.callbacks import on_prompt_toolkit_style PAGE_SIZE = 12 -# Colour palette (matches skills_menu palette) -_C_ENABLED = "fg:ansigreen" -_C_DISABLED = "fg:ansired" -_C_SELECTED_BG = "bold" -_C_DIM = "fg:ansibrightblack" -_C_CYAN = "fg:ansicyan" -_C_YELLOW = "fg:ansiyellow" -_C_MAGENTA = "fg:ansimagenta" -_C_HEADER = "dim cyan" -_C_GLOBAL = "fg:ansiblue" -_C_PROJECT = "fg:ansigreen" - class HooksMenu: """Interactive TUI for managing hooks from both global and project sources.""" @@ -230,16 +219,18 @@ def _render_list(self) -> List: project_count = sum(1 for e in self.entries if e.source == "project") global_count = sum(1 for e in self.entries if e.source == "global") - header_color = _C_ENABLED if enabled_count > 0 else _C_DISABLED + header_color = "class:tui.success" if enabled_count > 0 else "class:tui.error" lines.append((header_color, f" Hooks: {enabled_count}/{total} enabled")) lines.append(("", f" ({project_count} project, {global_count} global)\n\n")) if not self.entries: - lines.append((_C_YELLOW, " No hooks configured.")) + lines.append(("class:tui.warning", " No hooks configured.")) lines.append(("", "\n")) - lines.append((_C_DIM, " Add hooks to .claude/settings.json (project)")) + lines.append( + ("class:tui.muted", " Add hooks to .claude/settings.json (project)") + ) lines.append(("", "\n")) - lines.append((_C_DIM, " or ~/.code_puppy/hooks.json (global)")) + lines.append(("class:tui.muted", " or ~/.code_puppy/hooks.json (global)")) lines.append(("", "\n\n")) self._render_nav_hints(lines) return lines @@ -251,34 +242,40 @@ def _render_list(self) -> List: entry = self.entries[i] is_selected = i == self.selected_idx status_icon = "✓" if entry.enabled else "✗" - status_style = _C_ENABLED if entry.enabled else _C_DISABLED + status_style = "class:tui.success" if entry.enabled else "class:tui.error" source_indicator = "🌍" if entry.source == "global" else "📁" prefix = " > " if is_selected else " " if is_selected: - lines.append((_C_SELECTED_BG, prefix)) + lines.append(("class:tui.selected", prefix)) lines.append((status_style + " bold", status_icon)) lines.append( - (_C_SELECTED_BG, f" {source_indicator} [{entry.event_type}]") + ("class:tui.selected", f" {source_indicator} [{entry.event_type}]") ) - lines.append((_C_SELECTED_BG, f" {entry.display_matcher}")) + lines.append(("class:tui.selected", f" {entry.display_matcher}")) else: lines.append(("", prefix)) lines.append((status_style, status_icon)) - source_color = _C_GLOBAL if entry.source == "global" else _C_PROJECT + source_color = ( + "class:tui.header" + if entry.source == "global" + else "class:tui.success" + ) lines.append((source_color, f" {source_indicator}")) - lines.append((_C_DIM, f" [{entry.event_type}]")) + lines.append(("class:tui.muted", f" [{entry.event_type}]")) lines.append(("", f" {entry.display_matcher}")) lines.append(("", "\n")) lines.append(("", "\n")) - lines.append((_C_DIM, f" Page {self.current_page + 1}/{total_pages}")) + lines.append( + ("class:tui.muted", f" Page {self.current_page + 1}/{total_pages}") + ) lines.append(("", "\n")) # Status message (shows result of last action) if self.status_message: lines.append(("", "\n")) - lines.append((_C_CYAN, f" {self.status_message}")) + lines.append(("class:tui.header", f" {self.status_message}")) lines.append(("", "\n")) self._render_nav_hints(lines) @@ -287,39 +284,43 @@ def _render_list(self) -> List: def _render_nav_hints(self, lines: List) -> None: """Append keyboard shortcut hints to lines.""" lines.append(("", "\n")) - lines.append((_C_DIM, " ↑/↓ j/k ")) + lines.append(("class:tui.help-key", " ↑/↓ j/k ")) lines.append(("", "Navigate\n")) - lines.append((_C_ENABLED, " Enter ")) + lines.append(("class:tui.help-key", " Enter ")) lines.append(("", "Toggle enable/disable\n")) - lines.append((_C_DISABLED, " d ")) + lines.append(("class:tui.help-key", " d ")) lines.append(("", "Delete hook\n")) - lines.append((_C_YELLOW, " A ")) + lines.append(("class:tui.help-key", " A ")) lines.append(("", "Enable all\n")) - lines.append((_C_MAGENTA, " D ")) + lines.append(("class:tui.help-key", " D ")) lines.append(("", "Disable all\n")) - lines.append((_C_YELLOW, " r ")) + lines.append(("class:tui.help-key", " r ")) lines.append(("", "Refresh\n")) - lines.append((_C_DISABLED, " q/Esc ")) + lines.append(("class:tui.help-key", " q/Esc ")) lines.append(("", "Exit")) def _render_detail(self) -> List: """Render the right-hand hook detail panel.""" lines: List = [] - lines.append((_C_HEADER, " HOOK DETAILS")) + lines.append(("class:tui.title dim", " HOOK DETAILS")) lines.append(("", "\n\n")) entry = self._current_entry() if entry is None: - lines.append((_C_YELLOW, " No hook selected.")) + lines.append(("class:tui.warning", " No hook selected.")) lines.append(("", "\n\n")) - lines.append((_C_DIM, " Select a hook from the list")) + lines.append(("class:tui.muted", " Select a hook from the list")) lines.append(("", "\n")) - lines.append((_C_DIM, " to view its details.")) + lines.append(("class:tui.muted", " to view its details.")) return lines # Status badge status_text = "Enabled" if entry.enabled else "Disabled" - status_style = _C_ENABLED + " bold" if entry.enabled else _C_DISABLED + " bold" + status_style = ( + "class:tui.success" + " bold" + if entry.enabled + else "class:tui.error" + " bold" + ) lines.append(("bold", " Status: ")) lines.append((status_style, status_text)) lines.append(("", "\n\n")) @@ -330,27 +331,29 @@ def _render_detail(self) -> List: if entry.source == "global" else "Project (.claude/settings.json)" ) - source_color = _C_GLOBAL if entry.source == "global" else _C_PROJECT + source_color = ( + "class:tui.header" if entry.source == "global" else "class:tui.success" + ) lines.append(("bold", " Source: ")) lines.append((source_color, source_label)) lines.append(("", "\n\n")) # Event type lines.append(("bold", " Event: ")) - lines.append((_C_CYAN, entry.event_type)) + lines.append(("class:tui.header", entry.event_type)) lines.append(("", "\n\n")) # Matcher lines.append(("bold", " Matcher: ")) lines.append(("", "\n")) for chunk in _wrap(entry.matcher, 50): - lines.append((_C_YELLOW, f" {chunk}")) + lines.append(("class:tui.warning", f" {chunk}")) lines.append(("", "\n")) lines.append(("", "\n")) # Type lines.append(("bold", " Type: ")) - lines.append((_C_DIM, entry.hook_type)) + lines.append(("class:tui.muted", entry.hook_type)) lines.append(("", "\n\n")) # Command / prompt @@ -358,26 +361,29 @@ def _render_detail(self) -> List: lines.append(("bold", f" {label}")) lines.append(("", "\n")) for chunk in _wrap(entry.command, 50): - lines.append((_C_DIM, f" {chunk}")) + lines.append(("class:tui.muted", f" {chunk}")) lines.append(("", "\n")) lines.append(("", "\n")) # Timeout lines.append(("bold", " Timeout: ")) - lines.append((_C_DIM, f"{entry.timeout} ms")) + lines.append(("class:tui.muted", f"{entry.timeout} ms")) lines.append(("", "\n\n")) # Hook ID if entry.hook_id: lines.append(("bold", " ID: ")) - lines.append((_C_DIM, entry.hook_id)) + lines.append(("class:tui.muted", entry.hook_id)) lines.append(("", "\n\n")) # Config location hint - lines.append((_C_DIM, f" Stored in {source_label}")) + lines.append(("class:tui.muted", f" Stored in {source_label}")) lines.append(("", "\n")) lines.append( - (_C_DIM, f" group #{entry._group_index} hook #{entry._hook_index}") + ( + "class:tui.muted", + f" group #{entry._group_index} hook #{entry._hook_index}", + ) ) lines.append(("", "\n")) @@ -493,6 +499,7 @@ def _quit_ctrl_c(event): key_bindings=kb, full_screen=False, mouse_support=False, + style=on_prompt_toolkit_style(), ) try: 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/code_puppy/plugins/plugin_list/plugins_menu.py b/code_puppy/plugins/plugin_list/plugins_menu.py index 05f5f2e24..814d4d23c 100644 --- a/code_puppy/plugins/plugin_list/plugins_menu.py +++ b/code_puppy/plugins/plugin_list/plugins_menu.py @@ -43,6 +43,7 @@ render_list, ) from code_puppy.tools.command_runner import set_awaiting_user_input +from code_puppy.callbacks import on_prompt_toolkit_style PAGE_SIZE = 20 @@ -392,6 +393,7 @@ def run(self) -> Optional[str]: key_bindings=kb, full_screen=False, mouse_support=False, + style=on_prompt_toolkit_style(), ) # Live resize: prompt_toolkit re-renders on SIGWINCH automatically, but diff --git a/code_puppy/plugins/plugin_list/plugins_menu_render.py b/code_puppy/plugins/plugin_list/plugins_menu_render.py index 9170cc518..5f0969909 100644 --- a/code_puppy/plugins/plugin_list/plugins_menu_render.py +++ b/code_puppy/plugins/plugin_list/plugins_menu_render.py @@ -76,7 +76,7 @@ def render_list(menu: "PluginsMenu") -> Fragments: if menu.lock_builtin and menu.hidden_builtin_count: lines.append( ( - "fg:ansibrightblack", + "class:tui.muted", f" {menu.hidden_builtin_count} builtin plugins are " f"managed and hidden.", ) @@ -85,9 +85,9 @@ def render_list(menu: "PluginsMenu") -> Fragments: if not menu.plugins: if menu.lock_builtin and menu.hidden_builtin_count: - lines.append(("fg:ansiyellow", " No user or project plugins loaded.")) + lines.append(("class:tui.warning", " No user or project plugins loaded.")) else: - lines.append(("fg:ansiyellow", " No plugins loaded.")) + lines.append(("class:tui.warning", " No plugins loaded.")) lines.append(("", "\n")) _render_hints(lines) return lines @@ -104,26 +104,28 @@ def render_list(menu: "PluginsMenu") -> Fragments: if entry.status == "loaded": icon = "x" if is_disabled else "+" - icon_style = "fg:ansired" if is_disabled else "fg:ansigreen" + icon_style = "class:tui.error" if is_disabled else "class:tui.success" else: # Trust-gated project plugin: discovered but never imported. icon = "!" - icon_style = "fg:ansired" if entry.status == "error" else "fg:ansiyellow" + icon_style = ( + "class:tui.error" if entry.status == "error" else "class:tui.warning" + ) prefix = " > " if is_selected else " " if is_selected: - lines.append(("bold", prefix)) - lines.append((icon_style + " bold", icon)) - lines.append(("bold", f" {entry.name}")) + lines.append(("class:tui.selected", prefix)) + lines.append(("class:tui.selected", icon)) + lines.append(("class:tui.selected", f" {entry.name}")) else: lines.append(("", prefix)) lines.append((icon_style, icon)) - lines.append(("fg:ansibrightblack", f" {entry.name}")) + lines.append(("class:tui.muted", f" {entry.name}")) lines.append(("", "\n")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", f" Page {menu.current_page + 1}/{total_pages}")) + lines.append(("class:tui.muted", f" Page {menu.current_page + 1}/{total_pages}")) lines.append(("", "\n")) _render_hints(lines) @@ -135,12 +137,12 @@ def _render_hints(lines: Fragments) -> None: # across both split-pane inspectors. See _build_key_bindings docstring # for the full mental model. hints: List[Tuple[str, str, str]] = [ - ("fg:ansibrightblack", " up/down or j/k ", "Navigate"), - ("fg:ansibrightblack", " PgUp/PgDn ", "Page"), - ("fg:ansibrightblack", " g / G ", "First / Last"), - ("fg:ansibrightblack", " h/l or \u2190/\u2192 ", "Scroll details"), - ("fg:ansigreen", " Enter ", "Toggle / Enable"), - ("fg:ansired", " q / Esc ", "Exit"), + ("class:tui.help-key", " up/down or j/k ", "Navigate"), + ("class:tui.help-key", " PgUp/PgDn ", "Page"), + ("class:tui.help-key", " g / G ", "First / Last"), + ("class:tui.help-key", " h/l or \u2190/\u2192 ", "Scroll details"), + ("class:tui.help-key", " Enter ", "Toggle / Enable"), + ("class:tui.help-key", " q / Esc ", "Exit"), ] lines.append(("", "\n")) for i, (style, key_text, label) in enumerate(hints): @@ -156,23 +158,23 @@ def _render_hints(lines: Fragments) -> None: # ``{name}`` is substituted with the plugin name in the hint text. _GATE_STATUS_DETAILS = { "untrusted": ( - "fg:ansiyellow bold", + "class:tui.warning", "Not enabled (project plugins are disabled by default)", "Press Enter to review its files and enable it.", ), "changed": ( - "fg:ansiyellow bold", + "class:tui.warning", "Changed since you accepted it", "Its files were modified after you trusted it, so trust was " "revoked. Press Enter to re-review and re-enable.", ), "disabled": ( - "fg:ansired bold", + "class:tui.error", "Disabled (trusted, but not loaded)", "Press Enter to load it.", ), "error": ( - "fg:ansired bold", + "class:tui.error", "Failed to load", "Check the logs, then press Enter to retry.", ), @@ -183,15 +185,13 @@ def _render_gate_status(lines: Fragments, menu: "PluginsMenu", entry) -> None: """Status + hint for a project plugin that was never imported.""" style, label, hint = _GATE_STATUS_DETAILS.get( entry.status, - ("fg:ansiyellow bold", entry.status, "See '/plugins list' for details."), + ("class:tui.warning", entry.status, "See '/plugins list' for details."), ) lines.append(("bold", " Status: ")) lines.append((style, label)) lines.append(("", "\n")) inner = max(20, menu._detail_cols - 4) - _append_wrapped( - lines, "fg:ansibrightblack", " ", hint.format(name=entry.name), inner - ) + _append_wrapped(lines, "class:tui.muted", " ", hint.format(name=entry.name), inner) lines.append(("", "\n")) @@ -217,9 +217,7 @@ def render_trust_modal(menu: "PluginsMenu") -> Fragments: if entry.status == "changed" else "has never been enabled for this project" ) - _append_wrapped( - lines, "fg:ansiyellow bold", " ", f"'{entry.name}' {reason}.", inner - ) + _append_wrapped(lines, "class:tui.warning", " ", f"'{entry.name}' {reason}.", inner) lines.append(("", "\n")) _append_wrapped( lines, @@ -239,12 +237,12 @@ def render_trust_modal(menu: "PluginsMenu") -> Fragments: listing = plugin_file_listing(Path(menu.project_dir) / entry.name, limit=8) for row in listing.splitlines(): for piece in wrap_text(row.strip(), inner - 2): - lines.append(("fg:ansicyan", f" {piece}")) + lines.append(("class:tui.header", f" {piece}")) lines.append(("", "\n")) lines.append(("", "\n")) if menu.trust_error: - _append_wrapped(lines, "fg:ansired bold", " ", menu.trust_error, inner) + _append_wrapped(lines, "class:tui.error", " ", menu.trust_error, inner) lines.append(("", "\n")) _append_wrapped( @@ -260,18 +258,18 @@ def render_trust_modal(menu: "PluginsMenu") -> Fragments: def render_detail(menu: "PluginsMenu") -> Fragments: lines: Fragments = [] - lines.append(("dim cyan", " PLUGIN DETAILS")) + lines.append(("class:tui.title dim", " PLUGIN DETAILS")) lines.append(("", "\n\n")) # Outcome of the most recent enable/activate action, if any. if menu.trust_feedback: inner = max(20, menu._detail_cols - 4) - _append_wrapped(lines, "fg:ansigreen", " ", menu.trust_feedback, inner) + _append_wrapped(lines, "class:tui.success", " ", menu.trust_feedback, inner) lines.append(("", "\n")) entry = menu._current() if not entry: - lines.append(("fg:ansiyellow", " No plugin selected.")) + lines.append(("class:tui.warning", " No plugin selected.")) return lines from code_puppy.plugins.plugin_list import plugin_meta @@ -298,7 +296,7 @@ def render_detail(menu: "PluginsMenu") -> Fragments: lines.append(("bold", " Project:")) lines.append(("", "\n")) inner = max(20, menu._detail_cols - 6) - _append_wrapped(lines, "fg:ansibrightblack", " ", menu.project_dir, inner) + _append_wrapped(lines, "class:tui.muted", " ", menu.project_dir, inner) lines.append(("", "\n")) if entry.status != "loaded": @@ -309,20 +307,18 @@ def render_detail(menu: "PluginsMenu") -> Fragments: lines.append(("bold", " Path:")) lines.append(("", "\n")) inner = max(20, menu._detail_cols - 6) - _append_wrapped(lines, "fg:ansibrightblack", " ", path, inner) + _append_wrapped(lines, "class:tui.muted", " ", path, inner) return lines lines.append(("bold", " Status: ")) if is_disabled: - lines.append(("fg:ansired bold", "Disabled")) + lines.append(("class:tui.error", "Disabled")) lines.append(("", "\n")) - lines.append( - ("fg:ansibrightblack", " Callbacks are skipped at dispatch time.") - ) + lines.append(("class:tui.muted", " Callbacks are skipped at dispatch time.")) else: - lines.append(("fg:ansigreen bold", "Enabled")) + lines.append(("class:tui.success", "Enabled")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", " All callbacks are active.")) + lines.append(("class:tui.muted", " All callbacks are active.")) lines.append(("", "\n\n")) _render_contributions(lines, menu, entry) @@ -332,10 +328,10 @@ def render_detail(menu: "PluginsMenu") -> Fragments: lines.append(("", "\n")) if hooks: for hook in hooks: - lines.append(("fg:ansicyan", f" • {hook}")) + lines.append(("class:tui.header", f" • {hook}")) lines.append(("", "\n")) else: - lines.append(("fg:ansibrightblack", " (none registered)")) + lines.append(("class:tui.muted", " (none registered)")) lines.append(("", "\n")) lines.append(("", "\n")) @@ -345,7 +341,7 @@ def render_detail(menu: "PluginsMenu") -> Fragments: lines.append(("bold", " Path:")) lines.append(("", "\n")) inner = max(20, menu._detail_cols - 6) # -4 indent, -2 frame border - _append_wrapped(lines, "fg:ansibrightblack", " ", path, inner) + _append_wrapped(lines, "class:tui.muted", " ", path, inner) lines.append(("", "\n")) if menu._changed: @@ -356,7 +352,7 @@ def render_detail(menu: "PluginsMenu") -> Fragments: inner = max(20, menu._detail_cols - 4) # -2 indent, -2 frame border _append_wrapped( lines, - "fg:ansiyellow bold", + "class:tui.warning", " ", "Restart Code Puppy for changes to take effect.", inner, @@ -407,13 +403,13 @@ def _render_contributions( items = contributions.get(key) if not items: continue - lines.append(("fg:ansimagenta", f" {label}:")) + lines.append(("class:tui.title", f" {label}:")) lines.append(("", "\n")) for item in items: pieces = wrap_text(str(item), inner) for idx, piece in enumerate(pieces): prefix = bullet_prefix if idx == 0 else cont_prefix - lines.append(("fg:ansicyan", f"{prefix}{piece}")) + lines.append(("class:tui.header", f"{prefix}{piece}")) lines.append(("", "\n")) lines.append(("", "\n")) diff --git a/code_puppy/plugins/prune/prune_menu.py b/code_puppy/plugins/prune/prune_menu.py index 6c941f704..f0003ba8d 100644 --- a/code_puppy/plugins/prune/prune_menu.py +++ b/code_puppy/plugins/prune/prune_menu.py @@ -39,6 +39,7 @@ Row, ) from code_puppy.plugins.prune.prune_render import render_detail, render_list +from code_puppy.callbacks import on_prompt_toolkit_style class PruneMenu: @@ -274,6 +275,7 @@ def run(self) -> Optional[PruneSelection]: key_bindings=self._build_keybindings(), full_screen=False, mouse_support=False, + style=on_prompt_toolkit_style(), ) try: diff --git a/code_puppy/plugins/prune/prune_model.py b/code_puppy/plugins/prune/prune_model.py index 95af79389..246185bb6 100644 --- a/code_puppy/plugins/prune/prune_model.py +++ b/code_puppy/plugins/prune/prune_model.py @@ -9,24 +9,6 @@ from dataclasses import dataclass, field from typing import Any, List, Optional, Set -# ── palette ──────────────────────────────────────────────────────────────── -# Style strings for prompt_toolkit's formatted-text tuples. Centralised here -# so prune_menu.py and prune_render.py reference one source of truth. - -C_CURSOR = "bold fg:ansicyan" -C_USER = "fg:ansigreen" -C_ASSISTANT = "fg:ansiblue" -C_TOOL = "fg:ansiyellow" -C_WRITE = "fg:ansimagenta" -C_SHELL = "fg:ansired" -C_DIM = "fg:ansibrightblack" -C_HEADER = "dim cyan" -C_FOOTER_OK = "fg:ansigreen" -C_FOOTER_WARN = "fg:ansiyellow" -C_CHECKED = "bold fg:ansired" -C_IMPLIED = "fg:ansibrightblack" -C_SYSTEM = "bold fg:ansicyan" - # ── tool-name → side-effect classification ───────────────────────────────── _WRITE_TOOLS = { @@ -116,14 +98,14 @@ def is_locked(self) -> bool: @property def role_color(self) -> str: if self.role == "system": - return C_SYSTEM + return "class:tui.title" if self.role == "user": - return C_USER + return "class:tui.success" if self.role == "assistant": - return C_ASSISTANT + return "class:tui.header" if self.role == "tool-return": - return C_TOOL - return C_DIM + return "class:tui.warning" + return "class:tui.muted" @dataclass @@ -542,19 +524,6 @@ def annotate_context_window( __all__ = [ - "C_ASSISTANT", - "C_CHECKED", - "C_CURSOR", - "C_DIM", - "C_FOOTER_OK", - "C_FOOTER_WARN", - "C_HEADER", - "C_IMPLIED", - "C_SHELL", - "C_SYSTEM", - "C_TOOL", - "C_USER", - "C_WRITE", "ContextBudget", "MessageEntry", "PruneSelection", diff --git a/code_puppy/plugins/prune/prune_render.py b/code_puppy/plugins/prune/prune_render.py index 32901a9c4..3eaee808a 100644 --- a/code_puppy/plugins/prune/prune_render.py +++ b/code_puppy/plugins/prune/prune_render.py @@ -11,14 +11,6 @@ from typing import Any, List, Tuple from code_puppy.plugins.prune.prune_model import ( - C_CHECKED, - C_CURSOR, - C_DIM, - C_FOOTER_OK, - C_FOOTER_WARN, - C_HEADER, - C_SHELL, - C_TOOL, ContextBudget, MessageEntry, Row, @@ -32,10 +24,10 @@ def ctx_indicator(entry: MessageEntry) -> Tuple[str, str]: """Return (glyph, style) for the in-context indicator.""" if entry.in_context is True: - return ("●", C_FOOTER_OK) + return ("●", "class:tui.success") if entry.in_context is False: - return ("○", C_DIM) - return ("·", C_DIM) + return ("○", "class:tui.muted") + return ("·", "class:tui.muted") def tokens_str(entry: MessageEntry) -> str: @@ -138,23 +130,23 @@ def _format_value_lines(value: Any) -> List[str]: def render_budget_line(budget: ContextBudget) -> List[tuple]: if not budget.available or budget.context_length is None: - return [(C_DIM, "context: unavailable\n")] + return [("class:tui.muted", "context: unavailable\n")] total = budget.total_used or 0 pct = budget.percent_used or 0.0 if pct < 70: - style = C_FOOTER_OK + style = "class:tui.success" elif pct < 90: - style = C_FOOTER_WARN + style = "class:tui.warning" else: - style = C_SHELL + style = "class:tui.error" parts: List[tuple] = [ ( style, f"context: {total:,}/{budget.context_length:,} tokens ({pct:.0f}%)", ), - (C_DIM, f" overhead: {budget.overhead_tokens or 0:,}t\n"), + ("class:tui.muted", f" overhead: {budget.overhead_tokens or 0:,}t\n"), ] if budget.out_of_context_messages > 0: # On its own line — these tokens won't fit in this turn's window, @@ -163,7 +155,7 @@ def render_budget_line(budget: ContextBudget) -> List[tuple]: # history; pop or prune newer ones to slide them back in.) parts.append( ( - C_SHELL, + "class:tui.error", f" ↯ {budget.out_of_context_tokens:,}t in " f"{budget.out_of_context_messages} older msg(s) out of context\n", ) @@ -185,12 +177,12 @@ def render_legend() -> List[tuple]: budget are noisy ``●`` messages. """ return [ - (C_DIM, " legend: "), - (C_FOOTER_OK, "● in context"), - (C_DIM, " "), - (C_DIM, "○ out of context"), - (C_DIM, " "), - (C_DIM, "· unknown\n"), + ("class:tui.muted", " legend: "), + ("class:tui.success", "● in context"), + ("class:tui.muted", " "), + ("class:tui.muted", "○ out of context"), + ("class:tui.muted", " "), + ("class:tui.muted", "· unknown\n"), ] @@ -206,7 +198,7 @@ def render_list(menu: Any) -> List[tuple]: out: List[tuple] = [] out.append( ( - C_HEADER, + "class:tui.title dim", " prune ↓/↑ move space toggle a all c clear enter confirm q quit\n", ) ) @@ -222,7 +214,7 @@ def render_list(menu: Any) -> List[tuple]: hidden_below = max(0, total - bottom) if hidden_above: - out.append((C_DIM, f" ↑ {hidden_above} more above\n")) + out.append(("class:tui.muted", f" ↑ {hidden_above} more above\n")) else: out.append(("", "\n")) @@ -230,7 +222,7 @@ def render_list(menu: Any) -> List[tuple]: render_row(menu, idx, menu.rows[idx], out) if hidden_below: - out.append((C_DIM, f" ↓ {hidden_below} more below\n")) + out.append(("class:tui.muted", f" ↓ {hidden_below} more below\n")) else: out.append(("", "\n")) @@ -238,19 +230,19 @@ def render_list(menu: Any) -> List[tuple]: out.append(("", "\n")) out.append( ( - C_FOOTER_OK, + "class:tui.success", f" enter = remove {msg_count} message(s)\n", ) ) out.extend(render_legend()) - out.append((C_DIM, f" cursor {menu.cursor + 1}/{total}\n")) + out.append(("class:tui.muted", f" cursor {menu.cursor + 1}/{total}\n")) return out def render_row(menu: Any, idx: int, row: Row, out: List[tuple]) -> None: is_cursor = idx == menu.cursor cursor_marker = "▶ " if is_cursor else " " - cursor_style = C_CURSOR if is_cursor else "" + cursor_style = "class:tui.selected" if is_cursor else "" checked = menu._row_is_checked(row) entry = menu.entries[row.message_idx] @@ -258,11 +250,11 @@ def render_row(menu: Any, idx: int, row: Row, out: List[tuple]) -> None: # Locked rows (sys bundle or history[0]) get a distinct "locked" box # so the user can see at a glance that this isn't toggleable. if entry.is_locked: - box, box_style = "[-] ", C_DIM + box, box_style = "[-] ", "class:tui.muted" elif checked: - box, box_style = "[x] ", C_CHECKED + box, box_style = "[x] ", "class:tui.error" else: - box, box_style = "[ ] ", C_DIM + box, box_style = "[ ] ", "class:tui.muted" role_short = { "system": "sys ", @@ -271,7 +263,7 @@ def render_row(menu: Any, idx: int, row: Row, out: List[tuple]) -> None: "tool-return": "tool ", "unknown": "? ", }.get(entry.role, "? ") - row_style = C_CHECKED if checked else entry.role_color + row_style = "class:tui.error" if checked else entry.role_color out.append((cursor_style, cursor_marker)) out.append((box_style, box)) ctx_glyph, ctx_style = ctx_indicator(entry) @@ -285,7 +277,7 @@ def render_row(menu: Any, idx: int, row: Row, out: List[tuple]) -> None: ) tok = tokens_str(entry) if tok: - out.append((C_DIM, f" {tok}")) + out.append(("class:tui.muted", f" {tok}")) out.append(("", "\n")) @@ -306,7 +298,7 @@ def render_detail(menu: Any) -> List[tuple]: if menu._selection_has_side_effects(): out.append( ( - C_SHELL, + "class:tui.error", " ⚠ selection includes side-effecting tool calls; " "pruning does NOT roll them back\n", ) @@ -315,7 +307,7 @@ def render_detail(menu: Any) -> List[tuple]: def _render_message_detail(entry: MessageEntry, out: List[tuple]) -> None: - out.append((C_HEADER, f" {entry.role} (message)\n")) + out.append(("class:tui.title dim", f" {entry.role} (message)\n")) tok = f" · ~{entry.tokens}t" if entry.tokens is not None else "" location = f"history index {entry.history_index}" # Surface thinking presence in the metadata line so it's clear at a @@ -327,7 +319,7 @@ def _render_message_detail(entry: MessageEntry, out: List[tuple]) -> None: ) out.append( ( - C_DIM, + "class:tui.muted", f" {location} · " f"{len(entry.tool_calls)} tool call(s){tok}{ctx_detail_text(entry)}" f"{thinking_note}\n\n", @@ -338,17 +330,17 @@ def _render_message_detail(entry: MessageEntry, out: List[tuple]) -> None: if entry.tool_calls: out.append(("", "\n")) - out.append((C_HEADER, " tool calls (will go with message):\n")) + out.append(("class:tui.title dim", " tool calls (will go with message):\n")) for tc in entry.tool_calls: _render_tool_call_block(tc, out) def _render_tool_call_block(tc: ToolCallInfo, out: List[tuple]) -> None: """Inline tool-call block inside a message detail view, full args.""" - out.append((C_TOOL, f" {tc.icon} {tc.name}")) + out.append(("class:tui.warning", f" {tc.icon} {tc.name}")) out.append(("", "\n")) for line in format_args_full(tc.full_args, tc.args_preview): - out.append((C_DIM, f" {line}\n")) + out.append(("class:tui.muted", f" {line}\n")) __all__ = [ diff --git a/code_puppy/plugins/puppy_spinner/picker.py b/code_puppy/plugins/puppy_spinner/picker.py index 63714b24b..ddadee759 100644 --- a/code_puppy/plugins/puppy_spinner/picker.py +++ b/code_puppy/plugins/puppy_spinner/picker.py @@ -38,6 +38,7 @@ ) from . import spinners as sp +from code_puppy.callbacks import on_prompt_toolkit_style #: Invalidation cadence for the preview animation. Pinned to the speed #: floor so even a spinner dialed all the way down to MIN_INTERVAL @@ -77,7 +78,7 @@ def _format_menu( start, end = get_page_bounds(page, len(entries), PAGE_SIZE) lines: list[tuple[str, str]] = [ - ("bold", " Spinners"), + ("class:tui.header", " Spinners"), ("", "\n\n"), ] for i in range(start, end): @@ -87,17 +88,17 @@ def _format_menu( prefix = " > " if is_selected else " " if is_selected: - lines.append(("bold", prefix)) - lines.append(("fg:ansigreen bold", icon)) - lines.append(("bold", f" {spinner.name}")) + lines.append(("class:tui.selected", prefix)) + lines.append(("class:tui.selected", icon)) + lines.append(("class:tui.selected", f" {spinner.name}")) else: - lines.append(("", prefix)) - lines.append(("fg:ansigreen", icon)) - lines.append(("fg:ansibrightblack", f" {spinner.name}")) + lines.append(("class:tui.body", prefix)) + lines.append(("class:tui.success", icon)) + lines.append(("class:tui.muted", f" {spinner.name}")) lines.append(("", "\n")) lines.append(("", "\n")) - lines.append(("fg:ansibrightblack", f" Page {page + 1}/{total_pages}")) + lines.append(("class:tui.muted", f" Page {page + 1}/{total_pages}")) lines.append(("", "\n")) _render_hints(lines) @@ -108,13 +109,13 @@ def _format_menu( #: common cell width in ``_render_hints`` so the action column always #: lines up -- hand-padding drifted the moment a row gained an arrow. _HINTS = [ - ("fg:ansibrightblack", "up/down or j/k", "Navigate"), - ("fg:ansibrightblack", "PgUp/PgDn", "Page"), - ("fg:ansibrightblack", "g / G", "First / Last"), - ("fg:ansibrightblack", "-/+ or \u2190/\u2192", "Slower / Faster"), - ("fg:ansibrightblack", "i", "Init spinners.json"), - ("fg:ansigreen", "Enter", "Apply"), - ("fg:ansired", "q / Esc", "Exit"), + ("class:tui.help-key", "up/down or j/k", "Navigate"), + ("class:tui.help-key", "PgUp/PgDn", "Page"), + ("class:tui.help-key", "g / G", "First / Last"), + ("class:tui.help-key", "-/+ or \u2190/\u2192", "Slower / Faster"), + ("class:tui.help-key", "i", "Init spinners.json"), + ("class:tui.success", "Enter", "Apply"), + ("class:tui.error", "q / Esc", "Exit"), ] @@ -126,7 +127,9 @@ def _render_hints(lines: list[tuple[str, str]]) -> None: lines.append(("", "\n")) for i, (style, keys, label) in enumerate(_HINTS): lines.append((style, f" {keys.ljust(key_col)}")) - lines.append(("", label if i == len(_HINTS) - 1 else f"{label}\n")) + lines.append( + ("class:tui.help", label if i == len(_HINTS) - 1 else f"{label}\n") + ) def _format_preview( @@ -145,9 +148,9 @@ def _format_preview( elapsed = time.monotonic() - started_at frame = spinner.frames[int(elapsed / effective) % len(spinner.frames)] lines: list[tuple[str, str]] = [ - ("dim cyan", " LIVE PREVIEW"), + ("class:tui.title", " LIVE PREVIEW"), ("", "\n\n"), - ("bold", f" {spinner.name}"), + ("class:tui.label", f" {spinner.name}"), ("", "\n\n"), ] if spinner.description: @@ -155,32 +158,32 @@ def _format_preview( lines.append(("", "\n\n")) lines.extend( [ - ("bold", f" {frame}"), + ("class:tui.label", f" {frame}"), ("", "\n\n"), - ("bold", " Source: "), - ("", spinner.source), + ("class:tui.label", " Source: "), + ("class:tui.body", spinner.source), ("", "\n"), - ("bold", " Frames: "), - ("", str(len(spinner.frames))), + ("class:tui.label", " Frames: "), + ("class:tui.body", str(len(spinner.frames))), ("", "\n"), - ("bold", " Interval: "), - ("", f"{effective:.2f}s"), + ("class:tui.label", " Interval: "), + ("class:tui.body", f"{effective:.2f}s"), ( - "fg:ansiyellow", + "class:tui.warning", " (custom -- Enter saves it)" if interval is not None else "", ), ("", "\n\n"), - ("bold", " Custom spinners:"), + ("class:tui.label", " Custom spinners:"), ("", "\n"), - ("fg:ansibrightblack", f" {sp.USER_SPINNERS_FILE}"), + ("class:tui.muted", f" {sp.USER_SPINNERS_FILE}"), ("", "\n"), - ("fg:ansibrightblack", " (press i to write a starter file)"), + ("class:tui.muted", " (press i to write a starter file)"), ("", "\n"), ] ) if notice: lines.append(("", "\n")) - lines.append(("fg:ansiyellow", f" {notice}")) + lines.append(("class:tui.warning", f" {notice}")) lines.append(("", "\n")) return FormattedText(lines) @@ -349,6 +352,7 @@ def _(event): key_bindings=kb, full_screen=False, mouse_support=False, + style=on_prompt_toolkit_style(), ) async def _animate() -> None: 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/plugins/statusline/payload.py b/code_puppy/plugins/statusline/payload.py index da5935232..70f7ddb92 100644 --- a/code_puppy/plugins/statusline/payload.py +++ b/code_puppy/plugins/statusline/payload.py @@ -48,6 +48,8 @@ def detect_git_branch(cwd: str) -> Optional[str]: cwd=cwd, capture_output=True, text=True, + encoding="utf-8", # explicit UTF-8: prevents cp1252 crash on Windows + errors="replace", # never raise UnicodeDecodeError on branch names timeout=0.5, ) if out.returncode == 0: diff --git a/code_puppy/plugins/statusline/runner.py b/code_puppy/plugins/statusline/runner.py index ffeeaadd3..f77407a56 100644 --- a/code_puppy/plugins/statusline/runner.py +++ b/code_puppy/plugins/statusline/runner.py @@ -33,6 +33,8 @@ def _run_command_blocking(command: str) -> str: input=payload, capture_output=True, text=True, + encoding="utf-8", # explicit UTF-8: prevents cp1252 crash on Windows with umlauts + errors="replace", # never raise UnicodeDecodeError — bad chars become '?' timeout=get_timeout_ms() / 1000.0, ) except subprocess.TimeoutExpired: diff --git a/code_puppy/plugins/statusline/statusline_command.py b/code_puppy/plugins/statusline/statusline_command.py index 9f6790f8c..d8b438eeb 100644 --- a/code_puppy/plugins/statusline/statusline_command.py +++ b/code_puppy/plugins/statusline/statusline_command.py @@ -16,6 +16,8 @@ from __future__ import annotations import stat +import subprocess as _sp +import sys from pathlib import Path from typing import Any, List, Optional, Tuple @@ -25,6 +27,26 @@ _COMMAND_NAME = "statusline" +# PowerShell equivalent for Windows — no bash, no jq required. +_STARTER_SCRIPT_PS1 = """\ +# Code Puppy status line — PowerShell version. +# Receives session JSON on stdin; prints one line to stdout. +# Edit freely. Requires PowerShell 5+ (ships with Windows 10+). +$input_text = $input | Out-String +try { + $data = $input_text | ConvertFrom-Json +} catch { + $data = [PSCustomObject]@{} +} +$puppy = if ($data.puppy_name) { $data.puppy_name } else { "code-puppy" } +$model = if ($data.model.display_name) { $data.model.display_name } else { "model" } +$dir = if ($data.workspace.current_dir) { $data.workspace.current_dir } elseif ($data.cwd) { $data.cwd } else { "" } +$branch = if ($data.workspace.git_branch) { " ($($data.workspace.git_branch))" } else { "" } +$pct = if ($null -ne $data.context_window.used_percentage) { $data.context_window.used_percentage } else { 0 } +$base = if ($dir) { Split-Path $dir -Leaf } else { "" } +Write-Output "🐶 $puppy [$model] $base$branch ${pct}%ctx" # stdout, captured by parent process +""" + _STARTER_SCRIPT = """\ #!/usr/bin/env bash # Code Puppy status line. Receives session JSON on stdin; prints one line. @@ -60,6 +82,8 @@ def statusline_command_help() -> List[Tuple[str, str]]: def _default_script_path() -> Path: + if sys.platform == "win32": + return Path.home() / ".code_puppy" / "statusline.ps1" return Path.home() / ".code_puppy" / "statusline.sh" @@ -77,18 +101,35 @@ def _status_text() -> str: def _do_init() -> None: path = _default_script_path() path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(_STARTER_SCRIPT, encoding="utf-8") - path.chmod(path.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) - config.set_command(str(path)) - config.set_enabled(True) - runner.reset_cache() - emit_success(f"Wrote starter status line script: {path}") - emit_info("Enabled. Edit that file to customize. Preview with /statusline show.") - if not _has_jq(): - emit_warning( - "Note: the starter script uses `jq` (not found on PATH). Install jq " - "or rewrite the script to parse JSON another way." + if sys.platform == "win32": + # Windows: write a PowerShell script; invoke via powershell.exe so no + # bash/jq dependency and no chmod needed. + path.write_text(_STARTER_SCRIPT_PS1, encoding="utf-8") + # list2cmdline handles cmd.exe quoting rules for paths with spaces + ps1_cmd = _sp.list2cmdline( + ["powershell", "-ExecutionPolicy", "Bypass", "-File", str(path)] + ) + config.set_command(ps1_cmd) + runner.reset_cache() + emit_success(f"Wrote starter status line script: {path}") + emit_info( + "Enabled. Edit that file to customize. Preview with /statusline show." ) + else: + path.write_text(_STARTER_SCRIPT, encoding="utf-8") + path.chmod(path.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) + config.set_command(str(path)) + runner.reset_cache() + emit_success(f"Wrote starter status line script: {path}") + emit_info( + "Enabled. Edit that file to customize. Preview with /statusline show." + ) + if not _has_jq(): + emit_warning( + "Note: the starter script uses `jq` (not found on PATH). Install jq " + "or rewrite the script to parse JSON another way." + ) + config.set_enabled(True) def _has_jq() -> bool: 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..1c4d68e93 100644 --- a/code_puppy/plugins/steer_queue/queue_menu.py +++ b/code_puppy/plugins/steer_queue/queue_menu.py @@ -1,114 +1,429 @@ -"""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 +from code_puppy.callbacks import on_prompt_toolkit_style 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.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:tui.body", + ) + + body = VSplit( + [ + Frame( + Window( + self._list_control, + wrap_lines=False, + width=Dimension(weight=2), + ), + title="Queued prompts", + style="class:tui.body", + ), + Window(width=1, char=" "), + self._detail_frame, + ], + padding=0, + ) + root = HSplit( + [ + Window( + FormattedTextControl(self._render_header), + height=2, + style="class:tui.header", + ), + body, + Window( + FormattedTextControl(self._render_notice), + height=1, + style="class:tui.warning", + ), + Window( + FormattedTextControl(self._render_footer), + height=2, + style="class:tui.help", + ), + ] + ) + + 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=on_prompt_toolkit_style(), + ) + self._sync_editor() + + def _editor_prefix(self): + return [("class:tui.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:tui.title", " PROMPT QUEUE\n"), + ("class:tui.header", f" {count} item{'s' if count != 1 else ''} waiting"), + ] + + def _render_list(self): + items = self.state.items + if not items: + return [ + ( + "class:tui.muted", + "\n Queue is empty.\n\n Press A to add a prompt.", + ) + ] + fragments = [] + for index, text in enumerate(items): + style = ( + "class:tui.selected" + if index == self.state.selected + else "class:tui.body" + ) + marker = "▶" if index == self.state.selected else " " + fragments.extend( + [ + (style, f" {marker} "), + (f"{style} class:tui.muted", f"{index + 1:>2} "), + (style, _preview(text)), + (style, "\n"), + ] + ) + return fragments + + def _render_notice(self): + return [("class:tui.warning", f" {self.state.notice}")] + + def _render_footer(self): + if self.state.editing: + return [ + ("class:tui.help", " "), + ("class:tui.help-key", "Ctrl+S"), + ("class:tui.help", " save "), + ("class:tui.help-key", "Esc"), + ("class:tui.help", " cancel Multi-line editing enabled"), + ] + return [ + ("class:tui.help", " "), + ("class:tui.help-key", "↑↓/JK"), + ("class:tui.help", " select "), + ("class:tui.help-key", "Enter/E"), + ("class:tui.help", " edit "), + ("class:tui.help-key", "A"), + ("class:tui.help", " add "), + ("class:tui.help-key", "D D"), + ("class:tui.help", " delete\n "), + ("class:tui.help-key", "[ ]"), + ("class:tui.help", " reorder "), + ("class:tui.help-key", "Q/Esc"), + ("class:tui.help", " 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 +436,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/code_puppy/plugins/subagent_panel/register_callbacks.py b/code_puppy/plugins/subagent_panel/register_callbacks.py index b5189cf9b..82f61c769 100644 --- a/code_puppy/plugins/subagent_panel/register_callbacks.py +++ b/code_puppy/plugins/subagent_panel/register_callbacks.py @@ -569,11 +569,18 @@ async def _on_post_tool_call(tool_name, tool_args, result, duration_ms, context= if not sid: return try: - err = getattr(result, "error", None) - if err: - state.mark_failed(sid) + # Ordinary tool invocations stay frozen until their foreground root + # flushes, preserving nested tree structure. A detached /fork has no + # foreground root, so keeping it would leave a completed row hitched to + # every subsequent prompt. Remove detached forks at their own boundary. + if isinstance(context, dict) and context.get("detached_fork"): + state.finish(sid) else: - state.mark_done(sid) + err = getattr(result, "error", None) + if err: + state.mark_failed(sid) + else: + state.mark_done(sid) except Exception: pass _push_panel(force=True) 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..f05d9aa74 100644 --- a/code_puppy/plugins/theme/picker.py +++ b/code_puppy/plugins/theme/picker.py @@ -34,6 +34,7 @@ content_styles_for, terminal_palette_for, ) +from code_puppy.callbacks import on_prompt_toolkit_style # A few representative banners to show in the preview pane (keeps it readable). PREVIEW_BANNERS = [ @@ -58,6 +59,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,27 +192,51 @@ 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"), + ("class:tui.header", "Pick a Theme"), + ("class:tui.muted", 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 "" + style = "class:tui.selected" if i == selected_index else "class:tui.body" line = f"{prefix}{i + 1}. {theme['icon']} {theme['label']}" lines.append((style, line)) lines.append(("", "\n")) - lines.append(("fg:ansigray", f" {theme['blurb']}")) + lines.append(("class:tui.muted", f" {theme['blurb']}")) lines.append(("", "\n\n")) lines.append(("", "\n")) - lines.append( - ( - "fg:ansicyan", - "Up/Down Navigate | Enter Apply | Esc / Ctrl-C Cancel", - ) + lines.extend( + [ + ("class:tui.help-key", "Up/Down"), + ("class:tui.help", " Navigate | "), + ("class:tui.help-key", "PgUp/PgDn"), + ("class:tui.help", " Page\n"), + ("class:tui.help-key", "Enter"), + ("class:tui.help", " Apply | "), + ("class:tui.help-key", "Esc / Ctrl-C"), + ("class:tui.help", " Cancel"), + ] ) return FormattedText(lines) @@ -253,6 +280,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] @@ -307,6 +346,7 @@ def _current_preview_style() -> str: full_screen=False, mouse_support=False, color_depth="DEPTH_24_BIT", + style=on_prompt_toolkit_style(), ) try: diff --git a/code_puppy/plugins/theme/prompt_toolkit_theme.py b/code_puppy/plugins/theme/prompt_toolkit_theme.py new file mode 100644 index 000000000..45111b2bb --- /dev/null +++ b/code_puppy/plugins/theme/prompt_toolkit_theme.py @@ -0,0 +1,114 @@ +"""Central prompt-toolkit adapter for the active Code Puppy theme. + +TUI applications pass their local style through the ``prompt_toolkit_style`` +hook. The theme plugin layers this semantic base underneath it, so local menu +rules retain precedence and legacy fragments such as ``bold`` inherit a +readable foreground. +""" + +from __future__ import annotations + +from typing import Any + +from prompt_toolkit.styles import BaseStyle, Style, merge_styles + +from . import osc_palette + + +def _active_palette() -> dict[str, Any] | None: + palette = osc_palette.get_saved_palette() + ansi = palette.get("ansi") if palette else None + if not palette or not isinstance(ansi, list) or len(ansi) < 16: + return None + return palette + + +def _relative_luminance(color: str) -> float: + """Return WCAG relative luminance for a ``#rrggbb`` color.""" + 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) + + +def _muted_foreground(ansi: list[str], foreground: str, background: str) -> str: + """Choose the least-prominent candidate that still meets WCAG AA.""" + candidates = (ansi[8], ansi[14], ansi[15], foreground) + readable = [ + color for color in candidates if _contrast_ratio(color, background) >= 4.5 + ] + return min( + readable, + key=lambda color: _contrast_ratio(color, background), + default=foreground, + ) + + +def get_style_rules() -> dict[str, str]: + """Return semantic prompt-toolkit rules for the active theme.""" + palette = _active_palette() + if palette is None: + return {} + + ansi = palette["ansi"] + foreground = palette.get("fg", ansi[7]) + background = palette.get("bg", ansi[0]) + muted = _muted_foreground(ansi, foreground, background) + return { + "": f"fg:{foreground} bg:{background}", + "tui": f"fg:{foreground} bg:{background}", + "tui.header": f"fg:{ansi[12]} bold", + "tui.title": f"fg:{ansi[14]} bold", + "tui.body": f"fg:{foreground}", + "tui.label": f"fg:{foreground} bold", + "tui.muted": f"fg:{muted}", + "tui.border": f"fg:{ansi[12]}", + "tui.selected": f"fg:{background} bg:{ansi[12]} bold noreverse", + "tui.help": f"fg:{muted}", + "tui.help-key": f"fg:{ansi[10]} bold", + "tui.success": f"fg:{ansi[10]} bold", + "tui.warning": f"fg:{ansi[11]} bold", + "tui.error": f"fg:{ansi[9]} bold", + "tui.input": f"fg:{foreground}", + "tui.input.focused": f"fg:{background} bg:{ansi[12]} bold noreverse", + # prompt_toolkit's defaults hard-code grey on white for completion + # menus. Root semantic rules cannot beat those more-specific selectors, + # so adapt the standard widget classes here as part of the shared theme. + "completion-menu": f"fg:{muted} bg:{background} noreverse", + "completion-menu.completion": f"fg:{muted} bg:{background} noreverse", + "completion-menu.completion.current": ( + f"fg:{ansi[12]} bg:{background} bold noreverse" + ), + "completion-menu.meta.completion": ( + f"fg:{muted} bg:{background} italic noreverse" + ), + "completion-menu.meta.completion.current": ( + f"fg:{ansi[14]} bg:{background} italic noreverse" + ), + "completion-menu.multi-column-meta": f"bg:{background}", + "scrollbar.background": f"fg:{muted} bg:{background}", + "scrollbar.button": f"fg:{muted} bg:{background}", + } + + +def get_style() -> Style: + """Build the active semantic style, or an empty style without a theme.""" + return Style.from_dict(get_style_rules()) + + +def merge_with_active_style(style: BaseStyle | None) -> BaseStyle: + """Layer a menu's specialized style over the shared theme base.""" + themed = get_style() + return merge_styles([themed, style]) if style is not None else themed + + +__all__ = ["get_style", "get_style_rules", "merge_with_active_style"] diff --git a/code_puppy/plugins/theme/register_callbacks.py b/code_puppy/plugins/theme/register_callbacks.py index 1c112a842..dc5a5660e 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 @@ -33,6 +35,7 @@ from . import osc_palette as osc from . import rich_themes as rt from .picker import interactive_theme_picker +from .prompt_toolkit_theme import merge_with_active_style from .themes import ( MENU_BY_NAME, apply, @@ -44,6 +47,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 +95,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 +117,138 @@ 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"}: + # Preserve monochrome phosphor while giving token classes enough + # separation to remain useful on both addition and deletion lines. + token_styles = { + Token: "#72a85b", + Comment: "#456b4f", + Error: "#7dff68", + Keyword: "#39e75f", + Name.Function: "#75ff87", + Number: "#8acb72", + Operator: "#63b96a", + String: "#9bdc79", + } + elif active_theme in {"solarized-light", "solarized"}: + token_styles = { + Token: "#657b83", + Comment: "#586e75", + Error: "#dc322f", + Keyword: "#859900", + Name.Function: "#268bd2", + Number: "#2aa198", + Operator: "#657b83", + String: "#2aa198", + } + 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) + if active_theme in {"green-screen", "green", "crt"}: + # Added lines lean brighter/yellower; removed lines become cooler and + # more muted. Small shifts retain the token palette's internal contrast. + highlighter.diff_line_tints = { + "added": (10, 14, -8), + "removed": (-18, -4, 10), + } + 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 +296,10 @@ 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("prompt_toolkit_style", merge_with_active_style) register_callback("custom_command_help", _custom_help) register_callback("custom_command", _handle_theme) diff --git a/code_puppy/plugins/theme/rich_themes.py b/code_puppy/plugins/theme/rich_themes.py index abcdb254d..233818700 100644 --- a/code_puppy/plugins/theme/rich_themes.py +++ b/code_puppy/plugins/theme/rich_themes.py @@ -187,6 +187,7 @@ def make_remap( bright_blue: Optional[str] = None, bright_magenta: Optional[str] = None, white: Optional[str] = None, + bright_white: Optional[str] = None, ) -> dict[str, str]: """Build a clean remap dict from optional kwargs (None = skip).""" mapping = { @@ -197,6 +198,7 @@ def make_remap( "bright_blue": bright_blue, "bright_magenta": bright_magenta, "white": white, + "bright_white": bright_white, } # Drop entries Rich can't parse so a typo in a theme never crashes the UI. return {k: v for k, v in mapping.items() if v and _safe_parse(v) is not None} diff --git a/code_puppy/plugins/theme/themes.py b/code_puppy/plugins/theme/themes.py index b67077b69..ac97a2c2d 100644 --- a/code_puppy/plugins/theme/themes.py +++ b/code_puppy/plugins/theme/themes.py @@ -344,6 +344,39 @@ 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", + bright_white="#00ff00", + ), + "terminal_palette": bp.GREEN_SCREEN, + }, "deep-black": { "icon": "🌑", "label": "Deep Black", @@ -464,7 +497,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 +508,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 +519,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 +530,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 +548,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/code_puppy/plugins/yolo_cli/__init__.py b/code_puppy/plugins/yolo_cli/__init__.py new file mode 100644 index 000000000..9a80e7c5a --- /dev/null +++ b/code_puppy/plugins/yolo_cli/__init__.py @@ -0,0 +1 @@ +"""Runtime-only CLI and session controls for YOLO mode.""" diff --git a/code_puppy/plugins/yolo_cli/register_callbacks.py b/code_puppy/plugins/yolo_cli/register_callbacks.py new file mode 100644 index 000000000..b227dd2d7 --- /dev/null +++ b/code_puppy/plugins/yolo_cli/register_callbacks.py @@ -0,0 +1,41 @@ +"""Expose a runtime-only YOLO override through the CLI.""" + +from __future__ import annotations + +import argparse +from typing import Any, Optional + +from code_puppy.callbacks import register_callback +from code_puppy.config import set_cli_yolo_override + + +def _parse_bool(value: str) -> bool: + """Parse the explicit boolean syntax used by ``--yolo``.""" + normalized = value.strip().lower() + if normalized == "true": + return True + if normalized == "false": + return False + raise argparse.ArgumentTypeError("expected 'true' or 'false'") + + +def _register_cli_args(parser: Any) -> None: + parser.add_argument( + "--yolo", + type=_parse_bool, + default=None, + metavar="{true,false}", + help=( + "Override YOLO mode for this run without changing puppy.cfg " + "(/set yolo_mode takes precedence)" + ), + ) + + +def _handle_cli_args(args: Any) -> Optional[dict]: + set_cli_yolo_override(getattr(args, "yolo", None)) + return None + + +register_callback("register_cli_args", _register_cli_args) +register_callback("handle_cli_args", _handle_cli_args) 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/code_puppy/session_storage.py b/code_puppy/session_storage.py index 62c6bfad1..d8b9201c6 100644 --- a/code_puppy/session_storage.py +++ b/code_puppy/session_storage.py @@ -242,10 +242,8 @@ def render_page() -> None: summary = ( f" and {remaining} more" if (remaining > 0 and not is_last_page) else "" ) - next_label = ( - "Return to first page" if is_last_page else f"Next page{summary}" - ) - emit_system_message(f" [6] {next_label}") + label = "Return to first page" if is_last_page else f"Next page{summary}" + emit_system_message(f" [6] {label}") emit_system_message(" [Enter] Skip loading autosave") chosen_name: str | None = None @@ -317,9 +315,9 @@ def render_page() -> None: # Set current autosave session id so subsequent autosaves overwrite this session try: - from code_puppy.config import pin_current_session_name + from code_puppy.config import set_current_autosave_from_session_name - pin_current_session_name(chosen_name) + set_current_autosave_from_session_name(chosen_name) except Exception: pass diff --git a/code_puppy/tools/agent_tools.py b/code_puppy/tools/agent_tools.py index f472b14b6..d583799e0 100644 --- a/code_puppy/tools/agent_tools.py +++ b/code_puppy/tools/agent_tools.py @@ -1,14 +1,12 @@ # agent_tools.py import hashlib import json -import pickle import re from datetime import datetime from pathlib import Path from typing import List from pydantic import BaseModel - from pydantic_ai import RunContext from pydantic_ai.messages import ModelMessage @@ -22,6 +20,7 @@ get_session_context, set_session_context, ) +from code_puppy.session_storage import load_session, save_session from code_puppy.tools.common import atomic_write_text, generate_group_id @@ -124,13 +123,36 @@ def _save_session_history( sessions_dir = _get_subagent_sessions_dir() - # Save pickle file with message history (atomic: write a temp file then - # replace, so a crash mid-write can't corrupt an existing session pickle) - pkl_path = sessions_dir / f"{session_id}.pkl" - tmp_pkl = pkl_path.with_suffix(".tmp") - with open(tmp_pkl, "wb") as f: - pickle.dump(message_history, f) - tmp_pkl.replace(pkl_path) + # Save JSON session history using the shared session storage helpers. + from code_puppy.agents._history import estimate_tokens_for_message + + save_session( + history=message_history, + session_name=session_id, + base_dir=sessions_dir, + timestamp=datetime.now().isoformat(), + token_estimator=estimate_tokens_for_message, + ) + + # Backward-compat artifact: some tests and legacy tooling still look for + # `.json` in the subagent sessions directory. + legacy_json_path = sessions_dir / f"{session_id}.json" + try: + from pydantic_ai.messages import ModelMessagesTypeAdapter + + legacy_payload = ModelMessagesTypeAdapter.dump_json(message_history).decode( + "utf-8" + ) + except Exception: + try: + legacy_payload = json.dumps( + [str(msg) for msg in message_history], + ensure_ascii=False, + indent=2, + ) + except Exception: + legacy_payload = "[]" + atomic_write_text(str(legacy_json_path), legacy_payload) # Save or update txt file with metadata txt_path = sessions_dir / f"{session_id}.txt" @@ -172,16 +194,12 @@ def _load_session_history(session_id: str) -> List[ModelMessage]: _validate_session_id(session_id) sessions_dir = _get_subagent_sessions_dir() - pkl_path = sessions_dir / f"{session_id}.pkl" - - if not pkl_path.exists(): - return [] try: - with open(pkl_path, "rb") as f: - return pickle.load(f) + return load_session(session_id, sessions_dir) + except FileNotFoundError: + return [] except Exception: - # If pickle is corrupted or incompatible, return empty history return [] 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/ask_user_question/renderers.py b/code_puppy/tools/ask_user_question/renderers.py index b5b882a9f..61f5d2253 100644 --- a/code_puppy/tools/ask_user_question/renderers.py +++ b/code_puppy/tools/ask_user_question/renderers.py @@ -117,10 +117,11 @@ def _render_question_panel_unsafe( # Question text if question.multi_select: console.print( - f"{pad}[bold]? {safe_question}[/bold] [dim](select multiple)[/dim]" + f"{pad}[{colors.question}]? {safe_question}[/{colors.question}] " + f"[{colors.question_hint}](select multiple)[/{colors.question_hint}]" ) else: - console.print(f"{pad}[bold]? {safe_question}[/bold]") + console.print(f"{pad}[{colors.question}]? {safe_question}[/{colors.question}]") console.print() # Render options diff --git a/code_puppy/tools/ask_user_question/theme.py b/code_puppy/tools/ask_user_question/theme.py index e84bd21e0..08e1c63bb 100644 --- a/code_puppy/tools/ask_user_question/theme.py +++ b/code_puppy/tools/ask_user_question/theme.py @@ -55,50 +55,36 @@ def _apply_config_overrides(default: _T, config_map: Mapping[str, str]) -> _T: class TUIColors(NamedTuple): """Color scheme for the ask_user_question TUI.""" - # Header and title colors - header_bold: str = "bold cyan" - header_dim: str = "fg:ansicyan dim" + # Compatibility field names; values are shared prompt-toolkit semantic roles. + header_bold: str = "class:tui.header" + header_dim: str = "class:tui.help" - # Cursor and selection colors - cursor_active: str = "fg:ansigreen bold" - cursor_inactive: str = "fg:ansiwhite" - selected: str = "fg:ansicyan" - selected_check: str = "fg:ansigreen" + cursor_active: str = "class:tui.success" + cursor_inactive: str = "class:tui.body" + selected: str = "class:tui.selected" + selected_check: str = "class:tui.success" - # Text colors - text_normal: str = "" - text_dim: str = "fg:ansiwhite dim" - text_warning: str = "fg:ansiyellow bold" + text_normal: str = "class:tui.body" + text_dim: str = "class:tui.muted" + text_warning: str = "class:tui.warning" - # Help text colors - help_key: str = "fg:ansigreen" - help_text: str = "fg:ansiwhite dim" + help_key: str = "class:tui.help-key" + help_text: str = "class:tui.help" - # Error colors - error: str = "fg:ansired" + error: str = "class:tui.error" # Create defaults after class definitions _DEFAULT_TUI = TUIColors() -# Mapping of configurable TUI color fields to config keys -_TUI_CONFIG_MAP: dict[str, str] = { - "header_bold": "tui_header_color", - "cursor_active": "tui_cursor_color", - "selected": "tui_selected_color", -} - def get_tui_colors() -> TUIColors: - """Get the current TUI color scheme. - - Loads colors from code-puppy's configuration system for custom theming. - Falls back to defaults for any missing config values. + """Return shared semantic roles for prompt-toolkit rendering. - Returns: - TUIColors instance with the current theme. + The active theme resolves these classes centrally. Legacy per-tool color + configuration must not override that shared palette. """ - return _apply_config_overrides(_DEFAULT_TUI, _TUI_CONFIG_MAP) + return _DEFAULT_TUI # Rich console color mappings for the right panel @@ -107,29 +93,29 @@ class RichColors(NamedTuple): # Header colors (Rich markup format) header: str = "bold cyan" - progress: str = "dim" + progress: str = "italic" # Question text question: str = "bold" - question_hint: str = "dim" + question_hint: str = "italic" # Option colors cursor: str = "green bold" selected: str = "cyan" normal: str = "" - description: str = "dim" + description: str = "italic" # Input field input_label: str = "bold yellow" input_text: str = "green" - input_hint: str = "dim" + input_hint: str = "italic" # Help overlay help_border: str = "bold cyan" help_title: str = "bold cyan" help_section: str = "bold" help_key: str = "green" - help_close: str = "dim" + help_close: str = "italic" # Timeout warning timeout_warning: str = "bold yellow" @@ -137,19 +123,21 @@ class RichColors(NamedTuple): _DEFAULT_RICH = RichColors() -# Mapping of configurable Rich color fields to config keys -_RICH_CONFIG_MAP: dict[str, str] = { - "header": "tui_rich_header_color", - "cursor": "tui_rich_cursor_color", -} - def get_rich_colors() -> RichColors: - """Get Rich console colors for the question panel. - - Falls back to defaults for any missing config values. - - Returns: - RichColors instance with current theme. - """ - return _apply_config_overrides(_DEFAULT_RICH, _RICH_CONFIG_MAP) + """Return Rich styles backed by the shared prompt-toolkit palette.""" + from code_puppy.plugins.theme.prompt_toolkit_theme import get_style_rules + + muted_rule = get_style_rules().get("tui.muted", "").removeprefix("fg:") + if not muted_rule: + return _DEFAULT_RICH + muted = muted_rule.split(maxsplit=1)[0] + + muted_style = f"{muted} italic" + return _DEFAULT_RICH._replace( + progress=muted_style, + question_hint=muted_style, + description=muted_style, + input_hint=muted_style, + help_close=muted_style, + ) diff --git a/code_puppy/tools/ask_user_question/tui_loop.py b/code_puppy/tools/ask_user_question/tui_loop.py index 85be89ddb..ab8720949 100644 --- a/code_puppy/tools/ask_user_question/tui_loop.py +++ b/code_puppy/tools/ask_user_question/tui_loop.py @@ -34,6 +34,7 @@ ) from .renderers import render_question_panel from .theme import get_rich_colors, get_tui_colors +from code_puppy.callbacks import on_prompt_toolkit_style if TYPE_CHECKING: from .models import QuestionAnswer @@ -357,19 +358,21 @@ def get_left_panel_text() -> FormattedText: [ ("", "\n"), ("", pad), - (tui_colors.header_dim, f"{ARROW_LEFT}{ARROW_RIGHT} Switch question"), + (tui_colors.help_key, f"{ARROW_LEFT}{ARROW_RIGHT}"), + (tui_colors.help_text, " Switch question"), ("", "\n"), ("", pad), - (tui_colors.header_dim, f"{ARROW_UP}{ARROW_DOWN} Navigate options"), + (tui_colors.help_key, f"{ARROW_UP}{ARROW_DOWN}"), + (tui_colors.help_text, " Navigate options"), ("", "\n"), ("", "\n"), ("", pad), (tui_colors.help_key, "Ctrl+S"), - (tui_colors.header_dim, " Submit"), + (tui_colors.help_text, " Submit"), ("", "\n"), ("", pad), (tui_colors.help_key, "Tab"), - (tui_colors.header_dim, " Peek behind"), + (tui_colors.help_text, " Peek behind"), ] ) @@ -418,6 +421,7 @@ def get_right_panel_text() -> ANSI: mouse_support=False, color_depth=ColorDepth.DEPTH_24_BIT, output=output, + style=on_prompt_toolkit_style(), ) # Timeout checker background task 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..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,16 @@ 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] - # Get the user-configured absolute timeout for shell commands - from code_puppy.config import get_command_timeout_seconds - - ABSOLUTE_TIMEOUT_SECONDS = get_command_timeout_seconds() + # Foreground duration limit. Reaching it detaches rather than killing; + # inactivity remains the guard for genuinely wedged commands. + foreground_limit_seconds = get_command_timeout_seconds() stdout_lines = [] stderr_lines = [] @@ -911,8 +1081,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 +1099,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 +1120,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 +1144,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: @@ -1023,8 +1192,9 @@ def detach_to_background(): ) 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(): 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(): ) 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) @@ -1096,6 +1276,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: @@ -1151,9 +1348,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 +1375,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, @@ -1204,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 @@ -1218,24 +1429,26 @@ 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). - # 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: @@ -1455,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/code_puppy/tools/common.py b/code_puppy/tools/common.py index dd7a5e044..92ee602d0 100644 --- a/code_puppy/tools/common.py +++ b/code_puppy/tools/common.py @@ -11,7 +11,7 @@ from typing import Callable, Optional, Tuple from prompt_toolkit import Application -from prompt_toolkit.formatted_text import HTML +from prompt_toolkit.formatted_text import FormattedText from prompt_toolkit.key_binding import KeyBindings from prompt_toolkit.layout import Layout, Window from prompt_toolkit.layout.controls import FormattedTextControl @@ -20,6 +20,7 @@ from rich.panel import Panel from rich.prompt import Prompt from rich.text import Text +from code_puppy.callbacks import on_prompt_toolkit_style # ============================================================================= # Approval queueing locks @@ -173,7 +174,6 @@ def resolve_path(file_path: str) -> str: # Syntax highlighting imports for "syntax" diff mode try: - from pygments import lex from pygments.lexers import TextLexer, get_lexer_by_name from pygments.token import Token @@ -741,41 +741,43 @@ def _get_token_color(token_type) -> str: return "#cccccc" # Default light-grey for unmatched tokens -def _highlight_code_line(code: str, bg_color: str | None, lexer) -> Text: - """Highlight a line of code with syntax highlighting and optional background color. - - Args: - code: The code string to highlight - bg_color: Background color in hex format, or None for no background - lexer: Pygments lexer instance to use - - Returns: - Rich Text object with styling applied - """ +def _highlight_code_line( + code: str, bg_color: str | None, lexer, line_type: str = "context" +) -> Text: + """Highlight code using TermFlow's theme-aware highlighter.""" if not PYGMENTS_AVAILABLE or lexer is None: - # Fallback: just return text with optional background - if bg_color: - return Text(code, style=f"on {bg_color}") - return Text(code) - - text = Text() - - for token_type, value in lex(code, lexer): - # Strip trailing newlines that Pygments adds - # Pygments lexer always adds a \n at the end of the last token - value = value.rstrip("\n") - - # Skip if the value is now empty (was only whitespace/newlines) - if not value: - continue - - fg_color = _get_token_color(token_type) - # Apply foreground color and optional background - if bg_color: - text.append(value, style=f"{fg_color} on {bg_color}") - else: - text.append(value, style=fg_color) + return Text(code, style=f"on {bg_color}" if bg_color else None) + + from code_puppy.callbacks import on_termflow_highlighter + from termflow.syntax import Highlighter + + highlighter = on_termflow_highlighter(Highlighter()) + language = (getattr(lexer, "aliases", None) or ["text"])[0] + text = Text.from_ansi(highlighter.highlight_line(code, language)) + + # Themes may provide subtle per-diff-line RGB shifts. Keeping this metadata + # on the themed highlighter avoids hard-coding theme knowledge in tools. + tint = getattr(highlighter, "diff_line_tints", {}).get(line_type) + if tint: + from rich.style import Style + from rich.text import Span + + for index, span in enumerate(text.spans): + style = span.style + color = getattr(style, "color", None) + triplet = color.get_truecolor() if color else None + if triplet: + shifted = tuple( + max(0, min(255, channel + delta)) + for channel, delta in zip(triplet, tint, strict=True) + ) + text.spans[index] = Span( + span.start, span.end, style + Style(color=f"rgb{shifted}") + ) + if bg_color: + # Applying only a background preserves each token's themed foreground. + text.stylize(f"on {bg_color}") return text @@ -833,12 +835,11 @@ def _format_diff_with_syntax_highlighting( addition_color: str | None = None, deletion_color: str | None = None, ) -> Text: - """Format diff with full syntax highlighting using Pygments. + """Format a diff with theme-aware syntax highlighting via TermFlow. This renders diffs with: - - Syntax highlighting for code tokens + - Theme-aware syntax highlighting for code tokens - Colored backgrounds for context/added/removed lines - - Monokai color scheme - Optional custom colors for additions/deletions Args: @@ -910,7 +911,9 @@ def _format_diff_with_syntax_highlighting( result.append(prefix) # Add syntax-highlighted code - highlighted = _highlight_code_line(code, bg_colors[line_type], lexer) + highlighted = _highlight_code_line( + code, bg_colors[line_type], lexer, line_type + ) result.append_text(highlighted) # Add newline after each line except the last @@ -920,18 +923,24 @@ def _format_diff_with_syntax_highlighting( return result -def format_diff_with_colors(diff_text: str) -> Text: +def format_diff_with_colors( + diff_text: str, + addition_color: str | None = None, + deletion_color: str | None = None, +) -> Text: """Format diff text with beautiful syntax highlighting. This is the canonical diff formatting function used across the codebase. - It applies user-configurable color coding with full syntax highlighting using Pygments. + It applies user-configurable colors and TermFlow's theme-aware syntax highlighting. - The function respects user preferences from config: - - get_diff_addition_color(): Color for added lines (markers and backgrounds) - - get_diff_deletion_color(): Color for deleted lines (markers and backgrounds) + Colors default to the effective theme-aware/user-configured preferences. + Callers rendering a preview may pass colors directly, avoiding config + mutations just to draw transient UI. Args: diff_text: Raw diff text to format + addition_color: Optional addition background override. + deletion_color: Optional deletion background override. Returns: Rich Text object with syntax highlighting @@ -944,8 +953,8 @@ def format_diff_with_colors(diff_text: str) -> Text: if not diff_text or not diff_text.strip(): return Text("-- no diff available --", style="dim") - addition_base_color = get_diff_addition_color() - deletion_base_color = get_diff_deletion_color() + addition_base_color = addition_color or get_diff_addition_color() + deletion_base_color = deletion_color or get_diff_deletion_color() # Always use beautiful syntax highlighting! if not PYGMENTS_AVAILABLE: @@ -961,6 +970,60 @@ def format_diff_with_colors(diff_text: str) -> Text: ) +def _format_selector( + message: str, + choices: list[str], + selected_index: int, + preview_callback: Optional[Callable[[int], str]] = None, +) -> FormattedText: + """Build shared selector content from semantic, literal-text fragments.""" + import textwrap + + fragments: list[tuple[str, str]] = [ + ("class:tui.header", message), + ("", "\n\n"), + ] + for index, choice in enumerate(choices): + style = "class:tui.selected" if index == selected_index else "class:tui.body" + marker = "\u276f " if index == selected_index else " " + fragments.extend([(style, marker + choice), ("", "\n")]) + fragments.append(("", "\n")) + + preview_text = preview_callback(selected_index) if preview_callback else "" + if preview_text: + box_width = 60 + fragments.extend( + [ + ( + "class:tui.border", + "┌─ Preview " + "─" * (box_width - 10) + "┐\n", + ) + ] + ) + wrapped_lines = textwrap.wrap(preview_text, width=box_width - 2) or [""] + for wrapped_line in wrapped_lines: + fragments.append( + ("class:tui.muted", f"│ {wrapped_line.ljust(box_width - 2)} │\n") + ) + fragments.extend( + [ + ("class:tui.border", "└" + "─" * box_width + "┘\n"), + ("", "\n"), + ] + ) + + fragments.extend( + [ + ("class:tui.help", "("), + ("class:tui.help-key", "↑↓ or Ctrl+P/N"), + ("class:tui.help", " to select, "), + ("class:tui.help-key", "Enter"), + ("class:tui.help", " to confirm)"), + ] + ) + return FormattedText(fragments) + + async def arrow_select_async( message: str, choices: list[str], @@ -980,61 +1043,14 @@ async def arrow_select_async( Raises: KeyboardInterrupt: If user cancels with Ctrl-C """ - import html - selected_index = [0] # Mutable container for selected index result = [None] # Mutable container for result - def get_formatted_text(): - """Generate the formatted text for display.""" - # Escape XML special characters to prevent parsing errors - safe_message = html.escape(message) - lines = [f"{safe_message}", ""] - for i, choice in enumerate(choices): - safe_choice = html.escape(choice) - if i == selected_index[0]: - lines.append(f"❯ {safe_choice}") - else: - lines.append(f" {safe_choice}") - lines.append("") - - # Add preview section if callback provided - if preview_callback is not None: - preview_text = preview_callback(selected_index[0]) - if preview_text: - import textwrap - - # Box width (excluding borders and padding) - box_width = 60 - border_top = ( - "┌─ Preview " - + "─" * (box_width - 10) - + "┐" - ) - border_bottom = "└" + "─" * box_width + "┘" - - lines.append(border_top) - - # Wrap text to fit within box width (minus padding) - wrapped_lines = textwrap.wrap(preview_text, width=box_width - 2) - - # If no wrapped lines (empty text), add empty line - if not wrapped_lines: - wrapped_lines = [""] - - for wrapped_line in wrapped_lines: - safe_preview = html.escape(wrapped_line) - # Pad line to box width for consistent appearance - padded_line = safe_preview.ljust(box_width - 2) - lines.append(f"│ {padded_line} │") - - lines.append(border_bottom) - lines.append("") - - lines.append( - "(Use ↑↓ or Ctrl+P/N to select, Enter to confirm)" + def get_formatted_text() -> FormattedText: + """Generate semantic formatted text for display.""" + return _format_selector( + message, choices, selected_index[0], preview_callback=preview_callback ) - return HTML("\n".join(lines)) # Key bindings kb = KeyBindings() @@ -1070,6 +1086,7 @@ def cancel(event): layout=layout, key_bindings=kb, full_screen=False, + style=on_prompt_toolkit_style(), ) # Flush output before prompt_toolkit takes control @@ -1108,19 +1125,9 @@ def arrow_select(message: str, choices: list[str]) -> str: selected_index = [0] # Mutable container for selected index result = [None] # Mutable container for result - def get_formatted_text(): - """Generate the formatted text for display.""" - lines = [f"{message}", ""] - for i, choice in enumerate(choices): - if i == selected_index[0]: - lines.append(f"❯ {choice}") - else: - lines.append(f" {choice}") - lines.append("") - lines.append( - "(Use ↑↓ or Ctrl+P/N to select, Enter to confirm)" - ) - return HTML("\n".join(lines)) + def get_formatted_text() -> FormattedText: + """Generate semantic formatted text for display.""" + return _format_selector(message, choices, selected_index[0]) # Key bindings kb = KeyBindings() @@ -1156,6 +1163,7 @@ def cancel(event): layout=layout, key_bindings=kb, full_screen=False, + style=on_prompt_toolkit_style(), ) # Flush output before prompt_toolkit takes control diff --git a/code_puppy/tools/display.py b/code_puppy/tools/display.py index 4aa761577..17874ca38 100644 --- a/code_puppy/tools/display.py +++ b/code_puppy/tools/display.py @@ -32,6 +32,28 @@ def erase_progress_line(console: Console) -> None: console.control(_ERASE_LINE_CONTROL) +def render_markdown(content: str, console: Console) -> None: + """Render complete Markdown through the configured Termflow pipeline.""" + from termflow import Parser as TermflowParser + from termflow import Renderer as TermflowRenderer + from termflow.render.style import RenderFeatures, RenderStyle + from termflow.syntax import Highlighter + + from code_puppy.callbacks import on_termflow_highlighter, on_termflow_style + + parser = TermflowParser() + renderer = TermflowRenderer( + output=console.file, + width=console.width, + style=on_termflow_style(RenderStyle.default()), + features=RenderFeatures(clipboard=False), + highlighter=on_termflow_highlighter(Highlighter()), + ) + for line in content.split("\n"): + renderer.render_all(parser.parse_line(line)) + renderer.render_all(parser.finalize()) + + def display_non_streamed_result( content: str, console: Optional[Console] = None, @@ -61,9 +83,6 @@ def display_non_streamed_result( return from rich.text import Text - from termflow import Parser as TermflowParser - from termflow import Renderer as TermflowRenderer - from termflow.render.style import RenderFeatures if console is None: console = Console() @@ -79,22 +98,7 @@ def display_non_streamed_result( ) ) - # Use termflow for markdown rendering - parser = TermflowParser() - renderer = TermflowRenderer( - output=console.file, - width=console.width, - features=RenderFeatures(clipboard=False), - ) - - # Process content line by line - for line in content.split("\n"): - events = parser.parse_line(line) - renderer.render_all(events) - - # Finalize to close any open markdown blocks - final_events = parser.finalize() - renderer.render_all(final_events) + render_markdown(content, console) -__all__ = ["display_non_streamed_result", "erase_progress_line"] +__all__ = ["display_non_streamed_result", "erase_progress_line", "render_markdown"] 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/code_puppy/tools/subagent_invocation.py b/code_puppy/tools/subagent_invocation.py index ec298ead5..a3d4787b9 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 @@ -294,14 +295,26 @@ async def _invoke_agent_impl( async with AsyncExitStack() as stack: for cm in run_ctxs: await stack.enter_async_context(cm) - task = asyncio.create_task( - temp_agent.run( + # Wrap the model stream in streaming_retry so a transient + # provider hiccup (gateway 5xx delivered as an in-band SSE + # error, a dropped SSE socket, an overloaded upstream) gets + # the same slow spaced-out retry the top-level agent loop gets + # instead of crashing the whole sub-agent invocation. This + # path was previously the ONLY unprotected model-stream + # call -- run_agent_task uses @streaming_retry, but a raw + # temp_agent.run() here surfaced the 5xx straight to the REPL. + from code_puppy.agents._runtime import streaming_retry + + @streaming_retry() + async def _run_subagent(): + return await temp_agent.run( prompt, message_history=message_history, usage_limits=UsageLimits(request_limit=get_message_limit()), event_stream_handler=stream_handler, ) - ) + + task = asyncio.create_task(_run_subagent()) _active_subagent_tasks.add(task) try: @@ -343,7 +356,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/pyproject.toml b/pyproject.toml index c020c3847..9e5100641 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,11 +4,15 @@ build-backend = "hatchling.build" [project] name = "code-puppy" -version = "0.0.621" +version = "0.0.637" 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/agents/test_base_agent_configuration.py b/tests/agents/test_base_agent_configuration.py index 7aae7f460..550b7f91f 100644 --- a/tests/agents/test_base_agent_configuration.py +++ b/tests/agents/test_base_agent_configuration.py @@ -41,3 +41,24 @@ def test_non_reasoning_sections_unchanged(self, agent): "Continue autonomously", ]: assert expected in prompt, f"Missing prompt section: {expected}" + + +class TestBaseAgentSessionModelCompatibility: + @pytest.fixture + def agent(self): + return CodePuppyAgent() + + def test_session_model_override_roundtrip(self, agent): + original = agent.get_model_name() + agent.set_session_model("gpt-5.5") + assert agent.get_session_model() == "gpt-5.5" + assert agent.get_model_name() == "gpt-5.5" + agent.reset_session_model() + assert agent.get_session_model() is None + assert agent.get_model_name() == original + + def test_compacted_hash_compatibility_helpers(self, agent): + assert agent.get_compacted_message_hashes() == set() + agent.add_compacted_message_hash(123) + agent.add_compacted_message_hash(123) + assert agent.get_compacted_message_hashes() == {123} diff --git a/tests/agents/test_builder_model_fallback.py b/tests/agents/test_builder_model_fallback.py new file mode 100644 index 000000000..3b3fc201c --- /dev/null +++ b/tests/agents/test_builder_model_fallback.py @@ -0,0 +1,94 @@ +"""Regression tests for ``load_model_with_fallback`` model resolution.""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest + +from code_puppy.agents._builder import load_model_with_fallback + + +def _config() -> dict[str, dict[str, str]]: + return { + "broken": {"type": "openai", "name": "broken-model"}, + "fallback-none": {"type": "openai", "name": "fallback-none-model"}, + "fallback-good": {"type": "openai", "name": "fallback-good-model"}, + } + + +def test_fallback_skips_none_model_and_uses_next_candidate() -> None: + """A fallback candidate returning ``None`` must not be treated as success.""" + models_config = _config() + expected_model = object() + call_order: list[str] = [] + + def _fake_get_model(model_name: str, _cfg: dict[str, object]) -> object | None: + call_order.append(model_name) + if model_name == "broken": + raise ValueError( + "Model 'broken' was found in configuration but could not be instantiated " + "(handler returned None)." + ) + if model_name == "fallback-none": + return None + if model_name == "fallback-good": + return expected_model + raise AssertionError(f"Unexpected model lookup: {model_name}") + + with ( + patch( + "code_puppy.agents._builder.ModelFactory.get_model", + side_effect=_fake_get_model, + ), + patch( + "code_puppy.agents._builder.get_global_model_name", + return_value="fallback-none", + ), + patch("code_puppy.agents._builder.emit_warning"), + patch("code_puppy.agents._builder.emit_info") as emit_info, + ): + model, resolved_name = load_model_with_fallback("broken", models_config, "grp") + + assert model is expected_model + assert resolved_name == "fallback-good" + assert call_order == ["broken", "fallback-none", "fallback-good"] + emit_info.assert_called_once_with( + "Using fallback model: fallback-good", + message_group="grp", + ) + + +def test_fallback_raises_when_all_candidates_are_invalid_or_none() -> None: + """When every candidate fails, we should raise a clear ValueError.""" + models_config = { + "broken": {"type": "openai", "name": "broken-model"}, + "fallback-none": {"type": "openai", "name": "fallback-none-model"}, + "fallback-error": {"type": "openai", "name": "fallback-error-model"}, + } + + def _fake_get_model(model_name: str, _cfg: dict[str, object]) -> object | None: + if model_name == "broken": + raise ValueError("Model 'broken' could not be instantiated") + if model_name == "fallback-none": + return None + if model_name == "fallback-error": + raise ValueError("Model 'fallback-error' not found in configuration") + raise AssertionError(f"Unexpected model lookup: {model_name}") + + with ( + patch( + "code_puppy.agents._builder.ModelFactory.get_model", + side_effect=_fake_get_model, + ), + patch( + "code_puppy.agents._builder.get_global_model_name", + return_value="fallback-none", + ), + patch("code_puppy.agents._builder.emit_warning"), + patch("code_puppy.agents._builder.emit_error") as emit_error, + ): + with pytest.raises(ValueError, match="No valid model could be loaded"): + load_model_with_fallback("broken", models_config, "grp") + + emit_error.assert_called_once() 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/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/agents/test_streaming_retry.py b/tests/agents/test_streaming_retry.py index d05f32459..f826c38f1 100644 --- a/tests/agents/test_streaming_retry.py +++ b/tests/agents/test_streaming_retry.py @@ -433,3 +433,241 @@ 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" + ) + ) + + +class TestInBandSSE5xx: + """Regression for the *real* production 502 that defeated every prior fix. + + A gateway (puppy-backend ``custom_anthropic`` proxy) delivers an upstream + 5xx as an **in-band SSE ``error`` event over a connection that was itself + HTTP 200**. The Anthropic SDK builds an ``APIStatusError`` from the *stream* + response, so: + + * ``status_code`` is **200** (not 502), + * the exception is a plain ``APIStatusError`` (not ``InternalServerError``), + * ``body.error.type`` is a generic ``"internal_error"``, + * the only faithful trace of the failure is a ``[HTTP 502]`` marker baked + into the message text. + + The classifier used to check ``status_code``, snippet matches, and a fixed + set of error ``type`` values -- all of which miss this shape -- so a + perfectly transient gateway 5xx crashed the REPL. This descends into the + ``[HTTP 5xx]`` marker instead. + """ + + def _make_anthropic_inband_5xx(self, http_code: int): + from anthropic import Anthropic + + client = Anthropic(api_key="sk-fake") + req = httpx.Request("POST", "https://api.anthropic.com/v1/messages") + # The stream itself returned 200 -- the failure is in-band. + resp = httpx.Response(200, request=req) + body = { + "error": { + "message": f"[HTTP {http_code}] ... {http_code}. " + "That's an error. The server encountered a temporary error ...", + "type": "internal_error", + }, + "type": "error", + } + return client._make_status_error(f"{body}", body=body, response=resp) + + @pytest.mark.parametrize("http_code", [500, 502, 503, 504, 529]) + def test_inband_5xx_marker_retries(self, http_code): + # status_code is 200; the real status lives only in the [HTTP NNN] marker. + err = self._make_anthropic_inband_5xx(http_code) + assert err.status_code == 200 # sanity: proves the trap this guards + assert should_retry_streaming_exception(err) + + def test_inband_5xx_wrapped_in_exception_group_retries(self): + # The exact production shape: anyio wraps it in an ExceptionGroup. + err = self._make_anthropic_inband_5xx(502) + group = ExceptionGroup("unhandled errors in a TaskGroup", [err]) + assert should_retry_streaming_exception(group) + + def test_inband_internal_error_type_retries(self): + # Even without a marker match, type "internal_error" is a server-side + # Anthropic failure and must retry (belt-and-braces). + from anthropic import Anthropic + + client = Anthropic(api_key="sk-fake") + req = httpx.Request("POST", "https://api.anthropic.com/v1/messages") + resp = httpx.Response(200, request=req) + body = { + "error": {"message": "opaque", "type": "internal_error"}, + "type": "error", + } + err = client._make_status_error(f"{body}", body=body, response=resp) + assert should_retry_streaming_exception(err) + + @pytest.mark.parametrize("http_code", [400, 401, 403, 404, 422]) + def test_inband_4xx_marker_does_not_retry(self, http_code): + # A [HTTP 4xx] marker is a genuine client error -- must NOT retry. + # Uses a clean client-error body (no server-error wording) so the + # assertion stays robust even if the retryable-snippet list grows. + from anthropic import Anthropic + + client = Anthropic(api_key="sk-fake") + req = httpx.Request("POST", "https://api.anthropic.com/v1/messages") + resp = httpx.Response(200, request=req) + body = { + "error": { + "message": f"[HTTP {http_code}] invalid request payload", + "type": "invalid_request_error", + }, + "type": "error", + } + err = client._make_status_error(f"{body}", body=body, response=resp) + assert err.status_code == 200 # sanity: same 200-stream trap, 4xx marker + assert not should_retry_streaming_exception(err) + + +class TestRetryBudgetSpacing: + """The retry budget must outlast a gateway 5xx outage, not just a blip. + + A Google/gateway 502 ("try again in 30 seconds") outlasts a tight 1-2-4s + window, so a fast burst of retries just exhausts the budget before the + upstream recovers. The default is 3 attempts spaced (5, 30)s: a quick first + retry for instantaneous SSE blips, then a long gap that rides out the + advertised ~30s outage before giving up. + """ + + def test_default_budget_is_spaced_out(self): + import inspect + + from code_puppy.agents._runtime import streaming_retry + + params = inspect.signature(streaming_retry).parameters + assert params["max_attempts"].default == 3 + assert tuple(params["delays"].default) == (5, 30) + + def test_runner_sleeps_the_spaced_delays_then_gives_up(self): + # 3 attempts against a persistent transient error -> two sleeps (5, 30) + # then re-raise. We patch asyncio.sleep so the test stays instant while + # asserting the *real* runner uses the spaced-out delays. + from unittest.mock import AsyncMock, patch + + from code_puppy.agents._runtime import streaming_retry + + attempts = {"n": 0} + + @streaming_retry() + async def _always_transient(): + attempts["n"] += 1 + raise httpx.ConnectError("dropped socket") # always retryable + + with patch( + "code_puppy.agents._runtime.asyncio.sleep", new_callable=AsyncMock + ) as mock_sleep: + with pytest.raises(httpx.ConnectError): + asyncio.run(_always_transient()) + + assert attempts["n"] == 3 # all attempts consumed + assert [c.args[0] for c in mock_sleep.await_args_list] == [5, 30] 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/api/test_permissions_session_binding.py b/tests/api/test_permissions_session_binding.py new file mode 100644 index 000000000..b1f1204f6 --- /dev/null +++ b/tests/api/test_permissions_session_binding.py @@ -0,0 +1,115 @@ +from __future__ import annotations + +import asyncio + +import pytest + +from code_puppy.api.permissions import ( + PendingPermissionRequest, + handle_permission_response, + permission_futures, +) + + +@pytest.mark.asyncio +async def test_handle_permission_response_rejects_session_mismatch(): + future = asyncio.get_running_loop().create_future() + permission_futures["req-1"] = PendingPermissionRequest( + future=future, + session_id="session-a", + ) + try: + assert ( + handle_permission_response("req-1", True, session_id="session-b") is False + ) + assert future.done() is False + finally: + permission_futures.clear() + + +@pytest.mark.asyncio +async def test_handle_permission_response_allows_matching_session(): + future = asyncio.get_running_loop().create_future() + permission_futures["req-2"] = PendingPermissionRequest( + future=future, + session_id="session-a", + ) + try: + assert handle_permission_response("req-2", True, session_id="session-a") is True + assert future.done() is True + assert future.result() is True + finally: + permission_futures.clear() + + +def test_handle_permission_response_unknown_request_returns_false(): + permission_futures.clear() + assert handle_permission_response("missing", True, session_id="session-a") is False + + +@pytest.mark.asyncio +async def test_request_permission_auto_approves_when_yolo_enabled(monkeypatch): + from code_puppy.api.permissions import request_permission + import code_puppy.config as config + + class FakeWebSocket: + def __init__(self): + self.sent = [] + + async def send_json(self, payload): + self.sent.append(payload) + + websocket = FakeWebSocket() + monkeypatch.setattr(config, "get_yolo_mode", lambda: True) + + approved = await request_permission( + websocket=websocket, + session_id="session-a", + request_type="shell_command", + title="Execute Shell Command", + description="Run: pwd", + details={"command": "pwd"}, + timeout=1, + ) + + assert approved is True + assert websocket.sent == [] + + +@pytest.mark.asyncio +async def test_request_permission_registers_future_before_send(monkeypatch): + from code_puppy.api.permissions import request_permission + + class RaceWebSocket: + def __init__(self): + self.sent = [] + + async def send_json(self, payload): + self.sent.append(payload) + assert ( + handle_permission_response( + payload["request_id"], + True, + session_id=payload["session_id"], + ) + is True + ) + + websocket = RaceWebSocket() + permission_futures.clear() + + monkeypatch.setattr("code_puppy.config.get_yolo_mode", lambda: False) + + approved = await request_permission( + websocket=websocket, + session_id="session-a", + request_type="shell_command", + title="Execute Shell Command", + description="Run: pwd", + details={"command": "pwd"}, + timeout=1, + ) + + assert approved is True + assert len(websocket.sent) == 1 + assert permission_futures == {} diff --git a/tests/api/test_ws_sessions_router.py b/tests/api/test_ws_sessions_router.py new file mode 100644 index 000000000..d1d54aaf3 --- /dev/null +++ b/tests/api/test_ws_sessions_router.py @@ -0,0 +1,88 @@ +import importlib.util +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest +from starlette.routing import Router + + +_original_router_init = Router.__init__ + + +def _compat_router_init(self, *args, **kwargs): + kwargs.pop("on_startup", None) + kwargs.pop("on_shutdown", None) + kwargs.pop("lifespan", None) + return _original_router_init(self, *args, **kwargs) + + +Router.__init__ = _compat_router_init +try: + module_path = ( + Path(__file__).resolve().parents[2] + / "code_puppy" + / "api" + / "routers" + / "ws_sessions.py" + ) + spec = importlib.util.spec_from_file_location( + "test_ws_sessions_module", module_path + ) + assert spec is not None and spec.loader is not None + ws_sessions = importlib.util.module_from_spec(spec) + spec.loader.exec_module(ws_sessions) +finally: + Router.__init__ = _original_router_init + + +@pytest.mark.asyncio +async def test_get_ws_session_messages_keeps_plain_legacy_rows(monkeypatch): + monkeypatch.setattr( + ws_sessions, + "_validate_session_name", + lambda session_name, ws_dir: session_name, + ) + monkeypatch.setattr( + ws_sessions, + "get_active_messages", + AsyncMock( + return_value=[ + { + "role": "assistant", + "content": "plain legacy row", + "type": "message", + "agent_name": "code-puppy", + "model_name": "gpt-5", + "timestamp": "2026-01-01T00:00:00Z", + "thinking": None, + "clean_content": "plain legacy row", + "seq": 7, + "pydantic_json": None, + } + ] + ), + ) + monkeypatch.setattr( + ws_sessions, + "get_session_metadata", + AsyncMock(return_value=None), + ) + + result = await ws_sessions.get_ws_session_messages( + "session-1", + include_tool_calls=False, + ) + + assert result == [ + { + "role": "assistant", + "content": "plain legacy row", + "type": "message", + "agent_name": "code-puppy", + "model_name": "gpt-5", + "timestamp": "2026-01-01T00:00:00Z", + "thinking": None, + "clean_content": "plain legacy row", + "seq": 7, + } + ] diff --git a/tests/command_line/mcp/test_custom_server_form.py b/tests/command_line/mcp/test_custom_server_form.py index afa7a91d2..59d6319c4 100644 --- a/tests/command_line/mcp/test_custom_server_form.py +++ b/tests/command_line/mcp/test_custom_server_form.py @@ -303,6 +303,25 @@ def test_render_form_status_message_success(self): texts = "".join(t[1] for t in result) assert "Saved!" in texts + def test_render_form_uses_semantic_state_roles(self): + form = self._make_form() + form.focused_field = 1 + form.server_name = "bad name!" + form.validation_error = "Invalid config" + form.status_message = "Save failed" + form.status_is_error = True + + styles = {style for style, _text in form._render_form() if style} + + assert { + "class:tui.input.focused", + "class:tui.selected", + "class:tui.help-key", + "class:tui.warning", + "class:tui.error", + } <= styles + assert not any("fg:" in style or "ansi" in style for style in styles) + # --------------------------------------------------------------------------- # _render_preview @@ -345,6 +364,14 @@ def test_render_preview_contains_tips(self): assert "Tips" in texts assert "$ENV_VAR" in texts + def test_render_preview_uses_semantic_roles(self): + styles = { + style for style, _text in self._make_form()._render_preview() if style + } + + assert {"class:tui.header", "class:tui.label", "class:tui.muted"} <= styles + assert not any("fg:" in style or "ansi" in style for style in styles) + # --------------------------------------------------------------------------- # _install_server @@ -541,6 +568,8 @@ def side_run(**kwargs): result = form.run() assert result is False + assert "class:tui.input" in form.name_area.window.style + assert "class:tui.input" in form.json_area.window.style mock_set.assert_any_call(True) mock_set.assert_any_call(False) diff --git a/tests/command_line/mcp/test_install_menu.py b/tests/command_line/mcp/test_install_menu.py index c2dbf80af..91f5cab80 100644 --- a/tests/command_line/mcp/test_install_menu.py +++ b/tests/command_line/mcp/test_install_menu.py @@ -142,16 +142,16 @@ def test_get_current_server_out_of_range(self): class TestCategoryIcon: def test_custom_server_icon(self): menu = make_menu() - assert menu._get_category_icon("➕ Custom Server") == "➕" + assert menu._get_category_icon(menu.categories[0]) == "[+]" def test_known_category_icon(self): menu = make_menu() - assert menu._get_category_icon("Code") == "💻" - assert menu._get_category_icon("Storage") == "💾" + assert menu._get_category_icon("Code") == "[C]" + assert menu._get_category_icon("Storage") == "[S]" def test_unknown_category_icon(self): menu = make_menu() - assert menu._get_category_icon("Unknown") == "📁" + assert menu._get_category_icon("Unknown") == "[ ]" class TestIsCustomServerSelected: @@ -207,6 +207,19 @@ def test_renders_with_no_catalog(self): text = "".join(str(t[1]) for t in lines) assert "(0)" in text # zero server count when catalog is None + def test_uses_semantic_tui_roles(self): + menu = make_menu() + menu.selected_category_idx = 1 + + styles = {style for style, _text in menu._render_category_list() if style} + + assert { + "class:tui.header", + "class:tui.selected", + "class:tui.help-key", + } <= styles + assert not any("fg:" in style or "ansi" in style for style in styles) + class TestRenderServerList: def test_no_category_selected(self): @@ -252,6 +265,18 @@ def test_server_pagination(self): text = "".join(str(t[1]) for t in lines) assert "Page" in text + def test_server_details_use_semantic_status_roles(self, monkeypatch): + menu = make_menu() + menu.view_mode = "servers" + menu.current_servers = [FakeServer()] + menu.selected_server_idx = 0 + monkeypatch.delenv("API_KEY", raising=False) + + styles = {style for style, _text in menu._render_details() if style} + + assert {"class:tui.label", "class:tui.muted", "class:tui.warning"} <= styles + assert not any("fg:" in style or "ansi" in style for style in styles) + class TestRenderDetails: def test_no_category(self): diff --git a/tests/command_line/mcp/test_mcp_binding_menu.py b/tests/command_line/mcp/test_mcp_binding_menu.py new file mode 100644 index 000000000..5868e3348 --- /dev/null +++ b/tests/command_line/mcp/test_mcp_binding_menu.py @@ -0,0 +1,40 @@ +"""Semantic rendering tests for the MCP binding menus.""" + +from unittest.mock import patch + +from code_puppy.command_line.mcp_binding_menu import _render_menu, _render_preview + +MODULE = "code_puppy.command_line.mcp_binding_menu" + + +def test_binding_menu_uses_semantic_selection_and_status_roles(): + bindings = {"alpha": {"auto_start": True}} + + with patch(f"{MODULE}.get_bound_servers", return_value=bindings): + lines = _render_menu("code-puppy", [("alpha", "stdio", "running")], 0) + + styles = {style for style, _text in lines if style} + assert { + "class:tui.selected", + "class:tui.success", + "class:tui.warning", + "class:tui.help-key", + } <= styles + assert not any("fg:" in style or "ansi" in style for style in styles) + + +def test_binding_preview_uses_semantic_detail_roles(): + bindings = {"alpha": {"auto_start": True}} + + with patch(f"{MODULE}.get_bound_servers", return_value=bindings): + lines = _render_preview("code-puppy", [("alpha", "stdio", "running")], 0) + + styles = {style for style, _text in lines if style} + assert { + "class:tui.title", + "class:tui.label", + "class:tui.body", + "class:tui.success", + "class:tui.warning", + } <= styles + assert not any("fg:" in style or "ansi" in style for style in styles) diff --git a/tests/command_line/test_add_model_menu_coverage.py b/tests/command_line/test_add_model_menu_coverage.py index b577237bd..72691c671 100644 --- a/tests/command_line/test_add_model_menu_coverage.py +++ b/tests/command_line/test_add_model_menu_coverage.py @@ -162,12 +162,16 @@ def test_render_with_providers(self): assert "Bedrock" in text assert "Page" in text - def test_render_unsupported_provider_dimmed(self): - p = _make_provider(pid="amazon-bedrock", name="Bedrock") - menu = _make_menu_with_providers([p]) + def test_render_unsupported_provider_muted(self): + supported = _make_provider(pid="openai", name="OpenAI") + unsupported = _make_provider(pid="amazon-bedrock", name="Bedrock") + menu = _make_menu_with_providers([supported, unsupported]) lines = menu._render_provider_list() - styles = [s for s, _ in lines] - assert any("dim" in s for s in styles) + styles_by_text = {text: style for style, text in lines} + bedrock_style = next( + style for text, style in styles_by_text.items() if "Bedrock" in text + ) + assert bedrock_style == "class:tui.muted" # --------------- Render model list --------------- diff --git a/tests/command_line/test_agent_menu.py b/tests/command_line/test_agent_menu.py index d3af946fd..68db4cc86 100644 --- a/tests/command_line/test_agent_menu.py +++ b/tests/command_line/test_agent_menu.py @@ -606,57 +606,49 @@ def test_preview_panel_with_no_description_default(self): class TestMenuPanelStyling: """Test styling aspects of the menu panel.""" - def test_styling_includes_green_for_selection(self): - """Test that selection styling uses green color.""" + def test_selection_uses_semantic_style(self): + """Test that selection styling uses the shared semantic role.""" entries = [("agent1", "Agent One", "Description")] result = _render_menu_panel( entries, page=0, selected_idx=0, current_agent_name="" ) - # Check that green styling is applied somewhere styles = [style for style, _ in result] - has_green = any("green" in str(style).lower() for style in styles) - assert has_green, "Selection should use green styling" + assert "class:tui.selected" in styles - def test_styling_includes_cyan_for_current(self): - """Test that current agent marker uses cyan color.""" + def test_current_marker_uses_semantic_success_style(self): + """Test that the current marker uses the shared success role.""" entries = [("agent1", "Agent One", "Description")] result = _render_menu_panel( entries, page=0, selected_idx=0, current_agent_name="agent1" ) - # Check that cyan styling is used for current marker styles = [style for style, _ in result] - has_cyan = any("cyan" in str(style).lower() for style in styles) - assert has_cyan, "Current marker should use cyan styling" + assert "class:tui.success" in styles class TestPreviewPanelStyling: """Test styling aspects of the preview panel.""" def test_styling_for_active_status(self): - """Test that active status uses appropriate styling.""" + """Test that active status uses the semantic success role.""" entry = ("agent1", "Agent One", "Description") result = _render_preview_panel(entry, current_agent_name="agent1") - # Check for green styling on active status styles = [style for style, _ in result] - has_green = any("green" in str(style).lower() for style in styles) - assert has_green, "Active status should use green styling" + assert "class:tui.success" in styles def test_styling_for_inactive_status(self): - """Test that inactive status uses dimmed styling.""" + """Test that inactive status uses the semantic muted role.""" entry = ("agent1", "Agent One", "Description") result = _render_preview_panel(entry, current_agent_name="other_agent") - # Check for dimmed/bright black styling styles = [style for style, _ in result] - has_dim = any("bright" in str(style).lower() for style in styles) - assert has_dim, "Inactive status should use dimmed styling" + assert "class:tui.muted" in styles class TestGetPinnedModelWithJSONAgents: 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/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_diff_menu.py b/tests/command_line/test_diff_menu.py index 684bb085d..dd2b4f367 100644 --- a/tests/command_line/test_diff_menu.py +++ b/tests/command_line/test_diff_menu.py @@ -252,27 +252,17 @@ def test_preview_generation_with_mocked_config( result = _get_preview_text_for_prompt_toolkit(config) - # Should call config functions + # Construction reads config, but rendering passes transient preview + # colors directly instead of mutating persistent settings. mock_get_add.assert_called() mock_get_del.assert_called() - # Should set config color, then restore original (2 calls total) - from unittest.mock import call - - expected_calls = [ - call("#00aa00"), - call("#00ff00"), - ] # set config, restore original - assert mock_set_add.call_count == 2 - mock_set_add.assert_has_calls(expected_calls, any_order=False) - - expected_del_calls = [ - call("#aa0000"), - call("#ff0000"), - ] # set config, restore original - assert mock_set_del.call_count == 2 - mock_set_del.assert_has_calls(expected_del_calls, any_order=False) - - mock_format.assert_called() + mock_set_add.assert_not_called() + mock_set_del.assert_not_called() + mock_format.assert_called_once() + assert mock_format.call_args.kwargs == { + "addition_color": "#00aa00", + "deletion_color": "#aa0000", + } # Should return ANSI object assert hasattr(result, "__class__") @@ -415,6 +405,7 @@ def capture_app( full_screen=False, mouse_support=False, color_depth=None, + style=None, ): # Get the formatted text from the layout return mock_instance @@ -864,14 +855,14 @@ def test_full_preview_pipeline( result = _get_preview_text_for_prompt_toolkit(config) assert result is not None - # Verify mock calls - function sets config to current values then restores original + # Preview colors are direct arguments; config stays untouched. mock_format.assert_called() - # Should set to current config colors first - mock_set_add.assert_any_call(config.current_add_color) - mock_set_del.assert_any_call(config.current_del_color) - # Then restore original values - mock_set_add.assert_any_call(original_add) - mock_set_del.assert_any_call(original_del) + assert mock_format.call_args.kwargs == { + "addition_color": config.current_add_color, + "deletion_color": config.current_del_color, + } + mock_set_add.assert_not_called() + mock_set_del.assert_not_called() @pytest.mark.asyncio async def test_complete_interactive_workflow(self): diff --git a/tests/command_line/test_diff_menu_coverage.py b/tests/command_line/test_diff_menu_coverage.py index 93283903d..8f7583247 100644 --- a/tests/command_line/test_diff_menu_coverage.py +++ b/tests/command_line/test_diff_menu_coverage.py @@ -39,6 +39,7 @@ def capture_app( full_screen=False, mouse_support=False, color_depth=None, + style=None, ): captured_kb[0] = key_bindings mock_app = MagicMock() 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/command_line/test_onboarding_slides.py b/tests/command_line/test_onboarding_slides.py index 56fa2ba2c..1b43c3949 100644 --- a/tests/command_line/test_onboarding_slides.py +++ b/tests/command_line/test_onboarding_slides.py @@ -3,6 +3,10 @@ MODULE = "code_puppy.command_line.onboarding_slides" +def _plain(content): + return "".join(text for _, text in content) + + class TestModelOptions: def test_model_options_is_list(self): from code_puppy.command_line.onboarding_slides import MODEL_OPTIONS @@ -22,8 +26,9 @@ class TestGetNavFooter: def test_returns_string(self): from code_puppy.command_line.onboarding_slides import get_nav_footer - result = get_nav_footer() - assert isinstance(result, str) + content = get_nav_footer() + result = _plain(content) + assert isinstance(content, list) assert "Next" in result assert "Back" in result assert "ESC" in result @@ -33,8 +38,9 @@ class TestGetGradientBanner: def test_with_pyfiglet(self): from code_puppy.command_line.onboarding_slides import get_gradient_banner - result = get_gradient_banner() - assert isinstance(result, str) + content = get_gradient_banner() + result = _plain(content) + assert isinstance(content, list) # Should contain some content assert len(result) > 0 @@ -43,7 +49,8 @@ def test_without_pyfiglet(self): import code_puppy.command_line.onboarding_slides as mod # pyfiglet is available in this env, so normal path works - result = mod.get_gradient_banner() + content = mod.get_gradient_banner() + result = _plain(content) assert len(result) > 0 @@ -51,8 +58,9 @@ class TestSlideWelcome: def test_returns_string(self): from code_puppy.command_line.onboarding_slides import slide_welcome - result = slide_welcome() - assert isinstance(result, str) + content = slide_welcome() + result = _plain(content) + assert isinstance(content, list) assert "Welcome" in result assert "setup" in result.lower() or "quick" in result.lower() @@ -68,7 +76,8 @@ def test_with_options(self): ("openrouter", "OpenRouter"), ("skip", "Skip"), ] - result = slide_models(0, options) + content = slide_models(0, options) + result = _plain(content) assert "ChatGPT" in result assert "▶" in result # selected indicator @@ -76,41 +85,46 @@ def test_claude_selected(self): from code_puppy.command_line.onboarding_slides import slide_models options = [("chatgpt", "ChatGPT"), ("claude", "Claude")] - result = slide_models(1, options) + content = slide_models(1, options) + result = _plain(content) assert "Claude" in result def test_api_keys_context(self): from code_puppy.command_line.onboarding_slides import slide_models options = [("api_keys", "API Keys")] - result = slide_models(0, options) + content = slide_models(0, options) + result = _plain(content) assert "API Key" in result def test_openrouter_context(self): from code_puppy.command_line.onboarding_slides import slide_models options = [("openrouter", "OpenRouter")] - result = slide_models(0, options) + content = slide_models(0, options) + result = _plain(content) assert "OpenRouter" in result def test_skip_context(self): from code_puppy.command_line.onboarding_slides import slide_models options = [("skip", "Skip")] - result = slide_models(0, options) + content = slide_models(0, options) + result = _plain(content) assert "later" in result.lower() or "No worries" in result def test_empty_options(self): from code_puppy.command_line.onboarding_slides import slide_models - result = slide_models(0, []) - assert isinstance(result, str) + content = slide_models(0, []) + assert isinstance(content, list) def test_chatgpt_context(self): from code_puppy.command_line.onboarding_slides import slide_models options = [("chatgpt", "ChatGPT Plus")] - result = slide_models(0, options) + content = slide_models(0, options) + result = _plain(content) assert "ChatGPT" in result or "OAuth" in result @@ -118,8 +132,9 @@ class TestSlideMcp: def test_returns_string(self): from code_puppy.command_line.onboarding_slides import slide_mcp - result = slide_mcp() - assert isinstance(result, str) + content = slide_mcp() + result = _plain(content) + assert isinstance(content, list) assert "MCP" in result assert "/mcp" in result @@ -128,8 +143,9 @@ class TestSlideUseCases: def test_returns_string(self): from code_puppy.command_line.onboarding_slides import slide_use_cases - result = slide_use_cases() - assert isinstance(result, str) + content = slide_use_cases() + result = _plain(content) + assert isinstance(content, list) assert "Planning" in result assert "Code Puppy" in result @@ -138,20 +154,23 @@ class TestSlideDone: def test_without_oauth(self): from code_puppy.command_line.onboarding_slides import slide_done - result = slide_done(None) - assert isinstance(result, str) + content = slide_done(None) + result = _plain(content) + assert isinstance(content, list) assert "Ready" in result assert "/tutorial" in result def test_with_oauth_chatgpt(self): from code_puppy.command_line.onboarding_slides import slide_done - result = slide_done("chatgpt") + content = slide_done("chatgpt") + result = _plain(content) assert "Chatgpt" in result or "chatgpt" in result.lower() assert "OAuth" in result def test_with_oauth_claude(self): from code_puppy.command_line.onboarding_slides import slide_done - result = slide_done("claude") + content = slide_done("claude") + result = _plain(content) assert "Claude" in result or "claude" in result.lower() diff --git a/tests/command_line/test_onboarding_wizard.py b/tests/command_line/test_onboarding_wizard.py index 39d594e9d..f1ee664bd 100644 --- a/tests/command_line/test_onboarding_wizard.py +++ b/tests/command_line/test_onboarding_wizard.py @@ -128,31 +128,31 @@ def test_get_slide_content_slide_0(self): w = self._make_wizard() w.current_slide = 0 content = w.get_slide_content() - assert isinstance(content, str) + assert isinstance(content, list) def test_get_slide_content_slide_1(self): w = self._make_wizard() w.current_slide = 1 content = w.get_slide_content() - assert isinstance(content, str) + assert isinstance(content, list) def test_get_slide_content_slide_2(self): w = self._make_wizard() w.current_slide = 2 content = w.get_slide_content() - assert isinstance(content, str) + assert isinstance(content, list) def test_get_slide_content_slide_3(self): w = self._make_wizard() w.current_slide = 3 content = w.get_slide_content() - assert isinstance(content, str) + assert isinstance(content, list) def test_get_slide_content_slide_4(self): w = self._make_wizard() w.current_slide = 4 content = w.get_slide_content() - assert isinstance(content, str) + assert isinstance(content, list) def test_get_options_for_slide_1(self): w = self._make_wizard() @@ -268,8 +268,8 @@ def test_prev_option_no_options(self): class TestGetSlidePanelContent: - def test_returns_ansi(self): - from prompt_toolkit.formatted_text import ANSI + def test_returns_semantic_formatted_text(self): + from prompt_toolkit.formatted_text import FormattedText from code_puppy.command_line.onboarding_wizard import ( OnboardingWizard, @@ -278,7 +278,9 @@ def test_returns_ansi(self): w = OnboardingWizard() result = _get_slide_panel_content(w) - assert isinstance(result, ANSI) + assert isinstance(result, FormattedText) + styles = {style for style, _ in result} + assert {"class:tui.muted", "class:tui.header"} <= styles # --------------------------------------------------------------------------- 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_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_set_menu.py b/tests/command_line/test_set_menu.py index 2090f8817..a875e7359 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") @@ -202,6 +210,7 @@ async def test_prompt_passes_is_password_for_sensitive(self): class _FakeSession: def __init__(self, message, **kwargs): captured["is_password"] = kwargs.get("is_password") + captured["style"] = kwargs.get("style") async def prompt_async(self): return "secret" @@ -210,6 +219,7 @@ async def prompt_async(self): result = await _prompt_for_value(sensitive_setting, current_val=None) assert result == "secret" assert captured["is_password"] is True + assert "style" in captured @pytest.mark.asyncio async def test_prompt_no_password_for_non_sensitive(self): @@ -307,6 +317,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/command_line/test_set_menu_render.py b/tests/command_line/test_set_menu_render.py index 288cb4e58..b170dcd24 100644 --- a/tests/command_line/test_set_menu_render.py +++ b/tests/command_line/test_set_menu_render.py @@ -92,7 +92,7 @@ def test_prefix_absent_when_value_is_user_set(self): ) assert "(Default) " not in _flatten(panel) - def test_prefix_uses_dim_italic_style(self): + def test_prefix_uses_semantic_muted_style(self): entry = self._entry_for("true") with patch( "code_puppy.command_line.set_menu_render.is_default_value", @@ -110,7 +110,7 @@ def test_prefix_uses_dim_italic_style(self): total_pages_fn=get_total_pages, ) styles = _styles_with(panel, "(Default) ") - assert any("ansibrightblack" in s and "italic" in s for s in styles) + assert styles == ["class:tui.muted"] # --------------------------------------------------------------------------- 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/command_line/test_yolo_set_override.py b/tests/command_line/test_yolo_set_override.py new file mode 100644 index 000000000..2015b5f75 --- /dev/null +++ b/tests/command_line/test_yolo_set_override.py @@ -0,0 +1,44 @@ +"""Precedence tests for persistent ``/set yolo_mode`` changes.""" + +from unittest.mock import patch + +from code_puppy import config +from code_puppy.command_line.config_apply import apply_setting +from code_puppy.command_line.config_commands import handle_set_command + + +def setup_function(): + config.set_cli_yolo_override(None) + + +def teardown_function(): + config.set_cli_yolo_override(None) + + +def test_persistent_yolo_write_supersedes_cli_override(): + config.set_cli_yolo_override(False) + + with patch("code_puppy.config.set_config_value") as set_config: + result = apply_setting("yolo_mode", "true", reload_agent=False) + + assert result.ok is True + set_config.assert_called_once_with("yolo_mode", "true") + assert config.get_cli_yolo_override() is None + + +def test_yolo_config_clears_cli_override_without_changing_persisted_value(): + config.set_cli_yolo_override(False) + + with ( + patch("code_puppy.config.set_config_value") as set_value, + patch("code_puppy.config.reset_value") as reset_value, + patch("code_puppy.agents.get_current_agent") as current_agent, + patch("code_puppy.messaging.emit_success"), + patch("code_puppy.messaging.emit_info"), + ): + assert handle_set_command("/set yolo_mode config") is True + + set_value.assert_not_called() + reset_value.assert_not_called() + current_agent.return_value.reload_code_generation_agent.assert_called_once() + assert config.get_cli_yolo_override() is None diff --git a/tests/conftest.py b/tests/conftest.py index eda75f553..23822be86 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -9,14 +9,39 @@ import inspect import os import subprocess +import tempfile from copy import deepcopy from unittest.mock import MagicMock import pytest -from code_puppy import config as cp_config -from code_puppy import callbacks as cp_callbacks -from code_puppy.messaging import bottom_bar as cp_bottom_bar +# Config paths are resolved while code_puppy.config is imported, before any +# fixture can run. Point every XDG category at one session-scoped temp root now +# so collection, plugin imports, and tests cannot touch the developer's config. +_XDG_TEMP_DIR = tempfile.TemporaryDirectory(prefix="code_puppy_pytest_xdg_") +_XDG_ENV_VARS = ( + "XDG_CONFIG_HOME", + "XDG_DATA_HOME", + "XDG_CACHE_HOME", + "XDG_STATE_HOME", +) +_ORIGINAL_XDG_ENV = {name: os.environ.get(name) for name in _XDG_ENV_VARS} +for _xdg_name in _XDG_ENV_VARS: + os.environ[_xdg_name] = os.path.join(_XDG_TEMP_DIR.name, _xdg_name.lower()) + +from code_puppy import config as cp_config # noqa: E402 +from code_puppy import callbacks as cp_callbacks # noqa: E402 +from code_puppy.messaging import bottom_bar as cp_bottom_bar # noqa: E402 + + +def pytest_unconfigure(config): + """Restore the invoking shell's XDG environment and remove test state.""" + for name, original_value in _ORIGINAL_XDG_ENV.items(): + if original_value is None: + os.environ.pop(name, None) + else: + os.environ[name] = original_value + _XDG_TEMP_DIR.cleanup() class _InertStream: 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_bottom_bar.py b/tests/messaging/test_bottom_bar.py index 286025da0..9ccaca5f8 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,33 @@ 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 + + +def test_styled_prefix_restores_theme_foreground_for_user_text(bar, tty): + theme_sgr = "\x1b[38;2;76;79;105m" + with patch("code_puppy.callbacks.on_prompt_text_color", return_value="#4c4f69"): + bar.start() + drain(tty) + bar.set_prompt_text( + "> ", "readable latte", len("readable latte"), ["1;34", "1;34"] + ) + + output = written(tty) + prefix_reset = output.index("\x1b[0m") + user_text = output.index("readable latte") + + assert theme_sgr in output[prefix_reset:user_text] + + # ========================================================================= # Non-TTY: silent no-ops # ========================================================================= 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 c33c441db..845d77da7 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 # ========================================================================= @@ -194,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..7ab9ff280 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,39 @@ 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 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/messaging/test_rich_renderer.py b/tests/messaging/test_rich_renderer.py index 3d8e81abe..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): @@ -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/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/prune/test_prune_render.py b/tests/plugins/prune/test_prune_render.py index fff68d94c..264084f1b 100644 --- a/tests/plugins/prune/test_prune_render.py +++ b/tests/plugins/prune/test_prune_render.py @@ -7,7 +7,6 @@ from __future__ import annotations -from code_puppy.plugins.prune import prune_model from code_puppy.plugins.prune.prune_menu import PruneMenu from code_puppy.plugins.prune.prune_model import ( ContextBudget, @@ -120,27 +119,26 @@ def test_unavailable(self): out = render_budget_line(b) assert "unavailable" in out[0][1] - def test_green_under_70(self): + def test_success_under_70(self): b = ContextBudget( used_tokens=10, overhead_tokens=10, context_length=100, available=True ) # 20% used out = render_budget_line(b) - # Style is "fg:ansigreen" (matches C_FOOTER_OK) - assert out[0][0] == prune_model.C_FOOTER_OK + assert out[0][0] == "class:tui.success" - def test_yellow_70_to_90(self): + def test_warning_70_to_90(self): b = ContextBudget( used_tokens=40, overhead_tokens=40, context_length=100, available=True ) # 80% out = render_budget_line(b) - assert out[0][0] == prune_model.C_FOOTER_WARN + assert out[0][0] == "class:tui.warning" - def test_red_over_90(self): + def test_error_over_90(self): b = ContextBudget( used_tokens=50, overhead_tokens=45, context_length=100, available=True ) # 95% out = render_budget_line(b) - assert out[0][0] == prune_model.C_SHELL + assert out[0][0] == "class:tui.error" def test_shows_overflow_when_messages_dont_fit(self): b = ContextBudget( @@ -216,11 +214,11 @@ def test_legend_does_not_falsely_imply_overflow_is_prunable_or_gone(self): assert "removed" not in flat assert "gone" not in flat - def test_legend_uses_in_context_color_for_filled_dot(self): + def test_legend_uses_success_role_for_filled_dot(self): out = render_legend() - green_segments = [seg for seg in out if "●" in seg[1]] - assert green_segments, "Expected a segment containing the green dot" - assert green_segments[0][0] == prune_model.C_FOOTER_OK + success_segments = [seg for seg in out if "●" in seg[1]] + assert success_segments, "Expected a segment containing the filled dot" + assert success_segments[0][0] == "class:tui.success" # ─────────────────────────────────────────────────────────────────────────── diff --git a/tests/plugins/subagent_panel/test_post_tool_call_completion.py b/tests/plugins/subagent_panel/test_post_tool_call_completion.py index 30149bf26..fd9fb8234 100644 --- a/tests/plugins/subagent_panel/test_post_tool_call_completion.py +++ b/tests/plugins/subagent_panel/test_post_tool_call_completion.py @@ -67,6 +67,27 @@ async def test_success_marks_done(): assert not entry.get("failed") +async def test_detached_fork_completion_removes_panel_row(monkeypatch): + state.register("agent-fork-abc123", "fork-agent", model="gpt-5.4") + pushes = [] + monkeypatch.setattr( + register_callbacks, + "_push_panel", + lambda force=False: pushes.append(force), + ) + + await register_callbacks._on_post_tool_call( + tool_name="invoke_agent", + tool_args={}, + result=_StubInvokeResult(session_id="agent-fork-abc123"), + duration_ms=42, + context={"detached_fork": True}, + ) + + assert "agent-fork-abc123" not in state._AGENTS + assert pushes == [True] + + async def test_failure_marks_failed(): """Regression guard: the existing failure-renders-as-red behavior must survive any future edit to ``_on_post_tool_call`` that sits next to the diff --git a/tests/plugins/test_agent_skills.py b/tests/plugins/test_agent_skills.py index 564642087..6bb6bb59f 100644 --- a/tests/plugins/test_agent_skills.py +++ b/tests/plugins/test_agent_skills.py @@ -548,11 +548,14 @@ def test_parse_skill_metadata_nonexistent_path(self, tmp_path, caplog): """Test parsing skill metadata from nonexistent path.""" nonexistent = tmp_path / "does-not-exist" - with caplog.at_level(logging.WARNING): + with caplog.at_level(logging.DEBUG): metadata = parse_skill_metadata(nonexistent) assert metadata is None assert "Skill path does not exist" in caplog.text + # Missing paths are routine (e.g. project dirs without skills) — + # they must not emit user-visible warning noise. + assert not [r for r in caplog.records if r.levelno >= logging.WARNING] def test_parse_skill_metadata_with_file_io_error(self, tmp_path, monkeypatch): """Test parsing when file I/O fails.""" @@ -582,11 +585,12 @@ def test_load_full_skill_content_nonexistent_path(self, tmp_path, caplog): """Test loading full skill content from nonexistent path.""" nonexistent = tmp_path / "does-not-exist" - with caplog.at_level(logging.WARNING): + with caplog.at_level(logging.DEBUG): content = load_full_skill_content(nonexistent) assert content is None assert "Skill path does not exist" in caplog.text + assert not [r for r in caplog.records if r.levelno >= logging.WARNING] def test_get_skill_resources(self, skill_dir_with_resources): """Test getting resource files from skill directory.""" @@ -614,11 +618,12 @@ def test_get_skill_resources_nonexistent_path(self, tmp_path, caplog): """Test getting resources from nonexistent path.""" nonexistent = tmp_path / "does-not-exist" - with caplog.at_level(logging.WARNING): + with caplog.at_level(logging.DEBUG): resources = get_skill_resources(nonexistent) assert len(resources) == 0 assert "Skill path does not exist" in caplog.text + assert not [r for r in caplog.records if r.levelno >= logging.WARNING] # Tests for Prompt Builder Module diff --git a/tests/plugins/test_chatgpt_oauth_coverage.py b/tests/plugins/test_chatgpt_oauth_coverage.py index 5fb95ef3a..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: @@ -231,7 +237,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..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) @@ -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-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.5"]["supports_xhigh_reasoning"] is True + 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") @@ -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): @@ -318,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 b65bfa5d2..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, @@ -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): @@ -1050,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 ( @@ -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 = { + "codex-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 "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") @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,12 +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", - "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"] == [ @@ -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( + "codex-gpt-5.6-" + ) @patch("code_puppy.plugins.chatgpt_oauth.utils.save_chatgpt_models") @patch("code_puppy.plugins.chatgpt_oauth.utils.load_chatgpt_models") @@ -1136,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: @@ -1147,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", }, @@ -1171,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") @@ -1205,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", }, @@ -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/plugins/test_fork_plugin.py b/tests/plugins/test_fork_plugin.py index 1c14b5439..fae1d57c6 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,51 @@ 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 rendered.startswith("\n FORK #7 RESPONSE") + 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,10 +207,44 @@ 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) +async def test_fork_publishes_completion_lifecycle_event(): + lifecycle_events = [] + + async def capture(*args): + lifecycle_events.append(args) + + with ( + patch( + "code_puppy.tools.subagent_invocation._invoke_agent_impl", + new=_fake_impl(), + ), + patch("code_puppy.callbacks.on_post_tool_call", new=capture), + patch.object(rc, "_emit_info"), + patch.object(rc, "_emit_success"), + ): + rc._handle_fork("/fork @qa-kitten inspect lifecycle") + await _wait_for_forks() + + assert len(lifecycle_events) == 1 + tool_name, tool_args, result, duration_ms, context = lifecycle_events[0] + assert tool_name == "invoke_agent" + assert tool_args == { + "agent_name": "qa-kitten", + "prompt": "inspect lifecycle", + } + assert result.session_id == "sess-abc123" + assert duration_ms >= 0 + assert context == {"detached_fork": True} + + async def test_fork_defaults_to_current_agent(): with ( patch( @@ -200,7 +283,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 +311,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/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_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 "