ms2k/corago
SHA256
Watch
0
Code Releases 3 Activity

Add daemon-attached goal recovery operator

Browse source 
Resume goals after transient API failures with persisted exponential backoff while respecting pauses and interruptions.
This commit is contained in:
ms2k 2026-10-18 14:01:07 +09:00
commit e4b84d9802
12 changed files with 1595 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

33
codex-keep/README.md Normal file
View file

@ -0,0 +1,33 @@
```mermaid
stateDiagram-v2
[*] --> Running
Running --> [*]: operator stopped / close socket / retain checkpoint
state Running {
state Connection {
[*] --> Connect
Connect --> Reconcile: WebSocket over Unix socket / initialize
Reconcile --> Listen: enumerate loaded threads / subscribe / read goal and latest turn
Listen --> Reconcile: poll interval or retry deadline
Reconcile --> Reconnect: connection failure
Listen --> Reconnect: connection failure
Connect --> Reconnect: connection failure
Reconnect --> Connect: capped exponential backoff with jitter
}
--
state "Per loaded thread; checkpointed across connections" as Recovery {
[*] --> Observe
Observe --> Armed: active goal and in-progress turn observed
Armed --> Backoff: same goal blocked and same turn failed with transient API error
Armed --> Observe: successful turn / reset failure count
Armed --> Observe: human interrupt / pause / clear / replacement / limit / completion / non-transient error
Backoff --> Check: persisted retry deadline reached
Check --> Observe: goal changed / turn changed / thread active or unloaded
Check --> Resume: same blocked goal / failed turn / thread idle or systemError
Resume --> Armed: goal activated and new turn observed
Resume --> Backoff: RPC failure / persist next retry
Resume --> Observe: activation observed before next turn
Backoff --> Observe: human pause / clear / replacement / limit / completion
Observe --> Observe: existing blocked goal or non-transient failure
}
}
```

85
codex-keep/main.go Normal file
View file

@ -0,0 +1,85 @@
package main
import (
"context"
"errors"
"flag"
"fmt"
"io"
"os"
"os/signal"
"path/filepath"
"syscall"
"time"
)
type options struct {
socket string
state string
base time.Duration
cap time.Duration
poll time.Duration
}
func statePath() string {
root := os.Getenv("XDG_STATE_HOME")
if root == "" {
home, _ := os.UserHomeDir()
root = filepath.Join(home, ".local", "state")
}
return filepath.Join(root, "codex-keep", "state.json")
}
func socketPath() string {
home := os.Getenv("CODEX_HOME")
if home == "" {
root, _ := os.UserHomeDir()
home = filepath.Join(root, ".codex")
}
return filepath.Join(home, "app-server-control", "app-server-control.sock")
}
func parseOptions(args []string, output io.Writer) (options, error) {
var opts options
flags := flag.NewFlagSet("codex-keep", flag.ContinueOnError)
flags.SetOutput(output)
flags.StringVar(&opts.socket, "socket", socketPath(), "Codex daemon Unix socket")
flags.StringVar(&opts.state, "state", statePath(), "Recovery checkpoint")
flags.DurationVar(&opts.base, "base-delay", 5*time.Second, "Initial retry delay")
flags.DurationVar(&opts.cap, "max-delay", 5*time.Minute, "Retry delay ceiling")
flags.DurationVar(&opts.poll, "poll", 5*time.Second, "Reconciliation interval")
flags.Usage = func() {
_, _ = fmt.Fprintln(output, "Usage: codex-keep [options]")
_, _ = fmt.Fprintln(output, "Reconcile failed active goals through the shared Codex daemon.")
flags.PrintDefaults()
}
if err := flags.Parse(args); err != nil {
return opts, err
}
if len(flags.Args()) != 0 {
return opts, errors.New("unexpected positional arguments")
}
if opts.base <= 0 || opts.cap < opts.base || opts.poll <= 0 {
return opts, errors.New("delays must satisfy 0 < base-delay <= max-delay and poll > 0")
}
if opts.state == "" {
return opts, errors.New("state path is empty")
}
return opts, nil
}
func main() {
opts, err := parseOptions(os.Args[1:], os.Stderr)
if errors.Is(err, flag.ErrHelp) {
return
}
if err == nil {
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
defer stop()
err = run(ctx, opts, os.Stderr)
}
if err != nil {
fmt.Fprintln(os.Stderr, "codex-keep:", err)
os.Exit(1)
}
}

