promptkit/prepared_notification_contract_test.go

package promptkit_test

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"reflect"
	"strings"
	"testing"

	"go.maximumdirect.net/promptkit"
)

func TestPreparedNotificationSelectionAndOrdering(t *testing.T) {
	for _, split := range []bool{false, true} {
		t.Run(map[bool]string{false: "run", true: "prepared"}[split], func(t *testing.T) {
			calls := 0
			engine := observationEngine(t, observationClientFunc(func(_ context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
				if calls != 1 {
					t.Fatalf("generation before notification: %d", calls)
				}
				return &promptkit.GenerateResponse{Content: "ok"}, nil
			}))
			req := promptkit.RunRequest{PromptID: "prompt", OnPrepared: func(_ context.Context, details promptkit.PreparedRun) {
				calls++
				if details.PromptID != "prompt" || len(details.Messages) == 0 {
					t.Fatalf("details=%+v", details)
				}
			}}
			if split {
				handle, err := engine.PrepareExecution(context.Background(), req)
				if err != nil {
					t.Fatal(err)
				}
				if calls != 0 {
					t.Fatal("preparation notified")
				}
				req.OnPrepared = func(context.Context, promptkit.PreparedRun) { t.Fatal("selection changed") }
				_, err = engine.RunPrepared(context.Background(), handle)
				if err != nil {
					t.Fatal(err)
				}
			} else if _, err := engine.Run(context.Background(), req); err != nil {
				t.Fatal(err)
			}
			if calls != 1 {
				t.Fatalf("notifications=%d", calls)
			}
		})
	}
}

func TestNotificationRedactsEnvironmentCredential(t *testing.T) {
	t.Setenv("PROMPTKIT_NOTIFICATION_KEY", "synthetic-environment-key")
	config := simpleConfig("message")
	config.HTTPClient = &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
		if req.Header.Get("Authorization") != "Bearer synthetic-environment-key" {
			t.Error("environment credential did not reach provider")
		}
		return retryHTTPResponse(200, "ok"), nil
	})}
	engine, err := promptkit.NewEngine(config)
	if err != nil {
		t.Fatal(err)
	}
	req := simpleGenerationRequest()
	req.Execution = &promptkit.ExecutionTargetOverride{APIKeyEnv: "PROMPTKIT_NOTIFICATION_KEY"}
	req.OnPrepared = func(_ context.Context, details promptkit.PreparedRun) {
		if strings.Contains(fmt.Sprintf("%#v", details), "synthetic-environment-key") {
			t.Error("notification exposed credential")
		}
	}
	if _, err := engine.Run(context.Background(), req); err != nil {
		t.Fatal(err)
	}
}

func TestContinuationNotificationIsDetachedAndSurvivesFailure(t *testing.T) {
	failure := errors.New("provider failure")
	calls := 0
	engine := newSimpleEngine(t, "question", promptkit.WithLLMClient(observationClientFunc(func(_ context.Context, req promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
		calls++
		if calls == 1 {
			return &promptkit.GenerateResponse{Content: "answer"}, nil
		}
		if req.Prompt.Messages[0].Content != "question" || req.Prompt.Messages[3].Content != "follow-up" || req.APIKey != "synthetic-key" {
			t.Fatalf("generation request=%+v", req)
		}
		return nil, failure
	})))
	initial, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true})
	if err != nil {
		t.Fatal(err)
	}
	before := initial.Exchange.Details()
	var retained promptkit.PreparedRun
	req := promptkit.RunRequest{PromptID: "prompt", ContinueFrom: initial.Exchange, APIKey: "synthetic-key", SessionID: "session", AppendedMessages: []promptkit.RenderedMessage{{Role: promptkit.RoleUser, Content: "follow-up", CacheControl: &promptkit.CacheControl{Type: promptkit.CacheControlEphemeral}}}, OnPrepared: func(_ context.Context, details promptkit.PreparedRun) {
		retained = details
		details.Messages[0].Content = "changed"
		details.Messages[3].Content = "changed"
		details.Messages[3].CacheControl.Type = "changed"
	}}
	result, err := engine.Run(context.Background(), req)
	if result != nil || !errors.Is(err, failure) || retained.ParentRunID != initial.RunID || retained.SessionID != "session" || !reflect.DeepEqual(before, initial.Exchange.Details()) {
		t.Fatalf("result=%v error=%v retained=%+v", result, err, retained)
	}
}

