diff --git a/service/ccm/credential.go b/service/ccm/credential.go index 6f41ba128..89ad5bb97 100644 --- a/service/ccm/credential.go +++ b/service/ccm/credential.go @@ -120,7 +120,7 @@ type Credential interface { setStatusSubscriber(*observable.Subscriber[struct{}]) start() error - pollUsage(ctx context.Context) + pollUsage() lastUpdatedTime() time.Time pollBackoff(base time.Duration) time.Duration usageTrackerOrNil() *AggregatedUsage diff --git a/service/ccm/credential_default.go b/service/ccm/credential_default.go index c004d1e84..1003c7b36 100644 --- a/service/ccm/credential_default.go +++ b/service/ccm/credential_default.go @@ -139,6 +139,7 @@ func (c *defaultCredential) start() error { c.logger.Warn("load usage statistics for ", c.tag, ": ", err) } } + go c.pollUsage() return nil } @@ -516,7 +517,7 @@ func (c *defaultCredential) earliestReset() time.Time { return earliest } -func (c *defaultCredential) pollUsage(ctx context.Context) { +func (c *defaultCredential) pollUsage() { if !c.pollAccess.TryLock() { return } @@ -537,6 +538,7 @@ func (c *defaultCredential) pollUsage(ctx context.Context) { return } + ctx := c.serviceContext httpClient := &http.Client{ Transport: c.forwardHTTPClient.Transport, Timeout: 5 * time.Second, @@ -633,11 +635,12 @@ func (c *defaultCredential) pollUsage(ctx context.Context) { c.emitStatusUpdate() if needsProfileFetch { - c.fetchProfile(ctx, httpClient, accessToken) + c.fetchProfile(httpClient, accessToken) } } -func (c *defaultCredential) fetchProfile(ctx context.Context, httpClient *http.Client, accessToken string) { +func (c *defaultCredential) fetchProfile(httpClient *http.Client, accessToken string) { + ctx := c.serviceContext response, err := doHTTPWithRetry(ctx, httpClient, func() (*http.Request, error) { request, err := http.NewRequestWithContext(ctx, http.MethodGet, claudeAPIBaseURL+"/api/oauth/profile", nil) if err != nil { diff --git a/service/ccm/credential_external.go b/service/ccm/credential_external.go index f51205203..40d2c3677 100644 --- a/service/ccm/credential_external.go +++ b/service/ccm/credential_external.go @@ -574,13 +574,14 @@ func (c *externalCredential) doPollUsageRequest(ctx context.Context) (*http.Resp return nil, E.New("no transport available") } -func (c *externalCredential) pollUsage(ctx context.Context) { +func (c *externalCredential) pollUsage() { if !c.pollAccess.TryLock() { return } defer c.pollAccess.Unlock() defer c.markUsagePollAttempted() + ctx := c.getReverseContext() response, err := c.doPollUsageRequest(ctx) if err != nil { c.logger.Debug("poll usage for ", c.tag, ": ", err) @@ -919,7 +920,7 @@ func (c *externalCredential) setReverseSession(session *yamux.Session) bool { go c.statusStreamLoop() } if triggerUsageRefresh { - go c.pollUsage(c.getReverseContext()) + go c.pollUsage() } if emitStatus { c.emitStatusUpdate() diff --git a/service/ccm/credential_provider.go b/service/ccm/credential_provider.go index 8d993c6ca..9fac91b2c 100644 --- a/service/ccm/credential_provider.go +++ b/service/ccm/credential_provider.go @@ -16,7 +16,7 @@ type credentialProvider interface { selectCredential(sessionID string, selection credentialSelection) (Credential, bool, error) onRateLimited(sessionID string, credential Credential, resetAt time.Time, selection credentialSelection) Credential linkProviderInterrupt(credential Credential, selection credentialSelection, onInterrupt func()) func() bool - pollIfStale(ctx context.Context) + pollIfStale() allCredentials() []Credential close() } @@ -58,7 +58,7 @@ func (p *singleCredentialProvider) onRateLimited(_ string, credential Credential return nil } -func (p *singleCredentialProvider) pollIfStale(ctx context.Context) { +func (p *singleCredentialProvider) pollIfStale() { now := time.Now() p.sessionAccess.Lock() for id, createdAt := range p.sessions { @@ -69,7 +69,7 @@ func (p *singleCredentialProvider) pollIfStale(ctx context.Context) { p.sessionAccess.Unlock() if time.Since(p.credential.lastUpdatedTime()) > p.credential.pollBackoff(defaultPollInterval) { - p.credential.pollUsage(ctx) + p.credential.pollUsage() } } @@ -357,7 +357,7 @@ func (p *balancerProvider) pickRandom(filter func(Credential) bool) Credential { return usable[rand.IntN(len(usable))] } -func (p *balancerProvider) pollIfStale(ctx context.Context) { +func (p *balancerProvider) pollIfStale() { now := time.Now() p.sessionAccess.Lock() for id, entry := range p.sessions { @@ -377,7 +377,7 @@ func (p *balancerProvider) pollIfStale(ctx context.Context) { for _, credential := range p.credentials { if time.Since(credential.lastUpdatedTime()) > credential.pollBackoff(p.pollInterval) { - credential.pollUsage(ctx) + credential.pollUsage() } } } diff --git a/service/ccm/service_handler.go b/service/ccm/service_handler.go index 1ccbd83ff..6bface738 100644 --- a/service/ccm/service_handler.go +++ b/service/ccm/service_handler.go @@ -182,7 +182,7 @@ func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } - provider.pollIfStale(s.ctx) + provider.pollIfStale() anthropicBetaHeader := r.Header.Get("anthropic-beta") if isFastModeRequest(anthropicBetaHeader) { @@ -305,7 +305,7 @@ func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) { if response.StatusCode != http.StatusOK && response.StatusCode != http.StatusTooManyRequests { body, _ := io.ReadAll(response.Body) s.logger.ErrorContext(ctx, "upstream error from ", selectedCredential.tagName(), ": status ", response.StatusCode, " ", string(body)) - go selectedCredential.pollUsage(s.ctx) + go selectedCredential.pollUsage() writeJSONError(w, r, http.StatusInternalServerError, "api_error", "proxy request (status "+strconv.Itoa(response.StatusCode)+"): "+string(body)) return diff --git a/service/ccm/service_status.go b/service/ccm/service_status.go index 50ba7ffe4..424afa127 100644 --- a/service/ccm/service_status.go +++ b/service/ccm/service_status.go @@ -103,7 +103,7 @@ func (s *Service) handleStatusEndpoint(w http.ResponseWriter, r *http.Request) { return } - provider.pollIfStale(r.Context()) + provider.pollIfStale() status := s.computeAggregatedUtilization(provider, userConfig) w.Header().Set("Content-Type", "application/json") @@ -125,7 +125,7 @@ func (s *Service) handleStatusStream(w http.ResponseWriter, r *http.Request, pro } defer s.statusObserver.UnSubscribe(subscription) - provider.pollIfStale(r.Context()) + provider.pollIfStale() w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK) diff --git a/service/ocm/credential.go b/service/ocm/credential.go index 1478f5f19..b0226f2d0 100644 --- a/service/ocm/credential.go +++ b/service/ocm/credential.go @@ -122,7 +122,7 @@ type Credential interface { setOnBecameUnusable(fn func()) setStatusSubscriber(*observable.Subscriber[struct{}]) start() error - pollUsage(ctx context.Context) + pollUsage() lastUpdatedTime() time.Time pollBackoff(base time.Duration) time.Duration usageTrackerOrNil() *AggregatedUsage diff --git a/service/ocm/credential_default.go b/service/ocm/credential_default.go index dd610e88d..3622ac8ff 100644 --- a/service/ocm/credential_default.go +++ b/service/ocm/credential_default.go @@ -143,6 +143,7 @@ func (c *defaultCredential) start() error { c.logger.Warn("load usage statistics for ", c.tag, ": ", err) } } + go c.pollUsage() return nil } @@ -597,7 +598,7 @@ func (c *defaultCredential) ocmGetBaseURL() string { return c.getBaseURL() } -func (c *defaultCredential) pollUsage(ctx context.Context) { +func (c *defaultCredential) pollUsage() { if !c.pollAccess.TryLock() { return } @@ -621,6 +622,7 @@ func (c *defaultCredential) pollUsage(ctx context.Context) { return } + ctx := c.serviceContext usageURL := strings.TrimSuffix(chatGPTBackendURL, "/codex") + "/wham/usage" accountID := c.getAccountID() diff --git a/service/ocm/credential_external.go b/service/ocm/credential_external.go index 3bcae1391..0ee1595d9 100644 --- a/service/ocm/credential_external.go +++ b/service/ocm/credential_external.go @@ -615,13 +615,14 @@ func (c *externalCredential) doPollUsageRequest(ctx context.Context) (*http.Resp return nil, E.New("no transport available") } -func (c *externalCredential) pollUsage(ctx context.Context) { +func (c *externalCredential) pollUsage() { if !c.pollAccess.TryLock() { return } defer c.pollAccess.Unlock() defer c.markUsagePollAttempted() + ctx := c.getReverseContext() response, err := c.doPollUsageRequest(ctx) if err != nil { c.logger.Debug("poll usage for ", c.tag, ": ", err) @@ -988,7 +989,7 @@ func (c *externalCredential) setReverseSession(session *yamux.Session) bool { go c.statusStreamLoop() } if triggerUsageRefresh { - go c.pollUsage(c.getReverseContext()) + go c.pollUsage() } if emitStatus { c.emitStatusUpdate() diff --git a/service/ocm/credential_provider.go b/service/ocm/credential_provider.go index 421258cd6..5d67eb032 100644 --- a/service/ocm/credential_provider.go +++ b/service/ocm/credential_provider.go @@ -16,7 +16,7 @@ type credentialProvider interface { selectCredential(sessionID string, selection credentialSelection) (Credential, bool, error) onRateLimited(sessionID string, credential Credential, resetAt time.Time, selection credentialSelection) Credential linkProviderInterrupt(credential Credential, selection credentialSelection, onInterrupt func()) func() bool - pollIfStale(ctx context.Context) + pollIfStale() allCredentials() []Credential close() } @@ -58,7 +58,7 @@ func (p *singleCredentialProvider) onRateLimited(_ string, credential Credential return nil } -func (p *singleCredentialProvider) pollIfStale(ctx context.Context) { +func (p *singleCredentialProvider) pollIfStale() { now := time.Now() p.sessionAccess.Lock() for id, createdAt := range p.sessions { @@ -69,7 +69,7 @@ func (p *singleCredentialProvider) pollIfStale(ctx context.Context) { p.sessionAccess.Unlock() if time.Since(p.credential.lastUpdatedTime()) > p.credential.pollBackoff(defaultPollInterval) { - p.credential.pollUsage(ctx) + p.credential.pollUsage() } } @@ -384,7 +384,7 @@ func (p *balancerProvider) pickRandom(filter func(Credential) bool) Credential { return usable[rand.IntN(len(usable))] } -func (p *balancerProvider) pollIfStale(ctx context.Context) { +func (p *balancerProvider) pollIfStale() { now := time.Now() p.sessionAccess.Lock() for id, entry := range p.sessions { @@ -404,7 +404,7 @@ func (p *balancerProvider) pollIfStale(ctx context.Context) { for _, credential := range p.credentials { if time.Since(credential.lastUpdatedTime()) > credential.pollBackoff(p.pollInterval) { - credential.pollUsage(ctx) + credential.pollUsage() } } } diff --git a/service/ocm/service_handler.go b/service/ocm/service_handler.go index 52e35f39b..8b50f748a 100644 --- a/service/ocm/service_handler.go +++ b/service/ocm/service_handler.go @@ -131,7 +131,7 @@ func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } - provider.pollIfStale(s.ctx) + provider.pollIfStale() selection := credentialSelectionForUser(userConfig) @@ -285,7 +285,7 @@ func (s *Service) ServeHTTP(w http.ResponseWriter, r *http.Request) { if response.StatusCode != http.StatusOK && response.StatusCode != http.StatusTooManyRequests { body, _ := io.ReadAll(response.Body) s.logger.ErrorContext(ctx, "upstream error from ", selectedCredential.tagName(), ": status ", response.StatusCode, " ", string(body)) - go selectedCredential.pollUsage(s.ctx) + go selectedCredential.pollUsage() writeJSONError(w, r, http.StatusInternalServerError, "api_error", "proxy request (status "+strconv.Itoa(response.StatusCode)+"): "+string(body)) return diff --git a/service/ocm/service_status.go b/service/ocm/service_status.go index ebbc9ceaa..3e7e2ff65 100644 --- a/service/ocm/service_status.go +++ b/service/ocm/service_status.go @@ -103,7 +103,7 @@ func (s *Service) handleStatusEndpoint(w http.ResponseWriter, r *http.Request) { return } - provider.pollIfStale(r.Context()) + provider.pollIfStale() status := s.computeAggregatedUtilization(provider, userConfig) w.Header().Set("Content-Type", "application/json") @@ -125,7 +125,7 @@ func (s *Service) handleStatusStream(w http.ResponseWriter, r *http.Request, pro } defer s.statusObserver.UnSubscribe(subscription) - provider.pollIfStale(r.Context()) + provider.pollIfStale() w.Header().Set("Content-Type", "application/json") w.WriteHeader(http.StatusOK)