Merge pull request #2081 from Kilo-Org/beatlevic/fetch-openrouter-models-for-org

Add support for retrieving organization filtered models

Author: kevinvandijk

src/api/providers/fetchers/modelCache.ts

@@ -81,10 +81,12 @@ export const getModels = async (options: GetModelsOptions): Promise<ModelRecord>
 			// kilocode_change start
 			case "kilocode-openrouter":
 				models = await getOpenRouterModels({
-					openRouterBaseUrl: getKiloBaseUriFromToken(options.kilocodeToken ?? "") + "/api/openrouter",
-					headers: options.kilocodeOrganizationId
-						? { "X-KiloCode-OrganizationId": options.kilocodeOrganizationId }
-						: undefined,
+					openRouterBaseUrl:
+						getKiloBaseUriFromToken(options.kilocodeToken ?? "") +
+						(options.kilocodeOrganizationId
+							? `/api/organizations/${options.kilocodeOrganizationId}`
+							: "/api/openrouter"),
+					headers: options.kilocodeToken ? { Authorization: `Bearer ${options.kilocodeToken}` } : undefined,
 				})
 				break
 			case "cerebras":

src/core/webview/webviewMessageHandler.ts

@@ -1662,6 +1662,18 @@ export const webviewMessageHandler = async (
 					if (currentConfig.kilocodeToken !== message.apiConfiguration.kilocodeToken) {
 						configToSave = { ...message.apiConfiguration, kilocodeOrganizationId: undefined }
 					}
+					if (currentConfig.kilocodeOrganizationId !== message.apiConfiguration.kilocodeOrganizationId) {
+						await flushModels("kilocode-openrouter")
+						const models = await getModels({
+							provider: "kilocode-openrouter",
+							kilocodeOrganizationId: message.apiConfiguration.kilocodeOrganizationId,
+							kilocodeToken: message.apiConfiguration.kilocodeToken,
+						})
+						provider.postMessageToWebview({
+							type: "routerModels",
+							routerModels: { "kilocode-openrouter": models } as Record<RouterName, ModelRecord>,
+						})
+					}
 				} catch (error) {
 					// Config might not exist yet, that's fine
 				}

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

