Review LLM concurrency refactor

internal/framework/proposal_generation/generate_test.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 	"reflect"
+	"runtime"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -54,6 +55,8 @@ func (s *countingScheduler) Run(ctx context.Context, fn func(context.Context) er
 type sleepingStructuredClient struct {
 	inFlight    int32
 	maxInFlight int32
+	entered     chan struct{}
+	release     chan struct{}
 }
 
 func (c *sleepingStructuredClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
@@ -65,8 +68,11 @@ func (c *sleepingStructuredClient) CompleteStructured(ctx context.Context, req c
 			break
 		}
 	}
+	if c.entered != nil {
+		c.entered <- struct{}{}
+	}
 	select {
-	case <-time.After(20 * time.Millisecond):
+	case <-c.release:
 	case <-ctx.Done():
 		atomic.AddInt32(&c.inFlight, -1)
 		return contracts.StructuredCompletionResponse{}, ctx.Err()
@@ -309,7 +315,10 @@ func TestGenerateCandidatesRespectsSchedulerConcurrency(t *testing.T) {
 	if err != nil {
 		t.Fatalf("NewScheduler: %v", err)
 	}
-	client := &sleepingStructuredClient{}
+	client := &sleepingStructuredClient{
+		entered: make(chan struct{}, 16),
+		release: make(chan struct{}),
+	}
 	baseReq := defaultRequest(t)
 	baseReq.LLMClient = client
 	baseReq.Scheduler = scheduler
@@ -326,6 +335,9 @@ func TestGenerateCandidatesRespectsSchedulerConcurrency(t *testing.T) {
 			}
 		}(i)
 	}
+
+	waitForEntries(t, client.entered, 2)
+	close(client.release)
 	wg.Wait()
 
 	if got := atomic.LoadInt32(&client.maxInFlight); got > 2 {
@@ -335,3 +347,20 @@ func TestGenerateCandidatesRespectsSchedulerConcurrency(t *testing.T) {
 		t.Fatalf("expected observed concurrency of at least 2, got %d", got)
 	}
 }
+
+func waitForEntries(t *testing.T, entered <-chan struct{}, want int) {
+	t.Helper()
+	deadline := time.Now().Add(300 * time.Millisecond)
+	got := 0
+	for got < want && time.Now().Before(deadline) {
+		select {
+		case <-entered:
+			got++
+		default:
+			runtime.Gosched()
+		}
+	}
+	if got < want {
+		t.Fatalf("timed out waiting for %d entered calls (got %d)", want, got)
+	}
+}