Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(agent): close runtime context persistence gaps

Browse Source
This commit is contained in:
chengyongru
2026-07-12 00:35:17 +08:00
committed by Xubin Ren
parent f75d3519db
commit 2eb7398f34
10 changed files with 235 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/python-sdk.md
+7 -1
View File
@@ -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 |
nanobot/agent/loop.py
+1 -2
View File
@@ -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
+15
View File
@@ -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
+42 -3
View File
@@ -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
+23 -9
View File
@@ -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
+4 -2
View File
@@ -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
+8 -2
View File
@@ -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}]"
tests/agent/test_loop_image_generation_media.py
+55 -1
View File
@@ -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"
tests/agent/test_session_manager_history.py
+20
View File
@@ -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")
tests/test_nanobot_facade.py
+60
View File
@@ -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)