Test continuous execution budgets across collaborators

docs/roadmap/implementation.md

@@ -288,6 +288,8 @@ failure assertions belong to Stage 6, not duplicate tests here.
 
 ## Stage 5: Verify Continuous Budgets Across Execution Phases
 
+Status: complete.
+
 Outcome: tests demonstrate that one scope reaches all relevant collaborators,
 with correct attribution and observation on failure.
 

execution_budget_contract_test.go

@@ -0,0 +1,192 @@
+package promptkit_test
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"testing"
+	"testing/synctest"
+	"time"
+
+	"go.maximumdirect.net/promptkit"
+)
+
+func TestOperationBudgetCoversPreparationAndReleasesAdmission(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		release := make(chan struct{})
+		reader := &capacityArtifactReader{entered: make(chan struct{}, 2), release: release}
+		client := newCapacityGateClient(closedCapacityChannel(), 2)
+		engine := newBackendCapacityEngine(t, client, 1, capacityInt(0), reader)
+		parent := context.Background()
+		req := capacityInputRequest("http://example.invalid/v1")
+		req.OperationTimeout = time.Second
+		_, err := engine.Run(parent, req)
+		if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, context.DeadlineExceeded) || !errors.Is(err, promptkit.ErrArtifactLoad) || parent.Err() != nil {
+			t.Fatalf("preparation error=%v parent=%v", err, parent.Err())
+		}
+		close(release)
+		if result, err := engine.Run(parent, req); err != nil || result == nil {
+			t.Fatalf("capacity not released: %v", err)
+		}
+	})
+}
+
+func TestOperationBudgetCoversGenerationCapacityWait(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		release := make(chan struct{})
+		client := newCapacityGateClient(release, 2)
+		engine := newBackendCapacityEngine(t, client, 1, capacityInt(1), nil)
+		results := make(chan capacityRunResult, 1)
+		go runCapacityRequest(engine, context.Background(), promptkit.RunRequest{PromptID: "prompt"}, results)
+		<-client.started
+		_, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second})
+		if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, promptkit.ErrLLMGenerate) || !errors.Is(err, context.DeadlineExceeded) {
+			t.Fatalf("capacity wait error=%v", err)
+		}
+		close(release)
+		if result := <-results; result.err != nil {
+			t.Fatal(result.err)
+		}
+		if _, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt"}); err != nil {
+			t.Fatal(err)
+		}
+	})
+}
+
+func TestOperationBudgetInterruptsRetryAndRepair(t *testing.T) {
+	for _, repair := range []bool{false, true} {
+		t.Run(map[bool]string{false: "retry", true: "repair"}[repair], func(t *testing.T) {
+			synctest.Test(t, func(t *testing.T) {
+				calls := 0
+				config := simpleConfig("message")
+				one := 1
+				config.TransientRetries = &one
+				config.HTTPClient = &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
+					calls++
+					if repair && calls == 1 {
+						return retryHTTPResponse(200, "not json"), nil
+					}
+					if repair {
+						<-r.Context().Done()
+						return nil, r.Context().Err()
+					}
+					response := retryHTTPResponse(503, "")
+					response.Header = make(http.Header)
+					response.Header.Set("Retry-After", "5")
+					return response, nil
+				})}
+				engine, err := promptkit.NewEngine(config)
+				if err != nil {
+					t.Fatal(err)
+				}
+				req := simpleGenerationRequest()
+				req.OperationTimeout = time.Second
+				if repair {
+					req.Validation = &promptkit.OutputContract{Format: promptkit.FormatJSON, ValidationMode: promptkit.ValidationJSON, RepairAttempts: 1}
+				}
+				_, err = engine.Run(context.Background(), req)
+				if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, promptkit.ErrLLMGenerate) || !errors.Is(err, context.DeadlineExceeded) {
+					t.Fatalf("error=%v calls=%d", err, calls)
+				}
+			})
+		})
+	}
+}
+
+func TestOperationOutcomeFrozenBeforeTerminalObserver(t *testing.T) {
+	for _, split := range []bool{false, true} {
+		t.Run(map[bool]string{false: "run", true: "prepared"}[split], func(t *testing.T) {
+			synctest.Test(t, func(t *testing.T) {
+				engine := observationEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
+					return &promptkit.GenerateResponse{Content: "ok"}, nil
+				}))
+				var terminal promptkit.ExecutionEvent
+				req := promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second, Observer: func(event promptkit.ExecutionEvent) {
+					if event.Kind == promptkit.ExecutionFinished {
+						time.Sleep(2 * time.Second)
+						terminal = event
+					}
+				}}
+				var result *promptkit.RunResult
+				var err error
+				if split {
+					var handle *promptkit.PreparedExecution
+					handle, err = engine.PrepareExecution(context.Background(), req)
+					if err == nil {
+						result, err = engine.RunPrepared(context.Background(), handle)
+					}
+				} else {
+					result, err = engine.Run(context.Background(), req)
+				}
+				if err != nil || result == nil || terminal.Outcome != promptkit.ExecutionOutcomeSuccess {
+					t.Fatalf("result=%v error=%v terminal=%+v", result, err, terminal)
+				}
+			})
+		})
+	}
+}
+
+func TestNotificationConsumesOperationBudget(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		engine := observationEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
+			t.Fatal("generation ran")
+			return nil, nil
+		}))
+		_, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second, OnPrepared: func(ctx context.Context, _ promptkit.PreparedRun) { <-ctx.Done() }})
+		if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, promptkit.ErrLLMGenerate) || !errors.Is(err, context.DeadlineExceeded) {
+			t.Fatal(err)
+		}
+	})
+}
+
+func TestOperationTimeoutDoesNotRelabelOtherFailures(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		failure := errors.New("independent failure")
+		engine := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
+			<-ctx.Done()
+			return nil, failure
+		}))
+		var terminal promptkit.ExecutionEvent
+		_, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second, Observer: func(event promptkit.ExecutionEvent) {
+			if event.Kind == promptkit.ExecutionFinished {
+				terminal = event
+			}
+		}})
+		if !errors.Is(err, failure) || errors.Is(err, promptkit.ErrOperationTimeout) || terminal.Failure != promptkit.ExecutionFailureGeneration {
+			t.Fatalf("error=%v terminal=%+v", err, terminal)
+		}
+	})
+}
+
+func TestOperationTimeoutTerminalObserverSeesFinalFailure(t *testing.T) {
+	for _, split := range []bool{false, true} {
+		t.Run(map[bool]string{false: "run", true: "prepared"}[split], func(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()
+				}))
+				var terminal promptkit.ExecutionEvent
+				req := promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second, Observer: func(event promptkit.ExecutionEvent) {
+					if event.Kind == promptkit.ExecutionFinished {
+						time.Sleep(time.Second)
+						terminal = event
+					}
+				}}
+				var err error
+				if split {
+					var handle *promptkit.PreparedExecution
+					handle, err = engine.PrepareExecution(context.Background(), req)
+					if err == nil {
+						_, err = engine.RunPrepared(context.Background(), handle)
+					}
+				} else {
+					_, err = engine.Run(context.Background(), req)
+				}
+				if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, context.DeadlineExceeded) || terminal.Failure != promptkit.ExecutionFailureDeadline {
+					t.Fatalf("error=%v terminal=%+v", err, terminal)
+				}
+			})
+		})
+	}
+}

