initial commit after one day coding agent session

tests/server/tools/test_build_environment.py

@@ -0,0 +1,250 @@
+"""Tests for Build Custom Environment Tool."""
+
+import pytest
+from unittest.mock import Mock, AsyncMock
+from mcp.types import Tool, TextContent
+import json
+
+from mcp_forge.server.tools.build_environment import BuildEnvironmentTool
+from mcp_forge.builder.environment_builder import BuildResult
+
+
+@pytest.fixture
+def mock_environment_builder():
+    """Mock environment builder."""
+    builder = Mock()
+    builder.build_environment = AsyncMock(return_value=BuildResult(
+        success=True,
+        image_name="mcp-forge/custom:test-env",
+        image_id="sha256:abc123",
+        build_time=45.2,
+        size_bytes=524288000,  # 500MB
+        installed_packages=["numpy==1.24.0", "pandas==2.0.0"],
+        cache_hit=False
+    ))
+    return builder
+
+
+@pytest.fixture
+def mock_audit_logger():
+    """Mock audit logger."""
+    logger = Mock()
+    logger.log_environment_build = Mock()
+    return logger
+
+
+@pytest.fixture
+def build_environment_tool(mock_environment_builder, mock_audit_logger):
+    """Create BuildEnvironmentTool instance with mocked dependencies."""
+    return BuildEnvironmentTool(
+        environment_builder=mock_environment_builder,
+        audit_logger=mock_audit_logger
+    )
+
+
+@pytest.mark.asyncio
+async def test_get_tool_definition(build_environment_tool):
+    """Test that tool definition matches MCP spec."""
+    definition = build_environment_tool.get_tool_definition()
+    
+    assert isinstance(definition, Tool)
+    assert definition.name == "build_custom_environment"
+    assert definition.description is not None
+    assert "build" in definition.description.lower()
+    
+    # Verify required schema properties
+    schema = definition.inputSchema
+    assert schema["type"] == "object"
+    assert "name" in schema["properties"]
+    assert "packages" in schema["properties"]
+    assert "base_image" in schema["properties"]
+    assert "python_version" in schema["properties"]
+    assert "description" in schema["properties"]
+    assert "name" in schema["required"]
+    assert "packages" in schema["required"]
+
+
+@pytest.mark.asyncio
+async def test_build_with_valid_packages(
+    build_environment_tool,
+    mock_environment_builder,
+    mock_audit_logger
+):
+    """Test building environment with valid packages."""
+    arguments = {
+        "name": "test-ml-env",
+        "packages": ["numpy>=1.24.0", "pandas>=2.0.0"]
+    }
+    
+    result = await build_environment_tool.execute(arguments)
+    
+    assert isinstance(result, list)
+    assert len(result) == 1
+    assert isinstance(result[0], TextContent)
+    
+    # Parse JSON response
+    response = json.loads(result[0].text)
+    assert response["success"] is True
+    assert response["image_name"] == "mcp-forge/custom:test-env"
+    assert response["build_time"] == 45.2
+    assert len(response["installed_packages"]) == 2
+    
+    # Verify builder was called
+    mock_environment_builder.build_environment.assert_called_once()
+    
+    # Verify audit log was called
+    mock_audit_logger.log_environment_build.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_build_with_base_image(
+    build_environment_tool,
+    mock_environment_builder
+):
+    """Test building with custom base image."""
+    arguments = {
+        "name": "test-env",
+        "packages": ["requests"],
+        "base_image": "python:3.12"
+    }
+    
+    result = await build_environment_tool.execute(arguments)
+    
+    # Parse JSON response
+    response = json.loads(result[0].text)
+    assert response["success"] is True
+    
+    # Verify base_image was passed to builder
+    call_args = mock_environment_builder.build_environment.call_args
+    assert call_args.kwargs["base_image"] == "python:3.12"
+
+
+@pytest.mark.asyncio
+async def test_build_with_python_version(
+    build_environment_tool,
+    mock_environment_builder
+):
+    """Test building with specific Python version."""
+    arguments = {
+        "name": "test-env",
+        "packages": ["numpy"],
+        "python_version": "3.12"
+    }
+    
+    result = await build_environment_tool.execute(arguments)
+    
+    # Parse JSON response
+    response = json.loads(result[0].text)
+    assert response["success"] is True
+    
+    # Verify python_version was passed to builder
+    call_args = mock_environment_builder.build_environment.call_args
+    assert call_args.kwargs["python_version"] == "3.12"
+
+
+@pytest.mark.asyncio
+async def test_build_with_description(
+    build_environment_tool,
+    mock_environment_builder
+):
+    """Test building with description."""
+    arguments = {
+        "name": "test-env",
+        "packages": ["pandas"],
+        "description": "Environment for data analysis"
+    }
+    
+    result = await build_environment_tool.execute(arguments)
+    
+    # Parse JSON response
+    response = json.loads(result[0].text)
+    assert response["success"] is True
+
+
+@pytest.mark.asyncio
+async def test_validate_arguments_missing_name(build_environment_tool):
+    """Test that validation fails when name is missing."""
+    arguments = {
+        "packages": ["numpy"]
+    }
+    
+    with pytest.raises((ValueError, KeyError)):
+        await build_environment_tool.execute(arguments)
+
+
+@pytest.mark.asyncio
+async def test_validate_arguments_missing_packages(build_environment_tool):
+    """Test that validation fails when packages is missing."""
+    arguments = {
+        "name": "test-env"
+    }
+    
+    with pytest.raises((ValueError, KeyError)):
+        await build_environment_tool.execute(arguments)
+
+
+@pytest.mark.asyncio
+async def test_validate_arguments_invalid_name(build_environment_tool):
+    """Test that validation fails for invalid environment name."""
+    arguments = {
+        "name": "invalid name with spaces",
+        "packages": ["numpy"]
+    }
+    
+    with pytest.raises(ValueError, match="name"):
+        await build_environment_tool.execute(arguments)
+
+
+@pytest.mark.asyncio
+async def test_validate_arguments_invalid_packages_type(build_environment_tool):
+    """Test that validation fails when packages is not a list."""
+    arguments = {
+        "name": "test-env",
+        "packages": "not-a-list"
+    }
+    
+    with pytest.raises(ValueError, match="packages"):
+        await build_environment_tool.execute(arguments)
+
+
+@pytest.mark.asyncio
+async def test_validate_arguments_empty_packages(build_environment_tool):
+    """Test that validation fails when packages list is empty."""
+    arguments = {
+        "name": "test-env",
+        "packages": []
+    }
+    
+    with pytest.raises(ValueError, match="packages"):
+        await build_environment_tool.execute(arguments)
+
+
+@pytest.mark.asyncio
+async def test_build_error_handling(
+    build_environment_tool,
+    mock_environment_builder
+):
+    """Test that build errors are handled gracefully."""
+    mock_environment_builder.build_environment = AsyncMock(return_value=BuildResult(
+        success=False,
+        image_name="",
+        image_id="",
+        build_time=5.0,
+        size_bytes=0,
+        installed_packages=[],
+        cache_hit=False,
+        error="Package 'invalid-pkg' not found"
+    ))
+    
+    arguments = {
+        "name": "test-env",
+        "packages": ["invalid-pkg"]
+    }
+    
+    result = await build_environment_tool.execute(arguments)
+    
+    # Parse JSON response
+    response = json.loads(result[0].text)
+    assert response["success"] is False
+    assert "error" in response
+    assert "invalid-pkg" in response["error"]

