From c8d311a053711348b6b5ca09e2ab810e9bf0ceb7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 2 Apr 2026 05:08:16 -0700 Subject: [PATCH] feat(studio): display images from Python tool execution in chat UI (#4778) * feat(studio): display images from Python tool execution in chat UI When the model calls the Python tool to create a matplotlib plot or other image file, the image now displays inline in the chat output instead of being invisible to the user. Backend: - Detect new image files (png/jpg/gif/webp/bmp) after Python subprocess completes by diffing os.listdir before/after execution - Append __IMAGES__ sentinel to tool result for frontend consumption - Strip sentinel before injecting result into LLM context (role: tool) so the model never sees file paths - Add GET /sandbox/{session_id}/{filename} endpoint with JWT auth (header or query param), path traversal protection, extension allowlist, realpath containment check, and nosniff header Frontend: - Parse __IMAGES__ sentinel in tool_end SSE events, create structured result with text/images/sessionId - Render tags in Python tool UI pointing at the sandbox endpoint Also fixes a bug where SyntaxError in user code was misreported as "unsafe code detected" instead of showing the actual Python traceback. The _check_code_safety function now lets SyntaxError pass through to the subprocess for a proper error message. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio): improve SVG detection and strip XML preamble Handle declarations before tags in code fences, strip XML declaration from SVGs before data URI rendering, and update the sloth suggestion prompt to request showing code. * fix(studio): persist parentId so retries survive reload The append() handler was destructuring only { message } from ExportedMessageRepositoryItem and discarding parentId. When loading a saved thread, load() used ExportedMessageRepository.fromArray() which chains all messages sequentially, flattening retry branches into a linear list. Now append() writes parentId to the MessageRecord, and load() reconstructs the tree when parentIds are present. Old threads without parentId fall back to the existing fromArray() behavior. * fix(studio): address review findings for image display and retry persistence Image detection: - Use mtime comparison instead of filename-only diff so overwritten files (e.g. plt.savefig("chart.png") called twice) are detected Sentinel parsing: - Use rsplit/lastIndexOf instead of split/indexOf so user code that prints __IMAGES__: does not collide with the backend sentinel Mixed legacy/new threads: - For old messages without a stored parentId, infer sequential parent from the previous message instead of null, preventing multiple roots Sandbox endpoint: - Change Cache-Control from "public, max-age=3600" to "private, no-store" since these are authenticated responses --------- Co-authored-by: Daniel Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 7 +- studio/backend/core/inference/tools.py | 45 ++++++++- studio/backend/routes/inference.py | 91 ++++++++++++++++++- .../components/assistant-ui/markdown-text.tsx | 13 ++- .../src/components/assistant-ui/thread.tsx | 2 +- .../assistant-ui/tool-ui-python.tsx | 56 ++++++++++-- .../src/features/chat/api/chat-adapter.ts | 18 +++- .../src/features/chat/runtime-provider.tsx | 29 +++++- studio/frontend/src/features/chat/types.ts | 1 + 9 files changed, 245 insertions(+), 17 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 1b7fe548e4..80dda61b02 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2905,10 +2905,15 @@ class LlamaCppBackend: _error_prefixes ) _tool_call_history.append((_tc_key, _is_error)) + # Strip image sentinel before feeding result to the LLM + # (the full result with sentinel is still yielded via + # tool_end so the frontend can extract image paths). _result_content = result + if "\n__IMAGES__:" in _result_content: + _result_content = _result_content.rsplit("\n__IMAGES__:", 1)[0] if _is_error: _result_content = ( - result + "\n\nThe tool call encountered an issue. " + _result_content + "\n\nThe tool call encountered an issue. " "Please try a different approach or rephrase your request." ) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index d425daa49d..b23372b766 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -26,6 +26,10 @@ from loggers import get_logger logger = get_logger(__name__) _EXEC_TIMEOUT = 300 # 5 minutes + +# Strict raster-image allowlist for sandbox file serving. +# No .svg (XSS risk via embedded scripts), no .html, no .pdf. +_IMAGE_EXTS = frozenset({".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"}) _MAX_OUTPUT_CHARS = 8000 # truncate long output _BASH_BLOCKED_WORDS = {"rm", "sudo", "dd", "chmod", "mkfs", "shutdown", "reboot"} @@ -591,6 +595,12 @@ def _check_code_safety(code: str) -> str | None: """ safe, info = _check_signal_escape_patterns(code) if not safe: + # SyntaxError from ast.parse -- let these through so the subprocess + # produces a normal Python traceback instead of a misleading + # "unsafe code detected" message. + if info.get("error"): + return None + reasons = [ item.get("description", "") for item in info.get("signal_tampering", []) ] @@ -634,6 +644,17 @@ def _python_exec( tmp_path = None workdir = _get_workdir(session_id) + # Snapshot image mtimes so we detect both new and overwritten files. + _before: dict[str, int] = {} + if os.path.isdir(workdir): + for _name in os.listdir(workdir): + if os.path.splitext(_name)[1].lower() in _IMAGE_EXTS: + _p = os.path.join(workdir, _name) + if os.path.isfile(_p): + try: + _before[_name] = os.stat(_p).st_mtime_ns + except OSError: + pass try: fd, tmp_path = tempfile.mkstemp( suffix = ".py", prefix = "studio_exec_", dir = workdir @@ -669,7 +690,29 @@ def _python_exec( result = output or "" if proc.returncode != 0: result = f"Exit code {proc.returncode}:\n{result}" - return _truncate(result) if result.strip() else "(no output)" + result = _truncate(result) if result.strip() else "(no output)" + + # Detect new or overwritten image files and append sentinel for frontend + if session_id and os.path.isdir(workdir): + new_images = [] + for _name in os.listdir(workdir): + if os.path.splitext(_name)[1].lower() not in _IMAGE_EXTS: + continue + _p = os.path.join(workdir, _name) + if not os.path.isfile(_p): + continue + try: + _mtime = os.stat(_p).st_mtime_ns + except OSError: + continue + if _name not in _before or _mtime != _before[_name]: + new_images.append(_name) + if new_images: + import json as _json + + result += f"\n__IMAGES__:{_json.dumps(sorted(new_images))}" + + return result except Exception as e: return f"Execution error: {e}" diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 0c1d7e37f9..8f058bd0f2 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5,11 +5,12 @@ Inference API routes for model loading and text generation. """ +import os import sys import time import uuid from pathlib import Path -from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi import APIRouter, Depends, HTTPException, Request, status from fastapi.responses import StreamingResponse, JSONResponse from typing import Optional import json @@ -1680,6 +1681,94 @@ async def openai_chat_completions( raise HTTPException(status_code = 500, detail = str(e)) +# ===================================================================== +# Sandbox file serving (/sandbox/{session_id}/{filename}) +# ===================================================================== + +_SANDBOX_MEDIA_TYPES = { + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".gif": "image/gif", + ".webp": "image/webp", + ".bmp": "image/bmp", +} + + +@router.get("/sandbox/{session_id}/{filename}") +async def serve_sandbox_file( + session_id: str, + filename: str, + request: Request, + token: Optional[str] = None, +): + """ + Serve image files created by Python tool execution. + + Accepts auth via Authorization header OR ?token= query param + (needed because cannot send custom headers). + """ + from fastapi.responses import FileResponse + + # ── Authentication (header or query param) ────────────────── + auth_header = request.headers.get("authorization") + if auth_header and auth_header.lower().startswith("bearer "): + jwt_token = auth_header[7:] + elif token: + jwt_token = token + else: + raise HTTPException( + status_code = status.HTTP_401_UNAUTHORIZED, + detail = "Missing authentication token", + ) + from fastapi.security import HTTPAuthorizationCredentials + + creds = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = jwt_token) + await get_current_subject(creds) + + # ── Filename sanitization ─────────────────────────────────── + safe_filename = os.path.basename(filename) + if not safe_filename or safe_filename in (".", ".."): + raise HTTPException(status_code = 404, detail = "Not found") + + # ── Extension allowlist ───────────────────────────────────── + ext = os.path.splitext(safe_filename)[1].lower() + media_type = _SANDBOX_MEDIA_TYPES.get(ext) + if not media_type: + raise HTTPException( + status_code = status.HTTP_403_FORBIDDEN, + detail = "File type not allowed", + ) + + # ── Path containment check ────────────────────────────────── + home = os.path.expanduser("~") + sandbox_root = os.path.realpath(os.path.join(home, "studio_sandbox")) + safe_session = os.path.basename(session_id.replace("..", "")) + if not safe_session: + raise HTTPException(status_code = 404, detail = "Not found") + + file_path = os.path.realpath( + os.path.join(sandbox_root, safe_session, safe_filename) + ) + if not file_path.startswith(sandbox_root + os.sep): + raise HTTPException( + status_code = status.HTTP_403_FORBIDDEN, + detail = "Access denied", + ) + + if not os.path.isfile(file_path): + raise HTTPException(status_code = 404, detail = "Not found") + + return FileResponse( + path = file_path, + media_type = media_type, + headers = { + "Cache-Control": "private, no-store", + "X-Content-Type-Options": "nosniff", + }, + ) + + # ===================================================================== # OpenAI-Compatible Models Listing (/models → /v1/models) # ===================================================================== diff --git a/studio/frontend/src/components/assistant-ui/markdown-text.tsx b/studio/frontend/src/components/assistant-ui/markdown-text.tsx index 17a6f8f054..91ef78fcf9 100644 --- a/studio/frontend/src/components/assistant-ui/markdown-text.tsx +++ b/studio/frontend/src/components/assistant-ui/markdown-text.tsx @@ -100,20 +100,27 @@ function getCodeFilename(language: string | null) { function isSvgFence(codeFence: CodeFence): boolean { const lang = codeFence.language?.toLowerCase() ?? ""; if (lang === "svg") return true; - if ((lang === "xml" || lang === "html") && codeFence.source.trimStart().startsWith(" followed by ]|on\w+\s*=|javascript:|]|]|]|]/i; function sanitizeSvg(source: string): string | null { if (UNSAFE_SVG_RE.test(source)) return null; - return source; + // Strip XML declaration () -- not needed for data URI + // rendering and can cause issues with some renderers. + return source.replace(/^\s*<\?xml[^?]*\?>\s*/i, ""); } function SvgPreview({ source }: { source: string }) { diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 5170f6ff69..4db0eda0ad 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -124,7 +124,7 @@ const SUGGESTION_TOOLS: Record> = "How do you fine-tune an audio model with Unsloth?": ["thinking", "search"], "Create a live weather dashboard in HTML using no API key. Show me the code": ["thinking", "code", "search"], "Solve the integral of x·sin(x), and verify it step by step": ["thinking", "code"], - "Draw an SVG of a cute sloth": ["thinking", "code", "search"], + "Draw an SVG of a cute sloth & show the code": ["thinking", "code", "search"], }; const toolIconMap = { diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx index a510ed0d9e..6aa590ae11 100644 --- a/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx @@ -4,6 +4,7 @@ "use client"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; +import { getAuthToken } from "@/features/auth/session"; import type { ToolCallMessagePartComponent } from "@assistant-ui/react"; import { code as codePlugin } from "@streamdown/code"; import { CheckIcon, CodeIcon, CopyIcon, LoaderIcon } from "lucide-react"; @@ -15,6 +16,12 @@ import { ToolFallbackTrigger, } from "./tool-fallback"; +interface StructuredResult { + text: string; + images: string[]; + sessionId: string; +} + const MAX_DISPLAY = 10_000; const COPY_RESET_MS = 2000; const SHIKI_THEME = ["github-light", "github-dark"] as ["github-light", "github-dark"]; @@ -84,6 +91,16 @@ function HighlightedCode({ code: source, language }: { code: string; language: s ); } +function isStructuredResult(val: unknown): val is StructuredResult { + return ( + typeof val === "object" && + val !== null && + "text" in val && + "images" in val && + "sessionId" in val + ); +} + const PythonToolUIImpl: ToolCallMessagePartComponent = ({ args, result, @@ -92,12 +109,24 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({ const code = (args as { code?: string })?.code ?? ""; const firstLine = code.split("\n")[0]?.slice(0, 60) ?? ""; const isRunning = status?.type === "running"; - const output = - typeof result === "string" - ? result - : result - ? JSON.stringify(result, null, 2) - : ""; + + let output: string; + let images: string[] = []; + let sessionId = ""; + + if (isStructuredResult(result)) { + output = result.text; + images = result.images; + sessionId = result.sessionId; + } else if (typeof result === "string") { + output = result; + } else if (result) { + output = JSON.stringify(result, null, 2); + } else { + output = ""; + } + + const authToken = getAuthToken(); return ( @@ -133,6 +162,21 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({ ) : null} + + {/* Images from Python tool execution */} + {images.length > 0 && sessionId && ( +
+ {images.map((filename) => ( + {filename} + ))} +
+ )} diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 2b8a259930..e287daf33a 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -635,7 +635,23 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { toolCallParts[toolCallParts.length - 1]?.toolCallId || ""; const idx = toolCallParts.findIndex((p) => p.toolCallId === id); if (idx !== -1) { - toolCallParts[idx] = { ...toolCallParts[idx], result: toolEvent.result as string }; + const rawResult = (toolEvent.result as string) ?? ""; + const imgMarker = "\n__IMAGES__:"; + const imgIdx = rawResult.lastIndexOf(imgMarker); + let parsedResult: string | { text: string; images: string[]; sessionId: string }; + if (imgIdx !== -1) { + const text = rawResult.slice(0, imgIdx); + const sessionId = unstable_threadId || ""; + try { + const images = JSON.parse(rawResult.slice(imgIdx + imgMarker.length)) as string[]; + parsedResult = { text, images, sessionId }; + } catch { + parsedResult = rawResult; + } + } else { + parsedResult = rawResult; + } + toolCallParts[idx] = { ...toolCallParts[idx], result: parsedResult }; } } // Yield cumulative state so tool UI updates (tools first, text after) diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index 404271f896..02f792b509 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -47,9 +47,9 @@ const DEFAULT_SUGGESTIONS = [ prompt: "Solve the integral of x·sin(x), and verify it step by step", }, { - title: "Draw an SVG of a cute sloth", + title: "Draw an SVG of a cute sloth & show the code", label: "SVG sloth", - prompt: "Draw an SVG of a cute sloth", + prompt: "Draw an SVG of a cute sloth & show the code", }, ]; @@ -569,10 +569,32 @@ function ThreadHistoryProvider({ store.setContextUsage(savedUsage); } + // If any message has a stored parentId, reconstruct the tree + // so retries/regenerations load as branches instead of being + // unrolled into a flat list. For mixed legacy/new threads + // (old messages without parentId + new messages with), infer + // sequential parents for old messages to preserve the chain. + // Fall back to fromArray for fully legacy threads. + const hasParentIds = msgs.some((m) => "parentId" in m); + if (hasParentIds) { + let previousId: string | null = null; + return { + messages: msgs.map((m) => { + const parentId = "parentId" in m + ? (m.parentId ?? null) + : previousId; + previousId = m.id; + return { + parentId, + message: toThreadMessage(m), + }; + }), + }; + } return ExportedMessageRepository.fromArray(msgs.map(toThreadMessage)); }, - async append({ message }: ExportedMessageRepositoryItem) { + async append({ parentId, message }: ExportedMessageRepositoryItem) { const { remoteId } = await aui.threadListItem().initialize(); const content = cloneContent(message.content); const attachments = @@ -586,6 +608,7 @@ function ThreadHistoryProvider({ await db.messages.put({ id: message.id, threadId: remoteId, + parentId: parentId ?? null, role: message.role, content, ...(attachments.length > 0 && { attachments }), diff --git a/studio/frontend/src/features/chat/types.ts b/studio/frontend/src/features/chat/types.ts index 11e53d76d3..1f370b6ac1 100644 --- a/studio/frontend/src/features/chat/types.ts +++ b/studio/frontend/src/features/chat/types.ts @@ -20,6 +20,7 @@ export interface ThreadRecord { export interface MessageRecord { id: string; threadId: string; + parentId?: string | null; role: import("@assistant-ui/react").ThreadMessage["role"]; content: import("@assistant-ui/react").ThreadMessage["content"]; attachments?: import("@assistant-ui/react").ThreadMessage["attachments"];