diff --git a/backends.go b/backends.go index 0ecc039..8ae129b 100644 --- a/backends.go +++ b/backends.go @@ -31,6 +31,17 @@ type Backend struct { // service_tier, reasoning_effort, or response_format. An empty map supplies // no defaults. NewEngine deeply copies the map. ExtraParams map[string]any + // ConcurrencyLimit is the maximum number of simultaneous model-generation + // calls allowed for this backend within one Engine. Zero leaves the backend + // unlimited. A negative value makes NewEngine fail with ErrInvalidConfig. + ConcurrencyLimit int + // QueueCapacity controls how many additional Run calls may be admitted + // beyond ConcurrencyLimit. Nil uses 1024 when ConcurrencyLimit is positive; + // a pointer uses its exact value, including zero. The pointed-to value must + // be non-negative, and QueueCapacity must be nil when ConcurrencyLimit is + // zero. Their sum must fit in an int. WithBackend copies the value and does + // not retain the pointer. + QueueCapacity *int } // WithBackend adds one Backend registration to the constructed Engine. @@ -43,12 +54,20 @@ type Backend struct { // mutated after construction. WithBackend does not install package-global // state. func WithBackend(backend Backend) Option { + queueCapacity := 0 + queueCapacitySet := backend.QueueCapacity != nil + if queueCapacitySet { + queueCapacity = *backend.QueueCapacity + } return optionFunc(func(options *engineOptions) error { options.backends = append(options.backends, domain.Backend{ - ID: backend.ID, - Endpoint: backend.Endpoint, - APIKeyEnv: backend.APIKeyEnv, - ExtraParams: backend.ExtraParams, + ID: backend.ID, + Endpoint: backend.Endpoint, + APIKeyEnv: backend.APIKeyEnv, + ExtraParams: backend.ExtraParams, + ConcurrencyLimit: backend.ConcurrencyLimit, + QueueCapacity: queueCapacity, + QueueCapacitySet: queueCapacitySet, }) return nil }) diff --git a/docs/roadmap/implementation.md b/docs/roadmap/implementation.md index 9da45b6..511ecfe 100644 --- a/docs/roadmap/implementation.md +++ b/docs/roadmap/implementation.md @@ -464,7 +464,7 @@ behavior already protects them. ## Stage 1 — Backend Policy And Public Configuration -**Status:** Pending. +**Status:** Complete. ### Goal diff --git a/internal/backend/registry.go b/internal/backend/registry.go index da23af1..1382920 100644 --- a/internal/backend/registry.go +++ b/internal/backend/registry.go @@ -22,6 +22,9 @@ const ( openRouterEndpoint = "https://openrouter.ai/api/v1" openRouterAPIKeyEnv = "OPENROUTER_API_KEY" + + openRouterConcurrencyLimit = 16 + defaultQueueCapacity = 1024 ) // ErrBackendNotFound identifies a registry lookup for an unknown backend ID. @@ -43,9 +46,10 @@ func NewRegistry(additions []domain.Backend) (*Registry, error) { definitions := make([]domain.Backend, 0, len(additions)+1) definitions = append(definitions, domain.Backend{ - ID: OpenRouterID, - Endpoint: openRouterEndpoint, - APIKeyEnv: openRouterAPIKeyEnv, + ID: OpenRouterID, + Endpoint: openRouterEndpoint, + APIKeyEnv: openRouterAPIKeyEnv, + ConcurrencyLimit: openRouterConcurrencyLimit, }) definitions = append(definitions, additions...) @@ -85,6 +89,25 @@ func (r *Registry) GetBackend(id string) (domain.Backend, error) { return definition, nil } +// CapacityPolicies returns a copy of the normalized policies for limited +// backends. +func (r *Registry) CapacityPolicies() map[string]domain.BackendCapacityPolicy { + policies := make(map[string]domain.BackendCapacityPolicy) + if r == nil { + return policies + } + for id, definition := range r.backends { + if definition.ConcurrencyLimit == 0 { + continue + } + policies[id] = domain.BackendCapacityPolicy{ + ConcurrencyLimit: definition.ConcurrencyLimit, + QueueCapacity: definition.QueueCapacity, + } + } + return policies +} + func normalizeBackend(definition domain.Backend) (domain.Backend, error) { definition.Endpoint = strings.TrimSpace(definition.Endpoint) if err := validateEndpoint(definition.Endpoint); err != nil { @@ -100,6 +123,40 @@ func normalizeBackend(definition domain.Backend) (domain.Backend, error) { ) } + if definition.ConcurrencyLimit < 0 { + return domain.Backend{}, fmt.Errorf( + "backend %q concurrency limit must not be negative", + definition.ID, + ) + } + if definition.QueueCapacity < 0 { + return domain.Backend{}, fmt.Errorf( + "backend %q queue capacity must not be negative", + definition.ID, + ) + } + if definition.ConcurrencyLimit == 0 { + if definition.QueueCapacitySet { + return domain.Backend{}, fmt.Errorf( + "backend %q queue capacity requires a positive concurrency limit", + definition.ID, + ) + } + definition.QueueCapacity = 0 + } else { + if !definition.QueueCapacitySet { + definition.QueueCapacity = defaultQueueCapacity + definition.QueueCapacitySet = true + } + maxInt := int(^uint(0) >> 1) + if definition.QueueCapacity > maxInt-definition.ConcurrencyLimit { + return domain.Backend{}, fmt.Errorf( + "backend %q total capacity overflows int", + definition.ID, + ) + } + } + keys := make([]string, 0, len(definition.ExtraParams)) for key := range definition.ExtraParams { keys = append(keys, key) diff --git a/internal/backend/registry_test.go b/internal/backend/registry_test.go index 0911733..bc7b1c0 100644 --- a/internal/backend/registry_test.go +++ b/internal/backend/registry_test.go @@ -24,9 +24,20 @@ func TestRegistryIncludesExactOpenRouterDefinition(t *testing.T) { if definition.ID != "openrouter" || definition.Endpoint != "https://openrouter.ai/api/v1" || definition.APIKeyEnv != "OPENROUTER_API_KEY" || + definition.ConcurrencyLimit != 16 || + definition.QueueCapacity != 1024 || + !definition.QueueCapacitySet || definition.ExtraParams != nil { t.Fatalf("unexpected OpenRouter definition: %#v", definition) } + policies := registry.CapacityPolicies() + if len(policies) != 1 || + policies["openrouter"] != (domain.BackendCapacityPolicy{ + ConcurrencyLimit: 16, + QueueCapacity: 1024, + }) { + t.Fatalf("unexpected OpenRouter capacity policies: %#v", policies) + } } func TestRegistryNormalizesUniqueAdditionsAndIsolatesMutations(t *testing.T) { @@ -37,10 +48,13 @@ func TestRegistryNormalizesUniqueAdditionsAndIsolatesMutations(t *testing.T) { } registry, err := backend.NewRegistry([]domain.Backend{ { - ID: " custom ", - Endpoint: " https://custom.example/openai/v1 ", - APIKeyEnv: " CUSTOM_API_KEY ", - ExtraParams: extraParams, + ID: " custom ", + Endpoint: " https://custom.example/openai/v1 ", + APIKeyEnv: " CUSTOM_API_KEY ", + ExtraParams: extraParams, + ConcurrencyLimit: 3, + QueueCapacity: 2, + QueueCapacitySet: true, }, { ID: "Custom", @@ -60,7 +74,10 @@ func TestRegistryNormalizesUniqueAdditionsAndIsolatesMutations(t *testing.T) { } if got.ID != "custom" || got.Endpoint != "https://custom.example/openai/v1" || - got.APIKeyEnv != "CUSTOM_API_KEY" { + got.APIKeyEnv != "CUSTOM_API_KEY" || + got.ConcurrencyLimit != 3 || + got.QueueCapacity != 2 || + !got.QueueCapacitySet { t.Fatalf("unexpected normalized definition: %#v", got) } if count, ok := got.ExtraParams["count"].(int64); !ok || count != 7 { @@ -90,6 +107,132 @@ func TestRegistryNormalizesUniqueAdditionsAndIsolatesMutations(t *testing.T) { if _, err := registry.GetBackend("Custom"); err != nil { t.Fatalf("backend IDs should be case-sensitive: %v", err) } + + policies := registry.CapacityPolicies() + if len(policies) != 2 { + t.Fatalf("unexpected capacity policy count: %#v", policies) + } + policies["custom"] = domain.BackendCapacityPolicy{} + delete(policies, backend.OpenRouterID) + againPolicies := registry.CapacityPolicies() + if againPolicies["custom"] != (domain.BackendCapacityPolicy{ + ConcurrencyLimit: 3, + QueueCapacity: 2, + }) { + t.Fatalf("capacity policy map mutated registry state: %#v", againPolicies) + } + if _, ok := againPolicies[backend.OpenRouterID]; !ok { + t.Fatalf("capacity policy deletion mutated registry state: %#v", againPolicies) + } +} + +func TestNewRegistryNormalizesCapacityPolicy(t *testing.T) { + maxInt := int(^uint(0) >> 1) + tests := []struct { + name string + definition domain.Backend + want domain.BackendCapacityPolicy + wantSet bool + wantError bool + }{ + { + name: "unlimited when omitted", + definition: domain.Backend{}, + }, + { + name: "default queue", + definition: domain.Backend{ + ConcurrencyLimit: 2, + }, + want: domain.BackendCapacityPolicy{ + ConcurrencyLimit: 2, + QueueCapacity: 1024, + }, + wantSet: true, + }, + { + name: "explicit zero queue", + definition: domain.Backend{ + ConcurrencyLimit: 2, + QueueCapacitySet: true, + }, + want: domain.BackendCapacityPolicy{ + ConcurrencyLimit: 2, + }, + wantSet: true, + }, + { + name: "negative concurrency limit", + definition: domain.Backend{ + ConcurrencyLimit: -1, + }, + wantError: true, + }, + { + name: "negative queue capacity", + definition: domain.Backend{ + ConcurrencyLimit: 1, + QueueCapacity: -1, + QueueCapacitySet: true, + }, + wantError: true, + }, + { + name: "queue without limit", + definition: domain.Backend{ + QueueCapacitySet: true, + }, + wantError: true, + }, + { + name: "total overflow", + definition: domain.Backend{ + ConcurrencyLimit: maxInt, + QueueCapacity: 1, + QueueCapacitySet: true, + }, + wantError: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + tc.definition.ID = "custom" + tc.definition.Endpoint = validEndpoint + registry, err := backend.NewRegistry([]domain.Backend{tc.definition}) + if tc.wantError { + if err == nil { + t.Fatal("expected invalid capacity policy error") + } + return + } + if err != nil { + t.Fatalf("construct registry: %v", err) + } + + definition, err := registry.GetBackend("custom") + if err != nil { + t.Fatalf("look up custom backend: %v", err) + } + if definition.ConcurrencyLimit != tc.want.ConcurrencyLimit || + definition.QueueCapacity != tc.want.QueueCapacity || + definition.QueueCapacitySet != tc.wantSet { + t.Fatalf("normalized capacity=(%d, %d, %t), want (%d, %d, %t)", + definition.ConcurrencyLimit, + definition.QueueCapacity, + definition.QueueCapacitySet, + tc.want.ConcurrencyLimit, + tc.want.QueueCapacity, + tc.wantSet, + ) + } + policies := registry.CapacityPolicies() + got, ok := policies["custom"] + if ok != tc.wantSet || got != tc.want { + t.Fatalf("capacity policy=(%#v, %t), want (%#v, %t)", got, ok, tc.want, tc.wantSet) + } + }) + } } func TestNewRegistryRejectsDuplicateIDs(t *testing.T) { diff --git a/internal/domain/domain.go b/internal/domain/domain.go index 7294525..cbc2b33 100644 --- a/internal/domain/domain.go +++ b/internal/domain/domain.go @@ -163,10 +163,20 @@ type PromptMessageTemplate struct { // Backend describes reusable OpenAI-compatible connection defaults. type Backend struct { - ID string - Endpoint string - APIKeyEnv string - ExtraParams map[string]any + ID string + Endpoint string + APIKeyEnv string + ExtraParams map[string]any + ConcurrencyLimit int + QueueCapacity int + QueueCapacitySet bool +} + +// BackendCapacityPolicy describes normalized run and generation capacity for +// one limited backend. +type BackendCapacityPolicy struct { + ConcurrencyLimit int + QueueCapacity int } // ExecutionProfile describes how and where to execute a model. diff --git a/public_contract_test.go b/public_contract_test.go index 8168c64..0055451 100644 --- a/public_contract_test.go +++ b/public_contract_test.go @@ -303,6 +303,28 @@ func TestBackendOptionsAccumulateAndRegistrationsAreEngineLocal(t *testing.T) { } } +func TestWithBackendCopiesQueueCapacity(t *testing.T) { + queueCapacity := 4 + option := promptkit.WithBackend(promptkit.Backend{ + ID: "custom", + Endpoint: "http://custom.example/v1", + ConcurrencyLimit: 1, + QueueCapacity: &queueCapacity, + }) + queueCapacity = -1 + + _, err := promptkit.NewEngine(promptkit.Config{}, + promptkit.WithPromptFS(contractPromptFS("prompt", "profile", "message"), "."), + option, + promptkit.WithProfiles(promptkit.Profile{ + ID: "profile", BackendID: "custom", Model: "model", + }), + ) + if err != nil { + t.Fatalf("construct engine after mutating queue pointer: %v", err) + } +} + func TestBackendRegistrationRejectsInvalidAndDuplicateDefinitions(t *testing.T) { cycle := map[string]any{} cycle["self"] = cycle