feat(executor): add token refresh timeout and improve context handling during refresh

Introduced `tokenRefreshTimeout` constant for token refresh operations and enhanced context propagation for `refreshToken` by embedding roundtrip information if available. Adjusted `refreshAuth` to ensure default context initialization and handle cancellation errors appropriately.
This commit is contained in:
Luis Pater
2026-01-04 00:26:08 +08:00
parent 672e8549c0
commit 7a77b23f2d
2 changed files with 16 additions and 2 deletions

View File

@@ -45,6 +45,7 @@ const (
defaultAntigravityAgent = "antigravity/1.104.0 darwin/arm64" defaultAntigravityAgent = "antigravity/1.104.0 darwin/arm64"
antigravityAuthType = "antigravity" antigravityAuthType = "antigravity"
refreshSkew = 3000 * time.Second refreshSkew = 3000 * time.Second
tokenRefreshTimeout = 30 * time.Second
) )
var ( var (
@@ -914,7 +915,13 @@ func (e *AntigravityExecutor) ensureAccessToken(ctx context.Context, auth *clipr
if accessToken != "" && expiry.After(time.Now().Add(refreshSkew)) { if accessToken != "" && expiry.After(time.Now().Add(refreshSkew)) {
return accessToken, nil, nil return accessToken, nil, nil
} }
updated, errRefresh := e.refreshToken(ctx, auth.Clone()) refreshCtx := context.Background()
if ctx != nil {
if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
refreshCtx = context.WithValue(refreshCtx, "cliproxy.roundtripper", rt)
}
}
updated, errRefresh := e.refreshToken(refreshCtx, auth.Clone())
if errRefresh != nil { if errRefresh != nil {
return "", nil, errRefresh return "", nil, errRefresh
} }
@@ -944,7 +951,7 @@ func (e *AntigravityExecutor) refreshToken(ctx context.Context, auth *cliproxyau
httpReq.Header.Set("User-Agent", defaultAntigravityAgent) httpReq.Header.Set("User-Agent", defaultAntigravityAgent)
httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded") httpReq.Header.Set("Content-Type", "application/x-www-form-urlencoded")
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0) httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, tokenRefreshTimeout)
httpResp, errDo := httpClient.Do(httpReq) httpResp, errDo := httpClient.Do(httpReq)
if errDo != nil { if errDo != nil {
return auth, errDo return auth, errDo

View File

@@ -1536,6 +1536,9 @@ func (m *Manager) markRefreshPending(id string, now time.Time) bool {
} }
func (m *Manager) refreshAuth(ctx context.Context, id string) { func (m *Manager) refreshAuth(ctx context.Context, id string) {
if ctx == nil {
ctx = context.Background()
}
m.mu.RLock() m.mu.RLock()
auth := m.auths[id] auth := m.auths[id]
var exec ProviderExecutor var exec ProviderExecutor
@@ -1548,6 +1551,10 @@ func (m *Manager) refreshAuth(ctx context.Context, id string) {
} }
cloned := auth.Clone() cloned := auth.Clone()
updated, err := exec.Refresh(ctx, cloned) updated, err := exec.Refresh(ctx, cloned)
if err != nil && errors.Is(err, context.Canceled) {
log.Debugf("refresh canceled for %s, %s", auth.Provider, auth.ID)
return
}
log.Debugf("refreshed %s, %s, %v", auth.Provider, auth.ID, err) log.Debugf("refreshed %s, %s, %v", auth.Provider, auth.ID, err)
now := time.Now() now := time.Now()
if err != nil { if err != nil {