diff --git a/packages/core/src/credential.ts b/packages/core/src/credential.ts index 2c5d614cee4..11a987fac37 100644 --- a/packages/core/src/credential.ts +++ b/packages/core/src/credential.ts @@ -1,6 +1,6 @@ export * as Credential from "./credential" -import { asc, eq } from "drizzle-orm" +import { asc, desc, eq } from "drizzle-orm" // kilocode_change // kilocode_change start import { Context, Effect, Layer, Option, Schema, Semaphore } from "effect" // kilocode_change end @@ -122,30 +122,48 @@ export const legacyImportLayer = Layer.effectDiscard( const { db } = yield* Database.Service const fs = yield* FSUtil.Service const global = yield* Global.Service - // the v2 name re-runs the import because upstream migration 20260611192811 drops the credential table - const kiloName = "credential.kilo-account-json-v2" + // v3 repairs the active-only v2 import while remaining safe for users who already ran it. + const kiloName = "credential.kilo-account-json-v3" if (!(yield* db.select().from(DataMigrationTable).where(eq(DataMigrationTable.name, kiloName)).get())) { const current = yield* fs.readJson(path.join(global.data, "account.json")).pipe(Effect.option) const prior = yield* fs.readJson(path.join(global.data, "auth-v2.json")).pipe(Effect.option) const raw = Option.isSome(current) ? current.value : Option.getOrUndefined(prior) - // one credential per integration: only the active account of each service is imported - const values = parseKiloAccounts(raw).filter((item) => item.active) + const values = parseKiloAccounts(raw).toSorted( + (a, b) => a.connectorID.localeCompare(b.connectorID) || Number(a.active) - Number(b.active), + ) if (values.length > 0) { yield* db.transaction((tx) => Effect.gen(function* () { - const existing = new Set( - (yield* tx.select({ integrationID: CredentialTable.integration_id }).from(CredentialTable).all()).map( - (item) => item.integrationID, - ), - ) - for (const item of values) { + const existing = yield* tx.select().from(CredentialTable).all() + const used = new Set() + const created = Date.now() + for (const [index, item] of values.entries()) { const integration = IntegrationSchema.ID.make(item.connectorID.replace(/\/+$/, "")) - if (existing.has(integration)) continue + const value = legacyValue(integration, item.credential) + const current = existing.find( + (row) => + !used.has(row.id) && + row.integration_id === integration && + row.label === item.label && + JSON.stringify(row.value) === JSON.stringify(value), + ) + const time = created + index + if (current) { + used.add(current.id) + yield* tx + .update(CredentialTable) + .set({ time_created: time, time_updated: time }) + .where(eq(CredentialTable.id, current.id)) + .run() + continue + } yield* tx.insert(CredentialTable).values({ - id: ID.create(), + id: ID.make(`cred_kilo_${Buffer.from(item.id).toString("base64url")}`), integration_id: integration, label: item.label, - value: legacyValue(integration, item.credential), + value, + time_created: time, + time_updated: time, }) } yield* tx.insert(DataMigrationTable).values({ name: kiloName, time_completed: Date.now() }).run() @@ -171,6 +189,7 @@ export const legacyImportLayer = Layer.effectDiscard( .select() .from(CredentialTable) .where(eq(CredentialTable.integration_id, item.integration)) + .orderBy(desc(CredentialTable.time_created)) // kilocode_change - reconcile the active imported account .get() if (current) { yield* tx.update(CredentialTable).set({ value: item.value }).where(eq(CredentialTable.id, current.id)).run() @@ -276,6 +295,7 @@ export const layer = Layer.effect( .select() .from(CredentialTable) .where(eq(CredentialTable.integration_id, integration)) + .orderBy(desc(CredentialTable.time_created)) // kilocode_change - persist the active imported account .get() .pipe(Effect.orDie) delete data[integration + "/"] diff --git a/packages/core/src/kilocode/credential-migration.ts b/packages/core/src/kilocode/credential-migration.ts index 387c9dce0f6..590643ceb67 100644 --- a/packages/core/src/kilocode/credential-migration.ts +++ b/packages/core/src/kilocode/credential-migration.ts @@ -37,6 +37,7 @@ export function parse(input: unknown) { const fallback = !first.has(account.serviceID) first.add(account.serviceID) return { + id: account.id, connectorID: account.serviceID, label: account.description, credential: account.credential, diff --git a/packages/core/test/kilocode/account-auth-v2-migration.test.ts b/packages/core/test/kilocode/account-auth-v2-migration.test.ts index 22a7b7c159f..cabfb758c5c 100644 --- a/packages/core/test/kilocode/account-auth-v2-migration.test.ts +++ b/packages/core/test/kilocode/account-auth-v2-migration.test.ts @@ -1,12 +1,15 @@ import path from "path" import { describe, expect } from "bun:test" import { Effect, Layer } from "effect" +import { eq } from "drizzle-orm" import { IntegrationSchema } from "@opencode-ai/core/integration/schema" import { Credential } from "@opencode-ai/core/credential" import { Database } from "@opencode-ai/core/database/database" import { EventV2 } from "@opencode-ai/core/event" import { FSUtil } from "@opencode-ai/core/fs-util" import { Global } from "@opencode-ai/core/global" +import { DataMigrationTable } from "@opencode-ai/core/data-migration.sql" +import { CredentialTable } from "@opencode-ai/core/credential/sql" import { tmpdir } from "../fixture/tmpdir" import { it } from "../lib/effect" @@ -38,7 +41,7 @@ const auth = Effect.acquireRelease( ) describe("Credential auth-v2 migration", () => { - it.live("imports the active account with its Kilo organization", () => + it.live("imports every account with the active account ordered last", () => Effect.acquireRelease( Effect.promise(() => tmpdir()), (tmp) => Effect.promise(() => tmp[Symbol.asyncDispose]()), @@ -87,13 +90,13 @@ describe("Credential auth-v2 migration", () => { } }).pipe(Effect.provide(layer(tmp.path))) - // one credential per integration: only the active account is imported - expect(result.all.map((item) => item.label)).toEqual(["second"]) - expect(result.list.length).toBe(1) - expect(result.list[0]?.value.type).toBe("oauth") - if (result.list[0]?.value.type === "oauth") { - expect(result.list[0].value.access).toBe("access-second") - expect(result.list[0].value.metadata?.accountID).toBe("org-second") + expect(result.all.map((item) => item.label)).toEqual(["first", "second"]) + expect(result.list.length).toBe(2) + const active = result.list.at(-1) + expect(active?.value.type).toBe("oauth") + if (active?.value.type === "oauth") { + expect(active.value.access).toBe("access-second") + expect(active.value.metadata?.accountID).toBe("org-second") } }), ), @@ -101,4 +104,96 @@ describe("Credential auth-v2 migration", () => { ), ), ) + + it.live("repairs an active-only v2 import without duplicating the active account", () => + Effect.acquireRelease( + Effect.promise(() => tmpdir()), + (tmp) => Effect.promise(() => tmp[Symbol.asyncDispose]()), + ).pipe( + Effect.flatMap((tmp) => + auth.pipe( + Effect.flatMap(() => + Effect.gen(function* () { + const integration = IntegrationSchema.ID.make("kilo") + const active = new Credential.OAuth({ + type: "oauth", + methodID: IntegrationSchema.MethodID.make("oauth"), + refresh: "refresh-second", + access: "access-second", + expires: 2, + metadata: { accountID: "org-second" }, + }) + const database = Database.layerFromPath(path.join(tmp.path, "credential.db")).pipe(Layer.fresh) + yield* Effect.gen(function* () { + const { db } = yield* Database.Service + yield* db.insert(CredentialTable).values({ + id: Credential.ID.create(), + integration_id: integration, + label: "second", + value: active, + }) + yield* db.insert(DataMigrationTable).values({ + name: "credential.kilo-account-json-v2", + time_completed: Date.now(), + }) + }).pipe(Effect.provide(database)) + + yield* Effect.promise(() => + Bun.write( + path.join(tmp.path, "auth-v2.json"), + JSON.stringify({ + version: 2, + accounts: { + acc_first: { + id: "acc_first", + serviceID: "kilo", + description: "first", + credential: { + type: "oauth", + refresh: "refresh-first", + access: "access-first", + expires: 1, + accountId: "org-first", + }, + }, + acc_second: { + id: "acc_second", + serviceID: "kilo", + description: "second", + credential: { + type: "oauth", + refresh: "refresh-second", + access: "access-second", + expires: 2, + accountId: "org-second", + }, + }, + }, + active: { kilo: "acc_second" }, + }), + ), + ) + + const result = yield* Effect.gen(function* () { + const credentials = yield* Credential.Service + return yield* credentials.list(integration) + }).pipe(Effect.provide(layer(tmp.path))) + const repaired = yield* Effect.gen(function* () { + const { db } = yield* Database.Service + return yield* db + .select() + .from(DataMigrationTable) + .where(eq(DataMigrationTable.name, "credential.kilo-account-json-v3")) + .get() + }).pipe(Effect.provide(database)) + + expect(result.map((item) => item.label)).toEqual(["first", "second"]) + expect(result.filter((item) => item.label === "second")).toHaveLength(1) + expect(repaired).toBeDefined() + }), + ), + ), + ), + ), + ) }) diff --git a/packages/opencode/src/provider/provider.ts b/packages/opencode/src/provider/provider.ts index c6d6d2b3d67..8e78708465e 100644 --- a/packages/opencode/src/provider/provider.ts +++ b/packages/opencode/src/provider/provider.ts @@ -881,17 +881,23 @@ function custom(dep: CustomDep): Record { const account = env["SNOWFLAKE_ACCOUNT"] ?? (auth?.type === "api" ? auth.metadata?.account : undefined) ?? + (auth?.type === "oauth" ? auth.accountId : undefined) ?? input.options?.account - const pat = env["SNOWFLAKE_CORTEX_PAT"] ?? (auth?.type === "api" ? auth.key : undefined) ?? input.options?.apiKey + const envToken = env["SNOWFLAKE_CORTEX_TOKEN"] ?? env["SNOWFLAKE_CORTEX_PAT"] + const apiKeyToken = auth?.type === "api" ? auth.key : undefined + const oauthToken = auth?.type === "oauth" ? auth.access : undefined + const configToken = input.options?.token ?? input.options?.apiKey - if (!account || !pat) { - const missing = [!account && "SNOWFLAKE_ACCOUNT", !pat && "SNOWFLAKE_CORTEX_PAT"].filter(Boolean).join(", ") + const token = envToken ?? apiKeyToken ?? oauthToken ?? configToken + + if (!account || !token) { + const missing = [!account && "SNOWFLAKE_ACCOUNT", !token && "SNOWFLAKE_CORTEX_TOKEN"].filter(Boolean).join(", ") return { autoload: false, async getModel() { throw new Error( - `Snowflake Cortex: missing credentials (${missing}). Set via env var, kilo auth, or provider options.`, // kilocode_change + `Snowflake Cortex: missing credentials (${missing}). Provide a bearer token (OAuth, JWT, or PAT) via env var, Kilo auth, or provider options.`, // kilocode_change ) }, } @@ -899,12 +905,17 @@ function custom(dep: CustomDep): Record { const baseURL = `https://${account}.snowflakecomputing.com/api/v2/cortex/v1` - return { - autoload: input.source === "config", - options: { - baseURL, - apiKey: pat, - fetch: async (url: RequestInfo | URL, init?: RequestInit) => { + const options: Record = { baseURL, apiKey: token } + + // Only skip provider-level fetch when the token is from OAuth with no override. + // For OAuth tokens, the plugin auth loader's combined fetch handles + // OAuth refresh + snowflake transformations in one place. + // For env/config/API-key tokens, the provider fetch applies snowflake + // transformations directly. + const useOAuthHandler = + oauthToken !== undefined && envToken === undefined && apiKeyToken === undefined && configToken === undefined + if (!useOAuthHandler) { + options.fetch = async (url: RequestInfo | URL, init?: RequestInit) => { if (init?.body && typeof init.body === "string") { try { const body = JSON.parse(init.body) @@ -957,8 +968,12 @@ function custom(dep: CustomDep): Record { } return response - }, - }, + } + } + + return { + autoload: input.source === "config", + options, } }), } diff --git a/packages/opencode/test/kilocode/snowflake-cortex-provider.test.ts b/packages/opencode/test/kilocode/snowflake-cortex-provider.test.ts new file mode 100644 index 00000000000..85ade34248d --- /dev/null +++ b/packages/opencode/test/kilocode/snowflake-cortex-provider.test.ts @@ -0,0 +1,60 @@ +import { expect } from "bun:test" +import { Effect, Layer } from "effect" +import { CrossSpawnSpawner } from "@opencode-ai/core/cross-spawn-spawner" +import { ProviderV2 } from "@opencode-ai/core/provider" +import { Provider } from "@/provider/provider" +import { InstanceLayer } from "@/project/instance-layer" +import { Env } from "@/env" +import { Plugin } from "@/plugin" +import { provideInstanceEffect, tmpdirScoped } from "../fixture/fixture" +import { testEffect } from "../lib/effect" + +const it = testEffect( + Layer.mergeAll(Provider.defaultLayer, Env.defaultLayer, Plugin.defaultLayer, CrossSpawnSpawner.defaultLayer), +) + +it.effect("loads Snowflake Cortex from OAuth credentials", () => + Effect.acquireRelease( + Effect.sync(() => { + const value = process.env.KILO_AUTH_CONTENT + process.env.KILO_AUTH_CONTENT = JSON.stringify({ + "snowflake-cortex": { + type: "oauth", + refresh: "refresh-token", + access: "access-token", + expires: 1, + accountId: "test-account", + }, + }) + return value + }), + (value) => + Effect.sync(() => { + if (value === undefined) delete process.env.KILO_AUTH_CONTENT + else process.env.KILO_AUTH_CONTENT = value + }), + ).pipe( + Effect.flatMap(() => + Effect.gen(function* () { + const directory = yield* tmpdirScoped({ + config: { + provider: { + "snowflake-cortex": { + name: "Snowflake Cortex", + npm: "@ai-sdk/openai-compatible", + models: { test: { name: "Test" } }, + }, + }, + }, + }) + const provider = yield* Provider.use + .getProvider(ProviderV2.ID.make("snowflake-cortex")) + .pipe(provideInstanceEffect(directory), Effect.provide(InstanceLayer.layer)) + + expect(provider.options.baseURL).toBe("https://test-account.snowflakecomputing.com/api/v2/cortex/v1") + expect(provider.options.apiKey).toBe("access-token") + expect(provider.options.fetch).toBeFunction() + }), + ), + ), +)