byre/internal/codex/session_internal_test.go

package codex

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"image"
	"image/color"
	"image/draw"
	"image/png"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"strings"
	"sync"
	"testing"
	"time"

	"g.fc.run/misc/byre/internal/mcp"
)

const (
	testModel   = "gpt-5.6-luna"
	switchModel = "gpt-5.6-terra"
	gitName     = "Byre Test"
	gitEmail    = "test@byre.invalid"
)

type scenario struct {
	name string
	run  func(h *harness)
}

var scenarios = []scenario{
	{"CommitIdentity", scenarioCommitIdentity},
	{"TextDeltas", scenarioTextDeltas},
	{"Interrupt", scenarioInterrupt},
	{"ImageInput", scenarioImageInput},
	{"Resume", scenarioResume},
	{"ModelSwitch", scenarioModelSwitch},
	{"Compact", scenarioCompact},
	{"Tools", scenarioTools},
	{"Fork", scenarioFork},
	{"Steer", scenarioSteer},
}

func TestLive(t *testing.T) {
	for _, sc := range scenarios {
		t.Run(sc.name, func(t *testing.T) {
			t.Parallel()
			h := &harness{t: t, name: sc.name, live: true, ctx: testContext(t, 5*time.Minute), repo: scratchRepo(t)}
			sc.run(h)
			h.closeAll()
			for _, path := range h.files {
				checkTranscript(t, path)
			}
		})
	}
}

func TestContract(t *testing.T) {
	for _, sc := range scenarios {
		t.Run(sc.name, func(t *testing.T) {
			t.Parallel()
			h := &harness{t: t, name: sc.name, ctx: testContext(t, 30*time.Second), repo: t.TempDir()}
			sc.run(h)
			h.closeAll()
			for _, path := range h.files {
				checkTranscript(t, path)
			}
		})
	}
}

func scenarioCommitIdentity(h *harness) {
	s := h.start("", Options{})
	h.send(s, `Create a file hello.txt containing exactly the word hi, then git add and git commit it with the message "add hello". Reply done when finished.`)
	turn := h.waitTurn(s, nil)
	if turn.Status != TurnStatusCompleted {
		h.t.Fatalf("turn ended %s", turn.Status)
	}

	if !h.live {
		return
	}
	log := git(h.t, h.repo, "log", "-1", "--format=%an <%ae>%n%cn <%ce>%n%B")
	want := gitName + " <" + gitEmail + ">\n" + gitName + " <" + gitEmail + ">\nadd hello"
	if !strings.HasPrefix(log, want) {
		h.t.Fatalf("unexpected commit:\n%s", log)
	}
	lower := strings.ToLower(log)
	if strings.Contains(lower, "co-authored-by") || strings.Contains(lower, "codex") {
		h.t.Fatalf("agent attribution in commit:\n%s", log)
	}
}

func scenarioTextDeltas(h *harness) {
	s := h.start("", Options{})
	h.send(s, "Without using any tools, count from one to twenty in words, separated by spaces.")

	deltas := map[string][]string{}
	h.waitTurn(s, func(ev Event) {
		if ev.Method != MethodItemAgentMessageDelta {
			return
		}
		d := decode[*AgentMessageDeltaNotification](h.t, ev)
		deltas[d.ItemID] = append(deltas[d.ItemID], d.Delta)
	})

	msgs := h.messages[s]
	if len(msgs) == 0 {
		h.t.Fatalf("no agent message")
	}
	last := msgs[len(msgs)-1]
	got := deltas[last.ID]
	if len(got) < 2 {
		h.t.Fatalf("expected several text deltas for %s, got %d", last.ID, len(got))
	}
	if strings.Join(got, "") != last.Text {
		h.t.Fatalf("deltas %q don't add up to the final text %q", got, last.Text)
	}
}

