fix(session): preserve complete transcripts

This commit is contained in:
chengyongru
2026-08-19 18:40:20 +08:00
committed by chengyongru
parent 3c41d5e7f3
commit 9ef1e292ea
15 changed files with 160 additions and 770 deletions
+1 -19
View File
@@ -14,7 +14,6 @@ from collections.abc import Coroutine, Iterable, Mapping
from contextlib import AbstractContextManager, ExitStack, nullcontext, suppress
from dataclasses import dataclass, field
from enum import Enum, auto
from functools import partial
from pathlib import Path
from typing import TYPE_CHECKING, Any, Awaitable, Callable, TypeVar, cast
@@ -75,12 +74,7 @@ from nanobot.session.goal_state import (
)
from nanobot.session.history_visibility import HIDDEN_HISTORY_META
from nanobot.session.keys import UNIFIED_SESSION_KEY, remember_last_channel
from nanobot.session.manager import (
SESSION_CACHE_MAX_SIZE,
Session,
SessionManager,
replay_max_messages_for_context,
)
from nanobot.session.manager import SESSION_CACHE_MAX_SIZE, Session, SessionManager
from nanobot.session.model_selection import (
SESSION_MODEL_PRESET_METADATA_KEY,
model_preset_from_metadata,
@@ -386,7 +380,6 @@ class AgentLoop:
# WebUI and fork rollback paths. Observe that boundary once instead of
# duplicating cleanup in each consumer.
self.sessions.set_delete_observer(self._file_state_store.discard)
self.sessions.set_file_cap_archiver(self.context.memory.raw_archive)
self.tools = tool_registry if tool_registry is not None else ToolRegistry()
self._exec_session_manager = ExecSessionManager()
self.runner = AgentRunner()
@@ -1819,14 +1812,10 @@ class AgentLoop:
)
if ctx.on_runtime_admitted is not None:
await ctx.on_runtime_admitted(runtime)
replay_max_messages = replay_max_messages_for_context(
runtime.context_window_tokens
)
if not ctx.ephemeral:
await self.consolidator.maybe_consolidate_by_tokens(
session,
runtime=runtime,
replay_max_messages=replay_max_messages,
)
is_subagent = ctx.kind is TurnKind.SYSTEM and ctx.msg.sender_id == "subagent"
@@ -1835,7 +1824,6 @@ class AgentLoop:
message_tool.start_turn()
_hist_kwargs: dict[str, Any] = {
"max_messages": replay_max_messages,
"max_tokens": self._replay_token_budget(runtime),
"extend_to_user": is_subagent,
}
@@ -1999,16 +1987,10 @@ class AgentLoop:
)
ctx.delivery.record_latency(ctx.turn_latency_ms)
if not ctx.ephemeral:
session.enforce_file_cap(
on_archive=partial(self.context.memory.raw_archive, session_key=ctx.session_key)
)
self.schedule_background(
self.consolidator.maybe_consolidate_by_tokens(
session,
runtime=runtime,
replay_max_messages=replay_max_messages_for_context(
runtime.context_window_tokens
),
)
)
self._clear_pending_user_turn(session)
+4 -82
View File
@@ -25,7 +25,6 @@ from nanobot.session.manager import (
MIN_COMPACTED_REPLAY_MESSAGES,
Session,
SessionManager,
replay_max_messages_for_context,
)
from nanobot.session.summary import session_summary_from_metadata
from nanobot.utils.gitstore import GitStore
@@ -34,8 +33,6 @@ from nanobot.utils.helpers import (
ensure_dir,
estimate_message_tokens,
estimate_prompt_tokens_chain,
find_legal_message_start,
recent_message_start_index,
strip_think,
truncate_text,
)
@@ -868,74 +865,7 @@ class Consolidator:
"""Return all messages that can reach the next model prompt."""
if not session.messages:
return []
return session.get_history(max_messages=len(session.messages))
@staticmethod
def _replay_overflow_boundary(
session: Session,
replay_max_messages: int | None,
) -> int | None:
if not replay_max_messages or replay_max_messages <= 0:
return None
tail = list(enumerate(session.messages[session.last_consolidated:], session.last_consolidated))
if len(tail) <= replay_max_messages:
return None
tail_messages = [message for _idx, message in tail]
start_idx = recent_message_start_index(
tail_messages,
replay_max_messages,
extend_to_user=True,
)
sliced = tail[start_idx:]
for i, (_idx, message) in enumerate(sliced):
if message.get("role") == "user":
start = i
if i > 0 and sliced[i - 1][1].get("_channel_delivery"):
start = i - 1
sliced = sliced[start:]
break
legal_start = find_legal_message_start([message for _idx, message in sliced])
if legal_start:
sliced = sliced[legal_start:]
if not sliced:
return len(session.messages)
first_visible_idx = sliced[0][0]
if first_visible_idx <= session.last_consolidated:
return None
return first_visible_idx
async def _consolidate_replay_overflow(
self,
session: Session,
replay_max_messages: int | None,
*,
runtime: LLMRuntime,
) -> str | None:
"""Archive messages that would be hidden by the replay message window."""
end_idx = self._replay_overflow_boundary(session, replay_max_messages)
if end_idx is None:
return None
chunk = session.messages[session.last_consolidated:end_idx]
if not chunk:
return None
logger.info(
"Replay-window consolidation for {}: chunk={} msgs, replay_max={}",
session.key,
len(chunk),
replay_max_messages,
)
summary = await self.archive_session(
session,
archive_end=end_idx,
runtime=runtime,
)
session.last_consolidated = end_idx
session.provider_state = None
self.sessions.save(session)
return summary
return session.get_history()
def _persist_last_summary(self, session: Session, summary: str | None) -> None:
if summary and summary != "(nothing)":
@@ -1056,14 +986,11 @@ class Consolidator:
messages=list(session.messages[:archive_end]),
last_consolidated=session.last_consolidated,
)
history = prefix.get_history(
max_messages=replay_max_messages_for_context(runtime.context_window_tokens),
max_tokens=budget,
)
history = prefix.get_history(max_tokens=budget)
archive_history = Session(
key=session.key,
messages=messages,
).get_history(max_messages=len(messages))
).get_history()
if (
not archive_history
or history[-len(archive_history):] != archive_history
@@ -1125,7 +1052,6 @@ class Consolidator:
session: Session,
*,
runtime: LLMRuntime,
replay_max_messages: int | None = None,
) -> None:
"""Loop: archive old messages until prompt fits within safe budget.
@@ -1146,11 +1072,7 @@ class Consolidator:
budget = self._input_token_budget(runtime)
target = int(budget * self.consolidation_ratio)
last_summary = await self._consolidate_replay_overflow(
session,
replay_max_messages,
runtime=runtime,
)
last_summary: str | None = None
estimated, source = self.estimate_session_prompt_tokens(
session,
runtime=runtime,
+1 -5
View File
@@ -15,7 +15,6 @@ from nanobot.sdk.types import (
snapshot_from_payload,
snapshot_from_session,
)
from nanobot.session.manager import replay_max_messages_for_context
if TYPE_CHECKING:
from nanobot.agent.loop import AgentLoop
@@ -210,15 +209,12 @@ class RuntimeClient:
return self._loop.runtime_events.subscribe(handler, SessionTurnPersisted)
async def compact_session(self, session_key: str) -> SessionSnapshot:
"""Run token/replay-window consolidation for one session."""
"""Run token consolidation for one session."""
session = self._loop.sessions.get_or_create(session_key)
runtime = self._loop.runtime_for_session(session)
await self._loop.consolidator.maybe_consolidate_by_tokens(
session,
runtime=runtime,
replay_max_messages=replay_max_messages_for_context(
runtime.context_window_tokens
),
)
return snapshot_from_session(self._loop.sessions.get_or_create(session_key))
+15 -79
View File
@@ -39,11 +39,8 @@ from nanobot.utils.helpers import (
)
from nanobot.utils.subagent_channel_display import scrub_subagent_announce_body
FILE_MAX_MESSAGES = 2000
SESSION_CACHE_MAX_SIZE = 128
MIN_REPLAY_MAX_MESSAGES = 120
MIN_COMPACTED_REPLAY_MESSAGES = 8
REPLAY_TOKENS_PER_MESSAGE = 100
_MESSAGE_TIME_PREFIX_RE = re.compile(r"^\[Message Time: [^\]]+\]\n?")
_LOCAL_IMAGE_BREADCRUMB_RE = re.compile(r"^\[image: (?:/|~)[^\]]+\]\s*$")
_TOOL_CALL_ECHO_RE = re.compile(r'^\s*(?:generate_image|message)\([^)]*\)\s*$')
@@ -84,15 +81,6 @@ def _is_provider_state_record_line(line: str) -> bool:
return _PROVIDER_STATE_RECORD_PREFIX_RE.match(line) is not None
def replay_max_messages_for_context(context_window_tokens: int | None) -> int:
if not context_window_tokens or context_window_tokens <= 0:
return FILE_MAX_MESSAGES
return min(
FILE_MAX_MESSAGES,
max(MIN_REPLAY_MAX_MESSAGES, context_window_tokens // REPLAY_TOKENS_PER_MESSAGE),
)
def _sanitize_assistant_replay_text(content: str) -> str:
"""Remove internal replay artifacts that the model may have copied before.
@@ -209,7 +197,7 @@ class Session:
def get_history(
self,
max_messages: int = FILE_MAX_MESSAGES,
max_messages: int = 0,
*,
max_tokens: int = 0,
extend_to_user: bool = False,
@@ -217,8 +205,8 @@ class Session:
) -> list[dict[str, Any]]:
"""Return recent replayable messages for LLM input.
History is sliced by message count first (``max_messages``), then by
token budget from the tail (``max_tokens``) when provided.
A positive ``max_messages`` applies an explicit caller-owned count
limit. The normal model path relies on ``max_tokens`` instead.
"""
replay_start = self.last_consolidated
if replay_start:
@@ -233,18 +221,20 @@ class Session:
replay_start = min(replay_start, recent_start)
replayable = self.messages[replay_start:]
max_messages = max_messages if max_messages > 0 else FILE_MAX_MESSAGES
unarchived_count = len(self.messages) - self.last_consolidated
if replay_start < self.last_consolidated and unarchived_count < max_messages:
# The archived replay suffix can exceed the nominal count when one
# tool-heavy turn spans the boundary. Preserve that complete turn.
if max_messages <= 0:
start_idx = 0
else:
start_idx = recent_message_start_index(
replayable,
max_messages,
extend_to_user=extend_to_user,
)
unarchived_count = len(self.messages) - self.last_consolidated
if replay_start < self.last_consolidated and unarchived_count < max_messages:
# The archived replay suffix can exceed the nominal count when one
# tool-heavy turn spans the boundary. Preserve that complete turn.
start_idx = 0
else:
start_idx = recent_message_start_index(
replayable,
max_messages,
extend_to_user=extend_to_user,
)
sliced = replayable[start_idx:]
# Avoid starting mid-turn when possible, except for proactive
@@ -467,46 +457,6 @@ class Session:
already_consolidated_count=already_consolidated,
)
def enforce_file_cap(
self,
on_archive: Callable[[list[dict[str, Any]]], None] | None = None,
limit: int = FILE_MAX_MESSAGES,
) -> None:
"""Bound session message growth by archiving and trimming old prefixes."""
if limit <= 0 or len(self.messages) <= limit:
return
original_messages = self.messages
original_last_consolidated = self.last_consolidated
original_provider_state = self.provider_state
original_updated_at = self.updated_at
result = self.retain_recent_legal_suffix(limit)
if not result.dropped:
return
archive_chunk = result.dropped[result.already_consolidated_count:]
if archive_chunk and on_archive:
try:
on_archive(archive_chunk)
except BaseException:
# Retention runs before the archive callback so the callback can
# receive the exact dropped prefix. Restore the in-memory session
# if archival fails; otherwise a later save would persist the
# trimmed state and make that prefix impossible to retry.
self.messages = original_messages
self.last_consolidated = original_last_consolidated
self.provider_state = original_provider_state
self.updated_at = original_updated_at
raise
logger.info(
"Session file cap hit for {}: dropped {}, raw-archived {}, kept {}",
self.key,
len(result.dropped),
len(archive_chunk),
len(self.messages),
)
class SessionPayload(TypedDict):
key: str
created_at: str | None
@@ -1576,7 +1526,6 @@ class SessionManager:
# Preserve identity for sessions held by active callers without retaining idle ones.
self._overflow_cache: WeakValueDictionary[str, Session] = WeakValueDictionary()
self._max_cached_sessions = SESSION_CACHE_MAX_SIZE
self._file_cap_archiver: Callable[..., None] | None = None
self._delete_observer: Callable[[str], None] | None = None
def _remember(self, session: Session) -> None:
@@ -1603,10 +1552,6 @@ class SessionManager:
"""Return a cached session without creating or loading one from disk."""
return self._cached(key)
def set_file_cap_archiver(self, archiver: Callable[..., None]) -> None:
"""Archive unconsolidated overflow whenever a session is persisted."""
self._file_cap_archiver = archiver
def set_delete_observer(self, observer: Callable[[str], None]) -> None:
"""Observe explicit session deletion for process-local state cleanup."""
self._delete_observer = observer
@@ -1705,15 +1650,6 @@ class SessionManager:
if not session.policy.persist:
return
archiver = self._file_cap_archiver
if archiver is not None:
session.enforce_file_cap(
on_archive=lambda messages: archiver(
messages,
session_key=session.key,
)
)
self._store.save(session, fsync=fsync)
self._remember(session)