214
codex-keep/operator.go Normal file
View file

@ -0,0 +1,214 @@
package main
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"time"
)
type stateError struct{ err error }
func (e *stateError) Error() string { return "save checkpoint: " + e.err.Error() }
func (e *stateError) Unwrap() error { return e.err }
type operator struct {
opts options
store *store
conn *connection
log io.Writer
attached map[string]bool
jitter func(time.Duration) time.Duration
}
func (op *operator) save() error {
if err := op.store.save(); err != nil {
return &stateError{err}
}
return nil
}
func (op *operator) schedule(id string) {
w := op.store.watch(id)
if w.schedule(time.Now(), op.opts.base, op.opts.cap, op.jitter) {
_, _ = fmt.Fprintf(op.log, "codex-keep: %s retry %d at %s\n", id, w.Failures, w.Due.Format(time.RFC3339Nano))
}
}
func (op *operator) event(m message) error {
switch m.Method {
case "thread/goal/updated", "thread/goal/cleared", "turn/started", "turn/completed":
default:
return nil
}
var p params
if err := json.Unmarshal(m.Params, &p); err != nil {
return err
}
if p.ThreadID == "" {
return nil
}
w := op.store.watch(p.ThreadID)
switch m.Method {
case "thread/goal/updated":
w.update(p.Goal)
case "thread/goal/cleared":
w.update(nil)
case "turn/started":
w.observe(p.Turn)
case "turn/completed":
if w.Turn.ID != "" && w.Turn.ID != p.Turn.ID {
return nil
}
w.observe(p.Turn)
}
op.schedule(p.ThreadID)
return op.save()
}
func (op *operator) reconcile(ctx context.Context) error {
loaded := make(map[string]bool)
cursor := ""
for {
var page loadedResult
p := struct {
Cursor string `json:"cursor,omitempty"`
Limit int `json:"limit"`
}{cursor, 100}
if err := op.conn.call(ctx, "thread/loaded/list", p, &page); err != nil {
return err
}
for _, id := range page.Data {
loaded[id] = true
}
cursor = page.Next
if cursor == "" {
break
}
}
for id := range op.attached {
if !loaded[id] {
delete(op.attached, id)
}
}
for id := range loaded {
if err := op.inspect(ctx, id); err != nil {
var remote *rpcError
if !errors.As(err, &remote) {
return err
}
_, _ = fmt.Fprintf(op.log, "codex-keep: %s reconciliation failed: %v\n", id, err)
w := op.store.watch(id)
if w.eligible() && !w.Due.After(time.Now()) {
w.Due = time.Time{}
op.schedule(id)
if err := op.save(); err != nil {
return err
}
}
}
}
return nil
}
func (op *operator) inspect(ctx context.Context, id string) error {
w := op.store.watch(id)
revision := w.revision
var g goalResult
if err := op.conn.call(ctx, "thread/goal/get", threadParams{id}, &g); err != nil {
return err
}
if w.revision != revision {
return nil
}
w.update(g.Goal)
if g.Goal == nil {
return op.save()
}
if !op.attached[id] {
p := struct {
ThreadID string `json:"threadId"`
Exclude bool `json:"excludeTurns"`
}{id, true}
if err := op.conn.call(ctx, "thread/resume", p, nil); err != nil {
return err
}
op.attached[id] = true
}
revision = w.revision
var turns turnsResult
p := struct {
ThreadID string `json:"threadId"`
Limit int `json:"limit"`
View string `json:"itemsView"`
}{id, 1, "summary"}
if err := op.conn.call(ctx, "thread/turns/list", p, &turns); err != nil {
return err
}
if w.revision != revision {
return nil
}
if len(turns.Data) != 0 {
w.observe(turns.Data[0])
}
op.schedule(id)
if err := op.save(); err != nil {
return err
}
if w.eligible() && !w.Due.IsZero() && !w.Due.After(time.Now()) {
return op.resume(ctx, id)
}
return nil
}
func (op *operator) resume(ctx context.Context, id string) error {
w := op.store.watch(id)
revision := w.revision
var live struct {
Thread struct {
Status struct {
Type string `json:"type"`
} `json:"status"`
} `json:"thread"`
}
if err := op.conn.call(ctx, "thread/read", threadParams{id}, &live); err != nil {
return err
}
if w.revision != revision {
return nil
}
if live.Thread.Status.Type != "idle" && live.Thread.Status.Type != "systemError" {
w.stop()
return op.save()
}
var g goalResult
if err := op.conn.call(ctx, "thread/goal/get", threadParams{id}, &g); err != nil {
return err
}
if w.revision != revision {
return nil
}
w.update(g.Goal)
if !w.eligible() {
return op.save()
}
w.Due = time.Now().Add(min(op.opts.cap, op.jitter(backoff(op.opts.base, op.opts.cap, w.Failures+1))))
if err := op.save(); err != nil {
return err
}
revision = w.revision
if err := op.conn.call(ctx, "thread/goal/set", setParams{id, active, "user"}, &g); err != nil {
w.Failures++
if saveErr := op.save(); saveErr != nil {
return saveErr
}
return err
}
if w.revision == revision {
w.update(g.Goal)
}
_, _ = fmt.Fprintf(op.log, "codex-keep: %s resumed after API failure\n", id)
return op.save()
}

