mirror of
https://github.com/Kilo-Org/kilocode.git
synced 2026-09-24 16:02:55 +08:00
fix: use OAuth accountId for enterprise model list filtering (#193)
Enterprise users logging in via OAuth had their organization ID stored in auth.accountId but the model fetch code never read it, causing the model list to be fetched from the personal endpoint instead of the org-specific endpoint that applies the allow list.
This commit is contained in:
@@ -199,6 +199,11 @@ export namespace ModelCache {
|
||||
options.kilocodeToken = auth.key
|
||||
} else if (auth.type === "oauth") {
|
||||
options.kilocodeToken = auth.access
|
||||
// kilocode_change start - read org ID from OAuth accountId for enterprise model filtering
|
||||
if (auth.accountId) {
|
||||
options.kilocodeOrganizationId = auth.accountId
|
||||
}
|
||||
// kilocode_change end
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ import { Flag } from "../flag/flag"
|
||||
import { lazy } from "@/util/lazy"
|
||||
import { Config } from "../config/config" // kilocode_change
|
||||
import { ModelCache } from "./model-cache" // kilocode_change
|
||||
import { Auth } from "../auth" // kilocode_change
|
||||
import { KILO_OPENROUTER_BASE } from "@kilocode/kilo-gateway" // kilocode_change
|
||||
|
||||
// Try to import bundled snapshot (generated at build time)
|
||||
@@ -124,7 +125,11 @@ export namespace ModelsDev {
|
||||
if (!providers["kilo"]) {
|
||||
const config = await Config.get()
|
||||
const kiloOptions = config.provider?.kilo?.options
|
||||
const kiloOrgId = kiloOptions?.kilocodeOrganizationId
|
||||
// kilocode_change start - resolve org ID from auth (OAuth accountId) not just config
|
||||
const kiloAuth = await Auth.get("kilo")
|
||||
const kiloOrgId =
|
||||
kiloOptions?.kilocodeOrganizationId ?? (kiloAuth?.type === "oauth" ? kiloAuth.accountId : undefined)
|
||||
// kilocode_change end
|
||||
const normalizedBaseURL = normalizeKiloBaseURL(kiloOptions?.baseURL, kiloOrgId)
|
||||
const kiloFetchOptions = {
|
||||
...(normalizedBaseURL ? { baseURL: normalizedBaseURL } : {}),
|
||||
|
||||
@@ -0,0 +1,124 @@
|
||||
// kilocode_change - new file
|
||||
// Regression test: OAuth accountId must flow into model fetch as kilocodeOrganizationId
|
||||
// When a user logs in via OAuth and selects an enterprise organization, the model fetch
|
||||
// should use the organization-specific endpoint, not the personal endpoint.
|
||||
|
||||
import { test, expect, mock } from "bun:test"
|
||||
import path from "path"
|
||||
|
||||
// Capture the options passed to fetchKiloModels
|
||||
let captured: any = undefined
|
||||
|
||||
mock.module("@kilocode/kilo-gateway", () => ({
|
||||
fetchKiloModels: async (options: any) => {
|
||||
captured = options
|
||||
return {
|
||||
"test-model": {
|
||||
id: "test-model",
|
||||
name: "Test Model",
|
||||
cost: { input: 0.001, output: 0.002 },
|
||||
limit: { context: 128000, output: 4096 },
|
||||
},
|
||||
}
|
||||
},
|
||||
KILO_OPENROUTER_BASE: "https://api.kilo.ai/api/openrouter",
|
||||
}))
|
||||
|
||||
// Mock BunProc and default plugins to prevent actual installations during tests
|
||||
mock.module("../../src/bun/index", () => ({
|
||||
BunProc: {
|
||||
install: async (pkg: string) => {
|
||||
const lastAtIndex = pkg.lastIndexOf("@")
|
||||
return lastAtIndex > 0 ? pkg.substring(0, lastAtIndex) : pkg
|
||||
},
|
||||
run: async () => {
|
||||
throw new Error("BunProc.run should not be called in tests")
|
||||
},
|
||||
which: () => process.execPath,
|
||||
InstallFailedError: class extends Error {},
|
||||
},
|
||||
}))
|
||||
|
||||
const mockPlugin = () => ({})
|
||||
mock.module("opencode-copilot-auth", () => ({ default: mockPlugin }))
|
||||
mock.module("opencode-anthropic-auth", () => ({ default: mockPlugin }))
|
||||
mock.module("@gitlab/opencode-gitlab-auth", () => ({ default: mockPlugin }))
|
||||
|
||||
import { tmpdir } from "../fixture/fixture"
|
||||
import { Instance } from "../../src/project/instance"
|
||||
import { Auth } from "../../src/auth"
|
||||
import { ModelCache } from "../../src/provider/model-cache"
|
||||
|
||||
test("model fetch uses accountId from OAuth auth as kilocodeOrganizationId", async () => {
|
||||
await using tmp = await tmpdir({
|
||||
init: async (dir) => {
|
||||
await Bun.write(
|
||||
path.join(dir, "opencode.json"),
|
||||
JSON.stringify({
|
||||
$schema: "https://opencode.ai/config.json",
|
||||
}),
|
||||
)
|
||||
},
|
||||
})
|
||||
await Instance.provide({
|
||||
directory: tmp.path,
|
||||
init: async () => {
|
||||
// Simulate an OAuth login where user selected an enterprise organization
|
||||
await Auth.set("kilo", {
|
||||
type: "oauth",
|
||||
access: "test-oauth-token",
|
||||
refresh: "test-refresh-token",
|
||||
expires: Date.now() + 3600000,
|
||||
accountId: "org-enterprise-123",
|
||||
})
|
||||
},
|
||||
fn: async () => {
|
||||
// Reset captured and cache
|
||||
captured = undefined
|
||||
ModelCache.clear("kilo")
|
||||
|
||||
// Trigger model fetch through the cache
|
||||
await ModelCache.fetch("kilo")
|
||||
|
||||
// The fetchKiloModels call should have received the organization ID
|
||||
expect(captured).toBeDefined()
|
||||
expect(captured.kilocodeToken).toBe("test-oauth-token")
|
||||
expect(captured.kilocodeOrganizationId).toBe("org-enterprise-123")
|
||||
},
|
||||
})
|
||||
})
|
||||
|
||||
test("model fetch without OAuth accountId does not set kilocodeOrganizationId", async () => {
|
||||
await using tmp = await tmpdir({
|
||||
init: async (dir) => {
|
||||
await Bun.write(
|
||||
path.join(dir, "opencode.json"),
|
||||
JSON.stringify({
|
||||
$schema: "https://opencode.ai/config.json",
|
||||
}),
|
||||
)
|
||||
},
|
||||
})
|
||||
await Instance.provide({
|
||||
directory: tmp.path,
|
||||
init: async () => {
|
||||
// Simulate an OAuth login for a personal account (no accountId)
|
||||
await Auth.set("kilo", {
|
||||
type: "oauth",
|
||||
access: "test-personal-token",
|
||||
refresh: "test-refresh-token",
|
||||
expires: Date.now() + 3600000,
|
||||
})
|
||||
},
|
||||
fn: async () => {
|
||||
captured = undefined
|
||||
ModelCache.clear("kilo")
|
||||
|
||||
await ModelCache.fetch("kilo")
|
||||
|
||||
expect(captured).toBeDefined()
|
||||
expect(captured.kilocodeToken).toBe("test-personal-token")
|
||||
expect(captured.kilocodeOrganizationId).toBeUndefined()
|
||||
},
|
||||
})
|
||||
})
|
||||
Reference in New Issue
Block a user