diff --git a/internal/api/handlers/management/auth_files.go b/internal/api/handlers/management/auth_files.go index 0a4ed74f70..86878e9bf5 100644 --- a/internal/api/handlers/management/auth_files.go +++ b/internal/api/handlers/management/auth_files.go @@ -119,9 +119,9 @@ func (h *Handler) ListAuthFiles(c *gin.Context) { h.mu.Lock() host := h.pluginHost h.mu.Unlock() - var quotaSupportedProviders map[string]struct{} + quota := authFileQuotaContext{v8: c.GetBool(ConfigV8ContextKey)} if host != nil { - quotaSupportedProviders = host.QuotaSupportedProvidersSet(c.Request.Context()) + quota.pluginIDs = host.QuotaProviderPluginIDs(c.Request.Context()) } auths := h.authManager.List() observedAt := time.Now().UTC() @@ -141,7 +141,7 @@ func (h *Handler) ListAuthFiles(c *gin.Context) { start, end := pagination.bounds(total) files := make([]gin.H, 0, end-start) for _, auth := range matching[start:end] { - if entry := h.buildAuthFileEntry(auth, quotaSupportedProviders); entry != nil { + if entry := h.buildAuthFileEntry(auth, quota); entry != nil { entry["cooldowns"] = nil if cooldownsKnown { entry["cooldowns"] = coreauth.CooldownSnapshotForAuth(auth, observedAt) @@ -157,7 +157,7 @@ func (h *Handler) ListAuthFiles(c *gin.Context) { if !matchesAuthFileLookup(auth, nameFilter, authIndexFilter) { continue } - if entry := h.buildAuthFileEntry(auth, quotaSupportedProviders); entry != nil { + if entry := h.buildAuthFileEntry(auth, quota); entry != nil { entry["cooldowns"] = nil if cooldownsKnown { entry["cooldowns"] = coreauth.CooldownSnapshotForAuth(auth, observedAt) @@ -478,10 +478,19 @@ func (h *Handler) listAuthFilesFromDisk(c *gin.Context, pagination authFilesPagi c.JSON(200, authFilesListResponse(observedAt, files, pagination, total, end)) } -func (h *Handler) buildAuthFileEntry(auth *coreauth.Auth, quotaSupported ...map[string]struct{}) gin.H { +// authFileQuotaContext carries request-scoped quota capability data for auth file entries. +type authFileQuotaContext struct { + // pluginIDs maps normalized provider keys to the quota plugin selected for them. + // A nil map means no quota plugins are available. + pluginIDs map[string]string + // v8 hides declarative quota_probe support because v8 only serves plugin quota routes. + v8 bool +} + +func (h *Handler) buildAuthFileEntry(auth *coreauth.Auth, quota authFileQuotaContext) gin.H { authFileEntryMu.Lock() defer authFileEntryMu.Unlock() - return h.buildAuthFileEntryLocked(auth, quotaSupported...) + return h.buildAuthFileEntryLocked(auth, quota) } func isPersistentAuthFailure(auth *coreauth.Auth, now time.Time) bool { @@ -634,7 +643,7 @@ func reconcileAuthFileCooldownState(auth *coreauth.Auth, now time.Time) (unavail return unavailable, status, statusMessage, nextRetry } -func (h *Handler) buildAuthFileEntryLocked(auth *coreauth.Auth, quotaSupported ...map[string]struct{}) gin.H { +func (h *Handler) buildAuthFileEntryLocked(auth *coreauth.Auth, quota authFileQuotaContext) gin.H { if auth == nil { return nil } @@ -674,25 +683,15 @@ func (h *Handler) buildAuthFileEntryLocked(auth *coreauth.Auth, quotaSupported . if modelQuotas := modelQuotaObservationPayload(auth.Provider, auth.ModelStates); len(modelQuotas) > 0 { entry["model_quotas"] = modelQuotas } - var quotaSupportedMap map[string]struct{} - if len(quotaSupported) > 0 { - quotaSupportedMap = quotaSupported[0] - } - if quotaSupportedMap != nil { - if _, ok := quotaSupportedMap[strings.ToLower(strings.TrimSpace(auth.Provider))]; ok { - entry["supports_quota"] = true - entry["quota_provider"] = auth.Provider - } - } else { - h.mu.Lock() - host := h.pluginHost - h.mu.Unlock() - if host != nil && host.HasQuotaProvider(auth.Provider) { - entry["supports_quota"] = true - entry["quota_provider"] = auth.Provider + if pluginID := quota.pluginIDs[strings.ToLower(strings.TrimSpace(auth.Provider))]; pluginID != "" { + entry["supports_quota"] = true + entry["quota_provider"] = auth.Provider + if quota.v8 { + entry["quota_plugin_id"] = pluginID } } - if auth.Metadata != nil { + // Declarative quota_probe is only executed by the v0 quota/fetch route. + if !quota.v8 && auth.Metadata != nil { if probe, okProbe := auth.Metadata["quota_probe"]; okProbe && probe != nil { entry["supports_quota"] = true entry["quota_probe"] = probe diff --git a/internal/api/handlers/management/auth_files_filter_test.go b/internal/api/handlers/management/auth_files_filter_test.go index 23fef5c718..474c80442f 100644 --- a/internal/api/handlers/management/auth_files_filter_test.go +++ b/internal/api/handlers/management/auth_files_filter_test.go @@ -238,7 +238,7 @@ func TestAuthFileLookupAndEntryBuildConcurrentEnsureIndex(t *testing.T) { if !matchesAuthFileLookup(auth, fileName, "idx-concurrent") { t.Errorf("auth lookup did not match") } - entry := h.buildAuthFileEntry(auth) + entry := h.buildAuthFileEntry(auth, authFileQuotaContext{}) if entry == nil { t.Errorf("entry is nil") continue diff --git a/internal/api/handlers/management/plugin_quota_test.go b/internal/api/handlers/management/plugin_quota_test.go index d7abb3fe4e..a58b81f417 100644 --- a/internal/api/handlers/management/plugin_quota_test.go +++ b/internal/api/handlers/management/plugin_quota_test.go @@ -18,6 +18,7 @@ import ( type mockQuotaProvider struct { identifier string + supported []string fetchFn func(context.Context, pluginapi.QuotaFetchRequest) (pluginapi.QuotaFetchResponse, error) resetFn func(context.Context, pluginapi.QuotaResetRequest) (pluginapi.QuotaResetResponse, error) } @@ -27,8 +28,12 @@ func (m *mockQuotaProvider) Identifier() string { } func (m *mockQuotaProvider) DescribeQuota(ctx context.Context, req pluginapi.QuotaDescribeRequest) (pluginapi.QuotaDescribeResponse, error) { + supported := m.supported + if supported == nil { + supported = []string{m.identifier} + } return pluginapi.QuotaDescribeResponse{ - SupportedProviders: []string{m.identifier}, + SupportedProviders: supported, DisplayName: "Mock Quota " + m.identifier, SupportsReset: true, }, nil @@ -371,6 +376,9 @@ func TestAuthFilesList_IncludesQuotaSupport(t *testing.T) { if qp, ok := f["quota_provider"].(string); !ok || qp != "opencode-go" { t.Fatalf("expected quota_provider: opencode-go, got %#v", f["quota_provider"]) } + if _, ok := f["quota_plugin_id"]; ok { + t.Fatal("v0 auth files should not expose quota_plugin_id") + } case "custom-auth-list": foundProbe = true if supports, ok := f["supports_quota"].(bool); !ok || !supports { @@ -391,6 +399,94 @@ func TestAuthFilesList_IncludesQuotaSupport(t *testing.T) { } } +func TestAuthFilesListV8_QuotaPluginIDAndNoQuotaProbe(t *testing.T) { + manager := coreauth.NewManager(nil, nil, nil) + auths := []*coreauth.Auth{ + {ID: "foo-auth", FileName: "foo.json", Provider: "foo"}, + {ID: "bar-auth", FileName: "bar.json", Provider: "bar"}, + { + ID: "probe-auth", + FileName: "probe.json", + Provider: "custom", + Metadata: map[string]any{"quota_probe": map[string]any{"url": "https://api.example.com/quota"}}, + }, + } + for _, auth := range auths { + auth.Attributes = map[string]string{"runtime_only": "true"} + if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil { + t.Fatalf("failed to register auth %s: %v", auth.ID, errRegister) + } + } + + host := pluginhost.New() + // "a-multi-quota" sorts first, but a first-pass identifier match must still win for "foo". + host.RegisterPluginForTest("a-multi-quota", pluginapi.Plugin{ + Capabilities: pluginapi.Capabilities{ + QuotaProvider: &mockQuotaProvider{identifier: "multi-quota", supported: []string{"foo", "bar"}}, + }, + }) + host.RegisterPluginForTest("z-foo-quota", pluginapi.Plugin{ + Capabilities: pluginapi.Capabilities{ + QuotaProvider: &mockQuotaProvider{identifier: "foo"}, + }, + }) + + h := NewHandlerWithoutConfigFilePath(&config.Config{AuthDir: t.TempDir()}, manager) + h.SetPluginHost(host) + + listFiles := func(v8 bool) map[string]map[string]any { + t.Helper() + rec := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(rec) + ctx.Request = httptest.NewRequest(http.MethodGet, "/v8/management/credentials", nil) + if v8 { + ctx.Set(ConfigV8ContextKey, true) + } + h.ListAuthFiles(ctx) + if rec.Code != http.StatusOK { + t.Fatalf("expected status 200, got %d: %s", rec.Code, rec.Body.String()) + } + var resp struct { + Files []map[string]any `json:"files"` + } + if errUnmarshal := json.Unmarshal(rec.Body.Bytes(), &resp); errUnmarshal != nil { + t.Fatalf("failed to unmarshal files: %v", errUnmarshal) + } + out := make(map[string]map[string]any, len(resp.Files)) + for _, f := range resp.Files { + id, _ := f["id"].(string) + out[id] = f + } + return out + } + + v8Files := listFiles(true) + if got := v8Files["foo-auth"]["quota_plugin_id"]; got != "z-foo-quota" { + t.Fatalf("foo quota_plugin_id = %#v, want z-foo-quota", got) + } + if got := v8Files["bar-auth"]["quota_plugin_id"]; got != "a-multi-quota" { + t.Fatalf("bar quota_plugin_id = %#v, want a-multi-quota", got) + } + probeV8 := v8Files["probe-auth"] + if probeV8 == nil { + t.Fatal("probe auth missing from v8 list") + } + if _, ok := probeV8["supports_quota"]; ok { + t.Fatalf("v8 probe auth should not advertise supports_quota, got %#v", probeV8["supports_quota"]) + } + if _, ok := probeV8["quota_probe"]; ok { + t.Fatal("v8 probe auth should not expose quota_probe") + } + + v0Probe := listFiles(false)["probe-auth"] + if supports, ok := v0Probe["supports_quota"].(bool); !ok || !supports { + t.Fatalf("v0 probe auth should keep supports_quota: true, got %#v", v0Probe["supports_quota"]) + } + if v0Probe["quota_probe"] == nil { + t.Fatal("v0 probe auth should keep quota_probe") + } +} + func TestResetCredentialQuota_FailureDoesNotClearCooldown(t *testing.T) { manager := coreauth.NewManager(nil, nil, nil) next := time.Now().Add(time.Hour) diff --git a/internal/pluginhost/quota_provider.go b/internal/pluginhost/quota_provider.go index 5acf04f975..f366560d12 100644 --- a/internal/pluginhost/quota_provider.go +++ b/internal/pluginhost/quota_provider.go @@ -113,41 +113,67 @@ func (h *Host) HasQuotaProviderForPlugin(pluginID string) bool { return h.quotaProviderRecordByPlugin(pluginID) != nil } -// QuotaSupportedProvidersSet returns a point-in-time set of all provider keys supported by active quota providers. -// This allows callers (e.g. ListAuthFiles) to discover all quota-capable providers with zero per-credential RPCs. -func (h *Host) QuotaSupportedProvidersSet(ctx context.Context) map[string]struct{} { +// QuotaProviderPluginIDs returns a point-in-time map from normalized provider keys to the ID of the +// quota plugin that quotaProviderRecord would select for that provider, preserving its two-pass precedence. +// This allows callers (e.g. ListAuthFiles) to resolve quota plugins with zero per-credential RPCs. +func (h *Host) QuotaProviderPluginIDs(ctx context.Context) map[string]string { if h == nil { return nil } + if ctx == nil { + ctx = context.Background() + } snap := h.Snapshot() if snap == nil { return nil } records := h.activeRecordsFromSnapshot(snap) - out := make(map[string]struct{}) + // Keep candidates in selection order until probing finishes, because DescribeQuota may fuse a plugin. + candidates := make(map[string][]string) + addCandidate := func(provider, pluginID string) { + if provider != "" { + candidates[provider] = append(candidates[provider], pluginID) + } + } + quotaRecords := make([]capabilityRecord, 0, len(records)) + // First pass: QuotaProvider.Identifier(), plugin ID, or AuthProvider.Identifier(). for _, record := range records { if err := ctx.Err(); err != nil { - return out + break } quotaProvider := record.plugin.Capabilities.QuotaProvider if quotaProvider == nil || h.isPluginFused(record.id) { continue } - if identifier, okID := h.callQuotaIdentifier(record.id, quotaProvider); okID && identifier != "" { - out[identifier] = struct{}{} - } - if normID := normalizeProviderID(record.id); normID != "" { - out[normID] = struct{}{} + quotaRecords = append(quotaRecords, record) + if identifier, okID := h.callQuotaIdentifier(record.id, quotaProvider); okID { + addCandidate(identifier, record.id) } + addCandidate(normalizeProviderID(record.id), record.id) if record.plugin.Capabilities.AuthProvider != nil { - if authID, okAuth := h.callAuthProviderIdentifier(record.id, record.plugin.Capabilities.AuthProvider); okAuth && authID != "" { - out[authID] = struct{}{} + if authID, okAuth := h.callAuthProviderIdentifier(record.id, record.plugin.Capabilities.AuthProvider); okAuth { + addCandidate(authID, record.id) } } - supported := snap.cachedQuotaSupportedProviders(ctx, h, record, quotaProvider) - for _, p := range supported { - if clean := normalizeProviderID(p); clean != "" { - out[clean] = struct{}{} + } + // Second pass: SupportedProviders declared by DescribeQuota. + for _, record := range quotaRecords { + if err := ctx.Err(); err != nil { + break + } + if h.isPluginFused(record.id) { + continue + } + for _, p := range snap.cachedQuotaSupportedProviders(ctx, h, record, record.plugin.Capabilities.QuotaProvider) { + addCandidate(normalizeProviderID(p), record.id) + } + } + out := make(map[string]string, len(candidates)) + for provider, pluginIDs := range candidates { + for _, pluginID := range pluginIDs { + if !h.isPluginFused(pluginID) { + out[provider] = pluginID + break } } } diff --git a/internal/pluginhost/quota_provider_test.go b/internal/pluginhost/quota_provider_test.go index 16271bd1d4..3f6bc33156 100644 --- a/internal/pluginhost/quota_provider_test.go +++ b/internal/pluginhost/quota_provider_test.go @@ -493,3 +493,95 @@ func TestRegisterRPCPlugin_ContextCancellation(t *testing.T) { t.Fatal("expected error when registering with canceled context") } } + +func TestQuotaProviderPluginIDsSkipsPluginsFusedDuringDescribe(t *testing.T) { + for _, healthyIdentifier := range []string{"foo", "healthy-quota"} { + t.Run(healthyIdentifier, func(t *testing.T) { + broken := &testQuotaProvider{ + identifier: "foo", + describeFn: func(context.Context, pluginapi.QuotaDescribeRequest) (pluginapi.QuotaDescribeResponse, error) { + panic("broken quota description") + }, + } + healthy := &testQuotaProvider{ + identifier: healthyIdentifier, + describeFn: func(context.Context, pluginapi.QuotaDescribeRequest) (pluginapi.QuotaDescribeResponse, error) { + return pluginapi.QuotaDescribeResponse{SupportedProviders: []string{"foo"}}, nil + }, + } + host := newHostWithRecords( + capabilityRecord{id: "a-broken", plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{QuotaProvider: broken}}}, + capabilityRecord{id: "b-healthy", plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{QuotaProvider: healthy}}}, + ) + ctx := context.Background() + initial := host.quotaProviderRecord(ctx, "foo") + if initial == nil || initial.id != "a-broken" { + t.Fatalf("initial selection = %#v, want a-broken", initial) + } + + got := host.QuotaProviderPluginIDs(ctx) + if !host.isPluginFused("a-broken") { + t.Fatal("expected a-broken to be fused during DescribeQuota") + } + if got["foo"] != "b-healthy" { + t.Fatalf("quota plugin for foo = %q, want b-healthy", got["foo"]) + } + if _, exists := got["a-broken"]; exists { + t.Fatal("fused plugin's unique alias must be removed") + } + for provider, pluginID := range got { + if !host.HasQuotaProviderForPlugin(pluginID) { + t.Fatalf("provider %q references unavailable plugin %q", provider, pluginID) + } + record := host.quotaProviderRecord(ctx, provider) + if record == nil || record.id != pluginID { + t.Fatalf("quotaProviderRecord(%q) = %#v, want plugin %q", provider, record, pluginID) + } + } + }) + } +} + +func TestQuotaProviderPluginIDsMatchesQuotaProviderRecordPrecedence(t *testing.T) { + multi := &testQuotaProvider{ + identifier: "multi-quota", + describeFn: func(context.Context, pluginapi.QuotaDescribeRequest) (pluginapi.QuotaDescribeResponse, error) { + return pluginapi.QuotaDescribeResponse{SupportedProviders: []string{" FOO ", "bar"}}, nil + }, + } + host := newHostWithRecords( + capabilityRecord{id: "a-multi", plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{QuotaProvider: multi}}}, + // Same-pass collision: the earlier record wins "foo" over z-foo. + capabilityRecord{id: "b-foo", plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{QuotaProvider: &testQuotaProvider{identifier: "foo"}}}}, + // AuthProvider alias is a first-pass match and beats a-multi's SupportedProviders for "bar". + capabilityRecord{id: "m-bar-auth", plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + QuotaProvider: &testQuotaProvider{identifier: "m-quota"}, + AuthProvider: fakeAuthProvider{identifier: " Bar "}, + }}}, + capabilityRecord{id: "z-foo", plugin: pluginapi.Plugin{Capabilities: pluginapi.Capabilities{QuotaProvider: &testQuotaProvider{identifier: "foo"}}}}, + ) + + got := host.QuotaProviderPluginIDs(context.Background()) + want := map[string]string{ + "multi-quota": "a-multi", + "a-multi": "a-multi", + "foo": "b-foo", + "b-foo": "b-foo", + "bar": "m-bar-auth", + "m-quota": "m-bar-auth", + "m-bar-auth": "m-bar-auth", + "z-foo": "z-foo", + } + if len(got) != len(want) { + t.Fatalf("QuotaProviderPluginIDs() = %#v, want %#v", got, want) + } + for provider, pluginID := range want { + if got[provider] != pluginID { + t.Fatalf("QuotaProviderPluginIDs()[%q] = %q, want %q", provider, got[provider], pluginID) + } + record := host.quotaProviderRecord(context.Background(), provider) + if record == nil || record.id != pluginID { + t.Fatalf("quotaProviderRecord(%q) = %#v, want plugin %q", provider, record, pluginID) + } + } +} diff --git a/sdk/pluginhost/quota.go b/sdk/pluginhost/quota.go index 8c98905bbb..37b306da1b 100644 --- a/sdk/pluginhost/quota.go +++ b/sdk/pluginhost/quota.go @@ -18,6 +18,15 @@ func (h *Host) QuotaProviders(ctx context.Context) []RegisteredQuotaProviderInfo return h.inner.QuotaProviders(ctx) } +// QuotaProviderPluginIDs maps normalized provider keys to the selected active quota plugin. +// Selection preserves host precedence and excludes plugins fused during capability discovery. +func (h *Host) QuotaProviderPluginIDs(ctx context.Context) map[string]string { + if h == nil || h.inner == nil { + return nil + } + return h.inner.QuotaProviderPluginIDs(ctx) +} + // HasQuotaProviderContext reports whether a provider handles credential quotas. func (h *Host) HasQuotaProviderContext(ctx context.Context, provider string) bool { return h != nil && h.inner != nil && h.inner.HasQuotaProviderContext(ctx, provider) diff --git a/sdk/pluginhost/quota_test.go b/sdk/pluginhost/quota_test.go index fc28bd3f25..76014d807f 100644 --- a/sdk/pluginhost/quota_test.go +++ b/sdk/pluginhost/quota_test.go @@ -26,6 +26,60 @@ func (p quotaSDKFixture) ResetQuota(_ context.Context, req pluginapi.QuotaResetR return pluginapi.QuotaResetResponse{Success: true}, nil } +type quotaSelectionSDKFixture struct { + quotaSDKFixture + identifier string + broken bool +} + +func (p quotaSelectionSDKFixture) Identifier() string { return p.identifier } + +func (p quotaSelectionSDKFixture) DescribeQuota(context.Context, pluginapi.QuotaDescribeRequest) (pluginapi.QuotaDescribeResponse, error) { + if p.broken { + panic("broken quota description") + } + return pluginapi.QuotaDescribeResponse{SupportedProviders: []string{"foo"}}, nil +} + +func TestQuotaProviderPluginIDsSDK(t *testing.T) { + for _, healthyIdentifier := range []string{"foo", "healthy-quota"} { + t.Run(healthyIdentifier, func(t *testing.T) { + host := New() + host.inner.RegisterPluginForTest("a-broken", pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + QuotaProvider: quotaSelectionSDKFixture{identifier: "foo", broken: true}, + }}) + host.inner.RegisterPluginForTest("b-healthy", pluginapi.Plugin{Capabilities: pluginapi.Capabilities{ + QuotaProvider: quotaSelectionSDKFixture{identifier: healthyIdentifier}, + }}) + got := host.QuotaProviderPluginIDs(context.Background()) + if got["foo"] != "b-healthy" { + t.Fatalf("selected plugin = %q, want b-healthy", got["foo"]) + } + if _, exists := got["a-broken"]; exists { + t.Fatal("mapping includes a fused plugin's alias") + } + if host.HasQuotaProviderForPlugin("a-broken") { + t.Fatal("broken plugin must be unavailable through the public API") + } + for provider, pluginID := range got { + if !host.HasQuotaProviderForPlugin(pluginID) { + t.Fatalf("provider %q references unavailable plugin %q", provider, pluginID) + } + } + }) + } +} + +func TestQuotaProviderPluginIDsSDKWithoutHost(t *testing.T) { + for _, host := range []*Host{nil, {}, New()} { + for _, ctx := range []context.Context{nil, context.Background()} { + if got := host.QuotaProviderPluginIDs(ctx); len(got) != 0 { + t.Fatalf("expected no quota plugins, got %#v", got) + } + } + } +} + func TestQuotaSDKEmbedding(t *testing.T) { host := New() host.inner.RegisterPluginForTest("fixture-plugin", pluginapi.Plugin{Capabilities: pluginapi.Capabilities{QuotaProvider: quotaSDKFixture{t}}})