fix: BYOK

internal/api/chat/create_conversation_message_stream_v2.go

@@ -2,6 +2,7 @@ package chat
 
 import (
 	"context"
+	"fmt"
 	"paperdebugger/internal/api/mapper"
 	"paperdebugger/internal/libs/contextutil"
 	"paperdebugger/internal/libs/shared"
@@ -281,14 +282,21 @@ func (s *ChatServerV2) CreateConversationMessageStream(
 	var llmProvider *models.LLMProviderConfig
 	var customModel *models.CustomModel
 	customModel = nil
-	for i := range settings.CustomModels {
-		if settings.CustomModels[i].Slug == modelSlug {
-			customModel = &settings.CustomModels[i]
+
+	customModelID := req.GetCustomModelId()
+	if customModelID != "" {
+		for i := range settings.CustomModels {
+			if settings.CustomModels[i].Id.Hex() == customModelID {
+				customModel = &settings.CustomModels[i]
+				break
+			}
 		}
+		if customModel == nil {
+			return s.sendStreamError(stream, fmt.Errorf("custom model not found: %q", customModelID))
+		}
+		modelSlug = customModel.Slug
 	}
 
-	// Usage is the same as ChatCompletion, just passing the stream parameter
-
 	if customModel == nil {
 		// User did not specify API key for this model
 		llmProvider = &models.LLMProviderConfig{

internal/api/chat/list_supported_models_v2.go

@@ -222,7 +222,9 @@ func (s *ChatServerV2) ListSupportedModels(
 	var models []*chatv2.SupportedModel
 
 	for _, model := range settings.CustomModels {
+		modelID := model.Id.Hex()
 		models = append(models, &chatv2.SupportedModel{
+			Id:           &modelID,
 			Name:         model.Name,
 			Slug:         model.Slug,
 			TotalContext: int64(model.ContextWindow),

proto/chat/v2/chat.proto

@@ -137,6 +137,7 @@ message SupportedModel {
   bool disabled = 7; // If true, the model is disabled and cannot be used
   optional string disabled_reason = 8; // The reason why the model is disabled
   bool is_custom = 9;
+  optional string id = 10; // Custom model unique ID (empty for built-in models)
 }
 
 message ListSupportedModelsRequest {
@@ -223,6 +224,7 @@ message CreateConversationMessageStreamRequest {
   optional string user_selected_text = 5;
   optional ConversationType conversation_type = 6;
   optional string surrounding = 8;
+  optional string custom_model_id = 9; // Selected custom model ID
 }
 
 // Response for streaming a message within an existing conversation

webapp/_webapp/src/hooks/useLanguageModels.ts

@@ -5,6 +5,7 @@ import { useListSupportedModelsQuery } from "../query";
 import { useConversationUiStore } from "../stores/conversation/conversation-ui-store";
 
 export type Model = {
+  id?: string;
   name: string;
   slug: string;
   provider: string;
@@ -39,6 +40,7 @@ const fallbackModels: Model[] = [
 ];
 
 const mapSupportedModelToModel = (supportedModel: SupportedModel): Model => ({
+  id: supportedModel.id || undefined,
   name: supportedModel.name,
   slug: supportedModel.slug,
   provider: extractProvider(supportedModel.slug),
@@ -53,7 +55,7 @@ const mapSupportedModelToModel = (supportedModel: SupportedModel): Model => ({
 
 export const useLanguageModels = () => {
   const { currentConversation, setCurrentConversation } = useConversationStore();
-  const { setLastUsedModelSlug } = useConversationUiStore();
+  const { lastUsedCustomModelId, setLastUsedModelSlug, setLastUsedCustomModelId } = useConversationUiStore();
   const { data: supportedModelsResponse } = useListSupportedModelsQuery();
 
   const models: Model[] = useMemo(() => {
@@ -64,19 +66,25 @@ export const useLanguageModels = () => {
   }, [supportedModelsResponse]);
 
   const currentModel = useMemo(() => {
-    const model = models.find((m) => m.slug === currentConversation.modelSlug);
+    if (lastUsedCustomModelId) {
+      const customModel = models.find((m) => m.isCustom && m.id === lastUsedCustomModelId);
+      if (customModel) return customModel;
+    }
+
+    const model = models.find((m) => !m.isCustom && m.slug === currentConversation.modelSlug);
     return model || models[0];
-  }, [models, currentConversation.modelSlug]);
+  }, [models, currentConversation.modelSlug, lastUsedCustomModelId]);
 
   const setModel = useCallback(
     (model: Model) => {
       setLastUsedModelSlug(model.slug);
+      setLastUsedCustomModelId(model.isCustom ? (model.id ?? "") : "");
       setCurrentConversation({
         ...currentConversation,
         modelSlug: model.slug,
       });
     },
-    [setCurrentConversation, currentConversation, setLastUsedModelSlug],
+    [setCurrentConversation, currentConversation, setLastUsedModelSlug, setLastUsedCustomModelId],
   );
 
   return { models, currentModel, setModel };

webapp/_webapp/src/hooks/useSendMessageStream.ts

@@ -41,6 +41,7 @@ import { useAuthStore } from "../stores/auth-store";
 import { useDevtoolStore } from "../stores/devtool-store";
 import { useSelectionStore } from "../stores/selection-store";
 import { useSettingStore } from "../stores/setting-store";
+import { useConversationUiStore } from "../stores/conversation/conversation-ui-store";
 import { useSync } from "./useSync";
 import { useAdapter } from "../adapters";
 import { getProjectId } from "../libs/helpers";
@@ -86,6 +87,7 @@ export function useSendMessageStream(): UseSendMessageStreamResult {
   const surroundingText = useSelectionStore((s) => s.surroundingText);
   const alwaysSyncProject = useDevtoolStore((s) => s.alwaysSyncProject);
   const conversationMode = useSettingStore((s) => s.conversationMode);
+  const lastUsedCustomModelId = useConversationUiStore((s) => s.lastUsedCustomModelId);
 
   /**
    * Add the user message to the streaming state.
@@ -165,6 +167,7 @@ export function useSendMessageStream(): UseSendMessageStreamResult {
         projectId,
         conversationId: currentConversation.id,
         modelSlug: currentConversation.modelSlug,
+        customModelId: lastUsedCustomModelId || undefined,
         surroundingText: surroundingText ?? undefined,
         conversationMode: conversationMode === "debug" ? "debug" : "default",
       };
@@ -251,6 +254,7 @@ export function useSendMessageStream(): UseSendMessageStreamResult {
       alwaysSyncProject,
       conversationMode,
       surroundingText,
+      lastUsedCustomModelId,
       addUserMessageToStream,
       truncateConversationIfEditing,
     ],