refactor(cli): auth/workspace cleanup — record-backed token store

cli/src/auth/hosts.test.ts

@@ -1,5 +1,5 @@
 import type { AccountContext } from './hosts'
-import type { Key, Store } from '@/store/store'
+import type { TokenStore } from '@/store/token-store'
 import { mkdtemp, rm } from 'node:fs/promises'
 import { tmpdir } from 'node:os'
 import { join } from 'node:path'
@@ -53,6 +53,20 @@ describe('RegistrySchema', () => {
     })
     expect(ctx.external_subject?.issuer).toBe('https://issuer')
   })
+
+  it('strips a stale available_workspaces field from legacy contexts', () => {
+    const raw = {
+      account: { id: 'acct-1', email: 'bob@corp.com', name: 'Bob' },
+      workspace: { id: 'ws-1', name: 'Space', role: 'owner' },
+      available_workspaces: [
+        { id: 'ws-1', name: 'Space', role: 'owner' },
+        { id: '00000000-0000-0000-0000-000000000002', name: 'Other', role: 'normal' },
+      ],
+    } as unknown as Record<string, unknown>
+    const ctx = AccountContextSchema.parse(raw)
+    expect((ctx as Record<string, unknown>).available_workspaces).toBeUndefined()
+    expect(ctx.workspace?.id).toBe('ws-1')
+  })
 })
 
 describe('notLoggedInError', () => {
@@ -158,11 +172,12 @@ describe('Registry.load / Registry.save', () => {
   })
 })
 
