base-layer/promptkit
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

Retire the legacy capacity client wrapper

Browse Source
This commit is contained in:
Eric Rakestraw
2026-11-05 14:35:10 +00:00
parent cee9bbc7b2
commit 6f42e1e007
5 changed files with 28 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
docs/internal/capacity.md
View File

@@ -68,9 +68,8 @@ Preparation and validation never hold an active permit.
The manager's `Acquire` operation owns that active-permit acquisition and
returns an idempotent release function. The coordinator uses it around
individual attempts. The earlier one-shot wrapper still delegates to the same
operation for focused internal callers and tests, but engine assembly no
longer uses it. There is one queue and one active limit per limited backend pool.
individual attempts. There is one queue and one active limit per limited backend
pool.
When all active permits are occupied, calls join a mutex-protected FIFO waiter
list. Releasing a permit transfers it directly to the oldest remaining waiter
@@ -121,7 +120,7 @@ validation, bounded admission, idempotent release, context handling, and
unlimited admission. The
[client tests](../../internal/capacity/client_test.go) own peak enforcement,
FIFO transfer, canceled-waiter removal, grant/cancel races, independent pools,
unlimited calls, passthrough behavior, and panic release.
unlimited calls, and panic release using test-only one-shot calls.
The [runner tests](../../internal/usecase/runner_test.go) own ordinary early
admission, lease lifetime, failure release, and shared initial/repair

5
docs/roadmap/implementation.md
View File

@@ -826,6 +826,11 @@ focused public tests pass on Go 1.27.3; migration guidance notes hash/cache impa
exercise the real scheduling path. No production scheduling policy or public
API changes; no-repair runner construction and other deliberate fallbacks remain.
Completion: initial and repair scheduling use the real injected coordinator;
the production capacity wrapper was removed. Capacity tests retain one-shot
test-only scaffolding for manager scenarios. Race-enabled capacity, generation,
and use-case tests pass on Go 1.27.3, and internal capacity guidance is updated.
## Stage 33: Reconcile Documentation And Validate The End State
**Catalog:** all entries; complete the sprint, not a release publication.

40
internal/capacity/client.go
View File

@@ -1,40 +0,0 @@
package capacity
import (
"context"
"go.maximumdirect.net/promptkit/internal/domain"
"go.maximumdirect.net/promptkit/internal/llm"
)
type client struct {
manager *Manager
next llm.Client
}
// NewClient wraps next with configured active-generation limits. A nil manager
// leaves next unchanged.
func NewClient(manager *Manager, next llm.Client) llm.Client {
if manager == nil {
return next
}
return &client{
manager: manager,
next: next,
}
}
func (c *client) Generate(
ctx context.Context,
req domain.GenerateRequest,
) (*domain.GenerateResponse, error) {
if c.manager.getPool(req.Target.BackendID) == nil {
return c.next.Generate(ctx, req)
}
release, err := c.manager.Acquire(ctx, req.Target.BackendID)
if err != nil {
return nil, err
}
defer release()
return c.next.Generate(ctx, req)
}

82
internal/capacity/client_test.go
View File

