1package core
2
3import (
4 "context"
5 "encoding/json"
6 "sync"
7 "sync/atomic"
8 "testing"
9 "time"
10
11 "github.com/bnema/zut/packages/provider"
12)
13
14type queueFakeClient struct {
15 calls int32
16}
17
18func (c *queueFakeClient) Name() string { return "queue-fake" }
19
20func (c *queueFakeClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
21 call := atomic.AddInt32(&c.calls, 1)
22 out := make(chan provider.Event, 4)
23 go func() {
24 defer close(out)
25 out <- provider.EventStart{Provider: "queue-fake", Model: req.Model}
26 switch call {
27 case 1:
28 out <- provider.EventToolStart{ID: "t1", Name: "echo"}
29 out <- provider.EventToolEnd{ID: "t1"}
30 out <- provider.EventDone{Stop: provider.StopToolUse, Message: provider.Message{
31 Role: provider.RoleAssistant,
32 Content: []provider.Content{
33 provider.TextBlock{Text: "using tool"},
34 provider.ToolCallBlock{ID: "t1", Name: "echo", Arguments: json.RawMessage(`{}`)},
35 },
36 }}
37 case 2:
38 out <- provider.EventTextDelta{Delta: "saw queued"}
39 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
40 Role: provider.RoleAssistant,
41 Content: []provider.Content{provider.TextBlock{Text: "saw queued"}},
42 }}
43 default:
44 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
45 Role: provider.RoleAssistant,
46 Content: []provider.Content{provider.TextBlock{Text: "extra"}},
47 }}
48 }
49 }()
50 return out, nil
51}
52
53// blockingTool waits until the test has queued a message, then
54// returns. This pins the core behaviour: queued user text is delivered
55// after the current tool batch finishes and before the next model call.
56type blockingTool struct {
57 started chan struct{}
58 release chan struct{}
59}
60
61func (t *blockingTool) Name() string { return "echo" }
62func (t *blockingTool) Description() string { return "echoes" }
63func (t *blockingTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) }
64
65func (t *blockingTool) Execute(ctx context.Context, args json.RawMessage, progress func(string)) (ToolResult, error) {
66 close(t.started)
67 select {
68 case <-ctx.Done():
69 return ToolResult{Content: []provider.Content{provider.TextBlock{Text: ctx.Err().Error()}}, IsError: true}, ctx.Err()
70 case <-t.release:
71 }
72 return ToolResult{Content: []provider.Content{provider.TextBlock{Text: "tool ok"}}}, nil
73}
74
75func TestQueuedMessageInjectedAfterToolBatchBeforeNextModelCall(t *testing.T) {
76 client := &queueFakeClient{}
77 tool := &blockingTool{started: make(chan struct{}), release: make(chan struct{})}
78 a := NewAgent(client, "fake-model", "system", Registry{"echo": tool})
79
80 var (
81 mu sync.Mutex
82 texts []string
83 )
84 sink := func(ev AgentEvent) {
85 switch e := ev.(type) {
86 case EvUserMessage:
87 mu.Lock()
88 texts = append(texts, "user:"+extractText(e.Message))
89 mu.Unlock()
90 case EvAssistantMessage:
91 mu.Lock()
92 texts = append(texts, "asst:"+extractText(e.Message))
93 mu.Unlock()
94 }
95 }
96
97 done := make(chan error, 1)
98 go func() {
99 done <- a.Prompt(context.Background(), "do X", nil, sink)
100 }()
101
102 <-tool.started
103 if !a.QueueMessage("also do Y", nil) {
104 t.Fatal("QueueMessage returned false")
105 }
106 close(tool.release)
107
108 if err := <-done; err != nil {
109 t.Fatalf("Prompt returned %v", err)
110 }
111 if got := atomic.LoadInt32(&client.calls); got != 2 {
112 t.Fatalf("Stream calls = %d; want 2", got)
113 }
114
115 mu.Lock()
116 defer mu.Unlock()
117 if !queueTestContains(texts, "user:also do Y") {
118 t.Fatalf("queued message was not emitted as user message; texts=%v", texts)
119 }
120 if !queueTestContains(texts, "asst:saw queued") {
121 t.Fatalf("second assistant response missing; texts=%v", texts)
122 }
123}
124
125func queueTestContains(xs []string, want string) bool {
126 for _, x := range xs {
127 if x == want {
128 return true
129 }
130 }
131 return false
132}
133
134func TestQueuedImageMessageInjectedAtNextSafeBoundary(t *testing.T) {
135 client := &queueFakeClient{}
136 tool := &blockingTool{started: make(chan struct{}), release: make(chan struct{})}
137 a := NewAgent(client, "fake-model", "system", Registry{"echo": tool})
138
139 var queued provider.Message
140 done := make(chan error, 1)
141 go func() {
142 done <- a.Prompt(context.Background(), "inspect later", nil, func(ev AgentEvent) {
143 if user, ok := ev.(EvUserMessage); ok && len(user.Message.Content) > 0 {
144 if _, isImage := user.Message.Content[0].(provider.ImageBlock); isImage {
145 queued = user.Message
146 }
147 }
148 })
149 }()
150
151 <-tool.started
152 image := provider.ImageBlock{MimeType: "image/png", Data: []byte("png-bytes")}
153 if !a.QueueMessage("", []provider.ImageBlock{image}) {
154 t.Fatal("QueueMessage returned false for image-only prompt")
155 }
156 close(tool.release)
157
158 if err := <-done; err != nil {
159 t.Fatalf("Prompt returned %v", err)
160 }
161 if len(queued.Content) != 1 {
162 t.Fatalf("queued content = %#v, want one image", queued.Content)
163 }
164 got, ok := queued.Content[0].(provider.ImageBlock)
165 if !ok || got.MimeType != image.MimeType || string(got.Data) != string(image.Data) {
166 t.Fatalf("queued image = %#v, want %#v", queued.Content[0], image)
167 }
168}
169
170func TestQueuedMessageKeepsAcceptedTimeAcrossDrain(t *testing.T) {
171 a := NewAgent(nil, "fake", "", Registry{})
172 before := time.Now()
173 if !a.QueueMessage("queued", nil) {
174 t.Fatal("QueueMessage returned false")
175 }
176 after := time.Now()
177
178 queued := a.drainQueuedMessages()
179 if len(queued) != 1 {
180 t.Fatalf("drained %d messages, want 1", len(queued))
181 }
182 if queued[0].accepted.Before(before) || queued[0].accepted.After(after) {
183 t.Fatalf("accepted time %s outside queue submission window %s..%s", queued[0].accepted, before, after)
184 }
185}
186
187func TestQueuePromptPreservesHostProvenance(t *testing.T) {
188 a := NewAgent(nil, "fake", "", Registry{})
189 prompt := QueuedMessage{Text: "host evidence", HostEvent: true}
190 if !a.QueuePrompt(prompt) {
191 t.Fatal("QueuePrompt rejected host event")
192 }
193 if got := a.PendingQueuedMessages(); len(got) != 1 || !got[0].HostEvent {
194 t.Fatalf("snapshot lost host provenance: %#v", got)
195 }
196 got := a.DrainQueuedMessages()
197 if len(got) != 1 || !got[0].HostEvent {
198 t.Fatalf("drain lost host provenance: %#v", got)
199 }
200 a.QueuePrompt(got[0])
201 if popped, ok := a.PopQueuedMessage(); !ok || !popped.HostEvent {
202 t.Fatalf("requeued prompt lost host provenance: %#v, %v", popped, ok)
203 }
204}
205
206func TestPopQueuedUserMessagePreservesHostEvents(t *testing.T) {
207 a := NewAgent(nil, "fake", "", Registry{})
208 a.QueuePrompt(QueuedMessage{Text: "first report", HostEvent: true})
209 a.QueueMessage("older draft", nil)
210 a.QueueMessage("newest draft", []provider.ImageBlock{{MimeType: "image/png", Data: []byte("png-1")}})
211 a.QueuePrompt(QueuedMessage{Text: "latest report", HostEvent: true})
212 before := append([]queuedMessage(nil), a.queued...)
213
214 message, ok := a.PopQueuedUserMessage()
215 if !ok || message.Text != "newest draft" || len(message.Images) != 1 || string(message.Images[0].Data) != "png-1" {
216 t.Fatalf("popped = %#v, %v, want newest draft with image", message, ok)
217 }
218 for index, original := range []int{0, 1, 3} {
219 if !a.queued[index].accepted.Equal(before[original].accepted) {
220 t.Fatal("recall changed a remaining message's accepted time")
221 }
222 }
223 if message, ok := a.PopQueuedUserMessage(); !ok || message.Text != "older draft" {
224 t.Fatalf("second pop = %#v, %v, want older draft", message, ok)
225 }
226 if message, ok := a.PopQueuedUserMessage(); ok {
227 t.Fatalf("host-only pop = %#v, want no user message", message)
228 }
229 pending := a.DrainQueuedMessages()
230 if len(pending) != 2 || pending[0].Text != "first report" || pending[1].Text != "latest report" || !pending[0].HostEvent || !pending[1].HostEvent {
231 t.Fatalf("remaining reports = %#v, want original reports in order", pending)
232 }
233 if _, ok := a.PopQueuedUserMessage(); ok {
234 t.Fatal("empty pop succeeded")
235 }
236}
237
238func TestQueueMessageSnapshotPopAndDrain(t *testing.T) {
239 a := NewAgent(nil, "fake", "", Registry{})
240 if a.QueueMessage(" ", nil) {
241 t.Fatal("blank queue message accepted")
242 }
243 a.QueueMessage("one", nil)
244 a.QueueMessage("two", nil)
245 if got := a.PendingQueuedMessages(); len(got) != 2 || got[0].Text != "one" || got[1].Text != "two" {
246 t.Fatalf("PendingQueuedMessages = %v; want [one two]", got)
247 }
248 if message, ok := a.PopQueuedMessage(); !ok || message.Text != "two" {
249 t.Fatalf("PopQueuedMessage = %#v,%v; want two,true", message, ok)
250 }
251 if got := a.DrainQueuedMessages(); len(got) != 1 || got[0].Text != "one" {
252 t.Fatalf("DrainQueuedMessages = %v; want [one]", got)
253 }
254 if got := a.QueuedMessageCount(); got != 0 {
255 t.Fatalf("QueuedMessageCount = %d; want 0", got)
256 }
257}