1package bot
2
3import (
4 "context"
5 "sync"
6 "testing"
7 "time"
8
9 "github.com/bnema/zut/packages/core"
10 "github.com/bnema/zut/packages/provider"
11)
12
13type testAdapter struct{}
14
15func (testAdapter) Run(context.Context, func(InboundMessage), func(Command, InboundMessage)) error {
16 return nil
17}
18func (testAdapter) Send(context.Context, string, string, SendOptions) error {
19 return nil
20}
21func (testAdapter) IndicateWorking(context.Context, string) func() { return func() {} }
22func (testAdapter) StatusText() string { return "" }
23
24type blockingClient struct {
25 started chan struct{}
26 release chan struct{}
27
28 mu sync.Mutex
29 active int
30 maxActive int
31}
32
33func (c *blockingClient) Name() string { return "test" }
34
35func (c *blockingClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
36 c.mu.Lock()
37 c.active++
38 if c.active > c.maxActive {
39 c.maxActive = c.active
40 }
41 c.mu.Unlock()
42
43 select {
44 case c.started <- struct{}{}:
45 default:
46 }
47
48 out := make(chan provider.Event, 1)
49 go func() {
50 defer close(out)
51 defer func() {
52 c.mu.Lock()
53 c.active--
54 c.mu.Unlock()
55 }()
56 select {
57 case <-ctx.Done():
58 out <- provider.EventDone{Stop: provider.StopAborted, Err: ctx.Err()}
59 case <-c.release:
60 out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
61 Role: provider.RoleAssistant,
62 Content: []provider.Content{provider.TextBlock{Text: "ok"}},
63 }}
64 }
65 }()
66 return out, nil
67}
68
69func (c *blockingClient) max() int {
70 c.mu.Lock()
71 defer c.mu.Unlock()
72 return c.maxActive
73}
74
75func TestHandleMessageClaimsDrainSlotBeforeSpawningDrainer(t *testing.T) {
76 client := &blockingClient{started: make(chan struct{}, 2), release: make(chan struct{})}
77 r := NewRunner(testAdapter{}, core.NewAgent(client, "test-model", "", nil), Config{})
78 r.runCtx = context.Background()
79
80 var wg sync.WaitGroup
81 wg.Add(2)
82 go func() {
83 defer wg.Done()
84 r.handleMessage(InboundMessage{ChannelID: "c", Text: "one"})
85 }()
86 go func() {
87 defer wg.Done()
88 r.handleMessage(InboundMessage{ChannelID: "c", Text: "two"})
89 }()
90 wg.Wait()
91
92 select {
93 case <-client.started:
94 case <-time.After(2 * time.Second):
95 t.Fatal("first turn did not start")
96 }
97
98 // Give any accidentally spawned second drainer a chance to enter Stream.
99 time.Sleep(100 * time.Millisecond)
100 if got := client.max(); got != 1 {
101 t.Fatalf("concurrent provider streams = %d, want 1", got)
102 }
103
104 close(client.release)
105}
106
107type cacheOnlyUsageClient struct{}
108
109func (cacheOnlyUsageClient) Name() string { return "cache-only-usage" }
110
111func (cacheOnlyUsageClient) Stream(context.Context, provider.Request) (<-chan provider.Event, error) {
112 events := make(chan provider.Event, 2)
113 events <- provider.EventUsage{Usage: provider.Usage{
114 CacheReadTokens: 400,
115 CacheWriteTokens: 50,
116 }}
117 events <- provider.EventDone{Stop: provider.StopEnd}
118 close(events)
119 return events, nil
120}
121
122func TestRunTurnUpdatesContextUsageFromCacheOnlyTurn(t *testing.T) {
123 runner := NewRunner(testAdapter{}, core.NewAgent(cacheOnlyUsageClient{}, "test-model", "", nil), Config{})
124 runner.lastCtxInput = 99
125
126 runner.runTurn(context.Background(), queuedTurn{channelID: "channel", prompt: "prompt"})
127
128 runner.mu.Lock()
129 got := runner.lastCtxInput
130 runner.mu.Unlock()
131 if got != 450 {
132 t.Fatalf("context usage = %d, want 450", got)
133 }
134}