1package agent
2
3import (
4 "bufio"
5 "context"
6 "encoding/json"
7 "errors"
8 "fmt"
9 "io"
10 "os"
11 "strings"
12 "sync"
13 "time"
14
15 "github.com/patriceckhart/zot/packages/agent/extensions"
16 "github.com/patriceckhart/zot/packages/agent/extproto"
17 "github.com/patriceckhart/zot/packages/agent/modes"
18 "github.com/patriceckhart/zot/packages/core"
19 "github.com/patriceckhart/zot/packages/provider"
20)
21
22// runRPCMode implements the JSON-over-stdin/stdout RPC protocol.
23//
24// Wire format: one JSON object per line in both directions.
25//
26// Commands (stdin):
27//
28// {"id":"1","type":"prompt","message":"hello","images":[]}
29// {"id":"2","type":"abort"}
30// {"id":"3","type":"compact"}
31// {"id":"4","type":"get_state"}
32// {"id":"5","type":"set_model","model":"claude-opus-4-5"}
33// {"id":"6","type":"set_reasoning","reasoning":"max"}
34// {"id":"6","type":"get_messages"}
35// {"id":"7","type":"clear"}
36// {"id":"8","type":"get_models"}
37//
38// Responses (stdout): {"type":"response","id":"1","command":"prompt","success":true}
39// Events (stdout): one JSON object per AgentEvent (same schema as --json mode).
40//
41// Auth: if $ZOTCORE_RPC_TOKEN is set, the first command must be
42// {"type":"hello","token":"..."} or the connection is closed.
43func runRPCMode(ctx context.Context, args Args, version string) error {
44 if args.NoYolo {
45 fmt.Fprintln(os.Stderr, "warning: --no-yolo has no effect in rpc mode (no interactive prompt available); tools will run without confirmation")
46 }
47 r, err := Resolve(args, true)
48 if err != nil {
49 return err
50 }
51 reportSkillDiagnostics(os.Stderr, r.SkillDiagnostics)
52
53 // Extensions: same lifecycle as interactive mode, minus the
54 // host-hooks integration. Notify/Display calls from extensions
55 // emit RPC events instead of TUI lines so any consumer can react.
56 extHooks := &rpcExtHooks{}
57 extMgr := extensions.New(ZotHome(), r.CWD, version, r.Provider, r.Model, extHooks)
58 for _, e := range extMgr.LoadExplicit(ctx, args.Exts) {
59 fmt.Fprintln(os.Stderr, "extension load:", e)
60 }
61 if !args.NoExt {
62 for _, e := range extMgr.Discover(ctx) {
63 fmt.Fprintln(os.Stderr, "extension load:", e)
64 }
65 }
66 extMgr.WaitForReady(3 * time.Second)
67 defer extMgr.Stop(2 * time.Second)
68 r.MergeExtensionTools(&extToolAdapter{mgr: extMgr})
69
70 ag := r.NewAgent()
71 wireNonInteractiveAgentExtHooks(ctx, ag, extMgr)
72
73 // /reload-ext hot-reload callback (also triggered via rpc
74 // `reload_ext` if/when added). Rebuilds the tool registry on the
75 // current agent so freshly-registered extension tools become
76 // callable without restarting the rpc process.
77 adapter := &extToolAdapter{mgr: extMgr}
78 extMgr.SetOnReload(func() {
79 resolved, err := Resolve(args, true)
80 if err != nil {
81 return
82 }
83 resolved.MergeExtensionTools(adapter)
84 ag.SetTools(resolved.ToolRegistry)
85 })
86
87 startExtensionSession(extMgr, ag, r.CWD, "")
88
89 server := &rpcServer{
90 ctx: ctx,
91 args: args,
92 agent: ag,
93 provider: r.Provider,
94 model: r.Model,
95 out: os.Stdout,
96 version: version,
97 }
98 extHooks.server = server
99 return server.run(os.Stdin)
100}
101
102// rpcExtHooks implements extensions.HostHooks for the headless RPC
103// loop. Notify and Display surface as `event` frames so any RPC
104// client can render them; Submit and Insert are no-ops because the
105// RPC loop has no editor and the prompt comes from the client.
106type rpcExtHooks struct {
107 server *rpcServer
108}
109
110func (h *rpcExtHooks) Notify(extName, level, message string) {
111 if h.server != nil {
112 h.server.writeEvent(map[string]any{
113 "type": "ext_notify",
114 "extension": extName,
115 "level": level,
116 "message": message,
117 })
118 }
119}
120func (h *rpcExtHooks) Display(extName, text string) {
121 if h.server != nil {
122 h.server.writeEvent(map[string]any{
123 "type": "ext_display",
124 "extension": extName,
125 "text": text,
126 })
127 }
128}
129func (h *rpcExtHooks) ClearNotes(extName string) {
130 if h.server != nil {
131 h.server.writeEvent(map[string]any{
132 "type": "ext_clear_notes",
133 "extension": extName,
134 })
135 }
136}
137func (h *rpcExtHooks) Submit(string) {} // ignored in rpc mode
138func (h *rpcExtHooks) SubmitSlash(string) {} // ignored in rpc mode
139func (h *rpcExtHooks) Insert(string) {} // ignored in rpc mode
140func (h *rpcExtHooks) OpenPanel(string, extproto.PanelSpec) {}
141func (h *rpcExtHooks) UpdatePanel(string, string, string, []string, string) {}
142func (h *rpcExtHooks) ClosePanel(string, string) {}
143
144type rpcServer struct {
145 ctx context.Context
146 args Args
147 agent *core.Agent
148 provider string
149 model string
150 out io.Writer
151 version string
152
153 writeMu sync.Mutex
154 turnMu sync.Mutex // serialises one prompt at a time
155 activeCancel context.CancelFunc
156 authed bool
157
158 // inFlight tracks long-running command goroutines so run() can
159 // wait for them before returning when stdin closes. Without this,
160 // piping a single 'prompt' command into 'zot rpc' would race the
161 // process exit against the agent loop and the prompt would never
162 // produce output.
163 inFlight sync.WaitGroup
164}
165
166// run reads NDJSON commands from in and dispatches them. Returns when
167// in is closed AND every in-flight long-running command (prompt /
168// compact) has finished, so a quick `echo cmd | zot rpc` invocation
169// still produces full output before the process exits.
170func (s *rpcServer) run(in io.Reader) error {
171 requireToken := os.Getenv("ZOTCORE_RPC_TOKEN") != ""
172 s.authed = !requireToken
173
174 sc := bufio.NewScanner(in)
175 sc.Buffer(make([]byte, 0, 64*1024), 16*1024*1024)
176 for sc.Scan() {
177 line := strings.TrimSpace(sc.Text())
178 if line == "" {
179 continue
180 }
181 var head struct {
182 Type string `json:"type"`
183 ID string `json:"id,omitempty"`
184 }
185 if err := json.Unmarshal([]byte(line), &head); err != nil {
186 s.writeError("", "", fmt.Sprintf("malformed json: %v", err))
187 continue
188 }
189 if !s.authed {
190 if head.Type != "hello" {
191 s.writeError(head.ID, head.Type, "auth required: send hello with token first")
192 continue
193 }
194 var hello struct {
195 Token string `json:"token"`
196 }
197 _ = json.Unmarshal([]byte(line), &hello)
198 if hello.Token != os.Getenv("ZOTCORE_RPC_TOKEN") {
199 s.writeError(head.ID, head.Type, "invalid token")
200 return fmt.Errorf("rpc: bad auth token")
201 }
202 s.authed = true
203 s.writeResponse(head.ID, head.Type, map[string]any{
204 "protocol_version": 1,
205 "version": s.version,
206 "provider": s.provider,
207 "model": s.model,
208 })
209 continue
210 }
211 s.dispatch(head.Type, head.ID, []byte(line))
212 }
213 err := sc.Err()
214 s.inFlight.Wait()
215 return err
216}
217
218// dispatch routes a command. Long-running commands (prompt, compact)
219// run on their own goroutine so the read loop stays responsive.
220func (s *rpcServer) dispatch(cmd, id string, raw []byte) {
221 // Runtime mutations must not race with prompt preparation, model/tool
222 // execution, or compaction. Do not wait here: this is also the read loop
223 // that must keep accepting abort while an extension or provider is busy.
224 switch cmd {
225 case "clear", "set_model", "set_reasoning":
226 if !s.turnMu.TryLock() {
227 s.writeError(id, cmd, "agent is busy; wait for completion or abort before retrying")
228 return
229 }
230 defer s.turnMu.Unlock()
231 }
232
233 switch cmd {
234 case "hello":
235 s.writeResponse(id, cmd, map[string]any{
236 "protocol_version": 1,
237 "version": s.version,
238 "provider": s.provider,
239 "model": s.model,
240 })
241 case "prompt":
242 var req struct {
243 Message string `json:"message"`
244 Images []struct {
245 MimeType string `json:"mime_type"`
246 Data []byte `json:"data"`
247 } `json:"images"`
248 }
249 if err := json.Unmarshal(raw, &req); err != nil {
250 s.writeError(id, cmd, err.Error())
251 return
252 }
253 s.inFlight.Add(1)
254 go func() {
255 defer s.inFlight.Done()
256 s.runPrompt(id, req.Message, req.Images)
257 }()
258
259 case "abort":
260 if c := s.takeCancel(); c != nil {
261 c()
262 }
263 s.writeResponse(id, cmd, nil)
264
265 case "compact":
266 s.inFlight.Add(1)
267 go func() {
268 defer s.inFlight.Done()
269 s.runCompact(id)
270 }()
271
272 case "get_state":
273 s.writeResponse(id, cmd, s.snapshotState())
274
275 case "get_messages":
276 s.writeResponse(id, cmd, map[string]any{
277 "messages": messagesToJSON(s.agent.Messages()),
278 })
279
280 case "clear":
281 s.agent.SetMessages(nil)
282 s.writeResponse(id, cmd, nil)
283
284 case "set_model":
285 var req struct {
286 Model string `json:"model"`
287 }
288 if err := json.Unmarshal(raw, &req); err != nil {
289 s.writeError(id, cmd, err.Error())
290 return
291 }
292 m, err := provider.FindModel(s.provider, req.Model)
293 if err != nil {
294 s.writeError(id, cmd, err.Error())
295 return
296 }
297 s.agent.Model = req.Model
298 s.agent.MaxTokens = m.MaxOutput
299 s.model = req.Model
300 s.writeResponse(id, cmd, map[string]any{"model": req.Model})
301
302 case "set_reasoning":
303 var req struct {
304 Reasoning string `json:"reasoning"`
305 }
306 if err := json.Unmarshal(raw, &req); err != nil {
307 s.writeError(id, cmd, err.Error())
308 return
309 }
310 reasoning := provider.NormalizeReasoning(req.Reasoning)
311 switch reasoning {
312 case "", "minimum", "low", "medium", "high", "xhigh", "max":
313 default:
314 s.writeError(id, cmd, "reasoning must be off|minimum|low|medium|high|xhigh|max")
315 return
316 }
317 s.agent.Reasoning = reasoning
318 s.writeResponse(id, cmd, map[string]any{"reasoning": reasoning})
319
320 case "get_models":
321 out := []map[string]any{}
322 for _, m := range provider.ModelsForProvider(s.provider) {
323 out = append(out, map[string]any{
324 "id": m.ID,
325 "provider": m.Provider,
326 "context_window": m.ContextWindow,
327 "max_output": m.MaxOutput,
328 "reasoning": m.Reasoning,
329 })
330 }
331 s.writeResponse(id, cmd, map[string]any{"models": out})
332
333 case "ping":
334 s.writeResponse(id, cmd, map[string]any{"pong": true})
335
336 default:
337 s.writeError(id, cmd, "unknown command")
338 }
339}
340
341// runPrompt executes a single prompt turn and streams events out.
342// Holds turnMu so a second concurrent prompt blocks until this one
343// finishes; the user can abort with the abort command.
344func (s *rpcServer) runPrompt(id, message string, images []struct {
345 MimeType string `json:"mime_type"`
346 Data []byte `json:"data"`
347}) {
348 s.turnMu.Lock()
349 defer s.turnMu.Unlock()
350
351 subCtx, cancel := context.WithCancel(s.ctx)
352 s.setCancel(cancel)
353 defer s.setCancel(nil)
354
355 s.writeResponse(id, "prompt", map[string]any{"started": true})
356
357 imgs := make([]provider.ImageBlock, 0, len(images))
358 for _, im := range images {
359 imgs = append(imgs, provider.ImageBlock{MimeType: im.MimeType, Data: im.Data})
360 }
361
362 err := s.agent.Prompt(subCtx, message, imgs, func(ev core.AgentEvent) {
363 // EvDone is emitted by the agent loop and we re-emit our own
364 // 'done' below; suppressing it here avoids duplicate frames.
365 if _, ok := ev.(core.EvDone); ok {
366 return
367 }
368 s.writeEvent(modes.EventToJSON(ev))
369 })
370 // Don't emit a stand-alone error event for cancellation; the prior
371 // turn_end with stop=aborted already carries that signal.
372 if err != nil && !errors.Is(err, context.Canceled) {
373 s.writeEvent(map[string]any{"type": "error", "message": err.Error()})
374 }
375 s.writeEvent(map[string]any{"type": "done"})
376}
377
378// runCompact mirrors runPrompt but for compaction.
379func (s *rpcServer) runCompact(id string) {
380 s.turnMu.Lock()
381 defer s.turnMu.Unlock()
382
383 subCtx, cancel := context.WithCancel(s.ctx)
384 s.setCancel(cancel)
385 defer s.setCancel(nil)
386
387 s.writeResponse(id, "compact", map[string]any{"started": true})
388 summary, err := s.agent.Compact(subCtx, 4, nil)
389 if err != nil {
390 if !errors.Is(err, context.Canceled) {
391 s.writeEvent(map[string]any{"type": "error", "message": err.Error()})
392 }
393 return
394 }
395 s.writeEvent(map[string]any{
396 "type": "compact_done",
397 "summary": summary,
398 })
399}
400
401// snapshotState builds the get_state response.
402func (s *rpcServer) snapshotState() map[string]any {
403 cum := s.agent.Cost()
404 return map[string]any{
405 "provider": s.provider,
406 "model": s.model,
407 "reasoning": s.agent.Reasoning,
408 "cwd": s.args.CWD,
409 "message_count": len(s.agent.Messages()),
410 "busy": s.activeCancel != nil,
411 "usage": map[string]any{
412 "input": cum.InputTokens,
413 "output": cum.OutputTokens,
414 "reasoning": rpcReasoningTokens(cum),
415 "cache_read": cum.CacheReadTokens,
416 "cache_write": cum.CacheWriteTokens,
417 "cost_usd": cum.CostUSD,
418 },
419 }
420}
421
422func rpcReasoningTokens(usage provider.Usage) any {
423 if !usage.ReasoningTokensKnown {
424 return nil
425 }
426 return usage.ReasoningTokens
427}
428
429// ---- write helpers (single-line JSON, mutex-guarded) ----
430
431func (s *rpcServer) writeResponse(id, cmd string, data any) {
432 frame := map[string]any{
433 "type": "response",
434 "command": cmd,
435 "success": true,
436 }
437 if id != "" {
438 frame["id"] = id
439 }
440 if data != nil {
441 frame["data"] = data
442 }
443 s.write(frame)
444}
445
446func (s *rpcServer) writeError(id, cmd, msg string) {
447 frame := map[string]any{
448 "type": "response",
449 "command": cmd,
450 "success": false,
451 "error": msg,
452 }
453 if id != "" {
454 frame["id"] = id
455 }
456 s.write(frame)
457}
458
459func (s *rpcServer) writeEvent(payload map[string]any) {
460 s.write(payload)
461}
462
463func (s *rpcServer) write(v any) {
464 b, err := json.Marshal(v)
465 if err != nil {
466 return
467 }
468 s.writeMu.Lock()
469 defer s.writeMu.Unlock()
470 _, _ = s.out.Write(b)
471 _, _ = s.out.Write([]byte("\n"))
472}
473
474func (s *rpcServer) setCancel(c context.CancelFunc) {
475 s.writeMu.Lock()
476 s.activeCancel = c
477 s.writeMu.Unlock()
478}
479
480func (s *rpcServer) takeCancel() context.CancelFunc {
481 s.writeMu.Lock()
482 defer s.writeMu.Unlock()
483 c := s.activeCancel
484 s.activeCancel = nil
485 return c
486}
487
488// messagesToJSON serialises a transcript using the same schema as the
489// --json event mode for cross-format consistency.
490func messagesToJSON(msgs []provider.Message) []map[string]any {
491 out := make([]map[string]any, 0, len(msgs))
492 for _, m := range msgs {
493 out = append(out, map[string]any{
494 "role": string(m.Role),
495 "content": modes.ContentToJSON(m.Content),
496 "time": m.Time,
497 })
498 }
499 return out
500}