From 16e523a9c92e1b8172441b0de7f71adcb7e95488 Mon Sep 17 00:00:00 2001
From: Manuel Odendahl <wesen@ruinwesen.com>
Date: Tue, 21 Jan 2025 09:25:01 -0500
Subject: [PATCH] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20Simplify=20stdio=20server?=
 =?UTF-8?q?=20shutdown=20using=20only=20context?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 changelog.md                         | 20 ++++++++++++++++++-
 pkg/server/transports/stdio/stdio.go | 30 ++++++++--------------------
 2 files changed, 27 insertions(+), 23 deletions(-)

diff --git a/changelog.md b/changelog.md
index 0d6830e..9ec2e2b 100644
--- a/changelog.md
+++ b/changelog.md
@@ -149,4 +149,22 @@ Enhanced process group and signal handling in stdio transport:
 - Send signals to entire process group instead of just the main process
 - Added fallback to direct process signals if process group handling fails
 - Improved debug logging for process and signal management
-- Fixed issue with signals not being properly received by the server 
\ No newline at end of file
+- Fixed issue with signals not being properly received by the server 
+
+# Fixed Channel Close Race Condition
+
+Fixed a race condition in the stdio server where the done channel could be closed multiple times:
+
+- Added sync.Once to ensure done channel is only closed once
+- Fixed potential panic during server shutdown
+- Improved shutdown coordination between signal handling and Stop method
+- Enhanced logging around channel closing operations 
+
+# Simplified Server Shutdown
+
+Simplified stdio server shutdown by using only context-based coordination:
+
+- Removed redundant done channel in favor of context cancellation
+- Added dedicated scanner context for cleaner shutdown
+- Simplified shutdown logic and error handling
+- Improved logging messages for shutdown events 
\ No newline at end of file
diff --git a/pkg/server/transports/stdio/stdio.go b/pkg/server/transports/stdio/stdio.go
index 909ce95..1a1351b 100644
--- a/pkg/server/transports/stdio/stdio.go
+++ b/pkg/server/transports/stdio/stdio.go
@@ -25,7 +25,6 @@ type Server struct {
 	resourceService   services.ResourceService
 	toolService       services.ToolService
 	initializeService services.InitializeService
-	done              chan struct{}
 	signalChan        chan os.Signal
 }
 
@@ -54,7 +53,6 @@ func NewServer(logger zerolog.Logger, ps services.PromptService, rs services.Res
 		resourceService:   rs,
 		toolService:       ts,
 		initializeService: is,
-		done:              make(chan struct{}),
 		signalChan:        make(chan os.Signal, 1),
 	}
 }
@@ -70,19 +68,17 @@ func (s *Server) Start(ctx context.Context) error {
 	// Create a channel for scanner errors
 	scanErrChan := make(chan error, 1)
 
+	// Create a cancellable context for the scanner
+	scannerCtx, cancelScanner := context.WithCancel(ctx)
+	defer cancelScanner()
+
 	// Start scanning in a goroutine
 	go func() {
 		for s.scanner.Scan() {
 			select {
-			case <-s.done:
-				s.logger.Debug().Msg("Received done signal, stopping scanner")
-				scanErrChan <- nil
-				return
-			case <-ctx.Done():
-				s.logger.Debug().
-					Err(ctx.Err()).
-					Msg("Context cancelled, stopping scanner")
-				scanErrChan <- ctx.Err()
+			case <-scannerCtx.Done():
+				s.logger.Debug().Msg("Context cancelled, stopping scanner")
+				scanErrChan <- scannerCtx.Err()
 				return
 			default:
 				line := s.scanner.Text()
@@ -114,13 +110,12 @@ func (s *Server) Start(ctx context.Context) error {
 		s.logger.Debug().
 			Str("signal", sig.String()).
 			Msg("Received signal in stdio server")
-		close(s.done)
+		cancelScanner()
 		return nil
 	case <-ctx.Done():
 		s.logger.Debug().
 			Err(ctx.Err()).
 			Msg("Context cancelled in stdio server")
-		close(s.done)
 		return ctx.Err()
 	case err := <-scanErrChan:
 		if err == io.EOF {
@@ -138,15 +133,6 @@ func (s *Server) Start(ctx context.Context) error {
 func (s *Server) Stop(ctx context.Context) error {
 	s.logger.Info().Msg("Stopping stdio server")
 
-	select {
-	case <-s.done:
-		s.logger.Debug().Msg("Server already stopped")
-		return nil
-	default:
-		s.logger.Debug().Msg("Closing done channel to signal shutdown")
-		close(s.done)
-	}
-
 	// Wait for context to be done or timeout
 	select {
 	case <-ctx.Done():