Vectrum
Sign In
118 lines · 3.4 KB · Go
package agent

import (
	"bytes"
	"context"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"harness/internal/llm"
)

func TestSaveAndResumeSession(t *testing.T) {
	dir := t.TempDir()
	f := newFakeAPI(t,
		msg("tool_use", toolUseBlock("tu_1", "boom", `{}`)),
		msg("end_turn", textBlock("first answer")),
	)
	a := newTestAgent(t, f, t.TempDir(), &bytes.Buffer{})
	if _, err := a.Run(context.Background(), "the original task"); err != nil {
		t.Fatalf("Run: %v", err)
	}
	if err := a.Save(dir, "sess-1", "/repo", "claude-opus-5"); err != nil {
		t.Fatalf("Save: %v", err)
	}

	loaded, err := ResolveSession(dir, "sess-1")
	if err != nil {
		t.Fatalf("ResolveSession: %v", err)
	}
	if loaded.Model != "claude-opus-5" || loaded.Root != "/repo" || loaded.Turns != 2 {
		t.Errorf("session = %+v", loaded)
	}
	if len(loaded.Messages) != len(a.State().History()) {
		t.Errorf("%d messages saved, want %d", len(loaded.Messages), len(a.State().History()))
	}

	// A resumed agent continues the same conversation.
	f2 := newFakeAPI(t, msg("end_turn", textBlock("second answer")))
	b := newTestAgent(t, f2, t.TempDir(), &bytes.Buffer{})
	b.Restore(loaded)
	if _, err := b.Run(context.Background(), "and now this"); err != nil {
		t.Fatalf("Run after restore: %v", err)
	}
	sent := f2.rawRequest(0)
	if !strings.Contains(sent, "the original task") {
		t.Errorf("resumed request lost the earlier conversation:\n%s", sent)
	}
	if !strings.Contains(sent, "and now this") {
		t.Error("resumed request is missing the new prompt")
	}
	if _, turns := b.State().Usage(); turns != 3 {
		t.Errorf("turns = %d, want the restored 2 plus 1", turns)
	}
}

func TestSavePreservesCreatedTime(t *testing.T) {
	dir := t.TempDir()
	f := newFakeAPI(t, msg("end_turn", textBlock("x")), msg("end_turn", textBlock("y")))
	a := newTestAgent(t, f, t.TempDir(), &bytes.Buffer{})

	a.Run(context.Background(), "one")
	a.Save(dir, "s", "/r", "m")
	first, _ := LoadSession(filepath.Join(dir, "s.json"))

	time.Sleep(5 * time.Millisecond)
	a.Run(context.Background(), "two")
	a.Save(dir, "s", "/r", "m")
	second, _ := LoadSession(filepath.Join(dir, "s.json"))

	if !second.Created.Equal(first.Created) {
		t.Errorf("created moved from %v to %v across saves", first.Created, second.Created)
	}
	if !second.Updated.After(first.Updated) {
		t.Error("updated should move forward")
	}
}

func TestResolveLatestSession(t *testing.T) {
	dir := t.TempDir()
	for _, id := range []string{"older", "newer"} {
		s := Session{ID: id, Messages: []llm.Message{{Role: llm.RoleUser}}}
		writeSession(t, dir, s)
		time.Sleep(10 * time.Millisecond)
	}
	got, err := ResolveSession(dir, "latest")
	if err != nil {
		t.Fatalf("ResolveSession: %v", err)
	}
	if got.ID != "newer" {
		t.Errorf("resolved %q, want the most recently updated", got.ID)
	}
	if _, err := ResolveSession(t.TempDir(), "latest"); err == nil {
		t.Error("an empty session directory should be an error, not an empty session")
	}
	if _, err := ResolveSession(dir, "nonexistent"); err == nil {
		t.Error("an unknown id should be an error")
	}
}

func writeSession(t *testing.T, dir string, s Session) {
	t.Helper()
	a := &Agent{state: NewState()}
	a.state.Replace(s.Messages)
	if err := a.Save(dir, s.ID, "/r", "m"); err != nil {
		t.Fatal(err)
	}
}

func TestLoadSessionRejectsGarbage(t *testing.T) {
	dir := t.TempDir()
	os.WriteFile(filepath.Join(dir, "broken.json"), []byte("{not json"), 0o644)
	if _, err := ResolveSession(dir, "broken"); err == nil {
		t.Error("want an error for a corrupt session file")
	}
}