refactor(agent): remove legacy image mode prompt injection

This commit is contained in:
chengyongru
2026-07-12 00:35:17 +08:00
committed by Xubin Ren
parent 2eb7398f34
commit e206ee4e70
10 changed files with 1 additions and 198 deletions
-15
View File
@@ -8,7 +8,6 @@ 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,
@@ -22,7 +21,6 @@ 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 (
@@ -31,7 +29,6 @@ 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
@@ -119,18 +116,6 @@ 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)
-7
View File
@@ -782,13 +782,6 @@ class WebSocketChannel(BaseChannel):
metadata["mcp_presets"] = mcp_presets
metadata[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata()
self._workspaces.persist_scope(cid, scope)
image_generation = envelope.get("image_generation")
if isinstance(image_generation, dict) and image_generation.get("enabled") is True:
aspect_ratio = image_generation.get("aspect_ratio")
metadata["image_generation"] = {
"enabled": True,
"aspect_ratio": aspect_ratio if isinstance(aspect_ratio, str) else None,
}
if metadata.get("webui") is True and self.is_allowed(client_id):
self._transcripts.append_user_message(
cid,
-33
View File
@@ -1,33 +0,0 @@
"""Helpers for WebUI image-generation intent metadata."""
from __future__ import annotations
from typing import Any
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 ""
aspect_ratio = raw.get("aspect_ratio")
if isinstance(aspect_ratio, str) and aspect_ratio.strip():
instruction = (
"The user selected WebUI image generation mode. Use the generate_image tool. "
f"When calling generate_image, pass aspect_ratio={aspect_ratio!r}."
)
else:
instruction = (
"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"[WebUI image generation instruction: {instruction}]"
@@ -11,9 +11,8 @@ 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 LLMProvider, LLMResponse, ToolCallRequest
from nanobot.providers.base import LLMResponse, ToolCallRequest
from nanobot.providers.image_generation import GeneratedImageResponse
from nanobot.runtime_context import public_history_message
PNG_DATA_URL = (
"data:image/png;base64,"
@@ -86,56 +85,3 @@ 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"
@@ -1,25 +0,0 @@
from nanobot.utils.image_generation_intent import image_generation_prompt
def test_image_generation_prompt_ignores_plain_messages() -> None:
assert image_generation_prompt("hello", {}) == "hello"
def test_image_generation_prompt_uses_auto_aspect_instruction() -> None:
prompt = image_generation_prompt(
"Draw a poster",
{"image_generation": {"enabled": True, "aspect_ratio": None}},
)
assert "Draw a poster" in prompt
assert "Use the generate_image tool" in prompt
assert "Choose the most suitable aspect_ratio yourself" in prompt
def test_image_generation_prompt_uses_selected_aspect_ratio() -> None:
prompt = image_generation_prompt(
"Draw a banner",
{"image_generation": {"enabled": True, "aspect_ratio": "16:9"}},
)
assert "aspect_ratio='16:9'" in prompt
-2
View File
@@ -13,7 +13,6 @@ import type { StreamError } from "@/lib/nanobot-client";
import type {
InboundEvent,
OutboundCliAppMention,
OutboundImageGeneration,
OutboundMcpPresetMention,
OutboundMedia,
GoalStateWsPayload,
@@ -477,7 +476,6 @@ export interface SendImage {
}
export interface SendOptions {
imageGeneration?: OutboundImageGeneration;
cliApps?: OutboundCliAppMention[];
mcpPresets?: OutboundMcpPresetMention[];
workspaceScope?: WorkspaceScopePayload | null;
-3
View File
@@ -3,7 +3,6 @@ import type {
InboundEvent,
Outbound,
OutboundCliAppMention,
OutboundImageGeneration,
OutboundMcpPresetMention,
OutboundMedia,
GoalStateWsPayload,
@@ -385,7 +384,6 @@ export class NanobotClient {
content: string,
media?: OutboundMedia[],
options?: {
imageGeneration?: OutboundImageGeneration;
cliApps?: OutboundCliAppMention[];
mcpPresets?: OutboundMcpPresetMention[];
workspaceScope?: WorkspaceScopePayload | null;
@@ -398,7 +396,6 @@ export class NanobotClient {
chat_id: chatId,
content,
...(media && media.length > 0 ? { media } : {}),
...(options?.imageGeneration ? { image_generation: options.imageGeneration } : {}),
...(options?.cliApps?.length ? { cli_apps: options.cliApps } : {}),
...(options?.mcpPresets?.length ? { mcp_presets: options.mcpPresets } : {}),
...(options?.workspaceScope ? { workspace_scope: options.workspaceScope } : {}),
-6
View File
@@ -932,11 +932,6 @@ export interface OutboundMedia {
name?: string;
}
export interface OutboundImageGeneration {
enabled: true;
aspect_ratio?: string | null;
}
export interface OutboundCliAppMention {
name: string;
display_name?: string;
@@ -998,7 +993,6 @@ export type Outbound =
chat_id: string;
content: string;
media?: OutboundMedia[];
image_generation?: OutboundImageGeneration;
cli_apps?: OutboundCliAppMention[];
mcp_presets?: OutboundMcpPresetMention[];
workspace_scope?: WorkspaceScopePayload;
-27
View File
@@ -528,33 +528,6 @@ describe("NanobotClient", () => {
});
});
it("includes image generation options in outbound messages", () => {
const client = new NanobotClient({
url: "ws://test",
reconnect: false,
socketFactory: (url) => new FakeSocket(url) as unknown as WebSocket,
});
client.connect();
lastSocket().fakeOpen();
client.sendMessage(
"chat-img",
"draw a banner",
undefined,
{ imageGeneration: { enabled: true, aspect_ratio: "16:9" } },
);
expect(lastSocket().sent).toContain(
JSON.stringify({
type: "message",
chat_id: "chat-img",
content: "draw a banner",
image_generation: { enabled: true, aspect_ratio: "16:9" },
webui: true,
}),
);
});
it("includes CLI app attachments in outbound messages", () => {
const client = new NanobotClient({
url: "ws://test",
-25
View File
@@ -1640,31 +1640,6 @@ describe("useNanobotStream", () => {
expect(result.current.messages[0].media).toHaveLength(1);
});
it("passes image generation options to the websocket client", () => {
const fake = fakeClient();
const { result } = renderHook(() => useNanobotStream("chat-img", EMPTY_MESSAGES), {
wrapper: wrap(fake.client),
});
act(() => {
result.current.send(
"draw a square icon",
undefined,
{ imageGeneration: { enabled: true, aspect_ratio: "1:1" } },
);
});
expect(fake.client.sendMessage).toHaveBeenCalledWith(
"chat-img",
"draw a square icon",
undefined,
expect.objectContaining({
imageGeneration: { enabled: true, aspect_ratio: "1:1" },
turnId: expect.any(String),
}),
);
});
it("stops the active turn without adding a user slash command bubble", () => {
const fake = fakeClient();
const { result } = renderHook(() => useNanobotStream("chat-stop", EMPTY_MESSAGES), {