diff --git a/internal/proxy/claude_desktop.go b/internal/proxy/claude_desktop.go index 4f7486bb5..64d8b711f 100644 --- a/internal/proxy/claude_desktop.go +++ b/internal/proxy/claude_desktop.go @@ -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, diff --git a/internal/proxy/claude_desktop_test.go b/internal/proxy/claude_desktop_test.go index 27b55d520..b00b3828b 100644 --- a/internal/proxy/claude_desktop_test.go +++ b/internal/proxy/claude_desktop_test.go @@ -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) } }) }