promptkit/prepared_observation_contract_test.go

package promptkit_test

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

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

func TestPreparedObservationStartsOnlyAfterClaimAndKeepsObserver(t *testing.T) {
	engine := newPreparedContractEngine(t, &preparedRecordingClient{response: &promptkit.GenerateResponse{Content: "ok"}}, "hello")
	var events []promptkit.ExecutionEvent
	request := promptkit.RunRequest{PromptID: "prepared", Observer: func(event promptkit.ExecutionEvent) { events = append(events, event) }}
	prepared, err := engine.PrepareExecution(context.Background(), request)
	if err != nil {
		t.Fatal(err)
	}
	request.Observer = func(promptkit.ExecutionEvent) { t.Error("later observer was invoked") }
	if len(events) != 0 {
		t.Fatalf("preparation emitted events: %+v", events)
	}
	result, err := engine.RunPrepared(context.Background(), prepared)
	if err != nil {
		t.Fatal(err)
	}
	if len(events) != 3 || events[0].Kind != promptkit.ExecutionStarted || events[1].Kind != promptkit.ExecutionAttempt || events[2].Kind != promptkit.ExecutionFinished {
		t.Fatalf("events=%+v", events)
	}
	if events[0].PromptID != "prepared" || events[0].ProfileID != "profile" || events[0].Model != "model" || events[0].RunID != result.RunID || events[2].RunID != result.RunID || events[2].Attempts != 1 {
		t.Fatalf("result=%+v events=%+v", result, events)
	}
	if again, err := engine.RunPrepared(context.Background(), prepared); again != nil || !errors.Is(err, promptkit.ErrInvalidRequest) || len(events) != 3 {
		t.Fatalf("reused result=%+v error=%v events=%+v", again, err, events)
	}
}

func TestPreparedObservationInvalidForeignAndDiscardedHandlesAreSilent(t *testing.T) {
	engine := newPreparedContractEngine(t, &preparedRecordingClient{response: &promptkit.GenerateResponse{Content: "ok"}}, "hello")
	foreign := newPreparedContractEngine(t, &preparedRecordingClient{response: &promptkit.GenerateResponse{Content: "ok"}}, "hello")
	var events []promptkit.ExecutionEvent
	request := promptkit.RunRequest{PromptID: "prepared", Observer: func(event promptkit.ExecutionEvent) { events = append(events, event) }}
	prepared, err := engine.PrepareExecution(context.Background(), request)
	if err != nil {
		t.Fatal(err)
	}
	if result, err := foreign.RunPrepared(context.Background(), prepared); result != nil || !errors.Is(err, promptkit.ErrInvalidRequest) || len(events) != 0 {
		t.Fatalf("foreign result=%+v error=%v events=%+v", result, err, events)
	}
	prepared.Discard()
	if result, err := engine.RunPrepared(context.Background(), prepared); result != nil || !errors.Is(err, promptkit.ErrInvalidRequest) || len(events) != 0 {
		t.Fatalf("discarded result=%+v error=%v events=%+v", result, err, events)
	}
	if result, err := engine.RunPrepared(context.Background(), nil); result != nil || !errors.Is(err, promptkit.ErrInvalidRequest) || len(events) != 0 {
		t.Fatalf("nil result=%+v error=%v events=%+v", result, err, events)
	}
}

func TestPreparedObservationCountsCompletedUsageBeforeRepairFailure(t *testing.T) {
	client := &preparedRecordingClient{responses: []*promptkit.GenerateResponse{{Content: "invalid json", Usage: promptkit.TokenUsage{PromptTokens: 4}}, nil}}
	engine := newPreparedContractEngine(t, client, "hello")
	var events []promptkit.ExecutionEvent
	prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared", Validation: &promptkit.OutputContract{Format: promptkit.FormatJSON, ValidationMode: promptkit.ValidationJSON, RepairAttempts: 1}, Observer: func(event promptkit.ExecutionEvent) { events = append(events, event) }})
	if err != nil {
		t.Fatal(err)
	}
	result, err := engine.RunPrepared(context.Background(), prepared)
	if result != nil || !errors.Is(err, promptkit.ErrLLMGenerate) || len(events) != 4 {
		t.Fatalf("result=%+v error=%v events=%+v", result, err, events)
	}
	terminal := events[3]
	if terminal.Outcome != promptkit.ExecutionOutcomeError || terminal.Failure != promptkit.ExecutionFailureMalformedResponse || terminal.Attempts != 2 || terminal.RepairCalls != 1 || terminal.Usage.PromptTokens.Value != 4 || terminal.Usage.PromptTokens.Complete {
		t.Fatalf("terminal=%+v", terminal)
	}
}

