Add --fable-bedrock-primary to route Fable to Bedrock first (#65)

When enabled (flag or SUBROUTER_FABLE_BEDROCK_PRIMARY), Claude Fable
requests serve from the AWS Bedrock gateway before the subscription pool.
A non-2xx or unreachable Bedrock restores the request body and falls
through to the normal pool path, which keeps its own Bedrock/API-key
fallback, so Fable never hard-fails. Defaults off; only takes effect when
the Bedrock gateway is configured. Other Claude models are unaffected.

Author: lawrencecchen

cmd/subrouter/main.go

@@ -73,6 +73,17 @@ func newCLIFileLogHandler(path string) slog.Handler {
 	return slog.NewTextHandler(file, opts)
 }
 
+// envTrue reports whether an environment variable is set to a truthy value
+// ("1", "true", "yes", "on", case-insensitive).
+func envTrue(name string) bool {
+	switch strings.ToLower(strings.TrimSpace(os.Getenv(name))) {
+	case "1", "true", "yes", "on":
+		return true
+	default:
+		return false
+	}
+}
+
 func run(args []string) error {
 	return runForProgram("subrouter", args)
 }
@@ -182,6 +193,7 @@ func serve(args []string) error {
 	bedrockGatewayToken := flags.String("bedrock-gateway-token", "", "optional bearer token clients must present to the Bedrock gateway; defaults to SUBROUTER_BEDROCK_GATEWAY_TOKEN")
 	bedrockProfiles := flags.String("bedrock-profiles", "", "comma-separated AWS profiles for the Bedrock gateway; defaults to SUBROUTER_BEDROCK_PROFILES or discovered awN profiles")
 	bedrockAutoBump := flags.Bool("bedrock-autobump", false, "request a Service Quotas increase (2x, deduped) when Bedrock throttles Fable/Opus")
+	fableBedrockPrimary := flags.Bool("fable-bedrock-primary", false, "route Claude Fable to Bedrock first (before the subscription pool); defaults to SUBROUTER_FABLE_BEDROCK_PRIMARY")
 	if err := flags.Parse(args); err != nil {
 		return err
 	}
@@ -298,25 +310,26 @@ func serve(args []string) error {
 	})
 
 	server := proxy.Server{
-		Upstream:          upstream,
-		CodexUpstream:     codexUpstream,
-		APIUpstream:       apiUpstream,
-		ClaudeUpstream:    claudeUpstream,
-		KimiUpstream:      kimiUpstream,
-		ZAIUpstream:       zaiUpstream,
-		Accounts:          nil,
-		AccountRef:        accountRef,
-		Sessions:          store,
-		SchedulerRef:      schedulerRef,
-		UsageScoreTTL:     usageScoreTTLForServe(*fetchUsage, *usageScoreTTL),
-		Transport:         outboundTransport,
-		Logger:            slog.Default(),
-		Lifecycle:         proxy.NewLifecycle(),
-		AdminToken:        *adminToken,
-		MaxBodyBytes:      *maxBodyBytes,
-		Bedrock:           bedrockConfig,
-		ClaudeFableAPIKey: strings.TrimSpace(os.Getenv("SUBROUTER_CLAUDE_FABLE_API_KEY")),
-		Transcripts:       transcript.NewRecorder(*transcriptDir),
+		Upstream:            upstream,
+		CodexUpstream:       codexUpstream,
+		APIUpstream:         apiUpstream,
+		ClaudeUpstream:      claudeUpstream,
+		KimiUpstream:        kimiUpstream,
+		ZAIUpstream:         zaiUpstream,
+		Accounts:            nil,
+		AccountRef:          accountRef,
+		Sessions:            store,
+		SchedulerRef:        schedulerRef,
+		UsageScoreTTL:       usageScoreTTLForServe(*fetchUsage, *usageScoreTTL),
+		Transport:           outboundTransport,
+		Logger:              slog.Default(),
+		Lifecycle:           proxy.NewLifecycle(),
+		AdminToken:          *adminToken,
+		MaxBodyBytes:        *maxBodyBytes,
+		Bedrock:             bedrockConfig,
+		ClaudeFableAPIKey:   strings.TrimSpace(os.Getenv("SUBROUTER_CLAUDE_FABLE_API_KEY")),
+		FableBedrockPrimary: *fableBedrockPrimary || envTrue("SUBROUTER_FABLE_BEDROCK_PRIMARY"),
+		Transcripts:         transcript.NewRecorder(*transcriptDir),
 	}
 	transcriptGCSSyncer := transcript.NewGCSSyncer(transcript.GCSSyncerConfig{
 		SourceDir:      *transcriptDir,

internal/proxy/claude_fable.go

@@ -36,6 +36,56 @@ func (s Server) claudeFableRequest(r *http.Request) bool {
 	return claudeFableModel(session.ExtractModel(r, s.MaxBodyBytes))
 }
 
+// serveClaudeFableBedrockPrimary serves a Fable request from AWS Bedrock before
+// the subscription pool is consulted (FableBedrockPrimary mode). It streams the
+// response only on a 2xx from Bedrock; on any non-2xx status or a Bedrock error
+// it restores the request body and returns false so the caller continues to the
+// normal pool path (which keeps its own Bedrock/API-key fallback). This makes
+// Bedrock the primary Fable route without ever hard-failing when Bedrock is
+// throttled or unreachable.
+func (s Server) serveClaudeFableBedrockPrimary(w http.ResponseWriter, r *http.Request) bool {
+	body, err := io.ReadAll(io.LimitReader(r.Body, replayablePostMaxBodyBytes))
+	if err != nil {
+		return false
+	}
+	restore := func() {
+		r.Body = io.NopCloser(bytes.NewReader(body))
+		r.ContentLength = int64(len(body))
+		r.GetBody = func() (io.ReadCloser, error) { return io.NopCloser(bytes.NewReader(body)), nil }
+	}
+	resp, err := s.claudeFableBedrockResponse(r.Context(), body)
+	if err != nil {
+		if s.Logger != nil && r.Context().Err() == nil {
+			s.Logger.Warn("fable bedrock-primary request failed, falling through to pool", "error", err)
+		}
+		restore()
+		return false
+	}
+	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		if s.Logger != nil {
+			s.Logger.Warn("fable bedrock-primary non-2xx, falling through to pool", "status", resp.StatusCode)
+		}
+		_ = resp.Body.Close()
+		restore()
+		return false
+	}
+	if s.Logger != nil {
+		s.Logger.Info("serving claude fable via bedrock-primary", "status", resp.StatusCode)
+	}
+	defer resp.Body.Close()
+	for key, values := range resp.Header {
+		if isHopByHopHeader(key) {
+			continue
+		}
+		for _, value := range values {
+			w.Header().Add(key, value)
+		}
+	}
+	w.WriteHeader(resp.StatusCode)
+	flushingCopy(w, resp.Body, nil)
+	return true
+}
+
 // serveClaudeFableFallback serves a Fable request straight off the fallback
 // chain (Bedrock, then the dedicated API key). Used when the subscription pool
 // cannot even start: no usable Claude OAuth account or selection failed. It

internal/proxy/claude_fable_test.go

@@ -537,3 +537,57 @@ func TestClaudeFableBedrockStripsUnsupportedFields(t *testing.T) {
 		t.Fatalf("bedrock body missing anthropic_version: %s", forwarded)
 	}
 }
+
+// bedrockPrimaryServer builds a Server with the Bedrock gateway configured and
+// FableBedrockPrimary enabled, whose Bedrock upstream returns the given status
+// and body.
+func bedrockPrimaryServer(status int, body string) Server {
+	rt := bedrockRoundTripFunc(func(req *http.Request) (*http.Response, error) {
+		return &http.Response{
+			StatusCode: status,
+			Body:       io.NopCloser(strings.NewReader(body)),
+			Header:     http.Header{"Content-Type": []string{"application/json"}},
+		}, nil
+	})
+	return Server{
+		MaxBodyBytes:        1 << 20,
+		FableBedrockPrimary: true,
+		Bedrock:             &BedrockConfig{Regions: []string{"us-east-1"}, Credentials: staticBedrockCreds(), Transport: rt},
+	}
+}
+
+func TestServeClaudeFableBedrockPrimarySuccess(t *testing.T) {
+	s := bedrockPrimaryServer(200, `{"model":"claude-fable-5","content":[{"type":"text","text":"ok"}]}`)
+	req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"claude-fable-5","max_tokens":8,"messages":[]}`))
+	rec := httptest.NewRecorder()
+	if !s.serveClaudeFableBedrockPrimary(rec, req) {
+		t.Fatal("expected bedrock-primary to serve a 2xx Bedrock response")
+	}
+	if rec.Code != 200 {
+		t.Fatalf("status = %d, want 200 (body %q)", rec.Code, rec.Body.String())
+	}
+	if !strings.Contains(rec.Body.String(), "claude-fable-5") {
+		t.Fatalf("body = %q, want forwarded Bedrock body", rec.Body.String())
+	}
+}
+
+func TestServeClaudeFableBedrockPrimaryFallsThroughOnNon2xx(t *testing.T) {
+	s := bedrockPrimaryServer(429, `{"message":"Too many requests"}`)
+	bodyStr := `{"model":"claude-fable-5","max_tokens":8,"messages":[]}`
+	req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(bodyStr))
+	rec := httptest.NewRecorder()
+	if s.serveClaudeFableBedrockPrimary(rec, req) {
+		t.Fatal("expected bedrock-primary to fall through on a non-2xx Bedrock response")
+	}
+	if rec.Code != 200 { // httptest default; nothing should have been written
+		t.Fatalf("status = %d, want untouched recorder (nothing written)", rec.Code)
+	}
+	// Body must be restored so the pool path can read it.
+	restored, err := io.ReadAll(req.Body)
+	if err != nil {
+		t.Fatalf("read restored body: %v", err)
+	}
+	if string(restored) != bodyStr {
+		t.Fatalf("restored body = %q, want %q", string(restored), bodyStr)
+	}
+}

internal/proxy/proxy.go

@@ -65,6 +65,13 @@ type Server struct {
 	// ONLY to Fable; Opus/Sonnet/etc. continue to use the OAuth pool and never
 	// touch this key.
 	ClaudeFableAPIKey string
+	// FableBedrockPrimary, when true, routes Claude Fable requests to AWS Bedrock
+	// FIRST, before the subscription pool, instead of using Bedrock only as a
+	// fallback. It only takes effect when the Bedrock gateway is configured; a
+	// non-2xx Bedrock response (or an unreachable Bedrock) falls through to the
+	// normal pool path, which keeps its own Bedrock/API-key fallback. Applies
+	// ONLY to Fable; other Claude models are unaffected.
+	FableBedrockPrimary bool
 }
 
 type ActiveSessions struct {
@@ -1649,6 +1656,15 @@ func (s Server) proxyHandler() http.Handler {
 			fableFallbackConfigured = s.claudeFableEnabled() && claudeFableModel(requestModel) &&
 				r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/v1/messages")
 		}
+		// Bedrock-primary: when enabled, serve Fable straight from Bedrock before
+		// touching the subscription pool. A non-2xx or unreachable Bedrock restores
+		// the body and falls through to the normal pool path (which still carries
+		// its own Bedrock/API-key fallback), so this never hard-fails Fable.
+		if s.FableBedrockPrimary && fableFallbackConfigured && s.Bedrock != nil && s.Bedrock.configured() {
+			if s.serveClaudeFableBedrockPrimary(w, r) {
+				return
+			}
+		}
 		sessionAgentType := agentTypeForProviderSession(agentType, requestProvider)
 		account, sessionID, userEmail, err := s.accountForSessionProvider(requestProvider, sessionAgentType, sessionID, r)
 		if err != nil {