@@ -12,11 +12,27 @@ export const OrganizationSelector = ({ className, showLabel = false }: { classNa
 	const selectedOrg = organizations.find((o) => o.id === apiConfiguration?.kilocodeOrganizationId)
 	const containerRef = useRef<HTMLDivElement>(null)
 
+	const handleMessage = (event: MessageEvent<WebviewMessage>) => {
+		const message = event.data
+		if (message.type === "profileDataResponse") {
+			const payload = message.payload as ProfileDataResponsePayload
+			if (payload.success) {
+				setOrganizations(payload.data?.organizations ?? [])
+			} else {
+				console.error("Error fetching profile organizations data:", payload.error)
+				setOrganizations([])
+			}
+		} else if (message.type === "updateProfileData") {
+			vscode.postMessage({
+				type: "fetchProfileDataRequest",
+			})
+		}
+	}
+
 	useEffect(() => {
 		const onKeyDown = (e: KeyboardEvent) => {
 			if (e.key === "Escape") setIsOpen(false)
 		}
-		window.addEventListener("keydown", onKeyDown)
 
 		const onMouseDown = (e: MouseEvent) => {
 			if (!containerRef.current) return
@@ -25,11 +41,14 @@ export const OrganizationSelector = ({ className, showLabel = false }: { classNa
 			}
 		}
 
+		window.addEventListener("keydown", onKeyDown)
 		window.addEventListener("mousedown", onMouseDown)
+		window.addEventListener("message", handleMessage)
 
 		return () => {
 			window.removeEventListener("keydown", onKeyDown)
 			window.removeEventListener("mousedown", onMouseDown)
+			window.removeEventListener("message", handleMessage)
 		}
 	}, [])
 
@@ -41,29 +60,6 @@ export const OrganizationSelector = ({ className, showLabel = false }: { classNa
 		})
 	}, [apiConfiguration?.kilocodeToken])
 
-	useEffect(() => {
-		const handleMessage = (event: MessageEvent<WebviewMessage>) => {
-			const message = event.data
-			if (message.type === "profileDataResponse") {
-				const payload = message.payload as ProfileDataResponsePayload
-				if (payload.success) {
-					setOrganizations(payload.data?.organizations ?? [])
-				} else {
-					console.error("Error fetching profile organizations data:", payload.error)
-					setOrganizations([])
-				}
-			} else if (message.type === "updateProfileData") {
-				vscode.postMessage({
-					type: "fetchProfileDataRequest",
-				})
-			}
-		}
-
-		window.addEventListener("message", handleMessage)
-
-		return () => window.removeEventListener("message", handleMessage)
-	}, [apiConfiguration?.kilocodeToken])
-
 	const setSelectedOrganization = (organization: UserOrganizationWithApiKey | null) => {
 		if (organization === null) {
 			// Switch back to personal account - clear organization token

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

@@ -254,6 +254,7 @@ export const useProviderModels = (apiConfiguration?: ProviderSettings) => {
 	const routerModels = useRouterModels({
 		openRouterBaseUrl: apiConfiguration?.openRouterBaseUrl,
 		openRouterApiKey: apiConfiguration?.apiKey,
+		kilocodeOrganizationId: apiConfiguration?.kilocodeOrganizationId ?? "personal",
 	})
 
 	const { models, defaultModel } =

webview-ui/src/components/settings/ApiOptions.tsx

@@ -203,6 +203,7 @@ const ApiOptions = ({
 	const { data: routerModels, refetch: refetchRouterModels } = useRouterModels({
 		openRouterBaseUrl: apiConfiguration?.openRouterBaseUrl,
 		openRouterApiKey: apiConfiguration?.openRouterApiKey,
+		kilocodeOrganizationId: apiConfiguration?.kilocodeOrganizationId ?? "personal",
 	})
 
 	//const { data: openRouterModelProviders } = useOpenRouterModelProviders(

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

@@ -137,15 +137,16 @@ type UseOpenRouterModelProvidersOptions = Omit<
 	"queryKey" | "queryFn"
 >
 
-// kilocode_change start: baseUrl, apiKey
+// kilocode_change start: baseUrl, apiKey, organizationId
 export const useOpenRouterModelProviders = (
 	modelId?: string,
 	baseUrl?: string,
 	apiKey?: string,
+	organizationId?: string,
 	options?: UseOpenRouterModelProvidersOptions,
 ) =>
 	useQuery<Record<string, OpenRouterModelProvider>>({
-		queryKey: ["openrouter-model-providers", modelId, baseUrl, apiKey],
+		queryKey: ["openrouter-model-providers", modelId, baseUrl, apiKey, organizationId],
 		queryFn: () => (modelId ? getOpenRouterProvidersForModel(modelId, baseUrl, apiKey) : {}),
 		...options,
 	})

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

@@ -41,6 +41,7 @@ type RouterModelsQueryKey = {
 	openRouterApiKey?: string
 	lmStudioBaseUrl?: string
 	ollamaBaseUrl?: string
+	kilocodeOrganizationId?: string
 	// Requesty, Unbound, etc should perhaps also be here, but they already have their own hacks for reloading
 }
 

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

@@ -76,6 +76,8 @@ export const useModelProviders = (kilocodeDefaultModel: string, apiConfiguration
 				? (apiConfiguration?.openRouterModelId ?? openRouterDefaultModelId)
 				: undefined,
 		provider === "openrouter" ? apiConfiguration?.openRouterBaseUrl : undefined,
+		apiConfiguration?.apiKey,
+		apiConfiguration?.kilocodeOrganizationId ?? "personal",
 	)
 }
 // kilocode_change end
@@ -86,7 +88,8 @@ export const useSelectedModel = (apiConfiguration?: ProviderSettings) => {
 	const { kilocodeDefaultModel } = useExtensionState()
 	const routerModels = useRouterModels({
 		openRouterBaseUrl: apiConfiguration?.openRouterBaseUrl,
-		openRouterApiKey: apiConfiguration?.apiKey,
+		openRouterApiKey: apiConfiguration?.apiKey, // kilocode_change
+		kilocodeOrganizationId: apiConfiguration?.kilocodeOrganizationId, // kilocode_change
 	})
 	const openRouterModelProviders = useModelProviders(kilocodeDefaultModel, apiConfiguration)
 	// kilocode_change end