func scenarioInterrupt(h *harness) {
	s := h.start("", Options{})
	h.send(s, "Run the shell command `ping -c 120 127.0.0.1` and wait for it to finish, then say done.")

	interruptedAt := time.Time{}
	turn := h.waitTurn(s, func(ev Event) {
		if !interruptedAt.IsZero() || ev.Method != MethodItemStarted {
			return
		}
		item := decode[*ItemStartedNotification](h.t, ev).Item
		if item.Type != ThreadItemTypeCommandExecution {
			return
		}
		if h.live {
			time.Sleep(2 * time.Second)
		}
		interruptedAt = time.Now()
		err := s.Interrupt(h.ctx)
		if err != nil {
			h.t.Errorf("interrupt: %v", err)
		}
	})

	if interruptedAt.IsZero() {
		h.t.Fatalf("never saw a command to interrupt")
	}
	if turn.Status != TurnStatusInterrupted {
		h.t.Fatalf("turn ended %s, want interrupted", turn.Status)
	}
	elapsed := time.Since(interruptedAt)
	h.t.Logf("turn ended %s after interrupt", elapsed)
	if elapsed > 30*time.Second {
		h.t.Fatalf("interrupt took %s to end the turn", elapsed)
	}

	h.send(s, "Reply with the single word ok.")
	h.waitTurn(s, nil)
	if !strings.Contains(strings.ToLower(h.lastText(s)), "ok") {
		h.t.Fatalf("session unusable after interrupt; assistant said %q", h.lastText(s))
	}
}

func scenarioImageInput(h *harness) {
	s := h.start("", Options{})

	img := image.NewRGBA(image.Rect(0, 0, 64, 64))
	draw.Draw(img, img.Bounds(), &image.Uniform{C: color.RGBA{R: 255, A: 255}}, image.Point{}, draw.Src)
	buf := bytes.Buffer{}
	err := png.Encode(&buf, img)
	if err != nil {
		h.t.Fatalf("encode png: %v", err)
	}

	_, err = s.Send(h.ctx, "Without using any tools, what single colour fills the attached image? Reply with just the colour name.", Image{MediaType: "image/png", Data: buf.Bytes()})
	if err != nil {
		h.t.Fatalf("send: %v", err)
	}
	h.waitTurn(s, nil)

	if !strings.Contains(strings.ToLower(h.lastText(s)), "red") {
		h.t.Fatalf("model didn't see the image; assistant said %q", h.lastText(s))
	}
}

func scenarioResume(h *harness) {
	first := h.start("first", Options{})
	h.send(first, "Remember the codeword AZURE-17. Reply with the single word ok.")
	h.waitTurn(first, nil)
	id := first.ID()
	h.close(first)

	if id == "" {
		h.t.Fatalf("no thread id from first session")
	}

	resumed := h.start("resumed", Options{Resume: id})
	h.send(resumed, "What was the codeword? Reply with just the codeword.")
	h.waitTurn(resumed, nil)

	if !strings.Contains(h.lastText(resumed), "AZURE-17") {
		h.t.Fatalf("resumed session lost context; assistant said %q", h.lastText(resumed))
	}
	if resumed.ID() != id {
		h.t.Fatalf("resume changed thread id: want %s, got %s", id, resumed.ID())
	}
}

func scenarioModelSwitch(h *harness) {
	s := h.start("", Options{})
	h.send(s, "Reply with the single word ok.")
	h.waitTurn(s, nil)

	before := h.threadModel(s)
	if before != testModel {
		h.t.Fatalf("thread model %q before the switch, want %s", before, testModel)
	}

	s.SetModel(switchModel)
	h.send(s, "Reply with the single word ok.")
	turn := h.waitTurn(s, nil)
	if turn.Status != TurnStatusCompleted {
		h.t.Fatalf("turn on %s ended %s: %+v", switchModel, turn.Status, turn.Error)
	}

	after := h.threadModel(s)
	if after != switchModel {
		h.t.Fatalf("thread model %q after the switch, want %s", after, switchModel)
	}
}

