1package core
2
3import (
4 "context"
5 "fmt"
6 "testing"
7 "time"
8)
9
10func TestBeforeStartDiscardsSupersededPreparation(t *testing.T) {
11 for _, tc := range []struct {
12 name string
13 reset func(*Agent)
14 wantBase string
15 }{
16 {"new prompt", func(a *Agent) { a.SetSystem("new base") }, "new base"},
17 {"same prompt reset", func(a *Agent) { a.SetSystem("base") }, "base"},
18 {"ABA prompt reset", func(a *Agent) { a.SetSystem("temporary"); a.SetSystem("base") }, "base"},
19 {"multiple rebuilds", func(a *Agent) { a.SetSystem("intermediate"); a.SetSystem("latest") }, "latest"},
20 {"clear", func(a *Agent) { a.SetMessages(nil) }, "base"},
21 {"model change", func(a *Agent) { a.Model = "new model" }, "base"},
22 {"session change", func(a *Agent) { a.SessionID = "new session" }, "base"},
23 } {
24 t.Run(tc.name, func(t *testing.T) {
25 a := NewAgent(nil, "model", "base", nil)
26 entered, release := make(chan struct{}), make(chan struct{})
27 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
28 defer cancel()
29 var inputs []string
30 a.BeforeStart = func(ctx context.Context, system string) string {
31 inputs = append(inputs, system)
32 if len(inputs) == 1 {
33 close(entered)
34 select {
35 case <-release:
36 case <-ctx.Done():
37 }
38 }
39 return fmt.Sprintf("%s/replacement-%d", system, len(inputs))
40 }
41 a.BeforeTurn = func(int) (bool, string) { return false, "test" }
42 done := make(chan struct{})
43 var runErr error
44 go func() { defer close(done); runErr = a.Continue(ctx, nil) }()
45 t.Cleanup(func() { cancel(); <-done })
46 select {
47 case <-entered:
48 case <-ctx.Done():
49 t.Fatal("hook did not start")
50 }
51 // The hook is paused without the agent lock. Reset must remain responsive,
52 // including when the base string has not changed at all.
53 tc.reset(a)
54 close(release)
55 select {
56 case <-done:
57 if runErr != nil {
58 t.Fatal(runErr)
59 }
60 case <-ctx.Done():
61 t.Fatal("preparation did not finish")
62 }
63 if len(inputs) != 2 || inputs[1] != tc.wantBase {
64 t.Fatalf("hook inputs=%q, want [base %q]", inputs, tc.wantBase)
65 }
66 want := tc.wantBase + "/replacement-2"
67 if a.System != want || a.BaseSystem() != tc.wantBase {
68 t.Fatalf("system=%q base=%q", a.System, a.BaseSystem())
69 }
70 if a.startModel != a.Model || a.startSession != a.SessionID {
71 t.Fatal("prepared metadata does not match current runtime")
72 }
73 // The fresh result is cached; normal continuation must not run it again.
74 if err := a.Continue(ctx, nil); err != nil {
75 t.Fatal(err)
76 }
77 if len(inputs) != 2 {
78 t.Fatal("fresh preparation was not cached")
79 }
80 })
81 }
82}
83
84func TestBeforeStartCancellationDuringRepreparation(t *testing.T) {
85 a := NewAgent(nil, "model", "base", nil)
86 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
87 defer cancel()
88 first, second, release := make(chan struct{}), make(chan struct{}), make(chan struct{})
89 calls := 0
90 a.BeforeStart = func(ctx context.Context, system string) string {
91 calls++
92 if calls == 1 {
93 close(first)
94 select {
95 case <-release:
96 case <-ctx.Done():
97 }
98 } else {
99 close(second)
100 <-ctx.Done()
101 }
102 return "discard this replacement"
103 }
104 a.BeforeTurn = func(int) (bool, string) { return false, "test" }
105 done := make(chan struct{})
106 var runErr error
107 go func() { defer close(done); runErr = a.Continue(ctx, nil) }()
108 t.Cleanup(func() { cancel(); <-done })
109 select {
110 case <-first:
111 case <-ctx.Done():
112 t.Fatal("first hook did not start")
113 }
114 a.SetSystem("new base")
115 close(release)
116 select {
117 case <-second:
118 case <-done:
119 t.Fatal("stale preparation was committed instead of retried")
120 case <-ctx.Done():
121 t.Fatal("second hook did not start")
122 }
123 cancel()
124 <-done
125 if runErr != context.Canceled {
126 t.Fatalf("error=%v", runErr)
127 }
128 if a.System != "new base" || a.startPrepared {
129 t.Fatal("cancellation committed a stale result")
130 }
131}