From d219f96ba53510d057bcca6e0e3bcf80404f9bd5 Mon Sep 17 00:00:00 2001
From: Ethan <39577870+ethanndickson@users.noreply.github.com>
Date: Wed, 1 Jul 2026 00:16:25 +1000
Subject: [PATCH] fix(cli): join MCP reporter and watcher goroutines before
 exit (#26847)

## Problem

`TestExpMcpReporter/Reconnect` flakes under the race detector with a
data race on the shared `*serpent.Invocation`'s `inv.Stderr` field.

The MCP server's reporter and watcher goroutines write status warnings
via `cliui.Warnf(inv.Stderr, ...)`, but they were launched
fire-and-forget with nothing tying their lifetime to the command
handler. On shutdown, `startServer`'s deferred restore of
`inv.Stdin/Stdout/Stderr` could run concurrently with a still-running
goroutine reading `inv.Stderr`, which the race detector flags. The
reporter's error suppression only swallows `context.Canceled`, so a
shutdown error from an in-flight `UpdateAppStatus` RPC (a drpc "closed"
error, not `context.Canceled`) reaches the `Warnf` call and races the
restore.

## Fix

Track the reporter and watcher goroutines on a `sync.WaitGroup`. After
`startServer` returns, cancel the context, close the queue and socket
client, then `wg.Wait()` for the goroutines to exit before returning.
All three unblocks are needed: cancel stops the watcher retry loop and a
reporter blocked on `Pop`, `queue.Close` also unblocks `Pop`, and
`socketClient.Close` unblocks a reporter parked in an in-flight RPC.

This also removes the stdin/stdout/stderr save/restore in `startServer`,
which only ever wrote back identical values and was the racing write.

This mirrors the existing precedent in `cli/ssh.go`, where a
`sync.WaitGroup` guards against "logging while closing the log file in a
defer."

Verified with `go test ./cli -run 'TestExpMcpReporter/Reconnect' -race
-count=50` (the reproducer from the issue) plus a 240-execution parallel
stress run of the full `TestExpMcp` suite under `-race`, all green.

Closes CODAGT-710
Closes https://github.com/coder/internal/issues/1610
---
 cli/exp_mcp.go | 34 ++++++++++++++++++----------------
 1 file changed, 18 insertions(+), 16 deletions(-)

diff --git a/cli/exp_mcp.go b/cli/exp_mcp.go
index f0013afb52..6d72439b8c 100644
--- a/cli/exp_mcp.go
+++ b/cli/exp_mcp.go
@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"slices"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/mark3labs/mcp-go/mcp"
@@ -388,6 +389,9 @@ type mcpServer struct {
 	client           *codersdk.Client
 	aiAgentAPIClient *agentapi.Client
 	queue            *cliutil.Queue[taskReport]
+	// wg tracks the reporter and watcher goroutines, which write to
+	// inv.Stderr and must not outlive the handler.
+	wg sync.WaitGroup
 }
 
 func (r *RootCmd) mcpServer() *serpent.Command {
@@ -534,10 +538,6 @@ func (r *RootCmd) mcpServer() *serpent.Command {
 
 			ctx, cancel := context.WithCancel(inv.Context())
 			defer cancel()
-			defer srv.queue.Close()
-			if srv.socketClient != nil {
-				defer srv.socketClient.Close()
-			}
 
 			// Start the reporter, watcher, and server.  These are all tied to the
 			// lifetime of the MCP server, which is itself tied to the lifetime of the
@@ -548,7 +548,15 @@ func (r *RootCmd) mcpServer() *serpent.Command {
 					srv.startWatcher(ctx, inv)
 				}
 			}
-			return srv.startServer(ctx, inv, instructions, allowedTools)
+			serveErr := srv.startServer(ctx, inv, instructions, allowedTools)
+
+			cancel()
+			srv.queue.Close()
+			if srv.socketClient != nil {
+				_ = srv.socketClient.Close()
+			}
+			srv.wg.Wait()
+			return serveErr
 		},
 		Short: "Start the Coder MCP server.",
 		Options: []serpent.Option{
@@ -592,7 +600,9 @@ func (r *RootCmd) mcpServer() *serpent.Command {
 }
 
 func (s *mcpServer) startReporter(ctx context.Context, inv *serpent.Invocation) {
+	s.wg.Add(1)
 	go func() {
+		defer s.wg.Done()
 		for {
 			// TODO: Even with the queue, there is still the potential that a message
 			// from the screen watcher and a message from the AI agent could arrive
@@ -622,7 +632,9 @@ func (s *mcpServer) startReporter(ctx context.Context, inv *serpent.Invocation)
 }
 
 func (s *mcpServer) startWatcher(ctx context.Context, inv *serpent.Invocation) {
+	s.wg.Add(1)
 	go func() {
+		defer s.wg.Done()
 		for retrier := retry.New(time.Second, 30*time.Second); retrier.Wait(ctx); {
 			eventsCh, errCh, err := s.aiAgentAPIClient.SubscribeEvents(ctx)
 			if err == nil {
@@ -680,16 +692,6 @@ func (s *mcpServer) startServer(ctx context.Context, inv *serpent.Invocation, in
 		cliui.Infof(inv.Stderr, "Allowed Tools  : %v", allowedTools)
 	}
 
-	// Capture the original stdin, stdout, and stderr.
-	invStdin := inv.Stdin
-	invStdout := inv.Stdout
-	invStderr := inv.Stderr
-	defer func() {
-		inv.Stdin = invStdin
-		inv.Stdout = invStdout
-		inv.Stderr = invStderr
-	}()
-
 	mcpSrv := server.NewMCPServer(
 		"Coder Agent",
 		buildinfo.Version(),
@@ -756,7 +758,7 @@ func (s *mcpServer) startServer(ctx context.Context, inv *serpent.Invocation, in
 	done := make(chan error)
 	go func() {
 		defer close(done)
-		srvErr := srv.Listen(ctx, invStdin, invStdout)
+		srvErr := srv.Listen(ctx, inv.Stdin, inv.Stdout)
 		done <- srvErr
 	}()