Vectrum
Sign In
348 lines · 11.0 KB · Go
package serve

import (
	"encoding/json"
	"strings"
	"testing"
	"time"

	"codey/internal/acp"
	"codey/internal/agent"
)

// recorder collects the updates an emitter produces, and drives its clock, so
// the coalescing rules can be tested without sleeping.
type recorder struct {
	updates []acp.Update
	clock   time.Time
}

func newRecorded() (*emitter, *recorder) {
	r := &recorder{clock: time.Unix(0, 0)}
	e := newEmitter(func(u acp.Update) { r.updates = append(r.updates, u) })
	e.now = func() time.Time { return r.clock }
	return e, r
}

// text is everything the recorder saw as assistant text, concatenated.
func (r *recorder) text() string {
	var b strings.Builder
	for _, u := range r.updates {
		if u.SessionUpdate != acp.UpdateAgentMessageChunk {
			continue
		}
		if c, ok := u.Content.(*acp.ContentBlock); ok {
			b.WriteString(c.Text)
		}
	}
	return b.String()
}

func (r *recorder) kinds() []string {
	var out []string
	for _, u := range r.updates {
		out = append(out, u.SessionUpdate)
	}
	return out
}

func TestDeltasAreCoalescedButNothingIsLost(t *testing.T) {
	e, r := newRecorded()

	// Twenty small deltas inside one flush window: a client whose inbound
	// queue is 32 deep must not have to hold twenty notifications for one
	// sentence.
	for range 20 {
		e.handle(agent.Event{Kind: agent.EventTextDelta, Text: "ab"})
	}
	if got := len(r.updates); got != 0 {
		t.Fatalf("%d notifications went out inside the window, want 0", got)
	}

	e.Flush()
	if got, want := r.text(), strings.Repeat("ab", 20); got != want {
		t.Errorf("text = %q, want %q", got, want)
	}
	if got := len(r.updates); got != 1 {
		t.Errorf("%d notifications after the flush, want 1", got)
	}
}

func TestABigDeltaGoesOutWithoutWaiting(t *testing.T) {
	e, r := newRecorded()

	e.handle(agent.Event{Kind: agent.EventTextDelta, Text: strings.Repeat("x", flushBytes+1)})

	if len(r.updates) != 1 {
		t.Fatalf("a delta past the size limit was held back: %d notifications", len(r.updates))
	}
}

func TestTimePushesTextOut(t *testing.T) {
	e, r := newRecorded()

	e.handle(agent.Event{Kind: agent.EventTextDelta, Text: "slow"})
	if len(r.updates) != 0 {
		t.Fatal("the first delta should still be buffered")
	}

	r.clock = r.clock.Add(flushAfter)
	e.handle(agent.Event{Kind: agent.EventTextDelta, Text: "er"})

	if got := r.text(); got != "slower" {
		t.Errorf("text = %q, want slower", got)
	}
}

func TestBufferedTextIsSaidBeforeAToolRuns(t *testing.T) {
	e, r := newRecorded()

	e.handle(agent.Event{Kind: agent.EventTextDelta, Text: "let me look"})
	e.handle(agent.Event{Kind: agent.EventToolStart, Tool: "bash", ID: "t1",
		Input: json.RawMessage(`{"command":"go test ./..."}`)})

	// Otherwise a sentence written just before a long command would sit in the
	// buffer until the command finished.
	want := []string{acp.UpdateAgentMessageChunk, acp.UpdateToolCall}
	if got := r.kinds(); !equal(got, want) {
		t.Fatalf("updates = %v, want %v", got, want)
	}
	if r.text() != "let me look" {
		t.Errorf("text = %q", r.text())
	}
}

func TestAToolCallAndItsResultSharePairingID(t *testing.T) {
	e, r := newRecorded()

	e.handle(agent.Event{Kind: agent.EventToolStart, Tool: "read_file", ID: "toolu_42",
		Input: json.RawMessage(`{"path":"internal/serve/serve.go"}`)})
	e.handle(agent.Event{Kind: agent.EventToolResult, Tool: "read_file", ID: "toolu_42",
		Result: "package serve"})

	if len(r.updates) != 2 {
		t.Fatalf("want a call and an update, got %v", r.kinds())
	}
	call, update := r.updates[0], r.updates[1]

	if call.ToolCallID != "toolu_42" || update.ToolCallID != "toolu_42" {
		t.Errorf("ids did not pair: %q and %q", call.ToolCallID, update.ToolCallID)
	}
	if call.Kind != acp.KindRead {
		t.Errorf("kind = %q, want %q", call.Kind, acp.KindRead)
	}
	if call.Status != acp.StatusInProgress {
		t.Errorf("status = %q", call.Status)
	}
	if !strings.Contains(call.Title, "serve.go") {
		t.Errorf("title does not say what is being read: %q", call.Title)
	}
	if update.Status != acp.StatusCompleted {
		t.Errorf("result status = %q, want completed", update.Status)
	}
	if string(call.RawInput) == "" {
		t.Error("the raw input was dropped")
	}
}

