ab55cb6 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

fix(claude): add authentication refresh on 401 errors (#1581)

ab55cb6Kieran Klukas1 year ago

internal/agent/coordinator.go+67−24

compare with main

unifiedsplit

@@ -10,6 +10,7 @@ import (

10 "io"

11 "log/slog"

12 "maps"

13+ "net/http"

13 "os"

14 "os"

14 "slices"

15 "slices"

15 "strings"

16 "strings"

@@ -130,32 +131,42 @@ func (c *coordinator) Run(ctx context.Context, sessionID string, prompt string,

130

131

131 mergedOptions, temp, topP, topK, freqPenalty, presPenalty := mergeCallOptions(model, providerCfg)

132 mergedOptions, temp, topP, topK, freqPenalty, presPenalty := mergeCallOptions(model, providerCfg)

132

133

133− if providerCfg.OAuthToken != nil && providerCfg.OAuthToken.IsExpired() {

134+ run := func() (*fantasy.AgentResult, error) {

135+ return c.currentAgent.Run(ctx, SessionAgentCall{

136+ SessionID: sessionID,

137+ Prompt: prompt,

138+ Attachments: attachments,

139+ MaxOutputTokens: maxTokens,

140+ ProviderOptions: mergedOptions,

141+ Temperature: temp,

142+ TopP: topP,

143+ TopK: topK,

144+ FrequencyPenalty: freqPenalty,

145+ PresencePenalty: presPenalty,

146+ })

147+ }

148+ result, originalErr := run()

149+

150+ if c.isUnauthorized(originalErr) {

151+ switch {

152+ case providerCfg.OAuthToken != nil:

134− slog.Info("Detected expired OAuth token, attempting refresh", "provider", providerCfg.ID)

153+ slog.Info("Received 401. Refreshing token and retrying", "provider", providerCfg.ID)

135− if refreshErr := c.cfg.RefreshOAuthToken(ctx, providerCfg.ID); refreshErr != nil {

154+ if err := c.refreshOAuth2Token(ctx, providerCfg); err != nil {

155+ return nil, originalErr

156+ }

136− slog.Error("Failed to refresh OAuth token", "provider", providerCfg.ID, "error", refreshErr)

157+ slog.Info("Retrying request with refreshed OAuth token", "provider", providerCfg.ID)

137− return nil, refreshErr

158+ return run()

138− }

159+ case strings.Contains(providerCfg.APIKeyTemplate, "$"):

139−

160+ slog.Info("Received 401. Refreshing API Key template and retrying", "provider", providerCfg.ID)

140− // Rebuild models with refreshed token

161+ if err := c.refreshApiKeyTemplate(ctx, providerCfg); err != nil {

141− if updateErr := c.UpdateModels(ctx); updateErr != nil {

162+ return nil, originalErr

142− slog.Error("Failed to update models after token refresh", "error", updateErr)

163+ }

164+ slog.Info("Retrying request with refreshed API key", "provider", providerCfg.ID)

143− return nil, updateErr

165+ return run()

144 }

166 }

145 }

167 }

146− result, err := c.currentAgent.Run(ctx, SessionAgentCall{

168+

147− SessionID: sessionID,

148− Prompt: prompt,

149− Attachments: attachments,

150− MaxOutputTokens: maxTokens,

151− ProviderOptions: mergedOptions,

152− Temperature: temp,

153− TopP: topP,

154− TopK: topK,

155− FrequencyPenalty: freqPenalty,

156− PresencePenalty: presPenalty,

157− })

158− return result, err

169+ return result, originalErr

159}

170}

160

171

161func getProviderOptions(model Model, providerCfg config.ProviderConfig) fantasy.ProviderOptions {

172func getProviderOptions(model Model, providerCfg config.ProviderConfig) fantasy.ProviderOptions {

@@ -773,3 +784,35 @@ func (c *coordinator) Summarize(ctx context.Context, sessionID string) error {

773 }

784 }

774 return c.currentAgent.Summarize(ctx, sessionID, getProviderOptions(c.currentAgent.Model(), providerCfg))

785 return c.currentAgent.Summarize(ctx, sessionID, getProviderOptions(c.currentAgent.Model(), providerCfg))

775}

786}