func (h *harness) threadModel(s *Session) string {
	resp := ThreadReadResponse{}
	err := s.Call(h.ctx, MethodThreadRead, ThreadReadParams{ThreadID: s.ID()}, &resp)
	if err != nil {
		h.t.Fatalf("read thread: %v", err)
	}
	return resp.Thread.Model
}

func scenarioCompact(h *harness) {
	s := h.start("", Options{})
	h.send(s, "Remember the codeword CORAL-42. Reply with the single word ok.")
	h.waitTurn(s, nil)

	err := s.Compact(h.ctx)
	if err != nil {
		h.t.Fatalf("compact: %v", err)
	}
	compacted := false
	h.waitTurn(s, func(ev Event) {
		if ev.Method == MethodThreadCompacted {
			compacted = true
		}
		if ev.Method == MethodItemCompleted && decode[*ItemCompletedNotification](h.t, ev).Item.Type == ThreadItemTypeContextCompaction {
			compacted = true
		}
	})
	if !compacted {
		h.t.Fatalf("compaction turn ended without compacting")
	}

	h.send(s, "What was the codeword? Reply with just the codeword.")
	h.waitTurn(s, nil)
	if !strings.Contains(h.lastText(s), "CORAL-42") {
		h.t.Fatalf("compacted session lost context; assistant said %q", h.lastText(s))
	}
}

func lookupServer(calls *[]string) *mcp.Server {
	return &mcp.Server{Name: "byre", Version: "0.0.0", Tools: []mcp.Tool{{
		Name:        "lookup",
		Description: "Look up the current codeword for a topic.",
		InputSchema: json.RawMessage(`{"type":"object","properties":{"topic":{"type":"string"}},"required":["topic"]}`),
		Handler: func(ctx context.Context, args json.RawMessage) (string, error) {
			*calls = append(*calls, string(args))
			return "The codeword is TEAL-9.", nil
		},
	}}}
}

func scenarioTools(h *harness) {
	calls := []string{}
	srv := lookupServer(&calls)

	first := h.start("first", Options{MCPServers: []*mcp.Server{srv}, Reminder: "This is a byre test session."})
	h.send(first, "Call byre's lookup tool with topic birds, then reply with just the codeword it gives.")
	h.waitTurn(first, nil)
	if len(calls) != 1 || !strings.Contains(h.lastText(first), "TEAL-9") {
		h.t.Fatalf("tool calls %v; assistant said %q", calls, h.lastText(first))
	}
	id := first.ID()
	h.close(first)

	resumed := h.start("resumed", Options{Resume: id, MCPServers: []*mcp.Server{srv}})
	h.send(resumed, "Call byre's lookup tool again with topic fish, then reply with just the codeword it gives.")
	h.waitTurn(resumed, nil)
	if len(calls) != 2 || !strings.Contains(h.lastText(resumed), "TEAL-9") {
		h.t.Fatalf("resumed thread lost the tool: calls %v; assistant said %q", calls, h.lastText(resumed))
	}
}

func scenarioFork(h *harness) {
	src := h.start("source", Options{})
	h.send(src, "Remember the codeword AZURE-17. Reply with the single word ok.")
	keep := h.waitTurn(src, nil)
	h.send(src, "The codeword is now BEIGE-3; forget the old one. Reply with the single word ok.")
	h.waitTurn(src, nil)
	id := src.ID()
	h.close(src)

	fork := h.start("fork", Options{ForkFrom: id, ForkAt: keep.ID})
	if fork.ID() == "" || fork.ID() == id {
		h.t.Fatalf("fork thread id %q, source %q", fork.ID(), id)
	}
	h.send(fork, "What is the codeword? Reply with just the codeword.")
	h.waitTurn(fork, nil)
	text := h.lastText(fork)
	if !strings.Contains(text, "AZURE-17") || strings.Contains(text, "BEIGE-3") {
		h.t.Fatalf("fork didn't stop after the first turn; assistant said %q", text)
	}
}

