"""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 ""
