fix(session): preserve complete transcripts
This commit is contained in:
+1
-19
@@ -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
@@ -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,
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user