audita/internal/framework/llm/adapter_capacity_test.go

package llm

import (
	"context"
	"errors"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"gitea.maximumdirect.net/eric/audita/v3/internal/core/config"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/responseschema"
	ck "go.maximumdirect.net/configkit"
)

func capacityFixture(t *testing.T, resources string, transport http.RoundTripper) *Runtime {
	t.Helper()
	dir := t.TempDir()
	for name, content := range map[string]string{"config.yml": "version: 3\n", "prompts.yml": resources} {
		if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0600); err != nil {
			t.Fatal(err)
		}
	}
	loaded, err := config.Load(ck.Input{Source: ck.File(filepath.Join(dir, "config.yml"))}, "")
	if err != nil {
		t.Fatal(err)
	}
	runtime, err := NewRuntime(context.Background(), loaded, &http.Client{Transport: transport})
	if err != nil {
		t.Fatal(err)
	}
	return runtime
}

func capacityResponse(r *http.Request) *http.Response {
	return &http.Response{StatusCode: 200, Header: make(http.Header), Request: r, Body: io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"{\"corrections\":[]}"}}]}`))}
}

func TestRolesShareBackendCapacityAndReleaseAfterCancellation(t *testing.T) {
	entered := make(chan struct{}, 3)
	release := make(chan struct{})
	transport := transportFunc(func(r *http.Request) (*http.Response, error) {
		entered <- struct{}{}
		select {
		case <-release:
			return capacityResponse(r), nil
		case <-r.Context().Done():
			return nil, r.Context().Err()
		}
	})
	resources := "backends:\n  - id: shared\n    endpoint: https://example.invalid/v1\n    concurrency_limit: 1\n    queue_capacity: 1\nprofiles:\n  - id: proposal\n    backend: shared\n    model: local\n  - id: validation\n    backend: shared\n    model: local\n"
	runtime := capacityFixture(t, resources, transport)
	settings := config.LLM{OperationTimeout: 3 * time.Second}
	proposal, err := runtime.NewRoleClient("proposal", "proposal", settings, CallHooks{})
	if err != nil {
		t.Fatal(err)
	}
	validation, err := runtime.NewRoleClient("validation", "validation", settings, CallHooks{})
	if err != nil {
		t.Fatal(err)
	}
	request := adapterRequest(responseschema.CorrectionSetKey)
	done := make(chan error, 1)
	go func() {
		_, err := proposal.CompleteStructured(context.Background(), request, &map[string]any{})
		done <- err
	}()
	select {
	case <-entered:
	case <-time.After(2 * time.Second):
		t.Fatal("first call did not enter provider")
	}
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
	defer cancel()
	_, err = validation.CompleteStructured(ctx, request, &map[string]any{})
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("queued call did not honor deadline: %v", err)
	}
	select {
	case <-entered:
		t.Fatal("queued call dispatched")
	default:
	}
	close(release)
	if err := <-done; err != nil {
		t.Fatal(err)
	}
	if _, err := validation.CompleteStructured(context.Background(), request, &map[string]any{}); err != nil {
		t.Fatalf("cancellation leaked capacity: %v", err)
	}
}

func TestDistinctBackendsHaveIndependentCapacity(t *testing.T) {
	entered := make(chan struct{}, 2)
	release := make(chan struct{})
	transport := transportFunc(func(r *http.Request) (*http.Response, error) {
		entered <- struct{}{}
		<-release
		return capacityResponse(r), nil
	})
	resources := "backends:\n  - id: first\n    endpoint: https://example.invalid/v1\n    concurrency_limit: 1\n    queue_capacity: 0\n  - id: second\n    endpoint: https://example.invalid/v1\n    concurrency_limit: 1\n    queue_capacity: 0\nprofiles:\n  - id: proposal\n    backend: first\n    model: local\n  - id: validation\n    backend: second\n    model: local\n"
	runtime := capacityFixture(t, resources, transport)
	settings := config.LLM{OperationTimeout: 3 * time.Second}
	first, err := runtime.NewRoleClient("proposal", "proposal", settings, CallHooks{})
	if err != nil {
		t.Fatal(err)
	}
	second, err := runtime.NewRoleClient("validation", "validation", settings, CallHooks{})
	if err != nil {
		t.Fatal(err)
	}
	request := adapterRequest(responseschema.CorrectionSetKey)
	done := make(chan error, 2)
	go func() {
		_, err := first.CompleteStructured(context.Background(), request, &map[string]any{})
		done <- err
	}()
	select {
	case <-entered:
	case <-time.After(2 * time.Second):
		t.Fatal("first backend did not dispatch")
	}
	go func() {
		_, err := second.CompleteStructured(context.Background(), request, &map[string]any{})
		done <- err
	}()
	select {
	case <-entered:
	case <-time.After(2 * time.Second):
		close(release)
		t.Fatal("second backend was blocked by first")
	}
	close(release)
	for range 2 {
		if err := <-done; err != nil {
			t.Fatal(err)
		}
	}
}