Notify callers from the prepared execution snapshot

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(