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