coder · ibetitsmike · Aug 29, 2026 · Aug 29, 2026
diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go
@@ -3117,16 +3117,17 @@ func (api *API) streamChat(rw http.ResponseWriter, r *http.Request) {
 
 	// Subscribe before accepting the WebSocket so that failures
 	// can still be reported as normal HTTP errors.
-	snapshot, events, cancelSub, ok := api.chatDaemon.SubscribeAuthorized(ctx, chat, r.Header, afterMessageID)
-	// Defensive against future SubscribeAuthorized failure modes.
-	if !ok {
+	session := chatd.NewSession(api.chatDaemon.StreamSessionConfig(ctx, chat, r.Header, afterMessageID))
+	if session == nil {
 		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
 			Message: "Chat streaming is not available.",
 			Detail:  "Chat stream state is not configured.",
 		})
 		return
 	}
-	defer cancelSub()
+	defer session.Close()
+	snapshot := session.InitialSnapshot()
+	events := session.Events()
 
 	conn, err := websocket.Accept(rw, r, nil)
 	if err != nil {

diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md
@@ -1027,6 +1027,8 @@ Coder stores no hook-specific dispatch or decision state. Delivery is best-effor
 
 # Stream loop
 
+<!-- TODO(human) -->
+
 The stream loop powers the `GET /api/experimental/chats/{chat}/stream` endpoint. It is scoped to one chat and one client WebSocket. It's responsible for delivering a stream of chat updates to the client, including:
 
 - messages committed to the database; and
@@ -1295,6 +1297,8 @@ WHERE id = ANY($1::uuid[]);
 
 ## Relay mechanism
 
+<!-- TODO(human) -->
+
 We make use of a relay mechanism when there are multiple coderd replicas. If a client connects to the stream endpoint on replica A, but the chat worker that owns the chat is on replica B, the endpoint will connect to replica B and relay streaming message parts.
 
 There exists a `GET /api/v2/chats/{chat}/stream/parts` endpoint that is responsible exclusively for streaming message parts. That endpoint talks to the chat worker on the same replica to obtain the message parts and relay them to the client.

diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go
@@ -3290,21 +3290,6 @@ func (p *Server) Start() *Server {
 	return p
 }
 
-func subscribeWithInitialError(chatID uuid.UUID, message string) (
-	[]codersdk.ChatStreamEvent,
-	<-chan codersdk.ChatStreamEvent,
-	func(),
-	bool,
-) {
-	events := make(chan codersdk.ChatStreamEvent)
-	close(events)
-	return []codersdk.ChatStreamEvent{{
-		Type:   codersdk.ChatStreamEventTypeError,
-		ChatID: chatID,
-		Error:  &codersdk.ChatError{Message: message},
-	}}, events, func() {}, true
-}
-
 // publishChatPubsubEvents broadcasts a lifecycle event for each affected chat.
 func (p *Server) publishChatPubsubEvents(chats []database.Chat, kind codersdk.ChatWatchEventKind) {
 	for _, chat := range chats {

diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go
@@ -1598,63 +1598,6 @@ func TestTurnWorkspaceContext_EnsureWorkspaceAgentIgnoresCachedAgentForDifferent
 	require.Equal(t, updatedChat, currentChat)
 }
 
-func TestSubscribeRejectsUnauthorizedCallerBeforeSharedFetches(t *testing.T) {
-	t.Parallel()
-
-	ctx := testutil.Context(t, testutil.WaitShort)
-	ctrl := gomock.NewController(t)
-	db := dbmock.NewMockStore(ctrl)
-	server := newSubscribeTestServer(t, db)
-
-	chatID := uuid.New()
-	db.EXPECT().GetChatByID(gomock.Any(), chatID).
-		Return(database.Chat{}, dbauthz.NotAuthorizedError{Err: xerrors.New("not authorized")})
-
-	snapshot, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0)
-	require.False(t, ok)
-	require.Nil(t, snapshot)
-	require.Nil(t, events)
-	require.Nil(t, cancel)
-}
-
-func TestSubscribeSurfacesTransientLookupFailureAsInitialError(t *testing.T) {
-	t.Parallel()
-
-	ctx := testutil.Context(t, testutil.WaitShort)
-	ctrl := gomock.NewController(t)
-	db := dbmock.NewMockStore(ctrl)
-	server := newSubscribeTestServer(t, db)
-
-	chatID := uuid.New()
-	db.EXPECT().GetChatByID(gomock.Any(), chatID).
-		Return(database.Chat{}, xerrors.New("transient lookup failure"))
-
-	snapshot, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0)
-	require.True(t, ok)
-	require.NotNil(t, cancel)
-	require.Len(t, snapshot, 1)
-	require.Equal(t, codersdk.ChatStreamEventTypeError, snapshot[0].Type)
-	require.Equal(t, chatID, snapshot[0].ChatID)
-	require.Equal(t, "failed to load initial snapshot", snapshot[0].Error.Message)
-
-	_, open := <-events
-	require.False(t, open)
-}
-
-func newSubscribeTestServer(t *testing.T, db database.Store) *Server {
-	t.Helper()
-
-	poller := newStreamSyncPoller(context.Background(), db, nil, slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}))
-	t.Cleanup(poller.Close)
-	return &Server{
-		db:               db,
-		logger:           slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
-		pubsub:           dbpubsub.NewInMemory(),
-		clock:            quartz.NewReal(),
-		streamSyncPoller: poller,
-	}
-}
-
 func TestResolveUserCompactionThreshold(t *testing.T) {
 	t.Parallel()
 

diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go
@@ -2827,9 +2827,10 @@ func TestSubscribeSnapshotIncludesStatusEvent(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	snapshot, _, cancel, ok := replica.Subscribe(ctx, chat.ID, nil, 0)
-	require.True(t, ok)
-	t.Cleanup(cancel)
+	session := chatd.NewSession(replica.StreamSessionConfig(ctx, chat, nil, 0))
+	require.NotNil(t, session)
+	t.Cleanup(session.Close)
+	snapshot := session.InitialSnapshot()
 
 	// Passive server: status is always Pending.
 	require.NotEmpty(t, snapshot)
@@ -4431,9 +4432,11 @@ func TestSubscribeNoDuplicateMessageParts(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	snapshot, events, cancel, ok := replica.Subscribe(ctx, chat.ID, nil, 0)
-	require.True(t, ok)
-	t.Cleanup(cancel)
+	session := chatd.NewSession(replica.StreamSessionConfig(ctx, chat, nil, 0))
+	require.NotNil(t, session)
+	t.Cleanup(session.Close)
+	snapshot := session.InitialSnapshot()
+	events := session.Events()
 
 	// Snapshot should have events (at minimum: status + message).
 	require.NotEmpty(t, snapshot)
@@ -4513,23 +4516,21 @@ func TestSubscribeAfterMessageID(t *testing.T) {
 		Content:        thirdContent,
 	})
 
-	// Control: Subscribe with afterMessageID=0 returns ALL messages.
-	allSnapshot, _, cancelAll, ok := replica.Subscribe(ctx, chat.ID, nil, 0)
-	require.True(t, ok)
-	cancelAll()
-
-	allMessages := filterMessageEvents(allSnapshot)
-	require.Len(t, allMessages, 3, "afterMessageID=0 should return all three messages")
-
-	// Subscribe with afterMessageID set to the second message's ID.
-	// Only the third message (inserted after msg2) should appear.
-	partialSnapshot, _, cancelPartial, ok := replica.Subscribe(ctx, chat.ID, nil, msg2.ID)
-	require.True(t, ok)
-	cancelPartial()
-
-	partialMessages := filterMessageEvents(partialSnapshot)
-	require.Len(t, partialMessages, 1, "afterMessageID=msg2.ID should return only messages after msg2")
-	require.Equal(t, codersdk.ChatMessageRoleUser, partialMessages[0].Message.Role)
+	for _, tc := range []struct {
+		name    string
+		afterID int64
+		wantLen int
+	}{
+		{name: "all messages", afterID: 0, wantLen: 3},
+		{name: "only messages after msg2", afterID: msg2.ID, wantLen: 1},
+	} {
+		session := chatd.NewSession(replica.StreamSessionConfig(ctx, chat, nil, tc.afterID))
+		require.NotNil(t, session, tc.name)
+		messages := filterMessageEvents(session.InitialSnapshot())
+		session.Close()
+		require.Len(t, messages, tc.wantLen, tc.name)
+		require.Equal(t, codersdk.ChatMessageRoleUser, messages[len(messages)-1].Message.Role, tc.name)
+	}
 }
 
 // filterMessageEvents returns only the Message-type events from a
@@ -10623,9 +10624,10 @@ func TestProcessChat_RoutingUsesDelegatedAPIKey(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	_, events, cancel, ok := creator.Subscribe(ctx, chat.ID, nil, 0)
-	require.True(t, ok)
-	t.Cleanup(cancel)
+	session := chatd.NewSession(creator.StreamSessionConfig(ctx, chat, nil, 0))
+	require.NotNil(t, session)
+	t.Cleanup(session.Close)
+	events := session.Events()
 
 	_ = newActiveTestServer(t, db, ps, func(cfg *chatd.Config) {
 		cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(factory)
@@ -13922,8 +13924,9 @@ func TestAdvisorHappyPath_RootChat(t *testing.T) {
 	require.NoError(t, err)
 
 	// Advisor deltas are transient; a late subscriber misses them.
-	_, liveEvents, cancelLive, ok := server.Subscribe(ctx, chat.ID, nil, 0)
-	require.True(t, ok)
+	session := chatd.NewSession(server.StreamSessionConfig(ctx, chat, nil, 0))
+	require.NotNil(t, session)
+	liveEvents := session.Events()
 	liveCollectorDone := make(chan struct{})
 	go func() {
 		defer close(liveCollectorDone)
@@ -14016,7 +14019,7 @@ func TestAdvisorHappyPath_RootChat(t *testing.T) {
 	// new snapshots, so the assertion must use the live collector.
 	require.Eventually(t, liveDeltasCaptured, testutil.WaitLong, testutil.IntervalFast,
 		"advisor nested text deltas must stream into the parent tool card")
-	cancelLive()
+	session.Close()
 	<-liveCollectorDone
 	livePartsMu.Lock()
 	collectedAdvisorDeltas := append([]string(nil), liveAdvisorDeltas...)