git.iko.soy/crush
on main help sign in
refactor(coordinator): extract token refresh helpers to reduce duplication
8cd4786Andrey Nering
internal/agent/coordinator.go+32−39
compare with main
unifiedsplit
@@ -184,13 +184,10 @@ func (c *coordinator) Run(ctx context.Context, sessionID string, prompt string,
184
184
185 mergedOptions, temp, topP, topK, freqPenalty, presPenalty := mergeCallOptions(model, providerCfg)
185 mergedOptions, temp, topP, topK, freqPenalty, presPenalty := mergeCallOptions(model, providerCfg)
186
186
187− if providerCfg.OAuthToken != nil && providerCfg.OAuthToken.IsExpired() {
188− slog.Debug("Token needs to be refreshed", "provider", providerCfg.ID)
189− if err := c.refreshOAuth2Token(ctx, providerCfg); err != nil {
187+ if err := c.refreshTokenIfExpired(ctx, providerCfg); err != nil {
190− // NOTE(@andreynering): We don't return here because the event handling to ask the user to reauthenticate
188+ // NOTE(@andreynering): We don't return here because the event handling to ask the user to reauthenticate
191− // depends on the flow below. If refresh fails, proceed with the token we have.
189+ // depends on the flow below. If refresh fails, proceed with the token we have.
192− slog.Error("Failed to refresh OAuth2 token. Proceeding with existing token.", "error", err)
190+ slog.Error("Failed to refresh OAuth2 token. Proceeding with existing token.", "error", err)
193− }
194 }
191 }
195
192
196 run := func() (*fantasy.AgentResult, error) {
193 run := func() (*fantasy.AgentResult, error) {
@@ -212,20 +209,7 @@ func (c *coordinator) Run(ctx context.Context, sessionID string, prompt string,
212 logTurnSkillUsage(sessionID, prompt, c.activeSkills, c.skillTracker, beforeLoaded)
209 logTurnSkillUsage(sessionID, prompt, c.activeSkills, c.skillTracker, beforeLoaded)
213
210
214 if c.isUnauthorized(originalErr) {
211 if c.isUnauthorized(originalErr) {
215− switch {
216− case providerCfg.OAuthToken != nil:
217− slog.Debug("Received 401. Refreshing token and retrying", "provider", providerCfg.ID)
218− if err := c.refreshOAuth2Token(ctx, providerCfg); err != nil {
212+ if err := c.retryAfterUnauthorized(ctx, providerCfg); err == nil {
219− return nil, originalErr
220− }
221− slog.Debug("Retrying request with refreshed OAuth token", "provider", providerCfg.ID)
222− return run()
223− case strings.Contains(providerCfg.APIKeyTemplate, "$"):
224− slog.Debug("Received 401. Refreshing API Key template and retrying", "provider", providerCfg.ID)
225− if err := c.refreshApiKeyTemplate(ctx, providerCfg); err != nil {
226− return nil, originalErr
227− }
228− slog.Debug("Retrying request with refreshed API key", "provider", providerCfg.ID)
229 return run()
213 return run()
230 }
214 }
231 }
215 }
@@ -962,12 +946,8 @@ func (c *coordinator) Summarize(ctx context.Context, sessionID string) error {
962 return errModelProviderNotConfigured
946 return errModelProviderNotConfigured
963 }
947 }
964
948
965− // Proactively refresh OAuth token if expired, same as Run().
966− if providerCfg.OAuthToken != nil && providerCfg.OAuthToken.IsExpired() {
967− slog.Debug("Token needs to be refreshed before summarize", "provider", providerCfg.ID)
968− if err := c.refreshOAuth2Token(ctx, providerCfg); err != nil {
949+ if err := c.refreshTokenIfExpired(ctx, providerCfg); err != nil {
969− slog.Error("Failed to refresh OAuth2 token before summarize. Proceeding with existing token.", "error", err)
950+ slog.Error("Failed to refresh OAuth2 token before summarize. Proceeding with existing token.", "error", err)
970− }
971 }
951 }
972
952
973 summarize := func() error {
953 summarize := func() error {
@@ -976,18 +956,7 @@ func (c *coordinator) Summarize(ctx context.Context, sessionID string) error {
976
956
977 err := summarize()
957 err := summarize()
978 if err != nil && c.isUnauthorized(err) {
958 if err != nil && c.isUnauthorized(err) {
979− switch {
959+ if retryErr := c.retryAfterUnauthorized(ctx, providerCfg); retryErr == nil {
980− case providerCfg.OAuthToken != nil:
981− slog.Debug("Received 401 during summarize. Refreshing token and retrying", "provider", providerCfg.ID)
982− if refreshErr := c.refreshOAuth2Token(ctx, providerCfg); refreshErr != nil {
983− return err
984− }
985− return summarize()
986− case strings.Contains(providerCfg.APIKeyTemplate, "$"):
987− slog.Debug("Received 401 during summarize. Refreshing API Key template and retrying", "provider", providerCfg.ID)
988− if refreshErr := c.refreshApiKeyTemplate(ctx, providerCfg); refreshErr != nil {
989− return err
990− }
991 return summarize()
960 return summarize()
992 }
961 }
993 }
962 }
@@ -995,6 +964,30 @@ func (c *coordinator) Summarize(ctx context.Context, sessionID string) error {
995 return err
964 return err
996}
965}
997
966
967+// refreshTokenIfExpired proactively refreshes the OAuth token if it has expired.
968+func (c *coordinator) refreshTokenIfExpired(ctx context.Context, providerCfg config.ProviderConfig) error {
969+ if providerCfg.OAuthToken == nil || !providerCfg.OAuthToken.IsExpired() {
970+ return nil
971+ }
972+ slog.Debug("Token needs to be refreshed", "provider", providerCfg.ID)
973+ return c.refreshOAuth2Token(ctx, providerCfg)
974+}
975+
976+// retryAfterUnauthorized attempts to refresh credentials after receiving a 401
977+// and returns nil if retry should be attempted.
978+func (c *coordinator) retryAfterUnauthorized(ctx context.Context, providerCfg config.ProviderConfig) error {
979+ switch {
980+ case providerCfg.OAuthToken != nil:
981+ slog.Debug("Received 401. Refreshing token and retrying", "provider", providerCfg.ID)
982+ return c.refreshOAuth2Token(ctx, providerCfg)
983+ case strings.Contains(providerCfg.APIKeyTemplate, "$"):
984+ slog.Debug("Received 401. Refreshing API Key template and retrying", "provider", providerCfg.ID)
985+ return c.refreshApiKeyTemplate(ctx, providerCfg)
986+ default:
987+ return nil
988+ }
989+}
990+
998func (c *coordinator) isUnauthorized(err error) bool {
991func (c *coordinator) isUnauthorized(err error) bool {
999 var providerErr *fantasy.ProviderError
992 var providerErr *fantasy.ProviderError
1000 return errors.As(err, &providerErr) && providerErr.StatusCode == http.StatusUnauthorized
993 return errors.As(err, &providerErr) && providerErr.StatusCode == http.StatusUnauthorized