fix(rpc): wait for in-flight commands on stdin close, drop dup events

three real issues found by actually running zot rpc end-to-end:

1. piping a single command into the process and letting stdin close
   would race the agent loop against the process exit, swallowing
   the entire prompt response. fix: track in-flight prompt and
   compact goroutines in a sync.WaitGroup; run() now waits on it
   before returning.

2. each prompt emitted two consecutive done frames - one from the
   agent loop EvDone passing through EventToJSON, one added at the
   end of runPrompt. suppress EvDone in the prompt sink so only the
   explicit terminator remains.

3. cancelling a turn produced a spurious error frame on top of the
   turn_end stop=aborted that already carried the cancellation.
   suppress error frames when the underlying error is
   context.Canceled, in both prompt and compact paths.

verified manually: ping, get_state, get_models, set_model
(valid + invalid id), clear + get_messages, abort, malformed json,
unknown command, auth gate (missing/wrong/correct token),
stdin-close-while-prompt-running, in-process SDK Prompt, plus all
four reference clients parse cleanly and shell + python actually
drive the protocol.

commit: ebc5dad18c622b7ddc19aecf16a64f85aa4ef6b6
author: patriceckhart · 6mo ago · 2026-04-19 10:35 UTC
parent: a670442
signature: unsigned
refs: tag: v0.0.14

2 files · +35 -6

A examples/rpc/python/__pycache__/zot_client.cpython-314.pyc +0 -0 binary file: no text diff

binary file: no text diff

M internal/agent/rpc.go +35 -6


@@ -4,6 +4,7 @@ import (
4	4		"bufio"
5	5		"context"
6	6		"encoding/json"
	7	+	"errors"
7	8		"fmt"
8	9		"io"
9	10		"os"

@@ -67,10 +68,19 @@ type rpcServer struct {
67	68		turnMu sync.Mutex // serialises one prompt at a time
68	69		activeCancel context.CancelFunc
69	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
70	78		}
71	79
72	80		// run reads NDJSON commands from in and dispatches them. Returns when
73		-	// in is closed. Errors writing to stdout are fatal.
	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.
74	84		func (s *rpcServer) run(in io.Reader) error {
75	85		requireToken := os.Getenv("ZOTCORE_RPC_TOKEN") != ""
76	86		s.authed = !requireToken

@@ -114,7 +124,9 @@ func (s *rpcServer) run(in io.Reader) error {
114	124		}
115	125		s.dispatch(head.Type, head.ID, []byte(line))
116	126		}
117		-	return sc.Err()
	127	+	err := sc.Err()
	128	+	s.inFlight.Wait()
	129	+	return err
118	130		}
119	131
120	132		// dispatch routes a command. Long-running commands (prompt, compact)

@@ -140,7 +152,11 @@ func (s *rpcServer) dispatch(cmd, id string, raw []byte) {
140	152		s.writeError(id, cmd, err.Error())
141	153		return
142	154		}
143		-	go s.runPrompt(id, req.Message, req.Images)
	155	+	s.inFlight.Add(1)
	156	+	go func() {
	157	+	defer s.inFlight.Done()
	158	+	s.runPrompt(id, req.Message, req.Images)
	159	+	}()
144	160
145	161		case "abort":
146	162		if c := s.takeCancel(); c != nil {

@@ -149,7 +165,11 @@ func (s *rpcServer) dispatch(cmd, id string, raw []byte) {
149	165		s.writeResponse(id, cmd, nil)
150	166
151	167		case "compact":
152		-	go s.runCompact(id)
	168	+	s.inFlight.Add(1)
	169	+	go func() {
	170	+	defer s.inFlight.Done()
	171	+	s.runCompact(id)
	172	+	}()
153	173
154	174		case "get_state":
155	175		s.writeResponse(id, cmd, s.snapshotState())

@@ -222,9 +242,16 @@ func (s *rpcServer) runPrompt(id, message string, images []struct {
222	242		}
223	243
224	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	+	}
225	250		s.writeEvent(modes.EventToJSON(ev))
226	251		})
227		-	if err != nil {
	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) {
228	255		s.writeEvent(map[string]any{"type": "error", "message": err.Error()})
229	256		}
230	257		s.writeEvent(map[string]any{"type": "done"})