293
codex-keep/operator_test.go Normal file
View file

@ -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")
}
}

115
codex-keep/policy.go Normal file
View file

@ -0,0 +1,115 @@
package main
import "time"
type goalStatus string
const (
active goalStatus = "active"
paused goalStatus = "paused"
blocked goalStatus = "blocked"
)
type goal struct {
Objective string `json:"objective"`
Status goalStatus `json:"status"`
CreatedAt int64 `json:"createdAt"`
UpdatedAt int64 `json:"updatedAt"`
}
func sameGoal(a, b *goal) bool {
return a != nil && b != nil && a.Objective == b.Objective && a.CreatedAt == b.CreatedAt
}
type watch struct {
Goal *goal `json:"goal,omitempty"`
Turn turn `json:"turn"`
Armed bool `json:"armed"`
Failures int `json:"failures"`
Due time.Time `json:"due"`
revision uint64
}
func (w *watch) stop() {
w.Armed = false
w.Due = time.Time{}
}
func (w *watch) update(g *goal) {
w.revision++
if !sameGoal(w.Goal, g) {
w.stop()
w.Failures = 0
}
previous := w.Goal
w.Goal = g
if g == nil {
w.stop()
return
}
switch g.Status {
case active:
if previous != nil && previous.Status != active {
w.stop()
}
if w.Turn.Status == "inProgress" {
w.Armed = true
}
case blocked:
if !w.Due.IsZero() && previous != nil && previous.UpdatedAt != g.UpdatedAt {
w.stop()
}
default:
w.stop()
}
}
func (w *watch) observe(t turn) {
if t.ID == "" {
return
}
w.revision++
if t.ID != w.Turn.ID {
w.stop()
w.Turn = t
w.Armed = t.Status == "inProgress" && w.Goal != nil && w.Goal.Status == active
return
}
w.Turn = t
switch t.Status {
case "completed":
w.stop()
w.Failures = 0
case "interrupted":
w.stop()
case "failed":
if t.Error == nil || !t.Error.Info.retryable() {
w.stop()
}
}
}
func (w *watch) eligible() bool {
return w.Armed && w.Goal != nil && w.Goal.Status == blocked &&
w.Turn.Status == "failed" && w.Turn.Error != nil && w.Turn.Error.Info.retryable()
}
func (w *watch) schedule(now time.Time, base, cap time.Duration, jitter func(time.Duration) time.Duration) bool {
if !w.eligible() || !w.Due.IsZero() {
return false
}
w.Failures++
w.Due = now.Add(min(cap, jitter(backoff(base, cap, w.Failures))))
return true
}
func backoff(base, cap time.Duration, failures int) time.Duration {
delay := base
for i := 1; i < failures && delay < cap; i++ {
if delay > cap/2 {
return cap
}
delay *= 2
}
return min(delay, cap)
}

186
codex-keep/policy_test.go Normal file
View file

@ -0,0 +1,186 @@
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 TestChangedBlock(t *testing.T) {
w := failedWatch()
w.Due = time.Now()
g := testGoal(blocked)
g.UpdatedAt++
w.update(g)
if w.eligible() || !w.Due.IsZero() {
t.Fatal("resumed a changed blocked goal")
}
}
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")
}
}

