agent_start_test.go RawHistoryFolder

1package core
2
3import (
4	"context"
5	"testing"
6)
7
8func TestBeforeStartLifecycle(t *testing.T) {
9	a := NewAgent(&queueFakeClient{}, "model", "base", Registry{})
10	calls, turns := 0, 0
11	a.BeforeStart = func(ctx context.Context, system string) string {
12		calls++
13		if system != "base" {
14			t.Fatalf("reapplied to modified prompt %q", system)
15		}
16		return system + "/extension"
17	}
18	a.BeforeTurn = func(int) (bool, string) {
19		turns++
20		if calls != 1 || a.System != "base/extension" {
21			t.Fatal("turn ran before preparation")
22		}
23		return true, ""
24	}
25	for range 2 {
26		if err := a.Prompt(context.Background(), "hello", nil, nil); err != nil {
27			t.Fatal(err)
28		}
29	}
30	if calls != 1 || turns != 3 {
31		t.Fatalf("calls=%d turns=%d", calls, turns)
32	}
33	a.BeforeTurn = nil
34	for _, reset := range []func(){
35		func() { a.Model = "other" },
36		func() { a.SetMessages(nil) },
37		func() { a.SessionID = "new-session" },
38		func() { a.SetSystem("base") },
39	} {
40		reset()
41		before := calls
42		if err := a.Continue(context.Background(), nil); err != nil {
43			t.Fatal(err)
44		}
45		if calls != before+1 || a.System != "base/extension" {
46			t.Fatal("reset did not reprepare from base")
47		}
48	}
49	a.BeforeStart = func(context.Context, string) string { return "" }
50	a.SetSystem("base")
51	if err := a.Continue(context.Background(), nil); err != nil {
52		t.Fatal(err)
53	}
54	if a.System != "" {
55		t.Fatal("empty replacement lost")
56	}
57}
58
59func TestBeforeStartCancellationRetriesPreparation(t *testing.T) {
60	a := NewAgent(nil, "model", "base", nil)
61	ctx, cancel := context.WithCancel(context.Background())
62	a.BeforeStart = func(context.Context, string) string { cancel(); return "modified" }
63	if err := a.Continue(ctx, nil); err != context.Canceled {
64		t.Fatalf("error=%v", err)
65	}
66	if a.System != "base" || a.startPrepared {
67		t.Fatal("canceled preparation committed")
68	}
69	a.BeforeStart = func(context.Context, string) string { return "retry" }
70	a.BeforeTurn = func(int) (bool, string) { return false, "test" }
71	if err := a.Continue(context.Background(), nil); err != nil {
72		t.Fatal(err)
73	}
74	if a.System != "retry" {
75		t.Fatal("preparation not retried")
76	}
77}