func TestAFailedToolIsReportedAsFailed(t *testing.T) {
	e, r := newRecorded()

	e.handle(agent.Event{Kind: agent.EventToolStart, Tool: "bash", ID: "t1", Input: json.RawMessage(`{"command":"false"}`)})
	e.handle(agent.Event{Kind: agent.EventToolResult, Tool: "bash", ID: "t1", Result: "exit status 1", IsError: true})

	if got := r.updates[1].Status; got != acp.StatusFailed {
		t.Errorf("status = %q, want failed", got)
	}
}

func TestACallWithNoIDStillPairs(t *testing.T) {
	// The hosted tools do not always give the observer an id, and a result
	// whose call cannot be found is a tool row that never completes.
	e, r := newRecorded()

	e.handle(agent.Event{Kind: agent.EventToolStart, Tool: "web_search"})
	e.handle(agent.Event{Kind: agent.EventToolResult, Tool: "web_search", Result: "3 results"})

	if a, b := r.updates[0].ToolCallID, r.updates[1].ToolCallID; a == "" || a != b {
		t.Errorf("ids did not pair: %q and %q", a, b)
	}
}

func TestNoticesReachTheTranscript(t *testing.T) {
	e, r := newRecorded()

	e.handle(agent.Event{Kind: agent.EventNotice, Text: "the context window is full"})

	if !strings.Contains(r.text(), "window is full") {
		t.Errorf("a notice was swallowed: %v", r.kinds())
	}
}

func TestTheEventsWithNothingToSaySayNothing(t *testing.T) {
	e, r := newRecorded()

	for _, k := range []agent.EventKind{agent.EventTurnStart, agent.EventUsage, agent.EventPaused, agent.EventResumed} {
		e.handle(agent.Event{Kind: k})
	}

	if len(r.updates) != 0 {
		t.Errorf("%v went out and should not have", r.kinds())
	}
}

func TestTitleTrimsToOneLine(t *testing.T) {
	e, r := newRecorded()

	e.handle(agent.Event{Kind: agent.EventToolStart, Tool: "bash", ID: "t1",
		Input: json.RawMessage(`{"command":"echo one\necho two"}`)})

	title := r.updates[0].Title
	if strings.Contains(title, "\n") {
		t.Errorf("a multi-line title would break the row: %q", title)
	}
	if !strings.Contains(title, "echo one") {
		t.Errorf("title lost the command: %q", title)
	}
}

