diff --git a/sandbox/ipc/session.go b/sandbox/ipc/session.go index 3a6811b8..391db98c 100644 --- a/sandbox/ipc/session.go +++ b/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() diff --git a/sandbox/ipc/types.go b/sandbox/ipc/types.go index 4219bff2..b5c3c183 100644 --- a/sandbox/ipc/types.go +++ b/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