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