From c9b5b55623e28455f7b273cdab7381933fb61097 Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Thu, 28 May 2026 18:17:04 -0500 Subject: [PATCH 01/29] fix(session): add BaseAgent session-model compatibility methods --- code_puppy/agents/base_agent.py | 24 +++++++++++++++++++ tests/agents/test_base_agent_configuration.py | 21 ++++++++++++++++ 2 files changed, 45 insertions(+) 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/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} From 04316b65484f71abdc8816abcf3d1c42433e3e6b Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Thu, 28 May 2026 18:17:12 -0500 Subject: [PATCH 02/29] fix(ws): avoid cwd history corruption and harden config schema endpoint --- code_puppy/api/routers/config.py | 211 + .../api/tests/test_config_router_schema.py | 17 + code_puppy/api/ws/chat_handler.py | 3898 +++++++++++++++++ 3 files changed, 4126 insertions(+) create mode 100644 code_puppy/api/routers/config.py create mode 100644 code_puppy/api/tests/test_config_router_schema.py create mode 100644 code_puppy/api/ws/chat_handler.py 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/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/ws/chat_handler.py b/code_puppy/api/ws/chat_handler.py new file mode 100644 index 000000000..fbcb4c918 --- /dev/null +++ b/code_puppy/api/ws/chat_handler.py @@ -0,0 +1,3898 @@ +"""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 +import re +import uuid +from pathlib import Path +from typing import Any + +from fastapi import FastAPI, WebSocket, WebSocketDisconnect +from pydantic import TypeAdapter, ValidationError + +from code_puppy.api.db.queries import ( + update_session_working_directory, + write_error_message_to_sqlite, + write_system_message_to_sqlite, + write_turn_to_sqlite, +) +from code_puppy.api.error_parser import parse_api_error as _legacy_parse_api_error +from code_puppy.api.session_context import _validate_session_id, session_manager +from code_puppy.api.ws.attachments import build_file_context_and_attachments +from code_puppy.api.ws.background_save import ( + fire_and_track, + save_agent_result_in_background, +) +from code_puppy.api.ws.connection_manager import connection_manager +from code_puppy.api.ws.schemas import ( + PROTOCOL_VERSION, + ClientMessage, + ServerAgentInvoked, + ServerAssistantMessageDelta, + ServerAssistantMessageEnd, + ServerAssistantMessageStart, + ServerCancelled, + ServerCommandResult, + ServerConfigValue, + ServerError, + ServerMessage, + ServerResponse, + ServerSessionMetaUpdated, + ServerSessionRestored, + ServerSessionSwitched, + ServerStatus, + ServerStreamEnd, + ServerSystem, + ServerToolCall, + ServerToolResult, + ServerUserMessage, + ServerWorkingDirectoryChanged, +) +from code_puppy.config import get_global_model_name +from code_puppy.messaging.bus import get_message_bus +from code_puppy.model_errors import normalize_model_error as _normalize_model_error +from code_puppy.session_storage import generate_heuristic_title +from code_puppy.tools.command_runner import ( + cleanup_session_process_tracking, + init_session_process_tracking, +) + +# Import session context for desk-puppy working directory support +try: + from code_puppy.plugins.walmart_specific.session_context import ( + clear_session_working_directory, + set_session_working_directory, + ) + + HAS_SESSION_CONTEXT = True +except ImportError: + HAS_SESSION_CONTEXT = False + + def set_session_working_directory(d): + pass + + def clear_session_working_directory(): + pass + + +_ClientMessageAdapter = TypeAdapter(ClientMessage) + +logger = logging.getLogger(__name__) + +# === TOOL-CALL-PARITY BRANCH LOADED === +logger.warning( + "🐕 CHAT_HANDLER LOADED FROM TOOL-CALL-PARITY BRANCH - ToolReturnPart fix active!" +) +print("🐕 CHAT_HANDLER LOADED FROM TOOL-CALL-PARITY BRANCH", flush=True) + + +# --------------------------------------------------------------------------- +# WebSocket error frame helpers (see fix/orphaned-tool-use-id-history) +# --------------------------------------------------------------------------- + +_UNKNOWN_ERROR_TYPE = "unknown_error" +_TOOL_HISTORY_ERROR_TYPE = "tool_history_error" + +_CODE_TO_ERROR_TYPE: dict = { + "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: + """Convert an agent-run exception into a structured frontend error dict. + + Wraps ``normalize_model_error`` for primary classification (including the + new ``invalid_tool_history`` code for Anthropic HTTP 400 tool-mismatch + errors), falling back to the legacy ``_legacy_parse_api_error`` for all + other categories so existing behaviour is unchanged. + + Returns: + Dict with keys: user_message, error_type, technical_details, action_required. + """ + 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, + } + # Fall back to legacy parser for unclassified errors + return _legacy_parse_api_error(exc) + + +def has_streamed_content(collected_text) -> bool: + """Return True when the client has already received streaming output chunks. + + Args: + collected_text: List of text chunks accumulated during the agent run. + Elements may be None or empty strings. + + Returns: + True if at least one chunk has substantive (non-whitespace) content. + """ + return any((chunk or "").strip() for chunk in collected_text) + + +def build_error_response_frames( + agent_error: Exception, + collected_text, + session_id: str, +) -> list: + """Build the ordered WebSocket frames to send when an agent error occurs. + + When streaming output was already delivered to the client, a ``stream_end`` + frame (``success=False``) is prepended so the frontend can exit its + streaming state before processing the error frame. Without this handshake + the frontend hangs indefinitely waiting for a ``stream_end`` that never + arrives. + + Args: + agent_error: The exception raised by the agent run. + collected_text: Accumulated streaming chunks from the current turn. + session_id: The WebSocket session identifier. + + Returns: + List of JSON-serialisable dicts to send to the client in order. + """ + frames: list[dict] = [] + 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 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 response + {"type": "response", "content": "...", "done": true, + "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 + ) + + # Flag to track if WebSocket is still open + ws_closed = False + + async def persist_error_payload(data: dict[str, Any]) -> None: + """Persist structured error frames so they survive reloads.""" + if data.get("type") != "error" or not session_id: + return + + try: + await write_error_message_to_sqlite( + session_id=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=(ctx.agent_name if ctx else ""), + model_name=(ctx.model_name if ctx else ""), + timestamp=datetime.datetime.now(datetime.timezone.utc).isoformat(), + ) + except Exception: + logger.warning( + "[WS:%s] Failed to persist error payload to SQLite", + session_id, + exc_info=True, + ) + + async def safe_send_json(data: dict) -> bool: + """Safely send JSON to WebSocket, returns False if connection is closed.""" + nonlocal ws_closed + if ws_closed: + logger.debug( + "[WS:%s] safe_send_json skipped (ws_closed=True): type=%s", + session_id, + data.get("type"), + ) + return False + + msg_type = data.get("type") + if msg_type in { + "error", + "response", + "assistant_message_start", + "assistant_message_end", + }: + # Keep this low-noise but useful for tracing request lifecycle + logger.debug( + "[WS:%s] → send_json type=%s keys=%s", + 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", + session_id, + data.get("error_type"), + data.get("action_required"), + (data.get("error") or "")[:300], + ) + + try: + if msg_type == "error": + await persist_error_payload(data) + await websocket.send_json(data) + return True + except Exception as e: + logger.warning( + "[WS:%s] send_json failed for type=%s: %s", + session_id, + msg_type, + e, + exc_info=True, + ) + if "close message" in str(e).lower() or "closed" in str(e).lower(): + ws_closed = True + logger.debug("WebSocket closed, stopping sends") + return False + + async def send_typed(msg: ServerMessage) -> bool: + """Send a typed protocol message to the client.""" + return await safe_send_json(msg.model_dump(exclude_none=True)) + + async def send_typed_tool_lifecycle( + msg: ServerToolCall | ServerToolResult, + ) -> bool: + """Send tool lifecycle frames to the client.""" + return await send_typed(msg) + + ctx = None + session_title = "" + session_working_directory = "" # Will be set by client or loaded from metadata + session_pinned = False # Track pinned state for persistence + last_context_sent_directory = "" # Track when we last sent directory context + existing_history = None + + # Use provided session_id or generate new one + if session_id: + # Validate session_id to prevent path traversal + try: + _validate_session_id(session_id) + except ValueError as e: + logger.warning("Invalid session_id rejected: %r: %s", session_id, e) + await websocket.close(code=1008, reason="Invalid session ID") + return + + # Client wants to resume/continue a specific session + logger.debug("Client requested session: %s", session_id) + + # SQLite is the source of truth — check session existence there. + 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 e: + logger.warning("Failed to check session in SQLite: %s", e) + + if not existing_history: + logger.debug( + "Session %s not found in SQLite, will create new", session_id + ) + else: + # Generate new timestamp-based session ID + session_timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + session_id = f"WS_session_{session_timestamp}" + logger.debug("Generated new session ID: %s", session_id) + + try: + # --- SESSION ISOLATION: create or load via SessionManager --- + try: + if existing_history is not None: + # Session confirmed in SQLite — load it via SessionManager. + # get_or_load_session checks in-memory first, then falls + # back to a fresh SQLite load. SQLite is the sole source of truth. + ctx = await session_manager.get_or_load_session(session_id) + if ctx is None: + # SQLite confirmed this session exists but load returned None. + # This is unexpected (DB race, corrupt row, or schema mismatch). + # Log a warning and start fresh so the WebSocket can still + # function — silent data loss is worse than a blank session. + 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, + ) + ctx = await session_manager.create_session(session_id) + else: + # Update local state from loaded metadata + session_title = ctx.title + session_working_directory = ctx.working_directory + session_pinned = ctx.pinned + else: + ctx = await session_manager.create_session(session_id) + except Exception as e: + logger.warning("SessionManager init failed, falling back: %s", e) + try: + ctx = await session_manager.create_session(session_id) + except Exception: + logger.error("SessionManager fallback also failed", exc_info=True) + await websocket.close(code=1011, reason="Session init failed") + return + + # Mark session as active (cancels any pending cleanup from previous disconnect) + await session_manager.mark_session_active(session_id) + + # Initialize per-session process tracking (ContextVar isolation) + init_session_process_tracking() + + # Set MessageBus session context for this WS connection + try: + bus = get_message_bus() + bus.set_session_context(session_id) + except Exception: + logger.debug("MessageBus session context not available") + + # Convenience aliases — use ctx.agent everywhere instead of get_current_agent() + agent = ctx.agent + agent_name = ctx.agent_name + model_name = ctx.model_name + + # ── Persist initial session context to SQLite (new sessions only) ─────── + # Write a config row (agent + model) and, if a CWD is set, a directory + # row so the FE cold-load path sees them before the first user message. + # Skipped for resumed sessions — their init rows were written at creation. + if existing_history is None: + try: + _now_iso = datetime.datetime.now(datetime.timezone.utc).isoformat() + _init_agent = agent_name or "code-puppy" + _init_model = model_name or "unknown" + await write_system_message_to_sqlite( + session_id=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 session_working_directory: + _path_segments = session_working_directory.split("/")[-3:] + _relative = "/".join(s for s in _path_segments if s) + await write_system_message_to_sqlite( + session_id=session_id, + system_message_type="directory", + content=f"Starting in {_relative}", + system_message_path=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, + ) + + # Send welcome message + await send_typed( + ServerSystem( + content=f"Connected! Session: {session_id}", + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + resumed=existing_history is not None, + protocol_version=PROTOCOL_VERSION, + ) + ) + + # If we resumed an existing session, notify the client + if existing_history and ctx: + try: + # History is already loaded into ctx.agent by session_manager.load_session(). + # Pull it back out for the client-side UI metadata notification. + loaded_messages = ctx.agent.get_message_history() + message_count = len(loaded_messages) if loaded_messages else 0 + + # Notify client about restored session + await send_typed( + ServerSessionRestored( + session_id=session_id, + message_count=message_count, + title=session_title, + ui_metadata=[], + ) + ) + + # Replay system messages (agent/model switches, CWD banners) to UI. + # These are persisted with pydantic_json=NULL so _load_from_sqlite + # skips them — we must send them separately so the UI shows them. + try: + from code_puppy.api.db.queries import get_active_messages + + rows = await get_active_messages(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=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), + session_id, + ) + except Exception as sys_exc: + logger.warning( + "Failed to replay system messages for session %s: %s", + session_id, + sys_exc, + ) + + agent = ctx.agent # Refresh local alias + logger.debug("Restored %d messages to session agent", message_count) + except Exception as e: + logger.warning("Failed to restore session history: %s", e) + # Track active streaming task for cancellation + active_drain_task = None + # Track active agent task (run_with_mcp) for cancellation + active_agent_task: asyncio.Task | None = None + # Track stop_draining event for cancellation across iterations + stop_draining = asyncio.Event() + + 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") == "switch_agent": + agent_name = msg.get("agent_name") + if agent_name: + try: + new_agent = await session_manager.switch_agent( + session_id, agent_name + ) + agent = new_agent # Update local alias + ctx.agent = new_agent + model_name = ( + new_agent.get_model_name() + if new_agent + else "unknown" + ) + + # Persist to SQLite so the FE cold-load path sees it + try: + await write_system_message_to_sqlite( + session_id=session_id, + system_message_type="config", + content=f"🔄 Switched to {agent_name} ({model_name})", + agent_name=agent_name, + model_name=model_name, + ) + except Exception as _sw_exc: + logger.warning( + "Agent-switch SQLite write failed: %s", + _sw_exc, + exc_info=True, + ) + + await send_typed( + ServerSystem( + content=f"🔄 Switched to {agent_name} ({model_name})", + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + ) + ) + except Exception as e: + logger.error("Error switching agent: %s", e) + await send_typed( + ServerError( + error=f"Failed to switch to agent {agent_name}: {str(e)}", + session_id=session_id, + ) + ) + + elif msg.get("type") == "switch_model": + model_name = msg.get("model_name") or msg.get("model") + if model_name: + try: + await session_manager.switch_model( + session_id, model_name + ) + agent = ctx.agent # Refresh local alias + logger.debug("Switched model to: %s", model_name) + + # Persist to SQLite so the FE cold-load path sees it + _sw_agent = ctx.agent_name or agent_name or "code-puppy" + try: + await write_system_message_to_sqlite( + session_id=session_id, + system_message_type="config", + content=f"🔄 Switched to {_sw_agent} ({model_name})", + agent_name=_sw_agent, + model_name=model_name, + ) + except Exception as _sw_exc: + logger.warning( + "Model-switch SQLite write failed: %s", + _sw_exc, + exc_info=True, + ) + + await send_typed( + ServerSystem( + content=f"🔄 Switched to {_sw_agent} ({model_name})", + session_id=session_id, + model_name=model_name, + agent_name=_sw_agent, + ) + ) + except Exception as e: + logger.error("Error switching model: %s", e) + await send_typed( + ServerError( + error=f"Failed to switch to model {model_name}: {str(e)}", + session_id=session_id, + ) + ) + else: + await send_typed( + ServerError( + error="No model_name provided for switch_model", + session_id=session_id, + ) + ) + + elif msg.get("type") == "switch_session": + """ + Switch to a different session without reconnecting WebSocket. + + How it works: + 1. Client sends: {"type": "switch_session", "session_id": "WS_session_20260120_124356"} + 2. Server loads session from ~/.code_puppy/ws_sessions/{session_id}.pkl + 3. Server restores message history to the current agent + 4. Server responds with session_switched event containing message count and title + 5. Client can continue chatting with full conversation context + + This approach keeps a single WebSocket connection alive while allowing + users to switch between multiple sessions instantly. All session data + stays on the local filesystem for privacy. + """ + # Cancel any active streaming first + if active_drain_task and not active_drain_task.done(): + logger.debug( + "Cancelling active streaming due to session switch" + ) + stop_draining.set() # Signal drain to stop + active_drain_task.cancel() + try: + await active_drain_task + except asyncio.CancelledError: + pass + stop_draining.clear() # Reset for next streaming + + 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=session_id, + ) + ) + continue + + # Validate session_id to prevent path traversal + try: + _validate_session_id(new_session_id) + except ValueError as e: + logger.warning( + f"Invalid session_id in switch_session: {new_session_id!r}" + ) + await send_typed( + ServerError( + error=f"Invalid session ID: {e}", + session_id=session_id, + ) + ) + continue + + logger.debug("Switching to session: %s", new_session_id) + + try: + # Check if target session exists in SQLite + try: + from code_puppy.api.db.queries import ( + session_exists as _se, + ) + + _target_exists = await _se(new_session_id) + except Exception: + _target_exists = False + + if not _target_exists: + # Session doesn't exist - create new one with this ID + logger.debug( + f"Session {new_session_id} not found, creating new" + ) + + # Save current session and mark inactive (keep in memory for 15 min) + try: + await session_manager.save_session(session_id) + except Exception: + pass + await session_manager.mark_session_inactive(session_id) + + session_id = new_session_id + session_title = "" + session_working_directory = "" + session_pinned = False + + ctx = await session_manager.create_session(session_id) + # Mark the new session as active + await session_manager.mark_session_active(session_id) + agent = ctx.agent + agent_name = ctx.agent_name + model_name = ctx.model_name + + await send_typed( + ServerSessionSwitched( + session_id=new_session_id, + message_count=0, + title="", + created=True, + agent_name=agent_name, + model_name=model_name, + ) + ) + continue + + # Load existing session via SessionManager + # Save current session and mark inactive (keep in memory for 15 min) + try: + await session_manager.save_session(session_id) + except Exception: + pass + await session_manager.mark_session_inactive(session_id) + + session_id = new_session_id + + # Try loading via SessionManager (checks in-memory first, then SQLite) + loaded_ctx = await session_manager.get_or_load_session( + session_id + ) + if loaded_ctx is not None: + ctx = loaded_ctx + session_title = ctx.title + session_working_directory = ctx.working_directory + session_pinned = ctx.pinned + else: + # Legacy session or HMAC missing — load metadata + # from JSON (safe) and create fresh context + new_title = "" + new_working_directory = "" + new_pinned = False + try: + from code_puppy.api.db.queries import ( + get_session_metadata as _gsm, + ) + + _sm = await _gsm(session_id) or {} + new_title = _sm.get("title", "") + new_working_directory = _sm.get( + "working_directory", "" + ) + new_pinned = bool(_sm.get("pinned", False)) + except Exception: + pass + + ctx = await session_manager.create_session(session_id) + ctx.title = new_title + ctx.working_directory = new_working_directory + ctx.pinned = new_pinned + session_title = new_title + session_working_directory = new_working_directory + session_pinned = new_pinned + + # Mark the new session as active + await session_manager.mark_session_active(session_id) + + agent = ctx.agent + agent_name = ctx.agent_name + model_name = ctx.model_name + message_count = len(ctx.agent.get_message_history() or []) + logger.debug( + f"Restored {message_count} messages to session agent" + ) + + await send_typed( + ServerSessionSwitched( + session_id=new_session_id, + message_count=message_count, + title=session_title, + working_directory=session_working_directory, + created=False, + agent_name=agent_name, + model_name=model_name, + ) + ) + + logger.debug( + f"Switched to session {new_session_id} with {message_count} messages" + ) + + except Exception as e: + logger.error("Error switching session: %s", e) + await send_typed( + ServerError( + error=f"Failed to switch to session {new_session_id}: {str(e)}", + session_id=session_id, + ) + ) + + elif msg.get("type") == "set_working_directory": + new_directory = msg.get("directory", "") + if new_directory: + # Expand ~ and resolve symlinks before validation + new_directory = str( + Path(new_directory).expanduser().resolve() + ) + logger.info( + "[CWD DEBUG] set_working_directory received: new=%r, current=%r, session=%s", + new_directory, + session_working_directory, + session_id, + ) + # Validate the directory exists + if Path(new_directory).is_dir(): + # Skip if directory hasn't actually changed (avoids duplicate banners on reload) + if new_directory == session_working_directory: + logger.info( + "[CWD DEBUG] Skipping unchanged directory: %r, session=%s", + new_directory, + session_id, + ) + await send_typed( + ServerWorkingDirectoryChanged( + directory=new_directory, + success=True, + session_id=session_id, + unchanged=True, + ) + ) + continue # Skip to next message, don't write banner + + session_working_directory = new_directory + logger.debug( + "Working directory set to: %s", + session_working_directory, + ) + + # NOTE: Do NOT append raw dict entries into agent + # message history here. The runtime expects typed + # ModelMessage objects; dict injection can corrupt + # subsequent turns and lead to result=None failures. + + # Persist directory banner to SQLite so FE cold-load sees it + try: + # Bug 4: use ctx.agent_name / ctx.model_name directly + # so this works even when ctx.agent is None. + _cwd_agent = ( + ctx.agent_name if ctx else None + ) or "code-puppy" + _cwd_model = ( + ctx.model_name if ctx else None + ) or "unknown" + _cwd_segs = 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 _gpn + + _pname = _gpn() or "puppy" + logger.info( + "[CWD DEBUG] Writing CWD banner to SQLite: path=%r, session=%s", + session_working_directory, + session_id, + ) + await write_system_message_to_sqlite( + session_id=session_id, + system_message_type="directory", + content=f"{_pname} is now at {_cwd_rel}", + system_message_path=session_working_directory, + agent_name=_cwd_agent, + model_name=_cwd_model, + ) + # Bug 5: also stamp working_directory on the sessions row + _now_cwd = datetime.datetime.now( + datetime.timezone.utc + ).isoformat() + await update_session_working_directory( + session_id=session_id, + working_directory=session_working_directory, + updated_at=_now_cwd, + ) + except Exception as _cwd_exc: + logger.warning( + "CWD SQLite write failed: %s", + _cwd_exc, + exc_info=True, + ) + + await send_typed( + ServerWorkingDirectoryChanged( + directory=session_working_directory, + success=True, + session_id=session_id, + ) + ) + else: + await send_typed( + ServerWorkingDirectoryChanged( + directory=new_directory, + success=False, + error="Directory does not exist", + session_id=session_id, + ) + ) + else: + await send_typed( + ServerError( + error="No directory provided for set_working_directory", + session_id=session_id, + ) + ) + + elif msg.get("type") == "update_session_meta": + # Update session metadata (pinned, title) — persisted to SQLite. + try: + # Update in-memory state (local vars AND SessionContext) + if "pinned" in msg and isinstance(msg["pinned"], bool): + session_pinned = msg["pinned"] + if ctx is not None: + ctx.pinned = session_pinned + if "title" in msg and isinstance(msg["title"], str): + session_title = msg["title"] + if ctx is not None: + ctx.title = session_title + + # Persist to SQLite (single source of truth). + # Use a targeted UPDATE — NOT upsert_session() — so that + # message_count, total_tokens, and other stats are + # never zeroed out by a rename/pin operation. + 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=session_id, + title=session_title, + pinned=session_pinned, + updated_at=_dt.now(_tz.utc).isoformat(), + ) + logger.debug( + f"Updated session meta in SQLite for {session_id}: pinned={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=session_id, + pinned=session_pinned, + title=session_title, + ) + ) + except Exception as e: + logger.error("Error updating session meta: %s", e) + await send_typed( + ServerError( + error=f"Failed to update session metadata: {str(e)}", + session_id=session_id, + ) + ) + + elif msg.get("type") == "get_config": + 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=session_id, + ) + ) + else: + await send_typed( + ServerError( + error="No key provided for get_config", + session_id=session_id, + ) + ) + + elif msg.get("type") == "set_config": + 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( + f"Config set: {config_key} = {config_value}" + ) + await send_typed( + ServerConfigValue( + key=config_key, + value=config_value, + success=True, + session_id=session_id, + ) + ) + except Exception as e: + await send_typed( + ServerError( + error=f"Failed to set config: {e}", + session_id=session_id, + ) + ) + else: + await send_typed( + ServerError( + error="No key provided for set_config", + session_id=session_id, + ) + ) + + # Handle slash command execution + elif msg.get("type") == "command": + command_str = msg.get("command", "") + logger.debug("Command requested: %s", command_str) + + try: + from io import StringIO + + from rich.console import Console + + from code_puppy.command_line.command_handler import ( + get_commands_help, + handle_command, + ) + + output = None + captured_messages = [] + + # Special handling for /help - we can get the text directly + cmd_name = ( + command_str.strip().lstrip("/").split()[0] + if command_str.strip() + else "" + ) + if cmd_name in ("help", "h"): + # Get help text directly + help_text = get_commands_help() + # Render to plain text + 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: + # For other commands, just execute and report success/failure + # The actual output goes to the message queue/terminal + result = handle_command(command_str) + # Some commands might return a string as output + if isinstance(result, str): + output = result + result = True + + # Send command result + await send_typed( + ServerCommandResult( + command=command_str, + success=result is True or result is not False, + output=output, + messages=captured_messages, + result=str(result) + if result and result is not True + else None, + session_id=session_id, + ) + ) + logger.debug( + f"Command executed: {command_str} -> success={result is True or result is not False}, output_len={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, + ) + ) + continue + + # Handle cancel/interrupt request + elif msg.get("type") == "cancel": + logger.debug( + "Cancel request received - stopping active streaming and agent task" + ) + + # Cancel any active streaming + if active_drain_task and not active_drain_task.done(): + stop_draining.set() # Signal drain to stop + active_drain_task.cancel() + try: + await active_drain_task + except asyncio.CancelledError: + pass + stop_draining.clear() # Reset for next streaming + logger.debug("Active streaming cancelled") + + # Cancel the in-flight agent run (run_with_mcp) if present + if active_agent_task and not active_agent_task.done(): + logger.debug( + "Cancelling active agent task due to user interrupt" + ) + active_agent_task.cancel() + try: + await active_agent_task + except asyncio.CancelledError: + logger.debug("Active agent task cancelled successfully") + active_agent_task = None + + # Send confirmation + await send_typed( + ServerStatus( + status="cancelled", + session_id=session_id, + ) + ) + continue + + elif msg.get("type") == "permission_response": + # Handle permission response from user + from code_puppy.api.permissions import ( + handle_permission_response, + ) + + request_id = msg.get("request_id") + approved = msg.get("approved", False) + + if request_id: + handled = handle_permission_response(request_id, approved) + if handled: + logger.debug( + f"[Permission] ✅ Handled response: {request_id} = {approved}" + ) + else: + logger.warning( + f"[Permission] ❌ Unknown request: {request_id}" + ) + else: + logger.error( + "[WebSocket] ❌ No request_id in permission_response!" + ) + continue + + elif msg.get("type") == "message": + # Set WebSocket context for permission requests + from code_puppy.api.permission_plugin import ( + set_suppress_emitter_tool_events, + set_websocket_context, + ) + + set_websocket_context(websocket, 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 + event_queue = None + collected_text = [] + stop_draining.clear() # Reset for this message + drain_task = None + active_parts: dict[ + int, dict + ] = {} # Track message parts by index + tool_id_aliases: dict[str, str] = {} + tool_group_ids: dict[ + str, str + ] = {} # tool_id -> tool_group_id mapping + b1_streaming_used = ( + False # Track if B1 streaming sent any content + ) + agent_error = None # Will hold exception if agent run fails + + async def drain_events_concurrent( + ready_event: asyncio.Event = None, + ): + """Background task to drain events and send structured messages in real-time.""" + nonlocal collected_text, b1_streaming_used, ws_closed + 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" + ) + current_tool_name = None # Track active tool name + current_tool_group_id: str | None = ( + None # Track current tool batch group ID + ) + + # Track pending tool calls for tool_result generation + # {tool_id: {tool_name, start_time, part_index}} + pending_tool_calls: dict[str, dict] = {} + + def resolve_pending_tool_id( + *, + 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 pending_tool_calls + ): + return tool_call_id + + if tool_call_id: + for ( + _pending_id, + _pending_info, + ) in 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 pending_tool_calls.items(): + if ( + _pending_info.get("tool_name") + == tool_name + ): + return _pending_id + + return None + + def resolve_tool_group_id( + *, + tool_id: str | None = None, + pending_info: dict | 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. + + Priority order: + 1) pending tool info + 2) tool_id -> group map + 3) explicit fallback passed by caller + 4) current in-flight batch group + 5) synthetic deterministic-ish fallback (last resort) + """ + nonlocal current_tool_group_id + + 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 = 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 = 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: + tool_group_ids[tool_id] = group_id + if pending_info is not None: + pending_info["tool_group_id"] = group_id + elif tool_id in pending_tool_calls: + pending_tool_calls[tool_id][ + "tool_group_id" + ] = group_id + + if current_tool_group_id is None: + current_tool_group_id = group_id + + return group_id + + event_count = 0 + first_iteration = True + while not stop_draining.is_set(): + # Exit if WebSocket is closed + if 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( + 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 event_queue.empty(): + try: + event = 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 + + # 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": + tool_name = event_data.get( + "tool_name", "unknown" + ) + tool_args = event_data.get( + "tool_args", {} + ) + tool_id = str(uuid.uuid4())[:8] + + # Generate tool group ID for this batch if not already set + if current_tool_group_id is None: + current_tool_group_id = ( + f"tg-{str(uuid.uuid4())[:8]}" + ) + + logger.debug( + "[ws] tool_call: %s", tool_name + ) + + # Track current tool name + current_tool_name = tool_name + + # Register in pending_tool_calls so + # tool_call_complete can echo back the + # same tool_id the frontend already knows. + pending_tool_calls[tool_id] = { + "tool_name": tool_name, + "start_time": time_module.time(), + "tool_group_id": 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=current_agent_name, + model_name=current_model_name, + tool_group_id=current_tool_group_id, + ) + ) + + elif event_type == "tool_call_complete": + 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 + ) + + # Clear current tool name after completion + current_tool_name = None + + # Match the tool_id we generated at + # tool_call_start so the frontend can + # correlate result → call by ID, not + # just by name (names are not unique). + matching_tool_id = next( + ( + tid + for tid, info in pending_tool_calls.items() + if info["tool_name"] + == tool_name + ), + None, + ) + matching_pending_info = ( + pending_tool_calls.get( + matching_tool_id + ) + if matching_tool_id + else None + ) + tool_group_id_for_result = resolve_tool_group_id( + tool_id=matching_tool_id, + pending_info=matching_pending_info, + fallback_group_id=current_tool_group_id, + tool_name=tool_name, + source="tool_call_complete", + ) + + if matching_tool_id: + 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=current_agent_name, + model_name=current_model_name, + tool_group_id=tool_group_id_for_result, + ) + ) + + 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", {} + ) + initial_content = "" + + if ( + hasattr(part_obj, "content") + and part_obj.content + ): + initial_content = ( + part_obj.content + ) + elif isinstance( + part_obj, dict + ) and part_obj.get("content"): + initial_content = part_obj.get( + "content", "" + ) + + if part_type in ( + "TextPart", + "ThinkingPart", + ): + msg_type = ( + "thinking" + if part_type + == "ThinkingPart" + else "text" + ) + + # When a TextPart starts, any pending tool calls have completed + # Send status-only tool_result during streaming to update UI + # The actual result data will be sent later from extraction code + if ( + pending_tool_calls + and part_type == "TextPart" + ): + for ( + tool_id, + tool_info, + ) in list( + pending_tool_calls.items() + ): + if tool_info.get( + "status_only_sent" + ): + continue + duration_ms = ( + time_module.time() + - tool_info[ + "start_time" + ] + ) * 1000 + logger.debug( + f"[WebSocket] Sending status-only tool_result for: {tool_info['tool_name']} (duration: {duration_ms:.1f}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=current_agent_name, + model_name=current_model_name, + tool_group_id=resolve_tool_group_id( + tool_id=tool_id, + pending_info=tool_info, + fallback_group_id=current_tool_group_id, + tool_name=tool_info.get( + "tool_name" + ), + source="status_only_result", + ), + ) + ) + tool_info[ + "status_only_sent" + ] = True + current_tool_name = None + + # Check if part already exists (created by early delta) + if part_index in active_parts: + # Just update the type, keep existing message_id + active_parts[part_index][ + "type" + ] = msg_type + message_id = active_parts[ + part_index + ]["id"] + # If there's initial content, add it to the existing accumulated content + if initial_content: + active_parts[ + part_index + ]["content"] = ( + initial_content + + active_parts[ + part_index + ]["content"] + ) + collected_text.insert( + 0, initial_content + ) # Prepend to collected text too + logger.debug( + f"[Stream Debug] Part already exists, reusing message_id={message_id}" + ) + # Don't send another start event + continue + + message_id = f"msg-{int(time_module.time() * 1000)}-{part_index}" + active_parts[part_index] = { + "id": message_id, + "type": msg_type, + "content": initial_content, # Start with initial content if present + } + + # If there's initial content, add it to collected_text + if initial_content: + collected_text.append( + initial_content + ) + + # New assistant output part indicates turn boundary for tool grouping + if part_index == 0: + 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=current_agent_name, + model_name=current_model_name, + tool_name=current_tool_name, + ).model_dump( + exclude_none=True + ) + ) + + # If there's initial content, send it as a delta immediately + 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=current_agent_name, + model_name=current_model_name, + tool_name=current_tool_name, + ) + await safe_send_json( + _delta.model_dump( + exclude_none=True + ) + ) + # Mark that B1 streaming was used (nonlocal already declared above) + b1_streaming_used = True + + elif part_type == "ToolCallPart": + # Extract tool info from the part + tool_name = "unknown" + tool_call_id = None + tool_args_str = "" + + # Extract tool info from part_obj (dict or object) + if hasattr( + part_obj, "tool_name" + ): + tool_name = ( + part_obj.tool_name + ) + elif isinstance(part_obj, dict): + tool_name = part_obj.get( + "tool_name", "unknown" + ) + + if hasattr( + part_obj, "tool_call_id" + ): + tool_call_id = ( + part_obj.tool_call_id + ) + elif isinstance(part_obj, dict): + tool_call_id = part_obj.get( + "tool_call_id" + ) + + if hasattr(part_obj, "args"): + tool_args_str = ( + part_obj.args or "" + ) + elif isinstance(part_obj, dict): + tool_args_str = ( + part_obj.get("args", "") + or "" + ) + + tool_id = ( + tool_call_id + or str(uuid.uuid4())[:8] + ) + + # Track this tool call part with args buffer for accumulation + # Don't send tool_call yet - wait for part_end when args are complete + 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, # Buffer for accumulating args_delta + "start_time": time_module.time(), # Track when tool started + } + + # Track current tool name + current_tool_name = tool_name + + logger.debug( + f"[WebSocket] ToolCallPart started: {tool_name} (id: {tool_id})" + ) + # tool_call event will be sent on part_end when args are complete + + elif part_type == "ToolReturnPart": + logger.info( + f"[WebSocket] ToolReturnPart detected! part_index={part_index}" + ) + # Extract tool result info from the part + tool_call_id = None + tool_content = None + + # Extract tool_call_id + if hasattr( + part_obj, "tool_call_id" + ): + tool_call_id = ( + part_obj.tool_call_id + ) + elif isinstance(part_obj, dict): + tool_call_id = part_obj.get( + "tool_call_id" + ) + + # Extract content (the tool result) + if hasattr(part_obj, "content"): + tool_content = ( + part_obj.content + ) + elif isinstance(part_obj, dict): + tool_content = part_obj.get( + "content" + ) + + # Try to serialize complex result objects + if ( + tool_content + and not isinstance( + tool_content, + ( + str, + dict, + list, + int, + float, + bool, + type(None), + ), + ) + ): + try: + import json + + # Try to convert to dict first (for Pydantic models) + if hasattr( + tool_content, + "model_dump", + ): + tool_content = tool_content.model_dump() + elif hasattr( + tool_content, "dict" + ): + tool_content = tool_content.dict() + elif hasattr( + tool_content, + "__dict__", + ): + tool_content = tool_content.__dict__ + else: + tool_content = str( + tool_content + ) + except Exception as e: + logger.debug( + f"[WebSocket] Could not serialize tool result: {e}" + ) + tool_content = str( + tool_content + ) + + # Find the corresponding pending tool call, store and SEND the result + _result_sent = False + if tool_call_id: + resolved_pending_id = resolve_pending_tool_id( + tool_call_id=tool_call_id + ) + if resolved_pending_id: + _result_sent = True + pending_info = 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( + tool_id=resolved_pending_id, + pending_info=pending_info, + fallback_group_id=current_tool_group_id, + tool_name=_tool_name, + source="tool_return_resolved", + ) + logger.info( + f"[WebSocket] ToolReturnPart: Sending result for {_tool_name} (id: {resolved_pending_id}, raw: {tool_call_id})" + ) + # Send the REAL tool_result with actual content + 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=current_agent_name + or "code-puppy", + model_name=current_model_name + or "unknown", + tool_group_id=_group_id, + ) + ) + else: + logger.warning( + f"[WebSocket] ToolReturnPart: Could not resolve tool_call_id {tool_call_id}, pending keys: {list(pending_tool_calls.keys())}" + ) + + # Fallback: proximity-based matching if no tool_call_id or resolution failed + if not _result_sent: + for ( + pending_id, + pending_info, + ) in sorted( + pending_tool_calls.items(), + key=lambda x: abs( + x[1].get( + "part_index", + 9999, + ) + - part_index + ), + ): + if ( + abs( + pending_info.get( + "part_index", + 9999, + ) + - part_index + ) + <= 3 + ): + # Close enough - assume it's the result for this tool call + 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( + tool_id=pending_id, + pending_info=pending_info, + fallback_group_id=current_tool_group_id, + tool_name=_tool_name, + source="tool_return_proximity", + ) + logger.info( + f"[WebSocket] ToolReturnPart: Sending result (by proximity) for {_tool_name} (id: {pending_id})" + ) + # Send the REAL tool_result with actual content + 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=current_agent_name + or "code-puppy", + model_name=current_model_name + or "unknown", + tool_group_id=_group_id, + ) + ) + _result_sent = True + break + + # Warn if we couldn't send the result + if not _result_sent: + logger.warning( + f"[WebSocket] ToolReturnPart: Could NOT send result! tool_call_id={tool_call_id}, pending_tool_calls={list(pending_tool_calls.keys())}, part_index={part_index}" + ) + + # Track this part for cleanup + active_parts[part_index] = { + "id": f"tool-return-{part_index}", + "type": "tool_return", + "tool_call_id": tool_call_id, + "content": tool_content, + } + + 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", {} + ) + + # Handle ToolCallPartDelta - accumulate args + if ( + delta_type + == "ToolCallPartDelta" + ): + args_delta = "" + if hasattr( + delta_obj, "args_delta" + ): + args_delta = ( + delta_obj.args_delta + or "" + ) + elif isinstance( + delta_obj, dict + ): + args_delta = ( + delta_obj.get( + "args_delta", "" + ) + or "" + ) + + if ( + args_delta + and part_index + in active_parts + ): + part_info = active_parts[ + part_index + ] + if ( + part_info.get("type") + == "tool_call" + ): + part_info[ + "args_buffer" + ] = ( + part_info.get( + "args_buffer", + "", + ) + + args_delta + ) + continue # Don't process as text delta + + # Extract content_delta (supports both direct and nested formats) + content_delta = inner_data.get( + "content_delta", "" + ) + if not content_delta: + if isinstance(delta_obj, dict): + content_delta = ( + delta_obj.get( + "content_delta", "" + ) + ) + + if content_delta: + collected_text.append( + content_delta + ) + + # Ensure we have an entry for this part (handles delta before start) + if ( + part_index + not in active_parts + ): + # Create entry with consistent message_id + message_id = f"msg-{int(time_module.time() * 1000)}-{part_index}" + active_parts[part_index] = { + "id": message_id, + "type": "text", # Default, will be updated by part_start if it arrives + "content": "", + } + # Send a start event so client creates the message + logger.debug( + f"[Stream Debug] Creating part on first delta: message_id={message_id}" + ) + if part_index == 0: + 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=current_agent_name, + model_name=current_model_name, + tool_name=current_tool_name, + ).model_dump( + exclude_none=True + ) + ) + + part_info = active_parts[ + part_index + ] + message_id = part_info["id"] + + if part_index in active_parts: + active_parts[part_index][ + "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=current_agent_name, + model_name=current_model_name, + tool_name=current_tool_name, + ) + await safe_send_json( + _delta.model_dump( + exclude_none=True + ) + ) + # Mark that B1 streaming was used + b1_streaming_used = True + + elif inner_type == "part_end": + part_index = inner_data.get( + "index", 0 + ) + part_info = active_parts.get( + part_index, {} + ) + part_type_info = part_info.get( + "type", "text" + ) + message_id = part_info.get( + "id", f"msg-{part_index}" + ) + full_content = part_info.get( + "content", "" + ) + + # Handle tool_call parts - send tool_call event now that args are complete + if part_type_info == "tool_call": + 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" + ) + ) + + # Parse args if they're a JSON string + try: + import json + + args_dict = ( + json.loads( + tool_args_str + ) + if tool_args_str + else {} + ) + except ( + json.JSONDecodeError, + TypeError, + ): + args_dict = {} + + logger.debug( + f"[WebSocket] Sending tool_call (args complete): {tool_name}" + ) + + # Generate tool group ID for this batch if not already set + if ( + current_tool_group_id + is None + ): + current_tool_group_id = f"tg-{str(uuid.uuid4())[:8]}" + + # Send tool_call now that we have complete args + 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=current_agent_name, + model_name=current_model_name, + tool_group_id=current_tool_group_id, + ) + ) + + # Track as pending tool call for later tool_result + 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, # Will be populated by ToolReturnPart + "tool_group_id": current_tool_group_id, + } + if raw_tool_call_id: + tool_id_aliases[ + raw_tool_call_id + ] = tool_id + # Also track tool_group_id for pre-stream_end extraction + if current_tool_group_id: + tool_group_ids[tool_id] = ( + current_tool_group_id + ) + + # Clear current tool name tracking + current_tool_name = None + else: + await safe_send_json( + ServerAssistantMessageEnd( + message_id=message_id, + part_index=part_index, + full_content=full_content, + timestamp=time_module.time(), + session_id=session_id, + agent_name=current_agent_name, + model_name=current_model_name, + tool_name=current_tool_name, + ).model_dump( + exclude_none=True + ) + ) + + if part_index in active_parts: + del active_parts[part_index] + + except Exception as send_err: + error_msg = str(send_err).lower() + if ( + "close message" in error_msg + or "closed" in error_msg + ): + 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 = event_queue.get_nowait() + event_type = event.get("type", "") + event_data = event.get("data", {}) + final_count += 1 + + if event_type == "stream_event": + inner_type = event_data.get( + "event_type", "" + ) + inner_data = event_data.get( + "event_data", {} + ) + if inner_type == "part_delta": + content_delta = inner_data.get( + "content_delta", "" + ) + if not content_delta: + delta_obj = inner_data.get( + "delta", {} + ) + if isinstance(delta_obj, dict): + content_delta = delta_obj.get( + "content_delta", "" + ) + if content_delta: + collected_text.append(content_delta) + 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}" + ) + + # Event to signal drain task is ready + drain_ready = asyncio.Event() + + try: + from code_puppy.plugins.frontend_emitter.emitter import ( + subscribe, + unsubscribe, + ) + + event_queue = subscribe(session_id=session_id) + logger.debug( + "Subscribed to frontend emitter for streaming" + ) + + # Modify drain function to signal when ready + async def drain_events_with_signal(): + """Wrapper that signals readiness before starting drain loop.""" + await drain_events_concurrent(drain_ready) + + # Start concurrent drain task + drain_task = asyncio.create_task( + drain_events_with_signal() + ) + active_drain_task = ( + drain_task # Track for potential cancellation + ) + + # Wait for drain task to be ready before proceeding + await drain_ready.wait() + + except ImportError: + logger.warning("Frontend emitter not available") + + 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, + ) + + # Inject working directory as a system message in the + # agent's conversation history (once per directory change) + # so the LLM knows the CWD without polluting user messages. + message_to_send = user_message + if ( + session_working_directory + and session_working_directory + != last_context_sent_directory + ): + from pydantic_ai.messages import ( + ModelRequest, + SystemPromptPart, + ) + + wd_system_msg = ModelRequest( + parts=[ + SystemPromptPart( + content=( + f"The user's current working directory is updated to" + f" {session_working_directory}" + ) + ) + ] + ) + agent.append_to_message_history(wd_system_msg) + last_context_sent_directory = ( + session_working_directory + ) + logger.debug( + "Injected working directory system message: %s", + session_working_directory, + ) + + logger.debug( + f"Calling run_with_mcp with message: {message_to_send[:100]}..." + ) + + # Build file context and binary attachments + file_context, binary_attachments = ( + build_file_context_and_attachments(msg) + ) + + # CHANGE 1: Update attachment metadata for UI (variables already initialized) + if msg.get("attachments"): + from pathlib import Path as _AttachmentPath + + for raw_path in msg.get("attachments", []): + if ( + isinstance(raw_path, str) + and raw_path.strip() + ): + try: + file_path = _AttachmentPath( + 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( + f"Error building attachment metadata for '{raw_path}': {e}" + ) + + # Prepend file context to the message + if file_context: + message_to_send = ( + file_context + "\n\n" + message_to_send + ) + logger.debug( + f"Added file context ({len(file_context)} chars)" + ) + + run_kwargs = {} + if binary_attachments: + run_kwargs["attachments"] = binary_attachments + logger.debug( + f"Including {len(binary_attachments)} binary attachment(s)" + ) + + # ────────────────────────────────────────────────────────────────── + # 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, + ) + + # Run the agent in its own task so it can be cancelled independently. + # Suppress duplicate emitter lifecycle callbacks only during the + # run_with_mcp execution window (pre/post_tool_call hooks fire there). + set_suppress_emitter_tool_events(True) + active_agent_task = asyncio.create_task( + agent.run_with_mcp( + message_to_send, **run_kwargs + ) + ) + + # ===== CONCURRENT MESSAGE PROCESSING FIX ===== + # Instead of awaiting the agent task here (which blocks the receive loop), + # we use asyncio.wait() to handle BOTH the agent task AND incoming messages. + # This allows permission_response messages to be processed while the agent runs. + + result = None + agent_completed = False + # Deferred message for switch_session/create_session during streaming + _deferred_msg: dict | None = None + + while not agent_completed: + # Create a task for the next message (or wait for agent) + receive_task = asyncio.create_task( + websocket.receive_json() + ) + + # Wait for either the agent to complete OR a new message to arrive + done, pending = await asyncio.wait( + {active_agent_task, receive_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + + # Check if agent completed + if active_agent_task in done: + try: + result = await active_agent_task + logger.debug( + f"run_with_mcp completed, result type: {type(result)}" + ) + agent_completed = True + active_agent_task = None + except asyncio.CancelledError: + logger.debug( + "run_with_mcp task was cancelled by user" + ) + 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, + ) + agent_error = e + result = None + agent_completed = True + active_agent_task = None + + # Cancel the receive task if it's still pending + if receive_task in pending: + receive_task.cancel() + try: + await receive_task + except asyncio.CancelledError: + pass + + # Check if a new message arrived + elif receive_task in done: + try: + new_msg = await receive_task + + # Advisory validation — log but never reject + try: + _parsed_inner = _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" + }, + ) + + # Handle permission_response immediately + if ( + 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 + ) + ) + if handled: + logger.debug( + f"[Permission] ✅ Handled response: {request_id} = {approved}" + ) + else: + logger.warning( + f"[Permission] ❌ Unknown request: {request_id}" + ) + else: + logger.error( + "[WebSocket] ❌ No request_id in permission_response!" + ) + + # Handle cancel request + 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 # Exit the loop + ) + + # Handle session switch during streaming + elif new_msg.get("type") in ( + "switch_session", + "create_session", + ): + # User switched to a new chat while agent was running. + # The agent will finish in background and save to SQLite. + logger.debug( + "[WS:%s] Session switch during streaming — " + "agent continues in background, switching to: %s", + session_id, + new_msg.get("session_id"), + ) + + # Fire and forget — background task owns the agent result. + 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 # disown — background task now owns it + agent_completed = ( + True # exit inner loop + ) + + # For other message types, log and ignore (or queue for later) + else: + logger.warning( + f"[WebSocket] Received {new_msg.get('type')} message while agent running - ignoring" + ) + + except asyncio.CancelledError: + # Receive was cancelled, continue + pass + except WebSocketDisconnect: + # WebSocket gone — let agent finish and save to SQLite + logger.debug( + "[WS:%s] Disconnect during streaming — agent continues in background", + session_id, + ) + + # Fire and forget — background task owns the agent result. + 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(): + # Starlette raises RuntimeError when calling receive after disconnect + logger.debug( + "[WS:%s] WebSocket already disconnected: %s — agent continues in background", + session_id, + e, + ) + + # Fire and forget — background task owns the agent result. + 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( + f"RuntimeError processing message during agent execution: {e}" + ) + except Exception as e: + logger.error( + f"Error processing message during agent execution: {e}" + ) + + # ===== END CONCURRENT MESSAGE PROCESSING FIX ===== + + finally: + # End suppression scope once run_with_mcp window ends, + # including completion/cancel/background handoff paths. + set_suppress_emitter_tool_events(False) + # Clear session context for prompt generation + clear_session_working_directory() + + finally: + # Stop the drain task + if drain_task: + stop_draining.set() + try: + await asyncio.wait_for(drain_task, timeout=2.0) + except asyncio.TimeoutError: + drain_task.cancel() + try: + await drain_task + except asyncio.CancelledError: + pass + + # Unsubscribe from emitter + if event_queue: + unsubscribe(event_queue) + logger.debug("Unsubscribed from frontend emitter") + + # 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 + + # If agent errored, send error to GUI and skip success path + if agent_error == "cancelled": + await send_typed( + ServerCancelled( + session_id=session_id, + ) + ) + continue + elif agent_error is not None: + logger.debug( + "[WS:%s] agent_error -> sending frame(s) to client. type=%s", + session_id, + type(agent_error).__name__, + ) + # If streaming was already in progress, send stream_end first so + # the frontend exits its streaming state before the error frame. + # Without this handshake the UI hangs indefinitely. + for frame in build_error_response_frames( + agent_error, collected_text, session_id + ): + await safe_send_json(frame) + continue + + # Safety net: if the agent task completed without raising but returned + # no result and produced no streamed text, treat it as an error. + has_nonempty_stream = any( + (chunk or "").strip() for chunk in 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(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." + ) + ) + await send_typed( + 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, + ) + ) + continue + + # Extract final response text + response_text = "" + + # Priority 1: Use collected text from streaming events + if has_nonempty_stream: + response_text = "".join(collected_text) + logger.debug( + f"Using collected streaming text ({len(response_text)} chars)" + ) + # Priority 2: Extract from result.output (AgentRunResult) or result.data (legacy) + 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)" + ) + # Priority 3: Get last assistant message from history + elif agent: + messages = agent.get_message_history() + + for msg in reversed(messages): + if hasattr(msg, "role") and msg.role == "assistant": + if 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" + + # Extract token usage from result if available + tokens_used = None + if result: + # Try to get usage from 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 + ), + } + # Also try _usage or other common patterns + 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 we couldn't get usage from result, estimate from message history + 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 + + # Only send legacy 'response' if B1 streaming wasn't used + # B1 streaming already sent the content via assistant_message_end + logger.warning( + "[WebSocket] b1_streaming_used=%s before response/extraction", + b1_streaming_used, + ) + if not b1_streaming_used: + # For non-streaming models (like Gemini), extract and send thinking content first + thinking_text = "" + if agent: + try: + history = agent.get_message_history() + logger.debug( + f"[Thinking Debug] Checking history for thinking parts, {len(history) if history else 0} messages" + ) + if history: + # Look for ThinkingPart in the last ModelResponse + for i, msg in enumerate(reversed(history)): + msg_type = type(msg).__name__ + logger.debug( + f"[Thinking Debug] Message {i}: type={msg_type}, has_parts={hasattr(msg, 'parts')}" + ) + if "Response" in msg_type and hasattr( + msg, "parts" + ): + logger.debug( + f"[Thinking Debug] Found Response with {len(msg.parts)} 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( + f"[Thinking Debug] Part {j}: type={part_type}, content_preview={part_content_preview}" + ) + if ( + "Thinking" in part_type + and hasattr(part, "content") + ): + thinking_text = part.content + logger.debug( + f"[Thinking Debug] Found thinking content: {len(thinking_text)} chars" + ) + 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" + ) + + # 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" + ) + + # Send the main response + await send_typed( + ServerResponse( + content=response_text, + done=True, + 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, + ) + ) + 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: set = set() + logger.warning( + "[WebSocket] B1 Pre-stream_end extraction: result=%s, has_all_messages=%s", + type(result).__name__ if result else None, + hasattr(result, "all_messages") + if result + else False, + ) + if result and hasattr(result, "all_messages"): + try: + import time as _time_pre + + from pydantic_ai.messages import ( + ToolReturn, + ToolReturnPart, + ) + + _pre_msgs = list(result.all_messages()) + for _pre_msg in _pre_msgs: + if not hasattr(_pre_msg, "parts"): + continue + for _pre_part in _pre_msg.parts: + if not isinstance( + _pre_part, + (ToolReturnPart, ToolReturn), + ): + continue + _pre_tool_name = getattr( + _pre_part, "tool_name", "unknown" + ) + _pre_raw_tool_id = getattr( + _pre_part, "tool_call_id", None + ) + _pre_tool_id = ( + tool_id_aliases.get( + _pre_raw_tool_id, + _pre_raw_tool_id, + ) + if _pre_raw_tool_id + else "unknown" + ) + _pre_result = getattr( + _pre_part, "content", None + ) + # Serialize complex objects + if _pre_result and not isinstance( + _pre_result, + ( + str, + dict, + list, + int, + float, + bool, + type(None), + ), + ): + try: + if hasattr( + _pre_result, "model_dump" + ): + _pre_result = ( + _pre_result.model_dump() + ) + elif hasattr( + _pre_result, "dict" + ): + _pre_result = ( + _pre_result.dict() + ) + elif hasattr( + _pre_result, "__dict__" + ): + _pre_result = ( + _pre_result.__dict__ + ) + else: + _pre_result = str( + _pre_result + ) + except Exception: + _pre_result = str(_pre_result) + logger.warning( + "[WebSocket] Pre-stream_end tool result: %s (id: %s), content_preview=%s", + _pre_tool_name, + _pre_tool_id, + str(_pre_result)[:100] + if _pre_result + else None, + ) + _pre_group_id = tool_group_ids.get( + _pre_tool_id + ) + + await send_typed_tool_lifecycle( + ServerToolResult( + tool_id=_pre_tool_id, + tool_name=_pre_tool_name, + result=_pre_result, + success=True, + duration_ms=0, + timestamp=_time_pre.time(), + session_id=session_id, + agent_name=agent.name + if agent + else "code-puppy", + model_name=agent.get_model_name() + if agent + else "unknown", + tool_group_id=_pre_group_id, + ) + ) + _pre_sent_tool_ids.add(_pre_tool_id) + except Exception as _pre_e: + logger.warning( + "Pre-stream_end tool result extraction failed: %s", + _pre_e, + ) + + # 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 + # Track tool IDs already sent before stream_end to skip duplicates later + if "_pre_sent_tool_ids" not in dir(): + _pre_sent_tool_ids: set = set() + _save_history_snapshot = [] # Snapshot of history before await points + try: + # CRITICAL: Update message history from result to include final response + # The pydantic-ai result.all_messages() contains the complete conversation + # including the final assistant response that may not be in agent's history yet + if result and hasattr(result, "all_messages"): + try: + all_msgs = list(result.all_messages()) + if all_msgs: + agent.set_message_history(all_msgs) + _save_history_snapshot = list( + agent.get_message_history() + ) # Snapshot before any awaits can corrupt shared global state + logger.debug( + f"Updated message history from result.all_messages(): {len(all_msgs)} messages" + ) + + # Extract and send tool results from message history + try: + import time as _time + + from pydantic_ai.messages import ( + ToolReturn, + ToolReturnPart, + ) + + for msg in all_msgs: + if hasattr(msg, "parts"): + for part in msg.parts: + if isinstance( + part, + ( + ToolReturnPart, + ToolReturn, + ), + ): + tool_name = getattr( + part, + "tool_name", + "unknown", + ) + tool_call_id = getattr( + part, + "tool_call_id", + "unknown", + ) + # Serialize the result + result_data = getattr( + part, "content", None + ) + if ( + result_data + and not isinstance( + result_data, + ( + str, + dict, + list, + int, + float, + bool, + type(None), + ), + ) + ): + try: + if hasattr( + result_data, + "model_dump", + ): + result_data = result_data.model_dump() + elif hasattr( + result_data, + "dict", + ): + result_data = result_data.dict() + elif hasattr( + result_data, + "__dict__", + ): + result_data = result_data.__dict__ + else: + result_data = str( + result_data + ) + except Exception: + result_data = str( + result_data + ) + + # Log for shell commands + if ( + tool_name + == "agent_run_shell_command" + ): + stdout_val = ( + result_data.get( + "stdout", "N/A" + ) + if isinstance( + result_data, + dict, + ) + else "not dict" + ) + logger.debug( + "Extracted shell result: id=%s, stdout=%s", + tool_call_id, + stdout_val, + ) + + # Skip tool IDs already sent before stream_end + if ( + tool_call_id + in _pre_sent_tool_ids + ): + logger.debug( + "[WebSocket] Skipping duplicate tool result (pre-sent): %s", + tool_call_id, + ) + continue + logger.info( + f"[WebSocket] Sending extracted tool result for {tool_name} (id: {tool_call_id})" + ) + _post_group_id = ( + tool_group_ids.get( + tool_call_id + ) + ) + + await send_typed( + ServerToolResult( + tool_id=tool_call_id, + tool_name=tool_name, + result=result_data, + success=True, + duration_ms=0, + timestamp=_time.time(), + session_id=session_id, + agent_name=agent.name + if agent + else "code-puppy", + model_name=agent.get_model_name() + if agent + else "unknown", + tool_group_id=_post_group_id, + ) + ) + except Exception as e: + logger.warning( + f"Could not extract tool results from messages: {e}" + ) + + except Exception as e: + logger.warning( + f"Could not update history from result.all_messages(): {e}" + ) + + history = ( + _save_history_snapshot + if _save_history_snapshot + else agent.get_message_history() + ) # Use pre-await snapshot to avoid race condition + if history: + # Regenerate title if it's empty or still the default "untitled-session" + if ( + not session_title + or session_title == "untitled-session" + ): + session_title = generate_heuristic_title( + history + ) + + # Session name is just the ID (WS_session_timestamp format) + # Title is stored only in the meta file, not in the filename + session_name = session_id + + # Wrap each message with metadata for complete session information. + # Chain `or` fallbacks: agent.name can be "" before the agent + # object is fully configured, so we cascade to context defaults. + agent_name_meta = ( + (agent.name if agent else "") + or ctx.agent_name + or "code-puppy" + ) + model_name_meta = ( + (agent.get_model_name() if agent else "") + or ctx.model_name + or "unknown" + ) + + def _extract_message_timestamp( + raw_msg: Any, default_ts: str + ) -> str: + """Best-effort extraction of an existing timestamp for a message. + + This is important for WebSocket sessions where history may + already contain older messages. We don't want to overwrite + their original timestamps every time we auto-save. + + Precedence: + 1. If the message is a dict with a numeric epoch 'timestamp', + convert to ISO. + 2. If the message is a dict with an ISO-ish 'timestamp' str, + reuse it as-is. + 3. If the message is a dict with 'ts', reuse it. + 4. If the message object has a 'timestamp' attribute, try that. + 5. Fall back to the provided default_ts ("now" for new messages). + """ + # Dict-based histories (e.g. CLI format or older WS formats) + if isinstance(raw_msg, dict): + ts_val = raw_msg.get("timestamp") + + # Epoch seconds + if isinstance(ts_val, (int, float)): + try: + return ( + datetime.datetime.fromtimestamp( + ts_val + ).isoformat() + ) + except Exception: + pass + + # Already an ISO string + if isinstance(ts_val, str) and ts_val: + return ts_val + + # Our enhanced WS wrapper sometimes uses 'ts' + ts_field = raw_msg.get("ts") + if isinstance(ts_field, str) and ts_field: + return ts_field + + # Pydantic / custom objects that carry a timestamp attribute + 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 + + # Fallback: use provided default + return default_ts + + # Create enhanced history: list of dicts with message + metadata + # Each entry: {'msg': , 'agent': str, 'model': str, 'ts': str} + enhanced_history = [] + for idx, msg in enumerate(history): + # Check if already wrapped (for idempotency). If the wrapper + # already has a 'ts', leave it untouched so older sessions + # keep their original per-message timestamps. + if ( + isinstance(msg, dict) + and "msg" in msg + and "agent" in msg + ): + enhanced_history.append(msg) + else: + # Use existing timestamp if present; otherwise compute a default now() + current_timestamp = ( + datetime.datetime.now().isoformat() + ) + msg_ts = _extract_message_timestamp( + msg, current_timestamp + ) + wrapper = { + "msg": msg, + "agent": agent_name_meta, + "model": model_name_meta, + "ts": msg_ts, + } + + # Add clean_content and attachments to the user message we just processed. + # clean_content is only needed when attachments were injected into content + # (file blocks like --- File: auth.ts ---). Without attachments, content + # is already the user's words (FE strips [Session Context:] as fallback). + is_user_message_just_processed = ( + idx == len(history) - 2 + and len(history) >= 2 + and attachment_metadata # only needed when file content was injected + ) + + if is_user_message_just_processed: + wrapper["clean_content"] = ( + original_user_message + ) + wrapper["attachments"] = ( + attachment_metadata + ) + logger.debug( + "Added UI metadata to user message: %d attachment(s), " + "clean_content length: %d", + len(attachment_metadata), + len(original_user_message), + ) + + enhanced_history.append(wrapper) + + message_count = len(enhanced_history) + 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 + + # Write to SQLite for FE read path + try: + import datetime as _dt_mod + + _now_iso = _dt_mod.datetime.now( + _dt_mod.timezone.utc + ).isoformat() + await write_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, + updated_at=_now_iso, + created_at=ctx.created_at.isoformat(), + ctx=ctx, + ) + except Exception as _db_exc: + logger.debug( + "SQLite turn write skipped (DB not available): %s", + _db_exc, + ) + + # Send session metadata update to client + await websocket.send_json( + { + "type": "session_meta", + "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, + } + ) + + # Broadcast session update to session monitoring clients + session_update_data = { + "session_id": session_id, + "session_name": session_name, + "title": session_title, + "working_directory": session_working_directory, + "timestamp": datetime.datetime.now().isoformat(), + "message_count": message_count, + "total_tokens": total_tokens, + "auto_saved": True, + "pickle_path": "", + "metadata_path": "", + "action": "created" + if message_count == 1 + else "updated", + } + await connection_manager.broadcast_session_update( + session_update_data + ) + 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) + + 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 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: + 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) + except Exception: + pass From c32a8cfd7b38e3a4e890f244920a431920b05c3f Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Thu, 28 May 2026 18:18:45 -0500 Subject: [PATCH 03/29] docs(migration): track ws/model-switch regressions and fixes --- docs/puppy_desk_migration_issues.md | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 docs/puppy_desk_migration_issues.md diff --git a/docs/puppy_desk_migration_issues.md b/docs/puppy_desk_migration_issues.md new file mode 100644 index 000000000..1ccf1a9c2 --- /dev/null +++ b/docs/puppy_desk_migration_issues.md @@ -0,0 +1,21 @@ +# Puppy Desk Migration - Active Issue List + +## Open / Tracked + +1. **Model switch failure** + `Failed to switch to model gpt-5.5: CodePuppyAgent object has no attribute set_session_model` + - Status: **Fixed in commit `772e52ed`** + +2. **Post-CWD-change chat failure** + `Agent run failed (no result returned)` after `set_working_directory` + - Root cause: raw dict injection into agent message history + - Status: **Fix committed in `5c205a85`** + +3. **Config schema endpoint crash** + `ImportError: cannot import name 'CONFIG_SCHEMA' from code_puppy.config` on `/config/schema` + - Status: **Fix committed in `5c205a85`** + +## Residual Migration Gaps + +- `code_puppy/api/` and related ws/runtime files are still largely untracked in this branch and need staged migration/normalization. +- Additional end-to-end GUI validation is pending after full API migration commit set. From 4e3fc17a76acde44004ad84b96d7dc22546eef43 Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Thu, 28 May 2026 18:38:17 -0500 Subject: [PATCH 04/29] feat(api): migrate core api, db, routers, and session context --- code_puppy/api/__init__.py | 13 + code_puppy/api/app.py | 329 +++++++ code_puppy/api/compaction_tracker.py | 130 +++ code_puppy/api/db/__init__.py | 49 + code_puppy/api/db/backfill.py | 381 ++++++++ code_puppy/api/db/connection.py | 426 +++++++++ code_puppy/api/db/message_utils.py | 219 +++++ code_puppy/api/db/queries.py | 1269 +++++++++++++++++++++++++ code_puppy/api/db/seeder.py | 555 +++++++++++ code_puppy/api/error_parser.py | 239 +++++ code_puppy/api/main.py | 97 ++ code_puppy/api/permission_plugin.py | 216 +++++ code_puppy/api/permissions.py | 172 ++++ code_puppy/api/pty_manager.py | 453 +++++++++ code_puppy/api/routers/__init__.py | 12 + code_puppy/api/routers/agents.py | 116 +++ code_puppy/api/routers/commands.py | 217 +++++ code_puppy/api/routers/models.py | 23 + code_puppy/api/routers/protocol.py | 26 + code_puppy/api/routers/sessions.py | 274 ++++++ code_puppy/api/routers/ws_sessions.py | 302 ++++++ code_puppy/api/session_cache.py | 392 ++++++++ code_puppy/api/session_context.py | 675 +++++++++++++ code_puppy/api/tool_formatters.py | 138 +++ code_puppy/api/websocket.py | 51 + 25 files changed, 6774 insertions(+) create mode 100644 code_puppy/api/__init__.py create mode 100644 code_puppy/api/app.py create mode 100644 code_puppy/api/compaction_tracker.py create mode 100644 code_puppy/api/db/__init__.py create mode 100644 code_puppy/api/db/backfill.py create mode 100644 code_puppy/api/db/connection.py create mode 100644 code_puppy/api/db/message_utils.py create mode 100644 code_puppy/api/db/queries.py create mode 100644 code_puppy/api/db/seeder.py create mode 100644 code_puppy/api/error_parser.py create mode 100644 code_puppy/api/main.py create mode 100644 code_puppy/api/permission_plugin.py create mode 100644 code_puppy/api/permissions.py create mode 100644 code_puppy/api/pty_manager.py create mode 100644 code_puppy/api/routers/__init__.py create mode 100644 code_puppy/api/routers/agents.py create mode 100644 code_puppy/api/routers/commands.py create mode 100644 code_puppy/api/routers/models.py create mode 100644 code_puppy/api/routers/protocol.py create mode 100644 code_puppy/api/routers/sessions.py create mode 100644 code_puppy/api/routers/ws_sessions.py create mode 100644 code_puppy/api/session_cache.py create mode 100644 code_puppy/api/session_context.py create mode 100644 code_puppy/api/tool_formatters.py create mode 100644 code_puppy/api/websocket.py 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..a68b86c14 --- /dev/null +++ b/code_puppy/api/app.py @@ -0,0 +1,329 @@ +"""FastAPI application factory for Code Puppy API.""" + +import asyncio +import logging +import os +import time +from contextlib import asynccontextmanager +from pathlib import Path +from typing import AsyncGenerator + +from dbos import DBOS, DBOSConfig +from fastapi import FastAPI, Request +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import FileResponse, HTMLResponse, JSONResponse +from starlette.middleware.base import BaseHTTPMiddleware + +from code_puppy.config import DBOS_DATABASE_URL, get_use_dbos + +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['builtin']}, user={result['user']}, external={result['external']}" + ) + + # Initialise shared SQLite database + try: + from code_puppy.api.db.connection import init_db + from code_puppy.api.db.seeder import seed_from_pkl_dirs + + await init_db() + logger.info("✓ aiosqlite DB initialised") + + # Seed existing pkl sessions in the background — non-blocking + asyncio.create_task(seed_from_pkl_dirs()) + logger.info("✓ SQLite seeder task started") + except Exception as _db_exc: + logger.error("SQLite DB init failed (continuing without DB): %s", _db_exc) + + # Initialize DBOS if enabled (mirrors cli_runner.py; API server logs errors and continues degraded rather than exiting) + _dbos_initialized = [False] # use list to allow mutation across yield boundary + if get_use_dbos(): + from code_puppy import __version__ as current_version + + dbos_app_version = os.environ.get( + "DBOS_APP_VERSION", f"{current_version}-{int(time.time() * 1000)}" + ) + dbos_config: DBOSConfig = { + "name": "dbos-code-puppy", + "system_database_url": DBOS_DATABASE_URL, + "run_admin_server": False, + "conductor_key": os.environ.get("DBOS_CONDUCTOR_KEY"), + "log_level": os.environ.get("DBOS_LOG_LEVEL", "ERROR"), + "application_version": dbos_app_version, + } + try: + DBOS(config=dbos_config) + DBOS.launch() + _dbos_initialized[0] = True # only set after both calls succeed + logger.info("✓ DBOS initialized") + except Exception as e: + logger.error("Error initializing DBOS: %s", e) + + yield + # Shutdown: clean up all the things! + logger.info("🐶 Code Puppy API shutting down, cleaning up...") + + # 1. Close all PTY sessions + try: + from code_puppy.api.pty_manager import get_pty_manager + + pty_manager = get_pty_manager() + await pty_manager.close_all() + logger.info("✓ All PTY sessions closed") + except Exception as e: + logger.error("Error closing PTY sessions: %s", e) + + # 2. 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) + + # 3. 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) + + # Destroy DBOS if it was fully initialized + if _dbos_initialized[0]: + try: + DBOS.destroy() + logger.info("✓ DBOS destroyed") + except Exception as e: + logger.error("Error destroying DBOS: %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 (events + terminal) + 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 Terminal + + + Open Chat + + + View Sessions + + + API Docs + +
+

+ WebSocket: ws://localhost:8765/ws/terminal +

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

Terminal template not found

", + status_code=404, + ) + + @app.get("/sessions") + async def sessions_page(): + """Serve the sessions monitoring page.""" + html_file = templates_dir / "sessions.html" + if html_file.exists(): + return FileResponse(html_file, media_type="text/html") + return HTMLResponse( + content="

Sessions template not found

", + status_code=404, + ) + + @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.plugins.walmart_specific.auto_update import fetch_latest_version + from code_puppy.version_checker import versions_are_equal + + current_version = __version__ + latest_version = fetch_latest_version() + + # 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..f7805b048 --- /dev/null +++ b/code_puppy/api/db/__init__.py @@ -0,0 +1,49 @@ +"""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, +) +from code_puppy.api.db.seeder import seed_from_pkl_dirs + +__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", + "seed_from_pkl_dirs", + "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..5a22d69b7 --- /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 seeder.py (pkl import) and session_context.py (live WS saves). +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 (e.g. very old pkl message types). + """ + 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..b5adc75e2 --- /dev/null +++ b/code_puppy/api/db/queries.py @@ -0,0 +1,1269 @@ +"""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 + +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 = """ +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; +""" + + +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: + # Modify the query to exclude compacted messages + sql = """ +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; +""" + cursor = await db.execute(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..0c5a5b2d3 --- /dev/null +++ b/code_puppy/api/db/seeder.py @@ -0,0 +1,555 @@ +"""Seed ~/.puppy_desk/chat_messages.db from existing pkl session files. + +Scans both: + ~/.code_puppy/ws_sessions/ — live WS sessions (HMAC-signed pickle) + ~/.code_puppy/autosaves/ — CLI autosave sessions (raw or legacy-signed pickle) + +For each .pkl file not already in the DB: + 1. Deserialises the pickle natively (no subprocess, no JSON round-trip) + 2. Serialises each ModelMessage with ModelMessagesTypeAdapter.dump_json() + → stored verbatim in the pydantic_json column for perfect replay + 3. Extracts display fields (role, content, thinking, tool calls) from message parts + 4. Writes session + messages + tool_calls in a single transaction + +Idempotent: existing sessions are skipped. Safe to call on every startup. +""" + +from __future__ import annotations + +import hashlib +import hmac +import json +import logging +import pickle +import re +import time +import uuid +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Optional + +from code_puppy.api.db.message_utils import ( + extract_content, + extract_thinking, + get_message_timestamp, + get_role, + pydantic_json_for_message, +) + +logger = logging.getLogger(__name__) +# --------------------------------------------------------------------------- +# One-time migration marker — written after a successful complete seed. +# On subsequent startups the seeder skips entirely (no more pkl to import). +# --------------------------------------------------------------------------- +_SEEDER_DONE_MARKER = Path("~/.puppy_desk/.seeder_migration_done").expanduser() + + +# --------------------------------------------------------------------------- +# Legacy pickle header constants (from session_storage.py) +# --------------------------------------------------------------------------- + +_LEGACY_SIGNED_HEADER = b"CPSESSION\x01" +_LEGACY_SIGNATURE_SIZE = 32 + + +def _get_hmac_key() -> bytes: + """Return HMAC key for verifying WS session pickle files. + + Copied from session_context to make seeder.py self-contained after + _get_hmac_key was removed from session_context in desk-puppy-cr1x. + """ + import os + + key_env = os.environ.get("CODE_PUPPY_SESSION_KEY", "") + if key_env: + return key_env.encode() + # Fall back to a machine-stable key derived from the DB path + db_path = str(Path("~/.puppy_desk/chat_messages.db").expanduser()) + return db_path.encode() + + +def _extract_pickle_payload(raw: bytes) -> bytes: + """Strip legacy or new-format prefix and return the raw pickle bytes.""" + if raw.startswith(_LEGACY_SIGNED_HEADER): + offset = len(_LEGACY_SIGNED_HEADER) + _LEGACY_SIGNATURE_SIZE + return raw[offset:] + return raw + + +def _verify_and_load_ws_pickle(raw: bytes, hmac_key: bytes) -> Optional[list[Any]]: + """Load a WS session pickle after HMAC verification. + + WS session pkls are written by session_context.py: + file = 32-byte-HMAC-SHA256 + pickle(history) + + Returns None if verification fails or the file is too small. + """ + if len(raw) < 33: # 32-byte sig + at least 1 byte of pickle + return None + sig, blob = raw[:32], raw[32:] + expected = hmac.new(hmac_key, blob, hashlib.sha256).digest() + if not hmac.compare_digest(sig, expected): + logger.warning("HMAC verification failed — skipping") + return None + try: + data = pickle.loads(blob) # noqa: S301 + return data if isinstance(data, list) else None + except Exception as exc: + logger.warning("Pickle load error: %s", exc) + return None + + +def _load_autosave_pickle(raw: bytes) -> Optional[list[Any]]: + """Load an autosave pickle (no HMAC — strip legacy header if present).""" + payload = _extract_pickle_payload(raw) + try: + data = pickle.loads(payload) # noqa: S301 + return data if isinstance(data, list) else None + except Exception as exc: + logger.warning("Autosave pickle load error: %s", exc) + return None + + +# --------------------------------------------------------------------------- +# Message part extraction helpers (content/thinking/role moved to message_utils) +# --------------------------------------------------------------------------- + +_PART_TYPE_TOOL_CALL = {"ToolCallPart"} +_PART_TYPE_TOOL_RETURN = {"ToolReturnPart"} + + +def _extract_tool_calls(msg: Any) -> list[dict[str, Any]]: + """Extract ToolCallPart entries from a ModelResponse.""" + 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 _extract_tool_returns(msg: Any) -> list[dict[str, Any]]: + """Extract ToolReturnPart entries from a ModelRequest.""" + parts = getattr(msg, "parts", []) + result = [] + for part in parts: + if type(part).__name__ not in _PART_TYPE_TOOL_RETURN: + continue + content = getattr(part, "content", "") + try: + result_val = json.loads(content) if isinstance(content, str) else content + except Exception: + result_val = content + result.append( + { + "id": getattr(part, "tool_call_id", ""), + "name": getattr(part, "tool_name", "unknown"), + "result": result_val, + } + ) + return result + + +def _get_timestamp(msg: Any, wrapper: Optional[dict], fallback: str) -> str: + """Best-effort timestamp extraction.""" + # Try wrapper ts first + if wrapper: + ts = wrapper.get("ts", "") + if ts: + return str(ts) + # Try message-level timestamp + ts_attr_result = get_message_timestamp(msg) + if ts_attr_result is not None: + return ts_attr_result + return fallback + + +# --------------------------------------------------------------------------- +# Per-session import +# --------------------------------------------------------------------------- + +_SESSION_CONTEXT_RE = re.compile(r"^\[Session Context:[^\]]*\]\n\n?") + + +def _strip_session_context(raw: str) -> Optional[str]: + cleaned = _SESSION_CONTEXT_RE.sub("", raw).lstrip() + return cleaned if cleaned != raw else None + + +async def _import_session( + session_id: str, + history: list[Any], + meta: dict[str, Any], + now_iso: str, +) -> int: + """Write one session and all its messages/tool_calls to the DB. + + Returns the number of message rows written. + """ + from code_puppy.api.db.queries import ( + insert_messages_batch, + insert_tool_calls_batch, + upsert_session, + ) + + created_at = meta.get("timestamp") or now_iso + title = meta.get("title", "") + agent_name = meta.get("agent_name", "code-puppy") + model_name = meta.get("model_name", "") + working_directory = meta.get("working_directory", "") + pinned = bool(meta.get("pinned", False)) + total_tokens = meta.get("total_tokens", 0) + + 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=created_at, + message_count=len(history), + total_tokens=total_tokens, + deleted_at=None, + ) + + message_rows: list[dict[str, Any]] = [] + tool_call_rows: list[dict[str, Any]] = [] + seq = 0 + # Track pending tool calls by id for result matching + pending_tool_calls: dict[str, dict[str, Any]] = {} + + for item in history: + # ---- Unwrap the item ---------------------------------------- + wrapper: Optional[dict] = None + msg: Any + + if isinstance(item, dict) and item.get("msg") == "system": + # System message in WS format: {'msg': 'system', 'content': ..., 'path': ...} + seq += 1 + message_rows.append( + { + "session_id": session_id, + "seq": seq, + "role": "system", + "content": item.get("content", ""), + "type": "system", + "agent_name": item.get("agent", agent_name), + "model_name": item.get("model", model_name), + "timestamp": item.get("ts", now_iso), + "system_message_type": item.get("system_message_type", ""), + "system_message_path": item.get("path", ""), + "pydantic_json": None, + "token_count": 0, + } + ) + continue + + if isinstance(item, dict) and "msg" in item: + wrapper = item + msg = item["msg"] + else: + msg = item + + # ---- Skip if not a pydantic-ai message ---------------------- + if not hasattr(msg, "parts"): + continue + + role = get_role(msg) + ts = _get_timestamp(msg, wrapper, now_iso) + content = extract_content(msg) + thinking = extract_thinking(msg) + pydantic_json = pydantic_json_for_message(msg) + + msg_agent = wrapper.get("agent", agent_name) if wrapper else agent_name + msg_model = wrapper.get("model", model_name) if wrapper else model_name + + # clean_content: strip [Session Context: ...] for user messages + clean_content: Optional[str] = None + raw_clean = wrapper.get("clean_content") if wrapper else None + if raw_clean: + clean_content = raw_clean + elif role == "user" and content: + clean_content = _strip_session_context(content) + + # attachments_json + attachments = wrapper.get("attachments") if wrapper else None + attachments_json = json.dumps(attachments) if attachments else None + + seq += 1 + message_seq = seq + + # ---- Estimate token count (simple char/4 heuristic) --------- + token_count = max(1, len(content) // 4) + + message_rows.append( + { + "session_id": session_id, + "seq": message_seq, + "role": role, + "content": content, + "type": type(msg).__name__, + "agent_name": msg_agent, + "model_name": msg_model, + "timestamp": ts, + "thinking": thinking, + "attachments_json": attachments_json, + "clean_content": clean_content, + "pydantic_json": pydantic_json, + "token_count": token_count, + } + ) + + # ---- Tool calls from assistant messages --------------------- + if role == "assistant": + for tc in _extract_tool_calls(msg): + pending_tool_calls[tc["id"]] = { + "tc_id": tc["id"], + "name": tc["name"], + "args": tc["args"], + "parent_seq": message_seq, + "agent": msg_agent, + "model": msg_model, + "ts": ts, + } + + # ---- Tool returns from user messages (Anthropic API format) - + if role == "user": + for tr in _extract_tool_returns(msg): + tid = tr["id"] + if tid in pending_tool_calls: + tc = pending_tool_calls.pop(tid) + seq += 1 + try: + args_json = json.dumps(tc["args"]) + except Exception: + args_json = str(tc["args"]) + # Serialize result - handle Pydantic models and complex objects + result_val = tr["result"] + try: + if hasattr(result_val, "model_dump"): + result_json = json.dumps(result_val.model_dump()) + elif hasattr(result_val, "dict"): + result_json = json.dumps(result_val.dict()) + elif hasattr(result_val, "__dict__"): + result_json = json.dumps(vars(result_val)) + else: + result_json = json.dumps(result_val) + except Exception: + try: + result_json = json.dumps(str(result_val)) + except Exception: + result_json = json.dumps({"raw": str(result_val)}) + + tool_ts_str = tc.get("ts", ts) + try: + tool_ts = datetime.fromisoformat( + str(tool_ts_str).replace("Z", "+00:00") + ).timestamp() + except Exception: + tool_ts = time.time() + + tool_call_rows.append( + { + "id": tid, + "session_id": session_id, + "parent_message_seq": tc["parent_seq"], + "seq": seq, + "tool_name": tc["name"], + "args_json": args_json, + "result_json": result_json, + "status": "success", + "agent_name": tc["agent"], + "model_name": tc["model"], + "timestamp": tool_ts, + } + ) + + # Write everything in batches (each batch uses its own transaction) + await insert_messages_batch(message_rows) + await insert_tool_calls_batch(tool_call_rows) + return len(message_rows) + + +# --------------------------------------------------------------------------- +# Directory scanners +# --------------------------------------------------------------------------- + + +def _load_meta(meta_path: Path) -> dict[str, Any]: + try: + if meta_path.exists(): + return json.loads(meta_path.read_text(encoding="utf-8")) + except Exception: + pass + return {} + + +async def _seed_directory( + sessions_dir: Path, + is_ws: bool, + hmac_key: Optional[bytes], + now_iso: str, +) -> tuple[int, int, int]: + """Seed all pkl files from *sessions_dir*. Returns (imported, skipped, failed).""" + from code_puppy.api.db.queries import session_exists + + if not sessions_dir.exists(): + logger.debug("Sessions dir not found, skipping: %s", sessions_dir) + return 0, 0, 0 + + pkl_files = sorted(sessions_dir.glob("*.pkl")) + total = len(pkl_files) + if total == 0: + return 0, 0, 0 + + imported = skipped = failed = 0 + width = len(str(total)) + + for i, pkl_path in enumerate(pkl_files, 1): + session_id = pkl_path.stem + + # Skip tiny files (HMAC-only stubs, empty autosaves) + try: + if pkl_path.stat().st_size < 50: + skipped += 1 + continue + except OSError: + skipped += 1 + continue + + if await session_exists(session_id): + skipped += 1 + continue + + try: + raw = pkl_path.read_bytes() + except OSError as exc: + logger.warning( + "[%*d/%d] %s — read error: %s", width, i, total, session_id, exc + ) + failed += 1 + continue + + # Deserialise + if is_ws and hmac_key is not None: + history = _verify_and_load_ws_pickle(raw, hmac_key) + else: + history = _load_autosave_pickle(raw) + + if not history: + skipped += 1 + continue + + # Load metadata + meta_path = pkl_path.with_name(f"{session_id}_meta.json") + meta = _load_meta(meta_path) + + try: + n = await _import_session(session_id, history, meta, now_iso) + logger.info( + "[%*d/%d] %-50s %4d msgs — imported", + width, + i, + total, + session_id, + n, + ) + imported += 1 + except Exception as exc: + logger.error( + "[%*d/%d] %s — import error: %s", + width, + i, + total, + session_id, + exc, + exc_info=True, + ) + failed += 1 + + return imported, skipped, failed + + +# --------------------------------------------------------------------------- +# Public entry point +# --------------------------------------------------------------------------- + + +async def seed_from_pkl_dirs() -> None: + """Asynchronously seed the DB from both ws_sessions and autosaves directories. + + All database operations go through the shared aiosqlite connection so this + must run on the event loop that owns the DB connection. File I/O for + reading pickle files is blocking but acceptable at startup since the seeder + runs as a background asyncio.create_task(). + + Safe to call on every startup — already-imported sessions are skipped. + """ + # One-time migration guard: if we already migrated all pkl files, skip. + if _SEEDER_DONE_MARKER.exists(): + logger.debug("Seeder already completed (marker found) — skipping pkl scan") + return + + from code_puppy.config import AUTOSAVE_DIR, WS_SESSION_DIR + + now_iso = datetime.now(timezone.utc).isoformat() + + # ---- WS sessions (HMAC-signed) ----------------------------------- + ws_dir = Path(WS_SESSION_DIR) + hmac_key: Optional[bytes] = None + try: + hmac_key = _get_hmac_key() + except Exception as exc: + logger.warning("Could not load HMAC key, WS sessions will be skipped: %s", exc) + + if hmac_key is not None: + logger.info("Seeding WS sessions from %s ...", ws_dir) + wi, ws, wf = await _seed_directory( + ws_dir, is_ws=True, hmac_key=hmac_key, now_iso=now_iso + ) + logger.info("WS sessions: %d imported, %d skipped, %d failed", wi, ws, wf) + else: + wi = ws = wf = 0 + + # ---- Autosaves (no HMAC) ----------------------------------------- + auto_dir = Path(AUTOSAVE_DIR) + logger.info("Seeding autosave sessions from %s ...", auto_dir) + ai, as_, af = await _seed_directory( + auto_dir, is_ws=False, hmac_key=None, now_iso=now_iso + ) + logger.info("Autosave sessions: %d imported, %d skipped, %d failed", ai, as_, af) + + total_imported = wi + ai + total_skipped = ws + as_ + total_failed = wf + af + logger.info( + "\u2713 Seeding complete: %d imported, %d skipped, %d failed", + total_imported, + total_skipped, + total_failed, + ) + + # Write the one-time migration marker if seed succeeded with no failures. + if total_failed == 0: + try: + _SEEDER_DONE_MARKER.parent.mkdir(parents=True, exist_ok=True) + _SEEDER_DONE_MARKER.write_text( + f"Seeder migration completed. imported={total_imported} skipped={total_skipped}\n" + ) + logger.info("✓ Seeder migration marker written to %s", _SEEDER_DONE_MARKER) + except Exception as exc: + logger.warning("Could not write seeder marker: %s", exc) 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..411a4b799 --- /dev/null +++ b/code_puppy/api/permission_plugin.py @@ -0,0 +1,216 @@ +""" +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) + # On error, allow the tool (fail open for now) + return None + + +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) + # On error, allow the command (fail open for now) + return None + + +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..c820693dc --- /dev/null +++ b/code_puppy/api/permissions.py @@ -0,0 +1,172 @@ +""" +Permission System for Tool Call Execution + +This module provides a permission request system that integrates with the WebSocket API +to request user approval before executing potentially dangerous operations. + +Features: +- WebSocket-based permission requests with in-chat action buttons +- 5-minute timeout on permission waits +- Automatic approval in yolo mode +- CLI mode support via WebSocket + +Usage: + from code_puppy.api.permissions import request_permission + + approved = await request_permission( + websocket=ws, + session_id="session-123", + request_type="shell_command", + title="Execute Shell Command", + description="Run: ls -la", + details={"command": "ls -la", "cwd": "/home/user"} + ) +""" + +import asyncio +import logging +import uuid +from typing import Any, Dict, Optional + +from fastapi import WebSocket + +from code_puppy.api.ws.schemas import ServerPermissionRequest + +logger = logging.getLogger(__name__) + +# Global dictionary to track pending permission requests +# This is managed by the WebSocket API module +permission_futures: Dict[str, asyncio.Future] = {} + + +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. + + Args: + websocket: WebSocket connection to send request through + session_id: Current session ID + request_type: Type of permission (e.g., 'shell_command') + title: Permission dialog title + description: Human-readable description + details: Additional context (command, args, etc.) + timeout: Timeout in seconds (default 5 minutes) + + Returns: + bool: True if approved, False if denied or timeout + """ + # Check if yolo mode is enabled (auto-approve everything) + try: + import json + import os + + from code_puppy.config import CONFIG_FILE + + if os.path.exists(CONFIG_FILE): + try: + with open(CONFIG_FILE, "r") as f: + config = json.load(f) + yolo_mode = str(config.get("yolo_mode", "false")).lower() == "true" + if yolo_mode: + logger.info( + f"[Permission] YOLO mode enabled, auto-approving {request_type}" + ) + return True + except Exception: + pass # If config read fails, continue to request permission + except Exception: + pass # If we can't check, require permission + + # If no WebSocket connection, deny by default (fail-safe) + if websocket is None: + logger.warning("[Permission] No WebSocket connection, denying %s", request_type) + return False + + request_id = str(uuid.uuid4()) + + # Send permission request to frontend + 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: + logger.error("[Permission] Failed to send request: %s", e) + return False + + # Create a future to wait for the response + future = asyncio.Future() + permission_futures[request_id] = future + + try: + # Wait for user response with timeout + approved = await asyncio.wait_for(future, timeout=timeout) + logger.info( + f"[Permission] ✅ Got response for {request_id}: {'APPROVED' if approved else 'DENIED'}" + ) + return 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: + # Clean up the future + permission_futures.pop(request_id, None) + + +def handle_permission_response(request_id: str, approved: bool) -> bool: + """ + Handle a permission response from the client. + + Args: + request_id: The request ID from the permission_request + approved: True if approved, False if denied + + Returns: + bool: True if the response was handled, False if request_id not found + """ + logger.info( + f"[Permission] handle_permission_response called: request_id={request_id}, approved={approved}" + ) + logger.info( + f"[Permission] Current permission_futures keys: {list(permission_futures.keys())}" + ) + + future = permission_futures.get(request_id) + logger.info( + f"[Permission] Future found: {future is not None}, done: {future.done() if future else 'N/A'}" + ) + + if future and not future.done(): + logger.info("[Permission] Setting future result to: %s", approved) + future.set_result(approved) + logger.info("[Permission] Handled response for %s: %s", request_id, approved) + return True + else: + logger.warning( + f"[Permission] Received response for unknown/expired request: {request_id}" + ) + logger.warning( + f"[Permission] Future exists: {future is not None}, Future done: {future.done() if future else 'N/A'}" + ) + return False diff --git a/code_puppy/api/pty_manager.py b/code_puppy/api/pty_manager.py new file mode 100644 index 000000000..d57dacdc8 --- /dev/null +++ b/code_puppy/api/pty_manager.py @@ -0,0 +1,453 @@ +"""PTY Manager for terminal emulation with cross-platform support. + +Provides pseudo-terminal (PTY) functionality for interactive shell sessions +via WebSocket connections. Supports Unix (pty module) and Windows (pywinpty). +""" + +import asyncio +import logging +import os +import signal +import struct +import sys +from dataclasses import dataclass, field +from typing import Any, Callable, Optional + +logger = logging.getLogger(__name__) + +# Platform detection +IS_WINDOWS = sys.platform == "win32" + +# Conditional imports based on platform +if IS_WINDOWS: + try: + import winpty # type: ignore + + HAS_WINPTY = True + except ImportError: + HAS_WINPTY = False + winpty = None +else: + import fcntl + import pty + import termios + + HAS_WINPTY = False + + +@dataclass +class PTYSession: + """Represents an active PTY session.""" + + session_id: str + master_fd: Optional[int] = None # Unix only + slave_fd: Optional[int] = None # Unix only + pid: Optional[int] = None # Unix only + winpty_process: Any = None # Windows only + cols: int = 80 + rows: int = 24 + on_output: Optional[Callable[[bytes], None]] = None + _reader_task: Optional[asyncio.Task] = None # type: ignore + _running: bool = field(default=False, init=False) + + def is_alive(self) -> bool: + """Check if the PTY session is still active.""" + if IS_WINDOWS: + return self.winpty_process is not None and self.winpty_process.isalive() + else: + if self.pid is None: + return False + try: + pid_result, status = os.waitpid(self.pid, os.WNOHANG) + if pid_result == 0: + return True # Still running + return False # Exited + except ChildProcessError: + return False # Already reaped + + +class PTYManager: + """Manages PTY sessions for terminal emulation. + + Provides cross-platform terminal emulation with support for: + - Unix systems via the pty module + - Windows via pywinpty (optional dependency) + + Example: + manager = PTYManager() + session = await manager.create_session( + session_id="my-terminal", + on_output=lambda data: print(data.decode()) + ) + await manager.write(session.session_id, b"ls -la\n") + await manager.close_session(session.session_id) + """ + + def __init__(self) -> None: + self._sessions: dict[str, PTYSession] = {} + self._lock = asyncio.Lock() + + @property + def sessions(self) -> dict[str, PTYSession]: + """Get all active sessions.""" + return self._sessions.copy() + + async def create_session( + self, + session_id: str, + cols: int = 80, + rows: int = 24, + on_output: Optional[Callable[[bytes], None]] = None, + shell: Optional[str] = None, + ) -> PTYSession: + """Create a new PTY session. + + Args: + session_id: Unique identifier for the session + cols: Terminal width in columns + rows: Terminal height in rows + on_output: Callback for terminal output + shell: Shell to spawn (defaults to user's shell or /bin/bash) + + Returns: + PTYSession: The created session + + Raises: + RuntimeError: If session creation fails + """ + async with self._lock: + if session_id in self._sessions: + logger.warning(f"Session {session_id} already exists, closing old one") + await self._close_session_internal(session_id) + + if IS_WINDOWS: + session = await self._create_windows_session( + session_id, cols, rows, on_output, shell + ) + else: + session = await self._create_unix_session( + session_id, cols, rows, on_output, shell + ) + + self._sessions[session_id] = session + logger.info(f"Created PTY session: {session_id}") + return session + + async def _create_unix_session( + self, + session_id: str, + cols: int, + rows: int, + on_output: Optional[Callable[[bytes], None]], + shell: Optional[str], + ) -> PTYSession: + """Create a PTY session on Unix systems.""" + shell = shell or os.environ.get("SHELL", "/bin/bash") + + # Fork a new process with a PTY + pid, master_fd = pty.fork() + + if pid == 0: + # Child process - exec the shell + os.execlp(shell, shell, "-i") # noqa: S606 + else: + # Parent process + # Set terminal size + self._set_unix_winsize(master_fd, rows, cols) + + # Make master_fd non-blocking + flags = fcntl.fcntl(master_fd, fcntl.F_GETFL) + fcntl.fcntl(master_fd, fcntl.F_SETFL, flags | os.O_NONBLOCK) + + session = PTYSession( + session_id=session_id, + master_fd=master_fd, + pid=pid, + cols=cols, + rows=rows, + on_output=on_output, + ) + session._running = True + + # Start reader task + session._reader_task = asyncio.create_task(self._unix_reader_loop(session)) + + return session + + async def _create_windows_session( + self, + session_id: str, + cols: int, + rows: int, + on_output: Optional[Callable[[bytes], None]], + shell: Optional[str], + ) -> PTYSession: + """Create a PTY session on Windows systems.""" + if not HAS_WINPTY: + raise RuntimeError( + "pywinpty is required for Windows terminal support. " + "Install it with: pip install pywinpty" + ) + + shell = shell or os.environ.get("COMSPEC", "cmd.exe") + + # Create winpty process + winpty_process = winpty.PtyProcess.spawn( + shell, + dimensions=(rows, cols), + ) + + session = PTYSession( + session_id=session_id, + winpty_process=winpty_process, + cols=cols, + rows=rows, + on_output=on_output, + ) + session._running = True + + # Start reader task + session._reader_task = asyncio.create_task(self._windows_reader_loop(session)) + + return session + + def _set_unix_winsize(self, fd: int, rows: int, cols: int) -> None: + """Set the terminal window size on Unix.""" + winsize = struct.pack("HHHH", rows, cols, 0, 0) + fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize) + + async def _unix_reader_loop(self, session: PTYSession) -> None: + """Read output from Unix PTY and forward to callback.""" + loop = asyncio.get_running_loop() + + try: + while session._running and session.master_fd is not None: + try: + data = await loop.run_in_executor( + None, self._read_unix_pty, session.master_fd + ) + + if data is None: + # No data available, wait a bit + await asyncio.sleep(0.01) + continue + elif data == b"": + # EOF - process terminated + break + elif session.on_output: + session.on_output(data) + + except asyncio.CancelledError: + break + + except Exception as e: + logger.error(f"Unix reader loop error: {e}") + finally: + session._running = False + + def _read_unix_pty(self, fd: int) -> bytes | None: + """Read from Unix PTY file descriptor. + + Returns: + bytes: Data read from PTY + None: No data available (would block) + b'': EOF (process terminated) + """ + try: + data = os.read(fd, 4096) + return data + except BlockingIOError: + return None + except OSError: + return b"" + + async def _windows_reader_loop(self, session: PTYSession) -> None: + """Read output from Windows PTY and forward to callback.""" + loop = asyncio.get_running_loop() + + try: + while ( + session._running + and session.winpty_process is not None + and session.winpty_process.isalive() + ): + try: + data = await loop.run_in_executor( + None, session.winpty_process.read, 4096 + ) + if data and session.on_output: + session.on_output( + data.encode() if isinstance(data, str) else data + ) + except EOFError: + break + except asyncio.CancelledError: + break + + await asyncio.sleep(0.01) + + except Exception as e: + logger.error(f"Windows reader loop error: {e}") + finally: + session._running = False + + async def write(self, session_id: str, data: bytes) -> bool: + """Write data to a PTY session. + + Args: + session_id: The session to write to + data: Data to write + + Returns: + bool: True if write succeeded + """ + session = self._sessions.get(session_id) + if not session: + logger.warning(f"Session {session_id} not found") + return False + + try: + if IS_WINDOWS: + if session.winpty_process: + session.winpty_process.write( + data.decode() if isinstance(data, bytes) else data + ) + return True + else: + if session.master_fd is not None: + os.write(session.master_fd, data) + return True + except Exception as e: + logger.error(f"Write error for session {session_id}: {e}") + + return False + + async def resize(self, session_id: str, cols: int, rows: int) -> bool: + """Resize a PTY session. + + Args: + session_id: The session to resize + cols: New width in columns + rows: New height in rows + + Returns: + bool: True if resize succeeded + """ + session = self._sessions.get(session_id) + if not session: + logger.warning(f"Session {session_id} not found") + return False + + try: + if IS_WINDOWS: + if session.winpty_process: + session.winpty_process.setwinsize(rows, cols) + else: + if session.master_fd is not None: + self._set_unix_winsize(session.master_fd, rows, cols) + + session.cols = cols + session.rows = rows + logger.debug(f"Resized session {session_id} to {cols}x{rows}") + return True + + except Exception as e: + logger.error(f"Resize error for session {session_id}: {e}") + return False + + async def close_session(self, session_id: str) -> bool: + """Close a PTY session. + + Args: + session_id: The session to close + + Returns: + bool: True if session was closed + """ + async with self._lock: + return await self._close_session_internal(session_id) + + async def _close_session_internal(self, session_id: str) -> bool: + """Internal session close without lock.""" + session = self._sessions.pop(session_id, None) + if not session: + return False + + session._running = False + + # Cancel reader task + if session._reader_task: + session._reader_task.cancel() + try: + await session._reader_task + except asyncio.CancelledError: + pass + + # Clean up platform-specific resources + if IS_WINDOWS: + if session.winpty_process: + try: + session.winpty_process.terminate() + except Exception as e: + logger.debug(f"Error terminating winpty: {e}") + else: + # Close file descriptors + if session.master_fd is not None: + try: + os.close(session.master_fd) + except OSError: + pass + + # Terminate child process + if session.pid is not None: + try: + os.kill(session.pid, signal.SIGTERM) + # Use WNOHANG to avoid blocking the event loop + try: + os.waitpid(session.pid, os.WNOHANG) + except ChildProcessError: + pass + except (OSError, ChildProcessError): + pass + + logger.info(f"Closed PTY session: {session_id}") + return True + + async def close_all(self) -> None: + """Close all PTY sessions.""" + async with self._lock: + session_ids = list(self._sessions.keys()) + for session_id in session_ids: + await self._close_session_internal(session_id) + logger.info("Closed all PTY sessions") + + def get_session(self, session_id: str) -> Optional[PTYSession]: + """Get a session by ID. + + Args: + session_id: The session ID + + Returns: + PTYSession or None if not found + """ + return self._sessions.get(session_id) + + def list_sessions(self) -> list[str]: + """List all active session IDs. + + Returns: + List of session IDs + """ + return list(self._sessions.keys()) + + +# Global PTY manager instance +_pty_manager: Optional[PTYManager] = None + + +def get_pty_manager() -> PTYManager: + """Get or create the global PTY manager instance.""" + global _pty_manager + if _pty_manager is None: + _pty_manager = PTYManager() + return _pty_manager 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/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..bf046764c --- /dev/null +++ b/code_puppy/api/routers/sessions.py @@ -0,0 +1,274 @@ +"""Sessions API endpoints for retrieving subagent session data.""" + +import asyncio +import json +import pickle +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 + +# Thread pool for blocking file I/O +_executor = ThreadPoolExecutor(max_workers=2) + +# Timeout for file operations (seconds) +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: + """Get the subagent sessions directory. + + Returns: + Path to the subagent sessions directory + """ + from code_puppy.config import DATA_DIR + + return Path(DATA_DIR) / "subagent_sessions" + + +def _serialize_message(msg: Any) -> Dict[str, Any]: + """Serialize a pydantic-ai message to a 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. + + Args: + msg: A pydantic-ai message object or wrapped message dict + + Returns: + JSON-serializable dictionary representation of the message + """ + + 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 + result = _serialize_obj(msg) + if not isinstance(result, dict): + return {"content": str(result)} + return result + + +def _load_json_sync(file_path: Path) -> dict: + """Synchronous JSON file load (for use in executor).""" + with open(file_path, "r") as f: + return json.load(f) + + +def _load_pickle_sync(file_path: Path) -> Any: + """Synchronous pickle file load (for use in executor).""" + with open(file_path, "rb") as f: + return pickle.load(f) + + +@router.get("/") +async def list_sessions() -> List[SessionInfo]: + """List all available sessions. + + Returns: + List of SessionInfo objects for each session found + """ + sessions_dir = _get_sessions_dir() + if not sessions_dir.exists(): + return [] + + loop = asyncio.get_running_loop() + sessions = [] + + for txt_file in sessions_dir.glob("*.txt"): + session_id = txt_file.stem + try: + # Run blocking I/O in thread pool with timeout + 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: + # Timed out reading file, include basic info + sessions.append(SessionInfo(session_id=session_id)) + except Exception: + # If we can't parse metadata, still include basic session info + sessions.append(SessionInfo(session_id=session_id)) + + return sessions + + +@router.get("/{session_id}") +async def get_session(session_id: str) -> SessionInfo: + """Get session metadata. + + Args: + session_id: The session identifier + + Returns: + SessionInfo with metadata for the specified session + + Raises: + HTTPException: 404 if session not found, 504 on timeout + """ + 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]]: + """Get the full message history for a session. + + Args: + session_id: The session identifier + + Returns: + List of serialized message dictionaries + + Raises: + HTTPException: 404 if session messages not found, 500 on load error, 504 on timeout + """ + sessions_dir = _get_sessions_dir() + pkl_file = sessions_dir / f"{session_id}.pkl" + + if not pkl_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_pickle_sync, pkl_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]: + """Delete a session and its data. + + Args: + session_id: The session identifier + + Returns: + Success message dict + + Raises: + HTTPException: 404 if session not found + """ + sessions_dir = _get_sessions_dir() + txt_file = sessions_dir / f"{session_id}.txt" + pkl_file = sessions_dir / f"{session_id}.pkl" + + if not txt_file.exists() and not pkl_file.exists(): + raise HTTPException(404, f"Session '{session_id}' not found") + + if txt_file.exists(): + txt_file.unlink() + if pkl_file.exists(): + pkl_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..8a5cbb445 --- /dev/null +++ b/code_puppy/api/routers/ws_sessions.py @@ -0,0 +1,302 @@ +"""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"), + } + ) + + 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. + + Also attempts to clean up any legacy pkl files if they exist. + """ + ws_dir = _get_ws_sessions_dir() + safe_base = _validate_session_name(session_name, ws_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") + + # Cleanup legacy files (best effort) + try: + pkl_file = ws_dir / f"{safe_base}.pkl" + meta_file = ws_dir / f"{safe_base}_meta.json" + if pkl_file.exists(): + pkl_file.unlink() + if meta_file.exists(): + meta_file.unlink() + except Exception: + pass + + 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..235041418 --- /dev/null +++ b/code_puppy/api/session_cache.py @@ -0,0 +1,392 @@ +"""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 pickle +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, pkl_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 = pkl_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, + pkl_path: Path, + messages: List[Any], + json_messages: List[Dict[str, Any]], + metadata: Dict[str, Any], + ): + """Cache session data. + + Args: + session_id: Unique session identifier + pkl_path: Path to the pickle 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 = pkl_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_pickle_sync(pkl_path: Path) -> List[Any]: + """Synchronous pickle loading (for executor).""" + with open(pkl_path, "rb") as f: + return pickle.load(f) + + +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, pkl_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 + pkl_path: Path to the pickle file + timeout: Timeout for file operations + + Returns: + Tuple of (json_messages, metadata) + + Raises: + FileNotFoundError: If pickle file doesn't exist + asyncio.TimeoutError: If load times out + """ + cache = get_session_cache() + + # Try cache first + cached = await cache.get(session_id, pkl_path) + if cached is not None: + return cached + + # Cache miss - load from disk + if not pkl_path.exists(): + raise FileNotFoundError(f"Session file not found: {pkl_path}") + + loop = asyncio.get_running_loop() + + # Load pickle in thread pool + start_time = time.time() + messages = await asyncio.wait_for( + loop.run_in_executor(_executor, _load_pickle_sync, pkl_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, pkl_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..a8621e1bf --- /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 to {WS_SESSION_DIR}/{session_id}.pkl — that pkl 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 pkl 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/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..f6208f104 --- /dev/null +++ b/code_puppy/api/websocket.py @@ -0,0 +1,51 @@ +"""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/terminal - Interactive PTY terminal sessions +- /ws/health - Simple health check endpoint +- /ws/sessions - Real-time session monitoring + +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, + register_sessions_endpoint, + register_terminal_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_terminal_endpoint(app) + register_health_endpoint(app) + register_sessions_endpoint(app) + + logger.debug("All WebSocket endpoints registered") From 45c1b50e7dbf4fe254b45e1fb8aa52cbd56d0b0e Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Thu, 28 May 2026 18:38:32 -0500 Subject: [PATCH 05/29] feat(ws): migrate modular websocket handlers and runtime management --- code_puppy/api/ws/__init__.py | 30 + code_puppy/api/ws/attachments.py | 93 +++ code_puppy/api/ws/background_save.py | 244 +++++++ code_puppy/api/ws/connection_manager.py | 55 ++ code_puppy/api/ws/events_handler.py | 53 ++ code_puppy/api/ws/health_handler.py | 22 + code_puppy/api/ws/runtime/__init__.py | 14 + code_puppy/api/ws/runtime/session_runtime.py | 68 ++ .../api/ws/runtime/session_runtime_manager.py | 150 +++++ code_puppy/api/ws/schemas.py | 603 ++++++++++++++++++ code_puppy/api/ws/sessions_handler.py | 69 ++ code_puppy/api/ws/terminal_handler.py | 141 ++++ 12 files changed, 1542 insertions(+) create mode 100644 code_puppy/api/ws/__init__.py create mode 100644 code_puppy/api/ws/attachments.py create mode 100644 code_puppy/api/ws/background_save.py create mode 100644 code_puppy/api/ws/connection_manager.py create mode 100644 code_puppy/api/ws/events_handler.py create mode 100644 code_puppy/api/ws/health_handler.py create mode 100644 code_puppy/api/ws/runtime/__init__.py create mode 100644 code_puppy/api/ws/runtime/session_runtime.py create mode 100644 code_puppy/api/ws/runtime/session_runtime_manager.py create mode 100644 code_puppy/api/ws/schemas.py create mode 100644 code_puppy/api/ws/sessions_handler.py create mode 100644 code_puppy/api/ws/terminal_handler.py diff --git a/code_puppy/api/ws/__init__.py b/code_puppy/api/ws/__init__.py new file mode 100644 index 000000000..1649afd36 --- /dev/null +++ b/code_puppy/api/ws/__init__.py @@ -0,0 +1,30 @@ +"""WebSocket endpoint handlers package. + +Each handler module exposes a `register_*_endpoint(app)` function +that attaches one WebSocket route to the FastAPI application. + +Modules: + - chat_handler: Interactive chat with the Code Puppy agent (/ws/chat) + - events_handler: Server-sent events stream (/ws/events) + - terminal_handler: Interactive PTY terminal (/ws/terminal) + - health_handler: Simple health check (/ws/health) + - sessions_handler: Real-time session monitoring (/ws/sessions) + - connection_manager: Singleton for broadcasting session updates + - attachments: File attachment processing utilities +""" + +from code_puppy.api.ws.chat_handler import register_chat_endpoint +from code_puppy.api.ws.connection_manager import connection_manager +from code_puppy.api.ws.events_handler import register_events_endpoint +from code_puppy.api.ws.health_handler import register_health_endpoint +from code_puppy.api.ws.sessions_handler import register_sessions_endpoint +from code_puppy.api.ws.terminal_handler import register_terminal_endpoint + +__all__ = [ + "register_chat_endpoint", + "register_events_endpoint", + "register_terminal_endpoint", + "register_health_endpoint", + "register_sessions_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/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/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..05e56528d --- /dev/null +++ b/code_puppy/api/ws/schemas.py @@ -0,0 +1,603 @@ +"""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, 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" + + +class ServerMessageType(str, Enum): + """All ``type`` values the server may send over the WebSocket.""" + + SYSTEM = "system" + 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" + 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 + + +# ═══════════════════════════════════════════════════════════════════════════════ +# SERVER → CLIENT (21 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. + """ + + type: Literal["assistant_message_end"] = "assistant_message_end" + message_id: str + 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 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")], + ], + 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[ServerCancelled, Tag("cancelled")], + ], + Discriminator("type"), +] +"""Discriminated union of all server→client message types.""" 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/terminal_handler.py b/code_puppy/api/ws/terminal_handler.py new file mode 100644 index 000000000..b32bff451 --- /dev/null +++ b/code_puppy/api/ws/terminal_handler.py @@ -0,0 +1,141 @@ +"""WebSocket endpoint for interactive PTY terminal sessions.""" + +import asyncio +import base64 +import logging +import uuid + +from fastapi import FastAPI, WebSocket, WebSocketDisconnect + +logger = logging.getLogger(__name__) + + +def register_terminal_endpoint(app: FastAPI) -> None: + """Register the /ws/terminal WebSocket endpoint.""" + + @app.websocket("/ws/terminal") + async def websocket_terminal(websocket: WebSocket) -> None: + """Interactive terminal WebSocket endpoint.""" + await websocket.accept() + logger.debug("Terminal WebSocket client connected") + + from code_puppy.api.pty_manager import get_pty_manager + + manager = get_pty_manager() + session_id = str(uuid.uuid4())[:8] + session = None + + loop = asyncio.get_running_loop() + output_queue: asyncio.Queue[bytes] = asyncio.Queue() + + def on_output(data: bytes) -> None: + try: + loop.call_soon_threadsafe(output_queue.put_nowait, data) + except Exception as e: + logger.error("on_output error: %s", e) + + async def output_sender() -> None: + """Coroutine that sends queued output to WebSocket.""" + try: + while True: + data = await output_queue.get() + await websocket.send_json( + { + "type": "output", + "data": base64.b64encode(data).decode("ascii"), + } + ) + except asyncio.CancelledError: + pass + except Exception as e: + logger.error("output_sender error: %s", e) + + sender_task = None + + try: + session = await manager.create_session( + session_id=session_id, + on_output=on_output, + ) + + from code_puppy.api.session_context import session_manager + + try: + ctx = await session_manager.create_session(f"terminal_{session_id}") + current_agent_name = ctx.agent_name + except Exception: + current_agent_name = "code-puppy" + ctx = None + + await websocket.send_json( + { + "type": "session_started", + "session_id": session_id, + "current_agent": current_agent_name, + } + ) + + sender_task = asyncio.create_task(output_sender()) + + while True: + try: + msg = await websocket.receive_json() + + if msg.get("type") == "input": + data = msg.get("data", "") + if isinstance(data, str): + data = data.encode("utf-8") + await manager.write(session_id, data) + elif msg.get("type") == "resize": + cols = msg.get("cols", 80) + rows = msg.get("rows", 24) + await manager.resize(session_id, cols, rows) + elif msg.get("type") == "switch_agent": + agent_name = msg.get("agent_name") + if agent_name: + try: + await session_manager.switch_agent( + f"terminal_{session_id}", agent_name + ) + if ctx: + ctx = await session_manager.get_session( + f"terminal_{session_id}" + ) + await websocket.send_json( + {"type": "agent_changed", "agent_name": agent_name} + ) + except Exception as e: + logger.error("Error switching agent: %s", e) + await websocket.send_json( + { + "type": "error", + "message": f"Failed to switch to agent {agent_name}: {str(e)}", + } + ) + elif msg.get("type") == "get_current_agent": + if ctx: + current_agent_name = ctx.agent_name + else: + current_agent_name = "code-puppy" + await websocket.send_json( + {"type": "current_agent", "agent_name": current_agent_name} + ) + except WebSocketDisconnect: + break + except Exception as e: + logger.error("Terminal WebSocket error: %s", e) + break + except Exception as e: + logger.error("Terminal session error: %s", e) + finally: + if sender_task: + sender_task.cancel() + try: + await sender_task + except asyncio.CancelledError: + pass + if ctx: + await session_manager.destroy_session(f"terminal_{session_id}") + if session: + await manager.close_session(session_id) + logger.debug("Terminal WebSocket disconnected") From 796062f74525e287e384a12f8ad2b8b7bf6ae093 Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Thu, 28 May 2026 18:38:46 -0500 Subject: [PATCH 06/29] test(api): migrate api/ws test suites and desk templates --- code_puppy/api/templates/chat.html | 942 ++++++++++++++++++ code_puppy/api/templates/terminal.html | 382 +++++++ code_puppy/api/tests/__init__.py | 0 .../api/tests/test_backfill_pydantic_json.py | 349 +++++++ code_puppy/api/tests/test_session_load.py | 503 ++++++++++ .../api/tests/test_session_model_compat.py | 64 ++ code_puppy/api/ws/tests/__init__.py | 0 .../api/ws/tests/test_background_save.py | 559 +++++++++++ .../api/ws/tests/test_event_driven_drain.py | 556 +++++++++++ .../api/ws/tests/test_event_driven_polling.py | 337 +++++++ .../tests/test_tool_group_id_consistency.py | 67 ++ 11 files changed, 3759 insertions(+) create mode 100644 code_puppy/api/templates/chat.html create mode 100644 code_puppy/api/templates/terminal.html create mode 100644 code_puppy/api/tests/__init__.py create mode 100644 code_puppy/api/tests/test_backfill_pydantic_json.py create mode 100644 code_puppy/api/tests/test_session_load.py create mode 100644 code_puppy/api/tests/test_session_model_compat.py create mode 100644 code_puppy/api/ws/tests/__init__.py create mode 100644 code_puppy/api/ws/tests/test_background_save.py create mode 100644 code_puppy/api/ws/tests/test_event_driven_drain.py create mode 100644 code_puppy/api/ws/tests/test_event_driven_polling.py create mode 100644 code_puppy/api/ws/tests/test_tool_group_id_consistency.py diff --git a/code_puppy/api/templates/chat.html b/code_puppy/api/templates/chat.html new file mode 100644 index 000000000..b23d44951 --- /dev/null +++ b/code_puppy/api/templates/chat.html @@ -0,0 +1,942 @@ + + + + + + 🐶 Code Puppy Chat + + + + + + + + + + +
+ +
+
+ 🐶 +
+

Code Puppy Chat

+

Connecting...

+
+
+ +
+ +
+ + +
+ +
+ + +
+ + +
+
Working Dir:
+
+ + +
+
Not set
+
+ + +
+
+ Disconnected +
+ + +
+ + + Home + +
+
+
+ + +
+ +
+ + +
+
+ + +
+
+
+
+ + + + diff --git a/code_puppy/api/templates/terminal.html b/code_puppy/api/templates/terminal.html new file mode 100644 index 000000000..568dc7491 --- /dev/null +++ b/code_puppy/api/templates/terminal.html @@ -0,0 +1,382 @@ + + + + + + 🐶 Code Puppy Terminal + + + + + + + + + + + + + + + +
+ +
+
+ 🐶 +

Code Puppy Terminal

+
+ +
+ +
+
+ Disconnected +
+ + + + +
+ + +
+
+
+ + +
+
+
+ + +
+
+ 80x24 + • +
+ Agent: + +
+
+
+ Ctrl+C to interrupt + • + Ctrl+D to exit +
+
+
+ + + + diff --git a/code_puppy/api/tests/__init__.py b/code_puppy/api/tests/__init__.py new file mode 100644 index 000000000..e69de29bb 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_session_load.py b/code_puppy/api/tests/test_session_load.py new file mode 100644 index 000000000..674f2eae4 --- /dev/null +++ b/code_puppy/api/tests/test_session_load.py @@ -0,0 +1,503 @@ +"""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 \ No newline at end of file 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..0b008fd87 --- /dev/null +++ b/code_puppy/api/tests/test_session_model_compat.py @@ -0,0 +1,64 @@ +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/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_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_tool_group_id_consistency.py b/code_puppy/api/ws/tests/test_tool_group_id_consistency.py new file mode 100644 index 000000000..7175f67b0 --- /dev/null +++ b/code_puppy/api/ws/tests/test_tool_group_id_consistency.py @@ -0,0 +1,67 @@ +"""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``. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + +CHAT_HANDLER_PATH = Path(__file__).resolve().parents[1] / "chat_handler.py" + + +def _get_server_tool_result_calls(source: str) -> list[ast.Call]: + """Return every ``ServerToolResult(...)`` call in chat_handler 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. + """ + source = CHAT_HANDLER_PATH.read_text(encoding="utf-8") + calls = _get_server_tool_result_calls(source) + + assert calls, "Expected at least one ServerToolResult(...) call" + + missing = [] + 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(call.lineno) + + assert not missing, ( + f"ServerToolResult(...) calls missing tool_group_id keyword at lines: {missing}" + ) + + +def test_server_tool_result_calls_never_pass_explicit_none_for_tool_group_id() -> None: + """Guard against explicit ``tool_group_id=None`` regressions.""" + source = CHAT_HANDLER_PATH.read_text(encoding="utf-8") + calls = _get_server_tool_result_calls(source) + + explicit_none_lines: list[int] = [] + 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_lines.append(call.lineno) + + assert not explicit_none_lines, ( + "ServerToolResult(...) calls pass explicit tool_group_id=None at lines: " + f"{explicit_none_lines}" + ) From 1e5bf344e34e8ddf1d863ba85d67857c1f1d32aa Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Thu, 28 May 2026 18:38:52 -0500 Subject: [PATCH 07/29] feat(core): migrate backend logging and model error normalization modules --- code_puppy/logging_setup.py | 223 +++++++++++++++++++++++++++++ code_puppy/model_errors.py | 277 ++++++++++++++++++++++++++++++++++++ 2 files changed, 500 insertions(+) create mode 100644 code_puppy/logging_setup.py create mode 100644 code_puppy/model_errors.py 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/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, + ) From 52cd3b1ea581f24770eaaa78ebf55812a0cbd11b Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Mon, 1 Jun 2026 13:56:44 -0500 Subject: [PATCH 08/29] chore(migration): add legacy ws snapshot and migration docs --- .gitignore | 2 + code_puppy/api/ws/legacy/README.md | 11 + code_puppy/api/ws/legacy/__init__.py | 12 + code_puppy/api/ws/legacy/attachments.py | 93 + code_puppy/api/ws/legacy/background_save.py | 244 ++ code_puppy/api/ws/legacy/chat_handler.py | 3898 +++++++++++++++++ .../api/ws/legacy/connection_manager.py | 55 + code_puppy/api/ws/legacy/events_handler.py | 53 + code_puppy/api/ws/legacy/health_handler.py | 22 + code_puppy/api/ws/legacy/runtime/__init__.py | 14 + .../api/ws/legacy/runtime/session_runtime.py | 68 + .../legacy/runtime/session_runtime_manager.py | 150 + code_puppy/api/ws/legacy/schemas.py | 603 +++ code_puppy/api/ws/legacy/sessions_handler.py | 69 + .../puppy-desk-gate1-investigation-summary.md | 426 ++ .../puppy-desk-migration-plan-options.md | 825 ++++ .../puppy-desk-uncommitted-change-snapshot.md | 1567 +++++++ 17 files changed, 8112 insertions(+) create mode 100644 code_puppy/api/ws/legacy/README.md create mode 100644 code_puppy/api/ws/legacy/__init__.py create mode 100644 code_puppy/api/ws/legacy/attachments.py create mode 100644 code_puppy/api/ws/legacy/background_save.py create mode 100644 code_puppy/api/ws/legacy/chat_handler.py create mode 100644 code_puppy/api/ws/legacy/connection_manager.py create mode 100644 code_puppy/api/ws/legacy/events_handler.py create mode 100644 code_puppy/api/ws/legacy/health_handler.py create mode 100644 code_puppy/api/ws/legacy/runtime/__init__.py create mode 100644 code_puppy/api/ws/legacy/runtime/session_runtime.py create mode 100644 code_puppy/api/ws/legacy/runtime/session_runtime_manager.py create mode 100644 code_puppy/api/ws/legacy/schemas.py create mode 100644 code_puppy/api/ws/legacy/sessions_handler.py create mode 100644 docs/migration/puppy-desk-gate1-investigation-summary.md create mode 100644 docs/migration/puppy-desk-migration-plan-options.md create mode 100644 docs/migration/puppy-desk-uncommitted-change-snapshot.md diff --git a/.gitignore b/.gitignore index 56e9f3c35..603b7981f 100644 --- a/.gitignore +++ b/.gitignore @@ -45,3 +45,5 @@ code_puppy/bundled_skills/ .claude/hooks/ts-hooks/dist/ .json + +.maestro/* diff --git a/code_puppy/api/ws/legacy/README.md b/code_puppy/api/ws/legacy/README.md new file mode 100644 index 000000000..113e3a41f --- /dev/null +++ b/code_puppy/api/ws/legacy/README.md @@ -0,0 +1,11 @@ +# Legacy WS Snapshot — Puppy Desk Gate 2 + +This directory is a reference/rollback copy of `code_puppy.api.ws` captured +before further puppy-desk migration/refactor work. + +Rules: + +- It is not registered by `setup_websocket()`. +- Existing `/ws/chat` behavior remains provided by `code_puppy.api.ws.chat_handler`. +- Same-route swap must not happen until parity and GUGI validation pass. +- Keep this snapshot until user explicitly approves cleanup. diff --git a/code_puppy/api/ws/legacy/__init__.py b/code_puppy/api/ws/legacy/__init__.py new file mode 100644 index 000000000..fae8b03db --- /dev/null +++ b/code_puppy/api/ws/legacy/__init__.py @@ -0,0 +1,12 @@ +"""Legacy WebSocket implementation snapshot for puppy-desk migration. + +This package is a reference/rollback copy captured during Gate 2 of the +puppy-desk migration. It is intentionally not imported or registered by the +runtime application yet; the active routes continue to come from +``code_puppy.api.ws``. + +Do not edit this snapshot as part of feature work unless intentionally updating +the legacy rollback point. +""" + +LEGACY_SNAPSHOT_PURPOSE = "puppy-desk-gate2-reference-copy" diff --git a/code_puppy/api/ws/legacy/attachments.py b/code_puppy/api/ws/legacy/attachments.py new file mode 100644 index 000000000..f2db90726 --- /dev/null +++ b/code_puppy/api/ws/legacy/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/legacy/background_save.py b/code_puppy/api/ws/legacy/background_save.py new file mode 100644 index 000000000..4b1a1c6ad --- /dev/null +++ b/code_puppy/api/ws/legacy/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/legacy/chat_handler.py b/code_puppy/api/ws/legacy/chat_handler.py new file mode 100644 index 000000000..fbcb4c918 --- /dev/null +++ b/code_puppy/api/ws/legacy/chat_handler.py @@ -0,0 +1,3898 @@ +"""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 +import re +import uuid +from pathlib import Path +from typing import Any + +from fastapi import FastAPI, WebSocket, WebSocketDisconnect +from pydantic import TypeAdapter, ValidationError + +from code_puppy.api.db.queries import ( + update_session_working_directory, + write_error_message_to_sqlite, + write_system_message_to_sqlite, + write_turn_to_sqlite, +) +from code_puppy.api.error_parser import parse_api_error as _legacy_parse_api_error +from code_puppy.api.session_context import _validate_session_id, session_manager +from code_puppy.api.ws.attachments import build_file_context_and_attachments +from code_puppy.api.ws.background_save import ( + fire_and_track, + save_agent_result_in_background, +) +from code_puppy.api.ws.connection_manager import connection_manager +from code_puppy.api.ws.schemas import ( + PROTOCOL_VERSION, + ClientMessage, + ServerAgentInvoked, + ServerAssistantMessageDelta, + ServerAssistantMessageEnd, + ServerAssistantMessageStart, + ServerCancelled, + ServerCommandResult, + ServerConfigValue, + ServerError, + ServerMessage, + ServerResponse, + ServerSessionMetaUpdated, + ServerSessionRestored, + ServerSessionSwitched, + ServerStatus, + ServerStreamEnd, + ServerSystem, + ServerToolCall, + ServerToolResult, + ServerUserMessage, + ServerWorkingDirectoryChanged, +) +from code_puppy.config import get_global_model_name +from code_puppy.messaging.bus import get_message_bus +from code_puppy.model_errors import normalize_model_error as _normalize_model_error +from code_puppy.session_storage import generate_heuristic_title +from code_puppy.tools.command_runner import ( + cleanup_session_process_tracking, + init_session_process_tracking, +) + +# Import session context for desk-puppy working directory support +try: + from code_puppy.plugins.walmart_specific.session_context import ( + clear_session_working_directory, + set_session_working_directory, + ) + + HAS_SESSION_CONTEXT = True +except ImportError: + HAS_SESSION_CONTEXT = False + + def set_session_working_directory(d): + pass + + def clear_session_working_directory(): + pass + + +_ClientMessageAdapter = TypeAdapter(ClientMessage) + +logger = logging.getLogger(__name__) + +# === TOOL-CALL-PARITY BRANCH LOADED === +logger.warning( + "🐕 CHAT_HANDLER LOADED FROM TOOL-CALL-PARITY BRANCH - ToolReturnPart fix active!" +) +print("🐕 CHAT_HANDLER LOADED FROM TOOL-CALL-PARITY BRANCH", flush=True) + + +# --------------------------------------------------------------------------- +# WebSocket error frame helpers (see fix/orphaned-tool-use-id-history) +# --------------------------------------------------------------------------- + +_UNKNOWN_ERROR_TYPE = "unknown_error" +_TOOL_HISTORY_ERROR_TYPE = "tool_history_error" + +_CODE_TO_ERROR_TYPE: dict = { + "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: + """Convert an agent-run exception into a structured frontend error dict. + + Wraps ``normalize_model_error`` for primary classification (including the + new ``invalid_tool_history`` code for Anthropic HTTP 400 tool-mismatch + errors), falling back to the legacy ``_legacy_parse_api_error`` for all + other categories so existing behaviour is unchanged. + + Returns: + Dict with keys: user_message, error_type, technical_details, action_required. + """ + 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, + } + # Fall back to legacy parser for unclassified errors + return _legacy_parse_api_error(exc) + + +def has_streamed_content(collected_text) -> bool: + """Return True when the client has already received streaming output chunks. + + Args: + collected_text: List of text chunks accumulated during the agent run. + Elements may be None or empty strings. + + Returns: + True if at least one chunk has substantive (non-whitespace) content. + """ + return any((chunk or "").strip() for chunk in collected_text) + + +def build_error_response_frames( + agent_error: Exception, + collected_text, + session_id: str, +) -> list: + """Build the ordered WebSocket frames to send when an agent error occurs. + + When streaming output was already delivered to the client, a ``stream_end`` + frame (``success=False``) is prepended so the frontend can exit its + streaming state before processing the error frame. Without this handshake + the frontend hangs indefinitely waiting for a ``stream_end`` that never + arrives. + + Args: + agent_error: The exception raised by the agent run. + collected_text: Accumulated streaming chunks from the current turn. + session_id: The WebSocket session identifier. + + Returns: + List of JSON-serialisable dicts to send to the client in order. + """ + frames: list[dict] = [] + 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 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 response + {"type": "response", "content": "...", "done": true, + "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 + ) + + # Flag to track if WebSocket is still open + ws_closed = False + + async def persist_error_payload(data: dict[str, Any]) -> None: + """Persist structured error frames so they survive reloads.""" + if data.get("type") != "error" or not session_id: + return + + try: + await write_error_message_to_sqlite( + session_id=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=(ctx.agent_name if ctx else ""), + model_name=(ctx.model_name if ctx else ""), + timestamp=datetime.datetime.now(datetime.timezone.utc).isoformat(), + ) + except Exception: + logger.warning( + "[WS:%s] Failed to persist error payload to SQLite", + session_id, + exc_info=True, + ) + + async def safe_send_json(data: dict) -> bool: + """Safely send JSON to WebSocket, returns False if connection is closed.""" + nonlocal ws_closed + if ws_closed: + logger.debug( + "[WS:%s] safe_send_json skipped (ws_closed=True): type=%s", + session_id, + data.get("type"), + ) + return False + + msg_type = data.get("type") + if msg_type in { + "error", + "response", + "assistant_message_start", + "assistant_message_end", + }: + # Keep this low-noise but useful for tracing request lifecycle + logger.debug( + "[WS:%s] → send_json type=%s keys=%s", + 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", + session_id, + data.get("error_type"), + data.get("action_required"), + (data.get("error") or "")[:300], + ) + + try: + if msg_type == "error": + await persist_error_payload(data) + await websocket.send_json(data) + return True + except Exception as e: + logger.warning( + "[WS:%s] send_json failed for type=%s: %s", + session_id, + msg_type, + e, + exc_info=True, + ) + if "close message" in str(e).lower() or "closed" in str(e).lower(): + ws_closed = True + logger.debug("WebSocket closed, stopping sends") + return False + + async def send_typed(msg: ServerMessage) -> bool: + """Send a typed protocol message to the client.""" + return await safe_send_json(msg.model_dump(exclude_none=True)) + + async def send_typed_tool_lifecycle( + msg: ServerToolCall | ServerToolResult, + ) -> bool: + """Send tool lifecycle frames to the client.""" + return await send_typed(msg) + + ctx = None + session_title = "" + session_working_directory = "" # Will be set by client or loaded from metadata + session_pinned = False # Track pinned state for persistence + last_context_sent_directory = "" # Track when we last sent directory context + existing_history = None + + # Use provided session_id or generate new one + if session_id: + # Validate session_id to prevent path traversal + try: + _validate_session_id(session_id) + except ValueError as e: + logger.warning("Invalid session_id rejected: %r: %s", session_id, e) + await websocket.close(code=1008, reason="Invalid session ID") + return + + # Client wants to resume/continue a specific session + logger.debug("Client requested session: %s", session_id) + + # SQLite is the source of truth — check session existence there. + 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 e: + logger.warning("Failed to check session in SQLite: %s", e) + + if not existing_history: + logger.debug( + "Session %s not found in SQLite, will create new", session_id + ) + else: + # Generate new timestamp-based session ID + session_timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + session_id = f"WS_session_{session_timestamp}" + logger.debug("Generated new session ID: %s", session_id) + + try: + # --- SESSION ISOLATION: create or load via SessionManager --- + try: + if existing_history is not None: + # Session confirmed in SQLite — load it via SessionManager. + # get_or_load_session checks in-memory first, then falls + # back to a fresh SQLite load. SQLite is the sole source of truth. + ctx = await session_manager.get_or_load_session(session_id) + if ctx is None: + # SQLite confirmed this session exists but load returned None. + # This is unexpected (DB race, corrupt row, or schema mismatch). + # Log a warning and start fresh so the WebSocket can still + # function — silent data loss is worse than a blank session. + 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, + ) + ctx = await session_manager.create_session(session_id) + else: + # Update local state from loaded metadata + session_title = ctx.title + session_working_directory = ctx.working_directory + session_pinned = ctx.pinned + else: + ctx = await session_manager.create_session(session_id) + except Exception as e: + logger.warning("SessionManager init failed, falling back: %s", e) + try: + ctx = await session_manager.create_session(session_id) + except Exception: + logger.error("SessionManager fallback also failed", exc_info=True) + await websocket.close(code=1011, reason="Session init failed") + return + + # Mark session as active (cancels any pending cleanup from previous disconnect) + await session_manager.mark_session_active(session_id) + + # Initialize per-session process tracking (ContextVar isolation) + init_session_process_tracking() + + # Set MessageBus session context for this WS connection + try: + bus = get_message_bus() + bus.set_session_context(session_id) + except Exception: + logger.debug("MessageBus session context not available") + + # Convenience aliases — use ctx.agent everywhere instead of get_current_agent() + agent = ctx.agent + agent_name = ctx.agent_name + model_name = ctx.model_name + + # ── Persist initial session context to SQLite (new sessions only) ─────── + # Write a config row (agent + model) and, if a CWD is set, a directory + # row so the FE cold-load path sees them before the first user message. + # Skipped for resumed sessions — their init rows were written at creation. + if existing_history is None: + try: + _now_iso = datetime.datetime.now(datetime.timezone.utc).isoformat() + _init_agent = agent_name or "code-puppy" + _init_model = model_name or "unknown" + await write_system_message_to_sqlite( + session_id=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 session_working_directory: + _path_segments = session_working_directory.split("/")[-3:] + _relative = "/".join(s for s in _path_segments if s) + await write_system_message_to_sqlite( + session_id=session_id, + system_message_type="directory", + content=f"Starting in {_relative}", + system_message_path=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, + ) + + # Send welcome message + await send_typed( + ServerSystem( + content=f"Connected! Session: {session_id}", + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + resumed=existing_history is not None, + protocol_version=PROTOCOL_VERSION, + ) + ) + + # If we resumed an existing session, notify the client + if existing_history and ctx: + try: + # History is already loaded into ctx.agent by session_manager.load_session(). + # Pull it back out for the client-side UI metadata notification. + loaded_messages = ctx.agent.get_message_history() + message_count = len(loaded_messages) if loaded_messages else 0 + + # Notify client about restored session + await send_typed( + ServerSessionRestored( + session_id=session_id, + message_count=message_count, + title=session_title, + ui_metadata=[], + ) + ) + + # Replay system messages (agent/model switches, CWD banners) to UI. + # These are persisted with pydantic_json=NULL so _load_from_sqlite + # skips them — we must send them separately so the UI shows them. + try: + from code_puppy.api.db.queries import get_active_messages + + rows = await get_active_messages(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=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), + session_id, + ) + except Exception as sys_exc: + logger.warning( + "Failed to replay system messages for session %s: %s", + session_id, + sys_exc, + ) + + agent = ctx.agent # Refresh local alias + logger.debug("Restored %d messages to session agent", message_count) + except Exception as e: + logger.warning("Failed to restore session history: %s", e) + # Track active streaming task for cancellation + active_drain_task = None + # Track active agent task (run_with_mcp) for cancellation + active_agent_task: asyncio.Task | None = None + # Track stop_draining event for cancellation across iterations + stop_draining = asyncio.Event() + + 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") == "switch_agent": + agent_name = msg.get("agent_name") + if agent_name: + try: + new_agent = await session_manager.switch_agent( + session_id, agent_name + ) + agent = new_agent # Update local alias + ctx.agent = new_agent + model_name = ( + new_agent.get_model_name() + if new_agent + else "unknown" + ) + + # Persist to SQLite so the FE cold-load path sees it + try: + await write_system_message_to_sqlite( + session_id=session_id, + system_message_type="config", + content=f"🔄 Switched to {agent_name} ({model_name})", + agent_name=agent_name, + model_name=model_name, + ) + except Exception as _sw_exc: + logger.warning( + "Agent-switch SQLite write failed: %s", + _sw_exc, + exc_info=True, + ) + + await send_typed( + ServerSystem( + content=f"🔄 Switched to {agent_name} ({model_name})", + session_id=session_id, + agent_name=agent_name, + model_name=model_name, + ) + ) + except Exception as e: + logger.error("Error switching agent: %s", e) + await send_typed( + ServerError( + error=f"Failed to switch to agent {agent_name}: {str(e)}", + session_id=session_id, + ) + ) + + elif msg.get("type") == "switch_model": + model_name = msg.get("model_name") or msg.get("model") + if model_name: + try: + await session_manager.switch_model( + session_id, model_name + ) + agent = ctx.agent # Refresh local alias + logger.debug("Switched model to: %s", model_name) + + # Persist to SQLite so the FE cold-load path sees it + _sw_agent = ctx.agent_name or agent_name or "code-puppy" + try: + await write_system_message_to_sqlite( + session_id=session_id, + system_message_type="config", + content=f"🔄 Switched to {_sw_agent} ({model_name})", + agent_name=_sw_agent, + model_name=model_name, + ) + except Exception as _sw_exc: + logger.warning( + "Model-switch SQLite write failed: %s", + _sw_exc, + exc_info=True, + ) + + await send_typed( + ServerSystem( + content=f"🔄 Switched to {_sw_agent} ({model_name})", + session_id=session_id, + model_name=model_name, + agent_name=_sw_agent, + ) + ) + except Exception as e: + logger.error("Error switching model: %s", e) + await send_typed( + ServerError( + error=f"Failed to switch to model {model_name}: {str(e)}", + session_id=session_id, + ) + ) + else: + await send_typed( + ServerError( + error="No model_name provided for switch_model", + session_id=session_id, + ) + ) + + elif msg.get("type") == "switch_session": + """ + Switch to a different session without reconnecting WebSocket. + + How it works: + 1. Client sends: {"type": "switch_session", "session_id": "WS_session_20260120_124356"} + 2. Server loads session from ~/.code_puppy/ws_sessions/{session_id}.pkl + 3. Server restores message history to the current agent + 4. Server responds with session_switched event containing message count and title + 5. Client can continue chatting with full conversation context + + This approach keeps a single WebSocket connection alive while allowing + users to switch between multiple sessions instantly. All session data + stays on the local filesystem for privacy. + """ + # Cancel any active streaming first + if active_drain_task and not active_drain_task.done(): + logger.debug( + "Cancelling active streaming due to session switch" + ) + stop_draining.set() # Signal drain to stop + active_drain_task.cancel() + try: + await active_drain_task + except asyncio.CancelledError: + pass + stop_draining.clear() # Reset for next streaming + + 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=session_id, + ) + ) + continue + + # Validate session_id to prevent path traversal + try: + _validate_session_id(new_session_id) + except ValueError as e: + logger.warning( + f"Invalid session_id in switch_session: {new_session_id!r}" + ) + await send_typed( + ServerError( + error=f"Invalid session ID: {e}", + session_id=session_id, + ) + ) + continue + + logger.debug("Switching to session: %s", new_session_id) + + try: + # Check if target session exists in SQLite + try: + from code_puppy.api.db.queries import ( + session_exists as _se, + ) + + _target_exists = await _se(new_session_id) + except Exception: + _target_exists = False + + if not _target_exists: + # Session doesn't exist - create new one with this ID + logger.debug( + f"Session {new_session_id} not found, creating new" + ) + + # Save current session and mark inactive (keep in memory for 15 min) + try: + await session_manager.save_session(session_id) + except Exception: + pass + await session_manager.mark_session_inactive(session_id) + + session_id = new_session_id + session_title = "" + session_working_directory = "" + session_pinned = False + + ctx = await session_manager.create_session(session_id) + # Mark the new session as active + await session_manager.mark_session_active(session_id) + agent = ctx.agent + agent_name = ctx.agent_name + model_name = ctx.model_name + + await send_typed( + ServerSessionSwitched( + session_id=new_session_id, + message_count=0, + title="", + created=True, + agent_name=agent_name, + model_name=model_name, + ) + ) + continue + + # Load existing session via SessionManager + # Save current session and mark inactive (keep in memory for 15 min) + try: + await session_manager.save_session(session_id) + except Exception: + pass + await session_manager.mark_session_inactive(session_id) + + session_id = new_session_id + + # Try loading via SessionManager (checks in-memory first, then SQLite) + loaded_ctx = await session_manager.get_or_load_session( + session_id + ) + if loaded_ctx is not None: + ctx = loaded_ctx + session_title = ctx.title + session_working_directory = ctx.working_directory + session_pinned = ctx.pinned + else: + # Legacy session or HMAC missing — load metadata + # from JSON (safe) and create fresh context + new_title = "" + new_working_directory = "" + new_pinned = False + try: + from code_puppy.api.db.queries import ( + get_session_metadata as _gsm, + ) + + _sm = await _gsm(session_id) or {} + new_title = _sm.get("title", "") + new_working_directory = _sm.get( + "working_directory", "" + ) + new_pinned = bool(_sm.get("pinned", False)) + except Exception: + pass + + ctx = await session_manager.create_session(session_id) + ctx.title = new_title + ctx.working_directory = new_working_directory + ctx.pinned = new_pinned + session_title = new_title + session_working_directory = new_working_directory + session_pinned = new_pinned + + # Mark the new session as active + await session_manager.mark_session_active(session_id) + + agent = ctx.agent + agent_name = ctx.agent_name + model_name = ctx.model_name + message_count = len(ctx.agent.get_message_history() or []) + logger.debug( + f"Restored {message_count} messages to session agent" + ) + + await send_typed( + ServerSessionSwitched( + session_id=new_session_id, + message_count=message_count, + title=session_title, + working_directory=session_working_directory, + created=False, + agent_name=agent_name, + model_name=model_name, + ) + ) + + logger.debug( + f"Switched to session {new_session_id} with {message_count} messages" + ) + + except Exception as e: + logger.error("Error switching session: %s", e) + await send_typed( + ServerError( + error=f"Failed to switch to session {new_session_id}: {str(e)}", + session_id=session_id, + ) + ) + + elif msg.get("type") == "set_working_directory": + new_directory = msg.get("directory", "") + if new_directory: + # Expand ~ and resolve symlinks before validation + new_directory = str( + Path(new_directory).expanduser().resolve() + ) + logger.info( + "[CWD DEBUG] set_working_directory received: new=%r, current=%r, session=%s", + new_directory, + session_working_directory, + session_id, + ) + # Validate the directory exists + if Path(new_directory).is_dir(): + # Skip if directory hasn't actually changed (avoids duplicate banners on reload) + if new_directory == session_working_directory: + logger.info( + "[CWD DEBUG] Skipping unchanged directory: %r, session=%s", + new_directory, + session_id, + ) + await send_typed( + ServerWorkingDirectoryChanged( + directory=new_directory, + success=True, + session_id=session_id, + unchanged=True, + ) + ) + continue # Skip to next message, don't write banner + + session_working_directory = new_directory + logger.debug( + "Working directory set to: %s", + session_working_directory, + ) + + # NOTE: Do NOT append raw dict entries into agent + # message history here. The runtime expects typed + # ModelMessage objects; dict injection can corrupt + # subsequent turns and lead to result=None failures. + + # Persist directory banner to SQLite so FE cold-load sees it + try: + # Bug 4: use ctx.agent_name / ctx.model_name directly + # so this works even when ctx.agent is None. + _cwd_agent = ( + ctx.agent_name if ctx else None + ) or "code-puppy" + _cwd_model = ( + ctx.model_name if ctx else None + ) or "unknown" + _cwd_segs = 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 _gpn + + _pname = _gpn() or "puppy" + logger.info( + "[CWD DEBUG] Writing CWD banner to SQLite: path=%r, session=%s", + session_working_directory, + session_id, + ) + await write_system_message_to_sqlite( + session_id=session_id, + system_message_type="directory", + content=f"{_pname} is now at {_cwd_rel}", + system_message_path=session_working_directory, + agent_name=_cwd_agent, + model_name=_cwd_model, + ) + # Bug 5: also stamp working_directory on the sessions row + _now_cwd = datetime.datetime.now( + datetime.timezone.utc + ).isoformat() + await update_session_working_directory( + session_id=session_id, + working_directory=session_working_directory, + updated_at=_now_cwd, + ) + except Exception as _cwd_exc: + logger.warning( + "CWD SQLite write failed: %s", + _cwd_exc, + exc_info=True, + ) + + await send_typed( + ServerWorkingDirectoryChanged( + directory=session_working_directory, + success=True, + session_id=session_id, + ) + ) + else: + await send_typed( + ServerWorkingDirectoryChanged( + directory=new_directory, + success=False, + error="Directory does not exist", + session_id=session_id, + ) + ) + else: + await send_typed( + ServerError( + error="No directory provided for set_working_directory", + session_id=session_id, + ) + ) + + elif msg.get("type") == "update_session_meta": + # Update session metadata (pinned, title) — persisted to SQLite. + try: + # Update in-memory state (local vars AND SessionContext) + if "pinned" in msg and isinstance(msg["pinned"], bool): + session_pinned = msg["pinned"] + if ctx is not None: + ctx.pinned = session_pinned + if "title" in msg and isinstance(msg["title"], str): + session_title = msg["title"] + if ctx is not None: + ctx.title = session_title + + # Persist to SQLite (single source of truth). + # Use a targeted UPDATE — NOT upsert_session() — so that + # message_count, total_tokens, and other stats are + # never zeroed out by a rename/pin operation. + 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=session_id, + title=session_title, + pinned=session_pinned, + updated_at=_dt.now(_tz.utc).isoformat(), + ) + logger.debug( + f"Updated session meta in SQLite for {session_id}: pinned={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=session_id, + pinned=session_pinned, + title=session_title, + ) + ) + except Exception as e: + logger.error("Error updating session meta: %s", e) + await send_typed( + ServerError( + error=f"Failed to update session metadata: {str(e)}", + session_id=session_id, + ) + ) + + elif msg.get("type") == "get_config": + 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=session_id, + ) + ) + else: + await send_typed( + ServerError( + error="No key provided for get_config", + session_id=session_id, + ) + ) + + elif msg.get("type") == "set_config": + 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( + f"Config set: {config_key} = {config_value}" + ) + await send_typed( + ServerConfigValue( + key=config_key, + value=config_value, + success=True, + session_id=session_id, + ) + ) + except Exception as e: + await send_typed( + ServerError( + error=f"Failed to set config: {e}", + session_id=session_id, + ) + ) + else: + await send_typed( + ServerError( + error="No key provided for set_config", + session_id=session_id, + ) + ) + + # Handle slash command execution + elif msg.get("type") == "command": + command_str = msg.get("command", "") + logger.debug("Command requested: %s", command_str) + + try: + from io import StringIO + + from rich.console import Console + + from code_puppy.command_line.command_handler import ( + get_commands_help, + handle_command, + ) + + output = None + captured_messages = [] + + # Special handling for /help - we can get the text directly + cmd_name = ( + command_str.strip().lstrip("/").split()[0] + if command_str.strip() + else "" + ) + if cmd_name in ("help", "h"): + # Get help text directly + help_text = get_commands_help() + # Render to plain text + 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: + # For other commands, just execute and report success/failure + # The actual output goes to the message queue/terminal + result = handle_command(command_str) + # Some commands might return a string as output + if isinstance(result, str): + output = result + result = True + + # Send command result + await send_typed( + ServerCommandResult( + command=command_str, + success=result is True or result is not False, + output=output, + messages=captured_messages, + result=str(result) + if result and result is not True + else None, + session_id=session_id, + ) + ) + logger.debug( + f"Command executed: {command_str} -> success={result is True or result is not False}, output_len={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, + ) + ) + continue + + # Handle cancel/interrupt request + elif msg.get("type") == "cancel": + logger.debug( + "Cancel request received - stopping active streaming and agent task" + ) + + # Cancel any active streaming + if active_drain_task and not active_drain_task.done(): + stop_draining.set() # Signal drain to stop + active_drain_task.cancel() + try: + await active_drain_task + except asyncio.CancelledError: + pass + stop_draining.clear() # Reset for next streaming + logger.debug("Active streaming cancelled") + + # Cancel the in-flight agent run (run_with_mcp) if present + if active_agent_task and not active_agent_task.done(): + logger.debug( + "Cancelling active agent task due to user interrupt" + ) + active_agent_task.cancel() + try: + await active_agent_task + except asyncio.CancelledError: + logger.debug("Active agent task cancelled successfully") + active_agent_task = None + + # Send confirmation + await send_typed( + ServerStatus( + status="cancelled", + session_id=session_id, + ) + ) + continue + + elif msg.get("type") == "permission_response": + # Handle permission response from user + from code_puppy.api.permissions import ( + handle_permission_response, + ) + + request_id = msg.get("request_id") + approved = msg.get("approved", False) + + if request_id: + handled = handle_permission_response(request_id, approved) + if handled: + logger.debug( + f"[Permission] ✅ Handled response: {request_id} = {approved}" + ) + else: + logger.warning( + f"[Permission] ❌ Unknown request: {request_id}" + ) + else: + logger.error( + "[WebSocket] ❌ No request_id in permission_response!" + ) + continue + + elif msg.get("type") == "message": + # Set WebSocket context for permission requests + from code_puppy.api.permission_plugin import ( + set_suppress_emitter_tool_events, + set_websocket_context, + ) + + set_websocket_context(websocket, 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 + event_queue = None + collected_text = [] + stop_draining.clear() # Reset for this message + drain_task = None + active_parts: dict[ + int, dict + ] = {} # Track message parts by index + tool_id_aliases: dict[str, str] = {} + tool_group_ids: dict[ + str, str + ] = {} # tool_id -> tool_group_id mapping + b1_streaming_used = ( + False # Track if B1 streaming sent any content + ) + agent_error = None # Will hold exception if agent run fails + + async def drain_events_concurrent( + ready_event: asyncio.Event = None, + ): + """Background task to drain events and send structured messages in real-time.""" + nonlocal collected_text, b1_streaming_used, ws_closed + 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" + ) + current_tool_name = None # Track active tool name + current_tool_group_id: str | None = ( + None # Track current tool batch group ID + ) + + # Track pending tool calls for tool_result generation + # {tool_id: {tool_name, start_time, part_index}} + pending_tool_calls: dict[str, dict] = {} + + def resolve_pending_tool_id( + *, + 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 pending_tool_calls + ): + return tool_call_id + + if tool_call_id: + for ( + _pending_id, + _pending_info, + ) in 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 pending_tool_calls.items(): + if ( + _pending_info.get("tool_name") + == tool_name + ): + return _pending_id + + return None + + def resolve_tool_group_id( + *, + tool_id: str | None = None, + pending_info: dict | 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. + + Priority order: + 1) pending tool info + 2) tool_id -> group map + 3) explicit fallback passed by caller + 4) current in-flight batch group + 5) synthetic deterministic-ish fallback (last resort) + """ + nonlocal current_tool_group_id + + 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 = 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 = 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: + tool_group_ids[tool_id] = group_id + if pending_info is not None: + pending_info["tool_group_id"] = group_id + elif tool_id in pending_tool_calls: + pending_tool_calls[tool_id][ + "tool_group_id" + ] = group_id + + if current_tool_group_id is None: + current_tool_group_id = group_id + + return group_id + + event_count = 0 + first_iteration = True + while not stop_draining.is_set(): + # Exit if WebSocket is closed + if 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( + 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 event_queue.empty(): + try: + event = 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 + + # 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": + tool_name = event_data.get( + "tool_name", "unknown" + ) + tool_args = event_data.get( + "tool_args", {} + ) + tool_id = str(uuid.uuid4())[:8] + + # Generate tool group ID for this batch if not already set + if current_tool_group_id is None: + current_tool_group_id = ( + f"tg-{str(uuid.uuid4())[:8]}" + ) + + logger.debug( + "[ws] tool_call: %s", tool_name + ) + + # Track current tool name + current_tool_name = tool_name + + # Register in pending_tool_calls so + # tool_call_complete can echo back the + # same tool_id the frontend already knows. + pending_tool_calls[tool_id] = { + "tool_name": tool_name, + "start_time": time_module.time(), + "tool_group_id": 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=current_agent_name, + model_name=current_model_name, + tool_group_id=current_tool_group_id, + ) + ) + + elif event_type == "tool_call_complete": + 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 + ) + + # Clear current tool name after completion + current_tool_name = None + + # Match the tool_id we generated at + # tool_call_start so the frontend can + # correlate result → call by ID, not + # just by name (names are not unique). + matching_tool_id = next( + ( + tid + for tid, info in pending_tool_calls.items() + if info["tool_name"] + == tool_name + ), + None, + ) + matching_pending_info = ( + pending_tool_calls.get( + matching_tool_id + ) + if matching_tool_id + else None + ) + tool_group_id_for_result = resolve_tool_group_id( + tool_id=matching_tool_id, + pending_info=matching_pending_info, + fallback_group_id=current_tool_group_id, + tool_name=tool_name, + source="tool_call_complete", + ) + + if matching_tool_id: + 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=current_agent_name, + model_name=current_model_name, + tool_group_id=tool_group_id_for_result, + ) + ) + + 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", {} + ) + initial_content = "" + + if ( + hasattr(part_obj, "content") + and part_obj.content + ): + initial_content = ( + part_obj.content + ) + elif isinstance( + part_obj, dict + ) and part_obj.get("content"): + initial_content = part_obj.get( + "content", "" + ) + + if part_type in ( + "TextPart", + "ThinkingPart", + ): + msg_type = ( + "thinking" + if part_type + == "ThinkingPart" + else "text" + ) + + # When a TextPart starts, any pending tool calls have completed + # Send status-only tool_result during streaming to update UI + # The actual result data will be sent later from extraction code + if ( + pending_tool_calls + and part_type == "TextPart" + ): + for ( + tool_id, + tool_info, + ) in list( + pending_tool_calls.items() + ): + if tool_info.get( + "status_only_sent" + ): + continue + duration_ms = ( + time_module.time() + - tool_info[ + "start_time" + ] + ) * 1000 + logger.debug( + f"[WebSocket] Sending status-only tool_result for: {tool_info['tool_name']} (duration: {duration_ms:.1f}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=current_agent_name, + model_name=current_model_name, + tool_group_id=resolve_tool_group_id( + tool_id=tool_id, + pending_info=tool_info, + fallback_group_id=current_tool_group_id, + tool_name=tool_info.get( + "tool_name" + ), + source="status_only_result", + ), + ) + ) + tool_info[ + "status_only_sent" + ] = True + current_tool_name = None + + # Check if part already exists (created by early delta) + if part_index in active_parts: + # Just update the type, keep existing message_id + active_parts[part_index][ + "type" + ] = msg_type + message_id = active_parts[ + part_index + ]["id"] + # If there's initial content, add it to the existing accumulated content + if initial_content: + active_parts[ + part_index + ]["content"] = ( + initial_content + + active_parts[ + part_index + ]["content"] + ) + collected_text.insert( + 0, initial_content + ) # Prepend to collected text too + logger.debug( + f"[Stream Debug] Part already exists, reusing message_id={message_id}" + ) + # Don't send another start event + continue + + message_id = f"msg-{int(time_module.time() * 1000)}-{part_index}" + active_parts[part_index] = { + "id": message_id, + "type": msg_type, + "content": initial_content, # Start with initial content if present + } + + # If there's initial content, add it to collected_text + if initial_content: + collected_text.append( + initial_content + ) + + # New assistant output part indicates turn boundary for tool grouping + if part_index == 0: + 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=current_agent_name, + model_name=current_model_name, + tool_name=current_tool_name, + ).model_dump( + exclude_none=True + ) + ) + + # If there's initial content, send it as a delta immediately + 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=current_agent_name, + model_name=current_model_name, + tool_name=current_tool_name, + ) + await safe_send_json( + _delta.model_dump( + exclude_none=True + ) + ) + # Mark that B1 streaming was used (nonlocal already declared above) + b1_streaming_used = True + + elif part_type == "ToolCallPart": + # Extract tool info from the part + tool_name = "unknown" + tool_call_id = None + tool_args_str = "" + + # Extract tool info from part_obj (dict or object) + if hasattr( + part_obj, "tool_name" + ): + tool_name = ( + part_obj.tool_name + ) + elif isinstance(part_obj, dict): + tool_name = part_obj.get( + "tool_name", "unknown" + ) + + if hasattr( + part_obj, "tool_call_id" + ): + tool_call_id = ( + part_obj.tool_call_id + ) + elif isinstance(part_obj, dict): + tool_call_id = part_obj.get( + "tool_call_id" + ) + + if hasattr(part_obj, "args"): + tool_args_str = ( + part_obj.args or "" + ) + elif isinstance(part_obj, dict): + tool_args_str = ( + part_obj.get("args", "") + or "" + ) + + tool_id = ( + tool_call_id + or str(uuid.uuid4())[:8] + ) + + # Track this tool call part with args buffer for accumulation + # Don't send tool_call yet - wait for part_end when args are complete + 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, # Buffer for accumulating args_delta + "start_time": time_module.time(), # Track when tool started + } + + # Track current tool name + current_tool_name = tool_name + + logger.debug( + f"[WebSocket] ToolCallPart started: {tool_name} (id: {tool_id})" + ) + # tool_call event will be sent on part_end when args are complete + + elif part_type == "ToolReturnPart": + logger.info( + f"[WebSocket] ToolReturnPart detected! part_index={part_index}" + ) + # Extract tool result info from the part + tool_call_id = None + tool_content = None + + # Extract tool_call_id + if hasattr( + part_obj, "tool_call_id" + ): + tool_call_id = ( + part_obj.tool_call_id + ) + elif isinstance(part_obj, dict): + tool_call_id = part_obj.get( + "tool_call_id" + ) + + # Extract content (the tool result) + if hasattr(part_obj, "content"): + tool_content = ( + part_obj.content + ) + elif isinstance(part_obj, dict): + tool_content = part_obj.get( + "content" + ) + + # Try to serialize complex result objects + if ( + tool_content + and not isinstance( + tool_content, + ( + str, + dict, + list, + int, + float, + bool, + type(None), + ), + ) + ): + try: + import json + + # Try to convert to dict first (for Pydantic models) + if hasattr( + tool_content, + "model_dump", + ): + tool_content = tool_content.model_dump() + elif hasattr( + tool_content, "dict" + ): + tool_content = tool_content.dict() + elif hasattr( + tool_content, + "__dict__", + ): + tool_content = tool_content.__dict__ + else: + tool_content = str( + tool_content + ) + except Exception as e: + logger.debug( + f"[WebSocket] Could not serialize tool result: {e}" + ) + tool_content = str( + tool_content + ) + + # Find the corresponding pending tool call, store and SEND the result + _result_sent = False + if tool_call_id: + resolved_pending_id = resolve_pending_tool_id( + tool_call_id=tool_call_id + ) + if resolved_pending_id: + _result_sent = True + pending_info = 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( + tool_id=resolved_pending_id, + pending_info=pending_info, + fallback_group_id=current_tool_group_id, + tool_name=_tool_name, + source="tool_return_resolved", + ) + logger.info( + f"[WebSocket] ToolReturnPart: Sending result for {_tool_name} (id: {resolved_pending_id}, raw: {tool_call_id})" + ) + # Send the REAL tool_result with actual content + 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=current_agent_name + or "code-puppy", + model_name=current_model_name + or "unknown", + tool_group_id=_group_id, + ) + ) + else: + logger.warning( + f"[WebSocket] ToolReturnPart: Could not resolve tool_call_id {tool_call_id}, pending keys: {list(pending_tool_calls.keys())}" + ) + + # Fallback: proximity-based matching if no tool_call_id or resolution failed + if not _result_sent: + for ( + pending_id, + pending_info, + ) in sorted( + pending_tool_calls.items(), + key=lambda x: abs( + x[1].get( + "part_index", + 9999, + ) + - part_index + ), + ): + if ( + abs( + pending_info.get( + "part_index", + 9999, + ) + - part_index + ) + <= 3 + ): + # Close enough - assume it's the result for this tool call + 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( + tool_id=pending_id, + pending_info=pending_info, + fallback_group_id=current_tool_group_id, + tool_name=_tool_name, + source="tool_return_proximity", + ) + logger.info( + f"[WebSocket] ToolReturnPart: Sending result (by proximity) for {_tool_name} (id: {pending_id})" + ) + # Send the REAL tool_result with actual content + 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=current_agent_name + or "code-puppy", + model_name=current_model_name + or "unknown", + tool_group_id=_group_id, + ) + ) + _result_sent = True + break + + # Warn if we couldn't send the result + if not _result_sent: + logger.warning( + f"[WebSocket] ToolReturnPart: Could NOT send result! tool_call_id={tool_call_id}, pending_tool_calls={list(pending_tool_calls.keys())}, part_index={part_index}" + ) + + # Track this part for cleanup + active_parts[part_index] = { + "id": f"tool-return-{part_index}", + "type": "tool_return", + "tool_call_id": tool_call_id, + "content": tool_content, + } + + 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", {} + ) + + # Handle ToolCallPartDelta - accumulate args + if ( + delta_type + == "ToolCallPartDelta" + ): + args_delta = "" + if hasattr( + delta_obj, "args_delta" + ): + args_delta = ( + delta_obj.args_delta + or "" + ) + elif isinstance( + delta_obj, dict + ): + args_delta = ( + delta_obj.get( + "args_delta", "" + ) + or "" + ) + + if ( + args_delta + and part_index + in active_parts + ): + part_info = active_parts[ + part_index + ] + if ( + part_info.get("type") + == "tool_call" + ): + part_info[ + "args_buffer" + ] = ( + part_info.get( + "args_buffer", + "", + ) + + args_delta + ) + continue # Don't process as text delta + + # Extract content_delta (supports both direct and nested formats) + content_delta = inner_data.get( + "content_delta", "" + ) + if not content_delta: + if isinstance(delta_obj, dict): + content_delta = ( + delta_obj.get( + "content_delta", "" + ) + ) + + if content_delta: + collected_text.append( + content_delta + ) + + # Ensure we have an entry for this part (handles delta before start) + if ( + part_index + not in active_parts + ): + # Create entry with consistent message_id + message_id = f"msg-{int(time_module.time() * 1000)}-{part_index}" + active_parts[part_index] = { + "id": message_id, + "type": "text", # Default, will be updated by part_start if it arrives + "content": "", + } + # Send a start event so client creates the message + logger.debug( + f"[Stream Debug] Creating part on first delta: message_id={message_id}" + ) + if part_index == 0: + 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=current_agent_name, + model_name=current_model_name, + tool_name=current_tool_name, + ).model_dump( + exclude_none=True + ) + ) + + part_info = active_parts[ + part_index + ] + message_id = part_info["id"] + + if part_index in active_parts: + active_parts[part_index][ + "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=current_agent_name, + model_name=current_model_name, + tool_name=current_tool_name, + ) + await safe_send_json( + _delta.model_dump( + exclude_none=True + ) + ) + # Mark that B1 streaming was used + b1_streaming_used = True + + elif inner_type == "part_end": + part_index = inner_data.get( + "index", 0 + ) + part_info = active_parts.get( + part_index, {} + ) + part_type_info = part_info.get( + "type", "text" + ) + message_id = part_info.get( + "id", f"msg-{part_index}" + ) + full_content = part_info.get( + "content", "" + ) + + # Handle tool_call parts - send tool_call event now that args are complete + if part_type_info == "tool_call": + 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" + ) + ) + + # Parse args if they're a JSON string + try: + import json + + args_dict = ( + json.loads( + tool_args_str + ) + if tool_args_str + else {} + ) + except ( + json.JSONDecodeError, + TypeError, + ): + args_dict = {} + + logger.debug( + f"[WebSocket] Sending tool_call (args complete): {tool_name}" + ) + + # Generate tool group ID for this batch if not already set + if ( + current_tool_group_id + is None + ): + current_tool_group_id = f"tg-{str(uuid.uuid4())[:8]}" + + # Send tool_call now that we have complete args + 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=current_agent_name, + model_name=current_model_name, + tool_group_id=current_tool_group_id, + ) + ) + + # Track as pending tool call for later tool_result + 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, # Will be populated by ToolReturnPart + "tool_group_id": current_tool_group_id, + } + if raw_tool_call_id: + tool_id_aliases[ + raw_tool_call_id + ] = tool_id + # Also track tool_group_id for pre-stream_end extraction + if current_tool_group_id: + tool_group_ids[tool_id] = ( + current_tool_group_id + ) + + # Clear current tool name tracking + current_tool_name = None + else: + await safe_send_json( + ServerAssistantMessageEnd( + message_id=message_id, + part_index=part_index, + full_content=full_content, + timestamp=time_module.time(), + session_id=session_id, + agent_name=current_agent_name, + model_name=current_model_name, + tool_name=current_tool_name, + ).model_dump( + exclude_none=True + ) + ) + + if part_index in active_parts: + del active_parts[part_index] + + except Exception as send_err: + error_msg = str(send_err).lower() + if ( + "close message" in error_msg + or "closed" in error_msg + ): + 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 = event_queue.get_nowait() + event_type = event.get("type", "") + event_data = event.get("data", {}) + final_count += 1 + + if event_type == "stream_event": + inner_type = event_data.get( + "event_type", "" + ) + inner_data = event_data.get( + "event_data", {} + ) + if inner_type == "part_delta": + content_delta = inner_data.get( + "content_delta", "" + ) + if not content_delta: + delta_obj = inner_data.get( + "delta", {} + ) + if isinstance(delta_obj, dict): + content_delta = delta_obj.get( + "content_delta", "" + ) + if content_delta: + collected_text.append(content_delta) + 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}" + ) + + # Event to signal drain task is ready + drain_ready = asyncio.Event() + + try: + from code_puppy.plugins.frontend_emitter.emitter import ( + subscribe, + unsubscribe, + ) + + event_queue = subscribe(session_id=session_id) + logger.debug( + "Subscribed to frontend emitter for streaming" + ) + + # Modify drain function to signal when ready + async def drain_events_with_signal(): + """Wrapper that signals readiness before starting drain loop.""" + await drain_events_concurrent(drain_ready) + + # Start concurrent drain task + drain_task = asyncio.create_task( + drain_events_with_signal() + ) + active_drain_task = ( + drain_task # Track for potential cancellation + ) + + # Wait for drain task to be ready before proceeding + await drain_ready.wait() + + except ImportError: + logger.warning("Frontend emitter not available") + + 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, + ) + + # Inject working directory as a system message in the + # agent's conversation history (once per directory change) + # so the LLM knows the CWD without polluting user messages. + message_to_send = user_message + if ( + session_working_directory + and session_working_directory + != last_context_sent_directory + ): + from pydantic_ai.messages import ( + ModelRequest, + SystemPromptPart, + ) + + wd_system_msg = ModelRequest( + parts=[ + SystemPromptPart( + content=( + f"The user's current working directory is updated to" + f" {session_working_directory}" + ) + ) + ] + ) + agent.append_to_message_history(wd_system_msg) + last_context_sent_directory = ( + session_working_directory + ) + logger.debug( + "Injected working directory system message: %s", + session_working_directory, + ) + + logger.debug( + f"Calling run_with_mcp with message: {message_to_send[:100]}..." + ) + + # Build file context and binary attachments + file_context, binary_attachments = ( + build_file_context_and_attachments(msg) + ) + + # CHANGE 1: Update attachment metadata for UI (variables already initialized) + if msg.get("attachments"): + from pathlib import Path as _AttachmentPath + + for raw_path in msg.get("attachments", []): + if ( + isinstance(raw_path, str) + and raw_path.strip() + ): + try: + file_path = _AttachmentPath( + 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( + f"Error building attachment metadata for '{raw_path}': {e}" + ) + + # Prepend file context to the message + if file_context: + message_to_send = ( + file_context + "\n\n" + message_to_send + ) + logger.debug( + f"Added file context ({len(file_context)} chars)" + ) + + run_kwargs = {} + if binary_attachments: + run_kwargs["attachments"] = binary_attachments + logger.debug( + f"Including {len(binary_attachments)} binary attachment(s)" + ) + + # ────────────────────────────────────────────────────────────────── + # 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, + ) + + # Run the agent in its own task so it can be cancelled independently. + # Suppress duplicate emitter lifecycle callbacks only during the + # run_with_mcp execution window (pre/post_tool_call hooks fire there). + set_suppress_emitter_tool_events(True) + active_agent_task = asyncio.create_task( + agent.run_with_mcp( + message_to_send, **run_kwargs + ) + ) + + # ===== CONCURRENT MESSAGE PROCESSING FIX ===== + # Instead of awaiting the agent task here (which blocks the receive loop), + # we use asyncio.wait() to handle BOTH the agent task AND incoming messages. + # This allows permission_response messages to be processed while the agent runs. + + result = None + agent_completed = False + # Deferred message for switch_session/create_session during streaming + _deferred_msg: dict | None = None + + while not agent_completed: + # Create a task for the next message (or wait for agent) + receive_task = asyncio.create_task( + websocket.receive_json() + ) + + # Wait for either the agent to complete OR a new message to arrive + done, pending = await asyncio.wait( + {active_agent_task, receive_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + + # Check if agent completed + if active_agent_task in done: + try: + result = await active_agent_task + logger.debug( + f"run_with_mcp completed, result type: {type(result)}" + ) + agent_completed = True + active_agent_task = None + except asyncio.CancelledError: + logger.debug( + "run_with_mcp task was cancelled by user" + ) + 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, + ) + agent_error = e + result = None + agent_completed = True + active_agent_task = None + + # Cancel the receive task if it's still pending + if receive_task in pending: + receive_task.cancel() + try: + await receive_task + except asyncio.CancelledError: + pass + + # Check if a new message arrived + elif receive_task in done: + try: + new_msg = await receive_task + + # Advisory validation — log but never reject + try: + _parsed_inner = _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" + }, + ) + + # Handle permission_response immediately + if ( + 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 + ) + ) + if handled: + logger.debug( + f"[Permission] ✅ Handled response: {request_id} = {approved}" + ) + else: + logger.warning( + f"[Permission] ❌ Unknown request: {request_id}" + ) + else: + logger.error( + "[WebSocket] ❌ No request_id in permission_response!" + ) + + # Handle cancel request + 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 # Exit the loop + ) + + # Handle session switch during streaming + elif new_msg.get("type") in ( + "switch_session", + "create_session", + ): + # User switched to a new chat while agent was running. + # The agent will finish in background and save to SQLite. + logger.debug( + "[WS:%s] Session switch during streaming — " + "agent continues in background, switching to: %s", + session_id, + new_msg.get("session_id"), + ) + + # Fire and forget — background task owns the agent result. + 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 # disown — background task now owns it + agent_completed = ( + True # exit inner loop + ) + + # For other message types, log and ignore (or queue for later) + else: + logger.warning( + f"[WebSocket] Received {new_msg.get('type')} message while agent running - ignoring" + ) + + except asyncio.CancelledError: + # Receive was cancelled, continue + pass + except WebSocketDisconnect: + # WebSocket gone — let agent finish and save to SQLite + logger.debug( + "[WS:%s] Disconnect during streaming — agent continues in background", + session_id, + ) + + # Fire and forget — background task owns the agent result. + 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(): + # Starlette raises RuntimeError when calling receive after disconnect + logger.debug( + "[WS:%s] WebSocket already disconnected: %s — agent continues in background", + session_id, + e, + ) + + # Fire and forget — background task owns the agent result. + 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( + f"RuntimeError processing message during agent execution: {e}" + ) + except Exception as e: + logger.error( + f"Error processing message during agent execution: {e}" + ) + + # ===== END CONCURRENT MESSAGE PROCESSING FIX ===== + + finally: + # End suppression scope once run_with_mcp window ends, + # including completion/cancel/background handoff paths. + set_suppress_emitter_tool_events(False) + # Clear session context for prompt generation + clear_session_working_directory() + + finally: + # Stop the drain task + if drain_task: + stop_draining.set() + try: + await asyncio.wait_for(drain_task, timeout=2.0) + except asyncio.TimeoutError: + drain_task.cancel() + try: + await drain_task + except asyncio.CancelledError: + pass + + # Unsubscribe from emitter + if event_queue: + unsubscribe(event_queue) + logger.debug("Unsubscribed from frontend emitter") + + # 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 + + # If agent errored, send error to GUI and skip success path + if agent_error == "cancelled": + await send_typed( + ServerCancelled( + session_id=session_id, + ) + ) + continue + elif agent_error is not None: + logger.debug( + "[WS:%s] agent_error -> sending frame(s) to client. type=%s", + session_id, + type(agent_error).__name__, + ) + # If streaming was already in progress, send stream_end first so + # the frontend exits its streaming state before the error frame. + # Without this handshake the UI hangs indefinitely. + for frame in build_error_response_frames( + agent_error, collected_text, session_id + ): + await safe_send_json(frame) + continue + + # Safety net: if the agent task completed without raising but returned + # no result and produced no streamed text, treat it as an error. + has_nonempty_stream = any( + (chunk or "").strip() for chunk in 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(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." + ) + ) + await send_typed( + 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, + ) + ) + continue + + # Extract final response text + response_text = "" + + # Priority 1: Use collected text from streaming events + if has_nonempty_stream: + response_text = "".join(collected_text) + logger.debug( + f"Using collected streaming text ({len(response_text)} chars)" + ) + # Priority 2: Extract from result.output (AgentRunResult) or result.data (legacy) + 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)" + ) + # Priority 3: Get last assistant message from history + elif agent: + messages = agent.get_message_history() + + for msg in reversed(messages): + if hasattr(msg, "role") and msg.role == "assistant": + if 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" + + # Extract token usage from result if available + tokens_used = None + if result: + # Try to get usage from 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 + ), + } + # Also try _usage or other common patterns + 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 we couldn't get usage from result, estimate from message history + 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 + + # Only send legacy 'response' if B1 streaming wasn't used + # B1 streaming already sent the content via assistant_message_end + logger.warning( + "[WebSocket] b1_streaming_used=%s before response/extraction", + b1_streaming_used, + ) + if not b1_streaming_used: + # For non-streaming models (like Gemini), extract and send thinking content first + thinking_text = "" + if agent: + try: + history = agent.get_message_history() + logger.debug( + f"[Thinking Debug] Checking history for thinking parts, {len(history) if history else 0} messages" + ) + if history: + # Look for ThinkingPart in the last ModelResponse + for i, msg in enumerate(reversed(history)): + msg_type = type(msg).__name__ + logger.debug( + f"[Thinking Debug] Message {i}: type={msg_type}, has_parts={hasattr(msg, 'parts')}" + ) + if "Response" in msg_type and hasattr( + msg, "parts" + ): + logger.debug( + f"[Thinking Debug] Found Response with {len(msg.parts)} 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( + f"[Thinking Debug] Part {j}: type={part_type}, content_preview={part_content_preview}" + ) + if ( + "Thinking" in part_type + and hasattr(part, "content") + ): + thinking_text = part.content + logger.debug( + f"[Thinking Debug] Found thinking content: {len(thinking_text)} chars" + ) + 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" + ) + + # 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" + ) + + # Send the main response + await send_typed( + ServerResponse( + content=response_text, + done=True, + 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, + ) + ) + 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: set = set() + logger.warning( + "[WebSocket] B1 Pre-stream_end extraction: result=%s, has_all_messages=%s", + type(result).__name__ if result else None, + hasattr(result, "all_messages") + if result + else False, + ) + if result and hasattr(result, "all_messages"): + try: + import time as _time_pre + + from pydantic_ai.messages import ( + ToolReturn, + ToolReturnPart, + ) + + _pre_msgs = list(result.all_messages()) + for _pre_msg in _pre_msgs: + if not hasattr(_pre_msg, "parts"): + continue + for _pre_part in _pre_msg.parts: + if not isinstance( + _pre_part, + (ToolReturnPart, ToolReturn), + ): + continue + _pre_tool_name = getattr( + _pre_part, "tool_name", "unknown" + ) + _pre_raw_tool_id = getattr( + _pre_part, "tool_call_id", None + ) + _pre_tool_id = ( + tool_id_aliases.get( + _pre_raw_tool_id, + _pre_raw_tool_id, + ) + if _pre_raw_tool_id + else "unknown" + ) + _pre_result = getattr( + _pre_part, "content", None + ) + # Serialize complex objects + if _pre_result and not isinstance( + _pre_result, + ( + str, + dict, + list, + int, + float, + bool, + type(None), + ), + ): + try: + if hasattr( + _pre_result, "model_dump" + ): + _pre_result = ( + _pre_result.model_dump() + ) + elif hasattr( + _pre_result, "dict" + ): + _pre_result = ( + _pre_result.dict() + ) + elif hasattr( + _pre_result, "__dict__" + ): + _pre_result = ( + _pre_result.__dict__ + ) + else: + _pre_result = str( + _pre_result + ) + except Exception: + _pre_result = str(_pre_result) + logger.warning( + "[WebSocket] Pre-stream_end tool result: %s (id: %s), content_preview=%s", + _pre_tool_name, + _pre_tool_id, + str(_pre_result)[:100] + if _pre_result + else None, + ) + _pre_group_id = tool_group_ids.get( + _pre_tool_id + ) + + await send_typed_tool_lifecycle( + ServerToolResult( + tool_id=_pre_tool_id, + tool_name=_pre_tool_name, + result=_pre_result, + success=True, + duration_ms=0, + timestamp=_time_pre.time(), + session_id=session_id, + agent_name=agent.name + if agent + else "code-puppy", + model_name=agent.get_model_name() + if agent + else "unknown", + tool_group_id=_pre_group_id, + ) + ) + _pre_sent_tool_ids.add(_pre_tool_id) + except Exception as _pre_e: + logger.warning( + "Pre-stream_end tool result extraction failed: %s", + _pre_e, + ) + + # 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 + # Track tool IDs already sent before stream_end to skip duplicates later + if "_pre_sent_tool_ids" not in dir(): + _pre_sent_tool_ids: set = set() + _save_history_snapshot = [] # Snapshot of history before await points + try: + # CRITICAL: Update message history from result to include final response + # The pydantic-ai result.all_messages() contains the complete conversation + # including the final assistant response that may not be in agent's history yet + if result and hasattr(result, "all_messages"): + try: + all_msgs = list(result.all_messages()) + if all_msgs: + agent.set_message_history(all_msgs) + _save_history_snapshot = list( + agent.get_message_history() + ) # Snapshot before any awaits can corrupt shared global state + logger.debug( + f"Updated message history from result.all_messages(): {len(all_msgs)} messages" + ) + + # Extract and send tool results from message history + try: + import time as _time + + from pydantic_ai.messages import ( + ToolReturn, + ToolReturnPart, + ) + + for msg in all_msgs: + if hasattr(msg, "parts"): + for part in msg.parts: + if isinstance( + part, + ( + ToolReturnPart, + ToolReturn, + ), + ): + tool_name = getattr( + part, + "tool_name", + "unknown", + ) + tool_call_id = getattr( + part, + "tool_call_id", + "unknown", + ) + # Serialize the result + result_data = getattr( + part, "content", None + ) + if ( + result_data + and not isinstance( + result_data, + ( + str, + dict, + list, + int, + float, + bool, + type(None), + ), + ) + ): + try: + if hasattr( + result_data, + "model_dump", + ): + result_data = result_data.model_dump() + elif hasattr( + result_data, + "dict", + ): + result_data = result_data.dict() + elif hasattr( + result_data, + "__dict__", + ): + result_data = result_data.__dict__ + else: + result_data = str( + result_data + ) + except Exception: + result_data = str( + result_data + ) + + # Log for shell commands + if ( + tool_name + == "agent_run_shell_command" + ): + stdout_val = ( + result_data.get( + "stdout", "N/A" + ) + if isinstance( + result_data, + dict, + ) + else "not dict" + ) + logger.debug( + "Extracted shell result: id=%s, stdout=%s", + tool_call_id, + stdout_val, + ) + + # Skip tool IDs already sent before stream_end + if ( + tool_call_id + in _pre_sent_tool_ids + ): + logger.debug( + "[WebSocket] Skipping duplicate tool result (pre-sent): %s", + tool_call_id, + ) + continue + logger.info( + f"[WebSocket] Sending extracted tool result for {tool_name} (id: {tool_call_id})" + ) + _post_group_id = ( + tool_group_ids.get( + tool_call_id + ) + ) + + await send_typed( + ServerToolResult( + tool_id=tool_call_id, + tool_name=tool_name, + result=result_data, + success=True, + duration_ms=0, + timestamp=_time.time(), + session_id=session_id, + agent_name=agent.name + if agent + else "code-puppy", + model_name=agent.get_model_name() + if agent + else "unknown", + tool_group_id=_post_group_id, + ) + ) + except Exception as e: + logger.warning( + f"Could not extract tool results from messages: {e}" + ) + + except Exception as e: + logger.warning( + f"Could not update history from result.all_messages(): {e}" + ) + + history = ( + _save_history_snapshot + if _save_history_snapshot + else agent.get_message_history() + ) # Use pre-await snapshot to avoid race condition + if history: + # Regenerate title if it's empty or still the default "untitled-session" + if ( + not session_title + or session_title == "untitled-session" + ): + session_title = generate_heuristic_title( + history + ) + + # Session name is just the ID (WS_session_timestamp format) + # Title is stored only in the meta file, not in the filename + session_name = session_id + + # Wrap each message with metadata for complete session information. + # Chain `or` fallbacks: agent.name can be "" before the agent + # object is fully configured, so we cascade to context defaults. + agent_name_meta = ( + (agent.name if agent else "") + or ctx.agent_name + or "code-puppy" + ) + model_name_meta = ( + (agent.get_model_name() if agent else "") + or ctx.model_name + or "unknown" + ) + + def _extract_message_timestamp( + raw_msg: Any, default_ts: str + ) -> str: + """Best-effort extraction of an existing timestamp for a message. + + This is important for WebSocket sessions where history may + already contain older messages. We don't want to overwrite + their original timestamps every time we auto-save. + + Precedence: + 1. If the message is a dict with a numeric epoch 'timestamp', + convert to ISO. + 2. If the message is a dict with an ISO-ish 'timestamp' str, + reuse it as-is. + 3. If the message is a dict with 'ts', reuse it. + 4. If the message object has a 'timestamp' attribute, try that. + 5. Fall back to the provided default_ts ("now" for new messages). + """ + # Dict-based histories (e.g. CLI format or older WS formats) + if isinstance(raw_msg, dict): + ts_val = raw_msg.get("timestamp") + + # Epoch seconds + if isinstance(ts_val, (int, float)): + try: + return ( + datetime.datetime.fromtimestamp( + ts_val + ).isoformat() + ) + except Exception: + pass + + # Already an ISO string + if isinstance(ts_val, str) and ts_val: + return ts_val + + # Our enhanced WS wrapper sometimes uses 'ts' + ts_field = raw_msg.get("ts") + if isinstance(ts_field, str) and ts_field: + return ts_field + + # Pydantic / custom objects that carry a timestamp attribute + 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 + + # Fallback: use provided default + return default_ts + + # Create enhanced history: list of dicts with message + metadata + # Each entry: {'msg': , 'agent': str, 'model': str, 'ts': str} + enhanced_history = [] + for idx, msg in enumerate(history): + # Check if already wrapped (for idempotency). If the wrapper + # already has a 'ts', leave it untouched so older sessions + # keep their original per-message timestamps. + if ( + isinstance(msg, dict) + and "msg" in msg + and "agent" in msg + ): + enhanced_history.append(msg) + else: + # Use existing timestamp if present; otherwise compute a default now() + current_timestamp = ( + datetime.datetime.now().isoformat() + ) + msg_ts = _extract_message_timestamp( + msg, current_timestamp + ) + wrapper = { + "msg": msg, + "agent": agent_name_meta, + "model": model_name_meta, + "ts": msg_ts, + } + + # Add clean_content and attachments to the user message we just processed. + # clean_content is only needed when attachments were injected into content + # (file blocks like --- File: auth.ts ---). Without attachments, content + # is already the user's words (FE strips [Session Context:] as fallback). + is_user_message_just_processed = ( + idx == len(history) - 2 + and len(history) >= 2 + and attachment_metadata # only needed when file content was injected + ) + + if is_user_message_just_processed: + wrapper["clean_content"] = ( + original_user_message + ) + wrapper["attachments"] = ( + attachment_metadata + ) + logger.debug( + "Added UI metadata to user message: %d attachment(s), " + "clean_content length: %d", + len(attachment_metadata), + len(original_user_message), + ) + + enhanced_history.append(wrapper) + + message_count = len(enhanced_history) + 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 + + # Write to SQLite for FE read path + try: + import datetime as _dt_mod + + _now_iso = _dt_mod.datetime.now( + _dt_mod.timezone.utc + ).isoformat() + await write_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, + updated_at=_now_iso, + created_at=ctx.created_at.isoformat(), + ctx=ctx, + ) + except Exception as _db_exc: + logger.debug( + "SQLite turn write skipped (DB not available): %s", + _db_exc, + ) + + # Send session metadata update to client + await websocket.send_json( + { + "type": "session_meta", + "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, + } + ) + + # Broadcast session update to session monitoring clients + session_update_data = { + "session_id": session_id, + "session_name": session_name, + "title": session_title, + "working_directory": session_working_directory, + "timestamp": datetime.datetime.now().isoformat(), + "message_count": message_count, + "total_tokens": total_tokens, + "auto_saved": True, + "pickle_path": "", + "metadata_path": "", + "action": "created" + if message_count == 1 + else "updated", + } + await connection_manager.broadcast_session_update( + session_update_data + ) + 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) + + 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 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: + 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) + except Exception: + pass diff --git a/code_puppy/api/ws/legacy/connection_manager.py b/code_puppy/api/ws/legacy/connection_manager.py new file mode 100644 index 000000000..eec49804d --- /dev/null +++ b/code_puppy/api/ws/legacy/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/legacy/events_handler.py b/code_puppy/api/ws/legacy/events_handler.py new file mode 100644 index 000000000..b9bf647ec --- /dev/null +++ b/code_puppy/api/ws/legacy/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/legacy/health_handler.py b/code_puppy/api/ws/legacy/health_handler.py new file mode 100644 index 000000000..832f45e95 --- /dev/null +++ b/code_puppy/api/ws/legacy/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/legacy/runtime/__init__.py b/code_puppy/api/ws/legacy/runtime/__init__.py new file mode 100644 index 000000000..750f74d72 --- /dev/null +++ b/code_puppy/api/ws/legacy/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/legacy/runtime/session_runtime.py b/code_puppy/api/ws/legacy/runtime/session_runtime.py new file mode 100644 index 000000000..a6fda2741 --- /dev/null +++ b/code_puppy/api/ws/legacy/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/legacy/runtime/session_runtime_manager.py b/code_puppy/api/ws/legacy/runtime/session_runtime_manager.py new file mode 100644 index 000000000..82a8252d3 --- /dev/null +++ b/code_puppy/api/ws/legacy/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/legacy/schemas.py b/code_puppy/api/ws/legacy/schemas.py new file mode 100644 index 000000000..05e56528d --- /dev/null +++ b/code_puppy/api/ws/legacy/schemas.py @@ -0,0 +1,603 @@ +"""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, 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" + + +class ServerMessageType(str, Enum): + """All ``type`` values the server may send over the WebSocket.""" + + SYSTEM = "system" + 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" + 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 + + +# ═══════════════════════════════════════════════════════════════════════════════ +# SERVER → CLIENT (21 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. + """ + + type: Literal["assistant_message_end"] = "assistant_message_end" + message_id: str + 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 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")], + ], + 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[ServerCancelled, Tag("cancelled")], + ], + Discriminator("type"), +] +"""Discriminated union of all server→client message types.""" diff --git a/code_puppy/api/ws/legacy/sessions_handler.py b/code_puppy/api/ws/legacy/sessions_handler.py new file mode 100644 index 000000000..f0719648c --- /dev/null +++ b/code_puppy/api/ws/legacy/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/docs/migration/puppy-desk-gate1-investigation-summary.md b/docs/migration/puppy-desk-gate1-investigation-summary.md new file mode 100644 index 000000000..90fe9dd5b --- /dev/null +++ b/docs/migration/puppy-desk-gate1-investigation-summary.md @@ -0,0 +1,426 @@ +# Puppy Desk Migration — Gate 1 Investigation Summary + +Saved to avoid losing detailed context during summarization. + +## User Requirements and Constraints + +- Migrate divergent source branch `feature/puppy-desk` from `../../../../dev/puppy/code-puppy` into the current repo. +- Preserve features under `code_puppy/api/ws`, including: + - WebSocket session handling + - event streaming + - tool args/results + - streaming deltas + - strict message ordering + - permission handling + - CWD setup + - active agent/model handling + - local config updates + - async independent session management + - multiple parallel sessions +- Do **not touch `main`**. +- Work from a new branch off `main` / migration branch, not directly on `main`. +- Keep legacy code until user explicitly approves cleanup. +- Duplicate current legacy into a `legacy/` namespace before behavioral changes. +- Use the same WS route only after full parity tests pass. +- No application code changes until explicit approval. +- Prioritize safety over speed. +- Plan for migration, GUGI testing, and cleanup. +- Modularize long files, especially very large WS handlers. + +## Planning History + +- Generated migration task plans: + - `.maestro/tasks-20260601-093100-bzoj.json`, pack `zg6w`. + - `.maestro/tasks-20260601-093500-giep.json`, pack `kzh5`. +- User approved Gate 1: read-only inventory/migration matrix. +- Attempted BD issue creation from the plan; tool reported success but created `0` BD issues. +- Current continuation plan generated as: + - `.maestro/tasks-20260601-094423-9on6.json`, pack `vtyn`. + +## Local Repo / Worktree Observations + +- Working area root: `/Users/s0m0961/Documents/dev/personal/puppy/code_puppy.worktrees`. +- Candidate directories found: + - `external-ws-replica-test/` + - `h37o-input-notifier/` + - `maestro-checkpoints-20260526/` + - `puppy-desk-migration/` +- Important branches/statuses: + - `/code_puppy.worktrees/external-ws-replica-test` + - branch: `test/external-ws-migration-replica` + - clean + - no `code_puppy/api` directory + - `/code_puppy.worktrees/puppy-desk-migration` + - branch: `feature/puppy-desk-migration` + - modified files: + - `code_puppy/config.py` + - `code_puppy/plugins/frontend_emitter/emitter.py` + - `code_puppy/session_storage.py` + - `code_puppy/tools/command_runner.py` + - `pyproject.toml` + - `uv.lock` + - source repo `/Users/s0m0961/Documents/dev/puppy/code-puppy` + - branch: `feature/puppy-desk` + - modified files: + - `code_puppy/model_utils.py` + - `uv.lock` + - untracked: + - `scripts/run_ws_streaming_api_server.sh` + +## Branch Diffs and Scope + +- Source `feature/puppy-desk` vs `main` has 53 relevant changed/added files around: + - `code_puppy/api/**` + - `code_puppy/plugins/frontend_emitter/**` + - `code_puppy/config.py` + - `code_puppy/session_storage.py` + - `code_puppy/tools/command_runner.py` + - `pyproject.toml` + - `uv.lock` +- Diff stat source vs main: + - ~53 files changed + - ~15.5k insertions + - ~547 deletions +- Largest additions include: + - `api/ws/chat_handler.py` + - `api/db/queries.py` + - `api/templates/chat.html` + - API tests and WS tests + - config/session storage/tool runner/frontend emitter changes +- `puppy-desk-migration` vs `main` appears to already contain a port of many API files. +- `external-ws-replica-test` appears useful for newer frontend emitter work but does not contain the API/WS stack. + +## API / WS Stack Inventory + +Source `feature/puppy-desk` includes full API/WS stack: + +- `code_puppy/api/app.py` +- `code_puppy/api/db/` +- `code_puppy/api/permission_plugin.py` +- `code_puppy/api/permissions.py` +- routers +- `code_puppy/api/session_cache.py` +- `code_puppy/api/session_context.py` +- templates +- tests +- `code_puppy/api/websocket.py` +- `code_puppy/api/ws/` handlers, schemas, runtime, tests + +`puppy-desk-migration` has a similar stack and extra tests: + +- `test_config_router_schema.py` +- `test_session_model_compat.py` + +`external-ws-replica-test` has no `code_puppy/api` directory but has newer frontend emitter files. + +## Key Feature Search Findings + +Searched for terms including: + +- `assistant_message_delta` +- `ServerToolResult` +- `permission_response` +- `switch_session` +- `switch_model` +- `set_cwd` +- `working_directory` +- `b1_streaming_used` +- `stop_draining` +- `current_tool_group_id` + +Findings: + +- `schemas.py` defines protocol messages for: + - `switch_model` + - `switch_session` + - `set_working_directory` + - `permission_response` + - `assistant_message_delta` + - `tool_result` +- Tests assert `ServerToolResult(...)` includes `tool_group_id`. +- `background_save.py` handles `working_directory`. +- `session_runtime_manager.py` stores `working_directory`. + +## Frontend Emitter Findings + +Source `feature/puppy-desk` frontend emitter: + +- Files: + - `__init__.py` + - `emitter.py` + - `register_callbacks.py` +- No `session_context.py`. +- Supports: + - session-scoped subscriptions + - global legacy subscriptions + - per-session recent event buffers + - `emit_event(event_type, data, session_id=None)` + - `subscribe(session_id=None)` + - `unsubscribe(queue)` + - `get_recent_events(session_id=None)` + - `get_subscriber_count()` + - `clear_recent_events()` +- Global subscribers receive all events. +- Session subscribers receive only matching session events. + +`puppy-desk-migration` and `external-ws-replica-test` frontend emitter: + +- Include `session_context.py`. +- External replica has a newer design with: + - ContextVar fallback via `current_emitter_session_id` + - explicit `session_id` kwarg precedence + - wildcard and session-filtered subscribers + - `_Subscriber` dataclass records + - single recent events buffer + - legacy `_subscribers` compatibility + - double-delivery prevention. + +Frontend callback module findings: + +- Emits frontend events for: + - pre-tool call + - post-tool call + - stream events + - agent invocation +- Reads WebSocket context from permission plugin to attach session ID. +- Suppresses duplicate tool events during structured streaming. +- Sanitizes tool args and stream data. +- Serializes complex/Pydantic tool results. +- Preserves structured list/dict args when payload is small enough for frontend rendering. + +## Key File Size / Hash Observations + +Source `feature/puppy-desk`: + +- `frontend_emitter/emitter.py`: ~6.2 KB +- `register_callbacks.py`: ~12.1 KB +- no `session_context.py` +- `api/ws/chat_handler.py`: ~238 KB +- `api/ws/schemas.py`: ~19.5 KB +- `api/session_context.py`: ~24.2 KB +- `api/db/queries.py`: ~40.8 KB + +`puppy-desk-migration`: + +- `frontend_emitter/emitter.py`: ~10 KB +- `register_callbacks.py`: ~14.9 KB +- `session_context.py`: ~2.1 KB +- `api/ws/chat_handler.py`: ~235.8 KB +- `api/ws/schemas.py`: same hash as source +- `api/session_context.py`: ~24.9 KB +- `api/db/queries.py`: same hash as source + +`external-ws-replica-test`: + +- emitter files similar to migration +- no API/WS files + +## Source Git History Highlights + +Recent source commits touched: + +- daily backend logs +- DB empty-row guard +- `tool_group_id` consistency +- JSON serialization of tool results +- `RuntimeError` disconnect handling +- tracking `tool_group_id` outside drain scope +- pre-`stream_end` tool result extraction +- duck-typed tool return extraction +- emitter tool-event suppression scopes +- unified tool IDs / duplicate suppression +- schema migration of tool/terminal/permission messages +- schema migration of config/command/streaming messages +- advisory client validation +- protocol model Literal defaults +- `send_typed` helper migration + +## Key Source File Summaries + +### `code_puppy/api/websocket.py` + +- Thin endpoint registrar delegating to `code_puppy.api.ws` handlers: + - events + - chat + - terminal + - health + - sessions +- Re-exports attachment builder and connection manager for backward compatibility. + +### `code_puppy/api/ws/schemas.py` + +- Defines protocol version `1.0.0`. +- Defines typed Pydantic WebSocket protocol models. +- Client message types include: + - `message` + - `switch_agent` + - `switch_model` + - `switch_session` + - `set_working_directory` + - `update_session_meta` + - `get_config` + - `set_config` + - `command` + - `cancel` + - `permission_response` +- Server message types include: + - `system` + - `session_restored` + - `session_switched` + - `working_directory_changed` + - `session_meta_updated` + - `config_value` + - `command_result` + - `status` + - `user_message` + - `assistant_message_start` + - `assistant_message_delta` + - `assistant_message_end` + - `tool_call` + - `tool_result` + - `tool_return` + - `agent_invoked` + - `response` + - `stream_end` + - `error` + - `permission_request` + - `cancelled` + +### `api/ws/runtime/session_runtime_manager.py` + +- Async manager for per-session runtimes. +- Supports runtime creation, lookup, cancellation, removal, idle cleanup, and active task count. +- Designed for multiple hot sessions / multiplexed chat. + +### `api/ws/runtime/session_runtime.py` + +- Dataclass holding per-session runtime: + - active task + - cancel event + - lock + - created/last-used timestamps. + +### `api/permissions.py` + +- WebSocket permission request system. +- Maintains global `permission_futures`. +- `request_permission` sends typed permission request, waits up to 300s, honors YOLO mode, fails safe if no WS. +- `handle_permission_response` resolves pending futures. + +### `api/permission_plugin.py` + +- Uses task-scoped `ContextVar` for WS permission context. +- Has suppression flag for duplicate frontend emitter tool events. +- `pre_tool_call_permission` asks frontend for approval and blocks denied tools. +- Includes shell command permission handling. + +### `api/ws/chat_handler.py` + +Very large file, around 235–238 KB. Responsibilities include: + +- Session provisioning/restoration via SQLite and `session_manager`. +- Initial config/CWD/system/session-restored messages. +- WS message handlers for: + - switch agent + - switch model + - switch session + - set working directory + - update session metadata + - get/set config + - command + - cancel + - permission response + - user message +- CWD persistence and unchanged-CWD messages. +- Agent streaming plus concurrent frontend emitter draining. +- Typed assistant/tool lifecycle messages: + - `assistant_message_start` + - `assistant_message_delta` + - `assistant_message_end` + - `tool_call` + - `tool_result` + - `stream_end` +- Tracking of active parts, collected text, pending tool calls, tool ID aliases, and tool group IDs. +- Handling deltas before starts by creating starts first. +- Accumulating tool call args deltas. +- Extracting full tool results before `stream_end`. +- Accepting permission responses and cancel while streaming. +- Supporting session switch/create while streaming through background save. + +Specific inspected area around lines 2350–2530 showed: + +- JSON parsing of tool args. +- `tool_group_id` generation if missing. +- Sending `ServerToolCall`. +- Storing pending tool-call metadata. +- Mapping raw tool call IDs to normalized IDs. +- Cleaning active parts. +- Handling WebSocket close during streaming. +- Final-draining queued stream events. +- Subscribing to frontend emitter by `session_id`. +- Starting concurrent drain task. + +### `api/db/connection.py` + +- SQLite connection singleton using `aiosqlite`. +- DB path from `PUPPY_DESK_DB`, otherwise `~/.puppy_desk/chat_messages.db`. +- Backend can create/migrate schema if frontend has not. +- Schema includes: + - `sessions` + - `messages` + - `tool_calls` + - `compaction_log` +- Tracks schema versions/migrations up to v4. +- Concurrency handled by `aiosqlite`, no threading lock. + +### `api/db/queries.py` + +- Large async query helper module for shared `chat_messages.db`. +- Session operations include: + - `session_exists` + - `get_session_row` + - `get_session_metadata` + - `upsert_session` + - `update_session_stats` + - `update_session_working_directory` + - `update_session_meta_fields` + - `soft_delete_session` +- Defines message sequencing/insertion helpers. +- Uses `aiosqlite` and commits/rollbacks around writes. + +### `api/session_context.py` + +- Multi-session isolation layer. +- Defines `SessionContext` with: + - per-session agent + - agent name + - model name + - working directory + - title/pinned metadata + - websocket reference + - saved per-agent histories + - compacted hashes + - lazy async operation lock +- Defines validation for safe session IDs and known agent names. +- Defines `SessionManager` for: + - create/get/destroy sessions + - switch agents + - preserve/restore history across agent switches + - carry session model across agent reloads. + +### `api/session_cache.py` + +- LRU cache for deserialized session data. +- Uses async lock and thread pool for file-related work. +- Features: + - max cached sessions + - TTL expiry + - file modification invalidation + - pre-serialized JSON cache + - stats for hits/misses/evictions/expirations. + +## Current State / Important Notes + +- No application code was modified during read-only Gate 1 investigation. +- This file itself is a documentation checkpoint requested by the user. +- Final migration matrix, options, risks, and recommendation still need to be completed. diff --git a/docs/migration/puppy-desk-migration-plan-options.md b/docs/migration/puppy-desk-migration-plan-options.md new file mode 100644 index 000000000..aff3de7ae --- /dev/null +++ b/docs/migration/puppy-desk-migration-plan-options.md @@ -0,0 +1,825 @@ +# Puppy Desk Migration — Gate 1 Matrix, Options, and Approval Gates + +Status: **Gate 1 read-only planning + docs only** +Base branch target: **main** +Current working branch observed: `feature/puppy-desk-migration` +Source branch/repo: `/Users/s0m0961/Documents/dev/puppy/code-puppy`, branch `feature/puppy-desk` +Safety rule: **do not touch `main`; do not modify application code without explicit approval** + +Related checkpoint: `docs/migration/puppy-desk-gate1-investigation-summary.md` + +--- + +## 1. Migration Objective + +Migrate the divergent `feature/puppy-desk` backend/API/WebSocket work into this repo while preserving all GUI-facing behavior needed by GUGI / Puppy Desk: + +- Same final WS route, but only after parity tests pass. +- Legacy route/code preserved first in a `legacy/` namespace before behavior changes. +- No feature loss around streaming, tool calls, permissions, CWD, sessions, config, DB persistence, or multiple sessions. +- Long files must be modularized after parity is protected by tests. + +--- + +## 2. High-Level Inventory Matrix + +| Area | Source `feature/puppy-desk` | Current migration branch | Gap / Decision | Risk | +|---|---|---|---|---| +| API package | Full `code_puppy/api/**` stack added | Similar/full stack already present | Verify current branch contains all intended API files and no accidental omissions | Medium | +| WS route registrar | `api/websocket.py` delegates to `api.ws` handlers | Similar | Preserve legacy route until parity; same-route swap only after tests | High | +| WS schemas | Typed Pydantic protocol v1.0.0; source and migration same hash | Same hash | Treat as contract; freeze with schema tests | High | +| Huge chat handler | ~238 KB; central streaming/session/permission/CWD logic | ~230–236 KB; differs by small but important edits | Must first port/verify parity, then refactor behind tests | Very high | +| Runtime manager | Per-session async runtime manager | Same structure | Keep; validate parallel sessions and cancellation | High | +| Background save | Saves agent result/session metadata while switched away | Same file size/hash likely same | Validate session switch while streaming | High | +| DB connection/schema | SQLite `~/.puppy_desk/chat_messages.db`, schema/migrations to v4 | Same/similar | Need DB migration tests with temp `PUPPY_DESK_DB` | High | +| DB queries | Large async query helper; same hash in source/current | Same hash | Candidate to split after parity into session/message/tool modules | Medium | +| Session context | Source supports isolated agent/model/CWD/history | Current has compatibility helpers for missing `set_session_model` | Keep migration helper if needed; validate model switching | High | +| Permissions | WS permission futures + plugin context | Same/similar | Must validate approve/deny, timeout/fail-safe, YOLO | High | +| Frontend emitter | Source has session-scoped queues + per-session recent buffers | Current/external has newer ContextVar design | Choose final emitter design carefully; behavior must prevent cross-session leakage | High | +| Emitter callbacks | Source attaches session id via WS context and suppresses duplicate tool events | Current differs substantially | Need dedicated emitter parity tests | High | +| Config router | Source expects `CONFIG_SCHEMA` | Current has fallback if schema absent | Decide whether to preserve full schema or fallback; GUI likely benefits from full schema | Medium | +| Core config | Source contains larger `CONFIG_SCHEMA` and queue size default `10000`; current appears to remove schema and queue default returns `100` | Need explicit decision; queue size affects streaming reliability | High | +| Command runner | Source imports callback directly | Current has ImportError fallback | Keep fallback if current architecture requires; validate terminal streaming callbacks | Medium | +| Session storage | Same source/current | No action except integration validation | Low | +| Model utils | Source/current differ, source has uncommitted changes | Need inspect before migration; could affect active model handling | Medium | +| Dependencies | API requires FastAPI/aiosqlite/websockets/etc.; package files changed | `pyproject.toml` and `uv.lock` modified in current | Snyk required before testing if dependency changes are applied/retained | High | +| Tests | Source has API/WS tests; current adds config/model compatibility tests | Combine and strengthen | Must add GUGI/parity scenarios before same-route swap | High | + +--- + +## 3. Source vs Current Migration Differences Requiring Decisions + +### 3.1 `api/ws/chat_handler.py` + +Observed diff current-vs-source is small but semantically important: + +- Source appended a raw dict system message into agent in-memory history on CWD change. +- Current migration branch replaces that with a warning comment: do **not** append raw dicts because runtime expects typed `ModelMessage` objects and raw dict injection can corrupt later turns / `result=None` behavior. + +Recommendation: **Keep the safer current behavior**, but verify the GUI still receives/persists directory banners through SQLite/system messages. + +Required tests: + +- Change CWD and verify WS `working_directory_changed` message. +- Verify directory system banner persists and reloads. +- Verify next agent turn after CWD change does not fail from corrupt history. + +### 3.2 Frontend emitter design + +Source design: + +- Explicit `session_id` on `emit_event`. +- Global legacy subscribers receive all events. +- Session subscribers receive matching events only. +- Maintains global and per-session recent buffers. + +Current/external design: + +- Adds `ContextVar` fallback via `current_emitter_session_id`. +- Explicit kwarg takes precedence. +- Uses subscriber records to avoid double delivery. +- Single recent events buffer. + +Recommendation: **Use the current/external ContextVar-capable design**, but add tests for source behavior: + +- Explicit session id routing. +- ContextVar fallback routing. +- Explicit `session_id=None` opt-out if supported. +- Wildcard subscriber receives all events exactly once. +- Session subscriber does not receive unrelated/no-session events. +- No cross-session leakage with concurrent sessions. + +Open question: Do we need per-session recent buffers for GUGI replay, or is one global recent buffer acceptable? + +### 3.3 `config.py` and config schema + +Observed current-vs-source differences include: + +- Source appears to include a large `CONFIG_SCHEMA` used by `/api/config/schema`. +- Current `api/routers/config.py` has a fallback if `CONFIG_SCHEMA` is unavailable. +- Source has `frontend_emitter_queue_size` default around `10000` to support large responses. +- Current returns `100` if unset. + +Recommendation: **Do not silently lose `CONFIG_SCHEMA` or queue capacity**. Either: + +1. Restore/export full schema in modular form, or +2. Provide a new schema provider module consumed by router and GUI. + +Queue size should be intentionally chosen; for streaming-heavy GUI, `100` may be too small and risks event drops. + +### 3.4 `api/session_context.py` + +Current migration branch adds helper functions: + +- `_apply_session_model(agent, model_name)` +- `_reload_agent_if_supported(agent)` + +These make model switching robust when legacy/migrated agents lack `set_session_model` or reload hooks. + +Recommendation: **Keep compatibility helpers** unless deeper inspection proves all agents implement required methods. + +Required tests: + +- Switch model for code-puppy agent. +- Switch model for any older/simple agent lacking the method. +- Switch agent then model and verify session-local model is preserved. + +### 3.5 `tools/command_runner.py` + +Current migration branch adds fallback for missing `on_run_shell_command_output` callback import. + +Recommendation: Keep if current branch compatibility requires it, but test terminal command streaming and callback registration. + +--- + +## 4. Proposed Modular Architecture + +The current `api/ws/chat_handler.py` is too large and risky to maintain. Refactor only **after parity tests protect behavior**. + +Proposed module split: + +```text +code_puppy/api/ws/ + chat_handler.py # thin route orchestration only + chat/ + __init__.py + context.py # per-connection/session state object + lifecycle.py # connect/init/restore/cleanup + dispatch.py # client message dispatch table + session_ops.py # switch/create/restore session + agent_ops.py # switch agent/model + cwd_ops.py # set working directory + persistence + config_ops.py # get/set config WS handlers + permission_ops.py # permission response/cancel plumbing + streaming.py # run agent and coordinate stream lifecycle + drain.py # frontend emitter draining/event batching + tool_lifecycle.py # tool ids, args deltas, tool results/groups + persistence.py # save turn/session/tool/message data + errors.py # typed error helpers +``` + +DB query split after tests: + +```text +code_puppy/api/db/ + connection.py + sessions.py + messages.py + tool_calls.py + compaction.py + queries.py # compatibility re-export layer temporarily +``` + +Config split option: + +```text +code_puppy/config/ + schema.py # GUI schema metadata + accessors.py # get/set helpers or wrappers +``` + +Migration rule: keep old import paths via compatibility wrappers until GUI and tests are stable. + +--- + +## 5. Execution Options + +### Option A — Big-bang copy, then fix + +Copy source branch API/WS files wholesale over current repo, then resolve breakage. + +Pros: + +- Fastest initial copy. +- Lowest chance of missing a source file. + +Cons: + +- High regression risk. +- May overwrite already-needed frontend emitter fixes in this repo. +- Hard to isolate bugs. +- Does not respect modularity goal until later. + +Recommendation: **Do not choose** except as emergency baseline branch. + +### Option B — Current migration branch as base + targeted delta reconciliation + +Treat `feature/puppy-desk-migration` as existing partial migration. Reconcile only source-vs-current gaps, preserve current repo emitter/frontend fixes, then test. + +Pros: + +- Most efficient because API stack already exists in current branch. +- Preserves already-made frontend emitter changes. +- Smaller deltas to review. +- Safer than big-bang. + +Cons: + +- Requires careful audit that current branch did not accidentally drop source behavior. +- Existing uncommitted app-code changes need review/ownership. +- Long `chat_handler.py` remains until post-parity refactor. + +Recommendation: **Best default path**. + +### Option C — Clean-room modular port from source into new branch off `main` + +Start fresh from `main`, duplicate legacy into `legacy/`, port features module-by-module into modular architecture from the beginning. + +Pros: + +- Cleanest long-term architecture. +- Avoids carrying accidental partial migration decisions. +- Most aligned with modularity goal. + +Cons: + +- Slowest path. +- Higher risk of feature omissions during rewrite. +- GUGI validation delayed. + +Recommendation: Use only if current migration branch is considered untrustworthy. + +### Option D — Compatibility route side-by-side, then same-route swap + +Add migrated API under a temporary side route/namespace, keep legacy route active, run parity tests against both, then switch same route after approval. + +Pros: + +- Safest runtime rollout. +- Easy rollback. +- Lets GUGI compare old/new behavior before final route swap. + +Cons: + +- More temporary routing/code complexity. +- Requires duplicated route plumbing. + +Recommendation: Combine with **Option B** if route-level safety is important. + +--- + +## 6. Recommended Path + +Recommended execution path: **Option B + Option D safeguards** + +1. Continue from `feature/puppy-desk-migration`, or create a fresh branch off `main` and replay the current migration branch cleanly. +2. Before touching behavior, duplicate legacy/current WS route implementation into `legacy/` namespace. +3. Preserve current frontend emitter improvements unless parity tests show behavior loss. +4. Reconcile source-vs-current differences explicitly: + - Keep safer CWD history behavior. + - Restore or modularize `CONFIG_SCHEMA` if GUI needs it. + - Choose emitter recent-buffer policy. + - Keep session model compatibility helpers. + - Validate command runner callback fallback. +5. Add parity tests before same-route swap. +6. Run GUGI validation against temporary/new route. +7. Only after parity + GUGI pass, swap same route with approval. +8. Stabilize. +9. Refactor large files into modules under test coverage. +10. Cleanup legacy only after explicit user approval. + +--- + +## 7. Approval Gates + +### Gate 1 — Inventory and migration matrix + +Status: **in progress / docs only** + +Deliverables: + +- Investigation checkpoint saved. +- Migration matrix/options saved. +- User confirms preferred path and goals. + +Approval needed before Gate 2. + +### Gate 2 — Branch hygiene and legacy duplication + +No behavior changes except copying legacy/current route code into `legacy/` namespace. + +Acceptance: + +- Working branch confirmed not `main`. +- Legacy namespace created. +- Existing route behavior unchanged. +- Smoke tests pass. + +### Gate 3 — Parity test harness + +Acceptance: + +- Tests cover all high-risk GUI protocol features: + - connect/session restored + - user message -> assistant start/delta/end/stream_end order + - tool call args, result, `tool_group_id` + - permission request/approve/deny/timeout + - CWD set/unchanged/persistence + - switch agent/model/session + - config get/set/schema + - terminal command result/streaming if route active + - multiple sessions in parallel with no emitter leakage + - cancel during streaming + - switch session during streaming + background save +- Coverage matrix includes happy path, edge/boundary, negative/error, regression. + +### Gate 4 — Targeted migration reconciliation + +Acceptance: + +- Source-vs-current gaps resolved with documented choices. +- No known source feature loss. +- Tests pass. +- Snyk scan completed if dependency files are changed. + +### Gate 5 — GUGI validation on temporary/new route + +Acceptance: + +- GUGI can connect. +- Chat streams correctly. +- Tool panels show args/results. +- Permission UI works. +- Sessions list/load/switch correctly. +- CWD/model/agent/config UI flows work. +- Parallel sessions do not cross streams/events. + +### Gate 6 — Same-route swap + +Only after Gate 5 approval. + +Acceptance: + +- Same route uses migrated implementation. +- Legacy still retained for rollback. +- Full parity and GUGI smoke pass. + +### Gate 7 — Modular cleanup + +Only after migrated route is stable. + +Acceptance: + +- `chat_handler.py` split into modules. +- `db/queries.py` split or compatibility-layered. +- Import compatibility retained as needed. +- Tests remain green. + +### Gate 8 — Legacy removal + +Only after explicit user approval. + +Acceptance: + +- No GUGI dependency on legacy namespace. +- Rollback plan no longer needed or separately archived. +- Legacy code removed and docs updated. + +--- + +## 8. Test Coverage Matrix Proposal + +| ID | Requirement | Scenario | Risk Tag | Test Type | Status | +|---|---|---|---|---|---| +| T1 | WS connect/session restore | Existing session loads messages/meta | happy-path | automated WS | planned | +| T2 | Streaming order | start -> delta(s) -> end -> stream_end | regression | automated WS | planned | +| T3 | Delta before start robustness | handler emits synthetic start first | edge | automated unit/WS | planned | +| T4 | Tool call args | args deltas accumulate into full args | happy-path | automated WS | planned | +| T5 | Tool result | result emitted with `tool_group_id` | regression | automated WS | partially existing | +| T6 | Permission approve | tool proceeds after approval | happy-path | automated WS | planned | +| T7 | Permission deny | tool blocked and error/result surfaced safely | negative | automated WS | planned | +| T8 | Permission timeout/no WS | fail safe | negative | unit | planned | +| T9 | CWD set | path changes, persists, emits banner | happy-path | automated WS/DB | planned | +| T10 | CWD unchanged | unchanged flag, no duplicate corruption | edge | automated WS | planned | +| T11 | Switch model | session-local model changes and persists | happy-path | automated WS | planned | +| T12 | Agent without model setter | fallback does not crash | edge | unit | planned | +| T13 | Switch session during stream | old session saves in background | regression | automated WS | planned | +| T14 | Parallel sessions | no emitter leakage/crossed events | regression | automated concurrent WS | planned | +| T15 | Cancel stream | cancellation frame and cleanup | negative | automated WS | planned | +| T16 | Config schema | schema endpoint returns GUI-usable metadata | happy-path | API | planned/current extra test | +| T17 | Config missing schema fallback | endpoint still works | edge | API | current extra test | +| T18 | Terminal command | command result and output callback flow | happy-path | API/WS | planned | +| T19 | DB fresh init | empty DB creates schema v4 | happy-path | DB temp file | planned | +| T20 | DB migration/dedup | older schema upgrades safely | regression | DB temp file | planned | +| T21 | GUGI chat smoke | GUI sends prompt and renders stream | happy-path | manual/browser | planned | +| T22 | GUGI tools/permissions | GUI renders tool args/results + approval | high-risk manual | manual/browser | planned | + +--- + +## 9. Open Questions for User + +1. Should the existing `feature/puppy-desk-migration` branch be treated as the starting point, or should we create a brand-new clean branch off `main` and replay only approved changes? +2. Is a temporary side-by-side route acceptable for validation, or do you want legacy namespace only and then same-route swap? +3. For frontend emitter recent events, do you need per-session recent replay, or is global recent replay acceptable if session-filtered live delivery is correct? +4. Should `/api/config/schema` expose the full rich `CONFIG_SCHEMA` for GUGI, or is inferred fallback metadata enough? +5. What is the minimum GUGI validation flow you want before same-route swap? +6. Are the current uncommitted application-code changes on `feature/puppy-desk-migration` yours and intended to be preserved? + +--- + +## 10. Proposed Immediate Next Step After Approval + +If approved, proceed to **Gate 2**: + +- Confirm/create safe migration branch off `main`. +- Snapshot current status. +- Duplicate legacy/current WS route implementation into `legacy/` namespace. +- Make no behavior switch yet. +- Add initial parity tests around existing behavior. + + +--- + +## 11. User Decisions Recorded — 2026-06-01 + +Confirmed by user: + +1. **Use existing `feature/puppy-desk-migration` as the base.** + - Do not restart from a fully clean branch unless a later blocker requires it. + - Existing migration branch decisions and compatibility fixes are considered valuable. + +2. **Temporary side-by-side route is acceptable.** + - We may expose migrated/experimental WS/API behavior behind a temporary validation route while preserving legacy/current route behavior. + - Same-route swap remains gated by parity + GUGI validation + explicit approval. + +3. **Preserve uncommitted app-code changes on `feature/puppy-desk-migration`.** + - Current modified files must not be overwritten casually: + - `code_puppy/config.py` + - `code_puppy/plugins/frontend_emitter/emitter.py` + - `code_puppy/session_storage.py` + - `code_puppy/tools/command_runner.py` + - `pyproject.toml` + - `uv.lock` + - Before Gate 2 implementation, snapshot/diff these changes and treat them as intentional unless user says otherwise. + +## 12. GUGI Recent-Event Replay Decision Detail + +Open decision: whether GUGI needs **per-session recent-event replay** or whether a **global recent-event replay buffer** is acceptable when live delivery is correctly session-filtered. + +### What “recent-event replay” means + +The frontend emitter supports live event subscription, but it can also keep a small in-memory buffer of recent events. When GUGI reconnects, refreshes, or opens a panel late, it may ask for recent events to catch up on tool calls, stream chunks, permission events, or agent invocation events that occurred just before the subscriber attached. + +Live routing and replay routing are separate concerns: + +- **Live routing**: while connected, does session A only receive session A events? +- **Replay routing**: after reconnect/late subscribe, which older events are offered back to GUGI? + +A system can have correct live routing but still leak or confuse replayed events if the recent buffer is global and the replay API does not filter carefully. + +### Option R1 — Global recent buffer only + +All events are stored in one recent-events list, each event includes optional `session_id`. + +Expected behavior: + +- Wildcard callers can inspect all recent events. +- GUGI/session-specific callers must filter by `session_id` before display. + +Pros: + +- Simpler implementation. +- Lower memory overhead. +- Easier backward compatibility with older `get_recent_events()` semantics. +- Matches the newer external emitter design observed in `external-ws-replica-test`. + +Cons: + +- Higher risk of accidental cross-session replay if any consumer forgets to filter. +- More frontend/API discipline required. +- A busy session can evict another session’s recent events from the shared buffer. +- Reconnecting to a quiet session after a busy parallel session may miss relevant older events. + +Best if: + +- GUGI mostly relies on DB/session state for reloads, not emitter replay. +- Recent events are only diagnostic or best-effort. +- Event volume is moderate. +- All replay call sites are strongly typed and tested for session filtering. + +### Option R2 — Per-session recent buffers + +Each session has its own recent-events list; optionally a global buffer is retained for legacy/wildcard subscribers. + +Expected behavior: + +- `get_recent_events(session_id="A")` returns only session A events. +- `get_recent_events()` returns legacy/global events or all events depending on compatibility policy. + +Pros: + +- Safer isolation by default. +- Prevents cross-session replay leakage even if frontend forgets to filter. +- Busy session B cannot evict session A’s recent events from A’s own buffer. +- Better for multi-tab/multi-session GUI behavior. +- Better fit for strict session isolation requirement. + +Cons: + +- More state to manage and clean up. +- Need buffer cleanup when sessions are destroyed/idle. +- Slightly more complex compatibility semantics. +- Must decide whether no-session/global events should be visible to session subscribers. + +Best if: + +- GUGI uses recent replay to reconstruct transient UI such as active tool calls or streaming panels. +- Multiple sessions can be active in parallel. +- Strict no-cross-session leakage is more important than simplicity. +- Reconnect/refresh behavior matters. + +### Option R3 — Hybrid: global buffer plus per-session indexed replay + +Maintain a global recent buffer for backward compatibility and diagnostics, but expose session-filtered replay APIs that are safe by default. Internally this can be implemented either as: + +- a global list filtered by `session_id`, plus tests, or +- true per-session buffers plus a global legacy buffer. + +Pros: + +- Preserves backward compatibility. +- Allows safe GUGI code path: always request replay by session. +- Can start with global-filtered implementation and upgrade to true per-session if eviction becomes a problem. + +Cons: + +- More policy decisions to document. +- If implemented only as global-filtered, shared-buffer eviction risk remains. + +Recommended default for migration: + +**R3 with true per-session replay where practical.** + +Reasoning: + +- User explicitly requires async independent session management and multiple parallel sessions. +- GUI correctness depends on not crossing tool calls, deltas, permission prompts, or terminal output between sessions. +- The source branch already had per-session recent buffers, which is a useful safety feature. +- The newer current/external emitter adds ContextVar routing and double-delivery prevention, which is also valuable. +- Best final design is to combine both: + - ContextVar-capable emission/routing from current/external work. + - Per-session replay isolation from source behavior. + - Legacy wildcard/global support retained for compatibility. + +Suggested policy: + +1. Live session subscribers receive only exact `session_id` matches. +2. Wildcard subscribers receive all events exactly once. +3. `get_recent_events(session_id="A")` returns only session A events. +4. `get_recent_events()` retains old behavior for legacy/debug consumers. +5. Explicit `session_id=None` events are global/broadcast for wildcard/debug only, not replayed to session-specific callers unless explicitly requested. +6. Session recent buffers are capped and cleaned up when sessions expire. + +Required tests: + +- Two sessions emit interleaved events; each session replay returns only its own events. +- Wildcard replay returns all expected events with no duplicates. +- Busy session B does not evict session A’s per-session recent events. +- ContextVar fallback attaches the correct session id. +- Explicit kwarg overrides ContextVar. +- Explicit `session_id=None` behavior is documented and tested. +- Reconnect/late subscribe during tool call does not show another session’s tool event. + + +--- + +## 13. User Clarification — GUGI State Replay Source + +User clarified: + +- Recent replay/state reconstruction is handled by **React GUI state** or by the **DB**. +- Both React GUI state and DB-backed state are already separated by `session_id`. + +Updated implication: + +- The frontend emitter recent-events buffer is **not** the primary source of truth for GUGI session reconstruction. +- Therefore, we do **not** need to over-optimize emitter replay as the authoritative reconnect/reload mechanism. +- However, live event routing remains critical: active stream/tool/permission events must never cross sessions. + +Updated recommendation: + +Use the current/external emitter direction as the base: + +1. Keep **ContextVar-capable live routing**. +2. Keep explicit `session_id` routing and exact-match session subscribers. +3. Keep wildcard/global subscribers for legacy/debugging. +4. Treat emitter recent replay as **debug/best-effort compatibility**, not authoritative GUI state. +5. Ensure any replay API that GUGI might call can filter by `session_id`, even if internally stored in one global buffer. +6. Do not require true per-session recent buffers unless later GUGI testing shows replay eviction/leakage causes real UX issues. + +Revised preference: + +- **Global recent buffer is acceptable** if: + - live delivery is strictly session-filtered, + - replay consumers filter by `session_id`, + - DB/React remain the source of truth for state reconstruction, + - tests prove no live cross-session leakage. + +Still required tests: + +- Concurrent session live events do not cross. +- Wildcard subscribers receive all events exactly once. +- Session subscribers receive only matching session events. +- ContextVar fallback attaches correct session id. +- Explicit `session_id` overrides ContextVar. +- Replay filtering by `session_id` works if exposed/used. + +Priority change: + +- Per-session recent buffers move from **recommended default** to **optional enhancement**. +- Live session routing and DB/React session correctness are the hard requirements. + + +--- + +## 14. Gate 2 Implementation Shape — Proposed, Pending Explicit Approval + +User has approved the strategy direction, but application-code changes should still be gated explicitly. + +Confirmed base: + +- Continue on `feature/puppy-desk-migration`. +- Preserve uncommitted app-code changes. +- Temporary side-by-side route is acceptable. +- Do not touch `main`. + +### Current route shape observed + +- `code_puppy/api/app.py` calls `setup_websocket(app)`. +- `code_puppy/api/websocket.py` calls: + - `register_events_endpoint(app)` + - `register_chat_endpoint(app)` + - `register_terminal_endpoint(app)` + - `register_health_endpoint(app)` + - `register_sessions_endpoint(app)` +- `code_puppy/api/ws/chat_handler.py` defines `register_chat_endpoint(app)` and registers hardcoded route: + - `@app.websocket("/ws/chat")` + +### Gate 2 goal + +Preserve current behavior before any migration/refactor work by duplicating the current route implementation into a legacy namespace, then prepare for side-by-side validation. + +### Gate 2 code-change options + +#### G2-A — Archive-only legacy namespace + +Copy the current `code_puppy/api/ws` implementation into: + +```text +code_puppy/api/ws/legacy/ +``` + +Do not register the legacy copy yet. + +Pros: + +- Lowest behavioral risk. +- Existing `/ws/chat` route remains unchanged. +- Gives rollback/reference copy before edits. + +Cons: + +- Not enough by itself for side-by-side live validation. + +#### G2-B — Parameterized route registration + +Change endpoint registration functions to accept an optional path, defaulting to current paths: + +```python +def register_chat_endpoint(app: FastAPI, path: str = "/ws/chat") -> None: + @app.websocket(path) + async def websocket_chat(...): + ... +``` + +Then current behavior remains unchanged because default is still `/ws/chat`. + +Later, a temporary route can register the same or modular implementation at e.g.: + +```text +/ws/chat-migration +/ws/chat-next +/ws/chat-v2 +``` + +Pros: + +- Small, controlled change. +- Enables side-by-side validation without copying a 230KB handler again. +- Keeps same-route swap simple. + +Cons: + +- Touches the huge handler early. +- Must be covered by a smoke test verifying `/ws/chat` still registers. + +#### G2-C — New route module wrapping copied implementation + +Create a separate temporary route module that copies/wraps the current chat handler under a different path. + +Pros: + +- Avoids parameterizing current handler. +- Side-by-side route available immediately. + +Cons: + +- Duplicates very large code. +- Increases risk of divergence and merge pain. + +### Recommended Gate 2 path + +Recommended: **G2-A first, then G2-B only when we are ready to register side-by-side.** + +Sequence: + +1. Snapshot current uncommitted changes. ✅ done in docs +2. Copy current WS/API route implementation into `code_puppy/api/ws/legacy/` as a reference/rollback namespace. +3. Do not register legacy yet. +4. Add minimal smoke tests/import tests proving current `/ws/chat` registration still works. +5. Once tests are in place, parameterize `register_chat_endpoint(..., path="/ws/chat")` and register temporary route only when needed. + +### Proposed temporary route name + +Preferred temporary route: + +```text +/ws/chat-migration +``` + +Alternatives: + +```text +/ws/chat-next +/ws/chat-v2 +``` + +Recommendation: `/ws/chat-migration` because it is explicit and unlikely to be confused with final versioning. + +### Gate 2 acceptance checks + +Before marking Gate 2 complete: + +- `git branch --show-current` is not `main`. +- Existing uncommitted changes remain preserved. +- Legacy namespace exists as a copy/reference. +- `/ws/chat` behavior is unchanged. +- No same-route swap has occurred. +- Minimal import/route registration tests pass. +- No dependency/security scan required unless dependency files change beyond preserved existing state; if package files are touched, Snyk check is required before testing phase. + + +--- + +## 15. Gate 2 Execution Evidence — 2026-06-01 + +Gate 2 approved by user and executed on branch `feature/puppy-desk-migration`. + +Application/test changes made: + +- Created legacy WS snapshot namespace: + - `code_puppy/api/ws/legacy/` +- Copied current WS implementation files into that namespace as a rollback/reference snapshot. +- Replaced `legacy/__init__.py` with a safe non-registering marker module. +- Added `code_puppy/api/ws/legacy/README.md` documenting that it is not registered and must not be cleaned up without approval. +- Added smoke tests: + - `code_puppy/api/tests/test_gate2_legacy_ws_namespace.py` + +Behavioral guarantees for Gate 2: + +- Existing `/ws/chat` route remains registered from active `code_puppy.api.ws.chat_handler`. +- No temporary route is registered yet. +- No same-route swap occurred. +- Legacy snapshot is importable as a namespace marker only; importing `code_puppy.api.ws.legacy` does not register routes. +- Existing uncommitted app-code changes were preserved. + +Targeted test evidence: + +```text +uv run pytest code_puppy/api/tests/test_gate2_legacy_ws_namespace.py -q +2 passed in 19.57s +``` + +Security note: + +- Gate 2 did not intentionally modify dependency files. +- `pyproject.toml` and `uv.lock` were already modified before Gate 2 and are preserved as user-approved existing changes. +- Snyk scan should still be run before broader test/security signoff if those dependency changes remain part of the migration deliverable. + + +### Gate 2 reviewer sign-off + +Independent code review agent result: **APPROVE** + +Reviewer findings summary: + +- No `/ws/chat` behavior switch. +- No temporary route registered. +- Legacy package import is marker-only/non-registering. +- Existing uncommitted app-code changes remain preserved. +- Smoke tests are sufficient for Gate 2 scope. + +Non-blocking caution recorded: + +- The legacy snapshot is archive/reference quality, not yet a fully executable rollback route, because copied modules still contain absolute imports to active `code_puppy.api.ws.*` modules. This is acceptable for Gate 2 archive/non-registration, but must be addressed before using `legacy` as a live rollback route. + diff --git a/docs/migration/puppy-desk-uncommitted-change-snapshot.md b/docs/migration/puppy-desk-uncommitted-change-snapshot.md new file mode 100644 index 000000000..a4e325d90 --- /dev/null +++ b/docs/migration/puppy-desk-uncommitted-change-snapshot.md @@ -0,0 +1,1567 @@ +# Puppy Desk Migration — Uncommitted Change Snapshot + +Purpose: preserve and document existing uncommitted app-code changes on `feature/puppy-desk-migration` before Gate 2. + +Generated: 2026-06-01T15:15:33Z + +## Branch and status +```text +feature/puppy-desk-migration + M code_puppy/config.py + M code_puppy/plugins/frontend_emitter/emitter.py + M code_puppy/session_storage.py + M code_puppy/tools/command_runner.py + M pyproject.toml + M uv.lock +?? docs/migration/ +``` + +## Diff stat +```text + code_puppy/config.py | 44 ++- + code_puppy/plugins/frontend_emitter/emitter.py | 13 +- + code_puppy/session_storage.py | 464 ++++++++++--------------- + code_puppy/tools/command_runner.py | 372 ++++++++++++++------ + pyproject.toml | 6 + + uv.lock | 220 ++++++++++++ + 6 files changed, 715 insertions(+), 404 deletions(-) +``` + +## File-level notes + +| File | Observed status | Migration handling | +|---|---|---| +| `code_puppy/config.py` | Modified; includes config/schema/default differences vs source branch | Preserve and reconcile intentionally; do not overwrite wholesale | +| `code_puppy/plugins/frontend_emitter/emitter.py` | Modified; current/external ContextVar-capable emitter direction | Preserve as base; ensure strict live session filtering | +| `code_puppy/session_storage.py` | Modified relative to branch, but source/current comparison previously showed same as source | Preserve until inspected; likely session compatibility work | +| `code_puppy/tools/command_runner.py` | Modified; observed callback import fallback | Preserve unless tests show callback breakage | +| `pyproject.toml` | Modified dependency/project config | Preserve; requires Snyk scan before test phase if dependency changes remain | +| `uv.lock` | Modified lockfile | Preserve; requires Snyk scan before test phase if dependency changes remain | + +## Detailed diff excerpts / full current diff for tracked modified files + +> This section intentionally captures the current tracked diff so it can be recovered if later work accidentally overwrites it. + +```diff +diff --git a/code_puppy/config.py b/code_puppy/config.py +index 997d7935..a2d058b8 100644 +--- a/code_puppy/config.py ++++ b/code_puppy/config.py +@@ -48,18 +48,38 @@ EXTRA_MODELS_FILE = os.path.join(DATA_DIR, "extra_models.json") + AGENTS_DIR = os.path.join(DATA_DIR, "agents") + 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") ++_DEFAULT_SQLITE_FILE = os.path.join(DATA_DIR, "dbos_store.sqlite") + + # 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") ++DBOS_DATABASE_URL = os.environ.get( ++ "DBOS_SYSTEM_DATABASE_URL", f"sqlite:///{_DEFAULT_SQLITE_FILE}" ++) ++# DBOS enable switch is controlled solely via puppy.cfg using key 'enable_dbos'. ++# Default: True (DBOS enabled) unless explicitly disabled. ++ ++def get_use_dbos() -> bool: ++ """Return True if DBOS should be used based on 'enable_dbos' (default True).""" ++ cfg_val = get_value("enable_dbos") ++ if cfg_val is None: ++ return True ++ return str(cfg_val).strip().lower() in {"1", "true", "yes", "on"} ++ ++# 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: +@@ -184,6 +204,18 @@ def get_enable_streaming() -> bool: + return str(val).lower() in ("1", "true", "yes", "on") + + ++def get_command_timeout_seconds() -> int: ++ """Return shell command timeout in seconds (default: 30).""" ++ val = get_value("command_timeout_seconds") ++ if val is None: ++ return 30 ++ try: ++ parsed = int(str(val).strip()) ++ except (ValueError, TypeError): ++ return 30 ++ return max(1, parsed) ++ ++ + DEFAULT_SECTION = "puppy" + REQUIRED_KEYS = ["puppy_name", "owner_name"] + +diff --git a/code_puppy/plugins/frontend_emitter/emitter.py b/code_puppy/plugins/frontend_emitter/emitter.py +index 55836d90..cd0ffab6 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/session_storage.py b/code_puppy/session_storage.py +index d8b9201c..ab9a849b 100644 +--- a/code_puppy/session_storage.py ++++ b/code_puppy/session_storage.py +@@ -1,338 +1,238 @@ + """Shared helpers for persisting and restoring chat sessions. + +-This module centralises the pickle + metadata handling that used to live in +-both the CLI command handler and the auto-save feature. Keeping it here helps +-us avoid duplication while staying inside the Zen-of-Python sweet spot: simple +-is better than complex, nested side effects are worse than deliberate helpers. ++MIGRATION NOTE: Pickle-based storage has been removed. All sessions are now ++persisted exclusively to SQLite via code_puppy.api.db.queries. ++This module retains only stateless helpers (title generation) and stubs for ++removed functionality to maintain import compatibility during the transition. + """ + + from __future__ import annotations + + import json +-import pickle ++import re + from dataclasses import dataclass + from pathlib import Path +-from typing import Any, Callable, List ++from typing import Any, List + +- +-def _safe_loads(data: bytes) -> Any: +- """Deserialize pickle data.""" +- return pickle.loads(data) # noqa: S301 +- +- +-_LEGACY_SIGNED_HEADER = b"CPSESSION\x01" +-_LEGACY_SIGNATURE_SIZE = ( +- 32 # legacy signature bytes, retained only for backward-compat parsing +-) ++from pydantic import TypeAdapter ++from pydantic_ai.messages import ModelMessage + + SessionHistory = List[Any] +-TokenEstimator = Callable[[Any], int] +- + +-@dataclass(slots=True) +-class SessionPaths: +- pickle_path: Path +- metadata_path: Path + +- +-@dataclass(slots=True) ++@dataclass + class SessionMetadata: +- session_name: str +- timestamp: str ++ """Metadata returned after saving a session.""" ++ + message_count: int + total_tokens: int +- pickle_path: Path +- metadata_path: Path +- auto_saved: bool = False +- +- def as_serialisable(self) -> dict[str, Any]: +- return { +- "session_name": self.session_name, +- "timestamp": self.timestamp, +- "message_count": self.message_count, +- "total_tokens": self.total_tokens, +- "file_path": str(self.pickle_path), +- "auto_saved": self.auto_saved, +- } +- +- +-def _extract_pickle_payload(raw: bytes) -> bytes: +- """Return the pickle payload from raw session file bytes. +- +- New format is raw pickle bytes. +- Legacy format was: header + 32-byte signature + pickle payload. +- We no longer verify or generate signatures. +- """ +- if raw.startswith(_LEGACY_SIGNED_HEADER): +- offset = len(_LEGACY_SIGNED_HEADER) + _LEGACY_SIGNATURE_SIZE +- return raw[offset:] +- return raw ++ pickle_path: Path # Kept for compatibility, points to .json file ++ metadata_path: Path # Points to .json file (same as above) + + +-def ensure_directory(path: Path) -> Path: +- path.mkdir(parents=True, exist_ok=True) +- return path ++def generate_heuristic_title(history: SessionHistory, max_length: int = 50) -> str: ++ """Generate a short title from the first user message in the history. + ++ Extracts the first user message, takes the first ~50 chars, and converts ++ to a filename-safe format (lowercase, spaces to hyphens, remove special chars). + +-def build_session_paths(base_dir: Path, session_name: str) -> SessionPaths: +- pickle_path = base_dir / f"{session_name}.pkl" +- metadata_path = base_dir / f"{session_name}_meta.json" +- return SessionPaths(pickle_path=pickle_path, metadata_path=metadata_path) ++ Handles multiple message formats: ++ 1. Pydantic-ai format: msg.kind == 'request' with msg.parts[].content ++ 2. Enhanced/wrapped format: {'msg': , 'agent': str, ...} ++ 3. Simple dict format: {'role': 'user', 'content': str} ++ """ ++ ++ def extract_user_content(msg: Any) -> str | None: ++ """Extract user message content from various message formats.""" ++ # Handle wrapped/enhanced format: {'msg': , 'agent': ...} ++ if isinstance(msg, dict) and "msg" in msg: ++ msg = msg["msg"] ++ ++ # Handle pydantic-ai format: msg.kind == 'request' ++ if hasattr(msg, "kind") and msg.kind == "request": ++ for part in getattr(msg, "parts", []): ++ if hasattr(part, "content") and isinstance(part.content, str): ++ content = part.content.strip() ++ if content: ++ return content ++ ++ # Handle simple dict format: {'role': 'user', 'content': str} ++ if isinstance(msg, dict): ++ if msg.get("role") == "user" and isinstance(msg.get("content"), str): ++ content = msg["content"].strip() ++ if content: ++ return content ++ ++ return None ++ ++ def content_to_title(content: str) -> str: ++ """Convert content to a filename-safe kebab-case title.""" ++ # Take first line or first max_length chars ++ first_line = content.split("\n")[0][:max_length] ++ # Convert to kebab-case filename-safe format ++ title = first_line.lower() ++ title = re.sub(r"[^a-z0-9\s-]", "", title) # Remove special chars ++ title = re.sub(r"\s+", "-", title) # Spaces to hyphens ++ title = re.sub(r"-+", "-", title) # Collapse multiple hyphens ++ title = title.strip("-")[:max_length] ++ return title ++ ++ # Find first user message ++ for msg in history: ++ content = extract_user_content(msg) ++ if content: ++ title = content_to_title(content) ++ return title if title else "untitled-session" ++ ++ return "untitled-session" ++ ++ ++# --------------------------------------------------------------------------- ++# Deprecated / Removed Pickle Functionality -> Replaced with JSON for Export ++# --------------------------------------------------------------------------- + + + def save_session( +- *, + history: SessionHistory, + session_name: str, + base_dir: Path, +- timestamp: str, +- token_estimator: TokenEstimator, +- auto_saved: bool = False, ++ timestamp: str | None = None, ++ token_estimator: Any | None = None, ++ **kwargs: Any, + ) -> SessionMetadata: +- ensure_directory(base_dir) +- paths = build_session_paths(base_dir, session_name) +- +- pickle_data = pickle.dumps(history) +- tmp_pickle = paths.pickle_path.with_suffix(".tmp") +- with tmp_pickle.open("wb") as pickle_file: +- pickle_file.write(pickle_data) +- tmp_pickle.replace(paths.pickle_path) +- +- total_tokens = sum(token_estimator(message) for message in history) +- metadata = SessionMetadata( +- session_name=session_name, +- timestamp=timestamp, ++ """Save session history to a JSON file (replacing legacy pickle). ++ ++ This is used for 'pinning' or exporting sessions via /dump_context. ++ """ ++ base_dir.mkdir(parents=True, exist_ok=True) ++ file_path = base_dir / f"{session_name}.json" ++ ++ # Try to serialize the history, handling mixed types: ++ # - ModelMessage objects: serialize via Pydantic ++ # - System dicts: serialize as plain dicts ++ try: ++ from pydantic_ai.messages import ModelMessagesTypeAdapter ++ ++ history_data = [] ++ for item in history: ++ # Unwrap from {msg: ..., agent: ..., ts: ...} wrapper if present ++ if isinstance(item, dict) and "msg" in item: ++ inner = item["msg"] ++ wrapper_meta = {k: v for k, v in item.items() if k != "msg"} ++ ++ # Check if inner is a ModelMessage (has 'parts' attribute) ++ if hasattr(inner, "parts"): ++ # Serialize ModelMessage, then add wrapper metadata ++ try: ++ serialized = ModelMessagesTypeAdapter.dump_python( ++ [inner], mode="json" ++ )[0] ++ serialized["_wrapper"] = wrapper_meta ++ history_data.append(serialized) ++ except Exception: ++ # Fallback: keep as-is ++ history_data.append(item) ++ else: ++ # It's a system dict like {msg: 'system', ...} - keep as-is ++ history_data.append(item) ++ elif hasattr(item, "parts"): ++ # Direct ModelMessage without wrapper ++ try: ++ serialized = ModelMessagesTypeAdapter.dump_python( ++ [item], mode="json" ++ )[0] ++ history_data.append(serialized) ++ except Exception: ++ history_data.append({"_raw": str(item)}) ++ else: ++ # Unknown type - keep as-is ++ history_data.append(item) ++ ++ json_bytes = json.dumps(history_data, indent=2, default=str).encode("utf-8") ++ except Exception: ++ # Fallback: generic JSON dump ++ json_bytes = json.dumps(history, default=str, indent=2).encode("utf-8") ++ ++ file_path.write_bytes(json_bytes) ++ ++ # Calculate tokens if estimator provided ++ total_tokens = 0 ++ if token_estimator: ++ for msg in history: ++ try: ++ # estimator might expect the original object or dict ++ total_tokens += token_estimator(msg) ++ except Exception: ++ pass ++ ++ return SessionMetadata( + message_count=len(history), + total_tokens=total_tokens, +- pickle_path=paths.pickle_path, +- metadata_path=paths.metadata_path, +- auto_saved=auto_saved, ++ pickle_path=file_path, ++ metadata_path=file_path, + ) + +- tmp_metadata = paths.metadata_path.with_suffix(".tmp") +- with tmp_metadata.open("w", encoding="utf-8") as metadata_file: +- json.dump(metadata.as_serialisable(), metadata_file, indent=2) +- tmp_metadata.replace(paths.metadata_path) +- +- return metadata +- + + def load_session( + session_name: str, base_dir: Path, *, allow_legacy: bool = False + ) -> SessionHistory: +- # Kept for API compatibility; legacy loading is always supported now. +- _ = allow_legacy ++ """Load session history from a JSON file.""" ++ file_path = base_dir / f"{session_name}.json" ++ if not file_path.exists(): ++ # Fallback check for legacy .pkl if explicitly requested? ++ legacy_path = base_dir / f"{session_name}.pkl" ++ if allow_legacy and legacy_path.exists(): ++ raise NotImplementedError( ++ f"Found legacy pickle at {legacy_path} but pickle loading is removed." ++ ) ++ raise FileNotFoundError(f"Session file not found: {file_path}") + +- paths = build_session_paths(base_dir, session_name) +- if not paths.pickle_path.exists(): +- raise FileNotFoundError(paths.pickle_path) ++ json_data = file_path.read_bytes() + +- raw = paths.pickle_path.read_bytes() +- pickle_data = _extract_pickle_payload(raw) +- return _safe_loads(pickle_data) ++ # Try to load as ModelMessage objects ++ try: ++ adapter = TypeAdapter(List[ModelMessage]) ++ return adapter.validate_json(json_data) ++ except Exception: ++ # Fallback: return as list of dicts ++ return json.loads(json_data) + + + def list_sessions(base_dir: Path) -> List[str]: ++ """List available JSON sessions.""" + if not base_dir.exists(): + return [] +- return sorted(path.stem for path in base_dir.glob("*.pkl")) ++ return sorted([p.stem for p in base_dir.glob("*.json")]) + + + def cleanup_sessions(base_dir: Path, max_sessions: int) -> List[str]: +- if max_sessions <= 0: +- return [] ++ """Cleanup old sessions if count exceeds max_sessions. + +- if not base_dir.exists(): +- return [] +- +- candidate_paths = list(base_dir.glob("*.pkl")) +- if len(candidate_paths) <= max_sessions: +- return [] +- +- sorted_candidates = sorted( +- ((path.stat().st_mtime, path) for path in candidate_paths), +- key=lambda item: item[0], +- ) +- +- stale_entries = sorted_candidates[:-max_sessions] +- removed_sessions: List[str] = [] +- for _, pickle_path in stale_entries: +- metadata_path = base_dir / f"{pickle_path.stem}_meta.json" +- try: +- pickle_path.unlink(missing_ok=True) +- metadata_path.unlink(missing_ok=True) +- removed_sessions.append(pickle_path.stem) +- except OSError: +- continue +- +- return removed_sessions +- +- +-async def restore_autosave_interactively(base_dir: Path) -> None: +- """Prompt the user to load an autosave session from base_dir, if any exist. +- +- This helper is deliberately placed in session_storage to keep autosave +- restoration close to the persistence layer. It uses the same public APIs +- (list_sessions, load_session) and mirrors the interactive behaviours from +- the command handler. ++ Returns list of removed session names. + """ +- sessions = list_sessions(base_dir) +- if not sessions: +- return +- +- # Import locally to avoid pulling the messaging layer into storage modules +- from datetime import datetime ++ if not base_dir.exists() or max_sessions <= 0: ++ return [] + +- from prompt_toolkit.formatted_text import FormattedText ++ sessions = sorted(base_dir.glob("*.json"), key=lambda p: p.stat().st_mtime) + +- from code_puppy.agents.agent_manager import get_current_agent +- from code_puppy.command_line.prompt_toolkit_completion import ( +- get_input_with_combined_completion, +- ) +- from code_puppy.messaging import emit_success, emit_system_message, emit_warning +- +- entries = [] +- for name in sessions: +- meta_path = base_dir / f"{name}_meta.json" +- try: +- with meta_path.open("r", encoding="utf-8") as meta_file: +- data = json.load(meta_file) +- timestamp = data.get("timestamp") +- message_count = data.get("message_count") +- except Exception: +- timestamp = None +- message_count = None +- entries.append((name, timestamp, message_count)) +- +- def sort_key(entry): +- _, timestamp, _ = entry +- if timestamp: ++ removed = [] ++ if len(sessions) > max_sessions: ++ to_remove = sessions[: len(sessions) - max_sessions] ++ for p in to_remove: + try: +- return datetime.fromisoformat(timestamp) +- except ValueError: +- return datetime.min +- return datetime.min +- +- entries.sort(key=sort_key, reverse=True) +- +- PAGE_SIZE = 5 +- total = len(entries) +- page = 0 +- +- def render_page() -> None: +- start = page * PAGE_SIZE +- end = min(start + PAGE_SIZE, total) +- page_entries = entries[start:end] +- emit_system_message("Autosave Sessions Available:") +- for idx, (name, timestamp, message_count) in enumerate(page_entries, start=1): +- timestamp_display = timestamp or "unknown time" +- message_display = ( +- f"{message_count} messages" +- if message_count is not None +- else "unknown size" +- ) +- emit_system_message( +- f" [{idx}] {name} ({message_display}, saved at {timestamp_display})" +- ) +- # If there are more pages, offer next-page; show 'Return to first page' on last page +- if total > PAGE_SIZE: +- page_count = (total + PAGE_SIZE - 1) // PAGE_SIZE +- is_last_page = (page + 1) >= page_count +- remaining = total - (page * PAGE_SIZE + len(page_entries)) +- summary = ( +- f" and {remaining} more" if (remaining > 0 and not is_last_page) else "" +- ) +- 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 +- +- while True: +- render_page() +- try: +- selection = await get_input_with_combined_completion( +- FormattedText( +- [ +- ( +- "class:prompt", +- "Pick 1-5 to load, 6 for next, or name/Enter: ", +- ) +- ] +- ) +- ) +- except (KeyboardInterrupt, EOFError): +- emit_warning("Autosave selection cancelled") +- return +- +- selection = (selection or "").strip() +- if not selection: +- return +- +- # Numeric choice: 1-5 select within current page; 6 advances page +- if selection.isdigit(): +- num = int(selection) +- if num == 6 and total > PAGE_SIZE: +- page = (page + 1) % ((total + PAGE_SIZE - 1) // PAGE_SIZE) +- # loop and re-render next page +- continue +- if 1 <= num <= 5: +- start = page * PAGE_SIZE +- idx = start + (num - 1) +- if 0 <= idx < total: +- chosen_name = entries[idx][0] +- break +- else: +- emit_warning("Invalid selection for this page") +- continue +- emit_warning("Invalid selection; choose 1-5 or 6 for next") +- continue +- +- # Allow direct typing by exact session name +- for name, _ts, _mc in entries: +- if name == selection: +- chosen_name = name +- break +- if chosen_name: +- break +- emit_warning("No autosave loaded (invalid selection)") +- # keep looping and allow another try +- +- if not chosen_name: +- return ++ p.unlink() ++ removed.append(p.stem) ++ except Exception: ++ pass + +- try: +- history = load_session(chosen_name, base_dir) +- except FileNotFoundError: +- emit_warning(f"Autosave '{chosen_name}' could not be found") +- return +- except Exception as exc: +- emit_warning(f"Failed to load autosave '{chosen_name}': {exc}") +- return +- +- agent = get_current_agent() +- agent.set_message_history(history) +- +- # Set current autosave session id so subsequent autosaves overwrite this session +- try: +- from code_puppy.config import set_current_autosave_from_session_name ++ return removed + +- set_current_autosave_from_session_name(chosen_name) +- except Exception: +- pass +- +- total_tokens = sum(agent.estimate_tokens_for_message(msg) for msg in history) + +- session_path = base_dir / f"{chosen_name}.pkl" +- emit_success( +- f"✅ Autosave loaded: {len(history)} messages ({total_tokens} tokens)\n" +- f"📁 From: {session_path}" +- ) ++async def restore_autosave_interactively(base_dir: Path) -> None: ++ """Deprecated: No-op.""" ++ pass + +- # Display recent message history for context +- try: +- from code_puppy.command_line.autosave_menu import display_resumed_history + +- display_resumed_history(history) +- except Exception: +- pass # Don't fail if display doesn't work in non-TTY environment ++def build_session_paths(base_dir: Path, session_name: str) -> Any: ++ """Deprecated: Returns None or raises.""" ++ raise NotImplementedError("Session paths are no longer used.") +diff --git a/code_puppy/tools/command_runner.py b/code_puppy/tools/command_runner.py +index af67632c..ff03cd23 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 time + 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 + +@@ -18,6 +20,12 @@ from pydantic import BaseModel + from pydantic_ai import RunContext + from rich.text import Text + ++try: ++ from code_puppy.callbacks import on_run_shell_command_output ++except ImportError: ++ async def on_run_shell_command_output(*args, **kwargs): ++ return [] ++from code_puppy.config import get_command_timeout_seconds + from code_puppy.messaging import ( # Structured messaging types + AgentReasoningMessage, + ShellOutputMessage, +@@ -99,17 +107,104 @@ else: + + _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 +@@ -119,23 +214,97 @@ _ORIGINAL_SIGINT_HANDLER = None + _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 ++ ++ + # Thread pool for running blocking shell commands without blocking the event loop + # This allows multiple sub-agents to run shell commands in parallel + _SHELL_EXECUTOR = ThreadPoolExecutor(max_workers=16, thread_name_prefix="shell_cmd_") + + + 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: +@@ -211,13 +380,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: +@@ -235,7 +407,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 +@@ -243,32 +415,33 @@ 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 + def set_awaiting_user_input(awaiting=True): + """Set the flag indicating if user input is awaited.""" + if awaiting: +- _AWAITING_USER_INPUT.set() ++ _get_awaiting_input_event().set() + else: +- _AWAITING_USER_INPUT.clear() ++ _get_awaiting_input_event().clear() + + # When we're setting this flag, also pause/resume all active spinners + if awaiting: +@@ -330,27 +503,11 @@ def _listen_for_ctrl_x_windows( + stop_event: threading.Event, + on_escape: Callable[[], None], + ) -> None: +- """Windows-specific Ctrl-X listener. +- +- Pause-aware: while the agent is paused (Ctrl+T steering), we stop +- draining ``msvcrt.kbhit()`` so the steering editor can read input +- cleanly. See the POSIX sibling for the gory details. +- """ ++ """Windows-specific Ctrl-X listener.""" + import msvcrt + import time + +- def _is_agent_paused() -> bool: +- try: +- from code_puppy.messaging.pause_controller import get_pause_controller +- +- return get_pause_controller().is_paused() +- except Exception: +- return False +- + while not stop_event.is_set(): +- if _is_agent_paused(): +- time.sleep(0.05) +- continue + try: + if msvcrt.kbhit(): + try: +@@ -384,19 +541,10 @@ def _listen_for_ctrl_x_posix( + stop_event: threading.Event, + on_escape: Callable[[], None], + ) -> None: +- """POSIX-specific Ctrl-X listener. +- +- Pause-aware: while the ``PauseController`` is paused (Ctrl+T steering), +- we drop cbreak mode and stop reading stdin so the steering prompt's +- ``prompt_toolkit.Application`` can own the terminal cleanly. Without +- this, every other keystroke typed into the steer editor gets eaten by +- *this* listener's ``stdin.read(1)`` — the user sees half their input. +- On resume we re-acquire cbreak and continue. +- """ ++ """POSIX-specific Ctrl-X listener.""" + import select + import sys + import termios +- import time + import tty + + stdin = sys.stdin +@@ -409,50 +557,9 @@ def _listen_for_ctrl_x_posix( + except Exception: + return + +- cbreak_active = False +- +- def _enter_cbreak() -> None: +- nonlocal cbreak_active +- if not cbreak_active: +- tty.setcbreak(fd) +- cbreak_active = True +- +- def _exit_cbreak() -> None: +- nonlocal cbreak_active +- if cbreak_active: +- try: +- termios.tcsetattr(fd, termios.TCSADRAIN, original_attrs) +- except Exception: +- pass +- cbreak_active = False +- +- def _is_agent_paused() -> bool: +- """Lazy + exception-safe pause check — never crash the listener.""" +- try: +- from code_puppy.messaging.pause_controller import get_pause_controller +- +- return get_pause_controller().is_paused() +- except Exception: +- return False +- + try: +- _enter_cbreak() ++ tty.setcbreak(fd) + while not stop_event.is_set(): +- # Pause hand-off: release stdin and park until the steer +- # prompt finishes. Polling every 50ms keeps stop responsive. +- if _is_agent_paused(): +- _exit_cbreak() +- while _is_agent_paused() and not stop_event.is_set(): +- time.sleep(0.05) +- if stop_event.is_set(): +- return +- try: +- _enter_cbreak() +- except Exception: +- # Couldn't re-acquire raw mode — bail rather than spin. +- return +- continue +- + try: + read_ready, _, _ = select.select([stdin], [], [], 0.05) + except Exception: +@@ -470,11 +577,7 @@ def _listen_for_ctrl_x_posix( + "Ctrl+X handler raised unexpectedly; Ctrl+C still works." + ) + finally: +- _exit_cbreak() +- try: +- termios.tcsetattr(fd, termios.TCSADRAIN, original_attrs) +- except Exception: +- pass ++ termios.tcsetattr(fd, termios.TCSADRAIN, original_attrs) + + + def _spawn_ctrl_x_key_listener( +@@ -687,13 +790,15 @@ def run_shell_command_streaming( + silent: bool = False, + ): + stop_event = threading.Event() +- with _ACTIVE_STOP_EVENTS_LOCK: +- _ACTIVE_STOP_EVENTS.add(stop_event) ++ active_stop = _get_active_stop_events() ++ with _guarded_set(active_stop, _ACTIVE_STOP_EVENTS, _ACTIVE_STOP_EVENTS_LOCK): ++ active_stop.add(stop_event) + + start_time = time.time() + last_output_time = [start_time] + +- ABSOLUTE_TIMEOUT_SECONDS = 270 ++ # Get the user-configured absolute timeout for shell commands ++ ABSOLUTE_TIMEOUT_SECONDS = get_command_timeout_seconds() + + stdout_lines = [] + stderr_lines = [] +@@ -963,8 +1068,9 @@ def run_shell_command_streaming( + ) + 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) +@@ -978,7 +1084,7 @@ def run_shell_command_streaming( + 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( +@@ -992,8 +1098,9 @@ def run_shell_command_streaming( + ) + + 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, +@@ -1012,6 +1119,20 @@ 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: ++ try: ++ from code_puppy.plugins.walmart_specific.session_context import ( ++ get_session_working_directory, ++ ) ++ ++ cwd = get_session_working_directory() ++ except ImportError: ++ pass ++ + # Generate unique group_id for this command execution + group_id = generate_group_id("shell_command", command) + +@@ -1140,12 +1261,32 @@ async def run_shell_command( + # Check if we're running as a sub-agent (skip confirmation and run silently) + running_as_subagent = is_subagent() + ++ confirmation_lock_acquired = False ++ ++ # 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() ++ ): ++ confirmation_lock_acquired = _CONFIRMATION_LOCK.acquire(blocking=False) ++ if not confirmation_lock_acquired: ++ return ShellCommandOutput( ++ success=False, ++ command=command, ++ error="Another command is currently awaiting confirmation", ++ ) + + # Get puppy name for personalized messages + from code_puppy.config import get_puppy_name +@@ -1163,8 +1304,7 @@ async def run_shell_command( + panel_content.append("📂 Working directory: ", style="dim") + panel_content.append(cwd, style="dim cyan") + +- # Use the common approval function (async version). +- # Internal queueing means parallel calls wait their turn here. ++ # Use the common approval function (async version) + confirmed, user_feedback = await get_user_approval_async( + title="Shell Command", + content=panel_content, +@@ -1173,6 +1313,10 @@ async def run_shell_command( + puppy_name=puppy_name, + ) + ++ # Release lock after approval ++ if confirmation_lock_acquired: ++ _CONFIRMATION_LOCK.release() ++ + if not confirmed: + if user_feedback: + result = ShellCommandOutput( +@@ -1313,9 +1457,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: +@@ -1392,7 +1540,9 @@ def register_agent_run_shell_command(agent): + + Supports streaming output, timeout handling, and background execution. + """ +- return await run_shell_command(context, command, cwd, timeout, background) ++ result = await run_shell_command(context, command, cwd, timeout, background) ++ await on_run_shell_command_output(result) ++ return result + + + def register_agent_share_your_reasoning(agent): +diff --git a/pyproject.toml b/pyproject.toml +index a41eef50..1c433bcc 100644 +--- a/pyproject.toml ++++ b/pyproject.toml +@@ -9,6 +9,11 @@ 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", ++ "dbos>=2.11.0", + "pydantic-ai-slim[openai,anthropic,mcp]==1.56.0", + "typer>=0.12.0", + "mcp>=1.9.4", +@@ -63,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/uv.lock b/uv.lock +index 3d7e28f8..e23624b4 100644 +--- a/uv.lock ++++ b/uv.lock +@@ -2,6 +2,15 @@ version = 1 + revision = 3 + requires-python = ">=3.11, <3.15" + ++[[package]] ++name = "aiosqlite" ++version = "0.22.1" ++source = { registry = "https://pypi.org/simple" } ++sdist = { url = "https://files.pythonhosted.org/packages/4e/8a/64761f4005f17809769d23e518d915db74e6310474e733e3593cfc854ef1/aiosqlite-0.22.1.tar.gz", hash = "sha256:043e0bd78d32888c0a9ca90fc788b38796843360c855a7262a532813133a0650", size = 14821, upload-time = "2025-12-23T19:25:43.997Z" } ++wheels = [ ++ { url = "https://files.pythonhosted.org/packages/00/b7/e3bf5133d697a08128598c8d0abc5e16377b51465a33756de24fa7dee953/aiosqlite-0.22.1-py3-none-any.whl", hash = "sha256:21c002eb13823fad740196c5a2e9d8e62f6243bd9e7e4a1f87fb5e44ecb4fceb", size = 17405, upload-time = "2025-12-23T19:25:42.139Z" }, ++] ++ + [[package]] + name = "annotated-doc" + version = "0.0.4" +@@ -303,9 +312,12 @@ name = "code-puppy" + version = "0.0.527" + source = { editable = "." } + dependencies = [ ++ { name = "aiosqlite" }, + { name = "anthropic" }, + { name = "azure-identity" }, + { name = "boto3" }, ++ { name = "dbos" }, ++ { name = "fastapi" }, + { name = "httpx", extra = ["http2"] }, + { name = "json-repair" }, + { name = "mcp" }, +@@ -323,6 +335,8 @@ dependencies = [ + { name = "ripgrep" }, + { name = "termflow-md" }, + { name = "typer" }, ++ { name = "uvicorn", extra = ["standard"] }, ++ { name = "websockets" }, + ] + + [package.optional-dependencies] +@@ -344,11 +358,14 @@ dev = [ + + [package.metadata] + requires-dist = [ ++ { name = "aiosqlite", specifier = ">=0.22.1" }, + { name = "anthropic", specifier = "==0.79.0" }, + { name = "azure-identity", specifier = ">=1.15.0" }, + { name = "boto3", specifier = ">=1.43.9" }, + { name = "boto3", marker = "extra == 'bedrock'", specifier = ">=1.35.0" }, ++ { name = "dbos", specifier = ">=2.11.0" }, + { name = "dbos", marker = "extra == 'durable'", specifier = ">=2.11.0" }, ++ { name = "fastapi", specifier = ">=0.109.0" }, + { name = "httpx", extras = ["http2"], specifier = ">=0.24.1" }, + { name = "json-repair", specifier = ">=0.46.2" }, + { name = "mcp", specifier = ">=1.9.4" }, +@@ -366,6 +383,8 @@ requires-dist = [ + { name = "ripgrep", specifier = "==14.1.0" }, + { name = "termflow-md", specifier = ">=0.1.11" }, + { name = "typer", specifier = ">=0.12.0" }, ++ { name = "uvicorn", extras = ["standard"], specifier = ">=0.27.0" }, ++ { name = "websockets", specifier = ">=12.0" }, + ] + provides-extras = ["bedrock", "durable"] + +@@ -585,6 +604,22 @@ wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/5f/ed01f9a3cdffbd5a008556fc7b2a08ddb1cc6ace7effa7340604b1d16699/docstring_parser-0.18.0-py3-none-any.whl", hash = "sha256:b3fcbed555c47d8479be0796ef7e19c2670d428d72e96da63f3a40122860374b", size = 22484, upload-time = "2026-04-14T04:09:18.638Z" }, + ] + ++[[package]] ++name = "fastapi" ++version = "0.136.3" ++source = { registry = "https://pypi.org/simple" } ++dependencies = [ ++ { name = "annotated-doc" }, ++ { name = "pydantic" }, ++ { name = "starlette" }, ++ { name = "typing-extensions" }, ++ { name = "typing-inspection" }, ++] ++sdist = { url = "https://files.pythonhosted.org/packages/81/2d/ff8d91d7b564d464629a0fd50a4489c97fcb836ac230bf3a7269232a9b1f/fastapi-0.136.3.tar.gz", hash = "sha256:e487fae93ad408e6f47641ee4dfe389864fd7bec92e547ea8498fc13f43e83ab", size = 396410, upload-time = "2026-05-23T18:53:15.192Z" } ++wheels = [ ++ { url = "https://files.pythonhosted.org/packages/e0/82/45359b62a067409bd929ae8a56b8ed13e5a8c8a61194b3c236920999ab83/fastapi-0.136.3-py3-none-any.whl", hash = "sha256:3d2a69bdf04b7e9f3afa292c3bc7a98816bbfafa10bc9b45f3f3700d2f761620", size = 117481, upload-time = "2026-05-23T18:53:16.924Z" }, ++] ++ + [[package]] + name = "genai-prices" + version = "0.0.60" +@@ -734,6 +769,49 @@ wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, + ] + ++[[package]] ++name = "httptools" ++version = "0.8.0" ++source = { registry = "https://pypi.org/simple" } ++sdist = { url = "https://files.pythonhosted.org/packages/43/e5/d471fcb0e14523fe1c3f4ba58ca52480e7bd70ad7109a3846bc75892f7fb/httptools-0.8.0.tar.gz", hash = "sha256:6b2a32f18d97e16e90827d7a819ffa8dbd8cc245fc4e1fa9d1095b54ef4bd999", size = 271342, upload-time = "2026-05-25T22:17:48.841Z" } ++wheels = [ ++ { url = "https://files.pythonhosted.org/packages/f8/d2/c3eedaef57de65c3cc5f8dc244cf12d09c84ad258a479055aad6db23206c/httptools-0.8.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:ed377e64805bdba4943c82717333f8f8603a13b09aff9cead2717c6c817fb168", size = 208428, upload-time = "2026-05-25T22:16:59.717Z" }, ++ { url = "https://files.pythonhosted.org/packages/f1/94/dfe435d90d0ef61ec0f2cc3d480eef78c59727c6c2ce039f433882f6131a/httptools-0.8.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:9518c406d7b310f05adb1a37f80acabac40504a575d7c0da6d3e365c695ac20d", size = 113366, upload-time = "2026-05-25T22:17:00.795Z" }, ++ { url = "https://files.pythonhosted.org/packages/cc/d4/13025f1a56e615dcb331e0bbe2d9a1143212b58c263385fc5d2e558f5bac/httptools-0.8.0-cp311-cp311-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:57278e6fa0424c42a8a3e454828ab4f0aff27b40cddf9679579b98c6dce6a376", size = 464676, upload-time = "2026-05-25T22:17:02.014Z" }, ++ { url = "https://files.pythonhosted.org/packages/bf/95/4c1c26c0b985f8a3331682d802598f14e32dc41bf7509266eb2c04ad4801/httptools-0.8.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bbb8caadb2b742d293169d2b458b5c001ef70e3158704aa3d3ef9597624c5d1d", size = 464235, upload-time = "2026-05-25T22:17:03.109Z" }, ++ { url = "https://files.pythonhosted.org/packages/a2/82/6735be2b0ca527718c431cdb8e5f70c3862c0844a687df0f572c51e11497/httptools-0.8.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:52dd695b865fe96d9d2b16b64a895f3f57bf3cb064e8383cd3b5713a069e8085", size = 449809, upload-time = "2026-05-25T22:17:04.443Z" }, ++ { url = "https://files.pythonhosted.org/packages/b5/f9/5811c74f37a758c8a4aa3dc430375119d335947e883efc4664d8f3559a41/httptools-0.8.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:20b4aac66ff65f7db06a375808b78f42a94970aa22e826b3cb2b43eb09174124", size = 452174, upload-time = "2026-05-25T22:17:05.476Z" }, ++ { url = "https://files.pythonhosted.org/packages/cc/94/97b75870dea07b71e3ec535cebe525b08d723152e4c7d13fa887e51f4de2/httptools-0.8.0-cp311-cp311-win_amd64.whl", hash = "sha256:a1b4c8e7a489a0d750d91894e9a8cdc295838f1924c0ca903ae993456fddec07", size = 90991, upload-time = "2026-05-25T22:17:06.75Z" }, ++ { url = "https://files.pythonhosted.org/packages/14/88/1d21a36da8f5cb0fa49eafd4b169eba5608d57e75bbcf61845cbc6243216/httptools-0.8.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:880490234c10f70a9830743097e8958d6e4b9f5a0ffc24515023afeef984054d", size = 208247, upload-time = "2026-05-25T22:17:07.843Z" }, ++ { url = "https://files.pythonhosted.org/packages/a5/42/cc4feea2945cb3051038f090c9b36bd5b8a9d7f5a894a506a8983e33fd1c/httptools-0.8.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:5931891fb7b441b8a3853cf1b85c82c903defce084dd5f6771ca46e31bf862c5", size = 113064, upload-time = "2026-05-25T22:17:09.136Z" }, ++ { url = "https://files.pythonhosted.org/packages/e3/a6/febbb8b8db0f58b38e44ad6cb946e6a255ae49b55f2e8543408fb7501ccd/httptools-0.8.0-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:b15fc622b0f869d19207c4089a501d9bcc63ca5e071ffdd2f03f922df882dcb2", size = 523851, upload-time = "2026-05-25T22:17:10.106Z" }, ++ { url = "https://files.pythonhosted.org/packages/b7/e4/f90a0df0b83beff265b7e3b65f2a4cefd95792d4be0ac3e16049f2acd3c2/httptools-0.8.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:425f83884fd6343828d8c565f046cb72b6d19063f6924093e11bcd8e1548cd09", size = 518842, upload-time = "2026-05-25T22:17:11.218Z" }, ++ { url = "https://files.pythonhosted.org/packages/9e/2d/0c9ac76dd2c893841fbf6498d6acec4f2442e1b7067f6e3e316a80e494e8/httptools-0.8.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ef7c3c97f4311c7be57e2986629df89d49cb434dbff78eafcd48c2bff986b15a", size = 501238, upload-time = "2026-05-25T22:17:12.728Z" }, ++ { url = "https://files.pythonhosted.org/packages/ca/42/906adc91ae3a5fa9c59c0a2f21c139725bd7e5b41ae6acd485cd14123ebf/httptools-0.8.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a1afd7c9fbff0d9f5d489c4ce2768bd09c84a46ddefc7161e6aa82ae35c85745", size = 509567, upload-time = "2026-05-25T22:17:13.842Z" }, ++ { url = "https://files.pythonhosted.org/packages/05/0b/4240efeb672751ee5b9b380cb0e3fdc050bc05f68adc7a8aefc4fcd9a69a/httptools-0.8.0-cp312-cp312-win_amd64.whl", hash = "sha256:cd96f29b4bab1d42fa6e3d008711c75e0f79e94e06827330160e3a304227f150", size = 90918, upload-time = "2026-05-25T22:17:15.155Z" }, ++ { url = "https://files.pythonhosted.org/packages/5e/e5/8cfcabc5546e8022f168be28bcdaa128a240a0befdd03b59d558b4f18bd6/httptools-0.8.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:614ceea8ea606848bece2338ac03b3ce5324bcb4be8dc7d377ed708012fa4db8", size = 205148, upload-time = "2026-05-25T22:17:16.333Z" }, ++ { url = "https://files.pythonhosted.org/packages/2a/0e/0fb14848c19a686c8062ff9067c1a48793e3224b47bc5b201535b6036fce/httptools-0.8.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:2d689918c15a013c65ef52d9fd495d766893ab831a2c8d89f2ac5940a5df847c", size = 111368, upload-time = "2026-05-25T22:17:17.586Z" }, ++ { url = "https://files.pythonhosted.org/packages/2e/1b/46f1cecf06b9bbde8e4b8c88034ac7908989e5ff7a3a388ef38392949c1f/httptools-0.8.0-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:eb3028cca2fc0a6d720e52ef61d8ebb62fcbfeb1de56874546d858d3f25a26b7", size = 486447, upload-time = "2026-05-25T22:17:18.564Z" }, ++ { url = "https://files.pythonhosted.org/packages/77/00/258bfc0837221f81d9725c45f9b948a6a6b2994a147a4fb66e85100c668f/httptools-0.8.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:88bdd940f2b5d487b4d032c6afa5489a7dc4694410d43de3c38c4fb3af0dc45d", size = 482448, upload-time = "2026-05-25T22:17:19.912Z" }, ++ { url = "https://files.pythonhosted.org/packages/04/ab/d1cef3b5523f4d272a70f42a776c3169a2dddfe3a54de4b2ce4a36341528/httptools-0.8.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:6a43c9dd399758ccc0531acb0a3c4a6c299ee893ee9400e9c893b7bdcfae0681", size = 464460, upload-time = "2026-05-25T22:17:20.882Z" }, ++ { url = "https://files.pythonhosted.org/packages/ce/48/5d1d072442277bb2b3434e0e60690b8e8c23840ef7de8b6ea54040a536d3/httptools-0.8.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:0770728beb05094c809b98e814edff5fef69d26ad7d21185f2f6d5884a0ba683", size = 471312, upload-time = "2026-05-25T22:17:22.085Z" }, ++ { url = "https://files.pythonhosted.org/packages/0d/66/b96623b27e51a68199ef4efdda0613cced9233fe3062ac74e50749c5ad37/httptools-0.8.0-cp313-cp313-win_amd64.whl", hash = "sha256:7685df791fad561384bfb139e77fde27a1ffd93134e016f95a0db424ffbf77b1", size = 90117, upload-time = "2026-05-25T22:17:23.074Z" }, ++ { url = "https://files.pythonhosted.org/packages/1a/12/fa3fbf5f9517b273edea2dc982aa82a8c634091e67c590792b729017bc6f/httptools-0.8.0-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:de242a49b5d18e0a8776e654e9f6bf6d89f3875a5c35b425a0e7ce940feb3fd6", size = 206183, upload-time = "2026-05-25T22:17:24.004Z" }, ++ { url = "https://files.pythonhosted.org/packages/30/fc/5e7c4cb443370f2090a3aba0453a07384d29ff66b7435bb90e77e1037599/httptools-0.8.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:159e9ab5f701ccd42e555a12f1ad8ff69702910fc1c996cf2bb66e5fcb7a231b", size = 112079, upload-time = "2026-05-25T22:17:25.216Z" }, ++ { url = "https://files.pythonhosted.org/packages/ba/53/771bd891eb0f236f32145d6a1775777ec85745f3cc983a1f23d1a3b8ddfe/httptools-0.8.0-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:c4a9f1707e4823d54dfec6c33fa3697d302aed536ed352a7ebb5a061ddb869d0", size = 481596, upload-time = "2026-05-25T22:17:26.186Z" }, ++ { url = "https://files.pythonhosted.org/packages/62/42/94e15bc68ce3d423243c45d7f1b0c7561f13844f97dc52ae23182fb65628/httptools-0.8.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d76ad7b951387e3632c8716a9bb03ac5b45c5f16119aa409db0459520887944e", size = 480865, upload-time = "2026-05-25T22:17:27.542Z" }, ++ { url = "https://files.pythonhosted.org/packages/1c/7c/fe2980fc03723272e30f135b62360b075f513dfe7cc73aef36c7f04012bd/httptools-0.8.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:a3b7387147361c3fd47a0bde763c5c91b5b4cd4dc9989b8ece84ff436c99843b", size = 463189, upload-time = "2026-05-25T22:17:28.546Z" }, ++ { url = "https://files.pythonhosted.org/packages/15/1b/47fc5fff68acd1bfa20b4734059c9a06cadb88119dcd5258b5b0d21d91c8/httptools-0.8.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f256d6ce930c52ca1cb2a960b7da03548c454e7d28b06059ad41bfe789036ce0", size = 466610, upload-time = "2026-05-25T22:17:29.816Z" }, ++ { url = "https://files.pythonhosted.org/packages/60/bd/07b13c93ffd9bec9546e0d43f8e19378dd696dbd278511406bc07371ef1f/httptools-0.8.0-cp314-cp314-win_amd64.whl", hash = "sha256:19d1ee275bb59ba2643ba9a3a1e51cc0c788caf2b8df506368e03f56fdd08527", size = 92705, upload-time = "2026-05-25T22:17:31.133Z" }, ++ { url = "https://files.pythonhosted.org/packages/fd/c4/121648f68ce066d7bd762d6b6d97e620847642d38d54f3d90ff11d947629/httptools-0.8.0-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:de1ed58a974e75d56560acc7e7fed01a454994429456f65209789992e41f2568", size = 215023, upload-time = "2026-05-25T22:17:32.401Z" }, ++ { url = "https://files.pythonhosted.org/packages/b9/b0/312a062ae741ae3e8baa8c8bf20be81b2e67337b259ab4349bebc7b6142e/httptools-0.8.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:e93c227b595c6926c1acee96891dd9da4be338cfbe82e5cd3bb9d8dd7dc4ac0b", size = 117405, upload-time = "2026-05-25T22:17:33.742Z" }, ++ { url = "https://files.pythonhosted.org/packages/fc/37/fccd705f795386bb05bf413012fecff2a33e5aa8c2f069096de3e9fd8702/httptools-0.8.0-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:2a021c3a8e65cc125390d72f59b968afca3bdcaff25bd67965e0a055a14946ca", size = 558497, upload-time = "2026-05-25T22:17:34.732Z" }, ++ { url = "https://files.pythonhosted.org/packages/bd/39/f172e8003576de35f5ba77ff417cf0e34429d35dc014deef15afa337a72c/httptools-0.8.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:48774d39cbb70e2b1f71f88852a3087ae1d3a1eb80482bb48c13067ab080c14f", size = 571585, upload-time = "2026-05-25T22:17:35.813Z" }, ++ { url = "https://files.pythonhosted.org/packages/3e/b9/f5564760af99f3dbbf3f9104dc00e5da27e96cf433c6bdcf77617f70bf3f/httptools-0.8.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:88eead8ec8680a9f146c655bc88445a325bd7921cfd8194c7337e9467282427d", size = 543297, upload-time = "2026-05-25T22:17:37.08Z" }, ++ { url = "https://files.pythonhosted.org/packages/99/67/8d9f2c313618e161b82f3873188e7196126da1d6e29688df40eb3997c77a/httptools-0.8.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:2c032fa028f46871ec7e1fc59fc15e8023eab3e6bbe6ece786a1611719a5d081", size = 539535, upload-time = "2026-05-25T22:17:38.032Z" }, ++ { url = "https://files.pythonhosted.org/packages/48/63/b906c01e53f50d432c0defe43ce52764a111dc1bdd028bafbeb54dcfd008/httptools-0.8.0-cp314-cp314t-win_amd64.whl", hash = "sha256:384c17174464c8e873398b7af24f0b1f44d992c820328413951a625323155d77", size = 108209, upload-time = "2026-05-25T22:17:39.473Z" }, ++] ++ + [[package]] + name = "httpx" + version = "0.28.1" +@@ -2367,6 +2445,148 @@ wheels = [ + { url = "https://files.pythonhosted.org/packages/15/41/ac2dfdbc1f60c7af4f994c7a335cfa7040c01642b605d65f611cecc2a1e4/uvicorn-0.47.0-py3-none-any.whl", hash = "sha256:2c5715bc12d1892d84752049f400cd1c3cb018514967fdfeb97640443a6a9432", size = 71301, upload-time = "2026-05-14T18:16:51.762Z" }, + ] + ++[package.optional-dependencies] ++standard = [ ++ { name = "colorama", marker = "sys_platform == 'win32'" }, ++ { name = "httptools" }, ++ { name = "python-dotenv" }, ++ { name = "pyyaml" }, ++ { name = "uvloop", marker = "platform_python_implementation != 'PyPy' and sys_platform != 'cygwin' and sys_platform != 'win32'" }, ++ { name = "watchfiles" }, ++ { name = "websockets" }, ++] ++ ++[[package]] ++name = "uvloop" ++version = "0.22.1" ++source = { registry = "https://pypi.org/simple" } ++sdist = { url = "https://files.pythonhosted.org/packages/06/f0/18d39dbd1971d6d62c4629cc7fa67f74821b0dc1f5a77af43719de7936a7/uvloop-0.22.1.tar.gz", hash = "sha256:6c84bae345b9147082b17371e3dd5d42775bddce91f885499017f4607fdaf39f", size = 2443250, upload-time = "2025-10-16T22:17:19.342Z" } ++wheels = [ ++ { url = "https://files.pythonhosted.org/packages/c7/d5/69900f7883235562f1f50d8184bb7dd84a2fb61e9ec63f3782546fdbd057/uvloop-0.22.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:c60ebcd36f7b240b30788554b6f0782454826a0ed765d8430652621b5de674b9", size = 1352420, upload-time = "2025-10-16T22:16:21.187Z" }, ++ { url = "https://files.pythonhosted.org/packages/a8/73/c4e271b3bce59724e291465cc936c37758886a4868787da0278b3b56b905/uvloop-0.22.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:3b7f102bf3cb1995cfeaee9321105e8f5da76fdb104cdad8986f85461a1b7b77", size = 748677, upload-time = "2025-10-16T22:16:22.558Z" }, ++ { url = "https://files.pythonhosted.org/packages/86/94/9fb7fad2f824d25f8ecac0d70b94d0d48107ad5ece03769a9c543444f78a/uvloop-0.22.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:53c85520781d84a4b8b230e24a5af5b0778efdb39142b424990ff1ef7c48ba21", size = 3753819, upload-time = "2025-10-16T22:16:23.903Z" }, ++ { url = "https://files.pythonhosted.org/packages/74/4f/256aca690709e9b008b7108bc85fba619a2bc37c6d80743d18abad16ee09/uvloop-0.22.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:56a2d1fae65fd82197cb8c53c367310b3eabe1bbb9fb5a04d28e3e3520e4f702", size = 3804529, upload-time = "2025-10-16T22:16:25.246Z" }, ++ { url = "https://files.pythonhosted.org/packages/7f/74/03c05ae4737e871923d21a76fe28b6aad57f5c03b6e6bfcfa5ad616013e4/uvloop-0.22.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:40631b049d5972c6755b06d0bfe8233b1bd9a8a6392d9d1c45c10b6f9e9b2733", size = 3621267, upload-time = "2025-10-16T22:16:26.819Z" }, ++ { url = "https://files.pythonhosted.org/packages/75/be/f8e590fe61d18b4a92070905497aec4c0e64ae1761498cad09023f3f4b3e/uvloop-0.22.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:535cc37b3a04f6cd2c1ef65fa1d370c9a35b6695df735fcff5427323f2cd5473", size = 3723105, upload-time = "2025-10-16T22:16:28.252Z" }, ++ { url = "https://files.pythonhosted.org/packages/3d/ff/7f72e8170be527b4977b033239a83a68d5c881cc4775fca255c677f7ac5d/uvloop-0.22.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:fe94b4564e865d968414598eea1a6de60adba0c040ba4ed05ac1300de402cd42", size = 1359936, upload-time = "2025-10-16T22:16:29.436Z" }, ++ { url = "https://files.pythonhosted.org/packages/c3/c6/e5d433f88fd54d81ef4be58b2b7b0cea13c442454a1db703a1eea0db1a59/uvloop-0.22.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:51eb9bd88391483410daad430813d982010f9c9c89512321f5b60e2cddbdddd6", size = 752769, upload-time = "2025-10-16T22:16:30.493Z" }, ++ { url = "https://files.pythonhosted.org/packages/24/68/a6ac446820273e71aa762fa21cdcc09861edd3536ff47c5cd3b7afb10eeb/uvloop-0.22.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:700e674a166ca5778255e0e1dc4e9d79ab2acc57b9171b79e65feba7184b3370", size = 4317413, upload-time = "2025-10-16T22:16:31.644Z" }, ++ { url = "https://files.pythonhosted.org/packages/5f/6f/e62b4dfc7ad6518e7eff2516f680d02a0f6eb62c0c212e152ca708a0085e/uvloop-0.22.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7b5b1ac819a3f946d3b2ee07f09149578ae76066d70b44df3fa990add49a82e4", size = 4426307, upload-time = "2025-10-16T22:16:32.917Z" }, ++ { url = "https://files.pythonhosted.org/packages/90/60/97362554ac21e20e81bcef1150cb2a7e4ffdaf8ea1e5b2e8bf7a053caa18/uvloop-0.22.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e047cc068570bac9866237739607d1313b9253c3051ad84738cbb095be0537b2", size = 4131970, upload-time = "2025-10-16T22:16:34.015Z" }, ++ { url = "https://files.pythonhosted.org/packages/99/39/6b3f7d234ba3964c428a6e40006340f53ba37993f46ed6e111c6e9141d18/uvloop-0.22.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:512fec6815e2dd45161054592441ef76c830eddaad55c8aa30952e6fe1ed07c0", size = 4296343, upload-time = "2025-10-16T22:16:35.149Z" }, ++ { url = "https://files.pythonhosted.org/packages/89/8c/182a2a593195bfd39842ea68ebc084e20c850806117213f5a299dfc513d9/uvloop-0.22.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:561577354eb94200d75aca23fbde86ee11be36b00e52a4eaf8f50fb0c86b7705", size = 1358611, upload-time = "2025-10-16T22:16:36.833Z" }, ++ { url = "https://files.pythonhosted.org/packages/d2/14/e301ee96a6dc95224b6f1162cd3312f6d1217be3907b79173b06785f2fe7/uvloop-0.22.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:1cdf5192ab3e674ca26da2eada35b288d2fa49fdd0f357a19f0e7c4e7d5077c8", size = 751811, upload-time = "2025-10-16T22:16:38.275Z" }, ++ { url = "https://files.pythonhosted.org/packages/b7/02/654426ce265ac19e2980bfd9ea6590ca96a56f10c76e63801a2df01c0486/uvloop-0.22.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6e2ea3d6190a2968f4a14a23019d3b16870dd2190cd69c8180f7c632d21de68d", size = 4288562, upload-time = "2025-10-16T22:16:39.375Z" }, ++ { url = "https://files.pythonhosted.org/packages/15/c0/0be24758891ef825f2065cd5db8741aaddabe3e248ee6acc5e8a80f04005/uvloop-0.22.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0530a5fbad9c9e4ee3f2b33b148c6a64d47bbad8000ea63704fa8260f4cf728e", size = 4366890, upload-time = "2025-10-16T22:16:40.547Z" }, ++ { url = "https://files.pythonhosted.org/packages/d2/53/8369e5219a5855869bcee5f4d317f6da0e2c669aecf0ef7d371e3d084449/uvloop-0.22.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:bc5ef13bbc10b5335792360623cc378d52d7e62c2de64660616478c32cd0598e", size = 4119472, upload-time = "2025-10-16T22:16:41.694Z" }, ++ { url = "https://files.pythonhosted.org/packages/f8/ba/d69adbe699b768f6b29a5eec7b47dd610bd17a69de51b251126a801369ea/uvloop-0.22.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:1f38ec5e3f18c8a10ded09742f7fb8de0108796eb673f30ce7762ce1b8550cad", size = 4239051, upload-time = "2025-10-16T22:16:43.224Z" }, ++ { url = "https://files.pythonhosted.org/packages/90/cd/b62bdeaa429758aee8de8b00ac0dd26593a9de93d302bff3d21439e9791d/uvloop-0.22.1-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:3879b88423ec7e97cd4eba2a443aa26ed4e59b45e6b76aabf13fe2f27023a142", size = 1362067, upload-time = "2025-10-16T22:16:44.503Z" }, ++ { url = "https://files.pythonhosted.org/packages/0d/f8/a132124dfda0777e489ca86732e85e69afcd1ff7686647000050ba670689/uvloop-0.22.1-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:4baa86acedf1d62115c1dc6ad1e17134476688f08c6efd8a2ab076e815665c74", size = 752423, upload-time = "2025-10-16T22:16:45.968Z" }, ++ { url = "https://files.pythonhosted.org/packages/a3/94/94af78c156f88da4b3a733773ad5ba0b164393e357cc4bd0ab2e2677a7d6/uvloop-0.22.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:297c27d8003520596236bdb2335e6b3f649480bd09e00d1e3a99144b691d2a35", size = 4272437, upload-time = "2025-10-16T22:16:47.451Z" }, ++ { url = "https://files.pythonhosted.org/packages/b5/35/60249e9fd07b32c665192cec7af29e06c7cd96fa1d08b84f012a56a0b38e/uvloop-0.22.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c1955d5a1dd43198244d47664a5858082a3239766a839b2102a269aaff7a4e25", size = 4292101, upload-time = "2025-10-16T22:16:49.318Z" }, ++ { url = "https://files.pythonhosted.org/packages/02/62/67d382dfcb25d0a98ce73c11ed1a6fba5037a1a1d533dcbb7cab033a2636/uvloop-0.22.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:b31dc2fccbd42adc73bc4e7cdbae4fc5086cf378979e53ca5d0301838c5682c6", size = 4114158, upload-time = "2025-10-16T22:16:50.517Z" }, ++ { url = "https://files.pythonhosted.org/packages/f0/7a/f1171b4a882a5d13c8b7576f348acfe6074d72eaf52cccef752f748d4a9f/uvloop-0.22.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:93f617675b2d03af4e72a5333ef89450dfaa5321303ede6e67ba9c9d26878079", size = 4177360, upload-time = "2025-10-16T22:16:52.646Z" }, ++ { url = "https://files.pythonhosted.org/packages/79/7b/b01414f31546caf0919da80ad57cbfe24c56b151d12af68cee1b04922ca8/uvloop-0.22.1-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:37554f70528f60cad66945b885eb01f1bb514f132d92b6eeed1c90fd54ed6289", size = 1454790, upload-time = "2025-10-16T22:16:54.355Z" }, ++ { url = "https://files.pythonhosted.org/packages/d4/31/0bb232318dd838cad3fa8fb0c68c8b40e1145b32025581975e18b11fab40/uvloop-0.22.1-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:b76324e2dc033a0b2f435f33eb88ff9913c156ef78e153fb210e03c13da746b3", size = 796783, upload-time = "2025-10-16T22:16:55.906Z" }, ++ { url = "https://files.pythonhosted.org/packages/42/38/c9b09f3271a7a723a5de69f8e237ab8e7803183131bc57c890db0b6bb872/uvloop-0.22.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:badb4d8e58ee08dad957002027830d5c3b06aea446a6a3744483c2b3b745345c", size = 4647548, upload-time = "2025-10-16T22:16:57.008Z" }, ++ { url = "https://files.pythonhosted.org/packages/c1/37/945b4ca0ac27e3dc4952642d4c900edd030b3da6c9634875af6e13ae80e5/uvloop-0.22.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b91328c72635f6f9e0282e4a57da7470c7350ab1c9f48546c0f2866205349d21", size = 4467065, upload-time = "2025-10-16T22:16:58.206Z" }, ++ { url = "https://files.pythonhosted.org/packages/97/cc/48d232f33d60e2e2e0b42f4e73455b146b76ebe216487e862700457fbf3c/uvloop-0.22.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:daf620c2995d193449393d6c62131b3fbd40a63bf7b307a1527856ace637fe88", size = 4328384, upload-time = "2025-10-16T22:16:59.36Z" }, ++ { url = "https://files.pythonhosted.org/packages/e4/16/c1fd27e9549f3c4baf1dc9c20c456cd2f822dbf8de9f463824b0c0357e06/uvloop-0.22.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6cde23eeda1a25c75b2e07d39970f3374105d5eafbaab2a4482be82f272d5a5e", size = 4296730, upload-time = "2025-10-16T22:17:00.744Z" }, ++] ++ ++[[package]] ++name = "watchfiles" ++version = "1.2.0" ++source = { registry = "https://pypi.org/simple" } ++dependencies = [ ++ { name = "anyio" }, ++] ++sdist = { url = "https://files.pythonhosted.org/packages/cd/41/5e1a4bb12aac5f1493fa1bdc11154eca3b258ca4eba65d39c473fe19d8e9/watchfiles-1.2.0.tar.gz", hash = "sha256:c995fba777f1ea992f090f9236e9284cf7a5d1a0130dd5a3d82c598cacd76838", size = 108252, upload-time = "2026-05-18T04:32:04.251Z" } ++wheels = [ ++ { url = "https://files.pythonhosted.org/packages/fc/3d/8024c801df84d1587740d0359e7fdd80afeae3d159011f3d5376dd82f18e/watchfiles-1.2.0-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:704fd259e332e01f9b9c178f4bce9e49027e5587cc2600eeeaf8e76e1c846201", size = 400242, upload-time = "2026-05-18T04:31:19.014Z" }, ++ { url = "https://files.pythonhosted.org/packages/87/5b/f4dfd45323e949984a3a7f9dc31d1cbb049921e7d98253488dda72ccdaa9/watchfiles-1.2.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:6543cf55d170003296d185c0af981f3e1311564907e1f4e08671fc7693a890a5", size = 394562, upload-time = "2026-05-18T04:30:08.46Z" }, ++ { url = "https://files.pythonhosted.org/packages/98/d8/19483ef075d601c409bce8bcbb5c0f81a10876fff870400568f08ce484a1/watchfiles-1.2.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:89d8c2394a065ca86f5d2910ff263ae67c127e1376ccc4f9fc35c71db879f80a", size = 456611, upload-time = "2026-05-18T04:30:45.723Z" }, ++ { url = "https://files.pythonhosted.org/packages/b1/6a/cc81fbe7ee42f2f22e661a6e12def7807e01b14b2f39e0ff83fd373fd307/watchfiles-1.2.0-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:772b80df316480d894a0e3165fdd19cf77f5d17f9a787f94029465ad0e3529d1", size = 461379, upload-time = "2026-05-18T04:31:29.292Z" }, ++ { url = "https://files.pythonhosted.org/packages/b1/57/7e669002082c0a0f4fb5113bb70125f7110124b846b0a11bc5ae8e90eac1/watchfiles-1.2.0-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d158cd89df6053823533e06fb1d73c549133bff5f0396170c0e53d9559340717", size = 493556, upload-time = "2026-05-18T04:30:05.44Z" }, ++ { url = "https://files.pythonhosted.org/packages/45/7d/f60a2b19807b21fe8281f3a8da4f59eef0d5f96825ac4680ba2d4f2ebf91/watchfiles-1.2.0-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d516b3283a758e087841aedb8031549fb41ced08f3db10aa6d2bf32dc042525b", size = 575255, upload-time = "2026-05-18T04:30:40.568Z" }, ++ { url = "https://files.pythonhosted.org/packages/bd/49/77f5b5e6efbcd57482f74948ebb1b97e5c0046d6b61475042d830c84b3ff/watchfiles-1.2.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:53b2290c92e0506d102cd448fbc610d87079553f86caa39d67440856a8b8bba5", size = 467052, upload-time = "2026-05-18T04:31:17.942Z" }, ++ { url = "https://files.pythonhosted.org/packages/ee/5a/73e2959af1b97fd5d556f9a8bdba017be23ceeef731869d5eaa0a753d5a3/watchfiles-1.2.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a711b51aec4370d0dcda5b6c09463206f133a5759341d7744b953a7b62e1100e", size = 456858, upload-time = "2026-05-18T04:30:30.182Z" }, ++ { url = "https://files.pythonhosted.org/packages/50/57/1bc8c27fad7e6c19bddee15d276dbb6ab72480ec01c127afff1673aee417/watchfiles-1.2.0-cp311-cp311-manylinux_2_31_riscv64.whl", hash = "sha256:e2ca07fa7d89195ec0865d3d285666286740bfa83d83e5cee204043a31ecc165", size = 467579, upload-time = "2026-05-18T04:32:15.897Z" }, ++ { url = "https://files.pythonhosted.org/packages/09/6c/3c2e44edba3553c5e3c3b8c8a2a6dee6b9e12ae2cf4bd2378bebf9dc3038/watchfiles-1.2.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:e0618518f282c4ebff60f5e5b1247b6d91bb8b9f4476947563a1e74acc66f3c6", size = 633253, upload-time = "2026-05-18T04:31:37.123Z" }, ++ { url = "https://files.pythonhosted.org/packages/30/c2/d8c84a882ab39bbefcc4915ab3e91830b7a7e990c5570b0b69075aba3faf/watchfiles-1.2.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:0d191c054d0715c3c95c99df9b8dbf6fd096d8c1e021e8f212e1bd8bc444ccb5", size = 660713, upload-time = "2026-05-18T04:31:24.62Z" }, ++ { url = "https://files.pythonhosted.org/packages/a9/07/f97736a5fc605364fe67b25e9fa4a6965dfd4840d50c406ada507e9d735f/watchfiles-1.2.0-cp311-cp311-win32.whl", hash = "sha256:9342472aff9b093c5acd4f6d8f70ae0937964ab56542502bcf5579782da69ae8", size = 277222, upload-time = "2026-05-18T04:31:21.131Z" }, ++ { url = "https://files.pythonhosted.org/packages/cf/99/2b04981977fc2608afd60360d928c6aecf6b950292ca221d98f4005f6694/watchfiles-1.2.0-cp311-cp311-win_amd64.whl", hash = "sha256:dbd6c97045dad81227c8d040173da044c1de08de64a5ea8b555da4aee1d5fa22", size = 290274, upload-time = "2026-05-18T04:31:45.966Z" }, ++ { url = "https://files.pythonhosted.org/packages/3c/74/f7f58a7075ee9cf612b0cfcddb78b8cd8234f0742d6f0075cf0da2dde1c6/watchfiles-1.2.0-cp311-cp311-win_arm64.whl", hash = "sha256:57a2d9fa4fb4c2ecae57b13dfff2c7ab53e21a2ba674fe9f05506680fcdcc0d7", size = 283460, upload-time = "2026-05-18T04:31:39.126Z" }, ++ { url = "https://files.pythonhosted.org/packages/b8/2f/e42c992d2afda3108ea1c02acecc991b9f31d05c14adc2a7cee9ee211fc4/watchfiles-1.2.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:bc13eb17538be00c874699dc0abe4ee2bc8d50bb1166a6b9e175ef3fd7eb8f26", size = 400115, upload-time = "2026-05-18T04:32:02.06Z" }, ++ { url = "https://files.pythonhosted.org/packages/5f/8f/6af2ea19065c91d8b0ea3516fdfc8c0d349f407e8e9fbf4e5a17360de8ad/watchfiles-1.2.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:2d95ddc1eb6914154253d239089900813f6a767e174b8e6a50e7fdacb7e4236c", size = 393659, upload-time = "2026-05-18T04:30:50.951Z" }, ++ { url = "https://files.pythonhosted.org/packages/13/01/b32a967c56fb3e3e5be3db52c3d3b87fa4513aa367d8ed1ad96d42952e5f/watchfiles-1.2.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8f70d8b291ef6e88d19b1f297a6905ddb978888d9272b0d05e6f53309856bcfc", size = 453207, upload-time = "2026-05-18T04:31:04.231Z" }, ++ { url = "https://files.pythonhosted.org/packages/04/98/97557a812180338cb1abd32e1cffcc4588f59b5f23e0cb006b2ba95ba64a/watchfiles-1.2.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:56d8641cf834c2836922899105bd3ce3d0dfc69291d52edf0b4d0436829b34c0", size = 459273, upload-time = "2026-05-18T04:31:50.377Z" }, ++ { url = "https://files.pythonhosted.org/packages/e8/a8/b4b08dcb7653b8087c6586f7ce649505900e866bbcfe40dc9587af02e686/watchfiles-1.2.0-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2581a94056e55d7d0a31a823ea92bf73749c489ca2285bfdc0fbe6b2bb49d50c", size = 489927, upload-time = "2026-05-18T04:31:42.485Z" }, ++ { url = "https://files.pythonhosted.org/packages/50/94/3dceea03545d2e5ddfd839f0ddd5e1cecbf1697b5a428d5ba11cef6af95d/watchfiles-1.2.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:41bc1199f7523b3f82843c88cbb979180c949caef0342cf90968f178e5d49b01", size = 570476, upload-time = "2026-05-18T04:31:03.071Z" }, ++ { url = "https://files.pythonhosted.org/packages/cc/f2/d39a5450c3532092b91f81d274360e613c2371bc874a89c7a1a3c5e8d138/watchfiles-1.2.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:7571e4464cb6e434958f867f7f730b8ab0b75e3f8e5eac0499168486ab3c33a8", size = 465650, upload-time = "2026-05-18T04:30:12.701Z" }, ++ { url = "https://files.pythonhosted.org/packages/22/24/ed72f68cbc1333ca9b9f2200aa048bb6658ae41709bc1caad4310f4bdffd/watchfiles-1.2.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e53a384f76b631c3ae5334ce6a52f0baa3a911eb94a4eac7f160079868b716d5", size = 456398, upload-time = "2026-05-18T04:30:13.784Z" }, ++ { url = "https://files.pythonhosted.org/packages/0d/64/982ef4a4e5bab5b6e5b6becc8cd5e732f6130a78b855f0abec6439a9a135/watchfiles-1.2.0-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:d20029a60a71a052a24c4db7673bc4de39ab89adbaccbfb5d67987c5d73f424d", size = 465140, upload-time = "2026-05-18T04:31:52.111Z" }, ++ { url = "https://files.pythonhosted.org/packages/a0/0c/95282abf4ed680b6096010bcfc30c5fa7a041fc5aa5a2ad17a2cc6c75bba/watchfiles-1.2.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:2cb93af48550faf1cea04c303107c8b75833de7013e57ce27d3b8d21d8d0f58c", size = 630259, upload-time = "2026-05-18T04:31:25.676Z" }, ++ { url = "https://files.pythonhosted.org/packages/30/45/607c1de1530c4bdcf2cf1d1ecc2505ddba5d96bd43ba9f2b0e79876f850f/watchfiles-1.2.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:2995c176de7692b86a2e4c58d9ec718f753150a979cb4a754e2b4ffa38e70906", size = 659859, upload-time = "2026-05-18T04:30:24.333Z" }, ++ { url = "https://files.pythonhosted.org/packages/fa/08/d9e2e0f9e8e6791d33aefc694ad7eefa7f901f63caff84a81ded38692f9c/watchfiles-1.2.0-cp312-cp312-win32.whl", hash = "sha256:7a2cffd17d27d2ecbb310c2b1d8174f222a5495b1a721894afa88ec11e25b898", size = 275480, upload-time = "2026-05-18T04:30:31.307Z" }, ++ { url = "https://files.pythonhosted.org/packages/1c/e6/9d42569c0102645cc8cea5d8c7d8a1e9d4ada2cb7f05f75e554b8aa2202a/watchfiles-1.2.0-cp312-cp312-win_amd64.whl", hash = "sha256:f155b3a1b2a5fc89cdc70d47ee5d54e3b75e88efa34982028a35daef9ba00379", size = 288718, upload-time = "2026-05-18T04:32:10.745Z" }, ++ { url = "https://files.pythonhosted.org/packages/0a/26/88e0dc6ee3898169d7fa22bb6a69cabf2502d2ee25cb8c876d1262d204f8/watchfiles-1.2.0-cp312-cp312-win_arm64.whl", hash = "sha256:8fa585ede612ee9f9e91b18bebf9ba11b9ae29a4e3a0d0cf6fca3e382133f0d5", size = 281026, upload-time = "2026-05-18T04:30:22.23Z" }, ++ { url = "https://files.pythonhosted.org/packages/d1/4d/70a7feced9f87e2ff26dba42667290f41694fc64646c67261fbb8cab5d5c/watchfiles-1.2.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:01ea8d66f0693b9b60a6541c8d10263091ca9a9060d242f3c1f3143f9aad2c98", size = 399730, upload-time = "2026-05-18T04:31:38.162Z" }, ++ { url = "https://files.pythonhosted.org/packages/31/3a/0da302f2307aee316922806ebd5726c542cbd787c938271cf14a074c7daf/watchfiles-1.2.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:7ba0480b9a74af058f43b337e937a451e109295c420916d68ad24e3dc02f5e44", size = 392842, upload-time = "2026-05-18T04:30:27.051Z" }, ++ { url = "https://files.pythonhosted.org/packages/db/ef/d5bdb705c224dbc256aa0c1ec47bf4e61ec52558f2afb44a71a1fe4d7015/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4f34e26a19f91f710c08e0183429f0d1d15df734e6bc78c31e77b9ea9c433658", size = 452989, upload-time = "2026-05-18T04:31:11.945Z" }, ++ { url = "https://files.pythonhosted.org/packages/71/29/5495f2c1661949ef7a35e4d71111d129cfe7606414a26887a919d0a55406/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:b4e77f6a55f858504069abd35d336a637555c09bca453dde1ee1e5ada8a6a1fb", size = 458978, upload-time = "2026-05-18T04:30:52.606Z" }, ++ { url = "https://files.pythonhosted.org/packages/d5/8c/7f9c07c433811c2fffd93e13fdfb7135de9aab5f2ae41be08960fa0047dc/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0cb4d80e212f116474a545c21c912b445f16bb0cef9e6a73a498164223e14e2f", size = 490248, upload-time = "2026-05-18T04:31:36.003Z" }, ++ { url = "https://files.pythonhosted.org/packages/3c/11/d93632febc52fbc21be90231bb7c17fd5387f46c9076fd40a5f9c2ae6910/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b974946a10af379d425e2eef5b62f5c6ebeaccf91d45eaad6f5b27ecd4f91aa0", size = 571847, upload-time = "2026-05-18T04:31:10.862Z" }, ++ { url = "https://files.pythonhosted.org/packages/55/b4/383173e73aabb07ad1d9c7aa859d95437ac46a6d6a1e11005facda0c9d19/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:86bc13c25a8d1fcd70b51d0ce7c9b65e90de5666fcbfd3e34957cc73ee19aeb5", size = 465974, upload-time = "2026-05-18T04:30:17.006Z" }, ++ { url = "https://files.pythonhosted.org/packages/a7/6c/89b1a230a78f57c52dd8893adb1f92f94411721b6ec12596c56d98c74356/watchfiles-1.2.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ca148d73dea36c9763aaa351e4d7a51780ec1584217c45276f4fe8239c768b71", size = 454782, upload-time = "2026-05-18T04:30:35.656Z" }, ++ { url = "https://files.pythonhosted.org/packages/24/62/1732118367cfff0a9fce3bf62ff4bfded09ef5df21d9d446b858b3f70a96/watchfiles-1.2.0-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:c525543d91961c6955b2636b308569e84a1d1c5f5f2932041ab9ef46422f43e3", size = 465182, upload-time = "2026-05-18T04:30:20.846Z" }, ++ { url = "https://files.pythonhosted.org/packages/28/96/716f7e5f51339bf22963f3345f9f27d7f3b30e2eadc597e257c881dd3c53/watchfiles-1.2.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:a204794696ffb8f9b10fba6f7cb5216d42f3b2b71860ccac6b6e42f5f10973b0", size = 629841, upload-time = "2026-05-18T04:31:05.397Z" }, ++ { url = "https://files.pythonhosted.org/packages/4c/fe/c40783950fd771ccf66ab3ec2722d188a9af1c7f96c6e811f36e40c6e03f/watchfiles-1.2.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:10d86db20695afe7997ac9e1717637d6714a8d0220458c33f3d2061f54cec427", size = 658028, upload-time = "2026-05-18T04:31:48.22Z" }, ++ { url = "https://files.pythonhosted.org/packages/71/72/4508db1856d1d87fcbb3b63f4839bab1b5682cb0e8d224d122263c09654a/watchfiles-1.2.0-cp313-cp313-win32.whl", hash = "sha256:eb283ee99e21ad6443c8cdb06ac5b34b1308c329cbdf03fa02b445363714c799", size = 275183, upload-time = "2026-05-18T04:30:59.57Z" }, ++ { url = "https://files.pythonhosted.org/packages/f9/36/14b76ca57652e5cc5fd1c11f32a261292c08a0d19a00351013c2549cbfb2/watchfiles-1.2.0-cp313-cp313-win_amd64.whl", hash = "sha256:a0f27f01bee51861392bb6b7c4fdb290b27d1eb194e9e28788d68102a0e898d9", size = 288059, upload-time = "2026-05-18T04:32:07.937Z" }, ++ { url = "https://files.pythonhosted.org/packages/1b/8d/0a85e395398d8d20fadfe5c5d32c726eee17a519e78fb356f2cf7531bffe/watchfiles-1.2.0-cp313-cp313-win_arm64.whl", hash = "sha256:3651aa7058595e9cfb75d35dd5ada2bf9f48a5b8a0f3562821d3e210c507e077", size = 280186, upload-time = "2026-05-18T04:31:54.484Z" }, ++ { url = "https://files.pythonhosted.org/packages/37/68/36db056f1fdcc5f07302f56e631774d6835bcd6fa3ace402304621d5f9e5/watchfiles-1.2.0-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:faea288b6f0ab1902ef08f4ca6de005dccf856c4e0c4f21b8c5fce02d90a1b08", size = 399031, upload-time = "2026-05-18T04:30:44.576Z" }, ++ { url = "https://files.pythonhosted.org/packages/c1/64/01a9d6f66a82a5c101ce939274106cc72759d62427e153f01edd2b9f87c2/watchfiles-1.2.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:01859b11fd9fbca670f4d5da00fbac282cfea9bd67a2125d8b2833a3b5617ea9", size = 391205, upload-time = "2026-05-18T04:30:25.413Z" }, ++ { url = "https://files.pythonhosted.org/packages/84/2c/0a44fe058cb4bb7b8ede6b6670698bbb7c0400740e378d00022189b7b31d/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:fff610d7bb2256a317bb1e96f0d7862c7aa8076733ee5df0fd41bbe76a24a4f4", size = 451892, upload-time = "2026-05-18T04:32:14.005Z" }, ++ { url = "https://files.pythonhosted.org/packages/67/a1/351e0d56cd35e6488b5c8b4fb11a809a5bc923e8fe8fed9faf8920be0c89/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:b141a4891c995a039cd89e9a49e62df1dc8a559a5d1a6e4c7106d16c12777a55", size = 458867, upload-time = "2026-05-18T04:31:22.279Z" }, ++ { url = "https://files.pythonhosted.org/packages/d5/7d/9d09605187f1b838998624049fcf8bf47b73c1a3b76901fcac1782f62277/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f22943b7770483f6ea0721c6b11d022947a98eb0acae14694de034f4d0d38925", size = 490217, upload-time = "2026-05-18T04:31:43.657Z" }, ++ { url = "https://files.pythonhosted.org/packages/60/5d/a17a16eccb182f04188cd308ec24b1a71a9b5c4e7098269cf35d9fa56d02/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1bc6195825b7dcd217968bb1f801a60fd4c16e8eeab5bedc7fe917d7d5995ab4", size = 571458, upload-time = "2026-05-18T04:32:11.875Z" }, ++ { url = "https://files.pythonhosted.org/packages/d3/3d/4dd457062083ab1938e5dfd45032eb425cee2ac817287ca8ff4356183e5d/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d4a4b147f5dca2a5d325a06a832fb43f345751adfbc63204aec30e0d9ca965a2", size = 464707, upload-time = "2026-05-18T04:30:43.492Z" }, ++ { url = "https://files.pythonhosted.org/packages/c6/71/ea8c57b128f5383de74d0c7d2d9c57ad7c9a65a930c451bd25d524b295b7/watchfiles-1.2.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4543579a9bdb0c9560039b4ffddbdb39545707659fbc430ce4c10f3f68d557f9", size = 454663, upload-time = "2026-05-18T04:30:16.061Z" }, ++ { url = "https://files.pythonhosted.org/packages/53/fd/2e812bf938406d7db351f0703ddd3fc6c061cf30d96153a77bc79a943a44/watchfiles-1.2.0-cp313-cp313t-manylinux_2_31_riscv64.whl", hash = "sha256:20aa0e708b920bde876a4aa82dc7dd6ebea228a63a67cda6632c2fc87b787efa", size = 463537, upload-time = "2026-05-18T04:31:44.9Z" }, ++ { url = "https://files.pythonhosted.org/packages/86/56/d17a7f1dd1bc3035f1072694a551301272f1739c2d8e319c927cb9e29b38/watchfiles-1.2.0-cp313-cp313t-musllinux_1_1_aarch64.whl", hash = "sha256:d413349d565dab74297f2a63e84a097936be69bf8f3b3801f27f380e32040f44", size = 629194, upload-time = "2026-05-18T04:31:14.141Z" }, ++ { url = "https://files.pythonhosted.org/packages/be/06/f1ff66bf5cae50aa4062779a0ecd0bbaf15e466195719074078947d9a17d/watchfiles-1.2.0-cp313-cp313t-musllinux_1_1_x86_64.whl", hash = "sha256:f28b2725eb8cce327b9b3ab02415c853011dc55c95832fe90de6bc56f5315f72", size = 656194, upload-time = "2026-05-18T04:31:47.14Z" }, ++ { url = "https://files.pythonhosted.org/packages/e7/54/a9c7ea9a82a4ac65e7004c0a03920b5cdd2f9c3b678757d9cd425aa51d53/watchfiles-1.2.0-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:b8c8358484d5fa12ef34f05b7f4168eaf1932f408725ff6d023c33ec17bd79d4", size = 400205, upload-time = "2026-05-18T04:32:05.153Z" }, ++ { url = "https://files.pythonhosted.org/packages/aa/5d/c9ab3534374a4a67450696905d6ef16a04405448b8dc52bd752ae50423d4/watchfiles-1.2.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:9f04b092229ad2c50126dd3c922c8822e51e605993764a33058d4a791ab42281", size = 392508, upload-time = "2026-05-18T04:30:54.849Z" }, ++ { url = "https://files.pythonhosted.org/packages/26/ca/1ad30103535cf0cecd7b993e8d50edc5351b1820e38f2d22e3df58962feb/watchfiles-1.2.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7a7ce236284f002a156f70add88efe5c70879cccbb658be0822c54b1306fc09d", size = 452448, upload-time = "2026-05-18T04:30:53.727Z" }, ++ { url = "https://files.pythonhosted.org/packages/37/a1/ceee2cdf2afbd715fa07758d39c9859513eae411b23196f7fd039e5feedd/watchfiles-1.2.0-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:b9909cc2b48468b575eefa944919e1fe8a36c5849d5c7c168f80a8c1db69398e", size = 459605, upload-time = "2026-05-18T04:30:23.312Z" }, ++ { url = "https://files.pythonhosted.org/packages/e8/f6/421e30fd1cb3907a84ed92ab3f1983e37ba2dca015e9a894a048418417a2/watchfiles-1.2.0-cp314-cp314-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0a37faaed405c67e28e6be45a1fa4f206ef5a2860f27c237db9fa30704c38242", size = 490757, upload-time = "2026-05-18T04:30:47.358Z" }, ++ { url = "https://files.pythonhosted.org/packages/41/b0/55ed1b97ed08be7bba6f9a541cac15f2a858e1d74d2b07b6da70a82aab00/watchfiles-1.2.0-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9649193aa27bd9ff2e80ff29bfaa93085496c7a3a377592823cc58b77ee88add", size = 568672, upload-time = "2026-05-18T04:30:38.915Z" }, ++ { url = "https://files.pythonhosted.org/packages/d1/cf/d8ae8a80dd7bafab395ea7681c10237311bbf34d37704a8c744e7cf31fc7/watchfiles-1.2.0-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:4e4ff8e37f99cf1da89e255e07c9c4b37c214038c4283707bdec308cb1b0ea1f", size = 464197, upload-time = "2026-05-18T04:30:09.914Z" }, ++ { url = "https://files.pythonhosted.org/packages/7c/8a/3076c496ca8dafe0e8cd03fcebdfc47be4b1174b4e5b24ff6e396e6b3af2/watchfiles-1.2.0-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:054dc20fd2e3132b4c3883b4a00d72fd6e1f56fdaf89fccd12e8057d74cd74d7", size = 453181, upload-time = "2026-05-18T04:30:14.829Z" }, ++ { url = "https://files.pythonhosted.org/packages/e5/10/9745e17c98e7b8a86454df0a3c7b5686bd650383f1e9f26e4ebcbd6cc0c0/watchfiles-1.2.0-cp314-cp314-manylinux_2_31_riscv64.whl", hash = "sha256:e140ed30ebde76796b686e67c182cff10ea2fbab186fafd1560f74bb5a473a6e", size = 465109, upload-time = "2026-05-18T04:30:28.123Z" }, ++ { url = "https://files.pythonhosted.org/packages/8f/95/8ef4a95481d3e0cb52d62a06fa6e972e81424be2d9698b91a2fecca9904c/watchfiles-1.2.0-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:bb7e52ecf68ba46d22df23467b87cffeb2146908aa523ebfe803019618cfda06", size = 630653, upload-time = "2026-05-18T04:31:49.304Z" }, ++ { url = "https://files.pythonhosted.org/packages/fd/e4/3b3bf36b0f829b50c6ebcb8d031583863c59f923d6a6af3d485e470d0fac/watchfiles-1.2.0-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:23282a321c8baf9b3a3c4afff673f9fe65eb7fdc2338d765ccad9d3d1916a5ba", size = 657838, upload-time = "2026-05-18T04:31:06.497Z" }, ++ { url = "https://files.pythonhosted.org/packages/21/b1/6cbbb50c1f3002ab568777d44aa21206dfb8807a840990c4037523b51812/watchfiles-1.2.0-cp314-cp314-win32.whl", hash = "sha256:c0db965c5f79aa49fe672d297cf1febc5ad149b658594944f49a54a2b96270a7", size = 275108, upload-time = "2026-05-18T04:30:06.891Z" }, ++ { url = "https://files.pythonhosted.org/packages/92/45/190ce6db8dcb4536682cf75d3889ff1a27182a58cb519d343cb6d9ea63d8/watchfiles-1.2.0-cp314-cp314-win_amd64.whl", hash = "sha256:71283b39fd17e5408eb123bd37aeecfd9d54c81fc184421943208aadb879d103", size = 288441, upload-time = "2026-05-18T04:32:12.901Z" }, ++ { url = "https://files.pythonhosted.org/packages/74/0d/3eae1c2313ab08378431d907c3f8095ecca00f3eda33111cf4f0f2591799/watchfiles-1.2.0-cp314-cp314-win_arm64.whl", hash = "sha256:c5c19526f4e54a00f2666a6c0e9e40d582c09e865055ea7378bf0009aab857b3", size = 280684, upload-time = "2026-05-18T04:31:26.902Z" }, ++ { url = "https://files.pythonhosted.org/packages/b1/75/fb64e6c25d6b5ca636d03df34ffb1c6e9873303e76d27967e045f8df088f/watchfiles-1.2.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:d73a585accffa5ae39c17264c36ec3166d2fad7000c780f5ef83b2722afb9dd2", size = 398857, upload-time = "2026-05-18T04:32:17.108Z" }, ++ { url = "https://files.pythonhosted.org/packages/73/4e/9f7adf01754cbf81843722ccfec169d8f26c69778281a302855cecd2ee08/watchfiles-1.2.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:ae99b14c5f21e026e0e9d96f40e07d8570ebee6cafd9d8fc318354606daa7a28", size = 392413, upload-time = "2026-05-18T04:31:07.911Z" }, ++ { url = "https://files.pythonhosted.org/packages/47/c8/bec626bcc2d69f44b9acb24ce7d60ed7b16b73628eea747fcbd169d8edda/watchfiles-1.2.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4429f3b105524a10b72c3a819b091c495d2811d419c1e1e8df773a5a5974f831", size = 452409, upload-time = "2026-05-18T04:31:20.142Z" }, ++ { url = "https://files.pythonhosted.org/packages/00/b7/b6362068e81e7c556d155a34c35d40ac3ef42d747b06d7f6e5bf58e359c2/watchfiles-1.2.0-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:43d818978d06062d9b22c4fab2ebe44cf5213d42dc8e62bda8c2760cfa2eeb33", size = 458827, upload-time = "2026-05-18T04:32:06.219Z" }, ++ { url = "https://files.pythonhosted.org/packages/67/f8/9a813fa42afb1e0b4625e75f0479826644d3ee8dc287e093799bc01f390c/watchfiles-1.2.0-cp314-cp314t-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b9f732dc58b2dbe69e464ccf8fff7a03b0dd0be439da4c0720d3558527d3d6b4", size = 490104, upload-time = "2026-05-18T04:31:56.034Z" }, ++ { url = "https://files.pythonhosted.org/packages/2f/bf/27dfb6094ca4c9aad21298b5525b6c53cb36121ee454331d05161e58d130/watchfiles-1.2.0-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8f200104103feb097de4cab8fe4f5dd18a2026934c7dea98c55a2f5fd6d5a33b", size = 571360, upload-time = "2026-05-18T04:31:57.133Z" }, ++ { url = "https://files.pythonhosted.org/packages/fb/39/44a096d67270ea93df91d33877dbe91fbda3aa4f8ec2edf799d93eda8736/watchfiles-1.2.0-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:63ac26eefbf4af1741247d6fb68b11c49a25b2f7413fbd318a83a12aaa9cf666", size = 464644, upload-time = "2026-05-18T04:30:57.33Z" }, ++ { url = "https://files.pythonhosted.org/packages/0e/80/c7472203bad6268e3ef1ad260739704847898938ad7ea8b63a5131f46b50/watchfiles-1.2.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0c4997d4e4a55f0d02b6cde327322daf3a0400e5df6c6b15948994bf72497925", size = 454771, upload-time = "2026-05-18T04:30:48.736Z" }, ++ { url = "https://files.pythonhosted.org/packages/51/cf/3b10b268b4b7f0fc26e9debb5eef1998b515887840f444cd3ec80c688755/watchfiles-1.2.0-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:4c887eba18b7945ac73067a8b4a66f21cd46c2539b2bc68588f7be6c7eb6d26b", size = 463494, upload-time = "2026-05-18T04:31:33.826Z" }, ++ { url = "https://files.pythonhosted.org/packages/3d/3e/a4302545cd589262a0dc7d140e86f7688eba3f9c72776c27f7e23b8864c4/watchfiles-1.2.0-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:3416ff151bb6b5a8d8d11664974fbef4d9305b9b2957839ab5a270468fd8df30", size = 629383, upload-time = "2026-05-18T04:31:15.596Z" }, ++ { url = "https://files.pythonhosted.org/packages/db/99/d5649df0a9a410d45b7c882304d0b790903ac9b6e8f2cfd12114e0c6b9f2/watchfiles-1.2.0-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:0e831a271c035d89789cffc386b6aa1375f39f1cd25eb7ca0997e4970d152fc5", size = 656093, upload-time = "2026-05-18T04:31:58.707Z" }, ++ { url = "https://files.pythonhosted.org/packages/23/f4/7513ef1e85fc4c6331b59479d6d72661fc391fbe543678052ac72c8b6c19/watchfiles-1.2.0-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:4674d49eb94706dfe666c069fc0a1b646ffcf920473492e209f6d5f60d3f0cc2", size = 403050, upload-time = "2026-05-18T04:30:36.753Z" }, ++ { url = "https://files.pythonhosted.org/packages/27/0b/a54103cfd732bb703c7a749222011a0483ef3705948dae3b203158601119/watchfiles-1.2.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:094b9b70103d4e963499bdea001ee3c2697b144cd9ae6218a62c0f89ec9e31db", size = 396629, upload-time = "2026-05-18T04:32:03.268Z" }, ++ { url = "https://files.pythonhosted.org/packages/5e/2c/73f31a3b893886206c3f54d73e8ad8dee58cdb2f69ad2622e0a8a9e07f4e/watchfiles-1.2.0-pp311-pypy311_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b0ef001f8c25ad0fa9529f914c1600647ecd0f542d11c19b7894768c67b6acb7", size = 457318, upload-time = "2026-05-18T04:31:01.932Z" }, ++ { url = "https://files.pythonhosted.org/packages/e9/f9/45d021e4a5cc7b9dd567f7cbb06d3b75f751a690063fb6cc7ec60f4e46b7/watchfiles-1.2.0-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a88fc94e647bc4eec523f1caa540258eb71d14278b9daf72fa1e2658a98df0f0", size = 457771, upload-time = "2026-05-18T04:30:56.331Z" }, ++] ++ + [[package]] + name = "wcwidth" + version = "0.7.0" +``` From a47fe93eb39f4977491f0d3621ecb9195412dbfc Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Mon, 1 Jun 2026 13:56:44 -0500 Subject: [PATCH 09/29] feat(api): add puppy-desk backend runtime support --- code_puppy/config.py | 44 +- .../plugins/frontend_emitter/emitter.py | 13 +- code_puppy/session_storage.py | 466 +++++++----------- code_puppy/tools/command_runner.py | 285 +++++++++-- pyproject.toml | 6 + 5 files changed, 482 insertions(+), 332 deletions(-) diff --git a/code_puppy/config.py b/code_puppy/config.py index 62afbc9ea..7ddaf53a3 100644 --- a/code_puppy/config.py +++ b/code_puppy/config.py @@ -52,18 +52,38 @@ def _get_xdg_dir(env_var: str, fallback: str) -> str: AGENTS_DIR = os.path.join(DATA_DIR, "agents") 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") +_DEFAULT_SQLITE_FILE = os.path.join(DATA_DIR, "dbos_store.sqlite") # 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") +DBOS_DATABASE_URL = os.environ.get( + "DBOS_SYSTEM_DATABASE_URL", f"sqlite:///{_DEFAULT_SQLITE_FILE}" +) +# DBOS enable switch is controlled solely via puppy.cfg using key 'enable_dbos'. +# Default: True (DBOS enabled) unless explicitly disabled. + +def get_use_dbos() -> bool: + """Return True if DBOS should be used based on 'enable_dbos' (default True).""" + cfg_val = get_value("enable_dbos") + if cfg_val is None: + return True + return str(cfg_val).strip().lower() in {"1", "true", "yes", "on"} + +# 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: @@ -199,6 +219,18 @@ def get_suppress_directory_listing() -> bool: return str(val).lower() in ("1", "true", "yes", "on") +def get_command_timeout_seconds() -> int: + """Return shell command timeout in seconds (default: 30).""" + val = get_value("command_timeout_seconds") + if val is None: + return 30 + try: + parsed = int(str(val).strip()) + except (ValueError, TypeError): + return 30 + return max(1, parsed) + + DEFAULT_SECTION = "puppy" REQUIRED_KEYS = ["puppy_name", "owner_name"] 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/session_storage.py b/code_puppy/session_storage.py index 62c6bfad1..ab9a849bb 100644 --- a/code_puppy/session_storage.py +++ b/code_puppy/session_storage.py @@ -1,340 +1,238 @@ """Shared helpers for persisting and restoring chat sessions. -This module centralises the pickle + metadata handling that used to live in -both the CLI command handler and the auto-save feature. Keeping it here helps -us avoid duplication while staying inside the Zen-of-Python sweet spot: simple -is better than complex, nested side effects are worse than deliberate helpers. +MIGRATION NOTE: Pickle-based storage has been removed. All sessions are now +persisted exclusively to SQLite via code_puppy.api.db.queries. +This module retains only stateless helpers (title generation) and stubs for +removed functionality to maintain import compatibility during the transition. """ from __future__ import annotations import json -import pickle +import re from dataclasses import dataclass from pathlib import Path -from typing import Any, Callable, List +from typing import Any, List - -def _safe_loads(data: bytes) -> Any: - """Deserialize pickle data.""" - return pickle.loads(data) # noqa: S301 - - -_LEGACY_SIGNED_HEADER = b"CPSESSION\x01" -_LEGACY_SIGNATURE_SIZE = ( - 32 # legacy signature bytes, retained only for backward-compat parsing -) +from pydantic import TypeAdapter +from pydantic_ai.messages import ModelMessage SessionHistory = List[Any] -TokenEstimator = Callable[[Any], int] - -@dataclass(slots=True) -class SessionPaths: - pickle_path: Path - metadata_path: Path - -@dataclass(slots=True) +@dataclass class SessionMetadata: - session_name: str - timestamp: str + """Metadata returned after saving a session.""" + message_count: int total_tokens: int - pickle_path: Path - metadata_path: Path - auto_saved: bool = False - - def as_serialisable(self) -> dict[str, Any]: - return { - "session_name": self.session_name, - "timestamp": self.timestamp, - "message_count": self.message_count, - "total_tokens": self.total_tokens, - "file_path": str(self.pickle_path), - "auto_saved": self.auto_saved, - } - - -def _extract_pickle_payload(raw: bytes) -> bytes: - """Return the pickle payload from raw session file bytes. - - New format is raw pickle bytes. - Legacy format was: header + 32-byte signature + pickle payload. - We no longer verify or generate signatures. - """ - if raw.startswith(_LEGACY_SIGNED_HEADER): - offset = len(_LEGACY_SIGNED_HEADER) + _LEGACY_SIGNATURE_SIZE - return raw[offset:] - return raw + pickle_path: Path # Kept for compatibility, points to .json file + metadata_path: Path # Points to .json file (same as above) -def ensure_directory(path: Path) -> Path: - path.mkdir(parents=True, exist_ok=True) - return path +def generate_heuristic_title(history: SessionHistory, max_length: int = 50) -> str: + """Generate a short title from the first user message in the history. + Extracts the first user message, takes the first ~50 chars, and converts + to a filename-safe format (lowercase, spaces to hyphens, remove special chars). -def build_session_paths(base_dir: Path, session_name: str) -> SessionPaths: - pickle_path = base_dir / f"{session_name}.pkl" - metadata_path = base_dir / f"{session_name}_meta.json" - return SessionPaths(pickle_path=pickle_path, metadata_path=metadata_path) + Handles multiple message formats: + 1. Pydantic-ai format: msg.kind == 'request' with msg.parts[].content + 2. Enhanced/wrapped format: {'msg': , 'agent': str, ...} + 3. Simple dict format: {'role': 'user', 'content': str} + """ + + def extract_user_content(msg: Any) -> str | None: + """Extract user message content from various message formats.""" + # Handle wrapped/enhanced format: {'msg': , 'agent': ...} + if isinstance(msg, dict) and "msg" in msg: + msg = msg["msg"] + + # Handle pydantic-ai format: msg.kind == 'request' + if hasattr(msg, "kind") and msg.kind == "request": + for part in getattr(msg, "parts", []): + if hasattr(part, "content") and isinstance(part.content, str): + content = part.content.strip() + if content: + return content + + # Handle simple dict format: {'role': 'user', 'content': str} + if isinstance(msg, dict): + if msg.get("role") == "user" and isinstance(msg.get("content"), str): + content = msg["content"].strip() + if content: + return content + + return None + + def content_to_title(content: str) -> str: + """Convert content to a filename-safe kebab-case title.""" + # Take first line or first max_length chars + first_line = content.split("\n")[0][:max_length] + # Convert to kebab-case filename-safe format + title = first_line.lower() + title = re.sub(r"[^a-z0-9\s-]", "", title) # Remove special chars + title = re.sub(r"\s+", "-", title) # Spaces to hyphens + title = re.sub(r"-+", "-", title) # Collapse multiple hyphens + title = title.strip("-")[:max_length] + return title + + # Find first user message + for msg in history: + content = extract_user_content(msg) + if content: + title = content_to_title(content) + return title if title else "untitled-session" + + return "untitled-session" + + +# --------------------------------------------------------------------------- +# Deprecated / Removed Pickle Functionality -> Replaced with JSON for Export +# --------------------------------------------------------------------------- def save_session( - *, history: SessionHistory, session_name: str, base_dir: Path, - timestamp: str, - token_estimator: TokenEstimator, - auto_saved: bool = False, + timestamp: str | None = None, + token_estimator: Any | None = None, + **kwargs: Any, ) -> SessionMetadata: - ensure_directory(base_dir) - paths = build_session_paths(base_dir, session_name) - - pickle_data = pickle.dumps(history) - tmp_pickle = paths.pickle_path.with_suffix(".tmp") - with tmp_pickle.open("wb") as pickle_file: - pickle_file.write(pickle_data) - tmp_pickle.replace(paths.pickle_path) - - total_tokens = sum(token_estimator(message) for message in history) - metadata = SessionMetadata( - session_name=session_name, - timestamp=timestamp, + """Save session history to a JSON file (replacing legacy pickle). + + This is used for 'pinning' or exporting sessions via /dump_context. + """ + base_dir.mkdir(parents=True, exist_ok=True) + file_path = base_dir / f"{session_name}.json" + + # Try to serialize the history, handling mixed types: + # - ModelMessage objects: serialize via Pydantic + # - System dicts: serialize as plain dicts + try: + from pydantic_ai.messages import ModelMessagesTypeAdapter + + history_data = [] + for item in history: + # Unwrap from {msg: ..., agent: ..., ts: ...} wrapper if present + if isinstance(item, dict) and "msg" in item: + inner = item["msg"] + wrapper_meta = {k: v for k, v in item.items() if k != "msg"} + + # Check if inner is a ModelMessage (has 'parts' attribute) + if hasattr(inner, "parts"): + # Serialize ModelMessage, then add wrapper metadata + try: + serialized = ModelMessagesTypeAdapter.dump_python( + [inner], mode="json" + )[0] + serialized["_wrapper"] = wrapper_meta + history_data.append(serialized) + except Exception: + # Fallback: keep as-is + history_data.append(item) + else: + # It's a system dict like {msg: 'system', ...} - keep as-is + history_data.append(item) + elif hasattr(item, "parts"): + # Direct ModelMessage without wrapper + try: + serialized = ModelMessagesTypeAdapter.dump_python( + [item], mode="json" + )[0] + history_data.append(serialized) + except Exception: + history_data.append({"_raw": str(item)}) + else: + # Unknown type - keep as-is + history_data.append(item) + + json_bytes = json.dumps(history_data, indent=2, default=str).encode("utf-8") + except Exception: + # Fallback: generic JSON dump + json_bytes = json.dumps(history, default=str, indent=2).encode("utf-8") + + file_path.write_bytes(json_bytes) + + # Calculate tokens if estimator provided + total_tokens = 0 + if token_estimator: + for msg in history: + try: + # estimator might expect the original object or dict + total_tokens += token_estimator(msg) + except Exception: + pass + + return SessionMetadata( message_count=len(history), total_tokens=total_tokens, - pickle_path=paths.pickle_path, - metadata_path=paths.metadata_path, - auto_saved=auto_saved, + pickle_path=file_path, + metadata_path=file_path, ) - tmp_metadata = paths.metadata_path.with_suffix(".tmp") - with tmp_metadata.open("w", encoding="utf-8") as metadata_file: - json.dump(metadata.as_serialisable(), metadata_file, indent=2) - tmp_metadata.replace(paths.metadata_path) - - return metadata - def load_session( session_name: str, base_dir: Path, *, allow_legacy: bool = False ) -> SessionHistory: - # Kept for API compatibility; legacy loading is always supported now. - _ = allow_legacy + """Load session history from a JSON file.""" + file_path = base_dir / f"{session_name}.json" + if not file_path.exists(): + # Fallback check for legacy .pkl if explicitly requested? + legacy_path = base_dir / f"{session_name}.pkl" + if allow_legacy and legacy_path.exists(): + raise NotImplementedError( + f"Found legacy pickle at {legacy_path} but pickle loading is removed." + ) + raise FileNotFoundError(f"Session file not found: {file_path}") - paths = build_session_paths(base_dir, session_name) - if not paths.pickle_path.exists(): - raise FileNotFoundError(paths.pickle_path) + json_data = file_path.read_bytes() - raw = paths.pickle_path.read_bytes() - pickle_data = _extract_pickle_payload(raw) - return _safe_loads(pickle_data) + # Try to load as ModelMessage objects + try: + adapter = TypeAdapter(List[ModelMessage]) + return adapter.validate_json(json_data) + except Exception: + # Fallback: return as list of dicts + return json.loads(json_data) def list_sessions(base_dir: Path) -> List[str]: + """List available JSON sessions.""" if not base_dir.exists(): return [] - return sorted(path.stem for path in base_dir.glob("*.pkl")) + return sorted([p.stem for p in base_dir.glob("*.json")]) def cleanup_sessions(base_dir: Path, max_sessions: int) -> List[str]: - if max_sessions <= 0: - return [] - - if not base_dir.exists(): - return [] - - candidate_paths = list(base_dir.glob("*.pkl")) - if len(candidate_paths) <= max_sessions: - return [] - - sorted_candidates = sorted( - ((path.stat().st_mtime, path) for path in candidate_paths), - key=lambda item: item[0], - ) - - stale_entries = sorted_candidates[:-max_sessions] - removed_sessions: List[str] = [] - for _, pickle_path in stale_entries: - metadata_path = base_dir / f"{pickle_path.stem}_meta.json" - try: - pickle_path.unlink(missing_ok=True) - metadata_path.unlink(missing_ok=True) - removed_sessions.append(pickle_path.stem) - except OSError: - continue - - return removed_sessions + """Cleanup old sessions if count exceeds max_sessions. - -async def restore_autosave_interactively(base_dir: Path) -> None: - """Prompt the user to load an autosave session from base_dir, if any exist. - - This helper is deliberately placed in session_storage to keep autosave - restoration close to the persistence layer. It uses the same public APIs - (list_sessions, load_session) and mirrors the interactive behaviours from - the command handler. + Returns list of removed session names. """ - sessions = list_sessions(base_dir) - if not sessions: - return - - # Import locally to avoid pulling the messaging layer into storage modules - from datetime import datetime + if not base_dir.exists() or max_sessions <= 0: + return [] - from prompt_toolkit.formatted_text import FormattedText + sessions = sorted(base_dir.glob("*.json"), key=lambda p: p.stat().st_mtime) - from code_puppy.agents.agent_manager import get_current_agent - from code_puppy.command_line.prompt_toolkit_completion import ( - get_input_with_combined_completion, - ) - from code_puppy.messaging import emit_success, emit_system_message, emit_warning - - entries = [] - for name in sessions: - meta_path = base_dir / f"{name}_meta.json" - try: - with meta_path.open("r", encoding="utf-8") as meta_file: - data = json.load(meta_file) - timestamp = data.get("timestamp") - message_count = data.get("message_count") - except Exception: - timestamp = None - message_count = None - entries.append((name, timestamp, message_count)) - - def sort_key(entry): - _, timestamp, _ = entry - if timestamp: + removed = [] + if len(sessions) > max_sessions: + to_remove = sessions[: len(sessions) - max_sessions] + for p in to_remove: try: - return datetime.fromisoformat(timestamp) - except ValueError: - return datetime.min - return datetime.min - - entries.sort(key=sort_key, reverse=True) - - PAGE_SIZE = 5 - total = len(entries) - page = 0 - - def render_page() -> None: - start = page * PAGE_SIZE - end = min(start + PAGE_SIZE, total) - page_entries = entries[start:end] - emit_system_message("Autosave Sessions Available:") - for idx, (name, timestamp, message_count) in enumerate(page_entries, start=1): - timestamp_display = timestamp or "unknown time" - message_display = ( - f"{message_count} messages" - if message_count is not None - else "unknown size" - ) - emit_system_message( - f" [{idx}] {name} ({message_display}, saved at {timestamp_display})" - ) - # If there are more pages, offer next-page; show 'Return to first page' on last page - if total > PAGE_SIZE: - page_count = (total + PAGE_SIZE - 1) // PAGE_SIZE - is_last_page = (page + 1) >= page_count - remaining = total - (page * PAGE_SIZE + len(page_entries)) - 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}") - emit_system_message(" [Enter] Skip loading autosave") - - chosen_name: str | None = None - - while True: - render_page() - try: - selection = await get_input_with_combined_completion( - FormattedText( - [ - ( - "class:prompt", - "Pick 1-5 to load, 6 for next, or name/Enter: ", - ) - ] - ) - ) - except (KeyboardInterrupt, EOFError): - emit_warning("Autosave selection cancelled") - return - - selection = (selection or "").strip() - if not selection: - return - - # Numeric choice: 1-5 select within current page; 6 advances page - if selection.isdigit(): - num = int(selection) - if num == 6 and total > PAGE_SIZE: - page = (page + 1) % ((total + PAGE_SIZE - 1) // PAGE_SIZE) - # loop and re-render next page - continue - if 1 <= num <= 5: - start = page * PAGE_SIZE - idx = start + (num - 1) - if 0 <= idx < total: - chosen_name = entries[idx][0] - break - else: - emit_warning("Invalid selection for this page") - continue - emit_warning("Invalid selection; choose 1-5 or 6 for next") - continue - - # Allow direct typing by exact session name - for name, _ts, _mc in entries: - if name == selection: - chosen_name = name - break - if chosen_name: - break - emit_warning("No autosave loaded (invalid selection)") - # keep looping and allow another try - - if not chosen_name: - return + p.unlink() + removed.append(p.stem) + except Exception: + pass - try: - history = load_session(chosen_name, base_dir) - except FileNotFoundError: - emit_warning(f"Autosave '{chosen_name}' could not be found") - return - except Exception as exc: - emit_warning(f"Failed to load autosave '{chosen_name}': {exc}") - return - - agent = get_current_agent() - agent.set_message_history(history) - - # Set current autosave session id so subsequent autosaves overwrite this session - try: - from code_puppy.config import pin_current_session_name + return removed - pin_current_session_name(chosen_name) - except Exception: - pass - - total_tokens = sum(agent.estimate_tokens_for_message(msg) for msg in history) - session_path = base_dir / f"{chosen_name}.pkl" - emit_success( - f"✅ Autosave loaded: {len(history)} messages ({total_tokens} tokens)\n" - f"📁 From: {session_path}" - ) +async def restore_autosave_interactively(base_dir: Path) -> None: + """Deprecated: No-op.""" + pass - # Display recent message history for context - try: - from code_puppy.command_line.autosave_menu import display_resumed_history - display_resumed_history(history) - except Exception: - pass # Don't fail if display doesn't work in non-TTY environment +def build_session_paths(base_dir: Path, session_name: str) -> Any: + """Deprecated: Returns None or raises.""" + raise NotImplementedError("Session paths are no longer used.") diff --git a/code_puppy/tools/command_runner.py b/code_puppy/tools/command_runner.py index a24495035..90b779208 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,104 @@ 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 +228,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 +315,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 +408,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 +435,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 +443,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 +481,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,15 +848,14 @@ 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() stdout_lines = [] @@ -1023,8 +1191,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 +1207,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 +1221,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 +1242,20 @@ 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: + try: + from code_puppy.plugins.walmart_specific.session_context import ( + get_session_working_directory, + ) + + cwd = get_session_working_directory() + except ImportError: + pass + # Generate unique group_id for this command execution group_id = generate_group_id("shell_command", command) @@ -1204,12 +1388,32 @@ async def run_shell_command( # Check if we're running as a sub-agent (skip confirmation and run silently) running_as_subagent = is_subagent() + confirmation_lock_acquired = False + + # 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() + ): + confirmation_lock_acquired = _CONFIRMATION_LOCK.acquire(blocking=False) + if not confirmation_lock_acquired: + return ShellCommandOutput( + success=False, + command=command, + error="Another command is currently awaiting confirmation", + ) # Get puppy name for personalized messages from code_puppy.config import get_puppy_name @@ -1227,8 +1431,7 @@ async def run_shell_command( panel_content.append("📂 Working directory: ", style="dim") panel_content.append(cwd, style="dim cyan") - # Use the common approval function (async version). - # Internal queueing means parallel calls wait their turn here. + # Use the common approval function (async version) confirmed, user_feedback = await get_user_approval_async( title="Shell Command", content=panel_content, @@ -1237,6 +1440,10 @@ async def run_shell_command( puppy_name=puppy_name, ) + # Release lock after approval + if confirmation_lock_acquired: + _CONFIRMATION_LOCK.release() + if not confirmed: if user_feedback: result = ShellCommandOutput( @@ -1455,9 +1662,13 @@ async def _run_command_inner( try: # Run the blocking shell command in a thread pool to avoid blocking the event loop # This allows multiple sub-agents to run shell commands in parallel + # Copy context so ContextVar-based session tracking propagates to the worker thread + ctx = contextvars.copy_context() return await loop.run_in_executor( _SHELL_EXECUTOR, - partial(_run_command_sync, command, cwd, timeout, group_id, silent), + partial( + ctx.run, _run_command_sync, command, cwd, timeout, group_id, silent + ), ) except Exception as e: if not silent: diff --git a/pyproject.toml b/pyproject.toml index c020c3847..ba5ae44a4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,11 @@ 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", + "dbos>=2.11.0", "pydantic-ai-slim[openai,anthropic,mcp]==1.56.0", "typer>=0.12.0", "mcp>=1.9.4", @@ -64,6 +69,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 From 2e03e1d46f0134e91e04daeaeebda70f9c427e19 Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Mon, 1 Jun 2026 13:56:44 -0500 Subject: [PATCH 10/29] feat(ws): adapt complete responses to streaming frames --- .../test_streaming_protocol_transform.py | 93 +++++++++++++++ code_puppy/api/ws/chat_handler.py | 108 +++++++++++++++--- code_puppy/api/ws/schemas.py | 3 + 3 files changed, 186 insertions(+), 18 deletions(-) create mode 100644 code_puppy/api/tests/test_streaming_protocol_transform.py 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..a70d3f6cd --- /dev/null +++ b/code_puppy/api/tests/test_streaming_protocol_transform.py @@ -0,0 +1,93 @@ +"""Tests for streaming-only assistant response protocol adaptation.""" + +from code_puppy.api.ws.chat_handler import build_assistant_text_stream_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" diff --git a/code_puppy/api/ws/chat_handler.py b/code_puppy/api/ws/chat_handler.py index fbcb4c918..c8af429a7 100644 --- a/code_puppy/api/ws/chat_handler.py +++ b/code_puppy/api/ws/chat_handler.py @@ -53,7 +53,6 @@ ServerConfigValue, ServerError, ServerMessage, - ServerResponse, ServerSessionMetaUpdated, ServerSessionRestored, ServerSessionSwitched, @@ -201,6 +200,81 @@ def build_error_response_frames( 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. + + The GUI protocol is streaming-only: assistant text must be rendered from + ``assistant_message_start``/``assistant_message_delta``/ + ``assistant_message_end`` and finalized by ``stream_end``. Some upstream + model clients return a complete response rather than true token deltas; this + helper adapts those complete responses to the same wire shape so the GUI has + one rendering path. + + This intentionally emits a single full-content delta for non-streaming + upstream responses. Artificial local chunking can be added later without + changing the protocol contract. + """ + 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, + ), + ] + + def register_chat_endpoint(app: FastAPI) -> None: """Register the /ws/chat WebSocket endpoint.""" @@ -229,8 +303,8 @@ async def websocket_chat( {"type": "tool_result", "tool_name": "...", "result": "...", "success": true, "agent_name": "...", "model_name": "..."} - # Final response - {"type": "response", "content": "...", "done": true, + # Final stream marker + {"type": "stream_end", "success": true, "total_length": 123, "agent_name": "...", "model_name": "...", "tokens": {...}} # Errors @@ -3268,21 +3342,19 @@ async def drain_events_with_signal(): f"Sent thinking content ({len(thinking_text)} chars) for non-streaming model" ) - # Send the main response - await send_typed( - ServerResponse( - content=response_text, - done=True, - 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, - ) - ) + # 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. diff --git a/code_puppy/api/ws/schemas.py b/code_puppy/api/ws/schemas.py index 05e56528d..f7a198e65 100644 --- a/code_puppy/api/ws/schemas.py +++ b/code_puppy/api/ws/schemas.py @@ -385,10 +385,13 @@ 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 From ffdc3e0cc6bc6eea9d2fb403fd47b65e06a3b8d9 Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Mon, 1 Jun 2026 13:56:44 -0500 Subject: [PATCH 11/29] refactor(api): remove terminal and sessions server routes --- code_puppy/api/app.py | 46 +- code_puppy/api/pty_manager.py | 453 ------------------ code_puppy/api/templates/terminal.html | 382 --------------- .../tests/test_gate2_legacy_ws_namespace.py | 42 ++ code_puppy/api/websocket.py | 8 +- code_puppy/api/ws/__init__.py | 6 - code_puppy/api/ws/terminal_handler.py | 141 ------ 7 files changed, 47 insertions(+), 1031 deletions(-) delete mode 100644 code_puppy/api/pty_manager.py delete mode 100644 code_puppy/api/templates/terminal.html create mode 100644 code_puppy/api/tests/test_gate2_legacy_ws_namespace.py delete mode 100644 code_puppy/api/ws/terminal_handler.py diff --git a/code_puppy/api/app.py b/code_puppy/api/app.py index a68b86c14..8624db6fc 100644 --- a/code_puppy/api/app.py +++ b/code_puppy/api/app.py @@ -110,17 +110,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: # Shutdown: clean up all the things! logger.info("🐶 Code Puppy API shutting down, cleaning up...") - # 1. Close all PTY sessions - try: - from code_puppy.api.pty_manager import get_pty_manager - - pty_manager = get_pty_manager() - await pty_manager.close_all() - logger.info("✓ All PTY sessions closed") - except Exception as e: - logger.error("Error closing PTY sessions: %s", e) - - # 2. Shutdown session cache thread pool executor + # 1. Shutdown session cache thread pool executor try: from code_puppy.api.session_cache import shutdown_executor @@ -129,7 +119,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: except Exception as e: logger.error("Error shutting down session cache executor: %s", e) - # 3. Shutdown ws_sessions thread pool executor + # 2. Shutdown ws_sessions thread pool executor try: from code_puppy.api.routers import ws_sessions @@ -215,7 +205,7 @@ def create_app() -> FastAPI: app.include_router(protocol.router, prefix="/api/protocol", tags=["protocol"]) - # WebSocket endpoints (events + terminal) + # WebSocket endpoints from code_puppy.api.websocket import setup_websocket setup_websocket(app) @@ -239,21 +229,15 @@ async def root():

🐶

Code Puppy

- - Open Terminal - Open Chat - - View Sessions - API Docs

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

@@ -261,28 +245,6 @@ async def root(): """ ) - @app.get("/terminal") - async def terminal_page(): - """Serve the interactive terminal page.""" - html_file = templates_dir / "terminal.html" - if html_file.exists(): - return FileResponse(html_file, media_type="text/html") - return HTMLResponse( - content="

Terminal template not found

", - status_code=404, - ) - - @app.get("/sessions") - async def sessions_page(): - """Serve the sessions monitoring page.""" - html_file = templates_dir / "sessions.html" - if html_file.exists(): - return FileResponse(html_file, media_type="text/html") - return HTMLResponse( - content="

Sessions template not found

", - status_code=404, - ) - @app.get("/chat") async def chat_page(): """Serve the chat interface page.""" diff --git a/code_puppy/api/pty_manager.py b/code_puppy/api/pty_manager.py deleted file mode 100644 index d57dacdc8..000000000 --- a/code_puppy/api/pty_manager.py +++ /dev/null @@ -1,453 +0,0 @@ -"""PTY Manager for terminal emulation with cross-platform support. - -Provides pseudo-terminal (PTY) functionality for interactive shell sessions -via WebSocket connections. Supports Unix (pty module) and Windows (pywinpty). -""" - -import asyncio -import logging -import os -import signal -import struct -import sys -from dataclasses import dataclass, field -from typing import Any, Callable, Optional - -logger = logging.getLogger(__name__) - -# Platform detection -IS_WINDOWS = sys.platform == "win32" - -# Conditional imports based on platform -if IS_WINDOWS: - try: - import winpty # type: ignore - - HAS_WINPTY = True - except ImportError: - HAS_WINPTY = False - winpty = None -else: - import fcntl - import pty - import termios - - HAS_WINPTY = False - - -@dataclass -class PTYSession: - """Represents an active PTY session.""" - - session_id: str - master_fd: Optional[int] = None # Unix only - slave_fd: Optional[int] = None # Unix only - pid: Optional[int] = None # Unix only - winpty_process: Any = None # Windows only - cols: int = 80 - rows: int = 24 - on_output: Optional[Callable[[bytes], None]] = None - _reader_task: Optional[asyncio.Task] = None # type: ignore - _running: bool = field(default=False, init=False) - - def is_alive(self) -> bool: - """Check if the PTY session is still active.""" - if IS_WINDOWS: - return self.winpty_process is not None and self.winpty_process.isalive() - else: - if self.pid is None: - return False - try: - pid_result, status = os.waitpid(self.pid, os.WNOHANG) - if pid_result == 0: - return True # Still running - return False # Exited - except ChildProcessError: - return False # Already reaped - - -class PTYManager: - """Manages PTY sessions for terminal emulation. - - Provides cross-platform terminal emulation with support for: - - Unix systems via the pty module - - Windows via pywinpty (optional dependency) - - Example: - manager = PTYManager() - session = await manager.create_session( - session_id="my-terminal", - on_output=lambda data: print(data.decode()) - ) - await manager.write(session.session_id, b"ls -la\n") - await manager.close_session(session.session_id) - """ - - def __init__(self) -> None: - self._sessions: dict[str, PTYSession] = {} - self._lock = asyncio.Lock() - - @property - def sessions(self) -> dict[str, PTYSession]: - """Get all active sessions.""" - return self._sessions.copy() - - async def create_session( - self, - session_id: str, - cols: int = 80, - rows: int = 24, - on_output: Optional[Callable[[bytes], None]] = None, - shell: Optional[str] = None, - ) -> PTYSession: - """Create a new PTY session. - - Args: - session_id: Unique identifier for the session - cols: Terminal width in columns - rows: Terminal height in rows - on_output: Callback for terminal output - shell: Shell to spawn (defaults to user's shell or /bin/bash) - - Returns: - PTYSession: The created session - - Raises: - RuntimeError: If session creation fails - """ - async with self._lock: - if session_id in self._sessions: - logger.warning(f"Session {session_id} already exists, closing old one") - await self._close_session_internal(session_id) - - if IS_WINDOWS: - session = await self._create_windows_session( - session_id, cols, rows, on_output, shell - ) - else: - session = await self._create_unix_session( - session_id, cols, rows, on_output, shell - ) - - self._sessions[session_id] = session - logger.info(f"Created PTY session: {session_id}") - return session - - async def _create_unix_session( - self, - session_id: str, - cols: int, - rows: int, - on_output: Optional[Callable[[bytes], None]], - shell: Optional[str], - ) -> PTYSession: - """Create a PTY session on Unix systems.""" - shell = shell or os.environ.get("SHELL", "/bin/bash") - - # Fork a new process with a PTY - pid, master_fd = pty.fork() - - if pid == 0: - # Child process - exec the shell - os.execlp(shell, shell, "-i") # noqa: S606 - else: - # Parent process - # Set terminal size - self._set_unix_winsize(master_fd, rows, cols) - - # Make master_fd non-blocking - flags = fcntl.fcntl(master_fd, fcntl.F_GETFL) - fcntl.fcntl(master_fd, fcntl.F_SETFL, flags | os.O_NONBLOCK) - - session = PTYSession( - session_id=session_id, - master_fd=master_fd, - pid=pid, - cols=cols, - rows=rows, - on_output=on_output, - ) - session._running = True - - # Start reader task - session._reader_task = asyncio.create_task(self._unix_reader_loop(session)) - - return session - - async def _create_windows_session( - self, - session_id: str, - cols: int, - rows: int, - on_output: Optional[Callable[[bytes], None]], - shell: Optional[str], - ) -> PTYSession: - """Create a PTY session on Windows systems.""" - if not HAS_WINPTY: - raise RuntimeError( - "pywinpty is required for Windows terminal support. " - "Install it with: pip install pywinpty" - ) - - shell = shell or os.environ.get("COMSPEC", "cmd.exe") - - # Create winpty process - winpty_process = winpty.PtyProcess.spawn( - shell, - dimensions=(rows, cols), - ) - - session = PTYSession( - session_id=session_id, - winpty_process=winpty_process, - cols=cols, - rows=rows, - on_output=on_output, - ) - session._running = True - - # Start reader task - session._reader_task = asyncio.create_task(self._windows_reader_loop(session)) - - return session - - def _set_unix_winsize(self, fd: int, rows: int, cols: int) -> None: - """Set the terminal window size on Unix.""" - winsize = struct.pack("HHHH", rows, cols, 0, 0) - fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize) - - async def _unix_reader_loop(self, session: PTYSession) -> None: - """Read output from Unix PTY and forward to callback.""" - loop = asyncio.get_running_loop() - - try: - while session._running and session.master_fd is not None: - try: - data = await loop.run_in_executor( - None, self._read_unix_pty, session.master_fd - ) - - if data is None: - # No data available, wait a bit - await asyncio.sleep(0.01) - continue - elif data == b"": - # EOF - process terminated - break - elif session.on_output: - session.on_output(data) - - except asyncio.CancelledError: - break - - except Exception as e: - logger.error(f"Unix reader loop error: {e}") - finally: - session._running = False - - def _read_unix_pty(self, fd: int) -> bytes | None: - """Read from Unix PTY file descriptor. - - Returns: - bytes: Data read from PTY - None: No data available (would block) - b'': EOF (process terminated) - """ - try: - data = os.read(fd, 4096) - return data - except BlockingIOError: - return None - except OSError: - return b"" - - async def _windows_reader_loop(self, session: PTYSession) -> None: - """Read output from Windows PTY and forward to callback.""" - loop = asyncio.get_running_loop() - - try: - while ( - session._running - and session.winpty_process is not None - and session.winpty_process.isalive() - ): - try: - data = await loop.run_in_executor( - None, session.winpty_process.read, 4096 - ) - if data and session.on_output: - session.on_output( - data.encode() if isinstance(data, str) else data - ) - except EOFError: - break - except asyncio.CancelledError: - break - - await asyncio.sleep(0.01) - - except Exception as e: - logger.error(f"Windows reader loop error: {e}") - finally: - session._running = False - - async def write(self, session_id: str, data: bytes) -> bool: - """Write data to a PTY session. - - Args: - session_id: The session to write to - data: Data to write - - Returns: - bool: True if write succeeded - """ - session = self._sessions.get(session_id) - if not session: - logger.warning(f"Session {session_id} not found") - return False - - try: - if IS_WINDOWS: - if session.winpty_process: - session.winpty_process.write( - data.decode() if isinstance(data, bytes) else data - ) - return True - else: - if session.master_fd is not None: - os.write(session.master_fd, data) - return True - except Exception as e: - logger.error(f"Write error for session {session_id}: {e}") - - return False - - async def resize(self, session_id: str, cols: int, rows: int) -> bool: - """Resize a PTY session. - - Args: - session_id: The session to resize - cols: New width in columns - rows: New height in rows - - Returns: - bool: True if resize succeeded - """ - session = self._sessions.get(session_id) - if not session: - logger.warning(f"Session {session_id} not found") - return False - - try: - if IS_WINDOWS: - if session.winpty_process: - session.winpty_process.setwinsize(rows, cols) - else: - if session.master_fd is not None: - self._set_unix_winsize(session.master_fd, rows, cols) - - session.cols = cols - session.rows = rows - logger.debug(f"Resized session {session_id} to {cols}x{rows}") - return True - - except Exception as e: - logger.error(f"Resize error for session {session_id}: {e}") - return False - - async def close_session(self, session_id: str) -> bool: - """Close a PTY session. - - Args: - session_id: The session to close - - Returns: - bool: True if session was closed - """ - async with self._lock: - return await self._close_session_internal(session_id) - - async def _close_session_internal(self, session_id: str) -> bool: - """Internal session close without lock.""" - session = self._sessions.pop(session_id, None) - if not session: - return False - - session._running = False - - # Cancel reader task - if session._reader_task: - session._reader_task.cancel() - try: - await session._reader_task - except asyncio.CancelledError: - pass - - # Clean up platform-specific resources - if IS_WINDOWS: - if session.winpty_process: - try: - session.winpty_process.terminate() - except Exception as e: - logger.debug(f"Error terminating winpty: {e}") - else: - # Close file descriptors - if session.master_fd is not None: - try: - os.close(session.master_fd) - except OSError: - pass - - # Terminate child process - if session.pid is not None: - try: - os.kill(session.pid, signal.SIGTERM) - # Use WNOHANG to avoid blocking the event loop - try: - os.waitpid(session.pid, os.WNOHANG) - except ChildProcessError: - pass - except (OSError, ChildProcessError): - pass - - logger.info(f"Closed PTY session: {session_id}") - return True - - async def close_all(self) -> None: - """Close all PTY sessions.""" - async with self._lock: - session_ids = list(self._sessions.keys()) - for session_id in session_ids: - await self._close_session_internal(session_id) - logger.info("Closed all PTY sessions") - - def get_session(self, session_id: str) -> Optional[PTYSession]: - """Get a session by ID. - - Args: - session_id: The session ID - - Returns: - PTYSession or None if not found - """ - return self._sessions.get(session_id) - - def list_sessions(self) -> list[str]: - """List all active session IDs. - - Returns: - List of session IDs - """ - return list(self._sessions.keys()) - - -# Global PTY manager instance -_pty_manager: Optional[PTYManager] = None - - -def get_pty_manager() -> PTYManager: - """Get or create the global PTY manager instance.""" - global _pty_manager - if _pty_manager is None: - _pty_manager = PTYManager() - return _pty_manager diff --git a/code_puppy/api/templates/terminal.html b/code_puppy/api/templates/terminal.html deleted file mode 100644 index 568dc7491..000000000 --- a/code_puppy/api/templates/terminal.html +++ /dev/null @@ -1,382 +0,0 @@ - - - - - - 🐶 Code Puppy Terminal - - - - - - - - - - - - - - - -
- -
-
- 🐶 -

Code Puppy Terminal

-
- -
- -
-
- Disconnected -
- - - - -
- - -
-
-
- - -
-
-
- - -
-
- 80x24 - • -
- Agent: - -
-
-
- Ctrl+C to interrupt - • - Ctrl+D to exit -
-
-
- - - - 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..fff1673f8 --- /dev/null +++ b/code_puppy/api/tests/test_gate2_legacy_ws_namespace.py @@ -0,0 +1,42 @@ +"""Gate 2 smoke tests for puppy-desk migration safety. + +These tests ensure the legacy WS snapshot exists while the active WebSocket +registration still exposes the original `/ws/chat` route and does not expose a +migration side-route yet. +""" + +from pathlib import Path + +from fastapi import FastAPI + + +def test_setup_websocket_keeps_existing_chat_route_and_no_migration_route(): + from code_puppy.api.websocket import setup_websocket + + app = FastAPI() + 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/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_namespace_exists_without_route_registration(): + import code_puppy.api.ws.legacy as legacy_ws + + legacy_dir = Path(legacy_ws.__file__).parent + + assert legacy_ws.LEGACY_SNAPSHOT_PURPOSE == "puppy-desk-gate2-reference-copy" + assert (legacy_dir / "chat_handler.py").is_file() + assert (legacy_dir / "schemas.py").is_file() + assert (legacy_dir / "runtime" / "session_runtime_manager.py").is_file() + assert (legacy_dir / "README.md").is_file() diff --git a/code_puppy/api/websocket.py b/code_puppy/api/websocket.py index f6208f104..e4ecfca05 100644 --- a/code_puppy/api/websocket.py +++ b/code_puppy/api/websocket.py @@ -3,9 +3,7 @@ Provides real-time communication channels: - /ws/events - Server-sent events stream - /ws/chat - Interactive chat with the agent -- /ws/terminal - Interactive PTY terminal sessions - /ws/health - Simple health check endpoint -- /ws/sessions - Real-time session monitoring 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 @@ -20,8 +18,6 @@ register_chat_endpoint, register_events_endpoint, register_health_endpoint, - register_sessions_endpoint, - register_terminal_endpoint, ) # Re-export build_file_context_and_attachments for backward compatibility @@ -44,8 +40,6 @@ def setup_websocket(app: FastAPI) -> None: """ register_events_endpoint(app) register_chat_endpoint(app) - register_terminal_endpoint(app) register_health_endpoint(app) - register_sessions_endpoint(app) - logger.debug("All WebSocket endpoints registered") + logger.debug("WebSocket endpoints registered") diff --git a/code_puppy/api/ws/__init__.py b/code_puppy/api/ws/__init__.py index 1649afd36..55eaab749 100644 --- a/code_puppy/api/ws/__init__.py +++ b/code_puppy/api/ws/__init__.py @@ -6,9 +6,7 @@ Modules: - chat_handler: Interactive chat with the Code Puppy agent (/ws/chat) - events_handler: Server-sent events stream (/ws/events) - - terminal_handler: Interactive PTY terminal (/ws/terminal) - health_handler: Simple health check (/ws/health) - - sessions_handler: Real-time session monitoring (/ws/sessions) - connection_manager: Singleton for broadcasting session updates - attachments: File attachment processing utilities """ @@ -17,14 +15,10 @@ from code_puppy.api.ws.connection_manager import connection_manager from code_puppy.api.ws.events_handler import register_events_endpoint from code_puppy.api.ws.health_handler import register_health_endpoint -from code_puppy.api.ws.sessions_handler import register_sessions_endpoint -from code_puppy.api.ws.terminal_handler import register_terminal_endpoint __all__ = [ "register_chat_endpoint", "register_events_endpoint", - "register_terminal_endpoint", "register_health_endpoint", - "register_sessions_endpoint", "connection_manager", ] diff --git a/code_puppy/api/ws/terminal_handler.py b/code_puppy/api/ws/terminal_handler.py deleted file mode 100644 index b32bff451..000000000 --- a/code_puppy/api/ws/terminal_handler.py +++ /dev/null @@ -1,141 +0,0 @@ -"""WebSocket endpoint for interactive PTY terminal sessions.""" - -import asyncio -import base64 -import logging -import uuid - -from fastapi import FastAPI, WebSocket, WebSocketDisconnect - -logger = logging.getLogger(__name__) - - -def register_terminal_endpoint(app: FastAPI) -> None: - """Register the /ws/terminal WebSocket endpoint.""" - - @app.websocket("/ws/terminal") - async def websocket_terminal(websocket: WebSocket) -> None: - """Interactive terminal WebSocket endpoint.""" - await websocket.accept() - logger.debug("Terminal WebSocket client connected") - - from code_puppy.api.pty_manager import get_pty_manager - - manager = get_pty_manager() - session_id = str(uuid.uuid4())[:8] - session = None - - loop = asyncio.get_running_loop() - output_queue: asyncio.Queue[bytes] = asyncio.Queue() - - def on_output(data: bytes) -> None: - try: - loop.call_soon_threadsafe(output_queue.put_nowait, data) - except Exception as e: - logger.error("on_output error: %s", e) - - async def output_sender() -> None: - """Coroutine that sends queued output to WebSocket.""" - try: - while True: - data = await output_queue.get() - await websocket.send_json( - { - "type": "output", - "data": base64.b64encode(data).decode("ascii"), - } - ) - except asyncio.CancelledError: - pass - except Exception as e: - logger.error("output_sender error: %s", e) - - sender_task = None - - try: - session = await manager.create_session( - session_id=session_id, - on_output=on_output, - ) - - from code_puppy.api.session_context import session_manager - - try: - ctx = await session_manager.create_session(f"terminal_{session_id}") - current_agent_name = ctx.agent_name - except Exception: - current_agent_name = "code-puppy" - ctx = None - - await websocket.send_json( - { - "type": "session_started", - "session_id": session_id, - "current_agent": current_agent_name, - } - ) - - sender_task = asyncio.create_task(output_sender()) - - while True: - try: - msg = await websocket.receive_json() - - if msg.get("type") == "input": - data = msg.get("data", "") - if isinstance(data, str): - data = data.encode("utf-8") - await manager.write(session_id, data) - elif msg.get("type") == "resize": - cols = msg.get("cols", 80) - rows = msg.get("rows", 24) - await manager.resize(session_id, cols, rows) - elif msg.get("type") == "switch_agent": - agent_name = msg.get("agent_name") - if agent_name: - try: - await session_manager.switch_agent( - f"terminal_{session_id}", agent_name - ) - if ctx: - ctx = await session_manager.get_session( - f"terminal_{session_id}" - ) - await websocket.send_json( - {"type": "agent_changed", "agent_name": agent_name} - ) - except Exception as e: - logger.error("Error switching agent: %s", e) - await websocket.send_json( - { - "type": "error", - "message": f"Failed to switch to agent {agent_name}: {str(e)}", - } - ) - elif msg.get("type") == "get_current_agent": - if ctx: - current_agent_name = ctx.agent_name - else: - current_agent_name = "code-puppy" - await websocket.send_json( - {"type": "current_agent", "agent_name": current_agent_name} - ) - except WebSocketDisconnect: - break - except Exception as e: - logger.error("Terminal WebSocket error: %s", e) - break - except Exception as e: - logger.error("Terminal session error: %s", e) - finally: - if sender_task: - sender_task.cancel() - try: - await sender_task - except asyncio.CancelledError: - pass - if ctx: - await session_manager.destroy_session(f"terminal_{session_id}") - if session: - await manager.close_session(session_id) - logger.debug("Terminal WebSocket disconnected") From be3c13c23d666234e8a3b529341a0f74c7261b31 Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Tue, 2 Jun 2026 10:36:23 -0500 Subject: [PATCH 12/29] Clean up puppy-desk websocket migration --- .../tests/test_gate2_legacy_ws_namespace.py | 49 +- .../test_streaming_protocol_transform.py | 32 +- code_puppy/api/tests/test_ws_history_utils.py | 91 + code_puppy/api/ws/__init__.py | 44 +- code_puppy/api/ws/chat_handler.py | 333 +- code_puppy/api/ws/history_utils.py | 104 + code_puppy/api/ws/legacy/README.md | 11 - code_puppy/api/ws/legacy/__init__.py | 12 - code_puppy/api/ws/legacy/attachments.py | 93 - code_puppy/api/ws/legacy/background_save.py | 244 -- code_puppy/api/ws/legacy/chat_handler.py | 3898 ----------------- .../api/ws/legacy/connection_manager.py | 55 - code_puppy/api/ws/legacy/events_handler.py | 53 - code_puppy/api/ws/legacy/health_handler.py | 22 - code_puppy/api/ws/legacy/runtime/__init__.py | 14 - .../api/ws/legacy/runtime/session_runtime.py | 68 - .../legacy/runtime/session_runtime_manager.py | 150 - code_puppy/api/ws/legacy/schemas.py | 603 --- code_puppy/api/ws/legacy/sessions_handler.py | 69 - code_puppy/api/ws/response_frames.py | 151 + docs/migration/puppy-desk-cleanup-summary.md | 41 + docs/migration/puppy-desk-phase3-plan.md | 127 + 22 files changed, 630 insertions(+), 5634 deletions(-) create mode 100644 code_puppy/api/tests/test_ws_history_utils.py create mode 100644 code_puppy/api/ws/history_utils.py delete mode 100644 code_puppy/api/ws/legacy/README.md delete mode 100644 code_puppy/api/ws/legacy/__init__.py delete mode 100644 code_puppy/api/ws/legacy/attachments.py delete mode 100644 code_puppy/api/ws/legacy/background_save.py delete mode 100644 code_puppy/api/ws/legacy/chat_handler.py delete mode 100644 code_puppy/api/ws/legacy/connection_manager.py delete mode 100644 code_puppy/api/ws/legacy/events_handler.py delete mode 100644 code_puppy/api/ws/legacy/health_handler.py delete mode 100644 code_puppy/api/ws/legacy/runtime/__init__.py delete mode 100644 code_puppy/api/ws/legacy/runtime/session_runtime.py delete mode 100644 code_puppy/api/ws/legacy/runtime/session_runtime_manager.py delete mode 100644 code_puppy/api/ws/legacy/schemas.py delete mode 100644 code_puppy/api/ws/legacy/sessions_handler.py create mode 100644 code_puppy/api/ws/response_frames.py create mode 100644 docs/migration/puppy-desk-cleanup-summary.md create mode 100644 docs/migration/puppy-desk-phase3-plan.md diff --git a/code_puppy/api/tests/test_gate2_legacy_ws_namespace.py b/code_puppy/api/tests/test_gate2_legacy_ws_namespace.py index fff1673f8..d1aaedde0 100644 --- a/code_puppy/api/tests/test_gate2_legacy_ws_namespace.py +++ b/code_puppy/api/tests/test_gate2_legacy_ws_namespace.py @@ -1,20 +1,34 @@ -"""Gate 2 smoke tests for puppy-desk migration safety. - -These tests ensure the legacy WS snapshot exists while the active WebSocket -registration still exposes the original `/ws/chat` route and does not expose a -migration side-route yet. -""" +"""Post-cleanup WebSocket namespace tests for puppy-desk migration.""" from pathlib import Path -from fastapi import FastAPI +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 test_setup_websocket_keeps_existing_chat_route_and_no_migration_route(): - from code_puppy.api.websocket import setup_websocket + 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() - setup_websocket(app) + websocket_module.setup_websocket(app) websocket_paths = { getattr(route, "path", None) @@ -23,6 +37,8 @@ def test_setup_websocket_keeps_existing_chat_route_and_no_migration_route(): } 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 @@ -30,13 +46,6 @@ def test_setup_websocket_keeps_existing_chat_route_and_no_migration_route(): assert "/ws/chat-v2" not in websocket_paths -def test_legacy_ws_snapshot_namespace_exists_without_route_registration(): - import code_puppy.api.ws.legacy as legacy_ws - - legacy_dir = Path(legacy_ws.__file__).parent - - assert legacy_ws.LEGACY_SNAPSHOT_PURPOSE == "puppy-desk-gate2-reference-copy" - assert (legacy_dir / "chat_handler.py").is_file() - assert (legacy_dir / "schemas.py").is_file() - assert (legacy_dir / "runtime" / "session_runtime_manager.py").is_file() - assert (legacy_dir / "README.md").is_file() +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_streaming_protocol_transform.py b/code_puppy/api/tests/test_streaming_protocol_transform.py index a70d3f6cd..350b7ff3a 100644 --- a/code_puppy/api/tests/test_streaming_protocol_transform.py +++ b/code_puppy/api/tests/test_streaming_protocol_transform.py @@ -1,6 +1,9 @@ """Tests for streaming-only assistant response protocol adaptation.""" -from code_puppy.api.ws.chat_handler import build_assistant_text_stream_frames +from code_puppy.api.ws.response_frames import ( + build_assistant_text_stream_frames, + build_error_response_frames, +) def _dump_frames(**kwargs): @@ -91,3 +94,30 @@ def test_complete_response_transform_preserves_part_type_for_thinking_parts(): 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/ws/__init__.py b/code_puppy/api/ws/__init__.py index 55eaab749..8954576f2 100644 --- a/code_puppy/api/ws/__init__.py +++ b/code_puppy/api/ws/__init__.py @@ -1,20 +1,38 @@ """WebSocket endpoint handlers package. -Each handler module exposes a `register_*_endpoint(app)` function -that attaches one WebSocket route to the FastAPI application. - -Modules: - - chat_handler: Interactive chat with the Code Puppy agent (/ws/chat) - - events_handler: Server-sent events stream (/ws/events) - - health_handler: Simple health check (/ws/health) - - connection_manager: Singleton for broadcasting session updates - - attachments: File attachment processing utilities +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 code_puppy.api.ws.chat_handler import register_chat_endpoint -from code_puppy.api.ws.connection_manager import connection_manager -from code_puppy.api.ws.events_handler import register_events_endpoint -from code_puppy.api.ws.health_handler import register_health_endpoint +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", diff --git a/code_puppy/api/ws/chat_handler.py b/code_puppy/api/ws/chat_handler.py index c8af429a7..f22dc4280 100644 --- a/code_puppy/api/ws/chat_handler.py +++ b/code_puppy/api/ws/chat_handler.py @@ -33,7 +33,6 @@ write_system_message_to_sqlite, write_turn_to_sqlite, ) -from code_puppy.api.error_parser import parse_api_error as _legacy_parse_api_error from code_puppy.api.session_context import _validate_session_id, session_manager from code_puppy.api.ws.attachments import build_file_context_and_attachments from code_puppy.api.ws.background_save import ( @@ -41,6 +40,15 @@ save_agent_result_in_background, ) from code_puppy.api.ws.connection_manager import connection_manager +from code_puppy.api.ws.response_frames import ( + build_assistant_text_stream_frames, + build_error_response_frames, + parse_api_error, +) +from code_puppy.api.ws.history_utils import ( + build_enhanced_history, + estimate_total_tokens, +) from code_puppy.api.ws.schemas import ( PROTOCOL_VERSION, ClientMessage, @@ -66,7 +74,6 @@ ) from code_puppy.config import get_global_model_name from code_puppy.messaging.bus import get_message_bus -from code_puppy.model_errors import normalize_model_error as _normalize_model_error from code_puppy.session_storage import generate_heuristic_title from code_puppy.tools.command_runner import ( cleanup_session_process_tracking, @@ -95,185 +102,6 @@ def clear_session_working_directory(): logger = logging.getLogger(__name__) -# === TOOL-CALL-PARITY BRANCH LOADED === -logger.warning( - "🐕 CHAT_HANDLER LOADED FROM TOOL-CALL-PARITY BRANCH - ToolReturnPart fix active!" -) -print("🐕 CHAT_HANDLER LOADED FROM TOOL-CALL-PARITY BRANCH", flush=True) - - -# --------------------------------------------------------------------------- -# WebSocket error frame helpers (see fix/orphaned-tool-use-id-history) -# --------------------------------------------------------------------------- - -_UNKNOWN_ERROR_TYPE = "unknown_error" -_TOOL_HISTORY_ERROR_TYPE = "tool_history_error" - -_CODE_TO_ERROR_TYPE: dict = { - "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: - """Convert an agent-run exception into a structured frontend error dict. - - Wraps ``normalize_model_error`` for primary classification (including the - new ``invalid_tool_history`` code for Anthropic HTTP 400 tool-mismatch - errors), falling back to the legacy ``_legacy_parse_api_error`` for all - other categories so existing behaviour is unchanged. - - Returns: - Dict with keys: user_message, error_type, technical_details, action_required. - """ - 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, - } - # Fall back to legacy parser for unclassified errors - return _legacy_parse_api_error(exc) - - -def has_streamed_content(collected_text) -> bool: - """Return True when the client has already received streaming output chunks. - - Args: - collected_text: List of text chunks accumulated during the agent run. - Elements may be None or empty strings. - - Returns: - True if at least one chunk has substantive (non-whitespace) content. - """ - return any((chunk or "").strip() for chunk in collected_text) - - -def build_error_response_frames( - agent_error: Exception, - collected_text, - session_id: str, -) -> list: - """Build the ordered WebSocket frames to send when an agent error occurs. - - When streaming output was already delivered to the client, a ``stream_end`` - frame (``success=False``) is prepended so the frontend can exit its - streaming state before processing the error frame. Without this handshake - the frontend hangs indefinitely waiting for a ``stream_end`` that never - arrives. - - Args: - agent_error: The exception raised by the agent run. - collected_text: Accumulated streaming chunks from the current turn. - session_id: The WebSocket session identifier. - - Returns: - List of JSON-serialisable dicts to send to the client in order. - """ - frames: list[dict] = [] - 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. - - The GUI protocol is streaming-only: assistant text must be rendered from - ``assistant_message_start``/``assistant_message_delta``/ - ``assistant_message_end`` and finalized by ``stream_end``. Some upstream - model clients return a complete response rather than true token deltas; this - helper adapts those complete responses to the same wire shape so the GUI has - one rendering path. - - This intentionally emits a single full-content delta for non-streaming - upstream responses. Artificial local chunking can be added later without - changing the protocol contract. - """ - 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, - ), - ] - def register_chat_endpoint(app: FastAPI) -> None: """Register the /ws/chat WebSocket endpoint.""" @@ -3688,136 +3516,25 @@ async def drain_events_with_signal(): or "unknown" ) - def _extract_message_timestamp( - raw_msg: Any, default_ts: str - ) -> str: - """Best-effort extraction of an existing timestamp for a message. - - This is important for WebSocket sessions where history may - already contain older messages. We don't want to overwrite - their original timestamps every time we auto-save. - - Precedence: - 1. If the message is a dict with a numeric epoch 'timestamp', - convert to ISO. - 2. If the message is a dict with an ISO-ish 'timestamp' str, - reuse it as-is. - 3. If the message is a dict with 'ts', reuse it. - 4. If the message object has a 'timestamp' attribute, try that. - 5. Fall back to the provided default_ts ("now" for new messages). - """ - # Dict-based histories (e.g. CLI format or older WS formats) - if isinstance(raw_msg, dict): - ts_val = raw_msg.get("timestamp") - - # Epoch seconds - if isinstance(ts_val, (int, float)): - try: - return ( - datetime.datetime.fromtimestamp( - ts_val - ).isoformat() - ) - except Exception: - pass - - # Already an ISO string - if isinstance(ts_val, str) and ts_val: - return ts_val - - # Our enhanced WS wrapper sometimes uses 'ts' - ts_field = raw_msg.get("ts") - if isinstance(ts_field, str) and ts_field: - return ts_field - - # Pydantic / custom objects that carry a timestamp attribute - 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 - - # Fallback: use provided default - return default_ts - - # Create enhanced history: list of dicts with message + metadata - # Each entry: {'msg': , 'agent': str, 'model': str, 'ts': str} - enhanced_history = [] - for idx, msg in enumerate(history): - # Check if already wrapped (for idempotency). If the wrapper - # already has a 'ts', leave it untouched so older sessions - # keep their original per-message timestamps. - if ( - isinstance(msg, dict) - and "msg" in msg - and "agent" in msg - ): - enhanced_history.append(msg) - else: - # Use existing timestamp if present; otherwise compute a default now() - current_timestamp = ( - datetime.datetime.now().isoformat() - ) - msg_ts = _extract_message_timestamp( - msg, current_timestamp - ) - wrapper = { - "msg": msg, - "agent": agent_name_meta, - "model": model_name_meta, - "ts": msg_ts, - } - - # Add clean_content and attachments to the user message we just processed. - # clean_content is only needed when attachments were injected into content - # (file blocks like --- File: auth.ts ---). Without attachments, content - # is already the user's words (FE strips [Session Context:] as fallback). - is_user_message_just_processed = ( - idx == len(history) - 2 - and len(history) >= 2 - and attachment_metadata # only needed when file content was injected - ) - - if is_user_message_just_processed: - wrapper["clean_content"] = ( - original_user_message - ) - wrapper["attachments"] = ( - attachment_metadata - ) - logger.debug( - "Added UI metadata to user message: %d attachment(s), " - "clean_content length: %d", - len(attachment_metadata), - len(original_user_message), - ) + 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, + ) - enhanced_history.append(wrapper) + 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 = 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 + total_tokens = estimate_total_tokens( + enhanced_history, agent + ) # Write to SQLite for FE read path try: 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/legacy/README.md b/code_puppy/api/ws/legacy/README.md deleted file mode 100644 index 113e3a41f..000000000 --- a/code_puppy/api/ws/legacy/README.md +++ /dev/null @@ -1,11 +0,0 @@ -# Legacy WS Snapshot — Puppy Desk Gate 2 - -This directory is a reference/rollback copy of `code_puppy.api.ws` captured -before further puppy-desk migration/refactor work. - -Rules: - -- It is not registered by `setup_websocket()`. -- Existing `/ws/chat` behavior remains provided by `code_puppy.api.ws.chat_handler`. -- Same-route swap must not happen until parity and GUGI validation pass. -- Keep this snapshot until user explicitly approves cleanup. diff --git a/code_puppy/api/ws/legacy/__init__.py b/code_puppy/api/ws/legacy/__init__.py deleted file mode 100644 index fae8b03db..000000000 --- a/code_puppy/api/ws/legacy/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -"""Legacy WebSocket implementation snapshot for puppy-desk migration. - -This package is a reference/rollback copy captured during Gate 2 of the -puppy-desk migration. It is intentionally not imported or registered by the -runtime application yet; the active routes continue to come from -``code_puppy.api.ws``. - -Do not edit this snapshot as part of feature work unless intentionally updating -the legacy rollback point. -""" - -LEGACY_SNAPSHOT_PURPOSE = "puppy-desk-gate2-reference-copy" diff --git a/code_puppy/api/ws/legacy/attachments.py b/code_puppy/api/ws/legacy/attachments.py deleted file mode 100644 index f2db90726..000000000 --- a/code_puppy/api/ws/legacy/attachments.py +++ /dev/null @@ -1,93 +0,0 @@ -"""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/legacy/background_save.py b/code_puppy/api/ws/legacy/background_save.py deleted file mode 100644 index 4b1a1c6ad..000000000 --- a/code_puppy/api/ws/legacy/background_save.py +++ /dev/null @@ -1,244 +0,0 @@ -"""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/legacy/chat_handler.py b/code_puppy/api/ws/legacy/chat_handler.py deleted file mode 100644 index fbcb4c918..000000000 --- a/code_puppy/api/ws/legacy/chat_handler.py +++ /dev/null @@ -1,3898 +0,0 @@ -"""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 -import re -import uuid -from pathlib import Path -from typing import Any - -from fastapi import FastAPI, WebSocket, WebSocketDisconnect -from pydantic import TypeAdapter, ValidationError - -from code_puppy.api.db.queries import ( - update_session_working_directory, - write_error_message_to_sqlite, - write_system_message_to_sqlite, - write_turn_to_sqlite, -) -from code_puppy.api.error_parser import parse_api_error as _legacy_parse_api_error -from code_puppy.api.session_context import _validate_session_id, session_manager -from code_puppy.api.ws.attachments import build_file_context_and_attachments -from code_puppy.api.ws.background_save import ( - fire_and_track, - save_agent_result_in_background, -) -from code_puppy.api.ws.connection_manager import connection_manager -from code_puppy.api.ws.schemas import ( - PROTOCOL_VERSION, - ClientMessage, - ServerAgentInvoked, - ServerAssistantMessageDelta, - ServerAssistantMessageEnd, - ServerAssistantMessageStart, - ServerCancelled, - ServerCommandResult, - ServerConfigValue, - ServerError, - ServerMessage, - ServerResponse, - ServerSessionMetaUpdated, - ServerSessionRestored, - ServerSessionSwitched, - ServerStatus, - ServerStreamEnd, - ServerSystem, - ServerToolCall, - ServerToolResult, - ServerUserMessage, - ServerWorkingDirectoryChanged, -) -from code_puppy.config import get_global_model_name -from code_puppy.messaging.bus import get_message_bus -from code_puppy.model_errors import normalize_model_error as _normalize_model_error -from code_puppy.session_storage import generate_heuristic_title -from code_puppy.tools.command_runner import ( - cleanup_session_process_tracking, - init_session_process_tracking, -) - -# Import session context for desk-puppy working directory support -try: - from code_puppy.plugins.walmart_specific.session_context import ( - clear_session_working_directory, - set_session_working_directory, - ) - - HAS_SESSION_CONTEXT = True -except ImportError: - HAS_SESSION_CONTEXT = False - - def set_session_working_directory(d): - pass - - def clear_session_working_directory(): - pass - - -_ClientMessageAdapter = TypeAdapter(ClientMessage) - -logger = logging.getLogger(__name__) - -# === TOOL-CALL-PARITY BRANCH LOADED === -logger.warning( - "🐕 CHAT_HANDLER LOADED FROM TOOL-CALL-PARITY BRANCH - ToolReturnPart fix active!" -) -print("🐕 CHAT_HANDLER LOADED FROM TOOL-CALL-PARITY BRANCH", flush=True) - - -# --------------------------------------------------------------------------- -# WebSocket error frame helpers (see fix/orphaned-tool-use-id-history) -# --------------------------------------------------------------------------- - -_UNKNOWN_ERROR_TYPE = "unknown_error" -_TOOL_HISTORY_ERROR_TYPE = "tool_history_error" - -_CODE_TO_ERROR_TYPE: dict = { - "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: - """Convert an agent-run exception into a structured frontend error dict. - - Wraps ``normalize_model_error`` for primary classification (including the - new ``invalid_tool_history`` code for Anthropic HTTP 400 tool-mismatch - errors), falling back to the legacy ``_legacy_parse_api_error`` for all - other categories so existing behaviour is unchanged. - - Returns: - Dict with keys: user_message, error_type, technical_details, action_required. - """ - 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, - } - # Fall back to legacy parser for unclassified errors - return _legacy_parse_api_error(exc) - - -def has_streamed_content(collected_text) -> bool: - """Return True when the client has already received streaming output chunks. - - Args: - collected_text: List of text chunks accumulated during the agent run. - Elements may be None or empty strings. - - Returns: - True if at least one chunk has substantive (non-whitespace) content. - """ - return any((chunk or "").strip() for chunk in collected_text) - - -def build_error_response_frames( - agent_error: Exception, - collected_text, - session_id: str, -) -> list: - """Build the ordered WebSocket frames to send when an agent error occurs. - - When streaming output was already delivered to the client, a ``stream_end`` - frame (``success=False``) is prepended so the frontend can exit its - streaming state before processing the error frame. Without this handshake - the frontend hangs indefinitely waiting for a ``stream_end`` that never - arrives. - - Args: - agent_error: The exception raised by the agent run. - collected_text: Accumulated streaming chunks from the current turn. - session_id: The WebSocket session identifier. - - Returns: - List of JSON-serialisable dicts to send to the client in order. - """ - frames: list[dict] = [] - 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 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 response - {"type": "response", "content": "...", "done": true, - "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 - ) - - # Flag to track if WebSocket is still open - ws_closed = False - - async def persist_error_payload(data: dict[str, Any]) -> None: - """Persist structured error frames so they survive reloads.""" - if data.get("type") != "error" or not session_id: - return - - try: - await write_error_message_to_sqlite( - session_id=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=(ctx.agent_name if ctx else ""), - model_name=(ctx.model_name if ctx else ""), - timestamp=datetime.datetime.now(datetime.timezone.utc).isoformat(), - ) - except Exception: - logger.warning( - "[WS:%s] Failed to persist error payload to SQLite", - session_id, - exc_info=True, - ) - - async def safe_send_json(data: dict) -> bool: - """Safely send JSON to WebSocket, returns False if connection is closed.""" - nonlocal ws_closed - if ws_closed: - logger.debug( - "[WS:%s] safe_send_json skipped (ws_closed=True): type=%s", - session_id, - data.get("type"), - ) - return False - - msg_type = data.get("type") - if msg_type in { - "error", - "response", - "assistant_message_start", - "assistant_message_end", - }: - # Keep this low-noise but useful for tracing request lifecycle - logger.debug( - "[WS:%s] → send_json type=%s keys=%s", - 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", - session_id, - data.get("error_type"), - data.get("action_required"), - (data.get("error") or "")[:300], - ) - - try: - if msg_type == "error": - await persist_error_payload(data) - await websocket.send_json(data) - return True - except Exception as e: - logger.warning( - "[WS:%s] send_json failed for type=%s: %s", - session_id, - msg_type, - e, - exc_info=True, - ) - if "close message" in str(e).lower() or "closed" in str(e).lower(): - ws_closed = True - logger.debug("WebSocket closed, stopping sends") - return False - - async def send_typed(msg: ServerMessage) -> bool: - """Send a typed protocol message to the client.""" - return await safe_send_json(msg.model_dump(exclude_none=True)) - - async def send_typed_tool_lifecycle( - msg: ServerToolCall | ServerToolResult, - ) -> bool: - """Send tool lifecycle frames to the client.""" - return await send_typed(msg) - - ctx = None - session_title = "" - session_working_directory = "" # Will be set by client or loaded from metadata - session_pinned = False # Track pinned state for persistence - last_context_sent_directory = "" # Track when we last sent directory context - existing_history = None - - # Use provided session_id or generate new one - if session_id: - # Validate session_id to prevent path traversal - try: - _validate_session_id(session_id) - except ValueError as e: - logger.warning("Invalid session_id rejected: %r: %s", session_id, e) - await websocket.close(code=1008, reason="Invalid session ID") - return - - # Client wants to resume/continue a specific session - logger.debug("Client requested session: %s", session_id) - - # SQLite is the source of truth — check session existence there. - 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 e: - logger.warning("Failed to check session in SQLite: %s", e) - - if not existing_history: - logger.debug( - "Session %s not found in SQLite, will create new", session_id - ) - else: - # Generate new timestamp-based session ID - session_timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - session_id = f"WS_session_{session_timestamp}" - logger.debug("Generated new session ID: %s", session_id) - - try: - # --- SESSION ISOLATION: create or load via SessionManager --- - try: - if existing_history is not None: - # Session confirmed in SQLite — load it via SessionManager. - # get_or_load_session checks in-memory first, then falls - # back to a fresh SQLite load. SQLite is the sole source of truth. - ctx = await session_manager.get_or_load_session(session_id) - if ctx is None: - # SQLite confirmed this session exists but load returned None. - # This is unexpected (DB race, corrupt row, or schema mismatch). - # Log a warning and start fresh so the WebSocket can still - # function — silent data loss is worse than a blank session. - 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, - ) - ctx = await session_manager.create_session(session_id) - else: - # Update local state from loaded metadata - session_title = ctx.title - session_working_directory = ctx.working_directory - session_pinned = ctx.pinned - else: - ctx = await session_manager.create_session(session_id) - except Exception as e: - logger.warning("SessionManager init failed, falling back: %s", e) - try: - ctx = await session_manager.create_session(session_id) - except Exception: - logger.error("SessionManager fallback also failed", exc_info=True) - await websocket.close(code=1011, reason="Session init failed") - return - - # Mark session as active (cancels any pending cleanup from previous disconnect) - await session_manager.mark_session_active(session_id) - - # Initialize per-session process tracking (ContextVar isolation) - init_session_process_tracking() - - # Set MessageBus session context for this WS connection - try: - bus = get_message_bus() - bus.set_session_context(session_id) - except Exception: - logger.debug("MessageBus session context not available") - - # Convenience aliases — use ctx.agent everywhere instead of get_current_agent() - agent = ctx.agent - agent_name = ctx.agent_name - model_name = ctx.model_name - - # ── Persist initial session context to SQLite (new sessions only) ─────── - # Write a config row (agent + model) and, if a CWD is set, a directory - # row so the FE cold-load path sees them before the first user message. - # Skipped for resumed sessions — their init rows were written at creation. - if existing_history is None: - try: - _now_iso = datetime.datetime.now(datetime.timezone.utc).isoformat() - _init_agent = agent_name or "code-puppy" - _init_model = model_name or "unknown" - await write_system_message_to_sqlite( - session_id=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 session_working_directory: - _path_segments = session_working_directory.split("/")[-3:] - _relative = "/".join(s for s in _path_segments if s) - await write_system_message_to_sqlite( - session_id=session_id, - system_message_type="directory", - content=f"Starting in {_relative}", - system_message_path=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, - ) - - # Send welcome message - await send_typed( - ServerSystem( - content=f"Connected! Session: {session_id}", - session_id=session_id, - agent_name=agent_name, - model_name=model_name, - resumed=existing_history is not None, - protocol_version=PROTOCOL_VERSION, - ) - ) - - # If we resumed an existing session, notify the client - if existing_history and ctx: - try: - # History is already loaded into ctx.agent by session_manager.load_session(). - # Pull it back out for the client-side UI metadata notification. - loaded_messages = ctx.agent.get_message_history() - message_count = len(loaded_messages) if loaded_messages else 0 - - # Notify client about restored session - await send_typed( - ServerSessionRestored( - session_id=session_id, - message_count=message_count, - title=session_title, - ui_metadata=[], - ) - ) - - # Replay system messages (agent/model switches, CWD banners) to UI. - # These are persisted with pydantic_json=NULL so _load_from_sqlite - # skips them — we must send them separately so the UI shows them. - try: - from code_puppy.api.db.queries import get_active_messages - - rows = await get_active_messages(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=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), - session_id, - ) - except Exception as sys_exc: - logger.warning( - "Failed to replay system messages for session %s: %s", - session_id, - sys_exc, - ) - - agent = ctx.agent # Refresh local alias - logger.debug("Restored %d messages to session agent", message_count) - except Exception as e: - logger.warning("Failed to restore session history: %s", e) - # Track active streaming task for cancellation - active_drain_task = None - # Track active agent task (run_with_mcp) for cancellation - active_agent_task: asyncio.Task | None = None - # Track stop_draining event for cancellation across iterations - stop_draining = asyncio.Event() - - 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") == "switch_agent": - agent_name = msg.get("agent_name") - if agent_name: - try: - new_agent = await session_manager.switch_agent( - session_id, agent_name - ) - agent = new_agent # Update local alias - ctx.agent = new_agent - model_name = ( - new_agent.get_model_name() - if new_agent - else "unknown" - ) - - # Persist to SQLite so the FE cold-load path sees it - try: - await write_system_message_to_sqlite( - session_id=session_id, - system_message_type="config", - content=f"🔄 Switched to {agent_name} ({model_name})", - agent_name=agent_name, - model_name=model_name, - ) - except Exception as _sw_exc: - logger.warning( - "Agent-switch SQLite write failed: %s", - _sw_exc, - exc_info=True, - ) - - await send_typed( - ServerSystem( - content=f"🔄 Switched to {agent_name} ({model_name})", - session_id=session_id, - agent_name=agent_name, - model_name=model_name, - ) - ) - except Exception as e: - logger.error("Error switching agent: %s", e) - await send_typed( - ServerError( - error=f"Failed to switch to agent {agent_name}: {str(e)}", - session_id=session_id, - ) - ) - - elif msg.get("type") == "switch_model": - model_name = msg.get("model_name") or msg.get("model") - if model_name: - try: - await session_manager.switch_model( - session_id, model_name - ) - agent = ctx.agent # Refresh local alias - logger.debug("Switched model to: %s", model_name) - - # Persist to SQLite so the FE cold-load path sees it - _sw_agent = ctx.agent_name or agent_name or "code-puppy" - try: - await write_system_message_to_sqlite( - session_id=session_id, - system_message_type="config", - content=f"🔄 Switched to {_sw_agent} ({model_name})", - agent_name=_sw_agent, - model_name=model_name, - ) - except Exception as _sw_exc: - logger.warning( - "Model-switch SQLite write failed: %s", - _sw_exc, - exc_info=True, - ) - - await send_typed( - ServerSystem( - content=f"🔄 Switched to {_sw_agent} ({model_name})", - session_id=session_id, - model_name=model_name, - agent_name=_sw_agent, - ) - ) - except Exception as e: - logger.error("Error switching model: %s", e) - await send_typed( - ServerError( - error=f"Failed to switch to model {model_name}: {str(e)}", - session_id=session_id, - ) - ) - else: - await send_typed( - ServerError( - error="No model_name provided for switch_model", - session_id=session_id, - ) - ) - - elif msg.get("type") == "switch_session": - """ - Switch to a different session without reconnecting WebSocket. - - How it works: - 1. Client sends: {"type": "switch_session", "session_id": "WS_session_20260120_124356"} - 2. Server loads session from ~/.code_puppy/ws_sessions/{session_id}.pkl - 3. Server restores message history to the current agent - 4. Server responds with session_switched event containing message count and title - 5. Client can continue chatting with full conversation context - - This approach keeps a single WebSocket connection alive while allowing - users to switch between multiple sessions instantly. All session data - stays on the local filesystem for privacy. - """ - # Cancel any active streaming first - if active_drain_task and not active_drain_task.done(): - logger.debug( - "Cancelling active streaming due to session switch" - ) - stop_draining.set() # Signal drain to stop - active_drain_task.cancel() - try: - await active_drain_task - except asyncio.CancelledError: - pass - stop_draining.clear() # Reset for next streaming - - 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=session_id, - ) - ) - continue - - # Validate session_id to prevent path traversal - try: - _validate_session_id(new_session_id) - except ValueError as e: - logger.warning( - f"Invalid session_id in switch_session: {new_session_id!r}" - ) - await send_typed( - ServerError( - error=f"Invalid session ID: {e}", - session_id=session_id, - ) - ) - continue - - logger.debug("Switching to session: %s", new_session_id) - - try: - # Check if target session exists in SQLite - try: - from code_puppy.api.db.queries import ( - session_exists as _se, - ) - - _target_exists = await _se(new_session_id) - except Exception: - _target_exists = False - - if not _target_exists: - # Session doesn't exist - create new one with this ID - logger.debug( - f"Session {new_session_id} not found, creating new" - ) - - # Save current session and mark inactive (keep in memory for 15 min) - try: - await session_manager.save_session(session_id) - except Exception: - pass - await session_manager.mark_session_inactive(session_id) - - session_id = new_session_id - session_title = "" - session_working_directory = "" - session_pinned = False - - ctx = await session_manager.create_session(session_id) - # Mark the new session as active - await session_manager.mark_session_active(session_id) - agent = ctx.agent - agent_name = ctx.agent_name - model_name = ctx.model_name - - await send_typed( - ServerSessionSwitched( - session_id=new_session_id, - message_count=0, - title="", - created=True, - agent_name=agent_name, - model_name=model_name, - ) - ) - continue - - # Load existing session via SessionManager - # Save current session and mark inactive (keep in memory for 15 min) - try: - await session_manager.save_session(session_id) - except Exception: - pass - await session_manager.mark_session_inactive(session_id) - - session_id = new_session_id - - # Try loading via SessionManager (checks in-memory first, then SQLite) - loaded_ctx = await session_manager.get_or_load_session( - session_id - ) - if loaded_ctx is not None: - ctx = loaded_ctx - session_title = ctx.title - session_working_directory = ctx.working_directory - session_pinned = ctx.pinned - else: - # Legacy session or HMAC missing — load metadata - # from JSON (safe) and create fresh context - new_title = "" - new_working_directory = "" - new_pinned = False - try: - from code_puppy.api.db.queries import ( - get_session_metadata as _gsm, - ) - - _sm = await _gsm(session_id) or {} - new_title = _sm.get("title", "") - new_working_directory = _sm.get( - "working_directory", "" - ) - new_pinned = bool(_sm.get("pinned", False)) - except Exception: - pass - - ctx = await session_manager.create_session(session_id) - ctx.title = new_title - ctx.working_directory = new_working_directory - ctx.pinned = new_pinned - session_title = new_title - session_working_directory = new_working_directory - session_pinned = new_pinned - - # Mark the new session as active - await session_manager.mark_session_active(session_id) - - agent = ctx.agent - agent_name = ctx.agent_name - model_name = ctx.model_name - message_count = len(ctx.agent.get_message_history() or []) - logger.debug( - f"Restored {message_count} messages to session agent" - ) - - await send_typed( - ServerSessionSwitched( - session_id=new_session_id, - message_count=message_count, - title=session_title, - working_directory=session_working_directory, - created=False, - agent_name=agent_name, - model_name=model_name, - ) - ) - - logger.debug( - f"Switched to session {new_session_id} with {message_count} messages" - ) - - except Exception as e: - logger.error("Error switching session: %s", e) - await send_typed( - ServerError( - error=f"Failed to switch to session {new_session_id}: {str(e)}", - session_id=session_id, - ) - ) - - elif msg.get("type") == "set_working_directory": - new_directory = msg.get("directory", "") - if new_directory: - # Expand ~ and resolve symlinks before validation - new_directory = str( - Path(new_directory).expanduser().resolve() - ) - logger.info( - "[CWD DEBUG] set_working_directory received: new=%r, current=%r, session=%s", - new_directory, - session_working_directory, - session_id, - ) - # Validate the directory exists - if Path(new_directory).is_dir(): - # Skip if directory hasn't actually changed (avoids duplicate banners on reload) - if new_directory == session_working_directory: - logger.info( - "[CWD DEBUG] Skipping unchanged directory: %r, session=%s", - new_directory, - session_id, - ) - await send_typed( - ServerWorkingDirectoryChanged( - directory=new_directory, - success=True, - session_id=session_id, - unchanged=True, - ) - ) - continue # Skip to next message, don't write banner - - session_working_directory = new_directory - logger.debug( - "Working directory set to: %s", - session_working_directory, - ) - - # NOTE: Do NOT append raw dict entries into agent - # message history here. The runtime expects typed - # ModelMessage objects; dict injection can corrupt - # subsequent turns and lead to result=None failures. - - # Persist directory banner to SQLite so FE cold-load sees it - try: - # Bug 4: use ctx.agent_name / ctx.model_name directly - # so this works even when ctx.agent is None. - _cwd_agent = ( - ctx.agent_name if ctx else None - ) or "code-puppy" - _cwd_model = ( - ctx.model_name if ctx else None - ) or "unknown" - _cwd_segs = 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 _gpn - - _pname = _gpn() or "puppy" - logger.info( - "[CWD DEBUG] Writing CWD banner to SQLite: path=%r, session=%s", - session_working_directory, - session_id, - ) - await write_system_message_to_sqlite( - session_id=session_id, - system_message_type="directory", - content=f"{_pname} is now at {_cwd_rel}", - system_message_path=session_working_directory, - agent_name=_cwd_agent, - model_name=_cwd_model, - ) - # Bug 5: also stamp working_directory on the sessions row - _now_cwd = datetime.datetime.now( - datetime.timezone.utc - ).isoformat() - await update_session_working_directory( - session_id=session_id, - working_directory=session_working_directory, - updated_at=_now_cwd, - ) - except Exception as _cwd_exc: - logger.warning( - "CWD SQLite write failed: %s", - _cwd_exc, - exc_info=True, - ) - - await send_typed( - ServerWorkingDirectoryChanged( - directory=session_working_directory, - success=True, - session_id=session_id, - ) - ) - else: - await send_typed( - ServerWorkingDirectoryChanged( - directory=new_directory, - success=False, - error="Directory does not exist", - session_id=session_id, - ) - ) - else: - await send_typed( - ServerError( - error="No directory provided for set_working_directory", - session_id=session_id, - ) - ) - - elif msg.get("type") == "update_session_meta": - # Update session metadata (pinned, title) — persisted to SQLite. - try: - # Update in-memory state (local vars AND SessionContext) - if "pinned" in msg and isinstance(msg["pinned"], bool): - session_pinned = msg["pinned"] - if ctx is not None: - ctx.pinned = session_pinned - if "title" in msg and isinstance(msg["title"], str): - session_title = msg["title"] - if ctx is not None: - ctx.title = session_title - - # Persist to SQLite (single source of truth). - # Use a targeted UPDATE — NOT upsert_session() — so that - # message_count, total_tokens, and other stats are - # never zeroed out by a rename/pin operation. - 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=session_id, - title=session_title, - pinned=session_pinned, - updated_at=_dt.now(_tz.utc).isoformat(), - ) - logger.debug( - f"Updated session meta in SQLite for {session_id}: pinned={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=session_id, - pinned=session_pinned, - title=session_title, - ) - ) - except Exception as e: - logger.error("Error updating session meta: %s", e) - await send_typed( - ServerError( - error=f"Failed to update session metadata: {str(e)}", - session_id=session_id, - ) - ) - - elif msg.get("type") == "get_config": - 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=session_id, - ) - ) - else: - await send_typed( - ServerError( - error="No key provided for get_config", - session_id=session_id, - ) - ) - - elif msg.get("type") == "set_config": - 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( - f"Config set: {config_key} = {config_value}" - ) - await send_typed( - ServerConfigValue( - key=config_key, - value=config_value, - success=True, - session_id=session_id, - ) - ) - except Exception as e: - await send_typed( - ServerError( - error=f"Failed to set config: {e}", - session_id=session_id, - ) - ) - else: - await send_typed( - ServerError( - error="No key provided for set_config", - session_id=session_id, - ) - ) - - # Handle slash command execution - elif msg.get("type") == "command": - command_str = msg.get("command", "") - logger.debug("Command requested: %s", command_str) - - try: - from io import StringIO - - from rich.console import Console - - from code_puppy.command_line.command_handler import ( - get_commands_help, - handle_command, - ) - - output = None - captured_messages = [] - - # Special handling for /help - we can get the text directly - cmd_name = ( - command_str.strip().lstrip("/").split()[0] - if command_str.strip() - else "" - ) - if cmd_name in ("help", "h"): - # Get help text directly - help_text = get_commands_help() - # Render to plain text - 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: - # For other commands, just execute and report success/failure - # The actual output goes to the message queue/terminal - result = handle_command(command_str) - # Some commands might return a string as output - if isinstance(result, str): - output = result - result = True - - # Send command result - await send_typed( - ServerCommandResult( - command=command_str, - success=result is True or result is not False, - output=output, - messages=captured_messages, - result=str(result) - if result and result is not True - else None, - session_id=session_id, - ) - ) - logger.debug( - f"Command executed: {command_str} -> success={result is True or result is not False}, output_len={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, - ) - ) - continue - - # Handle cancel/interrupt request - elif msg.get("type") == "cancel": - logger.debug( - "Cancel request received - stopping active streaming and agent task" - ) - - # Cancel any active streaming - if active_drain_task and not active_drain_task.done(): - stop_draining.set() # Signal drain to stop - active_drain_task.cancel() - try: - await active_drain_task - except asyncio.CancelledError: - pass - stop_draining.clear() # Reset for next streaming - logger.debug("Active streaming cancelled") - - # Cancel the in-flight agent run (run_with_mcp) if present - if active_agent_task and not active_agent_task.done(): - logger.debug( - "Cancelling active agent task due to user interrupt" - ) - active_agent_task.cancel() - try: - await active_agent_task - except asyncio.CancelledError: - logger.debug("Active agent task cancelled successfully") - active_agent_task = None - - # Send confirmation - await send_typed( - ServerStatus( - status="cancelled", - session_id=session_id, - ) - ) - continue - - elif msg.get("type") == "permission_response": - # Handle permission response from user - from code_puppy.api.permissions import ( - handle_permission_response, - ) - - request_id = msg.get("request_id") - approved = msg.get("approved", False) - - if request_id: - handled = handle_permission_response(request_id, approved) - if handled: - logger.debug( - f"[Permission] ✅ Handled response: {request_id} = {approved}" - ) - else: - logger.warning( - f"[Permission] ❌ Unknown request: {request_id}" - ) - else: - logger.error( - "[WebSocket] ❌ No request_id in permission_response!" - ) - continue - - elif msg.get("type") == "message": - # Set WebSocket context for permission requests - from code_puppy.api.permission_plugin import ( - set_suppress_emitter_tool_events, - set_websocket_context, - ) - - set_websocket_context(websocket, 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 - event_queue = None - collected_text = [] - stop_draining.clear() # Reset for this message - drain_task = None - active_parts: dict[ - int, dict - ] = {} # Track message parts by index - tool_id_aliases: dict[str, str] = {} - tool_group_ids: dict[ - str, str - ] = {} # tool_id -> tool_group_id mapping - b1_streaming_used = ( - False # Track if B1 streaming sent any content - ) - agent_error = None # Will hold exception if agent run fails - - async def drain_events_concurrent( - ready_event: asyncio.Event = None, - ): - """Background task to drain events and send structured messages in real-time.""" - nonlocal collected_text, b1_streaming_used, ws_closed - 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" - ) - current_tool_name = None # Track active tool name - current_tool_group_id: str | None = ( - None # Track current tool batch group ID - ) - - # Track pending tool calls for tool_result generation - # {tool_id: {tool_name, start_time, part_index}} - pending_tool_calls: dict[str, dict] = {} - - def resolve_pending_tool_id( - *, - 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 pending_tool_calls - ): - return tool_call_id - - if tool_call_id: - for ( - _pending_id, - _pending_info, - ) in 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 pending_tool_calls.items(): - if ( - _pending_info.get("tool_name") - == tool_name - ): - return _pending_id - - return None - - def resolve_tool_group_id( - *, - tool_id: str | None = None, - pending_info: dict | 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. - - Priority order: - 1) pending tool info - 2) tool_id -> group map - 3) explicit fallback passed by caller - 4) current in-flight batch group - 5) synthetic deterministic-ish fallback (last resort) - """ - nonlocal current_tool_group_id - - 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 = 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 = 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: - tool_group_ids[tool_id] = group_id - if pending_info is not None: - pending_info["tool_group_id"] = group_id - elif tool_id in pending_tool_calls: - pending_tool_calls[tool_id][ - "tool_group_id" - ] = group_id - - if current_tool_group_id is None: - current_tool_group_id = group_id - - return group_id - - event_count = 0 - first_iteration = True - while not stop_draining.is_set(): - # Exit if WebSocket is closed - if 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( - 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 event_queue.empty(): - try: - event = 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 - - # 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": - tool_name = event_data.get( - "tool_name", "unknown" - ) - tool_args = event_data.get( - "tool_args", {} - ) - tool_id = str(uuid.uuid4())[:8] - - # Generate tool group ID for this batch if not already set - if current_tool_group_id is None: - current_tool_group_id = ( - f"tg-{str(uuid.uuid4())[:8]}" - ) - - logger.debug( - "[ws] tool_call: %s", tool_name - ) - - # Track current tool name - current_tool_name = tool_name - - # Register in pending_tool_calls so - # tool_call_complete can echo back the - # same tool_id the frontend already knows. - pending_tool_calls[tool_id] = { - "tool_name": tool_name, - "start_time": time_module.time(), - "tool_group_id": 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=current_agent_name, - model_name=current_model_name, - tool_group_id=current_tool_group_id, - ) - ) - - elif event_type == "tool_call_complete": - 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 - ) - - # Clear current tool name after completion - current_tool_name = None - - # Match the tool_id we generated at - # tool_call_start so the frontend can - # correlate result → call by ID, not - # just by name (names are not unique). - matching_tool_id = next( - ( - tid - for tid, info in pending_tool_calls.items() - if info["tool_name"] - == tool_name - ), - None, - ) - matching_pending_info = ( - pending_tool_calls.get( - matching_tool_id - ) - if matching_tool_id - else None - ) - tool_group_id_for_result = resolve_tool_group_id( - tool_id=matching_tool_id, - pending_info=matching_pending_info, - fallback_group_id=current_tool_group_id, - tool_name=tool_name, - source="tool_call_complete", - ) - - if matching_tool_id: - 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=current_agent_name, - model_name=current_model_name, - tool_group_id=tool_group_id_for_result, - ) - ) - - 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", {} - ) - initial_content = "" - - if ( - hasattr(part_obj, "content") - and part_obj.content - ): - initial_content = ( - part_obj.content - ) - elif isinstance( - part_obj, dict - ) and part_obj.get("content"): - initial_content = part_obj.get( - "content", "" - ) - - if part_type in ( - "TextPart", - "ThinkingPart", - ): - msg_type = ( - "thinking" - if part_type - == "ThinkingPart" - else "text" - ) - - # When a TextPart starts, any pending tool calls have completed - # Send status-only tool_result during streaming to update UI - # The actual result data will be sent later from extraction code - if ( - pending_tool_calls - and part_type == "TextPart" - ): - for ( - tool_id, - tool_info, - ) in list( - pending_tool_calls.items() - ): - if tool_info.get( - "status_only_sent" - ): - continue - duration_ms = ( - time_module.time() - - tool_info[ - "start_time" - ] - ) * 1000 - logger.debug( - f"[WebSocket] Sending status-only tool_result for: {tool_info['tool_name']} (duration: {duration_ms:.1f}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=current_agent_name, - model_name=current_model_name, - tool_group_id=resolve_tool_group_id( - tool_id=tool_id, - pending_info=tool_info, - fallback_group_id=current_tool_group_id, - tool_name=tool_info.get( - "tool_name" - ), - source="status_only_result", - ), - ) - ) - tool_info[ - "status_only_sent" - ] = True - current_tool_name = None - - # Check if part already exists (created by early delta) - if part_index in active_parts: - # Just update the type, keep existing message_id - active_parts[part_index][ - "type" - ] = msg_type - message_id = active_parts[ - part_index - ]["id"] - # If there's initial content, add it to the existing accumulated content - if initial_content: - active_parts[ - part_index - ]["content"] = ( - initial_content - + active_parts[ - part_index - ]["content"] - ) - collected_text.insert( - 0, initial_content - ) # Prepend to collected text too - logger.debug( - f"[Stream Debug] Part already exists, reusing message_id={message_id}" - ) - # Don't send another start event - continue - - message_id = f"msg-{int(time_module.time() * 1000)}-{part_index}" - active_parts[part_index] = { - "id": message_id, - "type": msg_type, - "content": initial_content, # Start with initial content if present - } - - # If there's initial content, add it to collected_text - if initial_content: - collected_text.append( - initial_content - ) - - # New assistant output part indicates turn boundary for tool grouping - if part_index == 0: - 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=current_agent_name, - model_name=current_model_name, - tool_name=current_tool_name, - ).model_dump( - exclude_none=True - ) - ) - - # If there's initial content, send it as a delta immediately - 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=current_agent_name, - model_name=current_model_name, - tool_name=current_tool_name, - ) - await safe_send_json( - _delta.model_dump( - exclude_none=True - ) - ) - # Mark that B1 streaming was used (nonlocal already declared above) - b1_streaming_used = True - - elif part_type == "ToolCallPart": - # Extract tool info from the part - tool_name = "unknown" - tool_call_id = None - tool_args_str = "" - - # Extract tool info from part_obj (dict or object) - if hasattr( - part_obj, "tool_name" - ): - tool_name = ( - part_obj.tool_name - ) - elif isinstance(part_obj, dict): - tool_name = part_obj.get( - "tool_name", "unknown" - ) - - if hasattr( - part_obj, "tool_call_id" - ): - tool_call_id = ( - part_obj.tool_call_id - ) - elif isinstance(part_obj, dict): - tool_call_id = part_obj.get( - "tool_call_id" - ) - - if hasattr(part_obj, "args"): - tool_args_str = ( - part_obj.args or "" - ) - elif isinstance(part_obj, dict): - tool_args_str = ( - part_obj.get("args", "") - or "" - ) - - tool_id = ( - tool_call_id - or str(uuid.uuid4())[:8] - ) - - # Track this tool call part with args buffer for accumulation - # Don't send tool_call yet - wait for part_end when args are complete - 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, # Buffer for accumulating args_delta - "start_time": time_module.time(), # Track when tool started - } - - # Track current tool name - current_tool_name = tool_name - - logger.debug( - f"[WebSocket] ToolCallPart started: {tool_name} (id: {tool_id})" - ) - # tool_call event will be sent on part_end when args are complete - - elif part_type == "ToolReturnPart": - logger.info( - f"[WebSocket] ToolReturnPart detected! part_index={part_index}" - ) - # Extract tool result info from the part - tool_call_id = None - tool_content = None - - # Extract tool_call_id - if hasattr( - part_obj, "tool_call_id" - ): - tool_call_id = ( - part_obj.tool_call_id - ) - elif isinstance(part_obj, dict): - tool_call_id = part_obj.get( - "tool_call_id" - ) - - # Extract content (the tool result) - if hasattr(part_obj, "content"): - tool_content = ( - part_obj.content - ) - elif isinstance(part_obj, dict): - tool_content = part_obj.get( - "content" - ) - - # Try to serialize complex result objects - if ( - tool_content - and not isinstance( - tool_content, - ( - str, - dict, - list, - int, - float, - bool, - type(None), - ), - ) - ): - try: - import json - - # Try to convert to dict first (for Pydantic models) - if hasattr( - tool_content, - "model_dump", - ): - tool_content = tool_content.model_dump() - elif hasattr( - tool_content, "dict" - ): - tool_content = tool_content.dict() - elif hasattr( - tool_content, - "__dict__", - ): - tool_content = tool_content.__dict__ - else: - tool_content = str( - tool_content - ) - except Exception as e: - logger.debug( - f"[WebSocket] Could not serialize tool result: {e}" - ) - tool_content = str( - tool_content - ) - - # Find the corresponding pending tool call, store and SEND the result - _result_sent = False - if tool_call_id: - resolved_pending_id = resolve_pending_tool_id( - tool_call_id=tool_call_id - ) - if resolved_pending_id: - _result_sent = True - pending_info = 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( - tool_id=resolved_pending_id, - pending_info=pending_info, - fallback_group_id=current_tool_group_id, - tool_name=_tool_name, - source="tool_return_resolved", - ) - logger.info( - f"[WebSocket] ToolReturnPart: Sending result for {_tool_name} (id: {resolved_pending_id}, raw: {tool_call_id})" - ) - # Send the REAL tool_result with actual content - 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=current_agent_name - or "code-puppy", - model_name=current_model_name - or "unknown", - tool_group_id=_group_id, - ) - ) - else: - logger.warning( - f"[WebSocket] ToolReturnPart: Could not resolve tool_call_id {tool_call_id}, pending keys: {list(pending_tool_calls.keys())}" - ) - - # Fallback: proximity-based matching if no tool_call_id or resolution failed - if not _result_sent: - for ( - pending_id, - pending_info, - ) in sorted( - pending_tool_calls.items(), - key=lambda x: abs( - x[1].get( - "part_index", - 9999, - ) - - part_index - ), - ): - if ( - abs( - pending_info.get( - "part_index", - 9999, - ) - - part_index - ) - <= 3 - ): - # Close enough - assume it's the result for this tool call - 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( - tool_id=pending_id, - pending_info=pending_info, - fallback_group_id=current_tool_group_id, - tool_name=_tool_name, - source="tool_return_proximity", - ) - logger.info( - f"[WebSocket] ToolReturnPart: Sending result (by proximity) for {_tool_name} (id: {pending_id})" - ) - # Send the REAL tool_result with actual content - 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=current_agent_name - or "code-puppy", - model_name=current_model_name - or "unknown", - tool_group_id=_group_id, - ) - ) - _result_sent = True - break - - # Warn if we couldn't send the result - if not _result_sent: - logger.warning( - f"[WebSocket] ToolReturnPart: Could NOT send result! tool_call_id={tool_call_id}, pending_tool_calls={list(pending_tool_calls.keys())}, part_index={part_index}" - ) - - # Track this part for cleanup - active_parts[part_index] = { - "id": f"tool-return-{part_index}", - "type": "tool_return", - "tool_call_id": tool_call_id, - "content": tool_content, - } - - 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", {} - ) - - # Handle ToolCallPartDelta - accumulate args - if ( - delta_type - == "ToolCallPartDelta" - ): - args_delta = "" - if hasattr( - delta_obj, "args_delta" - ): - args_delta = ( - delta_obj.args_delta - or "" - ) - elif isinstance( - delta_obj, dict - ): - args_delta = ( - delta_obj.get( - "args_delta", "" - ) - or "" - ) - - if ( - args_delta - and part_index - in active_parts - ): - part_info = active_parts[ - part_index - ] - if ( - part_info.get("type") - == "tool_call" - ): - part_info[ - "args_buffer" - ] = ( - part_info.get( - "args_buffer", - "", - ) - + args_delta - ) - continue # Don't process as text delta - - # Extract content_delta (supports both direct and nested formats) - content_delta = inner_data.get( - "content_delta", "" - ) - if not content_delta: - if isinstance(delta_obj, dict): - content_delta = ( - delta_obj.get( - "content_delta", "" - ) - ) - - if content_delta: - collected_text.append( - content_delta - ) - - # Ensure we have an entry for this part (handles delta before start) - if ( - part_index - not in active_parts - ): - # Create entry with consistent message_id - message_id = f"msg-{int(time_module.time() * 1000)}-{part_index}" - active_parts[part_index] = { - "id": message_id, - "type": "text", # Default, will be updated by part_start if it arrives - "content": "", - } - # Send a start event so client creates the message - logger.debug( - f"[Stream Debug] Creating part on first delta: message_id={message_id}" - ) - if part_index == 0: - 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=current_agent_name, - model_name=current_model_name, - tool_name=current_tool_name, - ).model_dump( - exclude_none=True - ) - ) - - part_info = active_parts[ - part_index - ] - message_id = part_info["id"] - - if part_index in active_parts: - active_parts[part_index][ - "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=current_agent_name, - model_name=current_model_name, - tool_name=current_tool_name, - ) - await safe_send_json( - _delta.model_dump( - exclude_none=True - ) - ) - # Mark that B1 streaming was used - b1_streaming_used = True - - elif inner_type == "part_end": - part_index = inner_data.get( - "index", 0 - ) - part_info = active_parts.get( - part_index, {} - ) - part_type_info = part_info.get( - "type", "text" - ) - message_id = part_info.get( - "id", f"msg-{part_index}" - ) - full_content = part_info.get( - "content", "" - ) - - # Handle tool_call parts - send tool_call event now that args are complete - if part_type_info == "tool_call": - 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" - ) - ) - - # Parse args if they're a JSON string - try: - import json - - args_dict = ( - json.loads( - tool_args_str - ) - if tool_args_str - else {} - ) - except ( - json.JSONDecodeError, - TypeError, - ): - args_dict = {} - - logger.debug( - f"[WebSocket] Sending tool_call (args complete): {tool_name}" - ) - - # Generate tool group ID for this batch if not already set - if ( - current_tool_group_id - is None - ): - current_tool_group_id = f"tg-{str(uuid.uuid4())[:8]}" - - # Send tool_call now that we have complete args - 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=current_agent_name, - model_name=current_model_name, - tool_group_id=current_tool_group_id, - ) - ) - - # Track as pending tool call for later tool_result - 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, # Will be populated by ToolReturnPart - "tool_group_id": current_tool_group_id, - } - if raw_tool_call_id: - tool_id_aliases[ - raw_tool_call_id - ] = tool_id - # Also track tool_group_id for pre-stream_end extraction - if current_tool_group_id: - tool_group_ids[tool_id] = ( - current_tool_group_id - ) - - # Clear current tool name tracking - current_tool_name = None - else: - await safe_send_json( - ServerAssistantMessageEnd( - message_id=message_id, - part_index=part_index, - full_content=full_content, - timestamp=time_module.time(), - session_id=session_id, - agent_name=current_agent_name, - model_name=current_model_name, - tool_name=current_tool_name, - ).model_dump( - exclude_none=True - ) - ) - - if part_index in active_parts: - del active_parts[part_index] - - except Exception as send_err: - error_msg = str(send_err).lower() - if ( - "close message" in error_msg - or "closed" in error_msg - ): - 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 = event_queue.get_nowait() - event_type = event.get("type", "") - event_data = event.get("data", {}) - final_count += 1 - - if event_type == "stream_event": - inner_type = event_data.get( - "event_type", "" - ) - inner_data = event_data.get( - "event_data", {} - ) - if inner_type == "part_delta": - content_delta = inner_data.get( - "content_delta", "" - ) - if not content_delta: - delta_obj = inner_data.get( - "delta", {} - ) - if isinstance(delta_obj, dict): - content_delta = delta_obj.get( - "content_delta", "" - ) - if content_delta: - collected_text.append(content_delta) - 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}" - ) - - # Event to signal drain task is ready - drain_ready = asyncio.Event() - - try: - from code_puppy.plugins.frontend_emitter.emitter import ( - subscribe, - unsubscribe, - ) - - event_queue = subscribe(session_id=session_id) - logger.debug( - "Subscribed to frontend emitter for streaming" - ) - - # Modify drain function to signal when ready - async def drain_events_with_signal(): - """Wrapper that signals readiness before starting drain loop.""" - await drain_events_concurrent(drain_ready) - - # Start concurrent drain task - drain_task = asyncio.create_task( - drain_events_with_signal() - ) - active_drain_task = ( - drain_task # Track for potential cancellation - ) - - # Wait for drain task to be ready before proceeding - await drain_ready.wait() - - except ImportError: - logger.warning("Frontend emitter not available") - - 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, - ) - - # Inject working directory as a system message in the - # agent's conversation history (once per directory change) - # so the LLM knows the CWD without polluting user messages. - message_to_send = user_message - if ( - session_working_directory - and session_working_directory - != last_context_sent_directory - ): - from pydantic_ai.messages import ( - ModelRequest, - SystemPromptPart, - ) - - wd_system_msg = ModelRequest( - parts=[ - SystemPromptPart( - content=( - f"The user's current working directory is updated to" - f" {session_working_directory}" - ) - ) - ] - ) - agent.append_to_message_history(wd_system_msg) - last_context_sent_directory = ( - session_working_directory - ) - logger.debug( - "Injected working directory system message: %s", - session_working_directory, - ) - - logger.debug( - f"Calling run_with_mcp with message: {message_to_send[:100]}..." - ) - - # Build file context and binary attachments - file_context, binary_attachments = ( - build_file_context_and_attachments(msg) - ) - - # CHANGE 1: Update attachment metadata for UI (variables already initialized) - if msg.get("attachments"): - from pathlib import Path as _AttachmentPath - - for raw_path in msg.get("attachments", []): - if ( - isinstance(raw_path, str) - and raw_path.strip() - ): - try: - file_path = _AttachmentPath( - 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( - f"Error building attachment metadata for '{raw_path}': {e}" - ) - - # Prepend file context to the message - if file_context: - message_to_send = ( - file_context + "\n\n" + message_to_send - ) - logger.debug( - f"Added file context ({len(file_context)} chars)" - ) - - run_kwargs = {} - if binary_attachments: - run_kwargs["attachments"] = binary_attachments - logger.debug( - f"Including {len(binary_attachments)} binary attachment(s)" - ) - - # ────────────────────────────────────────────────────────────────── - # 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, - ) - - # Run the agent in its own task so it can be cancelled independently. - # Suppress duplicate emitter lifecycle callbacks only during the - # run_with_mcp execution window (pre/post_tool_call hooks fire there). - set_suppress_emitter_tool_events(True) - active_agent_task = asyncio.create_task( - agent.run_with_mcp( - message_to_send, **run_kwargs - ) - ) - - # ===== CONCURRENT MESSAGE PROCESSING FIX ===== - # Instead of awaiting the agent task here (which blocks the receive loop), - # we use asyncio.wait() to handle BOTH the agent task AND incoming messages. - # This allows permission_response messages to be processed while the agent runs. - - result = None - agent_completed = False - # Deferred message for switch_session/create_session during streaming - _deferred_msg: dict | None = None - - while not agent_completed: - # Create a task for the next message (or wait for agent) - receive_task = asyncio.create_task( - websocket.receive_json() - ) - - # Wait for either the agent to complete OR a new message to arrive - done, pending = await asyncio.wait( - {active_agent_task, receive_task}, - return_when=asyncio.FIRST_COMPLETED, - ) - - # Check if agent completed - if active_agent_task in done: - try: - result = await active_agent_task - logger.debug( - f"run_with_mcp completed, result type: {type(result)}" - ) - agent_completed = True - active_agent_task = None - except asyncio.CancelledError: - logger.debug( - "run_with_mcp task was cancelled by user" - ) - 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, - ) - agent_error = e - result = None - agent_completed = True - active_agent_task = None - - # Cancel the receive task if it's still pending - if receive_task in pending: - receive_task.cancel() - try: - await receive_task - except asyncio.CancelledError: - pass - - # Check if a new message arrived - elif receive_task in done: - try: - new_msg = await receive_task - - # Advisory validation — log but never reject - try: - _parsed_inner = _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" - }, - ) - - # Handle permission_response immediately - if ( - 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 - ) - ) - if handled: - logger.debug( - f"[Permission] ✅ Handled response: {request_id} = {approved}" - ) - else: - logger.warning( - f"[Permission] ❌ Unknown request: {request_id}" - ) - else: - logger.error( - "[WebSocket] ❌ No request_id in permission_response!" - ) - - # Handle cancel request - 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 # Exit the loop - ) - - # Handle session switch during streaming - elif new_msg.get("type") in ( - "switch_session", - "create_session", - ): - # User switched to a new chat while agent was running. - # The agent will finish in background and save to SQLite. - logger.debug( - "[WS:%s] Session switch during streaming — " - "agent continues in background, switching to: %s", - session_id, - new_msg.get("session_id"), - ) - - # Fire and forget — background task owns the agent result. - 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 # disown — background task now owns it - agent_completed = ( - True # exit inner loop - ) - - # For other message types, log and ignore (or queue for later) - else: - logger.warning( - f"[WebSocket] Received {new_msg.get('type')} message while agent running - ignoring" - ) - - except asyncio.CancelledError: - # Receive was cancelled, continue - pass - except WebSocketDisconnect: - # WebSocket gone — let agent finish and save to SQLite - logger.debug( - "[WS:%s] Disconnect during streaming — agent continues in background", - session_id, - ) - - # Fire and forget — background task owns the agent result. - 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(): - # Starlette raises RuntimeError when calling receive after disconnect - logger.debug( - "[WS:%s] WebSocket already disconnected: %s — agent continues in background", - session_id, - e, - ) - - # Fire and forget — background task owns the agent result. - 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( - f"RuntimeError processing message during agent execution: {e}" - ) - except Exception as e: - logger.error( - f"Error processing message during agent execution: {e}" - ) - - # ===== END CONCURRENT MESSAGE PROCESSING FIX ===== - - finally: - # End suppression scope once run_with_mcp window ends, - # including completion/cancel/background handoff paths. - set_suppress_emitter_tool_events(False) - # Clear session context for prompt generation - clear_session_working_directory() - - finally: - # Stop the drain task - if drain_task: - stop_draining.set() - try: - await asyncio.wait_for(drain_task, timeout=2.0) - except asyncio.TimeoutError: - drain_task.cancel() - try: - await drain_task - except asyncio.CancelledError: - pass - - # Unsubscribe from emitter - if event_queue: - unsubscribe(event_queue) - logger.debug("Unsubscribed from frontend emitter") - - # 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 - - # If agent errored, send error to GUI and skip success path - if agent_error == "cancelled": - await send_typed( - ServerCancelled( - session_id=session_id, - ) - ) - continue - elif agent_error is not None: - logger.debug( - "[WS:%s] agent_error -> sending frame(s) to client. type=%s", - session_id, - type(agent_error).__name__, - ) - # If streaming was already in progress, send stream_end first so - # the frontend exits its streaming state before the error frame. - # Without this handshake the UI hangs indefinitely. - for frame in build_error_response_frames( - agent_error, collected_text, session_id - ): - await safe_send_json(frame) - continue - - # Safety net: if the agent task completed without raising but returned - # no result and produced no streamed text, treat it as an error. - has_nonempty_stream = any( - (chunk or "").strip() for chunk in 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(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." - ) - ) - await send_typed( - 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, - ) - ) - continue - - # Extract final response text - response_text = "" - - # Priority 1: Use collected text from streaming events - if has_nonempty_stream: - response_text = "".join(collected_text) - logger.debug( - f"Using collected streaming text ({len(response_text)} chars)" - ) - # Priority 2: Extract from result.output (AgentRunResult) or result.data (legacy) - 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)" - ) - # Priority 3: Get last assistant message from history - elif agent: - messages = agent.get_message_history() - - for msg in reversed(messages): - if hasattr(msg, "role") and msg.role == "assistant": - if 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" - - # Extract token usage from result if available - tokens_used = None - if result: - # Try to get usage from 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 - ), - } - # Also try _usage or other common patterns - 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 we couldn't get usage from result, estimate from message history - 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 - - # Only send legacy 'response' if B1 streaming wasn't used - # B1 streaming already sent the content via assistant_message_end - logger.warning( - "[WebSocket] b1_streaming_used=%s before response/extraction", - b1_streaming_used, - ) - if not b1_streaming_used: - # For non-streaming models (like Gemini), extract and send thinking content first - thinking_text = "" - if agent: - try: - history = agent.get_message_history() - logger.debug( - f"[Thinking Debug] Checking history for thinking parts, {len(history) if history else 0} messages" - ) - if history: - # Look for ThinkingPart in the last ModelResponse - for i, msg in enumerate(reversed(history)): - msg_type = type(msg).__name__ - logger.debug( - f"[Thinking Debug] Message {i}: type={msg_type}, has_parts={hasattr(msg, 'parts')}" - ) - if "Response" in msg_type and hasattr( - msg, "parts" - ): - logger.debug( - f"[Thinking Debug] Found Response with {len(msg.parts)} 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( - f"[Thinking Debug] Part {j}: type={part_type}, content_preview={part_content_preview}" - ) - if ( - "Thinking" in part_type - and hasattr(part, "content") - ): - thinking_text = part.content - logger.debug( - f"[Thinking Debug] Found thinking content: {len(thinking_text)} chars" - ) - 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" - ) - - # 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" - ) - - # Send the main response - await send_typed( - ServerResponse( - content=response_text, - done=True, - 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, - ) - ) - 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: set = set() - logger.warning( - "[WebSocket] B1 Pre-stream_end extraction: result=%s, has_all_messages=%s", - type(result).__name__ if result else None, - hasattr(result, "all_messages") - if result - else False, - ) - if result and hasattr(result, "all_messages"): - try: - import time as _time_pre - - from pydantic_ai.messages import ( - ToolReturn, - ToolReturnPart, - ) - - _pre_msgs = list(result.all_messages()) - for _pre_msg in _pre_msgs: - if not hasattr(_pre_msg, "parts"): - continue - for _pre_part in _pre_msg.parts: - if not isinstance( - _pre_part, - (ToolReturnPart, ToolReturn), - ): - continue - _pre_tool_name = getattr( - _pre_part, "tool_name", "unknown" - ) - _pre_raw_tool_id = getattr( - _pre_part, "tool_call_id", None - ) - _pre_tool_id = ( - tool_id_aliases.get( - _pre_raw_tool_id, - _pre_raw_tool_id, - ) - if _pre_raw_tool_id - else "unknown" - ) - _pre_result = getattr( - _pre_part, "content", None - ) - # Serialize complex objects - if _pre_result and not isinstance( - _pre_result, - ( - str, - dict, - list, - int, - float, - bool, - type(None), - ), - ): - try: - if hasattr( - _pre_result, "model_dump" - ): - _pre_result = ( - _pre_result.model_dump() - ) - elif hasattr( - _pre_result, "dict" - ): - _pre_result = ( - _pre_result.dict() - ) - elif hasattr( - _pre_result, "__dict__" - ): - _pre_result = ( - _pre_result.__dict__ - ) - else: - _pre_result = str( - _pre_result - ) - except Exception: - _pre_result = str(_pre_result) - logger.warning( - "[WebSocket] Pre-stream_end tool result: %s (id: %s), content_preview=%s", - _pre_tool_name, - _pre_tool_id, - str(_pre_result)[:100] - if _pre_result - else None, - ) - _pre_group_id = tool_group_ids.get( - _pre_tool_id - ) - - await send_typed_tool_lifecycle( - ServerToolResult( - tool_id=_pre_tool_id, - tool_name=_pre_tool_name, - result=_pre_result, - success=True, - duration_ms=0, - timestamp=_time_pre.time(), - session_id=session_id, - agent_name=agent.name - if agent - else "code-puppy", - model_name=agent.get_model_name() - if agent - else "unknown", - tool_group_id=_pre_group_id, - ) - ) - _pre_sent_tool_ids.add(_pre_tool_id) - except Exception as _pre_e: - logger.warning( - "Pre-stream_end tool result extraction failed: %s", - _pre_e, - ) - - # 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 - # Track tool IDs already sent before stream_end to skip duplicates later - if "_pre_sent_tool_ids" not in dir(): - _pre_sent_tool_ids: set = set() - _save_history_snapshot = [] # Snapshot of history before await points - try: - # CRITICAL: Update message history from result to include final response - # The pydantic-ai result.all_messages() contains the complete conversation - # including the final assistant response that may not be in agent's history yet - if result and hasattr(result, "all_messages"): - try: - all_msgs = list(result.all_messages()) - if all_msgs: - agent.set_message_history(all_msgs) - _save_history_snapshot = list( - agent.get_message_history() - ) # Snapshot before any awaits can corrupt shared global state - logger.debug( - f"Updated message history from result.all_messages(): {len(all_msgs)} messages" - ) - - # Extract and send tool results from message history - try: - import time as _time - - from pydantic_ai.messages import ( - ToolReturn, - ToolReturnPart, - ) - - for msg in all_msgs: - if hasattr(msg, "parts"): - for part in msg.parts: - if isinstance( - part, - ( - ToolReturnPart, - ToolReturn, - ), - ): - tool_name = getattr( - part, - "tool_name", - "unknown", - ) - tool_call_id = getattr( - part, - "tool_call_id", - "unknown", - ) - # Serialize the result - result_data = getattr( - part, "content", None - ) - if ( - result_data - and not isinstance( - result_data, - ( - str, - dict, - list, - int, - float, - bool, - type(None), - ), - ) - ): - try: - if hasattr( - result_data, - "model_dump", - ): - result_data = result_data.model_dump() - elif hasattr( - result_data, - "dict", - ): - result_data = result_data.dict() - elif hasattr( - result_data, - "__dict__", - ): - result_data = result_data.__dict__ - else: - result_data = str( - result_data - ) - except Exception: - result_data = str( - result_data - ) - - # Log for shell commands - if ( - tool_name - == "agent_run_shell_command" - ): - stdout_val = ( - result_data.get( - "stdout", "N/A" - ) - if isinstance( - result_data, - dict, - ) - else "not dict" - ) - logger.debug( - "Extracted shell result: id=%s, stdout=%s", - tool_call_id, - stdout_val, - ) - - # Skip tool IDs already sent before stream_end - if ( - tool_call_id - in _pre_sent_tool_ids - ): - logger.debug( - "[WebSocket] Skipping duplicate tool result (pre-sent): %s", - tool_call_id, - ) - continue - logger.info( - f"[WebSocket] Sending extracted tool result for {tool_name} (id: {tool_call_id})" - ) - _post_group_id = ( - tool_group_ids.get( - tool_call_id - ) - ) - - await send_typed( - ServerToolResult( - tool_id=tool_call_id, - tool_name=tool_name, - result=result_data, - success=True, - duration_ms=0, - timestamp=_time.time(), - session_id=session_id, - agent_name=agent.name - if agent - else "code-puppy", - model_name=agent.get_model_name() - if agent - else "unknown", - tool_group_id=_post_group_id, - ) - ) - except Exception as e: - logger.warning( - f"Could not extract tool results from messages: {e}" - ) - - except Exception as e: - logger.warning( - f"Could not update history from result.all_messages(): {e}" - ) - - history = ( - _save_history_snapshot - if _save_history_snapshot - else agent.get_message_history() - ) # Use pre-await snapshot to avoid race condition - if history: - # Regenerate title if it's empty or still the default "untitled-session" - if ( - not session_title - or session_title == "untitled-session" - ): - session_title = generate_heuristic_title( - history - ) - - # Session name is just the ID (WS_session_timestamp format) - # Title is stored only in the meta file, not in the filename - session_name = session_id - - # Wrap each message with metadata for complete session information. - # Chain `or` fallbacks: agent.name can be "" before the agent - # object is fully configured, so we cascade to context defaults. - agent_name_meta = ( - (agent.name if agent else "") - or ctx.agent_name - or "code-puppy" - ) - model_name_meta = ( - (agent.get_model_name() if agent else "") - or ctx.model_name - or "unknown" - ) - - def _extract_message_timestamp( - raw_msg: Any, default_ts: str - ) -> str: - """Best-effort extraction of an existing timestamp for a message. - - This is important for WebSocket sessions where history may - already contain older messages. We don't want to overwrite - their original timestamps every time we auto-save. - - Precedence: - 1. If the message is a dict with a numeric epoch 'timestamp', - convert to ISO. - 2. If the message is a dict with an ISO-ish 'timestamp' str, - reuse it as-is. - 3. If the message is a dict with 'ts', reuse it. - 4. If the message object has a 'timestamp' attribute, try that. - 5. Fall back to the provided default_ts ("now" for new messages). - """ - # Dict-based histories (e.g. CLI format or older WS formats) - if isinstance(raw_msg, dict): - ts_val = raw_msg.get("timestamp") - - # Epoch seconds - if isinstance(ts_val, (int, float)): - try: - return ( - datetime.datetime.fromtimestamp( - ts_val - ).isoformat() - ) - except Exception: - pass - - # Already an ISO string - if isinstance(ts_val, str) and ts_val: - return ts_val - - # Our enhanced WS wrapper sometimes uses 'ts' - ts_field = raw_msg.get("ts") - if isinstance(ts_field, str) and ts_field: - return ts_field - - # Pydantic / custom objects that carry a timestamp attribute - 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 - - # Fallback: use provided default - return default_ts - - # Create enhanced history: list of dicts with message + metadata - # Each entry: {'msg': , 'agent': str, 'model': str, 'ts': str} - enhanced_history = [] - for idx, msg in enumerate(history): - # Check if already wrapped (for idempotency). If the wrapper - # already has a 'ts', leave it untouched so older sessions - # keep their original per-message timestamps. - if ( - isinstance(msg, dict) - and "msg" in msg - and "agent" in msg - ): - enhanced_history.append(msg) - else: - # Use existing timestamp if present; otherwise compute a default now() - current_timestamp = ( - datetime.datetime.now().isoformat() - ) - msg_ts = _extract_message_timestamp( - msg, current_timestamp - ) - wrapper = { - "msg": msg, - "agent": agent_name_meta, - "model": model_name_meta, - "ts": msg_ts, - } - - # Add clean_content and attachments to the user message we just processed. - # clean_content is only needed when attachments were injected into content - # (file blocks like --- File: auth.ts ---). Without attachments, content - # is already the user's words (FE strips [Session Context:] as fallback). - is_user_message_just_processed = ( - idx == len(history) - 2 - and len(history) >= 2 - and attachment_metadata # only needed when file content was injected - ) - - if is_user_message_just_processed: - wrapper["clean_content"] = ( - original_user_message - ) - wrapper["attachments"] = ( - attachment_metadata - ) - logger.debug( - "Added UI metadata to user message: %d attachment(s), " - "clean_content length: %d", - len(attachment_metadata), - len(original_user_message), - ) - - enhanced_history.append(wrapper) - - message_count = len(enhanced_history) - 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 - - # Write to SQLite for FE read path - try: - import datetime as _dt_mod - - _now_iso = _dt_mod.datetime.now( - _dt_mod.timezone.utc - ).isoformat() - await write_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, - updated_at=_now_iso, - created_at=ctx.created_at.isoformat(), - ctx=ctx, - ) - except Exception as _db_exc: - logger.debug( - "SQLite turn write skipped (DB not available): %s", - _db_exc, - ) - - # Send session metadata update to client - await websocket.send_json( - { - "type": "session_meta", - "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, - } - ) - - # Broadcast session update to session monitoring clients - session_update_data = { - "session_id": session_id, - "session_name": session_name, - "title": session_title, - "working_directory": session_working_directory, - "timestamp": datetime.datetime.now().isoformat(), - "message_count": message_count, - "total_tokens": total_tokens, - "auto_saved": True, - "pickle_path": "", - "metadata_path": "", - "action": "created" - if message_count == 1 - else "updated", - } - await connection_manager.broadcast_session_update( - session_update_data - ) - 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) - - 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 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: - 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) - except Exception: - pass diff --git a/code_puppy/api/ws/legacy/connection_manager.py b/code_puppy/api/ws/legacy/connection_manager.py deleted file mode 100644 index eec49804d..000000000 --- a/code_puppy/api/ws/legacy/connection_manager.py +++ /dev/null @@ -1,55 +0,0 @@ -"""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/legacy/events_handler.py b/code_puppy/api/ws/legacy/events_handler.py deleted file mode 100644 index b9bf647ec..000000000 --- a/code_puppy/api/ws/legacy/events_handler.py +++ /dev/null @@ -1,53 +0,0 @@ -"""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/legacy/health_handler.py b/code_puppy/api/ws/legacy/health_handler.py deleted file mode 100644 index 832f45e95..000000000 --- a/code_puppy/api/ws/legacy/health_handler.py +++ /dev/null @@ -1,22 +0,0 @@ -"""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/legacy/runtime/__init__.py b/code_puppy/api/ws/legacy/runtime/__init__.py deleted file mode 100644 index 750f74d72..000000000 --- a/code_puppy/api/ws/legacy/runtime/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -"""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/legacy/runtime/session_runtime.py b/code_puppy/api/ws/legacy/runtime/session_runtime.py deleted file mode 100644 index a6fda2741..000000000 --- a/code_puppy/api/ws/legacy/runtime/session_runtime.py +++ /dev/null @@ -1,68 +0,0 @@ -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/legacy/runtime/session_runtime_manager.py b/code_puppy/api/ws/legacy/runtime/session_runtime_manager.py deleted file mode 100644 index 82a8252d3..000000000 --- a/code_puppy/api/ws/legacy/runtime/session_runtime_manager.py +++ /dev/null @@ -1,150 +0,0 @@ -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/legacy/schemas.py b/code_puppy/api/ws/legacy/schemas.py deleted file mode 100644 index 05e56528d..000000000 --- a/code_puppy/api/ws/legacy/schemas.py +++ /dev/null @@ -1,603 +0,0 @@ -"""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, 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" - - -class ServerMessageType(str, Enum): - """All ``type`` values the server may send over the WebSocket.""" - - SYSTEM = "system" - 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" - 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 - - -# ═══════════════════════════════════════════════════════════════════════════════ -# SERVER → CLIENT (21 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. - """ - - type: Literal["assistant_message_end"] = "assistant_message_end" - message_id: str - 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 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")], - ], - 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[ServerCancelled, Tag("cancelled")], - ], - Discriminator("type"), -] -"""Discriminated union of all server→client message types.""" diff --git a/code_puppy/api/ws/legacy/sessions_handler.py b/code_puppy/api/ws/legacy/sessions_handler.py deleted file mode 100644 index f0719648c..000000000 --- a/code_puppy/api/ws/legacy/sessions_handler.py +++ /dev/null @@ -1,69 +0,0 @@ -"""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/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/docs/migration/puppy-desk-cleanup-summary.md b/docs/migration/puppy-desk-cleanup-summary.md new file mode 100644 index 000000000..d8fba5c79 --- /dev/null +++ b/docs/migration/puppy-desk-cleanup-summary.md @@ -0,0 +1,41 @@ +# Puppy Desk Migration — Cleanup Summary + +## Current state + +- Legacy puppy-desk migration behavior is considered working. +- Cleanup approval was given to remove temporary migration scaffolding. +- Active WebSocket route remains `/ws/chat`. + +## Cleanup goals + +- Remove `code_puppy/api/ws/legacy/` snapshot now that rollback-by-copy is no longer required. +- Reduce `chat_handler.py` size by extracting reusable protocol/error frame helpers. +- Preserve existing wire behavior with focused regression tests. + +## This cleanup pass + +- Deleted the temporary `code_puppy/api/ws/legacy/` reference snapshot. +- Extracted response/error frame helpers into `code_puppy/api/ws/response_frames.py`. +- Kept `/ws/chat` registration unchanged. +- Replaced Gate 2 snapshot assertions with post-cleanup route/namespace assertions. + +## Follow-up modularization targets + +Remaining large modules that should be split in later passes: + +- `code_puppy/api/ws/chat_handler.py` +- `code_puppy/api/session_context.py` +- `code_puppy/api/db/queries.py` + +Recommended next split inside `chat_handler.py`: + +- drain/event lifecycle helpers +- session switch/create helpers +- persistence helpers +- typed send helpers + +## Phase 2 checkpoint + +- Extracted history wrapping/timestamp/token estimation helpers into `code_puppy/api/ws/history_utils.py`. +- Reduced inline persistence-preparation logic in `code_puppy/api/ws/chat_handler.py`. +- Added focused unit tests for history wrapping, timestamp extraction, attachment metadata backfill, and safe token estimation failure handling. diff --git a/docs/migration/puppy-desk-phase3-plan.md b/docs/migration/puppy-desk-phase3-plan.md new file mode 100644 index 000000000..6aa5ce00d --- /dev/null +++ b/docs/migration/puppy-desk-phase3-plan.md @@ -0,0 +1,127 @@ +# Puppy Desk Migration — Phase 3 Cleanup Plan + +## Objective + +Continue reducing `code_puppy/api/ws/chat_handler.py` without changing the live +`/ws/chat` wire contract or the already-working puppy-desk migration behavior. +This phase should focus on extracting cohesive helper modules that are easier to +unit test and safer to evolve. + +## Current baseline + +- Active route remains `/ws/chat` +- Legacy snapshot under `code_puppy/api/ws/legacy/` has been removed +- `code_puppy/api/ws/__init__.py` uses lazy imports +- Response-frame helpers live in `code_puppy/api/ws/response_frames.py` +- History wrapping/timestamp/token helpers live in `code_puppy/api/ws/history_utils.py` +- `chat_handler.py` is still large and contains multiple embedded responsibilities + +## Phase 3 scope + +### Primary extraction seam + +Extract session-save and client-broadcast preparation logic from +`chat_handler.py` into a focused helper module. + +Candidate module: + +- `code_puppy/api/ws/session_persistence.py` + +Candidate responsibilities: + +- build session meta payloads sent back to the client +- build session update payloads broadcast to session-monitoring clients +- normalize autosave timestamps/metadata fields +- centralize the SQLite turn-write call shape used by `/ws/chat` + +### Secondary extraction seam + +If the primary seam is clean and low risk, extract typed send/error helpers into +another focused helper module. + +Candidate module: + +- `code_puppy/api/ws/send_utils.py` + +Candidate responsibilities: + +- safe websocket JSON send wrapper +- structured error persistence wrapper +- typed message send adapters +- consistent closed-socket detection and logging + +## Explicit non-goals for phase 3 + +- No protocol redesign +- No frontend emitter behavior changes +- No changes to active route names or namespace layout +- No migration of logic into `code_puppy/command_line/` +- No DB schema changes unless a separate approval is given + +## Proposed execution order + +1. Identify the exact persistence/broadcast block inside `chat_handler.py` +2. Extract pure payload-building helpers first +3. Add focused unit tests for payload assembly and fallback handling +4. Move the SQLite turn-write call wrapper if the interface remains stable +5. Re-run focused WS regression tests +6. If still low risk, extract send/error helpers in a follow-up patch set + +## Validation strategy + +### Automated + +Minimum focused test set: + +- `code_puppy/api/tests/test_streaming_protocol_transform.py` +- `code_puppy/api/tests/test_ws_history_utils.py` +- new tests for session meta/update payload building +- `code_puppy/api/tests/test_gate2_legacy_ws_namespace.py` + +Add coverage for: + +- session meta payload fields +- broadcast payload action selection (`created` vs `updated`) +- persistence fallback behavior when token estimation or DB writes fail +- no eager import regressions from new helper modules + +### Manual / GUGI validation + +After phase 3 code lands: + +1. Start the backend in this repo +2. Connect the GUI/GUGI app to the migrated `/ws/chat` backend +3. Verify: + - session creation + - resumed session loading + - streaming assistant output + - tool lifecycle updates + - session metadata updates in the UI + - no regressions in frontend emitter behavior +4. Confirm session list updates still appear correctly after autosave + +## Risks + +### Low-risk + +- extracting pure payload builders +- moving timestamp/metadata assembly helpers + +### Medium-risk + +- extracting send/error helpers that close over websocket state +- moving persistence wrappers that depend on runtime context + +### Deferred/high-risk + +- splitting `drain_events_concurrent` +- deeper session lifecycle/state machine changes +- DB query modularization in the same patch + +## Done criteria for phase 3 + +- `chat_handler.py` shrinks further with no route/namespace changes +- extracted modules are independently importable and tested +- focused regression suite passes +- manual GUI/GUGI smoke test checklist is documented and run +- cleanup summary doc updated with a phase 3 checkpoint From fd02dac1f3196932d38c9bd445deabde0e9a4adc Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Tue, 2 Jun 2026 16:01:45 -0500 Subject: [PATCH 13/29] phase3: extract session_persistence.py and send_utils.py from chat_handler --- code_puppy/api/tests/test_send_utils.py | 189 ++++++++++++++++ .../api/tests/test_session_persistence.py | 177 +++++++++++++++ code_puppy/api/ws/chat_handler.py | 209 +++++------------- code_puppy/api/ws/send_utils.py | 151 +++++++++++++ code_puppy/api/ws/session_persistence.py | 143 ++++++++++++ docs/migration/puppy-desk-cleanup-summary.md | 50 +++++ 6 files changed, 764 insertions(+), 155 deletions(-) create mode 100644 code_puppy/api/tests/test_send_utils.py create mode 100644 code_puppy/api/tests/test_session_persistence.py create mode 100644 code_puppy/api/ws/send_utils.py create mode 100644 code_puppy/api/ws/session_persistence.py 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..9d40a4e71 --- /dev/null +++ b/code_puppy/api/tests/test_send_utils.py @@ -0,0 +1,189 @@ +"""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_persistence.py b/code_puppy/api/tests/test_session_persistence.py new file mode 100644 index 000000000..f63d04873 --- /dev/null +++ b/code_puppy/api/tests/test_session_persistence.py @@ -0,0 +1,177 @@ +"""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["pickle_path"] == "" + assert payload["metadata_path"] == "" + 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", + "pickle_path", + "metadata_path", + "action", + } + assert set(payload.keys()) == expected_keys diff --git a/code_puppy/api/ws/chat_handler.py b/code_puppy/api/ws/chat_handler.py index f22dc4280..195c5bf06 100644 --- a/code_puppy/api/ws/chat_handler.py +++ b/code_puppy/api/ws/chat_handler.py @@ -22,16 +22,13 @@ import re import uuid from pathlib import Path -from typing import Any from fastapi import FastAPI, WebSocket, WebSocketDisconnect from pydantic import TypeAdapter, ValidationError from code_puppy.api.db.queries import ( update_session_working_directory, - write_error_message_to_sqlite, write_system_message_to_sqlite, - write_turn_to_sqlite, ) from code_puppy.api.session_context import _validate_session_id, session_manager from code_puppy.api.ws.attachments import build_file_context_and_attachments @@ -49,6 +46,13 @@ build_enhanced_history, estimate_total_tokens, ) +from code_puppy.api.ws.send_utils import WebSocketSender +from code_puppy.api.ws.session_persistence import ( + build_session_meta_payload, + build_session_update_payload, + persist_turn_to_sqlite, + resolve_agent_model_meta, +) from code_puppy.api.ws.schemas import ( PROTOCOL_VERSION, ClientMessage, @@ -60,7 +64,6 @@ ServerCommandResult, ServerConfigValue, ServerError, - ServerMessage, ServerSessionMetaUpdated, ServerSessionRestored, ServerSessionSwitched, @@ -147,93 +150,15 @@ async def websocket_chat( "Chat WebSocket client connected (session_id param: %s)", session_id ) - # Flag to track if WebSocket is still open - ws_closed = False - - async def persist_error_payload(data: dict[str, Any]) -> None: - """Persist structured error frames so they survive reloads.""" - if data.get("type") != "error" or not session_id: - return - - try: - await write_error_message_to_sqlite( - session_id=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=(ctx.agent_name if ctx else ""), - model_name=(ctx.model_name if ctx else ""), - timestamp=datetime.datetime.now(datetime.timezone.utc).isoformat(), - ) - except Exception: - logger.warning( - "[WS:%s] Failed to persist error payload to SQLite", - session_id, - exc_info=True, - ) - - async def safe_send_json(data: dict) -> bool: - """Safely send JSON to WebSocket, returns False if connection is closed.""" - nonlocal ws_closed - if ws_closed: - logger.debug( - "[WS:%s] safe_send_json skipped (ws_closed=True): type=%s", - session_id, - data.get("type"), - ) - return False - - msg_type = data.get("type") - if msg_type in { - "error", - "response", - "assistant_message_start", - "assistant_message_end", - }: - # Keep this low-noise but useful for tracing request lifecycle - logger.debug( - "[WS:%s] → send_json type=%s keys=%s", - 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", - session_id, - data.get("error_type"), - data.get("action_required"), - (data.get("error") or "")[:300], - ) - - try: - if msg_type == "error": - await persist_error_payload(data) - await websocket.send_json(data) - return True - except Exception as e: - logger.warning( - "[WS:%s] send_json failed for type=%s: %s", - session_id, - msg_type, - e, - exc_info=True, - ) - if "close message" in str(e).lower() or "closed" in str(e).lower(): - ws_closed = True - logger.debug("WebSocket closed, stopping sends") - return False - - async def send_typed(msg: ServerMessage) -> bool: - """Send a typed protocol message to the client.""" - return await safe_send_json(msg.model_dump(exclude_none=True)) + # WebSocketSender encapsulates sender.ws_closed, safe_send_json, + # persist_error_payload, send_typed, and send_typed_tool_lifecycle. + sender = WebSocketSender(websocket, session_id) - async def send_typed_tool_lifecycle( - msg: ServerToolCall | ServerToolResult, - ) -> bool: - """Send tool lifecycle frames to the client.""" - return await send_typed(msg) + # 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 ctx = None session_title = "" @@ -289,6 +214,7 @@ async def send_typed_tool_lifecycle( # get_or_load_session checks in-memory first, then falls # back to a fresh SQLite load. SQLite is the sole source of truth. ctx = await session_manager.get_or_load_session(session_id) + sender.ctx = ctx if ctx is None: # SQLite confirmed this session exists but load returned None. # This is unexpected (DB race, corrupt row, or schema mismatch). @@ -301,6 +227,7 @@ async def send_typed_tool_lifecycle( session_id, ) ctx = await session_manager.create_session(session_id) + sender.ctx = ctx else: # Update local state from loaded metadata session_title = ctx.title @@ -308,10 +235,12 @@ async def send_typed_tool_lifecycle( session_pinned = ctx.pinned else: ctx = await session_manager.create_session(session_id) + sender.ctx = ctx except Exception as e: logger.warning("SessionManager init failed, falling back: %s", e) try: ctx = await session_manager.create_session(session_id) + sender.ctx = ctx except Exception: logger.error("SessionManager fallback also failed", exc_info=True) await websocket.close(code=1011, reason="Session init failed") @@ -650,6 +579,7 @@ async def send_typed_tool_lifecycle( session_pinned = False ctx = await session_manager.create_session(session_id) + sender.ctx = ctx # Mark the new session as active await session_manager.mark_session_active(session_id) agent = ctx.agent @@ -708,6 +638,7 @@ async def send_typed_tool_lifecycle( pass ctx = await session_manager.create_session(session_id) + sender.ctx = ctx ctx.title = new_title ctx.working_directory = new_working_directory ctx.pinned = new_pinned @@ -1229,7 +1160,7 @@ async def drain_events_concurrent( ready_event: asyncio.Event = None, ): """Background task to drain events and send structured messages in real-time.""" - nonlocal collected_text, b1_streaming_used, ws_closed + nonlocal collected_text, b1_streaming_used import time as time_module # Capture agent and model metadata at the start. @@ -1360,7 +1291,7 @@ def resolve_tool_group_id( first_iteration = True while not stop_draining.is_set(): # Exit if WebSocket is closed - if ws_closed: + if sender.ws_closed: logger.debug( "WebSocket closed, exiting drain loop" ) @@ -2302,7 +2233,7 @@ def resolve_tool_group_id( "close message" in error_msg or "closed" in error_msg ): - ws_closed = True + sender.ws_closed = True logger.debug( "WebSocket closed during streaming, stopping drain" ) @@ -3502,18 +3433,8 @@ async def drain_events_with_signal(): # Title is stored only in the meta file, not in the filename session_name = session_id - # Wrap each message with metadata for complete session information. - # Chain `or` fallbacks: agent.name can be "" before the agent - # object is fully configured, so we cascade to context defaults. - agent_name_meta = ( - (agent.name if agent else "") - or ctx.agent_name - or "code-puppy" - ) - model_name_meta = ( - (agent.get_model_name() if agent else "") - or ctx.model_name - or "unknown" + agent_name_meta, model_name_meta = resolve_agent_model_meta( + agent=agent, ctx=ctx ) enhanced_history = build_enhanced_history( @@ -3536,65 +3457,43 @@ async def drain_events_with_signal(): enhanced_history, agent ) - # Write to SQLite for FE read path - try: - import datetime as _dt_mod + 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, + ) - _now_iso = _dt_mod.datetime.now( - _dt_mod.timezone.utc - ).isoformat() - await write_turn_to_sqlite( + # Send session metadata update to client + await safe_send_json( + build_session_meta_payload( session_id=session_id, - enhanced_history=enhanced_history, + session_name=session_name, + total_tokens=total_tokens, + message_count=message_count, title=session_title, working_directory=session_working_directory, - pinned=session_pinned, agent_name=agent_name, model_name=model_name, - total_tokens=total_tokens, - updated_at=_now_iso, - created_at=ctx.created_at.isoformat(), - ctx=ctx, ) - except Exception as _db_exc: - logger.debug( - "SQLite turn write skipped (DB not available): %s", - _db_exc, - ) - - # Send session metadata update to client - await websocket.send_json( - { - "type": "session_meta", - "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, - } ) # Broadcast session update to session monitoring clients - session_update_data = { - "session_id": session_id, - "session_name": session_name, - "title": session_title, - "working_directory": session_working_directory, - "timestamp": datetime.datetime.now().isoformat(), - "message_count": message_count, - "total_tokens": total_tokens, - "auto_saved": True, - "pickle_path": "", - "metadata_path": "", - "action": "created" - if message_count == 1 - else "updated", - } await connection_manager.broadcast_session_update( - session_update_data + 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, + ) ) except Exception as save_err: logger.warning( @@ -3646,7 +3545,7 @@ async def drain_events_with_signal(): 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 ws_closed: + if sender.ws_closed: break try: _err_msg = ServerError( @@ -3662,7 +3561,7 @@ async def drain_events_with_signal(): break except WebSocketDisconnect: - ws_closed = True + sender.ws_closed = True logger.debug("Chat WebSocket client disconnected") except Exception as e: logger.error("Chat WebSocket error: %s", e, exc_info=True) diff --git a/code_puppy/api/ws/send_utils.py b/code_puppy/api/ws/send_utils.py new file mode 100644 index 000000000..dd0eed92a --- /dev/null +++ b/code_puppy/api/ws/send_utils.py @@ -0,0 +1,151 @@ +"""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 + + # -- 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..0bc356e42 --- /dev/null +++ b/code_puppy/api/ws/session_persistence.py @@ -0,0 +1,143 @@ +"""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 typing import Any, Optional + +logger = logging.getLogger(__name__) + + +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"``. Legacy ``pickle_path`` / ``metadata_path`` + fields are kept as empty strings for backwards-compat. + """ + 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, + "pickle_path": "", + "metadata_path": "", + "action": "created" if message_count == 1 else "updated", + } + + +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/docs/migration/puppy-desk-cleanup-summary.md b/docs/migration/puppy-desk-cleanup-summary.md index d8fba5c79..595aee545 100644 --- a/docs/migration/puppy-desk-cleanup-summary.md +++ b/docs/migration/puppy-desk-cleanup-summary.md @@ -39,3 +39,53 @@ Recommended next split inside `chat_handler.py`: - Extracted history wrapping/timestamp/token estimation helpers into `code_puppy/api/ws/history_utils.py`. - Reduced inline persistence-preparation logic in `code_puppy/api/ws/chat_handler.py`. - Added focused unit tests for history wrapping, timestamp extraction, attachment metadata backfill, and safe token estimation failure handling. + +## Phase 3 checkpoint + +### Extracted modules + +- **`code_puppy/api/ws/send_utils.py`** (151 lines) + - `WebSocketSender` class encapsulating `ws_closed`, `safe_send_json`, + `persist_error_payload`, `send_typed`, `send_typed_tool_lifecycle`. + - Replaces 4 nested closures (~90 lines) in `chat_handler.py`. + - Independently testable without a live WebSocket or database. + +- **`code_puppy/api/ws/session_persistence.py`** (143 lines) + - `resolve_agent_model_meta(agent, ctx)` — or-chain fallback pattern. + - `build_session_meta_payload(...)` — client `session_meta` frame. + - `build_session_update_payload(...)` — broadcast payload for monitoring clients. + - `persist_turn_to_sqlite(...)` — wraps DB write with try/except guard. + +### chat_handler.py changes + +- Removed inline closure definitions (safe_send_json, persist_error_payload, + send_typed, send_typed_tool_lifecycle). +- Replaced with `WebSocketSender` construction and method aliases. +- Replaced inline agent/model meta resolution with `resolve_agent_model_meta()`. +- Replaced inline SQLite turn write with `persist_turn_to_sqlite()`. +- Replaced inline session_meta dict with `build_session_meta_payload()`. +- Replaced inline broadcast dict with `build_session_update_payload()`. +- Removed unused imports (`typing.Any`, `write_error_message_to_sqlite`, + `write_turn_to_sqlite`, `ServerMessage`). +- `ws_closed` references replaced with `sender.ws_closed`. + +### Size reduction + +| File | Before | After | +|------|--------|-------| +| chat_handler.py | 3,687 lines / 220 KB | 3,586 lines / 215 KB | +| session_persistence.py | — | 143 lines | +| send_utils.py | — | 151 lines | + +### Test coverage + +- 22 new focused tests across `test_session_persistence.py` and `test_send_utils.py`. +- Import isolation tests confirm no eager loading of `chat_handler`. +- Full focused regression suite: 111 tests pass (0 failures). + +### No wire protocol changes + +- `/ws/chat` route unchanged. +- All session_meta, session_update, error persistence payloads match + previous inline implementations exactly. +- No frontend emitter behavior changes. From 4a3889b256123fd4b09338e37198b84ad2ac5df7 Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Tue, 2 Jun 2026 16:28:44 -0500 Subject: [PATCH 14/29] fix: use result.get() for plugin key access in lifespan load_plugin_callbacks() returns dict with 'builtin' and 'user' keys but not always 'external'. Use .get() with default to avoid KeyError. --- code_puppy/api/app.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/code_puppy/api/app.py b/code_puppy/api/app.py index 8624db6fc..24e4200fd 100644 --- a/code_puppy/api/app.py +++ b/code_puppy/api/app.py @@ -65,7 +65,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: result = plugins.load_plugin_callbacks() logger.info( - f"✓ Loaded plugins: builtin={result['builtin']}, user={result['user']}, external={result['external']}" + f"✓ Loaded plugins: builtin={result.get('builtin', [])}, user={result.get('user', [])}, external={result.get('external', [])}" ) # Initialise shared SQLite database From 0dd008dad61b4ead275102d000c74e7fda4d8dbd Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Tue, 2 Jun 2026 18:46:15 -0500 Subject: [PATCH 15/29] docs: phase 3 post-extraction fixes plan Root cause analysis for missing tool call streaming: - ContextVar bridge (current_emitter_session_id) never wired - Pre-existing gap from original puppy-desk branch - Events tagged session_id=None, dropped by session-filtered subscriber - b1_streaming_used stays False, GUGI gets no tool lifecycle frames Also documents: stale sender session_id, session switch sync, executor shutdown cosmetic error. Execution plan: P0 ContextVar fix, P1 sender sync, future extraction targets. --- docs/migration/puppy-desk-phase3-fixes.md | 195 ++++++++++++++++++++++ 1 file changed, 195 insertions(+) create mode 100644 docs/migration/puppy-desk-phase3-fixes.md diff --git a/docs/migration/puppy-desk-phase3-fixes.md b/docs/migration/puppy-desk-phase3-fixes.md new file mode 100644 index 000000000..cb70e7677 --- /dev/null +++ b/docs/migration/puppy-desk-phase3-fixes.md @@ -0,0 +1,195 @@ +# Phase 3 — Post-extraction fixes & next steps + +> Generated 2026-06-02 after Phase 3 extraction + manual GUGI smoke test. + +## Critical finding: tool calls not streaming + +**Symptom:** User reported no tool calls visible in GUGI during manual test. +**Server logs confirmed:** `b1_streaming_used=False` on every agent run. + +### Root cause: ContextVar bridge never wired + +The frontend emitter's `emit_event()` relies on an implicit ContextVar +(`current_emitter_session_id`) to tag each event with the correct session. +The subscribe call in `chat_handler.py` uses session-filtered subscription: + +```python +event_queue = subscribe(session_id=session_id) # line ~2306 +``` + +But `current_emitter_session_id` is **never set** anywhere in the codebase — +not in `chat_handler.py`, not on the original `feature/puppy-desk` branch. + +Event routing works like this: +1. `register_callbacks.py::on_pre_tool_call()` calls `emit_event("tool_call_start", ...)` + with **no explicit session_id** +2. `emit_event` resolves session_id via `current_emitter_session_id.get()` → `None` +3. Event is tagged `session_id=None` +4. Subscriber filtering: *"Events with session_id=None are NOT delivered + to session-filtered subscribers"* +5. **Event is dropped** — never reaches `drain_events_concurrent` +6. `b1_streaming_used` stays `False` + +**This is a pre-existing gap** — the original `feature/puppy-desk` branch also +never set the ContextVar. It was NOT caused by Phase 3 extraction. + +### Fix (P0) + +In `chat_handler.py`, before starting the agent run, set the ContextVar: + +```python +from code_puppy.plugins.frontend_emitter.session_context import ( + current_emitter_session_id, +) + +# Before agent.run_with_mcp / agent.run(): +_emitter_token = current_emitter_session_id.set(session_id) + +# In finally cleanup: +current_emitter_session_id.reset(_emitter_token) +``` + +Location: around line ~2340 (just before `run_with_mcp`) and in the +corresponding `finally` block around line ~2800. + +--- + +## Bug 2: Stale session_id on WebSocketSender + +### Problem + +`WebSocketSender` is constructed at line 155 with the initial `session_id` +parameter from the URL query string. For new sessions, this is `None`. +At line 206, `session_id` gets reassigned to a generated value like +`WS_session_20260602_162504`, but the sender still holds the original `None`. + +Evidence in server logs: +``` +[WS:None] send_json failed for type=error ... +``` + +### Impact + +- **Logging:** All send_utils log lines show `[WS:None]` — useless for debugging +- **Error persistence:** `persist_error_payload` checks + `if not self._session_id: return` — silently skips DB writes for new sessions +- **No impact on sends themselves** — `safe_send_json` doesn't gate on session_id + +### Fix (P1) + +Option A — Add a `session_id` setter to `WebSocketSender` and call it after +session_id is generated/resolved: + +```python +# In send_utils.py +@session_id.setter +def session_id(self, value: str) -> None: + self._session_id = value + +# In chat_handler.py after session_id is generated (line ~207): +sender._session_id = session_id # or sender.session_id = session_id +``` + +Option B — Defer `WebSocketSender` construction until after session_id is +resolved (move from line 155 to ~line 210). This would require reworking the +early error sends that happen before session creation. + +**Recommendation:** Option A — simpler, less risky. + +--- + +## Bug 3: Session switch doesn't update sender.session_id + +### Problem + +Session switching at line ~577 reassigns `session_id = new_session_id` but +the sender's `_session_id` stays stale. Same stale-logging and +error-persistence issue as Bug 2, but specifically during session switches. + +### Fix (P1) + +Wherever `session_id` is reassigned during session switching, also update +the sender: + +```python +session_id = new_session_id +sender._session_id = session_id # sync sender +``` + +Locations to patch (grep `session_id =` inside switch_session blocks): +- Line ~577 (new session creation during switch) +- Line ~638 (existing session load during switch) + +--- + +## Non-blocking issues found during investigation + +### 4. `b1_streaming_used=False` fallback path (informational) + +When B1 streaming doesn't fire, `chat_handler` falls through to the B2 +path which reconstructs response frames from the completed agent result. +B2 does NOT stream tool calls — it only sends the final response text. + +This means **the GUGI never sees tool lifecycle frames** when the ContextVar +bridge is missing. Fixing Bug 1 (ContextVar) resolves this entirely. + +### 5. `app.py` plugin key access (already fixed) + +`load_plugin_callbacks()` returns `{'builtin': [...], 'user': [...]}` but +`app.py` was accessing `result['external']` with bracket notation. +**Fixed:** commit `e8f6e56` — now uses `.get()` with defaults. + +### 6. `ws_sessions` executor shutdown error (cosmetic) + +``` +ERROR code_puppy.api.app: Error shutting down ws_sessions executor: + module 'code_puppy.api.routers.ws_sessions' has no attribute '_executor' +``` + +Non-blocking cosmetic error during server shutdown. Low priority. + +--- + +## Execution plan + +### Immediate (P0) — Fix tool call streaming + +| Step | Action | Risk | +|------|--------|------| +| 1 | Wire `current_emitter_session_id` ContextVar in `chat_handler.py` | Low — additive, no existing behavior changed | +| 2 | Add `sender.session_id` setter to `WebSocketSender` | Low — internal property | +| 3 | Sync `sender._session_id` after session generation + switch | Low — logging/persistence only | +| 4 | Restart server, test with GUGI — verify tool_call frames appear | Validation | +| 5 | Run full test suite (111 tests) | Regression check | + +### Follow-up (P1) — Tests for the fix + +| Step | Action | +|------|--------| +| 6 | Add integration test: connect → send message → assert tool_call frames received | +| 7 | Add unit test: verify ContextVar is set before agent run and reset in finally | +| 8 | Add unit test: verify sender.session_id updates after session switch | + +### Future phases — Further modularization + +| Target | Lines | Extraction | +|--------|-------|------------| +| `drain_events_concurrent` + `drain_events_with_signal` | ~200 | → `event_drain.py` | +| Session switch/create handler | ~150 | → `session_switch.py` | +| Tool group ID resolution helpers | ~70 | → merge into `send_utils.py` or `tool_lifecycle.py` | +| `session_context.py` cleanup | ~400 | Split by concern | +| `db/queries.py` cleanup | ~260 | Split by table/domain | + +### Phase 3 scorecard + +| Item | Status | +|------|--------| +| `session_persistence.py` extracted | ✅ Done | +| `send_utils.py` extracted | ✅ Done | +| `chat_handler.py` uses new modules | ✅ Verified (traceback + unit tests) | +| Unit tests (22 new) | ✅ All pass | +| Full regression suite (111) | ✅ All pass | +| Wire protocol unchanged | ✅ Verified | +| ContextVar bridge for emitter | ❌ **Not wired — P0 fix needed** | +| Sender session_id sync | ❌ **Stale after generation — P1 fix needed** | +| GUGI tool call streaming | ❌ **Blocked by ContextVar bug** | From aca870f03de39c3cf4c697e54420fdc8c8c34840 Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Tue, 2 Jun 2026 21:50:23 -0500 Subject: [PATCH 16/29] fix: wire emitter ContextVar + sync sender session_id MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P0 — Tool call streaming was broken because the emitter ContextVar (current_emitter_session_id) was never set. Events emitted by register_callbacks had session_id=None, which session-filtered subscribers dropped. Now set/reset around the agent run window. P1 — WebSocketSender.session_id was stale after session generation and session switching. Added setter property and synced at all 3 assignment sites (new session, switch-to-new, switch-to-existing). Pre-existing gap — original feature/puppy-desk also never set the ContextVar. Not a Phase 3 regression. Test results: 4095 passed, 29 skipped (3 pre-existing failures excluded) --- code_puppy/api/ws/chat_handler.py | 34 +++++++++++++++++++++++++++++-- code_puppy/api/ws/send_utils.py | 4 ++++ 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/code_puppy/api/ws/chat_handler.py b/code_puppy/api/ws/chat_handler.py index 195c5bf06..1771fd0e9 100644 --- a/code_puppy/api/ws/chat_handler.py +++ b/code_puppy/api/ws/chat_handler.py @@ -204,6 +204,7 @@ async def websocket_chat( # Generate new timestamp-based session ID 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) try: @@ -574,6 +575,7 @@ async def websocket_chat( await session_manager.mark_session_inactive(session_id) session_id = new_session_id + sender.session_id = session_id session_title = "" session_working_directory = "" session_pinned = False @@ -607,6 +609,7 @@ async def websocket_chat( await session_manager.mark_session_inactive(session_id) session_id = new_session_id + sender.session_id = session_id # Try loading via SessionManager (checks in-memory first, then SQLite) loaded_ctx = await session_manager.get_or_load_session( @@ -2534,6 +2537,20 @@ async def drain_events_with_signal(): # Run the agent in its own task so it can be cancelled independently. # Suppress duplicate emitter lifecycle callbacks only during the # run_with_mcp execution window (pre/post_tool_call hooks fire there). + # Wire the emitter ContextVar so events + # emitted by register_callbacks are tagged + # with this session and reach our subscriber. + 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) active_agent_task = asyncio.create_task( agent.run_with_mcp( @@ -2785,6 +2802,19 @@ async def drain_events_with_signal(): # End suppression scope once run_with_mcp window ends, # including completion/cancel/background handoff paths. set_suppress_emitter_tool_events(False) + # Reset emitter ContextVar so stale session_id + # doesn't leak into unrelated coroutines. + if _emitter_token is not None: + try: + from code_puppy.plugins.frontend_emitter.session_context import ( + current_emitter_session_id, + ) + + current_emitter_session_id.reset( + _emitter_token + ) + except ImportError: + pass # Clear session context for prompt generation clear_session_working_directory() @@ -3433,8 +3463,8 @@ async def drain_events_with_signal(): # Title is stored only in the meta file, not in the filename session_name = session_id - agent_name_meta, model_name_meta = resolve_agent_model_meta( - agent=agent, ctx=ctx + agent_name_meta, model_name_meta = ( + resolve_agent_model_meta(agent=agent, ctx=ctx) ) enhanced_history = build_enhanced_history( diff --git a/code_puppy/api/ws/send_utils.py b/code_puppy/api/ws/send_utils.py index dd0eed92a..fae384129 100644 --- a/code_puppy/api/ws/send_utils.py +++ b/code_puppy/api/ws/send_utils.py @@ -61,6 +61,10 @@ def ctx(self, value: Any) -> None: 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: From 865d7ca162d7bafc64be9a25c32e13a959eb3b12 Mon Sep 17 00:00:00 2001 From: maestro Date: Thu, 4 Jun 2026 15:42:33 -0500 Subject: [PATCH 17/29] style: sync import ordering with external repo (isort) Aligns chat_handler.py import order and removes extra blank lines in test files so ws/ layer is byte-identical across both repos. --- code_puppy/api/tests/test_send_utils.py | 2 - code_puppy/api/tests/test_session_load.py | 130 ++++++++++++++---- .../api/tests/test_session_model_compat.py | 6 +- .../api/tests/test_session_persistence.py | 1 - code_puppy/api/ws/chat_handler.py | 22 +-- 5 files changed, 117 insertions(+), 44 deletions(-) diff --git a/code_puppy/api/tests/test_send_utils.py b/code_puppy/api/tests/test_send_utils.py index 9d40a4e71..561bdc924 100644 --- a/code_puppy/api/tests/test_send_utils.py +++ b/code_puppy/api/tests/test_send_utils.py @@ -2,12 +2,10 @@ from __future__ import annotations - import pytest from code_puppy.api.ws.send_utils import WebSocketSender - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- diff --git a/code_puppy/api/tests/test_session_load.py b/code_puppy/api/tests/test_session_load.py index 674f2eae4..597cd3620 100644 --- a/code_puppy/api/tests/test_session_load.py +++ b/code_puppy/api/tests/test_session_load.py @@ -129,8 +129,14 @@ async def test_load_from_sqlite_session_not_in_db(): 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.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") @@ -160,7 +166,9 @@ async def test_load_from_sqlite_get_active_messages_raises_returns_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.session_exists", new=AsyncMock(return_value=True) + ), patch( "code_puppy.api.db.queries.get_active_messages", new=AsyncMock(side_effect=RuntimeError("db locked")), @@ -187,15 +195,22 @@ async def test_load_from_sqlite_all_null_pydantic_json_returns_empty_history(): 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.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" + assert messages == [], ( + "Empty history expected when all rows have NULL pydantic_json" + ) @pytest.mark.asyncio @@ -211,8 +226,13 @@ async def test_load_from_sqlite_mixed_null_and_valid_rows(): 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.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") @@ -237,8 +257,13 @@ async def test_load_from_sqlite_valid_rows_parsed_correctly(): 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.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") @@ -264,8 +289,13 @@ async def test_load_from_sqlite_malformed_pydantic_json_skipped(caplog): 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.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"), ): @@ -293,8 +323,13 @@ async def test_load_from_sqlite_metadata_populated(): 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.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") @@ -314,8 +349,13 @@ async def test_load_from_sqlite_no_session_row_gives_empty_meta(): 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.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") @@ -335,8 +375,13 @@ async def test_load_from_sqlite_sessions_table_query_fails_gives_empty_meta(): 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.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") @@ -390,7 +435,10 @@ async def test_load_session_happy_path_sets_history_and_registers_context(): 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"), + patch( + "code_puppy.api.session_context.get_global_model_name", + return_value="gpt-4o", + ), ): ctx = await mgr.load_session("my-chat-session") @@ -414,7 +462,10 @@ async def test_load_session_registers_context_in_sessions_dict(): 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"), + patch( + "code_puppy.api.session_context.get_global_model_name", + return_value="gpt-4o", + ), ): ctx = await mgr.load_session("registered-session") @@ -436,8 +487,14 @@ def _load_agent_side_effect(name): 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"), + 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") @@ -455,8 +512,13 @@ async def test_load_session_valid_created_at_parsed(): 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"), + 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") @@ -474,8 +536,13 @@ async def test_load_session_malformed_created_at_falls_back_to_now(): 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"), + 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") @@ -493,11 +560,16 @@ async def test_load_session_empty_created_at_falls_back_to_now(): 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"), + 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 \ No newline at end of file + 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 index 0b008fd87..1309badd6 100644 --- a/code_puppy/api/tests/test_session_model_compat.py +++ b/code_puppy/api/tests/test_session_model_compat.py @@ -4,7 +4,11 @@ import pytest -from code_puppy.api.session_context import SessionContext, SessionManager, _apply_session_model +from code_puppy.api.session_context import ( + SessionContext, + SessionManager, + _apply_session_model, +) class LegacyAgentWithoutSetter: diff --git a/code_puppy/api/tests/test_session_persistence.py b/code_puppy/api/tests/test_session_persistence.py index f63d04873..b97cbda54 100644 --- a/code_puppy/api/tests/test_session_persistence.py +++ b/code_puppy/api/tests/test_session_persistence.py @@ -8,7 +8,6 @@ resolve_agent_model_meta, ) - # --------------------------------------------------------------------------- # resolve_agent_model_meta # --------------------------------------------------------------------------- diff --git a/code_puppy/api/ws/chat_handler.py b/code_puppy/api/ws/chat_handler.py index 1771fd0e9..c04f2d9aa 100644 --- a/code_puppy/api/ws/chat_handler.py +++ b/code_puppy/api/ws/chat_handler.py @@ -37,21 +37,14 @@ save_agent_result_in_background, ) from code_puppy.api.ws.connection_manager import connection_manager -from code_puppy.api.ws.response_frames import ( - build_assistant_text_stream_frames, - build_error_response_frames, - parse_api_error, -) from code_puppy.api.ws.history_utils import ( build_enhanced_history, estimate_total_tokens, ) -from code_puppy.api.ws.send_utils import WebSocketSender -from code_puppy.api.ws.session_persistence import ( - build_session_meta_payload, - build_session_update_payload, - persist_turn_to_sqlite, - resolve_agent_model_meta, +from code_puppy.api.ws.response_frames import ( + build_assistant_text_stream_frames, + build_error_response_frames, + parse_api_error, ) from code_puppy.api.ws.schemas import ( PROTOCOL_VERSION, @@ -75,6 +68,13 @@ ServerUserMessage, ServerWorkingDirectoryChanged, ) +from code_puppy.api.ws.send_utils import WebSocketSender +from code_puppy.api.ws.session_persistence import ( + build_session_meta_payload, + build_session_update_payload, + persist_turn_to_sqlite, + resolve_agent_model_meta, +) from code_puppy.config import get_global_model_name from code_puppy.messaging.bus import get_message_bus from code_puppy.session_storage import generate_heuristic_title From a3281ebe91aea6aacb855c2b74142673c8bdb9e5 Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Wed, 10 Jun 2026 13:56:44 -0500 Subject: [PATCH 18/29] wip: preserve feature/puppy-desk-migration state for manual ws sync --- code_puppy/api/app.py | 5 - code_puppy/api/db/__init__.py | 2 - code_puppy/api/db/message_utils.py | 4 +- code_puppy/api/db/seeder.py | 550 +-- code_puppy/api/routers/sessions.py | 115 +- code_puppy/api/routers/ws_sessions.py | 19 +- code_puppy/api/session_cache.py | 39 +- code_puppy/api/session_context.py | 4 +- code_puppy/api/templates/chat.html | 135 +- .../tests/test_chat_template_session_state.py | 13 + .../api/tests/test_session_persistence.py | 4 - code_puppy/api/ws/chat_context.py | 72 + code_puppy/api/ws/chat_event_adapter.py | 243 ++ code_puppy/api/ws/chat_handler.py | 3145 ++--------------- code_puppy/api/ws/chat_tool_lifecycle.py | 520 +++ code_puppy/api/ws/chat_turn_runner.py | 273 ++ code_puppy/api/ws/chat_turn_state.py | 23 + code_puppy/api/ws/schemas.py | 1 + code_puppy/api/ws/session_persistence.py | 118 +- code_puppy/api/ws/tests/test_chat_context.py | 58 + .../api/ws/tests/test_chat_event_adapter.py | 164 + .../api/ws/tests/test_chat_tool_lifecycle.py | 231 ++ .../api/ws/tests/test_chat_turn_runner.py | 227 ++ .../api/ws/tests/test_session_persistence.py | 197 ++ .../tests/test_tool_group_id_consistency.py | 69 +- .../api/ws/tests/test_ws_command_handler.py | 108 + .../api/ws/tests/test_ws_control_messages.py | 217 ++ code_puppy/api/ws/tests/test_ws_post_run.py | 149 + .../api/ws/tests/test_ws_session_bootstrap.py | 235 ++ .../api/ws/tests/test_ws_stream_drain.py | 120 + .../api/ws/tests/test_ws_turn_finalization.py | 195 + .../api/ws/tests/test_ws_turn_preparation.py | 62 + code_puppy/api/ws/ws_chat_runtime.py | 49 + code_puppy/api/ws/ws_command_handler.py | 93 + code_puppy/api/ws/ws_control_messages.py | 606 ++++ code_puppy/api/ws/ws_post_run.py | 247 ++ code_puppy/api/ws/ws_session_bootstrap.py | 271 ++ code_puppy/api/ws/ws_stream_drain.py | 95 + code_puppy/api/ws/ws_turn_finalization.py | 218 ++ code_puppy/api/ws/ws_turn_preparation.py | 127 + code_puppy/cli_runner.py | 6 +- .../command_line/load_context_completion.py | 4 +- code_puppy/messaging/bus.py | 19 +- .../switch_agent_resume/register_callbacks.py | 10 +- code_puppy/session_storage.py | 232 +- code_puppy/tools/agent_tools.py | 25 +- .../ws-chat-handler-refactor-checklist.md | 464 +++ tests/test_agent_tools.py | 19 +- tests/test_load_context_completion.py | 18 +- tests/test_messaging_bus.py | 34 +- tests/test_session_storage.py | 16 +- tests/test_session_storage_edge_cases.py | 407 +-- tests/test_session_storage_json_only.py | 16 + 53 files changed, 6199 insertions(+), 4094 deletions(-) create mode 100644 code_puppy/api/tests/test_chat_template_session_state.py create mode 100644 code_puppy/api/ws/chat_context.py create mode 100644 code_puppy/api/ws/chat_event_adapter.py create mode 100644 code_puppy/api/ws/chat_tool_lifecycle.py create mode 100644 code_puppy/api/ws/chat_turn_runner.py create mode 100644 code_puppy/api/ws/chat_turn_state.py create mode 100644 code_puppy/api/ws/tests/test_chat_context.py create mode 100644 code_puppy/api/ws/tests/test_chat_event_adapter.py create mode 100644 code_puppy/api/ws/tests/test_chat_tool_lifecycle.py create mode 100644 code_puppy/api/ws/tests/test_chat_turn_runner.py create mode 100644 code_puppy/api/ws/tests/test_session_persistence.py create mode 100644 code_puppy/api/ws/tests/test_ws_command_handler.py create mode 100644 code_puppy/api/ws/tests/test_ws_control_messages.py create mode 100644 code_puppy/api/ws/tests/test_ws_post_run.py create mode 100644 code_puppy/api/ws/tests/test_ws_session_bootstrap.py create mode 100644 code_puppy/api/ws/tests/test_ws_stream_drain.py create mode 100644 code_puppy/api/ws/tests/test_ws_turn_finalization.py create mode 100644 code_puppy/api/ws/tests/test_ws_turn_preparation.py create mode 100644 code_puppy/api/ws/ws_chat_runtime.py create mode 100644 code_puppy/api/ws/ws_command_handler.py create mode 100644 code_puppy/api/ws/ws_control_messages.py create mode 100644 code_puppy/api/ws/ws_post_run.py create mode 100644 code_puppy/api/ws/ws_session_bootstrap.py create mode 100644 code_puppy/api/ws/ws_stream_drain.py create mode 100644 code_puppy/api/ws/ws_turn_finalization.py create mode 100644 code_puppy/api/ws/ws_turn_preparation.py create mode 100644 docs/migration/ws-chat-handler-refactor-checklist.md create mode 100644 tests/test_session_storage_json_only.py diff --git a/code_puppy/api/app.py b/code_puppy/api/app.py index 24e4200fd..47881e265 100644 --- a/code_puppy/api/app.py +++ b/code_puppy/api/app.py @@ -71,14 +71,9 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: # Initialise shared SQLite database try: from code_puppy.api.db.connection import init_db - from code_puppy.api.db.seeder import seed_from_pkl_dirs await init_db() logger.info("✓ aiosqlite DB initialised") - - # Seed existing pkl sessions in the background — non-blocking - asyncio.create_task(seed_from_pkl_dirs()) - logger.info("✓ SQLite seeder task started") except Exception as _db_exc: logger.error("SQLite DB init failed (continuing without DB): %s", _db_exc) diff --git a/code_puppy/api/db/__init__.py b/code_puppy/api/db/__init__.py index f7805b048..83cffcb97 100644 --- a/code_puppy/api/db/__init__.py +++ b/code_puppy/api/db/__init__.py @@ -24,7 +24,6 @@ write_system_message_to_sqlite, write_turn_to_sqlite, ) -from code_puppy.api.db.seeder import seed_from_pkl_dirs __all__ = [ "init_db", @@ -44,6 +43,5 @@ "get_active_messages", "mark_messages_compacted", "insert_compaction_log", - "seed_from_pkl_dirs", "write_turn_to_sqlite", ] diff --git a/code_puppy/api/db/message_utils.py b/code_puppy/api/db/message_utils.py index 5a22d69b7..8c38a1851 100644 --- a/code_puppy/api/db/message_utils.py +++ b/code_puppy/api/db/message_utils.py @@ -1,6 +1,6 @@ """Shared helpers for extracting display fields from pydantic-ai ModelMessages. -Used by seeder.py (pkl import) and session_context.py (live WS saves). +Used by the SQLite persistence helpers and session_context.py. Must not import from seeder.py or session_context.py to avoid circular imports. """ @@ -68,7 +68,7 @@ def pydantic_json_for_message(msg: Any) -> Optional[str]: Stores as a single-element list: `[msg_json]`. Deserialise with: `ModelMessagesTypeAdapter.validate_json(s)[0]` - Returns None if serialisation fails (e.g. very old pkl message types). + Returns None if serialisation fails for unsupported historical message shapes. """ try: from pydantic_ai.messages import ModelMessagesTypeAdapter diff --git a/code_puppy/api/db/seeder.py b/code_puppy/api/db/seeder.py index 0c5a5b2d3..d2d31ffde 100644 --- a/code_puppy/api/db/seeder.py +++ b/code_puppy/api/db/seeder.py @@ -1,555 +1,17 @@ -"""Seed ~/.puppy_desk/chat_messages.db from existing pkl session files. +"""Retired legacy session seeder. -Scans both: - ~/.code_puppy/ws_sessions/ — live WS sessions (HMAC-signed pickle) - ~/.code_puppy/autosaves/ — CLI autosave sessions (raw or legacy-signed pickle) - -For each .pkl file not already in the DB: - 1. Deserialises the pickle natively (no subprocess, no JSON round-trip) - 2. Serialises each ModelMessage with ModelMessagesTypeAdapter.dump_json() - → stored verbatim in the pydantic_json column for perfect replay - 3. Extracts display fields (role, content, thinking, tool calls) from message parts - 4. Writes session + messages + tool_calls in a single transaction - -Idempotent: existing sessions are skipped. Safe to call on every startup. +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 hashlib -import hmac -import json import logging -import pickle -import re -import time -import uuid -from datetime import datetime, timezone -from pathlib import Path -from typing import Any, Optional - -from code_puppy.api.db.message_utils import ( - extract_content, - extract_thinking, - get_message_timestamp, - get_role, - pydantic_json_for_message, -) logger = logging.getLogger(__name__) -# --------------------------------------------------------------------------- -# One-time migration marker — written after a successful complete seed. -# On subsequent startups the seeder skips entirely (no more pkl to import). -# --------------------------------------------------------------------------- -_SEEDER_DONE_MARKER = Path("~/.puppy_desk/.seeder_migration_done").expanduser() - - -# --------------------------------------------------------------------------- -# Legacy pickle header constants (from session_storage.py) -# --------------------------------------------------------------------------- - -_LEGACY_SIGNED_HEADER = b"CPSESSION\x01" -_LEGACY_SIGNATURE_SIZE = 32 - - -def _get_hmac_key() -> bytes: - """Return HMAC key for verifying WS session pickle files. - - Copied from session_context to make seeder.py self-contained after - _get_hmac_key was removed from session_context in desk-puppy-cr1x. - """ - import os - - key_env = os.environ.get("CODE_PUPPY_SESSION_KEY", "") - if key_env: - return key_env.encode() - # Fall back to a machine-stable key derived from the DB path - db_path = str(Path("~/.puppy_desk/chat_messages.db").expanduser()) - return db_path.encode() - - -def _extract_pickle_payload(raw: bytes) -> bytes: - """Strip legacy or new-format prefix and return the raw pickle bytes.""" - if raw.startswith(_LEGACY_SIGNED_HEADER): - offset = len(_LEGACY_SIGNED_HEADER) + _LEGACY_SIGNATURE_SIZE - return raw[offset:] - return raw - - -def _verify_and_load_ws_pickle(raw: bytes, hmac_key: bytes) -> Optional[list[Any]]: - """Load a WS session pickle after HMAC verification. - - WS session pkls are written by session_context.py: - file = 32-byte-HMAC-SHA256 + pickle(history) - - Returns None if verification fails or the file is too small. - """ - if len(raw) < 33: # 32-byte sig + at least 1 byte of pickle - return None - sig, blob = raw[:32], raw[32:] - expected = hmac.new(hmac_key, blob, hashlib.sha256).digest() - if not hmac.compare_digest(sig, expected): - logger.warning("HMAC verification failed — skipping") - return None - try: - data = pickle.loads(blob) # noqa: S301 - return data if isinstance(data, list) else None - except Exception as exc: - logger.warning("Pickle load error: %s", exc) - return None - - -def _load_autosave_pickle(raw: bytes) -> Optional[list[Any]]: - """Load an autosave pickle (no HMAC — strip legacy header if present).""" - payload = _extract_pickle_payload(raw) - try: - data = pickle.loads(payload) # noqa: S301 - return data if isinstance(data, list) else None - except Exception as exc: - logger.warning("Autosave pickle load error: %s", exc) - return None - - -# --------------------------------------------------------------------------- -# Message part extraction helpers (content/thinking/role moved to message_utils) -# --------------------------------------------------------------------------- - -_PART_TYPE_TOOL_CALL = {"ToolCallPart"} -_PART_TYPE_TOOL_RETURN = {"ToolReturnPart"} - - -def _extract_tool_calls(msg: Any) -> list[dict[str, Any]]: - """Extract ToolCallPart entries from a ModelResponse.""" - 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 _extract_tool_returns(msg: Any) -> list[dict[str, Any]]: - """Extract ToolReturnPart entries from a ModelRequest.""" - parts = getattr(msg, "parts", []) - result = [] - for part in parts: - if type(part).__name__ not in _PART_TYPE_TOOL_RETURN: - continue - content = getattr(part, "content", "") - try: - result_val = json.loads(content) if isinstance(content, str) else content - except Exception: - result_val = content - result.append( - { - "id": getattr(part, "tool_call_id", ""), - "name": getattr(part, "tool_name", "unknown"), - "result": result_val, - } - ) - return result - - -def _get_timestamp(msg: Any, wrapper: Optional[dict], fallback: str) -> str: - """Best-effort timestamp extraction.""" - # Try wrapper ts first - if wrapper: - ts = wrapper.get("ts", "") - if ts: - return str(ts) - # Try message-level timestamp - ts_attr_result = get_message_timestamp(msg) - if ts_attr_result is not None: - return ts_attr_result - return fallback - - -# --------------------------------------------------------------------------- -# Per-session import -# --------------------------------------------------------------------------- - -_SESSION_CONTEXT_RE = re.compile(r"^\[Session Context:[^\]]*\]\n\n?") - - -def _strip_session_context(raw: str) -> Optional[str]: - cleaned = _SESSION_CONTEXT_RE.sub("", raw).lstrip() - return cleaned if cleaned != raw else None - - -async def _import_session( - session_id: str, - history: list[Any], - meta: dict[str, Any], - now_iso: str, -) -> int: - """Write one session and all its messages/tool_calls to the DB. - - Returns the number of message rows written. - """ - from code_puppy.api.db.queries import ( - insert_messages_batch, - insert_tool_calls_batch, - upsert_session, - ) - - created_at = meta.get("timestamp") or now_iso - title = meta.get("title", "") - agent_name = meta.get("agent_name", "code-puppy") - model_name = meta.get("model_name", "") - working_directory = meta.get("working_directory", "") - pinned = bool(meta.get("pinned", False)) - total_tokens = meta.get("total_tokens", 0) - - 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=created_at, - message_count=len(history), - total_tokens=total_tokens, - deleted_at=None, - ) - - message_rows: list[dict[str, Any]] = [] - tool_call_rows: list[dict[str, Any]] = [] - seq = 0 - # Track pending tool calls by id for result matching - pending_tool_calls: dict[str, dict[str, Any]] = {} - - for item in history: - # ---- Unwrap the item ---------------------------------------- - wrapper: Optional[dict] = None - msg: Any - - if isinstance(item, dict) and item.get("msg") == "system": - # System message in WS format: {'msg': 'system', 'content': ..., 'path': ...} - seq += 1 - message_rows.append( - { - "session_id": session_id, - "seq": seq, - "role": "system", - "content": item.get("content", ""), - "type": "system", - "agent_name": item.get("agent", agent_name), - "model_name": item.get("model", model_name), - "timestamp": item.get("ts", now_iso), - "system_message_type": item.get("system_message_type", ""), - "system_message_path": item.get("path", ""), - "pydantic_json": None, - "token_count": 0, - } - ) - continue - - if isinstance(item, dict) and "msg" in item: - wrapper = item - msg = item["msg"] - else: - msg = item - - # ---- Skip if not a pydantic-ai message ---------------------- - if not hasattr(msg, "parts"): - continue - - role = get_role(msg) - ts = _get_timestamp(msg, wrapper, now_iso) - content = extract_content(msg) - thinking = extract_thinking(msg) - pydantic_json = pydantic_json_for_message(msg) - - msg_agent = wrapper.get("agent", agent_name) if wrapper else agent_name - msg_model = wrapper.get("model", model_name) if wrapper else model_name - - # clean_content: strip [Session Context: ...] for user messages - clean_content: Optional[str] = None - raw_clean = wrapper.get("clean_content") if wrapper else None - if raw_clean: - clean_content = raw_clean - elif role == "user" and content: - clean_content = _strip_session_context(content) - - # attachments_json - attachments = wrapper.get("attachments") if wrapper else None - attachments_json = json.dumps(attachments) if attachments else None - - seq += 1 - message_seq = seq - - # ---- Estimate token count (simple char/4 heuristic) --------- - token_count = max(1, len(content) // 4) - - message_rows.append( - { - "session_id": session_id, - "seq": message_seq, - "role": role, - "content": content, - "type": type(msg).__name__, - "agent_name": msg_agent, - "model_name": msg_model, - "timestamp": ts, - "thinking": thinking, - "attachments_json": attachments_json, - "clean_content": clean_content, - "pydantic_json": pydantic_json, - "token_count": token_count, - } - ) - - # ---- Tool calls from assistant messages --------------------- - if role == "assistant": - for tc in _extract_tool_calls(msg): - pending_tool_calls[tc["id"]] = { - "tc_id": tc["id"], - "name": tc["name"], - "args": tc["args"], - "parent_seq": message_seq, - "agent": msg_agent, - "model": msg_model, - "ts": ts, - } - - # ---- Tool returns from user messages (Anthropic API format) - - if role == "user": - for tr in _extract_tool_returns(msg): - tid = tr["id"] - if tid in pending_tool_calls: - tc = pending_tool_calls.pop(tid) - seq += 1 - try: - args_json = json.dumps(tc["args"]) - except Exception: - args_json = str(tc["args"]) - # Serialize result - handle Pydantic models and complex objects - result_val = tr["result"] - try: - if hasattr(result_val, "model_dump"): - result_json = json.dumps(result_val.model_dump()) - elif hasattr(result_val, "dict"): - result_json = json.dumps(result_val.dict()) - elif hasattr(result_val, "__dict__"): - result_json = json.dumps(vars(result_val)) - else: - result_json = json.dumps(result_val) - except Exception: - try: - result_json = json.dumps(str(result_val)) - except Exception: - result_json = json.dumps({"raw": str(result_val)}) - - tool_ts_str = tc.get("ts", ts) - try: - tool_ts = datetime.fromisoformat( - str(tool_ts_str).replace("Z", "+00:00") - ).timestamp() - except Exception: - tool_ts = time.time() - - tool_call_rows.append( - { - "id": tid, - "session_id": session_id, - "parent_message_seq": tc["parent_seq"], - "seq": seq, - "tool_name": tc["name"], - "args_json": args_json, - "result_json": result_json, - "status": "success", - "agent_name": tc["agent"], - "model_name": tc["model"], - "timestamp": tool_ts, - } - ) - - # Write everything in batches (each batch uses its own transaction) - await insert_messages_batch(message_rows) - await insert_tool_calls_batch(tool_call_rows) - return len(message_rows) - - -# --------------------------------------------------------------------------- -# Directory scanners -# --------------------------------------------------------------------------- - - -def _load_meta(meta_path: Path) -> dict[str, Any]: - try: - if meta_path.exists(): - return json.loads(meta_path.read_text(encoding="utf-8")) - except Exception: - pass - return {} - - -async def _seed_directory( - sessions_dir: Path, - is_ws: bool, - hmac_key: Optional[bytes], - now_iso: str, -) -> tuple[int, int, int]: - """Seed all pkl files from *sessions_dir*. Returns (imported, skipped, failed).""" - from code_puppy.api.db.queries import session_exists - - if not sessions_dir.exists(): - logger.debug("Sessions dir not found, skipping: %s", sessions_dir) - return 0, 0, 0 - - pkl_files = sorted(sessions_dir.glob("*.pkl")) - total = len(pkl_files) - if total == 0: - return 0, 0, 0 - - imported = skipped = failed = 0 - width = len(str(total)) - - for i, pkl_path in enumerate(pkl_files, 1): - session_id = pkl_path.stem - - # Skip tiny files (HMAC-only stubs, empty autosaves) - try: - if pkl_path.stat().st_size < 50: - skipped += 1 - continue - except OSError: - skipped += 1 - continue - - if await session_exists(session_id): - skipped += 1 - continue - - try: - raw = pkl_path.read_bytes() - except OSError as exc: - logger.warning( - "[%*d/%d] %s — read error: %s", width, i, total, session_id, exc - ) - failed += 1 - continue - - # Deserialise - if is_ws and hmac_key is not None: - history = _verify_and_load_ws_pickle(raw, hmac_key) - else: - history = _load_autosave_pickle(raw) - - if not history: - skipped += 1 - continue - - # Load metadata - meta_path = pkl_path.with_name(f"{session_id}_meta.json") - meta = _load_meta(meta_path) - - try: - n = await _import_session(session_id, history, meta, now_iso) - logger.info( - "[%*d/%d] %-50s %4d msgs — imported", - width, - i, - total, - session_id, - n, - ) - imported += 1 - except Exception as exc: - logger.error( - "[%*d/%d] %s — import error: %s", - width, - i, - total, - session_id, - exc, - exc_info=True, - ) - failed += 1 - - return imported, skipped, failed - - -# --------------------------------------------------------------------------- -# Public entry point -# --------------------------------------------------------------------------- async def seed_from_pkl_dirs() -> None: - """Asynchronously seed the DB from both ws_sessions and autosaves directories. - - All database operations go through the shared aiosqlite connection so this - must run on the event loop that owns the DB connection. File I/O for - reading pickle files is blocking but acceptable at startup since the seeder - runs as a background asyncio.create_task(). - - Safe to call on every startup — already-imported sessions are skipped. - """ - # One-time migration guard: if we already migrated all pkl files, skip. - if _SEEDER_DONE_MARKER.exists(): - logger.debug("Seeder already completed (marker found) — skipping pkl scan") - return - - from code_puppy.config import AUTOSAVE_DIR, WS_SESSION_DIR - - now_iso = datetime.now(timezone.utc).isoformat() - - # ---- WS sessions (HMAC-signed) ----------------------------------- - ws_dir = Path(WS_SESSION_DIR) - hmac_key: Optional[bytes] = None - try: - hmac_key = _get_hmac_key() - except Exception as exc: - logger.warning("Could not load HMAC key, WS sessions will be skipped: %s", exc) - - if hmac_key is not None: - logger.info("Seeding WS sessions from %s ...", ws_dir) - wi, ws, wf = await _seed_directory( - ws_dir, is_ws=True, hmac_key=hmac_key, now_iso=now_iso - ) - logger.info("WS sessions: %d imported, %d skipped, %d failed", wi, ws, wf) - else: - wi = ws = wf = 0 - - # ---- Autosaves (no HMAC) ----------------------------------------- - auto_dir = Path(AUTOSAVE_DIR) - logger.info("Seeding autosave sessions from %s ...", auto_dir) - ai, as_, af = await _seed_directory( - auto_dir, is_ws=False, hmac_key=None, now_iso=now_iso - ) - logger.info("Autosave sessions: %d imported, %d skipped, %d failed", ai, as_, af) - - total_imported = wi + ai - total_skipped = ws + as_ - total_failed = wf + af - logger.info( - "\u2713 Seeding complete: %d imported, %d skipped, %d failed", - total_imported, - total_skipped, - total_failed, - ) - - # Write the one-time migration marker if seed succeeded with no failures. - if total_failed == 0: - try: - _SEEDER_DONE_MARKER.parent.mkdir(parents=True, exist_ok=True) - _SEEDER_DONE_MARKER.write_text( - f"Seeder migration completed. imported={total_imported} skipped={total_skipped}\n" - ) - logger.info("✓ Seeder migration marker written to %s", _SEEDER_DONE_MARKER) - except Exception as exc: - logger.warning("Could not write seeder marker: %s", exc) + """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/routers/sessions.py b/code_puppy/api/routers/sessions.py index bf046764c..89fc1a1a5 100644 --- a/code_puppy/api/routers/sessions.py +++ b/code_puppy/api/routers/sessions.py @@ -2,7 +2,6 @@ import asyncio import json -import pickle from concurrent.futures import ThreadPoolExecutor from pathlib import Path from typing import Any, Dict, List, Optional @@ -10,10 +9,9 @@ from fastapi import APIRouter, HTTPException from pydantic import BaseModel -# Thread pool for blocking file I/O -_executor = ThreadPoolExecutor(max_workers=2) +from code_puppy.session_storage import get_session_file_path, load_session -# Timeout for file operations (seconds) +_executor = ThreadPoolExecutor(max_workers=2) FILE_IO_TIMEOUT = 10.0 router = APIRouter() @@ -45,63 +43,25 @@ class SessionDetail(SessionInfo): def _get_sessions_dir() -> Path: - """Get the subagent sessions directory. - - Returns: - Path to the subagent sessions directory - """ from code_puppy.config import DATA_DIR return Path(DATA_DIR) / "subagent_sessions" def _serialize_message(msg: Any) -> Dict[str, Any]: - """Serialize a pydantic-ai message to a 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. - - Args: - msg: A pydantic-ai message object or wrapped message dict - - Returns: - JSON-serializable dictionary representation of the message - """ - 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 {k: _serialize_obj(v) for k, v in obj.__dict__.items()} 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 { @@ -111,7 +71,6 @@ def _serialize_obj(obj: Any) -> Any: "ts": msg.get("ts"), } - # Handle unwrapped messages result = _serialize_obj(msg) if not isinstance(result, dict): return {"content": str(result)} @@ -119,35 +78,26 @@ def _serialize_obj(obj: Any) -> Any: def _load_json_sync(file_path: Path) -> dict: - """Synchronous JSON file load (for use in executor).""" - with open(file_path, "r") as f: + with open(file_path, "r", encoding="utf-8") as f: return json.load(f) -def _load_pickle_sync(file_path: Path) -> Any: - """Synchronous pickle file load (for use in executor).""" - with open(file_path, "rb") as f: - return pickle.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]: - """List all available sessions. - - Returns: - List of SessionInfo objects for each session found - """ sessions_dir = _get_sessions_dir() if not sessions_dir.exists(): return [] loop = asyncio.get_running_loop() - sessions = [] + sessions: list[SessionInfo] = [] for txt_file in sessions_dir.glob("*.txt"): session_id = txt_file.stem try: - # Run blocking I/O in thread pool with timeout metadata = await asyncio.wait_for( loop.run_in_executor(_executor, _load_json_sync, txt_file), timeout=FILE_IO_TIMEOUT, @@ -163,10 +113,8 @@ async def list_sessions() -> List[SessionInfo]: ) ) except asyncio.TimeoutError: - # Timed out reading file, include basic info sessions.append(SessionInfo(session_id=session_id)) except Exception: - # If we can't parse metadata, still include basic session info sessions.append(SessionInfo(session_id=session_id)) return sessions @@ -174,17 +122,6 @@ async def list_sessions() -> List[SessionInfo]: @router.get("/{session_id}") async def get_session(session_id: str) -> SessionInfo: - """Get session metadata. - - Args: - session_id: The session identifier - - Returns: - SessionInfo with metadata for the specified session - - Raises: - HTTPException: 404 if session not found, 504 on timeout - """ sessions_dir = _get_sessions_dir() txt_file = sessions_dir / f"{session_id}.txt" @@ -192,7 +129,6 @@ async def get_session(session_id: str) -> SessionInfo: 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), @@ -213,28 +149,16 @@ async def get_session(session_id: str) -> SessionInfo: @router.get("/{session_id}/messages") async def get_session_messages(session_id: str) -> List[Dict[str, Any]]: - """Get the full message history for a session. - - Args: - session_id: The session identifier - - Returns: - List of serialized message dictionaries - - Raises: - HTTPException: 404 if session messages not found, 500 on load error, 504 on timeout - """ sessions_dir = _get_sessions_dir() - pkl_file = sessions_dir / f"{session_id}.pkl" + session_file = get_session_file_path(sessions_dir, session_id) - if not pkl_file.exists(): + 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_pickle_sync, pkl_file), + loop.run_in_executor(_executor, _load_session_history_sync, session_file), timeout=FILE_IO_TIMEOUT, ) return [_serialize_message(msg) for msg in messages] @@ -248,27 +172,16 @@ async def get_session_messages(session_id: str) -> List[Dict[str, Any]]: @router.delete("/{session_id}") async def delete_session(session_id: str) -> Dict[str, str]: - """Delete a session and its data. - - Args: - session_id: The session identifier - - Returns: - Success message dict - - Raises: - HTTPException: 404 if session not found - """ sessions_dir = _get_sessions_dir() txt_file = sessions_dir / f"{session_id}.txt" - pkl_file = sessions_dir / f"{session_id}.pkl" + session_file = get_session_file_path(sessions_dir, session_id) - if not txt_file.exists() and not pkl_file.exists(): + 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 pkl_file.exists(): - pkl_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 index 8a5cbb445..44dd07bfe 100644 --- a/code_puppy/api/routers/ws_sessions.py +++ b/code_puppy/api/routers/ws_sessions.py @@ -222,12 +222,8 @@ async def get_ws_session_tool_calls(session_name: str) -> List[Dict[str, Any]]: @router.delete("/{session_name}") async def delete_ws_session(session_name: str) -> Dict[str, str]: - """Soft-delete a WebSocket session in SQLite. - - Also attempts to clean up any legacy pkl files if they exist. - """ - ws_dir = _get_ws_sessions_dir() - safe_base = _validate_session_name(session_name, ws_dir) + """Soft-delete a WebSocket session in SQLite.""" + _validate_session_name(session_name, _get_ws_sessions_dir()) # Soft delete in SQLite try: @@ -239,17 +235,6 @@ async def delete_ws_session(session_name: str) -> Dict[str, str]: logger.error("delete_ws_session: SQLite error for %s: %s", session_name, e) raise HTTPException(503, "Service unavailable: database error") - # Cleanup legacy files (best effort) - try: - pkl_file = ws_dir / f"{safe_base}.pkl" - meta_file = ws_dir / f"{safe_base}_meta.json" - if pkl_file.exists(): - pkl_file.unlink() - if meta_file.exists(): - meta_file.unlink() - except Exception: - pass - return {"message": f"WebSocket session '{session_name}' deleted"} diff --git a/code_puppy/api/session_cache.py b/code_puppy/api/session_cache.py index 235041418..bbcd4e858 100644 --- a/code_puppy/api/session_cache.py +++ b/code_puppy/api/session_cache.py @@ -6,7 +6,6 @@ import asyncio import logging -import pickle import time from collections import OrderedDict from concurrent.futures import ThreadPoolExecutor @@ -81,7 +80,7 @@ def __init__(self, max_size: int = MAX_CACHED_SESSIONS): self._stats = {"hits": 0, "misses": 0, "evictions": 0, "expirations": 0} async def get( - self, session_id: str, pkl_path: Path + self, session_id: str, session_file_path: Path ) -> Optional[Tuple[List[Dict[str, Any]], Dict[str, Any]]]: """Get cached session messages (pre-serialized JSON). @@ -104,7 +103,7 @@ async def get( # Check if file was modified try: - current_mtime = pkl_path.stat().st_mtime + current_mtime = session_file_path.stat().st_mtime if current_mtime > entry.file_mtime: self._stats["expirations"] += 1 del self._cache[session_id] @@ -131,7 +130,7 @@ async def get( async def put( self, session_id: str, - pkl_path: Path, + session_file_path: Path, messages: List[Any], json_messages: List[Dict[str, Any]], metadata: Dict[str, Any], @@ -140,7 +139,7 @@ async def put( Args: session_id: Unique session identifier - pkl_path: Path to the pickle file (for mtime tracking) + 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 @@ -156,7 +155,7 @@ async def put( # Get file mtime try: - file_mtime = pkl_path.stat().st_mtime + file_mtime = session_file_path.stat().st_mtime except FileNotFoundError: file_mtime = time.time() @@ -214,10 +213,11 @@ def get_session_cache() -> SessionCache: return _session_cache -def _load_pickle_sync(pkl_path: Path) -> List[Any]: - """Synchronous pickle loading (for executor).""" - with open(pkl_path, "rb") as f: - return pickle.load(f) +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]: @@ -284,7 +284,7 @@ def _serialize_messages_sync(messages: List[Any]) -> List[Dict[str, Any]]: async def load_session_cached( - session_id: str, pkl_path: Path, timeout: float = 10.0 + session_id: str, session_file_path: Path, timeout: float = 10.0 ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]: """Load session messages with caching. @@ -293,33 +293,34 @@ async def load_session_cached( Args: session_id: Unique session identifier - pkl_path: Path to the pickle file + session_file_path: Path to the session JSON file timeout: Timeout for file operations Returns: Tuple of (json_messages, metadata) Raises: - FileNotFoundError: If pickle file doesn't exist + 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, pkl_path) + cached = await cache.get(session_id, session_file_path) if cached is not None: return cached # Cache miss - load from disk - if not pkl_path.exists(): - raise FileNotFoundError(f"Session file not found: {pkl_path}") + if not session_file_path.exists(): + raise FileNotFoundError(f"Session file not found: {session_file_path}") loop = asyncio.get_running_loop() - # Load pickle in thread pool + # Load session file in thread pool start_time = time.time() messages = await asyncio.wait_for( - loop.run_in_executor(_executor, _load_pickle_sync, pkl_path), timeout=timeout + loop.run_in_executor(_executor, _load_session_sync, session_file_path), + timeout=timeout, ) load_time = time.time() - start_time @@ -335,7 +336,7 @@ async def load_session_cached( } # Cache for next time - await cache.put(session_id, pkl_path, messages, json_messages, metadata) + 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" diff --git a/code_puppy/api/session_context.py b/code_puppy/api/session_context.py index a8621e1bf..f341f735e 100644 --- a/code_puppy/api/session_context.py +++ b/code_puppy/api/session_context.py @@ -441,7 +441,7 @@ async def _write_to_sqlite( async def save_session(self, session_id: str) -> None: """Persist session history and metadata to SQLite (single source of truth). - Previously wrote to {WS_SESSION_DIR}/{session_id}.pkl — that pkl write + 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. @@ -582,7 +582,7 @@ async def load_session(self, session_id: str) -> Optional[SessionContext]: self._sessions[session_id] = ctx return ctx - # SQLite is the single source of truth — no pkl fallback. + # 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 diff --git a/code_puppy/api/templates/chat.html b/code_puppy/api/templates/chat.html index b23d44951..0ab87812a 100644 --- a/code_puppy/api/templates/chat.html +++ b/code_puppy/api/templates/chat.html @@ -233,6 +233,43 @@

Code Puppy Chat

const workingDirInput = document.getElementById('working-dir-input'); const btnSetDir = document.getElementById('btn-set-dir'); const currentWorkingDirEl = document.getElementById('current-working-dir'); + const sessionState = { + sessionId: CONFIG.sessionId, + title: '', + messageCount: null, + totalTokens: null, + pinned: false, + workingDirectory: '', + config: {} + }; + window.codePuppySessionState = sessionState; + + function renderSessionInfo() { + const parts = []; + if (sessionState.sessionId) parts.push(`Session: ${sessionState.sessionId}`); + if (sessionState.title) parts.push(sessionState.title); + if (Number.isInteger(sessionState.messageCount)) parts.push(`${sessionState.messageCount} msgs`); + if (Number.isInteger(sessionState.totalTokens)) parts.push(`${sessionState.totalTokens} tokens`); + sessionInfo.textContent = parts.length ? parts.join(' • ') : 'Connected'; + } + + function applyWorkingDirectory(directory, error = null) { + const normalized = directory || ''; + workingDirInput.value = normalized; + sessionState.workingDirectory = normalized; + if (error) { + currentWorkingDirEl.textContent = `Error: ${error}`; + currentWorkingDirEl.className = 'font-semibold text-red-400 text-xs mt-1'; + return; + } + if (normalized) { + currentWorkingDirEl.textContent = normalized; + currentWorkingDirEl.className = 'font-semibold text-green-400 text-xs mt-1'; + } else { + currentWorkingDirEl.textContent = 'Not set'; + currentWorkingDirEl.className = 'font-semibold text-yellow-400 text-xs mt-1'; + } + } // Initialize document.addEventListener('DOMContentLoaded', () => { @@ -364,6 +401,7 @@

Code Puppy Chat

setStatus('connected'); reconnectAttempts = 0; console.log('Connected to chat WebSocket'); + renderSessionInfo(); loadAgentsAndModels(); }; @@ -405,6 +443,21 @@

Code Puppy Chat

case 'system': handleSystemMessage(message); break; + case 'session_meta': + handleSessionMeta(message); + break; + case 'session_meta_updated': + handleSessionMetaUpdated(message); + break; + case 'session_restored': + handleSessionRestored(message); + break; + case 'session_switched': + handleSessionSwitched(message); + break; + case 'config_value': + handleConfigValue(message); + break; case 'working_directory_changed': handleWorkingDirectoryChanged(message); break; @@ -445,7 +498,9 @@

Code Puppy Chat

} function handleSystemMessage(message) { - sessionInfo.textContent = message.content; + if (message.content && (!sessionState.sessionId || sessionInfo.textContent === 'Connecting...')) { + sessionInfo.textContent = message.content; + } if (message.agent_name) { agentSelector.value = message.agent_name; } @@ -455,14 +510,84 @@

Code Puppy Chat

addSystemMessage(message.content); } + function handleSessionMeta(message) { + if (message.session_id) { + sessionState.sessionId = message.session_id; + CONFIG.sessionId = message.session_id; + } + sessionState.title = message.title || ''; + sessionState.messageCount = Number.isInteger(message.message_count) ? message.message_count : null; + sessionState.totalTokens = Number.isInteger(message.total_tokens) ? message.total_tokens : null; + if (typeof message.pinned === 'boolean') { + sessionState.pinned = message.pinned; + } + if (message.agent_name) { + agentSelector.value = message.agent_name; + } + if (message.model_name) { + modelSelector.value = message.model_name; + } + applyWorkingDirectory(message.working_directory || ''); + renderSessionInfo(); + } + + function handleSessionMetaUpdated(message) { + if (typeof message.title === 'string') { + sessionState.title = message.title; + } + if (typeof message.pinned === 'boolean') { + sessionState.pinned = message.pinned; + } + renderSessionInfo(); + } + + function handleSessionRestored(message) { + if (message.session_id) { + sessionState.sessionId = message.session_id; + CONFIG.sessionId = message.session_id; + } + if (typeof message.title === 'string') { + sessionState.title = message.title; + } + if (Number.isInteger(message.message_count)) { + sessionState.messageCount = message.message_count; + } + renderSessionInfo(); + } + + function handleSessionSwitched(message) { + if (message.session_id) { + sessionState.sessionId = message.session_id; + CONFIG.sessionId = message.session_id; + } + if (typeof message.title === 'string') { + sessionState.title = message.title; + } + if (Number.isInteger(message.message_count)) { + sessionState.messageCount = message.message_count; + } + if (message.agent_name) { + agentSelector.value = message.agent_name; + } + if (message.model_name) { + modelSelector.value = message.model_name; + } + applyWorkingDirectory(message.working_directory || ''); + renderSessionInfo(); + } + + function handleConfigValue(message) { + if (!message.key) return; + sessionState.config[message.key] = message.value; + window.codePuppyConfig = sessionState.config; + } + function handleWorkingDirectoryChanged(message) { if (message.success) { - currentWorkingDirEl.textContent = message.directory; - currentWorkingDirEl.className = 'font-semibold text-green-400 text-xs mt-1'; + applyWorkingDirectory(message.directory || ''); addSystemMessage(`Working directory set to: ${message.directory}`); } else { - currentWorkingDirEl.textContent = `Error: ${message.error}`; - currentWorkingDirEl.className = 'font-semibold text-red-400 text-xs mt-1'; + applyWorkingDirectory(message.directory || '', message.error || 'Unknown error'); addSystemMessage(`Failed to set working directory: ${message.error}`, 'error'); } } 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_session_persistence.py b/code_puppy/api/tests/test_session_persistence.py index b97cbda54..2a5b828b4 100644 --- a/code_puppy/api/tests/test_session_persistence.py +++ b/code_puppy/api/tests/test_session_persistence.py @@ -132,8 +132,6 @@ def test_session_update_action_created_for_first_message(): ) assert payload["action"] == "created" assert payload["auto_saved"] is True - assert payload["pickle_path"] == "" - assert payload["metadata_path"] == "" assert payload["timestamp"] == "2026-01-01T00:00:00" @@ -169,8 +167,6 @@ def test_session_update_payload_has_expected_keys(): "message_count", "total_tokens", "auto_saved", - "pickle_path", - "metadata_path", "action", } assert set(payload.keys()) == expected_keys 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 index c04f2d9aa..94fa92267 100644 --- a/code_puppy/api/ws/chat_handler.py +++ b/code_puppy/api/ws/chat_handler.py @@ -19,69 +19,66 @@ import datetime import json import logging -import re -import uuid -from pathlib import Path from fastapi import FastAPI, WebSocket, WebSocketDisconnect from pydantic import TypeAdapter, ValidationError -from code_puppy.api.db.queries import ( - update_session_working_directory, - write_system_message_to_sqlite, +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.session_context import _validate_session_id, session_manager -from code_puppy.api.ws.attachments import build_file_context_and_attachments -from code_puppy.api.ws.background_save import ( - fire_and_track, - save_agent_result_in_background, +from code_puppy.api.ws.ws_post_run import resolve_post_run_resolution +from code_puppy.api.ws.ws_turn_finalization import ( + emit_pre_stream_end_tool_results, + finalize_turn_history, ) -from code_puppy.api.ws.connection_manager import connection_manager -from code_puppy.api.ws.history_utils import ( - build_enhanced_history, - estimate_total_tokens, +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, - build_error_response_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 ( - PROTOCOL_VERSION, ClientMessage, ServerAgentInvoked, ServerAssistantMessageDelta, ServerAssistantMessageEnd, ServerAssistantMessageStart, ServerCancelled, - ServerCommandResult, - ServerConfigValue, ServerError, - ServerSessionMetaUpdated, - ServerSessionRestored, - ServerSessionSwitched, ServerStatus, ServerStreamEnd, - ServerSystem, - ServerToolCall, - ServerToolResult, ServerUserMessage, - ServerWorkingDirectoryChanged, ) from code_puppy.api.ws.send_utils import WebSocketSender -from code_puppy.api.ws.session_persistence import ( - build_session_meta_payload, - build_session_update_payload, - persist_turn_to_sqlite, - resolve_agent_model_meta, +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.session_storage import generate_heuristic_title -from code_puppy.tools.command_runner import ( - cleanup_session_process_tracking, - init_session_process_tracking, -) +from code_puppy.tools.command_runner import cleanup_session_process_tracking # Import session context for desk-puppy working directory support try: @@ -160,224 +157,52 @@ async def websocket_chat( send_typed_tool_lifecycle = sender.send_typed_tool_lifecycle persist_error_payload = sender.persist_error_payload - ctx = None - session_title = "" - session_working_directory = "" # Will be set by client or loaded from metadata - session_pinned = False # Track pinned state for persistence - last_context_sent_directory = "" # Track when we last sent directory context - existing_history = None - - # Use provided session_id or generate new one - if session_id: - # Validate session_id to prevent path traversal - try: - _validate_session_id(session_id) - except ValueError as e: - logger.warning("Invalid session_id rejected: %r: %s", session_id, e) - await websocket.close(code=1008, reason="Invalid session ID") - return - - # Client wants to resume/continue a specific session - logger.debug("Client requested session: %s", session_id) - - # SQLite is the source of truth — check session existence there. - 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 e: - logger.warning("Failed to check session in SQLite: %s", e) - - if not existing_history: - logger.debug( - "Session %s not found in SQLite, will create new", session_id - ) - else: - # Generate new timestamp-based session ID - 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) - - try: - # --- SESSION ISOLATION: create or load via SessionManager --- - try: - if existing_history is not None: - # Session confirmed in SQLite — load it via SessionManager. - # get_or_load_session checks in-memory first, then falls - # back to a fresh SQLite load. SQLite is the sole source of truth. - ctx = await session_manager.get_or_load_session(session_id) - sender.ctx = ctx - if ctx is None: - # SQLite confirmed this session exists but load returned None. - # This is unexpected (DB race, corrupt row, or schema mismatch). - # Log a warning and start fresh so the WebSocket can still - # function — silent data loss is worse than a blank session. - 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, - ) - ctx = await session_manager.create_session(session_id) - sender.ctx = ctx - else: - # Update local state from loaded metadata - session_title = ctx.title - session_working_directory = ctx.working_directory - session_pinned = ctx.pinned - else: - ctx = await session_manager.create_session(session_id) - sender.ctx = ctx - except Exception as e: - logger.warning("SessionManager init failed, falling back: %s", e) - try: - ctx = await session_manager.create_session(session_id) - sender.ctx = ctx - except Exception: - logger.error("SessionManager fallback also failed", exc_info=True) - await websocket.close(code=1011, reason="Session init failed") - return - - # Mark session as active (cancels any pending cleanup from previous disconnect) - await session_manager.mark_session_active(session_id) - - # Initialize per-session process tracking (ContextVar isolation) - init_session_process_tracking() - - # Set MessageBus session context for this WS connection - try: - bus = get_message_bus() - bus.set_session_context(session_id) - except Exception: - logger.debug("MessageBus session context not available") - - # Convenience aliases — use ctx.agent everywhere instead of get_current_agent() - agent = ctx.agent - agent_name = ctx.agent_name - model_name = ctx.model_name - - # ── Persist initial session context to SQLite (new sessions only) ─────── - # Write a config row (agent + model) and, if a CWD is set, a directory - # row so the FE cold-load path sees them before the first user message. - # Skipped for resumed sessions — their init rows were written at creation. - if existing_history is None: - try: - _now_iso = datetime.datetime.now(datetime.timezone.utc).isoformat() - _init_agent = agent_name or "code-puppy" - _init_model = model_name or "unknown" - await write_system_message_to_sqlite( - session_id=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 session_working_directory: - _path_segments = session_working_directory.split("/")[-3:] - _relative = "/".join(s for s in _path_segments if s) - await write_system_message_to_sqlite( - session_id=session_id, - system_message_type="directory", - content=f"Starting in {_relative}", - system_message_path=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, - ) - - # Send welcome message - await send_typed( - ServerSystem( - content=f"Connected! Session: {session_id}", - session_id=session_id, - agent_name=agent_name, - model_name=model_name, - resumed=existing_history is not None, - protocol_version=PROTOCOL_VERSION, - ) + 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 + + 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, ) - # If we resumed an existing session, notify the client - if existing_history and ctx: - try: - # History is already loaded into ctx.agent by session_manager.load_session(). - # Pull it back out for the client-side UI metadata notification. - loaded_messages = ctx.agent.get_message_history() - message_count = len(loaded_messages) if loaded_messages else 0 - - # Notify client about restored session - await send_typed( - ServerSessionRestored( - session_id=session_id, - message_count=message_count, - title=session_title, - ui_metadata=[], - ) - ) - - # Replay system messages (agent/model switches, CWD banners) to UI. - # These are persisted with pydantic_json=NULL so _load_from_sqlite - # skips them — we must send them separately so the UI shows them. - try: - from code_puppy.api.db.queries import get_active_messages - - rows = await get_active_messages(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=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), - session_id, - ) - except Exception as sys_exc: - logger.warning( - "Failed to replay system messages for session %s: %s", - session_id, - sys_exc, - ) - - agent = ctx.agent # Refresh local alias - logger.debug("Restored %d messages to session agent", message_count) - except Exception as e: - logger.warning("Failed to restore session history: %s", e) - # Track active streaming task for cancellation - active_drain_task = None - # Track active agent task (run_with_mcp) for cancellation - active_agent_task: asyncio.Task | None = None - # Track stop_draining event for cancellation across iterations - stop_draining = asyncio.Event() + await _send_session_meta_snapshot( + runtime=runtime, + safe_send_json=safe_send_json, + ) + try: while True: try: msg = await websocket.receive_json() @@ -396,671 +221,53 @@ async def websocket_chat( }, ) - if msg.get("type") == "switch_agent": - agent_name = msg.get("agent_name") - if agent_name: - try: - new_agent = await session_manager.switch_agent( - session_id, agent_name - ) - agent = new_agent # Update local alias - ctx.agent = new_agent - model_name = ( - new_agent.get_model_name() - if new_agent - else "unknown" - ) - - # Persist to SQLite so the FE cold-load path sees it - try: - await write_system_message_to_sqlite( - session_id=session_id, - system_message_type="config", - content=f"🔄 Switched to {agent_name} ({model_name})", - agent_name=agent_name, - model_name=model_name, - ) - except Exception as _sw_exc: - logger.warning( - "Agent-switch SQLite write failed: %s", - _sw_exc, - exc_info=True, - ) - - await send_typed( - ServerSystem( - content=f"🔄 Switched to {agent_name} ({model_name})", - session_id=session_id, - agent_name=agent_name, - model_name=model_name, - ) - ) - except Exception as e: - logger.error("Error switching agent: %s", e) - await send_typed( - ServerError( - error=f"Failed to switch to agent {agent_name}: {str(e)}", - session_id=session_id, - ) - ) - - elif msg.get("type") == "switch_model": - model_name = msg.get("model_name") or msg.get("model") - if model_name: - try: - await session_manager.switch_model( - session_id, model_name - ) - agent = ctx.agent # Refresh local alias - logger.debug("Switched model to: %s", model_name) - - # Persist to SQLite so the FE cold-load path sees it - _sw_agent = ctx.agent_name or agent_name or "code-puppy" - try: - await write_system_message_to_sqlite( - session_id=session_id, - system_message_type="config", - content=f"🔄 Switched to {_sw_agent} ({model_name})", - agent_name=_sw_agent, - model_name=model_name, - ) - except Exception as _sw_exc: - logger.warning( - "Model-switch SQLite write failed: %s", - _sw_exc, - exc_info=True, - ) - - await send_typed( - ServerSystem( - content=f"🔄 Switched to {_sw_agent} ({model_name})", - session_id=session_id, - model_name=model_name, - agent_name=_sw_agent, - ) - ) - except Exception as e: - logger.error("Error switching model: %s", e) - await send_typed( - ServerError( - error=f"Failed to switch to model {model_name}: {str(e)}", - session_id=session_id, - ) - ) - else: - await send_typed( - ServerError( - error="No model_name provided for switch_model", - session_id=session_id, - ) - ) - - elif msg.get("type") == "switch_session": - """ - Switch to a different session without reconnecting WebSocket. - - How it works: - 1. Client sends: {"type": "switch_session", "session_id": "WS_session_20260120_124356"} - 2. Server loads session from ~/.code_puppy/ws_sessions/{session_id}.pkl - 3. Server restores message history to the current agent - 4. Server responds with session_switched event containing message count and title - 5. Client can continue chatting with full conversation context - - This approach keeps a single WebSocket connection alive while allowing - users to switch between multiple sessions instantly. All session data - stays on the local filesystem for privacy. - """ - # Cancel any active streaming first - if active_drain_task and not active_drain_task.done(): - logger.debug( - "Cancelling active streaming due to session switch" - ) - stop_draining.set() # Signal drain to stop - active_drain_task.cancel() - try: - await active_drain_task - except asyncio.CancelledError: - pass - stop_draining.clear() # Reset for next streaming - - 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=session_id, - ) - ) - continue - - # Validate session_id to prevent path traversal - try: - _validate_session_id(new_session_id) - except ValueError as e: - logger.warning( - f"Invalid session_id in switch_session: {new_session_id!r}" - ) - await send_typed( - ServerError( - error=f"Invalid session ID: {e}", - session_id=session_id, - ) - ) - continue - - logger.debug("Switching to session: %s", new_session_id) - - try: - # Check if target session exists in SQLite - try: - from code_puppy.api.db.queries import ( - session_exists as _se, - ) - - _target_exists = await _se(new_session_id) - except Exception: - _target_exists = False - - if not _target_exists: - # Session doesn't exist - create new one with this ID - logger.debug( - f"Session {new_session_id} not found, creating new" - ) - - # Save current session and mark inactive (keep in memory for 15 min) - try: - await session_manager.save_session(session_id) - except Exception: - pass - await session_manager.mark_session_inactive(session_id) - - session_id = new_session_id - sender.session_id = session_id - session_title = "" - session_working_directory = "" - session_pinned = False - - ctx = await session_manager.create_session(session_id) - sender.ctx = ctx - # Mark the new session as active - await session_manager.mark_session_active(session_id) - agent = ctx.agent - agent_name = ctx.agent_name - model_name = ctx.model_name - - await send_typed( - ServerSessionSwitched( - session_id=new_session_id, - message_count=0, - title="", - created=True, - agent_name=agent_name, - model_name=model_name, - ) - ) - continue - - # Load existing session via SessionManager - # Save current session and mark inactive (keep in memory for 15 min) - try: - await session_manager.save_session(session_id) - except Exception: - pass - await session_manager.mark_session_inactive(session_id) - - session_id = new_session_id - sender.session_id = session_id - - # Try loading via SessionManager (checks in-memory first, then SQLite) - loaded_ctx = await session_manager.get_or_load_session( - session_id - ) - if loaded_ctx is not None: - ctx = loaded_ctx - session_title = ctx.title - session_working_directory = ctx.working_directory - session_pinned = ctx.pinned - else: - # Legacy session or HMAC missing — load metadata - # from JSON (safe) and create fresh context - new_title = "" - new_working_directory = "" - new_pinned = False - try: - from code_puppy.api.db.queries import ( - get_session_metadata as _gsm, - ) - - _sm = await _gsm(session_id) or {} - new_title = _sm.get("title", "") - new_working_directory = _sm.get( - "working_directory", "" - ) - new_pinned = bool(_sm.get("pinned", False)) - except Exception: - pass - - ctx = await session_manager.create_session(session_id) - sender.ctx = ctx - ctx.title = new_title - ctx.working_directory = new_working_directory - ctx.pinned = new_pinned - session_title = new_title - session_working_directory = new_working_directory - session_pinned = new_pinned - - # Mark the new session as active - await session_manager.mark_session_active(session_id) - - agent = ctx.agent - agent_name = ctx.agent_name - model_name = ctx.model_name - message_count = len(ctx.agent.get_message_history() or []) - logger.debug( - f"Restored {message_count} messages to session agent" - ) - - await send_typed( - ServerSessionSwitched( - session_id=new_session_id, - message_count=message_count, - title=session_title, - working_directory=session_working_directory, - created=False, - agent_name=agent_name, - model_name=model_name, - ) - ) - - logger.debug( - f"Switched to session {new_session_id} with {message_count} messages" - ) - - except Exception as e: - logger.error("Error switching session: %s", e) - await send_typed( - ServerError( - error=f"Failed to switch to session {new_session_id}: {str(e)}", - session_id=session_id, - ) - ) - - elif msg.get("type") == "set_working_directory": - new_directory = msg.get("directory", "") - if new_directory: - # Expand ~ and resolve symlinks before validation - new_directory = str( - Path(new_directory).expanduser().resolve() - ) - logger.info( - "[CWD DEBUG] set_working_directory received: new=%r, current=%r, session=%s", - new_directory, - session_working_directory, - session_id, - ) - # Validate the directory exists - if Path(new_directory).is_dir(): - # Skip if directory hasn't actually changed (avoids duplicate banners on reload) - if new_directory == session_working_directory: - logger.info( - "[CWD DEBUG] Skipping unchanged directory: %r, session=%s", - new_directory, - session_id, - ) - await send_typed( - ServerWorkingDirectoryChanged( - directory=new_directory, - success=True, - session_id=session_id, - unchanged=True, - ) - ) - continue # Skip to next message, don't write banner - - session_working_directory = new_directory - logger.debug( - "Working directory set to: %s", - session_working_directory, - ) - - # NOTE: Do NOT append raw dict entries into agent - # message history here. The runtime expects typed - # ModelMessage objects; dict injection can corrupt - # subsequent turns and lead to result=None failures. - - # Persist directory banner to SQLite so FE cold-load sees it - try: - # Bug 4: use ctx.agent_name / ctx.model_name directly - # so this works even when ctx.agent is None. - _cwd_agent = ( - ctx.agent_name if ctx else None - ) or "code-puppy" - _cwd_model = ( - ctx.model_name if ctx else None - ) or "unknown" - _cwd_segs = 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 _gpn - - _pname = _gpn() or "puppy" - logger.info( - "[CWD DEBUG] Writing CWD banner to SQLite: path=%r, session=%s", - session_working_directory, - session_id, - ) - await write_system_message_to_sqlite( - session_id=session_id, - system_message_type="directory", - content=f"{_pname} is now at {_cwd_rel}", - system_message_path=session_working_directory, - agent_name=_cwd_agent, - model_name=_cwd_model, - ) - # Bug 5: also stamp working_directory on the sessions row - _now_cwd = datetime.datetime.now( - datetime.timezone.utc - ).isoformat() - await update_session_working_directory( - session_id=session_id, - working_directory=session_working_directory, - updated_at=_now_cwd, - ) - except Exception as _cwd_exc: - logger.warning( - "CWD SQLite write failed: %s", - _cwd_exc, - exc_info=True, - ) - - await send_typed( - ServerWorkingDirectoryChanged( - directory=session_working_directory, - success=True, - session_id=session_id, - ) - ) - else: - await send_typed( - ServerWorkingDirectoryChanged( - directory=new_directory, - success=False, - error="Directory does not exist", - session_id=session_id, - ) - ) - else: - await send_typed( - ServerError( - error="No directory provided for set_working_directory", - session_id=session_id, - ) - ) - - elif msg.get("type") == "update_session_meta": - # Update session metadata (pinned, title) — persisted to SQLite. - try: - # Update in-memory state (local vars AND SessionContext) - if "pinned" in msg and isinstance(msg["pinned"], bool): - session_pinned = msg["pinned"] - if ctx is not None: - ctx.pinned = session_pinned - if "title" in msg and isinstance(msg["title"], str): - session_title = msg["title"] - if ctx is not None: - ctx.title = session_title - - # Persist to SQLite (single source of truth). - # Use a targeted UPDATE — NOT upsert_session() — so that - # message_count, total_tokens, and other stats are - # never zeroed out by a rename/pin operation. - 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=session_id, - title=session_title, - pinned=session_pinned, - updated_at=_dt.now(_tz.utc).isoformat(), - ) - logger.debug( - f"Updated session meta in SQLite for {session_id}: pinned={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=session_id, - pinned=session_pinned, - title=session_title, - ) - ) - except Exception as e: - logger.error("Error updating session meta: %s", e) - await send_typed( - ServerError( - error=f"Failed to update session metadata: {str(e)}", - session_id=session_id, - ) - ) - - elif msg.get("type") == "get_config": - 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=session_id, - ) - ) - else: - await send_typed( - ServerError( - error="No key provided for get_config", - session_id=session_id, - ) - ) - - elif msg.get("type") == "set_config": - 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( - f"Config set: {config_key} = {config_value}" - ) - await send_typed( - ServerConfigValue( - key=config_key, - value=config_value, - success=True, - session_id=session_id, - ) - ) - except Exception as e: - await send_typed( - ServerError( - error=f"Failed to set config: {e}", - session_id=session_id, - ) - ) - else: - await send_typed( - ServerError( - error="No key provided for set_config", - session_id=session_id, - ) - ) - - # Handle slash command execution - elif msg.get("type") == "command": - command_str = msg.get("command", "") - logger.debug("Command requested: %s", command_str) - - try: - from io import StringIO - - from rich.console import Console - - from code_puppy.command_line.command_handler import ( - get_commands_help, - handle_command, - ) - - output = None - captured_messages = [] - - # Special handling for /help - we can get the text directly - cmd_name = ( - command_str.strip().lstrip("/").split()[0] - if command_str.strip() - else "" - ) - if cmd_name in ("help", "h"): - # Get help text directly - help_text = get_commands_help() - # Render to plain text - 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: - # For other commands, just execute and report success/failure - # The actual output goes to the message queue/terminal - result = handle_command(command_str) - # Some commands might return a string as output - if isinstance(result, str): - output = result - result = True - - # Send command result - await send_typed( - ServerCommandResult( - command=command_str, - success=result is True or result is not False, - output=output, - messages=captured_messages, - result=str(result) - if result and result is not True - else None, - session_id=session_id, - ) - ) - logger.debug( - f"Command executed: {command_str} -> success={result is True or result is not False}, output_len={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, - ) - ) - continue - - # Handle cancel/interrupt request - elif msg.get("type") == "cancel": - logger.debug( - "Cancel request received - stopping active streaming and agent task" - ) - - # Cancel any active streaming - if active_drain_task and not active_drain_task.done(): - stop_draining.set() # Signal drain to stop - active_drain_task.cancel() - try: - await active_drain_task - except asyncio.CancelledError: - pass - stop_draining.clear() # Reset for next streaming - logger.debug("Active streaming cancelled") - - # Cancel the in-flight agent run (run_with_mcp) if present - if active_agent_task and not active_agent_task.done(): - logger.debug( - "Cancelling active agent task due to user interrupt" - ) - active_agent_task.cancel() - try: - await active_agent_task - except asyncio.CancelledError: - logger.debug("Active agent task cancelled successfully") - active_agent_task = None - - # Send confirmation - await send_typed( - ServerStatus( - status="cancelled", - session_id=session_id, - ) - ) + if await handle_command_message( + msg=msg, + session_id=session_id, + send_typed=send_typed, + ): continue - elif msg.get("type") == "permission_response": - # Handle permission response from user - from code_puppy.api.permissions import ( - handle_permission_response, + 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 ) - - request_id = msg.get("request_id") - approved = msg.get("approved", False) - - if request_id: - handled = handle_permission_response(request_id, approved) - if handled: - logger.debug( - f"[Permission] ✅ Handled response: {request_id} = {approved}" - ) - else: - logger.warning( - f"[Permission] ❌ Unknown request: {request_id}" - ) - else: - logger.error( - "[WebSocket] ❌ No request_id in permission_response!" - ) + 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 - from code_puppy.api.permission_plugin import ( - set_suppress_emitter_tool_events, - set_websocket_context, + setup_message_context( + websocket=websocket, session_id=session_id ) - set_websocket_context(websocket, session_id) - user_message = msg.get("content", "") # Initialize attachment tracking variables at message scope @@ -1143,27 +350,16 @@ async def websocket_chat( continue # Subscribe to frontend emitter and run drain task CONCURRENTLY - event_queue = None - collected_text = [] + turn_state = WebSocketTurnState() stop_draining.clear() # Reset for this message drain_task = None - active_parts: dict[ - int, dict - ] = {} # Track message parts by index - tool_id_aliases: dict[str, str] = {} - tool_group_ids: dict[ - str, str - ] = {} # tool_id -> tool_group_id mapping - b1_streaming_used = ( - False # Track if B1 streaming sent any content - ) - agent_error = None # Will hold exception if agent run fails + 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.""" - nonlocal collected_text, b1_streaming_used import time as time_module # Capture agent and model metadata at the start. @@ -1180,115 +376,18 @@ async def drain_events_concurrent( or ctx.model_name or "unknown" ) - current_tool_name = None # Track active tool name - current_tool_group_id: str | None = ( - None # Track current tool batch group ID - ) - # Track pending tool calls for tool_result generation - # {tool_id: {tool_name, start_time, part_index}} - pending_tool_calls: dict[str, dict] = {} - - def resolve_pending_tool_id( - *, - 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 pending_tool_calls - ): - return tool_call_id - - if tool_call_id: - for ( - _pending_id, - _pending_info, - ) in 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 pending_tool_calls.items(): - if ( - _pending_info.get("tool_name") - == tool_name - ): - return _pending_id - - return None - - def resolve_tool_group_id( - *, - tool_id: str | None = None, - pending_info: dict | 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. - - Priority order: - 1) pending tool info - 2) tool_id -> group map - 3) explicit fallback passed by caller - 4) current in-flight batch group - 5) synthetic deterministic-ish fallback (last resort) - """ - nonlocal current_tool_group_id - - 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 = 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 = 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, + 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, ) - - if tool_id: - tool_group_ids[tool_id] = group_id - if pending_info is not None: - pending_info["tool_group_id"] = group_id - elif tool_id in pending_tool_calls: - pending_tool_calls[tool_id][ - "tool_group_id" - ] = group_id - - if current_tool_group_id is None: - current_tool_group_id = group_id - - return group_id + ) event_count = 0 first_iteration = True @@ -1313,16 +412,16 @@ def resolve_tool_group_id( try: # Wait for first event with 10ms timeout (blocks efficiently) first_event = await asyncio.wait_for( - event_queue.get(), timeout=0.01 + 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 event_queue.empty(): + while not stream_event_queue.empty(): try: - event = event_queue.get_nowait() + event = stream_event_queue.get_nowait() events_to_send.append(event) event_count += 1 except Exception: @@ -1361,119 +460,24 @@ def resolve_tool_group_id( try: # Handle tool call events if event_type == "tool_call_start": - tool_name = event_data.get( - "tool_name", "unknown" - ) - tool_args = event_data.get( - "tool_args", {} - ) - tool_id = str(uuid.uuid4())[:8] - - # Generate tool group ID for this batch if not already set - if current_tool_group_id is None: - current_tool_group_id = ( - f"tg-{str(uuid.uuid4())[:8]}" - ) - - logger.debug( - "[ws] tool_call: %s", tool_name - ) - - # Track current tool name - current_tool_name = tool_name - - # Register in pending_tool_calls so - # tool_call_complete can echo back the - # same tool_id the frontend already knows. - pending_tool_calls[tool_id] = { - "tool_name": tool_name, - "start_time": time_module.time(), - "tool_group_id": 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=current_agent_name, - model_name=current_model_name, - tool_group_id=current_tool_group_id, - ) + 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": - 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 - ) - - # Clear current tool name after completion - current_tool_name = None - - # Match the tool_id we generated at - # tool_call_start so the frontend can - # correlate result → call by ID, not - # just by name (names are not unique). - matching_tool_id = next( - ( - tid - for tid, info in pending_tool_calls.items() - if info["tool_name"] - == tool_name - ), - None, - ) - matching_pending_info = ( - pending_tool_calls.get( - matching_tool_id - ) - if matching_tool_id - else None - ) - tool_group_id_for_result = resolve_tool_group_id( - tool_id=matching_tool_id, - pending_info=matching_pending_info, - fallback_group_id=current_tool_group_id, - tool_name=tool_name, - source="tool_call_complete", - ) - - if matching_tool_id: - 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=current_agent_name, - model_name=current_model_name, - tool_group_id=tool_group_id_for_result, - ) + 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": @@ -1524,455 +528,38 @@ def resolve_tool_group_id( part_obj = inner_data.get( "part", {} ) - initial_content = "" - - if ( - hasattr(part_obj, "content") - and part_obj.content - ): - initial_content = ( - part_obj.content - ) - elif isinstance( - part_obj, dict - ) and part_obj.get("content"): - initial_content = part_obj.get( - "content", "" - ) - if part_type in ( - "TextPart", - "ThinkingPart", + 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, ): - msg_type = ( - "thinking" - if part_type - == "ThinkingPart" - else "text" - ) - - # When a TextPart starts, any pending tool calls have completed - # Send status-only tool_result during streaming to update UI - # The actual result data will be sent later from extraction code - if ( - pending_tool_calls - and part_type == "TextPart" - ): - for ( - tool_id, - tool_info, - ) in list( - pending_tool_calls.items() - ): - if tool_info.get( - "status_only_sent" - ): - continue - duration_ms = ( - time_module.time() - - tool_info[ - "start_time" - ] - ) * 1000 - logger.debug( - f"[WebSocket] Sending status-only tool_result for: {tool_info['tool_name']} (duration: {duration_ms:.1f}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=current_agent_name, - model_name=current_model_name, - tool_group_id=resolve_tool_group_id( - tool_id=tool_id, - pending_info=tool_info, - fallback_group_id=current_tool_group_id, - tool_name=tool_info.get( - "tool_name" - ), - source="status_only_result", - ), - ) - ) - tool_info[ - "status_only_sent" - ] = True - current_tool_name = None - - # Check if part already exists (created by early delta) - if part_index in active_parts: - # Just update the type, keep existing message_id - active_parts[part_index][ - "type" - ] = msg_type - message_id = active_parts[ - part_index - ]["id"] - # If there's initial content, add it to the existing accumulated content - if initial_content: - active_parts[ - part_index - ]["content"] = ( - initial_content - + active_parts[ - part_index - ]["content"] - ) - collected_text.insert( - 0, initial_content - ) # Prepend to collected text too - logger.debug( - f"[Stream Debug] Part already exists, reusing message_id={message_id}" - ) - # Don't send another start event - continue - - message_id = f"msg-{int(time_module.time() * 1000)}-{part_index}" - active_parts[part_index] = { - "id": message_id, - "type": msg_type, - "content": initial_content, # Start with initial content if present - } - - # If there's initial content, add it to collected_text - if initial_content: - collected_text.append( - initial_content - ) - - # New assistant output part indicates turn boundary for tool grouping - if part_index == 0: - 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=current_agent_name, - model_name=current_model_name, - tool_name=current_tool_name, - ).model_dump( - exclude_none=True - ) - ) - - # If there's initial content, send it as a delta immediately - 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=current_agent_name, - model_name=current_model_name, - tool_name=current_tool_name, - ) - await safe_send_json( - _delta.model_dump( - exclude_none=True - ) - ) - # Mark that B1 streaming was used (nonlocal already declared above) - b1_streaming_used = True - - elif part_type == "ToolCallPart": - # Extract tool info from the part - tool_name = "unknown" - tool_call_id = None - tool_args_str = "" - - # Extract tool info from part_obj (dict or object) - if hasattr( - part_obj, "tool_name" - ): - tool_name = ( - part_obj.tool_name - ) - elif isinstance(part_obj, dict): - tool_name = part_obj.get( - "tool_name", "unknown" - ) - - if hasattr( - part_obj, "tool_call_id" - ): - tool_call_id = ( - part_obj.tool_call_id - ) - elif isinstance(part_obj, dict): - tool_call_id = part_obj.get( - "tool_call_id" - ) - - if hasattr(part_obj, "args"): - tool_args_str = ( - part_obj.args or "" - ) - elif isinstance(part_obj, dict): - tool_args_str = ( - part_obj.get("args", "") - or "" - ) - - tool_id = ( - tool_call_id - or str(uuid.uuid4())[:8] - ) - - # Track this tool call part with args buffer for accumulation - # Don't send tool_call yet - wait for part_end when args are complete - 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, # Buffer for accumulating args_delta - "start_time": time_module.time(), # Track when tool started - } - - # Track current tool name - current_tool_name = tool_name - - logger.debug( - f"[WebSocket] ToolCallPart started: {tool_name} (id: {tool_id})" + continue + if part_type == "ToolCallPart": + lifecycle_start_tool_call_part( + turn_state=turn_state, + part_index=part_index, + part_obj=part_obj, + logger=logger, ) - # tool_call event will be sent on part_end when args are complete - elif part_type == "ToolReturnPart": - logger.info( - f"[WebSocket] ToolReturnPart detected! part_index={part_index}" + 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, ) - # Extract tool result info from the part - tool_call_id = None - tool_content = None - - # Extract tool_call_id - if hasattr( - part_obj, "tool_call_id" - ): - tool_call_id = ( - part_obj.tool_call_id - ) - elif isinstance(part_obj, dict): - tool_call_id = part_obj.get( - "tool_call_id" - ) - - # Extract content (the tool result) - if hasattr(part_obj, "content"): - tool_content = ( - part_obj.content - ) - elif isinstance(part_obj, dict): - tool_content = part_obj.get( - "content" - ) - - # Try to serialize complex result objects - if ( - tool_content - and not isinstance( - tool_content, - ( - str, - dict, - list, - int, - float, - bool, - type(None), - ), - ) - ): - try: - import json - - # Try to convert to dict first (for Pydantic models) - if hasattr( - tool_content, - "model_dump", - ): - tool_content = tool_content.model_dump() - elif hasattr( - tool_content, "dict" - ): - tool_content = tool_content.dict() - elif hasattr( - tool_content, - "__dict__", - ): - tool_content = tool_content.__dict__ - else: - tool_content = str( - tool_content - ) - except Exception as e: - logger.debug( - f"[WebSocket] Could not serialize tool result: {e}" - ) - tool_content = str( - tool_content - ) - - # Find the corresponding pending tool call, store and SEND the result - _result_sent = False - if tool_call_id: - resolved_pending_id = resolve_pending_tool_id( - tool_call_id=tool_call_id - ) - if resolved_pending_id: - _result_sent = True - pending_info = 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( - tool_id=resolved_pending_id, - pending_info=pending_info, - fallback_group_id=current_tool_group_id, - tool_name=_tool_name, - source="tool_return_resolved", - ) - logger.info( - f"[WebSocket] ToolReturnPart: Sending result for {_tool_name} (id: {resolved_pending_id}, raw: {tool_call_id})" - ) - # Send the REAL tool_result with actual content - 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=current_agent_name - or "code-puppy", - model_name=current_model_name - or "unknown", - tool_group_id=_group_id, - ) - ) - else: - logger.warning( - f"[WebSocket] ToolReturnPart: Could not resolve tool_call_id {tool_call_id}, pending keys: {list(pending_tool_calls.keys())}" - ) - - # Fallback: proximity-based matching if no tool_call_id or resolution failed - if not _result_sent: - for ( - pending_id, - pending_info, - ) in sorted( - pending_tool_calls.items(), - key=lambda x: abs( - x[1].get( - "part_index", - 9999, - ) - - part_index - ), - ): - if ( - abs( - pending_info.get( - "part_index", - 9999, - ) - - part_index - ) - <= 3 - ): - # Close enough - assume it's the result for this tool call - 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( - tool_id=pending_id, - pending_info=pending_info, - fallback_group_id=current_tool_group_id, - tool_name=_tool_name, - source="tool_return_proximity", - ) - logger.info( - f"[WebSocket] ToolReturnPart: Sending result (by proximity) for {_tool_name} (id: {pending_id})" - ) - # Send the REAL tool_result with actual content - 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=current_agent_name - or "code-puppy", - model_name=current_model_name - or "unknown", - tool_group_id=_group_id, - ) - ) - _result_sent = True - break - - # Warn if we couldn't send the result - if not _result_sent: - logger.warning( - f"[WebSocket] ToolReturnPart: Could NOT send result! tool_call_id={tool_call_id}, pending_tool_calls={list(pending_tool_calls.keys())}, part_index={part_index}" - ) - - # Track this part for cleanup - active_parts[part_index] = { - "id": f"tool-return-{part_index}", - "type": "tool_return", - "tool_call_id": tool_call_id, - "content": tool_content, - } elif inner_type == "part_delta": part_index = inner_data.get( @@ -1984,252 +571,68 @@ def resolve_tool_group_id( delta_obj = inner_data.get( "delta", {} ) - - # Handle ToolCallPartDelta - accumulate args if ( delta_type == "ToolCallPartDelta" ): - args_delta = "" - if hasattr( - delta_obj, "args_delta" - ): - args_delta = ( - delta_obj.args_delta - or "" - ) - elif isinstance( - delta_obj, dict - ): - args_delta = ( - delta_obj.get( - "args_delta", "" - ) - or "" - ) - - if ( - args_delta - and part_index - in active_parts + if lifecycle_accumulate_tool_call_part_delta( + turn_state=turn_state, + part_index=part_index, + delta_obj=delta_obj, ): - part_info = active_parts[ - part_index - ] - if ( - part_info.get("type") - == "tool_call" - ): - part_info[ - "args_buffer" - ] = ( - part_info.get( - "args_buffer", - "", - ) - + args_delta - ) - continue # Don't process as text delta - - # Extract content_delta (supports both direct and nested formats) - content_delta = inner_data.get( - "content_delta", "" - ) - if not content_delta: - if isinstance(delta_obj, dict): - content_delta = ( - delta_obj.get( - "content_delta", "" - ) - ) - - if content_delta: - collected_text.append( - content_delta - ) + continue # Don't process as text delta - # Ensure we have an entry for this part (handles delta before start) - if ( - part_index - not in active_parts - ): - # Create entry with consistent message_id - message_id = f"msg-{int(time_module.time() * 1000)}-{part_index}" - active_parts[part_index] = { - "id": message_id, - "type": "text", # Default, will be updated by part_start if it arrives - "content": "", - } - # Send a start event so client creates the message - logger.debug( - f"[Stream Debug] Creating part on first delta: message_id={message_id}" - ) - if part_index == 0: - 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=current_agent_name, - model_name=current_model_name, - tool_name=current_tool_name, - ).model_dump( - exclude_none=True - ) - ) - - part_info = active_parts[ - part_index - ] - message_id = part_info["id"] - - if part_index in active_parts: - active_parts[part_index][ - "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=current_agent_name, - model_name=current_model_name, - tool_name=current_tool_name, - ) - await safe_send_json( - _delta.model_dump( - exclude_none=True - ) - ) - # Mark that B1 streaming was used - b1_streaming_used = True + 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 = active_parts.get( - part_index, {} + part_info = ( + turn_state.active_parts.get( + part_index, {} + ) ) part_type_info = part_info.get( "type", "text" ) - message_id = part_info.get( - "id", f"msg-{part_index}" - ) - full_content = part_info.get( - "content", "" - ) - # Handle tool_call parts - send tool_call event now that args are complete + 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": - 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" - ) - ) - - # Parse args if they're a JSON string - try: - import json - - args_dict = ( - json.loads( - tool_args_str - ) - if tool_args_str - else {} - ) - except ( - json.JSONDecodeError, - TypeError, - ): - args_dict = {} - - logger.debug( - f"[WebSocket] Sending tool_call (args complete): {tool_name}" - ) - - # Generate tool group ID for this batch if not already set - if ( - current_tool_group_id - is None - ): - current_tool_group_id = f"tg-{str(uuid.uuid4())[:8]}" - - # Send tool_call now that we have complete args - 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=current_agent_name, - model_name=current_model_name, - tool_group_id=current_tool_group_id, - ) + 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, ) - - # Track as pending tool call for later tool_result - 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, # Will be populated by ToolReturnPart - "tool_group_id": current_tool_group_id, - } - if raw_tool_call_id: - tool_id_aliases[ - raw_tool_call_id - ] = tool_id - # Also track tool_group_id for pre-stream_end extraction - if current_tool_group_id: - tool_group_ids[tool_id] = ( - current_tool_group_id - ) - - # Clear current tool name tracking - current_tool_name = None else: - await safe_send_json( - ServerAssistantMessageEnd( - message_id=message_id, - part_index=part_index, - full_content=full_content, - timestamp=time_module.time(), - session_id=session_id, - agent_name=current_agent_name, - model_name=current_model_name, - tool_name=current_tool_name, - ).model_dump( - exclude_none=True - ) + turn_state.active_parts.pop( + part_index, None ) - if part_index in active_parts: - del active_parts[part_index] - except Exception as send_err: error_msg = str(send_err).lower() if ( @@ -2256,32 +659,16 @@ def resolve_tool_group_id( final_count = 0 while True: try: - event = event_queue.get_nowait() + event = stream_event_queue.get_nowait() event_type = event.get("type", "") event_data = event.get("data", {}) final_count += 1 if event_type == "stream_event": - inner_type = event_data.get( - "event_type", "" + collect_final_stream_text_delta( + turn_state=turn_state, + event=event, ) - inner_data = event_data.get( - "event_data", {} - ) - if inner_type == "part_delta": - content_delta = inner_data.get( - "content_delta", "" - ) - if not content_delta: - delta_obj = inner_data.get( - "delta", {} - ) - if isinstance(delta_obj, dict): - content_delta = delta_obj.get( - "content_delta", "" - ) - if content_delta: - collected_text.append(content_delta) except Exception: break @@ -2292,43 +679,19 @@ def resolve_tool_group_id( 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}" - ) - - # Event to signal drain task is ready - drain_ready = asyncio.Event() - - try: - from code_puppy.plugins.frontend_emitter.emitter import ( - subscribe, - unsubscribe, - ) - - event_queue = subscribe(session_id=session_id) - logger.debug( - "Subscribed to frontend emitter for streaming" - ) - - # Modify drain function to signal when ready - async def drain_events_with_signal(): - """Wrapper that signals readiness before starting drain loop.""" - await drain_events_concurrent(drain_ready) - - # Start concurrent drain task - drain_task = asyncio.create_task( - drain_events_with_signal() - ) - active_drain_task = ( - drain_task # Track for potential cancellation + f"[{session_id}] Drain complete: " + f"{event_count} events during run, {final_count} final, " + f"batching efficiency: {avg_batch_size:.2f}" ) - # Wait for drain task to be ready before proceeding - await drain_ready.wait() - - except ImportError: - logger.warning("Frontend emitter not available") + 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 @@ -2356,91 +719,21 @@ async def drain_events_with_signal(): session_working_directory, ) - # Inject working directory as a system message in the - # agent's conversation history (once per directory change) - # so the LLM knows the CWD without polluting user messages. - message_to_send = user_message - if ( - session_working_directory - and session_working_directory - != last_context_sent_directory - ): - from pydantic_ai.messages import ( - ModelRequest, - SystemPromptPart, - ) - - wd_system_msg = ModelRequest( - parts=[ - SystemPromptPart( - content=( - f"The user's current working directory is updated to" - f" {session_working_directory}" - ) - ) - ] - ) - agent.append_to_message_history(wd_system_msg) - last_context_sent_directory = ( - session_working_directory - ) - logger.debug( - "Injected working directory system message: %s", - session_working_directory, - ) - - logger.debug( - f"Calling run_with_mcp with message: {message_to_send[:100]}..." + _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, ) - - # Build file context and binary attachments - file_context, binary_attachments = ( - build_file_context_and_attachments(msg) + 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 ) - - # CHANGE 1: Update attachment metadata for UI (variables already initialized) - if msg.get("attachments"): - from pathlib import Path as _AttachmentPath - - for raw_path in msg.get("attachments", []): - if ( - isinstance(raw_path, str) - and raw_path.strip() - ): - try: - file_path = _AttachmentPath( - 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( - f"Error building attachment metadata for '{raw_path}': {e}" - ) - - # Prepend file context to the message - if file_context: - message_to_send = ( - file_context + "\n\n" + message_to_send - ) - logger.debug( - f"Added file context ({len(file_context)} chars)" - ) - - run_kwargs = {} - if binary_attachments: - run_kwargs["attachments"] = binary_attachments - logger.debug( - f"Including {len(binary_attachments)} binary attachment(s)" - ) # ────────────────────────────────────────────────────────────────── # Phase 7: Create session + user message in SQLite BEFORE streaming @@ -2534,307 +827,33 @@ async def drain_events_with_signal(): exc_info=True, ) - # Run the agent in its own task so it can be cancelled independently. - # Suppress duplicate emitter lifecycle callbacks only during the - # run_with_mcp execution window (pre/post_tool_call hooks fire there). - # Wire the emitter ContextVar so events - # emitted by register_callbacks are tagged - # with this session and reach our subscriber. - 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) - active_agent_task = asyncio.create_task( - agent.run_with_mcp( - message_to_send, **run_kwargs - ) + _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, ) - - # ===== CONCURRENT MESSAGE PROCESSING FIX ===== - # Instead of awaiting the agent task here (which blocks the receive loop), - # we use asyncio.wait() to handle BOTH the agent task AND incoming messages. - # This allows permission_response messages to be processed while the agent runs. - - result = None - agent_completed = False - # Deferred message for switch_session/create_session during streaming - _deferred_msg: dict | None = None - - while not agent_completed: - # Create a task for the next message (or wait for agent) - receive_task = asyncio.create_task( - websocket.receive_json() - ) - - # Wait for either the agent to complete OR a new message to arrive - done, pending = await asyncio.wait( - {active_agent_task, receive_task}, - return_when=asyncio.FIRST_COMPLETED, - ) - - # Check if agent completed - if active_agent_task in done: - try: - result = await active_agent_task - logger.debug( - f"run_with_mcp completed, result type: {type(result)}" - ) - agent_completed = True - active_agent_task = None - except asyncio.CancelledError: - logger.debug( - "run_with_mcp task was cancelled by user" - ) - 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, - ) - agent_error = e - result = None - agent_completed = True - active_agent_task = None - - # Cancel the receive task if it's still pending - if receive_task in pending: - receive_task.cancel() - try: - await receive_task - except asyncio.CancelledError: - pass - - # Check if a new message arrived - elif receive_task in done: - try: - new_msg = await receive_task - - # Advisory validation — log but never reject - try: - _parsed_inner = _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" - }, - ) - - # Handle permission_response immediately - if ( - 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 - ) - ) - if handled: - logger.debug( - f"[Permission] ✅ Handled response: {request_id} = {approved}" - ) - else: - logger.warning( - f"[Permission] ❌ Unknown request: {request_id}" - ) - else: - logger.error( - "[WebSocket] ❌ No request_id in permission_response!" - ) - - # Handle cancel request - 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 # Exit the loop - ) - - # Handle session switch during streaming - elif new_msg.get("type") in ( - "switch_session", - "create_session", - ): - # User switched to a new chat while agent was running. - # The agent will finish in background and save to SQLite. - logger.debug( - "[WS:%s] Session switch during streaming — " - "agent continues in background, switching to: %s", - session_id, - new_msg.get("session_id"), - ) - - # Fire and forget — background task owns the agent result. - 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 # disown — background task now owns it - agent_completed = ( - True # exit inner loop - ) - - # For other message types, log and ignore (or queue for later) - else: - logger.warning( - f"[WebSocket] Received {new_msg.get('type')} message while agent running - ignoring" - ) - - except asyncio.CancelledError: - # Receive was cancelled, continue - pass - except WebSocketDisconnect: - # WebSocket gone — let agent finish and save to SQLite - logger.debug( - "[WS:%s] Disconnect during streaming — agent continues in background", - session_id, - ) - - # Fire and forget — background task owns the agent result. - 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(): - # Starlette raises RuntimeError when calling receive after disconnect - logger.debug( - "[WS:%s] WebSocket already disconnected: %s — agent continues in background", - session_id, - e, - ) - - # Fire and forget — background task owns the agent result. - 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( - f"RuntimeError processing message during agent execution: {e}" - ) - except Exception as e: - logger.error( - f"Error processing message during agent execution: {e}" - ) - - # ===== END CONCURRENT MESSAGE PROCESSING FIX ===== + result = _turn_run.result + _deferred_msg = _turn_run.deferred_msg finally: - # End suppression scope once run_with_mcp window ends, - # including completion/cancel/background handoff paths. - set_suppress_emitter_tool_events(False) - # Reset emitter ContextVar so stale session_id - # doesn't leak into unrelated coroutines. - if _emitter_token is not None: - try: - from code_puppy.plugins.frontend_emitter.session_context import ( - current_emitter_session_id, - ) - - current_emitter_session_id.reset( - _emitter_token - ) - except ImportError: - pass - # Clear session context for prompt generation - clear_session_working_directory() + pass finally: - # Stop the drain task - if drain_task: - stop_draining.set() - try: - await asyncio.wait_for(drain_task, timeout=2.0) - except asyncio.TimeoutError: - drain_task.cancel() - try: - await drain_task - except asyncio.CancelledError: - pass - - # Unsubscribe from emitter - if event_queue: - unsubscribe(event_queue) - logger.debug("Unsubscribed from frontend emitter") + 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: @@ -2844,236 +863,39 @@ async def drain_events_with_signal(): # The outer while True loop will handle switch_session/create_session continue - # If agent errored, send error to GUI and skip success path - if agent_error == "cancelled": + 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 - elif agent_error is not None: - logger.debug( - "[WS:%s] agent_error -> sending frame(s) to client. type=%s", - session_id, - type(agent_error).__name__, - ) - # If streaming was already in progress, send stream_end first so - # the frontend exits its streaming state before the error frame. - # Without this handshake the UI hangs indefinitely. - for frame in build_error_response_frames( - agent_error, collected_text, session_id - ): + if post_run.error_frames is not None: + for frame in post_run.error_frames: await safe_send_json(frame) continue - - # Safety net: if the agent task completed without raising but returned - # no result and produced no streamed text, treat it as an error. - has_nonempty_stream = any( - (chunk or "").strip() for chunk in 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(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." - ) - ) - await send_typed( - 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, - ) - ) + if post_run.no_result_error is not None: + await send_typed(post_run.no_result_error) continue - # Extract final response text - response_text = "" - - # Priority 1: Use collected text from streaming events - if has_nonempty_stream: - response_text = "".join(collected_text) - logger.debug( - f"Using collected streaming text ({len(response_text)} chars)" - ) - # Priority 2: Extract from result.output (AgentRunResult) or result.data (legacy) - 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)" - ) - # Priority 3: Get last assistant message from history - elif agent: - messages = agent.get_message_history() - - for msg in reversed(messages): - if hasattr(msg, "role") and msg.role == "assistant": - if 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" - - # Extract token usage from result if available - tokens_used = None - if result: - # Try to get usage from 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 - ), - } - # Also try _usage or other common patterns - 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 we couldn't get usage from result, estimate from message history - 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 + 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] b1_streaming_used=%s before response/extraction", - b1_streaming_used, + "[WebSocket] turn_state.b1_streaming_used=%s before response/extraction", + turn_state.b1_streaming_used, ) - if not b1_streaming_used: - # For non-streaming models (like Gemini), extract and send thinking content first - thinking_text = "" - if agent: - try: - history = agent.get_message_history() - logger.debug( - f"[Thinking Debug] Checking history for thinking parts, {len(history) if history else 0} messages" - ) - if history: - # Look for ThinkingPart in the last ModelResponse - for i, msg in enumerate(reversed(history)): - msg_type = type(msg).__name__ - logger.debug( - f"[Thinking Debug] Message {i}: type={msg_type}, has_parts={hasattr(msg, 'parts')}" - ) - if "Response" in msg_type and hasattr( - msg, "parts" - ): - logger.debug( - f"[Thinking Debug] Found Response with {len(msg.parts)} 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( - f"[Thinking Debug] Part {j}: type={part_type}, content_preview={part_content_preview}" - ) - if ( - "Thinking" in part_type - and hasattr(part, "content") - ): - thinking_text = part.content - logger.debug( - f"[Thinking Debug] Found thinking content: {len(thinking_text)} chars" - ) - 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" - ) - + if not turn_state.b1_streaming_used: # Send thinking content as B1 message if available if thinking_text: import time as time_module @@ -3147,124 +969,17 @@ async def drain_events_with_signal(): 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: set = set() - logger.warning( - "[WebSocket] B1 Pre-stream_end extraction: result=%s, has_all_messages=%s", - type(result).__name__ if result else None, - hasattr(result, "all_messages") - if result - else False, + _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, ) - if result and hasattr(result, "all_messages"): - try: - import time as _time_pre - - from pydantic_ai.messages import ( - ToolReturn, - ToolReturnPart, - ) - - _pre_msgs = list(result.all_messages()) - for _pre_msg in _pre_msgs: - if not hasattr(_pre_msg, "parts"): - continue - for _pre_part in _pre_msg.parts: - if not isinstance( - _pre_part, - (ToolReturnPart, ToolReturn), - ): - continue - _pre_tool_name = getattr( - _pre_part, "tool_name", "unknown" - ) - _pre_raw_tool_id = getattr( - _pre_part, "tool_call_id", None - ) - _pre_tool_id = ( - tool_id_aliases.get( - _pre_raw_tool_id, - _pre_raw_tool_id, - ) - if _pre_raw_tool_id - else "unknown" - ) - _pre_result = getattr( - _pre_part, "content", None - ) - # Serialize complex objects - if _pre_result and not isinstance( - _pre_result, - ( - str, - dict, - list, - int, - float, - bool, - type(None), - ), - ): - try: - if hasattr( - _pre_result, "model_dump" - ): - _pre_result = ( - _pre_result.model_dump() - ) - elif hasattr( - _pre_result, "dict" - ): - _pre_result = ( - _pre_result.dict() - ) - elif hasattr( - _pre_result, "__dict__" - ): - _pre_result = ( - _pre_result.__dict__ - ) - else: - _pre_result = str( - _pre_result - ) - except Exception: - _pre_result = str(_pre_result) - logger.warning( - "[WebSocket] Pre-stream_end tool result: %s (id: %s), content_preview=%s", - _pre_tool_name, - _pre_tool_id, - str(_pre_result)[:100] - if _pre_result - else None, - ) - _pre_group_id = tool_group_ids.get( - _pre_tool_id - ) - - await send_typed_tool_lifecycle( - ServerToolResult( - tool_id=_pre_tool_id, - tool_name=_pre_tool_name, - result=_pre_result, - success=True, - duration_ms=0, - timestamp=_time_pre.time(), - session_id=session_id, - agent_name=agent.name - if agent - else "code-puppy", - model_name=agent.get_model_name() - if agent - else "unknown", - tool_group_id=_pre_group_id, - ) - ) - _pre_sent_tool_ids.add(_pre_tool_id) - except Exception as _pre_e: - logger.warning( - "Pre-stream_end tool result extraction failed: %s", - _pre_e, - ) # Send stream_end AFTER real tool results are delivered await send_typed( @@ -3281,250 +996,46 @@ async def drain_events_with_signal(): session_id=session_id, ) ) - # Save session after each response - # Track tool IDs already sent before stream_end to skip duplicates later - if "_pre_sent_tool_ids" not in dir(): - _pre_sent_tool_ids: set = set() - _save_history_snapshot = [] # Snapshot of history before await points try: - # CRITICAL: Update message history from result to include final response - # The pydantic-ai result.all_messages() contains the complete conversation - # including the final assistant response that may not be in agent's history yet - if result and hasattr(result, "all_messages"): - try: - all_msgs = list(result.all_messages()) - if all_msgs: - agent.set_message_history(all_msgs) - _save_history_snapshot = list( - agent.get_message_history() - ) # Snapshot before any awaits can corrupt shared global state - logger.debug( - f"Updated message history from result.all_messages(): {len(all_msgs)} messages" - ) - - # Extract and send tool results from message history - try: - import time as _time - - from pydantic_ai.messages import ( - ToolReturn, - ToolReturnPart, - ) - - for msg in all_msgs: - if hasattr(msg, "parts"): - for part in msg.parts: - if isinstance( - part, - ( - ToolReturnPart, - ToolReturn, - ), - ): - tool_name = getattr( - part, - "tool_name", - "unknown", - ) - tool_call_id = getattr( - part, - "tool_call_id", - "unknown", - ) - # Serialize the result - result_data = getattr( - part, "content", None - ) - if ( - result_data - and not isinstance( - result_data, - ( - str, - dict, - list, - int, - float, - bool, - type(None), - ), - ) - ): - try: - if hasattr( - result_data, - "model_dump", - ): - result_data = result_data.model_dump() - elif hasattr( - result_data, - "dict", - ): - result_data = result_data.dict() - elif hasattr( - result_data, - "__dict__", - ): - result_data = result_data.__dict__ - else: - result_data = str( - result_data - ) - except Exception: - result_data = str( - result_data - ) - - # Log for shell commands - if ( - tool_name - == "agent_run_shell_command" - ): - stdout_val = ( - result_data.get( - "stdout", "N/A" - ) - if isinstance( - result_data, - dict, - ) - else "not dict" - ) - logger.debug( - "Extracted shell result: id=%s, stdout=%s", - tool_call_id, - stdout_val, - ) - - # Skip tool IDs already sent before stream_end - if ( - tool_call_id - in _pre_sent_tool_ids - ): - logger.debug( - "[WebSocket] Skipping duplicate tool result (pre-sent): %s", - tool_call_id, - ) - continue - logger.info( - f"[WebSocket] Sending extracted tool result for {tool_name} (id: {tool_call_id})" - ) - _post_group_id = ( - tool_group_ids.get( - tool_call_id - ) - ) - - await send_typed( - ServerToolResult( - tool_id=tool_call_id, - tool_name=tool_name, - result=result_data, - success=True, - duration_ms=0, - timestamp=_time.time(), - session_id=session_id, - agent_name=agent.name - if agent - else "code-puppy", - model_name=agent.get_model_name() - if agent - else "unknown", - tool_group_id=_post_group_id, - ) - ) - except Exception as e: - logger.warning( - f"Could not extract tool results from messages: {e}" - ) - - except Exception as e: - logger.warning( - f"Could not update history from result.all_messages(): {e}" - ) + 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 = ( - _save_history_snapshot - if _save_history_snapshot + finalized_turn.history_snapshot + if finalized_turn.history_snapshot else agent.get_message_history() ) # Use pre-await snapshot to avoid race condition - if history: - # Regenerate title if it's empty or still the default "untitled-session" - if ( - not session_title - or session_title == "untitled-session" - ): - session_title = generate_heuristic_title( - history - ) - - # Session name is just the ID (WS_session_timestamp format) - # Title is stored only in the meta file, not in the filename - 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, - ) - - # Send session metadata update to client - 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, - ) - ) - - # Broadcast session update to session monitoring clients - 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, - ) - ) + 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}" @@ -3559,6 +1070,10 @@ async def drain_events_with_signal(): 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 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..dca09762f --- /dev/null +++ b/code_puppy/api/ws/chat_turn_runner.py @@ -0,0 +1,273 @@ +"""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 + +_ClientMessageAdapter = TypeAdapter(ClientMessage) +logger = logging.getLogger(__name__) + + +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 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) + 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/schemas.py b/code_puppy/api/ws/schemas.py index f7a198e65..ee67f9087 100644 --- a/code_puppy/api/ws/schemas.py +++ b/code_puppy/api/ws/schemas.py @@ -47,6 +47,7 @@ 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" diff --git a/code_puppy/api/ws/session_persistence.py b/code_puppy/api/ws/session_persistence.py index 0bc356e42..bc1a7be0d 100644 --- a/code_puppy/api/ws/session_persistence.py +++ b/code_puppy/api/ws/session_persistence.py @@ -12,7 +12,15 @@ import datetime import logging -from typing import Any, Optional +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, +) +from code_puppy.session_storage import generate_heuristic_title logger = logging.getLogger(__name__) @@ -80,8 +88,7 @@ def build_session_update_payload( """Build the broadcast payload for ``connection_manager.broadcast_session_update``. ``action`` is ``"created"`` when ``message_count == 1`` (first turn), - otherwise ``"updated"``. Legacy ``pickle_path`` / ``metadata_path`` - fields are kept as empty strings for backwards-compat. + otherwise ``"updated"``. """ return { "session_id": session_id, @@ -92,12 +99,113 @@ def build_session_update_payload( "message_count": message_count, "total_tokens": total_tokens, "auto_saved": True, - "pickle_path": "", - "metadata_path": "", "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, 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..a2ca6ad34 --- /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): + handled.append((request_id, approved)) + 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)] + 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_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 index 7175f67b0..2d1170ee7 100644 --- a/code_puppy/api/ws/tests/test_tool_group_id_consistency.py +++ b/code_puppy/api/ws/tests/test_tool_group_id_consistency.py @@ -2,6 +2,8 @@ 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 @@ -9,11 +11,16 @@ import ast from pathlib import Path -CHAT_HANDLER_PATH = Path(__file__).resolve().parents[1] / "chat_handler.py" +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 chat_handler source.""" + """Return every ``ServerToolResult(...)`` call in source.""" tree = ast.parse(source) calls: list[ast.Call] = [] @@ -32,36 +39,42 @@ def test_server_tool_result_calls_always_include_tool_group_id_keyword() -> None This protects against regressions where a new emission path forgets to include the grouping field and breaks frontend tool-message grouping. """ - source = CHAT_HANDLER_PATH.read_text(encoding="utf-8") - calls = _get_server_tool_result_calls(source) - - assert calls, "Expected at least one ServerToolResult(...) call" - - missing = [] - 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(call.lineno) - + 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 lines: {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.""" - source = CHAT_HANDLER_PATH.read_text(encoding="utf-8") - calls = _get_server_tool_result_calls(source) - - explicit_none_lines: list[int] = [] - 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_lines.append(call.lineno) - - assert not explicit_none_lines, ( - "ServerToolResult(...) calls pass explicit tool_group_id=None at lines: " - f"{explicit_none_lines}" + 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..105e575e1 --- /dev/null +++ b/code_puppy/api/ws/tests/test_ws_command_handler.py @@ -0,0 +1,108 @@ +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" 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..063807702 --- /dev/null +++ b/code_puppy/api/ws/tests/test_ws_control_messages.py @@ -0,0 +1,217 @@ +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: permission_calls.append((request_id, approved)) 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)] + 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_session_bootstrap.py b/code_puppy/api/ws/tests/test_ws_session_bootstrap.py new file mode 100644 index 000000000..c144b7b5d --- /dev/null +++ b/code_puppy/api/ws/tests/test_ws_session_bootstrap.py @@ -0,0 +1,235 @@ +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..bca4f3095 --- /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 or result is not False, + 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 or result is not False, + 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..915c40566 --- /dev/null +++ b/code_puppy/api/ws/ws_control_messages.py @@ -0,0 +1,606 @@ +"""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) + ) + + +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]) -> 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) + 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_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/cli_runner.py b/code_puppy/cli_runner.py index 9fb6d170d..825268524 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", @@ -1051,7 +1051,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/load_context_completion.py b/code_puppy/command_line/load_context_completion.py index 5b8157a6b..05d80c259 100644 --- a/code_puppy/command_line/load_context_completion.py +++ b/code_puppy/command_line/load_context_completion.py @@ -38,8 +38,8 @@ 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 + for pkl_file in contexts_dir.glob("*.json"): + session_name = pkl_file.stem # removes .json extension if session_name.startswith(session_filter): yield Completion( session_name, diff --git a/code_puppy/messaging/bus.py b/code_puppy/messaging/bus.py index f53206ad8..2b8196fc8 100644 --- a/code_puppy/messaging/bus.py +++ b/code_puppy/messaging/bus.py @@ -35,6 +35,7 @@ import asyncio import queue import threading +from contextvars import ContextVar from typing import Any, Dict, List, Optional, Tuple from uuid import uuid4 @@ -89,8 +90,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 +112,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 +195,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 +203,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) diff --git a/code_puppy/plugins/switch_agent_resume/register_callbacks.py b/code_puppy/plugins/switch_agent_resume/register_callbacks.py index 0625b4091..671e7500c 100644 --- a/code_puppy/plugins/switch_agent_resume/register_callbacks.py +++ b/code_puppy/plugins/switch_agent_resume/register_callbacks.py @@ -29,7 +29,7 @@ def _do_switch_and_resume(agent_name: str) -> bool: ) from code_puppy.messaging import emit_info, emit_success, emit_warning from code_puppy.session_storage import ( - build_session_paths, + get_session_file_path, load_session, save_session, ) @@ -91,11 +91,11 @@ def _do_switch_and_resume(agent_name: str) -> bool: if last_session_to_resume: try: autosave_dir = Path(AUTOSAVE_DIR).resolve() - session_pickle = build_session_paths( + session_file = get_session_file_path( autosave_dir, last_session_to_resume - ).pickle_path.resolve() - if autosave_dir not in session_pickle.parents: - raise FileNotFoundError(session_pickle) + ).resolve() + if autosave_dir not in session_file.parents: + raise FileNotFoundError(session_file) history = load_session(last_session_to_resume, autosave_dir) new_agent.set_message_history(history) pin_current_session_name(last_session_to_resume) diff --git a/code_puppy/session_storage.py b/code_puppy/session_storage.py index ab9a849bb..4e5c1e67e 100644 --- a/code_puppy/session_storage.py +++ b/code_puppy/session_storage.py @@ -1,54 +1,43 @@ -"""Shared helpers for persisting and restoring chat sessions. +"""Shared helpers for persisting and restoring JSON chat sessions. -MIGRATION NOTE: Pickle-based storage has been removed. All sessions are now -persisted exclusively to SQLite via code_puppy.api.db.queries. -This module retains only stateless helpers (title generation) and stubs for -removed functionality to maintain import compatibility during the transition. +All session export/autosave helpers now use JSON files. Pickle compatibility +and legacy path shims have been removed from active runtime code. """ from __future__ import annotations import json import re +from datetime import datetime from dataclasses import dataclass from pathlib import Path from typing import Any, List -from pydantic import TypeAdapter -from pydantic_ai.messages import ModelMessage SessionHistory = List[Any] -@dataclass +@dataclass(slots=True) class SessionMetadata: - """Metadata returned after saving a session.""" + """Metadata returned after saving a session export/autosave.""" message_count: int total_tokens: int - pickle_path: Path # Kept for compatibility, points to .json file - metadata_path: Path # Points to .json file (same as above) + session_file_path: Path -def generate_heuristic_title(history: SessionHistory, max_length: int = 50) -> str: - """Generate a short title from the first user message in the history. +def get_session_file_path(base_dir: Path, session_name: str) -> Path: + """Return the canonical JSON session file path for *session_name*.""" + return base_dir / f"{session_name}.json" - Extracts the first user message, takes the first ~50 chars, and converts - to a filename-safe format (lowercase, spaces to hyphens, remove special chars). - Handles multiple message formats: - 1. Pydantic-ai format: msg.kind == 'request' with msg.parts[].content - 2. Enhanced/wrapped format: {'msg': , 'agent': str, ...} - 3. Simple dict format: {'role': 'user', 'content': str} - """ +def generate_heuristic_title(history: SessionHistory, max_length: int = 50) -> str: + """Generate a short title from the first user message in the history.""" def extract_user_content(msg: Any) -> str | None: - """Extract user message content from various message formats.""" - # Handle wrapped/enhanced format: {'msg': , 'agent': ...} if isinstance(msg, dict) and "msg" in msg: msg = msg["msg"] - # Handle pydantic-ai format: msg.kind == 'request' if hasattr(msg, "kind") and msg.kind == "request": for part in getattr(msg, "parts", []): if hasattr(part, "content") and isinstance(part.content, str): @@ -56,7 +45,6 @@ def extract_user_content(msg: Any) -> str | None: if content: return content - # Handle simple dict format: {'role': 'user', 'content': str} if isinstance(msg, dict): if msg.get("role") == "user" and isinstance(msg.get("content"), str): content = msg["content"].strip() @@ -66,18 +54,14 @@ def extract_user_content(msg: Any) -> str | None: return None def content_to_title(content: str) -> str: - """Convert content to a filename-safe kebab-case title.""" - # Take first line or first max_length chars first_line = content.split("\n")[0][:max_length] - # Convert to kebab-case filename-safe format title = first_line.lower() - title = re.sub(r"[^a-z0-9\s-]", "", title) # Remove special chars - title = re.sub(r"\s+", "-", title) # Spaces to hyphens - title = re.sub(r"-+", "-", title) # Collapse multiple hyphens + 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 - # Find first user message for msg in history: content = extract_user_content(msg) if content: @@ -87,11 +71,6 @@ def content_to_title(content: str) -> str: return "untitled-session" -# --------------------------------------------------------------------------- -# Deprecated / Removed Pickle Functionality -> Replaced with JSON for Export -# --------------------------------------------------------------------------- - - def save_session( history: SessionHistory, session_name: str, @@ -100,67 +79,55 @@ def save_session( token_estimator: Any | None = None, **kwargs: Any, ) -> SessionMetadata: - """Save session history to a JSON file (replacing legacy pickle). + """Save session history to a JSON file.""" + del timestamp, kwargs - This is used for 'pinning' or exporting sessions via /dump_context. - """ base_dir.mkdir(parents=True, exist_ok=True) - file_path = base_dir / f"{session_name}.json" + file_path = get_session_file_path(base_dir, session_name) - # Try to serialize the history, handling mixed types: - # - ModelMessage objects: serialize via Pydantic - # - System dicts: serialize as plain dicts try: from pydantic_ai.messages import ModelMessagesTypeAdapter - history_data = [] - for item in history: - # Unwrap from {msg: ..., agent: ..., ts: ...} wrapper if present - if isinstance(item, dict) and "msg" in item: - inner = item["msg"] - wrapper_meta = {k: v for k, v in item.items() if k != "msg"} - - # Check if inner is a ModelMessage (has 'parts' attribute) - if hasattr(inner, "parts"): - # Serialize ModelMessage, then add wrapper metadata + if history and all(hasattr(item, "parts") for item in history): + json_bytes = ModelMessagesTypeAdapter.dump_json(history, indent=2) + else: + history_data = [] + for item in history: + if isinstance(item, dict) and "msg" in item: + inner = item["msg"] + wrapper_meta = {k: v for k, v in item.items() if k != "msg"} + if hasattr(inner, "parts"): + try: + serialized = ModelMessagesTypeAdapter.dump_python( + [inner], mode="json" + )[0] + serialized["_wrapper"] = wrapper_meta + history_data.append(serialized) + except Exception: + history_data.append(item) + else: + history_data.append(item) + elif hasattr(item, "parts"): try: serialized = ModelMessagesTypeAdapter.dump_python( - [inner], mode="json" + [item], mode="json" )[0] - serialized["_wrapper"] = wrapper_meta history_data.append(serialized) except Exception: - # Fallback: keep as-is - history_data.append(item) + history_data.append({"_raw": str(item)}) else: - # It's a system dict like {msg: 'system', ...} - keep as-is history_data.append(item) - elif hasattr(item, "parts"): - # Direct ModelMessage without wrapper - try: - serialized = ModelMessagesTypeAdapter.dump_python( - [item], mode="json" - )[0] - history_data.append(serialized) - except Exception: - history_data.append({"_raw": str(item)}) - else: - # Unknown type - keep as-is - history_data.append(item) - json_bytes = json.dumps(history_data, indent=2, default=str).encode("utf-8") + json_bytes = json.dumps(history_data, indent=2, default=str).encode("utf-8") except Exception: - # Fallback: generic JSON dump json_bytes = json.dumps(history, default=str, indent=2).encode("utf-8") file_path.write_bytes(json_bytes) - # Calculate tokens if estimator provided total_tokens = 0 if token_estimator: for msg in history: try: - # estimator might expect the original object or dict total_tokens += token_estimator(msg) except Exception: pass @@ -168,71 +135,114 @@ def save_session( return SessionMetadata( message_count=len(history), total_tokens=total_tokens, - pickle_path=file_path, - metadata_path=file_path, + session_file_path=file_path, ) -def load_session( - session_name: str, base_dir: Path, *, allow_legacy: bool = False -) -> SessionHistory: +def load_session(session_name: str, base_dir: Path) -> SessionHistory: """Load session history from a JSON file.""" - file_path = base_dir / f"{session_name}.json" + file_path = get_session_file_path(base_dir, session_name) if not file_path.exists(): - # Fallback check for legacy .pkl if explicitly requested? - legacy_path = base_dir / f"{session_name}.pkl" - if allow_legacy and legacy_path.exists(): - raise NotImplementedError( - f"Found legacy pickle at {legacy_path} but pickle loading is removed." - ) raise FileNotFoundError(f"Session file not found: {file_path}") json_data = file_path.read_bytes() - # Try to load as ModelMessage objects try: - adapter = TypeAdapter(List[ModelMessage]) - return adapter.validate_json(json_data) + from pydantic_ai.messages import ModelMessagesTypeAdapter + + return ModelMessagesTypeAdapter.validate_json(json_data) + except Exception: + raw = json.loads(json_data) + + if not isinstance(raw, list): + return raw + + try: + from pydantic_ai.messages import ModelRequest, ModelResponse, TextPart + from pydantic_ai.usage import RequestUsage + + restored: list[Any] = [] + for item in raw: + wrapper = None + if isinstance(item, dict): + wrapper = item.pop("_wrapper", None) + + if isinstance(item, dict) and item.get("kind") in {"request", "response"}: + parts = [] + for part in item.get("parts", []): + if isinstance(part, dict) and part.get("part_kind") == "text": + parts.append( + TextPart( + content=part.get("content", ""), + id=part.get("id"), + provider_details=part.get("provider_details"), + ) + ) + else: + parts.append(part) + + if item.get("kind") == "request": + msg = ModelRequest( + parts=parts, + instructions=item.get("instructions"), + run_id=item.get("run_id"), + metadata=item.get("metadata"), + ) + else: + usage_data = item.get("usage") or {} + timestamp = item.get("timestamp") + if isinstance(timestamp, str): + timestamp = datetime.fromisoformat( + timestamp.replace("Z", "+00:00") + ) + msg = ModelResponse( + parts=parts, + usage=RequestUsage(**usage_data), + model_name=item.get("model_name"), + timestamp=timestamp or datetime.now().astimezone(), + provider_name=item.get("provider_name"), + provider_details=item.get("provider_details"), + provider_response_id=item.get("provider_response_id"), + finish_reason=item.get("finish_reason"), + run_id=item.get("run_id"), + metadata=item.get("metadata"), + ) + + restored.append( + {**wrapper, "msg": msg} if isinstance(wrapper, dict) else msg + ) + else: + restored.append(item) + + return restored except Exception: - # Fallback: return as list of dicts - return json.loads(json_data) + return raw def list_sessions(base_dir: Path) -> List[str]: """List available JSON sessions.""" if not base_dir.exists(): return [] - return sorted([p.stem for p in base_dir.glob("*.json")]) + return sorted(path.stem for path in base_dir.glob("*.json")) def cleanup_sessions(base_dir: Path, max_sessions: int) -> List[str]: - """Cleanup old sessions if count exceeds max_sessions. - - Returns list of removed session names. - """ + """Cleanup old sessions if count exceeds *max_sessions*.""" if not base_dir.exists() or max_sessions <= 0: return [] - sessions = sorted(base_dir.glob("*.json"), key=lambda p: p.stat().st_mtime) - - removed = [] + sessions = sorted(base_dir.glob("*.json"), key=lambda path: path.stat().st_mtime) + removed: list[str] = [] if len(sessions) > max_sessions: - to_remove = sessions[: len(sessions) - max_sessions] - for p in to_remove: + for path in sessions[: len(sessions) - max_sessions]: try: - p.unlink() - removed.append(p.stem) + path.unlink() + removed.append(path.stem) except Exception: pass - return removed async def restore_autosave_interactively(base_dir: Path) -> None: - """Deprecated: No-op.""" - pass - - -def build_session_paths(base_dir: Path, session_name: str) -> Any: - """Deprecated: Returns None or raises.""" - raise NotImplementedError("Session paths are no longer used.") + """Deprecated compatibility shim retained as a no-op.""" + del base_dir diff --git a/code_puppy/tools/agent_tools.py b/code_puppy/tools/agent_tools.py index f472b14b6..37af1dd9f 100644 --- a/code_puppy/tools/agent_tools.py +++ b/code_puppy/tools/agent_tools.py @@ -1,7 +1,6 @@ # agent_tools.py import hashlib import json -import pickle import re from datetime import datetime from pathlib import Path @@ -22,6 +21,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 +124,12 @@ 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. + save_session( + history=message_history, + session_name=session_id, + base_dir=sessions_dir, + ) # Save or update txt file with metadata txt_path = sessions_dir / f"{session_id}.txt" @@ -172,16 +171,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/docs/migration/ws-chat-handler-refactor-checklist.md b/docs/migration/ws-chat-handler-refactor-checklist.md new file mode 100644 index 000000000..08ba00cda --- /dev/null +++ b/docs/migration/ws-chat-handler-refactor-checklist.md @@ -0,0 +1,464 @@ +# WS `chat_handler.py` modularization checklist + +Base branch: `feature/puppy-desk-migration` +Maestro owner: `maestro-687542` +Primary bead chain: `code_puppy-40g` → `code_puppy-79a` → `code_puppy-8ql` → `code_puppy-e7v` + +Goal: split `code_puppy/api/ws/chat_handler.py` into smaller modules **without any protocol, persistence, session, or streaming behavior regressions**. + +Last reconciled against code: 2026-06-10 + +--- + +## Non-negotiable safety invariants + +- [ ] WebSocket frame shapes stay identical for all existing message types. +- [ ] Frame ordering stays identical, especially around: + - `thinking` + - assistant `part_start` / `part_delta` / `part_end` + - tool call/result events + - `stream_end` + - final `status=done` +- [ ] Session bootstrap semantics stay identical for both new and restored sessions. +- [ ] SQLite remains the source of truth for resumed sessions. +- [ ] `switch_session` behavior stays identical, including create-if-missing flow. +- [ ] `cancel` and disconnect behavior stays identical, including background-save behavior. +- [ ] Permission request/response flow stays identical. +- [ ] Working-directory banners + session row updates stay identical. +- [ ] Command execution behavior stays identical. +- [ ] Agent/model switch persistence stays identical. +- [ ] No shared mutable WS state leaks across sessions. +- [ ] No changes to `code_puppy/command_line/`. + +--- + +## Current extracted modules already in place + +These are already good extraction boundaries and should be preserved: + +- [x] `code_puppy/api/ws/chat_context.py` + - per-message / per-agent-run ContextVar setup & cleanup +- [x] `code_puppy/api/ws/chat_turn_state.py` + - per-turn mutable streaming/tool state +- [x] `code_puppy/api/ws/chat_turn_runner.py` + - concurrent `run_with_mcp` execution while receiving control messages +- [x] `code_puppy/api/ws/ws_turn_preparation.py` + - working-directory prompt injection + attachment packaging +- [x] `code_puppy/api/ws/ws_stream_drain.py` + - drain task lifecycle shell (subscribe/start/stop/cancel) +- [x] `code_puppy/api/ws/chat_event_adapter.py` + - assistant text/thinking frame adaptation +- [x] `code_puppy/api/ws/chat_tool_lifecycle.py` + - tool call/result reconciliation helpers +- [x] `code_puppy/api/ws/ws_post_run.py` + - post-run cancelled/error/no-result resolution +- [x] `code_puppy/api/ws/ws_turn_finalization.py` + - pre/post `stream_end` tool result emission + history snapshot logic +- [x] `code_puppy/api/ws/session_persistence.py` + - session meta payloads + final persistence/broadcast orchestration + +--- + +## What still lives inside `chat_handler.py` + +`chat_handler.py` is still carrying too many responsibilities: + +- [ ] WebSocket/session bootstrap +- [ ] resumed-session replay of persisted system banners +- [ ] non-message control routing (`switch_agent`, `switch_model`, `switch_session`, etc.) +- [ ] slash-command execution +- [ ] permission-response top-level handling +- [ ] per-message orchestration for `type=message` +- [ ] inline `drain_events_concurrent()` implementation +- [ ] pre-stream SQLite upsert + user-message write +- [ ] final send/persist/status wiring +- [ ] disconnect/final cleanup + +That means the next safe split should be about **removing orchestration branches**, not changing lower-level behavior. + +--- + +## Target file layout + +### Keep as-is +- `chat_context.py` +- `chat_turn_state.py` +- `chat_turn_runner.py` +- `chat_event_adapter.py` +- `chat_tool_lifecycle.py` +- `ws_turn_preparation.py` +- `ws_stream_drain.py` +- `ws_post_run.py` +- `ws_turn_finalization.py` +- `session_persistence.py` +- `send_utils.py` +- `response_frames.py` +- `schemas.py` + +### Add next +- [ ] `code_puppy/api/ws/ws_session_bootstrap.py` +- [ ] `code_puppy/api/ws/ws_control_messages.py` +- [ ] `code_puppy/api/ws/ws_command_handler.py` +- [ ] `code_puppy/api/ws/ws_stream_processor.py` +- [ ] `code_puppy/api/ws/ws_prestream_persistence.py` +- [ ] `code_puppy/api/ws/ws_message_orchestrator.py` +- [ ] `code_puppy/api/ws/ws_chat_runtime.py` + +### End state for `chat_handler.py` +- [ ] endpoint registration only +- [ ] create runtime state +- [ ] call bootstrap helper +- [ ] receive loop delegates to control/message router +- [ ] final cleanup helper calls only + +--- + +## File-by-file checklist + +### 1) `code_puppy/api/ws/ws_chat_runtime.py` — create connection-scoped runtime state + +Purpose: remove the long list of mutable local variables from `websocket_chat()`. + +Create a runtime dataclass that holds only per-connection mutable state: + +- [ ] `session_id` +- [ ] `ctx` +- [ ] `session_title` +- [ ] `session_working_directory` +- [ ] `session_pinned` +- [ ] `last_context_sent_directory` +- [ ] `existing_history` +- [ ] `agent` +- [ ] `agent_name` +- [ ] `model_name` +- [ ] `active_drain_task` +- [ ] `active_agent_task` (if still needed at outer level) +- [ ] `stop_draining` + +Rules: +- [ ] instantiate inside `websocket_chat()` only +- [ ] no module-level mutable state +- [ ] use `field(default_factory=...)` for mutable fields like `asyncio.Event` +- [ ] keep it as a plain state bag, not a behavior-heavy class + +Acceptance criteria: +- [ ] zero behavior change +- [ ] easier helper signatures because they take `runtime` instead of 8-12 loose arguments + +--- + +### 2) `code_puppy/api/ws/ws_session_bootstrap.py` — extract connect / restore / init flow + +Move the following logic out of `chat_handler.py`: + +- [ ] validate incoming `session_id` +- [ ] generate timestamp-based session id when absent +- [ ] check SQLite for existing session +- [ ] create or load session via `session_manager` +- [ ] set `sender.session_id` and `sender.ctx` +- [ ] mark session active +- [ ] initialize process tracking +- [ ] set MessageBus session context +- [ ] derive initial `agent`, `agent_name`, `model_name` +- [ ] write initial config/directory system messages for brand-new sessions only +- [ ] send initial `ServerSystem(Connected!)` +- [ ] send session meta snapshot +- [ ] send `ServerSessionRestored` for resumed sessions +- [ ] replay persisted system rows (`init`, `config`, `directory`) + +Recommended public helpers: +- [ ] `initialize_ws_session(...) -> WebSocketChatRuntime` +- [ ] `send_session_meta_snapshot(...)` +- [ ] `replay_restored_system_messages(...)` + +Safety notes: +- [ ] preserve fallback behavior when `get_or_load_session()` returns `None` +- [ ] preserve SQLite-first semantics +- [ ] preserve warning-only behavior for bootstrap persistence failures + +Suggested tests: +- [ ] new session bootstrap writes initial config row +- [ ] resumed session sends `session_restored` +- [ ] resumed session replays persisted system rows +- [ ] invalid `session_id` closes with policy violation behavior unchanged + +--- + +### 3) `code_puppy/api/ws/ws_control_messages.py` — extract non-chat control routing + +This file should handle all top-level message types except `type=message`. + +Move branches for: +- [ ] `switch_agent` +- [ ] `switch_model` +- [ ] `switch_session` +- [ ] `set_working_directory` +- [ ] `update_session_meta` +- [ ] `get_config` +- [ ] `set_config` +- [ ] `cancel` +- [ ] `permission_response` + +Recommended shape: +- [ ] one public dispatcher: `handle_control_message(...) -> bool` +- [ ] returns `True` when the message was handled and outer loop should continue +- [ ] returns `False` when caller should treat it as a chat message or unknown type + +Internal helper split: +- [ ] `_handle_switch_agent(...)` +- [ ] `_handle_switch_model(...)` +- [ ] `_handle_switch_session(...)` +- [ ] `_handle_set_working_directory(...)` +- [ ] `_handle_update_session_meta(...)` +- [ ] `_handle_get_config(...)` +- [ ] `_handle_set_config(...)` +- [ ] `_handle_cancel(...)` +- [ ] `_handle_permission_response(...)` + +Safety notes: +- [ ] preserve current send frames exactly +- [ ] preserve `cancel_active_streaming(...)` calls before session switch/cancel +- [ ] preserve create-if-missing semantics in `switch_session` +- [ ] preserve exact SQLite writes for directory changes and meta updates +- [ ] preserve `continue` behavior after handled control messages + +Suggested tests: +- [ ] switch session existing target +- [ ] switch session create-new target +- [ ] set unchanged working directory returns `unchanged=True` +- [ ] invalid working directory returns failure payload +- [ ] update session meta does not zero counters in SQLite +- [ ] cancel cancels drain + in-flight task + sends cancelled status +- [ ] permission response routes request id unchanged + +--- + +### 4) `code_puppy/api/ws/ws_command_handler.py` — isolate slash-command execution + +Move the `type=command` branch into its own helper. + +Move logic for: +- [ ] `/help` special rendering path +- [ ] generic `handle_command(command_str)` path +- [ ] `ServerCommandResult` success/error payload creation +- [ ] traceback logging on failure + +Recommended public helper: +- [ ] `handle_command_message(...) -> bool` + +Safety notes: +- [ ] do not change command backend +- [ ] do not modify `code_puppy/command_line/` +- [ ] preserve `/help` rich-to-plain rendering behavior + +Suggested tests: +- [ ] `/help` returns rendered text +- [ ] normal command returns success payload +- [ ] exception returns `ServerCommandResult(success=False)` + +--- + +### 5) `code_puppy/api/ws/ws_stream_processor.py` — extract `drain_events_concurrent()` internals + +This is the largest remaining extraction still embedded inline. + +Move from nested function into a module-level helper: +- [ ] event batching with `asyncio.wait_for(..., timeout=0.01)` +- [ ] `agent_invoked` forwarding +- [ ] `tool_call_start` forwarding +- [ ] `tool_call_complete` forwarding +- [ ] stream `part_start` routing +- [ ] stream `part_delta` routing +- [ ] stream `part_end` routing +- [ ] final queue drain collecting last text deltas +- [ ] websocket-close detection during sends +- [ ] batch/debug logging + +Recommended public helper: +- [ ] `drain_stream_events(...)` + +Recommended support dataclass if signatures become too wide: +- [ ] `StreamProcessorDeps` holding sender callbacks, logger, metadata, etc. + +Safety notes: +- [ ] preserve exact event ordering +- [ ] preserve tool alias/group-id bookkeeping +- [ ] preserve final drain behavior after stop signal +- [ ] preserve current close-detection behavior from send exceptions + +Suggested tests: +- [ ] part_start text/thinking adaptation still identical +- [ ] tool-call start/delta/end path still identical +- [ ] final drain still captures trailing deltas +- [ ] websocket-close error stops drain gracefully + +--- + +### 6) `code_puppy/api/ws/ws_prestream_persistence.py` — isolate pre-stream SQLite writes + +Move the pre-run persistence block out of `chat_handler.py`: + +- [ ] session `upsert_session(...)` +- [ ] `ModelRequest(UserPromptPart(...))` construction +- [ ] `get_next_seq(...)` +- [ ] `insert_message(...)` +- [ ] attachment metadata JSON persistence +- [ ] warning-only failure handling + +Recommended public helper: +- [ ] `persist_user_message_before_stream(...)` + +Safety notes: +- [ ] preserve current `seq` behavior +- [ ] preserve warning-only semantics on failure +- [ ] preserve use of `original_user_message` not file-context-expanded content + +Suggested tests: +- [ ] writes user message after existing system rows +- [ ] stores attachments JSON when provided +- [ ] pre-stream failure remains non-fatal + +--- + +### 7) `code_puppy/api/ws/ws_message_orchestrator.py` — extract the `type=message` happy path + +This module should own the high-level orchestration of a user chat turn. + +Move orchestration for: +- [ ] setup of per-message context +- [ ] requested per-message model switch +- [ ] frontend `model_settings` application +- [ ] empty-message rejection +- [ ] user echo frame +- [ ] `WebSocketTurnState()` creation +- [ ] stream drain lifecycle start/stop +- [ ] send `status=thinking` +- [ ] working directory prompt context setup +- [ ] `prepare_turn_input(...)` +- [ ] `persist_user_message_before_stream(...)` +- [ ] `execute_turn_runner(...)` +- [ ] deferred switch/create-session handoff +- [ ] `resolve_post_run_resolution(...)` +- [ ] non-streaming assistant frame fallback +- [ ] B1 `stream_end` path +- [ ] `finalize_turn_history(...)` +- [ ] `persist_session_turn_and_broadcast(...)` +- [ ] final `status=done` +- [ ] error handling + `persist_error_payload(...)` +- [ ] final `cleanup_message_context(...)` + +Recommended public helper: +- [ ] `handle_chat_message(...) -> dict | None` + - returns deferred control message when one arrives during streaming + - returns `None` otherwise + +Safety notes: +- [ ] preserve exact post-run decision tree order +- [ ] preserve `stream_end` placement relative to tool results +- [ ] preserve history snapshot-before-await behavior +- [ ] preserve final `status=done` placement + +Suggested tests: +- [ ] no-result, cancelled, exception, and success paths +- [ ] non-streaming fallback still emits start/delta/end triplet +- [ ] B1 streaming emits tool results before `stream_end` +- [ ] deferred `switch_session` during active run still background-saves and returns control + +--- + +### 8) `code_puppy/api/ws/chat_handler.py` — reduce to thin endpoint wiring + +After the extractions above, `chat_handler.py` should only: + +- [ ] accept websocket +- [ ] construct `WebSocketSender` +- [ ] call session bootstrap helper +- [ ] loop on `receive_json()` +- [ ] route to `handle_command_message(...)`, `handle_control_message(...)`, and `handle_chat_message(...)` +- [ ] if chat helper returns deferred message, re-dispatch it +- [ ] handle top-level disconnect/runtime fallback +- [ ] call final cleanup helper(s) + +Hard stop criteria before calling this “done”: +- [ ] no nested `drain_events_concurrent()` function remains +- [ ] no 150+ line control-message branch chain remains +- [ ] top-level loop is readable without scrolling through protocol details + +Target outcome: +- [ ] `chat_handler.py` becomes endpoint composition/orchestration only +- [ ] most behavior-specific tests point at extracted modules, not the giant handler + +--- + +## Recommended execution order + +### Bead `code_puppy-40g` — planning + bootstrap/control split +- [ ] add `ws_chat_runtime.py` +- [ ] add `ws_session_bootstrap.py` +- [ ] add `ws_control_messages.py` +- [ ] trim corresponding branches from `chat_handler.py` + +### Bead `code_puppy-79a` — stream/message execution split +- [ ] add `ws_stream_processor.py` +- [ ] add `ws_prestream_persistence.py` +- [ ] add `ws_message_orchestrator.py` +- [ ] trim `type=message` branch from `chat_handler.py` + +### Bead `code_puppy-8ql` — regression coverage +- [ ] add focused tests for every new helper module +- [ ] keep existing WS tests green +- [ ] cover all four risk categories: happy-path / edge / boundary / negative + +### Bead `code_puppy-e7v` — final docs/validation +- [ ] update this checklist with final landed module set +- [ ] capture residual risks +- [ ] record exact validation commands + outcomes + +--- + +## Validation checklist + +### Minimum automated validation +- [ ] `ruff check --fix code_puppy/api/ws` +- [ ] `ruff format code_puppy/api/ws docs/migration/ws-chat-handler-refactor-checklist.md` +- [ ] `pytest -q code_puppy/api/ws/tests` +- [ ] `pytest -q tests/test_messaging_bus.py` + +### Coverage matrix requirement for this refactor + +| Risk tag | Required coverage | Status | +|---|---|---| +| happy-path | normal new session + message + stream success | [ ] | +| edge | restored session replay + switch-session during stream | [ ] | +| boundary | unchanged CWD, empty message, unknown config key | [ ] | +| negative | invalid session id, command failure, send failure, cancel/disconnect | [ ] | + +### Adversarial review prompts +- [ ] What bug would still pass if frame order changed around `stream_end`? +- [ ] What bug would still pass if resumed-session system banners stopped replaying? +- [ ] What bug would still pass if `switch_session` forgot to mark old session inactive? +- [ ] What bug would still pass if pre-stream SQLite insert wrote the wrong message content? +- [ ] What bug would still pass if a close/send exception did not stop the drain loop? + +### Manual spot checks +- [ ] resume an existing WS session and confirm restored messages + system banners appear +- [ ] switch sessions without reconnecting +- [ ] interrupt a running command/tool +- [ ] set working directory twice and confirm second event is `unchanged=True` +- [ ] run at least one slash command and `/help` + +--- + +## Definition of done + +This refactor is only done when all of the following are true: + +- [ ] `chat_handler.py` is mostly endpoint wiring, not business logic +- [ ] each extracted file has a single narrow responsibility +- [ ] existing frontend-visible behavior is unchanged +- [ ] focused tests cover the extracted helpers +- [ ] WS regression suite passes +- [ ] messaging regression suite passes +- [ ] residual risks are documented explicitly +- [ ] this checklist is updated to reflect the final landed state + diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 04ed5b3a8..218bbe999 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -374,8 +374,8 @@ def test_load_with_invalid_session_id_raises_error(self, temp_session_dir): with pytest.raises(ValueError, match="must be kebab-case"): _load_session_history("Invalid_Session") - def test_save_creates_pkl_and_txt_files(self, temp_session_dir, mock_messages): - """Test that save creates both .pkl and .txt files.""" + def test_save_creates_json_and_txt_files(self, temp_session_dir, mock_messages): + """Test that save creates both .json and .txt files.""" session_id = "test-session" agent_name = "test-agent" initial_prompt = "Test prompt" @@ -392,9 +392,9 @@ def test_save_creates_pkl_and_txt_files(self, temp_session_dir, mock_messages): ) # Check that both files exist - pkl_file = temp_session_dir / f"{session_id}.pkl" + session_file = temp_session_dir / f"{session_id}.json" txt_file = temp_session_dir / f"{session_id}.txt" - assert pkl_file.exists() + assert session_file.exists() assert txt_file.exists() def test_txt_file_contains_readable_metadata(self, temp_session_dir, mock_messages): @@ -465,18 +465,17 @@ def test_txt_file_updates_on_subsequent_saves( # last_updated should exist assert "last_updated" in metadata - def test_load_handles_corrupted_pickle(self, temp_session_dir): - """Test that loading a corrupted pickle file returns empty list.""" + def test_load_handles_corrupted_session_json(self, temp_session_dir): + """Test that loading a corrupted JSON session file returns empty list.""" session_id = "corrupted-session" with patch( "code_puppy.tools.agent_tools._get_subagent_sessions_dir", return_value=temp_session_dir, ): - # Create a corrupted pickle file - pkl_file = temp_session_dir / f"{session_id}.pkl" - with open(pkl_file, "wb") as f: - f.write(b"This is not a valid pickle file!") + # Create a corrupted JSON session file + session_file = temp_session_dir / f"{session_id}.json" + session_file.write_text("{not valid json", encoding="utf-8") # Should return empty list instead of crashing loaded_messages = _load_session_history(session_id) diff --git a/tests/test_load_context_completion.py b/tests/test_load_context_completion.py index 54ce0ceee..d769a08b5 100644 --- a/tests/test_load_context_completion.py +++ b/tests/test_load_context_completion.py @@ -49,21 +49,21 @@ def test_session_name_completion(self): contexts_dir.mkdir() # Create test context files - (contexts_dir / "session1.pkl").touch() - (contexts_dir / "session2.pkl").touch() - (contexts_dir / "another_session.pkl").touch() - (contexts_dir / "not_a_pkl.txt").touch() # Should be ignored + (contexts_dir / "session1.json").touch() + (contexts_dir / "session2.json").touch() + (contexts_dir / "another_session.json").touch() + (contexts_dir / "not_a_session.txt").touch() # Should be ignored # Test completion with space doc = Document("/load_context ") completions = list(self.completer.get_completions(doc, None)) - # Should suggest all .pkl files (without extension) + # Should suggest all .json files (without extension) completion_texts = [c.text for c in completions] assert "session1" in completion_texts assert "session2" in completion_texts assert "another_session" in completion_texts - assert "not_a_pkl" not in completion_texts # .txt files ignored + assert "not_a_session" not in completion_texts # .txt files ignored # All should have proper metadata for completion in completions: @@ -81,9 +81,9 @@ def test_partial_session_name_completion(self): contexts_dir.mkdir() # Create test context files - (contexts_dir / "session1.pkl").touch() - (contexts_dir / "session2.pkl").touch() - (contexts_dir / "another_session.pkl").touch() + (contexts_dir / "session1.json").touch() + (contexts_dir / "session2.json").touch() + (contexts_dir / "another_session.json").touch() # Test completion with partial match doc = Document("/load_context sess") diff --git a/tests/test_messaging_bus.py b/tests/test_messaging_bus.py index 6ee0c6282..3ce770d84 100644 --- a/tests/test_messaging_bus.py +++ b/tests/test_messaging_bus.py @@ -34,7 +34,7 @@ def test_initialization_default(self): assert bus._maxsize == 1000 assert isinstance(bus._outgoing, queue.Queue) assert isinstance(bus._incoming, queue.Queue) - assert bus._current_session_id is None + assert bus.get_session_context() is None assert not bus._has_active_renderer assert bus._startup_buffer == [] @@ -587,3 +587,35 @@ def update_context(session_id): t.join() assert len(contexts) == 5 + + @pytest.mark.asyncio + async def test_session_context_is_task_local(self): + """Each asyncio task should see its own session context.""" + bus = MessageBus() + bus._has_active_renderer = True + + ready_first = asyncio.Event() + ready_second = asyncio.Event() + release = asyncio.Event() + + async def worker(session_id, ready_to_set, ready_done): + if ready_to_set is not None: + await ready_to_set.wait() + bus.set_session_context(session_id) + ready_done.set() + await release.wait() + msg = TextMessage(level=MessageLevel.INFO, text=session_id) + bus.emit(msg) + return bus.get_session_context(), msg.session_id + + first = asyncio.create_task(worker("session-1", None, ready_first)) + second = asyncio.create_task(worker("session-2", ready_first, ready_second)) + + await ready_second.wait() + release.set() + + first_ctx, second_ctx = await asyncio.gather(first, second) + + assert first_ctx == ("session-1", "session-1") + assert second_ctx == ("session-2", "session-2") + assert bus.get_session_context() is None diff --git a/tests/test_session_storage.py b/tests/test_session_storage.py index 339f9dc26..5b82f8f18 100644 --- a/tests/test_session_storage.py +++ b/tests/test_session_storage.py @@ -1,6 +1,5 @@ from __future__ import annotations -import json import os from pathlib import Path from typing import Callable, List @@ -27,25 +26,18 @@ def token_estimator() -> Callable[[object], int]: def test_save_and_load_session(tmp_path: Path, history: List[str], token_estimator): session_name = "demo_session" - timestamp = "2024-01-01T00:00:00" metadata = save_session( history=history, session_name=session_name, base_dir=tmp_path, - timestamp=timestamp, + timestamp="2024-01-01T00:00:00", token_estimator=token_estimator, ) - assert metadata.session_name == session_name assert metadata.message_count == len(history) assert metadata.total_tokens == sum(token_estimator(m) for m in history) - assert metadata.pickle_path.exists() - assert metadata.metadata_path.exists() - - with metadata.metadata_path.open() as meta_file: - stored = json.load(meta_file) - assert stored["session_name"] == session_name - assert stored["auto_saved"] is False + assert metadata.session_file_path.exists() + assert metadata.session_file_path.name == f"{session_name}.json" loaded_history = load_session(session_name, tmp_path) assert loaded_history == history @@ -75,7 +67,7 @@ def test_cleanup_sessions(tmp_path: Path, history: List[str], token_estimator): timestamp="2024-01-01T00:00:00", token_estimator=token_estimator, ) - os.utime(metadata.pickle_path, (0, index)) + os.utime(metadata.session_file_path, (0, index)) removed = cleanup_sessions(tmp_path, 2) assert removed == ["session_earliest"] diff --git a/tests/test_session_storage_edge_cases.py b/tests/test_session_storage_edge_cases.py index 4bb29e44f..0178fab41 100644 --- a/tests/test_session_storage_edge_cases.py +++ b/tests/test_session_storage_edge_cases.py @@ -1,89 +1,53 @@ -"""Comprehensive edge case tests for session_storage.py. +"""Edge case tests for JSON-only session storage.""" -Focuses on: -- Session corruption recovery -- Concurrent access scenarios -- Cleanup logic edge cases -- Error handling in save/load -- Metadata handling -""" +from __future__ import annotations -import json -import pickle +import os from pathlib import Path -from unittest.mock import patch import pytest from code_puppy import session_storage -from code_puppy.session_storage import _LEGACY_SIGNATURE_SIZE, _LEGACY_SIGNED_HEADER class TestSessionPathEdgeCases: - """Test session path building edge cases.""" - - def test_build_session_paths_with_special_characters(self, tmp_path): - """Test that session names with special characters are handled.""" - # Special characters that might appear in session names + def test_get_session_file_path_with_special_characters(self, tmp_path: Path): session_name = "session-with_special.chars" - paths = session_storage.build_session_paths(tmp_path, session_name) - - assert session_name in str(paths.pickle_path) - assert session_name in str(paths.metadata_path) - - def test_ensure_directory_with_nested_paths(self, tmp_path): - """Test that ensure_directory creates nested paths.""" - nested_path = tmp_path / "a" / "b" / "c" / "d" - assert not nested_path.exists() - - result = session_storage.ensure_directory(nested_path) + path = session_storage.get_session_file_path(tmp_path, session_name) - assert nested_path.exists() - assert result == nested_path + assert session_name in str(path) + assert path.suffix == ".json" class TestSessionSaveEdgeCases: - """Test edge cases in session saving.""" - - def test_save_session_with_empty_history(self, tmp_path): - """Test saving an empty session.""" - - def mock_token_estimator(msg): - return 0 - + def test_save_session_with_empty_history(self, tmp_path: Path): metadata = session_storage.save_session( history=[], session_name="empty", base_dir=tmp_path, timestamp="2024-01-01T00:00:00", - token_estimator=mock_token_estimator, + token_estimator=lambda _msg: 0, ) assert metadata.message_count == 0 assert metadata.total_tokens == 0 - assert (tmp_path / "empty.pkl").exists() + assert (tmp_path / "empty.json").exists() - def test_save_session_with_large_messages(self, tmp_path): - """Test saving large messages in history.""" - large_content = "x" * 10000 # 10KB message + def test_save_session_with_large_messages(self, tmp_path: Path): + large_content = "x" * 10000 history = [{"role": "user", "content": large_content}] - def mock_token_estimator(msg): - return len(msg.get("content", "")) - metadata = session_storage.save_session( history=history, session_name="large", base_dir=tmp_path, - timestamp="2024-01-01T00:00:00", - token_estimator=mock_token_estimator, + token_estimator=lambda msg: len(msg.get("content", "")), ) assert metadata.message_count == 1 assert metadata.total_tokens == 10000 - def test_save_session_with_complex_objects(self, tmp_path): - """Test saving messages with complex nested objects.""" + def test_save_and_load_complex_nested_objects(self, tmp_path: Path): history = [ { "role": "user", @@ -95,122 +59,32 @@ def test_save_session_with_complex_objects(self, tmp_path): } ] - def mock_token_estimator(msg): - return 10 - session_storage.save_session( - history=history, - session_name="complex", - base_dir=tmp_path, - timestamp="2024-01-01T00:00:00", - token_estimator=mock_token_estimator, + history=history, session_name="complex", base_dir=tmp_path ) - - # Load and verify loaded = session_storage.load_session("complex", tmp_path) - assert loaded[0]["metadata"]["nested"]["deeply"]["structured"] == "data" - def test_save_session_with_none_token_values(self, tmp_path): - """Test saving with token estimator returning None values.""" - history = [{"role": "user", "content": "test"}] - - def mock_token_estimator(msg): - # Simulate unpredictable token estimation - return 0 # Return 0 instead of None - - metadata = session_storage.save_session( - history=history, - session_name="test", - base_dir=tmp_path, - timestamp="2024-01-01T00:00:00", - token_estimator=mock_token_estimator, - ) - - assert metadata.total_tokens == 0 - - def test_save_session_marks_auto_saved_correctly(self, tmp_path): - """Test that auto_saved flag is persisted correctly.""" - - def mock_token_estimator(msg): - return 10 - - # Save as auto-saved - metadata1 = session_storage.save_session( - history=[{"role": "user", "content": "msg"}], - session_name="auto", - base_dir=tmp_path, - timestamp="2024-01-01T00:00:00", - token_estimator=mock_token_estimator, - auto_saved=True, - ) - assert metadata1.auto_saved is True - - # Save as manual - metadata2 = session_storage.save_session( - history=[{"role": "user", "content": "msg"}], - session_name="manual", - base_dir=tmp_path, - timestamp="2024-01-01T00:00:00", - token_estimator=mock_token_estimator, - auto_saved=False, - ) - assert metadata2.auto_saved is False + assert loaded[0]["metadata"]["nested"]["deeply"]["structured"] == "data" class TestSessionLoadEdgeCases: - """Test edge cases in session loading.""" + def test_load_session_with_missing_file_raises(self, tmp_path: Path): + with pytest.raises(FileNotFoundError): + session_storage.load_session("missing", tmp_path) - def test_load_session_with_corrupted_pickle(self, tmp_path): - """Test that loading corrupted pickle raises appropriate error.""" - # Create a pickle file with corrupted data - pickle_path = tmp_path / "corrupted.pkl" - with open(pickle_path, "wb") as f: - f.write(b"this is not valid pickle data") + def test_load_session_with_invalid_json_raises(self, tmp_path: Path): + (tmp_path / "corrupted.json").write_text("{not valid json", encoding="utf-8") - with pytest.raises(Exception): # pickle.UnpicklingError or similar + with pytest.raises(Exception): session_storage.load_session("corrupted", tmp_path) - def test_load_session_with_empty_pickle(self, tmp_path): - """Test loading empty pickle file.""" - pkl_data = pickle.dumps([]) - pickle_path = tmp_path / "empty.pkl" - pickle_path.write_bytes( - _LEGACY_SIGNED_HEADER + (b"x" * _LEGACY_SIGNATURE_SIZE) + pkl_data - ) - - loaded = session_storage.load_session("empty", tmp_path) - assert loaded == [] - - def test_load_session_returns_exact_history(self, tmp_path): - """Test that loaded history matches saved history exactly.""" - original = [ - {"id": 1, "data": [1, 2, 3]}, - {"id": 2, "data": {"nested": True}}, - None, # Test None values - "string", # Test string values - ] - - pkl_data = pickle.dumps(original) - pickle_path = tmp_path / "test.pkl" - pickle_path.write_bytes( - _LEGACY_SIGNED_HEADER + (b"x" * _LEGACY_SIGNATURE_SIZE) + pkl_data - ) - - loaded = session_storage.load_session("test", tmp_path) - assert loaded == original - class TestSessionListingEdgeCases: - """Test edge cases in session listing.""" - - def test_list_sessions_ignores_non_pkl_files(self, tmp_path): - """Test that only .pkl files are counted.""" - # Create various files - (tmp_path / "session_1.pkl").touch() - (tmp_path / "session_2.pkl").touch() - (tmp_path / "session_1_meta.json").touch() # Should be ignored - (tmp_path / "random.txt").touch() # Should be ignored - (tmp_path / "session_2.bak").touch() # Should be ignored + def test_list_sessions_ignores_non_json_files(self, tmp_path: Path): + (tmp_path / "session_1.json").touch() + (tmp_path / "session_2.json").touch() + (tmp_path / "random.txt").touch() + (tmp_path / "session_2.bak").touch() result = session_storage.list_sessions(tmp_path) @@ -218,232 +92,27 @@ def test_list_sessions_ignores_non_pkl_files(self, tmp_path): assert "session_1" in result assert "session_2" in result - def test_list_sessions_returns_sorted_names(self, tmp_path): - """Test that session names are returned sorted.""" + def test_list_sessions_returns_sorted_names(self, tmp_path: Path): names = ["z_session", "a_session", "m_session"] for name in names: - (tmp_path / f"{name}.pkl").touch() + (tmp_path / f"{name}.json").touch() - result = session_storage.list_sessions(tmp_path) - - assert result == ["a_session", "m_session", "z_session"] - - def test_list_sessions_handles_malformed_names(self, tmp_path): - """Test that sessions with special names are handled.""" - special_names = ["normal", "with-dash", "with_underscore"] - for name in special_names: - (tmp_path / f"{name}.pkl").touch() - - result = session_storage.list_sessions(tmp_path) - - assert len(result) == 3 - assert all(name in result for name in special_names) + assert session_storage.list_sessions(tmp_path) == [ + "a_session", + "m_session", + "z_session", + ] class TestSessionCleanupEdgeCases: - """Test edge cases in session cleanup.""" - - def test_cleanup_sessions_respects_max_limit(self, tmp_path): - """Test that cleanup respects the max_sessions limit.""" - # Create 10 sessions + def test_cleanup_sessions_respects_max_limit(self, tmp_path: Path): for i in range(10): - session_path = tmp_path / f"session_{i:02d}.pkl" + session_path = tmp_path / f"session_{i:02d}.json" session_path.touch() - # Artificially age them - import os - os.utime(session_path, (i, i)) - # Keep only 3 removed = session_storage.cleanup_sessions(tmp_path, max_sessions=3) - remaining = session_storage.list_sessions(tmp_path) - assert len(remaining) == 3 - assert len(removed) == 7 - - def test_cleanup_sessions_handles_metadata_file_errors(self, tmp_path): - """Test that cleanup continues even if metadata can't be deleted.""" - # Create session files - (tmp_path / "old.pkl").touch() - (tmp_path / "old_meta.json").touch() - (tmp_path / "new.pkl").touch() - - import os - - # Age the old one - os.utime(tmp_path / "old.pkl", (1, 1)) - os.utime(tmp_path / "old_meta.json", (1, 1)) - os.utime(tmp_path / "new.pkl", (999, 999)) - # Should clean up without raising errors - removed = session_storage.cleanup_sessions(tmp_path, max_sessions=1) - assert "old" in removed - - def test_cleanup_sessions_handles_permission_errors(self, tmp_path): - """Test that cleanup gracefully handles permission errors.""" - # Create a session - pickle_path = tmp_path / "session.pkl" - pickle_path.touch() - - # Mock the unlink method to raise OSError - with patch.object(Path, "unlink", side_effect=OSError("Permission denied")): - # Should not raise, but continue - session_storage.cleanup_sessions(tmp_path, max_sessions=0) - # The session should still be listed since delete failed - assert "session" in session_storage.list_sessions(tmp_path) - - -class TestSessionMetadataEdgeCases: - """Test edge cases in session metadata.""" - - def test_session_metadata_serialization_with_paths(self): - """Test metadata serialization includes file path.""" - metadata = session_storage.SessionMetadata( - session_name="test", - timestamp="2024-01-01T00:00:00", - message_count=5, - total_tokens=100, - pickle_path=Path("/tmp/test.pkl"), - metadata_path=Path("/tmp/test_meta.json"), - ) - - serialized = metadata.as_serialisable() - - assert serialized["file_path"] == "/tmp/test.pkl" - assert "metadata_path" not in serialized # Should not be in serialization - - def test_session_metadata_preserves_all_fields(self): - """Test that metadata serialization preserves all important fields.""" - metadata = session_storage.SessionMetadata( - session_name="comprehensive_test", - timestamp="2024-06-15T12:30:45.123456", - message_count=42, - total_tokens=9999, - pickle_path=Path("/custom/path/session.pkl"), - metadata_path=Path("/custom/path/session_meta.json"), - auto_saved=True, - ) - - serialized = metadata.as_serialisable() - - assert serialized["session_name"] == "comprehensive_test" - assert serialized["timestamp"] == "2024-06-15T12:30:45.123456" - assert serialized["message_count"] == 42 - assert serialized["total_tokens"] == 9999 - assert serialized["auto_saved"] is True - - -class TestSessionRoundTrip: - """Test complete save and load round-trips.""" - - def test_save_and_load_preserves_data(self, tmp_path): - """Test that data is preserved through save and load.""" - original_history = [ - {"role": "system", "content": "You are helpful"}, - {"role": "user", "content": "What's 2+2?"}, - {"role": "assistant", "content": "4"}, - {"role": "user", "content": "Prove it"}, - {"role": "assistant", "content": "2+2=4 because..."}, - ] - - def mock_token_estimator(msg): - return len(msg.get("content", "").split()) - - # Save - saved_metadata = session_storage.save_session( - history=original_history, - session_name="roundtrip", - base_dir=tmp_path, - timestamp="2024-01-01T00:00:00", - token_estimator=mock_token_estimator, - auto_saved=False, - ) - - # Load - loaded_history = session_storage.load_session("roundtrip", tmp_path) - - # Verify - assert loaded_history == original_history - assert saved_metadata.message_count == len(original_history) - - def test_multiple_sessions_independent(self, tmp_path): - """Test that multiple sessions don't interfere.""" - - def mock_token_estimator(msg): - return 10 - - # Create session 1 - session_storage.save_session( - history=[{"role": "user", "content": "Session 1"}], - session_name="session_1", - base_dir=tmp_path, - timestamp="2024-01-01T00:00:00", - token_estimator=mock_token_estimator, - ) - - # Create session 2 - session_storage.save_session( - history=[{"role": "user", "content": "Session 2"}], - session_name="session_2", - base_dir=tmp_path, - timestamp="2024-01-02T00:00:00", - token_estimator=mock_token_estimator, - ) - - # Load and verify they're different - hist1 = session_storage.load_session("session_1", tmp_path) - hist2 = session_storage.load_session("session_2", tmp_path) - - assert hist1[0]["content"] == "Session 1" - assert hist2[0]["content"] == "Session 2" - - -class TestSessionMetadataFileHandling: - """Test metadata JSON file handling.""" - - def test_metadata_file_creation(self, tmp_path): - """Test that metadata JSON file is properly created.""" - - def mock_token_estimator(msg): - return 10 - - session_storage.save_session( - history=[{"role": "user", "content": "test"}], - session_name="meta_test", - base_dir=tmp_path, - timestamp="2024-01-01T10:30:45", - token_estimator=mock_token_estimator, - auto_saved=True, - ) - - # Check metadata file exists and is valid JSON - meta_path = tmp_path / "meta_test_meta.json" - assert meta_path.exists() - - with open(meta_path) as f: - metadata = json.load(f) - - assert metadata["session_name"] == "meta_test" - assert metadata["timestamp"] == "2024-01-01T10:30:45" - assert metadata["auto_saved"] is True - - def test_metadata_file_is_readable_json(self, tmp_path): - """Test that metadata files are valid, readable JSON.""" - - def mock_token_estimator(msg): - return 5 - - session_storage.save_session( - history=[{"role": "assistant", "content": "Hello there"}], - session_name="json_test", - base_dir=tmp_path, - timestamp="2024-03-15T08:00:00", - token_estimator=mock_token_estimator, - ) - - # Verify we can read it back - meta_path = tmp_path / "json_test_meta.json" - metadata = json.loads(meta_path.read_text()) - - assert metadata["message_count"] == 1 - assert metadata["total_tokens"] == 5 + assert len(removed) == 7 + assert len(remaining) == 3 diff --git a/tests/test_session_storage_json_only.py b/tests/test_session_storage_json_only.py new file mode 100644 index 000000000..78ceab8ce --- /dev/null +++ b/tests/test_session_storage_json_only.py @@ -0,0 +1,16 @@ +from pathlib import Path + +from code_puppy.session_storage import get_session_file_path, save_session + + +def test_save_session_returns_json_file_path(tmp_path: Path): + metadata = save_session( + history=[{"role": "user", "content": "hello"}], + session_name="demo", + base_dir=tmp_path, + ) + + assert metadata.message_count == 1 + assert metadata.session_file_path == get_session_file_path(tmp_path, "demo") + assert metadata.session_file_path.suffix == ".json" + assert metadata.session_file_path.exists() From 534225d21279af58d7b62e18c9fded14aca0db2d Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Wed, 10 Jun 2026 19:20:45 -0500 Subject: [PATCH 19/29] fix(ws): add one-shot resume recovery with sqlite reload --- code_puppy/api/ws/chat_handler.py | 94 +++++++++++++++- .../api/ws/tests/test_ws_resume_recovery.py | 42 ++++++++ code_puppy/api/ws/ws_resume_recovery.py | 102 ++++++++++++++++++ 3 files changed, 236 insertions(+), 2 deletions(-) create mode 100644 code_puppy/api/ws/tests/test_ws_resume_recovery.py create mode 100644 code_puppy/api/ws/ws_resume_recovery.py diff --git a/code_puppy/api/ws/chat_handler.py b/code_puppy/api/ws/chat_handler.py index 94fa92267..6da5d1778 100644 --- a/code_puppy/api/ws/chat_handler.py +++ b/code_puppy/api/ws/chat_handler.py @@ -29,6 +29,9 @@ 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, @@ -64,6 +67,7 @@ ServerCancelled, ServerError, ServerStatus, + ServerSystem, ServerStreamEnd, ServerUserMessage, ) @@ -882,8 +886,94 @@ async def send_status_only_pending_tool_results(): await safe_send_json(frame) continue if post_run.no_result_error is not None: - await send_typed(post_run.no_result_error) - continue + 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 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..9d1cd35bf --- /dev/null +++ b/code_puppy/api/ws/tests/test_ws_resume_recovery.py @@ -0,0 +1,42 @@ +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/ws_resume_recovery.py b/code_puppy/api/ws/ws_resume_recovery.py new file mode 100644 index 000000000..034e3c4f0 --- /dev/null +++ b/code_puppy/api/ws/ws_resume_recovery.py @@ -0,0 +1,102 @@ +"""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", +] From 80a37aebea761ea969e6a98e1477fe11d384352c Mon Sep 17 00:00:00 2001 From: Code Puppy Maestro Date: Sat, 13 Jun 2026 01:08:03 -0500 Subject: [PATCH 20/29] fix(codex): dedupe response input items by id before API request --- code_puppy/chatgpt_codex_client.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/code_puppy/chatgpt_codex_client.py b/code_puppy/chatgpt_codex_client.py index 7e00716e4..26e606b7f 100644 --- a/code_puppy/chatgpt_codex_client.py +++ b/code_puppy/chatgpt_codex_client.py @@ -211,6 +211,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"] From af4aaf057c3aec60cf9ed677b21725ca90e77fd5 Mon Sep 17 00:00:00 2001 From: Code Puppy Maestro Date: Sat, 13 Jun 2026 01:16:47 -0500 Subject: [PATCH 21/29] test(codex): cover duplicate input item id dedupe behavior --- tests/test_chatgpt_codex_client.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/test_chatgpt_codex_client.py b/tests/test_chatgpt_codex_client.py index 15c62a950..737dd1ffd 100644 --- a/tests/test_chatgpt_codex_client.py +++ b/tests/test_chatgpt_codex_client.py @@ -249,6 +249,34 @@ def test_remove_all_unsupported_params(self): assert "verbosity" not in data assert data["temperature"] == 0.7 # Preserved + def test_dedupes_duplicate_input_item_ids(self): + """Duplicate input item IDs should be removed while preserving first-seen order.""" + body = json.dumps( + { + "model": "gpt-5", + "input": [ + {"id": "rs_1", "type": "reasoning", "content": "first"}, + {"id": "rs_1", "type": "reasoning", "content": "duplicate"}, + {"id": "rs_2", "type": "reasoning", "content": "second"}, + {"type": "message", "content": "no-id item"}, + ], + } + ).encode() + + result, _ = ChatGPTCodexAsyncClient._inject_codex_fields(body) + + assert result is not None + data = json.loads(result) + input_items = data["input"] + + # rs_1 duplicate should be removed; first occurrence preserved. + assert [item.get("id") for item in input_items if isinstance(item, dict) and item.get("id")] == [ + "rs_1", + "rs_2", + ] + # non-id items should still be present + assert any(isinstance(item, dict) and item.get("type") == "message" for item in input_items) + def test_invalid_json_returns_none(self): """Test that invalid JSON returns (None, False).""" body = b"not valid json" From 55c20613cc8a710cb3b74c5dbace04275337c326 Mon Sep 17 00:00:00 2001 From: shravya maddipudi Date: Sat, 13 Jun 2026 12:08:28 -0500 Subject: [PATCH 22/29] feat(desk): refresh chat UI and session recovery Improve the desk chat template with wireframe-inspired message and tool-call rendering, including safer markdown/tool-card handling and collapsed tool call payloads. Strengthen resumed-session recovery, session CWD handling, permission auto-approval behavior, SQL query loading, and command runner cleanup coverage. Validation: node --check /tmp/chat_inline_pr_handoff_check.js; ruff check targeted files; pytest -q --no-cov targeted tests (11 passed). --- .gitignore | 2 - code_puppy/api/db/queries.py | 133 +- .../api/db/sql/session_history_parity.sql | 59 + .../session_history_parity_no_compacted.sql | 59 + code_puppy/api/db/sql_loader.py | 22 + code_puppy/api/permission_plugin.py | 14 +- code_puppy/api/permissions.py | 168 +- code_puppy/api/session_cwd.py | 39 + code_puppy/api/templates/chat.html | 1357 +++++++++++--- .../tests/test_chat_template_xss_guards.py | 21 + code_puppy/api/tests/test_db_sql_loader.py | 22 + .../test_permission_plugin_fail_closed.py | 53 + .../tests/test_permissions_session_binding.py | 76 + code_puppy/api/ws/chat_handler.py | 25 +- code_puppy/api/ws/chat_turn_runner.py | 4 +- .../api/ws/tests/test_chat_turn_runner.py | 6 +- .../api/ws/tests/test_ws_command_handler.py | 38 +- .../api/ws/tests/test_ws_control_messages.py | 72 +- code_puppy/api/ws/ws_command_handler.py | 4 +- code_puppy/api/ws/ws_control_messages.py | 8 +- code_puppy/tools/command_runner.py | 32 +- docs/migration/puppy-desk-cleanup-summary.md | 91 - .../puppy-desk-gate1-investigation-summary.md | 426 ----- .../puppy-desk-migration-plan-options.md | 825 --------- docs/migration/puppy-desk-phase3-fixes.md | 195 -- docs/migration/puppy-desk-phase3-plan.md | 127 -- .../puppy-desk-uncommitted-change-snapshot.md | 1567 ----------------- .../ws-chat-handler-refactor-checklist.md | 464 ----- .../tools/test_command_runner_lock_release.py | 47 + 29 files changed, 1685 insertions(+), 4271 deletions(-) create mode 100644 code_puppy/api/db/sql/session_history_parity.sql create mode 100644 code_puppy/api/db/sql/session_history_parity_no_compacted.sql create mode 100644 code_puppy/api/db/sql_loader.py create mode 100644 code_puppy/api/session_cwd.py create mode 100644 code_puppy/api/tests/test_chat_template_xss_guards.py create mode 100644 code_puppy/api/tests/test_db_sql_loader.py create mode 100644 code_puppy/api/tests/test_permission_plugin_fail_closed.py create mode 100644 code_puppy/api/tests/test_permissions_session_binding.py delete mode 100644 docs/migration/puppy-desk-cleanup-summary.md delete mode 100644 docs/migration/puppy-desk-gate1-investigation-summary.md delete mode 100644 docs/migration/puppy-desk-migration-plan-options.md delete mode 100644 docs/migration/puppy-desk-phase3-fixes.md delete mode 100644 docs/migration/puppy-desk-phase3-plan.md delete mode 100644 docs/migration/puppy-desk-uncommitted-change-snapshot.md delete mode 100644 docs/migration/ws-chat-handler-refactor-checklist.md create mode 100644 tests/tools/test_command_runner_lock_release.py diff --git a/.gitignore b/.gitignore index 603b7981f..56e9f3c35 100644 --- a/.gitignore +++ b/.gitignore @@ -45,5 +45,3 @@ code_puppy/bundled_skills/ .claude/hooks/ts-hooks/dist/ .json - -.maestro/* diff --git a/code_puppy/api/db/queries.py b/code_puppy/api/db/queries.py index b5adc75e2..1a26ad0b1 100644 --- a/code_puppy/api/db/queries.py +++ b/code_puppy/api/db/queries.py @@ -11,6 +11,7 @@ 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__) @@ -1080,67 +1081,10 @@ async def write_turn_to_sqlite( # --------------------------------------------------------------------------- # SQL query for interleaved messages + tool_calls, ordered by seq -_SESSION_HISTORY_PARITY_SQL = """ -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; -""" +_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( @@ -1180,69 +1124,10 @@ async def get_session_history_parity( (session_id, session_id), ) else: - # Modify the query to exclude compacted messages - sql = """ -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; -""" - cursor = await db.execute(sql, (session_id, session_id)) + 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] 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/permission_plugin.py b/code_puppy/api/permission_plugin.py index 411a4b799..91a2e9f3d 100644 --- a/code_puppy/api/permission_plugin.py +++ b/code_puppy/api/permission_plugin.py @@ -131,8 +131,11 @@ async def pre_tool_call_permission( except Exception as e: logger.error("[Permission] Error requesting permission: %s", e) - # On error, allow the tool (fail open for now) - return None + return { + "error": "Permission system error", + "blocked": True, + "tool_name": tool_name, + } async def shell_command_permission( @@ -195,8 +198,11 @@ async def shell_command_permission( except Exception as e: logger.error("[Permission] Error requesting permission for command: %s", e) - # On error, allow the command (fail open for now) - return None + return { + "blocked": True, + "error": "Permission system error", + "reasoning": f"Permission request failed for command: {command}", + } def register_permission_callbacks(): diff --git a/code_puppy/api/permissions.py b/code_puppy/api/permissions.py index c820693dc..c666f082d 100644 --- a/code_puppy/api/permissions.py +++ b/code_puppy/api/permissions.py @@ -1,31 +1,14 @@ """ -Permission System for Tool Call Execution - -This module provides a permission request system that integrates with the WebSocket API -to request user approval before executing potentially dangerous operations. - -Features: -- WebSocket-based permission requests with in-chat action buttons -- 5-minute timeout on permission waits -- Automatic approval in yolo mode -- CLI mode support via WebSocket - -Usage: - from code_puppy.api.permissions import request_permission - - approved = await request_permission( - websocket=ws, - session_id="session-123", - request_type="shell_command", - title="Execute Shell Command", - description="Run: ls -la", - details={"command": "ls -la", "cwd": "/home/user"} - ) +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 @@ -34,9 +17,15 @@ logger = logging.getLogger(__name__) + +@dataclass(slots=True) +class PendingPermissionRequest: + future: asyncio.Future + session_id: str + + # Global dictionary to track pending permission requests -# This is managed by the WebSocket API module -permission_futures: Dict[str, asyncio.Future] = {} +permission_futures: Dict[str, PendingPermissionRequest] = {} async def request_permission( @@ -48,51 +37,33 @@ async def request_permission( details: Dict[str, Any], timeout: int = 300, ) -> bool: - """ - Request permission from user via WebSocket. - - Args: - websocket: WebSocket connection to send request through - session_id: Current session ID - request_type: Type of permission (e.g., 'shell_command') - title: Permission dialog title - description: Human-readable description - details: Additional context (command, args, etc.) - timeout: Timeout in seconds (default 5 minutes) - - Returns: - bool: True if approved, False if denied or timeout - """ - # Check if yolo mode is enabled (auto-approve everything) + """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: - import json - import os - - from code_puppy.config import CONFIG_FILE - - if os.path.exists(CONFIG_FILE): - try: - with open(CONFIG_FILE, "r") as f: - config = json.load(f) - yolo_mode = str(config.get("yolo_mode", "false")).lower() == "true" - if yolo_mode: - logger.info( - f"[Permission] YOLO mode enabled, auto-approving {request_type}" - ) - return True - except Exception: - pass # If config read fails, continue to request permission - except Exception: - pass # If we can't check, require permission - - # If no WebSocket connection, deny by default (fail-safe) + 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()) - # Send permission request to frontend try: await websocket.send_json( ServerPermissionRequest( @@ -110,17 +81,20 @@ async def request_permission( logger.error("[Permission] Failed to send request: %s", e) return False - # Create a future to wait for the response - future = asyncio.Future() - permission_futures[request_id] = future + future: asyncio.Future = asyncio.Future() + permission_futures[request_id] = PendingPermissionRequest( + future=future, + session_id=session_id, + ) try: - # Wait for user response with timeout approved = await asyncio.wait_for(future, timeout=timeout) logger.info( - f"[Permission] ✅ Got response for {request_id}: {'APPROVED' if approved else 'DENIED'}" + "[Permission] Got response for %s: %s", + request_id, + "APPROVED" if approved else "DENIED", ) - return approved + return bool(approved) except asyncio.TimeoutError: logger.warning( "[Permission] Request %s timed out after %ss", request_id, timeout @@ -130,43 +104,41 @@ async def request_permission( logger.error("[Permission] Error waiting for permission: %s", e) return False finally: - # Clean up the future permission_futures.pop(request_id, None) -def handle_permission_response(request_id: str, approved: bool) -> bool: - """ - Handle a permission response from the client. - - Args: - request_id: The request ID from the permission_request - approved: True if approved, False if denied +def handle_permission_response( + request_id: str, + approved: bool, + *, + session_id: Optional[str] = None, +) -> bool: + """Handle a permission response from the client. - Returns: - bool: True if the response was handled, False if request_id not found + If ``session_id`` is provided, response handling is restricted to matching + pending requests from the same session. """ - logger.info( - f"[Permission] handle_permission_response called: request_id={request_id}, approved={approved}" - ) - logger.info( - f"[Permission] Current permission_futures keys: {list(permission_futures.keys())}" - ) - - future = permission_futures.get(request_id) - logger.info( - f"[Permission] Future found: {future is not None}, done: {future.done() if future else 'N/A'}" - ) - - if future and not future.done(): - logger.info("[Permission] Setting future result to: %s", approved) - future.set_result(approved) - logger.info("[Permission] Handled response for %s: %s", request_id, approved) - return True - else: + pending = permission_futures.get(request_id) + if pending is None: logger.warning( - f"[Permission] Received response for unknown/expired request: {request_id}" + "[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( - f"[Permission] Future exists: {future is not None}, Future done: {future.done() if future else 'N/A'}" + "[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/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 index 0ab87812a..ca0cd200c 100644 --- a/code_puppy/api/templates/chat.html +++ b/code_puppy/api/templates/chat.html @@ -5,13 +5,61 @@ 🐶 Code Puppy Chat + + + + + +
- -
-
- 🐶 -
-

Code Puppy Chat

-

Connecting...

-
+ +
+ + +
+ + + + Connecting… + + +
+ Agent +
- -
- -
- - -
- -
- - -
- - -
-
Working Dir:
-
- - -
-
Not set
-
- - -
-
- Disconnected -
- - -
- - - Home - -
+ + +
+ Model + +
+ + +
+
+ Disconnected
+ + + + Home +
+ + +
+ + +
+ + Not set +
+ + +
+ + + + + + + + + + +
+ + + +
+
-
+
@@ -207,7 +794,7 @@

Code Puppy Chat