mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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:
+18
-16
@@ -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
|
||||
}()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user