Refactor session cleanup logic to ensure single execution and improve resource management

- Replace direct cleanup calls in the Close method with a sync.Once mechanism to guarantee that cleanup runs only once.
- Remove the cleanup method and integrate its functionality into the Close method, streamlining resource management.
- Update the serve method to call Close for cleanup, enhancing clarity and consistency in session termination.

sandbox/ipc/session.go

@@ -24,23 +24,24 @@ func (s *Session) SetContext(ctx *AgentContext) {
 
 // Close closes the session and cleans up resources
 func (s *Session) Close() error {
-	if s.cancel != nil {
-		s.cancel()
-	}
-	if s.Conn != nil {
-		s.Conn.Close()
-	}
-	if s.Listener != nil {
-		s.Listener.Close()
-	}
-	// Remove socket file
-	os.Remove(s.SocketPath)
+	s.closeOnce.Do(func() {
+		if s.cancel != nil {
+			s.cancel()
+		}
+		if s.Conn != nil {
+			s.Conn.Close()
+		}
+		if s.Listener != nil {
+			s.Listener.Close()
+		}
+		os.Remove(s.SocketPath)
+	})
 	return nil
 }
 
 // serve handles incoming connections
 func (s *Session) serve(ctx context.Context) {
-	defer s.cleanup()
+	defer s.Close()
 
 	for {
 		select {
@@ -49,10 +50,8 @@ func (s *Session) serve(ctx context.Context) {
 		default:
 		}
 
-		// Accept connection with deadline to allow context cancellation check
 		conn, err := s.Listener.Accept()
 		if err != nil {
-			// Check if context was cancelled
 			select {
 			case <-ctx.Done():
 				return
@@ -66,17 +65,6 @@ func (s *Session) serve(ctx context.Context) {
 	}
 }
 
-// cleanup cleans up session resources
-func (s *Session) cleanup() {
-	if s.Conn != nil {
-		s.Conn.Close()
-	}
-	if s.Listener != nil {
-		s.Listener.Close()
-	}
-	os.Remove(s.SocketPath)
-}
-
 // handleConnection handles a single connection
 func (s *Session) handleConnection(ctx context.Context, conn net.Conn) {
 	defer conn.Close()

sandbox/ipc/types.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"net"
+	"sync"
 )
 
 // Session represents an IPC session for a sandbox container
@@ -15,6 +16,7 @@ type Session struct {
 	Context    *AgentContext       // Agent context
 	MCPTools   map[string]*MCPTool // Authorized MCP tools
 	cancel     context.CancelFunc  // Cancel function for cleanup
+	closeOnce  sync.Once           // Ensures cleanup runs exactly once
 }
 
 // AgentContext holds context information for the agent