1package bot
2
3import (
4 "context"
5 "fmt"
6 "io"
7 "os"
8 "strings"
9 "sync"
10
11 "github.com/patriceckhart/zot/packages/core"
12 "github.com/patriceckhart/zot/packages/provider"
13)
14
15// stderr is a tiny hook so tests can redirect bot logging.
16var stderr = func() io.Writer { return os.Stderr }
17
18// Config holds runner-level settings that are protocol-independent.
19type Config struct {
20 ZotHome string
21 Provider string
22 Model string
23 AuthMethod string
24 CWD string
25 RefreshCreds func() error
26}
27
28// queuedTurn is an inbound message waiting to become a prompt.
29type queuedTurn struct {
30 channelID string
31 messageID string
32 prompt string
33 images []provider.ImageBlock
34}
35
36// Runner is the protocol-agnostic bot engine. It owns the turn queue,
37// dispatches prompts to the agent, and streams replies back through
38// the BotAdapter.
39type Runner struct {
40 agent *core.Agent
41 adapter BotAdapter
42 cfg Config
43
44 mu sync.Mutex
45 busy bool
46 activeCtx context.CancelFunc
47 queue []queuedTurn
48 lastCtxInput int
49 runCtx context.Context // set at Run entry; used by goroutines
50}
51
52// NewRunner creates a Runner wired to the given adapter and agent.
53func NewRunner(adapter BotAdapter, agent *core.Agent, cfg Config) *Runner {
54 return &Runner{
55 agent: agent,
56 adapter: adapter,
57 cfg: cfg,
58 }
59}
60
61// UpdateRuntimeConfig updates provider/model/auth/cwd at runtime (e.g. after
62// credential refresh). This is thread-safe.
63func (r *Runner) UpdateRuntimeConfig(provider, model, authMethod, cwd string) {
64 r.mu.Lock()
65 defer r.mu.Unlock()
66 r.cfg.Provider = provider
67 r.cfg.Model = model
68 r.cfg.AuthMethod = authMethod
69 r.cfg.CWD = cwd
70}
71
72// Run starts the adapter's polling loop and blocks until ctx cancels.
73func (r *Runner) Run(ctx context.Context) error {
74 r.mu.Lock()
75 r.runCtx = ctx
76 r.mu.Unlock()
77
78 return r.adapter.Run(ctx, r.handleMessage, r.handleCommand)
79}
80
81// handleMessage is called by the adapter for every normal inbound message.
82func (r *Runner) handleMessage(msg InboundMessage) {
83 if msg.Text == "" && len(msg.Images) == 0 {
84 return
85 }
86
87 r.mu.Lock()
88 r.queue = append(r.queue, queuedTurn{
89 channelID: msg.ChannelID,
90 messageID: msg.MessageID,
91 prompt: msg.Text,
92 images: msg.Images,
93 })
94 idle := !r.busy
95 if idle {
96 r.busy = true
97 }
98 r.mu.Unlock()
99
100 if idle {
101 go r.drainQueue()
102 }
103}
104
105// handleCommand is called by the adapter for built-in commands.
106func (r *Runner) handleCommand(cmd Command, msg InboundMessage) {
107 switch cmd {
108 case CmdStart, CmdHelp:
109 _ = r.adapter.Send(context.Background(), msg.ChannelID,
110 "send me any message and i'll forward it to zot. attach an image and i'll pass it to the model. commands: /status, /stop, or plain stop.",
111 SendOptions{ReplyToMessageID: msg.MessageID})
112 case CmdStatus:
113 r.sendStatus(msg.ChannelID, msg.MessageID)
114 case CmdStop:
115 r.cancelActiveTurn(msg.ChannelID, msg.MessageID)
116 }
117}
118
119// drainQueue runs queued turns one at a time until the queue is empty.
120func (r *Runner) drainQueue() {
121 r.mu.Lock()
122 parent := r.runCtx
123 r.mu.Unlock()
124
125 for {
126 r.mu.Lock()
127 if len(r.queue) == 0 {
128 r.busy = false
129 r.activeCtx = nil
130 r.mu.Unlock()
131 return
132 }
133 t := r.queue[0]
134 r.queue = r.queue[1:]
135 turnCtx, cancel := context.WithCancel(parent)
136 r.activeCtx = cancel
137 r.mu.Unlock()
138
139 if r.cfg.RefreshCreds != nil {
140 if err := r.cfg.RefreshCreds(); err != nil {
141 fmt.Fprintln(stderr(), "bot: refresh creds:", err)
142 }
143 }
144 r.runTurn(turnCtx, t)
145 cancel()
146 }
147}
148
149// runTurn sends the queued prompt to the agent and streams the reply.
150func (r *Runner) runTurn(ctx context.Context, t queuedTurn) {
151 stopWorking := r.adapter.IndicateWorking(ctx, t.channelID)
152 defer stopWorking()
153
154 var replyBuilder strings.Builder
155 var lastAssistantText string
156 var turnErr error
157
158 sink := func(ev core.AgentEvent) {
159 switch e := ev.(type) {
160 case core.EvTextDelta:
161 replyBuilder.WriteString(e.Delta)
162 case core.EvUsage:
163 r.mu.Lock()
164 if e.Usage.InputTokens > 0 {
165 r.lastCtxInput = e.Usage.InputTokens + e.Usage.CacheReadTokens + e.Usage.CacheWriteTokens
166 }
167 r.mu.Unlock()
168 case core.EvAssistantMessage:
169 var sb strings.Builder
170 for _, c := range e.Message.Content {
171 if tb, ok := c.(provider.TextBlock); ok {
172 if sb.Len() > 0 {
173 sb.WriteString("\n")
174 }
175 sb.WriteString(tb.Text)
176 }
177 }
178 if sb.Len() > 0 {
179 lastAssistantText = sb.String()
180 }
181 replyBuilder.Reset()
182 case core.EvTurnEnd:
183 if e.Err != nil {
184 turnErr = e.Err
185 }
186 }
187 }
188
189 if err := r.agent.Prompt(ctx, t.prompt, t.images, sink); err != nil {
190 turnErr = err
191 }
192
193 reply := strings.TrimSpace(lastAssistantText)
194 if reply == "" {
195 reply = strings.TrimSpace(replyBuilder.String())
196 }
197 if turnErr != nil && ctx.Err() == nil {
198 reply = "error: " + turnErr.Error()
199 }
200 if reply == "" {
201 reply = "(no reply)"
202 }
203
204 // Adapter.Send is responsible for chunking to protocol limits.
205 if err := r.adapter.Send(context.Background(), t.channelID, reply, SendOptions{}); err != nil {
206 fmt.Fprintln(stderr(), "bot: send reply:", err)
207 }
208}
209
210// cancelActiveTurn aborts the currently running turn, if any.
211func (r *Runner) cancelActiveTurn(channelID, messageID string) {
212 r.mu.Lock()
213 cancel := r.activeCtx
214 r.mu.Unlock()
215 opts := SendOptions{ReplyToMessageID: messageID}
216 if cancel != nil {
217 cancel()
218 _ = r.adapter.Send(context.Background(), channelID, "cancelled the current turn.", opts)
219 } else {
220 _ = r.adapter.Send(context.Background(), channelID, "nothing running.", opts)
221 }
222}
223
224// sendStatus describes agent state to the user.
225func (r *Runner) sendStatus(channelID, messageID string) {
226 r.mu.Lock()
227 busy := r.busy
228 queued := len(r.queue)
229 ctxUsed := r.lastCtxInput
230 providerName := r.cfg.Provider
231 model := r.cfg.Model
232 authMethod := r.cfg.AuthMethod
233 cwd := r.cfg.CWD
234 r.mu.Unlock()
235
236 ctxMax := 0
237 if m, err := provider.FindModel(providerName, model); err == nil {
238 ctxMax = m.ContextWindow
239 }
240
241 status := FormatStatus(StatusSnapshot{
242 Provider: providerName,
243 Model: model,
244 CWD: cwd,
245 Usage: r.agent.Cost(),
246 Subscription: authMethod == "oauth",
247 ContextUsed: ctxUsed,
248 ContextMax: ctxMax,
249 Busy: busy,
250 Queued: queued,
251 })
252
253 if extra := r.adapter.StatusText(); extra != "" {
254 status += "\n" + extra
255 }
256
257 _ = r.adapter.Send(context.Background(), channelID, status, SendOptions{ReplyToMessageID: messageID})
258}