promptkit/reliable_execution_contract_test.go

package promptkit_test

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"strings"
	"sync/atomic"
	"testing"
	"testing/fstest"
	"time"

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

func observedHTTPResponse(status int, body string) *http.Response {
	return &http.Response{StatusCode: status, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))}
}

func TestPublicRetryRepairSequenceKeepsUsageAndEventsConsistent(t *testing.T) {
	responses := []*http.Response{
		observedHTTPResponse(503, `{"error":{"message":"private-provider-marker"}}`),
		observedHTTPResponse(200, `{"choices":[{"message":{"content":"not json"}}],"usage":{"prompt_tokens":0,"completion_tokens":2}}`),
		observedHTTPResponse(503, `{"error":{"message":"private-provider-marker"}}`),
		observedHTTPResponse(200, `{"choices":[{"message":{"content":"{\"ok\":true}"}}],"usage":{"prompt_tokens":3,"total_tokens":5,"cache_write_tokens":0}}`),
	}
	var calls int
	one := 1
	config := contractConfig(frameworkSchemaDir)
	config.TransientRetries = &one
	config.HTTPClient = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
		if calls >= len(responses) {
			t.Error("extra HTTP attempt")
			return nil, errors.New("extra attempt")
		}
		r := responses[calls]
		calls++
		return r, nil
	})}
	engine, err := testfixture.NewEngine(config)
	if err != nil {
		t.Fatal(err)
	}
	var events []promptkit.ExecutionEvent
	req := generationErrorRunRequest()
	req.Validation = &promptkit.OutputContract{Format: promptkit.FormatJSON, ValidationMode: promptkit.ValidationJSON, RepairAttempts: 1}
	req.Observer = func(event promptkit.ExecutionEvent) { events = append(events, event) }
	result, err := engine.Run(context.Background(), req)
	if err != nil || result == nil || !result.Validation.IsValid || calls != 4 {
		t.Fatalf("result=%+v error=%v calls=%d", result, err, calls)
	}
	if len(events) != 8 {
		t.Fatalf("events=%+v", events)
	}
	for index, want := range []promptkit.ExecutionEventKind{promptkit.ExecutionStarted, promptkit.ExecutionAttempt, promptkit.ExecutionRetryDelay, promptkit.ExecutionAttempt, promptkit.ExecutionAttempt, promptkit.ExecutionRetryDelay, promptkit.ExecutionAttempt, promptkit.ExecutionFinished} {
		if events[index].Kind != want {
			t.Fatalf("event %d = %+v", index, events[index])
		}
	}
	if events[1].Generation != 1 || events[3].Attempt != 2 || events[4].Generation != 2 || events[4].Attempt != 1 || events[6].Generation != 2 || events[6].Attempt != 2 {
		t.Fatalf("attempt indexes=%+v", events)
	}
	terminal := events[7]
	if terminal.RunID != result.RunID || terminal.Attempts != 4 || terminal.TransientRetries != 2 || terminal.RepairCalls != 1 || terminal.Outcome != promptkit.ExecutionOutcomeSuccess || terminal.RetriesExhausted {
		t.Fatalf("terminal=%+v result=%+v", terminal, result)
	}
	if terminal.Usage.PromptTokens.Value != 3 || !terminal.Usage.PromptTokens.Known || terminal.Usage.PromptTokens.Complete || terminal.Usage.CompletionTokens.Value != 2 || terminal.Usage.CompletionTokens.Complete || terminal.Usage.TotalTokens.Value != 5 || terminal.Usage.TotalTokens.Complete || !events[3].Usage.PromptTokens.Known || events[3].Usage.PromptTokens.Value != 0 {
		t.Fatalf("observed usage=%+v", terminal.Usage)
	}
	if result.Usage.PromptTokens != 3 || result.Usage.CompletionTokens != 2 || result.Usage.TotalTokens != 5 || result.Usage.CacheWriteTokens != 0 {
		t.Fatalf("result usage=%+v", result.Usage)
	}
	encoded, err := json.Marshal(events)
	if err != nil {
		t.Fatal(err)
	}
	if strings.Contains(string(encoded), "private-provider-marker") || strings.Contains(string(encoded), "not json") {
		t.Fatalf("events leaked content: %s", encoded)
	}
}

