fix(cli): revalidate auth when token changes

- Key auth validation query by sha256(apiKey) to avoid stale cache and avoid putting the raw token in the cache key

- Resolve CODEBUFF_API_KEY via CiEnv and make auth token resolution injectable

Author: brandonkachen

cli/src/hooks/__tests__/use-auth-query.test.ts

@@ -0,0 +1,22 @@
+import { describe, expect, test } from 'bun:test'
+
+import { authQueryKeys } from '../use-auth-query'
+
+describe('authQueryKeys.validation', () => {
+  test('changes when api key changes', () => {
+    const firstKey = authQueryKeys.validation('token-1')
+    const secondKey = authQueryKeys.validation('token-2')
+
+    expect(firstKey).not.toEqual(secondKey)
+  })
+
+  test('does not include the raw api key', () => {
+    const token = 'secret-token-123'
+    const key = authQueryKeys.validation(token)
+    const [, , apiKeyHash] = key
+
+    expect(key).not.toContain(token)
+    expect(apiKeyHash).toMatch(/^[0-9a-f]{64}$/)
+  })
+})
+

cli/src/hooks/use-auth-query.ts

@@ -1,4 +1,6 @@
-import { API_KEY_ENV_VAR } from '@codebuff/common/old-constants'
+import { createHash } from 'crypto'
+
+import { getCiEnv } from '@codebuff/common/env-ci'
 import {
   AuthenticationError,
   ErrorCodes,
@@ -22,12 +24,16 @@ import { logger as defaultLogger, loggerContext } from '../utils/logger'
 import type { GetUserInfoFromApiKeyFn } from '@codebuff/common/types/contracts/database'
 import type { Logger } from '@codebuff/common/types/contracts/logger'
 
+const getApiKeyHash = (apiKey: string): string => {
+  return createHash('sha256').update(apiKey).digest('hex')
+}
+
 // Query keys for type-safe cache management
 export const authQueryKeys = {
   all: ['auth'] as const,
   user: () => [...authQueryKeys.all, 'user'] as const,
   validation: (apiKey: string) =>
-    [...authQueryKeys.all, 'validation', apiKey] as const,
+    [...authQueryKeys.all, 'validation', getApiKeyHash(apiKey)] as const,
 }
 
 interface ValidateAuthParams {
@@ -122,8 +128,7 @@ export function useAuthQuery(deps: UseAuthQueryDeps = {}) {
   } = deps
 
   const userCredentials = getUserCredentials()
-  const apiKey =
-    userCredentials?.authToken || process.env[API_KEY_ENV_VAR] || ''
+  const apiKey = userCredentials?.authToken || getCiEnv().CODEBUFF_API_KEY || ''
 
   return useQuery({
     queryKey: authQueryKeys.validation(apiKey),

cli/src/hooks/use-auth-state.ts

@@ -10,6 +10,17 @@ import { loggerContext } from '../utils/logger'
 import type { MultilineInputHandle } from '../components/multiline-input'
 import type { User } from '../utils/auth'
 
+const setAuthLoggerContext = (params: { userId: string; email: string }) => {
+  loggerContext.userId = params.userId
+  loggerContext.userEmail = params.email
+  identifyUser(params.userId, { email: params.email })
+}
+
+const clearAuthLoggerContext = () => {
+  delete loggerContext.userId
+  delete loggerContext.userEmail
+}
+
 interface UseAuthStateOptions {
   requireAuth: boolean | null
   inputRef: React.MutableRefObject<MultilineInputHandle | null>
@@ -27,8 +38,7 @@ export const useAuthState = ({
   const logoutMutation = useLogoutMutation()
   const { resetLoginState } = useLoginStore()
 
-  const initialAuthState =
-    requireAuth === false ? true : requireAuth === true ? false : null
+  const initialAuthState = requireAuth === null ? null : !requireAuth
   const [isAuthenticated, setIsAuthenticated] = useState<boolean | null>(
     initialAuthState,
   )
@@ -55,23 +65,15 @@ export const useAuthState = ({
           authToken: userCredentials?.authToken || '',
         }
         setUser(userData)
-        
-        // Set logger context for analytics
-        loggerContext.userId = authQuery.data.id
-        loggerContext.userEmail = authQuery.data.email
-        
-        // Identify user with PostHog
-        identifyUser(authQuery.data.id, {
-          email: authQuery.data.email,
+        setAuthLoggerContext({
+          userId: authQuery.data.id,
+          email: authQuery.data.email || '',
         })
       }
     } else if (authQuery.isError) {
       setIsAuthenticated(false)
       setUser(null)
-      
-      // Clear logger context on auth error
-      delete loggerContext.userId
-      delete loggerContext.userEmail
+      clearAuthLoggerContext()
     }
   }, [authQuery.isSuccess, authQuery.isError, authQuery.data, user])
 
@@ -85,14 +87,10 @@ export const useAuthState = ({
       setInputFocused(true)
       setUser(loggedInUser)
       setIsAuthenticated(true)
-      
-      // Set logger context for analytics
+
       if (loggedInUser.id && loggedInUser.email) {
-        loggerContext.userId = loggedInUser.id
-        loggerContext.userEmail = loggedInUser.email
-        
-        // Identify user with PostHog
-        identifyUser(loggedInUser.id, {
+        setAuthLoggerContext({
+          userId: loggedInUser.id,
           email: loggedInUser.email,
         })
       }

cli/src/index.tsx

@@ -3,7 +3,6 @@
 import { promises as fs } from 'fs'
 import { createRequire } from 'module'
 
-import { API_KEY_ENV_VAR } from '@codebuff/common/old-constants'
 import { getProjectFileTree } from '@codebuff/common/project-file-tree'
 import { createCliRenderer } from '@opentui/core'
 import { createRoot } from '@opentui/react'
@@ -21,7 +20,8 @@ import { handlePublish } from './commands/publish'
 import { initializeApp } from './init/init-app'
 import { getProjectRoot } from './project-files'
 import { initAnalytics } from './utils/analytics'
-import { getUserCredentials } from './utils/auth'
+import { getAuthTokenDetails } from './utils/auth'
+import { getCliEnv } from './utils/env'
 import { initializeAgentRegistry } from './utils/local-agent-registry'
 import { clearLogFile, logger } from './utils/logger'
 import { detectTerminalTheme } from './utils/terminal-color-detection'
@@ -33,8 +33,9 @@ import type { AgentMode } from './utils/constants'
 const require = createRequire(import.meta.url)
 
 function loadPackageVersion(): string {
-  if (process.env.CODEBUFF_CLI_VERSION) {
-    return process.env.CODEBUFF_CLI_VERSION
+  const env = getCliEnv()
+  if (env.CODEBUFF_CLI_VERSION) {
+    return env.CODEBUFF_CLI_VERSION
   }
 
   try {
@@ -215,9 +216,7 @@ async function main(): Promise<void> {
     const [fileTree, setFileTree] = React.useState<FileTreeNode[]>([])
 
     React.useEffect(() => {
-      const userCredentials = getUserCredentials()
-      const apiKey =
-        userCredentials?.authToken || process.env[API_KEY_ENV_VAR] || ''
+      const apiKey = getAuthTokenDetails().token ?? ''
 
       if (!apiKey) {
         setRequireAuth(true)

cli/src/utils/auth.ts

@@ -2,10 +2,12 @@ import fs from 'fs'
 import os from 'os'
 import path from 'path'
 
+import { getCiEnv } from '@codebuff/common/env-ci'
 import { env } from '@codebuff/common/env'
-import { API_KEY_ENV_VAR } from '@codebuff/common/old-constants'
 import { z } from 'zod'
 
+import type { CiEnv } from '@codebuff/common/types/contracts/env'
+
 import { getApiClient, setApiClientAuthToken } from './codebuff-api'
 import { logger } from './logger'
 
@@ -107,13 +109,15 @@ export interface AuthTokenDetails {
 /**
  * Resolve the auth token and track where it came from.
  */
-export const getAuthTokenDetails = (): AuthTokenDetails => {
+export const getAuthTokenDetails = (
+  ciEnv: CiEnv = getCiEnv(),
+): AuthTokenDetails => {
   const userCredentials = getUserCredentials()
   if (userCredentials?.authToken) {
     return { token: userCredentials.authToken, source: 'credentials' }
   }
 
-  const envToken = process.env[API_KEY_ENV_VAR]
+  const envToken = ciEnv.CODEBUFF_API_KEY
   if (envToken) {
     return { token: envToken, source: 'environment' }
   }