Verify prepared notification ownership and cleanup

prepared_notification_contract_test.go

@@ -3,6 +3,10 @@ package promptkit_test
 import (
 	"context"
 	"errors"
+	"fmt"
+	"net/http"
+	"reflect"
+	"strings"
 	"testing"
 
 	"go.maximumdirect.net/promptkit"
@@ -47,6 +51,160 @@ func TestPreparedNotificationSelectionAndOrdering(t *testing.T) {
 	}
 }
 
+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)
+	}
+	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)
+	}
+}
+
+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()