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