1package agent
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "os"
9 "sync"
10 "testing"
11 "time"
12
13 "github.com/bnema/zut/packages/agent/modes"
14 "github.com/bnema/zut/packages/core"
15 "github.com/bnema/zut/packages/provider"
16)
17
18func TestTrimMessagesForResumeCarriesDeferredToolActivation(t *testing.T) {
19 msgs := make([]provider.Message, 0, 101)
20 msgs = append(msgs, provider.Message{
21 Role: provider.RoleTool,
22 AddedToolNames: []string{"lookup_weather"},
23 Content: []provider.Content{provider.ToolResultBlock{CallID: "old-call"}},
24 })
25 for idx := 1; idx < 101; idx++ {
26 msgs = append(msgs, provider.Message{
27 Role: provider.RoleUser,
28 Content: []provider.Content{provider.TextBlock{Text: "message"}},
29 })
30 }
31 trimmed := trimMessagesForResume(msgs, 100)
32 if len(trimmed) != 100 {
33 t.Fatalf("trimmed message count = %d, want 100", len(trimmed))
34 }
35 if len(trimmed[0].AddedToolNames) != 1 || trimmed[0].AddedToolNames[0] != "lookup_weather" {
36 t.Fatalf("trimmed activation markers = %v, want lookup_weather", trimmed[0].AddedToolNames)
37 }
38}
39
40func TestTrimMessagesForResumeKeepsCompactionSummaryAtTailBoundary(t *testing.T) {
41 for _, total := range []int{100, 101, 102} {
42 t.Run(fmt.Sprintf("%d messages", total), func(t *testing.T) {
43 msgs := make([]provider.Message, 0, total)
44 msgs = append(msgs, provider.Message{
45 Role: provider.RoleUser,
46 Meta: map[string]string{"compaction": "true"},
47 Content: []provider.Content{
48 provider.TextBlock{Text: "## Context Summary (compacted)"},
49 },
50 })
51 for idx := 1; idx < total; idx++ {
52 msgs = append(msgs, provider.Message{
53 Role: provider.RoleUser,
54 Content: []provider.Content{provider.TextBlock{Text: fmt.Sprintf("message-%d", idx)}},
55 })
56 }
57
58 trimmed := trimMessagesForResume(msgs, 100)
59 if len(trimmed) > 100 {
60 t.Fatalf("trimmed message count = %d, want at most 100", len(trimmed))
61 }
62 if len(trimmed) == 0 || trimmed[0].Meta["compaction"] != "true" {
63 t.Fatalf("trimmed transcript lost compaction summary: %+v", trimmed)
64 }
65 })
66 }
67}
68
69func TestLiveInteractiveAgentUsesReplacementAgentForSessionResume(t *testing.T) {
70 startup := core.NewAgent(nil, "startup-model", "", nil)
71 startup.SetMessages([]provider.Message{{
72 Role: provider.RoleUser,
73 Content: []provider.Content{provider.TextBlock{Text: "startup transcript"}},
74 }})
75
76 replacement := core.NewAgent(nil, "replacement-model", "", nil)
77 iv := modes.NewInteractive(modes.InteractiveConfig{Agent: replacement})
78
79 resumed := []provider.Message{{
80 Role: provider.RoleUser,
81 Content: []provider.Content{provider.TextBlock{Text: "resumed transcript"}},
82 }}
83 liveInteractiveAgent(iv, startup).SetMessages(resumed)
84
85 if got := firstMessageText(replacement.Messages()); got != "resumed transcript" {
86 t.Fatalf("replacement agent transcript = %q, want resumed transcript", got)
87 }
88 if got := firstMessageText(startup.Messages()); got != "startup transcript" {
89 t.Fatalf("startup agent transcript changed to %q", got)
90 }
91}
92
93func TestLiveInteractiveAgentFallsBackBeforeInteractiveConstruction(t *testing.T) {
94 startup := core.NewAgent(nil, "startup-model", "", nil)
95 if got := liveInteractiveAgent(nil, startup); got != startup {
96 t.Fatalf("fallback agent = %p, want %p", got, startup)
97 }
98}
99
100func TestPersistModelCallbackDoesNotReenterSessionTransition(t *testing.T) {
101 t.Setenv("ZUT_HOME", t.TempDir())
102 sess, err := core.NewSession(t.TempDir(), t.TempDir(), "old-provider", "old-model", "test")
103 if err != nil {
104 t.Fatal(err)
105 }
106 defer sess.Close()
107
108 var persistMu sync.Mutex
109 activeProvider, activeModel := "old-provider", "old-model"
110 persistModel := newPersistModelCallback(&persistMu, &sess, &activeProvider, &activeModel, nil)
111
112 var transitionMu sync.RWMutex
113 sessionTransition := newSessionTransition(&transitionMu)
114 done := make(chan struct{})
115 go func() {
116 sessionTransition(func() {
117 persistModel("new-provider", "new-model")
118 })
119 close(done)
120 }()
121 select {
122 case <-done:
123 case <-time.After(time.Second):
124 t.Fatal("model persistence deadlocked inside the session transition")
125 }
126 if sess.Meta.Provider != "new-provider" || sess.Meta.Model != "new-model" {
127 t.Fatalf("session model = %q/%q, want new-provider/new-model", sess.Meta.Provider, sess.Meta.Model)
128 }
129 if activeProvider != "new-provider" || activeModel != "new-model" {
130 t.Fatalf("active model = %q/%q, want new-provider/new-model", activeProvider, activeModel)
131 }
132}
133
134func firstMessageText(messages []provider.Message) string {
135 if len(messages) == 0 || len(messages[0].Content) == 0 {
136 return ""
137 }
138 text, _ := messages[0].Content[0].(provider.TextBlock)
139 return text.Text
140}
141
142func syntheticSession(t *testing.T, providerName, model string, usage provider.Usage) string {
143 t.Helper()
144 root := t.TempDir()
145 cwd := t.TempDir()
146 sess, err := core.NewSession(root, cwd, providerName, model, "test")
147 if err != nil {
148 t.Fatal(err)
149 }
150 if err := sess.AppendMessage(provider.Message{
151 Role: provider.RoleUser,
152 Content: []provider.Content{provider.TextBlock{Text: "resumed transcript"}},
153 }); err != nil {
154 t.Fatal(err)
155 }
156 if usage != (provider.Usage{}) {
157 if err := sess.AppendUsage(usage, usage); err != nil {
158 t.Fatal(err)
159 }
160 }
161 if err := sess.Close(); err != nil {
162 t.Fatal(err)
163 }
164 return sess.Path
165}
166
167func TestApplyInitialSessionResumeKeepsFreshEmptySessionOwned(t *testing.T) {
168 root := t.TempDir()
169 sess, err := core.NewSession(root, "/workspace", "provider", "model", "test")
170 if err != nil {
171 t.Fatal(err)
172 }
173 ag := core.NewAgent(nil, "model", "", nil)
174 gotSess, gotAgent, providerName, model, err := applyInitialSessionResume(context.Background(), Args{}, Resolved{Provider: "provider", Model: "model"}, nil, sess, ag)
175 if err != nil {
176 t.Fatal(err)
177 }
178 if gotSess != sess || gotAgent != ag || providerName != "provider" || model != "model" {
179 t.Fatalf("empty startup resume = session %p/%p agent %p/%p provider/model %q/%q", gotSess, sess, gotAgent, ag, providerName, model)
180 }
181 if _, err := os.Stat(sess.Path); err != nil {
182 t.Fatalf("fresh session disappeared before close: %v", err)
183 }
184 if err := sess.Close(); err != nil {
185 t.Fatal(err)
186 }
187 if _, err := os.Stat(sess.Path); !errors.Is(err, os.ErrNotExist) {
188 t.Fatalf("fresh empty session remained after close: %v", err)
189 }
190}
191
192func TestPrepareSessionResumeRebuildsForStoredProviderAndModel(t *testing.T) {
193 old := core.NewAgent(nil, "old-model", "", nil)
194 old.SetMessages([]provider.Message{{
195 Role: provider.RoleUser,
196 Content: []provider.Content{provider.TextBlock{Text: "current transcript"}},
197 }})
198 path := syntheticSession(t, "stored-provider", "stored-model", provider.Usage{
199 InputTokens: 12,
200 OutputTokens: 7,
201 })
202
203 var gotProvider, gotModel string
204 candidate, err := prepareSessionResume(path, old, "current-provider", "old-model", func(providerName, model string) (*core.Agent, string, string, error) {
205 gotProvider, gotModel = providerName, model
206 return core.NewAgent(nil, model, "", nil), providerName, model, nil
207 })
208 if err != nil {
209 t.Fatal(err)
210 }
211 defer candidate.session.Close()
212
213 if gotProvider != "stored-provider" || gotModel != "stored-model" {
214 t.Fatalf("builder selection = %q/%q, want stored-provider/stored-model", gotProvider, gotModel)
215 }
216 if !candidate.rebuilt || candidate.agent == old {
217 t.Fatalf("candidate rebuilt = %v, agent=%p, old=%p", candidate.rebuilt, candidate.agent, old)
218 }
219 if got := firstMessageText(candidate.agent.Messages()); got != "resumed transcript" {
220 t.Fatalf("candidate transcript = %q, want resumed transcript", got)
221 }
222 if got := candidate.agent.Cost(); got.InputTokens != 12 || got.OutputTokens != 7 {
223 t.Fatalf("candidate usage = %+v, want input=12 output=7", got)
224 }
225 if got := old.Messages(); len(got) != 1 || firstMessageText(got) != "current transcript" {
226 t.Fatalf("current transcript changed during preparation: %+v", got)
227 }
228}
229
230func TestPrepareSessionResumePreservesLegacyMissingMetadata(t *testing.T) {
231 old := core.NewAgent(nil, "old-model", "", nil)
232 old.SetMessages([]provider.Message{{
233 Role: provider.RoleUser,
234 Content: []provider.Content{provider.TextBlock{Text: "current transcript"}},
235 }})
236 path := syntheticSession(t, "", "", provider.Usage{})
237
238 candidate, err := prepareSessionResume(path, old, "current-provider", "old-model", func(string, string) (*core.Agent, string, string, error) {
239 t.Fatal("legacy session unexpectedly requested a rebuild")
240 return nil, "", "", nil
241 })
242 if err != nil {
243 t.Fatal(err)
244 }
245 defer candidate.session.Close()
246
247 if candidate.rebuilt || candidate.agent != old {
248 t.Fatalf("legacy candidate = rebuilt %v, agent=%p; want current agent without rebuild", candidate.rebuilt, candidate.agent)
249 }
250 if got := firstMessageText(old.Messages()); got != "current transcript" {
251 t.Fatalf("current transcript changed before commit: %q", got)
252 }
253 if got := firstMessageText(candidate.messages); got != "resumed transcript" {
254 t.Fatalf("legacy candidate transcript = %q, want resumed transcript", got)
255 }
256}
257
258func TestPrepareSessionResumePreservesCompactHandoffMetadata(t *testing.T) {
259 path := syntheticSession(t, "stored-provider", "stored-model", provider.Usage{})
260 session, _, err := core.OpenSession(path)
261 if err != nil {
262 t.Fatal(err)
263 }
264 handoff := json.RawMessage(`{"version":1,"reason":"status_rescue","rescue_attempts":1}`)
265 if err := session.UpdateCompactHandoff(handoff); err != nil {
266 t.Fatal(err)
267 }
268 if err := session.Close(); err != nil {
269 t.Fatal(err)
270 }
271
272 current := core.NewAgent(nil, "current-model", "", nil)
273 candidate, err := prepareSessionResume(path, current, "stored-provider", "stored-model", nil)
274 if err != nil {
275 t.Fatal(err)
276 }
277 defer candidate.session.Close()
278 if got := string(candidate.session.Meta.CompactHandoff); got != string(handoff) {
279 t.Fatalf("candidate compact handoff = %q, want %q", got, handoff)
280 }
281}
282
283func TestPrepareSessionResumeFailureLeavesCurrentAndCandidateFileUsable(t *testing.T) {
284 old := core.NewAgent(nil, "old-model", "", nil)
285 old.SetMessages([]provider.Message{{
286 Role: provider.RoleUser,
287 Content: []provider.Content{provider.TextBlock{Text: "current transcript"}},
288 }})
289 path := syntheticSession(t, "stored-provider", "stored-model", provider.Usage{})
290 wantErr := errors.New("synthetic builder failure")
291
292 if _, err := prepareSessionResume(path, old, "current-provider", "old-model", func(string, string) (*core.Agent, string, string, error) {
293 return nil, "", "", wantErr
294 }); !errors.Is(err, wantErr) {
295 t.Fatalf("prepare error = %v, want %v", err, wantErr)
296 }
297 if got := firstMessageText(old.Messages()); got != "current transcript" {
298 t.Fatalf("current transcript changed after failed preparation: %q", got)
299 }
300
301 // The failed candidate must have released its append handle, leaving
302 // the selected file readable and reopenable for a later retry.
303 reopened, _, err := core.OpenSession(path)
304 if err != nil {
305 t.Fatalf("reopen candidate after failed preparation: %v", err)
306 }
307 if err := reopened.Close(); err != nil {
308 t.Fatal(err)
309 }
310 if _, err := os.Stat(path); err != nil {
311 t.Fatalf("candidate session disappeared after failed preparation: %v", err)
312 }
313}
314
315func TestPrepareSessionResumeHonorsExplicitProviderModelFields(t *testing.T) {
316 for _, tc := range []struct {
317 name string
318 explicitProvider bool
319 explicitModel bool
320 wantProvider string
321 wantModel string
322 wantBuild bool
323 }{
324 {name: "both", explicitProvider: true, explicitModel: true, wantProvider: "current-provider", wantModel: "current-model"},
325 {name: "provider-only", explicitProvider: true, wantProvider: "current-provider", wantModel: "stored-model", wantBuild: true},
326 {name: "model-only", explicitModel: true, wantProvider: "stored-provider", wantModel: "current-model", wantBuild: true},
327 } {
328 t.Run(tc.name, func(t *testing.T) {
329 old := core.NewAgent(nil, "old-model", "", nil)
330 path := syntheticSession(t, "stored-provider", "stored-model", provider.Usage{})
331 var gotProvider, gotModel string
332 candidate, err := prepareSessionResumeWithOptions(path, old, "current-provider", "current-model", tc.explicitProvider, tc.explicitModel, func(providerName, model string) (*core.Agent, string, string, error) {
333 gotProvider, gotModel = providerName, model
334 return core.NewAgent(nil, model, "", nil), providerName, model, nil
335 })
336 if err != nil {
337 t.Fatal(err)
338 }
339 defer candidate.session.Close()
340 if tc.wantBuild {
341 if gotProvider != tc.wantProvider || gotModel != tc.wantModel {
342 t.Fatalf("builder selection = %q/%q, want %q/%q", gotProvider, gotModel, tc.wantProvider, tc.wantModel)
343 }
344 } else if candidate.provider != tc.wantProvider || candidate.model != tc.wantModel || candidate.rebuilt {
345 t.Fatalf("candidate selection = %q/%q rebuilt=%v, want %q/%q without rebuild", candidate.provider, candidate.model, candidate.rebuilt, tc.wantProvider, tc.wantModel)
346 }
347 })
348 }
349}