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

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)},