func TestProfilesSharingBackendCoordinateRetryRepairAndAdmission(t *testing.T) {
	firstA := make(chan struct{})
	enteredB := make(chan struct{})
	releaseB := make(chan struct{})
	var callsA, callsB, active atomic.Int64
	transport := roundTripFunc(func(req *http.Request) (*http.Response, error) {
		if active.Add(1) != 1 {
			t.Error("overlapping HTTP attempts on limited backend")
		}
		defer active.Add(-1)
		var payload struct {
			Model string `json:"model"`
		}
		if err := json.NewDecoder(req.Body).Decode(&payload); err != nil {
			return nil, err
		}
		switch payload.Model {
		case "model-a":
			switch callsA.Add(1) {
			case 1:
				response := observedHTTPResponse(503, `{"error":{}}`)
				response.Header.Set("Retry-After", "1")
				close(firstA)
				return response, nil
			case 2:
				return observedHTTPResponse(200, `{"choices":[{"message":{"content":"invalid json"}}]}`), nil
			default:
				return observedHTTPResponse(200, `{"choices":[{"message":{"content":"{\"ok\":true}"}}]}`), nil
			}
		case "model-b":
			callsB.Add(1)
			close(enteredB)
			select {
			case <-releaseB:
			case <-req.Context().Done():
				return nil, req.Context().Err()
			}
			return observedHTTPResponse(200, `{"choices":[{"message":{"content":"done"}}]}`), nil
		default:
			return nil, errors.New("unexpected model")
		}
	})
	one := 1
	engine, err := testfixture.NewEngine(testfixture.Config{HTTPClient: &http.Client{Transport: transport}, TransientRetries: &one},
		testfixture.WithPromptFS(contractPromptFS("prompt", "profile-a", "hello"), "."),
		testfixture.WithBackend(testfixture.Backend{ID: "limited", Endpoint: "https://example.invalid/v1", ConcurrencyLimit: 1, QueueCapacity: &one}),
		testfixture.WithProfiles(testfixture.Profile{ID: "profile-a", BackendID: "limited", Model: "model-a"}, testfixture.Profile{ID: "profile-b", BackendID: "limited", Model: "model-b"}))
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 8*time.Second)
	defer cancel()
	resultA := make(chan capacityRunResult, 1)
	requestA := promptkit.RunRequest{PromptID: "prompt", Validation: &promptkit.OutputContract{Format: promptkit.FormatJSON, ValidationMode: promptkit.ValidationJSON, RepairAttempts: 1}}
	go runCapacityRequest(engine, ctx, requestA, resultA)
	select {
	case <-firstA:
	case <-ctx.Done():
		t.Fatal("first model attempt did not start")
	}
	resultB := make(chan capacityRunResult, 1)
	go runCapacityRequest(engine, ctx, promptkit.RunRequest{PromptID: "prompt", ProfileID: "profile-b"}, resultB)
	select {
	case <-enteredB:
	case <-ctx.Done():
		t.Fatal("second profile did not acquire permit during backoff")
	}
	third, thirdErr := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", ProfileID: "profile-b"})
	if third != nil || !errors.Is(thirdErr, promptkit.ErrCapacityExceeded) {
		t.Fatalf("admission during two active runs: result=%+v error=%v", third, thirdErr)
	}
	close(releaseB)
	if outcome := <-resultB; outcome.err != nil || outcome.result == nil {
		t.Fatalf("second profile: %+v", outcome)
	}
	if outcome := <-resultA; outcome.err != nil || outcome.result == nil || !outcome.result.Validation.IsValid {
		t.Fatalf("retry and repair: %+v", outcome)
	}
	if callsA.Load() != 3 || callsB.Load() != 1 {
		t.Fatalf("calls A=%d B=%d", callsA.Load(), callsB.Load())
	}
}