113
codex-keep/protocol.go Normal file
View file

@ -0,0 +1,113 @@
package main
import "encoding/json"
type rpcError struct {
Code int `json:"code"`
Message string `json:"message"`
}
func (e *rpcError) Error() string { return e.Message }
type message struct {
ID json.RawMessage `json:"id"`
Method string `json:"method"`
Params json.RawMessage `json:"params"`
Result json.RawMessage `json:"result"`
Error *rpcError `json:"error"`
}
type turnError struct {
Info errorInfo `json:"codexErrorInfo"`
}
type errorInfo struct {
Kind string
HTTP *int
}
func (e errorInfo) MarshalJSON() ([]byte, error) {
if e.HTTP == nil {
return json.Marshal(e.Kind)
}
return json.Marshal(map[string]struct {
HTTP *int `json:"httpStatusCode"`
}{e.Kind: {e.HTTP}})
}
func (e *errorInfo) UnmarshalJSON(data []byte) error {
*e = errorInfo{}
if string(data) == "null" {
return nil
}
if json.Unmarshal(data, &e.Kind) == nil {
return nil
}
var variants map[string]struct {
HTTP *int `json:"httpStatusCode"`
}
if err := json.Unmarshal(data, &variants); err != nil {
return err
}
if len(variants) == 1 {
for kind, value := range variants {
e.Kind, e.HTTP = kind, value.HTTP
}
}
return nil
}
func (e errorInfo) retryable() bool {
switch e.Kind {
case "serverOverloaded", "rateLimitExceeded", "internalServerError", "flexUnavailable":
return true
case "httpConnectionFailed", "responseStreamConnectionFailed", "responseStreamDisconnected", "responseTooManyFailedAttempts":
return e.HTTP == nil || *e.HTTP == 408 || *e.HTTP == 429 || *e.HTTP >= 500 && *e.HTTP <= 599
}
return false
}
type turn struct {
ID string `json:"id"`
Status string `json:"status"`
Error *turnError `json:"error,omitempty"`
}
type params struct {
ThreadID string `json:"threadId"`
Thread struct {
ID string `json:"id"`
} `json:"thread"`
Turn turn `json:"turn"`
Goal *goal `json:"goal"`
}
type goalResult struct {
Goal *goal `json:"goal"`
}
type turnsResult struct {
Data []turn `json:"data"`
}
type loadedResult struct {
Data []string `json:"data"`
Next string `json:"nextCursor"`
}
type threadParams struct {
ThreadID string `json:"threadId"`
}
type setParams struct {
ThreadID string `json:"threadId"`
Status goalStatus `json:"status"`
Origin string `json:"origin"`
}
type initParams struct {
Client struct {
Name string `json:"name"`
Version string `json:"version"`
} `json:"clientInfo"`
Capabilities struct {
Experimental bool `json:"experimentalApi"`
} `json:"capabilities"`
}

126
codex-keep/rpc.go Normal file
View file

