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