diff --git a/internal/proxy/failover_integration_test.go b/internal/proxy/failover_integration_test.go index d506c9cb3..cafd3c8cd 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,126 @@ 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}, + }) + 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: {}}) + + 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) + 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") + 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") + 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 +// 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() + + 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, telemetry, + ).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"}]}`) + + 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") + 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") + 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 9ad50fc56..4a3822610 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{} @@ -2944,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). @@ -2978,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) @@ -4736,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). @@ -4766,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) 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