fix

README.md

@@ -6,7 +6,7 @@ colorTo: red
 sdk: static
 pinned: true
 app_build_command: npm run build
-app_file: build/index.html
+app_file: dist/index.html
 license: apache-2.0
 short_description: Use MCP and WebGPU-based LLMs with tool calling
 ---

package.json

@@ -10,12 +10,14 @@
     "preview": "vite preview"
   },
   "dependencies": {
-    "@huggingface/transformers": "^3.7.1",
+    "@huggingface/transformers": "^3.7.5",
     "@modelcontextprotocol/sdk": "^1.17.3",
     "@monaco-editor/react": "^4.7.0",
     "@tailwindcss/vite": "^4.1.11",
+    "dompurify": "^3.2.7",
     "idb": "^8.0.3",
     "lucide-react": "^0.535.0",
+    "marked": "^16.3.0",
     "react": "^19.1.0",
     "react-dom": "^19.1.0",
     "react-router-dom": "^7.8.0",
@@ -23,6 +25,8 @@
   },
   "devDependencies": {
     "@eslint/js": "^9.30.1",
+    "@tailwindcss/typography": "^0.5.19",
+    "@types/dompurify": "^3.0.5",
     "@types/react": "^19.1.8",
     "@types/react-dom": "^19.1.6",
     "@vitejs/plugin-react": "^4.6.0",

src/App.tsx

@@ -15,7 +15,10 @@ import {
   X,
   PanelRightClose,
   PanelRightOpen,
+  StopCircle,
 } from "lucide-react";
+import { marked } from "marked";
+import DOMPurify from "dompurify";
 import { useLLM } from "./hooks/useLLM";
 import { useMCP } from "./hooks/useMCP";
 
@@ -29,7 +32,6 @@ import {
   extractToolCallContent,
   mapArgsToNamedParams,
   getErrorMessage,
-  isMobileOrTablet,
 } from "./utils";
 
 import { DEFAULT_SYSTEM_PROMPT } from "./constants/systemPrompt";
@@ -80,6 +82,10 @@ async function getDB(): Promise<IDBPDatabase> {
   });
 }
 
