recovery_orchestrator.go

package claudesdk

import (
	"context"
	"fmt"
	"iter"
	"time"
)

// RecoveryOrchestratorOptions configure host-side recovery planning.
type RecoveryOrchestratorOptions struct {
	RetryPolicy          RetryPolicy
	MaxAttempts          int
	EnableMediaStrip     bool
	EnableFallbackSwitch bool
}

// RecoveryPlan summarizes the next recovery step after a failed result.
type RecoveryPlan struct {
	Attempt        int
	MaxAttempts    int
	Classification ErrorClassification
	Retry          bool
	Delay          time.Duration
	Action         RecoveryAction
	Reason         string
	Content        UserMessageContent
	ContentChanged bool
	OverrideModel  *string
}

// RecoveryOrchestrator plans host-side retries and safe content/model adjustments.
type RecoveryOrchestrator struct {
	retryPolicy          RetryPolicy
	maxAttempts          int
	enableMediaStrip     bool
	enableFallbackSwitch bool
}

// NewRecoveryOrchestrator creates a host-side recovery planner.
func NewRecoveryOrchestrator(options RecoveryOrchestratorOptions) *RecoveryOrchestrator {
	maxAttempts := options.MaxAttempts
	if maxAttempts <= 0 {
		maxAttempts = 3
	}

	return &RecoveryOrchestrator{
		retryPolicy:          normalizeRetryPolicy(options.RetryPolicy),
		maxAttempts:          maxAttempts,
		enableMediaStrip:     options.EnableMediaStrip,
		enableFallbackSwitch: options.EnableFallbackSwitch,
	}
}

// PlanResult returns the next recovery step for a terminal result.
func (o *RecoveryOrchestrator) PlanResult(
	result *ResultMessage,
	content UserMessageContent,
	options *ClaudeAgentOptions,
	attempt int,
) RecoveryPlan {
	plan := RecoveryPlan{
		Attempt:     attempt,
		MaxAttempts: 1,
		Content:     content,
	}

	if o == nil {
		plan.MaxAttempts = 1

		return plan
	}

	plan.MaxAttempts = o.maxAttempts
	if result == nil || !result.IsError {
		return plan
	}

	classification := ClassifyResultError(result)
	plan.Classification = classification
	plan.Action = classification.RecoveryAction
	plan.Reason = string(classification.Class)

	if attempt >= o.maxAttempts {
		return plan
	}

	switch classification.RecoveryAction {
	case RecoveryActionStripMedia:
		if o.enableMediaStrip {
			if stripped, changed := StripMediaContent(content); changed {
				plan.Retry = true
				plan.Content = stripped
				plan.ContentChanged = true
				plan.Reason = "strip_media"

				return plan
			}
		}
	case RecoveryActionSwitchModel:
		if o.enableFallbackSwitch && options != nil && options.FallbackModel != "" && options.FallbackModel != options.Model {
			model := options.FallbackModel
			plan.Retry = true
			plan.OverrideModel = &model
			plan.Reason = "switch_model"

			return plan
		}
	}

	decision := EvaluateResultRetry(result, attempt, o.retryPolicy)
	if decision.Action == RetryDecisionActionRetry && decision.Retryable {
		plan.Retry = true
		plan.Delay = decision.RecommendedDelay
		plan.Reason = decision.Reason
	}

	return plan
}

// ResilientQuery reruns a one-shot query when the recovery orchestrator can
// safely retry the request.
func ResilientQuery(
	ctx context.Context,
	content UserMessageContent,
	orchestrator *RecoveryOrchestrator,
	opts ...Option,
) iter.Seq2[Message, error] {
	return func(yield func(Message, error) bool) {
		if orchestrator == nil {
			for msg, err := range Query(ctx, content, opts...) {
				if !yield(msg, err) {
					return
				}
			}

			return
		}

		currentContent := content

		baseOptions := append([]Option(nil), opts...)

		for attempt := 1; attempt <= orchestrator.maxAttempts; attempt++ {
			var terminal *ResultMessage

			for msg, err := range Query(ctx, currentContent, baseOptions...) {
				if err != nil {
					if !yield(nil, err) {
						return
					}

					return
				}

				if result, ok := msg.(*ResultMessage); ok {
					terminal = result
				}

				if !yield(msg, nil) {
					return
				}
			}

			plan := orchestrator.PlanResult(terminal, currentContent, applyAgentOptions(baseOptions), attempt)
			if !plan.Retry {
				return
			}

			currentContent = plan.Content
			if plan.OverrideModel != nil {
				baseOptions = append(baseOptions, WithModel(*plan.OverrideModel))
			}

			if plan.Delay > 0 {
				timer := time.NewTimer(plan.Delay)
				select {
				case <-ctx.Done():
					timer.Stop()
					yield(nil, ctx.Err())

					return
				case <-timer.C:
				}
			}
		}
	}
}

// StripMediaContent removes inline image and document blocks from block-based
// user content so a retry can proceed without large attachments.
func StripMediaContent(content UserMessageContent) (UserMessageContent, bool) {
	if content.IsString() {
		return content, false
	}

	blocks := content.Blocks()
	if len(blocks) == 0 {
		return content, false
	}

	filtered := make([]ContentBlock, 0, len(blocks))
	changed := false

	for _, block := range blocks {
		switch block.(type) {
		case *InputImageBlock, *InputDocumentBlock:
			changed = true

			continue
		default:
			filtered = append(filtered, block)
		}
	}

	if !changed {
		return content, false
	}

	return Blocks(filtered...), true
}

// MustNewRecoveryOrchestrator is a convenience helper for examples/tests.
func MustNewRecoveryOrchestrator(options RecoveryOrchestratorOptions) *RecoveryOrchestrator {
	return NewRecoveryOrchestrator(options)
}

// Validate ensures the recovery plan is internally consistent.
func (p RecoveryPlan) Validate() error {
	if p.Retry && p.Attempt >= p.MaxAttempts {
		return fmt.Errorf("retry plan exceeds max attempts")
	}

	return nil
}