proxy: continue requests when the model catalog changes (#18058)

2026-09-21 13:38:14 -05:00 · 2026-08-26 19:39:05 -07:00
parent a67fe8c537
commit d366f4868a
2 changed files with 109 additions and 30 deletions
@@ -319,7 +319,7 @@ func (p *ClaudeDesktop) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusNoContent)
 		return
 	}
-	generation, models := p.modelSnapshotWithGeneration()
+	_, models := p.modelSnapshot()
 	switch r.URL.Path {
 	case "/v1/models":
 		if r.Method != http.MethodGet {
@@ -330,7 +330,7 @@ func (p *ClaudeDesktop) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			if err := p.refreshModelCatalog(r.Context()); err != nil {
 				p.logger.Debug("could not refresh Claude model catalog", "error", err)
 			}
-			_, models = p.modelSnapshotWithGeneration()
+			_, models = p.modelSnapshot()
 		}
 		p.serveModels(w, r.Context(), models)
 		return
@@ -339,14 +339,14 @@ func (p *ClaudeDesktop) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			methodNotAllowed(w, http.MethodPost)
 			return
 		}
-		p.serveTokenCount(w, r, generation, models)
+		p.serveTokenCount(w, r, models)
 		return
 	case "/v1/messages":
 		if r.Method != http.MethodPost {
 			methodNotAllowed(w, http.MethodPost)
 			return
 		}
