Agent-Squad/squad/step_retry.go

package squad

import (
	"os"
	"path/filepath"
	"strings"
	"unicode"
)

// MaxTransientStepInvocations is the product-layer bound: the first
// invocation plus one fresh re-invocation. SDK in-process auto_retry is
// separate and must not be confused with this.
const MaxTransientStepInvocations = 2

type FailureClass int

const (
	FailureNone FailureClass = iota
	FailureTransient
	FailurePermanent
)

func (c FailureClass) String() string {
	switch c {
	case FailureTransient:
		return "transient"
	case FailurePermanent:
		return "permanent"
	default:
		return "none"
	}
}

type TraceSignals struct {
	ProviderTransient    bool
	Settled              bool
	RetryExhausted       bool
	HasUsableDeliverable bool // always false: write() args in traces are not deliverables
}

type StepFailureSignals struct {
	Err           error
	Result        *ExecResult
	WorkDir       string
	ExtraText     string
	TargetMissing bool
	// TargetExists means this invocation wrote a non-empty file that differs
	// from the pre-invoke baseline. Do not retry or overwrite it. The step
	// still fails unless the backend already reported a trusted success.
	TargetExists bool
}

type FailureAssessment struct {
	Class  FailureClass
	Reason string
}

func ScanTraceText(text string) TraceSignals {
	var out TraceSignals
	if text == "" {
		return out
	}
	// Line scan only. Never parse write.arguments.content into a body.
	for _, line := range strings.Split(text, "\n") {
		if line == "" {
			continue
		}
		lower := strings.ToLower(line)
		if strings.Contains(lower, `"type":"agent_settled"`) || strings.Contains(lower, `"type": "agent_settled"`) {
			out.Settled = true
		}
		if strings.Contains(lower, `"type":"auto_retry_end"`) && strings.Contains(lower, `"success":false`) {
			out.ProviderTransient = true
			out.RetryExhausted = true
		}
		if strings.Contains(lower, `"type":"auto_retry_start"`) {
			out.ProviderTransient = true
		}
		if strings.Contains(lower, `"stopreason":"error"`) {
			out.ProviderTransient = true
		}
		if strings.Contains(lower, "connection error") {
			out.ProviderTransient = true
		}
		if strings.Contains(lower, `"errormessage":"terminated"`) || strings.Contains(lower, `"errormessage": "terminated"`) {
			out.ProviderTransient = true
		}
	}
	return out
}

func ReadWorkDirTraces(workDir string) string {
	if strings.TrimSpace(workDir) == "" {
		return ""
	}
	p := filepath.Join(workDir, "execution.log")
	b, err := os.ReadFile(p)
	if err != nil {
		return ""
	}
	const capBytes = 512 << 10
	if len(b) > capBytes {
		b = b[len(b)-capBytes:]
	}
	return string(b)
}

func AssessStepFailure(sig StepFailureSignals) FailureAssessment {
	reason := combinedFailureText(sig)
	lower := strings.ToLower(reason)
	trace := ScanTraceText(sig.ExtraText + "\n" + ReadWorkDirTraces(sig.WorkDir))
	missing := sig.TargetMissing || strings.Contains(lower, "missing target deliverable")
	if sig.TargetExists {
		return FailureAssessment{Class: FailurePermanent, Reason: "target-already-present: do not overwrite"}
	}
	if isPermanentReason(lower) {
		return FailureAssessment{Class: FailurePermanent, Reason: "permanent: " + compactReason(reason)}
	}

	runnerFailed := (sig.Result != nil && sig.Result.Retryable) || isRunnerExplicitFailure(lower)
	provider := trace.ProviderTransient || hasProviderTermOrConn(lower)
	if runnerFailed {
		return FailureAssessment{Class: FailureTransient, Reason: "runner-explicit-failure: " + compactReason(reason)}
	}
	if provider && missing {
		kind := "provider-transient-missing-target"
		if trace.Settled {
			kind = "settled-failure-missing-target"
		}
		return FailureAssessment{Class: FailureTransient, Reason: kind + ": " + compactReason(reason)}
	}
	// Model settled cleanly but never wrote the target (e.g. text-only reply,
	// no write tool). One product-layer re-invocation with a corrective prompt
	// is worth trying; still bounded by MaxTransientStepInvocations.
	if missing && trace.Settled && !provider {
		return FailureAssessment{Class: FailureTransient, Reason: "settled-no-deliverable: " + compactReason(reason)}
	}
	if missing {
		return FailureAssessment{Class: FailurePermanent, Reason: "permanent-missing: " + compactReason(reason)}
	}
	return FailureAssessment{Class: FailurePermanent, Reason: "fail-closed: " + compactReason(reason)}
}

func combinedFailureText(sig StepFailureSignals) string {
	var parts []string
	if sig.Err != nil {
		parts = append(parts, sig.Err.Error())
	}
	if sig.Result != nil {
		if sig.Result.ErrorReason != "" {
			parts = append(parts, sig.Result.ErrorReason)
		}
		if sig.Result.Logs != "" {
			parts = append(parts, sig.Result.Logs)
		}
	}
	return strings.Join(parts, " ")
}

func isPermanentReason(lower string) bool {
	return strings.Contains(lower, "prompt_empty") ||
		strings.Contains(lower, "unauthorized") ||
		strings.Contains(lower, "budget_exceeded") ||
		strings.Contains(lower, "治理强制终止") ||
		strings.Contains(lower, "治理拦截")
}

func hasProviderTermOrConn(lower string) bool {
	return strings.Contains(lower, "connection error") ||
		strings.Contains(lower, "terminated") ||
		strings.Contains(lower, "econnreset") ||
		strings.Contains(lower, "socket hang up")
}

func isRunnerExplicitFailure(lower string) bool {
	return strings.Contains(lower, "pi remote runner call failed") ||
		strings.Contains(lower, "auto_retry_failed") ||
		strings.Contains(lower, "auto_retry_end") ||
		(hasProviderTermOrConn(lower) && !strings.Contains(lower, "missing target deliverable"))
}

func compactReason(s string) string {
	s = strings.Join(strings.FieldsFunc(s, func(r rune) bool {
		return r == '\n' || r == '\r'
	}), " ")
	s = strings.TrimSpace(s)
	if s == "" {
		return "unknown"
	}
	// Drop accidental write-argument bodies if a caller stuffed logs here.
	if len(s) > 240 {
		s = s[:240]
	}
	return strings.Map(func(r rune) rune {
		if unicode.IsControl(r) && r != '\t' {
			return -1
		}
		return r
	}, s)
}