func TestPreparedNotificationInvalidHandlesAndPreparationOnlyCalls(t *testing.T) {
	calls := 0
	engine := newPreparedContractEngine(t, &preparedRecordingClient{response: &promptkit.GenerateResponse{Content: "ok"}}, "content")
	foreign := newPreparedContractEngine(t, &preparedRecordingClient{response: &promptkit.GenerateResponse{Content: "ok"}}, "other")
	req := promptkit.RunRequest{PromptID: "prepared", OnPrepared: func(context.Context, promptkit.PreparedRun) { calls++ }}
	if _, err := engine.Prepare(context.Background(), req); err != nil {
		t.Fatal(err)
	}
	if err := engine.ValidateRequest(context.Background(), req); err != nil {
		t.Fatal(err)
	}
	handle, err := engine.PrepareExecution(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}
	_ = handle.Details()
	if _, err := foreign.RunPrepared(context.Background(), handle); !errors.Is(err, promptkit.ErrInvalidRequest) {
		t.Fatal(err)
	}
	if _, err := engine.RunPrepared(context.Background(), nil); !errors.Is(err, promptkit.ErrInvalidRequest) {
		t.Fatal(err)
	}
	handle.Discard()
	if _, err := engine.RunPrepared(context.Background(), handle); !errors.Is(err, promptkit.ErrInvalidRequest) {
		t.Fatal(err)
	}
	if _, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "missing", OnPrepared: req.OnPrepared}); err == nil {
		t.Fatal("expected preparation failure")
	}
	if calls != 0 {
		t.Fatalf("unexpected notifications=%d", calls)
	}
	handle, err = engine.PrepareExecution(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := engine.RunPrepared(context.Background(), handle); err != nil {
		t.Fatal(err)
	}
	if _, err := engine.RunPrepared(context.Background(), handle); !errors.Is(err, promptkit.ErrInvalidRequest) || calls != 1 {
		t.Fatalf("reused handle error=%v notifications=%d", err, calls)
	}
}

func TestNotificationHoldsAdmission(t *testing.T) {
	engine := newBackendCapacityEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
		return &promptkit.GenerateResponse{Content: "ok"}, nil
	}), 1, capacityInt(0), nil)
	_, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", OnPrepared: func(ctx context.Context, _ promptkit.PreparedRun) {
		if _, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt"}); !errors.Is(err, promptkit.ErrCapacityExceeded) {
			t.Errorf("nested admission error=%v", err)
		}
	}})
	if err != nil {
		t.Fatal(err)
	}
}

func TestNotificationPanicReleasesCapacityAndConsumesHandle(t *testing.T) {
	engine := newBackendCapacityEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
		return &promptkit.GenerateResponse{Content: "ok"}, nil
	}), 1, capacityInt(0), nil)
	for _, split := range []bool{false, true} {
		req := promptkit.RunRequest{PromptID: "prompt", OnPrepared: func(context.Context, promptkit.PreparedRun) { panic("notification") }}
		var handle *promptkit.PreparedExecution
		if split {
			var err error
			handle, err = engine.PrepareExecution(context.Background(), req)
			if err != nil {
				t.Fatal(err)
			}
		}
		func() {
			defer func() {
				if recover() == nil {
					t.Error("expected panic")
				}
			}()
			if split {
				_, _ = engine.RunPrepared(context.Background(), handle)
			} else {
				_, _ = engine.Run(context.Background(), req)
			}
		}()
		if split {
			if _, err := engine.RunPrepared(context.Background(), handle); !errors.Is(err, promptkit.ErrInvalidRequest) {
				t.Fatal("handle reusable after panic")
			}
		}
		if _, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt"}); err != nil {
			t.Fatalf("capacity not released: %v", err)
		}
	}
}

func TestNotificationSkippedOnAdmissionAndPreExecutionCancellation(t *testing.T) {
	release := make(chan struct{})
	client := newCapacityGateClient(release, 1)
	engine := newBackendCapacityEngine(t, client, 1, capacityInt(0), nil)
	results := make(chan capacityRunResult, 1)
	go runCapacityRequest(engine, context.Background(), promptkit.RunRequest{PromptID: "prompt"}, results)
	<-client.started
	calls := 0
	req := promptkit.RunRequest{PromptID: "prompt", OnPrepared: func(context.Context, promptkit.PreparedRun) { calls++ }}
	if _, err := engine.Run(context.Background(), req); !errors.Is(err, promptkit.ErrCapacityExceeded) {
		t.Fatalf("admission error=%v", err)
	}
	close(release)
	if result := <-results; result.err != nil {
		t.Fatal(result.err)
	}
	handle, err := engine.PrepareExecution(context.Background(), req)
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if _, err = engine.RunPrepared(ctx, handle); !errors.Is(err, context.Canceled) {
		t.Fatalf("cancellation error=%v", err)
	}
	if calls != 0 {
		t.Fatalf("notifications=%d", calls)
	}
}

func TestCancellationInPreparedNotificationPreventsGeneration(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	engine := observationEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
		t.Fatal("generation after cancellation")
		return nil, nil
	}))
	_, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", OnPrepared: func(context.Context, promptkit.PreparedRun) { cancel() }})
	if !errors.Is(err, context.Canceled) || !errors.Is(err, promptkit.ErrLLMGenerate) {
		t.Fatalf("error=%v", err)
	}
}