From 78a62f494c7028965140c9ff1b01ebe3a304559a Mon Sep 17 00:00:00 2001 From: steventohme Date: Wed, 29 Jul 2026 14:55:37 -0700 Subject: [PATCH 1/2] feat(proxy): warn Claude Code when a subscription turn fails over to the billable Weave key Munir's Claude subscription hit its 5h unified rate limit mid-session (anthropic-ratelimit-unified-5h-status: rejected) while a long, cache-heavy session was running under a native-web-search Anthropic-only constraint. The router correctly failed the turn over to the billable Weave/BYOK Anthropic key -- both credential-suppression paths already do this -- but did so silently, so the customer had no signal that ~$800 of subsequent turns would be billed instead of served on their own plan. Add subscriptionFailoverWarningMarker, a new routing-marker constant distinct from subscriptionOnlyWarningMarker: that one signals a permanent, org-level depleted-credits state with a top-up CTA; this one is a per-turn, self-healing condition ("this turn", not "until you add credits"), so it carries no CTA and no opt-out gate -- same billing-state-change rule as the existing subscription-only warning. Wire it into both places a subscription-served Anthropic turn can fail over: - Pre-emptive (service.go ~2457): a prior observer snapshot already reads exhausted before dispatch. Thread subscriptionFailingOver through to the marker assembly so the turn's single marker choice reflects it. - Reactive (service.go ~2886): a live 429/OAuth-rejection triggers the subscription-credit retry onto the Weave key. Since this branch only runs pre-commit (preludeBuf.Committed() == false), swapping the retry attempt's marker is safe -- nothing under the original marker has hit the wire yet. Codex/OpenAI is out of scope: there is no codexSubscriptionExhausted equivalent gating that path today, so a Codex subscription failover stays silent until that symmetric guard exists. Tests: two new full-flow ProxyMessages integration tests in failover_integration_test.go covering both paths end-to-end (including, for the reactive path, that the router's own same-binding retry runs first and the failover marker only appears on the eventual Weave-key attempt). --- internal/proxy/failover_integration_test.go | 115 ++++++++++++++++++++ internal/proxy/service.go | 17 ++- internal/proxy/usage_bypass.go | 9 ++ 3 files changed, 138 insertions(+), 3 deletions(-) diff --git a/internal/proxy/failover_integration_test.go b/internal/proxy/failover_integration_test.go index d506c9cb3..d2aea04b4 100644 --- a/internal/proxy/failover_integration_test.go +++ b/internal/proxy/failover_integration_test.go @@ -9,12 +9,14 @@ import ( "strings" "sync" "testing" + "time" "workweave/router/internal/providers" "workweave/router/internal/providers/anthropic" "workweave/router/internal/providers/openai" "workweave/router/internal/providers/openaicompat" "workweave/router/internal/proxy" + "workweave/router/internal/proxy/usage" "workweave/router/internal/router" "github.com/stretchr/testify/assert" @@ -525,3 +527,116 @@ func TestProxyMessages_GeminiNon400NotRetried(t *testing.T) { assert.Len(t, client.bodies, 1, "a 503 is not a VALIDATED-schema 400 — no AUTO retry") } + +// TestProxyMessages_PreemptiveSubscriptionFailoverWarnsUser guards the +// pre-emptive path (an observer snapshot already reads exhausted before +// dispatch): the turn must serve on the deployment key AND the client must see +// the billable-failover warning marker instead of the normal routing marker. +func TestProxyMessages_PreemptiveSubscriptionFailoverWarnsUser(t *testing.T) { + fr := &fakeRouter{decision: router.Decision{Provider: providers.ProviderAnthropic, Model: bypassScorerPickMdl}} + p := &fakeProvider{proxyResponse: func(w http.ResponseWriter) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"role\":\"assistant\",\"content\":[],\"model\":\""+bypassScorerPickMdl+"\",\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\n") + _, _ = io.WriteString(w, "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n") + _, _ = io.WriteString(w, "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}\n\n") + _, _ = io.WriteString(w, "event: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}\n\n") + _, _ = io.WriteString(w, "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":1}}\n\n") + _, _ = io.WriteString(w, "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n") + }} + // Observer seeded EXHAUSTED on the weekly window, mirroring + // TestSubscriptionExhausted_ServesOnDeploymentKey. + obs := usage.NewObserver([]byte("salt"), 10*time.Minute, time.Now) + obs.Record(obs.Key([]byte(bypassSubToken)), usage.Snapshot{ + Secondary: usage.Window{UsedPercent: 1.0, WindowMinutes: 10080}, + }) + svc := proxy.NewService(fr, map[string]providers.Client{providers.ProviderAnthropic: p}, nil, false, nil, nil, false, providers.ProviderAnthropic, bypassScorerPickMdl, nil). + WithSubscriptionAwareRouting(obs, 0.05, 2.0). + WithDeploymentKeyedProviders(map[string]struct{}{providers.ProviderAnthropic: {}}) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader("")) + body := []byte(`{"model":"` + bypassScorerPickMdl + `","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + + // bypassCtx sets the usage-bypass gate; the routed path is what we're + // exercising here, so use a plain subscription context instead. + ctx := context.WithValue(context.Background(), proxy.AnthropicSubscriptionContextKey{}, bypassSubToken) + require.NoError(t, svc.ProxyMessages(ctx, body, rec, req)) + + require.Len(t, p.proxyCreds, 1, "the turn must be dispatched once, on the deployment key") + creds := p.proxyCreds[0] + if creds != nil { + assert.False(t, creds.OAuth, "the exhausted subscription must not be forwarded") + } + respBody := rec.Body.String() + assert.Contains(t, respBody, "your Claude subscription hit its usage limit", + "the client must see the billable-failover warning, not the normal routing marker") + assert.NotContains(t, respBody, "· "+"best pick for this turn", + "the failover warning replaces the routing marker, it doesn't append to it") +} + +// TestProxyMessages_ReactiveSubscriptionFailoverWarnsUser guards the reactive +// path: a subscription-served Anthropic turn hits a live retryable error +// (429), so the router retries the SAME model on the Weave key. The retry +// must carry the billable-failover warning marker, not the original routing +// marker, and the retry must authenticate via x-api-key (not the spent OAuth +// bearer). +func TestProxyMessages_ReactiveSubscriptionFailoverWarnsUser(t *testing.T) { + var ( + mu sync.Mutex + calls int + ) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + mu.Lock() + calls++ + attempt := calls + mu.Unlock() + if attempt <= 3 { + // Attempts 1-3: subscription OAuth bearer. dispatchWithFallback + // retries a sole binding in place up to maxSameBindingRetries (2) + // before giving up, so the 429 must persist through all three + // same-binding attempts to reach the subscription-failover retry. + assert.Equal(t, "Bearer "+bypassSubToken, r.Header.Get("Authorization"), + "same-binding retries must still authenticate with the subscription bearer") + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = io.WriteString(w, `{"type":"error","error":{"type":"rate_limit_error","message":"5h limit reached"}}`) + return + } + // Attempt 4: the subscription-failover retry, on the deployment/BYOK + // key rather than the spent subscription bearer. + assert.Empty(t, r.Header.Get("Authorization"), "the retry must not reuse the spent subscription bearer") + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_recovered\",\"role\":\"assistant\",\"content\":[],\"model\":\""+bypassScorerPickMdl+"\",\"usage\":{\"input_tokens\":5,\"output_tokens\":0}}}\n\n") + _, _ = io.WriteString(w, "event: content_block_start\ndata: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n") + _, _ = io.WriteString(w, "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"recovered\"}}\n\n") + _, _ = io.WriteString(w, "event: content_block_stop\ndata: {\"type\":\"content_block_stop\",\"index\":0}\n\n") + _, _ = io.WriteString(w, "event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":1}}\n\n") + _, _ = io.WriteString(w, "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n") + })) + defer upstream.Close() + + svc := proxy.NewService( + &fakeRouter{decision: router.Decision{Provider: providers.ProviderAnthropic, Model: bypassScorerPickMdl}}, + map[string]providers.Client{providers.ProviderAnthropic: anthropic.NewClient("test-key", upstream.URL)}, + nil, false, nil, nil, false, providers.ProviderAnthropic, bypassScorerPickMdl, nil, + ).WithDeploymentKeyedProviders(map[string]struct{}{providers.ProviderAnthropic: {}}) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader("")) + req.Header.Set("Authorization", "Bearer "+bypassSubToken) + body := []byte(`{"model":"` + bypassScorerPickMdl + `","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + + require.NoError(t, svc.ProxyMessages(context.Background(), body, rec, req)) + + mu.Lock() + assert.Equal(t, 4, calls, "3 same-binding attempts on the subscription, then 1 subscription-failover retry on the Weave key") + mu.Unlock() + assert.Equal(t, http.StatusOK, rec.Code) + respBody := rec.Body.String() + assert.Contains(t, respBody, "recovered", "the client receives the retried attempt's content") + assert.Contains(t, respBody, "your Claude subscription hit its usage limit", + "the retry must carry the billable-failover warning") + assert.NotContains(t, respBody, "rate_limit_error", "the failed subscription attempt must never commit bytes") +} diff --git a/internal/proxy/service.go b/internal/proxy/service.go index 9ad50fc56..6d9d8b906 100644 --- a/internal/proxy/service.go +++ b/internal/proxy/service.go @@ -2453,8 +2453,10 @@ func (s *Service) ProxyMessages(ctx context.Context, body []byte, w http.Respons // another turn on it (429 until reset). Suppress the spent token so // resolution falls through to the deployment/BYOK key — the turn serves on // the Weave key (full cost) instead of hard-failing. Only fires once the - // observer has recorded exhaustion and a fallback key exists. - if s.claudeSubscriptionExhausted(ctx, r.Header) { + // observer has recorded exhaustion and a fallback key exists. Recorded so + // the marker built below can surface the billable-failover warning. + subscriptionFailingOver := s.claudeSubscriptionExhausted(ctx, r.Header) + if subscriptionFailingOver { ctx = withSuppressedClaudeSubscription(ctx) } ctx = resolveAndInjectCredentials(ctx, decision.Provider, r.Header) @@ -2535,6 +2537,12 @@ func (s *Service) ProxyMessages(ctx context.Context, body []byte, w http.Respons // turn reaching here is served free and should carry the top-up CTA. if billing.SubscriptionOnlyFromContext(ctx) { marker = subscriptionOnlyWarningMarker + } else if subscriptionFailingOver { + // Subscription hit its 5h/7d limit pre-dispatch: the turn is about to + // serve on the billable Weave/BYOK key instead. Not gated by the + // routing-marker opt-out — same billing-state-change rule as the + // subscription-only warning above. + marker = subscriptionFailoverWarningMarker } // toolValidator compiles the request's tool schemas once (LRU-cached); // translators validate/repair model tool calls against it. Nil if no tools. @@ -2875,7 +2883,10 @@ func (s *Service) ProxyMessages(ctx context.Context, body []byte, w http.Respons "model", decision.Model, "err", proxyErr, "upstream_status", upstreamStatus(proxyErr)) - subAttempt := s.anthropicNativeAttempt(env, r, subPrep, sink, preludeBuf, marker, setExtractor) + // Swap in the billable-failover warning for this retry attempt — safe + // because reaching here already required preludeBuf.Committed() == + // false, so nothing under the original marker has hit the wire yet. + subAttempt := s.anthropicNativeAttempt(env, r, subPrep, sink, preludeBuf, subscriptionFailoverWarningMarker, setExtractor) crossFormat = false respSummary = translate.ResponseSummary{} reqStats = providers.RequestMutationStats{} diff --git a/internal/proxy/usage_bypass.go b/internal/proxy/usage_bypass.go index 55ae28fe8..dfe803e5b 100644 --- a/internal/proxy/usage_bypass.go +++ b/internal/proxy/usage_bypass.go @@ -213,6 +213,15 @@ const subscriptionOnlyWarningMarkerCodex = routingMarkerPrefix + "your Weave router credits are depleted, so this turn is running on your own ChatGPT (Codex) subscription and paid model fallback is disabled. Add credits to restore full routing: " + topUpURL + "\n\n" +// subscriptionFailoverWarningMarker is prepended when a subscription-served +// Anthropic turn fails over to the billable Weave/BYOK key mid-session (5h/7d +// unified rate-limit rejection, or an OAuth credential rejection). Distinct +// from subscriptionOnlyWarningMarker: that one signals a permanent, org-level +// state (credits depleted, until top-up); this one is a per-turn, self-healing +// condition, so the copy says "this turn" rather than pointing at a CTA. +const subscriptionFailoverWarningMarker = routingMarkerPrefix + + "your Claude subscription hit its usage limit, so this turn is running on the Weave router key and will be billed. Full routing resumes automatically once your subscription window resets.\n\n" + // ErrCreditsExhaustedSubscriptionUnavailable is returned by ProxyMessages and // ProxyOpenAIChatCompletion when the org is in subscription-only mode but the // turn cannot be served on the caller's own From b03b8681330f59c5e2a13f7725f2c3ed7c86b09a Mon Sep 17 00:00:00 2001 From: steventohme Date: Wed, 29 Jul 2026 15:08:21 -0700 Subject: [PATCH 2/2] feat(proxy): surface credential.source + a combined subscription-failover bit on router.call Downstream consumers of the OTLP router.call log record (WorkWeave's backend ingest) have no way to detect a subscription->billable-key failover today: credential_source only reaches the router's OWN Postgres telemetry insert (InsertTelemetryParams), never the OTel span/log attrs that recordCallLog reuses as its metadata base. Add credential.source to both ProxyMessages' and ProxyOpenAIChatCompletion's upstream span builders, computed once via the existing credentialKeyParts(ctx) (moved earlier so it's in scope for the builder, not just the later telemetry-insert call site). Also add dispatch.subscription_failover_billable, the OR of the pre-emptive (subscriptionFailingOver, an already-exhausted observer snapshot skipping dispatch entirely) and reactive (subscriptionFailoverUsed, a live 429/OAuth-rejection retry) subscription-failover paths on the Anthropic side. The existing dispatch.subscription_failover only covers the reactive branch -- on its own it would under-report, since the pre-emptive path is the more common trigger in practice. Extends the two failover_integration_test.go tests added in the prior commit to also assert CredentialSource on the captured telemetry row. --- internal/proxy/failover_integration_test.go | 16 ++++++++++++--- internal/proxy/service.go | 22 +++++++++++++++++++-- 2 files changed, 33 insertions(+), 5 deletions(-) diff --git a/internal/proxy/failover_integration_test.go b/internal/proxy/failover_integration_test.go index d2aea04b4..cafd3c8cd 100644 --- a/internal/proxy/failover_integration_test.go +++ b/internal/proxy/failover_integration_test.go @@ -550,7 +550,8 @@ func TestProxyMessages_PreemptiveSubscriptionFailoverWarnsUser(t *testing.T) { obs.Record(obs.Key([]byte(bypassSubToken)), usage.Snapshot{ Secondary: usage.Window{UsedPercent: 1.0, WindowMinutes: 10080}, }) - svc := proxy.NewService(fr, map[string]providers.Client{providers.ProviderAnthropic: p}, nil, false, nil, nil, false, providers.ProviderAnthropic, bypassScorerPickMdl, nil). + telemetry := newCaptureTelemetry() + svc := proxy.NewService(fr, map[string]providers.Client{providers.ProviderAnthropic: p}, nil, false, nil, nil, false, providers.ProviderAnthropic, bypassScorerPickMdl, telemetry). WithSubscriptionAwareRouting(obs, 0.05, 2.0). WithDeploymentKeyedProviders(map[string]struct{}{providers.ProviderAnthropic: {}}) @@ -561,6 +562,7 @@ func TestProxyMessages_PreemptiveSubscriptionFailoverWarnsUser(t *testing.T) { // bypassCtx sets the usage-bypass gate; the routed path is what we're // exercising here, so use a plain subscription context instead. ctx := context.WithValue(context.Background(), proxy.AnthropicSubscriptionContextKey{}, bypassSubToken) + ctx = context.WithValue(ctx, proxy.InstallationIDContextKey{}, "11111111-1111-1111-1111-111111111111") require.NoError(t, svc.ProxyMessages(ctx, body, rec, req)) require.Len(t, p.proxyCreds, 1, "the turn must be dispatched once, on the deployment key") @@ -573,6 +575,9 @@ func TestProxyMessages_PreemptiveSubscriptionFailoverWarnsUser(t *testing.T) { "the client must see the billable-failover warning, not the normal routing marker") assert.NotContains(t, respBody, "· "+"best pick for this turn", "the failover warning replaces the routing marker, it doesn't append to it") + row := telemetry.firstRow(t) + assert.NotEqual(t, "subscription", row.CredentialSource, + "the telemetry row must not attribute this billable turn to the spent subscription") } // TestProxyMessages_ReactiveSubscriptionFailoverWarnsUser guards the reactive @@ -617,10 +622,11 @@ func TestProxyMessages_ReactiveSubscriptionFailoverWarnsUser(t *testing.T) { })) defer upstream.Close() + telemetry := newCaptureTelemetry() svc := proxy.NewService( &fakeRouter{decision: router.Decision{Provider: providers.ProviderAnthropic, Model: bypassScorerPickMdl}}, map[string]providers.Client{providers.ProviderAnthropic: anthropic.NewClient("test-key", upstream.URL)}, - nil, false, nil, nil, false, providers.ProviderAnthropic, bypassScorerPickMdl, nil, + nil, false, nil, nil, false, providers.ProviderAnthropic, bypassScorerPickMdl, telemetry, ).WithDeploymentKeyedProviders(map[string]struct{}{providers.ProviderAnthropic: {}}) rec := httptest.NewRecorder() @@ -628,7 +634,8 @@ func TestProxyMessages_ReactiveSubscriptionFailoverWarnsUser(t *testing.T) { req.Header.Set("Authorization", "Bearer "+bypassSubToken) body := []byte(`{"model":"` + bypassScorerPickMdl + `","stream":true,"messages":[{"role":"user","content":"hi"}]}`) - require.NoError(t, svc.ProxyMessages(context.Background(), body, rec, req)) + ctx := context.WithValue(context.Background(), proxy.InstallationIDContextKey{}, "11111111-1111-1111-1111-111111111111") + require.NoError(t, svc.ProxyMessages(ctx, body, rec, req)) mu.Lock() assert.Equal(t, 4, calls, "3 same-binding attempts on the subscription, then 1 subscription-failover retry on the Weave key") @@ -639,4 +646,7 @@ func TestProxyMessages_ReactiveSubscriptionFailoverWarnsUser(t *testing.T) { assert.Contains(t, respBody, "your Claude subscription hit its usage limit", "the retry must carry the billable-failover warning") assert.NotContains(t, respBody, "rate_limit_error", "the failed subscription attempt must never commit bytes") + row := telemetry.firstRow(t) + assert.NotEqual(t, "subscription", row.CredentialSource, + "the telemetry row must attribute this billable retry to the Weave key, not the spent subscription") } diff --git a/internal/proxy/service.go b/internal/proxy/service.go index 6d9d8b906..4a3822610 100644 --- a/internal/proxy/service.go +++ b/internal/proxy/service.go @@ -2955,6 +2955,11 @@ func (s *Service) ProxyMessages(ctx context.Context, body []byte, w http.Respons in, out := extractor.Tokens() cacheCreation, cacheRead := extractor.CacheTokens() + // Computed here (rather than only at the later Postgres-telemetry site) + // so credential.source and the combined subscription-failover flag reach + // the router.call OTLP log record — WorkWeave's ingest reads that record, + // not this router's own Postgres telemetry table. + _, _, credSourceForAttrs := s.credentialKeyParts(ctx) upstreamBuilder := otel.NewAttrBuilder(40). String("request_id", requestID). String("external_id", externalID). @@ -2989,7 +2994,16 @@ func (s *Service) ProxyMessages(ctx context.Context, body []byte, w http.Respons Int64("dispatch.fallback_attempts", int64(winnerIdx)). Bool("dispatch.failover_used", finalProvider != primaryProvider || subscriptionFailoverUsed). Bool("dispatch.baseline_failover", baselineFailoverUsed). - Bool("dispatch.subscription_failover", subscriptionFailoverUsed) + Bool("dispatch.subscription_failover", subscriptionFailoverUsed). + // Combined pre-emptive (subscriptionFailingOver, an exhausted observer + // snapshot skipping dispatch entirely) OR reactive + // (subscriptionFailoverUsed, a live 429/OAuth-rejection retry) signal — + // either means this turn is billed on the Weave/BYOK key instead of the + // caller's own subscription. dispatch.subscription_failover above only + // covers the reactive branch; this one is the authoritative bit for + // admin alerting downstream. + Bool("dispatch.subscription_failover_billable", subscriptionFailingOver || subscriptionFailoverUsed). + String("credential.source", credSourceForAttrs) applyPlannerAttrs(upstreamBuilder, routeRes) applyRoutingStateAttrs(upstreamBuilder, routeRes, decision.Model, sessionKey) addTimingAttrs(ctx, upstreamBuilder) @@ -4747,6 +4761,9 @@ func (s *Service) ProxyOpenAIChatCompletion(ctx context.Context, body []byte, w in, out := extractor.Tokens() cacheCreation, cacheRead := extractor.CacheTokens() + // See ProxyMessages for why this is computed before the builder rather + // than only at the later Postgres-telemetry site. + _, _, openaiCredSourceForAttrs := s.credentialKeyParts(ctx) openaiUpstreamBuilder := otel.NewAttrBuilder(40). String("request_id", requestID). String("external_id", externalID). @@ -4777,7 +4794,8 @@ func (s *Service) ProxyOpenAIChatCompletion(ctx context.Context, body []byte, w String("dispatch.primary_provider", primaryProvider). String("dispatch.final_provider", finalProvider). Int64("dispatch.fallback_attempts", int64(winnerIdx)). - Bool("dispatch.failover_used", finalProvider != primaryProvider) + Bool("dispatch.failover_used", finalProvider != primaryProvider). + String("credential.source", openaiCredSourceForAttrs) applyPlannerAttrs(openaiUpstreamBuilder, routeRes) applyRoutingStateAttrs(openaiUpstreamBuilder, routeRes, decision.Model, sessionKey) addTimingAttrs(ctx, openaiUpstreamBuilder)