1ae0519 fix: address bot refactor review issues · zut

M packages/agent/botcmd.go +10 -4


@@ -10,7 +10,6 @@ import (
10	10		"os/exec"
11	11		"os/signal"
12	12		"path/filepath"
13		-	"strconv"
14	13		"strings"
15	14		"syscall"
16	15		"time"

@@ -18,6 +17,7 @@ import (
18	17		"github.com/patriceckhart/zot/packages/agent/modes/bot"
19	18		"github.com/patriceckhart/zot/packages/agent/modes/telegram"
20	19		"github.com/patriceckhart/zot/packages/core"
	20	+	"github.com/patriceckhart/zot/packages/provider"
21	21		)
22	22
23	23		// detachChild configures cmd to run in its own process group so tty

@@ -357,6 +357,15 @@ func botRun(rawTail []string, version string) error {
357	357		s, _, serr := openOrCreateSessionForBot(args, resolved, agent, version)
358	358		if serr == nil {
359	359		sess = s
	360	+	agent.OnMessageAppended = func(msg provider.Message) {
	361	+	_ = sess.AppendMessage(msg)
	362	+	}
	363	+	agent.OnUsage = func(u provider.Usage) {
	364	+	_ = sess.AppendUsage(u, u)
	365	+	}
	366	+	agent.OnTranscriptCompacted = func(msgs []provider.Message) {
	367	+	_ = sess.AppendCompaction(msgs)
	368	+	}
360	369		defer sess.Close()
361	370		} else {
362	371		fmt.Fprintln(os.Stderr, "session:", serr)

@@ -445,6 +454,3 @@ func maskToken(tok string) string {
445	454		}
446	455		return tok[:i+1] + body[:3] + "..." + body[len(body)-3:]
447	456		}
448		-
449		-	// _ compile-time hint so the strconv import stays if we later add numeric parsing.
450		-	var _ = strconv.Itoa

 	"os/exec"
 	"os/signal"
 	"path/filepath"
-	"strconv"
 	"strings"
 	"syscall"
 	"time"
 	"github.com/patriceckhart/zot/packages/agent/modes/bot"
 	"github.com/patriceckhart/zot/packages/agent/modes/telegram"
 	"github.com/patriceckhart/zot/packages/core"
+	"github.com/patriceckhart/zot/packages/provider"
 )
 // detachChild configures cmd to run in its own process group so tty
 		s, _, serr := openOrCreateSessionForBot(args, resolved, agent, version)
 		if serr == nil {
 			sess = s
+			agent.OnMessageAppended = func(msg provider.Message) {
+				_ = sess.AppendMessage(msg)
+			}
+			agent.OnUsage = func(u provider.Usage) {
+				_ = sess.AppendUsage(u, u)
+			}
+			agent.OnTranscriptCompacted = func(msgs []provider.Message) {
+				_ = sess.AppendCompaction(msgs)
+			}
 			defer sess.Close()
 		} else {
 			fmt.Fprintln(os.Stderr, "session:", serr)
 	}
 	return tok[:i+1] + body[:3] + "..." + body[len(body)-3:]
 }
-// _ compile-time hint so the strconv import stays if we later add numeric parsing.
-var _ = strconv.Itoa

M packages/agent/modes/bot/runner.go +3 -3


@@ -90,6 +90,9 @@ func (r *Runner) handleMessage(msg InboundMessage) {
90	90		images: msg.Images,
91	91		})
92	92		idle := !r.busy
	93	+	if idle {
	94	+	r.busy = true
	95	+	}
93	96		r.mu.Unlock()
94	97
95	98		if idle {

@@ -126,7 +129,6 @@ func (r *Runner) drainQueue() {
126	129		}
127	130		t := r.queue[0]
128	131		r.queue = r.queue[1:]
129		-	r.busy = true
130	132		turnCtx, cancel := context.WithCancel(parent)
131	133		r.activeCtx = cancel
132	134		r.mu.Unlock()

@@ -250,5 +252,3 @@ func (r *Runner) sendStatus(channelID string) {
250	252
251	253		_ = r.adapter.Send(context.Background(), channelID, status)
252	254		}
253		-
254		-

 		images:    msg.Images,
 	})
 	idle := !r.busy