internal/usecase/operation_budget_test.go

@@ -0,0 +1,54 @@
+package usecase
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"testing/synctest"
+	"time"
+
+	"go.maximumdirect.net/promptkit/internal/domain"
+	"go.maximumdirect.net/promptkit/internal/validate"
+)
+
+type deadlineValidation struct{ deadline time.Time }
+
+func (v *deadlineValidation) PrepareValidation(ctx context.Context, _ domain.OutputContract) (validate.PreparedValidation, error) {
+	v.deadline, _ = ctx.Deadline()
+	return v, nil
+}
+func (v *deadlineValidation) Validate(ctx context.Context, _ *domain.Artifact) (domain.ValidationResult, error) {
+	deadline, _ := ctx.Deadline()
+	if deadline != v.deadline {
+		return domain.ValidationResult{}, errors.New("validation deadline changed")
+	}
+	<-ctx.Done()
+	return domain.ValidationResult{}, ctx.Err()
+}
+func (*deadlineValidation) SchemaDocument() any { return nil }
+
+type deadlineModel struct{ validation *deadlineValidation }
+
+func (m deadlineModel) Generate(ctx context.Context, _ domain.GenerateRequest) (*domain.GenerateResponse, error) {
+	deadline, _ := ctx.Deadline()
+	if deadline != m.validation.deadline {
+		return nil, errors.New("generation deadline changed")
+	}
+	return &domain.GenerateResponse{Content: "ok"}, nil
+}
+
+func TestOperationDeadlineReachesPreparationGenerationAndValidation(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		validation := &deadlineValidation{}
+		runner := NewRunner(&fakePromptRepo{def: promptDef(domain.FormatText, domain.ValidationBasic, 0)},
+			&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
+			defaultArtifactReader(), defaultRenderer(), deadlineModel{validation}, validation, nil)
+		scope := NewOperationScope(context.Background(), time.Second)
+		defer scope.Close()
+		_, err := runner.Run(scope.Context, domain.RunRequest{PromptID: "p", ProfileID: "exec", Inputs: singleInputRef()})
+		err = scope.Finalize(err)
+		if !errors.Is(err, ErrOperationTimeout) || !errors.Is(err, ErrValidation) || !errors.Is(err, context.DeadlineExceeded) {
+			t.Fatalf("validation error=%v", err)
+		}
+	})
+}