Add daemon-attached goal recovery operator
Resume goals after transient API failures with persisted exponential backoff while respecting pauses and interruptions.
This commit is contained in:
parent
a4522e5ed6
commit
e4b84d9802
12 changed files with 1595 additions and 0 deletions
33
codex-keep/README.md
Normal file
33
codex-keep/README.md
Normal 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
85
codex-keep/main.go
Normal 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
214
codex-keep/operator.go
Normal 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
293
codex-keep/operator_test.go
Normal 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
115
codex-keep/policy.go
Normal 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
186
codex-keep/policy_test.go
Normal 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
113
codex-keep/protocol.go
Normal 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
126
codex-keep/rpc.go
Normal 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
106
codex-keep/run.go
Normal 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
183
codex-keep/run_test.go
Normal 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
78
codex-keep/store.go
Normal 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
63
codex-keep/store_test.go
Normal 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")
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue