package streamgate import ( "context" "errors" "io" "slices" "testing" "time" ) func mustRecoveryContributor(t *testing.T) RecoveryContributor { t.Helper() contributor, err := NewRecoveryContributor("consumer.one", "filter.one", "rule.one") if err != nil { t.Fatalf("NewRecoveryContributor: %v", err) } return contributor } func mustRecoveryPolicy(t *testing.T, total, perStrategy int) RecoveryPolicySnapshot { t.Helper() policy, err := NewRecoveryPolicySnapshot(total, map[RecoveryStrategy]int{ RecoveryStrategyExactReplay: perStrategy, RecoveryStrategyContinuationRepair: perStrategy, RecoveryStrategySchemaRepair: perStrategy, }) if err != nil { t.Fatalf("NewRecoveryPolicySnapshot: %v", err) } return policy } func mustRecoveryUsage(t *testing.T, total int, byStrategy map[RecoveryStrategy]int) RecoveryUsageSnapshot { t.Helper() usage, err := NewRecoveryUsageSnapshot(total, byStrategy) if err != nil { t.Fatalf("NewRecoveryUsageSnapshot: %v", err) } return usage } func mustManagedBudget( t *testing.T, attemptCap, contextWindow, reserve, rebuiltPrompt int, callerCap *int, callerUsed int, ) ManagedTrajectoryBudget { t.Helper() budget, err := NewManagedTrajectoryBudget( attemptCap, contextWindow, reserve, rebuiltPrompt, callerCap, callerUsed, ) if err != nil { t.Fatalf("NewManagedTrajectoryBudget: %v", err) } return budget } func recoveryPlanInput(t *testing.T, strategy RecoveryStrategy, commit CommitState) RecoveryEligibilityInput { t.Helper() var directive RecoveryDirective var err error switch strategy { case RecoveryStrategyExactReplay: directive, err = NewRecoveryDirectiveExact("request.one") case RecoveryStrategySchemaRepair: directive, err = NewRecoveryDirectiveSchema("schema.one", "patch.one") case RecoveryStrategyContinuationRepair, RecoveryStrategyManagedContinuation: directive, err = NewRecoveryDirectiveContinuation(17, "snapshot.one") default: t.Fatalf("unsupported recovery test strategy %q", strategy) } if err != nil { t.Fatalf("new directive: %v", err) } intent, err := NewRecoveryIntent(strategy, directive, "reason.one", 10) if err != nil { t.Fatalf("NewRecoveryIntent: %v", err) } input := RecoveryEligibilityInput{ PlanID: "plan.one", IdempotencyKey: "request.one:cycle.one", Intent: intent, Contributors: []RecoveryContributor{mustRecoveryContributor(t)}, CommitState: commit, Policy: mustRecoveryPolicy(t, 3, 3), Usage: mustRecoveryUsage(t, 0, nil), RequiredCapabilities: []string{"stream.normalized", "assistant.prefill"}, } if strategy == RecoveryStrategyManagedContinuation { budget := mustManagedBudget(t, 64, 1_000, 100, 700, nil, 0) input.ManagedTrajectory = &budget input.TerminalReason = TerminalReasonOutputCap } return input } func TestManagedContinuationStrategyAndTerminalMetadata(t *testing.T) { directive, err := NewRecoveryDirectiveContinuation(9, "snapshot.ref") if err != nil { t.Fatalf("NewRecoveryDirectiveContinuation: %v", err) } if !directive.MatchStrategy(RecoveryStrategyManagedContinuation) { t.Fatal("continuation directive does not match managed continuation") } intent, err := NewRecoveryIntent( RecoveryStrategyManagedContinuation, directive, "provider.output_cap", 3, ) if err != nil { t.Fatalf("NewRecoveryIntent(managed): %v", err) } if err := intent.Validate(); err != nil { t.Fatalf("managed intent Validate: %v", err) } ts := time.Date(2026, 7, 26, 0, 0, 0, 0, time.UTC) event, err := NewTerminalEventWithReason("ch", TerminalReasonOutputCap, ts) if err != nil { t.Fatalf("NewTerminalEventWithReason: %v", err) } reason, err := event.TerminalReason() if err != nil { t.Fatalf("TerminalReason: %v", err) } if reason != TerminalReasonOutputCap { t.Fatalf("TerminalReason = %q, want %q", reason, TerminalReasonOutputCap) } metadata, err := event.AsTerminalMetadata() if err != nil { t.Fatalf("AsTerminalMetadata: %v", err) } if metadata.Reason() != TerminalReasonOutputCap { t.Fatalf("metadata reason = %q", metadata.Reason()) } defaultEvent, err := NewTerminalEvent("ch", ts) if err != nil { t.Fatalf("NewTerminalEvent: %v", err) } defaultReason, err := defaultEvent.TerminalReason() if err != nil || defaultReason != TerminalReasonCompleted { t.Fatalf("default terminal reason = %q, err=%v", defaultReason, err) } if _, err := NewTerminalEventWithReason("ch", TerminalReason("length"), ts); err == nil { t.Fatal("raw provider terminal reason was accepted") } textEvent, _ := NewTextDeltaEvent("ch", "text", ts) if _, err := textEvent.TerminalReason(); err == nil { t.Fatal("non-terminal event returned terminal metadata") } } func TestRecoveryPolicyAndUsageValidationMatrix(t *testing.T) { for _, total := range []int{0, 1, 3} { t.Run("total_valid_"+string(rune('0'+total)), func(t *testing.T) { policy, err := NewRecoveryPolicySnapshot(total, nil) if err != nil { t.Fatalf("NewRecoveryPolicySnapshot(%d): %v", total, err) } for _, strategy := range faultRecoveryStrategies { if got := policy.StrategyLimit(strategy); got != total { t.Fatalf("StrategyLimit(%s) = %d, want %d", strategy, got, total) } } }) } if _, err := NewRecoveryPolicySnapshot(4, nil); err == nil { t.Fatal("request-total cap 4 was accepted") } if _, err := NewRecoveryPolicySnapshot(1, map[RecoveryStrategy]int{ RecoveryStrategyExactReplay: 2, }); err == nil { t.Fatal("strategy cap above request-total cap was accepted") } if _, err := NewRecoveryPolicySnapshot(3, map[RecoveryStrategy]int{ RecoveryStrategyManagedContinuation: 1, }); err == nil { t.Fatal("managed continuation was accepted in fault policy") } if _, err := NewRecoveryUsageSnapshot(1, nil); err == nil { t.Fatal("unaccounted request usage was accepted") } if _, err := NewRecoveryUsageSnapshot(1, map[RecoveryStrategy]int{ RecoveryStrategyManagedContinuation: 1, }); err == nil { t.Fatal("managed continuation was accepted in fault usage") } limits := map[RecoveryStrategy]int{RecoveryStrategyExactReplay: 1} policy, err := NewRecoveryPolicySnapshot(3, limits) if err != nil { t.Fatalf("NewRecoveryPolicySnapshot(partial): %v", err) } limits[RecoveryStrategyExactReplay] = 3 if got := policy.StrategyLimit(RecoveryStrategyExactReplay); got != 1 { t.Fatalf("policy changed after input map mutation: %d", got) } returned := policy.StrategyLimits() returned[RecoveryStrategyExactReplay] = 3 if got := policy.StrategyLimit(RecoveryStrategyExactReplay); got != 1 { t.Fatalf("policy changed after accessor map mutation: %d", got) } } func TestRecoveryPlanCommitEligibilityMatrix(t *testing.T) { tests := []struct { name string strategy RecoveryStrategy commit CommitState wantMode RecoveryResumeMode wantErr error }{ {"exact_uncommitted", RecoveryStrategyExactReplay, CommitStateTransportUncommitted, RecoveryResumeModeReplaceAttempt, nil}, {"schema_uncommitted", RecoveryStrategySchemaRepair, CommitStateTransportUncommitted, RecoveryResumeModeReplaceAttempt, nil}, {"continuation_open", RecoveryStrategyContinuationRepair, CommitStateStreamOpen, RecoveryResumeModeContinueStream, nil}, {"managed_open", RecoveryStrategyManagedContinuation, CommitStateStreamOpen, RecoveryResumeModeContinueStream, nil}, {"exact_open", RecoveryStrategyExactReplay, CommitStateStreamOpen, "", ErrRecoveryCommitIneligible}, {"schema_open", RecoveryStrategySchemaRepair, CommitStateStreamOpen, "", ErrRecoveryCommitIneligible}, {"continuation_uncommitted", RecoveryStrategyContinuationRepair, CommitStateTransportUncommitted, "", ErrRecoveryCommitIneligible}, {"managed_uncommitted", RecoveryStrategyManagedContinuation, CommitStateTransportUncommitted, "", ErrRecoveryCommitIneligible}, {"terminal_committed", RecoveryStrategyContinuationRepair, CommitStateTerminalCommitted, "", ErrRecoveryTerminalCommitted}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { plan, err := NewRecoveryPlan(recoveryPlanInput(t, tt.strategy, tt.commit)) if tt.wantErr != nil { if !errors.Is(err, tt.wantErr) { t.Fatalf("error = %v, want %v", err, tt.wantErr) } return } if err != nil { t.Fatalf("NewRecoveryPlan: %v", err) } if plan.ResumeMode() != tt.wantMode { t.Fatalf("ResumeMode = %q, want %q", plan.ResumeMode(), tt.wantMode) } if err := plan.Validate(); err != nil { t.Fatalf("Validate: %v", err) } }) } } func TestRecoveryPlanEligibilityGuards(t *testing.T) { tests := []struct { name string mutate func(*RecoveryEligibilityInput) wantErr error }{ {"caller_cancel", func(in *RecoveryEligibilityInput) { in.CallerCanceled = true }, ErrRecoveryCallerCanceled}, {"tool_side_effect", func(in *RecoveryEligibilityInput) { in.HasToolSideEffect = true }, ErrRecoverySideEffect}, {"complete_tool_call", func(in *RecoveryEligibilityInput) { in.HasCompleteToolCall = true }, ErrRecoverySideEffect}, {"same_plan_reentry", func(in *RecoveryEligibilityInput) { in.PreviouslyUsedPlanIDs = []string{"other", in.PlanID} }, ErrRecoveryPlanReentry}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { input := recoveryPlanInput(t, RecoveryStrategyContinuationRepair, CommitStateStreamOpen) tt.mutate(&input) if _, err := NewRecoveryPlan(input); !errors.Is(err, tt.wantErr) { t.Fatalf("error = %v, want %v", err, tt.wantErr) } }) } } func TestRecoveryFaultBudgetTotalPrecedesStrategyAndConsumesAtDispatch(t *testing.T) { input := recoveryPlanInput(t, RecoveryStrategyExactReplay, CommitStateTransportUncommitted) input.Usage = mustRecoveryUsage(t, 3, map[RecoveryStrategy]int{ RecoveryStrategyExactReplay: 1, RecoveryStrategyContinuationRepair: 1, RecoveryStrategySchemaRepair: 1, }) if _, err := NewRecoveryPlan(input); !errors.Is(err, ErrRecoveryRequestBudgetExhausted) { t.Fatalf("total-exhausted error = %v", err) } input = recoveryPlanInput(t, RecoveryStrategyExactReplay, CommitStateTransportUncommitted) input.Policy = mustRecoveryPolicy(t, 3, 1) input.Usage = mustRecoveryUsage(t, 1, map[RecoveryStrategy]int{ RecoveryStrategyExactReplay: 1, }) if _, err := NewRecoveryPlan(input); !errors.Is(err, ErrRecoveryStrategyBudgetExhausted) { t.Fatalf("strategy-exhausted error = %v", err) } input = recoveryPlanInput(t, RecoveryStrategyExactReplay, CommitStateTransportUncommitted) originalUsage := input.Usage plan, err := NewRecoveryPlan(input) if err != nil { t.Fatalf("NewRecoveryPlan: %v", err) } if originalUsage.RequestFaultRecoveries() != 0 || plan.UsageSnapshot().RequestFaultRecoveries() != 0 { t.Fatal("plan creation consumed fault budget") } if plan.RequestAttempt() != 1 || plan.StrategyAttempt() != 1 { t.Fatalf("attempt ordinals = request %d strategy %d", plan.RequestAttempt(), plan.StrategyAttempt()) } after, err := plan.UsageAfterDispatch() if err != nil { t.Fatalf("UsageAfterDispatch: %v", err) } if after.RequestFaultRecoveries() != 1 || after.StrategyUsage(RecoveryStrategyExactReplay) != 1 { t.Fatalf("usage after dispatch = total %d exact %d", after.RequestFaultRecoveries(), after.StrategyUsage(RecoveryStrategyExactReplay)) } if plan.UsageSnapshot().RequestFaultRecoveries() != 0 { t.Fatal("UsageAfterDispatch mutated plan snapshot") } zero := recoveryPlanInput(t, RecoveryStrategyExactReplay, CommitStateTransportUncommitted) zero.Policy = mustRecoveryPolicy(t, 0, 0) if _, err := NewRecoveryPlan(zero); !errors.Is(err, ErrRecoveryRequestBudgetExhausted) { t.Fatalf("zero-policy error = %v", err) } } func TestRecoveryRequestSnapshotRefBounds(t *testing.T) { ref, err := NewRecoveryRequestSnapshotRef("snapshot.ref", 100, 200, 500) if err != nil { t.Fatalf("NewRecoveryRequestSnapshotRef: %v", err) } if ref.SnapshotRef() != "snapshot.ref" || ref.RetainedBytes() != 100 || ref.PeakBytes() != 200 || ref.MaxBytes() != 500 { t.Fatalf("getters mismatch: %+v", ref) } if _, err := NewRecoveryRequestSnapshotRef("snapshot.ref", 300, 200, 500); !errors.Is(err, ErrRecoverySnapshotLimitExceeded) { t.Fatalf("retained > peak error = %v", err) } if _, err := NewRecoveryRequestSnapshotRef("snapshot.ref", 100, 600, 500); !errors.Is(err, ErrRecoverySnapshotLimitExceeded) { t.Fatalf("peak > max error = %v", err) } if _, err := NewRecoveryRequestSnapshotRef("snapshot.ref", 100, 200, 0); err == nil { t.Fatal("zero max bytes accepted") } } func TestManagedTrajectoryPostRebuildFinalization(t *testing.T) { callerCap := 500 tests := []struct { name string baseBudget ManagedTrajectoryBudget rebuiltPrompt int wantAllowance int wantFinalErr error }{ {"attempt_cap", mustManagedBudget(t, 64, 2_000, 100, 0, nil, 600), 700, 64, nil}, {"caller_remaining", mustManagedBudget(t, 300, 2_000, 100, 0, &callerCap, 250), 700, 250, nil}, {"context_remaining", mustManagedBudget(t, 300, 1_000, 100, 0, nil, 600), 700, 200, nil}, {"caller_exhausted", mustManagedBudget(t, 300, 2_000, 100, 0, &callerCap, 500), 700, 0, ErrManagedTrajectoryExhausted}, {"context_exhausted", mustManagedBudget(t, 300, 800, 100, 0, nil, 0), 700, 0, ErrManagedTrajectoryExhausted}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { input := recoveryPlanInput(t, RecoveryStrategyManagedContinuation, CommitStateStreamOpen) input.ManagedTrajectory = &tt.baseBudget provisional, err := NewRecoveryPlan(input) if err != nil { t.Fatalf("NewRecoveryPlan: %v", err) } if provisional.ReadyForDispatch() { t.Fatal("provisional managed plan must not be ready for dispatch") } if provisional.ManagedAllowance() != 0 { t.Fatalf("provisional managed allowance = %d, want 0", provisional.ManagedAllowance()) } if _, err := provisional.UsageAfterDispatch(); !errors.Is(err, ErrRecoveryPlanNotReady) { t.Fatalf("provisional UsageAfterDispatch error = %v, want ErrRecoveryPlanNotReady", err) } draft, err := NewRebuiltRequestDraftWithIdempotency( provisional.PlanID(), provisional.IdempotencyKey(), "req.ref", "chat", "openai", 10, 20, 100, tt.rebuiltPrompt, provisional.RequiredCapabilities(), ) if err != nil { t.Fatalf("NewRebuiltRequestDraftWithIdempotency: %v", err) } recording := &recordingDispatcher{} finalPlan, req, err := provisional.FinalizeRebuiltRequest(draft) if tt.wantFinalErr != nil { if !errors.Is(err, tt.wantFinalErr) { t.Fatalf("FinalizeRebuiltRequest error = %v, want %v", err, tt.wantFinalErr) } if recording.calls != 0 { t.Fatalf("exhausted trajectory dispatched request: %d", recording.calls) } return } if err != nil { t.Fatalf("FinalizeRebuiltRequest: %v", err) } if !finalPlan.ReadyForDispatch() { t.Fatal("finalized managed plan should be ready for dispatch") } if finalPlan.ManagedAllowance() != tt.wantAllowance { t.Fatalf("final ManagedAllowance = %d, want %d", finalPlan.ManagedAllowance(), tt.wantAllowance) } if req.ManagedAllowance() != tt.wantAllowance { t.Fatalf("request ManagedAllowance = %d, want %d", req.ManagedAllowance(), tt.wantAllowance) } if provisional.ManagedAllowance() != 0 || provisional.ReadyForDispatch() { t.Fatal("FinalizeRebuiltRequest mutated provisional plan") } binding, err := recording.DispatchAttempt(context.Background(), req) if err != nil { t.Fatalf("recording DispatchAttempt: %v", err) } if recording.calls != 1 || recording.dispatchedRequest.ManagedAllowance() != tt.wantAllowance { t.Fatalf("dispatcher observed wrong request/allowance: %d, %d", recording.calls, recording.dispatchedRequest.ManagedAllowance()) } if !slices.Equal(recording.dispatchedRequest.RequiredCapabilities(), provisional.RequiredCapabilities()) { t.Fatalf("dispatcher observed wrong capabilities: %v, want %v", recording.dispatchedRequest.RequiredCapabilities(), provisional.RequiredCapabilities()) } if recording.dispatchedRequest.PlanID() != provisional.PlanID() || recording.dispatchedRequest.IdempotencyKey() != provisional.IdempotencyKey() { t.Fatalf("dispatcher observed wrong plan/idempotency: %s, %s", recording.dispatchedRequest.PlanID(), recording.dispatchedRequest.IdempotencyKey()) } if binding.AttemptID() == "" { t.Fatal("binding missing attempt ID") } after, err := finalPlan.UsageAfterDispatch() if err != nil { t.Fatalf("final UsageAfterDispatch: %v", err) } if after.RequestFaultRecoveries() != 0 { t.Fatalf("managed dispatch changed fault ledger: %d", after.RequestFaultRecoveries()) } }) } // Identity mismatch check input := recoveryPlanInput(t, RecoveryStrategyManagedContinuation, CommitStateStreamOpen) provisional, err := NewRecoveryPlan(input) if err != nil { t.Fatalf("NewRecoveryPlan: %v", err) } mismatchedDraft, err := NewRebuiltRequestDraftWithIdempotency( "other.plan", provisional.IdempotencyKey(), "req.ref", "chat", "openai", 10, 20, 100, 100, provisional.RequiredCapabilities(), ) if err != nil { t.Fatalf("NewRebuiltRequestDraftWithIdempotency: %v", err) } if _, _, err := provisional.FinalizeRebuiltRequest(mismatchedDraft); err == nil { t.Fatal("FinalizeRebuiltRequest accepted mismatched plan ID") } // Fault strategy check (allowance is 0, request is returned and dispatchable) faultInput := recoveryPlanInput(t, RecoveryStrategyContinuationRepair, CommitStateStreamOpen) faultPlan, err := NewRecoveryPlan(faultInput) if err != nil { t.Fatalf("NewRecoveryPlan(fault): %v", err) } faultDraft, err := NewRebuiltRequestDraftWithIdempotency( faultPlan.PlanID(), faultPlan.IdempotencyKey(), "req.ref", "chat", "openai", 10, 20, 100, 100, faultPlan.RequiredCapabilities(), ) if err != nil { t.Fatalf("NewRebuiltRequestDraftWithIdempotency: %v", err) } faultFinalPlan, faultReq, err := faultPlan.FinalizeRebuiltRequest(faultDraft) if err != nil { t.Fatalf("FinalizeRebuiltRequest on fault plan: %v", err) } if faultReq.ManagedAllowance() != 0 { t.Fatalf("fault request ManagedAllowance = %d, want 0", faultReq.ManagedAllowance()) } if !faultFinalPlan.ReadyForDispatch() { t.Fatal("fault plan should be ready for dispatch") } // Preparer-required managed plan lifecycle prepInput := recoveryPlanInput(t, RecoveryStrategyManagedContinuation, CommitStateStreamOpen) prepInput.PreparerID = "preparer.one" prepInput.PreparationSnapshotRef = "prep.ref.one" prepProvisional, err := NewRecoveryPlan(prepInput) if err != nil { t.Fatalf("NewRecoveryPlan(prep): %v", err) } if prepProvisional.ReadyForRebuild() || prepProvisional.ReadyForDispatch() { t.Fatal("unprepared managed plan must not be ready for rebuild or dispatch") } validDraft, err := NewRebuiltRequestDraftWithIdempotency( prepProvisional.PlanID(), prepProvisional.IdempotencyKey(), "req.ref", "chat", "openai", 10, 20, 100, 100, prepProvisional.RequiredCapabilities(), ) if err != nil { t.Fatalf("NewRebuiltRequestDraftWithIdempotency: %v", err) } if _, _, err := prepProvisional.FinalizeRebuiltRequest(validDraft); !errors.Is(err, ErrRecoveryPlanNotReady) { t.Fatalf("unprepared FinalizeRebuiltRequest error = %v, want ErrRecoveryPlanNotReady", err) } directive, err := NewRecoveryDirectiveContinuation(18, "snapshot.prepared") if err != nil { t.Fatalf("NewRecoveryDirectiveContinuation: %v", err) } preparedProvisional, err := prepProvisional.WithPreparedDirective(directive) if err != nil { t.Fatalf("WithPreparedDirective: %v", err) } if !preparedProvisional.ReadyForRebuild() { t.Fatal("prepared managed plan must be ready for rebuild") } if preparedProvisional.ReadyForDispatch() { t.Fatal("prepared managed plan must not be ready for dispatch before measurement") } finalPrepPlan, prepReq, err := preparedProvisional.FinalizeRebuiltRequest(validDraft) if err != nil { t.Fatalf("FinalizeRebuiltRequest after prep: %v", err) } if !finalPrepPlan.ReadyForDispatch() { t.Fatal("prepared and measured managed plan must be ready for dispatch") } if prepReq.ManagedAllowance() <= 0 { t.Fatalf("prepReq ManagedAllowance = %d, want > 0", prepReq.ManagedAllowance()) } } func TestRecoveryPlanDefensiveCopiesAndPreparationContract(t *testing.T) { input := recoveryPlanInput(t, RecoveryStrategyContinuationRepair, CommitStateStreamOpen) input.PreparerID = "preparer.one" input.PreparationSnapshotRef = "prepare.snapshot.one" cause, err := NewFailureCause("recovery", "retryable", "consumer.one", "filter.one", "rule.one") if err != nil { t.Fatalf("NewFailureCause: %v", err) } input.FailureCauses, err = NewFailureCauseChain([]FailureCause{cause}) if err != nil { t.Fatalf("NewFailureCauseChain: %v", err) } plan, err := NewRecoveryPlan(input) if err != nil { t.Fatalf("NewRecoveryPlan: %v", err) } if !plan.RequiresPreparation() || plan.PreparerID() != "preparer.one" { t.Fatal("preparation contract missing") } if plan.PreparationStatus() != RecoveryPreparationRequired || plan.ReadyForRebuild() { t.Fatal("new prepared-required plan was dispatchable") } if _, err := plan.UsageAfterDispatch(); !errors.Is(err, ErrRecoveryPlanNotReady) { t.Fatalf("unprepared UsageAfterDispatch error = %v", err) } preparedDirective, err := NewRecoveryDirectiveContinuation(18, "snapshot.prepared") if err != nil { t.Fatalf("NewRecoveryDirectiveContinuation(prepared): %v", err) } preparedPlan, err := plan.WithPreparedDirective(preparedDirective) if err != nil { t.Fatalf("WithPreparedDirective: %v", err) } if preparedPlan.PreparationStatus() != RecoveryPreparationPrepared || !preparedPlan.ReadyForRebuild() { t.Fatal("prepared plan was not rebuild-ready") } if plan.PreparationStatus() != RecoveryPreparationRequired { t.Fatal("WithPreparedDirective mutated original plan") } if _, err := preparedPlan.UsageAfterDispatch(); err != nil { t.Fatalf("prepared UsageAfterDispatch: %v", err) } input.Contributors[0] = RecoveryContributor{} input.RequiredCapabilities[0] = "mutated" if plan.Contributors()[0].FilterID() != "filter.one" { t.Fatal("input contributor mutation reached plan") } if plan.RequiredCapabilities()[0] != "stream.normalized" { t.Fatal("input capability mutation reached plan") } contributors := plan.Contributors() contributors[0] = RecoveryContributor{} capabilities := plan.RequiredCapabilities() capabilities[0] = "mutated" causes := plan.FailureCauses() _, _ = causes.Append(cause) if plan.Contributors()[0].FilterID() != "filter.one" || plan.RequiredCapabilities()[0] != "stream.normalized" || plan.FailureCauses().Len() != 1 { t.Fatal("plan accessor leaked mutable state") } unpaired := recoveryPlanInput(t, RecoveryStrategyContinuationRepair, CommitStateStreamOpen) unpaired.PreparerID = "preparer.one" if _, err := NewRecoveryPlan(unpaired); err == nil { t.Fatal("unpaired preparer id was accepted") } } func TestManagedBudgetAndPlanDefensiveCopies(t *testing.T) { callerCap := 500 budget := mustManagedBudget(t, 300, 2_000, 100, 700, &callerCap, 250) callerCap = 250 if capValue, ok := budget.CallerLogicalOutputCap(); !ok || capValue != 500 { t.Fatalf("budget caller cap changed through input pointer: %d, %v", capValue, ok) } input := recoveryPlanInput(t, RecoveryStrategyManagedContinuation, CommitStateStreamOpen) input.ManagedTrajectory = &budget plan, err := NewRecoveryPlan(input) if err != nil { t.Fatalf("NewRecoveryPlan: %v", err) } returned, ok := plan.ManagedTrajectoryBudget() if !ok { t.Fatal("managed trajectory missing") } returned.providerAttemptCap = 1 if plan.ManagedAllowance() != 0 { t.Fatalf("managed plan changed through accessor copy: %d", plan.ManagedAllowance()) } } type recordingDispatcher struct { dispatchedRequest RebuiltRequest calls int } func (d *recordingDispatcher) DispatchAttempt(_ context.Context, request RebuiltRequest) (AttemptBinding, error) { d.dispatchedRequest = request d.calls++ return NewAttemptBinding("attempt.one", "model", "provider", "normalized", recoveryTestEventSource{}, &recoveryTestController{}) } type recoveryTestEventSource struct{} func (recoveryTestEventSource) NextEvent(context.Context) (NormalizedEvent, error) { return NormalizedEvent{}, io.EOF } type recoveryTestController struct{ calls int } func (c *recoveryTestController) AbortAttempt(context.Context) error { c.calls++ return nil } type recoveryTestDispatcher struct{} func (recoveryTestDispatcher) DispatchAttempt(context.Context, RebuiltRequest) (AttemptBinding, error) { return NewAttemptBinding("attempt.one", "model", "provider", "normalized", recoveryTestEventSource{}, &recoveryTestController{}) } type recoveryTestPreparationSnapshot struct { ref string released bool } func (s *recoveryTestPreparationSnapshot) SnapshotRef() string { return s.ref } func (s *recoveryTestPreparationSnapshot) Release() error { s.released = true return nil } type recoveryTestPreparer struct{} func (recoveryTestPreparer) PrepareRecoveryPlan( _ context.Context, plan RecoveryPlan, _ RecoveryPreparationSnapshot, ) (RecoveryDirective, error) { return plan.Directive(), nil } type recoveryTestRebuilder struct{} func (recoveryTestRebuilder) RebuildRequest(_ context.Context, _ RecoveryRequestSnapshotRef, plan RecoveryPlan) (RebuiltRequestDraft, error) { return NewRebuiltRequestDraftWithIdempotency( plan.PlanID(), plan.IdempotencyKey(), "request.rebuilt", "chat", "openai", 8, 8, 16, 700, plan.RequiredCapabilities(), ) } var ( _ NormalizedEventSource = recoveryTestEventSource{} _ AttemptController = (*recoveryTestController)(nil) _ AttemptDispatcher = recoveryTestDispatcher{} _ RecoveryPreparationSnapshot = (*recoveryTestPreparationSnapshot)(nil) _ RecoveryPlanPreparer = recoveryTestPreparer{} _ RequestRebuilder = recoveryTestRebuilder{} ) func TestRebuiltRequestDraftLifecycleAndIdentity(t *testing.T) { capabilities := []string{"stream.normalized"} draft, err := NewRebuiltRequestDraftWithIdempotency( "plan.one", "idemp.one", "request.ref.one", "chat", "openai", 10, 20, 100, 250, capabilities, ) if err != nil { t.Fatalf("NewRebuiltRequestDraftWithIdempotency: %v", err) } if draft.PlanID() != "plan.one" || draft.IdempotencyKey() != "idemp.one" || draft.RequestRef() != "request.ref.one" { t.Fatalf("draft identity mismatch: %s, %s, %s", draft.PlanID(), draft.IdempotencyKey(), draft.RequestRef()) } if draft.Endpoint() != "chat" || draft.Family() != "openai" || draft.RebuiltPromptTokens() != 250 { t.Fatalf("draft metadata mismatch: %s, %s, %d", draft.Endpoint(), draft.Family(), draft.RebuiltPromptTokens()) } if draft.RetainedBytes() != 10 || draft.PeakBytes() != 20 || draft.MaxBytes() != 100 { t.Fatalf("draft bounds mismatch: %d, %d, %d", draft.RetainedBytes(), draft.PeakBytes(), draft.MaxBytes()) } capabilities[0] = "mutated" returnedCaps := draft.RequiredCapabilities() returnedCaps[0] = "also.mutated" if draft.RequiredCapabilities()[0] != "stream.normalized" { t.Fatal("draft capability slice is mutable") } input := recoveryPlanInput(t, RecoveryStrategyContinuationRepair, CommitStateStreamOpen) input.PlanID = "plan.one" input.IdempotencyKey = "idemp.one" input.RequiredCapabilities = []string{"stream.normalized"} plan, err := NewRecoveryPlan(input) if err != nil { t.Fatalf("NewRecoveryPlan: %v", err) } _, req, err := plan.FinalizeRebuiltRequest(draft) if err != nil { t.Fatalf("FinalizeRebuiltRequest: %v", err) } returnedReqCaps := req.RequiredCapabilities() returnedReqCaps[0] = "mutated.req" if req.RequiredCapabilities()[0] != "stream.normalized" { t.Fatal("final request capability slice is mutable") } } func TestRebuiltRequestFinalizationAuthorityAndCapabilities(t *testing.T) { input := recoveryPlanInput(t, RecoveryStrategyManagedContinuation, CommitStateStreamOpen) plan, err := NewRecoveryPlan(input) if err != nil { t.Fatalf("NewRecoveryPlan: %v", err) } draft, err := NewRebuiltRequestDraftWithIdempotency( plan.PlanID(), plan.IdempotencyKey(), "req.ref", "chat", "openai", 10, 20, 100, 100, plan.RequiredCapabilities(), ) if err != nil { t.Fatalf("NewRebuiltRequestDraftWithIdempotency: %v", err) } finalPlan, req, err := plan.FinalizeRebuiltRequest(draft) if err != nil { t.Fatalf("FinalizeRebuiltRequest success: %v", err) } if !slices.Equal(req.RequiredCapabilities(), plan.RequiredCapabilities()) { t.Fatalf("req RequiredCapabilities = %v, want %v", req.RequiredCapabilities(), plan.RequiredCapabilities()) } if req.PlanID() != plan.PlanID() || req.IdempotencyKey() != plan.IdempotencyKey() { t.Fatalf("req identity mismatch: %s, %s", req.PlanID(), req.IdempotencyKey()) } caps := req.RequiredCapabilities() caps[0] = "mutated.cap" if req.RequiredCapabilities()[0] != "stream.normalized" { t.Fatal("final request capability slice is mutable") } recording := &recordingDispatcher{} binding, err := recording.DispatchAttempt(context.Background(), req) if err != nil { t.Fatalf("recording DispatchAttempt: %v", err) } if recording.calls != 1 { t.Fatalf("recording calls = %d, want 1", recording.calls) } if recording.dispatchedRequest.ManagedAllowance() != finalPlan.ManagedAllowance() { t.Fatalf("dispatched ManagedAllowance = %d, want %d", recording.dispatchedRequest.ManagedAllowance(), finalPlan.ManagedAllowance()) } if !slices.Equal(recording.dispatchedRequest.RequiredCapabilities(), plan.RequiredCapabilities()) { t.Fatalf("dispatched RequiredCapabilities = %v, want %v", recording.dispatchedRequest.RequiredCapabilities(), plan.RequiredCapabilities()) } if binding.AttemptID() == "" { t.Fatal("binding missing attempt ID") } mismatches := []struct { name string capabilities []string }{ {"nil_caps", nil}, {"empty_caps", []string{}}, {"missing_one", []string{"stream.normalized"}}, {"reordered", []string{"assistant.prefill", "stream.normalized"}}, {"mutated_extra", []string{"stream.normalized", "assistant.prefill", "extra"}}, {"different", []string{"stream.normalized", "other"}}, } for _, tc := range mismatches { t.Run(tc.name, func(t *testing.T) { mDraft, err := NewRebuiltRequestDraftWithIdempotency( plan.PlanID(), plan.IdempotencyKey(), "req.ref", "chat", "openai", 10, 20, 100, 100, tc.capabilities, ) if err != nil { t.Fatalf("NewRebuiltRequestDraftWithIdempotency(%s): %v", tc.name, err) } rec := &recordingDispatcher{} outPlan, outReq, err := plan.FinalizeRebuiltRequest(mDraft) if err == nil { t.Fatalf("FinalizeRebuiltRequest accepted capability mismatch %s", tc.name) } if outPlan.PlanID() != "" || outReq.PlanID() != "" { t.Fatalf("FinalizeRebuiltRequest error returned non-zero plan/request for %s", tc.name) } if rec.calls != 0 { t.Fatalf("dispatcher called on mismatch %s: %d", tc.name, rec.calls) } }) } } func TestRecoveryHostSeamsAndRawFreeReferences(t *testing.T) { capabilities := []string{"stream.normalized"} draft, err := NewRebuiltRequestDraftWithIdempotency( "plan.one", "request.one:cycle.one", "request.rebuilt", "chat", "openai", 8, 10, 16, 700, capabilities, ) if err != nil { t.Fatalf("NewRebuiltRequestDraftWithIdempotency: %v", err) } capabilities[0] = "mutated" returned := draft.RequiredCapabilities() returned[0] = "also.mutated" if draft.RequiredCapabilities()[0] != "stream.normalized" { t.Fatal("rebuilt request draft capability slice is mutable") } if draft.IdempotencyKey() != "request.one:cycle.one" { t.Fatalf("rebuilt idempotency key = %q", draft.IdempotencyKey()) } if _, err := NewRebuiltRequestDraft("plan.one", "request.rebuilt", "chat", "openai", 8, 17, 16, 700, nil); !errors.Is(err, ErrRecoverySnapshotLimitExceeded) { t.Fatalf("rebuilt peak overflow error = %v", err) } controller := &recoveryTestController{} binding, err := NewAttemptBinding( "attempt.one", "actual-model", "actual-provider", "normalized", recoveryTestEventSource{}, controller, ) if err != nil { t.Fatalf("NewAttemptBinding: %v", err) } if binding.Model() != "actual-model" || binding.Provider() != "actual-provider" || binding.ExecutionPath() != "normalized" { t.Fatalf("attempt binding metadata mismatch: %+v", binding) } if err := binding.Controller().AbortAttempt(context.Background()); err != nil { t.Fatalf("AbortAttempt: %v", err) } if controller.calls != 1 { t.Fatalf("controller calls = %d, want 1", controller.calls) } snapshot := &recoveryTestPreparationSnapshot{ref: "prepare.snapshot.one"} if snapshot.SnapshotRef() != "prepare.snapshot.one" { t.Fatalf("preparation snapshot ref = %q", snapshot.SnapshotRef()) } if err := snapshot.Release(); err != nil { t.Fatalf("preparation snapshot Release: %v", err) } if err := snapshot.Release(); err != nil { t.Fatalf("preparation snapshot second Release: %v", err) } if !snapshot.released { t.Fatal("preparation snapshot was not released") } }