fix beta model access (#47)

Author: code-crusher

src/core/webview/webviewMessageHandler.ts

@@ -3661,6 +3661,7 @@ ${comment.suggestion}
 				})
 			}
 			break
+
 		case "fetchBetaModelsRequest": // New handler for beta models
 			try {
 				const { apiConfiguration } = await provider.getState()
@@ -3681,11 +3682,22 @@ ${comment.suggestion}
 				}
 
 				const url = "https://api.matterai.so/axoncode/betamodels"
-				const response = await axios.get<{ enabled: boolean }>(url, { headers })
+				const response = await axios.get(url, { headers, timeout: 10000 })
+
+				// Handle different response formats
+				let enabled = false
+				if (response.data === true) {
+					// API returns boolean true when beta models are enabled
+					enabled = true
+				} else if (response.data && typeof response.data === "object") {
+					// Try different possible field names
+					enabled =
+						response.data.enabled ?? response.data.betaModelsEnabled ?? response.data.isEnabled ?? false
+				}
 
 				provider.postMessageToWebview({
 					type: "betaModelsResponse",
-					payload: { success: true, enabled: response.data.enabled },
+					payload: { success: true, enabled },
 				})
 			} catch (error: any) {
 				const errorMessage =

src/package.json

@@ -3,7 +3,7 @@
 	"displayName": "%extension.displayName%",
 	"description": "%extension.description%",
 	"publisher": "matterai",
-	"version": "5.3.2",
+	"version": "5.3.3",
 	"icon": "assets/icons/matterai-ic.png",
 	"galleryBanner": {
 		"color": "#FFFFFF",

src/shared/ExtensionMessage.ts

@@ -381,6 +381,7 @@ export type ExtensionState = Pick<
 	| "reasoningBlockCollapsed"
 	| "codeReviewSettings"
 > & {
+	betaModelsEnabled?: boolean // kilocode_change: Beta models availability
 	version: string
 	clineMessages: ClineMessage[]
 	currentTaskItem?: HistoryItem

webview-ui/src/App.tsx

@@ -251,13 +251,29 @@ const App = () => {
 	// Tell the extension that we are ready to receive messages.
 	useEffect(() => vscode.postMessage({ type: "webviewDidLaunch" }), [])
 
-	// kilocode_change start: Fetch beta models availability on mount
+	// kilocode_change start: Fetch beta models availability on mount and when API config changes
 	useEffect(() => {
 		if (didHydrateState && apiConfiguration?.kilocodeToken) {
 			vscode.postMessage({ type: "fetchBetaModelsRequest" })
 		}
 	}, [didHydrateState, apiConfiguration?.kilocodeToken])
 
+	// Additional effect to refresh beta models when switching between API configurations
+	useEffect(() => {
+		if (didHydrateState && apiConfiguration?.kilocodeToken) {
+			// Add a small delay to ensure the API configuration is fully updated
+			const timeoutId = setTimeout(() => {
+				vscode.postMessage({ type: "fetchBetaModelsRequest" })
+			}, 500)
+			return () => clearTimeout(timeoutId)
+		}
+	}, [
+		didHydrateState,
+		apiConfiguration?.kilocodeToken,
+		apiConfiguration?.apiProvider,
+		apiConfiguration?.kilocodeOrganizationId,
+	])
+
 	// Initialize source map support for better error reporting
 	useEffect(() => {
 		// Initialize source maps for better error reporting in production

webview-ui/src/components/kilocode/hooks/useProviderModels.ts

@@ -309,6 +309,7 @@ export const useProviderModels = (apiConfiguration?: ProviderSettings) => {
 		chutesApiKey: apiConfiguration?.chutesApiKey,
 		geminiApiKey: apiConfiguration?.geminiApiKey,
 		googleGeminiBaseUrl: apiConfiguration?.googleGeminiBaseUrl,
+		betaModelsEnabled, // kilocode_change: Beta models availability
 		// kilocode_change end
 	})
 
@@ -324,12 +325,12 @@ export const useProviderModels = (apiConfiguration?: ProviderSettings) => {
 
 	// kilocode_change start: Filter out axon-code-2-pro if beta models are not enabled
 	const filteredModels = useMemo(() => {
-		if (betaModelsEnabled) {
-			return models
+		if (!betaModelsEnabled) {
+			// Filter out axon-code-2-pro when beta models are not enabled
+			const { "axon-code-2-pro": _, ...rest } = models as ModelRecord
+			return rest
 		}
-		// Filter out axon-code-2-pro when beta models are not enabled
-		const { "axon-code-2-pro": _, ...rest } = models as ModelRecord
-		return rest
+		return models
 	}, [models, betaModelsEnabled])
 	// kilocode_change end
 

webview-ui/src/components/kilocode/settings/providers/KiloCode.tsx

@@ -8,6 +8,9 @@ import { KiloCodeWrapperProperties } from "../../../../../../src/shared/kilocode
 import { ModelPicker } from "../../../settings/ModelPicker"
 import { OrganizationSelector } from "../../common/OrganizationSelector"
 import { getKiloCodeBackendSignInUrl } from "../../helpers"
+import { useExtensionState } from "@/context/ExtensionStateContext"
+import { useMemo } from "react"
+import type { ModelRecord } from "@roo/api"
 
 type KiloCodeProps = {
 	apiConfiguration: ProviderSettings
@@ -35,6 +38,18 @@ export const KiloCode = ({
 	kilocodeDefaultModel,
 }: KiloCodeProps) => {
 	const { t } = useAppTranslation()
+	const { betaModelsEnabled } = useExtensionState()
+
+	// Filter out axon-code-2-pro if beta models are not enabled
+	const filteredModels = useMemo(() => {
+		const models = routerModels?.["kilocode-openrouter"] ?? {}
+		if (!betaModelsEnabled) {
+			// Filter out axon-code-2-pro when beta models are not enabled
+			const { "axon-code-2-pro": _, ...rest } = models as ModelRecord
+			return rest
+		}
+		return models
+	}, [routerModels, betaModelsEnabled])
 
 	// const handleInputChange = useCallback(
 	// 	<K extends keyof ProviderSettings, E>(
@@ -111,7 +126,7 @@ export const KiloCode = ({
 				apiConfiguration={apiConfiguration}
 				setApiConfigurationField={setApiConfigurationField}
 				defaultModelId={kilocodeDefaultModel}
-				models={routerModels?.["kilocode-openrouter"] ?? {}}
+				models={filteredModels}
 				modelIdKey="kilocodeModel"
 				serviceName="Axon Code"
 				serviceUrl={getAppUrl()}

webview-ui/src/components/ui/hooks/useRouterModels.ts

@@ -46,6 +46,7 @@ type RouterModelsQueryKey = {
 	geminiApiKey?: string
 	googleGeminiBaseUrl?: string
 	chutesApiKey?: string
+	betaModelsEnabled?: boolean // kilocode_change: Beta models availability
 	// Requesty, Unbound, etc should perhaps also be here, but they already have their own hacks for reloading
 }
 

webview-ui/src/context/ExtensionStateContext.tsx

@@ -335,6 +335,7 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
 		openRouterImageApiKey: "",
 		kiloCodeImageApiKey: "",
 		openRouterImageGenerationSelectedModel: "",
+		betaModelsEnabled: false, // kilocode_change: Default to false
 	})
 
 	const [didHydrateState, setDidHydrateState] = useState(false)
@@ -408,13 +409,20 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode
 					break
 				}
 				case "betaModelsResponse": {
-					if (message.payload && "enabled" in message.payload) {
-						const betaPayload = message.payload as { enabled?: boolean }
-						setState((prevState) => ({
-							...prevState,
-							betaModelsEnabled: betaPayload.enabled,
-						}))
-					}
+					const payload = message.payload as { enabled?: boolean; success?: boolean }
+					const enabled = payload?.enabled ?? payload?.success ?? false
+
+					setState((prevState) => ({
+						...prevState,
+						betaModelsEnabled: !!enabled,
+					}))
+
+					// Flush router models cache to force refresh when beta models status changes
+					vscode.postMessage({ type: "flushRouterModels", text: "kilocode-openrouter" })
+					// Add a small delay to ensure cache flush completes before requesting new models
+					setTimeout(() => {
+						vscode.postMessage({ type: "requestRouterModels" })
+					}, 100)
 					break
 				}
 				case "action": {

webview-ui/src/utils/prettyModelName.ts

@@ -3,7 +3,7 @@ const AXON_MODEL_CREDITS: Record<string, string> = {
 	"axon-mini": "(0.5x)",
 	"axon-code": "(0.8x)",
 	"axon-code-2": "(1x)",
-	"axon-code-2-pro": "(2x)",
+	"axon-code-2-pro": "(1.5x)",
 }
 
 export const prettyModelName = (modelId: string): string => {