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) {
		_, _ = 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()
}