From d173377a2bbe47b89eb7438286ed5650b272bf25 Mon Sep 17 00:00:00 2001 From: Marius Date: Mon, 9 Feb 2026 16:08:19 +0100 Subject: [PATCH] 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. --- packages/opencode/src/provider/model-cache.ts | 5 + packages/opencode/src/provider/models.ts | 7 +- .../test/kilocode/model-cache-org.test.ts | 124 ++++++++++++++++++ 3 files changed, 135 insertions(+), 1 deletion(-) create mode 100644 packages/opencode/test/kilocode/model-cache-org.test.ts 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() + }, + }) +})