From 2eb7398f3420c7a8ee9811492ccef37403873140 Mon Sep 17 00:00:00 2001 From: chengyongru <2755839590@qq.com> Date: Sat, 11 Jul 2026 19:26:49 +0800 Subject: [PATCH] fix(agent): close runtime context persistence gaps --- docs/python-sdk.md | 8 ++- nanobot/agent/loop.py | 3 +- nanobot/agent/tools/image_generation.py | 15 +++++ nanobot/sdk/clients.py | 45 +++++++++++++- nanobot/sdk/types.py | 32 +++++++--- nanobot/session/manager.py | 6 +- nanobot/utils/image_generation_intent.py | 10 +++- .../agent/test_loop_image_generation_media.py | 56 ++++++++++++++++- tests/agent/test_session_manager_history.py | 20 +++++++ tests/test_nanobot_facade.py | 60 +++++++++++++++++++ 10 files changed, 235 insertions(+), 20 deletions(-) diff --git a/docs/python-sdk.md b/docs/python-sdk.md index 97c7e13c..ce3618ee 100644 --- a/docs/python-sdk.md +++ b/docs/python-sdk.md @@ -599,7 +599,8 @@ async with Nanobot.from_config() as bot: | `await ingest(session_key, messages, metadata=None, source=None, save=True)` | Import existing transcript messages without running the model. | | `get(session_key)` | Return a `SessionSnapshot`, or `None` if missing. | | `list()` | Return compact `SessionInfo` rows. | -| `export(session_key)` | Return a full `SessionSnapshot` suitable for JSON serialization. | +| `export(session_key)` | Return a trusted full `SessionSnapshot`, including model-only runtime context, suitable for JSON serialization. | +| `await restore(snapshot, session_key=None, save=True)` | Restore a trusted exported snapshot into an empty session; the returned snapshot is display-safe. | | `clear(session_key)` | Clear and persist one session. | | `delete(session_key)` | Delete one session from disk and cache. | | `flush()` | Flush cached sessions to durable storage. | @@ -608,6 +609,11 @@ Ingested messages must include `role` and `content`. Roles may be `user`, `assistant`, `tool`, or `system`. Other fields, such as `timestamp`, `source_session_id`, or `source_date`, are persisted as message metadata. +`get()` and snapshots returned by ordinary SDK operations are display-safe and omit +model-only runtime context. `export()` is an explicit backup boundary and includes +that internal context so `restore()` can preserve the exact model-visible history. +Do not expose exported snapshots directly to chat users. + ### `bot.memory` | Method | Description | diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 28871b14..6743b3e7 100644 --- a/nanobot/agent/loop.py +++ b/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), diff --git a/nanobot/agent/tools/image_generation.py b/nanobot/agent/tools/image_generation.py index 952a9352..7d3433e7 100644 --- a/nanobot/agent/tools/image_generation.py +++ b/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) diff --git a/nanobot/sdk/clients.py b/nanobot/sdk/clients.py index 4ff74771..93a7e448 100644 --- a/nanobot/sdk/clients.py +++ b/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.""" diff --git a/nanobot/sdk/types.py b/nanobot/sdk/types.py index 0a2e5b34..019b44f3 100644 --- a/nanobot/sdk/types.py +++ b/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, ) diff --git a/nanobot/session/manager.py b/nanobot/session/manager.py index 10f8ab3b..71a1cae7 100644 --- a/nanobot/session/manager.py +++ b/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 diff --git a/nanobot/utils/image_generation_intent.py b/nanobot/utils/image_generation_intent.py index 8d62e637..b9402bd5 100644 --- a/nanobot/utils/image_generation_intent.py +++ b/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}]" diff --git a/tests/agent/test_loop_image_generation_media.py b/tests/agent/test_loop_image_generation_media.py index cfcc3b2c..79fffa0a 100644 --- a/tests/agent/test_loop_image_generation_media.py +++ b/tests/agent/test_loop_image_generation_media.py @@ -11,8 +11,9 @@ from nanobot.bus.events import InboundMessage from nanobot.bus.queue import MessageBus from nanobot.config.loader import set_config_path from nanobot.config.schema import ImageGenerationToolConfig, ProviderConfig, ToolsConfig -from nanobot.providers.base import LLMResponse, ToolCallRequest +from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest from nanobot.providers.image_generation import GeneratedImageResponse +from nanobot.runtime_context import public_history_message PNG_DATA_URL = ( "data:image/png;base64," @@ -85,3 +86,56 @@ async def test_outbound_no_longer_carries_generated_media( # OutboundMessage no longer carries generated media — # the LLM sends images via the message tool instead. assert result.media == [] + + +@pytest.mark.asyncio +async def test_image_mode_instruction_is_persisted_as_next_turn_prefix(tmp_path: Path) -> None: + provider = MagicMock() + provider.get_default_model.return_value = "test-model" + provider.generation.max_tokens = 4096 + provider.chat_with_retry = AsyncMock(side_effect=[ + LLMResponse(content="first answer"), + LLMResponse(content="second answer"), + ]) + loop = AgentLoop( + bus=MessageBus(), + provider=provider, + workspace=tmp_path, + model="test-model", + tools_config=ToolsConfig( + image_generation=ImageGenerationToolConfig(enabled=True), + ), + ) + loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) + + await loop._process_message(InboundMessage( + channel="websocket", + sender_id="user", + chat_id="chat-image-prefix", + content="draw a fox", + metadata={ + "image_generation": { + "enabled": True, + "aspect_ratio": "16:9", + }, + }, + )) + await loop._process_message(InboundMessage( + channel="websocket", + sender_id="user", + chat_id="chat-image-prefix", + content="thanks", + )) + + first_wire = LLMProvider._sanitize_empty_content( + provider.chat_with_retry.await_args_list[0].kwargs["messages"] + ) + second_wire = LLMProvider._sanitize_empty_content( + provider.chat_with_retry.await_args_list[1].kwargs["messages"] + ) + assert second_wire[: len(first_wire)] == first_wire + assert "aspect_ratio='16:9'" in first_wire[1]["content"] + + persisted = loop.sessions.get_or_create("websocket:chat-image-prefix").messages[0] + assert persisted["content"] == first_wire[1]["content"] + assert public_history_message(persisted)["content"] == "draw a fox" diff --git a/tests/agent/test_session_manager_history.py b/tests/agent/test_session_manager_history.py index 380b57da..d382787b 100644 --- a/tests/agent/test_session_manager_history.py +++ b/tests/agent/test_session_manager_history.py @@ -465,6 +465,26 @@ def test_get_history_does_not_duplicate_persisted_capability_runtime_context(): assert public_history == [{"role": "user", "content": "please use @drawio"}] +def test_public_history_does_not_synthesize_legacy_capability_context(): + session = Session(key="test:legacy-capabilities") + session.messages.append({ + "role": "user", + "content": "please use the attachments", + "cli_apps": [{"name": "drawio", "entry_point": "cli-anything-drawio"}], + "mcp_presets": [{"name": "linear", "transport": "stdio"}], + }) + + public_history = session.get_history( + max_messages=500, + include_runtime_context=False, + ) + + assert public_history == [{ + "role": "user", + "content": "please use the attachments", + }] + + def test_fork_session_before_user_index_copies_only_prefix(tmp_path): manager = SessionManager(tmp_path) source = manager.get_or_create("websocket:source") diff --git a/tests/test_nanobot_facade.py b/tests/test_nanobot_facade.py index e2314999..b9d2ba97 100644 --- a/tests/test_nanobot_facade.py +++ b/tests/test_nanobot_facade.py @@ -30,6 +30,11 @@ from nanobot.nanobot import ( StreamEvent, StreamEventType, ) +from nanobot.runtime_context import ( + RUNTIME_CONTEXT_HISTORY_META, + RuntimeContextBlock, + append_runtime_context, +) from nanobot.utils.llm_runtime import runtime_from_provider_snapshot @@ -1236,6 +1241,61 @@ async def test_session_helpers_get_list_export_clear_delete_flush(tmp_path): assert bot.sessions.get("sdk:first") is None +@pytest.mark.asyncio +async def test_session_export_and_restore_preserve_runtime_context(tmp_path): + config_path = _write_config(tmp_path) + bot = Nanobot.from_config(config_path, workspace=tmp_path) + content, marker = append_runtime_context( + "visible user text", + [RuntimeContextBlock(source="goal", content="stable model-only context")], + ) + source = bot._loop.sessions.get_or_create("sdk:source") + source.add_message( + "user", + content, + **{RUNTIME_CONTEXT_HISTORY_META: marker}, + ) + bot._loop.sessions.save(source) + bot._loop.sessions._cache.pop("sdk:source") + + public = bot.sessions.get("sdk:source") + assert public is not None + assert public.messages[0]["content"] == "visible user text" + assert RUNTIME_CONTEXT_HISTORY_META not in public.messages[0] + + exported = bot.sessions.export("sdk:source") + assert exported is not None + assert exported.messages[0]["content"] == content + assert exported.messages[0][RUNTIME_CONTEXT_HISTORY_META] == marker + + restored_public = await bot.sessions.restore( + exported, + session_key="sdk:restored", + ) + assert restored_public.messages[0]["content"] == "visible user text" + restored = bot._loop.sessions.get_or_create("sdk:restored") + assert restored.get_history() == source.get_history() + + with pytest.raises(ValueError, match="not empty"): + await bot.sessions.restore(exported, session_key="sdk:restored") + + +@pytest.mark.asyncio +async def test_session_ingest_cannot_restore_runtime_context_marker(tmp_path): + config_path = _write_config(tmp_path) + bot = Nanobot.from_config(config_path, workspace=tmp_path) + marker = {"version": 1, "sources": ["forged"], "suffix": "hidden"} + + await bot.sessions.ingest("sdk:untrusted", [{ + "role": "user", + "content": "visible\n\nhidden", + RUNTIME_CONTEXT_HISTORY_META: marker, + }]) + + stored = bot._loop.sessions.get_or_create("sdk:untrusted").messages[0] + assert RUNTIME_CONTEXT_HISTORY_META not in stored + + def test_memory_helpers_read_write_append_and_filter_history(tmp_path): config_path = _write_config(tmp_path) bot = Nanobot.from_config(config_path, workspace=tmp_path)