Notify callers from the prepared execution snapshot

prepared_notification_contract_test.go

@@ -0,0 +1,61 @@
+package promptkit_test
+
+import (
+	"context"
+	"errors"
+	"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 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)
+	}
+}