1package core
2
3import (
4 "context"
5 "strings"
6 "testing"
7
8 "github.com/bnema/zut/packages/provider"
9)
10
11type compactLifecycleClient struct {
12 hadLifecycle bool
13 req provider.Request
14 summary string
15 usage []provider.Usage
16}
17
18func (c *compactLifecycleClient) Name() string { return "compact-lifecycle" }
19
20func (c *compactLifecycleClient) Stream(_ context.Context, req provider.Request) (<-chan provider.Event, error) {
21 c.req = req
22 c.hadLifecycle = req.Lifecycle != nil
23 if req.Lifecycle != nil {
24 req.Lifecycle.RequestAttempt(1, 1)
25 }
26 summary := c.summary
27 if summary == "" {
28 summary = "summary"
29 }
30 out := make(chan provider.Event, 2+len(c.usage))
31 out <- provider.EventTextDelta{Delta: summary}
32 for _, usage := range c.usage {
33 out <- provider.EventUsage{Usage: usage}
34 }
35 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
36 Role: provider.RoleAssistant,
37 Content: []provider.Content{provider.TextBlock{Text: summary}},
38 }}
39 close(out)
40 return out, nil
41}
42
43func TestCompactWithEventsAccumulatesAndPersistsUsage(t *testing.T) {
44 client := &compactLifecycleClient{usage: []provider.Usage{{InputTokens: 2}, {OutputTokens: 3, CostUSD: 0.25}}}
45 agent := NewAgent(client, "compact-model", "system", Registry{})
46 agent.SetMessages([]provider.Message{{
47 Role: provider.RoleUser,
48 Content: []provider.Content{provider.TextBlock{Text: "source"}},
49 }})
50 root := t.TempDir()
51 sess, err := NewSession(root, t.TempDir(), "provider", "compact-model", "test")
52 if err != nil {
53 t.Fatal(err)
54 }
55 t.Cleanup(func() { _ = sess.Close() })
56 agent.OnUsage = func(cumulative provider.Usage) {
57 if err := sess.AppendUsage(cumulative, cumulative); err != nil {
58 t.Fatalf("persist usage: %v", err)
59 }
60 }
61 var events []AgentEvent
62 if _, err := agent.CompactWithEvents(context.Background(), 0, func(event AgentEvent) {
63 events = append(events, event)
64 }); err != nil {
65 t.Fatalf("CompactWithEvents: %v", err)
66 }
67 if got, want := agent.Cost(), (provider.Usage{InputTokens: 2, OutputTokens: 3, CostUSD: 0.25}); got != want {
68 t.Fatalf("agent cost = %+v, want %+v", got, want)
69 }
70 cum, _, err := SessionUsageDetail(sess.Path)
71 if err != nil {
72 t.Fatal(err)
73 }
74 if cum != agent.Cost() {
75 t.Fatalf("persisted cumulative usage = %+v, want %+v", cum, agent.Cost())
76 }
77 var cumulative []provider.Usage
78 for _, event := range events {
79 if usage, ok := event.(EvUsage); ok {
80 cumulative = append(cumulative, usage.Cumulative)
81 }
82 }
83 if len(cumulative) != 2 || cumulative[0].InputTokens != 2 || cumulative[1] != agent.Cost() {
84 t.Fatalf("usage event cumulative values = %+v, want first input=2 and final=%+v", cumulative, agent.Cost())
85 }
86}
87
88func TestCompactWithEventsEmitsRequestLifecycle(t *testing.T) {
89 client := &compactLifecycleClient{}
90 agent := NewAgent(client, "compact-model", "system", Registry{})
91 agent.SetMessages([]provider.Message{{
92 Role: provider.RoleUser,
93 Content: []provider.Content{provider.TextBlock{Text: "source"}},
94 }})
95
96 var events []AgentEvent
97 summary, err := agent.CompactWithEvents(context.Background(), 0, func(event AgentEvent) {
98 events = append(events, event)
99 })
100 if err != nil {
101 t.Fatalf("CompactWithEvents: %v", err)
102 }
103 if summary != "summary" {
104 t.Fatalf("summary = %q, want summary", summary)
105 }
106 if !client.hadLifecycle {
107 t.Fatal("compaction request did not receive a lifecycle observer")
108 }
109
110 if len(events) != 4 {
111 t.Fatalf("events = %#v, want agent request, provider request, assistant start, text delta", events)
112 }
113 agentRequest, ok := events[0].(EvRequestStarted)
114 if !ok || agentRequest.Scope != RetryScopeAgent || agentRequest.Attempt != 1 || agentRequest.MaxAttempts != 1 {
115 t.Fatalf("first event = %#v, want initial agent request", events[0])
116 }
117 providerRequest, ok := events[1].(EvRequestStarted)
118 if !ok || providerRequest.Scope != RetryScopeProvider || providerRequest.Attempt != 1 || providerRequest.MaxAttempts != 1 {
119 t.Fatalf("second event = %#v, want initial provider request", events[1])
120 }
121 if _, ok := events[2].(EvAssistantStart); !ok {
122 t.Fatalf("third event = %#v, want assistant start", events[2])
123 }
124 if delta, ok := events[3].(EvTextDelta); !ok || delta.Delta != "summary" {
125 t.Fatalf("fourth event = %#v, want summary text delta", events[3])
126 }
127}
128
129func TestCompactPromptPreservesActiveInstructions(t *testing.T) {
130 const activeInstruction = "delegate independent work to subagents"
131 client := &compactLifecycleClient{summary: "Active instruction still in force: " + activeInstruction}
132 agent := NewAgent(client, "compact-model", "system", Registry{})
133 agent.SetMessages([]provider.Message{
134 {
135 Role: provider.RoleUser,
136 Content: []provider.Content{provider.TextBlock{Text: "When coding, delegate independent work to subagents."}},
137 },
138 {
139 Role: provider.RoleAssistant,
140 Content: []provider.Content{provider.TextBlock{Text: "I will do that."}},
141 },
142 })
143
144 if _, err := agent.Compact(context.Background(), 0, nil); err != nil {
145 t.Fatalf("Compact: %v", err)
146 }
147 messages := agent.Messages()
148 if len(messages) != 1 {
149 t.Fatalf("compacted messages = %d, want 1", len(messages))
150 }
151 if messages[0].Role != provider.RoleUser {
152 t.Fatalf("compacted message role = %q, want user", messages[0].Role)
153 }
154 if messages[0].Meta["compaction"] != "true" {
155 t.Fatalf("compacted message meta = %#v, want compaction=true", messages[0].Meta)
156 }
157 compacted, ok := messages[0].Content[0].(provider.TextBlock)
158 if !ok {
159 t.Fatalf("compacted message content = %#v, want TextBlock", messages[0].Content[0])
160 }
161 if !strings.Contains(compacted.Text, "## Context Summary (compacted)") || !strings.Contains(compacted.Text, activeInstruction) {
162 t.Fatalf("compacted transcript did not preserve active instruction:\n%s", compacted.Text)
163 }
164 if !strings.Contains(client.req.System, "Preserve active user instructions") {
165 t.Fatalf("compaction system prompt does not preserve active instructions: %q", client.req.System)
166 }
167 if len(client.req.Messages) != 1 {
168 t.Fatalf("compaction request messages = %d, want 1", len(client.req.Messages))
169 }
170 text, ok := client.req.Messages[0].Content[0].(provider.TextBlock)
171 if !ok {
172 t.Fatalf("compaction request content = %#v, want TextBlock", client.req.Messages[0].Content[0])
173 }
174 for _, want := range []string{"## Active Instructions & Preferences", "delegation/subagent", activeInstruction} {
175 if !strings.Contains(text.Text, want) {
176 t.Fatalf("compaction prompt missing %q:\n%s", want, text.Text)
177 }
178 }
179}