func TestDifferentBackendPoolsCanDispatchTogether(t *testing.T) {
	entered := make(chan string, 2)
	release := make(chan struct{})
	var active atomic.Int64
	transport := roundTripFunc(func(req *http.Request) (*http.Response, error) {
		var payload struct {
			Model string `json:"model"`
		}
		if err := json.NewDecoder(req.Body).Decode(&payload); err != nil {
			return nil, err
		}
		active.Add(1)
		defer active.Add(-1)
		entered <- payload.Model
		select {
		case <-release:
		case <-req.Context().Done():
			return nil, req.Context().Err()
		}
		return observedHTTPResponse(200, `{"choices":[{"message":{"content":"ok"}}]}`), nil
	})
	zero := 0
	engine, err := testfixture.NewEngine(testfixture.Config{HTTPClient: &http.Client{Transport: transport}, TransientRetries: &zero},
		testfixture.WithPromptFS(contractPromptFS("prompt", "profile-a", "hello"), "."),
		testfixture.WithBackend(testfixture.Backend{ID: "backend-a", Endpoint: "https://example.invalid/v1", ConcurrencyLimit: 1, QueueCapacity: &zero}),
		testfixture.WithBackend(testfixture.Backend{ID: "backend-b", Endpoint: "https://example.invalid/v1", ConcurrencyLimit: 1, QueueCapacity: &zero}),
		testfixture.WithProfiles(testfixture.Profile{ID: "profile-a", BackendID: "backend-a", Model: "model-a"}, testfixture.Profile{ID: "profile-b", BackendID: "backend-b", Model: "model-b"}))
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	results := make(chan capacityRunResult, 2)
	go runCapacityRequest(engine, ctx, promptkit.RunRequest{PromptID: "prompt"}, results)
	select {
	case model := <-entered:
		if model != "model-a" {
			t.Fatalf("first model=%q", model)
		}
	case <-ctx.Done():
		t.Fatal("first backend did not dispatch")
	}
	go runCapacityRequest(engine, ctx, promptkit.RunRequest{PromptID: "prompt", ProfileID: "profile-b"}, results)
	select {
	case model := <-entered:
		if model != "model-b" {
			t.Fatalf("second model=%q", model)
		}
	case <-ctx.Done():
		t.Fatal("independent backend did not dispatch")
	}
	if active.Load() != 2 {
		t.Fatalf("active HTTP attempts=%d", active.Load())
	}
	close(release)
	for range 2 {
		if outcome := <-results; outcome.err != nil || outcome.result == nil {
			t.Fatalf("outcome=%+v", outcome)
		}
	}
}

func TestPreparedRetriesFreezeSourceAndCredentialAtExecution(t *testing.T) {
	const envName = "PROMPTKIT_PREPARED_RETRY_KEY"
	t.Setenv(envName, "first-key")
	source := contractPromptFS("prompt", "profile", "original message")
	var calls int
	var auth, content []string
	transport := roundTripFunc(func(req *http.Request) (*http.Response, error) {
		calls++
		auth = append(auth, req.Header.Get("Authorization"))
		var payload struct {
			Messages []struct {
				Content string `json:"content"`
			} `json:"messages"`
		}
		if err := json.NewDecoder(req.Body).Decode(&payload); err != nil {
			return nil, err
		}
		content = append(content, payload.Messages[0].Content)
		if calls == 1 {
			t.Setenv(envName, "later-key")
			return observedHTTPResponse(503, `{"error":{}}`), nil
		}
		return observedHTTPResponse(200, `{"choices":[{"message":{"content":"ok"}}]}`), nil
	})
	one := 1
	engine, err := testfixture.NewEngine(testfixture.Config{HTTPClient: &http.Client{Transport: transport}, TransientRetries: &one},
		testfixture.WithPromptFS(source, "."), testfixture.WithProfiles(testfixture.Profile{ID: "profile", Endpoint: "https://example.invalid/v1", Model: "model"}))
	if err != nil {
		t.Fatal(err)
	}
	var events []promptkit.ExecutionEvent
	prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prompt", Execution: &promptkit.ExecutionTargetOverride{APIKeyEnv: envName}, Observer: func(event promptkit.ExecutionEvent) { events = append(events, event) }})
	if err != nil {
		t.Fatal(err)
	}
	source["prompt.yaml"] = &fstest.MapFile{Data: []byte(`id: prompt
version: "1"
default_profile: profile
messages:
  - role: user
    content: changed message
output:
  format: text
  validation_mode: none
`)}
	result, err := engine.RunPrepared(context.Background(), prepared)
	if err != nil || result == nil || calls != 2 {
		t.Fatalf("result=%+v error=%v calls=%d", result, err, calls)
	}
	for index := range auth {
		if auth[index] != "Bearer first-key" || content[index] != "original message" {
			t.Fatalf("attempt %d auth=%q content=%q", index, auth[index], content[index])
		}
	}
	if events[len(events)-1].Attempts != 2 || events[len(events)-1].TransientRetries != 1 {
		t.Fatalf("terminal=%+v", events[len(events)-1])
	}
}

