Merge pull request #16 from TGlide/improvements

model picker & other improvements
2025-06-18 10:47:13 +01:00 · 2025-06-18 10:47:13 +01:00 · 7247fe9322
commit 7247fe9322
parent e8b5f94456 aff1901765
9 changed files with 594 additions and 44 deletions
--- a/README.md
+++ b/README.md
@ -54,12 +54,14 @@ TODO: add instructions
 - [x] conversation title generation
 - [ ] kbd powered popover model picker
 - [x] autosize
+- [x] AbortController for message generation

 ### Extra

 - [ ] Web Search
 - [ ] MCPs
 - [ ] Chat branching
+- [ ] Regenerate message
 - [ ] Image generation
 - [ ] Chat sharing
 - [ ] 404 page/redirect
--- a/package.json
+++ b/package.json
@ -77,6 +77,7 @@
 		"@fontsource-variable/nunito-sans": "^5.2.6",
 		"@fontsource-variable/open-sans": "^5.2.6",
 		"better-auth": "^1.2.9",
+		"bits-ui": "^2.8.0",
 		"convex-helpers": "^0.1.94",
 		"markdown-it-async": "^2.2.0",
 		"openai": "^5.3.0",
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@ -35,6 +35,9 @@ importers:
      better-auth:
        specifier: ^1.2.9
        version: 1.2.9
+      bits-ui:
+        specifier: ^2.8.0
+        version: 2.8.0(@internationalized/date@3.8.2)(svelte@5.34.1)
      convex-helpers:
        specifier: ^0.1.94
        version: 0.1.94(convex@1.24.8)(typescript@5.8.3)(zod@3.25.64)
@ -662,6 +665,9 @@ packages:
  '@iconify/utils@2.3.0':
    resolution: {integrity: sha512-GmQ78prtwYW6EtzXRU1rY+KwOKfz32PD7iJh6Iyqw68GiKuoZ2A6pRtzWONz5VQJbp50mEjXh/7NkumtrAgRKA==}

+  '@internationalized/date@3.8.2':
+    resolution: {integrity: sha512-/wENk7CbvLbkUvX1tu0mwq49CVkkWpkXubGel6birjRPyo6uQ4nQpnq5xZu823zRCwwn82zgHrvgF1vZyvmVgA==}
+
  '@isaacs/fs-minipass@4.0.1':
    resolution: {integrity: sha512-wgm9Ehl2jpeqP3zw/7mo3kRHFp5MEDhqAdwy1fTGkHAwnkGOVsgpvQhL8B5n1qlb01jV3n/bI0ZfZp5lWA1k4w==}
    engines: {node: '>=18.0.0'}
@ -906,6 +912,9 @@ packages:
      svelte: ^5.0.0
      vite: ^6.0.0

+  '@swc/helpers@0.5.17':
+    resolution: {integrity: sha512-5IKx/Y13RsYd+sauPb2x+U/xZikHjolzfuDgTAl/Tdf3Q8rslRvC19NKDLgAJQ6wsqADk10ntlv08nPFw/gO/A==}
+
  '@tailwindcss/node@4.1.10':
    resolution: {integrity: sha512-2ACf1znY5fpRBwRhMgj9ZXvb2XZW8qs+oTfotJ2C5xR0/WNL7UHZ7zXl6s+rUqedL1mNi+0O+WQr5awGowS3PQ==}

@ -1223,6 +1232,13 @@ packages:
  better-call@1.0.9:
    resolution: {integrity: sha512-Qfm0gjk0XQz0oI7qvTK1hbqTsBY4xV2hsHAxF8LZfUYl3RaECCIifXuVqtPpZJWvlCCMlQSvkvhhyuApGUba6g==}

+  bits-ui@2.8.0:
+    resolution: {integrity: sha512-WiTZcCbYLm4Cx6/67NqXVSD0BkfNmdX8Abs84HpIaplX/wRRbg8tkMtJYlLw7mepgGvwGR3enLi6tFkcHU3JXA==}
+    engines: {node: '>=20', pnpm: '>=8.7.0'}
+    peerDependencies:
+      '@internationalized/date': ^3.8.1
+      svelte: ^5.33.0
+
  brace-expansion@1.1.12:
    resolution: {integrity: sha512-9T9UjW3r0UW5c1Q7GTwllptXwhvYmEzFhzMfZ9H7FQWt+uZePjZPjBP/W1ZEyZ1twGWom5/56TF4lPcqjnDHcg==}

@ -2381,6 +2397,12 @@ packages:
    peerDependencies:
      svelte: ^5.0.0

