feat: auto-refresh OAuth tokens before each API call
Wraps OAuth clients with a RefreshingClient that checks token expiry before every Stream call. Refreshes transparently and rebuilds the underlying client with the fresh token. Fixes sessions silently dying after the 1-hour token lifetime.
- commit
- 8841800acdec3c2b73de434f7cc49e4a480e2b3e
- author
- patriceckhart · · 2026-04-24 17:37 UTC
- parent
- cb9de10
- signature
- unsigned
- refs
- tag: v0.1.9
3 files · +110 -2
M internal/agent/build.go +33 -2
| @@ -1,6 +1,7 @@ | |||
| 1 | 1 | package agent | |
| 2 | 2 | ||
| 3 | 3 | import ( | |
| 4 | + | "context" | |
| 4 | 5 | "fmt" | |
| 5 | 6 | "os" | |
| 6 | 7 | "path/filepath" | |
| @@ -337,17 +338,47 @@ func (r Resolved) NewClient() provider.Client { | |||
| 337 | 338 | return provider.NewOpenAI(r.Credential, r.BaseURL) | |
| 338 | 339 | case "openai": | |
| 339 | 340 | if r.AuthMethod == "oauth" { | |
| 340 | - | return provider.NewOpenAICodex(r.Credential, r.AccountID, r.BaseURL) | |
| 341 | + | inner := provider.NewOpenAICodex(r.Credential, r.AccountID, r.BaseURL) | |
| 342 | + | return r.wrapWithRefresh(inner) | |
| 341 | 343 | } | |
| 342 | 344 | return provider.NewOpenAI(r.Credential, r.BaseURL) | |
| 343 | 345 | default: | |
| 344 | 346 | if r.AuthMethod == "oauth" { | |
| 345 | - | return provider.NewAnthropicOAuth(r.Credential, r.BaseURL) | |
| 347 | + | inner := provider.NewAnthropicOAuth(r.Credential, r.BaseURL) | |
| 348 | + | return r.wrapWithRefresh(inner) | |
| 346 | 349 | } | |
| 347 | 350 | return provider.NewAnthropic(r.Credential, r.BaseURL) | |
| 348 | 351 | } | |
| 349 | 352 | } | |
| 350 | 353 | ||
| 354 | + | // wrapWithRefresh wraps an OAuth client so the access token is | |
| 355 | + | // refreshed automatically before each API call. Without this, long | |
| 356 | + | // sessions (hours) silently fail when the 1-hour token expires. | |
| 357 | + | func (r Resolved) wrapWithRefresh(inner provider.Client) provider.Client { | |
| 358 | + | provName := r.Provider | |
| 359 | + | baseURL := r.BaseURL | |
| 360 | + | accountID := r.AccountID | |
| 361 | + | ||
| 362 | + | refreshFn := func(ctx context.Context) (string, error) { | |
| 363 | + | tok, err := refreshIfExpired(provName, loadOAuthToken(provName)) | |
| 364 | + | if err != nil { | |
| 365 | + | return "", err | |
| 366 | + | } | |
| 367 | + | return tok.AccessToken, nil | |
| 368 | + | } | |
| 369 | + | ||
| 370 | + | factory := func(token string) provider.Client { | |
| 371 | + | switch provName { | |
| 372 | + | case "openai": | |
| 373 | + | return provider.NewOpenAICodex(token, accountID, baseURL) | |
| 374 | + | default: | |
| 375 | + | return provider.NewAnthropicOAuth(token, baseURL) | |
| 376 | + | } | |
| 377 | + | } | |
| 378 | + | ||
| 379 | + | return provider.NewRefreshingClient(inner, refreshFn, factory) | |
| 380 | + | } | |
| 381 | + | ||
| 351 | 382 | // UseSandbox replaces the sandbox pointer that every tool in r's | |
| 352 | 383 | // registry references. Used to keep the /jail state stable across | |
| 353 | 384 | // agent rebuilds (e.g. /login, /model switching providers). | |
| @@ -1,6 +1,7 @@ | |||||
| 1 | package agent | 1 | package agent | ||
| 2 | 2 | ||||
| 3 | import ( | 3 | import ( | ||
| 4 | + | "context" | |||
| 4 | "fmt" | 5 | "fmt" | ||
| 5 | "os" | 6 | "os" | ||
| 6 | "path/filepath" | 7 | "path/filepath" | ||
| @@ -337,17 +338,47 @@ func (r Resolved) NewClient() provider.Client { | |||||
| 337 | return provider.NewOpenAI(r.Credential, r.BaseURL) | 338 | return provider.NewOpenAI(r.Credential, r.BaseURL) | ||
| 338 | case "openai": | 339 | case "openai": | ||
| 339 | if r.AuthMethod == "oauth" { | 340 | if r.AuthMethod == "oauth" { | ||
| 340 | - | return provider.NewOpenAICodex(r.Credential, r.AccountID, r.BaseURL) | 341 | + | inner := provider.NewOpenAICodex(r.Credential, r.AccountID, r.BaseURL) |
| 342 | + | return r.wrapWithRefresh(inner) | |||
| 341 | } | 343 | } | ||
| 342 | return provider.NewOpenAI(r.Credential, r.BaseURL) | 344 | return provider.NewOpenAI(r.Credential, r.BaseURL) | ||
| 343 | default: | 345 | default: | ||
| 344 | if r.AuthMethod == "oauth" { | 346 | if r.AuthMethod == "oauth" { | ||
| 345 | - | return provider.NewAnthropicOAuth(r.Credential, r.BaseURL) | 347 | + | inner := provider.NewAnthropicOAuth(r.Credential, r.BaseURL) |
| 348 | + | return r.wrapWithRefresh(inner) | |||
| 346 | } | 349 | } | ||
| 347 | return provider.NewAnthropic(r.Credential, r.BaseURL) | 350 | return provider.NewAnthropic(r.Credential, r.BaseURL) | ||
| 348 | } | 351 | } | ||
| 349 | } | 352 | } | ||
| 350 | 353 | ||||
| 354 | + | // wrapWithRefresh wraps an OAuth client so the access token is | |||
| 355 | + | // refreshed automatically before each API call. Without this, long | |||
| 356 | + | // sessions (hours) silently fail when the 1-hour token expires. | |||
| 357 | + | func (r Resolved) wrapWithRefresh(inner provider.Client) provider.Client { | |||
| 358 | + | provName := r.Provider | |||
| 359 | + | baseURL := r.BaseURL | |||
| 360 | + | accountID := r.AccountID | |||
| 361 | + | ||||
| 362 | + | refreshFn := func(ctx context.Context) (string, error) { | |||
| 363 | + | tok, err := refreshIfExpired(provName, loadOAuthToken(provName)) | |||
| 364 | + | if err != nil { | |||
| 365 | + | return "", err | |||
| 366 | + | } | |||
| 367 | + | return tok.AccessToken, nil | |||
| 368 | + | } | |||
| 369 | + | ||||
| 370 | + | factory := func(token string) provider.Client { | |||
| 371 | + | switch provName { | |||
| 372 | + | case "openai": | |||
| 373 | + | return provider.NewOpenAICodex(token, accountID, baseURL) | |||
| 374 | + | default: | |||
| 375 | + | return provider.NewAnthropicOAuth(token, baseURL) | |||
| 376 | + | } | |||
| 377 | + | } | |||
| 378 | + | ||||
| 379 | + | return provider.NewRefreshingClient(inner, refreshFn, factory) | |||
| 380 | + | } | |||
| 381 | + | ||||
| 351 | // UseSandbox replaces the sandbox pointer that every tool in r's | 382 | // UseSandbox replaces the sandbox pointer that every tool in r's | ||
| 352 | // registry references. Used to keep the /jail state stable across | 383 | // registry references. Used to keep the /jail state stable across | ||
| 353 | // agent rebuilds (e.g. /login, /model switching providers). | 384 | // agent rebuilds (e.g. /login, /model switching providers). | ||
M internal/agent/config.go +20 -0
| @@ -155,6 +155,26 @@ func ResolveCredentialFull(provider, explicit string) (cred, method, accountID s | |||
| 155 | 155 | return "", "", "", fmt.Errorf("no credential for %s", provider) | |
| 156 | 156 | } | |
| 157 | 157 | ||
| 158 | + | // loadOAuthToken reads the current OAuth token from auth.json for the | |
| 159 | + | // given provider. Returns nil if no token is stored. | |
| 160 | + | func loadOAuthToken(providerName string) *auth.OAuthToken { | |
| 161 | + | c, err := AuthStoreFor().Load() | |
| 162 | + | if err != nil { | |
| 163 | + | return nil | |
| 164 | + | } | |
| 165 | + | switch providerName { | |
| 166 | + | case "anthropic": | |
| 167 | + | if c.Anthropic.OAuth != nil { | |
| 168 | + | return c.Anthropic.OAuth | |
| 169 | + | } | |
| 170 | + | case "openai": | |
| 171 | + | if c.OpenAI.OAuth != nil { | |
| 172 | + | return c.OpenAI.OAuth | |
| 173 | + | } | |
| 174 | + | } | |
| 175 | + | return nil | |
| 176 | + | } | |
| 177 | + | ||
| 158 | 178 | // refreshIfExpired returns a usable OAuth token for the given provider, | |
| 159 | 179 | // refreshing it synchronously when it's past (or near) expiry. The | |
| 160 | 180 | // refreshed token is persisted to auth.json. | |
| @@ -155,6 +155,26 @@ func ResolveCredentialFull(provider, explicit string) (cred, method, accountID s | |||||
| 155 | return "", "", "", fmt.Errorf("no credential for %s", provider) | 155 | return "", "", "", fmt.Errorf("no credential for %s", provider) | ||
| 156 | } | 156 | } | ||
| 157 | 157 | ||||
| 158 | + | // loadOAuthToken reads the current OAuth token from auth.json for the | |||
| 159 | + | // given provider. Returns nil if no token is stored. | |||
| 160 | + | func loadOAuthToken(providerName string) *auth.OAuthToken { | |||
| 161 | + | c, err := AuthStoreFor().Load() | |||
| 162 | + | if err != nil { | |||
| 163 | + | return nil | |||
| 164 | + | } | |||
| 165 | + | switch providerName { | |||
| 166 | + | case "anthropic": | |||
| 167 | + | if c.Anthropic.OAuth != nil { | |||
| 168 | + | return c.Anthropic.OAuth | |||
| 169 | + | } | |||
| 170 | + | case "openai": | |||
| 171 | + | if c.OpenAI.OAuth != nil { | |||
| 172 | + | return c.OpenAI.OAuth | |||
| 173 | + | } | |||
| 174 | + | } | |||
| 175 | + | return nil | |||
| 176 | + | } | |||
| 177 | + | ||||
| 158 | // refreshIfExpired returns a usable OAuth token for the given provider, | 178 | // refreshIfExpired returns a usable OAuth token for the given provider, | ||
| 159 | // refreshing it synchronously when it's past (or near) expiry. The | 179 | // refreshing it synchronously when it's past (or near) expiry. The | ||
| 160 | // refreshed token is persisted to auth.json. | 180 | // refreshed token is persisted to auth.json. | ||
A internal/provider/refreshing.go +57 -0
| @@ -0,0 +1,57 @@ | |||
| 1 | + | package provider | |
| 2 | + | ||
| 3 | + | import ( | |
| 4 | + | "context" | |
| 5 | + | "sync" | |
| 6 | + | ) | |
| 7 | + | ||
| 8 | + | // TokenRefresher is a callback that checks whether the current token | |
| 9 | + | // is still valid and returns a fresh one if needed. The returned | |
| 10 | + | // string is the new access token; if empty the old one is still fine. | |
| 11 | + | // An error means refresh failed (network down, refresh token expired, | |
| 12 | + | // etc.) — the caller should proceed with the stale token and let the | |
| 13 | + | // API return 401 naturally. | |
| 14 | + | type TokenRefresher func(ctx context.Context) (newToken string, err error) | |
| 15 | + | ||
| 16 | + | // RefreshingClient wraps a Client and calls a TokenRefresher before | |
| 17 | + | // every Stream call. When the refresher returns a new token, a fresh | |
| 18 | + | // underlying client is built via the factory function. | |
| 19 | + | type RefreshingClient struct { | |
| 20 | + | mu sync.Mutex | |
| 21 | + | inner Client | |
| 22 | + | refresh TokenRefresher | |
| 23 | + | factory func(token string) Client | |
| 24 | + | } | |
| 25 | + | ||
| 26 | + | // NewRefreshingClient wraps inner with automatic token refresh. | |
| 27 | + | // refreshFn is called before each Stream; if it returns a non-empty | |
| 28 | + | // token the factory rebuilds the underlying client with the new token. | |
| 29 | + | func NewRefreshingClient(inner Client, refreshFn TokenRefresher, factory func(token string) Client) *RefreshingClient { | |
| 30 | + | return &RefreshingClient{ | |
| 31 | + | inner: inner, | |
| 32 | + | refresh: refreshFn, | |
| 33 | + | factory: factory, | |
| 34 | + | } | |
| 35 | + | } | |
| 36 | + | ||
| 37 | + | func (c *RefreshingClient) Name() string { | |
| 38 | + | c.mu.Lock() | |
| 39 | + | defer c.mu.Unlock() | |
| 40 | + | return c.inner.Name() | |
| 41 | + | } | |
| 42 | + | ||
| 43 | + | func (c *RefreshingClient) Stream(ctx context.Context, req Request) (<-chan Event, error) { | |
| 44 | + | if c.refresh != nil { | |
| 45 | + | if newToken, err := c.refresh(ctx); err == nil && newToken != "" { | |
| 46 | + | c.mu.Lock() | |
| 47 | + | c.inner = c.factory(newToken) | |
| 48 | + | c.mu.Unlock() | |
| 49 | + | } | |
| 50 | + | // On refresh error: proceed with the current client. | |
| 51 | + | // The stale token will 401 and the user sees a clear error. | |
| 52 | + | } | |
| 53 | + | c.mu.Lock() | |
| 54 | + | inner := c.inner | |
| 55 | + | c.mu.Unlock() | |
| 56 | + | return inner.Stream(ctx, req) | |
| 57 | + | } | |
| @@ -0,0 +1,57 @@ | |||||
| 1 | + | package provider | |||
| 2 | + | ||||
| 3 | + | import ( | |||
| 4 | + | "context" | |||
| 5 | + | "sync" | |||
| 6 | + | ) | |||
| 7 | + | ||||
| 8 | + | // TokenRefresher is a callback that checks whether the current token | |||
| 9 | + | // is still valid and returns a fresh one if needed. The returned | |||
| 10 | + | // string is the new access token; if empty the old one is still fine. | |||
| 11 | + | // An error means refresh failed (network down, refresh token expired, | |||
| 12 | + | // etc.) — the caller should proceed with the stale token and let the | |||
| 13 | + | // API return 401 naturally. | |||
| 14 | + | type TokenRefresher func(ctx context.Context) (newToken string, err error) | |||
| 15 | + | ||||
| 16 | + | // RefreshingClient wraps a Client and calls a TokenRefresher before | |||
| 17 | + | // every Stream call. When the refresher returns a new token, a fresh | |||
| 18 | + | // underlying client is built via the factory function. | |||
| 19 | + | type RefreshingClient struct { | |||
| 20 | + | mu sync.Mutex | |||
| 21 | + | inner Client | |||
| 22 | + | refresh TokenRefresher | |||
| 23 | + | factory func(token string) Client | |||
| 24 | + | } | |||
| 25 | + | ||||
| 26 | + | // NewRefreshingClient wraps inner with automatic token refresh. | |||
| 27 | + | // refreshFn is called before each Stream; if it returns a non-empty | |||
| 28 | + | // token the factory rebuilds the underlying client with the new token. | |||
| 29 | + | func NewRefreshingClient(inner Client, refreshFn TokenRefresher, factory func(token string) Client) *RefreshingClient { | |||
| 30 | + | return &RefreshingClient{ | |||
| 31 | + | inner: inner, | |||
| 32 | + | refresh: refreshFn, | |||
| 33 | + | factory: factory, | |||
| 34 | + | } | |||
| 35 | + | } | |||
| 36 | + | ||||
| 37 | + | func (c *RefreshingClient) Name() string { | |||
| 38 | + | c.mu.Lock() | |||
| 39 | + | defer c.mu.Unlock() | |||
| 40 | + | return c.inner.Name() | |||
| 41 | + | } | |||
| 42 | + | ||||
| 43 | + | func (c *RefreshingClient) Stream(ctx context.Context, req Request) (<-chan Event, error) { | |||
| 44 | + | if c.refresh != nil { | |||
| 45 | + | if newToken, err := c.refresh(ctx); err == nil && newToken != "" { | |||
| 46 | + | c.mu.Lock() | |||
| 47 | + | c.inner = c.factory(newToken) | |||
| 48 | + | c.mu.Unlock() | |||
| 49 | + | } | |||
| 50 | + | // On refresh error: proceed with the current client. | |||
| 51 | + | // The stale token will 401 and the user sees a clear error. | |||
| 52 | + | } | |||
| 53 | + | c.mu.Lock() | |||
| 54 | + | inner := c.inner | |||
| 55 | + | c.mu.Unlock() | |||
| 56 | + | return inner.Stream(ctx, req) | |||
| 57 | + | } | |||