1package agent
2
3import (
4 "bytes"
5 "fmt"
6 "slices"
7 "strings"
8 "testing"
9
10 "github.com/bnema/zut/packages/agent/modes"
11 "github.com/bnema/zut/packages/core"
12 "github.com/bnema/zut/packages/provider"
13)
14
15// sessionWithPlan writes a session whose persisted meta carries steps, then
16// closes it and returns the path for a resume.
17func sessionWithPlan(t *testing.T, providerName, model string, steps []core.PlanStep) string {
18 t.Helper()
19 sess, err := core.NewSession(t.TempDir(), t.TempDir(), providerName, model, "test")
20 if err != nil {
21 t.Fatal(err)
22 }
23 if err := sess.AppendMessage(provider.Message{
24 Role: provider.RoleUser,
25 Content: []provider.Content{provider.TextBlock{Text: "resumed transcript"}},
26 }); err != nil {
27 t.Fatal(err)
28 }
29 if err := sess.UpdatePlan(&core.SessionPlan{Steps: steps}); err != nil {
30 t.Fatal(err)
31 }
32 if err := sess.Close(); err != nil {
33 t.Fatal(err)
34 }
35 return sess.Path
36}
37
38// TestResumeSeedsAgentPlanFromSessionMeta covers the merge-blocking bug: without
39// seeding, a resumed session's file holds steps while core holds none.
40func TestResumeSeedsAgentPlanFromSessionMeta(t *testing.T) {
41 want := []core.PlanStep{
42 {Step: "one", Status: core.PlanPending},
43 {Step: "two", Status: core.PlanInProgress},
44 }
45 path := sessionWithPlan(t, "stored-provider", "stored-model", want)
46 sess, _, err := core.OpenSession(path)
47 if err != nil {
48 t.Fatal(err)
49 }
50 defer sess.Close()
51
52 current := core.NewAgent(nil, "current-model", "", nil)
53 candidate, err := applySessionResume(sess, current, "current-provider", "current-model", false, false, func(providerName, model string) (*core.Agent, string, string, error) {
54 return core.NewAgent(nil, model, "", nil), providerName, model, nil
55 })
56 if err != nil {
57 t.Fatal(err)
58 }
59 defer candidate.session.Close()
60 if !candidate.rebuilt {
61 t.Fatal("resume did not rebuild the agent")
62 }
63 if got := candidate.agent.CurrentPlan(); !slices.Equal(got, want) {
64 t.Fatalf("resumed plan = %#v, want %#v", got, want)
65 }
66}
67
68// TestResumeSeedsReusedAgentPlan covers the same-provider path, where the live
69// agent is reused rather than rebuilt: it still has to be seeded.
70func TestResumeSeedsReusedAgentPlan(t *testing.T) {
71 want := []core.PlanStep{{Step: "only", Status: core.PlanInProgress}}
72 path := sessionWithPlan(t, "provider", "model", want)
73 sess, _, err := core.OpenSession(path)
74 if err != nil {
75 t.Fatal(err)
76 }
77 defer sess.Close()
78
79 current := core.NewAgent(nil, "model", "", nil)
80 candidate, err := applySessionResume(sess, current, "provider", "model", false, false, func(string, string) (*core.Agent, string, string, error) {
81 t.Fatal("same provider/model unexpectedly rebuilt the agent")
82 return nil, "", "", nil
83 })
84 if err != nil {
85 t.Fatal(err)
86 }
87 defer candidate.session.Close()
88 if candidate.agent != current {
89 t.Fatalf("reused candidate agent = %p, want the live agent %p", candidate.agent, current)
90 }
91 if got := current.CurrentPlan(); !slices.Equal(got, want) {
92 t.Fatalf("reused agent plan = %#v, want %#v", got, want)
93 }
94}
95
96func TestRPCClearDropsPlanState(t *testing.T) {
97 var out bytes.Buffer
98 ag := &core.Agent{}
99 ag.SetPlan([]core.PlanStep{{Step: "one", Status: core.PlanPending}})
100 s := &rpcServer{agent: ag, out: &out}
101
102 s.dispatch("clear", "1", []byte(`{"type":"clear","id":"1"}`))
103
104 if got := ag.CurrentPlan(); got != nil {
105 t.Fatalf("plan after rpc clear = %#v, want nil", got)
106 }
107 if !strings.Contains(out.String(), `"success":true`) {
108 t.Fatalf("rpc clear response = %q", out.String())
109 }
110}
111
112// appendPlanCall appends a plan tool call paired with a successful result. The
113// load-time replay only counts a call that has a non-error result.
114func appendPlanCall(t *testing.T, sess *core.Session, callID, name, args string) {
115 t.Helper()
116 if err := sess.AppendMessage(provider.Message{
117 Role: provider.RoleAssistant,
118 Content: []provider.Content{provider.ToolCallBlock{ID: callID, Name: name, Arguments: []byte(args)}},
119 }); err != nil {
120 t.Fatal(err)
121 }
122 if err := sess.AppendMessage(provider.Message{
123 Role: provider.RoleTool,
124 Content: []provider.Content{provider.ToolResultBlock{
125 CallID: callID,
126 Content: []provider.Content{provider.TextBlock{Text: "Plan updated"}},
127 }},
128 }); err != nil {
129 t.Fatal(err)
130 }
131}
132
133func newResumeBuilder(t *testing.T) func(providerName, model string) (*core.Agent, string, string, error) {
134 t.Helper()
135 return func(providerName, model string) (*core.Agent, string, string, error) {
136 return core.NewAgent(nil, model, "", nil), providerName, model, nil
137 }
138}
139
140// TestResumePlanSnapshotsUseFullTranscriptBeyondTrimWindow guards GO-001: at
141// resume only the trim window reaches the agent, so replaying its messages
142// turns an in-window delta into a rebuild from an empty base. The reconstructed
143// snapshots must instead come from the untrimmed session transcript.
144func TestResumePlanSnapshotsUseFullTranscriptBeyondTrimWindow(t *testing.T) {
145 sess, err := core.NewSession(t.TempDir(), t.TempDir(), "stored-provider", "stored-model", "test")
146 if err != nil {
147 t.Fatal(err)
148 }
149 appendPlanCall(t, sess, "set-1", "update_plan",
150 `{"plan":[{"step":"one","status":"completed"},{"step":"two","status":"in_progress"},{"step":"three","status":"pending"}]}`)
151 for i := 0; i < 120; i++ {
152 if err := sess.AppendMessage(provider.Message{
153 Role: provider.RoleUser,
154 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("filler-%d", i)}},
155 }); err != nil {
156 t.Fatal(err)
157 }
158 }
159 appendPlanCall(t, sess, "add-four", "plan", `{"action":"add","steps":[{"step":"four"}]}`)
160 appendPlanCall(t, sess, "upd-two", "plan", `{"action":"update","index":2,"status":"completed"}`)
161 final := []core.PlanStep{
162 {Step: "one", Status: core.PlanCompleted},
163 {Step: "two", Status: core.PlanCompleted},
164 {Step: "three", Status: core.PlanPending},
165 {Step: "four", Status: core.PlanPending},
166 }
167 if err := sess.UpdatePlan(&core.SessionPlan{Steps: final}); err != nil {
168 t.Fatal(err)
169 }
170 path := sess.Path
171 if err := sess.Close(); err != nil {
172 t.Fatal(err)
173 }
174
175 resumed, _, err := core.OpenSession(path)
176 if err != nil {
177 t.Fatal(err)
178 }
179 defer resumed.Close()
180 current := core.NewAgent(nil, "current-model", "", nil)
181 candidate, err := applySessionResume(resumed, current, "current-provider", "current-model", false, false, newResumeBuilder(t))
182 if err != nil {
183 t.Fatal(err)
184 }
185 defer candidate.session.Close()
186
187 if got := candidate.agent.CurrentPlan(); !slices.Equal(got, final) {
188 t.Fatalf("resumed plan = %#v, want %#v", got, final)
189 }
190 if len(candidate.planSnapshots) != 3 {
191 t.Fatalf("snapshots = %#v, want the full transcript's replays", candidate.planSnapshots)
192 }
193 if got := candidate.planSnapshots["set-1"].Plan; !slices.Equal(got, []core.PlanStep{
194 {Step: "one", Status: core.PlanCompleted},
195 {Step: "two", Status: core.PlanInProgress},
196 {Step: "three", Status: core.PlanPending},
197 }) {
198 t.Fatalf("set snapshot = %#v", got)
199 }
200 if got := candidate.planSnapshots["add-four"].Plan; !slices.Equal(got, []core.PlanStep{
201 {Step: "one", Status: core.PlanCompleted},
202 {Step: "two", Status: core.PlanInProgress},
203 {Step: "three", Status: core.PlanPending},
204 {Step: "four", Status: core.PlanPending},
205 }) {
206 t.Fatalf("add snapshot = %#v", got)
207 }
208 if got := candidate.planSnapshots["upd-two"].Plan; !slices.Equal(got, final) {
209 t.Fatalf("update snapshot = %#v, want %#v", got, final)
210 }
211}
212
213// TestResumeSeedsInteractivePlanCountersAndSnapshots covers the CLI resume
214// wiring the session picker shares: the candidate agent is seeded from
215// SessionMeta.Plan, and the Interactive reconstructs its counters and per-call
216// snapshot map from the resume result.
217func TestResumeSeedsInteractivePlanCountersAndSnapshots(t *testing.T) {
218 want := []core.PlanStep{
219 {Step: "one", Status: core.PlanCompleted},
220 {Step: "two", Status: core.PlanInProgress},
221 }
222 sess, err := core.NewSession(t.TempDir(), t.TempDir(), "stored-provider", "stored-model", "test")
223 if err != nil {
224 t.Fatal(err)
225 }
226 appendPlanCall(t, sess, "set-1", "plan", `{"action":"set","steps":[{"step":"one","status":"completed"},{"step":"two","status":"in_progress"}]}`)
227 if err := sess.UpdatePlan(&core.SessionPlan{Steps: want}); err != nil {
228 t.Fatal(err)
229 }
230 path := sess.Path
231 if err := sess.Close(); err != nil {
232 t.Fatal(err)
233 }
234
235 resumed, _, err := core.OpenSession(path)
236 if err != nil {
237 t.Fatal(err)
238 }
239 defer resumed.Close()
240 current := core.NewAgent(nil, "current-model", "", nil)
241 candidate, err := applySessionResume(resumed, current, "current-provider", "current-model", false, false, newResumeBuilder(t))
242 if err != nil {
243 t.Fatal(err)
244 }
245 defer candidate.session.Close()
246
247 iv := modes.NewInteractive(modes.InteractiveConfig{
248 Agent: candidate.agent,
249 CurrentPlan: candidate.agent.CurrentPlan,
250 InitialPlanSnapshots: candidate.planSnapshots,
251 })
252 if got := candidate.agent.CurrentPlan(); !slices.Equal(got, want) {
253 t.Fatalf("resumed plan = %#v, want %#v", got, want)
254 }
255 if currentStep, total, snapshots := iv.PlanStatus(); currentStep != 2 || total != 2 || snapshots != 1 {
256 t.Fatalf("interactive plan status = %d/%d with %d snapshots, want 2/2 with 1", currentStep, total, snapshots)
257 }
258}