mirror/nanobot
1
0
Fork 0
mirror of https://github.com/HKUDS/nanobot.git synced 2026-09-04 10:11:46 +03:00
Code Issues Packages Projects Releases Wiki Activity

refactor(agent): let runner own context compaction (#5568)

Browse Source 
* refactor(agent): consolidate accepted history under pressure

* fix(agent): align provider and session compaction

* refactor(agent): simplify runner context compaction

* refactor(agent): remove background token consolidation

* fix(agent): keep injected transcript messages distinct

* refactor(agent): unify native compaction summaries

* fix(agent): preserve native compaction boundary

* fix(agent): unify context compaction paths

* fix(agent): preserve exact compaction request boundaries
This commit is contained in:
chengyongru
2026-09-02 18:05:54 +08:00
committed by GitHub
parent da96c5c6eb
commit d81aa5a4ab
44 changed files with 1914 additions and 1666 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/context.py
+3 -1
View File
@@ -235,7 +235,7 @@ class ContextBuilder:
def build_messages(
self,
history: list[dict[str, Any]],
current_message: str,
current_message: str | None,
*,
media: list[str] | None = None,
channel: str | None = None,
@@ -259,6 +259,8 @@ class ContextBuilder:
workspace=workspace,
include_memory=include_memory,
)
if current_message is None:
return messages
current = messages[-1]
if len(messages) < 2 or messages[-2].get("role") != current.get("role"):
return messages
nanobot/agent/context_governance.py
+502 -39
View File
@@ -1,19 +1,43 @@
"""Model-message governance for agent runner requests.
"""Model-message governance and compaction for agent runner requests.
This module owns model-facing message shaping and tool-result content normalization.
It may return copied messages or persisted-result placeholders, but it must not
mutate an existing session history list in place.
This module owns model-facing message shaping, request pressure, H/delta
compaction state, and tool-result content normalization. It may return copied
messages or persisted-result placeholders, but it must not mutate an existing
session history list in place.
"""
from __future__ import annotations
from dataclasses import dataclass
from collections.abc import Awaitable, Callable
from copy import deepcopy
from dataclasses import dataclass, replace
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast
from loguru import logger
from nanobot.providers.base import LLMUsage
from nanobot.agent.context import TranscriptInput
from nanobot.providers.base import (
LLMResponse,
LLMUsage,
ProviderCallContext,
ProviderConversationState,
)
from nanobot.providers.conversation_state import (
ProviderConversationStateController,
allows_conversation_message_merge,
)
from nanobot.runtime_context import (
RUNTIME_CONTEXT_MESSAGE_META,
detach_runtime_context,
reattach_runtime_context,
)
from nanobot.session.history_visibility import is_hidden_history_message
from nanobot.session.summary import (
SUMMARY_CONTINUATION_TEXT,
SessionSummaryCheckpoint,
)
from nanobot.utils.helpers import (
estimate_message_tokens,
estimate_prompt_tokens_chain,
@@ -27,6 +51,16 @@ if TYPE_CHECKING:
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.providers.base import LLMProvider
TranscriptBuilder = Callable[[TranscriptInput], list[dict[str, Any]]]
HistoryConsolidator = Callable[
[list[dict[str, Any]], str | None],
Awaitable[str | None],
]
ProviderCompactionConsolidator = Callable[
[ProviderConversationState, list[dict[str, Any]], str | None],
Awaitable[str | None],
]
SNIP_SAFETY_BUFFER = 1024
# read_file is the recovery path for persisted results; exempting it prevents persist->read->persist loops.
TOOL_RESULT_OFFLOAD_EXEMPT_TOOLS = frozenset({"read_file"})
@@ -85,8 +119,204 @@ class ContextGovernanceConfig:
max_tokens: int | None = None
@dataclass(slots=True)
class ContextCompactionState:
"""Track accepted provider input H separately from the unsent delta."""
raw_messages: list[dict[str, Any]]
accepted_messages: list[dict[str, Any]]
raw_accepted_boundary: int
active_summary: str | None
transcript_input: TranscriptInput
transcript_builder: TranscriptBuilder
consolidate_history: HistoryConsolidator
consolidate_provider_compaction: ProviderCompactionConsolidator | None
summary_checkpoint: SessionSummaryCheckpoint | None = None
@classmethod
def from_transcript(
cls,
transcript_input: TranscriptInput,
transcript_builder: TranscriptBuilder,
consolidate_history: HistoryConsolidator | None,
consolidate_provider_compaction: ProviderCompactionConsolidator | None,
) -> tuple[list[dict[str, Any]], ContextCompactionState | None]:
"""Build the raw transcript and its initial H/delta boundary."""
messages = list(transcript_builder(transcript_input))
if consolidate_history is None:
return messages, None
accepted_history_boundary = 1 + len(transcript_input.history)
return messages, cls(
raw_messages=messages,
accepted_messages=deepcopy(messages[:accepted_history_boundary]),
raw_accepted_boundary=accepted_history_boundary,
active_summary=(
transcript_input.session_summary["text"]
if transcript_input.session_summary is not None
else None
),
transcript_input=transcript_input,
transcript_builder=transcript_builder,
consolidate_history=consolidate_history,
consolidate_provider_compaction=consolidate_provider_compaction,
)
def request_messages(
self,
raw_messages: list[dict[str, Any]],
) -> list[dict[str, Any]]:
return [
*deepcopy(self.accepted_messages),
*deepcopy(raw_messages[self.raw_accepted_boundary:]),
]
def delta_after_accepted(
self,
request_messages: list[dict[str, Any]],
) -> list[dict[str, Any]]:
return deepcopy(request_messages[len(self.accepted_messages):])
def accept_request(
self,
model_messages: list[dict[str, Any]],
*,
raw_boundary: int,
) -> None:
"""Advance H after the provider has received one request."""
self.accepted_messages = deepcopy(model_messages)
self.raw_accepted_boundary = raw_boundary
@dataclass(slots=True)
class ModelRequestState:
"""Context state shared by every provider request in one runner turn."""
config: ContextGovernanceConfig
conversation: ProviderConversationStateController
usage: LLMUsage | None = None
messages: list[dict[str, Any]] | None = None
tool_definitions: list[dict[str, Any]] | None = None
compaction: ContextCompactionState | None = None
provider_compaction_applied: bool = False
class ContextGovernor:
"""Prepare model-copy messages while preserving persisted history."""
"""Own model-request context while preserving persisted history."""
@staticmethod
def _merge_message_content(left: Any, right: Any) -> str | list[dict[str, Any]]:
if isinstance(left, str) and isinstance(right, str):
return f"{left}\n\n{right}" if left else right
def _to_blocks(value: Any) -> list[dict[str, Any]]:
if isinstance(value, list):
return [
cast(dict[str, Any], item)
if isinstance(item, dict)
else {"type": "text", "text": str(item)}
for item in cast(list[Any], value)
]
if value is None:
return []
return [{"type": "text", "text": str(value)}]
return _to_blocks(left) + _to_blocks(right)
@classmethod
def _merge_adjacent_user_messages_for_model(
cls,
messages: list[dict[str, Any]],
) -> list[dict[str, Any]]:
"""Merge adjacent visible user messages only in the model-facing copy."""
prepared: list[dict[str, Any]] = []
for source in messages:
injection = deepcopy(source)
if (
prepared
and injection.get("role") == "user"
and prepared[-1].get("role") == "user"
and injection.get("content") != SUMMARY_CONTINUATION_TEXT
and prepared[-1].get("content") != SUMMARY_CONTINUATION_TEXT
and not is_hidden_history_message(injection)
and not is_hidden_history_message(prepared[-1])
and allows_conversation_message_merge(injection)
and allows_conversation_message_merge(prepared[-1])
):
merged = dict(prepared[-1])
left_meta = merged.get("_meta")
right_meta = injection.get("_meta")
left_meta_dict = (
cast(dict[str, Any], left_meta) if isinstance(left_meta, dict) else None
)
right_meta_dict = (
cast(dict[str, Any], right_meta) if isinstance(right_meta, dict) else None
)
left_marker = (
left_meta_dict.get(RUNTIME_CONTEXT_MESSAGE_META)
if left_meta_dict is not None
else None
)
right_marker = (
right_meta_dict.get(RUNTIME_CONTEXT_MESSAGE_META)
if right_meta_dict is not None
else None
)
left_marker_dict = (
cast(dict[str, Any], left_marker) if isinstance(left_marker, dict) else None
)
right_marker_dict = (
cast(dict[str, Any], right_marker) if isinstance(right_marker, dict) else None
)
empty_sources: list[str] = []
empty_blocks: list[dict[str, Any]] = []
detached_left = (
detach_runtime_context(merged.get("content"), left_marker_dict)
if left_marker_dict is not None
else (merged.get("content"), empty_sources, empty_blocks)
)
detached_right = (
detach_runtime_context(injection.get("content"), right_marker_dict)
if right_marker_dict is not None
else (injection.get("content"), empty_sources, empty_blocks)
)
if detached_left is not None and detached_right is not None:
left_content, left_sources, left_blocks = detached_left
right_content, right_sources, right_blocks = detached_right
merged_content = cls._merge_message_content(left_content, right_content)
context_blocks = [*left_blocks, *right_blocks]
if context_blocks:
merged_content, marker = reattach_runtime_context(
merged_content,
[*left_sources, *right_sources],
context_blocks,
)
internal_meta = (
dict(left_meta_dict) if left_meta_dict is not None else {}
)
if right_meta_dict is not None:
for key, value in right_meta_dict.items():
internal_meta.setdefault(key, value)
internal_meta[RUNTIME_CONTEXT_MESSAGE_META] = marker
merged["_meta"] = internal_meta
merged["content"] = merged_content
else:
merged["content"] = cls._merge_message_content(
merged.get("content"),
injection.get("content"),
)
prepared[-1] = merged
continue
prepared.append(injection)
return prepared
def prepare_messages_for_model(
self,
config: ContextGovernanceConfig,
messages: list[dict[str, Any]],
) -> list[dict[str, Any]]:
"""Build the normalized model-facing copy of a raw transcript."""
governed = self.prepare_for_model(config, messages)
return self._merge_adjacent_user_messages_for_model(governed)
def prepare_for_model(
self,
@@ -115,17 +345,31 @@ class ContextGovernor:
)
updated = self.drop_orphan_tool_results(updated)
updated = self.backfill_missing_tool_results(updated)
return self.ensure_request_fits(
config,
updated,
tool_definitions=tool_definitions,
)
def ensure_request_fits(
self,
config: ContextGovernanceConfig,
messages: list[dict[str, Any]],
*,
tool_definitions: list[dict[str, Any]] | None,
) -> list[dict[str, Any]]:
"""Validate an exact model request without dropping any messages."""
if not config.context_window_tokens:
return updated
return messages
budget = self.input_budget(config)
estimated, source = estimate_prompt_tokens_chain(
config.provider,
config.model,
updated,
messages,
tool_definitions,
)
if budget > 0 and estimated <= budget:
return updated
return messages
raise ContextWindowExceededError(
session_key=config.session_key,
estimated_tokens=estimated,
@@ -133,6 +377,41 @@ class ContextGovernor:
source=source,
)
def request_pressure(
self,
config: ContextGovernanceConfig,
messages: list[dict[str, Any]],
usage: LLMUsage | None,
*,
usage_matches_messages: bool,
tool_definitions: list[dict[str, Any]] | None,
request_context_tokens: int | None = None,
) -> tuple[int, str] | None:
"""Return the authoritative measurement when a request is pressured."""
if not config.context_window_tokens:
return None
budget = self.input_budget(config)
if request_context_tokens is not None:
measured = request_context_tokens
source = "resumed provider state plus pending messages"
elif (
usage_matches_messages
and usage is not None
and usage.context_tokens is not None
):
measured = usage.context_tokens
source = "matching provider usage"
else:
measured, source = estimate_prompt_tokens_chain(
config.provider,
config.model,
messages,
tool_definitions,
)
if budget > 0 and measured < budget:
return None
return measured, source
def fit_request(
self,
config: ContextGovernanceConfig,
@@ -144,27 +423,15 @@ class ContextGovernor:
request_context_tokens: int | None = None,
) -> tuple[list[dict[str, Any]], bool]:
"""Fit the request when its measured or estimated input is pressured."""
if not config.context_window_tokens:
return messages, False
budget = self.input_budget(config)
if (
request_context_tokens is None
and usage_matches_messages
and usage is not None
and usage.context_tokens is not None
):
pressured = budget <= 0 or usage.context_tokens >= budget
else:
estimated, _ = estimate_prompt_tokens_chain(
config.provider,
config.model,
messages,
tool_definitions,
)
if request_context_tokens is not None:
estimated = max(estimated, request_context_tokens)
pressured = budget <= 0 or estimated >= budget
if not pressured:
pressure = self.request_pressure(
config,
messages,
usage,
usage_matches_messages=usage_matches_messages,
tool_definitions=tool_definitions,
request_context_tokens=request_context_tokens,
)
if pressure is None:
return messages, False
return self.fit_to_budget(
config,
@@ -172,6 +439,201 @@ class ContextGovernor:
tool_definitions=tool_definitions,
), True
@staticmethod
def _summary_transcript(
compaction: ContextCompactionState,
summary: str,
) -> list[dict[str, Any]]:
"""Rebuild only the stable system prefix around a replacement summary."""
return compaction.transcript_builder(
replace(
compaction.transcript_input,
history=[],
current_message=None,
media=None,
session_summary={
"text": summary,
"last_active": datetime.now().astimezone().isoformat(),
},
runtime_context_blocks=None,
)
)
async def summarize_provider_compaction(
self,
state: ModelRequestState,
response: LLMResponse,
*,
current_request_boundary: int | None,
) -> None:
"""Materialize the exact input replaced by provider-native compaction."""
compaction = state.compaction
if (
not response.provider_compaction_applied
or response.provider_compaction_state is None
or compaction is None
or compaction.consolidate_provider_compaction is None
):
return
if response.provider_compaction_scope == "prior_context":
accepted_messages = compaction.accepted_messages
transcript_boundary = compaction.raw_accepted_boundary
elif (
response.provider_compaction_scope == "current_request"
and state.messages is not None
and current_request_boundary is not None
):
accepted_messages = state.messages
transcript_boundary = current_request_boundary
else:
logger.warning(
"Ignoring provider compaction with missing request-boundary scope for {}",
state.config.session_key or "default",
)
return
summary = await compaction.consolidate_provider_compaction(
response.provider_compaction_state,
deepcopy(accepted_messages),
compaction.active_summary,
)
if not summary:
return
compaction.active_summary = summary
compaction.summary_checkpoint = SessionSummaryCheckpoint(
summary=summary,
transcript_boundary=transcript_boundary,
)
async def _compact_request_history(
self,
state: ModelRequestState,
compaction: ContextCompactionState,
messages: list[dict[str, Any]],
pressure: tuple[int, str],
*,
tool_definitions: list[dict[str, Any]] | None,
) -> list[dict[str, Any]]:
"""Replace accepted history H with a checkpoint while preserving delta."""
delta_messages = compaction.delta_after_accepted(messages)
consolidation_prefix = self.prepare_messages_for_model(
state.config,
compaction.accepted_messages,
)
summary = await compaction.consolidate_history(
deepcopy(consolidation_prefix),
compaction.active_summary,
)
if not summary:
measured, source = pressure
raise ContextWindowExceededError(
session_key=state.config.session_key,
estimated_tokens=measured,
input_budget=self.input_budget(state.config),
source=source,
)
compaction.active_summary = summary
prepared = self.prepare_messages_for_model(
state.config,
[
*self._summary_transcript(compaction, summary),
{"role": "user", "content": SUMMARY_CONTINUATION_TEXT},
*delta_messages,
],
)
# Responses-style state is append-only. Replacing H with a
# checkpoint requires a fresh request; a successful response may
# establish a new provider-owned state at the rewritten boundary.
state.conversation.replace_transcript(compaction.raw_messages)
state.usage = None
prepared = self.ensure_request_fits(
state.config,
prepared,
tool_definitions=tool_definitions,
)
compaction.summary_checkpoint = SessionSummaryCheckpoint(
summary=summary,
transcript_boundary=compaction.raw_accepted_boundary,
)
return prepared
async def prepare_request(
self,
state: ModelRequestState,
messages: list[dict[str, Any]],
*,
tool_definitions: list[dict[str, Any]] | None,
transcript: list[dict[str, Any]] | None = None,
) -> tuple[list[dict[str, Any]], ProviderCallContext | None]:
"""Prepare, compact or fit, and record the exact provider payload."""
prepared = self.prepare_messages_for_model(state.config, messages)
model_messages: list[dict[str, Any]] | None = prepared
supplemental_messages: list[dict[str, Any]] | None = None
request_context_tokens = None
if transcript is not None:
if tool_definitions is None:
model_messages = None
supplemental_messages = [prepared[-1]]
request_context_tokens = state.conversation.estimate_request_context_tokens(
transcript,
model_messages=model_messages,
supplemental_messages=supplemental_messages,
tool_definitions=tool_definitions,
)
usage_matches_messages = (
state.messages is not None
and prepared == state.messages
and tool_definitions == state.tool_definitions
)
request_was_fitted = False
compaction = state.compaction
if compaction is None:
prepared, request_was_fitted = self.fit_request(
state.config,
prepared,
state.usage,
usage_matches_messages=usage_matches_messages,
tool_definitions=tool_definitions,
request_context_tokens=request_context_tokens,
)
else:
pressure = self.request_pressure(
state.config,
prepared,
state.usage,
usage_matches_messages=usage_matches_messages,
tool_definitions=tool_definitions,
request_context_tokens=request_context_tokens,
)
if pressure is not None:
prepared = await self._compact_request_history(
state,
compaction,
messages,
pressure,
tool_definitions=tool_definitions,
)
model_messages = prepared
supplemental_messages = None
provider_context = (
state.conversation.prepare_request(
transcript,
context_window_tokens=state.config.context_window_tokens,
model_messages=model_messages,
supplemental_messages=supplemental_messages,
resume_state=not request_was_fitted,
)
if transcript is not None
else state.conversation.independent_request_context(
context_window_tokens=state.config.context_window_tokens,
)
)
state.messages = deepcopy(prepared)
state.tool_definitions = deepcopy(tool_definitions)
return prepared, provider_context
@staticmethod
def input_budget(config: ContextGovernanceConfig) -> int:
if not config.context_window_tokens:
@@ -424,14 +886,15 @@ class ContextGovernor:
if budget <= 0:
return messages
estimate, _ = estimate_prompt_tokens_chain(
config.provider,
config.model,
messages,
tool_definitions,
)
if not force and estimate <= budget:
return messages
if not force:
estimate, _ = estimate_prompt_tokens_chain(
config.provider,
config.model,
messages,
tool_definitions,
)
if estimate <= budget:
return messages
system_messages = [dict(msg) for msg in messages if msg.get("role") == "system"]
non_system = [dict(msg) for msg in messages if msg.get("role") != "system"]
nanobot/agent/loop.py
+119 -38
View File
@@ -13,6 +13,7 @@ import weakref
from collections.abc import Coroutine, Iterable, Mapping
from contextlib import AbstractContextManager, ExitStack, nullcontext, suppress
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum, auto
from functools import partial
from pathlib import Path
@@ -93,7 +94,11 @@ from nanobot.session.recovery import (
restore_pending_interruption,
restore_runtime_checkpoint,
)
from nanobot.session.summary import SessionSummary
from nanobot.session.summary import (
SUMMARY_CONTINUATION_TEXT,
SessionSummary,
SessionSummaryCheckpoint,
)
from nanobot.triggers.local_turns import LocalTriggerTurnCoordinator
from nanobot.utils.cancellation import task_is_cancelling
from nanobot.utils.document import reference_non_image_attachments
@@ -161,6 +166,8 @@ class TurnContext:
pending_queue: asyncio.Queue[InboundMessage] | None = None
pending_summary: SessionSummary | None = None
summary_checkpoint: SessionSummaryCheckpoint | None = None
provider_compaction_applied: bool = False
ephemeral: bool = False
run_extra_hooks_for_ephemeral: bool = False
@@ -923,19 +930,6 @@ class AgentLoop:
return
remember_last_channel(session.metadata, msg.channel, msg.chat_id)
@staticmethod
def _replay_token_budget(runtime: LLMRuntime) -> int:
"""Derive a token budget for session history replay from the context window."""
if runtime.context_window_tokens <= 0:
return 0
max_output = runtime.generation.max_tokens
try:
reserved_output = int(max_output)
except (TypeError, ValueError):
reserved_output = 4096
budget = runtime.context_window_tokens - max(1, reserved_output) - 1024
return budget if budget > 0 else max(128, runtime.context_window_tokens // 2)
async def _run_agent_loop(
self,
transcript_input: TranscriptInput,
@@ -1186,6 +1180,26 @@ class AgentLoop:
provider_retry_mode=self.provider_retry_mode,
retry_wait_callback=on_retry_wait,
checkpoint_callback=_checkpoint,
consolidate_history=(
partial(
self.consolidator.summarize_transcript,
runtime=runtime,
session_key=session.key,
tools=effective_tools.get_definitions(),
)
if session is not None and not ephemeral
else None
),
consolidate_provider_compaction=(
partial(
self.consolidator.summarize_provider_compaction,
runtime=runtime,
session_key=session.key,
tools=effective_tools.get_definitions(),
)
if session is not None and not ephemeral
else None
),
injection_callback=_drain_pending,
terminal_injection_callback=_wait_for_pending,
# Sustained goals may legitimately exceed NANOBOT_LLM_TIMEOUT_S; idle stall
@@ -1886,12 +1900,6 @@ class AgentLoop:
if ctx.on_runtime_admitted is not None:
await ctx.on_runtime_admitted(runtime)
if not ctx.ephemeral:
await self.consolidator.maybe_consolidate_by_tokens(
session,
runtime=runtime,
)
# Token consolidation may have committed a replacement checkpoint
# after the compact stage captured its summary for this request.
ctx.session, ctx.pending_summary = self.auto_compact.prepare_session(
session,
ctx.session_key,
@@ -1899,11 +1907,7 @@ class AgentLoop:
session = ctx.require_session()
is_subagent = ctx.kind is TurnKind.SYSTEM and ctx.msg.sender_id == "subagent"
_hist_kwargs: dict[str, Any] = {
"max_tokens": self._replay_token_budget(runtime),
"extend_to_user": is_subagent,
}
ctx.history = session.get_history(**_hist_kwargs)
ctx.history = session.get_history(extend_to_user=is_subagent)
stored_state = session.provider_state
subagent_followup_persisted = False
if is_subagent:
@@ -2021,6 +2025,8 @@ class AgentLoop:
)
ctx.final_content = result.final_content
ctx.all_messages = result.messages
ctx.summary_checkpoint = result.summary_checkpoint
ctx.provider_compaction_applied = result.provider_compaction_applied
ctx.stop_reason = result.stop_reason
if (
ctx.kind is TurnKind.USER
@@ -2034,7 +2040,6 @@ class AgentLoop:
await turn_continuation.maybe_continue_turn(ctx)
async def _persist_turn(self, ctx: TurnContext) -> None:
runtime = ctx.require_runtime()
session = ctx.require_session()
turn_continuation.prepare_save_boundary(ctx)
@@ -2060,15 +2065,18 @@ class AgentLoop:
self._save_turn(
session, ctx.all_messages, ctx.save_skip,
turn_latency_ms=ctx.turn_latency_ms,
summary_checkpoint=ctx.summary_checkpoint,
input_persisted_early=ctx.input_persisted_early,
)
if (
not ctx.ephemeral
and ctx.provider_compaction_applied
and ctx.summary_checkpoint is not None
):
# The next request must rebuild from the portable checkpoint;
# the opaque continuation predates that transcript rewrite.
session.provider_state = None
ctx.delivery.record_latency(ctx.turn_latency_ms)
if not ctx.ephemeral:
self.schedule_background(
self.consolidator.maybe_consolidate_by_tokens(
session,
runtime=runtime,
)
)
self._clear_pending_user_turn(session)
self._clear_runtime_checkpoint(session)
self.sessions.save(session)
@@ -2142,6 +2150,55 @@ class AgentLoop:
return filtered
@staticmethod
def _insert_summary_checkpoint(
session: Session,
checkpoint: SessionSummaryCheckpoint,
*,
insert_at: int | None = None,
) -> None:
"""Commit a replacement summary and its hidden transcript boundary."""
hint = {
"role": "user",
"content": SUMMARY_CONTINUATION_TEXT,
HIDDEN_HISTORY_META: True,
"timestamp": datetime.now().isoformat(),
}
if insert_at is None:
session.messages.append(hint)
checkpoint_session_index = len(session.messages) - 1
else:
session.messages.insert(insert_at, hint)
checkpoint_session_index = insert_at
session.metadata["_last_summary"] = {
"text": checkpoint.summary,
"last_active": session.updated_at.isoformat(),
}
session.last_archived = checkpoint_session_index
@staticmethod
def _validated_checkpoint_boundary(
checkpoint: SessionSummaryCheckpoint | None,
*,
skip: int,
message_count: int,
session_key: str,
) -> int | None:
"""Return a checkpoint boundary only when it belongs to this turn."""
if checkpoint is None:
return None
boundary = checkpoint.transcript_boundary
if skip - 1 <= boundary <= message_count:
return boundary
logger.warning(
"Ignoring invalid summary boundary {} outside [{}, {}] for {}",
boundary,
skip - 1,
message_count,
session_key,
)
return None
def _save_turn(
self,
session: Session,
@@ -2149,10 +2206,10 @@ class AgentLoop:
skip: int,
*,
turn_latency_ms: int | None = None,
summary_checkpoint: SessionSummaryCheckpoint | None = None,
input_persisted_early: bool = False,
) -> None:
"""Save new-turn messages into session, truncating large tool results."""
from datetime import datetime
"""Commit new-turn messages and an optional summary boundary."""
declared_tool_call_ids = {
str(tc["id"])
for m in session.messages
@@ -2169,8 +2226,30 @@ class AgentLoop:
}
last_assistant_idx: int | None = None
saved_followup_ids: set[str] = set()
for m in messages[skip:]:
entry = dict(m)
checkpoint_boundary = self._validated_checkpoint_boundary(
summary_checkpoint,
skip=skip,
message_count=len(messages),
session_key=session.key,
)
# The trigger input may already be the session tail while still being
# the first message after the replacement checkpoint.
if summary_checkpoint is not None and checkpoint_boundary == skip - 1:
insert_at = len(session.messages) - (1 if input_persisted_early else 0)
self._insert_summary_checkpoint(
session,
summary_checkpoint,
insert_at=insert_at,
)
for message_index, message in enumerate(messages[skip:], start=skip):
# Insert against the raw transcript index before filtering the
# message so persistence cleanup cannot shift the H/Δ boundary.
if summary_checkpoint is not None and checkpoint_boundary == message_index:
self._insert_summary_checkpoint(session, summary_checkpoint)
entry = dict(message)
followup_id_value = cast(object, entry.pop(PENDING_FOLLOWUP_ID_KEY, None))
followup_ids = (
[followup_id_value]
@@ -2249,6 +2328,8 @@ class AgentLoop:
for tc in (cast(dict[str, Any], tc_value),)
if tc.get("id")
)
if summary_checkpoint is not None and checkpoint_boundary == len(messages):
self._insert_summary_checkpoint(session, summary_checkpoint)
if turn_latency_ms is not None and last_assistant_idx is not None:
session.messages[last_assistant_idx]["latency_ms"] = int(turn_latency_ms)
if saved_followup_ids:
nanobot/agent/memory.py
+164 -124
View File
@@ -1,4 +1,4 @@
"""Memory storage, transcript archiving, and legacy consolidation coordination."""
"""Memory storage, transcript archiving, and session checkpoint consolidation."""
# Tool schemas are installed by the ``@tool_parameters`` class decorator at
# runtime; static analyzers cannot observe that it clears ``parameters`` from
@@ -21,6 +21,7 @@ from typing import TYPE_CHECKING, Any, Callable, Iterator, cast
from loguru import logger
from nanobot.llm_usage.context import llm_usage_source
from nanobot.providers.base import ProviderCallContext, ProviderConversationState
from nanobot.runtime_context import public_history_messages
from nanobot.session.manager import (
MIN_COMPACTED_REPLAY_MESSAGES,
@@ -740,7 +741,7 @@ class MemoryStore:
# ---------------------------------------------------------------------------
# Memory ingestion and legacy context-pressure coordination
# Memory ingestion and context-pressure coordination
# ---------------------------------------------------------------------------
# Raw fallbacks use a tighter cap. Completed model summaries may scale with the
@@ -780,6 +781,20 @@ class MemoryArchiver:
) -> str:
"""Persist the failed chunk and return a bounded replacement checkpoint."""
raw = self.store.raw_archive(messages, session_key=session_key)
return self._combine_raw_checkpoint(
raw,
previous_summary=previous_summary,
max_tokens=max_tokens,
)
@staticmethod
def _combine_raw_checkpoint(
raw: str,
*,
previous_summary: str | None,
max_tokens: int,
) -> str:
"""Return a bounded checkpoint that preserves prior and newly archived context."""
token_limit = max(1, max_tokens)
if not previous_summary:
return truncate_text_to_tokens(raw, token_limit)
@@ -806,35 +821,94 @@ class MemoryArchiver:
async def archive(
self,
messages: list[dict[str, Any]],
source_messages: list[dict[str, Any]],
*,
runtime: LLMRuntime,
session_key: str,
request_messages: list[dict[str, Any]],
history: list[dict[str, Any]],
request_tools: list[dict[str, Any]],
previous_summary: str | None = None,
input_token_budget: int | None = None,
fallback_max_tokens: int | None = None,
provider_state: ProviderConversationState | None = None,
) -> str | None:
"""Execute a prepared archive request and persist its result."""
if not messages:
"""Append the archive prompt to H and persist its summary."""
if not source_messages:
return None
def raw_fallback() -> str:
return self._raw_checkpoint(
messages,
source_messages,
session_key=session_key,
previous_summary=previous_summary,
max_tokens=runtime.generation.max_tokens,
max_tokens=(
fallback_max_tokens
if fallback_max_tokens is not None
else runtime.generation.max_tokens
),
)
prompt = render_template(
"agent/consolidator_archive.md",
strip=True,
archive_count=len(source_messages),
)
prompt_message = {"role": "user", "content": prompt}
provider_context = None
call_tools = request_tools
if provider_state is not None:
if not runtime.provider.can_resume_conversation_state(
provider_state,
runtime.model,
):
return raw_fallback()
instruction_messages: list[dict[str, Any]] = []
for message in history:
if message.get("role") not in {"system", "developer"}:
break
instruction_messages.append(dict(message))
request_messages = [*instruction_messages, prompt_message]
provider_context = ProviderCallContext(
conversation_state=provider_state.with_pending_messages([
*provider_state.pending_messages,
prompt_message,
]),
context_window_tokens=runtime.context_window_tokens,
session_id=session_key,
)
call_tools = []
else:
request_messages = [
*[dict(message) for message in history],
prompt_message,
]
if input_token_budget is not None and provider_context is None:
estimated, source = estimate_prompt_tokens_chain(
runtime.provider,
runtime.model,
request_messages,
call_tools,
)
if input_token_budget <= 0 or estimated > input_token_budget:
logger.debug(
"Memory archive input does not fit for {}: {}/{} via {}; raw-dumping",
session_key,
estimated,
input_token_budget,
source,
)
return raw_fallback()
try:
with llm_usage_source("dream"):
response = await runtime.provider.chat_with_retry(
model=runtime.model,
messages=request_messages,
tools=request_tools,
tools=call_tools,
temperature=runtime.generation.temperature,
max_tokens=runtime.generation.max_tokens,
reasoning_effort=runtime.generation.reasoning_effort,
provider_context=provider_context,
)
except Exception:
logger.warning("Memory archive provider call failed, raw-dumping to history")
@@ -879,20 +953,17 @@ class MemoryArchiver:
)
previous_summary = session_summary["text"] if session_summary else None
def raw_fallback() -> str:
if input_token_budget <= 0:
logger.debug(
"Memory archive has no safe input budget for {}; raw-dumping",
session.key,
)
return self._raw_checkpoint(
messages,
session_key=session.key,
previous_summary=previous_summary,
max_tokens=runtime.generation.max_tokens,
)
if input_token_budget <= 0:
logger.debug(
"Memory archive has no safe input budget for {}; raw-dumping",
session.key,
)
return raw_fallback()
prefix = Session(
key=session.key,
messages=list(session.messages[:archive_end]),
@@ -908,47 +979,37 @@ class MemoryArchiver:
"Memory archive cannot replay the full chunk for {}; raw-dumping",
session.key,
)
return raw_fallback()
prompt = render_template("agent/consolidator_archive.md", strip=True)
return self._raw_checkpoint(
messages,
session_key=session.key,
previous_summary=previous_summary,
max_tokens=runtime.generation.max_tokens,
)
channel = session.key.split(":", 1)[0] if ":" in session.key else None
workspace: Path | None = None
if self._resolve_prompt_context is not None:
channel, workspace = self._resolve_prompt_context(session)
request_messages = self._build_messages(
history_messages = self._build_messages(
history=history,
current_message=prompt,
current_message=None,
channel=channel,
session_summary=session_summary,
workspace=workspace,
)
tools = self._get_tool_definitions()
estimated, source = estimate_prompt_tokens_chain(
runtime.provider,
runtime.model,
request_messages,
tools,
)
if estimated > input_token_budget:
logger.debug(
"Memory archive prefix exceeds budget for {}; raw-dumping: {}/{} via {}",
session.key,
estimated,
input_token_budget,
source,
)
return raw_fallback()
return await self.archive(
messages,
runtime=runtime,
session_key=session.key,
request_messages=request_messages,
history=history_messages,
request_tools=tools,
previous_summary=previous_summary,
input_token_budget=input_token_budget,
)
class Consolidator:
"""Legacy context-pressure coordinator backed by a MemoryArchiver."""
"""Coordinate session Memory checkpoints through ``MemoryArchiver``."""
_SAFETY_BUFFER = 1024 # extra headroom for tokenizer estimation drift
@@ -978,22 +1039,73 @@ class Consolidator:
"""Return the shared consolidation lock for one session."""
return self._locks.setdefault(session_key, asyncio.Lock())
def pick_consolidation_boundary(
async def summarize_transcript(
self,
session: Session,
) -> int | None:
"""Return the fixed user-led boundary before the recent replay tail."""
if not session.messages:
accepted_messages: list[dict[str, Any]],
previous_summary: str | None,
*,
runtime: LLMRuntime,
session_key: str,
tools: list[dict[str, Any]],
provider_state: ProviderConversationState | None = None,
) -> str | None:
"""Summarize the exact transcript prefix already accepted by the model."""
source_messages = [
dict(message)
for message in accepted_messages
if message.get("role") != "system"
]
if not source_messages:
return None
boundary = max(0, len(session.messages) - MIN_COMPACTED_REPLAY_MESSAGES)
while boundary > 0 and session.messages[boundary].get("role") != "user":
boundary -= 1
if (
boundary <= session.last_archived
or session.messages[boundary].get("role") != "user"
):
max_output_tokens = max(0, runtime.generation.max_tokens)
input_token_budget = runtime.context_window_tokens - max_output_tokens
checkpoint_tokens = min(
max_output_tokens,
max(1, (input_token_budget - self._SAFETY_BUFFER) // 2),
)
summary = await self.archiver.archive(
source_messages,
runtime=runtime,
session_key=session_key,
history=accepted_messages,
request_tools=tools,
previous_summary=previous_summary,
input_token_budget=input_token_budget,
fallback_max_tokens=max(1, checkpoint_tokens),
provider_state=provider_state,
)
if summary == "(nothing)":
summary = self.archiver._raw_checkpoint(
source_messages,
session_key=session_key,
previous_summary=previous_summary,
max_tokens=max_output_tokens,
)
if summary is None:
return None
return boundary
return truncate_text_to_tokens(summary, max(1, max_output_tokens))
async def summarize_provider_compaction(
self,
state: ProviderConversationState,
fallback_messages: list[dict[str, Any]],
previous_summary: str | None,
*,
runtime: LLMRuntime,
session_key: str,
tools: list[dict[str, Any]],
) -> str | None:
"""Prompt a native compacted state without replaying its raw history."""
return await self.summarize_transcript(
fallback_messages,
previous_summary,
runtime=runtime,
session_key=session_key,
tools=tools,
provider_state=state,
)
@staticmethod
def _full_replay_history(
@@ -1058,7 +1170,7 @@ class Consolidator:
archive_end: int,
runtime: LLMRuntime,
) -> str | None:
"""Compatibility wrapper for the extracted MemoryArchiver."""
"""Archive one captured session range through the shared Memory path."""
return await self.archiver.archive_session(
session,
archive_end=archive_end,
@@ -1066,78 +1178,6 @@ class Consolidator:
input_token_budget=self._input_token_budget(runtime),
)
async def maybe_consolidate_by_tokens(
self,
session: Session,
*,
runtime: LLMRuntime,
) -> None:
"""Archive one fixed old prefix when the prompt exceeds the safe budget.
The budget reserves space for completion tokens and a safety buffer
so the LLM request never exceeds the context window.
"""
lock = self.get_lock(session.key)
async with lock:
# Refresh session reference: AutoCompact may have replaced it.
fresh = self.sessions.get_or_create(session.key)
if fresh is not session:
session = fresh
if runtime.context_window_tokens <= 0:
return
if not session.messages:
return
budget = self._input_token_budget(runtime)
estimated, source = self.estimate_session_prompt_tokens(
session,
runtime=runtime,
)
if estimated <= 0:
return
if estimated < budget:
unarchived_count = len(session.messages) - session.last_archived
logger.debug(
"Token consolidation idle {}: {}/{} via {}, msgs={}",
session.key,
estimated,
runtime.context_window_tokens,
source,
unarchived_count,
)
return
end_idx = self.pick_consolidation_boundary(session)
if end_idx is None:
logger.debug(
"Token consolidation: no safe fixed boundary for {}",
session.key,
)
return
chunk = session.messages[session.last_archived:end_idx]
if not chunk:
return
logger.info(
"Token consolidation for {}: {}/{} via {}, chunk={} msgs",
session.key,
estimated,
runtime.context_window_tokens,
source,
len(chunk),
)
summary = await self.archive_session(
session,
archive_end=end_idx,
runtime=runtime,
)
if summary is None:
return
self._set_last_summary(session, summary)
session.last_archived = end_idx
self.sessions.save(session)
async def compact_idle_session(
self,
session_key: str,
nanobot/agent/runner.py
+81 -205
View File
@@ -16,8 +16,13 @@ from loguru import logger
from nanobot.agent.context import TranscriptInput
from nanobot.agent.context_governance import (
ContextCompactionState,
ContextGovernanceConfig,
ContextGovernor,
HistoryConsolidator,
ModelRequestState,
ProviderCompactionConsolidator,
TranscriptBuilder,
)
from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext
from nanobot.agent.tools.execution import execute_tool_calls
@@ -32,20 +37,10 @@ from nanobot.providers.base import (
LLMProvider,
LLMResponse,
LLMUsage,
ProviderCallContext,
ProviderConversationState,
)
from nanobot.providers.conversation_state import (
ProviderConversationStateController,
allows_conversation_message_merge,
)
from nanobot.runtime_context import (
RUNTIME_CONTEXT_MESSAGE_META,
detach_runtime_context,
reattach_runtime_context,
)
from nanobot.session.history_visibility import is_hidden_history_message
from nanobot.session.recovery import PENDING_FOLLOWUP_ID_KEY
from nanobot.providers.conversation_state import ProviderConversationStateController
from nanobot.session.summary import SessionSummaryCheckpoint
from nanobot.utils.helpers import (
build_assistant_message,
estimate_message_tokens,
@@ -67,7 +62,6 @@ ContinuationCallback = Callable[[], str | None]
RetryWaitCallback = Callable[[str], Awaitable[None]]
CheckpointCallback = Callable[[dict[str, Any]], Awaitable[None]]
InjectionCallback = Callable[..., Awaitable[Iterable[Any] | None]]
TranscriptBuilder = Callable[[TranscriptInput], list[dict[str, Any]]]
_DEFAULT_ERROR_MESSAGE = "Sorry, I encountered an error calling the AI model."
_ARREARAGE_ERROR_MESSAGE = (
@@ -113,6 +107,8 @@ class AgentRunSpec:
provider_retry_mode: str = "standard"
retry_wait_callback: RetryWaitCallback | None = None
checkpoint_callback: CheckpointCallback | None = None
consolidate_history: HistoryConsolidator | None = None
consolidate_provider_compaction: ProviderCompactionConsolidator | None = None
injection_callback: InjectionCallback | None = None
terminal_injection_callback: InjectionCallback | None = None
llm_timeout_s: float | None = None
@@ -137,17 +133,8 @@ class AgentRunResult:
# Terminal tail to emit when the preceding final-content prefix was already streamed.
pending_stream_content: str | None = None
provider_state: ProviderConversationState | None = field(default=None, repr=False)
@dataclass(slots=True)
class _ModelRequestState:
"""Per-run state used to govern the next provider request."""
config: ContextGovernanceConfig
conversation: ProviderConversationStateController
usage: LLMUsage | None = None
messages: list[dict[str, Any]] | None = None
tool_definitions: list[dict[str, Any]] | None = None
summary_checkpoint: SessionSummaryCheckpoint | None = field(default=None, repr=False)
provider_compaction_applied: bool = field(default=False, repr=False)
class AgentRunner:
@@ -157,118 +144,12 @@ class AgentRunner:
self.context_governor = ContextGovernor()
@staticmethod
def _merge_message_content(left: Any, right: Any) -> str | list[dict[str, Any]]:
if isinstance(left, str) and isinstance(right, str):
return f"{left}\n\n{right}" if left else right
def _to_blocks(value: Any) -> list[dict[str, Any]]:
if isinstance(value, list):
return [
cast(dict[str, Any], item)
if isinstance(item, dict)
else {"type": "text", "text": str(item)}
for item in cast(list[Any], value)
]
if value is None:
return []
return [{"type": "text", "text": str(value)}]
return _to_blocks(left) + _to_blocks(right)
@classmethod
def _append_injected_messages(
cls,
messages: list[dict[str, Any]],
injections: list[dict[str, Any]],
) -> None:
"""Append injected user messages while preserving role alternation."""
for injection in injections:
if (
messages
and injection.get("role") == "user"
and messages[-1].get("role") == "user"
and not is_hidden_history_message(injection)
and not is_hidden_history_message(messages[-1])
and allows_conversation_message_merge(messages[-1])
):
merged = dict(messages[-1])
left_meta = merged.get("_meta")
right_meta = injection.get("_meta")
left_meta_dict = cast(dict[str, Any], left_meta) if isinstance(left_meta, dict) else None
right_meta_dict = (
cast(dict[str, Any], right_meta) if isinstance(right_meta, dict) else None
)
left_marker = (
left_meta_dict.get(RUNTIME_CONTEXT_MESSAGE_META)
if left_meta_dict is not None
else None
)
right_marker = (
right_meta_dict.get(RUNTIME_CONTEXT_MESSAGE_META)
if right_meta_dict is not None
else None
)
left_marker_dict = (
cast(dict[str, Any], left_marker) if isinstance(left_marker, dict) else None
)
right_marker_dict = (
cast(dict[str, Any], right_marker) if isinstance(right_marker, dict) else None
)
empty_sources: list[str] = []
empty_blocks: list[dict[str, Any]] = []
detached_left = (
detach_runtime_context(merged.get("content"), left_marker_dict)
if left_marker_dict is not None
else (merged.get("content"), empty_sources, empty_blocks)
)
detached_right = (
detach_runtime_context(injection.get("content"), right_marker_dict)
if right_marker_dict is not None
else (injection.get("content"), empty_sources, empty_blocks)
)
if detached_left is not None and detached_right is not None:
left_content, left_sources, left_blocks = detached_left
right_content, right_sources, right_blocks = detached_right
merged_content = cls._merge_message_content(left_content, right_content)
context_blocks = [*left_blocks, *right_blocks]
if context_blocks:
merged_content, marker = reattach_runtime_context(
merged_content,
[*left_sources, *right_sources],
context_blocks,
)
internal_meta = dict(left_meta_dict) if left_meta_dict is not None else {}
if right_meta_dict is not None:
for key, value in right_meta_dict.items():
internal_meta.setdefault(key, value)
internal_meta[RUNTIME_CONTEXT_MESSAGE_META] = marker
merged["_meta"] = internal_meta
merged["content"] = merged_content
else:
merged["content"] = cls._merge_message_content(
merged.get("content"),
injection.get("content"),
)
followup_id = injection.get(PENDING_FOLLOWUP_ID_KEY)
if isinstance(followup_id, str) and followup_id:
existing = cast(object, merged.get(PENDING_FOLLOWUP_ID_KEY))
followup_ids = (
[existing]
if isinstance(existing, str)
else [
item
for item in cast(list[object], existing)
if isinstance(item, str)
]
if isinstance(existing, list)
else []
)
if followup_id not in followup_ids:
followup_ids.append(followup_id)
merged[PENDING_FOLLOWUP_ID_KEY] = followup_ids
messages[-1] = merged
continue
messages.append(injection)
"""Append injected messages without rewriting the raw transcript."""
messages.extend(injections)
async def _try_drain_injections(
self,
@@ -425,7 +306,7 @@ class AgentRunner:
async def run(self, spec: AgentRunSpec) -> AgentRunResult:
hook = spec.hook or AgentHook()
messages = self._initial_transcript(spec)
messages, compaction = self._initial_transcript_and_compaction(spec)
context = AgentRunHookContext(messages=deepcopy(messages))
llm_usage_source_token = bind_llm_usage_source(
spec.llm_usage_source or source_from_session_key(spec.session_key)
@@ -433,7 +314,7 @@ class AgentRunner:
try:
await hook.before_run(context)
result = await self._run_core(spec, hook, messages)
result = await self._run_core(spec, hook, messages, compaction)
except asyncio.CancelledError as exc:
context.messages = deepcopy(messages)
context.stop_reason = "cancelled"
@@ -478,23 +359,35 @@ class AgentRunner:
reset_llm_usage_source(llm_usage_source_token)
@staticmethod
def _initial_transcript(spec: AgentRunSpec) -> list[dict[str, Any]]:
"""Resolve exactly one supported source for the initial model transcript."""
if spec.transcript_input is not None:
def _initial_transcript_and_compaction(
spec: AgentRunSpec,
) -> tuple[list[dict[str, Any]], ContextCompactionState | None]:
"""Build the initial transcript and its optional compaction state."""
transcript_input = spec.transcript_input
if transcript_input is not None:
if spec.initial_messages is not None:
raise ValueError("provide either transcript_input or initial_messages, not both")
if spec.transcript_builder is None:
transcript_builder = spec.transcript_builder
if transcript_builder is None:
raise ValueError("transcript_builder is required with transcript_input")
return list(spec.transcript_builder(spec.transcript_input))
return ContextCompactionState.from_transcript(
transcript_input,
transcript_builder,
spec.consolidate_history,
spec.consolidate_provider_compaction,
)
if spec.initial_messages is None:
raise ValueError("initial_messages is required without transcript_input")
return list(spec.initial_messages)
if spec.consolidate_history is not None:
raise ValueError("consolidate_history requires transcript_input")
return list(spec.initial_messages), None
async def _run_core(
self,
spec: AgentRunSpec,
hook: AgentHook,
messages: list[dict[str, Any]],
compaction: ContextCompactionState | None,
) -> AgentRunResult:
final_content: str | None = None
tools_used: list[str] = []
@@ -530,9 +423,10 @@ class AgentRunner:
context_block_limit=spec.context_block_limit,
max_tokens=spec.runtime.generation.max_tokens,
)
request_state = _ModelRequestState(
request_state = ModelRequestState(
config=governance_config,
conversation=conversation_state,
compaction=compaction,
)
for iteration in range(spec.max_iterations):
@@ -542,9 +436,15 @@ class AgentRunner:
session_key=spec.session_key,
)
await hook.before_iteration(context)
request_message_count = len(messages)
request_messages = (
request_state.compaction.request_messages(messages)
if request_state.compaction is not None
else messages
)
response = await self._request_model(
spec,
messages,
request_messages,
hook,
context,
request_state=request_state,
@@ -553,6 +453,11 @@ class AgentRunner:
assert request_state.messages is not None
messages_for_model = request_state.messages
conversation_state.observe_response(response, messages)
if request_state.compaction is not None:
request_state.compaction.accept_request(
messages_for_model,
raw_boundary=request_message_count,
)
context.response = response
context.tool_calls = list(response.tool_calls)
@@ -634,7 +539,7 @@ class AgentRunner:
messages.append(tool_message)
completed_tool_results.append(tool_message)
checkpoint_model_messages = (
self.context_governor.prepare_for_model(
self.context_governor.prepare_messages_for_model(
governance_config,
messages,
)
@@ -920,6 +825,12 @@ class AgentRunner:
had_injections=had_injections,
pending_stream_content=pending_stream_content,
provider_state=conversation_state.finish(messages),
summary_checkpoint=(
request_state.compaction.summary_checkpoint
if request_state.compaction is not None
else None
),
provider_compaction_applied=request_state.provider_compaction_applied,
)
def _build_request_kwargs(
@@ -942,60 +853,6 @@ class AgentRunner:
kwargs["reasoning_effort"] = generation.reasoning_effort
return kwargs
def _prepare_model_request(
self,
state: _ModelRequestState,
messages: list[dict[str, Any]],
*,
tool_definitions: list[dict[str, Any]] | None,
transcript: list[dict[str, Any]] | None = None,
) -> tuple[list[dict[str, Any]], ProviderCallContext | None]:
"""Prepare, fit, and record the exact payload sent to a provider."""
prepared = self.context_governor.prepare_for_model(state.config, messages)
supplemental_messages = (
[prepared[-1]] if transcript is not None and tool_definitions is None else None
)
model_messages = None if supplemental_messages is not None else prepared
request_context_tokens = (
state.conversation.estimate_request_context_tokens(
transcript,
model_messages=model_messages,
supplemental_messages=supplemental_messages,
tool_definitions=tool_definitions,
)
if transcript is not None
else None
)
usage_matches_messages = (
state.messages is not None
and prepared == state.messages
and tool_definitions == state.tool_definitions
)
prepared, fitted = self.context_governor.fit_request(
state.config,
prepared,
state.usage,
usage_matches_messages=usage_matches_messages,
tool_definitions=tool_definitions,
request_context_tokens=request_context_tokens,
)
provider_context = (
state.conversation.prepare_request(
transcript,
context_window_tokens=state.config.context_window_tokens,
model_messages=model_messages,
supplemental_messages=supplemental_messages,
resume_state=not fitted,
)
if transcript is not None
else state.conversation.independent_request_context(
context_window_tokens=state.config.context_window_tokens,
)
)
state.messages = deepcopy(prepared)
state.tool_definitions = deepcopy(tool_definitions)
return prepared, provider_context
async def _request_model(
self,
spec: AgentRunSpec,
@@ -1003,13 +860,13 @@ class AgentRunner:
hook: AgentHook,
context: AgentHookContext,
*,
request_state: _ModelRequestState,
request_state: ModelRequestState,
malformed_retry: bool = False,
transcript: list[dict[str, Any]] | None,
) -> LLMResponse:
timeout_s = self._resolve_llm_timeout_s(spec)
tool_definitions = spec.tools.get_definitions()
messages, provider_context = self._prepare_model_request(
messages, provider_context = await self.context_governor.prepare_request(
request_state,
messages,
tool_definitions=tool_definitions,
@@ -1169,6 +1026,12 @@ class AgentRunner:
response.ttft_ms = max(0, round((first_output_at - request_started_at) * 1000))
if generation_elapsed_s > 0:
response.generation_ms = max(1, round(generation_elapsed_s * 1000))
await self.context_governor.summarize_provider_compaction(
request_state,
response,
current_request_boundary=(len(transcript) if transcript is not None else None),
)
request_state.provider_compaction_applied |= response.provider_compaction_applied
# chat_stream_with_retry may recover internally, so only fail unfinished
# hosted calls after the provider returns its final error response.
if response.finish_reason == "error":
@@ -1283,7 +1146,7 @@ class AgentRunner:
spec: AgentRunSpec,
messages: list[dict[str, Any]],
*,
request_state: _ModelRequestState,
request_state: ModelRequestState,
transcript: list[dict[str, Any]],
) -> LLMResponse:
retry_messages = self._finalization_retry_messages(messages)
@@ -1313,14 +1176,21 @@ class AgentRunner:
messages: list[dict[str, Any]],
usage: LLMUsage | None,
*,
request_state: _ModelRequestState,
request_state: ModelRequestState,
) -> tuple[str | None, LLMUsage | None]:
retry_messages = self._budget_exhausted_finalization_messages(messages)
compaction = request_state.compaction
request_messages = (
compaction.request_messages(messages)
if compaction is not None
else messages
)
retry_messages = self._budget_exhausted_finalization_messages(request_messages)
try:
response = await self._request_no_tools(
spec,
retry_messages,
request_state=request_state,
transcript=messages if compaction is not None else None,
)
except Exception:
logger.exception(
@@ -1358,10 +1228,10 @@ class AgentRunner:
spec: AgentRunSpec,
messages: list[dict[str, Any]],
*,
request_state: _ModelRequestState,
request_state: ModelRequestState,
transcript: list[dict[str, Any]] | None = None,
) -> LLMResponse:
messages, provider_context = self._prepare_model_request(
messages, provider_context = await self.context_governor.prepare_request(
request_state,
messages,
tool_definitions=None,
@@ -1389,6 +1259,12 @@ class AgentRunner:
finish_reason="error",
error_kind="timeout",
)
await self.context_governor.summarize_provider_compaction(
request_state,
response,
current_request_boundary=(len(transcript) if transcript is not None else None),
)
request_state.provider_compaction_applied |= response.provider_compaction_applied
return response
@staticmethod
@@ -1453,7 +1329,7 @@ class AgentRunner:
def _record_request_usage(
self,
spec: AgentRunSpec,
state: _ModelRequestState,
state: ModelRequestState,
response: LLMResponse,
) -> LLMUsage | None:
assert state.messages is not None
nanobot/cli/gateway_runtime.py
-5
View File
@@ -662,11 +662,6 @@ def _run_gateway(
if isinstance(message_tool, MessageTool) and suppress_token is not None:
message_tool.reset_suppress_delivery(suppress_token)
# Keep a small tail of heartbeat history so the loop stays bounded.
session = agent.sessions.get_or_create("heartbeat")
session.retain_recent_legal_suffix(hb_cfg.keep_recent_messages)
agent.sessions.save(session)
if not resp or not resp.content:
return
nanobot/config/schema.py
-1
View File
@@ -329,7 +329,6 @@ class HeartbeatConfig(Base):
enabled: bool = True
interval_s: int = 30 * 60 # 30 minutes
keep_recent_messages: int = 8
class ApiConfig(Base):
nanobot/providers/azure_openai_provider.py
+11
View File
@@ -34,6 +34,7 @@ from nanobot.providers.base import (
)
from nanobot.providers.openai_responses import (
ResponsesStreamCapture,
build_responses_compaction_state,
build_responses_state,
consume_sdk_stream,
convert_tools,
@@ -410,6 +411,16 @@ class AzureOpenAIProvider(LLMProvider):
output_items=capture.output_items,
usage=usage,
)
result.provider_compaction_state = build_responses_compaction_state(
provider=self._responses_state_provider(),
model=str(body["model"]),
output_items=capture.output_items,
)
result.provider_compaction_applied = (
result.provider_compaction_state is not None
)
if result.provider_compaction_applied:
result.provider_compaction_scope = "current_request"
return result
except Exception as e:
return self._handle_error(e)
nanobot/providers/base.py
+17
View File
@@ -31,6 +31,7 @@ RETRY_AFTER_BUFFER = 1
RetryEventCallback = Callable[[str], Awaitable[None]]
LLMCallObserver = Callable[["LLMCallRecord"], None]
ProviderCompactionScope = Literal["prior_context", "current_request"]
def resolve_stream_idle_timeout_s(
@@ -563,6 +564,22 @@ class LLMResponse:
reasoning_content: str | None = None # Kimi, DeepSeek-R1, MiMo etc.
thinking_blocks: list[dict[str, Any]] | None = None # Anthropic extended thinking
provider_state: ProviderConversationState | None = field(default=None, repr=False)
# True only when this response installed a new provider-native compaction
# boundary. Replaying an older compaction item does not set this flag.
provider_compaction_applied: bool = field(default=False, repr=False)
# State immediately after native compaction, before the normal response
# continues. An archive prompt can resume this state without replaying H.
provider_compaction_state: ProviderConversationState | None = field(
default=None,
repr=False,
)
# Which model input the native compaction state replaces. Providers that
# compact before attaching the current request delta report
# ``prior_context``; in-request compaction reports ``current_request``.
provider_compaction_scope: ProviderCompactionScope | None = field(
default=None,
repr=False,
)
# Routing wrappers may preserve or discard an incoming provider-owned
# continuation independently of the final fallback error's retry policy.
preserve_provider_state_on_error: bool | None = field(default=None, repr=False)
nanobot/providers/conversation_state.py
+6
View File
@@ -101,6 +101,12 @@ class ProviderConversationStateController:
)
return context_tokens + max(0, delta_tokens)
def replace_transcript(self, messages: list[dict[str, Any]]) -> None:
"""Discard append-only provider state after a transcript rewrite."""
self._state = None
self._boundary = len(messages)
self._request_messages = []
def prepare_request(
self,
messages: list[dict[str, Any]],
nanobot/providers/openai_codex_provider.py
+22 -3
View File
@@ -31,6 +31,7 @@ from nanobot.providers.oauth_model_catalog import (
)
from nanobot.providers.openai_responses import (
ResponsesStreamCapture,
build_responses_compaction_state,
build_responses_state,
consume_sse_with_reasoning,
convert_tools,
@@ -137,6 +138,8 @@ class OpenAICodexProvider(LLMProvider):
body.update(self._extra_body)
stage = "oauth_token"
native_compaction_applied = False
native_compaction_state: ProviderConversationState | None = None
try:
token = await asyncio.to_thread(get_codex_token, proxy=self.proxy)
headers = _build_headers(cast(str, token.account_id), token.access)
@@ -187,9 +190,11 @@ class OpenAICodexProvider(LLMProvider):
and responses_state_context_tokens(sanitized_state) >= compact_threshold
):
stage = "codex_compaction"
history_items = responses_state_items(sanitized_state) or []
delta_items = input_items[len(history_items):]
compact_body = {
**body,
"input": [*input_items, {"type": "compaction_trigger"}],
"input": [*history_items, {"type": "compaction_trigger"}],
}
try:
compact_result = await _send(compact_body, emit_deltas=False)
@@ -205,9 +210,16 @@ class OpenAICodexProvider(LLMProvider):
}:
raise RuntimeError("Codex compaction returned no compaction item")
body["input"] = [
*_retained_compaction_messages(input_items),
*_retained_compaction_messages(history_items),
*compact_items,
*delta_items,
]
native_compaction_state = build_responses_compaction_state(
provider=self._responses_state_provider(),
model=_strip_model_prefix(model),
output_items=compact_items,
)
native_compaction_applied = True
except Exception as compact_error:
if is_compaction_compatibility_error(compact_error):
self._native_compaction_available = False
@@ -220,7 +232,14 @@ class OpenAICodexProvider(LLMProvider):
)
stage = "codex_request"
return await _send(body, emit_deltas=True)
result = await _send(body, emit_deltas=True)
result.provider_compaction_applied = (
result.provider_compaction_applied or native_compaction_applied
)
if native_compaction_state is not None:
result.provider_compaction_state = native_compaction_state
result.provider_compaction_scope = "prior_context"
return result
except Exception as e:
response = _codex_error_response(e)
exc_type = "CodexHTTPError" if isinstance(e, _CodexHTTPError) else type(e).__name__
nanobot/providers/openai_compat_provider.py
+13
View File
@@ -36,6 +36,7 @@ from nanobot.providers.base import (
)
from nanobot.providers.openai_responses import (
ResponsesStreamCapture,
build_responses_compaction_state,
build_responses_state,
consume_sdk_stream,
convert_tools,
@@ -2049,6 +2050,18 @@ class OpenAICompatProvider(LLMProvider):
output_items=capture.output_items,
usage=usage,
)
result.provider_compaction_state = (
build_responses_compaction_state(
provider=self._responses_state_provider(),
model=str(body["model"]),
output_items=capture.output_items,
)
)
result.provider_compaction_applied = (
result.provider_compaction_state is not None
)
if result.provider_compaction_applied:
result.provider_compaction_scope = "current_request"
return result
except Exception as responses_error:
if self._spec and self._spec.name == "github_copilot":
nanobot/providers/openai_responses/__init__.py
+2
View File
@@ -18,6 +18,7 @@ from nanobot.providers.openai_responses.parsing import (
parse_response_output,
)
from nanobot.providers.openai_responses.state import (
build_responses_compaction_state,
build_responses_state,
is_compaction_compatibility_error,
prepare_responses_input,
@@ -40,6 +41,7 @@ __all__ = [
"is_replayable_finish_reason",
"map_finish_reason",
"parse_response_output",
"build_responses_compaction_state",
"build_responses_state",
"is_compaction_compatibility_error",
"prepare_responses_input",
nanobot/providers/openai_responses/parsing.py
+12 -1
View File
@@ -11,7 +11,10 @@ import httpx
from loguru import logger
from nanobot.providers.base import LLMResponse, LLMUsage, ToolCallRequest, parse_tool_arguments
from nanobot.providers.openai_responses.state import build_responses_state
from nanobot.providers.openai_responses.state import (
build_responses_compaction_state,
build_responses_state,
)
FINISH_REASON_MAP = {
"completed": "stop",
@@ -655,6 +658,14 @@ def parse_response_output(
output_items=output,
usage=usage,
)
result.provider_compaction_state = build_responses_compaction_state(
provider=state_provider,
model=state_model,
output_items=output,
)
result.provider_compaction_applied = result.provider_compaction_state is not None
if result.provider_compaction_applied:
result.provider_compaction_scope = "current_request"
return result
nanobot/providers/openai_responses/state.py
+22
View File
@@ -108,6 +108,28 @@ def build_responses_state(
)
def build_responses_compaction_state(
*,
provider: str,
model: str,
output_items: list[dict[str, Any]],
) -> ProviderConversationState | None:
"""Return the state at the latest native compaction output boundary."""
latest = None
for index, item in enumerate(output_items):
if item.get("type") in _COMPACTION_ITEM_TYPES:
latest = index
if latest is None:
return None
return ProviderConversationState(
kind=RESPONSES_STATE_KIND,
provider=provider,
model=model,
version=RESPONSES_STATE_VERSION,
payload={_ITEMS_KEY: [deepcopy(output_items[latest])]},
)
def responses_state_items(
state: ProviderConversationState,
) -> list[dict[str, Any]] | None:
nanobot/sdk/clients.py
+3 -3
View File
@@ -209,11 +209,11 @@ class RuntimeClient:
return self._loop.runtime_events.subscribe(handler, SessionTurnPersisted)
async def compact_session(self, session_key: str) -> SessionSnapshot:
"""Run token consolidation for one session."""
"""Archive one session through the shared idle-compaction path."""
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,
await self._loop.consolidator.compact_idle_session(
session_key,
runtime=runtime,
)
return snapshot_from_session(self._loop.sessions.get_or_create(session_key))
nanobot/session/manager.py
+14 -121
View File
@@ -27,7 +27,9 @@ from nanobot.runtime_context import (
RUNTIME_CONTEXT_HISTORY_META,
public_history_message,
)
from nanobot.session.history_visibility import is_hidden_history_message
from nanobot.session.model_selection import SESSION_MODEL_PRESET_METADATA_KEY
from nanobot.session.summary import SUMMARY_CONTINUATION_TEXT
from nanobot.utils.helpers import (
content_with_media_breadcrumbs,
ensure_dir,
@@ -262,12 +264,6 @@ def _metadata_title(metadata: object) -> str:
return strip_think(title)
@dataclass
class RetentionResult:
dropped: list[dict[str, Any]]
already_consolidated_count: int
@dataclass(frozen=True)
class SessionPolicy:
"""Runtime rules that do not belong in durable session data."""
@@ -286,9 +282,7 @@ class Session:
created_at: datetime = field(default_factory=datetime.now)
updated_at: datetime = field(default_factory=datetime.now)
metadata: dict[str, Any] = field(default_factory=dict)
# Legacy storage name for the Memory ingestion watermark. New code should
# use ``last_archived`` so this progress is not confused with model-context
# compaction. Keep the field while persisted sessions and SDK callers migrate.
# Keep the legacy storage name while persisted sessions and SDK callers migrate.
last_consolidated: int = 0
provider_state: ProviderConversationState | None = field(default=None, repr=False)
policy: SessionPolicy = field(default_factory=SessionPolicy, repr=False, compare=False)
@@ -309,7 +303,7 @@ class Session:
@property
def last_archived(self) -> int:
"""Number of transcript messages already written to the Memory journal."""
"""End of the latest committed Memory checkpoint."""
return self.last_consolidated
@last_archived.setter
@@ -337,14 +331,17 @@ class Session:
) -> list[dict[str, Any]]:
"""Return recent replayable messages for LLM input.
A positive ``max_messages`` applies an explicit caller-owned count
limit. The normal model path relies on ``max_tokens`` instead.
A committed in-turn checkpoint replaces its old prefix with the stored
summary and resumes replay at a hidden continuation marker. A positive
``max_messages`` applies an additional caller-owned count limit.
"""
replay_start = self.last_archived
if replay_start:
# ``last_archived`` is archive progress, not a replay boundary.
# Keep a small raw suffix for continuity, extending back to the user
# that started an assistant/tool sequence when necessary.
resumes_from_checkpoint = (
replay_start < len(self.messages)
and is_hidden_history_message(self.messages[replay_start])
and self.messages[replay_start].get("content") == SUMMARY_CONTINUATION_TEXT
)
if replay_start and not resumes_from_checkpoint:
recent_start = recent_message_start_index(
self.messages,
MIN_COMPACTED_REPLAY_MESSAGES,
@@ -485,110 +482,6 @@ class Session:
self.updated_at = datetime.now()
self.metadata.pop("_last_summary", None)
def retain_recent_legal_suffix(
self,
max_messages: int,
*,
extend_to_user: bool = False,
) -> RetentionResult:
"""Keep a legal recent suffix, optionally extending it back to a user turn.
Returns a RetentionResult with dropped messages and how many of those
were in the already-consolidated prefix. This method mutates
self.messages and self.last_archived in place.
"""
if max_messages <= 0:
dropped = list(self.messages)
lc = self.last_archived
self.clear()
return RetentionResult(
dropped=dropped,
already_consolidated_count=min(lc, len(dropped)),
)
if len(self.messages) <= max_messages:
return RetentionResult(
dropped=[],
already_consolidated_count=0,
)
original = list(self.messages)
before_lc = self.last_archived
start_idx = max(0, len(self.messages) - max_messages)
if extend_to_user:
recovered_user = next(
(i for i in range(start_idx, -1, -1) if self.messages[i].get("role") == "user"),
None,
)
if recovered_user is not None:
start_idx = recovered_user
if start_idx > 0 and self.messages[start_idx - 1].get("_channel_delivery"):
start_idx -= 1
retained = self.messages[start_idx:]
# Prefer starting at a user turn (or its preceding _channel_delivery) when one exists within the retained window.
first_user = next((i for i, m in enumerate(retained) if m.get("role") == "user"), None)
if first_user is not None:
if first_user > 0 and retained[first_user - 1].get("_channel_delivery"):
retained = retained[first_user - 1:]
else:
retained = retained[first_user:]
elif not extend_to_user:
# If the hard-capped tail is assistant/tool-only, anchor to the
# latest user in the full session and take a capped forward window.
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 = 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 unless the caller requested user-turn extension.
if not extend_to_user and len(retained) > max_messages:
retained = retained[-max_messages:]
start = find_legal_message_start(retained)
if start:
retained = retained[start:]
# Compute actually-dropped messages using identity comparison so that
# even when retained is a non-contiguous slice of original (the else
# branch above), we never duplicate or lose messages.
retained_ids = set(id(m) for m in retained)
dropped = [m for m in original if id(m) not in retained_ids]
# Count how many dropped messages were in the already-consolidated
# prefix of the original list. This cannot be a simple min() because
# dropped may include messages from *after* the consolidated prefix
# (e.g. in the else branch).
already_consolidated = sum(
1 for i, m in enumerate(original)
if i < before_lc and id(m) not in retained_ids
)
# New last_archived = count of retained messages that were inside
# the old consolidated prefix.
new_lc = sum(
1 for i, m in enumerate(original)
if i < before_lc and id(m) in retained_ids
)
self.messages = retained
self.last_archived = new_lc
if dropped:
self.provider_state = None
self.updated_at = datetime.now()
return RetentionResult(
dropped=dropped,
already_consolidated_count=already_consolidated,
)
class SessionPayload(TypedDict):
key: str
created_at: str | None
@@ -2010,7 +1903,7 @@ class SessionManager:
user_index = 0
found_target = False
for message in source.messages:
if message.get("role") == "user":
if message.get("role") == "user" and not is_hidden_history_message(message):
if user_index == before_user_index:
found_target = True
break
nanobot/session/summary.py
+12
View File
@@ -3,15 +3,27 @@
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime
from typing import TypedDict, cast
SUMMARY_CONTINUATION_TEXT = (
"Continue the active task from the working-memory checkpoint above."
)
class SessionSummary(TypedDict):
text: str
last_active: str
@dataclass(frozen=True, slots=True)
class SessionSummaryCheckpoint:
"""A replacement summary and the raw transcript boundary it covers."""
summary: str
transcript_boundary: int
def session_summary_from_metadata(
metadata: Mapping[str, object] | None,
*,
nanobot/webui/settings_system.py
-1
View File
@@ -114,7 +114,6 @@ def system_settings_payload(
"heartbeat": {
"enabled": config.gateway.heartbeat.enabled,
"interval_s": config.gateway.heartbeat.interval_s,
"keep_recent_messages": config.gateway.heartbeat.keep_recent_messages,
},
"dream": {
"schedule": defaults.dream.describe_schedule(),