From d8d6f701f0a3a1ef0f1bab69fa0f35705d3a785c Mon Sep 17 00:00:00 2001 From: sskmlm Date: Mon, 13 Jul 2026 18:38:37 -0500 Subject: [PATCH 1/4] split: db/query layer --- code_puppy/api/compaction_tracker.py | 130 ++ code_puppy/api/db/__init__.py | 47 + 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 | 1154 +++++++++++++++++ code_puppy/api/db/seeder.py | 17 + .../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/error_parser.py | 239 ++++ code_puppy/api/routers/ws_sessions.py | 290 +++++ .../api/tests/test_backfill_pydantic_json.py | 349 +++++ code_puppy/api/tests/test_db_sql_loader.py | 22 + code_puppy/api/tool_formatters.py | 138 ++ code_puppy/config.py | 21 +- tests/api/test_ws_sessions_router.py | 88 ++ 17 files changed, 3655 insertions(+), 6 deletions(-) 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/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/error_parser.py create mode 100644 code_puppy/api/routers/ws_sessions.py create mode 100644 code_puppy/api/tests/test_backfill_pydantic_json.py create mode 100644 code_puppy/api/tests/test_db_sql_loader.py create mode 100644 code_puppy/api/tool_formatters.py create mode 100644 tests/api/test_ws_sessions_router.py diff --git a/code_puppy/api/compaction_tracker.py b/code_puppy/api/compaction_tracker.py new file mode 100644 index 000000000..8145b39c9 --- /dev/null +++ b/code_puppy/api/compaction_tracker.py @@ -0,0 +1,130 @@ +"""Compaction tracking — writes compaction metadata to SQLite.""" + +import logging +from datetime import datetime, timezone +from typing import Any, List + +logger = logging.getLogger(__name__) + + +def track_compaction_event( + session_id: str, + messages_before: List[Any], + messages_after: List[Any], + strategy: str, + agent: Any | None = None, +) -> None: + """ + Track a compaction event in SQLite. + + Called directly from message_history_processor after compaction completes. + + Args: + session_id: Session identifier + messages_before: Message history before compaction + messages_after: Message history after compaction + strategy: 'summarization' or 'truncation' + agent: Agent instance (optional, used for token estimation) + """ + # Only proceed if compaction actually happened (message count decreased) + if len(messages_after) >= len(messages_before): + logger.debug( + "Compaction tracker: no size reduction (%d -> %d), skipping", + len(messages_before), + len(messages_after), + ) + return + + try: + from code_puppy.api.db.connection import get_db, get_write_lock + + # Estimate tokens (use agent's method if available, else default) + def fallback_estimate(m: Any) -> int: + return 100 + + if agent and hasattr(agent, "estimate_tokens_for_message"): + estimate_tokens = agent.estimate_tokens_for_message + else: + estimate_tokens = fallback_estimate + + # Calculate compaction metadata + compacted_count = len(messages_before) - len(messages_after) + + # Estimate tokens for source messages (messages that were compacted) + # Assume first message is system prompt, compaction starts at index 1 + source_messages = messages_before[1 : compacted_count + 1] + source_tokens = sum(estimate_tokens(m) for m in source_messages) + + # Extract summary text from new messages (compare object IDs) + before_ids = {id(m) for m in messages_before} + new_messages = [m for m in messages_after if id(m) not in before_ids] + + summary_text = "(no summary found)" + summary_tokens = 0 + if new_messages: + summary_msg = new_messages[0] + text_parts = [] + for part in getattr(summary_msg, "parts", []) or []: + if hasattr(part, "content"): + text_parts.append(str(part.content)) + summary_text = " ".join(text_parts) if text_parts else "(empty summary)" + summary_tokens = estimate_tokens(summary_msg) + + # Assume compaction removes from the start (after system message at seq 1) + source_start = 2 # First message after system prompt + source_end = source_start + compacted_count - 1 + + # Create compaction_log row + created_at = datetime.now(timezone.utc).isoformat() + + with get_write_lock(): + db = get_db() + cursor = db.execute( + """ + INSERT INTO compaction_log + (session_id, summary_text, source_start, source_end, + source_count, source_tokens, summary_tokens, strategy, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + session_id, + summary_text[:1000], # Truncate very long summaries + source_start, + source_end, + compacted_count, + source_tokens, + summary_tokens, + strategy, + created_at, + ), + ) + compaction_log_id = cursor.lastrowid + + # Mark compacted messages in the messages table + rows_updated = db.execute( + """ + UPDATE messages + SET compacted = 1, compaction_log_id = ? + WHERE session_id = ? AND seq >= ? AND seq <= ? AND compacted = 0 + """, + (compaction_log_id, session_id, source_start, source_end), + ).rowcount + + db.commit() + + logger.info( + "Compaction tracked: session=%s log_id=%d seq=%d-%d count=%d marked=%d strategy=%s", + session_id, + compaction_log_id, + source_start, + source_end, + compacted_count, + rows_updated, + strategy, + ) + + except Exception as e: + # Non-fatal — compaction still worked, we just didn't track it + logger.warning( + "Compaction tracker failed for %s: %s", session_id, e, exc_info=True + ) diff --git a/code_puppy/api/db/__init__.py b/code_puppy/api/db/__init__.py new file mode 100644 index 000000000..83cffcb97 --- /dev/null +++ b/code_puppy/api/db/__init__.py @@ -0,0 +1,47 @@ +"""SQLite persistence layer for the Code Puppy WS API. + +Stores session history in ~/.puppy_desk/chat_messages.db — the same +database that the Desk Puppy Electron FE reads from. + +All database operations are async coroutines backed by aiosqlite. +""" + +from code_puppy.api.db.connection import close_db, get_db, init_db +from code_puppy.api.db.queries import ( + get_active_messages, + get_next_seq, + get_session_metadata, + get_session_row, + insert_compaction_log, + insert_message, + insert_tool_call, + mark_messages_compacted, + session_exists, + soft_delete_session, + update_session_stats, + update_session_working_directory, + upsert_session, + write_system_message_to_sqlite, + write_turn_to_sqlite, +) + +__all__ = [ + "init_db", + "close_db", + "get_db", + "get_session_row", + "get_session_metadata", + "session_exists", + "upsert_session", + "soft_delete_session", + "update_session_stats", + "update_session_working_directory", + "write_system_message_to_sqlite", + "get_next_seq", + "insert_message", + "insert_tool_call", + "get_active_messages", + "mark_messages_compacted", + "insert_compaction_log", + "write_turn_to_sqlite", +] diff --git a/code_puppy/api/db/backfill.py b/code_puppy/api/db/backfill.py new file mode 100644 index 000000000..142c0f471 --- /dev/null +++ b/code_puppy/api/db/backfill.py @@ -0,0 +1,381 @@ +"""Backfill helpers for the messages table. + +Currently contains one migration helper: + + backfill_pydantic_json(db) + Reconstructs pydantic_json from the plain `content` column for any + message row where pydantic_json IS NULL. Used by the v4→v5 schema + migration in connection.py. + +Design notes +------------ +* Only 'user' and 'assistant' roles are processed — 'system' rows have no + pydantic_json and the load path already skips them. +* The reconstruction is lossy for assistant messages that originally contained + tool calls or extended thinking, because that structure was never stored in + the `content` column. The backfill preserves the visible text so sessions + resume with *something* rather than blank history. +* The function is idempotent: it filters on ``pydantic_json IS NULL`` so + running it twice is safe. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import aiosqlite + +logger = logging.getLogger(__name__) + + +def _build_message(role: str, content: str, model_name: str) -> Any: + """Build a pydantic-ai ModelMessage from a plain content string. + + Args: + role: 'user' or 'assistant'. + content: Plain text content extracted from the messages row. + model_name: Model name stored on the row (used for ModelResponse). + + Returns: + A ModelRequest or ModelResponse instance ready for serialisation. + + Raises: + ValueError: If *role* is not 'user' or 'assistant'. + """ + from pydantic_ai.messages import ( + ModelRequest, + ModelResponse, + TextPart, + UserPromptPart, + ) + + if role == "user": + return ModelRequest(parts=[UserPromptPart(content=content)]) + if role == "assistant": + return ModelResponse( + parts=[TextPart(content=content)], + model_name=model_name or "unknown", + ) + raise ValueError(f"Unsupported role for backfill: {role!r}") + + +def _serialize(msg: Any) -> str: + """Serialise a single ModelMessage to the pydantic_json wire format. + + Stores as a single-element JSON array — identical to what + pydantic_json_for_message() produces during normal chat saves. + """ + from pydantic_ai.messages import ModelMessagesTypeAdapter + + return ModelMessagesTypeAdapter.dump_json([msg]).decode("utf-8") + + +async def backfill_pydantic_json( + db: "aiosqlite.Connection", + *, + batch_size: int = 200, +) -> tuple[int, int]: + """Backfill pydantic_json for message rows where it is currently NULL. + + Reads all eligible rows in one query, groups UPDATE statements into + batches of *batch_size* to keep transactions small, and commits after + each batch so a mid-run crash leaves the DB consistent. + + Args: + db: Open aiosqlite connection (must have row_factory set). + batch_size: Number of rows updated per transaction. + + Returns: + ``(updated, skipped)`` — row counts for rows successfully backfilled + and rows skipped due to serialisation errors. + """ + cursor = await db.execute( + """ + SELECT id, session_id, seq, role, content, model_name + FROM messages + WHERE pydantic_json IS NULL + AND compacted = 0 + AND role IN ('user', 'assistant') + ORDER BY session_id, seq + """ + ) + rows = await cursor.fetchall() + + if not rows: + logger.info("backfill_pydantic_json: no NULL rows found — nothing to do") + return 0, 0 + + logger.info("backfill_pydantic_json: %d rows to backfill", len(rows)) + + updated = 0 + skipped = 0 + batch: list[tuple[str, int]] = [] # (pydantic_json, id) + + for row in rows: + row_id: int = row["id"] + role: str = row["role"] + content: str = row["content"] or "" + model_name: str = row["model_name"] or "unknown" + session_id: str = row["session_id"] + seq: int = row["seq"] + + try: + msg = _build_message(role, content, model_name) + pj = _serialize(msg) + batch.append((pj, row_id)) + except Exception as exc: + logger.warning( + "backfill_pydantic_json: skipping id=%s (session=%s seq=%s role=%s): %s", + row_id, + session_id, + seq, + role, + exc, + ) + skipped += 1 + + if len(batch) >= batch_size: + await _flush_batch(db, batch) + updated += len(batch) + batch = [] + + if batch: + await _flush_batch(db, batch) + updated += len(batch) + + logger.info( + "backfill_pydantic_json: done — %d updated, %d skipped", + updated, + skipped, + ) + return updated, skipped + + +async def _flush_batch( + db: "aiosqlite.Connection", + batch: list[tuple[str, int]], +) -> None: + """Write one batch of (pydantic_json, id) pairs and commit.""" + await db.executemany( + "UPDATE messages SET pydantic_json = ? WHERE id = ?", + batch, + ) + await db.commit() + + +# --------------------------------------------------------------------------- +# Tool-call backfill (added to fix the early-continue bug in +# write_turn_to_sqlite that prevented tool_calls rows from ever being written) +# --------------------------------------------------------------------------- + + +async def backfill_tool_calls( + db: "aiosqlite.Connection", + *, + session_id: str | None = None, + batch_size: int = 100, +) -> tuple[int, int]: + """Backfill tool_call rows from pydantic_json for sessions saved without them. + + Sessions written before the fix for the early-continue bug in + ``write_turn_to_sqlite`` have no rows in ``tool_calls`` even though the + pydantic_json column contains full ToolCallPart / ToolReturnPart data. + + This function replays each session's message history in seq order, + matches ToolCallPart entries in assistant messages with ToolReturnPart + entries in the following user messages, and inserts the reconstructed + rows into ``tool_calls``. Uses ``INSERT OR IGNORE`` — fully idempotent. + + Args: + db: Open aiosqlite connection with ``row_factory`` set to + ``aiosqlite.Row``. + session_id: If provided, only process that one session; otherwise + process every session that has pydantic_json rows. + batch_size: Number of ``tool_calls`` rows inserted per commit. + + Returns: + ``(inserted, skipped)`` row counts. + """ + import json + import time + import uuid + + from pydantic_ai.messages import ModelMessagesTypeAdapter + + where_clause = "WHERE m.pydantic_json IS NOT NULL AND m.compacted = 0" + params: tuple = () + if session_id: + where_clause += " AND m.session_id = ?" + params = (session_id,) + + cursor = await db.execute( + f""" + SELECT m.id, m.session_id, m.seq, m.role, + m.agent_name, m.model_name, m.timestamp, m.pydantic_json + FROM messages m + {where_clause} + ORDER BY m.session_id, m.seq + """, + params, + ) + rows = await cursor.fetchall() + + if not rows: + logger.info("backfill_tool_calls: no eligible messages found — nothing to do") + return 0, 0 + + logger.info( + "backfill_tool_calls: processing %d message rows%s", + len(rows), + f" for session={session_id}" if session_id else "", + ) + + inserted = 0 + skipped = 0 + batch: list[tuple] = [] + + # Per-session state: pending tool calls keyed by tool_call_id + current_session: str | None = None + # tool_call_id → {name, args, parent_seq, agent, model, ts} + pending: dict[str, dict] = {} + + for row in rows: + row_session = row["session_id"] + if row_session != current_session: + # Reset pending state when we enter a new session + current_session = row_session + pending = {} + + try: + parsed = ModelMessagesTypeAdapter.validate_json(row["pydantic_json"]) + except Exception as exc: + logger.warning( + "backfill_tool_calls: skip seq=%s session=%s parse error: %s", + row["seq"], + row["session_id"], + exc, + ) + skipped += 1 + continue + + if not parsed: + continue + + msg = parsed[0] + role: str = row["role"] + seq: int = row["seq"] + agent_name: str = row["agent_name"] or "" + model_name: str = row["model_name"] or "" + + # Collect ToolCallParts from assistant (ModelResponse) messages + if role == "assistant": + for part in getattr(msg, "parts", []): + if type(part).__name__ != "ToolCallPart": + continue + tc_id = getattr(part, "tool_call_id", "") or str(uuid.uuid4()) + args = getattr(part, "args", {}) + if hasattr(args, "model_dump"): + args = args.model_dump() + elif hasattr(args, "__dict__"): + args = dict(args.__dict__) + pending[tc_id] = { + "name": getattr(part, "tool_name", "unknown"), + "args": args, + "parent_seq": seq, + "agent": agent_name, + "model": model_name, + "ts": row["timestamp"], + } + + # Match ToolReturnParts from user (ModelRequest) messages + elif role == "user": + for part in getattr(msg, "parts", []): + if type(part).__name__ != "ToolReturnPart": + continue + tc_id = getattr(part, "tool_call_id", "") + if not tc_id or tc_id not in pending: + continue + + tc = pending.pop(tc_id) + content = getattr(part, "content", "") + try: + result_val = ( + json.loads(content) if isinstance(content, str) else content + ) + except Exception: + result_val = content + + try: + args_json_str = json.dumps(tc["args"]) + except Exception: + args_json_str = str(tc["args"]) + try: + result_json_str = json.dumps(result_val) + except Exception: + result_json_str = str(result_val) + + try: + from datetime import datetime as _dt + + tool_ts = _dt.fromisoformat( + str(tc["ts"]).replace("Z", "+00:00") + ).timestamp() + except Exception: + tool_ts = time.time() + + batch.append( + ( + tc_id, # id + row_session, # session_id + tc["parent_seq"], # parent_message_seq + tc[ + "parent_seq" + ], # seq (best approximation without original seq) + tc["name"], # tool_name + args_json_str, # args_json + result_json_str, # result_json + "success", # status + None, # duration_ms + None, # error_text + tc["agent"], # agent_name + tc["model"], # model_name + tool_ts, # timestamp + ) + ) + + if len(batch) >= batch_size: + n = await _flush_tool_call_batch(db, batch) + inserted += n + batch = [] + + if batch: + n = await _flush_tool_call_batch(db, batch) + inserted += n + + logger.info( + "backfill_tool_calls: done — %d inserted, %d skipped", + inserted, + skipped, + ) + return inserted, skipped + + +async def _flush_tool_call_batch( + db: "aiosqlite.Connection", + batch: list[tuple], +) -> int: + """INSERT OR IGNORE a batch of tool_call rows and commit.""" + await db.executemany( + """ + INSERT OR IGNORE INTO tool_calls + (id, session_id, parent_message_seq, seq, tool_name, + args_json, result_json, status, duration_ms, error_text, + agent_name, model_name, timestamp) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + batch, + ) + await db.commit() + return len(batch) diff --git a/code_puppy/api/db/connection.py b/code_puppy/api/db/connection.py new file mode 100644 index 000000000..d60f9e2b2 --- /dev/null +++ b/code_puppy/api/db/connection.py @@ -0,0 +1,426 @@ +"""SQLite connection singleton using aiosqlite for non-blocking async access. + +The Desk Puppy Electron FE owns the schema (creates it at user_version=2). +This module opens the same database from the Python BE, runs any needed +migrations if the FE hasn't run yet, and exposes get_db() / init_db() / +close_db() for use by queries.py and the seeder. + +Schema versions: + 0 = no schema yet (fresh DB) + 1 = added token_count, compacted, pydantic_json, compaction_log_id to messages + 2 = added deleted_at to sessions (soft delete) + 3 = rebuilt messages table with UNIQUE(session_id, seq) + dedup + 4 = backfill blank agent_name / model_name in messages + tool_calls from session defaults + +Concurrency: + aiosqlite serialises all DB operations through its own internal background + thread — no threading.Lock required. All callers simply await the relevant + coroutine; the event loop is never blocked. +""" + +from __future__ import annotations + +import logging +import os +from pathlib import Path +from typing import Optional + +import aiosqlite + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Path resolution +# --------------------------------------------------------------------------- + +_PUPPY_DESK_DB_ENV = "PUPPY_DESK_DB" + + +def get_db_path() -> Path: + """Return path to the shared SQLite database. + + Override with PUPPY_DESK_DB env var for testing. + """ + env = os.environ.get(_PUPPY_DESK_DB_ENV) + if env: + return Path(env) + return Path.home() / ".puppy_desk" / "chat_messages.db" + + +# --------------------------------------------------------------------------- +# Full schema (v4) — created on first BE run if FE hasn't seeded it yet +# --------------------------------------------------------------------------- + +_SCHEMA_SQL = """ +CREATE TABLE IF NOT EXISTS sessions ( + session_id TEXT PRIMARY KEY, + title TEXT DEFAULT '', + agent_name TEXT DEFAULT 'code-puppy', + model_name TEXT DEFAULT '', + working_directory TEXT DEFAULT '', + pinned INTEGER DEFAULT 0, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + message_count INTEGER DEFAULT 0, + total_tokens INTEGER DEFAULT 0, + deleted_at TEXT +); +CREATE INDEX IF NOT EXISTS idx_sessions_updated ON sessions(updated_at DESC); +CREATE INDEX IF NOT EXISTS idx_sessions_deleted ON sessions(deleted_at); + +CREATE TABLE IF NOT EXISTS compaction_log ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + summary_text TEXT NOT NULL, + source_start INTEGER NOT NULL, + source_end INTEGER NOT NULL, + source_count INTEGER NOT NULL, + source_tokens INTEGER NOT NULL, + summary_tokens INTEGER NOT NULL, + strategy TEXT DEFAULT 'summarization', + created_at TEXT NOT NULL, + FOREIGN KEY (session_id) REFERENCES sessions(session_id) ON DELETE CASCADE +); +CREATE INDEX IF NOT EXISTS idx_compaction_session ON compaction_log(session_id); + +CREATE TABLE IF NOT EXISTS messages ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + seq INTEGER NOT NULL, + role TEXT NOT NULL, + content TEXT DEFAULT '', + type TEXT DEFAULT '', + agent_name TEXT DEFAULT '', + model_name TEXT DEFAULT '', + timestamp TEXT NOT NULL, + thinking TEXT, + attachments_json TEXT, + clean_content TEXT, + system_message_type TEXT, + system_message_path TEXT, + token_count INTEGER DEFAULT 0, + compacted INTEGER DEFAULT 0, + pydantic_json TEXT, + compaction_log_id INTEGER REFERENCES compaction_log(id) ON DELETE SET NULL, + FOREIGN KEY (session_id) REFERENCES sessions(session_id) ON DELETE CASCADE, + UNIQUE(session_id, seq) +); +CREATE INDEX IF NOT EXISTS idx_messages_session_seq ON messages(session_id, seq); +CREATE INDEX IF NOT EXISTS idx_messages_active ON messages(session_id, compacted); + +CREATE TABLE IF NOT EXISTS tool_calls ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + parent_message_seq INTEGER, + seq INTEGER NOT NULL, + tool_name TEXT NOT NULL, + args_json TEXT, + result_json TEXT, + status TEXT DEFAULT 'running', + duration_ms INTEGER, + error_text TEXT, + agent_name TEXT DEFAULT '', + model_name TEXT DEFAULT '', + timestamp REAL NOT NULL, + FOREIGN KEY (session_id) REFERENCES sessions(session_id) ON DELETE CASCADE +); +CREATE INDEX IF NOT EXISTS idx_tool_calls_session_seq ON tool_calls(session_id, seq); +CREATE INDEX IF NOT EXISTS idx_tool_calls_parent ON tool_calls(parent_message_seq); +""" + +# Migrations applied to existing databases that are below SCHEMA_VERSION +_MIGRATION_V0_TO_V1: list[str] = [ + "ALTER TABLE messages ADD COLUMN token_count INTEGER DEFAULT 0", + "ALTER TABLE messages ADD COLUMN compacted INTEGER DEFAULT 0", + "ALTER TABLE messages ADD COLUMN pydantic_json TEXT", + "ALTER TABLE messages ADD COLUMN compaction_log_id INTEGER REFERENCES compaction_log(id) ON DELETE SET NULL", +] + +_MIGRATION_V1_TO_V2: list[str] = [ + "ALTER TABLE sessions ADD COLUMN deleted_at TEXT", +] + +# v3: Rebuild messages table to add UNIQUE(session_id, seq). +# SQLite cannot add UNIQUE constraints via ALTER TABLE, so we must rebuild. +_MIGRATION_V2_TO_V3_SQL = """ +PRAGMA foreign_keys = OFF; +BEGIN; + +CREATE TABLE messages_v3 ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + seq INTEGER NOT NULL, + role TEXT NOT NULL, + content TEXT DEFAULT '', + type TEXT DEFAULT '', + agent_name TEXT DEFAULT '', + model_name TEXT DEFAULT '', + timestamp TEXT NOT NULL, + thinking TEXT, + attachments_json TEXT, + clean_content TEXT, + system_message_type TEXT, + system_message_path TEXT, + token_count INTEGER DEFAULT 0, + compacted INTEGER DEFAULT 0, + pydantic_json TEXT, + compaction_log_id INTEGER REFERENCES compaction_log(id) ON DELETE SET NULL, + FOREIGN KEY (session_id) REFERENCES sessions(session_id) ON DELETE CASCADE, + UNIQUE(session_id, seq) +); + +INSERT INTO messages_v3 +SELECT id, session_id, seq, role, content, type, agent_name, model_name, + timestamp, thinking, attachments_json, clean_content, + system_message_type, system_message_path, + token_count, compacted, pydantic_json, compaction_log_id +FROM messages AS m1 +WHERE id = ( + SELECT id FROM messages AS m2 + WHERE m2.session_id = m1.session_id AND m2.seq = m1.seq + ORDER BY + (m2.clean_content IS NOT NULL) DESC, + (m2.pydantic_json IS NOT NULL) DESC, + m2.id DESC + LIMIT 1 +); + +DROP TABLE messages; +ALTER TABLE messages_v3 RENAME TO messages; + +CREATE INDEX IF NOT EXISTS idx_messages_session_seq ON messages(session_id, seq); +CREATE INDEX IF NOT EXISTS idx_messages_active ON messages(session_id, compacted); + +COMMIT; +PRAGMA foreign_keys = ON; +""".strip() + +SCHEMA_VERSION = 5 + +# --------------------------------------------------------------------------- +# Async singleton (aiosqlite) +# --------------------------------------------------------------------------- + +_aconn: Optional[aiosqlite.Connection] = None + + +async def init_db() -> None: + """Open the database and ensure schema is at SCHEMA_VERSION. + + Safe to call multiple times — subsequent calls are no-ops. + """ + global _aconn + + if _aconn is not None: + return # already initialised + + db_path = get_db_path() + db_path.parent.mkdir(parents=True, exist_ok=True) + + logger.info("Opening aiosqlite DB at %s", db_path) + _aconn = await aiosqlite.connect(str(db_path)) + _aconn.row_factory = aiosqlite.Row + + # Performance + safety pragmas + await _aconn.execute("PRAGMA journal_mode = WAL") + await _aconn.execute("PRAGMA foreign_keys = ON") + await _aconn.execute("PRAGMA busy_timeout = 5000") + + cursor = await _aconn.execute("PRAGMA user_version") + row = await cursor.fetchone() + current_version: int = row[0] + + if current_version == 0: + # Fresh database — create everything at the current schema version + await _aconn.executescript(_SCHEMA_SQL) + await _aconn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") + await _aconn.commit() + logger.info( + "\u2713 aiosqlite DB created (schema v%d) at %s", SCHEMA_VERSION, db_path + ) + else: + # Existing database — run incremental migrations + await _run_alters_and_migrate(current_version) + + +async def _run_alters_and_migrate(current_version: int) -> None: + """Apply all pending schema migrations in order.""" + assert _aconn is not None + + if current_version < 1: + await _run_alters(_MIGRATION_V0_TO_V1) + await _aconn.execute("PRAGMA user_version = 1") + await _aconn.commit() + logger.info("\u2713 aiosqlite DB migrated v0 \u2192 v1") + + if current_version < 2: + await _run_alters(_MIGRATION_V1_TO_V2) + # Also create compaction_log and its index if missing (v0 DBs) + await _aconn.execute(""" + CREATE TABLE IF NOT EXISTS compaction_log ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + summary_text TEXT NOT NULL, + source_start INTEGER NOT NULL, + source_end INTEGER NOT NULL, + source_count INTEGER NOT NULL, + source_tokens INTEGER NOT NULL, + summary_tokens INTEGER NOT NULL, + strategy TEXT DEFAULT 'summarization', + created_at TEXT NOT NULL, + FOREIGN KEY (session_id) REFERENCES sessions(session_id) ON DELETE CASCADE + ) + """) + await _aconn.execute( + "CREATE INDEX IF NOT EXISTS idx_compaction_session ON compaction_log(session_id)" + ) + await _aconn.execute( + "CREATE INDEX IF NOT EXISTS idx_sessions_deleted ON sessions(deleted_at)" + ) + await _aconn.execute( + "CREATE INDEX IF NOT EXISTS idx_messages_active ON messages(session_id, compacted)" + ) + await _aconn.execute("PRAGMA user_version = 2") + await _aconn.commit() + logger.info("\u2713 aiosqlite DB migrated to v2") + + if current_version < 3: + logger.info( + "Running v2\u21923 migration: rebuilding messages table with UNIQUE(session_id, seq)..." + ) + await _aconn.executescript(_MIGRATION_V2_TO_V3_SQL) + await _aconn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") + await _aconn.commit() + cursor = await _aconn.execute("SELECT COUNT(*) FROM messages") + row = await cursor.fetchone() + count = row[0] + logger.info( + "\u2713 aiosqlite DB migrated to v%d (%d messages after dedup)", + SCHEMA_VERSION, + count, + ) + + if current_version < 4: + logger.info( + "Running v3\u2192v4 migration: backfilling blank agent_name/model_name \u2026" + ) + await _aconn.executescript(""" + -- messages: blank agent_name → session.agent_name + UPDATE messages + SET agent_name = ( + SELECT s.agent_name + FROM sessions s + WHERE s.session_id = messages.session_id + ) + WHERE (agent_name IS NULL OR agent_name = '') + AND ( + SELECT COALESCE(s.agent_name, '') + FROM sessions s + WHERE s.session_id = messages.session_id + ) != ''; + + -- messages: blank model_name → session.model_name + UPDATE messages + SET model_name = ( + SELECT s.model_name + FROM sessions s + WHERE s.session_id = messages.session_id + ) + WHERE (model_name IS NULL OR model_name = '') + AND ( + SELECT COALESCE(s.model_name, '') + FROM sessions s + WHERE s.session_id = messages.session_id + ) != ''; + + -- tool_calls: blank agent_name → session.agent_name + UPDATE tool_calls + SET agent_name = ( + SELECT s.agent_name + FROM sessions s + WHERE s.session_id = tool_calls.session_id + ) + WHERE (agent_name IS NULL OR agent_name = '') + AND ( + SELECT COALESCE(s.agent_name, '') + FROM sessions s + WHERE s.session_id = tool_calls.session_id + ) != ''; + + -- tool_calls: blank model_name → session.model_name + UPDATE tool_calls + SET model_name = ( + SELECT s.model_name + FROM sessions s + WHERE s.session_id = tool_calls.session_id + ) + WHERE (model_name IS NULL OR model_name = '') + AND ( + SELECT COALESCE(s.model_name, '') + FROM sessions s + WHERE s.session_id = tool_calls.session_id + ) != ''; + """) + await _aconn.execute("PRAGMA user_version = 4") + await _aconn.commit() + cursor = await _aconn.execute( + "SELECT COUNT(*) FROM messages WHERE agent_name != '' AND model_name != ''" + ) + row = await cursor.fetchone() + msg_count = row[0] + logger.info( + "\u2713 aiosqlite DB migrated to v4 (%d messages with agent+model)", + msg_count, + ) + + if current_version < 5: + logger.info( + "Running v4\u2192v5 migration: backfilling pydantic_json for NULL rows \u2026" + ) + from code_puppy.api.db.backfill import backfill_pydantic_json + + updated, skipped = await backfill_pydantic_json(_aconn) + await _aconn.execute(f"PRAGMA user_version = {SCHEMA_VERSION}") + await _aconn.commit() + logger.info( + "\u2713 aiosqlite DB migrated to v5 " + "(%d pydantic_json rows backfilled, %d skipped)", + updated, + skipped, + ) + + logger.info("\u2713 aiosqlite DB ready (schema v%d)", SCHEMA_VERSION) + + +async def _run_alters(statements: list[str]) -> None: + """Run ALTER TABLE statements, skipping 'duplicate column' errors gracefully.""" + assert _aconn is not None + for sql in statements: + try: + await _aconn.execute(sql) + except Exception as exc: + if "duplicate column" in str(exc).lower(): + logger.debug("Skipping already-applied migration: %s", sql) + else: + raise + + +def get_db() -> aiosqlite.Connection: + """Return the open aiosqlite connection. + + Raises RuntimeError if init_db() has not been awaited. + """ + if _aconn is None: + raise RuntimeError( + "aiosqlite DB not initialised — call `await init_db()` first" + ) + return _aconn + + +async def close_db() -> None: + """Close the database connection. Called during app shutdown.""" + global _aconn + if _aconn is not None: + await _aconn.close() + _aconn = None + logger.info("aiosqlite DB closed") diff --git a/code_puppy/api/db/message_utils.py b/code_puppy/api/db/message_utils.py new file mode 100644 index 000000000..8c38a1851 --- /dev/null +++ b/code_puppy/api/db/message_utils.py @@ -0,0 +1,219 @@ +"""Shared helpers for extracting display fields from pydantic-ai ModelMessages. + +Used by the SQLite persistence helpers and session_context.py. +Must not import from seeder.py or session_context.py to avoid circular imports. +""" + +from __future__ import annotations + +import logging +import warnings +from typing import Any, Optional + +logger = logging.getLogger(__name__) + +_PART_TYPE_TEXT = {"TextPart", "UserPromptPart", "SystemPromptPart"} +_PART_TYPE_THINKING = {"ThinkingPart"} + + +def get_role(msg: Any) -> str: + """Map a ModelMessage to a DB role string.""" + kind = getattr(msg, "kind", None) + if kind == "request": + return "user" + if kind == "response": + return "assistant" + return "system" + + +def extract_content(msg: Any) -> str: + """Extract human-readable text from a ModelMessage's parts.""" + parts = getattr(msg, "parts", []) + texts = [] + for part in parts: + if type(part).__name__ in _PART_TYPE_TEXT: + c = getattr(part, "content", "") + if isinstance(c, str) and c.strip(): + texts.append(c) + return "\n".join(texts) + + +def extract_thinking(msg: Any) -> Optional[str]: + """Extract extended thinking text from a ModelResponse's parts.""" + parts = getattr(msg, "parts", []) + texts = [ + getattr(p, "content", "") + for p in parts + if type(p).__name__ in _PART_TYPE_THINKING + and isinstance(getattr(p, "content", ""), str) + ] + joined = "\n".join(t for t in texts if t) + return joined or None + + +def get_message_timestamp(msg: Any) -> Optional[str]: + """Return the message's own timestamp as an ISO string, or None.""" + ts = getattr(msg, "timestamp", None) + if ts is None: + return None + try: + return ts.isoformat() + except Exception: + return str(ts) + + +def pydantic_json_for_message(msg: Any) -> Optional[str]: + """Serialise a single ModelMessage to a JSON string via ModelMessagesTypeAdapter. + + Stores as a single-element list: `[msg_json]`. + Deserialise with: `ModelMessagesTypeAdapter.validate_json(s)[0]` + + Returns None if serialisation fails for unsupported historical message shapes. + """ + try: + from pydantic_ai.messages import ModelMessagesTypeAdapter + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + return ModelMessagesTypeAdapter.dump_json([msg]).decode("utf-8") + except Exception as exc: + logger.debug("pydantic_json serialisation failed: %s", exc) + return None + + +_PART_TYPE_TOOL_CALL = {"ToolCallPart"} +# Extended set of tool return part type names for runtime compatibility +# Different versions of pydantic-ai and streaming may use different class names +_PART_TYPE_TOOL_RETURN = {"ToolReturnPart", "ToolReturn", "RetryPromptPart"} + + +def extract_tool_calls(msg: Any) -> list[dict[str, Any]]: + """Extract ToolCallPart entries from a ModelResponse. + + Returns list of dicts: {id, name, args_dict} + """ + import uuid + + parts = getattr(msg, "parts", []) + result = [] + for part in parts: + if type(part).__name__ not in _PART_TYPE_TOOL_CALL: + continue + args = getattr(part, "args", {}) + if hasattr(args, "model_dump"): + args = args.model_dump() + elif hasattr(args, "__dict__"): + args = dict(args.__dict__) + result.append( + { + "id": getattr(part, "tool_call_id", "") or str(uuid.uuid4()), + "name": getattr(part, "tool_name", "unknown"), + "args": args, + } + ) + return result + + +def _is_tool_return_like(part: Any) -> bool: + """Check if a part is a tool return using duck-typing. + + Handles various runtime representations: + - ToolReturnPart (standard pydantic-ai) + - ToolReturn (streaming/runtime variant) + - Any object with tool_call_id + (content or return_value) + """ + class_name = type(part).__name__ + + # Fast path: known class names + if class_name in _PART_TYPE_TOOL_RETURN: + return True + + # Duck-typing fallback: must have tool_call_id and some result content + has_tool_call_id = hasattr(part, "tool_call_id") and getattr( + part, "tool_call_id", None + ) + has_content = hasattr(part, "content") or hasattr(part, "return_value") + + return bool(has_tool_call_id and has_content) + + +def _extract_tool_return_data(part: Any) -> dict[str, Any]: + """Extract tool return data from a part using multiple field strategies. + + Handles different field naming conventions: + - content (ToolReturnPart) + - return_value (some runtime variants) + - result (potential future variants) + """ + import json + + # Get the result content - try multiple field names + raw_content = None + for field in ("content", "return_value", "result"): + if hasattr(part, field): + raw_content = getattr(part, field, None) + if raw_content is not None: + break + + # Parse JSON if string, otherwise use as-is + if raw_content is None: + result_val = None + elif isinstance(raw_content, str): + try: + result_val = json.loads(raw_content) + except (json.JSONDecodeError, TypeError): + result_val = raw_content + else: + result_val = raw_content + + # Get tool name - try multiple field names + tool_name = "unknown" + for field in ("tool_name", "name"): + if hasattr(part, field): + name = getattr(part, field, None) + if name: + tool_name = name + break + + return { + "id": getattr(part, "tool_call_id", "") or "", + "name": tool_name, + "result": result_val, + } + + +def extract_tool_returns(msg: Any) -> list[dict[str, Any]]: + """Extract tool return entries from a ModelRequest. + + Returns list of dicts: {id, name, result} + + Handles multiple runtime representations: + - ToolReturnPart (standard pydantic-ai) + - ToolReturn (streaming variant) + - Duck-typed objects with tool_call_id + content/return_value + + This extended compatibility ensures tool results are persisted to SQLite + regardless of which pydantic-ai version or streaming mode is in use. + """ + parts = getattr(msg, "parts", []) + result = [] + + for part in parts: + if not _is_tool_return_like(part): + continue + + try: + data = _extract_tool_return_data(part) + if data["id"]: # Only add if we have a valid tool_call_id + result.append(data) + else: + logger.debug( + "Skipping tool return with empty tool_call_id: %s", + type(part).__name__, + ) + except Exception as exc: + logger.warning( + "Failed to extract tool return from %s: %s", type(part).__name__, exc + ) + + return result diff --git a/code_puppy/api/db/queries.py b/code_puppy/api/db/queries.py new file mode 100644 index 000000000..1a26ad0b1 --- /dev/null +++ b/code_puppy/api/db/queries.py @@ -0,0 +1,1154 @@ +"""Typed async query helpers for the shared chat_messages.db. + +All functions are coroutines. aiosqlite serialises writes internally via its +own background thread — no threading.Lock is needed here. +""" + +from __future__ import annotations + +import json +import logging +from typing import Any, Optional + +from code_puppy.api.db.connection import get_db +from code_puppy.api.db.sql_loader import load_sql + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Sessions +# --------------------------------------------------------------------------- + + +async def session_exists(session_id: str) -> bool: + """Return True if a session row already exists (regardless of deleted_at).""" + db = get_db() + cursor = await db.execute( + "SELECT 1 FROM sessions WHERE session_id = ? LIMIT 1", (session_id,) + ) + row = await cursor.fetchone() + return row is not None + + +async def get_session_row(session_id: str) -> Optional[dict]: + """Return the full session row as a plain dict, or None if not found. + + Non-fatal: returns None on any exception (DB unavailable, etc.). + """ + try: + db = get_db() + cursor = await db.execute( + "SELECT * FROM sessions WHERE session_id = ?", (session_id,) + ) + row = await cursor.fetchone() + if row is None: + return None + return dict(row) + except Exception: + return None + + +async def get_session_metadata(session_id: str) -> Optional[dict]: + """Return a subset of session columns useful for the chat handler. + + Returns keys: session_id, title, agent_name, model_name, + working_directory, pinned, created_at — or None if not found. + Non-fatal: returns None on any exception. + """ + try: + db = get_db() + cursor = await db.execute( + """ + SELECT session_id, title, agent_name, model_name, + working_directory, pinned, created_at + FROM sessions + WHERE session_id = ? + """, + (session_id,), + ) + row = await cursor.fetchone() + if row is None: + return None + return dict(row) + except Exception: + return None + + +async def upsert_session( + *, + session_id: str, + title: str = "", + agent_name: str = "code-puppy", + model_name: str = "", + working_directory: str = "", + pinned: bool = False, + created_at: str, + updated_at: str, + message_count: int = 0, + total_tokens: int = 0, + deleted_at: Optional[str] = None, +) -> None: + """Insert or replace a session row.""" + db = get_db() + try: + await db.execute( + """ + INSERT INTO sessions + (session_id, title, agent_name, model_name, working_directory, + pinned, created_at, updated_at, message_count, total_tokens, deleted_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(session_id) DO UPDATE SET + title = excluded.title, + agent_name = excluded.agent_name, + model_name = excluded.model_name, + working_directory = excluded.working_directory, + pinned = excluded.pinned, + updated_at = excluded.updated_at, + message_count = excluded.message_count, + total_tokens = excluded.total_tokens, + deleted_at = COALESCE(sessions.deleted_at, excluded.deleted_at) + """, + ( + session_id, + title, + agent_name, + model_name, + working_directory, + 1 if pinned else 0, + created_at, + updated_at, + message_count, + total_tokens, + deleted_at, + ), + ) + await db.commit() + except Exception: + await db.rollback() + raise + + +async def update_session_stats( + session_id: str, + *, + message_count: int, + total_tokens: int, + updated_at: str, +) -> None: + """Bump message_count, total_tokens, and updated_at after a save.""" + db = get_db() + try: + await db.execute( + """ + UPDATE sessions + SET message_count = ?, total_tokens = ?, updated_at = ? + WHERE session_id = ? + """, + (message_count, total_tokens, updated_at, session_id), + ) + await db.commit() + except Exception: + await db.rollback() + raise + + +async def update_session_working_directory( + session_id: str, + working_directory: str, + updated_at: str, +) -> None: + """Update only the working_directory and updated_at columns for a session. + + Narrow UPDATE — does not touch title, agent, model, pinned, counts, etc. + Called after a successful set_working_directory WS message. + """ + db = get_db() + await db.execute( + "UPDATE sessions SET working_directory = ?, updated_at = ? WHERE session_id = ?", + (working_directory, updated_at, session_id), + ) + await db.commit() + + +async def update_session_meta_fields( + session_id: str, + *, + title: str, + pinned: bool, + updated_at: str, +) -> None: + """Update only title, pinned, and updated_at for an existing session. + + Narrow UPDATE — does NOT touch message_count, total_tokens, agent_name, + model_name, or working_directory. Called by the update_session_meta WS + message handler (user renames or pins/unpins a session). + """ + db = get_db() + await db.execute( + "UPDATE sessions SET title = ?, pinned = ?, updated_at = ? WHERE session_id = ?", + (title, 1 if pinned else 0, updated_at, session_id), + ) + await db.commit() + + +async def soft_delete_session(session_id: str, deleted_at: str) -> None: + """Set deleted_at on a session (soft delete). Data is preserved.""" + db = get_db() + await db.execute( + "UPDATE sessions SET deleted_at = ? WHERE session_id = ?", + (deleted_at, session_id), + ) + await db.commit() + + +# --------------------------------------------------------------------------- +# Messages +# --------------------------------------------------------------------------- + + +async def get_next_seq(session_id: str) -> int: + """Return the next available seq number for a session (1-based).""" + db = get_db() + cursor = await db.execute( + "SELECT COALESCE(MAX(seq), 0) FROM messages WHERE session_id = ?", + (session_id,), + ) + row = await cursor.fetchone() + return (row[0] or 0) + 1 + + +async def _insert_immediate_message_and_sync_session( + *, + session_id: str, + role: str, + content: str, + type: str, + agent_name: str = "", + model_name: str = "", + timestamp: str, + thinking: Optional[str] = None, + attachments_json: Optional[str] = None, + clean_content: Optional[str] = None, + system_message_type: Optional[str] = None, + system_message_path: Optional[str] = None, + token_count: int = 0, + increment_message_count: bool = True, +) -> int: + """Insert an immediate WS message and sync session metadata atomically. + + This path is used for system/error frames that must persist immediately, + outside the normal end-of-turn batch save. The session row is created if + needed, the next seq is allocated atomically, the message row is inserted, + and session metadata is updated in the same transaction so + message_count/updated_at cannot drift. + + increment_message_count controls whether the sessions.message_count should + track this row. Persisted system banners/config rows should not affect the + user-visible chat turn count, while persisted error rows should. + """ + db = get_db() + try: + await db.execute("BEGIN IMMEDIATE") + cursor = await db.execute( + "SELECT 1 FROM sessions WHERE session_id = ?", (session_id,) + ) + existing = await cursor.fetchone() + if not existing: + await db.execute( + """ + INSERT INTO sessions + (session_id, title, agent_name, model_name, working_directory, + pinned, created_at, updated_at, message_count, total_tokens, deleted_at) + VALUES (?, '', ?, ?, '', 0, ?, ?, 0, 0, NULL) + """, + ( + session_id, + agent_name or "code-puppy", + model_name or "", + timestamp, + timestamp, + ), + ) + + cursor = await db.execute( + "SELECT COALESCE(MAX(seq), 0) + 1 AS next_seq FROM messages WHERE session_id = ?", + (session_id,), + ) + seq_row = await cursor.fetchone() + seq = int(seq_row["next_seq"] if seq_row is not None else 1) + + await db.execute( + """ + INSERT INTO messages + (session_id, seq, role, content, type, agent_name, model_name, + timestamp, thinking, attachments_json, clean_content, + system_message_type, system_message_path, + token_count, compacted, pydantic_json, compaction_log_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0, NULL, NULL) + """, + ( + session_id, + seq, + role, + content, + type, + agent_name, + model_name, + timestamp, + thinking, + attachments_json, + clean_content, + system_message_type, + system_message_path, + token_count, + ), + ) + await db.execute( + """ + UPDATE sessions + SET message_count = COALESCE(message_count, 0) + ?, + total_tokens = COALESCE(total_tokens, 0) + ?, + updated_at = ? + WHERE session_id = ? + """, + ( + 1 if increment_message_count else 0, + token_count, + timestamp, + session_id, + ), + ) + await db.commit() + return seq + except Exception: + await db.rollback() + raise + + +async def write_system_message_to_sqlite( + *, + session_id: str, + system_message_type: str, + content: str, + system_message_path: str = "", + agent_name: str = "", + model_name: str = "", + timestamp: Optional[str] = None, +) -> None: + """Write a single system-event row to SQLite immediately. + + Used by the WS chat handler to persist agent/model switches, CWD changes, + and new-session initialisation messages so the FE cold-load path sees them + without relying on ephemeral WS-only state. + + For directory banners, deduplicates to avoid multiple banners for the same path. + """ + import datetime as _dt + + ts = timestamp or _dt.datetime.now(_dt.timezone.utc).isoformat() + + # Dedup: for directory system messages, check if one already exists with same path + if system_message_type == "directory" and system_message_path: + try: + db = get_db() + cursor = await db.execute( + """ + SELECT 1 FROM messages + WHERE session_id = ? + AND system_message_type = 'directory' + AND system_message_path = ? + LIMIT 1 + """, + (session_id, system_message_path), + ) + existing = await cursor.fetchone() + if existing: + import logging + + logging.getLogger(__name__).debug( + "Skipping duplicate directory banner: session=%s path=%s", + session_id, + system_message_path, + ) + return # Already have a banner for this path + except Exception as e: + import logging + + logging.getLogger(__name__).warning("Directory dedup check failed: %s", e) + # Continue with insert on error + + try: + await _insert_immediate_message_and_sync_session( + session_id=session_id, + role="system", + content=content, + type="system", + agent_name=agent_name, + model_name=model_name, + timestamp=ts, + system_message_type=system_message_type, + system_message_path=system_message_path, + increment_message_count=False, + ) + except Exception as exc: + import logging + + logging.getLogger(__name__).warning( + "write_system_message_to_sqlite failed (session=%s type=%s): %s", + session_id, + system_message_type, + exc, + exc_info=True, + ) + + +async def write_error_message_to_sqlite( + *, + session_id: str, + error: str, + error_type: str = "unknown", + technical_details: str = "", + action_required: Optional[str] = None, + agent_name: str = "", + model_name: str = "", + timestamp: Optional[str] = None, +) -> None: + """Write a structured error row to SQLite immediately. + + Error rows are stored as regular message rows with type='error' so they + survive session reloads. The structured payload is serialized into + attachments_json for the read path to reconstruct for the frontend. + """ + import datetime as _dt + + ts = timestamp or _dt.datetime.now(_dt.timezone.utc).isoformat() + payload = { + "error": error, + "error_type": error_type, + "technical_details": technical_details, + "action_required": action_required, + "session_id": session_id, + } + + try: + await _insert_immediate_message_and_sync_session( + session_id=session_id, + role="system", + content=error, + type="error", + agent_name=agent_name, + model_name=model_name, + timestamp=ts, + clean_content=error, + attachments_json=json.dumps(payload), + ) + except Exception as exc: + logging.getLogger(__name__).warning( + "write_error_message_to_sqlite failed (session=%s error_type=%s): %s", + session_id, + error_type, + exc, + exc_info=True, + ) + + +async def insert_message( + *, + session_id: str, + seq: int, + role: str, + content: str = "", + type: str = "", + agent_name: str = "", + model_name: str = "", + timestamp: str, + thinking: Optional[str] = None, + attachments_json: Optional[str] = None, + clean_content: Optional[str] = None, + system_message_type: Optional[str] = None, + system_message_path: Optional[str] = None, + token_count: int = 0, + compacted: int = 0, + pydantic_json: Optional[str] = None, + compaction_log_id: Optional[int] = None, +) -> None: + """Insert a single message row.""" + # Guard: skip empty visible content for user/assistant roles + # unless attachments are present (non-text payloads are still valid). + if ( + role in ("user", "assistant") + and not content.strip() + and not attachments_json + and not pydantic_json + ): + logger.debug( + "insert_message: skipping empty %s message for session=%s seq=%d", + role, + session_id, + seq, + ) + return + + db = get_db() + await db.execute( + """ + INSERT OR IGNORE INTO messages + (session_id, seq, role, content, type, agent_name, model_name, + timestamp, thinking, attachments_json, clean_content, + system_message_type, system_message_path, + token_count, compacted, pydantic_json, compaction_log_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + session_id, + seq, + role, + content, + type, + agent_name, + model_name, + timestamp, + thinking, + attachments_json, + clean_content, + system_message_type, + system_message_path, + token_count, + compacted, + pydantic_json, + compaction_log_id, + ), + ) + await db.commit() + + +async def insert_messages_batch(rows: list[dict[str, Any]]) -> None: + """Insert multiple message rows in a single transaction. + + Each dict in *rows* must have all the same keys as insert_message kwargs. + Uses OR IGNORE so duplicate (session_id, seq) pairs are skipped silently. + Rolls back on any failure so later immediate error/system writes can still + open a fresh transaction. + """ + if not rows: + return + + # Guard: filter out empty visible content for user/assistant roles, + # unless attachments are present (non-text payloads are still valid). + original_count = len(rows) + rows = [ + r + for r in rows + if not ( + r.get("role") in ("user", "assistant") + and not r.get("content", "").strip() + and not r.get("attachments_json") + and not r.get("pydantic_json") + ) + ] + dropped = original_count - len(rows) + if dropped > 0: + logger.info( + "insert_messages_batch: filtered %d empty user/assistant rows", + dropped, + ) + if not rows: + return + + db = get_db() + params = [ + ( + r["session_id"], + r["seq"], + r["role"], + r.get("content", ""), + r.get("type", ""), + r.get("agent_name", ""), + r.get("model_name", ""), + r["timestamp"], + r.get("thinking"), + r.get("attachments_json"), + r.get("clean_content"), + r.get("system_message_type"), + r.get("system_message_path"), + r.get("token_count", 0), + r.get("compacted", 0), + r.get("pydantic_json"), + r.get("compaction_log_id"), + ) + for r in rows + ] + try: + await db.executemany( + """ + INSERT OR IGNORE INTO messages + (session_id, seq, role, content, type, agent_name, model_name, + timestamp, thinking, attachments_json, clean_content, + system_message_type, system_message_path, + token_count, compacted, pydantic_json, compaction_log_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + params, + ) + await db.commit() + except Exception: + await db.rollback() + raise + + +async def get_active_messages(session_id: str) -> list[dict[str, Any]]: + """Return all non-compacted message rows for a session, ordered by seq. + + This is what the agent uses to reconstruct its context via pydantic_json. + """ + db = get_db() + cursor = await db.execute( + """ + SELECT session_id, seq, role, content, type, agent_name, model_name, + timestamp, thinking, attachments_json, clean_content, + system_message_type, system_message_path, + token_count, compacted, pydantic_json, compaction_log_id + FROM messages + WHERE session_id = ? AND compacted = 0 + ORDER BY seq + """, + (session_id,), + ) + rows = await cursor.fetchall() + return [dict(r) for r in rows] + + +# --------------------------------------------------------------------------- +# Tool calls +# --------------------------------------------------------------------------- + + +async def insert_tool_call( + *, + id: str, + session_id: str, + parent_message_seq: Optional[int], + seq: int, + tool_name: str, + args_json: Optional[str] = None, + result_json: Optional[str] = None, + status: str = "success", + duration_ms: Optional[int] = None, + error_text: Optional[str] = None, + agent_name: str = "", + model_name: str = "", + timestamp: float, +) -> None: + db = get_db() + await db.execute( + """ + INSERT OR IGNORE INTO tool_calls + (id, session_id, parent_message_seq, seq, tool_name, + args_json, result_json, status, duration_ms, error_text, + agent_name, model_name, timestamp) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + id, + session_id, + parent_message_seq, + seq, + tool_name, + args_json, + result_json, + status, + duration_ms, + error_text, + agent_name, + model_name, + timestamp, + ), + ) + await db.commit() + + +async def insert_tool_calls_batch(rows: list[dict[str, Any]]) -> None: + """Insert multiple tool call rows in a single transaction.""" + if not rows: + return + db = get_db() + params = [ + ( + r["id"], + r["session_id"], + r.get("parent_message_seq"), + r["seq"], + r["tool_name"], + r.get("args_json"), + r.get("result_json"), + r.get("status", "success"), + r.get("duration_ms"), + r.get("error_text"), + r.get("agent_name", ""), + r.get("model_name", ""), + r["timestamp"], + ) + for r in rows + ] + await db.executemany( + """ + INSERT OR IGNORE INTO tool_calls + (id, session_id, parent_message_seq, seq, tool_name, + args_json, result_json, status, duration_ms, error_text, + agent_name, model_name, timestamp) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + params, + ) + await db.commit() + + +# --------------------------------------------------------------------------- +# Compaction +# --------------------------------------------------------------------------- + + +async def mark_messages_compacted( + session_id: str, start_seq: int, end_seq: int +) -> None: + """Soft-delete a range of messages by setting compacted=1.""" + db = get_db() + await db.execute( + """ + UPDATE messages SET compacted = 1 + WHERE session_id = ? AND seq BETWEEN ? AND ? AND compacted = 0 + """, + (session_id, start_seq, end_seq), + ) + await db.commit() + + +async def insert_compaction_log( + *, + session_id: str, + summary_text: str, + source_start: int, + source_end: int, + source_count: int, + source_tokens: int, + summary_tokens: int, + strategy: str = "summarization", + created_at: str, +) -> int: + """Insert a compaction log entry and return its id.""" + db = get_db() + cursor = await db.execute( + """ + INSERT INTO compaction_log + (session_id, summary_text, source_start, source_end, + source_count, source_tokens, summary_tokens, strategy, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + session_id, + summary_text, + source_start, + source_end, + source_count, + source_tokens, + summary_tokens, + strategy, + created_at, + ), + ) + await db.commit() + return cursor.lastrowid # type: ignore[return-value] + + +# --------------------------------------------------------------------------- +# Per-turn chat write +# --------------------------------------------------------------------------- + + +async def write_turn_to_sqlite( + *, + session_id: str, + enhanced_history: list[Any], + title: str = "", + working_directory: str = "", + pinned: bool = False, + agent_name: str = "code-puppy", + model_name: str = "", + total_tokens: int = 0, + updated_at: str, + created_at: str, + ctx: Any, +) -> None: + """Write an entire enhanced_history to SQLite after a chat turn. + + Called from chat_handler.py after save_session() succeeds. + Uses INSERT OR IGNORE throughout — fully idempotent. + + Message seq: For non-user messages, we query the current max seq from the DB + and increment from there. User messages are skipped (already written pre-stream + via insert_message() with get_next_seq()). + Tool call seq: max_message_seq + tc_global_idx + 1 (after all messages). + Tool calls carry parent_message_seq so the FE doesn't need to infer it. + + Args: + enhanced_history: List of dicts (already-wrapped format from chat_handler): + {'msg': ModelMessage, 'agent': str, 'model': str, 'ts': str} + Possibly also: {'clean_content': str, 'attachments': [...]} + ctx: SessionContext — used for estimate_tokens_for_message(). + Pass None to use a simple len(content)//4 fallback. + """ + import re + import time + from datetime import datetime + + from code_puppy.api.db.message_utils import ( + extract_content, + extract_thinking, + extract_tool_calls, + extract_tool_returns, + get_message_timestamp, + get_role, + pydantic_json_for_message, + ) + + # Upsert session first + try: + await upsert_session( + session_id=session_id, + title=title, + agent_name=agent_name, + model_name=model_name, + working_directory=working_directory, + pinned=pinned, + created_at=created_at, + updated_at=updated_at, + message_count=len(enhanced_history), + total_tokens=total_tokens, + deleted_at=None, + ) + except Exception as exc: + logger.warning("write_turn_to_sqlite: upsert_session failed: %s", exc) + return + + # ------------------------------------------------------------------ # + # Build message rows and collect pending tool calls # + # ------------------------------------------------------------------ # + message_rows: list[dict[str, Any]] = [] + # pending_tool_calls: tool_call_id → {name, args, parent_seq, agent, model, ts} + pending_tool_calls: dict[str, dict[str, Any]] = {} + tool_call_rows_pending: list[dict[str, Any]] = [] + + _ctx_re = re.compile(r"^\[Session Context:[^\]]*\]\n\n?") + + # Get the current max seq from the DB so we can assign correct seq numbers + # to non-user messages. User messages are already written pre-stream with + # get_next_seq(), so we must not collide with them. + current_max_seq = await get_next_seq(session_id) - 1 # get_next_seq returns max+1 + + for idx, item in enumerate(enhanced_history): + # Unwrap the enhanced wrapper + if isinstance(item, dict) and "msg" in item: + actual_msg = item["msg"] + wrapper = item + else: + actual_msg = item + wrapper = {} + + if not hasattr(actual_msg, "parts"): + continue # skip system dict entries + + role = get_role(actual_msg) + ts = wrapper.get("ts") or get_message_timestamp(actual_msg) or updated_at + content = extract_content(actual_msg) + + if role in {"user", "assistant"} and not content.strip(): + diag_tool_calls_present = bool(extract_tool_calls(actual_msg)) + diag_tool_returns_present = bool(extract_tool_returns(actual_msg)) + logger.debug( + "write_turn_to_sqlite empty visible content: " + "session=%s role=%s history_idx=%d has_tool_calls=%s has_tool_returns=%s", + session_id, + role, + idx, + diag_tool_calls_present, + diag_tool_returns_present, + ) + + # Extract tool returns from user (ModelRequest) messages BEFORE skipping. + # BUG FIX: previously this block was placed AFTER the early `continue` + # below, making it dead code — tool_call_rows_pending was never populated + # and insert_tool_calls_batch() always received an empty list. + if role == "user": + for tr in extract_tool_returns(actual_msg): + tid = tr["id"] + if tid in pending_tool_calls: + tc = pending_tool_calls.pop(tid) + try: + args_json_str = json.dumps(tc["args"]) + except Exception: + args_json_str = str(tc["args"]) + # Serialize result - handle Pydantic models and complex objects + result_val = tr["result"] + try: + # Try Pydantic model_dump first + if hasattr(result_val, "model_dump"): + result_json_str = json.dumps(result_val.model_dump()) + elif hasattr(result_val, "dict"): + # Older Pydantic v1 style + result_json_str = json.dumps(result_val.dict()) + elif hasattr(result_val, "__dict__"): + # Generic object - convert to dict + result_json_str = json.dumps(vars(result_val)) + else: + result_json_str = json.dumps(result_val) + except Exception: + # Last resort - try to make it JSON-serializable + try: + result_json_str = json.dumps(str(result_val)) + except Exception: + result_json_str = json.dumps({"raw": str(result_val)}) + + try: + tool_ts = datetime.fromisoformat( + str(tc["ts"]).replace("Z", "+00:00") + ).timestamp() + except Exception: + tool_ts = time.time() + + tool_call_rows_pending.append( + { + "id": tid, + "session_id": session_id, + "parent_message_seq": tc["parent_seq"], + "tool_name": tc["name"], + "args_json": args_json_str, + "result_json": result_json_str, + "status": "success", + "agent_name": tc["agent"], + "model_name": tc["model"], + "timestamp": tool_ts, + } + ) + # Skip writing user message rows — already pre-written by + # chat_handler.py via insert_message()/get_next_seq(). Writing + # them again here would create duplicates. + continue + + # Assign seq for non-user messages: increment from current max + current_max_seq += 1 + seq = current_max_seq + + thinking = extract_thinking(actual_msg) + pj = pydantic_json_for_message(actual_msg) + + msg_agent = wrapper.get("agent") or agent_name or "code-puppy" + msg_model = wrapper.get("model") or model_name or "unknown" + + # clean_content: prefer wrapper value, then strip session context for user + clean_content: str | None = wrapper.get("clean_content") or None + if clean_content is None and role == "user" and content: + stripped = _ctx_re.sub("", content).lstrip() + if stripped != content: + clean_content = stripped + + # attachments_json: from wrapper + raw_attachments = wrapper.get("attachments") + attachments_json: str | None = ( + json.dumps(raw_attachments) if raw_attachments else None + ) + + # token count + try: + token_count = ( + ctx.agent.estimate_tokens_for_message(actual_msg) + if ctx + else max(1, len(content) // 4) + ) + except Exception: + token_count = max(1, len(content) // 4) + + assistant_tool_calls = ( + extract_tool_calls(actual_msg) if role == "assistant" else [] + ) + + # Preserve non-text assistant payloads (tool-call-only responses). + # This keeps the row from being filtered by insert_messages_batch. + if role == "assistant" and assistant_tool_calls and not attachments_json: + attachments_json = json.dumps({"tool_calls": assistant_tool_calls}) + + # Guard: skip only truly empty assistant messages + # (no text AND no tool calls). + if role == "assistant" and not content.strip() and not assistant_tool_calls: + logger.debug( + "write_turn_to_sqlite: skipping empty assistant message " + "session=%s history_idx=%d (no text, no tool calls)", + session_id, + idx, + ) + continue # truly empty - skip + + message_rows.append( + { + "session_id": session_id, + "seq": seq, + "role": role, + "content": content, + "type": type(actual_msg).__name__, + "agent_name": msg_agent, + "model_name": msg_model, + "timestamp": ts, + "thinking": thinking, + "attachments_json": attachments_json, + "clean_content": clean_content, + "token_count": token_count, + "pydantic_json": pj, + } + ) + + # Collect tool calls from assistant messages + if role == "assistant": + for tc in assistant_tool_calls: + pending_tool_calls[tc["id"]] = { + "name": tc["name"], + "args": tc["args"], + "parent_seq": seq, + "agent": msg_agent, + "model": msg_model, + "ts": ts, + } + + # Assign seqs to tool calls: current_max_seq + 1, current_max_seq + 2, ... + # (current_max_seq was already incremented for each message we added) + tool_call_rows: list[dict[str, Any]] = [] + for tc_idx, tc_row in enumerate(tool_call_rows_pending): + tc_row["seq"] = current_max_seq + tc_idx + 1 + tool_call_rows.append(tc_row) + + # Write everything + try: + await insert_messages_batch(message_rows) + except Exception as exc: + logger.warning("write_turn_to_sqlite: insert_messages_batch failed: %s", exc) + + if tool_call_rows: + try: + await insert_tool_calls_batch(tool_call_rows) + except Exception as exc: + logger.warning( + "write_turn_to_sqlite: insert_tool_calls_batch failed: %s", exc + ) + + # Update session stats + computed_tokens = sum(r.get("token_count", 0) for r in message_rows) + final_tokens = total_tokens if total_tokens > 0 else computed_tokens + try: + await update_session_stats( + session_id, + message_count=len(message_rows), + total_tokens=final_tokens, + updated_at=updated_at, + ) + except Exception as exc: + logger.warning("write_turn_to_sqlite: update_session_stats failed: %s", exc) + + # Log with diagnostic counts for forensic debugging + unmatched_calls = len(pending_tool_calls) + logger.info( + "write_turn_to_sqlite: session=%s messages=%d tool_calls=%d pending_unmatched=%d", + session_id, + len(message_rows), + len(tool_call_rows), + unmatched_calls, + ) + + # Warn if tool calls were detected but not matched to returns + if unmatched_calls > 0: + logger.warning( + "write_turn_to_sqlite: %d tool calls not matched to returns for session=%s: %s", + unmatched_calls, + session_id, + list(pending_tool_calls.keys()), + ) + + +# --------------------------------------------------------------------------- +# Interleaved session history query (Full Parity Model) +# +# Returns messages + tool_calls in a single ordered result set with a +# row_type discriminator, matching the desktop GUI's SESSION_LOAD_SQL. +# This enables the browser API to return the same data shape as the desktop. +# --------------------------------------------------------------------------- + +# SQL query for interleaved messages + tool_calls, ordered by seq +_SESSION_HISTORY_PARITY_SQL = load_sql("session_history_parity.sql") +_SESSION_HISTORY_PARITY_NO_COMPACTED_SQL = load_sql( + "session_history_parity_no_compacted.sql" +) + + +async def get_session_history_parity( + session_id: str, + *, + include_compacted: bool = True, +) -> list[dict[str, Any]]: + """Return interleaved messages + tool_calls for a session. + + This is the Full Parity Model query — returns the same row shape as the + desktop GUI's SESSION_LOAD_SQL, enabling consistent behavior across: + - Desktop GUI (Electron) + - Browser chat UI + - API consumers + + Args: + session_id: The session to load + include_compacted: If False, filter out compacted message rows + (tool_call rows are never compacted) + + Returns: + List of dicts with keys: + - row_type: 'message' | 'tool_call' + - row_id: str (message id or tool_call id) + - seq: int (ordering key) + - role: str + - content, type, agent_name, model_name, timestamp, ... + - tool_name, args_json, result_json, status, ... (for tool_call rows) + - parent_message_seq: int | None (for tool_call rows) + """ + db = get_db() + + if include_compacted: + # Use the standard parity SQL + cursor = await db.execute( + _SESSION_HISTORY_PARITY_SQL, + (session_id, session_id), + ) + else: + cursor = await db.execute( + _SESSION_HISTORY_PARITY_NO_COMPACTED_SQL, + (session_id, session_id), + ) + + rows = await cursor.fetchall() + return [dict(r) for r in rows] + + +async def get_session_tool_calls(session_id: str) -> list[dict[str, Any]]: + """Return all tool_call rows for a session, ordered by seq. + + Useful for diagnostics and validation. + """ + db = get_db() + cursor = await db.execute( + """ + SELECT id, session_id, parent_message_seq, seq, tool_name, + args_json, result_json, status, duration_ms, error_text, + agent_name, model_name, timestamp + FROM tool_calls + WHERE session_id = ? + ORDER BY seq + """, + (session_id,), + ) + rows = await cursor.fetchall() + return [dict(r) for r in rows] diff --git a/code_puppy/api/db/seeder.py b/code_puppy/api/db/seeder.py new file mode 100644 index 000000000..d2d31ffde --- /dev/null +++ b/code_puppy/api/db/seeder.py @@ -0,0 +1,17 @@ +"""Retired legacy session seeder. + +The application now treats SQLite + JSON session files as the only supported +runtime storage format. The old pickle-to-SQLite migration path has been +removed; this module remains only as a compatibility import target. +""" + +from __future__ import annotations + +import logging + +logger = logging.getLogger(__name__) + + +async def seed_from_pkl_dirs() -> None: + """No-op compatibility shim for the retired pickle migration path.""" + logger.debug("Legacy session seeder is retired; nothing to migrate") diff --git a/code_puppy/api/db/sql/session_history_parity.sql b/code_puppy/api/db/sql/session_history_parity.sql new file mode 100644 index 000000000..84c85295f --- /dev/null +++ b/code_puppy/api/db/sql/session_history_parity.sql @@ -0,0 +1,59 @@ +SELECT + 'message' AS row_type, + CAST(m.id AS TEXT) AS row_id, + m.seq, + m.role, + m.content, + m.type, + m.agent_name, + m.model_name, + m.timestamp, + m.thinking, + m.attachments_json, + m.clean_content, + m.system_message_type, + m.system_message_path, + m.token_count, + m.compacted, + m.compaction_log_id, + NULL AS tool_name, + NULL AS args_json, + NULL AS result_json, + NULL AS status, + NULL AS duration_ms, + NULL AS error_text, + NULL AS parent_message_seq +FROM messages m +WHERE m.session_id = ? + +UNION ALL + +SELECT + 'tool_call' AS row_type, + tc.id AS row_id, + tc.seq, + 'tool' AS role, + NULL AS content, + NULL AS type, + tc.agent_name, + tc.model_name, + CAST(tc.timestamp AS TEXT) AS timestamp, + NULL AS thinking, + NULL AS attachments_json, + NULL AS clean_content, + NULL AS system_message_type, + NULL AS system_message_path, + 0 AS token_count, + 0 AS compacted, + NULL AS compaction_log_id, + tc.tool_name, + tc.args_json, + tc.result_json, + tc.status, + tc.duration_ms, + tc.error_text, + tc.parent_message_seq +FROM tool_calls tc +WHERE tc.session_id = ? + +ORDER BY seq; diff --git a/code_puppy/api/db/sql/session_history_parity_no_compacted.sql b/code_puppy/api/db/sql/session_history_parity_no_compacted.sql new file mode 100644 index 000000000..fad44d5f5 --- /dev/null +++ b/code_puppy/api/db/sql/session_history_parity_no_compacted.sql @@ -0,0 +1,59 @@ +SELECT + 'message' AS row_type, + CAST(m.id AS TEXT) AS row_id, + m.seq, + m.role, + m.content, + m.type, + m.agent_name, + m.model_name, + m.timestamp, + m.thinking, + m.attachments_json, + m.clean_content, + m.system_message_type, + m.system_message_path, + m.token_count, + m.compacted, + m.compaction_log_id, + NULL AS tool_name, + NULL AS args_json, + NULL AS result_json, + NULL AS status, + NULL AS duration_ms, + NULL AS error_text, + NULL AS parent_message_seq +FROM messages m +WHERE m.session_id = ? AND m.compacted = 0 + +UNION ALL + +SELECT + 'tool_call' AS row_type, + tc.id AS row_id, + tc.seq, + 'tool' AS role, + NULL AS content, + NULL AS type, + tc.agent_name, + tc.model_name, + CAST(tc.timestamp AS TEXT) AS timestamp, + NULL AS thinking, + NULL AS attachments_json, + NULL AS clean_content, + NULL AS system_message_type, + NULL AS system_message_path, + 0 AS token_count, + 0 AS compacted, + NULL AS compaction_log_id, + tc.tool_name, + tc.args_json, + tc.result_json, + tc.status, + tc.duration_ms, + tc.error_text, + tc.parent_message_seq +FROM tool_calls tc +WHERE tc.session_id = ? + +ORDER BY seq; diff --git a/code_puppy/api/db/sql_loader.py b/code_puppy/api/db/sql_loader.py new file mode 100644 index 000000000..47ac4a1a7 --- /dev/null +++ b/code_puppy/api/db/sql_loader.py @@ -0,0 +1,22 @@ +"""Helpers for loading external SQL query files. + +Keeps complex SQL out of Python source while preserving parameterized execution. +""" + +from __future__ import annotations + +from functools import lru_cache +from pathlib import Path + +_SQL_DIR = Path(__file__).with_name("sql") + + +@lru_cache(maxsize=64) +def load_sql(name: str) -> str: + """Load an SQL file from ``code_puppy/api/db/sql``. + + Args: + name: Filename like ``session_history_parity.sql``. + """ + sql_path = _SQL_DIR / name + return sql_path.read_text(encoding="utf-8") diff --git a/code_puppy/api/error_parser.py b/code_puppy/api/error_parser.py new file mode 100644 index 000000000..ffd60fafe --- /dev/null +++ b/code_puppy/api/error_parser.py @@ -0,0 +1,239 @@ +"""Error parsing utilities for API errors. + +Provides human-readable error messages and actionable guidance for common API errors. +""" + +import logging +import re +from typing import Any, Dict + +logger = logging.getLogger(__name__) + + +def parse_api_error(error: Exception) -> Dict[str, Any]: + """Parse API errors into user-friendly messages with actionable guidance. + + Debugging: this function logs the chosen classification at DEBUG level + to help trace WS error propagation to the UI. + + Args: + error: The exception that was raised + + Returns: + Dictionary with: + - error_type: Category of error (quota_exceeded, rate_limit, auth_error, network_error, unknown) + - user_message: Human-readable message for the user + - technical_details: Technical error information + - action_required: Specific action the user should take (optional) + - original_error: String representation of the original error + """ + error_str = str(error) + error_type_name = type(error).__name__ + + logger.debug( + "parse_api_error: incoming type=%s msg=%r", + error_type_name, + error_str[:400], + ) + + # Check for quota exceeded errors (HTTP 429 or explicit quota messages) + if _is_quota_exceeded_error(error_str): + parsed = { + "error_type": "quota_exceeded", + "user_message": "Your API quota has been exceeded for the current model. Please switch to a different model and click continue to retry your request.", + "technical_details": _extract_quota_details(error_str), + "action_required": "switch_model", + "original_error": error_str, + } + logger.debug("parse_api_error: classified as %s", parsed["error_type"]) + return parsed + + # Check for rate limit errors + if _is_rate_limit_error(error_str): + parsed = { + "error_type": "rate_limit", + "user_message": "You're making requests too quickly. Please wait a moment and try again, or switch to a different model.", + "technical_details": error_str, + "action_required": "wait_or_switch_model", + "original_error": error_str, + } + logger.debug("parse_api_error: classified as %s", parsed["error_type"]) + return parsed + + # Check for authentication/permission errors + if _is_auth_error(error_str): + return { + "error_type": "auth_error", + "user_message": "Authentication failed. Please check your credentials or API key configuration.", + "technical_details": error_str, + "action_required": "check_credentials", + "original_error": error_str, + } + + # Check for network/timeout errors + if _is_network_error(error_str, error_type_name): + return { + "error_type": "network_error", + "user_message": "Network connection failed. Please check your internet connection and try again.", + "technical_details": error_str, + "action_required": "retry", + "original_error": error_str, + } + + # Check for model not found errors + if _is_model_not_found_error(error_str): + return { + "error_type": "model_not_found", + "user_message": "The requested model is not available. Please switch to a different model.", + "technical_details": error_str, + "action_required": "switch_model", + "original_error": error_str, + } + + # Check for Claude extended thinking temperature errors + if _is_claude_temperature_thinking_error(error_str): + parsed = { + "error_type": "claude_temperature_error", + "user_message": "Claude's extended thinking mode requires temperature to be set to 1.0. Please adjust your model settings or disable extended thinking.", + "technical_details": error_str, + "action_required": "adjust_temperature", + "original_error": error_str, + } + logger.debug("parse_api_error: classified as %s", parsed["error_type"]) + return parsed + + # Generic error fallback + parsed = { + "error_type": "unknown", + "user_message": f"An unexpected error occurred: {error_str}", + "technical_details": error_str, + "action_required": None, + "original_error": error_str, + } + logger.debug("parse_api_error: classified as %s", parsed["error_type"]) + return parsed + + +def _is_quota_exceeded_error(error_str: str) -> bool: + """Check if error is a quota exceeded error.""" + quota_patterns = [ + r"quota\s+exceeded", + r"HTTP\s+429", + r"aiplatform\.googleapis\.com/online_prediction_requests_per_base_model", + r"resource_exhausted", + r"too\s+many\s+requests.*quota", + ] + + for pattern in quota_patterns: + if re.search(pattern, error_str, re.IGNORECASE): + return True + return False + + +def _is_rate_limit_error(error_str: str) -> bool: + """Check if error is a rate limit error.""" + rate_limit_patterns = [ + r"rate\s+limit", + r"too\s+many\s+requests", + r"throttle", + r"requests.*per.*second", + r"requests.*per.*minute", + ] + + for pattern in rate_limit_patterns: + if re.search(pattern, error_str, re.IGNORECASE): + # Don't double-match quota errors + if not _is_quota_exceeded_error(error_str): + return True + return False + + +def _is_auth_error(error_str: str) -> bool: + """Check if error is an authentication error.""" + auth_patterns = [ + r"authentication\s+failed", + r"unauthorized", + r"HTTP\s+401", + r"HTTP\s+403", + r"permission\s+denied", + r"invalid.*api.*key", + r"invalid.*credentials", + r"access\s+denied", + ] + + for pattern in auth_patterns: + if re.search(pattern, error_str, re.IGNORECASE): + return True + return False + + +def _is_network_error(error_str: str, error_type_name: str) -> bool: + """Check if error is a network/connectivity error.""" + # Check error type names + network_error_types = [ + "TimeoutError", + "ConnectionError", + "ConnectTimeout", + "ReadTimeout", + "HTTPError", + ] + + if any(err_type in error_type_name for err_type in network_error_types): + return True + + # Check error message patterns + network_patterns = [ + r"timeout", + r"connection\s+refused", + r"connection\s+reset", + r"connection\s+failed", + r"network\s+error", + r"unable\s+to\s+connect", + r"host\s+unreachable", + ] + + for pattern in network_patterns: + if re.search(pattern, error_str, re.IGNORECASE): + return True + return False + + +def _is_model_not_found_error(error_str: str) -> bool: + """Check if error is a model not found error.""" + # Use simple patterns that match anywhere in the string + if re.search(r"model.*not\s+found", error_str, re.IGNORECASE): + return True + if re.search(r"invalid\s+model", error_str, re.IGNORECASE): + return True + if re.search(r"unknown\s+model", error_str, re.IGNORECASE): + return True + if re.search(r"model.*does\s+not\s+exist", error_str, re.IGNORECASE): + return True + if re.search(r"HTTP\s+404.*model", error_str, re.IGNORECASE): + return True + + return False + + +def _is_claude_temperature_thinking_error(error_str: str) -> bool: + """Check if error is Claude's temperature/thinking configuration error.""" + return ( + "temperature" in error_str.lower() + and "thinking" in error_str.lower() + and ("may only be set to 1" in error_str or "must be 1" in error_str) + ) + + +def _extract_quota_details(error_str: str) -> str: + """Extract specific quota details from error message.""" + # Try to extract the specific quota metric if present + match = re.search(r"quota\s+exceeded\s+for\s+([\w\./]+)", error_str, re.IGNORECASE) + if match: + return f"Quota exceeded for: {match.group(1)}" + + # Try to extract model name if present + match = re.search(r"base\s+model[:\s]+([\w\-]+)", error_str, re.IGNORECASE) + if match: + return f"Quota exceeded for base model: {match.group(1)}" + + return error_str diff --git a/code_puppy/api/routers/ws_sessions.py b/code_puppy/api/routers/ws_sessions.py new file mode 100644 index 000000000..dff3a5c84 --- /dev/null +++ b/code_puppy/api/routers/ws_sessions.py @@ -0,0 +1,290 @@ +"""WebSocket session management router. + +Handles CRUD operations for chat sessions. +Now exclusively uses SQLite as the source of truth. +""" + +import json +import logging +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, List, Optional + +from fastapi import APIRouter, HTTPException, Query + +from code_puppy.api.db.queries import ( + get_active_messages, + get_session_history_parity, + get_session_metadata, + get_session_tool_calls, + soft_delete_session, + update_session_meta_fields, +) +from code_puppy.config import get_ws_sessions_dir + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +def _get_ws_sessions_dir() -> Path: + """Get the WebSocket sessions directory.""" + return get_ws_sessions_dir() + + +def _validate_session_name(session_name: str, ws_dir: Path) -> str: + """Validate session name prevents path traversal.""" + if not session_name or ".." in session_name or "/" in session_name: + raise HTTPException(400, "Invalid session name") + return session_name + + +@router.get("/{session_name}/messages") +async def get_ws_session_messages( + session_name: str, + include_tool_calls: Optional[bool] = Query( + default=None, + description="Include tool_call rows interleaved with messages (Full Parity Model). " + "When true, returns unified row format with row_type discriminator.", + ), + include_compacted: bool = Query( + default=True, + description="Include compacted (summarized) messages. Only applies when include_tool_calls=true.", + ), +) -> List[Dict[str, Any]]: + """Get the full message history for a WebSocket session from SQLite. + + Args: + session_name: The session name / session_id + include_tool_calls: Feature flag for Full Parity Model. + - None/False: Legacy behavior, returns only message rows + - True: Returns interleaved message + tool_call rows with row_type discriminator + include_compacted: When using parity mode, whether to include compacted messages + + Returns: + List of serialized message dictionaries. + + When include_tool_calls=true, each row has: + - row_type: 'message' | 'tool_call' + - row_id: unique identifier + - seq: ordering key + - Plus type-specific fields (content for messages, result_json for tool_calls, etc.) + """ + _validate_session_name(session_name, _get_ws_sessions_dir()) + + try: + # ────────────────────────────────────────────────────────────────── + # Full Parity Model: interleaved messages + tool_calls + # ────────────────────────────────────────────────────────────────── + if include_tool_calls: + rows = await get_session_history_parity( + session_name, + include_compacted=include_compacted, + ) + + if rows: + logger.debug( + "get_ws_session_messages (parity): session=%s total_rows=%d " + "messages=%d tool_calls=%d", + session_name, + len(rows), + sum(1 for r in rows if r.get("row_type") == "message"), + sum(1 for r in rows if r.get("row_type") == "tool_call"), + ) + return rows + + # Check if session exists but has no data + meta = await get_session_metadata(session_name) + if meta: + return [] + + raise HTTPException(404, f"WebSocket session '{session_name}' not found") + + # ────────────────────────────────────────────────────────────────── + # Legacy behavior: message rows only (backward compatible) + # ────────────────────────────────────────────────────────────────── + rows = await get_active_messages(session_name) + serialized_rows: List[Dict[str, Any]] = [] + + for row in rows: + base_row = { + "role": row["role"], + "content": row["content"], + "type": row["type"], + "agent_name": row["agent_name"], + "model_name": row["model_name"], + "timestamp": row["timestamp"], + "thinking": row["thinking"], + "clean_content": row["clean_content"], + "seq": row["seq"], + } + + if row.get("pydantic_json"): + serialized_rows.append(base_row) + continue + + if row.get("type") == "error": + payload: Dict[str, Any] = {} + raw_payload = row.get("attachments_json") + if raw_payload: + try: + payload = json.loads(raw_payload) + except (TypeError, ValueError): + logger.warning( + "get_ws_session_messages: invalid error payload JSON for %s seq=%s", + session_name, + row.get("seq"), + ) + + serialized_rows.append( + { + **base_row, + "type": "error", + "error": payload.get("error", row["content"]), + "error_type": payload.get("error_type", "unknown"), + "technical_details": payload.get("technical_details", ""), + "action_required": payload.get("action_required"), + "session_id": payload.get("session_id", session_name), + } + ) + continue + + if row.get("type") == "system": + serialized_rows.append( + { + **base_row, + "system_message_type": row.get("system_message_type"), + "system_message_path": row.get("system_message_path"), + } + ) + continue + + serialized_rows.append(base_row) + + if serialized_rows: + return serialized_rows + + # If no messages found, check if session exists at all + meta = await get_session_metadata(session_name) + if meta: + # Session exists but has no messages yet - return empty list + return [] + + raise HTTPException(404, f"WebSocket session '{session_name}' not found") + + except HTTPException: + raise + except Exception as e: + logger.error( + "get_ws_session_messages: SQLite error for %s: %s", session_name, e + ) + raise HTTPException(503, "Service unavailable: database error") + + +@router.get("/{session_name}/tool-calls") +async def get_ws_session_tool_calls(session_name: str) -> List[Dict[str, Any]]: + """Get all tool calls for a session. + + This is a diagnostic/debugging endpoint that returns just the tool_calls + table rows for a session, useful for verifying persistence. + + Args: + session_name: The session name / session_id + + Returns: + List of tool_call rows with id, tool_name, args_json, result_json, status, etc. + """ + _validate_session_name(session_name, _get_ws_sessions_dir()) + + try: + rows = await get_session_tool_calls(session_name) + + if rows: + logger.debug( + "get_ws_session_tool_calls: session=%s count=%d", + session_name, + len(rows), + ) + return rows + + # Check if session exists + meta = await get_session_metadata(session_name) + if meta: + return [] # Session exists but has no tool calls + + raise HTTPException(404, f"WebSocket session '{session_name}' not found") + + except HTTPException: + raise + except Exception as e: + logger.error( + "get_ws_session_tool_calls: SQLite error for %s: %s", session_name, e + ) + raise HTTPException(503, "Service unavailable: database error") + + +@router.delete("/{session_name}") +async def delete_ws_session(session_name: str) -> Dict[str, str]: + """Soft-delete a WebSocket session in SQLite.""" + _validate_session_name(session_name, _get_ws_sessions_dir()) + + # Soft delete in SQLite + try: + await soft_delete_session(session_name, datetime.now(timezone.utc).isoformat()) + # We don't check if it existed, soft_delete is idempotent-ish + # But we might want to know if we actually updated anything. + # For now, we assume success if no error. + except Exception as e: + logger.error("delete_ws_session: SQLite error for %s: %s", session_name, e) + raise HTTPException(503, "Service unavailable: database error") + + return {"message": f"WebSocket session '{session_name}' deleted"} + + +@router.patch("/{session_name}") +async def update_ws_session( + session_name: str, updates: Dict[str, Any] +) -> Dict[str, Any]: + """Update WebSocket session metadata in SQLite. + + Supports updating: title, pinned. + """ + _validate_session_name(session_name, _get_ws_sessions_dir()) + + # Get current metadata + current_meta = await get_session_metadata(session_name) + if not current_meta: + raise HTTPException(404, f"WebSocket session '{session_name}' not found") + + # Prepare updates + new_title = current_meta["title"] + new_pinned = bool(current_meta["pinned"]) + updated_at = datetime.now(timezone.utc).isoformat() + + if "title" in updates or "name" in updates: + val = updates.get("title") or updates.get("name") + if val and isinstance(val, str) and val.strip(): + new_title = val.strip() + + if "pinned" in updates: + if isinstance(updates["pinned"], bool): + new_pinned = updates["pinned"] + + # Write to SQLite + try: + await update_session_meta_fields( + session_name, + title=new_title, + pinned=new_pinned, + updated_at=updated_at, + ) + except Exception as e: + logger.error("update_ws_session: SQLite error for %s: %s", session_name, e) + raise HTTPException(503, "Service unavailable: database error") + + return { + "session_id": session_name, + "title": new_title, + "pinned": new_pinned, + "updated_at": updated_at, + } diff --git a/code_puppy/api/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_db_sql_loader.py b/code_puppy/api/tests/test_db_sql_loader.py new file mode 100644 index 000000000..d60a3c912 --- /dev/null +++ b/code_puppy/api/tests/test_db_sql_loader.py @@ -0,0 +1,22 @@ +from pathlib import Path + + +def test_external_session_history_sql_files_exist_and_contain_expected_clauses(): + sql_dir = Path("code_puppy/api/db/sql") + parity_sql = (sql_dir / "session_history_parity.sql").read_text(encoding="utf-8") + no_compacted_sql = (sql_dir / "session_history_parity_no_compacted.sql").read_text( + encoding="utf-8" + ) + + assert "FROM messages m" in parity_sql + assert "FROM tool_calls tc" in parity_sql + assert "ORDER BY seq" in parity_sql + + assert "m.compacted = 0" in no_compacted_sql + assert "ORDER BY seq" in no_compacted_sql + + +def test_queries_module_uses_sql_loader_for_complex_session_query(): + queries_src = Path("code_puppy/api/db/queries.py").read_text(encoding="utf-8") + assert 'load_sql("session_history_parity.sql")' in queries_src + assert "session_history_parity_no_compacted.sql" in queries_src diff --git a/code_puppy/api/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/config.py b/code_puppy/config.py index 10d1326c6..df9932239 100644 --- a/code_puppy/config.py +++ b/code_puppy/config.py @@ -53,18 +53,27 @@ def _get_xdg_dir(env_var: str, fallback: str) -> str: SKILLS_DIR = os.path.join(DATA_DIR, "skills") CONTEXTS_DIR = os.path.join(DATA_DIR, "contexts") -# OAuth plugin model files (XDG_DATA_HOME) -GEMINI_MODELS_FILE = os.path.join(DATA_DIR, "gemini_models.json") -CHATGPT_MODELS_FILE = os.path.join(DATA_DIR, "chatgpt_models.json") -CLAUDE_MODELS_FILE = os.path.join(DATA_DIR, "claude_models.json") -COPILOT_MODELS_FILE = os.path.join(DATA_DIR, "copilot_models.json") - # Cache files (XDG_CACHE_HOME) AUTOSAVE_DIR = os.path.join(CACHE_DIR, "autosaves") +WS_SESSION_DIR = os.path.join(CACHE_DIR, "ws_sessions") + + +def get_ws_sessions_dir() -> pathlib.Path: + """Return the ws_sessions directory path and ensure it exists.""" + p = pathlib.Path(WS_SESSION_DIR) + p.mkdir(parents=True, exist_ok=True, mode=0o700) + return p + # State files (XDG_STATE_HOME) COMMAND_HISTORY_FILE = os.path.join(STATE_DIR, "command_history.txt") +# OAuth plugin model files (XDG_DATA_HOME) +GEMINI_MODELS_FILE = os.path.join(DATA_DIR, "gemini_models.json") +CHATGPT_MODELS_FILE = os.path.join(DATA_DIR, "chatgpt_models.json") +CLAUDE_MODELS_FILE = os.path.join(DATA_DIR, "claude_models.json") +COPILOT_MODELS_FILE = os.path.join(DATA_DIR, "copilot_models.json") + def get_subagent_verbose() -> bool: """Return True if sub-agent verbose output is enabled (default False). diff --git a/tests/api/test_ws_sessions_router.py b/tests/api/test_ws_sessions_router.py new file mode 100644 index 000000000..d1d54aaf3 --- /dev/null +++ b/tests/api/test_ws_sessions_router.py @@ -0,0 +1,88 @@ +import importlib.util +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest +from starlette.routing import Router + + +_original_router_init = Router.__init__ + + +def _compat_router_init(self, *args, **kwargs): + kwargs.pop("on_startup", None) + kwargs.pop("on_shutdown", None) + kwargs.pop("lifespan", None) + return _original_router_init(self, *args, **kwargs) + + +Router.__init__ = _compat_router_init +try: + module_path = ( + Path(__file__).resolve().parents[2] + / "code_puppy" + / "api" + / "routers" + / "ws_sessions.py" + ) + spec = importlib.util.spec_from_file_location( + "test_ws_sessions_module", module_path + ) + assert spec is not None and spec.loader is not None + ws_sessions = importlib.util.module_from_spec(spec) + spec.loader.exec_module(ws_sessions) +finally: + Router.__init__ = _original_router_init + + +@pytest.mark.asyncio +async def test_get_ws_session_messages_keeps_plain_legacy_rows(monkeypatch): + monkeypatch.setattr( + ws_sessions, + "_validate_session_name", + lambda session_name, ws_dir: session_name, + ) + monkeypatch.setattr( + ws_sessions, + "get_active_messages", + AsyncMock( + return_value=[ + { + "role": "assistant", + "content": "plain legacy row", + "type": "message", + "agent_name": "code-puppy", + "model_name": "gpt-5", + "timestamp": "2026-01-01T00:00:00Z", + "thinking": None, + "clean_content": "plain legacy row", + "seq": 7, + "pydantic_json": None, + } + ] + ), + ) + monkeypatch.setattr( + ws_sessions, + "get_session_metadata", + AsyncMock(return_value=None), + ) + + result = await ws_sessions.get_ws_session_messages( + "session-1", + include_tool_calls=False, + ) + + assert result == [ + { + "role": "assistant", + "content": "plain legacy row", + "type": "message", + "agent_name": "code-puppy", + "model_name": "gpt-5", + "timestamp": "2026-01-01T00:00:00Z", + "thinking": None, + "clean_content": "plain legacy row", + "seq": 7, + } + ] From c6f31ed1863e3e0ebb4eb8b03d27dcf6c6e6370d Mon Sep 17 00:00:00 2001 From: sskmlm Date: Mon, 13 Jul 2026 19:41:33 -0500 Subject: [PATCH 2/4] refactor: simplify ws session schema and add project id --- code_puppy/api/db/connection.py | 335 +++------------------ code_puppy/api/db/queries.py | 24 +- code_puppy/api/routers/ws_sessions.py | 12 +- code_puppy/api/tests/test_db_connection.py | 42 +++ tests/api/test_ws_sessions_router.py | 43 +++ 5 files changed, 151 insertions(+), 305 deletions(-) create mode 100644 code_puppy/api/tests/test_db_connection.py diff --git a/code_puppy/api/db/connection.py b/code_puppy/api/db/connection.py index d60f9e2b2..e852b4eca 100644 --- a/code_puppy/api/db/connection.py +++ b/code_puppy/api/db/connection.py @@ -1,21 +1,9 @@ -"""SQLite connection singleton using aiosqlite for non-blocking async access. +"""SQLite connection singleton for the pre-release WS session database. -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. +This feature has not shipped to users yet, so we intentionally keep the schema +policy boring: one canonical schema, one user_version, zero migration archaeology. +If a developer has an older pre-release database, they should delete it and let +Code Puppy recreate it from scratch. """ from __future__ import annotations @@ -29,44 +17,27 @@ 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_VERSION = 1 _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 + session_id TEXT PRIMARY KEY, + title TEXT DEFAULT '', + project_id 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 INDEX IF NOT EXISTS idx_sessions_project_id ON sessions(project_id); CREATE TABLE IF NOT EXISTS compaction_log ( id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -128,91 +99,26 @@ def get_db_path() -> Path: 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() +_aconn: Optional[aiosqlite.Connection] = None -SCHEMA_VERSION = 5 -# --------------------------------------------------------------------------- -# Async singleton (aiosqlite) -# --------------------------------------------------------------------------- +def get_db_path() -> Path: + """Return path to the shared SQLite database. -_aconn: Optional[aiosqlite.Connection] = None + 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" 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. - """ + """Open the database and ensure the canonical pre-release schema exists.""" global _aconn if _aconn is not None: - return # already initialised + return db_path = get_db_path() db_path.parent.mkdir(parents=True, exist_ok=True) @@ -221,188 +127,27 @@ async def init_db() -> None: _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) - + current_version = int(row[0]) if row else 0 -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" + if current_version not in {0, SCHEMA_VERSION}: + raise RuntimeError( + "Unsupported pre-release WS session DB schema " + f"v{current_version} at {db_path}. Delete the database and restart " + "to recreate it with the current schema." ) - from code_puppy.api.db.backfill import backfill_pydantic_json - updated, skipped = await backfill_pydantic_json(_aconn) + await _aconn.executescript(_SCHEMA_SQL) + if current_version != SCHEMA_VERSION: 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) - + await _aconn.commit() -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 + logger.info(" aiosqlite DB ready (schema v%d)", SCHEMA_VERSION) def get_db() -> aiosqlite.Connection: diff --git a/code_puppy/api/db/queries.py b/code_puppy/api/db/queries.py index 1a26ad0b1..b0e4eba94 100644 --- a/code_puppy/api/db/queries.py +++ b/code_puppy/api/db/queries.py @@ -52,7 +52,7 @@ async def get_session_row(session_id: str) -> Optional[dict]: 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, + Returns keys: session_id, title, project_id, agent_name, model_name, working_directory, pinned, created_at — or None if not found. Non-fatal: returns None on any exception. """ @@ -60,7 +60,7 @@ async def get_session_metadata(session_id: str) -> Optional[dict]: db = get_db() cursor = await db.execute( """ - SELECT session_id, title, agent_name, model_name, + SELECT session_id, title, project_id, agent_name, model_name, working_directory, pinned, created_at FROM sessions WHERE session_id = ? @@ -82,6 +82,7 @@ async def upsert_session( agent_name: str = "code-puppy", model_name: str = "", working_directory: str = "", + project_id: str = "", pinned: bool = False, created_at: str, updated_at: str, @@ -95,11 +96,12 @@ async def upsert_session( await db.execute( """ INSERT INTO sessions - (session_id, title, agent_name, model_name, working_directory, + (session_id, title, project_id, agent_name, model_name, working_directory, pinned, created_at, updated_at, message_count, total_tokens, deleted_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(session_id) DO UPDATE SET title = excluded.title, + project_id = excluded.project_id, agent_name = excluded.agent_name, model_name = excluded.model_name, working_directory = excluded.working_directory, @@ -112,6 +114,7 @@ async def upsert_session( ( session_id, title, + project_id, agent_name, model_name, working_directory, @@ -175,19 +178,20 @@ async def update_session_meta_fields( session_id: str, *, title: str, + project_id: str, pinned: bool, updated_at: str, ) -> None: - """Update only title, pinned, and updated_at for an existing session. + """Update only title, project_id, pinned, and updated_at for a 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). + model_name, or working_directory. Called by the update_session_meta WS + message handler when the user renames, tags, 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), + "UPDATE sessions SET title = ?, project_id = ?, pinned = ?, updated_at = ? WHERE session_id = ?", + (title, project_id, 1 if pinned else 0, updated_at, session_id), ) await db.commit() @@ -773,6 +777,7 @@ async def write_turn_to_sqlite( title: str = "", working_directory: str = "", pinned: bool = False, + project_id: str = "", agent_name: str = "code-puppy", model_name: str = "", total_tokens: int = 0, @@ -821,6 +826,7 @@ async def write_turn_to_sqlite( model_name=model_name, working_directory=working_directory, pinned=pinned, + project_id=project_id, created_at=created_at, updated_at=updated_at, message_count=len(enhanced_history), diff --git a/code_puppy/api/routers/ws_sessions.py b/code_puppy/api/routers/ws_sessions.py index dff3a5c84..21f361c7b 100644 --- a/code_puppy/api/routers/ws_sessions.py +++ b/code_puppy/api/routers/ws_sessions.py @@ -247,7 +247,7 @@ async def update_ws_session( ) -> Dict[str, Any]: """Update WebSocket session metadata in SQLite. - Supports updating: title, pinned. + Supports updating: title, project_id, pinned. """ _validate_session_name(session_name, _get_ws_sessions_dir()) @@ -258,6 +258,7 @@ async def update_ws_session( # Prepare updates new_title = current_meta["title"] + new_project_id = str(current_meta.get("project_id") or "") new_pinned = bool(current_meta["pinned"]) updated_at = datetime.now(timezone.utc).isoformat() @@ -266,6 +267,13 @@ async def update_ws_session( if val and isinstance(val, str) and val.strip(): new_title = val.strip() + if "project_id" in updates: + value = updates.get("project_id") + if value is None: + new_project_id = "" + elif isinstance(value, str): + new_project_id = value.strip() + if "pinned" in updates: if isinstance(updates["pinned"], bool): new_pinned = updates["pinned"] @@ -275,6 +283,7 @@ async def update_ws_session( await update_session_meta_fields( session_name, title=new_title, + project_id=new_project_id, pinned=new_pinned, updated_at=updated_at, ) @@ -285,6 +294,7 @@ async def update_ws_session( return { "session_id": session_name, "title": new_title, + "project_id": new_project_id, "pinned": new_pinned, "updated_at": updated_at, } diff --git a/code_puppy/api/tests/test_db_connection.py b/code_puppy/api/tests/test_db_connection.py new file mode 100644 index 000000000..b84f48bbd --- /dev/null +++ b/code_puppy/api/tests/test_db_connection.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +import sqlite3 + +import pytest + + +@pytest.mark.asyncio +async def test_init_db_creates_sessions_project_id_column(tmp_path, monkeypatch): + from code_puppy.api.db.connection import close_db, get_db, init_db + + db_path = tmp_path / "chat_messages.db" + monkeypatch.setenv("PUPPY_DESK_DB", str(db_path)) + + await init_db() + db = get_db() + cursor = await db.execute("PRAGMA table_info(sessions)") + rows = await cursor.fetchall() + await close_db() + + column_names = {row[1] for row in rows} + assert "project_id" in column_names + + +@pytest.mark.asyncio +async def test_init_db_rejects_old_prerelease_schema(tmp_path, monkeypatch): + from code_puppy.api.db.connection import close_db, init_db + + db_path = tmp_path / "chat_messages.db" + conn = sqlite3.connect(db_path) + conn.execute("PRAGMA user_version = 5") + conn.commit() + conn.close() + + monkeypatch.setenv("PUPPY_DESK_DB", str(db_path)) + + with pytest.raises( + RuntimeError, match="Unsupported pre-release WS session DB schema" + ): + await init_db() + + await close_db() diff --git a/tests/api/test_ws_sessions_router.py b/tests/api/test_ws_sessions_router.py index d1d54aaf3..5ade742af 100644 --- a/tests/api/test_ws_sessions_router.py +++ b/tests/api/test_ws_sessions_router.py @@ -86,3 +86,46 @@ async def test_get_ws_session_messages_keeps_plain_legacy_rows(monkeypatch): "seq": 7, } ] + + +@pytest.mark.asyncio +async def test_update_ws_session_supports_project_id(monkeypatch): + monkeypatch.setattr( + ws_sessions, + "_validate_session_name", + lambda session_name, ws_dir: session_name, + ) + monkeypatch.setattr( + ws_sessions, + "get_session_metadata", + AsyncMock( + return_value={ + "session_id": "session-1", + "title": "Original", + "project_id": "old-project", + "pinned": 0, + } + ), + ) + update_mock = AsyncMock() + monkeypatch.setattr(ws_sessions, "update_session_meta_fields", update_mock) + + result = await ws_sessions.update_ws_session( + "session-1", + { + "title": " Renamed ", + "project_id": " project-alpha ", + "pinned": True, + }, + ) + + update_mock.assert_awaited_once() + _, kwargs = update_mock.await_args + assert kwargs["title"] == "Renamed" + assert kwargs["project_id"] == "project-alpha" + assert kwargs["pinned"] is True + + assert result["session_id"] == "session-1" + assert result["title"] == "Renamed" + assert result["project_id"] == "project-alpha" + assert result["pinned"] is True From 835fe151045d25afd56b5c8dbd4a25f0d47ad2ea Mon Sep 17 00:00:00 2001 From: sskmlm Date: Tue, 14 Jul 2026 10:14:19 -0500 Subject: [PATCH 3/4] test: stub fastapi in ws sessions router test --- tests/api/test_ws_sessions_router.py | 31 ++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/api/test_ws_sessions_router.py b/tests/api/test_ws_sessions_router.py index 5ade742af..2e150a9a4 100644 --- a/tests/api/test_ws_sessions_router.py +++ b/tests/api/test_ws_sessions_router.py @@ -1,4 +1,6 @@ import importlib.util +import sys +import types from pathlib import Path from unittest.mock import AsyncMock @@ -6,6 +8,35 @@ from starlette.routing import Router +class _HTTPException(Exception): + def __init__(self, status_code, detail): + super().__init__(detail) + self.status_code = status_code + self.detail = detail + + +class _APIRouter: + def get(self, *args, **kwargs): + return lambda fn: fn + + def delete(self, *args, **kwargs): + return lambda fn: fn + + def patch(self, *args, **kwargs): + return lambda fn: fn + + +def _query(*args, **kwargs): + return kwargs.get("default") + + +_fastapi_stub = types.ModuleType("fastapi") +_fastapi_stub.APIRouter = _APIRouter +_fastapi_stub.HTTPException = _HTTPException +_fastapi_stub.Query = _query +sys.modules.setdefault("fastapi", _fastapi_stub) + + _original_router_init = Router.__init__ From 422fe3079b2f7cb86f40ab5c131a7abe8a969269 Mon Sep 17 00:00:00 2001 From: sskmlm Date: Tue, 14 Jul 2026 10:26:08 -0500 Subject: [PATCH 4/4] test: isolate ws sessions router imports --- tests/api/test_ws_sessions_router.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/api/test_ws_sessions_router.py b/tests/api/test_ws_sessions_router.py index 2e150a9a4..f58798bef 100644 --- a/tests/api/test_ws_sessions_router.py +++ b/tests/api/test_ws_sessions_router.py @@ -37,6 +37,25 @@ def _query(*args, **kwargs): sys.modules.setdefault("fastapi", _fastapi_stub) +async def _unused_async(*args, **kwargs): + raise AssertionError("test stub should be monkeypatched before use") + + +_queries_stub = types.ModuleType("code_puppy.api.db.queries") +_queries_stub.get_active_messages = _unused_async +_queries_stub.get_session_history_parity = _unused_async +_queries_stub.get_session_metadata = _unused_async +_queries_stub.get_session_tool_calls = _unused_async +_queries_stub.soft_delete_session = _unused_async +_queries_stub.update_session_meta_fields = _unused_async +sys.modules.setdefault("code_puppy.api.db.queries", _queries_stub) + + +_config_stub = types.ModuleType("code_puppy.config") +_config_stub.get_ws_sessions_dir = lambda: Path("/tmp/code-puppy-test-ws-sessions") +sys.modules.setdefault("code_puppy.config", _config_stub) + + _original_router_init = Router.__init__