@ -0,0 +1,126 @@
package main
import (
"context"
"encoding/json"
"fmt"
"net"
"net/http"
"strconv"
"time"
"github.com/coder/websocket"
)
type frame struct {
message message
err error
}
type connection struct {
write func([]byte) error
frames chan frame
serial uint64
onEvent func(message) error
stop func()
}
func connect(ctx context.Context, socket string) (*connection, error) {
ctx, cancel := context.WithCancel(ctx)
transport := &http.Transport{DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) {
return (&net.Dialer{}).DialContext(ctx, "unix", socket)
}}
handshake, stop := context.WithTimeout(ctx, 5*time.Second)
defer stop()
ws, _, err := websocket.Dial(handshake, "ws://localhost/", &websocket.DialOptions{HTTPClient: &http.Client{Transport: transport}})
if err != nil {
cancel()
transport.CloseIdleConnections()
return nil, err
}
ws.SetReadLimit(-1)
c := &connection{frames: make(chan frame, 16)}
c.write = func(data []byte) error { return ws.Write(ctx, websocket.MessageText, data) }
done := make(chan struct{})
go func() {
defer close(done)
for {
_, data, err := ws.Read(ctx)
var m message
if err == nil {
err = json.Unmarshal(data, &m)
}
select {
case c.frames <- frame{m, err}:
case <-ctx.Done():
return
}
if err != nil {
return
}
}
}()
c.stop = func() {
cancel()
_ = ws.CloseNow()
transport.CloseIdleConnections()
<-done
}
return c, nil
}
func (c *connection) send(id uint64, method string, params any) error {
data, err := json.Marshal(struct {
ID uint64 `json:"id,omitempty"`
Method string `json:"method"`
Params any `json:"params"`
}{id, method, params})
if err != nil {
return err
}
return c.write(data)
}
func (c *connection) event(m message) error {
// Approval and input requests belong to the interactive clients.
if len(m.ID) != 0 || m.Method == "" || c.onEvent == nil {
return nil
}
return c.onEvent(m)
}
func (c *connection) call(ctx context.Context, method string, params, result any) error {
ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
defer cancel()
c.serial++
id := c.serial
if err := c.send(id, method, params); err != nil {
return err
}
for {
select {
case <-ctx.Done():
return ctx.Err()
case f := <-c.frames:
if f.err != nil {
return f.err
}
m := f.message
if m.Method == "" && string(m.ID) == strconv.FormatUint(id, 10) {
if m.Error != nil {
return m.Error
}
if result == nil {
return nil
}
if err := json.Unmarshal(m.Result, result); err != nil {
return fmt.Errorf("%s: %w", method, err)
}
return nil
}
if err := c.event(m); err != nil {
return err
}
}
}
}

106
codex-keep/run.go Normal file
View file

@ -0,0 +1,106 @@
package main
import (
"context"
"errors"
"fmt"
"io"
mathrand "math/rand/v2"
"time"
)
func jitter(d time.Duration) time.Duration {
return time.Duration(float64(d) * (0.9 + mathrand.Float64()*0.2))
}
func run(ctx context.Context, opts options, log io.Writer) (runErr error) {
s, err := openStore(opts.state)
if err != nil {
return err
}
defer func() { runErr = errors.Join(runErr, s.close()) }()
failures := 0
for ctx.Err() == nil {
started := time.Now()
err := session(ctx, opts, s, log)
if ctx.Err() != nil {
return nil
}
var diskErr *stateError
if errors.As(err, &diskErr) {
return err
}
if time.Since(started) >= opts.cap {
failures = 0
}
failures++
delay := min(opts.cap, jitter(backoff(opts.base, opts.cap, failures)))
_, _ = fmt.Fprintf(log, "codex-keep: connection ended (%v); reconnecting in %s\n", err, delay.Round(time.Millisecond))
timer := time.NewTimer(delay)
select {
case <-ctx.Done():
timer.Stop()
return nil
case <-timer.C:
}
}
return nil
}
func session(ctx context.Context, opts options, s *store, log io.Writer) error {
c, err := connect(ctx, opts.socket)
if err != nil {
return err
}
defer c.stop()
op := &operator{opts: opts, store: s, conn: c, log: log, attached: make(map[string]bool), jitter: jitter}
c.onEvent = op.event
var init initParams
init.Client.Name, init.Client.Version = "codex-keep", "0.1.0"
init.Capabilities.Experimental = true
if err := c.call(ctx, "initialize", init, nil); err != nil {
return err
}
if err := c.send(0, "initialized", struct{}{}); err != nil {
return err
}
_, _ = fmt.Fprintln(log, "codex-keep: connected to shared daemon")
nextPoll := time.Now()
for {
now := time.Now()
if !nextPoll.After(now) {
if err := op.reconcile(ctx); err != nil {
return err
}
nextPoll = time.Now().Add(opts.poll)
}
wake := nextPoll
for id := range op.attached {
w := s.watch(id)
if !w.Due.IsZero() && w.eligible() && w.Due.Before(wake) {
wake = w.Due
}
}
if !wake.After(time.Now()) {
if err := op.reconcile(ctx); err != nil {
return err
}
continue
}
timer := time.NewTimer(time.Until(wake))
select {
case <-ctx.Done():
timer.Stop()
return ctx.Err()
case <-timer.C:
case f := <-c.frames:
timer.Stop()
if f.err != nil {
return f.err
}
if err := c.event(f.message); err != nil {
return err
}
}
}
}

183
codex-keep/run_test.go Normal file
View file

