From 6cd7063682e16aceeb45b637bb0141adb8974fd9 Mon Sep 17 00:00:00 2001 From: chengyongru <61816729+chengyongru@users.noreply.github.com> Date: Mon, 31 Aug 2026 00:06:15 +0800 Subject: [PATCH] refactor(agent): defer transcript assembly to runner (#5608) * refactor(agent): defer transcript assembly to runner Keep persisted history and the fresh turn as explicit inputs until the Runner assembles the provider transcript. Preserve ContextBuilder and direct AgentRunner compatibility while making the save boundary structural. Refs NAN-81. * fix(providers): preserve mixed adjacent user content --- nanobot/agent/context.py | 93 ++++++++++++++----- nanobot/agent/loop.py | 38 ++++---- nanobot/agent/runner.py | 23 ++++- nanobot/providers/base.py | 21 +++++ nanobot/session/turn_continuation.py | 10 +- tests/agent/test_attachment_references.py | 6 +- tests/agent/test_context_builder.py | 25 ++++- tests/agent/test_hook_composite.py | 14 ++- tests/agent/test_loop_progress.py | 19 +++- tests/agent/test_loop_runner_integration.py | 21 +++-- tests/agent/test_loop_save_turn.py | 69 +++++++++----- tests/agent/test_loop_tool_context.py | 5 +- tests/agent/test_runner_core.py | 30 ++++++ tests/agent/test_runner_injections.py | 12 ++- tests/agent/tools/test_subagent_tools.py | 12 ++- tests/cli/test_restart_command.py | 11 ++- .../test_enforce_role_alternation.py | 45 ++++++++- tests/session/test_turn_continuation.py | 9 +- tests/tools/test_message_tool_suppress.py | 5 +- 19 files changed, 355 insertions(+), 113 deletions(-) diff --git a/nanobot/agent/context.py b/nanobot/agent/context.py index a2aad8351..0c8f604ce 100644 --- a/nanobot/agent/context.py +++ b/nanobot/agent/context.py @@ -75,6 +75,23 @@ class PersistedPromptContextResolver: return channel, scope.project_path +@dataclass(frozen=True, slots=True) +class TranscriptInput: + """Raw turn inputs from which ``ContextBuilder`` assembles a transcript.""" + + history: list[dict[str, Any]] + current_message: str | None + media: Sequence[str] | None = None + current_role: str = "user" + session_summary: SessionSummary | None = None + runtime_context_blocks: Sequence[RuntimeContextBlock] | None = None + + @property + def message_count(self) -> int: + """Number of boundary-preserving messages in the assembled transcript.""" + return 1 + len(self.history) + (self.current_message is not None) + + class ContextBuilder: """Builds the context (system prompt + messages) for the agent.""" @@ -282,14 +299,58 @@ class ContextBuilder: session_key: str | None = None, unified_session: bool = False, ) -> list[dict[str, Any]]: - """Build the complete message list for an LLM call.""" + """Compatibility wrapper for callers that need merged adjacent roles.""" + messages = self.build_transcript( + TranscriptInput( + history=history, + current_message=current_message, + media=media, + current_role=current_role, + session_summary=session_summary, + runtime_context_blocks=runtime_context_blocks, + ), + channel=channel, + workspace=workspace, + include_memory=include_memory, + include_memory_recent_history=include_memory_recent_history, + session_key=session_key, + unified_session=unified_session, + ) + current = messages[-1] + if len(messages) < 2 or messages[-2].get("role") != current.get("role"): + return messages + + merged = dict(messages[-2]) + merged["content"] = self._merge_message_content( + merged.get("content"), + current.get("content"), + ) + current_meta = current.get("_meta") + if current.get("role") == "user" and isinstance(current_meta, dict): + internal_meta = dict(merged.get("_meta") or {}) + internal_meta.update(cast(dict[str, Any], current_meta)) + merged["_meta"] = internal_meta + return [*messages[:-2], merged] + + def build_transcript( + self, + transcript: TranscriptInput, + *, + channel: str | None = None, + workspace: Path | None = None, + include_memory: bool = True, + include_memory_recent_history: bool = True, + session_key: str | None = None, + unified_session: bool = False, + ) -> list[dict[str, Any]]: + """Build a model transcript while preserving the fresh-turn boundary.""" root = workspace or self.workspace messages: list[dict[str, Any]] = [ { "role": "system", "content": self.build_system_prompt( channel=channel, - session_summary=session_summary, + session_summary=transcript.session_summary, workspace=root, include_memory=include_memory, include_memory_recent_history=include_memory_recent_history, @@ -297,27 +358,17 @@ class ContextBuilder: unified_session=unified_session, ), }, - *history, + *transcript.history, ] - current = self.build_current_message( - current_message, - media=media, - current_role=current_role, - runtime_context_blocks=runtime_context_blocks, - ) - if messages[-1].get("role") == current_role: - last = dict(messages[-1]) - last["content"] = self._merge_message_content( - last.get("content"), - current.get("content"), - ) - current_meta = current.get("_meta") - if current_role == "user" and isinstance(current_meta, dict): - internal_meta = dict(last.get("_meta") or {}) - internal_meta.update(cast(dict[str, Any], current_meta)) - last["_meta"] = internal_meta - messages[-1] = last + if transcript.current_message is None: return messages + + current = self.build_current_message( + transcript.current_message, + media=list(transcript.media) if transcript.media else None, + current_role=transcript.current_role, + runtime_context_blocks=transcript.runtime_context_blocks, + ) messages.append(current) return messages diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 2887e6564..9243504e0 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -14,6 +14,7 @@ from collections.abc import Coroutine, Iterable, Mapping from contextlib import AbstractContextManager, ExitStack, nullcontext, suppress from dataclasses import dataclass, field from enum import Enum, auto +from functools import partial from pathlib import Path from typing import TYPE_CHECKING, Any, Awaitable, Callable, TypeVar, cast @@ -23,7 +24,7 @@ from nanobot.agent import context as agent_context from nanobot.agent import model_presets as preset_helpers from nanobot.agent.autocompact import AutoCompact from nanobot.agent.automation_turns import publish_next_deferred_turn -from nanobot.agent.context import ContextBuilder, PersistedPromptContextResolver +from nanobot.agent.context import ContextBuilder, PersistedPromptContextResolver, TranscriptInput from nanobot.agent.cron_turns import CronTurnCoordinator from nanobot.agent.hook import AgentHook, AgentTurnHookFactory from nanobot.agent.memory import Consolidator @@ -135,7 +136,7 @@ class TurnContext: session: Session | None = None history: list[dict[str, Any]] = field(default_factory=list) - initial_messages: list[dict[str, Any]] = field(default_factory=list) + transcript_input: TranscriptInput | None = None provider_state: ProviderConversationState | None = field(default=None, repr=False) request_context: RequestContext | None = None runtime_context_blocks: list[RuntimeContextBlock] = field(default_factory=list) @@ -723,22 +724,15 @@ class AgentLoop: return True return False - def _build_initial_messages(self, ctx: TurnContext) -> list[dict[str, Any]]: - """Build the initial message list for the LLM turn.""" + def _build_transcript_input(self, ctx: TurnContext) -> TranscriptInput: + """Capture the persisted history and fresh input as separate transcript parts.""" assert ctx.session is not None - scope = self.workspace_scopes.for_message(ctx.msg, ctx.session.metadata) - return self.context.build_messages( + return TranscriptInput( history=ctx.history, current_message=ctx.msg.content, media=ctx.msg.media if ctx.kind is TurnKind.USER and ctx.msg.media else None, - channel=ctx.delivery.route.channel, session_summary=ctx.pending_summary, - workspace=scope.project_path, runtime_context_blocks=ctx.runtime_context_blocks, - include_memory=ctx.session.policy.persist, - include_memory_recent_history=not ctx.ephemeral, - session_key=ctx.session.key, - unified_session=self._unified_session, ) def _request_context_for_turn(self, ctx: TurnContext) -> RequestContext: @@ -929,7 +923,7 @@ class AgentLoop: async def _run_agent_loop( self, - initial_messages: list[dict[str, Any]], + transcript_input: TranscriptInput, on_progress: Callable[..., Awaitable[None]] | None = None, on_stream: Callable[[str], Awaitable[None]] | None = None, on_stream_end: Callable[..., Awaitable[None]] | None = None, @@ -1110,6 +1104,15 @@ class AgentLoop: message_metadata=request_metadata, session_metadata=session.metadata if session is not None else None, ) + transcript_builder = partial( + self.context.build_transcript, + channel=request_ctx.channel, + workspace=effective_scope.project_path, + include_memory=session.policy.persist if session is not None else True, + include_memory_recent_history=not ephemeral, + session_key=session.key if session is not None else request_ctx.session_key, + unified_session=self._unified_session, + ) if request_context is None: request_ctx = dataclasses.replace( request_ctx, @@ -1156,11 +1159,13 @@ class AgentLoop: run_extra_hooks_for_ephemeral=run_extra_hooks_for_ephemeral, )) result = await self.runner.run(AgentRunSpec( - initial_messages=initial_messages, + initial_messages=None, tools=effective_tools, runtime=runtime, max_iterations=self.max_iterations, max_tool_result_chars=self.max_tool_result_chars, + transcript_input=transcript_input, + transcript_builder=transcript_builder, hook=hook, concurrent_tools=True, workspace=effective_scope.project_path, @@ -1968,7 +1973,7 @@ class AgentLoop: # Upgrade the replay-safe baseline to the resumable state before # prompt assembly and the first model checkpoint. self.sessions.save(session) - ctx.initial_messages = self._build_initial_messages(ctx) + ctx.transcript_input = self._build_transcript_input(ctx) if ctx.on_progress is None: ctx.on_progress = ctx.delivery.progress_callback() @@ -1980,9 +1985,10 @@ class AgentLoop: if ctx.visible_run_started_at is None: ctx.visible_run_started_at = time.time() await ctx.delivery.running(started_at=ctx.visible_run_started_at) + assert ctx.transcript_input is not None with capture_message_deliveries() as message_sends: result = await self._run_agent_loop( - ctx.initial_messages, + ctx.transcript_input, runtime=runtime, on_progress=ctx.on_progress, on_stream=ctx.on_stream, diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index 63f7e7670..36d5303f7 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -14,6 +14,7 @@ from typing import Any, cast from loguru import logger +from nanobot.agent.context import TranscriptInput from nanobot.agent.context_governance import ( ContextGovernanceConfig, ContextGovernor, @@ -66,6 +67,7 @@ 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 = ( @@ -94,11 +96,13 @@ def _restore_outer_whitespace(content: str, original: str | None) -> str: class AgentRunSpec: """Configuration for a single agent execution.""" - initial_messages: list[dict[str, Any]] + initial_messages: list[dict[str, Any]] | None tools: ToolRegistry runtime: LLMRuntime max_iterations: int max_tool_result_chars: int + transcript_input: TranscriptInput | None = None + transcript_builder: TranscriptBuilder | None = None hook: AgentHook | None = None error_message: str | None = _DEFAULT_ERROR_MESSAGE max_iterations_message: str | None = None @@ -410,7 +414,7 @@ class AgentRunner: async def run(self, spec: AgentRunSpec) -> AgentRunResult: hook = spec.hook or AgentHook() - messages = list(spec.initial_messages) + messages = self._initial_transcript(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) @@ -462,6 +466,19 @@ class AgentRunner: finally: 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: + if spec.initial_messages is not None: + raise ValueError("provide either transcript_input or initial_messages, not both") + if spec.transcript_builder is None: + raise ValueError("transcript_builder is required with transcript_input") + return list(spec.transcript_builder(spec.transcript_input)) + if spec.initial_messages is None: + raise ValueError("initial_messages is required without transcript_input") + return list(spec.initial_messages) + async def _run_core( self, spec: AgentRunSpec, @@ -502,7 +519,7 @@ class AgentRunner: context_window_tokens=spec.runtime.context_window_tokens, context_block_limit=spec.context_block_limit, max_tokens=spec.runtime.generation.max_tokens, - inflight_start_index=len(spec.initial_messages), + inflight_start_index=len(messages), ) for iteration in range(spec.max_iterations): diff --git a/nanobot/providers/base.py b/nanobot/providers/base.py index 24f4c0df3..0b9b9ef5e 100644 --- a/nanobot/providers/base.py +++ b/nanobot/providers/base.py @@ -1029,6 +1029,20 @@ class LLMProvider(ABC): # Unknown 429 defaults to WAIT+retry. return True + @staticmethod + def _content_as_blocks(content: Any) -> list[dict[str, Any]]: + """Convert message content to blocks so mixed user content can be merged.""" + if isinstance(content, list): + return [ + dict(cast(dict[str, Any], item)) + if isinstance(item, dict) + else {"type": "text", "text": str(item)} + for item in cast(list[object], content) + ] + if content is None: + return [] + return [{"type": "text", "text": str(content)}] + @staticmethod def _enforce_role_alternation(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: """Merge consecutive same-role messages and drop trailing assistant messages. @@ -1063,6 +1077,13 @@ class LLMProvider(ABC): curr_content = msg.get("content") or "" if isinstance(prev_content, str) and isinstance(curr_content, str): prev["content"] = (prev_content + "\n\n" + curr_content).strip() + elif role == "user": + combined = dict(msg) + combined["content"] = [ + *LLMProvider._content_as_blocks(prev_content), + *LLMProvider._content_as_blocks(curr_content), + ] + merged[-1] = combined else: merged[-1] = dict(msg) else: diff --git a/nanobot/session/turn_continuation.py b/nanobot/session/turn_continuation.py index cc98ea08e..3b4736bc7 100644 --- a/nanobot/session/turn_continuation.py +++ b/nanobot/session/turn_continuation.py @@ -147,10 +147,10 @@ def prepare_save_boundary(ctx: TurnContext) -> None: if ctx.session is not None: clear_internal_continuation_state(ctx.session.metadata) + assert ctx.transcript_input is not None ctx.save_skip = _save_skip_for_turn( message_metadata=ctx.msg.metadata, - initial_message_count=len(ctx.initial_messages), - history_count=len(ctx.history), + initial_message_count=ctx.transcript_input.message_count, input_persisted_early=ctx.input_persisted_early, ) @@ -185,7 +185,6 @@ def _save_skip_for_turn( *, message_metadata: Mapping[str, Any] | None, initial_message_count: int, - history_count: int, input_persisted_early: bool, ) -> int: """Return the persisted-message append boundary for this turn.""" @@ -193,10 +192,7 @@ def _save_skip_for_turn( return initial_message_count if internal_continuation_inbound(message_metadata): return initial_message_count - # build_messages may merge the current message into a same-role history tail. - # Runner-appended messages start at initial_message_count in either shape. - has_standalone_current = initial_message_count > 1 + history_count - if has_standalone_current and not input_persisted_early: + if not input_persisted_early: return initial_message_count - 1 return initial_message_count diff --git a/tests/agent/test_attachment_references.py b/tests/agent/test_attachment_references.py index 15a3cf172..465a7ea02 100644 --- a/tests/agent/test_attachment_references.py +++ b/tests/agent/test_attachment_references.py @@ -5,6 +5,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from nanobot.agent.context import TranscriptInput from nanobot.agent.loop import AgentLoop, TurnContext, TurnKind from nanobot.agent.tools.context import RequestContext from nanobot.agent.tools.filesystem import ReadFileTool @@ -148,7 +149,10 @@ async def test_pending_document_attachment_keeps_body_out_of_prompt( runtime = loop.llm_runtime() result = await loop._run_agent_loop( - [{"role": "user", "content": "hello"}], + TranscriptInput( + history=[{"role": "user", "content": "hello"}], + current_message=None, + ), runtime=runtime, request_context=RequestContext(channel="cli", chat_id="c", runtime=runtime), pending_queue=pending_queue, diff --git a/tests/agent/test_context_builder.py b/tests/agent/test_context_builder.py index 13dd5ee52..55c014ba8 100644 --- a/tests/agent/test_context_builder.py +++ b/tests/agent/test_context_builder.py @@ -4,7 +4,7 @@ from pathlib import Path import pytest -from nanobot.agent.context import ContextBuilder +from nanobot.agent.context import ContextBuilder, TranscriptInput from nanobot.runtime_context import RuntimeContextBlock # --------------------------------------------------------------------------- @@ -403,6 +403,15 @@ class TestBuildMessages: assert "user-only runtime context" not in messages[-1]["content"] assert "_meta" not in messages[-1] + def test_compatibility_builder_merges_system_role_without_history(self, tmp_path): + builder = _builder(tmp_path) + + messages = builder.build_messages([], "system event", current_role="system") + + assert len(messages) == 1 + assert messages[0]["role"] == "system" + assert str(messages[0]["content"]).endswith("system event") + def test_explicit_skill_reference_loads_full_instructions_for_this_turn(self, tmp_path): skill_dir = tmp_path / "skills" / "review" skill_dir.mkdir(parents=True) @@ -472,6 +481,20 @@ class TestBuildMessages: assert "previous user message" in str(messages[1]["content"]) assert "new message" in str(messages[1]["content"]) + def test_structured_transcript_preserves_fresh_turn_boundary(self, tmp_path): + builder = _builder(tmp_path) + transcript = TranscriptInput( + history=[{"role": "user", "content": "previous user message"}], + current_message="new message", + ) + + messages = builder.build_transcript(transcript) + + assert [message["role"] for message in messages] == ["system", "user", "user"] + assert messages[-2]["content"] == "previous user message" + assert messages[-1]["content"] == "new message" + assert transcript.message_count == 3 + def test_current_message_can_be_built_without_history_merge(self, tmp_path): builder = _builder(tmp_path) current = builder.build_current_message( diff --git a/tests/agent/test_hook_composite.py b/tests/agent/test_hook_composite.py index 4ddb000fc..3cde0a2c4 100644 --- a/tests/agent/test_hook_composite.py +++ b/tests/agent/test_hook_composite.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from nanobot.agent.context import TranscriptInput from nanobot.agent.hook import ( AgentHook, AgentHookContext, @@ -459,7 +460,7 @@ async def test_agent_loop_extra_hook_receives_calls(tmp_path): loop.tools.get_definitions = MagicMock(return_value=[]) result = await loop._run_agent_loop( - [{"role": "user", "content": "hi"}], + TranscriptInput(history=[{"role": "user", "content": "hi"}], current_message=None), runtime=loop.llm_runtime(), ) @@ -504,7 +505,7 @@ async def test_agent_loop_turn_hook_factories_receive_context(tmp_path): runtime = loop.llm_runtime() await loop._run_agent_loop( - [{"role": "user", "content": "hi"}], + TranscriptInput(history=[{"role": "user", "content": "hi"}], current_message=None), runtime=runtime, on_progress=on_progress, request_context=RequestContext( @@ -551,7 +552,7 @@ async def test_agent_loop_extra_hook_error_isolation(tmp_path): loop.tools.get_definitions = MagicMock(return_value=[]) result = await loop._run_agent_loop( - [{"role": "user", "content": "hi"}], + TranscriptInput(history=[{"role": "user", "content": "hi"}], current_message=None), runtime=loop.llm_runtime(), ) @@ -577,7 +578,9 @@ async def test_agent_loop_extra_hooks_do_not_swallow_loop_hook_errors(tmp_path): with pytest.raises(RuntimeError, match="progress failed"): await loop._run_agent_loop( - [], runtime=loop.llm_runtime(), on_progress=bad_progress + TranscriptInput(history=[], current_message=None), + runtime=loop.llm_runtime(), + on_progress=bad_progress, ) @@ -596,7 +599,8 @@ async def test_agent_loop_no_hooks_backward_compat(tmp_path): loop.max_iterations = 2 result = await loop._run_agent_loop( - [], runtime=loop.llm_runtime() + TranscriptInput(history=[], current_message=None), + runtime=loop.llm_runtime(), ) assert result.final_content == ( "I reached the maximum number of tool call iterations (2) " diff --git a/tests/agent/test_loop_progress.py b/tests/agent/test_loop_progress.py index ca9a8a61e..b286fb569 100644 --- a/tests/agent/test_loop_progress.py +++ b/tests/agent/test_loop_progress.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from nanobot.agent.context import TranscriptInput from nanobot.agent.hooks import create_file_edit_activity_hook from nanobot.agent.loop import AgentLoop from nanobot.agent.tools.context import current_request_context @@ -84,7 +85,9 @@ class TestToolEventProgress: progress.append((content, tool_hint, tool_events)) result = await loop._run_agent_loop( - [], runtime=loop.llm_runtime(), on_progress=on_progress + TranscriptInput(history=[], current_message=None), + runtime=loop.llm_runtime(), + on_progress=on_progress, ) assert result.final_content == "Done" @@ -155,7 +158,9 @@ class TestToolEventProgress: file_events.extend(file_edit_events) result = await loop._run_agent_loop( - [], runtime=loop.llm_runtime(), on_progress=on_progress + TranscriptInput(history=[], current_message=None), + runtime=loop.llm_runtime(), + on_progress=on_progress, ) assert result.final_content == "Done" @@ -225,7 +230,9 @@ class TestToolEventProgress: ) result = await loop._run_agent_loop( - [], runtime=loop.llm_runtime(), on_progress=on_progress + TranscriptInput(history=[], current_message=None), + runtime=loop.llm_runtime(), + on_progress=on_progress, ) assert result.final_content == "Done" @@ -263,7 +270,9 @@ class TestToolEventProgress: file_events.extend(file_edit_events) await loop._run_agent_loop( - [], runtime=loop.llm_runtime(), on_progress=on_progress + TranscriptInput(history=[], current_message=None), + runtime=loop.llm_runtime(), + on_progress=on_progress, ) assert file_events == [] @@ -1019,7 +1028,7 @@ class TestToolEventProgress: progress.append((content, tool_hint, tool_events)) result = await loop._run_agent_loop( - [], + TranscriptInput(history=[], current_message=None), runtime=loop.llm_runtime(), on_progress=on_progress, on_stream=on_stream, diff --git a/tests/agent/test_loop_runner_integration.py b/tests/agent/test_loop_runner_integration.py index cfd9f7cf6..af2776da3 100644 --- a/tests/agent/test_loop_runner_integration.py +++ b/tests/agent/test_loop_runner_integration.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from nanobot.agent.context import TranscriptInput from nanobot.agent.goal_permission import goal_mutation_allowed, goal_mutation_permission from nanobot.agent.tools.context import RequestContext from nanobot.bus.outbound_events import StreamedResponseEvent @@ -55,7 +56,7 @@ async def test_ephemeral_runner_enters_and_restores_turn_scopes(tmp_path): loop.tools.get_definitions = MagicMock(return_value=[]) await loop._run_agent_loop( - [], + TranscriptInput(history=[], current_message=None), runtime=loop.llm_runtime(), ephemeral=True, turn_scopes=[goal_mutation_permission(True)], @@ -340,7 +341,8 @@ async def test_loop_max_iterations_message_stays_stable(tmp_path): loop.max_iterations = 2 result = await loop._run_agent_loop( - [], runtime=loop.llm_runtime() + TranscriptInput(history=[], current_message=None), + runtime=loop.llm_runtime(), ) assert result.final_content == ( @@ -362,7 +364,7 @@ async def test_loop_goal_turn_uses_standard_iteration_budget(tmp_path): runtime = loop.llm_runtime() result = await loop._run_agent_loop( - [], + TranscriptInput(history=[], current_message=None), runtime=runtime, request_context=RequestContext( channel="cli", @@ -401,7 +403,7 @@ async def test_loop_stream_filter_handles_think_only_prefix_without_crashing(tmp endings.append(resuming) result = await loop._run_agent_loop( - [], + TranscriptInput(history=[], current_message=None), runtime=loop.llm_runtime(), on_stream=on_stream, on_stream_end=on_stream_end, @@ -428,7 +430,9 @@ async def test_loop_stream_filter_hides_partial_trailing_think_prefix(tmp_path): deltas.append(delta) result = await loop._run_agent_loop( - [], runtime=loop.llm_runtime(), on_stream=on_stream + TranscriptInput(history=[], current_message=None), + runtime=loop.llm_runtime(), + on_stream=on_stream, ) assert result.final_content == "Hello World" @@ -451,7 +455,9 @@ async def test_loop_stream_filter_hides_complete_trailing_think_tag(tmp_path): deltas.append(delta) result = await loop._run_agent_loop( - [], runtime=loop.llm_runtime(), on_stream=on_stream + TranscriptInput(history=[], current_message=None), + runtime=loop.llm_runtime(), + on_stream=on_stream, ) assert result.final_content == "Hello World" @@ -472,7 +478,8 @@ async def test_loop_retries_think_only_final_response(tmp_path): loop.provider.chat_with_retry = chat_with_retry result = await loop._run_agent_loop( - [], runtime=loop.llm_runtime() + TranscriptInput(history=[], current_message=None), + runtime=loop.llm_runtime(), ) assert result.final_content == "Recovered answer" diff --git a/tests/agent/test_loop_save_turn.py b/tests/agent/test_loop_save_turn.py index ce41c5206..7059da77b 100644 --- a/tests/agent/test_loop_save_turn.py +++ b/tests/agent/test_loop_save_turn.py @@ -7,7 +7,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest from loguru import logger -from nanobot.agent.context import ContextBuilder +from nanobot.agent.context import ContextBuilder, TranscriptInput from nanobot.agent.loop import AgentLoop from nanobot.agent.runner import AgentRunResult from nanobot.agent.tools.context import RequestContext, request_context @@ -79,6 +79,13 @@ def _agent_run_result( ) +def _assembled_messages( + builder: ContextBuilder, + transcript_input: TranscriptInput, +) -> list[dict]: + return builder.build_transcript(transcript_input, include_memory=False) + + def _mk_loop() -> AgentLoop: loop = AgentLoop.__new__(AgentLoop) from nanobot.config.schema import AgentDefaults @@ -930,10 +937,13 @@ async def test_runtime_checkpoint_keeps_provider_state_out_of_public_metadata( session = loop.sessions.get_or_create("cli:private-checkpoint") await loop._run_agent_loop( - [ - {"role": "system", "content": "system"}, - {"role": "user", "content": "question"}, - ], + TranscriptInput( + history=[ + {"role": "system", "content": "system"}, + {"role": "user", "content": "question"}, + ], + current_message=None, + ), runtime=loop.llm_runtime(), session=session, ) @@ -1008,7 +1018,7 @@ async def test_subagent_followup_state_is_durable_before_prompt_assembly( loop = _make_full_loop(tmp_path) loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] loop.provider.can_resume_conversation_state.return_value = True - loop._build_initial_messages = MagicMock( # type: ignore[method-assign] + loop.context.build_system_prompt = MagicMock( # type: ignore[method-assign] side_effect=RuntimeError("prompt boom"), ) session = loop.sessions.get_or_create("cli:subagent-prompt-crash") @@ -1041,8 +1051,8 @@ async def test_subagent_redelivery_does_not_duplicate_staged_provider_input( loop = _make_full_loop(tmp_path) loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] loop.provider.can_resume_conversation_state.return_value = True - build_initial_messages = loop._build_initial_messages - loop._build_initial_messages = MagicMock( # type: ignore[method-assign] + build_system_prompt = loop.context.build_system_prompt + loop.context.build_system_prompt = MagicMock( # type: ignore[method-assign] side_effect=RuntimeError("prompt boom"), ) session = loop.sessions.get_or_create("cli:subagent-redelivery") @@ -1066,7 +1076,7 @@ async def test_subagent_redelivery_does_not_duplicate_staged_provider_input( message.get("content") for message in persisted.provider_state.pending_messages ].count("subagent result") == 1 - loop._build_initial_messages = build_initial_messages # type: ignore[method-assign] + loop.context.build_system_prompt = build_system_prompt # type: ignore[method-assign] loop._run_agent_loop = AsyncMock( # type: ignore[method-assign] side_effect=RuntimeError("provider boom"), ) @@ -1319,7 +1329,8 @@ async def test_internal_continuation_queues_turn_without_fake_user_history( calls: list[dict] = [] - async def fake_run_agent_loop(initial_messages, *, metadata=None, **_kwargs): + async def fake_run_agent_loop(transcript_input, *, metadata=None, **_kwargs): + initial_messages = _assembled_messages(loop.context, transcript_input) calls.append({"initial_messages": initial_messages, "metadata": metadata}) if len(calls) == 1: return _agent_run_result( @@ -1387,8 +1398,9 @@ async def test_internal_continuation_preserves_streaming_route_metadata( calls = 0 - async def fake_run_agent_loop(initial_messages, *, on_stream=None, on_stream_end=None, **_kwargs): + async def fake_run_agent_loop(transcript_input, *, on_stream=None, on_stream_end=None, **_kwargs): nonlocal calls + initial_messages = _assembled_messages(loop.context, transcript_input) calls += 1 if calls == 1: return _agent_run_result( @@ -1460,8 +1472,9 @@ async def test_websocket_internal_continuation_keeps_single_visible_run( calls = 0 - async def fake_run_agent_loop(initial_messages, **_kwargs): + async def fake_run_agent_loop(transcript_input, **_kwargs): nonlocal calls + initial_messages = _assembled_messages(loop.context, transcript_input) calls += 1 if calls == 1: return _agent_run_result( @@ -1623,7 +1636,7 @@ async def test_run_agent_loop_continuation_reads_latest_goal_metadata( runtime = loop.llm_runtime() await loop._run_agent_loop( - [], + TranscriptInput(history=[], current_message=None), runtime=runtime, session=session, request_context=RequestContext( @@ -1753,7 +1766,7 @@ async def test_stop_preserves_runtime_checkpoint_for_next_turn(tmp_path: Path) - checkpoint_saved = asyncio.Event() - async def interrupted_run_agent_loop(_initial_messages, *, session=None, **_kwargs): + async def interrupted_run_agent_loop(_transcript_input, *, session=None, **_kwargs): assert session is not None loop._set_runtime_checkpoint( session, @@ -1813,7 +1826,8 @@ async def test_stop_preserves_runtime_checkpoint_for_next_turn(tmp_path: Path) - assert interrupted.metadata.get(AgentLoop._PENDING_USER_TURN_KEY) is True assert interrupted.metadata.get(AgentLoop._RUNTIME_CHECKPOINT_KEY) is not None - async def resumed_run_agent_loop(initial_messages, **_kwargs): + async def resumed_run_agent_loop(transcript_input, **_kwargs): + initial_messages = _assembled_messages(loop.context, transcript_input) return _agent_run_result( "next answer", [*initial_messages, {"role": "assistant", "content": "next answer"}], @@ -1864,7 +1878,8 @@ async def test_system_subagent_followup_is_persisted_before_prompt_assembly(tmp_ record_runtime = MagicMock(wraps=loop.runtime_event_publisher.record_turn_runtime) loop.runtime_event_publisher.record_turn_runtime = record_runtime - async def fake_run_agent_loop(initial_messages, **kwargs): + async def fake_run_agent_loop(transcript_input, **kwargs): + initial_messages = _assembled_messages(loop.context, transcript_input) seen["initial_messages"] = initial_messages seen["runtime"] = kwargs["runtime"] seen["request_context"] = kwargs["request_context"] @@ -1940,7 +1955,8 @@ async def test_turn_usage_is_persisted_with_the_saved_session(tmp_path: Path) -> loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] turn_usage = LLMUsage.reported(input_tokens=64, output_tokens=9) - async def fake_run_agent_loop(initial_messages, **_kwargs): + async def fake_run_agent_loop(transcript_input, **_kwargs): + initial_messages = _assembled_messages(loop.context, transcript_input) return _agent_run_result( "done", [*initial_messages, {"role": "assistant", "content": "done"}], @@ -1966,7 +1982,8 @@ async def test_system_subagent_followup_does_not_log_content(tmp_path: Path) -> return_value=False ) - async def fake_run_agent_loop(initial_messages, **_kwargs): + async def fake_run_agent_loop(transcript_input, **_kwargs): + initial_messages = _assembled_messages(loop.context, transcript_input) return _agent_run_result( "done", [*initial_messages, {"role": "assistant", "content": "done"}], @@ -2022,7 +2039,8 @@ async def test_system_subagent_followup_uses_common_turn_lifecycle(tmp_path: Pat setattr(loop, name, record) - async def fake_run_agent_loop(initial_messages, **_kwargs): + async def fake_run_agent_loop(transcript_input, **_kwargs): + initial_messages = _assembled_messages(loop.context, transcript_input) return _agent_run_result( "done", [*initial_messages, {"role": "assistant", "content": "done"}], @@ -2065,7 +2083,8 @@ async def test_multiple_subagent_followups_all_persist_as_standalone_history(tmp loop = _make_full_loop(tmp_path) loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] - async def fake_run_agent_loop(initial_messages, **_kwargs): + async def fake_run_agent_loop(transcript_input, **_kwargs): + initial_messages = _assembled_messages(loop.context, transcript_input) return _agent_run_result( "ack", [*initial_messages, {"role": "assistant", "content": "ack"}], @@ -2196,7 +2215,8 @@ async def test_system_subagent_followup_uses_thread_session_and_slack_metadata(t seen: dict[str, object] = {} - async def fake_run_agent_loop(initial_messages, **kwargs): + async def fake_run_agent_loop(transcript_input, **kwargs): + initial_messages = _assembled_messages(loop.context, transcript_input) seen["initial_messages"] = initial_messages seen["request_context"] = kwargs["request_context"] return _agent_run_result( @@ -2252,8 +2272,11 @@ async def test_turn_after_unanswered_user_keeps_tool_call_pairing(tmp_path: Path session.add_message("user", "earlier question that never got an answer") loop.sessions.save(session) - async def fake_run_agent_loop(initial_messages, **_kwargs): - assert [m["role"] for m in initial_messages] == ["system", "user"] + async def fake_run_agent_loop(transcript_input, **_kwargs): + initial_messages = _assembled_messages(loop.context, transcript_input) + assert [m["role"] for m in initial_messages] == ["system", "user", "user"] + assert initial_messages[-2]["content"] == "earlier question that never got an answer" + assert initial_messages[-1]["content"] == "and another thing" return _agent_run_result( "done", [ diff --git a/tests/agent/test_loop_tool_context.py b/tests/agent/test_loop_tool_context.py index bf4b00601..5553fb5ca 100644 --- a/tests/agent/test_loop_tool_context.py +++ b/tests/agent/test_loop_tool_context.py @@ -5,6 +5,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from nanobot.agent.context import TranscriptInput from nanobot.agent.loop import AgentLoop from nanobot.agent.tools.context import ( RequestContext, @@ -133,7 +134,7 @@ async def test_loop_binds_request_context_for_tool_execution(tmp_path: Path) -> metadata = {"slack": {"thread_ts": "111.222", "channel_type": "channel"}} runtime = loop.llm_runtime() await loop._run_agent_loop( - [], + TranscriptInput(history=[], current_message=None), runtime=runtime, request_context=RequestContext( channel="slack", @@ -234,7 +235,7 @@ async def test_agent_loop_restores_outer_request_context_after_runner_exception( try: with pytest.raises(RuntimeError, match="runner failed"): await loop._run_agent_loop( - [], + TranscriptInput(history=[], current_message=None), runtime=runtime, request_context=RequestContext( channel="slack", diff --git a/tests/agent/test_runner_core.py b/tests/agent/test_runner_core.py index f69dd2247..62f7d0a60 100644 --- a/tests/agent/test_runner_core.py +++ b/tests/agent/test_runner_core.py @@ -10,6 +10,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from agent.runner_helpers import make_run_spec +from nanobot.agent.context import TranscriptInput from nanobot.config.schema import AgentDefaults from nanobot.providers.base import ( LLMProvider, @@ -34,6 +35,35 @@ def _make_usage_spec(provider, tools): ) +def test_initial_transcript_is_built_from_structured_turn_input() -> None: + from nanobot.agent.runner import AgentRunner + + provider = MagicMock(spec=LLMProvider) + transcript_input = TranscriptInput( + history=[{"role": "user", "content": "earlier"}], + current_message="fresh", + ) + expected = [ + {"role": "system", "content": "system"}, + {"role": "user", "content": "earlier"}, + {"role": "user", "content": "fresh"}, + ] + transcript_builder = MagicMock(return_value=expected) + spec = make_run_spec( + provider, + initial_messages=None, + transcript_input=transcript_input, + transcript_builder=transcript_builder, + tools=MagicMock(), + model="test-model", + max_iterations=1, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + ) + + assert AgentRunner._initial_transcript(spec) == expected + transcript_builder.assert_called_once_with(transcript_input) + + def test_usage_or_estimate_replaces_reported_zero_for_content(monkeypatch) -> None: from nanobot.agent.runner import AgentRunner diff --git a/tests/agent/test_runner_injections.py b/tests/agent/test_runner_injections.py index a7979e73a..c7763f9de 100644 --- a/tests/agent/test_runner_injections.py +++ b/tests/agent/test_runner_injections.py @@ -10,6 +10,7 @@ import pytest from agent.runner_helpers import make_run_spec from nanobot.agent.automation_turns import publish_next_deferred_turn +from nanobot.agent.context import TranscriptInput from nanobot.agent.tools.context import RequestContext from nanobot.config.schema import AgentDefaults from nanobot.providers.base import LLMResponse, ToolCallRequest @@ -617,7 +618,7 @@ async def test_loop_injected_followup_preserves_image_media(tmp_path): runtime = loop.llm_runtime() result = await loop._run_agent_loop( - [{"role": "user", "content": "hello"}], + TranscriptInput(history=[{"role": "user", "content": "hello"}], current_message=None), runtime=runtime, request_context=RequestContext(channel="cli", chat_id="c", runtime=runtime), pending_queue=pending_queue, @@ -711,7 +712,10 @@ async def test_pending_injection_resolves_its_own_runtime_context(tmp_path): runtime = loop.llm_runtime() result = await loop._run_agent_loop( - [{"role": "user", "content": "initial message from user A"}], + TranscriptInput( + history=[{"role": "user", "content": "initial message from user A"}], + current_message=None, + ), runtime=runtime, session=session, request_context=RequestContext( @@ -812,7 +816,7 @@ async def test_subagent_pending_injection_is_hidden_history_and_not_merged(tmp_p runtime = loop.llm_runtime() result = await loop._run_agent_loop( - [{"role": "user", "content": "hello"}], + TranscriptInput(history=[{"role": "user", "content": "hello"}], current_message=None), runtime=runtime, request_context=RequestContext(channel="cli", chat_id="c", runtime=runtime), pending_queue=pending_queue, @@ -1476,7 +1480,7 @@ async def test_pending_queue_preserves_overflow_for_next_injection_cycle(tmp_pat runtime = loop.llm_runtime() result = await loop._run_agent_loop( - [{"role": "user", "content": "hello"}], + TranscriptInput(history=[{"role": "user", "content": "hello"}], current_message=None), runtime=runtime, request_context=RequestContext(channel="cli", chat_id="c", runtime=runtime), pending_queue=pending_queue, diff --git a/tests/agent/tools/test_subagent_tools.py b/tests/agent/tools/test_subagent_tools.py index ded12e721..9ee375399 100644 --- a/tests/agent/tools/test_subagent_tools.py +++ b/tests/agent/tools/test_subagent_tools.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from nanobot.agent.context import TranscriptInput from nanobot.agent.tools.context import RequestContext from nanobot.config.schema import AgentDefaults from nanobot.providers.base import GenerationSettings @@ -568,7 +569,10 @@ async def test_agent_loop_syncs_updated_max_iterations_before_run(tmp_path): loop.runner.run = AsyncMock(side_effect=fake_run) loop.max_iterations = 55 - await loop._run_agent_loop([], runtime=loop.llm_runtime()) + await loop._run_agent_loop( + TranscriptInput(history=[], current_message=None), + runtime=loop.llm_runtime(), + ) loop.runner.run.assert_awaited_once() @@ -609,7 +613,7 @@ async def test_drain_pending_no_block_when_no_subagents(tmp_path): runtime = loop.llm_runtime() await loop._run_agent_loop( - [{"role": "user", "content": "test"}], + TranscriptInput(history=[{"role": "user", "content": "test"}], current_message=None), runtime=runtime, session=None, request_context=RequestContext(channel="test", chat_id="c1", runtime=runtime), @@ -668,7 +672,7 @@ async def test_terminal_drain_timeout(tmp_path): runtime = loop.llm_runtime() await loop._run_agent_loop( - [{"role": "user", "content": "test"}], + TranscriptInput(history=[{"role": "user", "content": "test"}], current_message=None), runtime=runtime, session=session, request_context=RequestContext( @@ -742,7 +746,7 @@ async def test_terminal_drain_reuses_one_timeout_budget(tmp_path): loop.subagents._running_tasks["sub-deadline-1"] = hang_task await loop._run_agent_loop( - [{"role": "user", "content": "test"}], + TranscriptInput(history=[{"role": "user", "content": "test"}], current_message=None), runtime=loop.llm_runtime(), session=session, pending_queue=pending_queue, diff --git a/tests/cli/test_restart_command.py b/tests/cli/test_restart_command.py index 85370b935..13fe3515a 100644 --- a/tests/cli/test_restart_command.py +++ b/tests/cli/test_restart_command.py @@ -11,6 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from nanobot.agent.context import TranscriptInput from nanobot.bus.events import InboundMessage from nanobot.providers.base import LLMResponse, LLMUsage @@ -311,10 +312,16 @@ class TestRestartCommand: LLMResponse(content="second", usage=None), ]) - first = await loop._run_agent_loop([], runtime=loop.llm_runtime()) + first = await loop._run_agent_loop( + TranscriptInput(history=[], current_message=None), + runtime=loop.llm_runtime(), + ) assert first.usage == LLMUsage.reported(input_tokens=9, output_tokens=4) - second = await loop._run_agent_loop([], runtime=loop.llm_runtime()) + second = await loop._run_agent_loop( + TranscriptInput(history=[], current_message=None), + runtime=loop.llm_runtime(), + ) assert second.usage == LLMUsage.estimated(input_tokens=123, output_tokens=7) @pytest.mark.asyncio diff --git a/tests/providers/test_enforce_role_alternation.py b/tests/providers/test_enforce_role_alternation.py index 332a5fcbe..ac458e5d3 100644 --- a/tests/providers/test_enforce_role_alternation.py +++ b/tests/providers/test_enforce_role_alternation.py @@ -112,14 +112,49 @@ class TestEnforceRoleAlternation: assert result[1]["content"] is None assert result[2]["role"] == "tool" - def test_non_string_content_uses_latest(self): + def test_consecutive_user_messages_preserve_text_before_multimodal_content(self): + image = { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,aW1hZ2U="}, + } msgs = [ - {"role": "user", "content": [{"type": "text", "text": "A"}]}, - {"role": "user", "content": "B"}, + {"role": "user", "content": "Earlier unanswered question"}, + { + "role": "user", + "content": [image, {"type": "text", "text": "The error is here"}], + }, ] result = LLMProvider._enforce_role_alternation(msgs) - assert len(result) == 1 - assert result[0]["content"] == "B" + assert result == [{ + "role": "user", + "content": [ + {"type": "text", "text": "Earlier unanswered question"}, + image, + {"type": "text", "text": "The error is here"}, + ], + }] + + def test_consecutive_user_messages_preserve_multimodal_content_before_text(self): + image = { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,aW1hZ2U="}, + } + msgs = [ + { + "role": "user", + "content": [image, {"type": "text", "text": "First question"}], + }, + {"role": "user", "content": "Follow-up detail"}, + ] + result = LLMProvider._enforce_role_alternation(msgs) + assert result == [{ + "role": "user", + "content": [ + image, + {"type": "text", "text": "First question"}, + {"type": "text", "text": "Follow-up detail"}, + ], + }] def test_original_messages_not_mutated(self): msgs = [ diff --git a/tests/session/test_turn_continuation.py b/tests/session/test_turn_continuation.py index 862d3415d..090275e3d 100644 --- a/tests/session/test_turn_continuation.py +++ b/tests/session/test_turn_continuation.py @@ -141,14 +141,13 @@ def test_internal_continuation_requires_budget_boundary_and_queue(): ) -def test_save_skip_matches_prefix_when_current_message_merged(): +def test_save_skip_matches_prefix_when_current_message_was_persisted(): skip = _save_skip_for_turn( message_metadata=None, - initial_message_count=2, # [system, merged user] - history_count=1, + initial_message_count=3, # [system, history user, current user] input_persisted_early=True, ) - assert skip == 2 + assert skip == 3 def test_save_skip_unchanged_for_standalone_current_message(): @@ -156,12 +155,10 @@ def test_save_skip_unchanged_for_standalone_current_message(): assert _save_skip_for_turn( message_metadata=None, initial_message_count=3, - history_count=1, input_persisted_early=True, ) == 3 assert _save_skip_for_turn( message_metadata=None, initial_message_count=3, - history_count=1, input_persisted_early=False, ) == 2 diff --git a/tests/tools/test_message_tool_suppress.py b/tests/tools/test_message_tool_suppress.py index 80198eff0..1a04ccbd4 100644 --- a/tests/tools/test_message_tool_suppress.py +++ b/tests/tools/test_message_tool_suppress.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from nanobot.agent.context import TranscriptInput from nanobot.agent.loop import AgentLoop from nanobot.agent.tools.message import MessageTool from nanobot.bus.events import InboundMessage, OutboundMessage @@ -178,7 +179,9 @@ class TestMessageToolSuppressLogic: progress.append((content, tool_hint)) result = await loop._run_agent_loop( - [], runtime=loop.llm_runtime(), on_progress=on_progress + TranscriptInput(history=[], current_message=None), + runtime=loop.llm_runtime(), + on_progress=on_progress, ) assert result.final_content == "Done"