Skip to content

Commit d3c77df

Browse files
authored
fix: add mutex to SSEServer to avoid data race between Start and Shutdown; fix test error on Windows (#166 #172) (#170)
* Add mutex to avoid data race bewteen Start and Shutdown in SSEServer struct * delete the mutex for session * Added checks for Windows * Update sse.go * Update stdio_test.go
1 parent 71b910b commit d3c77df

File tree

2 files changed

+24
-5
lines changed

2 files changed

+24
-5
lines changed

client/transport/stdio_test.go

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"os"
88
"os/exec"
99
"path/filepath"
10+
"runtime"
1011
"sync"
1112
"testing"
1213
"time"
@@ -31,6 +32,10 @@ func compileTestServer(outputPath string) error {
3132
func TestStdio(t *testing.T) {
3233
// Compile mock server
3334
mockServerPath := filepath.Join(os.TempDir(), "mockstdio_server")
35+
// Add .exe suffix on Windows
36+
if runtime.GOOS == "windows" {
37+
mockServerPath += ".exe"
38+
}
3439
if err := compileTestServer(mockServerPath); err != nil {
3540
t.Fatalf("Failed to compile mock server: %v", err)
3641
}
@@ -302,16 +307,19 @@ func TestStdioErrors(t *testing.T) {
302307
})
303308

304309
t.Run("RequestBeforeStart", func(t *testing.T) {
305-
// 创建一个新的 Stdio 实例但不调用 Start 方法
306310
mockServerPath := filepath.Join(os.TempDir(), "mockstdio_server")
311+
// Add .exe suffix on Windows
312+
if runtime.GOOS == "windows" {
313+
mockServerPath += ".exe"
314+
}
307315
if err := compileTestServer(mockServerPath); err != nil {
308316
t.Fatalf("Failed to compile mock server: %v", err)
309317
}
310318
defer os.Remove(mockServerPath)
311319

312320
uninitiatedStdio := NewStdio(mockServerPath, nil)
313321

314-
// 准备一个请求
322+
// Prepare a request
315323
request := JSONRPCRequest{
316324
JSONRPC: "2.0",
317325
ID: 99,
@@ -331,6 +339,10 @@ func TestStdioErrors(t *testing.T) {
331339
t.Run("RequestAfterClose", func(t *testing.T) {
332340
// Compile mock server
333341
mockServerPath := filepath.Join(os.TempDir(), "mockstdio_server")
342+
// Add .exe suffix on Windows
343+
if runtime.GOOS == "windows" {
344+
mockServerPath += ".exe"
345+
}
334346
if err := compileTestServer(mockServerPath); err != nil {
335347
t.Fatalf("Failed to compile mock server: %v", err)
336348
}

server/sse.go

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,8 @@ type SSEServer struct {
6565

6666
keepAlive bool
6767
keepAliveInterval time.Duration
68+
69+
mu sync.RWMutex
6870
}
6971

7072
// SSEOption defines a function type for configuring SSEServer
@@ -189,18 +191,24 @@ func NewTestServer(server *MCPServer, opts ...SSEOption) *httptest.Server {
189191
// Start begins serving SSE connections on the specified address.
190192
// It sets up HTTP handlers for SSE and message endpoints.
191193
func (s *SSEServer) Start(addr string) error {
194+
s.mu.Lock()
192195
s.srv = &http.Server{
193196
Addr: addr,
194197
Handler: s,
195198
}
199+
s.mu.Unlock()
196200

197201
return s.srv.ListenAndServe()
198202
}
199203

200204
// Shutdown gracefully stops the SSE server, closing all active sessions
201205
// and shutting down the HTTP server.
202206
func (s *SSEServer) Shutdown(ctx context.Context) error {
203-
if s.srv != nil {
207+
s.mu.RLock()
208+
srv := s.srv
209+
s.mu.RUnlock()
210+
211+
if srv != nil {
204212
s.sessions.Range(func(key, value interface{}) bool {
205213
if session, ok := value.(*sseSession); ok {
206214
close(session.done)
@@ -209,7 +217,7 @@ func (s *SSEServer) Shutdown(ctx context.Context) error {
209217
return true
210218
})
211219

212-
return s.srv.Shutdown(ctx)
220+
return srv.Shutdown(ctx)
213221
}
214222
return nil
215223
}
@@ -335,7 +343,6 @@ func (s *SSEServer) handleMessage(w http.ResponseWriter, r *http.Request) {
335343
s.writeJSONRPCError(w, nil, mcp.INVALID_PARAMS, "Missing sessionId")
336344
return
337345
}
338-
339346
sessionI, ok := s.sessions.Load(sessionID)
340347
if !ok {
341348
s.writeJSONRPCError(w, nil, mcp.INVALID_PARAMS, "Invalid session ID")

0 commit comments

Comments
 (0)