func equal(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

// contentText is the text of a tool call update's content, which is a list
// rather than the single block a chunk carries.
func contentText(t *testing.T, u acp.Update) string {
	t.Helper()
	items, ok := u.Content.([]acp.ToolCallContent)
	if !ok {
		t.Fatalf("content is not tool call content: %#v", u.Content)
	}
	var b strings.Builder
	for _, it := range items {
		if it.Content != nil {
			b.WriteString(it.Content.Text)
		}
	}
	return b.String()
}

func TestToolOutputGoesOutWhileTheToolIsStillRunning(t *testing.T) {
	e, r := newRecorded()

	e.handle(agent.Event{Kind: agent.EventToolStart, Tool: "bash", ID: "t1",
		Input: json.RawMessage(`{"command":"go test ./..."}`)})
	e.handle(agent.Event{Kind: agent.EventToolDelta, Tool: "bash", ID: "t1", Text: strings.Repeat("ok\n", 200)})

	if len(r.updates) != 2 {
		t.Fatalf("the output did not go out until the tool finished: %v", r.kinds())
	}
	live := r.updates[1]
	if live.SessionUpdate != acp.UpdateToolCallUpdate || live.ToolCallID != "t1" {
		t.Fatalf("update = %+v", live)
	}
	if live.Status != acp.StatusInProgress {
		t.Errorf("status = %q, want in_progress — a completed row stops updating", live.Status)
	}
	if !strings.Contains(contentText(t, live), "ok") {
		t.Errorf("content = %q", contentText(t, live))
	}
}

func TestProgressUpdatesCarryEverythingSoFar(t *testing.T) {
	// A tool_call_update replaces the content of the call it names rather than
	// appending to it, so a delta-shaped update would show only the last line.
	e, r := newRecorded()

	e.handle(agent.Event{Kind: agent.EventToolStart, Tool: "bash", ID: "t1", Input: json.RawMessage(`{"command":"x"}`)})
	e.handle(agent.Event{Kind: agent.EventToolDelta, Tool: "bash", ID: "t1", Text: strings.Repeat("a", progressBytes)})
	r.clock = r.clock.Add(progressAfter)
	e.handle(agent.Event{Kind: agent.EventToolDelta, Tool: "bash", ID: "t1", Text: "b"})

	last := contentText(t, r.updates[len(r.updates)-1])
	if !strings.HasPrefix(last, "aaa") || !strings.HasSuffix(last, "b") {
		t.Errorf("the update did not carry the whole output: %q…%q", last[:min(8, len(last))], last[max(0, len(last)-4):])
	}
}

func TestProgressIsThrottled(t *testing.T) {
	e, r := newRecorded()

	e.handle(agent.Event{Kind: agent.EventToolStart, Tool: "bash", ID: "t1", Input: json.RawMessage(`{"command":"x"}`)})
	before := len(r.updates)

	// A build writing a line at a time, inside one window.
	for range 50 {
		e.handle(agent.Event{Kind: agent.EventToolDelta, Tool: "bash", ID: "t1", Text: "line\n"})
	}

	// 50 lines of five bytes is 250, just under the size trigger, and no time
	// has passed: one notification per line would be fifty.
	if got := len(r.updates) - before; got > 1 {
		t.Errorf("%d notifications for 50 lines inside one window", got)
	}
}

func TestProgressIsBoundedToATail(t *testing.T) {
	e, r := newRecorded()

	e.handle(agent.Event{Kind: agent.EventToolStart, Tool: "bash", ID: "t1", Input: json.RawMessage(`{"command":"x"}`)})
	e.handle(agent.Event{Kind: agent.EventToolDelta, Tool: "bash", ID: "t1",
		Text: strings.Repeat("x", progressMaxBytes*2) + "END"})

	body := contentText(t, r.updates[len(r.updates)-1])
	if len(body) > progressMaxBytes+64 {
		t.Errorf("an unbounded %d bytes went over the wire", len(body))
	}
	if !strings.HasSuffix(body, "END") {
		t.Error("the tail is the useful part and it was the part dropped")
	}
	if !strings.Contains(body, "truncated") {
		t.Error("the truncation was silent")
	}
}

func TestTheFinalResultReplacesTheLiveView(t *testing.T) {
	e, r := newRecorded()

	e.handle(agent.Event{Kind: agent.EventToolStart, Tool: "bash", ID: "t1", Input: json.RawMessage(`{"command":"x"}`)})
	e.handle(agent.Event{Kind: agent.EventToolDelta, Tool: "bash", ID: "t1", Text: strings.Repeat("noise\n", 100)})
	e.handle(agent.Event{Kind: agent.EventToolResult, Tool: "bash", ID: "t1", Result: "exit code: 0 (1ms)"})

	last := r.updates[len(r.updates)-1]
	if last.Status != acp.StatusCompleted {
		t.Fatalf("status = %q", last.Status)
	}
	if got := contentText(t, last); got != "exit code: 0 (1ms)" {
		t.Errorf("content = %q, want the finished report", got)
	}
}

func TestProgressForAFinishedCallStartsCleanNextTime(t *testing.T) {
	// Two commands in one turn must not show each other's output.
	e, r := newRecorded()

	for _, id := range []string{"t1", "t2"} {
		e.handle(agent.Event{Kind: agent.EventToolStart, Tool: "bash", ID: id, Input: json.RawMessage(`{"command":"x"}`)})
		e.handle(agent.Event{Kind: agent.EventToolDelta, Tool: "bash", ID: id, Text: strings.Repeat(id, 200)})
		e.handle(agent.Event{Kind: agent.EventToolResult, Tool: "bash", ID: id, Result: "done " + id})
		r.clock = r.clock.Add(progressAfter)
	}

	for _, u := range r.updates {
		if u.ToolCallID != "t2" || u.SessionUpdate != acp.UpdateToolCallUpdate || u.Status != acp.StatusInProgress {
			continue
		}
		if strings.Contains(contentText(t, u), "t1") {
			t.Error("the second command's progress carried the first command's output")
		}
	}
}