diff --git a/packages/opencode/src/provider/model-cache.ts b/packages/opencode/src/provider/model-cache.ts index deba50254e3..1d1a2d4e129 100644 --- a/packages/opencode/src/provider/model-cache.ts +++ b/packages/opencode/src/provider/model-cache.ts @@ -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 } } diff --git a/packages/opencode/src/provider/models.ts b/packages/opencode/src/provider/models.ts index 7ed5ac7e7e2..bbc95913b03 100644 --- a/packages/opencode/src/provider/models.ts +++ b/packages/opencode/src/provider/models.ts @@ -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 } : {}), diff --git a/packages/opencode/test/kilocode/model-cache-org.test.ts b/packages/opencode/test/kilocode/model-cache-org.test.ts new file mode 100644 index 00000000000..e18bc35580e --- /dev/null +++ b/packages/opencode/test/kilocode/model-cache-org.test.ts @@ -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() + }, + }) +})