on main help sign in
fix(agent): centralize 401 retry logic and fix notify-on-success bug
9019775Kieran Klukas
internal/agent/coordinator.go+41−29
@@ -259,19 +259,21 @@ func (c *coordinator) Run(ctx context.Context, sessionID string, prompt string,
259 })
259 })
260 }
260 }
261 beforeLoaded := c.skillTracker.LoadedNames()
261 beforeLoaded := c.skillTracker.LoadedNames()
262− result, originalErr := run()
262+ var result *fantasy.AgentResult
263+ var originalErr error = c.runWithUnauthorizedRetry(ctx, providerCfg, func() error {
264+ var err error
265+ result, err = run()
266+ return err
267+ })
263 logTurnSkillUsage(sessionID, prompt, c.activeSkills, c.skillTracker, beforeLoaded)
268 logTurnSkillUsage(sessionID, prompt, c.activeSkills, c.skillTracker, beforeLoaded)
264
269
265− if c.isUnauthorized(originalErr) {
270+ // Notify only if still unauthorized after retry — a successful
266− if err := c.retryAfterUnauthorized(ctx, providerCfg); err == nil {
271+ // retry means the user doesn't need to re-authenticate.
267− result, originalErr = run()
268− }
269− if c.notify != nil && model.ModelCfg.Provider == hyper.Name {
272+ if originalErr != nil && c.isUnauthorized(originalErr) && c.notify != nil && model.ModelCfg.Provider == hyper.Name {
270− c.notify.Publish(pubsub.CreatedEvent, notify.Notification{
273+ c.notify.Publish(pubsub.CreatedEvent, notify.Notification{
271− Type: notify.TypeReAuthenticate,
274+ Type: notify.TypeReAuthenticate,
272− ProviderID: model.ModelCfg.Provider,
275+ ProviderID: model.ModelCfg.Provider,
273− })
274− }
276+ })
275 }
277 }
276
278
277 if hasLatest && c.runComplete != nil {
279 if hasLatest && c.runComplete != nil {
@@ -1062,14 +1064,7 @@ func (c *coordinator) Summarize(ctx context.Context, sessionID string) error {
1062 return c.currentAgent.Summarize(ctx, sessionID, getProviderOptions(c.currentAgent.Model(), providerCfg))
1064 return c.currentAgent.Summarize(ctx, sessionID, getProviderOptions(c.currentAgent.Model(), providerCfg))
1063 }
1065 }
1064
1066
1065− err := summarize()
1067+ return c.runWithUnauthorizedRetry(ctx, providerCfg, summarize)
1066− if err != nil && c.isUnauthorized(err) {
1067− if retryErr := c.retryAfterUnauthorized(ctx, providerCfg); retryErr == nil {
1068− return summarize()
1069− }
1070− }
1071−
1072− return err
1073}
1068}
1074
1069
1075// refreshTokenIfExpired proactively refreshes the OAuth token if it has expired.
1070// refreshTokenIfExpired proactively refreshes the OAuth token if it has expired.
@@ -1081,6 +1076,21 @@ func (c *coordinator) refreshTokenIfExpired(ctx context.Context, providerCfg con
1081 return c.refreshOAuth2Token(ctx, providerCfg)
1076 return c.refreshOAuth2Token(ctx, providerCfg)
1082}
1077}
1083
1078
1079+// runWithUnauthorizedRetry executes fn. If fn returns a 401 error, it
1080+// attempts to refresh credentials and re-runs fn once. Returns the
1081+// final error: from the retry if a retry was attempted, otherwise from
1082+// the original run. Callers that need to notify the user on persistent
1083+// failure should check isUnauthorized on the returned error.
1084+func (c *coordinator) runWithUnauthorizedRetry(ctx context.Context, providerCfg config.ProviderConfig, fn func() error) error {
1085+ err := fn()
1086+ if err != nil && c.isUnauthorized(err) {
1087+ if retryErr := c.retryAfterUnauthorized(ctx, providerCfg); retryErr == nil {
1088+ return fn()
1089+ }
1090+ }
1091+ return err
1092+}
1093+
1084// retryAfterUnauthorized attempts to refresh credentials after receiving a 401
1094// retryAfterUnauthorized attempts to refresh credentials after receiving a 401
1085// and returns nil if retry should be attempted.
1095// and returns nil if retry should be attempted.
1086func (c *coordinator) retryAfterUnauthorized(ctx context.Context, providerCfg config.ProviderConfig) error {
1096func (c *coordinator) retryAfterUnauthorized(ctx context.Context, providerCfg config.ProviderConfig) error {
@@ -1184,16 +1194,18 @@ func (c *coordinator) runSubAgent(ctx context.Context, params subAgentParams) (f
1184 NonInteractive: true,
1194 NonInteractive: true,
1185 })
1195 })
1186 }
1196 }
1187− result, err := run()
1197+ var result *fantasy.AgentResult
1188− if err != nil && c.isUnauthorized(err) {
1198+ err = c.runWithUnauthorizedRetry(ctx, providerCfg, func() error {
1189− if retryErr := c.retryAfterUnauthorized(ctx, providerCfg); retryErr == nil {
1199+ var runErr error
1190− result, err = run()
1200+ result, runErr = run()
1201+ return runErr
1202+ })
1203+ // Notify only if still unauthorized after retry.
1191− } else if c.notify != nil && model.ModelCfg.Provider == hyper.Name {
1204+ if err != nil && c.isUnauthorized(err) && c.notify != nil && model.ModelCfg.Provider == hyper.Name {
1192− c.notify.Publish(pubsub.CreatedEvent, notify.Notification{
1205+ c.notify.Publish(pubsub.CreatedEvent, notify.Notification{
1193− Type: notify.TypeReAuthenticate,
1206+ Type: notify.TypeReAuthenticate,
1194− ProviderID: model.ModelCfg.Provider,
1207+ ProviderID: model.ModelCfg.Provider,
1195− })
1196− }
1208+ })
1197 }
1209 }
1198 if err != nil {
1210 if err != nil {
1199 return fantasy.NewTextErrorResponse(fmt.Sprintf("Failed to generate response: %s", err)), nil
1211 return fantasy.NewTextErrorResponse(fmt.Sprintf("Failed to generate response: %s", err)), nil