func scenarioSteer(h *harness) {
	s := h.start("", Options{})
	h.send(s, "Run the shell command `sleep 5 && echo finished`, then reply with what it printed.")

	steered := false
	h.waitTurn(s, func(ev Event) {
		if steered || ev.Method != MethodItemStarted {
			return
		}
		if decode[*ItemStartedNotification](h.t, ev).Item.Type != ThreadItemTypeCommandExecution {
			return
		}
		steered = true
		err := s.Steer(h.ctx, "Also end your reply with the word PELICAN.")
		if err != nil {
			h.t.Errorf("steer: %v", err)
		}
	})

	if !steered {
		h.t.Fatalf("never saw a command to steer during")
	}
	if !strings.Contains(h.lastText(s), "PELICAN") {
		h.t.Fatalf("steer didn't reach the turn; assistant said %q", h.lastText(s))
	}
}

type harness struct {
	t    *testing.T
	name string
	live bool
	ctx  context.Context
	repo string

	files    []string
	sessions []*Session
	closed   map[*Session]bool
	replays  map[*Session]*replay
	messages map[*Session][]ThreadItem
}

func (h *harness) path(dir, suffix string) string {
	name := h.name
	if suffix != "" {
		name += "-" + suffix
	}
	return filepath.Join("testdata", dir, name+".jsonl")
}

func (h *harness) start(suffix string, opts Options) *Session {
	opts.Cwd = h.repo
	opts.Model = testModel
	opts.GitName = gitName
	opts.GitEmail = gitEmail
	opts.OnBadLine = func(line []byte, reason string) {
		h.t.Errorf("bad line (%s): %s", reason, line)
	}

	if h.messages == nil {
		h.messages = map[*Session][]ThreadItem{}
		h.closed = map[*Session]bool{}
		h.replays = map[*Session]*replay{}
	}

	if h.live {
		path := h.path("transcripts", suffix)
		err := os.MkdirAll(filepath.Dir(path), 0o755)
		if err != nil {
			h.t.Fatalf("mkdir transcripts: %v", err)
		}
		f, err := os.Create(path)
		if err != nil {
			h.t.Fatalf("create transcript: %v", err)
		}
		h.t.Cleanup(func() { f.Close() })
		opts.Transcript = f
		h.files = append(h.files, path)

		s, err := Start(h.ctx, opts)
		if err != nil {
			h.t.Fatalf("start: %v", err)
		}
		h.sessions = append(h.sessions, s)
		return s
	}

	path := h.path("contract", suffix)
	h.files = append(h.files, path)
	r := newReplay(h.t, path)
	s := newSession(opts, r.clientStdin, r.clientStdout)
	h.sessions = append(h.sessions, s)
	h.replays[s] = r
	err := s.open(h.ctx, opts)
	if err != nil {
		h.t.Fatalf("open: %v", err)
	}
	return s
}

func (h *harness) close(s *Session) {
	if h.closed[s] {
		return
	}
	h.closed[s] = true

	if h.live {
		err := s.Close()
		if err != nil {
			h.t.Logf("codex exit: %v", err)
		}
		return
	}

	s.stdin.Close()
	<-s.done
	h.replays[s].wait()
}

func (h *harness) closeAll() {
	for _, s := range h.sessions {
		h.close(s)
	}
}

func (h *harness) send(s *Session, text string) {
	_, err := s.Send(h.ctx, text)
	if err != nil {
		h.t.Fatalf("send: %v", err)
	}
}

func (h *harness) waitTurn(s *Session, onEvent func(Event)) Turn {
	for {
		select {
		case ev, ok := <-s.Events():
			if !ok {
				h.t.Fatalf("codex exited before the turn completed")
			}
			if onEvent != nil {
				onEvent(ev)
			}
			if ev.Method == MethodItemCompleted {
				item := decode[*ItemCompletedNotification](h.t, ev).Item
				if item.Type == ThreadItemTypeAgentMessage {
					h.messages[s] = append(h.messages[s], item)
				}
			}
			if ev.Method == MethodTurnCompleted {
				turn := decode[*TurnCompletedNotification](h.t, ev).Turn
				h.t.Logf("turn %s %s", turn.ID, turn.Status)
				return turn
			}
		case <-h.ctx.Done():
			h.t.Fatalf("timed out waiting for the turn to complete")
		}
	}
}

