test: cover Kilo account selection defaults

packages/kilo-gateway/src/index.ts

@@ -75,6 +75,7 @@ export {
   getCloudSessions,
   getNotifications,
   getProfile,
+  resolveProfileOrganization,
   getToken,
   normalizeClawStatus,
   setOrganization,

packages/kilo-gateway/src/server/handlers.ts

@@ -36,6 +36,11 @@ export interface OrganizationDeps {
   dispose(): Promise<void>
 }
 
+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<KiloAuth, { type: "oauth" }>): 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<KiloProfileResult> {
   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<KiloProfileResult> {
     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) {

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<KiloAuth, { type: "oauth" }>> = {}): Extract<KiloAuth, { type: "oauth" }> => ({
+  type: "oauth",
+  access: "token",
+  refresh: "token",
+  expires: 1,
+  ...input,
+})
+
+const profile = (input: Partial<KilocodeProfile> = {}): 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" })
+  })
+})

packages/kilo-vscode/webview-ui/src/components/profile/ProfileView.tsx

@@ -44,6 +44,8 @@ const ProfileView: Component<ProfileViewProps> = (props) => {
   const language = useLanguage()
   const [target, setTarget] = createSignal<string | null>(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<ProfileViewProps> = (props) => {
   const orgOptions = createMemo<OrgOption[]>(() => {
     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(() => {

packages/opencode/src/kilocode/components/dialog-kilo-auto-method.tsx

@@ -83,12 +83,12 @@ export function KiloAutoMethod(props: KiloAutoMethodProps) {
         await sync.bootstrap()
 
         dialog.replace(() => (
-            <DialogKiloOrganization
-              organizations={profile.organizations!}
-              userEmail={profile.email}
-              providerID={props.providerID}
-              hasPersonalAccount={profile.hasPersonalAccount !== false}
-              useSDK={props.useSDK}
+          <DialogKiloOrganization
+            organizations={profile.organizations!}
+            userEmail={profile.email}
+            providerID={props.providerID}
+            hasPersonalAccount={profile.hasPersonalAccount !== false}
+            useSDK={props.useSDK}
             useTheme={props.useTheme}
             DialogModel={props.DialogModel}
           />

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* () {