Fix goal recovery

codex-keep/usage_test.go

@@ -0,0 +1,105 @@
+package main
+
+import (
+	"testing"
+	"time"
+)
+
+func TestUsagePoll(t *testing.T) {
+	exhausted := false
+	for _, allowed := range []*bool{nil, &exhausted} {
+		f := newFake(t, nil)
+		f.usage = allowed
+		if err := f.reconcile(); err != nil {
+			t.Fatal(err)
+		}
+		w := f.op.store.watch("thread-1")
+		w.Failures = 3
+		f.fail()
+		f.goal.Status = "usageLimited"
+		for poll := 1; poll <= 3; poll++ {
+			w.Due = time.Now().Add(-time.Second)
+			start := time.Now()
+			if err := f.reconcile(); err != nil {
+				t.Fatal(err)
+			}
+			if f.count("account/rateLimits/read") != poll || f.count("thread/goal/set") != 0 || w.Failures != 3 {
+				t.Fatal("usage wait attempted work or increased backoff")
+			}
+			if w.Due.Before(start.Add(f.op.opts.poll)) || w.Due.After(time.Now().Add(f.op.opts.poll)) {
+				t.Fatal("usage poll did not retain the fixed interval")
+			}
+			if err := f.reconcile(); err != nil {
+				t.Fatal(err)
+			}
+			if f.count("account/rateLimits/read") != poll {
+				t.Fatal("usage was queried before the next poll")
+			}
+		}
+		ready := true
+		f.usage = &ready
+		w.Due = time.Now().Add(-time.Second)
+		if err := f.reconcile(); err != nil {
+			t.Fatal(err)
+		}
+		if f.count("thread/goal/set") != 1 || f.goal.Status != active {
+			t.Fatal("early usage reset did not resume work")
+		}
+	}
+}
+
+func TestUsageReadStop(t *testing.T) {
+	for _, status := range []string{"paused", "interrupted"} {
+		t.Run(status, func(t *testing.T) {
+			f := newFake(t, nil)
+			if err := f.reconcile(); err != nil {
+				t.Fatal(err)
+			}
+			f.fail()
+			f.goal.Status = "usageLimited"
+			ready := true
+			f.usage = &ready
+			f.before = func(m message) {
+				if m.Method != "account/rateLimits/read" {
+					return
+				}
+				f.before = nil
+				if status == "interrupted" {
+					f.turn.Status = status
+					f.notify("turn/completed", params{ThreadID: "thread-1", Turn: f.turn})
+				} else {
+					f.goal = testGoal(paused)
+					f.human(f.goal)
+				}
+			}
+			for range 3 {
+				if err := f.reconcile(); err != nil {
+					t.Fatal(err)
+				}
+			}
+			if w := f.op.store.watch("thread-1"); w.eligible() || !w.Due.IsZero() || f.count("thread/goal/set") != 0 {
+				t.Fatal("usage response revoked human cancellation")
+			}
+		})
+	}
+}
+
+func TestUsageReadError(t *testing.T) {
+	f := newFake(t, nil)
+	if err := f.reconcile(); err != nil {
+		t.Fatal(err)
+	}
+	f.fail()
+	f.goal.Status = "usageLimited"
+	f.usageErr = &rpcError{Code: -32603, Message: "synthetic usage error"}
+	for poll := 1; poll <= 2; poll++ {
+		w := f.op.store.watch("thread-1")
+		w.Due = time.Now().Add(-time.Second)
+		if err := f.reconcile(); err != nil {
+			t.Fatal(err)
+		}
+		if f.count("account/rateLimits/read") != poll || f.count("thread/goal/set") != 0 || !w.Due.After(time.Now()) {
+			t.Fatal("usage read error lost polling or attempted work")
+		}
+	}
+}