1package core
2
3import (
4 "context"
5 "reflect"
6 "strings"
7 "testing"
8 "time"
9
10 "github.com/bnema/zut/packages/provider"
11)
12
13type compactLifecycleClient struct {
14 hadLifecycle bool
15 reportRetry bool
16 req provider.Request
17 summary string
18 usage []provider.Usage
19}
20
21func (c *compactLifecycleClient) Name() string { return "compact-lifecycle" }
22
23func (c *compactLifecycleClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
24 c.req = req
25 c.hadLifecycle = req.Lifecycle != nil
26 if req.Lifecycle != nil {
27 maxAttempts := 1
28 if c.reportRetry {
29 maxAttempts = 2
30 }
31 req.Lifecycle.RequestAttempt(1, maxAttempts)
32 if c.reportRetry {
33 if failures, ok := req.Lifecycle.(provider.RequestFailureLifecycle); ok {
34 failures.RequestFailed(1, 2, provider.RequestFailureOverload, false)
35 }
36 req.Lifecycle.RetryScheduled(2, 2, 250*time.Millisecond)
37 }
38 }
39 summary := c.summary
40 if summary == "" {
41 summary = "summary"
42 }
43 out := make(chan provider.Event, 2+len(c.usage))
44 out <- provider.EventTextDelta{Delta: summary}
45 for _, usage := range c.usage {
46 out <- provider.EventUsage{Usage: usage}
47 }
48 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
49 Role: provider.RoleAssistant,
50 Content: []provider.Content{provider.TextBlock{Text: summary}},
51 }}
52 close(out)
53 return out, nil
54}
55
56func TestCompactPreservesOnlyLatestInternalContext(t *testing.T) {
57 client := &compactLifecycleClient{}
58 agent := NewAgent(client, "compact-model", "system", Registry{})
59 agent.SetMessages([]provider.Message{
60 {Role: provider.RoleDeveloper, Content: []provider.Content{provider.TextBlock{Text: "old"}}, Meta: map[string]string{internalContextMarker: "true"}},
61 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "first"}}},
62 {Role: provider.RoleDeveloper, Content: []provider.Content{provider.TextBlock{Text: "new"}}, Meta: map[string]string{internalContextMarker: "true"}},
63 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "answer"}}},
64 })
65 if _, err := agent.Compact(context.Background(), 0, nil); err != nil {
66 t.Fatal(err)
67 }
68 messages := agent.Messages()
69 if len(messages) != 2 || messages[0].Role != provider.RoleDeveloper || internalContextText(messages[0]) != "new" {
70 t.Fatalf("compacted messages = %#v, want latest internal context then summary", messages)
71 }
72}
73
74func TestCompactStartsProviderRequestIdentity(t *testing.T) {
75 client := &compactLifecycleClient{}
76 agent := NewAgent(client, "compact-model", "system", Registry{})
77 agent.SetMessages([]provider.Message{{
78 Role: provider.RoleUser,
79 Content: []provider.Content{provider.TextBlock{Text: "source"}},
80 }})
81 if _, err := agent.Compact(context.Background(), 0, nil); err != nil {
82 t.Fatal(err)
83 }
84 if client.req.Context.CacheSessionID == "" || client.req.Context.ThreadID == "" || client.req.Context.TurnID == "" {
85 t.Fatalf("compaction request context = %#v, want all identities", client.req.Context)
86 }
87 if err := agent.BindRequestIdentity("other-cache", "other-thread"); err == nil {
88 t.Fatal("BindRequestIdentity succeeded after compaction provider request")
89 }
90}
91
92func TestCompactWithEventsAccumulatesAndPersistsUsage(t *testing.T) {
93 client := &compactLifecycleClient{usage: []provider.Usage{{InputTokens: 2}, {OutputTokens: 3, CostUSD: 0.25}}}
94 agent := NewAgent(client, "compact-model", "system", Registry{})
95 agent.SetMessages([]provider.Message{{
96 Role: provider.RoleUser,
97 Content: []provider.Content{provider.TextBlock{Text: "source"}},
98 }})
99 root := t.TempDir()
100 sess, err := NewSession(root, t.TempDir(), "provider", "compact-model", "test")
101 if err != nil {
102 t.Fatal(err)
103 }
104 t.Cleanup(func() { _ = sess.Close() })
105 agent.OnUsage = func(cumulative provider.Usage) {
106 if err := sess.AppendUsage(cumulative, cumulative); err != nil {
107 t.Fatalf("persist usage: %v", err)
108 }
109 }
110 var events []AgentEvent
111 if _, err := agent.CompactWithEvents(context.Background(), 0, func(event AgentEvent) {
112 events = append(events, event)
113 }); err != nil {
114 t.Fatalf("CompactWithEvents: %v", err)
115 }
116 if got, want := agent.Cost(), (provider.Usage{InputTokens: 2, OutputTokens: 3, CostUSD: 0.25}); got != want {
117 t.Fatalf("agent cost = %+v, want %+v", got, want)
118 }
119 cum, _, err := SessionUsageDetail(sess.Path)
120 if err != nil {
121 t.Fatal(err)
122 }
123 if cum != agent.Cost() {
124 t.Fatalf("persisted cumulative usage = %+v, want %+v", cum, agent.Cost())
125 }
126 var cumulative []provider.Usage
127 for _, event := range events {
128 if usage, ok := event.(EvUsage); ok {
129 cumulative = append(cumulative, usage.Cumulative)
130 }
131 }
132 if len(cumulative) != 2 || cumulative[0].InputTokens != 2 || cumulative[1] != agent.Cost() {
133 t.Fatalf("usage event cumulative values = %+v, want first input=2 and final=%+v", cumulative, agent.Cost())
134 }
135}
136
137func TestCompactWithEventsEmitsRequestLifecycle(t *testing.T) {
138 client := &compactLifecycleClient{}
139 agent := NewAgent(client, "compact-model", "system", Registry{})
140 agent.SetMessages([]provider.Message{{
141 Role: provider.RoleUser,
142 Content: []provider.Content{provider.TextBlock{Text: "source"}},
143 }})
144
145 var events []AgentEvent
146 summary, err := agent.CompactWithEvents(context.Background(), 0, func(event AgentEvent) {
147 events = append(events, event)
148 })
149 if err != nil {
150 t.Fatalf("CompactWithEvents: %v", err)
151 }
152 if summary != "summary" {
153 t.Fatalf("summary = %q, want summary", summary)
154 }
155 if !client.hadLifecycle {
156 t.Fatal("compaction request did not receive a lifecycle observer")
157 }
158
159 if len(events) != 4 {
160 t.Fatalf("events = %#v, want agent request, provider request, assistant start, text delta", events)
161 }
162 agentRequest, ok := events[0].(EvRequestStarted)
163 if !ok || agentRequest.Scope != RetryScopeAgent || agentRequest.Attempt != 1 || agentRequest.MaxAttempts != 1 {
164 t.Fatalf("first event = %#v, want initial agent request", events[0])
165 }
166 providerRequest, ok := events[1].(EvRequestStarted)
167 if !ok || providerRequest.Scope != RetryScopeProvider || providerRequest.Attempt != 1 || providerRequest.MaxAttempts != 1 {
168 t.Fatalf("second event = %#v, want initial provider request", events[1])
169 }
170 if _, ok := events[2].(EvAssistantStart); !ok {
171 t.Fatalf("third event = %#v, want assistant start", events[2])
172 }
173 if delta, ok := events[3].(EvTextDelta); !ok || delta.Delta != "summary" {
174 t.Fatalf("fourth event = %#v, want summary text delta", events[3])
175 }
176}
177
178func TestCompactPersistsProviderRetryLifecycleWithoutEventSink(t *testing.T) {
179 client := &compactLifecycleClient{reportRetry: true}
180 agent := NewAgent(client, "compact-model", "system", Registry{})
181 agent.SetMessages([]provider.Message{textMessage(provider.RoleUser, "source")})
182
183 var records []RetryLifecycleRecord
184 agent.OnRetryLifecycle = func(record RetryLifecycleRecord) {
185 records = append(records, record)
186 }
187 if _, err := agent.Compact(context.Background(), 0, nil); err != nil {
188 t.Fatalf("Compact: %v", err)
189 }
190 if !client.hadLifecycle {
191 t.Fatal("event-less compaction request did not receive a lifecycle observer")
192 }
193 want := []RetryLifecycleRecord{
194 {
195 Event: RetryLifecycleRequestFailed, Scope: RetryScopeProvider,
196 Attempt: 1, MaxAttempts: 2, Reason: RetryReasonOverload,
197 },
198 {
199 Event: RetryLifecycleRetryScheduled, Scope: RetryScopeProvider,
200 Attempt: 2, MaxAttempts: 2, Reason: RetryReasonOverload, DelayMS: 250,
201 },
202 }
203 if !reflect.DeepEqual(records, want) {
204 t.Fatalf("retry lifecycle = %#v, want %#v", records, want)
205 }
206}
207
208func TestCompactPromptPreservesActiveInstructions(t *testing.T) {
209 const activeInstruction = "delegate independent work to subagents"
210 client := &compactLifecycleClient{summary: "Active instruction still in force: " + activeInstruction}
211 agent := NewAgent(client, "compact-model", "system", Registry{})
212 agent.SetMessages([]provider.Message{
213 {
214 Role: provider.RoleUser,
215 Content: []provider.Content{provider.TextBlock{Text: "When coding, delegate independent work to subagents."}},
216 },
217 {
218 Role: provider.RoleAssistant,
219 Content: []provider.Content{provider.TextBlock{Text: "I will do that."}},
220 },
221 })
222
223 if _, err := agent.Compact(context.Background(), 0, nil); err != nil {
224 t.Fatalf("Compact: %v", err)
225 }
226 messages := agent.Messages()
227 if len(messages) != 1 {
228 t.Fatalf("compacted messages = %d, want 1", len(messages))
229 }
230 if messages[0].Role != provider.RoleUser {
231 t.Fatalf("compacted message role = %q, want user", messages[0].Role)
232 }
233 if messages[0].Meta["compaction"] != "true" {
234 t.Fatalf("compacted message meta = %#v, want compaction=true", messages[0].Meta)
235 }
236 compacted, ok := messages[0].Content[0].(provider.TextBlock)
237 if !ok {
238 t.Fatalf("compacted message content = %#v, want TextBlock", messages[0].Content[0])
239 }
240 if !strings.Contains(compacted.Text, "## Context Summary (compacted)") || !strings.Contains(compacted.Text, activeInstruction) {
241 t.Fatalf("compacted transcript did not preserve active instruction:\n%s", compacted.Text)
242 }
243 if !strings.Contains(client.req.System, "Preserve active user instructions") {
244 t.Fatalf("compaction system prompt does not preserve active instructions: %q", client.req.System)
245 }
246 if len(client.req.Messages) != 1 {
247 t.Fatalf("compaction request messages = %d, want 1", len(client.req.Messages))
248 }
249 text, ok := client.req.Messages[0].Content[0].(provider.TextBlock)
250 if !ok {
251 t.Fatalf("compaction request content = %#v, want TextBlock", client.req.Messages[0].Content[0])
252 }
253 for _, want := range []string{"## Active Instructions & Preferences", "delegation/subagent", activeInstruction} {
254 if !strings.Contains(text.Text, want) {
255 t.Fatalf("compaction prompt missing %q:\n%s", want, text.Text)
256 }
257 }
258}