@ -0,0 +1,183 @@
package main
import (
"context"
"encoding/json"
"errors"
"io"
"net"
"net/http"
"os"
"path/filepath"
"sync/atomic"
"testing"
"time"
"github.com/coder/websocket"
)
func TestSocketReconnect(t *testing.T) {
root := t.TempDir()
socket := filepath.Join(root, "rpc.sock")
listener, err := net.Listen("unix", socket)
if err != nil {
t.Fatal(err)
}
var connects atomic.Int32
var live atomic.Int32
ready := make(chan struct{}, 1)
server := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ws, err := websocket.Accept(w, r, nil)
if err != nil {
return
}
live.Add(1)
defer live.Add(-1)
defer func() { _ = ws.CloseNow() }()
drop := connects.Add(1) == 1
for {
_, data, err := ws.Read(r.Context())
if err != nil {
return
}
var m message
if json.Unmarshal(data, &m) != nil {
return
}
if len(m.ID) == 0 {
continue
}
var result any = struct{}{}
switch m.Method {
case "initialize":
var init initParams
if json.Unmarshal(m.Params, &init) != nil || !init.Capabilities.Experimental {
t.Error("experimental API not enabled")
return
}
case "thread/loaded/list":
if drop {
return
}
result = loadedResult{Data: []string{}}
select {
case ready <- struct{}{}:
default:
}
default:
t.Errorf("unexpected RPC: %s", m.Method)
return
}
response, err := json.Marshal(struct {
ID json.RawMessage `json:"id"`
Result any `json:"result"`
}{m.ID, result})
if err != nil {
t.Error(err)
return
}
if ws.Write(r.Context(), websocket.MessageText, response) != nil {
return
}
}
})}
served := make(chan error, 1)
go func() { served <- server.Serve(listener) }()
defer func() {
if err := server.Close(); err != nil {
t.Error(err)
}
if err := <-served; err != nil && !errors.Is(err, http.ErrServerClosed) {
t.Error(err)
}
}()
opts := options{socket: socket, state: filepath.Join(root, "state.json"), base: 10 * time.Millisecond, cap: 40 * time.Millisecond, poll: 10 * time.Millisecond}
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
defer cancel()
done := make(chan error, 1)
go func() { done <- run(ctx, opts, io.Discard) }()
select {
case <-ready:
case err := <-done:
t.Fatalf("operator exited early: %v", err)
case <-ctx.Done():
if err := <-done; err != nil {
t.Error(err)
}
t.Fatal("operator did not reconnect")
}
cancel()
if err := <-done; err != nil {
t.Fatal(err)
}
deadline := time.Now().Add(time.Second)
for live.Load() != 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if connects.Load() < 2 || live.Load() != 0 {
t.Fatalf("connections: total=%d live=%d", connects.Load(), live.Load())
}
}
func TestOptions(t *testing.T) {
for _, args := range [][]string{{"--base-delay", "0s"}, {"--base-delay", "2m", "--max-delay", "1m"}, {"--poll", "0s"}, {"--state", ""}, {"resume", "thread-1"}} {
if _, err := parseOptions(args, io.Discard); err == nil {
t.Fatal(args)
}
}
opts, err := parseOptions([]string{"--base-delay", "1s", "--max-delay", "1m"}, io.Discard)
if err != nil || opts.base != time.Second || opts.cap != time.Minute {
t.Fatalf("options = %+v, %v", opts, err)
}
t.Setenv("CODEX_HOME", "/tmp/synthetic-codex")
if socketPath() != "/tmp/synthetic-codex/app-server-control/app-server-control.sock" {
t.Fatal(socketPath())
}
}
func TestDaemonRead(t *testing.T) {
socket := os.Getenv("KEEP_LIVE_SOCKET")
if socket == "" {
t.Skip("KEEP_LIVE_SOCKET is unset")
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
c, err := connect(ctx, socket)
if err != nil {
t.Fatal(err)
}
defer c.stop()
var init initParams
init.Client.Name, init.Client.Version = "codex-keep-probe", "0.1.0"
init.Capabilities.Experimental = true
if err := c.call(ctx, "initialize", init, nil); err != nil {
t.Fatal(err)
}
if err := c.send(0, "initialized", struct{}{}); err != nil {
t.Fatal(err)
}
var page loadedResult
if err := c.call(ctx, "thread/loaded/list", struct{}{}, &page); err != nil {
t.Fatal(err)
}
statuses := make(map[goalStatus]int)
for _, id := range page.Data {
var result goalResult
if err := c.call(ctx, "thread/goal/get", threadParams{id}, &result); err != nil {
t.Fatal(err)
}
if result.Goal != nil {
statuses[result.Goal.Status]++
var turns turnsResult
p := struct {
ThreadID string `json:"threadId"`
Limit int `json:"limit"`
View string `json:"itemsView"`
}{id, 1, "summary"}
if err := c.call(ctx, "thread/turns/list", p, &turns); err != nil {
t.Fatal(err)
}
}
}
t.Logf("loaded=%d goals=%v; read-only", len(page.Data), statuses)
}

78
codex-keep/store.go Normal file
View file

@ -0,0 +1,78 @@
package main
import (
"encoding/json"
"errors"
"fmt"
"os"
"path/filepath"
"syscall"
)
type store struct {
path string
lock *os.File
Watches map[string]*watch `json:"watches"`
}
func openStore(path string) (*store, error) {
if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
return nil, err
}
lock, err := os.OpenFile(path+".lock", os.O_CREATE|os.O_RDWR, 0600)
if err != nil {
return nil, err
}
if err := syscall.Flock(int(lock.Fd()), syscall.LOCK_EX|syscall.LOCK_NB); err != nil {
return nil, fmt.Errorf("checkpoint is already in use: %w", errors.Join(err, lock.Close()))
}
s := &store{path: path, lock: lock, Watches: make(map[string]*watch)}
data, err := os.ReadFile(path)
if errors.Is(err, os.ErrNotExist) {
return s, nil
}
if err == nil {
err = json.Unmarshal(data, s)
}
if err == nil && s.Watches == nil {
err = errors.New("checkpoint has no watches")
}
if err != nil {
return nil, errors.Join(err, s.close())
}
return s, nil
}
func (s *store) close() error { return s.lock.Close() }
func (s *store) save() error {
data, err := json.Marshal(s)
if err != nil {
return err
}
f, err := os.CreateTemp(filepath.Dir(s.path), ".state-*")
if err != nil {
return err
}
defer func() { _ = os.Remove(f.Name()) }()
if _, err = f.Write(data); err == nil {
err = f.Sync()
}
closeErr := f.Close()
if err != nil {
return err
}
if closeErr != nil {
return closeErr
}
return os.Rename(f.Name(), s.path)
}
func (s *store) watch(id string) *watch {
w := s.Watches[id]
if w == nil {
w = &watch{}
s.Watches[id] = w
}
return w
}