+	if idle {
+		r.busy = true
+	}
 	r.mu.Unlock()
 	if idle {
 		}
 		t := r.queue[0]
 		r.queue = r.queue[1:]
-		r.busy = true
 		turnCtx, cancel := context.WithCancel(parent)
 		r.activeCtx = cancel
 		r.mu.Unlock()
 	_ = r.adapter.Send(context.Background(), channelID, status)
 }

A packages/agent/modes/bot/runner_test.go +103 -0

@@ -0,0 +1,103 @@
	1	+	package bot
	2	+
	3	+	import (
	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	+
	13	+	type testAdapter struct{}
	14	+
	15	+	func (testAdapter) Run(context.Context, func(InboundMessage), func(Command, InboundMessage)) error {
	16	+	return nil
	17	+	}
	18	+	func (testAdapter) Send(context.Context, string, string) error { return nil }
	19	+	func (testAdapter) IndicateWorking(context.Context, string) func() { return func() {} }
	20	+	func (testAdapter) StatusText() string { return "" }
	21	+
	22	+	type blockingClient struct {
	23	+	started chan struct{}
	24	+	release chan struct{}
	25	+
	26	+	mu sync.Mutex
	27	+	active int
	28	+	maxActive int
	29	+	}
	30	+
	31	+	func (c *blockingClient) Name() string { return "test" }
	32	+
	33	+	func (c *blockingClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
	34	+	c.mu.Lock()
	35	+	c.active++
	36	+	if c.active > c.maxActive {
	37	+	c.maxActive = c.active
	38	+	}
	39	+	c.mu.Unlock()
	40	+
	41	+	select {
	42	+	case c.started <- struct{}{}:
	43	+	default:
	44	+	}
	45	+
	46	+	out := make(chan provider.Event, 1)
	47	+	go func() {
	48	+	defer close(out)
	49	+	defer func() {
	50	+	c.mu.Lock()
	51	+	c.active--
	52	+	c.mu.Unlock()
	53	+	}()
	54	+	select {
	55	+	case <-ctx.Done():
	56	+	out <- provider.EventDone{Stop: provider.StopAborted, Err: ctx.Err()}
	57	+	case <-c.release:
	58	+	out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
	59	+	Role: provider.RoleAssistant,
	60	+	Content: []provider.Content{provider.TextBlock{Text: "ok"}},
	61	+	}}
	62	+	}
	63	+	}()
	64	+	return out, nil
	65	+	}
	66	+
	67	+	func (c *blockingClient) max() int {
	68	+	c.mu.Lock()
	69	+	defer c.mu.Unlock()
	70	+	return c.maxActive
	71	+	}
	72	+
	73	+	func TestHandleMessageClaimsDrainSlotBeforeSpawningDrainer(t *testing.T) {
	74	+	client := &blockingClient{started: make(chan struct{}, 2), release: make(chan struct{})}
	75	+	r := NewRunner(testAdapter{}, core.NewAgent(client, "test-model", "", nil), Config{})
	76	+	r.runCtx = context.Background()
	77	+
	78	+	var wg sync.WaitGroup
	79	+	wg.Add(2)
	80	+	go func() {
	81	+	defer wg.Done()
	82	+	r.handleMessage(InboundMessage{ChannelID: "c", Text: "one"})
	83	+	}()
	84	+	go func() {
	85	+	defer wg.Done()
	86	+	r.handleMessage(InboundMessage{ChannelID: "c", Text: "two"})
	87	+	}()
	88	+	wg.Wait()
	89	+
	90	+	select {
	91	+	case <-client.started:
	92	+	case <-time.After(2 * time.Second):
	93	+	t.Fatal("first turn did not start")
	94	+	}
	95	+
	96	+	// Give any accidentally spawned second drainer a chance to enter Stream.
	97	+	time.Sleep(100 * time.Millisecond)
	98	+	if got := client.max(); got != 1 {
	99	+	t.Fatalf("concurrent provider streams = %d, want 1", got)
	100	+	}
	101	+
	102	+	close(client.release)
	103	+	}