@@ -242,7 +269,9 @@ func (s *rpcServer) runCompact(id string) {
242	269		s.writeResponse(id, "compact", map[string]any{"started": true})
243	270		summary, err := s.agent.Compact(subCtx, 4, nil)
244	271		if err != nil {
245		-	s.writeEvent(map[string]any{"type": "error", "message": err.Error()})
	272	+	if !errors.Is(err, context.Canceled) {
	273	+	s.writeEvent(map[string]any{"type": "error", "message": err.Error()})
	274	+	}
246	275		return
247	276		}
248	277		s.writeEvent(map[string]any{

 	"bufio"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"os"
 	turnMu       sync.Mutex // serialises one prompt at a time
 	activeCancel context.CancelFunc
 	authed       bool
+	// inFlight tracks long-running command goroutines so run() can
+	// wait for them before returning when stdin closes. Without this,
+	// piping a single 'prompt' command into 'zot rpc' would race the
+	// process exit against the agent loop and the prompt would never
+	// produce output.
+	inFlight sync.WaitGroup
 }
 // run reads NDJSON commands from in and dispatches them. Returns when
-// in is closed. Errors writing to stdout are fatal.
+// in is closed AND every in-flight long-running command (prompt /
+// compact) has finished, so a quick `echo cmd | zot rpc` invocation
+// still produces full output before the process exits.
 func (s *rpcServer) run(in io.Reader) error {
 	requireToken := os.Getenv("ZOTCORE_RPC_TOKEN") != ""
 	s.authed = !requireToken
 		}
 		s.dispatch(head.Type, head.ID, []byte(line))
 	}
-	return sc.Err()
+	err := sc.Err()
+	s.inFlight.Wait()
+	return err
 }
 // dispatch routes a command. Long-running commands (prompt, compact)
 			s.writeError(id, cmd, err.Error())
 			return
 		}
-		go s.runPrompt(id, req.Message, req.Images)
+		s.inFlight.Add(1)
+		go func() {
+			defer s.inFlight.Done()
+			s.runPrompt(id, req.Message, req.Images)
+		}()
 	case "abort":
 		if c := s.takeCancel(); c != nil {
 		s.writeResponse(id, cmd, nil)
 	case "compact":
-		go s.runCompact(id)
+		s.inFlight.Add(1)
+		go func() {
+			defer s.inFlight.Done()
+			s.runCompact(id)
+		}()
 	case "get_state":
 		s.writeResponse(id, cmd, s.snapshotState())
 	}
 	err := s.agent.Prompt(subCtx, message, imgs, func(ev core.AgentEvent) {
+		// EvDone is emitted by the agent loop and we re-emit our own
+		// 'done' below; suppressing it here avoids duplicate frames.
+		if _, ok := ev.(core.EvDone); ok {
+			return
+		}
 		s.writeEvent(modes.EventToJSON(ev))
 	})
-	if err != nil {
+	// Don't emit a stand-alone error event for cancellation; the prior
+	// turn_end with stop=aborted already carries that signal.
+	if err != nil && !errors.Is(err, context.Canceled) {
 		s.writeEvent(map[string]any{"type": "error", "message": err.Error()})
 	}
 	s.writeEvent(map[string]any{"type": "done"})
 	s.writeResponse(id, "compact", map[string]any{"started": true})
 	summary, err := s.agent.Compact(subCtx, 4, nil)
 	if err != nil {
-		s.writeEvent(map[string]any{"type": "error", "message": err.Error()})
+		if !errors.Is(err, context.Canceled) {
+			s.writeEvent(map[string]any{"type": "error", "message": err.Error()})
+		}
 		return
 	}
 	s.writeEvent(map[string]any{