149c6c3 fix(agent): preserve alert delivery during startup · zut

M packages/agent/cli.go +43 -10


@@ -30,17 +30,40 @@ import (
30	30
31	31		// interactiveExtHooks is a tiny adapter that lets the extension
32	32		// manager call back into the Interactive instance built later in
33		-	// runInteractive. The forward-declared *modes.Interactive is filled
34		-	// in immediately after manager construction.
	33	+	// runInteractive. Alerts are buffered until the instance is attached so
	34	+	// an extension that signals attention during startup is not silently lost.
35	35		type interactiveExtHooks struct {
36		-	ivPtr **modes.Interactive
	36	+	mu sync.Mutex
	37	+	interactive *modes.Interactive
	38	+	pendingAlerts []interactiveAlert
	39	+	}
	40	+
	41	+	type interactiveAlert struct {
	42	+	extName string
	43	+	alert extproto.AlertRequest
37	44		}
38	45
39	46		func (h interactiveExtHooks) iv() modes.Interactive {
40		-	if h == nil \|\| h.ivPtr == nil {
	47	+	if h == nil {
41	48		return nil
42	49		}
43		-	return *h.ivPtr
	50	+	h.mu.Lock()
	51	+	defer h.mu.Unlock()
	52	+	return h.interactive
	53	+	}
	54	+
	55	+	func (h interactiveExtHooks) attachInteractive(iv modes.Interactive) {
	56	+	if h == nil \|\| iv == nil {
	57	+	return
	58	+	}
	59	+	h.mu.Lock()
	60	+	h.interactive = iv
	61	+	pending := h.pendingAlerts
	62	+	h.pendingAlerts = nil
	63	+	for _, item := range pending {
	64	+	iv.Alert(item.extName, item.alert)
	65	+	}
	66	+	h.mu.Unlock()
44	67		}
45	68
46	69		func (h *interactiveExtHooks) Notify(extName, level, message string) {

@@ -49,9 +72,18 @@ func (h *interactiveExtHooks) Notify(extName, level, message string) {
49	72		}
50	73		}
51	74		func (h *interactiveExtHooks) Alert(extName string, alert extproto.AlertRequest) {
52		-	if iv := h.iv(); iv != nil {
53		-	iv.Alert(extName, alert)
	75	+	if h == nil {
	76	+	return
	77	+	}
	78	+	h.mu.Lock()
	79	+	if h.interactive == nil {
	80	+	h.pendingAlerts = append(h.pendingAlerts, interactiveAlert{extName: extName, alert: alert})
	81	+	h.mu.Unlock()
	82	+	return
54	83		}
	84	+	iv := h.interactive
	85	+	h.mu.Unlock()
	86	+	iv.Alert(extName, alert)
55	87		}
56	88		func (h *interactiveExtHooks) Submit(text string) {
57	89		if iv := h.iv(); iv != nil {

@@ -659,10 +691,10 @@ func runInteractive(ctx context.Context, args Args, version string) error {
659	691		sharedSandbox := r.Sandbox
660	692
661	693		// Build the extension manager BEFORE the agent so we can fold
662		-	// extension-defined tools into the registry. Forward-declare iv so
663		-	// the host hooks adapter can dereference it after construction.
	694	+	// extension-defined tools into the registry. Attach the interactive
	695	+	// host after constructing it below so startup alerts can be buffered.
664	696		var iv *modes.Interactive
665		-	extHooks := &interactiveExtHooks{ivPtr: &iv}
	697	+	extHooks := &interactiveExtHooks{}
666	698		extMgr := extensions.New(ZotHome(), r.CWD, version, r.Provider, r.Model, extHooks)
667	699		var startupExtensionErrors []string
668	700		// --ext paths first so they win against installed extensions of

@@ -1447,6 +1479,7 @@ func runInteractive(ctx context.Context, args Args, version string) error {
1447	1479		}
1448	1480		},
1449	1481		})
	1482	+	extHooks.attachInteractive(iv)
1450	1483
1451	1484		// Bind the interactive TUI as the Confirmer. We deferred this
1452	1485		// until now because the gate is constructed before the TUI

 // interactiveExtHooks is a tiny adapter that lets the extension
 // manager call back into the Interactive instance built later in
-// runInteractive. The forward-declared *modes.Interactive is filled
+// runInteractive. Alerts are buffered until the instance is attached so
-// in immediately after manager construction.
+// an extension that signals attention during startup is not silently lost.
 type interactiveExtHooks struct {
-	ivPtr **modes.Interactive
+	mu            sync.Mutex
+	interactive   *modes.Interactive
+	pendingAlerts []interactiveAlert
+}
+type interactiveAlert struct {
+	extName string
+	alert   extproto.AlertRequest
 }
 func (h *interactiveExtHooks) iv() *modes.Interactive {
-	if h == nil || h.ivPtr == nil {
+	if h == nil {
 		return nil
 	}
-	return *h.ivPtr
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	return h.interactive
+}
+func (h *interactiveExtHooks) attachInteractive(iv *modes.Interactive) {
+	if h == nil || iv == nil {
+		return
+	}
+	h.mu.Lock()
+	h.interactive = iv
+	pending := h.pendingAlerts
+	h.pendingAlerts = nil
+	for _, item := range pending {
+		iv.Alert(item.extName, item.alert)
+	}
+	h.mu.Unlock()
 }
 func (h *interactiveExtHooks) Notify(extName, level, message string) {
 	}
 }
 func (h *interactiveExtHooks) Alert(extName string, alert extproto.AlertRequest) {
-	if iv := h.iv(); iv != nil {
+	if h == nil {
-		iv.Alert(extName, alert)
+		return
+	}
+	h.mu.Lock()
+	if h.interactive == nil {
+		h.pendingAlerts = append(h.pendingAlerts, interactiveAlert{extName: extName, alert: alert})
+		h.mu.Unlock()
+		return
 	}
+	iv := h.interactive
+	h.mu.Unlock()
+	iv.Alert(extName, alert)
 }
 func (h *interactiveExtHooks) Submit(text string) {
 	if iv := h.iv(); iv != nil {
 	sharedSandbox := r.Sandbox
 	// Build the extension manager BEFORE the agent so we can fold
-	// extension-defined tools into the registry. Forward-declare iv so
+	// extension-defined tools into the registry. Attach the interactive
-	// the host hooks adapter can dereference it after construction.
+	// host after constructing it below so startup alerts can be buffered.
 	var iv *modes.Interactive
-	extHooks := &interactiveExtHooks{ivPtr: &iv}
+	extHooks := &interactiveExtHooks{}
 	extMgr := extensions.New(ZotHome(), r.CWD, version, r.Provider, r.Model, extHooks)
 	var startupExtensionErrors []string
 	// --ext paths first so they win against installed extensions of
 			}
 		},
 	})
+	extHooks.attachInteractive(iv)
 	// Bind the interactive TUI as the Confirmer. We deferred this
 	// until now because the gate is constructed before the TUI

A packages/agent/cli_alert_test.go +30 -0

@@ -0,0 +1,30 @@
	1	+	package agent
	2	+
	3	+	import (
	4	+	"testing"
	5	+
	6	+	"github.com/patriceckhart/zot/packages/agent/extproto"
	7	+	"github.com/patriceckhart/zot/packages/agent/modes"
	8	+	)
	9	+
	10	+	func TestInteractiveExtHooksBufferStartupAlerts(t *testing.T) {
	11	+	hooks := &interactiveExtHooks{}
	12	+	hooks.Alert("question-ext", extproto.AlertRequest{Kind: extproto.AlertKindBell, Reason: "question_ready"})
	13	+
	14	+	hooks.mu.Lock()
	15	+	pending := len(hooks.pendingAlerts)
	16	+	hooks.mu.Unlock()
	17	+	if pending != 1 {
	18	+	t.Fatalf("pending startup alerts = %d, want 1", pending)
	19	+	}
	20	+
	21	+	iv := &modes.Interactive{}
	22	+	hooks.attachInteractive(iv)
	23	+	hooks.mu.Lock()
	24	+	pending = len(hooks.pendingAlerts)
	25	+	attached := hooks.interactive == iv
	26	+	hooks.mu.Unlock()
	27	+	if pending != 0 \|\| !attached {
	28	+	t.Fatalf("after attach: pending=%d attached=%v, want no pending alert and attached host", pending, attached)
	29	+	}
	30	+	}

@@ -0,0 +1,30 @@
	1	+	package agent
	2	+
	3	+	import (
	4	+	"testing"
	5	+
	6	+	"github.com/patriceckhart/zot/packages/agent/extproto"
	7	+	"github.com/patriceckhart/zot/packages/agent/modes"
	8	+	)
	9	+
	10	+	func TestInteractiveExtHooksBufferStartupAlerts(t *testing.T) {
	11	+	hooks := &interactiveExtHooks{}
	12	+	hooks.Alert("question-ext", extproto.AlertRequest{Kind: extproto.AlertKindBell, Reason: "question_ready"})
	13	+
	14	+	hooks.mu.Lock()
	15	+	pending := len(hooks.pendingAlerts)
	16	+	hooks.mu.Unlock()
	17	+	if pending != 1 {
	18	+	t.Fatalf("pending startup alerts = %d, want 1", pending)
	19	+	}
	20	+
	21	+	iv := &modes.Interactive{}
	22	+	hooks.attachInteractive(iv)
	23	+	hooks.mu.Lock()
	24	+	pending = len(hooks.pendingAlerts)
	25	+	attached := hooks.interactive == iv
	26	+	hooks.mu.Unlock()
	27	+	if pending != 0 \|\| !attached {
	28	+	t.Fatalf("after attach: pending=%d attached=%v, want no pending alert and attached host", pending, attached)
	29	+	}
	30	+	}

M packages/agent/modes/interactive.go +5 -9


@@ -3014,21 +3014,17 @@ func (i *Interactive) emitAlertLocked(alert extproto.AlertRequest) {
3014	3014		_ = tui.WriteBell(i.cfg.Terminal)
3015	3015		}
3016	3016
3017		-	// scheduleMainAlert emits a main-session alert now when no paced text remains,
3018		-	// or defers it until the pacer drains and the next redraw commits the final frame.
	3017	+	// scheduleMainAlert defers a main-session alert until the next redraw
	3018	+	// commits the final frame. Keeping this deferred even when no paced text is
	3019	+	// pending avoids racing the pacer's final state transition.
3019	3020		func (i *Interactive) scheduleMainAlert(reason string) {
3020	3021		if reason == "" {
3021	3022		return
3022	3023		}
3023	3024		i.mu.Lock()
3024		-	if len(i.streamPending) > 0 {
3025		-	i.pendingAlert = &extproto.AlertRequest{Kind: extproto.AlertKindBell, Reason: reason}
3026		-	i.mu.Unlock()
3027		-	i.invalidate()
3028		-	return
3029		-	}
3030		-	i.emitAlertLocked(extproto.AlertRequest{Kind: extproto.AlertKindBell, Reason: reason})
	3025	+	i.pendingAlert = &extproto.AlertRequest{Kind: extproto.AlertKindBell, Reason: reason}
3031	3026		i.mu.Unlock()
	3027	+	i.invalidate()
3032	3028		}
3033	3029
3034	3030		// ClearNotes removes every note line owned by extName from the

 	_ = tui.WriteBell(i.cfg.Terminal)
 }
-// scheduleMainAlert emits a main-session alert now when no paced text remains,
+// scheduleMainAlert defers a main-session alert until the next redraw
-// or defers it until the pacer drains and the next redraw commits the final frame.
+// commits the final frame. Keeping this deferred even when no paced text is
+// pending avoids racing the pacer's final state transition.
 func (i *Interactive) scheduleMainAlert(reason string) {
 	if reason == "" {
 		return
 	}
 	i.mu.Lock()
-	if len(i.streamPending) > 0 {
+	i.pendingAlert = &extproto.AlertRequest{Kind: extproto.AlertKindBell, Reason: reason}
-		i.pendingAlert = &extproto.AlertRequest{Kind: extproto.AlertKindBell, Reason: reason}
-		i.mu.Unlock()
-		i.invalidate()
-		return
-	}
-	i.emitAlertLocked(extproto.AlertRequest{Kind: extproto.AlertKindBell, Reason: reason})
 	i.mu.Unlock()
+	i.invalidate()
 }
 // ClearNotes removes every note line owned by extName from the

M packages/agent/rpc.go +89 -27


@@ -49,10 +49,22 @@ func runRPCMode(ctx context.Context, args Args, version string) error {
49	49		return err
50	50		}
51	51
	52	+	// Construct the RPC event sink before extension startup. Alerts can
	53	+	// arrive as soon as an extension is ready; rpcServer queues them until
	54	+	// the command loop has started and, when configured, authenticated.
	55	+	server := &rpcServer{
	56	+	ctx: ctx,
	57	+	args: args,
	58	+	provider: r.Provider,
	59	+	model: r.Model,
	60	+	out: os.Stdout,
	61	+	version: version,
	62	+	}
	63	+
52	64		// Extensions: same lifecycle as interactive mode, minus the
53	65		// host-hooks integration. Notify/Display calls from extensions
54	66		// emit RPC events instead of TUI lines so any consumer can react.
55		-	extHooks := &rpcExtHooks{}
	67	+	extHooks := &rpcExtHooks{server: server}
56	68		extMgr := extensions.New(ZotHome(), r.CWD, version, r.Provider, r.Model, extHooks)
57	69		for _, e := range extMgr.LoadExplicit(ctx, args.Exts) {
58	70		fmt.Fprintln(os.Stderr, "extension load:", e)

@@ -67,6 +79,7 @@ func runRPCMode(ctx context.Context, args Args, version string) error {
67	79		r.MergeExtensionTools(&extToolAdapter{mgr: extMgr})
68	80
69	81		ag := r.NewAgent()
	82	+	server.agent = ag
70	83		ag.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
71	84		r := extMgr.InterceptToolCall(ctx, call.ID, call.Name, call.Arguments)
72	85		if r.Block {

@@ -103,16 +116,6 @@ func runRPCMode(ctx context.Context, args Args, version string) error {
103	116
104	117		extMgr.EmitEvent(extproto.EventFromHost{Event: "session_start"})
105	118
106		-	server := &rpcServer{
107		-	ctx: ctx,
108		-	args: args,
109		-	agent: ag,
110		-	provider: r.Provider,
111		-	model: r.Model,
112		-	out: os.Stdout,
113		-	version: version,
114		-	}
115		-	extHooks.server = server
116	119		return server.run(os.Stdin)
117	120		}
118	121

@@ -126,7 +129,7 @@ type rpcExtHooks struct {
126	129
127	130		func (h *rpcExtHooks) Notify(extName, level, message string) {
128	131		if h.server != nil {
129		-	h.server.writeEvent(map[string]any{
	132	+	h.server.writeExtensionEvent(map[string]any{
130	133		"type": "ext_notify",
131	134		"extension": extName,
132	135		"level": level,

@@ -136,7 +139,7 @@ func (h *rpcExtHooks) Notify(extName, level, message string) {
136	139		}
137	140		func (h *rpcExtHooks) Alert(extName string, alert extproto.AlertRequest) {
138	141		if h.server != nil {
139		-	h.server.writeEvent(map[string]any{
	142	+	h.server.writeExtensionEvent(map[string]any{
140	143		"type": "ext_alert",
141	144		"extension": extName,
142	145		"kind": alert.Kind,

@@ -146,7 +149,7 @@ func (h *rpcExtHooks) Alert(extName string, alert extproto.AlertRequest) {
146	149		}
147	150		func (h *rpcExtHooks) Display(extName, text string) {
148	151		if h.server != nil {
149		-	h.server.writeEvent(map[string]any{
	152	+	h.server.writeExtensionEvent(map[string]any{
150	153		"type": "ext_display",
151	154		"extension": extName,
152	155		"text": text,

@@ -155,7 +158,7 @@ func (h *rpcExtHooks) Display(extName, text string) {
155	158		}
156	159		func (h *rpcExtHooks) ClearNotes(extName string) {
157	160		if h.server != nil {
158		-	h.server.writeEvent(map[string]any{
	161	+	h.server.writeExtensionEvent(map[string]any{
159	162		"type": "ext_clear_notes",
160	163		"extension": extName,
161	164		})

@@ -177,10 +180,12 @@ type rpcServer struct {
177	180		out io.Writer
178	181		version string
179	182
180		-	writeMu sync.Mutex
181		-	turnMu sync.Mutex // serialises one prompt at a time
182		-	activeCancel context.CancelFunc
183		-	authed bool
	183	+	writeMu sync.Mutex
	184	+	turnMu sync.Mutex // serialises one prompt at a time
	185	+	activeCancel context.CancelFunc
	186	+	started bool
	187	+	authed bool
	188	+	pendingExtEvents [][]byte
184	189
185	190		// inFlight tracks long-running command goroutines so run() can
186	191		// wait for them before returning when stdin closes. Without this,

@@ -196,7 +201,13 @@ type rpcServer struct {
196	201		// still produces full output before the process exits.
197	202		func (s *rpcServer) run(in io.Reader) error {
198	203		requireToken := os.Getenv("ZOTCORE_RPC_TOKEN") != ""
	204	+	s.writeMu.Lock()
	205	+	s.started = true
199	206		s.authed = !requireToken
	207	+	if s.authed {
	208	+	s.flushPendingExtEventsLocked()
	209	+	}
	210	+	s.writeMu.Unlock()
200	211
201	212		sc := bufio.NewScanner(in)
202	213		sc.Buffer(make([]byte, 0, 641024), 161024*1024)

@@ -213,7 +224,7 @@ func (s *rpcServer) run(in io.Reader) error {
213	224		s.writeError("", "", fmt.Sprintf("malformed json: %v", err))
214	225		continue
215	226		}
216		-	if !s.authed {
	227	+	if !s.isAuthed() {
217	228		if head.Type != "hello" {
218	229		s.writeError(head.ID, head.Type, "auth required: send hello with token first")
219	230		continue

@@ -226,13 +237,7 @@ func (s *rpcServer) run(in io.Reader) error {
226	237		s.writeError(head.ID, head.Type, "invalid token")
227	238		return fmt.Errorf("rpc: bad auth token")
228	239		}
229		-	s.authed = true
230		-	s.writeResponse(head.ID, head.Type, map[string]any{
231		-	"protocol_version": 1,
232		-	"version": s.version,
233		-	"provider": s.provider,
234		-	"model": s.model,
235		-	})
	240	+	s.authenticate(head.ID, head.Type)
236	241		continue
237	242		}
238	243		s.dispatch(head.Type, head.ID, []byte(line))

@@ -473,6 +478,59 @@ func (s *rpcServer) writeEvent(payload map[string]any) {
473	478		s.write(payload)
474	479		}
475	480
	481	+	func (s *rpcServer) writeExtensionEvent(payload map[string]any) {
	482	+	b, err := json.Marshal(payload)
	483	+	if err != nil {
	484	+	return
	485	+	}
	486	+	s.writeMu.Lock()
	487	+	defer s.writeMu.Unlock()
	488	+	if !s.started \|\| !s.authed {
	489	+	s.pendingExtEvents = append(s.pendingExtEvents, b)
	490	+	return
	491	+	}
	492	+	s.writeLocked(b)
	493	+	}
	494	+
	495	+	func (s *rpcServer) flushPendingExtEventsLocked() {
	496	+	for _, b := range s.pendingExtEvents {
	497	+	s.writeLocked(b)
	498	+	}
	499	+	s.pendingExtEvents = nil
	500	+	}
	501	+
	502	+	func (s *rpcServer) isAuthed() bool {
	503	+	s.writeMu.Lock()
	504	+	defer s.writeMu.Unlock()
	505	+	return s.authed
	506	+	}
	507	+
	508	+	func (s *rpcServer) authenticate(id, cmd string) {
	509	+	frame := map[string]any{
	510	+	"type": "response",
	511	+	"command": cmd,
	512	+	"success": true,
	513	+	"data": map[string]any{
	514	+	"protocol_version": 1,
	515	+	"version": s.version,
	516	+	"provider": s.provider,
	517	+	"model": s.model,
	518	+	},
	519	+	}
	520	+	if id != "" {
	521	+	frame["id"] = id
	522	+	}
	523	+	b, err := json.Marshal(frame)
	524	+	if err != nil {
	525	+	return
	526	+	}
	527	+	s.writeMu.Lock()
	528	+	defer s.writeMu.Unlock()
	529	+	s.authed = true
	530	+	s.writeLocked(b)
	531	+	s.flushPendingExtEventsLocked()
	532	+	}
	533	+
476	534		func (s *rpcServer) write(v any) {
477	535		b, err := json.Marshal(v)
478	536		if err != nil {

@@ -480,6 +538,10 @@ func (s *rpcServer) write(v any) {
480	538		}
481	539		s.writeMu.Lock()
482	540		defer s.writeMu.Unlock()
	541	+	s.writeLocked(b)
	542	+	}
	543	+
	544	+	func (s *rpcServer) writeLocked(b []byte) {
483	545		_, _ = s.out.Write(b)
484	546		_, _ = s.out.Write([]byte("\n"))
485	547		}

 		return err
 	}
+	// Construct the RPC event sink before extension startup. Alerts can
+	// arrive as soon as an extension is ready; rpcServer queues them until
+	// the command loop has started and, when configured, authenticated.
+	server := &rpcServer{
+		ctx:      ctx,
+		args:     args,
+		provider: r.Provider,
+		model:    r.Model,
+		out:      os.Stdout,
+		version:  version,
+	}
 	// Extensions: same lifecycle as interactive mode, minus the
 	// host-hooks integration. Notify/Display calls from extensions
 	// emit RPC events instead of TUI lines so any consumer can react.
-	extHooks := &rpcExtHooks{}
+	extHooks := &rpcExtHooks{server: server}
 	extMgr := extensions.New(ZotHome(), r.CWD, version, r.Provider, r.Model, extHooks)
 	for _, e := range extMgr.LoadExplicit(ctx, args.Exts) {
 		fmt.Fprintln(os.Stderr, "extension load:", e)
 	r.MergeExtensionTools(&extToolAdapter{mgr: extMgr})
 	ag := r.NewAgent()
+	server.agent = ag
 	ag.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
 		r := extMgr.InterceptToolCall(ctx, call.ID, call.Name, call.Arguments)
 		if r.Block {
 	extMgr.EmitEvent(extproto.EventFromHost{Event: "session_start"})
-	server := &rpcServer{
-		ctx:      ctx,
-		args:     args,
-		agent:    ag,
-		provider: r.Provider,
-		model:    r.Model,
-		out:      os.Stdout,
-		version:  version,
-	}
-	extHooks.server = server
 	return server.run(os.Stdin)
 }
 func (h *rpcExtHooks) Notify(extName, level, message string) {
 	if h.server != nil {
-		h.server.writeEvent(map[string]any{
+		h.server.writeExtensionEvent(map[string]any{
 			"type":      "ext_notify",
 			"extension": extName,
 			"level":     level,
 }
 func (h *rpcExtHooks) Alert(extName string, alert extproto.AlertRequest) {
 	if h.server != nil {
-		h.server.writeEvent(map[string]any{
+		h.server.writeExtensionEvent(map[string]any{
 			"type":      "ext_alert",
 			"extension": extName,
 			"kind":      alert.Kind,
 }
 func (h *rpcExtHooks) Display(extName, text string) {
 	if h.server != nil {
-		h.server.writeEvent(map[string]any{
+		h.server.writeExtensionEvent(map[string]any{
 			"type":      "ext_display",
 			"extension": extName,
 			"text":      text,
 }
 func (h *rpcExtHooks) ClearNotes(extName string) {
 	if h.server != nil {
-		h.server.writeEvent(map[string]any{
+		h.server.writeExtensionEvent(map[string]any{
 			"type":      "ext_clear_notes",
 			"extension": extName,
 		})
 	out      io.Writer
 	version  string
-	writeMu      sync.Mutex
+	writeMu          sync.Mutex
-	turnMu       sync.Mutex // serialises one prompt at a time
+	turnMu           sync.Mutex // serialises one prompt at a time
-	activeCancel context.CancelFunc
+	activeCancel     context.CancelFunc
-	authed       bool
+	started          bool
+	authed           bool
+	pendingExtEvents [][]byte
 	// inFlight tracks long-running command goroutines so run() can
 	// wait for them before returning when stdin closes. Without this,
 // still produces full output before the process exits.
 func (s *rpcServer) run(in io.Reader) error {
 	requireToken := os.Getenv("ZOTCORE_RPC_TOKEN") != ""
+	s.writeMu.Lock()
+	s.started = true
 	s.authed = !requireToken
+	if s.authed {
+		s.flushPendingExtEventsLocked()
+	}
+	s.writeMu.Unlock()
 	sc := bufio.NewScanner(in)
 	sc.Buffer(make([]byte, 0, 64*1024), 16*1024*1024)
 			s.writeError("", "", fmt.Sprintf("malformed json: %v", err))
 			continue
 		}
-		if !s.authed {
+		if !s.isAuthed() {
 			if head.Type != "hello" {
 				s.writeError(head.ID, head.Type, "auth required: send hello with token first")
 				continue
 				s.writeError(head.ID, head.Type, "invalid token")
 				return fmt.Errorf("rpc: bad auth token")
 			}
-			s.authed = true
+			s.authenticate(head.ID, head.Type)
-			s.writeResponse(head.ID, head.Type, map[string]any{
-				"protocol_version": 1,
-				"version":          s.version,
-				"provider":         s.provider,
-				"model":            s.model,
-			})
 			continue
 		}
 		s.dispatch(head.Type, head.ID, []byte(line))
 	s.write(payload)
 }
+func (s *rpcServer) writeExtensionEvent(payload map[string]any) {
+	b, err := json.Marshal(payload)
+	if err != nil {
+		return
+	}
+	s.writeMu.Lock()
+	defer s.writeMu.Unlock()
+	if !s.started || !s.authed {
+		s.pendingExtEvents = append(s.pendingExtEvents, b)
+		return
+	}
+	s.writeLocked(b)
+}
+func (s *rpcServer) flushPendingExtEventsLocked() {
+	for _, b := range s.pendingExtEvents {
+		s.writeLocked(b)
+	}
+	s.pendingExtEvents = nil
+}
+func (s *rpcServer) isAuthed() bool {
+	s.writeMu.Lock()
+	defer s.writeMu.Unlock()
+	return s.authed
+}
+func (s *rpcServer) authenticate(id, cmd string) {
+	frame := map[string]any{
+		"type":    "response",
+		"command": cmd,
+		"success": true,
+		"data": map[string]any{
+			"protocol_version": 1,
+			"version":          s.version,
+			"provider":         s.provider,
+			"model":            s.model,
+		},
+	}
+	if id != "" {
+		frame["id"] = id
+	}
+	b, err := json.Marshal(frame)
+	if err != nil {
+		return
+	}
+	s.writeMu.Lock()
+	defer s.writeMu.Unlock()
+	s.authed = true
+	s.writeLocked(b)
+	s.flushPendingExtEventsLocked()
+}
 func (s *rpcServer) write(v any) {
 	b, err := json.Marshal(v)
 	if err != nil {
 	}
 	s.writeMu.Lock()
 	defer s.writeMu.Unlock()
+	s.writeLocked(b)
+}
+func (s *rpcServer) writeLocked(b []byte) {
 	_, _ = s.out.Write(b)
 	_, _ = s.out.Write([]byte("\n"))
 }

M packages/agent/rpc_alert_test.go +21 -1


@@ -11,7 +11,7 @@ import (
11	11
12	12		func TestRPCExtAlertIsJSONOnly(t *testing.T) {
13	13		var out bytes.Buffer
14		-	hooks := &rpcExtHooks{server: &rpcServer{out: &out}}
	14	+	hooks := &rpcExtHooks{server: &rpcServer{out: &out, started: true, authed: true}}
15	15		hooks.Alert("question-ext", extproto.AlertRequest{Kind: extproto.AlertKindBell, Reason: "question_ready"})
16	16
17	17		if strings.ContainsRune(out.String(), '\a') {

@@ -25,3 +25,23 @@ func TestRPCExtAlertIsJSONOnly(t *testing.T) {
25	25		t.Fatalf("RPC alert frame = %#v", frame)
26	26		}
27	27		}
	28	+
	29	+	func TestRPCExtAlertWaitsForAuthentication(t *testing.T) {
	30	+	t.Setenv("ZOTCORE_RPC_TOKEN", "synthetic-token")
	31	+
	32	+	var out bytes.Buffer
	33	+	server := &rpcServer{out: &out, provider: "test", model: "test-model", version: "test-version"}
	34	+	hooks := &rpcExtHooks{server: server}
	35	+	hooks.Alert("question-ext", extproto.AlertRequest{Kind: extproto.AlertKindBell, Reason: "question_ready"})
	36	+	if out.Len() != 0 {
	37	+	t.Fatalf("pre-auth RPC output = %q, want empty", out.String())
	38	+	}
	39	+
	40	+	if err := server.run(strings.NewReader(`{"type":"hello","id":"1","token":"synthetic-token"}` + "\n")); err != nil {
	41	+	t.Fatalf("run RPC server: %v", err)
	42	+	}
	43	+	lines := strings.Split(strings.TrimSpace(out.String()), "\n")
	44	+	if len(lines) != 2 \|\| !strings.Contains(lines[0], `"type":"response"`) \|\| !strings.Contains(lines[1], `"type":"ext_alert"`) {
	45	+	t.Fatalf("authenticated RPC frames = %q, want response followed by ext_alert", out.String())
	46	+	}
	47	+	}


@@ -11,7 +11,7 @@ import (
11			11
12		func TestRPCExtAlertIsJSONOnly(t *testing.T) {	12		func TestRPCExtAlertIsJSONOnly(t *testing.T) {
13		var out bytes.Buffer	13		var out bytes.Buffer
14	-	hooks := &rpcExtHooks{server: &rpcServer{out: &out}}	14	+	hooks := &rpcExtHooks{server: &rpcServer{out: &out, started: true, authed: true}}
15		hooks.Alert("question-ext", extproto.AlertRequest{Kind: extproto.AlertKindBell, Reason: "question_ready"})	15		hooks.Alert("question-ext", extproto.AlertRequest{Kind: extproto.AlertKindBell, Reason: "question_ready"})
16			16
17		if strings.ContainsRune(out.String(), '\a') {	17		if strings.ContainsRune(out.String(), '\a') {

@@ -25,3 +25,23 @@ func TestRPCExtAlertIsJSONOnly(t *testing.T) {
25		t.Fatalf("RPC alert frame = %#v", frame)	25		t.Fatalf("RPC alert frame = %#v", frame)
26		}	26		}
27		}	27		}
			28	+
			29	+	func TestRPCExtAlertWaitsForAuthentication(t *testing.T) {
			30	+	t.Setenv("ZOTCORE_RPC_TOKEN", "synthetic-token")
			31	+
			32	+	var out bytes.Buffer
			33	+	server := &rpcServer{out: &out, provider: "test", model: "test-model", version: "test-version"}
			34	+	hooks := &rpcExtHooks{server: server}
			35	+	hooks.Alert("question-ext", extproto.AlertRequest{Kind: extproto.AlertKindBell, Reason: "question_ready"})
			36	+	if out.Len() != 0 {
			37	+	t.Fatalf("pre-auth RPC output = %q, want empty", out.String())
			38	+	}
			39	+
			40	+	if err := server.run(strings.NewReader(`{"type":"hello","id":"1","token":"synthetic-token"}` + "\n")); err != nil {
			41	+	t.Fatalf("run RPC server: %v", err)
			42	+	}
			43	+	lines := strings.Split(strings.TrimSpace(out.String()), "\n")
			44	+	if len(lines) != 2 \|\| !strings.Contains(lines[0], `"type":"response"`) \|\| !strings.Contains(lines[1], `"type":"ext_alert"`) {
			45	+	t.Fatalf("authenticated RPC frames = %q, want response followed by ext_alert", out.String())
			46	+	}
			47	+	}