Add explicit LLM concurrency controls

internal/framework/llm/scheduler_test.go

@@ -3,12 +3,64 @@ package llm
 import (
 	"context"
 	"errors"
+	"reflect"
 	"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)
+			}
+		}()
+		time.Sleep(10 * time.Millisecond)
+	}
+
+	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 {
@@ -79,7 +131,6 @@ func TestSchedulerReleasesPermitOnError(t *testing.T) {
 		t.Fatalf("expected %v, got %v", expectedErr, err)
 	}
 
-	// Must be able to run again after an error, proving permit release.
 	if err := s.Run(context.Background(), func(ctx context.Context) error {
 		_ = ctx
 		return nil
@@ -88,7 +139,7 @@ func TestSchedulerReleasesPermitOnError(t *testing.T) {
 	}
 }
 
-func TestSchedulerRespectsContextCancellation(t *testing.T) {
+func TestSchedulerContextCancellationWhileQueued(t *testing.T) {
 	s, err := NewScheduler(1)
 	if err != nil {
 		t.Fatalf("NewScheduler: %v", err)
@@ -111,3 +162,27 @@ func TestSchedulerRespectsContextCancellation(t *testing.T) {
 		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)
+	}
+}