From e6d6ca5c56987649cecabf039637509ccd3575eb Mon Sep 17 00:00:00 2001
From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com>
Date: Tue, 8 Oct 2024 02:43:41 -0400
Subject: [PATCH] Fixes

---
 .../settings/OpenRouterModelPicker.tsx          |  1 +
 .../src/context/ExtensionStateContext.tsx       | 17 ++++++++++++++---
 webview-ui/src/utils/validate.ts                |  5 +++--
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/webview-ui/src/components/settings/OpenRouterModelPicker.tsx b/webview-ui/src/components/settings/OpenRouterModelPicker.tsx
index 5f7f1bf..bf9f247 100644
--- a/webview-ui/src/components/settings/OpenRouterModelPicker.tsx
+++ b/webview-ui/src/components/settings/OpenRouterModelPicker.tsx
@@ -21,6 +21,7 @@ const OpenRouterModelPicker: React.FC = () => {
 	const dropdownListRef = useRef<HTMLDivElement>(null)
 
 	const handleModelChange = (newModelId: string) => {
+		// could be setting invalid model id/undefined info but validation will catch it
 		setApiConfiguration({
 			...apiConfiguration,
 			openRouterModelId: newModelId,
diff --git a/webview-ui/src/context/ExtensionStateContext.tsx b/webview-ui/src/context/ExtensionStateContext.tsx
index 8381ee8..16206ae 100644
--- a/webview-ui/src/context/ExtensionStateContext.tsx
+++ b/webview-ui/src/context/ExtensionStateContext.tsx
@@ -1,7 +1,12 @@
 import React, { createContext, useCallback, useContext, useEffect, useState } from "react"
 import { useEvent } from "react-use"
 import { ExtensionMessage, ExtensionState } from "../../../src/shared/ExtensionMessage"
-import { ApiConfiguration, ModelInfo } from "../../../src/shared/api"
+import {
+	ApiConfiguration,
+	ModelInfo,
+	openRouterDefaultModelId,
+	openRouterDefaultModelInfo,
+} from "../../../src/shared/api"
 import { vscode } from "../utils/vscode"
 import { convertTextMateToHljs } from "../utils/textMateToHljs"
 import { findLastIndex } from "../../../src/shared/array"
@@ -31,7 +36,9 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
 	const [showWelcome, setShowWelcome] = useState(false)
 	const [theme, setTheme] = useState<any>(undefined)
 	const [filePaths, setFilePaths] = useState<string[]>([])
-	const [openRouterModels, setOpenRouterModels] = useState<Record<string, ModelInfo>>({})
+	const [openRouterModels, setOpenRouterModels] = useState<Record<string, ModelInfo>>({
+		[openRouterDefaultModelId]: openRouterDefaultModelInfo,
+	})
 
 	const handleMessage = useCallback((event: MessageEvent) => {
 		const message: ExtensionMessage = event.data
@@ -80,7 +87,11 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
 				break
 			}
 			case "openRouterModels": {
-				setOpenRouterModels(message.openRouterModels ?? {})
+				const updatedModels = message.openRouterModels ?? {}
+				setOpenRouterModels({
+					[openRouterDefaultModelId]: openRouterDefaultModelInfo, // in case the extension sent a model list without the default model
+					...updatedModels,
+				})
 				break
 			}
 		}
diff --git a/webview-ui/src/utils/validate.ts b/webview-ui/src/utils/validate.ts
index 92e9e69..dd3ec58 100644
--- a/webview-ui/src/utils/validate.ts
+++ b/webview-ui/src/utils/validate.ts
@@ -1,4 +1,4 @@
-import { ApiConfiguration } from "../../../src/shared/api"
+import { ApiConfiguration, openRouterDefaultModelId } from "../../../src/shared/api"
 import { ModelInfo } from "../../../src/shared/api"
 export function validateApiConfiguration(apiConfiguration?: ApiConfiguration): string | undefined {
 	if (apiConfiguration) {
@@ -59,11 +59,12 @@ export function validateModelId(
 	if (apiConfiguration) {
 		switch (apiConfiguration.apiProvider) {
 			case "openrouter":
-				const modelId = apiConfiguration.openRouterModelId
+				const modelId = apiConfiguration.openRouterModelId ?? openRouterDefaultModelId // in case the user hasn't changed the model id, it will be undefined by default
 				if (!modelId) {
 					return "You must provide a model ID."
 				}
 				if (openRouterModels && !Object.keys(openRouterModels).includes(modelId)) {
+					// even if the model list endpoint failed, extensionstatecontext will always have the default model info
 					return "The model ID you provided is not available. Please choose a different model."
 				}
 				break