feat(proxy): warn Claude Code when a subscription turn fails over to the billable Weave key

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")
+}

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)

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