1package core
2
3import (
4 "context"
5 "encoding/json"
6 "sync"
7 "sync/atomic"
8 "testing"
9
10 "github.com/patriceckhart/zot/internal/provider"
11)
12
13type queueFakeClient struct {
14 calls int32
15}
16
17func (c *queueFakeClient) Name() string { return "queue-fake" }
18
19func (c *queueFakeClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
20 call := atomic.AddInt32(&c.calls, 1)
21 out := make(chan provider.Event, 4)
22 go func() {
23 defer close(out)
24 out <- provider.EventStart{Provider: "queue-fake", Model: req.Model}
25 switch call {
26 case 1:
27 out <- provider.EventToolStart{ID: "t1", Name: "echo"}
28 out <- provider.EventToolEnd{ID: "t1"}
29 out <- provider.EventDone{Stop: provider.StopToolUse, Message: provider.Message{
30 Role: provider.RoleAssistant,
31 Content: []provider.Content{
32 provider.TextBlock{Text: "using tool"},
33 provider.ToolCallBlock{ID: "t1", Name: "echo", Arguments: json.RawMessage(`{}`)},
34 },
35 }}
36 case 2:
37 out <- provider.EventTextDelta{Delta: "saw queued"}
38 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
39 Role: provider.RoleAssistant,
40 Content: []provider.Content{provider.TextBlock{Text: "saw queued"}},
41 }}
42 default:
43 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
44 Role: provider.RoleAssistant,
45 Content: []provider.Content{provider.TextBlock{Text: "extra"}},
46 }}
47 }
48 }()
49 return out, nil
50}
51
52// blockingTool waits until the test has queued a message, then
53// returns. This pins the core behaviour: queued user text is delivered
54// after the current tool batch finishes and before the next model call.
55type blockingTool struct {
56 started chan struct{}
57 release chan struct{}
58}
59
60func (t *blockingTool) Name() string { return "echo" }
61func (t *blockingTool) Description() string { return "echoes" }
62func (t *blockingTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) }
63
64func (t *blockingTool) Execute(ctx context.Context, args json.RawMessage, progress func(string)) (ToolResult, error) {
65 close(t.started)
66 select {
67 case <-ctx.Done():
68 return ToolResult{Content: []provider.Content{provider.TextBlock{Text: ctx.Err().Error()}}, IsError: true}, ctx.Err()
69 case <-t.release:
70 }
71 return ToolResult{Content: []provider.Content{provider.TextBlock{Text: "tool ok"}}}, nil
72}
73
74func TestQueuedMessageInjectedAfterToolBatchBeforeNextModelCall(t *testing.T) {
75 client := &queueFakeClient{}
76 tool := &blockingTool{started: make(chan struct{}), release: make(chan struct{})}
77 a := NewAgent(client, "fake-model", "system", Registry{"echo": tool})
78
79 var (
80 mu sync.Mutex
81 texts []string
82 )
83 sink := func(ev AgentEvent) {
84 switch e := ev.(type) {
85 case EvUserMessage:
86 mu.Lock()
87 texts = append(texts, "user:"+extractText(e.Message))
88 mu.Unlock()
89 case EvAssistantMessage:
90 mu.Lock()
91 texts = append(texts, "asst:"+extractText(e.Message))
92 mu.Unlock()
93 }
94 }
95
96 done := make(chan error, 1)
97 go func() {
98 done <- a.Prompt(context.Background(), "do X", nil, sink)
99 }()
100
101 <-tool.started
102 if !a.QueueMessage("also do Y") {
103 t.Fatal("QueueMessage returned false")
104 }
105 close(tool.release)
106
107 if err := <-done; err != nil {
108 t.Fatalf("Prompt returned %v", err)
109 }
110 if got := atomic.LoadInt32(&client.calls); got != 2 {
111 t.Fatalf("Stream calls = %d; want 2", got)
112 }
113
114 mu.Lock()
115 defer mu.Unlock()
116 if !queueTestContains(texts, "user:also do Y") {
117 t.Fatalf("queued message was not emitted as user message; texts=%v", texts)
118 }
119 if !queueTestContains(texts, "asst:saw queued") {
120 t.Fatalf("second assistant response missing; texts=%v", texts)
121 }
122}
123
124func queueTestContains(xs []string, want string) bool {
125 for _, x := range xs {
126 if x == want {
127 return true
128 }
129 }
130 return false
131}
132
133func TestQueueMessageSnapshotPopAndDrain(t *testing.T) {
134 a := NewAgent(nil, "fake", "", Registry{})
135 if a.QueueMessage(" ") {
136 t.Fatal("blank queue message accepted")
137 }
138 a.QueueMessage("one")
139 a.QueueMessage("two")
140 if got := a.PendingQueuedMessages(); len(got) != 2 || got[0] != "one" || got[1] != "two" {
141 t.Fatalf("PendingQueuedMessages = %v; want [one two]", got)
142 }
143 if text, ok := a.PopQueuedMessage(); !ok || text != "two" {
144 t.Fatalf("PopQueuedMessage = %q,%v; want two,true", text, ok)
145 }
146 if got := a.DrainQueuedMessages(); len(got) != 1 || got[0] != "one" {
147 t.Fatalf("DrainQueuedMessages = %v; want [one]", got)
148 }
149 if got := a.QueuedMessageCount(); got != 0 {
150 t.Fatalf("QueuedMessageCount = %d; want 0", got)
151 }
152}