diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py index 7bd4a7d6e9..c6b8acfdc4 100644 --- a/studio/backend/core/inference/mcp_client.py +++ b/studio/backend/core/inference/mcp_client.py @@ -298,20 +298,48 @@ def invalidate_tool_cache(server_id: Optional[str] = None) -> None: _probe_cooloff_until.pop(server_id, None) +MCP_IMAGES_SENTINEL = "__MCP_IMAGES__:" +MAX_IMAGE_PAYLOAD_CHARS = 12_000_000 + + def _flatten_result(result: Any) -> str: parts = [] + images = [] + omitted = 0 + budget = MAX_IMAGE_PAYLOAD_CHARS for block in getattr(result, "content", None) or []: text = getattr(block, "text", None) if text: parts.append(str(text)) + continue + data = getattr(block, "data", None) + mime = getattr(block, "mimeType", None) + if data and isinstance(mime, str) and mime.startswith("image/"): + data = str(data) + if len(data) > budget: + omitted += 1 + continue + budget -= len(data) + images.append({"data": data, "mimeType": mime}) body = "\n".join(parts) if not body: structured = getattr(result, "structured_content", None) body = str(structured) if structured is not None else "" + if images or omitted: + notes = [] + if images: + n = len(images) + notes.append(f"{n} image{'s' if n > 1 else ''} attached; displayed to the user") + if omitted: + notes.append(f"{omitted} image{'s' if omitted > 1 else ''} omitted (too large)") + note = f"[{'; '.join(notes)}]" + body = f"{body}\n{note}" if body else note if getattr(result, "is_error", False): # "Error: " prefix triggers tool_call_parser's TOOL_ERROR_PREFIXES nudge. - return f"Error: {body}" if body else "Error: tool returned no content" + body = f"Error: {body}" if body else "Error: tool returned no content" + if images: + body += "\n" + MCP_IMAGES_SENTINEL + json.dumps(images) return body @@ -333,7 +361,10 @@ def call_tool_sync( async def _call() -> Any: async with _client(url, headers, use_oauth) as client: - return await client.call_tool(name, args) + # raise_on_error=False lets an is_error result (which may still carry + # image content) reach _flatten_result instead of FastMCP raising ToolError + # and dropping the images. Transport failures still raise (handled below). + return await client.call_tool(name, args, raise_on_error = False) async def _watch_cancel() -> None: # 50 ms cadence keeps cancellation responsive without busy-looping; diff --git a/studio/backend/core/inference/tool_loop_controller.py b/studio/backend/core/inference/tool_loop_controller.py index cb751ede3d..f595531b90 100644 --- a/studio/backend/core/inference/tool_loop_controller.py +++ b/studio/backend/core/inference/tool_loop_controller.py @@ -233,9 +233,33 @@ def is_tool_error(result: str) -> bool: return isinstance(result, str) and result.lstrip().startswith(TOOL_ERROR_PREFIXES) +def _strip_mcp_image_suffix(result: str) -> str: + """Drop a trailing __MCP_IMAGES__ envelope only when it is the valid JSON + image array appended by _flatten_result, so legit tool text that merely + mentions the marker is not truncated.""" + head, sep, payload = result.rpartition("\n__MCP_IMAGES__:") + if not sep: + return result + try: + images = json.loads(payload) + except (ValueError, RecursionError): + return result + if not isinstance(images, list) or not images: + return result + if not all( + isinstance(img, dict) + and isinstance(img.get("data"), str) + and isinstance(img.get("mimeType"), str) + for img in images + ): + return result + return head.rstrip() + + def strip_result_for_model(result: str) -> str: """Remove frontend-only sentinels (image paths, RAG source map) before feeding the result back to the model.""" + result = _strip_mcp_image_suffix(result) for sentinel in ("__IMAGES__:", "__RAG_SOURCES__:"): if sentinel in result: result = result.split(sentinel, 1)[0].rstrip() diff --git a/studio/backend/tests/test_mcp_flatten_result.py b/studio/backend/tests/test_mcp_flatten_result.py new file mode 100644 index 0000000000..7daee799f9 --- /dev/null +++ b/studio/backend/tests/test_mcp_flatten_result.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import contextlib +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from core.inference import mcp_client +from core.inference.mcp_client import ( + MAX_IMAGE_PAYLOAD_CHARS, + MCP_IMAGES_SENTINEL, + _flatten_result, + call_tool_sync, +) +from core.inference.tool_loop_controller import is_tool_error, strip_result_for_model + +PNG_B64 = "iVBORw0KGgoAAAANSUhEUg==" + + +def _text(value: str) -> SimpleNamespace: + return SimpleNamespace(type = "text", text = value) + + +def _image(data: str = PNG_B64, mime: str = "image/png") -> SimpleNamespace: + return SimpleNamespace(type = "image", data = data, mimeType = mime) + + +def _result( + *blocks, + is_error = False, + structured = None, +) -> SimpleNamespace: + return SimpleNamespace( + content = list(blocks), + is_error = is_error, + structured_content = structured, + ) + + +def test_text_only_result_unchanged(): + assert _flatten_result(_result(_text("hello"))) == "hello" + + +def test_image_only_result_keeps_image_and_notes_model(): + flat = _flatten_result(_result(_image())) + body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1) + assert body == "[1 image attached; displayed to the user]" + assert json.loads(payload) == [{"data": PNG_B64, "mimeType": "image/png"}] + + +def test_text_plus_image_keeps_both(): + flat = _flatten_result(_result(_text("Took a screenshot"), _image())) + body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1) + assert body == "Took a screenshot\n[1 image attached; displayed to the user]" + assert json.loads(payload)[0]["mimeType"] == "image/png" + + +def test_multiple_images_pluralized(): + flat = _flatten_result(_result(_image(), _image(mime = "image/jpeg"))) + body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1) + assert "[2 images attached; displayed to the user]" in body + assert [img["mimeType"] for img in json.loads(payload)] == ["image/png", "image/jpeg"] + + +def test_strip_result_for_model_drops_image_payload(): + flat = _flatten_result(_result(_text("Took a screenshot"), _image())) + stripped = strip_result_for_model(flat) + assert stripped == "Took a screenshot\n[1 image attached; displayed to the user]" + assert PNG_B64 not in stripped + + +def test_strip_preserves_literal_mcp_sentinel_in_text(): + # A tool that legitimately returns text containing the marker (e.g. reading + # source/docs that quote it) must not be truncated: the suffix is not a + # valid JSON image array. + text = "before\n__MCP_IMAGES__: literal from source\nafter" + assert strip_result_for_model(text) == text + + +def test_strip_preserves_non_image_json_after_marker(): + text = 'log line\n__MCP_IMAGES__:["not", "image", "dicts"]' + assert strip_result_for_model(text) == text + + +def test_strip_removes_only_valid_terminal_envelope(): + text = ( + "Earlier mention: __MCP_IMAGES__: is documented here" + "\n[1 image attached; displayed to the user]" + '\n__MCP_IMAGES__:[{"data": "AAAA", "mimeType": "image/png"}]' + ) + assert strip_result_for_model(text) == ( + "Earlier mention: __MCP_IMAGES__: is documented here" + "\n[1 image attached; displayed to the user]" + ) + + +def test_strip_still_handles_images_and_rag_sentinels(): + assert strip_result_for_model("output\n__IMAGES__:['a.png']") == "output" + assert strip_result_for_model("answer\n__RAG_SOURCES__:[{}]") == "answer" + + +def test_error_result_keeps_error_prefix_and_images(): + flat = _flatten_result(_result(_text("boom"), _image(), is_error = True)) + assert flat.startswith("Error: boom") + assert is_tool_error(flat) + assert MCP_IMAGES_SENTINEL in flat + + +def test_image_only_error_no_longer_reports_no_content(): + flat = _flatten_result(_result(_image(), is_error = True)) + assert flat.startswith("Error: [1 image attached") + assert "tool returned no content" not in flat + + +def test_oversized_image_omitted_with_note(): + huge = "A" * (MAX_IMAGE_PAYLOAD_CHARS + 1) + flat = _flatten_result(_result(_image(data = huge))) + assert flat == "[1 image omitted (too large)]" + assert MCP_IMAGES_SENTINEL not in flat + + +def test_oversized_budget_shared_across_images(): + big = "A" * (MAX_IMAGE_PAYLOAD_CHARS - 10) + flat = _flatten_result(_result(_image(data = big), _image())) + body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1) + assert "1 image attached" in body + assert "1 image omitted (too large)" in body + images = json.loads(payload) + assert len(images) == 1 and images[0]["data"] == big + + +def test_non_image_binary_block_still_ignored(): + flat = _flatten_result( + _result(SimpleNamespace(type = "audio", data = PNG_B64, mimeType = "audio/wav")) + ) + assert flat == "" + + +def test_structured_content_fallback_still_used(): + flat = _flatten_result(_result(structured = {"ok": True})) + assert flat == "{'ok': True}" + + +def test_call_tool_sync_passes_raise_on_error_false_and_keeps_error_images(monkeypatch): + # Guards that call_tool_sync passes raise_on_error=False, so an is_error result + # with image content reaches _flatten_result instead of FastMCP raising ToolError. + seen = {} + + class _FakeClient: + async def call_tool( + self, + name, + args, + raise_on_error = True, + ): + seen["raise_on_error"] = raise_on_error + return _result(_text("boom"), _image(), is_error = True) + + @contextlib.asynccontextmanager + async def _fake_client(url, headers, use_oauth): + yield _FakeClient() + + monkeypatch.setattr(mcp_client, "_client", _fake_client) + out = call_tool_sync("http://x", None, "take_screenshot", {}) + + assert seen["raise_on_error"] is False + assert out.startswith("Error: boom") + assert MCP_IMAGES_SENTINEL in out + assert is_tool_error(out) diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index 6d26d075cf..24784ca34c 100644 --- a/studio/backend/tests/test_mcp_servers.py +++ b/studio/backend/tests/test_mcp_servers.py @@ -198,7 +198,12 @@ def test_call_tool_sync_respects_pre_set_cancel_event(monkeypatch): async def __aexit__(self, *args): return False - async def call_tool(self, name, args): + async def call_tool( + self, + name, + args, + raise_on_error = True, + ): import asyncio as _asyncio await _asyncio.sleep(30) # never finishes during the test @@ -520,7 +525,12 @@ def test_call_tool_sync_short_circuits_on_pre_set_cancel(monkeypatch): async def __aexit__(self, *args): return False - async def call_tool(self, name, args): + async def call_tool( + self, + name, + args, + raise_on_error = True, + ): return "ran" monkeypatch.setattr(mcp_client, "_client", lambda *a, **kw: _StubClient()) diff --git a/studio/backend/tests/test_mcp_stdio_pr5863.py b/studio/backend/tests/test_mcp_stdio_pr5863.py index 15fe553fb2..1cb1211cf2 100644 --- a/studio/backend/tests/test_mcp_stdio_pr5863.py +++ b/studio/backend/tests/test_mcp_stdio_pr5863.py @@ -93,7 +93,12 @@ class _RecordingClient: async def list_tools(self): return [_FakeTool("list_directory"), _FakeTool("write_file")] - async def call_tool(self, name, args): + async def call_tool( + self, + name, + args, + raise_on_error = True, + ): return _FakeResult(f"called {name}") diff --git a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx index 20a22c3a4a..d0bc12706e 100644 --- a/studio/frontend/src/components/assistant-ui/tool-fallback.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-fallback.tsx @@ -260,6 +260,30 @@ function ToolFallbackArgs({ ); } +interface McpImageResult { + text: string; + images: { data: string; mimeType: string }[]; +} + +function isMcpImageResult(val: unknown): val is McpImageResult { + if (typeof val !== "object" || val === null) { + return false; + } + const v = val as { text?: unknown; images?: unknown }; + return ( + typeof v.text === "string" && + Array.isArray(v.images) && + v.images.length > 0 && + v.images.every( + (img: unknown) => + typeof img === "object" && + img !== null && + typeof (img as { data?: unknown }).data === "string" && + typeof (img as { mimeType?: unknown }).mimeType === "string", + ) + ); +} + function ToolFallbackResult({ result, className, @@ -271,6 +295,8 @@ function ToolFallbackResult({ return null; } + const imageResult = isMcpImageResult(result) ? result : null; + return (