+function renderMarkdown(text: string): string {
+  return DOMPurify.sanitize(marked.parse(text) as string);
+}
+
 const App: React.FC = () => {
   const [systemPrompt, setSystemPrompt] = useState<string>(
     DEFAULT_SYSTEM_PROMPT
@@ -91,14 +97,14 @@ const App: React.FC = () => {
   const [tools, setTools] = useState<Tool[]>([]);
   const [input, setInput] = useState<string>("");
   const [isGenerating, setIsGenerating] = useState<boolean>(false);
-  const isMobile = useMemo(isMobileOrTablet, []);
+  // const isMobile = useMemo(isMobileOrTablet, []);
   const [selectedModelId, setSelectedModelId] = useState<string>(
-    isMobile ? "350M" : "1.2B"
+    "onnx-community/granite-4.0-micro-ONNX-web"
   );
   const [isModelDropdownOpen, setIsModelDropdownOpen] =
     useState<boolean>(false);
   const [isMCPManagerOpen, setIsMCPManagerOpen] = useState<boolean>(false);
-  const [isToolsPanelVisible, setIsToolsPanelVisible] = useState<boolean>(true);
+  const [isToolsPanelVisible, setIsToolsPanelVisible] = useState<boolean>(false);
   const chatContainerRef = useRef<HTMLDivElement>(null);
   const debounceTimers = useRef<Record<number, NodeJS.Timeout>>({});
   const toolsContainerRef = useRef<HTMLDivElement>(null);
@@ -111,6 +117,9 @@ const App: React.FC = () => {
     loadModel,
     generateResponse,
     clearPastKeyValues,
+    interruptGeneration,
+    tokensPerSecond,
+    numTokens,
   } = useLLM(selectedModelId);
 
   // MCP integration
@@ -708,14 +717,50 @@ const App: React.FC = () => {
             >
               {messages.length === 0 && isReady ? (
                 <ExamplePrompts
-                  examples={tools
-                    .filter((tool) => tool.enabled)
-                    .map((tool) => ({
-                      icon: "🛠️",
-                      displayText: tool.name,
-                      messageText: `${tool.name}()`,
-                    }))
-                    .filter((ex) => ex.displayText)}
+                  examples={(() => {
+                    const enabledTools = tools.filter((tool) => tool.enabled);
+
+                    // Group tools by server (MCP tools have mcpServerId in their code)
+                    const toolsByServer = enabledTools.reduce((acc, tool) => {
+                      const mcpServerMatch = tool.code?.match(/mcpServerId: "([^"]+)"/);
+                      const serverId = mcpServerMatch ? mcpServerMatch[1] : 'local';
+                      if (!acc[serverId]) acc[serverId] = [];
+                      acc[serverId].push(tool);
+                      return acc;
+                    }, {} as Record<string, typeof enabledTools>);
+
+                    // Pick one tool from each server (up to 3 servers)
+                    const serverIds = Object.keys(toolsByServer).slice(0, 3);
+                    const selectedTools = serverIds.map(serverId => {
+                      const serverTools = toolsByServer[serverId];
+                      return serverTools[Math.floor(Math.random() * serverTools.length)];
+                    });
+
+                    return selectedTools.map((tool) => {
+                      const schema = generateSchemaFromCode(tool.code);
+                      const description = schema.description || tool.name;
+
+                      // Create a cleaner natural language prompt
+                      let displayText = description;
+                      if (description !== tool.name) {
+                        // If there's a description, make it conversational
+                        displayText = description.charAt(0).toUpperCase() + description.slice(1);
+                        if (!displayText.endsWith('?') && !displayText.endsWith('.')) {
+                          displayText += '?';
+                        }
+                      } else {
+                        // Fallback to tool name in a readable format
+                        displayText = tool.name.replace(/_/g, ' ');
+                        displayText = displayText.charAt(0).toUpperCase() + displayText.slice(1);
+                      }
+
+                      return {
+                        icon: "🛠️",
+                        displayText,
+                        messageText: displayText,
+                      };
+                    });
+                  })()}
                   onExampleClick={handleExampleClick}
                 />
               ) : (
@@ -733,9 +778,9 @@ const App: React.FC = () => {
                       </div>
                     );
                   } else if (msg.role === "assistant") {
-                    const isToolCall = msg.content.includes(
-                      "<|tool_call_start|>"
-                    );
+                    const isToolCall =
+                      msg.content.includes("<|tool_call_start|>") ||
+                      msg.content.includes("<tool_call>");
 
                     if (isToolCall) {
                       const nextMessage = messages[index + 1];
@@ -762,9 +807,26 @@ const App: React.FC = () => {
                     return (
                       <div key={key} className="flex justify-start">
                         <div className="p-3 rounded-lg max-w-md bg-gray-700">
-                          <p className="text-sm whitespace-pre-wrap">
-                            {msg.content}
-                          </p>
+                          {msg.content.length > 0 ? (
+                            <div
+                              className="text-sm prose prose-invert prose-sm max-w-none"
+                              dangerouslySetInnerHTML={{
+                                __html: renderMarkdown(msg.content),
+                              }}
+                            />
+                          ) : (
+                            <div className="flex items-center gap-1 h-6">
+                              <span className="w-2 h-2 bg-gray-400 rounded-full animate-pulse"></span>
+                              <span
+                                className="w-2 h-2 bg-gray-400 rounded-full animate-pulse"
+                                style={{ animationDelay: "0.2s" }}
+                              ></span>
+                              <span
+                                className="w-2 h-2 bg-gray-400 rounded-full animate-pulse"
+                                style={{ animationDelay: "0.4s" }}
+                              ></span>
+                            </div>
+                          )}
                         </div>
                       </div>
                     );
@@ -806,33 +868,56 @@ const App: React.FC = () => {
               )}
             </div>
 
-            <div className="flex">
-              <input
-                ref={inputRef}
-                type="text"
-                value={input}
-                onChange={(e) => setInput(e.target.value)}
-                onKeyDown={(e) =>
-                  e.key === "Enter" &&
-                  !isGenerating &&
-                  isReady &&
-                  handleSendMessage()
-                }
-                disabled={isGenerating || !isReady}
-                className="flex-grow bg-gray-700 rounded-l-lg p-3 focus:outline-none focus:ring-2 focus:ring-indigo-500 disabled:opacity-50"
-                placeholder={
-                  isReady
-                    ? "Type your message here..."
-                    : "Load model first to enable chat"
-                }
-              />
-              <button
-                onClick={handleSendMessage}
-                disabled={isGenerating || !isReady}
-                className="bg-indigo-600 hover:bg-indigo-700 disabled:bg-gray-600 disabled:cursor-not-allowed text-white font-bold p-3 rounded-r-lg transition-colors"
-              >
-                <Play size={20} />
-              </button>
+            <div>
+              {/* TPS Display */}
+              {isGenerating && tokensPerSecond !== null && (
+                <div className="mb-2 text-sm text-gray-400 flex items-center gap-2">
+                  <span>
+                    {tokensPerSecond.toFixed(1)} tokens/sec
+                  </span>
+                  <span>•</span>
+                  <span>{numTokens} tokens</span>
+                </div>
+              )}
+
+              <div className="flex gap-2">
+                <input
+                  ref={inputRef}
+                  type="text"
+                  value={input}
+                  onChange={(e) => setInput(e.target.value)}
+                  onKeyDown={(e) =>
+                    e.key === "Enter" &&
+                    !isGenerating &&
+                    isReady &&
+                    handleSendMessage()
+                  }
+                  disabled={isGenerating || !isReady}
+                  className="flex-grow bg-gray-700 rounded-lg p-3 focus:outline-none focus:ring-2 focus:ring-indigo-500 disabled:opacity-50"
+                  placeholder={
+                    isReady
+                      ? "Type your message here..."
+                      : "Load model first to enable chat"
+                  }
+                />
+                {isGenerating ? (
+                  <button
+                    onClick={interruptGeneration}
+                    className="bg-red-600 hover:bg-red-700 text-white font-bold px-4 py-3 rounded-lg transition-colors flex items-center gap-2"
+                  >
+                    <StopCircle size={20} />
+                    <span className="hidden sm:inline">Stop</span>
+                  </button>
+                ) : (
+                  <button
+                    onClick={handleSendMessage}
+                    disabled={!isReady}
+                    className="bg-indigo-600 hover:bg-indigo-700 disabled:bg-gray-600 disabled:cursor-not-allowed text-white font-bold px-4 py-3 rounded-lg transition-colors"
+                  >
+                    <Play size={20} />
+                  </button>
+                )}
+              </div>
             </div>
           </div>
 

src/components/ExamplePrompts.tsx

@@ -30,17 +30,17 @@ const ExamplePrompts: React.FC<ExamplePromptsProps> = ({
         <p className="text-sm text-gray-500">Click one to get started</p>
       </div>
 
-      <div className="grid grid-cols-1 sm:grid-cols-2 gap-3 max-w-2xl w-full px-4">
+      <div className="grid grid-cols-1 sm:grid-cols-2 lg:grid-cols-3 gap-3 max-w-4xl w-full px-4">
         {dynamicExamples.map((example, index) => (
           <button
             key={index}
             onClick={() => onExampleClick(example.messageText)}
-            className="flex items-center gap-3 p-4 bg-gray-700 hover:bg-gray-600 rounded-lg transition-colors text-left group cursor-pointer"
+            className="flex items-start gap-3 p-4 bg-gray-700 hover:bg-gray-600 rounded-lg transition-colors text-left group cursor-pointer"
           >
             <span className="text-xl flex-shrink-0 group-hover:scale-110 transition-transform">
               {example.icon}
             </span>
-            <span className="text-sm text-gray-200 group-hover:text-white transition-colors">
+            <span className="text-sm text-gray-200 group-hover:text-white transition-colors break-words line-clamp-3">
               {example.displayText}
             </span>
           </button>

src/components/LoadingScreen.tsx

@@ -359,15 +359,12 @@ export const LoadingScreen = ({
             to get started.
           </p>
 
-          <div className="relative">
-            <div
-              ref={wrapperRef} // anchor for dropdown centering
-              className="flex rounded-2xl shadow-2xl overflow-hidden"
-            >
+          <div className="relative inline-block">
+            <div className="flex rounded-2xl shadow-2xl overflow-hidden">
               <button
                 onClick={isLoading ? undefined : loadSelectedModel}
                 disabled={isLoading}
-                className={`flex items-center justify-center font-bold transition-all text-lg flex-1 ${
+                className={`flex items-center justify-center font-bold transition-all text-lg ${
                   isLoading
                     ? "bg-gray-700 text-gray-400 cursor-not-allowed"
                     : "bg-gradient-to-r from-indigo-600 to-purple-600 hover:from-indigo-500 hover:to-purple-500 text-white shadow-lg hover:shadow-xl transform hover:scale-[1.01] active:scale-[0.99]"
@@ -408,20 +405,8 @@ export const LoadingScreen = ({
                     );
                   }
                 }}
-                onKeyDown={(e) => {
-                  if (isLoading) return;
-                  if (
-                    e.key === " " ||
-                    e.key === "Enter" ||
-                    e.key === "ArrowDown"
-                  ) {
-                    e.preventDefault();
-                    if (!isModelDropdownOpen) setIsModelDropdownOpen(true);
-                  }
-                }}
                 aria-haspopup="menu"
                 aria-expanded={isModelDropdownOpen}
-                aria-controls="model-dropdown"
                 aria-label="Select model"
                 className={`px-4 py-4 border-l border-white/20 transition-all ${
                   isLoading
@@ -439,17 +424,28 @@ export const LoadingScreen = ({
               </button>
             </div>
 
-            {/* Dropdown (Portal) */}
-            {isModelDropdownOpen &&
-              typeof document !== "undefined" &&
-              ReactDOM.createPortal(
+            {/* Dropdown - render in portal to avoid z-index issues */}
+            {isModelDropdownOpen && typeof document !== "undefined" && ReactDOM.createPortal(
+              <>
+                {/* Backdrop */}
+                <div
+                  className="fixed inset-0 z-[999]"
+                  onClick={() => setIsModelDropdownOpen(false)}
+                />
+                {/* Dropdown */}
                 <div
-                  id="model-dropdown"
                   ref={dropdownRef}
-                  style={portalStyle}
                   role="menu"
                   aria-label="Model options"
-                  className="bg-gray-800/95 border border-gray-600/50 rounded-2xl shadow-2xl overflow-hidden animate-in slide-in-from-top-2 duration-200 dropdown-z30"
+                  style={{
+                    position: 'fixed',
+                    bottom: dropdownBtnRef.current ? `${window.innerHeight - dropdownBtnRef.current.getBoundingClientRect().top}px` : 'auto',
+                    left: dropdownBtnRef.current ? `${dropdownBtnRef.current.getBoundingClientRect().left}px` : 'auto',
+                    width: dropdownBtnRef.current ? `${dropdownBtnRef.current.getBoundingClientRect().width + 200}px` : '320px',
+                    zIndex: 1000,
+                  }}
+                  className="bg-gray-900 border-2 border-gray-600 rounded-2xl shadow-2xl overflow-y-auto max-h-[300px] min-w-[320px]"
+                  onClick={(e) => e.stopPropagation()}
                 >
                   {MODEL_OPTIONS.map((option, index) => {
                     const selected = selectedModelId === option.id;
@@ -460,43 +456,47 @@ export const LoadingScreen = ({
                         role="menuitem"
                         aria-checked={selected}
                         onMouseEnter={() => setActiveIndex(index)}
-                        onClick={() => {
+                        onClick={(e) => {
+                          e.stopPropagation();
                           handleModelSelect(option.id);
                           setIsModelDropdownOpen(false);
                           dropdownBtnRef.current?.focus();
                         }}
-                        className={`w-full px-6 py-4 text-left transition-all duration-200 relative group outline-none ${
+                        className={`w-full px-6 py-4 text-left transition-all duration-150 relative outline-none border-b border-gray-700/50 last:border-b-0 cursor-pointer ${
                           selected
-                            ? "bg-gradient-to-r from-indigo-600/50 to-purple-600/50 text-white border-l-4 border-indigo-400"
-                            : "text-gray-200 hover:bg-white/10 hover:text-white"
+                            ? "bg-indigo-600 text-white hover:bg-indigo-500"
+                            : "bg-gray-800 text-gray-300 hover:bg-gray-700 hover:text-white active:bg-gray-600"
                         } ${index === 0 ? "rounded-t-2xl" : ""} ${
                           index === MODEL_OPTIONS.length - 1
                             ? "rounded-b-2xl"
                             : ""
-                        } ${isActive && !selected ? "bg-white/5" : ""}`}
+                        } ${isActive && !selected ? "bg-gray-700" : ""}`}
                       >
                         <div className="flex items-center justify-between">
-                          <div>
-                            <div className="font-semibold text-lg">
+                          <div className="flex-1">
+                            <div className={`font-semibold text-base mb-1 ${selected ? "text-white" : "text-gray-100"}`}>
                               {option.label}
                             </div>
-                            <div className="text-sm text-gray-400 mt-1">
+                            <div className={`text-sm ${selected ? "text-indigo-200" : "text-gray-500"}`}>
                               {option.size}
                             </div>
                           </div>
                           {selected && (
-                            <div className="w-2 h-2 bg-indigo-400 rounded-full" />
+                            <div className="flex items-center gap-2 ml-4">
+                              <span className="text-xs font-medium text-indigo-200">Selected</span>
+                              <svg className="w-5 h-5 text-white" fill="currentColor" viewBox="0 0 20 20">
+                                <path fillRule="evenodd" d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z" clipRule="evenodd" />
+                              </svg>
+                            </div>
                           )}
                         </div>
-                        {!selected && (
-                          <div className="absolute inset-0 bg-gradient-to-r from-indigo-500/10 to-purple-500/10 opacity-0 group-hover:opacity-100 transition-opacity rounded-2xl" />
-                        )}
                       </button>
                     );
                   })}
-                </div>,
-                document.body
-              )}
+                </div>
+              </>,
+              document.body
+            )}
           </div>
         </div>
 
@@ -517,22 +517,6 @@ export const LoadingScreen = ({
         )}
       </div>
 
-      {/* Click-away fallback for touch devices */}
-      {isModelDropdownOpen && (
-        <div
-          className="fixed inset-0 z-40 bg-black/20"
-          onClick={(e) => {
-            const target = e.target as Node;
-            if (
-              dropdownRef.current &&
-              !dropdownRef.current.contains(target) &&
-              !dropdownBtnRef.current?.contains(target)
-            ) {
-              setIsModelDropdownOpen(false);
-            }
-          }}
-        />
-      )}
     </div>
   );
 };

src/config/constants.ts

@@ -4,9 +4,9 @@
 
 // MCP Client Configuration
 export const MCP_CLIENT_CONFIG = {
-  NAME: "LFM2-WebGPU",
+  NAME: "WebGPU-MCP",
   VERSION: "1.0.0",
-  TEST_CLIENT_NAME: "LFM2-WebGPU-Test",
+  TEST_CLIENT_NAME: "WebGPU-MCP-Test",
 } as const;
 
 // Storage Keys

src/constants/models.ts

@@ -1,5 +1,22 @@
 export const MODEL_OPTIONS = [
-  { id: "350M", label: "LFM2-350M", size: "350M parameters (312 MB)" },
-  { id: "700M", label: "LFM2-700M", size: "700M parameters (579 MB)" },
-  { id: "1.2B", label: "LFM2-1.2B", size: "1.2B parameters (868 MB)" },
+  {
+    id: "onnx-community/granite-4.0-micro-ONNX-web",
+    label: "Granite 4.0 Micro",
+    size: "400M parameters"
+  },
+  {
+    id: "onnx-community/LFM2-350M-ONNX",
+    label: "LFM2-350M",
+    size: "350M parameters (312 MB)"
+  },
+  {
+    id: "onnx-community/LFM2-700M-ONNX",
+    label: "LFM2-700M",
+    size: "700M parameters (579 MB)"
+  },
+  {
+    id: "onnx-community/LFM2-1.2B-ONNX",
+    label: "LFM2-1.2B",
+    size: "1.2B parameters (868 MB)"
+  },
 ];

src/hooks/useLLM.ts

@@ -10,6 +10,8 @@ interface LLMState {
   isReady: boolean;
   error: string | null;
   progress: number;
+  tokensPerSecond: number | null;
+  numTokens: number;
 }
 
 interface LLMInstance {
@@ -30,6 +32,8 @@ export const useLLM = (modelId?: string) => {
     isReady: false,
     error: null,
     progress: 0,
+    tokensPerSecond: null,
+    numTokens: 0,
   });
 
   const instanceRef = useRef<LLMInstance | null>(null);
@@ -37,13 +41,14 @@ export const useLLM = (modelId?: string) => {
 
   const abortControllerRef = useRef<AbortController | null>(null);
   const pastKeyValuesRef = useRef<any>(null);
+  const generationAbortControllerRef = useRef<AbortController | null>(null);
 
   const loadModel = useCallback(async () => {
     if (!modelId) {
       throw new Error("Model ID is required");
     }
 
-    const MODEL_ID = `onnx-community/LFM2-${modelId}-ONNX`;
+    const MODEL_ID = modelId;
 
     if (!moduleCache[modelId]) {
       moduleCache[modelId] = {
@@ -99,7 +104,7 @@ export const useLLM = (modelId?: string) => {
             progress.file.endsWith(".onnx_data")
           ) {
             const percentage = Math.round(
-              (progress.loaded / progress.total) * 100,
+              (progress.loaded / progress.total) * 100
             );
             setState((prev) => ({ ...prev, progress: percentage }));
           }
@@ -115,6 +120,20 @@ export const useLLM = (modelId?: string) => {
           progress_callback: progressCallback,
         });
 
+        // Pre-warm the model with a dummy input for shader compilation
+        console.log("Pre-warming model...");
+        const dummyInput = tokenizer("Hello", {
+          return_tensors: "pt",
+          padding: false,
+          truncation: false,
+        });
+        await model.generate({
+          ...dummyInput,
+          max_new_tokens: 1,
+          do_sample: false,
+        });
+        console.log("Model pre-warmed");
+
         const instance = { model, tokenizer };
         instanceRef.current = instance;
         cache.instance = instance;
@@ -150,7 +169,7 @@ export const useLLM = (modelId?: string) => {
     async (
       messages: Array<{ role: string; content: string }>,
       tools: Array<any>,
-      onToken?: (token: string) => void,
+      onToken?: (token: string) => void
     ): Promise<string> => {
       const instance = instanceRef.current;
       if (!instance) {
@@ -159,6 +178,9 @@ export const useLLM = (modelId?: string) => {
 
       const { model, tokenizer } = instance;
 
+      // Create abort controller for this generation
+      generationAbortControllerRef.current = new AbortController();
+
       // Apply chat template with tools
       const input = tokenizer.apply_chat_template(messages, {
         tools,
@@ -166,39 +188,62 @@ export const useLLM = (modelId?: string) => {
         return_dict: true,
       });
 
+      // Track tokens and timing
+      const startTime = performance.now();
+      let tokenCount = 0;
+
       const streamer = onToken
         ? new TextStreamer(tokenizer, {
             skip_prompt: true,
             skip_special_tokens: false,
             callback_function: (token: string) => {
+              tokenCount++;
+              const elapsed = (performance.now() - startTime) / 1000;
+              const tps = tokenCount / elapsed;
+              setState((prev) => ({
+                ...prev,
+                tokensPerSecond: tps,
+                numTokens: tokenCount,
+              }));
               onToken(token);
             },
           })
         : undefined;
 
-      // Generate the response
-      const { sequences, past_key_values } = await model.generate({
-        ...input,
-        past_key_values: pastKeyValuesRef.current,
-        max_new_tokens: 512,
-        do_sample: false,
-        streamer,
-        return_dict_in_generate: true,
-      });
-      pastKeyValuesRef.current = past_key_values;
+      try {
+        // Generate the response
+        const { sequences, past_key_values } = await model.generate({
+          ...input,
+          past_key_values: pastKeyValuesRef.current,
+          max_new_tokens: 1024,
+          do_sample: false,
+          streamer,
+          return_dict_in_generate: true,
+        });
+        pastKeyValuesRef.current = past_key_values;
 
-      // Decode the generated text with special tokens preserved (except final <|im_end|>) for tool call detection
-      const response = tokenizer
-        .batch_decode(sequences.slice(null, [input.input_ids.dims[1], null]), {
-          skip_special_tokens: false,
-        })[0]
-        .replace(/<\|im_end\|>$/, "");
+        // Decode the generated text with special tokens preserved (except end tokens) for tool call detection
+        const response = tokenizer
+          .batch_decode(sequences.slice(null, [input.input_ids.dims[1], null]), {
+            skip_special_tokens: false,
+          })[0]
+          .replace(/<\|im_end\|>$/, "")
+          .replace(/<\|end_of_text\|>$/, "");
 
-      return response;
+        return response;
+      } finally {
+        generationAbortControllerRef.current = null;
+      }
     },
-    [],
+    []
   );
 
+  const interruptGeneration = useCallback(() => {
+    if (generationAbortControllerRef.current) {
+      generationAbortControllerRef.current.abort();
+    }
+  }, []);
+
   const clearPastKeyValues = useCallback(() => {
     pastKeyValuesRef.current = null;
   }, []);
@@ -230,5 +275,6 @@ export const useLLM = (modelId?: string) => {
     generateResponse,
     clearPastKeyValues,
     cleanup,
+    interruptGeneration,
   };
 };

src/index.css

@@ -1 +1,2 @@
 @import "tailwindcss";
+@plugin "@tailwindcss/typography";

src/utils.ts

@@ -71,8 +71,17 @@ export const extractPythonicCalls = (toolCallContent: string): string[] => {
   try {
     const cleanContent = toolCallContent.trim();
 
+    // Try to parse as Granite format (JSON object with name and arguments)
     try {
       const parsed = JSON.parse(cleanContent);
+      if (parsed && typeof parsed === 'object' && parsed.name) {
+        // Convert Granite JSON format to Pythonic format
+        const args = parsed.arguments || {};
+        const argPairs = Object.entries(args).map(([key, value]) =>
+          `${key}=${JSON.stringify(value)}`
+        );
+        return [`${parsed.name}(${argPairs.join(', ')})`];
+      }
       if (Array.isArray(parsed)) {
         return parsed;
       }
@@ -305,12 +314,25 @@ export const generateSchemaFromCode = (code: string): Schema => {
 
 /**
  * Extracts tool call content from a string using the tool call markers.
+ * Supports both LFM2 format (<|tool_call_start|>...<|tool_call_end|>)
+ * and Granite format (<tool_call>...</tool_call>)
  */
 export const extractToolCallContent = (content: string): string | null => {
-  const toolCallMatch = content.match(
+  // Try LFM2 format first
+  const lfm2Match = content.match(
     /<\|tool_call_start\|>(.*?)<\|tool_call_end\|>/s,
   );
-  return toolCallMatch ? toolCallMatch[1].trim() : null;
+  if (lfm2Match) {
+    return lfm2Match[1].trim();
+  }
+
+  // Try Granite format (XML-style)
+  const graniteMatch = content.match(/<tool_call>(.*?)<\/tool_call>/s);
+  if (graniteMatch) {
+    return graniteMatch[1].trim();
+  }
+
+  return null;
 };
 
 /**