Merge PR #3459: feat(session): enforce replay/file-cap invariants for history lifecycle
feat(session): enforce replay/file-cap invariants for history lifecycle
This commit is contained in:
+22
-2
@@ -477,6 +477,18 @@ class AgentLoop:
|
||||
return UNIFIED_SESSION_KEY
|
||||
return msg.session_key
|
||||
|
||||
def _replay_token_budget(self) -> int:
|
||||
"""Derive a token budget for session history replay from the context window."""
|
||||
if self.context_window_tokens <= 0:
|
||||
return 0
|
||||
max_output = getattr(getattr(self.provider, "generation", None), "max_tokens", 4096)
|
||||
try:
|
||||
reserved_output = int(max_output)
|
||||
except (TypeError, ValueError):
|
||||
reserved_output = 4096
|
||||
budget = self.context_window_tokens - max(1, reserved_output) - 1024
|
||||
return budget if budget > 0 else max(128, self.context_window_tokens // 2)
|
||||
|
||||
async def _run_agent_loop(
|
||||
self,
|
||||
initial_messages: list[dict],
|
||||
@@ -867,7 +879,10 @@ class AgentLoop:
|
||||
channel, chat_id, msg.metadata.get("message_id"),
|
||||
msg.metadata, session_key=key,
|
||||
)
|
||||
history = session.get_history(max_messages=0, include_timestamps=True)
|
||||
history = session.get_history(
|
||||
max_tokens=self._replay_token_budget(),
|
||||
include_timestamps=True,
|
||||
)
|
||||
current_role = "assistant" if is_subagent else "user"
|
||||
|
||||
# Subagent content is already in `history` above; passing it again
|
||||
@@ -888,6 +903,7 @@ class AgentLoop:
|
||||
pending_queue=pending_queue,
|
||||
)
|
||||
self._save_turn(session, all_msgs, 1 + len(history))
|
||||
session.enforce_file_cap(on_archive=self.context.memory.raw_archive)
|
||||
self._clear_runtime_checkpoint(session)
|
||||
self.sessions.save(session)
|
||||
self._schedule_background(self.consolidator.maybe_consolidate_by_tokens(session))
|
||||
@@ -950,7 +966,10 @@ class AgentLoop:
|
||||
if isinstance(message_tool, MessageTool):
|
||||
message_tool.start_turn()
|
||||
|
||||
history = session.get_history(max_messages=0, include_timestamps=True)
|
||||
history = session.get_history(
|
||||
max_tokens=self._replay_token_budget(),
|
||||
include_timestamps=True,
|
||||
)
|
||||
|
||||
pending_ask_id = pending_ask_user_id(history)
|
||||
if pending_ask_id:
|
||||
@@ -1038,6 +1057,7 @@ class AgentLoop:
|
||||
# Skip the already-persisted user message when saving the turn
|
||||
save_skip = 1 + len(history) + (1 if user_persisted_early else 0)
|
||||
self._save_turn(session, all_msgs, save_skip)
|
||||
session.enforce_file_cap(on_archive=self.context.memory.raw_archive)
|
||||
self._clear_pending_user_turn(session)
|
||||
self._clear_runtime_checkpoint(session)
|
||||
self.sessions.save(session)
|
||||
|
||||
@@ -12,6 +12,7 @@ from loguru import logger
|
||||
|
||||
from nanobot.config.paths import get_legacy_sessions_dir
|
||||
from nanobot.utils.helpers import (
|
||||
estimate_message_tokens,
|
||||
ensure_dir,
|
||||
find_legal_message_start,
|
||||
image_placeholder_text,
|
||||
@@ -19,6 +20,10 @@ from nanobot.utils.helpers import (
|
||||
)
|
||||
|
||||
|
||||
HISTORY_MAX_MESSAGES = 120
|
||||
FILE_MAX_MESSAGES = 2000
|
||||
|
||||
|
||||
@dataclass
|
||||
class Session:
|
||||
"""A conversation session."""
|
||||
@@ -69,11 +74,16 @@ class Session:
|
||||
|
||||
def get_history(
|
||||
self,
|
||||
max_messages: int = 500,
|
||||
max_messages: int = HISTORY_MAX_MESSAGES,
|
||||
*,
|
||||
max_tokens: int = 0,
|
||||
include_timestamps: bool = False,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Return unconsolidated messages for LLM input, aligned to a legal tool-call boundary."""
|
||||
"""Return unconsolidated messages for LLM input.
|
||||
|
||||
History is sliced by message count first (``max_messages``), then by
|
||||
token budget from the tail (``max_tokens``) when provided.
|
||||
"""
|
||||
unconsolidated = self.messages[self.last_consolidated:]
|
||||
sliced = unconsolidated[-max_messages:]
|
||||
|
||||
@@ -113,6 +123,38 @@ class Session:
|
||||
if key in message:
|
||||
entry[key] = message[key]
|
||||
out.append(entry)
|
||||
|
||||
if max_tokens > 0 and out:
|
||||
kept: list[dict[str, Any]] = []
|
||||
used = 0
|
||||
for message in reversed(out):
|
||||
tokens = estimate_message_tokens(message)
|
||||
if kept and used + tokens > max_tokens:
|
||||
break
|
||||
kept.append(message)
|
||||
used += tokens
|
||||
kept.reverse()
|
||||
|
||||
# Keep history aligned to the first visible user turn.
|
||||
first_user = next((i for i, m in enumerate(kept) if m.get("role") == "user"), None)
|
||||
if first_user is not None:
|
||||
kept = kept[first_user:]
|
||||
else:
|
||||
# Tight token budgets can otherwise leave assistant-only tails.
|
||||
# If a user turn exists in the unsliced output, recover the
|
||||
# nearest one even if it slightly exceeds the token budget.
|
||||
recovered_user = next(
|
||||
(i for i in range(len(out) - 1, -1, -1) if out[i].get("role") == "user"),
|
||||
None,
|
||||
)
|
||||
if recovered_user is not None:
|
||||
kept = out[recovered_user:]
|
||||
|
||||
# And keep a legal tool-call boundary at the front.
|
||||
start = find_legal_message_start(kept)
|
||||
if start:
|
||||
kept = kept[start:]
|
||||
out = kept
|
||||
return out
|
||||
|
||||
def clear(self) -> None:
|
||||
@@ -122,31 +164,77 @@ class Session:
|
||||
self.updated_at = datetime.now()
|
||||
|
||||
def retain_recent_legal_suffix(self, max_messages: int) -> None:
|
||||
"""Keep a legal recent suffix, mirroring get_history boundary rules."""
|
||||
"""Keep a legal recent suffix constrained by a hard message cap."""
|
||||
if max_messages <= 0:
|
||||
self.clear()
|
||||
return
|
||||
if len(self.messages) <= max_messages:
|
||||
return
|
||||
|
||||
start_idx = max(0, len(self.messages) - max_messages)
|
||||
retained = list(self.messages[-max_messages:])
|
||||
|
||||
# If the cutoff lands mid-turn, extend backward to the nearest user turn.
|
||||
while start_idx > 0 and self.messages[start_idx].get("role") != "user":
|
||||
start_idx -= 1
|
||||
|
||||
retained = self.messages[start_idx:]
|
||||
# Prefer starting at a user turn when one exists within the tail.
|
||||
first_user = next((i for i, m in enumerate(retained) if m.get("role") == "user"), None)
|
||||
if first_user is not None:
|
||||
retained = retained[first_user:]
|
||||
else:
|
||||
# If the tail is assistant/tool-only, anchor to the latest user in
|
||||
# the full session and take a capped forward window from there.
|
||||
latest_user = next(
|
||||
(i for i in range(len(self.messages) - 1, -1, -1)
|
||||
if self.messages[i].get("role") == "user"),
|
||||
None,
|
||||
)
|
||||
if latest_user is not None:
|
||||
retained = list(self.messages[latest_user: latest_user + max_messages])
|
||||
|
||||
# Mirror get_history(): avoid persisting orphan tool results at the front.
|
||||
start = find_legal_message_start(retained)
|
||||
if start:
|
||||
retained = retained[start:]
|
||||
|
||||
# Hard-cap guarantee: never keep more than max_messages.
|
||||
if len(retained) > max_messages:
|
||||
retained = retained[-max_messages:]
|
||||
start = find_legal_message_start(retained)
|
||||
if start:
|
||||
retained = retained[start:]
|
||||
|
||||
dropped = len(self.messages) - len(retained)
|
||||
self.messages = retained
|
||||
self.last_consolidated = max(0, self.last_consolidated - dropped)
|
||||
self.updated_at = datetime.now()
|
||||
|
||||
def enforce_file_cap(
|
||||
self,
|
||||
on_archive: Any = 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
|
||||
|
||||
before = list(self.messages)
|
||||
before_last_consolidated = self.last_consolidated
|
||||
before_count = len(before)
|
||||
self.retain_recent_legal_suffix(limit)
|
||||
dropped_count = before_count - len(self.messages)
|
||||
if dropped_count <= 0:
|
||||
return
|
||||
|
||||
dropped = before[:dropped_count]
|
||||
already_consolidated = min(before_last_consolidated, dropped_count)
|
||||
archive_chunk = dropped[already_consolidated:]
|
||||
if archive_chunk and on_archive:
|
||||
on_archive(archive_chunk)
|
||||
logger.info(
|
||||
"Session file cap hit for {}: dropped {}, raw-archived {}, kept {}",
|
||||
self.key,
|
||||
dropped_count,
|
||||
len(archive_chunk),
|
||||
len(self.messages),
|
||||
)
|
||||
|
||||
|
||||
class SessionManager:
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user