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") {
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 TestQueuedMessageKeepsAcceptedTimeAcrossDrain(t *testing.T) {
135 a := NewAgent(nil, "fake", "", Registry{})
136 before := time.Now()
137 if !a.QueueMessage("queued") {
138 t.Fatal("QueueMessage returned false")
139 }
140 after := time.Now()
141
142 queued := a.drainQueuedMessages()
143 if len(queued) != 1 {
144 t.Fatalf("drained %d messages, want 1", len(queued))
145 }
146 if queued[0].accepted.Before(before) || queued[0].accepted.After(after) {
147 t.Fatalf("accepted time %s outside queue submission window %s..%s", queued[0].accepted, before, after)
148 }
149}
150
151func TestQueueMessageSnapshotPopAndDrain(t *testing.T) {
152 a := NewAgent(nil, "fake", "", Registry{})
153 if a.QueueMessage(" ") {
154 t.Fatal("blank queue message accepted")
155 }
156 a.QueueMessage("one")
157 a.QueueMessage("two")
158 if got := a.PendingQueuedMessages(); len(got) != 2 || got[0] != "one" || got[1] != "two" {
159 t.Fatalf("PendingQueuedMessages = %v; want [one two]", got)
160 }
161 if text, ok := a.PopQueuedMessage(); !ok || text != "two" {
162 t.Fatalf("PopQueuedMessage = %q,%v; want two,true", text, ok)
163 }
164 if got := a.DrainQueuedMessages(); len(got) != 1 || got[0] != "one" {
165 t.Fatalf("DrainQueuedMessages = %v; want [one]", got)
166 }
167 if got := a.QueuedMessageCount(); got != 0 {
168 t.Fatalf("QueuedMessageCount = %d; want 0", got)
169 }
170}