test(distributed): wait for the log subscriber instead of racing it

The WebSocket log handler writes its "initial" batch before it calls
Subscribe, so a line appended the instant that batch arrives lands in the
circular buffer with no subscriber to receive it. Three backend-logs specs
append exactly there and then wait out a 5s read deadline; once a gorilla
read hits its deadline the connection is unusable, so the spec cannot retry.
`--focus='Worker WebSocket log streaming' --repeat=25` failed on attempt 17
with nothing else running, which is far too often to wire into CI.

Add BackendLogStore.SubscriberCount, resolving a model ID by the same
exact-key and replica-prefix rules Subscribe uses, and have the specs poll it
until the handler has attached. Nothing in production calls it and no
assertion is weakened; the handler's own snapshot/subscribe window is left as
it is, being a production streaming question rather than a test one.

Verified with 60 repeats of the WebSocket specs and three consecutive
--randomize-all runs of the whole distributed suite, all at
--flake-attempts 1: 239 of 240 specs pass in about 80 seconds.

Assisted-by: Claude Opus 5 [claude-code]
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

pkg/model/backend_log_store.go

@@ -344,3 +344,42 @@ func (s *BackendLogStore) Subscribe(modelID string) (chan BackendLogLine, func()
 
 	return ch, unsubscribe
 }
+
+// SubscriberCount reports how many live subscriptions exist for modelID,
+// resolving the ID with the same exact-key / replica-prefix rules as Subscribe.
+//
+// Streaming handlers send a GetLines snapshot before they call Subscribe, so a
+// line appended between those two calls reaches the buffer but no channel. A
+// caller that has to observe a line it appends itself must therefore wait for
+// the subscription to exist rather than assume the handler got there first.
+func (s *BackendLogStore) SubscriberCount(modelID string) int {
+	s.mu.RLock()
+	exact, exactOK := s.buffers[modelID]
+	var replicas []*backendLogBuffer
+	if !strings.Contains(modelID, replicaSeparator) {
+		prefix := modelID + replicaSeparator
+		for k, b := range s.buffers {
+			if strings.HasPrefix(k, prefix) {
+				replicas = append(replicas, b)
+			}
+		}
+	}
+	s.mu.RUnlock()
+
+	// Counted after releasing s.mu: no other path takes s.mu and a buffer lock
+	// together, and keeping it that way costs nothing here.
+	count := func(buf *backendLogBuffer) int {
+		buf.mu.Lock()
+		defer buf.mu.Unlock()
+		return len(buf.subscribers)
+	}
+
+	total := 0
+	if exactOK {
+		total += count(exact)
+	}
+	for _, b := range replicas {
+		total += count(b)
+	}
+	return total
+}

pkg/model/backend_log_store_test.go

@@ -76,6 +76,38 @@ var _ = Describe("BackendLogStore", func() {
 		})
 	})
 
+	Describe("SubscriberCount", func() {
+		It("reports zero before anyone subscribes and drops back after unsubscribe", func() {
+			s.AppendLine("model-a", "stderr", "preload")
+			Expect(s.SubscriberCount("model-a")).To(Equal(0))
+
+			_, unsubscribe := s.Subscribe("model-a")
+			Expect(s.SubscriberCount("model-a")).To(Equal(1))
+
+			unsubscribe()
+			Expect(s.SubscriberCount("model-a")).To(Equal(0))
+		})
+
+		// Subscribe resolves a bare model ID across every replica buffer, so the
+		// count has to follow the same rule or a caller waiting on it would give
+		// up while a perfectly good subscription was in place.
+		It("sums the replica buffers a bare model ID resolves to", func() {
+			s.AppendLine("model-a#0", "stderr", "preload-r0")
+			s.AppendLine("model-a#1", "stderr", "preload-r1")
+
+			_, unsubscribe := s.Subscribe("model-a")
+			defer unsubscribe()
+
+			Expect(s.SubscriberCount("model-a")).To(Equal(2))
+			Expect(s.SubscriberCount("model-a#0")).To(Equal(1))
+			Expect(s.SubscriberCount("model-b")).To(Equal(0))
+		})
+
+		It("returns zero for a model that has no buffer at all", func() {
+			Expect(s.SubscriberCount("never-seen")).To(Equal(0))
+		})
+	})
+
 	Describe("Subscribe", func() {
 		// Confirms the WebSocket streaming path (the live tail UI) receives
 		// lines from every replica when the caller subscribes by bare modelID.

tests/e2e/distributed/backend_logs_test.go

@@ -25,6 +25,21 @@ import (
 	"gorm.io/gorm/logger"
 )
 
+// waitForLogSubscriber blocks until the worker's WebSocket log handler has
+// registered its subscription on the store.
+//
+// The handler writes the "initial" batch first and subscribes only afterwards,
+// so a line appended the instant that batch lands is buffered but never
+// streamed, and the spec then waits out its full read deadline. Measured at
+// roughly one run in seventeen with `--repeat`, which is far too often for CI.
+// Waiting on the subscription removes the race from the spec; the handler's own
+// snapshot/subscribe window is a separate production question.
+func waitForLogSubscriber(logStore *model.BackendLogStore, modelID string) {
+	GinkgoHelper()
+	Eventually(func() int { return logStore.SubscriberCount(modelID) }, "10s", "5ms").
+		Should(BeNumerically(">", 0), "the WebSocket handler never subscribed to %q", modelID)
+}
+
 var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func() {
 
 	Context("Worker HTTP log endpoints", func() {
@@ -212,6 +227,7 @@ var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func
 			Expect(initialLines[1].Text).To(Equal("line-2"))
 
 			// Now append a new line and verify it arrives via WebSocket
+			waitForLogSubscriber(logStore, "ws-model")
 			logStore.AppendLine("ws-model", "stdout", "line-3-realtime")
 
 			conn.SetReadDeadline(time.Now().Add(5 * time.Second))
@@ -280,6 +296,7 @@ var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func
 			Expect(conn.ReadJSON(&initialMsg)).To(Succeed())
 
 			// Append line to a different model
+			waitForLogSubscriber(logStore, "ws-model")
 			logStore.AppendLine("other-model", "stdout", "should not appear")
 			// Append line to our model
 			logStore.AppendLine("ws-model", "stdout", "should appear")
@@ -475,6 +492,7 @@ var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func
 			Expect(initialLines[0].Text).To(Equal("initial line from worker"))
 
 			// Append a new line on the worker's log store
+			waitForLogSubscriber(logStore, "proxy-model")
 			logStore.AppendLine("proxy-model", "stderr", "realtime via proxy")
 
 			// Read the streamed line through the proxy