runner_test.go RawHistoryFolder

1package bot
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)
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}