iop/packages/go/streamgate/recovery_plan_test.go

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")
	}
}