promptkit/capacity_contract_test.go

package promptkit_test

import (
	"context"
	"errors"
	"sync"
	"testing"
	"time"

	"gitea.maximumdirect.net/eric/promptkit"
)

func TestEngineLimitsInjectedClientConcurrency(t *testing.T) {
	release := make(chan struct{})
	client := newCapacityGateClient(release, 8)
	engine := newBackendCapacityEngine(t, client, 2, capacityInt(4), nil)

	results := make(chan capacityRunResult, 6)
	for i := 0; i < 6; i++ {
		go runCapacityRequest(engine, context.Background(), promptkit.RunRequest{
			PromptID: "prompt",
			Execution: &promptkit.ExecutionTargetOverride{
				Endpoint: "http://request.example/v1",
			},
		}, results)
	}

	first := awaitCapacityRequest(t, client.started)
	second := awaitCapacityRequest(t, client.started)
	if first.Target.BackendID != "limited" || second.Target.BackendID != "limited" {
		t.Fatalf("endpoint override changed backend pool: first=%q second=%q",
			first.Target.BackendID, second.Target.BackendID)
	}
	if active, peak, _ := client.snapshot(); active != 2 || peak != 2 {
		t.Fatalf("client concurrency before release=(active=%d peak=%d), want 2", active, peak)
	}

	close(release)
	for i := 0; i < 6; i++ {
		outcome := awaitCapacityRun(t, results)
		if outcome.err != nil || outcome.result == nil {
			t.Fatalf("run outcome=(%+v, %v), want success", outcome.result, outcome.err)
		}
	}
	if _, peak, calls := client.snapshot(); peak > 2 || calls != 6 {
		t.Fatalf("client observations=(peak=%d calls=%d), want peak <= 2 and 6 calls", peak, calls)
	}
}

func TestEngineRejectsRunBeforeCompletionWhenAdmissionIsFull(t *testing.T) {
	artifactRelease := make(chan struct{})
	reader := &capacityArtifactReader{
		entered: make(chan struct{}, 2),
		release: artifactRelease,
	}
	client := newCapacityGateClient(closedCapacityChannel(), 2)
	engine := newBackendCapacityEngine(t, client, 1, capacityInt(0), reader)
	firstResult := make(chan capacityRunResult, 1)
	go runCapacityRequest(engine, context.Background(), capacityInputRequest("http://first.example/v1"), firstResult)

	awaitCapacitySignal(t, reader.entered, "first artifact read")

	canceledContext, cancel := context.WithCancel(context.Background())
	cancel()
	result, err := engine.Run(canceledContext, capacityInputRequest("http://canceled.example/v1"))
	if result != nil || !errors.Is(err, context.Canceled) {
		t.Fatalf("canceled capacity admission=(%+v, %v), want context cancellation", result, err)
	}
	var canceledCapacityErr *promptkit.CapacityError
	if errors.Is(err, promptkit.ErrCapacityExceeded) || errors.As(err, &canceledCapacityErr) {
		t.Fatalf("canceled admission exposed capacity rejection: %v", err)
	}

	result, err = engine.Run(context.Background(), capacityInputRequest("http://second.example/v1"))
	if result != nil {
		t.Fatalf("capacity rejection returned partial result: %+v", result)
	}
	if !errors.Is(err, promptkit.ErrCapacityExceeded) {
		t.Fatalf("capacity rejection=%v, want ErrCapacityExceeded", err)
	}
	if errors.Is(err, promptkit.ErrInvalidRequest) || errors.Is(err, promptkit.ErrLLMGenerate) {
		t.Fatalf("capacity rejection had an unrelated category: %v", err)
	}
	var capacityErr *promptkit.CapacityError
	if !errors.As(err, &capacityErr) || capacityErr == nil {
		t.Fatalf("capacity rejection=%v, want CapacityError", err)
	}
	if capacityErr.BackendID != "limited" {
		t.Fatalf("capacity backend ID=%q, want limited", capacityErr.BackendID)
	}
	capacityErr.BackendID = "changed"

	result, err = engine.Run(context.Background(), capacityInputRequest("http://third.example/v1"))
	var subsequentCapacityErr *promptkit.CapacityError
	if result != nil || !errors.As(err, &subsequentCapacityErr) ||
		subsequentCapacityErr == nil || subsequentCapacityErr.BackendID != "limited" {
		t.Fatalf("subsequent capacity rejection=(%+v, %v), want independent limited CapacityError", result, err)
	}
	if calls := reader.callCount(); calls != 1 {
		t.Fatalf("artifact calls=%d, want only the admitted run", calls)
	}
	if _, _, calls := client.snapshot(); calls != 0 {
		t.Fatalf("client calls=%d before admitted run was released, want 0", calls)
	}

	close(artifactRelease)
	outcome := awaitCapacityRun(t, firstResult)
	if outcome.err != nil || outcome.result == nil {
		t.Fatalf("first run outcome=(%+v, %v), want success", outcome.result, outcome.err)
	}
}

