corago/codex-keep/policy_test.go

package main

import (
	"encoding/json"
	"math"
	"testing"
	"time"
)

func testGoal(status goalStatus) *goal {
	return &goal{Objective: "Synthetic goal", Status: status, CreatedAt: 1, UpdatedAt: 1}
}

func testTurn(status string) turn {
	t := turn{ID: "turn-1", Status: status}
	if status == "failed" {
		t.Error = &turnError{errorInfo{Kind: "serverOverloaded"}}
	}
	return t
}

func activeWatch() *watch {
	w := &watch{}
	w.update(testGoal(active))
	w.observe(testTurn("inProgress"))
	return w
}

func failedWatch() *watch {
	w := activeWatch()
	w.update(testGoal(blocked))
	w.observe(testTurn("failed"))
	return w
}

func TestRecoveryPolicy(t *testing.T) {
	for _, tc := range []struct {
		name    string
		initial goalStatus
		final   goalStatus
		turn    string
		want    bool
	}{
		{"capacity", active, blocked, "failed", true},
		{"human stop", active, paused, "interrupted", false},
		{"interrupt alone", active, active, "interrupted", false},
		{"paused failure", active, paused, "failed", false},
		{"old blocked", blocked, blocked, "failed", false},
		{"success", active, active, "completed", false},
		{"budget", active, "budgetLimited", "failed", false},
		{"usage", active, "usageLimited", "failed", false},
		{"complete", active, "complete", "failed", false},
	} {
		t.Run(tc.name, func(t *testing.T) {
			w := &watch{}
			w.update(testGoal(tc.initial))
			w.observe(testTurn("inProgress"))
			w.update(testGoal(tc.final))
			w.observe(testTurn(tc.turn))
			if w.eligible() != tc.want {
				t.Fatalf("eligible = %v", w.eligible())
			}
		})
	}
}

func TestGoalChanges(t *testing.T) {
	for _, g := range []*goal{nil, testGoal(paused), testGoal("budgetLimited"), testGoal("usageLimited"), testGoal("complete"), {Objective: "Replacement", Status: blocked, CreatedAt: 2}} {
		w := failedWatch()
		w.schedule(time.Now(), time.Second, time.Minute, func(d time.Duration) time.Duration { return d })
		w.update(g)
		if w.eligible() || !w.Due.IsZero() {
			t.Fatalf("retained retry after %+v", g)
		}
	}
}

func TestCompletionOrder(t *testing.T) {
	for _, failedFirst := range []bool{true, false} {
		w := activeWatch()
		if failedFirst {
			w.observe(testTurn("failed"))
		}
		w.update(testGoal(blocked))
		if !failedFirst {
			w.observe(testTurn("failed"))
		}
		if !w.eligible() {
			t.Fatalf("lost failure, failedFirst=%v", failedFirst)
		}
	}
}

func TestSchedule(t *testing.T) {
	w := failedWatch()
	now := time.Unix(100, 0)
	identity := func(d time.Duration) time.Duration { return d }
	if !w.schedule(now, time.Second, time.Minute, identity) {
		t.Fatal("not scheduled")
	}
	due := w.Due
	if w.schedule(now, time.Second, time.Minute, identity) || w.Due != due || w.Failures != 1 {
		t.Fatal("duplicate changed backoff")
	}
	w.update(testGoal(active))
	w.observe(turn{ID: "turn-2", Status: "inProgress"})
	w.update(testGoal(blocked))
	failure := testTurn("failed")
	failure.ID = "turn-2"
	w.observe(failure)
	w.schedule(now, time.Second, time.Minute, identity)
	if w.Failures != 2 || w.Due.Sub(now) != 2*time.Second {
		t.Fatalf("backoff = %+v", w)
	}
	w.update(testGoal(active))
	w.observe(turn{ID: "turn-3", Status: "inProgress"})
	w.observe(turn{ID: "turn-3", Status: "completed"})
	if w.Failures != 0 {
		t.Fatal("success did not reset backoff")
	}
}

func TestSnapshotReset(t *testing.T) {
	w := failedWatch()
	w.Failures = 3
	w.Due = time.Now()
	w.update(testGoal(active))
	w.observe(turn{ID: "turn-2", Status: "completed"})
	if w.Failures != 0 || !w.Due.IsZero() || w.eligible() {
		t.Fatalf("successful snapshot retained recovery state: %+v", w)
	}
}

func TestErrorTypes(t *testing.T) {
	for _, tc := range []struct {
		raw  string
		want bool
	}{
		{`"serverOverloaded"`, true},
		{`"rateLimitExceeded"`, true},
		{`"internalServerError"`, true},
		{`{"responseTooManyFailedAttempts":{"httpStatusCode":503}}`, true},
		{`{"responseStreamDisconnected":{"httpStatusCode":null}}`, true},
		{`{"httpConnectionFailed":{"httpStatusCode":401}}`, false},
		{`"unauthorized"`, false},
		{`"badRequest"`, false},
		{`"misalignmentPolicyViolation"`, false},
		{`"usageLimitExceeded"`, false},
		{`"other"`, false},
		{`{"futureError":{"httpStatusCode":503}}`, false},
		{`null`, false},
	} {
		var info errorInfo
		if err := json.Unmarshal([]byte(tc.raw), &info); err != nil {
			t.Fatal(err)
		}
		if info.retryable() != tc.want {
			t.Fatalf("%s retryable = %v", tc.raw, info.retryable())
		}
		data, err := json.Marshal(info)
		if err != nil {
			t.Fatal(err)
		}
		var restored errorInfo
		if err := json.Unmarshal(data, &restored); err != nil {
			t.Fatal(err)
		}
		if restored.retryable() != tc.want {
			t.Fatalf("lost error on checkpoint: %s", data)
		}
	}
}

func TestBackoff(t *testing.T) {
	for _, tc := range []struct {
		count int
		want  time.Duration
	}{{1, time.Second}, {2, 2 * time.Second}, {3, 4 * time.Second}, {50, 5 * time.Second}} {
		if got := backoff(time.Second, 5*time.Second, tc.count); got != tc.want {
			t.Fatalf("retry %d = %s", tc.count, got)
		}
	}
	if backoff(time.Duration(math.MaxInt64/2+1), time.Duration(math.MaxInt64), 3) != time.Duration(math.MaxInt64) {
		t.Fatal("duration overflow")
	}
}