func TestPublicNonretryableAndInjectedAttemptsStaySingle(t *testing.T) {
	three := 3
	engine, calls := retryTestEngine(t, &three, func(int) *http.Response { return retryHTTPResponse(400, "") })
	var events []promptkit.ExecutionEvent
	req := generationErrorRunRequest()
	req.Observer = func(event promptkit.ExecutionEvent) { events = append(events, event) }
	result, err := engine.Run(context.Background(), req)
	if result != nil || !errors.Is(err, promptkit.ErrLLMGenerate) || errors.Is(err, promptkit.ErrRetriesExhausted) || *calls != 1 || len(events) != 3 {
		t.Fatalf("result=%+v error=%v calls=%d events=%+v", result, err, *calls, events)
	}
	if terminal := events[2]; terminal.Attempts != 1 || terminal.TransientRetries != 0 || terminal.RetriesExhausted || terminal.Failure != promptkit.ExecutionFailureProviderStatus || terminal.StatusCode != 400 {
		t.Fatalf("terminal=%+v", terminal)
	}

	client := &fakeLLMClient{err: errors.New("opaque injected failure")}
	injected := observationEngine(t, client)
	events = nil
	request := promptkit.RunRequest{PromptID: "prompt", TransientRetries: &three, Observer: func(event promptkit.ExecutionEvent) { events = append(events, event) }}
	result, err = injected.Run(context.Background(), request)
	if result != nil || !errors.Is(err, client.err) || len(client.requests) != 1 || len(events) != 3 || events[2].Attempts != 1 || events[2].TransientRetries != 0 || events[2].Failure != promptkit.ExecutionFailureGeneration {
		t.Fatalf("result=%+v error=%v events=%+v", result, err, events)
	}
}

func TestPublicExhaustionObservationRetainsProviderCause(t *testing.T) {
	one := 1
	engine, calls := retryTestEngine(t, &one, func(int) *http.Response { return retryHTTPResponse(503, "") })
	var events []promptkit.ExecutionEvent
	req := generationErrorRunRequest()
	req.Observer = func(event promptkit.ExecutionEvent) { events = append(events, event) }
	result, err := engine.Run(context.Background(), req)
	var exhausted *promptkit.RetryExhaustedError
	var provider *promptkit.GenerationError
	if result != nil || *calls != 2 || !errors.As(err, &exhausted) || !errors.As(err, &provider) || exhausted.Attempts() != 2 || provider.StatusCode() != 503 {
		t.Fatalf("result=%+v error=%v calls=%d", result, err, *calls)
	}
	terminal := events[len(events)-1]
	if !terminal.RetriesExhausted || terminal.Failure != promptkit.ExecutionFailureProviderStatus || terminal.StatusCode != 503 || terminal.Attempts != 2 || terminal.TransientRetries != 1 {
		t.Fatalf("terminal=%+v", terminal)
	}
}

func TestPublicObserverCancellationStopsBeforeRetry(t *testing.T) {
	one := 1
	engine, calls := retryTestEngine(t, &one, func(int) *http.Response { return retryHTTPResponse(503, "") })
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	var events []promptkit.ExecutionEvent
	req := generationErrorRunRequest()
	req.Observer = func(event promptkit.ExecutionEvent) {
		events = append(events, event)
		if event.Kind == promptkit.ExecutionAttempt {
			cancel()
		}
	}
	result, err := engine.Run(ctx, req)
	if result != nil || !errors.Is(err, context.Canceled) || *calls != 1 || len(events) != 3 || events[2].Failure != promptkit.ExecutionFailureCanceled || events[2].Attempts != 1 || events[2].RetriesExhausted {
		t.Fatalf("result=%+v error=%v calls=%d events=%+v", result, err, *calls, events)
	}
}