func TestBackendCapacityIsIndependentBetweenEngines(t *testing.T) {
	firstRelease := make(chan struct{})
	firstClient := newCapacityGateClient(firstRelease, 1)
	firstEngine := newBackendCapacityEngine(t, firstClient, 1, capacityInt(0), nil)
	secondClient := newCapacityGateClient(closedCapacityChannel(), 1)
	secondEngine := newBackendCapacityEngine(t, secondClient, 1, capacityInt(0), nil)

	firstResult := make(chan capacityRunResult, 1)
	go runCapacityRequest(firstEngine, context.Background(), promptkit.RunRequest{PromptID: "prompt"}, firstResult)
	awaitCapacityRequest(t, firstClient.started)

	result, err := secondEngine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt"})
	if err != nil || result == nil {
		t.Fatalf("second engine run=(%+v, %v), want independent success", result, err)
	}
	if _, _, calls := secondClient.snapshot(); calls != 1 {
		t.Fatalf("second engine client calls=%d, want 1", calls)
	}

	close(firstRelease)
	outcome := awaitCapacityRun(t, firstResult)
	if outcome.err != nil || outcome.result == nil {
		t.Fatalf("first engine run=(%+v, %v), want success", outcome.result, outcome.err)
	}
}

func TestUnlimitedBackendsRetainInjectedClientConcurrency(t *testing.T) {
	tests := []struct {
		name      string
		configure func(*testing.T, promptkit.LLMClient) *promptkit.Engine
	}{
		{
			name: "custom backend",
			configure: func(t *testing.T, client promptkit.LLMClient) *promptkit.Engine {
				return newBackendCapacityEngine(t, client, 0, nil, nil)
			},
		},
		{
			name: "endpoint-only profile",
			configure: func(t *testing.T, client promptkit.LLMClient) *promptkit.Engine {
				engine, err := promptkit.NewEngine(promptkit.Config{},
					promptkit.WithPromptFS(contractPromptFS("prompt", "profile", "message"), "."),
					promptkit.WithProfiles(promptkit.Profile{
						ID: "profile", Endpoint: "http://endpoint.example/v1", Model: "model",
					}),
					promptkit.WithLLMClient(client),
				)
				if err != nil {
					t.Fatalf("construct endpoint-only engine: %v", err)
				}
				return engine
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			release := make(chan struct{})
			client := newCapacityGateClient(release, 2)
			engine := tc.configure(t, client)
			results := make(chan capacityRunResult, 2)
			for i := 0; i < 2; i++ {
				go runCapacityRequest(
					engine,
					context.Background(),
					promptkit.RunRequest{PromptID: "prompt"},
					results,
				)
			}

			awaitCapacityRequest(t, client.started)
			awaitCapacityRequest(t, client.started)
			if active, peak, _ := client.snapshot(); active != 2 || peak != 2 {
				t.Fatalf("unlimited concurrency=(active=%d peak=%d), want 2", active, peak)
			}
			close(release)
			for i := 0; i < 2; i++ {
				outcome := awaitCapacityRun(t, results)
				if outcome.err != nil || outcome.result == nil {
					t.Fatalf("run outcome=(%+v, %v), want success", outcome.result, outcome.err)
				}
			}
		})
	}
}

func TestCapacityExceededSentinelContract(t *testing.T) {
	if promptkit.ErrCapacityExceeded == nil {
		t.Fatal("ErrCapacityExceeded is nil")
	}
	var nilCapacityErr *promptkit.CapacityError
	zeroCapacityErr := &promptkit.CapacityError{}
	populatedCapacityErr := &promptkit.CapacityError{BackendID: "limited"}
	for _, capacityErr := range []error{nilCapacityErr, zeroCapacityErr} {
		if !errors.Is(capacityErr, promptkit.ErrCapacityExceeded) {
			t.Fatalf("capacity error=%v, want ErrCapacityExceeded", capacityErr)
		}
	}
	var discoveredCapacityErr *promptkit.CapacityError
	if !errors.As(populatedCapacityErr, &discoveredCapacityErr) || discoveredCapacityErr != populatedCapacityErr {
		t.Fatalf("populated capacity error is not discoverable: %v", populatedCapacityErr)
	}

	for _, unrelated := range []error{
		promptkit.ErrInvalidConfig,
		promptkit.ErrInvalidRequest,
		promptkit.ErrLLMGenerate,
		promptkit.ErrValidation,
	} {
		if errors.Is(promptkit.ErrCapacityExceeded, unrelated) ||
			errors.Is(unrelated, promptkit.ErrCapacityExceeded) ||
			errors.Is(populatedCapacityErr, unrelated) {
			t.Fatalf("ErrCapacityExceeded aliases unrelated sentinel %v", unrelated)
		}
	}
}