func TestPreparedObservationClaimRaceAndPanicCleanup(t *testing.T) {
	engine := newPreparedContractEngine(t, &preparedRecordingClient{response: &promptkit.GenerateResponse{Content: "ok"}}, "hello")
	var mu sync.Mutex
	var events []promptkit.ExecutionEvent
	prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared", Observer: func(event promptkit.ExecutionEvent) { mu.Lock(); events = append(events, event); mu.Unlock() }})
	if err != nil {
		t.Fatal(err)
	}
	var wg sync.WaitGroup
	var successes, invalid int
	var countMu sync.Mutex
	for range 8 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			_, err := engine.RunPrepared(context.Background(), prepared)
			countMu.Lock()
			defer countMu.Unlock()
			if err == nil {
				successes++
			} else if errors.Is(err, promptkit.ErrInvalidRequest) {
				invalid++
			} else {
				t.Errorf("unexpected error: %v", err)
			}
		}()
	}
	wg.Wait()
	if successes != 1 || invalid != 7 || len(events) != 3 {
		t.Fatalf("successes=%d invalid=%d events=%+v", successes, invalid, events)
	}

	panicking, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared", Observer: func(event promptkit.ExecutionEvent) {
		if event.Kind == promptkit.ExecutionStarted {
			panic("callback")
		}
	}})
	if err != nil {
		t.Fatal(err)
	}
	func() {
		defer func() {
			if recover() == nil {
				t.Error("expected panic")
			}
		}()
		_, _ = engine.RunPrepared(context.Background(), panicking)
	}()
	if result, err := engine.RunPrepared(context.Background(), panicking); result != nil || !errors.Is(err, promptkit.ErrInvalidRequest) {
		t.Fatalf("reused after panic: result=%+v error=%v", result, err)
	}
}

func TestPreparedObservationFinishesAfterCredentialAndCapacityFailures(t *testing.T) {
	const envName = "PROMPTKIT_PREPARED_OBSERVATION_KEY"
	t.Setenv(envName, "available")
	engine, err := promptkit.NewEngine(promptkit.Config{},
		promptkit.WithPromptFS(contractPromptFS("prepared", "profile", "hello"), "."),
		promptkit.WithProfiles(promptkit.Profile{ID: "profile", Endpoint: "https://example.invalid/v1", Model: "model", APIKeyRequired: true}),
		promptkit.WithLLMClient(&fakeLLMClient{response: &promptkit.GenerateResponse{Content: "ok"}}))
	if err != nil {
		t.Fatal(err)
	}
	var events []promptkit.ExecutionEvent
	prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared", Execution: &promptkit.ExecutionTargetOverride{APIKeyEnv: envName}, Observer: func(event promptkit.ExecutionEvent) { events = append(events, event) }})
	if err != nil {
		t.Fatal(err)
	}
	t.Setenv(envName, "")
	result, err := engine.RunPrepared(context.Background(), prepared)
	if result != nil || !errors.Is(err, promptkit.ErrInvalidRequest) || len(events) != 2 || events[1].Failure != promptkit.ExecutionFailureInvalidRequest || events[1].Attempts != 0 {
		t.Fatalf("credential result=%+v error=%v events=%+v", result, err, events)
	}

	release := make(chan struct{})
	client := newCapacityGateClient(release, 2)
	limited := newBackendCapacityEngine(t, client, 1, capacityInt(0), nil)
	first := make(chan capacityRunResult, 1)
	go runCapacityRequest(limited, context.Background(), promptkit.RunRequest{PromptID: "prompt"}, first)
	awaitCapacityRequest(t, client.started)
	events = nil
	prepared, err = limited.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prompt", Observer: func(event promptkit.ExecutionEvent) { events = append(events, event) }})
	if err != nil {
		t.Fatal(err)
	}
	result, err = limited.RunPrepared(context.Background(), prepared)
	close(release)
	if firstResult := awaitCapacityRun(t, first); firstResult.err != nil {
		t.Fatalf("first: %v", firstResult.err)
	}
	if result != nil || !errors.Is(err, promptkit.ErrCapacityExceeded) || len(events) != 2 || events[1].Failure != promptkit.ExecutionFailureCapacity || events[1].Attempts != 0 {
		t.Fatalf("capacity result=%+v error=%v events=%+v", result, err, events)
	}
}

func TestPreparedObservationCanceledClaimStillFinishes(t *testing.T) {
	engine := newPreparedContractEngine(t, &preparedRecordingClient{response: &promptkit.GenerateResponse{Content: "ok"}}, "hello")
	var events []promptkit.ExecutionEvent
	prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared", Observer: func(event promptkit.ExecutionEvent) { events = append(events, event) }})
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	result, err := engine.RunPrepared(ctx, prepared)
	if result != nil || !errors.Is(err, context.Canceled) || len(events) != 3 || events[1].Dispatched || events[2].Failure != promptkit.ExecutionFailureCanceled || events[2].Attempts != 0 {
		t.Fatalf("result=%+v error=%v events=%+v", result, err, events)
	}
}