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