diff --git a/nanobot/agent/context_governance.py b/nanobot/agent/context_governance.py new file mode 100644 index 00000000..148b40e6 --- /dev/null +++ b/nanobot/agent/context_governance.py @@ -0,0 +1,391 @@ +"""Model-message governance 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. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from loguru import logger + +from nanobot.utils.helpers import ( + estimate_message_tokens, + estimate_prompt_tokens_chain, + find_legal_message_start, + maybe_persist_tool_result, + truncate_text, +) +from nanobot.utils.runtime import ensure_nonempty_tool_result + +if TYPE_CHECKING: + from nanobot.providers.base import LLMProvider + +SNIP_SAFETY_BUFFER = 1024 +MICROCOMPACT_KEEP_RECENT = 10 +MICROCOMPACT_MIN_CHARS = 500 +INFLIGHT_COMPACT_TARGET_RATIO = 0.85 +COMPACTABLE_TOOLS = frozenset({ + "read_file", "exec", "grep", "find_files", + "web_search", "web_fetch", "list_dir", "list_exec_sessions", +}) +# read_file is the recovery path for persisted results; exempting it prevents persist->read->persist loops. +TOOL_RESULT_OFFLOAD_EXEMPT_TOOLS = frozenset({"read_file"}) +BACKFILL_CONTENT = "[Tool result unavailable — call was interrupted or lost]" + + +@dataclass(slots=True) +class ContextGovernanceConfig: + provider: LLMProvider + model: str + tools: Any + workspace: Path | None + session_key: str | None + max_tool_result_chars: int + context_window_tokens: int | None = None + context_block_limit: int | None = None + max_tokens: int | None = None + inflight_start_index: int = 0 + + +class ContextGovernor: + """Prepare model-copy messages while preserving persisted history.""" + + def prepare_for_model( + self, + config: ContextGovernanceConfig, + messages: list[dict[str, Any]], + compacted_tool_call_ids: set[str], + ) -> list[dict[str, Any]]: + updated = self.drop_orphan_tool_results(messages) + updated = self.backfill_missing_tool_results(updated) + updated = self.apply_tool_result_budget(config, updated) + updated = self.compact_inflight_overflow(config, updated, compacted_tool_call_ids) + updated = self.snip_history(config, updated) + updated = self.drop_orphan_tool_results(updated) + return self.backfill_missing_tool_results(updated) + + @staticmethod + def input_budget(config: ContextGovernanceConfig) -> int: + if not config.context_window_tokens: + return 0 + + provider_max_tokens = getattr( + getattr(config.provider, "generation", None), + "max_tokens", + 4096, + ) + max_output = config.max_tokens if isinstance(config.max_tokens, int) else ( + provider_max_tokens if isinstance(provider_max_tokens, int) else 4096 + ) + budget = config.context_block_limit or ( + config.context_window_tokens - max_output - SNIP_SAFETY_BUFFER + ) + return budget if budget > 0 else 0 + + @staticmethod + def normalize_tool_result( + config: ContextGovernanceConfig, + tool_call_id: str, + tool_name: str, + result: Any, + ) -> Any: + result = ensure_nonempty_tool_result(tool_name, result) + if tool_name in TOOL_RESULT_OFFLOAD_EXEMPT_TOOLS: + return result + try: + content = maybe_persist_tool_result( + config.workspace, + config.session_key, + tool_call_id, + result, + max_chars=config.max_tool_result_chars, + ) + except Exception: + logger.exception( + "Tool result persist failed for {} in {}; using raw result", + tool_call_id, + config.session_key or "default", + ) + content = result + if isinstance(content, str) and len(content) > config.max_tool_result_chars: + return truncate_text(content, config.max_tool_result_chars) + return content + + @staticmethod + def drop_orphan_tool_results( + messages: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + """Drop tool results that have no matching assistant tool_call earlier in history.""" + declared: set[str] = set() + updated: list[dict[str, Any]] | None = None + for idx, msg in enumerate(messages): + role = msg.get("role") + if role == "assistant": + for tc in msg.get("tool_calls") or []: + if isinstance(tc, dict) and tc.get("id"): + declared.add(str(tc["id"])) + if role == "tool": + tid = msg.get("tool_call_id") + if tid and str(tid) not in declared: + if updated is None: + updated = [dict(m) for m in messages[:idx]] + continue + if updated is not None: + updated.append(dict(msg)) + + if updated is None: + return messages + return updated + + @staticmethod + def backfill_missing_tool_results( + messages: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + """Insert synthetic error results for assistant tool_calls with missing tool outputs.""" + declared: list[tuple[int, str, str]] = [] + fulfilled: set[str] = set() + for idx, msg in enumerate(messages): + role = msg.get("role") + if role == "assistant": + for tc in msg.get("tool_calls") or []: + if isinstance(tc, dict) and tc.get("id"): + name = "" + func = tc.get("function") + if isinstance(func, dict): + name = func.get("name", "") + declared.append((idx, str(tc["id"]), name)) + elif role == "tool": + tid = msg.get("tool_call_id") + if tid: + fulfilled.add(str(tid)) + + missing = [(ai, cid, name) for ai, cid, name in declared if cid not in fulfilled] + if not missing: + return messages + + updated = list(messages) + offset = 0 + for assistant_idx, call_id, name in missing: + insert_at = assistant_idx + 1 + offset + while insert_at < len(updated) and updated[insert_at].get("role") == "tool": + insert_at += 1 + updated.insert(insert_at, { + "role": "tool", + "tool_call_id": call_id, + "name": name, + "content": BACKFILL_CONTENT, + }) + offset += 1 + return updated + + def apply_tool_result_budget( + self, + config: ContextGovernanceConfig, + messages: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + updated = messages + for idx, message in enumerate(messages): + if message.get("role") != "tool": + continue + normalized = self.normalize_tool_result( + config, + str(message.get("tool_call_id") or f"tool_{idx}"), + str(message.get("name") or "tool"), + message.get("content"), + ) + if normalized != message.get("content"): + if updated is messages: + updated = [dict(m) for m in messages] + updated[idx]["content"] = normalized + return updated + + def compact_inflight_overflow( + self, + config: ContextGovernanceConfig, + messages: list[dict[str, Any]], + compacted_tool_call_ids: set[str], + ) -> list[dict[str, Any]]: + """Compact in-flight tool results only when the request would overflow.""" + budget = self.input_budget(config) + if budget <= 0: + return messages + + tools = config.tools.get_definitions() + updated = self._apply_recorded_compactions(messages, compacted_tool_call_ids) + estimate, source = estimate_prompt_tokens_chain( + config.provider, + config.model, + updated, + tools, + ) + if estimate <= budget: + return updated + + target = int(budget * INFLIGHT_COMPACT_TARGET_RATIO) + candidates = self._inflight_compaction_candidates( + config, + updated, + compacted_tool_call_ids, + ) + if not candidates: + return updated + + for candidate_idx, (idx, tool_call_id) in enumerate(candidates): + is_newest_candidate = candidate_idx == len(candidates) - 1 + if is_newest_candidate and estimate <= budget: + break + if tool_call_id in compacted_tool_call_ids: + continue + if updated is messages: + updated = [dict(m) for m in messages] + compacted_tool_call_ids.add(tool_call_id) + self._compact_tool_result_at(updated, idx) + estimate, source = estimate_prompt_tokens_chain( + config.provider, + config.model, + updated, + tools, + ) + if estimate <= target: + break + + logger.debug( + "In-flight context compaction for {}: prompt={} budget={} target={} via {}, ids={}", + config.session_key or "default", + estimate, + budget, + target, + source, + len(compacted_tool_call_ids), + ) + return updated + + def snip_history( + self, + config: ContextGovernanceConfig, + messages: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + if not messages or not config.context_window_tokens: + return messages + + budget = self.input_budget(config) + if budget <= 0: + return messages + + tools = config.tools.get_definitions() + estimate, _ = estimate_prompt_tokens_chain( + config.provider, + config.model, + messages, + tools, + ) + 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"] + if not non_system: + return messages + + system_tokens = sum(estimate_message_tokens(msg) for msg in system_messages) + fixed_tokens, _ = estimate_prompt_tokens_chain( + config.provider, + config.model, + system_messages, + tools, + ) + remaining_budget = max(0, budget - max(system_tokens, fixed_tokens)) + kept: list[dict[str, Any]] = [] + kept_tokens = 0 + for message in reversed(non_system): + msg_tokens = estimate_message_tokens(message) + if kept and kept_tokens + msg_tokens > remaining_budget: + break + kept.append(message) + kept_tokens += msg_tokens + kept.reverse() + + return system_messages + self._legal_history_tail(kept, non_system) + + @staticmethod + def _summary_for(message: dict[str, Any]) -> str: + name = message.get("name", "tool") + return f"[{name} result omitted from context]" + + def _legal_history_tail( + self, + kept: list[dict[str, Any]], + non_system: list[dict[str, Any]], + ) -> list[dict[str, Any]]: + fallback = kept if kept else (non_system[-1:] if non_system else []) + kept = self._user_tail(kept) or self._user_tail(non_system, last=True) or fallback + + start = find_legal_message_start(kept) + return kept[start:] if start else kept + + @staticmethod + def _user_tail(messages: list[dict[str, Any]], *, last: bool = False) -> list[dict[str, Any]]: + indexes = range(len(messages) - 1, -1, -1) if last else range(len(messages)) + for idx in indexes: + if messages[idx].get("role") == "user": + return messages[idx:] + return [] + + def _apply_recorded_compactions( + self, + messages: list[dict[str, Any]], + compacted_tool_call_ids: set[str], + ) -> list[dict[str, Any]]: + if not compacted_tool_call_ids: + return messages + updated = messages + for idx, msg in enumerate(messages): + if msg.get("role") != "tool": + continue + tool_call_id = msg.get("tool_call_id") + if not tool_call_id or str(tool_call_id) not in compacted_tool_call_ids: + continue + summary = self._summary_for(msg) + if msg.get("content") == summary: + continue + if updated is messages: + updated = [dict(m) for m in messages] + updated[idx]["content"] = summary + return updated + + def _inflight_compaction_candidates( + self, + config: ContextGovernanceConfig, + messages: list[dict[str, Any]], + compacted_tool_call_ids: set[str], + ) -> list[tuple[int, str]]: + compactable: list[tuple[int, str]] = [] + for idx, msg in enumerate(messages): + if idx < config.inflight_start_index: + continue + if msg.get("role") != "tool" or msg.get("name") not in COMPACTABLE_TOOLS: + continue + tool_call_id = msg.get("tool_call_id") + if not tool_call_id or str(tool_call_id) in compacted_tool_call_ids: + continue + content = msg.get("content") + if not isinstance(content, str) or len(content) < MICROCOMPACT_MIN_CHARS: + continue + compactable.append((idx, str(tool_call_id))) + + if not compactable: + return [] + primary_count = max(0, len(compactable) - MICROCOMPACT_KEEP_RECENT) + primary = compactable[:primary_count] + # Hard overflow beats the keep-recent preference. Return recent results + # after stale ones so the newest result is naturally last. + fallback = compactable[primary_count:] + return primary + fallback + + def _compact_tool_result_at(self, messages: list[dict[str, Any]], idx: int) -> None: + messages[idx]["content"] = self._summary_for(messages[idx]) diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index d124156f..436a1a12 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -13,6 +13,10 @@ from typing import Any, Callable from loguru import logger +from nanobot.agent.context_governance import ( + ContextGovernanceConfig, + ContextGovernor, +) from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext from nanobot.agent.tools.registry import ToolRegistry from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest @@ -32,11 +36,8 @@ from nanobot.utils.helpers import ( estimate_message_tokens, estimate_prompt_tokens_chain, extract_reasoning, - find_legal_message_start, - maybe_persist_tool_result, strip_reasoning_tags, strip_think, - truncate_text, ) from nanobot.utils.progress_events import ( invoke_file_edit_progress, @@ -49,7 +50,6 @@ from nanobot.utils.runtime import ( build_finalization_retry_message, build_goal_continue_message, build_length_recovery_message, - ensure_nonempty_tool_result, is_blank_text, repeated_external_lookup_error, repeated_workspace_violation_error, @@ -67,17 +67,6 @@ _MAX_EMPTY_RETRIES = 2 _MAX_LENGTH_RECOVERIES = 3 _MAX_INJECTIONS_PER_TURN = 3 _MAX_INJECTION_CYCLES = 5 -_SNIP_SAFETY_BUFFER = 1024 -_MICROCOMPACT_KEEP_RECENT = 10 -_MICROCOMPACT_MIN_CHARS = 500 -_COMPACTABLE_TOOLS = frozenset({ - "read_file", "exec", "grep", "find_files", - "web_search", "web_fetch", "list_dir", "list_exec_sessions", -}) -# read_file is the recovery path for persisted results; exempting it prevents persist->read->persist loops. -_TOOL_RESULT_OFFLOAD_EXEMPT_TOOLS = frozenset({"read_file"}) -_BACKFILL_CONTENT = "[Tool result unavailable — call was interrupted or lost]" - # Backward-compatible module attribute for tests/extensions that monkeypatch # the former single-file tracker hook. Runtime uses prepare_file_edit_trackers. prepare_file_edit_tracker = _prepare_file_edit_tracker @@ -135,6 +124,7 @@ class AgentRunner: def __init__(self, provider: LLMProvider): self.provider = provider + self.context_governor = ContextGovernor() @staticmethod def _merge_message_content(left: Any, right: Any) -> str | list[dict[str, Any]]: @@ -367,6 +357,19 @@ class AgentRunner: length_recovery_count = 0 had_injections = False injection_cycles = 0 + compacted_tool_call_ids: set[str] = set() + governance_config = ContextGovernanceConfig( + provider=self.provider, + model=spec.model, + tools=spec.tools, + workspace=spec.workspace, + session_key=spec.session_key, + max_tool_result_chars=spec.max_tool_result_chars, + context_window_tokens=spec.context_window_tokens, + context_block_limit=spec.context_block_limit, + max_tokens=spec.max_tokens, + inflight_start_index=len(spec.initial_messages), + ) for iteration in range(spec.max_iterations): try: @@ -374,14 +377,11 @@ class AgentRunner: # may repair or compact historical messages for the model, but # those synthetic edits must not shift the append boundary used # later when the caller saves only the new turn. - messages_for_model = self._drop_orphan_tool_results(messages) - messages_for_model = self._backfill_missing_tool_results(messages_for_model) - messages_for_model = self._microcompact(messages_for_model) - messages_for_model = self._apply_tool_result_budget(spec, messages_for_model) - messages_for_model = self._snip_history(spec, messages_for_model) - # Snipping may have created new orphans; clean them up. - messages_for_model = self._drop_orphan_tool_results(messages_for_model) - messages_for_model = self._backfill_missing_tool_results(messages_for_model) + messages_for_model = self.context_governor.prepare_for_model( + governance_config, + messages, + compacted_tool_call_ids, + ) except Exception: logger.exception( "Context governance failed on turn {} for {}; applying minimal repair", @@ -389,8 +389,10 @@ class AgentRunner: spec.session_key or "default", ) try: - messages_for_model = self._drop_orphan_tool_results(messages) - messages_for_model = self._backfill_missing_tool_results(messages_for_model) + messages_for_model = ContextGovernor.drop_orphan_tool_results(messages) + messages_for_model = ContextGovernor.backfill_missing_tool_results( + messages_for_model + ) except Exception: messages_for_model = messages context = AgentHookContext( @@ -463,8 +465,8 @@ class AgentRunner: "role": "tool", "tool_call_id": tool_call.id, "name": tool_call.name, - "content": self._normalize_tool_result( - spec, + "content": self.context_governor.normalize_tool_result( + governance_config, tool_call.id, tool_call.name, result, @@ -1334,225 +1336,6 @@ class AgentRunner: return messages.append(build_assistant_message(_PERSISTED_MODEL_ERROR_PLACEHOLDER)) - def _normalize_tool_result( - self, - spec: AgentRunSpec, - tool_call_id: str, - tool_name: str, - result: Any, - ) -> Any: - result = ensure_nonempty_tool_result(tool_name, result) - if tool_name in _TOOL_RESULT_OFFLOAD_EXEMPT_TOOLS: - # Exempt tools bound their own output; skip generic offload and truncation. - return result - try: - content = maybe_persist_tool_result( - spec.workspace, - spec.session_key, - tool_call_id, - result, - max_chars=spec.max_tool_result_chars, - ) - except Exception: - logger.exception( - "Tool result persist failed for {} in {}; using raw result", - tool_call_id, - spec.session_key or "default", - ) - content = result - if isinstance(content, str) and len(content) > spec.max_tool_result_chars: - return truncate_text(content, spec.max_tool_result_chars) - return content - - @staticmethod - def _drop_orphan_tool_results( - messages: list[dict[str, Any]], - ) -> list[dict[str, Any]]: - """Drop tool results that have no matching assistant tool_call earlier in the history.""" - declared: set[str] = set() - updated: list[dict[str, Any]] | None = None - for idx, msg in enumerate(messages): - role = msg.get("role") - if role == "assistant": - for tc in msg.get("tool_calls") or []: - if isinstance(tc, dict) and tc.get("id"): - declared.add(str(tc["id"])) - if role == "tool": - tid = msg.get("tool_call_id") - if tid and str(tid) not in declared: - if updated is None: - updated = [dict(m) for m in messages[:idx]] - continue - if updated is not None: - updated.append(dict(msg)) - - if updated is None: - return messages - return updated - - @staticmethod - def _backfill_missing_tool_results( - messages: list[dict[str, Any]], - ) -> list[dict[str, Any]]: - """Insert synthetic error results for orphaned tool_use blocks.""" - declared: list[tuple[int, str, str]] = [] # (assistant_idx, call_id, name) - fulfilled: set[str] = set() - for idx, msg in enumerate(messages): - role = msg.get("role") - if role == "assistant": - for tc in msg.get("tool_calls") or []: - if isinstance(tc, dict) and tc.get("id"): - name = "" - func = tc.get("function") - if isinstance(func, dict): - name = func.get("name", "") - declared.append((idx, str(tc["id"]), name)) - elif role == "tool": - tid = msg.get("tool_call_id") - if tid: - fulfilled.add(str(tid)) - - missing = [(ai, cid, name) for ai, cid, name in declared if cid not in fulfilled] - if not missing: - return messages - - updated = list(messages) - offset = 0 - for assistant_idx, call_id, name in missing: - insert_at = assistant_idx + 1 + offset - while insert_at < len(updated) and updated[insert_at].get("role") == "tool": - insert_at += 1 - updated.insert(insert_at, { - "role": "tool", - "tool_call_id": call_id, - "name": name, - "content": _BACKFILL_CONTENT, - }) - offset += 1 - return updated - - @staticmethod - def _microcompact(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: - """Replace old compactable tool results with one-line summaries.""" - compactable_indices: list[int] = [] - for idx, msg in enumerate(messages): - if msg.get("role") == "tool" and msg.get("name") in _COMPACTABLE_TOOLS: - compactable_indices.append(idx) - - if len(compactable_indices) <= _MICROCOMPACT_KEEP_RECENT: - return messages - - stale = compactable_indices[: len(compactable_indices) - _MICROCOMPACT_KEEP_RECENT] - updated: list[dict[str, Any]] | None = None - for idx in stale: - msg = messages[idx] - content = msg.get("content") - if not isinstance(content, str) or len(content) < _MICROCOMPACT_MIN_CHARS: - continue - name = msg.get("name", "tool") - summary = f"[{name} result omitted from context]" - if updated is None: - updated = [dict(m) for m in messages] - updated[idx]["content"] = summary - - return updated if updated is not None else messages - - def _apply_tool_result_budget( - self, - spec: AgentRunSpec, - messages: list[dict[str, Any]], - ) -> list[dict[str, Any]]: - updated = messages - for idx, message in enumerate(messages): - if message.get("role") != "tool": - continue - normalized = self._normalize_tool_result( - spec, - str(message.get("tool_call_id") or f"tool_{idx}"), - str(message.get("name") or "tool"), - message.get("content"), - ) - if normalized != message.get("content"): - if updated is messages: - updated = [dict(m) for m in messages] - updated[idx]["content"] = normalized - return updated - - def _snip_history( - self, - spec: AgentRunSpec, - messages: list[dict[str, Any]], - ) -> list[dict[str, Any]]: - if not messages or not spec.context_window_tokens: - return messages - - provider_max_tokens = getattr(getattr(self.provider, "generation", None), "max_tokens", 4096) - max_output = spec.max_tokens if isinstance(spec.max_tokens, int) else ( - provider_max_tokens if isinstance(provider_max_tokens, int) else 4096 - ) - budget = spec.context_block_limit or ( - spec.context_window_tokens - max_output - _SNIP_SAFETY_BUFFER - ) - if budget <= 0: - return messages - - estimate, _ = estimate_prompt_tokens_chain( - self.provider, - spec.model, - messages, - spec.tools.get_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"] - if not non_system: - return messages - - system_tokens = sum(estimate_message_tokens(msg) for msg in system_messages) - fixed_tokens, _ = estimate_prompt_tokens_chain( - self.provider, - spec.model, - system_messages, - spec.tools.get_definitions(), - ) - remaining_budget = max(0, budget - max(system_tokens, fixed_tokens)) - kept: list[dict[str, Any]] = [] - kept_tokens = 0 - for message in reversed(non_system): - msg_tokens = estimate_message_tokens(message) - if kept and kept_tokens + msg_tokens > remaining_budget: - break - kept.append(message) - kept_tokens += msg_tokens - kept.reverse() - - if kept: - for i, message in enumerate(kept): - if message.get("role") == "user": - kept = kept[i:] - break - else: - # Recover nearest user message from outside the kept window; - # GLM rejects system→assistant (error 1214). Budget is - # intentionally exceeded — oversized beats invalid. - for idx in range(len(non_system) - 1, -1, -1): - if non_system[idx].get("role") == "user": - kept = non_system[idx:] - break - # If no user exists at all, _enforce_role_alternation - # will insert a synthetic one as a safety net. - start = find_legal_message_start(kept) - if start: - kept = kept[start:] - if not kept: - kept = non_system[-min(len(non_system), 4) :] - start = find_legal_message_start(kept) - if start: - kept = kept[start:] - return system_messages + kept - def _partition_tool_batches( self, spec: AgentRunSpec, diff --git a/tests/agent/test_runner_governance.py b/tests/agent/test_runner_governance.py index 901afc71..01f1c210 100644 --- a/tests/agent/test_runner_governance.py +++ b/tests/agent/test_runner_governance.py @@ -2,16 +2,45 @@ from __future__ import annotations +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest +from nanobot.agent.context_governance import ( + BACKFILL_CONTENT, + MICROCOMPACT_KEEP_RECENT, + ContextGovernanceConfig, + ContextGovernor, +) +from nanobot.agent.runner import AgentRunSpec from nanobot.config.schema import AgentDefaults -from nanobot.providers.base import LLMResponse, ToolCallRequest +from nanobot.providers.base import LLMResponse _MAX_TOOL_RESULT_CHARS = AgentDefaults().max_tool_result_chars +def _governance_config( + provider, + tools, + spec: AgentRunSpec, + *, + inflight_start_index: int = 0, +) -> ContextGovernanceConfig: + return ContextGovernanceConfig( + provider=provider, + model=spec.model, + tools=tools, + workspace=spec.workspace, + session_key=spec.session_key, + max_tool_result_chars=spec.max_tool_result_chars, + context_window_tokens=spec.context_window_tokens, + context_block_limit=spec.context_block_limit, + max_tokens=spec.max_tokens, + inflight_start_index=inflight_start_index, + ) + + def _make_loop(tmp_path): from nanobot.agent.loop import AgentLoop from nanobot.bus.queue import MessageBus @@ -22,13 +51,14 @@ def _make_loop(tmp_path): with patch("nanobot.agent.loop.ContextBuilder"), \ patch("nanobot.agent.loop.SessionManager"), \ - patch("nanobot.agent.loop.SubagentManager") as MockSubMgr: - MockSubMgr.return_value.cancel_by_session = AsyncMock(return_value=0) + patch("nanobot.agent.loop.SubagentManager") as mock_sub_mgr: + mock_sub_mgr.return_value.cancel_by_session = AsyncMock(return_value=0) loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path) return loop + async def test_runner_uses_raw_messages_when_context_governance_fails(): - from nanobot.agent.runner import AgentRunSpec, AgentRunner + from nanobot.agent.runner import AgentRunner provider = MagicMock() captured_messages: list[dict] = [] @@ -46,7 +76,9 @@ async def test_runner_uses_raw_messages_when_context_governance_fails(): ] runner = AgentRunner(provider) - runner._snip_history = MagicMock(side_effect=RuntimeError("boom")) # type: ignore[method-assign] + runner.context_governor.prepare_for_model = MagicMock( # type: ignore[method-assign] + side_effect=RuntimeError("boom") + ) result = await runner.run(AgentRunSpec( initial_messages=initial_messages, tools=tools, @@ -57,13 +89,12 @@ async def test_runner_uses_raw_messages_when_context_governance_fails(): assert result.final_content == "done" assert captured_messages == initial_messages -def test_snip_history_drops_orphaned_tool_results_from_trimmed_slice(monkeypatch): - from nanobot.agent.runner import AgentRunSpec, AgentRunner + +def test_snip_history_drops_orphaned_tool_results_from_trimmed_slice(monkeypatch): provider = MagicMock() tools = MagicMock() tools.get_definitions.return_value = [] - runner = AgentRunner(provider) messages = [ {"role": "system", "content": "system"}, {"role": "user", "content": "old user"}, @@ -85,7 +116,10 @@ def test_snip_history_drops_orphaned_tool_results_from_trimmed_slice(monkeypatch context_block_limit=100, ) - monkeypatch.setattr("nanobot.agent.runner.estimate_prompt_tokens_chain", lambda *_args, **_kwargs: (500, None)) + monkeypatch.setattr( + "nanobot.agent.context_governance.estimate_prompt_tokens_chain", + lambda *_args, **_kwargs: (500, None), + ) token_sizes = { "old user": 120, "tool call": 120, @@ -94,11 +128,11 @@ def test_snip_history_drops_orphaned_tool_results_from_trimmed_slice(monkeypatch "system": 0, } monkeypatch.setattr( - "nanobot.agent.runner.estimate_message_tokens", + "nanobot.agent.context_governance.estimate_message_tokens", lambda msg: token_sizes.get(str(msg.get("content")), 40), ) - trimmed = runner._snip_history(spec, messages) + trimmed = ContextGovernor().snip_history(_governance_config(provider, tools, spec), messages) # After the fix, the user message is recovered so the sequence is valid # for providers that require system → user (e.g. GLM error 1214). @@ -108,12 +142,9 @@ def test_snip_history_drops_orphaned_tool_results_from_trimmed_slice(monkeypatch def test_snip_history_reserves_budget_for_tool_definitions(monkeypatch): - from nanobot.agent.runner import AgentRunSpec, AgentRunner - provider = MagicMock() tools = MagicMock() tools.get_definitions.return_value = [{"type": "function", "function": {"name": "large_tool"}}] - runner = AgentRunner(provider) messages = [ {"role": "system", "content": "system"}, {"role": "user", "content": "old user"}, @@ -139,7 +170,7 @@ def test_snip_history_reserves_budget_for_tool_definitions(monkeypatch): assert estimate_tools == tools.get_definitions.return_value return 350, None - monkeypatch.setattr("nanobot.agent.runner.estimate_prompt_tokens_chain", _estimate) + monkeypatch.setattr("nanobot.agent.context_governance.estimate_prompt_tokens_chain", _estimate) token_sizes = { "system": 50, "old user": 200, @@ -149,11 +180,11 @@ def test_snip_history_reserves_budget_for_tool_definitions(monkeypatch): "recent two": 200, } monkeypatch.setattr( - "nanobot.agent.runner.estimate_message_tokens", + "nanobot.agent.context_governance.estimate_message_tokens", lambda msg: token_sizes.get(str(msg.get("content")), 40), ) - trimmed = runner._snip_history(spec, messages) + trimmed = ContextGovernor().snip_history(_governance_config(provider, tools, spec), messages) contents = [message.get("content") for message in trimmed] assert contents == ["system", "recent two"] @@ -161,7 +192,6 @@ def test_snip_history_reserves_budget_for_tool_definitions(monkeypatch): async def test_backfill_missing_tool_results_inserts_error(): """Orphaned tool_use (no matching tool_result) should get a synthetic error.""" - from nanobot.agent.runner import AgentRunner, _BACKFILL_CONTENT messages = [ {"role": "user", "content": "hi"}, @@ -175,18 +205,16 @@ async def test_backfill_missing_tool_results_inserts_error(): }, {"role": "tool", "tool_call_id": "call_a", "name": "exec", "content": "ok"}, ] - result = AgentRunner._backfill_missing_tool_results(messages) + result = ContextGovernor.backfill_missing_tool_results(messages) tool_msgs = [m for m in result if m.get("role") == "tool"] assert len(tool_msgs) == 2 backfilled = [m for m in tool_msgs if m.get("tool_call_id") == "call_b"] assert len(backfilled) == 1 - assert backfilled[0]["content"] == _BACKFILL_CONTENT + assert backfilled[0]["content"] == BACKFILL_CONTENT assert backfilled[0]["name"] == "read_file" def test_drop_orphan_tool_results_removes_unmatched_tool_messages(): - from nanobot.agent.runner import AgentRunner - messages = [ {"role": "system", "content": "system"}, {"role": "user", "content": "old user"}, @@ -202,7 +230,7 @@ def test_drop_orphan_tool_results_removes_unmatched_tool_messages(): {"role": "assistant", "content": "after tool"}, ] - cleaned = AgentRunner._drop_orphan_tool_results(messages) + cleaned = ContextGovernor.drop_orphan_tool_results(messages) assert cleaned == [ {"role": "system", "content": "system"}, @@ -222,8 +250,6 @@ def test_drop_orphan_tool_results_removes_unmatched_tool_messages(): @pytest.mark.asyncio async def test_backfill_noop_when_complete(): """Complete message chains should not be modified.""" - from nanobot.agent.runner import AgentRunner - messages = [ {"role": "user", "content": "hi"}, { @@ -236,13 +262,13 @@ async def test_backfill_noop_when_complete(): {"role": "tool", "tool_call_id": "call_x", "name": "exec", "content": "done"}, {"role": "assistant", "content": "all good"}, ] - result = AgentRunner._backfill_missing_tool_results(messages) + result = ContextGovernor.backfill_missing_tool_results(messages) assert result is messages # same object — no copy @pytest.mark.asyncio async def test_runner_drops_orphan_tool_results_before_model_request(): - from nanobot.agent.runner import AgentRunSpec, AgentRunner + from nanobot.agent.runner import AgentRunner provider = MagicMock() captured_messages: list[dict] = [] @@ -283,7 +309,6 @@ async def test_runner_drops_orphan_tool_results_before_model_request(): async def test_backfill_repairs_model_context_without_shifting_save_turn_boundary(tmp_path): """Historical backfill should not duplicate old tail messages on persist.""" from nanobot.agent.loop import AgentLoop - from nanobot.agent.runner import _BACKFILL_CONTENT from nanobot.bus.events import InboundMessage from nanobot.bus.queue import MessageBus @@ -335,7 +360,7 @@ async def test_backfill_repairs_model_context_without_shifting_save_turn_boundar if message.get("role") == "tool" and message.get("tool_call_id") == "call_missing" ] assert len(synthetic) == 1 - assert synthetic[0]["content"] == _BACKFILL_CONTENT + assert synthetic[0]["content"] == BACKFILL_CONTENT session_after = loop.sessions.get_or_create("cli:test") assert [ @@ -367,7 +392,7 @@ async def test_backfill_repairs_model_context_without_shifting_save_turn_boundar @pytest.mark.asyncio async def test_runner_backfill_only_mutates_model_context_not_returned_messages(): """Runner should repair orphaned tool calls for the model without rewriting result.messages.""" - from nanobot.agent.runner import AgentRunSpec, AgentRunner, _BACKFILL_CONTENT + from nanobot.agent.runner import AgentRunner provider = MagicMock() captured_messages: list[dict] = [] @@ -413,7 +438,7 @@ async def test_runner_backfill_only_mutates_model_context_not_returned_messages( if message.get("role") == "tool" and message.get("tool_call_id") == "call_missing" ] assert len(synthetic) == 1 - assert synthetic[0]["content"] == _BACKFILL_CONTENT + assert synthetic[0]["content"] == BACKFILL_CONTENT assert [ { @@ -447,96 +472,254 @@ async def test_runner_backfill_only_mutates_model_context_not_returned_messages( # --------------------------------------------------------------------------- -@pytest.mark.asyncio -async def test_microcompact_replaces_old_tool_results(): - """Tool results beyond _MICROCOMPACT_KEEP_RECENT should be summarized.""" - from nanobot.agent.runner import AgentRunner, _MICROCOMPACT_KEEP_RECENT - - total = _MICROCOMPACT_KEEP_RECENT + 5 - long_content = "x" * 600 +def _microcompact_messages(*, total: int, tool_name: str, content: str) -> list[dict]: messages: list[dict] = [{"role": "system", "content": "sys"}] for i in range(total): messages.append({ "role": "assistant", "content": "", - "tool_calls": [{"id": f"c{i}", "type": "function", "function": {"name": "read_file", "arguments": "{}"}}], + "tool_calls": [{ + "id": f"c{i}", + "type": "function", + "function": {"name": tool_name, "arguments": "{}"}, + }], }) messages.append({ - "role": "tool", "tool_call_id": f"c{i}", "name": "read_file", - "content": long_content, + "role": "tool", + "tool_call_id": f"c{i}", + "name": tool_name, + "content": content, }) + return messages - result = AgentRunner._microcompact(messages) + +def test_microcompact_skips_when_prompt_under_hard_budget(monkeypatch): + """Cache-friendly path: in-flight tool results stay stable while prompt fits.""" + provider = MagicMock() + provider.generation = SimpleNamespace(max_tokens=0) + tools = MagicMock() + tools.get_definitions.return_value = [] + + total = MICROCOMPACT_KEEP_RECENT + 5 + long_content = "x" * 600 + messages = _microcompact_messages(total=total, tool_name="read_file", content=long_content) + spec = AgentRunSpec( + initial_messages=messages, + tools=tools, + model="test-model", + max_iterations=1, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + max_tokens=0, + context_window_tokens=20_000, + ) + + monkeypatch.setattr( + "nanobot.agent.context_governance.estimate_prompt_tokens_chain", + lambda *_args, **_kwargs: (1000, "test"), + ) + + result = ContextGovernor().compact_inflight_overflow( + _governance_config(provider, tools, spec), + messages, + set(), + ) + + assert result is messages + + +def test_microcompact_overflow_compacts_to_low_watermark(monkeypatch): + """Overflow path: compact in-flight stale results with headroom for later calls.""" + provider = MagicMock() + provider.generation = SimpleNamespace(max_tokens=0) + tools = MagicMock() + tools.get_definitions.return_value = [] + + total = MICROCOMPACT_KEEP_RECENT + 8 + long_content = "x" * 600 + messages = _microcompact_messages(total=total, tool_name="read_file", content=long_content) + spec = AgentRunSpec( + initial_messages=messages, + tools=tools, + model="test-model", + max_iterations=1, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + max_tokens=0, + context_window_tokens=2224, # input budget 1200, low target 1020 + ) + + def estimate(_provider, _model, msgs, _tools): + return sum( + 100 if (content := msg.get("content")) == long_content + else 1 if isinstance(content, str) and "omitted from context" in content + else 0 + for msg in msgs + if msg.get("role") == "tool" + ), "test" + + monkeypatch.setattr("nanobot.agent.context_governance.estimate_prompt_tokens_chain", estimate) + + result = ContextGovernor().compact_inflight_overflow( + _governance_config(provider, tools, spec), + messages, + set(), + ) tool_msgs = [m for m in result if m.get("role") == "tool"] - stale_count = total - _MICROCOMPACT_KEEP_RECENT compacted = [m for m in tool_msgs if "omitted from context" in str(m.get("content", ""))] preserved = [m for m in tool_msgs if m.get("content") == long_content] - assert len(compacted) == stale_count - assert len(preserved) == _MICROCOMPACT_KEEP_RECENT + + assert len(compacted) == 8 + assert len(preserved) == total - 8 + assert [m["tool_call_id"] for m in compacted] == [f"c{i}" for i in range(8)] -@pytest.mark.asyncio -async def test_microcompact_preserves_short_results(): - """Short tool results (< _MICROCOMPACT_MIN_CHARS) should not be replaced.""" - from nanobot.agent.runner import AgentRunner, _MICROCOMPACT_KEEP_RECENT +def test_microcompact_compacts_newest_when_it_alone_overflows(monkeypatch): + """The newest result is preserved only while the request can still fit.""" + provider = MagicMock() + provider.generation = SimpleNamespace(max_tokens=0) + tools = MagicMock() + tools.get_definitions.return_value = [] - total = _MICROCOMPACT_KEEP_RECENT + 5 - messages: list[dict] = [] - for i in range(total): - messages.append({ - "role": "assistant", - "content": "", - "tool_calls": [{"id": f"c{i}", "type": "function", "function": {"name": "exec", "arguments": "{}"}}], - }) - messages.append({ - "role": "tool", "tool_call_id": f"c{i}", "name": "exec", - "content": "short", - }) + long_content = "x" * 600 + messages = _microcompact_messages(total=1, tool_name="read_file", content=long_content) + spec = AgentRunSpec( + initial_messages=messages, + tools=tools, + model="test-model", + max_iterations=1, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + max_tokens=0, + context_window_tokens=2000, + context_block_limit=500, + ) - result = AgentRunner._microcompact(messages) + def estimate(_provider, _model, msgs, _tools): + return sum( + 1000 if msg.get("content") == long_content else 1 + for msg in msgs + if msg.get("role") == "tool" + ), "test" + + monkeypatch.setattr("nanobot.agent.context_governance.estimate_prompt_tokens_chain", estimate) + + compacted_tool_call_ids: set[str] = set() + result = ContextGovernor().compact_inflight_overflow( + _governance_config(provider, tools, spec), + messages, + compacted_tool_call_ids, + ) + + tool_msg = next(m for m in result if m.get("role") == "tool") + assert "omitted from context" in tool_msg["content"] + assert compacted_tool_call_ids == {"c0"} + + +def test_context_governor_keeps_compaction_boundary_stable(monkeypatch): + provider = MagicMock() + provider.generation = SimpleNamespace(max_tokens=0) + tools = MagicMock() + tools.get_definitions.return_value = [] + + total = MICROCOMPACT_KEEP_RECENT + 8 + long_content = "x" * 600 + messages = _microcompact_messages(total=total, tool_name="read_file", content=long_content) + spec = AgentRunSpec( + initial_messages=messages, + tools=tools, + model="test-model", + max_iterations=1, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + max_tokens=0, + context_window_tokens=2224, + ) + + def estimate(_provider, _model, msgs, _tools): + return sum( + 100 if msg.get("content") == long_content else 1 + for msg in msgs + if msg.get("role") == "tool" + ), "test" + + monkeypatch.setattr("nanobot.agent.context_governance.estimate_prompt_tokens_chain", estimate) + + governor = ContextGovernor() + compacted_tool_call_ids: set[str] = set() + config = _governance_config(provider, tools, spec, inflight_start_index=0) + first = governor.compact_inflight_overflow(config, messages, compacted_tool_call_ids) + first_ids = set(compacted_tool_call_ids) + + second = governor.compact_inflight_overflow(config, messages, compacted_tool_call_ids) + + assert compacted_tool_call_ids == first_ids + assert [m.get("content") for m in second] == [m.get("content") for m in first] + + +def test_microcompact_preserves_short_results(monkeypatch): + """Short tool results below the compaction threshold should not be replaced.""" + provider = MagicMock() + provider.generation = SimpleNamespace(max_tokens=0) + tools = MagicMock() + tools.get_definitions.return_value = [] + + total = MICROCOMPACT_KEEP_RECENT + 5 + messages = _microcompact_messages(total=total, tool_name="exec", content="short") + spec = AgentRunSpec( + initial_messages=messages, + tools=tools, + model="test-model", + max_iterations=1, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + max_tokens=0, + context_window_tokens=2024, + ) + + monkeypatch.setattr( + "nanobot.agent.context_governance.estimate_prompt_tokens_chain", + lambda *_args, **_kwargs: (2000, "test"), + ) + + result = ContextGovernor().compact_inflight_overflow( + _governance_config(provider, tools, spec), + messages, + set(), + ) assert result is messages # no copy needed — all stale results are short -@pytest.mark.asyncio -async def test_microcompact_skips_non_compactable_tools(): +def test_microcompact_skips_non_compactable_tools(monkeypatch): """Non-compactable tools (e.g. 'message') should never be replaced.""" - from nanobot.agent.runner import AgentRunner, _MICROCOMPACT_KEEP_RECENT + provider = MagicMock() + provider.generation = SimpleNamespace(max_tokens=0) + tools = MagicMock() + tools.get_definitions.return_value = [] - total = _MICROCOMPACT_KEEP_RECENT + 5 + total = MICROCOMPACT_KEEP_RECENT + 5 long_content = "y" * 1000 - messages: list[dict] = [] - for i in range(total): - messages.append({ - "role": "assistant", - "content": "", - "tool_calls": [{"id": f"c{i}", "type": "function", "function": {"name": "message", "arguments": "{}"}}], - }) - messages.append({ - "role": "tool", "tool_call_id": f"c{i}", "name": "message", - "content": long_content, - }) + messages = _microcompact_messages(total=total, tool_name="message", content=long_content) + spec = AgentRunSpec( + initial_messages=messages, + tools=tools, + model="test-model", + max_iterations=1, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + max_tokens=0, + context_window_tokens=2024, + ) - result = AgentRunner._microcompact(messages) + monkeypatch.setattr( + "nanobot.agent.context_governance.estimate_prompt_tokens_chain", + lambda *_args, **_kwargs: (2000, "test"), + ) + + result = ContextGovernor().compact_inflight_overflow( + _governance_config(provider, tools, spec), + messages, + set(), + ) assert result is messages # no compactable tools found def test_governance_repairs_orphans_after_snip(): - """After _snip_history clips an assistant+tool_calls, the second - _drop_orphan_tool_results pass must clean up the resulting orphans.""" - from nanobot.agent.runner import AgentRunner - - messages = [ - {"role": "system", "content": "system"}, - {"role": "user", "content": "old msg"}, - {"role": "assistant", "content": None, - "tool_calls": [{"id": "tc_old", "type": "function", - "function": {"name": "search", "arguments": "{}"}}]}, - {"role": "tool", "tool_call_id": "tc_old", "name": "search", - "content": "old result"}, - {"role": "assistant", "content": "old answer"}, - {"role": "user", "content": "new msg"}, - ] - + """After snipping clips an assistant+tool_calls, orphan repair cleans up the tail.""" # Simulate snipping that keeps only the tail: drop the assistant with # tool_calls but keep its tool result (orphan). snipped = [ @@ -547,7 +730,7 @@ def test_governance_repairs_orphans_after_snip(): {"role": "user", "content": "new msg"}, ] - cleaned = AgentRunner._drop_orphan_tool_results(snipped) + cleaned = ContextGovernor.drop_orphan_tool_results(snipped) # The orphan tool result should be removed. assert not any( m.get("role") == "tool" and m.get("tool_call_id") == "tc_old" @@ -556,10 +739,7 @@ def test_governance_repairs_orphans_after_snip(): def test_governance_fallback_still_repairs_orphans(): - """When full governance fails, the fallback must still run - _drop_orphan_tool_results and _backfill_missing_tool_results.""" - from nanobot.agent.runner import AgentRunner - + """When full governance fails, the fallback must still repair orphans.""" # Messages with an orphan tool result (no matching assistant tool_call). messages = [ {"role": "user", "content": "hello"}, @@ -568,10 +748,12 @@ def test_governance_fallback_still_repairs_orphans(): {"role": "assistant", "content": "hi"}, ] - repaired = AgentRunner._drop_orphan_tool_results(messages) - repaired = AgentRunner._backfill_missing_tool_results(repaired) + repaired = ContextGovernor.drop_orphan_tool_results(messages) + repaired = ContextGovernor.backfill_missing_tool_results(repaired) # Orphan tool result should be gone. assert not any(m.get("tool_call_id") == "orphan_tc" for m in repaired) + + def test_snip_history_preserves_user_message_after_truncation(monkeypatch): """When _snip_history truncates messages and the only user message ends up outside the kept window, the method must recover the nearest user message @@ -585,12 +767,9 @@ def test_snip_history_preserves_user_message_after_truncation(monkeypatch): - _snip_history activates, keeping only recent assistant/tool pairs. - The injected user message is in the truncated prefix and gets lost. """ - from nanobot.agent.runner import AgentRunSpec, AgentRunner - provider = MagicMock() tools = MagicMock() tools.get_definitions.return_value = [] - runner = AgentRunner(provider) messages = [ {"role": "system", "content": "system"}, @@ -621,7 +800,10 @@ def test_snip_history_preserves_user_message_after_truncation(monkeypatch): ) # Make estimate_prompt_tokens_chain report above budget so _snip_history activates. - monkeypatch.setattr("nanobot.agent.runner.estimate_prompt_tokens_chain", lambda *_a, **_kw: (500, None)) + monkeypatch.setattr( + "nanobot.agent.context_governance.estimate_prompt_tokens_chain", + lambda *_a, **_kw: (500, None), + ) # Make kept window small: only the last 2 messages fit the budget. token_sizes = { "system": 0, @@ -631,11 +813,11 @@ def test_snip_history_preserves_user_message_after_truncation(monkeypatch): "tool output 2": 80, } monkeypatch.setattr( - "nanobot.agent.runner.estimate_message_tokens", + "nanobot.agent.context_governance.estimate_message_tokens", lambda msg: token_sizes.get(str(msg.get("content")), 100), ) - trimmed = runner._snip_history(spec, messages) + trimmed = ContextGovernor().snip_history(_governance_config(provider, tools, spec), messages) # The first non-system message MUST be user (not assistant). non_system = [m for m in trimmed if m.get("role") != "system"] @@ -649,12 +831,9 @@ def test_snip_history_preserves_user_message_after_truncation(monkeypatch): def test_snip_history_no_user_at_all_falls_back_gracefully(monkeypatch): """Edge case: if non_system has zero user messages, _snip_history should still return a valid sequence (not crash or produce system→assistant).""" - from nanobot.agent.runner import AgentRunSpec, AgentRunner - provider = MagicMock() tools = MagicMock() tools.get_definitions.return_value = [] - runner = AgentRunner(provider) messages = [ {"role": "system", "content": "system"}, @@ -674,13 +853,16 @@ def test_snip_history_no_user_at_all_falls_back_gracefully(monkeypatch): context_block_limit=100, ) - monkeypatch.setattr("nanobot.agent.runner.estimate_prompt_tokens_chain", lambda *_a, **_kw: (500, None)) monkeypatch.setattr( - "nanobot.agent.runner.estimate_message_tokens", + "nanobot.agent.context_governance.estimate_prompt_tokens_chain", + lambda *_a, **_kw: (500, None), + ) + monkeypatch.setattr( + "nanobot.agent.context_governance.estimate_message_tokens", lambda msg: 100, ) - trimmed = runner._snip_history(spec, messages) + trimmed = ContextGovernor().snip_history(_governance_config(provider, tools, spec), messages) # Should not crash. The result should still be a valid list. assert isinstance(trimmed, list) diff --git a/tests/agent/test_runner_persistence.py b/tests/agent/test_runner_persistence.py index 3c943175..24ac1051 100644 --- a/tests/agent/test_runner_persistence.py +++ b/tests/agent/test_runner_persistence.py @@ -6,15 +6,13 @@ import os import time from unittest.mock import AsyncMock, MagicMock, patch -import pytest - from nanobot.config.schema import AgentDefaults from nanobot.providers.base import LLMResponse, ToolCallRequest _MAX_TOOL_RESULT_CHARS = AgentDefaults().max_tool_result_chars async def test_runner_persists_large_tool_results_for_follow_up_calls(tmp_path): - from nanobot.agent.runner import AgentRunSpec, AgentRunner + from nanobot.agent.runner import AgentRunner, AgentRunSpec provider = MagicMock() captured_second_call: list[dict] = [] @@ -172,7 +170,7 @@ async def test_read_file_result_is_not_offloaded(tmp_path): async def test_runner_keeps_going_when_tool_result_persistence_fails(): - from nanobot.agent.runner import AgentRunSpec, AgentRunner + from nanobot.agent.runner import AgentRunner, AgentRunSpec provider = MagicMock() captured_second_call: list[dict] = [] @@ -195,7 +193,10 @@ async def test_runner_keeps_going_when_tool_result_persistence_fails(): tools.execute = AsyncMock(return_value="tool result") runner = AgentRunner(provider) - with patch("nanobot.agent.runner.maybe_persist_tool_result", side_effect=RuntimeError("disk full")): + with patch( + "nanobot.agent.context_governance.maybe_persist_tool_result", + side_effect=RuntimeError("disk full"), + ): result = await runner.run(AgentRunSpec( initial_messages=[{"role": "user", "content": "do task"}], tools=tools,