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