corago/codex-keep/run_test.go

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