@@ -10,6 +10,7 @@ import (
1010 "io"
1111 "log/slog"
1212 "maps"
13+ "net/http"
1314 "os"
1415 "slices"
1516 "strings"
@@ -130,32 +131,42 @@ func (c *coordinator) Run(ctx context.Context, sessionID string, prompt string,
130131
131132 mergedOptions, temp, topP, topK, freqPenalty, presPenalty := mergeCallOptions(model, providerCfg)
132133
133− if providerCfg.OAuthToken != nil && providerCfg.OAuthToken.IsExpired() {
134− slog.Info("Detected expired OAuth token, attempting refresh", "provider", providerCfg.ID)
135− if refreshErr := c.cfg.RefreshOAuthToken(ctx, providerCfg.ID); refreshErr != nil {
136− slog.Error("Failed to refresh OAuth token", "provider", providerCfg.ID, "error", refreshErr)
137− return nil, refreshErr
138− }
139−
140− // Rebuild models with refreshed token
141− if updateErr := c.UpdateModels(ctx); updateErr != nil {
142− slog.Error("Failed to update models after token refresh", "error", updateErr)
143− return nil, updateErr
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:
153+ slog.Info("Received 401. Refreshing token and retrying", "provider", providerCfg.ID)
154+ if err := c.refreshOAuth2Token(ctx, providerCfg); err != nil {
155+ return nil, originalErr
156+ }
157+ slog.Info("Retrying request with refreshed OAuth token", "provider", providerCfg.ID)
158+ return run()
159+ case strings.Contains(providerCfg.APIKeyTemplate, "$"):
160+ slog.Info("Received 401. Refreshing API Key template and retrying", "provider", providerCfg.ID)
161+ if err := c.refreshApiKeyTemplate(ctx, providerCfg); err != nil {
162+ return nil, originalErr
163+ }
164+ slog.Info("Retrying request with refreshed API key", "provider", providerCfg.ID)
165+ return run()
144166 }
145167 }
146− result, err := c.currentAgent.Run(ctx, SessionAgentCall{
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
168+
169+ return result, originalErr
159170}
160171
161172func getProviderOptions(model Model, providerCfg config.ProviderConfig) fantasy.ProviderOptions {
@@ -773,3 +784,35 @@ func (c *coordinator) Summarize(ctx context.Context, sessionID string) error {
773784 }
774785 return c.currentAgent.Summarize(ctx, sessionID, getProviderOptions(c.currentAgent.Model(), providerCfg))
775786}
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+}