diff --git a/pkg/tools/session.go b/pkg/tools/session.go index 8c7584254..65aa924fd 100644 --- a/pkg/tools/session.go +++ b/pkg/tools/session.go @@ -171,25 +171,44 @@ func (s *ProcessSession) ToSessionInfo() SessionInfo { type SessionManager struct { mu sync.RWMutex sessions map[string]*ProcessSession + stopCh chan struct{} } 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 { - sm.cleanupOldSessions() + for { + select { + case <-sm.stopCh: + return + case <-ticker.C: + sm.cleanupOldSessions() + } } }() return sm } +// Stop shuts down the background cleanup goroutine. Safe to call multiple times. +// After Stop returns, the SessionManager is still usable — only the cleanup +// goroutine is terminated. +func (sm *SessionManager) Stop() { + select { + case <-sm.stopCh: + // already closed + default: + close(sm.stopCh) + } +} + // cleanupOldSessions removes sessions that are done and older than 30 minutes func (sm *SessionManager) cleanupOldSessions() { sm.mu.Lock() diff --git a/pkg/tools/session_test.go b/pkg/tools/session_test.go index 6cfe72a10..96e3410e7 100644 --- a/pkg/tools/session_test.go +++ b/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", diff --git a/pkg/tools/shell_test.go b/pkg/tools/shell_test.go index 14d5a6697..4ad1f4ee8 100644 --- a/pkg/tools/shell_test.go +++ b/pkg/tools/shell_test.go @@ -1,4 +1,4 @@ -package tools +package tools import ( "context" @@ -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")