type capacityRunResult struct {
	result *promptkit.RunResult
	err    error
}

func runCapacityRequest(
	engine *promptkit.Engine,
	ctx context.Context,
	request promptkit.RunRequest,
	results chan<- capacityRunResult,
) {
	result, err := engine.Run(ctx, request)
	results <- capacityRunResult{result: result, err: err}
}

func newBackendCapacityEngine(
	t *testing.T,
	client promptkit.LLMClient,
	limit int,
	queueCapacity *int,
	reader promptkit.ArtifactReader,
) *promptkit.Engine {
	t.Helper()
	promptFS := contractPromptFS("prompt", "profile", "message")
	if reader != nil {
		promptFS = contractInputPromptFS()
	}
	options := []promptkit.Option{
		promptkit.WithPromptFS(promptFS, "."),
		promptkit.WithBackend(promptkit.Backend{
			ID:               "limited",
			Endpoint:         "http://backend.example/v1",
			ConcurrencyLimit: limit,
			QueueCapacity:    queueCapacity,
		}),
		promptkit.WithProfiles(promptkit.Profile{
			ID: "profile", BackendID: "limited", Model: "model",
		}),
		promptkit.WithLLMClient(client),
	}
	if reader != nil {
		options = append(options, promptkit.WithArtifactReader(reader))
	}
	engine, err := promptkit.NewEngine(promptkit.Config{}, options...)
	if err != nil {
		t.Fatalf("construct capacity engine: %v", err)
	}
	return engine
}

func capacityInputRequest(endpoint string) promptkit.RunRequest {
	return promptkit.RunRequest{
		PromptID: "input-prompt",
		Inputs: map[string]promptkit.ArtifactRef{
			"input": promptkit.Inline("input"),
		},
		Execution: &promptkit.ExecutionTargetOverride{Endpoint: endpoint},
	}
}

type capacityGateClient struct {
	mu      sync.Mutex
	active  int
	peak    int
	calls   int
	started chan promptkit.GenerateRequest
	release <-chan struct{}
}

func newCapacityGateClient(release <-chan struct{}, buffer int) *capacityGateClient {
	return &capacityGateClient{
		started: make(chan promptkit.GenerateRequest, buffer),
		release: release,
	}
}

func (c *capacityGateClient) Generate(
	ctx context.Context,
	request promptkit.GenerateRequest,
) (*promptkit.GenerateResponse, error) {
	c.mu.Lock()
	c.calls++
	c.active++
	if c.active > c.peak {
		c.peak = c.active
	}
	c.mu.Unlock()
	defer func() {
		c.mu.Lock()
		c.active--
		c.mu.Unlock()
	}()

	c.started <- request
	select {
	case <-c.release:
		return &promptkit.GenerateResponse{Content: "ok"}, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

func (c *capacityGateClient) snapshot() (active, peak, calls int) {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.active, c.peak, c.calls
}

type capacityArtifactReader struct {
	mu      sync.Mutex
	calls   int
	entered chan struct{}
	release <-chan struct{}
}

func (r *capacityArtifactReader) Read(
	ctx context.Context,
	_ promptkit.ArtifactRef,
) (*promptkit.Artifact, error) {
	r.mu.Lock()
	r.calls++
	r.mu.Unlock()
	r.entered <- struct{}{}
	select {
	case <-r.release:
		return &promptkit.Artifact{Body: []byte("input")}, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

func (r *capacityArtifactReader) callCount() int {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.calls
}

func awaitCapacityRequest(
	t *testing.T,
	requests <-chan promptkit.GenerateRequest,
) promptkit.GenerateRequest {
	t.Helper()
	select {
	case request := <-requests:
		return request
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for client invocation")
		return promptkit.GenerateRequest{}
	}
}

func awaitCapacityRun(t *testing.T, results <-chan capacityRunResult) capacityRunResult {
	t.Helper()
	select {
	case result := <-results:
		return result
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for Run")
		return capacityRunResult{}
	}
}

func awaitCapacitySignal(t *testing.T, signal <-chan struct{}, name string) {
	t.Helper()
	select {
	case <-signal:
	case <-time.After(5 * time.Second):
		t.Fatalf("timed out waiting for %s", name)
	}
}

func capacityInt(value int) *int {
	return &value
}

func closedCapacityChannel() <-chan struct{} {
	channel := make(chan struct{})
	close(channel)
	return channel
}