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