From 059ed7ab5cf3b9780f6c95b399e9040607a6a9b5 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@coder.com>
Date: Tue, 3 Mar 2026 11:34:08 -0500
Subject: [PATCH] fix(chatd): return chat to pending when server shuts down
 during successful completion (#22559)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Problem

Flaky test:
`TestCloseDuringShutdownContextCanceledShouldRetryOnNewReplica`
(coder/internal#1371)

The test intermittently fails because the chat ends up in `waiting`
status instead of `pending` after server shutdown.

## Root Cause

There is a race condition in `processChat` where `runChat` completes
successfully just as the server context is being canceled during
`Close()`. The sequence:

1. Server calls `Close()`, canceling the server context.
2. The LLM HTTP response has already been fully written by the mock
server (the stream closes normally before context cancellation
propagates to the HTTP client).
3. `runChat` returns `nil` (success) instead of `context.Canceled`.
4. The existing `isShutdownCancellation` check only runs when `runChat`
returns an error, so the shutdown is not detected.
5. `processChat`'s deferred cleanup marks the chat as `waiting` instead
of `pending`.
6. The test's assertion that the chat is `pending` never becomes true.

This race is timing-dependent — it only triggers when the mock server's
HTTP response completes in the narrow window between context
cancellation being initiated and it propagating through the HTTP
transport layer.

## Fix

Add a server context check after `runChat` returns successfully. If the
server is shutting down (`ctx.Err() != nil`), override the status to
`pending` so another replica can pick up the chat.

This is the same pattern already used for the error path
(`isShutdownCancellation`), extended to cover the success path.
---
 coderd/chatd/chatd.go | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/coderd/chatd/chatd.go b/coderd/chatd/chatd.go
index 133f400192..4fff70af4f 100644
--- a/coderd/chatd/chatd.go
+++ b/coderd/chatd/chatd.go
@@ -1846,6 +1846,21 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) {
 		status = database.ChatStatusError
 		return
 	}
+
+	// If runChat completed successfully but the server context was
+	// canceled (e.g. during Close()), the chat should be returned
+	// to pending so another replica can pick it up. There is a
+	// race where the LLM stream finishes just as the server is
+	// shutting down — the HTTP response completes before context
+	// cancellation propagates, so runChat returns nil instead of
+	// a context.Canceled error. Without this check the chat would
+	// be marked "waiting" and never retried.
+	if ctx.Err() != nil {
+		logger.Info(ctx, "chat completed during shutdown; returning to pending")
+		status = database.ChatStatusPending
+		lastError = ""
+		return
+	}
 }
 
 func isShutdownCancellation(