From 12e49c18a5b0acac82bc35abd251449b73e147b9 Mon Sep 17 00:00:00 2001 From: Cian Johnston Date: Mon, 20 Apr 2026 09:19:17 +0100 Subject: [PATCH] fix(enterprise/coderd/x/chatd): reduce relay reconnect spam (#24495) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Replaces the hard-coded 500ms reconnect timer for dialing chat relays with exponential backoff via `coder/retry`. - `dialRelay` drops the `codersdk.ExperimentalClient.StreamChat` wrapper and calls `websocket.Dial` directly so we can capture `*http.Response.StatusCode` without parsing error strings. - Adds `RelayDialError` that exposes the HTTP status from `websocket.Dial` - Modifies retry logic: 401/403 tear the stream down immediately, 5xx/network/timeouts retry then tear down on cap. Outer stream closes cleanly so the browser SDK reconnects with a fresh cookie. - Retry state resets on successful dial and on target-worker change, not on every `closeRelay()`. > πŸ€– Generated by Coder Agents. --- enterprise/coderd/x/chatd/chatd.go | 413 ++++++--- enterprise/coderd/x/chatd/chatd_retry_test.go | 796 ++++++++++++++++++ 2 files changed, 1112 insertions(+), 97 deletions(-) create mode 100644 enterprise/coderd/x/chatd/chatd_retry_test.go diff --git a/enterprise/coderd/x/chatd/chatd.go b/enterprise/coderd/x/chatd/chatd.go index 22c63d3722..ce1a002723 100644 --- a/enterprise/coderd/x/chatd/chatd.go +++ b/enterprise/coderd/x/chatd/chatd.go @@ -2,9 +2,12 @@ package chatd import ( "context" + "errors" + "fmt" "math" "net/http" "net/url" + "strconv" "strings" "time" @@ -13,11 +16,12 @@ import ( "cdr.dev/slog/v3" "github.com/coder/coder/v2/coderd/database" - "github.com/coder/coder/v2/coderd/util/ptr" osschatd "github.com/coder/coder/v2/coderd/x/chatd" "github.com/coder/coder/v2/codersdk" "github.com/coder/quartz" + "github.com/coder/retry" "github.com/coder/websocket" + "github.com/coder/websocket/wsjson" ) // RelaySourceHeader marks replica-relayed stream requests. @@ -32,8 +36,39 @@ const ( // buffered snapshot events time to be forwarded before // the relay is torn down. relayDrainTimeout = 200 * time.Millisecond + + // Retry knobs for the cross-replica relay handshake. Uses the + // github.com/coder/retry defaults (Ο†-growth, no jitter) but drives + // the delay manually because retry.Retrier.Wait uses time.After, + // which isn't compatible with quartz.Clock determinism in tests. + relayRetryFloor = 500 * time.Millisecond // first retry matches old fixed delay + relayRetryCeil = 15 * time.Second // cap stall before tear-down + // After this many reconnect retries the relay leg is torn down. + // Total dial attempts = 1 initial dial + relayMaxRetries. + relayMaxRetries = 6 ) +// RelayDialError wraps a failed relay handshake. HTTPStatus is 0 +// when the failure happened before a response (DNS, TCP, TLS, +// timeout, context cancel); otherwise it carries the peer's status +// code for the reconnect loop to classify. +type RelayDialError struct { + HTTPStatus int + Err error +} + +func (e *RelayDialError) Error() string { return e.Err.Error() } +func (e *RelayDialError) Unwrap() error { return e.Err } + +// IsUnrecoverable reports whether retrying with the same captured +// session token is futile. Only 401/403 qualify - the token is dead +// or the peer won't authorize it. 5xx, 429, network, and context +// errors fall through to backoff. +func (e *RelayDialError) IsUnrecoverable() bool { + return e.HTTPStatus == http.StatusUnauthorized || + e.HTTPStatus == http.StatusForbidden +} + // MultiReplicaSubscribeConfig holds the dependencies for multi-replica chat // subscription. ReplicaIDFn is called lazily because the // replica ID may not be known at construction time. @@ -62,10 +97,8 @@ type MultiReplicaSubscribeConfig struct { Clock quartz.Clock } -// dial returns the dialer function to use for relay connections. -// If DialerFn is set (e.g. in tests), it takes precedence. -// Otherwise, dialRelay is used with the real MultiReplicaSubscribeConfig dependencies. -// Returns nil when no relay capability is configured. +// dial returns the configured dialer, preferring DialerFn (tests) +// over the real dialRelay. Returns nil when relay is not configured. func (c MultiReplicaSubscribeConfig) dial() func( ctx context.Context, chatID uuid.UUID, @@ -160,9 +193,16 @@ func NewMultiReplicaSubscribeFn( parts <-chan codersdk.ChatStreamEvent cancel func() workerID uuid.UUID // the worker this dial targeted + // err and parts are mutually exclusive: success sets + // parts; failure sets err (unwrap to *RelayDialError + // for classification). + err error } relayReadyCh := make(chan relayResult, 4) + // Reset on successful dial or when the relay target + // changes, so a fresh target starts at the floor delay. + retryState := newRelayRetryState() // Per-dial context so in-flight dials can be canceled when // a new dial is initiated or the relay is closed. var dialCancel context.CancelFunc @@ -236,6 +276,12 @@ func NewMultiReplicaSubscribeFn( if cfg.dial() == nil { return } + // Scoped here (not in closeRelay) so repeated dials + // against the same worker keep the attempt counter and + // correctly trip the cap. + if workerID != expectedWorkerID { + retryState.reset() + } closeRelay() // Create a per-dial context so this goroutine is // canceled if closeRelay() or openRelayAsync() is @@ -250,23 +296,31 @@ func NewMultiReplicaSubscribeFn( // since they are expected when a dial is // superseded by a newer one. if dialCtx.Err() == nil { - logger.Warn(ctx, "failed to open relay for message parts", + fields := []slog.Field{ slog.F("chat_id", chatID), slog.F("worker_id", workerID), slog.Error(err), - ) + } + // Surface the peer's HTTP status (when we + // got one) as a structured field so + // operators can filter 401/403 spam + // separately from 5xx/network warnings. + var dialErr *RelayDialError + if errors.As(err, &dialErr) && dialErr.HTTPStatus != 0 { + fields = append(fields, slog.F("http_status", dialErr.HTTPStatus)) + } + logger.Warn(ctx, "failed to open relay for message parts", fields...) } - // Send an empty result so the merge loop - // can schedule a reconnect attempt. + // Hand the error to the merge loop, which will + // classify it and either back off or tear down. select { - case relayReadyCh <- relayResult{workerID: workerID}: + case relayReadyCh <- relayResult{workerID: workerID, err: err}: case <-dialCtx.Done(): } return - } // If the dial context was canceled while the - // dial was in progress, discard the result to - // avoid starting a wrappedParts goroutine for - // a stale connection. + } + // Discard stale dials so we don't start a + // wrappedParts goroutine on a canceled connection. if dialCtx.Err() != nil { cancel() return @@ -274,7 +328,7 @@ func NewMultiReplicaSubscribeFn( // Wrap the relay channel so snapshot parts // are delivered through the same channel as // live parts. This goroutine only forwards - // events β€” it does not own the relay + // events - it does not own the relay // lifecycle. When dialCtx is canceled it // simply returns, closing wrappedParts via // its defer. The cancel() is called by @@ -316,20 +370,35 @@ func NewMultiReplicaSubscribeFn( }() } - // 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() { + // scheduleRelayReconnect arms a timer so the select loop + // can re-check chat status and reopen the relay. Callers + // pass the delay from retryState so the failed-dial branch + // gets backoff while transient branches stay at the floor. + scheduleRelayReconnect := func(delay time.Duration) { if cfg.dial() == nil { return } if reconnectTimer != nil { reconnectTimer.Stop() } - reconnectTimer = cfg.clock().NewTimer(500*time.Millisecond, "reconnect") + reconnectTimer = cfg.clock().NewTimer(delay, "reconnect") reconnectCh = reconnectTimer.C } + // sendRelayTerminalError enqueues one error event for the + // subscriber; callers return afterwards so the deferred + // close(mergedEvents) fires and the OSS merge loop tears + // the relay leg down while pubsub/local sources keep going. + sendRelayTerminalError := func(msg string) { + select { + case mergedEvents <- codersdk.ChatStreamEvent{ + Type: codersdk.ChatStreamEventTypeError, + ChatID: chatID, + Error: &codersdk.ChatStreamError{Message: msg}, + }: + case <-ctx.Done(): + } + } statusNotifications := params.StatusNotifications go func() { defer close(mergedEvents) @@ -360,20 +429,59 @@ func NewMultiReplicaSubscribeFn( continue } // A nil parts channel signals the dial - // failed β€” schedule a retry. + // failed - classify the error to decide + // whether to schedule a backoff retry, emit a + // terminal error and tear the relay leg down + // (unrecoverable / cap reached), or simply + // drop the stale drain. if result.parts == nil { if drainAndClose { // Dial failed and we were only - // waiting to drain β€” nothing to do. + // waiting to drain - nothing to do. drainAndClose = false - } else { - scheduleRelayReconnect() + continue } + var dialErr *RelayDialError + if errors.As(result.err, &dialErr) && dialErr.IsUnrecoverable() { + logger.Warn(ctx, "relay dial unrecoverable; tearing down relay leg", + slog.F("chat_id", chatID), + slog.F("worker_id", result.workerID), + slog.F("http_status", dialErr.HTTPStatus), + ) + sendRelayTerminalError(fmt.Sprintf( + "relay authentication failed (status %d)", + dialErr.HTTPStatus, + )) + return + } + delay, giveUp := retryState.next() + if giveUp { + logger.Warn(ctx, "relay dial retry cap reached; tearing down relay leg", + slog.F("chat_id", chatID), + slog.F("worker_id", result.workerID), + slog.F("max_retries", relayMaxRetries), + ) + sendRelayTerminalError(fmt.Sprintf( + "relay connection failed after %d retries", + relayMaxRetries, + )) + return + } + scheduleRelayReconnect(delay) continue - } // An async relay dial completed; swap - // in the new relay channel. + } + // An async relay dial completed. Swap in the + // new relay channel. We deliberately do NOT + // reset the retry counter here: a peer that + // accepts the handshake and immediately drops + // the stream would otherwise keep reconnecting + // forever, since each success would zero the + // counter before the next drop re-incremented + // it. The counter only resets when the target + // worker changes (see openRelayAsync). if relayCancel != nil { relayCancel() + relayCancel = nil } relayParts = result.parts relayCancel = result.cancel @@ -398,7 +506,7 @@ func NewMultiReplicaSubscribeFn( // openRelayAsync handle the new one. closeRelay() } else { - // Chat is still idle β€” drain the + // Chat is still idle - drain the // buffered snapshot before closing. if drainTimer != nil { drainTimer.Stop() @@ -421,8 +529,24 @@ func NewMultiReplicaSubscribeFn( ) // Retry on transient DB errors to // avoid permanently stalling the - // stream. - scheduleRelayReconnect() + // stream. The same retry state + // bounds the DB-error loop too so a + // persistently broken DB eventually + // tears the relay down instead of + // spinning forever. + delay, giveUp := retryState.next() + if giveUp { + logger.Warn(ctx, "relay reconnect retry cap reached; tearing down relay leg", + slog.F("chat_id", chatID), + slog.F("max_retries", relayMaxRetries), + ) + sendRelayTerminalError(fmt.Sprintf( + "relay connection failed after %d retries", + relayMaxRetries, + )) + return + } + scheduleRelayReconnect(delay) continue } if currentChat.Status == database.ChatStatusRunning && @@ -459,9 +583,6 @@ func NewMultiReplicaSubscribeFn( drainTimerCh = nil drainTimer = nil closeRelay() - drainTimerCh = nil - drainTimer = nil - closeRelay() case event, ok := <-relayPartsCh: if !ok { if relayCancel != nil { @@ -469,9 +590,21 @@ func NewMultiReplicaSubscribeFn( relayCancel = nil } relayParts = nil - // Schedule reconnection instead of - // giving up. - scheduleRelayReconnect() + // Reuse the retry state so a relay that + // repeatedly drops eventually tears down. + delay, giveUp := retryState.next() + if giveUp { + logger.Warn(ctx, "relay drop retry cap reached; tearing down relay leg", + slog.F("chat_id", chatID), + slog.F("max_retries", relayMaxRetries), + ) + sendRelayTerminalError(fmt.Sprintf( + "relay connection failed after %d retries", + relayMaxRetries, + )) + return + } + scheduleRelayReconnect(delay) continue } // Only forward message_part events from @@ -495,12 +628,60 @@ func NewMultiReplicaSubscribeFn( } } -// dialRelay opens a WebSocket relay connection to the replica -// identified by workerID and returns a snapshot of buffered -// message_part events plus a live channel of subsequent events. -// It passes afterID=MaxInt64 so the remote replica skips the -// full message history snapshot, since the relay only needs -// live message_part events. +// relayRetryState drives the retry policy for the relay reconnect +// loop. Wraps github.com/coder/retry to reuse its Ο†-growth defaults +// but computes the delay without blocking so the merge loop can +// schedule its own quartz.Clock timer. +// +// Not safe for concurrent use. +type relayRetryState struct { + retrier *retry.Retrier + attempts int +} + +func newRelayRetryState() *relayRetryState { + return &relayRetryState{ + retrier: retry.New(relayRetryFloor, relayRetryCeil), + } +} + +// next returns the delay before the next dial and sets giveUp once +// attempts exceed relayMaxRetries. Adapts the math from +// retry.Retrier.Wait (github.com/coder/retry/retrier.go) without +// blocking: the library's Wait returns 0 on the first call and sets +// Delay to Floor only after the sleep, so we clamp to Floor up +// front. +func (s *relayRetryState) next() (delay time.Duration, giveUp bool) { + s.attempts++ + if s.attempts > relayMaxRetries { + return 0, true + } + r := s.retrier + d := time.Duration(float64(r.Delay) * r.Rate) + if d > r.Ceil { + d = r.Ceil + } + if d < r.Floor { + d = r.Floor + } + r.Delay = d + return d, false +} + +// reset returns the state to the floor delay and zero attempts. +// Called after a successful dial or a relay target change. +func (s *relayRetryState) reset() { + s.retrier.Reset() + s.attempts = 0 +} + +// dialRelay opens a WebSocket to the replica owning chatID and +// returns any buffered message_part snapshot plus a live channel of +// subsequent events. Handshake failures return an error unwrapping +// to *RelayDialError so callers can classify via IsUnrecoverable. +// +// websocket.Dial is called directly (not via the SDK wrapper) so we +// can read *http.Response.StatusCode for classification. func dialRelay( ctx context.Context, chatID uuid.UUID, @@ -516,32 +697,79 @@ func dialRelay( ) { address, ok := cfg.ResolveReplicaAddress(ctx, workerID) if !ok { - return nil, nil, nil, xerrors.New("worker replica not found") + return nil, nil, nil, &RelayDialError{ + Err: xerrors.New("dial relay stream: worker replica not found"), + } } - baseURL, err := url.Parse(address) + wsURL, err := buildRelayURL(address, chatID) if err != nil { - return nil, nil, nil, xerrors.Errorf("parse relay address %q: %w", address, err) + return nil, nil, nil, &RelayDialError{ + Err: xerrors.Errorf("dial relay stream: %w", err), + } } + replicaID := cfg.ReplicaIDFn() + headers := make(http.Header, 2) + headers.Set(codersdk.SessionTokenHeader, extractSessionToken(requestHeader)) + headers.Set(RelaySourceHeader, replicaID.String()) + relayCtx, relayCancel := context.WithCancel(ctx) - sdkClient := codersdk.New(baseURL) - sdkClient.HTTPClient = cfg.ReplicaHTTPClient - sdkClient.SessionTokenProvider = relayTokenProvider{ - token: extractSessionToken(requestHeader), - replicaID: replicaID, - } - expClient := codersdk.NewExperimentalClient(sdkClient) - sourceEvents, sourceStream, err := expClient.StreamChat(relayCtx, chatID, &codersdk.StreamChatOptions{ - AfterID: ptr.Ref(int64(math.MaxInt64)), + conn, resp, dialErr := websocket.Dial(relayCtx, wsURL, &websocket.DialOptions{ + HTTPClient: cfg.ReplicaHTTPClient, + HTTPHeader: headers, + CompressionMode: websocket.CompressionDisabled, }) - if err != nil { - relayCancel() - return nil, nil, nil, xerrors.Errorf("dial relay stream: %w", err) + status := 0 + if resp != nil { + status = resp.StatusCode + // The websocket library closes resp.Body on success; on + // failure we close it ourselves so we don't leak the TCP + // connection. + if dialErr != nil && resp.Body != nil { + _ = resp.Body.Close() + } } + if dialErr != nil { + relayCancel() + return nil, nil, nil, &RelayDialError{ + HTTPStatus: status, + Err: xerrors.Errorf("dial relay stream: %w", dialErr), + } + } + // Match the server's 4 MiB read limit in codersdk.StreamChat so + // large message_part batches don't trip the default 32 KiB cap. + conn.SetReadLimit(1 << 22) snapshot = make([]codersdk.ChatStreamEvent, 0, 100) + // sourceEvents is the flattened batchβ†’event channel. A small + // goroutine reads batches off the websocket and fans them out; + // callers see a single event stream identical to the shape the + // old SDK call produced. + sourceEvents := make(chan codersdk.ChatStreamEvent, 128) + go func() { + defer close(sourceEvents) + for { + var batch []codersdk.ChatStreamEvent + if readErr := wsjson.Read(relayCtx, conn, &batch); readErr != nil { + return + } + for _, event := range batch { + select { + case sourceEvents <- event: + case <-relayCtx.Done(): + return + } + } + } + }() + + closeSource := func() { + relayCancel() + _ = conn.Close(websocket.StatusNormalClosure, "") + } + // 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. @@ -553,9 +781,10 @@ drainInitial: for len(snapshot) < cap(snapshot) { select { case <-relayCtx.Done(): - _ = sourceStream.Close() - relayCancel() - return nil, nil, nil, xerrors.Errorf("dial relay stream: %w", relayCtx.Err()) + closeSource() + return nil, nil, nil, &RelayDialError{ + Err: xerrors.Errorf("dial relay stream: %w", relayCtx.Err()), + } case event, ok := <-sourceEvents: if !ok { break drainInitial @@ -577,12 +806,9 @@ drainInitial: go func() { defer close(events) - defer relayCancel() - defer func() { - _ = sourceStream.Close() - }() + defer closeSource() - // No need to re-send snapshot events β€” they're + // No need to re-send snapshot events - they're // returned to the caller directly. for { select { @@ -604,40 +830,33 @@ drainInitial: } }() - cancelFn := func() { - relayCancel() - _ = sourceStream.Close() + return snapshot, events, closeSource, nil +} + +// buildRelayURL builds the websocket URL for the chat stream +// endpoint on a peer replica. It maps http(s) schemes to ws(s). +func buildRelayURL(address string, chatID uuid.UUID) (string, error) { + u, err := url.Parse(address) + if err != nil { + return "", xerrors.Errorf("parse relay address %q: %w", address, err) } - return snapshot, events, cancelFn, nil -} - -// relayTokenProvider authenticates relay requests to the worker -// replica using the session token extracted from the original -// browser request. It also stamps each request with the relay -// source header so the worker can identify it as an inter-replica -// call. -type relayTokenProvider struct { - token string - replicaID uuid.UUID -} - -func (p relayTokenProvider) AsRequestOption() codersdk.RequestOption { - return func(req *http.Request) { - req.Header.Set(codersdk.SessionTokenHeader, p.token) - req.Header.Set(RelaySourceHeader, p.replicaID.String()) + switch u.Scheme { + case "http": + u.Scheme = "ws" + case "https": + u.Scheme = "wss" + case "ws", "wss": + // already a websocket URL, leave as-is. + default: + return "", xerrors.Errorf("unsupported relay address scheme %q", u.Scheme) } -} - -func (p relayTokenProvider) SetDialOption(opts *websocket.DialOptions) { - if opts.HTTPHeader == nil { - opts.HTTPHeader = make(http.Header) - } - opts.HTTPHeader.Set(codersdk.SessionTokenHeader, p.token) - opts.HTTPHeader.Set(RelaySourceHeader, p.replicaID.String()) -} - -func (p relayTokenProvider) GetSessionToken() string { - return p.token + u.Path = fmt.Sprintf("/api/experimental/chats/%s/stream", chatID) + q := u.Query() + // Relays only need live message_part events, not the full + // history; pass after_id=MaxInt64 so the peer skips its snapshot. + q.Set("after_id", strconv.FormatInt(math.MaxInt64, 10)) + u.RawQuery = q.Encode() + return u.String(), nil } // extractSessionToken returns the session token carried by the @@ -648,7 +867,7 @@ func extractSessionToken(header http.Header) string { if header == nil { return "" } - // Cookie (browser WebSocket upgrade β€” most common relay case). + // Cookie (browser WebSocket upgrade - most common relay case). if raw := header.Get(cookieHeader); raw != "" { r := &http.Request{Header: http.Header{cookieHeader: {raw}}} if c, err := r.Cookie(codersdk.SessionTokenCookie); err == nil && c.Value != "" { diff --git a/enterprise/coderd/x/chatd/chatd_retry_test.go b/enterprise/coderd/x/chatd/chatd_retry_test.go new file mode 100644 index 0000000000..3135796116 --- /dev/null +++ b/enterprise/coderd/x/chatd/chatd_retry_test.go @@ -0,0 +1,796 @@ +package chatd_test + +import ( + "context" + "database/sql" + "encoding/json" + "io" + "math" + "net/http" + "net/http/httptest" + "regexp" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "cdr.dev/slog/v3/sloggers/slogtest" + "github.com/coder/coder/v2/coderd/database" + "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" + osschatd "github.com/coder/coder/v2/coderd/x/chatd" + "github.com/coder/coder/v2/codersdk" + entchatd "github.com/coder/coder/v2/enterprise/coderd/x/chatd" + "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" +) + +// mulPhi multiplies a duration by math.Phi to compute the next +// step in retry.Retrier's Ο†-growth backoff sequence. If +// TestRelayReconnectUsesExponentialBackoff starts failing after a +// retry library bump, check whether the growth factor has changed. +func mulPhi(d time.Duration) time.Duration { + return time.Duration(float64(d) * math.Phi) +} + +// setChatRunningAndPublish marks the chat row as running on workerID +// and publishes a matching status notification. It keeps the DB row +// and pubsub notification in sync so the async reconnect loop +// re-dials on each timer fire (the reconnect branch re-checks DB +// status before calling openRelayAsync). +func setChatRunningAndPublish( + ctx context.Context, + t *testing.T, + db database.Store, + ps dbpubsub.Pubsub, + chatID, workerID uuid.UUID, +) { + t.Helper() + now := time.Now() + _, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ + ID: chatID, + Status: database.ChatStatusRunning, + WorkerID: uuid.NullUUID{UUID: workerID, Valid: true}, + StartedAt: sql.NullTime{Time: now, Valid: true}, + HeartbeatAt: sql.NullTime{Time: now, Valid: true}, + }) + require.NoError(t, err) + payload, err := json.Marshal(coderdpubsub.ChatStreamNotifyMessage{ + Status: string(database.ChatStatusRunning), + WorkerID: workerID.String(), + }) + require.NoError(t, err) + require.NoError(t, ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chatID), payload)) +} + +// TestRelayDialErrorIsUnrecoverable locks the classification policy. +// Adding a new HTTP status to the unrecoverable set should force a +// test edit too. +func TestRelayDialErrorIsUnrecoverable(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + status int + want bool + }{ + {"unauthorized", http.StatusUnauthorized, true}, + {"forbidden", http.StatusForbidden, true}, + {"internal_server", http.StatusInternalServerError, false}, + {"bad_gateway", http.StatusBadGateway, false}, + {"service_unavailable", http.StatusServiceUnavailable, false}, + {"too_many_requests", http.StatusTooManyRequests, false}, + {"pre_response", 0, false}, + {"bad_request", http.StatusBadRequest, false}, + {"not_found", http.StatusNotFound, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + e := &entchatd.RelayDialError{HTTPStatus: tc.status, Err: io.EOF} + require.Equal(t, tc.want, e.IsUnrecoverable(), + "status=%d", tc.status) + }) + } +} + +// TestRelayReconnectUsesExponentialBackoff asserts that the reconnect +// timer follows the Ο†-growth sequence produced by +// github.com/coder/retry's defaults, floored at relayRetryFloor. +func TestRelayReconnectUsesExponentialBackoff(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + workerID := uuid.New() + subscriberID := uuid.New() + + var failCount atomic.Int32 + dialer := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( + []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, + ) { + failCount.Add(1) + return nil, nil, nil, &entchatd.RelayDialError{ + HTTPStatus: http.StatusBadGateway, + Err: io.EOF, + } + } + + mclk := quartz.NewMock(t) + trapReconnect := mclk.Trap().NewTimer("reconnect") + defer trapReconnect.Close() + + subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) + + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(ctx, t, db) + chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-backoff") + + _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) + require.True(t, ok) + t.Cleanup(cancel) + + // Kick the async relay loop and keep the DB row in sync so + // each reconnect timer fire triggers another dial. + setChatRunningAndPublish(ctx, t, db, ps, chat.ID, workerID) + // Expected sequence from retry.Retrier math: + // attempt 1 β†’ floor (500ms) + // attempt n β†’ prev Γ— Ο† (capped at ceil) + floor := 500 * time.Millisecond + expected := []time.Duration{ + floor, + mulPhi(floor), + mulPhi(mulPhi(floor)), + mulPhi(mulPhi(mulPhi(floor))), + mulPhi(mulPhi(mulPhi(mulPhi(floor)))), + } + + for i, want := range expected { + call := trapReconnect.MustWait(ctx) + require.Equal(t, want, call.Duration, + "attempt %d: want %v got %v", i+1, want, call.Duration) + call.MustRelease(ctx) + mclk.Advance(want).MustWait(ctx) + } + + // We expect 1 initial attempt + 5 reconnects fired by the + // trapped timer = 6 dials before the cap-check runs. Use + // Eventually so we don't race the final dial goroutine that + // the last Advance kicked off. + require.Eventually(t, func() bool { + return failCount.Load() >= 6 + }, testutil.WaitShort, testutil.IntervalFast, + "expected 6 dials, got %d", failCount.Load()) + + // The events channel must remain open - we're still under the + // cap. + select { + case ev, open := <-events: + if !open { + t.Fatalf("events channel closed prematurely; retries should continue below cap") + } + // Allow through events that might have been queued; just + // confirm it's not a terminal error. + if ev.Type == codersdk.ChatStreamEventTypeError { + t.Fatalf("unexpected terminal error: %v", ev.Error) + } + default: + } +} + +// TestRelayReconnectResetsOnSuccess exercises the path where a +// successful dial resets the retry state so the next failure starts +// over at the floor delay. +// TestRelayRepeatedDropsHitCap verifies the cap covers a peer that +// accepts the handshake and immediately drops it. Without a proper +// cap, such a peer would produce one reconnect per floor delay +// forever. The retry counter must accumulate across dial-success / +// parts-close cycles so the cap trips. +func TestRelayRepeatedDropsHitCap(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + workerID := uuid.New() + subscriberID := uuid.New() + + opened := make(chan chan codersdk.ChatStreamEvent, 32) + var call atomic.Int32 + dialer := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( + []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, + ) { + call.Add(1) + ch := make(chan codersdk.ChatStreamEvent, 1) + opened <- ch + return nil, ch, func() {}, nil + } + + mclk := quartz.NewMock(t) + trapReconnect := mclk.Trap().NewTimer("reconnect") + defer trapReconnect.Close() + + subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) + + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(ctx, t, db) + chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-drops") + + _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) + require.True(t, ok) + t.Cleanup(cancel) + + // Kick off the first async dial. + setChatRunningAndPublish(ctx, t, db, ps, chat.ID, workerID) + + // Close the first dial's parts channel so the merge loop + // schedules a reconnect. Then advance 6 reconnect timers, + // closing the parts channel each time so the cycle is: + // dial -> success -> parts-close -> next() -> reconnect. + // 1 initial dial + 6 timer-driven dials = 7 total; the 7th + // parts-close trips the cap. + for i := 0; i < 7; i++ { + var ch chan codersdk.ChatStreamEvent + select { + case ch = <-opened: + case <-ctx.Done(): + t.Fatalf("timed out waiting for dial %d", i+1) + } + // Closing the parts channel triggers the relayPartsCh + // close branch, which calls retryState.next() and + // schedules the next reconnect. + close(ch) + if i == 6 { + // 7th parts-close should trip the cap; no more + // reconnect timers. + break + } + call := trapReconnect.MustWait(ctx) + call.MustRelease(ctx) + mclk.Advance(call.Duration).MustWait(ctx) + } + + // A terminal error event must arrive on the events channel. + var errEvent *codersdk.ChatStreamEvent + require.Eventually(t, func() bool { + select { + case ev, open := <-events: + if !open { + return errEvent != nil + } + if ev.Type == codersdk.ChatStreamEventTypeError { + errEvent = &ev + return true + } + return false + default: + return false + } + }, testutil.WaitShort, testutil.IntervalFast, + "expected a terminal error event after repeated drops hit cap") + require.NotNil(t, errEvent.Error) + require.Contains(t, errEvent.Error.Message, "relay connection failed") + + // We should have observed exactly 7 dials before tear-down. + require.Equal(t, int32(7), call.Load(), + "expected 7 dials (1 initial + 6 reconnect retries) before cap") +} + +// TestRelayStopsAfterIntermittentCap verifies the cap-reached +// tear-down path: after N intermittent failures the merge loop emits +// one error event, closes the events channel, and stops dialing. +func TestRelayStopsAfterIntermittentCap(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + workerID := uuid.New() + subscriberID := uuid.New() + + var callCount atomic.Int32 + dialer := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( + []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, + ) { + callCount.Add(1) + return nil, nil, nil, &entchatd.RelayDialError{ + HTTPStatus: http.StatusBadGateway, + Err: io.EOF, + } + } + + mclk := quartz.NewMock(t) + trapReconnect := mclk.Trap().NewTimer("reconnect") + defer trapReconnect.Close() + + subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) + + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(ctx, t, db) + chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-cap") + + _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) + require.True(t, ok) + t.Cleanup(cancel) + + setChatRunningAndPublish(ctx, t, db, ps, chat.ID, workerID) + // Advance through N consecutive reconnect timers. Each one + // triggers a dial, which fails and schedules the next timer. + // After the Nth failure the retry state says giveUp=true on + // the next .next() call, so the merge loop tears down. + for i := 0; i < 6; i++ { + call := trapReconnect.MustWait(ctx) + call.MustRelease(ctx) + mclk.Advance(call.Duration).MustWait(ctx) + } + + // Wait for the terminal error event to arrive. mergedEvents + // closes inside the enterprise merge goroutine, but OSS only + // nil-outs relayEvents on close - the outer events channel + // stays open for pubsub/local, so we wait for the error event + // itself rather than channel closure. + var errEvent *codersdk.ChatStreamEvent + require.Eventually(t, func() bool { + select { + case ev, open := <-events: + if !open { + return errEvent != nil + } + if ev.Type == codersdk.ChatStreamEventTypeError { + errEvent = &ev + return true + } + return false + default: + return false + } + }, testutil.WaitShort, testutil.IntervalFast, + "expected a terminal error event") + require.NotNil(t, errEvent, "expected a terminal error event") + require.NotNil(t, errEvent.Error) + require.Contains(t, errEvent.Error.Message, "relay connection failed") + require.Contains(t, errEvent.Error.Message, "6") + + // Ensure the cap fires at attempt N+1 - the retry state allows + // relayMaxRetries successful next() calls before flipping + // giveUp. With one initial dial + 6 reconnect-timer fires the + // 7th .next() trips the cap and tears down, so we see 7 dials + // total and nothing further. + totalDials := callCount.Load() + require.Equal(t, int32(7), totalDials, + "expected exactly relayMaxRetries+1 dials before cap; got %d", totalDials) +} + +// chatByIDErrorStore wraps a database.Store and forces GetChatByID +// to return a caller-supplied error once after N successful calls. +// This lets the initial Subscribe call succeed (OSS's initial state +// load needs a real Chat to wire up the relay) while subsequent +// reconnect-branch calls exercise the DB-error retry path. +type chatByIDErrorStore struct { + database.Store + err error + okRemain atomic.Int32 // number of calls allowed to delegate before erroring. +} + +func (s *chatByIDErrorStore) GetChatByID(ctx context.Context, id uuid.UUID) (database.Chat, error) { + if s.okRemain.Add(-1) >= 0 { + return s.Store.GetChatByID(ctx, id) + } + return database.Chat{}, s.err +} + +// TestRelayReconnectStopsAfterDBErrorCap verifies the reconnect-timer +// branch's DB-error path shares the same retry budget as dial +// failures and trips the cap after enough consecutive DB errors. +func TestRelayReconnectStopsAfterDBErrorCap(t *testing.T) { + t.Parallel() + + realDB, ps := dbtestutil.NewDB(t) + workerID := uuid.New() + subscriberID := uuid.New() + + var callCount atomic.Int32 + dialer := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( + []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, + ) { + callCount.Add(1) + return nil, nil, nil, &entchatd.RelayDialError{ + HTTPStatus: http.StatusBadGateway, + Err: io.EOF, + } + } + + mclk := quartz.NewMock(t) + trapReconnect := mclk.Trap().NewTimer("reconnect") + defer trapReconnect.Close() + + // The server sees a DB whose GetChatByID always errors after + // the initial Subscribe snapshot load. Other methods delegate + // to the real DB, so seeding below still works. + failingDB := &chatByIDErrorStore{ + Store: realDB, + err: xerrors.New("mock: GetChatByID always fails"), + } + // Allow one successful GetChatByID (the Subscribe preamble's + // initial state load). All subsequent calls return the mock + // error, exercising the reconnect-branch DB-error path. + failingDB.okRemain.Store(1) + + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(ctx, t, realDB) + chat := seedWaitingChat(ctx, t, realDB, org.ID, user, model, "relay-db-error") + + subscriber := newTestServer(t, failingDB, ps, subscriberID, dialer, mclk) + _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) + require.True(t, ok) + t.Cleanup(cancel) + + // Flip to running so the merge loop starts an async dial. The + // dial fails (attempts=1, reconnect scheduled). From there each + // reconnect timer fires, the merge loop calls GetChatByID, the + // failing DB returns an error, and retryState.next() increments. + // + // Budget: 1 dial-failure + 6 DB-failures = 7 next() calls; the + // 7th trips the cap. + setChatRunningAndPublish(ctx, t, realDB, ps, chat.ID, workerID) + for i := 0; i < 6; i++ { + call := trapReconnect.MustWait(ctx) + call.MustRelease(ctx) + mclk.Advance(call.Duration).MustWait(ctx) + } + + var errEvent *codersdk.ChatStreamEvent + require.Eventually(t, func() bool { + select { + case ev, open := <-events: + if !open { + return errEvent != nil + } + if ev.Type == codersdk.ChatStreamEventTypeError { + errEvent = &ev + return true + } + return false + default: + return false + } + }, testutil.WaitShort, testutil.IntervalFast, + "expected terminal error event after DB-error cap") + require.NotNil(t, errEvent.Error) + require.Contains(t, errEvent.Error.Message, "relay connection failed") + require.Contains(t, errEvent.Error.Message, "6") + + // Exactly 1 dial fired: the one that triggered the initial + // reconnect schedule. All subsequent next() calls come from the + // DB-error branch without calling the dialer. + require.Equal(t, int32(1), callCount.Load(), + "expected exactly 1 dial; reconnects should short-circuit on DB error") +} + +// TestRelayStopsImmediatelyOnUnauthorized tests the unrecoverable +// branch and its table of status codes. +func TestRelayStopsImmediatelyOnUnauthorized(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + status int + wantUnrecoverable bool + wantMsgContains string + }{ + {"401", http.StatusUnauthorized, true, "401"}, + {"403", http.StatusForbidden, true, "403"}, + {"500_intermittent", http.StatusInternalServerError, false, ""}, + {"zero_intermittent", 0, false, ""}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + workerID := uuid.New() + subscriberID := uuid.New() + + var callCount atomic.Int32 + dialer := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( + []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, + ) { + callCount.Add(1) + return nil, nil, nil, &entchatd.RelayDialError{ + HTTPStatus: tc.status, + Err: io.EOF, + } + } + + mclk := quartz.NewMock(t) + trapReconnect := mclk.Trap().NewTimer("reconnect") + defer trapReconnect.Close() + + subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) + + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(ctx, t, db) + chat := seedWaitingChat(ctx, t, db, org.ID, user, model, + "relay-unrec-"+tc.name) + + _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) + require.True(t, ok) + t.Cleanup(cancel) + + setChatRunningAndPublish(ctx, t, db, ps, chat.ID, workerID) + if tc.wantUnrecoverable { + // First dial should tear the relay down. + var errEvent *codersdk.ChatStreamEvent + require.Eventually(t, func() bool { + select { + case ev, open := <-events: + if !open { + return errEvent != nil + } + if ev.Type == codersdk.ChatStreamEventTypeError { + errEvent = &ev + return true + } + return false + default: + return false + } + }, testutil.WaitShort, testutil.IntervalFast, + "expected terminal error event") + require.NotNil(t, errEvent) + require.Contains(t, errEvent.Error.Message, "relay authentication failed") + require.Contains(t, errEvent.Error.Message, tc.wantMsgContains) + require.Equal(t, int32(1), callCount.Load(), + "unrecoverable errors must not retry; got %d dials", callCount.Load()) + } else { + // Intermittent: fire one reconnect timer + // and confirm the dialer is called again. + call := trapReconnect.MustWait(ctx) + call.MustRelease(ctx) + mclk.Advance(call.Duration).MustWait(ctx) + require.Eventually(t, func() bool { + return callCount.Load() >= 2 + }, testutil.WaitShort, testutil.IntervalFast, + "intermittent should retry at least once") + } + }) + } +} + +// TestRelayBackoffResetsOnStatusChange checks that closeRelay (driven +// by a status notification) resets the retry counter so subsequent +// dials against a new target start at the floor delay. +func TestRelayBackoffResetsOnStatusChange(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + workerID1 := uuid.New() + workerID2 := uuid.New() + subscriberID := uuid.New() + + dialer := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( + []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, + ) { + return nil, nil, nil, &entchatd.RelayDialError{ + HTTPStatus: http.StatusBadGateway, + Err: io.EOF, + } + } + + mclk := quartz.NewMock(t) + trapReconnect := mclk.Trap().NewTimer("reconnect") + defer trapReconnect.Close() + + subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) + + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(ctx, t, db) + chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-reset-on-status") + + _, _, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) + require.True(t, ok) + t.Cleanup(cancel) + + // Drive the async openRelayAsync path with workerID1. + setChatRunningAndPublish(ctx, t, db, ps, chat.ID, workerID1) + + // Drive 3 intermittent failures so attempts=3 and the delay + // has grown past the floor. After each loop iteration the 4th + // reconnect timer is queued - consume it too so our later + // assertion sees the reset's timer, not a stale one. + for i := 0; i < 3; i++ { + call := trapReconnect.MustWait(ctx) + call.MustRelease(ctx) + mclk.Advance(call.Duration).MustWait(ctx) + } + // Grab the next trapped timer (the grown one scheduled after + // the 3rd dial fails) but don't advance it - we want to see it + // replaced by a fresh floor-delay timer after the reset. + grown := trapReconnect.MustWait(ctx) + require.Greater(t, grown.Duration, 500*time.Millisecond, + "sanity: pre-reset delay should have grown past the floor") + grown.MustRelease(ctx) + + // Flip the chat to waiting; closeRelay runs (because the + // status notification no longer points at a running peer) and + // should reset the retry state. + _, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ + ID: chat.ID, + Status: database.ChatStatusWaiting, + }) + require.NoError(t, err) + waitingPayload, err := json.Marshal(coderdpubsub.ChatStreamNotifyMessage{ + Status: string(database.ChatStatusWaiting), + }) + require.NoError(t, err) + require.NoError(t, ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chat.ID), waitingPayload)) + + // Flip back to running on a different worker. This triggers a + // fresh openRelayAsync which fails, arming a reconnect timer. + // That timer's delay must be the floor, proving the reset. + setChatRunningAndPublish(ctx, t, db, ps, chat.ID, workerID2) + + call := trapReconnect.MustWait(ctx) + require.Equal(t, 500*time.Millisecond, call.Duration, + "retry state must reset after status change; got grown delay %v", call.Duration) + call.MustRelease(ctx) +} + +// TestRelayBackoffRespectsContextCancel is a regression guard: the +// reconnect timer must respect ctx cancellation promptly. +func TestRelayBackoffRespectsContextCancel(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + workerID := uuid.New() + subscriberID := uuid.New() + + dialer := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( + []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, + ) { + return nil, nil, nil, &entchatd.RelayDialError{ + HTTPStatus: http.StatusBadGateway, + Err: io.EOF, + } + } + + mclk := quartz.NewMock(t) + trapReconnect := mclk.Trap().NewTimer("reconnect") + defer trapReconnect.Close() + + subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) + + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(ctx, t, db) + chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-cancel") + + subCtx, subCancel := context.WithCancel(ctx) + _, events, cancel, ok := subscriber.Subscribe(subCtx, chat.ID, nil, 0) + require.True(t, ok) + t.Cleanup(cancel) + + setChatRunningAndPublish(ctx, t, db, ps, chat.ID, workerID) + + // Wait for the first reconnect timer to arm. + call := trapReconnect.MustWait(ctx) + call.MustRelease(ctx) + + // Cancel the subscriber context. The events channel should + // close promptly (the merge goroutine's select exits on + // ctx.Done). + subCancel() + + done := make(chan struct{}) + go func() { + defer close(done) + for { + if _, open := <-events; !open { + return + } + } + }() + select { + case <-done: + case <-time.After(testutil.WaitShort): + t.Fatal("events channel did not close after ctx cancel") + } +} + +// TestDialRelayReal401 exercises the real dialRelay path against an +// httptest server that returns 401 on the stream endpoint. It +// validates that the websocket library's handshake failure +// propagates through as *RelayDialError with HTTPStatus == 401. +// +// This is the one test that uses the real coder/websocket library +// on the failure path - a safety net against library upgrades +// silently breaking status capture. +func TestDialRelayReal401(t *testing.T) { + t.Parallel() + + // An httptest server that 401s every request on the stream + // endpoint. Any other path gets a 404. + srv := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { + if !streamPathRE.MatchString(r.URL.Path) { + http.NotFound(rw, r) + return + } + rw.Header().Set("Content-Type", "application/json") + rw.WriteHeader(http.StatusUnauthorized) + _, _ = rw.Write([]byte(`{"message":"unauthorized"}`)) + })) + t.Cleanup(srv.Close) + + db, _ := dbtestutil.NewDB(t) + workerID := uuid.New() + subscriberID := uuid.New() + + // Wire real config (no DialerFn override) so dialRelay runs + // end-to-end against the httptest server. Seeding a waiting + // chat (below) keeps Subscribe's initial synchronous dial a + // no-op; we then push a running status notification to the + // merge loop so it invokes dialRelay via the async path, where + // the 401 tear-down logic lives. + cfg := entchatd.MultiReplicaSubscribeConfig{ + ResolveReplicaAddress: func(_ context.Context, _ uuid.UUID) (string, bool) { + return srv.URL, true + }, + ReplicaHTTPClient: srv.Client(), + ReplicaIDFn: func() uuid.UUID { return subscriberID }, + } + subscribeFn := entchatd.NewMultiReplicaSubscribeFn(cfg) + + ctx := testutil.Context(t, testutil.WaitMedium) + user, org, model := seedChatDependencies(ctx, t, db) + // Seed a waiting chat - no sync dial - then push a running + // status notification to trigger the async dial via the real + // dialRelay path. + chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-real-401") + + statusCh := make(chan osschatd.StatusNotification, 1) + evs := subscribeFn(ctx, osschatd.SubscribeFnParams{ + ChatID: chat.ID, + Chat: chat, + WorkerID: subscriberID, + StatusNotifications: statusCh, + RequestHeader: http.Header{codersdk.SessionTokenHeader: {"test-token"}}, + DB: db, + Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + }) + + statusCh <- osschatd.StatusNotification{ + Status: database.ChatStatusRunning, + WorkerID: workerID, + } + + // Wait for a terminal error event. On a real 401 handshake, + // the classifier flags it unrecoverable β†’ one dial, then + // error event, then channel close. + var errEvent *codersdk.ChatStreamEvent + deadline := time.After(testutil.WaitMedium) +waitErr: + for { + select { + case ev, open := <-evs: + if !open { + break waitErr + } + if ev.Type == codersdk.ChatStreamEventTypeError { + errEvent = &ev + } + case <-deadline: + break waitErr + } + } + + require.NotNil(t, errEvent, "expected terminal error event from real 401 dial") + require.NotNil(t, errEvent.Error) + require.Contains(t, errEvent.Error.Message, "relay authentication failed") + require.Contains(t, errEvent.Error.Message, "401") +} + +// streamPathRE matches the chat stream endpoint path built by +// buildRelayURL. Compiled at package scope so the httptest handler +// below doesn't pay regexp.Compile per request. +var streamPathRE = regexp.MustCompile( + `^/api/experimental/chats/[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}/stream$`, +)