@@ -0,0 +1,103 @@
	1	+	package bot
	2	+
	3	+	import (
	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	+
	13	+	type testAdapter struct{}
	14	+
	15	+	func (testAdapter) Run(context.Context, func(InboundMessage), func(Command, InboundMessage)) error {
	16	+	return nil
	17	+	}
	18	+	func (testAdapter) Send(context.Context, string, string) error { return nil }
	19	+	func (testAdapter) IndicateWorking(context.Context, string) func() { return func() {} }
	20	+	func (testAdapter) StatusText() string { return "" }
	21	+
	22	+	type blockingClient struct {
	23	+	started chan struct{}
	24	+	release chan struct{}
	25	+
	26	+	mu sync.Mutex
	27	+	active int
	28	+	maxActive int
	29	+	}
	30	+
	31	+	func (c *blockingClient) Name() string { return "test" }
	32	+
	33	+	func (c *blockingClient) Stream(ctx context.Context, req provider.Request) (<-chan provider.Event, error) {
	34	+	c.mu.Lock()
	35	+	c.active++
	36	+	if c.active > c.maxActive {
	37	+	c.maxActive = c.active
	38	+	}
	39	+	c.mu.Unlock()
	40	+
	41	+	select {
	42	+	case c.started <- struct{}{}:
	43	+	default:
	44	+	}
	45	+
	46	+	out := make(chan provider.Event, 1)
	47	+	go func() {
	48	+	defer close(out)
	49	+	defer func() {
	50	+	c.mu.Lock()
	51	+	c.active--
	52	+	c.mu.Unlock()
	53	+	}()
	54	+	select {
	55	+	case <-ctx.Done():
	56	+	out <- provider.EventDone{Stop: provider.StopAborted, Err: ctx.Err()}
	57	+	case <-c.release:
	58	+	out <- provider.EventDone{Stop: provider.StopEnd, Message: provider.Message{
	59	+	Role: provider.RoleAssistant,
	60	+	Content: []provider.Content{provider.TextBlock{Text: "ok"}},
	61	+	}}
	62	+	}
	63	+	}()
	64	+	return out, nil
	65	+	}
	66	+
	67	+	func (c *blockingClient) max() int {
	68	+	c.mu.Lock()
	69	+	defer c.mu.Unlock()
	70	+	return c.maxActive
	71	+	}
	72	+
	73	+	func TestHandleMessageClaimsDrainSlotBeforeSpawningDrainer(t *testing.T) {
	74	+	client := &blockingClient{started: make(chan struct{}, 2), release: make(chan struct{})}
	75	+	r := NewRunner(testAdapter{}, core.NewAgent(client, "test-model", "", nil), Config{})
	76	+	r.runCtx = context.Background()
	77	+
	78	+	var wg sync.WaitGroup
	79	+	wg.Add(2)
	80	+	go func() {
	81	+	defer wg.Done()
	82	+	r.handleMessage(InboundMessage{ChannelID: "c", Text: "one"})
	83	+	}()
	84	+	go func() {
	85	+	defer wg.Done()
	86	+	r.handleMessage(InboundMessage{ChannelID: "c", Text: "two"})
	87	+	}()
	88	+	wg.Wait()
	89	+
	90	+	select {
	91	+	case <-client.started:
	92	+	case <-time.After(2 * time.Second):
	93	+	t.Fatal("first turn did not start")
	94	+	}
	95	+
	96	+	// Give any accidentally spawned second drainer a chance to enter Stream.
	97	+	time.Sleep(100 * time.Millisecond)
	98	+	if got := client.max(); got != 1 {
	99	+	t.Fatalf("concurrent provider streams = %d, want 1", got)
	100	+	}
	101	+
	102	+	close(client.release)
	103	+	}

M packages/agent/modes/telegram/bot.go +36 -51

@@ -1,83 +1,68 @@
1	1		package telegram
2	2
3	3		import (
4		-	"context"
5	4		"strings"
6		-
7		-	"github.com/patriceckhart/zot/packages/agent/modes/bot"
8		-	"github.com/patriceckhart/zot/packages/core"
	5	+	"unicode/utf8"
9	6		)
10	7
11		-	// Bot is a thin wrapper around bot.Runner + Adapter. It exists to
12		-	// keep the exported API stable for botcmd.go.
13		-	//
14		-	// Deprecated: Use telegram.NewAdapter + bot.NewRunner directly.
15		-	// Bot is kept for backward compatibility but is no longer used internally.
16		-	type Bot struct {
17		-	Client *Client
18		-	Agent *core.Agent
19		-	Config Config
20		-	ZotHome string
21		-	Provider string
22		-	AuthMethod string
23		-	CWD string
24		-	Save func(Config) error
25		-	RefreshCreds func() error
26		-
27		-	runner *bot.Runner
28		-	adapter *Adapter
29		-	}
30		-
31		-	// Run starts the bot. It constructs the adapter and runner on first
32		-	// call, then delegates to runner.Run.
33		-	func (b *Bot) Run(ctx context.Context) error {
34		-	if b.adapter == nil {
35		-	b.adapter = NewAdapter(b.Client, &b.Config, b.Save)
36		-	}
37		-	if b.runner == nil {
38		-	b.runner = bot.NewRunner(b.adapter, b.Agent, bot.Config{
39		-	ZotHome: b.ZotHome,
40		-	Provider: b.Provider,
41		-	AuthMethod: b.AuthMethod,
42		-	CWD: b.CWD,
43		-	RefreshCreds: b.RefreshCreds,
44		-	})
45		-	}
46		-	return b.runner.Run(ctx)
47		-	}
48		-
49	8		// chunkMessage splits s into chunks no larger than limit runes, on line
50	9		// boundaries when possible.
51	10		func chunkMessage(s string, limit int) []string {
52		-	if len(s) <= limit {
	11	+	if limit <= 0 \|\| utf8.RuneCountInString(s) <= limit {
53	12		return []string{s}
54	13		}
55	14		var out []string
56	15		lines := strings.Split(s, "\n")
57	16		var cur strings.Builder
	17	+	curRunes := 0
58	18		for _, l := range lines {
59		-	if cur.Len()+len(l)+1 > limit && cur.Len() > 0 {
	19	+	lineRunes := utf8.RuneCountInString(l)
	20	+	sepRunes := 0
	21	+	if curRunes > 0 {
	22	+	sepRunes = 1
	23	+	}
	24	+	if curRunes+sepRunes+lineRunes > limit && curRunes > 0 {
60	25		out = append(out, cur.String())
61	26		cur.Reset()
	27	+	curRunes = 0
	28	+	sepRunes = 0
62	29		}
63		-	if len(l) > limit {
64		-	// Line itself too long; hard-split.
65		-	for len(l) > limit {
66		-	out = append(out, l[:limit])
67		-	l = l[limit:]
	30	+	if lineRunes > limit {
	31	+	// Line itself too long; hard-split on rune boundaries.
	32	+	for lineRunes > limit {
	33	+	i := byteIndexAfterRunes(l, limit)
	34	+	out = append(out, l[:i])
	35	+	l = l[i:]
	36	+	lineRunes = utf8.RuneCountInString(l)
68	37		}
69	38		}
70		-	if cur.Len() > 0 {
	39	+	if curRunes > 0 {
71	40		cur.WriteString("\n")
	41	+	curRunes++
72	42		}
73	43		cur.WriteString(l)
	44	+	curRunes += utf8.RuneCountInString(l)
74	45		}
75		-	if cur.Len() > 0 {
	46	+	if curRunes > 0 {
76	47		out = append(out, cur.String())
77	48		}
78	49		return out
79	50		}
80	51
	52	+	func byteIndexAfterRunes(s string, n int) int {
	53	+	if n <= 0 {
	54	+	return 0
	55	+	}
	56	+	count := 0
	57	+	for i := range s {
	58	+	if count == n {
	59	+	return i
	60	+	}
	61	+	count++
	62	+	}
	63	+	return len(s)
	64	+	}
	65	+
81	66		// isImageMIME returns true for MIME types the model can probably ingest
82	67		// as a vision input.
83	68		func isImageMIME(m string) bool {

@@ -1,83 +1,68 @@
1		package telegram	1		package telegram
2			2
3		import (	3		import (
4	-	"context"
5		"strings"	4		"strings"
6	-		5	+	"unicode/utf8"
7	-	"github.com/patriceckhart/zot/packages/agent/modes/bot"
8	-	"github.com/patriceckhart/zot/packages/core"
9		)	6		)
10			7
11	-	// Bot is a thin wrapper around bot.Runner + Adapter. It exists to
12	-	// keep the exported API stable for botcmd.go.
13	-	//
14	-	// Deprecated: Use telegram.NewAdapter + bot.NewRunner directly.
15	-	// Bot is kept for backward compatibility but is no longer used internally.
16	-	type Bot struct {
17	-	Client *Client
18	-	Agent *core.Agent
19	-	Config Config
20	-	ZotHome string
21	-	Provider string
22	-	AuthMethod string
23	-	CWD string
24	-	Save func(Config) error
25	-	RefreshCreds func() error
26	-
27	-	runner *bot.Runner
28	-	adapter *Adapter
29	-	}
30	-
31	-	// Run starts the bot. It constructs the adapter and runner on first
32	-	// call, then delegates to runner.Run.
33	-	func (b *Bot) Run(ctx context.Context) error {
34	-	if b.adapter == nil {
35	-	b.adapter = NewAdapter(b.Client, &b.Config, b.Save)
36	-	}
37	-	if b.runner == nil {
38	-	b.runner = bot.NewRunner(b.adapter, b.Agent, bot.Config{
39	-	ZotHome: b.ZotHome,
40	-	Provider: b.Provider,
41	-	AuthMethod: b.AuthMethod,
42	-	CWD: b.CWD,
43	-	RefreshCreds: b.RefreshCreds,
44	-	})
45	-	}
46	-	return b.runner.Run(ctx)
47	-	}
48	-
49		// chunkMessage splits s into chunks no larger than limit runes, on line	8		// chunkMessage splits s into chunks no larger than limit runes, on line
50		// boundaries when possible.	9		// boundaries when possible.
51		func chunkMessage(s string, limit int) []string {	10		func chunkMessage(s string, limit int) []string {
52	-	if len(s) <= limit {	11	+	if limit <= 0 \|\| utf8.RuneCountInString(s) <= limit {
53		return []string{s}	12		return []string{s}
54		}	13		}
55		var out []string	14		var out []string
56		lines := strings.Split(s, "\n")	15		lines := strings.Split(s, "\n")
57		var cur strings.Builder	16		var cur strings.Builder
			17	+	curRunes := 0
58		for _, l := range lines {	18		for _, l := range lines {
59	-	if cur.Len()+len(l)+1 > limit && cur.Len() > 0 {	19	+	lineRunes := utf8.RuneCountInString(l)
			20	+	sepRunes := 0
			21	+	if curRunes > 0 {
			22	+	sepRunes = 1
			23	+	}
			24	+	if curRunes+sepRunes+lineRunes > limit && curRunes > 0 {
60		out = append(out, cur.String())	25		out = append(out, cur.String())
61		cur.Reset()	26		cur.Reset()
			27	+	curRunes = 0
			28	+	sepRunes = 0
62		}	29		}
63	-	if len(l) > limit {	30	+	if lineRunes > limit {
64	-	// Line itself too long; hard-split.	31	+	// Line itself too long; hard-split on rune boundaries.
65	-	for len(l) > limit {	32	+	for lineRunes > limit {
66	-	out = append(out, l[:limit])	33	+	i := byteIndexAfterRunes(l, limit)
67	-	l = l[limit:]	34	+	out = append(out, l[:i])
			35	+	l = l[i:]
			36	+	lineRunes = utf8.RuneCountInString(l)
68		}	37		}
69		}	38		}
70	-	if cur.Len() > 0 {	39	+	if curRunes > 0 {
71		cur.WriteString("\n")	40		cur.WriteString("\n")
			41	+	curRunes++
72		}	42		}
73		cur.WriteString(l)	43		cur.WriteString(l)
			44	+	curRunes += utf8.RuneCountInString(l)
74		}	45		}
75	-	if cur.Len() > 0 {	46	+	if curRunes > 0 {
76		out = append(out, cur.String())	47		out = append(out, cur.String())
77		}	48		}
78		return out	49		return out
79		}	50		}
80			51
			52	+	func byteIndexAfterRunes(s string, n int) int {
			53	+	if n <= 0 {
			54	+	return 0
			55	+	}
			56	+	count := 0
			57	+	for i := range s {
			58	+	if count == n {
			59	+	return i
			60	+	}
			61	+	count++
			62	+	}
			63	+	return len(s)
			64	+	}
			65	+
81		// isImageMIME returns true for MIME types the model can probably ingest	66		// isImageMIME returns true for MIME types the model can probably ingest
82		// as a vision input.	67		// as a vision input.
83		func isImageMIME(m string) bool {	68		func isImageMIME(m string) bool {

A packages/agent/modes/telegram/bot_test.go +31 -0

@@ -0,0 +1,31 @@
	1	+	package telegram
	2	+
	3	+	import (
	4	+	"strings"
	5	+	"testing"
	6	+	"unicode/utf8"
	7	+	)
	8	+
	9	+	func TestChunkMessageUsesRuneLimit(t *testing.T) {
	10	+	s := strings.Repeat("界", 4001)
	11	+	chunks := chunkMessage(s, 4000)
	12	+	if len(chunks) != 2 {
	13	+	t.Fatalf("len(chunks) = %d, want 2", len(chunks))
	14	+	}
	15	+	if got := utf8.RuneCountInString(chunks[0]); got != 4000 {
	16	+	t.Fatalf("first chunk runes = %d, want 4000", got)
	17	+	}
	18	+	if got := utf8.RuneCountInString(chunks[1]); got != 1 {
	19	+	t.Fatalf("second chunk runes = %d, want 1", got)
	20	+	}
	21	+	}
	22	+
	23	+	func TestChunkMessageDoesNotSplitMultiByteRune(t *testing.T) {
	24	+	chunks := chunkMessage("🙂🙂🙂", 2)
	25	+	if len(chunks) != 2 {
	26	+	t.Fatalf("len(chunks) = %d, want 2", len(chunks))
	27	+	}
	28	+	if chunks[0] != "🙂🙂" \|\| chunks[1] != "🙂" {
	29	+	t.Fatalf("chunks = %#v", chunks)
	30	+	}
	31	+	}

@@ -0,0 +1,31 @@
	1	+	package telegram
	2	+
	3	+	import (
	4	+	"strings"
	5	+	"testing"
	6	+	"unicode/utf8"
	7	+	)
	8	+
	9	+	func TestChunkMessageUsesRuneLimit(t *testing.T) {
	10	+	s := strings.Repeat("界", 4001)
	11	+	chunks := chunkMessage(s, 4000)
	12	+	if len(chunks) != 2 {
	13	+	t.Fatalf("len(chunks) = %d, want 2", len(chunks))
	14	+	}
	15	+	if got := utf8.RuneCountInString(chunks[0]); got != 4000 {
	16	+	t.Fatalf("first chunk runes = %d, want 4000", got)
	17	+	}
	18	+	if got := utf8.RuneCountInString(chunks[1]); got != 1 {
	19	+	t.Fatalf("second chunk runes = %d, want 1", got)
	20	+	}
	21	+	}
	22	+
	23	+	func TestChunkMessageDoesNotSplitMultiByteRune(t *testing.T) {
	24	+	chunks := chunkMessage("🙂🙂🙂", 2)
	25	+	if len(chunks) != 2 {
	26	+	t.Fatalf("len(chunks) = %d, want 2", len(chunks))
	27	+	}
	28	+	if chunks[0] != "🙂🙂" \|\| chunks[1] != "🙂" {
	29	+	t.Fatalf("chunks = %#v", chunks)
	30	+	}
	31	+	}