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) Submit(string) {} // ignored in rpc mode
146func (h *rpcExtHooks) SubmitSlash(string) {} // ignored in rpc mode
147func (h *rpcExtHooks) Insert(string) {} // ignored in rpc mode
148func (h *rpcExtHooks) OpenPanel(string, extproto.PanelSpec) {}
149func (h *rpcExtHooks) UpdatePanel(string, string, string, []string, string) {}
150func (h *rpcExtHooks) ClosePanel(string, string) {}
151
152type rpcServer struct {
153 ctx context.Context
154 args Args
155 agent *core.Agent
156 provider string
157 model string
158 out io.Writer
159 version string
160
161 writeMu sync.Mutex
162 turnMu sync.Mutex // serialises one prompt at a time
163 activeCancel context.CancelFunc
164 authed bool
165
166 // inFlight tracks long-running command goroutines so run() can
167 // wait for them before returning when stdin closes. Without this,
168 // piping a single 'prompt' command into 'zot rpc' would race the
169 // process exit against the agent loop and the prompt would never
170 // produce output.
171 inFlight sync.WaitGroup
172}
173
174// run reads NDJSON commands from in and dispatches them. Returns when
175// in is closed AND every in-flight long-running command (prompt /
176// compact) has finished, so a quick `echo cmd | zot rpc` invocation
177// still produces full output before the process exits.
178func (s *rpcServer) run(in io.Reader) error {
179 requireToken := os.Getenv("ZOTCORE_RPC_TOKEN") != ""
180 s.authed = !requireToken
181
182 sc := bufio.NewScanner(in)
183 sc.Buffer(make([]byte, 0, 64*1024), 16*1024*1024)
184 for sc.Scan() {
185 line := strings.TrimSpace(sc.Text())
186 if line == "" {
187 continue
188 }
189 var head struct {
190 Type string `json:"type"`
191 ID string `json:"id,omitempty"`
192 }
193 if err := json.Unmarshal([]byte(line), &head); err != nil {
194 s.writeError("", "", fmt.Sprintf("malformed json: %v", err))
195 continue
196 }
197 if !s.authed {
198 if head.Type != "hello" {
199 s.writeError(head.ID, head.Type, "auth required: send hello with token first")
200 continue
201 }
202 var hello struct {
203 Token string `json:"token"`
204 }
205 _ = json.Unmarshal([]byte(line), &hello)
206 if hello.Token != os.Getenv("ZOTCORE_RPC_TOKEN") {
207 s.writeError(head.ID, head.Type, "invalid token")
208 return fmt.Errorf("rpc: bad auth token")
209 }
210 s.authed = true
211 s.writeResponse(head.ID, head.Type, map[string]any{
212 "protocol_version": 1,
213 "version": s.version,
214 "provider": s.provider,
215 "model": s.model,
216 })
217 continue
218 }
219 s.dispatch(head.Type, head.ID, []byte(line))
220 }
221 err := sc.Err()
222 s.inFlight.Wait()
223 return err
224}
225
226// dispatch routes a command. Long-running commands (prompt, compact)
227// run on their own goroutine so the read loop stays responsive.
228func (s *rpcServer) dispatch(cmd, id string, raw []byte) {
229 switch cmd {
230 case "hello":
231 s.writeResponse(id, cmd, map[string]any{
232 "protocol_version": 1,
233 "version": s.version,
234 "provider": s.provider,
235 "model": s.model,
236 })
237 case "prompt":
238 var req struct {
239 Message string `json:"message"`
240 Images []struct {
241 MimeType string `json:"mime_type"`
242 Data []byte `json:"data"`
243 } `json:"images"`
244 }
245 if err := json.Unmarshal(raw, &req); err != nil {
246 s.writeError(id, cmd, err.Error())
247 return
248 }
249 s.inFlight.Add(1)
250 go func() {
251 defer s.inFlight.Done()
252 s.runPrompt(id, req.Message, req.Images)
253 }()
254
255 case "abort":
256 if c := s.takeCancel(); c != nil {
257 c()
258 }
259 s.writeResponse(id, cmd, nil)
260
261 case "compact":
262 s.inFlight.Add(1)
263 go func() {
264 defer s.inFlight.Done()
265 s.runCompact(id)
266 }()
267
268 case "get_state":
269 s.writeResponse(id, cmd, s.snapshotState())
270
271 case "get_messages":
272 s.writeResponse(id, cmd, map[string]any{
273 "messages": messagesToJSON(s.agent.Messages()),
274 })
275
276 case "clear":
277 s.agent.SetMessages(nil)
278 s.writeResponse(id, cmd, nil)
279
280 case "set_model":
281 var req struct {
282 Model string `json:"model"`
283 }
284 if err := json.Unmarshal(raw, &req); err != nil {
285 s.writeError(id, cmd, err.Error())
286 return
287 }
288 if _, err := provider.FindModel(s.provider, req.Model); err != nil {
289 s.writeError(id, cmd, err.Error())
290 return
291 }
292 s.agent.Model = req.Model
293 s.model = req.Model
294 s.writeResponse(id, cmd, map[string]any{"model": req.Model})
295
296 case "get_models":
297 out := []map[string]any{}
298 for _, m := range provider.ModelsForProvider(s.provider) {
299 out = append(out, map[string]any{
300 "id": m.ID,
301 "provider": m.Provider,
302 "context_window": m.ContextWindow,
303 "max_output": m.MaxOutput,
304 "reasoning": m.Reasoning,
305 })
306 }
307 s.writeResponse(id, cmd, map[string]any{"models": out})
308
309 case "ping":
310 s.writeResponse(id, cmd, map[string]any{"pong": true})
311
312 default:
313 s.writeError(id, cmd, "unknown command")
314 }
315}
316
317// runPrompt executes a single prompt turn and streams events out.
318// Holds turnMu so a second concurrent prompt blocks until this one
319// finishes; the user can abort with the abort command.
320func (s *rpcServer) runPrompt(id, message string, images []struct {
321 MimeType string `json:"mime_type"`
322 Data []byte `json:"data"`
323}) {
324 s.turnMu.Lock()
325 defer s.turnMu.Unlock()
326
327 subCtx, cancel := context.WithCancel(s.ctx)
328 s.setCancel(cancel)
329 defer s.setCancel(nil)
330
331 s.writeResponse(id, "prompt", map[string]any{"started": true})
332
333 imgs := make([]provider.ImageBlock, 0, len(images))
334 for _, im := range images {
335 imgs = append(imgs, provider.ImageBlock{MimeType: im.MimeType, Data: im.Data})
336 }
337
338 err := s.agent.Prompt(subCtx, message, imgs, func(ev core.AgentEvent) {
339 // EvDone is emitted by the agent loop and we re-emit our own
340 // 'done' below; suppressing it here avoids duplicate frames.
341 if _, ok := ev.(core.EvDone); ok {
342 return
343 }
344 s.writeEvent(modes.EventToJSON(ev))
345 })
346 // Don't emit a stand-alone error event for cancellation; the prior
347 // turn_end with stop=aborted already carries that signal.
348 if err != nil && !errors.Is(err, context.Canceled) {
349 s.writeEvent(map[string]any{"type": "error", "message": err.Error()})
350 }
351 s.writeEvent(map[string]any{"type": "done"})
352}
353
354// runCompact mirrors runPrompt but for compaction.
355func (s *rpcServer) runCompact(id string) {
356 s.turnMu.Lock()
357 defer s.turnMu.Unlock()
358
359 subCtx, cancel := context.WithCancel(s.ctx)
360 s.setCancel(cancel)
361 defer s.setCancel(nil)
362
363 s.writeResponse(id, "compact", map[string]any{"started": true})
364 summary, err := s.agent.Compact(subCtx, 4, nil)
365 if err != nil {
366 if !errors.Is(err, context.Canceled) {
367 s.writeEvent(map[string]any{"type": "error", "message": err.Error()})
368 }
369 return
370 }
371 s.writeEvent(map[string]any{
372 "type": "compact_done",
373 "summary": summary,
374 })
375}
376
377// snapshotState builds the get_state response.
378func (s *rpcServer) snapshotState() map[string]any {
379 cum := s.agent.Cost()
380 return map[string]any{
381 "provider": s.provider,
382 "model": s.model,
383 "cwd": s.args.CWD,
384 "message_count": len(s.agent.Messages()),
385 "busy": s.activeCancel != nil,
386 "usage": map[string]any{
387 "input": cum.InputTokens,
388 "output": cum.OutputTokens,
389 "cache_read": cum.CacheReadTokens,
390 "cache_write": cum.CacheWriteTokens,
391 "cost_usd": cum.CostUSD,
392 },
393 }
394}
395
396// ---- write helpers (single-line JSON, mutex-guarded) ----
397
398func (s *rpcServer) writeResponse(id, cmd string, data any) {
399 frame := map[string]any{
400 "type": "response",
401 "command": cmd,
402 "success": true,
403 }
404 if id != "" {
405 frame["id"] = id
406 }
407 if data != nil {
408 frame["data"] = data
409 }
410 s.write(frame)
411}
412
413func (s *rpcServer) writeError(id, cmd, msg string) {
414 frame := map[string]any{
415 "type": "response",
416 "command": cmd,
417 "success": false,
418 "error": msg,
419 }
420 if id != "" {
421 frame["id"] = id
422 }
423 s.write(frame)
424}
425
426func (s *rpcServer) writeEvent(payload map[string]any) {
427 s.write(payload)
428}
429
430func (s *rpcServer) write(v any) {
431 b, err := json.Marshal(v)
432 if err != nil {
433 return
434 }
435 s.writeMu.Lock()
436 defer s.writeMu.Unlock()
437 _, _ = s.out.Write(b)
438 _, _ = s.out.Write([]byte("\n"))
439}
440
441func (s *rpcServer) setCancel(c context.CancelFunc) {
442 s.writeMu.Lock()
443 s.activeCancel = c
444 s.writeMu.Unlock()
445}
446
447func (s *rpcServer) takeCancel() context.CancelFunc {
448 s.writeMu.Lock()
449 defer s.writeMu.Unlock()
450 c := s.activeCancel
451 s.activeCancel = nil
452 return c
453}
454
455// messagesToJSON serialises a transcript using the same schema as the
456// --json event mode for cross-format consistency.
457func messagesToJSON(msgs []provider.Message) []map[string]any {
458 out := make([]map[string]any, 0, len(msgs))
459 for _, m := range msgs {
460 out = append(out, map[string]any{
461 "role": string(m.Role),
462 "content": modes.ContentToJSON(m.Content),
463 "time": m.Time,
464 })
465 }
466 return out
467}