+  svelte-toolbelt@0.9.2:
+    resolution: {integrity: sha512-REb1cENGnFbhNSmIdCb1SDIpjEa3n1kXhNVHqGNEesjmPX3bG87gUZiCG8cqOt9AAarqzTzOtI2jEEWr/ZbHwA==}
+    engines: {node: '>=18', pnpm: '>=8.7.0'}
+    peerDependencies:
+      svelte: ^5.30.2
+
  svelte@5.34.1:
    resolution: {integrity: sha512-jWNnN2hZFNtnzKPptCcJHBWrD9CtbHPDwIRIODufOYaWkR0kLmAIlM384lMt4ucwuIRX4hCJwD2D8ZtEcGJQ0Q==}
    engines: {node: '>=18'}
@ -2388,6 +2410,9 @@ packages:
  symbol-tree@3.2.4:
    resolution: {integrity: sha512-9QNk5KwDF+Bvz+PyObkmSYjI5ksVUYtjW7AU22r2NKcfLJcXp96hkDWU3+XndOsUb+AQ9QhfzfCT2O+CNWT5Tw==}

+  tabbable@6.2.0:
+    resolution: {integrity: sha512-Cat63mxsVJlzYvN51JmVXIgNoUokrIaT2zLclCXjRd8boZ0004U4KCs/sToJ75C6sdlByWxpYnb5Boif1VSFew==}
+
  tailwind-merge@3.0.2:
    resolution: {integrity: sha512-l7z+OYZ7mu3DTqrL88RiKrKIqO3NcpEO8V/Od04bNpvk0kiIFndGEoqfuzvj4yuhRkHKjRkII2z+KS2HfPcSxw==}

@ -3061,6 +3086,10 @@ snapshots:
    transitivePeerDependencies:
      - supports-color

+  '@internationalized/date@3.8.2':
+    dependencies:
+      '@swc/helpers': 0.5.17
+
  '@isaacs/fs-minipass@4.0.1':
    dependencies:
      minipass: 7.1.2
@ -3306,6 +3335,10 @@ snapshots:
    transitivePeerDependencies:
      - supports-color

+  '@swc/helpers@0.5.17':
+    dependencies:
+      tslib: 2.8.1
+
  '@tailwindcss/node@4.1.10':
    dependencies:
      '@ampproject/remapping': 2.3.0
@ -3651,6 +3684,18 @@ snapshots:
      set-cookie-parser: 2.7.1
      uncrypto: 0.1.3

+  bits-ui@2.8.0(@internationalized/date@3.8.2)(svelte@5.34.1):
+    dependencies:
+      '@floating-ui/core': 1.7.1
+      '@floating-ui/dom': 1.7.1
+      '@internationalized/date': 3.8.2
+      css.escape: 1.5.1
+      esm-env: 1.2.2
+      runed: 0.28.0(svelte@5.34.1)
+      svelte: 5.34.1
+      svelte-toolbelt: 0.9.2(svelte@5.34.1)
+      tabbable: 6.2.0
+
  brace-expansion@1.1.12:
    dependencies:
      balanced-match: 1.0.2
@ -4744,6 +4789,13 @@ snapshots:
      style-to-object: 1.0.9
      svelte: 5.34.1

+  svelte-toolbelt@0.9.2(svelte@5.34.1):
+    dependencies:
+      clsx: 2.1.1
+      runed: 0.28.0(svelte@5.34.1)
+      style-to-object: 1.0.9
+      svelte: 5.34.1
+
  svelte@5.34.1:
    dependencies:
      '@ampproject/remapping': 2.3.0
@ -4763,6 +4815,8 @@ snapshots:

  symbol-tree@3.2.4: {}

+  tabbable@6.2.0: {}
+
  tailwind-merge@3.0.2: {}

  tailwind-merge@3.3.1: {}