tests/server/tools/test_document_state.py

@@ -0,0 +1,188 @@
+"""Tests for Document State Tool."""
+
+import pytest
+from unittest.mock import Mock, AsyncMock
+from mcp.types import Tool, TextContent
+import json
+
+from mcp_forge.server.tools.document_state import DocumentStateTool
+
+
+@pytest.fixture
+def mock_session_manager():
+    """Mock session manager."""
+    manager = Mock()
+    # Mock get_session_state to return a session with state
+    mock_session = Mock()
+    mock_session.state = Mock()
+    mock_session.state.to_dict = Mock(return_value={
+        "variables": {"x": 1, "y": 2},
+        "documented_variables": {},
+        "note": None
+    })
+    manager.get_session_state = Mock(return_value=mock_session)
+    manager.document_variables = AsyncMock(return_value={"success": True, "documented_count": 2})
+    manager.session_exists = Mock(return_value=True)
+    return manager
+
+
+@pytest.fixture
+def document_state_tool(mock_session_manager):
+    """Create DocumentStateTool instance with mocked dependencies."""
+    return DocumentStateTool(session_manager=mock_session_manager)
+
+
+@pytest.mark.asyncio
+async def test_get_tool_definition(document_state_tool):
+    """Test that tool definition matches MCP spec."""
+    definition = document_state_tool.get_tool_definition()
+    
+    assert isinstance(definition, Tool)
+    assert definition.name == "document_state"
+    assert definition.description is not None
+    assert "document" in definition.description.lower()
+    
+    # Verify required schema properties
+    schema = definition.inputSchema
+    assert schema["type"] == "object"
+    assert "session_id" in schema["properties"]
+    assert "variables" in schema["properties"]
+    assert "note" in schema["properties"]
+    assert "clear" in schema["properties"]
+    assert "session_id" in schema["required"]
+    assert "variables" in schema["required"]
+
+
+@pytest.mark.asyncio
+async def test_document_variables(
+    document_state_tool,
+    mock_session_manager
+):
+    """Test documenting variables in a session."""
+    arguments = {
+        "session_id": "test-session",
+        "variables": {
+            "df": "Customer data with 1000 rows",
+            "model": "Trained RandomForest classifier"
+        }
+    }
+    
+    result = await document_state_tool.execute(arguments)
+    
+    assert isinstance(result, list)
+    assert len(result) == 1
+    assert isinstance(result[0], TextContent)
+    
+    # Parse JSON response
+    response = json.loads(result[0].text)
+    assert response["success"] is True
+    assert response["documented_count"] == 2
+    assert response["session_id"] == "test-session"
+    
+    # Verify session manager was called
+    mock_session_manager.document_variables.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_document_with_note(
+    document_state_tool,
+    mock_session_manager
+):
+    """Test documenting with a note."""
+    arguments = {
+        "session_id": "test-session",
+        "variables": {
+            "result": "Final analysis output"
+        },
+        "note": "Analysis complete, ready for reporting"
+    }
+    
+    result = await document_state_tool.execute(arguments)
+    
+    # Parse JSON response
+    response = json.loads(result[0].text)
+    assert response["success"] is True
+    
+    # Verify note was passed
+    call_args = mock_session_manager.document_variables.call_args
+    assert call_args is not None
+    assert call_args.kwargs["note"] == "Analysis complete, ready for reporting"
+
+
+@pytest.mark.asyncio
+async def test_clear_existing_documentation(
+    document_state_tool,
+    mock_session_manager
+):
+    """Test clearing existing documentation."""
+    arguments = {
+        "session_id": "test-session",
+        "variables": {
+            "new_var": "New variable"
+        },
+        "clear": True
+    }
+    
+    result = await document_state_tool.execute(arguments)
+    
+    # Parse JSON response
+    response = json.loads(result[0].text)
+    assert response["success"] is True
+    
+    # Verify clear flag was passed
+    call_args = mock_session_manager.document_variables.call_args
+    assert call_args is not None
+    assert call_args.kwargs["clear"] is True
+
+
+@pytest.mark.asyncio
+async def test_session_validation(
+    document_state_tool,
+    mock_session_manager
+):
+    """Test that validation fails for non-existent session."""
+    mock_session_manager.session_exists = Mock(return_value=False)
+    
+    arguments = {
+        "session_id": "nonexistent",
+        "variables": {
+            "x": "test"
+        }
+    }
+    
+    with pytest.raises(ValueError, match="Session"):
+        await document_state_tool.execute(arguments)
+
+
+@pytest.mark.asyncio
+async def test_validate_arguments_missing_session_id(document_state_tool):
+    """Test that validation fails when session_id is missing."""
+    arguments = {
+        "variables": {"x": "test"}
+    }
+    
+    with pytest.raises((ValueError, KeyError)):
+        await document_state_tool.execute(arguments)
+
+
+@pytest.mark.asyncio
+async def test_validate_arguments_missing_variables(document_state_tool):
+    """Test that validation fails when variables is missing."""
+    arguments = {
+        "session_id": "test-session"
+    }
+    
+    with pytest.raises((ValueError, KeyError)):
+        await document_state_tool.execute(arguments)
+
+
+@pytest.mark.asyncio
+async def test_validate_arguments_invalid_variables_type(document_state_tool):
+    """Test that validation fails when variables is not a dict."""
+    arguments = {
+        "session_id": "test-session",
+        "variables": ["not", "a", "dict"]
+    }
+    
+    with pytest.raises(ValueError, match="variables"):
+        await document_state_tool.execute(arguments)

