Support follow-up edits for generated images (#5712)

* Support follow-up edits for generated images

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix generated image edit references

* Fix image generation tool guard

* Replay OpenAI image reasoning refs

* Capture streamed OpenAI reasoning refs

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trigger Gather Town test

* Use OpenAI response context for image edits

* Studio: improve generated image card UI

* Studio: refine generated image overlay and edit context

* Studio: animate generated image loading state

* Studio: bind generated-image edits to selected image

* Studio: drop empty external assistant payloads

* Studio: guard generated image clipboard MIME

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Simplify image generation edit wiring

* Remove temporary image generation test changes

* Preserve explicit image edit references

* Scope image edit references to threads

* Harden OpenAI image edit replay handling

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update image generation tool event test

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

studio/backend/core/inference/external_provider.py

@@ -67,6 +67,50 @@ _ANTHROPIC_4_7_SAMPLING_REMOVED = re.compile(
     r"^claude-(?:opus|sonnet|haiku)-4-7(?:[-.]|$)"
 )
 _OPENAI_REASONING_SUMMARY_UNSUPPORTED = re.compile(r"^o3(?:[-.]|$)")
+_OPENAI_REASONING_STATUSES = {"in_progress", "completed", "incomplete"}
+
+
+def _openai_image_replay_requires_reasoning(model: str) -> bool:
+    normalized = model.strip().lower()
+    return normalized.startswith("gpt-5") or normalized.startswith("o")
+
+
+def _sanitize_openai_reasoning_replay_item(
+    item: Any,
+) -> Optional[dict[str, Any]]:
+    """Return a Responses input-safe reasoning item, if ``item`` is one.
+
+    OpenAI's image-generation docs allow follow-up edits by sending the
+    previous ``image_generation_call`` id. Reasoning models can additionally
+    require the paired ``reasoning`` output item in manually managed context,
+    so keep the public replay fields only and drop everything else.
+    """
+    if not isinstance(item, dict) or item.get("type") != "reasoning":
+        return None
+    item_id = item.get("id")
+    if not isinstance(item_id, str) or not item_id:
+        return None
+    summary_parts: list[dict[str, str]] = []
+    summary = item.get("summary")
+    if isinstance(summary, list):
+        for part in summary:
+            if not isinstance(part, dict):
+                continue
+            if part.get("type") != "summary_text":
+                continue
+            text = part.get("text")
+            if isinstance(text, str):
+                summary_parts.append({"type": "summary_text", "text": text})
+    replay_item: dict[str, Any] = {
+        "type": "reasoning",
+        "id": item_id,
+        "summary": summary_parts,
+    }
+    status = item.get("status")
+    if isinstance(status, str) and status in _OPENAI_REASONING_STATUSES:
+        replay_item["status"] = status
+    return replay_item
+
 
 # OpenAI Responses inline citation markers: `citeSOURCE_ID[id2...][LOCATOR]`
 # using private-use codepoints (see
@@ -733,8 +777,7 @@ class ExternalProviderClient:
                     plugins.append({"id": "web"})
                 body["plugins"] = plugins
                 logger.info(
-                    "OpenRouter web_search: attached plugins=[{id: 'web'}] "
-                    "(model=%s)",
+                    "OpenRouter web_search: attached plugins=[{id: 'web'}] (model=%s)",
                     body.get("model"),
                 )
 
@@ -1808,7 +1851,7 @@ class ExternalProviderClient:
             and compaction_threshold > 0
             and _anthropic_supports_compaction(model)
         )
-        if compaction_active:
+        if compaction_active and compaction_threshold is not None:
             trigger_value = max(
                 int(compaction_threshold),
                 _ANTHROPIC_COMPACTION_MIN,
@@ -2864,10 +2907,17 @@ class ExternalProviderClient:
         """
         import json as _json
 
+        is_openai_cloud = _is_openai_family_cloud(self.base_url)
+        image_generation_requested = bool(
+            enabled_tools and "image_generation" in enabled_tools and is_openai_cloud
+        )
+
         # Split system messages out into a single `instructions` string and
         # translate user/assistant messages into the Responses input shape.
         instructions_parts: list[str] = []
         input_items: list[dict[str, Any]] = []
+        openai_replay_items: list[dict[str, Any]] = []
+        previous_response_id: Optional[str] = None
         for msg in messages:
             role = msg.get("role")
             content = msg.get("content", "")
@@ -2888,6 +2938,7 @@ class ExternalProviderClient:
 
             if isinstance(content, list):
                 translated_parts: list[dict[str, Any]] = []
+                used_previous_response_id = False
                 for part in content:
                     part_type = part.get("type")
                     if part_type == "text":
@@ -2902,6 +2953,36 @@ class ExternalProviderClient:
                             translated_parts.append(
                                 {"type": "input_image", "image_url": url}
                             )
+                    elif (
+                        part_type == "reasoning"
+                        and role == "assistant"
+                        and image_generation_requested
+                    ):
+                        replay_item = _sanitize_openai_reasoning_replay_item(part)
+                        if replay_item:
+                            openai_replay_items.append(replay_item)
+                    elif (
+                        part_type == "image_generation_call"
+                        and role == "assistant"
+                        and image_generation_requested
+                    ):
+                        response_id = (
+                            part.get("response_id")
+                            or part.get("openai_response_id")
+                            or part.get("previous_response_id")
+                        )
+                        call_id = part.get("id") or part.get("image_generation_call_id")
+                        if isinstance(call_id, str) and call_id:
+                            if isinstance(response_id, str) and response_id:
+                                previous_response_id = response_id
+                                input_items = []
+                                translated_parts = []
+                                used_previous_response_id = True
+                            else:
+                                previous_response_id = None
+                            openai_replay_items.append(
+                                {"type": "image_generation_call", "id": call_id}
+                            )
                     elif part_type == "input_document":
                         # OpenAI Responses accepts PDFs / docs as
                         # `{type:"input_file", file_data:"data:application/pdf;base64,..."}`
@@ -2939,9 +3020,59 @@ class ExternalProviderClient:
                         if filename:
                             block["filename"] = filename
                         translated_parts.append(block)
-                if translated_parts:
+                if translated_parts and not used_previous_response_id:
                     input_items.append({"role": role, "content": translated_parts})
 
+        if previous_response_id:
+            # OpenAI's documented multi-turn image generation path can use
+            # `previous_response_id` to carry the prior generated image and
+            # paired reasoning state. Prefer that over manual item replay when
+            # we captured the response id; keep replay below as a fallback for
+            # older stored turns that only have an image_generation_call id.
+            openai_replay_items = []
+        elif (
+            _openai_image_replay_requires_reasoning(model)
+            and reasoning_effort != "none"
+            and enable_thinking is not False
+        ):
+            filtered_replay_items: list[dict[str, Any]] = []
+            has_reasoning_replay = False
+            dropped_image_replay_without_reasoning = False
+            for item in openai_replay_items:
+                if item.get("type") == "reasoning":
+                    has_reasoning_replay = True
+                    filtered_replay_items.append(item)
+                elif item.get("type") == "image_generation_call":
+                    if has_reasoning_replay:
+                        filtered_replay_items.append(item)
+                    else:
+                        dropped_image_replay_without_reasoning = True
+                else:
+                    filtered_replay_items.append(item)
+            openai_replay_items = filtered_replay_items
+            if dropped_image_replay_without_reasoning:
+                yield _error_sse_line(
+                    400,
+                    "OpenAI image edit reference is missing paired reasoning state. "
+                    "Regenerate the image, then retry the edit.",
+                    self.provider_type,
+                )
+                return
+        image_generation_has_reference = bool(
+            previous_response_id
+            or any(
+                isinstance(item, dict) and item.get("type") == "image_generation_call"
+                for item in openai_replay_items
+            )
+        )
+        if openai_replay_items:
+            insert_at = len(input_items)
+            for index in range(len(input_items) - 1, -1, -1):
+                if input_items[index].get("role") == "user":
+                    insert_at = index
+                    break
+            input_items[insert_at:insert_at] = openai_replay_items
+
         # NOTE: gpt-5.x / o3 / gpt-4.5 are reasoning-class models. They reject
         # temperature and top_p with `Unsupported parameter` 400s on
         # /v1/responses (and on /v1/chat/completions for the same families).
@@ -2957,6 +3088,8 @@ class ExternalProviderClient:
             "input": input_items,
             "stream": True,
         }
+        if previous_response_id:
+            body["previous_response_id"] = previous_response_id
         # `summary: "auto"` is what makes /v1/responses emit reasoning
         # summary events — without it OpenAI returns no thinking text on
         # most reasoning models, the SSE handler has no <think>…</think>
@@ -3013,7 +3146,6 @@ class ExternalProviderClient:
         # (ollama / llama.cpp / vLLM / "custom" preset) hit /v1/responses
         # without these extensions and would 400 on the unknown body
         # fields, so they intentionally fall outside this gate.
-        is_openai_cloud = _is_openai_family_cloud(self.base_url)
         if is_openai_cloud and enable_prompt_caching is not False:
             body["prompt_cache_retention"] = "24h"
 
@@ -3059,9 +3191,18 @@ class ExternalProviderClient:
         # plus gpt-4.1 / gpt-4o / o3 per the docs; restrict to cloud
         # OpenAI because the local llama.cpp / ollama backends don't
         # implement it and would 400.
-        image_generation_enabled_openai = bool(
-            enabled_tools and "image_generation" in enabled_tools and is_openai_cloud
-        )
+        image_generation_enabled_openai = image_generation_requested
+
+        def _openai_image_generation_tool() -> dict[str, Any]:
+            tool: dict[str, Any] = {"type": "image_generation"}
+            if image_generation_has_reference:
+                # OpenAI's Responses image tool defaults to `auto`. For
+                # Studio's explicit follow-up edit flow, force edit mode so
+                # the provider uses the previous response / call id as image
+                # context instead of treating the text as a fresh generation.
+                tool["action"] = "edit"
+            return tool
+
         if enabled_tools:
             tools_array: list[dict[str, Any]] = []
             if "web_search" in enabled_tools:
@@ -3089,7 +3230,7 @@ class ExternalProviderClient:
                     shell_env = {"type": "container_auto"}
                 tools_array.append({"type": "shell", "environment": shell_env})
             if image_generation_enabled_openai:
-                tools_array.append({"type": "image_generation"})
+                tools_array.append(_openai_image_generation_tool())
             if tools_array:
                 body["tools"] = tools_array
 
@@ -3121,7 +3262,7 @@ class ExternalProviderClient:
                         {"type": "shell", "environment": env_attempt}
                     )
                 if image_generation_enabled_openai:
-                    tools_array_attempt.append({"type": "image_generation"})
+                    tools_array_attempt.append(_openai_image_generation_tool())
                 if tools_array_attempt:
                     attempt_body["tools"] = tools_array_attempt
                 else:
@@ -3212,7 +3353,7 @@ class ExternalProviderClient:
                     # to a specific search invocation. Hence the shared list.
                     # web_search_calls: { item_id -> {query} }
                     web_search_calls: dict[str, dict[str, Any]] = {}
-                    all_url_citations: list[dict[str, str]] = []
+                    all_url_citations: list[dict[str, Any]] = []
                     # Shell-tool (code execution) state. OpenAI emits
                     # `shell_call` items (model requesting a command list)
                     # paired with `shell_call_output` items (execution
@@ -3235,6 +3376,10 @@ class ExternalProviderClient:
                     # see.
                     latched_container_id: Optional[str] = None
                     container_id_emitted = False
+                    current_openai_response_id: Optional[str] = None
+                    last_openai_reasoning_replay_item: Optional[dict[str, Any]] = None
+                    openai_reasoning_replay_items: dict[str, dict[str, Any]] = {}
+                    image_generation_calls_started: set[str] = set()
                     # Buffer for a citation marker straddling two delta events;
                     # prepended onto the next delta. See _split_pending_citation_tail.
                     pending_marker_tail: str = ""
@@ -3245,6 +3390,18 @@ class ExternalProviderClient:
                     # with leftover private-use codepoints stripped.
                     pending_citation_segments: list[str] = []
 
+                    def _record_openai_response_id(payload: dict[str, Any]) -> None:
+                        nonlocal current_openai_response_id
+                        response_obj = payload.get("response")
+                        candidates: list[Any] = []
+                        if isinstance(response_obj, dict):
+                            candidates.append(response_obj.get("id"))
+                        candidates.append(payload.get("response_id"))
+                        for candidate in candidates:
+                            if isinstance(candidate, str) and candidate:
+                                current_openai_response_id = candidate
+                                return
+
                     def _drain_pending_segments(force: bool) -> str:
                         """Re-attempt resolution on buffered segments in order.
                         Stops at the first still-unresolved segment unless
@@ -3392,6 +3549,80 @@ class ExternalProviderClient:
                             }
                         )
 
