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"];