63
codex-keep/store_test.go Normal file
View file

@ -0,0 +1,63 @@
package main
import (
"errors"
"os"
"path/filepath"
"testing"
"time"
)
func closeStore(t *testing.T, s *store) {
t.Helper()
if err := s.close(); err != nil && !errors.Is(err, os.ErrClosed) {
t.Error(err)
}
}
func TestCheckpoint(t *testing.T) {
path := filepath.Join(t.TempDir(), "state.json")
s, err := openStore(path)
if err != nil {
t.Fatal(err)
}
w := failedWatch()
w.Due = time.Now().Add(time.Minute).UTC()
w.Failures = 3
s.Watches["thread-1"] = w
if err := s.save(); err != nil {
t.Fatal(err)
}
if other, err := openStore(path); err == nil {
closeStore(t, other)
t.Fatal("two operators acquired one checkpoint")
}
closeStore(t, s)
restored, err := openStore(path)
if err != nil {
t.Fatal(err)
}
defer closeStore(t, restored)
r := restored.watch("thread-1")
if !r.eligible() || !r.Due.Equal(w.Due) || r.Failures != 3 {
t.Fatalf("lost recovery: %+v", r)
}
info, err := os.Stat(path)
if err != nil {
t.Fatal(err)
}
if info.Mode().Perm() != 0600 {
t.Fatal(info.Mode())
}
}
func TestBadCheckpoint(t *testing.T) {
path := filepath.Join(t.TempDir(), "state.json")
if err := os.WriteFile(path, []byte(`{"watches":`), 0600); err != nil {
t.Fatal(err)
}
if s, err := openStore(path); err == nil {
closeStore(t, s)
t.Fatal("ignored corrupt state")
}
}