Vectrum
Sign In
248 lines · 6.5 KB · Go
package tui

import (
	"context"
	"encoding/json"
	"sync"
	"testing"
	"time"

	tea "github.com/charmbracelet/bubbletea"

	"github.com/johannSo/codey/internal/agent"
	"github.com/johannSo/codey/internal/approval"
	"github.com/johannSo/codey/internal/ask"
)

// recorder stands in for tea.Program: it captures what the bridge sends
// without needing a terminal.
type recorder struct {
	mu   sync.Mutex
	msgs []tea.Msg
	// onSend runs for each message, so a test can answer an approval from the
	// "UI side" the way the Update loop would.
	onSend func(tea.Msg)
}

func (r *recorder) Send(m tea.Msg) {
	r.mu.Lock()
	r.msgs = append(r.msgs, m)
	f := r.onSend
	r.mu.Unlock()
	if f != nil {
		f(m)
	}
}

func (r *recorder) all() []tea.Msg {
	r.mu.Lock()
	defer r.mu.Unlock()
	return append([]tea.Msg(nil), r.msgs...)
}

func TestObserveTranslatesEvents(t *testing.T) {
	r := &recorder{}
	b := &bridge{p: r}

	b.Observe(agent.Event{Kind: agent.EventTurnStart, Turn: 2})
	b.Observe(agent.Event{Kind: agent.EventTextDelta, Text: "hi"})
	b.Observe(agent.Event{Kind: agent.EventToolStart, Tool: "grep", Input: json.RawMessage(`{"pattern":"x"}`)})
	b.Observe(agent.Event{Kind: agent.EventToolResult, Tool: "grep", Result: "x.go:1", IsError: false})
	b.Observe(agent.Event{Kind: agent.EventCompact})

	msgs := r.all()
	if len(msgs) != 5 {
		t.Fatalf("got %d messages, want 5", len(msgs))
	}
	if m, ok := msgs[0].(evTurnStart); !ok || m.turn != 2 {
		t.Errorf("msgs[0] = %#v, want evTurnStart{2}", msgs[0])
	}
	if m, ok := msgs[2].(evToolStart); !ok || m.input != "pattern=x" {
		t.Errorf("msgs[2] = %#v, want summarised tool input", msgs[2])
	}
	if _, ok := msgs[4].(evCompact); !ok {
		t.Errorf("msgs[4] = %#v, want evCompact", msgs[4])
	}
}

// The agent goroutine must block until the UI answers, and must get exactly
// the decision the user picked.
func TestPromptBlocksUntilAnswered(t *testing.T) {
	r := &recorder{}
	r.onSend = func(m tea.Msg) {
		if am, ok := m.(approvalMsg); ok {
			go func() {
				time.Sleep(10 * time.Millisecond)
				am.reply <- approval.AlwaysThisTool
			}()
		}
	}
	b := &bridge{p: r}

	start := time.Now()
	got, err := b.Prompt(context.Background(), approval.Request{Tool: "bash", Summary: "run a shell command"})
	if err != nil {
		t.Fatalf("Prompt: %v", err)
	}
	if got != approval.AlwaysThisTool {
		t.Errorf("decision = %v, want AlwaysThisTool", got)
	}
	if time.Since(start) < 10*time.Millisecond {
		t.Error("Prompt returned before the UI answered")
	}
}

// A cancelled turn must not strand the agent on a channel no one will read.
// This is the deadlock the whole design has to avoid.
func TestPromptUnblocksOnCancel(t *testing.T) {
	r := &recorder{} // never answers
	b := &bridge{p: r}

	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		time.Sleep(10 * time.Millisecond)
		cancel()
	}()

	done := make(chan struct{})
	var got approval.Decision
	var err error
	go func() {
		got, err = b.Prompt(ctx, approval.Request{Tool: "bash"})
		close(done)
	}()

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("Prompt did not return after its context was cancelled")
	}
	if got != approval.Deny {
		t.Errorf("decision = %v, want Deny on cancellation", got)
	}
	if err == nil {
		t.Error("want a non-nil error on cancellation")
	}
}

func TestPromptRefusesAnAlreadyCancelledContext(t *testing.T) {
	r := &recorder{}
	b := &bridge{p: r}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	got, err := b.Prompt(ctx, approval.Request{Tool: "bash"})
	if got != approval.Deny || err == nil {
		t.Errorf("Prompt = (%v, %v), want (Deny, error)", got, err)
	}
	if len(r.all()) != 0 {
		t.Error("a cancelled prompt still bothered the UI")
	}
}

// The reply channel is buffered, so a UI that answers a request whose agent
// has already given up must not block forever.
func TestReplyChannelNeverBlocksTheUI(t *testing.T) {
	r := &recorder{}
	b := &bridge{p: r}
	ctx, cancel := context.WithCancel(context.Background())

	var captured approvalMsg
	got := make(chan struct{})
	r.onSend = func(m tea.Msg) {
		if am, ok := m.(approvalMsg); ok {
			captured = am
			close(got)
		}
	}
	go b.Prompt(ctx, approval.Request{Tool: "bash"})
	<-got
	cancel()
	time.Sleep(20 * time.Millisecond) // let Prompt give up

	done := make(chan struct{})
	go func() {
		captured.reply <- approval.Once // must not block
		close(done)
	}()
	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("answering an abandoned approval blocked the UI")
	}
}

// The bridge must satisfy the interfaces the agent and policy expect. If these
// signatures drift, this stops compiling, which is the point.
var (
	_ = func(b *bridge) { var _ func(agent.Event) = b.Observe }
	_ = func(b *bridge) { var _ approval.Prompt = b.Prompt }
)

// A question blocks the agent the same way an approval does, and comes back
// with exactly what the user picked.
func TestAskBlocksUntilAnswered(t *testing.T) {
	r := &recorder{}
	r.onSend = func(m tea.Msg) {
		if am, ok := m.(askMsg); ok {
			go func() {
				time.Sleep(10 * time.Millisecond)
				am.reply <- ask.Response{Answers: []ask.Answer{{Chosen: []string{"Postgres"}}}}
			}()
		}
	}
	b := &bridge{p: r}

	start := time.Now()
	got, err := b.Ask(context.Background(), ask.Request{Questions: []ask.Question{{Header: "Storage"}}})
	if err != nil {
		t.Fatalf("Ask: %v", err)
	}
	if len(got.Answers) != 1 || got.Answers[0].Chosen[0] != "Postgres" {
		t.Errorf("response = %+v", got)
	}
	if time.Since(start) < 10*time.Millisecond {
		t.Error("Ask returned before the UI answered")
	}
}

// The same deadlock the approval path has to avoid: an interrupted turn must
// not leave the agent parked on a channel nobody will read.
func TestAskUnblocksOnCancel(t *testing.T) {
	b := &bridge{p: &recorder{}} // never answers
	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		time.Sleep(10 * time.Millisecond)
		cancel()
	}()

	done := make(chan struct{})
	var err error
	go func() {
		_, err = b.Ask(ctx, ask.Request{})
		close(done)
	}()

	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("Ask did not return after its context was cancelled")
	}
	if err == nil {
		t.Error("a cancelled question is an error, not an empty answer")
	}
}

func TestAskRefusesAnAlreadyCancelledContext(t *testing.T) {
	r := &recorder{}
	b := &bridge{p: r}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	if _, err := b.Ask(ctx, ask.Request{}); err == nil {
		t.Error("want an error for a context that was cancelled before the question")
	}
	if len(r.all()) != 0 {
		t.Error("a cancelled question still bothered the UI")
	}
}