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
This commit is contained in:
Ethan
2026-07-01 00:16:25 +10:00
committed by GitHub
parent d179266cc7
commit d219f96ba5
+18 -16
View File
@@ -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
}()