-class MemStore implements Store {
-  readonly entries = new Map<string, unknown>()
-  get<T>(key: Key<T>): T { return (this.entries.get(key.key) as T | undefined) ?? key.default }
-  set<T>(key: Key<T>, value: T): void { this.entries.set(key.key, value) }
-  unset<T>(key: Key<T>): void { this.entries.delete(key.key) }
+class MemStore implements TokenStore {
+  readonly entries = new Map<string, string>()
+  private k(host: string, email: string): string { return `${host} ${email}` }
+  read(host: string, email: string): string { return this.entries.get(this.k(host, email)) ?? '' }
+  write(host: string, email: string, bearer: string): void { this.entries.set(this.k(host, email), bearer) }
+  remove(host: string, email: string): void { this.entries.delete(this.k(host, email)) }
 }
 
 describe('Registry.forget', () => {
@@ -188,12 +203,12 @@ describe('Registry.forget', () => {
     reg.setHost('h1')
     reg.setAccount('a@x')
     reg.save()
-    store.set({ key: 'tokens.h1.a@x', default: '' }, 'dfoa_a')
+    store.write('h1', 'a@x', 'dfoa_a')
 
     const active = reg.resolveActive()!
     reg.forget(active, store)
 
-    expect(store.get({ key: 'tokens.h1.a@x', default: '' })).toBe('')
+    expect(store.read('h1', 'a@x')).toBe('')
     const after = Registry.load()
     expect(after?.hosts.h1?.accounts['a@x']).toBeUndefined()
     expect(after?.hosts.h1?.accounts['b@x']).toBeDefined()

cli/src/auth/hosts.ts

@@ -1,11 +1,14 @@
-import type { Store } from '@/store/store'
+import type { StorageMode } from '@/store/store'
+import type { TokenStore } from '@/store/token-store'
 import { z } from 'zod'
 import { BaseError } from '@/errors/base'
 import { ErrorCode } from '@/errors/codes'
-import { getHostStore, tokenKey } from '@/store/manager'
+import { getHostStore } from '@/store/manager'
+import { STORAGE_MODES } from '@/store/store'
 
-const StorageModeSchema = z.enum(['keychain', 'file'])
-export type StorageMode = z.infer<typeof StorageModeSchema>
+const StorageModeSchema = z.enum(STORAGE_MODES)
+
+export type { StorageMode } from '@/store/store'
 
 export const AccountSchema = z.object({
   id: z.string().optional(),
@@ -30,7 +33,6 @@ export type ExternalSubject = z.infer<typeof ExternalSubjectSchema>
 export const AccountContextSchema = z.object({
   account: AccountSchema,
   workspace: WorkspaceSchema.optional(),
-  available_workspaces: z.array(WorkspaceSchema).optional(),
   token_id: z.string().optional(),
   token_expires_at: z.string().optional(),
   external_subject: ExternalSubjectSchema.optional(),
@@ -163,9 +165,9 @@ export class Registry {
 
   // Teardown for "this credential is gone": drop the token, drop the context
   // (unsets pointers when active), persist. Logout + self-revoke share it.
-  forget(active: ActiveContext, store: Store): void {
+  forget(active: ActiveContext, store: TokenStore): void {
     try {
-      store.unset(tokenKey(active.host, active.email))
+      store.remove(active.host, active.email)
     }
     catch { /* best-effort */ }
     this.remove(active.host, active.email)

cli/src/commands/_shared/authed-command.ts

@@ -2,7 +2,7 @@ import type { ActiveContext } from '@/auth/hosts'
 import type { AppInfoCache } from '@/cache/app-info'
 import type { Command } from '@/framework/command'
 import type { HttpClient } from '@/http/types'
-import type { Store } from '@/store/store'
+import type { TokenStore } from '@/store/token-store'
 import type { IOStreams } from '@/sys/io/streams'
 import { META_PROBE_TIMEOUT_MS, MetaClient } from '@/api/meta'
 import { notLoggedInError, Registry } from '@/auth/hosts'
@@ -11,7 +11,7 @@ import { loadNudgeStore } from '@/cache/nudge-store'
 import { getEnv } from '@/env/registry'
 import { formatErrorForCli } from '@/errors/format'
 import { createHttpClient } from '@/http/client'
-import { getTokenStore, tokenKey } from '@/store/manager'
+import { getTokenStore } from '@/store/manager'
 import { realStreams } from '@/sys/io/streams'
 import { hostWithScheme, openAPIBase } from '@/util/host'
 import { versionInfo } from '@/version/info'
@@ -21,7 +21,7 @@ import { resolveRetryAttempts } from './global-flags.js'
 export type AuthedContext = {
   readonly reg: Registry
   readonly active: ActiveContext
-  readonly store: Store
+  readonly store: TokenStore
   readonly http: HttpClient
   readonly host: string
   readonly io: IOStreams
@@ -44,8 +44,8 @@ export async function buildAuthedContext(
   if (active === undefined)
     fail(cmd, opts, io)
 
-  const { store } = getTokenStore()
-  const bearer = store.get(tokenKey(active.host, active.email))
+  const store = getTokenStore(reg.token_storage)
+  const bearer = store.read(active.host, active.email)
   if (bearer === '')
     fail(cmd, opts, io)
 

cli/src/commands/auth/devices/_shared/devices.test.ts

@@ -2,7 +2,7 @@ import type { SessionListResponse, SessionRow } from '@dify/contracts/api/openap
 import type { DifyMock } from '@test/fixtures/dify-mock/server'
 import type { AccountSessionsClient } from '@/api/account-sessions'
 import type { ActiveContext } from '@/auth/hosts'
-import type { Key, Store } from '@/store/store'
+import type { TokenStore } from '@/store/token-store'
 import { mkdtemp, rm } from 'node:fs/promises'
 import { tmpdir } from 'node:os'
 import { join } from 'node:path'
@@ -11,22 +11,25 @@ import { testHttpClient } from '@test/fixtures/http-client'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import { Registry } from '@/auth/hosts'
 import { ENV_CONFIG_DIR } from '@/store/dir'
-import { tokenKey } from '@/store/manager'
 import { bufferStreams } from '@/sys/io/streams'
 import { listAllSessions, runDevicesList, runDevicesRevoke } from './devices.js'
 
-class MemStore implements Store {
-  readonly entries = new Map<string, unknown>()
-  get<T>(key: Key<T>): T {
-    return (this.entries.get(key.key) as T | undefined) ?? key.default
+class MemStore implements TokenStore {
+  readonly entries = new Map<string, string>()
+  private k(host: string, email: string): string {
+    return `${host} ${email}`
   }
 
-  set<T>(key: Key<T>, value: T): void {
-    this.entries.set(key.key, value)
+  read(host: string, email: string): string {
+    return this.entries.get(this.k(host, email)) ?? ''
   }
 
-  unset<T>(key: Key<T>): void {
-    this.entries.delete(key.key)
+  write(host: string, email: string, bearer: string): void {
+    this.entries.set(this.k(host, email), bearer)
+  }
+
+  remove(host: string, email: string): void {
+    this.entries.delete(this.k(host, email))
   }
 }
 
@@ -35,10 +38,6 @@ function buildRegistry(host: string, email: string, tokenId: string): { reg: Reg
   reg.upsert(host, email, {
     account: { id: 'acct-1', email, name: 'Test Tester' },
     workspace: { id: 'ws-1', name: 'Default', role: 'owner' },
-    available_workspaces: [
-      { id: 'ws-1', name: 'Default', role: 'owner' },
-      { id: 'ws-2', name: 'Other', role: 'normal' },
-    ],
     token_id: tokenId,
   })
   reg.setHost(host)
@@ -103,7 +102,7 @@ describe('runDevicesRevoke', () => {
     const io = bufferStreams()
     const store = new MemStore()
     const { reg, active } = buildRegistry(mock.url, 'tester@dify.ai', 'tok-1')
-    store.set(tokenKey(mock.url, 'tester@dify.ai'), 'dfoa_test')
+    store.write(mock.url, 'tester@dify.ai', 'dfoa_test')
     reg.save()
     const http = testHttpClient(mock.url, 'dfoa_test')
 
@@ -168,7 +167,7 @@ describe('runDevicesRevoke', () => {
     const io = bufferStreams()
     const store = new MemStore()
     const { reg, active } = buildRegistry(mock.url, 'tester@dify.ai', 'tok-1')
-    store.set(tokenKey(mock.url, 'tester@dify.ai'), 'dfoa_test')
+    store.write(mock.url, 'tester@dify.ai', 'dfoa_test')
     reg.save()
     const http = testHttpClient(mock.url, 'dfoa_test')
 

cli/src/commands/auth/devices/_shared/devices.ts

@@ -1,7 +1,7 @@
 import type { SessionRow } from '@dify/contracts/api/openapi/types.gen'
 import type { ActiveContext, Registry } from '@/auth/hosts'
 import type { HttpClient } from '@/http/types'
-import type { Store } from '@/store/store'
+import type { TokenStore } from '@/store/token-store'
 import type { IOStreams } from '@/sys/io/streams'
 import { AccountSessionsClient } from '@/api/account-sessions'
 import { BaseError } from '@/errors/base'
@@ -71,7 +71,7 @@ export type DevicesRevokeOptions = {
   readonly io: IOStreams
   readonly reg: Registry
   readonly active: ActiveContext
-  readonly store: Store
+  readonly store: TokenStore
   readonly http: HttpClient
   readonly target?: string
   readonly all: boolean

cli/src/commands/auth/login/login.test.ts

@@ -1,6 +1,6 @@
 import type { DifyMock } from '@test/fixtures/dify-mock/server'
 import type { Clock } from './device-flow.js'
-import type { Key, Store } from '@/store/store'
+import type { TokenStore } from '@/store/token-store'
 import { mkdtemp, readFile, rm } from 'node:fs/promises'
 import { tmpdir } from 'node:os'
 import { join } from 'node:path'
@@ -10,7 +10,6 @@ import { afterEach, beforeEach, describe, expect, it } from 'vitest'
 import { DeviceFlowApi } from '@/api/oauth-device'
 import { createHttpClient } from '@/http/client'
 import { ENV_CONFIG_DIR } from '@/store/dir'
-import { tokenKey } from '@/store/manager'
 import { bufferStreams } from '@/sys/io/streams'
 import { openAPIBase } from '@/util/host'
 import { runLogin } from './login.js'
@@ -22,18 +21,22 @@ const noopClock: Clock = {
 
 const noopBrowser = async (): Promise<void> => { /* skip OS open */ }
 
-class MemStore implements Store {
-  readonly entries = new Map<string, unknown>()
-  get<T>(key: Key<T>): T {
-    return (this.entries.get(key.key) as T | undefined) ?? key.default
+class MemStore implements TokenStore {
+  readonly entries = new Map<string, string>()
+  private k(host: string, email: string): string {
+    return `${host} ${email}`
   }
 
-  set<T>(key: Key<T>, value: T): void {
-    this.entries.set(key.key, value)
+  read(host: string, email: string): string {
+    return this.entries.get(this.k(host, email)) ?? ''
   }
 
-  unset<T>(key: Key<T>): void {
-    this.entries.delete(key.key)
+  write(host: string, email: string, bearer: string): void {
+    this.entries.set(this.k(host, email), bearer)
+  }
+
+  remove(host: string, email: string): void {
+    this.entries.delete(this.k(host, email))
   }
 }
 
@@ -75,8 +78,7 @@ describe('runLogin', () => {
     const active = reg.resolveActive()
     expect(active?.ctx.account.email).toBe('tester@dify.ai')
     expect(active?.ctx.workspace?.id).toBe('550e8400-e29b-41d4-a716-446655440000')
-    expect(active?.ctx.available_workspaces).toHaveLength(2)
-    expect(store.get(tokenKey(active!.host, 'tester@dify.ai'))).toBe('dfoa_test')
+    expect(store.read(active!.host, 'tester@dify.ai')).toBe('dfoa_test')
 
     const hostsRaw = await readFile(join(configDir, 'hosts.yml'), 'utf8')
     expect(hostsRaw).toContain('current_host:')
@@ -109,7 +111,7 @@ describe('runLogin', () => {
     expect(active?.ctx.external_subject?.email).toBe('sso@dify.ai')
     expect(active?.ctx.external_subject?.issuer).toBe('https://issuer.example')
     expect(active?.ctx.account.email).toBe('')
-    expect(store.get(tokenKey(active!.host, 'sso@dify.ai'))).toBe('dfoe_test')
+    expect(store.read(active!.host, 'sso@dify.ai')).toBe('dfoe_test')
     expect(io.outBuf()).toContain('external SSO')
     expect(io.outBuf()).toContain('sso@dify.ai')
   })

cli/src/commands/auth/login/login.ts

@@ -1,7 +1,8 @@
 import type { Clock } from './device-flow.js'
 import type { CodeResponse, PollSuccess } from '@/api/oauth-device'
-import type { AccountContext, Workspace } from '@/auth/hosts'
-import type { StorageMode, Store } from '@/store/store'
+import type { AccountContext } from '@/auth/hosts'
+import type { StorageMode } from '@/store/store'
+import type { TokenStore } from '@/store/token-store'
 import type { ParseResult } from '@/sys/io/prompt'
 import type { IOStreams } from '@/sys/io/streams'
 import type { BrowserEnv, BrowserOpener } from '@/util/browser'
@@ -11,7 +12,7 @@ import { Registry } from '@/auth/hosts'
 import { BaseError, isBaseError } from '@/errors/base'
 import { ErrorCode } from '@/errors/codes'
 import { createHttpClient } from '@/http/client'
-import { getTokenStore, tokenKey } from '@/store/manager'
+import { detectTokenStore } from '@/store/manager'
 import { colorEnabled, colorScheme } from '@/sys/io/color'
 import { promptText } from '@/sys/io/prompt'
 import { startSpinner } from '@/sys/io/spinner'
@@ -25,7 +26,7 @@ export type LoginOptions = {
   readonly noBrowser?: boolean
   readonly insecure?: boolean
   readonly deviceLabel?: string
-  readonly store?: { readonly store: Store, readonly mode: StorageMode }
+  readonly store?: { readonly store: TokenStore, readonly mode: StorageMode }
   readonly api?: DeviceFlowApi
   readonly browserEnv?: BrowserEnv
   readonly browserOpener?: BrowserOpener
@@ -69,12 +70,12 @@ export async function runLogin(opts: LoginOptions): Promise<Registry> {
     spinner.stop()
   }
 
-  const storeBundle = opts.store ?? getTokenStore()
+  const storeBundle = opts.store ?? detectTokenStore()
   const display = bareHost(host)
   const email = accountEmail(success)
   const ctx = contextFromSuccess(success)
 
-  storeBundle.store.set(tokenKey(display, email), success.token)
+  storeBundle.store.write(display, email, success.token)
 
   const reg = Registry.load()
   reg.token_storage = storeBundle.mode
@@ -187,9 +188,6 @@ function contextFromSuccess(s: PollSuccess): AccountContext {
   const def = findDefaultWorkspace(s)
   if (def !== undefined)
     ctx.workspace = def
-  if (s.workspaces !== undefined && s.workspaces.length > 0) {
-    ctx.available_workspaces = s.workspaces.map<Workspace>(w => ({ id: w.id, name: w.name, role: w.role }))
-  }
   return ctx
 }
 

cli/src/commands/auth/logout/index.ts

@@ -3,7 +3,7 @@ import type { HttpClient } from '@/http/types'
 import { Registry } from '@/auth/hosts'
 import { DifyCommand } from '@/commands/_shared/dify-command'
 import { createHttpClient } from '@/http/client'
-import { getTokenStore, tokenKey } from '@/store/manager'
+import { getTokenStore } from '@/store/manager'
 import { runWithSpinner } from '@/sys/io/spinner'
 import { realStreams } from '@/sys/io/streams'
 import { hostWithScheme, openAPIBase } from '@/util/host'
@@ -26,7 +26,11 @@ export default class Logout extends DifyCommand {
 
     let http: HttpClient | undefined
     if (active !== undefined) {
-      const bearer = getTokenStore().store.get(tokenKey(active.host, active.email))
+      let bearer = ''
+      try {
+        bearer = getTokenStore(reg.token_storage).read(active.host, active.email)
+      }
+      catch { /* keyring locked — skip remote revocation, local cleanup still runs */ }
       if (bearer !== '') {
         http = createHttpClient({ baseURL: openAPIBase(hostWithScheme(active.host, active.scheme)), bearer, retryAttempts: 0 })
       }