Notify callers from the prepared execution snapshot

docs/roadmap/implementation.md

@@ -254,6 +254,8 @@ Notification invocation itself is completed in Stage 4.
 
 ## Stage 4: Notify From The Actual Shared Prepared Snapshot
 
+Status: complete.
+
 Outcome: ordinary and prepared execution expose one detached logical request
 before first generation, without repeating preparation.
 

internal/usecase/prepared_execution.go

@@ -152,7 +152,7 @@ func (r *Runner) RunPrepared(ctx context.Context, prepared *PreparedExecution) (
 	}
 	defer release()
 
-	return r.executePreparedRun(ctx, payload.prepared, payload.capture, payload.parent, payload.directKey, payload.transientRetries, runID, start, func(
+	return r.executePreparedRun(ctx, payload.prepared, payload.onPrepared, payload.capture, payload.parent, payload.directKey, payload.transientRetries, runID, start, func(
 		ctx context.Context,
 		artifact *domain.Artifact,
 		attemptsUsed int,

internal/usecase/prepared_notification_test.go

@@ -0,0 +1,20 @@
+package usecase
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"go.maximumdirect.net/promptkit/internal/domain"
+)
+
+func TestPreparedNotificationChecksContextBeforeInvocation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	called := false
+	_, err := (&Runner{}).executePreparedRun(ctx, &domain.PreparedRun{}, func(context.Context, *domain.PreparedRun) { called = true }, false, nil, "", 0, "", time.Time{}, nil)
+	if called || !errors.Is(err, context.Canceled) || !errors.Is(err, ErrLLMGenerate) {
+		t.Fatalf("called=%v error=%v", called, err)
+	}
+}

internal/usecase/runner.go

@@ -166,7 +166,7 @@ func (r *Runner) Run(ctx context.Context, req domain.RunRequest) (*domain.RunRes
 	}
 	directAPIKey := state.effectiveModel.APIKey
 
-	return r.executePreparedRun(ctx, operation.run, operation.capture, operation.parent, directAPIKey, state.transientRetries, runID, start, func(
+	return r.executePreparedRun(ctx, operation.run, req.OnPrepared, operation.capture, operation.parent, directAPIKey, state.transientRetries, runID, start, func(
 		ctx context.Context,
 		artifact *domain.Artifact,
 		attemptsUsed int,
@@ -186,6 +186,7 @@ type preparedValidationFunc func(
 func (r *Runner) executePreparedRun(
 	ctx context.Context,
 	prepared *domain.PreparedRun,
+	onPrepared func(context.Context, *domain.PreparedRun),
 	capture bool,
 	parent *domain.Exchange,
 	directAPIKey string,
@@ -194,6 +195,15 @@ func (r *Runner) executePreparedRun(
 	start time.Time,
 	validateArtifact preparedValidationFunc,
 ) (*domain.RunResult, error) {
+	if onPrepared != nil {
+		if err := ctx.Err(); err != nil {
+			return nil, wrapGenerationError(err)
+		}
+		onPrepared(ctx, prepared)
+		if err := ctx.Err(); err != nil {
+			return nil, wrapGenerationError(err)
+		}
+	}
 	executionTarget := prepared.EffectiveModelParams
 	executionTarget.APIKey = directAPIKey
 	genResp, err := r.llm.Generate(ctx, newGenerationRequest(

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)
+	}
+}