diff --git a/README.md b/README.md
index 4427ec6..15baa6f 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,6 @@
 # open-gsio
 [![Tests](https://github.com/geoffsee/open-gsio/actions/workflows/test.yml/badge.svg)](https://github.com/geoffsee/open-gsio/actions/workflows/test.yml)
 [![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT)
-![image](https://github.com/user-attachments/assets/a323d373-6241-4b76-b564-f0d080ff93f7)
 
 ### Stack:
 - vike
@@ -51,6 +50,18 @@ sed -i '' '/^OPENAI_API_ENDPOINT=/d' .dev.vars; echo 'OPENAI_API_ENDPOINT=http:/
 ### Restart open-gsio server so it uses the new variables
 bun run server:dev
 ~~~
+## Adding models for local inference
+~~~console
+MODEL_TO_ADD=mlx-community/gemma-3-4b-it-8bit
+# Chat completions endpoint
+curl http://localhost:10240/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model: ,
+    "messages": [{"role": "user", "content": "Hello"}]
+  }'
+~~~  
+
 History
 ---
 
diff --git a/package.json b/package.json
index 0ea62ab..8d9e804 100644
--- a/package.json
+++ b/package.json
@@ -19,9 +19,9 @@
     "tail:analytics-service": "wrangler tail -c workers/analytics/wrangler-analytics.toml",
     "tail:session-proxy": "wrangler tail -c workers/session-proxy/wrangler-session-proxy.toml --env production",
     "openai:local": "./scripts/start_inference_server.sh",
-    "test": "NODE_OPTIONS=--no-experimental-fetch vitest run",
-    "test:watch": "NODE_OPTIONS=--no-experimental-fetch vitest",
-    "test:coverage": "NODE_OPTIONS=--no-experimental-fetch vitest run --coverage"
+    "test": "vitest run",
+    "test:watch": "vitest",
+    "test:coverage": "vitest run --coverage"
   },
   "devDependencies": {
     "@anthropic-ai/sdk": "^0.32.1",
diff --git a/scripts/start_inference_server.sh b/scripts/start_inference_server.sh
index 111e57c..0f7d04b 100755
--- a/scripts/start_inference_server.sh
+++ b/scripts/start_inference_server.sh
@@ -5,4 +5,4 @@ SERVER_TYPE="mlx-omni-server"
 printf "Starting Inference Server: %s\n" ${SERVER_TYPE}
 
 
-mlx-omni-server
\ No newline at end of file
+mlx-omni-server --log-level debug
\ No newline at end of file
diff --git a/src/components/chat/messages/MessageEditorComponent.tsx b/src/components/chat/messages/MessageEditorComponent.tsx
index 11699ac..ff7f0d7 100644
--- a/src/components/chat/messages/MessageEditorComponent.tsx
+++ b/src/components/chat/messages/MessageEditorComponent.tsx
@@ -1,11 +1,10 @@
-import React, { KeyboardEvent, useState } from "react";
+import React, { KeyboardEvent, useEffect } from "react";
 import { Box, Flex, IconButton, Textarea } from "@chakra-ui/react";
 import { Check, X } from "lucide-react";
 import { observer } from "mobx-react-lite";
 import { Instance } from "mobx-state-tree";
 import Message from "../../../models/Message";
-import clientChatStore from "../../../stores/ClientChatStore";
-import UserOptionsStore from "../../../stores/UserOptionsStore";
+import messageEditorStore from "../../../stores/MessageEditorStore";
 
 interface MessageEditorProps {
   message: Instance<typeof Message>;
@@ -13,113 +12,38 @@ interface MessageEditorProps {
 }
 
 const MessageEditor = observer(({ message, onCancel }: MessageEditorProps) => {
-  const [editedContent, setEditedContent] = useState(message.content);
+  useEffect(() => {
+    messageEditorStore.setMessage(message);
 
-  const handleSave = async () => {
-    message.setContent(editedContent);
-
-    // Find the index of the edited message
-    const messageIndex = clientChatStore.items.indexOf(message);
-    if (messageIndex !== -1) {
-      // Remove all messages after the edited message
-      clientChatStore.removeAfter(messageIndex);
-
-      // Set follow mode and loading state
-      UserOptionsStore.setFollowModeEnabled(true);
-      clientChatStore.setIsLoading(true);
-
-      try {
-        // Add a small delay before adding the assistant message (for better UX)
-        await new Promise((r) => setTimeout(r, 500));
-
-        // Add an empty assistant message
-        clientChatStore.add(Message.create({ content: "", role: "assistant" }));
-        const payload = { messages: clientChatStore.items.slice(), model: clientChatStore.model };
-
-        // Make API call
-        const response = await fetch("/api/chat", {
-          method: "POST",
-          headers: { "Content-Type": "application/json" },
-          body: JSON.stringify(payload),
-        });
-
-        if (response.status === 429) {
-          clientChatStore.updateLast("Too many requests • please slow down.");
-          clientChatStore.setIsLoading(false);
-          UserOptionsStore.setFollowModeEnabled(false);
-          return;
-        }
-        if (response.status > 200) {
-          clientChatStore.updateLast("Error • something went wrong.");
-          clientChatStore.setIsLoading(false);
-          UserOptionsStore.setFollowModeEnabled(false);
-          return;
-        }
-
-        const { streamUrl } = await response.json();
-        const eventSource = new EventSource(streamUrl);
-
-        eventSource.onmessage = (event) => {
-          try {
-            const parsed = JSON.parse(event.data);
-            if (parsed.type === "error") {
-              clientChatStore.updateLast(parsed.error);
-              clientChatStore.setIsLoading(false);
-              UserOptionsStore.setFollowModeEnabled(false);
-              eventSource.close();
-              return;
-            }
-
-            if (parsed.type === "chat" && parsed.data.choices[0]?.finish_reason === "stop") {
-              clientChatStore.appendLast(parsed.data.choices[0]?.delta?.content ?? "");
-              clientChatStore.setIsLoading(false);
-              UserOptionsStore.setFollowModeEnabled(false);
-              eventSource.close();
-              return;
-            }
-
-            if (parsed.type === "chat") {
-              clientChatStore.appendLast(parsed.data.choices[0]?.delta?.content ?? "");
-            }
-          } catch (err) {
-            console.error("stream parse error", err);
-          }
-        };
-
-        eventSource.onerror = () => {
-          clientChatStore.updateLast("Error • connection lost.");
-          clientChatStore.setIsLoading(false);
-          UserOptionsStore.setFollowModeEnabled(false);
-          eventSource.close();
-        };
-      } catch (err) {
-        console.error("sendMessage", err);
-        clientChatStore.updateLast("Sorry • network error.");
-        clientChatStore.setIsLoading(false);
-        UserOptionsStore.setFollowModeEnabled(false);
-      }
-    }
+    return () => {
+      messageEditorStore.onCancel();
+    };
+  }, [message]);
 
+  const handleCancel = () => {
+    messageEditorStore.onCancel();
     onCancel();
   };
 
+
+
   const handleKeyDown = (e: KeyboardEvent<HTMLTextAreaElement>) => {
     if (e.key === "Enter" && (e.metaKey || e.ctrlKey)) {
       e.preventDefault();
-      handleSave();
+      messageEditorStore.handleSave();
     }
 
     if (e.key === "Escape") {
       e.preventDefault();
-      onCancel();
+      handleCancel();
     }
   };
 
   return (
     <Box width="100%">
       <Textarea
-        value={editedContent}
-        onChange={(e) => setEditedContent(e.target.value)}
+        value={messageEditorStore.editedContent}
+        onChange={(e) => messageEditorStore.setEditedContent(e.target.value)}
         onKeyDown={handleKeyDown}
         minHeight="100px"
         bg="transparent"
@@ -134,7 +58,7 @@ const MessageEditor = observer(({ message, onCancel }: MessageEditorProps) => {
         <IconButton
           aria-label="Cancel edit"
           icon={<X />}
-          onClick={onCancel}
+          onClick={handleCancel}
           size="sm"
           variant="ghost"
           color={"accent.danger"}
@@ -142,7 +66,7 @@ const MessageEditor = observer(({ message, onCancel }: MessageEditorProps) => {
         <IconButton
           aria-label="Save edit"
           icon={<Check />}
-          onClick={handleSave}
+          onClick={() => messageEditorStore.handleSave()}
           size="sm"
           variant="ghost"
           color={"accent.confirm"}
diff --git a/src/components/chat/messages/__tests__/MessageEditorComponent.test.tsx b/src/components/chat/messages/__tests__/MessageEditorComponent.test.tsx
index fd3806b..ad2c065 100644
--- a/src/components/chat/messages/__tests__/MessageEditorComponent.test.tsx
+++ b/src/components/chat/messages/__tests__/MessageEditorComponent.test.tsx
@@ -3,8 +3,9 @@ import { render, screen, fireEvent } from '@testing-library/react';
 import React from 'react';
 import MessageEditor from '../MessageEditorComponent';
 
-// Import the mocked clientChatStore
+// Import the mocked stores
 import clientChatStore from '../../../../stores/ClientChatStore';
+import messageEditorStore from '../../../../stores/MessageEditorStore';
 
 // Mock the Message model
 vi.mock('../../../../models/Message', () => {
@@ -15,13 +16,22 @@ vi.mock('../../../../models/Message', () => {
   };
 });
 
+// Mock fetch globally
+globalThis.fetch = vi.fn(() =>
+    Promise.resolve({
+      ok: true,
+      json: () => Promise.resolve({})
+    })
+);
+
 // Mock the ClientChatStore
 vi.mock('../../../../stores/ClientChatStore', () => {
   const mockStore = {
     items: [],
     removeAfter: vi.fn(),
     sendMessage: vi.fn(),
-    setIsLoading: vi.fn()
+    setIsLoading: vi.fn(),
+    editMessage: vi.fn().mockReturnValue(true)
   };
 
   // Add the mockUserMessage to the items array
@@ -32,6 +42,25 @@ vi.mock('../../../../stores/ClientChatStore', () => {
   };
 });
 
+// Mock the MessageEditorStore
+vi.mock('../../../../stores/MessageEditorStore', () => {
+  const mockStore = {
+    editedContent: 'Test message', // Set initial value to match the test expectation
+    message: null,
+    setEditedContent: vi.fn(),
+    setMessage: vi.fn((message) => {
+      mockStore.message = message;
+      mockStore.editedContent = message.content;
+    }),
+    onCancel: vi.fn(),
+    handleSave: vi.fn()
+  };
+
+  return {
+    default: mockStore
+  };
+});
+
 describe('MessageEditor', () => {
   // Create a message object with a setContent method
   const mockUserMessage = { 
@@ -51,6 +80,7 @@ describe('MessageEditor', () => {
     const textarea = screen.getByRole('textbox');
     expect(textarea).toBeInTheDocument();
     expect(textarea).toHaveValue('Test message');
+    expect(messageEditorStore.setMessage).toHaveBeenCalledWith(mockUserMessage);
   });
 
   it('should update the content when typing', () => {
@@ -59,22 +89,16 @@ describe('MessageEditor', () => {
     const textarea = screen.getByRole('textbox');
     fireEvent.change(textarea, { target: { value: 'Updated message' } });
 
-    expect(textarea).toHaveValue('Updated message');
+    expect(messageEditorStore.setEditedContent).toHaveBeenCalledWith('Updated message');
   });
 
-  it('should call setContent, removeAfter, sendMessage, and onCancel when save button is clicked', () => {
-    render(<MessageEditor message={mockUserMessage} onCancel={mockOnCancel} />);
-
-    const textarea = screen.getByRole('textbox');
-    fireEvent.change(textarea, { target: { value: 'Updated message' } });
+  it('should call handleSave when save button is clicked', () => {
+    render(<MessageEditor message={mockUserMessage} onCancel={mockOnCancel}/>);
 
     const saveButton = screen.getByLabelText('Save edit');
     fireEvent.click(saveButton);
 
-    expect(mockUserMessage.setContent).toHaveBeenCalledWith('Updated message');
-    expect(clientChatStore.removeAfter).toHaveBeenCalledWith(0);
-    expect(clientChatStore.sendMessage).toHaveBeenCalled();
-    expect(mockOnCancel).toHaveBeenCalled();
+    expect(messageEditorStore.handleSave).toHaveBeenCalled();
   });
 
   it('should call onCancel when cancel button is clicked', () => {
@@ -83,43 +107,35 @@ describe('MessageEditor', () => {
     const cancelButton = screen.getByLabelText('Cancel edit');
     fireEvent.click(cancelButton);
 
+    expect(messageEditorStore.onCancel).toHaveBeenCalled();
     expect(mockOnCancel).toHaveBeenCalled();
-    expect(mockUserMessage.setContent).not.toHaveBeenCalled();
   });
 
-  it('should save when Ctrl+Enter is pressed', () => {
+  it('should call handleSave when Ctrl+Enter is pressed', () => {
     render(<MessageEditor message={mockUserMessage} onCancel={mockOnCancel} />);
 
     const textarea = screen.getByRole('textbox');
-    fireEvent.change(textarea, { target: { value: 'Updated message' } });
     fireEvent.keyDown(textarea, { key: 'Enter', ctrlKey: true });
 
-    expect(mockUserMessage.setContent).toHaveBeenCalledWith('Updated message');
-    expect(clientChatStore.removeAfter).toHaveBeenCalledWith(0);
-    expect(clientChatStore.sendMessage).toHaveBeenCalled();
-    expect(mockOnCancel).toHaveBeenCalled();
+    expect(messageEditorStore.handleSave).toHaveBeenCalled();
   });
 
-  it('should save when Meta+Enter is pressed', () => {
+  it('should call handleSave when Meta+Enter is pressed', () => {
     render(<MessageEditor message={mockUserMessage} onCancel={mockOnCancel} />);
 
     const textarea = screen.getByRole('textbox');
-    fireEvent.change(textarea, { target: { value: 'Updated message' } });
     fireEvent.keyDown(textarea, { key: 'Enter', metaKey: true });
 
-    expect(mockUserMessage.setContent).toHaveBeenCalledWith('Updated message');
-    expect(clientChatStore.removeAfter).toHaveBeenCalledWith(0);
-    expect(clientChatStore.sendMessage).toHaveBeenCalled();
-    expect(mockOnCancel).toHaveBeenCalled();
+    expect(messageEditorStore.handleSave).toHaveBeenCalled();
   });
 
-  it('should cancel when Escape is pressed', () => {
+  it('should call onCancel when Escape is pressed', () => {
     render(<MessageEditor message={mockUserMessage} onCancel={mockOnCancel} />);
 
     const textarea = screen.getByRole('textbox');
     fireEvent.keyDown(textarea, { key: 'Escape' });
 
+    expect(messageEditorStore.onCancel).toHaveBeenCalled();
     expect(mockOnCancel).toHaveBeenCalled();
-    expect(mockUserMessage.setContent).not.toHaveBeenCalled();
   });
 });
diff --git a/src/components/markdown/MarkdownComponent.tsx b/src/components/markdown/MarkdownComponent.tsx
index 51aab63..0867309 100644
--- a/src/components/markdown/MarkdownComponent.tsx
+++ b/src/components/markdown/MarkdownComponent.tsx
@@ -238,7 +238,7 @@ const CustomTable: React.FC<{
 };
 
 const CustomHtmlBlock: React.FC<{ content: string }> = ({content}) => {
-    return <Box dangerouslySetInnerHTML={{__html: content}} mb={2}/>;
+    return <Box as="span" display="inline" dangerouslySetInnerHTML={{__html: content}} mb={2}/>;
 };
 
 const CustomText: React.FC<{ text: React.ReactNode }> = ({text}) => {
diff --git a/src/models/Message.ts b/src/models/Message.ts
index c2c9b42..9f72336 100644
--- a/src/models/Message.ts
+++ b/src/models/Message.ts
@@ -1,7 +1,13 @@
 import { types } from "mobx-state-tree";
 
+// Simple function to generate a unique ID
+export const generateId = () => {
+  return Date.now().toString(36) + Math.random().toString(36).substring(2);
+};
+
 export default types
   .model("Message", {
+    id: types.optional(types.identifier, generateId),
     content: types.string,
     role: types.enumeration(["user", "assistant"]),
   })
diff --git a/src/stores/MessageEditorStore.ts b/src/stores/MessageEditorStore.ts
new file mode 100644
index 0000000..3bf65b2
--- /dev/null
+++ b/src/stores/MessageEditorStore.ts
@@ -0,0 +1,148 @@
+import {types, Instance} from "mobx-state-tree";
+import clientChatStore from "./ClientChatStore";
+import UserOptionsStore from "./UserOptionsStore";
+import Message from "../models/Message";
+import {MessagesStore} from "./MessagesStore";
+
+export const MessageEditorStore = types
+    .compose(
+        MessagesStore,
+        types.model("MessageEditorStore", {
+            editedContent: types.optional(types.string, ""),
+            messageId: types.optional(types.string, "")
+        })
+    )
+    .views((self) => ({
+        getMessage() {
+            // Find the message in clientChatStore by ID
+            if (!self.messageId) return null;
+
+            const message = clientChatStore.items.find(item => item.id === self.messageId);
+            return message || null;
+        }
+    }))
+    .actions((self) => ({
+        setEditedContent(content: string) {
+            self.editedContent = content;
+        },
+        setMessage(message: Instance<typeof Message>) {
+            self.messageId = message.id;
+            self.editedContent = message.content;
+        },
+        onCancel() {
+            self.messageId = "";
+            self.editedContent = "";
+        },
+        handleSave: async () => {
+            // Get the message using the ID
+            const message = self.getMessage();
+
+            // Check if message reference is still valid
+            if (!message) {
+                // Message reference is no longer valid, just cancel the edit
+                self.onCancel();
+                return;
+            }
+
+            // Store the content we want to update
+            const contentToUpdate = self.editedContent;
+
+            try {
+                // Use the editMessage function from MessagesStore
+                const success = clientChatStore.editMessage(message, contentToUpdate);
+                if (!success) {
+                    // Message not found in the items array, just cancel
+                    self.onCancel();
+                    return;
+                }
+
+                // Set follow mode and loading state
+                UserOptionsStore.setFollowModeEnabled(true);
+                clientChatStore.setIsLoading(true);
+
+                try {
+                    // Add a small delay before adding the assistant message (for better UX)
+                    await new Promise((r) => setTimeout(r, 500));
+
+                    // Add an empty assistant message to clientChatStore's items
+                    clientChatStore.add(Message.create({content: "", role: "assistant"}));
+                    // Use clientChatStore for the API call since it has the model property
+                    const payload = {messages: clientChatStore.items.slice(), model: clientChatStore.model};
+
+                    // Make API call
+                    const response = await fetch("/api/chat", {
+                        method: "POST",
+                        headers: {"Content-Type": "application/json"},
+                        body: JSON.stringify(payload),
+                    });
+
+                    if (response.status === 429) {
+                        clientChatStore.updateLast("Too many requests • please slow down.");
+                        clientChatStore.setIsLoading(false);
+                        UserOptionsStore.setFollowModeEnabled(false);
+                        return;
+                    }
+                    if (response.status > 200) {
+                        clientChatStore.updateLast("Error • something went wrong.");
+                        clientChatStore.setIsLoading(false);
+                        UserOptionsStore.setFollowModeEnabled(false);
+                        return;
+                    }
+
+                    const {streamUrl} = await response.json();
+                    const eventSource = new EventSource(streamUrl);
+
+                    eventSource.onmessage = (event) => {
+                        try {
+                            const parsed = JSON.parse(event.data);
+                            if (parsed.type === "error") {
+                                clientChatStore.updateLast(parsed.error);
+                                clientChatStore.setIsLoading(false);
+                                UserOptionsStore.setFollowModeEnabled(false);
+                                eventSource.close();
+                                return;
+                            }
+
+                            if (parsed.type === "chat" && parsed.data.choices[0]?.finish_reason === "stop") {
+                                clientChatStore.appendLast(parsed.data.choices[0]?.delta?.content ?? "");
+                                clientChatStore.setIsLoading(false);
+                                UserOptionsStore.setFollowModeEnabled(false);
+                                eventSource.close();
+                                return;
+                            }
+
+                            if (parsed.type === "chat") {
+                                clientChatStore.appendLast(parsed.data.choices[0]?.delta?.content ?? "");
+                            }
+                        } catch (err) {
+                            console.error("stream parse error", err);
+                        }
+                    };
+
+                    eventSource.onerror = () => {
+                        clientChatStore.updateLast("Error • connection lost.");
+                        clientChatStore.setIsLoading(false);
+                        UserOptionsStore.setFollowModeEnabled(false);
+                        eventSource.close();
+                    };
+                } catch (err) {
+                    console.error("sendMessage", err);
+                    clientChatStore.updateLast("Sorry • network error.");
+                    clientChatStore.setIsLoading(false);
+                    UserOptionsStore.setFollowModeEnabled(false);
+                }
+            } catch (err) {
+                console.error("Error in handleSave:", err);
+                // If any error occurs, just cancel the edit
+                self.onCancel();
+                return;
+            }
+
+            // Always clean up at the end
+            self.onCancel();
+        }
+    }));
+
+const messageEditorStore = MessageEditorStore.create();
+
+export default messageEditorStore;
diff --git a/src/stores/MessagesStore.ts b/src/stores/MessagesStore.ts
index b9c73d6..b43108f 100644
--- a/src/stores/MessagesStore.ts
+++ b/src/stores/MessagesStore.ts
@@ -30,6 +30,22 @@ export const MessagesStore = types
         reset() {
             self.items.clear();
         },
+        editMessage(message: Instance<typeof Message>, newContent: string) {
+            // Find the index of the message in the items array
+            const messageIndex = self.items.indexOf(message);
+            if (messageIndex === -1) {
+                // Message not found in the items array
+                return false;
+            }
+
+            // Update the message content
+            message.setContent(newContent);
+
+            // Remove all messages after the edited message
+            self.removeAfter(messageIndex);
+
+            return true;
+        },
     }));
 
 export interface IMessagesStore extends Instance<typeof MessagesStore> {}
diff --git a/src/stores/StreamStore.ts b/src/stores/StreamStore.ts
index a640256..d03112b 100644
--- a/src/stores/StreamStore.ts
+++ b/src/stores/StreamStore.ts
@@ -3,6 +3,7 @@ import {
     Instance,
     flow,
     types,
+    applyAction,
 } from "mobx-state-tree";
 import type { IMessagesStore } from "./MessagesStore";
 import type { IUIStore } from "./UIStore";
@@ -28,7 +29,7 @@ export const StreamStore = types
         }
 
         function setEventSource(source: EventSource | null) {
-            self.eventSource = source;
+                self.eventSource = source;
         }
 
         function cleanup() {
diff --git a/vite.config.ts b/vite.config.ts
index ccb9541..38fadbf 100644
--- a/vite.config.ts
+++ b/vite.config.ts
@@ -71,7 +71,7 @@ export default defineConfig(({command}) => {
             },
         },
         esbuild: {
-            drop: ["console"]
+            // drop: ["console"]
         },
         build: {
             emitAssets: false,
diff --git a/workers/site/services/ChatService.ts b/workers/site/services/ChatService.ts
index 6b4cabc..76909a7 100644
--- a/workers/site/services/ChatService.ts
+++ b/workers/site/services/ChatService.ts
@@ -121,7 +121,10 @@ const ChatService = types
                     console.log("getting local models")
                     const openaiClient = new OpenAI({baseURL: self.env.OPENAI_API_ENDPOINT})
                     const models = await openaiClient.models.list();
-                    return Response.json(models.data.map(model => model.id));
+                    return Response.json(
+                        models.data
+                            .filter(model => model.id.includes("mlx"))
+                            .map(model => model.id));
                 }
                 return Response.json(SUPPORTED_MODELS);
             },
@@ -186,8 +189,10 @@ const ChatService = types
                 }
 
                 const handler = useModelHandler();
+                
 
                 if (handler) {
+                    console.log(`Using provider: ${modelFamily}`);
                     try {
                         await handler(streamParams, handleStreamData(controller, encoder));