@@ -3,7 +3,6 @@ package capacity
import (
"context"
"errors"
"reflect"
"runtime"
"sync"
"sync/atomic"
@@ -24,6 +23,18 @@ type clientFunc func(
domain.GenerateRequest,
) (*domain.GenerateResponse, error)
// newTestClient exercises one acquisition around a controlled call.
func newTestClient(manager *Manager, next llm.Client) llm.Client {
return clientFunc(func(ctx context.Context, req domain.GenerateRequest) (*domain.GenerateResponse, error) {
release, err := manager.Acquire(ctx, req.Target.BackendID)
if err != nil {
return nil, err
}
defer release()
return next.Generate(ctx, req)
})
}
func (f clientFunc) Generate(
ctx context.Context,
req domain.GenerateRequest,
@@ -167,7 +178,7 @@ func TestClientLimitsPeakConcurrencyAndServesWaitersFIFO(t *testing.T) {
"second": secondRelease,
"third": thirdRelease,
})
client := NewClient(manager, next)
client := newTestClient(manager, next)
first := generateAsync(client, context.Background(), "limited", "first")
if got := receiveStarted(t, next.started); got != "first" {
@@ -212,7 +223,7 @@ func TestClientPeakConcurrencyDoesNotExceedConfiguredLimit(t *testing.T) {
releases[string(rune('a'+i))] = gate
}
next := newBlockingClient(releases)
client := NewClient(manager, next)
client := newTestClient(manager, next)
results := make([]<-chan generateResult, 0, len(releases))
for id := range releases {
@@ -257,7 +268,7 @@ func TestClientRemovesCanceledWaiters(t *testing.T) {
"three": make(chan struct{}),
}
next := newBlockingClient(releases)
client := NewClient(manager, next)
client := newTestClient(manager, next)
holder := generateAsync(client, context.Background(), "limited", "holder")
if got := receiveStarted(t, next.started); got != "holder" {
@@ -321,7 +332,7 @@ func TestClientGrantCancellationRaceDoesNotLeakPermit(t *testing.T) {
}
return &domain.GenerateResponse{Content: req.Prompt.SessionID}, nil
})
client := NewClient(manager, next)
client := newTestClient(manager, next)
holder := generateAsync(client, context.Background(), "limited", "holder")
waitForActiveCount(t, manager, "limited", 1)
@@ -394,7 +405,7 @@ func TestClientUsesIndependentPoolsAndUnlimitedFastPaths(t *testing.T) {
"alpha": alphaRelease,
"beta": betaRelease,
})
client := NewClient(manager, next)
client := newTestClient(manager, next)
alpha := generateAsync(client, context.Background(), "alpha", "alpha")
beta := generateAsync(client, context.Background(), "beta", "beta")
@@ -423,63 +434,6 @@ func TestClientUsesIndependentPoolsAndUnlimitedFastPaths(t *testing.T) {
t.Fatalf("unlimited backend %q response=(%#v, %v)", backendID, response, err)
}
}
if got := NewClient(nil, next); got != next {
t.Fatal("nil manager did not return the wrapped client unchanged")
}
}
func TestClientPreservesRequestsResponsesAndErrors(t *testing.T) {
manager := newTestManager(t, map[string]domain.BackendCapacityPolicy{
"limited": {ConcurrencyLimit: 1},
})
request := domain.GenerateRequest{
Prompt: domain.RenderedPrompt{
SessionID: "session",
Messages: []domain.RenderedMessage{
{Role: "user", Content: "content"},
},
},
Target: domain.ExecutionTarget{
BackendID: "limited",
Model: "model",
ExtraParams: map[string]any{"key": "value"},
},
}
response := &domain.GenerateResponse{
Content: "output",
Usage: domain.TokenUsage{TotalTokens: 7},
}
collaboratorErr := errors.New("collaborator failure")
tests := []struct {
name string
response *domain.GenerateResponse
err error
}{
{name: "successful response", response: response},
{name: "nil response"},
{name: "collaborator error", response: response, err: collaboratorErr},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
var captured domain.GenerateRequest
next := clientFunc(func(
_ context.Context,
req domain.GenerateRequest,
) (*domain.GenerateResponse, error) {
captured = req
return tc.response, tc.err
})
gotResponse, gotErr := NewClient(manager, next).Generate(context.Background(), request)
if !reflect.DeepEqual(captured, request) {
t.Fatalf("request changed: %#v", captured)
}
if gotResponse != tc.response || gotErr != tc.err {
t.Fatalf("response=(%p, %v), want (%p, %v)",
gotResponse, gotErr, tc.response, tc.err)
}
})
}
}
func TestClientReleasesPermitDuringPanicUnwinding(t *testing.T) {
@@ -496,7 +450,7 @@ func TestClientReleasesPermitDuringPanicUnwinding(t *testing.T) {
}
return &domain.GenerateResponse{Content: "recovered"}, nil
})
client := NewClient(manager, next)
client := newTestClient(manager, next)
request := domain.GenerateRequest{
Target: domain.ExecutionTarget{BackendID: "limited"},
}

3
internal/usecase/runner_test.go
View File

@@ -18,6 +18,7 @@ import (
"go.maximumdirect.net/promptkit/internal/capacity"
"go.maximumdirect.net/promptkit/internal/defaults"
"go.maximumdirect.net/promptkit/internal/domain"
"go.maximumdirect.net/promptkit/internal/generation"
"go.maximumdirect.net/promptkit/internal/llm"
"go.maximumdirect.net/promptkit/internal/prompt"
"go.maximumdirect.net/promptkit/internal/validate"
@@ -2398,7 +2399,7 @@ func TestRunnerSchedulesInitialAndRepairGenerationThroughOneBackendPool(t *testi
baseClient := &controlledRepairLLM{
events: make(chan controlledGeneration, 4),
}
scheduledClient := capacity.NewClient(manager, baseClient)
scheduledClient := generation.NewInjected(manager, baseClient)
admitter := &recordingRunAdmitter{next: manager}
runner := NewRunnerWithRepairer(
staticPromptRepo{def: promptDef(domain.FormatJSON, domain.ValidationJSON, 1)},