Skip to content

Commit e029967

Browse files
jahoomajahooma
committed
Refresh claude oauth token as needed
1 parent 9ba28ea commit e029967

File tree

4 files changed

+183
-65
lines changed

4 files changed

+183
-65
lines changed

‎cli/src/utils/claude-oauth.ts‎

Lines changed: 13 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,7 @@
44

55
import crypto from 'crypto'
66
import open from 'open'
7-
import {
8-
CLAUDE_OAUTH_CLIENT_ID,
9-
CLAUDE_OAUTH_AUTHORIZE_URL,
10-
CLAUDE_OAUTH_TOKEN_URL,
11-
} from '@codebuff/common/constants/claude-oauth'
7+
import { CLAUDE_OAUTH_CLIENT_ID } from '@codebuff/common/constants/claude-oauth'
128
import {
139
saveClaudeOAuthCredentials,
1410
clearClaudeOAuthCredentials,
@@ -51,7 +47,7 @@ let pendingState: string | null = null
5147
export function startOAuthFlow(): { codeVerifier: string; authUrl: string } {
5248
const codeVerifier = generateCodeVerifier()
5349
const codeChallenge = generateCodeChallenge(codeVerifier)
54-
50+
5551
// Generate a random state parameter for CSRF protection
5652
const state = crypto.randomBytes(16).toString('hex')
5753

@@ -65,8 +61,14 @@ export function startOAuthFlow(): { codeVerifier: string; authUrl: string } {
6561
authUrl.searchParams.set('code', 'true')
6662
authUrl.searchParams.set('client_id', CLAUDE_OAUTH_CLIENT_ID)
6763
authUrl.searchParams.set('response_type', 'code')
68-
authUrl.searchParams.set('redirect_uri', 'https://console.anthropic.com/oauth/code/callback')
69-
authUrl.searchParams.set('scope', 'org:create_api_key user:profile user:inference')
64+
authUrl.searchParams.set(
65+
'redirect_uri',
66+
'https://console.anthropic.com/oauth/code/callback',
67+
)
68+
authUrl.searchParams.set(
69+
'scope',
70+
'org:create_api_key user:profile user:inference',
71+
)
7072
authUrl.searchParams.set('code_challenge', codeChallenge)
7173
authUrl.searchParams.set('code_challenge_method', 'S256')
7274
authUrl.searchParams.set('state', codeVerifier) // opencode uses verifier as state
@@ -92,7 +94,9 @@ export async function exchangeCodeForTokens(
9294
): Promise<ClaudeOAuthCredentials> {
9395
const verifier = codeVerifier ?? pendingCodeVerifier
9496
if (!verifier) {
95-
throw new Error('No code verifier found. Please start the OAuth flow again.')
97+
throw new Error(
98+
'No code verifier found. Please start the OAuth flow again.',
99+
)
96100
}
97101

98102
// The authorization code from claude.ai comes in format: code#state
@@ -140,55 +144,6 @@ export async function exchangeCodeForTokens(
140144
return credentials
141145
}
142146

143-
/**
144-
* Refresh the access token using the refresh token.
145-
*/
146-
export async function refreshAccessToken(): Promise<ClaudeOAuthCredentials | null> {
147-
const credentials = getClaudeOAuthCredentials()
148-
if (!credentials?.refreshToken) {
149-
return null
150-
}
151-
152-
try {
153-
// Use the v1 OAuth token endpoint (same as opencode)
154-
const response = await fetch('https://console.anthropic.com/v1/oauth/token', {
155-
method: 'POST',
156-
headers: {
157-
'Content-Type': 'application/json',
158-
},
159-
body: JSON.stringify({
160-
grant_type: 'refresh_token',
161-
refresh_token: credentials.refreshToken,
162-
client_id: CLAUDE_OAUTH_CLIENT_ID,
163-
}),
164-
})
165-
166-
if (!response.ok) {
167-
// Refresh failed, clear credentials
168-
clearClaudeOAuthCredentials()
169-
return null
170-
}
171-
172-
const data = await response.json()
173-
174-
const newCredentials: ClaudeOAuthCredentials = {
175-
accessToken: data.access_token,
176-
refreshToken: data.refresh_token ?? credentials.refreshToken,
177-
expiresAt: Date.now() + data.expires_in * 1000,
178-
connectedAt: credentials.connectedAt,
179-
}
180-
181-
// Save updated credentials
182-
saveClaudeOAuthCredentials(newCredentials)
183-
184-
return newCredentials
185-
} catch {
186-
// Refresh failed, clear credentials
187-
clearClaudeOAuthCredentials()
188-
return null
189-
}
190-
}
191-
192147
/**
193148
* Disconnect from Claude OAuth (clear credentials).
194149
*/

‎sdk/src/credentials.ts‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import path from 'node:path'
33
import os from 'os'
44

55
import { env } from '@codebuff/common/env'
6+
import { CLAUDE_OAUTH_CLIENT_ID } from '@codebuff/common/constants/claude-oauth'
67
import { userSchema } from '@codebuff/common/util/credentials'
78
import { z } from 'zod/v4'
89

@@ -211,3 +212,103 @@ export const isClaudeOAuthValid = (
211212
const bufferMs = 5 * 60 * 1000
212213
return credentials.expiresAt > Date.now() + bufferMs
213214
}
215+
216+
// Mutex to prevent concurrent refresh attempts
217+
let refreshPromise: Promise<ClaudeOAuthCredentials | null> | null = null
218+
219+
/**
220+
* Refresh the Claude OAuth access token using the refresh token.
221+
* Returns the new credentials if successful, null if refresh fails.
222+
* Uses a mutex to prevent concurrent refresh attempts.
223+
*/
224+
export const refreshClaudeOAuthToken = async (
225+
clientEnv: ClientEnv = env,
226+
): Promise<ClaudeOAuthCredentials | null> => {
227+
// If a refresh is already in progress, wait for it
228+
if (refreshPromise) {
229+
return refreshPromise
230+
}
231+
232+
const credentials = getClaudeOAuthCredentials(clientEnv)
233+
if (!credentials?.refreshToken) {
234+
return null
235+
}
236+
237+
// Start the refresh and store the promise
238+
refreshPromise = (async () => {
239+
try {
240+
const response = await fetch('https://console.anthropic.com/v1/oauth/token', {
241+
method: 'POST',
242+
headers: {
243+
'Content-Type': 'application/json',
244+
},
245+
body: JSON.stringify({
246+
grant_type: 'refresh_token',
247+
refresh_token: credentials.refreshToken,
248+
client_id: CLAUDE_OAUTH_CLIENT_ID,
249+
}),
250+
})
251+
252+
if (!response.ok) {
253+
// Refresh failed, clear credentials
254+
clearClaudeOAuthCredentials(clientEnv)
255+
return null
256+
}
257+
258+
const data = await response.json()
259+
260+
const newCredentials: ClaudeOAuthCredentials = {
261+
accessToken: data.access_token,
262+
refreshToken: data.refresh_token ?? credentials.refreshToken,
263+
expiresAt: Date.now() + data.expires_in * 1000,
264+
connectedAt: credentials.connectedAt,
265+
}
266+
267+
// Save updated credentials
268+
saveClaudeOAuthCredentials(newCredentials, clientEnv)
269+
270+
return newCredentials
271+
} catch {
272+
// Refresh failed, clear credentials
273+
clearClaudeOAuthCredentials(clientEnv)
274+
return null
275+
} finally {
276+
// Clear the mutex after completion
277+
refreshPromise = null
278+
}
279+
})()
280+
281+
return refreshPromise
282+
}
283+
284+
/**
285+
* Get valid Claude OAuth credentials, refreshing if necessary.
286+
* This is the main function to use when you need credentials for an API call.
287+
*
288+
* - Returns credentials immediately if valid (>5 min until expiry)
289+
* - Attempts refresh if token is expired or near-expiry
290+
* - Returns null if no credentials or refresh fails
291+
*/
292+
export const getValidClaudeOAuthCredentials = async (
293+
clientEnv: ClientEnv = env,
294+
): Promise<ClaudeOAuthCredentials | null> => {
295+
const credentials = getClaudeOAuthCredentials(clientEnv)
296+
if (!credentials) {
297+
return null
298+
}
299+
300+
// Check if token is from environment variable (synthetic credentials, no refresh needed)
301+
if (!credentials.refreshToken) {
302+
// Environment variable tokens are assumed valid
303+
return credentials
304+
}
305+
306+
// Check if token is valid with 5 minute buffer
307+
const bufferMs = 5 * 60 * 1000
308+
if (credentials.expiresAt > Date.now() + bufferMs) {
309+
return credentials
310+
}
311+
312+
// Token is expired or expiring soon, try to refresh
313+
return refreshClaudeOAuthToken(clientEnv)
314+
}

‎sdk/src/impl/llm.ts‎

Lines changed: 62 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,43 @@ function isClaudeOAuthRateLimitError(error: unknown): boolean {
143143
return false
144144
}
145145

146+
/**
147+
* Check if an error is a Claude OAuth authentication error (expired/invalid token).
148+
* This indicates we should try refreshing the token.
149+
*/
150+
function isClaudeOAuthAuthError(error: unknown): boolean {
151+
if (!error || typeof error !== 'object') return false
152+
153+
const err = error as {
154+
statusCode?: number
155+
message?: string
156+
responseBody?: string
157+
}
158+
159+
// 401 Unauthorized or 403 Forbidden typically indicate auth issues
160+
if (err.statusCode === 401 || err.statusCode === 403) return true
161+
162+
const message = (err.message || '').toLowerCase()
163+
const responseBody = (err.responseBody || '').toLowerCase()
164+
165+
if (message.includes('unauthorized') || message.includes('invalid_token'))
166+
return true
167+
if (message.includes('authentication') || message.includes('expired'))
168+
return true
169+
if (
170+
responseBody.includes('unauthorized') ||
171+
responseBody.includes('invalid_token')
172+
)
173+
return true
174+
if (
175+
responseBody.includes('authentication') ||
176+
responseBody.includes('expired')
177+
)
178+
return true
179+
180+
return false
181+
}
182+
146183
export async function* promptAiSdkStream(
147184
params: ParamsOf<PromptAiSdkStreamFn> & {
148185
skipClaudeOAuth?: boolean
@@ -169,7 +206,7 @@ export async function* promptAiSdkStream(
169206
model: params.model,
170207
skipClaudeOAuth: params.skipClaudeOAuth,
171208
}
172-
const { model: aiSDKModel, isClaudeOAuth } = getModelForRequest(modelParams)
209+
const { model: aiSDKModel, isClaudeOAuth } = await getModelForRequest(modelParams)
173210

174211
// Notify about Claude OAuth usage
175212
if (isClaudeOAuth && params.onClaudeOAuthStatusChange) {
@@ -377,6 +414,28 @@ export async function* promptAiSdkStream(
377414
return fallbackResult
378415
}
379416

417+
// Check if this is a Claude OAuth authentication error (expired token) - only fall back if no content yielded yet
418+
if (
419+
isClaudeOAuth &&
420+
!params.skipClaudeOAuth &&
421+
!hasYieldedContent &&
422+
isClaudeOAuthAuthError(chunkValue.error)
423+
) {
424+
logger.info(
425+
{ error: getErrorObject(chunkValue.error) },
426+
'Claude OAuth auth error during stream, falling back to Codebuff backend',
427+
)
428+
if (params.onClaudeOAuthStatusChange) {
429+
params.onClaudeOAuthStatusChange(false)
430+
}
431+
// Retry with Codebuff backend (skipClaudeOAuth will bypass the failed OAuth)
432+
const fallbackResult = yield* promptAiSdkStream({
433+
...params,
434+
skipClaudeOAuth: true,
435+
})
436+
return fallbackResult
437+
}
438+
380439
logger.error(
381440
{
382441
chunk: { ...chunkValue, error: undefined },
@@ -495,7 +554,7 @@ export async function promptAiSdk(
495554
model: params.model,
496555
skipClaudeOAuth: true, // Always use Codebuff backend for non-streaming
497556
}
498-
const { model: aiSDKModel } = getModelForRequest(modelParams)
557+
const { model: aiSDKModel } = await getModelForRequest(modelParams)
499558

500559
const response = await generateText({
501560
...params,
@@ -552,7 +611,7 @@ export async function promptAiSdkStructured<T>(
552611
model: params.model,
553612
skipClaudeOAuth: true, // Always use Codebuff backend for non-streaming
554613
}
555-
const { model: aiSDKModel } = getModelForRequest(modelParams)
614+
const { model: aiSDKModel } = await getModelForRequest(modelParams)
556615

557616
const response = await generateObject<z.ZodType<T>, 'object'>({
558617
...params,

‎sdk/src/impl/model-provider.ts‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ import {
2121
} from '@codebuff/internal/openai-compatible/index'
2222

2323
import { WEBSITE_URL } from '../constants'
24-
import { getClaudeOAuthCredentials, isClaudeOAuthValid } from '../credentials'
24+
import { getValidClaudeOAuthCredentials } from '../credentials'
2525
import { getByokOpenrouterApiKeyFromEnv } from '../env'
2626

2727
import type { LanguageModel } from 'ai'
@@ -61,13 +61,16 @@ type OpenRouterUsageAccounting = {
6161
*
6262
* If Claude OAuth credentials are available and the model is a Claude model,
6363
* returns an Anthropic direct model. Otherwise, returns the Codebuff backend model.
64+
*
65+
* This function is async because it may need to refresh the OAuth token.
6466
*/
65-
export function getModelForRequest(params: ModelRequestParams): ModelResult {
67+
export async function getModelForRequest(params: ModelRequestParams): Promise<ModelResult> {
6668
const { apiKey, model, skipClaudeOAuth } = params
6769

6870
// Check if we should use Claude OAuth direct
69-
if (!skipClaudeOAuth && isClaudeModel(model) && isClaudeOAuthValid()) {
70-
const claudeOAuthCredentials = getClaudeOAuthCredentials()
71+
if (!skipClaudeOAuth && isClaudeModel(model)) {
72+
// Get valid credentials (will refresh if needed)
73+
const claudeOAuthCredentials = await getValidClaudeOAuthCredentials()
7174
if (claudeOAuthCredentials) {
7275
return {
7376
model: createAnthropicOAuthModel(

0 commit comments

Comments
 (0)