--- a/src/routes/api/cancel-generation/+server.ts
+++ b/src/routes/api/cancel-generation/+server.ts
@ -0,0 +1,88 @@
+import { api } from '$lib/backend/convex/_generated/api';
+import type { Id } from '$lib/backend/convex/_generated/dataModel';
+import { error, json, type RequestHandler } from '@sveltejs/kit';
+import { ConvexHttpClient } from 'convex/browser';
+import { ResultAsync } from 'neverthrow';
+import { z } from 'zod/v4';
+import { getSessionCookie } from 'better-auth/cookies';
+import { PUBLIC_CONVEX_URL } from '$env/static/public';
+
+// Import the global cache from generate-message
+import { generationAbortControllers } from '../generate-message/cache.js';
+
+const client = new ConvexHttpClient(PUBLIC_CONVEX_URL);
+
+const reqBodySchema = z.object({
+	conversation_id: z.string(),
+	session_token: z.string(),
+});
+
+export type CancelGenerationRequestBody = z.infer<typeof reqBodySchema>;
+
+export type CancelGenerationResponse = {
+	ok: true;
+	cancelled: boolean;
+};
+
+function response(res: CancelGenerationResponse) {
+	return json(res);
+}
+
+export const POST: RequestHandler = async ({ request }) => {
+	const bodyResult = await ResultAsync.fromPromise(
+		request.json(),
+		() => 'Failed to parse request body'
+	);
+
+	if (bodyResult.isErr()) {
+		return error(400, 'Failed to parse request body');
+	}
+
+	const parsed = reqBodySchema.safeParse(bodyResult.value);
+	if (!parsed.success) {
+		return error(400, parsed.error);
+	}
+	const args = parsed.data;
+
+	const cookie = getSessionCookie(request.headers);
+	const sessionToken = cookie?.split('.')[0] ?? null;
+
+	if (!sessionToken || sessionToken !== args.session_token) {
+		return error(401, 'Unauthorized');
+	}
+
+	// Verify the user owns this conversation
+	const conversationResult = await ResultAsync.fromPromise(
+		client.query(api.conversations.getById, {
+			conversation_id: args.conversation_id as Id<'conversations'>,
+			session_token: sessionToken,
+		}),
+		(e) => `Failed to get conversation: ${e}`
+	);
+
+	if (conversationResult.isErr()) {
+		return error(403, 'Conversation not found or unauthorized');
+	}
+
+	// Try to cancel the generation
+	const abortController = generationAbortControllers.get(args.conversation_id);
+	let cancelled = false;
+
+	if (abortController) {
+		abortController.abort();
+		generationAbortControllers.delete(args.conversation_id);
+		cancelled = true;
+
+		// Update conversation generating status to false
+		await ResultAsync.fromPromise(
+			client.mutation(api.conversations.updateGenerating, {
+				conversation_id: args.conversation_id as Id<'conversations'>,
+				generating: false,
+				session_token: sessionToken,
+			}),
+			(e) => `Failed to update generating status: ${e}`
+		);
+	}
+
+	return response({ ok: true, cancelled });
+};
--- a/src/routes/api/cancel-generation/call.ts
+++ b/src/routes/api/cancel-generation/call.ts
@ -0,0 +1,17 @@
+import { ResultAsync } from 'neverthrow';
+import type { CancelGenerationRequestBody, CancelGenerationResponse } from './+server';
+
+export async function callCancelGeneration(args: CancelGenerationRequestBody) {
+	const res = ResultAsync.fromPromise(
+		fetch('/api/cancel-generation', {
+			method: 'POST',
+			headers: {
+				'Content-Type': 'application/json',
+			},
+			body: JSON.stringify(args),
+		}),
+		(e) => e
+	).map((r) => r.json() as Promise<CancelGenerationResponse>);
+
+	return res;
+}
--- a/src/routes/api/generate-message/+server.ts
+++ b/src/routes/api/generate-message/+server.ts
@ -3,13 +3,13 @@ import { api } from '$lib/backend/convex/_generated/api';
 import type { Doc, Id } from '$lib/backend/convex/_generated/dataModel';
 import { Provider } from '$lib/types';
 import { error, json, type RequestHandler } from '@sveltejs/kit';
+import { waitUntil } from '@vercel/functions';
+import { getSessionCookie } from 'better-auth/cookies';
 import { ConvexHttpClient } from 'convex/browser';
 import { err, ok, Result, ResultAsync } from 'neverthrow';
 import OpenAI from 'openai';
-import { waitUntil } from '@vercel/functions';
 import { z } from 'zod/v4';
-import type { ChatCompletionSystemMessageParam } from 'openai/resources';
-import { getSessionCookie } from 'better-auth/cookies';
+import { generationAbortControllers } from './cache.js';

 // Set to true to enable debug logging
 const ENABLE_LOGGING = true;
