feat: add beta models gating for axon-code-2-pro

- Add backend API endpoint integration for beta model availability
- Implement betaModelsEnabled state in ExtensionStateContext
- Filter axon-code-2-pro from model list when beta not enabled
- Add message types for fetchBetaModelsRequest and betaModelsResponse
- Fix openrouter baseURL configuration (remove hardcoded localhost)
- Fix indentation in openrouter.ts CompletionUsage interface

Author: code-crusher

src/api/providers/openrouter.ts

@@ -143,7 +143,7 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH
 		const baseURL = this.options.openRouterBaseUrl || "https://api2.matterai.so/v1/web"
 		const apiKey = this.options.openRouterApiKey ?? "not-provided"
 
-		this.client = new OpenAI({ baseURL: "http://localhost:4064/v1/web", apiKey, defaultHeaders: DEFAULT_HEADERS })
+		this.client = new OpenAI({ baseURL, apiKey, defaultHeaders: DEFAULT_HEADERS })
 	}
 
 	// kilocode_change start

src/core/webview/webviewMessageHandler.ts

@@ -3661,6 +3661,42 @@ ${comment.suggestion}
 				})
 			}
 			break
+		case "fetchBetaModelsRequest": // New handler for beta models
+			try {
+				const { apiConfiguration } = await provider.getState()
+				const kilocodeToken = apiConfiguration?.kilocodeToken
+
+				if (!kilocodeToken) {
+					provider.log("KiloCode token not found in extension state for beta models.")
+					provider.postMessageToWebview({
+						type: "betaModelsResponse",
+						payload: { success: false, error: "KiloCode API token not configured." },
+					})
+					break
+				}
+
+				const headers: Record<string, string> = {
+					Authorization: `Bearer ${kilocodeToken}`,
+					"Content-Type": "application/json",
+				}
+
+				const url = "https://api.matterai.so/axoncode/betamodels"
+				const response = await axios.get<{ enabled: boolean }>(url, { headers })
+
+				provider.postMessageToWebview({
+					type: "betaModelsResponse",
+					payload: { success: true, enabled: response.data.enabled },
+				})
+			} catch (error: any) {
+				const errorMessage =
+					error.response?.data?.message || error.message || "Failed to fetch beta models data from backend."
+				provider.log(`Error fetching beta models data: ${errorMessage}`)
+				provider.postMessageToWebview({
+					type: "betaModelsResponse",
+					payload: { success: false, error: errorMessage },
+				})
+			}
+			break
 
 		case "fetchMcpMarketplace": {
 			await provider.fetchMcpMarketplace(message.bool)

src/shared/ExtensionMessage.ts

@@ -25,6 +25,7 @@ import { Mode } from "./modes"
 // kilocode_change start
 import {
 	BalanceDataResponsePayload,
+	BetaModelsResponsePayload,
 	ProfileDataResponsePayload,
 	TaskHistoryResponsePayload,
 	TasksByIdResponsePayload,
@@ -130,6 +131,7 @@ export interface ExtensionMessage {
 		| "profileDataResponse" // kilocode_change
 		| "balanceDataResponse" // kilocode_change
 		| "updateProfileData" // kilocode_change
+		| "betaModelsResponse" // kilocode_change
 		| "authenticatedUser"
 		| "condenseTaskContextResponse"
 		| "singleRouterModelFetchResponse"
@@ -165,6 +167,7 @@ export interface ExtensionMessage {
 	payload?:
 		| ProfileDataResponsePayload
 		| BalanceDataResponsePayload
+		| BetaModelsResponsePayload
 		| TasksByIdResponsePayload
 		| TaskHistoryResponsePayload
 		| ImplementPlanPayload

src/shared/WebviewMessage.ts

@@ -229,6 +229,8 @@ export interface WebviewMessage {
 		| "shopBuyCredits" // kilocode_change
 		| "balanceDataResponse" // kilocode_change
 		| "updateProfileData" // kilocode_change
+		| "fetchBetaModelsRequest" // kilocode_change
+		| "betaModelsResponse" // kilocode_change
 		| "condense" // kilocode_change
 		| "toggleWorkflow" // kilocode_change
 		| "refreshRules" // kilocode_change
@@ -472,6 +474,13 @@ export interface BalanceDataResponsePayload {
 	error?: string
 }
 
+export interface BetaModelsResponsePayload {
+	// Payload for beta models availability
+	success: boolean
+	enabled?: boolean
+	error?: string
+}
+
 export interface SeeNewChangesPayload {
 	commitRange: CommitRange
 }
@@ -593,6 +602,7 @@ export type WebViewMessagePayload =
 	// kilocode_change start
 	| ProfileDataResponsePayload
 	| BalanceDataResponsePayload
+	| BetaModelsResponsePayload
 	| SeeNewChangesPayload
 	| TasksByIdRequestPayload
 	| TaskHistoryRequestPayload

webview-ui/src/App.tsx

@@ -251,6 +251,13 @@ 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
+	useEffect(() => {
+		if (didHydrateState && apiConfiguration?.kilocodeToken) {
+			vscode.postMessage({ type: "fetchBetaModelsRequest" })
+		}
+	}, [didHydrateState, apiConfiguration?.kilocodeToken])
+
 	// Initialize source map support for better error reporting
 	useEffect(() => {
 		// Initialize source maps for better error reporting in production

webview-ui/src/components/kilocode/common/OrganizationSelector.tsx

@@ -77,6 +77,9 @@ export const OrganizationSelector = ({ className, showLabel = false }: { classNa
 					apiKey: apiConfiguration?.kilocodeToken,
 				},
 			})
+			vscode.postMessage({
+				type: "fetchBetaModelsRequest",
+			})
 		} else {
 			vscode.postMessage({
 				type: "upsertApiConfiguration",
@@ -89,6 +92,9 @@ export const OrganizationSelector = ({ className, showLabel = false }: { classNa
 			vscode.postMessage({
 				type: "fetchBalanceDataRequest",
 			})
+			vscode.postMessage({
+				type: "fetchBetaModelsRequest",
+			})
 		}
 	}
 

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

@@ -1,3 +1,4 @@
+import { useMemo } from "react"
 import {
 	type ProviderName,
 	type ProviderSettings,
@@ -64,10 +65,12 @@ export const getModelsByProvider = ({
 	provider,
 	routerModels,
 	kilocodeDefaultModel,
+	_betaModelsEnabled,
 }: {
 	provider: ProviderName
 	routerModels: RouterModels
 	kilocodeDefaultModel: string
+	_betaModelsEnabled?: boolean
 }): { models: ModelRecord; defaultModel: string } => {
 	switch (provider) {
 		case "openrouter": {
@@ -296,7 +299,7 @@ export const getModelsByProvider = ({
 export const useProviderModels = (apiConfiguration?: ProviderSettings) => {
 	const provider = apiConfiguration?.apiProvider || "anthropic"
 
-	const { kilocodeDefaultModel } = useExtensionState()
+	const { kilocodeDefaultModel, betaModelsEnabled } = useExtensionState()
 
 	const routerModels = useRouterModels({
 		openRouterBaseUrl: apiConfiguration?.openRouterBaseUrl,
@@ -315,12 +318,24 @@ export const useProviderModels = (apiConfiguration?: ProviderSettings) => {
 					provider,
 					routerModels: routerModels.data,
 					kilocodeDefaultModel,
+					_betaModelsEnabled: betaModelsEnabled,
 				})
 			: FALLBACK_MODELS
 
+	// kilocode_change start: Filter out axon-code-2-pro if beta models are not enabled
+	const filteredModels = useMemo(() => {
+		if (betaModelsEnabled) {
+			return models
+		}
+		// Filter out axon-code-2-pro when beta models are not enabled
+		const { "axon-code-2-pro": _, ...rest } = models as ModelRecord
+		return rest
+	}, [models, betaModelsEnabled])
+	// kilocode_change end
+
 	return {
 		provider,
-		providerModels: models as ModelRecord,
+		providerModels: filteredModels as ModelRecord,
 		providerDefaultModel: defaultModel,
 		isLoading: routerModels.isLoading,
 		isError: routerModels.isError,

webview-ui/src/components/kilocode/profile/ProfileView.tsx

@@ -34,6 +34,7 @@ const ProfileView: React.FC<ProfileViewProps> = ({ onDone }) => {
 	useEffect(() => {
 		vscode.postMessage({ type: "fetchProfileDataRequest" })
 		vscode.postMessage({ type: "fetchBalanceDataRequest" })
+		vscode.postMessage({ type: "fetchBetaModelsRequest" })
 	}, [apiConfiguration?.kilocodeToken, organizationId])
 
 	useEffect(() => {
@@ -64,6 +65,9 @@ const ProfileView: React.FC<ProfileViewProps> = ({ onDone }) => {
 				vscode.postMessage({
 					type: "fetchBalanceDataRequest",
 				})
+				vscode.postMessage({
+					type: "fetchBetaModelsRequest",
+				})
 			}
 		}
 

webview-ui/src/context/ExtensionStateContext.tsx

@@ -205,6 +205,7 @@ export interface ExtensionStateContextType extends ExtensionState {
 		currentTokens: number
 		maxTokens: number
 	} // kilocode_change: Track context window usage
+	betaModelsEnabled?: boolean // kilocode_change: Beta models availability
 }
 
 export const ExtensionStateContext = createContext<ExtensionStateContextType | undefined>(undefined)
@@ -406,6 +407,16 @@ 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,
+						}))
+					}
+					break
+				}
 				case "action": {
 					if (message.action === "toggleAutoApprove") {
 						// Toggle the auto-approval state