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/bnema/zut/packages/agent/extensions"
16 "github.com/bnema/zut/packages/agent/extproto"
17 "github.com/bnema/zut/packages/agent/modes"
18 "github.com/bnema/zut/packages/agent/tools"
19 "github.com/bnema/zut/packages/core"
20 "github.com/bnema/zut/packages/provider"
21)
22
23// runRPCMode implements the JSON-over-stdin/stdout RPC protocol.
24//
25// Wire format: one JSON object per line in both directions.
26//
27// Commands (stdin):
28//
29// {"id":"1","type":"prompt","message":"hello","images":[]}
30// {"id":"2","type":"abort"}
31// {"id":"3","type":"compact"}
32// {"id":"4","type":"get_state"}
33// {"id":"5","type":"set_model","model":"claude-opus-4-5"}
34// {"id":"6","type":"set_reasoning","reasoning":"max"}
35// {"id":"6","type":"get_messages"}
36// {"id":"7","type":"clear"}
37// {"id":"8","type":"get_models"}
38//
39// Responses (stdout): {"type":"response","id":"1","command":"prompt","success":true}
40// Events (stdout): one JSON object per AgentEvent (same schema as --json mode).
41// Extensions may additionally emit ext_status with extension/key/level/text,
42// ext_widget with extension/id/position/title/lines, and ext_widget_clear
43// with extension/id when loaded extensions publish persistent UI state.
44//
45// Auth: if $ZUTCORE_RPC_TOKEN is set, the first command must be
46// {"type":"hello","token":"..."} or the connection is closed.
47func runRPCMode(ctx context.Context, args Args, version string) error {
48 if args.NoYolo {
49 fmt.Fprintln(os.Stderr, "warning: --no-yolo has no effect in rpc mode (no interactive prompt available); tools will run without confirmation")
50 }
51 r, err := Resolve(args, true)
52 if err != nil {
53 return err
54 }
55
56 // Construct the RPC event sink before extension startup. Alerts can
57 // arrive as soon as an extension is ready; rpcServer queues them until
58 // the command loop has started and, when configured, authenticated.
59 server := &rpcServer{
60 ctx: ctx,
61 args: args,
62 provider: r.Provider,
63 model: r.Model,
64 out: os.Stdout,
65 version: version,
66 }
67
68 // Extensions: same lifecycle as interactive mode, minus the
69 // host-hooks integration. Notify/Display calls from extensions
70 // emit RPC events instead of TUI lines so any consumer can react.
71 extHooks := &rpcExtHooks{server: server}
72 extMgr := extensions.New(ZutHome(), r.CWD, version, r.Provider, r.Model, extHooks)
73 for _, e := range extMgr.LoadExplicit(ctx, args.Exts) {
74 fmt.Fprintln(os.Stderr, "extension load:", e)
75 }
76 if !args.NoExt {
77 for _, e := range extMgr.Discover(ctx) {
78 fmt.Fprintln(os.Stderr, "extension load:", e)
79 }
80 }
81 extMgr.WaitForReady(3 * time.Second)
82 defer extMgr.Stop(2 * time.Second)
83 r.MergeExtensionTools(&extToolAdapter{mgr: extMgr})
84
85 ag := r.NewAgent()
86 defer closeAgentLSP(ag)
87 server.agent = ag
88 ag.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
89 r := extMgr.InterceptToolCall(ctx, call.ID, call.Name, call.Arguments)
90 if r.Block {
91 return false, r.Reason, nil
92 }
93 return true, "", r.ModifiedArgs
94 }
95 ag.BeforeTurnContext = func(turnCtx context.Context, step int) (bool, string, string) {
96 r := extMgr.InterceptTurnStart(turnCtx, step)
97 return !r.Block, r.Reason, r.Context
98 }
99 ag.BeforeAssistantMessage = func(text string) (bool, string, string) {
100 r := extMgr.InterceptAssistantMessage(ctx, text)
101 if r.Block {
102 return false, r.Reason, ""
103 }
104 return true, "", r.ReplaceText
105 }
106 ag.OnEvent = func(ev core.AgentEvent) { fanoutAgentEvent(extMgr, ev) }
107
108 // /reload-ext hot-reload callback (also triggered via rpc
109 // `reload_ext` if/when added). Rebuilds the tool registry on the
110 // current agent so freshly-registered extension tools become
111 // callable without restarting the rpc process.
112 adapter := &extToolAdapter{mgr: extMgr}
113 extMgr.SetOnReload(func() {
114 resolved, err := Resolve(args, true)
115 if err != nil {
116 return
117 }
118 resolved.MergeExtensionTools(adapter)
119 oldTools := ag.SetPromptConfig(resolved.SystemPrompt, resolved.ToolRegistry)
120 if closeErr := tools.CloseLSPManagers(oldTools); closeErr != nil {
121 fmt.Fprintln(os.Stderr, "reload LSP cleanup:", closeErr)
122 }
123 })
124
125 extMgr.EmitSessionEvent("session_start", nil, nil)
126
127 return server.run(os.Stdin)
128}
129
130// rpcExtHooks implements extensions.HostHooks for the headless RPC
131// loop. Notify and Display surface as `event` frames so any RPC
132// client can render them; Submit and Insert are no-ops because the
133// RPC loop has no editor and the prompt comes from the client.
134type rpcExtHooks struct {
135 server *rpcServer
136}
137
138func (h *rpcExtHooks) Notify(extName, level, message string) {
139 if h.server != nil {
140 h.server.writeExtensionEvent(map[string]any{
141 "type": "ext_notify",
142 "extension": extName,
143 "level": level,
144 "message": message,
145 })
146 }
147}
148func (h *rpcExtHooks) Alert(extName string, alert extproto.AlertRequest) {
149 if h.server != nil {
150 h.server.writeExtensionEvent(map[string]any{
151 "type": "ext_alert",
152 "extension": extName,
153 "kind": alert.Kind,
154 "reason": alert.Reason,
155 })
156 }
157}
158func (h *rpcExtHooks) Display(extName, text string) {
159 if h.server != nil {
160 h.server.writeExtensionEvent(map[string]any{
161 "type": "ext_display",
162 "extension": extName,
163 "text": text,
164 })
165 }
166}
167func (h *rpcExtHooks) ClearNotes(extName string) {
168 if h.server != nil {
169 h.server.writeExtensionEvent(map[string]any{
170 "type": "ext_clear_notes",
171 "extension": extName,
172 })
173 }
174}
175func (h *rpcExtHooks) Submit(string) {} // ignored in rpc mode
176func (h *rpcExtHooks) SubmitSlash(string) {} // ignored in rpc mode
177func (h *rpcExtHooks) Insert(string) {} // ignored in rpc mode
178func (h *rpcExtHooks) OpenPanel(string, extproto.PanelSpec) {}
179func (h *rpcExtHooks) UpdatePanel(string, string, string, []string, string) {}
180func (h *rpcExtHooks) ClosePanel(string, string) {}
181func (h *rpcExtHooks) SetStatus(extName, key, level, text string) {
182 if h.server != nil {
183 h.server.writeExtensionEvent(map[string]any{
184 "type": "ext_status", "extension": extName, "key": key, "level": level, "text": text,
185 })
186 }
187}
188func (h *rpcExtHooks) SetWidget(extName, id, position, title string, lines []string) {
189 if h.server != nil {
190 h.server.writeExtensionEvent(map[string]any{
191 "type": "ext_widget", "extension": extName, "id": id, "position": position, "title": title, "lines": lines,
192 })
193 }
194}
195func (h *rpcExtHooks) ClearWidget(extName, id string) {
196 if h.server != nil {
197 h.server.writeExtensionEvent(map[string]any{
198 "type": "ext_widget_clear", "extension": extName, "id": id,
199 })
200 }
201}
202
203func (h *rpcExtHooks) ClearExtensionChrome(extName string) {
204 if h.server != nil {
205 h.server.writeExtensionEvent(map[string]any{
206 "type": "ext_chrome_clear", "extension": extName,
207 })
208 }
209}
210
211const maxPendingRPCExtEvents = 64
212
213type rpcServer struct {
214 ctx context.Context
215 args Args
216 agent *core.Agent
217 provider string
218 model string
219 out io.Writer
220 version string
221
222 writeMu sync.Mutex
223 turnMu sync.Mutex // serialises one prompt at a time
224 activeCancel context.CancelFunc
225 started bool
226 authed bool
227 pendingExtEvents [][]byte
228
229 // inFlight tracks long-running command goroutines so run() can
230 // wait for them before returning when stdin closes. Without this,
231 // piping a single 'prompt' command into 'zut rpc' would race the
232 // process exit against the agent loop and the prompt would never
233 // produce output.
234 inFlight sync.WaitGroup
235}
236
237// run reads NDJSON commands from in and dispatches them. Returns when
238// in is closed AND every in-flight long-running command (prompt /
239// compact) has finished, so a quick `echo cmd | zut rpc` invocation
240// still produces full output before the process exits.
241func (s *rpcServer) run(in io.Reader) error {
242 requireToken := os.Getenv("ZUTCORE_RPC_TOKEN") != ""
243 s.writeMu.Lock()
244 s.started = true
245 s.authed = !requireToken
246 if s.authed {
247 s.flushPendingExtEventsLocked()
248 }
249 s.writeMu.Unlock()
250
251 sc := bufio.NewScanner(in)
252 sc.Buffer(make([]byte, 0, 64*1024), 16*1024*1024)
253 for sc.Scan() {
254 line := strings.TrimSpace(sc.Text())
255 if line == "" {
256 continue
257 }
258 var head struct {
259 Type string `json:"type"`
260 ID string `json:"id,omitempty"`
261 }
262 if err := json.Unmarshal([]byte(line), &head); err != nil {
263 s.writeError("", "", fmt.Sprintf("malformed json: %v", err))
264 continue
265 }
266 if !s.isAuthed() {
267 if head.Type != "hello" {
268 s.writeError(head.ID, head.Type, "auth required: send hello with token first")
269 continue
270 }
271 var hello struct {
272 Token string `json:"token"`
273 }
274 _ = json.Unmarshal([]byte(line), &hello)
275 if hello.Token != os.Getenv("ZUTCORE_RPC_TOKEN") {
276 s.writeError(head.ID, head.Type, "invalid token")
277 return fmt.Errorf("rpc: bad auth token")
278 }
279 s.authenticate(head.ID, head.Type)
280 continue
281 }
282 s.dispatch(head.Type, head.ID, []byte(line))
283 }
284 err := sc.Err()
285 s.inFlight.Wait()
286 return err
287}
288
289func resolveRPCModel(providerName, modelID string) (provider.Model, error) {
290 model, err := provider.FindModel(providerName, modelID)
291 if err == nil {
292 return model, nil
293 }
294 if providerName == provider.ProviderOpenCodeGo && provider.AcceptsUnlistedModels(providerName) {
295 return provider.DynamicOpenCodeGoModel(providerName, modelID, ""), nil
296 }
297 return provider.Model{}, err
298}
299
300// dispatch routes a command. Long-running commands (prompt, compact)
301// run on their own goroutine so the read loop stays responsive.
302func (s *rpcServer) dispatch(cmd, id string, raw []byte) {
303 switch cmd {
304 case "hello":
305 s.writeResponse(id, cmd, map[string]any{
306 "protocol_version": 1,
307 "version": s.version,
308 "provider": s.provider,
309 "model": s.model,
310 })
311 case "prompt":
312 var req struct {
313 Message string `json:"message"`
314 Images []struct {
315 MimeType string `json:"mime_type"`
316 Data []byte `json:"data"`
317 } `json:"images"`
318 }
319 if err := json.Unmarshal(raw, &req); err != nil {
320 s.writeError(id, cmd, err.Error())
321 return
322 }
323 s.inFlight.Add(1)
324 go func() {
325 defer s.inFlight.Done()
326 s.runPrompt(id, req.Message, req.Images)
327 }()
328
329 case "abort":
330 if c := s.takeCancel(); c != nil {
331 c()
332 }
333 s.writeResponse(id, cmd, nil)
334
335 case "compact":
336 s.inFlight.Add(1)
337 go func() {
338 defer s.inFlight.Done()
339 s.runCompact(id)
340 }()
341
342 case "get_state":
343 s.writeResponse(id, cmd, s.snapshotState())
344
345 case "get_messages":
346 s.writeResponse(id, cmd, map[string]any{
347 "messages": messagesToJSON(s.agent.Messages()),
348 })
349
350 case "clear":
351 s.agent.SetMessages(nil)
352 s.agent.ResetRepetitionGuard()
353 s.agent.SetPlan(nil)
354 s.writeResponse(id, cmd, nil)
355
356 case "set_model":
357 if !s.turnMu.TryLock() {
358 s.writeError(id, cmd, "runtime is busy")
359 return
360 }
361 defer s.turnMu.Unlock()
362
363 var req struct {
364 Model string `json:"model"`
365 }
366 if err := json.Unmarshal(raw, &req); err != nil {
367 s.writeError(id, cmd, err.Error())
368 return
369 }
370 model := strings.TrimSpace(req.Model)
371 if model == "" {
372 s.writeError(id, cmd, "model must not be empty")
373 return
374 }
375 metadata, err := resolveRPCModel(s.provider, model)
376 if err != nil {
377 s.writeError(id, cmd, err.Error())
378 return
379 }
380 if setter, ok := s.agent.Client.(provider.ModelMetadataSetter); ok {
381 setter.SetModelMetadata(metadata)
382 }
383 if metadata.ContextWindow > 0 {
384 s.agent.ContextWindow = metadata.ContextWindow
385 }
386 if metadata.MaxOutput > 0 {
387 s.agent.MaxTokens = metadata.MaxOutput
388 }
389 s.agent.Model = model
390 s.model = model
391 s.writeResponse(id, cmd, map[string]any{"model": model})
392
393 case "set_reasoning":
394 var req struct {
395 Reasoning string `json:"reasoning"`
396 }
397 if err := json.Unmarshal(raw, &req); err != nil {
398 s.writeError(id, cmd, err.Error())
399 return
400 }
401 reasoning := provider.NormalizeReasoning(req.Reasoning)
402 switch reasoning {
403 case "", "minimum", "low", "medium", "high", "xhigh", "max":
404 default:
405 s.writeError(id, cmd, "reasoning must be off|minimum|low|medium|high|xhigh|max")
406 return
407 }
408 s.agent.Reasoning = reasoning
409 s.writeResponse(id, cmd, map[string]any{"reasoning": reasoning})
410
411 case "get_models":
412 out := []map[string]any{}
413 for _, m := range provider.ModelsForProvider(s.provider) {
414 out = append(out, map[string]any{
415 "id": m.ID,
416 "provider": m.Provider,
417 "context_window": m.ContextWindow,
418 "max_output": m.MaxOutput,
419 "reasoning": m.Reasoning,
420 })
421 }
422 s.writeResponse(id, cmd, map[string]any{"models": out})
423
424 case "ping":
425 s.writeResponse(id, cmd, map[string]any{"pong": true})
426
427 default:
428 s.writeError(id, cmd, "unknown command")
429 }
430}
431
432// runPrompt executes a single prompt turn and streams events out.
433// Holds turnMu so a second concurrent prompt blocks until this one
434// finishes; the user can abort with the abort command.
435func (s *rpcServer) runPrompt(id, message string, images []struct {
436 MimeType string `json:"mime_type"`
437 Data []byte `json:"data"`
438}) {
439 s.turnMu.Lock()
440 defer s.turnMu.Unlock()
441
442 subCtx, cancel := context.WithCancel(s.ctx)
443 s.setCancel(cancel)
444 defer s.setCancel(nil)
445
446 s.writeResponse(id, "prompt", map[string]any{"started": true})
447
448 imgs := make([]provider.ImageBlock, 0, len(images))
449 for _, im := range images {
450 imgs = append(imgs, provider.ImageBlock{MimeType: im.MimeType, Data: im.Data})
451 }
452
453 err := s.agent.Prompt(subCtx, message, imgs, func(ev core.AgentEvent) {
454 // EvDone is emitted by the agent loop and we re-emit our own
455 // 'done' below; suppressing it here avoids duplicate frames.
456 if _, ok := ev.(core.EvDone); ok {
457 return
458 }
459 s.writeEvent(modes.EventToJSON(ev))
460 })
461 // Don't emit a stand-alone error event for cancellation; the prior
462 // turn_end with stop=aborted already carries that signal.
463 if err != nil && !errors.Is(err, context.Canceled) {
464 s.writeEvent(map[string]any{"type": "error", "message": err.Error()})
465 }
466 s.writeEvent(map[string]any{"type": "done"})
467}
468
469// runCompact mirrors runPrompt but for compaction.
470func (s *rpcServer) runCompact(id string) {
471 s.turnMu.Lock()
472 defer s.turnMu.Unlock()
473
474 subCtx, cancel := context.WithCancel(s.ctx)
475 s.setCancel(cancel)
476 defer s.setCancel(nil)
477
478 s.writeResponse(id, "compact", map[string]any{"started": true})
479 summary, err := s.agent.CompactWithEvents(subCtx, 4, func(ev core.AgentEvent) {
480 s.writeEvent(modes.EventToJSON(ev))
481 })
482 if err != nil {
483 if !errors.Is(err, context.Canceled) {
484 s.writeEvent(map[string]any{"type": "error", "message": err.Error()})
485 }
486 return
487 }
488 s.writeEvent(map[string]any{
489 "type": "compact_done",
490 "summary": summary,
491 })
492}
493
494// snapshotState builds the get_state response.
495func (s *rpcServer) snapshotState() map[string]any {
496 cum := s.agent.Cost()
497 return map[string]any{
498 "provider": s.provider,
499 "model": s.model,
500 "reasoning": s.agent.Reasoning,
501 "cwd": s.args.CWD,
502 "message_count": len(s.agent.Messages()),
503 "busy": s.activeCancel != nil,
504 "usage": map[string]any{
505 "input": cum.InputTokens,
506 "output": cum.OutputTokens,
507 "reasoning": rpcReasoningTokens(cum),
508 "cache_read": cum.CacheReadTokens,
509 "cache_write": cum.CacheWriteTokens,
510 "cache_measured_prompt": cum.CacheMeasuredPromptTokens,
511 "cache_measured_read": cum.CacheMeasuredReadTokens,
512 "cost_usd": cum.CostUSD,
513 },
514 }
515}
516
517func rpcReasoningTokens(usage provider.Usage) any {
518 if !usage.ReasoningTokensKnown {
519 return nil
520 }
521 return usage.ReasoningTokens
522}
523
524// ---- write helpers (single-line JSON, mutex-guarded) ----
525
526func (s *rpcServer) writeResponse(id, cmd string, data any) {
527 frame := map[string]any{
528 "type": "response",
529 "command": cmd,
530 "success": true,
531 }
532 if id != "" {
533 frame["id"] = id
534 }
535 if data != nil {
536 frame["data"] = data
537 }
538 s.write(frame)
539}
540
541func (s *rpcServer) writeError(id, cmd, msg string) {
542 frame := map[string]any{
543 "type": "response",
544 "command": cmd,
545 "success": false,
546 "error": msg,
547 }
548 if id != "" {
549 frame["id"] = id
550 }
551 s.write(frame)
552}
553
554func (s *rpcServer) writeEvent(payload map[string]any) {
555 s.write(payload)
556}
557
558func (s *rpcServer) writeExtensionEvent(payload map[string]any) {
559 b, err := json.Marshal(payload)
560 if err != nil {
561 return
562 }
563 s.writeMu.Lock()
564 defer s.writeMu.Unlock()
565 if !s.started || !s.authed {
566 if len(s.pendingExtEvents) < maxPendingRPCExtEvents {
567 s.pendingExtEvents = append(s.pendingExtEvents, b)
568 }
569 return
570 }
571 s.writeLocked(b)
572}
573
574func (s *rpcServer) flushPendingExtEventsLocked() {
575 for _, b := range s.pendingExtEvents {
576 s.writeLocked(b)
577 }
578 s.pendingExtEvents = nil
579}
580
581func (s *rpcServer) isAuthed() bool {
582 s.writeMu.Lock()
583 defer s.writeMu.Unlock()
584 return s.authed
585}
586
587func (s *rpcServer) authenticate(id, cmd string) {
588 frame := map[string]any{
589 "type": "response",
590 "command": cmd,
591 "success": true,
592 "data": map[string]any{
593 "protocol_version": 1,
594 "version": s.version,
595 "provider": s.provider,
596 "model": s.model,
597 },
598 }
599 if id != "" {
600 frame["id"] = id
601 }
602 b, err := json.Marshal(frame)
603 if err != nil {
604 return
605 }
606 s.writeMu.Lock()
607 defer s.writeMu.Unlock()
608 s.authed = true
609 s.writeLocked(b)
610 s.flushPendingExtEventsLocked()
611}
612
613func (s *rpcServer) write(v any) {
614 b, err := json.Marshal(v)
615 if err != nil {
616 return
617 }
618 s.writeMu.Lock()
619 defer s.writeMu.Unlock()
620 s.writeLocked(b)
621}
622
623func (s *rpcServer) writeLocked(b []byte) {
624 _, _ = s.out.Write(b)
625 _, _ = s.out.Write([]byte("\n"))
626}
627
628func (s *rpcServer) setCancel(c context.CancelFunc) {
629 s.writeMu.Lock()
630 s.activeCancel = c
631 s.writeMu.Unlock()
632}
633
634func (s *rpcServer) takeCancel() context.CancelFunc {
635 s.writeMu.Lock()
636 defer s.writeMu.Unlock()
637 c := s.activeCancel
638 s.activeCancel = nil
639 return c
640}
641
642// messagesToJSON serialises a transcript using the same schema as the
643// --json event mode for cross-format consistency.
644func messagesToJSON(msgs []provider.Message) []map[string]any {
645 out := make([]map[string]any, 0, len(msgs))
646 for _, m := range msgs {
647 out = append(out, map[string]any{
648 "role": string(m.Role),
649 "content": modes.ContentToJSON(m.Content),
650 "time": m.Time,
651 })
652 }
653 return out
654}