+                    def _record_openai_reasoning_replay_item(
+                        payload: Any,
+                    ) -> Optional[dict[str, Any]]:
+                        if not isinstance(payload, dict):
+                            return None
+                        item_id = payload.get("id") or payload.get("item_id")
+                        if not isinstance(item_id, str) or not item_id:
+                            return None
+                        existing = openai_reasoning_replay_items.setdefault(
+                            item_id,
+                            {
+                                "type": "reasoning",
+                                "id": item_id,
+                                "summary": [],
+                                "status": "completed",
+                            },
+                        )
+                        if payload.get("type") == "reasoning":
+                            sanitized = _sanitize_openai_reasoning_replay_item(payload)
+                            if sanitized:
+                                existing.update(sanitized)
+                                return existing
+                        summary_text = ""
+                        part = payload.get("part")
+                        if (
+                            isinstance(part, dict)
+                            and part.get("type") == "summary_text"
+                        ):
+                            text = part.get("text")
+                            if isinstance(text, str):
+                                summary_text = text
+                        elif (
+                            payload.get("type")
+                            == "response.reasoning_summary_text.done"
+                        ):
+                            text = payload.get("text")
+                            if isinstance(text, str):
+                                summary_text = text
+                        if summary_text:
+                            summary_index = payload.get("summary_index")
+                            summary = existing.setdefault("summary", [])
+                            if isinstance(summary, list):
+                                summary_part = {
+                                    "type": "summary_text",
+                                    "text": summary_text,
+                                }
+                                if (
+                                    isinstance(summary_index, int)
+                                    and summary_index >= 0
+                                ):
+                                    while len(summary) <= summary_index:
+                                        summary.append(
+                                            {"type": "summary_text", "text": ""}
+                                        )
+                                    summary[summary_index] = summary_part
+                                else:
+                                    summary.append(summary_part)
+                        return existing
+
+                    def _image_generation_arguments(
+                        prompt: str,
+                        raw_item_id: Any,
+                    ) -> dict[str, Any]:
+                        arguments: dict[str, Any] = {"kind": "image", "prompt": prompt}
+                        if isinstance(raw_item_id, str) and raw_item_id:
+                            arguments["openai_image_generation_call_id"] = raw_item_id
+                        if current_openai_response_id:
+                            arguments["openai_response_id"] = current_openai_response_id
+                        if last_openai_reasoning_replay_item:
+                            arguments["openai_reasoning_item"] = (
+                                last_openai_reasoning_replay_item
+                            )
+                        return arguments
+
                     def _extract_reasoning_text(payload: Any) -> str:
                         if payload is None:
                             return ""
@@ -3478,6 +3709,7 @@ class ExternalProviderClient:
                                 continue
 
                             event_type = event.get("type")
+                            _record_openai_response_id(event)
 
                             if event_type == "response.output_text.delta":
                                 delta_text = event.get("delta", "")
@@ -3534,11 +3766,6 @@ class ExternalProviderClient:
                                     yield _chunk_with_text(flushed)
 
                             elif event_type == "response.output_item.added":
-                                # Track the call early but do NOT emit tool_start
-                                # yet — action.query is not reliably populated on
-                                # added across OpenAI API versions, and the
-                                # frontend's tool_start is a one-shot push (no
-                                # update mechanism). Wait for output_item.done.
                                 item = event.get("item", {})
                                 if (
                                     isinstance(item, dict)
@@ -3579,12 +3806,34 @@ class ExternalProviderClient:
                                             and latched_container_id is None
                                         ):
                                             latched_container_id = probe
+                                if (
+                                    isinstance(item, dict)
+                                    and item.get("type") == "image_generation_call"
+                                ):
+                                    raw_item_id = item.get("id")
+                                    if isinstance(raw_item_id, str) and raw_item_id:
+                                        arguments = _image_generation_arguments(
+                                            "",
+                                            raw_item_id,
+                                        )
+                                        image_generation_calls_started.add(raw_item_id)
+                                        yield _emit_tool_event(
+                                            {
+                                                "type": "tool_start",
+                                                "tool_name": "image_generation",
+                                                "tool_call_id": raw_item_id,
+                                                "arguments": arguments,
+                                            }
+                                        )
 
                             elif event_type == "response.output_item.done":
                                 item = event.get("item", {})
                                 if not isinstance(item, dict):
                                     continue
                                 if item.get("type") == "reasoning":
+                                    last_openai_reasoning_replay_item = (
+                                        _record_openai_reasoning_replay_item(item)
+                                    )
                                     summary_text = _extract_reasoning_text(
                                         item.get("summary")
                                     )
@@ -3716,25 +3965,26 @@ class ExternalProviderClient:
                                     # millisecond resolution so synthesised
                                     # ids stay unique even when two image
                                     # generations resolve in the same ms.
-                                    item_id = item.get("id", "") or (
-                                        f"img_{time.time_ns()}"
-                                    )
+                                    raw_item_id = item.get("id")
+                                    item_id = raw_item_id or f"img_{time.time_ns()}"
                                     prompt_in = (
                                         item.get("revised_prompt")
                                         or item.get("prompt")
                                         or ""
                                     )
-                                    yield _emit_tool_event(
-                                        {
-                                            "type": "tool_start",
-                                            "tool_name": "image_generation",
-                                            "tool_call_id": item_id,
-                                            "arguments": {
-                                                "kind": "image",
-                                                "prompt": prompt_in,
-                                            },
-                                        }
+                                    done_arguments = _image_generation_arguments(
+                                        prompt_in,
+                                        raw_item_id,
                                     )
+                                    if item_id not in image_generation_calls_started:
+                                        yield _emit_tool_event(
+                                            {
+                                                "type": "tool_start",
+                                                "tool_name": "image_generation",
+                                                "tool_call_id": item_id,
+                                                "arguments": done_arguments,
+                                            }
+                                        )
                                     b64 = (
                                         item.get("result") or item.get("b64_json") or ""
                                     )
@@ -3744,11 +3994,13 @@ class ExternalProviderClient:
                                             "type": "tool_end",
                                             "tool_call_id": item_id,
                                             "result": "",
+                                            "arguments": done_arguments,
                                             "image_b64": b64,
                                             "image_mime": (f"image/{output_format}"),
                                             "size": item.get("size"),
                                             "quality": item.get("quality"),
                                             "background": item.get("background"),
+                                            "prompt": prompt_in,
                                         }
                                     )
 
@@ -3756,6 +4008,13 @@ class ExternalProviderClient:
                                 isinstance(event_type, str)
                                 and "reasoning" in event_type
                             ):
+                                recorded_reasoning = (
+                                    _record_openai_reasoning_replay_item(event)
+                                )
+                                if recorded_reasoning:
+                                    last_openai_reasoning_replay_item = (
+                                        recorded_reasoning
+                                    )
                                 reasoning_delta = _extract_reasoning_text(event)
                                 if reasoning_delta:
                                     if not reasoning_open:
@@ -3848,8 +4107,7 @@ class ExternalProviderClient:
                                     blocks: list[str] = []
                                     for cit in all_url_citations:
                                         line = (
-                                            f"Title: {cit['title']}\n"
-                                            f"URL: {cit['url']}"
+                                            f"Title: {cit['title']}\nURL: {cit['url']}"
                                         )
                                         if cit.get("snippet"):
                                             line += f"\nSnippet: {cit['snippet']}"
@@ -3927,8 +4185,7 @@ class ExternalProviderClient:
                                     blocks = []
                                     for cit in all_url_citations:
                                         line = (
-                                            f"Title: {cit['title']}\n"
-                                            f"URL: {cit['url']}"
+                                            f"Title: {cit['title']}\nURL: {cit['url']}"
                                         )
                                         if cit.get("snippet"):
                                             line += f"\nSnippet: {cit['snippet']}"

studio/backend/models/inference.py

@@ -471,6 +471,40 @@ class InputDocumentContentPart(BaseModel):
     )
 
 
+class OpenAIReasoningContentPart(BaseModel):
+    """OpenAI Responses reasoning item paired with a tool output.
+
+    Reasoning models can require the previous ``reasoning`` output item
+    to be replayed immediately before an ``image_generation_call`` id
+    when manually managing Responses context. This part is OpenAI-only;
+    routes strip it for every other provider before proxying.
+    """
+
+    type: Literal["reasoning"]
+    id: str = Field(..., description = "OpenAI reasoning output item id.")
+    summary: list[dict[str, Any]] = Field(default_factory = list)
+    status: Optional[Literal["in_progress", "completed", "incomplete"]] = None
+
+
+class ImageGenerationCallContentPart(BaseModel):
+    """OpenAI Responses image_generation call reference.
+
+    OpenAI accepts prior ``image_generation_call`` items in the next
+    Responses ``input`` array so follow-up prompts can edit or refine a
+    generated image without resending the base64 payload. The frontend
+    forwards this as a synthetic assistant content part when building
+    the next OpenAI Responses request; ``external_provider`` translates
+    it back to the provider-specific top-level input item.
+    """
+
+    type: Literal["image_generation_call"]
+    id: str = Field(..., description = "OpenAI image_generation_call output item id.")
+    response_id: Optional[str] = Field(
+        None,
+        description = "OpenAI Responses response id to use as previous_response_id for follow-up edits.",
+    )
+
+
 class CompactionContentPart(BaseModel):
     """Anthropic server-side compaction state, attached to an assistant
     message for round-tripping on the next turn.
@@ -504,6 +538,8 @@ ContentPart = Annotated[
         Annotated[TextContentPart, Tag("text")],
         Annotated[ImageContentPart, Tag("image_url")],
         Annotated[InputDocumentContentPart, Tag("input_document")],
+        Annotated[OpenAIReasoningContentPart, Tag("reasoning")],
+        Annotated[ImageGenerationCallContentPart, Tag("image_generation_call")],
         Annotated[CompactionContentPart, Tag("compaction")],
     ],
     Discriminator(_content_part_discriminator),

studio/backend/routes/inference.py

@@ -1709,6 +1709,12 @@ def _build_external_messages(
       see ``_INPUT_DOCUMENT_PROVIDERS``). For every other provider the
       part is stripped so the unknown content type doesn't reach generic
       /chat/completions passthrough and 400 the request.
+    - `reasoning`: OpenAI-only Responses reasoning item paired with a
+      prior tool output. Forwarded ONLY when provider_type=="openai"
+      so follow-up image edits can replay the required reasoning item.
+    - `image_generation_call`: OpenAI-only Responses image reference.
+      Forwarded ONLY when provider_type=="openai" so follow-up image
+      edits can reference prior generated images.
     - `compaction`: Anthropic-only synthetic part (round-trips server-side
       compaction state). Forwarded ONLY when provider_type=="anthropic";
       stripped for every other provider so the unknown part doesn't
@@ -1717,6 +1723,7 @@ def _build_external_messages(
     """
     document_provider = provider_type in _INPUT_DOCUMENT_PROVIDERS
     anthropic = provider_type == "anthropic"
+    openai = provider_type == "openai"
     result = []
     for msg in messages:
         if isinstance(msg.content, str):
@@ -1737,6 +1744,30 @@ def _build_external_messages(
                                 "image_url": {"url": part.image_url.url},
                             }
                         )
+                    elif (
+                        part.type == "reasoning" and openai and msg.role == "assistant"
+                    ):
+                        reasoning: dict[str, Any] = {
+                            "type": "reasoning",
+                            "id": part.id,
+                            "summary": part.summary,
+                        }
+                        if part.status:
+                            reasoning["status"] = part.status
+                        parts.append(reasoning)
+                    elif (
+                        part.type == "image_generation_call"
+                        and openai
+                        and msg.role == "assistant"
+                    ):
+                        # ExternalProviderClient maps this onto a top-level
+                        # Responses input item after the current user prompt,
+                        # or onto `previous_response_id` when response_id is
+                        # available from the prior Responses turn.
+                        image_ref = {"type": "image_generation_call", "id": part.id}
+                        if getattr(part, "response_id", None):
+                            image_ref["response_id"] = part.response_id
+                        parts.append(image_ref)
                     elif part.type == "input_document" and document_provider:
                         # ExternalProviderClient maps this onto
                         # Anthropic's `document` or OpenAI Responses'
