git.iko.soy/crush
on main help sign in
fix(client): prevent event subscription panic on cancellation
ed870ecChristian Rocha
internal/client/proto.go+45−14
compare with main
unifiedsplit
@@ -131,6 +131,7 @@ func (c *Client) SubscribeEvents(ctx context.Context, id string) (<-chan any, er
131
131
132 go func() {
132 go func() {
133 defer rsp.Body.Close()
133 defer rsp.Body.Close()
134+ defer close(events)
134
135
135 scr := bufio.NewReader(rsp.Body)
136 scr := bufio.NewReader(rsp.Body)
136 for {
137 for {
@@ -139,8 +140,15 @@ func (c *Client) SubscribeEvents(ctx context.Context, id string) (<-chan any, er
139 break
140 break
140 }
141 }
141 if err != nil {
142 if err != nil {
143+ if ctx.Err() != nil {
144+ return
145+ }
142 slog.Error("Reading from events stream", "error", err)
146 slog.Error("Reading from events stream", "error", err)
147+ select {
143− time.Sleep(time.Second * 2)
148+ case <-time.After(time.Second * 2):
149+ case <-ctx.Done():
150+ return
151+ }
144 continue
152 continue
145 }
153 }
146 line = bytes.TrimSpace(line)
154 line = bytes.TrimSpace(line)
@@ -166,43 +174,63 @@ func (c *Client) SubscribeEvents(ctx context.Context, id string) (<-chan any, er
166 case pubsub.PayloadTypeLSPEvent:
174 case pubsub.PayloadTypeLSPEvent:
167 var e pubsub.Event[proto.LSPEvent]
175 var e pubsub.Event[proto.LSPEvent]
168 _ = json.Unmarshal(p.Payload, &e)
176 _ = json.Unmarshal(p.Payload, &e)
169− sendEvent(ctx, events, e)
177+ if !sendEvent(ctx, events, e) {
178+ return
179+ }
170 case pubsub.PayloadTypeMCPEvent:
180 case pubsub.PayloadTypeMCPEvent:
171 var e pubsub.Event[proto.MCPEvent]
181 var e pubsub.Event[proto.MCPEvent]
172 _ = json.Unmarshal(p.Payload, &e)
182 _ = json.Unmarshal(p.Payload, &e)
173− sendEvent(ctx, events, e)
183+ if !sendEvent(ctx, events, e) {
184+ return
185+ }
174 case pubsub.PayloadTypePermissionRequest:
186 case pubsub.PayloadTypePermissionRequest:
175 var e pubsub.Event[proto.PermissionRequest]
187 var e pubsub.Event[proto.PermissionRequest]
176 _ = json.Unmarshal(p.Payload, &e)
188 _ = json.Unmarshal(p.Payload, &e)
177− sendEvent(ctx, events, e)
189+ if !sendEvent(ctx, events, e) {
190+ return
191+ }
178 case pubsub.PayloadTypePermissionNotification:
192 case pubsub.PayloadTypePermissionNotification:
179 var e pubsub.Event[proto.PermissionNotification]
193 var e pubsub.Event[proto.PermissionNotification]
180 _ = json.Unmarshal(p.Payload, &e)
194 _ = json.Unmarshal(p.Payload, &e)
181− sendEvent(ctx, events, e)
195+ if !sendEvent(ctx, events, e) {
196+ return
197+ }
182 case pubsub.PayloadTypeMessage:
198 case pubsub.PayloadTypeMessage:
183 var e pubsub.Event[proto.Message]
199 var e pubsub.Event[proto.Message]
184 _ = json.Unmarshal(p.Payload, &e)
200 _ = json.Unmarshal(p.Payload, &e)
185− sendEvent(ctx, events, e)
201+ if !sendEvent(ctx, events, e) {
202+ return
203+ }
186 case pubsub.PayloadTypeSession:
204 case pubsub.PayloadTypeSession:
187 var e pubsub.Event[proto.Session]
205 var e pubsub.Event[proto.Session]
188 _ = json.Unmarshal(p.Payload, &e)
206 _ = json.Unmarshal(p.Payload, &e)
189− sendEvent(ctx, events, e)
207+ if !sendEvent(ctx, events, e) {
208+ return
209+ }
190 case pubsub.PayloadTypeFile:
210 case pubsub.PayloadTypeFile:
191 var e pubsub.Event[proto.File]
211 var e pubsub.Event[proto.File]
192 _ = json.Unmarshal(p.Payload, &e)
212 _ = json.Unmarshal(p.Payload, &e)
193− sendEvent(ctx, events, e)
213+ if !sendEvent(ctx, events, e) {
214+ return
215+ }
194 case pubsub.PayloadTypeAgentEvent:
216 case pubsub.PayloadTypeAgentEvent:
195 var e pubsub.Event[proto.AgentEvent]
217 var e pubsub.Event[proto.AgentEvent]
196 _ = json.Unmarshal(p.Payload, &e)
218 _ = json.Unmarshal(p.Payload, &e)
197− sendEvent(ctx, events, e)
219+ if !sendEvent(ctx, events, e) {
220+ return
221+ }
198 case pubsub.PayloadTypeConfigChanged:
222 case pubsub.PayloadTypeConfigChanged:
199 var e pubsub.Event[proto.ConfigChanged]
223 var e pubsub.Event[proto.ConfigChanged]
200 _ = json.Unmarshal(p.Payload, &e)
224 _ = json.Unmarshal(p.Payload, &e)
201− sendEvent(ctx, events, e)
225+ if !sendEvent(ctx, events, e) {
226+ return
227+ }
202 case pubsub.PayloadTypeSkillsEvent:
228 case pubsub.PayloadTypeSkillsEvent:
203 var e pubsub.Event[proto.SkillsEvent]
229 var e pubsub.Event[proto.SkillsEvent]
204 _ = json.Unmarshal(p.Payload, &e)
230 _ = json.Unmarshal(p.Payload, &e)
205− sendEvent(ctx, events, e)
231+ if !sendEvent(ctx, events, e) {
232+ return
233+ }
206 default:
234 default:
207 slog.Warn("Unknown event type", "type", p.Type)
235 slog.Warn("Unknown event type", "type", p.Type)
208 continue
236 continue
@@ -213,12 +241,15 @@ func (c *Client) SubscribeEvents(ctx context.Context, id string) (<-chan any, er
213 return events, nil
241 return events, nil
214}
242}
215
243
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+ }
217 select {
248 select {
218 case evc <- ev:
249 case evc <- ev:
250+ return true
219 case <-ctx.Done():
251 case <-ctx.Done():
220− close(evc)
221− return
252+ return false
222 }
253 }
223}
254}
224
255
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+}