Add daemon-attached goal recovery operator

Resume goals after transient API failures with persisted exponential backoff while respecting pauses and interruptions.

codex-keep/operator_test.go

@@ -0,0 +1,293 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"path/filepath"
+	"testing"
+	"time"
+)
+
+type fakeDaemon struct {
+	op        *operator
+	t         *testing.T
+	goal      *goal
+	turn      turn
+	status    string
+	requests  []message
+	before    func(message)
+	lostReply bool
+}
+
+func newFake(t *testing.T, s *store) *fakeDaemon {
+	t.Helper()
+	if s == nil {
+		var err error
+		s, err = openStore(filepath.Join(t.TempDir(), "state.json"))
+		if err != nil {
+			t.Fatal(err)
+		}
+		t.Cleanup(func() { closeStore(t, s) })
+	}
+	f := &fakeDaemon{t: t, goal: testGoal(active), turn: testTurn("inProgress"), status: "systemError"}
+	c := &connection{frames: make(chan frame, 32)}
+	op := &operator{
+		opts: options{base: time.Second, cap: time.Minute}, store: s, conn: c,
+		log: io.Discard, attached: make(map[string]bool), jitter: func(d time.Duration) time.Duration { return d },
+	}
+	f.op = op
+	c.write = f.write
+	c.onEvent = op.event
+	return f
+}
+
+func (f *fakeDaemon) enqueue(m message) { f.op.conn.frames <- frame{message: m} }
+
+func (f *fakeDaemon) notify(method string, value any) {
+	data, err := json.Marshal(value)
+	if err != nil {
+		f.t.Fatal(err)
+	}
+	f.enqueue(message{Method: method, Params: data})
+}
+
+func (f *fakeDaemon) write(data []byte) error {
+	var m message
+	if err := json.Unmarshal(data, &m); err != nil {
+		return err
+	}
+	f.requests = append(f.requests, m)
+	if f.before != nil {
+		f.before(m)
+	}
+	var result any
+	switch m.Method {
+	case "thread/loaded/list":
+		result = loadedResult{Data: []string{"thread-1"}}
+	case "thread/goal/get":
+		result = goalResult{f.goal}
+	case "thread/resume":
+		result = struct{}{}
+	case "thread/turns/list":
+		result = turnsResult{[]turn{f.turn}}
+	case "thread/read":
+		result = map[string]any{"thread": map[string]any{"status": map[string]string{"type": f.status}}}
+	case "thread/goal/set":
+		var p map[string]any
+		if err := json.Unmarshal(m.Params, &p); err != nil {
+			return err
+		}
+		if len(p) != 3 || p["threadId"] != "thread-1" || p["status"] != "active" || p["origin"] != "user" {
+			f.t.Fatalf("goal changed beyond status: %s", m.Params)
+		}
+		f.goal = testGoal(active)
+		f.turn = turn{ID: "turn-2", Status: "inProgress"}
+		if f.lostReply {
+			f.op.conn.frames <- frame{err: io.EOF}
+			return nil
+		}
+		result = goalResult{f.goal}
+	default:
+		f.t.Fatalf("unexpected RPC: %s", data)
+	}
+	raw, err := json.Marshal(result)
+	if err != nil {
+		return err
+	}
+	f.enqueue(message{ID: m.ID, Result: raw})
+	return nil
+}
+
+func (f *fakeDaemon) reconcile() error {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	return f.op.reconcile(ctx)
+}
+
+func (f *fakeDaemon) count(method string) int {
+	count := 0
+	for _, m := range f.requests {
+		if m.Method == method {
+			count++
+		}
+	}
+	return count
+}
+
+func (f *fakeDaemon) fail() {
+	f.goal, f.turn = testGoal(blocked), testTurn("failed")
+}
+
+func TestRecovery(t *testing.T) {
+	f := newFake(t, nil)
+	if err := f.reconcile(); err != nil {
+		t.Fatal(err)
+	}
+	if !f.op.store.watch("thread-1").Armed || f.count("thread/resume") != 1 {
+		t.Fatal("did not attach active goal")
+	}
+	f.fail()
+	if err := f.reconcile(); err != nil {
+		t.Fatal(err)
+	}
+	w := f.op.store.watch("thread-1")
+	if !w.eligible() || w.Failures != 1 || f.count("thread/goal/set") != 0 {
+		t.Fatal("failed to wait for backoff")
+	}
+	w.Due = time.Now().Add(-time.Second)
+	if err := f.reconcile(); err != nil {
+		t.Fatal(err)
+	}
+	if f.count("thread/goal/set") != 1 || w.Failures != 1 {
+		t.Fatal("resume count or backoff changed")
+	}
+	if err := f.reconcile(); err != nil {
+		t.Fatal(err)
+	}
+	if f.count("thread/goal/set") != 1 {
+		t.Fatal("resumed active work twice")
+	}
+	if f.count("thread/resume") != 1 {
+		t.Fatal("subscribed repeatedly")
+	}
+	f.fail()
+	f.turn.ID = "turn-2"
+	if err := f.reconcile(); err != nil {
+		t.Fatal(err)
+	}
+	if w.Failures != 2 || time.Until(w.Due) > 2*time.Second || time.Until(w.Due) < 1900*time.Millisecond {
+		t.Fatalf("backoff did not double: %+v", w)
+	}
+}
+
+func TestOldBlocked(t *testing.T) {
+	f := newFake(t, nil)
+	f.fail()
+	if err := f.reconcile(); err != nil {
+		t.Fatal(err)
+	}
+	if f.op.store.watch("thread-1").eligible() || f.count("thread/goal/set") != 0 {
+		t.Fatal("adopted an old blocked goal")
+	}
+}
+
+func TestPauseDuringRead(t *testing.T) {
+	for _, method := range []string{"thread/goal/get", "thread/turns/list", "thread/read"} {
+		t.Run(method, func(t *testing.T) {
+			f := newFake(t, nil)
+			f.fail()
+			w := failedWatch()
+			w.Due = time.Now().Add(-time.Second)
+			f.op.store.Watches["thread-1"] = w
+			f.before = func(m message) {
+				if m.Method == method {
+					f.notify("thread/goal/updated", map[string]any{"threadId": "thread-1", "goal": testGoal(paused)})
+					f.before = nil
+				}
+			}
+			if err := f.reconcile(); err != nil {
+				t.Fatal(err)
+			}
+			if w.Goal.Status != paused || w.eligible() || f.count("thread/goal/set") != 0 {
+				t.Fatal("stale snapshot overrode human pause")
+			}
+		})
+	}
+}
+
+func TestLiveActivity(t *testing.T) {
+	f := newFake(t, nil)
+	f.fail()
+	f.status = "active"
+	w := failedWatch()
+	w.Due = time.Now().Add(-time.Second)
+	f.op.store.Watches["thread-1"] = w
+	f.before = func(m message) {
+		if m.Method == "thread/read" {
+			f.enqueue(message{ID: json.RawMessage(`"approval-1"`), Method: "item/commandExecution/requestApproval", Params: json.RawMessage(`{"threadId":"thread-1"}`)})
+		}
+	}
+	if err := f.reconcile(); err != nil {
+		t.Fatal(err)
+	}
+	if f.count("thread/goal/set") != 0 || w.eligible() {
+		t.Fatal("competed with active turn or approval")
+	}
+}
+
+func TestRestartPending(t *testing.T) {
+	f := newFake(t, nil)
+	w := failedWatch()
+	w.Due = time.Now().Add(-time.Second)
+	w.Failures = 3
+	f.op.store.Watches["thread-1"] = w
+	if err := f.op.save(); err != nil {
+		t.Fatal(err)
+	}
+	path := f.op.store.path
+	closeStore(t, f.op.store)
+	s, err := openStore(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer closeStore(t, s)
+	f = newFake(t, s)
+	f.fail()
+	if err := f.reconcile(); err != nil {
+		t.Fatal(err)
+	}
+	if f.count("thread/goal/set") != 1 || s.watch("thread-1").Failures != 3 {
+		t.Fatal("lost pending recovery across restart")
+	}
+}
+
+func TestLostSetReply(t *testing.T) {
+	f := newFake(t, nil)
+	f.fail()
+	f.lostReply = true
+	w := failedWatch()
+	w.Due = time.Now().Add(-time.Second)
+	w.Failures = 1
+	f.op.store.Watches["thread-1"] = w
+	if err := f.reconcile(); err != io.EOF {
+		t.Fatalf("connection error = %v", err)
+	}
+	if f.count("thread/goal/set") != 1 {
+		t.Fatal("missing mutation")
+	}
+	path := f.op.store.path
+	closeStore(t, f.op.store)
+	s, err := openStore(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer closeStore(t, s)
+	reconnected := newFake(t, s)
+	reconnected.goal, reconnected.turn = f.goal, f.turn
+	if err := reconnected.reconcile(); err != nil {
+		t.Fatal(err)
+	}
+	if reconnected.count("thread/goal/set") != 0 || !s.watch("thread-1").Due.IsZero() {
+		t.Fatal("duplicated a successful mutation after lost reply")
+	}
+}
+
+func TestNoCompetingRetry(t *testing.T) {
+	f := newFake(t, nil)
+	if err := f.reconcile(); err != nil {
+		t.Fatal(err)
+	}
+	for _, method := range []string{"error", "item/agentMessage/delta"} {
+		data := json.RawMessage(`{"threadId":"thread-1","willRetry":true}`)
+		if err := f.op.event(message{Method: method, Params: data}); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if err := f.reconcile(); err != nil {
+		t.Fatal(err)
+	}
+	if f.count("thread/goal/set") != 0 {
+		t.Fatal("competed with internal retry")
+	}
+}