787+

788+func (c *coordinator) isUnauthorized(err error) bool {

789+ var providerErr *fantasy.ProviderError

790+ return errors.As(err, &providerErr) && providerErr.StatusCode == http.StatusUnauthorized

791+}

792+

793+func (c *coordinator) refreshOAuth2Token(ctx context.Context, providerCfg config.ProviderConfig) error {

794+ if err := c.cfg.RefreshOAuthToken(ctx, providerCfg.ID); err != nil {

795+ slog.Error("Failed to refresh OAuth token after 401 error", "provider", providerCfg.ID, "error", err)

796+ return err

797+ }

798+ if err := c.UpdateModels(ctx); err != nil {

799+ return err

800+ }

801+ return nil

802+}

803+

804+func (c *coordinator) refreshApiKeyTemplate(ctx context.Context, providerCfg config.ProviderConfig) error {

805+ newAPIKey, err := c.cfg.Resolve(providerCfg.APIKeyTemplate)

806+ if err != nil {

807+ slog.Error("Failed to re-resolve API key after 401 error", "provider", providerCfg.ID, "error", err)

808+ return err

809+ }

810+

811+ providerCfg.APIKey = newAPIKey

812+ c.cfg.Providers.Set(providerCfg.ID, providerCfg)

813+

814+ if err := c.UpdateModels(ctx); err != nil {

815+ return err

816+ }

817+ return nil

818+}

internal/config/config.go+5−2

@@ -95,6 +95,8 @@ type ProviderConfig struct {

95 Type catwalk.Type `json:"type,omitempty" jsonschema:"description=Provider type that determines the API format,enum=openai,enum=openai-compat,enum=anthropic,enum=gemini,enum=azure,enum=vertexai,default=openai"`

96 // The provider's API key.

97 APIKey string `json:"api_key,omitempty" jsonschema:"description=API key for authentication with the provider,example=$OPENAI_API_KEY"`

98+ // The original API key template before resolution (for re-resolution on auth errors).

99+ APIKeyTemplate string `json:"-"`

98 // OAuthToken for providers that use OAuth2 authentication.

100 // OAuthToken for providers that use OAuth2 authentication.

99 OAuthToken *oauth.Token `json:"oauth,omitempty" jsonschema:"description=OAuth2 token for authentication with the provider"`

101 OAuthToken *oauth.Token `json:"oauth,omitempty" jsonschema:"description=OAuth2 token for authentication with the provider"`

100 // Marks the provider as disabled.

102 // Marks the provider as disabled.

@@ -469,6 +471,7 @@ func (c *Config) SetConfigField(key string, value any) error {

469 return nil

471 return nil

470}

472}

471

473

474+// RefreshOAuthToken refreshes the OAuth token for the given provider.

472func (c *Config) RefreshOAuthToken(ctx context.Context, providerID string) error {

475func (c *Config) RefreshOAuthToken(ctx context.Context, providerID string) error {

473 providerConfig, exists := c.Providers.Get(providerID)

476 providerConfig, exists := c.Providers.Get(providerID)

474 if !exists {

477 if !exists {

@@ -479,7 +482,7 @@ func (c *Config) RefreshOAuthToken(ctx context.Context, providerID string) error

479 return fmt.Errorf("provider %s does not have an OAuth token", providerID)

482 return fmt.Errorf("provider %s does not have an OAuth token", providerID)

480 }

483 }

481

484

482− // Only Anthropic provider uses OAuth for now

485+ // Only Anthropic provider uses OAuth for now.

483 if providerID != string(catwalk.InferenceProviderAnthropic) {

486 if providerID != string(catwalk.InferenceProviderAnthropic) {

484 return fmt.Errorf("OAuth refresh not supported for provider %s", providerID)

487 return fmt.Errorf("OAuth refresh not supported for provider %s", providerID)

485 }

488 }

@@ -489,7 +492,7 @@ func (c *Config) RefreshOAuthToken(ctx context.Context, providerID string) error

489 return fmt.Errorf("failed to refresh OAuth token for provider %s: %w", providerID, err)

492 return fmt.Errorf("failed to refresh OAuth token for provider %s: %w", providerID, err)

490 }

493 }