@@ -1758,6 +1789,8 @@ def _build_external_messages(
                         # provider would 400 on the unknown part, so
                         # gate by provider_type.
                         parts.append({"type": "compaction", "content": part.content})
+                if msg.role == "assistant" and not parts:
+                    continue
                 result.append({"role": msg.role, "content": parts})
             else:
                 # Non-vision provider: strip images / documents, keep
@@ -1769,8 +1802,28 @@ def _build_external_messages(
                 for p in msg.content:
                     if p.type == "text":
                         preserved.append({"type": "text", "text": p.text})
+                    elif p.type == "reasoning" and openai and msg.role == "assistant":
+                        reasoning: dict[str, Any] = {
+                            "type": "reasoning",
+                            "id": p.id,
+                            "summary": p.summary,
+                        }
+                        if p.status:
+                            reasoning["status"] = p.status
+                        preserved.append(reasoning)
+                    elif (
+                        p.type == "image_generation_call"
+                        and openai
+                        and msg.role == "assistant"
+                    ):
+                        image_ref = {"type": "image_generation_call", "id": p.id}
+                        if getattr(p, "response_id", None):
+                            image_ref["response_id"] = p.response_id
+                        preserved.append(image_ref)
                     elif p.type == "compaction" and anthropic:
                         preserved.append({"type": "compaction", "content": p.content})
+                if msg.role == "assistant" and not preserved:
+                    continue
                 if len(preserved) == 1 and preserved[0]["type"] == "text":
                     # Single text part collapses back to a string for
                     # providers that don't accept content arrays.

studio/backend/tests/test_openai_image_generation.py

@@ -210,6 +210,7 @@ def test_image_generation_done_emits_tool_event_chunks(monkeypatch):
     assert starts[0]["arguments"] == {
         "kind": "image",
         "prompt": "A photorealistic cat sitting",
+        "openai_image_generation_call_id": "img_abc",
     }
     assert ends[0]["image_b64"] == "AAAA"
     assert ends[0]["image_mime"] == "image/png"

studio/frontend/src/components/assistant-ui/generated-image-overlay-context.tsx

@@ -0,0 +1,77 @@
+// 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 ReactNode,
+  createContext,
+  useCallback,
+  useContext,
+  useMemo,
+  useState,
+} from "react";
+
+export type GeneratedImageOverlayState = {
+  image: string;
+  title: string;
+  metadata: string;
+  filename?: string;
+  openaiImageGenerationCallId?: string;
+  openaiResponseId?: string;
+  openaiReasoningItem?: unknown;
+  threadId?: string | null;
+};
+
+type GeneratedImageOverlayContextValue = {
+  overlay: GeneratedImageOverlayState | null;
+  openOverlay: (overlay: GeneratedImageOverlayState) => void;
+  closeOverlay: () => void;
+};
+
+const GeneratedImageOverlayContext =
+  createContext<GeneratedImageOverlayContextValue | null>(null);
+
+export function GeneratedImageOverlayProvider({
+  children,
+  threadId = null,
+}: {
+  children: ReactNode;
+  threadId?: string | null;
+}) {
+  const [overlay, setOverlay] = useState<GeneratedImageOverlayState | null>(
+    null,
+  );
+
+  const openOverlay = useCallback(
+    (nextOverlay: GeneratedImageOverlayState) => {
+      setOverlay({ ...nextOverlay, threadId: nextOverlay.threadId ?? threadId });
+    },
+    [threadId],
+  );
+
+  const closeOverlay = useCallback(() => {
+    setOverlay(null);
+  }, []);
+
+  const value = useMemo(
+    () => ({ overlay, openOverlay, closeOverlay }),
+    [closeOverlay, openOverlay, overlay],
+  );
+
+  return (
+    <GeneratedImageOverlayContext.Provider value={value}>
+      {children}
+    </GeneratedImageOverlayContext.Provider>
+  );
+}
+
+export function useGeneratedImageOverlay(): GeneratedImageOverlayContextValue {
+  const context = useContext(GeneratedImageOverlayContext);
+  if (!context) {
+    throw new Error(
+      "useGeneratedImageOverlay must be used within GeneratedImageOverlayProvider.",
+    );
+  }
+  return context;
+}

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

@@ -0,0 +1,510 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+//
+// Portions adapted from assistant-ui packages/ui/src/components/assistant-ui/image.tsx
+// MIT License, Copyright (c) 2025 AgentbaseAI Inc.
+// Source: https://github.com/assistant-ui/assistant-ui/blob/main/packages/ui/src/components/assistant-ui/image.tsx
+
+"use client";
+
+import { cn } from "@/lib/utils";
+import type {
+  ImageMessagePart,
+  ImageMessagePartComponent,
+} from "@assistant-ui/react";
+import { type VariantProps, cva } from "class-variance-authority";
+import {
+  CopyIcon,
+  DownloadIcon,
+  ImageIcon,
+  ImageOffIcon,
+  Loader2Icon,
+  RefreshCwIcon,
+  ShieldAlertIcon,
+} from "lucide-react";
+import {
+  type ComponentProps,
+  type PropsWithChildren,
+  memo,
+  useEffect,
+  useRef,
+  useState,
+} from "react";
+import { createPortal } from "react-dom";
+
+const extensionForMimeType = (mimeType?: string): string => {
+  switch (mimeType) {
+    case "image/png":
+      return "png";
+    case "image/jpeg":
+    case "image/jpg":
+      return "jpg";
+    case "image/webp":
+      return "webp";
+    case "image/gif":
+      return "gif";
+    case "image/svg+xml":
+      return "svg";
+    default:
+      return "png";
+  }
+};
+
+const DATA_URI_MIME_RE = /data:([^;]+)/;
+const DATA_URI_BASE64_RE = /;base64/i;
+const IMAGE_DATA_URI_MIME_RE = /^data:([^;,]+)/;
+
+export const dataUriToBlob = (dataUri: string): Blob => {
+  const [meta, data] = dataUri.split(",");
+  const mime = meta?.match(DATA_URI_MIME_RE)?.[1] ?? "application/octet-stream";
+  if (!DATA_URI_BASE64_RE.test(meta ?? "")) {
+    return new Blob([decodeURIComponent(data ?? "")], { type: mime });
+  }
+  const bytes = atob(data ?? "");
+  const arr = new Uint8Array(bytes.length);
+  for (let i = 0; i < bytes.length; i += 1) {
+    arr[i] = bytes.charCodeAt(i);
+  }
+  return new Blob([arr], { type: mime });
+};
+
+const mimeFromImage = (image: string): string | undefined =>
+  image.match(IMAGE_DATA_URI_MIME_RE)?.[1];
+
+export const downloadImagePart = (
+  part: Pick<ImageMessagePart, "image" | "filename">,
+): void => {
+  if (typeof document === "undefined") {
+    return;
+  }
+  const ext = extensionForMimeType(mimeFromImage(part.image));
+  const filename = part.filename ?? `image.${ext}`;
+  const isDataUri = part.image.startsWith("data:");
+  const objectUrl = isDataUri
+    ? URL.createObjectURL(dataUriToBlob(part.image))
+    : null;
+  const href = objectUrl ?? part.image;
+  const a = document.createElement("a");
+  a.href = href;
+  a.download = filename;
+  a.rel = "noopener";
+  document.body.appendChild(a);
+  a.click();
+  document.body.removeChild(a);
+  if (objectUrl) {
+    URL.revokeObjectURL(objectUrl);
+  }
+};
+
+export const copyImagePart = async (
+  part: Pick<ImageMessagePart, "image">,
+): Promise<void> => {
+  if (
+    typeof navigator === "undefined" ||
+    !navigator.clipboard ||
+    typeof ClipboardItem === "undefined"
+  ) {
+    throw new Error("Clipboard API is not available in this environment.");
+  }
+  const blob = part.image.startsWith("data:")
+    ? dataUriToBlob(part.image)
+    : await fetch(part.image).then((r) => r.blob());
+  const mime = mimeFromImage(part.image) || blob.type || "image/png";
+  await navigator.clipboard.write([new ClipboardItem({ [mime]: blob })]);
+};
+
+const reportImageCopyError = (error: unknown): void => {
+  if (typeof window === "undefined") {
+    return;
+  }
+  window.dispatchEvent(
+    new CustomEvent("assistant-ui:image-copy-error", { detail: error }),
+  );
+};
+
+const imageVariants = cva(
+  "aui-image-root relative overflow-hidden rounded-lg",
+  {
+    variants: {
+      variant: {
+        outline: "border border-border",
+        ghost: "",
+        muted: "bg-muted/50",
+      },
+      size: {
+        sm: "max-w-64",
+        default: "max-w-96",
+        lg: "max-w-[512px]",
+        full: "w-full",
+      },
+    },
+    defaultVariants: {
+      variant: "outline",
+      size: "default",
+    },
+  },
+);
+
+export type ImageRootProps = ComponentProps<"div"> &
+  VariantProps<typeof imageVariants>;
+
+function ImageRoot({
+  className,
+  variant,
+  size,
+  children,
+  ...props
+}: ImageRootProps) {
+  return (
+    <div
+      data-slot="image-root"
+      data-variant={variant}
+      data-size={size}
+      className={cn(imageVariants({ variant, size, className }))}
+      {...props}
+    >
+      {children}
+    </div>
+  );
+}
+
+type ImagePreviewProps = Omit<ComponentProps<"img">, "children"> & {
+  containerClassName?: string;
+};
+
+function ImagePreview({
+  className,
+  containerClassName,
+  onLoad,
+  onError,
+  alt = "Image content",
+  src,
+  ...props
+}: ImagePreviewProps) {
+  const imgRef = useRef<HTMLImageElement>(null);
+  const [loadedSrc, setLoadedSrc] = useState<string | undefined>(undefined);
+  const [errorSrc, setErrorSrc] = useState<string | undefined>(undefined);
+
+  const loaded = loadedSrc === src;
+  const error = errorSrc === src;
+
+  useEffect(() => {
+    if (
+      typeof src === "string" &&
+      imgRef.current?.complete &&
+      imgRef.current.naturalWidth > 0
+    ) {
+      setLoadedSrc(src);
+    }
+  }, [src]);
+
+  return (
+    <div
+      data-slot="image-preview"
+      className={cn("relative min-h-32", containerClassName)}
+    >
+      {!(loaded || error) && (
+        <div
+          data-slot="image-preview-loading"
+          className="absolute inset-0 flex items-center justify-center bg-muted/50"
+        >
+          <ImageIcon className="size-8 animate-pulse text-muted-foreground" />
+        </div>
+      )}
+      {error ? (
+        <div
+          data-slot="image-preview-error"
+          className="flex min-h-32 items-center justify-center bg-muted/50 p-4"
+        >
+          <ImageOffIcon className="size-8 text-muted-foreground" />
+        </div>
+      ) : (
+        <img
+          {...props}
+          ref={imgRef}
+          src={src}
+          alt={alt}
+          className={cn(
+            "block h-auto w-full object-contain",
+            !loaded && "invisible",
+            className,
+          )}
+          onLoad={(e) => {
+            if (typeof src === "string") {
+              setLoadedSrc(src);
+            }
+            onLoad?.(e);
+          }}
+          onError={(e) => {
+            if (typeof src === "string") {
+              setErrorSrc(src);
+            }
+            onError?.(e);
+          }}
+        />
+      )}
+    </div>
+  );
+}
+
+function ImageFilename({
+  className,
+  children,
+  ...props
+}: ComponentProps<"span">) {
+  if (!children) {
+    return null;
+  }
+
+  return (
+    <span
+      data-slot="image-filename"
+      className={cn(
+        "block truncate px-2 py-1.5 text-muted-foreground text-xs",
+        className,
+      )}
+      {...props}
+    >
+      {children}
+    </span>
+  );
+}
+
+type ImageZoomProps = PropsWithChildren<{
+  src: string;
+  alt?: string;
+}>;
+
+function ImageZoom({ src, alt = "Image preview", children }: ImageZoomProps) {
+  const [isOpen, setIsOpen] = useState(false);
+
+  const handleOpen = () => setIsOpen(true);
+  const handleClose = () => setIsOpen(false);
+
+  useEffect(() => {
+    if (!isOpen) {
+      return;
+    }
+    const handleKeyDown = (e: KeyboardEvent) => {
+      if (e.key === "Escape") {
+        setIsOpen(false);
+      }
+    };
+    document.addEventListener("keydown", handleKeyDown);
+    return () => document.removeEventListener("keydown", handleKeyDown);
+  }, [isOpen]);
+
+  useEffect(() => {
+    if (!isOpen) {
+      return;
+    }
+    const originalOverflow = document.body.style.overflow;
+    document.body.style.overflow = "hidden";
+    return () => {
+      document.body.style.overflow = originalOverflow;
+    };
+  }, [isOpen]);
+
+  return (
+    <>
+      <button
+        type="button"
+        onClick={handleOpen}
+        className="aui-image-zoom-trigger w-full cursor-zoom-in border-0 bg-transparent p-0 text-left"
+        aria-label="Click to zoom image"
+      >
+        {children}
+      </button>
+      {isOpen &&
+        typeof document !== "undefined" &&
+        createPortal(
+          <button
+            type="button"
+            data-slot="image-zoom-overlay"
+            className="aui-image-zoom-overlay fade-in fixed inset-0 z-50 flex animate-in items-center justify-center border-0 bg-black/80 p-0 duration-200"
+            onClick={handleClose}
+            aria-label="Close zoomed image"
+          >
+            <img
+              data-slot="image-zoom-content"
+              src={src}
+              alt={alt}
+              className="aui-image-zoom-content fade-in zoom-in-95 max-h-[90vh] max-w-[90vw] animate-in cursor-zoom-out object-contain duration-200"
+            />
+          </button>,
+          document.body,
+        )}
+    </>
+  );
+}
+
+function ImageGenerating({ className }: { className?: string }) {
+  return (
+    <div
+      data-slot="image-generating"
+      className={cn(
+        "flex min-h-32 items-center justify-center bg-muted/50 p-4",
+        className,
+      )}
+    >
+      <Loader2Icon className="size-8 animate-spin text-muted-foreground" />
+      <span className="sr-only">Generating image…</span>
+    </div>
+  );
+}
+
+function ImageContentFilterError({
+  className,
+  reason,
+}: {
+  className?: string;
+  reason?: string;
+}) {
+  return (
+    <div
+      data-slot="image-content-filter-error"
+      className={cn(
+        "flex min-h-32 flex-col items-center justify-center gap-2 bg-muted/50 p-4 text-center",
+        className,
+      )}
+    >
+      <ShieldAlertIcon className="size-8 text-muted-foreground" />
+      <p className="font-medium text-sm">Image could not be generated</p>
+      {reason && <p className="text-muted-foreground text-xs">{reason}</p>}
+    </div>
+  );
+}
+
+export type ImageActionsProps = {
+  part: ImageMessagePart;
+  /**
+   * Wire to your own generation call to show a regenerate button. The button
+   * renders only when this is set and the part carries a `prompt`.
+   */
+  onRegenerate?: () => void | Promise<void>;
+  className?: string;
+};
+
+function RegenerateButton({
+  onRegenerate,
+}: {
+  onRegenerate: () => void | Promise<void>;
+}) {
+  const [isRegenerating, setIsRegenerating] = useState(false);
+  return (
+    <button
+      type="button"
+      onClick={async () => {
+        setIsRegenerating(true);
+        try {
+          await onRegenerate();
+        } finally {
+          setIsRegenerating(false);
+        }
+      }}
+      disabled={isRegenerating}
+      data-slot="image-regenerate"
+      aria-label="Regenerate image"
+      className="inline-flex size-7 items-center justify-center rounded hover:bg-muted disabled:opacity-50"
+    >
+      <RefreshCwIcon
+        className={cn("size-4", isRegenerating && "animate-spin")}
+      />
+    </button>
+  );
+}
+
+function ImageActions({ part, onRegenerate, className }: ImageActionsProps) {
+  return (
+    <div
+      data-slot="image-actions"
+      className={cn("flex items-center gap-1 p-1", className)}
+    >
+      <button
+        type="button"
+        onClick={() => downloadImagePart(part)}
+        data-slot="image-download"
+        aria-label="Download image"
+        className="inline-flex size-7 items-center justify-center rounded hover:bg-muted"
+      >
+        <DownloadIcon className="size-4" />
+      </button>
+      <button
+        type="button"
+        onClick={() => {
+          copyImagePart(part).catch((error) => {
+            reportImageCopyError(error);
+          });
+        }}
+        data-slot="image-copy"
+        aria-label="Copy image"
+        className="inline-flex size-7 items-center justify-center rounded hover:bg-muted"
+      >
+        <CopyIcon className="size-4" />
+      </button>
+      {onRegenerate && <RegenerateButton onRegenerate={onRegenerate} />}
+    </div>
+  );
+}
+
+const ImageImpl: ImageMessagePartComponent = (props) => {
+  const { image, filename, status } = props;
+  const alt = filename || "Image content";
+
+  if (status?.type === "running") {
+    return (
+      <ImageRoot>
+        <ImageGenerating />
+        <ImageFilename>{filename}</ImageFilename>
+      </ImageRoot>
+    );
+  }
+
+  if (status?.type === "incomplete" && status.reason === "content-filter") {
+    return (
+      <ImageRoot>
+        <ImageContentFilterError reason="The provider blocked this image." />
+      </ImageRoot>
+    );
+  }
+
+  return (
+    <ImageRoot>
+      <ImageZoom src={image} alt={alt}>
+        <ImagePreview src={image} alt={alt} />
+      </ImageZoom>
+      <ImageFilename>{filename}</ImageFilename>
+    </ImageRoot>
+  );
+};
+
+const Image = memo(ImageImpl) as unknown as ImageMessagePartComponent & {
+  Root: typeof ImageRoot;
+  Preview: typeof ImagePreview;
+  Filename: typeof ImageFilename;
+  Zoom: typeof ImageZoom;
+  Actions: typeof ImageActions;
+  Generating: typeof ImageGenerating;
+  ContentFilterError: typeof ImageContentFilterError;
+};
+
+Image.displayName = "Image";
+Image.Root = ImageRoot;
+Image.Preview = ImagePreview;
+Image.Filename = ImageFilename;
+Image.Zoom = ImageZoom;
+Image.Actions = ImageActions;
+Image.Generating = ImageGenerating;
+Image.ContentFilterError = ImageContentFilterError;
+
+export {
+  Image,
+  ImageRoot,
+  ImagePreview,
+  ImageFilename,
+  ImageZoom,
+  ImageActions,
+  ImageGenerating,
+  ImageContentFilterError,
+  imageVariants,
+};

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

@@ -7,6 +7,11 @@ import {
   UserMessageAttachments,
 } from "@/components/assistant-ui/attachment";
 import { CodeToggleIcon } from "@/components/assistant-ui/code-toggle-icon";
+import {
+  GeneratedImageOverlayProvider,
+  useGeneratedImageOverlay,
+} from "@/components/assistant-ui/generated-image-overlay-context";
+import { downloadImagePart } from "@/components/assistant-ui/image";
 import { MarkdownText } from "@/components/assistant-ui/markdown-text";
 import { MessageTiming } from "@/components/assistant-ui/message-timing";
 import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning";
@@ -39,13 +44,14 @@ import {
 import { sentAudioNames } from "@/features/chat/api/chat-adapter";
 import { parseExternalModelId } from "@/features/chat/external-providers";
 import { getExternalReasoningCapabilities } from "@/features/chat/provider-capabilities";
-import { useExternalProvidersStore } from "@/features/chat/stores/external-providers-store";
 import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
+import { useExternalProvidersStore } from "@/features/chat/stores/external-providers-store";
+import { deleteThreadMessage } from "@/features/chat/utils/delete-thread-message";
 import { applyQwenThinkingParams } from "@/features/chat/utils/qwen-params";
 import { isTauri } from "@/lib/api-base";
-import { deleteThreadMessage } from "@/features/chat/utils/delete-thread-message";
 import { AUDIO_ACCEPT, MAX_AUDIO_SIZE, fileToBase64 } from "@/lib/audio-utils";
 import { copyToClipboard } from "@/lib/copy-to-clipboard";
+import { toast } from "@/lib/toast";
 import { cn } from "@/lib/utils";
 import {
   ActionBarMorePrimitive,
@@ -79,30 +85,30 @@ import {
   TerminalIcon,
   XIcon,
 } from "lucide-react";
-import { Copy01Icon, Delete02Icon, Edit03Icon, Tick02Icon } from "@hugeicons/core-free-icons";
+import {
+  Copy01Icon,
+  Delete02Icon,
+  Edit03Icon,
+  Tick02Icon,
+} from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import {
   type ChangeEvent,
+  type ComponentProps,
   type CompositionEvent,
   type FC,
-  type FormEvent,
   type KeyboardEvent,
   useCallback,
   useEffect,
   useRef,
   useState,
 } from "react";
-import { toast } from "@/lib/toast";
 
 export const Thread: FC<{
   hideComposer?: boolean;
   hideWelcome?: boolean;
   targetThreadId?: string;
-}> = ({
-  hideComposer,
-  hideWelcome,
-  targetThreadId,
-}) => {
+}> = ({ hideComposer, hideWelcome, targetThreadId }) => {
   // Intent-aware autoscroll: replaces assistant-ui's built-in autoscroll
   // to prevent the streaming-mutation race that makes the viewport snap
   // back to the bottom while the user is scrolling up (see the hook for
@@ -113,85 +119,204 @@ export const Thread: FC<{
   const isComposerAttachPending = useAuiState(({ threads }) =>
     targetThreadId ? threads.mainThreadId !== targetThreadId : false,
   );
+  const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
+  const threadId = targetThreadId ?? activeThreadId ?? null;
 
   return (
-    <ThreadPrimitive.Root
-      className="aui-root aui-thread-root @container relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-hidden"
-      style={{
-        ["--thread-max-width" as string]: "48rem",
-        ["--thread-content-max-width" as string]:
-          "calc(var(--thread-max-width) - 1.5rem)",
-      }}
-    >
-      <IntentAwareScrollProvider value={autoScrollContext}>
-        <ThreadPrimitive.Viewport
-          ref={viewportRef}
-          autoScroll={false}
-          scrollToBottomOnRunStart={false}
-          scrollToBottomOnInitialize={false}
-          scrollToBottomOnThreadSwitch={false}
-          className={cn(
-            "aui-thread-viewport aui-stream-viewport relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-x-auto overflow-y-auto scroll-smooth px-5",
-            hideComposer ? "pt-4" : "pt-[48px]",
-          )}
-        >
-          {!hideWelcome && (
-            <AuiIf condition={({ thread }) => thread.isEmpty && !thread.isLoading}>
-              <ThreadWelcome hideComposer={hideComposer} />
-            </AuiIf>
-          )}
+    <GeneratedImageOverlayProvider key={threadId ?? "default"} threadId={threadId}>
+      <ThreadPrimitive.Root
+        className="aui-root aui-thread-root @container relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-hidden"
+        style={{
+          ["--thread-max-width" as string]: "48rem",
+          ["--thread-content-max-width" as string]:
+            "calc(var(--thread-max-width) - 1.5rem)",
+        }}
+      >
+        <IntentAwareScrollProvider value={autoScrollContext}>
+          <ThreadPrimitive.Viewport
+            ref={viewportRef}
+            autoScroll={false}
+            scrollToBottomOnRunStart={false}
+            scrollToBottomOnInitialize={false}
+            scrollToBottomOnThreadSwitch={false}
+            className={cn(
+              "aui-thread-viewport aui-stream-viewport relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-x-auto overflow-y-auto scroll-smooth px-5",
+              hideComposer ? "pt-4" : "pt-[48px]",
+            )}
+          >
+            {!hideWelcome && (
+              <AuiIf
+                condition={({ thread }) => thread.isEmpty && !thread.isLoading}
+              >
+                <ThreadWelcome hideComposer={hideComposer} threadId={threadId} />
+              </AuiIf>
+            )}
 
-          <ThreadPrimitive.Messages
-            components={{
-              UserMessage,
-              EditComposer,
-              AssistantMessage,
-            }}
-          />
+            <ThreadPrimitive.Messages
+              components={{
+                UserMessage,
+                EditComposer,
+                AssistantMessage,
+              }}
+            />
 
-          {/* Bottom slack so the last message has breathing room above the
+            {/* Bottom slack so the last message has breathing room above the
             sticky scroll-to-bottom button (and the floating composer in
             single mode). Without this, content would butt against the
             sticky footer and feel cramped. */}
-          <AuiIf condition={({ thread }) => hideWelcome || !thread.isEmpty}>
-            <div
-              className={cn("shrink-0", hideComposer ? "h-16" : "h-40")}
-              aria-hidden={true}
-            />
-          </AuiIf>
-
-          <AuiIf condition={({ thread }) => hideWelcome || !thread.isEmpty}>
-            <ThreadPrimitive.ViewportFooter
-              className={cn(
-                "aui-thread-viewport-footer pointer-events-none sticky z-20 flex w-full justify-center bg-transparent",
-                hideComposer ? "bottom-3" : "bottom-[140px]",
-              )}
-            >
-              <ThreadScrollToBottom />
-            </ThreadPrimitive.ViewportFooter>
-          </AuiIf>
-        </ThreadPrimitive.Viewport>
-
-        {!hideComposer && (
-          <AuiIf condition={({ thread }) => hideWelcome || !thread.isEmpty}>
-            <div className="aui-thread-composer-dock pointer-events-none absolute bottom-0 left-0 right-0 md:right-[10px] z-20">
+            <AuiIf condition={({ thread }) => hideWelcome || !thread.isEmpty}>
               <div
+                className={cn("shrink-0", hideComposer ? "h-16" : "h-40")}
                 aria-hidden={true}
-                className="absolute inset-x-0 bottom-0 top-[10px] bg-background"
               />
-              <div className="relative px-5 pb-2">
-                <div className="pointer-events-auto mx-auto w-full max-w-(--thread-max-width)">
-                  <ComposerAnimated disabled={isComposerAttachPending} />
-                </div>
-                <p className="composer-footer-note">
-                  LLMs can make mistakes. Double-check responses.
-                </p>
-              </div>
-            </div>
-          </AuiIf>
+            </AuiIf>
+
+            <AuiIf condition={({ thread }) => hideWelcome || !thread.isEmpty}>
+              <ThreadPrimitive.ViewportFooter
+                className={cn(
+                  "aui-thread-viewport-footer pointer-events-none sticky z-20 flex w-full justify-center bg-transparent",
+                  hideComposer ? "bottom-3" : "bottom-[140px]",
+                )}
+              >
+                <ThreadScrollToBottom />
+              </ThreadPrimitive.ViewportFooter>
+            </AuiIf>
+          </ThreadPrimitive.Viewport>
+
+          <GeneratedImageViewportOverlay hideComposer={hideComposer} />
+
+          {!hideComposer && (
+            <AuiIf condition={({ thread }) => hideWelcome || !thread.isEmpty}>
+              <ThreadComposerDock
+                disabled={isComposerAttachPending}
+                threadId={threadId}
+              />
+            </AuiIf>
+          )}
+        </IntentAwareScrollProvider>
+      </ThreadPrimitive.Root>
+    </GeneratedImageOverlayProvider>
+  );
+};
+
+const GeneratedImageViewportOverlay: FC<{ hideComposer?: boolean }> = ({
+  hideComposer,
+}) => {
+  const { overlay, closeOverlay } = useGeneratedImageOverlay();
+
+  useEffect(() => {
+    if (!overlay) {
+      return;
+    }
+    document.querySelector<HTMLTextAreaElement>(".aui-composer-input")?.focus();
+  }, [overlay]);
+
+  if (!overlay) {
+    return null;
+  }
+
+  return (
+    <div className="pointer-events-none absolute inset-0 z-30">
+      <button
+        type="button"
+        className="pointer-events-auto absolute inset-0 bg-background/65 backdrop-blur-[1px] dark:bg-background/55"
+        onClick={closeOverlay}
+        aria-label="Close generated image preview"
+      />
+      <section
+        className={cn(
+          "pointer-events-none absolute inset-x-5 top-[48px] flex flex-col items-center",
+          hideComposer ? "bottom-4" : "bottom-[150px]",
         )}
-      </IntentAwareScrollProvider>
-    </ThreadPrimitive.Root>
+        aria-label="Generated image preview"
+      >
+        <div className="pointer-events-auto relative flex min-h-0 w-full max-w-[1100px] flex-1 flex-col items-center justify-center gap-3 rounded-3xl bg-muted/10 p-3 ring-1 ring-border/20">
+          <div className="absolute inset-x-3 top-3 z-10 flex justify-end">
+            <div className="flex shrink-0 items-center gap-1 rounded-full bg-background/70 p-1 ring-1 ring-border/20 backdrop-blur-sm">
+              <Button
+                type="button"
+                variant="ghost"
+                size="icon-sm"
+                className="size-7 rounded-full"
+                onClick={() =>
+                  downloadImagePart({
+                    image: overlay.image,
+                    filename: overlay.filename,
+                  })
+                }
+                aria-label="Download generated image"
+              >
+                <DownloadIcon className="size-3.5" />
+              </Button>
+              <Button
+                type="button"
+                variant="ghost"
+                size="icon-sm"
+                className="size-7 rounded-full"
+                onClick={closeOverlay}
+                aria-label="Close generated image preview"
+              >
+                <XIcon className="size-3.5" />
+              </Button>
+            </div>
+          </div>
+          <div className="flex min-h-0 flex-1 items-center justify-center pt-1">
+            <img
+              src={overlay.image}
+              alt={overlay.title}
+              className="max-h-full max-w-full object-contain"
+            />
+          </div>
+          <div
+            className="w-full max-w-[min(100%,46rem)] shrink-0 text-center"
+            title={overlay.title}
+          >
+            <p className="truncate font-medium text-foreground/70 text-xs">
+              Generated image
+            </p>
+            {overlay.metadata ? (
+              <p className="truncate text-[11px] text-muted-foreground/75">
+                {overlay.metadata}
+              </p>
+            ) : null}
+            {hideComposer ? null : (
+              <p className="mt-1 text-muted-foreground text-xs">
+                Type edits below, then send.
+              </p>
+            )}
+          </div>
+        </div>
+      </section>
+    </div>
+  );
+};
+
+const ThreadComposerDock: FC<{
+  disabled?: boolean;
+  threadId?: string | null;
+}> = ({ disabled, threadId }) => {
+  const { overlay } = useGeneratedImageOverlay();
+
+  return (
+    <div
+      className={cn(
+        "aui-thread-composer-dock pointer-events-none absolute bottom-0 left-0 right-0 md:right-[10px]",
+        overlay ? "z-40" : "z-20",
+      )}
+    >
+      <div
+        aria-hidden={true}
+        className="absolute inset-x-0 bottom-0 top-[10px] bg-background"
+      />
+      <div className="relative px-5 pb-2">
+        <div className="pointer-events-auto mx-auto w-full max-w-(--thread-max-width)">
+          <ComposerAnimated disabled={disabled} threadId={threadId} />
+        </div>
+        <p className="composer-footer-note">
+          LLMs can make mistakes. Double-check responses.
+        </p>
+      </div>
+    </div>
   );
 };
 
@@ -219,13 +344,17 @@ const ThreadScrollToBottom: FC = () => {
   );
 };
 
-const ThreadWelcome: FC<{ hideComposer?: boolean }> = ({ hideComposer }) => {
+const ThreadWelcome: FC<{
+  hideComposer?: boolean;
+  threadId?: string | null;
+}> = ({ hideComposer, threadId }) => {
   const [currentEmoji, setCurrentEmoji] = useState("large sloth drink.png");
 
   useEffect(() => {
     const hour = new Date().getHours();
     if (hour >= 6 && hour < 12) setCurrentEmoji("large sloth drink.png");
-    else if (hour >= 12 && hour < 17) setCurrentEmoji("sloth magnify final.png");
+    else if (hour >= 12 && hour < 17)
+      setCurrentEmoji("sloth magnify final.png");
     else if (hour >= 17 && hour < 21) setCurrentEmoji("sloth shy large.png");
     else setCurrentEmoji("unsloth-gem.png");
   }, []);
@@ -240,11 +369,7 @@ const ThreadWelcome: FC<{ hideComposer?: boolean }> = ({ hideComposer }) => {
       <div className="aui-thread-welcome-center flex w-full grow flex-col items-center justify-center pb-[48px]">
         <div className="aui-thread-welcome-message flex w-full flex-col justify-center gap-6 px-4">
           <div className="flex flex-col items-center gap-2 text-center">
-            <img
-              src={currentEmojiSrc}
-              alt="Sloth mascot"
-              className="size-20"
-            />
+            <img src={currentEmojiSrc} alt="Sloth mascot" className="size-20" />
             <h1 className="aui-thread-welcome-message-inner fade-in slide-in-from-bottom-1 animate-in font-heading font-semibold text-2xl tracking-[-0.02em] duration-200">
               Chat with your model
             </h1>
@@ -252,18 +377,21 @@ const ThreadWelcome: FC<{ hideComposer?: boolean }> = ({ hideComposer }) => {
               Run GGUFs, safetensors, vision and audio models
             </p>
           </div>
-          {!hideComposer && <ComposerAnimated />}
+          {!hideComposer && <ComposerAnimated threadId={threadId} />}
         </div>
       </div>
     </div>
   );
 };
 
-const ComposerAnimated: FC<{ disabled?: boolean }> = ({ disabled }) => {
+const ComposerAnimated: FC<{
+  disabled?: boolean;
+  threadId?: string | null;
+}> = ({ disabled, threadId }) => {
   return (
     <div className="relative mx-auto min-w-0 w-full max-w-(--thread-max-width)">
       <div className="relative z-10 w-full">
-        <Composer disabled={disabled} />
+        <Composer disabled={disabled} threadId={threadId} />
       </div>
     </div>
   );
@@ -293,8 +421,21 @@ const PendingAudioChip: FC = () => {
   );
 };
 
-const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => {
-  const { inputProps, isComposing, isComposingRef } = useImeComposerInputHandlers();
+const Composer: FC<{
+  disabled?: boolean;
+  threadId?: string | null;
+}> = ({ disabled, threadId }) => {
+  const aui = useAui();
+  const { overlay, closeOverlay } = useGeneratedImageOverlay();
+  const setImageToolsEnabled = useChatRuntimeStore(
+    (s) => s.setImageToolsEnabled,
+  );
+  const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
+  const setPendingImageEditReference = useChatRuntimeStore(
+    (s) => s.setPendingImageEditReference,
+  );
+  const { inputProps, isComposing, isComposingRef } =
+    useImeComposerInputHandlers();
   const composerText = useAuiState(({ composer }) => composer.text);
   const hasAttachments = useAuiState(
     ({ composer }) => composer.attachments.length > 0,
@@ -304,22 +445,78 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => {
       (attachment) => attachment.status.type === "running",
     ),
   );
-  const hasPendingAudio = useChatRuntimeStore((s) => Boolean(s.pendingAudioName));
+  const hasPendingAudio = useChatRuntimeStore((s) =>
+    Boolean(s.pendingAudioName),
+  );
+  const referenceThreadId = threadId ?? activeThreadId ?? null;
   const hasSendableContent =
     composerText.trim().length > 0 || hasAttachments || hasPendingAudio;
+  const shouldBlockSend = useCallback(
+    () =>
+      !hasSendableContent || isComposingRef.current || hasPendingAttachments,
+    [hasPendingAttachments, hasSendableContent, isComposingRef],
+  );
 
   const handleSubmit = useCallback(
-    (event: FormEvent<HTMLFormElement>) => {
-      if (
-        disabled ||
-        !hasSendableContent ||
-        isComposingRef.current ||
-        hasPendingAttachments
-      ) {
+    (event: Parameters<NonNullable<ComponentProps<"form">["onSubmit"]>>[0]) => {
+      if (disabled || shouldBlockSend()) {
         event.preventDefault();
+        return;
+      }
+
+      if (overlay) {
+        const trimmed = composerText.trim();
+        if (!trimmed) {
+          event.preventDefault();
+          return;
+        }
+        if (!overlay.openaiImageGenerationCallId) {
+          event.preventDefault();
+          toast.error("This generated image cannot be edited", {
+            description:
+              "The original image reference is missing. Generate the image again, then retry the edit.",
+          });
+          closeOverlay();
+          return;
+        }
+        if ((overlay.threadId ?? null) !== referenceThreadId) {
+          event.preventDefault();
+          toast.error("This generated image belongs to another chat", {
+            description: "Open the original chat and retry the edit.",
+          });
+          closeOverlay();
+          return;
+        }
+        setImageToolsEnabled(true);
+        setPendingImageEditReference({
+          threadId: overlay.threadId ?? referenceThreadId,
+          openaiImageGenerationCallId: overlay.openaiImageGenerationCallId,
+          ...(overlay.openaiResponseId
+            ? { openaiResponseId: overlay.openaiResponseId }
+            : {}),
+          openaiReasoningItem: overlay.openaiReasoningItem,
+        });
+        flushResourcesSync(() => {
+          aui
+            .composer()
+            .setText(
+              `Use the selected generated image as the reference and apply this edit: ${trimmed}. Preserve everything else exactly.`,
+            );
+        });
+        closeOverlay();
       }
     },
-    [disabled, hasPendingAttachments, hasSendableContent, isComposingRef],
+    [
+      aui,
+      closeOverlay,
+      composerText,
+      disabled,
+      overlay,
+      referenceThreadId,
+      setImageToolsEnabled,
+      setPendingImageEditReference,
+      shouldBlockSend,
+    ],
   );
 
   const composerContent = (
@@ -342,11 +539,12 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => {
       />
       <ComposerAction
         disabled={
-          disabled || !hasSendableContent || isComposing || hasPendingAttachments
-        }
-        blockSend={() =>
-          !hasSendableContent || isComposingRef.current || hasPendingAttachments
+          disabled ||
+          !hasSendableContent ||
+          isComposing ||
+          hasPendingAttachments
         }
+        shouldBlockSend={shouldBlockSend}
       />
     </>
   );
@@ -553,7 +751,6 @@ const ComposerAudioUpload: FC = () => {
   );
 };
 
-
 const ReasoningToggle: FC = () => {
   const modelLoaded = useChatRuntimeStore(
     (s) => !!s.params.checkpoint && !s.modelLoading,
@@ -565,8 +762,12 @@ const ReasoningToggle: FC = () => {
   const setReasoningEnabled = useChatRuntimeStore((s) => s.setReasoningEnabled);
   const reasoningStyle = useChatRuntimeStore((s) => s.reasoningStyle);
   const reasoningEffort = useChatRuntimeStore((s) => s.reasoningEffort);
-  const supportsReasoningOff = useChatRuntimeStore((s) => s.supportsReasoningOff);
-  const reasoningEffortLevels = useChatRuntimeStore((s) => s.reasoningEffortLevels);
+  const supportsReasoningOff = useChatRuntimeStore(
+    (s) => s.supportsReasoningOff,
+  );
+  const reasoningEffortLevels = useChatRuntimeStore(
+    (s) => s.reasoningEffortLevels,
+  );
   const setReasoningEffort = useChatRuntimeStore((s) => s.setReasoningEffort);
   const lastOpenRouterChosenModel = useChatRuntimeStore(
     (s) => s.lastOpenRouterChosenModel,
@@ -619,7 +820,8 @@ const ReasoningToggle: FC = () => {
     effectiveReasoningEnabled && reasoningEffort !== "none";
   const disabled = !(modelLoaded && effectiveSupportsReasoning);
   const formatEffortLabel = (level: typeof reasoningEffort): string => {
-    if (level !== "xhigh") return level.charAt(0).toUpperCase() + level.slice(1);
+    if (level !== "xhigh")
+      return level.charAt(0).toUpperCase() + level.slice(1);
     const normalized = externalSelection?.modelId?.trim().toLowerCase() ?? "";
     if (
       normalized.startsWith("claude-opus-4-6") ||
@@ -677,23 +879,25 @@ const ReasoningToggle: FC = () => {
           {effectiveReasoningEffortLevels
             .filter((level) => level !== "none")
             .map((level) => (
-            <DropdownMenuItem
-              key={level}
-              onSelect={() => {
-                setReasoningEffort(level);
-                setReasoningEnabled(true);
-                applyQwenThinkingParams(true);
-                // Kimi's $web_search builtin forbids thinking, so
-                // enabling thinking flips the Search pill off.
-                if (isKimiExternal && toolsEnabled) {
-                  setToolsEnabled(false);
-                }
-              }}
-            >
-              {formatEffortLabel(level)}
-              {effectiveReasoningVisualEnabled && reasoningEffort === level ? " \u2713" : ""}
-            </DropdownMenuItem>
-          ))}
+              <DropdownMenuItem
+                key={level}
+                onSelect={() => {
+                  setReasoningEffort(level);
+                  setReasoningEnabled(true);
+                  applyQwenThinkingParams(true);
+                  // Kimi's $web_search builtin forbids thinking, so
+                  // enabling thinking flips the Search pill off.
+                  if (isKimiExternal && toolsEnabled) {
+                    setToolsEnabled(false);
+                  }
+                }}
+              >
+                {formatEffortLabel(level)}
+                {effectiveReasoningVisualEnabled && reasoningEffort === level
+                  ? " \u2713"
+                  : ""}
+              </DropdownMenuItem>
+            ))}
         </DropdownMenuContent>
       </DropdownMenu>
     );
@@ -808,8 +1012,7 @@ const WebSearchToggle: FC = () => {
       ? externalProviders.find((p) => p.id === externalSelection.providerId)
       : undefined;
   const isKimiExternal = selectedExternalProvider?.providerType === "kimi";
-  const disabled =
-    !modelLoaded || !(supportsTools || supportsBuiltinWebSearch);
+  const disabled = !modelLoaded || !(supportsTools || supportsBuiltinWebSearch);
 
   return (
     <button
@@ -899,7 +1102,9 @@ const ImagesToggle: FC = () => {
       className="composer-pill-btn"
       data-active={imageToolsEnabled && !disabled ? "true" : "false"}
       aria-label={
-        imageToolsEnabled ? "Disable image generation" : "Enable image generation"
+        imageToolsEnabled
+          ? "Disable image generation"
+          : "Enable image generation"
       }
     >
       <ImageIcon className="size-3.5" />
@@ -966,10 +1171,10 @@ const ToolStatusDisplay: FC = () => {
   );
 };
 
-const ComposerAction: FC<{ disabled?: boolean; blockSend?: () => boolean }> = ({
-  disabled,
-  blockSend,
-}) => {
+const ComposerAction: FC<{
+  disabled?: boolean;
+  shouldBlockSend?: () => boolean;
+}> = ({ disabled, shouldBlockSend }) => {
   return (
     <div className="aui-composer-action-wrapper composer-action-wrapper">
       <div className="flex items-center gap-0.5">
@@ -1016,7 +1221,7 @@ const ComposerAction: FC<{ disabled?: boolean; blockSend?: () => boolean }> = ({
               size="icon"
               disabled={disabled}
               onClick={(event) => {
-                if (blockSend?.()) {
+                if (shouldBlockSend?.()) {
                   event.preventDefault();
                 }
               }}
@@ -1284,7 +1489,11 @@ const UserActionBar: FC = () => {
       <CopyButton />
       <ActionBarPrimitive.Edit asChild={true}>
         <TooltipIconButton tooltip="Edit" className="aui-user-action-edit">
-          <HugeiconsIcon icon={Edit03Icon} strokeWidth={1.75} className="size-icon" />
+          <HugeiconsIcon
+            icon={Edit03Icon}
+            strokeWidth={1.75}
+            className="size-icon"
+          />
         </TooltipIconButton>
       </ActionBarPrimitive.Edit>
       <DeleteMessageButton />

studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx

@@ -3,9 +3,14 @@
 
 "use client";
 
-import { type ToolCallMessagePartComponent, useAuiState } from "@assistant-ui/react";
-import { ImageIcon, LoaderIcon } from "lucide-react";
-import { memo, useEffect, useState } from "react";
+import { Button } from "@/components/ui/button";
+import { cn } from "@/lib/utils";
+import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
+import { DownloadIcon, ImageIcon, PencilIcon } from "lucide-react";
+import type { CSSProperties, MouseEvent } from "react";
+import { memo, useState } from "react";
+import { useGeneratedImageOverlay } from "./generated-image-overlay-context";
+import { Image, downloadImagePart } from "./image";
 import {
   ToolFallbackContent,
   ToolFallbackRoot,
@@ -38,6 +43,9 @@ import {
 interface ImageGenerationArgs {
   prompt?: string;
   kind?: string;
+  openai_image_generation_call_id?: unknown;
+  openai_response_id?: unknown;
+  openai_reasoning_item?: unknown;
 }
 
 interface ImageGenerationResult {
@@ -46,6 +54,83 @@ interface ImageGenerationResult {
   size?: string;
   quality?: string;
   background?: string;
+  prompt?: string;
+}
+
+type GeneratedImagePart = {
+  type: "image";
+  image: string;
+  filename?: string;
+};
+
+const extensionForMime = (mime: string): string => {
+  switch (mime.toLowerCase()) {
+    case "image/jpeg":
+    case "image/jpg":
+      return "jpg";
+    case "image/webp":
+      return "webp";
+    case "image/gif":
+      return "gif";
+    case "image/svg+xml":
+      return "svg";
+    default:
+      return "png";
+  }
+};
+
+const imageFilenameFromPrompt = (prompt: string, mime: string): string => {
+  const slug = prompt
+    .trim()
+    .toLowerCase()
+    .replace(/[^a-z0-9]+/g, "-")
+    .replace(/^-|-$/g, "")
+    .slice(0, 48);
+  return `${slug || "generated-image"}.${extensionForMime(mime)}`;
+};
+
+const formatGeneratedImageLabel = (prompt: string): string => {
+  if (!prompt) {
+    return "Generated image";
+  }
+  return prompt.length > 80
+    ? `Generated image: ${prompt.slice(0, 80)}…`
+    : `Generated image: ${prompt}`;
+};
+
+const loadingDots = Array.from({ length: 64 }, (_, index) => {
+  const row = Math.floor(index / 8);
+  const col = index % 8;
+  return (
+    <span
+      key={index}
+      className="generated-image-loading-dot"
+      style={
+        {
+          "--dot-row": row,
+          "--dot-col": col,
+        } as CSSProperties
+      }
+    />
+  );
+});
+
+function GeneratedImagePlaceholder({ label }: { label: string }) {
+  return (
+    <div
+      className={cn(
+        "generated-image-loading-card flex aspect-square w-[480px] max-w-full items-center justify-center rounded-2xl border border-border/70 bg-muted/15",
+      )}
+      aria-busy="true"
+      aria-label={label}
+      aria-live="polite"
+    >
+      <span className="sr-only">{label}</span>
+      <div className="generated-image-loading-wave" aria-hidden={true}>
+        {loadingDots}
+      </div>
+    </div>
+  );
 }
 
 const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({
@@ -53,6 +138,7 @@ const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({
   result,
   status,
 }) => {
+  const { openOverlay } = useGeneratedImageOverlay();
   const parsedArgs = (args as ImageGenerationArgs) ?? {};
   const prompt = parsedArgs.prompt ?? "";
   const isRunning = status?.type === "running";
@@ -66,33 +152,74 @@ const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({
   const imageSrc = imageResult?.image_b64
     ? `data:${mime};base64,${imageResult.image_b64}`
     : null;
+  const imageTitle =
+    imageResult?.prompt?.trim() || prompt.trim() || "Generated image";
+  const imageMetadata = [imageResult?.size, imageResult?.quality, mime]
+    .filter(Boolean)
+    .join(" · ");
+  const openaiImageGenerationCallId =
+    typeof parsedArgs.openai_image_generation_call_id === "string"
+      ? parsedArgs.openai_image_generation_call_id
+      : undefined;
+  const openaiResponseId =
+    typeof parsedArgs.openai_response_id === "string"
+      ? parsedArgs.openai_response_id
+      : undefined;
+  const imagePart: GeneratedImagePart | null = imageSrc
+    ? {
+        type: "image",
+        image: imageSrc,
+        filename: imageFilenameFromPrompt(prompt, mime),
+      }
+    : 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 isPendingImage = !imagePart && status?.type === "running";
 
   const runningLabel = "Generating image…";
-  const completedLabel = prompt
-    ? prompt.length > 80
-      ? `Generated image: ${prompt.slice(0, 80)}…`
-      : `Generated image: ${prompt}`
-    : "Generated image";
+  const completedLabel = formatGeneratedImageLabel(prompt);
+
+  const showPreview = () => {
+    if (!imagePart) {
+      return;
+    }
+    openOverlay({
+      image: imagePart.image,
+      title: imageTitle,
+      metadata: imageMetadata,
+      filename: imagePart.filename,
+      openaiImageGenerationCallId,
+      openaiResponseId,
+      openaiReasoningItem: parsedArgs.openai_reasoning_item,
+    });
+  };
+
+  const stopOverlayActionPropagation = (
+    event: MouseEvent<HTMLButtonElement>,
+  ) => {
+    event.preventDefault();
+    event.stopPropagation();
+  };
+
+  const handleDownload = (event: MouseEvent<HTMLButtonElement>) => {
+    stopOverlayActionPropagation(event);
+    if (imagePart) {
+      downloadImagePart(imagePart);
+    }
+  };
+
+  const handleEditClick = (event: MouseEvent<HTMLButtonElement>) => {
+    stopOverlayActionPropagation(event);
+    showPreview();
+  };
+
+  if (isPendingImage) {
+    return (
+      <div className="aui-tool-fallback-root w-full py-1">
+        <GeneratedImagePlaceholder label={runningLabel} />
+      </div>
+    );
+  }
 
   return (
     <ToolFallbackRoot open={open} onOpenChange={setOpen}>
@@ -102,20 +229,47 @@ const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({
         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"
-            />
+        {imagePart ? (
+          <figure className="m-0 flex flex-col gap-2">
+            <div className="group/generated-image relative aspect-square w-[480px] max-w-full overflow-hidden rounded-2xl border border-border/70 bg-muted/30 shadow-sm">
+              <button
+                type="button"
+                className="block size-full cursor-zoom-in overflow-hidden rounded-2xl focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 focus-visible:ring-offset-background"
+                onClick={showPreview}
+                aria-label="Open generated image preview"
+              >
+                <Image.Preview
+                  src={imagePart.image}
+                  alt={imageTitle}
+                  containerClassName="size-full min-h-0 bg-background"
+                  className="size-full object-cover"
+                />
+              </button>
+              <div className="pointer-events-none absolute inset-x-0 bottom-0 flex items-end justify-between gap-2 bg-gradient-to-t from-black/55 via-black/20 to-transparent p-3 opacity-100 transition-opacity sm:opacity-0 sm:group-hover/generated-image:opacity-100 sm:group-focus-within/generated-image:opacity-100">
+                <Button
+                  type="button"
+                  variant="dark"
+                  size="sm"
+                  className="pointer-events-auto h-8 rounded-full bg-black/70 text-white hover:bg-black/85"
+                  onClick={handleEditClick}
+                >
+                  <PencilIcon className="size-3.5" />
+                  Edit
+                </Button>
+                <Button
+                  type="button"
+                  variant="dark"
+                  size="icon-sm"
+                  className="pointer-events-auto rounded-full bg-black/70 text-white hover:bg-black/85"
+                  onClick={handleDownload}
+                  aria-label="Download generated image"
+                >
+                  <DownloadIcon className="size-4" />
+                </Button>
+              </div>
+            </div>
             {prompt ? (
-              <figcaption className="text-xs leading-snug text-muted-foreground">
+              <figcaption className="max-w-[480px] text-xs leading-snug text-muted-foreground">
                 {prompt}
               </figcaption>
             ) : null}

studio/frontend/src/features/chat/api/chat-adapter.ts

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: AGPL-3.0-only
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
-import { getAuthToken } from "@/features/auth/session";
+import { getAuthToken } from "@/features/auth";
 import { apiUrl } from "@/lib/api-base";
 import { toast } from "@/lib/toast";
 import type { MessageTiming, ToolCallMessagePart } from "@assistant-ui/core";
@@ -30,12 +30,17 @@ import {
   providerSupportsBuiltinWebSearch,
   providerSupportsFastMode,
 } from "../provider-capabilities";
-import { useChatRuntimeStore } from "../stores/chat-runtime-store";
+import {
+  type PendingImageEditReference,
+  useChatRuntimeStore,
+} from "../stores/chat-runtime-store";
 import { useExternalProvidersStore } from "../stores/external-providers-store";
 import { isMultimodalResponse } from "../types/api";
 import type {
   OpenAIChatCompletionsRequest,
+  OpenAIChatMessage,
   OpenAIMessageContent,
+  OpenAIReasoningContentPart,
 } from "../types/api";
 import type { ChatModelSummary } from "../types/runtime";
 import { getImageInputUnavailableReason } from "../utils/image-input-support";
@@ -399,37 +404,30 @@ function collectImageParts(
   message: RunMessage,
 ): Array<{ type: "image_url"; image_url: { url: string } }> {
   const parts: Array<{ type: "image_url"; image_url: { url: string } }> = [];
+  const pushImagePart = (part: { type: string }) => {
+    if (part.type !== "image" || !("image" in part)) {
+      return;
+    }
+    const src = (part as { image: string }).image;
+    if (!src) {
+      return;
+    }
+    parts.push({
+      type: "image_url",
+      image_url: {
+        url: src.startsWith("data:") ? src : `data:image/png;base64,${src}`,
+      },
+    });
+  };
 
   for (const part of message.content ?? []) {
-    if (part.type === "image" && "image" in part) {
-      const src = (part as { image: string }).image;
-      if (src) {
-        parts.push({
-          type: "image_url",
-          image_url: {
-            url: src.startsWith("data:") ? src : `data:image/png;base64,${src}`,
-          },
-        });
-      }
-    }
+    pushImagePart(part);
   }
 
   if ("attachments" in message && (message.attachments?.length ?? 0) > 0) {
     for (const attachment of message.attachments ?? []) {
       for (const part of attachment.content ?? []) {
-        if (part.type === "image" && "image" in part) {
-          const src = (part as { image: string }).image;
-          if (src) {
-            parts.push({
-              type: "image_url",
-              image_url: {
-                url: src.startsWith("data:")
-                  ? src
-                  : `data:image/png;base64,${src}`,
-              },
-            });
-          }
-        }
+        pushImagePart(part);
       }
     }
   }
@@ -437,6 +435,66 @@ function collectImageParts(
   return parts;
 }
 
+function normalizeOpenAIReasoningItem(
+  value: unknown,
+): OpenAIReasoningContentPart | null {
+  if (!value || typeof value !== "object") {
+    return null;
+  }
+  const item = value as Record<string, unknown>;
+  if (item.type !== "reasoning" || typeof item.id !== "string" || !item.id) {
+    return null;
+  }
+  const summary = Array.isArray(item.summary)
+    ? item.summary.flatMap((part) => {
+        if (!part || typeof part !== "object") {
+          return [];
+        }
+        const summaryPart = part as Record<string, unknown>;
+        return summaryPart.type === "summary_text" &&
+          typeof summaryPart.text === "string"
+          ? [{ type: "summary_text" as const, text: summaryPart.text }]
+          : [];
+      })
+    : [];
+  const normalized: OpenAIReasoningContentPart = {
+    type: "reasoning",
+    id: item.id,
+    summary,
+  };
+  if (
+    item.status === "in_progress" ||
+    item.status === "completed" ||
+    item.status === "incomplete"
+  ) {
+    normalized.status = item.status;
+  }
+  return normalized;
+}
+
+function toOpenAIImageEditReferenceMessage(
+  reference: PendingImageEditReference,
+): OpenAIChatMessage | null {
+  if (!reference.openaiImageGenerationCallId) {
+    return null;
+  }
+  const content: Exclude<OpenAIMessageContent, string> = [];
+  const reasoningItem = normalizeOpenAIReasoningItem(
+    reference.openaiReasoningItem,
+  );
+  if (reasoningItem) {
+    content.push(reasoningItem);
+  }
+  content.push({
+    type: "image_generation_call",
+    id: reference.openaiImageGenerationCallId,
+    ...(reference.openaiResponseId
+      ? { response_id: reference.openaiResponseId }
+      : {}),
+  });
+  return { role: "assistant", content };
+}
+
 // Refusal flag stamped on assistant metadata when the backend emits the
 // `anthropic_refusal` _toolEvent. We drop the refused pair from the next
 // request body (Anthropic guidance: leaving refusals in context keeps
@@ -480,10 +538,16 @@ function toOpenAIMessage(message: RunMessage): {
   if (imageParts.length > 0) {
     return {
       role: message.role,
-      content: [{ type: "text", text: textContent }, ...imageParts],
+      content: [
+        ...(textContent ? [{ type: "text" as const, text: textContent }] : []),
+        ...imageParts,
+      ],
     };
   }
 
+  if (!textContent) {
+    return null;
+  }
   return { role: message.role, content: textContent };
 }
 
@@ -918,17 +982,52 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
       // the user switches chats while waiting for model load / auto-load.
       const resolvedThreadId =
         (unstable_threadId ?? runtime.activeThreadId) || undefined;
+      const resolvedThreadKey = resolvedThreadId ?? null;
+      const pendingImageEditReferenceForRun = runtime.pendingImageEditReference;
+      const selectedImageEditReference =
+        (pendingImageEditReferenceForRun?.threadId ?? null) ===
+        resolvedThreadKey
+          ? pendingImageEditReferenceForRun
+          : null;
+      const clearSelectedImageEditReference = () => {
+        if (!selectedImageEditReference) {
+          return;
+        }
+        const store = useChatRuntimeStore.getState();
+        const pending = store.pendingImageEditReference;
+        if (
+          pending?.openaiImageGenerationCallId ===
+            selectedImageEditReference.openaiImageGenerationCallId &&
+          pending.openaiResponseId ===
+            selectedImageEditReference.openaiResponseId &&
+          (pending.threadId ?? null) ===
+            (selectedImageEditReference.threadId ?? null)
+        ) {
+          store.clearPendingImageEditReference();
+        }
+      };
 
       // Wait for in-progress model load to finish before inferring
       if (runtime.modelLoading) {
         toast.info("Waiting for model to finish loading…");
-        await waitForModelReady(abortSignal);
+        try {
+          await waitForModelReady(abortSignal);
+        } catch (error) {
+          clearSelectedImageEditReference();
+          throw error;
+        }
       }
 
       if (!useChatRuntimeStore.getState().params.checkpoint) {
         // Auto-load the smallest downloaded model
-        const { loaded, blockedByTrustRemoteCode } =
-          await autoLoadSmallestModel();
+        let loaded: boolean;
+        let blockedByTrustRemoteCode: boolean;
+        try {
+          ({ loaded, blockedByTrustRemoteCode } = await autoLoadSmallestModel());
+        } catch (error) {
+          clearSelectedImageEditReference();
+          throw error;
+        }
         if (!loaded) {
           toast.error(
             blockedByTrustRemoteCode
@@ -940,6 +1039,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                 : "Pick a model in the top bar, then retry.",
             },
           );
+          clearSelectedImageEditReference();
           throw new Error("Load a model first.");
         }
       }
@@ -964,6 +1064,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
           description:
             "Turn on Enable connections in Settings → Connections to use hosted models.",
         });
+        clearSelectedImageEditReference();
         throw new Error("Connections disabled.");
       }
       const externalProvider = isExternalRequest
@@ -979,6 +1080,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
         toast.error("Connection not found.", {
           description: "Open Settings → Connections and add it again.",
         });
+        clearSelectedImageEditReference();
         throw new Error("Connection not found.");
       }
       // Local providers (llama.cpp / vLLM / Ollama) allow an empty key — only block hosted providers.
@@ -989,36 +1091,34 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
         toast.error("Missing API key for selected connection.", {
           description: "Open Settings → Connections and set the API key again.",
         });
+        clearSelectedImageEditReference();
         throw new Error("Missing connection API key.");
       }
 
-      const webSearchEnabledForThisTurn =
-        Boolean(
-          externalProvider &&
-            toolsEnabled &&
-            providerSupportsBuiltinWebSearch(externalProvider.providerType),
-        );
-      const codeExecEnabledForThisTurn =
-        Boolean(
-          externalProvider &&
-            externalSelection &&
-            codeToolsEnabled &&
-            providerSupportsBuiltinCodeExecution(
-              externalProvider.providerType,
-              externalSelection.modelId,
-              externalProvider.baseUrl,
-            ),
-        );
+      const webSearchEnabledForThisTurn = Boolean(
+        externalProvider &&
+          toolsEnabled &&
+          providerSupportsBuiltinWebSearch(externalProvider.providerType),
+      );
+      const codeExecEnabledForThisTurn = Boolean(
+        externalProvider &&
+          externalSelection &&
+          codeToolsEnabled &&
+          providerSupportsBuiltinCodeExecution(
+            externalProvider.providerType,
+            externalSelection.modelId,
+            externalProvider.baseUrl,
+          ),
+      );
       // Fetch pill is independent of Search (Anthropic bills web_fetch
       // separately from web_search). Sourced from `webFetchToolsEnabled`;
       // on providers without web_fetch the toggle is forced off in
       // chat-page's runtime setState.
-      const webFetchEnabledForThisTurn =
-        Boolean(
-          externalProvider &&
-            webFetchToolsEnabled &&
-            providerSupportsBuiltinWebFetch(externalProvider.providerType),
-        );
+      const webFetchEnabledForThisTurn = Boolean(
+        externalProvider &&
+          webFetchToolsEnabled &&
+          providerSupportsBuiltinWebFetch(externalProvider.providerType),
+      );
       const providerShipsWebFetch = Boolean(
         externalProvider &&
           providerSupportsBuiltinWebFetch(externalProvider.providerType),
@@ -1038,6 +1138,15 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
           ),
       );
 
+      if (selectedImageEditReference && !imageGenerationEnabledForThisTurn) {
+        clearSelectedImageEditReference();
+        toast.error("Image editing is unavailable", {
+          description:
+            "Select an OpenAI image-generation model, then retry the edit.",
+        });
+        throw new Error("Image generation edit unavailable.");
+      }
+
       // Two-pass build: a refused assistant turn also drops the user
       // prompt that triggered it (leaving it in context re-triggers
       // the classifier). Refusal flag rides assistant
@@ -1060,6 +1169,27 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
         .filter((message): message is NonNullable<typeof message> =>
           Boolean(message),
         );
+      if (selectedImageEditReference) {
+        const referenceMessage = toOpenAIImageEditReferenceMessage(
+          selectedImageEditReference,
+        );
+        if (!referenceMessage) {
+          clearSelectedImageEditReference();
+          toast.error("This generated image cannot be edited", {
+            description:
+              "The original image reference is missing. Generate the image again, then retry the edit.",
+          });
+          throw new Error("Generated image edit reference missing.");
+        }
+        let insertAt = outboundMessages.length;
+        for (let i = outboundMessages.length - 1; i >= 0; i -= 1) {
+          if (outboundMessages[i]?.role === "user") {
+            insertAt = i;
+            break;
+          }
+        }
+        outboundMessages.splice(insertAt, 0, referenceMessage);
+      }
 
       const safeSystemPrompt =
         typeof params.systemPrompt === "string" ? params.systemPrompt : "";
@@ -1084,24 +1214,45 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
         // was on and suppressed live web_fetch calls.
         const anyWebEnabledForThisTurn =
           webSearchEnabledForThisTurn || webFetchEnabledForThisTurn;
-        if (!anyWebEnabledForThisTurn && !codeExecEnabledForThisTurn) {
+        if (
+          !anyWebEnabledForThisTurn &&
+          !codeExecEnabledForThisTurn &&
+          !imageGenerationEnabledForThisTurn
+        ) {
+          disabledToolGuard =
+            `You do not have ${webLabel}, code execution, or image generation tools in this conversation. ` +
+            "Answer from your own knowledge. " +
+            "If a request genuinely requires tool use, live data fetch, running code, or image generation, " +
+            "inform the user that you do not have access to these capabilities. " +
+            "Do not return tool-call syntax inside your response.";
+        } else if (!anyWebEnabledForThisTurn && !codeExecEnabledForThisTurn) {
           disabledToolGuard =
             `You do not have ${webLabel} or code execution tools in this conversation. ` +
-            "Answer from your own knowledge. " +
-            "If a request genuinely requires tool use, live data fetch or running code, " +
+            "You may still use image generation tools when they are available and useful. " +
+            "If a request genuinely requires live data fetch or running code, " +
             "inform the user that you do not have access to these capabilities. " +
             "Do not return tool-call syntax inside your response.";
         } else if (!anyWebEnabledForThisTurn) {
+          const availableTools = [
+            codeExecEnabledForThisTurn ? "code execution" : null,
+            imageGenerationEnabledForThisTurn ? "image generation" : null,
+          ].filter(Boolean);
           disabledToolGuard =
             `You do not have ${webLabel} tools in this conversation. ` +
-            "You may still use code execution tools when they are available and useful. " +
+            (availableTools.length > 0
+              ? `You may still use ${availableTools.join(" and ")} tools when they are available and useful. `
+              : "") +
             "If a request genuinely requires live data fetch or web search tool use, " +
             "inform the user that you do not have access to these capabilities. " +
             "Do not return tool-call syntax inside your response.";
         } else if (!codeExecEnabledForThisTurn) {
+          const availableTools = [
+            webLabel,
+            imageGenerationEnabledForThisTurn ? "image generation" : null,
+          ].filter(Boolean);
           disabledToolGuard =
             "You do not have code execution tools in this conversation. " +
-            `You may still use ${webLabel} tools when they are available and useful. ` +
+            `You may still use ${availableTools.join(" and ")} tools when they are available and useful. ` +
             "If a request genuinely requires running code or code execution tool use, " +
             "inform the user that you do not have access to these capabilities. " +
             "Do not return tool-call syntax inside your response.";
@@ -1163,6 +1314,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
           const gatedThreadKey = resolvedThreadId || "__default";
           runtime.setThreadRunning(gatedThreadKey, true);
           runtime.setThreadRunning(gatedThreadKey, false);
+          clearSelectedImageEditReference();
           throw new Error(imageGateReason);
         }
       }
@@ -1474,8 +1626,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                     ) {
                       void updateStoredChatThreadEventually(t.id, {
                         openaiCodeExecContainerId: null,
-                      })
-                        .catch(() => {});
+                      }).catch(() => {});
                       continue;
                     }
                     openaiCodeExecContainerId = t.openaiCodeExecContainerId;
@@ -1519,8 +1670,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                   openaiCodeExecContainerId = created.id;
                   void updateStoredChatThreadEventually(resolvedThreadId, {
                     openaiCodeExecContainerId: created.id,
-                  })
-                    .catch(() => {});
+                  }).catch(() => {});
                 } catch {
                   // Fall back to backend's container_auto path on
                   // failure — keeps the chat moving; the next turn
@@ -1628,7 +1778,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
               // attaches `cache_control.ttl` when the value is one of
               // "5m" / "1h" (see external_provider.py near line 1375),
               // so unknown values are a no-op end-to-end.
-              ...(supportsProviderPromptCacheTtl(externalProvider.providerType) &&
+              ...(supportsProviderPromptCacheTtl(
+                externalProvider.providerType,
+              ) &&
               (externalProvider.enablePromptCaching ?? true) &&
               isPromptCacheTtl(externalProvider.promptCacheTtl)
                 ? { prompt_cache_ttl: externalProvider.promptCacheTtl }
@@ -1706,10 +1858,15 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
         let retriedWithRefreshedKey = false;
         while (true) {
           try {
-            const stream = streamChatCompletions(
-              await buildRequestPayload(retriedWithRefreshedKey),
-              abortSignal,
-            );
+            let requestPayload: OpenAIChatCompletionsRequest;
+            try {
+              requestPayload = await buildRequestPayload(retriedWithRefreshedKey);
+            } catch (error) {
+              clearSelectedImageEditReference();
+              throw error;
+            }
+            clearSelectedImageEditReference();
+            const stream = streamChatCompletions(requestPayload, abortSignal);
 
             for await (const chunk of stream) {
               // Handle tool status events
@@ -1777,8 +1934,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                         : "openaiCodeExecContainerId";
                     void updateStoredChatThreadEventually(resolvedThreadId, {
                       [field]: null,
-                    })
-                      .catch(() => {});
+                    }).catch(() => {});
                   }
                   continue;
                 }
@@ -1822,6 +1978,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                           size?: string;
                           quality?: string;
                           background?: string;
+                          prompt?: string;
                         };
                     const imageB64 = toolEvent.image_b64 as string | undefined;
                     if (
@@ -1843,6 +2000,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                         size: toolEvent.size as string | undefined,
                         quality: toolEvent.quality as string | undefined,
                         background: toolEvent.background as string | undefined,
+                        prompt: toolEvent.prompt as string | undefined,
                       };
                     } else if (imgIdx !== -1) {
                       const text = rawResult.slice(0, imgIdx);
@@ -1860,8 +2018,20 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
                     } else {
                       parsedResult = rawResult;
                     }
+                    const nextArgs =
+                      toolEvent.arguments &&
+                      typeof toolEvent.arguments === "object"
+                        ? (toolEvent.arguments as ToolCallMessagePart["args"])
+                        : undefined;
+                    const mergedArgs = nextArgs
+                      ? { ...(toolCallParts[idx].args ?? {}), ...nextArgs }
+                      : toolCallParts[idx].args;
                     toolCallParts[idx] = {
                       ...toolCallParts[idx],
+                      args: mergedArgs,
+                      argsText: mergedArgs
+                        ? JSON.stringify(mergedArgs)
+                        : toolCallParts[idx].argsText,
                       result: parsedResult,
                     };
                   }

studio/frontend/src/features/chat/provider-capabilities.ts

@@ -207,15 +207,21 @@ const OPENAI_CODE_EXECUTION_MODEL_PREFIXES = [
 
 /**
  * Strict check that a provider configuration points at OpenAI's
- * managed cloud (api.openai.com), as opposed to a custom OpenAI-compat
- * backend (ollama / llama.cpp / vLLM / generic "custom" preset). The
- * shell tool ONLY exists on OpenAI cloud; sending it to anything else
- * 400s the request. Mirror of the backend's
- * `is_openai_cloud = "api.openai.com" in self.base_url` guard.
+ * managed cloud (api.openai.com) or Azure OpenAI Foundry
+ * (*.openai.azure.com), as opposed to a custom OpenAI-compat backend
+ * (ollama / llama.cpp / vLLM / generic "custom" preset). The shell and
+ * image-generation tools only exist on cloud backends; sending them to
+ * anything else 400s the request. Mirror of the backend's
+ * `_is_openai_family_cloud` host check.
  */
 function isOpenAICloudBaseUrl(baseUrl: string | null | undefined): boolean {
   if (!baseUrl) return true; // No override → uses the default openai.com base.
-  return baseUrl.trim().toLowerCase().includes("api.openai.com");
+  try {
+    const host = new URL(baseUrl).hostname.toLowerCase();
+    return host === "api.openai.com" || host.endsWith(".openai.azure.com");
+  } catch {
+    return false;
+  }
 }
 
 export function providerSupportsBuiltinCodeExecution(

studio/frontend/src/features/chat/stores/chat-runtime-store.ts

@@ -64,6 +64,12 @@ function saveLastExternalCheckpoint(value: string | null): void {
 }
 
 export type ReasoningStyle = "enable_thinking" | "reasoning_effort";
+export type PendingImageEditReference = {
+  threadId: string | null;
+  openaiImageGenerationCallId: string;
+  openaiResponseId?: string;
+  openaiReasoningItem?: unknown;
+};
 export type ReasoningEffort =
   | "none"
   | "minimal"
@@ -300,6 +306,7 @@ type ChatRuntimeStore = {
   settingsPanelOpen: boolean;
   pendingAudioBase64: string | null;
   pendingAudioName: string | null;
+  pendingImageEditReference: PendingImageEditReference | null;
   contextUsage: {
     promptTokens: number;
     completionTokens: number;
@@ -353,6 +360,10 @@ type ChatRuntimeStore = {
   setChatTemplateOverride: (template: string | null) => void;
   setPendingAudio: (base64: string, name: string) => void;
   clearPendingAudio: () => void;
+  setPendingImageEditReference: (
+    reference: PendingImageEditReference | null,
+  ) => void;
+  clearPendingImageEditReference: () => void;
   setContextUsage: (usage: ChatRuntimeStore["contextUsage"]) => void;
 };
 
@@ -607,6 +618,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
   settingsPanelOpen: false,
   pendingAudioBase64: null,
   pendingAudioName: null,
+  pendingImageEditReference: null,
   contextUsage: null,
   modelLoading: false,
   activeNativePathToken: null,
@@ -793,6 +805,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
       defaultChatTemplate: null,
       chatTemplateOverride: null,
       loadedChatTemplateOverride: null,
+      pendingImageEditReference: null,
     }));
   },
   setReasoningEnabled: (reasoningEnabled, options) =>
@@ -884,5 +897,9 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
     set({ pendingAudioBase64: base64, pendingAudioName: name }),
   clearPendingAudio: () =>
     set({ pendingAudioBase64: null, pendingAudioName: null }),
+  setPendingImageEditReference: (pendingImageEditReference) =>
+    set({ pendingImageEditReference }),
+  clearPendingImageEditReference: () =>
+    set({ pendingImageEditReference: null }),
   setContextUsage: (contextUsage) => set({ contextUsage }),
 }));

studio/frontend/src/features/chat/types/api.ts

@@ -192,12 +192,31 @@ export interface AudioGenerationResponse {
   }>;
 }
 
-export type OpenAIMessageContent =
-  | string
-  | Array<
-      | { type: "text"; text: string }
-      | { type: "image_url"; image_url: { url: string } }
-    >;
+export type OpenAIReasoningSummaryPart = {
+  type: "summary_text";
+  text: string;
+};
+
+export type OpenAIReasoningContentPart = {
+  type: "reasoning";
+  id: string;
+  summary: OpenAIReasoningSummaryPart[];
+  status?: "in_progress" | "completed" | "incomplete";
+};
+
+export type OpenAIImageGenerationCallContentPart = {
+  type: "image_generation_call";
+  id: string;
+  response_id?: string;
+};
+
+export type OpenAIMessageContentPart =
+  | { type: "text"; text: string }
+  | { type: "image_url"; image_url: { url: string } }
+  | OpenAIReasoningContentPart
+  | OpenAIImageGenerationCallContentPart;
+
+export type OpenAIMessageContent = string | OpenAIMessageContentPart[];
 
 export interface OpenAIChatMessage {
   role: "system" | "user" | "assistant";

studio/frontend/src/index.css

@@ -1188,6 +1188,53 @@
 	border-color: var(--border) !important;
 }
 
+.generated-image-loading-card {
+	position: relative;
+	overflow: hidden;
+	contain: paint;
+}
+
+.generated-image-loading-wave {
+	position: relative;
+	display: grid;
+	grid-template-columns: repeat(8, minmax(0, 1fr));
+	gap: 14px;
+	width: min(66%, 18rem);
+	padding: 1.5rem;
+	border-radius: 1.5rem;
+}
+
+.generated-image-loading-dot {
+	width: 7px;
+	height: 7px;
+	border-radius: 9999px;
+	background: color-mix(in oklch, var(--muted-foreground) 82%, var(--primary));
+	opacity: 0.12;
+	transform: translate3d(0, 4px, 0) scale(0.72);
+	animation: generated-image-dot-wave 1850ms var(--ease-out-quart) infinite;
+	animation-delay: calc((var(--dot-row) * 72ms) + (var(--dot-col) * 72ms));
+	will-change: transform, opacity;
+}
+
+@keyframes generated-image-dot-wave {
+	0%,
+	22%,
+	100% {
+		opacity: 0.1;
+		transform: translate3d(0, 4px, 0) scale(0.72);
+	}
+
+	46% {
+		opacity: 0.46;
+		transform: translate3d(0, -3px, 0) scale(0.96);
+	}
+
+	66% {
+		opacity: 0.2;
+		transform: translate3d(0, 0, 0) scale(0.82);
+	}
+}
+
 /*
  * prefers-reduced-motion: honour the OS-level "reduce motion" preference.
  * Tailwind animate-in/out, Radix open/close transforms, infinite shine/pulse