Add image lightbox

## Motivation

Kimi produces its own tool id. It gets confused when we generate our own
id.

## Changes

Add id to tool call item and parse Kimi id properly.

## Test Plan

### Manual Testing
<img width="3198" height="522" alt="image"
src="https://github.com/user-attachments/assets/d71ec2be-7f57-49dc-a569-d304cc430f4d"
/>

Long running Kimi K2.5 cluster querying itself through OpenCode running
on the same Kimi K2.5 instance.

dashboard/src/lib/components/ChatMessages.svelte

@@ -13,6 +13,7 @@
   import type { MessageAttachment } from "$lib/stores/app.svelte";
   import MarkdownContent from "./MarkdownContent.svelte";
   import TokenHeatmap from "./TokenHeatmap.svelte";
+  import ImageLightbox from "./ImageLightbox.svelte";
 
   interface Props {
     class?: string;
@@ -101,6 +102,9 @@
   let copiedMessageId = $state<string | null>(null);
   let expandedThinkingMessageIds = $state<Set<string>>(new Set());
 
+  // Lightbox state
+  let expandedImageSrc = $state<string | null>(null);
+
   // Uncertainty heatmap toggle
   let heatmapMessageIds = $state<Set<string>>(new Set());
 
@@ -389,10 +393,15 @@
                         class="flex items-center gap-2 bg-exo-dark-gray/60 border border-exo-yellow/20 rounded px-2 py-1 text-xs font-mono"
                       >
                         {#if attachment.type === "image" && attachment.preview}
+                          <!-- svelte-ignore a11y_no_noninteractive_element_interactions, a11y_click_events_have_key_events -->
                           <img
                             src={attachment.preview}
                             alt={attachment.name}
-                            class="w-12 h-12 object-cover rounded border border-exo-yellow/20"
+                            class="w-12 h-12 object-cover rounded border border-exo-yellow/20 cursor-pointer hover:border-exo-yellow/50 transition-colors"
+                            onclick={() => {
+                              if (attachment.preview)
+                                expandedImageSrc = attachment.preview;
+                            }}
                           />
                         {:else}
                           <span>{getAttachmentIcon(attachment)}</span>
@@ -466,15 +475,44 @@
                   <div class="mb-3">
                     {#each message.attachments.filter((a) => a.type === "generated-image") as attachment}
                       <div class="relative group/img inline-block">
+                        <!-- svelte-ignore a11y_no_noninteractive_element_interactions, a11y_click_events_have_key_events -->
                         <img
                           src={attachment.preview}
                           alt=""
-                          class="max-w-full max-h-[512px] rounded-lg border border-exo-yellow/20 shadow-lg shadow-black/20"
+                          class="max-w-full max-h-[512px] rounded-lg border border-exo-yellow/20 shadow-lg shadow-black/20 cursor-pointer"
+                          onclick={() => {
+                            if (attachment.preview)
+                              expandedImageSrc = attachment.preview;
+                          }}
                         />
                         <!-- Button overlay -->
                         <div
                           class="absolute top-2 right-2 flex gap-1 opacity-0 group-hover/img:opacity-100 transition-opacity"
                         >
+                          <!-- Expand button -->
+                          <button
+                            type="button"
+                            class="p-2 rounded-lg bg-exo-dark-gray/80 border border-exo-yellow/30 text-exo-yellow hover:bg-exo-dark-gray hover:border-exo-yellow/50 cursor-pointer"
+                            onclick={() => {
+                              if (attachment.preview)
+                                expandedImageSrc = attachment.preview;
+                            }}
+                            title="Expand image"
+                          >
+                            <svg
+                              class="w-4 h-4"
+                              fill="none"
+                              viewBox="0 0 24 24"
+                              stroke="currentColor"
+                              stroke-width="2"
+                            >
+                              <path
+                                stroke-linecap="round"
+                                stroke-linejoin="round"
+                                d="M4 8V4m0 0h4M4 4l5 5m11-1V4m0 0h-4m4 0l-5 5M4 16v4m0 0h4m-4 0l5-5m11 5l-5-5m5 5v-4m0 4h-4"
+                              />
+                            </svg>
+                          </button>
                           <!-- Edit button -->
                           <button
                             type="button"
@@ -789,3 +827,8 @@
     </button>
   {/if}
 </div>
+
+<ImageLightbox
+  src={expandedImageSrc}
+  onclose={() => (expandedImageSrc = null)}
+/>

dashboard/src/lib/components/ImageLightbox.svelte

@@ -0,0 +1,96 @@
+<script lang="ts">
+  import { fade, fly } from "svelte/transition";
+  import { cubicOut } from "svelte/easing";
+
+  interface Props {
+    src: string | null;
+    onclose: () => void;
+  }
+
+  let { src, onclose }: Props = $props();
+
+  function handleKeydown(e: KeyboardEvent) {
+    if (e.key === "Escape") {
+      onclose();
+    }
+  }
+
+  function extensionFromSrc(dataSrc: string): string {
+    const match = dataSrc.match(/^data:image\/(\w+)/);
+    if (match) return match[1] === "jpeg" ? "jpg" : match[1];
+    const urlMatch = dataSrc.match(/\.(\w+)(?:\?|$)/);
+    if (urlMatch) return urlMatch[1];
+    return "png";
+  }
+
+  function handleDownload(e: MouseEvent) {
+    e.stopPropagation();
+    if (!src) return;
+    const link = document.createElement("a");
+    link.href = src;
+    link.download = `image-${Date.now()}.${extensionFromSrc(src)}`;
+    link.click();
+  }
+
+  function handleClose(e: MouseEvent) {
+    e.stopPropagation();
+    onclose();
+  }
+</script>
+
+<svelte:window onkeydown={src ? handleKeydown : undefined} />
+
+{#if src}
+  <div
+    class="fixed inset-0 z-50 bg-black/90 backdrop-blur-sm flex items-center justify-center"
+    transition:fade={{ duration: 200 }}
+    onclick={onclose}
+    role="presentation"
+    onintrostart={() => (document.body.style.overflow = "hidden")}
+    onoutroend={() => (document.body.style.overflow = "")}
+  >
+    <div class="absolute top-4 right-4 flex gap-2 z-10">
+      <button
+        type="button"
+        class="p-2 rounded-lg bg-exo-dark-gray/80 border border-exo-yellow/30 text-exo-yellow hover:bg-exo-dark-gray hover:border-exo-yellow/50 cursor-pointer transition-colors"
+        onclick={handleDownload}
+        title="Download image"
+      >
+        <svg
+          class="w-5 h-5"
+          fill="none"
+          viewBox="0 0 24 24"
+          stroke="currentColor"
+          stroke-width="2"
+        >
+          <path
+            stroke-linecap="round"
+            stroke-linejoin="round"
+            d="M4 16v1a3 3 0 003 3h10a3 3 0 003-3v-1m-4-4l-4 4m0 0l-4-4m4 4V4"
+          />
+        </svg>
+      </button>
+      <button
+        type="button"
+        class="p-2 rounded-lg bg-exo-dark-gray/80 border border-exo-yellow/30 text-exo-yellow hover:bg-exo-dark-gray hover:border-exo-yellow/50 cursor-pointer transition-colors"
+        onclick={handleClose}
+        title="Close"
+      >
+        <svg class="w-5 h-5" viewBox="0 0 24 24" fill="currentColor">
+          <path
+            d="M19 6.41L17.59 5 12 10.59 6.41 5 5 6.41 10.59 12 5 17.59 6.41 19 12 13.41 17.59 19 19 17.59 13.41 12 19 6.41z"
+          />
+        </svg>
+      </button>
+    </div>
+
+    <!-- svelte-ignore a11y_no_noninteractive_element_interactions, a11y_click_events_have_key_events -->
+    <img
+      {src}
+      alt=""
+      class="max-w-[90vw] max-h-[90vh] object-contain rounded-lg shadow-2xl"
+      transition:fly={{ y: 20, duration: 300, easing: cubicOut }}
+      onclick={(e) => e.stopPropagation()}
+    />
+  </div>
+{/if}

src/exo/master/adapters/chat_completions.py

@@ -3,7 +3,6 @@
 import time
 from collections.abc import AsyncGenerator
 from typing import Any
-from uuid import uuid4
 
 from exo.shared.types.api import (
     ChatCompletionChoice,
@@ -141,7 +140,7 @@ async def generate_chat_stream(
         if isinstance(chunk, ToolCallChunk):
             tool_call_deltas = [
                 ToolCall(
-                    id=str(uuid4()),
+                    id=tool.id,
                     index=i,
                     function=tool,
                 )
@@ -207,7 +206,7 @@ async def collect_chat_response(
         if isinstance(chunk, ToolCallChunk):
             tool_calls.extend(
                 ToolCall(
-                    id=str(uuid4()),
+                    id=tool.id,
                     index=i,
                     function=tool,
                 )

src/exo/master/adapters/claude.py

@@ -3,7 +3,6 @@
 import json
 from collections.abc import AsyncGenerator
 from typing import Any
-from uuid import uuid4
 
 from exo.shared.types.api import FinishReason
 from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
@@ -179,7 +178,7 @@ async def collect_claude_response(
             for tool in chunk.tool_calls:
                 tool_use_blocks.append(
                     ClaudeToolUseBlock(
-                        id=f"toolu_{uuid4().hex[:24]}",
+                        id=f"toolu_{tool.id}",
                         name=tool.name,
                         input=json.loads(tool.arguments),  # pyright: ignore[reportAny]
                     )
@@ -264,7 +263,7 @@ async def generate_claude_stream(
 
             # Emit tool_use content blocks
             for tool in chunk.tool_calls:
-                tool_id = f"toolu_{uuid4().hex[:24]}"
+                tool_id = f"toolu_{tool.id}"
                 tool_input_json = tool.arguments
 
                 # content_block_start for tool_use

src/exo/master/adapters/responses.py

@@ -3,7 +3,6 @@
 from collections.abc import AsyncGenerator
 from itertools import count
 from typing import Any
-from uuid import uuid4
 
 from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
 from exo.shared.types.common import CommandId
@@ -140,8 +139,8 @@ async def collect_responses_response(
             for tool in chunk.tool_calls:
                 function_call_items.append(
                     ResponseFunctionCallItem(
-                        id=f"fc_{uuid4().hex[:24]}",
-                        call_id=f"call_{uuid4().hex[:24]}",
+                        id=f"fc_{tool.id}",
+                        call_id=f"call_{tool.id}",
                         name=tool.name,
                         arguments=tool.arguments,
                     )
@@ -246,8 +245,8 @@ async def generate_responses_stream(
         if isinstance(chunk, ToolCallChunk):
             last_stats = chunk.stats or last_stats
             for tool in chunk.tool_calls:
-                fc_id = f"fc_{uuid4().hex[:24]}"
-                call_id = f"call_{uuid4().hex[:24]}"
+                fc_id = f"fc_{tool.id}"
+                call_id = f"call_{tool.id}"
 
                 # response.output_item.added for function_call
                 fc_item = ResponseFunctionCallItem(

src/exo/shared/types/api.py

@@ -1,6 +1,7 @@
 import time
 from collections.abc import Generator
 from typing import Annotated, Any, Literal
+from uuid import uuid4
 
 from pydantic import BaseModel, Field, field_validator
 from pydantic_core import PydanticUseDefault
@@ -60,6 +61,7 @@ class ChatCompletionMessageText(BaseModel):
 
 
 class ToolCallItem(BaseModel):
+    id: str = Field(default_factory=lambda: str(uuid4()))
     name: str
     arguments: str
 

src/exo/worker/engines/mlx/generator/generate.py

@@ -298,6 +298,9 @@ def mlx_generate(
     )
     max_stop_len = max((len(s) for s in stop_sequences), default=0)
 
+    mx_barrier(group)
+    logger.info("Ready to prefill")
+
     # Prefill cache with all tokens except the last one
     prefill_tps, prefill_tokens, ssm_snapshots_list = prefill(
         model,

src/exo/worker/runner/runner.py

@@ -810,8 +810,9 @@ def patch_kimi_tokenizer(tokenizer: TokenizerWrapper):
 
     # kimi has a fixed function naming scheme, with a json formatted arg
     #   functions.multiply:0 <|tool_call_argument_begin|> {"a": 2, "b": 3}
+    #   Also needs to handle tools like call_0<|tool_call_argument_begin|>{"filePath": "..."}
     _func_name_regex = re.compile(
-        r"^\s*(.+):\d+\s*<\|tool_call_argument_begin\|>", re.DOTALL
+        r"^\s*(.+)[:](\d+)\s*<\|tool_call_argument_begin\|>", re.DOTALL
     )
     _func_arg_regex = re.compile(r"<\|tool_call_argument_begin\|>\s*(.*)\s*", re.DOTALL)
 
@@ -835,9 +836,10 @@ def patch_kimi_tokenizer(tokenizer: TokenizerWrapper):
         func_name_match = _func_name_regex.search(text)
         if func_name_match is None:
             raise ValueError(f"Could not parse function name from tool call: {text!r}")
-        func_name = func_name_match.group(1)
+        original_func_name = func_name_match.group(1)
+        tool_id = func_name_match.group(2)
         # strip off the `functions.` prefix, if it exists.
-        func_name = func_name[func_name.find(".") + 1 :]
+        func_name = original_func_name[original_func_name.find(".") + 1 :]
 
         func_args_match = _func_arg_regex.search(text)
         if func_args_match is None:
@@ -846,7 +848,11 @@ def patch_kimi_tokenizer(tokenizer: TokenizerWrapper):
         # the args should be valid json - no need to check against our tools to deserialize
         arg_dct = _deserialize(func_args)  # pyright: ignore[reportAny]
 
-        return dict(name=func_name, arguments=arg_dct)  # pyright: ignore[reportAny]
+        return dict(
+            id=f"{original_func_name}:{tool_id}",
+            name=func_name,
+            arguments=arg_dct,  # pyright: ignore[reportAny]
+        )
 
     tokenizer._tool_call_start = tool_call_start
     tokenizer._tool_call_end = tool_call_end
@@ -929,7 +935,13 @@ def _validate_single_tool(obj: dict[str, Any]) -> ToolCallItem:
         and ((args := obj.get("arguments")) is not None)
         and isinstance(name, str)
     ):
-        return ToolCallItem(name=name, arguments=json.dumps(args))
+        raw_id: object = obj.get("id")
+        extra = {"id": str(raw_id)} if raw_id is not None else {}
+        return ToolCallItem(
+            **extra,
+            name=name,
+            arguments=json.dumps(args),
+        )
     else:
         raise ValidationError