fix(chatd): harden cross-replica relay for chat stream parts (#22533)

## Problem Subscribers connecting to a different replica than the one running the chat see full messages appear but no streaming partials (`message_part` events). The relay mechanism that forwards ephemeral parts across replicas had several bugs. ## Root Causes 1. **`openRelay()` blocked the event loop** — The WebSocket dial (TCP + TLS + HTTP upgrade) to the worker replica ran synchronously inside the select loop. While dialing, no events could be processed, channels filled up, and parts were silently dropped. 2. **Relay drops were permanent** — When the relay WebSocket closed mid-stream, `relayParts` was set to nil and never reopened. No status notification would re-trigger it since the chat was still running on the same worker. 3. **`drainInitial` snapshot race** — The `default` case in the initial drain loop caused the snapshot to be empty if the remote hadn't flushed data yet (common immediately after WebSocket connect). 4. **Duplicate event delivery** — The `preloaded` slice caused snapshot events to be sent both in the return value and re-sent through the channel goroutine. ## Fixes ### `coderd/chatd/chatd.go` (Subscribe method) - **Async relay dial**: `openRelayAsync()` spawns a goroutine to dial the remote replica. The result (channel + cancel func) is delivered on a `relayReadyCh` channel that the select loop reads without blocking. - **Relay reconnection**: When the relay channel closes, a 500ms timer fires. The handler re-checks chat status from the DB and reopens the relay if the chat is still running on a remote worker. - **Snapshot parts via channel**: Relay snapshot + live parts are wrapped into a single channel so they flow through the same path, avoiding races with the select loop. ### `enterprise/coderd/chats.go` (newRemotePartsProvider) - **Timer-based drain**: Replaced `default` with a 1-second timer. After the first event, `Reset(0)` switches to non-blocking drain for remaining buffered events. - **Remove preloaded duplication**: The goroutine now only forwards new events; snapshot events are returned to the caller directly. ## Testing All existing tests pass: - `TestInterruptChatBroadcastsStatusAcrossInstances` - `TestSubscribeSnapshotIncludesStatusEvent` - `TestSubscribeNoPubsubNoDuplicateMessageParts` - `TestSubscribeAfterMessageID` - `TestChatStreamRelay/RelayMessagePartsAcrossReplicas`
2026-09-24 15:04:27 +08:00 · 2026-03-02 19:57:13 -05:00
parent b8a74a4fcb
commit b7a7683ac0
4 changed files with 584 additions and 34 deletions
@@ -1094,42 +1094,92 @@ func (p *Server) Subscribe(
 		allCancels = append(allCancels, relayCancel)
 	}

-	// Helper to close relay
+	// Channel for async relay establishment.
+	type relayResult struct {
+		parts  <-chan codersdk.ChatStreamEvent
+		cancel func()
+	}
+	relayReadyCh := make(chan relayResult, 1)
+
+	// Reconnect timer state.
+	var reconnectTimer *time.Timer
+	var reconnectCh <-chan time.Time
+
+	// Helper to close relay and stop any pending reconnect timer.
 	closeRelay := func() {
 		if relayCancel != nil {
 			relayCancel()
 			relayCancel = nil
 		}
 		relayParts = nil
+		if reconnectTimer != nil {
+			reconnectTimer.Stop()
+			reconnectTimer = nil
+			reconnectCh = nil
+		}
 	}

-	// Helper to open relay to a worker
-	openRelay := func(workerID uuid.UUID) {
+	// openRelayAsync dials the remote replica in a background
+	// goroutine and delivers the result on relayReadyCh so the
+	// main select loop is never blocked by network I/O.
+	openRelayAsync := func(workerID uuid.UUID) {
 		if p.remotePartsProvider == nil {
 			return
 		}
 		closeRelay()
-		snapshot, parts, cancel, err := p.remotePartsProvider(mergedCtx, chatID, workerID, requestHeader)
-		if err != nil {
-			p.logger.Warn(mergedCtx, "failed to open relay for message parts",
-				slog.F("chat_id", chatID),
-				slog.F("worker_id", workerID),
-				slog.Error(err),
-			)
+		go func() {
+			snapshot, parts, cancel, err := p.remotePartsProvider(mergedCtx, chatID, workerID, requestHeader)
+			if err != nil {
+				p.logger.Warn(mergedCtx, "failed to open relay for message parts",
+					slog.F("chat_id", chatID),
+					slog.F("worker_id", workerID),
+					slog.Error(err),
+				)
+				return
+			}
+			// Wrap the relay channel so snapshot parts are
+			// delivered through the same channel as live parts.
+			wrappedParts := make(chan codersdk.ChatStreamEvent, 128)
+			go func() {
+				defer close(wrappedParts)
+				for _, event := range snapshot {
+					if event.Type == codersdk.ChatStreamEventTypeMessagePart {
+						select {
+						case wrappedParts <- event:
+						case <-mergedCtx.Done():
+							cancel()
+							return
+						}
+					}
+				}
+				for event := range parts {
+					select {
+					case wrappedParts <- event:
+					case <-mergedCtx.Done():
+						return
+					}
+				}
+			}()
+			select {
+			case relayReadyCh <- relayResult{parts: wrappedParts, cancel: cancel}:
+			case <-mergedCtx.Done():
+				cancel()
+			}
+		}()
+	}
+
+	// scheduleRelayReconnect arms a short timer so the select
+	// loop can re-check chat status and reopen the relay without
+	// spinning in a tight loop.
+	scheduleRelayReconnect := func() {
+		if p.remotePartsProvider == nil {
 			return
 		}
-		relayParts = parts
-		relayCancel = cancel
-		// Send relay snapshot message_parts
-		for _, event := range snapshot {
-			if event.Type == codersdk.ChatStreamEventTypeMessagePart {
-				select {
-				case <-mergedCtx.Done():
-					return
-				case mergedEvents <- event:
-				}
-			}
+		if reconnectTimer != nil {
+			reconnectTimer.Stop()
 		}
+		reconnectTimer = time.NewTimer(500 * time.Millisecond)
+		reconnectCh = reconnectTimer.C
 	}

 	//nolint:nestif
@@ -1195,6 +1245,21 @@ func (p *Server) Subscribe(
 						},
 					}
 					return
+				case result := <-relayReadyCh:
+					// An async relay dial completed; swap in the
+					// new relay channel.
+					closeRelay()
+					relayParts = result.parts
+					relayCancel = result.cancel
+				case <-reconnectCh:
+					reconnectCh = nil
+					// Re-check whether the chat is still running
+					// on a remote worker before reconnecting.
+					currentChat, chatErr := p.db.GetChatByID(mergedCtx, chatID)
+					if chatErr == nil && currentChat.Status == database.ChatStatusRunning &&
+						currentChat.WorkerID.Valid && currentChat.WorkerID.UUID != p.workerID {
+						openRelayAsync(currentChat.WorkerID.UUID)
+					}
 				case notify := <-notifications:
 					// Handle different notification types
 					if notify.AfterMessageID > 0 {
@@ -1230,11 +1295,11 @@ func (p *Server) Subscribe(
 							Status: &codersdk.ChatStreamStatus{Status: codersdk.ChatStatus(status)},
 						}:
 						}
-						// Manage relay lifecycle based on status
+						// Manage relay lifecycle based on status.
 						if status == database.ChatStatusRunning && notify.WorkerID != "" {
 							workerID, err := uuid.Parse(notify.WorkerID)
 							if err == nil && workerID != p.workerID {
-								openRelay(workerID)
+								openRelayAsync(workerID)
 							} else if workerID == p.workerID {
 								closeRelay()
 							}
@@ -1285,6 +1350,8 @@ func (p *Server) Subscribe(
 				case event, ok := <-relayPartsCh:
 					if !ok {
 						relayParts = nil
+						// Schedule reconnection instead of giving up.
+						scheduleRelayReconnect()
 						continue
 					}
 					// Only forward message_part events from relay (durable events come via pubsub)
@@ -1320,6 +1387,9 @@ func (p *Server) Subscribe(
 				cancelFn()
 			}
 		}
+		if reconnectTimer != nil {
+			reconnectTimer.Stop()
+		}
 	}

 	return initialSnapshot, mergedEvents, cancel, true
@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"net/http"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -27,6 +28,7 @@ import (
 	"github.com/coder/coder/v2/coderd/database/dbgen"
 	"github.com/coder/coder/v2/coderd/database/dbtestutil"
 	dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub"
+	coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
 	"github.com/coder/coder/v2/coderd/util/slice"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/provisioner/echo"
@@ -977,6 +979,30 @@ func newTestServer(
 	return server
 }

+func newTestServerWithRelay(
+	t *testing.T,
+	db database.Store,
+	ps dbpubsub.Pubsub,
+	replicaID uuid.UUID,
+	provider chatd.RemotePartsProvider,
+) *chatd.Server {
+	t.Helper()
+
+	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
+	server := chatd.New(chatd.Config{
+		Logger:                     logger,
+		Database:                   db,
+		ReplicaID:                  replicaID,
+		Pubsub:                     ps,
+		RemotePartsProvider:        provider,
+		PendingChatAcquireInterval: testutil.WaitSuperLong,
+	})
+	t.Cleanup(func() {
+		require.NoError(t, server.Close())
+	})
+	return server
+}
+
 func seedChatDependencies(
 	ctx context.Context,
 	t *testing.T,
@@ -1033,6 +1059,293 @@ func setOpenAIProviderBaseURL(
 	require.NoError(t, err)
 }

+func TestSubscribeRelayReconnectsOnDrop(t *testing.T) {
+	t.Parallel()
+
+	db, ps := dbtestutil.NewDB(t)
+	workerID := uuid.New()
+	subscriberID := uuid.New()
+
+	var callCount atomic.Int32
+
+	provider := func(ctx context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) (
+		[]codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error,
+	) {
+		call := callCount.Add(1)
+		ch := make(chan codersdk.ChatStreamEvent, 10)
+		if call == 1 {
+			// First relay: send a part then close to simulate a drop.
+			ch <- codersdk.ChatStreamEvent{
+				Type: codersdk.ChatStreamEventTypeMessagePart,
+				MessagePart: &codersdk.ChatStreamMessagePart{
+					Role: "assistant",
+					Part: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: "first-relay"},
+				},
+			}
+			close(ch)
+		} else {
+			// Second relay: send a different part, keep open.
+			ch <- codersdk.ChatStreamEvent{
+				Type: codersdk.ChatStreamEventTypeMessagePart,
+				MessagePart: &codersdk.ChatStreamMessagePart{
+					Role: "assistant",
+					Part: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: "second-relay"},
+				},
+			}
+			// Don't close — keep alive so the subscriber stays connected.
+		}
+		return nil, ch, func() {}, nil
+	}
+
+	subscriber := newTestServerWithRelay(t, db, ps, subscriberID, provider)
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	user, model := seedChatDependencies(ctx, t, db)
+
+	// Create a chat and mark it as running on a remote worker.
+	chat, err := subscriber.CreateChat(ctx, chatd.CreateOptions{
+		OwnerID:            user.ID,
+		Title:              "relay-reconnect",
+		ModelConfigID:      model.ID,
+		InitialUserContent: []fantasy.Content{fantasy.TextContent{Text: "hello"}},
+	})
+	require.NoError(t, err)
+
+	chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{
+		ID:          chat.ID,
+		Status:      database.ChatStatusRunning,
+		WorkerID:    uuid.NullUUID{UUID: workerID, Valid: true},
+		StartedAt:   sql.NullTime{Time: time.Now(), Valid: true},
+		HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true},
+	})
+	require.NoError(t, err)
+
+	_, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
+	require.True(t, ok)
+	t.Cleanup(cancel)
+
+	// Should get the first relay part.
+	require.Eventually(t, func() bool {
+		select {
+		case event := <-events:
+			if event.Type == codersdk.ChatStreamEventTypeMessagePart &&
+				event.MessagePart != nil &&
+				event.MessagePart.Part.Text == "first-relay" {
+				return true
+			}
+			return false
+		default:
+			return false
+		}
+	}, testutil.WaitMedium, testutil.IntervalFast)
+
+	// After the first relay closes, a reconnection should happen and
+	// deliver the second relay part.
+	require.Eventually(t, func() bool {
+		select {
+		case event := <-events:
+			if event.Type == codersdk.ChatStreamEventTypeMessagePart &&
+				event.MessagePart != nil &&
+				event.MessagePart.Part.Text == "second-relay" {
+				return true
+			}
+			return false
+		default:
+			return false
+		}
+	}, testutil.WaitMedium, testutil.IntervalFast)
+
+	require.GreaterOrEqual(t, int(callCount.Load()), 2)
+}
+
+func TestSubscribeRelayAsyncDoesNotBlock(t *testing.T) {
+	t.Parallel()
+
+	db, ps := dbtestutil.NewDB(t)
+	workerID := uuid.New()
+	subscriberID := uuid.New()
+
+	dialStarted := make(chan struct{})
+	dialContinue := make(chan struct{})
+
+	provider := func(ctx context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) (
+		[]codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error,
+	) {
+		// Signal that the dial has started, then block until released.
+		select {
+		case <-dialStarted:
+		default:
+			close(dialStarted)
+		}
+		select {
+		case <-dialContinue:
+		case <-ctx.Done():
+			return nil, nil, nil, ctx.Err()
+		}
+		ch := make(chan codersdk.ChatStreamEvent, 10)
+		return nil, ch, func() {}, nil
+	}
+
+	subscriber := newTestServerWithRelay(t, db, ps, subscriberID, provider)
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	user, model := seedChatDependencies(ctx, t, db)
+
+	// Create a chat in pending status.
+	chat, err := subscriber.CreateChat(ctx, chatd.CreateOptions{
+		OwnerID:            user.ID,
+		Title:              "relay-async-nonblock",
+		ModelConfigID:      model.ID,
+		InitialUserContent: []fantasy.Content{fantasy.TextContent{Text: "hello"}},
+	})
+	require.NoError(t, err)
+
+	// Subscribe before the chat is marked running so the relay opens
+	// via pubsub notification (openRelayAsync path).
+	_, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
+	require.True(t, ok)
+	t.Cleanup(cancel)
+
+	// Now mark the chat as running on a remote worker. This publishes
+	// a status notification which triggers openRelayAsync on the
+	// subscriber.
+	notify := coderdpubsub.ChatStreamNotifyMessage{
+		Status:   string(database.ChatStatusRunning),
+		WorkerID: workerID.String(),
+	}
+	payload, err := json.Marshal(notify)
+	require.NoError(t, err)
+	err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chat.ID), payload)
+	require.NoError(t, err)
+
+	// Wait for the relay dial to actually start (blocking in the
+	// provider).
+	select {
+	case <-dialStarted:
+	case <-ctx.Done():
+		t.Fatal("timed out waiting for relay dial to start")
+	}
+
+	// While the relay is still dialing (provider is blocked), publish
+	// another status change. If openRelayAsync blocked the select loop
+	// this event would never arrive.
+	statusNotify := coderdpubsub.ChatStreamNotifyMessage{
+		Status: string(database.ChatStatusWaiting),
+	}
+	statusPayload, err := json.Marshal(statusNotify)
+	require.NoError(t, err)
+	err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chat.ID), statusPayload)
+	require.NoError(t, err)
+
+	// The waiting status event should arrive promptly despite the
+	// relay still dialing.
+	require.Eventually(t, func() bool {
+		select {
+		case event := <-events:
+			return event.Type == codersdk.ChatStreamEventTypeStatus &&
+				event.Status != nil &&
+				event.Status.Status == codersdk.ChatStatusWaiting
+		default:
+			return false
+		}
+	}, testutil.WaitShort, testutil.IntervalFast)
+
+	// Unblock the relay dial so the test can clean up.
+	close(dialContinue)
+}
+
+func TestSubscribeRelaySnapshotDelivered(t *testing.T) {
+	t.Parallel()
+
+	db, ps := dbtestutil.NewDB(t)
+	workerID := uuid.New()
+	subscriberID := uuid.New()
+
+	provider := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) (
+		[]codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error,
+	) {
+		// Return a non-empty snapshot with two parts.
+		snapshot := []codersdk.ChatStreamEvent{
+			{
+				Type: codersdk.ChatStreamEventTypeMessagePart,
+				MessagePart: &codersdk.ChatStreamMessagePart{
+					Role: "assistant",
+					Part: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: "snap-one"},
+				},
+			},
+			{
+				Type: codersdk.ChatStreamEventTypeMessagePart,
+				MessagePart: &codersdk.ChatStreamMessagePart{
+					Role: "assistant",
+					Part: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: "snap-two"},
+				},
+			},
+		}
+		ch := make(chan codersdk.ChatStreamEvent, 10)
+		// Also send a live part after the snapshot.
+		ch <- codersdk.ChatStreamEvent{
+			Type: codersdk.ChatStreamEventTypeMessagePart,
+			MessagePart: &codersdk.ChatStreamMessagePart{
+				Role: "assistant",
+				Part: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: "live-part"},
+			},
+		}
+		return snapshot, ch, func() {}, nil
+	}
+
+	subscriber := newTestServerWithRelay(t, db, ps, subscriberID, provider)
+
+	ctx := testutil.Context(t, testutil.WaitLong)
+	user, model := seedChatDependencies(ctx, t, db)
+
+	// Create a chat already running on a remote worker.
+	chat, err := subscriber.CreateChat(ctx, chatd.CreateOptions{
+		OwnerID:            user.ID,
+		Title:              "relay-snapshot",
+		ModelConfigID:      model.ID,
+		InitialUserContent: []fantasy.Content{fantasy.TextContent{Text: "hello"}},
+	})
+	require.NoError(t, err)
+
+	_, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{
+		ID:          chat.ID,
+		Status:      database.ChatStatusRunning,
+		WorkerID:    uuid.NullUUID{UUID: workerID, Valid: true},
+		StartedAt:   sql.NullTime{Time: time.Now(), Valid: true},
+		HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true},
+	})
+	require.NoError(t, err)
+
+	initialSnapshot, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0)
+	require.True(t, ok)
+	t.Cleanup(cancel)
+
+	// The initial snapshot should contain the two relay snapshot parts.
+	var snapshotTexts []string
+	for _, event := range initialSnapshot {
+		if event.Type == codersdk.ChatStreamEventTypeMessagePart && event.MessagePart != nil {
+			snapshotTexts = append(snapshotTexts, event.MessagePart.Part.Text)
+		}
+	}
+	require.Contains(t, snapshotTexts, "snap-one")
+	require.Contains(t, snapshotTexts, "snap-two")
+
+	// The live part should arrive on the events channel.
+	require.Eventually(t, func() bool {
+		select {
+		case event := <-events:
+			if event.Type == codersdk.ChatStreamEventTypeMessagePart &&
+				event.MessagePart != nil &&
+				event.MessagePart.Part.Text == "live-part" {
+				return true
+			}
+			return false
+		default:
+			return false
+		}
+	}, testutil.WaitMedium, testutil.IntervalFast)
+}
+
 func TestCloseDuringShutdownContextCanceledShouldRetryOnNewReplica(t *testing.T) {
 	t.Parallel()

@@ -4,6 +4,7 @@ import (
 	"context"
 	"net/http"
 	"net/url"
+	"time"

 	"github.com/google/uuid"
 	"golang.org/x/xerrors"
@@ -62,7 +63,14 @@ func newRemotePartsProvider(
 		}

 		snapshot := make([]codersdk.ChatStreamEvent, 0, 100)
-		preloaded := make([]codersdk.ChatStreamEvent, 0, 100)
+
+		// Wait briefly for the first event to handle the common
+		// case where the remote side has buffered parts but hasn't
+		// flushed them to the WebSocket yet.
+		const drainTimeout = time.Second
+		drainTimer := time.NewTimer(drainTimeout)
+		defer drainTimer.Stop()
+
 	drainInitial:
 		for len(snapshot) < cap(snapshot) {
 			select {
@@ -78,8 +86,11 @@ func newRemotePartsProvider(
 					continue
 				}
 				snapshot = append(snapshot, event)
-				preloaded = append(preloaded, event)
-			default:
+				// After getting the first event, switch to
+				// non-blocking drain for remaining buffered events.
+				drainTimer.Stop()
+				drainTimer.Reset(0)
+			case <-drainTimer.C:
 				break drainInitial
 			}
 		}
@@ -93,14 +104,8 @@ func newRemotePartsProvider(
 				_ = sourceStream.Close()
 			}()

-			for _, event := range preloaded {
-				select {
-				case events <- event:
-				case <-relayCtx.Done():
-					return
-				}
-			}
-
+			// No need to re-send snapshot events — they're
+			// returned to the caller directly.
 			for {
 				select {
 				case <-relayCtx.Done():
@@ -165,6 +165,168 @@ func TestChatStreamRelay(t *testing.T) {

 		close(streamingChunks)
 	})
+
+	t.Run("RelaySnapshotIncludesBufferedParts", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitLong)
+
+		db, pubsub := dbtestutil.NewDB(t)
+		firstClient, _ := coderdenttest.New(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{
+				Database: db,
+				Pubsub:   pubsub,
+			},
+			LicenseOptions: &coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureHighAvailability: 1,
+				},
+			},
+		})
+
+		secondClient, _ := coderdenttest.New(t, &coderdenttest.Options{
+			Options: &coderdtest.Options{
+				Database: db,
+				Pubsub:   pubsub,
+			},
+			DontAddLicense:   true,
+			DontAddFirstUser: true,
+		})
+		secondClient.SetSessionToken(firstClient.SessionToken())
+
+		// Verify we have two replicas.
+		replicas, err := secondClient.Replicas(ctx)
+		require.NoError(t, err)
+		require.Len(t, replicas, 2)
+		firstReplicaID := replicaIDForClientURL(t, firstClient.URL, replicas)
+		secondReplicaID := replicaIDForClientURL(t, secondClient.URL, replicas)
+
+		streamingChunks := make(chan chattest.OpenAIChunk, 8)
+		chatStreamStarted := make(chan struct{}, 1)
+		openai := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
+			if req.Stream {
+				select {
+				case chatStreamStarted <- struct{}{}:
+				default:
+				}
+				return chattest.OpenAIResponse{StreamingChunks: streamingChunks}
+			}
+			return chattest.OpenAINonStreamingResponse("ok")
+		})
+
+		//nolint:gocritic // Test uses owner client to configure chat providers.
+		provider, err := firstClient.CreateChatProvider(ctx, codersdk.CreateChatProviderConfigRequest{
+			Provider:    "openai",
+			DisplayName: "OpenAI",
+			APIKey:      "test",
+			BaseURL:     openai,
+		})
+		require.NoError(t, err)
+
+		model, err := firstClient.CreateChatModelConfig(ctx, codersdk.CreateChatModelConfigRequest{
+			Provider:             provider.Provider,
+			Model:                "gpt-4",
+			DisplayName:          "GPT-4",
+			ContextLimit:         &[]int64{1000}[0],
+			CompressionThreshold: &[]int32{70}[0],
+		})
+		require.NoError(t, err)
+
+		// Create a chat on the first replica.
+		chat, err := firstClient.CreateChat(ctx, codersdk.CreateChatRequest{
+			Content: []codersdk.ChatInputPart{{
+				Type: codersdk.ChatInputPartTypeText,
+				Text: "Test chat for buffered relay",
+			}},
+			ModelConfigID: &model.ID,
+		})
+		require.NoError(t, err)
+		require.Equal(t, codersdk.ChatStatusPending, chat.Status)
+
+		var runningChat database.Chat
+		require.Eventually(t, func() bool {
+			current, getErr := db.GetChatByID(ctx, chat.ID)
+			if getErr != nil {
+				return false
+			}
+			if current.Status != database.ChatStatusRunning || !current.WorkerID.Valid {
+				return false
+			}
+			runningChat = current
+			return true
+		}, testutil.WaitLong, testutil.IntervalFast)
+
+		var localClient *codersdk.Client
+		var relayClient *codersdk.Client
+		switch runningChat.WorkerID.UUID {
+		case firstReplicaID:
+			localClient = firstClient
+			relayClient = secondClient
+		case secondReplicaID:
+			localClient = secondClient
+			relayClient = firstClient
+		default:
+			require.FailNowf(
+				t,
+				"worker replica was not recognized",
+				"worker %s was not one of %s or %s",
+				runningChat.WorkerID.UUID,
+				firstReplicaID,
+				secondReplicaID,
+			)
+		}
+
+		// Subscribe on the local (worker) replica so the stream is
+		// consumed and chunks flow through the pipeline.
+		localEvents, localStream, err := localClient.StreamChat(ctx, chat.ID)
+		require.NoError(t, err)
+		defer localStream.Close()
+
+		// Wait for the OpenAI handler to start serving the stream.
+		select {
+		case <-chatStreamStarted:
+		case <-ctx.Done():
+			require.FailNowf(
+				t,
+				"timed out waiting for OpenAI stream request",
+				"chat stream request did not start before context deadline: %v",
+				ctx.Err(),
+			)
+		}
+
+		// Send multiple chunks BEFORE the relay subscriber connects.
+		// This is the key difference from the existing test: we
+		// buffer several parts so the drainInitial timer in
+		// newRemotePartsProvider must collect them all.
+		bufferedTexts := []string{"buffered-one", "buffered-two", "buffered-three"}
+		for _, text := range bufferedTexts {
+			streamingChunks <- chattest.OpenAITextChunks(text)[0]
+			// Confirm each part arrives on the local subscriber so
+			// we know it has been processed by the worker.
+			waitForStreamTextPart(ctx, t, localEvents, text)
+		}
+
+		// NOW connect the relay subscriber on the non-worker replica.
+		// The relay must pick up all three buffered parts in its
+		// initial snapshot via the drainInitial loop.
+		relayEvents, relayStream, err := relayClient.StreamChat(ctx, chat.ID)
+		require.NoError(t, err)
+		defer relayStream.Close()
+
+		// Verify every buffered part arrives on the relay subscriber.
+		for _, text := range bufferedTexts {
+			event := waitForStreamTextPart(ctx, t, relayEvents, text)
+			require.Equal(t, "assistant", event.MessagePart.Role)
+		}
+
+		// Send one more chunk after the relay subscriber is connected
+		// and verify it arrives through the live channel.
+		liveText := "live-after-relay"
+		streamingChunks <- chattest.OpenAITextChunks(liveText)[0]
+		waitForStreamTextPart(ctx, t, localEvents, liveText)
+		waitForStreamTextPart(ctx, t, relayEvents, liveText)
+
+		close(streamingChunks)
+	})
 }

 func waitForStreamTextPart(