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