Connect processing to shared PromptKit backend capacity

internal/framework/llm/adapter_capacity_test.go

@@ -0,0 +1,214 @@
+package llm
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+	"time"
+
+	ck "gitea.maximumdirect.net/eric/ConfigKit"
+	"gitea.maximumdirect.net/eric/audita/internal/core/config"
+	"gitea.maximumdirect.net/eric/audita/internal/framework/responseschema"
+	"gitea.maximumdirect.net/eric/audita/internal/framework/structuredoutput"
+	"gitea.maximumdirect.net/eric/audita/internal/prompts"
+	pk "gitea.maximumdirect.net/eric/promptkit"
+)
+
+func TestRolesShareBackendCapacityAndReleaseAfterCancellation(t *testing.T) {
+	entered := make(chan struct{}, 3)
+	release := make(chan struct{})
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		entered <- struct{}{}
+		<-release
+		content, _ := json.Marshal(`{"corrections":[]}`)
+		_, _ = fmt.Fprintf(w, `{"choices":[{"message":{"content":%s}}]}`, content)
+	}))
+	defer server.Close()
+	dir := t.TempDir()
+	if err := os.Mkdir(filepath.Join(dir, "profiles"), 0700); err != nil {
+		t.Fatal(err)
+	}
+	main := filepath.Join(dir, "config.yml")
+	body := fmt.Sprintf("version: 2\nbackends:\n  shared:\n    endpoint: %s/v1\n    concurrency: 1\n    queue_capacity: 0\nllm:\n  proposal_profile: proposal\n  validation_profile: validation\n  transient_retries: 0\n  repair_attempts: 0\n", server.URL)
+	if err := os.WriteFile(main, []byte(body), 0600); err != nil {
+		t.Fatal(err)
+	}
+	for _, name := range []string{"proposal", "validation"} {
+		profile := "id: " + name + "\nbackend: shared\nmodel: local\n"
+		if err := os.WriteFile(filepath.Join(dir, "profiles", name+".yaml"), []byte(profile), 0600); err != nil {
+			t.Fatal(err)
+		}
+	}
+	loaded, err := config.Load(ck.File(main), "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	runtime, err := NewRuntime(context.Background(), loaded.Config, main, server.Client())
+	if err != nil {
+		t.Fatal(err)
+	}
+	proposal, err := runtime.NewRoleClient("proposal", CallHooks{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	validation, err := runtime.NewRoleClient("validation", CallHooks{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	request := structuredoutput.Request{PromptID: prompts.PromptIDModuleGlossaryProposal, PromptVersion: prompts.VersionV1, Variables: map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}, Schema: responseschema.CorrectionSetKey}
+	firstDone := make(chan error, 1)
+	go func() {
+		var output map[string]any
+		_, callErr := proposal.CompleteStructured(context.Background(), request, &output)
+		firstDone <- callErr
+	}()
+	select {
+	case <-entered:
+	case <-time.After(2 * time.Second):
+		t.Fatal("first call did not enter provider")
+	}
+	var output map[string]any
+	_, err = validation.CompleteStructured(context.Background(), request, &output)
+	if !errors.Is(err, pk.ErrCapacityExceeded) {
+		t.Fatalf("other role did not share capacity: %v", err)
+	}
+	select {
+	case <-entered:
+		t.Fatal("second role reached provider despite full backend")
+	default:
+	}
+	close(release)
+	if err := <-firstDone; err != nil {
+		t.Fatal(err)
+	}
+	_, err = validation.CompleteStructured(context.Background(), request, &output)
+	if err != nil {
+		t.Fatalf("capacity was not released: %v", err)
+	}
+	for len(entered) > 0 {
+		<-entered
+	}
+	release = make(chan struct{})
+	if err := os.WriteFile(main, []byte(strings.Replace(body, "queue_capacity: 0", "queue_capacity: 1", 1)), 0600); err != nil {
+		t.Fatal(err)
+	}
+	loaded, err = config.Load(ck.File(main), "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	runtime, err = NewRuntime(context.Background(), loaded.Config, main, server.Client())
+	if err != nil {
+		t.Fatal(err)
+	}
+	proposal, err = runtime.NewRoleClient("proposal", CallHooks{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	validation, err = runtime.NewRoleClient("validation", CallHooks{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	go func() {
+		var out map[string]any
+		_, callErr := proposal.CompleteStructured(context.Background(), request, &out)
+		firstDone <- callErr
+	}()
+	select {
+	case <-entered:
+	case <-time.After(2 * time.Second):
+		t.Fatal("first queued-capacity call did not dispatch")
+	}
+	queuedCtx, cancelQueued := context.WithTimeout(context.Background(), 30*time.Millisecond)
+	defer cancelQueued()
+	_, err = validation.CompleteStructured(queuedCtx, request, &output)
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("queued call did not honor deadline: %v", err)
+	}
+	close(release)
+	if err := <-firstDone; err != nil {
+		t.Fatal(err)
+	}
+	_, err = validation.CompleteStructured(context.Background(), request, &output)
+	if err != nil {
+		t.Fatalf("queued cancellation leaked capacity: %v", err)
+	}
+}
+
+func TestDistinctBackendsHaveIndependentCapacity(t *testing.T) {
+	entered := make(chan struct{}, 2)
+	release := make(chan struct{})
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		entered <- struct{}{}
+		<-release
+		schema := map[string]any{}
+		_ = json.NewDecoder(r.Body).Decode(&schema)
+		content, _ := json.Marshal(`{"corrections":[]}`)
+		_, _ = fmt.Fprintf(w, `{"choices":[{"message":{"content":%s}}]}`, content)
+	}))
+	defer server.Close()
+	dir := t.TempDir()
+	if err := os.Mkdir(filepath.Join(dir, "profiles"), 0700); err != nil {
+		t.Fatal(err)
+	}
+	main := filepath.Join(dir, "config.yml")
+	body := fmt.Sprintf("version: 2\nbackends:\n  first:\n    endpoint: %s/v1\n    concurrency: 1\n    queue_capacity: 0\n  second:\n    endpoint: %s/v1\n    concurrency: 1\n    queue_capacity: 0\nllm:\n  proposal_profile: proposal\n  validation_profile: validation\n  transient_retries: 0\n  repair_attempts: 0\n", server.URL, server.URL)
+	if err := os.WriteFile(main, []byte(body), 0600); err != nil {
+		t.Fatal(err)
+	}
+	for name, backend := range map[string]string{"proposal": "first", "validation": "second"} {
+		if err := os.WriteFile(filepath.Join(dir, "profiles", name+".yaml"), []byte("id: "+name+"\nbackend: "+backend+"\nmodel: local\n"), 0600); err != nil {
+			t.Fatal(err)
+		}
+	}
+	loaded, err := config.Load(ck.File(main), "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	runtime, err := NewRuntime(context.Background(), loaded.Config, main, server.Client())
+	if err != nil {
+		t.Fatal(err)
+	}
+	proposal, err := runtime.NewRoleClient("proposal", CallHooks{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	validation, err := runtime.NewRoleClient("validation", CallHooks{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	request := structuredoutput.Request{PromptID: prompts.PromptIDModuleGlossaryProposal, PromptVersion: prompts.VersionV1, Variables: map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}, Schema: responseschema.CorrectionSetKey}
+	done := make(chan error, 2)
+	go func() {
+		var out map[string]any
+		_, e := proposal.CompleteStructured(context.Background(), request, &out)
+		done <- e
+	}()
+	select {
+	case <-entered:
+	case <-time.After(2 * time.Second):
+		t.Fatal("first backend did not dispatch")
+	}
+	go func() {
+		var out map[string]any
+		_, e := validation.CompleteStructured(context.Background(), request, &out)
+		done <- e
+	}()
+	select {
+	case <-entered:
+	case <-time.After(2 * time.Second):
+		t.Fatal("second backend was blocked by first")
+	}
+	close(release)
+	for range 2 {
+		if err := <-done; err != nil {
+			t.Fatal(err)
+		}
+	}
+}

internal/framework/llm/adapter_test.go

@@ -11,6 +11,7 @@ import (
 	"path/filepath"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 
 	ck "gitea.maximumdirect.net/eric/ConfigKit"
@@ -26,17 +27,17 @@ type reply struct {
 	content string
 }
 
-func adapterFixture(t *testing.T, replies []reply, retries, repairs int, hook func(*http.Request)) (*Runtime, *int, string) {
+func adapterFixture(t *testing.T, replies []reply, retries, repairs int, hook func(*http.Request)) (*Runtime, *atomic.Int32, string) {
 	t.Helper()
-	count := 0
+	var count atomic.Int32
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		count++
+		call := int(count.Add(1))
 		if hook != nil {
 			hook(r)
 		}
 		step := replies[len(replies)-1]
-		if count <= len(replies) {
-			step = replies[count-1]
+		if call <= len(replies) {
+			step = replies[call-1]
 		}
 		if step.status == 0 {
 			connection, _, hijackErr := w.(http.Hijacker).Hijack()
@@ -128,8 +129,8 @@ func TestPreparedAdapterRepairsAndClassifies(t *testing.T) {
 				}
 				var out map[string]any
 				response, err := client.CompleteStructured(context.Background(), adapterRequest(schema), &out)
-				if *count != tc.attempts {
-					t.Fatalf("HTTP attempts=%d, want %d", *count, tc.attempts)
+				if int(count.Load()) != tc.attempts {
+					t.Fatalf("HTTP attempts=%d, want %d", count.Load(), tc.attempts)
 				}
 				if tc.malformed != structuredoutput.IsMalformedError(err) {
 					t.Fatalf("malformed classification: %v", err)
@@ -188,7 +189,7 @@ func TestPreparedAdapterMetadataAndTargetGuard(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	if !requestSeen || *count != 1 || response.Prepared == nil || len(response.Prepared.Messages) == 0 {
+	if !requestSeen || count.Load() != 1 || response.Prepared == nil || len(response.Prepared.Messages) == 0 {
 		t.Fatal("prepared call missing")
 	}
 	if len(posted.Messages) != len(response.Prepared.Messages) {
@@ -222,8 +223,8 @@ func TestPreparedAdapterRejectsTargetDriftAndPreparationErrors(t *testing.T) {
 	req := adapterRequest(responseschema.CorrectionSetKey)
 	req.PromptID = "missing.prompt"
 	var out map[string]any
-	if _, err := client.CompleteStructured(context.Background(), req, &out); err == nil || *count != 0 {
-		t.Fatalf("preparation failure dispatched: %v, %d", err, *count)
+	if _, err := client.CompleteStructured(context.Background(), req, &out); err == nil || count.Load() != 0 {
+		t.Fatalf("preparation failure dispatched: %v, %d", err, count.Load())
 	}
 	// PrepareExecution reloads the profile, so a changed target must be blocked.
 	if err := os.WriteFile(filepath.Join(dir, "profiles", "base.yaml"), []byte("id: base\nbackend: offline\nendpoint: http://127.0.0.1:1/v1\nmodel: full-model\n"), 0600); err != nil {
@@ -231,8 +232,8 @@ func TestPreparedAdapterRejectsTargetDriftAndPreparationErrors(t *testing.T) {
 	}
 	req = adapterRequest(responseschema.CorrectionSetKey)
 	_, err = client.CompleteStructured(context.Background(), req, &out)
-	if err == nil || !strings.Contains(err.Error(), "target changed") || *count != 0 {
-		t.Fatalf("target drift dispatched: %v, %d", err, *count)
+	if err == nil || !strings.Contains(err.Error(), "target changed") || count.Load() != 0 {
+		t.Fatalf("target drift dispatched: %v, %d", err, count.Load())
 	}
 }
 
@@ -251,7 +252,7 @@ func TestPreparedAdapterCancellationStopsRetryWait(t *testing.T) {
 	}
 	var out map[string]any
 	_, err = client.CompleteStructured(ctx, adapterRequest(responseschema.CorrectionSetKey), &out)
-	if !errors.Is(err, context.Canceled) || *count != 1 {
-		t.Fatalf("cancellation: %v, HTTP calls=%d", err, *count)
+	if !errors.Is(err, context.Canceled) || count.Load() != 1 {
+		t.Fatalf("cancellation: %v, HTTP calls=%d", err, count.Load())
 	}
 }

internal/framework/llm/scheduler.go

@@ -1,117 +0,0 @@
-package llm
-
-import (
-	"context"
-	"fmt"
-	"sync"
-)
-
-// Scheduler bounds concurrent backend LLM calls.
-type Scheduler struct {
-	maxConcurrency int
-	mu             sync.Mutex
-	inFlight       int
-	queue          []*waiter
-}
-
-type waiter struct {
-	ready   chan struct{}
-	queued  bool
-	granted bool
-}
-
-// NewScheduler creates a scheduler with a fixed concurrency limit.
-func NewScheduler(maxConcurrency int) (*Scheduler, error) {
-	if maxConcurrency <= 0 {
-		return nil, fmt.Errorf("max concurrency must be greater than zero")
-	}
-	return &Scheduler{
-		maxConcurrency: maxConcurrency,
-	}, nil
-}
-
-// Acquire blocks until a permit is available or the context is canceled.
-// The returned release function is safe to call multiple times.
-func (s *Scheduler) Acquire(ctx context.Context) (func(), error) {
-	if err := ctx.Err(); err != nil {
-		return nil, err
-	}
-
-	s.mu.Lock()
-	if s.inFlight < s.maxConcurrency && len(s.queue) == 0 {
-		s.inFlight++
-		s.mu.Unlock()
-		return s.releaseFunc(), nil
-	}
-	w := &waiter{
-		ready:  make(chan struct{}),
-		queued: true,
-	}
-	s.queue = append(s.queue, w)
-	s.mu.Unlock()
-
-	select {
-	case <-w.ready:
-		return s.releaseFunc(), nil
-	case <-ctx.Done():
-		s.mu.Lock()
-		if w.queued {
-			s.removeQueuedWaiterLocked(w)
-			s.mu.Unlock()
-			return nil, ctx.Err()
-		}
-		if w.granted {
-			// Permit was granted concurrently with cancellation; release it to
-			// avoid leaks and unblock queued callers.
-			s.inFlight--
-			s.grantQueuedLocked()
-		}
-		s.mu.Unlock()
-		return nil, ctx.Err()
-	}
-}
-
-func (s *Scheduler) releaseFunc() func() {
-	var once sync.Once
-	return func() {
-		once.Do(func() {
-			s.mu.Lock()
-			if s.inFlight > 0 {
-				s.inFlight--
-				s.grantQueuedLocked()
-			}
-			s.mu.Unlock()
-		})
-	}
-}
-
-func (s *Scheduler) grantQueuedLocked() {
-	for s.inFlight < s.maxConcurrency && len(s.queue) > 0 {
-		w := s.queue[0]
-		s.queue = s.queue[1:]
-		w.queued = false
-		w.granted = true
-		s.inFlight++
-		close(w.ready)
-	}
-}
-
-func (s *Scheduler) removeQueuedWaiterLocked(target *waiter) {
-	for i, w := range s.queue {
-		if w == target {
-			w.queued = false
-			s.queue = append(s.queue[:i], s.queue[i+1:]...)
-			return
-		}
-	}
-}
-
-// Run acquires a permit, executes fn, and releases the permit.
-func (s *Scheduler) Run(ctx context.Context, fn func(context.Context) error) error {
-	release, err := s.Acquire(ctx)
-	if err != nil {
-		return err
-	}
-	defer release()
-	return fn(ctx)
-}

internal/framework/llm/scheduler_test.go

@@ -1,223 +0,0 @@
-package llm
-
-import (
-	"context"
-	"errors"
-	"reflect"
-	"runtime"
-	"sync"
-	"sync/atomic"
-	"testing"
-	"time"
-)
-
-func TestSchedulerFIFOOrdering(t *testing.T) {
-	s, err := NewScheduler(1)
-	if err != nil {
-		t.Fatalf("NewScheduler: %v", err)
-	}
-
-	firstRelease, err := s.Acquire(context.Background())
-	if err != nil {
-		t.Fatalf("Acquire(first): %v", err)
-	}
-
-	gotOrder := make(chan int, 3)
-	waitChans := []chan struct{}{
-		make(chan struct{}),
-		make(chan struct{}),
-		make(chan struct{}),
-	}
-	var wg sync.WaitGroup
-	for i := 0; i < 3; i++ {
-		id := i
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			runErr := s.Run(context.Background(), func(context.Context) error {
-				gotOrder <- id
-				<-waitChans[id]
-				return nil
-			})
-			if runErr != nil {
-				t.Errorf("Run[%d] error: %v", id, runErr)
-			}
-		}()
-		waitForQueueDepth(t, s, i+1)
-	}
-
-	firstRelease()
-
-	order := make([]int, 0, 3)
-	for i := 0; i < 3; i++ {
-		id := <-gotOrder
-		order = append(order, id)
-		close(waitChans[id])
-	}
-
-	wg.Wait()
-
-	if !reflect.DeepEqual(order, []int{0, 1, 2}) {
-		t.Fatalf("expected FIFO order [0 1 2], got %v", order)
-	}
-}
-
-func TestSchedulerEnforcesMaxConcurrency(t *testing.T) {
-	s, err := NewScheduler(2)
-	if err != nil {
-		t.Fatalf("NewScheduler: %v", err)
-	}
-
-	var inFlight int32
-	var maxInFlight int32
-	release := make(chan struct{})
-	var wg sync.WaitGroup
-
-	for i := 0; i < 12; i++ {
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			err := s.Run(context.Background(), func(ctx context.Context) error {
-				_ = ctx
-				current := atomic.AddInt32(&inFlight, 1)
-				for {
-					seen := atomic.LoadInt32(&maxInFlight)
-					if current <= seen || atomic.CompareAndSwapInt32(&maxInFlight, seen, current) {
-						break
-					}
-				}
-				<-release
-				atomic.AddInt32(&inFlight, -1)
-				return nil
-			})
-			if err != nil {
-				t.Errorf("Run error: %v", err)
-			}
-		}()
-	}
-
-	waitForMinInFlight(t, &maxInFlight, 2)
-	close(release)
-
-	wg.Wait()
-	if got := atomic.LoadInt32(&maxInFlight); got > 2 {
-		t.Fatalf("expected max in-flight <= 2, got %d", got)
-	}
-}
-
-func TestSchedulerReleasesPermitOnSuccess(t *testing.T) {
-	s, err := NewScheduler(1)
-	if err != nil {
-		t.Fatalf("NewScheduler: %v", err)
-	}
-
-	for i := 0; i < 3; i++ {
-		if err := s.Run(context.Background(), func(ctx context.Context) error {
-			_ = ctx
-			return nil
-		}); err != nil {
-			t.Fatalf("Run[%d] error: %v", i, err)
-		}
-	}
-}
-
-func TestSchedulerReleasesPermitOnError(t *testing.T) {
-	s, err := NewScheduler(1)
-	if err != nil {
-		t.Fatalf("NewScheduler: %v", err)
-	}
-
-	expectedErr := errors.New("boom")
-	err = s.Run(context.Background(), func(ctx context.Context) error {
-		_ = ctx
-		return expectedErr
-	})
-	if !errors.Is(err, expectedErr) {
-		t.Fatalf("expected %v, got %v", expectedErr, err)
-	}
-
-	if err := s.Run(context.Background(), func(ctx context.Context) error {
-		_ = ctx
-		return nil
-	}); err != nil {
-		t.Fatalf("second Run error: %v", err)
-	}
-}
-
-func TestSchedulerContextCancellationWhileQueued(t *testing.T) {
-	s, err := NewScheduler(1)
-	if err != nil {
-		t.Fatalf("NewScheduler: %v", err)
-	}
-
-	release, err := s.Acquire(context.Background())
-	if err != nil {
-		t.Fatalf("Acquire: %v", err)
-	}
-	defer release()
-
-	ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond)
-	defer cancel()
-
-	_, err = s.Acquire(ctx)
-	if err == nil {
-		t.Fatalf("expected context cancellation error")
-	}
-	if !errors.Is(err, context.DeadlineExceeded) {
-		t.Fatalf("expected deadline exceeded, got %v", err)
-	}
-}
-
-func TestSchedulerNoPermitLeakAfterQueuedCancellation(t *testing.T) {
-	s, err := NewScheduler(1)
-	if err != nil {
-		t.Fatalf("NewScheduler: %v", err)
-	}
-
-	firstRelease, err := s.Acquire(context.Background())
-	if err != nil {
-		t.Fatalf("Acquire(first): %v", err)
-	}
-
-	ctx, cancel := context.WithCancel(context.Background())
-	cancel()
-	if _, err := s.Acquire(ctx); !errors.Is(err, context.Canceled) {
-		t.Fatalf("expected context canceled, got %v", err)
-	}
-
-	firstRelease()
-
-	if err := s.Run(context.Background(), func(context.Context) error { return nil }); err != nil {
-		t.Fatalf("expected scheduler to accept new work after cancellation, got %v", err)
-	}
-}
-
-func waitForQueueDepth(t *testing.T, s *Scheduler, want int) {
-	t.Helper()
-	deadline := time.Now().Add(250 * time.Millisecond)
-	for time.Now().Before(deadline) {
-		s.mu.Lock()
-		depth := len(s.queue)
-		s.mu.Unlock()
-		if depth >= want {
-			return
-		}
-		runtime.Gosched()
-	}
-	s.mu.Lock()
-	depth := len(s.queue)
-	s.mu.Unlock()
-	t.Fatalf("timed out waiting for queue depth >= %d (got %d)", want, depth)
-}
-
-func waitForMinInFlight(t *testing.T, maxInFlight *int32, want int32) {
-	t.Helper()
-	deadline := time.Now().Add(250 * time.Millisecond)
-	for time.Now().Before(deadline) {
-		if atomic.LoadInt32(maxInFlight) >= want {
-			return
-		}
-		runtime.Gosched()
-	}
-	t.Fatalf("timed out waiting for max in-flight >= %d (got %d)", want, atomic.LoadInt32(maxInFlight))
-}