"""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, match_snippet
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),
)
def message_fingerprint(m: Message) -> tuple:
"""The mutable columns ``save()`` refreshes, as one comparable value.
``thinking`` grows during a turn; length plus the (cached) string hash is
enough to tell "same as last write" from "changed" without keeping a copy.
"""
thinking = m.thinking or ""
return (
m.is_context,
m.is_display,
m.is_summary,
m.is_plan,
m.is_compression,
m.is_recall,
len(thinking),
hash(thinking),
m.elapsed_seconds,
m.tool_call_count,
m.token_count,
)
def saved_fingerprints(messages: list[Message]) -> dict[int, tuple]:
"""Baseline for an in-memory session: what the DB already holds."""
return {
m.sequence_number: message_fingerprint(m)
for m in messages
if m.sequence_number >= 0
}
def preview_of(messages: list[Message]) -> str:
"""The sidebar snippet: the last user/assistant message, 60 chars."""
for msg in reversed(messages):
if msg.role in ("user", "assistant") and msg.content:
return msg.content[:60]
return ""
async def _load_match_info(
conn: asyncpg.Connection, session_ids: list[str], search: str
) -> dict[str, tuple[list[int], str]]:
"""Where the search term appears in each session, without loading history.
``idx`` reproduces the index the caller would get from the hydrated display
list: display rows only, in ``sequence_number`` order, above the session's
archive threshold. The preview comes from the first hit only, so nothing
but one content per session crosses the wire.
The positions are numbered over a *projection* — ``session_id``,
``sequence_number``, ``is_display`` — that ``idx_session_messages_counts``
covers, so the window runs as an index-only scan. Content is then read for
the hits the trigram index returns, not for every display row of the
session: a common term in big sessions cost ~1.7 s on prod the other way.
"""
if not session_ids or not search:
return {}
rows = await conn.fetch(
"""
WITH hot AS (
SELECT m.session_id,
m.sequence_number,
row_number() OVER (
PARTITION BY m.session_id ORDER BY m.sequence_number
) - 1 AS idx
FROM session_messages m
JOIN sessions s ON s.id = m.session_id
WHERE m.session_id = ANY($1)
AND m.is_display
AND m.sequence_number >= coalesce(s.archive_threshold, 0)
), matched AS (
SELECT hot.session_id, hot.idx, hits.content
FROM hot
JOIN (
SELECT m.session_id, m.sequence_number, m.content
FROM session_messages m
WHERE m.session_id = ANY($1)
AND m.is_display
AND m.content ILIKE $2
) hits
ON hits.session_id = hot.session_id
AND hits.sequence_number = hot.sequence_number
)
SELECT session_id,
array_agg(idx ORDER BY idx) AS indices,
(SELECT content
FROM matched x
WHERE x.session_id = matched.session_id
ORDER BY idx
LIMIT 1) AS first_content
FROM matched
GROUP BY session_id
""",
session_ids,
f"%{search}%",
)
return {
r["session_id"]: (list(r["indices"]), match_snippet(r["first_content"] or "", search))
for r in rows
}
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
def _session_head(row: asyncpg.Record, *, parse_json: bool = True) -> Session:
"""The `sessions` row alone — no message payload.
``parse_json=False`` skips `planning_logs` and `session_metadata`, which no
list view reads (the planning endpoint goes through `get()`). Those two
columns reach ~95 kB on prod, and parsing one per session on every sidebar
load was a few tens of milliseconds of pure JSON work.
"""
planning_logs: list = []
if parse_json:
planning_logs_raw = row["planning_logs"]
planning_logs = json.loads(planning_logs_raw) if planning_logs_raw else []
s = Session(
id=row["id"],
profile_id=row["profile_id"],
user_id=row["user_id"],
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.archive_threshold = row.get("archive_threshold", 0) or 0
if parse_json and row.get("session_metadata"):
s.session_metadata = json.loads(row["session_metadata"])
return s
def _next_sequence(row: asyncpg.Record, hot_msgs: list[Message]) -> int:
next_seq = row.get("next_sequence", 0) or 0
if next_seq == 0:
next_seq = max((m.sequence_number for m in hot_msgs), default=-1) + 1
return next_seq
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]
s = _session_head(row)
s.messages = messages
s.context = context
s.db_message_count = len(hot_msgs)
s.db_next_sequence = _next_sequence(row, hot_msgs)
s.display_message_count = len(messages)
s.preview = preview_of(messages)
# What the DB already holds: save() compares against this and rewrites
# only the rows whose mutable state changed (usually the tail).
s.saved_fingerprints = saved_fingerprints(hot_msgs)
sessions.append(s)
return sessions
async def _load_summary_stats(
conn: asyncpg.Connection, session_ids: list[str]
) -> dict[str, tuple[int, int, str]]:
"""Display/total message counts and sidebar preview per session, in one query."""
if not session_ids:
return {}
rows = await conn.fetch(
"""
SELECT m.session_id,
count(*) AS total_count,
count(*) FILTER (WHERE m.is_display) AS display_count,
(SELECT left(m2.content, 60)
FROM session_messages m2
WHERE m2.session_id = m.session_id
AND m2.is_display
AND m2.role IN ('user', 'assistant')
AND coalesce(m2.content, '') <> ''
ORDER BY m2.sequence_number DESC
LIMIT 1) AS preview
FROM session_messages m
JOIN sessions s ON s.id = m.session_id
WHERE m.session_id = ANY($1)
AND m.sequence_number >= s.archive_threshold
GROUP BY m.session_id
""",
session_ids,
)
return {
r["session_id"]: (
int(r["total_count"] or 0),
int(r["display_count"] or 0),
r["preview"] or "",
)
for r in rows
}
async def _build_session_summaries(
conn: asyncpg.Connection,
rows: list[asyncpg.Record],
) -> list[Session]:
"""Session rows with only what the sidebar shows: count and preview.
``GET /sessions`` used to hydrate every hot message of every listed
session (content, thinking, tool-call JSON) just to read the last line and
a count. This keeps the list query proportional to the number of sessions,
not to their history. The heads carry no planning logs and no session
metadata — a list view reads neither; `get()` returns the full session.
"""
stats = await _load_summary_stats(conn, [r["id"] for r in rows])
sessions: list[Session] = []
for row in rows:
total, display, preview = stats.get(row["id"], (0, 0, ""))
s = _session_head(row, parse_json=False)
s.db_message_count = total
s.db_next_sequence = row.get("next_sequence", 0) or 0
s.display_message_count = display
s.preview = preview
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 ""
