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