Studio: render generated images inline for the Images pill

The pill wired the request end of the loop but the response was lost on the client: the backend emits a `tool_end` _toolEvent carrying the base64 PNG on `image_b64` / `image_mime`, but the chat-adapter only read the `result` string and the generic ToolFallback printed the prompt as JSON args with an empty Result block -- the "I see no image" symptom in the chat. - chat-adapter: when the closing `tool_end` is for `image_generation`, repackage `image_b64` + `image_mime` (+ size/quality/background) into a structured result object instead of dropping them. - New `ImageGenerationToolUI` reads that result and renders the image inline via `<img src="data:image/...;base64,...">` with the prompt as a caption. Falls back to a spinner while the request is still running. - Register the component under `image_generation` in thread.tsx's tools.by_name map so it preempts ToolFallback for this tool only.
Merge remote-tracking branch 'origin/main' into feat/openai-image-generation-pill
2026-05-22 14:20:23 +00:00 · 2026-05-22 14:07:55 +00:00 · 2026-05-22 13:59:48 +00:00
7 changed files with 354 additions and 9 deletions
--- a/studio/frontend/src/components/assistant-ui/thread.tsx
+++ b/studio/frontend/src/components/assistant-ui/thread.tsx
@ -18,6 +18,7 @@ import {
 import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
 import { ToolGroup } from "@/components/assistant-ui/tool-group";
 import { CodeExecutionToolUI } from "@/components/assistant-ui/tool-ui-code-execution";
+import { ImageGenerationToolUI } from "@/components/assistant-ui/tool-ui-image-generation";
 import { PythonToolUI } from "@/components/assistant-ui/tool-ui-python";
 import { TerminalToolUI } from "@/components/assistant-ui/tool-ui-terminal";
 import { WebSearchToolUI } from "@/components/assistant-ui/tool-ui-web-search";
@ -68,6 +69,7 @@ import {
  DownloadIcon,
  GlobeIcon,
  HeadphonesIcon,
+  ImageIcon,
  LightbulbIcon,
  LightbulbOffIcon,
  MicIcon,
@ -870,6 +872,42 @@ const CodeToolsToggle: FC = () => {
  );
 };

+const ImagesToggle: FC = () => {
+  const modelLoaded = useChatRuntimeStore(
+    (s) => !!s.params.checkpoint && !s.modelLoading,
+  );
+  // OpenAI cloud Responses-API models advertise image_generation as a
+  // server-side tool; no local runtime fallback exists. Mirror of
+  // shared-composer's imageDisabled / showImagePill so the in-thread
+  // composer surfaces the same control as the empty-state composer.
+  const supportsBuiltinImageGeneration = useChatRuntimeStore(
+    (s) => s.supportsBuiltinImageGeneration,
+  );
+  const imageToolsEnabled = useChatRuntimeStore((s) => s.imageToolsEnabled);
+  const setImageToolsEnabled = useChatRuntimeStore(
+    (s) => s.setImageToolsEnabled,
+  );
+  if (!supportsBuiltinImageGeneration) {
+    return null;
+  }
+  const disabled = !modelLoaded;
+  return (
+    <button
+      type="button"
+      disabled={disabled}
+      onClick={() => setImageToolsEnabled(!imageToolsEnabled)}
+      className="composer-pill-btn"
+      data-active={imageToolsEnabled && !disabled ? "true" : "false"}
+      aria-label={
+        imageToolsEnabled ? "Disable image generation" : "Enable image generation"
+      }
+    >
+      <ImageIcon className="size-3.5" />
+      <span>Images</span>
+    </button>
+  );
+};
+
 const ToolStatusDisplay: FC = () => {
  const toolStatus = useChatRuntimeStore((s) => s.toolStatus);
  const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);
@ -941,6 +979,7 @@ const ComposerAction: FC<{ disabled?: boolean; blockSend?: () => boolean }> = ({
        <PreserveThinkingToggle />
        <WebSearchToggle />
        <CodeToolsToggle />
+        <ImagesToggle />
      </div>
      <div className="flex items-center gap-1">
        <ComposerPrimitive.If dictation={false}>
@ -1067,6 +1106,7 @@ const AssistantMessage: FC = () => {
                python: PythonToolUI,
                terminal: TerminalToolUI,
                code_execution: CodeExecutionToolUI,
+                image_generation: ImageGenerationToolUI,
              },
              Fallback: ToolFallback,
            },
--- a/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx
+++ b/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx
@ -0,0 +1,132 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"use client";
+
+import { type ToolCallMessagePartComponent, useAuiState } from "@assistant-ui/react";
+import { ImageIcon, LoaderIcon } from "lucide-react";
+import { memo, useEffect, useState } from "react";
+import {
+  ToolFallbackContent,
+  ToolFallbackRoot,
+  ToolFallbackTrigger,
+} from "./tool-fallback";
+
+/**
+ * Renders the synthetic `_toolEvent` chunks emitted by
+ * `_stream_openai_responses` when OpenAI's Responses-API
+ * `image_generation` tool fires. The backend stashes the base64
+ * PNG/WebP/JPEG (the gpt-image backbone output) on an `image_b64`
+ * field of the tool_end event so the JSON result stays small, and the
+ * adapter repackages it into a structured `result` shape:
+ *
+ *   {
+ *     image_b64: string,
+ *     image_mime: string,        // e.g. "image/png"
+ *     size?: string,             // "1024x1024" etc
+ *     quality?: string,
+ *     background?: string,
+ *   }
+ *
+ * The corresponding `tool_start` carries the prompt as
+ * `args.prompt` (after gpt-image's revision pass) plus `args.kind:
+ * "image"`. Without this component the generic ToolFallback would
+ * print the prompt as JSON args text with an empty Result block --
+ * which is exactly the "no image" symptom users hit before this UI
+ * landed.
+ */
+interface ImageGenerationArgs {
+  prompt?: string;
+  kind?: string;
+}
+
+interface ImageGenerationResult {
+  image_b64?: string;
+  image_mime?: string;
+  size?: string;
+  quality?: string;
+  background?: string;
+}
+
+const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({
+  args,
+  result,
+  status,
+}) => {
+  const parsedArgs = (args as ImageGenerationArgs) ?? {};
+  const prompt = parsedArgs.prompt ?? "";
+  const isRunning = status?.type === "running";
+
+  const isImageResult =
+    !!result &&
+    typeof result === "object" &&
+    typeof (result as ImageGenerationResult).image_b64 === "string";
+  const imageResult = isImageResult ? (result as ImageGenerationResult) : null;
+  const mime = imageResult?.image_mime || "image/png";
+  const imageSrc = imageResult?.image_b64
+    ? `data:${mime};base64,${imageResult.image_b64}`
+    : null;
+
+  // Collapse the card once the model has resumed streaming prose
+  // after the image. Mirrors CodeExecutionToolUI so the inline image
+  // doesn't collapse mid-stream and the user can click to re-expand.
+  const hasText = useAuiState(({ message }) =>
+    message.content.some(
+      (p) =>
+        p.type === "text" &&
+        "text" in p &&
+        (p as { text: string }).text.length > 0,
+    ),
+  );
+  const [open, setOpen] = useState(true);
+  useEffect(() => {
+    if (isRunning) {
+      setOpen(true);
+    } else if (hasText && !imageSrc) {
+      setOpen(false);
+    }
+  }, [isRunning, hasText, imageSrc]);
+
+  const runningLabel = "Generating image…";
+  const completedLabel = prompt
+    ? prompt.length > 80
+      ? `Generated image: ${prompt.slice(0, 80)}…`
+      : `Generated image: ${prompt}`
+    : "Generated image";
+
+  return (
+    <ToolFallbackRoot open={open} onOpenChange={setOpen}>
+      <ToolFallbackTrigger
+        toolName={isRunning ? runningLabel : completedLabel}
+        status={status}
+        icon={ImageIcon}
+      />
+      <ToolFallbackContent>
+        {isRunning && !imageSrc ? (
+          <div className="flex items-center gap-2 text-sm text-muted-foreground">
+            <LoaderIcon className="size-3.5 animate-spin" />
+            <span>{runningLabel}</span>
+          </div>
+        ) : imageSrc ? (
+          <figure className="m-0 flex flex-col gap-1.5">
+            <img
+              src={imageSrc}
+              alt={prompt || "Generated image"}
+              className="max-w-full rounded-md border border-border/60"
+            />
+            {prompt ? (
+              <figcaption className="text-xs leading-snug text-muted-foreground">
+                {prompt}
+              </figcaption>
+            ) : null}
+          </figure>
+        ) : null}
+      </ToolFallbackContent>
+    </ToolFallbackRoot>
+  );
+};
+
+export const ImageGenerationToolUI = memo(
+  ImageGenerationToolUIImpl,
+) as unknown as ToolCallMessagePartComponent;
+ImageGenerationToolUI.displayName = "ImageGenerationToolUI";
--- a/studio/frontend/src/features/chat/api/chat-adapter.ts
+++ b/studio/frontend/src/features/chat/api/chat-adapter.ts
@ -23,6 +23,7 @@ import {
  getExternalReasoningCapabilities,
  getProviderCapabilities,
  providerSupportsBuiltinCodeExecution,
+  providerSupportsBuiltinImageGeneration,
  providerSupportsBuiltinWebFetch,
  providerSupportsBuiltinWebSearch,
 } from "../provider-capabilities";
@ -830,7 +831,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
      // Re-read store after potential auto-load / model ready wait
      runtime = useChatRuntimeStore.getState();
      const { params } = runtime;
-      const { supportsTools, toolsEnabled, codeToolsEnabled } = runtime;
+      const { supportsTools, toolsEnabled, codeToolsEnabled, imageToolsEnabled } = runtime;
      const externalSelection = parseExternalModelId(params.checkpoint);
      const isExternalRequest = externalSelection !== null;
      if (
@ -900,6 +901,20 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
        externalProvider &&
          providerSupportsBuiltinWebFetch(externalProvider.providerType),
      );
+      // OpenAI Responses-API image_generation server tool. Pill is
+      // gated on OpenAI cloud + a Responses-API model id; the backend
+      // additionally re-checks is_openai_cloud before appending
+      // {type:"image_generation"} to the request tools array.
+      const imageGenerationEnabledForThisTurn = Boolean(
+        externalProvider &&
+          externalSelection &&
+          imageToolsEnabled &&
+          providerSupportsBuiltinImageGeneration(
+            externalProvider.providerType,
+            externalSelection.modelId,
+            externalProvider.baseUrl,
+          ),
+      );

      const outboundMessages = messages
        .map(toOpenAIMessage)
@ -1396,7 +1411,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
              // body["tools"] inside _stream_anthropic.
              ...(webSearchEnabledForThisTurn ||
              webFetchEnabledForThisTurn ||
-              codeExecEnabledForThisTurn
+              codeExecEnabledForThisTurn ||
+              imageGenerationEnabledForThisTurn
                ? {
                    enable_tools: true,
                    enabled_tools: [
@ -1410,6 +1426,13 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                      // tool. There is no separate UI toggle yet.
                      ...(webFetchEnabledForThisTurn ? ["web_fetch"] : []),
                      ...(codeExecEnabledForThisTurn ? ["code_execution"] : []),
+                      // OpenAI Responses-API only: `image_generation`
+                      // returns inline image_generation_call output
+                      // items; the backend's _stream_openai_responses
+                      // path translates them to assistant tool events.
+                      ...(imageGenerationEnabledForThisTurn
+                        ? ["image_generation"]
+                        : []),
                    ],
                  }
                : {}),
@ -1587,8 +1610,36 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                    const imgIdx = rawResult.lastIndexOf(imgMarker);
                    let parsedResult:
                      | string
-                      | { text: string; images: string[]; sessionId: string };
-                    if (imgIdx !== -1) {
+                      | { text: string; images: string[]; sessionId: string }
+                      | {
+                          image_b64: string;
+                          image_mime: string;
+                          size?: string;
+                          quality?: string;
+                          background?: string;
+                        };
+                    const imageB64 = toolEvent.image_b64 as string | undefined;
+                    if (
+                      toolCallParts[idx].toolName === "image_generation" &&
+                      typeof imageB64 === "string" &&
+                      imageB64
+                    ) {
+                      // OpenAI Responses image_generation_call: the
+                      // backend stashes the base64 PNG/WebP/JPEG on
+                      // separate `image_b64` / `image_mime` fields on
+                      // the synthetic _toolEvent so the JSON result
+                      // string stays small enough to log. Repackage as
+                      // a structured result for the dedicated tool UI.
+                      parsedResult = {
+                        image_b64: imageB64,
+                        image_mime:
+                          (toolEvent.image_mime as string | undefined) ??
+                          "image/png",
+                        size: toolEvent.size as string | undefined,
+                        quality: toolEvent.quality as string | undefined,
+                        background: toolEvent.background as string | undefined,
+                      };
+                    } else if (imgIdx !== -1) {
                      const text = rawResult.slice(0, imgIdx);
                      // Fall back to "_default" to match the backend sandbox directory
                      // used when no session_id is provided (see tools.py _get_workdir).
--- a/studio/frontend/src/features/chat/chat-page.tsx
+++ b/studio/frontend/src/features/chat/chat-page.tsx
@ -56,6 +56,7 @@ import {
  getExternalReasoningCapabilities,
  getProviderCapabilities,
  providerSupportsBuiltinCodeExecution,
+  providerSupportsBuiltinImageGeneration,
  providerSupportsBuiltinWebSearch,
 } from "./provider-capabilities";
 import { ChatRuntimeProvider } from "./runtime-provider";
@ -68,6 +69,7 @@ import {
 } from "./shared-composer";
 import {
  CHAT_CODE_TOOLS_ENABLED_KEY,
+  CHAT_IMAGE_TOOLS_ENABLED_KEY,
  CHAT_TOOLS_ENABLED_KEY,
  loadOptionalBool,
  useChatRuntimeStore,
@ -771,6 +773,12 @@ export function ChatPage(): ReactElement {
      selection.modelId,
      provider?.baseUrl,
    );
+    const supportsBuiltinImageGeneration =
+      providerSupportsBuiltinImageGeneration(
+        provider?.providerType,
+        selection.modelId,
+        provider?.baseUrl,
+      );
    // Kimi's k2.6/k2.5 default to thinking enabled on the server side
    // (per https://platform.kimi.ai/docs/models). Mirror that default
    // in the UI so the Think pill comes up clicked when the user picks
@ -790,6 +798,9 @@ export function ChatPage(): ReactElement {
        provider?.providerType === "openai");
    const storedToolsEnabled = loadOptionalBool(CHAT_TOOLS_ENABLED_KEY);
    const storedCodeToolsEnabled = loadOptionalBool(CHAT_CODE_TOOLS_ENABLED_KEY);
+    const storedImageToolsEnabled = loadOptionalBool(
+      CHAT_IMAGE_TOOLS_ENABLED_KEY,
+    );
    const nextToolsEnabled = supportsBuiltinWebSearch
      ? isKimi
        ? false
@ -811,19 +822,25 @@ export function ChatPage(): ReactElement {
        : state.reasoningEnabled,
      supportsPreserveThinking: false,
      // External models never give us a local tool runtime (no
-      // python sandbox), so `supportsTools` must be false. The two
+      // python sandbox), so `supportsTools` must be false. The three
      // `supportsBuiltin*` flags pick up the slack for providers that
      // run the tool server-side: `supportsBuiltinWebSearch` lights
      // up the Search pill (OpenAI / Anthropic / OpenRouter / Kimi),
      // `supportsBuiltinCodeExecution` lights up the Code pill
-      // (Anthropic Claude 4.x only, today).
+      // (Anthropic Claude 4.x and OpenAI gpt-5.5), and
+      // `supportsBuiltinImageGeneration` lights up the Images pill
+      // (OpenAI cloud Responses-API models only).
      supportsTools: false,
      supportsBuiltinWebSearch,
      supportsBuiltinCodeExecution,
+      supportsBuiltinImageGeneration,
      toolsEnabled: nextToolsEnabled,
      codeToolsEnabled: supportsBuiltinCodeExecution
        ? (storedCodeToolsEnabled ?? false)
        : false,
+      imageToolsEnabled: supportsBuiltinImageGeneration
+        ? (storedImageToolsEnabled ?? false)
+        : false,
    });
  }, [externalProvidersForChat, inferenceParams.checkpoint]);
  const canCompare = useMemo(() => {
@ -988,6 +1005,12 @@ export function ChatPage(): ReactElement {
          selectedExternal?.modelId,
          selectedProvider?.baseUrl,
        );
+        const supportsBuiltinImageGeneration =
+          providerSupportsBuiltinImageGeneration(
+            selectedProvider?.providerType,
+            selectedExternal?.modelId,
+            selectedProvider?.baseUrl,
+          );
        // See sibling useEffect above: Kimi's k2.x default to thinking
        // enabled, so the Think pill comes up clicked. Search pill stays
        // off by default; mutual exclusion flips them via the composer.
@ -1003,6 +1026,9 @@ export function ChatPage(): ReactElement {
        const storedCodeToolsEnabled = loadOptionalBool(
          CHAT_CODE_TOOLS_ENABLED_KEY,
        );
+        const storedImageToolsEnabled = loadOptionalBool(
+          CHAT_IMAGE_TOOLS_ENABLED_KEY,
+        );
        const nextToolsEnabled = supportsBuiltinWebSearch
          ? isKimi
            ? false
@ -1029,18 +1055,24 @@ export function ChatPage(): ReactElement {
            : store.reasoningEnabled,
          supportsPreserveThinking: false,
          // External models have no local tool runtime → supportsTools
-          // stays false. The two supportsBuiltin* flags carry the
+          // stays false. The three supportsBuiltin* flags carry the
          // server-side capability info for each pill:
          //   - Search → providerSupportsBuiltinWebSearch
          //   - Code   → providerSupportsBuiltinCodeExecution
-          //              (Anthropic Claude 4.x only, today)
+          //              (Anthropic Claude 4.x + OpenAI gpt-5.5)
+          //   - Images → providerSupportsBuiltinImageGeneration
+          //              (OpenAI cloud Responses-API models)
          supportsTools: false,
          supportsBuiltinWebSearch,
          supportsBuiltinCodeExecution,
+          supportsBuiltinImageGeneration,
          toolsEnabled: nextToolsEnabled,
          codeToolsEnabled: supportsBuiltinCodeExecution
            ? (storedCodeToolsEnabled ?? false)
            : false,
+          imageToolsEnabled: supportsBuiltinImageGeneration
+            ? (storedImageToolsEnabled ?? false)
+            : false,
          ...(stillOnOpenRouterFree ? {} : { lastOpenRouterChosenModel: null }),
        });
        return;
--- a/studio/frontend/src/features/chat/provider-capabilities.ts
+++ b/studio/frontend/src/features/chat/provider-capabilities.ts
@ -217,6 +217,44 @@ export function providerSupportsBuiltinCodeExecution(
  return false;
 }

+/**
+ * Whether the selected external provider/model exposes OpenAI's
+ * Responses-API server-side image_generation tool. Lit on for OpenAI
+ * cloud (`api.openai.com`) when the picked model is a Responses-API
+ * family id (gpt-5.x today). The backend additionally gates on
+ * `is_openai_cloud`; mirror that here so the pill is hidden on custom
+ * OpenAI-compat backends (ollama / llama.cpp / vLLM) that report
+ * `provider_type="openai"` but would 400 on a `{type:"image_generation"}`
+ * tool. See backend/core/inference/external_provider.py near line 2770
+ * for the dispatch and backend/tests/test_openai_image_generation.py
+ * for the round-trip coverage.
+ */
+const OPENAI_IMAGE_GENERATION_MODEL_PREFIXES = [
+  "gpt-5.5-pro",
+  "gpt-5.5",
+  "gpt-5.4-pro",
+  "gpt-5.4",
+  "gpt-5.3",
+  "gpt-5.2",
+  "gpt-5.1",
+  "gpt-5",
+  "o3",
+] as const;
+
+export function providerSupportsBuiltinImageGeneration(
+  providerType: string | null | undefined,
+  modelId: string | null | undefined,
+  baseUrl?: string | null,
+): boolean {
+  if (providerType !== "openai") return false;
+  if (!isOpenAICloudBaseUrl(baseUrl)) return false;
+  const normalized = modelId?.trim().toLowerCase() ?? "";
+  if (!normalized) return false;
+  return OPENAI_IMAGE_GENERATION_MODEL_PREFIXES.some((prefix) =>
+    normalized.startsWith(prefix),
+  );
+}
+
 /**
 * Per-provider minimum on the outbound max_tokens. Kimi's docs require
 * `max_tokens >= 16000` whenever a thinking model is in use so the
--- a/studio/frontend/src/features/chat/shared-composer.tsx
+++ b/studio/frontend/src/features/chat/shared-composer.tsx
@ -21,7 +21,7 @@ import { isTauri } from "@/lib/api-base";
 import { isMultimodalResponse } from "./types/api";
 import { getImageInputUnavailableReason } from "./utils/image-input-support";
 import { useAui } from "@assistant-ui/react";
-import { ArrowUpIcon, GlobeIcon, HeadphonesIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react";
+import { ArrowUpIcon, GlobeIcon, HeadphonesIcon, ImageIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react";
 import { toast } from "@/lib/toast";
 import { loadModel, validateModel } from "./api/chat-api";
 import { parseExternalModelId, providerTypeSupportsVision } from "./external-providers";
@ -33,6 +33,7 @@ import {
 import {
  getExternalReasoningCapabilities,
  providerSupportsBuiltinCodeExecution,
+  providerSupportsBuiltinImageGeneration,
 } from "./provider-capabilities";
 import {
  type CompositionEvent,
@ -331,6 +332,10 @@ export function SharedComposer({
  const setToolsEnabled = useChatRuntimeStore((s) => s.setToolsEnabled);
  const codeToolsEnabled = useChatRuntimeStore((s) => s.codeToolsEnabled);
  const setCodeToolsEnabled = useChatRuntimeStore((s) => s.setCodeToolsEnabled);
+  const imageToolsEnabled = useChatRuntimeStore((s) => s.imageToolsEnabled);
+  const setImageToolsEnabled = useChatRuntimeStore(
+    (s) => s.setImageToolsEnabled,
+  );
  const lastOpenRouterChosenModel = useChatRuntimeStore(
    (s) => s.lastOpenRouterChosenModel,
  );
@ -416,10 +421,22 @@ export function SharedComposer({
    effectiveExternalModelId,
    selectedExternalProvider?.baseUrl,
  );
+  const supportsBuiltinImageGeneration = providerSupportsBuiltinImageGeneration(
+    selectedExternalProvider?.providerType,
+    effectiveExternalModelId,
+    selectedExternalProvider?.baseUrl,
+  );
  const searchDisabled =
    !modelLoaded || !(supportsTools || supportsBuiltinWebSearch);
  const codeDisabled =
    !modelLoaded || !(supportsTools || supportsBuiltinCodeExecution);
+  // Images pill is only ever lit on OpenAI cloud's Responses-API models.
+  // No local tool runtime fallback because the only image-generation
+  // server tool we wire today is OpenAI's; local models cannot dispatch
+  // it. Hidden entirely when the active model does not advertise it so
+  // the pill row stays compact for providers without the capability.
+  const imageDisabled = !modelLoaded || !supportsBuiltinImageGeneration;
+  const showImagePill = supportsBuiltinImageGeneration;
  // Backwards-compatible alias for any other call site that may still
  // reference `toolsDisabled` (rare; both pills used it before).
  const toolsDisabled = codeDisabled;
@ -1074,6 +1091,21 @@ export function SharedComposer({
            <CodeToggleIcon className="size-3.5" />
            <span>Code</span>
          </button>
+          {showImagePill && (
+            <button
+              type="button"
+              disabled={imageDisabled}
+              onClick={() => setImageToolsEnabled(!imageToolsEnabled)}
+              className="composer-pill-btn"
+              data-active={imageToolsEnabled && !imageDisabled ? "true" : "false"}
+              aria-label={
+                imageToolsEnabled ? "Disable image generation" : "Enable image generation"
+              }
+            >
+              <ImageIcon className="size-3.5" />
+              <span>Images</span>
+            </button>
+          )}
        </div>
        <div className="flex items-center gap-1">
          {dictationSupported && (
--- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
+++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
@ -24,6 +24,8 @@ const HF_TOKEN_KEY = "unsloth_hf_token";
 export const CHAT_REASONING_ENABLED_KEY = "unsloth_chat_reasoning_enabled";
 export const CHAT_TOOLS_ENABLED_KEY = "unsloth_chat_tools_enabled";
 export const CHAT_CODE_TOOLS_ENABLED_KEY = "unsloth_chat_code_tools_enabled";
+export const CHAT_IMAGE_TOOLS_ENABLED_KEY = "unsloth_chat_image_tools_enabled";
+
 // External provider selection is encoded into `params.checkpoint` as
 // `external::<providerId>::<modelId>`. PersistedChatSettings deliberately
 // Omits `checkpoint` because the local-model side is mirrored by the
@ -253,8 +255,16 @@ type ChatRuntimeStore = {
   * execution server-side. Read by both composers' Code pill gate.
   */
  supportsBuiltinCodeExecution: boolean;
+  /**
+   * Whether the active external provider exposes a server-side
+   * image-generation tool (OpenAI's Responses-API `image_generation`
+   * today). Gates the chat composer's Images pill. Local models never
+   * receive the tool because their runtime cannot dispatch it.
+   */
+  supportsBuiltinImageGeneration: boolean;
  toolsEnabled: boolean;
  codeToolsEnabled: boolean;
+  imageToolsEnabled: boolean;
  toolStatus: string | null;
  generatingStatus: string | null;
  autoHealToolCalls: boolean;
@ -313,6 +323,7 @@ type ChatRuntimeStore = {
  setPreserveThinking: (value: boolean) => void;
  setToolsEnabled: (enabled: boolean, options?: { persist?: boolean }) => void;
  setCodeToolsEnabled: (enabled: boolean) => void;
+  setImageToolsEnabled: (enabled: boolean) => void;
  setToolStatus: (status: string | null) => void;
  setGeneratingStatus: (status: string | null) => void;
  setAutoHealToolCalls: (enabled: boolean) => void;
@ -552,8 +563,10 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
  supportsTools: false,
  supportsBuiltinWebSearch: false,
  supportsBuiltinCodeExecution: false,
+  supportsBuiltinImageGeneration: false,
  toolsEnabled: loadBool(CHAT_TOOLS_ENABLED_KEY, false),
  codeToolsEnabled: loadBool(CHAT_CODE_TOOLS_ENABLED_KEY, false),
+  imageToolsEnabled: loadBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, false),
  toolStatus: null,
  generatingStatus: null,
  autoHealToolCalls: true,
@ -730,8 +743,10 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
      supportsTools: false,
      supportsBuiltinWebSearch: false,
      supportsBuiltinCodeExecution: false,
+      supportsBuiltinImageGeneration: false,
      toolsEnabled: false,
      codeToolsEnabled: false,
+      imageToolsEnabled: false,
      toolStatus: null,
      kvCacheDtype: null,
      loadedKvCacheDtype: null,
@ -786,6 +801,11 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
      saveBool(CHAT_CODE_TOOLS_ENABLED_KEY, codeToolsEnabled);
      return { codeToolsEnabled };
    }),
+  setImageToolsEnabled: (imageToolsEnabled) =>
+    set(() => {
+      saveBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, imageToolsEnabled);
+      return { imageToolsEnabled };
+    }),
  setToolStatus: (toolStatus) => set({ toolStatus }),
  setGeneratingStatus: (generatingStatus) => set({ generatingStatus }),
  setAutoHealToolCalls: (autoHealToolCalls) =>