diff --git a/navi/core/pg_session_store.py b/navi/core/pg_session_store.py index 816fcc2..ab0ba09 100644 --- a/navi/core/pg_session_store.py +++ b/navi/core/pg_session_store.py @@ -1,1164 +1,53 @@ -"""PostgreSQL-backed session store using asyncpg connection pool. +"""Compatibility shim — the PostgreSQL session store lives in `navi.core.sessions`. -Normalized storage: -- session_messages — hot messages (recent N per session) -- session_messages_archive — old messages moved here to keep the hot table small -- session_images — base64 images referenced by session_messages (future use) - - sessions.next_sequence — monotonic global seq for the session (next free number) - sessions.archive_threshold — all seq < threshold live in the archive table - -All session message I/O goes through the normalized tables. The legacy JSON -columns (sessions.messages, sessions.context) remain in the schema for backward -compatibility but are no longer read or written. +Existing imports (`from navi.core.pg_session_store import PgSessionStore`) keep +working; new code should import from `navi.core.sessions`. """ -import asyncio -import json -import secrets -from datetime import datetime, timezone - -import asyncpg -import structlog - -from navi.llm.base import Message, ToolCallRequest - -from .session import Session, SessionStore - -_DDL = """ -CREATE TABLE IF NOT EXISTS sessions ( - id TEXT PRIMARY KEY, - profile_id TEXT NOT NULL, - user_id TEXT REFERENCES navi_users(id) ON DELETE SET NULL, - messages TEXT NOT NULL DEFAULT '[]', - context TEXT NOT NULL DEFAULT '', - pinned BOOLEAN NOT NULL DEFAULT FALSE, - special BOOLEAN NOT NULL DEFAULT FALSE, - created_at TIMESTAMPTZ NOT NULL, - last_active TIMESTAMPTZ NOT NULL, - context_token_count INTEGER NOT NULL DEFAULT 0, - name TEXT, - session_metadata TEXT NOT NULL DEFAULT '{}' +from .sessions import PgSessionStore, pending_sweep_loop, repair_dangling_tool_calls +from .sessions.ddl import ( + _DDL, + _MIGRATE, + _SESSION_ARCHIVE_DDL, + _SESSION_IMAGES_DDL, + _SESSION_MESSAGES_DDL, + _SESSION_SHARES_DDL, + _SYNAPSE_INDEX_DDL, + _ensure_normalized_tables, + _migrate_to_normalized, ) -""" +from .sessions.rows import ( + _build_sessions, + _deserialize, + _load_messages_map, + _message_key, + _row_to_message, + _serialize, + _special_and, + _special_frag, + _special_where, +) -_MIGRATE = """ -ALTER TABLE sessions ADD COLUMN IF NOT EXISTS name TEXT; -ALTER TABLE sessions ADD COLUMN IF NOT EXISTS planning_logs TEXT NOT NULL DEFAULT '[]'; -ALTER TABLE sessions ADD COLUMN IF NOT EXISTS user_id TEXT REFERENCES navi_users(id) ON DELETE SET NULL; -ALTER TABLE sessions ADD COLUMN IF NOT EXISTS next_sequence INTEGER NOT NULL DEFAULT 0; -ALTER TABLE sessions ADD COLUMN IF NOT EXISTS archive_threshold INTEGER NOT NULL DEFAULT 0; -ALTER TABLE sessions ADD COLUMN IF NOT EXISTS session_metadata TEXT NOT NULL DEFAULT '{}'; -ALTER TABLE sessions ADD COLUMN IF NOT EXISTS special BOOLEAN NOT NULL DEFAULT FALSE -""" - -# Reaction threads are found by (owner, conversation key) among special sessions -# only — a tiny slice of the table, so the index is partial. The expression is the -# exact one the lookup uses; `session_metadata` is a TEXT column holding JSON. -_SYNAPSE_INDEX_DDL = """ -CREATE INDEX IF NOT EXISTS idx_sessions_synapse_session_key - ON sessions (user_id, (session_metadata::jsonb -> 'synapse' ->> 'session_key')) - WHERE special -""" - -_SESSION_MESSAGES_DDL = """ -CREATE TABLE IF NOT EXISTS session_messages ( - id SERIAL PRIMARY KEY, - session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, - sequence_number INT NOT NULL, - role TEXT NOT NULL, - content TEXT, - images TEXT, -- JSON array of base64 strings - tool_calls TEXT, -- JSON - tool_call_id TEXT, - name TEXT, - created_at TIMESTAMPTZ, - is_summary BOOLEAN NOT NULL DEFAULT FALSE, - thinking TEXT, - is_plan BOOLEAN NOT NULL DEFAULT FALSE, - is_compression BOOLEAN NOT NULL DEFAULT FALSE, - is_context BOOLEAN NOT NULL DEFAULT TRUE, - is_display BOOLEAN NOT NULL DEFAULT TRUE, - elapsed_seconds FLOAT, - tool_call_count INT, - token_count INT, - files TEXT, -- JSON - metadata TEXT, -- JSON - is_recall BOOLEAN NOT NULL DEFAULT FALSE, - UNIQUE(session_id, sequence_number) -); - -CREATE INDEX IF NOT EXISTS idx_session_messages_session_seq ON session_messages(session_id, sequence_number); -CREATE INDEX IF NOT EXISTS idx_session_messages_context ON session_messages(session_id, is_context, sequence_number); -""" - -_SESSION_IMAGES_DDL = """ -CREATE TABLE IF NOT EXISTS session_images ( - id SERIAL PRIMARY KEY, - session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, - message_id INT REFERENCES session_messages(id) ON DELETE CASCADE, - base64 TEXT NOT NULL, - filename TEXT -); - -CREATE INDEX IF NOT EXISTS idx_session_images_session ON session_images(session_id); -CREATE INDEX IF NOT EXISTS idx_session_images_message ON session_images(message_id); -""" - -_SESSION_ARCHIVE_DDL = """ -CREATE TABLE IF NOT EXISTS session_messages_archive ( - id SERIAL PRIMARY KEY, - session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, - sequence_number INT NOT NULL, - role TEXT NOT NULL, - content TEXT, - images TEXT, - tool_calls TEXT, - tool_call_id TEXT, - name TEXT, - created_at TIMESTAMPTZ, - is_summary BOOLEAN NOT NULL DEFAULT FALSE, - thinking TEXT, - is_plan BOOLEAN NOT NULL DEFAULT FALSE, - is_compression BOOLEAN NOT NULL DEFAULT FALSE, - is_context BOOLEAN NOT NULL DEFAULT TRUE, - is_display BOOLEAN NOT NULL DEFAULT TRUE, - elapsed_seconds FLOAT, - tool_call_count INT, - token_count INT, - files TEXT, - metadata TEXT, - is_recall BOOLEAN NOT NULL DEFAULT FALSE, - UNIQUE(session_id, sequence_number) -); -CREATE INDEX IF NOT EXISTS idx_session_messages_archive_session_seq ON session_messages_archive(session_id, sequence_number); -""" - -_SESSION_SHARES_DDL = """ -CREATE TABLE IF NOT EXISTS session_shares ( - session_id TEXT PRIMARY KEY REFERENCES sessions(id) ON DELETE CASCADE, - token TEXT NOT NULL UNIQUE, - owner_id TEXT, - created_at TIMESTAMPTZ NOT NULL DEFAULT now() -); -""" - - -def _serialize(messages: list[Message]) -> str: - return json.dumps( - [m.model_dump(mode="json", exclude_none=True) for m in messages], - ensure_ascii=False, - ) - - -def _deserialize(raw: str) -> list[Message]: - if not raw: - return [] - return [Message.model_validate(m) for m in json.loads(raw)] - - -def repair_dangling_tool_calls(messages: list[Message]) -> int: - """Synthesize placeholder results for assistant tool_calls that lost their - role=tool pair (a crash mid-batch, e.g. a parallel batch killed by a server - restart). Returns the number of placeholders added. - - Placeholders are inserted right after the last tool message of the affected - batch (before the next non-tool message) and marked ``is_context=False`` — - display history stays coherent, the LLM context is untouched. - """ - answered = { - m.tool_call_id for m in messages if m.role == "tool" and m.tool_call_id - } - out: list[Message] = [] - pending: list[Message] = [] - for m in messages: - if m.role != "tool" and pending: - out.extend(pending) - pending = [] - out.append(m) - if m.role == "assistant" and m.tool_calls: - for tc in m.tool_calls: - if not tc.id or tc.id in answered: - continue - pending.append( - Message( - role="tool", - content="[Interrupted before result]", - tool_call_id=tc.id, - name=tc.name, - metadata={}, - is_context=False, - ) - ) - out.extend(pending) - if len(out) != len(messages): - added = len(out) - len(messages) - messages[:] = out - return added - return 0 - - -def _message_key(m: Message) -> tuple: - """Stable key for matching a message between messages[] and context[]. - - Used only by the one-shot boot migration for legacy JSON rows. - """ - return ( - m.role, - m.content, - m.tool_call_id, - m.name, - m.is_summary, - m.is_plan, - m.is_compression, - m.is_recall, - m.thinking, - m.created_at.isoformat() if m.created_at else None, - json.dumps([tc.model_dump(mode="json") for tc in m.tool_calls], ensure_ascii=False) if m.tool_calls else None, - json.dumps(m.files, ensure_ascii=False) if m.files else None, - json.dumps(m.metadata, ensure_ascii=False) if m.metadata else None, - ) - - -def _row_to_message(row: asyncpg.Record) -> Message: - images = json.loads(row["images"]) if row.get("images") else None - tool_calls = None - raw_tool_calls = row.get("tool_calls") - if raw_tool_calls: - tool_calls = [ToolCallRequest.model_validate(tc) for tc in json.loads(raw_tool_calls)] - files = json.loads(row["files"]) if row.get("files") else None - metadata = json.loads(row["metadata"]) if row.get("metadata") else {} - - return Message( - role=row["role"], - content=row["content"], - images=images, - tool_calls=tool_calls, - tool_call_id=row.get("tool_call_id"), - name=row.get("name"), - created_at=row.get("created_at"), - is_summary=bool(row.get("is_summary", False)), - thinking=row.get("thinking"), - is_plan=bool(row.get("is_plan", False)), - is_compression=bool(row.get("is_compression", False)), - is_context=bool(row.get("is_context", True)), - is_display=bool(row.get("is_display", True)), - elapsed_seconds=row.get("elapsed_seconds"), - tool_call_count=row.get("tool_call_count"), - token_count=row.get("token_count"), - files=files, - metadata=metadata, - is_recall=bool(row.get("is_recall", False)), - sequence_number=row.get("sequence_number", 0), - ) - - -async def _ensure_normalized_tables(conn: asyncpg.Connection) -> None: - await conn.execute(_SESSION_MESSAGES_DDL) - await conn.execute(_SESSION_IMAGES_DDL) - await conn.execute(_SESSION_ARCHIVE_DDL) - await conn.execute(_SESSION_SHARES_DDL) - # Best-effort: a row whose session_metadata is not valid JSON would make the - # expression index fail and take the whole schema init down with it. The - # lookup still works without the index, so a failure is a slowdown, not an - # outage — log it and carry on. - try: - await conn.execute(_SYNAPSE_INDEX_DDL) - except Exception as e: # noqa: BLE001 — never block startup over an index - structlog.get_logger().warning( - "pg_session_store.synapse_index_failed", error=str(e) - ) - - -async def _migrate_to_normalized(conn: asyncpg.Connection) -> None: - """One-shot migration: copy existing JSON messages/context into session_messages.""" - migrated = await conn.fetchval("SELECT COUNT(*) FROM session_messages") - if migrated: - return - - rows = await conn.fetch("SELECT id, messages, context FROM sessions") - for row in rows: - session_id = row["id"] - messages = _deserialize(row["messages"] or "[]") - context = _deserialize(row["context"] or "") - if not context and messages: - context = list(messages) - - context_set = {_message_key(m) for m in context} - - for seq, m in enumerate(messages): - is_ctx = _message_key(m) in context_set - await conn.execute( - """ - INSERT INTO session_messages - (session_id, sequence_number, role, content, images, tool_calls, tool_call_id, name, - created_at, is_summary, thinking, is_plan, is_compression, is_context, is_display, - elapsed_seconds, tool_call_count, token_count, files, metadata, is_recall) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21) - ON CONFLICT (session_id, sequence_number) DO NOTHING - """, - session_id, - seq, - m.role, - m.content, - json.dumps(m.images, ensure_ascii=False) if m.images else None, - json.dumps([tc.model_dump(mode="json") for tc in m.tool_calls], ensure_ascii=False) if m.tool_calls else None, - m.tool_call_id, - m.name, - m.created_at, - m.is_summary, - m.thinking, - m.is_plan, - m.is_compression, - is_ctx, - True, - m.elapsed_seconds, - m.tool_call_count, - m.token_count, - json.dumps(m.files, ensure_ascii=False) if m.files else None, - json.dumps(m.metadata, ensure_ascii=False) if m.metadata else None, - m.is_recall, - ) - - -async def _load_messages_map(conn: asyncpg.Connection, session_ids: list[str]) -> dict[str, list[Message]]: - """Batch-load all messages for the given session IDs.""" - if not session_ids: - return {} - rows = await conn.fetch( - "SELECT * FROM session_messages WHERE session_id = ANY($1) ORDER BY sequence_number", - session_ids, - ) - result: dict[str, list[Message]] = {sid: [] for sid in session_ids} - for row in rows: - sid = row["session_id"] - result[sid].append(_row_to_message(row)) - return result - - -async def _build_sessions( - conn: asyncpg.Connection, - rows: list[asyncpg.Record], -) -> list[Session]: - """Hydrate session rows with hot (non-archived) messages.""" - session_ids = [r["id"] for r in rows] - messages_map = await _load_messages_map(conn, session_ids) - sessions: list[Session] = [] - for row in rows: - all_msgs = messages_map.get(row["id"], []) - archive_threshold = row.get("archive_threshold", 0) or 0 - hot_msgs = [m for m in all_msgs if m.sequence_number >= archive_threshold] - messages = [m for m in hot_msgs if m.is_display] - context = [m for m in hot_msgs if m.is_context] - planning_logs_raw = row["planning_logs"] - planning_logs = json.loads(planning_logs_raw) if planning_logs_raw else [] - - next_seq = row.get("next_sequence", 0) or 0 - max_seq = max((m.sequence_number for m in hot_msgs), default=-1) - if next_seq == 0: - next_seq = max_seq + 1 - - s = Session( - id=row["id"], - profile_id=row["profile_id"], - user_id=row["user_id"], - messages=messages, - context=context, - pinned=bool(row["pinned"]), - special=bool(row.get("special", False)), - name=row["name"], - created_at=row["created_at"], - last_active=row["last_active"], - context_token_count=row["context_token_count"] or 0, - planning_logs=planning_logs, - ) - session_metadata = json.loads(row.get("session_metadata") or "{}") if row.get("session_metadata") else {} - s.db_message_count = len(hot_msgs) - s.db_next_sequence = next_seq - s.archive_threshold = archive_threshold - s.session_metadata = session_metadata - sessions.append(s) - return sessions - - -def _special_frag(special: bool | None) -> str: - """SQL fragment for the special filter; '' when unfiltered.""" - if special is None: - return "" - return f"special = {str(bool(special)).upper()}" - - -def _special_and(special: bool | None) -> str: - frag = _special_frag(special) - return f" AND {frag}" if frag else "" - - -def _special_where(special: bool | None) -> str: - frag = _special_frag(special) - return f" WHERE {frag}" if frag else "" - - -class PgSessionStore(SessionStore): - def __init__(self, pool: asyncpg.Pool) -> None: - self._pool = pool - self._initialized = False - self._lock = asyncio.Lock() - # Sessions created via create() but not yet persisted to the DB. - # They live only in memory until the first save() upserts the row - # (on the first user message / meaningful state change), so empty - # sessions that never receive a message never reach the DB. - self._pending: dict[str, Session] = {} - self._pending_lock = asyncio.Lock() - - async def _get_pool(self) -> asyncpg.Pool: - if not self._initialized: - async with self._lock: - if not self._initialized: - async with self._pool.acquire() as conn: - await conn.execute(_DDL) - await conn.execute(_MIGRATE) - await _ensure_normalized_tables(conn) - await _migrate_to_normalized(conn) - self._initialized = True - return self._pool - - async def create(self, profile_id: str, user_id: str | None = None) -> Session: - # Lazy persistence: do NOT insert a DB row here. Register the session - # in memory; the row is upserted on the first save(). This keeps empty - # sessions (created but never sent a message) out of the DB entirely. - session = Session(profile_id=profile_id, user_id=user_id) - async with self._pending_lock: - self._pending[session.id] = session - return session - - async def get(self, session_id: str) -> Session | None: - # Pending (not-yet-persisted) sessions live only in memory. - async with self._pending_lock: - pending = self._pending.get(session_id) - if pending is not None: - return pending - - pool = await self._get_pool() - async with pool.acquire() as conn: - row = await conn.fetchrow( - "SELECT id, profile_id, user_id, pinned, special, created_at, last_active, context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata " - "FROM sessions WHERE id = $1", - session_id, - ) - if not row: - return None - - archive_threshold = row["archive_threshold"] or 0 - # Load hot (non-archived) messages so messages[] and context[] share - # the same Python objects (id() matching in the agent works). - all_rows = await conn.fetch( - "SELECT * FROM session_messages WHERE session_id = $1 AND sequence_number >= $2 ORDER BY sequence_number", - session_id, - archive_threshold, - ) - all_messages = [_row_to_message(r) for r in all_rows] - # A crash mid-tool-batch (e.g. parallel batch) can leave assistant - # tool_calls without their role=tool pair — patch display history. - repair_dangling_tool_calls(all_messages) - messages = [m for m in all_messages if m.is_display] - context = [m for m in all_messages if m.is_context] - - planning_logs_raw = row["planning_logs"] - planning_logs = json.loads(planning_logs_raw) if planning_logs_raw else [] - - next_seq = row["next_sequence"] or 0 - max_seq = max((m.sequence_number for m in all_messages), default=-1) - if next_seq == 0: - next_seq = max_seq + 1 - - session_metadata = json.loads(row.get("session_metadata") or "{}") if row.get("session_metadata") else {} - s = Session( - id=row["id"], - profile_id=row["profile_id"], - user_id=row["user_id"], - messages=messages, - context=context, - pinned=bool(row["pinned"]), - special=bool(row.get("special", False)), - name=row["name"], - created_at=row["created_at"], - last_active=row["last_active"], - context_token_count=row["context_token_count"] or 0, - planning_logs=planning_logs, - ) - s.db_message_count = len(all_messages) - s.db_next_sequence = next_seq - s.archive_threshold = archive_threshold - s.session_metadata = session_metadata - return s - - async def save(self, session: Session) -> None: - session.last_active = datetime.now(timezone.utc) - pool = await self._get_pool() - # One transaction: the sequence range handed out below and the rows that - # consume it must commit together, or a crash in between burns numbers - # and leaves the counter ahead of the messages. - async with pool.acquire() as conn, conn.transaction(): - # Upsert the session row. On the first save() after a lazy create() - # the row does not exist yet — INSERT ... ON CONFLICT creates it so - # the session_messages FK is satisfied before messages are inserted. - # created_at / pinned / name / next_sequence / archive_threshold are - # set on insert and not overwritten on conflict (managed elsewhere). - # profile_id is on that list too, for the same reason: set_profile() - # owns the column after creation, while a run's Session can still - # carry the profile it started with — writing it here would silently - # undo a switch_profile made mid-run (agent.py then reloads the row - # it just clobbered and sees no change at all). - await conn.execute( - "INSERT INTO sessions " - "(id, profile_id, user_id, pinned, special, created_at, last_active, " - " context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata) " - "VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) " - "ON CONFLICT (id) DO UPDATE SET " - " user_id = EXCLUDED.user_id, " - " last_active = EXCLUDED.last_active, " - " context_token_count = EXCLUDED.context_token_count, " - " planning_logs = EXCLUDED.planning_logs, " - " session_metadata = EXCLUDED.session_metadata", - session.id, session.profile_id, session.user_id, session.pinned, - session.special, session.created_at, session.last_active, - session.context_token_count, - session.name, - json.dumps(session.planning_logs, ensure_ascii=False), - session.db_next_sequence, session.archive_threshold, - json.dumps(session.session_metadata, ensure_ascii=False), - ) - # Row now guaranteed to exist — drop the pending entry if any. - async with self._pending_lock: - self._pending.pop(session.id, None) - - messages = session.messages - - # 1. Update mutable flags for already-persisted rows (sequence_number >= 0) - existing = [m for m in messages if m.sequence_number >= 0] - if existing: - update_rows = [] - for m in existing: - update_rows.append( - ( - session.id, - m.sequence_number, - m.is_context, - m.is_display, - m.is_summary, - m.is_plan, - m.is_compression, - m.is_recall, - m.thinking, - m.elapsed_seconds, - m.tool_call_count, - m.token_count, - ) - ) - await conn.executemany( - """ - UPDATE session_messages - SET is_context = $3, - is_display = $4, - is_summary = $5, - is_plan = $6, - is_compression = $7, - is_recall = $8, - thinking = $9, - elapsed_seconds = $10, - tool_call_count = $11, - token_count = $12 - WHERE session_id = $1 AND sequence_number = $2 - """, - update_rows, - ) - - # 2. Insert new messages (sequence_number < 0 means "not yet persisted") - new_msgs = [m for m in messages if m.sequence_number < 0] - if new_msgs: - # Claim the range in the DB, never from session.db_next_sequence: - # another writer of the same session may have consumed those - # numbers since this object was loaded. switch_profile used to do - # exactly that — it loaded the session mid-run and saved it back — - # and the running turn then inserted onto numbers already taken, - # dying on UNIQUE(session_id, sequence_number). GREATEST() also - # heals rows whose counter is still 0 (pre-counter sessions). - base = await conn.fetchval( - """ - UPDATE sessions s - SET next_sequence = GREATEST( - s.next_sequence, - COALESCE((SELECT MAX(m.sequence_number) + 1 - FROM session_messages m - WHERE m.session_id = s.id), 0) - ) + $2 - WHERE s.id = $1 - RETURNING s.next_sequence - $2 - """, - session.id, - len(new_msgs), - ) - insert_rows = [] - for i, m in enumerate(new_msgs): - seq = base + i - m.sequence_number = seq - insert_rows.append( - ( - session.id, - seq, - m.role, - m.content, - json.dumps(m.images, ensure_ascii=False) if m.images else None, - json.dumps([tc.model_dump(mode="json") for tc in m.tool_calls], ensure_ascii=False) if m.tool_calls else None, - m.tool_call_id, - m.name, - m.created_at, - m.is_summary, - m.thinking, - m.is_plan, - m.is_compression, - m.is_context, - m.is_display, - m.elapsed_seconds, - m.tool_call_count, - m.token_count, - json.dumps(m.files, ensure_ascii=False) if m.files else None, - json.dumps(m.metadata, ensure_ascii=False) if m.metadata else None, - m.is_recall, - ) - ) - await conn.executemany( - """ - INSERT INTO session_messages - (session_id, sequence_number, role, content, images, tool_calls, tool_call_id, name, - created_at, is_summary, thinking, is_plan, is_compression, is_context, is_display, - elapsed_seconds, tool_call_count, token_count, files, metadata, is_recall) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21) - """, - insert_rows, - ) - session.db_next_sequence = base + len(new_msgs) - - session.db_message_count = len(messages) - - async def sweep_pending(self, max_age_seconds: int = 3600) -> int: - """Drop pending (not-yet-persisted) sessions older than max_age_seconds. - - Bounds memory used by clients that create a session and never send a - message. Returns the number of entries removed. - """ - cutoff = datetime.now(timezone.utc).timestamp() - max_age_seconds - removed = 0 - async with self._pending_lock: - for sid in list(self._pending): - s = self._pending[sid] - if s.created_at.timestamp() < cutoff: - del self._pending[sid] - removed += 1 - return removed - - async def set_context_flag( - self, session_id: str, sequence_numbers: list[int], is_context: bool - ) -> int: - """Persist a context-flag flip for rows save() cannot see. - - ``save()`` updates flags only for rows in ``session.messages``. A - summary is ``is_display=False``: after a reload it lives in - ``session.context`` alone, so compression's demotion of a superseded - summary never reached the database — every reload resurrected it, and - prod sessions accumulated up to 15 live summaries (80% of the context). - """ - if not sequence_numbers: - return 0 - pool = await self._get_pool() - async with pool.acquire() as conn: - result = await conn.execute( - "UPDATE session_messages SET is_context = $3 " - "WHERE session_id = $1 AND sequence_number = ANY($2::int[])", - session_id, - list(sequence_numbers), - is_context, - ) - parts = result.split() - return int(parts[-1]) if len(parts) >= 2 else 0 - - async def archive_old_messages(self, session_id: str, keep_seq_threshold: int) -> int: - """Move messages older than keep_seq_threshold from hot to archive table. - - Returns number of rows archived. - """ - pool = await self._get_pool() - async with pool.acquire() as conn: - # Copy old rows to archive - copied = await conn.execute( - """ - INSERT INTO session_messages_archive - SELECT * FROM session_messages - WHERE session_id = $1 AND sequence_number < $2 - ON CONFLICT (session_id, sequence_number) DO NOTHING - """, - session_id, keep_seq_threshold, - ) - # Delete from hot table - await conn.execute( - "DELETE FROM session_messages WHERE session_id = $1 AND sequence_number < $2", - session_id, keep_seq_threshold, - ) - # Update threshold on session - await conn.execute( - "UPDATE sessions SET archive_threshold = $1 WHERE id = $2", - keep_seq_threshold, session_id, - ) - # asyncpg execute returns 'INSERT 0 N' — extract N - parts = copied.split() - return int(parts[-1]) if len(parts) >= 2 else 0 - - async def get_archived_messages( - self, session_id: str, before_seq: int | None = None, limit: int = 50 - ) -> list[Message]: - """Return archived messages for a session, newest first.""" - pool = await self._get_pool() - async with pool.acquire() as conn: - if before_seq is not None: - rows = await conn.fetch( - "SELECT * FROM session_messages_archive " - "WHERE session_id = $1 AND sequence_number < $2 " - "ORDER BY sequence_number DESC LIMIT $3", - session_id, before_seq, limit, - ) - else: - rows = await conn.fetch( - "SELECT * FROM session_messages_archive " - "WHERE session_id = $1 " - "ORDER BY sequence_number DESC LIMIT $2", - session_id, limit, - ) - # Return oldest-first so the caller can prepend them in order - return [_row_to_message(r) for r in reversed(rows)] - - async def get_meta(self, session_id: str) -> Session | None: - """Session row without any message payload — cheap head for the UI. - - Unlike get() this never reads session_messages; message counts come - from a single aggregate. For pending (not-yet-persisted) sessions it - returns the in-memory object like get() does. - """ - async with self._pending_lock: - pending = self._pending.get(session_id) - if pending is not None: - return pending - - pool = await self._get_pool() - async with pool.acquire() as conn: - row = await conn.fetchrow( - "SELECT id, profile_id, user_id, pinned, special, created_at, last_active, context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata " - "FROM sessions WHERE id = $1", - session_id, - ) - if not row: - return None - msg_count = await conn.fetchval( - "SELECT COUNT(*) FROM session_messages WHERE session_id = $1", - session_id, - ) - planning_logs_raw = row["planning_logs"] - planning_logs = json.loads(planning_logs_raw) if planning_logs_raw else [] - metadata_raw = row.get("session_metadata") - s = Session( - id=row["id"], - profile_id=row["profile_id"], - user_id=row["user_id"], - messages=[], - context=[], - pinned=bool(row["pinned"]), - special=bool(row.get("special", False)), - name=row["name"], - created_at=row["created_at"], - last_active=row["last_active"], - context_token_count=row["context_token_count"] or 0, - planning_logs=planning_logs, - ) - s.db_message_count = int(msg_count or 0) - next_seq = row["next_sequence"] or (s.db_message_count + 1) - s.db_next_sequence = int(next_seq) - s.archive_threshold = row["archive_threshold"] or 0 - s.session_metadata = json.loads(metadata_raw or "{}") if metadata_raw else {} - return s - - async def get_history_page( - self, session_id: str, before_seq: int | None = None, limit: int = 150 - ) -> tuple[list[Message], bool]: - """One page of full history (archive ∪ hot), display messages only. - - Returns (messages oldest-first, has_more). With before_seq=None the - page starts at the newest message — the tail the UI loads initially; - with before_seq set it continues below that sequence number (scroll-up - pagination). UNION (not ALL) collapses rows doubled across hot/archive - by a crash between the archive copy and delete. - - Each message carries ``metadata.display_index`` — its 0-based rank in - the full display history — so paged clients build ids that match a full - reload (feedback keys and message jumps depend on that stability). - """ - pool = await self._get_pool() - async with pool.acquire() as conn: - rows = await conn.fetch( - """ - WITH hist AS ( - SELECT * FROM session_messages WHERE session_id = $1 AND is_display - UNION - SELECT * FROM session_messages_archive WHERE session_id = $1 AND is_display - ), ranked AS ( - SELECT *, ROW_NUMBER() OVER (ORDER BY sequence_number) - 1 AS display_index - FROM hist - ) - SELECT * FROM ranked - WHERE ($2::int IS NULL OR sequence_number < $2) - ORDER BY sequence_number DESC - LIMIT $3 - """, - session_id, before_seq, limit + 1, - ) - has_more = len(rows) > limit - # SQL rank of each row: 0-based position in the full display history. - rank_by_seq = {r["sequence_number"]: r["display_index"] for r in rows[:limit]} - msgs = [_row_to_message(r) for r in rows[:limit]] - msgs.reverse() # oldest-first - # A page boundary can split a tool batch — patch the fragment like get() does. - repair_dangling_tool_calls(msgs) - # Synthetic placeholders from the repair have no sequence number and no - # rank — give them the next rank after the last real message so client - # ids stay collision-free. - last_known: int | None = None - for m in msgs: - di = rank_by_seq.get(m.sequence_number) if m.sequence_number >= 0 else None - if di is None and last_known is not None: - di = last_known + 1 - if di is not None: - m.metadata = dict(m.metadata or {}) | {"display_index": di} - last_known = di - return msgs, has_more - - async def set_pinned(self, session_id: str, pinned: bool) -> bool: - pool = await self._get_pool() - async with pool.acquire() as conn: - result = await conn.execute( - "UPDATE sessions SET pinned = $1 WHERE id = $2", - pinned, session_id, - ) - return result == "UPDATE 1" - - async def set_name(self, session_id: str, name: str) -> bool: - # A never-persisted session lives only in _pending (lazy persistence): - # the bare UPDATE would touch 0 rows and silently drop the name. - async with self._pending_lock: - pending = self._pending.get(session_id) - if pending is not None: - pending.name = name - return True - pool = await self._get_pool() - async with pool.acquire() as conn: - result = await conn.execute( - "UPDATE sessions SET name = $1 WHERE id = $2", - name, session_id, - ) - return result == "UPDATE 1" - - async def set_profile(self, session_id: str, profile_id: str) -> bool: - """Repoint a session at another profile — one narrow UPDATE, no save(). - - switch_profile runs *inside* a live agent run that holds its own Session - object. Loading the session and saving it back (as it once did) makes the - tool a second writer of the same session, and the run's next insert then - lands on a sequence number this tool already took. Touching only the - profile column keeps the two out of each other's way; the run picks the - change up through its post-switch profile reload. - """ - # A never-persisted session lives only in _pending (lazy persistence): - # the bare UPDATE would touch 0 rows and silently drop the switch. - async with self._pending_lock: - pending = self._pending.get(session_id) - if pending is not None: - pending.profile_id = profile_id - return True - pool = await self._get_pool() - async with pool.acquire() as conn: - result = await conn.execute( - "UPDATE sessions SET profile_id = $1 WHERE id = $2", - profile_id, session_id, - ) - return result == "UPDATE 1" - - async def find_reaction_session( - self, - *, - user_id: str, - session_key: str, - not_active_before: datetime | None = None, - ) -> Session | None: - """The newest special session of `user_id` carrying this conversation key. - - One indexed lookup, then `get()` for the hydration — rather than repeating - the message/context loading here. Sessions still in `_pending` are skipped - on purpose: the reaction handler saves the session (and therefore persists - the key) before it starts the run, so a key that is not in the table yet - belongs to a session whose reaction has not begun. - """ - pool = await self._get_pool() - async with pool.acquire() as conn: - row = await conn.fetchrow( - """ - SELECT id FROM sessions - WHERE special - AND user_id = $1 - AND session_metadata::jsonb -> 'synapse' ->> 'session_key' = $2 - AND ($3::timestamptz IS NULL OR last_active >= $3) - ORDER BY last_active DESC - LIMIT 1 - """, - user_id, session_key, not_active_before, - ) - if row is None: - return None - return await self.get(row["id"]) - - async def list_all( - self, user_id: str | None = None, is_admin: bool = False, special: bool | None = None - ) -> list[Session]: - """Sessions visible to the caller. - - Scoping: admins see everything; a named owner sees their own rows; - ``user_id=None`` (local mode, the anonymous local user) sees only - NULL-owner sessions — never other users' chats. - ``special``: None — no filter, False — regular only, True — special only. - """ - pool = await self._get_pool() - async with pool.acquire() as conn: - if is_admin: - rows = await conn.fetch( - "SELECT id, profile_id, user_id, pinned, special, created_at, last_active, context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata " - f"FROM sessions{_special_where(special)} ORDER BY pinned DESC, last_active DESC" - ) - elif user_id is not None: - rows = await conn.fetch( - "SELECT id, profile_id, user_id, pinned, special, created_at, last_active, context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata " - f"FROM sessions WHERE user_id = $1{_special_and(special)} ORDER BY pinned DESC, last_active DESC", - user_id, - ) - else: - rows = await conn.fetch( - "SELECT id, profile_id, user_id, pinned, special, created_at, last_active, context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata " - f"FROM sessions WHERE user_id IS NULL{_special_and(special)} ORDER BY pinned DESC, last_active DESC" - ) - return await _build_sessions(conn, rows) - - async def list_page( - self, - *, - limit: int, - offset: int, - profile_id: str | None = None, - user_id: str | None = None, - is_admin: bool = False, - special: bool | None = None, - ) -> list[Session]: - pool = await self._get_pool() - async with pool.acquire() as conn: - conditions = [] - params: list = [] - param_idx = 0 - - def add_param(value): - nonlocal param_idx - param_idx += 1 - params.append(value) - return f"${param_idx}" - - if not is_admin and user_id is None: - # Local mode: the anonymous local user owns NULL sessions only. - conditions.append("user_id IS NULL") - elif not is_admin: - conditions.append(f"user_id = {add_param(user_id)}") - if profile_id: - conditions.append(f"profile_id = {add_param(profile_id)}") - frag = _special_frag(special) - if frag: - conditions.append(frag) - - where = "WHERE " + " AND ".join(conditions) if conditions else "" - order_limit = f"ORDER BY pinned DESC, last_active DESC LIMIT {add_param(limit)} OFFSET {add_param(offset)}" - - rows = await conn.fetch( - "SELECT id, profile_id, user_id, pinned, special, created_at, last_active, context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata " - f"FROM sessions {where} {order_limit}", - *params, - ) - return await _build_sessions(conn, rows) - - async def delete(self, session_id: str) -> bool: - # A never-persisted session lives only in _pending (lazy persistence: - # create() defers the row until the first save()), so the bare DB - # DELETE alone reports 0 rows for a very real session. Pop the - # pending entry first, then count the DB result. - async with self._pending_lock: - pending = self._pending.pop(session_id, None) - pool = await self._get_pool() - async with pool.acquire() as conn: - result = await conn.execute("DELETE FROM sessions WHERE id = $1", session_id) - # Raw chunks have no FK on purpose (memory DDL may run before the - # sessions table exists), so deleting a chat must clean them here — - # otherwise a deleted conversation stays searchable. Best-effort: - # the session row is already gone, and a failure here must not turn - # a successful delete into a 500 (or leave the caller thinking the - # delete did not happen). - try: - await conn.execute( - "DELETE FROM memory_chunks WHERE session_id = $1", session_id - ) - except asyncpg.UndefinedTableError: - pass - except Exception: - structlog.get_logger().warning( - "session.chunk_cleanup_failed", - session_id=session_id, - exc_info=True, - ) - return pending is not None or result == "DELETE 1" - - async def get_or_create_share(self, session_id: str, owner_id: str | None) -> dict | None: - pool = await self._get_pool() - async with pool.acquire() as conn: - try: - row = await conn.fetchrow( - """ - INSERT INTO session_shares (session_id, token, owner_id) - VALUES ($1, $2, $3) - ON CONFLICT (session_id) DO NOTHING - RETURNING token, session_id, owner_id, created_at - """, - session_id, secrets.token_urlsafe(32), owner_id, - ) - except asyncpg.ForeignKeyViolationError: - # A pending session has no row yet — there is nothing to share. - return None - if row is None: # already shared: hand back the existing link - row = await conn.fetchrow( - "SELECT token, session_id, owner_id, created_at " - "FROM session_shares WHERE session_id = $1", - session_id, - ) - return dict(row) if row is not None else None - - async def get_share(self, token: str) -> dict | None: - pool = await self._get_pool() - async with pool.acquire() as conn: - row = await conn.fetchrow( - "SELECT token, session_id, owner_id, created_at " - "FROM session_shares WHERE token = $1", - token, - ) - return dict(row) if row is not None else None - - async def count_all( - self, - *, - user_id: str | None = None, - is_admin: bool = False, - search: str | None = None, - special: bool | None = None, - ) -> int: - pool = await self._get_pool() - async with pool.acquire() as conn: - conditions = [] - params: list = [] - param_idx = 0 - - def add_param(value): - nonlocal param_idx - param_idx += 1 - params.append(value) - return f"${param_idx}" - - if not is_admin and user_id is None: - # Local mode: the anonymous local user owns NULL sessions only. - conditions.append("user_id IS NULL") - elif not is_admin: - conditions.append(f"user_id = {add_param(user_id)}") - if search: - like = f"%{search}%" - conditions.append( - f"(id ILIKE {add_param(like)} OR name ILIKE {add_param(like)} OR user_id ILIKE {add_param(like)} OR profile_id ILIKE {add_param(like)} OR EXISTS (SELECT 1 FROM session_messages m WHERE m.session_id = sessions.id AND m.is_display = true AND m.content ILIKE {add_param(like)}) OR EXISTS (SELECT 1 FROM session_messages_archive a WHERE a.session_id = sessions.id AND a.is_display = true AND a.content ILIKE {add_param(like)}))" - ) - frag = _special_frag(special) - if frag: - conditions.append(frag) - - where = "WHERE " + " AND ".join(conditions) if conditions else "" - row = await conn.fetchrow(f"SELECT COUNT(*) FROM sessions {where}", *params) - return row["count"] if row else 0 - - async def search_list( - self, - *, - limit: int, - offset: int, - user_id: str | None = None, - is_admin: bool = False, - search: str | None = None, - sort_by: str = "last_active", - sort_order: str = "desc", - special: bool | None = None, - ) -> list[Session]: - pool = await self._get_pool() - async with pool.acquire() as conn: - conditions = [] - params: list = [] - param_idx = 0 - - def add_param(value): - nonlocal param_idx - param_idx += 1 - params.append(value) - return f"${param_idx}" - - if not is_admin and user_id is None: - # Local mode: the anonymous local user owns NULL sessions only. - conditions.append("user_id IS NULL") - elif not is_admin: - conditions.append(f"user_id = {add_param(user_id)}") - if search: - like = f"%{search}%" - conditions.append( - f"(id ILIKE {add_param(like)} OR name ILIKE {add_param(like)} OR user_id ILIKE {add_param(like)} OR profile_id ILIKE {add_param(like)} OR EXISTS (SELECT 1 FROM session_messages m WHERE m.session_id = sessions.id AND m.is_display = true AND m.content ILIKE {add_param(like)}) OR EXISTS (SELECT 1 FROM session_messages_archive a WHERE a.session_id = sessions.id AND a.is_display = true AND a.content ILIKE {add_param(like)}))" - ) - frag = _special_frag(special) - if frag: - conditions.append(frag) - - where = "WHERE " + " AND ".join(conditions) if conditions else "" - - allowed_cols = {"created_at", "last_active", "name", "profile_id", "user_id", "pinned"} - col = sort_by if sort_by in allowed_cols else "last_active" - order = "DESC" if sort_order == "desc" else "ASC" - # secondary sort by pinned DESC for stability - order_clause = f"ORDER BY pinned DESC, {col} {order} LIMIT {add_param(limit)} OFFSET {add_param(offset)}" - - rows = await conn.fetch( - "SELECT id, profile_id, user_id, pinned, special, created_at, last_active, context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata " - f"FROM sessions {where} {order_clause}", - *params, - ) - return await _build_sessions(conn, rows) - - -async def pending_sweep_loop(store: "PgSessionStore", interval: int = 3600) -> None: - """Background task: drop abandoned pending sessions every `interval` seconds.""" - while True: - try: - await store.sweep_pending() - except asyncio.CancelledError: - raise - except Exception: - log = structlog.get_logger() - log.exception("session_store.pending_sweep_error") - await asyncio.sleep(interval) +__all__ = [ + "_DDL", + "_MIGRATE", + "_SESSION_ARCHIVE_DDL", + "_SESSION_IMAGES_DDL", + "_SESSION_MESSAGES_DDL", + "_SESSION_SHARES_DDL", + "_SYNAPSE_INDEX_DDL", + "PgSessionStore", + "_build_sessions", + "_deserialize", + "_ensure_normalized_tables", + "_load_messages_map", + "_message_key", + "_migrate_to_normalized", + "_row_to_message", + "_serialize", + "_special_and", + "_special_frag", + "_special_where", + "pending_sweep_loop", + "repair_dangling_tool_calls", +] diff --git a/navi/core/sessions/__init__.py b/navi/core/sessions/__init__.py new file mode 100644 index 0000000..4db86a9 --- /dev/null +++ b/navi/core/sessions/__init__.py @@ -0,0 +1,14 @@ +"""Session persistence — the PostgreSQL store, split by concern. + +* `ddl` — schema definitions and the migrations run on first use. +* `rows` — `Message` <-> row mapping and archive-aware session hydration. +* `store` — `PgSessionStore`; every public behavior lives here. + +`navi.core.pg_session_store` remains as a compatibility shim; new code should +import from here. +""" + +from .rows import repair_dangling_tool_calls +from .store import PgSessionStore, pending_sweep_loop + +__all__ = ["PgSessionStore", "pending_sweep_loop", "repair_dangling_tool_calls"] diff --git a/navi/core/sessions/ddl.py b/navi/core/sessions/ddl.py new file mode 100644 index 0000000..9d5043c --- /dev/null +++ b/navi/core/sessions/ddl.py @@ -0,0 +1,196 @@ +"""Schema DDL and the schema-time migrations (tables, columns, legacy JSON rows).""" + +import json + +import asyncpg +import structlog + +from .rows import _deserialize, _message_key + +_DDL = """ +CREATE TABLE IF NOT EXISTS sessions ( + id TEXT PRIMARY KEY, + profile_id TEXT NOT NULL, + user_id TEXT REFERENCES navi_users(id) ON DELETE SET NULL, + messages TEXT NOT NULL DEFAULT '[]', + context TEXT NOT NULL DEFAULT '', + pinned BOOLEAN NOT NULL DEFAULT FALSE, + special BOOLEAN NOT NULL DEFAULT FALSE, + created_at TIMESTAMPTZ NOT NULL, + last_active TIMESTAMPTZ NOT NULL, + context_token_count INTEGER NOT NULL DEFAULT 0, + name TEXT, + session_metadata TEXT NOT NULL DEFAULT '{}' +) +""" + +_MIGRATE = """ +ALTER TABLE sessions ADD COLUMN IF NOT EXISTS name TEXT; +ALTER TABLE sessions ADD COLUMN IF NOT EXISTS planning_logs TEXT NOT NULL DEFAULT '[]'; +ALTER TABLE sessions ADD COLUMN IF NOT EXISTS user_id TEXT REFERENCES navi_users(id) ON DELETE SET NULL; +ALTER TABLE sessions ADD COLUMN IF NOT EXISTS next_sequence INTEGER NOT NULL DEFAULT 0; +ALTER TABLE sessions ADD COLUMN IF NOT EXISTS archive_threshold INTEGER NOT NULL DEFAULT 0; +ALTER TABLE sessions ADD COLUMN IF NOT EXISTS session_metadata TEXT NOT NULL DEFAULT '{}'; +ALTER TABLE sessions ADD COLUMN IF NOT EXISTS special BOOLEAN NOT NULL DEFAULT FALSE +""" + +# Reaction threads are found by (owner, conversation key) among special sessions +# only — a tiny slice of the table, so the index is partial. The expression is the +# exact one the lookup uses; `session_metadata` is a TEXT column holding JSON. +_SYNAPSE_INDEX_DDL = """ +CREATE INDEX IF NOT EXISTS idx_sessions_synapse_session_key + ON sessions (user_id, (session_metadata::jsonb -> 'synapse' ->> 'session_key')) + WHERE special +""" + +_SESSION_MESSAGES_DDL = """ +CREATE TABLE IF NOT EXISTS session_messages ( + id SERIAL PRIMARY KEY, + session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, + sequence_number INT NOT NULL, + role TEXT NOT NULL, + content TEXT, + images TEXT, -- JSON array of base64 strings + tool_calls TEXT, -- JSON + tool_call_id TEXT, + name TEXT, + created_at TIMESTAMPTZ, + is_summary BOOLEAN NOT NULL DEFAULT FALSE, + thinking TEXT, + is_plan BOOLEAN NOT NULL DEFAULT FALSE, + is_compression BOOLEAN NOT NULL DEFAULT FALSE, + is_context BOOLEAN NOT NULL DEFAULT TRUE, + is_display BOOLEAN NOT NULL DEFAULT TRUE, + elapsed_seconds FLOAT, + tool_call_count INT, + token_count INT, + files TEXT, -- JSON + metadata TEXT, -- JSON + is_recall BOOLEAN NOT NULL DEFAULT FALSE, + UNIQUE(session_id, sequence_number) +); + +CREATE INDEX IF NOT EXISTS idx_session_messages_session_seq ON session_messages(session_id, sequence_number); +CREATE INDEX IF NOT EXISTS idx_session_messages_context ON session_messages(session_id, is_context, sequence_number); +""" + +_SESSION_IMAGES_DDL = """ +CREATE TABLE IF NOT EXISTS session_images ( + id SERIAL PRIMARY KEY, + session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, + message_id INT REFERENCES session_messages(id) ON DELETE CASCADE, + base64 TEXT NOT NULL, + filename TEXT +); + +CREATE INDEX IF NOT EXISTS idx_session_images_session ON session_images(session_id); +CREATE INDEX IF NOT EXISTS idx_session_images_message ON session_images(message_id); +""" + +_SESSION_ARCHIVE_DDL = """ +CREATE TABLE IF NOT EXISTS session_messages_archive ( + id SERIAL PRIMARY KEY, + session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, + sequence_number INT NOT NULL, + role TEXT NOT NULL, + content TEXT, + images TEXT, + tool_calls TEXT, + tool_call_id TEXT, + name TEXT, + created_at TIMESTAMPTZ, + is_summary BOOLEAN NOT NULL DEFAULT FALSE, + thinking TEXT, + is_plan BOOLEAN NOT NULL DEFAULT FALSE, + is_compression BOOLEAN NOT NULL DEFAULT FALSE, + is_context BOOLEAN NOT NULL DEFAULT TRUE, + is_display BOOLEAN NOT NULL DEFAULT TRUE, + elapsed_seconds FLOAT, + tool_call_count INT, + token_count INT, + files TEXT, + metadata TEXT, + is_recall BOOLEAN NOT NULL DEFAULT FALSE, + UNIQUE(session_id, sequence_number) +); +CREATE INDEX IF NOT EXISTS idx_session_messages_archive_session_seq ON session_messages_archive(session_id, sequence_number); +""" + +_SESSION_SHARES_DDL = """ +CREATE TABLE IF NOT EXISTS session_shares ( + session_id TEXT PRIMARY KEY REFERENCES sessions(id) ON DELETE CASCADE, + token TEXT NOT NULL UNIQUE, + owner_id TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() +); +""" + + + +async def _ensure_normalized_tables(conn: asyncpg.Connection) -> None: + await conn.execute(_SESSION_MESSAGES_DDL) + await conn.execute(_SESSION_IMAGES_DDL) + await conn.execute(_SESSION_ARCHIVE_DDL) + await conn.execute(_SESSION_SHARES_DDL) + # Best-effort: a row whose session_metadata is not valid JSON would make the + # expression index fail and take the whole schema init down with it. The + # lookup still works without the index, so a failure is a slowdown, not an + # outage — log it and carry on. + try: + await conn.execute(_SYNAPSE_INDEX_DDL) + except Exception as e: # noqa: BLE001 — never block startup over an index + structlog.get_logger().warning( + "pg_session_store.synapse_index_failed", error=str(e) + ) + + +async def _migrate_to_normalized(conn: asyncpg.Connection) -> None: + """One-shot migration: copy existing JSON messages/context into session_messages.""" + migrated = await conn.fetchval("SELECT COUNT(*) FROM session_messages") + if migrated: + return + + rows = await conn.fetch("SELECT id, messages, context FROM sessions") + for row in rows: + session_id = row["id"] + messages = _deserialize(row["messages"] or "[]") + context = _deserialize(row["context"] or "") + if not context and messages: + context = list(messages) + + context_set = {_message_key(m) for m in context} + + for seq, m in enumerate(messages): + is_ctx = _message_key(m) in context_set + await conn.execute( + """ + INSERT INTO session_messages + (session_id, sequence_number, role, content, images, tool_calls, tool_call_id, name, + created_at, is_summary, thinking, is_plan, is_compression, is_context, is_display, + elapsed_seconds, tool_call_count, token_count, files, metadata, is_recall) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21) + ON CONFLICT (session_id, sequence_number) DO NOTHING + """, + session_id, + seq, + m.role, + m.content, + json.dumps(m.images, ensure_ascii=False) if m.images else None, + json.dumps([tc.model_dump(mode="json") for tc in m.tool_calls], ensure_ascii=False) if m.tool_calls else None, + m.tool_call_id, + m.name, + m.created_at, + m.is_summary, + m.thinking, + m.is_plan, + m.is_compression, + is_ctx, + True, + m.elapsed_seconds, + m.tool_call_count, + m.token_count, + json.dumps(m.files, ensure_ascii=False) if m.files else None, + json.dumps(m.metadata, ensure_ascii=False) if m.metadata else None, + m.is_recall, + ) + diff --git a/navi/core/sessions/rows.py b/navi/core/sessions/rows.py new file mode 100644 index 0000000..8a3a33b --- /dev/null +++ b/navi/core/sessions/rows.py @@ -0,0 +1,200 @@ +"""Message rows: the `Message` <-> table-column mapping and session hydration. + +Covers the column mapping, the archive-aware loading of sessions, and the +one-shot repair of tool-call batches torn apart by a crash (or by a page +boundary). Nothing here touches the `sessions` row itself — `store` does that. +""" + +import json + +import asyncpg + +from navi.llm.base import Message, ToolCallRequest + +from ..session import Session + + +def _serialize(messages: list[Message]) -> str: + return json.dumps( + [m.model_dump(mode="json", exclude_none=True) for m in messages], + ensure_ascii=False, + ) + + +def _deserialize(raw: str) -> list[Message]: + if not raw: + return [] + return [Message.model_validate(m) for m in json.loads(raw)] + + +def repair_dangling_tool_calls(messages: list[Message]) -> int: + """Synthesize placeholder results for assistant tool_calls that lost their + role=tool pair (a crash mid-batch, e.g. a parallel batch killed by a server + restart). Returns the number of placeholders added. + + Placeholders are inserted right after the last tool message of the affected + batch (before the next non-tool message) and marked ``is_context=False`` — + display history stays coherent, the LLM context is untouched. + """ + answered = { + m.tool_call_id for m in messages if m.role == "tool" and m.tool_call_id + } + out: list[Message] = [] + pending: list[Message] = [] + for m in messages: + if m.role != "tool" and pending: + out.extend(pending) + pending = [] + out.append(m) + if m.role == "assistant" and m.tool_calls: + for tc in m.tool_calls: + if not tc.id or tc.id in answered: + continue + pending.append( + Message( + role="tool", + content="[Interrupted before result]", + tool_call_id=tc.id, + name=tc.name, + metadata={}, + is_context=False, + ) + ) + out.extend(pending) + if len(out) != len(messages): + added = len(out) - len(messages) + messages[:] = out + return added + return 0 + + +def _message_key(m: Message) -> tuple: + """Stable key for matching a message between messages[] and context[]. + + Used only by the one-shot boot migration for legacy JSON rows. + """ + return ( + m.role, + m.content, + m.tool_call_id, + m.name, + m.is_summary, + m.is_plan, + m.is_compression, + m.is_recall, + m.thinking, + m.created_at.isoformat() if m.created_at else None, + json.dumps([tc.model_dump(mode="json") for tc in m.tool_calls], ensure_ascii=False) if m.tool_calls else None, + json.dumps(m.files, ensure_ascii=False) if m.files else None, + json.dumps(m.metadata, ensure_ascii=False) if m.metadata else None, + ) + + +def _row_to_message(row: asyncpg.Record) -> Message: + images = json.loads(row["images"]) if row.get("images") else None + tool_calls = None + raw_tool_calls = row.get("tool_calls") + if raw_tool_calls: + tool_calls = [ToolCallRequest.model_validate(tc) for tc in json.loads(raw_tool_calls)] + files = json.loads(row["files"]) if row.get("files") else None + metadata = json.loads(row["metadata"]) if row.get("metadata") else {} + + return Message( + role=row["role"], + content=row["content"], + images=images, + tool_calls=tool_calls, + tool_call_id=row.get("tool_call_id"), + name=row.get("name"), + created_at=row.get("created_at"), + is_summary=bool(row.get("is_summary", False)), + thinking=row.get("thinking"), + is_plan=bool(row.get("is_plan", False)), + is_compression=bool(row.get("is_compression", False)), + is_context=bool(row.get("is_context", True)), + is_display=bool(row.get("is_display", True)), + elapsed_seconds=row.get("elapsed_seconds"), + tool_call_count=row.get("tool_call_count"), + token_count=row.get("token_count"), + files=files, + metadata=metadata, + is_recall=bool(row.get("is_recall", False)), + sequence_number=row.get("sequence_number", 0), + ) + + +async def _load_messages_map(conn: asyncpg.Connection, session_ids: list[str]) -> dict[str, list[Message]]: + """Batch-load all messages for the given session IDs.""" + if not session_ids: + return {} + rows = await conn.fetch( + "SELECT * FROM session_messages WHERE session_id = ANY($1) ORDER BY sequence_number", + session_ids, + ) + result: dict[str, list[Message]] = {sid: [] for sid in session_ids} + for row in rows: + sid = row["session_id"] + result[sid].append(_row_to_message(row)) + return result + + +async def _build_sessions( + conn: asyncpg.Connection, + rows: list[asyncpg.Record], +) -> list[Session]: + """Hydrate session rows with hot (non-archived) messages.""" + session_ids = [r["id"] for r in rows] + messages_map = await _load_messages_map(conn, session_ids) + sessions: list[Session] = [] + for row in rows: + all_msgs = messages_map.get(row["id"], []) + archive_threshold = row.get("archive_threshold", 0) or 0 + hot_msgs = [m for m in all_msgs if m.sequence_number >= archive_threshold] + messages = [m for m in hot_msgs if m.is_display] + context = [m for m in hot_msgs if m.is_context] + planning_logs_raw = row["planning_logs"] + planning_logs = json.loads(planning_logs_raw) if planning_logs_raw else [] + + next_seq = row.get("next_sequence", 0) or 0 + max_seq = max((m.sequence_number for m in hot_msgs), default=-1) + if next_seq == 0: + next_seq = max_seq + 1 + + s = Session( + id=row["id"], + profile_id=row["profile_id"], + user_id=row["user_id"], + messages=messages, + context=context, + pinned=bool(row["pinned"]), + special=bool(row.get("special", False)), + name=row["name"], + created_at=row["created_at"], + last_active=row["last_active"], + context_token_count=row["context_token_count"] or 0, + planning_logs=planning_logs, + ) + session_metadata = json.loads(row.get("session_metadata") or "{}") if row.get("session_metadata") else {} + s.db_message_count = len(hot_msgs) + s.db_next_sequence = next_seq + s.archive_threshold = archive_threshold + s.session_metadata = session_metadata + sessions.append(s) + return sessions + + +def _special_frag(special: bool | None) -> str: + """SQL fragment for the special filter; '' when unfiltered.""" + if special is None: + return "" + return f"special = {str(bool(special)).upper()}" + + +def _special_and(special: bool | None) -> str: + frag = _special_frag(special) + return f" AND {frag}" if frag else "" + + +def _special_where(special: bool | None) -> str: + frag = _special_frag(special) + return f" WHERE {frag}" if frag else "" diff --git a/navi/core/sessions/store.py b/navi/core/sessions/store.py new file mode 100644 index 0000000..a119e09 --- /dev/null +++ b/navi/core/sessions/store.py @@ -0,0 +1,788 @@ +"""`PgSessionStore` — the PostgreSQL-backed implementation of `SessionStore`.""" + +import asyncio +import json +import secrets +from datetime import datetime, timezone + +import asyncpg +import structlog + +from navi.llm.base import Message + +from ..session import Session, SessionStore +from .ddl import _DDL, _MIGRATE, _ensure_normalized_tables, _migrate_to_normalized +from .rows import ( + _build_sessions, + _row_to_message, + _special_and, + _special_frag, + _special_where, + repair_dangling_tool_calls, +) + + +class PgSessionStore(SessionStore): + def __init__(self, pool: asyncpg.Pool) -> None: + self._pool = pool + self._initialized = False + self._lock = asyncio.Lock() + # Sessions created via create() but not yet persisted to the DB. + # They live only in memory until the first save() upserts the row + # (on the first user message / meaningful state change), so empty + # sessions that never receive a message never reach the DB. + self._pending: dict[str, Session] = {} + self._pending_lock = asyncio.Lock() + + async def _get_pool(self) -> asyncpg.Pool: + if not self._initialized: + async with self._lock: + if not self._initialized: + async with self._pool.acquire() as conn: + await conn.execute(_DDL) + await conn.execute(_MIGRATE) + await _ensure_normalized_tables(conn) + await _migrate_to_normalized(conn) + self._initialized = True + return self._pool + + async def create(self, profile_id: str, user_id: str | None = None) -> Session: + # Lazy persistence: do NOT insert a DB row here. Register the session + # in memory; the row is upserted on the first save(). This keeps empty + # sessions (created but never sent a message) out of the DB entirely. + session = Session(profile_id=profile_id, user_id=user_id) + async with self._pending_lock: + self._pending[session.id] = session + return session + + async def get(self, session_id: str) -> Session | None: + # Pending (not-yet-persisted) sessions live only in memory. + async with self._pending_lock: + pending = self._pending.get(session_id) + if pending is not None: + return pending + + pool = await self._get_pool() + async with pool.acquire() as conn: + row = await conn.fetchrow( + "SELECT id, profile_id, user_id, pinned, special, created_at, last_active, context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata " + "FROM sessions WHERE id = $1", + session_id, + ) + if not row: + return None + + archive_threshold = row["archive_threshold"] or 0 + # Load hot (non-archived) messages so messages[] and context[] share + # the same Python objects (id() matching in the agent works). + all_rows = await conn.fetch( + "SELECT * FROM session_messages WHERE session_id = $1 AND sequence_number >= $2 ORDER BY sequence_number", + session_id, + archive_threshold, + ) + all_messages = [_row_to_message(r) for r in all_rows] + # A crash mid-tool-batch (e.g. parallel batch) can leave assistant + # tool_calls without their role=tool pair — patch display history. + repair_dangling_tool_calls(all_messages) + messages = [m for m in all_messages if m.is_display] + context = [m for m in all_messages if m.is_context] + + planning_logs_raw = row["planning_logs"] + planning_logs = json.loads(planning_logs_raw) if planning_logs_raw else [] + + next_seq = row["next_sequence"] or 0 + max_seq = max((m.sequence_number for m in all_messages), default=-1) + if next_seq == 0: + next_seq = max_seq + 1 + + session_metadata = json.loads(row.get("session_metadata") or "{}") if row.get("session_metadata") else {} + s = Session( + id=row["id"], + profile_id=row["profile_id"], + user_id=row["user_id"], + messages=messages, + context=context, + pinned=bool(row["pinned"]), + special=bool(row.get("special", False)), + name=row["name"], + created_at=row["created_at"], + last_active=row["last_active"], + context_token_count=row["context_token_count"] or 0, + planning_logs=planning_logs, + ) + s.db_message_count = len(all_messages) + s.db_next_sequence = next_seq + s.archive_threshold = archive_threshold + s.session_metadata = session_metadata + return s + + async def save(self, session: Session) -> None: + session.last_active = datetime.now(timezone.utc) + pool = await self._get_pool() + # One transaction: the sequence range handed out below and the rows that + # consume it must commit together, or a crash in between burns numbers + # and leaves the counter ahead of the messages. + async with pool.acquire() as conn, conn.transaction(): + # Upsert the session row. On the first save() after a lazy create() + # the row does not exist yet — INSERT ... ON CONFLICT creates it so + # the session_messages FK is satisfied before messages are inserted. + # created_at / pinned / name / next_sequence / archive_threshold are + # set on insert and not overwritten on conflict (managed elsewhere). + # profile_id is on that list too, for the same reason: set_profile() + # owns the column after creation, while a run's Session can still + # carry the profile it started with — writing it here would silently + # undo a switch_profile made mid-run (agent.py then reloads the row + # it just clobbered and sees no change at all). + await conn.execute( + "INSERT INTO sessions " + "(id, profile_id, user_id, pinned, special, created_at, last_active, " + " context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata) " + "VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) " + "ON CONFLICT (id) DO UPDATE SET " + " user_id = EXCLUDED.user_id, " + " last_active = EXCLUDED.last_active, " + " context_token_count = EXCLUDED.context_token_count, " + " planning_logs = EXCLUDED.planning_logs, " + " session_metadata = EXCLUDED.session_metadata", + session.id, session.profile_id, session.user_id, session.pinned, + session.special, session.created_at, session.last_active, + session.context_token_count, + session.name, + json.dumps(session.planning_logs, ensure_ascii=False), + session.db_next_sequence, session.archive_threshold, + json.dumps(session.session_metadata, ensure_ascii=False), + ) + # Row now guaranteed to exist — drop the pending entry if any. + async with self._pending_lock: + self._pending.pop(session.id, None) + + messages = session.messages + + # 1. Update mutable flags for already-persisted rows (sequence_number >= 0) + existing = [m for m in messages if m.sequence_number >= 0] + if existing: + update_rows = [] + for m in existing: + update_rows.append( + ( + session.id, + m.sequence_number, + m.is_context, + m.is_display, + m.is_summary, + m.is_plan, + m.is_compression, + m.is_recall, + m.thinking, + m.elapsed_seconds, + m.tool_call_count, + m.token_count, + ) + ) + await conn.executemany( + """ + UPDATE session_messages + SET is_context = $3, + is_display = $4, + is_summary = $5, + is_plan = $6, + is_compression = $7, + is_recall = $8, + thinking = $9, + elapsed_seconds = $10, + tool_call_count = $11, + token_count = $12 + WHERE session_id = $1 AND sequence_number = $2 + """, + update_rows, + ) + + # 2. Insert new messages (sequence_number < 0 means "not yet persisted") + new_msgs = [m for m in messages if m.sequence_number < 0] + if new_msgs: + # Claim the range in the DB, never from session.db_next_sequence: + # another writer of the same session may have consumed those + # numbers since this object was loaded. switch_profile used to do + # exactly that — it loaded the session mid-run and saved it back — + # and the running turn then inserted onto numbers already taken, + # dying on UNIQUE(session_id, sequence_number). GREATEST() also + # heals rows whose counter is still 0 (pre-counter sessions). + base = await conn.fetchval( + """ + UPDATE sessions s + SET next_sequence = GREATEST( + s.next_sequence, + COALESCE((SELECT MAX(m.sequence_number) + 1 + FROM session_messages m + WHERE m.session_id = s.id), 0) + ) + $2 + WHERE s.id = $1 + RETURNING s.next_sequence - $2 + """, + session.id, + len(new_msgs), + ) + insert_rows = [] + for i, m in enumerate(new_msgs): + seq = base + i + m.sequence_number = seq + insert_rows.append( + ( + session.id, + seq, + m.role, + m.content, + json.dumps(m.images, ensure_ascii=False) if m.images else None, + json.dumps([tc.model_dump(mode="json") for tc in m.tool_calls], ensure_ascii=False) if m.tool_calls else None, + m.tool_call_id, + m.name, + m.created_at, + m.is_summary, + m.thinking, + m.is_plan, + m.is_compression, + m.is_context, + m.is_display, + m.elapsed_seconds, + m.tool_call_count, + m.token_count, + json.dumps(m.files, ensure_ascii=False) if m.files else None, + json.dumps(m.metadata, ensure_ascii=False) if m.metadata else None, + m.is_recall, + ) + ) + await conn.executemany( + """ + INSERT INTO session_messages + (session_id, sequence_number, role, content, images, tool_calls, tool_call_id, name, + created_at, is_summary, thinking, is_plan, is_compression, is_context, is_display, + elapsed_seconds, tool_call_count, token_count, files, metadata, is_recall) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21) + """, + insert_rows, + ) + session.db_next_sequence = base + len(new_msgs) + + session.db_message_count = len(messages) + + async def sweep_pending(self, max_age_seconds: int = 3600) -> int: + """Drop pending (not-yet-persisted) sessions older than max_age_seconds. + + Bounds memory used by clients that create a session and never send a + message. Returns the number of entries removed. + """ + cutoff = datetime.now(timezone.utc).timestamp() - max_age_seconds + removed = 0 + async with self._pending_lock: + for sid in list(self._pending): + s = self._pending[sid] + if s.created_at.timestamp() < cutoff: + del self._pending[sid] + removed += 1 + return removed + + async def set_context_flag( + self, session_id: str, sequence_numbers: list[int], is_context: bool + ) -> int: + """Persist a context-flag flip for rows save() cannot see. + + ``save()`` updates flags only for rows in ``session.messages``. A + summary is ``is_display=False``: after a reload it lives in + ``session.context`` alone, so compression's demotion of a superseded + summary never reached the database — every reload resurrected it, and + prod sessions accumulated up to 15 live summaries (80% of the context). + """ + if not sequence_numbers: + return 0 + pool = await self._get_pool() + async with pool.acquire() as conn: + result = await conn.execute( + "UPDATE session_messages SET is_context = $3 " + "WHERE session_id = $1 AND sequence_number = ANY($2::int[])", + session_id, + list(sequence_numbers), + is_context, + ) + parts = result.split() + return int(parts[-1]) if len(parts) >= 2 else 0 + + async def archive_old_messages(self, session_id: str, keep_seq_threshold: int) -> int: + """Move messages older than keep_seq_threshold from hot to archive table. + + Returns number of rows archived. + """ + pool = await self._get_pool() + async with pool.acquire() as conn: + # Copy old rows to archive + copied = await conn.execute( + """ + INSERT INTO session_messages_archive + SELECT * FROM session_messages + WHERE session_id = $1 AND sequence_number < $2 + ON CONFLICT (session_id, sequence_number) DO NOTHING + """, + session_id, keep_seq_threshold, + ) + # Delete from hot table + await conn.execute( + "DELETE FROM session_messages WHERE session_id = $1 AND sequence_number < $2", + session_id, keep_seq_threshold, + ) + # Update threshold on session + await conn.execute( + "UPDATE sessions SET archive_threshold = $1 WHERE id = $2", + keep_seq_threshold, session_id, + ) + # asyncpg execute returns 'INSERT 0 N' — extract N + parts = copied.split() + return int(parts[-1]) if len(parts) >= 2 else 0 + + async def get_archived_messages( + self, session_id: str, before_seq: int | None = None, limit: int = 50 + ) -> list[Message]: + """Return archived messages for a session, newest first.""" + pool = await self._get_pool() + async with pool.acquire() as conn: + if before_seq is not None: + rows = await conn.fetch( + "SELECT * FROM session_messages_archive " + "WHERE session_id = $1 AND sequence_number < $2 " + "ORDER BY sequence_number DESC LIMIT $3", + session_id, before_seq, limit, + ) + else: + rows = await conn.fetch( + "SELECT * FROM session_messages_archive " + "WHERE session_id = $1 " + "ORDER BY sequence_number DESC LIMIT $2", + session_id, limit, + ) + # Return oldest-first so the caller can prepend them in order + return [_row_to_message(r) for r in reversed(rows)] + + async def get_meta(self, session_id: str) -> Session | None: + """Session row without any message payload — cheap head for the UI. + + Unlike get() this never reads session_messages; message counts come + from a single aggregate. For pending (not-yet-persisted) sessions it + returns the in-memory object like get() does. + """ + async with self._pending_lock: + pending = self._pending.get(session_id) + if pending is not None: + return pending + + pool = await self._get_pool() + async with pool.acquire() as conn: + row = await conn.fetchrow( + "SELECT id, profile_id, user_id, pinned, special, created_at, last_active, context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata " + "FROM sessions WHERE id = $1", + session_id, + ) + if not row: + return None + msg_count = await conn.fetchval( + "SELECT COUNT(*) FROM session_messages WHERE session_id = $1", + session_id, + ) + planning_logs_raw = row["planning_logs"] + planning_logs = json.loads(planning_logs_raw) if planning_logs_raw else [] + metadata_raw = row.get("session_metadata") + s = Session( + id=row["id"], + profile_id=row["profile_id"], + user_id=row["user_id"], + messages=[], + context=[], + pinned=bool(row["pinned"]), + special=bool(row.get("special", False)), + name=row["name"], + created_at=row["created_at"], + last_active=row["last_active"], + context_token_count=row["context_token_count"] or 0, + planning_logs=planning_logs, + ) + s.db_message_count = int(msg_count or 0) + next_seq = row["next_sequence"] or (s.db_message_count + 1) + s.db_next_sequence = int(next_seq) + s.archive_threshold = row["archive_threshold"] or 0 + s.session_metadata = json.loads(metadata_raw or "{}") if metadata_raw else {} + return s + + async def get_history_page( + self, session_id: str, before_seq: int | None = None, limit: int = 150 + ) -> tuple[list[Message], bool]: + """One page of full history (archive ∪ hot), display messages only. + + Returns (messages oldest-first, has_more). With before_seq=None the + page starts at the newest message — the tail the UI loads initially; + with before_seq set it continues below that sequence number (scroll-up + pagination). UNION (not ALL) collapses rows doubled across hot/archive + by a crash between the archive copy and delete. + + Each message carries ``metadata.display_index`` — its 0-based rank in + the full display history — so paged clients build ids that match a full + reload (feedback keys and message jumps depend on that stability). + """ + pool = await self._get_pool() + async with pool.acquire() as conn: + rows = await conn.fetch( + """ + WITH hist AS ( + SELECT * FROM session_messages WHERE session_id = $1 AND is_display + UNION + SELECT * FROM session_messages_archive WHERE session_id = $1 AND is_display + ), ranked AS ( + SELECT *, ROW_NUMBER() OVER (ORDER BY sequence_number) - 1 AS display_index + FROM hist + ) + SELECT * FROM ranked + WHERE ($2::int IS NULL OR sequence_number < $2) + ORDER BY sequence_number DESC + LIMIT $3 + """, + session_id, before_seq, limit + 1, + ) + has_more = len(rows) > limit + # SQL rank of each row: 0-based position in the full display history. + rank_by_seq = {r["sequence_number"]: r["display_index"] for r in rows[:limit]} + msgs = [_row_to_message(r) for r in rows[:limit]] + msgs.reverse() # oldest-first + # A page boundary can split a tool batch — patch the fragment like get() does. + repair_dangling_tool_calls(msgs) + # Synthetic placeholders from the repair have no sequence number and no + # rank — give them the next rank after the last real message so client + # ids stay collision-free. + last_known: int | None = None + for m in msgs: + di = rank_by_seq.get(m.sequence_number) if m.sequence_number >= 0 else None + if di is None and last_known is not None: + di = last_known + 1 + if di is not None: + m.metadata = dict(m.metadata or {}) | {"display_index": di} + last_known = di + return msgs, has_more + + async def set_pinned(self, session_id: str, pinned: bool) -> bool: + pool = await self._get_pool() + async with pool.acquire() as conn: + result = await conn.execute( + "UPDATE sessions SET pinned = $1 WHERE id = $2", + pinned, session_id, + ) + return result == "UPDATE 1" + + async def set_name(self, session_id: str, name: str) -> bool: + # A never-persisted session lives only in _pending (lazy persistence): + # the bare UPDATE would touch 0 rows and silently drop the name. + async with self._pending_lock: + pending = self._pending.get(session_id) + if pending is not None: + pending.name = name + return True + pool = await self._get_pool() + async with pool.acquire() as conn: + result = await conn.execute( + "UPDATE sessions SET name = $1 WHERE id = $2", + name, session_id, + ) + return result == "UPDATE 1" + + async def set_profile(self, session_id: str, profile_id: str) -> bool: + """Repoint a session at another profile — one narrow UPDATE, no save(). + + switch_profile runs *inside* a live agent run that holds its own Session + object. Loading the session and saving it back (as it once did) makes the + tool a second writer of the same session, and the run's next insert then + lands on a sequence number this tool already took. Touching only the + profile column keeps the two out of each other's way; the run picks the + change up through its post-switch profile reload. + """ + # A never-persisted session lives only in _pending (lazy persistence): + # the bare UPDATE would touch 0 rows and silently drop the switch. + async with self._pending_lock: + pending = self._pending.get(session_id) + if pending is not None: + pending.profile_id = profile_id + return True + pool = await self._get_pool() + async with pool.acquire() as conn: + result = await conn.execute( + "UPDATE sessions SET profile_id = $1 WHERE id = $2", + profile_id, session_id, + ) + return result == "UPDATE 1" + + async def find_reaction_session( + self, + *, + user_id: str, + session_key: str, + not_active_before: datetime | None = None, + ) -> Session | None: + """The newest special session of `user_id` carrying this conversation key. + + One indexed lookup, then `get()` for the hydration — rather than repeating + the message/context loading here. Sessions still in `_pending` are skipped + on purpose: the reaction handler saves the session (and therefore persists + the key) before it starts the run, so a key that is not in the table yet + belongs to a session whose reaction has not begun. + """ + pool = await self._get_pool() + async with pool.acquire() as conn: + row = await conn.fetchrow( + """ + SELECT id FROM sessions + WHERE special + AND user_id = $1 + AND session_metadata::jsonb -> 'synapse' ->> 'session_key' = $2 + AND ($3::timestamptz IS NULL OR last_active >= $3) + ORDER BY last_active DESC + LIMIT 1 + """, + user_id, session_key, not_active_before, + ) + if row is None: + return None + return await self.get(row["id"]) + + async def list_all( + self, user_id: str | None = None, is_admin: bool = False, special: bool | None = None + ) -> list[Session]: + """Sessions visible to the caller. + + Scoping: admins see everything; a named owner sees their own rows; + ``user_id=None`` (local mode, the anonymous local user) sees only + NULL-owner sessions — never other users' chats. + ``special``: None — no filter, False — regular only, True — special only. + """ + pool = await self._get_pool() + async with pool.acquire() as conn: + if is_admin: + rows = await conn.fetch( + "SELECT id, profile_id, user_id, pinned, special, created_at, last_active, context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata " + f"FROM sessions{_special_where(special)} ORDER BY pinned DESC, last_active DESC" + ) + elif user_id is not None: + rows = await conn.fetch( + "SELECT id, profile_id, user_id, pinned, special, created_at, last_active, context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata " + f"FROM sessions WHERE user_id = $1{_special_and(special)} ORDER BY pinned DESC, last_active DESC", + user_id, + ) + else: + rows = await conn.fetch( + "SELECT id, profile_id, user_id, pinned, special, created_at, last_active, context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata " + f"FROM sessions WHERE user_id IS NULL{_special_and(special)} ORDER BY pinned DESC, last_active DESC" + ) + return await _build_sessions(conn, rows) + + async def list_page( + self, + *, + limit: int, + offset: int, + profile_id: str | None = None, + user_id: str | None = None, + is_admin: bool = False, + special: bool | None = None, + ) -> list[Session]: + pool = await self._get_pool() + async with pool.acquire() as conn: + conditions = [] + params: list = [] + param_idx = 0 + + def add_param(value): + nonlocal param_idx + param_idx += 1 + params.append(value) + return f"${param_idx}" + + if not is_admin and user_id is None: + # Local mode: the anonymous local user owns NULL sessions only. + conditions.append("user_id IS NULL") + elif not is_admin: + conditions.append(f"user_id = {add_param(user_id)}") + if profile_id: + conditions.append(f"profile_id = {add_param(profile_id)}") + frag = _special_frag(special) + if frag: + conditions.append(frag) + + where = "WHERE " + " AND ".join(conditions) if conditions else "" + order_limit = f"ORDER BY pinned DESC, last_active DESC LIMIT {add_param(limit)} OFFSET {add_param(offset)}" + + rows = await conn.fetch( + "SELECT id, profile_id, user_id, pinned, special, created_at, last_active, context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata " + f"FROM sessions {where} {order_limit}", + *params, + ) + return await _build_sessions(conn, rows) + + async def delete(self, session_id: str) -> bool: + # A never-persisted session lives only in _pending (lazy persistence: + # create() defers the row until the first save()), so the bare DB + # DELETE alone reports 0 rows for a very real session. Pop the + # pending entry first, then count the DB result. + async with self._pending_lock: + pending = self._pending.pop(session_id, None) + pool = await self._get_pool() + async with pool.acquire() as conn: + result = await conn.execute("DELETE FROM sessions WHERE id = $1", session_id) + # Raw chunks have no FK on purpose (memory DDL may run before the + # sessions table exists), so deleting a chat must clean them here — + # otherwise a deleted conversation stays searchable. Best-effort: + # the session row is already gone, and a failure here must not turn + # a successful delete into a 500 (or leave the caller thinking the + # delete did not happen). + try: + await conn.execute( + "DELETE FROM memory_chunks WHERE session_id = $1", session_id + ) + except asyncpg.UndefinedTableError: + pass + except Exception: + structlog.get_logger().warning( + "session.chunk_cleanup_failed", + session_id=session_id, + exc_info=True, + ) + return pending is not None or result == "DELETE 1" + + async def get_or_create_share(self, session_id: str, owner_id: str | None) -> dict | None: + pool = await self._get_pool() + async with pool.acquire() as conn: + try: + row = await conn.fetchrow( + """ + INSERT INTO session_shares (session_id, token, owner_id) + VALUES ($1, $2, $3) + ON CONFLICT (session_id) DO NOTHING + RETURNING token, session_id, owner_id, created_at + """, + session_id, secrets.token_urlsafe(32), owner_id, + ) + except asyncpg.ForeignKeyViolationError: + # A pending session has no row yet — there is nothing to share. + return None + if row is None: # already shared: hand back the existing link + row = await conn.fetchrow( + "SELECT token, session_id, owner_id, created_at " + "FROM session_shares WHERE session_id = $1", + session_id, + ) + return dict(row) if row is not None else None + + async def get_share(self, token: str) -> dict | None: + pool = await self._get_pool() + async with pool.acquire() as conn: + row = await conn.fetchrow( + "SELECT token, session_id, owner_id, created_at " + "FROM session_shares WHERE token = $1", + token, + ) + return dict(row) if row is not None else None + + async def count_all( + self, + *, + user_id: str | None = None, + is_admin: bool = False, + search: str | None = None, + special: bool | None = None, + ) -> int: + pool = await self._get_pool() + async with pool.acquire() as conn: + conditions = [] + params: list = [] + param_idx = 0 + + def add_param(value): + nonlocal param_idx + param_idx += 1 + params.append(value) + return f"${param_idx}" + + if not is_admin and user_id is None: + # Local mode: the anonymous local user owns NULL sessions only. + conditions.append("user_id IS NULL") + elif not is_admin: + conditions.append(f"user_id = {add_param(user_id)}") + if search: + like = f"%{search}%" + conditions.append( + f"(id ILIKE {add_param(like)} OR name ILIKE {add_param(like)} OR user_id ILIKE {add_param(like)} OR profile_id ILIKE {add_param(like)} OR EXISTS (SELECT 1 FROM session_messages m WHERE m.session_id = sessions.id AND m.is_display = true AND m.content ILIKE {add_param(like)}) OR EXISTS (SELECT 1 FROM session_messages_archive a WHERE a.session_id = sessions.id AND a.is_display = true AND a.content ILIKE {add_param(like)}))" + ) + frag = _special_frag(special) + if frag: + conditions.append(frag) + + where = "WHERE " + " AND ".join(conditions) if conditions else "" + row = await conn.fetchrow(f"SELECT COUNT(*) FROM sessions {where}", *params) + return row["count"] if row else 0 + + async def search_list( + self, + *, + limit: int, + offset: int, + user_id: str | None = None, + is_admin: bool = False, + search: str | None = None, + sort_by: str = "last_active", + sort_order: str = "desc", + special: bool | None = None, + ) -> list[Session]: + pool = await self._get_pool() + async with pool.acquire() as conn: + conditions = [] + params: list = [] + param_idx = 0 + + def add_param(value): + nonlocal param_idx + param_idx += 1 + params.append(value) + return f"${param_idx}" + + if not is_admin and user_id is None: + # Local mode: the anonymous local user owns NULL sessions only. + conditions.append("user_id IS NULL") + elif not is_admin: + conditions.append(f"user_id = {add_param(user_id)}") + if search: + like = f"%{search}%" + conditions.append( + f"(id ILIKE {add_param(like)} OR name ILIKE {add_param(like)} OR user_id ILIKE {add_param(like)} OR profile_id ILIKE {add_param(like)} OR EXISTS (SELECT 1 FROM session_messages m WHERE m.session_id = sessions.id AND m.is_display = true AND m.content ILIKE {add_param(like)}) OR EXISTS (SELECT 1 FROM session_messages_archive a WHERE a.session_id = sessions.id AND a.is_display = true AND a.content ILIKE {add_param(like)}))" + ) + frag = _special_frag(special) + if frag: + conditions.append(frag) + + where = "WHERE " + " AND ".join(conditions) if conditions else "" + + allowed_cols = {"created_at", "last_active", "name", "profile_id", "user_id", "pinned"} + col = sort_by if sort_by in allowed_cols else "last_active" + order = "DESC" if sort_order == "desc" else "ASC" + # secondary sort by pinned DESC for stability + order_clause = f"ORDER BY pinned DESC, {col} {order} LIMIT {add_param(limit)} OFFSET {add_param(offset)}" + + rows = await conn.fetch( + "SELECT id, profile_id, user_id, pinned, special, created_at, last_active, context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata " + f"FROM sessions {where} {order_clause}", + *params, + ) + return await _build_sessions(conn, rows) + + +async def pending_sweep_loop(store: "PgSessionStore", interval: int = 3600) -> None: + """Background task: drop abandoned pending sessions every `interval` seconds.""" + while True: + try: + await store.sweep_pending() + except asyncio.CancelledError: + raise + except Exception: + log = structlog.get_logger() + log.exception("session_store.pending_sweep_error") + await asyncio.sleep(interval)