-		if err := p.routeModel(r, generation, models); err != nil {
+		if err := p.routeModel(r, models); err != nil {
 			var accessErr *claudeDesktopAccessError
 			if errors.As(err, &accessErr) {
 				writeAnthropicError(w, accessErr.status, accessErr)
@@ -554,7 +554,7 @@ func (p *ClaudeDesktop) serveModels(w http.ResponseWriter, ctx context.Context,
 	}
 }

-func (p *ClaudeDesktop) serveTokenCount(w http.ResponseWriter, r *http.Request, generation uint64, models []ClaudeDesktopModel) {
+func (p *ClaudeDesktop) serveTokenCount(w http.ResponseWriter, r *http.Request, models []ClaudeDesktopModel) {
 	body, err := readRequestBody(r)
 	if err != nil {
 		writeAnthropicError(w, http.StatusBadRequest, err)
@@ -575,10 +575,8 @@ func (p *ClaudeDesktop) serveTokenCount(w http.ResponseWriter, r *http.Request,
 			p.logger.Debug("could not refresh Claude model catalog", "error", err)
 		}
 	}
-	if !p.modelGenerationMatches(generation) {
-		writeAnthropicError(w, http.StatusConflict, errors.New("Claude model catalog changed; try again"))
-		return
-	}
+	// Continue with the model snapshot that admitted this request. A concurrent
+	// catalog refresh must not turn a valid token-count request into an error.
 	if access := p.modelAccess(r.Context(), selected); access.Availability != ClaudeDesktopAvailabilityAvailable {
 		accessErr := newClaudeDesktopAccessError(selected, access)
 		writeAnthropicError(w, accessErr.status, accessErr)
@@ -593,7 +591,7 @@ func (p *ClaudeDesktop) serveTokenCount(w http.ResponseWriter, r *http.Request,
 	}
 }

-func (p *ClaudeDesktop) routeModel(r *http.Request, generation uint64, models []ClaudeDesktopModel) error {
+func (p *ClaudeDesktop) routeModel(r *http.Request, models []ClaudeDesktopModel) error {
 	body, err := readRequestBody(r)
 	if err != nil {
 		return err
@@ -615,9 +613,8 @@ func (p *ClaudeDesktop) routeModel(r *http.Request, generation uint64, models []
 			p.logger.Debug("could not refresh Claude model catalog", "error", err)
 		}
 	}
-	if !p.modelGenerationMatches(generation) {
-		return errors.New("Claude model catalog changed; try again")
-	}
+	// Continue with the model snapshot that admitted this request. A concurrent
+	// catalog refresh may affect later requests, but not one already in flight.
 	if access := p.modelAccess(r.Context(), selected); access.Availability != ClaudeDesktopAvailabilityAvailable {
 		return newClaudeDesktopAccessError(selected, access)
 	}
@@ -750,12 +747,6 @@ func claudeDesktopModelForID(models []ClaudeDesktopModel, id string) (ClaudeDesk
 	return ClaudeDesktopModel{}, fmt.Errorf("unknown Claude model %q", id)
 }

-func (p *ClaudeDesktop) modelGenerationMatches(generation uint64) bool {
-	p.modelsMu.RLock()
-	defer p.modelsMu.RUnlock()
-	return p.modelsGeneration == generation
-}
-
 func (p *ClaudeDesktop) accessFacts(ctx context.Context) (ClaudeDesktopAccessState, map[string]struct{}, bool) {
 	state := ClaudeDesktopAccessState{
 		Cloud:   ClaudeDesktopCloudOn,
@@ -812,7 +812,7 @@ func TestGatewayRefreshesEntitlementsWithoutRestart(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	assertResponseContains(t, resp, http.StatusBadRequest, "model catalog changed")
+	assertResponseContains(t, resp, http.StatusOK, `"ok":true`)

 	resp, err = http.Post(
 		"http://"+p.Addr()+"/v1/messages",
@@ -860,7 +860,7 @@ func TestGatewayDiscardsRefreshThatRacesWithModelUpdate(t *testing.T) {
 	}
 }

-func TestGatewayRejectsRequestWhenSlotChangesDuringRefresh(t *testing.T) {
+func TestGatewayContinuesRequestWhenSlotChangesDuringRefresh(t *testing.T) {
 	var upstreamCalls atomic.Int32
 	upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
 		upstreamCalls.Add(1)
@@ -929,18 +929,98 @@ func TestGatewayRejectsRequestWhenSlotChangesDuringRefresh(t *testing.T) {
 	case err := <-requestErr:
 		t.Fatal(err)
 	case resp := <-response:
-		if resp.status != http.StatusBadRequest || !strings.Contains(resp.body, "model catalog changed") {
-			t.Fatalf("response = (%d, %q), want status %d containing %q", resp.status, resp.body, http.StatusBadRequest, "model catalog changed")
+		if resp.status != http.StatusOK {
+			t.Fatalf("response = (%d, %q), want status %d", resp.status, resp.body, http.StatusOK)
 		}
 	case <-time.After(2 * time.Second):
 		t.Fatal("timed out waiting for request")
 	}
-	if got := upstreamCalls.Load(); got != 0 {
-		t.Fatalf("upstream calls = %d, want 0 after a slot change", got)
+	if got := upstreamCalls.Load(); got != 1 {
+		t.Fatalf("upstream calls = %d, want 1 after a slot change", got)
 	}
 }

-func TestGatewayRejectsAdmittedRequestAfterSlotReassignment(t *testing.T) {
+func TestGatewayContinuesTokenCountWhenSlotChangesDuringRefresh(t *testing.T) {
+	initial := ClaudeDesktopModelsFromRecommendations([]api.ModelRecommendation{{Model: "model-a:cloud", RequiredPlan: "free"}})
+	updated := ClaudeDesktopModelsFromRecommendations([]api.ModelRecommendation{{Model: "model-b:cloud", RequiredPlan: "free"}})
+	started := make(chan struct{})
+	release := make(chan struct{})
+	p, err := NewClaudeDesktop(ClaudeDesktopConfig{
+		ListenAddr: "127.0.0.1:0",
+		OllamaURL:  "http://127.0.0.1:11434",
+		Model:      initial[0].OllamaModel,
+		Models:     initial,
+		RefreshModels: func(_ context.Context, current []ClaudeDesktopModel) ([]ClaudeDesktopModel, error) {
+			close(started)
+			<-release
+			return current, nil
+		},
+		ResolveAccessState: func(context.Context) (ClaudeDesktopAccessState, error) {
+			return ClaudeDesktopAccessState{Cloud: ClaudeDesktopCloudOn, Account: ClaudeDesktopAccountSignedIn, Plan: "free"}, nil
+		},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := p.Start(); err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		_ = p.Close(ctx)
+	})
+
+	type tokenCountResponse struct {
+		status int
+		body   []byte
+	}
+	response := make(chan tokenCountResponse, 1)
+	requestErr := make(chan error, 1)
+	go func() {
+		resp, err := http.Post(
+			"http://"+p.Addr()+"/v1/messages/count_tokens",
+			"application/json",
+			strings.NewReader(`{"model":"claude-fable-5","messages":[{"role":"user","content":"hello"}]}`),
+		)
+		if err != nil {
+			requestErr <- err
+			return
+		}
+		body, err := io.ReadAll(resp.Body)
+		_ = resp.Body.Close()
+		if err != nil {
+			requestErr <- err
+			return
+		}
+		response <- tokenCountResponse{status: resp.StatusCode, body: body}
+	}()
+	<-started
+	if err := p.SetModels(updated); err != nil {
+		t.Fatal(err)
+	}
+	close(release)
+
+	select {
+	case err := <-requestErr:
+		t.Fatal(err)
+	case resp := <-response:
+		if resp.status != http.StatusOK {
+			t.Fatalf("response = (%d, %q), want status %d", resp.status, resp.body, http.StatusOK)
+		}
+		var result anthropic.CountTokensResponse
+		if err := json.Unmarshal(resp.body, &result); err != nil {
+			t.Fatal(err)
+		}
+		if result.InputTokens <= 0 {
+			t.Fatalf("input_tokens = %d, want positive estimate", result.InputTokens)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("timed out waiting for token-count response")
+	}
+}
+
+func TestGatewayRoutesAdmittedRequestAfterSlotReassignment(t *testing.T) {
 	for _, tt := range []struct {
 		name    string
 		initial []ClaudeDesktopModel
@@ -967,7 +1047,7 @@ func TestGatewayRejectsAdmittedRequestAfterSlotReassignment(t *testing.T) {
 			if err != nil {
 				t.Fatal(err)
 			}
-			generation, admitted := p.modelSnapshotWithGeneration()
+			_, admitted := p.modelSnapshot()
 			if err := p.SetModels(tt.updated); err != nil {
 				t.Fatal(err)
 			}
@@ -976,9 +1056,17 @@ func TestGatewayRejectsAdmittedRequestAfterSlotReassignment(t *testing.T) {
 				"http://127.0.0.1/v1/messages",
 				strings.NewReader(`{"model":"claude-fable-5","messages":[]}`),
 			)
-			err = p.routeModel(request, generation, admitted)
-			if err == nil || !strings.Contains(err.Error(), "model catalog changed") {
-				t.Fatalf("route error = %v, want catalog-changed rejection", err)
+			if err := p.routeModel(request, admitted); err != nil {
+				t.Fatal(err)
+			}
+			var payload struct {
+				Model string `json:"model"`
+			}
+			if err := json.NewDecoder(request.Body).Decode(&payload); err != nil {
+				t.Fatal(err)
+			}
+			if payload.Model != tt.initial[0].OllamaModel {
+				t.Fatalf("routed model = %q, want admitted model %q", payload.Model, tt.initial[0].OllamaModel)
 			}
 		})
 	}