diff --git a/desktop/src/components/settings/ClaudeOfficialLogin.tsx b/desktop/src/components/settings/ClaudeOfficialLogin.tsx index c6e1aa00..9d5bfb63 100644 --- a/desktop/src/components/settings/ClaudeOfficialLogin.tsx +++ b/desktop/src/components/settings/ClaudeOfficialLogin.tsx @@ -11,8 +11,16 @@ import { useTranslation } from '../../i18n' export function ClaudeOfficialLogin() { const t = useTranslation() - const { status, isLoading, error, fetchStatus, login, logout, stopPolling } = - useHahaOAuthStore() + const { + status, + isLoading, + error, + fetchStatus, + login, + logout, + startPolling, + stopPolling, + } = useHahaOAuthStore() useEffect(() => { fetchStatus() @@ -24,6 +32,7 @@ export function ClaudeOfficialLogin() { const { authorizeUrl } = await login() try { await shellOpen(authorizeUrl) + startPolling() } catch (err) { console.error('[ClaudeOfficialLogin] shellOpen failed:', err) useHahaOAuthStore.setState({ diff --git a/desktop/src/stores/hahaOAuthStore.test.ts b/desktop/src/stores/hahaOAuthStore.test.ts new file mode 100644 index 00000000..4a230279 --- /dev/null +++ b/desktop/src/stores/hahaOAuthStore.test.ts @@ -0,0 +1,77 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +const { startMock, statusMock, logoutMock } = vi.hoisted(() => ({ + startMock: vi.fn(), + statusMock: vi.fn(), + logoutMock: vi.fn(), +})) + +vi.mock('../api/hahaOAuth', () => ({ + hahaOAuthApi: { + start: startMock, + status: statusMock, + logout: logoutMock, + }, +})) + +import { useHahaOAuthStore } from './hahaOAuthStore' + +const initialState = useHahaOAuthStore.getState() + +describe('hahaOAuthStore', () => { + beforeEach(() => { + vi.useFakeTimers() + startMock.mockReset() + statusMock.mockReset() + logoutMock.mockReset() + useHahaOAuthStore.setState({ + ...initialState, + status: null, + isPolling: false, + isLoading: false, + error: null, + }) + }) + + afterEach(() => { + useHahaOAuthStore.getState().stopPolling() + useHahaOAuthStore.setState(initialState) + vi.useRealTimers() + }) + + it('login does not start polling until the browser launch succeeds', async () => { + startMock.mockResolvedValue({ + authorizeUrl: 'http://localhost:3456/api/haha-oauth/callback', + state: 'state-123', + }) + + const result = await useHahaOAuthStore.getState().login() + + expect(result.authorizeUrl).toContain('/api/haha-oauth/callback') + expect(useHahaOAuthStore.getState().isPolling).toBe(false) + }) + + it('startPolling stops after the status becomes logged in', async () => { + statusMock + .mockResolvedValueOnce({ loggedIn: false }) + .mockResolvedValueOnce({ + loggedIn: true, + expiresAt: Date.now() + 60_000, + scopes: ['user:inference'], + subscriptionType: 'max', + }) + + useHahaOAuthStore.getState().startPolling() + expect(useHahaOAuthStore.getState().isPolling).toBe(true) + + await vi.advanceTimersByTimeAsync(2_000) + expect(useHahaOAuthStore.getState().isPolling).toBe(true) + + await vi.advanceTimersByTimeAsync(2_000) + expect(useHahaOAuthStore.getState().status).toMatchObject({ + loggedIn: true, + subscriptionType: 'max', + }) + expect(useHahaOAuthStore.getState().isPolling).toBe(false) + }) +}) diff --git a/desktop/src/stores/hahaOAuthStore.ts b/desktop/src/stores/hahaOAuthStore.ts index df6f9f98..64642eb3 100644 --- a/desktop/src/stores/hahaOAuthStore.ts +++ b/desktop/src/stores/hahaOAuthStore.ts @@ -41,7 +41,6 @@ export const useHahaOAuthStore = create((set, get) => { try { const res = await hahaOAuthApi.start() set({ isLoading: false }) - get().startPolling() return { authorizeUrl: res.authorizeUrl } } catch (err) { set({ diff --git a/src/server/__tests__/conversation-service.test.ts b/src/server/__tests__/conversation-service.test.ts index 63d5c7aa..210846e3 100644 --- a/src/server/__tests__/conversation-service.test.ts +++ b/src/server/__tests__/conversation-service.test.ts @@ -11,6 +11,7 @@ describe('ConversationService', () => { let originalBaseUrl: string | undefined let originalModel: string | undefined let originalEntrypoint: string | undefined + let originalOAuthToken: string | undefined beforeEach(async () => { tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), 'cc-haha-conversation-service-')) @@ -19,11 +20,13 @@ describe('ConversationService', () => { originalBaseUrl = process.env.ANTHROPIC_BASE_URL originalModel = process.env.ANTHROPIC_MODEL originalEntrypoint = process.env.CLAUDE_CODE_ENTRYPOINT + originalOAuthToken = process.env.CLAUDE_CODE_OAUTH_TOKEN process.env.CLAUDE_CONFIG_DIR = tmpDir process.env.ANTHROPIC_AUTH_TOKEN = 'test-token' process.env.ANTHROPIC_BASE_URL = 'https://example.invalid/anthropic' process.env.ANTHROPIC_MODEL = 'test-model' + process.env.CLAUDE_CODE_OAUTH_TOKEN = 'inherited-parent-oauth-token' // Clear inherited CLAUDE_CODE_ENTRYPOINT so tests can assert whether // buildChildEnv injects it or not without interference from the shell env. delete process.env.CLAUDE_CODE_ENTRYPOINT @@ -45,6 +48,9 @@ describe('ConversationService', () => { if (originalEntrypoint === undefined) delete process.env.CLAUDE_CODE_ENTRYPOINT else process.env.CLAUDE_CODE_ENTRYPOINT = originalEntrypoint + if (originalOAuthToken === undefined) delete process.env.CLAUDE_CODE_OAUTH_TOKEN + else process.env.CLAUDE_CODE_OAUTH_TOKEN = originalOAuthToken + await fs.rm(tmpDir, { recursive: true, force: true }) }) @@ -124,6 +130,22 @@ describe('ConversationService', () => { expect(env.CLAUDE_CODE_ENTRYPOINT).toBeUndefined() }) + test('buildChildEnv does not leak inherited CLAUDE_CODE_OAUTH_TOKEN when official token is unavailable', async () => { + const ccHahaDir = path.join(tmpDir, 'cc-haha') + await fs.mkdir(ccHahaDir, { recursive: true }) + await fs.writeFile( + path.join(ccHahaDir, 'settings.json'), + JSON.stringify({ env: {} }), + 'utf-8', + ) + + const service = new ConversationService() as any + const env = (await service.buildChildEnv('/tmp')) as Record + + expect(env.CLAUDE_CODE_ENTRYPOINT).toBe('claude-desktop') + expect(env.CLAUDE_CODE_OAUTH_TOKEN).toBeUndefined() + }) + test('uses bun entrypoint fallback on Windows dev mode', () => { const service = new ConversationService() as any const args = service.resolveCliArgs(['--print']) diff --git a/src/server/__tests__/haha-oauth-api.test.ts b/src/server/__tests__/haha-oauth-api.test.ts index 69f992ef..422a4486 100644 --- a/src/server/__tests__/haha-oauth-api.test.ts +++ b/src/server/__tests__/haha-oauth-api.test.ts @@ -104,6 +104,25 @@ describe('GET /api/haha-oauth/status', () => { expect(JSON.stringify(data)).not.toContain('sk-ant-oat01') expect(JSON.stringify(data)).not.toContain('sk-ant-ort01') }) + + test('returns loggedIn=false when stored token is expired and refresh fails', async () => { + await hahaOAuthService.saveTokens({ + accessToken: 'expired-token', + refreshToken: 'revoked-refresh-token', + expiresAt: Date.now() - 1_000, + scopes: ['user:inference'], + subscriptionType: 'max', + }) + hahaOAuthService.setRefreshFn(async () => { + throw new Error('refresh revoked') + }) + + const { req, url, segments } = buildReq('GET', '/api/haha-oauth/status') + const res = await handleHahaOAuthApi(req, url, segments) + + expect(res.status).toBe(200) + expect(await res.json()).toEqual({ loggedIn: false }) + }) }) describe('DELETE /api/haha-oauth', () => { diff --git a/src/server/__tests__/haha-oauth-service.test.ts b/src/server/__tests__/haha-oauth-service.test.ts index 90c9e055..659991ee 100644 --- a/src/server/__tests__/haha-oauth-service.test.ts +++ b/src/server/__tests__/haha-oauth-service.test.ts @@ -101,6 +101,24 @@ describe('HahaOAuthService — session management', () => { expect(service.consumeSession(session.state)).not.toBeNull() expect(service.getSession(session.state)).toBeNull() }) + + test('completeSession stores subscription type fetched from profile info', async () => { + const session = service.startSession({ serverPort: 54321 }) + ;(service as any).exchangeWithCustomCallback = async () => ({ + access_token: 'fresh-access-token', + refresh_token: 'fresh-refresh-token', + expires_in: 3600, + scope: 'user:inference', + }) + service.setFetchProfileFn(async () => ({ + subscriptionType: 'team', + })) + + const tokens = await service.completeSession('authorization-code', session.state) + + expect(tokens.subscriptionType).toBe('team') + expect((await service.loadTokens())?.subscriptionType).toBe('team') + }) }) describe('HahaOAuthService — ensureFreshAccessToken', () => { diff --git a/src/server/api/haha-oauth.ts b/src/server/api/haha-oauth.ts index dfe54d7b..dbedb661 100644 --- a/src/server/api/haha-oauth.ts +++ b/src/server/api/haha-oauth.ts @@ -73,7 +73,7 @@ export async function handleHahaOAuthApi( } if ((action === undefined || action === 'status') && req.method === 'GET') { - const tokens = await hahaOAuthService.loadTokens() + const tokens = await hahaOAuthService.ensureFreshTokens() if (!tokens) { return Response.json({ loggedIn: false }) } diff --git a/src/server/services/conversationService.ts b/src/server/services/conversationService.ts index 4b9ce831..2ae4dd62 100644 --- a/src/server/services/conversationService.ts +++ b/src/server/services/conversationService.ts @@ -490,6 +490,7 @@ export class ConversationService { ] as const const cleanEnv = { ...process.env } + delete cleanEnv.CLAUDE_CODE_OAUTH_TOKEN if (this.shouldStripInheritedProviderEnv()) { for (const key of PROVIDER_ENV_KEYS) { delete cleanEnv[key] diff --git a/src/server/services/hahaOAuthService.ts b/src/server/services/hahaOAuthService.ts index 227f59fd..9fb77478 100644 --- a/src/server/services/hahaOAuthService.ts +++ b/src/server/services/hahaOAuthService.ts @@ -19,6 +19,7 @@ import { } from '../../services/oauth/crypto.js' import { buildAuthUrl, + fetchProfileInfo, refreshOAuthToken, isOAuthTokenExpired, parseScopes, @@ -47,17 +48,25 @@ export type OAuthSession = { } type RefreshFn = (refreshToken: string, opts?: { scopes?: string[] }) => Promise +type FetchProfileFn = ( + accessToken: string, +) => Promise<{ subscriptionType: SubscriptionType | null }> const SESSION_TTL_MS = 5 * 60 * 1000 export class HahaOAuthService { private sessions = new Map() private refreshFn: RefreshFn = refreshOAuthToken + private fetchProfileFn: FetchProfileFn = fetchProfileInfo setRefreshFn(fn: RefreshFn): void { this.refreshFn = fn } + setFetchProfileFn(fn: FetchProfileFn): void { + this.fetchProfileFn = fn + } + private getOAuthFilePath(): string { const configDir = process.env.CLAUDE_CONFIG_DIR || path.join(os.homedir(), '.claude') @@ -166,13 +175,14 @@ export class HahaOAuthService { session.codeVerifier, session.serverPort, ) + const profile = await this.fetchProfileFn(response.access_token) const tokens: StoredOAuthTokens = { accessToken: response.access_token, refreshToken: response.refresh_token ?? null, expiresAt: Date.now() + response.expires_in * 1000, scopes: parseScopes(response.scope), - subscriptionType: null, + subscriptionType: profile.subscriptionType, } await this.saveTokens(tokens) return tokens @@ -214,13 +224,13 @@ export class HahaOAuthService { return (await res.json()) as OAuthTokenExchangeResponse } - async ensureFreshAccessToken(): Promise { + async ensureFreshTokens(): Promise { const tokens = await this.loadTokens() if (!tokens) return null - if (tokens.expiresAt === null) return tokens.accessToken + if (tokens.expiresAt === null) return tokens - if (!isOAuthTokenExpired(tokens.expiresAt)) return tokens.accessToken + if (!isOAuthTokenExpired(tokens.expiresAt)) return tokens if (!tokens.refreshToken) return null @@ -236,7 +246,7 @@ export class HahaOAuthService { subscriptionType: refreshed.subscriptionType ?? tokens.subscriptionType, } await this.saveTokens(updated) - return updated.accessToken + return updated } catch (err) { console.error( '[HahaOAuthService] token refresh failed:', @@ -245,6 +255,11 @@ export class HahaOAuthService { return null } } + + async ensureFreshAccessToken(): Promise { + const tokens = await this.ensureFreshTokens() + return tokens?.accessToken ?? null + } } export const hahaOAuthService = new HahaOAuthService()