491

494

492− slog.Info("Successfully refreshed OAuth token in background", "provider", providerID)

495+ slog.Info("Successfully refreshed OAuth token", "provider", providerID)

493 providerConfig.OAuthToken = newToken

496 providerConfig.OAuthToken = newToken

494 providerConfig.APIKey = fmt.Sprintf("Bearer %s", newToken.AccessToken)

497 providerConfig.APIKey = fmt.Sprintf("Bearer %s", newToken.AccessToken)

495 providerConfig.SetupClaudeCode()

498 providerConfig.SetupClaudeCode()

internal/config/load.go+1−22

@@ -1,7 +1,6 @@

1package config

3import (

4− "cmp"

5 "context"

4 "context"

6 "encoding/json"

5 "encoding/json"

7 "fmt"

6 "fmt"

@@ -19,11 +18,9 @@ import (

19 "github.com/charmbracelet/catwalk/pkg/catwalk"

18 "github.com/charmbracelet/catwalk/pkg/catwalk"

20 "github.com/charmbracelet/crush/internal/csync"

19 "github.com/charmbracelet/crush/internal/csync"

21 "github.com/charmbracelet/crush/internal/env"

20 "github.com/charmbracelet/crush/internal/env"

22− "github.com/charmbracelet/crush/internal/event"

23 "github.com/charmbracelet/crush/internal/fsext"

21 "github.com/charmbracelet/crush/internal/fsext"

24 "github.com/charmbracelet/crush/internal/home"

22 "github.com/charmbracelet/crush/internal/home"

25 "github.com/charmbracelet/crush/internal/log"

23 "github.com/charmbracelet/crush/internal/log"

26− "github.com/charmbracelet/crush/internal/oauth/claude"

27 powernapConfig "github.com/charmbracelet/x/powernap/pkg/config"

24 powernapConfig "github.com/charmbracelet/x/powernap/pkg/config"

28)

25)

@@ -189,6 +186,7 @@ func (c *Config) configureProviders(env env.Env, resolver VariableResolver, know

189 Name: p.Name,

186 Name: p.Name,

190 BaseURL: p.APIEndpoint,

187 BaseURL: p.APIEndpoint,

191 APIKey: p.APIKey,

188 APIKey: p.APIKey,

189+ APIKeyTemplate: p.APIKey, // Store original template for re-resolution

192 OAuthToken: config.OAuthToken,

190 OAuthToken: config.OAuthToken,

193 Type: p.Type,

191 Type: p.Type,

194 Disable: config.Disable,

192 Disable: config.Disable,

@@ -200,25 +198,6 @@ func (c *Config) configureProviders(env env.Env, resolver VariableResolver, know

200 }

198 }

201

199

202 if p.ID == catwalk.InferenceProviderAnthropic && config.OAuthToken != nil {

200 if p.ID == catwalk.InferenceProviderAnthropic && config.OAuthToken != nil {

203− if config.OAuthToken.IsExpired() {

204− newToken, err := claude.RefreshToken(context.TODO(), config.OAuthToken.RefreshToken)

205− if err == nil {

206− slog.Info("Successfully refreshed Anthropic OAuth token")

207− config.OAuthToken = newToken

208− prepared.OAuthToken = newToken

209− if err := cmp.Or(

210− c.SetConfigField("providers.anthropic.api_key", newToken.AccessToken),

211− c.SetConfigField("providers.anthropic.oauth", newToken),

212− ); err != nil {

213− return err

214− }

215− } else {

216− slog.Error("Failed to refresh Anthropic OAuth token", "error", err)

217− event.Error(err)

218− }

219− } else {

220− slog.Info("Using existing non-expired Anthropic OAuth token")

221− }

222 prepared.SetupClaudeCode()

201 prepared.SetupClaudeCode()

223 }

202 }

224

203