Kilo-Org_kilocode/packages/opencode/test/kilocode/model-cache-org.test.ts

// 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"
import { Log } from "../../src/util"

Log.init({ print: false })

// 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 default plugins to prevent actual installations during tests
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://app.kilo.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://app.kilo.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()
    },
  })
})

test("ModelCache.clear removes cached entry so next fetch hits the network", async () => {
  await using tmp = await tmpdir({
    init: async (dir) => {
      await Bun.write(
        path.join(dir, "opencode.json"),
        JSON.stringify({
          $schema: "https://app.kilo.ai/config.json",
        }),
      )
    },
  })
  await Instance.provide({
    directory: tmp.path,
    init: async () => {
      await Auth.set("kilo", {
        type: "oauth",
        access: "token-clear-test",
        refresh: "refresh-clear",
        expires: Date.now() + 3600000,
        accountId: "org-clear",
      })
    },
    fn: async () => {
      // Populate cache
      captured = undefined
      ModelCache.clear("kilo")
      await ModelCache.fetch("kilo")
      expect(captured).toBeDefined()

      // Verify cache is populated — second fetch should NOT call fetchKiloModels
      captured = undefined
      await ModelCache.fetch("kilo")
      expect(captured).toBeUndefined()
      expect(ModelCache.get("kilo")).toBeDefined()

      // Clear the cache
      ModelCache.clear("kilo")

      // get() should return undefined after clear
      expect(ModelCache.get("kilo")).toBeUndefined()

      // Next fetch should call fetchKiloModels again
      captured = undefined
      await ModelCache.fetch("kilo")
      expect(captured).toBeDefined()
    },
  })
})