"""`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_session_summaries,
_build_sessions,
_load_match_info,
_row_to_message,
_special_and,
_special_frag,
_special_where,
message_fingerprint,
repair_dangling_tool_calls,
saved_fingerprints,
)
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
# Baseline for save(): the state these rows already have in the DB.
s.saved_fingerprints = saved_fingerprints(all_messages)
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]
# Only the rows whose mutable state changed since the last write.
# save() runs after every streaming step; rewriting the whole hot
# history each time was one UPDATE per message per call and grew
# with the session (plus WAL and table churn for untouched rows).
fingerprints: dict[int, tuple] = {}
changed = []
for m in existing:
fp = message_fingerprint(m)
fingerprints[m.sequence_number] = fp
if session.saved_fingerprints.get(m.sequence_number) != fp:
changed.append(m)
if changed:
update_rows = []
for m in changed:
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,
)
# Rebuild rather than merge: rows archived since the last save drop
# out of `messages` and their fingerprints must not linger.
session.saved_fingerprints = fingerprints
# 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,
)
for m in new_msgs:
session.saved_fingerprints[m.sequence_number] = message_fingerprint(m)
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:
# The count rides along in the same round trip: this head is fetched
# on every session open and on every WebSocket reconnect.
row = await conn.fetchrow(
"SELECT s.id, s.profile_id, s.user_id, s.pinned, s.special, s.created_at, s.last_active, "
"s.context_token_count, s.name, s.planning_logs, s.next_sequence, s.archive_threshold, s.session_metadata, "
"(SELECT COUNT(*) FROM session_messages m WHERE m.session_id = s.id) AS message_count "
"FROM sessions s WHERE s.id = $1",
session_id,
)
if not row:
return None
msg_count = row["message_count"]
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 display_index_of(
self, session_id: str, sequence_number: int
) -> int | None:
"""0-based rank of a display row — the same value `get_history_page`
stamps as ``metadata.display_index``.
The client uses it as the stable id of the live turn's assistant block
(``h_<index>``) and as the feedback key, so a live turn and the same
turn after a reload must agree. Counting display rows below the row is
the same ORDER BY the pagination query builds.
"""
pool = await self._get_pool()
async with pool.acquire() as conn:
return await conn.fetchval(
"""
WITH hist AS (
SELECT sequence_number FROM session_messages
WHERE session_id = $1 AND is_display
UNION
SELECT sequence_number FROM session_messages_archive
WHERE session_id = $1 AND is_display
)
SELECT COUNT(*) FROM hist WHERE sequence_number < $2
""",
session_id,
sequence_number,
)
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, with their hot history hydrated.
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.
Callers that only need name/count/preview should ask for
``list_all_summaries`` instead — this one loads every message.
"""
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await self._list_all_rows(conn, user_id=user_id, is_admin=is_admin, special=special)
return await _build_sessions(conn, rows)
async def list_all_summaries(
self, user_id: str | None = None, is_admin: bool = False, special: bool | None = None
) -> list[Session]:
"""``list_all`` without the history: counts and previews come from SQL.
Same rows, same order, same scoping — the difference is that no message
body is read or deserialized, which is what the sidebar needs.
"""
pool = await self._get_pool()
async with pool.acquire() as conn:
rows = await self._list_all_rows(conn, user_id=user_id, is_admin=is_admin, special=special)
return await _build_session_summaries(conn, rows)
async def _list_all_rows(
self, conn: asyncpg.Connection, user_id: str | None, is_admin: bool, special: bool | None
) -> list[asyncpg.Record]:
select = (
"SELECT id, profile_id, user_id, pinned, special, created_at, last_active, "
"context_token_count, name, planning_logs, next_sequence, archive_threshold, session_metadata "
)
if is_admin:
return await conn.fetch(
select + f"FROM sessions{_special_where(special)} ORDER BY pinned DESC, last_active DESC"
)
if user_id is not None:
return await conn.fetch(
select + f"FROM sessions WHERE user_id = $1{_special_and(special)} "
"ORDER BY pinned DESC, last_active DESC",
user_id,
)
return await conn.fetch(
select + f"FROM sessions WHERE user_id IS NULL{_special_and(special)} "
"ORDER BY pinned DESC, last_active DESC"
)
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,
)
# The sidebar needs a count and a 60-char snippet, not the history.
return await _build_session_summaries(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,
)
# Search results are a list view like any other: the caller needs
# counts and previews, and `search_matches` answers "where" from
# SQL. Hydrating every matched history is what made this ~500 ms.
return await _build_session_summaries(conn, rows)
async def search_matches(
self, session_ids: list[str], search: str
) -> dict[str, tuple[list[int], str]]:
"""``{session_id: (display indices, snippet)}`` for a search result page."""
pool = await self._get_pool()
async with pool.acquire() as conn:
return await _load_match_info(conn, session_ids, search)
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)
