Merge pull request #2986 from chengzhichao-xydt/codex/session-manager-stop-cleanup

fix(tools): add Stop() to SessionManager to prevent goroutine leak

pkg/tools/session.go

@@ -171,25 +171,42 @@ func (s *ProcessSession) ToSessionInfo() SessionInfo {
 type SessionManager struct {
 	mu       sync.RWMutex
 	sessions map[string]*ProcessSession
+	stopCh   chan struct{}
+	stopOnce sync.Once
 }
 
 func NewSessionManager() *SessionManager {
 	sm := &SessionManager{
 		sessions: make(map[string]*ProcessSession),
+		stopCh:   make(chan struct{}),
 	}
 
 	// Start cleaner goroutine - runs every 5 minutes, cleans up sessions done for >30 minutes
 	go func() {
 		ticker := time.NewTicker(5 * time.Minute)
 		defer ticker.Stop()
-		for range ticker.C {
+		for {
-			sm.cleanupOldSessions()
+			select {
+			case <-sm.stopCh:
+				return
+			case <-ticker.C:
+				sm.cleanupOldSessions()
+			}
 		}
 	}()
 
 	return sm
 }
 
+// Stop shuts down the background cleanup goroutine. Safe to call multiple
+// times from concurrent goroutines. After Stop returns, the SessionManager
+// is still usable — only the cleanup goroutine is terminated.
+func (sm *SessionManager) Stop() {
+	sm.stopOnce.Do(func() {
+		close(sm.stopCh)
+	})
+}
+
 // cleanupOldSessions removes sessions that are done and older than 30 minutes
 func (sm *SessionManager) cleanupOldSessions() {
 	sm.mu.Lock()

pkg/tools/session_test.go

@@ -8,6 +8,7 @@ import (
 
 func TestSessionManager_AddGet(t *testing.T) {
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	session := &ProcessSession{
 		ID:        "test-1",
 		Command:   "echo hello",
@@ -24,6 +25,7 @@ func TestSessionManager_AddGet(t *testing.T) {
 
 func TestSessionManager_Remove(t *testing.T) {
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	session := &ProcessSession{
 		ID:        "test-1",
 		Command:   "echo hello",
@@ -39,6 +41,7 @@ func TestSessionManager_Remove(t *testing.T) {
 
 func TestSessionManager_List(t *testing.T) {
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	sm.Add(&ProcessSession{
 		ID:        "test-1",
 		Command:   "echo hello",

pkg/tools/shell_test.go

@@ -907,6 +907,7 @@ func TestShellTool_List_Empty(t *testing.T) {
 	require.NoError(t, err)
 
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	tool.sessionManager = sm
 
 	ctx := context.Background()
@@ -922,6 +923,7 @@ func TestShellTool_RunBackground_List(t *testing.T) {
 	require.NoError(t, err)
 
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	tool.sessionManager = sm
 
 	ctx := WithToolContext(context.Background(), "cli", "test")
@@ -961,6 +963,7 @@ func TestShellTool_Read_Output(t *testing.T) {
 	require.NoError(t, err)
 
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	tool.sessionManager = sm
 
 	ctx := WithToolContext(context.Background(), "cli", "test")
@@ -995,6 +998,7 @@ func TestShellTool_Kill(t *testing.T) {
 	require.NoError(t, err)
 
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	tool.sessionManager = sm
 
 	ctx := WithToolContext(context.Background(), "cli", "test")
@@ -1034,6 +1038,7 @@ func TestShellTool_PTY_AllowedCommands(t *testing.T) {
 	require.NoError(t, err)
 
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	tool.sessionManager = sm
 
 	ctx := WithToolContext(context.Background(), "cli", "test")
@@ -1069,6 +1074,7 @@ func TestShellTool_PTY_WriteRead(t *testing.T) {
 	require.NoError(t, err)
 
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	tool.sessionManager = sm
 
 	ctx := WithToolContext(context.Background(), "cli", "test")
@@ -1128,6 +1134,7 @@ func TestShellTool_PTY_Poll(t *testing.T) {
 	require.NoError(t, err)
 
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	tool.sessionManager = sm
 
 	ctx := WithToolContext(context.Background(), "cli", "test")
@@ -1181,6 +1188,7 @@ func TestShellTool_PTY_Kill(t *testing.T) {
 	require.NoError(t, err)
 
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	tool.sessionManager = sm
 
 	ctx := WithToolContext(context.Background(), "cli", "test")
@@ -1226,6 +1234,7 @@ func TestShellTool_Write_Read_NonPTY(t *testing.T) {
 	require.NoError(t, err)
 
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	tool.sessionManager = sm
 
 	ctx := WithToolContext(context.Background(), "cli", "test")
@@ -1279,6 +1288,7 @@ func TestShellTool_Read_NonPTY_Running(t *testing.T) {
 	require.NoError(t, err)
 
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	tool.sessionManager = sm
 
 	ctx := WithToolContext(context.Background(), "cli", "test")
@@ -1347,6 +1357,7 @@ func TestShellTool_ProcessGroupKill(t *testing.T) {
 	require.NoError(t, err)
 
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	tool.sessionManager = sm
 
 	ctx := WithToolContext(context.Background(), "cli", "test")
@@ -1407,6 +1418,7 @@ func TestShellTool_PTY_ProcessGroupKill(t *testing.T) {
 	require.NoError(t, err)
 
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	tool.sessionManager = sm
 
 	ctx := WithToolContext(context.Background(), "cli", "test")
@@ -1459,6 +1471,7 @@ func TestShellTool_PTY_Background_Read(t *testing.T) {
 	require.NoError(t, err)
 
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	tool.sessionManager = sm
 
 	ctx := WithToolContext(context.Background(), "cli", "test")
@@ -1499,6 +1512,7 @@ func TestShellTool_PTY_Background_ReadNoBlock(t *testing.T) {
 	require.NoError(t, err)
 
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	tool.sessionManager = sm
 
 	ctx := WithToolContext(context.Background(), "cli", "test")
@@ -1543,6 +1557,7 @@ func TestShellTool_Poll_Status(t *testing.T) {
 	require.NoError(t, err)
 
 	sm := NewSessionManager()
+	t.Cleanup(sm.Stop)
 	tool.sessionManager = sm
 
 	ctx := WithToolContext(context.Background(), "cli", "test")