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