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