From d750e64608998f6f3a03928bba18428f576b412f Mon Sep 17 00:00:00 2001 From: Xavier Bouthillier Date: Thu, 26 Jun 2025 18:25:11 -0400 Subject: [PATCH] feat: add MCP tracking to sessions (#19) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add mcps field to Session model to track active MCP servers and populate it from container labels in ContainerManager. Enhance MCP remove command to warn when removing servers used by active sessions. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-authored-by: Claude --- cubbi/container.py | 9 ++ cubbi/models.py | 1 + tests/test_mcp_commands.py | 217 ++++++++++++++++++++++++++++++++++--- 3 files changed, 214 insertions(+), 13 deletions(-) diff --git a/cubbi/container.py b/cubbi/container.py index dbf75cc..cb565fd 100644 --- a/cubbi/container.py +++ b/cubbi/container.py @@ -107,12 +107,21 @@ class ContainerManager: elif container.status == "created": status = SessionStatus.CREATING + # Get MCP list from container labels + mcps_str = labels.get("cubbi.mcps", "") + mcps = ( + [mcp.strip() for mcp in mcps_str.split(",") if mcp.strip()] + if mcps_str + else [] + ) + session = Session( id=session_id, name=labels.get("cubbi.session.name", f"cubbi-{session_id}"), image=labels.get("cubbi.image", "unknown"), status=status, container_id=container_id, + mcps=mcps, ) # Get port mappings diff --git a/cubbi/models.py b/cubbi/models.py index c0e2cc6..8b310be 100644 --- a/cubbi/models.py +++ b/cubbi/models.py @@ -102,6 +102,7 @@ class Session(BaseModel): status: SessionStatus container_id: Optional[str] = None ports: Dict[int, int] = Field(default_factory=dict) + mcps: List[str] = Field(default_factory=list) class Config(BaseModel): diff --git a/tests/test_mcp_commands.py b/tests/test_mcp_commands.py index 7821a86..76e4679 100644 --- a/tests/test_mcp_commands.py +++ b/tests/test_mcp_commands.py @@ -93,21 +93,212 @@ def test_mcp_remove(cli_runner, patched_config_manager): ], ) - # Mock the get_mcp and remove_mcp methods - with patch("cubbi.cli.mcp_manager.get_mcp") as mock_get_mcp: - # First make get_mcp return our MCP - mock_get_mcp.return_value = { - "name": "test-mcp", - "type": "remote", - "url": "http://test-server.com/sse", - "headers": {"Authorization": "Bearer test-token"}, - } + # Mock the container_manager.list_sessions to return sessions without MCPs + with patch("cubbi.cli.container_manager.list_sessions") as mock_list_sessions: + mock_list_sessions.return_value = [] - # Remove the MCP server - result = cli_runner.invoke(app, ["mcp", "remove", "test-mcp"]) + # Mock the remove_mcp method + with patch("cubbi.cli.mcp_manager.remove_mcp") as mock_remove_mcp: + # Make remove_mcp return True (successful removal) + mock_remove_mcp.return_value = True - # Just check it ran successfully with exit code 0 - assert result.exit_code == 0 + # Remove the MCP server + result = cli_runner.invoke(app, ["mcp", "remove", "test-mcp"]) + + # Just check it ran successfully with exit code 0 + assert result.exit_code == 0 + assert "Removed MCP server 'test-mcp'" in result.stdout + + +def test_mcp_remove_with_active_sessions(cli_runner, patched_config_manager): + """Test removing an MCP server that is used by active sessions.""" + from cubbi.models import Session, SessionStatus + + # Add a remote MCP server + patched_config_manager.set( + "mcps", + [ + { + "name": "test-mcp", + "type": "remote", + "url": "http://test-server.com/sse", + "headers": {"Authorization": "Bearer test-token"}, + } + ], + ) + + # Create mock sessions that use the MCP + mock_sessions = [ + Session( + id="session-1", + name="test-session-1", + image="goose", + status=SessionStatus.RUNNING, + container_id="container-1", + mcps=["test-mcp", "other-mcp"], + ), + Session( + id="session-2", + name="test-session-2", + image="goose", + status=SessionStatus.RUNNING, + container_id="container-2", + mcps=["other-mcp"], # This one doesn't use test-mcp + ), + Session( + id="session-3", + name="test-session-3", + image="goose", + status=SessionStatus.RUNNING, + container_id="container-3", + mcps=["test-mcp"], # This one uses test-mcp + ), + ] + + # Mock the container_manager.list_sessions to return our sessions + with patch("cubbi.cli.container_manager.list_sessions") as mock_list_sessions: + mock_list_sessions.return_value = mock_sessions + + # Mock the remove_mcp method + with patch("cubbi.cli.mcp_manager.remove_mcp") as mock_remove_mcp: + # Make remove_mcp return True (successful removal) + mock_remove_mcp.return_value = True + + # Remove the MCP server + result = cli_runner.invoke(app, ["mcp", "remove", "test-mcp"]) + + # Check it ran successfully with exit code 0 + assert result.exit_code == 0 + assert "Removed MCP server 'test-mcp'" in result.stdout + # Check warning about affected sessions + assert ( + "Warning: Found 2 active sessions using MCP 'test-mcp'" in result.stdout + ) + assert "session-1" in result.stdout + assert "session-3" in result.stdout + # session-2 should not be mentioned since it doesn't use test-mcp + assert "session-2" not in result.stdout + + +def test_mcp_remove_nonexistent(cli_runner, patched_config_manager): + """Test removing a non-existent MCP server.""" + # No MCPs configured + patched_config_manager.set("mcps", []) + + # Mock the container_manager.list_sessions to return empty list + with patch("cubbi.cli.container_manager.list_sessions") as mock_list_sessions: + mock_list_sessions.return_value = [] + + # Mock the remove_mcp method to return False (MCP not found) + with patch("cubbi.cli.mcp_manager.remove_mcp") as mock_remove_mcp: + mock_remove_mcp.return_value = False + + # Try to remove a non-existent MCP server + result = cli_runner.invoke(app, ["mcp", "remove", "nonexistent-mcp"]) + + # Check it ran successfully but reported not found + assert result.exit_code == 0 + assert "MCP server 'nonexistent-mcp' not found" in result.stdout + + +def test_session_mcps_attribute(): + """Test that Session model has mcps attribute and can be populated correctly.""" + from cubbi.models import Session, SessionStatus + + # Test that Session can be created with mcps attribute + session = Session( + id="test-session", + name="test-session", + image="goose", + status=SessionStatus.RUNNING, + container_id="test-container", + mcps=["mcp1", "mcp2"], + ) + + assert session.mcps == ["mcp1", "mcp2"] + + # Test that Session can be created with empty mcps list + session_empty = Session( + id="test-session-2", + name="test-session-2", + image="goose", + status=SessionStatus.RUNNING, + container_id="test-container-2", + ) + + assert session_empty.mcps == [] # Should default to empty list + + +def test_session_mcps_from_container_labels(): + """Test that Session mcps are correctly populated from container labels.""" + from unittest.mock import Mock + from cubbi.container import ContainerManager + + # Mock a container with MCP labels + mock_container = Mock() + mock_container.id = "test-container-id" + mock_container.status = "running" + mock_container.labels = { + "cubbi.session": "true", + "cubbi.session.id": "test-session", + "cubbi.session.name": "test-session-name", + "cubbi.image": "goose", + "cubbi.mcps": "mcp1,mcp2,mcp3", # Test with multiple MCPs + } + mock_container.attrs = {"NetworkSettings": {"Ports": {}}} + + # Mock Docker client + mock_client = Mock() + mock_client.containers.list.return_value = [mock_container] + + # Create container manager with mocked client + with patch("cubbi.container.docker.from_env") as mock_docker: + mock_docker.return_value = mock_client + mock_client.ping.return_value = True + + container_manager = ContainerManager() + sessions = container_manager.list_sessions() + + assert len(sessions) == 1 + session = sessions[0] + assert session.id == "test-session" + assert session.mcps == ["mcp1", "mcp2", "mcp3"] + + +def test_session_mcps_from_empty_container_labels(): + """Test that Session mcps are correctly handled when container has no MCP labels.""" + from unittest.mock import Mock + from cubbi.container import ContainerManager + + # Mock a container without MCP labels + mock_container = Mock() + mock_container.id = "test-container-id" + mock_container.status = "running" + mock_container.labels = { + "cubbi.session": "true", + "cubbi.session.id": "test-session", + "cubbi.session.name": "test-session-name", + "cubbi.image": "goose", + # No cubbi.mcps label + } + mock_container.attrs = {"NetworkSettings": {"Ports": {}}} + + # Mock Docker client + mock_client = Mock() + mock_client.containers.list.return_value = [mock_container] + + # Create container manager with mocked client + with patch("cubbi.container.docker.from_env") as mock_docker: + mock_docker.return_value = mock_client + mock_client.ping.return_value = True + + container_manager = ContainerManager() + sessions = container_manager.list_sessions() + + assert len(sessions) == 1 + session = sessions[0] + assert session.id == "test-session" + assert session.mcps == [] # Should be empty list when no MCPs @pytest.mark.requires_docker