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

@@ -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
 

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,

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):

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:

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