@ -162,6 +162,7 @@ async function generateAIResponse({
 	modelResultPromise,
 	keyResultPromise,
 	rulesResultPromise,
+	abortSignal,
 }: {
 	conversationId: string;
 	sessionToken: string;
@ -169,9 +170,15 @@ async function generateAIResponse({
 	keyResultPromise: ResultAsync<string | null, string>;
 	modelResultPromise: ResultAsync<Doc<'user_enabled_models'> | null, string>;
 	rulesResultPromise: ResultAsync<Doc<'user_rules'>[], string>;
+	abortSignal?: AbortSignal;
 }) {
 	log('Starting AI response generation in background', startTime);

+	if (abortSignal?.aborted) {
+		log('AI response generation aborted before starting', startTime);
+		return;
+	}
+
 	const [modelResult, keyResult, messagesQueryResult, rulesResult] = await Promise.all([
 		modelResultPromise,
 		keyResultPromise,
@ -241,13 +248,6 @@ async function generateAIResponse({

 	log(`Background: ${attachedRules.length} rules attached`, startTime);

-	const systemMessage: ChatCompletionSystemMessageParam = {
-		role: 'system',
-		content: `Respond in markdown format. The user may have mentioned one or more rules to follow with the @<rule_name> syntax. Please follow these rules.
-Rules to follow:
-${attachedRules.map((r) => `- ${r.name}: ${r.rule}`).join('\n')}`,
-	};
-
 	const openai = new OpenAI({
 		baseURL: 'https://openrouter.ai/api/v1',
 		apiKey: key,
@ -272,13 +272,37 @@ ${attachedRules.map((r) => `- ${r.name}: ${r.rule}`).join('\n')}`,
 		};
 	});

+	// Only include system message if there are rules to follow
+	const messagesToSend =
+		attachedRules.length > 0
+			? [
+					...formattedMessages,
+					{
+						role: 'system' as const,
+						content: `Respond in markdown format. The user may have mentioned one or more rules to follow with the @<rule_name> syntax. Please follow these rules.
+Rules to follow:
+${attachedRules.map((r) => `- ${r.name}: ${r.rule}`).join('\n')}`,
+					},
+				]
+			: formattedMessages;
+
+	if (abortSignal?.aborted) {
+		log('AI response generation aborted before OpenAI call', startTime);
+		return;
+	}
+
 	const streamResult = await ResultAsync.fromPromise(
-		openai.chat.completions.create({
-			model: model.model_id,
-			messages: [...formattedMessages, systemMessage],
-			temperature: 0.7,
-			stream: true,
-		}),
+		openai.chat.completions.create(
+			{
+				model: model.model_id,
+				messages: messagesToSend,
+				temperature: 0.7,
+				stream: true,
+			},
+			{
+				signal: abortSignal,
+			}
+		),
 		(e) => `OpenAI API call failed: ${e}`
 	);

@ -317,6 +341,11 @@ ${attachedRules.map((r) => `- ${r.name}: ${r.rule}`).join('\n')}`,

 	try {
 		for await (const chunk of stream) {
+			if (abortSignal?.aborted) {
+				log('AI response generation aborted during streaming', startTime);
+				break;
+			}
+
 			chunkCount++;
 			content += chunk.choices[0]?.delta?.content || '';
 			if (!content) continue;
@ -420,6 +449,10 @@ ${attachedRules.map((r) => `- ${r.name}: ${r.rule}`).join('\n')}`,
 		log('Background: Cost usd updated', startTime);
 	} catch (error) {
 		log(`Background stream processing error: ${error}`, startTime);
+	} finally {
+		// Clean up the cached AbortController
+		generationAbortControllers.delete(conversationId);
+		log('Background: Cleaned up abort controller', startTime);
 	}
 }

@ -537,6 +570,25 @@ export const POST: RequestHandler = async ({ request }) => {
 		log('User message created', startTime);
 	}

+	// Set generating status to true before starting background generation
+	const setGeneratingResult = await ResultAsync.fromPromise(
+		client.mutation(api.conversations.updateGenerating, {
+			conversation_id: conversationId as Id<'conversations'>,
+			generating: true,
+			session_token: sessionToken,
+		}),
+		(e) => `Failed to set generating status: ${e}`
+	);
+
+	if (setGeneratingResult.isErr()) {
+		log(`Failed to set generating status: ${setGeneratingResult.error}`, startTime);
+		return error(500, 'Failed to set generating status');
+	}
+
+	// Create and cache AbortController for this generation
+	const abortController = new AbortController();
+	generationAbortControllers.set(conversationId, abortController);
+
 	// Start AI response generation in background - don't await
 	waitUntil(
 		generateAIResponse({
@ -546,9 +598,25 @@ export const POST: RequestHandler = async ({ request }) => {
 			modelResultPromise,
 			keyResultPromise,
 			rulesResultPromise,
-		}).catch((error) => {
-			log(`Background AI response generation error: ${error}`, startTime);
+			abortSignal: abortController.signal,
 		})
+			.catch(async (error) => {
+				log(`Background AI response generation error: ${error}`, startTime);
+				// Reset generating status on error
+				try {
+					await client.mutation(api.conversations.updateGenerating, {
+						conversation_id: conversationId as Id<'conversations'>,
+						generating: false,
+						session_token: sessionToken,
+					});
+				} catch (e) {
+					log(`Failed to reset generating status after error: ${e}`, startTime);
+				}
+			})
+			.finally(() => {
+				// Clean up the cached AbortController
+				generationAbortControllers.delete(conversationId);
+			})
 	);

 	log('Response sent, AI generation started in background', startTime);
--- a/src/routes/api/generate-message/cache.ts
+++ b/src/routes/api/generate-message/cache.ts
@ -0,0 +1,2 @@
+// Global cache for AbortControllers keyed by conversation ID
+export const generationAbortControllers = new Map<string, AbortController>();
--- a/src/routes/chat/+layout.svelte
+++ b/src/routes/chat/+layout.svelte
@ -28,6 +28,7 @@
 	import { Avatar } from 'melt/components';
 	import { Debounced, ElementSize, IsMounted, ScrollState } from 'runed';
 	import SendIcon from '~icons/lucide/arrow-up';
+	import StopIcon from '~icons/lucide/square';
 	import ChevronDownIcon from '~icons/lucide/chevron-down';
 	import ImageIcon from '~icons/lucide/image';
 	import LoaderCircleIcon from '~icons/lucide/loader-circle';
@ -38,6 +39,7 @@
 	import UploadIcon from '~icons/lucide/upload';
 	import XIcon from '~icons/lucide/x';
 	import { callGenerateMessage } from '../api/generate-message/call.js';
+	import { callCancelGeneration } from '../api/cancel-generation/call.js';
 	import ModelPicker from './model-picker.svelte';

 	const client = useConvexClient();
@ -46,7 +48,42 @@

 	let form = $state<HTMLFormElement>();
 	let textarea = $state<HTMLTextAreaElement>();
+	let abortController = $state<AbortController | null>(null);
+
+	const currentConversationQuery = useCachedQuery(api.conversations.getById, () => ({
+		conversation_id: page.params.id as Id<'conversations'>,
+		session_token: session.current?.session.token ?? '',
+	}));
+
+	const isGenerating = $derived(Boolean(currentConversationQuery.data?.generating));
+
+	async function stopGeneration() {
+		if (!page.params.id || !session.current?.session.token) return;
+
+		try {
+			const result = await callCancelGeneration({
+				conversation_id: page.params.id,
+				session_token: session.current.session.token,
+			});
+
+			if (result.isErr()) {
+				console.error('Failed to cancel generation:', result.error);
+			} else {
+				console.log('Generation cancelled:', result.value.cancelled);
+			}
+		} catch (error) {
+			console.error('Error cancelling generation:', error);
+		}
+
+		// Clear local abort controller if it exists
+		if (abortController) {
+			abortController = null;
+		}
+	}
+
 	async function handleSubmit() {
+		if (isGenerating) return;
+
 		const formData = new FormData(form);
 		const message = formData.get('message');

@ -58,19 +95,26 @@
 		const imagesCopy = [...selectedImages];
 		selectedImages = [];

-		const res = await callGenerateMessage({
-			message: messageCopy,
-			session_token: session.current?.session.token,
-			conversation_id: page.params.id ?? undefined,
-			model_id: settings.modelId,
-			images: imagesCopy.length > 0 ? imagesCopy : undefined,
-		});
-		if (res.isErr()) return; // TODO: Handle error
+		try {
+			const res = await callGenerateMessage({
+				message: messageCopy,
+				session_token: session.current?.session.token,
+				conversation_id: page.params.id ?? undefined,
+				model_id: settings.modelId,
+				images: imagesCopy.length > 0 ? imagesCopy : undefined,
+			});

-		const cid = res.value.conversation_id;
+			if (res.isErr()) {
+				return; // TODO: Handle error
+			}

-		if (page.params.id !== cid) {
-			goto(`/chat/${cid}`);
+			const cid = res.value.conversation_id;
+
+			if (page.params.id !== cid) {
+				goto(`/chat/${cid}`);
+			}
+		} catch (error) {
+			console.error('Error generating message:', error);
 		}
 	}

@ -586,12 +630,21 @@
 			</div>

 			<div
-				class="abs-x-center absolute -bottom-px left-1/2 mt-auto flex w-full max-w-3xl flex-col gap-1"
+				class="abs-x-center group absolute -bottom-px left-1/2 mt-auto flex w-full max-w-3xl flex-col gap-1"
 				bind:this={textareaWrapper}
 			>
-				<div class="border-reflect bg-background/80 rounded-t-[20px] p-2 pb-0 backdrop-blur-lg">
+				<div
+					class={[
+						'border-reflect bg-background/80 rounded-t-[20px] p-2 pb-0 backdrop-blur-lg',
+						'[--opacity:50%] group-focus-within:[--opacity:100%]',
+					]}
+				>
 					<form
-						class="bg-background/50 text-foreground outline-primary/10 dark:bg-secondary/20 relative flex w-full flex-col items-stretch gap-2 rounded-t-xl border border-b-0 border-white/70 px-3 pt-3 pb-3 outline-8 dark:border-white/10"
+						class={[
+							'bg-background/50 text-foreground dark:bg-secondary/20 relative flex w-full flex-col items-stretch gap-2 rounded-t-xl border border-b-0 border-white/70 px-3 pt-3 pb-3 outline-8 dark:border-white/10',
+							'transition duration-200',
+							'outline-primary/1 group-focus-within:outline-primary/10',
+						]}
 						style="box-shadow: rgba(0, 0, 0, 0.1) 0px 80px 50px 0px, rgba(0, 0, 0, 0.07) 0px 50px 30px 0px, rgba(0, 0, 0, 0.06) 0px 30px 15px 0px, rgba(0, 0, 0, 0.04) 0px 15px 8px, rgba(0, 0, 0, 0.04) 0px 6px 4px, rgba(0, 0, 0, 0.02) 0px 2px 2px;"
 						onsubmit={(e) => {
 							e.preventDefault();
@ -670,9 +723,11 @@
 								<textarea
 									{...pick(popover.trigger, ['id', 'style', 'onfocusout', 'onfocus'])}
 									bind:this={textarea}
-									disabled={!openRouterKeyQuery.data}
+									disabled={!openRouterKeyQuery.data || isGenerating}
 									class="text-foreground placeholder:text-muted-foreground/60 max-h-64 min-h-[60px] w-full resize-none !overflow-y-auto bg-transparent text-base leading-6 outline-none disabled:cursor-not-allowed disabled:opacity-50"
-									placeholder="Type your message here... Tag rules with @"
+									placeholder={isGenerating
+										? 'Generating response...'
+										: 'Type your message here... Tag rules with @'}
 									name="message"
 									onkeydown={(e) => {
 										if (e.key === 'Enter' && !e.shiftKey && !popover.open) {
@ -715,14 +770,20 @@
 									<Tooltip placement="top">
 										{#snippet trigger(tooltip)}
 											<button
-												type="submit"
-												class="border-reflect button-reflect hover:bg-primary/90 active:bg-primary text-primary-foreground relative h-9 w-9 rounded-lg p-2 font-semibold shadow transition"
+												type={isGenerating ? 'button' : 'submit'}
+												onclick={isGenerating ? stopGeneration : undefined}
+												disabled={isGenerating ? false : !message.trim()}
+												class="border-reflect button-reflect hover:bg-primary/90 active:bg-primary text-primary-foreground relative h-9 w-9 rounded-lg p-2 font-semibold shadow transition disabled:cursor-not-allowed disabled:opacity-50"
 												{...tooltip.trigger}
 											>
-												<SendIcon class="!size-5" />
+												{#if isGenerating}
+													<StopIcon class="!size-5" />
+												{:else}
+													<SendIcon class="!size-5" />
+												{/if}
 											</button>
 										{/snippet}
-										Send message
+										{isGenerating ? 'Stop generation' : 'Send message'}
 									</Tooltip>
 								</div>
 								<div class="flex flex-col gap-2 pr-2 sm:flex-row sm:items-center">
@ -741,7 +802,7 @@
 											{:else}
 												<ImageIcon class="!size-3" />
 											{/if}
-											<span>Attach image</span>
+											<span class="whitespace-nowrap">Attach image</span>
 										</button>
 									{/if}
 								</div>
--- a/src/routes/chat/model-picker.svelte
+++ b/src/routes/chat/model-picker.svelte
@ -4,6 +4,26 @@
 	import { session } from '$lib/state/session.svelte';
 	import { settings } from '$lib/state/settings.svelte';
 	import { cn } from '$lib/utils/utils';
+	import { Command } from 'bits-ui';
+	import CheckIcon from '~icons/lucide/check';
+	import SearchIcon from '~icons/lucide/search';
+	import ChevronDownIcon from '~icons/lucide/chevron-down';
+	// Company icons from simple-icons
+	import GoogleIcon from '~icons/simple-icons/google';
+	import MetaIcon from '~icons/simple-icons/meta';
+	import MicrosoftIcon from '~icons/simple-icons/microsoft';
+	import OpenaiIcon from '~icons/simple-icons/openai';
+	import XIcon from '~icons/simple-icons/x';
+	// Fallback to lucide icons for companies without simple-icons
+	import RobotIcon from '~icons/lucide/bot';
+	import BrainIcon from '~icons/lucide/brain';
+	import CpuIcon from '~icons/lucide/cpu';
+	import ZapIcon from '~icons/lucide/zap';
+	// Model-specific icons
+	import LogosClaudeIcon from '~icons/logos/claude-icon';
+	import LogosMistralAiIcon from '~icons/logos/mistral-ai-icon';
+	import MaterialIconThemeGeminiAi from '~icons/material-icon-theme/gemini-ai';
+	import { Popover } from 'melt/builders';

 	type Props = {
 		class?: string;
@ -17,15 +37,252 @@

 	const enabledArr = $derived(Object.values(enabledModelsQuery.data ?? {}));

+	// Company icon mapping
+	const companyIcons: Record<string, typeof OpenaiIcon> = {
+		openai: OpenaiIcon,
+		anthropic: BrainIcon,
+		google: GoogleIcon,
+		meta: MetaIcon,
+		mistral: ZapIcon,
+		'x-ai': XIcon,
+		microsoft: MicrosoftIcon,
+		qwen: CpuIcon,
+		deepseek: RobotIcon,
+		cohere: CpuIcon,
+	};
+
+	// Function to get model-specific icon
+	function getModelIcon(modelId: string): typeof LogosClaudeIcon | null {
+		const id = modelId.toLowerCase();
+
+		if (id.includes('claude') || id.includes('anthropic')) return LogosClaudeIcon;
+		if (id.includes('gemini') || id.includes('gemma')) return MaterialIconThemeGeminiAi;
+		if (id.includes('mistral') || id.includes('mixtral')) return LogosMistralAiIcon;
+
+		return null;
+	}
+
+	// Function to extract company from model ID
+	function getCompanyFromModelId(modelId: string): string {
+		const id = modelId.toLowerCase();
+
+		// OpenAI models
+		if (id.includes('gpt') || id.includes('o1') || id.includes('openai')) return 'openai';
+
+		// Anthropic models
+		if (id.includes('claude') || id.includes('anthropic')) return 'anthropic';
+
+		// Google models
+		if (
+			id.includes('gemini') ||
+			id.includes('gemma') ||
+			id.includes('google') ||
+			id.includes('palm')
+		)
+			return 'google';
+
+		// Meta models
+		if (id.includes('llama') || id.includes('meta')) return 'meta';
+
+		// Mistral models
+		if (id.includes('mistral') || id.includes('mixtral')) return 'mistral';
+
+		// xAI models
+		if (id.includes('grok') || id.includes('x-ai')) return 'x-ai';
+
+		// Microsoft models
+		if (id.includes('phi') || id.includes('microsoft')) return 'microsoft';
+
+		// Qwen models
+		if (id.includes('qwen') || id.includes('alibaba')) return 'qwen';
+
+		// DeepSeek models
+		if (id.includes('deepseek')) return 'deepseek';
+
+		// Cohere models
+		if (id.includes('command') || id.includes('cohere')) return 'cohere';
+
+		// Try to extract from model path (e.g., "anthropic/claude-3")
+		const pathParts = modelId.split('/');
+		if (pathParts.length > 1) {
+			const provider = pathParts[0]?.toLowerCase();
+			if (provider && companyIcons[provider]) return provider;
+		}
+
+		return 'other';
+	}
+
+	// Group models by company
+	const groupedModels = $derived.by(() => {
+		console.log('📊 enabledArr:', enabledArr);
+		const groups: Record<string, typeof enabledArr> = {};
+
+		enabledArr.forEach((model) => {
+			const company = getCompanyFromModelId(model.model_id);
+			console.log(`🏢 Model ${model.model_id} -> Company: ${company}`);
+			if (!groups[company]) {
+				groups[company] = [];
+			}
+			groups[company].push(model);
+		});
+
+		console.log('📋 Groups:', groups);
+
+		// Sort companies with known icons first
+		const result = Object.entries(groups).sort(([a], [b]) => {
+			const aHasIcon = companyIcons[a] ? 0 : 1;
+			const bHasIcon = companyIcons[b] ? 0 : 1;
+			return aHasIcon - bHasIcon || a.localeCompare(b);
+		});
+
+		console.log('🎯 Final grouped models:', result);
+		return result;
+	});
+
+	// Find current model details
+	const currentModel = $derived(enabledArr.find((m) => m.model_id === settings.modelId));
+	const currentCompany = $derived(
+		currentModel ? getCompanyFromModelId(currentModel.model_id) : 'other'
+	);
+
 	$effect(() => {
 		if (!enabledArr.find((m) => m.model_id === settings.modelId) && enabledArr.length > 0) {
 			settings.modelId = enabledArr[0]!.model_id;
 		}
 	});
+
+	function selectModel(modelId: string) {
+		settings.modelId = modelId;
+		popover.open = false;
+	}
+
+	let open = $state(true);
+	const popover = new Popover({
+		open: () => open,
+		onOpenChange: (v) => {
+			console.log('📋 popover open:', v);
+			if (v === open) return;
+			open = v;
+			console.log('assigned', v);
+			if (v) return;
+			console.log('attempting to focus');
+			document.getElementById(popover.trigger.id)?.focus();
+		},
+		floatingConfig: {
+			computePosition: { placement: 'top-start' },
+		},
+	});
+
+	// Name splitter. splits -,_,:
+	function splitName(name: string) {
+		return name.split(/[-_,:]/);
+	}
+
+	function getModelTitle(modelId: string) {
+		const sn = splitName(modelId.replace(/^[^/]+\//, ''));
+		return sn[0] + (sn.length > 1 ? ' ' + sn.slice(1).join(' ') : '');
+	}
 </script>

-<select bind:value={settings.modelId} class={cn('border-border border', className)}>
-	{#each enabledArr as model (model._id)}
-		<option value={model.model_id}>{model.model_id}</option>
-	{/each}
-</select>
+{#if enabledArr.length === 0}
+	<!-- Fallback to original select if no models are loaded -->
+	<select bind:value={settings.modelId} class={cn('border-border border', className)}>
+		<option value="">Loading models...</option>
+	</select>
+{:else}
+	<button
+		{...popover.trigger}
+		class={cn(
+			'border-input bg-background ring-offset-background placeholder:text-muted-foreground focus:ring-ring flex w-full items-center justify-between rounded-md border px-3 py-2 text-sm focus:ring-2 focus:ring-offset-2 focus:outline-none disabled:cursor-not-allowed disabled:opacity-50',
+			className
+		)}
+		aria-expanded={open}
+	>
+		<div class="flex items-center gap-2">
+			{#if companyIcons[currentCompany]}
+				{@const IconComponent = companyIcons[currentCompany]}
+				<IconComponent class="size-4" />
+			{/if}
+			<span class="truncate capitalize">
+				{currentModel ? getModelTitle(currentModel.model_id) : 'Select model'}
+			</span>
+		</div>
+		<ChevronDownIcon class="size-4 opacity-50" />
+	</button>
+
+	<div
+		{...popover.content}
+		class="border-border bg-popover mt-1 max-h-200 min-w-80 flex-col overflow-hidden rounded-xl border p-0 backdrop-blur-sm data-[open]:flex"
+	>
+		<Command.Root class="flex h-full flex-col overflow-hidden" columns={4}>
+			<label class="group/label relative flex items-center gap-2 px-4 py-3 text-sm">
+				<SearchIcon class="text-muted-foreground" />
+				<Command.Input
+					class="w-full outline-none"
+					placeholder="Search models..."
+					{@attach (node) => {
+						if (popover.open) {
+							node.focus();
+						}
+						return () => {
+							node.value = '';
+						};
+					}}
+				/>
+				<div
+					class="border-border/50 group-focus-within/label:border-foreground/30 absolute inset-x-2 bottom-0 h-1 border-b"
+					aria-hidden="true"
+				></div>
+			</label>
+			<Command.List class="overflow-y-auto">
+				<Command.Viewport>
+					<Command.Empty class="text-muted-foreground p-4 text-sm">
+						No models available. Enable some models in the account settings.
+					</Command.Empty>
+					{#each groupedModels as [company, models] (company)}
+						<Command.Group class="space-y-2">
+							<Command.GroupHeading
+								class="text-heading/75 flex items-center gap-2 px-3 pt-3 pb-1 text-xs font-semibold tracking-wide capitalize"
+							>
+								{company}
+							</Command.GroupHeading>
+							<Command.GroupItems class="grid grid-cols-4 gap-3 px-3 pb-3">
+								{#each models as model (model._id)}
+									{@const isSelected = settings.modelId === model.model_id}
+									{@const sn = splitName(model.model_id.replace(/^[^/]+\//, ''))}
+									<Command.Item
+										value={model.model_id}
+										onSelect={() => selectModel(model.model_id)}
+										class={cn(
+											'border-border flex h-40 w-32 flex-col items-center justify-center rounded-lg border p-2',
+											'select-none',
+											'data-selected:bg-accent/50 data-selected:text-accent-foreground',
+											isSelected && 'border-reflect border-none',
+											'scroll-m-10'
+										)}
+									>
+										{#if getModelIcon(model.model_id)}
+											{@const ModelIcon = getModelIcon(model.model_id)}
+											<ModelIcon class="size-6 shrink-0" />
+										{:else if companyIcons[getCompanyFromModelId(model.model_id)]}
+											{@const CompanyIcon = companyIcons[getCompanyFromModelId(model.model_id)]}
+											<CompanyIcon class="size-6 shrink-0" />
+										{/if}
+										<p
+											class="font-fake-proxima mt-2 text-center leading-tight font-bold capitalize"
+										>
+											{sn[0]}
+										</p>
+										<p class="mt-0 text-center text-xs leading-tight font-medium">
+											{sn.slice(1).join(' ')}
+										</p>
+									</Command.Item>
+								{/each}
+							</Command.GroupItems>
+						</Command.Group>
+					{/each}
+				</Command.Viewport>
+			</Command.List>
+		</Command.Root>
+	</div>
+{/if}