corago/codex-keep/operator.go

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) && w.Goal.Status != "usageLimited" {
		_, _ = 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)
	if m.Method == "turn/completed" && w.Turn.ID != "" && w.Turn.ID != p.Turn.ID {
		return nil
	}
	var inputErr error
	if m.Method == "turn/completed" && p.Turn.Status == "interrupted" && w.Source != nil {
		inputErr = w.Source.seal()
	}

	switch m.Method {
	case "thread/goal/updated":
		w.notice(p.Goal)
	case "thread/goal/cleared":
		w.notice(nil)
	default:
		w.observe(p.Turn)
	}

	op.schedule(p.ThreadID)
	return errors.Join(inputErr, 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)
	if w.Source == nil {
		var probe goalResult
		if err := op.conn.call(ctx, "thread/goal/get", threadParams{id}, &probe); err != nil {
			return err
		}
		if probe.Goal == nil {
			w.update(nil)
			w.Pending = false
			return op.save()
		}
		var info struct {
			Thread struct {
				Path string `json:"path"`
			} `json:"thread"`
		}
		if err := op.conn.call(ctx, "thread/read", threadParams{id}, &info); err != nil {
			return err
		}
		source, err := openIntent(info.Thread.Path)
		if err != nil {
			return err
		}
		w.Source = source
		if !w.Pending {
			w.Hold = nil
		}
	}
	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)
	w.Pending = false
	if g.Goal == nil {
		return op.save()
	}
	end, err := w.Source.end()
	if err != nil {
		return err
	}
	if !op.attached[id] {
		w.echo = nil
	}
	p := struct {
		ThreadID string `json:"threadId"`
		Exclude  bool   `json:"excludeTurns"`
	}{id, true}
	var resumed struct {
		Thread struct {
			Path string `json:"path"`
		} `json:"thread"`
	}
	revision = w.revision
	if err := op.conn.call(ctx, "thread/resume", p, &resumed); err != nil {
		return err
	}
	op.attached[id] = true
	// The listener orders this reply after earlier notifications and before its goal snapshot.
	w.seed = g.Goal
	if w.Source.Path != resumed.Thread.Path {
		source, err := openIntent(resumed.Thread.Path)
		if err != nil {
			return err
		}
		w.Source, end = source, source.Pos
	}
	if w.revision != revision {
		return op.save()
	}
	if err := op.inputs(w, end); err != nil {
		return err
	}
	if w.Hold != nil && w.Hold.Status == active && g.Goal.Status.waiting() {
		w.Hold = nil
		w.Goal = g.Goal
		w.Known, w.Armed = true, true
	}
	if w.Hold != nil {
		if w.echo != nil {
			return op.save()
		}
		if g.Goal.Status != w.Hold.Status {
			return op.setGoal(ctx, id, w.Hold.Status)
		}
		w.update(g.Goal)
	}
	revision = w.revision
	var turns turnsResult
	turnPage := 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", turnPage, &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)
	if w.echo != nil || w.seed != nil {
		return nil
	}
	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.Due = time.Time{}
		w.Turn = turn{}
		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()
	}
	if w.Goal.Status == "usageLimited" {
		revision = w.revision
		w.Due = time.Now().Add(op.opts.poll)
		var usage usageResult
		p := struct {
			Exclude bool `json:"excludeResetCreditDetails"`
		}{true}
		if err := op.conn.call(ctx, "account/rateLimits/read", p, &usage); err != nil {
			if saveErr := op.save(); saveErr != nil {
				return saveErr
			}
			return err
		}
		if w.revision != revision || usage.Allowed == nil || !*usage.Allowed {
			return op.save()
		}
	} else {
		w.Due = time.Now().Add(min(op.opts.cap, op.jitter(backoff(op.opts.base, op.opts.cap, w.Failures+1))))
	}
	if err := op.setGoal(ctx, id, active); err != nil {
		return err
	}
	_, _ = fmt.Fprintf(op.log, "codex-keep: %s resumed after failed turn\n", id)
	return nil
}

func (op *operator) setGoal(ctx context.Context, id string, status goalStatus) error {
	w := op.store.watch(id)
	revision := w.revision
	w.Pending = true
	if err := op.save(); err != nil {
		return err
	}
	var g goalResult
	err := op.conn.call(ctx, "thread/goal/set", setParams{id, status, "automatic"}, &g)
	var remote *rpcError
	if err == nil || errors.As(err, &remote) {
		w.Pending = false
	}
	if err != nil {
		if status == active && w.revision == revision && w.eligible() && w.Goal.Status != "usageLimited" {
			w.Failures++
		}
	} else {
		w.echo = g.Goal
		if w.revision == revision {
			w.update(g.Goal)
		}
	}
	if saveErr := op.save(); saveErr != nil {
		return saveErr
	}
	return err
}

func (op *operator) inputs(w *watch, end int64) error {
	if w.Source == nil {
		return nil
	}
	status, err := w.Source.read(end)
	if err != nil {
		return err
	}
	if status == "" {
		return nil
	}
	w.revision++
	w.Due = time.Time{}
	if status == "cleared" {
		w.stop()
		w.Hold = nil
		w.Known = false
		return nil
	}
	if w.Goal == nil {
		return nil
	}
	if status == active && (w.Goal.Status == active || w.Goal.Status.waiting()) {
		w.Hold = nil
		w.Known, w.Armed = true, true
		return nil
	}
	desired := *w.Goal
	desired.Status = status
	w.Goal, w.Hold = &desired, &desired
	w.stop()
	w.Known = status == active
	return nil
}