git.iko.soy/crush
on main help sign in
fix(client): prevent event subscription panic on cancellation

Co-Authored-By: Charm Crush <crush@charm.land>

commited870ecb91c53a98a9114d930cf62fa9c9170993 authorChristian Rocha <christian@rocha.is> dateTue 26 May 2026, 19:39 (5 months ago) parentd14f3b1
2 files changed+153−14
compare with main
unifiedsplit
internal/client/proto.go+45−14 internal/client/proto_test.go+108
internal/client/proto.go+45−14
@@ -131,6 +131,7 @@ func (c *Client) SubscribeEvents(ctx context.Context, id string) (<-chan any, er
131131
132132 go func() {
133133 defer rsp.Body.Close()
134+ defer close(events)
134135
135136 scr := bufio.NewReader(rsp.Body)
136137 for {
@@ -139,8 +140,15 @@ func (c *Client) SubscribeEvents(ctx context.Context, id string) (<-chan any, er
139140 break
140141 }
141142 if err != nil {
143+ if ctx.Err() != nil {
144+ return
145+ }
142146 slog.Error("Reading from events stream", "error", err)
143− time.Sleep(time.Second * 2)
147+ select {
148+ case <-time.After(time.Second * 2):
149+ case <-ctx.Done():
150+ return
151+ }
144152 continue
145153 }
146154 line = bytes.TrimSpace(line)
@@ -166,43 +174,63 @@ func (c *Client) SubscribeEvents(ctx context.Context, id string) (<-chan any, er
166174 case pubsub.PayloadTypeLSPEvent:
167175 var e pubsub.Event[proto.LSPEvent]
168176 _ = json.Unmarshal(p.Payload, &e)
169− sendEvent(ctx, events, e)
177+ if !sendEvent(ctx, events, e) {
178+ return
179+ }
170180 case pubsub.PayloadTypeMCPEvent:
171181 var e pubsub.Event[proto.MCPEvent]
172182 _ = json.Unmarshal(p.Payload, &e)
173− sendEvent(ctx, events, e)
183+ if !sendEvent(ctx, events, e) {
184+ return
185+ }
174186 case pubsub.PayloadTypePermissionRequest:
175187 var e pubsub.Event[proto.PermissionRequest]
176188 _ = json.Unmarshal(p.Payload, &e)
177− sendEvent(ctx, events, e)
189+ if !sendEvent(ctx, events, e) {
190+ return
191+ }
178192 case pubsub.PayloadTypePermissionNotification:
179193 var e pubsub.Event[proto.PermissionNotification]
180194 _ = json.Unmarshal(p.Payload, &e)
181− sendEvent(ctx, events, e)
195+ if !sendEvent(ctx, events, e) {
196+ return
197+ }
182198 case pubsub.PayloadTypeMessage:
183199 var e pubsub.Event[proto.Message]
184200 _ = json.Unmarshal(p.Payload, &e)
185− sendEvent(ctx, events, e)
201+ if !sendEvent(ctx, events, e) {
202+ return
203+ }
186204 case pubsub.PayloadTypeSession:
187205 var e pubsub.Event[proto.Session]
188206 _ = json.Unmarshal(p.Payload, &e)
189− sendEvent(ctx, events, e)
207+ if !sendEvent(ctx, events, e) {
208+ return
209+ }
190210 case pubsub.PayloadTypeFile:
191211 var e pubsub.Event[proto.File]
192212 _ = json.Unmarshal(p.Payload, &e)
193− sendEvent(ctx, events, e)
213+ if !sendEvent(ctx, events, e) {
214+ return
215+ }
194216 case pubsub.PayloadTypeAgentEvent:
195217 var e pubsub.Event[proto.AgentEvent]
196218 _ = json.Unmarshal(p.Payload, &e)
197− sendEvent(ctx, events, e)
219+ if !sendEvent(ctx, events, e) {
220+ return
221+ }
198222 case pubsub.PayloadTypeConfigChanged:
199223 var e pubsub.Event[proto.ConfigChanged]
200224 _ = json.Unmarshal(p.Payload, &e)
201− sendEvent(ctx, events, e)
225+ if !sendEvent(ctx, events, e) {
226+ return
227+ }
202228 case pubsub.PayloadTypeSkillsEvent:
203229 var e pubsub.Event[proto.SkillsEvent]
204230 _ = json.Unmarshal(p.Payload, &e)
205− sendEvent(ctx, events, e)
231+ if !sendEvent(ctx, events, e) {
232+ return
233+ }
206234 default:
207235 slog.Warn("Unknown event type", "type", p.Type)
208236 continue
@@ -213,12 +241,15 @@ func (c *Client) SubscribeEvents(ctx context.Context, id string) (<-chan any, er
213241 return events, nil
214242}
215243
216−func sendEvent(ctx context.Context, evc chan any, ev any) {
244+func sendEvent(ctx context.Context, evc chan any, ev any) bool {
245+ if ctx.Err() != nil {
246+ return false
247+ }
217248 select {
218249 case evc <- ev:
250+ return true
219251 case <-ctx.Done():
220− close(evc)
221− return
252+ return false
222253 }
223254}
224255
internal/client/proto_test.gonew file+108
@@ -0,0 +1,108 @@
1+package client
2+
3+import (
4+ "context"
5+ "encoding/json"
6+ "fmt"
7+ "net/http"
8+ "net/http/httptest"
9+ "testing"
10+ "time"
11+
12+ "github.com/charmbracelet/crush/internal/proto"
13+ "github.com/charmbracelet/crush/internal/pubsub"
14+ "github.com/stretchr/testify/require"
15+)
16+
17+func TestSendEventAfterContextCancelIsIdempotent(t *testing.T) {
18+ t.Parallel()
19+
20+ ctx, cancel := context.WithCancel(context.Background())
21+ cancel()
22+
23+ events := make(chan any, 1)
24+ require.False(t, sendEvent(ctx, events, "one"))
25+ require.False(t, sendEvent(ctx, events, "two"))
26+
27+ select {
28+ case ev := <-events:
29+ require.Failf(t, "unexpected event", "event: %v", ev)
30+ default:
31+ }
32+}
33+
34+func TestSubscribeEventsContextCancelClosesEvents(t *testing.T) {
35+ t.Parallel()
36+
37+ payload := marshalSSEPayload(t)
38+ firstEventSent := make(chan struct{})
39+ writeSecondEvent := make(chan struct{})
40+
41+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
42+ w.Header().Set("Content-Type", "text/event-stream")
43+ flusher, ok := w.(http.Flusher)
44+ require.True(t, ok)
45+
46+ _, err := fmt.Fprintf(w, "data: %s\n\n", payload)
47+ require.NoError(t, err)
48+ flusher.Flush()
49+ close(firstEventSent)
50+
51+ select {
52+ case <-writeSecondEvent:
53+ case <-time.After(5 * time.Second):
54+ return
55+ }
56+ _, _ = fmt.Fprintf(w, "data: %s\n\n", payload)
57+ flusher.Flush()
58+ }))
59+ defer srv.Close()
60+
61+ ctx, cancel := context.WithCancel(context.Background())
62+ defer cancel()
63+
64+ c := captureClient(t, srv)
65+ events, err := c.SubscribeEvents(ctx, "ws1")
66+ require.NoError(t, err)
67+
68+ select {
69+ case <-firstEventSent:
70+ case <-time.After(5 * time.Second):
71+ require.Fail(t, "timed out waiting for server event")
72+ }
73+
74+ select {
75+ case <-events:
76+ case <-time.After(5 * time.Second):
77+ require.Fail(t, "timed out waiting for first event")
78+ }
79+
80+ cancel()
81+ close(writeSecondEvent)
82+
83+ select {
84+ case _, ok := <-events:
85+ require.False(t, ok)
86+ case <-time.After(5 * time.Second):
87+ require.Fail(t, "timed out waiting for event channel close")
88+ }
89+}
90+
91+func marshalSSEPayload(t *testing.T) []byte {
92+ t.Helper()
93+
94+ eventPayload, err := json.Marshal(pubsub.Event[proto.AgentEvent]{
95+ Type: pubsub.CreatedEvent,
96+ Payload: proto.AgentEvent{
97+ Type: proto.AgentEventTypeResponse,
98+ },
99+ })
100+ require.NoError(t, err)
101+
102+ payload, err := json.Marshal(pubsub.Payload{
103+ Type: pubsub.PayloadTypeAgentEvent,
104+ Payload: eventPayload,
105+ })
106+ require.NoError(t, err)
107+ return payload
108+}