From d70977fd92dba79a52176a1a7ddaa4cd11a7de00 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 7 Jul 2026 14:38:02 +0200 Subject: [PATCH] test: cover Kilo account selection defaults --- packages/kilo-gateway/src/index.ts | 1 + packages/kilo-gateway/src/server/handlers.ts | 36 +++++++---- .../kilo-gateway/test/server/handlers.test.ts | 59 +++++++++++++++++++ .../src/components/profile/ProfileView.tsx | 8 +-- .../components/dialog-kilo-auto-method.tsx | 12 ++-- .../server/httpapi/handlers/kilo-gateway.ts | 16 ++--- 6 files changed, 102 insertions(+), 30 deletions(-) create mode 100644 packages/kilo-gateway/test/server/handlers.test.ts diff --git a/packages/kilo-gateway/src/index.ts b/packages/kilo-gateway/src/index.ts index 850047c4dd..034902c614 100644 --- a/packages/kilo-gateway/src/index.ts +++ b/packages/kilo-gateway/src/index.ts @@ -75,6 +75,7 @@ export { getCloudSessions, getNotifications, getProfile, + resolveProfileOrganization, getToken, normalizeClawStatus, setOrganization, diff --git a/packages/kilo-gateway/src/server/handlers.ts b/packages/kilo-gateway/src/server/handlers.ts index ac8b5d2907..dc0a5424fd 100644 --- a/packages/kilo-gateway/src/server/handlers.ts +++ b/packages/kilo-gateway/src/server/handlers.ts @@ -36,6 +36,11 @@ export interface OrganizationDeps { dispose(): Promise } +export interface OrganizationResolution { + currentOrgId: string | null + persistOrgId: string | null +} + export interface CloudSessionsInput { cursor?: string limit?: number @@ -64,6 +69,22 @@ export function getOrganizationId(auth: KiloAuth | undefined) { return undefined } +export function resolveProfileOrganization(profile: KilocodeProfile, auth: Extract): OrganizationResolution { + const selected = profile.selectedOrganizationId + const orgs = profile.organizations ?? [] + const valid = selected && orgs.some((org) => org.id === selected) ? selected : undefined + const cloud = valid ?? (profile.hasPersonalAccount === false ? orgs[0]?.id : undefined) + if (auth.accountSelection === "manual" && (auth.accountId || profile.hasPersonalAccount !== false)) { + return { currentOrgId: auth.accountId ?? null, persistOrgId: null } + } + + const local = auth.accountSelection === "cloud" ? undefined : auth.accountId + const currentOrgId = local ?? cloud ?? auth.accountId ?? null + const invalid = auth.accountSelection === "manual" && !auth.accountId && profile.hasPersonalAccount === false + const persistOrgId = currentOrgId && !local && (auth.accountSelection !== "manual" || invalid) && currentOrgId !== auth.accountId ? currentOrgId : null + return { currentOrgId, persistOrgId } +} + export async function getProfile(auth: AuthStore): Promise { const info = await auth.get("kilo") if (!info || info.type !== "oauth") throw new UnauthorizedError("Not authenticated with Kilo Gateway") @@ -73,25 +94,20 @@ export async function getProfile(auth: AuthStore): Promise { fetchKiloPassState(info.access), ]) - const selected = profile.selectedOrganizationId - const orgs = profile.organizations ?? [] - const valid = selected && orgs.some((org) => org.id === selected) ? selected : undefined - const cloud = valid ?? (profile.hasPersonalAccount === false ? orgs[0]?.id : undefined) - const local = info.accountSelection === "cloud" ? undefined : info.accountId - const currentOrgId = info.accountSelection === "manual" ? (info.accountId ?? null) : (local ?? cloud ?? info.accountId ?? null) - if (currentOrgId && !local && info.accountSelection !== "manual" && currentOrgId !== info.accountId) { + const org = resolveProfileOrganization(profile, info) + if (org.persistOrgId) { await auth.set("kilo", { type: "oauth", refresh: info.refresh, access: info.access, expires: info.expires, - accountId: currentOrgId, + accountId: org.persistOrgId, accountSelection: "cloud", }).catch((err) => console.warn("Failed to persist cloud account selection:", err)) } - const balance = await fetchBalance(info.access, currentOrgId ?? undefined) - return { profile, balance, kiloPass, currentOrgId } + const balance = await fetchBalance(info.access, org.currentOrgId ?? undefined) + return { profile, balance, kiloPass, currentOrgId: org.currentOrgId } } export async function getNotifications(auth: AuthStore) { diff --git a/packages/kilo-gateway/test/server/handlers.test.ts b/packages/kilo-gateway/test/server/handlers.test.ts new file mode 100644 index 0000000000..0b888c9ca8 --- /dev/null +++ b/packages/kilo-gateway/test/server/handlers.test.ts @@ -0,0 +1,59 @@ +import { describe, expect, test } from "bun:test" +import { resolveProfileOrganization, type KiloAuth } from "../../src/server/handlers.js" +import type { KilocodeProfile } from "../../src/types.js" + +const oauth = (input: Partial> = {}): Extract => ({ + type: "oauth", + access: "token", + refresh: "token", + expires: 1, + ...input, +}) + +const profile = (input: Partial = {}): KilocodeProfile => ({ + email: "user@example.com", + organizations: [{ id: "org_1", name: "Acme", role: "MEMBER" }], + ...input, +}) + +describe("resolveProfileOrganization", () => { + test("uses cloud selected organization only when no local selection exists", () => { + expect(resolveProfileOrganization(profile({ selectedOrganizationId: "org_1" }), oauth())).toEqual({ + currentOrgId: "org_1", + persistOrgId: "org_1", + }) + }) + + test("preserves a manual organization selection over the cloud default", () => { + expect( + resolveProfileOrganization( + profile({ + selectedOrganizationId: "org_1", + organizations: [ + { id: "org_1", name: "Acme", role: "MEMBER" }, + { id: "org_2", name: "Beta", role: "MEMBER" }, + ], + }), + oauth({ accountId: "org_2", accountSelection: "manual" }), + ), + ).toEqual({ currentOrgId: "org_2", persistOrgId: null }) + }) + + test("preserves explicit personal selection when personal is available", () => { + expect( + resolveProfileOrganization( + profile({ selectedOrganizationId: "org_1", hasPersonalAccount: true }), + oauth({ accountSelection: "manual" }), + ), + ).toEqual({ currentOrgId: null, persistOrgId: null }) + }) + + test("uses the first organization when personal is unavailable and no valid cloud selection exists", () => { + expect( + resolveProfileOrganization( + profile({ selectedOrganizationId: "missing", hasPersonalAccount: false }), + oauth({ accountSelection: "manual" }), + ), + ).toEqual({ currentOrgId: "org_1", persistOrgId: "org_1" }) + }) +}) diff --git a/packages/kilo-vscode/webview-ui/src/components/profile/ProfileView.tsx b/packages/kilo-vscode/webview-ui/src/components/profile/ProfileView.tsx index bed2bfbc98..436784c75c 100644 --- a/packages/kilo-vscode/webview-ui/src/components/profile/ProfileView.tsx +++ b/packages/kilo-vscode/webview-ui/src/components/profile/ProfileView.tsx @@ -44,6 +44,8 @@ const ProfileView: Component = (props) => { const language = useLanguage() const [target, setTarget] = createSignal(null) + const personal = createMemo(() => props.profileData?.profile.hasPersonalAccount !== false) + // Always fetch fresh profile+balance when navigating to this view onMount(() => { vscode.postMessage({ type: "refreshProfile" }) @@ -58,16 +60,14 @@ const ProfileView: Component = (props) => { const orgOptions = createMemo(() => { const orgs = props.profileData?.profile.organizations ?? [] if (orgs.length === 0) return [] - const personal = props.profileData?.profile.hasPersonalAccount !== false return [ - ...(personal ? [{ value: PERSONAL, label: language.t("profile.personalAccount") }] : []), + ...(personal() ? [{ value: PERSONAL, label: language.t("profile.personalAccount") }] : []), ...orgs.map((org) => ({ value: org.id, label: org.name, description: org.role })), ] }) const currentId = createMemo(() => { - const personal = props.profileData?.profile.hasPersonalAccount !== false - return props.profileData?.currentOrgId ?? (personal ? PERSONAL : orgOptions()[0]?.value) + return props.profileData?.currentOrgId ?? (personal() ? PERSONAL : orgOptions()[0]?.value) }) const switching = createMemo(() => { diff --git a/packages/opencode/src/kilocode/components/dialog-kilo-auto-method.tsx b/packages/opencode/src/kilocode/components/dialog-kilo-auto-method.tsx index 533571d02e..239ea52eec 100644 --- a/packages/opencode/src/kilocode/components/dialog-kilo-auto-method.tsx +++ b/packages/opencode/src/kilocode/components/dialog-kilo-auto-method.tsx @@ -83,12 +83,12 @@ export function KiloAutoMethod(props: KiloAutoMethodProps) { await sync.bootstrap() dialog.replace(() => ( - diff --git a/packages/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.ts b/packages/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.ts index bd610d3595..7e4d9968c3 100644 --- a/packages/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.ts +++ b/packages/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.ts @@ -7,6 +7,7 @@ import { getToken, importSessionToDb, normalizeClawStatus, + resolveProfileOrganization, } from "@kilocode/kilo-gateway" import { HEADER_FEATURE, @@ -76,28 +77,23 @@ export const kiloGatewayHandlers = HttpApiBuilder.group(InstanceHttpApi, "kilo", catch: () => new HttpApiError.BadRequest({}), }) - const selected = profile.selectedOrganizationId - const orgs = profile.organizations ?? [] - const valid = selected && orgs.some((org) => org.id === selected) ? selected : undefined - const cloud = valid ?? (profile.hasPersonalAccount === false ? orgs[0]?.id : undefined) - const local = info.accountSelection === "cloud" ? undefined : info.accountId - const currentOrgId = info.accountSelection === "manual" ? (info.accountId ?? null) : (local ?? cloud ?? info.accountId ?? null) - if (currentOrgId && !local && info.accountSelection !== "manual" && currentOrgId !== info.accountId) { + const org = resolveProfileOrganization(profile, info) + if (org.persistOrgId) { yield* auth.set("kilo", { type: "oauth", refresh: info.refresh, access: info.access, expires: info.expires, - accountId: currentOrgId, + accountId: org.persistOrgId, accountSelection: "cloud", }).pipe(Effect.catch((err) => Effect.sync(() => log.warn("failed to persist cloud account selection", { err })))) } const balance = yield* Effect.tryPromise({ - try: () => fetchBalance(info.access, currentOrgId ?? undefined), + try: () => fetchBalance(info.access, org.currentOrgId ?? undefined), catch: () => new HttpApiError.BadRequest({}), }) - return { profile, balance, kiloPass, currentOrgId } + return { profile, balance, kiloPass, currentOrgId: org.currentOrgId } }) const authStatus = Effect.fn("KiloGatewayHttpApi.authStatus")(function* () {