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 TestCompactWithEventsAccumulatesAndPersistsUsage(t *testing.T) {
57 client := &compactLifecycleClient{usage: []provider.Usage{{InputTokens: 2}, {OutputTokens: 3, CostUSD: 0.25}}}
58 agent := NewAgent(client, "compact-model", "system", Registry{})
59 agent.SetMessages([]provider.Message{{
60 Role: provider.RoleUser,
61 Content: []provider.Content{provider.TextBlock{Text: "source"}},
62 }})
63 root := t.TempDir()
64 sess, err := NewSession(root, t.TempDir(), "provider", "compact-model", "test")
65 if err != nil {
66 t.Fatal(err)
67 }
68 t.Cleanup(func() { _ = sess.Close() })
69 agent.OnUsage = func(cumulative provider.Usage) {
70 if err := sess.AppendUsage(cumulative, cumulative); err != nil {
71 t.Fatalf("persist usage: %v", err)
72 }
73 }
74 var events []AgentEvent
75 if _, err := agent.CompactWithEvents(context.Background(), 0, func(event AgentEvent) {
76 events = append(events, event)
77 }); err != nil {
78 t.Fatalf("CompactWithEvents: %v", err)
79 }
80 if got, want := agent.Cost(), (provider.Usage{InputTokens: 2, OutputTokens: 3, CostUSD: 0.25}); got != want {
81 t.Fatalf("agent cost = %+v, want %+v", got, want)
82 }
83 cum, _, err := SessionUsageDetail(sess.Path)
84 if err != nil {
85 t.Fatal(err)
86 }
87 if cum != agent.Cost() {
88 t.Fatalf("persisted cumulative usage = %+v, want %+v", cum, agent.Cost())
89 }
90 var cumulative []provider.Usage
91 for _, event := range events {
92 if usage, ok := event.(EvUsage); ok {
93 cumulative = append(cumulative, usage.Cumulative)
94 }
95 }
96 if len(cumulative) != 2 || cumulative[0].InputTokens != 2 || cumulative[1] != agent.Cost() {
97 t.Fatalf("usage event cumulative values = %+v, want first input=2 and final=%+v", cumulative, agent.Cost())
98 }
99}
100
101func TestCompactWithEventsEmitsRequestLifecycle(t *testing.T) {
102 client := &compactLifecycleClient{}
103 agent := NewAgent(client, "compact-model", "system", Registry{})
104 agent.SetMessages([]provider.Message{{
105 Role: provider.RoleUser,
106 Content: []provider.Content{provider.TextBlock{Text: "source"}},
107 }})
108
109 var events []AgentEvent
110 summary, err := agent.CompactWithEvents(context.Background(), 0, func(event AgentEvent) {
111 events = append(events, event)
112 })
113 if err != nil {
114 t.Fatalf("CompactWithEvents: %v", err)
115 }
116 if summary != "summary" {
117 t.Fatalf("summary = %q, want summary", summary)
118 }
119 if !client.hadLifecycle {
120 t.Fatal("compaction request did not receive a lifecycle observer")
121 }
122
123 if len(events) != 4 {
124 t.Fatalf("events = %#v, want agent request, provider request, assistant start, text delta", events)
125 }
126 agentRequest, ok := events[0].(EvRequestStarted)
127 if !ok || agentRequest.Scope != RetryScopeAgent || agentRequest.Attempt != 1 || agentRequest.MaxAttempts != 1 {
128 t.Fatalf("first event = %#v, want initial agent request", events[0])
129 }
130 providerRequest, ok := events[1].(EvRequestStarted)
131 if !ok || providerRequest.Scope != RetryScopeProvider || providerRequest.Attempt != 1 || providerRequest.MaxAttempts != 1 {
132 t.Fatalf("second event = %#v, want initial provider request", events[1])
133 }
134 if _, ok := events[2].(EvAssistantStart); !ok {
135 t.Fatalf("third event = %#v, want assistant start", events[2])
136 }
137 if delta, ok := events[3].(EvTextDelta); !ok || delta.Delta != "summary" {
138 t.Fatalf("fourth event = %#v, want summary text delta", events[3])
139 }
140}
141
142func TestCompactPersistsProviderRetryLifecycleWithoutEventSink(t *testing.T) {
143 client := &compactLifecycleClient{reportRetry: true}
144 agent := NewAgent(client, "compact-model", "system", Registry{})
145 agent.SetMessages([]provider.Message{textMessage(provider.RoleUser, "source")})
146
147 var records []RetryLifecycleRecord
148 agent.OnRetryLifecycle = func(record RetryLifecycleRecord) {
149 records = append(records, record)
150 }
151 if _, err := agent.Compact(context.Background(), 0, nil); err != nil {
152 t.Fatalf("Compact: %v", err)
153 }
154 if !client.hadLifecycle {
155 t.Fatal("event-less compaction request did not receive a lifecycle observer")
156 }
157 want := []RetryLifecycleRecord{
158 {
159 Event: RetryLifecycleRequestFailed, Scope: RetryScopeProvider,
160 Attempt: 1, MaxAttempts: 2, Reason: RetryReasonOverload,
161 },
162 {
163 Event: RetryLifecycleRetryScheduled, Scope: RetryScopeProvider,
164 Attempt: 2, MaxAttempts: 2, Reason: RetryReasonOverload, DelayMS: 250,
165 },
166 }
167 if !reflect.DeepEqual(records, want) {
168 t.Fatalf("retry lifecycle = %#v, want %#v", records, want)
169 }
170}
171
172func TestCompactPromptPreservesActiveInstructions(t *testing.T) {
173 const activeInstruction = "delegate independent work to subagents"
174 client := &compactLifecycleClient{summary: "Active instruction still in force: " + activeInstruction}
175 agent := NewAgent(client, "compact-model", "system", Registry{})
176 agent.SetMessages([]provider.Message{
177 {
178 Role: provider.RoleUser,
179 Content: []provider.Content{provider.TextBlock{Text: "When coding, delegate independent work to subagents."}},
180 },
181 {
182 Role: provider.RoleAssistant,
183 Content: []provider.Content{provider.TextBlock{Text: "I will do that."}},
184 },
185 })
186
187 if _, err := agent.Compact(context.Background(), 0, nil); err != nil {
188 t.Fatalf("Compact: %v", err)
189 }
190 messages := agent.Messages()
191 if len(messages) != 1 {
192 t.Fatalf("compacted messages = %d, want 1", len(messages))
193 }
194 if messages[0].Role != provider.RoleUser {
195 t.Fatalf("compacted message role = %q, want user", messages[0].Role)
196 }
197 if messages[0].Meta["compaction"] != "true" {
198 t.Fatalf("compacted message meta = %#v, want compaction=true", messages[0].Meta)
199 }
200 compacted, ok := messages[0].Content[0].(provider.TextBlock)
201 if !ok {
202 t.Fatalf("compacted message content = %#v, want TextBlock", messages[0].Content[0])
203 }
204 if !strings.Contains(compacted.Text, "## Context Summary (compacted)") || !strings.Contains(compacted.Text, activeInstruction) {
205 t.Fatalf("compacted transcript did not preserve active instruction:\n%s", compacted.Text)
206 }
207 if !strings.Contains(client.req.System, "Preserve active user instructions") {
208 t.Fatalf("compaction system prompt does not preserve active instructions: %q", client.req.System)
209 }
210 if len(client.req.Messages) != 1 {
211 t.Fatalf("compaction request messages = %d, want 1", len(client.req.Messages))
212 }
213 text, ok := client.req.Messages[0].Content[0].(provider.TextBlock)
214 if !ok {
215 t.Fatalf("compaction request content = %#v, want TextBlock", client.req.Messages[0].Content[0])
216 }
217 for _, want := range []string{"## Active Instructions & Preferences", "delegation/subagent", activeInstruction} {
218 if !strings.Contains(text.Text, want) {
219 t.Fatalf("compaction prompt missing %q:\n%s", want, text.Text)
220 }
221 }
222}