Result:

-
-        {typeof result === "string" ? result : JSON.stringify(result, null, 2)}
-      
+ {imageResult ? ( + <> + {imageResult.text && ( +
+              {imageResult.text}
+            
+ )} +
+ {imageResult.images.map((img, i) => ( + {`Tool + ))} +
+ + ) : ( +
+          {typeof result === "string" ? result : JSON.stringify(result, null, 2)}
+        
+ )}
); } diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index ca062bf93c..4fb7d9bd7e 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -901,6 +901,33 @@ function serializeAssistantToolCallPart( return entry; } +export interface McpImageToolResult { + text: string; + images: { data: string; mimeType: string }[]; +} + +export function isMcpImageToolResult( + val: unknown, +): val is McpImageToolResult { + if (typeof val !== "object" || val === null) { + return false; + } + const v = val as { text?: unknown; images?: unknown; sessionId?: unknown }; + return ( + typeof v.text === "string" && + v.sessionId === undefined && + Array.isArray(v.images) && + v.images.length > 0 && + v.images.every( + (img: unknown) => + typeof img === "object" && + img !== null && + typeof (img as { data?: unknown }).data === "string" && + typeof (img as { mimeType?: unknown }).mimeType === "string", + ) + ); +} + function serializeToolResultPart( part: ToolCallMessagePart, ): SerializedToolResult | null { @@ -920,6 +947,8 @@ function serializeToolResultPart( // content; serialise a sentinel JSON so legitimately empty tool // outputs still round-trip the follow-up turn to the provider. content = result.length > 0 ? result : JSON.stringify({ result: "" }); + } else if (isMcpImageToolResult(result)) { + content = result.text.length > 0 ? result.text : JSON.stringify({ result: "" }); } else { try { content = JSON.stringify(result); @@ -3196,9 +3225,12 @@ export function createOpenAIStreamAdapter( const rawResult = (toolEvent.result as string) ?? ""; const imgMarker = "\n__IMAGES__:"; const imgIdx = rawResult.lastIndexOf(imgMarker); + const mcpImgMarker = "\n__MCP_IMAGES__:"; + const mcpImgIdx = rawResult.lastIndexOf(mcpImgMarker); let parsedResult: | string | { text: string; images: string[]; sessionId: string } + | McpImageToolResult | { image_b64: string; image_mime: string; @@ -3208,6 +3240,24 @@ export function createOpenAIStreamAdapter( prompt?: string; }; const imageB64 = toolEvent.image_b64 as string | undefined; + // A valid MCP image envelope wins; an invalid marker falls + // through so a sandbox __IMAGES__ suffix still renders and + // legit text round-trips unchanged. + let mcpImages: McpImageToolResult | null = null; + if (mcpImgIdx !== -1) { + try { + const images = JSON.parse( + rawResult.slice(mcpImgIdx + mcpImgMarker.length), + ); + const candidate = { + text: rawResult.slice(0, mcpImgIdx), + images, + }; + if (isMcpImageToolResult(candidate)) mcpImages = candidate; + } catch { + // Not a valid envelope; fall through below. + } + } if ( toolCallParts[idx].toolName === "image_generation" && typeof imageB64 === "string" && @@ -3225,6 +3275,8 @@ export function createOpenAIStreamAdapter( background: toolEvent.background as string | undefined, prompt: toolEvent.prompt as string | undefined, }; + } else if (mcpImages !== null) { + parsedResult = mcpImages; } else if (imgIdx !== -1) { const text = rawResult.slice(0, imgIdx); // Fall back to "_default" to match the backend sandbox diff --git a/studio/frontend/src/features/chat/hooks/use-chat-search-index.ts b/studio/frontend/src/features/chat/hooks/use-chat-search-index.ts index 841f2c2a2f..9badd43bc3 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-search-index.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-search-index.ts @@ -28,13 +28,43 @@ const SEARCH_REBUILD_DEBOUNCE_MS = 300; // Keys whose values are base64 image/audio payloads, not searchable text. const BINARY_KEY = /b64|base64|^(images?|audio|video)$/i; +// Drop a trailing __MCP_IMAGES__ envelope only when it is the valid JSON image +// array appended by the backend, so legit tool text that merely mentions the +// marker stays searchable. (base64 runs below are scrubbed regardless.) +function stripMcpImageSuffix(value: string): string { + const marker = "\n__MCP_IMAGES__:"; + const idx = value.lastIndexOf(marker); + if (idx === -1) return value; + try { + const images: unknown = JSON.parse(value.slice(idx + marker.length)); + if ( + Array.isArray(images) && + images.length > 0 && + images.every( + (img) => + typeof img === "object" && + img !== null && + typeof (img as Record).data === "string" && + typeof (img as Record).mimeType === "string", + ) + ) { + return value.slice(0, idx); + } + } catch { + // Not a valid envelope; leave the text intact. + } + return value; +} + // Readable text from tool args/results, dropping base64 image/audio blobs so // they never bloat the index (object fields by key, plus data URLs / long // base64 runs and the "__IMAGES__" suffix inside strings). function searchableText(value: unknown, depth = 0): string { if (typeof value === "string") { - const cut = value.indexOf("\n__IMAGES__:"); - return (cut === -1 ? value : value.slice(0, cut)) + let text = stripMcpImageSuffix(value); + const cut = text.indexOf("\n__IMAGES__:"); + if (cut !== -1) text = text.slice(0, cut); + return text .replace(/data:[^;,\s]+;base64,[A-Za-z0-9+/=]+/g, " ") .replace(/[A-Za-z0-9+/]{120,}={0,2}/g, " "); } diff --git a/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx b/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx index ee3a49526f..f4546a01a8 100644 --- a/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx +++ b/studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx @@ -54,6 +54,7 @@ import { syncStoredChatMessages, } from "../utils/chat-history-storage"; import { notifyChatHistoryUpdated } from "../api/chat-api"; +import { isMcpImageToolResult } from "../api/chat-adapter"; import { usePlusMenuPrefsStore } from "../stores/plus-menu-prefs-store"; import type { ThreadRecord, MessageRecord } from "../types"; @@ -170,11 +171,14 @@ function contentBlocksToText(content: unknown): string { parts.push("[thinking]\n" + thinkText + "\n[/thinking]"); } } else if (p.type === "tool-call") { + // Keep base64 image payloads out of every export format: use the + // model-visible text for MCP image results (matches chat replay). + const result = isMcpImageToolResult(p.result) ? p.result.text : p.result; parts.push( JSON.stringify({ tool_call: p.toolName, args: p.args, - result: p.result, + result, }), ); } else if (p.type === "image") { @@ -299,7 +303,15 @@ function messageToOpenAI(msg: { role: unknown; content: unknown; attachments?: u const argsStr = p.args != null ? JSON.stringify(p.args) : (typeof p.argsText === "string" ? p.argsText : "{}"); toolCalls.push({ id, type: "function", function: { name, arguments: argsStr } }); if (p.result !== undefined && p.result !== null) { - const resultStr = typeof p.result === "string" ? p.result : JSON.stringify(p.result); + // Keep base64 image payloads out of exports: MCP image results carry + // their model-visible text alongside the data, so serialize the text + // (matching chat replay) instead of the full object. + const resultStr = + typeof p.result === "string" + ? p.result + : isMcpImageToolResult(p.result) + ? p.result.text + : JSON.stringify(p.result); toolResults.push({ role: "tool", tool_call_id: id, name, content: resultStr }); } }