diff --git a/docs/roadmap/implementation.md b/docs/roadmap/implementation.md index 26430a1..36b41a3 100644 --- a/docs/roadmap/implementation.md +++ b/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. diff --git a/execution_timeout_contract_test.go b/execution_timeout_contract_test.go index 8771163..d9ba559 100644 --- a/execution_timeout_contract_test.go +++ b/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) + } + }) +} diff --git a/internal/usecase/prepared_execution.go b/internal/usecase/prepared_execution.go index 22fa2f5..9d07c4e 100644 --- a/internal/usecase/prepared_execution.go +++ b/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 {