Retain execution controls in prepared handles

docs/roadmap/implementation.md

@@ -218,6 +218,8 @@ errors while the operation child is active must not gain the marker.
 
 ## Stage 3: Retain Controls Across Split Preparation
 
+Status: complete.
+
 Outcome: prepared execution carries the selected duration and notification,
 without keeping a live operation scope while idle.
 

execution_timeout_contract_test.go

@@ -38,3 +38,44 @@ func TestRunProviderDeadlineIsNotOperationTimeout(t *testing.T) {
 		t.Fatalf("provider deadline attribution: %v", err)
 	}
 }
+
+func TestPreparedExecutionRetainsTimeoutWithoutIdleTimer(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		engine := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
+			<-ctx.Done()
+			return nil, ctx.Err()
+		}))
+		req := promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second}
+		handle, err := engine.PrepareExecution(context.Background(), req)
+		if err != nil {
+			t.Fatal(err)
+		}
+		req.OperationTimeout = 0
+		time.Sleep(2 * time.Second)
+		start := time.Now()
+		_, err = engine.RunPrepared(context.Background(), handle)
+		if time.Since(start) != time.Second || !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, context.DeadlineExceeded) || !errors.Is(err, promptkit.ErrLLMGenerate) {
+			t.Fatalf("elapsed=%v error=%v", time.Since(start), err)
+		}
+	})
+}
+
+func TestPreparedExecutionRespectsCallerDeadline(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		engine := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
+			<-ctx.Done()
+			return nil, ctx.Err()
+		}))
+		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+		defer cancel()
+		handle, err := engine.PrepareExecution(ctx, promptkit.RunRequest{PromptID: "prompt", OperationTimeout: 5 * time.Second})
+		if err != nil {
+			t.Fatal(err)
+		}
+		time.Sleep(time.Second)
+		_, err = engine.RunPrepared(ctx, handle)
+		if !errors.Is(err, context.DeadlineExceeded) || errors.Is(err, promptkit.ErrOperationTimeout) {
+			t.Fatalf("caller deadline attribution: %v", err)
+		}
+	})
+}

internal/usecase/prepared_execution.go

@@ -29,6 +29,8 @@ type PreparedExecution struct {
 }
 
 type preparedExecutionPayload struct {
+	operationTimeout time.Duration
+	onPrepared       func(context.Context, *domain.PreparedRun)
 	prepared         *domain.PreparedRun
 	capture          bool
 	parent           *domain.Exchange
@@ -56,6 +58,8 @@ func (r *Runner) PrepareExecution(ctx context.Context, req domain.RunRequest) (*
 		state:   preparedExecutionReady,
 		details: operation.run,
 		payload: &preparedExecutionPayload{
+			operationTimeout: req.OperationTimeout,
+			onPrepared:       req.OnPrepared,
 			prepared:         operation.run,
 			capture:          operation.capture,
 			parent:           operation.parent,
@@ -100,12 +104,16 @@ func (p *PreparedExecution) Discard() {
 }
 
 // RunPrepared claims and executes one prepared execution owned by this runner.
-func (r *Runner) RunPrepared(ctx context.Context, prepared *PreparedExecution) (*domain.RunResult, error) {
+func (r *Runner) RunPrepared(ctx context.Context, prepared *PreparedExecution) (result *domain.RunResult, err error) {
 	payload, err := prepared.claim(r)
 	if err != nil {
 		return nil, err
 	}
 	defer payload.clear()
+	scope := NewOperationScope(ctx, payload.operationTimeout)
+	defer scope.Close()
+	defer func() { err = scope.Finalize(err) }()
+	ctx = scope.Context
 
 	var recorder *observation.Recorder
 	var runID string
@@ -185,6 +193,8 @@ func (p *preparedExecutionPayload) clear() {
 	p.directKey = ""
 	p.transientRetries = 0
 	p.observer = nil
+	p.operationTimeout = 0
+	p.onPrepared = nil
 }
 
 type noOpPreparedValidation struct {