diff --git a/navi/core/pg_session_store.py b/navi/core/pg_session_store.py index 479331f..9df770c 100644 --- a/navi/core/pg_session_store.py +++ b/navi/core/pg_session_store.py @@ -18,6 +18,7 @@ from datetime import datetime, timezone import asyncpg +import structlog from navi.llm.base import Message, ToolCallRequest @@ -311,6 +312,12 @@ 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: @@ -325,18 +332,21 @@ 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) - pool = await self._get_pool() - async with pool.acquire() as conn: - await conn.execute( - "INSERT INTO sessions " - "(id, profile_id, user_id, pinned, created_at, last_active, context_token_count) " - "VALUES ($1, $2, $3, FALSE, $4, $5, 0)", - session.id, session.profile_id, session.user_id, session.created_at, session.last_active, - ) + 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( @@ -391,14 +401,33 @@ session.last_active = datetime.now(timezone.utc) pool = await self._get_pool() async with pool.acquire() as conn: + # 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). await conn.execute( - "UPDATE sessions SET profile_id = $1, user_id = $2, " - "last_active = $3, context_token_count = $4, planning_logs = $5, session_metadata = $6 WHERE id = $7", - session.profile_id, session.user_id, - session.last_active, session.context_token_count, + "INSERT INTO sessions " + "(id, profile_id, user_id, pinned, 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) " + "ON CONFLICT (id) DO UPDATE SET " + " profile_id = EXCLUDED.profile_id, " + " 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.created_at, session.last_active, session.context_token_count, + session.name, json.dumps(session.planning_logs, ensure_ascii=False), - json.dumps(session.session_metadata, ensure_ascii=False), session.id, + 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) db_next = session.db_next_sequence messages = session.messages @@ -493,6 +522,22 @@ 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 archive_old_messages(self, session_id: str, keep_seq_threshold: int) -> int: """Move messages older than keep_seq_threshold from hot to archive table. @@ -699,3 +744,16 @@ *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) diff --git a/navi/main.py b/navi/main.py index f0b2d4e..7d6093c 100644 --- a/navi/main.py +++ b/navi/main.py @@ -53,6 +53,7 @@ from navi.profiles._overrides import ensure_table, load_overrides from navi.api.routes.health import _check_embed from navi.core.scheduler import recall_scheduler_loop + from navi.core.pg_session_store import pending_sweep_loop # Ensure auth tables first (navi_users is referenced by other DDL). for attempt in range(1, 6): @@ -95,6 +96,7 @@ scheduler_task = asyncio.create_task( recall_scheduler_loop(container.scheduler, container.session_store, container.orchestrator) ) + pending_sweep_task = asyncio.create_task(pending_sweep_loop(container.session_store)) yield @@ -109,6 +111,11 @@ await cleanup_task except asyncio.CancelledError: pass + pending_sweep_task.cancel() + try: + await pending_sweep_task + except asyncio.CancelledError: + pass from navi.tools.ssh_exec import close_all_connections close_all_connections() diff --git a/tests/unit/core/test_pg_session_store.py b/tests/unit/core/test_pg_session_store.py index 0072d41..3668431 100644 --- a/tests/unit/core/test_pg_session_store.py +++ b/tests/unit/core/test_pg_session_store.py @@ -11,7 +11,7 @@ @pytest.mark.asyncio async def test_save_assigns_sequence_numbers_to_new_messages(): conn = FakeConnection() - conn.enqueue("OK") # UPDATE sessions + conn.enqueue("OK") # upsert sessions row (INSERT ... ON CONFLICT) conn.enqueue(None) # executemany UPDATE (empty) conn.enqueue("OK") # INSERT executemany conn.enqueue("OK") # UPDATE next_sequence @@ -38,7 +38,7 @@ @pytest.mark.asyncio async def test_save_updates_existing_rows_by_sequence_number(): conn = FakeConnection() - conn.enqueue("OK") # UPDATE sessions + conn.enqueue("OK") # upsert sessions row (INSERT ... ON CONFLICT) conn.enqueue(None) # executemany UPDATE conn.enqueue("OK") # INSERT (empty) conn.enqueue("OK") # UPDATE next_sequence @@ -79,3 +79,125 @@ execute_calls = [c for c in conn.calls if c[0] == "execute"] assert any("DELETE FROM session_messages" in c[1] for c in execute_calls) assert any("UPDATE sessions SET archive_threshold" in c[1] for c in execute_calls) + + +@pytest.mark.asyncio +async def test_create_registers_pending_without_db(): + """create() must NOT touch the DB — the session lives only in _pending.""" + conn = FakeConnection() + pool = FakePool(conn) + store = PgSessionStore(pool) + store._initialized = True + + session = await store.create(profile_id="test", user_id="u1") + + assert session.profile_id == "test" + assert session.user_id == "u1" + # Registered as pending, no DB row yet. + assert session.id in store._pending + assert store._pending[session.id] is session + # No execute/fetch happened on the connection. + assert conn.calls == [] + + +@pytest.mark.asyncio +async def test_get_returns_pending_session_without_db(): + """get() returns the in-memory pending session without hitting the DB.""" + conn = FakeConnection() + pool = FakePool(conn) + store = PgSessionStore(pool) + store._initialized = True + + created = await store.create(profile_id="test", user_id="u1") + fetched = await store.get(created.id) + + assert fetched is created + assert conn.calls == [] # no DB lookup for a pending session + + +@pytest.mark.asyncio +async def test_get_falls_back_to_db_for_unknown_session(): + """get() returns None for an id that is neither pending nor in the DB.""" + conn = FakeConnection() + conn.enqueue(None) # fetchrow for the sessions row -> not found + pool = FakePool(conn) + store = PgSessionStore(pool) + store._initialized = True + + fetched = await store.get("does-not-exist") + assert fetched is None + + +@pytest.mark.asyncio +async def test_save_upserts_and_pops_pending(): + """First save() upserts the sessions row and removes it from _pending.""" + conn = FakeConnection() + conn.enqueue("INSERT 0 1") # upsert sessions row + conn.enqueue("OK") # INSERT executemany (one new message) + conn.enqueue("OK") # UPDATE next_sequence + pool = FakePool(conn) + store = PgSessionStore(pool) + store._initialized = True + + session = await store.create(profile_id="test", user_id="u1") + assert session.id in store._pending + + session.messages.append(Message(role="user", content="hi")) + await store.save(session) + + # The pending entry is gone after the first save. + assert session.id not in store._pending + # The first execute is the upsert (INSERT ... ON CONFLICT), not a plain UPDATE. + execute_calls = [c for c in conn.calls if c[0] == "execute"] + assert "INSERT INTO sessions" in execute_calls[0][1] + assert "ON CONFLICT" in execute_calls[0][1] + # Message got a real sequence number. + assert session.messages[0].sequence_number == 0 + + +@pytest.mark.asyncio +async def test_list_all_excludes_pending(): + """Pending sessions must not appear in list_all (DB query only).""" + conn = FakeConnection() + conn.enqueue([]) # fetch sessions -> DB has no rows + pool = FakePool(conn) + store = PgSessionStore(pool) + store._initialized = True + + await store.create(profile_id="test", user_id="u1") + sessions = await store.list_all(user_id="u1") + + assert sessions == [] + # The pending session is still in memory, just not listed. + assert len(store._pending) == 1 + + +@pytest.mark.asyncio +async def test_sweep_pending_removes_old_entries(): + """sweep_pending drops entries older than the age threshold.""" + from datetime import datetime, timedelta, timezone + + pool = FakePool(FakeConnection()) + store = PgSessionStore(pool) + store._initialized = True + + session = await store.create(profile_id="test") + # Backdate creation so it's older than a 1s threshold. + session.created_at = datetime.now(timezone.utc) - timedelta(hours=2) + + removed = await store.sweep_pending(max_age_seconds=1) + assert removed == 1 + assert session.id not in store._pending + + +@pytest.mark.asyncio +async def test_sweep_pending_keeps_recent_entries(): + """sweep_pending leaves recently-created entries alone.""" + pool = FakePool(FakeConnection()) + store = PgSessionStore(pool) + store._initialized = True + + session = await store.create(profile_id="test") + removed = await store.sweep_pending(max_age_seconds=3600) + assert removed == 0 + assert session.id in store._pending