diff --git a/src/lib/Chat.svelte b/src/lib/Chat.svelte
index 9becf1e..0e1f2e1 100644
--- a/src/lib/Chat.svelte
+++ b/src/lib/Chat.svelte
@@ -3,6 +3,7 @@
 
   import { apiKeyStorage, chatsStorage, addMessage, clearMessages } from "./Storage.svelte";
   import type { Request, Response, Message, Settings } from "./Types.svelte";
+  import Code from "./Code.svelte";
 
   import { afterUpdate, onMount } from "svelte";
   import SvelteMarkdown from "svelte-markdown";
@@ -330,7 +331,7 @@
           source={message.content}
           options={markedownOptions}
           renderers={{
-            /*code: Code*/
+            code: Code,
           }}
         />
       </div>
@@ -342,7 +343,7 @@
           source={message.content}
           options={markedownOptions}
           renderers={{
-            /*code: Code*/
+            code: Code,
           }}
         />
       </div>
@@ -354,7 +355,7 @@
           source={message.content}
           options={markedownOptions}
           renderers={{
-            /*code: Code*/
+            code: Code,
           }}
         />
         {#if message.usage}
diff --git a/src/lib/Code.svelte b/src/lib/Code.svelte
index a843903..cddd000 100644
--- a/src/lib/Code.svelte
+++ b/src/lib/Code.svelte
@@ -1,23 +1,77 @@
 <script lang="ts">
-  // Disabled for now - it blows up the size too much
-  import { HighlightAuto } from "svelte-highlight";
-  import style from "svelte-highlight/styles/edge-light";
+  import { Highlight } from "svelte-highlight";
 
-  export let type: "code";
-  export let raw: string;
-  export let codeBlockStyle: "indented" | undefined = undefined;
+  // Import both dark and light styles
+  import { github, githubDark } from "svelte-highlight/styles";
+
+  // Style depends on system theme
+  const style = window.matchMedia("(prefers-color-scheme: dark)").matches ? githubDark : github;
+
+  // Import all supported languages
+  import {
+    javascript,
+    python,
+    typescript,
+    ruby,
+    go,
+    java,
+    sql,
+    shell,
+    php,
+    plaintext,
+    type LanguageType,
+  } from "svelte-highlight/languages";
+
+  export const type: "code" = "code";
+  export const raw: string = "";
+  export const codeBlockStyle: "indented" | undefined = undefined;
   export let lang: string | undefined = undefined;
   export let text: string;
 
-  // Disable warnings
-  type = "code";
-  raw;
-  codeBlockStyle;
-  lang;
+  // Map lang string to LanguageType
+  let language: LanguageType<string>;
+  switch (lang) {
+    case "js":
+    case "javascript":
+      language = javascript;
+      break;
+    case "py":
+    case "python":
+      language = python;
+      break;
+    case "ts":
+    case "typescript":
+      language = typescript;
+      break;
+    case "rb":
+    case "ruby":
+      language = ruby;
+      break;
+    case "go":
+    case "golang":
+      language = go;
+      break;
+    case "java":
+      language = java;
+      break;
+    case "sql":
+      language = sql;
+      break;
+    case "sh":
+    case "shell":
+    case "bash":
+      language = shell;
+      break;
+    case "php":
+      language = php;
+      break;
+    default:
+      language = plaintext;
+  }
 </script>
 
 <svelte:head>
   {@html style}
 </svelte:head>
 
-<HighlightAuto code={text} />
+<Highlight code={text} {language} />