Fix runstream subscriber close race

server/internal/runstream/broker.go

@@ -10,7 +10,7 @@ import (
 const DefaultBufferSize = 64
 
 type Broker struct {
-	mu         sync.Mutex
+	mu         sync.Mutex // serializes subscriber sends and closes as well as run state
 	bufferSize int
 	runs       map[string]*runState
 }
@@ -33,21 +33,16 @@ func (b *Broker) Publish(runID string, ev connector.PromptEvent) {
 		return
 	}
 	b.mu.Lock()
+	defer b.mu.Unlock()
 	st := b.stateLocked(runID)
 	if st.closed {
-		b.mu.Unlock()
 		return
 	}
 	st.events = append(st.events, ev)
 	if len(st.events) > b.bufferSize {
 		st.events = append([]connector.PromptEvent(nil), st.events[len(st.events)-b.bufferSize:]...)
 	}
-	subs := make([]chan connector.PromptEvent, 0, len(st.subscribers))
 	for ch := range st.subscribers {
-		subs = append(subs, ch)
-	}
-	b.mu.Unlock()
-	for _, ch := range subs {
 		select {
 		case ch <- ev:
 		default:
@@ -63,30 +58,19 @@ func (b *Broker) Subscribe(ctx context.Context, runID string) <-chan connector.P
 	}
 	b.mu.Lock()
 	st := b.stateLocked(runID)
-	replay := append([]connector.PromptEvent(nil), st.events...)
-	closed := st.closed
-	if !closed {
-		st.subscribers[out] = struct{}{}
+	for _, ev := range st.events {
+		out <- ev
 	}
+	if st.closed {
+		close(out)
+		b.mu.Unlock()
+		return out
+	}
+	st.subscribers[out] = struct{}{}
 	b.mu.Unlock()
 	go func() {
-		defer func() {
-			if !closed {
-				b.unsubscribe(runID, out)
-			}
-		}()
-		for _, ev := range replay {
-			select {
-			case <-ctx.Done():
-				return
-			case out <- ev:
-			}
-		}
-		if closed {
-			close(out)
-			return
-		}
 		<-ctx.Done()
+		b.unsubscribe(runID, out)
 	}()
 	return out
 }
@@ -96,19 +80,14 @@ func (b *Broker) Finish(runID string) {
 		return
 	}
 	b.mu.Lock()
+	defer b.mu.Unlock()
 	st, ok := b.runs[runID]
 	if !ok || st.closed {
-		b.mu.Unlock()
 		return
 	}
 	st.closed = true
-	subs := make([]chan connector.PromptEvent, 0, len(st.subscribers))
 	for ch := range st.subscribers {
-		subs = append(subs, ch)
 		delete(st.subscribers, ch)
-	}
-	b.mu.Unlock()
-	for _, ch := range subs {
 		close(ch)
 	}
 }

server/internal/runstream/broker_test.go

@@ -2,6 +2,7 @@ package runstream
 
 import (
 	"context"
+	"runtime"
 	"testing"
 	"time"
 
@@ -42,6 +43,36 @@ func TestCancelCleansUp(t *testing.T) {
 	}
 }
 
+func TestPublishConcurrentWithCancelDoesNotPanic(t *testing.T) {
+	b := NewBroker(1)
+	ctx, cancel := context.WithCancel(context.Background())
+	const runID = "run-publish-cancel-race"
+	const subscriberCount = 4096
+	for range subscriberCount {
+		_ = b.Subscribe(ctx, runID)
+	}
+
+	start := make(chan struct{})
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		<-start
+		b.Publish(runID, connector.PromptEvent{Type: connector.EventDelta, Delta: "x"})
+	}()
+	close(start)
+	runtime.Gosched()
+	cancel()
+	<-done
+
+	deadline := time.Now().Add(5 * time.Second)
+	for b.SubscriberCount(runID) != 0 && time.Now().Before(deadline) {
+		runtime.Gosched()
+	}
+	if got := b.SubscriberCount(runID); got != 0 {
+		t.Fatalf("subscribers after cancel = %d, want 0", got)
+	}
+}
+
 func TestMultiSubscriber(t *testing.T) {
 	b := NewBroker(64)
 	ctx, cancel := context.WithCancel(context.Background())