diff --git a/coderd/chatd/chatd.go b/coderd/chatd/chatd.go index 94d4a9902d..5ec44b5c3e 100644 --- a/coderd/chatd/chatd.go +++ b/coderd/chatd/chatd.go @@ -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 diff --git a/coderd/chatd/chatd_test.go b/coderd/chatd/chatd_test.go index 1c494ed75d..3f82586bd8 100644 --- a/coderd/chatd/chatd_test.go +++ b/coderd/chatd/chatd_test.go @@ -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() diff --git a/enterprise/coderd/chats.go b/enterprise/coderd/chats.go index b919d0188c..204eac588d 100644 --- a/enterprise/coderd/chats.go +++ b/enterprise/coderd/chats.go @@ -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(): diff --git a/enterprise/coderd/chats_test.go b/enterprise/coderd/chats_test.go index 628c289dd7..d920e6d4c6 100644 --- a/enterprise/coderd/chats_test.go +++ b/enterprise/coderd/chats_test.go @@ -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(