tests/server/tools/test_execute_python.py

@@ -0,0 +1,389 @@
+"""Tests for Execute Python Tool."""
+
+import pytest
+from unittest.mock import Mock, AsyncMock, patch
+from mcp.types import Tool, TextContent
+import json
+
+from mcp_forge.server.tools.execute_python import ExecutePythonTool
+from mcp_forge.execution.simple.backend import ExecutionResult
+
+
+@pytest.fixture
+def mock_simple_backend():
+    """Mock simple backend."""
+    backend = Mock()
+    backend.execute = AsyncMock(return_value=ExecutionResult(
+        success=True,
+        stdout="Hello World",
+        stderr="",
+        result="42",
+        execution_time=0.5,
+        exit_code=0
+    ))
+    return backend
+
+
+@pytest.fixture
+def mock_jupyter_backend():
+    """Mock jupyter backend."""
+    backend = Mock()
+    backend.execute = AsyncMock(return_value=ExecutionResult(
+        success=True,
+        stdout="Stateful execution",
+        stderr="",
+        result="session-123",
+        execution_time=1.2,
+        exit_code=0
+    ))
+    backend.execute_in_session = AsyncMock(return_value=ExecutionResult(
+        success=True,
+        stdout="Using existing session",
+        stderr="",
+        result="[1, 2, 3]",
+        execution_time=0.3,
+        exit_code=0
+    ))
+    return backend
+
+
+@pytest.fixture
+def mock_client_manager():
+    """Mock MCP client manager."""
+    manager = Mock()
+    manager.list_all_tools = AsyncMock(return_value=[
+        {"name": "github_search_repos", "description": "Search GitHub repositories"},
+        {"name": "filesystem_read", "description": "Read file contents"}
+    ])
+    return manager
+
+
+@pytest.fixture
+def mock_bridge_server():
+    """Mock tool bridge server."""
+    server = Mock()
+    server.socket_path = "/tmp/mcp-bridge.sock"
+    server.is_running = Mock(return_value=True)
+    return server
+
+
+@pytest.fixture
+def mock_injection_generator():
+    """Mock tool injection generator."""
+    generator = Mock()
+    generator.generate_injection_code = Mock(return_value="""
+# MCP Tool Injection
+def github_search_repos(**kwargs):
+    import socket
+    # ... tool implementation
+    pass
+""")
+    return generator
+
+
+@pytest.fixture
+def mock_config():
+    """Mock forge configuration."""
+    config = Mock()
+    config.execution = Mock()
+    config.execution.default_backend = "simple"
+    config.execution.default_timeout = 300
+    config.execution.max_timeout = 1800
+    config.execution.default_memory = "512m"
+    config.execution.max_memory = "2g"
+    return config
+
+
+@pytest.fixture
+def execute_python_tool(
+    mock_simple_backend,
+    mock_jupyter_backend,
+    mock_client_manager,
+    mock_bridge_server,
+    mock_injection_generator,
+    mock_config
+):
+    """Create ExecutePythonTool instance with mocked dependencies."""
+    return ExecutePythonTool(
+        simple_backend=mock_simple_backend,
+        jupyter_backend=mock_jupyter_backend,
+        client_manager=mock_client_manager,
+        bridge_server=mock_bridge_server,
+        injection_generator=mock_injection_generator,
+        config=mock_config
+    )
+
+
+@pytest.mark.asyncio
+async def test_get_tool_definition(execute_python_tool):
+    """Test that tool definition matches MCP spec."""
+    definition = execute_python_tool.get_tool_definition()
+    
+    assert isinstance(definition, Tool)
+    assert definition.name == "execute_python"
+    assert definition.description is not None
+    assert "Execute Python code" in definition.description
+    
+    # Verify required schema properties
+    schema = definition.inputSchema
+    assert schema["type"] == "object"
+    assert "code" in schema["properties"]
+    assert "mcp_tools" in schema["properties"]
+    assert "session_id" in schema["properties"]
+    assert "backend" in schema["properties"]
+    assert "timeout" in schema["properties"]
+    assert "custom_image" in schema["properties"]
+    assert "environment" in schema["properties"]
+    assert schema["required"] == ["code"]
+
+
+@pytest.mark.asyncio
+async def test_execute_simple_backend_stateless(
+    execute_python_tool,
+    mock_simple_backend
+):
+    """Test execution with simple backend (stateless)."""
+    arguments = {
+        "code": "print('Hello World'); 42"
+    }
+    
+    result = await execute_python_tool.execute(arguments)
+    
+    assert isinstance(result, list)
+    assert len(result) == 1
+    assert isinstance(result[0], TextContent)
+    
+    # Parse JSON response
+    response = json.loads(result[0].text)
+    assert response["success"] is True
+    assert response["stdout"] == "Hello World"
+    assert response["result"] == "42"
+    assert response["execution_time"] == 0.5
+    assert "session_id" not in response or response["session_id"] is None
+    
+    # Verify simple backend was called
+    mock_simple_backend.execute.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_execute_jupyter_backend_stateful(
+    execute_python_tool,
+    mock_jupyter_backend
+):
+    """Test execution with jupyter backend (stateful)."""
+    arguments = {
+        "code": "x = 42; print('Stateful')",
+        "session_id": "test-session",
+        "backend": "jupyter"
+    }
+    
+    result = await execute_python_tool.execute(arguments)
+    
+    # Parse JSON response
+    response = json.loads(result[0].text)
+    assert response["success"] is True
+    assert "session_id" in response
+    
+    # Verify jupyter backend was called
+    mock_jupyter_backend.execute_in_session.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_execute_with_mcp_tools(
+    execute_python_tool,
+    mock_simple_backend,
+    mock_injection_generator,
+    mock_bridge_server
+):
+    """Test execution with MCP tool injection."""
+    arguments = {
+        "code": "repos = github_search_repos(query='test', max_results=10)",
+        "mcp_tools": ["github_search_repos"]
+    }
+    
+    result = await execute_python_tool.execute(arguments)
+    
+    # Parse JSON response
+    response = json.loads(result[0].text)
+    assert response["success"] is True
+    assert "available_tools" in response
+    assert "github_search_repos" in response["available_tools"]
+    
+    # Verify injection generator was called
+    mock_injection_generator.generate_injection_code.assert_called_once_with(
+        tool_names=["github_search_repos"],
+        socket_path="/tmp/mcp-bridge.sock"
+    )
+
+
+@pytest.mark.asyncio
+async def test_execute_with_custom_image(
+    execute_python_tool,
+    mock_simple_backend
+):
+    """Test execution with custom image."""
+    arguments = {
+        "code": "import pandas as pd; pd.DataFrame()",
+        "custom_image": "mcp-forge/custom:my-ml-env"
+    }
+    
+    _result = await execute_python_tool.execute(arguments)
+    
+    # Verify backend was called with custom image
+    call_args = mock_simple_backend.execute.call_args
+    assert call_args is not None
+    # Check that custom_image was passed in execution options
+    assert "image" in call_args.kwargs or "custom_image" in call_args.kwargs
+
+
+@pytest.mark.asyncio
+async def test_execute_with_environment_template(
+    execute_python_tool,
+    mock_simple_backend
+):
+    """Test execution with environment template."""
+    arguments = {
+        "code": "import numpy as np; np.array([1,2,3])",
+        "environment": "datascience"
+    }
+    
+    _result = await execute_python_tool.execute(arguments)
+    
+    # Verify backend was called with environment specification
+    mock_simple_backend.execute.assert_called_once()
+    call_args = mock_simple_backend.execute.call_args
+    assert call_args is not None
+
+
+@pytest.mark.asyncio
+async def test_execute_with_timeout(
+    execute_python_tool,
+    mock_simple_backend
+):
+    """Test execution with custom timeout."""
+    arguments = {
+        "code": "import time; time.sleep(10)",
+        "timeout": 5
+    }
+    
+    _result = await execute_python_tool.execute(arguments)
+    
+    # Verify timeout was passed to backend
+    call_args = mock_simple_backend.execute.call_args
+    assert call_args is not None
+
+
+@pytest.mark.asyncio
+async def test_validate_arguments_missing_code(execute_python_tool):
+    """Test that validation fails when code is missing."""
+    arguments = {}
+    
+    with pytest.raises((ValueError, KeyError)):
+        await execute_python_tool.execute(arguments)
+
+
+@pytest.mark.asyncio
+async def test_validate_arguments_invalid_backend(execute_python_tool):
+    """Test that validation fails for invalid backend."""
+    arguments = {
+        "code": "print('test')",
+        "backend": "invalid"
+    }
+    
+    with pytest.raises(ValueError, match="backend"):
+        await execute_python_tool.execute(arguments)
+
+
+@pytest.mark.asyncio
+async def test_validate_arguments_invalid_timeout(
+    execute_python_tool,
+    mock_config
+):
+    """Test that validation fails for timeout exceeding max."""
+    arguments = {
+        "code": "print('test')",
+        "timeout": 3600  # Exceeds max_timeout of 1800
+    }
+    
+    with pytest.raises(ValueError, match="Timeout"):
+        await execute_python_tool.execute(arguments)
+
+
+@pytest.mark.asyncio
+async def test_backend_selection_with_session_id(execute_python_tool):
+    """Test that jupyter backend is selected when session_id provided."""
+    backend = execute_python_tool._select_backend(
+        session_id="test-session",
+        backend=None
+    )
+    
+    assert backend == "jupyter"
+
+
+@pytest.mark.asyncio
+async def test_backend_selection_default(
+    execute_python_tool,
+    mock_config
+):
+    """Test that default backend is used when no session_id."""
+    backend = execute_python_tool._select_backend(
+        session_id=None,
+        backend=None
+    )
+    
+    assert backend == "simple"
+
+
+@pytest.mark.asyncio
+async def test_backend_selection_explicit(execute_python_tool):
+    """Test that explicit backend is respected."""
+    backend = execute_python_tool._select_backend(
+        session_id=None,
+        backend="jupyter"
+    )
+    
+    assert backend == "jupyter"
+
+
+@pytest.mark.asyncio
+async def test_execution_error_handling(
+    execute_python_tool,
+    mock_simple_backend
+):
+    """Test that execution errors are handled gracefully."""
+    mock_simple_backend.execute = AsyncMock(return_value=ExecutionResult(
+        success=False,
+        stdout="",
+        stderr="NameError: name 'undefined_variable' is not defined",
+        result=None,
+        execution_time=0.1,
+        exit_code=1
+    ))
+    
+    arguments = {
+        "code": "print(undefined_variable)"
+    }
+    
+    result = await execute_python_tool.execute(arguments)
+    
+    # Parse JSON response
+    response = json.loads(result[0].text)
+    assert response["success"] is False
+    assert "NameError" in response["stderr"]
+
+
+@pytest.mark.asyncio
+async def test_bridge_server_not_running(
+    execute_python_tool,
+    mock_bridge_server
+):
+    """Test that error is raised if bridge server not running."""
+    mock_bridge_server.is_running = Mock(return_value=False)
+    
+    arguments = {
+        "code": "repos = github_search_repos(query='test')",
+        "mcp_tools": ["github_search_repos"]
+    }
+    
+    with pytest.raises(RuntimeError, match="bridge server"):
+        await execute_python_tool.execute(arguments)