func (h *harness) lastText(s *Session) string {
	msgs := h.messages[s]
	if len(msgs) == 0 {
		return ""
	}
	return msgs[len(msgs)-1].Text
}

func decode[T any](t *testing.T, ev Event) T {
	v, err := ev.Decode()
	if err != nil {
		t.Fatalf("decode %s: %v", ev.Method, err)
	}
	out, ok := v.(T)
	if !ok {
		t.Fatalf("decode %s: got %T", ev.Method, v)
	}
	return out
}

type entry struct {
	Dir string          `json:"dir"`
	Msg json.RawMessage `json:"msg"`
}

func readTranscript(t *testing.T, path string) []entry {
	b, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("read transcript: %v", err)
	}
	out := []entry{}
	for _, line := range bytes.Split(bytes.TrimSpace(b), []byte("\n")) {
		e := entry{}
		err := json.Unmarshal(line, &e)
		if err != nil {
			t.Fatalf("parse transcript line %s: %v", line, err)
		}
		out = append(out, e)
	}
	return out
}

func strict(raw json.RawMessage, v any) error {
	d := json.NewDecoder(bytes.NewReader(raw))
	d.DisallowUnknownFields()
	return d.Decode(v)
}

func checkTranscript(t *testing.T, path string) {
	methods := map[string]string{}
	for i, e := range readTranscript(t, path) {
		m := message{}
		err := json.Unmarshal(e.Msg, &m)
		if err != nil {
			t.Fatalf("%s:%d: %v", path, i+1, err)
		}

		if e.Dir == "send" {
			if m.Method != "" && len(m.ID) > 0 {
				methods[string(m.ID)] = m.Method
			}
			continue
		}

		var ctor func() any
		var raw json.RawMessage
		label := m.Method
		switch {
		case m.Method == "":
			label = methods[string(m.ID)] + " response"
			c, ok := clientResults[methods[string(m.ID)]]
			if !ok {
				t.Errorf("%s:%d: response to unknown request %s", path, i+1, m.ID)
				continue
			}
			if m.Error != nil {
				continue
			}
			ctor, raw = c, m.Result
		case len(m.ID) > 0:
			c, ok := serverRequestParams[m.Method]
			if !ok {
				t.Errorf("%s:%d: unknown server request %s", path, i+1, m.Method)
				continue
			}
			ctor, raw = c, m.Params
		default:
			c, ok := notificationParams[m.Method]
			if !ok {
				t.Errorf("%s:%d: unknown notification %s", path, i+1, m.Method)
				continue
			}
			ctor, raw = c, m.Params
		}
		if ctor == nil {
			continue
		}
		err = strict(raw, ctor())
		if err != nil {
			t.Errorf("%s:%d: %s doesn't match the generated types: %v", path, i+1, label, err)
		}
	}
}

type replay struct {
	t       *testing.T
	entries []entry

	clientStdin  io.WriteCloser
	clientStdout io.Reader
	serverIn     *bufio.Reader
	serverOut    io.WriteCloser

	ids  map[string]json.RawMessage
	done chan struct{}
	once sync.Once
}

func newReplay(t *testing.T, path string) *replay {
	inR, inW := io.Pipe()
	outR, outW := io.Pipe()
	r := &replay{
		t:            t,
		entries:      readTranscript(t, path),
		clientStdin:  inW,
		clientStdout: outR,
		serverIn:     bufio.NewReader(inR),
		serverOut:    outW,
		ids:          map[string]json.RawMessage{},
		done:         make(chan struct{}),
	}
	go r.run()
	return r
}

func (r *replay) wait() {
	<-r.done
}

