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 <img> 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 <?xml ...?> declarations before <svg> 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 <danielhanchen@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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("<svg")) return true;
+  if (lang === "xml" || lang === "html") {
+    const trimmed = codeFence.source.trimStart();
+    // Match <svg directly or <?xml ...?> followed by <svg
+    if (trimmed.startsWith("<svg")) return true;
+    if (trimmed.startsWith("<?xml") && trimmed.includes("<svg")) return true;
+  }
   return false;
 }
 
 function isHtmlFence(codeFence: CodeFence): boolean {
   const lang = codeFence.language?.toLowerCase() ?? "";
-  return lang === "html" && !codeFence.source.trimStart().startsWith("<svg");
+  return lang === "html" && !isSvgFence(codeFence);
 }
 
 const UNSAFE_SVG_RE = /<script[\s>]|on\w+\s*=|javascript:|<foreignObject[\s>]|<iframe[\s>]|<embed[\s>]|<object[\s>]/i;
 
 function sanitizeSvg(source: string): string | null {
   if (UNSAFE_SVG_RE.test(source)) return null;
-  return source;
+  // Strip XML declaration (<?xml ...?>) -- 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 }) {

studio/frontend/src/components/assistant-ui/thread.tsx

@@ -124,7 +124,7 @@ const SUGGESTION_TOOLS: Record<string, Array<"thinking" | "search" | "code">> =
   "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 = {

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 (
     <ToolFallbackRoot>
@@ -133,6 +162,21 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({
               </pre>
             </div>
           ) : null}
+
+          {/* Images from Python tool execution */}
+          {images.length > 0 && sessionId && (
+            <div className="mt-2 flex flex-col gap-2">
+              {images.map((filename) => (
+                <img
+                  key={filename}
+                  src={`/api/inference/sandbox/${encodeURIComponent(sessionId)}/${encodeURIComponent(filename)}${authToken ? `?token=${encodeURIComponent(authToken)}` : ""}`}
+                  alt={filename}
+                  loading="lazy"
+                  className="max-w-full rounded border border-border"
+                />
+              ))}
+            </div>
+          )}
         </div>
       </ToolFallbackContent>
     </ToolFallbackRoot>