fix(agent): close runtime context persistence gaps

nanobot/agent/loop.py

@@ -86,7 +86,6 @@ from nanobot.triggers.local_turns import LocalTriggerTurnCoordinator
 from nanobot.utils.document import extract_documents, reference_non_image_attachments
 from nanobot.utils.helpers import image_placeholder_text
 from nanobot.utils.helpers import truncate_text as truncate_text_fn
-from nanobot.utils.image_generation_intent import image_generation_prompt
 from nanobot.utils.llm_runtime import LLMRuntime
 from nanobot.utils.runtime import (
     EMPTY_FINAL_RESPONSE_MESSAGE,
@@ -674,7 +673,7 @@ class AgentLoop:
         scope = self.workspace_scopes.for_message(msg, session.metadata)
         return self.context.build_messages(
             history=history,
-            current_message=image_generation_prompt(msg.content, msg.metadata),
+            current_message=msg.content,
             media=msg.media if msg.media else None,
             channel=msg.channel,
             chat_id=self._runtime_chat_id(msg),

nanobot/agent/tools/image_generation.py

@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any
 from pydantic import Field
 
 from nanobot.agent.tools.base import Tool, ToolResult, tool_parameters
+from nanobot.agent.tools.context import RequestContext
 from nanobot.agent.tools.schema import (
     ArraySchema,
     IntegerSchema,
@@ -21,6 +22,7 @@ from nanobot.providers.image_generation import (
     ImageGenerationProvider,
     get_image_gen_provider,
 )
+from nanobot.runtime_context import RuntimeContextBlock
 from nanobot.security.workspace_access import current_tool_workspace
 from nanobot.security.workspace_policy import WorkspaceBoundaryError, resolve_allowed_path
 from nanobot.utils.artifacts import (
@@ -29,6 +31,7 @@ from nanobot.utils.artifacts import (
     store_generated_image_artifact,
 )
 from nanobot.utils.helpers import detect_image_mime
+from nanobot.utils.image_generation_intent import image_generation_runtime_context
 
 if TYPE_CHECKING:
     from nanobot.config.schema import ProviderConfig
@@ -116,6 +119,18 @@ class ImageGenerationTool(Tool):
             "or user image paths as reference_images."
         )
 
+    def runtime_context_provider(self):
+        return self._provide_runtime_context
+
+    async def _provide_runtime_context(
+        self,
+        request: RequestContext,
+    ) -> RuntimeContextBlock | None:
+        content = image_generation_runtime_context(request.metadata)
+        if not content:
+            return None
+        return RuntimeContextBlock(source="image_generation", content=content)
+
     def _provider_config(self) -> ProviderConfig | None:
         return self.provider_configs.get(self.config.provider)
 

nanobot/sdk/clients.py

@@ -65,7 +65,7 @@ class SessionClient:
         return snapshot_from_session(session)
 
     def get(self, session_key: str) -> SessionSnapshot | None:
-        """Return a session snapshot without creating a new session on disk."""
+        """Return a display-safe snapshot without creating a new session on disk."""
         cached = self._loop.sessions._cache.get(session_key)
         if cached is not None:
             return snapshot_from_session(cached)
@@ -89,8 +89,47 @@ class SessionClient:
         ]
 
     def export(self, session_key: str) -> SessionSnapshot | None:
-        """Return a full session snapshot suitable for JSON serialization."""
-        return self.get(session_key)
+        """Return a trusted full snapshot, including model-only runtime context."""
+        cached = self._loop.sessions._cache.get(session_key)
+        if cached is not None:
+            return snapshot_from_session(cached, include_runtime_context=True)
+        payload = self._loop.sessions.read_session_file(session_key)
+        if payload is None:
+            return None
+        return snapshot_from_payload(payload, include_runtime_context=True)
+
+    async def restore(
+        self,
+        snapshot: SessionSnapshot,
+        *,
+        session_key: str | None = None,
+        save: bool = True,
+    ) -> SessionSnapshot:
+        """Restore a trusted snapshot into an empty session."""
+        key = session_key or snapshot.key
+        if not key:
+            raise ValueError("restored snapshots must include a session key")
+        session = self._loop.sessions.get_or_create(key)
+        if session.messages:
+            raise ValueError(f"restore target session is not empty: {key}")
+        session.metadata.update(deepcopy(snapshot.metadata))
+
+        for raw in snapshot.messages:
+            if "role" not in raw or "content" not in raw:
+                raise ValueError("restored messages must include role and content")
+            role = str(raw["role"]).strip()
+            if role not in self._VALID_ROLES:
+                raise ValueError(f"unsupported message role: {role!r}")
+            extra = {
+                field: deepcopy(value)
+                for field, value in raw.items()
+                if field not in {"role", "content"}
+            }
+            session.add_message(role, deepcopy(raw["content"]), **extra)
+
+        if save:
+            self._loop.sessions.save(session)
+        return snapshot_from_session(session)
 
     def clear(self, session_key: str) -> SessionSnapshot:
         """Clear one session and persist the empty session."""

nanobot/sdk/types.py

@@ -79,7 +79,7 @@ class StreamEvent:
 
 @dataclass(slots=True)
 class SessionSnapshot:
-    """A durable snapshot of one nanobot session."""
+    """A serializable session snapshot; trusted exports may include internal context."""
 
     key: str
     messages: list[dict[str, Any]]
@@ -121,27 +121,41 @@ class SessionInfo:
         }
 
 
-def snapshot_from_session(session: Any) -> SessionSnapshot:
+def snapshot_from_session(
+    session: Any,
+    *,
+    include_runtime_context: bool = False,
+) -> SessionSnapshot:
+    messages = deepcopy(session.messages)
+    if not include_runtime_context:
+        messages = public_history_messages(messages)
     return SessionSnapshot(
         key=session.key,
         created_at=session.created_at.isoformat(),
         updated_at=session.updated_at.isoformat(),
         metadata=deepcopy(session.metadata),
-        messages=public_history_messages(session.messages),
+        messages=messages,
     )
 
 
-def snapshot_from_payload(payload: Mapping[str, Any]) -> SessionSnapshot:
+def snapshot_from_payload(
+    payload: Mapping[str, Any],
+    *,
+    include_runtime_context: bool = False,
+) -> SessionSnapshot:
+    messages = [
+        deepcopy(dict(message))
+        for message in list(payload.get("messages") or [])
+        if isinstance(message, Mapping)
+    ]
+    if not include_runtime_context:
+        messages = public_history_messages(messages)
     return SessionSnapshot(
         key=str(payload.get("key") or ""),
         created_at=payload.get("created_at"),
         updated_at=payload.get("updated_at"),
         metadata=deepcopy(dict(payload.get("metadata") or {})),
-        messages=public_history_messages(
-            message
-            for message in list(payload.get("messages") or [])
-            if isinstance(message, Mapping)
-        ),
+        messages=messages,
     )
 
 

nanobot/session/manager.py

@@ -216,7 +216,8 @@ class Session:
                 content = f"{content}\n{breadcrumbs}" if content else breadcrumbs
             cli_apps = message.get("cli_apps")
             if (
-                not has_persisted_runtime_context
+                include_runtime_context
+                and not has_persisted_runtime_context
                 and role == "user"
                 and isinstance(cli_apps, list)
                 and cli_apps
@@ -239,7 +240,8 @@ class Session:
                     content = f"{content}\n{breadcrumbs}" if content else breadcrumbs
             mcp_presets = message.get("mcp_presets")
             if (
-                not has_persisted_runtime_context
+                include_runtime_context
+                and not has_persisted_runtime_context
                 and role == "user"
                 and isinstance(mcp_presets, list)
                 and mcp_presets

nanobot/utils/image_generation_intent.py

@@ -9,9 +9,15 @@ IMAGE_GENERATION_METADATA_KEY = "image_generation"
 
 def image_generation_prompt(content: str, metadata: dict[str, Any] | None) -> str:
     """Decorate a user prompt when WebUI image mode is enabled."""
+    runtime_context = image_generation_runtime_context(metadata)
+    return f"{content}\n\n{runtime_context}" if runtime_context else content
+
+
+def image_generation_runtime_context(metadata: dict[str, Any] | None) -> str:
+    """Return the model-only instruction for WebUI image generation mode."""
     raw = (metadata or {}).get(IMAGE_GENERATION_METADATA_KEY)
     if not isinstance(raw, dict) or raw.get("enabled") is not True:
-        return content
+        return ""
 
     aspect_ratio = raw.get("aspect_ratio")
     if isinstance(aspect_ratio, str) and aspect_ratio.strip():
@@ -24,4 +30,4 @@ def image_generation_prompt(content: str, metadata: dict[str, Any] | None) -> st
             "The user selected WebUI image generation mode. Use the generate_image tool. "
             "Choose the most suitable aspect_ratio yourself from the prompt and intended use."
         )
-    return f"{content}\n\n[WebUI image generation instruction: {instruction}]"
+    return f"[WebUI image generation instruction: {instruction}]"