1package modes
2
3import (
4 "context"
5 "sync"
6 "testing"
7 "time"
8
9 "github.com/patriceckhart/zot/packages/core"
10 "github.com/patriceckhart/zot/packages/provider"
11 "github.com/patriceckhart/zot/packages/tui"
12)
13
14type compactQueueClient struct {
15 compactionStarted chan struct{}
16 releaseCompaction chan struct{}
17 followUpRequest chan provider.Request
18
19 mu sync.Mutex
20 calls int
21}
22
23func (c *compactQueueClient) Name() string { return "compact-queue-test" }
24
25func (c *compactQueueClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
26 c.mu.Lock()
27 c.calls++
28 call := c.calls
29 c.mu.Unlock()
30
31 out := make(chan provider.Event, 2)
32 go func() {
33 defer close(out)
34 if call == 1 {
35 close(c.compactionStarted)
36 select {
37 case <-ctx.Done():
38 out <- provider.EventDone{Stop: provider.StopAborted, Err: ctx.Err()}
39 case <-c.releaseCompaction:
40 out <- provider.EventTextDelta{Delta: "summary"}
41 out <- provider.EventDone{Stop: provider.StopEnd}
42 }
43 return
44 }
45
46 c.followUpRequest <- req
47 out <- provider.EventDone{
48 Stop: provider.StopEnd,
49 Message: provider.Message{
50 Role: provider.RoleAssistant,
51 Content: []provider.Content{provider.TextBlock{Text: "done"}},
52 },
53 }
54 }()
55 return out, nil
56}
57
58func TestPromptSubmittedDuringCompactionStartsFollowUpTurn(t *testing.T) {
59 client := &compactQueueClient{
60 compactionStarted: make(chan struct{}),
61 releaseCompaction: make(chan struct{}),
62 followUpRequest: make(chan provider.Request, 1),
63 }
64 agent := core.NewAgent(client, "test-model", "", nil)
65 agent.SetMessages([]provider.Message{
66 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "one"}}},
67 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "two"}}},
68 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "three"}}},
69 {Role: provider.RoleAssistant, Content: []provider.Content{provider.TextBlock{Text: "four"}}},
70 {Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "five"}}},
71 })
72 interactive := NewInteractive(InteractiveConfig{Agent: agent})
73 interactive.runCtx = context.Background()
74
75 interactive.runCompact(context.Background(), false)
76 select {
77 case <-client.compactionStarted:
78 case <-time.After(2 * time.Second):
79 t.Fatal("compaction did not start")
80 }
81
82 interactive.ed.SetValue("follow up")
83 interactive.handleKey(context.Background(), tui.Key{Kind: tui.KeyEnter})
84 close(client.releaseCompaction)
85
86 select {
87 case req := <-client.followUpRequest:
88 if !requestContainsUserText(req, "follow up") {
89 t.Fatalf("follow-up request does not contain queued prompt: %#v", req.Messages)
90 }
91 case <-time.After(2 * time.Second):
92 t.Fatal("queued prompt did not start a turn after compaction")
93 }
94}
95
96func requestContainsUserText(req provider.Request, want string) bool {
97 for _, message := range req.Messages {
98 if message.Role != provider.RoleUser {
99 continue
100 }
101 for _, content := range message.Content {
102 if text, ok := content.(provider.TextBlock); ok && text.Text == want {
103 return true
104 }
105 }
106 }
107 return false
108}