Wire backend capacity into engine runtime

internal/usecase/runner_test.go

@@ -10,7 +10,9 @@ import (
 	"reflect"
 	"regexp"
 	"strings"
+	"sync"
 	"testing"
+	"time"
 
 	"gitea.maximumdirect.net/eric/promptkit/internal/capacity"
 	"gitea.maximumdirect.net/eric/promptkit/internal/defaults"
@@ -39,6 +41,52 @@ type fakeBackendResolver struct {
 	backends map[string]domain.Backend
 }
 
+type staticPromptRepo struct {
+	def *domain.PromptDefinition
+}
+
+func (r staticPromptRepo) GetPromptDefinition(
+	context.Context,
+	string,
+	string,
+) (*domain.PromptDefinition, error) {
+	return r.def, nil
+}
+
+type staticExecutionProfileRepo struct {
+	profile *domain.ExecutionProfile
+}
+
+func (r staticExecutionProfileRepo) GetProfile(
+	context.Context,
+	string,
+) (*domain.ExecutionProfile, error) {
+	value := *r.profile
+	return &value, nil
+}
+
+type staticArtifactReader struct{}
+
+func (staticArtifactReader) Read(
+	context.Context,
+	domain.ArtifactRef,
+) (*domain.Artifact, error) {
+	return &domain.Artifact{Body: []byte("artifact"), Hash: hashString("artifact")}, nil
+}
+
+type staticRenderer struct{}
+
+func (staticRenderer) Render(
+	context.Context,
+	*domain.PromptDefinition,
+	map[string]*domain.Artifact,
+	map[string]string,
+) (*domain.RenderedPrompt, error) {
+	return &domain.RenderedPrompt{
+		Messages: []domain.RenderedMessage{{Role: "user", Content: "hello"}},
+	}, nil
+}
+
 func (f fakeBackendResolver) GetBackend(id string) (domain.Backend, error) {
 	value, ok := f.backends[id]
 	if !ok {
@@ -166,6 +214,81 @@ type fakeRunAdmitter struct {
 	releaseCalls int
 }
 
+type recordingRunAdmitter struct {
+	mu         sync.Mutex
+	next       RunAdmitter
+	backendIDs []string
+}
+
+func (a *recordingRunAdmitter) Admit(ctx context.Context, backendID string) (func(), error) {
+	a.mu.Lock()
+	a.backendIDs = append(a.backendIDs, backendID)
+	a.mu.Unlock()
+	return a.next.Admit(ctx, backendID)
+}
+
+func (a *recordingRunAdmitter) admittedBackendIDs() []string {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+	return append([]string(nil), a.backendIDs...)
+}
+
+type controlledRepairLLM struct {
+	mu          sync.Mutex
+	active      int
+	peak        int
+	calls       int
+	repairCalls int
+	events      chan controlledGeneration
+	backendIDs  []string
+}
+
+type controlledGeneration struct {
+	release chan struct{}
+}
+
+func (c *controlledRepairLLM) Generate(
+	ctx context.Context,
+	req domain.GenerateRequest,
+) (*domain.GenerateResponse, error) {
+	isRepair := len(req.Prompt.Messages) > 0 &&
+		strings.HasPrefix(req.Prompt.Messages[0].Content, "You repair invalid JSON")
+	c.mu.Lock()
+	c.calls++
+	c.active++
+	if c.active > c.peak {
+		c.peak = c.active
+	}
+	if isRepair {
+		c.repairCalls++
+	}
+	c.backendIDs = append(c.backendIDs, req.Target.BackendID)
+	c.mu.Unlock()
+	defer func() {
+		c.mu.Lock()
+		c.active--
+		c.mu.Unlock()
+	}()
+
+	event := controlledGeneration{release: make(chan struct{})}
+	c.events <- event
+	select {
+	case <-event.release:
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+	if isRepair {
+		return &domain.GenerateResponse{Content: `{}`}, nil
+	}
+	return &domain.GenerateResponse{Content: `{"broken":`}, nil
+}
+
+func (c *controlledRepairLLM) snapshot() (peak, calls, repairCalls int, backendIDs []string) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return c.peak, c.calls, c.repairCalls, append([]string(nil), c.backendIDs...)
+}
+
 func (f *fakeRunAdmitter) Admit(_ context.Context, backendID string) (func(), error) {
 	f.backendIDs = append(f.backendIDs, backendID)
 	if f.err != nil {
@@ -1957,6 +2080,97 @@ func TestRunnerRunStructuredRepairRemainsBoundedAndUsesEffectiveModelSettings(t
 	}
 }
 
+func TestRunnerSchedulesInitialAndRepairGenerationThroughOneBackendPool(t *testing.T) {
+	manager, err := capacity.NewManager(map[string]domain.BackendCapacityPolicy{
+		"custom": {
+			ConcurrencyLimit: 1,
+			QueueCapacity:    1,
+		},
+	})
+	if err != nil {
+		t.Fatalf("construct capacity manager: %v", err)
+	}
+	baseClient := &controlledRepairLLM{
+		events: make(chan controlledGeneration, 4),
+	}
+	scheduledClient := capacity.NewClient(manager, baseClient)
+	admitter := &recordingRunAdmitter{next: manager}
+	runner := NewRunnerWithRepairer(
+		staticPromptRepo{def: promptDef(domain.FormatJSON, domain.ValidationJSON, 1)},
+		staticExecutionProfileRepo{profile: &domain.ExecutionProfile{
+			ID: "exec", BackendID: "custom", Model: "model",
+		}},
+		fakeBackendResolver{backends: map[string]domain.Backend{
+			"custom": {ID: "custom", Endpoint: "http://backend.example/v1"},
+		}},
+		staticArtifactReader{},
+		staticRenderer{},
+		scheduledClient,
+		validate.NewStandardValidator("."),
+		NewDefaultOutputRepairer(scheduledClient),
+		admitter,
+	)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	type runOutcome struct {
+		result *domain.RunResult
+		err    error
+	}
+	outcomes := make(chan runOutcome, 2)
+	request := domain.RunRequest{
+		PromptID:  "p",
+		ProfileID: "exec",
+		Inputs:    singleInputRef(),
+	}
+	for i := 0; i < 2; i++ {
+		go func() {
+			result, runErr := runner.Run(ctx, request)
+			outcomes <- runOutcome{result: result, err: runErr}
+		}()
+	}
+
+	for i := 0; i < 4; i++ {
+		select {
+		case event := <-baseClient.events:
+			close(event.release)
+		case <-ctx.Done():
+			t.Fatalf("timed out waiting for generation %d: %v", i+1, ctx.Err())
+		}
+	}
+	for i := 0; i < 2; i++ {
+		select {
+		case outcome := <-outcomes:
+			if outcome.err != nil || outcome.result == nil {
+				t.Fatalf("run outcome=(%+v, %v), want success", outcome.result, outcome.err)
+			}
+			if outcome.result.Validation.Status != domain.ValidationPassed ||
+				outcome.result.Validation.RepairAttempts != 1 {
+				t.Fatalf("unexpected repaired result: %+v", outcome.result.Validation)
+			}
+		case <-ctx.Done():
+			t.Fatalf("timed out waiting for repaired run: %v", ctx.Err())
+		}
+	}
+
+	admitted := admitter.admittedBackendIDs()
+	if !reflect.DeepEqual(admitted, []string{"custom", "custom"}) {
+		t.Fatalf("admission backend IDs=%#v, want one admission per run", admitted)
+	}
+	peak, calls, repairCalls, backendIDs := baseClient.snapshot()
+	if peak != 1 || calls != 4 || repairCalls != 2 {
+		t.Fatalf(
+			"generation observations=(peak=%d calls=%d repairs=%d), want (1, 4, 2)",
+			peak,
+			calls,
+			repairCalls,
+		)
+	}
+	if !reflect.DeepEqual(backendIDs, []string{"custom", "custom", "custom", "custom"}) {
+		t.Fatalf("generation backend IDs=%#v, want custom for initial and repair calls", backendIDs)
+	}
+}
+
 func TestRunnerRunRepairCarriesEffectiveSessionID(t *testing.T) {
 	llmClient := &fakeLLM{resp: &domain.GenerateResponse{Content: `{"broken":`}}
 	runner := NewRunnerWithRepairer(