func (r *replay) run() {
	defer close(r.done)
	defer r.serverOut.Close()

	for i, e := range r.entries {
		if e.Dir == "recv" {
			_, err := r.serverOut.Write(append(r.rewrite(e.Msg), '\n'))
			if err != nil {
				r.t.Errorf("replay entry %d: client stopped reading: %v", i+1, err)
				return
			}
			continue
		}

		line, err := r.serverIn.ReadBytes('\n')
		if err != nil {
			r.t.Errorf("replay entry %d: client closed instead of sending %s", i+1, e.Msg)
			return
		}
		r.match(i+1, e.Msg, bytes.TrimSpace(line))
	}

	extra, err := r.serverIn.ReadBytes('\n')
	if len(bytes.TrimSpace(extra)) > 0 {
		r.t.Errorf("client sent more than the transcript: %s", extra)
	}
	if err == nil {
		r.t.Errorf("client kept stdin open after the transcript ended")
	}
}

func (r *replay) rewrite(raw json.RawMessage) json.RawMessage {
	m := message{}
	err := json.Unmarshal(raw, &m)
	if err != nil || m.Method != "" || len(m.ID) == 0 {
		return raw
	}
	id, ok := r.ids[string(m.ID)]
	if !ok {
		return raw
	}
	m.ID = id
	b, err := json.Marshal(m)
	if err != nil {
		r.t.Errorf("rewrite response: %v", err)
		return raw
	}
	return b
}

func (r *replay) match(n int, want, got json.RawMessage) {
	w, g := message{}, message{}
	err := json.Unmarshal(want, &w)
	if err != nil {
		r.t.Errorf("replay entry %d: %v", n, err)
		return
	}
	err = json.Unmarshal(got, &g)
	if err != nil {
		r.t.Errorf("replay entry %d: client sent %s: %v", n, got, err)
		return
	}

	if w.Method != "" && len(w.ID) > 0 {
		r.ids[string(w.ID)] = g.ID
	} else if !bytes.Equal(w.ID, g.ID) {
		r.t.Errorf("replay entry %d: id %s, want %s", n, g.ID, w.ID)
	}

	if w.Method != g.Method {
		r.t.Errorf("replay entry %d: client sent %s, want %s", n, got, want)
		return
	}
	if !sameJSON(w.Params, g.Params) || !sameJSON(w.Result, g.Result) {
		r.t.Errorf("replay entry %d: client sent\n%s\nwant\n%s", n, got, want)
	}
}

func sameJSON(a, b json.RawMessage) bool {
	if len(a) == 0 || len(b) == 0 {
		return len(a) == len(b)
	}
	var x, y any
	if json.Unmarshal(a, &x) != nil || json.Unmarshal(b, &y) != nil {
		return false
	}
	return reflect.DeepEqual(dropCwd(x), dropCwd(y))
}

func dropCwd(v any) any {
	m, ok := v.(map[string]any)
	if !ok {
		return v
	}
	delete(m, "cwd")
	return m
}

func testContext(t *testing.T, d time.Duration) context.Context {
	ctx, cancel := context.WithTimeout(context.Background(), d)
	t.Cleanup(cancel)
	return ctx
}

func scratchRepo(t *testing.T) string {
	dir := t.TempDir()
	git(t, dir, "init", "--initial-branch=main")
	git(t, dir, "commit", "--allow-empty", "--message", "root")
	return dir
}

func git(t *testing.T, dir string, args ...string) string {
	cmd := exec.Command("git", args...)
	cmd.Dir = dir
	cmd.Env = append(os.Environ(),
		"GIT_AUTHOR_NAME="+gitName,
		"GIT_AUTHOR_EMAIL="+gitEmail,
		"GIT_COMMITTER_NAME="+gitName,
		"GIT_COMMITTER_EMAIL="+gitEmail,
	)
	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Fatalf("git %v: %v\n%s", args, err, out)
	}
	return string(out)
}