byre/internal/claude/session_test.go

package claude_test

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

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

const (
	testModel = "sonnet"
	gitName   = "Byre Test"
	gitEmail  = "test@byre.invalid"
)

func TestCommitIdentityHooksAndMCP(t *testing.T) {
	t.Parallel()
	ctx := testContext(t)
	repo := scratchRepo(t)

	finished := 0
	mu := sync.Mutex{}
	finish := mcp.Tool{
		Name:        "finish",
		Description: "Call this once the task is complete.",
		InputSchema: json.RawMessage(`{"type": "object", "properties": {}}`),
		Handler: func(ctx context.Context, args json.RawMessage) (string, error) {
			mu.Lock()
			defer mu.Unlock()
			finished++
			return "noted", nil
		},
	}

	tools := []string{}
	s := start(t, ctx, claude.Options{
		Cwd:        repo,
		MCPServers: []*mcp.Server{{Name: "byre", Version: "0.0.0", Tools: []mcp.Tool{finish}}},
		Hooks: map[string][]claude.HookMatcher{
			"PostToolUse": {{Func: func(ctx context.Context, input json.RawMessage) (map[string]any, error) {
				in := struct {
					ToolName string `json:"tool_name"`
				}{}
				err := json.Unmarshal(input, &in)
				if err != nil {
					return nil, err
				}
				mu.Lock()
				defer mu.Unlock()
				tools = append(tools, in.ToolName)
				return map[string]any{}, nil
			}}},
		},
	})

	send(t, s, `Create a file hello.txt containing exactly the word hi, then git add and git commit it with the message "add hello", then call the finish tool.`)
	waitResult(t, ctx, s, nil)

	mu.Lock()
	defer mu.Unlock()
	t.Logf("tools=%v finished=%d", tools, finished)

	if finished == 0 {
		t.Fatalf("the MCP finish tool was never called")
	}
	if len(tools) == 0 {
		t.Fatalf("PostToolUse hook never fired")
	}
	if slices.Contains(tools, "ToolSearch") {
		t.Fatalf("byre MCP tools were deferred behind ToolSearch")
	}

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

func TestStopHookBlocks(t *testing.T) {
	t.Parallel()
	ctx := testContext(t)

	blocked := false
	mu := sync.Mutex{}
	s := start(t, ctx, claude.Options{
		Cwd: scratchRepo(t),
		Hooks: map[string][]claude.HookMatcher{
			"Stop": {{Func: func(ctx context.Context, input json.RawMessage) (map[string]any, error) {
				mu.Lock()
				defer mu.Unlock()
				if blocked {
					return map[string]any{}, nil
				}
				blocked = true
				return map[string]any{"decision": "block", "reason": "Before stopping, reply with the single word PINEAPPLE."}, nil
			}}},
		},
	})

	send(t, s, "Reply with the single word ok.")
	text := waitResult(t, ctx, s, nil)

	if !blocked {
		t.Fatalf("Stop hook never fired")
	}
	if !strings.Contains(text, "PINEAPPLE") {
		t.Fatalf("Stop hook block reason not acted on; assistant said %q", text)
	}
}

func TestTextDeltas(t *testing.T) {
	t.Parallel()
	ctx := testContext(t)
	s := start(t, ctx, claude.Options{Cwd: scratchRepo(t)})

	send(t, s, "Without using any tools, count from one to twenty in words, separated by spaces.")
	deltas := []string{}
	text := waitResult(t, ctx, s, func(ev claude.Event) {
		if ev.Type != "stream_event" {
			return
		}
		d, err := ev.TextDelta()
		if err != nil {
			t.Fatalf("parse stream event: %v", err)
		}
		if d != "" {
			deltas = append(deltas, d)
		}
	})

	if len(deltas) < 2 {
		t.Fatalf("expected several text deltas, got %d", len(deltas))
	}
	if strings.Join(deltas, "") != text {
		t.Fatalf("deltas %q don't add up to the final text %q", deltas, text)
	}
}

func TestInterrupt(t *testing.T) {
	t.Parallel()
	ctx := testContext(t)
	s := start(t, ctx, claude.Options{Cwd: scratchRepo(t)})

	send(t, s, "Run the bash command `ping -c 120 127.0.0.1` in the foreground (not in the background) and wait for it to finish, then say done.")

	interruptedAt := time.Time{}
	waitResult(t, ctx, s, func(ev claude.Event) {
		if ev.Type == "user" {
			blocks, err := ev.Blocks()
			if err != nil {
				t.Fatalf("parse user: %v", err)
			}
			for _, b := range blocks {
				if b.Type == "tool_result" && (!b.IsError || strings.Contains(string(b.Content), "Blocked")) {
					t.Fatalf("command finished or was refused instead of being interrupted: %s", b.Content)
				}
			}
		}

		if !interruptedAt.IsZero() || !hasToolUse(t, ev) {
			return
		}
		time.Sleep(2 * time.Second)
		interruptedAt = time.Now()
		err := s.Interrupt(ctx)
		if err != nil {
			t.Errorf("interrupt: %v", err)
		}
	})

	if interruptedAt.IsZero() {
		t.Fatalf("never saw a tool use to interrupt")
	}
	elapsed := time.Since(interruptedAt)
	t.Logf("turn ended %s after interrupt", elapsed)
	if elapsed > 30*time.Second {
		t.Fatalf("interrupt took %s to end the turn", elapsed)
	}

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

func TestImageInput(t *testing.T) {
	t.Parallel()
	ctx := testContext(t)
	s := start(t, ctx, claude.Options{Cwd: scratchRepo(t)})

	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 {
		t.Fatalf("encode png: %v", err)
	}

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

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

func TestResume(t *testing.T) {
	t.Parallel()
	ctx := testContext(t)
	repo := scratchRepo(t)

	first, err := claude.Start(ctx, options(t, "first", claude.Options{Cwd: repo}))
	if err != nil {
		t.Fatalf("start: %v", err)
	}
	send(t, first, "Remember the codeword AZURE-17. Reply with the single word ok.")
	waitResult(t, ctx, first, nil)
	id := first.ID()
	closeSession(t, first)

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

	resumed, err := claude.Start(ctx, options(t, "resumed", claude.Options{Cwd: repo, Resume: id}))
	if err != nil {
		t.Fatalf("resume: %v", err)
	}
	t.Cleanup(func() { closeSession(t, resumed) })

	send(t, resumed, "What was the codeword? Reply with just the codeword.")
	text := waitResult(t, ctx, resumed, nil)

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

func TestForkAt(t *testing.T) {
	t.Parallel()
	ctx := testContext(t)

	first, err := claude.Start(ctx, options(t, "first", claude.Options{Cwd: scratchRepo(t)}))
	if err != nil {
		t.Fatalf("start: %v", err)
	}
	send(t, first, "Run the bash command `echo APPLE-1`. After it finishes, run the bash command `echo BANANA-2`. Then reply with the single word ok.")
	cut := ""
	waitResult(t, ctx, first, func(ev claude.Event) {
		if cut != "" || ev.Type != "user" {
			return
		}
		blocks, err := ev.Blocks()
		if err != nil {
			t.Fatalf("parse user: %v", err)
		}
		for _, b := range blocks {
			if b.Type == "tool_result" && strings.Contains(string(b.Content), "APPLE-1") {
				cut, err = ev.UUID()
				if err != nil {
					t.Fatalf("parse uuid: %v", err)
				}
			}
		}
	})
	id := first.ID()
	closeSession(t, first)

	if cut == "" {
		t.Fatalf("never saw the APPLE-1 tool result")
	}

	forkID := "8a1b2c3d-0000-4000-8000-" + strings.ReplaceAll(time.Now().Format("150405.000000"), ".", "")
	fork, err := claude.Start(ctx, options(t, "fork", claude.Options{
		Cwd:         scratchRepo(t),
		SessionID:   forkID,
		Resume:      id,
		ResumeAt:    cut,
		ForkSession: true,
	}))
	if err != nil {
		t.Fatalf("fork: %v", err)
	}
	t.Cleanup(func() { closeSession(t, fork) })

	send(t, fork, "Without running anything, list the exact output of every bash command you have run in this conversation, one per line, and nothing else.")
	text := waitResult(t, ctx, fork, nil)

	if !strings.Contains(text, "APPLE-1") {
		t.Fatalf("fork lost the conversation before the cut; assistant said %q", text)
	}
	if strings.Contains(text, "BANANA-2") {
		t.Fatalf("fork kept the conversation after the cut; assistant said %q", text)
	}
	if fork.ID() != forkID {
		t.Fatalf("fork has session id %s, want %s", fork.ID(), forkID)
	}
}

func testContext(t *testing.T) context.Context {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
	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 options(t *testing.T, suffix string, opts claude.Options) claude.Options {
	name := t.Name()
	if suffix != "" {
		name += "-" + suffix
	}

	err := os.MkdirAll(filepath.Join("testdata", "transcripts"), 0o755)
	if err != nil {
		t.Fatalf("mkdir transcripts: %v", err)
	}

	f, err := os.Create(filepath.Join("testdata", "transcripts", name+".jsonl"))
	if err != nil {
		t.Fatalf("create transcript: %v", err)
	}
	t.Cleanup(func() { f.Close() })

	opts.Model = testModel
	opts.GitName = gitName
	opts.GitEmail = gitEmail
	opts.Transcript = f
	return opts
}

func start(t *testing.T, ctx context.Context, opts claude.Options) *claude.Session {
	s, err := claude.Start(ctx, options(t, "", opts))
	if err != nil {
		t.Fatalf("start: %v", err)
	}
	t.Cleanup(func() { closeSession(t, s) })
	return s
}

func closeSession(t *testing.T, s *claude.Session) {
	err := s.Close()
	if err != nil {
		t.Logf("claude exit: %v", err)
	}
}

func send(t *testing.T, s *claude.Session, text string) {
	err := s.Send(text)
	if err != nil {
		t.Fatalf("send: %v", err)
	}
}

func waitResult(t *testing.T, ctx context.Context, s *claude.Session, onEvent func(claude.Event)) string {
	text := []string{}
	for {
		select {
		case ev, ok := <-s.Events():
			if !ok {
				t.Fatalf("claude exited before result")
			}
			if onEvent != nil {
				onEvent(ev)
			}
			if ev.Type == "assistant" {
				blocks, err := ev.Blocks()
				if err != nil {
					t.Fatalf("parse assistant: %v", err)
				}
				for _, b := range blocks {
					if b.Type == "text" {
						text = append(text, b.Text)
					}
				}
			}
			if ev.Type == "result" {
				res, err := ev.Result()
				if err != nil {
					t.Fatalf("parse result: %v", err)
				}
				t.Logf("result %s: %d turns, %dms, $%.4f", res.Subtype, res.NumTurns, res.DurationMS, res.TotalCostUSD)
				return strings.Join(text, "\n")
			}
		case <-ctx.Done():
			t.Fatalf("timed out waiting for result")
		}
	}
}

func hasToolUse(t *testing.T, ev claude.Event) bool {
	if ev.Type != "assistant" {
		return false
	}
	blocks, err := ev.Blocks()
	if err != nil {
		t.Fatalf("parse assistant: %v", err)
	}
	for _, b := range blocks {
		if b.Type == "tool_use" {
			return true
		}
	}
	return false
}

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