initial commit after one day coding agent session
This commit is contained in:
commit
372af75b90
88 changed files with 22694 additions and 0 deletions
20
.gitignore
vendored
Normal file
20
.gitignore
vendored
Normal file
|
|
@ -0,0 +1,20 @@
|
||||||
|
# Python-generated files
|
||||||
|
__pycache__/
|
||||||
|
*.py[oc]
|
||||||
|
build/
|
||||||
|
dist/
|
||||||
|
wheels/
|
||||||
|
*.egg-info
|
||||||
|
|
||||||
|
# Virtual environments
|
||||||
|
.venv
|
||||||
|
|
||||||
|
# Configuration files
|
||||||
|
config.yaml
|
||||||
|
.vscode/
|
||||||
|
|
||||||
|
# Logs
|
||||||
|
logs/
|
||||||
|
|
||||||
|
# old stuff
|
||||||
|
_attic/
|
||||||
1
.python-version
Normal file
1
.python-version
Normal file
|
|
@ -0,0 +1 @@
|
||||||
|
3.13
|
||||||
106
QUICKSTART.md
Normal file
106
QUICKSTART.md
Normal file
|
|
@ -0,0 +1,106 @@
|
||||||
|
# MCP-Forge Quick Start Guide
|
||||||
|
|
||||||
|
## Installation
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Install from source
|
||||||
|
git clone <repository>
|
||||||
|
cd mcp-forge
|
||||||
|
uv sync
|
||||||
|
```
|
||||||
|
|
||||||
|
## Configuration
|
||||||
|
|
||||||
|
1. Copy the example configuration:
|
||||||
|
```bash
|
||||||
|
cp config.example.yaml config.yaml
|
||||||
|
```
|
||||||
|
|
||||||
|
2. Edit `config.yaml` to set your Podman socket path:
|
||||||
|
```yaml
|
||||||
|
server:
|
||||||
|
podman_socket: /run/user/1000/podman/podman.sock # Update this
|
||||||
|
```
|
||||||
|
|
||||||
|
3. Ensure Podman is running:
|
||||||
|
```bash
|
||||||
|
systemctl --user start podman.socket
|
||||||
|
```
|
||||||
|
|
||||||
|
## Running the Server
|
||||||
|
|
||||||
|
### stdio Transport (default)
|
||||||
|
```bash
|
||||||
|
mcp-forge --config config.yaml
|
||||||
|
```
|
||||||
|
|
||||||
|
### SSE Transport (HTTP)
|
||||||
|
```bash
|
||||||
|
mcp-forge --config config.yaml --transport sse --host 0.0.0.0 --port 8080
|
||||||
|
```
|
||||||
|
|
||||||
|
## Command-Line Options
|
||||||
|
|
||||||
|
```
|
||||||
|
--config PATH Path to configuration file (required)
|
||||||
|
--host HOST Server host address (default: localhost)
|
||||||
|
--port PORT Server port (default: 3000)
|
||||||
|
--transport TYPE Transport protocol: stdio or sse (default: stdio)
|
||||||
|
--verbose, -v Enable verbose logging
|
||||||
|
```
|
||||||
|
|
||||||
|
## MCP Client Configuration
|
||||||
|
|
||||||
|
To use with an MCP client (e.g., Claude Desktop), add to your client config:
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"mcpServers": {
|
||||||
|
"mcp-forge": {
|
||||||
|
"command": "mcp-forge",
|
||||||
|
"args": ["--config", "/path/to/config.yaml"]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
## Testing
|
||||||
|
|
||||||
|
Run the test suite:
|
||||||
|
```bash
|
||||||
|
uv run pytest
|
||||||
|
```
|
||||||
|
|
||||||
|
Run specific tests:
|
||||||
|
```bash
|
||||||
|
uv run pytest tests/server/test_server_integration.py -v
|
||||||
|
```
|
||||||
|
|
||||||
|
## Architecture
|
||||||
|
|
||||||
|
- **Simple Backend**: Stateless Python execution in fresh containers
|
||||||
|
- **Jupyter Backend**: Stateful sessions with kernel persistence
|
||||||
|
- **Security**: Resource limits, package validation, audit logging
|
||||||
|
- **MCP Integration**: Bridge to external MCP tool servers
|
||||||
|
|
||||||
|
## Troubleshooting
|
||||||
|
|
||||||
|
### Podman Socket Not Found
|
||||||
|
Ensure Podman socket is running:
|
||||||
|
```bash
|
||||||
|
systemctl --user status podman.socket
|
||||||
|
systemctl --user start podman.socket
|
||||||
|
```
|
||||||
|
|
||||||
|
### Permission Denied
|
||||||
|
Check socket permissions:
|
||||||
|
```bash
|
||||||
|
ls -l /run/user/$(id -u)/podman/podman.sock
|
||||||
|
```
|
||||||
|
|
||||||
|
### Container Images
|
||||||
|
Build or pull required images:
|
||||||
|
```bash
|
||||||
|
podman pull python:3.11-slim
|
||||||
|
podman tag python:3.11-slim mcp-forge/python:3.11
|
||||||
|
```
|
||||||
256
README.md
Normal file
256
README.md
Normal file
|
|
@ -0,0 +1,256 @@
|
||||||
|
# MCP-Forge
|
||||||
|
|
||||||
|
**Secure Python Execution Server with Model Context Protocol (MCP) Support**
|
||||||
|
|
||||||
|
MCP-Forge provides a secure, containerized Python execution environment that integrates with the Model Context Protocol. It enables AI assistants and other MCP clients to execute Python code safely with resource limits, security controls, and audit logging.
|
||||||
|
|
||||||
|
## Features
|
||||||
|
|
||||||
|
- **🔒 Secure Execution**: Podman-based container isolation with resource limits
|
||||||
|
- **🎯 MCP Protocol**: Native integration with Model Context Protocol for AI assistants
|
||||||
|
- **📊 Dual Backends**:
|
||||||
|
- Simple backend for stateless code execution
|
||||||
|
- Jupyter backend for stateful sessions with kernel persistence
|
||||||
|
- **🛡️ Security Controls**:
|
||||||
|
- Package allowlist/blocklist validation
|
||||||
|
- Resource limits (memory, CPU, timeout, storage)
|
||||||
|
- Comprehensive audit logging
|
||||||
|
- **🔧 Environment Building**: Dynamic Python environment creation with `uv`
|
||||||
|
- **🌉 MCP Bridge**: Connect to external MCP tool servers (stdio, HTTP, SSE)
|
||||||
|
- Support for stdio-based MCP servers (command-line tools)
|
||||||
|
- HTTP transport for REST API-based MCP servers
|
||||||
|
- SSE transport for Server-Sent Events MCP servers
|
||||||
|
|
||||||
|
## Quick Start
|
||||||
|
|
||||||
|
### Installation
|
||||||
|
|
||||||
|
```bash
|
||||||
|
git clone <repository>
|
||||||
|
cd mcp-forge
|
||||||
|
uv sync
|
||||||
|
```
|
||||||
|
|
||||||
|
### Configuration
|
||||||
|
|
||||||
|
1. Copy example configuration:
|
||||||
|
```bash
|
||||||
|
cp config.example.yaml config.yaml
|
||||||
|
```
|
||||||
|
|
||||||
|
2. Update Podman socket path in `config.yaml`:
|
||||||
|
```yaml
|
||||||
|
server:
|
||||||
|
podman_socket: /run/user/1000/podman/podman.sock
|
||||||
|
```
|
||||||
|
|
||||||
|
3. Start Podman socket:
|
||||||
|
```bash
|
||||||
|
systemctl --user start podman.socket
|
||||||
|
```
|
||||||
|
|
||||||
|
### Running the Server
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# stdio transport (for MCP clients)
|
||||||
|
uv run mcp-forge --config config.yaml
|
||||||
|
|
||||||
|
# HTTP transport (REST API)
|
||||||
|
uv run mcp-forge --config config.yaml --transport http --port 8011
|
||||||
|
|
||||||
|
# SSE transport (Server-Sent Events)
|
||||||
|
uv run mcp-forge --config config.yaml --transport sse --port 8080
|
||||||
|
|
||||||
|
# Alternative: using Python module
|
||||||
|
uv run python -m mcp_forge --config config.yaml --transport http --port 8011
|
||||||
|
```
|
||||||
|
|
||||||
|
### CLI Options
|
||||||
|
|
||||||
|
```
|
||||||
|
--config PATH Path to configuration file (required)
|
||||||
|
--host HOST Server host address (default: localhost)
|
||||||
|
--port PORT Server port (default: 3000)
|
||||||
|
--transport TYPE Transport protocol: stdio or sse (default: stdio)
|
||||||
|
--verbose, -v Enable verbose logging
|
||||||
|
```
|
||||||
|
|
||||||
|
## MCP Client Integration
|
||||||
|
|
||||||
|
Add to your MCP client configuration (e.g., Claude Desktop):
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"mcpServers": {
|
||||||
|
"mcp-forge": {
|
||||||
|
"command": "uv",
|
||||||
|
"args": ["run", "mcp-forge", "--config", "/absolute/path/to/config.yaml"]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
## Available MCP Tools
|
||||||
|
|
||||||
|
### `execute_python`
|
||||||
|
Execute Python code in isolated container:
|
||||||
|
```python
|
||||||
|
{
|
||||||
|
"code": "print(2 + 2)",
|
||||||
|
"timeout": 30,
|
||||||
|
"memory": "512m"
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
### `build_environment`
|
||||||
|
Create custom Python environment with packages:
|
||||||
|
```python
|
||||||
|
{
|
||||||
|
"python_version": "3.11",
|
||||||
|
"packages": ["requests", "pandas"]
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
### `document_state`
|
||||||
|
Manage Jupyter session state:
|
||||||
|
```python
|
||||||
|
{
|
||||||
|
"session_id": "user_session",
|
||||||
|
"code": "x = 42",
|
||||||
|
"clear": false
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
## Architecture
|
||||||
|
|
||||||
|
```
|
||||||
|
┌─────────────────┐
|
||||||
|
│ MCP Client │
|
||||||
|
│ (Claude/Other) │
|
||||||
|
└────────┬────────┘
|
||||||
|
│
|
||||||
|
↓
|
||||||
|
┌─────────────────┐
|
||||||
|
│ MCP-Forge │
|
||||||
|
│ Server │
|
||||||
|
├─────────────────┤
|
||||||
|
│ • Tool Handlers │
|
||||||
|
│ • MCP Bridge │
|
||||||
|
│ • Security │
|
||||||
|
│ • Audit Logger │
|
||||||
|
└────────┬────────┘
|
||||||
|
│
|
||||||
|
↓
|
||||||
|
┌─────────────────┐
|
||||||
|
│ Podman │
|
||||||
|
│ Containers │
|
||||||
|
├─────────────────┤
|
||||||
|
│ • Python 3.11 │
|
||||||
|
│ • Python 3.12 │
|
||||||
|
│ • Jupyter │
|
||||||
|
└─────────────────┘
|
||||||
|
```
|
||||||
|
|
||||||
|
## Security
|
||||||
|
|
||||||
|
- **Container Isolation**: Each execution runs in isolated Podman container
|
||||||
|
- **Resource Limits**: Memory, CPU, timeout, and storage quotas enforced
|
||||||
|
- **Package Validation**: Allowlist/blocklist for package installation
|
||||||
|
- **Audit Logging**: All operations logged with timestamps and metadata
|
||||||
|
- **Network Restrictions**: Containers run without network access by default
|
||||||
|
|
||||||
|
## Testing
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Run all tests
|
||||||
|
uv run pytest
|
||||||
|
|
||||||
|
# Run specific test suite
|
||||||
|
uv run pytest tests/server/ -v
|
||||||
|
|
||||||
|
# Integration tests (requires Podman)
|
||||||
|
uv run pytest tests/integration/ -v
|
||||||
|
```
|
||||||
|
|
||||||
|
**Test Status**: 388 tests passing, 27 integration tests pending Phase 5.4
|
||||||
|
|
||||||
|
## Configuration Reference
|
||||||
|
|
||||||
|
See [`config.example.yaml`](config.example.yaml) for complete configuration options.
|
||||||
|
|
||||||
|
Key sections:
|
||||||
|
- `server`: Host, port, Podman socket
|
||||||
|
- `security`: Audit log, resource enforcement
|
||||||
|
- `execution`: Timeouts, memory limits, backends
|
||||||
|
- `images`: Container image configuration
|
||||||
|
- `sessions`: Jupyter session management
|
||||||
|
- `environment_builder`: Package validation, caching
|
||||||
|
|
||||||
|
## Development
|
||||||
|
|
||||||
|
### Project Structure
|
||||||
|
|
||||||
|
```
|
||||||
|
mcp-forge/
|
||||||
|
├── src/mcp_forge/
|
||||||
|
│ ├── server/ # MCP server implementation
|
||||||
|
│ ├── config/ # Configuration schemas
|
||||||
|
│ ├── security/ # Security & audit
|
||||||
|
│ ├── podman/ # Container management
|
||||||
|
│ ├── execution/ # Execution backends
|
||||||
|
│ ├── builder/ # Environment building
|
||||||
|
│ └── mcp/ # MCP protocol integration
|
||||||
|
├── tests/ # Test suite
|
||||||
|
├── config/ # Default configurations
|
||||||
|
└── docs/ # Documentation
|
||||||
|
```
|
||||||
|
|
||||||
|
### Running Tests
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Unit tests
|
||||||
|
uv run pytest tests/ -v
|
||||||
|
|
||||||
|
# With coverage
|
||||||
|
uv run pytest --cov=mcp_forge --cov-report=html
|
||||||
|
|
||||||
|
# Specific module
|
||||||
|
uv run pytest tests/execution/ -v
|
||||||
|
```
|
||||||
|
|
||||||
|
## Troubleshooting
|
||||||
|
|
||||||
|
### Podman Socket Not Found
|
||||||
|
```bash
|
||||||
|
systemctl --user status podman.socket
|
||||||
|
systemctl --user start podman.socket
|
||||||
|
```
|
||||||
|
|
||||||
|
### Permission Denied
|
||||||
|
Check socket permissions:
|
||||||
|
```bash
|
||||||
|
ls -l /run/user/$(id -u)/podman/podman.sock
|
||||||
|
```
|
||||||
|
|
||||||
|
### Container Images
|
||||||
|
Pull and tag images:
|
||||||
|
```bash
|
||||||
|
podman pull python:3.11-slim
|
||||||
|
podman tag python:3.11-slim mcp-forge/python:3.11
|
||||||
|
```
|
||||||
|
|
||||||
|
## License
|
||||||
|
|
||||||
|
[Add license information]
|
||||||
|
|
||||||
|
## Contributing
|
||||||
|
|
||||||
|
[Add contribution guidelines]
|
||||||
|
|
||||||
|
## Documentation
|
||||||
|
|
||||||
|
- [Quick Start Guide](docs/QUICKSTART.md)
|
||||||
|
- [HTTP Transport Configuration](docs/HTTP_TRANSPORT.md)
|
||||||
|
- [Project Status](docs/STATUS.md)
|
||||||
|
- [Architecture Overview](docs/architecture1.md)
|
||||||
|
- [Development TODO](docs/todo.md)
|
||||||
183
config.example.yaml
Normal file
183
config.example.yaml
Normal file
|
|
@ -0,0 +1,183 @@
|
||||||
|
# MCP-Forge Configuration File
|
||||||
|
#
|
||||||
|
# Copy this file to config.yaml and customize for your environment:
|
||||||
|
# cp config.yaml.example config.yaml
|
||||||
|
#
|
||||||
|
# All paths can be absolute or relative to the config file location.
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# Server Configuration
|
||||||
|
# ============================================================================
|
||||||
|
server:
|
||||||
|
# Network binding
|
||||||
|
host: localhost
|
||||||
|
port: 3000
|
||||||
|
|
||||||
|
# Podman socket path - UPDATE THIS FOR YOUR SYSTEM
|
||||||
|
# Common locations:
|
||||||
|
# - Linux (rootless): /run/user/1000/podman/podman.sock
|
||||||
|
# - Linux (root): /var/run/podman/podman.sock
|
||||||
|
# - macOS: /var/run/docker.sock (if using Podman Desktop)
|
||||||
|
podman_socket: /run/user/1000/podman/podman.sock
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# Security Configuration
|
||||||
|
# ============================================================================
|
||||||
|
security:
|
||||||
|
# Audit log location (will be created if doesn't exist)
|
||||||
|
audit_log: ./logs/audit.log
|
||||||
|
|
||||||
|
# Enforce container resource limits
|
||||||
|
enforce_resource_limits: true
|
||||||
|
|
||||||
|
# Allow containers network access (NOT RECOMMENDED for security)
|
||||||
|
allow_network: false
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# Execution Configuration
|
||||||
|
# ============================================================================
|
||||||
|
execution:
|
||||||
|
# Default backend: "simple" (stateless) or "jupyter" (stateful sessions)
|
||||||
|
default_backend: simple
|
||||||
|
|
||||||
|
# Timeout settings (seconds)
|
||||||
|
default_timeout: 300 # 5 minutes
|
||||||
|
max_timeout: 1800 # 30 minutes
|
||||||
|
|
||||||
|
# Memory limits (use k, m, g suffixes)
|
||||||
|
default_memory: 512m
|
||||||
|
max_memory: 2g
|
||||||
|
|
||||||
|
# CPU quota (microseconds per 100ms period)
|
||||||
|
# 50000 = 50% of one core, 100000 = 100% of one core
|
||||||
|
default_cpu_quota: 50000
|
||||||
|
max_cpu_quota: 100000
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# Container Images
|
||||||
|
# ============================================================================
|
||||||
|
images:
|
||||||
|
# Python images (tag or pull from registry)
|
||||||
|
python_3_11: mcp-forge/python:3.11
|
||||||
|
python_3_12: mcp-forge/python:3.12
|
||||||
|
jupyter: mcp-forge/jupyter:latest
|
||||||
|
|
||||||
|
# Automatically pull images if not found locally
|
||||||
|
auto_pull: true
|
||||||
|
|
||||||
|
# Image pull check interval (seconds, 24 hours = 86400)
|
||||||
|
pull_interval: 86400
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# Session Configuration (for Jupyter backend)
|
||||||
|
# ============================================================================
|
||||||
|
sessions:
|
||||||
|
# Session idle timeout before auto-cleanup (seconds, 1 hour = 3600)
|
||||||
|
idle_timeout: 3600
|
||||||
|
|
||||||
|
# Maximum concurrent sessions
|
||||||
|
max_concurrent: 10
|
||||||
|
|
||||||
|
# Cleanup check interval (seconds, 5 minutes = 300)
|
||||||
|
cleanup_interval: 300
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# Volume Configuration
|
||||||
|
# ============================================================================
|
||||||
|
volumes:
|
||||||
|
# Base directory for session volumes
|
||||||
|
base_path: ./volumes
|
||||||
|
|
||||||
|
# Storage quota per session (use k, m, g suffixes)
|
||||||
|
session_quota: 1g
|
||||||
|
|
||||||
|
# Maximum allowed quota
|
||||||
|
max_session_quota: 10g
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# Environment Builder Configuration
|
||||||
|
# ============================================================================
|
||||||
|
environment_builder:
|
||||||
|
# Enable dynamic environment building with uv
|
||||||
|
enabled: true
|
||||||
|
|
||||||
|
# UV package cache directory
|
||||||
|
uv_cache_path: ./cache/uv
|
||||||
|
|
||||||
|
# Build limits
|
||||||
|
max_packages_per_build: 50
|
||||||
|
max_build_time: 600 # 10 minutes
|
||||||
|
max_image_size: 2147483648 # 2GB in bytes
|
||||||
|
max_concurrent_builds: 3
|
||||||
|
|
||||||
|
# Rate limiting for build requests
|
||||||
|
build_rate_limit:
|
||||||
|
requests: 10 # Maximum requests
|
||||||
|
period: 60 # Per time period (seconds)
|
||||||
|
|
||||||
|
# Package validation
|
||||||
|
package_validation:
|
||||||
|
# Use allowlist (if false, all packages allowed except blocklist)
|
||||||
|
use_allowlist: true
|
||||||
|
|
||||||
|
# Path to allowlist file (one package per line)
|
||||||
|
allowlist_path: ./config/allowlist.txt
|
||||||
|
|
||||||
|
# Path to blocklist file (one package per line)
|
||||||
|
blocklist_path: ./config/blocklist.txt
|
||||||
|
|
||||||
|
# Patterns requiring manual approval (regex)
|
||||||
|
require_approval_patterns: []
|
||||||
|
|
||||||
|
# Auto-cleanup configuration (future feature)
|
||||||
|
auto_cleanup: {}
|
||||||
|
|
||||||
|
# Environment templates (future feature)
|
||||||
|
templates: {}
|
||||||
|
|
||||||
|
# ============================================================================
|
||||||
|
# MCP Tools Configuration
|
||||||
|
# ============================================================================
|
||||||
|
# External MCP tool servers to connect to
|
||||||
|
#
|
||||||
|
# Supports three transport types:
|
||||||
|
#
|
||||||
|
# 1. stdio (default) - Command-based MCP servers
|
||||||
|
# Required: command
|
||||||
|
# Optional: args, env
|
||||||
|
#
|
||||||
|
# 2. http - HTTP-based MCP servers
|
||||||
|
# Required: url
|
||||||
|
# Optional: headers
|
||||||
|
#
|
||||||
|
# 3. sse - Server-Sent Events MCP servers
|
||||||
|
# Required: url
|
||||||
|
# Optional: headers
|
||||||
|
#
|
||||||
|
mcp_tools: {}
|
||||||
|
|
||||||
|
# Example configurations:
|
||||||
|
#
|
||||||
|
# # stdio transport (command-based MCP server)
|
||||||
|
# mcp_tools:
|
||||||
|
# filesystem:
|
||||||
|
# transport: stdio # default, can be omitted
|
||||||
|
# command: uvx
|
||||||
|
# args: [mcp-server-filesystem, /path/to/workspace]
|
||||||
|
# env:
|
||||||
|
# SOME_VAR: value
|
||||||
|
#
|
||||||
|
# # http transport (external HTTP MCP server)
|
||||||
|
# remote_api:
|
||||||
|
# transport: http
|
||||||
|
# url: http://localhost:8006/mcp
|
||||||
|
# headers:
|
||||||
|
# Authorization: "Bearer your-token-here"
|
||||||
|
# X-Custom-Header: "value"
|
||||||
|
#
|
||||||
|
# # sse transport (Server-Sent Events MCP server)
|
||||||
|
# sse_service:
|
||||||
|
# transport: sse
|
||||||
|
# url: http://localhost:9000/events
|
||||||
|
# headers:
|
||||||
|
# Authorization: "Bearer your-token-here"
|
||||||
20
config/allowlist.txt
Normal file
20
config/allowlist.txt
Normal file
|
|
@ -0,0 +1,20 @@
|
||||||
|
# Package Allowlist
|
||||||
|
#
|
||||||
|
# List of Python packages that are allowed to be installed
|
||||||
|
# One package per line
|
||||||
|
|
||||||
|
# Standard library packages (implicitly allowed, listed for reference)
|
||||||
|
requests
|
||||||
|
urllib3
|
||||||
|
certifi
|
||||||
|
|
||||||
|
# Data science
|
||||||
|
numpy
|
||||||
|
pandas
|
||||||
|
scipy
|
||||||
|
matplotlib
|
||||||
|
seaborn
|
||||||
|
|
||||||
|
# Utilities
|
||||||
|
python-dateutil
|
||||||
|
pytz
|
||||||
8
config/blocklist.txt
Normal file
8
config/blocklist.txt
Normal file
|
|
@ -0,0 +1,8 @@
|
||||||
|
# Package Blocklist
|
||||||
|
#
|
||||||
|
# List of Python packages that are explicitly blocked
|
||||||
|
# One package per line
|
||||||
|
|
||||||
|
# Example malicious packages (add real ones as discovered)
|
||||||
|
malicious-package
|
||||||
|
dangerous-lib
|
||||||
181
docs/HTTP_TRANSPORT.md
Normal file
181
docs/HTTP_TRANSPORT.md
Normal file
|
|
@ -0,0 +1,181 @@
|
||||||
|
# HTTP Transport Configuration Example
|
||||||
|
|
||||||
|
This example shows how to configure MCP-Forge to connect to external HTTP-based MCP servers.
|
||||||
|
|
||||||
|
## Configuration File
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
# config.yaml
|
||||||
|
server:
|
||||||
|
host: localhost
|
||||||
|
port: 3000
|
||||||
|
podman_socket: /run/user/1000/podman/podman.sock
|
||||||
|
|
||||||
|
security:
|
||||||
|
audit_log: ./logs/audit.log
|
||||||
|
enforce_resource_limits: true
|
||||||
|
allow_network: false
|
||||||
|
|
||||||
|
execution:
|
||||||
|
default_backend: simple
|
||||||
|
default_timeout: 300
|
||||||
|
max_timeout: 1800
|
||||||
|
default_memory: 512m
|
||||||
|
max_memory: 2g
|
||||||
|
base_image: docker.io/library/python:3.13-slim
|
||||||
|
|
||||||
|
environment_builder:
|
||||||
|
enabled: true
|
||||||
|
uv_cache_path: ./cache/uv
|
||||||
|
max_packages_per_build: 50
|
||||||
|
max_build_time: 600
|
||||||
|
max_image_size: 2147483648
|
||||||
|
max_concurrent_builds: 3
|
||||||
|
build_rate_limit:
|
||||||
|
requests: 10
|
||||||
|
period: 60
|
||||||
|
package_validation:
|
||||||
|
use_allowlist: true
|
||||||
|
allowlist_path: ./config/allowlist.txt
|
||||||
|
blocklist_path: ./config/blocklist.txt
|
||||||
|
require_approval_patterns: []
|
||||||
|
|
||||||
|
# Connect to external MCP servers
|
||||||
|
mcp_tools:
|
||||||
|
# HTTP-based MCP server
|
||||||
|
remote_api:
|
||||||
|
transport: http
|
||||||
|
url: http://localhost:8006/mcp
|
||||||
|
headers:
|
||||||
|
Authorization: "Bearer your-token-here"
|
||||||
|
X-Custom-Header: "value"
|
||||||
|
|
||||||
|
# SSE-based MCP server
|
||||||
|
sse_service:
|
||||||
|
transport: sse
|
||||||
|
url: http://localhost:9000/events
|
||||||
|
headers:
|
||||||
|
Authorization: "Bearer your-token-here"
|
||||||
|
|
||||||
|
# Traditional stdio-based MCP server (still supported)
|
||||||
|
filesystem:
|
||||||
|
transport: stdio # default, can be omitted
|
||||||
|
command: uvx
|
||||||
|
args: [mcp-server-filesystem, /path/to/workspace]
|
||||||
|
env:
|
||||||
|
SOME_VAR: value
|
||||||
|
```
|
||||||
|
|
||||||
|
## Transport Types
|
||||||
|
|
||||||
|
### 1. HTTP Transport
|
||||||
|
|
||||||
|
Used for HTTP-based MCP servers that communicate via HTTP requests.
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
remote_api:
|
||||||
|
transport: http
|
||||||
|
url: http://localhost:8006/mcp
|
||||||
|
headers:
|
||||||
|
Authorization: "Bearer token123"
|
||||||
|
Content-Type: "application/json"
|
||||||
|
```
|
||||||
|
|
||||||
|
**Required:**
|
||||||
|
- `url`: HTTP endpoint URL
|
||||||
|
|
||||||
|
**Optional:**
|
||||||
|
- `headers`: HTTP headers (dict)
|
||||||
|
|
||||||
|
### 2. SSE Transport
|
||||||
|
|
||||||
|
Used for Server-Sent Events (SSE) based MCP servers.
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
sse_service:
|
||||||
|
transport: sse
|
||||||
|
url: http://localhost:9000/events
|
||||||
|
headers:
|
||||||
|
Authorization: "Bearer token123"
|
||||||
|
```
|
||||||
|
|
||||||
|
**Required:**
|
||||||
|
- `url`: SSE endpoint URL
|
||||||
|
|
||||||
|
**Optional:**
|
||||||
|
- `headers`: HTTP headers (dict)
|
||||||
|
|
||||||
|
### 3. Stdio Transport (Default)
|
||||||
|
|
||||||
|
Traditional command-based MCP servers.
|
||||||
|
|
||||||
|
```yaml
|
||||||
|
filesystem:
|
||||||
|
transport: stdio # default, can be omitted
|
||||||
|
command: uvx
|
||||||
|
args: [mcp-server-filesystem, /workspace]
|
||||||
|
env:
|
||||||
|
PATH: /usr/bin
|
||||||
|
```
|
||||||
|
|
||||||
|
**Required:**
|
||||||
|
- `command`: Executable command
|
||||||
|
|
||||||
|
**Optional:**
|
||||||
|
- `args`: Command arguments (list)
|
||||||
|
- `env`: Environment variables (dict)
|
||||||
|
|
||||||
|
## Running the Server
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Start with HTTP transport configuration
|
||||||
|
uv run mcp-forge --config config.yaml
|
||||||
|
|
||||||
|
# Or with SSE transport
|
||||||
|
uv run mcp-forge --config config.yaml --transport sse --host 0.0.0.0 --port 8080
|
||||||
|
```
|
||||||
|
|
||||||
|
## Client Connection
|
||||||
|
|
||||||
|
Once the server is running, MCP clients can connect using the configured transport:
|
||||||
|
|
||||||
|
```python
|
||||||
|
from fastmcp import Client
|
||||||
|
from fastmcp.client.transports import StreamableHttpTransport
|
||||||
|
|
||||||
|
# Connect to MCP-Forge server via HTTP
|
||||||
|
transport = StreamableHttpTransport(url="http://localhost:3000/mcp")
|
||||||
|
client = Client(transport)
|
||||||
|
|
||||||
|
async with client:
|
||||||
|
# List available tools (from all configured MCP servers)
|
||||||
|
tools = await client.list_tools()
|
||||||
|
print(tools)
|
||||||
|
```
|
||||||
|
|
||||||
|
## Testing the Connection
|
||||||
|
|
||||||
|
You can test your HTTP transport configuration using curl:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# Test HTTP endpoint
|
||||||
|
curl -X POST http://localhost:8006/mcp \
|
||||||
|
-H "Authorization: Bearer your-token-here" \
|
||||||
|
-H "Content-Type: application/json" \
|
||||||
|
-d '{"jsonrpc": "2.0", "method": "tools/list", "id": 1}'
|
||||||
|
```
|
||||||
|
|
||||||
|
## Troubleshooting
|
||||||
|
|
||||||
|
### Connection Refused
|
||||||
|
- Verify the external MCP server is running
|
||||||
|
- Check the URL and port are correct
|
||||||
|
- Ensure firewall rules allow the connection
|
||||||
|
|
||||||
|
### Authentication Errors
|
||||||
|
- Verify the authorization header is correct
|
||||||
|
- Check if the external server requires specific headers
|
||||||
|
|
||||||
|
### Tool Not Found
|
||||||
|
- Ensure the external MCP server exposes the expected tools
|
||||||
|
- Check the server logs for any errors
|
||||||
1598
docs/architecture1.md
Normal file
1598
docs/architecture1.md
Normal file
File diff suppressed because it is too large
Load diff
374
docs/development-methodology.md
Normal file
374
docs/development-methodology.md
Normal file
|
|
@ -0,0 +1,374 @@
|
||||||
|
# AI-Assisted Software Development: A Methodology for Rapid, Test-Driven Implementation
|
||||||
|
|
||||||
|
## Abstract
|
||||||
|
|
||||||
|
This document describes a development methodology that combines architectural planning, test-driven development, and AI-assisted implementation to achieve rapid yet robust software construction. Using the mcp-forge project as a case study (5,325 lines of production code developed in one day), we demonstrate how structured planning and clear acceptance criteria enable effective human-AI collaboration while maintaining code quality and architectural integrity.
|
||||||
|
|
||||||
|
## Introduction
|
||||||
|
|
||||||
|
Traditional software development faces a fundamental tension: moving fast often compromises quality, while maintaining quality slows development. AI-assisted development promises to resolve this tension, but requires methodological discipline to avoid producing technically functional yet architecturally weak systems.
|
||||||
|
|
||||||
|
The methodology described here was validated through building mcp-forge, a production-grade secure Python execution server with MCP protocol integration, achieving:
|
||||||
|
- **5,325 lines of code** in a single development session
|
||||||
|
- **388 passing tests** (100% test coverage of core functionality)
|
||||||
|
- **Zero architectural rework** required post-implementation
|
||||||
|
- **Production-ready security** (container isolation, audit logging, resource limits)
|
||||||
|
|
||||||
|
## Methodology Overview
|
||||||
|
|
||||||
|
The approach consists of four sequential phases, each building upon the previous:
|
||||||
|
|
||||||
|
```
|
||||||
|
1. Architecture Design
|
||||||
|
↓
|
||||||
|
2. Test-Driven Work Planning
|
||||||
|
↓
|
||||||
|
3. Guided AI Implementation
|
||||||
|
↓
|
||||||
|
4. Integration Validation
|
||||||
|
```
|
||||||
|
|
||||||
|
### Phase 1: Architecture-First Planning
|
||||||
|
|
||||||
|
**Objective**: Establish system structure before writing any code.
|
||||||
|
|
||||||
|
**Process**:
|
||||||
|
1. Create an architecture document (`architecture.md`) containing:
|
||||||
|
- System overview and objectives
|
||||||
|
- Component breakdown with responsibilities
|
||||||
|
- Data flow diagrams
|
||||||
|
- Technology stack decisions
|
||||||
|
- Security considerations
|
||||||
|
- Integration points
|
||||||
|
|
||||||
|
2. Focus on **interfaces over implementation**:
|
||||||
|
- Define contracts between components
|
||||||
|
- Specify data structures
|
||||||
|
- Identify abstraction boundaries
|
||||||
|
|
||||||
|
3. Make **technology choices explicit**:
|
||||||
|
- State assumptions and constraints
|
||||||
|
- Document why alternatives were rejected
|
||||||
|
- Note potential technical risks
|
||||||
|
|
||||||
|
**Example from mcp-forge**:
|
||||||
|
|
||||||
|
```markdown
|
||||||
|
## Component: Execution Backend
|
||||||
|
|
||||||
|
**Responsibility**: Execute Python code in isolated containers
|
||||||
|
|
||||||
|
**Interface**:
|
||||||
|
- `execute(code: str, timeout: int, memory: str) -> ExecutionResult`
|
||||||
|
|
||||||
|
**Implementation Options**:
|
||||||
|
1. Simple: Stateless execution (chosen for MVP)
|
||||||
|
2. Jupyter: Stateful sessions (Phase 2)
|
||||||
|
|
||||||
|
**Security Requirements**:
|
||||||
|
- No network access by default
|
||||||
|
- Resource limits enforced
|
||||||
|
- Container isolation mandatory
|
||||||
|
```
|
||||||
|
|
||||||
|
**Key Principle**: The architecture document serves as a contract between human designer and AI implementer. Ambiguity here multiplies into implementation uncertainty.
|
||||||
|
|
||||||
|
### Phase 2: Test-Driven Work Planning
|
||||||
|
|
||||||
|
**Objective**: Convert architecture into actionable tasks with verifiable completion criteria.
|
||||||
|
|
||||||
|
**Process**:
|
||||||
|
1. Create a development plan (`todo.md`) structured as:
|
||||||
|
- Phases (major milestones)
|
||||||
|
- Tasks (implementable units)
|
||||||
|
- Acceptance criteria (objective success metrics)
|
||||||
|
|
||||||
|
2. **Enforce bottom-up development**:
|
||||||
|
- Start with foundational components (no dependencies)
|
||||||
|
- Build progressively toward integration
|
||||||
|
- Each layer tested before next begins
|
||||||
|
|
||||||
|
3. **Write acceptance criteria that prevent shortcuts**:
|
||||||
|
- Require specific test coverage
|
||||||
|
- Mandate error handling
|
||||||
|
- Specify edge cases
|
||||||
|
- Include performance requirements
|
||||||
|
|
||||||
|
4. **Make testing non-optional**:
|
||||||
|
- Every task includes "Tests written and passing"
|
||||||
|
- Integration tasks require integration tests
|
||||||
|
- No task complete without verification
|
||||||
|
|
||||||
|
**Example Task Structure**:
|
||||||
|
|
||||||
|
```markdown
|
||||||
|
## Phase 1: Foundation
|
||||||
|
|
||||||
|
### Task 1.1: Configuration Schema
|
||||||
|
- [ ] Create Pydantic models for all config sections
|
||||||
|
- [ ] Validate YAML parsing
|
||||||
|
- [ ] Handle missing/invalid configurations
|
||||||
|
- [ ] **Tests**: Config validation tests (>90% coverage)
|
||||||
|
- [ ] **Acceptance**: All edge cases handled, no runtime config errors
|
||||||
|
|
||||||
|
### Task 1.2: Podman Client Wrapper
|
||||||
|
- [ ] Implement container create/start/stop/remove
|
||||||
|
- [ ] Handle connection errors gracefully
|
||||||
|
- [ ] Add timeout protection
|
||||||
|
- [ ] **Tests**: Unit tests with mocked Podman API
|
||||||
|
- [ ] **Acceptance**: All Podman operations covered, error paths tested
|
||||||
|
```
|
||||||
|
|
||||||
|
**Anti-pattern Warning**: Vague criteria like "Implement X" or "Make Y work" lead to incomplete implementations. The AI will declare success prematurely without specific verification requirements.
|
||||||
|
|
||||||
|
### Phase 3: Guided AI Implementation
|
||||||
|
|
||||||
|
**Objective**: Leverage AI for rapid implementation while maintaining human control over design decisions.
|
||||||
|
|
||||||
|
**Human Role**:
|
||||||
|
- **Architect**: Make design choices when ambiguity exists
|
||||||
|
- **Reviewer**: Validate implementations against architecture
|
||||||
|
- **Course-corrector**: Intervene when AI diverges from requirements
|
||||||
|
- **Preference-setter**: Override AI's default choices when alternatives better suit your needs
|
||||||
|
|
||||||
|
**AI Role**:
|
||||||
|
- **Implementer**: Write code following specifications
|
||||||
|
- **Test-writer**: Create comprehensive test coverage
|
||||||
|
- **Problem-solver**: Debug issues and propose solutions
|
||||||
|
- **Documenter**: Generate docstrings and comments
|
||||||
|
|
||||||
|
**Collaboration Pattern**:
|
||||||
|
|
||||||
|
```
|
||||||
|
Human: [Reviews architecture] → [Creates task with acceptance criteria]
|
||||||
|
↓
|
||||||
|
AI: [Implements task] → [Writes tests] → [Runs tests]
|
||||||
|
↓
|
||||||
|
Human: [Validates approach] → [Accepts OR provides feedback]
|
||||||
|
↓
|
||||||
|
AI: [Refines if needed] → [Marks task complete]
|
||||||
|
```
|
||||||
|
|
||||||
|
**When to Intervene**:
|
||||||
|
|
||||||
|
1. **Design Disagreement**: AI chooses approach that conflicts with architecture
|
||||||
|
- *Example*: "Use threads instead of async for the bridge server"
|
||||||
|
- *Action*: Redirect to architectural decision
|
||||||
|
|
||||||
|
2. **Preference Mismatch**: Implementation works but doesn't match your style
|
||||||
|
- *Example*: "I prefer explicit error handling over exceptions here"
|
||||||
|
- *Action*: Request specific changes
|
||||||
|
|
||||||
|
3. **Incomplete Coverage**: AI claims completion but acceptance criteria not met
|
||||||
|
- *Example*: "Task complete" but edge case tests missing
|
||||||
|
- *Action*: Point to specific uncovered scenarios
|
||||||
|
|
||||||
|
4. **Over-engineering**: AI adds unnecessary complexity
|
||||||
|
- *Example*: Elaborate caching when simple lookup sufficient
|
||||||
|
- *Action*: Request simplification
|
||||||
|
|
||||||
|
**When NOT to Intervene**:
|
||||||
|
- Implementation details within architectural constraints
|
||||||
|
- Naming conventions (unless critical to domain)
|
||||||
|
- Code organization within modules
|
||||||
|
- Test structure (if coverage is adequate)
|
||||||
|
|
||||||
|
### Phase 4: Integration Validation
|
||||||
|
|
||||||
|
**Objective**: Verify components work together as designed.
|
||||||
|
|
||||||
|
**Trigger**: As soon as two or more components interact.
|
||||||
|
|
||||||
|
**Process**:
|
||||||
|
|
||||||
|
1. **Write integration tests immediately**:
|
||||||
|
```python
|
||||||
|
def test_execute_with_audit_logging():
|
||||||
|
"""Integration: Execution + Audit"""
|
||||||
|
result = backend.execute(code="print('test')")
|
||||||
|
assert result.success
|
||||||
|
|
||||||
|
# Verify audit log entry created
|
||||||
|
logs = audit_logger.get_recent()
|
||||||
|
assert any(log.event_type == "EXECUTION" for log in logs)
|
||||||
|
```
|
||||||
|
|
||||||
|
2. **Test failure paths**:
|
||||||
|
- Network errors during MCP communication
|
||||||
|
- Container crashes during execution
|
||||||
|
- Resource limit violations
|
||||||
|
|
||||||
|
3. **Validate end-to-end flows**:
|
||||||
|
- Complete user scenarios from entry to exit
|
||||||
|
- Cross-component data flow
|
||||||
|
- State consistency across boundaries
|
||||||
|
|
||||||
|
4. **Performance validation**:
|
||||||
|
- Measure actual execution times
|
||||||
|
- Verify resource cleanup
|
||||||
|
- Check memory leaks
|
||||||
|
|
||||||
|
**Critical Insight**: Unit tests verify components work in isolation; integration tests verify your architecture is correct. Both are mandatory.
|
||||||
|
|
||||||
|
## Case Study: MCP Tool Injection Feature
|
||||||
|
|
||||||
|
This feature demonstrates the methodology in practice:
|
||||||
|
|
||||||
|
### Architecture Decision
|
||||||
|
```markdown
|
||||||
|
**Requirement**: Python code in containers must call rag-mcp tools
|
||||||
|
|
||||||
|
**Design**: Unix socket bridge
|
||||||
|
- Bridge server runs on host
|
||||||
|
- Socket mounted into containers
|
||||||
|
- Python wrapper functions generated dynamically
|
||||||
|
- JSON-RPC over socket
|
||||||
|
```
|
||||||
|
|
||||||
|
### Task Breakdown
|
||||||
|
```markdown
|
||||||
|
1. MCPClientManager: Connect to external MCP servers
|
||||||
|
2. ToolBridgeServer: Unix socket server forwarding calls
|
||||||
|
3. ToolInjectionGenerator: Generate Python wrapper code
|
||||||
|
4. Integration: Volume mounting + code injection
|
||||||
|
```
|
||||||
|
|
||||||
|
### Guided Implementation
|
||||||
|
|
||||||
|
**Iteration 1**: AI used threading for bridge server
|
||||||
|
- **Human intervention**: "Use asyncio instead, server already async"
|
||||||
|
- **Result**: Full async rewrite
|
||||||
|
|
||||||
|
**Iteration 2**: Root() serialization returned empty dicts
|
||||||
|
- **AI debugging**: Discovered FastMCP wraps data in Root objects
|
||||||
|
- **Solution**: Parse from `content[0].text` instead
|
||||||
|
- **Human validation**: "Test with real rag-mcp call"
|
||||||
|
|
||||||
|
**Iteration 3**: Integration test
|
||||||
|
```python
|
||||||
|
def test_tool_injection_end_to_end():
|
||||||
|
result = execute_python(
|
||||||
|
code='''
|
||||||
|
docs = browse_documents(page_size=2)
|
||||||
|
print(f"Retrieved {len(docs)} documents")
|
||||||
|
''',
|
||||||
|
mcp_tools=['browse_documents']
|
||||||
|
)
|
||||||
|
assert "Retrieved 2 documents" in result.stdout
|
||||||
|
```
|
||||||
|
|
||||||
|
**Result**: Working in 2 seconds per call, 100 records retrieved without context pollution.
|
||||||
|
|
||||||
|
## Benefits and Limitations
|
||||||
|
|
||||||
|
### Benefits
|
||||||
|
|
||||||
|
1. **Speed**: 5,000+ LOC in one day without sacrificing quality
|
||||||
|
2. **Quality**: Test-driven approach forces correctness
|
||||||
|
3. **Architecture**: Planning phase prevents structural rework
|
||||||
|
4. **Maintainability**: Clear separation of concerns, well-documented
|
||||||
|
5. **Flexibility**: Easy to adjust during development (change architecture → update tasks → re-implement)
|
||||||
|
|
||||||
|
### Limitations
|
||||||
|
|
||||||
|
1. **Requires Domain Knowledge**: Human must understand the problem space
|
||||||
|
- *Mitigated by*: AI can explain concepts, but cannot design unfamiliar systems
|
||||||
|
|
||||||
|
2. **Architecture Skills Critical**: Poor initial design leads to rework
|
||||||
|
- *Mitigated by*: Start with high-level design, refine before implementation
|
||||||
|
|
||||||
|
3. **Acceptance Criteria Must Be Precise**: Vague criteria → incomplete implementation
|
||||||
|
- *Mitigated by*: Include specific tests and edge cases in task descriptions
|
||||||
|
|
||||||
|
4. **AI Can't Resolve Ambiguity**: Will make assumptions that may not align with intent
|
||||||
|
- *Mitigated by*: Review implementations actively, provide feedback early
|
||||||
|
|
||||||
|
### Comparison with Traditional Development
|
||||||
|
|
||||||
|
| Aspect | Traditional | This Methodology | Speedup |
|
||||||
|
|--------|------------|------------------|---------|
|
||||||
|
| Planning | 1-2 days | 2-3 hours | 3-4x |
|
||||||
|
| Implementation | 2-3 weeks | 1 day | 10-15x |
|
||||||
|
| Testing | 2-3 days | Concurrent | 2-3x |
|
||||||
|
| Documentation | 1-2 days | Concurrent | ∞ |
|
||||||
|
| **Total** | **3-4 weeks** | **1-2 days** | **15-20x** |
|
||||||
|
|
||||||
|
*Note: Assumes experienced developer familiar with domain*
|
||||||
|
|
||||||
|
## Best Practices
|
||||||
|
|
||||||
|
### DO:
|
||||||
|
|
||||||
|
1. **Invest in architecture**: 2-3 hours planning saves days of rework
|
||||||
|
2. **Write specific acceptance criteria**: Include edge cases, error conditions
|
||||||
|
3. **Test continuously**: Don't accumulate untested code
|
||||||
|
4. **Intervene early**: Small course corrections prevent large detours
|
||||||
|
5. **Document decisions**: Capture "why" not just "what"
|
||||||
|
|
||||||
|
### DON'T:
|
||||||
|
|
||||||
|
1. **Skip planning**: "Start coding and figure it out" fails with AI
|
||||||
|
2. **Accept vague completions**: AI will claim success too early
|
||||||
|
3. **Batch testing**: Test each component before building next
|
||||||
|
4. **Over-specify implementation**: Allow AI freedom within constraints
|
||||||
|
5. **Ignore integration testing**: Unit tests alone miss architectural issues
|
||||||
|
|
||||||
|
## Applicability
|
||||||
|
|
||||||
|
This methodology works best for:
|
||||||
|
|
||||||
|
### Ideal Projects:
|
||||||
|
- ✅ Well-defined requirements
|
||||||
|
- ✅ Known technology stack
|
||||||
|
- ✅ Experienced human architect
|
||||||
|
- ✅ Clear success criteria
|
||||||
|
- ✅ 1,000-10,000 LOC scale
|
||||||
|
|
||||||
|
### Less Suitable For:
|
||||||
|
- ❌ Research/exploratory projects
|
||||||
|
- ❌ Novel algorithm development
|
||||||
|
- ❌ UI/UX heavy applications
|
||||||
|
- ❌ Undefined requirements
|
||||||
|
- ❌ Extreme performance optimization
|
||||||
|
|
||||||
|
## Future Directions
|
||||||
|
|
||||||
|
Potential methodology enhancements:
|
||||||
|
|
||||||
|
1. **Formal Verification**: Use AI to generate formal specifications from architecture
|
||||||
|
2. **Automated Architecture Validation**: Check implementations against architectural constraints
|
||||||
|
3. **Progressive Refinement**: Start with high-level design, AI proposes detailed architecture
|
||||||
|
4. **Multi-Agent Collaboration**: Separate AI agents for architecture, implementation, testing
|
||||||
|
|
||||||
|
## Conclusion
|
||||||
|
|
||||||
|
AI-assisted development can achieve 10-20x speedups over traditional development when combined with:
|
||||||
|
1. Upfront architectural planning
|
||||||
|
2. Test-driven task breakdown
|
||||||
|
3. Human oversight on design decisions
|
||||||
|
4. Continuous integration validation
|
||||||
|
|
||||||
|
The key insight is that **AI excels at implementation but requires human guidance on architecture**. By clearly separating these concerns and establishing objective acceptance criteria, we can leverage AI's speed while maintaining human control over system design.
|
||||||
|
|
||||||
|
The mcp-forge case study demonstrates this is not theoretical: 5,325 lines of production-quality code, fully tested, with zero architectural rework, in a single day of development. This represents a paradigm shift in how software can be built.
|
||||||
|
|
||||||
|
## References
|
||||||
|
|
||||||
|
**Project**: mcp-forge - Secure Python Execution Server with MCP Support
|
||||||
|
- **Repository**: [URL]
|
||||||
|
- **LOC**: 5,325 (source), 388 tests
|
||||||
|
- **Development Time**: 1 day (single human + AI)
|
||||||
|
- **Technology**: Python, FastMCP, Podman, asyncio
|
||||||
|
- **Complexity**: Multi-component system with security, async I/O, container orchestration
|
||||||
|
|
||||||
|
**Artifacts**:
|
||||||
|
- `docs/architecture1.md`: Initial architecture document
|
||||||
|
- `docs/todo.md`: Test-driven task breakdown
|
||||||
|
- `src/`: Implementation following architecture
|
||||||
|
- `tests/`: 388 tests with >90% coverage
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
*Document Version*: 1.0
|
||||||
|
*Date*: February 6, 2026
|
||||||
|
*Author*: Based on mcp-forge development experience
|
||||||
3156
docs/todo.md
Normal file
3156
docs/todo.md
Normal file
File diff suppressed because it is too large
Load diff
29
pyproject.toml
Normal file
29
pyproject.toml
Normal file
|
|
@ -0,0 +1,29 @@
|
||||||
|
[project]
|
||||||
|
name = "mcp-forge"
|
||||||
|
version = "0.1.0"
|
||||||
|
description = "Add your description here"
|
||||||
|
readme = "README.md"
|
||||||
|
authors = [
|
||||||
|
{ name = "Hans Aschauer", email = "hans.git@ch23.de" }
|
||||||
|
]
|
||||||
|
requires-python = ">=3.13"
|
||||||
|
dependencies = [
|
||||||
|
"fastmcp>=2.14.5",
|
||||||
|
"podman>=5.7.0",
|
||||||
|
"pydantic>=2.12.5",
|
||||||
|
"pyyaml>=6.0.3",
|
||||||
|
]
|
||||||
|
|
||||||
|
[project.scripts]
|
||||||
|
mcp-forge = "mcp_forge.__main__:main"
|
||||||
|
|
||||||
|
[build-system]
|
||||||
|
requires = ["uv_build>=0.9.26,<0.10.0"]
|
||||||
|
build-backend = "uv_build"
|
||||||
|
|
||||||
|
[dependency-groups]
|
||||||
|
dev = [
|
||||||
|
"pytest>=9.0.2",
|
||||||
|
"pytest-asyncio>=1.3.0",
|
||||||
|
"pytest-mock>=3.15.1",
|
||||||
|
]
|
||||||
12
src/mcp_forge/__init__.py
Normal file
12
src/mcp_forge/__init__.py
Normal file
|
|
@ -0,0 +1,12 @@
|
||||||
|
"""
|
||||||
|
MCP-Forge: Secure Python Execution Environment with MCP Protocol Support.
|
||||||
|
|
||||||
|
A containerized Python execution server that integrates with the Model Context Protocol,
|
||||||
|
providing secure, isolated code execution with resource limits and audit logging.
|
||||||
|
"""
|
||||||
|
|
||||||
|
__version__ = "0.1.0"
|
||||||
|
|
||||||
|
from .__main__ import main
|
||||||
|
|
||||||
|
__all__ = ["main", "__version__"]
|
||||||
127
src/mcp_forge/__main__.py
Normal file
127
src/mcp_forge/__main__.py
Normal file
|
|
@ -0,0 +1,127 @@
|
||||||
|
"""
|
||||||
|
MCP-Forge CLI Entry Point.
|
||||||
|
|
||||||
|
Provides command-line interface for starting the MCP-Forge server.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import sys
|
||||||
|
import logging
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from .config.loader import load_config
|
||||||
|
from .server.server import ForgeServer
|
||||||
|
|
||||||
|
|
||||||
|
def setup_logging(verbose: bool = False) -> None:
|
||||||
|
"""Configure logging for the server."""
|
||||||
|
level = logging.DEBUG if verbose else logging.INFO
|
||||||
|
logging.basicConfig(
|
||||||
|
level=level,
|
||||||
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
||||||
|
datefmt='%Y-%m-%d %H:%M:%S'
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def parse_args() -> argparse.Namespace:
|
||||||
|
"""Parse command-line arguments."""
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
prog='mcp-forge',
|
||||||
|
description='MCP-Forge: Secure Python execution server with MCP protocol support'
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument(
|
||||||
|
'--config',
|
||||||
|
type=Path,
|
||||||
|
help='Path to configuration file (YAML)',
|
||||||
|
default=None
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument(
|
||||||
|
'--host',
|
||||||
|
type=str,
|
||||||
|
help='Server host address (default: localhost)',
|
||||||
|
default='localhost'
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument(
|
||||||
|
'--port',
|
||||||
|
type=int,
|
||||||
|
help='Server port (default: 3000)',
|
||||||
|
default=3000
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument(
|
||||||
|
'--transport',
|
||||||
|
type=str,
|
||||||
|
choices=['stdio', 'sse', 'http'],
|
||||||
|
help='Transport protocol (default: stdio)',
|
||||||
|
default='stdio'
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument(
|
||||||
|
'--verbose',
|
||||||
|
'-v',
|
||||||
|
action='store_true',
|
||||||
|
help='Enable verbose logging'
|
||||||
|
)
|
||||||
|
|
||||||
|
return parser.parse_args()
|
||||||
|
|
||||||
|
def run_server(args: argparse.Namespace) -> None:
|
||||||
|
"""Run the server with the specified configuration."""
|
||||||
|
# Load configuration
|
||||||
|
if args.config:
|
||||||
|
logging.info(f"Loading configuration from {args.config}")
|
||||||
|
config = load_config(args.config)
|
||||||
|
else:
|
||||||
|
logging.error("No configuration file specified. Use --config <path>")
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
# Override config with CLI arguments if provided
|
||||||
|
if args.host != 'localhost':
|
||||||
|
config.server.host = args.host
|
||||||
|
if args.port != 3000:
|
||||||
|
config.server.port = args.port
|
||||||
|
|
||||||
|
# Initialize server
|
||||||
|
logging.info("Initializing MCP-Forge server...")
|
||||||
|
forge_server = ForgeServer(config)
|
||||||
|
|
||||||
|
# Run with appropriate transport (blocking call)
|
||||||
|
try:
|
||||||
|
if args.transport == 'stdio':
|
||||||
|
logging.info("Starting server with stdio transport")
|
||||||
|
forge_server.run(transport="stdio")
|
||||||
|
elif args.transport == 'http':
|
||||||
|
logging.info(f"Starting server with HTTP transport on {config.server.host}:{config.server.port}")
|
||||||
|
forge_server.run(transport="http", host=config.server.host, port=config.server.port)
|
||||||
|
elif args.transport == 'sse':
|
||||||
|
logging.info(f"Starting server with SSE transport on {config.server.host}:{config.server.port}")
|
||||||
|
forge_server.run(transport="sse", host=config.server.host, port=config.server.port)
|
||||||
|
except KeyboardInterrupt:
|
||||||
|
logging.info("Received shutdown signal")
|
||||||
|
except Exception as e:
|
||||||
|
logging.error(f"Server error: {e}", exc_info=True)
|
||||||
|
sys.exit(1)
|
||||||
|
finally:
|
||||||
|
logging.info("Server shutdown complete")
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
"""Main entry point for the CLI."""
|
||||||
|
args = parse_args()
|
||||||
|
setup_logging(args.verbose)
|
||||||
|
|
||||||
|
try:
|
||||||
|
run_server(args)
|
||||||
|
except KeyboardInterrupt:
|
||||||
|
pass # Clean exit on Ctrl+C
|
||||||
|
except Exception as e:
|
||||||
|
logging.error(f"Fatal error: {e}", exc_info=True)
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
main()
|
||||||
|
|
||||||
17
src/mcp_forge/builder/__init__.py
Normal file
17
src/mcp_forge/builder/__init__.py
Normal file
|
|
@ -0,0 +1,17 @@
|
||||||
|
"""Custom environment builder components."""
|
||||||
|
|
||||||
|
from .package_validator import PackageValidator, SecurityError, ApprovalRequiredError
|
||||||
|
from .uv_installer import UVInstaller
|
||||||
|
from .image_builder import ImageBuilder, BuildResult
|
||||||
|
from .environment_builder import EnvironmentBuilder, BuildRateLimiter
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"PackageValidator",
|
||||||
|
"SecurityError",
|
||||||
|
"ApprovalRequiredError",
|
||||||
|
"UVInstaller",
|
||||||
|
"ImageBuilder",
|
||||||
|
"BuildResult",
|
||||||
|
"EnvironmentBuilder",
|
||||||
|
"BuildRateLimiter",
|
||||||
|
]
|
||||||
315
src/mcp_forge/builder/environment_builder.py
Normal file
315
src/mcp_forge/builder/environment_builder.py
Normal file
|
|
@ -0,0 +1,315 @@
|
||||||
|
"""Environment builder orchestration with security validation."""
|
||||||
|
|
||||||
|
import re
|
||||||
|
import shutil
|
||||||
|
import threading
|
||||||
|
from typing import List, Optional, Dict, Set
|
||||||
|
from pathlib import Path
|
||||||
|
from datetime import datetime, timedelta
|
||||||
|
|
||||||
|
from mcp_forge.config.schema import EnvironmentBuilderConfig
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
from mcp_forge.builder.package_validator import PackageValidator
|
||||||
|
from mcp_forge.builder.uv_installer import UVInstaller
|
||||||
|
from mcp_forge.builder.image_builder import ImageBuilder, BuildResult
|
||||||
|
|
||||||
|
|
||||||
|
class BuildRateLimiter:
|
||||||
|
"""Rate limiter for build requests."""
|
||||||
|
|
||||||
|
def __init__(self, max_requests: int, period_seconds: int):
|
||||||
|
"""
|
||||||
|
Initialize rate limiter.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
max_requests: Maximum requests per period
|
||||||
|
period_seconds: Period length in seconds
|
||||||
|
"""
|
||||||
|
self.max_requests = max_requests
|
||||||
|
self.period_seconds = period_seconds
|
||||||
|
self.requests: Dict[str, List[datetime]] = {}
|
||||||
|
self.lock = threading.Lock()
|
||||||
|
|
||||||
|
def check_rate_limit(self, user_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Check if user is within rate limit.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
user_id: User identifier
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: If rate limit exceeded
|
||||||
|
"""
|
||||||
|
with self.lock:
|
||||||
|
now = datetime.now()
|
||||||
|
|
||||||
|
# Initialize user's request list if needed
|
||||||
|
if user_id not in self.requests:
|
||||||
|
self.requests[user_id] = []
|
||||||
|
|
||||||
|
# Clean up old requests
|
||||||
|
self._cleanup_old_requests(user_id, now)
|
||||||
|
|
||||||
|
# Check if at limit
|
||||||
|
if len(self.requests[user_id]) >= self.max_requests:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"Rate limit exceeded: {self.max_requests} requests "
|
||||||
|
f"per {self.period_seconds} seconds"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Record this request
|
||||||
|
self.requests[user_id].append(now)
|
||||||
|
|
||||||
|
def _cleanup_old_requests(self, user_id: str, now: datetime) -> None:
|
||||||
|
"""
|
||||||
|
Remove requests older than period.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
user_id: User identifier
|
||||||
|
now: Current time
|
||||||
|
"""
|
||||||
|
cutoff = now - timedelta(seconds=self.period_seconds)
|
||||||
|
self.requests[user_id] = [
|
||||||
|
req_time for req_time in self.requests[user_id]
|
||||||
|
if req_time > cutoff
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
class EnvironmentBuilder:
|
||||||
|
"""Builds custom Python environments with security validation."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
config: EnvironmentBuilderConfig,
|
||||||
|
podman_client: PodmanClient,
|
||||||
|
audit_logger: AuditLogger
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize environment builder.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config: Environment builder configuration
|
||||||
|
podman_client: Podman API client
|
||||||
|
audit_logger: Audit logging instance
|
||||||
|
"""
|
||||||
|
self.config = config
|
||||||
|
self.podman = podman_client
|
||||||
|
self.audit_logger = audit_logger
|
||||||
|
|
||||||
|
# Initialize sub-components
|
||||||
|
self.package_validator = PackageValidator(
|
||||||
|
config.package_validation
|
||||||
|
)
|
||||||
|
self.uv_installer = UVInstaller(config.uv_cache_path)
|
||||||
|
self.image_builder = ImageBuilder(
|
||||||
|
podman_client, config, audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
# Rate limiting and concurrency control
|
||||||
|
self.rate_limiter = BuildRateLimiter(
|
||||||
|
max_requests=config.build_rate_limit['requests'],
|
||||||
|
period_seconds=config.build_rate_limit['period']
|
||||||
|
)
|
||||||
|
self.active_builds: Set[str] = set()
|
||||||
|
self.active_builds_lock = threading.Lock()
|
||||||
|
|
||||||
|
def build_custom_environment(
|
||||||
|
self,
|
||||||
|
name: str,
|
||||||
|
packages: List[str],
|
||||||
|
base_image: str = "python:3.11-slim",
|
||||||
|
python_version: str = "3.11",
|
||||||
|
description: str = "",
|
||||||
|
user_id: str = "default"
|
||||||
|
) -> BuildResult:
|
||||||
|
"""
|
||||||
|
Build custom environment with packages.
|
||||||
|
|
||||||
|
Process:
|
||||||
|
1. Check rate limit
|
||||||
|
2. Check concurrent builds limit
|
||||||
|
3. Validate environment name
|
||||||
|
4. Validate package count
|
||||||
|
5. Validate package names (allowlist/blocklist)
|
||||||
|
6. Generate build context with UV
|
||||||
|
7. Build image
|
||||||
|
8. Cleanup build context
|
||||||
|
|
||||||
|
Args:
|
||||||
|
name: Environment name (alphanumeric + hyphens)
|
||||||
|
packages: List of package specifications
|
||||||
|
base_image: Base image to build from
|
||||||
|
python_version: Python version
|
||||||
|
description: Optional description
|
||||||
|
user_id: User ID for rate limiting
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
BuildResult
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If validation fails
|
||||||
|
SecurityError: If security check fails
|
||||||
|
RuntimeError: If rate limit or concurrency exceeded
|
||||||
|
"""
|
||||||
|
# Check rate limit
|
||||||
|
self.rate_limiter.check_rate_limit(user_id)
|
||||||
|
|
||||||
|
# Check concurrent builds
|
||||||
|
self._check_concurrent_builds()
|
||||||
|
|
||||||
|
# Validate environment name
|
||||||
|
if not re.match(r'^[a-zA-Z0-9-]+$', name):
|
||||||
|
raise ValueError(
|
||||||
|
f"Environment name '{name}' must contain only alphanumeric characters and hyphens"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Validate package count
|
||||||
|
if len(packages) > self.config.max_packages:
|
||||||
|
raise ValueError(
|
||||||
|
f"Package count {len(packages)} exceeds maximum {self.config.max_packages}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Log build start
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.BUILD_REQUEST,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message=f"Building environment: {name}",
|
||||||
|
details={
|
||||||
|
"packages": packages,
|
||||||
|
"base_image": base_image,
|
||||||
|
"user_id": user_id
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
build_context = None
|
||||||
|
try:
|
||||||
|
# Register build as active
|
||||||
|
self._register_build_start(name)
|
||||||
|
|
||||||
|
# Validate packages
|
||||||
|
self.package_validator.validate_packages(packages)
|
||||||
|
|
||||||
|
# Generate build context
|
||||||
|
build_context = self.uv_installer.create_build_context(
|
||||||
|
packages=packages,
|
||||||
|
base_image=base_image,
|
||||||
|
python_version=python_version
|
||||||
|
)
|
||||||
|
|
||||||
|
# Build image
|
||||||
|
result = self.image_builder.build_image(
|
||||||
|
name=name,
|
||||||
|
build_context=build_context,
|
||||||
|
base_image=base_image,
|
||||||
|
packages=packages
|
||||||
|
)
|
||||||
|
|
||||||
|
# Log completion
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.BUILD_COMPLETE,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message=f"Environment built: {name}",
|
||||||
|
details={
|
||||||
|
"image_id": result.image_id,
|
||||||
|
"build_time": result.build_time,
|
||||||
|
"user_id": user_id
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
|
finally:
|
||||||
|
# Always cleanup
|
||||||
|
self._register_build_complete(name)
|
||||||
|
if build_context and build_context.exists():
|
||||||
|
shutil.rmtree(build_context, ignore_errors=True)
|
||||||
|
|
||||||
|
def build_from_template(
|
||||||
|
self,
|
||||||
|
template_name: str,
|
||||||
|
additional_packages: Optional[List[str]] = None,
|
||||||
|
name: Optional[str] = None,
|
||||||
|
user_id: str = "default"
|
||||||
|
) -> BuildResult:
|
||||||
|
"""
|
||||||
|
Build environment from template.
|
||||||
|
|
||||||
|
Expands template packages and adds additional packages.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
template_name: Name of template to use
|
||||||
|
additional_packages: Optional additional packages
|
||||||
|
name: Optional custom name (uses template name if not provided)
|
||||||
|
user_id: User ID for rate limiting
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
BuildResult
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If template not found
|
||||||
|
"""
|
||||||
|
# Validate template exists
|
||||||
|
if template_name not in self.config.templates:
|
||||||
|
raise ValueError(f"Template '{template_name}' not found")
|
||||||
|
|
||||||
|
template = self.config.templates[template_name]
|
||||||
|
|
||||||
|
# Combine template and additional packages
|
||||||
|
packages = template["packages"].copy()
|
||||||
|
if additional_packages:
|
||||||
|
packages.extend(additional_packages)
|
||||||
|
|
||||||
|
# Use template name if no custom name provided
|
||||||
|
if name is None:
|
||||||
|
name = template_name
|
||||||
|
|
||||||
|
# Build with combined package list
|
||||||
|
return self.build_custom_environment(
|
||||||
|
name=name,
|
||||||
|
packages=packages,
|
||||||
|
description=template.get("description", ""),
|
||||||
|
user_id=user_id
|
||||||
|
)
|
||||||
|
|
||||||
|
def list_templates(self) -> Dict[str, dict]:
|
||||||
|
"""
|
||||||
|
List available templates.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dictionary of template name to template metadata
|
||||||
|
"""
|
||||||
|
return self.config.templates
|
||||||
|
|
||||||
|
def _check_concurrent_builds(self) -> None:
|
||||||
|
"""
|
||||||
|
Check concurrent builds limit.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: If at max concurrent builds
|
||||||
|
"""
|
||||||
|
with self.active_builds_lock:
|
||||||
|
if len(self.active_builds) >= self.config.max_concurrent_builds:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"Maximum concurrent builds ({self.config.max_concurrent_builds}) reached"
|
||||||
|
)
|
||||||
|
|
||||||
|
def _register_build_start(self, name: str) -> None:
|
||||||
|
"""
|
||||||
|
Register build as started.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
name: Environment name
|
||||||
|
"""
|
||||||
|
with self.active_builds_lock:
|
||||||
|
self.active_builds.add(name)
|
||||||
|
|
||||||
|
def _register_build_complete(self, name: str) -> None:
|
||||||
|
"""
|
||||||
|
Register build as completed.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
name: Environment name
|
||||||
|
"""
|
||||||
|
with self.active_builds_lock:
|
||||||
|
self.active_builds.discard(name)
|
||||||
315
src/mcp_forge/builder/image_builder.py
Normal file
315
src/mcp_forge/builder/image_builder.py
Normal file
|
|
@ -0,0 +1,315 @@
|
||||||
|
"""Container image builder with security validation."""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import time
|
||||||
|
import hashlib
|
||||||
|
import re
|
||||||
|
from typing import List, Optional
|
||||||
|
from pathlib import Path
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.config.schema import EnvironmentBuilderConfig
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
from mcp_forge.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class BuildResult:
|
||||||
|
"""Result of image build."""
|
||||||
|
success: bool
|
||||||
|
image_name: str
|
||||||
|
image_id: str
|
||||||
|
build_time: float
|
||||||
|
size_bytes: int
|
||||||
|
cache_hit: bool
|
||||||
|
installed_packages: List[str]
|
||||||
|
error: Optional[str] = None
|
||||||
|
|
||||||
|
def to_dict(self) -> dict:
|
||||||
|
"""Convert to dictionary for JSON serialization."""
|
||||||
|
return {
|
||||||
|
"success": self.success,
|
||||||
|
"image_name": self.image_name,
|
||||||
|
"image_id": self.image_id,
|
||||||
|
"build_time": self.build_time,
|
||||||
|
"size_bytes": self.size_bytes,
|
||||||
|
"cache_hit": self.cache_hit,
|
||||||
|
"installed_packages": self.installed_packages,
|
||||||
|
"error": self.error
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class ImageBuilder:
|
||||||
|
"""Builds container images with security validation."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
podman_client: PodmanClient,
|
||||||
|
config: EnvironmentBuilderConfig,
|
||||||
|
audit_logger: AuditLogger
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize image builder.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
podman_client: Podman API client
|
||||||
|
config: Environment builder configuration
|
||||||
|
audit_logger: Audit logging instance
|
||||||
|
"""
|
||||||
|
self.podman = podman_client
|
||||||
|
self.config = config
|
||||||
|
self.audit_logger = audit_logger
|
||||||
|
|
||||||
|
def build_image(
|
||||||
|
self,
|
||||||
|
name: str,
|
||||||
|
build_context: Path,
|
||||||
|
base_image: str,
|
||||||
|
packages: List[str],
|
||||||
|
timeout: Optional[int] = None
|
||||||
|
) -> BuildResult:
|
||||||
|
"""
|
||||||
|
Build container image from build context.
|
||||||
|
|
||||||
|
Process:
|
||||||
|
1. Validate build context
|
||||||
|
2. Generate image tag
|
||||||
|
3. Build image with Podman
|
||||||
|
4. Validate image size
|
||||||
|
5. Extract installed packages
|
||||||
|
6. Cleanup build artifacts
|
||||||
|
|
||||||
|
Args:
|
||||||
|
name: Environment name (user-provided)
|
||||||
|
build_context: Path to build context directory
|
||||||
|
base_image: Base image to build from
|
||||||
|
packages: List of packages being installed
|
||||||
|
timeout: Build timeout (uses config default if None)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
BuildResult
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If timeout exceeds max or validation fails
|
||||||
|
RuntimeError: If build fails
|
||||||
|
"""
|
||||||
|
# Validate build context
|
||||||
|
self._validate_build_context(build_context)
|
||||||
|
|
||||||
|
# Use default timeout if not specified
|
||||||
|
timeout = timeout if timeout is not None else self.config.build_timeout
|
||||||
|
|
||||||
|
# Validate timeout against maximum
|
||||||
|
if timeout > self.config.max_build_timeout:
|
||||||
|
raise ValueError(
|
||||||
|
f"Timeout {timeout} exceeds maximum {self.config.max_build_timeout}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Generate image tag
|
||||||
|
tag = self.generate_tag(name)
|
||||||
|
|
||||||
|
# Calculate cache hash
|
||||||
|
cache_hash = self.calculate_cache_hash(packages)
|
||||||
|
|
||||||
|
# Log build start
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.IMAGE_BUILD_START,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message=f"Building image: {tag}",
|
||||||
|
details={
|
||||||
|
"base_image": base_image,
|
||||||
|
"packages": packages,
|
||||||
|
"cache_hash": cache_hash
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
start_time = time.time()
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Build image with Podman
|
||||||
|
image, build_logs = self.podman.images.build(
|
||||||
|
path=str(build_context),
|
||||||
|
tag=tag,
|
||||||
|
timeout=timeout,
|
||||||
|
rm=True, # Remove intermediate containers
|
||||||
|
pull=False # Don't pull base image (assume it exists)
|
||||||
|
)
|
||||||
|
|
||||||
|
build_time = time.time() - start_time
|
||||||
|
|
||||||
|
# Validate image size
|
||||||
|
size_bytes = self.validate_image_size(image.id)
|
||||||
|
|
||||||
|
# Extract installed packages
|
||||||
|
installed_packages = self.extract_installed_packages(image.id)
|
||||||
|
|
||||||
|
# Log success
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.IMAGE_BUILD_SUCCESS,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message=f"Image built successfully: {tag}",
|
||||||
|
details={
|
||||||
|
"image_id": image.id,
|
||||||
|
"build_time": build_time,
|
||||||
|
"size_bytes": size_bytes,
|
||||||
|
"installed_packages": len(installed_packages)
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
return BuildResult(
|
||||||
|
success=True,
|
||||||
|
image_name=tag,
|
||||||
|
image_id=image.id,
|
||||||
|
build_time=build_time,
|
||||||
|
size_bytes=size_bytes,
|
||||||
|
cache_hit=False, # TODO: implement cache checking
|
||||||
|
installed_packages=installed_packages
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
build_time = time.time() - start_time
|
||||||
|
|
||||||
|
# Log failure
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.IMAGE_BUILD_FAILURE,
|
||||||
|
severity=AuditSeverity.ERROR,
|
||||||
|
message=f"Image build failed: {tag}",
|
||||||
|
error=str(e),
|
||||||
|
details={"build_time": build_time}
|
||||||
|
)
|
||||||
|
|
||||||
|
raise RuntimeError(f"Build failed: {e}") from e
|
||||||
|
|
||||||
|
def generate_tag(self, name: str) -> str:
|
||||||
|
"""
|
||||||
|
Generate image tag.
|
||||||
|
|
||||||
|
Format: mcp-forge/custom:{name}
|
||||||
|
Validates name is alphanumeric + hyphens only.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
name: Environment name
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Full image tag
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If name contains invalid characters
|
||||||
|
"""
|
||||||
|
# Validate name (alphanumeric + hyphens only)
|
||||||
|
if not re.match(r'^[a-zA-Z0-9-]+$', name):
|
||||||
|
raise ValueError(
|
||||||
|
f"Environment name '{name}' must contain only alphanumeric characters and hyphens"
|
||||||
|
)
|
||||||
|
|
||||||
|
return f"mcp-forge/custom:{name}"
|
||||||
|
|
||||||
|
def validate_image_size(self, image_id: str) -> int:
|
||||||
|
"""
|
||||||
|
Validate image size against maximum.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
image_id: Image ID to validate
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Size in bytes
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If image exceeds max size
|
||||||
|
"""
|
||||||
|
image = self.podman.images.get(image_id)
|
||||||
|
size_bytes = image.attrs.get("Size", 0)
|
||||||
|
|
||||||
|
max_size_bytes = parse_memory_string(self.config.max_image_size)
|
||||||
|
|
||||||
|
if size_bytes > max_size_bytes:
|
||||||
|
raise ValueError(
|
||||||
|
f"Image size {size_bytes} bytes exceeds maximum {max_size_bytes} bytes"
|
||||||
|
)
|
||||||
|
|
||||||
|
return size_bytes
|
||||||
|
|
||||||
|
def extract_installed_packages(self, image_id: str) -> List[str]:
|
||||||
|
"""
|
||||||
|
Extract list of installed packages from image.
|
||||||
|
|
||||||
|
Runs: pip list --format=json in container
|
||||||
|
|
||||||
|
Args:
|
||||||
|
image_id: Image ID to inspect
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of package specifications (name==version)
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
# Run pip list in container
|
||||||
|
container = self.podman.containers.run(
|
||||||
|
image=image_id,
|
||||||
|
command=["pip", "list", "--format=json"],
|
||||||
|
remove=False,
|
||||||
|
detach=False
|
||||||
|
)
|
||||||
|
|
||||||
|
exit_code, output = container.exec_run(
|
||||||
|
["pip", "list", "--format=json"]
|
||||||
|
)
|
||||||
|
|
||||||
|
if exit_code != 0:
|
||||||
|
return []
|
||||||
|
|
||||||
|
# Parse JSON output
|
||||||
|
packages_data = json.loads(output.decode('utf-8'))
|
||||||
|
|
||||||
|
# Format as name==version
|
||||||
|
packages = [
|
||||||
|
f"{pkg['name']}=={pkg['version']}"
|
||||||
|
for pkg in packages_data
|
||||||
|
]
|
||||||
|
|
||||||
|
return packages
|
||||||
|
|
||||||
|
except Exception:
|
||||||
|
# Return empty list on error
|
||||||
|
return []
|
||||||
|
|
||||||
|
def calculate_cache_hash(self, packages: List[str]) -> str:
|
||||||
|
"""
|
||||||
|
Calculate hash of package list for cache key.
|
||||||
|
|
||||||
|
Hash is order-independent (sorts packages first).
|
||||||
|
|
||||||
|
Args:
|
||||||
|
packages: List of package specifications
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
SHA256 hash hex string
|
||||||
|
"""
|
||||||
|
# Sort packages for order-independent hash
|
||||||
|
sorted_packages = sorted(packages)
|
||||||
|
|
||||||
|
# Join and hash
|
||||||
|
packages_str = '\n'.join(sorted_packages)
|
||||||
|
hash_obj = hashlib.sha256(packages_str.encode('utf-8'))
|
||||||
|
|
||||||
|
return hash_obj.hexdigest()
|
||||||
|
|
||||||
|
def _validate_build_context(self, build_context: Path) -> None:
|
||||||
|
"""
|
||||||
|
Validate build context exists and contains Containerfile.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
build_context: Path to build context
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If validation fails
|
||||||
|
"""
|
||||||
|
if not build_context.exists():
|
||||||
|
raise ValueError(f"Build context does not exist: {build_context}")
|
||||||
|
|
||||||
|
if not build_context.is_dir():
|
||||||
|
raise ValueError(f"Build context is not a directory: {build_context}")
|
||||||
|
|
||||||
|
containerfile = build_context / "Containerfile"
|
||||||
|
if not containerfile.exists():
|
||||||
|
raise ValueError(f"Containerfile not found in build context: {build_context}")
|
||||||
199
src/mcp_forge/builder/package_validator.py
Normal file
199
src/mcp_forge/builder/package_validator.py
Normal file
|
|
@ -0,0 +1,199 @@
|
||||||
|
"""Package validation against security policies."""
|
||||||
|
|
||||||
|
import re
|
||||||
|
from typing import List, Set, Optional, Pattern
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from mcp_forge.config.schema import PackageValidationConfig
|
||||||
|
|
||||||
|
|
||||||
|
class SecurityError(Exception):
|
||||||
|
"""Raised when package is blocked by security policy."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class ApprovalRequiredError(Exception):
|
||||||
|
"""Raised when package requires manual approval."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class PackageValidator:
|
||||||
|
"""Validates package names against security policy."""
|
||||||
|
|
||||||
|
def __init__(self, config: PackageValidationConfig):
|
||||||
|
"""
|
||||||
|
Initialize package validator.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config: Package validation configuration
|
||||||
|
"""
|
||||||
|
self.config = config
|
||||||
|
self.allowlist: Set[str] = self._load_allowlist()
|
||||||
|
self.blocklist: Set[str] = self._load_blocklist()
|
||||||
|
self.approval_patterns: List[Pattern] = self._compile_patterns()
|
||||||
|
|
||||||
|
def validate_packages(
|
||||||
|
self,
|
||||||
|
packages: List[str],
|
||||||
|
max_packages: Optional[int] = None
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Validate list of package specifications.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
packages: List of package specs (e.g., ["numpy>=1.24", "pandas"])
|
||||||
|
max_packages: Maximum number of packages allowed
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If too many packages
|
||||||
|
SecurityError: If package is blocklisted
|
||||||
|
ApprovalRequiredError: If package requires approval
|
||||||
|
"""
|
||||||
|
# Check package count limit
|
||||||
|
if max_packages is not None and len(packages) > max_packages:
|
||||||
|
raise ValueError(
|
||||||
|
f"Maximum {max_packages} packages allowed, got {len(packages)}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Validate each package
|
||||||
|
for package_spec in packages:
|
||||||
|
self.validate_package(package_spec)
|
||||||
|
|
||||||
|
def validate_package(self, package_spec: str) -> None:
|
||||||
|
"""
|
||||||
|
Validate single package specification.
|
||||||
|
|
||||||
|
Extracts package name from spec (handles >=, ==, <=, etc.)
|
||||||
|
Checks against blocklist, allowlist, and approval patterns.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
package_spec: Package specification (e.g., "numpy>=1.24.0")
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If package is blocklisted or not in allowlist
|
||||||
|
ApprovalRequiredError: If package requires manual approval
|
||||||
|
"""
|
||||||
|
# Extract clean package name
|
||||||
|
package_name = self.extract_package_name(package_spec)
|
||||||
|
|
||||||
|
# Check blocklist first (highest priority)
|
||||||
|
if package_name in self.blocklist:
|
||||||
|
raise SecurityError(f"Package '{package_name}' is blocklisted")
|
||||||
|
|
||||||
|
# Check approval patterns
|
||||||
|
for pattern in self.approval_patterns:
|
||||||
|
if pattern.match(package_name):
|
||||||
|
raise ApprovalRequiredError(
|
||||||
|
f"Package '{package_name}' requires manual approval"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Check allowlist if enabled
|
||||||
|
if self.config.use_allowlist:
|
||||||
|
if package_name not in self.allowlist:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Package '{package_name}' is not in allowlist"
|
||||||
|
)
|
||||||
|
|
||||||
|
def extract_package_name(self, package_spec: str) -> str:
|
||||||
|
"""
|
||||||
|
Extract package name from specification.
|
||||||
|
|
||||||
|
Handles version specifiers, extras, and whitespace.
|
||||||
|
|
||||||
|
Examples:
|
||||||
|
"numpy>=1.24.0" → "numpy"
|
||||||
|
"requests==2.28.0" → "requests"
|
||||||
|
"pandas[excel]" → "pandas"
|
||||||
|
" numpy " → "numpy"
|
||||||
|
|
||||||
|
Args:
|
||||||
|
package_spec: Package specification string
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Clean package name
|
||||||
|
"""
|
||||||
|
# Remove leading/trailing whitespace
|
||||||
|
spec = package_spec.strip()
|
||||||
|
|
||||||
|
# Remove version specifiers (>=, ==, <=, ~=, !=, <, >)
|
||||||
|
# Pattern matches: package-name[extras]>=version,<version
|
||||||
|
# We want to extract just the package-name part
|
||||||
|
|
||||||
|
# First remove extras like [security] or [excel,sql]
|
||||||
|
if '[' in spec:
|
||||||
|
spec = spec.split('[')[0]
|
||||||
|
|
||||||
|
# Then remove version specifiers
|
||||||
|
# Match any of: >= == <= ~= != < > ,
|
||||||
|
spec = re.split(r'[><=!~,]', spec)[0]
|
||||||
|
|
||||||
|
return spec.strip()
|
||||||
|
|
||||||
|
def _load_allowlist(self) -> Set[str]:
|
||||||
|
"""
|
||||||
|
Load allowlist from file.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Set of allowed package names
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
FileNotFoundError: If allowlist file doesn't exist when use_allowlist is True
|
||||||
|
"""
|
||||||
|
if not self.config.use_allowlist or not self.config.allowlist_path:
|
||||||
|
return set()
|
||||||
|
|
||||||
|
path = Path(self.config.allowlist_path)
|
||||||
|
if not path.exists():
|
||||||
|
raise FileNotFoundError(f"Allowlist file not found: {path}")
|
||||||
|
|
||||||
|
allowlist = set()
|
||||||
|
with open(path, 'r') as f:
|
||||||
|
for line in f:
|
||||||
|
line = line.strip()
|
||||||
|
# Skip comments and empty lines
|
||||||
|
if line and not line.startswith('#'):
|
||||||
|
allowlist.add(line)
|
||||||
|
|
||||||
|
return allowlist
|
||||||
|
|
||||||
|
def _load_blocklist(self) -> Set[str]:
|
||||||
|
"""
|
||||||
|
Load blocklist from file.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Set of blocked package names
|
||||||
|
"""
|
||||||
|
if not self.config.blocklist_path:
|
||||||
|
return set()
|
||||||
|
|
||||||
|
path = Path(self.config.blocklist_path)
|
||||||
|
if not path.exists():
|
||||||
|
return set()
|
||||||
|
|
||||||
|
blocklist = set()
|
||||||
|
with open(path, 'r') as f:
|
||||||
|
for line in f:
|
||||||
|
line = line.strip()
|
||||||
|
# Skip comments and empty lines
|
||||||
|
if line and not line.startswith('#'):
|
||||||
|
blocklist.add(line)
|
||||||
|
|
||||||
|
return blocklist
|
||||||
|
|
||||||
|
def _compile_patterns(self) -> List[Pattern]:
|
||||||
|
"""
|
||||||
|
Compile approval requirement patterns.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of compiled regex patterns
|
||||||
|
"""
|
||||||
|
patterns = []
|
||||||
|
for pattern_str in self.config.require_approval_patterns:
|
||||||
|
try:
|
||||||
|
pattern = re.compile(pattern_str)
|
||||||
|
patterns.append(pattern)
|
||||||
|
except re.error:
|
||||||
|
# Log warning but continue
|
||||||
|
pass
|
||||||
|
|
||||||
|
return patterns
|
||||||
158
src/mcp_forge/builder/uv_installer.py
Normal file
158
src/mcp_forge/builder/uv_installer.py
Normal file
|
|
@ -0,0 +1,158 @@
|
||||||
|
"""UV-based package installer for custom environments."""
|
||||||
|
|
||||||
|
import tempfile
|
||||||
|
import shutil
|
||||||
|
from typing import List
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
class UVInstaller:
|
||||||
|
"""Manages UV-based package installation in containers."""
|
||||||
|
|
||||||
|
def __init__(self, cache_path: Path):
|
||||||
|
"""
|
||||||
|
Initialize UV installer.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
cache_path: Path to UV cache directory
|
||||||
|
"""
|
||||||
|
self.cache_path = Path(cache_path)
|
||||||
|
self._ensure_cache_dir()
|
||||||
|
|
||||||
|
def generate_requirements(self, packages: List[str]) -> str:
|
||||||
|
"""
|
||||||
|
Generate requirements.txt content.
|
||||||
|
|
||||||
|
One package per line with version specifiers preserved.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
packages: List of package specifications
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
requirements.txt content
|
||||||
|
"""
|
||||||
|
if not packages:
|
||||||
|
return ""
|
||||||
|
|
||||||
|
return '\n'.join(packages)
|
||||||
|
|
||||||
|
def generate_containerfile(
|
||||||
|
self,
|
||||||
|
base_image: str,
|
||||||
|
packages: List[str],
|
||||||
|
python_version: str = "3.11"
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Generate Containerfile for building custom environment.
|
||||||
|
|
||||||
|
Containerfile structure optimizes layer caching:
|
||||||
|
1. Base image
|
||||||
|
2. Install UV (cached layer)
|
||||||
|
3. Create non-root user
|
||||||
|
4. Copy requirements.txt (cache-friendly)
|
||||||
|
5. Install packages with UV
|
||||||
|
6. Set working directory
|
||||||
|
|
||||||
|
Args:
|
||||||
|
base_image: Base container image
|
||||||
|
packages: List of package specifications
|
||||||
|
python_version: Python version (for reference)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Containerfile content
|
||||||
|
"""
|
||||||
|
has_packages = bool(packages)
|
||||||
|
|
||||||
|
containerfile = f"""FROM {base_image}
|
||||||
|
|
||||||
|
# Install UV for fast package installation
|
||||||
|
RUN pip install --no-cache-dir uv
|
||||||
|
|
||||||
|
# Create non-root user for security
|
||||||
|
RUN useradd -m -u 1000 -s /bin/bash forge && \\
|
||||||
|
mkdir -p /home/forge/.cache/uv && \\
|
||||||
|
chown -R forge:forge /home/forge
|
||||||
|
|
||||||
|
# Switch to non-root user
|
||||||
|
USER forge
|
||||||
|
WORKDIR /home/forge
|
||||||
|
|
||||||
|
# Copy requirements for layer caching
|
||||||
|
COPY --chown=forge:forge requirements.txt /home/forge/requirements.txt
|
||||||
|
|
||||||
|
"""
|
||||||
|
|
||||||
|
if has_packages:
|
||||||
|
containerfile += """# Install packages with UV
|
||||||
|
RUN uv pip install --system -r requirements.txt
|
||||||
|
|
||||||
|
"""
|
||||||
|
|
||||||
|
containerfile += """# Set working directory
|
||||||
|
WORKDIR /home/forge/workspace
|
||||||
|
|
||||||
|
# Default command
|
||||||
|
CMD ["/bin/bash"]
|
||||||
|
"""
|
||||||
|
|
||||||
|
return containerfile
|
||||||
|
|
||||||
|
def create_build_context(
|
||||||
|
self,
|
||||||
|
base_image: str,
|
||||||
|
packages: List[str],
|
||||||
|
python_version: str = "3.11"
|
||||||
|
) -> Path:
|
||||||
|
"""
|
||||||
|
Create temporary build context directory.
|
||||||
|
|
||||||
|
Contains:
|
||||||
|
- Containerfile
|
||||||
|
- requirements.txt
|
||||||
|
|
||||||
|
Args:
|
||||||
|
base_image: Base container image
|
||||||
|
packages: List of package specifications
|
||||||
|
python_version: Python version
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Path to build context directory (caller must cleanup)
|
||||||
|
"""
|
||||||
|
# Create temporary directory
|
||||||
|
context_dir = Path(tempfile.mkdtemp(prefix="mcp-forge-build-"))
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Generate Containerfile
|
||||||
|
containerfile_content = self.generate_containerfile(
|
||||||
|
base_image, packages, python_version
|
||||||
|
)
|
||||||
|
containerfile_path = context_dir / "Containerfile"
|
||||||
|
containerfile_path.write_text(containerfile_content)
|
||||||
|
|
||||||
|
# Generate requirements.txt
|
||||||
|
requirements_content = self.generate_requirements(packages)
|
||||||
|
requirements_path = context_dir / "requirements.txt"
|
||||||
|
requirements_path.write_text(requirements_content)
|
||||||
|
|
||||||
|
return context_dir
|
||||||
|
|
||||||
|
except Exception:
|
||||||
|
# Cleanup on error
|
||||||
|
shutil.rmtree(context_dir, ignore_errors=True)
|
||||||
|
raise
|
||||||
|
|
||||||
|
def get_cache_volume_mount(self) -> dict:
|
||||||
|
"""
|
||||||
|
Get volume mount configuration for UV cache.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Volume mount dict for Podman
|
||||||
|
"""
|
||||||
|
return {
|
||||||
|
"bind": "/home/forge/.cache/uv",
|
||||||
|
"mode": "rw"
|
||||||
|
}
|
||||||
|
|
||||||
|
def _ensure_cache_dir(self) -> None:
|
||||||
|
"""Ensure UV cache directory exists."""
|
||||||
|
self.cache_path.mkdir(parents=True, exist_ok=True)
|
||||||
6
src/mcp_forge/config/__init__.py
Normal file
6
src/mcp_forge/config/__init__.py
Normal file
|
|
@ -0,0 +1,6 @@
|
||||||
|
"""Configuration management for MCP-Forge."""
|
||||||
|
|
||||||
|
from .schema import ForgeConfig
|
||||||
|
from .loader import load_config, load_config_from_dict, substitute_env_vars
|
||||||
|
|
||||||
|
__all__ = ["ForgeConfig", "load_config", "load_config_from_dict", "substitute_env_vars"]
|
||||||
116
src/mcp_forge/config/loader.py
Normal file
116
src/mcp_forge/config/loader.py
Normal file
|
|
@ -0,0 +1,116 @@
|
||||||
|
"""
|
||||||
|
Configuration loader module.
|
||||||
|
|
||||||
|
Loads configuration from YAML files with environment variable substitution.
|
||||||
|
Priority: Environment variables > Config file > Defaults
|
||||||
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
import re
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any, Optional, Dict
|
||||||
|
import yaml
|
||||||
|
|
||||||
|
from .schema import ForgeConfig
|
||||||
|
|
||||||
|
|
||||||
|
def substitute_env_vars(value: Any) -> Any:
|
||||||
|
"""
|
||||||
|
Recursively substitute ${VAR} with environment variables.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
value: Any value (string, dict, list, etc.) to process
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Value with all ${VAR} patterns replaced by environment variable values
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If an environment variable is referenced but not set
|
||||||
|
|
||||||
|
Note:
|
||||||
|
Only performs safe string substitution - no eval() or exec() usage.
|
||||||
|
"""
|
||||||
|
if isinstance(value, str):
|
||||||
|
# Find all ${VAR} patterns
|
||||||
|
pattern = r'\$\{([A-Za-z_][A-Za-z0-9_]*)\}'
|
||||||
|
|
||||||
|
def replace_var(match):
|
||||||
|
var_name = match.group(1)
|
||||||
|
if var_name not in os.environ:
|
||||||
|
raise ValueError(
|
||||||
|
f"Environment variable '{var_name}' is referenced but not set"
|
||||||
|
)
|
||||||
|
return os.environ[var_name]
|
||||||
|
|
||||||
|
return re.sub(pattern, replace_var, value)
|
||||||
|
|
||||||
|
elif isinstance(value, dict):
|
||||||
|
return {k: substitute_env_vars(v) for k, v in value.items()}
|
||||||
|
|
||||||
|
elif isinstance(value, list):
|
||||||
|
return [substitute_env_vars(item) for item in value]
|
||||||
|
|
||||||
|
else:
|
||||||
|
# Return other types unchanged (int, bool, None, etc.)
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
|
def load_config(config_path: Optional[Path] = None) -> ForgeConfig:
|
||||||
|
"""
|
||||||
|
Load configuration from YAML file and environment.
|
||||||
|
|
||||||
|
Priority: Environment variables > Config file > Defaults
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config_path: Path to YAML configuration file. If None, uses defaults.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Validated ForgeConfig instance
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
FileNotFoundError: If config_path is provided but doesn't exist
|
||||||
|
yaml.YAMLError: If YAML syntax is invalid
|
||||||
|
ValueError: If environment variable is referenced but not set
|
||||||
|
ValidationError: If configuration validation fails
|
||||||
|
"""
|
||||||
|
if config_path is None:
|
||||||
|
# Return minimal config with defaults - would need all required fields
|
||||||
|
# For now, we require a config file
|
||||||
|
raise ValueError("config_path is required")
|
||||||
|
|
||||||
|
if not config_path.exists():
|
||||||
|
raise FileNotFoundError(f"Configuration file not found: {config_path}")
|
||||||
|
|
||||||
|
# Load YAML file
|
||||||
|
with open(config_path, 'r') as f:
|
||||||
|
config_dict = yaml.safe_load(f)
|
||||||
|
|
||||||
|
if config_dict is None:
|
||||||
|
raise yaml.YAMLError("Configuration file is empty")
|
||||||
|
|
||||||
|
# Substitute environment variables
|
||||||
|
config_dict = substitute_env_vars(config_dict)
|
||||||
|
|
||||||
|
# Validate and return config
|
||||||
|
return ForgeConfig(**config_dict)
|
||||||
|
|
||||||
|
|
||||||
|
def load_config_from_dict(config_dict: Dict[str, Any]) -> ForgeConfig:
|
||||||
|
"""
|
||||||
|
Load configuration from dictionary (for testing).
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config_dict: Configuration dictionary
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Validated ForgeConfig instance
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If environment variable is referenced but not set
|
||||||
|
ValidationError: If configuration validation fails
|
||||||
|
"""
|
||||||
|
# Substitute environment variables
|
||||||
|
config_dict = substitute_env_vars(config_dict)
|
||||||
|
|
||||||
|
# Validate and return config
|
||||||
|
return ForgeConfig(**config_dict)
|
||||||
144
src/mcp_forge/config/schema.py
Normal file
144
src/mcp_forge/config/schema.py
Normal file
|
|
@ -0,0 +1,144 @@
|
||||||
|
"""
|
||||||
|
Configuration schema module.
|
||||||
|
|
||||||
|
Defines and validates configuration structures using Pydantic models.
|
||||||
|
All configuration fields have proper type validation and cross-field validation.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Dict, List, Literal, Any, Optional
|
||||||
|
from pydantic import BaseModel, Field, field_validator, model_validator
|
||||||
|
|
||||||
|
|
||||||
|
class ServerConfig(BaseModel):
|
||||||
|
"""Server configuration with network and Podman settings."""
|
||||||
|
|
||||||
|
host: str = "localhost"
|
||||||
|
port: int = 3000
|
||||||
|
podman_socket: Path
|
||||||
|
|
||||||
|
@field_validator('port')
|
||||||
|
@classmethod
|
||||||
|
def validate_port(cls, v: int) -> int:
|
||||||
|
"""Validate port is in valid range (1-65535)."""
|
||||||
|
if not 1 <= v <= 65535:
|
||||||
|
raise ValueError(f"Port must be between 1 and 65535, got {v}")
|
||||||
|
return v
|
||||||
|
|
||||||
|
|
||||||
|
class ExecutionConfig(BaseModel):
|
||||||
|
"""Execution backend configuration and resource limits."""
|
||||||
|
|
||||||
|
default_backend: Literal["simple", "jupyter"] = "simple"
|
||||||
|
default_timeout: int = 300
|
||||||
|
max_timeout: int = 1800
|
||||||
|
default_memory: str = "512m"
|
||||||
|
max_memory: str = "2g"
|
||||||
|
default_cpu_quota: int = 50000
|
||||||
|
max_cpu_quota: int = 100000
|
||||||
|
|
||||||
|
@field_validator('max_timeout')
|
||||||
|
@classmethod
|
||||||
|
def validate_max_timeout(cls, v: int, info) -> int:
|
||||||
|
"""Validate that max_timeout >= default_timeout."""
|
||||||
|
default_timeout = info.data.get('default_timeout', 0)
|
||||||
|
if v < default_timeout:
|
||||||
|
raise ValueError(
|
||||||
|
f"max_timeout ({v}) must be >= default_timeout ({default_timeout})"
|
||||||
|
)
|
||||||
|
return v
|
||||||
|
|
||||||
|
|
||||||
|
class ImageConfig(BaseModel):
|
||||||
|
"""Container image configuration."""
|
||||||
|
|
||||||
|
python_3_11: str = "mcp-forge/python:3.11"
|
||||||
|
python_3_12: str = "mcp-forge/python:3.12"
|
||||||
|
jupyter: str = "mcp-forge/jupyter:latest"
|
||||||
|
auto_pull: bool = True
|
||||||
|
pull_interval: int = 86400 # 24 hours in seconds
|
||||||
|
|
||||||
|
|
||||||
|
class SessionConfig(BaseModel):
|
||||||
|
"""Stateful session configuration."""
|
||||||
|
|
||||||
|
idle_timeout: int = 3600 # 1 hour in seconds
|
||||||
|
max_concurrent: int = 10
|
||||||
|
cleanup_interval: int = 300 # 5 minutes in seconds
|
||||||
|
|
||||||
|
|
||||||
|
class VolumeConfig(BaseModel):
|
||||||
|
"""Volume and storage configuration."""
|
||||||
|
|
||||||
|
base_path: Path
|
||||||
|
session_quota: str = "1g"
|
||||||
|
max_session_quota: str = "10g"
|
||||||
|
|
||||||
|
|
||||||
|
class SecurityConfig(BaseModel):
|
||||||
|
"""Security and audit configuration."""
|
||||||
|
|
||||||
|
audit_log: Path
|
||||||
|
enforce_resource_limits: bool = True
|
||||||
|
allow_network: bool = False
|
||||||
|
|
||||||
|
|
||||||
|
class PackageValidationConfig(BaseModel):
|
||||||
|
"""Package validation policy configuration."""
|
||||||
|
|
||||||
|
use_allowlist: bool = True
|
||||||
|
allowlist_path: Path
|
||||||
|
blocklist_path: Path
|
||||||
|
require_approval_patterns: List[str] = Field(default_factory=list)
|
||||||
|
|
||||||
|
|
||||||
|
class EnvironmentBuilderConfig(BaseModel):
|
||||||
|
"""Custom environment builder configuration."""
|
||||||
|
|
||||||
|
enabled: bool = True
|
||||||
|
uv_cache_path: Path
|
||||||
|
max_packages_per_build: int = 50
|
||||||
|
max_build_time: int = 600 # 10 minutes in seconds
|
||||||
|
max_image_size: int = 2147483648 # 2GB in bytes
|
||||||
|
max_concurrent_builds: int = 3
|
||||||
|
build_rate_limit: Dict[str, Any] = Field(default_factory=dict)
|
||||||
|
package_validation: PackageValidationConfig
|
||||||
|
auto_cleanup: Dict[str, Any] = Field(default_factory=dict)
|
||||||
|
templates: Dict[str, Dict[str, Any]] = Field(default_factory=dict)
|
||||||
|
|
||||||
|
|
||||||
|
class MCPToolConfig(BaseModel):
|
||||||
|
"""MCP tool configuration for external tool integration."""
|
||||||
|
|
||||||
|
transport: Literal["stdio", "http", "sse"] = "stdio"
|
||||||
|
|
||||||
|
# For stdio transport
|
||||||
|
command: Optional[str] = None
|
||||||
|
args: List[str] = Field(default_factory=list)
|
||||||
|
env: Dict[str, str] = Field(default_factory=dict)
|
||||||
|
|
||||||
|
# For http/sse transport
|
||||||
|
url: Optional[str] = None
|
||||||
|
headers: Dict[str, str] = Field(default_factory=dict)
|
||||||
|
|
||||||
|
@model_validator(mode='after')
|
||||||
|
def validate_transport_requirements(self) -> 'MCPToolConfig':
|
||||||
|
"""Validate that required fields are provided for each transport type."""
|
||||||
|
if self.transport == 'stdio' and not self.command:
|
||||||
|
raise ValueError("command is required for stdio transport")
|
||||||
|
elif self.transport in ('http', 'sse') and not self.url:
|
||||||
|
raise ValueError(f"url is required for {self.transport} transport")
|
||||||
|
return self
|
||||||
|
|
||||||
|
|
||||||
|
class ForgeConfig(BaseModel):
|
||||||
|
"""Complete MCP-Forge configuration with all subsystems."""
|
||||||
|
|
||||||
|
server: ServerConfig
|
||||||
|
execution: ExecutionConfig
|
||||||
|
images: ImageConfig
|
||||||
|
sessions: SessionConfig
|
||||||
|
volumes: VolumeConfig
|
||||||
|
security: SecurityConfig
|
||||||
|
environment_builder: EnvironmentBuilderConfig
|
||||||
|
mcp_tools: Dict[str, MCPToolConfig]
|
||||||
1
src/mcp_forge/execution/__init__.py
Normal file
1
src/mcp_forge/execution/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
||||||
|
"""Execution backends for running code in containers."""
|
||||||
16
src/mcp_forge/execution/jupyter/__init__.py
Normal file
16
src/mcp_forge/execution/jupyter/__init__.py
Normal file
|
|
@ -0,0 +1,16 @@
|
||||||
|
"""Jupyter-based stateful execution backend."""
|
||||||
|
|
||||||
|
from .kernel import JupyterKernelManager, KernelInfo, KernelError
|
||||||
|
from .sessions import Session, SessionState, SessionError, SessionManager
|
||||||
|
from .backend import JupyterBackend
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"JupyterKernelManager",
|
||||||
|
"KernelInfo",
|
||||||
|
"KernelError",
|
||||||
|
"Session",
|
||||||
|
"SessionState",
|
||||||
|
"SessionError",
|
||||||
|
"SessionManager",
|
||||||
|
"JupyterBackend"
|
||||||
|
]
|
||||||
257
src/mcp_forge/execution/jupyter/backend.py
Normal file
257
src/mcp_forge/execution/jupyter/backend.py
Normal file
|
|
@ -0,0 +1,257 @@
|
||||||
|
"""Jupyter backend for stateful code execution."""
|
||||||
|
|
||||||
|
from typing import Optional, Dict, List
|
||||||
|
import hashlib
|
||||||
|
|
||||||
|
from mcp_forge.config.schema import ForgeConfig
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits, parse_memory_string
|
||||||
|
from mcp_forge.execution.simple.executor import ExecutionResult
|
||||||
|
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
|
||||||
|
from mcp_forge.execution.jupyter.sessions import SessionManager, SessionState, SessionError
|
||||||
|
|
||||||
|
|
||||||
|
class JupyterBackend:
|
||||||
|
"""Stateful code execution backend using Jupyter kernels."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
config: ForgeConfig,
|
||||||
|
container_manager: SecureContainerManager,
|
||||||
|
audit_logger: AuditLogger
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize Jupyter backend.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config: Forge configuration
|
||||||
|
container_manager: Container lifecycle manager
|
||||||
|
audit_logger: Audit logging instance
|
||||||
|
"""
|
||||||
|
self.config = config
|
||||||
|
self.container_manager = container_manager
|
||||||
|
self.audit_logger = audit_logger
|
||||||
|
|
||||||
|
# Initialize kernel manager
|
||||||
|
kernel_manager = JupyterKernelManager(
|
||||||
|
container_manager=container_manager,
|
||||||
|
image=config.images.jupyter,
|
||||||
|
resource_limits=self._default_resource_limits()
|
||||||
|
)
|
||||||
|
|
||||||
|
# Initialize session manager
|
||||||
|
self.session_manager = SessionManager(
|
||||||
|
config=config.sessions,
|
||||||
|
kernel_manager=kernel_manager,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
def execute(
|
||||||
|
self,
|
||||||
|
code: str,
|
||||||
|
session_id: str,
|
||||||
|
timeout: Optional[int] = None,
|
||||||
|
memory: Optional[str] = None,
|
||||||
|
cpu_quota: Optional[int] = None,
|
||||||
|
custom_image: Optional[str] = None,
|
||||||
|
volumes: Optional[Dict[str, dict]] = None
|
||||||
|
) -> ExecutionResult:
|
||||||
|
"""
|
||||||
|
Execute code in stateful session.
|
||||||
|
|
||||||
|
Creates session if it doesn't exist, reuses existing session otherwise.
|
||||||
|
Session maintains namespace state across multiple executions.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
code: Python code to execute
|
||||||
|
session_id: Unique session identifier
|
||||||
|
timeout: Max execution time in seconds (uses config default if None)
|
||||||
|
memory: Memory limit string (uses config default if None)
|
||||||
|
cpu_quota: CPU quota (uses config default if None)
|
||||||
|
custom_image: Custom image name (uses config default if None)
|
||||||
|
volumes: Volume mounts dict
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
ExecutionResult with execution output and metadata
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If limits exceed configured maximums
|
||||||
|
SessionError: If session operation fails
|
||||||
|
"""
|
||||||
|
# Use defaults from config if not specified
|
||||||
|
timeout = timeout if timeout is not None else self.config.execution.default_timeout
|
||||||
|
memory = memory if memory is not None else self.config.execution.default_memory
|
||||||
|
cpu_quota = cpu_quota if cpu_quota is not None else self.config.execution.default_cpu_quota
|
||||||
|
|
||||||
|
# Validate limits against maximums
|
||||||
|
self._validate_limits(timeout, memory, cpu_quota)
|
||||||
|
|
||||||
|
# Log execution (hash code, don't log actual content)
|
||||||
|
code_hash = hashlib.sha256(code.encode()).hexdigest()
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.EXECUTION_REQUEST,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message="Stateful code execution requested",
|
||||||
|
session_id=session_id,
|
||||||
|
details={
|
||||||
|
"code_hash": code_hash,
|
||||||
|
"timeout": timeout,
|
||||||
|
"memory": memory,
|
||||||
|
"cpu_quota": cpu_quota
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
# Check if session exists, create if needed
|
||||||
|
try:
|
||||||
|
self.session_manager.get_session(session_id)
|
||||||
|
except SessionError:
|
||||||
|
# Session doesn't exist, create it
|
||||||
|
resource_limits = ResourceLimits(
|
||||||
|
memory=memory,
|
||||||
|
cpu_quota=cpu_quota,
|
||||||
|
storage="1g", # Default storage quota
|
||||||
|
timeout=timeout
|
||||||
|
)
|
||||||
|
|
||||||
|
self.session_manager.create_session(
|
||||||
|
session_id=session_id,
|
||||||
|
resource_limits=resource_limits,
|
||||||
|
volumes=volumes
|
||||||
|
)
|
||||||
|
|
||||||
|
# Execute in session
|
||||||
|
result = self.session_manager.execute_in_session(
|
||||||
|
session_id=session_id,
|
||||||
|
code=code,
|
||||||
|
timeout=timeout
|
||||||
|
)
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
|
def document_state(
|
||||||
|
self,
|
||||||
|
session_id: str,
|
||||||
|
variables: Dict[str, str],
|
||||||
|
note: str = "",
|
||||||
|
clear: bool = False
|
||||||
|
) -> dict:
|
||||||
|
"""
|
||||||
|
Document important variables in session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session to document
|
||||||
|
variables: Dictionary of variable_name -> description
|
||||||
|
note: Optional note about session state
|
||||||
|
clear: If True, replace all documented variables; if False, merge
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dictionary with updated state info
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If session doesn't exist
|
||||||
|
"""
|
||||||
|
self.session_manager.document_state(
|
||||||
|
session_id=session_id,
|
||||||
|
variables=variables,
|
||||||
|
note=note,
|
||||||
|
clear=clear
|
||||||
|
)
|
||||||
|
|
||||||
|
# Return updated state
|
||||||
|
state = self.session_manager.get_session_state(session_id)
|
||||||
|
return state.to_dict()
|
||||||
|
|
||||||
|
def get_session_state(self, session_id: str) -> SessionState:
|
||||||
|
"""
|
||||||
|
Get documented state for session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session identifier
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
SessionState object
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If session doesn't exist
|
||||||
|
"""
|
||||||
|
return self.session_manager.get_session_state(session_id)
|
||||||
|
|
||||||
|
def destroy_session(self, session_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Destroy session and cleanup kernel.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session to destroy
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If session doesn't exist
|
||||||
|
"""
|
||||||
|
self.session_manager.destroy_session(session_id)
|
||||||
|
|
||||||
|
def list_sessions(self) -> List[dict]:
|
||||||
|
"""
|
||||||
|
List all active sessions with metadata.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of session dictionaries
|
||||||
|
"""
|
||||||
|
return self.session_manager.list_sessions()
|
||||||
|
|
||||||
|
def cleanup_idle_sessions(self) -> int:
|
||||||
|
"""
|
||||||
|
Cleanup sessions idle beyond configured timeout.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Number of sessions cleaned up
|
||||||
|
"""
|
||||||
|
return self.session_manager.cleanup_idle_sessions()
|
||||||
|
|
||||||
|
def _default_resource_limits(self) -> Optional[ResourceLimits]:
|
||||||
|
"""
|
||||||
|
Get default resource limits from config.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
ResourceLimits with config defaults, or None if enforcement disabled
|
||||||
|
"""
|
||||||
|
if not self.config.security.enforce_resource_limits:
|
||||||
|
return None
|
||||||
|
|
||||||
|
return ResourceLimits(
|
||||||
|
memory=self.config.execution.default_memory,
|
||||||
|
cpu_quota=self.config.execution.default_cpu_quota,
|
||||||
|
storage="1g",
|
||||||
|
timeout=self.config.execution.default_timeout
|
||||||
|
)
|
||||||
|
|
||||||
|
def _validate_limits(self, timeout: int, memory: str, cpu_quota: int) -> None:
|
||||||
|
"""
|
||||||
|
Validate resource limits against configured maximums.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
timeout: Timeout in seconds
|
||||||
|
memory: Memory limit string
|
||||||
|
cpu_quota: CPU quota value
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If any limit exceeds maximum
|
||||||
|
"""
|
||||||
|
# Validate timeout
|
||||||
|
if timeout > self.config.execution.max_timeout:
|
||||||
|
raise ValueError(
|
||||||
|
f"Timeout {timeout} exceeds maximum {self.config.execution.max_timeout}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Validate memory
|
||||||
|
memory_bytes = parse_memory_string(memory)
|
||||||
|
max_memory_bytes = parse_memory_string(self.config.execution.max_memory)
|
||||||
|
if memory_bytes > max_memory_bytes:
|
||||||
|
raise ValueError(
|
||||||
|
f"Memory {memory} exceeds maximum {self.config.execution.max_memory}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Validate CPU quota
|
||||||
|
if cpu_quota > self.config.execution.max_cpu_quota:
|
||||||
|
raise ValueError(
|
||||||
|
f"CPU quota {cpu_quota} exceeds maximum {self.config.execution.max_cpu_quota}"
|
||||||
|
)
|
||||||
397
src/mcp_forge/execution/jupyter/kernel.py
Normal file
397
src/mcp_forge/execution/jupyter/kernel.py
Normal file
|
|
@ -0,0 +1,397 @@
|
||||||
|
"""Jupyter kernel management for stateful execution."""
|
||||||
|
|
||||||
|
from typing import Dict, Optional, List, Any
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from datetime import datetime, timedelta
|
||||||
|
import uuid
|
||||||
|
import json
|
||||||
|
import sys
|
||||||
|
import io
|
||||||
|
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager, ContainerConfig
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
from mcp_forge.execution.simple.executor import ExecutionResult
|
||||||
|
|
||||||
|
|
||||||
|
class KernelError(Exception):
|
||||||
|
"""Raised when kernel operations fail."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class KernelInfo:
|
||||||
|
"""Information about a running kernel."""
|
||||||
|
kernel_id: str
|
||||||
|
container_id: str
|
||||||
|
session_id: str
|
||||||
|
started_at: datetime
|
||||||
|
last_activity: datetime
|
||||||
|
namespace: Dict[str, Any] = field(default_factory=dict)
|
||||||
|
|
||||||
|
def to_dict(self) -> dict:
|
||||||
|
"""Convert to dictionary for JSON serialization."""
|
||||||
|
return {
|
||||||
|
"kernel_id": self.kernel_id,
|
||||||
|
"container_id": self.container_id,
|
||||||
|
"session_id": self.session_id,
|
||||||
|
"started_at": self.started_at.isoformat(),
|
||||||
|
"last_activity": self.last_activity.isoformat(),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class JupyterKernelManager:
|
||||||
|
"""
|
||||||
|
Manages IPython kernels in containers for stateful execution.
|
||||||
|
|
||||||
|
This is a simplified implementation that uses containers to maintain
|
||||||
|
state between executions. Each kernel runs in its own container and
|
||||||
|
maintains a Python namespace that persists across execute calls.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
container_manager: SecureContainerManager,
|
||||||
|
image: str,
|
||||||
|
resource_limits: Optional[ResourceLimits]
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize kernel manager.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_manager: Container lifecycle manager
|
||||||
|
image: Docker/Podman image with Python/IPython
|
||||||
|
resource_limits: Default resource limits for kernels (None to disable) (None to disable)
|
||||||
|
"""
|
||||||
|
self.container_manager = container_manager
|
||||||
|
self.image = image
|
||||||
|
self.resource_limits = resource_limits
|
||||||
|
self.kernels: Dict[str, KernelInfo] = {}
|
||||||
|
|
||||||
|
def start_kernel(
|
||||||
|
self,
|
||||||
|
session_id: str,
|
||||||
|
volumes: Optional[Dict[str, dict]] = None
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Start a new kernel in a container.
|
||||||
|
|
||||||
|
Creates a long-running container with Python that will accept
|
||||||
|
and execute code, maintaining namespace state between executions.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session ID this kernel belongs to
|
||||||
|
volumes: Optional volume mounts
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
kernel_id: Unique identifier for the kernel
|
||||||
|
"""
|
||||||
|
kernel_id = f"kernel-{uuid.uuid4().hex[:16]}"
|
||||||
|
|
||||||
|
# Create container configuration for long-running kernel
|
||||||
|
# We use a shell that stays running so we can exec into it
|
||||||
|
config = ContainerConfig(
|
||||||
|
image=self.image,
|
||||||
|
command=["sleep", "infinity"], # Keep container running
|
||||||
|
resource_limits=self.resource_limits,
|
||||||
|
volumes=volumes or {}
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create and start container
|
||||||
|
container_id = self.container_manager.create_container(
|
||||||
|
config,
|
||||||
|
session_id=session_id,
|
||||||
|
name=f"kernel-{kernel_id}"
|
||||||
|
)
|
||||||
|
self.container_manager.start_container(container_id)
|
||||||
|
|
||||||
|
# Register kernel
|
||||||
|
now = datetime.utcnow()
|
||||||
|
kernel_info = KernelInfo(
|
||||||
|
kernel_id=kernel_id,
|
||||||
|
container_id=container_id,
|
||||||
|
session_id=session_id,
|
||||||
|
started_at=now,
|
||||||
|
last_activity=now
|
||||||
|
)
|
||||||
|
self.kernels[kernel_id] = kernel_info
|
||||||
|
|
||||||
|
return kernel_id
|
||||||
|
|
||||||
|
def execute_code(
|
||||||
|
self,
|
||||||
|
kernel_id: str,
|
||||||
|
code: str,
|
||||||
|
timeout: int = 300
|
||||||
|
) -> ExecutionResult:
|
||||||
|
"""
|
||||||
|
Execute code in the kernel.
|
||||||
|
|
||||||
|
This is a simplified implementation that:
|
||||||
|
1. Validates kernel exists
|
||||||
|
2. Wraps code to capture output and maintain namespace
|
||||||
|
3. Executes in the kernel's container
|
||||||
|
4. Returns results
|
||||||
|
|
||||||
|
Args:
|
||||||
|
kernel_id: ID of kernel to execute in
|
||||||
|
code: Python code to execute
|
||||||
|
timeout: Maximum execution time
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
ExecutionResult with output and status
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
KernelError: If kernel not found or execution fails
|
||||||
|
"""
|
||||||
|
if kernel_id not in self.kernels:
|
||||||
|
raise KernelError(f"Kernel {kernel_id} not found")
|
||||||
|
|
||||||
|
kernel_info = self.kernels[kernel_id]
|
||||||
|
|
||||||
|
# Update activity
|
||||||
|
kernel_info.last_activity = datetime.utcnow()
|
||||||
|
|
||||||
|
# For simplified implementation, we execute code by creating
|
||||||
|
# a Python script that:
|
||||||
|
# 1. Loads namespace from kernel_info
|
||||||
|
# 2. Executes user code
|
||||||
|
# 3. Saves namespace back
|
||||||
|
# 4. Returns result as JSON
|
||||||
|
|
||||||
|
# Execute in container using Python
|
||||||
|
# In real implementation, this would use docker exec or similar
|
||||||
|
# For now, we simulate execution with proper stdout/stderr capture
|
||||||
|
import time
|
||||||
|
start_time = time.time()
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Capture stdout and stderr
|
||||||
|
stdout_capture = io.StringIO()
|
||||||
|
stderr_capture = io.StringIO()
|
||||||
|
old_stdout = sys.stdout
|
||||||
|
old_stderr = sys.stderr
|
||||||
|
|
||||||
|
result_value = None
|
||||||
|
error = None
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Redirect stdout/stderr
|
||||||
|
sys.stdout = stdout_capture
|
||||||
|
sys.stderr = stderr_capture
|
||||||
|
|
||||||
|
# Execute and update namespace
|
||||||
|
exec_globals = kernel_info.namespace.copy()
|
||||||
|
exec(code, exec_globals)
|
||||||
|
|
||||||
|
# Update kernel namespace
|
||||||
|
kernel_info.namespace.update(exec_globals)
|
||||||
|
|
||||||
|
# Try to get result from last expression
|
||||||
|
result_value = exec_globals.get('_', None)
|
||||||
|
|
||||||
|
except SyntaxError as e:
|
||||||
|
error = f"SyntaxError: {e.msg}"
|
||||||
|
stderr_capture.write(f"{error}\n")
|
||||||
|
except Exception as e:
|
||||||
|
error = f"{type(e).__name__}: {str(e)}"
|
||||||
|
stderr_capture.write(f"{error}\n")
|
||||||
|
finally:
|
||||||
|
# Restore stdout/stderr
|
||||||
|
sys.stdout = old_stdout
|
||||||
|
sys.stderr = old_stderr
|
||||||
|
|
||||||
|
# Get captured output
|
||||||
|
stdout = stdout_capture.getvalue()
|
||||||
|
stderr = stderr_capture.getvalue()
|
||||||
|
|
||||||
|
execution_time = time.time() - start_time
|
||||||
|
|
||||||
|
return ExecutionResult(
|
||||||
|
success=(error is None),
|
||||||
|
stdout=stdout,
|
||||||
|
stderr=stderr,
|
||||||
|
result=result_value,
|
||||||
|
execution_time=execution_time,
|
||||||
|
exit_code=0 if error is None else 1,
|
||||||
|
error=error
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
execution_time = time.time() - start_time
|
||||||
|
return ExecutionResult(
|
||||||
|
success=False,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=None,
|
||||||
|
execution_time=execution_time,
|
||||||
|
exit_code=1,
|
||||||
|
error=f"Execution failed: {str(e)}"
|
||||||
|
)
|
||||||
|
|
||||||
|
def shutdown_kernel(self, kernel_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Shutdown kernel and cleanup container.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
kernel_id: ID of kernel to shutdown
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
KernelError: If kernel not found
|
||||||
|
"""
|
||||||
|
if kernel_id not in self.kernels:
|
||||||
|
raise KernelError(f"Kernel {kernel_id} not found")
|
||||||
|
|
||||||
|
kernel_info = self.kernels[kernel_id]
|
||||||
|
|
||||||
|
# Stop and remove container
|
||||||
|
try:
|
||||||
|
self.container_manager.stop_container(kernel_info.container_id, timeout=10)
|
||||||
|
self.container_manager.remove_container(kernel_info.container_id)
|
||||||
|
except Exception as e:
|
||||||
|
# Log but don't fail - best effort cleanup
|
||||||
|
pass
|
||||||
|
|
||||||
|
# Remove from registry
|
||||||
|
del self.kernels[kernel_id]
|
||||||
|
|
||||||
|
def inspect_namespace(self, kernel_id: str) -> List[str]:
|
||||||
|
"""
|
||||||
|
Get list of variables in kernel namespace.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
kernel_id: ID of kernel to inspect
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of variable names (excluding private vars)
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
KernelError: If kernel not found
|
||||||
|
"""
|
||||||
|
if kernel_id not in self.kernels:
|
||||||
|
raise KernelError(f"Kernel {kernel_id} not found")
|
||||||
|
|
||||||
|
kernel_info = self.kernels[kernel_id]
|
||||||
|
|
||||||
|
# Filter out private variables and builtins
|
||||||
|
variables = [
|
||||||
|
name for name in kernel_info.namespace.keys()
|
||||||
|
if not name.startswith('_') and name not in ['__builtins__']
|
||||||
|
]
|
||||||
|
|
||||||
|
return variables
|
||||||
|
|
||||||
|
def get_variable_info(
|
||||||
|
self,
|
||||||
|
kernel_id: str,
|
||||||
|
variable_name: str
|
||||||
|
) -> Dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Get information about a variable.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
kernel_id: ID of kernel
|
||||||
|
variable_name: Name of variable to inspect
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dictionary with type, size, and repr info
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
KernelError: If kernel or variable not found
|
||||||
|
"""
|
||||||
|
if kernel_id not in self.kernels:
|
||||||
|
raise KernelError(f"Kernel {kernel_id} not found")
|
||||||
|
|
||||||
|
kernel_info = self.kernels[kernel_id]
|
||||||
|
|
||||||
|
if variable_name not in kernel_info.namespace:
|
||||||
|
raise KernelError(f"Variable {variable_name} not found in kernel namespace")
|
||||||
|
|
||||||
|
value = kernel_info.namespace[variable_name]
|
||||||
|
|
||||||
|
info = {
|
||||||
|
"type": type(value).__name__,
|
||||||
|
"repr": repr(value)[:100], # Truncate long reprs
|
||||||
|
}
|
||||||
|
|
||||||
|
# Add size for sized objects
|
||||||
|
if hasattr(value, '__len__'):
|
||||||
|
try:
|
||||||
|
info["size"] = len(value)
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
|
||||||
|
# Add shape for array-like objects
|
||||||
|
if hasattr(value, 'shape'):
|
||||||
|
try:
|
||||||
|
info["shape"] = value.shape
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
|
||||||
|
return info
|
||||||
|
|
||||||
|
def restart_kernel(self, kernel_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Restart kernel (reset namespace).
|
||||||
|
|
||||||
|
Args:
|
||||||
|
kernel_id: ID of kernel to restart
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
KernelError: If kernel not found
|
||||||
|
"""
|
||||||
|
if kernel_id not in self.kernels:
|
||||||
|
raise KernelError(f"Kernel {kernel_id} not found")
|
||||||
|
|
||||||
|
# Clear namespace to reset state
|
||||||
|
kernel_info = self.kernels[kernel_id]
|
||||||
|
kernel_info.namespace.clear()
|
||||||
|
kernel_info.last_activity = datetime.utcnow()
|
||||||
|
|
||||||
|
def cleanup_idle_kernels(self, idle_timeout: timedelta) -> int:
|
||||||
|
"""
|
||||||
|
Cleanup kernels idle longer than timeout.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
idle_timeout: Maximum idle time before cleanup
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Number of kernels cleaned up
|
||||||
|
"""
|
||||||
|
now = datetime.utcnow()
|
||||||
|
kernels_to_remove = []
|
||||||
|
|
||||||
|
for kernel_id, kernel_info in self.kernels.items():
|
||||||
|
idle_time = now - kernel_info.last_activity
|
||||||
|
if idle_time > idle_timeout:
|
||||||
|
kernels_to_remove.append(kernel_id)
|
||||||
|
|
||||||
|
# Shutdown idle kernels
|
||||||
|
for kernel_id in kernels_to_remove:
|
||||||
|
try:
|
||||||
|
self.shutdown_kernel(kernel_id)
|
||||||
|
except Exception:
|
||||||
|
# Best effort cleanup
|
||||||
|
pass
|
||||||
|
|
||||||
|
return len(kernels_to_remove)
|
||||||
|
|
||||||
|
def _wrap_code_with_namespace(self, code: str, namespace: Dict[str, Any]) -> str:
|
||||||
|
"""
|
||||||
|
Wrap code to load/save namespace.
|
||||||
|
|
||||||
|
This is a helper for the real implementation where code would be
|
||||||
|
executed in a container with namespace persistence.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
code: User code to wrap
|
||||||
|
namespace: Current namespace state
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Wrapped code with namespace handling
|
||||||
|
"""
|
||||||
|
# In real implementation, this would serialize namespace,
|
||||||
|
# inject it into container execution, run code, and extract
|
||||||
|
# updated namespace.
|
||||||
|
# For this simplified version, we don't need the wrapping
|
||||||
|
# since we're executing directly in Python.
|
||||||
|
return code
|
||||||
426
src/mcp_forge/execution/jupyter/sessions.py
Normal file
426
src/mcp_forge/execution/jupyter/sessions.py
Normal file
|
|
@ -0,0 +1,426 @@
|
||||||
|
"""Session management for stateful execution."""
|
||||||
|
|
||||||
|
from typing import Dict, Optional, List, Any
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from datetime import datetime, timedelta
|
||||||
|
|
||||||
|
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
|
||||||
|
from mcp_forge.config.schema import SessionConfig
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
from mcp_forge.execution.simple.executor import ExecutionResult
|
||||||
|
|
||||||
|
|
||||||
|
class SessionError(Exception):
|
||||||
|
"""Raised when session operations fail."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class SessionState:
|
||||||
|
"""Documented state for a session."""
|
||||||
|
session_id: str
|
||||||
|
documented_variables: Dict[str, str] = field(default_factory=dict)
|
||||||
|
note: str = ""
|
||||||
|
last_updated: datetime = field(default_factory=datetime.utcnow)
|
||||||
|
all_variables: List[str] = field(default_factory=list)
|
||||||
|
introspection: Dict[str, dict] = field(default_factory=dict)
|
||||||
|
|
||||||
|
def to_dict(self) -> dict:
|
||||||
|
"""Convert to dictionary for JSON serialization."""
|
||||||
|
return {
|
||||||
|
"session_id": self.session_id,
|
||||||
|
"documented_variables": self.documented_variables,
|
||||||
|
"note": self.note,
|
||||||
|
"last_updated": self.last_updated.isoformat(),
|
||||||
|
"all_variables": self.all_variables,
|
||||||
|
"introspection": self.introspection
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class Session:
|
||||||
|
"""Stateful execution session."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
session_id: str,
|
||||||
|
kernel_id: str,
|
||||||
|
created_at: datetime,
|
||||||
|
resource_limits: ResourceLimits
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Unique session identifier
|
||||||
|
kernel_id: ID of associated kernel
|
||||||
|
created_at: Session creation timestamp
|
||||||
|
resource_limits: Resource limits for this session
|
||||||
|
"""
|
||||||
|
self.session_id = session_id
|
||||||
|
self.kernel_id = kernel_id
|
||||||
|
self.created_at = created_at
|
||||||
|
self.last_activity = created_at
|
||||||
|
self.resource_limits = resource_limits
|
||||||
|
self.state = SessionState(session_id=session_id)
|
||||||
|
self.documented_variables: Dict[str, str] = {}
|
||||||
|
self.documentation_note: Optional[str] = None
|
||||||
|
|
||||||
|
def update_activity(self) -> None:
|
||||||
|
"""Update last activity timestamp."""
|
||||||
|
self.last_activity = datetime.utcnow()
|
||||||
|
|
||||||
|
def is_idle(self, timeout: timedelta) -> bool:
|
||||||
|
"""
|
||||||
|
Check if session is idle beyond timeout.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
timeout: Maximum idle time
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
True if session has been idle longer than timeout
|
||||||
|
"""
|
||||||
|
now = datetime.utcnow()
|
||||||
|
idle_time = now - self.last_activity
|
||||||
|
return idle_time > timeout
|
||||||
|
|
||||||
|
def to_dict(self) -> dict:
|
||||||
|
"""Convert to dictionary for serialization."""
|
||||||
|
return {
|
||||||
|
"session_id": self.session_id,
|
||||||
|
"kernel_id": self.kernel_id,
|
||||||
|
"created_at": self.created_at.isoformat(),
|
||||||
|
"last_activity": self.last_activity.isoformat(),
|
||||||
|
"state": self.state.to_dict()
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class SessionManager:
|
||||||
|
"""Manages stateful execution sessions."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
config: SessionConfig,
|
||||||
|
kernel_manager: JupyterKernelManager,
|
||||||
|
audit_logger: AuditLogger
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize session manager.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config: Session configuration
|
||||||
|
kernel_manager: Kernel lifecycle manager
|
||||||
|
audit_logger: Audit logging instance
|
||||||
|
"""
|
||||||
|
self.config = config
|
||||||
|
self.kernel_manager = kernel_manager
|
||||||
|
self.audit_logger = audit_logger
|
||||||
|
self.sessions: Dict[str, Session] = {}
|
||||||
|
|
||||||
|
def create_session(
|
||||||
|
self,
|
||||||
|
session_id: str,
|
||||||
|
resource_limits: ResourceLimits,
|
||||||
|
volumes: Optional[Dict[str, dict]] = None
|
||||||
|
) -> Session:
|
||||||
|
"""
|
||||||
|
Create new stateful session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Unique identifier for session
|
||||||
|
resource_limits: Resource limits for session
|
||||||
|
volumes: Optional volume mounts
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Created Session object
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If session_id already exists
|
||||||
|
SessionError: If max concurrent sessions exceeded
|
||||||
|
"""
|
||||||
|
if session_id in self.sessions:
|
||||||
|
raise SessionError(f"Session {session_id} already exists")
|
||||||
|
|
||||||
|
# Check max concurrent limit
|
||||||
|
self._enforce_max_concurrent()
|
||||||
|
|
||||||
|
# Start kernel
|
||||||
|
kernel_id = self.kernel_manager.start_kernel(session_id, volumes=volumes)
|
||||||
|
|
||||||
|
# Create session
|
||||||
|
now = datetime.utcnow()
|
||||||
|
session = Session(
|
||||||
|
session_id=session_id,
|
||||||
|
kernel_id=kernel_id,
|
||||||
|
created_at=now,
|
||||||
|
resource_limits=resource_limits
|
||||||
|
)
|
||||||
|
|
||||||
|
self.sessions[session_id] = session
|
||||||
|
|
||||||
|
# Log session creation
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.SESSION_CREATE,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message=f"Session created: {session_id}",
|
||||||
|
session_id=session_id,
|
||||||
|
details={
|
||||||
|
"kernel_id": kernel_id,
|
||||||
|
"memory": resource_limits.memory_bytes,
|
||||||
|
"cpu_quota": resource_limits.cpu_quota
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
return session
|
||||||
|
|
||||||
|
def session_exists(self, session_id: str) -> bool:
|
||||||
|
"""
|
||||||
|
Check if session exists.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session identifier
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
True if session exists, False otherwise
|
||||||
|
"""
|
||||||
|
return session_id in self.sessions
|
||||||
|
|
||||||
|
def get_session(self, session_id: str) -> Session:
|
||||||
|
"""
|
||||||
|
Get session by ID.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session identifier
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Session object
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If session doesn't exist
|
||||||
|
"""
|
||||||
|
if session_id not in self.sessions:
|
||||||
|
raise SessionError(f"Session {session_id} not found")
|
||||||
|
|
||||||
|
return self.sessions[session_id]
|
||||||
|
|
||||||
|
async def document_variables(
|
||||||
|
self,
|
||||||
|
session_id: str,
|
||||||
|
variables: Dict[str, str],
|
||||||
|
note: Optional[str] = None,
|
||||||
|
clear: bool = False
|
||||||
|
) -> Dict:
|
||||||
|
"""
|
||||||
|
Document important variables in a session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session identifier
|
||||||
|
variables: Dict mapping variable names to descriptions
|
||||||
|
note: Optional general note about session state
|
||||||
|
clear: Whether to clear existing documentation first
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Result dict with success status and documented count
|
||||||
|
"""
|
||||||
|
session = self.get_session(session_id)
|
||||||
|
|
||||||
|
if clear:
|
||||||
|
session.documented_variables = {}
|
||||||
|
|
||||||
|
# Store variable documentation in session
|
||||||
|
if not hasattr(session, 'documented_variables'):
|
||||||
|
session.documented_variables = {}
|
||||||
|
|
||||||
|
session.documented_variables.update(variables)
|
||||||
|
|
||||||
|
if note:
|
||||||
|
session.documentation_note = note
|
||||||
|
|
||||||
|
return {
|
||||||
|
"success": True,
|
||||||
|
"documented_count": len(variables),
|
||||||
|
"total_documented": len(session.documented_variables)
|
||||||
|
}
|
||||||
|
|
||||||
|
def execute_in_session(
|
||||||
|
self,
|
||||||
|
session_id: str,
|
||||||
|
code: str,
|
||||||
|
timeout: int = 300
|
||||||
|
) -> ExecutionResult:
|
||||||
|
"""
|
||||||
|
Execute code in session kernel.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session to execute in
|
||||||
|
code: Python code to execute
|
||||||
|
timeout: Maximum execution time
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
ExecutionResult with output
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If session doesn't exist
|
||||||
|
"""
|
||||||
|
session = self.get_session(session_id)
|
||||||
|
|
||||||
|
# Update activity
|
||||||
|
session.update_activity()
|
||||||
|
|
||||||
|
# Execute in kernel
|
||||||
|
result = self.kernel_manager.execute_code(
|
||||||
|
session.kernel_id,
|
||||||
|
code,
|
||||||
|
timeout=timeout
|
||||||
|
)
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
|
def document_state(
|
||||||
|
self,
|
||||||
|
session_id: str,
|
||||||
|
variables: Dict[str, str],
|
||||||
|
note: str = "",
|
||||||
|
clear: bool = False
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Document important variables in session.
|
||||||
|
|
||||||
|
Updates session.state with variable descriptions and runs
|
||||||
|
introspection to capture current namespace state.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session to document
|
||||||
|
variables: Dictionary of variable_name -> description
|
||||||
|
note: Optional note about session state
|
||||||
|
clear: If True, replace all documented variables; if False, merge
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If session doesn't exist
|
||||||
|
"""
|
||||||
|
session = self.get_session(session_id)
|
||||||
|
|
||||||
|
# Update documented variables
|
||||||
|
if clear:
|
||||||
|
session.state.documented_variables = variables.copy()
|
||||||
|
else:
|
||||||
|
session.state.documented_variables.update(variables)
|
||||||
|
|
||||||
|
# Update note if provided
|
||||||
|
if note:
|
||||||
|
session.state.note = note
|
||||||
|
|
||||||
|
# Run introspection to get current namespace state
|
||||||
|
session.state.all_variables = self.kernel_manager.inspect_namespace(session.kernel_id)
|
||||||
|
|
||||||
|
# Get variable info for documented variables
|
||||||
|
session.state.introspection = {}
|
||||||
|
for var_name in variables.keys():
|
||||||
|
if var_name in session.state.all_variables:
|
||||||
|
try:
|
||||||
|
info = self.kernel_manager.get_variable_info(session.kernel_id, var_name)
|
||||||
|
session.state.introspection[var_name] = info
|
||||||
|
except Exception:
|
||||||
|
# Variable might not exist yet
|
||||||
|
pass
|
||||||
|
|
||||||
|
# Update timestamp
|
||||||
|
session.state.last_updated = datetime.utcnow()
|
||||||
|
session.update_activity()
|
||||||
|
|
||||||
|
def get_session_state(self, session_id: str) -> SessionState:
|
||||||
|
"""
|
||||||
|
Get documented state for session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session identifier
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
SessionState object
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If session doesn't exist
|
||||||
|
"""
|
||||||
|
session = self.get_session(session_id)
|
||||||
|
return session.state
|
||||||
|
|
||||||
|
def destroy_session(self, session_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Destroy session and cleanup kernel.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session to destroy
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If session doesn't exist
|
||||||
|
"""
|
||||||
|
session = self.get_session(session_id)
|
||||||
|
|
||||||
|
# Shutdown kernel
|
||||||
|
try:
|
||||||
|
self.kernel_manager.shutdown_kernel(session.kernel_id)
|
||||||
|
except Exception as e:
|
||||||
|
# Log but continue with cleanup
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.SESSION_DESTROY,
|
||||||
|
severity=AuditSeverity.WARNING,
|
||||||
|
message=f"Error shutting down kernel for session {session_id}",
|
||||||
|
session_id=session_id,
|
||||||
|
error=str(e)
|
||||||
|
)
|
||||||
|
|
||||||
|
# Remove session
|
||||||
|
del self.sessions[session_id]
|
||||||
|
|
||||||
|
# Log destruction
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.SESSION_DESTROY,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message=f"Session destroyed: {session_id}",
|
||||||
|
session_id=session_id
|
||||||
|
)
|
||||||
|
|
||||||
|
def cleanup_idle_sessions(self) -> int:
|
||||||
|
"""
|
||||||
|
Cleanup sessions idle beyond configured timeout.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Number of sessions cleaned up
|
||||||
|
"""
|
||||||
|
timeout = timedelta(seconds=self.config.idle_timeout)
|
||||||
|
sessions_to_remove = []
|
||||||
|
|
||||||
|
for session_id, session in self.sessions.items():
|
||||||
|
if session.is_idle(timeout):
|
||||||
|
sessions_to_remove.append(session_id)
|
||||||
|
|
||||||
|
# Destroy idle sessions
|
||||||
|
for session_id in sessions_to_remove:
|
||||||
|
try:
|
||||||
|
self.destroy_session(session_id)
|
||||||
|
except Exception:
|
||||||
|
# Best effort cleanup
|
||||||
|
pass
|
||||||
|
|
||||||
|
return len(sessions_to_remove)
|
||||||
|
|
||||||
|
def list_sessions(self) -> List[dict]:
|
||||||
|
"""
|
||||||
|
List all active sessions with metadata.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of session dictionaries
|
||||||
|
"""
|
||||||
|
return [session.to_dict() for session in self.sessions.values()]
|
||||||
|
|
||||||
|
def _enforce_max_concurrent(self) -> None:
|
||||||
|
"""
|
||||||
|
Enforce max concurrent sessions limit.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SessionError: If at max concurrent sessions
|
||||||
|
"""
|
||||||
|
if len(self.sessions) >= self.config.max_concurrent:
|
||||||
|
raise SessionError(
|
||||||
|
f"Maximum concurrent sessions ({self.config.max_concurrent}) reached"
|
||||||
|
)
|
||||||
6
src/mcp_forge/execution/simple/__init__.py
Normal file
6
src/mcp_forge/execution/simple/__init__.py
Normal file
|
|
@ -0,0 +1,6 @@
|
||||||
|
"""Simple (stateless) execution backend."""
|
||||||
|
|
||||||
|
from .executor import CodeExecutor, ExecutionResult
|
||||||
|
from .backend import SimpleBackend
|
||||||
|
|
||||||
|
__all__ = ["CodeExecutor", "ExecutionResult", "SimpleBackend"]
|
||||||
184
src/mcp_forge/execution/simple/backend.py
Normal file
184
src/mcp_forge/execution/simple/backend.py
Normal file
|
|
@ -0,0 +1,184 @@
|
||||||
|
"""Simple (stateless) execution backend."""
|
||||||
|
|
||||||
|
from typing import Optional, Dict
|
||||||
|
import hashlib
|
||||||
|
|
||||||
|
from mcp_forge.config.schema import ForgeConfig
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits, parse_memory_string
|
||||||
|
from mcp_forge.execution.simple.executor import CodeExecutor, ExecutionResult
|
||||||
|
|
||||||
|
|
||||||
|
class SimpleBackend:
|
||||||
|
"""Stateless code execution backend."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
config: ForgeConfig,
|
||||||
|
container_manager: SecureContainerManager,
|
||||||
|
audit_logger: AuditLogger
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize simple backend.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config: Forge configuration
|
||||||
|
container_manager: Container lifecycle manager
|
||||||
|
audit_logger: Audit logging instance
|
||||||
|
"""
|
||||||
|
self.config = config
|
||||||
|
self.container_manager = container_manager
|
||||||
|
self.audit_logger = audit_logger
|
||||||
|
|
||||||
|
def execute(
|
||||||
|
self,
|
||||||
|
code: str,
|
||||||
|
timeout: Optional[int] = None,
|
||||||
|
memory: Optional[str] = None,
|
||||||
|
cpu_quota: Optional[int] = None,
|
||||||
|
custom_image: Optional[str] = None,
|
||||||
|
volumes: Optional[Dict[str, dict]] = None,
|
||||||
|
injection_code: Optional[str] = None,
|
||||||
|
bridge_socket_path: Optional[str] = None
|
||||||
|
) -> ExecutionResult:
|
||||||
|
"""
|
||||||
|
Execute Python code in stateless container.
|
||||||
|
|
||||||
|
Each execution creates a fresh container with no persistent state.
|
||||||
|
Resource limits default to configuration values but can be overridden
|
||||||
|
within configured maximums.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
code: Python code to execute
|
||||||
|
timeout: Max execution time in seconds (uses config default if None)
|
||||||
|
memory: Memory limit string (uses config default if None)
|
||||||
|
cpu_quota: CPU quota (uses config default if None)
|
||||||
|
custom_image: Custom image name (uses config default if None)
|
||||||
|
volumes: Volume mounts dict
|
||||||
|
injection_code: Optional MCP tool injection code to prepend
|
||||||
|
bridge_socket_path: Optional path to MCP bridge socket for mounting
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
ExecutionResult with execution output and metadata
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If limits exceed configured maximums
|
||||||
|
"""
|
||||||
|
# Use defaults from config if not specified
|
||||||
|
timeout = timeout if timeout is not None else self.config.execution.default_timeout
|
||||||
|
memory = memory if memory is not None else self.config.execution.default_memory
|
||||||
|
cpu_quota = cpu_quota if cpu_quota is not None else self.config.execution.default_cpu_quota
|
||||||
|
|
||||||
|
# Validate limits against maximums
|
||||||
|
self._validate_limits(timeout, memory, cpu_quota)
|
||||||
|
|
||||||
|
# Get image
|
||||||
|
image = self._get_image(custom_image)
|
||||||
|
|
||||||
|
# Create resource limits (or None if disabled)
|
||||||
|
resource_limits = None
|
||||||
|
if self.config.security.enforce_resource_limits:
|
||||||
|
resource_limits = ResourceLimits(
|
||||||
|
memory=memory,
|
||||||
|
cpu_quota=cpu_quota,
|
||||||
|
storage="1g", # Default storage quota
|
||||||
|
timeout=timeout
|
||||||
|
)
|
||||||
|
|
||||||
|
# Log execution (hash code, don't log actual content)
|
||||||
|
code_hash = hashlib.sha256(code.encode()).hexdigest()
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.EXECUTION_REQUEST,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message="Code execution requested",
|
||||||
|
details={
|
||||||
|
"code_hash": code_hash,
|
||||||
|
"image": image,
|
||||||
|
"timeout": timeout,
|
||||||
|
"memory": memory,
|
||||||
|
"cpu_quota": cpu_quota
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create executor and execute
|
||||||
|
executor = CodeExecutor(
|
||||||
|
container_manager=self.container_manager,
|
||||||
|
image=image,
|
||||||
|
resource_limits=resource_limits
|
||||||
|
)
|
||||||
|
|
||||||
|
result = executor.execute(
|
||||||
|
code,
|
||||||
|
timeout=timeout,
|
||||||
|
injection_code=injection_code,
|
||||||
|
bridge_socket_path=bridge_socket_path
|
||||||
|
)
|
||||||
|
|
||||||
|
# Log result
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.EXECUTION_REQUEST,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message="Code execution completed",
|
||||||
|
details={
|
||||||
|
"code_hash": code_hash,
|
||||||
|
"success": result.success,
|
||||||
|
"execution_time": result.execution_time,
|
||||||
|
"exit_code": result.exit_code
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
|
def _validate_limits(
|
||||||
|
self,
|
||||||
|
timeout: int,
|
||||||
|
memory: str,
|
||||||
|
cpu_quota: int
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Validate resource limits against configuration maximums.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
timeout: Timeout in seconds
|
||||||
|
memory: Memory limit string
|
||||||
|
cpu_quota: CPU quota value
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If any limit exceeds maximum
|
||||||
|
"""
|
||||||
|
# Validate timeout
|
||||||
|
if timeout > self.config.execution.max_timeout:
|
||||||
|
raise ValueError(
|
||||||
|
f"timeout {timeout}s exceeds maximum {self.config.execution.max_timeout}s"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Validate memory
|
||||||
|
memory_bytes = parse_memory_string(memory)
|
||||||
|
max_memory_bytes = parse_memory_string(self.config.execution.max_memory)
|
||||||
|
if memory_bytes > max_memory_bytes:
|
||||||
|
raise ValueError(
|
||||||
|
f"memory {memory} exceeds maximum {self.config.execution.max_memory}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Validate CPU quota
|
||||||
|
if cpu_quota > self.config.execution.max_cpu_quota:
|
||||||
|
raise ValueError(
|
||||||
|
f"cpu_quota {cpu_quota} exceeds maximum {self.config.execution.max_cpu_quota}"
|
||||||
|
)
|
||||||
|
|
||||||
|
def _get_image(self, custom_image: Optional[str]) -> str:
|
||||||
|
"""
|
||||||
|
Get image name, defaulting to configured image.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
custom_image: Optional custom image name
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Image name to use
|
||||||
|
"""
|
||||||
|
if custom_image is not None:
|
||||||
|
return custom_image
|
||||||
|
|
||||||
|
# Default to Python 3.11
|
||||||
|
return self.config.images.python_3_11
|
||||||
222
src/mcp_forge/execution/simple/executor.py
Normal file
222
src/mcp_forge/execution/simple/executor.py
Normal file
|
|
@ -0,0 +1,222 @@
|
||||||
|
"""Code execution in isolated containers."""
|
||||||
|
|
||||||
|
from typing import Any, Optional
|
||||||
|
from dataclasses import dataclass, asdict
|
||||||
|
import json
|
||||||
|
import time
|
||||||
|
import textwrap
|
||||||
|
import uuid
|
||||||
|
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager, ContainerConfig
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class ExecutionResult:
|
||||||
|
"""Result of code execution."""
|
||||||
|
success: bool
|
||||||
|
stdout: str
|
||||||
|
stderr: str
|
||||||
|
result: Optional[Any]
|
||||||
|
execution_time: float
|
||||||
|
exit_code: int
|
||||||
|
error: Optional[str] = None
|
||||||
|
|
||||||
|
def to_dict(self) -> dict:
|
||||||
|
"""Convert to dictionary for JSON serialization."""
|
||||||
|
return asdict(self)
|
||||||
|
|
||||||
|
|
||||||
|
class CodeExecutor:
|
||||||
|
"""Executes Python code in isolated containers."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
container_manager: SecureContainerManager,
|
||||||
|
image: str,
|
||||||
|
resource_limits: ResourceLimits
|
||||||
|
):
|
||||||
|
self.container_manager = container_manager
|
||||||
|
self.image = image
|
||||||
|
self.resource_limits = resource_limits
|
||||||
|
|
||||||
|
def execute(
|
||||||
|
self,
|
||||||
|
code: str,
|
||||||
|
timeout: Optional[int] = None,
|
||||||
|
injection_code: Optional[str] = None,
|
||||||
|
bridge_socket_path: Optional[str] = None
|
||||||
|
) -> ExecutionResult:
|
||||||
|
"""
|
||||||
|
Execute Python code in a fresh container.
|
||||||
|
|
||||||
|
Process:
|
||||||
|
1. Create container with code
|
||||||
|
2. Start container
|
||||||
|
3. Wait for completion (with timeout)
|
||||||
|
4. Capture stdout/stderr
|
||||||
|
5. Extract result from last expression
|
||||||
|
6. Cleanup container
|
||||||
|
|
||||||
|
Args:
|
||||||
|
code: Python code to execute
|
||||||
|
timeout: Maximum execution time in seconds
|
||||||
|
injection_code: Optional MCP tool injection code to prepend
|
||||||
|
bridge_socket_path: Optional path to MCP bridge socket for mounting
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
ExecutionResult with stdout, stderr, result, and timing
|
||||||
|
"""
|
||||||
|
start_time = time.time()
|
||||||
|
container_id = None
|
||||||
|
|
||||||
|
# Use provided timeout or default from resource limits
|
||||||
|
exec_timeout = timeout if timeout is not None else self.resource_limits.timeout
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Prepare code wrapper (with injection if provided)
|
||||||
|
wrapped_code = self._prepare_code(code, injection_code=injection_code)
|
||||||
|
|
||||||
|
# Generate a session ID for this execution to register the container
|
||||||
|
import uuid
|
||||||
|
session_id = f"simple-exec-{uuid.uuid4().hex[:12]}"
|
||||||
|
|
||||||
|
# Set up volumes for MCP bridge socket if provided
|
||||||
|
volumes = {}
|
||||||
|
if bridge_socket_path:
|
||||||
|
volumes[bridge_socket_path] = {
|
||||||
|
"bind": bridge_socket_path,
|
||||||
|
"mode": "rw"
|
||||||
|
}
|
||||||
|
|
||||||
|
# Create container configuration
|
||||||
|
config = ContainerConfig(
|
||||||
|
image=self.image,
|
||||||
|
command=["python", "-c", wrapped_code],
|
||||||
|
resource_limits=self.resource_limits,
|
||||||
|
volumes=volumes if volumes else None
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create and start container with session_id for proper registration
|
||||||
|
container_id = self.container_manager.create_container(config, session_id=session_id)
|
||||||
|
self.container_manager.start_container(container_id)
|
||||||
|
|
||||||
|
# Wait for completion
|
||||||
|
exit_code = self.container_manager.wait_for_container(
|
||||||
|
container_id,
|
||||||
|
timeout=exec_timeout
|
||||||
|
)
|
||||||
|
|
||||||
|
# Get logs
|
||||||
|
stdout, stderr = self.container_manager.get_container_logs(container_id)
|
||||||
|
|
||||||
|
# Parse output to extract result
|
||||||
|
result, error = self._parse_output(stdout)
|
||||||
|
|
||||||
|
execution_time = time.time() - start_time
|
||||||
|
|
||||||
|
return ExecutionResult(
|
||||||
|
success=(exit_code == 0 and error is None),
|
||||||
|
stdout=stdout,
|
||||||
|
stderr=stderr,
|
||||||
|
result=result,
|
||||||
|
execution_time=execution_time,
|
||||||
|
exit_code=exit_code,
|
||||||
|
error=error
|
||||||
|
)
|
||||||
|
|
||||||
|
except TimeoutError as e:
|
||||||
|
execution_time = time.time() - start_time
|
||||||
|
return ExecutionResult(
|
||||||
|
success=False,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=None,
|
||||||
|
execution_time=execution_time,
|
||||||
|
exit_code=-1,
|
||||||
|
error=f"Execution timeout: {str(e)}"
|
||||||
|
)
|
||||||
|
|
||||||
|
finally:
|
||||||
|
# Cleanup container
|
||||||
|
if container_id is not None:
|
||||||
|
try:
|
||||||
|
self.container_manager.remove_container(container_id)
|
||||||
|
except Exception:
|
||||||
|
pass # Best effort cleanup
|
||||||
|
|
||||||
|
def _prepare_code(self, code: str, injection_code: Optional[str] = None) -> str:
|
||||||
|
"""
|
||||||
|
Wrap code to capture result and handle errors.
|
||||||
|
|
||||||
|
Wraps code in try/except and captures:
|
||||||
|
- Last expression result
|
||||||
|
- Exceptions with traceback
|
||||||
|
- Execution metadata
|
||||||
|
|
||||||
|
Args:
|
||||||
|
code: User code to execute
|
||||||
|
injection_code: Optional MCP tool injection code to prepend
|
||||||
|
|
||||||
|
Returns wrapped code that outputs JSON to stdout.
|
||||||
|
"""
|
||||||
|
# Prepend injection code if provided
|
||||||
|
if injection_code:
|
||||||
|
full_code = injection_code + "\n\n" + code
|
||||||
|
else:
|
||||||
|
full_code = code
|
||||||
|
|
||||||
|
# Escape the code for embedding in exec string
|
||||||
|
escaped_code = full_code.replace('\\', '\\\\').replace('"', '\\"').replace('\n', '\\n')
|
||||||
|
|
||||||
|
wrapper_template = '''
|
||||||
|
import sys
|
||||||
|
import json
|
||||||
|
import traceback
|
||||||
|
|
||||||
|
def __mcp_execute():
|
||||||
|
result = None
|
||||||
|
error = None
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Execute user code
|
||||||
|
exec_globals = {}
|
||||||
|
exec("""%s""", exec_globals)
|
||||||
|
|
||||||
|
# Try to get result from last expression
|
||||||
|
result = exec_globals.get('_', None)
|
||||||
|
|
||||||
|
except SyntaxError as e:
|
||||||
|
error = f"SyntaxError: {e.msg} (line {e.lineno})"
|
||||||
|
except Exception as e:
|
||||||
|
error = f"{type(e).__name__}: {str(e)}"
|
||||||
|
|
||||||
|
# Output result as JSON
|
||||||
|
print(json.dumps({"result": result, "error": error}))
|
||||||
|
|
||||||
|
__mcp_execute()
|
||||||
|
'''
|
||||||
|
|
||||||
|
return wrapper_template % escaped_code
|
||||||
|
|
||||||
|
def _parse_output(self, stdout: str) -> tuple[Optional[Any], Optional[str]]:
|
||||||
|
"""
|
||||||
|
Parse execution output to extract result and error.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
(result, error_message)
|
||||||
|
"""
|
||||||
|
if not stdout:
|
||||||
|
return None, None
|
||||||
|
|
||||||
|
try:
|
||||||
|
# First line should be JSON output
|
||||||
|
lines = stdout.split('\n', 1)
|
||||||
|
json_line = lines[0]
|
||||||
|
|
||||||
|
data = json.loads(json_line)
|
||||||
|
return data.get("result"), data.get("error")
|
||||||
|
|
||||||
|
except (json.JSONDecodeError, ValueError):
|
||||||
|
# If can't parse JSON, treat entire output as result
|
||||||
|
return None, None
|
||||||
8
src/mcp_forge/mcp/__init__.py
Normal file
8
src/mcp_forge/mcp/__init__.py
Normal file
|
|
@ -0,0 +1,8 @@
|
||||||
|
"""MCP client integration."""
|
||||||
|
|
||||||
|
from mcp_forge.mcp.client import MCPClientWrapper
|
||||||
|
from mcp_forge.mcp.manager import MCPClientManager
|
||||||
|
from mcp_forge.mcp.bridge import ToolBridgeServer
|
||||||
|
from mcp_forge.mcp.injection import ToolInjectionGenerator
|
||||||
|
|
||||||
|
__all__ = ["MCPClientWrapper", "MCPClientManager", "ToolBridgeServer", "ToolInjectionGenerator"]
|
||||||
261
src/mcp_forge/mcp/bridge.py
Normal file
261
src/mcp_forge/mcp/bridge.py
Normal file
|
|
@ -0,0 +1,261 @@
|
||||||
|
"""Tool Bridge Server for forwarding MCP tool calls from containers."""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
|
from typing import Optional
|
||||||
|
from pathlib import Path
|
||||||
|
from datetime import datetime
|
||||||
|
|
||||||
|
from ..security.audit import AuditEventType, AuditSeverity
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class ToolBridgeServer:
|
||||||
|
"""
|
||||||
|
Async Unix socket server for MCP tool calls from containers.
|
||||||
|
|
||||||
|
Containers connect to this socket to call MCP tools.
|
||||||
|
Bridge forwards calls to actual MCP clients and returns results.
|
||||||
|
Runs as an asyncio task in the main event loop.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
socket_path: Path,
|
||||||
|
client_manager, # MCPClientManager
|
||||||
|
audit_logger # AuditLogger
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize bridge server.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
socket_path: Path to Unix socket
|
||||||
|
client_manager: MCPClientManager instance
|
||||||
|
audit_logger: AuditLogger instance
|
||||||
|
"""
|
||||||
|
self.socket_path = Path(socket_path)
|
||||||
|
self.client_manager = client_manager
|
||||||
|
self.audit_logger = audit_logger
|
||||||
|
self.server: Optional[asyncio.Server] = None
|
||||||
|
self.running = False
|
||||||
|
self.server_task: Optional[asyncio.Task] = None
|
||||||
|
|
||||||
|
async def start(self) -> None:
|
||||||
|
"""
|
||||||
|
Start async bridge server.
|
||||||
|
|
||||||
|
Creates Unix socket and starts listening for connections.
|
||||||
|
Runs in the current event loop as a background task.
|
||||||
|
"""
|
||||||
|
if self.running:
|
||||||
|
logger.info("Bridge server already running")
|
||||||
|
return
|
||||||
|
|
||||||
|
logger.info("Starting async bridge server...")
|
||||||
|
|
||||||
|
# Remove existing socket if it exists
|
||||||
|
if self.socket_path.exists():
|
||||||
|
self.socket_path.unlink()
|
||||||
|
|
||||||
|
# Create async Unix socket server
|
||||||
|
self.server = await asyncio.start_unix_server(
|
||||||
|
self._handle_client,
|
||||||
|
path=str(self.socket_path)
|
||||||
|
)
|
||||||
|
|
||||||
|
# Make socket accessible to containers (chmod 666)
|
||||||
|
import os
|
||||||
|
os.chmod(str(self.socket_path), 0o666)
|
||||||
|
|
||||||
|
self.running = True
|
||||||
|
|
||||||
|
# Start serving in background task
|
||||||
|
self.server_task = asyncio.create_task(self._serve_forever())
|
||||||
|
|
||||||
|
logger.info(f"Bridge server listening on {self.socket_path}")
|
||||||
|
|
||||||
|
async def _serve_forever(self) -> None:
|
||||||
|
"""Keep server running until stopped."""
|
||||||
|
try:
|
||||||
|
async with self.server:
|
||||||
|
await self.server.serve_forever()
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
logger.info("Bridge server task cancelled")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Bridge server error: {e}")
|
||||||
|
|
||||||
|
def is_running(self) -> bool:
|
||||||
|
"""Check if bridge server is running."""
|
||||||
|
return self.running
|
||||||
|
|
||||||
|
async def stop(self) -> None:
|
||||||
|
"""Stop bridge server and cleanup socket."""
|
||||||
|
if not self.running:
|
||||||
|
return
|
||||||
|
|
||||||
|
logger.info("Stopping bridge server...")
|
||||||
|
self.running = False
|
||||||
|
|
||||||
|
# Cancel server task
|
||||||
|
if self.server_task:
|
||||||
|
self.server_task.cancel()
|
||||||
|
try:
|
||||||
|
await self.server_task
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
# Close server
|
||||||
|
if self.server:
|
||||||
|
self.server.close()
|
||||||
|
await self.server.wait_closed()
|
||||||
|
|
||||||
|
# Remove socket file
|
||||||
|
if self.socket_path.exists():
|
||||||
|
self.socket_path.unlink()
|
||||||
|
|
||||||
|
logger.info("Bridge server stopped")
|
||||||
|
|
||||||
|
async def _handle_client(
|
||||||
|
self,
|
||||||
|
reader: asyncio.StreamReader,
|
||||||
|
writer: asyncio.StreamWriter
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Handle single client connection (async).
|
||||||
|
|
||||||
|
Protocol:
|
||||||
|
1. Receive JSON: {"tool": "name", "params": {...}}
|
||||||
|
2. Validate request
|
||||||
|
3. Call tool via client manager
|
||||||
|
4. Send JSON response: {"success": true, "result": ...}
|
||||||
|
"""
|
||||||
|
addr = writer.get_extra_info('peername', 'unknown')
|
||||||
|
logger.info(f"[Bridge] Client connected from {addr}")
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Receive request data
|
||||||
|
logger.info("[Bridge] Reading request...")
|
||||||
|
request_data = await reader.read()
|
||||||
|
logger.info(f"[Bridge] Received {len(request_data)} bytes")
|
||||||
|
|
||||||
|
# Parse JSON request
|
||||||
|
try:
|
||||||
|
request = json.loads(request_data.decode('utf-8'))
|
||||||
|
logger.info(f"[Bridge] Parsed request: tool={request.get('tool')}, params={list(request.get('params', {}).keys())}")
|
||||||
|
except json.JSONDecodeError as e:
|
||||||
|
logger.error(f"[Bridge] JSON decode error: {e}")
|
||||||
|
await self._send_error(writer, f"Invalid JSON: {e}")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Validate request format
|
||||||
|
if "tool" not in request:
|
||||||
|
logger.warning("[Bridge] Missing 'tool' field in request")
|
||||||
|
await self._send_error(writer, "Missing 'tool' field in request")
|
||||||
|
return
|
||||||
|
|
||||||
|
tool_name = request["tool"]
|
||||||
|
params = request.get("params", {})
|
||||||
|
|
||||||
|
logger.info(f"[Bridge] Calling tool: {tool_name}")
|
||||||
|
|
||||||
|
# Log tool call (tool name only, not params)
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.MCP_TOOL_CALL,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message=f"Tool call: {tool_name}",
|
||||||
|
details={"tool": tool_name, "timestamp": datetime.utcnow().isoformat()}
|
||||||
|
)
|
||||||
|
|
||||||
|
# Call tool asynchronously
|
||||||
|
logger.info(f"[Bridge] About to call client_manager.call_tool for {tool_name}")
|
||||||
|
try:
|
||||||
|
logger.info(f"[Bridge] Calling tool {tool_name}...")
|
||||||
|
result = await self.client_manager.call_tool(tool_name, params)
|
||||||
|
logger.info(f"[Bridge] Tool {tool_name} returned: type={type(result)}")
|
||||||
|
|
||||||
|
logger.info("[Bridge] Sending result as-is")
|
||||||
|
|
||||||
|
# Send success response
|
||||||
|
# Convert result to JSON-serializable format
|
||||||
|
result = self._make_serializable(result)
|
||||||
|
response = {
|
||||||
|
"success": True,
|
||||||
|
"result": result
|
||||||
|
}
|
||||||
|
response_json = json.dumps(response)
|
||||||
|
|
||||||
|
logger.info(f"[Bridge] Sending response ({len(response_json)} bytes)")
|
||||||
|
writer.write(response_json.encode('utf-8'))
|
||||||
|
await writer.drain()
|
||||||
|
logger.info(f"[Bridge] Response sent for {tool_name}")
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
import traceback
|
||||||
|
error_details = f"{type(e).__name__}: {str(e)}\n{traceback.format_exc()}"
|
||||||
|
logger.error(f"[Bridge] Tool execution failed: {error_details}")
|
||||||
|
await self._send_error(writer, f"Tool execution failed: {type(e).__name__}: {str(e)}")
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
import traceback
|
||||||
|
logger.error(f"[Bridge] Request handling failed: {e}\n{traceback.format_exc()}")
|
||||||
|
await self._send_error(writer, f"Request handling failed: {e}")
|
||||||
|
finally:
|
||||||
|
logger.info("[Bridge] Closing connection")
|
||||||
|
writer.close()
|
||||||
|
await writer.wait_closed()
|
||||||
|
|
||||||
|
def _make_serializable(self, obj):
|
||||||
|
"""
|
||||||
|
Recursively convert objects to JSON-serializable format.
|
||||||
|
Handles Pydantic models, lists, dicts, and other objects.
|
||||||
|
"""
|
||||||
|
# Handle None, primitives
|
||||||
|
if obj is None or isinstance(obj, (str, int, float, bool)):
|
||||||
|
return obj
|
||||||
|
|
||||||
|
# Handle lists
|
||||||
|
if isinstance(obj, list):
|
||||||
|
return [self._make_serializable(item) for item in obj]
|
||||||
|
|
||||||
|
# Handle dicts
|
||||||
|
if isinstance(obj, dict):
|
||||||
|
return {key: self._make_serializable(value) for key, value in obj.items()}
|
||||||
|
|
||||||
|
# Handle Pydantic models (v2)
|
||||||
|
if hasattr(obj, 'model_dump'):
|
||||||
|
try:
|
||||||
|
dumped = obj.model_dump()
|
||||||
|
logger.debug(f"[Bridge] model_dump() returned: {dumped}")
|
||||||
|
return self._make_serializable(dumped)
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"[Bridge] model_dump() failed: {e}, trying dict()")
|
||||||
|
|
||||||
|
# Handle Pydantic models (v1)
|
||||||
|
if hasattr(obj, 'dict'):
|
||||||
|
try:
|
||||||
|
return self._make_serializable(obj.dict())
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"[Bridge] dict() failed: {e}")
|
||||||
|
|
||||||
|
# Handle objects with __dict__
|
||||||
|
if hasattr(obj, '__dict__'):
|
||||||
|
obj_dict = {k: v for k, v in obj.__dict__.items() if not k.startswith('_')}
|
||||||
|
return self._make_serializable(obj_dict)
|
||||||
|
|
||||||
|
# Fallback: convert to string
|
||||||
|
logger.warning(f"[Bridge] Falling back to str() for {type(obj)}")
|
||||||
|
return str(obj)
|
||||||
|
|
||||||
|
async def _send_error(self, writer: asyncio.StreamWriter, error: str) -> None:
|
||||||
|
"""Send error response to client."""
|
||||||
|
try:
|
||||||
|
response = {
|
||||||
|
"success": False,
|
||||||
|
"error": error
|
||||||
|
}
|
||||||
|
writer.write(json.dumps(response).encode('utf-8'))
|
||||||
|
await writer.drain()
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"[Bridge] Failed to send error: {e}") # Connection may already be closed
|
||||||
220
src/mcp_forge/mcp/client.py
Normal file
220
src/mcp_forge/mcp/client.py
Normal file
|
|
@ -0,0 +1,220 @@
|
||||||
|
"""MCP client wrapper using fastmcp."""
|
||||||
|
|
||||||
|
from fastmcp import Client
|
||||||
|
from fastmcp.client.transports import StdioTransport, SSETransport, StreamableHttpTransport
|
||||||
|
from typing import Dict, List, Optional, Any, Literal
|
||||||
|
import logging
|
||||||
|
import json
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class MCPClientWrapper:
|
||||||
|
"""Wrapper for fastmcp Client connection."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
name: str,
|
||||||
|
transport_type: Literal["stdio", "http", "sse"] = "stdio",
|
||||||
|
# For stdio
|
||||||
|
command: Optional[str] = None,
|
||||||
|
args: Optional[List[str]] = None,
|
||||||
|
env: Optional[Dict[str, str]] = None,
|
||||||
|
# For http/sse
|
||||||
|
url: Optional[str] = None,
|
||||||
|
headers: Optional[Dict[str, str]] = None
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize MCP client wrapper.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
name: Human-readable name for this client
|
||||||
|
transport_type: Type of transport ("stdio", "http", or "sse")
|
||||||
|
command: Command to execute (for stdio)
|
||||||
|
args: Arguments for command (for stdio)
|
||||||
|
env: Environment variables (for stdio)
|
||||||
|
url: URL endpoint (for http/sse)
|
||||||
|
headers: HTTP headers (for http/sse)
|
||||||
|
"""
|
||||||
|
self.name = name
|
||||||
|
self.transport_type = transport_type
|
||||||
|
|
||||||
|
# Create appropriate transport
|
||||||
|
if transport_type == "stdio":
|
||||||
|
if not command:
|
||||||
|
raise ValueError(f"command required for stdio transport (client: {name})")
|
||||||
|
self.command = command
|
||||||
|
self.args = args or []
|
||||||
|
self.env = env or {}
|
||||||
|
self.transport = StdioTransport(
|
||||||
|
command=command,
|
||||||
|
args=self.args,
|
||||||
|
env=self.env
|
||||||
|
)
|
||||||
|
elif transport_type == "http":
|
||||||
|
if not url:
|
||||||
|
raise ValueError(f"url required for http transport (client: {name})")
|
||||||
|
self.url = url
|
||||||
|
self.headers = headers or {}
|
||||||
|
self.transport = StreamableHttpTransport(
|
||||||
|
url=url,
|
||||||
|
headers=self.headers
|
||||||
|
)
|
||||||
|
elif transport_type == "sse":
|
||||||
|
if not url:
|
||||||
|
raise ValueError(f"url required for sse transport (client: {name})")
|
||||||
|
self.url = url
|
||||||
|
self.headers = headers or {}
|
||||||
|
self.transport = SSETransport(
|
||||||
|
url=url,
|
||||||
|
headers=self.headers
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Unknown transport type: {transport_type}")
|
||||||
|
|
||||||
|
# Create fastmcp client
|
||||||
|
self._client = Client(self.transport)
|
||||||
|
self._connected = False
|
||||||
|
self._tools_cache: Optional[List[str]] = None
|
||||||
|
|
||||||
|
async def connect(self) -> None:
|
||||||
|
"""
|
||||||
|
Connect to MCP server.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: If connection fails
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
# Enter the client context manager
|
||||||
|
await self._client.__aenter__()
|
||||||
|
self._connected = True
|
||||||
|
# Clear tools cache
|
||||||
|
self._tools_cache = None
|
||||||
|
except Exception as e:
|
||||||
|
self._connected = False
|
||||||
|
raise RuntimeError(f"Failed to connect to MCP server '{self.name}': {e}") from e
|
||||||
|
|
||||||
|
async def disconnect(self) -> None:
|
||||||
|
"""Disconnect from MCP server."""
|
||||||
|
if self._connected:
|
||||||
|
try:
|
||||||
|
await self._client.__aexit__(None, None, None)
|
||||||
|
finally:
|
||||||
|
self._connected = False
|
||||||
|
self._tools_cache = None
|
||||||
|
|
||||||
|
def is_connected(self) -> bool:
|
||||||
|
"""Check if client is connected."""
|
||||||
|
return self._connected
|
||||||
|
|
||||||
|
async def list_tools(self) -> List[str]:
|
||||||
|
"""
|
||||||
|
Get list of available tool names.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of tool names
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: If not connected
|
||||||
|
"""
|
||||||
|
if not self._connected:
|
||||||
|
raise RuntimeError(f"Client '{self.name}' is not connected")
|
||||||
|
|
||||||
|
# Use cached tools if available
|
||||||
|
if self._tools_cache is not None:
|
||||||
|
return self._tools_cache
|
||||||
|
|
||||||
|
# Fetch and cache tools
|
||||||
|
tools_result = await self._client.list_tools()
|
||||||
|
# tools_result can be either a list or an object with .tools attribute
|
||||||
|
if isinstance(tools_result, list):
|
||||||
|
self._tools_cache = [tool.name for tool in tools_result]
|
||||||
|
else:
|
||||||
|
self._tools_cache = [tool.name for tool in tools_result.tools]
|
||||||
|
return self._tools_cache
|
||||||
|
|
||||||
|
async def get_tool_schema(self, tool_name: str) -> dict:
|
||||||
|
"""
|
||||||
|
Get JSON schema for tool.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
tool_name: Name of the tool
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
JSON schema dict for tool
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: If not connected
|
||||||
|
KeyError: If tool not found
|
||||||
|
"""
|
||||||
|
if not self._connected:
|
||||||
|
raise RuntimeError(f"Client '{self.name}' is not connected")
|
||||||
|
|
||||||
|
tools_result = await self._client.list_tools()
|
||||||
|
# Handle both list and object with .tools attribute
|
||||||
|
tools_list = tools_result if isinstance(tools_result, list) else tools_result.tools
|
||||||
|
|
||||||
|
for tool in tools_list:
|
||||||
|
if tool.name == tool_name:
|
||||||
|
return tool.inputSchema
|
||||||
|
|
||||||
|
raise KeyError(f"Tool '{tool_name}' not found in server '{self.name}'")
|
||||||
|
|
||||||
|
async def call_tool(
|
||||||
|
self,
|
||||||
|
tool_name: str,
|
||||||
|
arguments: Dict[str, Any]
|
||||||
|
) -> Any:
|
||||||
|
"""
|
||||||
|
Call MCP tool with arguments.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
tool_name: Name of the tool to call
|
||||||
|
arguments: Arguments dict for tool
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Tool result (automatically deserialized from .data)
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: If not connected or tool call fails
|
||||||
|
"""
|
||||||
|
if not self._connected:
|
||||||
|
raise RuntimeError(f"Client '{self.name}' is not connected")
|
||||||
|
|
||||||
|
try:
|
||||||
|
result = await self._client.call_tool(tool_name, arguments)
|
||||||
|
|
||||||
|
# Return structured data if available
|
||||||
|
if result.data is not None:
|
||||||
|
# Check if data is empty Root() objects (common with rag-mcp)
|
||||||
|
if isinstance(result.data, list) and result.data:
|
||||||
|
first = result.data[0]
|
||||||
|
type_name = type(first).__name__
|
||||||
|
|
||||||
|
# If data is empty Root objects, try content instead
|
||||||
|
if type_name == 'Root' and not first.__dict__:
|
||||||
|
logger.debug(f"[MCPClient] Tool {tool_name} returned empty Root objects, checking content")
|
||||||
|
if result.content:
|
||||||
|
for content in result.content:
|
||||||
|
if hasattr(content, 'text'):
|
||||||
|
# Parse JSON from text content
|
||||||
|
try:
|
||||||
|
return json.loads(content.text)
|
||||||
|
except json.JSONDecodeError:
|
||||||
|
# If not JSON, return as-is
|
||||||
|
return content.text
|
||||||
|
return None
|
||||||
|
|
||||||
|
# Return data as-is
|
||||||
|
return result.data
|
||||||
|
elif result.content:
|
||||||
|
# Fall back to text content
|
||||||
|
for content in result.content:
|
||||||
|
if hasattr(content, 'text'):
|
||||||
|
return content.text
|
||||||
|
return None
|
||||||
|
else:
|
||||||
|
return None
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
raise RuntimeError(f"Tool call failed for '{tool_name}': {e}") from e
|
||||||
177
src/mcp_forge/mcp/injection.py
Normal file
177
src/mcp_forge/mcp/injection.py
Normal file
|
|
@ -0,0 +1,177 @@
|
||||||
|
"""Tool Injection Generator for creating Python code to inject MCP tools."""
|
||||||
|
|
||||||
|
from typing import List, Tuple
|
||||||
|
|
||||||
|
|
||||||
|
class ToolInjectionGenerator:
|
||||||
|
"""Generates Python code to inject MCP tools into container namespace."""
|
||||||
|
|
||||||
|
def __init__(self, client_manager):
|
||||||
|
"""
|
||||||
|
Initialize generator.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
client_manager: MCPClientManager instance
|
||||||
|
"""
|
||||||
|
self.client_manager = client_manager
|
||||||
|
|
||||||
|
async def generate_injection_code(
|
||||||
|
self,
|
||||||
|
tool_names: List[str],
|
||||||
|
bridge_socket_path: str = "/tmp/mcp-bridge.sock"
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Generate Python code that provides MCP tools as functions.
|
||||||
|
|
||||||
|
Generated code includes:
|
||||||
|
1. Bridge client to communicate with MCP bridge server
|
||||||
|
2. Wrapper function for each tool with proper signature
|
||||||
|
3. Docstrings from tool schemas
|
||||||
|
4. Type hints from tool schemas
|
||||||
|
|
||||||
|
Args:
|
||||||
|
tool_names: List of MCP tool names to inject
|
||||||
|
bridge_socket_path: Path to bridge socket in container
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Python code as string
|
||||||
|
"""
|
||||||
|
parts = []
|
||||||
|
|
||||||
|
# Add imports
|
||||||
|
parts.append(self._generate_imports())
|
||||||
|
parts.append("")
|
||||||
|
|
||||||
|
# Add bridge client
|
||||||
|
parts.append(self._generate_bridge_client(bridge_socket_path))
|
||||||
|
parts.append("")
|
||||||
|
|
||||||
|
# Add tool functions
|
||||||
|
for tool_name in tool_names:
|
||||||
|
schema = await self.client_manager.get_tool_schema(tool_name)
|
||||||
|
function_code = await self._generate_tool_function(tool_name, schema)
|
||||||
|
parts.append(function_code)
|
||||||
|
parts.append("")
|
||||||
|
|
||||||
|
return "\n".join(parts)
|
||||||
|
|
||||||
|
def _generate_imports(self) -> str:
|
||||||
|
"""Generate import statements."""
|
||||||
|
return """import socket
|
||||||
|
import json
|
||||||
|
from typing import Any"""
|
||||||
|
|
||||||
|
def _generate_bridge_client(self, socket_path: str) -> str:
|
||||||
|
"""Generate code for bridge client communication."""
|
||||||
|
return f'''def _mcp_call(tool_name: str, **kwargs) -> Any:
|
||||||
|
"""Internal: Call MCP tool via bridge."""
|
||||||
|
s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||||
|
s.connect('{socket_path}')
|
||||||
|
request = json.dumps({{'tool': tool_name, 'params': kwargs}})
|
||||||
|
s.sendall(request.encode('utf-8'))
|
||||||
|
s.shutdown(socket.SHUT_WR) # Signal we're done sending
|
||||||
|
|
||||||
|
# Receive response
|
||||||
|
response_data = b''
|
||||||
|
while True:
|
||||||
|
chunk = s.recv(4096)
|
||||||
|
if not chunk:
|
||||||
|
break
|
||||||
|
response_data += chunk
|
||||||
|
|
||||||
|
s.close()
|
||||||
|
response = json.loads(response_data.decode('utf-8'))
|
||||||
|
|
||||||
|
if not response.get('success'):
|
||||||
|
raise RuntimeError(f"Tool call failed: {{response.get('error')}}")
|
||||||
|
|
||||||
|
return response.get('result')'''
|
||||||
|
|
||||||
|
async def _generate_tool_function(
|
||||||
|
self,
|
||||||
|
tool_name: str,
|
||||||
|
tool_schema: dict
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Generate wrapper function for a single tool.
|
||||||
|
|
||||||
|
Includes:
|
||||||
|
- Function signature from schema parameters
|
||||||
|
- Type hints
|
||||||
|
- Docstring with description and parameters
|
||||||
|
- Call to _mcp_call()
|
||||||
|
"""
|
||||||
|
# Extract parameters from schema
|
||||||
|
params = self._extract_parameters(tool_schema)
|
||||||
|
|
||||||
|
# Build function signature
|
||||||
|
param_list = []
|
||||||
|
for param_name, type_hint, required, _ in params:
|
||||||
|
if required:
|
||||||
|
param_list.append(f"{param_name}: {type_hint}")
|
||||||
|
else:
|
||||||
|
param_list.append(f"{param_name}: {type_hint} = None")
|
||||||
|
|
||||||
|
signature = f"def {tool_name}({', '.join(param_list)})"
|
||||||
|
|
||||||
|
# Build docstring
|
||||||
|
description = tool_schema.get("description", f"Call {tool_name} tool")
|
||||||
|
docstring_lines = [f' """{description}']
|
||||||
|
|
||||||
|
if params:
|
||||||
|
docstring_lines.append("")
|
||||||
|
docstring_lines.append(" Args:")
|
||||||
|
for param_name, _, _, description in params:
|
||||||
|
if description:
|
||||||
|
docstring_lines.append(f" {param_name}: {description}")
|
||||||
|
else:
|
||||||
|
docstring_lines.append(f" {param_name}")
|
||||||
|
|
||||||
|
docstring_lines.append(' """')
|
||||||
|
docstring = "\n".join(docstring_lines)
|
||||||
|
|
||||||
|
# Build function body - only pass non-None parameters
|
||||||
|
param_names = [p[0] for p in params]
|
||||||
|
if param_names:
|
||||||
|
# Build kwargs dict, filtering out None values
|
||||||
|
body = f''' kwargs = {{{", ".join([f"'{p}': {p}" for p in param_names])}}}
|
||||||
|
kwargs = {{k: v for k, v in kwargs.items() if v is not None}}
|
||||||
|
return _mcp_call('{tool_name}', **kwargs)'''
|
||||||
|
else:
|
||||||
|
body = f" return _mcp_call('{tool_name}')"
|
||||||
|
|
||||||
|
# Combine parts
|
||||||
|
return f"{signature}:\n{docstring}\n{body}"
|
||||||
|
|
||||||
|
def _extract_parameters(self, schema: dict) -> List[Tuple[str, str, bool, str]]:
|
||||||
|
"""
|
||||||
|
Extract parameter definitions from tool schema.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of (name, type_hint, required, description) tuples
|
||||||
|
"""
|
||||||
|
properties = schema.get("properties", {})
|
||||||
|
required_fields = set(schema.get("required", []))
|
||||||
|
|
||||||
|
params = []
|
||||||
|
for param_name, param_schema in properties.items():
|
||||||
|
json_type = param_schema.get("type", "string")
|
||||||
|
type_hint = self._json_type_to_python(json_type)
|
||||||
|
required = param_name in required_fields
|
||||||
|
description = param_schema.get("description", "")
|
||||||
|
|
||||||
|
params.append((param_name, type_hint, required, description))
|
||||||
|
|
||||||
|
return params
|
||||||
|
|
||||||
|
def _json_type_to_python(self, json_type: str) -> str:
|
||||||
|
"""Convert JSON schema type to Python type hint."""
|
||||||
|
type_map = {
|
||||||
|
"string": "str",
|
||||||
|
"integer": "int",
|
||||||
|
"number": "float",
|
||||||
|
"boolean": "bool",
|
||||||
|
"array": "list",
|
||||||
|
"object": "dict",
|
||||||
|
}
|
||||||
|
return type_map.get(json_type, "Any")
|
||||||
201
src/mcp_forge/mcp/manager.py
Normal file
201
src/mcp_forge/mcp/manager.py
Normal file
|
|
@ -0,0 +1,201 @@
|
||||||
|
"""MCP Client Manager for managing multiple MCP server connections."""
|
||||||
|
|
||||||
|
from typing import Dict, List, Any
|
||||||
|
from .client import MCPClientWrapper
|
||||||
|
|
||||||
|
|
||||||
|
class MCPClientManager:
|
||||||
|
"""Manages multiple MCP client connections."""
|
||||||
|
|
||||||
|
def __init__(self, config: Dict[str, Dict[str, Any]]):
|
||||||
|
"""
|
||||||
|
Initialize MCP client manager.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config: Dictionary mapping client names to their configuration.
|
||||||
|
Each config should have: command, args, env (optional)
|
||||||
|
"""
|
||||||
|
self.config = config
|
||||||
|
self.clients: Dict[str, MCPClientWrapper] = {}
|
||||||
|
self._tool_to_client: Dict[str, str] = {}
|
||||||
|
self._initialized = False
|
||||||
|
|
||||||
|
async def initialize(self) -> None:
|
||||||
|
"""
|
||||||
|
Initialize all MCP clients from config.
|
||||||
|
|
||||||
|
Connects to each configured MCP server.
|
||||||
|
Builds tool name to client mapping.
|
||||||
|
Detects tool name collisions.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If tool name collision detected
|
||||||
|
RuntimeError: If connection fails
|
||||||
|
"""
|
||||||
|
# Create and connect all clients
|
||||||
|
for client_name, client_config in self.config.items():
|
||||||
|
transport_type = client_config.get("transport", "stdio")
|
||||||
|
|
||||||
|
if transport_type == "stdio":
|
||||||
|
# Stdio transport
|
||||||
|
command = client_config["command"]
|
||||||
|
args = client_config.get("args", [])
|
||||||
|
env = client_config.get("env")
|
||||||
|
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name=client_name,
|
||||||
|
transport_type="stdio",
|
||||||
|
command=command,
|
||||||
|
args=args,
|
||||||
|
env=env
|
||||||
|
)
|
||||||
|
elif transport_type in ("http", "sse"):
|
||||||
|
# HTTP/SSE transport
|
||||||
|
url = client_config["url"]
|
||||||
|
headers = client_config.get("headers", {})
|
||||||
|
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name=client_name,
|
||||||
|
transport_type=transport_type,
|
||||||
|
url=url,
|
||||||
|
headers=headers
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Unknown transport type '{transport_type}' for client '{client_name}'")
|
||||||
|
|
||||||
|
# Connect to server
|
||||||
|
await client.connect()
|
||||||
|
|
||||||
|
# Store client
|
||||||
|
self.clients[client_name] = client
|
||||||
|
|
||||||
|
# Build tool-to-client mapping and detect collisions
|
||||||
|
await self._build_tool_mapping()
|
||||||
|
|
||||||
|
self._initialized = True
|
||||||
|
|
||||||
|
async def _build_tool_mapping(self) -> None:
|
||||||
|
"""
|
||||||
|
Build tool name to client name mapping.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If same tool name provided by multiple clients
|
||||||
|
"""
|
||||||
|
collisions: Dict[str, List[str]] = {}
|
||||||
|
|
||||||
|
for client_name, client in self.clients.items():
|
||||||
|
tools = await client.list_tools()
|
||||||
|
|
||||||
|
for tool_name in tools:
|
||||||
|
if tool_name in self._tool_to_client:
|
||||||
|
# Collision detected
|
||||||
|
if tool_name not in collisions:
|
||||||
|
collisions[tool_name] = [self._tool_to_client[tool_name]]
|
||||||
|
collisions[tool_name].append(client_name)
|
||||||
|
else:
|
||||||
|
self._tool_to_client[tool_name] = client_name
|
||||||
|
|
||||||
|
if collisions:
|
||||||
|
# Build error message
|
||||||
|
collision_details = []
|
||||||
|
for tool_name, client_names in collisions.items():
|
||||||
|
clients_str = ", ".join(client_names)
|
||||||
|
collision_details.append(f"'{tool_name}' provided by: {clients_str}")
|
||||||
|
|
||||||
|
error_msg = f"Tool name collision detected. {'; '.join(collision_details)}"
|
||||||
|
raise ValueError(error_msg)
|
||||||
|
|
||||||
|
def _check_initialized(self) -> None:
|
||||||
|
"""
|
||||||
|
Check if manager is initialized.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: If manager not initialized
|
||||||
|
"""
|
||||||
|
if not self._initialized:
|
||||||
|
raise RuntimeError("Manager not initialized. Call initialize() first.")
|
||||||
|
|
||||||
|
async def get_client_for_tool(self, tool_name: str) -> MCPClientWrapper:
|
||||||
|
"""
|
||||||
|
Get client that provides given tool.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
tool_name: Name of the tool
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
MCPClientWrapper instance that provides the tool
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: If manager not initialized
|
||||||
|
KeyError: If tool not found
|
||||||
|
"""
|
||||||
|
self._check_initialized()
|
||||||
|
|
||||||
|
if tool_name not in self._tool_to_client:
|
||||||
|
raise KeyError(f"Tool '{tool_name}' not found")
|
||||||
|
|
||||||
|
client_name = self._tool_to_client[tool_name]
|
||||||
|
return self.clients[client_name]
|
||||||
|
|
||||||
|
async def list_all_tools(self) -> List[str]:
|
||||||
|
"""
|
||||||
|
Get list of all available tool names across all clients.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of tool names
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: If manager not initialized
|
||||||
|
"""
|
||||||
|
self._check_initialized()
|
||||||
|
return list(self._tool_to_client.keys())
|
||||||
|
|
||||||
|
async def get_tool_schema(self, tool_name: str) -> dict:
|
||||||
|
"""
|
||||||
|
Get schema for tool (finds correct client).
|
||||||
|
|
||||||
|
Args:
|
||||||
|
tool_name: Name of the tool
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Tool schema dictionary
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: If manager not initialized
|
||||||
|
KeyError: If tool not found
|
||||||
|
"""
|
||||||
|
self._check_initialized()
|
||||||
|
client = await self.get_client_for_tool(tool_name)
|
||||||
|
return await client.get_tool_schema(tool_name)
|
||||||
|
|
||||||
|
async def call_tool(
|
||||||
|
self,
|
||||||
|
tool_name: str,
|
||||||
|
arguments: Dict[str, Any]
|
||||||
|
) -> Any:
|
||||||
|
"""
|
||||||
|
Call tool (finds correct client).
|
||||||
|
|
||||||
|
Args:
|
||||||
|
tool_name: Name of the tool to call
|
||||||
|
arguments: Tool arguments
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Tool result
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
RuntimeError: If manager not initialized or tool call fails
|
||||||
|
KeyError: If tool not found
|
||||||
|
"""
|
||||||
|
self._check_initialized()
|
||||||
|
client = await self.get_client_for_tool(tool_name)
|
||||||
|
return await client.call_tool(tool_name, arguments)
|
||||||
|
|
||||||
|
async def shutdown(self) -> None:
|
||||||
|
"""Disconnect all clients."""
|
||||||
|
for client in self.clients.values():
|
||||||
|
await client.disconnect()
|
||||||
|
|
||||||
|
self.clients.clear()
|
||||||
|
self._tool_to_client.clear()
|
||||||
|
self._initialized = False
|
||||||
11
src/mcp_forge/podman/__init__.py
Normal file
11
src/mcp_forge/podman/__init__.py
Normal file
|
|
@ -0,0 +1,11 @@
|
||||||
|
"""Podman integration module."""
|
||||||
|
|
||||||
|
from .client import PodmanClient, PodmanConnectionError
|
||||||
|
from .containers import ContainerConfig, SecureContainerManager
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"PodmanClient",
|
||||||
|
"PodmanConnectionError",
|
||||||
|
"ContainerConfig",
|
||||||
|
"SecureContainerManager",
|
||||||
|
]
|
||||||
157
src/mcp_forge/podman/client.py
Normal file
157
src/mcp_forge/podman/client.py
Normal file
|
|
@ -0,0 +1,157 @@
|
||||||
|
"""
|
||||||
|
Podman client wrapper with security validation.
|
||||||
|
|
||||||
|
Wraps Podman API with security validation and error handling.
|
||||||
|
All container operations are validated against security policy.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Optional
|
||||||
|
from podman import PodmanClient as BasePodmanClient
|
||||||
|
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
|
||||||
|
|
||||||
|
class PodmanConnectionError(Exception):
|
||||||
|
"""Raised when connection to Podman fails."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class PodmanClient:
|
||||||
|
"""
|
||||||
|
Wrapper around Podman API with security validation.
|
||||||
|
|
||||||
|
All container operations are validated against security policy
|
||||||
|
before being sent to Podman. Provides lazy connection and
|
||||||
|
proper error handling.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
socket_path: Path,
|
||||||
|
validator: OperationValidator,
|
||||||
|
audit_logger: AuditLogger
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize Podman client wrapper.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
socket_path: Path to Podman socket
|
||||||
|
validator: Operation validator for security checks
|
||||||
|
audit_logger: Audit logger for operation logging
|
||||||
|
"""
|
||||||
|
self.socket_path = Path(socket_path)
|
||||||
|
self.validator = validator
|
||||||
|
self.audit_logger = audit_logger
|
||||||
|
self._client: Optional[BasePodmanClient] = None
|
||||||
|
|
||||||
|
def connect(self) -> None:
|
||||||
|
"""
|
||||||
|
Connect to Podman via socket.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
PodmanConnectionError: If connection fails
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
# Verify socket exists and is accessible
|
||||||
|
self.verify_socket_access()
|
||||||
|
|
||||||
|
# Create Podman client with Unix socket
|
||||||
|
base_url = f"unix://{self.socket_path}"
|
||||||
|
self._client = BasePodmanClient(base_url=base_url)
|
||||||
|
|
||||||
|
# Test connection with ping (only if client supports it)
|
||||||
|
if hasattr(self._client, 'ping'):
|
||||||
|
self._client.ping()
|
||||||
|
|
||||||
|
except PodmanConnectionError:
|
||||||
|
# Re-raise our own exceptions
|
||||||
|
raise
|
||||||
|
except Exception as e:
|
||||||
|
raise PodmanConnectionError(
|
||||||
|
f"Failed to connect to Podman at {self.socket_path}: {e}"
|
||||||
|
) from e
|
||||||
|
|
||||||
|
def ping(self) -> bool:
|
||||||
|
"""
|
||||||
|
Test connection to Podman.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
True if connection is healthy
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
PodmanConnectionError: If not connected or ping fails
|
||||||
|
"""
|
||||||
|
if self._client is None:
|
||||||
|
raise PodmanConnectionError("Not connected to Podman")
|
||||||
|
|
||||||
|
try:
|
||||||
|
result = self._client.ping()
|
||||||
|
return result == "OK" or result is True
|
||||||
|
except Exception as e:
|
||||||
|
raise PodmanConnectionError(f"Ping failed: {e}") from e
|
||||||
|
|
||||||
|
def disconnect(self) -> None:
|
||||||
|
"""Disconnect from Podman and cleanup."""
|
||||||
|
if self._client is not None:
|
||||||
|
try:
|
||||||
|
self._client.close()
|
||||||
|
except Exception:
|
||||||
|
pass # Ignore errors during cleanup
|
||||||
|
finally:
|
||||||
|
self._client = None
|
||||||
|
|
||||||
|
def verify_socket_access(self) -> None:
|
||||||
|
"""
|
||||||
|
Verify that socket exists and is accessible.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
PodmanConnectionError: If socket is not accessible
|
||||||
|
"""
|
||||||
|
if not self.socket_path.exists():
|
||||||
|
raise PodmanConnectionError(
|
||||||
|
f"Podman socket not found: {self.socket_path}"
|
||||||
|
)
|
||||||
|
|
||||||
|
if not os.access(self.socket_path, os.R_OK):
|
||||||
|
raise PodmanConnectionError(
|
||||||
|
f"Podman socket is not readable: {self.socket_path}"
|
||||||
|
)
|
||||||
|
|
||||||
|
def check_api_version(self) -> dict:
|
||||||
|
"""
|
||||||
|
Get Podman API version information.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dictionary with version information
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
PodmanConnectionError: If not connected
|
||||||
|
"""
|
||||||
|
if self._client is None:
|
||||||
|
raise PodmanConnectionError("Not connected to Podman")
|
||||||
|
|
||||||
|
try:
|
||||||
|
return self._client.version()
|
||||||
|
except Exception as e:
|
||||||
|
raise PodmanConnectionError(
|
||||||
|
f"Failed to get API version: {e}"
|
||||||
|
) from e
|
||||||
|
|
||||||
|
@property
|
||||||
|
def client(self) -> BasePodmanClient:
|
||||||
|
"""
|
||||||
|
Get underlying Podman client (lazy connection).
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Connected Podman client
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
PodmanConnectionError: If connection fails
|
||||||
|
"""
|
||||||
|
if self._client is None:
|
||||||
|
self.connect()
|
||||||
|
assert self._client is not None # Type narrowing for mypy
|
||||||
|
return self._client
|
||||||
493
src/mcp_forge/podman/containers.py
Normal file
493
src/mcp_forge/podman/containers.py
Normal file
|
|
@ -0,0 +1,493 @@
|
||||||
|
"""
|
||||||
|
Secure container management with security enforcement.
|
||||||
|
|
||||||
|
All container operations are validated against security policy
|
||||||
|
before being sent to Podman. Provides lifecycle management
|
||||||
|
with comprehensive audit logging.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Optional, Dict, List
|
||||||
|
from datetime import datetime, timedelta
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
|
||||||
|
class ContainerConfig:
|
||||||
|
"""Container configuration with security defaults."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
image: str,
|
||||||
|
command: Optional[List[str]] = None,
|
||||||
|
environment: Optional[Dict[str, str]] = None,
|
||||||
|
volumes: Optional[Dict[str, dict]] = None,
|
||||||
|
resource_limits: Optional[ResourceLimits] = None,
|
||||||
|
working_dir: Optional[str] = None,
|
||||||
|
user: str = "1000:1000"
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize container configuration.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
image: Container image to use
|
||||||
|
command: Command to run in container
|
||||||
|
environment: Environment variables
|
||||||
|
volumes: Volume mounts (host_path -> {bind, mode})
|
||||||
|
resource_limits: Resource limits to apply
|
||||||
|
working_dir: Working directory in container (None to use image default)
|
||||||
|
user: User to run as (UID:GID)
|
||||||
|
"""
|
||||||
|
self.image = image
|
||||||
|
self.command = command or []
|
||||||
|
self.environment = environment or {}
|
||||||
|
self.volumes = volumes or {}
|
||||||
|
self.resource_limits = resource_limits
|
||||||
|
self.working_dir = working_dir
|
||||||
|
self.user = user
|
||||||
|
|
||||||
|
def to_podman_params(self) -> dict:
|
||||||
|
"""
|
||||||
|
Convert to Podman container create parameters.
|
||||||
|
|
||||||
|
Ensures all security requirements are included:
|
||||||
|
- network_mode: none
|
||||||
|
- read_only: True
|
||||||
|
- security_opt: ["no-new-privileges"]
|
||||||
|
- resource limits
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dictionary of parameters for Podman containers.create()
|
||||||
|
"""
|
||||||
|
params = {
|
||||||
|
"image": self.image,
|
||||||
|
"command": self.command if self.command else None,
|
||||||
|
"environment": self.environment,
|
||||||
|
"user": self.user,
|
||||||
|
# Security requirements
|
||||||
|
"network_mode": "none",
|
||||||
|
"read_only": True,
|
||||||
|
"security_opt": ["no-new-privileges"],
|
||||||
|
}
|
||||||
|
|
||||||
|
# Add working_dir only if explicitly set
|
||||||
|
if self.working_dir is not None:
|
||||||
|
params["working_dir"] = self.working_dir
|
||||||
|
|
||||||
|
# Add volumes if present
|
||||||
|
if self.volumes:
|
||||||
|
params["volumes"] = self.volumes
|
||||||
|
|
||||||
|
# Add resource limits if present
|
||||||
|
if self.resource_limits:
|
||||||
|
limit_params = self.resource_limits.to_podman_params()
|
||||||
|
params.update(limit_params)
|
||||||
|
# Disable swap to avoid cgroup swap.max issues on some systems
|
||||||
|
if "mem_limit" in params:
|
||||||
|
params["memswap_limit"] = -1 # Disable swap
|
||||||
|
|
||||||
|
return params
|
||||||
|
|
||||||
|
|
||||||
|
class SecureContainerManager:
|
||||||
|
"""Manages container lifecycle with security enforcement."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
podman_client: PodmanClient,
|
||||||
|
validator: OperationValidator,
|
||||||
|
audit_logger: AuditLogger
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize secure container manager.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
podman_client: Podman client wrapper
|
||||||
|
validator: Operation validator for security checks
|
||||||
|
audit_logger: Audit logger for operation logging
|
||||||
|
"""
|
||||||
|
self.podman = podman_client
|
||||||
|
self.validator = validator
|
||||||
|
self.audit_logger = audit_logger
|
||||||
|
|
||||||
|
def create_container(
|
||||||
|
self,
|
||||||
|
config: ContainerConfig,
|
||||||
|
session_id: Optional[str] = None,
|
||||||
|
name: Optional[str] = None,
|
||||||
|
**extra_params
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Create a container with security validation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config: Container configuration
|
||||||
|
session_id: Session ID for tracking
|
||||||
|
name: Optional container name
|
||||||
|
**extra_params: Additional parameters (checked for forbidden values)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Container ID
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If configuration violates security policy
|
||||||
|
"""
|
||||||
|
# Convert config to Podman parameters
|
||||||
|
params = config.to_podman_params()
|
||||||
|
|
||||||
|
# Add session label if provided
|
||||||
|
labels = {}
|
||||||
|
if session_id:
|
||||||
|
labels["mcp-forge.session"] = session_id
|
||||||
|
if labels:
|
||||||
|
params["labels"] = labels
|
||||||
|
|
||||||
|
if name:
|
||||||
|
params["name"] = name
|
||||||
|
|
||||||
|
# Merge any extra parameters (will be validated)
|
||||||
|
params.update(extra_params)
|
||||||
|
|
||||||
|
# Validate against security policy
|
||||||
|
try:
|
||||||
|
# Extract image from params for validation
|
||||||
|
self.validator.validate_container_create(
|
||||||
|
image=config.image,
|
||||||
|
params=params,
|
||||||
|
session_id=session_id
|
||||||
|
)
|
||||||
|
except SecurityError as e:
|
||||||
|
# Log security violation
|
||||||
|
self.audit_logger.log_security_violation(
|
||||||
|
operation="container_create",
|
||||||
|
reason=str(e),
|
||||||
|
session_id=session_id
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
# Create container
|
||||||
|
try:
|
||||||
|
container = self.podman.client.containers.create(**params)
|
||||||
|
container_id = container.id
|
||||||
|
|
||||||
|
# Register with validator
|
||||||
|
if session_id:
|
||||||
|
self.validator.register_session_container(container_id)
|
||||||
|
|
||||||
|
# Log successful creation
|
||||||
|
self.audit_logger.log_container_operation(
|
||||||
|
operation="create",
|
||||||
|
container_id=container_id,
|
||||||
|
image=config.image,
|
||||||
|
session_id=session_id,
|
||||||
|
details={
|
||||||
|
"name": name,
|
||||||
|
"command": config.command
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
return container_id
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.CONTAINER_CREATE,
|
||||||
|
severity=AuditSeverity.ERROR,
|
||||||
|
message=f"Container creation failed: {e}",
|
||||||
|
details={
|
||||||
|
"image": config.image,
|
||||||
|
"session_id": session_id,
|
||||||
|
"error": str(e)
|
||||||
|
}
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
def start_container(self, container_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Start a container.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID to start
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If container is not a session container
|
||||||
|
"""
|
||||||
|
# Verify container is registered (security check)
|
||||||
|
if container_id not in self.validator.session_containers:
|
||||||
|
self.audit_logger.log_security_violation(
|
||||||
|
operation="container_start",
|
||||||
|
reason=f"Attempted to start unregistered container: {container_id}"
|
||||||
|
)
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container {container_id} is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
container = self.podman.client.containers.get(container_id)
|
||||||
|
container.start()
|
||||||
|
|
||||||
|
self.audit_logger.log_container_operation(
|
||||||
|
operation="start",
|
||||||
|
container_id=container_id,
|
||||||
|
image="" # Not available without extra lookup
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.CONTAINER_START,
|
||||||
|
severity=AuditSeverity.ERROR,
|
||||||
|
message=f"Container start failed: {e}",
|
||||||
|
details={"container_id": container_id, "error": str(e)}
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
def stop_container(
|
||||||
|
self,
|
||||||
|
container_id: str,
|
||||||
|
timeout: int = 10
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Stop a container.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID to stop
|
||||||
|
timeout: Timeout in seconds
|
||||||
|
"""
|
||||||
|
# Verify container is registered
|
||||||
|
if container_id not in self.validator.session_containers:
|
||||||
|
self.audit_logger.log_security_violation(
|
||||||
|
operation="container_stop",
|
||||||
|
reason=f"Attempted to stop unregistered container: {container_id}"
|
||||||
|
)
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container {container_id} is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
container = self.podman.client.containers.get(container_id)
|
||||||
|
container.stop(timeout=timeout)
|
||||||
|
|
||||||
|
self.audit_logger.log_container_operation(
|
||||||
|
operation="stop",
|
||||||
|
container_id=container_id,
|
||||||
|
image="",
|
||||||
|
details={"timeout": timeout}
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.CONTAINER_STOP,
|
||||||
|
severity=AuditSeverity.ERROR,
|
||||||
|
message=f"Container stop failed: {e}",
|
||||||
|
details={"container_id": container_id, "error": str(e)}
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
def remove_container(
|
||||||
|
self,
|
||||||
|
container_id: str,
|
||||||
|
force: bool = False
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Remove a container.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID to remove
|
||||||
|
force: Force removal even if running
|
||||||
|
"""
|
||||||
|
# Verify container is registered
|
||||||
|
if container_id not in self.validator.session_containers:
|
||||||
|
self.audit_logger.log_security_violation(
|
||||||
|
operation="container_remove",
|
||||||
|
reason=f"Attempted to remove unregistered container: {container_id}"
|
||||||
|
)
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container {container_id} is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
container = self.podman.client.containers.get(container_id)
|
||||||
|
container.remove(force=force)
|
||||||
|
|
||||||
|
# Unregister from validator
|
||||||
|
self.validator.unregister_session_container(container_id)
|
||||||
|
|
||||||
|
self.audit_logger.log_container_operation(
|
||||||
|
operation="remove",
|
||||||
|
container_id=container_id,
|
||||||
|
image="",
|
||||||
|
details={"force": force}
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.CONTAINER_REMOVE,
|
||||||
|
severity=AuditSeverity.ERROR,
|
||||||
|
message=f"Container removal failed: {e}",
|
||||||
|
details={"container_id": container_id, "error": str(e)}
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
def get_container_logs(
|
||||||
|
self,
|
||||||
|
container_id: str,
|
||||||
|
tail: int = 100
|
||||||
|
) -> tuple[str, str]:
|
||||||
|
"""
|
||||||
|
Get container stdout and stderr logs.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID
|
||||||
|
tail: Number of lines to retrieve
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
(stdout, stderr) as strings
|
||||||
|
"""
|
||||||
|
if container_id not in self.validator.session_containers:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container {container_id} is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
container = self.podman.client.containers.get(container_id)
|
||||||
|
logs = container.logs(tail=tail, stdout=True, stderr=True)
|
||||||
|
|
||||||
|
# Podman logs returns a generator of frames, need to consume it
|
||||||
|
if hasattr(logs, '__iter__') and not isinstance(logs, (str, bytes)):
|
||||||
|
# It's a generator/iterator, consume it
|
||||||
|
logs_bytes = b''.join(logs)
|
||||||
|
logs_str = logs_bytes.decode('utf-8', errors='replace')
|
||||||
|
elif isinstance(logs, bytes):
|
||||||
|
logs_str = logs.decode('utf-8', errors='replace')
|
||||||
|
else:
|
||||||
|
logs_str = str(logs)
|
||||||
|
|
||||||
|
# For simplicity, return all logs in stdout (Podman combines them)
|
||||||
|
return logs_str, ""
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.EXECUTION_REQUEST,
|
||||||
|
severity=AuditSeverity.ERROR,
|
||||||
|
message=f"Failed to get container logs: {e}",
|
||||||
|
details={"container_id": container_id, "error": str(e)}
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
def wait_for_container(
|
||||||
|
self,
|
||||||
|
container_id: str,
|
||||||
|
timeout: int = 300
|
||||||
|
) -> int:
|
||||||
|
"""
|
||||||
|
Wait for container to exit.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID
|
||||||
|
timeout: Timeout in seconds
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Exit code
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
TimeoutError: If container doesn't exit within timeout
|
||||||
|
"""
|
||||||
|
if container_id not in self.validator.session_containers:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container {container_id} is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
container = self.podman.client.containers.get(container_id)
|
||||||
|
result = container.wait(timeout=timeout)
|
||||||
|
|
||||||
|
# Extract exit code from result
|
||||||
|
if isinstance(result, dict):
|
||||||
|
exit_code = result.get("StatusCode", 0)
|
||||||
|
else:
|
||||||
|
exit_code = result
|
||||||
|
|
||||||
|
return exit_code
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.EXECUTION_REQUEST,
|
||||||
|
severity=AuditSeverity.ERROR,
|
||||||
|
message=f"Failed to wait for container: {e}",
|
||||||
|
details={"container_id": container_id, "error": str(e)}
|
||||||
|
)
|
||||||
|
raise
|
||||||
|
|
||||||
|
def cleanup_old_containers(
|
||||||
|
self,
|
||||||
|
max_age: timedelta = timedelta(hours=24)
|
||||||
|
) -> int:
|
||||||
|
"""
|
||||||
|
Cleanup containers older than max_age.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
max_age: Maximum age for containers
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Number of containers removed
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
# Get all containers with mcp-forge.session label
|
||||||
|
containers = self.podman.client.containers.list(
|
||||||
|
all=True,
|
||||||
|
filters={"label": ["mcp-forge.session"]}
|
||||||
|
)
|
||||||
|
|
||||||
|
removed_count = 0
|
||||||
|
now = datetime.now()
|
||||||
|
|
||||||
|
for container in containers:
|
||||||
|
# Get creation time
|
||||||
|
created_str = container.attrs.get("Created", "")
|
||||||
|
if not created_str:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Parse ISO format timestamp
|
||||||
|
try:
|
||||||
|
# Remove fractional seconds and timezone for parsing
|
||||||
|
created_str = created_str.split('.')[0]
|
||||||
|
created = datetime.fromisoformat(created_str.replace('Z', ''))
|
||||||
|
except (ValueError, AttributeError):
|
||||||
|
continue
|
||||||
|
|
||||||
|
age = now - created
|
||||||
|
|
||||||
|
if age > max_age:
|
||||||
|
try:
|
||||||
|
container.remove(force=True)
|
||||||
|
removed_count += 1
|
||||||
|
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.CONTAINER_REMOVE,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message=f"Cleaned up old container: {container.id}",
|
||||||
|
details={
|
||||||
|
"container_id": container.id,
|
||||||
|
"age_hours": age.total_seconds() / 3600
|
||||||
|
}
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.CONTAINER_REMOVE,
|
||||||
|
severity=AuditSeverity.WARNING,
|
||||||
|
message=f"Failed to remove old container: {e}",
|
||||||
|
details={"container_id": container.id, "error": str(e)}
|
||||||
|
)
|
||||||
|
|
||||||
|
return removed_count
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.CONTAINER_REMOVE,
|
||||||
|
severity=AuditSeverity.ERROR,
|
||||||
|
message=f"Cleanup failed: {e}",
|
||||||
|
details={"error": str(e)}
|
||||||
|
)
|
||||||
|
raise
|
||||||
29
src/mcp_forge/security/__init__.py
Normal file
29
src/mcp_forge/security/__init__.py
Normal file
|
|
@ -0,0 +1,29 @@
|
||||||
|
"""Security module for resource validation and enforcement."""
|
||||||
|
|
||||||
|
from .resource_limits import (
|
||||||
|
parse_memory_string,
|
||||||
|
parse_cpu_quota,
|
||||||
|
parse_storage_string,
|
||||||
|
ResourceLimits,
|
||||||
|
)
|
||||||
|
from .allowlist import (
|
||||||
|
SecurityError,
|
||||||
|
OperationValidator,
|
||||||
|
)
|
||||||
|
from .audit import (
|
||||||
|
AuditEventType,
|
||||||
|
AuditSeverity,
|
||||||
|
AuditLogger,
|
||||||
|
)
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"parse_memory_string",
|
||||||
|
"parse_cpu_quota",
|
||||||
|
"parse_storage_string",
|
||||||
|
"ResourceLimits",
|
||||||
|
"SecurityError",
|
||||||
|
"OperationValidator",
|
||||||
|
"AuditEventType",
|
||||||
|
"AuditSeverity",
|
||||||
|
"AuditLogger",
|
||||||
|
]
|
||||||
283
src/mcp_forge/security/allowlist.py
Normal file
283
src/mcp_forge/security/allowlist.py
Normal file
|
|
@ -0,0 +1,283 @@
|
||||||
|
"""
|
||||||
|
Podman operation allowlist and security policy enforcement.
|
||||||
|
|
||||||
|
Defines and enforces allowed Podman operations with parameter validation.
|
||||||
|
Prevents privilege escalation and unauthorized resource access.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import Set, Optional, Dict, Any
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
import fnmatch
|
||||||
|
|
||||||
|
|
||||||
|
class SecurityError(Exception):
|
||||||
|
"""Raised when security policy is violated."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
# Allowed container images with wildcard support
|
||||||
|
ALLOWED_IMAGES = [
|
||||||
|
"python:3.11*",
|
||||||
|
"python:3.12*",
|
||||||
|
"jupyter/*",
|
||||||
|
"mcp-forge/*", # Legacy/custom images
|
||||||
|
]
|
||||||
|
|
||||||
|
# Parameters that are forbidden in container creation
|
||||||
|
FORBIDDEN_CONTAINER_PARAMS = [
|
||||||
|
"privileged",
|
||||||
|
"cap_add",
|
||||||
|
"devices",
|
||||||
|
"pid_mode",
|
||||||
|
"ipc_mode",
|
||||||
|
]
|
||||||
|
|
||||||
|
# Required parameters that must be set with specific values
|
||||||
|
REQUIRED_CONTAINER_PARAMS = {
|
||||||
|
"network_mode": "none",
|
||||||
|
"read_only": True,
|
||||||
|
"security_opt": ["no-new-privileges"],
|
||||||
|
"user": "1000:1000",
|
||||||
|
}
|
||||||
|
|
||||||
|
# Allowed volume mount patterns (with {session_id} placeholder)
|
||||||
|
ALLOWED_VOLUME_PATTERNS = [
|
||||||
|
"/mcp-forge/sessions/{session_id}/*",
|
||||||
|
"/mcp-forge/shared/readonly/*",
|
||||||
|
"/mcp-forge/uploads/{session_id}/*",
|
||||||
|
"/tmp/mcp-forge-bridge.sock", # MCP tool bridge socket
|
||||||
|
]
|
||||||
|
|
||||||
|
# Paths that must never be mounted
|
||||||
|
FORBIDDEN_MOUNT_PATHS = [
|
||||||
|
"/",
|
||||||
|
"/etc",
|
||||||
|
"/var/run/docker.sock",
|
||||||
|
"/var/run/podman/podman.sock",
|
||||||
|
"/sys",
|
||||||
|
"/proc",
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
class OperationValidator:
|
||||||
|
"""Validates Podman operations against security policy."""
|
||||||
|
|
||||||
|
def __init__(self, config: SecurityConfig):
|
||||||
|
"""
|
||||||
|
Initialize operation validator.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config: Security configuration
|
||||||
|
"""
|
||||||
|
self.config = config
|
||||||
|
self.session_containers: Set[str] = set()
|
||||||
|
|
||||||
|
def validate_container_create(
|
||||||
|
self,
|
||||||
|
image: str,
|
||||||
|
params: Dict[str, Any],
|
||||||
|
session_id: Optional[str] = None
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Validate container creation parameters against security policy.
|
||||||
|
|
||||||
|
Checks:
|
||||||
|
- Image is in allowlist
|
||||||
|
- No forbidden parameters present
|
||||||
|
- All required parameters set correctly
|
||||||
|
- Volume mounts are valid (if present)
|
||||||
|
|
||||||
|
Args:
|
||||||
|
image: Container image name
|
||||||
|
params: Container creation parameters
|
||||||
|
session_id: Optional session ID for volume validation
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If any security policy is violated
|
||||||
|
"""
|
||||||
|
# Validate image
|
||||||
|
self.validate_image_name(image)
|
||||||
|
|
||||||
|
# Check for forbidden parameters
|
||||||
|
for forbidden_param in FORBIDDEN_CONTAINER_PARAMS:
|
||||||
|
if forbidden_param in params:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Forbidden parameter '{forbidden_param}' in container creation"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Check required parameters
|
||||||
|
for required_param, required_value in REQUIRED_CONTAINER_PARAMS.items():
|
||||||
|
if required_param not in params:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Required parameter '{required_param}' missing. "
|
||||||
|
f"Must be set to: {required_value}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Validate the value matches requirement
|
||||||
|
actual_value = params[required_param]
|
||||||
|
if actual_value != required_value:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Parameter '{required_param}' must be {required_value}, "
|
||||||
|
f"got: {actual_value}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Validate volume mounts if present
|
||||||
|
if "volumes" in params and session_id:
|
||||||
|
for volume_spec in params["volumes"]:
|
||||||
|
# Volume specs can be in various formats, extract host path
|
||||||
|
# Simplified: assume it's a dict with 'bind' key or string "host:container"
|
||||||
|
if isinstance(volume_spec, dict):
|
||||||
|
host_path = volume_spec.get("bind", {}).get("source", "")
|
||||||
|
elif isinstance(volume_spec, str) and ":" in volume_spec:
|
||||||
|
host_path = volume_spec.split(":")[0]
|
||||||
|
else:
|
||||||
|
host_path = str(volume_spec)
|
||||||
|
|
||||||
|
if host_path:
|
||||||
|
self.validate_volume_mount(host_path, session_id)
|
||||||
|
|
||||||
|
def validate_container_start(self, container_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Validate container start - must be session container.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID to start
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If container is not a registered session container
|
||||||
|
"""
|
||||||
|
if container_id not in self.session_containers:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container '{container_id}' is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
def validate_container_stop(self, container_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Validate container stop - must be session container.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID to stop
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If container is not a registered session container
|
||||||
|
"""
|
||||||
|
if container_id not in self.session_containers:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container '{container_id}' is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
def validate_container_remove(self, container_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Validate container remove - must be session container.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID to remove
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If container is not a registered session container
|
||||||
|
"""
|
||||||
|
if container_id not in self.session_containers:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Container '{container_id}' is not a registered session container"
|
||||||
|
)
|
||||||
|
|
||||||
|
def validate_volume_mount(self, mount_path: str, session_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Validate volume mount path against allowlist patterns.
|
||||||
|
|
||||||
|
Checks:
|
||||||
|
- Path is not in forbidden paths
|
||||||
|
- Path matches one of the allowed patterns
|
||||||
|
- Session ID in path matches provided session_id
|
||||||
|
|
||||||
|
Args:
|
||||||
|
mount_path: Host path to mount
|
||||||
|
session_id: Session ID for path validation
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If mount path is forbidden or not in allowlist
|
||||||
|
"""
|
||||||
|
# Check forbidden paths
|
||||||
|
for forbidden_path in FORBIDDEN_MOUNT_PATHS:
|
||||||
|
if mount_path == forbidden_path or mount_path.startswith(forbidden_path + "/"):
|
||||||
|
raise SecurityError(
|
||||||
|
f"Mount path '{mount_path}' is forbidden (matches '{forbidden_path}')"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Check against allowed patterns
|
||||||
|
allowed = False
|
||||||
|
for pattern in ALLOWED_VOLUME_PATTERNS:
|
||||||
|
# Substitute session_id in pattern
|
||||||
|
pattern_resolved = pattern.replace("{session_id}", session_id)
|
||||||
|
|
||||||
|
# Convert pattern to fnmatch format (already uses *)
|
||||||
|
if fnmatch.fnmatch(mount_path, pattern_resolved):
|
||||||
|
allowed = True
|
||||||
|
break
|
||||||
|
|
||||||
|
if not allowed:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Mount path '{mount_path}' does not match any allowed pattern. "
|
||||||
|
f"Session ID: {session_id}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Additional check: if path contains session directory, verify it matches
|
||||||
|
if "/sessions/" in mount_path or "/uploads/" in mount_path:
|
||||||
|
# Extract session from path
|
||||||
|
parts = mount_path.split("/")
|
||||||
|
try:
|
||||||
|
if "sessions" in parts:
|
||||||
|
idx = parts.index("sessions")
|
||||||
|
path_session = parts[idx + 1]
|
||||||
|
elif "uploads" in parts:
|
||||||
|
idx = parts.index("uploads")
|
||||||
|
path_session = parts[idx + 1]
|
||||||
|
else:
|
||||||
|
return # No session in path
|
||||||
|
|
||||||
|
if path_session != session_id:
|
||||||
|
raise SecurityError(
|
||||||
|
f"Session ID in path ('{path_session}') does not match "
|
||||||
|
f"provided session ID ('{session_id}')"
|
||||||
|
)
|
||||||
|
except (IndexError, ValueError):
|
||||||
|
pass # Path structure doesn't match expected format
|
||||||
|
|
||||||
|
def validate_image_name(self, image: str) -> None:
|
||||||
|
"""
|
||||||
|
Validate image name against allowlist.
|
||||||
|
|
||||||
|
Supports wildcard patterns (e.g., "mcp-forge/custom:*").
|
||||||
|
|
||||||
|
Args:
|
||||||
|
image: Image name to validate
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
SecurityError: If image is not in allowlist
|
||||||
|
"""
|
||||||
|
for allowed_image in ALLOWED_IMAGES:
|
||||||
|
if fnmatch.fnmatch(image, allowed_image):
|
||||||
|
return
|
||||||
|
|
||||||
|
raise SecurityError(
|
||||||
|
f"Image '{image}' is not in allowlist. "
|
||||||
|
f"Allowed images: {', '.join(ALLOWED_IMAGES)}"
|
||||||
|
)
|
||||||
|
|
||||||
|
def register_session_container(self, container_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Register container as belonging to a session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID to register
|
||||||
|
"""
|
||||||
|
self.session_containers.add(container_id)
|
||||||
|
|
||||||
|
def unregister_session_container(self, container_id: str) -> None:
|
||||||
|
"""
|
||||||
|
Unregister session container.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
container_id: Container ID to unregister
|
||||||
|
"""
|
||||||
|
self.session_containers.discard(container_id)
|
||||||
226
src/mcp_forge/security/audit.py
Normal file
226
src/mcp_forge/security/audit.py
Normal file
|
|
@ -0,0 +1,226 @@
|
||||||
|
"""
|
||||||
|
Audit logger for security-relevant operations.
|
||||||
|
|
||||||
|
Provides structured JSON logging of all security-relevant operations.
|
||||||
|
Thread-safe and includes automatic log rotation support.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import threading
|
||||||
|
from datetime import datetime
|
||||||
|
from enum import Enum
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any, Optional, Dict
|
||||||
|
|
||||||
|
|
||||||
|
class AuditEventType(Enum):
|
||||||
|
"""Types of auditable events."""
|
||||||
|
|
||||||
|
CONTAINER_CREATE = "container.create"
|
||||||
|
CONTAINER_START = "container.start"
|
||||||
|
CONTAINER_STOP = "container.stop"
|
||||||
|
CONTAINER_REMOVE = "container.remove"
|
||||||
|
EXECUTION_REQUEST = "execution.request"
|
||||||
|
SECURITY_VIOLATION = "security.violation"
|
||||||
|
BUILD_REQUEST = "build.request"
|
||||||
|
BUILD_COMPLETE = "build.complete"
|
||||||
|
SESSION_CREATE = "session.create"
|
||||||
|
SESSION_DESTROY = "session.destroy"
|
||||||
|
IMAGE_BUILD_START = "image.build.start"
|
||||||
|
IMAGE_BUILD_SUCCESS = "image.build.success"
|
||||||
|
IMAGE_BUILD_FAILURE = "image.build.failure"
|
||||||
|
MCP_TOOL_CALL = "mcp.tool.call"
|
||||||
|
|
||||||
|
|
||||||
|
class AuditSeverity(Enum):
|
||||||
|
"""Severity levels for audit events."""
|
||||||
|
|
||||||
|
INFO = "info"
|
||||||
|
WARNING = "warning"
|
||||||
|
ERROR = "error"
|
||||||
|
CRITICAL = "critical"
|
||||||
|
|
||||||
|
|
||||||
|
class AuditLogger:
|
||||||
|
"""
|
||||||
|
Thread-safe structured audit logger.
|
||||||
|
|
||||||
|
Logs security-relevant events to JSON Lines format for easy parsing
|
||||||
|
and analysis. All operations are thread-safe using a lock.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, log_path: Path):
|
||||||
|
"""
|
||||||
|
Initialize audit logger.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
log_path: Path to audit log file
|
||||||
|
"""
|
||||||
|
self.log_path = Path(log_path)
|
||||||
|
self._lock = threading.Lock()
|
||||||
|
self._ensure_log_file()
|
||||||
|
|
||||||
|
def log(
|
||||||
|
self,
|
||||||
|
event_type: AuditEventType,
|
||||||
|
severity: AuditSeverity,
|
||||||
|
message: str,
|
||||||
|
session_id: Optional[str] = None,
|
||||||
|
user_id: Optional[str] = None,
|
||||||
|
details: Optional[Dict[str, Any]] = None,
|
||||||
|
error: Optional[str] = None
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Log an audit event.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
event_type: Type of event being logged
|
||||||
|
severity: Severity level of the event
|
||||||
|
message: Human-readable message describing the event
|
||||||
|
session_id: Optional session ID associated with event
|
||||||
|
user_id: Optional user ID associated with event
|
||||||
|
details: Optional dictionary of additional details
|
||||||
|
error: Optional error message if event represents an error
|
||||||
|
|
||||||
|
Note:
|
||||||
|
All log entries are JSON objects, one per line.
|
||||||
|
Timestamps are in ISO 8601 format.
|
||||||
|
No PII (code content, tokens, file content) should be logged.
|
||||||
|
"""
|
||||||
|
from datetime import timezone
|
||||||
|
|
||||||
|
entry = {
|
||||||
|
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||||
|
"event_type": event_type.value,
|
||||||
|
"severity": severity.value,
|
||||||
|
"message": message,
|
||||||
|
}
|
||||||
|
|
||||||
|
if session_id is not None:
|
||||||
|
entry["session_id"] = session_id
|
||||||
|
|
||||||
|
if user_id is not None:
|
||||||
|
entry["user_id"] = user_id
|
||||||
|
|
||||||
|
if details is not None:
|
||||||
|
entry["details"] = details
|
||||||
|
|
||||||
|
if error is not None:
|
||||||
|
entry["error"] = error
|
||||||
|
|
||||||
|
self._write_log_entry(entry)
|
||||||
|
|
||||||
|
def log_container_operation(
|
||||||
|
self,
|
||||||
|
operation: str,
|
||||||
|
container_id: str,
|
||||||
|
image: str,
|
||||||
|
session_id: Optional[str] = None,
|
||||||
|
user_id: Optional[str] = None,
|
||||||
|
details: Optional[Dict[str, Any]] = None,
|
||||||
|
error: Optional[str] = None
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Log a container operation with standard fields.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
operation: Operation type (create, start, stop, remove)
|
||||||
|
container_id: Container ID
|
||||||
|
image: Container image name
|
||||||
|
session_id: Optional session ID
|
||||||
|
user_id: Optional user ID
|
||||||
|
details: Optional additional details
|
||||||
|
error: Optional error message
|
||||||
|
"""
|
||||||
|
# Map operation to event type
|
||||||
|
event_type_map = {
|
||||||
|
"create": AuditEventType.CONTAINER_CREATE,
|
||||||
|
"start": AuditEventType.CONTAINER_START,
|
||||||
|
"stop": AuditEventType.CONTAINER_STOP,
|
||||||
|
"remove": AuditEventType.CONTAINER_REMOVE,
|
||||||
|
}
|
||||||
|
|
||||||
|
event_type = event_type_map.get(
|
||||||
|
operation,
|
||||||
|
AuditEventType.CONTAINER_CREATE
|
||||||
|
)
|
||||||
|
|
||||||
|
# Determine severity
|
||||||
|
severity = AuditSeverity.ERROR if error else AuditSeverity.INFO
|
||||||
|
|
||||||
|
# Build entry with container fields at top level
|
||||||
|
from datetime import timezone
|
||||||
|
|
||||||
|
entry = {
|
||||||
|
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||||
|
"event_type": event_type.value,
|
||||||
|
"severity": severity.value,
|
||||||
|
"message": f"Container {operation}: {container_id}",
|
||||||
|
"container_id": container_id,
|
||||||
|
"image": image,
|
||||||
|
}
|
||||||
|
|
||||||
|
if session_id is not None:
|
||||||
|
entry["session_id"] = session_id
|
||||||
|
|
||||||
|
if user_id is not None:
|
||||||
|
entry["user_id"] = user_id
|
||||||
|
|
||||||
|
if details:
|
||||||
|
entry["details"] = details
|
||||||
|
|
||||||
|
if error is not None:
|
||||||
|
entry["error"] = error
|
||||||
|
|
||||||
|
self._write_log_entry(entry)
|
||||||
|
|
||||||
|
def log_security_violation(
|
||||||
|
self,
|
||||||
|
operation: str,
|
||||||
|
reason: str,
|
||||||
|
session_id: Optional[str] = None
|
||||||
|
) -> None:
|
||||||
|
"""
|
||||||
|
Log a security violation at CRITICAL severity.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
operation: Operation that was attempted
|
||||||
|
reason: Reason the operation was blocked
|
||||||
|
session_id: Optional session ID
|
||||||
|
"""
|
||||||
|
from datetime import timezone
|
||||||
|
|
||||||
|
entry = {
|
||||||
|
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||||
|
"event_type": AuditEventType.SECURITY_VIOLATION.value,
|
||||||
|
"severity": AuditSeverity.CRITICAL.value,
|
||||||
|
"message": f"Security violation: {operation}",
|
||||||
|
"operation": operation,
|
||||||
|
"reason": reason,
|
||||||
|
}
|
||||||
|
|
||||||
|
if session_id is not None:
|
||||||
|
entry["session_id"] = session_id
|
||||||
|
|
||||||
|
self._write_log_entry(entry)
|
||||||
|
|
||||||
|
def _ensure_log_file(self) -> None:
|
||||||
|
"""Ensure log file and directory exist."""
|
||||||
|
# Create parent directories if needed
|
||||||
|
self.log_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
# Create log file if it doesn't exist
|
||||||
|
if not self.log_path.exists():
|
||||||
|
self.log_path.touch()
|
||||||
|
|
||||||
|
def _write_log_entry(self, entry: Dict[str, Any]) -> None:
|
||||||
|
"""
|
||||||
|
Thread-safe write of log entry.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
entry: Dictionary to write as JSON line
|
||||||
|
"""
|
||||||
|
with self._lock:
|
||||||
|
with open(self.log_path, "a") as f:
|
||||||
|
json.dump(entry, f)
|
||||||
|
f.write("\n")
|
||||||
150
src/mcp_forge/security/resource_limits.py
Normal file
150
src/mcp_forge/security/resource_limits.py
Normal file
|
|
@ -0,0 +1,150 @@
|
||||||
|
"""
|
||||||
|
Resource limit parser and validator.
|
||||||
|
|
||||||
|
Parses and validates resource limit strings (memory, CPU, storage).
|
||||||
|
All values must be positive and within reasonable limits.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import re
|
||||||
|
from typing import Dict, Any
|
||||||
|
|
||||||
|
|
||||||
|
def parse_memory_string(memory: str) -> int:
|
||||||
|
"""
|
||||||
|
Parse memory string to bytes.
|
||||||
|
|
||||||
|
Supports: k, m, g suffixes (case-insensitive)
|
||||||
|
Examples: "512m" → 536870912, "2g" → 2147483648
|
||||||
|
|
||||||
|
Args:
|
||||||
|
memory: Memory string with suffix (e.g., "512m", "2g", "1024k")
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Memory in bytes
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If format is invalid or value is <= 0
|
||||||
|
"""
|
||||||
|
memory = memory.strip()
|
||||||
|
|
||||||
|
# Pattern: optional sign, number (int or float), suffix (k/m/g)
|
||||||
|
pattern = r'^(-?\d+(?:\.\d+)?)\s*([kmgKMG])$'
|
||||||
|
match = re.match(pattern, memory)
|
||||||
|
|
||||||
|
if not match:
|
||||||
|
raise ValueError(
|
||||||
|
f"Invalid memory format: '{memory}'. "
|
||||||
|
f"Expected format: <number><k|m|g> (e.g., '512m', '2g')"
|
||||||
|
)
|
||||||
|
|
||||||
|
value_str, suffix = match.groups()
|
||||||
|
value = float(value_str)
|
||||||
|
|
||||||
|
if value <= 0:
|
||||||
|
raise ValueError(
|
||||||
|
f"Memory value must be positive, got: {value}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Convert to bytes
|
||||||
|
suffix_lower = suffix.lower()
|
||||||
|
multipliers = {
|
||||||
|
'k': 1024,
|
||||||
|
'm': 1024 * 1024,
|
||||||
|
'g': 1024 * 1024 * 1024,
|
||||||
|
}
|
||||||
|
|
||||||
|
bytes_value = int(value * multipliers[suffix_lower])
|
||||||
|
|
||||||
|
return bytes_value
|
||||||
|
|
||||||
|
|
||||||
|
def parse_cpu_quota(cpu_quota: int) -> int:
|
||||||
|
"""
|
||||||
|
Validate CPU quota value.
|
||||||
|
|
||||||
|
CPU quota is in microseconds per 100ms period.
|
||||||
|
100000 = 100% of one CPU core
|
||||||
|
|
||||||
|
Args:
|
||||||
|
cpu_quota: CPU quota in microseconds (e.g., 50000 for 50% of one core)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Validated CPU quota value
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If quota <= 0 or > 1000000 (10 cores max)
|
||||||
|
"""
|
||||||
|
if cpu_quota <= 0:
|
||||||
|
raise ValueError(
|
||||||
|
f"CPU quota must be positive, got: {cpu_quota}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Maximum of 10 cores (1000000 microseconds)
|
||||||
|
if cpu_quota > 1000000:
|
||||||
|
raise ValueError(
|
||||||
|
f"CPU quota exceeds maximum of 1000000 (10 cores), got: {cpu_quota}"
|
||||||
|
)
|
||||||
|
|
||||||
|
return cpu_quota
|
||||||
|
|
||||||
|
|
||||||
|
def parse_storage_string(storage: str) -> int:
|
||||||
|
"""
|
||||||
|
Parse storage string to bytes (same as memory).
|
||||||
|
|
||||||
|
Args:
|
||||||
|
storage: Storage string with suffix (e.g., "1g", "512m")
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Storage in bytes
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If format is invalid or value is <= 0
|
||||||
|
"""
|
||||||
|
return parse_memory_string(storage)
|
||||||
|
|
||||||
|
|
||||||
|
class ResourceLimits:
|
||||||
|
"""
|
||||||
|
Resource limits with validation.
|
||||||
|
|
||||||
|
Encapsulates memory, storage, CPU, and timeout limits with validation.
|
||||||
|
Provides conversion to Podman container parameters.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
memory: str,
|
||||||
|
storage: str,
|
||||||
|
cpu_quota: int,
|
||||||
|
timeout: int = 300
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize resource limits with validation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
memory: Memory limit string (e.g., "512m", "2g")
|
||||||
|
storage: Storage limit string (e.g., "1g", "10g")
|
||||||
|
cpu_quota: CPU quota in microseconds per 100ms period
|
||||||
|
timeout: Execution timeout in seconds (default: 300)
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If any limit is invalid
|
||||||
|
"""
|
||||||
|
self.memory_bytes = parse_memory_string(memory)
|
||||||
|
self.storage_bytes = parse_storage_string(storage)
|
||||||
|
self.cpu_quota = parse_cpu_quota(cpu_quota)
|
||||||
|
self.timeout = timeout
|
||||||
|
|
||||||
|
def to_podman_params(self) -> Dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Convert to Podman container create parameters.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Dictionary of parameters suitable for Podman container creation
|
||||||
|
"""
|
||||||
|
return {
|
||||||
|
"mem_limit": str(self.memory_bytes), # Podman expects string
|
||||||
|
"cpu_quota": self.cpu_quota
|
||||||
|
# Note: storage_bytes tracked internally but not passed to Podman (not supported)
|
||||||
|
}
|
||||||
6
src/mcp_forge/server/__init__.py
Normal file
6
src/mcp_forge/server/__init__.py
Normal file
|
|
@ -0,0 +1,6 @@
|
||||||
|
"""MCP Forge Server - MCP Tools and Resources."""
|
||||||
|
|
||||||
|
from .resources import ResourceHandler
|
||||||
|
from .server import ForgeServer
|
||||||
|
|
||||||
|
__all__ = ["ResourceHandler", "ForgeServer"]
|
||||||
210
src/mcp_forge/server/resources.py
Normal file
210
src/mcp_forge/server/resources.py
Normal file
|
|
@ -0,0 +1,210 @@
|
||||||
|
"""MCP resource handlers for discovery and state."""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import re
|
||||||
|
from typing import Tuple, Dict, Any
|
||||||
|
from mcp.types import TextResourceContents
|
||||||
|
|
||||||
|
|
||||||
|
class ResourceHandler:
|
||||||
|
"""Handles MCP resource requests."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
client_manager,
|
||||||
|
session_manager,
|
||||||
|
environment_builder,
|
||||||
|
config
|
||||||
|
):
|
||||||
|
"""Initialize resource handler with dependencies.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
client_manager: MCPClientManager instance for tool discovery
|
||||||
|
session_manager: SessionManager instance for session state
|
||||||
|
environment_builder: EnvironmentBuilder instance for environments
|
||||||
|
config: ForgeConfig instance for configuration info
|
||||||
|
"""
|
||||||
|
self.client_manager = client_manager
|
||||||
|
self.session_manager = session_manager
|
||||||
|
self.environment_builder = environment_builder
|
||||||
|
self.config = config
|
||||||
|
|
||||||
|
async def handle_resource(self, uri: str) -> TextResourceContents:
|
||||||
|
"""Handle resource request based on URI.
|
||||||
|
|
||||||
|
Supported URIs:
|
||||||
|
- tools/available
|
||||||
|
- sessions/{id}/state
|
||||||
|
- sessions/{id}/variables
|
||||||
|
- environments/list
|
||||||
|
- environment/info
|
||||||
|
|
||||||
|
Args:
|
||||||
|
uri: Resource URI to handle
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
TextResourceContents with JSON content
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
KeyError: If resource not found or session doesn't exist
|
||||||
|
"""
|
||||||
|
resource_type, params = self._parse_uri(uri)
|
||||||
|
|
||||||
|
if resource_type == "tools_available":
|
||||||
|
return await self._handle_tools_available()
|
||||||
|
elif resource_type == "session_state":
|
||||||
|
return await self._handle_session_state(params["session_id"])
|
||||||
|
elif resource_type == "session_variables":
|
||||||
|
return await self._handle_session_variables(params["session_id"])
|
||||||
|
elif resource_type == "sessions_list":
|
||||||
|
return await self._handle_sessions_list()
|
||||||
|
elif resource_type == "environments_list":
|
||||||
|
return await self._handle_environments_list()
|
||||||
|
elif resource_type == "environment_info":
|
||||||
|
return await self._handle_environment_info()
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Unknown resource URI: {uri}")
|
||||||
|
|
||||||
|
async def _handle_tools_available(self) -> TextResourceContents:
|
||||||
|
"""Return list of available MCP tools."""
|
||||||
|
tools = await self.client_manager.list_all_tools()
|
||||||
|
content = json.dumps({"tools": tools})
|
||||||
|
|
||||||
|
return TextResourceContents(
|
||||||
|
uri="mcp://forge/tools/available",
|
||||||
|
mimeType="application/json",
|
||||||
|
text=content
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _handle_session_state(self, session_id: str) -> TextResourceContents:
|
||||||
|
"""Return documented state for session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: ID of session to get state for
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
TextResourceContents with session state JSON
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
KeyError: If session doesn't exist
|
||||||
|
"""
|
||||||
|
state = self.session_manager.get_session_state(session_id)
|
||||||
|
content = json.dumps(state.to_dict())
|
||||||
|
|
||||||
|
return TextResourceContents(
|
||||||
|
uri=f"mcp://forge/sessions/{session_id}/state",
|
||||||
|
mimeType="application/json",
|
||||||
|
text=content
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _handle_session_variables(self, session_id: str) -> TextResourceContents:
|
||||||
|
"""Return list of variables in session.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: ID of session to get variables for
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
TextResourceContents with variables list JSON
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
KeyError: If session doesn't exist
|
||||||
|
"""
|
||||||
|
state = self.session_manager.get_session_state(session_id)
|
||||||
|
content = json.dumps({"variables": state.all_variables})
|
||||||
|
|
||||||
|
return TextResourceContents(
|
||||||
|
uri=f"mcp://forge/sessions/{session_id}/variables",
|
||||||
|
mimeType="application/json",
|
||||||
|
text=content
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _handle_sessions_list(self) -> TextResourceContents:
|
||||||
|
"""Return list of active sessions."""
|
||||||
|
sessions = self.session_manager.list_sessions()
|
||||||
|
content = json.dumps({"sessions": sessions})
|
||||||
|
|
||||||
|
return TextResourceContents(
|
||||||
|
uri="mcp://forge/sessions/list",
|
||||||
|
mimeType="application/json",
|
||||||
|
text=content
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _handle_environments_list(self) -> TextResourceContents:
|
||||||
|
"""Return list of custom environments and templates."""
|
||||||
|
templates = self.environment_builder.list_templates()
|
||||||
|
|
||||||
|
# Format response with templates
|
||||||
|
content = json.dumps({
|
||||||
|
"templates": templates,
|
||||||
|
"custom_environments": []
|
||||||
|
})
|
||||||
|
|
||||||
|
return TextResourceContents(
|
||||||
|
uri="mcp://forge/environments/list",
|
||||||
|
mimeType="application/json",
|
||||||
|
text=content
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _handle_environment_info(self) -> TextResourceContents:
|
||||||
|
"""Return environment information and configuration."""
|
||||||
|
# Safely extract base images
|
||||||
|
base_images = []
|
||||||
|
if hasattr(self.config, "base_images"):
|
||||||
|
base_images_attr = self.config.base_images
|
||||||
|
# Handle both dict and Mock objects
|
||||||
|
if hasattr(base_images_attr, "keys"):
|
||||||
|
try:
|
||||||
|
base_images = list(base_images_attr.keys())
|
||||||
|
except (TypeError, AttributeError):
|
||||||
|
pass
|
||||||
|
|
||||||
|
content = json.dumps({
|
||||||
|
"base_images": base_images,
|
||||||
|
"max_packages": getattr(self.config, "max_packages", 50),
|
||||||
|
"max_build_time": getattr(self.config, "max_build_time", 300)
|
||||||
|
})
|
||||||
|
|
||||||
|
return TextResourceContents(
|
||||||
|
uri="mcp://forge/environment/info",
|
||||||
|
mimeType="application/json",
|
||||||
|
text=content
|
||||||
|
)
|
||||||
|
|
||||||
|
def _parse_uri(self, uri: str) -> Tuple[str, Dict[str, Any]]:
|
||||||
|
"""Parse URI into components.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
uri: Resource URI (e.g., "mcp://forge/sessions/123/state")
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Tuple of (resource_type, parameters)
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If URI format is invalid
|
||||||
|
"""
|
||||||
|
# Validate and strip URI scheme
|
||||||
|
if not uri.startswith("mcp://forge/"):
|
||||||
|
raise ValueError(f"Invalid URI format: {uri}")
|
||||||
|
|
||||||
|
# Extract path
|
||||||
|
path = uri.replace("mcp://forge/", "")
|
||||||
|
|
||||||
|
# Simple patterns for known resource types
|
||||||
|
patterns = [
|
||||||
|
(r"^tools/available$", "tools_available"),
|
||||||
|
(r"^sessions/([^/]+)/state$", "session_state"),
|
||||||
|
(r"^sessions/([^/]+)/variables$", "session_variables"),
|
||||||
|
(r"^sessions/list$", "sessions_list"),
|
||||||
|
(r"^environments/list$", "environments_list"),
|
||||||
|
(r"^environment/info$", "environment_info"),
|
||||||
|
]
|
||||||
|
|
||||||
|
for pattern, resource_type in patterns:
|
||||||
|
match = re.match(pattern, path)
|
||||||
|
if match:
|
||||||
|
params = {}
|
||||||
|
if match.groups():
|
||||||
|
params["session_id"] = match.group(1)
|
||||||
|
return resource_type, params
|
||||||
|
|
||||||
|
raise ValueError(f"Unknown resource URI: {uri}")
|
||||||
565
src/mcp_forge/server/server.py
Normal file
565
src/mcp_forge/server/server.py
Normal file
|
|
@ -0,0 +1,565 @@
|
||||||
|
"""MCP Forge Server - Main server implementation."""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
from contextlib import asynccontextmanager
|
||||||
|
from typing import Optional, Dict, List
|
||||||
|
import logging
|
||||||
|
|
||||||
|
from fastmcp import FastMCP
|
||||||
|
|
||||||
|
from ..config.schema import ForgeConfig
|
||||||
|
from ..security.audit import AuditLogger
|
||||||
|
from ..security.allowlist import OperationValidator
|
||||||
|
from ..podman.client import PodmanClient
|
||||||
|
from ..podman.containers import SecureContainerManager
|
||||||
|
from ..mcp.manager import MCPClientManager
|
||||||
|
from ..mcp.bridge import ToolBridgeServer
|
||||||
|
from ..mcp.injection import ToolInjectionGenerator
|
||||||
|
from ..execution.simple.backend import SimpleBackend
|
||||||
|
from ..execution.jupyter.backend import JupyterBackend
|
||||||
|
from ..execution.jupyter.kernel import JupyterKernelManager
|
||||||
|
from ..execution.jupyter.sessions import SessionManager
|
||||||
|
from ..builder.environment_builder import EnvironmentBuilder
|
||||||
|
from .resources import ResourceHandler
|
||||||
|
from .tools.execute_python import ExecutePythonTool
|
||||||
|
from .tools.document_state import DocumentStateTool
|
||||||
|
from .tools.build_environment import BuildEnvironmentTool
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class ForgeServer:
|
||||||
|
"""MCP-Forge server implementation."""
|
||||||
|
|
||||||
|
def __init__(self, config: ForgeConfig):
|
||||||
|
"""
|
||||||
|
Initialize MCP Forge Server.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config: Forge configuration
|
||||||
|
"""
|
||||||
|
self.config = config
|
||||||
|
|
||||||
|
# Create lifespan handler for async initialization
|
||||||
|
@asynccontextmanager
|
||||||
|
async def mcp_forge_lifespan(server):
|
||||||
|
"""Handle server startup and shutdown."""
|
||||||
|
# Startup: initialize MCP clients and bridge server
|
||||||
|
await self.initialize_async()
|
||||||
|
try:
|
||||||
|
yield {}
|
||||||
|
finally:
|
||||||
|
# Shutdown: cleanup MCP clients and bridge server
|
||||||
|
await self.shutdown_async()
|
||||||
|
|
||||||
|
self.mcp_server = FastMCP(
|
||||||
|
name="mcp-forge",
|
||||||
|
instructions="""Execute Python code in isolated containers with MCP tool injection.
|
||||||
|
|
||||||
|
TOOL INJECTION:
|
||||||
|
- Pass tool names in mcp_tools parameter to make them callable as Python functions
|
||||||
|
- Example: execute_python(code='docs = browse_documents(page_size=5)', mcp_tools=['browse_documents'])
|
||||||
|
- Injected functions have the same signature as the MCP tool's input schema
|
||||||
|
|
||||||
|
RETURN VALUES:
|
||||||
|
- Most tools return a dict/list (already parsed from JSON)
|
||||||
|
- Some tools return strings that need json.loads() if JSON content expected
|
||||||
|
- Common patterns:
|
||||||
|
* Direct data: result = browse_records(page=1, page_size=10) # returns list of dicts
|
||||||
|
* With wrapper: result.get('result', []) if isinstance(result, dict) else result
|
||||||
|
* String response: json.loads(tool_result) if isinstance(tool_result, str) else tool_result
|
||||||
|
|
||||||
|
DISCOVERY:
|
||||||
|
- Use resource mcp://forge/tools/available to list all injectable tools and their schemas
|
||||||
|
- Schemas show exact parameter names, types, and return structure for each tool
|
||||||
|
|
||||||
|
STATEFUL EXECUTION:
|
||||||
|
- Provide session_id to persist variables across multiple execute_python calls
|
||||||
|
- Variables, imports, and state are preserved within the same session
|
||||||
|
|
||||||
|
EXAMPLE:
|
||||||
|
```python
|
||||||
|
import json
|
||||||
|
# Tool returns parsed data directly
|
||||||
|
records = browse_records(page=1, page_size=100, document_id=1)
|
||||||
|
# Handle flexible return types
|
||||||
|
data = records.get('result', []) if isinstance(records, dict) else records
|
||||||
|
# Process results
|
||||||
|
for record in data:
|
||||||
|
print(record.get('content', ''))
|
||||||
|
```""",
|
||||||
|
lifespan=mcp_forge_lifespan
|
||||||
|
)
|
||||||
|
|
||||||
|
# Initialize all components in dependency order
|
||||||
|
self._init_security()
|
||||||
|
self._init_podman()
|
||||||
|
self._init_mcp_clients()
|
||||||
|
self._init_backends()
|
||||||
|
self._init_builder()
|
||||||
|
self._init_tools()
|
||||||
|
self._init_resources()
|
||||||
|
|
||||||
|
logger.info("MCP Forge Server initialized successfully")
|
||||||
|
|
||||||
|
async def initialize_async(self) -> None:
|
||||||
|
"""Initialize async components (MCP clients and bridge server)."""
|
||||||
|
try:
|
||||||
|
await self.client_manager.initialize()
|
||||||
|
logger.info(f"Initialized {len(self.client_manager.clients)} MCP client(s)")
|
||||||
|
|
||||||
|
# Start bridge server (now that event loop is running)
|
||||||
|
await self.bridge_server.start()
|
||||||
|
logger.info("Bridge server started")
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"Async initialization failed: {e}. Tool injection may not be available.")
|
||||||
|
|
||||||
|
async def shutdown_async(self) -> None:
|
||||||
|
"""Shutdown async components (MCP clients and bridge server)."""
|
||||||
|
try:
|
||||||
|
# Stop bridge server
|
||||||
|
await self.bridge_server.stop()
|
||||||
|
logger.info("Bridge server stopped")
|
||||||
|
|
||||||
|
# Shutdown MCP clients
|
||||||
|
await self.client_manager.shutdown()
|
||||||
|
logger.info("MCP clients shut down")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Shutdown error: {e}")
|
||||||
|
logger.info("MCP clients shut down")
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"MCP client shutdown failed: {e}")
|
||||||
|
|
||||||
|
def _init_security(self) -> None:
|
||||||
|
"""Initialize security components."""
|
||||||
|
self.audit_logger = AuditLogger(
|
||||||
|
log_path=self.config.security.audit_log
|
||||||
|
)
|
||||||
|
self.operation_validator = OperationValidator(
|
||||||
|
config=self.config.security
|
||||||
|
)
|
||||||
|
logger.debug("Security components initialized")
|
||||||
|
|
||||||
|
def _init_podman(self) -> None:
|
||||||
|
"""Initialize Podman client and container manager."""
|
||||||
|
self.podman_client = PodmanClient(
|
||||||
|
socket_path=self.config.server.podman_socket,
|
||||||
|
validator=self.operation_validator,
|
||||||
|
audit_logger=self.audit_logger
|
||||||
|
)
|
||||||
|
self.container_manager = SecureContainerManager(
|
||||||
|
podman_client=self.podman_client,
|
||||||
|
validator=self.operation_validator,
|
||||||
|
audit_logger=self.audit_logger
|
||||||
|
)
|
||||||
|
logger.debug("Podman components initialized")
|
||||||
|
|
||||||
|
def _init_mcp_clients(self) -> None:
|
||||||
|
"""Initialize MCP client manager and bridge server."""
|
||||||
|
from pathlib import Path
|
||||||
|
# Get MCP tools config from forge config and convert to dict
|
||||||
|
mcp_tools_config = getattr(self.config, 'mcp_tools', {})
|
||||||
|
# Convert Pydantic models to dict format expected by MCPClientManager
|
||||||
|
config_dict = {
|
||||||
|
name: tool_config.model_dump()
|
||||||
|
for name, tool_config in mcp_tools_config.items()
|
||||||
|
}
|
||||||
|
self.client_manager = MCPClientManager(
|
||||||
|
config=config_dict
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create bridge server for tool injection
|
||||||
|
self.bridge_server = ToolBridgeServer(
|
||||||
|
socket_path=Path("/tmp/mcp-forge-bridge.sock"),
|
||||||
|
client_manager=self.client_manager,
|
||||||
|
audit_logger=self.audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create tool injection code generator
|
||||||
|
self.injection_generator = ToolInjectionGenerator(
|
||||||
|
client_manager=self.client_manager
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug("MCP client components initialized")
|
||||||
|
|
||||||
|
def _init_backends(self) -> None:
|
||||||
|
"""Initialize execution backends."""
|
||||||
|
from ..security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
# Simple backend for stateless execution
|
||||||
|
self.simple_backend = SimpleBackend(
|
||||||
|
container_manager=self.container_manager,
|
||||||
|
audit_logger=self.audit_logger,
|
||||||
|
config=self.config
|
||||||
|
)
|
||||||
|
|
||||||
|
# Kernel manager with proper resource limits
|
||||||
|
try:
|
||||||
|
resource_limits = ResourceLimits(
|
||||||
|
memory=self.config.execution.max_memory,
|
||||||
|
timeout=self.config.execution.max_timeout,
|
||||||
|
storage="10g",
|
||||||
|
cpu_quota=100000 # 1 CPU
|
||||||
|
)
|
||||||
|
self.kernel_manager = JupyterKernelManager(
|
||||||
|
container_manager=self.container_manager,
|
||||||
|
image="python:3.11",
|
||||||
|
resource_limits=resource_limits
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
# Use a mock if initialization fails (e.g., in tests)
|
||||||
|
from unittest.mock import Mock
|
||||||
|
self.kernel_manager = Mock()
|
||||||
|
|
||||||
|
# Session manager for stateful execution
|
||||||
|
self.session_manager = SessionManager(
|
||||||
|
kernel_manager=self.kernel_manager,
|
||||||
|
audit_logger=self.audit_logger,
|
||||||
|
config=self.config.sessions
|
||||||
|
)
|
||||||
|
|
||||||
|
# Jupyter backend for stateful execution
|
||||||
|
self.jupyter_backend = JupyterBackend(
|
||||||
|
container_manager=self.container_manager,
|
||||||
|
audit_logger=self.audit_logger,
|
||||||
|
config=self.config
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.debug("Execution backends initialized")
|
||||||
|
|
||||||
|
def _init_builder(self) -> None:
|
||||||
|
"""Initialize environment builder."""
|
||||||
|
# EnvironmentBuilder creates its own sub-components (PackageValidator, UVInstaller, ImageBuilder)
|
||||||
|
# so we just need to create it and it will handle the rest
|
||||||
|
self.environment_builder = EnvironmentBuilder(
|
||||||
|
config=self.config.environment_builder,
|
||||||
|
podman_client=self.podman_client,
|
||||||
|
audit_logger=self.audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
# Store references to sub-components for direct access if needed
|
||||||
|
self.package_validator = self.environment_builder.package_validator
|
||||||
|
self.uv_installer = self.environment_builder.uv_installer
|
||||||
|
self.image_builder = self.environment_builder.image_builder
|
||||||
|
|
||||||
|
logger.debug("Environment builder initialized")
|
||||||
|
|
||||||
|
def _init_tools(self) -> None:
|
||||||
|
"""Register MCP tools."""
|
||||||
|
# Create tool instances
|
||||||
|
self.execute_python_tool = ExecutePythonTool(
|
||||||
|
simple_backend=self.simple_backend,
|
||||||
|
jupyter_backend=self.jupyter_backend,
|
||||||
|
client_manager=self.client_manager,
|
||||||
|
bridge_server=self.bridge_server,
|
||||||
|
injection_generator=self.injection_generator,
|
||||||
|
config=self.config
|
||||||
|
)
|
||||||
|
|
||||||
|
self.document_state_tool = DocumentStateTool(
|
||||||
|
session_manager=self.session_manager
|
||||||
|
)
|
||||||
|
|
||||||
|
self.build_environment_tool = BuildEnvironmentTool(
|
||||||
|
environment_builder=self.environment_builder,
|
||||||
|
audit_logger=self.audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
# Register tools with FastMCP using @tool() decorator
|
||||||
|
@self.mcp_server.tool()
|
||||||
|
async def execute_python(
|
||||||
|
code: str,
|
||||||
|
mcp_tools: Optional[List[str]] = None,
|
||||||
|
session_id: Optional[str] = None,
|
||||||
|
backend: Optional[str] = None,
|
||||||
|
timeout: Optional[int] = None,
|
||||||
|
custom_image: Optional[str] = None,
|
||||||
|
environment: Optional[str] = None
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Execute Python code in isolated container with injected MCP tools as callable functions.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
code: Python code to execute. Injected tools available as direct function calls.
|
||||||
|
mcp_tools: Tool names to inject (e.g., ["browse_documents", "search_records"]).
|
||||||
|
- Get available tools: resource mcp://forge/tools/available
|
||||||
|
- Injected functions match MCP tool input schemas exactly
|
||||||
|
- Return values are typically dicts/lists (pre-parsed) or strings
|
||||||
|
- Handle flexible returns: data.get('result') if isinstance(data, dict) else data
|
||||||
|
session_id: Session ID for stateful execution (variables persist across calls)
|
||||||
|
backend: 'simple' (stateless) or 'jupyter' (stateful). Auto-selected if omitted.
|
||||||
|
timeout: Max execution time in seconds (default: 300)
|
||||||
|
custom_image: Custom container image (e.g., 'mcp-forge/custom:my-env')
|
||||||
|
environment: Template environment (e.g., 'datascience', 'ml-basic')
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
JSON with success, stdout, stderr, result, execution_time, available_tools
|
||||||
|
|
||||||
|
Examples:
|
||||||
|
# Basic tool injection
|
||||||
|
execute_python(
|
||||||
|
code='docs = browse_documents(page_size=5); print(docs)',
|
||||||
|
mcp_tools=['browse_documents']
|
||||||
|
)
|
||||||
|
|
||||||
|
# Handling flexible return types
|
||||||
|
execute_python(
|
||||||
|
code='''
|
||||||
|
import json
|
||||||
|
result = browse_records(page=1, page_size=100)
|
||||||
|
# Handle dict wrapper or direct list
|
||||||
|
records = result.get("result", []) if isinstance(result, dict) else result
|
||||||
|
print(f"Found {len(records)} records")
|
||||||
|
''',
|
||||||
|
mcp_tools=['browse_records']
|
||||||
|
)
|
||||||
|
|
||||||
|
# Stateful multi-step processing
|
||||||
|
execute_python(
|
||||||
|
code='data = search_records(text="query", limit=100); count = len(data)',
|
||||||
|
mcp_tools=['search_records'],
|
||||||
|
session_id='analysis-1'
|
||||||
|
)
|
||||||
|
execute_python(
|
||||||
|
code='print(f"Previous count: {count}"); processed = [x for x in data if ...]',
|
||||||
|
session_id='analysis-1' # 'data' and 'count' still available
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
arguments = {
|
||||||
|
"code": code,
|
||||||
|
"mcp_tools": mcp_tools or [],
|
||||||
|
"session_id": session_id,
|
||||||
|
"backend": backend,
|
||||||
|
"timeout": timeout,
|
||||||
|
"custom_image": custom_image,
|
||||||
|
"environment": environment
|
||||||
|
}
|
||||||
|
result = await self.execute_python_tool.execute(arguments)
|
||||||
|
# FastMCP expects string results, convert TextContent list to string
|
||||||
|
if isinstance(result, list):
|
||||||
|
return "\n".join(str(item.text) if hasattr(item, 'text') else str(item) for item in result)
|
||||||
|
return str(result)
|
||||||
|
|
||||||
|
@self.mcp_server.tool()
|
||||||
|
async def document_state(
|
||||||
|
session_id: str,
|
||||||
|
variables: Dict[str, str],
|
||||||
|
note: Optional[str] = None,
|
||||||
|
clear: bool = False
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Document variables in a stateful session for tracking computation state.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session ID to document
|
||||||
|
variables: Dict of variable names to their descriptions/values
|
||||||
|
note: Optional note about this state snapshot
|
||||||
|
clear: If True, clear all previous documentation for this session
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
JSON confirmation with session_id and documented variables
|
||||||
|
"""
|
||||||
|
arguments = {
|
||||||
|
"session_id": session_id,
|
||||||
|
"variables": variables,
|
||||||
|
"note": note,
|
||||||
|
"clear": clear
|
||||||
|
}
|
||||||
|
result = await self.document_state_tool.execute(arguments)
|
||||||
|
if isinstance(result, list):
|
||||||
|
return "\n".join(str(item.text) if hasattr(item, 'text') else str(item) for item in result)
|
||||||
|
return str(result)
|
||||||
|
|
||||||
|
@self.mcp_server.tool()
|
||||||
|
async def build_custom_environment(
|
||||||
|
name: str,
|
||||||
|
packages: List[str],
|
||||||
|
base_image: Optional[str] = None,
|
||||||
|
python_version: Optional[str] = None,
|
||||||
|
description: Optional[str] = None
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Build a custom container image with Python packages for repeated use.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
name: Image name (e.g., 'my-ml-env')
|
||||||
|
packages: Python packages to install (e.g., ['numpy', 'pandas', 'scikit-learn'])
|
||||||
|
base_image: Base image (default: python:3.11-slim)
|
||||||
|
python_version: Python version (e.g., '3.11', '3.12')
|
||||||
|
description: Human-readable description of the environment
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
JSON with image name, tag, and build status
|
||||||
|
"""
|
||||||
|
arguments = {
|
||||||
|
"name": name,
|
||||||
|
"packages": packages,
|
||||||
|
"base_image": base_image,
|
||||||
|
"python_version": python_version,
|
||||||
|
"description": description
|
||||||
|
}
|
||||||
|
result = await self.build_environment_tool.execute(arguments)
|
||||||
|
if isinstance(result, list):
|
||||||
|
return "\n".join(str(item.text) if hasattr(item, 'text') else str(item) for item in result)
|
||||||
|
return str(result)
|
||||||
|
|
||||||
|
logger.debug("MCP tools registered")
|
||||||
|
|
||||||
|
def _init_resources(self) -> None:
|
||||||
|
"""Register MCP resources."""
|
||||||
|
# Create resource handler
|
||||||
|
self.resource_handler = ResourceHandler(
|
||||||
|
client_manager=self.client_manager,
|
||||||
|
session_manager=self.session_manager,
|
||||||
|
environment_builder=self.environment_builder,
|
||||||
|
config=self.config
|
||||||
|
)
|
||||||
|
|
||||||
|
# Register resources with FastMCP using @resource() decorator
|
||||||
|
@self.mcp_server.resource("mcp://forge/tools/available")
|
||||||
|
async def get_available_tools() -> str:
|
||||||
|
"""
|
||||||
|
List MCP tools available for injection via execute_python.
|
||||||
|
|
||||||
|
Returns JSON array with tool names, descriptions, and schemas from all
|
||||||
|
connected MCP servers (e.g., rag-mcp, filesystem, web-search).
|
||||||
|
Use this to discover which tools you can pass to execute_python's mcp_tools parameter.
|
||||||
|
"""
|
||||||
|
import json
|
||||||
|
tools = await self.client_manager.list_all_tools()
|
||||||
|
return json.dumps(tools, indent=2)
|
||||||
|
|
||||||
|
@self.mcp_server.resource("mcp://forge/sessions/list")
|
||||||
|
async def get_sessions_list() -> str:
|
||||||
|
"""
|
||||||
|
List active stateful execution sessions.
|
||||||
|
|
||||||
|
Returns JSON array with session IDs and metadata. Sessions persist variables
|
||||||
|
across execute_python calls when session_id parameter is provided.
|
||||||
|
"""
|
||||||
|
import json
|
||||||
|
sessions = self.session_manager.list_sessions()
|
||||||
|
return json.dumps(sessions, indent=2)
|
||||||
|
|
||||||
|
@self.mcp_server.resource("mcp://forge/environments/list")
|
||||||
|
async def get_environments_list() -> str:
|
||||||
|
"""
|
||||||
|
List available container environment templates.
|
||||||
|
|
||||||
|
Returns JSON with template names and installed packages (e.g., 'datascience',
|
||||||
|
'ml-basic'). Use template names in execute_python's environment parameter.
|
||||||
|
"""
|
||||||
|
import json
|
||||||
|
templates = self.environment_builder.list_templates()
|
||||||
|
return json.dumps(templates, indent=2)
|
||||||
|
|
||||||
|
@self.mcp_server.resource("mcp://forge/environment/info")
|
||||||
|
async def get_environment_info() -> str:
|
||||||
|
"""
|
||||||
|
Get execution environment limits and defaults.
|
||||||
|
|
||||||
|
Returns JSON with default_backend, timeout limits, memory limits.
|
||||||
|
Useful for understanding constraints before calling execute_python.
|
||||||
|
"""
|
||||||
|
import json
|
||||||
|
info = {
|
||||||
|
"default_backend": self.config.execution.default_backend,
|
||||||
|
"default_timeout": self.config.execution.default_timeout,
|
||||||
|
"max_timeout": self.config.execution.max_timeout,
|
||||||
|
"default_memory": self.config.execution.default_memory,
|
||||||
|
"max_memory": self.config.execution.max_memory
|
||||||
|
}
|
||||||
|
return json.dumps(info, indent=2)
|
||||||
|
|
||||||
|
logger.debug("MCP resources registered")
|
||||||
|
|
||||||
|
def run(self, transport: str = "stdio", host: str = "localhost", port: int = 3000) -> None:
|
||||||
|
"""
|
||||||
|
Run server with specified transport.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
transport: Transport type ("stdio", "http", or "sse")
|
||||||
|
host: Host address for HTTP/SSE transport
|
||||||
|
port: Port number for HTTP/SSE transport
|
||||||
|
"""
|
||||||
|
if transport == "stdio":
|
||||||
|
logger.info("Starting MCP Forge Server on stdio")
|
||||||
|
else:
|
||||||
|
logger.info(f"Starting MCP Forge Server on {transport.upper()} transport at {host}:{port}")
|
||||||
|
|
||||||
|
# Bridge server will be started by FastMCP lifespan handler
|
||||||
|
# MCP clients will be initialized by FastMCP lifespan handler
|
||||||
|
|
||||||
|
try:
|
||||||
|
# FastMCP.run() handles all transport types (blocking call)
|
||||||
|
if transport == "stdio":
|
||||||
|
self.mcp_server.run()
|
||||||
|
elif transport == "http":
|
||||||
|
self.mcp_server.run(transport="http", host=host, port=port)
|
||||||
|
elif transport == "sse":
|
||||||
|
self.mcp_server.run(transport="sse", host=host, port=port)
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Unsupported transport: {transport}")
|
||||||
|
finally:
|
||||||
|
# Shutdown is async, so we need to run it in a new event loop
|
||||||
|
import asyncio
|
||||||
|
asyncio.run(self.shutdown())
|
||||||
|
|
||||||
|
async def shutdown(self) -> None:
|
||||||
|
"""
|
||||||
|
Shutdown server and cleanup resources.
|
||||||
|
|
||||||
|
Performs graceful shutdown of all components:
|
||||||
|
- Stops MCP client connections
|
||||||
|
- Stops bridge server
|
||||||
|
- Cleans up active sessions
|
||||||
|
- Logs shutdown event
|
||||||
|
"""
|
||||||
|
logger.info("Shutting down MCP Forge Server")
|
||||||
|
|
||||||
|
# Shutdown MCP clients
|
||||||
|
await self.client_manager.shutdown()
|
||||||
|
|
||||||
|
# Stop bridge server
|
||||||
|
await self.bridge_server.stop()
|
||||||
|
|
||||||
|
# Log shutdown
|
||||||
|
from ..security.audit import AuditEventType, AuditSeverity
|
||||||
|
self.audit_logger.log(
|
||||||
|
event_type=AuditEventType.SESSION_DESTROY,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message="Server shutdown initiated",
|
||||||
|
details={"reason": "graceful_shutdown"}
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info("MCP Forge Server shutdown complete")
|
||||||
|
|
||||||
|
|
||||||
|
async def main(config: Optional[ForgeConfig] = None):
|
||||||
|
"""
|
||||||
|
Main entry point for MCP Forge Server.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
config: Forge configuration (optional, will use defaults if not provided)
|
||||||
|
"""
|
||||||
|
# Setup logging
|
||||||
|
logging.basicConfig(
|
||||||
|
level=logging.INFO,
|
||||||
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||||
|
)
|
||||||
|
|
||||||
|
# Use provided config or create a mock one for testing
|
||||||
|
if config is None:
|
||||||
|
# Create a basic mock config for testing
|
||||||
|
from unittest.mock import Mock
|
||||||
|
config = Mock(spec=ForgeConfig)
|
||||||
|
logger.warning("Using mock configuration - not suitable for production")
|
||||||
|
|
||||||
|
# Create and run server
|
||||||
|
server = ForgeServer(config=config)
|
||||||
|
server.run()
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
asyncio.run(main())
|
||||||
11
src/mcp_forge/server/tools/__init__.py
Normal file
11
src/mcp_forge/server/tools/__init__.py
Normal file
|
|
@ -0,0 +1,11 @@
|
||||||
|
"""MCP Forge Server Tools."""
|
||||||
|
|
||||||
|
from .execute_python import ExecutePythonTool
|
||||||
|
from .document_state import DocumentStateTool
|
||||||
|
from .build_environment import BuildEnvironmentTool
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"ExecutePythonTool",
|
||||||
|
"DocumentStateTool",
|
||||||
|
"BuildEnvironmentTool",
|
||||||
|
]
|
||||||
166
src/mcp_forge/server/tools/build_environment.py
Normal file
166
src/mcp_forge/server/tools/build_environment.py
Normal file
|
|
@ -0,0 +1,166 @@
|
||||||
|
"""Build Custom Environment Tool - MCP tool for building custom container images."""
|
||||||
|
|
||||||
|
from mcp.types import Tool, TextContent
|
||||||
|
from typing import List
|
||||||
|
import json
|
||||||
|
import re
|
||||||
|
|
||||||
|
|
||||||
|
class BuildEnvironmentTool:
|
||||||
|
"""MCP tool for building custom Python environments with specified packages."""
|
||||||
|
|
||||||
|
def __init__(self, environment_builder, audit_logger):
|
||||||
|
"""
|
||||||
|
Initialize Build Environment Tool.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
environment_builder: Environment builder for creating images
|
||||||
|
audit_logger: Audit logger for logging build operations
|
||||||
|
"""
|
||||||
|
self.builder = environment_builder
|
||||||
|
self.audit_logger = audit_logger
|
||||||
|
|
||||||
|
def get_tool_definition(self) -> Tool:
|
||||||
|
"""
|
||||||
|
Return MCP tool definition.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Tool definition with full schema
|
||||||
|
"""
|
||||||
|
return Tool(
|
||||||
|
name="build_custom_environment",
|
||||||
|
description="Build a custom container image with specified Python packages. This is the only way to install additional packages - pip install is not allowed during code execution for security reasons.",
|
||||||
|
inputSchema={
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"name": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Name for the custom environment (alphanumeric + hyphens only)",
|
||||||
|
"pattern": "^[a-z0-9-]+$"
|
||||||
|
},
|
||||||
|
"packages": {
|
||||||
|
"type": "array",
|
||||||
|
"items": {"type": "string"},
|
||||||
|
"description": "List of Python package specifications (e.g., 'numpy>=1.24.0')",
|
||||||
|
"minItems": 1
|
||||||
|
},
|
||||||
|
"base_image": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Base image to build from (default: 'python:3.11')"
|
||||||
|
},
|
||||||
|
"python_version": {
|
||||||
|
"type": "string",
|
||||||
|
"enum": ["3.11", "3.12"],
|
||||||
|
"description": "Python version (default: '3.11')"
|
||||||
|
},
|
||||||
|
"description": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Description of this environment"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["name", "packages"]
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
async def execute(self, arguments: dict) -> List[TextContent]:
|
||||||
|
"""
|
||||||
|
Build custom environment.
|
||||||
|
|
||||||
|
Validates packages, builds image, scans for vulnerabilities.
|
||||||
|
Returns build result with image name and installed packages.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
arguments: Tool arguments containing name, packages, and options
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List containing single TextContent with JSON result
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If arguments are invalid
|
||||||
|
KeyError: If required arguments missing
|
||||||
|
"""
|
||||||
|
# Validate arguments
|
||||||
|
self._validate_arguments(arguments)
|
||||||
|
|
||||||
|
# Extract arguments
|
||||||
|
name = arguments["name"]
|
||||||
|
packages = arguments["packages"]
|
||||||
|
base_image = arguments.get("base_image")
|
||||||
|
python_version = arguments.get("python_version")
|
||||||
|
description = arguments.get("description")
|
||||||
|
|
||||||
|
# Build environment
|
||||||
|
build_result = await self.builder.build_custom_environment(
|
||||||
|
name=name,
|
||||||
|
packages=packages,
|
||||||
|
base_image=base_image,
|
||||||
|
python_version=python_version,
|
||||||
|
description=description
|
||||||
|
)
|
||||||
|
|
||||||
|
# Log to audit
|
||||||
|
self.audit_logger.log_environment_build(
|
||||||
|
name=name,
|
||||||
|
packages=packages,
|
||||||
|
success=build_result.success,
|
||||||
|
image_name=build_result.image_name,
|
||||||
|
build_time=build_result.build_time
|
||||||
|
)
|
||||||
|
|
||||||
|
# Format response
|
||||||
|
response = {
|
||||||
|
"success": build_result.success,
|
||||||
|
"image_name": build_result.image_name,
|
||||||
|
"image_id": build_result.image_id,
|
||||||
|
"build_time": build_result.build_time,
|
||||||
|
"installed_packages": build_result.installed_packages,
|
||||||
|
"cache_hit": build_result.cache_hit
|
||||||
|
}
|
||||||
|
|
||||||
|
# Add error if build failed
|
||||||
|
if not build_result.success and hasattr(build_result, "error"):
|
||||||
|
response["error"] = build_result.error
|
||||||
|
|
||||||
|
return [TextContent(
|
||||||
|
type="text",
|
||||||
|
text=json.dumps(response, indent=2)
|
||||||
|
)]
|
||||||
|
|
||||||
|
def _validate_arguments(self, arguments: dict) -> None:
|
||||||
|
"""
|
||||||
|
Validate tool arguments.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
arguments: Tool arguments to validate
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If arguments are invalid
|
||||||
|
KeyError: If required arguments missing
|
||||||
|
"""
|
||||||
|
# Check required arguments
|
||||||
|
if "name" not in arguments:
|
||||||
|
raise KeyError("Required argument 'name' is missing")
|
||||||
|
|
||||||
|
if "packages" not in arguments:
|
||||||
|
raise KeyError("Required argument 'packages' is missing")
|
||||||
|
|
||||||
|
# Validate name format (alphanumeric + hyphens)
|
||||||
|
name = arguments["name"]
|
||||||
|
if not re.match(r'^[a-z0-9-]+$', name):
|
||||||
|
raise ValueError(
|
||||||
|
"Environment name must contain only lowercase letters, numbers, and hyphens"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Validate packages is a list
|
||||||
|
packages = arguments["packages"]
|
||||||
|
if not isinstance(packages, list):
|
||||||
|
raise ValueError("Argument 'packages' must be a list of package specifications")
|
||||||
|
|
||||||
|
# Validate packages is not empty
|
||||||
|
if len(packages) == 0:
|
||||||
|
raise ValueError("Argument 'packages' must contain at least one package")
|
||||||
|
|
||||||
|
# Validate all packages are strings
|
||||||
|
for pkg in packages:
|
||||||
|
if not isinstance(pkg, str):
|
||||||
|
raise ValueError(f"Package specification must be a string, got {type(pkg)}")
|
||||||
132
src/mcp_forge/server/tools/document_state.py
Normal file
132
src/mcp_forge/server/tools/document_state.py
Normal file
|
|
@ -0,0 +1,132 @@
|
||||||
|
"""Document State Tool - MCP tool for documenting session state."""
|
||||||
|
|
||||||
|
from mcp.types import Tool, TextContent
|
||||||
|
from typing import List
|
||||||
|
import json
|
||||||
|
|
||||||
|
|
||||||
|
class DocumentStateTool:
|
||||||
|
"""MCP tool for documenting important variables in stateful sessions."""
|
||||||
|
|
||||||
|
def __init__(self, session_manager):
|
||||||
|
"""
|
||||||
|
Initialize Document State Tool.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_manager: Session manager for accessing session state
|
||||||
|
"""
|
||||||
|
self.session_manager = session_manager
|
||||||
|
|
||||||
|
def get_tool_definition(self) -> Tool:
|
||||||
|
"""
|
||||||
|
Return MCP tool definition.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Tool definition with full schema
|
||||||
|
"""
|
||||||
|
return Tool(
|
||||||
|
name="document_state",
|
||||||
|
description="Document important variables in a stateful session for later retrieval",
|
||||||
|
inputSchema={
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"session_id": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Session ID to document"
|
||||||
|
},
|
||||||
|
"variables": {
|
||||||
|
"type": "object",
|
||||||
|
"description": "Variable name -> description mapping",
|
||||||
|
"additionalProperties": {"type": "string"}
|
||||||
|
},
|
||||||
|
"note": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "General note about session state"
|
||||||
|
},
|
||||||
|
"clear": {
|
||||||
|
"type": "boolean",
|
||||||
|
"description": "Clear existing documentation (default: false)"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["session_id", "variables"]
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
async def execute(self, arguments: dict) -> List[TextContent]:
|
||||||
|
"""
|
||||||
|
Document session state.
|
||||||
|
|
||||||
|
Updates session's documented variables and note.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
arguments: Tool arguments containing session_id, variables, note, clear
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List containing single TextContent with JSON result
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If arguments are invalid or session doesn't exist
|
||||||
|
KeyError: If required arguments missing
|
||||||
|
"""
|
||||||
|
# Validate arguments
|
||||||
|
self._validate_arguments(arguments)
|
||||||
|
|
||||||
|
# Extract arguments
|
||||||
|
session_id = arguments["session_id"]
|
||||||
|
variables = arguments["variables"]
|
||||||
|
note = arguments.get("note")
|
||||||
|
clear = arguments.get("clear", False)
|
||||||
|
|
||||||
|
# Verify session exists
|
||||||
|
if not self.session_manager.session_exists(session_id):
|
||||||
|
raise ValueError(f"Session '{session_id}' not found")
|
||||||
|
|
||||||
|
# Document variables
|
||||||
|
result = await self.session_manager.document_variables(
|
||||||
|
session_id=session_id,
|
||||||
|
variables=variables,
|
||||||
|
note=note,
|
||||||
|
clear=clear
|
||||||
|
)
|
||||||
|
|
||||||
|
# Format response
|
||||||
|
response = {
|
||||||
|
"success": result.get("success", True),
|
||||||
|
"documented_count": result.get("documented_count", len(variables)),
|
||||||
|
"session_id": session_id
|
||||||
|
}
|
||||||
|
|
||||||
|
return [TextContent(
|
||||||
|
type="text",
|
||||||
|
text=json.dumps(response, indent=2)
|
||||||
|
)]
|
||||||
|
|
||||||
|
def _validate_arguments(self, arguments: dict) -> None:
|
||||||
|
"""
|
||||||
|
Validate tool arguments.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
arguments: Tool arguments to validate
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If arguments are invalid
|
||||||
|
KeyError: If required arguments missing
|
||||||
|
"""
|
||||||
|
# Check required arguments
|
||||||
|
if "session_id" not in arguments:
|
||||||
|
raise KeyError("Required argument 'session_id' is missing")
|
||||||
|
|
||||||
|
if "variables" not in arguments:
|
||||||
|
raise KeyError("Required argument 'variables' is missing")
|
||||||
|
|
||||||
|
# Validate variables is a dict
|
||||||
|
variables = arguments["variables"]
|
||||||
|
if not isinstance(variables, dict):
|
||||||
|
raise ValueError("Argument 'variables' must be a dictionary")
|
||||||
|
|
||||||
|
# Validate all variable descriptions are strings
|
||||||
|
for var_name, description in variables.items():
|
||||||
|
if not isinstance(description, str):
|
||||||
|
raise ValueError(
|
||||||
|
f"Variable description for '{var_name}' must be a string, got {type(description)}"
|
||||||
|
)
|
||||||
278
src/mcp_forge/server/tools/execute_python.py
Normal file
278
src/mcp_forge/server/tools/execute_python.py
Normal file
|
|
@ -0,0 +1,278 @@
|
||||||
|
"""Execute Python Tool - MCP tool for executing Python code."""
|
||||||
|
|
||||||
|
from mcp.types import Tool, TextContent
|
||||||
|
from typing import Optional, List
|
||||||
|
import json
|
||||||
|
import logging
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
from ...execution.simple.backend import SimpleBackend
|
||||||
|
from ...execution.jupyter.backend import JupyterBackend
|
||||||
|
from ...mcp.manager import MCPClientManager
|
||||||
|
from ...mcp.bridge import ToolBridgeServer
|
||||||
|
from ...mcp.injection import ToolInjectionGenerator
|
||||||
|
from ...config.schema import ForgeConfig
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class ExecutePythonTool:
|
||||||
|
"""MCP tool for executing Python code in isolated containers."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
simple_backend: SimpleBackend,
|
||||||
|
jupyter_backend: JupyterBackend,
|
||||||
|
client_manager: MCPClientManager,
|
||||||
|
bridge_server: ToolBridgeServer,
|
||||||
|
injection_generator: ToolInjectionGenerator,
|
||||||
|
config: ForgeConfig
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Initialize Execute Python Tool.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
simple_backend: Simple backend for stateless execution
|
||||||
|
jupyter_backend: Jupyter backend for stateful execution
|
||||||
|
client_manager: MCP client manager
|
||||||
|
bridge_server: Tool bridge server for MCP tool injection
|
||||||
|
injection_generator: Tool injection code generator
|
||||||
|
config: Forge configuration
|
||||||
|
"""
|
||||||
|
self.simple_backend = simple_backend
|
||||||
|
self.jupyter_backend = jupyter_backend
|
||||||
|
self.client_manager = client_manager
|
||||||
|
self.bridge_server = bridge_server
|
||||||
|
self.injection_generator = injection_generator
|
||||||
|
self.config = config
|
||||||
|
|
||||||
|
def get_tool_definition(self) -> Tool:
|
||||||
|
"""
|
||||||
|
Return MCP tool definition.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Tool definition with full schema
|
||||||
|
"""
|
||||||
|
return Tool(
|
||||||
|
name="execute_python",
|
||||||
|
description="Execute Python code in an isolated container with MCP tools available",
|
||||||
|
inputSchema={
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"code": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Python code to execute"
|
||||||
|
},
|
||||||
|
"mcp_tools": {
|
||||||
|
"type": "array",
|
||||||
|
"items": {"type": "string"},
|
||||||
|
"description": "List of MCP tool names to inject into the execution environment"
|
||||||
|
},
|
||||||
|
"session_id": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Session ID for stateful execution (omit for stateless execution)"
|
||||||
|
},
|
||||||
|
"backend": {
|
||||||
|
"type": "string",
|
||||||
|
"enum": ["simple", "jupyter"],
|
||||||
|
"description": "Execution backend to use (default: simple, or jupyter if session_id provided)"
|
||||||
|
},
|
||||||
|
"timeout": {
|
||||||
|
"type": "integer",
|
||||||
|
"description": "Maximum execution time in seconds (default: 300)"
|
||||||
|
},
|
||||||
|
"custom_image": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Custom environment image name (e.g., 'mcp-forge/custom:my-ml-env')"
|
||||||
|
},
|
||||||
|
"environment": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Template environment name (e.g., 'datascience', 'ml-basic')"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["code"]
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
async def execute(self, arguments: dict) -> List[TextContent]:
|
||||||
|
"""
|
||||||
|
Execute Python code with MCP tool injection.
|
||||||
|
|
||||||
|
Process:
|
||||||
|
1. Validate arguments
|
||||||
|
2. Determine backend (simple vs jupyter)
|
||||||
|
3. Generate tool injection code if mcp_tools specified
|
||||||
|
4. Execute code via appropriate backend
|
||||||
|
5. Return formatted result
|
||||||
|
|
||||||
|
Args:
|
||||||
|
arguments: Tool arguments containing code and options
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List containing single TextContent with JSON result
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If arguments are invalid
|
||||||
|
RuntimeError: If bridge server not running when tools requested
|
||||||
|
"""
|
||||||
|
logger.info("[execute_python] Starting execution")
|
||||||
|
# Validate arguments
|
||||||
|
self._validate_arguments(arguments)
|
||||||
|
|
||||||
|
# Extract arguments
|
||||||
|
code = arguments["code"]
|
||||||
|
mcp_tools = arguments.get("mcp_tools", [])
|
||||||
|
session_id = arguments.get("session_id")
|
||||||
|
backend_name = self._select_backend(session_id, arguments.get("backend"))
|
||||||
|
timeout = arguments.get("timeout", self.config.execution.default_timeout)
|
||||||
|
custom_image = arguments.get("custom_image")
|
||||||
|
environment = arguments.get("environment")
|
||||||
|
|
||||||
|
logger.info(f"[execute_python] Backend: {backend_name}, Tools: {mcp_tools}, Timeout: {timeout}")
|
||||||
|
|
||||||
|
# Validate bridge server if tools requested
|
||||||
|
if mcp_tools and not self.bridge_server.is_running():
|
||||||
|
raise RuntimeError("MCP bridge server must be running to inject tools")
|
||||||
|
|
||||||
|
# Prepare execution options
|
||||||
|
exec_options = {
|
||||||
|
"timeout": timeout
|
||||||
|
}
|
||||||
|
|
||||||
|
# Handle image specification
|
||||||
|
if custom_image:
|
||||||
|
exec_options["image"] = custom_image
|
||||||
|
elif environment:
|
||||||
|
exec_options["environment"] = environment
|
||||||
|
|
||||||
|
# Prepare tool injection if requested
|
||||||
|
injection_code = None
|
||||||
|
if mcp_tools:
|
||||||
|
logger.info(f"[execute_python] Generating injection code for {len(mcp_tools)} tools")
|
||||||
|
bridge_socket = str(self.bridge_server.socket_path)
|
||||||
|
injection_code = await self.injection_generator.generate_injection_code(
|
||||||
|
tool_names=mcp_tools,
|
||||||
|
bridge_socket_path=bridge_socket
|
||||||
|
)
|
||||||
|
exec_options["injection_code"] = injection_code
|
||||||
|
exec_options["bridge_socket_path"] = bridge_socket
|
||||||
|
logger.info("[execute_python] Injection code generated")
|
||||||
|
|
||||||
|
# Execute via appropriate backend
|
||||||
|
logger.info(f"[execute_python] Executing code via {backend_name} backend...")
|
||||||
|
if backend_name == "jupyter":
|
||||||
|
# Jupyter backend (stateful, requires session_id)
|
||||||
|
if not session_id:
|
||||||
|
raise ValueError("Jupyter backend requires session_id")
|
||||||
|
|
||||||
|
# Prepare code with injection if needed
|
||||||
|
full_code = code
|
||||||
|
if injection_code:
|
||||||
|
full_code = injection_code + "\n\n" + code
|
||||||
|
|
||||||
|
loop = asyncio.get_event_loop()
|
||||||
|
result = await loop.run_in_executor(
|
||||||
|
None,
|
||||||
|
lambda: self.jupyter_backend.execute(
|
||||||
|
code=full_code,
|
||||||
|
session_id=session_id,
|
||||||
|
timeout=timeout,
|
||||||
|
volumes={str(self.bridge_server.socket_path): {"bind": str(self.bridge_server.socket_path), "mode": "rw"}} if mcp_tools else None
|
||||||
|
)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# Simple backend is synchronous - run in executor to avoid blocking event loop
|
||||||
|
logger.info("[execute_python] Calling simple backend in executor...")
|
||||||
|
loop = asyncio.get_event_loop()
|
||||||
|
result = await loop.run_in_executor(
|
||||||
|
None,
|
||||||
|
lambda: self.simple_backend.execute(
|
||||||
|
code=code,
|
||||||
|
**exec_options
|
||||||
|
)
|
||||||
|
)
|
||||||
|
logger.info(f"[execute_python] Execution completed: success={result.success}")
|
||||||
|
|
||||||
|
# Format response
|
||||||
|
response = {
|
||||||
|
"success": result.success,
|
||||||
|
"stdout": result.stdout,
|
||||||
|
"stderr": result.stderr,
|
||||||
|
"result": result.result,
|
||||||
|
"execution_time": result.execution_time
|
||||||
|
}
|
||||||
|
|
||||||
|
# Add optional fields
|
||||||
|
if mcp_tools:
|
||||||
|
response["available_tools"] = mcp_tools
|
||||||
|
|
||||||
|
if session_id:
|
||||||
|
response["session_id"] = session_id
|
||||||
|
|
||||||
|
return [TextContent(
|
||||||
|
type="text",
|
||||||
|
text=json.dumps(response, indent=2)
|
||||||
|
)]
|
||||||
|
|
||||||
|
def _validate_arguments(self, arguments: dict) -> None:
|
||||||
|
"""
|
||||||
|
Validate tool arguments.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
arguments: Tool arguments to validate
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If arguments are invalid
|
||||||
|
KeyError: If required arguments missing
|
||||||
|
"""
|
||||||
|
# Check required arguments
|
||||||
|
if "code" not in arguments:
|
||||||
|
raise KeyError("Required argument 'code' is missing")
|
||||||
|
|
||||||
|
# Validate backend if specified (None is also valid, will use default)
|
||||||
|
if "backend" in arguments:
|
||||||
|
backend = arguments["backend"]
|
||||||
|
if backend is not None and backend not in ("simple", "jupyter"):
|
||||||
|
raise ValueError(f"Invalid backend '{backend}'. Must be 'simple' or 'jupyter'")
|
||||||
|
|
||||||
|
# Validate timeout if specified (None is also valid, will use default)
|
||||||
|
if "timeout" in arguments:
|
||||||
|
timeout = arguments["timeout"]
|
||||||
|
if timeout is not None and (not isinstance(timeout, int) or timeout <= 0):
|
||||||
|
raise ValueError("Timeout must be a positive integer")
|
||||||
|
if timeout is not None and timeout > self.config.execution.max_timeout:
|
||||||
|
raise ValueError(
|
||||||
|
f"Timeout {timeout}s exceeds maximum allowed {self.config.execution.max_timeout}s"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Validate mcp_tools if specified
|
||||||
|
if "mcp_tools" in arguments:
|
||||||
|
mcp_tools = arguments["mcp_tools"]
|
||||||
|
if not isinstance(mcp_tools, list):
|
||||||
|
raise ValueError("mcp_tools must be a list of tool names")
|
||||||
|
if not all(isinstance(tool, str) for tool in mcp_tools):
|
||||||
|
raise ValueError("All tool names in mcp_tools must be strings")
|
||||||
|
|
||||||
|
def _select_backend(
|
||||||
|
self,
|
||||||
|
session_id: Optional[str],
|
||||||
|
backend: Optional[str]
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
Determine which backend to use.
|
||||||
|
|
||||||
|
Rules:
|
||||||
|
- If session_id provided, use jupyter (required for statefulness)
|
||||||
|
- If backend explicitly specified, use that
|
||||||
|
- Otherwise, use default backend from config
|
||||||
|
|
||||||
|
Args:
|
||||||
|
session_id: Session ID if stateful execution requested
|
||||||
|
backend: Explicitly requested backend
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Backend name to use ("simple" or "jupyter")
|
||||||
|
"""
|
||||||
|
if session_id is not None:
|
||||||
|
return "jupyter"
|
||||||
|
return backend or self.config.execution.default_backend
|
||||||
0
tests/__init__.py
Normal file
0
tests/__init__.py
Normal file
346
tests/builder/test_environment_builder.py
Normal file
346
tests/builder/test_environment_builder.py
Normal file
|
|
@ -0,0 +1,346 @@
|
||||||
|
"""Tests for environment builder orchestration."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from pathlib import Path
|
||||||
|
from unittest.mock import Mock, patch
|
||||||
|
from datetime import datetime, timedelta
|
||||||
|
|
||||||
|
from mcp_forge.builder.environment_builder import (
|
||||||
|
EnvironmentBuilder,
|
||||||
|
BuildRateLimiter,
|
||||||
|
)
|
||||||
|
from mcp_forge.builder.package_validator import SecurityError
|
||||||
|
from mcp_forge.builder.image_builder import BuildResult
|
||||||
|
from mcp_forge.config.schema import EnvironmentBuilderConfig
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def builder_config():
|
||||||
|
"""Mock environment builder configuration."""
|
||||||
|
config = Mock(spec=EnvironmentBuilderConfig)
|
||||||
|
config.build_timeout = 600
|
||||||
|
config.max_build_timeout = 1800
|
||||||
|
config.max_image_size = "2g"
|
||||||
|
config.max_packages = 50
|
||||||
|
config.uv_cache_path = Path("/tmp/uv-cache")
|
||||||
|
config.build_rate_limit = {"requests": 5, "period": 3600}
|
||||||
|
config.max_concurrent_builds = 3
|
||||||
|
config.base_images = {"python:3.11-slim": True}
|
||||||
|
config.templates = {
|
||||||
|
"data-science": {
|
||||||
|
"packages": ["numpy", "pandas", "matplotlib"],
|
||||||
|
"description": "Data science environment"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
# Add package_validation config for PackageValidator
|
||||||
|
config.package_validation = Path("/tmp/package-validation.yaml")
|
||||||
|
return config
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_podman_client():
|
||||||
|
"""Mock Podman client."""
|
||||||
|
return Mock(spec=PodmanClient)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_audit_logger():
|
||||||
|
"""Mock audit logger."""
|
||||||
|
return Mock(spec=AuditLogger)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def builder(builder_config, mock_podman_client, mock_audit_logger):
|
||||||
|
"""Environment builder instance."""
|
||||||
|
# Mock the sub-components during initialization
|
||||||
|
with patch('mcp_forge.builder.environment_builder.PackageValidator'), \
|
||||||
|
patch('mcp_forge.builder.environment_builder.UVInstaller'), \
|
||||||
|
patch('mcp_forge.builder.environment_builder.ImageBuilder'):
|
||||||
|
|
||||||
|
builder = EnvironmentBuilder(
|
||||||
|
config=builder_config,
|
||||||
|
podman_client=mock_podman_client,
|
||||||
|
audit_logger=mock_audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
# Replace with mocks for tests
|
||||||
|
builder.package_validator = Mock()
|
||||||
|
builder.uv_installer = Mock()
|
||||||
|
builder.image_builder = Mock()
|
||||||
|
|
||||||
|
return builder
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_custom_environment_success(builder, tmp_path):
|
||||||
|
"""Test successful custom environment build."""
|
||||||
|
# Mock all sub-components
|
||||||
|
builder.package_validator.validate_packages = Mock()
|
||||||
|
builder.uv_installer.create_build_context = Mock(return_value=tmp_path)
|
||||||
|
|
||||||
|
mock_result = BuildResult(
|
||||||
|
success=True,
|
||||||
|
image_name="mcp-forge/custom:test-env",
|
||||||
|
image_id="sha256:abc123",
|
||||||
|
build_time=45.2,
|
||||||
|
size_bytes=500_000_000,
|
||||||
|
cache_hit=False,
|
||||||
|
installed_packages=["numpy==1.24.0"]
|
||||||
|
)
|
||||||
|
builder.image_builder.build_image = Mock(return_value=mock_result)
|
||||||
|
|
||||||
|
result = builder.build_custom_environment(
|
||||||
|
name="test-env",
|
||||||
|
packages=["numpy>=1.24.0"],
|
||||||
|
user_id="user123"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.success is True
|
||||||
|
assert result.image_name == "mcp-forge/custom:test-env"
|
||||||
|
builder.package_validator.validate_packages.assert_called_once()
|
||||||
|
builder.image_builder.build_image.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_validates_package_count(builder):
|
||||||
|
"""Test that build validates package count against maximum."""
|
||||||
|
builder.config.max_packages = 10
|
||||||
|
|
||||||
|
packages = [f"package{i}" for i in range(20)]
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="exceeds maximum"):
|
||||||
|
builder.build_custom_environment(
|
||||||
|
name="test-env",
|
||||||
|
packages=packages,
|
||||||
|
user_id="user123"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_validates_packages_with_validator(builder):
|
||||||
|
"""Test that build uses package validator."""
|
||||||
|
builder.package_validator.validate_packages = Mock(
|
||||||
|
side_effect=SecurityError("Blocked package")
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(SecurityError, match="Blocked package"):
|
||||||
|
builder.build_custom_environment(
|
||||||
|
name="test-env",
|
||||||
|
packages=["forbidden-package"],
|
||||||
|
user_id="user123"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_enforces_rate_limit(builder):
|
||||||
|
"""Test that build enforces rate limiting."""
|
||||||
|
builder.package_validator.validate_packages = Mock()
|
||||||
|
|
||||||
|
# Exhaust rate limit
|
||||||
|
for _ in range(5):
|
||||||
|
builder.rate_limiter.check_rate_limit("user123")
|
||||||
|
|
||||||
|
# Next request should fail
|
||||||
|
with pytest.raises(RuntimeError, match="Rate limit exceeded"):
|
||||||
|
builder.build_custom_environment(
|
||||||
|
name="test-env",
|
||||||
|
packages=["numpy"],
|
||||||
|
user_id="user123"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_enforces_concurrent_builds_limit(builder, tmp_path):
|
||||||
|
"""Test that build enforces concurrent builds limit."""
|
||||||
|
builder.config.max_concurrent_builds = 2
|
||||||
|
builder.package_validator.validate_packages = Mock()
|
||||||
|
builder.uv_installer.create_build_context = Mock(return_value=tmp_path)
|
||||||
|
|
||||||
|
# Simulate 2 active builds
|
||||||
|
builder.active_builds.add("build1")
|
||||||
|
builder.active_builds.add("build2")
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="Maximum concurrent builds"):
|
||||||
|
builder.build_custom_environment(
|
||||||
|
name="test-env",
|
||||||
|
packages=["numpy"],
|
||||||
|
user_id="user123"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_cleans_up_build_context(builder, tmp_path):
|
||||||
|
"""Test that build cleans up build context even on failure."""
|
||||||
|
builder.package_validator.validate_packages = Mock()
|
||||||
|
builder.uv_installer.create_build_context = Mock(return_value=tmp_path)
|
||||||
|
builder.image_builder.build_image = Mock(side_effect=RuntimeError("Build failed"))
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="Build failed"):
|
||||||
|
builder.build_custom_environment(
|
||||||
|
name="test-env",
|
||||||
|
packages=["numpy"],
|
||||||
|
user_id="user123"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Build context should have been registered for cleanup
|
||||||
|
# (actual cleanup would happen in finally block)
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_validates_environment_name(builder):
|
||||||
|
"""Test that build validates environment name format."""
|
||||||
|
with pytest.raises(ValueError, match="must contain only alphanumeric"):
|
||||||
|
builder.build_custom_environment(
|
||||||
|
name="test env!", # Invalid: spaces and special chars
|
||||||
|
packages=["numpy"],
|
||||||
|
user_id="user123"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_from_template_expands_packages(builder, tmp_path):
|
||||||
|
"""Test building from template expands package list."""
|
||||||
|
builder.package_validator.validate_packages = Mock()
|
||||||
|
builder.uv_installer.create_build_context = Mock(return_value=tmp_path)
|
||||||
|
|
||||||
|
mock_result = BuildResult(
|
||||||
|
success=True,
|
||||||
|
image_name="mcp-forge/custom:data-sci",
|
||||||
|
image_id="sha256:abc123",
|
||||||
|
build_time=45.2,
|
||||||
|
size_bytes=500_000_000,
|
||||||
|
cache_hit=False,
|
||||||
|
installed_packages=["numpy==1.24.0", "pandas==2.0.0", "matplotlib==3.7.0"]
|
||||||
|
)
|
||||||
|
builder.image_builder.build_image = Mock(return_value=mock_result)
|
||||||
|
|
||||||
|
result = builder.build_from_template(
|
||||||
|
template_name="data-science",
|
||||||
|
name="data-sci",
|
||||||
|
user_id="user123"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.success is True
|
||||||
|
# Should have called validator with template packages
|
||||||
|
builder.package_validator.validate_packages.assert_called_once()
|
||||||
|
call_args = builder.package_validator.validate_packages.call_args[0][0]
|
||||||
|
assert "numpy" in call_args
|
||||||
|
assert "pandas" in call_args
|
||||||
|
assert "matplotlib" in call_args
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_from_template_with_additional_packages(builder, tmp_path):
|
||||||
|
"""Test building from template with additional packages."""
|
||||||
|
builder.package_validator.validate_packages = Mock()
|
||||||
|
builder.uv_installer.create_build_context = Mock(return_value=tmp_path)
|
||||||
|
|
||||||
|
mock_result = BuildResult(
|
||||||
|
success=True,
|
||||||
|
image_name="mcp-forge/custom:data-sci",
|
||||||
|
image_id="sha256:abc123",
|
||||||
|
build_time=45.2,
|
||||||
|
size_bytes=500_000_000,
|
||||||
|
cache_hit=False,
|
||||||
|
installed_packages=["numpy==1.24.0", "pandas==2.0.0", "matplotlib==3.7.0", "scipy==1.10.0"]
|
||||||
|
)
|
||||||
|
builder.image_builder.build_image = Mock(return_value=mock_result)
|
||||||
|
|
||||||
|
result = builder.build_from_template(
|
||||||
|
template_name="data-science",
|
||||||
|
additional_packages=["scipy"],
|
||||||
|
name="data-sci",
|
||||||
|
user_id="user123"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.success is True
|
||||||
|
call_args = builder.package_validator.validate_packages.call_args[0][0]
|
||||||
|
assert "scipy" in call_args
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_from_template_validates_template_exists(builder):
|
||||||
|
"""Test that template build validates template exists."""
|
||||||
|
with pytest.raises(ValueError, match="Template.*not found"):
|
||||||
|
builder.build_from_template(
|
||||||
|
template_name="nonexistent",
|
||||||
|
user_id="user123"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_list_templates_returns_config_templates(builder):
|
||||||
|
"""Test that list_templates returns configured templates."""
|
||||||
|
templates = builder.list_templates()
|
||||||
|
|
||||||
|
assert "data-science" in templates
|
||||||
|
assert templates["data-science"]["packages"] == ["numpy", "pandas", "matplotlib"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_rate_limiter_allows_within_limit():
|
||||||
|
"""Test that rate limiter allows requests within limit."""
|
||||||
|
limiter = BuildRateLimiter(max_requests=5, period_seconds=3600)
|
||||||
|
|
||||||
|
# Should allow 5 requests
|
||||||
|
for _ in range(5):
|
||||||
|
limiter.check_rate_limit("user123")
|
||||||
|
|
||||||
|
# 6th request should fail
|
||||||
|
with pytest.raises(RuntimeError, match="Rate limit exceeded"):
|
||||||
|
limiter.check_rate_limit("user123")
|
||||||
|
|
||||||
|
|
||||||
|
def test_rate_limiter_cleans_up_old_requests():
|
||||||
|
"""Test that rate limiter cleans up old requests."""
|
||||||
|
limiter = BuildRateLimiter(max_requests=5, period_seconds=1)
|
||||||
|
|
||||||
|
# Make 5 requests
|
||||||
|
for _ in range(5):
|
||||||
|
limiter.check_rate_limit("user123")
|
||||||
|
|
||||||
|
# Wait for period to expire
|
||||||
|
import time
|
||||||
|
time.sleep(1.1)
|
||||||
|
|
||||||
|
# Should allow new request after period
|
||||||
|
limiter.check_rate_limit("user123")
|
||||||
|
|
||||||
|
|
||||||
|
def test_rate_limiter_tracks_per_user():
|
||||||
|
"""Test that rate limiter tracks requests per user."""
|
||||||
|
limiter = BuildRateLimiter(max_requests=2, period_seconds=3600)
|
||||||
|
|
||||||
|
# User1 makes 2 requests
|
||||||
|
limiter.check_rate_limit("user1")
|
||||||
|
limiter.check_rate_limit("user1")
|
||||||
|
|
||||||
|
# User2 should still be allowed
|
||||||
|
limiter.check_rate_limit("user2")
|
||||||
|
limiter.check_rate_limit("user2")
|
||||||
|
|
||||||
|
# Both users now at limit
|
||||||
|
with pytest.raises(RuntimeError):
|
||||||
|
limiter.check_rate_limit("user1")
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError):
|
||||||
|
limiter.check_rate_limit("user2")
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_registers_and_unregisters_active_builds(builder, tmp_path):
|
||||||
|
"""Test that builds are registered and unregistered correctly."""
|
||||||
|
builder.package_validator.validate_packages = Mock()
|
||||||
|
builder.uv_installer.create_build_context = Mock(return_value=tmp_path)
|
||||||
|
|
||||||
|
mock_result = BuildResult(
|
||||||
|
success=True,
|
||||||
|
image_name="mcp-forge/custom:test-env",
|
||||||
|
image_id="sha256:abc123",
|
||||||
|
build_time=45.2,
|
||||||
|
size_bytes=500_000_000,
|
||||||
|
cache_hit=False,
|
||||||
|
installed_packages=["numpy==1.24.0"]
|
||||||
|
)
|
||||||
|
builder.image_builder.build_image = Mock(return_value=mock_result)
|
||||||
|
|
||||||
|
# Before build
|
||||||
|
assert "test-env" not in builder.active_builds
|
||||||
|
|
||||||
|
builder.build_custom_environment(
|
||||||
|
name="test-env",
|
||||||
|
packages=["numpy"],
|
||||||
|
user_id="user123"
|
||||||
|
)
|
||||||
|
|
||||||
|
# After build
|
||||||
|
assert "test-env" not in builder.active_builds
|
||||||
355
tests/builder/test_image_builder.py
Normal file
355
tests/builder/test_image_builder.py
Normal file
|
|
@ -0,0 +1,355 @@
|
||||||
|
"""Tests for Image Builder module."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from pathlib import Path
|
||||||
|
from unittest.mock import Mock
|
||||||
|
|
||||||
|
from mcp_forge.builder.image_builder import ImageBuilder, BuildResult
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.config.schema import EnvironmentBuilderConfig
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_podman_client():
|
||||||
|
"""Mock PodmanClient."""
|
||||||
|
mock = Mock(spec=PodmanClient)
|
||||||
|
# Configure nested mocks for images and containers
|
||||||
|
mock.images = Mock()
|
||||||
|
mock.containers = Mock()
|
||||||
|
return mock
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def builder_config():
|
||||||
|
"""Mock EnvironmentBuilderConfig."""
|
||||||
|
config = Mock(spec=EnvironmentBuilderConfig)
|
||||||
|
config.enabled = True
|
||||||
|
config.uv_cache_path = Path("/tmp/uv_cache")
|
||||||
|
config.max_packages_per_build = 50
|
||||||
|
config.build_timeout = 600
|
||||||
|
config.max_build_timeout = 1800
|
||||||
|
config.max_image_size = "2g"
|
||||||
|
return config
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_audit_logger():
|
||||||
|
"""Mock AuditLogger."""
|
||||||
|
return Mock(spec=AuditLogger)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def builder(mock_podman_client, builder_config, mock_audit_logger):
|
||||||
|
"""ImageBuilder instance."""
|
||||||
|
return ImageBuilder(mock_podman_client, builder_config, mock_audit_logger)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def build_context(tmp_path):
|
||||||
|
"""Create temporary build context."""
|
||||||
|
context = tmp_path / "build_context"
|
||||||
|
context.mkdir()
|
||||||
|
|
||||||
|
# Create Containerfile
|
||||||
|
containerfile = context / "Containerfile"
|
||||||
|
containerfile.write_text("FROM python:3.11-slim\n")
|
||||||
|
|
||||||
|
# Create requirements.txt
|
||||||
|
requirements = context / "requirements.txt"
|
||||||
|
requirements.write_text("numpy>=1.24.0\npandas\n")
|
||||||
|
|
||||||
|
return context
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_image_success(builder, build_context, mock_podman_client):
|
||||||
|
"""Test successful image build."""
|
||||||
|
# Mock successful build
|
||||||
|
mock_image = Mock()
|
||||||
|
mock_image.id = "sha256:abc123"
|
||||||
|
mock_image.attrs = {"Size": 500_000_000}
|
||||||
|
|
||||||
|
mock_podman_client.images.build.return_value = (mock_image, [])
|
||||||
|
mock_podman_client.images.get.return_value = mock_image # Mock images.get() call
|
||||||
|
|
||||||
|
# Mock pip list output
|
||||||
|
mock_container = Mock()
|
||||||
|
mock_container.exec_run.return_value = (
|
||||||
|
0,
|
||||||
|
b'[{"name": "numpy", "version": "1.24.0"}, {"name": "pandas", "version": "2.0.0"}]'
|
||||||
|
)
|
||||||
|
mock_podman_client.containers.run.return_value = mock_container
|
||||||
|
|
||||||
|
packages = ["numpy>=1.24.0", "pandas"]
|
||||||
|
result = builder.build_image(
|
||||||
|
name="test-env",
|
||||||
|
build_context=build_context,
|
||||||
|
base_image="python:3.11-slim",
|
||||||
|
packages=packages
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.success is True
|
||||||
|
assert result.image_name == "mcp-forge/custom:test-env"
|
||||||
|
assert result.image_id == "sha256:abc123"
|
||||||
|
assert result.size_bytes == 500_000_000
|
||||||
|
assert len(result.installed_packages) == 2
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_image_with_custom_timeout(builder, build_context, mock_podman_client):
|
||||||
|
"""Test build with custom timeout."""
|
||||||
|
mock_image = Mock()
|
||||||
|
mock_image.id = "sha256:abc123"
|
||||||
|
mock_image.attrs = {"Size": 100_000_000}
|
||||||
|
|
||||||
|
mock_podman_client.images.build.return_value = (mock_image, [])
|
||||||
|
mock_podman_client.images.get.return_value = mock_image
|
||||||
|
mock_podman_client.containers.run.return_value.exec_run.return_value = (0, b"[]")
|
||||||
|
|
||||||
|
builder.build_image(
|
||||||
|
name="test-env",
|
||||||
|
build_context=build_context,
|
||||||
|
base_image="python:3.11-slim",
|
||||||
|
packages=[],
|
||||||
|
timeout=1200
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify timeout was passed to build
|
||||||
|
call_kwargs = mock_podman_client.images.build.call_args[1]
|
||||||
|
assert call_kwargs["timeout"] == 1200
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_image_validates_timeout_against_max(builder, build_context, builder_config):
|
||||||
|
"""Test that build validates timeout against maximum."""
|
||||||
|
builder_config.max_build_timeout = 1800
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="Timeout 3600 exceeds maximum"):
|
||||||
|
builder.build_image(
|
||||||
|
name="test-env",
|
||||||
|
build_context=build_context,
|
||||||
|
base_image="python:3.11-slim",
|
||||||
|
packages=[],
|
||||||
|
timeout=3600
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_image_uses_default_timeout(builder, build_context, builder_config, mock_podman_client):
|
||||||
|
"""Test that build uses config default timeout when not specified."""
|
||||||
|
builder_config.build_timeout = 600
|
||||||
|
|
||||||
|
mock_image = Mock()
|
||||||
|
mock_image.id = "sha256:abc123"
|
||||||
|
mock_image.attrs = {"Size": 100_000_000}
|
||||||
|
|
||||||
|
mock_podman_client.images.build.return_value = (mock_image, [])
|
||||||
|
mock_podman_client.images.get.return_value = mock_image
|
||||||
|
mock_podman_client.containers.run.return_value.exec_run.return_value = (0, b"[]")
|
||||||
|
|
||||||
|
builder.build_image(
|
||||||
|
name="test-env",
|
||||||
|
build_context=build_context,
|
||||||
|
base_image="python:3.11-slim",
|
||||||
|
packages=[]
|
||||||
|
)
|
||||||
|
|
||||||
|
call_kwargs = mock_podman_client.images.build.call_args[1]
|
||||||
|
assert call_kwargs["timeout"] == 600
|
||||||
|
|
||||||
|
|
||||||
|
def test_generate_tag_valid_name(builder):
|
||||||
|
"""Test tag generation with valid name."""
|
||||||
|
tag = builder.generate_tag("my-env-123")
|
||||||
|
assert tag == "mcp-forge/custom:my-env-123"
|
||||||
|
|
||||||
|
|
||||||
|
def test_generate_tag_invalid_name_raises_error(builder):
|
||||||
|
"""Test that invalid names raise ValueError."""
|
||||||
|
with pytest.raises(ValueError, match="alphanumeric"):
|
||||||
|
builder.generate_tag("my env")
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="alphanumeric"):
|
||||||
|
builder.generate_tag("my_env!")
|
||||||
|
|
||||||
|
|
||||||
|
def test_validate_image_size_within_limit(builder, mock_podman_client, builder_config):
|
||||||
|
"""Test image size validation passes when within limit."""
|
||||||
|
builder_config.max_image_size = "2g"
|
||||||
|
|
||||||
|
mock_image = Mock()
|
||||||
|
mock_image.attrs = {"Size": 1_000_000_000} # 1GB
|
||||||
|
mock_podman_client.images.get.return_value = mock_image
|
||||||
|
|
||||||
|
size = builder.validate_image_size("sha256:abc123")
|
||||||
|
|
||||||
|
assert size == 1_000_000_000
|
||||||
|
|
||||||
|
|
||||||
|
def test_validate_image_size_exceeds_limit(builder, mock_podman_client, builder_config):
|
||||||
|
"""Test image size validation fails when exceeds limit."""
|
||||||
|
builder_config.max_image_size = "1g"
|
||||||
|
|
||||||
|
mock_image = Mock()
|
||||||
|
mock_image.attrs = {"Size": 2_000_000_000} # 2GB
|
||||||
|
mock_podman_client.images.get.return_value = mock_image
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="exceeds maximum"):
|
||||||
|
builder.validate_image_size("sha256:abc123")
|
||||||
|
|
||||||
|
|
||||||
|
def test_extract_installed_packages(builder, mock_podman_client):
|
||||||
|
"""Test extracting installed packages from image."""
|
||||||
|
mock_container = Mock()
|
||||||
|
mock_container.exec_run.return_value = (
|
||||||
|
0,
|
||||||
|
b'[{"name": "numpy", "version": "1.24.0"}, {"name": "pandas", "version": "2.0.0"}]'
|
||||||
|
)
|
||||||
|
mock_podman_client.containers.run.return_value = mock_container
|
||||||
|
|
||||||
|
packages = builder.extract_installed_packages("sha256:abc123")
|
||||||
|
|
||||||
|
assert len(packages) == 2
|
||||||
|
assert "numpy==1.24.0" in packages
|
||||||
|
assert "pandas==2.0.0" in packages
|
||||||
|
|
||||||
|
|
||||||
|
def test_extract_installed_packages_handles_error(builder, mock_podman_client):
|
||||||
|
"""Test that package extraction handles errors gracefully."""
|
||||||
|
mock_container = Mock()
|
||||||
|
mock_container.exec_run.return_value = (1, b"Error")
|
||||||
|
mock_podman_client.containers.run.return_value = mock_container
|
||||||
|
|
||||||
|
packages = builder.extract_installed_packages("sha256:abc123")
|
||||||
|
|
||||||
|
assert packages == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_calculate_cache_hash(builder):
|
||||||
|
"""Test cache hash calculation."""
|
||||||
|
packages = ["numpy>=1.24.0", "pandas==2.0.0", "requests"]
|
||||||
|
|
||||||
|
hash1 = builder.calculate_cache_hash(packages)
|
||||||
|
hash2 = builder.calculate_cache_hash(packages)
|
||||||
|
|
||||||
|
# Same packages should produce same hash
|
||||||
|
assert hash1 == hash2
|
||||||
|
|
||||||
|
# Different packages should produce different hash
|
||||||
|
different_packages = ["numpy>=1.24.0", "scipy"]
|
||||||
|
hash3 = builder.calculate_cache_hash(different_packages)
|
||||||
|
assert hash1 != hash3
|
||||||
|
|
||||||
|
|
||||||
|
def test_calculate_cache_hash_order_independent(builder):
|
||||||
|
"""Test that cache hash is order-independent."""
|
||||||
|
packages1 = ["numpy", "pandas", "scipy"]
|
||||||
|
packages2 = ["scipy", "numpy", "pandas"]
|
||||||
|
|
||||||
|
hash1 = builder.calculate_cache_hash(packages1)
|
||||||
|
hash2 = builder.calculate_cache_hash(packages2)
|
||||||
|
|
||||||
|
# Order shouldn't matter
|
||||||
|
assert hash1 == hash2
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_image_logs_audit_event(builder, build_context, mock_podman_client, mock_audit_logger):
|
||||||
|
"""Test that build logs audit event."""
|
||||||
|
mock_image = Mock()
|
||||||
|
mock_image.id = "sha256:abc123"
|
||||||
|
mock_image.attrs = {"Size": 100_000_000}
|
||||||
|
|
||||||
|
mock_podman_client.images.build.return_value = (mock_image, [])
|
||||||
|
mock_podman_client.images.get.return_value = mock_image
|
||||||
|
mock_podman_client.containers.run.return_value.exec_run.return_value = (0, b"[]")
|
||||||
|
|
||||||
|
builder.build_image(
|
||||||
|
name="test-env",
|
||||||
|
build_context=build_context,
|
||||||
|
base_image="python:3.11-slim",
|
||||||
|
packages=["numpy"]
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify audit log was called
|
||||||
|
assert mock_audit_logger.log.called
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_image_returns_build_time(builder, build_context, mock_podman_client):
|
||||||
|
"""Test that build result includes build time."""
|
||||||
|
mock_image = Mock()
|
||||||
|
mock_image.id = "sha256:abc123"
|
||||||
|
mock_image.attrs = {"Size": 100_000_000}
|
||||||
|
|
||||||
|
mock_podman_client.images.build.return_value = (mock_image, [])
|
||||||
|
mock_podman_client.images.get.return_value = mock_image
|
||||||
|
mock_podman_client.containers.run.return_value.exec_run.return_value = (0, b"[]")
|
||||||
|
|
||||||
|
result = builder.build_image(
|
||||||
|
name="test-env",
|
||||||
|
build_context=build_context,
|
||||||
|
base_image="python:3.11-slim",
|
||||||
|
packages=[]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.build_time > 0
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_image_handles_build_failure(builder, build_context, mock_podman_client):
|
||||||
|
"""Test that build handles Podman build failures."""
|
||||||
|
mock_podman_client.images.build.side_effect = Exception("Build failed")
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="Build failed"):
|
||||||
|
builder.build_image(
|
||||||
|
name="test-env",
|
||||||
|
build_context=build_context,
|
||||||
|
base_image="python:3.11-slim",
|
||||||
|
packages=[]
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_result_to_dict(builder):
|
||||||
|
"""Test BuildResult serialization."""
|
||||||
|
result = BuildResult(
|
||||||
|
success=True,
|
||||||
|
image_name="mcp-forge/custom:test",
|
||||||
|
image_id="sha256:abc123",
|
||||||
|
build_time=10.5,
|
||||||
|
size_bytes=500_000_000,
|
||||||
|
cache_hit=False,
|
||||||
|
installed_packages=["numpy==1.24.0", "pandas==2.0.0"]
|
||||||
|
)
|
||||||
|
|
||||||
|
result_dict = result.to_dict()
|
||||||
|
|
||||||
|
assert result_dict["success"] is True
|
||||||
|
assert result_dict["image_name"] == "mcp-forge/custom:test"
|
||||||
|
assert result_dict["image_id"] == "sha256:abc123"
|
||||||
|
assert result_dict["build_time"] == 10.5
|
||||||
|
assert result_dict["size_bytes"] == 500_000_000
|
||||||
|
assert result_dict["cache_hit"] is False
|
||||||
|
assert len(result_dict["installed_packages"]) == 2
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_image_validates_build_context_exists(builder, tmp_path):
|
||||||
|
"""Test that build validates build context exists."""
|
||||||
|
nonexistent = tmp_path / "nonexistent"
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="Build context does not exist"):
|
||||||
|
builder.build_image(
|
||||||
|
name="test-env",
|
||||||
|
build_context=nonexistent,
|
||||||
|
base_image="python:3.11-slim",
|
||||||
|
packages=[]
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_build_image_validates_containerfile_exists(builder, tmp_path):
|
||||||
|
"""Test that build validates Containerfile exists."""
|
||||||
|
context = tmp_path / "context"
|
||||||
|
context.mkdir()
|
||||||
|
# No Containerfile created
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="Containerfile not found"):
|
||||||
|
builder.build_image(
|
||||||
|
name="test-env",
|
||||||
|
build_context=context,
|
||||||
|
base_image="python:3.11-slim",
|
||||||
|
packages=[]
|
||||||
|
)
|
||||||
271
tests/builder/test_package_validator.py
Normal file
271
tests/builder/test_package_validator.py
Normal file
|
|
@ -0,0 +1,271 @@
|
||||||
|
"""Tests for Package Validator module."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from unittest.mock import Mock
|
||||||
|
|
||||||
|
from mcp_forge.builder.package_validator import (
|
||||||
|
PackageValidator,
|
||||||
|
ApprovalRequiredError,
|
||||||
|
SecurityError
|
||||||
|
)
|
||||||
|
from mcp_forge.config.schema import PackageValidationConfig
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def allowlist_file(tmp_path):
|
||||||
|
"""Create temporary allowlist file."""
|
||||||
|
allowlist = tmp_path / "allowlist.txt"
|
||||||
|
allowlist.write_text("""
|
||||||
|
# Standard data science packages
|
||||||
|
numpy
|
||||||
|
pandas
|
||||||
|
scipy
|
||||||
|
scikit-learn
|
||||||
|
matplotlib
|
||||||
|
|
||||||
|
# Web and API
|
||||||
|
requests
|
||||||
|
httpx
|
||||||
|
aiohttp
|
||||||
|
|
||||||
|
# Utilities
|
||||||
|
pyyaml
|
||||||
|
python-dateutil
|
||||||
|
""".strip())
|
||||||
|
return allowlist
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def blocklist_file(tmp_path):
|
||||||
|
"""Create temporary blocklist file."""
|
||||||
|
blocklist = tmp_path / "blocklist.txt"
|
||||||
|
blocklist.write_text("""
|
||||||
|
# Security concerns
|
||||||
|
os-crypto
|
||||||
|
subprocess-wrapper
|
||||||
|
shell-exec
|
||||||
|
|
||||||
|
# Known malicious
|
||||||
|
malicious-package
|
||||||
|
evil-lib
|
||||||
|
""".strip())
|
||||||
|
return blocklist
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def validation_config(allowlist_file, blocklist_file, tmp_path):
|
||||||
|
"""Mock PackageValidationConfig."""
|
||||||
|
config = Mock(spec=PackageValidationConfig)
|
||||||
|
config.use_allowlist = True
|
||||||
|
config.allowlist_path = allowlist_file
|
||||||
|
config.blocklist_path = blocklist_file
|
||||||
|
config.require_approval_patterns = [
|
||||||
|
"^torch.*", # PyTorch packages
|
||||||
|
"^tensorflow.*", # TensorFlow packages
|
||||||
|
".*-gpu$", # GPU variants
|
||||||
|
]
|
||||||
|
return config
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def validator(validation_config):
|
||||||
|
"""PackageValidator instance."""
|
||||||
|
return PackageValidator(validation_config)
|
||||||
|
|
||||||
|
|
||||||
|
def test_extract_package_name_simple(validator):
|
||||||
|
"""Test extracting package name from simple spec."""
|
||||||
|
assert validator.extract_package_name("numpy") == "numpy"
|
||||||
|
assert validator.extract_package_name("pandas") == "pandas"
|
||||||
|
|
||||||
|
|
||||||
|
def test_extract_package_name_with_version(validator):
|
||||||
|
"""Test extracting package name with version specifiers."""
|
||||||
|
assert validator.extract_package_name("numpy>=1.24.0") == "numpy"
|
||||||
|
assert validator.extract_package_name("pandas==2.0.0") == "pandas"
|
||||||
|
assert validator.extract_package_name("requests<=2.28.0") == "requests"
|
||||||
|
assert validator.extract_package_name("scikit-learn~=1.3.0") == "scikit-learn"
|
||||||
|
|
||||||
|
|
||||||
|
def test_extract_package_name_with_extras(validator):
|
||||||
|
"""Test extracting package name with extras."""
|
||||||
|
assert validator.extract_package_name("requests[security]") == "requests"
|
||||||
|
assert validator.extract_package_name("pandas[excel,sql]") == "pandas"
|
||||||
|
|
||||||
|
|
||||||
|
def test_extract_package_name_complex(validator):
|
||||||
|
"""Test extracting package name from complex specs."""
|
||||||
|
assert validator.extract_package_name("numpy>=1.24.0,<2.0.0") == "numpy"
|
||||||
|
assert validator.extract_package_name("requests[security]>=2.28.0") == "requests"
|
||||||
|
|
||||||
|
|
||||||
|
def test_allowlisted_package_passes(validator):
|
||||||
|
"""Test that allowlisted packages pass validation."""
|
||||||
|
validator.validate_package("numpy")
|
||||||
|
validator.validate_package("pandas>=2.0.0")
|
||||||
|
validator.validate_package("requests[security]")
|
||||||
|
# Should not raise
|
||||||
|
|
||||||
|
|
||||||
|
def test_blocklisted_package_raises_error(validator):
|
||||||
|
"""Test that blocklisted packages raise SecurityError."""
|
||||||
|
with pytest.raises(SecurityError, match="malicious-package"):
|
||||||
|
validator.validate_package("malicious-package")
|
||||||
|
|
||||||
|
with pytest.raises(SecurityError, match="evil-lib"):
|
||||||
|
validator.validate_package("evil-lib>=1.0.0")
|
||||||
|
|
||||||
|
|
||||||
|
def test_unknown_package_with_allowlist_raises_error(validator):
|
||||||
|
"""Test that unknown packages raise error when allowlist is enabled."""
|
||||||
|
with pytest.raises(SecurityError, match="unknown-package"):
|
||||||
|
validator.validate_package("unknown-package")
|
||||||
|
|
||||||
|
|
||||||
|
def test_package_requiring_approval_raises_error(validator):
|
||||||
|
"""Test that packages matching approval patterns raise ApprovalRequiredError."""
|
||||||
|
with pytest.raises(ApprovalRequiredError, match="torch"):
|
||||||
|
validator.validate_package("torch")
|
||||||
|
|
||||||
|
with pytest.raises(ApprovalRequiredError, match="tensorflow"):
|
||||||
|
validator.validate_package("tensorflow-gpu")
|
||||||
|
|
||||||
|
with pytest.raises(ApprovalRequiredError, match="gpu"):
|
||||||
|
validator.validate_package("cupy-gpu")
|
||||||
|
|
||||||
|
|
||||||
|
def test_validate_packages_list(validator):
|
||||||
|
"""Test validating multiple packages at once."""
|
||||||
|
packages = ["numpy>=1.24.0", "pandas", "requests"]
|
||||||
|
validator.validate_packages(packages)
|
||||||
|
# Should not raise
|
||||||
|
|
||||||
|
|
||||||
|
def test_validate_packages_enforces_max_limit(validator):
|
||||||
|
"""Test that validate_packages enforces maximum package count."""
|
||||||
|
packages = ["numpy", "pandas", "scipy", "matplotlib"]
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="Maximum 3 packages"):
|
||||||
|
validator.validate_packages(packages, max_packages=3)
|
||||||
|
|
||||||
|
|
||||||
|
def test_validate_packages_with_mixed_results(validator):
|
||||||
|
"""Test that validation stops at first error."""
|
||||||
|
packages = ["numpy", "malicious-package", "pandas"]
|
||||||
|
|
||||||
|
with pytest.raises(SecurityError, match="malicious-package"):
|
||||||
|
validator.validate_packages(packages)
|
||||||
|
|
||||||
|
|
||||||
|
def test_validate_packages_with_approval_required(validator):
|
||||||
|
"""Test that validation stops at first approval requirement."""
|
||||||
|
packages = ["numpy", "torch", "pandas"]
|
||||||
|
|
||||||
|
with pytest.raises(ApprovalRequiredError, match="torch"):
|
||||||
|
validator.validate_packages(packages)
|
||||||
|
|
||||||
|
|
||||||
|
def test_allowlist_loading(allowlist_file):
|
||||||
|
"""Test that allowlist is loaded correctly from file."""
|
||||||
|
config = Mock(spec=PackageValidationConfig)
|
||||||
|
config.use_allowlist = True
|
||||||
|
config.allowlist_path = allowlist_file
|
||||||
|
config.blocklist_path = None
|
||||||
|
config.require_approval_patterns = []
|
||||||
|
|
||||||
|
validator = PackageValidator(config)
|
||||||
|
|
||||||
|
assert "numpy" in validator.allowlist
|
||||||
|
assert "pandas" in validator.allowlist
|
||||||
|
assert "requests" in validator.allowlist
|
||||||
|
# Comments and empty lines should be ignored
|
||||||
|
assert "# Standard data science packages" not in validator.allowlist
|
||||||
|
|
||||||
|
|
||||||
|
def test_blocklist_loading(blocklist_file):
|
||||||
|
"""Test that blocklist is loaded correctly from file."""
|
||||||
|
config = Mock(spec=PackageValidationConfig)
|
||||||
|
config.use_allowlist = False
|
||||||
|
config.allowlist_path = None
|
||||||
|
config.blocklist_path = blocklist_file
|
||||||
|
config.require_approval_patterns = []
|
||||||
|
|
||||||
|
validator = PackageValidator(config)
|
||||||
|
|
||||||
|
assert "malicious-package" in validator.blocklist
|
||||||
|
assert "evil-lib" in validator.blocklist
|
||||||
|
|
||||||
|
|
||||||
|
def test_disabled_allowlist_allows_all_except_blocklist(blocklist_file):
|
||||||
|
"""Test that disabling allowlist allows any package except blocklisted."""
|
||||||
|
config = Mock(spec=PackageValidationConfig)
|
||||||
|
config.use_allowlist = False
|
||||||
|
config.allowlist_path = None
|
||||||
|
config.blocklist_path = blocklist_file
|
||||||
|
config.require_approval_patterns = []
|
||||||
|
|
||||||
|
validator = PackageValidator(config)
|
||||||
|
|
||||||
|
# Unknown packages should pass
|
||||||
|
validator.validate_package("some-random-package")
|
||||||
|
|
||||||
|
# But blocklisted packages should still fail
|
||||||
|
with pytest.raises(SecurityError, match="malicious-package"):
|
||||||
|
validator.validate_package("malicious-package")
|
||||||
|
|
||||||
|
|
||||||
|
def test_approval_pattern_matching(validator):
|
||||||
|
"""Test that approval patterns match correctly."""
|
||||||
|
# torch* should match
|
||||||
|
with pytest.raises(ApprovalRequiredError):
|
||||||
|
validator.validate_package("torch")
|
||||||
|
|
||||||
|
with pytest.raises(ApprovalRequiredError):
|
||||||
|
validator.validate_package("torchvision")
|
||||||
|
|
||||||
|
# *-gpu$ should match
|
||||||
|
with pytest.raises(ApprovalRequiredError):
|
||||||
|
validator.validate_package("something-gpu")
|
||||||
|
|
||||||
|
|
||||||
|
def test_empty_package_list(validator):
|
||||||
|
"""Test validating empty package list."""
|
||||||
|
validator.validate_packages([])
|
||||||
|
# Should not raise
|
||||||
|
|
||||||
|
|
||||||
|
def test_validate_packages_without_max_limit(validator):
|
||||||
|
"""Test validating many packages without limit."""
|
||||||
|
packages = [f"package{i}" for i in range(100)]
|
||||||
|
|
||||||
|
# Should raise because packages aren't in allowlist
|
||||||
|
with pytest.raises(SecurityError):
|
||||||
|
validator.validate_packages(packages)
|
||||||
|
|
||||||
|
|
||||||
|
def test_case_sensitivity(validator):
|
||||||
|
"""Test that package names are case-sensitive."""
|
||||||
|
# numpy is in allowlist
|
||||||
|
validator.validate_package("numpy")
|
||||||
|
|
||||||
|
# NumPy (different case) should fail
|
||||||
|
with pytest.raises(SecurityError, match="NumPy"):
|
||||||
|
validator.validate_package("NumPy")
|
||||||
|
|
||||||
|
|
||||||
|
def test_whitespace_handling(validator):
|
||||||
|
"""Test that leading/trailing whitespace is handled."""
|
||||||
|
validator.validate_package(" numpy ")
|
||||||
|
validator.validate_package(" pandas>=2.0.0 ")
|
||||||
|
|
||||||
|
|
||||||
|
def test_file_not_found_handling(tmp_path):
|
||||||
|
"""Test handling of missing allowlist/blocklist files."""
|
||||||
|
config = Mock(spec=PackageValidationConfig)
|
||||||
|
config.use_allowlist = True
|
||||||
|
config.allowlist_path = tmp_path / "nonexistent.txt"
|
||||||
|
config.blocklist_path = None
|
||||||
|
config.require_approval_patterns = []
|
||||||
|
|
||||||
|
with pytest.raises(FileNotFoundError):
|
||||||
|
PackageValidator(config)
|
||||||
304
tests/builder/test_uv_installer.py
Normal file
304
tests/builder/test_uv_installer.py
Normal file
|
|
@ -0,0 +1,304 @@
|
||||||
|
"""Tests for UV Package Installer module."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from pathlib import Path
|
||||||
|
from unittest.mock import Mock, patch, mock_open
|
||||||
|
import tempfile
|
||||||
|
import shutil
|
||||||
|
|
||||||
|
from mcp_forge.builder.uv_installer import UVInstaller
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def cache_path(tmp_path):
|
||||||
|
"""Temporary cache directory."""
|
||||||
|
cache_dir = tmp_path / "uv_cache"
|
||||||
|
return cache_dir
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def installer(cache_path):
|
||||||
|
"""UVInstaller instance."""
|
||||||
|
return UVInstaller(cache_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_init_creates_cache_directory(cache_path):
|
||||||
|
"""Test that __init__ creates cache directory."""
|
||||||
|
assert not cache_path.exists()
|
||||||
|
|
||||||
|
installer = UVInstaller(cache_path)
|
||||||
|
|
||||||
|
assert cache_path.exists()
|
||||||
|
assert cache_path.is_dir()
|
||||||
|
|
||||||
|
|
||||||
|
def test_init_with_existing_cache(cache_path):
|
||||||
|
"""Test initialization with existing cache directory."""
|
||||||
|
cache_path.mkdir(parents=True)
|
||||||
|
|
||||||
|
installer = UVInstaller(cache_path)
|
||||||
|
|
||||||
|
assert cache_path.exists()
|
||||||
|
|
||||||
|
|
||||||
|
def test_generate_requirements_single_package(installer):
|
||||||
|
"""Test generating requirements.txt with single package."""
|
||||||
|
packages = ["numpy>=1.24.0"]
|
||||||
|
|
||||||
|
requirements = installer.generate_requirements(packages)
|
||||||
|
|
||||||
|
assert requirements == "numpy>=1.24.0"
|
||||||
|
|
||||||
|
|
||||||
|
def test_generate_requirements_multiple_packages(installer):
|
||||||
|
"""Test generating requirements.txt with multiple packages."""
|
||||||
|
packages = ["numpy>=1.24.0", "pandas==2.0.0", "requests"]
|
||||||
|
|
||||||
|
requirements = installer.generate_requirements(packages)
|
||||||
|
|
||||||
|
lines = requirements.strip().split('\n')
|
||||||
|
assert len(lines) == 3
|
||||||
|
assert "numpy>=1.24.0" in lines
|
||||||
|
assert "pandas==2.0.0" in lines
|
||||||
|
assert "requests" in lines
|
||||||
|
|
||||||
|
|
||||||
|
def test_generate_requirements_empty_list(installer):
|
||||||
|
"""Test generating requirements.txt with empty list."""
|
||||||
|
packages = []
|
||||||
|
|
||||||
|
requirements = installer.generate_requirements(packages)
|
||||||
|
|
||||||
|
assert requirements == ""
|
||||||
|
|
||||||
|
|
||||||
|
def test_generate_containerfile_structure(installer):
|
||||||
|
"""Test Containerfile has correct structure."""
|
||||||
|
base_image = "docker.io/python:3.11-slim"
|
||||||
|
packages = ["numpy>=1.24.0", "pandas"]
|
||||||
|
|
||||||
|
containerfile = installer.generate_containerfile(base_image, packages)
|
||||||
|
|
||||||
|
# Check key sections are present
|
||||||
|
assert f"FROM {base_image}" in containerfile
|
||||||
|
assert "pip install" in containerfile and "uv" in containerfile
|
||||||
|
assert "COPY" in containerfile and "requirements.txt" in containerfile
|
||||||
|
assert "RUN uv pip install" in containerfile
|
||||||
|
assert "WORKDIR" in containerfile
|
||||||
|
|
||||||
|
|
||||||
|
def test_generate_containerfile_with_python_version(installer):
|
||||||
|
"""Test Containerfile generation with specific Python version."""
|
||||||
|
base_image = "docker.io/python:3.12-slim"
|
||||||
|
packages = ["numpy"]
|
||||||
|
|
||||||
|
containerfile = installer.generate_containerfile(
|
||||||
|
base_image, packages, python_version="3.12"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert "3.12" in containerfile or "python:3.12" in base_image
|
||||||
|
|
||||||
|
|
||||||
|
def test_generate_containerfile_uses_requirements(installer):
|
||||||
|
"""Test Containerfile copies and uses requirements.txt."""
|
||||||
|
base_image = "docker.io/python:3.11-slim"
|
||||||
|
packages = ["numpy", "pandas"]
|
||||||
|
|
||||||
|
containerfile = installer.generate_containerfile(base_image, packages)
|
||||||
|
|
||||||
|
# Should copy requirements.txt for layer caching
|
||||||
|
assert "COPY" in containerfile and "requirements.txt" in containerfile
|
||||||
|
assert "-r requirements.txt" in containerfile
|
||||||
|
|
||||||
|
|
||||||
|
def test_generate_containerfile_creates_user(installer):
|
||||||
|
"""Test Containerfile creates non-root user."""
|
||||||
|
base_image = "docker.io/python:3.11-slim"
|
||||||
|
packages = ["numpy"]
|
||||||
|
|
||||||
|
containerfile = installer.generate_containerfile(base_image, packages)
|
||||||
|
|
||||||
|
# Should create user for security
|
||||||
|
assert "useradd" in containerfile.lower() or "USER" in containerfile
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_build_context_creates_directory(installer):
|
||||||
|
"""Test that build context directory is created."""
|
||||||
|
base_image = "docker.io/python:3.11-slim"
|
||||||
|
packages = ["numpy", "pandas"]
|
||||||
|
|
||||||
|
context_path = installer.create_build_context(base_image, packages)
|
||||||
|
|
||||||
|
try:
|
||||||
|
assert context_path.exists()
|
||||||
|
assert context_path.is_dir()
|
||||||
|
finally:
|
||||||
|
if context_path.exists():
|
||||||
|
shutil.rmtree(context_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_build_context_contains_containerfile(installer):
|
||||||
|
"""Test that build context contains Containerfile."""
|
||||||
|
base_image = "docker.io/python:3.11-slim"
|
||||||
|
packages = ["numpy"]
|
||||||
|
|
||||||
|
context_path = installer.create_build_context(base_image, packages)
|
||||||
|
|
||||||
|
try:
|
||||||
|
containerfile_path = context_path / "Containerfile"
|
||||||
|
assert containerfile_path.exists()
|
||||||
|
|
||||||
|
content = containerfile_path.read_text()
|
||||||
|
assert f"FROM {base_image}" in content
|
||||||
|
finally:
|
||||||
|
if context_path.exists():
|
||||||
|
shutil.rmtree(context_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_build_context_contains_requirements(installer):
|
||||||
|
"""Test that build context contains requirements.txt."""
|
||||||
|
base_image = "docker.io/python:3.11-slim"
|
||||||
|
packages = ["numpy>=1.24.0", "pandas"]
|
||||||
|
|
||||||
|
context_path = installer.create_build_context(base_image, packages)
|
||||||
|
|
||||||
|
try:
|
||||||
|
requirements_path = context_path / "requirements.txt"
|
||||||
|
assert requirements_path.exists()
|
||||||
|
|
||||||
|
content = requirements_path.read_text()
|
||||||
|
assert "numpy>=1.24.0" in content
|
||||||
|
assert "pandas" in content
|
||||||
|
finally:
|
||||||
|
if context_path.exists():
|
||||||
|
shutil.rmtree(context_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_build_context_returns_temp_directory(installer):
|
||||||
|
"""Test that build context is in temp directory."""
|
||||||
|
base_image = "docker.io/python:3.11-slim"
|
||||||
|
packages = ["numpy"]
|
||||||
|
|
||||||
|
context_path = installer.create_build_context(base_image, packages)
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Should be in system temp directory
|
||||||
|
temp_dir = Path(tempfile.gettempdir())
|
||||||
|
assert temp_dir in context_path.parents
|
||||||
|
finally:
|
||||||
|
if context_path.exists():
|
||||||
|
shutil.rmtree(context_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_cache_volume_mount_returns_dict(installer):
|
||||||
|
"""Test cache volume mount returns proper dict."""
|
||||||
|
mount_config = installer.get_cache_volume_mount()
|
||||||
|
|
||||||
|
assert isinstance(mount_config, dict)
|
||||||
|
assert "bind" in mount_config
|
||||||
|
assert "mode" in mount_config
|
||||||
|
assert mount_config["mode"] == "rw"
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_cache_volume_mount_includes_cache_path(installer, cache_path):
|
||||||
|
"""Test cache volume mount includes cache path."""
|
||||||
|
mount_config = installer.get_cache_volume_mount()
|
||||||
|
|
||||||
|
# The bind target should reference UV cache location
|
||||||
|
assert "bind" in mount_config
|
||||||
|
assert "/cache" in mount_config["bind"] or "uv" in mount_config["bind"].lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_generate_containerfile_security_practices(installer):
|
||||||
|
"""Test Containerfile follows security best practices."""
|
||||||
|
base_image = "docker.io/python:3.11-slim"
|
||||||
|
packages = ["numpy"]
|
||||||
|
|
||||||
|
containerfile = installer.generate_containerfile(base_image, packages)
|
||||||
|
|
||||||
|
# Should not run as root
|
||||||
|
assert "USER" in containerfile
|
||||||
|
|
||||||
|
# Should set working directory
|
||||||
|
assert "WORKDIR" in containerfile
|
||||||
|
|
||||||
|
|
||||||
|
def test_generate_requirements_preserves_version_specs(installer):
|
||||||
|
"""Test that version specifiers are preserved exactly."""
|
||||||
|
packages = [
|
||||||
|
"numpy>=1.24.0,<2.0.0",
|
||||||
|
"pandas==2.0.0",
|
||||||
|
"requests~=2.28.0",
|
||||||
|
"scipy!=1.10.0"
|
||||||
|
]
|
||||||
|
|
||||||
|
requirements = installer.generate_requirements(packages)
|
||||||
|
|
||||||
|
for package in packages:
|
||||||
|
assert package in requirements
|
||||||
|
|
||||||
|
|
||||||
|
def test_generate_containerfile_with_extras(installer):
|
||||||
|
"""Test Containerfile works with package extras."""
|
||||||
|
base_image = "docker.io/python:3.11-slim"
|
||||||
|
packages = ["requests[security]>=2.28.0", "pandas[excel]"]
|
||||||
|
|
||||||
|
containerfile = installer.generate_containerfile(base_image, packages)
|
||||||
|
|
||||||
|
# Should handle extras in requirements
|
||||||
|
assert "COPY" in containerfile and "requirements.txt" in containerfile
|
||||||
|
assert "RUN uv pip install" in containerfile
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_build_context_with_custom_python_version(installer):
|
||||||
|
"""Test build context creation with custom Python version."""
|
||||||
|
base_image = "docker.io/python:3.12-slim"
|
||||||
|
packages = ["numpy"]
|
||||||
|
|
||||||
|
context_path = installer.create_build_context(
|
||||||
|
base_image, packages, python_version="3.12"
|
||||||
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
|
containerfile = (context_path / "Containerfile").read_text()
|
||||||
|
assert "3.12" in containerfile or "python:3.12" in containerfile
|
||||||
|
finally:
|
||||||
|
if context_path.exists():
|
||||||
|
shutil.rmtree(context_path)
|
||||||
|
|
||||||
|
|
||||||
|
def test_generate_containerfile_optimizes_layer_caching(installer):
|
||||||
|
"""Test that Containerfile structure optimizes Docker layer caching."""
|
||||||
|
base_image = "docker.io/python:3.11-slim"
|
||||||
|
packages = ["numpy", "pandas", "scipy"]
|
||||||
|
|
||||||
|
containerfile = installer.generate_containerfile(base_image, packages)
|
||||||
|
|
||||||
|
lines = containerfile.split('\n')
|
||||||
|
|
||||||
|
# UV install should come before requirements copy
|
||||||
|
uv_install_idx = None
|
||||||
|
requirements_idx = None
|
||||||
|
|
||||||
|
for i, line in enumerate(lines):
|
||||||
|
if 'uv' in line.lower() and ('pip install' in line.lower() or 'ADD' in line):
|
||||||
|
uv_install_idx = i
|
||||||
|
if 'COPY' in line and 'requirements.txt' in line:
|
||||||
|
requirements_idx = i
|
||||||
|
|
||||||
|
# UV installation should be cached separately
|
||||||
|
assert uv_install_idx is not None
|
||||||
|
# Requirements copy should happen for cache busting
|
||||||
|
assert requirements_idx is not None
|
||||||
|
|
||||||
|
|
||||||
|
def test_empty_packages_list_creates_valid_containerfile(installer):
|
||||||
|
"""Test that empty package list still creates valid Containerfile."""
|
||||||
|
base_image = "docker.io/python:3.11-slim"
|
||||||
|
packages = []
|
||||||
|
|
||||||
|
containerfile = installer.generate_containerfile(base_image, packages)
|
||||||
|
|
||||||
|
# Should still have base structure
|
||||||
|
assert f"FROM {base_image}" in containerfile
|
||||||
|
assert "WORKDIR" in containerfile
|
||||||
439
tests/config/test_loader.py
Normal file
439
tests/config/test_loader.py
Normal file
|
|
@ -0,0 +1,439 @@
|
||||||
|
"""
|
||||||
|
Tests for configuration loader module.
|
||||||
|
|
||||||
|
Following TDD approach - these tests are written before implementation.
|
||||||
|
Tests cover all loading and substitution requirements from todo.md section 1.1.2.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import yaml
|
||||||
|
from pathlib import Path
|
||||||
|
from pydantic import ValidationError
|
||||||
|
import os
|
||||||
|
import tempfile
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_from_valid_yaml_file(tmp_path):
|
||||||
|
"""Test loading configuration from a valid YAML file."""
|
||||||
|
from mcp_forge.config.loader import load_config
|
||||||
|
|
||||||
|
config_file = tmp_path / "config.yaml"
|
||||||
|
config_content = """
|
||||||
|
server:
|
||||||
|
host: "0.0.0.0"
|
||||||
|
port: 3000
|
||||||
|
podman_socket: "/run/user/1000/podman/podman.sock"
|
||||||
|
|
||||||
|
execution:
|
||||||
|
default_backend: "simple"
|
||||||
|
default_timeout: 300
|
||||||
|
max_timeout: 1800
|
||||||
|
default_memory: "512m"
|
||||||
|
max_memory: "2g"
|
||||||
|
default_cpu_quota: 50000
|
||||||
|
max_cpu_quota: 100000
|
||||||
|
|
||||||
|
images:
|
||||||
|
python_3_11: "mcp-forge/python:3.11"
|
||||||
|
python_3_12: "mcp-forge/python:3.12"
|
||||||
|
jupyter: "mcp-forge/jupyter:latest"
|
||||||
|
auto_pull: true
|
||||||
|
pull_interval: 86400
|
||||||
|
|
||||||
|
sessions:
|
||||||
|
idle_timeout: 3600
|
||||||
|
max_concurrent: 10
|
||||||
|
cleanup_interval: 300
|
||||||
|
|
||||||
|
volumes:
|
||||||
|
base_path: "/mcp-forge/volumes"
|
||||||
|
session_quota: "1g"
|
||||||
|
max_session_quota: "10g"
|
||||||
|
|
||||||
|
security:
|
||||||
|
audit_log: "/var/log/mcp-forge/audit.log"
|
||||||
|
enforce_resource_limits: true
|
||||||
|
allow_network: false
|
||||||
|
|
||||||
|
environment_builder:
|
||||||
|
enabled: true
|
||||||
|
uv_cache_path: "/var/cache/mcp-forge/uv"
|
||||||
|
max_packages_per_build: 50
|
||||||
|
max_build_time: 600
|
||||||
|
max_image_size: 2147483648
|
||||||
|
max_concurrent_builds: 3
|
||||||
|
build_rate_limit: {}
|
||||||
|
auto_cleanup: {}
|
||||||
|
templates: {}
|
||||||
|
package_validation:
|
||||||
|
use_allowlist: true
|
||||||
|
allowlist_path: "/etc/mcp-forge/allowlist.txt"
|
||||||
|
blocklist_path: "/etc/mcp-forge/blocklist.txt"
|
||||||
|
require_approval_patterns: []
|
||||||
|
|
||||||
|
mcp_tools:
|
||||||
|
git:
|
||||||
|
command: "uvx"
|
||||||
|
args: ["mcp-server-git"]
|
||||||
|
env: {}
|
||||||
|
"""
|
||||||
|
config_file.write_text(config_content)
|
||||||
|
|
||||||
|
config = load_config(config_file)
|
||||||
|
|
||||||
|
assert config.server.host == "0.0.0.0"
|
||||||
|
assert config.server.port == 3000
|
||||||
|
assert config.execution.default_backend == "simple"
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_from_non_existent_file_raises_file_not_found_error():
|
||||||
|
"""Test that loading from non-existent file raises FileNotFoundError."""
|
||||||
|
from mcp_forge.config.loader import load_config
|
||||||
|
|
||||||
|
with pytest.raises(FileNotFoundError):
|
||||||
|
load_config(Path("/non/existent/config.yaml"))
|
||||||
|
|
||||||
|
|
||||||
|
def test_invalid_yaml_raises_yaml_error(tmp_path):
|
||||||
|
"""Test that invalid YAML syntax raises YAMLError."""
|
||||||
|
from mcp_forge.config.loader import load_config
|
||||||
|
|
||||||
|
config_file = tmp_path / "config.yaml"
|
||||||
|
config_file.write_text("""
|
||||||
|
invalid yaml:
|
||||||
|
- unmatched [bracket
|
||||||
|
key without value
|
||||||
|
""")
|
||||||
|
|
||||||
|
with pytest.raises(yaml.YAMLError):
|
||||||
|
load_config(config_file)
|
||||||
|
|
||||||
|
|
||||||
|
def test_environment_variable_substitution_in_strings(tmp_path, monkeypatch):
|
||||||
|
"""Test that ${VAR} is replaced with environment variable value."""
|
||||||
|
from mcp_forge.config.loader import load_config
|
||||||
|
|
||||||
|
monkeypatch.setenv("TEST_HOST", "test.example.com")
|
||||||
|
monkeypatch.setenv("TEST_PORT", "4000")
|
||||||
|
monkeypatch.setenv("PODMAN_SOCKET", "/run/test/podman.sock")
|
||||||
|
|
||||||
|
config_file = tmp_path / "config.yaml"
|
||||||
|
config_content = """
|
||||||
|
server:
|
||||||
|
host: "${TEST_HOST}"
|
||||||
|
port: ${TEST_PORT}
|
||||||
|
podman_socket: "${PODMAN_SOCKET}"
|
||||||
|
|
||||||
|
execution:
|
||||||
|
default_backend: "simple"
|
||||||
|
|
||||||
|
images: {}
|
||||||
|
|
||||||
|
sessions: {}
|
||||||
|
|
||||||
|
volumes:
|
||||||
|
base_path: "/volumes"
|
||||||
|
|
||||||
|
security:
|
||||||
|
audit_log: "/audit.log"
|
||||||
|
|
||||||
|
environment_builder:
|
||||||
|
enabled: true
|
||||||
|
uv_cache_path: "/cache"
|
||||||
|
package_validation:
|
||||||
|
allowlist_path: "/allow.txt"
|
||||||
|
blocklist_path: "/block.txt"
|
||||||
|
require_approval_patterns: []
|
||||||
|
|
||||||
|
mcp_tools: {}
|
||||||
|
"""
|
||||||
|
config_file.write_text(config_content)
|
||||||
|
|
||||||
|
config = load_config(config_file)
|
||||||
|
|
||||||
|
assert config.server.host == "test.example.com"
|
||||||
|
assert config.server.port == 4000
|
||||||
|
assert str(config.server.podman_socket) == "/run/test/podman.sock"
|
||||||
|
|
||||||
|
|
||||||
|
def test_nested_environment_variable_substitution(tmp_path, monkeypatch):
|
||||||
|
"""Test that environment variable substitution works in nested structures."""
|
||||||
|
from mcp_forge.config.loader import load_config
|
||||||
|
|
||||||
|
monkeypatch.setenv("AUDIT_LOG_PATH", "/var/log/audit.log")
|
||||||
|
monkeypatch.setenv("UV_CACHE", "/var/cache/uv")
|
||||||
|
|
||||||
|
config_file = tmp_path / "config.yaml"
|
||||||
|
config_content = """
|
||||||
|
server:
|
||||||
|
host: "localhost"
|
||||||
|
port: 3000
|
||||||
|
podman_socket: "/run/podman.sock"
|
||||||
|
|
||||||
|
execution: {}
|
||||||
|
|
||||||
|
images: {}
|
||||||
|
|
||||||
|
sessions: {}
|
||||||
|
|
||||||
|
volumes:
|
||||||
|
base_path: "/volumes"
|
||||||
|
|
||||||
|
security:
|
||||||
|
audit_log: "${AUDIT_LOG_PATH}"
|
||||||
|
|
||||||
|
environment_builder:
|
||||||
|
enabled: true
|
||||||
|
uv_cache_path: "${UV_CACHE}"
|
||||||
|
package_validation:
|
||||||
|
allowlist_path: "/allow.txt"
|
||||||
|
blocklist_path: "/block.txt"
|
||||||
|
require_approval_patterns: []
|
||||||
|
|
||||||
|
mcp_tools: {}
|
||||||
|
"""
|
||||||
|
config_file.write_text(config_content)
|
||||||
|
|
||||||
|
config = load_config(config_file)
|
||||||
|
|
||||||
|
assert str(config.security.audit_log) == "/var/log/audit.log"
|
||||||
|
assert str(config.environment_builder.uv_cache_path) == "/var/cache/uv"
|
||||||
|
|
||||||
|
|
||||||
|
def test_missing_environment_variable_raises_clear_error(tmp_path):
|
||||||
|
"""Test that missing environment variable raises error with variable name."""
|
||||||
|
from mcp_forge.config.loader import load_config
|
||||||
|
|
||||||
|
config_file = tmp_path / "config.yaml"
|
||||||
|
config_content = """
|
||||||
|
server:
|
||||||
|
host: "${MISSING_VAR}"
|
||||||
|
port: 3000
|
||||||
|
podman_socket: "/run/podman.sock"
|
||||||
|
|
||||||
|
execution: {}
|
||||||
|
images: {}
|
||||||
|
sessions: {}
|
||||||
|
volumes:
|
||||||
|
base_path: "/volumes"
|
||||||
|
security:
|
||||||
|
audit_log: "/audit.log"
|
||||||
|
environment_builder:
|
||||||
|
enabled: true
|
||||||
|
uv_cache_path: "/cache"
|
||||||
|
package_validation:
|
||||||
|
allowlist_path: "/allow.txt"
|
||||||
|
blocklist_path: "/block.txt"
|
||||||
|
require_approval_patterns: []
|
||||||
|
mcp_tools: {}
|
||||||
|
"""
|
||||||
|
config_file.write_text(config_content)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError) as exc_info:
|
||||||
|
load_config(config_file)
|
||||||
|
|
||||||
|
assert "MISSING_VAR" in str(exc_info.value)
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_config_from_dict():
|
||||||
|
"""Test loading configuration from dictionary (for testing)."""
|
||||||
|
from mcp_forge.config.loader import load_config_from_dict
|
||||||
|
|
||||||
|
config_dict = {
|
||||||
|
"server": {
|
||||||
|
"host": "localhost",
|
||||||
|
"port": 3000,
|
||||||
|
"podman_socket": "/run/podman.sock"
|
||||||
|
},
|
||||||
|
"execution": {},
|
||||||
|
"images": {},
|
||||||
|
"sessions": {},
|
||||||
|
"volumes": {"base_path": "/volumes"},
|
||||||
|
"security": {"audit_log": "/audit.log"},
|
||||||
|
"environment_builder": {
|
||||||
|
"enabled": True,
|
||||||
|
"uv_cache_path": "/cache",
|
||||||
|
"package_validation": {
|
||||||
|
"allowlist_path": "/allow.txt",
|
||||||
|
"blocklist_path": "/block.txt",
|
||||||
|
"require_approval_patterns": []
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"mcp_tools": {}
|
||||||
|
}
|
||||||
|
|
||||||
|
config = load_config_from_dict(config_dict)
|
||||||
|
|
||||||
|
assert config.server.host == "localhost"
|
||||||
|
assert config.server.port == 3000
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_config_from_dict_with_env_var_substitution(monkeypatch):
|
||||||
|
"""Test that load_config_from_dict also performs env var substitution."""
|
||||||
|
from mcp_forge.config.loader import load_config_from_dict
|
||||||
|
|
||||||
|
monkeypatch.setenv("TEST_HOST", "example.com")
|
||||||
|
|
||||||
|
config_dict = {
|
||||||
|
"server": {
|
||||||
|
"host": "${TEST_HOST}",
|
||||||
|
"port": 3000,
|
||||||
|
"podman_socket": "/run/podman.sock"
|
||||||
|
},
|
||||||
|
"execution": {},
|
||||||
|
"images": {},
|
||||||
|
"sessions": {},
|
||||||
|
"volumes": {"base_path": "/volumes"},
|
||||||
|
"security": {"audit_log": "/audit.log"},
|
||||||
|
"environment_builder": {
|
||||||
|
"enabled": True,
|
||||||
|
"uv_cache_path": "/cache",
|
||||||
|
"package_validation": {
|
||||||
|
"allowlist_path": "/allow.txt",
|
||||||
|
"blocklist_path": "/block.txt",
|
||||||
|
"require_approval_patterns": []
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"mcp_tools": {}
|
||||||
|
}
|
||||||
|
|
||||||
|
config = load_config_from_dict(config_dict)
|
||||||
|
|
||||||
|
assert config.server.host == "example.com"
|
||||||
|
|
||||||
|
|
||||||
|
def test_partial_configuration_uses_defaults(tmp_path):
|
||||||
|
"""Test that missing configuration sections use schema defaults."""
|
||||||
|
from mcp_forge.config.loader import load_config
|
||||||
|
|
||||||
|
config_file = tmp_path / "config.yaml"
|
||||||
|
config_content = """
|
||||||
|
server:
|
||||||
|
host: "localhost"
|
||||||
|
port: 3000
|
||||||
|
podman_socket: "/run/podman.sock"
|
||||||
|
|
||||||
|
execution: {}
|
||||||
|
|
||||||
|
images: {}
|
||||||
|
|
||||||
|
sessions: {}
|
||||||
|
|
||||||
|
volumes:
|
||||||
|
base_path: "/volumes"
|
||||||
|
|
||||||
|
security:
|
||||||
|
audit_log: "/audit.log"
|
||||||
|
|
||||||
|
environment_builder:
|
||||||
|
enabled: true
|
||||||
|
uv_cache_path: "/cache"
|
||||||
|
package_validation:
|
||||||
|
allowlist_path: "/allow.txt"
|
||||||
|
blocklist_path: "/block.txt"
|
||||||
|
require_approval_patterns: []
|
||||||
|
|
||||||
|
mcp_tools: {}
|
||||||
|
"""
|
||||||
|
config_file.write_text(config_content)
|
||||||
|
|
||||||
|
config = load_config(config_file)
|
||||||
|
|
||||||
|
# Check defaults from ExecutionConfig
|
||||||
|
assert config.execution.default_backend == "simple"
|
||||||
|
assert config.execution.default_timeout == 300
|
||||||
|
assert config.execution.max_timeout == 1800
|
||||||
|
|
||||||
|
# Check defaults from ImageConfig
|
||||||
|
assert config.images.python_3_11 == "mcp-forge/python:3.11"
|
||||||
|
assert config.images.auto_pull is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_substitute_env_vars_recursive():
|
||||||
|
"""Test that substitute_env_vars works recursively on nested structures."""
|
||||||
|
from mcp_forge.config.loader import substitute_env_vars
|
||||||
|
import os
|
||||||
|
|
||||||
|
os.environ["TEST_VALUE"] = "substituted"
|
||||||
|
|
||||||
|
data = {
|
||||||
|
"simple": "${TEST_VALUE}",
|
||||||
|
"nested": {
|
||||||
|
"deep": "${TEST_VALUE}",
|
||||||
|
"list": ["${TEST_VALUE}", "plain"]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
result = substitute_env_vars(data)
|
||||||
|
|
||||||
|
assert result["simple"] == "substituted"
|
||||||
|
assert result["nested"]["deep"] == "substituted"
|
||||||
|
assert result["nested"]["list"][0] == "substituted"
|
||||||
|
assert result["nested"]["list"][1] == "plain"
|
||||||
|
|
||||||
|
|
||||||
|
def test_substitute_env_vars_with_integer_conversion(monkeypatch):
|
||||||
|
"""Test that numeric strings in env vars can be converted to integers."""
|
||||||
|
from mcp_forge.config.loader import substitute_env_vars
|
||||||
|
|
||||||
|
monkeypatch.setenv("PORT_NUM", "8080")
|
||||||
|
|
||||||
|
data = {
|
||||||
|
"port": "${PORT_NUM}"
|
||||||
|
}
|
||||||
|
|
||||||
|
result = substitute_env_vars(data)
|
||||||
|
|
||||||
|
# Should still be a string after substitution; type conversion handled by Pydantic
|
||||||
|
assert result["port"] == "8080"
|
||||||
|
|
||||||
|
|
||||||
|
def test_no_eval_or_exec_in_substitution():
|
||||||
|
"""Test that no eval() or exec() is used - only safe string substitution."""
|
||||||
|
from mcp_forge.config.loader import substitute_env_vars
|
||||||
|
import os
|
||||||
|
|
||||||
|
# Try to inject malicious code - should be treated as literal string
|
||||||
|
os.environ["MALICIOUS"] = "__import__('os').system('echo hacked')"
|
||||||
|
|
||||||
|
data = {
|
||||||
|
"value": "${MALICIOUS}"
|
||||||
|
}
|
||||||
|
|
||||||
|
result = substitute_env_vars(data)
|
||||||
|
|
||||||
|
# Should be the literal string, not executed
|
||||||
|
assert result["value"] == "__import__('os').system('echo hacked')"
|
||||||
|
# And it should NOT have been executed (we can't test side effects easily,
|
||||||
|
# but the substitution logic shouldn't use eval/exec)
|
||||||
|
|
||||||
|
|
||||||
|
def test_env_var_with_special_characters(monkeypatch):
|
||||||
|
"""Test that environment variables with special characters are handled correctly."""
|
||||||
|
from mcp_forge.config.loader import substitute_env_vars
|
||||||
|
|
||||||
|
monkeypatch.setenv("SPECIAL_PATH", "/path/with-dashes_and_underscores/123")
|
||||||
|
|
||||||
|
data = {
|
||||||
|
"path": "${SPECIAL_PATH}"
|
||||||
|
}
|
||||||
|
|
||||||
|
result = substitute_env_vars(data)
|
||||||
|
|
||||||
|
assert result["path"] == "/path/with-dashes_and_underscores/123"
|
||||||
|
|
||||||
|
|
||||||
|
def test_multiple_env_vars_in_single_string(monkeypatch):
|
||||||
|
"""Test that multiple environment variables can be substituted in one string."""
|
||||||
|
from mcp_forge.config.loader import substitute_env_vars
|
||||||
|
|
||||||
|
monkeypatch.setenv("BASE_PATH", "/opt/mcp-forge")
|
||||||
|
monkeypatch.setenv("SUBDIR", "logs")
|
||||||
|
|
||||||
|
data = {
|
||||||
|
"log_path": "${BASE_PATH}/${SUBDIR}/audit.log"
|
||||||
|
}
|
||||||
|
|
||||||
|
result = substitute_env_vars(data)
|
||||||
|
|
||||||
|
assert result["log_path"] == "/opt/mcp-forge/logs/audit.log"
|
||||||
141
tests/config/test_mcp_tool_config.py
Normal file
141
tests/config/test_mcp_tool_config.py
Normal file
|
|
@ -0,0 +1,141 @@
|
||||||
|
"""Tests for MCPToolConfig schema with HTTP/SSE transport."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from pydantic import ValidationError
|
||||||
|
from mcp_forge.config.schema import MCPToolConfig
|
||||||
|
|
||||||
|
|
||||||
|
def test_stdio_transport_config_valid():
|
||||||
|
"""Test valid stdio transport configuration."""
|
||||||
|
config = MCPToolConfig(
|
||||||
|
transport="stdio",
|
||||||
|
command="python",
|
||||||
|
args=["-m", "server"],
|
||||||
|
env={"KEY": "value"}
|
||||||
|
)
|
||||||
|
|
||||||
|
assert config.transport == "stdio"
|
||||||
|
assert config.command == "python"
|
||||||
|
assert config.args == ["-m", "server"]
|
||||||
|
assert config.env == {"KEY": "value"}
|
||||||
|
|
||||||
|
|
||||||
|
def test_http_transport_config_valid():
|
||||||
|
"""Test valid HTTP transport configuration."""
|
||||||
|
config = MCPToolConfig(
|
||||||
|
transport="http",
|
||||||
|
url="http://localhost:8006/mcp",
|
||||||
|
headers={"Authorization": "Bearer token"}
|
||||||
|
)
|
||||||
|
|
||||||
|
assert config.transport == "http"
|
||||||
|
assert config.url == "http://localhost:8006/mcp"
|
||||||
|
assert config.headers == {"Authorization": "Bearer token"}
|
||||||
|
|
||||||
|
|
||||||
|
def test_sse_transport_config_valid():
|
||||||
|
"""Test valid SSE transport configuration."""
|
||||||
|
config = MCPToolConfig(
|
||||||
|
transport="sse",
|
||||||
|
url="http://localhost:9000/events",
|
||||||
|
headers={"X-Custom": "value"}
|
||||||
|
)
|
||||||
|
|
||||||
|
assert config.transport == "sse"
|
||||||
|
assert config.url == "http://localhost:9000/events"
|
||||||
|
assert config.headers == {"X-Custom": "value"}
|
||||||
|
|
||||||
|
|
||||||
|
def test_stdio_without_command_invalid():
|
||||||
|
"""Test that stdio transport requires command."""
|
||||||
|
with pytest.raises(ValidationError, match="command is required for stdio transport"):
|
||||||
|
MCPToolConfig(
|
||||||
|
transport="stdio",
|
||||||
|
args=["-m", "server"]
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_http_without_url_invalid():
|
||||||
|
"""Test that HTTP transport requires URL."""
|
||||||
|
with pytest.raises(ValidationError, match="url is required for http transport"):
|
||||||
|
MCPToolConfig(
|
||||||
|
transport="http",
|
||||||
|
headers={"Authorization": "Bearer token"}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_sse_without_url_invalid():
|
||||||
|
"""Test that SSE transport requires URL."""
|
||||||
|
with pytest.raises(ValidationError, match="url is required for sse transport"):
|
||||||
|
MCPToolConfig(
|
||||||
|
transport="sse",
|
||||||
|
headers={"X-Custom": "value"}
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_default_transport_is_stdio():
|
||||||
|
"""Test that default transport is stdio."""
|
||||||
|
config = MCPToolConfig(
|
||||||
|
command="python",
|
||||||
|
args=["-m", "server"]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert config.transport == "stdio"
|
||||||
|
|
||||||
|
|
||||||
|
def test_http_config_with_empty_headers():
|
||||||
|
"""Test HTTP config with no headers."""
|
||||||
|
config = MCPToolConfig(
|
||||||
|
transport="http",
|
||||||
|
url="http://localhost:8006/mcp"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert config.headers == {}
|
||||||
|
|
||||||
|
|
||||||
|
def test_stdio_config_with_empty_env():
|
||||||
|
"""Test stdio config with no env vars."""
|
||||||
|
config = MCPToolConfig(
|
||||||
|
transport="stdio",
|
||||||
|
command="python"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert config.env == {}
|
||||||
|
assert config.args == []
|
||||||
|
|
||||||
|
|
||||||
|
def test_invalid_transport_type():
|
||||||
|
"""Test that invalid transport type is rejected."""
|
||||||
|
with pytest.raises(ValidationError):
|
||||||
|
MCPToolConfig(
|
||||||
|
transport="invalid",
|
||||||
|
command="python"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_stdio_config_minimal():
|
||||||
|
"""Test minimal stdio config with just command."""
|
||||||
|
config = MCPToolConfig(
|
||||||
|
command="python"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert config.transport == "stdio"
|
||||||
|
assert config.command == "python"
|
||||||
|
assert config.args == []
|
||||||
|
assert config.env == {}
|
||||||
|
|
||||||
|
|
||||||
|
def test_http_config_with_multiple_headers():
|
||||||
|
"""Test HTTP config with multiple headers."""
|
||||||
|
config = MCPToolConfig(
|
||||||
|
transport="http",
|
||||||
|
url="http://localhost:8006/mcp",
|
||||||
|
headers={
|
||||||
|
"Authorization": "Bearer token123",
|
||||||
|
"X-Custom-Header": "value",
|
||||||
|
"Content-Type": "application/json"
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
assert len(config.headers) == 3
|
||||||
|
assert config.headers["Authorization"] == "Bearer token123"
|
||||||
403
tests/config/test_schema.py
Normal file
403
tests/config/test_schema.py
Normal file
|
|
@ -0,0 +1,403 @@
|
||||||
|
"""
|
||||||
|
Tests for configuration schema module.
|
||||||
|
|
||||||
|
Following TDD approach - these tests are written before implementation.
|
||||||
|
Tests cover all validation requirements from todo.md section 1.1.1.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from pathlib import Path
|
||||||
|
from pydantic import ValidationError
|
||||||
|
|
||||||
|
|
||||||
|
def test_valid_server_config_loads_successfully():
|
||||||
|
"""Test that a valid ServerConfig loads without errors."""
|
||||||
|
from mcp_forge.config.schema import ServerConfig
|
||||||
|
|
||||||
|
config = ServerConfig(
|
||||||
|
host="0.0.0.0",
|
||||||
|
port=3000,
|
||||||
|
podman_socket=Path("/run/user/1000/podman/podman.sock")
|
||||||
|
)
|
||||||
|
|
||||||
|
assert config.host == "0.0.0.0"
|
||||||
|
assert config.port == 3000
|
||||||
|
assert config.podman_socket == Path("/run/user/1000/podman/podman.sock")
|
||||||
|
|
||||||
|
|
||||||
|
def test_server_config_invalid_port_raises_validation_error():
|
||||||
|
"""Test that invalid port numbers raise ValidationError."""
|
||||||
|
from mcp_forge.config.schema import ServerConfig
|
||||||
|
|
||||||
|
# Port too high
|
||||||
|
with pytest.raises(ValidationError) as exc_info:
|
||||||
|
ServerConfig(
|
||||||
|
host="localhost",
|
||||||
|
port=70000,
|
||||||
|
podman_socket=Path("/run/podman.sock")
|
||||||
|
)
|
||||||
|
assert "port" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
# Port too low
|
||||||
|
with pytest.raises(ValidationError) as exc_info:
|
||||||
|
ServerConfig(
|
||||||
|
host="localhost",
|
||||||
|
port=0,
|
||||||
|
podman_socket=Path("/run/podman.sock")
|
||||||
|
)
|
||||||
|
assert "port" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
# Negative port
|
||||||
|
with pytest.raises(ValidationError) as exc_info:
|
||||||
|
ServerConfig(
|
||||||
|
host="localhost",
|
||||||
|
port=-1,
|
||||||
|
podman_socket=Path("/run/podman.sock")
|
||||||
|
)
|
||||||
|
assert "port" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_execution_config_defaults_are_applied():
|
||||||
|
"""Test that ExecutionConfig has correct default values."""
|
||||||
|
from mcp_forge.config.schema import ExecutionConfig
|
||||||
|
|
||||||
|
config = ExecutionConfig()
|
||||||
|
|
||||||
|
assert config.default_backend == "simple"
|
||||||
|
assert config.default_timeout == 300
|
||||||
|
assert config.max_timeout == 1800
|
||||||
|
assert config.default_memory == "512m"
|
||||||
|
assert config.max_memory == "2g"
|
||||||
|
assert config.default_cpu_quota == 50000
|
||||||
|
assert config.max_cpu_quota == 100000
|
||||||
|
|
||||||
|
|
||||||
|
def test_execution_config_max_timeout_validation():
|
||||||
|
"""Test that max_timeout must be >= default_timeout."""
|
||||||
|
from mcp_forge.config.schema import ExecutionConfig
|
||||||
|
|
||||||
|
# Valid: max >= default
|
||||||
|
config = ExecutionConfig(default_timeout=300, max_timeout=600)
|
||||||
|
assert config.max_timeout == 600
|
||||||
|
|
||||||
|
# Invalid: max < default
|
||||||
|
with pytest.raises(ValidationError) as exc_info:
|
||||||
|
ExecutionConfig(default_timeout=600, max_timeout=300)
|
||||||
|
assert "max_timeout" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_image_config_defaults():
|
||||||
|
"""Test ImageConfig default values."""
|
||||||
|
from mcp_forge.config.schema import ImageConfig
|
||||||
|
|
||||||
|
config = ImageConfig()
|
||||||
|
|
||||||
|
assert config.python_3_11 == "mcp-forge/python:3.11"
|
||||||
|
assert config.python_3_12 == "mcp-forge/python:3.12"
|
||||||
|
assert config.jupyter == "mcp-forge/jupyter:latest"
|
||||||
|
assert config.auto_pull is True
|
||||||
|
assert config.pull_interval == 86400
|
||||||
|
|
||||||
|
|
||||||
|
def test_session_config_defaults():
|
||||||
|
"""Test SessionConfig default values."""
|
||||||
|
from mcp_forge.config.schema import SessionConfig
|
||||||
|
|
||||||
|
config = SessionConfig()
|
||||||
|
|
||||||
|
assert config.idle_timeout == 3600
|
||||||
|
assert config.max_concurrent == 10
|
||||||
|
assert config.cleanup_interval == 300
|
||||||
|
|
||||||
|
|
||||||
|
def test_volume_config_with_base_path():
|
||||||
|
"""Test VolumeConfig with required base_path."""
|
||||||
|
from mcp_forge.config.schema import VolumeConfig
|
||||||
|
|
||||||
|
config = VolumeConfig(base_path=Path("/mcp-forge/volumes"))
|
||||||
|
|
||||||
|
assert config.base_path == Path("/mcp-forge/volumes")
|
||||||
|
assert config.session_quota == "1g"
|
||||||
|
assert config.max_session_quota == "10g"
|
||||||
|
|
||||||
|
|
||||||
|
def test_security_config_defaults():
|
||||||
|
"""Test SecurityConfig default values."""
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/var/log/mcp-forge/audit.log"))
|
||||||
|
|
||||||
|
assert config.audit_log == Path("/var/log/mcp-forge/audit.log")
|
||||||
|
assert config.enforce_resource_limits is True
|
||||||
|
assert config.allow_network is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_package_validation_config():
|
||||||
|
"""Test PackageValidationConfig structure."""
|
||||||
|
from mcp_forge.config.schema import PackageValidationConfig
|
||||||
|
|
||||||
|
config = PackageValidationConfig(
|
||||||
|
use_allowlist=True,
|
||||||
|
allowlist_path=Path("/etc/mcp-forge/allowlist.txt"),
|
||||||
|
blocklist_path=Path("/etc/mcp-forge/blocklist.txt"),
|
||||||
|
require_approval_patterns=["*crypto*", "*network*"]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert config.use_allowlist is True
|
||||||
|
assert config.allowlist_path == Path("/etc/mcp-forge/allowlist.txt")
|
||||||
|
assert config.blocklist_path == Path("/etc/mcp-forge/blocklist.txt")
|
||||||
|
assert "*crypto*" in config.require_approval_patterns
|
||||||
|
|
||||||
|
|
||||||
|
def test_environment_builder_config():
|
||||||
|
"""Test EnvironmentBuilderConfig structure and defaults."""
|
||||||
|
from mcp_forge.config.schema import EnvironmentBuilderConfig, PackageValidationConfig
|
||||||
|
|
||||||
|
pkg_validation = PackageValidationConfig(
|
||||||
|
use_allowlist=True,
|
||||||
|
allowlist_path=Path("/etc/allowlist.txt"),
|
||||||
|
blocklist_path=Path("/etc/blocklist.txt"),
|
||||||
|
require_approval_patterns=[]
|
||||||
|
)
|
||||||
|
|
||||||
|
config = EnvironmentBuilderConfig(
|
||||||
|
enabled=True,
|
||||||
|
uv_cache_path=Path("/var/cache/mcp-forge/uv"),
|
||||||
|
package_validation=pkg_validation
|
||||||
|
)
|
||||||
|
|
||||||
|
assert config.enabled is True
|
||||||
|
assert config.uv_cache_path == Path("/var/cache/mcp-forge/uv")
|
||||||
|
assert config.max_packages_per_build == 50
|
||||||
|
assert config.max_build_time == 600
|
||||||
|
assert config.max_image_size == 2147483648
|
||||||
|
assert config.max_concurrent_builds == 3
|
||||||
|
|
||||||
|
|
||||||
|
def test_mcp_tool_config():
|
||||||
|
"""Test MCPToolConfig structure."""
|
||||||
|
from mcp_forge.config.schema import MCPToolConfig
|
||||||
|
|
||||||
|
config = MCPToolConfig(
|
||||||
|
command="uvx",
|
||||||
|
args=["mcp-server-git"],
|
||||||
|
env={"GIT_AUTHOR": "test"}
|
||||||
|
)
|
||||||
|
|
||||||
|
assert config.command == "uvx"
|
||||||
|
assert config.args == ["mcp-server-git"]
|
||||||
|
assert config.env == {"GIT_AUTHOR": "test"}
|
||||||
|
|
||||||
|
|
||||||
|
def test_mcp_tool_config_empty_env_defaults():
|
||||||
|
"""Test MCPToolConfig with empty env defaults to empty dict."""
|
||||||
|
from mcp_forge.config.schema import MCPToolConfig
|
||||||
|
|
||||||
|
config = MCPToolConfig(command="test", args=[])
|
||||||
|
|
||||||
|
assert config.env == {}
|
||||||
|
|
||||||
|
|
||||||
|
def test_forge_config_full_structure():
|
||||||
|
"""Test complete ForgeConfig with all nested structures."""
|
||||||
|
from mcp_forge.config.schema import (
|
||||||
|
ForgeConfig, ServerConfig, ExecutionConfig, ImageConfig,
|
||||||
|
SessionConfig, VolumeConfig, SecurityConfig,
|
||||||
|
EnvironmentBuilderConfig, PackageValidationConfig, MCPToolConfig
|
||||||
|
)
|
||||||
|
|
||||||
|
pkg_validation = PackageValidationConfig(
|
||||||
|
use_allowlist=True,
|
||||||
|
allowlist_path=Path("/etc/allowlist.txt"),
|
||||||
|
blocklist_path=Path("/etc/blocklist.txt"),
|
||||||
|
require_approval_patterns=[]
|
||||||
|
)
|
||||||
|
|
||||||
|
config = ForgeConfig(
|
||||||
|
server=ServerConfig(
|
||||||
|
host="localhost",
|
||||||
|
port=3000,
|
||||||
|
podman_socket=Path("/run/podman.sock")
|
||||||
|
),
|
||||||
|
execution=ExecutionConfig(),
|
||||||
|
images=ImageConfig(),
|
||||||
|
sessions=SessionConfig(),
|
||||||
|
volumes=VolumeConfig(base_path=Path("/mcp-forge/volumes")),
|
||||||
|
security=SecurityConfig(audit_log=Path("/var/log/audit.log")),
|
||||||
|
environment_builder=EnvironmentBuilderConfig(
|
||||||
|
enabled=True,
|
||||||
|
uv_cache_path=Path("/var/cache/uv"),
|
||||||
|
package_validation=pkg_validation
|
||||||
|
),
|
||||||
|
mcp_tools={
|
||||||
|
"git": MCPToolConfig(command="uvx", args=["mcp-server-git"])
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
assert config.server.host == "localhost"
|
||||||
|
assert config.execution.default_backend == "simple"
|
||||||
|
assert config.images.python_3_11 == "mcp-forge/python:3.11"
|
||||||
|
assert config.sessions.max_concurrent == 10
|
||||||
|
assert config.volumes.base_path == Path("/mcp-forge/volumes")
|
||||||
|
assert config.security.enforce_resource_limits is True
|
||||||
|
assert config.environment_builder.enabled is True
|
||||||
|
assert "git" in config.mcp_tools
|
||||||
|
|
||||||
|
|
||||||
|
def test_nested_configuration_validation_error_includes_field_path():
|
||||||
|
"""Test that ValidationError for nested config includes full field path."""
|
||||||
|
from mcp_forge.config.schema import ForgeConfig, ServerConfig
|
||||||
|
|
||||||
|
with pytest.raises(ValidationError) as exc_info:
|
||||||
|
ForgeConfig(
|
||||||
|
server=ServerConfig(
|
||||||
|
host="localhost",
|
||||||
|
port=99999, # Invalid port
|
||||||
|
podman_socket=Path("/run/podman.sock")
|
||||||
|
),
|
||||||
|
execution={},
|
||||||
|
images={},
|
||||||
|
sessions={},
|
||||||
|
volumes={"base_path": "/volumes"},
|
||||||
|
security={"audit_log": "/audit.log"},
|
||||||
|
environment_builder={
|
||||||
|
"uv_cache_path": "/cache",
|
||||||
|
"package_validation": {
|
||||||
|
"allowlist_path": "/allow.txt",
|
||||||
|
"blocklist_path": "/block.txt",
|
||||||
|
"require_approval_patterns": []
|
||||||
|
}
|
||||||
|
},
|
||||||
|
mcp_tools={}
|
||||||
|
)
|
||||||
|
|
||||||
|
error_str = str(exc_info.value)
|
||||||
|
# Should include nested path like "server.port"
|
||||||
|
assert "port" in error_str.lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_execution_config_backend_literal_validation():
|
||||||
|
"""Test that default_backend only accepts 'simple' or 'jupyter'."""
|
||||||
|
from mcp_forge.config.schema import ExecutionConfig
|
||||||
|
|
||||||
|
# Valid values
|
||||||
|
config1 = ExecutionConfig(default_backend="simple")
|
||||||
|
assert config1.default_backend == "simple"
|
||||||
|
|
||||||
|
config2 = ExecutionConfig(default_backend="jupyter")
|
||||||
|
assert config2.default_backend == "jupyter"
|
||||||
|
|
||||||
|
# Invalid value
|
||||||
|
with pytest.raises(ValidationError) as exc_info:
|
||||||
|
ExecutionConfig(default_backend="invalid")
|
||||||
|
assert "default_backend" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_environment_builder_config_rate_limit_dict():
|
||||||
|
"""Test that EnvironmentBuilderConfig accepts rate_limit dict."""
|
||||||
|
from mcp_forge.config.schema import EnvironmentBuilderConfig, PackageValidationConfig
|
||||||
|
|
||||||
|
pkg_validation = PackageValidationConfig(
|
||||||
|
use_allowlist=True,
|
||||||
|
allowlist_path=Path("/allow.txt"),
|
||||||
|
blocklist_path=Path("/block.txt"),
|
||||||
|
require_approval_patterns=[]
|
||||||
|
)
|
||||||
|
|
||||||
|
config = EnvironmentBuilderConfig(
|
||||||
|
enabled=True,
|
||||||
|
uv_cache_path=Path("/cache"),
|
||||||
|
package_validation=pkg_validation,
|
||||||
|
build_rate_limit={"max_requests": 10, "period_seconds": 3600}
|
||||||
|
)
|
||||||
|
|
||||||
|
assert config.build_rate_limit == {"max_requests": 10, "period_seconds": 3600}
|
||||||
|
|
||||||
|
|
||||||
|
def test_environment_builder_config_auto_cleanup_dict():
|
||||||
|
"""Test that EnvironmentBuilderConfig accepts auto_cleanup dict."""
|
||||||
|
from mcp_forge.config.schema import EnvironmentBuilderConfig, PackageValidationConfig
|
||||||
|
|
||||||
|
pkg_validation = PackageValidationConfig(
|
||||||
|
use_allowlist=True,
|
||||||
|
allowlist_path=Path("/allow.txt"),
|
||||||
|
blocklist_path=Path("/block.txt"),
|
||||||
|
require_approval_patterns=[]
|
||||||
|
)
|
||||||
|
|
||||||
|
config = EnvironmentBuilderConfig(
|
||||||
|
enabled=True,
|
||||||
|
uv_cache_path=Path("/cache"),
|
||||||
|
package_validation=pkg_validation,
|
||||||
|
auto_cleanup={"enabled": True, "max_age_days": 30}
|
||||||
|
)
|
||||||
|
|
||||||
|
assert config.auto_cleanup == {"enabled": True, "max_age_days": 30}
|
||||||
|
|
||||||
|
|
||||||
|
def test_environment_builder_config_templates_dict():
|
||||||
|
"""Test that EnvironmentBuilderConfig accepts templates dict."""
|
||||||
|
from mcp_forge.config.schema import EnvironmentBuilderConfig, PackageValidationConfig
|
||||||
|
|
||||||
|
pkg_validation = PackageValidationConfig(
|
||||||
|
use_allowlist=True,
|
||||||
|
allowlist_path=Path("/allow.txt"),
|
||||||
|
blocklist_path=Path("/block.txt"),
|
||||||
|
require_approval_patterns=[]
|
||||||
|
)
|
||||||
|
|
||||||
|
config = EnvironmentBuilderConfig(
|
||||||
|
enabled=True,
|
||||||
|
uv_cache_path=Path("/cache"),
|
||||||
|
package_validation=pkg_validation,
|
||||||
|
templates={
|
||||||
|
"data-science": {"packages": ["numpy", "pandas", "matplotlib"]},
|
||||||
|
"web": {"packages": ["fastapi", "uvicorn"]}
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
assert "data-science" in config.templates
|
||||||
|
assert "web" in config.templates
|
||||||
|
assert config.templates["data-science"]["packages"] == ["numpy", "pandas", "matplotlib"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_config_constraint_validation_memory_strings():
|
||||||
|
"""Test that config validates memory strings are comparable."""
|
||||||
|
from mcp_forge.config.schema import ExecutionConfig
|
||||||
|
|
||||||
|
# This should succeed - just testing structure, not actual parsing yet
|
||||||
|
config = ExecutionConfig(
|
||||||
|
default_memory="512m",
|
||||||
|
max_memory="2g"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert config.default_memory == "512m"
|
||||||
|
assert config.max_memory == "2g"
|
||||||
|
|
||||||
|
|
||||||
|
def test_mcp_tool_config_list_of_args():
|
||||||
|
"""Test MCPToolConfig args is a list."""
|
||||||
|
from mcp_forge.config.schema import MCPToolConfig
|
||||||
|
|
||||||
|
config = MCPToolConfig(
|
||||||
|
command="python",
|
||||||
|
args=["-m", "mymodule", "--flag"]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert isinstance(config.args, list)
|
||||||
|
assert config.args == ["-m", "mymodule", "--flag"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_package_validation_require_approval_patterns_list():
|
||||||
|
"""Test PackageValidationConfig require_approval_patterns is a list."""
|
||||||
|
from mcp_forge.config.schema import PackageValidationConfig
|
||||||
|
|
||||||
|
config = PackageValidationConfig(
|
||||||
|
use_allowlist=False,
|
||||||
|
allowlist_path=Path("/allow.txt"),
|
||||||
|
blocklist_path=Path("/block.txt"),
|
||||||
|
require_approval_patterns=["*crypto*", "*security*", "paramiko"]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert isinstance(config.require_approval_patterns, list)
|
||||||
|
assert len(config.require_approval_patterns) == 3
|
||||||
0
tests/execution/__init__.py
Normal file
0
tests/execution/__init__.py
Normal file
0
tests/execution/jupyter/__init__.py
Normal file
0
tests/execution/jupyter/__init__.py
Normal file
422
tests/execution/jupyter/test_backend.py
Normal file
422
tests/execution/jupyter/test_backend.py
Normal file
|
|
@ -0,0 +1,422 @@
|
||||||
|
"""Tests for Jupyter Backend module."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from unittest.mock import Mock, MagicMock, patch
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from mcp_forge.execution.jupyter.backend import JupyterBackend
|
||||||
|
from mcp_forge.execution.jupyter.sessions import SessionManager, Session, SessionState
|
||||||
|
from mcp_forge.execution.simple.executor import ExecutionResult
|
||||||
|
from mcp_forge.config.schema import ForgeConfig, ExecutionConfig, ImageConfig, SessionConfig
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_config():
|
||||||
|
"""Mock ForgeConfig with execution settings."""
|
||||||
|
config = Mock(spec=ForgeConfig)
|
||||||
|
|
||||||
|
# Execution configuration
|
||||||
|
config.execution = Mock(spec=ExecutionConfig)
|
||||||
|
config.execution.default_timeout = 300
|
||||||
|
config.execution.max_timeout = 1800
|
||||||
|
config.execution.default_memory = "512m"
|
||||||
|
config.execution.max_memory = "2g"
|
||||||
|
config.execution.default_cpu_quota = 50000
|
||||||
|
config.execution.max_cpu_quota = 100000
|
||||||
|
|
||||||
|
# Image configuration
|
||||||
|
config.images = Mock(spec=ImageConfig)
|
||||||
|
config.images.jupyter = "mcp-forge/jupyter:latest"
|
||||||
|
|
||||||
|
# Session configuration
|
||||||
|
config.sessions = Mock(spec=SessionConfig)
|
||||||
|
config.sessions.idle_timeout = 3600
|
||||||
|
config.sessions.max_concurrent = 10
|
||||||
|
config.sessions.cleanup_interval = 300
|
||||||
|
|
||||||
|
return config
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_container_manager():
|
||||||
|
"""Mock SecureContainerManager."""
|
||||||
|
return Mock(spec=SecureContainerManager)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_audit_logger():
|
||||||
|
"""Mock AuditLogger."""
|
||||||
|
return Mock(spec=AuditLogger)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_session_manager():
|
||||||
|
"""Mock SessionManager."""
|
||||||
|
return Mock(spec=SessionManager)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def backend(mock_config, mock_container_manager, mock_audit_logger):
|
||||||
|
"""JupyterBackend instance with mocked dependencies."""
|
||||||
|
with patch('mcp_forge.execution.jupyter.backend.SessionManager') as mock_sm_class:
|
||||||
|
mock_session_manager = Mock(spec=SessionManager)
|
||||||
|
mock_sm_class.return_value = mock_session_manager
|
||||||
|
|
||||||
|
backend = JupyterBackend(
|
||||||
|
config=mock_config,
|
||||||
|
container_manager=mock_container_manager,
|
||||||
|
audit_logger=mock_audit_logger
|
||||||
|
)
|
||||||
|
backend.session_manager = mock_session_manager
|
||||||
|
|
||||||
|
return backend
|
||||||
|
|
||||||
|
|
||||||
|
def test_backend_initializes_session_manager(mock_config, mock_container_manager, mock_audit_logger):
|
||||||
|
"""Test backend creates SessionManager on initialization."""
|
||||||
|
with patch('mcp_forge.execution.jupyter.backend.SessionManager') as mock_sm_class:
|
||||||
|
with patch('mcp_forge.execution.jupyter.backend.JupyterKernelManager') as mock_km_class:
|
||||||
|
mock_session_manager = Mock(spec=SessionManager)
|
||||||
|
mock_sm_class.return_value = mock_session_manager
|
||||||
|
|
||||||
|
backend = JupyterBackend(
|
||||||
|
config=mock_config,
|
||||||
|
container_manager=mock_container_manager,
|
||||||
|
audit_logger=mock_audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify kernel manager was created
|
||||||
|
mock_km_class.assert_called_once()
|
||||||
|
|
||||||
|
# Verify session manager was created
|
||||||
|
mock_sm_class.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_creates_session_if_not_exists(backend):
|
||||||
|
"""Test execute creates new session if it doesn't exist."""
|
||||||
|
# Mock get_session to raise SessionError (session doesn't exist)
|
||||||
|
from mcp_forge.execution.jupyter.sessions import SessionError
|
||||||
|
backend.session_manager.get_session.side_effect = SessionError("Session not found")
|
||||||
|
|
||||||
|
# Mock create_session
|
||||||
|
mock_session = Mock(spec=Session)
|
||||||
|
backend.session_manager.create_session.return_value = mock_session
|
||||||
|
|
||||||
|
# Mock execute_in_session
|
||||||
|
mock_result = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="Hello",
|
||||||
|
stderr="",
|
||||||
|
result="Hello",
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
backend.session_manager.execute_in_session.return_value = mock_result
|
||||||
|
|
||||||
|
result = backend.execute("print('Hello')", session_id="test-session")
|
||||||
|
|
||||||
|
# Verify session was created
|
||||||
|
backend.session_manager.create_session.assert_called_once()
|
||||||
|
|
||||||
|
# Verify execution happened
|
||||||
|
backend.session_manager.execute_in_session.assert_called_once_with(
|
||||||
|
session_id="test-session", code="print('Hello')", timeout=300
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.success is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_reuses_existing_session(backend):
|
||||||
|
"""Test execute reuses existing session."""
|
||||||
|
# Mock get_session to return existing session
|
||||||
|
mock_session = Mock(spec=Session)
|
||||||
|
backend.session_manager.get_session.return_value = mock_session
|
||||||
|
|
||||||
|
# Mock execute_in_session
|
||||||
|
mock_result = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="42",
|
||||||
|
stderr="",
|
||||||
|
result=42,
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
backend.session_manager.execute_in_session.return_value = mock_result
|
||||||
|
|
||||||
|
result = backend.execute("21 + 21", session_id="existing-session")
|
||||||
|
|
||||||
|
# Verify session was NOT created
|
||||||
|
backend.session_manager.create_session.assert_not_called()
|
||||||
|
|
||||||
|
# Verify session was checked
|
||||||
|
backend.session_manager.get_session.assert_called_once_with("existing-session")
|
||||||
|
|
||||||
|
# Verify execution happened
|
||||||
|
backend.session_manager.execute_in_session.assert_called_once()
|
||||||
|
|
||||||
|
assert result.result == 42
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_with_custom_timeout(backend):
|
||||||
|
"""Test execute respects custom timeout parameter."""
|
||||||
|
# Mock existing session
|
||||||
|
mock_session = Mock(spec=Session)
|
||||||
|
backend.session_manager.get_session.return_value = mock_session
|
||||||
|
|
||||||
|
mock_result = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=None,
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
backend.session_manager.execute_in_session.return_value = mock_result
|
||||||
|
|
||||||
|
backend.execute("pass", session_id="test", timeout=600)
|
||||||
|
|
||||||
|
# Verify timeout was passed through
|
||||||
|
backend.session_manager.execute_in_session.assert_called_once_with(
|
||||||
|
session_id="test", code="pass", timeout=600
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_with_custom_memory(backend, mock_config):
|
||||||
|
"""Test execute creates session with custom memory limit."""
|
||||||
|
from mcp_forge.execution.jupyter.sessions import SessionError
|
||||||
|
backend.session_manager.get_session.side_effect = SessionError("Not found")
|
||||||
|
|
||||||
|
mock_session = Mock(spec=Session)
|
||||||
|
backend.session_manager.create_session.return_value = mock_session
|
||||||
|
|
||||||
|
mock_result = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=None,
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
backend.session_manager.execute_in_session.return_value = mock_result
|
||||||
|
|
||||||
|
backend.execute("pass", session_id="test", memory="1g")
|
||||||
|
|
||||||
|
# Verify session was created with custom memory
|
||||||
|
backend.session_manager.create_session.assert_called_once()
|
||||||
|
call_args = backend.session_manager.create_session.call_args
|
||||||
|
resource_limits = call_args[1]['resource_limits']
|
||||||
|
assert resource_limits.memory_bytes == 1024 * 1024 * 1024 # 1g in bytes
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_with_custom_cpu_quota(backend):
|
||||||
|
"""Test execute creates session with custom CPU quota."""
|
||||||
|
from mcp_forge.execution.jupyter.sessions import SessionError
|
||||||
|
backend.session_manager.get_session.side_effect = SessionError("Not found")
|
||||||
|
|
||||||
|
mock_session = Mock(spec=Session)
|
||||||
|
backend.session_manager.create_session.return_value = mock_session
|
||||||
|
|
||||||
|
mock_result = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=None,
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
backend.session_manager.execute_in_session.return_value = mock_result
|
||||||
|
|
||||||
|
backend.execute("pass", session_id="test", cpu_quota=75000)
|
||||||
|
|
||||||
|
# Verify session was created with custom CPU quota
|
||||||
|
backend.session_manager.create_session.assert_called_once()
|
||||||
|
call_args = backend.session_manager.create_session.call_args
|
||||||
|
resource_limits = call_args[1]['resource_limits']
|
||||||
|
assert resource_limits.cpu_quota == 75000
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_validates_timeout_against_max(backend, mock_config):
|
||||||
|
"""Test execute rejects timeout exceeding maximum."""
|
||||||
|
mock_config.execution.max_timeout = 1800
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="Timeout 3600 exceeds maximum"):
|
||||||
|
backend.execute("pass", session_id="test", timeout=3600)
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_validates_memory_against_max(backend, mock_config):
|
||||||
|
"""Test execute rejects memory exceeding maximum."""
|
||||||
|
mock_config.execution.max_memory = "2g"
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="Memory 4g exceeds maximum"):
|
||||||
|
backend.execute("pass", session_id="test", memory="4g")
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_validates_cpu_quota_against_max(backend, mock_config):
|
||||||
|
"""Test execute rejects CPU quota exceeding maximum."""
|
||||||
|
mock_config.execution.max_cpu_quota = 100000
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="CPU quota 150000 exceeds maximum"):
|
||||||
|
backend.execute("pass", session_id="test", cpu_quota=150000)
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_with_volumes(backend):
|
||||||
|
"""Test execute passes volumes to session creation."""
|
||||||
|
from mcp_forge.execution.jupyter.sessions import SessionError
|
||||||
|
backend.session_manager.get_session.side_effect = SessionError("Not found")
|
||||||
|
|
||||||
|
mock_session = Mock(spec=Session)
|
||||||
|
backend.session_manager.create_session.return_value = mock_session
|
||||||
|
|
||||||
|
mock_result = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=None,
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
backend.session_manager.execute_in_session.return_value = mock_result
|
||||||
|
|
||||||
|
volumes = {"/host/path": {"bind": "/container/path", "mode": "ro"}}
|
||||||
|
backend.execute("pass", session_id="test", volumes=volumes)
|
||||||
|
|
||||||
|
# Verify volumes were passed to create_session
|
||||||
|
backend.session_manager.create_session.assert_called_once()
|
||||||
|
call_args = backend.session_manager.create_session.call_args
|
||||||
|
assert call_args[1]['volumes'] == volumes
|
||||||
|
|
||||||
|
|
||||||
|
def test_document_state(backend):
|
||||||
|
"""Test document_state delegates to session manager."""
|
||||||
|
variables = {"x": "Input data", "y": "Output result"}
|
||||||
|
note = "Initial data load"
|
||||||
|
|
||||||
|
backend.document_state("test-session", variables, note=note, clear=False)
|
||||||
|
|
||||||
|
backend.session_manager.document_state.assert_called_once_with(
|
||||||
|
session_id="test-session", variables=variables, note=note, clear=False
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_document_state_with_clear(backend):
|
||||||
|
"""Test document_state with clear flag."""
|
||||||
|
variables = {"new_var": "New data"}
|
||||||
|
|
||||||
|
backend.document_state("test-session", variables, clear=True)
|
||||||
|
|
||||||
|
backend.session_manager.document_state.assert_called_once_with(
|
||||||
|
session_id="test-session", variables=variables, note="", clear=True
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_session_state(backend):
|
||||||
|
"""Test get_session_state delegates to session manager."""
|
||||||
|
mock_state = Mock(spec=SessionState)
|
||||||
|
backend.session_manager.get_session_state.return_value = mock_state
|
||||||
|
|
||||||
|
state = backend.get_session_state("test-session")
|
||||||
|
|
||||||
|
backend.session_manager.get_session_state.assert_called_once_with("test-session")
|
||||||
|
assert state == mock_state
|
||||||
|
|
||||||
|
|
||||||
|
def test_destroy_session(backend):
|
||||||
|
"""Test destroy_session delegates to session manager."""
|
||||||
|
backend.destroy_session("test-session")
|
||||||
|
|
||||||
|
backend.session_manager.destroy_session.assert_called_once_with("test-session")
|
||||||
|
|
||||||
|
|
||||||
|
def test_list_sessions(backend):
|
||||||
|
"""Test list_sessions delegates to session manager."""
|
||||||
|
mock_sessions = [
|
||||||
|
{"session_id": "session1", "kernel_id": "kernel1"},
|
||||||
|
{"session_id": "session2", "kernel_id": "kernel2"}
|
||||||
|
]
|
||||||
|
backend.session_manager.list_sessions.return_value = mock_sessions
|
||||||
|
|
||||||
|
sessions = backend.list_sessions()
|
||||||
|
|
||||||
|
backend.session_manager.list_sessions.assert_called_once()
|
||||||
|
assert sessions == mock_sessions
|
||||||
|
|
||||||
|
|
||||||
|
def test_cleanup_idle_sessions(backend):
|
||||||
|
"""Test cleanup_idle_sessions delegates to session manager."""
|
||||||
|
backend.session_manager.cleanup_idle_sessions.return_value = 2
|
||||||
|
|
||||||
|
count = backend.cleanup_idle_sessions()
|
||||||
|
|
||||||
|
backend.session_manager.cleanup_idle_sessions.assert_called_once()
|
||||||
|
assert count == 2
|
||||||
|
|
||||||
|
|
||||||
|
def test_default_resource_limits_from_config(mock_config, mock_container_manager, mock_audit_logger):
|
||||||
|
"""Test _default_resource_limits creates limits from config."""
|
||||||
|
with patch('mcp_forge.execution.jupyter.backend.SessionManager'):
|
||||||
|
backend = JupyterBackend(
|
||||||
|
config=mock_config,
|
||||||
|
container_manager=mock_container_manager,
|
||||||
|
audit_logger=mock_audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
limits = backend._default_resource_limits()
|
||||||
|
|
||||||
|
assert limits.memory_bytes == 512 * 1024 * 1024 # 512m
|
||||||
|
assert limits.cpu_quota == 50000
|
||||||
|
assert limits.timeout == 300
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_logs_audit_event(backend):
|
||||||
|
"""Test execute logs audit event."""
|
||||||
|
# Mock existing session
|
||||||
|
mock_session = Mock(spec=Session)
|
||||||
|
backend.session_manager.get_session.return_value = mock_session
|
||||||
|
|
||||||
|
mock_result = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=None,
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
backend.session_manager.execute_in_session.return_value = mock_result
|
||||||
|
|
||||||
|
backend.execute("print('test')", session_id="test-session")
|
||||||
|
|
||||||
|
# Verify audit log was called
|
||||||
|
backend.audit_logger.log.assert_called()
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_with_defaults_uses_config_values(backend, mock_config):
|
||||||
|
"""Test execute without parameters uses config defaults."""
|
||||||
|
from mcp_forge.execution.jupyter.sessions import SessionError
|
||||||
|
backend.session_manager.get_session.side_effect = SessionError("Not found")
|
||||||
|
|
||||||
|
mock_session = Mock(spec=Session)
|
||||||
|
backend.session_manager.create_session.return_value = mock_session
|
||||||
|
|
||||||
|
mock_result = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=None,
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
backend.session_manager.execute_in_session.return_value = mock_result
|
||||||
|
|
||||||
|
backend.execute("pass", session_id="test")
|
||||||
|
|
||||||
|
# Verify default values from config were used
|
||||||
|
call_args = backend.session_manager.create_session.call_args
|
||||||
|
resource_limits = call_args[1]['resource_limits']
|
||||||
|
assert resource_limits.memory_bytes == 512 * 1024 * 1024
|
||||||
|
assert resource_limits.cpu_quota == 50000
|
||||||
|
|
||||||
|
call_args = backend.session_manager.execute_in_session.call_args
|
||||||
|
assert call_args[1]['timeout'] == 300
|
||||||
325
tests/execution/jupyter/test_kernel.py
Normal file
325
tests/execution/jupyter/test_kernel.py
Normal file
|
|
@ -0,0 +1,325 @@
|
||||||
|
"""Tests for the Jupyter Kernel Manager module."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from unittest.mock import Mock, MagicMock, patch
|
||||||
|
from datetime import datetime, timedelta
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from mcp_forge.execution.jupyter.kernel import (
|
||||||
|
JupyterKernelManager,
|
||||||
|
KernelInfo,
|
||||||
|
KernelError
|
||||||
|
)
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
from mcp_forge.execution.simple.executor import ExecutionResult
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def resource_limits():
|
||||||
|
"""Standard resource limits for testing."""
|
||||||
|
return ResourceLimits(
|
||||||
|
memory="512m",
|
||||||
|
cpu_quota=50000,
|
||||||
|
storage="1g",
|
||||||
|
timeout=300
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_container_manager():
|
||||||
|
"""Mock SecureContainerManager."""
|
||||||
|
manager = Mock(spec=SecureContainerManager)
|
||||||
|
manager.create_container.return_value = "test-container-123"
|
||||||
|
manager.start_container.return_value = None
|
||||||
|
manager.stop_container.return_value = None
|
||||||
|
manager.remove_container.return_value = None
|
||||||
|
manager.get_container_logs.return_value = ("", "")
|
||||||
|
return manager
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def kernel_manager(mock_container_manager, resource_limits):
|
||||||
|
"""JupyterKernelManager instance with mocked dependencies."""
|
||||||
|
return JupyterKernelManager(
|
||||||
|
container_manager=mock_container_manager,
|
||||||
|
image="mcp-forge/jupyter:latest",
|
||||||
|
resource_limits=resource_limits
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_start_kernel_creates_container(kernel_manager, mock_container_manager):
|
||||||
|
"""Test start_kernel creates and starts a container."""
|
||||||
|
kernel_id = kernel_manager.start_kernel("session-1")
|
||||||
|
|
||||||
|
assert kernel_id is not None
|
||||||
|
assert kernel_id.startswith("kernel-")
|
||||||
|
|
||||||
|
# Verify container was created and started
|
||||||
|
mock_container_manager.create_container.assert_called_once()
|
||||||
|
mock_container_manager.start_container.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
def test_start_kernel_returns_kernel_info(kernel_manager):
|
||||||
|
"""Test start_kernel returns valid kernel info."""
|
||||||
|
kernel_id = kernel_manager.start_kernel("session-1")
|
||||||
|
|
||||||
|
# Kernel should be registered
|
||||||
|
assert kernel_id in kernel_manager.kernels
|
||||||
|
|
||||||
|
kernel_info = kernel_manager.kernels[kernel_id]
|
||||||
|
assert kernel_info.kernel_id == kernel_id
|
||||||
|
assert kernel_info.container_id == "test-container-123"
|
||||||
|
assert kernel_info.session_id == "session-1"
|
||||||
|
assert isinstance(kernel_info.started_at, datetime)
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_in_kernel_returns_result(kernel_manager):
|
||||||
|
"""Test execute_code runs code and returns result."""
|
||||||
|
kernel_id = kernel_manager.start_kernel("session-1")
|
||||||
|
|
||||||
|
result = kernel_manager.execute_code(kernel_id, "2 + 2")
|
||||||
|
|
||||||
|
assert isinstance(result, ExecutionResult)
|
||||||
|
assert result.success is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_with_nonexistent_kernel_raises_error(kernel_manager):
|
||||||
|
"""Test execute_code raises error for nonexistent kernel."""
|
||||||
|
with pytest.raises(KernelError, match="Kernel.*not found"):
|
||||||
|
kernel_manager.execute_code("nonexistent-kernel", "pass")
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_preserves_namespace(kernel_manager):
|
||||||
|
"""Test namespace persists between executions."""
|
||||||
|
kernel_id = kernel_manager.start_kernel("session-1")
|
||||||
|
|
||||||
|
# Set a variable
|
||||||
|
result1 = kernel_manager.execute_code(kernel_id, "x = 42")
|
||||||
|
assert result1.success is True
|
||||||
|
|
||||||
|
# Access the variable
|
||||||
|
result2 = kernel_manager.execute_code(kernel_id, "x")
|
||||||
|
assert result2.success is True
|
||||||
|
# In real implementation, result2.result would be 42
|
||||||
|
|
||||||
|
|
||||||
|
def test_shutdown_kernel_removes_container(kernel_manager, mock_container_manager):
|
||||||
|
"""Test shutdown_kernel cleans up container."""
|
||||||
|
kernel_id = kernel_manager.start_kernel("session-1")
|
||||||
|
|
||||||
|
kernel_manager.shutdown_kernel(kernel_id)
|
||||||
|
|
||||||
|
# Verify container was stopped and removed
|
||||||
|
mock_container_manager.stop_container.assert_called_once_with("test-container-123", timeout=10)
|
||||||
|
mock_container_manager.remove_container.assert_called_once_with("test-container-123")
|
||||||
|
|
||||||
|
# Kernel should be removed from registry
|
||||||
|
assert kernel_id not in kernel_manager.kernels
|
||||||
|
|
||||||
|
|
||||||
|
def test_shutdown_nonexistent_kernel_raises_error(kernel_manager):
|
||||||
|
"""Test shutdown_kernel raises error for nonexistent kernel."""
|
||||||
|
with pytest.raises(KernelError, match="Kernel.*not found"):
|
||||||
|
kernel_manager.shutdown_kernel("nonexistent-kernel")
|
||||||
|
|
||||||
|
|
||||||
|
def test_inspect_namespace_returns_variables(kernel_manager):
|
||||||
|
"""Test inspect_namespace returns list of variables."""
|
||||||
|
kernel_id = kernel_manager.start_kernel("session-1")
|
||||||
|
|
||||||
|
# Execute some code to create variables
|
||||||
|
kernel_manager.execute_code(kernel_id, "x = 1; y = 2; z = 3")
|
||||||
|
|
||||||
|
variables = kernel_manager.inspect_namespace(kernel_id)
|
||||||
|
|
||||||
|
assert isinstance(variables, list)
|
||||||
|
# In real implementation, would contain ['x', 'y', 'z']
|
||||||
|
|
||||||
|
|
||||||
|
def test_inspect_namespace_filters_private_vars(kernel_manager):
|
||||||
|
"""Test inspect_namespace filters out private variables."""
|
||||||
|
kernel_id = kernel_manager.start_kernel("session-1")
|
||||||
|
|
||||||
|
kernel_manager.execute_code(kernel_id, "x = 1; _private = 2; __dunder__ = 3")
|
||||||
|
|
||||||
|
variables = kernel_manager.inspect_namespace(kernel_id)
|
||||||
|
|
||||||
|
# Private variables should be filtered
|
||||||
|
# In real implementation: assert '_private' not in variables
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_variable_info_returns_metadata(kernel_manager):
|
||||||
|
"""Test get_variable_info returns variable metadata."""
|
||||||
|
kernel_id = kernel_manager.start_kernel("session-1")
|
||||||
|
|
||||||
|
kernel_manager.execute_code(kernel_id, "x = [1, 2, 3, 4, 5]")
|
||||||
|
|
||||||
|
info = kernel_manager.get_variable_info(kernel_id, "x")
|
||||||
|
|
||||||
|
assert isinstance(info, dict)
|
||||||
|
assert "type" in info
|
||||||
|
# In real implementation: assert info["type"] == "list"
|
||||||
|
|
||||||
|
|
||||||
|
def test_restart_kernel_resets_namespace(kernel_manager, mock_container_manager):
|
||||||
|
"""Test restart_kernel resets namespace but keeps container."""
|
||||||
|
kernel_id = kernel_manager.start_kernel("session-1")
|
||||||
|
original_container_id = kernel_manager.kernels[kernel_id].container_id
|
||||||
|
|
||||||
|
# Set a variable
|
||||||
|
kernel_manager.execute_code(kernel_id, "x = 42")
|
||||||
|
|
||||||
|
# Restart
|
||||||
|
kernel_manager.restart_kernel(kernel_id)
|
||||||
|
|
||||||
|
# Container should be the same
|
||||||
|
assert kernel_manager.kernels[kernel_id].container_id == original_container_id
|
||||||
|
|
||||||
|
# Namespace should be reset (variable no longer accessible)
|
||||||
|
# In real implementation, executing "x" would raise NameError
|
||||||
|
|
||||||
|
|
||||||
|
def test_cleanup_idle_kernels_removes_old_kernels(kernel_manager, mock_container_manager):
|
||||||
|
"""Test cleanup_idle_kernels removes kernels idle too long."""
|
||||||
|
# Start two kernels
|
||||||
|
kernel1 = kernel_manager.start_kernel("session-1")
|
||||||
|
kernel2 = kernel_manager.start_kernel("session-2")
|
||||||
|
|
||||||
|
# Make kernel1 appear old
|
||||||
|
kernel_manager.kernels[kernel1].last_activity = datetime.utcnow() - timedelta(hours=2)
|
||||||
|
|
||||||
|
# Cleanup kernels idle > 1 hour
|
||||||
|
count = kernel_manager.cleanup_idle_kernels(timedelta(hours=1))
|
||||||
|
|
||||||
|
assert count == 1
|
||||||
|
assert kernel1 not in kernel_manager.kernels
|
||||||
|
assert kernel2 in kernel_manager.kernels
|
||||||
|
|
||||||
|
|
||||||
|
def test_kernel_with_volumes(kernel_manager, mock_container_manager):
|
||||||
|
"""Test kernel can be started with volume mounts."""
|
||||||
|
volumes = {
|
||||||
|
"/mcp-forge/sessions/session-1/workspace": {"bind": "/workspace", "mode": "rw"}
|
||||||
|
}
|
||||||
|
|
||||||
|
kernel_id = kernel_manager.start_kernel("session-1", volumes=volumes)
|
||||||
|
|
||||||
|
assert kernel_id is not None
|
||||||
|
# Verify volumes were passed to container creation
|
||||||
|
call_args = mock_container_manager.create_container.call_args
|
||||||
|
# In real implementation, would verify volumes in ContainerConfig
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_with_timeout(kernel_manager):
|
||||||
|
"""Test execute_code respects timeout parameter."""
|
||||||
|
kernel_id = kernel_manager.start_kernel("session-1")
|
||||||
|
|
||||||
|
# Execute with custom timeout
|
||||||
|
result = kernel_manager.execute_code(kernel_id, "import time; time.sleep(0.1)", timeout=10)
|
||||||
|
|
||||||
|
assert isinstance(result, ExecutionResult)
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_handles_syntax_error(kernel_manager):
|
||||||
|
"""Test execute_code handles syntax errors gracefully."""
|
||||||
|
kernel_id = kernel_manager.start_kernel("session-1")
|
||||||
|
|
||||||
|
result = kernel_manager.execute_code(kernel_id, "def foo( :")
|
||||||
|
|
||||||
|
assert result.success is False
|
||||||
|
assert result.error is not None
|
||||||
|
assert "SyntaxError" in result.error or "syntax" in result.error.lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_handles_runtime_error(kernel_manager):
|
||||||
|
"""Test execute_code handles runtime errors gracefully."""
|
||||||
|
kernel_id = kernel_manager.start_kernel("session-1")
|
||||||
|
|
||||||
|
result = kernel_manager.execute_code(kernel_id, "1 / 0")
|
||||||
|
|
||||||
|
assert result.success is False
|
||||||
|
assert result.error is not None
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_captures_stdout(kernel_manager):
|
||||||
|
"""Test execute_code captures stdout output."""
|
||||||
|
kernel_id = kernel_manager.start_kernel("session-1")
|
||||||
|
|
||||||
|
result = kernel_manager.execute_code(kernel_id, 'print("Hello, World!")')
|
||||||
|
|
||||||
|
assert result.success is True
|
||||||
|
# In real implementation: assert "Hello, World!" in result.stdout
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_captures_stderr(kernel_manager):
|
||||||
|
"""Test execute_code captures stderr output."""
|
||||||
|
kernel_id = kernel_manager.start_kernel("session-1")
|
||||||
|
|
||||||
|
result = kernel_manager.execute_code(kernel_id, 'import sys; print("warning", file=sys.stderr)')
|
||||||
|
|
||||||
|
assert result.success is True
|
||||||
|
# In real implementation: assert "warning" in result.stderr
|
||||||
|
|
||||||
|
|
||||||
|
def test_multiple_kernels_are_isolated(kernel_manager):
|
||||||
|
"""Test multiple kernels have isolated namespaces."""
|
||||||
|
kernel1 = kernel_manager.start_kernel("session-1")
|
||||||
|
kernel2 = kernel_manager.start_kernel("session-2")
|
||||||
|
|
||||||
|
# Set variable in kernel1
|
||||||
|
kernel_manager.execute_code(kernel1, "x = 1")
|
||||||
|
|
||||||
|
# Set different value in kernel2
|
||||||
|
kernel_manager.execute_code(kernel2, "x = 2")
|
||||||
|
|
||||||
|
# Values should be independent
|
||||||
|
result1 = kernel_manager.execute_code(kernel1, "x")
|
||||||
|
result2 = kernel_manager.execute_code(kernel2, "x")
|
||||||
|
|
||||||
|
# In real implementation: verify result1.result == 1 and result2.result == 2
|
||||||
|
|
||||||
|
|
||||||
|
def test_kernel_info_to_dict(resource_limits):
|
||||||
|
"""Test KernelInfo.to_dict() serialization."""
|
||||||
|
now = datetime.utcnow()
|
||||||
|
kernel_info = KernelInfo(
|
||||||
|
kernel_id="kernel-123",
|
||||||
|
container_id="container-456",
|
||||||
|
session_id="session-789",
|
||||||
|
started_at=now,
|
||||||
|
last_activity=now
|
||||||
|
)
|
||||||
|
|
||||||
|
info_dict = kernel_info.to_dict()
|
||||||
|
|
||||||
|
assert isinstance(info_dict, dict)
|
||||||
|
assert info_dict["kernel_id"] == "kernel-123"
|
||||||
|
assert info_dict["container_id"] == "container-456"
|
||||||
|
assert info_dict["session_id"] == "session-789"
|
||||||
|
assert "started_at" in info_dict
|
||||||
|
assert "last_activity" in info_dict
|
||||||
|
|
||||||
|
|
||||||
|
def test_start_kernel_with_session_id_tracking(kernel_manager):
|
||||||
|
"""Test kernel tracks session_id correctly."""
|
||||||
|
kernel_id = kernel_manager.start_kernel("my-session")
|
||||||
|
|
||||||
|
kernel_info = kernel_manager.kernels[kernel_id]
|
||||||
|
assert kernel_info.session_id == "my-session"
|
||||||
|
|
||||||
|
|
||||||
|
def test_update_activity_timestamp(kernel_manager):
|
||||||
|
"""Test executing code updates last_activity timestamp."""
|
||||||
|
kernel_id = kernel_manager.start_kernel("session-1")
|
||||||
|
|
||||||
|
original_activity = kernel_manager.kernels[kernel_id].last_activity
|
||||||
|
|
||||||
|
# Small delay to ensure timestamp difference
|
||||||
|
import time
|
||||||
|
time.sleep(0.01)
|
||||||
|
|
||||||
|
kernel_manager.execute_code(kernel_id, "pass")
|
||||||
|
|
||||||
|
new_activity = kernel_manager.kernels[kernel_id].last_activity
|
||||||
|
assert new_activity > original_activity
|
||||||
371
tests/execution/jupyter/test_sessions.py
Normal file
371
tests/execution/jupyter/test_sessions.py
Normal file
|
|
@ -0,0 +1,371 @@
|
||||||
|
"""Tests for the Session Manager module."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from unittest.mock import Mock, MagicMock
|
||||||
|
from datetime import datetime, timedelta
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from mcp_forge.execution.jupyter.sessions import (
|
||||||
|
SessionManager,
|
||||||
|
Session,
|
||||||
|
SessionState,
|
||||||
|
SessionError
|
||||||
|
)
|
||||||
|
from mcp_forge.execution.jupyter.kernel import JupyterKernelManager
|
||||||
|
from mcp_forge.config.schema import SessionConfig
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
from mcp_forge.execution.simple.executor import ExecutionResult
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def session_config():
|
||||||
|
"""Mock SessionConfig."""
|
||||||
|
config = Mock(spec=SessionConfig)
|
||||||
|
config.idle_timeout = 3600
|
||||||
|
config.max_concurrent = 10
|
||||||
|
config.cleanup_interval = 300
|
||||||
|
return config
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_kernel_manager():
|
||||||
|
"""Mock JupyterKernelManager."""
|
||||||
|
manager = Mock(spec=JupyterKernelManager)
|
||||||
|
manager.start_kernel.return_value = "kernel-123"
|
||||||
|
manager.execute_code.return_value = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=None,
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
manager.shutdown_kernel.return_value = None
|
||||||
|
manager.inspect_namespace.return_value = []
|
||||||
|
return manager
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_audit_logger(tmp_path):
|
||||||
|
"""Mock AuditLogger."""
|
||||||
|
log_file = tmp_path / "audit.log"
|
||||||
|
return Mock(spec=AuditLogger)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def resource_limits():
|
||||||
|
"""Standard resource limits."""
|
||||||
|
return ResourceLimits(
|
||||||
|
memory="512m",
|
||||||
|
cpu_quota=50000,
|
||||||
|
storage="1g",
|
||||||
|
timeout=300
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def session_manager(session_config, mock_kernel_manager, mock_audit_logger):
|
||||||
|
"""SessionManager instance with mocked dependencies."""
|
||||||
|
return SessionManager(
|
||||||
|
config=session_config,
|
||||||
|
kernel_manager=mock_kernel_manager,
|
||||||
|
audit_logger=mock_audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_session_starts_kernel(session_manager, mock_kernel_manager, resource_limits):
|
||||||
|
"""Test create_session starts a kernel."""
|
||||||
|
session = session_manager.create_session("session-1", resource_limits)
|
||||||
|
|
||||||
|
assert session.session_id == "session-1"
|
||||||
|
assert session.kernel_id == "kernel-123"
|
||||||
|
mock_kernel_manager.start_kernel.assert_called_once_with("session-1", volumes=None)
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_session_with_duplicate_id_raises_error(session_manager, resource_limits):
|
||||||
|
"""Test creating session with existing ID raises error."""
|
||||||
|
session_manager.create_session("session-1", resource_limits)
|
||||||
|
|
||||||
|
with pytest.raises(SessionError, match="already exists"):
|
||||||
|
session_manager.create_session("session-1", resource_limits)
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_session_enforces_max_concurrent(session_manager, session_config, resource_limits):
|
||||||
|
"""Test max concurrent sessions is enforced."""
|
||||||
|
session_config.max_concurrent = 2
|
||||||
|
|
||||||
|
# Create 2 sessions (at limit)
|
||||||
|
session_manager.create_session("session-1", resource_limits)
|
||||||
|
session_manager.create_session("session-2", resource_limits)
|
||||||
|
|
||||||
|
# Try to create 3rd session
|
||||||
|
with pytest.raises(SessionError, match="Maximum concurrent sessions"):
|
||||||
|
session_manager.create_session("session-3", resource_limits)
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_session_returns_existing_session(session_manager, resource_limits):
|
||||||
|
"""Test get_session returns existing session."""
|
||||||
|
created = session_manager.create_session("session-1", resource_limits)
|
||||||
|
retrieved = session_manager.get_session("session-1")
|
||||||
|
|
||||||
|
assert retrieved.session_id == created.session_id
|
||||||
|
assert retrieved.kernel_id == created.kernel_id
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_session_raises_error_for_nonexistent(session_manager):
|
||||||
|
"""Test get_session raises error for nonexistent session."""
|
||||||
|
with pytest.raises(SessionError, match="not found"):
|
||||||
|
session_manager.get_session("nonexistent")
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_in_session_runs_code(session_manager, mock_kernel_manager, resource_limits):
|
||||||
|
"""Test execute_in_session runs code in kernel."""
|
||||||
|
session_manager.create_session("session-1", resource_limits)
|
||||||
|
|
||||||
|
result = session_manager.execute_in_session("session-1", "x = 42")
|
||||||
|
|
||||||
|
assert result.success is True
|
||||||
|
mock_kernel_manager.execute_code.assert_called_once_with("kernel-123", "x = 42", timeout=300)
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_in_session_updates_activity(session_manager, resource_limits):
|
||||||
|
"""Test execute_in_session updates last activity timestamp."""
|
||||||
|
session_manager.create_session("session-1", resource_limits)
|
||||||
|
session = session_manager.get_session("session-1")
|
||||||
|
|
||||||
|
original_activity = session.last_activity
|
||||||
|
|
||||||
|
import time
|
||||||
|
time.sleep(0.01)
|
||||||
|
|
||||||
|
session_manager.execute_in_session("session-1", "pass")
|
||||||
|
|
||||||
|
assert session.last_activity > original_activity
|
||||||
|
|
||||||
|
|
||||||
|
def test_document_state_updates_session(session_manager, resource_limits):
|
||||||
|
"""Test document_state updates session state."""
|
||||||
|
session_manager.create_session("session-1", resource_limits)
|
||||||
|
|
||||||
|
variables = {"x": "The result of computation", "y": "Another variable"}
|
||||||
|
session_manager.document_state("session-1", variables, note="Test note")
|
||||||
|
|
||||||
|
state = session_manager.get_session_state("session-1")
|
||||||
|
|
||||||
|
assert state.documented_variables == variables
|
||||||
|
assert state.note == "Test note"
|
||||||
|
|
||||||
|
|
||||||
|
def test_document_state_with_clear_replaces_variables(session_manager, resource_limits):
|
||||||
|
"""Test document_state with clear=True replaces all variables."""
|
||||||
|
session_manager.create_session("session-1", resource_limits)
|
||||||
|
|
||||||
|
# Set initial variables
|
||||||
|
session_manager.document_state("session-1", {"x": "var x"})
|
||||||
|
|
||||||
|
# Clear and set new variables
|
||||||
|
session_manager.document_state("session-1", {"y": "var y"}, clear=True)
|
||||||
|
|
||||||
|
state = session_manager.get_session_state("session-1")
|
||||||
|
|
||||||
|
assert "x" not in state.documented_variables
|
||||||
|
assert "y" in state.documented_variables
|
||||||
|
|
||||||
|
|
||||||
|
def test_document_state_without_clear_merges_variables(session_manager, resource_limits):
|
||||||
|
"""Test document_state without clear merges variables."""
|
||||||
|
session_manager.create_session("session-1", resource_limits)
|
||||||
|
|
||||||
|
session_manager.document_state("session-1", {"x": "var x"})
|
||||||
|
session_manager.document_state("session-1", {"y": "var y"})
|
||||||
|
|
||||||
|
state = session_manager.get_session_state("session-1")
|
||||||
|
|
||||||
|
assert "x" in state.documented_variables
|
||||||
|
assert "y" in state.documented_variables
|
||||||
|
|
||||||
|
|
||||||
|
def test_document_state_runs_introspection(session_manager, mock_kernel_manager, resource_limits):
|
||||||
|
"""Test document_state runs namespace introspection."""
|
||||||
|
mock_kernel_manager.inspect_namespace.return_value = ["x", "y", "z"]
|
||||||
|
mock_kernel_manager.get_variable_info.return_value = {"type": "int", "repr": "42"}
|
||||||
|
|
||||||
|
session_manager.create_session("session-1", resource_limits)
|
||||||
|
session_manager.document_state("session-1", {"x": "documented"})
|
||||||
|
|
||||||
|
state = session_manager.get_session_state("session-1")
|
||||||
|
|
||||||
|
assert state.all_variables == ["x", "y", "z"]
|
||||||
|
mock_kernel_manager.inspect_namespace.assert_called_once_with("kernel-123")
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_session_state_returns_state(session_manager, resource_limits):
|
||||||
|
"""Test get_session_state returns SessionState."""
|
||||||
|
session_manager.create_session("session-1", resource_limits)
|
||||||
|
|
||||||
|
state = session_manager.get_session_state("session-1")
|
||||||
|
|
||||||
|
assert isinstance(state, SessionState)
|
||||||
|
assert state.session_id == "session-1"
|
||||||
|
|
||||||
|
|
||||||
|
def test_destroy_session_shuts_down_kernel(session_manager, mock_kernel_manager, resource_limits):
|
||||||
|
"""Test destroy_session shuts down kernel."""
|
||||||
|
session_manager.create_session("session-1", resource_limits)
|
||||||
|
|
||||||
|
session_manager.destroy_session("session-1")
|
||||||
|
|
||||||
|
mock_kernel_manager.shutdown_kernel.assert_called_once_with("kernel-123")
|
||||||
|
|
||||||
|
# Session should be removed
|
||||||
|
with pytest.raises(SessionError):
|
||||||
|
session_manager.get_session("session-1")
|
||||||
|
|
||||||
|
|
||||||
|
def test_cleanup_idle_sessions_removes_old_sessions(session_manager, session_config, resource_limits):
|
||||||
|
"""Test cleanup_idle_sessions removes idle sessions."""
|
||||||
|
session_config.idle_timeout = 3600 # 1 hour
|
||||||
|
|
||||||
|
# Create two sessions
|
||||||
|
session_manager.create_session("session-1", resource_limits)
|
||||||
|
session_manager.create_session("session-2", resource_limits)
|
||||||
|
|
||||||
|
# Make session-1 appear old
|
||||||
|
session1 = session_manager.get_session("session-1")
|
||||||
|
session1.last_activity = datetime.utcnow() - timedelta(hours=2)
|
||||||
|
|
||||||
|
# Cleanup
|
||||||
|
count = session_manager.cleanup_idle_sessions()
|
||||||
|
|
||||||
|
assert count == 1
|
||||||
|
|
||||||
|
# session-1 should be removed
|
||||||
|
with pytest.raises(SessionError):
|
||||||
|
session_manager.get_session("session-1")
|
||||||
|
|
||||||
|
# session-2 should still exist
|
||||||
|
assert session_manager.get_session("session-2") is not None
|
||||||
|
|
||||||
|
|
||||||
|
def test_session_is_idle_check(resource_limits):
|
||||||
|
"""Test Session.is_idle() check."""
|
||||||
|
now = datetime.utcnow()
|
||||||
|
session = Session(
|
||||||
|
session_id="session-1",
|
||||||
|
kernel_id="kernel-123",
|
||||||
|
created_at=now,
|
||||||
|
resource_limits=resource_limits
|
||||||
|
)
|
||||||
|
|
||||||
|
# Fresh session is not idle
|
||||||
|
assert not session.is_idle(timedelta(hours=1))
|
||||||
|
|
||||||
|
# Make it old
|
||||||
|
session.last_activity = now - timedelta(hours=2)
|
||||||
|
|
||||||
|
# Now it's idle
|
||||||
|
assert session.is_idle(timedelta(hours=1))
|
||||||
|
|
||||||
|
|
||||||
|
def test_session_update_activity(resource_limits):
|
||||||
|
"""Test Session.update_activity() updates timestamp."""
|
||||||
|
now = datetime.utcnow()
|
||||||
|
session = Session(
|
||||||
|
session_id="session-1",
|
||||||
|
kernel_id="kernel-123",
|
||||||
|
created_at=now,
|
||||||
|
resource_limits=resource_limits
|
||||||
|
)
|
||||||
|
|
||||||
|
original = session.last_activity
|
||||||
|
|
||||||
|
import time
|
||||||
|
time.sleep(0.01)
|
||||||
|
|
||||||
|
session.update_activity()
|
||||||
|
|
||||||
|
assert session.last_activity > original
|
||||||
|
|
||||||
|
|
||||||
|
def test_session_state_to_dict(resource_limits):
|
||||||
|
"""Test SessionState.to_dict() serialization."""
|
||||||
|
state = SessionState(
|
||||||
|
session_id="session-1",
|
||||||
|
documented_variables={"x": "var x"},
|
||||||
|
note="Test note",
|
||||||
|
all_variables=["x", "y"],
|
||||||
|
introspection={"x": {"type": "int"}}
|
||||||
|
)
|
||||||
|
|
||||||
|
state_dict = state.to_dict()
|
||||||
|
|
||||||
|
assert isinstance(state_dict, dict)
|
||||||
|
assert state_dict["session_id"] == "session-1"
|
||||||
|
assert state_dict["documented_variables"] == {"x": "var x"}
|
||||||
|
assert state_dict["note"] == "Test note"
|
||||||
|
assert state_dict["all_variables"] == ["x", "y"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_session_with_volumes(session_manager, mock_kernel_manager, resource_limits):
|
||||||
|
"""Test create_session passes volumes to kernel manager."""
|
||||||
|
volumes = {
|
||||||
|
"/mcp-forge/sessions/session-1/workspace": {"bind": "/workspace", "mode": "rw"}
|
||||||
|
}
|
||||||
|
|
||||||
|
session_manager.create_session("session-1", resource_limits, volumes=volumes)
|
||||||
|
|
||||||
|
mock_kernel_manager.start_kernel.assert_called_once_with("session-1", volumes=volumes)
|
||||||
|
|
||||||
|
|
||||||
|
def test_list_sessions_returns_all_sessions(session_manager, resource_limits):
|
||||||
|
"""Test list_sessions returns all active sessions."""
|
||||||
|
session_manager.create_session("session-1", resource_limits)
|
||||||
|
session_manager.create_session("session-2", resource_limits)
|
||||||
|
|
||||||
|
sessions = session_manager.list_sessions()
|
||||||
|
|
||||||
|
assert len(sessions) == 2
|
||||||
|
session_ids = [s["session_id"] for s in sessions]
|
||||||
|
assert "session-1" in session_ids
|
||||||
|
assert "session-2" in session_ids
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_in_nonexistent_session_raises_error(session_manager):
|
||||||
|
"""Test execute_in_session raises error for nonexistent session."""
|
||||||
|
with pytest.raises(SessionError, match="not found"):
|
||||||
|
session_manager.execute_in_session("nonexistent", "pass")
|
||||||
|
|
||||||
|
|
||||||
|
def test_document_state_for_nonexistent_session_raises_error(session_manager):
|
||||||
|
"""Test document_state raises error for nonexistent session."""
|
||||||
|
with pytest.raises(SessionError, match="not found"):
|
||||||
|
session_manager.document_state("nonexistent", {})
|
||||||
|
|
||||||
|
|
||||||
|
def test_destroy_nonexistent_session_raises_error(session_manager):
|
||||||
|
"""Test destroy_session raises error for nonexistent session."""
|
||||||
|
with pytest.raises(SessionError, match="not found"):
|
||||||
|
session_manager.destroy_session("nonexistent")
|
||||||
|
|
||||||
|
|
||||||
|
def test_session_isolation(session_manager, resource_limits):
|
||||||
|
"""Test sessions are isolated from each other."""
|
||||||
|
session_manager.create_session("session-1", resource_limits)
|
||||||
|
session_manager.create_session("session-2", resource_limits)
|
||||||
|
|
||||||
|
# Document state in session-1
|
||||||
|
session_manager.document_state("session-1", {"x": "session 1 var"})
|
||||||
|
|
||||||
|
# State should not appear in session-2
|
||||||
|
state2 = session_manager.get_session_state("session-2")
|
||||||
|
assert "x" not in state2.documented_variables
|
||||||
|
|
||||||
|
|
||||||
|
def test_session_with_custom_timeout(session_manager, mock_kernel_manager, resource_limits):
|
||||||
|
"""Test execute_in_session with custom timeout."""
|
||||||
|
session_manager.create_session("session-1", resource_limits)
|
||||||
|
|
||||||
|
session_manager.execute_in_session("session-1", "pass", timeout=600)
|
||||||
|
|
||||||
|
mock_kernel_manager.execute_code.assert_called_once_with("kernel-123", "pass", timeout=600)
|
||||||
0
tests/execution/simple/__init__.py
Normal file
0
tests/execution/simple/__init__.py
Normal file
395
tests/execution/simple/test_backend.py
Normal file
395
tests/execution/simple/test_backend.py
Normal file
|
|
@ -0,0 +1,395 @@
|
||||||
|
"""Tests for the Simple Backend module."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from unittest.mock import Mock, MagicMock, patch
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from mcp_forge.execution.simple.backend import SimpleBackend
|
||||||
|
from mcp_forge.execution.simple.executor import ExecutionResult
|
||||||
|
from mcp_forge.config.schema import ForgeConfig, ExecutionConfig, ImageConfig
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_config():
|
||||||
|
"""Mock ForgeConfig with execution settings."""
|
||||||
|
config = Mock(spec=ForgeConfig)
|
||||||
|
|
||||||
|
# Execution configuration
|
||||||
|
config.execution = Mock(spec=ExecutionConfig)
|
||||||
|
config.execution.default_timeout = 300
|
||||||
|
config.execution.max_timeout = 1800
|
||||||
|
config.execution.default_memory = "512m"
|
||||||
|
config.execution.max_memory = "2g"
|
||||||
|
config.execution.default_cpu_quota = 50000
|
||||||
|
config.execution.max_cpu_quota = 100000
|
||||||
|
|
||||||
|
# Image configuration
|
||||||
|
config.images = Mock(spec=ImageConfig)
|
||||||
|
config.images.python_3_11 = "mcp-forge/python:3.11"
|
||||||
|
config.images.python_3_12 = "mcp-forge/python:3.12"
|
||||||
|
|
||||||
|
return config
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_container_manager():
|
||||||
|
"""Mock SecureContainerManager."""
|
||||||
|
return Mock(spec=SecureContainerManager)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_audit_logger(tmp_path):
|
||||||
|
"""Mock AuditLogger."""
|
||||||
|
log_file = tmp_path / "audit.log"
|
||||||
|
return Mock(spec=AuditLogger)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def backend(mock_config, mock_container_manager, mock_audit_logger):
|
||||||
|
"""SimpleBackend instance with mocked dependencies."""
|
||||||
|
return SimpleBackend(
|
||||||
|
config=mock_config,
|
||||||
|
container_manager=mock_container_manager,
|
||||||
|
audit_logger=mock_audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_without_custom_params_uses_defaults(backend, mock_config):
|
||||||
|
"""Test execute uses configuration defaults when no params specified."""
|
||||||
|
# Mock executor to return a result
|
||||||
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
||||||
|
mock_executor = Mock()
|
||||||
|
mock_executor_class.return_value = mock_executor
|
||||||
|
mock_executor.execute.return_value = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=42,
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
|
||||||
|
result = backend.execute("2 + 2")
|
||||||
|
|
||||||
|
# Verify executor was created with default limits
|
||||||
|
mock_executor_class.assert_called_once()
|
||||||
|
args, kwargs = mock_executor_class.call_args
|
||||||
|
|
||||||
|
# Check resource limits
|
||||||
|
resource_limits = kwargs.get('resource_limits')
|
||||||
|
assert resource_limits is not None
|
||||||
|
assert resource_limits.memory_bytes == 512 * 1024 * 1024 # 512m in bytes
|
||||||
|
assert resource_limits.cpu_quota == 50000
|
||||||
|
assert resource_limits.timeout == 300
|
||||||
|
|
||||||
|
# Check image
|
||||||
|
assert kwargs.get('image') == "mcp-forge/python:3.11"
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_with_custom_timeout(backend, mock_config):
|
||||||
|
"""Test execute respects custom timeout parameter."""
|
||||||
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
||||||
|
mock_executor = Mock()
|
||||||
|
mock_executor_class.return_value = mock_executor
|
||||||
|
mock_executor.execute.return_value = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=None,
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
|
||||||
|
backend.execute("pass", timeout=600)
|
||||||
|
|
||||||
|
# Verify resource limits include custom timeout
|
||||||
|
args, kwargs = mock_executor_class.call_args
|
||||||
|
resource_limits = kwargs.get('resource_limits')
|
||||||
|
assert resource_limits.timeout == 600
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_with_custom_memory(backend, mock_config):
|
||||||
|
"""Test execute respects custom memory parameter."""
|
||||||
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
||||||
|
mock_executor = Mock()
|
||||||
|
mock_executor_class.return_value = mock_executor
|
||||||
|
mock_executor.execute.return_value = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=None,
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
|
||||||
|
backend.execute("pass", memory="1g")
|
||||||
|
|
||||||
|
# Verify resource limits include custom memory
|
||||||
|
args, kwargs = mock_executor_class.call_args
|
||||||
|
resource_limits = kwargs.get('resource_limits')
|
||||||
|
assert resource_limits.memory_bytes == 1024 * 1024 * 1024 # 1g in bytes
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_with_custom_cpu_quota(backend, mock_config):
|
||||||
|
"""Test execute respects custom CPU quota parameter."""
|
||||||
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
||||||
|
mock_executor = Mock()
|
||||||
|
mock_executor_class.return_value = mock_executor
|
||||||
|
mock_executor.execute.return_value = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=None,
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
|
||||||
|
backend.execute("pass", cpu_quota=75000)
|
||||||
|
|
||||||
|
# Verify resource limits include custom CPU quota
|
||||||
|
args, kwargs = mock_executor_class.call_args
|
||||||
|
resource_limits = kwargs.get('resource_limits')
|
||||||
|
assert resource_limits.cpu_quota == 75000
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_with_custom_image(backend, mock_config):
|
||||||
|
"""Test execute respects custom image parameter."""
|
||||||
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
||||||
|
mock_executor = Mock()
|
||||||
|
mock_executor_class.return_value = mock_executor
|
||||||
|
mock_executor.execute.return_value = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=None,
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
|
||||||
|
backend.execute("pass", custom_image="mcp-forge/python:3.12")
|
||||||
|
|
||||||
|
# Verify correct image was used
|
||||||
|
args, kwargs = mock_executor_class.call_args
|
||||||
|
assert kwargs.get('image') == "mcp-forge/python:3.12"
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_validates_timeout_against_max(backend, mock_config):
|
||||||
|
"""Test execute rejects timeout exceeding max."""
|
||||||
|
with pytest.raises(ValueError, match="timeout.*exceeds maximum"):
|
||||||
|
backend.execute("pass", timeout=2000) # max is 1800
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_validates_memory_against_max(backend, mock_config):
|
||||||
|
"""Test execute rejects memory exceeding max."""
|
||||||
|
with pytest.raises(ValueError, match="memory.*exceeds maximum"):
|
||||||
|
backend.execute("pass", memory="4g") # max is 2g
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_validates_cpu_quota_against_max(backend, mock_config):
|
||||||
|
"""Test execute rejects CPU quota exceeding max."""
|
||||||
|
with pytest.raises(ValueError, match="cpu_quota.*exceeds maximum"):
|
||||||
|
backend.execute("pass", cpu_quota=150000) # max is 100000
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_logs_to_audit(backend, mock_audit_logger):
|
||||||
|
"""Test execute logs execution to audit log."""
|
||||||
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
||||||
|
mock_executor = Mock()
|
||||||
|
mock_executor_class.return_value = mock_executor
|
||||||
|
mock_executor.execute.return_value = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=42,
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
|
||||||
|
backend.execute("x = 2 + 2")
|
||||||
|
|
||||||
|
# Verify audit log was called
|
||||||
|
mock_audit_logger.log.assert_called()
|
||||||
|
call_args = mock_audit_logger.log.call_args
|
||||||
|
|
||||||
|
# Check that code hash is logged, not actual code
|
||||||
|
log_data = call_args[1]
|
||||||
|
assert 'code_hash' in log_data or 'details' in log_data
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_with_volumes(backend, mock_container_manager):
|
||||||
|
"""Test execute passes volume configuration to container manager."""
|
||||||
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
||||||
|
mock_executor = Mock()
|
||||||
|
mock_executor_class.return_value = mock_executor
|
||||||
|
mock_executor.execute.return_value = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=None,
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
|
||||||
|
volumes = {
|
||||||
|
"/mcp-forge/sessions/test-session/workspace": {"bind": "/workspace", "mode": "rw"}
|
||||||
|
}
|
||||||
|
|
||||||
|
backend.execute("pass", volumes=volumes)
|
||||||
|
|
||||||
|
# Verify volumes were passed through
|
||||||
|
args, kwargs = mock_executor_class.call_args
|
||||||
|
# Volumes should be passed to container_manager through executor
|
||||||
|
# This is verified through the executor initialization
|
||||||
|
assert mock_executor_class.called
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_returns_result(backend):
|
||||||
|
"""Test execute returns ExecutionResult from executor."""
|
||||||
|
expected_result = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="Hello\n",
|
||||||
|
stderr="",
|
||||||
|
result=42,
|
||||||
|
execution_time=0.5,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
||||||
|
mock_executor = Mock()
|
||||||
|
mock_executor_class.return_value = mock_executor
|
||||||
|
mock_executor.execute.return_value = expected_result
|
||||||
|
|
||||||
|
result = backend.execute('print("Hello"); 42')
|
||||||
|
|
||||||
|
assert result == expected_result
|
||||||
|
assert result.success is True
|
||||||
|
assert result.result == 42
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_handles_executor_errors(backend):
|
||||||
|
"""Test execute propagates executor errors."""
|
||||||
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
||||||
|
mock_executor = Mock()
|
||||||
|
mock_executor_class.return_value = mock_executor
|
||||||
|
mock_executor.execute.side_effect = RuntimeError("Container failed")
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="Container failed"):
|
||||||
|
backend.execute("pass")
|
||||||
|
|
||||||
|
|
||||||
|
def test_validate_limits_accepts_valid_limits(backend):
|
||||||
|
"""Test _validate_limits accepts limits within maximums."""
|
||||||
|
# Should not raise
|
||||||
|
backend._validate_limits(
|
||||||
|
timeout=1000,
|
||||||
|
memory="1g",
|
||||||
|
cpu_quota=75000
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_validate_limits_rejects_excessive_timeout(backend):
|
||||||
|
"""Test _validate_limits rejects excessive timeout."""
|
||||||
|
with pytest.raises(ValueError, match="timeout"):
|
||||||
|
backend._validate_limits(
|
||||||
|
timeout=2000,
|
||||||
|
memory="512m",
|
||||||
|
cpu_quota=50000
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_validate_limits_rejects_excessive_memory(backend):
|
||||||
|
"""Test _validate_limits rejects excessive memory."""
|
||||||
|
with pytest.raises(ValueError, match="memory"):
|
||||||
|
backend._validate_limits(
|
||||||
|
timeout=300,
|
||||||
|
memory="4g",
|
||||||
|
cpu_quota=50000
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_validate_limits_rejects_excessive_cpu_quota(backend):
|
||||||
|
"""Test _validate_limits rejects excessive CPU quota."""
|
||||||
|
with pytest.raises(ValueError, match="cpu_quota"):
|
||||||
|
backend._validate_limits(
|
||||||
|
timeout=300,
|
||||||
|
memory="512m",
|
||||||
|
cpu_quota=150000
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_image_returns_custom_when_provided(backend):
|
||||||
|
"""Test _get_image returns custom image when provided."""
|
||||||
|
image = backend._get_image("mcp-forge/custom:latest")
|
||||||
|
assert image == "mcp-forge/custom:latest"
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_image_returns_default_when_none(backend, mock_config):
|
||||||
|
"""Test _get_image returns default image when None provided."""
|
||||||
|
image = backend._get_image(None)
|
||||||
|
assert image == mock_config.images.python_3_11
|
||||||
|
|
||||||
|
|
||||||
|
def test_concurrent_executions_are_independent(backend):
|
||||||
|
"""Test multiple concurrent executions don't interfere."""
|
||||||
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
||||||
|
# Create separate mock executors for each call
|
||||||
|
executor1 = Mock()
|
||||||
|
executor2 = Mock()
|
||||||
|
mock_executor_class.side_effect = [executor1, executor2]
|
||||||
|
|
||||||
|
executor1.execute.return_value = ExecutionResult(
|
||||||
|
success=True, stdout="", stderr="", result=1,
|
||||||
|
execution_time=0.1, exit_code=0
|
||||||
|
)
|
||||||
|
executor2.execute.return_value = ExecutionResult(
|
||||||
|
success=True, stdout="", stderr="", result=2,
|
||||||
|
execution_time=0.1, exit_code=0
|
||||||
|
)
|
||||||
|
|
||||||
|
result1 = backend.execute("1")
|
||||||
|
result2 = backend.execute("2")
|
||||||
|
|
||||||
|
assert result1.result == 1
|
||||||
|
assert result2.result == 2
|
||||||
|
|
||||||
|
# Each execution should create its own executor
|
||||||
|
assert mock_executor_class.call_count == 2
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_with_all_custom_params(backend):
|
||||||
|
"""Test execute with all parameters customized."""
|
||||||
|
with patch('mcp_forge.execution.simple.backend.CodeExecutor') as mock_executor_class:
|
||||||
|
mock_executor = Mock()
|
||||||
|
mock_executor_class.return_value = mock_executor
|
||||||
|
mock_executor.execute.return_value = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="",
|
||||||
|
stderr="",
|
||||||
|
result=None,
|
||||||
|
execution_time=0.1,
|
||||||
|
exit_code=0
|
||||||
|
)
|
||||||
|
|
||||||
|
volumes = {"/mcp-forge/sessions/test/work": {"bind": "/workspace", "mode": "rw"}}
|
||||||
|
|
||||||
|
backend.execute(
|
||||||
|
"pass",
|
||||||
|
timeout=600,
|
||||||
|
memory="1g",
|
||||||
|
cpu_quota=75000,
|
||||||
|
custom_image="mcp-forge/python:3.12",
|
||||||
|
volumes=volumes
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify all parameters were applied
|
||||||
|
args, kwargs = mock_executor_class.call_args
|
||||||
|
|
||||||
|
resource_limits = kwargs.get('resource_limits')
|
||||||
|
assert resource_limits.timeout == 600
|
||||||
|
assert resource_limits.memory_bytes == 1024 * 1024 * 1024 # 1g in bytes
|
||||||
|
assert resource_limits.cpu_quota == 75000
|
||||||
|
|
||||||
|
assert kwargs.get('image') == "mcp-forge/python:3.12"
|
||||||
299
tests/execution/simple/test_executor.py
Normal file
299
tests/execution/simple/test_executor.py
Normal file
|
|
@ -0,0 +1,299 @@
|
||||||
|
"""Tests for the Code Executor module."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from unittest.mock import Mock
|
||||||
|
import json
|
||||||
|
|
||||||
|
from mcp_forge.execution.simple.executor import CodeExecutor, ExecutionResult
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def resource_limits():
|
||||||
|
"""Standard resource limits for testing."""
|
||||||
|
return ResourceLimits(
|
||||||
|
memory="512m",
|
||||||
|
cpu_quota=50000,
|
||||||
|
storage="1g",
|
||||||
|
timeout=30
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_container_manager():
|
||||||
|
"""Mock SecureContainerManager."""
|
||||||
|
manager = Mock(spec=SecureContainerManager)
|
||||||
|
|
||||||
|
# Mock container lifecycle
|
||||||
|
manager.create_container.return_value = "test-container-123"
|
||||||
|
manager.start_container.return_value = None
|
||||||
|
manager.stop_container.return_value = None
|
||||||
|
manager.remove_container.return_value = None
|
||||||
|
manager.wait_for_container.return_value = 0 # exit code
|
||||||
|
manager.get_container_logs.return_value = ("", "") # (stdout, stderr)
|
||||||
|
|
||||||
|
return manager
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def executor(mock_container_manager, resource_limits):
|
||||||
|
"""CodeExecutor instance with mocked dependencies."""
|
||||||
|
return CodeExecutor(
|
||||||
|
container_manager=mock_container_manager,
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
resource_limits=resource_limits
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_simple_python_code_returns_result(executor, mock_container_manager):
|
||||||
|
"""Test executing simple Python code returns the result."""
|
||||||
|
# Mock successful execution with result
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": 42, "error": None}),
|
||||||
|
""
|
||||||
|
)
|
||||||
|
|
||||||
|
result = executor.execute("2 + 2")
|
||||||
|
|
||||||
|
assert result.success is True
|
||||||
|
assert result.result == 42
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert result.error is None
|
||||||
|
|
||||||
|
# Verify container lifecycle
|
||||||
|
mock_container_manager.create_container.assert_called_once()
|
||||||
|
mock_container_manager.start_container.assert_called_once_with("test-container-123")
|
||||||
|
mock_container_manager.wait_for_container.assert_called_once_with("test-container-123", timeout=30)
|
||||||
|
mock_container_manager.remove_container.assert_called_once_with("test-container-123")
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_with_stdout_capture(executor, mock_container_manager):
|
||||||
|
"""Test code execution captures stdout."""
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": None}) + "\n" + "Hello, World!",
|
||||||
|
""
|
||||||
|
)
|
||||||
|
|
||||||
|
result = executor.execute('print("Hello, World!")')
|
||||||
|
|
||||||
|
assert result.success is True
|
||||||
|
assert "Hello, World!" in result.stdout
|
||||||
|
assert result.stderr == ""
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_with_stderr_capture(executor, mock_container_manager):
|
||||||
|
"""Test code execution captures stderr."""
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": None}),
|
||||||
|
"Warning: something happened"
|
||||||
|
)
|
||||||
|
|
||||||
|
result = executor.execute('import sys; print("warning", file=sys.stderr)')
|
||||||
|
|
||||||
|
assert result.success is True
|
||||||
|
assert result.stderr == "Warning: something happened"
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_timeout_enforcement(executor, mock_container_manager):
|
||||||
|
"""Test code execution enforces timeout."""
|
||||||
|
# Simulate timeout by having wait_for_container take too long
|
||||||
|
mock_container_manager.wait_for_container.side_effect = TimeoutError("Container exceeded timeout")
|
||||||
|
|
||||||
|
result = executor.execute("import time; time.sleep(60)", timeout=1)
|
||||||
|
|
||||||
|
assert result.success is False
|
||||||
|
assert result.error is not None
|
||||||
|
assert "timeout" in result.error.lower()
|
||||||
|
|
||||||
|
# Verify cleanup still happens
|
||||||
|
mock_container_manager.remove_container.assert_called_once_with("test-container-123")
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_with_exception_handling(executor, mock_container_manager):
|
||||||
|
"""Test code execution handles exceptions gracefully."""
|
||||||
|
error_msg = "ZeroDivisionError: division by zero"
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": error_msg}),
|
||||||
|
""
|
||||||
|
)
|
||||||
|
mock_container_manager.wait_for_container.return_value = 1 # non-zero exit
|
||||||
|
|
||||||
|
result = executor.execute("1 / 0")
|
||||||
|
|
||||||
|
assert result.success is False
|
||||||
|
assert result.error == error_msg
|
||||||
|
assert result.exit_code == 1
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_with_syntax_error_returns_clear_error(executor, mock_container_manager):
|
||||||
|
"""Test code with syntax error returns clear error message."""
|
||||||
|
error_msg = "SyntaxError: invalid syntax"
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": error_msg}),
|
||||||
|
""
|
||||||
|
)
|
||||||
|
mock_container_manager.wait_for_container.return_value = 1
|
||||||
|
|
||||||
|
result = executor.execute("def foo( :")
|
||||||
|
|
||||||
|
assert result.success is False
|
||||||
|
assert "SyntaxError" in result.error
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_code_with_runtime_error_returns_clear_error(executor, mock_container_manager):
|
||||||
|
"""Test code with runtime error returns clear error with traceback."""
|
||||||
|
error_msg = "NameError: name 'undefined_var' is not defined"
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": error_msg}),
|
||||||
|
""
|
||||||
|
)
|
||||||
|
mock_container_manager.wait_for_container.return_value = 1
|
||||||
|
|
||||||
|
result = executor.execute("print(undefined_var)")
|
||||||
|
|
||||||
|
assert result.success is False
|
||||||
|
assert "NameError" in result.error
|
||||||
|
|
||||||
|
|
||||||
|
def test_result_serialization_json_compatible_types(executor, mock_container_manager):
|
||||||
|
"""Test execution result contains only JSON-serializable data."""
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": [1, 2, {"key": "value"}], "error": None}),
|
||||||
|
""
|
||||||
|
)
|
||||||
|
|
||||||
|
result = executor.execute('[1, 2, {"key": "value"}]')
|
||||||
|
|
||||||
|
# Verify result can be serialized to JSON
|
||||||
|
result_dict = result.to_dict()
|
||||||
|
json_str = json.dumps(result_dict)
|
||||||
|
assert json_str is not None
|
||||||
|
|
||||||
|
# Verify result data
|
||||||
|
assert result.result == [1, 2, {"key": "value"}]
|
||||||
|
|
||||||
|
|
||||||
|
def test_large_output_handling(executor, mock_container_manager):
|
||||||
|
"""Test execution handles large output without issues."""
|
||||||
|
large_output = "x" * 10000 # 10KB of output
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": None}) + "\n" + large_output,
|
||||||
|
""
|
||||||
|
)
|
||||||
|
|
||||||
|
result = executor.execute('print("x" * 10000)')
|
||||||
|
|
||||||
|
assert result.success is True
|
||||||
|
assert len(result.stdout) >= 10000
|
||||||
|
|
||||||
|
|
||||||
|
def test_execution_result_to_dict(resource_limits):
|
||||||
|
"""Test ExecutionResult.to_dict() returns proper dictionary."""
|
||||||
|
result = ExecutionResult(
|
||||||
|
success=True,
|
||||||
|
stdout="output",
|
||||||
|
stderr="",
|
||||||
|
result=42,
|
||||||
|
execution_time=0.5,
|
||||||
|
exit_code=0,
|
||||||
|
error=None
|
||||||
|
)
|
||||||
|
|
||||||
|
result_dict = result.to_dict()
|
||||||
|
|
||||||
|
assert isinstance(result_dict, dict)
|
||||||
|
assert result_dict["success"] is True
|
||||||
|
assert result_dict["stdout"] == "output"
|
||||||
|
assert result_dict["stderr"] == ""
|
||||||
|
assert result_dict["result"] == 42
|
||||||
|
assert result_dict["execution_time"] == 0.5
|
||||||
|
assert result_dict["exit_code"] == 0
|
||||||
|
assert result_dict["error"] is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_prepare_code_wraps_code_properly(executor):
|
||||||
|
"""Test _prepare_code wraps code to capture result."""
|
||||||
|
code = "x = 2 + 2\nx"
|
||||||
|
wrapped = executor._prepare_code(code)
|
||||||
|
|
||||||
|
# Wrapped code should be executable Python
|
||||||
|
assert "import" in wrapped
|
||||||
|
assert "json" in wrapped
|
||||||
|
assert code in wrapped or "2 + 2" in wrapped
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_output_extracts_result_and_error(executor):
|
||||||
|
"""Test _parse_output correctly extracts result and error from JSON."""
|
||||||
|
# Test successful result
|
||||||
|
stdout = json.dumps({"result": 42, "error": None})
|
||||||
|
result, error = executor._parse_output(stdout)
|
||||||
|
assert result == 42
|
||||||
|
assert error is None
|
||||||
|
|
||||||
|
# Test error
|
||||||
|
stdout = json.dumps({"result": None, "error": "ValueError: invalid"})
|
||||||
|
result, error = executor._parse_output(stdout)
|
||||||
|
assert result is None
|
||||||
|
assert error == "ValueError: invalid"
|
||||||
|
|
||||||
|
|
||||||
|
def test_cleanup_happens_even_on_create_failure(executor, mock_container_manager):
|
||||||
|
"""Test container cleanup happens even if create fails."""
|
||||||
|
mock_container_manager.create_container.side_effect = Exception("Create failed")
|
||||||
|
|
||||||
|
with pytest.raises(Exception, match="Create failed"):
|
||||||
|
executor.execute("print('test')")
|
||||||
|
|
||||||
|
# No container to remove since create failed
|
||||||
|
mock_container_manager.remove_container.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
def test_cleanup_happens_even_on_start_failure(executor, mock_container_manager):
|
||||||
|
"""Test container cleanup happens even if start fails."""
|
||||||
|
mock_container_manager.start_container.side_effect = Exception("Start failed")
|
||||||
|
|
||||||
|
with pytest.raises(Exception, match="Start failed"):
|
||||||
|
executor.execute("print('test')")
|
||||||
|
|
||||||
|
# Container should still be removed
|
||||||
|
mock_container_manager.remove_container.assert_called_once_with("test-container-123")
|
||||||
|
|
||||||
|
|
||||||
|
def test_execution_time_tracking(executor, mock_container_manager):
|
||||||
|
"""Test execution time is tracked accurately."""
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": None}),
|
||||||
|
""
|
||||||
|
)
|
||||||
|
|
||||||
|
result = executor.execute("pass")
|
||||||
|
|
||||||
|
assert result.execution_time >= 0
|
||||||
|
assert isinstance(result.execution_time, float)
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_with_custom_timeout(executor, mock_container_manager):
|
||||||
|
"""Test execute respects custom timeout parameter."""
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": None}),
|
||||||
|
""
|
||||||
|
)
|
||||||
|
|
||||||
|
executor.execute("pass", timeout=60)
|
||||||
|
|
||||||
|
# Verify wait was called with custom timeout
|
||||||
|
mock_container_manager.wait_for_container.assert_called_with("test-container-123", timeout=60)
|
||||||
|
|
||||||
|
|
||||||
|
def test_execute_uses_default_timeout_from_resource_limits(executor, mock_container_manager):
|
||||||
|
"""Test execute uses default timeout from resource limits when not specified."""
|
||||||
|
mock_container_manager.get_container_logs.return_value = (
|
||||||
|
json.dumps({"result": None, "error": None}),
|
||||||
|
""
|
||||||
|
)
|
||||||
|
|
||||||
|
executor.execute("pass") # No timeout specified
|
||||||
|
|
||||||
|
# Should use resource_limits.timeout (30)
|
||||||
|
mock_container_manager.wait_for_container.assert_called_with("test-container-123", timeout=30)
|
||||||
59
tests/integration/conftest.py
Normal file
59
tests/integration/conftest.py
Normal file
|
|
@ -0,0 +1,59 @@
|
||||||
|
"""Shared fixtures for integration tests."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from mcp_forge.config.schema import (
|
||||||
|
ForgeConfig, ServerConfig, SecurityConfig, ExecutionConfig, SessionConfig,
|
||||||
|
ImageConfig, VolumeConfig, EnvironmentBuilderConfig, PackageValidationConfig
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def real_config(tmp_path):
|
||||||
|
"""Real configuration with all required fields for integration tests."""
|
||||||
|
# Create required files
|
||||||
|
(tmp_path / "allowlist.txt").write_text("requests\npandas\nnumpy\n")
|
||||||
|
(tmp_path / "blocklist.txt").write_text("")
|
||||||
|
|
||||||
|
config = ForgeConfig(
|
||||||
|
server=ServerConfig(
|
||||||
|
host="localhost",
|
||||||
|
port=3000,
|
||||||
|
podman_socket=Path("/run/user/1000/podman/podman.sock")
|
||||||
|
),
|
||||||
|
security=SecurityConfig(
|
||||||
|
audit_log=tmp_path / "audit.log",
|
||||||
|
enforce_resource_limits=True,
|
||||||
|
allow_network=False
|
||||||
|
),
|
||||||
|
execution=ExecutionConfig(
|
||||||
|
default_backend="simple",
|
||||||
|
default_timeout=300,
|
||||||
|
max_timeout=1800,
|
||||||
|
default_memory="512m",
|
||||||
|
max_memory="2g"
|
||||||
|
),
|
||||||
|
images=ImageConfig(
|
||||||
|
python_3_11="mcp-forge/python:3.11",
|
||||||
|
python_3_12="mcp-forge/python:3.12",
|
||||||
|
auto_pull=False
|
||||||
|
),
|
||||||
|
sessions=SessionConfig(
|
||||||
|
max_concurrent=10,
|
||||||
|
idle_timeout=3600
|
||||||
|
),
|
||||||
|
volumes=VolumeConfig(
|
||||||
|
base_path=tmp_path / "volumes"
|
||||||
|
),
|
||||||
|
environment_builder=EnvironmentBuilderConfig(
|
||||||
|
uv_cache_path=tmp_path / "cache",
|
||||||
|
build_rate_limit={"requests": 5, "period": 60},
|
||||||
|
package_validation=PackageValidationConfig(
|
||||||
|
allowlist_path=tmp_path / "allowlist.txt",
|
||||||
|
blocklist_path=tmp_path / "blocklist.txt"
|
||||||
|
)
|
||||||
|
),
|
||||||
|
mcp_tools={}
|
||||||
|
)
|
||||||
|
return config
|
||||||
56
tests/integration/test_environment_build.py
Normal file
56
tests/integration/test_environment_build.py
Normal file
|
|
@ -0,0 +1,56 @@
|
||||||
|
"""Integration tests for environment building workflow."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_environment_build_end_to_end(tmp_path):
|
||||||
|
"""Test: Package list → validation → UV install → image build → container creation."""
|
||||||
|
# Setup: Create environment builder with all dependencies
|
||||||
|
# Execute:
|
||||||
|
# 1. Request environment with packages: ["requests", "pandas"]
|
||||||
|
# 2. Validate packages (should pass)
|
||||||
|
# 3. Install with UV
|
||||||
|
# 4. Build container image
|
||||||
|
# 5. Create running container
|
||||||
|
# Verify: Container has packages installed and importable
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Environment build flow")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_blocked_package_prevents_build(tmp_path):
|
||||||
|
"""Test: Security validation prevents building environments with blocked packages."""
|
||||||
|
# Setup: Config with blocked packages
|
||||||
|
# Execute: Try to build environment with blocked package
|
||||||
|
# Verify: Build fails at validation stage, no container created
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Security validation integration")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_invalid_package_name_fails_gracefully(tmp_path):
|
||||||
|
"""Test: Invalid package names are caught early."""
|
||||||
|
# Setup: Environment builder
|
||||||
|
# Execute: Request build with non-existent package
|
||||||
|
# Verify: Validation or install fails with clear error message
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Error handling in build flow")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_environment_caching(tmp_path):
|
||||||
|
"""Test: Building same environment twice uses cache."""
|
||||||
|
# Setup: Environment builder with caching enabled
|
||||||
|
# Execute:
|
||||||
|
# 1. Build environment with ["requests"]
|
||||||
|
# 2. Build same environment again
|
||||||
|
# Verify: Second build is faster (uses cached image)
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Build caching")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_concurrent_builds_respect_limits(tmp_path):
|
||||||
|
"""Test: Multiple simultaneous builds respect max_parallel_builds limit."""
|
||||||
|
# Setup: Environment builder with max_parallel_builds=2
|
||||||
|
# Execute: Trigger 5 builds simultaneously
|
||||||
|
# Verify: Only 2 run at once, others wait
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Build rate limiting")
|
||||||
48
tests/integration/test_error_handling.py
Normal file
48
tests/integration/test_error_handling.py
Normal file
|
|
@ -0,0 +1,48 @@
|
||||||
|
"""Integration tests for error handling and recovery."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_container_crash_cleanup(tmp_path):
|
||||||
|
"""Test: If container crashes, resources are cleaned up."""
|
||||||
|
# Setup: Create session with container
|
||||||
|
# Execute: Crash the container (kill -9)
|
||||||
|
# Verify: Session marked as failed, container removed, resources freed
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Crash recovery")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_execution_error_propagation(tmp_path):
|
||||||
|
"""Test: Python errors in execution are returned with full traceback."""
|
||||||
|
# Setup: Create backend
|
||||||
|
# Execute: Code with syntax error or runtime error
|
||||||
|
# Verify: Error returned to caller with traceback, doesn't crash server
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Error propagation")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_podman_connection_loss_handling(tmp_path):
|
||||||
|
"""Test: If Podman socket disconnects, errors are clear."""
|
||||||
|
# Setup: Create system connected to Podman
|
||||||
|
# Execute: Simulate socket disconnect
|
||||||
|
# Verify: Operations fail with clear "Podman unavailable" error
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Connection loss handling")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_concurrent_session_cleanup(tmp_path):
|
||||||
|
"""Test: Cleaning up many sessions simultaneously doesn't deadlock."""
|
||||||
|
# Setup: Create 100 sessions
|
||||||
|
# Execute: Cleanup all simultaneously
|
||||||
|
# Verify: All cleaned up without deadlock or resource leaks
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Concurrent cleanup")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_invalid_mcp_tool_call_error_handling(tmp_path):
|
||||||
|
"""Test: Calling MCP tool with wrong arguments returns clear error."""
|
||||||
|
# Setup: Register MCP tool
|
||||||
|
# Execute: Call with invalid arguments
|
||||||
|
# Verify: Returns validation error, doesn't crash bridge
|
||||||
|
pytest.skip("TODO: Phase 5.3 - MCP error handling")
|
||||||
33
tests/integration/test_execution_flow.py
Normal file
33
tests/integration/test_execution_flow.py
Normal file
|
|
@ -0,0 +1,33 @@
|
||||||
|
"""Integration tests for end-to-end execution flows."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_simple_backend_execution_flow(real_config):
|
||||||
|
"""Test complete flow: code submission → execution → result return."""
|
||||||
|
pytest.skip("TODO: Phase 5.4 - Requires proper container registration and security setup")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_jupyter_session_lifecycle(real_config):
|
||||||
|
"""Test: Create session → execute multiple code blocks → cleanup."""
|
||||||
|
pytest.skip("TODO: Phase 5.4 - Complex Jupyter session testing requires more setup")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_jupyter_session_isolation(real_config):
|
||||||
|
"""Test: Two sessions don't share state."""
|
||||||
|
pytest.skip("TODO: Phase 5.4 - Complex Jupyter session testing requires more setup")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_execution_with_timeout(real_config):
|
||||||
|
"""Test: Long-running code gets killed after timeout."""
|
||||||
|
pytest.skip("TODO: Phase 5.4 - Requires proper container registration and security setup")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_execution_with_memory_limit(real_config):
|
||||||
|
"""Test: Memory-intensive code respects limits."""
|
||||||
|
pytest.skip("TODO: Phase 5.4 - Requires proper container registration and security setup")
|
||||||
66
tests/integration/test_mcp_integration.py
Normal file
66
tests/integration/test_mcp_integration.py
Normal file
|
|
@ -0,0 +1,66 @@
|
||||||
|
"""Integration tests for MCP tool and bridge integration."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_mcp_tool_bridge_connection(tmp_path):
|
||||||
|
"""Test: MCP client → bridge server → tool execution."""
|
||||||
|
# Setup: Start bridge server, create MCP client, register tools
|
||||||
|
# Execute: Client calls tool through bridge
|
||||||
|
# Verify: Tool executes and returns result through bridge
|
||||||
|
pytest.skip("TODO: Phase 5.3 - MCP bridge integration")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_tool_injection_into_execution(tmp_path):
|
||||||
|
"""Test: MCP tools are injected into Python execution environment."""
|
||||||
|
# Setup: Create session with MCP tools available
|
||||||
|
# Execute: Python code that calls MCP tool (e.g., mcp_tools.search_web())
|
||||||
|
# Verify: Tool is callable and returns expected result
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Tool injection integration")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_multiple_mcp_clients_no_collision(tmp_path):
|
||||||
|
"""Test: Multiple MCP clients with same tool names don't conflict."""
|
||||||
|
# Setup: Register two clients, both with "search" tool
|
||||||
|
# Execute: Call search tool
|
||||||
|
# Verify: Collision detected and handled (namespacing or error)
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Tool collision detection")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_mcp_tool_handler_execute_python(tmp_path):
|
||||||
|
"""Test: execute_python MCP tool end-to-end."""
|
||||||
|
# Setup: Create ForgeServer
|
||||||
|
# Execute: Call execute_python tool with simple code
|
||||||
|
# Verify: Code executes and returns stdout/result
|
||||||
|
pytest.skip("TODO: Phase 5.3 - ExecutePythonTool integration")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_mcp_tool_handler_document_state(tmp_path):
|
||||||
|
"""Test: document_state MCP tool shows session variables."""
|
||||||
|
# Setup: Create session, execute code that sets variables
|
||||||
|
# Execute: Call document_state tool
|
||||||
|
# Verify: Returns list of variables and their values
|
||||||
|
pytest.skip("TODO: Phase 5.3 - DocumentStateTool integration")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_mcp_tool_handler_build_environment(tmp_path):
|
||||||
|
"""Test: build_environment MCP tool creates custom environment."""
|
||||||
|
# Setup: Create ForgeServer
|
||||||
|
# Execute: Call build_environment with package list
|
||||||
|
# Verify: Environment built and can be used for execution
|
||||||
|
pytest.skip("TODO: Phase 5.3 - BuildEnvironmentTool integration")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_mcp_resource_handlers(tmp_path):
|
||||||
|
"""Test: MCP resource handlers return correct data."""
|
||||||
|
# Setup: Create ForgeServer with sessions and environments
|
||||||
|
# Execute: Read resources (tools/available, sessions/list, environments/list)
|
||||||
|
# Verify: Resources return expected data
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Resource handler integration")
|
||||||
48
tests/integration/test_security_enforcement.py
Normal file
48
tests/integration/test_security_enforcement.py
Normal file
|
|
@ -0,0 +1,48 @@
|
||||||
|
"""Integration tests for security enforcement across components."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_blocked_package_import_prevented(tmp_path):
|
||||||
|
"""Test: Attempting to import blocked package fails."""
|
||||||
|
# Setup: Config with blocked packages (e.g., ["subprocess", "os"])
|
||||||
|
# Execute: Try to execute code that imports blocked package
|
||||||
|
# Verify: Execution blocked or import fails
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Package blocking enforcement")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_resource_limits_enforced_in_container(tmp_path):
|
||||||
|
"""Test: Container actually respects memory/CPU/timeout limits."""
|
||||||
|
# Setup: Create container with strict limits
|
||||||
|
# Execute: Run code that tries to exceed limits
|
||||||
|
# Verify: Container killed or limited appropriately
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Resource limit enforcement")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_audit_logging_across_operations(tmp_path):
|
||||||
|
"""Test: All operations are logged to audit trail."""
|
||||||
|
# Setup: Create system with audit logging
|
||||||
|
# Execute: Multiple operations (create session, execute code, build env)
|
||||||
|
# Verify: All operations appear in audit log with correct metadata
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Audit trail integration")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_network_isolation_in_containers(tmp_path):
|
||||||
|
"""Test: Containers cannot access external network (if configured)."""
|
||||||
|
# Setup: Create container with network disabled
|
||||||
|
# Execute: Try to make HTTP request
|
||||||
|
# Verify: Request fails (network isolated)
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Network isolation")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_filesystem_isolation_in_containers(tmp_path):
|
||||||
|
"""Test: Containers cannot access host filesystem outside mounts."""
|
||||||
|
# Setup: Create container
|
||||||
|
# Execute: Try to read /etc/passwd or other host files
|
||||||
|
# Verify: Access denied
|
||||||
|
pytest.skip("TODO: Phase 5.3 - Filesystem isolation")
|
||||||
359
tests/mcp/test_bridge.py
Normal file
359
tests/mcp/test_bridge.py
Normal file
|
|
@ -0,0 +1,359 @@
|
||||||
|
"""Tests for Tool Bridge Server."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import socket
|
||||||
|
import json
|
||||||
|
import tempfile
|
||||||
|
from pathlib import Path
|
||||||
|
from unittest.mock import AsyncMock, Mock, patch
|
||||||
|
from mcp_forge.mcp.bridge import ToolBridgeServer
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def temp_socket_path():
|
||||||
|
"""Create temporary socket path."""
|
||||||
|
with tempfile.TemporaryDirectory() as tmpdir:
|
||||||
|
yield Path(tmpdir) / "test_bridge.sock"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_client_manager():
|
||||||
|
"""Create mock MCP client manager."""
|
||||||
|
manager = AsyncMock()
|
||||||
|
manager.call_tool = AsyncMock(return_value={"result": "success"})
|
||||||
|
return manager
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_audit_logger():
|
||||||
|
"""Create mock audit logger."""
|
||||||
|
logger = Mock()
|
||||||
|
logger.log = Mock()
|
||||||
|
return logger
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_bridge_server_starts_and_stops(temp_socket_path, mock_client_manager, mock_audit_logger):
|
||||||
|
"""Test that bridge server starts and stops cleanly."""
|
||||||
|
bridge = ToolBridgeServer(
|
||||||
|
socket_path=temp_socket_path,
|
||||||
|
client_manager=mock_client_manager,
|
||||||
|
audit_logger=mock_audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
# Start server
|
||||||
|
bridge.start()
|
||||||
|
|
||||||
|
# Verify socket was created
|
||||||
|
assert temp_socket_path.exists()
|
||||||
|
|
||||||
|
# Stop server
|
||||||
|
bridge.stop()
|
||||||
|
|
||||||
|
# Verify socket was removed
|
||||||
|
assert not temp_socket_path.exists()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_receive_and_forward_tool_call(temp_socket_path, mock_client_manager, mock_audit_logger):
|
||||||
|
"""Test receiving tool call request and forwarding to client."""
|
||||||
|
bridge = ToolBridgeServer(
|
||||||
|
socket_path=temp_socket_path,
|
||||||
|
client_manager=mock_client_manager,
|
||||||
|
audit_logger=mock_audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
bridge.start()
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Connect as client and send tool call request
|
||||||
|
client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||||
|
client_sock.connect(str(temp_socket_path))
|
||||||
|
|
||||||
|
request = {
|
||||||
|
"tool": "test_tool",
|
||||||
|
"params": {"arg1": "value1", "arg2": 42}
|
||||||
|
}
|
||||||
|
|
||||||
|
client_sock.sendall(json.dumps(request).encode('utf-8'))
|
||||||
|
client_sock.shutdown(socket.SHUT_WR)
|
||||||
|
|
||||||
|
# Receive response
|
||||||
|
response_data = b''
|
||||||
|
while True:
|
||||||
|
chunk = client_sock.recv(4096)
|
||||||
|
if not chunk:
|
||||||
|
break
|
||||||
|
response_data += chunk
|
||||||
|
|
||||||
|
response = json.loads(response_data.decode('utf-8'))
|
||||||
|
|
||||||
|
# Verify response
|
||||||
|
assert response["success"] is True
|
||||||
|
assert response["result"] == {"result": "success"}
|
||||||
|
|
||||||
|
# Verify tool was called with correct arguments
|
||||||
|
mock_client_manager.call_tool.assert_called_once_with(
|
||||||
|
"test_tool",
|
||||||
|
{"arg1": "value1", "arg2": 42}
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify audit log was called
|
||||||
|
mock_audit_logger.log.assert_called()
|
||||||
|
|
||||||
|
client_sock.close()
|
||||||
|
finally:
|
||||||
|
bridge.stop()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_handle_tool_call_error(temp_socket_path, mock_client_manager, mock_audit_logger):
|
||||||
|
"""Test handling tool call errors."""
|
||||||
|
# Make client manager raise error
|
||||||
|
mock_client_manager.call_tool = AsyncMock(side_effect=RuntimeError("Tool failed"))
|
||||||
|
|
||||||
|
bridge = ToolBridgeServer(
|
||||||
|
socket_path=temp_socket_path,
|
||||||
|
client_manager=mock_client_manager,
|
||||||
|
audit_logger=mock_audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
bridge.start()
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Connect and send request
|
||||||
|
client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||||
|
client_sock.connect(str(temp_socket_path))
|
||||||
|
|
||||||
|
request = {"tool": "failing_tool", "params": {}}
|
||||||
|
client_sock.sendall(json.dumps(request).encode('utf-8'))
|
||||||
|
client_sock.shutdown(socket.SHUT_WR)
|
||||||
|
|
||||||
|
# Receive response
|
||||||
|
response_data = b''
|
||||||
|
while True:
|
||||||
|
chunk = client_sock.recv(4096)
|
||||||
|
if not chunk:
|
||||||
|
break
|
||||||
|
response_data += chunk
|
||||||
|
|
||||||
|
response = json.loads(response_data.decode('utf-8'))
|
||||||
|
|
||||||
|
# Verify error response
|
||||||
|
assert response["success"] is False
|
||||||
|
assert "Tool failed" in response["error"]
|
||||||
|
|
||||||
|
client_sock.close()
|
||||||
|
finally:
|
||||||
|
bridge.stop()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_handle_invalid_json(temp_socket_path, mock_client_manager, mock_audit_logger):
|
||||||
|
"""Test handling invalid JSON in request."""
|
||||||
|
bridge = ToolBridgeServer(
|
||||||
|
socket_path=temp_socket_path,
|
||||||
|
client_manager=mock_client_manager,
|
||||||
|
audit_logger=mock_audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
bridge.start()
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Connect and send invalid JSON
|
||||||
|
client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||||
|
client_sock.connect(str(temp_socket_path))
|
||||||
|
|
||||||
|
client_sock.sendall(b"not valid json")
|
||||||
|
client_sock.shutdown(socket.SHUT_WR)
|
||||||
|
|
||||||
|
# Receive response
|
||||||
|
response_data = b''
|
||||||
|
while True:
|
||||||
|
chunk = client_sock.recv(4096)
|
||||||
|
if not chunk:
|
||||||
|
break
|
||||||
|
response_data += chunk
|
||||||
|
|
||||||
|
response = json.loads(response_data.decode('utf-8'))
|
||||||
|
|
||||||
|
# Verify error response
|
||||||
|
assert response["success"] is False
|
||||||
|
assert "Invalid JSON" in response["error"]
|
||||||
|
|
||||||
|
client_sock.close()
|
||||||
|
finally:
|
||||||
|
bridge.stop()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_handle_missing_tool_field(temp_socket_path, mock_client_manager, mock_audit_logger):
|
||||||
|
"""Test handling request missing 'tool' field."""
|
||||||
|
bridge = ToolBridgeServer(
|
||||||
|
socket_path=temp_socket_path,
|
||||||
|
client_manager=mock_client_manager,
|
||||||
|
audit_logger=mock_audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
bridge.start()
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Connect and send request without 'tool' field
|
||||||
|
client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||||
|
client_sock.connect(str(temp_socket_path))
|
||||||
|
|
||||||
|
request = {"params": {"arg1": "value1"}} # Missing 'tool'
|
||||||
|
client_sock.sendall(json.dumps(request).encode('utf-8'))
|
||||||
|
client_sock.shutdown(socket.SHUT_WR)
|
||||||
|
|
||||||
|
# Receive response
|
||||||
|
response_data = b''
|
||||||
|
while True:
|
||||||
|
chunk = client_sock.recv(4096)
|
||||||
|
if not chunk:
|
||||||
|
break
|
||||||
|
response_data += chunk
|
||||||
|
|
||||||
|
response = json.loads(response_data.decode('utf-8'))
|
||||||
|
|
||||||
|
# Verify error response
|
||||||
|
assert response["success"] is False
|
||||||
|
assert "tool" in response["error"].lower()
|
||||||
|
|
||||||
|
client_sock.close()
|
||||||
|
finally:
|
||||||
|
bridge.stop()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_concurrent_requests(temp_socket_path, mock_client_manager, mock_audit_logger):
|
||||||
|
"""Test handling multiple concurrent requests."""
|
||||||
|
import threading
|
||||||
|
|
||||||
|
# Track call counts
|
||||||
|
call_count = {"count": 0}
|
||||||
|
|
||||||
|
async def mock_call_tool(tool_name, arguments):
|
||||||
|
call_count["count"] += 1
|
||||||
|
return {"result": f"success_{call_count['count']}"}
|
||||||
|
|
||||||
|
mock_client_manager.call_tool = mock_call_tool
|
||||||
|
|
||||||
|
bridge = ToolBridgeServer(
|
||||||
|
socket_path=temp_socket_path,
|
||||||
|
client_manager=mock_client_manager,
|
||||||
|
audit_logger=mock_audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
bridge.start()
|
||||||
|
|
||||||
|
try:
|
||||||
|
results = []
|
||||||
|
|
||||||
|
def make_request(tool_name):
|
||||||
|
client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||||
|
client_sock.connect(str(temp_socket_path))
|
||||||
|
|
||||||
|
request = {"tool": tool_name, "params": {}}
|
||||||
|
client_sock.sendall(json.dumps(request).encode('utf-8'))
|
||||||
|
client_sock.shutdown(socket.SHUT_WR)
|
||||||
|
|
||||||
|
response_data = b''
|
||||||
|
while True:
|
||||||
|
chunk = client_sock.recv(4096)
|
||||||
|
if not chunk:
|
||||||
|
break
|
||||||
|
response_data += chunk
|
||||||
|
|
||||||
|
response = json.loads(response_data.decode('utf-8'))
|
||||||
|
results.append(response)
|
||||||
|
client_sock.close()
|
||||||
|
|
||||||
|
# Make 3 concurrent requests
|
||||||
|
threads = []
|
||||||
|
for i in range(3):
|
||||||
|
thread = threading.Thread(target=make_request, args=(f"tool_{i}",))
|
||||||
|
threads.append(thread)
|
||||||
|
thread.start()
|
||||||
|
|
||||||
|
# Wait for all threads
|
||||||
|
for thread in threads:
|
||||||
|
thread.join()
|
||||||
|
|
||||||
|
# Verify all requests succeeded
|
||||||
|
assert len(results) == 3
|
||||||
|
for result in results:
|
||||||
|
assert result["success"] is True
|
||||||
|
|
||||||
|
# Verify all were processed
|
||||||
|
assert call_count["count"] == 3
|
||||||
|
finally:
|
||||||
|
bridge.stop()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_audit_logging_tool_name_only(temp_socket_path, mock_client_manager, mock_audit_logger):
|
||||||
|
"""Test that audit log only logs tool name, not parameters."""
|
||||||
|
bridge = ToolBridgeServer(
|
||||||
|
socket_path=temp_socket_path,
|
||||||
|
client_manager=mock_client_manager,
|
||||||
|
audit_logger=mock_audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
bridge.start()
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Send request with sensitive parameters
|
||||||
|
client_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||||
|
client_sock.connect(str(temp_socket_path))
|
||||||
|
|
||||||
|
request = {
|
||||||
|
"tool": "sensitive_tool",
|
||||||
|
"params": {"password": "secret123", "token": "abc123"}
|
||||||
|
}
|
||||||
|
|
||||||
|
client_sock.sendall(json.dumps(request).encode('utf-8'))
|
||||||
|
client_sock.shutdown(socket.SHUT_WR)
|
||||||
|
|
||||||
|
# Receive response
|
||||||
|
response_data = b''
|
||||||
|
while True:
|
||||||
|
chunk = client_sock.recv(4096)
|
||||||
|
if not chunk:
|
||||||
|
break
|
||||||
|
response_data += chunk
|
||||||
|
|
||||||
|
client_sock.close()
|
||||||
|
|
||||||
|
# Verify audit log was called
|
||||||
|
mock_audit_logger.log.assert_called()
|
||||||
|
|
||||||
|
# Get the log call arguments
|
||||||
|
log_call = mock_audit_logger.log.call_args
|
||||||
|
|
||||||
|
# Verify tool name is in log
|
||||||
|
log_str = str(log_call)
|
||||||
|
assert "sensitive_tool" in log_str
|
||||||
|
|
||||||
|
# Verify sensitive parameters are NOT in log
|
||||||
|
assert "secret123" not in log_str
|
||||||
|
assert "abc123" not in log_str
|
||||||
|
finally:
|
||||||
|
bridge.stop()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_socket_cleanup_on_error(temp_socket_path, mock_client_manager, mock_audit_logger):
|
||||||
|
"""Test that socket is cleaned up even if server encounters error."""
|
||||||
|
bridge = ToolBridgeServer(
|
||||||
|
socket_path=temp_socket_path,
|
||||||
|
client_manager=mock_client_manager,
|
||||||
|
audit_logger=mock_audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
bridge.start()
|
||||||
|
assert temp_socket_path.exists()
|
||||||
|
|
||||||
|
# Stop should cleanup
|
||||||
|
bridge.stop()
|
||||||
|
assert not temp_socket_path.exists()
|
||||||
249
tests/mcp/test_client.py
Normal file
249
tests/mcp/test_client.py
Normal file
|
|
@ -0,0 +1,249 @@
|
||||||
|
"""Tests for MCP Client Wrapper."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from unittest.mock import Mock, AsyncMock, patch
|
||||||
|
from mcp_forge.mcp.client import MCPClientWrapper
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_fastmcp_client():
|
||||||
|
"""Mock fastmcp Client."""
|
||||||
|
client = AsyncMock()
|
||||||
|
|
||||||
|
# Create tool mocks with actual string names (not Mock.name)
|
||||||
|
tool1 = Mock()
|
||||||
|
tool1.name = "tool1"
|
||||||
|
tool1.description = "Tool 1"
|
||||||
|
tool1.inputSchema = {"type": "object", "properties": {}}
|
||||||
|
|
||||||
|
tool2 = Mock()
|
||||||
|
tool2.name = "tool2"
|
||||||
|
tool2.description = "Tool 2"
|
||||||
|
tool2.inputSchema = {"type": "object", "properties": {}}
|
||||||
|
|
||||||
|
# Mock list_tools to return tool objects
|
||||||
|
client.list_tools = AsyncMock(return_value=Mock(tools=[tool1, tool2]))
|
||||||
|
|
||||||
|
# Mock call_tool to return result with data
|
||||||
|
client.call_tool = AsyncMock(return_value=Mock(
|
||||||
|
data="result",
|
||||||
|
content=[Mock(text="result")],
|
||||||
|
is_error=False
|
||||||
|
))
|
||||||
|
|
||||||
|
# Mock context manager
|
||||||
|
client.__aenter__ = AsyncMock(return_value=client)
|
||||||
|
client.__aexit__ = AsyncMock(return_value=None)
|
||||||
|
|
||||||
|
return client
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_connect_success(mock_fastmcp_client):
|
||||||
|
"""Test successful connection to MCP server."""
|
||||||
|
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test-server",
|
||||||
|
command="python",
|
||||||
|
args=["-m", "test_server"]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert not client.is_connected()
|
||||||
|
|
||||||
|
await client.connect()
|
||||||
|
|
||||||
|
assert client.is_connected()
|
||||||
|
mock_fastmcp_client.__aenter__.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_connect_with_env(mock_fastmcp_client):
|
||||||
|
"""Test connection with environment variables."""
|
||||||
|
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||||
|
env = {"API_KEY": "test123", "DEBUG": "true"}
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test-server",
|
||||||
|
command="python",
|
||||||
|
args=["-m", "test_server"],
|
||||||
|
env=env
|
||||||
|
)
|
||||||
|
|
||||||
|
await client.connect()
|
||||||
|
|
||||||
|
assert client.is_connected()
|
||||||
|
assert client.env == env
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_disconnect(mock_fastmcp_client):
|
||||||
|
"""Test disconnect from MCP server."""
|
||||||
|
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test-server",
|
||||||
|
command="python",
|
||||||
|
args=["-m", "test_server"]
|
||||||
|
)
|
||||||
|
|
||||||
|
await client.connect()
|
||||||
|
assert client.is_connected()
|
||||||
|
|
||||||
|
await client.disconnect()
|
||||||
|
|
||||||
|
assert not client.is_connected()
|
||||||
|
mock_fastmcp_client.__aexit__.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_list_tools(mock_fastmcp_client):
|
||||||
|
"""Test listing available tools."""
|
||||||
|
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test-server",
|
||||||
|
command="python",
|
||||||
|
args=["-m", "test_server"]
|
||||||
|
)
|
||||||
|
|
||||||
|
await client.connect()
|
||||||
|
tools = await client.list_tools()
|
||||||
|
|
||||||
|
assert tools == ["tool1", "tool2"]
|
||||||
|
mock_fastmcp_client.list_tools.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_list_tools_not_connected():
|
||||||
|
"""Test list_tools raises error when not connected."""
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test-server",
|
||||||
|
command="python",
|
||||||
|
args=["-m", "test_server"]
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="not connected"):
|
||||||
|
await client.list_tools()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_tool_schema(mock_fastmcp_client):
|
||||||
|
"""Test getting tool schema."""
|
||||||
|
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test-server",
|
||||||
|
command="python",
|
||||||
|
args=["-m", "test_server"]
|
||||||
|
)
|
||||||
|
|
||||||
|
await client.connect()
|
||||||
|
schema = await client.get_tool_schema("tool1")
|
||||||
|
|
||||||
|
assert schema == {"type": "object", "properties": {}}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_tool_schema_not_found(mock_fastmcp_client):
|
||||||
|
"""Test get_tool_schema raises error for unknown tool."""
|
||||||
|
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test-server",
|
||||||
|
command="python",
|
||||||
|
args=["-m", "test_server"]
|
||||||
|
)
|
||||||
|
|
||||||
|
await client.connect()
|
||||||
|
|
||||||
|
with pytest.raises(KeyError, match="Tool 'unknown' not found"):
|
||||||
|
await client.get_tool_schema("unknown")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_call_tool_success(mock_fastmcp_client):
|
||||||
|
"""Test successful tool call."""
|
||||||
|
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test-server",
|
||||||
|
command="python",
|
||||||
|
args=["-m", "test_server"]
|
||||||
|
)
|
||||||
|
|
||||||
|
await client.connect()
|
||||||
|
result = await client.call_tool("tool1", {"param": "value"})
|
||||||
|
|
||||||
|
assert result == "result"
|
||||||
|
mock_fastmcp_client.call_tool.assert_called_once_with("tool1", {"param": "value"})
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_call_tool_not_connected():
|
||||||
|
"""Test call_tool raises error when not connected."""
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test-server",
|
||||||
|
command="python",
|
||||||
|
args=["-m", "test_server"]
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="not connected"):
|
||||||
|
await client.call_tool("tool1", {})
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_call_tool_failure(mock_fastmcp_client):
|
||||||
|
"""Test tool call failure handling."""
|
||||||
|
mock_fastmcp_client.call_tool = AsyncMock(side_effect=Exception("Tool failed"))
|
||||||
|
|
||||||
|
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test-server",
|
||||||
|
command="python",
|
||||||
|
args=["-m", "test_server"]
|
||||||
|
)
|
||||||
|
|
||||||
|
await client.connect()
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="Tool call failed.*Tool failed"):
|
||||||
|
await client.call_tool("tool1", {})
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_connection_failure():
|
||||||
|
"""Test connection failure handling."""
|
||||||
|
failing_client = AsyncMock()
|
||||||
|
failing_client.__aenter__ = AsyncMock(side_effect=Exception("Connection failed"))
|
||||||
|
|
||||||
|
with patch('mcp_forge.mcp.client.Client', return_value=failing_client):
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test-server",
|
||||||
|
command="python",
|
||||||
|
args=["-m", "test_server"]
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="Failed to connect.*Connection failed"):
|
||||||
|
await client.connect()
|
||||||
|
|
||||||
|
assert not client.is_connected()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_reconnect(mock_fastmcp_client):
|
||||||
|
"""Test reconnection after disconnect."""
|
||||||
|
with patch('mcp_forge.mcp.client.Client', return_value=mock_fastmcp_client):
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test-server",
|
||||||
|
command="python",
|
||||||
|
args=["-m", "test_server"]
|
||||||
|
)
|
||||||
|
|
||||||
|
# First connection
|
||||||
|
await client.connect()
|
||||||
|
assert client.is_connected()
|
||||||
|
|
||||||
|
# Disconnect
|
||||||
|
await client.disconnect()
|
||||||
|
assert not client.is_connected()
|
||||||
|
|
||||||
|
# Reconnect
|
||||||
|
await client.connect()
|
||||||
|
assert client.is_connected()
|
||||||
|
|
||||||
|
# Should be able to use tools
|
||||||
|
tools = await client.list_tools()
|
||||||
|
assert tools == ["tool1", "tool2"]
|
||||||
114
tests/mcp/test_http_transport.py
Normal file
114
tests/mcp/test_http_transport.py
Normal file
|
|
@ -0,0 +1,114 @@
|
||||||
|
"""Tests for HTTP and SSE transport support in MCP client."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from mcp_forge.mcp.client import MCPClientWrapper
|
||||||
|
|
||||||
|
|
||||||
|
def test_stdio_transport_creation():
|
||||||
|
"""Test creating a client with stdio transport."""
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test_stdio",
|
||||||
|
transport_type="stdio",
|
||||||
|
command="python",
|
||||||
|
args=["-m", "server"],
|
||||||
|
env={"KEY": "value"}
|
||||||
|
)
|
||||||
|
|
||||||
|
assert client.name == "test_stdio"
|
||||||
|
assert client.transport_type == "stdio"
|
||||||
|
assert client.command == "python"
|
||||||
|
assert client.args == ["-m", "server"]
|
||||||
|
assert client.env == {"KEY": "value"}
|
||||||
|
assert client.transport is not None
|
||||||
|
|
||||||
|
|
||||||
|
def test_http_transport_creation():
|
||||||
|
"""Test creating a client with HTTP transport."""
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test_http",
|
||||||
|
transport_type="http",
|
||||||
|
url="http://localhost:8006/mcp",
|
||||||
|
headers={"Authorization": "Bearer token123"}
|
||||||
|
)
|
||||||
|
|
||||||
|
assert client.name == "test_http"
|
||||||
|
assert client.transport_type == "http"
|
||||||
|
assert client.url == "http://localhost:8006/mcp"
|
||||||
|
assert client.headers == {"Authorization": "Bearer token123"}
|
||||||
|
assert client.transport is not None
|
||||||
|
|
||||||
|
|
||||||
|
def test_sse_transport_creation():
|
||||||
|
"""Test creating a client with SSE transport."""
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test_sse",
|
||||||
|
transport_type="sse",
|
||||||
|
url="http://localhost:9000/events",
|
||||||
|
headers={"X-Custom": "value"}
|
||||||
|
)
|
||||||
|
|
||||||
|
assert client.name == "test_sse"
|
||||||
|
assert client.transport_type == "sse"
|
||||||
|
assert client.url == "http://localhost:9000/events"
|
||||||
|
assert client.headers == {"X-Custom": "value"}
|
||||||
|
assert client.transport is not None
|
||||||
|
|
||||||
|
|
||||||
|
def test_stdio_without_command_raises_error():
|
||||||
|
"""Test that stdio transport requires a command."""
|
||||||
|
with pytest.raises(ValueError, match="command required for stdio transport"):
|
||||||
|
MCPClientWrapper(
|
||||||
|
name="test_stdio",
|
||||||
|
transport_type="stdio"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_http_without_url_raises_error():
|
||||||
|
"""Test that HTTP transport requires a URL."""
|
||||||
|
with pytest.raises(ValueError, match="url required for http transport"):
|
||||||
|
MCPClientWrapper(
|
||||||
|
name="test_http",
|
||||||
|
transport_type="http"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_sse_without_url_raises_error():
|
||||||
|
"""Test that SSE transport requires a URL."""
|
||||||
|
with pytest.raises(ValueError, match="url required for sse transport"):
|
||||||
|
MCPClientWrapper(
|
||||||
|
name="test_sse",
|
||||||
|
transport_type="sse"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_invalid_transport_type_raises_error():
|
||||||
|
"""Test that an invalid transport type raises an error."""
|
||||||
|
with pytest.raises(ValueError, match="Unknown transport type"):
|
||||||
|
MCPClientWrapper(
|
||||||
|
name="test_invalid",
|
||||||
|
transport_type="invalid"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_http_transport_with_empty_headers():
|
||||||
|
"""Test HTTP transport with empty headers dict."""
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test_http",
|
||||||
|
transport_type="http",
|
||||||
|
url="http://localhost:8006/mcp"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert client.headers == {}
|
||||||
|
assert client.transport is not None
|
||||||
|
|
||||||
|
|
||||||
|
def test_default_stdio_transport():
|
||||||
|
"""Test that stdio is the default transport type."""
|
||||||
|
client = MCPClientWrapper(
|
||||||
|
name="test_default",
|
||||||
|
command="python",
|
||||||
|
args=["-m", "server"]
|
||||||
|
)
|
||||||
|
|
||||||
|
assert client.transport_type == "stdio"
|
||||||
|
assert client.command == "python"
|
||||||
227
tests/mcp/test_injection.py
Normal file
227
tests/mcp/test_injection.py
Normal file
|
|
@ -0,0 +1,227 @@
|
||||||
|
"""Tests for Tool Injection Generator."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import ast
|
||||||
|
from unittest.mock import AsyncMock
|
||||||
|
from mcp_forge.mcp.injection import ToolInjectionGenerator
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_client_manager():
|
||||||
|
"""Create mock MCP client manager with tools."""
|
||||||
|
manager = AsyncMock()
|
||||||
|
|
||||||
|
# Tool schemas
|
||||||
|
manager.get_tool_schema = AsyncMock(side_effect=lambda tool_name: {
|
||||||
|
"read_file": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": {"type": "string", "description": "File path to read"},
|
||||||
|
"encoding": {"type": "string", "description": "File encoding"}
|
||||||
|
},
|
||||||
|
"required": ["path"]
|
||||||
|
},
|
||||||
|
"write_file": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"path": {"type": "string", "description": "File path to write"},
|
||||||
|
"content": {"type": "string", "description": "Content to write"},
|
||||||
|
"mode": {"type": "string", "description": "Write mode"}
|
||||||
|
},
|
||||||
|
"required": ["path", "content"]
|
||||||
|
},
|
||||||
|
"calculate": {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"expression": {"type": "string", "description": "Math expression"},
|
||||||
|
"precision": {"type": "integer", "description": "Decimal precision"}
|
||||||
|
},
|
||||||
|
"required": ["expression"]
|
||||||
|
}
|
||||||
|
}[tool_name])
|
||||||
|
|
||||||
|
return manager
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_generate_injection_code_is_valid_python(mock_client_manager):
|
||||||
|
"""Test that generated code is valid Python."""
|
||||||
|
generator = ToolInjectionGenerator(mock_client_manager)
|
||||||
|
|
||||||
|
code = await generator.generate_injection_code(
|
||||||
|
tool_names=["read_file", "write_file"],
|
||||||
|
bridge_socket_path="/tmp/bridge.sock"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Try to parse the generated code
|
||||||
|
try:
|
||||||
|
ast.parse(code)
|
||||||
|
except SyntaxError as e:
|
||||||
|
pytest.fail(f"Generated code has syntax error: {e}")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_generated_code_includes_bridge_client(mock_client_manager):
|
||||||
|
"""Test that generated code includes bridge client."""
|
||||||
|
generator = ToolInjectionGenerator(mock_client_manager)
|
||||||
|
|
||||||
|
code = await generator.generate_injection_code(
|
||||||
|
tool_names=["read_file"],
|
||||||
|
bridge_socket_path="/tmp/bridge.sock"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify bridge client function is included
|
||||||
|
assert "_mcp_call" in code
|
||||||
|
assert "socket.socket" in code
|
||||||
|
assert "socket.AF_UNIX" in code
|
||||||
|
assert "/tmp/bridge.sock" in code
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_generated_code_includes_tool_functions(mock_client_manager):
|
||||||
|
"""Test that generated code includes wrapper functions for each tool."""
|
||||||
|
generator = ToolInjectionGenerator(mock_client_manager)
|
||||||
|
|
||||||
|
code = await generator.generate_injection_code(
|
||||||
|
tool_names=["read_file", "write_file", "calculate"],
|
||||||
|
bridge_socket_path="/tmp/bridge.sock"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify tool functions are defined
|
||||||
|
assert "def read_file(" in code
|
||||||
|
assert "def write_file(" in code
|
||||||
|
assert "def calculate(" in code
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_function_signatures_match_schemas(mock_client_manager):
|
||||||
|
"""Test that function signatures match tool schemas."""
|
||||||
|
generator = ToolInjectionGenerator(mock_client_manager)
|
||||||
|
|
||||||
|
code = await generator.generate_injection_code(
|
||||||
|
tool_names=["read_file", "write_file"],
|
||||||
|
bridge_socket_path="/tmp/bridge.sock"
|
||||||
|
)
|
||||||
|
|
||||||
|
# read_file has path (required) and encoding (optional)
|
||||||
|
assert "def read_file(path: str, encoding: str = None)" in code
|
||||||
|
|
||||||
|
# write_file has path, content (required) and mode (optional)
|
||||||
|
assert "def write_file(path: str, content: str, mode: str = None)" in code
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_generated_functions_have_docstrings(mock_client_manager):
|
||||||
|
"""Test that generated functions have docstrings."""
|
||||||
|
generator = ToolInjectionGenerator(mock_client_manager)
|
||||||
|
|
||||||
|
code = await generator.generate_injection_code(
|
||||||
|
tool_names=["read_file"],
|
||||||
|
bridge_socket_path="/tmp/bridge.sock"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify docstring is present (contains parameter descriptions)
|
||||||
|
assert '"""' in code
|
||||||
|
assert "File path to read" in code or "path:" in code.lower()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_generated_functions_call_bridge(mock_client_manager):
|
||||||
|
"""Test that generated functions call _mcp_call."""
|
||||||
|
generator = ToolInjectionGenerator(mock_client_manager)
|
||||||
|
|
||||||
|
code = await generator.generate_injection_code(
|
||||||
|
tool_names=["read_file"],
|
||||||
|
bridge_socket_path="/tmp/bridge.sock"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify function calls _mcp_call with tool name
|
||||||
|
lines = code.split('\n')
|
||||||
|
in_read_file = False
|
||||||
|
found_call = False
|
||||||
|
|
||||||
|
for line in lines:
|
||||||
|
if "def read_file(" in line:
|
||||||
|
in_read_file = True
|
||||||
|
if in_read_file and "_mcp_call" in line and "read_file" in line:
|
||||||
|
found_call = True
|
||||||
|
break
|
||||||
|
|
||||||
|
assert found_call, "Generated function should call _mcp_call with tool name"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_type_hints_from_schema(mock_client_manager):
|
||||||
|
"""Test that type hints are generated from schema types."""
|
||||||
|
generator = ToolInjectionGenerator(mock_client_manager)
|
||||||
|
|
||||||
|
code = await generator.generate_injection_code(
|
||||||
|
tool_names=["calculate"],
|
||||||
|
bridge_socket_path="/tmp/bridge.sock"
|
||||||
|
)
|
||||||
|
|
||||||
|
# calculate has string expression and integer precision
|
||||||
|
assert "expression: str" in code
|
||||||
|
assert "precision: int" in code or "precision: " in code
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_required_vs_optional_parameters(mock_client_manager):
|
||||||
|
"""Test that required and optional parameters are handled correctly."""
|
||||||
|
generator = ToolInjectionGenerator(mock_client_manager)
|
||||||
|
|
||||||
|
code = await generator.generate_injection_code(
|
||||||
|
tool_names=["read_file"],
|
||||||
|
bridge_socket_path="/tmp/bridge.sock"
|
||||||
|
)
|
||||||
|
|
||||||
|
# path is required (no default), encoding is optional (has default)
|
||||||
|
assert "def read_file(path: str, encoding: str = None)" in code or \
|
||||||
|
"def read_file(path: str, encoding: " in code
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_generated_code_has_imports(mock_client_manager):
|
||||||
|
"""Test that generated code includes necessary imports."""
|
||||||
|
generator = ToolInjectionGenerator(mock_client_manager)
|
||||||
|
|
||||||
|
code = await generator.generate_injection_code(
|
||||||
|
tool_names=["read_file"],
|
||||||
|
bridge_socket_path="/tmp/bridge.sock"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify imports
|
||||||
|
assert "import socket" in code
|
||||||
|
assert "import json" in code
|
||||||
|
assert "from typing import Any" in code or "typing" in code
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_empty_tool_list(mock_client_manager):
|
||||||
|
"""Test handling of empty tool list."""
|
||||||
|
generator = ToolInjectionGenerator(mock_client_manager)
|
||||||
|
|
||||||
|
code = await generator.generate_injection_code(
|
||||||
|
tool_names=[],
|
||||||
|
bridge_socket_path="/tmp/bridge.sock"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Should still include bridge client
|
||||||
|
assert "_mcp_call" in code
|
||||||
|
# But no tool functions
|
||||||
|
assert "def read_file(" not in code
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_custom_socket_path(mock_client_manager):
|
||||||
|
"""Test that custom socket path is used correctly."""
|
||||||
|
generator = ToolInjectionGenerator(mock_client_manager)
|
||||||
|
|
||||||
|
custom_path = "/custom/path/to/socket.sock"
|
||||||
|
code = await generator.generate_injection_code(
|
||||||
|
tool_names=["read_file"],
|
||||||
|
bridge_socket_path=custom_path
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify custom path is in generated code
|
||||||
|
assert custom_path in code
|
||||||
245
tests/mcp/test_manager.py
Normal file
245
tests/mcp/test_manager.py
Normal file
|
|
@ -0,0 +1,245 @@
|
||||||
|
"""Tests for MCP Client Manager."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from unittest.mock import AsyncMock, patch
|
||||||
|
from mcp_forge.mcp.manager import MCPClientManager
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_client1():
|
||||||
|
"""Create mock MCP client 1."""
|
||||||
|
client = AsyncMock()
|
||||||
|
client.name = "client1"
|
||||||
|
client.is_connected.return_value = False
|
||||||
|
client.connect = AsyncMock()
|
||||||
|
client.disconnect = AsyncMock()
|
||||||
|
client.list_tools = AsyncMock(return_value=["tool1", "tool2"])
|
||||||
|
client.get_tool_schema = AsyncMock(return_value={
|
||||||
|
"type": "object",
|
||||||
|
"properties": {"param1": {"type": "string"}}
|
||||||
|
})
|
||||||
|
client.call_tool = AsyncMock(return_value={"result": "success"})
|
||||||
|
return client
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_client2():
|
||||||
|
"""Create mock MCP client 2."""
|
||||||
|
client = AsyncMock()
|
||||||
|
client.name = "client2"
|
||||||
|
client.is_connected.return_value = False
|
||||||
|
client.connect = AsyncMock()
|
||||||
|
client.disconnect = AsyncMock()
|
||||||
|
client.list_tools = AsyncMock(return_value=["tool3", "tool4"])
|
||||||
|
client.get_tool_schema = AsyncMock(return_value={
|
||||||
|
"type": "object",
|
||||||
|
"properties": {"param2": {"type": "number"}}
|
||||||
|
})
|
||||||
|
client.call_tool = AsyncMock(return_value={"result": "success2"})
|
||||||
|
return client
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def client_config():
|
||||||
|
"""Create test client configuration."""
|
||||||
|
return {
|
||||||
|
"client1": {
|
||||||
|
"command": "python",
|
||||||
|
"args": ["server1.py"],
|
||||||
|
"env": {"KEY1": "value1"}
|
||||||
|
},
|
||||||
|
"client2": {
|
||||||
|
"command": "python",
|
||||||
|
"args": ["server2.py"],
|
||||||
|
"env": {"KEY2": "value2"}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_initialize_clients_from_config(mock_client1, mock_client2, client_config):
|
||||||
|
"""Test initializing clients from configuration."""
|
||||||
|
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
||||||
|
# Setup mock to return different clients for different configs
|
||||||
|
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
||||||
|
|
||||||
|
manager = MCPClientManager(client_config)
|
||||||
|
await manager.initialize()
|
||||||
|
|
||||||
|
# Verify clients were created with correct config
|
||||||
|
assert mock_wrapper.call_count == 2
|
||||||
|
|
||||||
|
# Verify first client created with correct params
|
||||||
|
call1 = mock_wrapper.call_args_list[0]
|
||||||
|
assert call1[1]["name"] == "client1"
|
||||||
|
assert call1[1]["command"] == "python"
|
||||||
|
assert call1[1]["args"] == ["server1.py"]
|
||||||
|
assert call1[1]["env"] == {"KEY1": "value1"}
|
||||||
|
|
||||||
|
# Verify second client created with correct params
|
||||||
|
call2 = mock_wrapper.call_args_list[1]
|
||||||
|
assert call2[1]["name"] == "client2"
|
||||||
|
assert call2[1]["command"] == "python"
|
||||||
|
assert call2[1]["args"] == ["server2.py"]
|
||||||
|
assert call2[1]["env"] == {"KEY2": "value2"}
|
||||||
|
|
||||||
|
# Verify clients were connected
|
||||||
|
mock_client1.connect.assert_called_once()
|
||||||
|
mock_client2.connect.assert_called_once()
|
||||||
|
|
||||||
|
# Verify tools were listed
|
||||||
|
mock_client1.list_tools.assert_called_once()
|
||||||
|
mock_client2.list_tools.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_client_for_tool(mock_client1, mock_client2, client_config):
|
||||||
|
"""Test getting client that provides a specific tool."""
|
||||||
|
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
||||||
|
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
||||||
|
|
||||||
|
manager = MCPClientManager(client_config)
|
||||||
|
await manager.initialize()
|
||||||
|
|
||||||
|
# Get client for tool1 (from client1)
|
||||||
|
client = await manager.get_client_for_tool("tool1")
|
||||||
|
assert client == mock_client1
|
||||||
|
|
||||||
|
# Get client for tool3 (from client2)
|
||||||
|
client = await manager.get_client_for_tool("tool3")
|
||||||
|
assert client == mock_client2
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_client_for_unknown_tool(mock_client1, mock_client2, client_config):
|
||||||
|
"""Test getting client for tool that doesn't exist."""
|
||||||
|
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
||||||
|
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
||||||
|
|
||||||
|
manager = MCPClientManager(client_config)
|
||||||
|
await manager.initialize()
|
||||||
|
|
||||||
|
# Try to get client for non-existent tool
|
||||||
|
with pytest.raises(KeyError, match="Tool 'unknown_tool' not found"):
|
||||||
|
await manager.get_client_for_tool("unknown_tool")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_list_all_tools(mock_client1, mock_client2, client_config):
|
||||||
|
"""Test listing all tools across all clients."""
|
||||||
|
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
||||||
|
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
||||||
|
|
||||||
|
manager = MCPClientManager(client_config)
|
||||||
|
await manager.initialize()
|
||||||
|
|
||||||
|
tools = await manager.list_all_tools()
|
||||||
|
|
||||||
|
# Should have all tools from both clients
|
||||||
|
assert set(tools) == {"tool1", "tool2", "tool3", "tool4"}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_detect_tool_name_collision():
|
||||||
|
"""Test detection of tool name collisions across clients."""
|
||||||
|
# Create clients with overlapping tool names
|
||||||
|
client1 = AsyncMock()
|
||||||
|
client1.name = "client1"
|
||||||
|
client1.connect = AsyncMock()
|
||||||
|
client1.list_tools = AsyncMock(return_value=["tool1", "tool2"])
|
||||||
|
|
||||||
|
client2 = AsyncMock()
|
||||||
|
client2.name = "client2"
|
||||||
|
client2.connect = AsyncMock()
|
||||||
|
client2.list_tools = AsyncMock(return_value=["tool2", "tool3"]) # tool2 collision!
|
||||||
|
|
||||||
|
config = {
|
||||||
|
"client1": {"command": "python", "args": ["server1.py"]},
|
||||||
|
"client2": {"command": "python", "args": ["server2.py"]}
|
||||||
|
}
|
||||||
|
|
||||||
|
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
||||||
|
mock_wrapper.side_effect = [client1, client2]
|
||||||
|
|
||||||
|
manager = MCPClientManager(config)
|
||||||
|
|
||||||
|
# Should raise ValueError about collision during initialization
|
||||||
|
with pytest.raises(ValueError, match="Tool name collision.*tool2.*client1.*client2"):
|
||||||
|
await manager.initialize()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_tool_schema(mock_client1, mock_client2, client_config):
|
||||||
|
"""Test getting tool schema via manager."""
|
||||||
|
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
||||||
|
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
||||||
|
|
||||||
|
manager = MCPClientManager(client_config)
|
||||||
|
await manager.initialize()
|
||||||
|
|
||||||
|
# Get schema for tool1 (from client1)
|
||||||
|
schema = await manager.get_tool_schema("tool1")
|
||||||
|
assert schema == {"type": "object", "properties": {"param1": {"type": "string"}}}
|
||||||
|
mock_client1.get_tool_schema.assert_called_once_with("tool1")
|
||||||
|
|
||||||
|
# Get schema for tool3 (from client2)
|
||||||
|
schema = await manager.get_tool_schema("tool3")
|
||||||
|
assert schema == {"type": "object", "properties": {"param2": {"type": "number"}}}
|
||||||
|
mock_client2.get_tool_schema.assert_called_once_with("tool3")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_call_tool(mock_client1, mock_client2, client_config):
|
||||||
|
"""Test calling tool via manager."""
|
||||||
|
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
||||||
|
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
||||||
|
|
||||||
|
manager = MCPClientManager(client_config)
|
||||||
|
await manager.initialize()
|
||||||
|
|
||||||
|
# Call tool1 (from client1)
|
||||||
|
result = await manager.call_tool("tool1", {"param1": "value"})
|
||||||
|
assert result == {"result": "success"}
|
||||||
|
mock_client1.call_tool.assert_called_once_with("tool1", {"param1": "value"})
|
||||||
|
|
||||||
|
# Call tool3 (from client2)
|
||||||
|
result = await manager.call_tool("tool3", {"param2": 42})
|
||||||
|
assert result == {"result": "success2"}
|
||||||
|
mock_client2.call_tool.assert_called_once_with("tool3", {"param2": 42})
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_shutdown_all_clients(mock_client1, mock_client2, client_config):
|
||||||
|
"""Test shutting down all clients."""
|
||||||
|
with patch('mcp_forge.mcp.manager.MCPClientWrapper') as mock_wrapper:
|
||||||
|
mock_wrapper.side_effect = [mock_client1, mock_client2]
|
||||||
|
|
||||||
|
manager = MCPClientManager(client_config)
|
||||||
|
await manager.initialize()
|
||||||
|
|
||||||
|
# Shutdown all clients
|
||||||
|
await manager.shutdown()
|
||||||
|
|
||||||
|
# Verify both clients were disconnected
|
||||||
|
mock_client1.disconnect.assert_called_once()
|
||||||
|
mock_client2.disconnect.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_manager_before_initialization():
|
||||||
|
"""Test that manager methods fail before initialization."""
|
||||||
|
config = {"client1": {"command": "python", "args": ["server.py"]}}
|
||||||
|
manager = MCPClientManager(config)
|
||||||
|
|
||||||
|
# Should raise RuntimeError if not initialized
|
||||||
|
with pytest.raises(RuntimeError, match="Manager not initialized"):
|
||||||
|
await manager.list_all_tools()
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="Manager not initialized"):
|
||||||
|
await manager.get_client_for_tool("tool1")
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="Manager not initialized"):
|
||||||
|
await manager.get_tool_schema("tool1")
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="Manager not initialized"):
|
||||||
|
await manager.call_tool("tool1", {})
|
||||||
433
tests/podman/test_containers.py
Normal file
433
tests/podman/test_containers.py
Normal file
|
|
@ -0,0 +1,433 @@
|
||||||
|
"""
|
||||||
|
Tests for Secure Container Manager.
|
||||||
|
|
||||||
|
Following TDD approach - these tests are written before implementation.
|
||||||
|
Tests cover all requirements from todo.md section 1.3.2.
|
||||||
|
All tests use mocked Podman client (no actual containers needed).
|
||||||
|
"""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from datetime import datetime, timedelta
|
||||||
|
from unittest.mock import MagicMock, Mock, patch, call
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_container_with_valid_params_succeeds():
|
||||||
|
"""Test that container creation with valid params succeeds."""
|
||||||
|
from mcp_forge.podman.containers import ContainerConfig, SecureContainerManager
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
# Setup mocks
|
||||||
|
mock_podman_client = MagicMock(spec=PodmanClient)
|
||||||
|
mock_container = MagicMock()
|
||||||
|
mock_container.id = "abc123"
|
||||||
|
mock_podman_client.client.containers.create.return_value = mock_container
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(Path("/tmp/test_audit.log"))
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=Path("/tmp/test_audit.log")))
|
||||||
|
|
||||||
|
manager = SecureContainerManager(
|
||||||
|
podman_client=mock_podman_client,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
config = ContainerConfig(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
command=["python", "-c", "print('hello')"],
|
||||||
|
resource_limits=ResourceLimits(memory="512m", storage="1g", cpu_quota=100000)
|
||||||
|
)
|
||||||
|
|
||||||
|
container_id = manager.create_container(config, session_id="test-session-1")
|
||||||
|
|
||||||
|
assert container_id == "abc123"
|
||||||
|
mock_podman_client.client.containers.create.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_container_with_forbidden_params_raises_security_error():
|
||||||
|
"""Test that forbidden parameters raise SecurityError."""
|
||||||
|
from mcp_forge.podman.containers import ContainerConfig, SecureContainerManager
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
mock_podman_client = MagicMock(spec=PodmanClient)
|
||||||
|
audit_logger = AuditLogger(Path("/tmp/test_audit.log"))
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=Path("/tmp/test_audit.log")))
|
||||||
|
|
||||||
|
manager = SecureContainerManager(
|
||||||
|
podman_client=mock_podman_client,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
config = ContainerConfig(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
command=["python", "-c", "print('hello')"]
|
||||||
|
)
|
||||||
|
|
||||||
|
# Try to override security params (should be caught in to_podman_params or validation)
|
||||||
|
with pytest.raises(SecurityError):
|
||||||
|
# This should fail validation
|
||||||
|
manager.create_container(config, session_id="test-session", privileged=True)
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_container_with_invalid_image_raises_security_error():
|
||||||
|
"""Test that invalid/disallowed images raise SecurityError."""
|
||||||
|
from mcp_forge.podman.containers import ContainerConfig, SecureContainerManager
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
mock_podman_client = MagicMock(spec=PodmanClient)
|
||||||
|
audit_logger = AuditLogger(Path("/tmp/test_audit.log"))
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=Path("/tmp/test_audit.log")))
|
||||||
|
|
||||||
|
manager = SecureContainerManager(
|
||||||
|
podman_client=mock_podman_client,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
config = ContainerConfig(
|
||||||
|
image="evil/malicious:latest",
|
||||||
|
command=["python", "-c", "print('hello')"]
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(SecurityError) as exc_info:
|
||||||
|
manager.create_container(config, session_id="test-session")
|
||||||
|
assert "image" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_container_enforces_required_parameters():
|
||||||
|
"""Test that required security parameters are enforced."""
|
||||||
|
from mcp_forge.podman.containers import ContainerConfig
|
||||||
|
|
||||||
|
config = ContainerConfig(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
command=["python", "-c", "print('hello')"]
|
||||||
|
)
|
||||||
|
|
||||||
|
params = config.to_podman_params()
|
||||||
|
|
||||||
|
# Check required security parameters
|
||||||
|
assert params["network_mode"] == "none"
|
||||||
|
assert params["read_only"] is True
|
||||||
|
assert "no-new-privileges" in params["security_opt"]
|
||||||
|
assert params["user"] == "1000:1000"
|
||||||
|
|
||||||
|
|
||||||
|
def test_resource_limits_are_applied_correctly():
|
||||||
|
"""Test that resource limits are correctly applied."""
|
||||||
|
from mcp_forge.podman.containers import ContainerConfig
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
limits = ResourceLimits(
|
||||||
|
memory="1g",
|
||||||
|
storage="2g",
|
||||||
|
cpu_quota=200000
|
||||||
|
)
|
||||||
|
|
||||||
|
config = ContainerConfig(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
resource_limits=limits
|
||||||
|
)
|
||||||
|
|
||||||
|
params = config.to_podman_params()
|
||||||
|
|
||||||
|
assert params["mem_limit"] == "1073741824" # 1GB in bytes as string
|
||||||
|
assert params["cpu_quota"] == 200000
|
||||||
|
assert params["storage_opt"]["size"] == "2147483648" # 2GB in bytes as string
|
||||||
|
|
||||||
|
|
||||||
|
def test_volume_mounts_are_validated():
|
||||||
|
"""Test that volume mounts are validated against allowlist."""
|
||||||
|
from mcp_forge.podman.containers import ContainerConfig, SecureContainerManager
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
mock_podman_client = MagicMock(spec=PodmanClient)
|
||||||
|
audit_logger = AuditLogger(Path("/tmp/test_audit.log"))
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=Path("/tmp/test_audit.log")))
|
||||||
|
|
||||||
|
manager = SecureContainerManager(
|
||||||
|
podman_client=mock_podman_client,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
# Valid mount (session path)
|
||||||
|
config = ContainerConfig(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
volumes={"/mcp-forge/sessions/test-session-1/workspace": {"bind": "/workspace", "mode": "rw"}}
|
||||||
|
)
|
||||||
|
|
||||||
|
# This should succeed (valid session path)
|
||||||
|
mock_container = MagicMock()
|
||||||
|
mock_container.id = "abc123"
|
||||||
|
mock_podman_client.client.containers.create.return_value = mock_container
|
||||||
|
container_id = manager.create_container(config, session_id="test-session-1")
|
||||||
|
assert container_id == "abc123"
|
||||||
|
|
||||||
|
# Invalid mount (forbidden path)
|
||||||
|
config_bad = ContainerConfig(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
volumes={"/etc/passwd": {"bind": "/tmp/passwd", "mode": "r"}}
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(SecurityError):
|
||||||
|
manager.create_container(config_bad, session_id="test-session-1")
|
||||||
|
|
||||||
|
|
||||||
|
def test_start_container_on_session_container_succeeds():
|
||||||
|
"""Test that starting a session container succeeds."""
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
mock_podman_client = MagicMock(spec=PodmanClient)
|
||||||
|
mock_container = MagicMock()
|
||||||
|
mock_podman_client.client.containers.get.return_value = mock_container
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(Path("/tmp/test_audit.log"))
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=Path("/tmp/test_audit.log")))
|
||||||
|
|
||||||
|
# Register container with validator
|
||||||
|
validator.register_session_container("abc123")
|
||||||
|
|
||||||
|
manager = SecureContainerManager(
|
||||||
|
podman_client=mock_podman_client,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
manager.start_container("abc123")
|
||||||
|
mock_container.start.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
def test_start_container_on_non_session_container_raises_security_error():
|
||||||
|
"""Test that starting a non-session container raises SecurityError."""
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
mock_podman_client = MagicMock(spec=PodmanClient)
|
||||||
|
audit_logger = AuditLogger(Path("/tmp/test_audit.log"))
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=Path("/tmp/test_audit.log")))
|
||||||
|
|
||||||
|
manager = SecureContainerManager(
|
||||||
|
podman_client=mock_podman_client,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
# Try to start container not registered with validator
|
||||||
|
with pytest.raises(SecurityError):
|
||||||
|
manager.start_container("unknown123")
|
||||||
|
|
||||||
|
|
||||||
|
def test_stop_container_works():
|
||||||
|
"""Test that stopping a container works."""
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
mock_podman_client = MagicMock(spec=PodmanClient)
|
||||||
|
mock_container = MagicMock()
|
||||||
|
mock_podman_client.client.containers.get.return_value = mock_container
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(Path("/tmp/test_audit.log"))
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=Path("/tmp/test_audit.log")))
|
||||||
|
|
||||||
|
# Register container
|
||||||
|
validator.register_session_container("abc123")
|
||||||
|
|
||||||
|
manager = SecureContainerManager(
|
||||||
|
podman_client=mock_podman_client,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
manager.stop_container("abc123", timeout=10)
|
||||||
|
mock_container.stop.assert_called_once_with(timeout=10)
|
||||||
|
|
||||||
|
|
||||||
|
def test_remove_container_works():
|
||||||
|
"""Test that removing a container works."""
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
mock_podman_client = MagicMock(spec=PodmanClient)
|
||||||
|
mock_container = MagicMock()
|
||||||
|
mock_podman_client.client.containers.get.return_value = mock_container
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(Path("/tmp/test_audit.log"))
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=Path("/tmp/test_audit.log")))
|
||||||
|
|
||||||
|
# Register container
|
||||||
|
validator.register_session_container("abc123")
|
||||||
|
|
||||||
|
manager = SecureContainerManager(
|
||||||
|
podman_client=mock_podman_client,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
manager.remove_container("abc123", force=True)
|
||||||
|
mock_container.remove.assert_called_once_with(force=True)
|
||||||
|
|
||||||
|
|
||||||
|
def test_cleanup_old_containers():
|
||||||
|
"""Test cleanup of old containers."""
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
mock_podman_client = MagicMock(spec=PodmanClient)
|
||||||
|
|
||||||
|
# Create mock old and new containers
|
||||||
|
old_container = MagicMock()
|
||||||
|
old_container.id = "old123"
|
||||||
|
old_container.attrs = {
|
||||||
|
"Created": (datetime.now() - timedelta(hours=25)).isoformat(),
|
||||||
|
"Labels": {"mcp-forge.session": "old-session"}
|
||||||
|
}
|
||||||
|
|
||||||
|
new_container = MagicMock()
|
||||||
|
new_container.id = "new123"
|
||||||
|
new_container.attrs = {
|
||||||
|
"Created": datetime.now().isoformat(),
|
||||||
|
"Labels": {"mcp-forge.session": "new-session"}
|
||||||
|
}
|
||||||
|
|
||||||
|
mock_podman_client.client.containers.list.return_value = [old_container, new_container]
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(Path("/tmp/test_audit.log"))
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=Path("/tmp/test_audit.log")))
|
||||||
|
|
||||||
|
manager = SecureContainerManager(
|
||||||
|
podman_client=mock_podman_client,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
count = manager.cleanup_old_containers(max_age=timedelta(hours=24))
|
||||||
|
|
||||||
|
assert count == 1
|
||||||
|
old_container.remove.assert_called_once_with(force=True)
|
||||||
|
new_container.remove.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_container_logs():
|
||||||
|
"""Test getting container logs."""
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
mock_podman_client = MagicMock(spec=PodmanClient)
|
||||||
|
mock_container = MagicMock()
|
||||||
|
mock_container.logs.return_value = b"stdout output\nstderr output"
|
||||||
|
mock_podman_client.client.containers.get.return_value = mock_container
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(Path("/tmp/test_audit.log"))
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=Path("/tmp/test_audit.log")))
|
||||||
|
|
||||||
|
# Register container
|
||||||
|
validator.register_session_container("abc123")
|
||||||
|
|
||||||
|
manager = SecureContainerManager(
|
||||||
|
podman_client=mock_podman_client,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
stdout, stderr = manager.get_container_logs("abc123", tail=100)
|
||||||
|
|
||||||
|
assert "output" in stdout or "output" in stderr
|
||||||
|
mock_container.logs.assert_called()
|
||||||
|
|
||||||
|
|
||||||
|
def test_wait_for_container():
|
||||||
|
"""Test waiting for container to exit."""
|
||||||
|
from mcp_forge.podman.containers import SecureContainerManager
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
mock_podman_client = MagicMock(spec=PodmanClient)
|
||||||
|
mock_container = MagicMock()
|
||||||
|
mock_container.wait.return_value = {"StatusCode": 0}
|
||||||
|
mock_podman_client.client.containers.get.return_value = mock_container
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(Path("/tmp/test_audit.log"))
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=Path("/tmp/test_audit.log")))
|
||||||
|
|
||||||
|
# Register container
|
||||||
|
validator.register_session_container("abc123")
|
||||||
|
|
||||||
|
manager = SecureContainerManager(
|
||||||
|
podman_client=mock_podman_client,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
exit_code = manager.wait_for_container("abc123", timeout=300)
|
||||||
|
|
||||||
|
assert exit_code == 0
|
||||||
|
mock_container.wait.assert_called_once_with(timeout=300)
|
||||||
|
|
||||||
|
|
||||||
|
def test_container_config_to_podman_params_includes_all_security_settings():
|
||||||
|
"""Test that ContainerConfig.to_podman_params includes all required settings."""
|
||||||
|
from mcp_forge.podman.containers import ContainerConfig
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
config = ContainerConfig(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
command=["python", "-m", "test"],
|
||||||
|
environment={"VAR1": "value1"},
|
||||||
|
volumes={"/mcp-forge/sessions/sess-1/work": {"bind": "/workspace", "mode": "rw"}},
|
||||||
|
resource_limits=ResourceLimits(memory="512m", storage="1g", cpu_quota=100000),
|
||||||
|
working_dir="/workspace",
|
||||||
|
user="1000:1000"
|
||||||
|
)
|
||||||
|
|
||||||
|
params = config.to_podman_params()
|
||||||
|
|
||||||
|
# Required security settings
|
||||||
|
assert params["network_mode"] == "none"
|
||||||
|
assert params["read_only"] is True
|
||||||
|
assert "no-new-privileges" in params["security_opt"]
|
||||||
|
assert params["user"] == "1000:1000"
|
||||||
|
|
||||||
|
# Configuration passthrough
|
||||||
|
assert params["image"] == "mcp-forge/python:3.11"
|
||||||
|
assert params["command"] == ["python", "-m", "test"]
|
||||||
|
assert params["environment"] == {"VAR1": "value1"}
|
||||||
|
assert params["working_dir"] == "/workspace"
|
||||||
|
|
||||||
|
# Resource limits
|
||||||
|
assert "mem_limit" in params
|
||||||
|
assert "cpu_quota" in params
|
||||||
324
tests/podman/test_podman_client.py
Normal file
324
tests/podman/test_podman_client.py
Normal file
|
|
@ -0,0 +1,324 @@
|
||||||
|
"""
|
||||||
|
Tests for Podman client wrapper.
|
||||||
|
|
||||||
|
Following TDD approach - these tests are written before implementation.
|
||||||
|
Tests cover all requirements from todo.md section 1.3.1.
|
||||||
|
All tests use mocked Podman client (no actual Podman needed).
|
||||||
|
"""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from pathlib import Path
|
||||||
|
from unittest.mock import Mock, MagicMock, patch
|
||||||
|
|
||||||
|
|
||||||
|
def test_connection_to_podman_socket_succeeds(tmp_path):
|
||||||
|
"""Test that connection to Podman socket succeeds."""
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
socket_path = tmp_path / "podman.sock"
|
||||||
|
socket_path.touch() # Create fake socket file
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(tmp_path / "audit.log")
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=tmp_path / "audit.log"))
|
||||||
|
|
||||||
|
client = PodmanClient(
|
||||||
|
socket_path=socket_path,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch('mcp_forge.podman.client.BasePodmanClient') as mock_podman:
|
||||||
|
mock_client_instance = MagicMock()
|
||||||
|
mock_client_instance.ping.return_value = "OK"
|
||||||
|
mock_podman.return_value = mock_client_instance
|
||||||
|
|
||||||
|
client.connect()
|
||||||
|
assert client._client is not None
|
||||||
|
mock_podman.assert_called_once_with(base_url=f"unix://{socket_path}")
|
||||||
|
|
||||||
|
|
||||||
|
def test_connection_failure_raises_clear_error(tmp_path):
|
||||||
|
"""Test that connection failure raises clear error."""
|
||||||
|
from mcp_forge.podman.client import PodmanClient, PodmanConnectionError
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
socket_path = tmp_path / "nonexistent.sock"
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(tmp_path / "audit.log")
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=tmp_path / "audit.log"))
|
||||||
|
|
||||||
|
client = PodmanClient(
|
||||||
|
socket_path=socket_path,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch('mcp_forge.podman.client.BasePodmanClient', side_effect=Exception("Connection failed")):
|
||||||
|
with pytest.raises(PodmanConnectionError) as exc_info:
|
||||||
|
client.connect()
|
||||||
|
assert "connection" in str(exc_info.value).lower() or "failed" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_socket_path_validation(tmp_path):
|
||||||
|
"""Test that socket path is validated before connecting."""
|
||||||
|
from mcp_forge.podman.client import PodmanClient, PodmanConnectionError
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
nonexistent_socket = tmp_path / "nonexistent.sock"
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(tmp_path / "audit.log")
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=tmp_path / "audit.log"))
|
||||||
|
|
||||||
|
client = PodmanClient(
|
||||||
|
socket_path=nonexistent_socket,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(PodmanConnectionError) as exc_info:
|
||||||
|
client.verify_socket_access()
|
||||||
|
assert "not found" in str(exc_info.value).lower() or "does not exist" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_socket_permissions_check(tmp_path):
|
||||||
|
"""Test that socket permissions are checked."""
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
socket_path = tmp_path / "podman.sock"
|
||||||
|
socket_path.touch()
|
||||||
|
socket_path.chmod(0o000) # Remove all permissions
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(tmp_path / "audit.log")
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=tmp_path / "audit.log"))
|
||||||
|
|
||||||
|
client = PodmanClient(
|
||||||
|
socket_path=socket_path,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
# verify_socket_access should check readability
|
||||||
|
# Depending on implementation, might raise error or just warn
|
||||||
|
try:
|
||||||
|
client.verify_socket_access()
|
||||||
|
except Exception as e:
|
||||||
|
# Should mention permissions or access
|
||||||
|
assert "permission" in str(e).lower() or "access" in str(e).lower() or "readable" in str(e).lower()
|
||||||
|
finally:
|
||||||
|
socket_path.chmod(0o644) # Restore for cleanup
|
||||||
|
|
||||||
|
|
||||||
|
def test_api_version_compatibility_check(tmp_path):
|
||||||
|
"""Test that API version is checked."""
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
socket_path = tmp_path / "podman.sock"
|
||||||
|
socket_path.touch()
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(tmp_path / "audit.log")
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=tmp_path / "audit.log"))
|
||||||
|
|
||||||
|
client = PodmanClient(
|
||||||
|
socket_path=socket_path,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch('mcp_forge.podman.client.BasePodmanClient') as mock_podman:
|
||||||
|
mock_client = MagicMock()
|
||||||
|
mock_client.version.return_value = {
|
||||||
|
"Version": "4.5.0",
|
||||||
|
"ApiVersion": "4.5.0"
|
||||||
|
}
|
||||||
|
mock_client.ping.return_value = "OK"
|
||||||
|
mock_podman.return_value = mock_client
|
||||||
|
|
||||||
|
client.connect()
|
||||||
|
version_info = client.check_api_version()
|
||||||
|
|
||||||
|
assert "Version" in version_info or "ApiVersion" in version_info
|
||||||
|
|
||||||
|
|
||||||
|
def test_ping_health_check(tmp_path):
|
||||||
|
"""Test that ping/health check works."""
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
socket_path = tmp_path / "podman.sock"
|
||||||
|
socket_path.touch()
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(tmp_path / "audit.log")
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=tmp_path / "audit.log"))
|
||||||
|
|
||||||
|
client = PodmanClient(
|
||||||
|
socket_path=socket_path,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch('mcp_forge.podman.client.BasePodmanClient') as mock_podman:
|
||||||
|
mock_client = MagicMock()
|
||||||
|
mock_client.ping.return_value = "OK"
|
||||||
|
mock_podman.return_value = mock_client
|
||||||
|
|
||||||
|
client.connect()
|
||||||
|
result = client.ping()
|
||||||
|
|
||||||
|
assert result is True or result == "OK"
|
||||||
|
|
||||||
|
|
||||||
|
def test_lazy_connection(tmp_path):
|
||||||
|
"""Test that connection is lazy (only connects when needed)."""
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
socket_path = tmp_path / "podman.sock"
|
||||||
|
socket_path.touch()
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(tmp_path / "audit.log")
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=tmp_path / "audit.log"))
|
||||||
|
|
||||||
|
# Creating client should not connect
|
||||||
|
client = PodmanClient(
|
||||||
|
socket_path=socket_path,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
assert client._client is None # Not connected yet
|
||||||
|
|
||||||
|
# Accessing client property should trigger connection
|
||||||
|
with patch('mcp_forge.podman.client.BasePodmanClient') as mock_podman:
|
||||||
|
mock_client_instance = MagicMock()
|
||||||
|
mock_client_instance.ping.return_value = "OK"
|
||||||
|
mock_podman.return_value = mock_client_instance
|
||||||
|
_ = client.client
|
||||||
|
assert client._client is not None
|
||||||
|
|
||||||
|
|
||||||
|
def test_disconnect_cleanup(tmp_path):
|
||||||
|
"""Test that disconnect cleans up properly."""
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
socket_path = tmp_path / "podman.sock"
|
||||||
|
socket_path.touch()
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(tmp_path / "audit.log")
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=tmp_path / "audit.log"))
|
||||||
|
|
||||||
|
client = PodmanClient(
|
||||||
|
socket_path=socket_path,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch('mcp_forge.podman.client.BasePodmanClient') as mock_podman:
|
||||||
|
mock_client = MagicMock()
|
||||||
|
mock_client.ping.return_value = "OK"
|
||||||
|
mock_podman.return_value = mock_client
|
||||||
|
|
||||||
|
client.connect()
|
||||||
|
assert client._client is not None
|
||||||
|
|
||||||
|
client.disconnect()
|
||||||
|
assert client._client is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_connection_error_includes_socket_path(tmp_path):
|
||||||
|
"""Test that connection errors include the socket path for debugging."""
|
||||||
|
from mcp_forge.podman.client import PodmanClient, PodmanConnectionError
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
socket_path = tmp_path / "test.sock"
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(tmp_path / "audit.log")
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=tmp_path / "audit.log"))
|
||||||
|
|
||||||
|
client = PodmanClient(
|
||||||
|
socket_path=socket_path,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
with pytest.raises(PodmanConnectionError) as exc_info:
|
||||||
|
client.verify_socket_access()
|
||||||
|
|
||||||
|
assert str(socket_path) in str(exc_info.value) or socket_path.name in str(exc_info.value)
|
||||||
|
|
||||||
|
|
||||||
|
def test_client_property_auto_connects(tmp_path):
|
||||||
|
"""Test that accessing client property auto-connects if not connected."""
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
socket_path = tmp_path / "podman.sock"
|
||||||
|
socket_path.touch()
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(tmp_path / "audit.log")
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=tmp_path / "audit.log"))
|
||||||
|
|
||||||
|
client = PodmanClient(
|
||||||
|
socket_path=socket_path,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch('mcp_forge.podman.client.BasePodmanClient') as mock_podman:
|
||||||
|
mock_client = MagicMock()
|
||||||
|
mock_client.ping.return_value = "OK"
|
||||||
|
mock_podman.return_value = mock_client
|
||||||
|
|
||||||
|
# First access should trigger connect
|
||||||
|
_ = client.client
|
||||||
|
assert mock_podman.called
|
||||||
|
|
||||||
|
# Second access should reuse connection
|
||||||
|
mock_podman.reset_mock()
|
||||||
|
_ = client.client
|
||||||
|
assert not mock_podman.called # Should not connect again
|
||||||
|
|
||||||
|
|
||||||
|
def test_validator_and_audit_logger_stored(tmp_path):
|
||||||
|
"""Test that validator and audit logger are stored for later use."""
|
||||||
|
from mcp_forge.podman.client import PodmanClient
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
|
||||||
|
socket_path = tmp_path / "podman.sock"
|
||||||
|
|
||||||
|
audit_logger = AuditLogger(tmp_path / "audit.log")
|
||||||
|
validator = OperationValidator(SecurityConfig(audit_log=tmp_path / "audit.log"))
|
||||||
|
|
||||||
|
client = PodmanClient(
|
||||||
|
socket_path=socket_path,
|
||||||
|
validator=validator,
|
||||||
|
audit_logger=audit_logger
|
||||||
|
)
|
||||||
|
|
||||||
|
assert client.validator is validator
|
||||||
|
assert client.audit_logger is audit_logger
|
||||||
486
tests/security/test_allowlist.py
Normal file
486
tests/security/test_allowlist.py
Normal file
|
|
@ -0,0 +1,486 @@
|
||||||
|
"""
|
||||||
|
Tests for Podman operation allowlist and validation.
|
||||||
|
|
||||||
|
Following TDD approach - these tests are written before implementation.
|
||||||
|
Tests cover all security validation requirements from todo.md section 1.2.2.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
def test_allowed_operation_with_valid_params_passes():
|
||||||
|
"""Test that allowed operation with valid parameters passes validation."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(
|
||||||
|
audit_log=Path("/var/log/audit.log"),
|
||||||
|
enforce_resource_limits=True,
|
||||||
|
allow_network=False
|
||||||
|
)
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
# Should not raise
|
||||||
|
validator.validate_container_create(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
params={
|
||||||
|
"network_mode": "none",
|
||||||
|
"read_only": True,
|
||||||
|
"security_opt": ["no-new-privileges"],
|
||||||
|
"user": "1000:1000",
|
||||||
|
"memory": "536870912",
|
||||||
|
"cpu_quota": 50000
|
||||||
|
},
|
||||||
|
session_id="test-session-123"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_allowed_operation_with_forbidden_params_raises_security_error():
|
||||||
|
"""Test that allowed operation with forbidden parameters raises SecurityError."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(
|
||||||
|
audit_log=Path("/var/log/audit.log"),
|
||||||
|
enforce_resource_limits=True,
|
||||||
|
allow_network=False
|
||||||
|
)
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
# Try to add privileged mode
|
||||||
|
with pytest.raises(SecurityError) as exc_info:
|
||||||
|
validator.validate_container_create(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
params={
|
||||||
|
"privileged": True, # FORBIDDEN
|
||||||
|
"network_mode": "none",
|
||||||
|
"read_only": True,
|
||||||
|
"security_opt": ["no-new-privileges"],
|
||||||
|
"user": "1000:1000"
|
||||||
|
},
|
||||||
|
session_id="test-session-123"
|
||||||
|
)
|
||||||
|
assert "privileged" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_forbidden_param_cap_add_raises_security_error():
|
||||||
|
"""Test that cap_add parameter is rejected."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
with pytest.raises(SecurityError) as exc_info:
|
||||||
|
validator.validate_container_create(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
params={
|
||||||
|
"cap_add": ["SYS_ADMIN"], # FORBIDDEN
|
||||||
|
"network_mode": "none",
|
||||||
|
"read_only": True
|
||||||
|
},
|
||||||
|
session_id="test-session"
|
||||||
|
)
|
||||||
|
assert "cap_add" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_forbidden_param_devices_raises_security_error():
|
||||||
|
"""Test that devices parameter is rejected."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
with pytest.raises(SecurityError) as exc_info:
|
||||||
|
validator.validate_container_create(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
params={
|
||||||
|
"devices": ["/dev/sda"], # FORBIDDEN
|
||||||
|
"network_mode": "none",
|
||||||
|
"read_only": True
|
||||||
|
},
|
||||||
|
session_id="test-session"
|
||||||
|
)
|
||||||
|
assert "devices" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_required_parameters_validation():
|
||||||
|
"""Test that required parameters are enforced."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
# Missing network_mode
|
||||||
|
with pytest.raises(SecurityError) as exc_info:
|
||||||
|
validator.validate_container_create(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
params={
|
||||||
|
"read_only": True,
|
||||||
|
"security_opt": ["no-new-privileges"],
|
||||||
|
"user": "1000:1000"
|
||||||
|
},
|
||||||
|
session_id="test-session"
|
||||||
|
)
|
||||||
|
assert "network_mode" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
# Missing read_only
|
||||||
|
with pytest.raises(SecurityError) as exc_info:
|
||||||
|
validator.validate_container_create(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
params={
|
||||||
|
"network_mode": "none",
|
||||||
|
"security_opt": ["no-new-privileges"],
|
||||||
|
"user": "1000:1000"
|
||||||
|
},
|
||||||
|
session_id="test-session"
|
||||||
|
)
|
||||||
|
assert "read_only" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_image_allowlist_enforcement():
|
||||||
|
"""Test that only allowed images can be used."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
# Valid image
|
||||||
|
validator.validate_image_name("mcp-forge/python:3.11")
|
||||||
|
validator.validate_image_name("mcp-forge/python:3.12")
|
||||||
|
validator.validate_image_name("mcp-forge/jupyter:latest")
|
||||||
|
validator.validate_image_name("mcp-forge/custom:my-env")
|
||||||
|
|
||||||
|
# Invalid image
|
||||||
|
with pytest.raises(SecurityError) as exc_info:
|
||||||
|
validator.validate_image_name("ubuntu:latest")
|
||||||
|
assert "image" in str(exc_info.value).lower() or "allowlist" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
with pytest.raises(SecurityError):
|
||||||
|
validator.validate_image_name("malicious/image:latest")
|
||||||
|
|
||||||
|
|
||||||
|
def test_volume_mount_path_validation():
|
||||||
|
"""Test that volume mount paths are validated against allowlist."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
# Valid session path
|
||||||
|
validator.validate_volume_mount(
|
||||||
|
"/mcp-forge/sessions/test-session-123/workdir",
|
||||||
|
"test-session-123"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Valid shared readonly path
|
||||||
|
validator.validate_volume_mount(
|
||||||
|
"/mcp-forge/shared/readonly/data",
|
||||||
|
"test-session-123"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Valid upload path
|
||||||
|
validator.validate_volume_mount(
|
||||||
|
"/mcp-forge/uploads/test-session-123/file.txt",
|
||||||
|
"test-session-123"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Invalid: root path
|
||||||
|
with pytest.raises(SecurityError) as exc_info:
|
||||||
|
validator.validate_volume_mount("/", "test-session")
|
||||||
|
assert "forbidden" in str(exc_info.value).lower() or "root" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
# Invalid: /etc
|
||||||
|
with pytest.raises(SecurityError):
|
||||||
|
validator.validate_volume_mount("/etc/passwd", "test-session")
|
||||||
|
|
||||||
|
# Invalid: docker socket
|
||||||
|
with pytest.raises(SecurityError):
|
||||||
|
validator.validate_volume_mount("/var/run/docker.sock", "test-session")
|
||||||
|
|
||||||
|
# Invalid: podman socket
|
||||||
|
with pytest.raises(SecurityError):
|
||||||
|
validator.validate_volume_mount("/var/run/podman/podman.sock", "test-session")
|
||||||
|
|
||||||
|
|
||||||
|
def test_capability_restrictions():
|
||||||
|
"""Test that capability restrictions are enforced."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
# cap_add is forbidden
|
||||||
|
with pytest.raises(SecurityError):
|
||||||
|
validator.validate_container_create(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
params={
|
||||||
|
"cap_add": ["NET_ADMIN"],
|
||||||
|
"network_mode": "none",
|
||||||
|
"read_only": True
|
||||||
|
},
|
||||||
|
session_id="test-session"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_network_mode_enforcement():
|
||||||
|
"""Test that network mode is enforced as 'none'."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
# Wrong network mode
|
||||||
|
with pytest.raises(SecurityError) as exc_info:
|
||||||
|
validator.validate_container_create(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
params={
|
||||||
|
"network_mode": "bridge", # Must be "none"
|
||||||
|
"read_only": True,
|
||||||
|
"security_opt": ["no-new-privileges"],
|
||||||
|
"user": "1000:1000"
|
||||||
|
},
|
||||||
|
session_id="test-session"
|
||||||
|
)
|
||||||
|
assert "network_mode" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_privileged_mode_always_rejected():
|
||||||
|
"""Test that privileged mode is always rejected regardless of other params."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
with pytest.raises(SecurityError) as exc_info:
|
||||||
|
validator.validate_container_create(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
params={
|
||||||
|
"privileged": True,
|
||||||
|
"network_mode": "none",
|
||||||
|
"read_only": True,
|
||||||
|
"security_opt": ["no-new-privileges"],
|
||||||
|
"user": "1000:1000"
|
||||||
|
},
|
||||||
|
session_id="test-session"
|
||||||
|
)
|
||||||
|
assert "privileged" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_session_container_tracking():
|
||||||
|
"""Test that session containers are tracked and validated."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
# Register a session container
|
||||||
|
validator.register_session_container("container-123")
|
||||||
|
|
||||||
|
# Should be able to operate on registered container
|
||||||
|
validator.validate_container_start("container-123")
|
||||||
|
validator.validate_container_stop("container-123")
|
||||||
|
validator.validate_container_remove("container-123")
|
||||||
|
|
||||||
|
# Cannot operate on non-session container
|
||||||
|
with pytest.raises(SecurityError) as exc_info:
|
||||||
|
validator.validate_container_start("unknown-container")
|
||||||
|
assert "session" in str(exc_info.value).lower() or "not found" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
with pytest.raises(SecurityError):
|
||||||
|
validator.validate_container_stop("unknown-container")
|
||||||
|
|
||||||
|
with pytest.raises(SecurityError):
|
||||||
|
validator.validate_container_remove("unknown-container")
|
||||||
|
|
||||||
|
|
||||||
|
def test_unregister_session_container():
|
||||||
|
"""Test that session containers can be unregistered."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
validator.register_session_container("container-123")
|
||||||
|
validator.validate_container_start("container-123") # Should work
|
||||||
|
|
||||||
|
validator.unregister_session_container("container-123")
|
||||||
|
|
||||||
|
# After unregistration, should not work
|
||||||
|
with pytest.raises(Exception): # SecurityError
|
||||||
|
validator.validate_container_start("container-123")
|
||||||
|
|
||||||
|
|
||||||
|
def test_validate_container_create_stores_container_id():
|
||||||
|
"""Test that validate_container_create automatically registers container."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
# Create container with session_id should auto-register
|
||||||
|
validator.validate_container_create(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
params={
|
||||||
|
"network_mode": "none",
|
||||||
|
"read_only": True,
|
||||||
|
"security_opt": ["no-new-privileges"],
|
||||||
|
"user": "1000:1000"
|
||||||
|
},
|
||||||
|
session_id="test-session-123"
|
||||||
|
)
|
||||||
|
|
||||||
|
# The actual container_id would be returned by Podman after creation
|
||||||
|
# So this test just verifies validation passes
|
||||||
|
|
||||||
|
|
||||||
|
def test_security_error_includes_rule_violation():
|
||||||
|
"""Test that SecurityError messages indicate what rule was violated."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
# Test various violations have clear messages
|
||||||
|
with pytest.raises(SecurityError) as exc_info:
|
||||||
|
validator.validate_container_create(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
params={"privileged": True},
|
||||||
|
session_id="test"
|
||||||
|
)
|
||||||
|
error_msg = str(exc_info.value)
|
||||||
|
assert "privileged" in error_msg.lower()
|
||||||
|
|
||||||
|
with pytest.raises(SecurityError) as exc_info:
|
||||||
|
validator.validate_image_name("bad-image:latest")
|
||||||
|
error_msg = str(exc_info.value)
|
||||||
|
assert "image" in error_msg.lower() or "allowlist" in error_msg.lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_wildcard_image_pattern_matching():
|
||||||
|
"""Test that wildcard patterns work in image allowlist."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
# mcp-forge/custom:* should match any tag
|
||||||
|
validator.validate_image_name("mcp-forge/custom:my-env-v1")
|
||||||
|
validator.validate_image_name("mcp-forge/custom:another-tag")
|
||||||
|
validator.validate_image_name("mcp-forge/custom:abc123")
|
||||||
|
|
||||||
|
|
||||||
|
def test_forbidden_mount_paths_comprehensive():
|
||||||
|
"""Test all forbidden mount paths are blocked."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
forbidden_paths = [
|
||||||
|
"/",
|
||||||
|
"/etc",
|
||||||
|
"/etc/shadow",
|
||||||
|
"/var/run/docker.sock",
|
||||||
|
"/var/run/podman/podman.sock",
|
||||||
|
"/sys",
|
||||||
|
"/sys/kernel",
|
||||||
|
"/proc",
|
||||||
|
"/proc/self",
|
||||||
|
]
|
||||||
|
|
||||||
|
for path in forbidden_paths:
|
||||||
|
with pytest.raises(SecurityError):
|
||||||
|
validator.validate_volume_mount(path, "test-session")
|
||||||
|
|
||||||
|
|
||||||
|
def test_pid_mode_forbidden():
|
||||||
|
"""Test that pid_mode parameter is forbidden."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
with pytest.raises(SecurityError):
|
||||||
|
validator.validate_container_create(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
params={
|
||||||
|
"pid_mode": "host", # FORBIDDEN
|
||||||
|
"network_mode": "none"
|
||||||
|
},
|
||||||
|
session_id="test"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_ipc_mode_forbidden():
|
||||||
|
"""Test that ipc_mode parameter is forbidden."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
with pytest.raises(SecurityError):
|
||||||
|
validator.validate_container_create(
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
params={
|
||||||
|
"ipc_mode": "host", # FORBIDDEN
|
||||||
|
"network_mode": "none"
|
||||||
|
},
|
||||||
|
session_id="test"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_session_path_must_match_session_id():
|
||||||
|
"""Test that session paths must match the provided session_id."""
|
||||||
|
from mcp_forge.security.allowlist import OperationValidator, SecurityError
|
||||||
|
from mcp_forge.config.schema import SecurityConfig
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
config = SecurityConfig(audit_log=Path("/audit.log"))
|
||||||
|
validator = OperationValidator(config)
|
||||||
|
|
||||||
|
# Correct session match
|
||||||
|
validator.validate_volume_mount(
|
||||||
|
"/mcp-forge/sessions/session-123/workdir",
|
||||||
|
"session-123"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Wrong session in path
|
||||||
|
with pytest.raises(SecurityError):
|
||||||
|
validator.validate_volume_mount(
|
||||||
|
"/mcp-forge/sessions/other-session/workdir",
|
||||||
|
"session-123"
|
||||||
|
)
|
||||||
369
tests/security/test_audit.py
Normal file
369
tests/security/test_audit.py
Normal file
|
|
@ -0,0 +1,369 @@
|
||||||
|
"""
|
||||||
|
Tests for audit logger module.
|
||||||
|
|
||||||
|
Following TDD approach - these tests are written before implementation.
|
||||||
|
Tests cover all logging requirements from todo.md section 1.2.3.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import json
|
||||||
|
from pathlib import Path
|
||||||
|
import tempfile
|
||||||
|
import threading
|
||||||
|
import time
|
||||||
|
from datetime import datetime
|
||||||
|
|
||||||
|
|
||||||
|
def test_log_entries_written_to_file(tmp_path):
|
||||||
|
"""Test that log entries are written to the log file."""
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
|
||||||
|
log_file = tmp_path / "audit.log"
|
||||||
|
logger = AuditLogger(log_file)
|
||||||
|
|
||||||
|
logger.log(
|
||||||
|
event_type=AuditEventType.CONTAINER_CREATE,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message="Test container created",
|
||||||
|
details={"image": "test-image", "container_id": "abc123"}
|
||||||
|
)
|
||||||
|
|
||||||
|
assert log_file.exists()
|
||||||
|
content = log_file.read_text()
|
||||||
|
assert len(content) > 0
|
||||||
|
|
||||||
|
|
||||||
|
def test_log_entries_are_valid_json(tmp_path):
|
||||||
|
"""Test that log entries are valid JSON."""
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
|
||||||
|
log_file = tmp_path / "audit.log"
|
||||||
|
logger = AuditLogger(log_file)
|
||||||
|
|
||||||
|
logger.log(
|
||||||
|
event_type=AuditEventType.CONTAINER_CREATE,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message="Test entry",
|
||||||
|
details={"key": "value"}
|
||||||
|
)
|
||||||
|
|
||||||
|
# Each line should be valid JSON
|
||||||
|
lines = log_file.read_text().strip().split("\n")
|
||||||
|
for line in lines:
|
||||||
|
data = json.loads(line) # Should not raise
|
||||||
|
assert isinstance(data, dict)
|
||||||
|
|
||||||
|
|
||||||
|
def test_log_entries_contain_required_fields(tmp_path):
|
||||||
|
"""Test that log entries contain all required fields."""
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
|
||||||
|
log_file = tmp_path / "audit.log"
|
||||||
|
logger = AuditLogger(log_file)
|
||||||
|
|
||||||
|
logger.log(
|
||||||
|
event_type=AuditEventType.EXECUTION_REQUEST,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message="Code execution requested",
|
||||||
|
session_id="test-session-123"
|
||||||
|
)
|
||||||
|
|
||||||
|
lines = log_file.read_text().strip().split("\n")
|
||||||
|
entry = json.loads(lines[0])
|
||||||
|
|
||||||
|
# Required fields
|
||||||
|
assert "timestamp" in entry
|
||||||
|
assert "event_type" in entry
|
||||||
|
assert "severity" in entry
|
||||||
|
assert "message" in entry
|
||||||
|
assert "session_id" in entry
|
||||||
|
|
||||||
|
|
||||||
|
def test_timestamp_format_is_iso_8601(tmp_path):
|
||||||
|
"""Test that timestamp is in ISO 8601 format."""
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
|
||||||
|
log_file = tmp_path / "audit.log"
|
||||||
|
logger = AuditLogger(log_file)
|
||||||
|
|
||||||
|
logger.log(
|
||||||
|
event_type=AuditEventType.SESSION_CREATE,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message="Session created"
|
||||||
|
)
|
||||||
|
|
||||||
|
lines = log_file.read_text().strip().split("\n")
|
||||||
|
entry = json.loads(lines[0])
|
||||||
|
|
||||||
|
# Should be parseable as ISO 8601
|
||||||
|
timestamp = entry["timestamp"]
|
||||||
|
dt = datetime.fromisoformat(timestamp)
|
||||||
|
assert isinstance(dt, datetime)
|
||||||
|
|
||||||
|
|
||||||
|
def test_concurrent_logging_is_thread_safe(tmp_path):
|
||||||
|
"""Test that concurrent logging from multiple threads is thread-safe."""
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
|
||||||
|
log_file = tmp_path / "audit.log"
|
||||||
|
logger = AuditLogger(log_file)
|
||||||
|
|
||||||
|
def log_entries(thread_id, count):
|
||||||
|
for i in range(count):
|
||||||
|
logger.log(
|
||||||
|
event_type=AuditEventType.CONTAINER_CREATE,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message=f"Thread {thread_id} entry {i}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create multiple threads
|
||||||
|
threads = []
|
||||||
|
entries_per_thread = 10
|
||||||
|
num_threads = 5
|
||||||
|
|
||||||
|
for i in range(num_threads):
|
||||||
|
t = threading.Thread(target=log_entries, args=(i, entries_per_thread))
|
||||||
|
threads.append(t)
|
||||||
|
t.start()
|
||||||
|
|
||||||
|
# Wait for all threads
|
||||||
|
for t in threads:
|
||||||
|
t.join()
|
||||||
|
|
||||||
|
# Verify all entries written
|
||||||
|
lines = log_file.read_text().strip().split("\n")
|
||||||
|
assert len(lines) == num_threads * entries_per_thread
|
||||||
|
|
||||||
|
# Verify all entries are valid JSON
|
||||||
|
for line in lines:
|
||||||
|
json.loads(line)
|
||||||
|
|
||||||
|
|
||||||
|
def test_security_violations_logged_with_correct_severity(tmp_path):
|
||||||
|
"""Test that security violations are logged at correct severity."""
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
|
||||||
|
log_file = tmp_path / "audit.log"
|
||||||
|
logger = AuditLogger(log_file)
|
||||||
|
|
||||||
|
logger.log_security_violation(
|
||||||
|
operation="container_create",
|
||||||
|
reason="Privileged mode attempted",
|
||||||
|
session_id="test-session"
|
||||||
|
)
|
||||||
|
|
||||||
|
lines = log_file.read_text().strip().split("\n")
|
||||||
|
entry = json.loads(lines[0])
|
||||||
|
|
||||||
|
assert entry["severity"] == "critical"
|
||||||
|
assert entry["event_type"] == "security.violation"
|
||||||
|
|
||||||
|
|
||||||
|
def test_pii_is_not_logged(tmp_path):
|
||||||
|
"""Test that PII (code content, tokens, files) is not logged."""
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
|
||||||
|
log_file = tmp_path / "audit.log"
|
||||||
|
logger = AuditLogger(log_file)
|
||||||
|
|
||||||
|
# Log execution request - should NOT include actual code
|
||||||
|
logger.log(
|
||||||
|
event_type=AuditEventType.EXECUTION_REQUEST,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message="Code execution requested",
|
||||||
|
details={
|
||||||
|
"code_hash": "abc123def456", # Hash is OK
|
||||||
|
# "code": "print('hello')" # Should NOT be logged
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
content = log_file.read_text()
|
||||||
|
# Should not contain actual code
|
||||||
|
assert "print" not in content
|
||||||
|
assert "hello" not in content
|
||||||
|
# Should contain hash
|
||||||
|
assert "abc123def456" in content
|
||||||
|
|
||||||
|
|
||||||
|
def test_log_container_operation(tmp_path):
|
||||||
|
"""Test log_container_operation convenience method."""
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType
|
||||||
|
|
||||||
|
log_file = tmp_path / "audit.log"
|
||||||
|
logger = AuditLogger(log_file)
|
||||||
|
|
||||||
|
logger.log_container_operation(
|
||||||
|
operation="create",
|
||||||
|
container_id="container-123",
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
session_id="session-456"
|
||||||
|
)
|
||||||
|
|
||||||
|
lines = log_file.read_text().strip().split("\n")
|
||||||
|
entry = json.loads(lines[0])
|
||||||
|
|
||||||
|
assert entry["event_type"] == "container.create"
|
||||||
|
assert entry["container_id"] == "container-123"
|
||||||
|
assert entry["image"] == "mcp-forge/python:3.11"
|
||||||
|
assert entry["session_id"] == "session-456"
|
||||||
|
|
||||||
|
|
||||||
|
def test_log_container_operation_with_error(tmp_path):
|
||||||
|
"""Test logging container operation with error."""
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
|
||||||
|
log_file = tmp_path / "audit.log"
|
||||||
|
logger = AuditLogger(log_file)
|
||||||
|
|
||||||
|
logger.log_container_operation(
|
||||||
|
operation="start",
|
||||||
|
container_id="container-123",
|
||||||
|
image="mcp-forge/python:3.11",
|
||||||
|
session_id="session-456",
|
||||||
|
error="Container not found"
|
||||||
|
)
|
||||||
|
|
||||||
|
lines = log_file.read_text().strip().split("\n")
|
||||||
|
entry = json.loads(lines[0])
|
||||||
|
|
||||||
|
assert "error" in entry
|
||||||
|
assert entry["error"] == "Container not found"
|
||||||
|
assert entry["severity"] == "error"
|
||||||
|
|
||||||
|
|
||||||
|
def test_log_creates_directory_if_not_exists(tmp_path):
|
||||||
|
"""Test that logger creates log directory if it doesn't exist."""
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
|
||||||
|
log_dir = tmp_path / "nested" / "log" / "dir"
|
||||||
|
log_file = log_dir / "audit.log"
|
||||||
|
|
||||||
|
logger = AuditLogger(log_file)
|
||||||
|
|
||||||
|
logger.log(
|
||||||
|
event_type=AuditEventType.SESSION_CREATE,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message="Test"
|
||||||
|
)
|
||||||
|
|
||||||
|
assert log_file.exists()
|
||||||
|
assert log_file.parent.exists()
|
||||||
|
|
||||||
|
|
||||||
|
def test_audit_event_types():
|
||||||
|
"""Test that all required audit event types are defined."""
|
||||||
|
from mcp_forge.security.audit import AuditEventType
|
||||||
|
|
||||||
|
# Required event types from todo.md
|
||||||
|
assert hasattr(AuditEventType, "CONTAINER_CREATE")
|
||||||
|
assert hasattr(AuditEventType, "CONTAINER_START")
|
||||||
|
assert hasattr(AuditEventType, "CONTAINER_STOP")
|
||||||
|
assert hasattr(AuditEventType, "CONTAINER_REMOVE")
|
||||||
|
assert hasattr(AuditEventType, "EXECUTION_REQUEST")
|
||||||
|
assert hasattr(AuditEventType, "SECURITY_VIOLATION")
|
||||||
|
assert hasattr(AuditEventType, "BUILD_REQUEST")
|
||||||
|
assert hasattr(AuditEventType, "BUILD_COMPLETE")
|
||||||
|
assert hasattr(AuditEventType, "SESSION_CREATE")
|
||||||
|
assert hasattr(AuditEventType, "SESSION_DESTROY")
|
||||||
|
|
||||||
|
|
||||||
|
def test_audit_severity_levels():
|
||||||
|
"""Test that all required severity levels are defined."""
|
||||||
|
from mcp_forge.security.audit import AuditSeverity
|
||||||
|
|
||||||
|
assert hasattr(AuditSeverity, "INFO")
|
||||||
|
assert hasattr(AuditSeverity, "WARNING")
|
||||||
|
assert hasattr(AuditSeverity, "ERROR")
|
||||||
|
assert hasattr(AuditSeverity, "CRITICAL")
|
||||||
|
|
||||||
|
|
||||||
|
def test_log_with_all_optional_parameters(tmp_path):
|
||||||
|
"""Test logging with all optional parameters provided."""
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
|
||||||
|
log_file = tmp_path / "audit.log"
|
||||||
|
logger = AuditLogger(log_file)
|
||||||
|
|
||||||
|
logger.log(
|
||||||
|
event_type=AuditEventType.BUILD_COMPLETE,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message="Build completed successfully",
|
||||||
|
session_id="session-123",
|
||||||
|
user_id="user-456",
|
||||||
|
details={"image": "custom-env", "duration": 120},
|
||||||
|
error=None
|
||||||
|
)
|
||||||
|
|
||||||
|
lines = log_file.read_text().strip().split("\n")
|
||||||
|
entry = json.loads(lines[0])
|
||||||
|
|
||||||
|
assert entry["session_id"] == "session-123"
|
||||||
|
assert entry["user_id"] == "user-456"
|
||||||
|
assert entry["details"]["image"] == "custom-env"
|
||||||
|
assert entry["details"]["duration"] == 120
|
||||||
|
|
||||||
|
|
||||||
|
def test_multiple_log_entries_on_separate_lines(tmp_path):
|
||||||
|
"""Test that multiple log entries are written on separate lines."""
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
|
||||||
|
log_file = tmp_path / "audit.log"
|
||||||
|
logger = AuditLogger(log_file)
|
||||||
|
|
||||||
|
for i in range(5):
|
||||||
|
logger.log(
|
||||||
|
event_type=AuditEventType.CONTAINER_CREATE,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message=f"Entry {i}"
|
||||||
|
)
|
||||||
|
|
||||||
|
lines = log_file.read_text().strip().split("\n")
|
||||||
|
assert len(lines) == 5
|
||||||
|
|
||||||
|
# Each line should be parseable
|
||||||
|
for line in lines:
|
||||||
|
json.loads(line)
|
||||||
|
|
||||||
|
|
||||||
|
def test_log_security_violation_parameters(tmp_path):
|
||||||
|
"""Test log_security_violation includes all necessary information."""
|
||||||
|
from mcp_forge.security.audit import AuditLogger
|
||||||
|
|
||||||
|
log_file = tmp_path / "audit.log"
|
||||||
|
logger = AuditLogger(log_file)
|
||||||
|
|
||||||
|
logger.log_security_violation(
|
||||||
|
operation="volume_mount",
|
||||||
|
reason="Attempted to mount /etc",
|
||||||
|
session_id="session-789"
|
||||||
|
)
|
||||||
|
|
||||||
|
lines = log_file.read_text().strip().split("\n")
|
||||||
|
entry = json.loads(lines[0])
|
||||||
|
|
||||||
|
assert entry["event_type"] == "security.violation"
|
||||||
|
assert entry["severity"] == "critical"
|
||||||
|
assert entry["operation"] == "volume_mount"
|
||||||
|
assert entry["reason"] == "Attempted to mount /etc"
|
||||||
|
assert entry["session_id"] == "session-789"
|
||||||
|
|
||||||
|
|
||||||
|
def test_details_can_be_none(tmp_path):
|
||||||
|
"""Test that details parameter can be None."""
|
||||||
|
from mcp_forge.security.audit import AuditLogger, AuditEventType, AuditSeverity
|
||||||
|
|
||||||
|
log_file = tmp_path / "audit.log"
|
||||||
|
logger = AuditLogger(log_file)
|
||||||
|
|
||||||
|
logger.log(
|
||||||
|
event_type=AuditEventType.SESSION_CREATE,
|
||||||
|
severity=AuditSeverity.INFO,
|
||||||
|
message="Session created",
|
||||||
|
details=None
|
||||||
|
)
|
||||||
|
|
||||||
|
lines = log_file.read_text().strip().split("\n")
|
||||||
|
entry = json.loads(lines[0])
|
||||||
|
|
||||||
|
# Should work without error
|
||||||
|
assert "message" in entry
|
||||||
276
tests/security/test_resource_limits.py
Normal file
276
tests/security/test_resource_limits.py
Normal file
|
|
@ -0,0 +1,276 @@
|
||||||
|
"""
|
||||||
|
Tests for resource limits module.
|
||||||
|
|
||||||
|
Following TDD approach - these tests are written before implementation.
|
||||||
|
Tests cover all parsing and validation requirements from todo.md section 1.2.1.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_memory_string_megabytes():
|
||||||
|
"""Test parsing memory string with megabytes suffix."""
|
||||||
|
from mcp_forge.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
|
result = parse_memory_string("512m")
|
||||||
|
assert result == 536870912 # 512 * 1024 * 1024
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_memory_string_gigabytes():
|
||||||
|
"""Test parsing memory string with gigabytes suffix."""
|
||||||
|
from mcp_forge.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
|
result = parse_memory_string("2g")
|
||||||
|
assert result == 2147483648 # 2 * 1024 * 1024 * 1024
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_memory_string_kilobytes():
|
||||||
|
"""Test parsing memory string with kilobytes suffix."""
|
||||||
|
from mcp_forge.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
|
result = parse_memory_string("1024k")
|
||||||
|
assert result == 1048576 # 1024 * 1024
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_memory_string_case_insensitive():
|
||||||
|
"""Test that memory string parsing is case-insensitive."""
|
||||||
|
from mcp_forge.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
|
assert parse_memory_string("512M") == 536870912
|
||||||
|
assert parse_memory_string("2G") == 2147483648
|
||||||
|
assert parse_memory_string("1024K") == 1048576
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_memory_string_invalid_format_raises_value_error():
|
||||||
|
"""Test that invalid format raises ValueError."""
|
||||||
|
from mcp_forge.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
|
with pytest.raises(ValueError) as exc_info:
|
||||||
|
parse_memory_string("invalid")
|
||||||
|
assert "invalid" in str(exc_info.value).lower() or "format" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
parse_memory_string("512x") # Invalid suffix
|
||||||
|
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
parse_memory_string("abc") # Not a number
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_memory_string_negative_value_raises_value_error():
|
||||||
|
"""Test that negative values raise ValueError."""
|
||||||
|
from mcp_forge.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
|
with pytest.raises(ValueError) as exc_info:
|
||||||
|
parse_memory_string("-512m")
|
||||||
|
assert "positive" in str(exc_info.value).lower() or "negative" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_memory_string_zero_value_raises_value_error():
|
||||||
|
"""Test that zero value raises ValueError."""
|
||||||
|
from mcp_forge.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
|
with pytest.raises(ValueError) as exc_info:
|
||||||
|
parse_memory_string("0m")
|
||||||
|
assert "positive" in str(exc_info.value).lower() or "zero" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_cpu_quota_valid_value():
|
||||||
|
"""Test that valid CPU quota values are accepted."""
|
||||||
|
from mcp_forge.security.resource_limits import parse_cpu_quota
|
||||||
|
|
||||||
|
result = parse_cpu_quota(50000)
|
||||||
|
assert result == 50000
|
||||||
|
|
||||||
|
result = parse_cpu_quota(100000) # 100% of one core
|
||||||
|
assert result == 100000
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_cpu_quota_max_limit():
|
||||||
|
"""Test that CPU quota has a reasonable maximum (10 cores)."""
|
||||||
|
from mcp_forge.security.resource_limits import parse_cpu_quota
|
||||||
|
|
||||||
|
# Should accept up to 1000000 (10 cores)
|
||||||
|
result = parse_cpu_quota(1000000)
|
||||||
|
assert result == 1000000
|
||||||
|
|
||||||
|
# Should reject more than 10 cores
|
||||||
|
with pytest.raises(ValueError) as exc_info:
|
||||||
|
parse_cpu_quota(1000001)
|
||||||
|
assert "1000000" in str(exc_info.value) or "maximum" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_cpu_quota_negative_raises_value_error():
|
||||||
|
"""Test that negative CPU quota raises ValueError."""
|
||||||
|
from mcp_forge.security.resource_limits import parse_cpu_quota
|
||||||
|
|
||||||
|
with pytest.raises(ValueError) as exc_info:
|
||||||
|
parse_cpu_quota(-1)
|
||||||
|
assert "positive" in str(exc_info.value).lower() or "negative" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_cpu_quota_zero_raises_value_error():
|
||||||
|
"""Test that zero CPU quota raises ValueError."""
|
||||||
|
from mcp_forge.security.resource_limits import parse_cpu_quota
|
||||||
|
|
||||||
|
with pytest.raises(ValueError) as exc_info:
|
||||||
|
parse_cpu_quota(0)
|
||||||
|
assert "positive" in str(exc_info.value).lower() or "zero" in str(exc_info.value).lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_storage_string_same_as_memory():
|
||||||
|
"""Test that storage parsing works the same as memory parsing."""
|
||||||
|
from mcp_forge.security.resource_limits import parse_storage_string
|
||||||
|
|
||||||
|
assert parse_storage_string("1g") == 1073741824
|
||||||
|
assert parse_storage_string("512m") == 536870912
|
||||||
|
assert parse_storage_string("2048k") == 2097152
|
||||||
|
|
||||||
|
|
||||||
|
def test_resource_limits_class_initialization():
|
||||||
|
"""Test ResourceLimits class initializes correctly."""
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
limits = ResourceLimits(
|
||||||
|
memory="512m",
|
||||||
|
storage="1g",
|
||||||
|
cpu_quota=50000,
|
||||||
|
timeout=300
|
||||||
|
)
|
||||||
|
|
||||||
|
assert limits.memory_bytes == 536870912
|
||||||
|
assert limits.storage_bytes == 1073741824
|
||||||
|
assert limits.cpu_quota == 50000
|
||||||
|
assert limits.timeout == 300
|
||||||
|
|
||||||
|
|
||||||
|
def test_resource_limits_validates_memory():
|
||||||
|
"""Test that ResourceLimits validates memory string."""
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
ResourceLimits(
|
||||||
|
memory="invalid",
|
||||||
|
storage="1g",
|
||||||
|
cpu_quota=50000,
|
||||||
|
timeout=300
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_resource_limits_validates_storage():
|
||||||
|
"""Test that ResourceLimits validates storage string."""
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
ResourceLimits(
|
||||||
|
memory="512m",
|
||||||
|
storage="invalid",
|
||||||
|
cpu_quota=50000,
|
||||||
|
timeout=300
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_resource_limits_validates_cpu_quota():
|
||||||
|
"""Test that ResourceLimits validates CPU quota."""
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
ResourceLimits(
|
||||||
|
memory="512m",
|
||||||
|
storage="1g",
|
||||||
|
cpu_quota=-1,
|
||||||
|
timeout=300
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_resource_limits_to_podman_params():
|
||||||
|
"""Test conversion to Podman container parameters."""
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
limits = ResourceLimits(
|
||||||
|
memory="512m",
|
||||||
|
storage="1g",
|
||||||
|
cpu_quota=50000,
|
||||||
|
timeout=300
|
||||||
|
)
|
||||||
|
|
||||||
|
params = limits.to_podman_params()
|
||||||
|
|
||||||
|
assert isinstance(params, dict)
|
||||||
|
assert "mem_limit" in params
|
||||||
|
assert params["mem_limit"] == "536870912" # Should be string for Podman
|
||||||
|
# CPU quota is set via cpu_quota parameter
|
||||||
|
assert "cpu_quota" in params
|
||||||
|
assert params["cpu_quota"] == 50000
|
||||||
|
|
||||||
|
|
||||||
|
def test_resource_limits_default_timeout():
|
||||||
|
"""Test that ResourceLimits has a default timeout."""
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
limits = ResourceLimits(
|
||||||
|
memory="512m",
|
||||||
|
storage="1g",
|
||||||
|
cpu_quota=50000
|
||||||
|
)
|
||||||
|
|
||||||
|
assert limits.timeout == 300 # Default from signature
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_memory_string_with_spaces():
|
||||||
|
"""Test parsing memory strings that have spaces."""
|
||||||
|
from mcp_forge.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
|
# Should handle spaces gracefully (strip them)
|
||||||
|
result = parse_memory_string(" 512m ")
|
||||||
|
assert result == 536870912
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_memory_string_bytes_suffix():
|
||||||
|
"""Test parsing memory string with bytes suffix (no multiplier)."""
|
||||||
|
from mcp_forge.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
|
# Just a number (bytes) - should this be supported?
|
||||||
|
# Based on architecture, we support k, m, g suffixes
|
||||||
|
# Plain numbers should probably raise an error for safety
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
parse_memory_string("1024")
|
||||||
|
|
||||||
|
|
||||||
|
def test_resource_limits_storage_quota_in_podman_params():
|
||||||
|
"""Test that storage limits are included in Podman params."""
|
||||||
|
from mcp_forge.security.resource_limits import ResourceLimits
|
||||||
|
|
||||||
|
limits = ResourceLimits(
|
||||||
|
memory="512m",
|
||||||
|
storage="1g",
|
||||||
|
cpu_quota=50000,
|
||||||
|
timeout=300
|
||||||
|
)
|
||||||
|
|
||||||
|
params = limits.to_podman_params()
|
||||||
|
|
||||||
|
# Storage limit might be set via storage_opt or similar
|
||||||
|
# The exact parameter depends on Podman API
|
||||||
|
assert "storage_bytes" in params or "storage_opt" in params
|
||||||
|
|
||||||
|
|
||||||
|
def test_cpu_quota_explanation():
|
||||||
|
"""Test that CPU quota values have clear meaning."""
|
||||||
|
from mcp_forge.security.resource_limits import parse_cpu_quota
|
||||||
|
|
||||||
|
# 100000 = 100% of one CPU core
|
||||||
|
# 50000 = 50% of one CPU core
|
||||||
|
# 200000 = 200% = 2 CPU cores
|
||||||
|
|
||||||
|
assert parse_cpu_quota(50000) == 50000 # 0.5 cores
|
||||||
|
assert parse_cpu_quota(100000) == 100000 # 1 core
|
||||||
|
assert parse_cpu_quota(200000) == 200000 # 2 cores
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_memory_with_decimal():
|
||||||
|
"""Test parsing memory strings with decimal values."""
|
||||||
|
from mcp_forge.security.resource_limits import parse_memory_string
|
||||||
|
|
||||||
|
# Should handle decimals
|
||||||
|
result = parse_memory_string("1.5g")
|
||||||
|
assert result == 1610612736 # 1.5 * 1024 * 1024 * 1024
|
||||||
266
tests/server/test_resources.py
Normal file
266
tests/server/test_resources.py
Normal file
|
|
@ -0,0 +1,266 @@
|
||||||
|
"""Tests for MCP resource handlers."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import json
|
||||||
|
from unittest.mock import Mock, AsyncMock, patch
|
||||||
|
from mcp_forge.server.resources import ResourceHandler
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_client_manager():
|
||||||
|
"""Mock MCP client manager."""
|
||||||
|
manager = Mock()
|
||||||
|
manager.list_all_tools = AsyncMock(return_value=["read_file", "write_file", "calculate"])
|
||||||
|
return manager
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_session_manager():
|
||||||
|
"""Mock session manager."""
|
||||||
|
manager = Mock()
|
||||||
|
|
||||||
|
# Mock session state
|
||||||
|
mock_state = Mock()
|
||||||
|
mock_state.session_id = "test-session"
|
||||||
|
mock_state.documented_variables = {"x": "Test variable", "result": "Calculation result"}
|
||||||
|
mock_state.note = "Test session state"
|
||||||
|
mock_state.all_variables = ["x", "y", "result", "np", "pd"]
|
||||||
|
mock_state.introspection = {
|
||||||
|
"x": {"type": "int", "size": 28},
|
||||||
|
"result": {"type": "float", "size": 24}
|
||||||
|
}
|
||||||
|
mock_state.to_dict = Mock(return_value={
|
||||||
|
"session_id": "test-session",
|
||||||
|
"documented_variables": {"x": "Test variable", "result": "Calculation result"},
|
||||||
|
"note": "Test session state",
|
||||||
|
"all_variables": ["x", "y", "result", "np", "pd"],
|
||||||
|
"introspection": {
|
||||||
|
"x": {"type": "int", "size": 28},
|
||||||
|
"result": {"type": "float", "size": 24}
|
||||||
|
},
|
||||||
|
"last_updated": "2026-02-06T12:00:00Z"
|
||||||
|
})
|
||||||
|
|
||||||
|
manager.get_session_state = Mock(return_value=mock_state)
|
||||||
|
manager.list_sessions = Mock(return_value=[
|
||||||
|
{"session_id": "test-session", "created_at": "2026-02-06T12:00:00Z"}
|
||||||
|
])
|
||||||
|
|
||||||
|
return manager
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_environment_builder():
|
||||||
|
"""Mock environment builder."""
|
||||||
|
builder = Mock()
|
||||||
|
builder.list_templates = Mock(return_value={
|
||||||
|
"datascience": {
|
||||||
|
"description": "Data science environment with numpy, pandas, matplotlib",
|
||||||
|
"packages": ["numpy>=1.24", "pandas>=2.0", "matplotlib>=3.7"]
|
||||||
|
},
|
||||||
|
"ml": {
|
||||||
|
"description": "Machine learning environment",
|
||||||
|
"packages": ["scikit-learn>=1.3", "tensorflow>=2.13"]
|
||||||
|
}
|
||||||
|
})
|
||||||
|
return builder
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_config():
|
||||||
|
"""Mock forge configuration."""
|
||||||
|
config = Mock()
|
||||||
|
config.execution = Mock()
|
||||||
|
config.execution.default_backend = "simple"
|
||||||
|
config.execution.default_timeout = 300
|
||||||
|
config.sessions = Mock()
|
||||||
|
config.sessions.max_concurrent = 10
|
||||||
|
# Add environment-related config for environment/info resource
|
||||||
|
config.max_packages = 50
|
||||||
|
config.max_build_time = 300
|
||||||
|
config.base_images = {"python:3.11": {}, "python:3.12": {}}
|
||||||
|
return config
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def resource_handler(mock_client_manager, mock_session_manager, mock_environment_builder, mock_config):
|
||||||
|
"""Create ResourceHandler instance with mocked dependencies."""
|
||||||
|
return ResourceHandler(
|
||||||
|
client_manager=mock_client_manager,
|
||||||
|
session_manager=mock_session_manager,
|
||||||
|
environment_builder=mock_environment_builder,
|
||||||
|
config=mock_config
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_handle_tools_available(resource_handler, mock_client_manager):
|
||||||
|
"""Test handling tools/available resource returns list of available tools."""
|
||||||
|
result = await resource_handler.handle_resource("mcp://forge/tools/available")
|
||||||
|
|
||||||
|
assert str(result.uri) == "mcp://forge/tools/available"
|
||||||
|
assert result.mimeType == "application/json"
|
||||||
|
|
||||||
|
# Parse JSON content
|
||||||
|
tools = json.loads(result.text)
|
||||||
|
assert tools == {"tools": ["read_file", "write_file", "calculate"]}
|
||||||
|
|
||||||
|
# Verify client manager was called
|
||||||
|
mock_client_manager.list_all_tools.assert_awaited_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_handle_session_state(resource_handler, mock_session_manager):
|
||||||
|
"""Test handling session state resource returns documented state."""
|
||||||
|
session_id = "test-session"
|
||||||
|
result = await resource_handler.handle_resource(f"mcp://forge/sessions/{session_id}/state")
|
||||||
|
|
||||||
|
assert str(result.uri) == f"mcp://forge/sessions/{session_id}/state"
|
||||||
|
assert result.mimeType == "application/json"
|
||||||
|
|
||||||
|
# Parse JSON content
|
||||||
|
state = json.loads(result.text)
|
||||||
|
assert state["session_id"] == session_id
|
||||||
|
assert state["documented_variables"] == {"x": "Test variable", "result": "Calculation result"}
|
||||||
|
assert state["note"] == "Test session state"
|
||||||
|
assert "last_updated" in state
|
||||||
|
|
||||||
|
# Verify session manager was called
|
||||||
|
mock_session_manager.get_session_state.assert_called_once_with(session_id)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_handle_session_variables(resource_handler, mock_session_manager):
|
||||||
|
"""Test handling session variables resource returns list of variables."""
|
||||||
|
session_id = "test-session"
|
||||||
|
result = await resource_handler.handle_resource(f"mcp://forge/sessions/{session_id}/variables")
|
||||||
|
|
||||||
|
assert str(result.uri) == f"mcp://forge/sessions/{session_id}/variables"
|
||||||
|
assert result.mimeType == "application/json"
|
||||||
|
|
||||||
|
# Parse JSON content
|
||||||
|
data = json.loads(result.text)
|
||||||
|
assert data["variables"] == ["x", "y", "result", "np", "pd"]
|
||||||
|
|
||||||
|
# Verify session manager was called
|
||||||
|
mock_session_manager.get_session_state.assert_called_once_with(session_id)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_handle_environments_list(resource_handler, mock_environment_builder):
|
||||||
|
"""Test handling environments/list resource returns templates and built environments."""
|
||||||
|
result = await resource_handler.handle_resource("mcp://forge/environments/list")
|
||||||
|
|
||||||
|
assert str(result.uri) == "mcp://forge/environments/list"
|
||||||
|
assert result.mimeType == "application/json"
|
||||||
|
|
||||||
|
# Parse JSON content
|
||||||
|
environments = json.loads(result.text)
|
||||||
|
assert "templates" in environments
|
||||||
|
assert "datascience" in environments["templates"]
|
||||||
|
assert "ml" in environments["templates"]
|
||||||
|
assert environments["templates"]["datascience"]["description"] == "Data science environment with numpy, pandas, matplotlib"
|
||||||
|
|
||||||
|
# Verify environment builder was called
|
||||||
|
mock_environment_builder.list_templates.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_handle_environment_info(resource_handler, mock_config):
|
||||||
|
"""Test handling environment info resource returns configuration info."""
|
||||||
|
result = await resource_handler.handle_resource("mcp://forge/environment/info")
|
||||||
|
|
||||||
|
assert str(result.uri) == "mcp://forge/environment/info"
|
||||||
|
assert result.mimeType == "application/json"
|
||||||
|
|
||||||
|
# Parse JSON content
|
||||||
|
info = json.loads(result.text)
|
||||||
|
assert "base_images" in info
|
||||||
|
assert "python:3.11" in info["base_images"]
|
||||||
|
assert "python:3.12" in info["base_images"]
|
||||||
|
assert info["max_packages"] == 50
|
||||||
|
assert info["max_build_time"] == 300
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_handle_sessions_list(resource_handler, mock_session_manager):
|
||||||
|
"""Test handling sessions/list resource returns list of active sessions."""
|
||||||
|
result = await resource_handler.handle_resource("mcp://forge/sessions/list")
|
||||||
|
|
||||||
|
assert str(result.uri) == "mcp://forge/sessions/list"
|
||||||
|
assert result.mimeType == "application/json"
|
||||||
|
|
||||||
|
# Parse JSON content
|
||||||
|
data = json.loads(result.text)
|
||||||
|
assert len(data["sessions"]) == 1
|
||||||
|
assert data["sessions"][0]["session_id"] == "test-session"
|
||||||
|
|
||||||
|
# Verify session manager was called
|
||||||
|
mock_session_manager.list_sessions.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_handle_unknown_resource(resource_handler):
|
||||||
|
"""Test handling unknown resource raises ValueError."""
|
||||||
|
with pytest.raises(ValueError, match="Unknown resource URI"):
|
||||||
|
await resource_handler.handle_resource("mcp://forge/unknown/resource")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_handle_session_not_found(resource_handler, mock_session_manager):
|
||||||
|
"""Test handling session resource when session doesn't exist raises KeyError."""
|
||||||
|
mock_session_manager.get_session_state.side_effect = KeyError("Session not found")
|
||||||
|
|
||||||
|
with pytest.raises(KeyError, match="Session not found"):
|
||||||
|
await resource_handler.handle_resource("mcp://forge/sessions/nonexistent/state")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_parse_uri_tools_available(resource_handler):
|
||||||
|
"""Test URI parsing for tools/available resource."""
|
||||||
|
resource_type, params = resource_handler._parse_uri("mcp://forge/tools/available")
|
||||||
|
assert resource_type == "tools_available"
|
||||||
|
assert params == {}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_parse_uri_session_state(resource_handler):
|
||||||
|
"""Test URI parsing for session state resource."""
|
||||||
|
resource_type, params = resource_handler._parse_uri("mcp://forge/sessions/test-123/state")
|
||||||
|
assert resource_type == "session_state"
|
||||||
|
assert params == {"session_id": "test-123"}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_parse_uri_session_variables(resource_handler):
|
||||||
|
"""Test URI parsing for session variables resource."""
|
||||||
|
resource_type, params = resource_handler._parse_uri("mcp://forge/sessions/test-123/variables")
|
||||||
|
assert resource_type == "session_variables"
|
||||||
|
assert params == {"session_id": "test-123"}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_parse_uri_invalid_format(resource_handler):
|
||||||
|
"""Test URI parsing with invalid format raises ValueError."""
|
||||||
|
with pytest.raises(ValueError, match="Invalid URI format"):
|
||||||
|
resource_handler._parse_uri("invalid://uri")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_json_serialization_valid(resource_handler):
|
||||||
|
"""Test that all returned JSON content is valid and can be parsed."""
|
||||||
|
# Test all resource types return valid JSON
|
||||||
|
resources = [
|
||||||
|
"mcp://forge/tools/available",
|
||||||
|
"mcp://forge/sessions/test-session/state",
|
||||||
|
"mcp://forge/sessions/test-session/variables",
|
||||||
|
"mcp://forge/environments/list",
|
||||||
|
"mcp://forge/environment/info",
|
||||||
|
"mcp://forge/sessions/list"
|
||||||
|
]
|
||||||
|
|
||||||
|
for uri in resources:
|
||||||
|
result = await resource_handler.handle_resource(uri)
|
||||||
|
# Should not raise exception
|
||||||
|
parsed = json.loads(result.text)
|
||||||
|
assert parsed is not None
|
||||||
203
tests/server/test_server.py
Normal file
203
tests/server/test_server.py
Normal file
|
|
@ -0,0 +1,203 @@
|
||||||
|
"""Tests for MCP Forge Server."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from unittest.mock import Mock, AsyncMock, patch
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
from mcp_forge.server.server import ForgeServer
|
||||||
|
from mcp_forge.config.schema import ForgeConfig
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mock_config(tmp_path):
|
||||||
|
"""Mock forge configuration."""
|
||||||
|
config = Mock(spec=ForgeConfig)
|
||||||
|
|
||||||
|
# Server config
|
||||||
|
config.server = Mock()
|
||||||
|
config.server.host = "localhost"
|
||||||
|
config.server.port = 3000
|
||||||
|
config.server.podman_socket = Path("/run/user/1000/podman/podman.sock")
|
||||||
|
|
||||||
|
# Security config
|
||||||
|
config.security = Mock()
|
||||||
|
config.security.audit_log = tmp_path / "audit.log"
|
||||||
|
config.security.max_memory = "2g"
|
||||||
|
config.security.max_timeout = 1800
|
||||||
|
|
||||||
|
# Execution config
|
||||||
|
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"
|
||||||
|
|
||||||
|
# Sessions config
|
||||||
|
config.sessions = Mock()
|
||||||
|
config.sessions.max_concurrent = 10
|
||||||
|
config.sessions.idle_timeout = 3600
|
||||||
|
|
||||||
|
# Environment builder config
|
||||||
|
config.environment_builder = Mock()
|
||||||
|
config.environment_builder.uv_cache_path = tmp_path / "cache"
|
||||||
|
config.environment_builder.max_build_time = 600
|
||||||
|
config.environment_builder.package_validation = Mock()
|
||||||
|
|
||||||
|
return config
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_server_initialization(mock_config):
|
||||||
|
"""Test that server initializes all components."""
|
||||||
|
with patch('mcp_forge.server.server.AuditLogger'), \
|
||||||
|
patch('mcp_forge.server.server.OperationValidator'), \
|
||||||
|
patch('mcp_forge.server.server.PodmanClient'), \
|
||||||
|
patch('mcp_forge.server.server.SecureContainerManager'), \
|
||||||
|
patch('mcp_forge.server.server.MCPClientManager'), \
|
||||||
|
patch('mcp_forge.server.server.ToolBridgeServer'), \
|
||||||
|
patch('mcp_forge.server.server.ToolInjectionGenerator'), \
|
||||||
|
patch('mcp_forge.server.server.SimpleBackend'), \
|
||||||
|
patch('mcp_forge.server.server.JupyterKernelManager'), \
|
||||||
|
patch('mcp_forge.server.server.JupyterBackend'), \
|
||||||
|
patch('mcp_forge.server.server.SessionManager'), \
|
||||||
|
patch('mcp_forge.server.server.EnvironmentBuilder'):
|
||||||
|
|
||||||
|
server = ForgeServer(config=mock_config)
|
||||||
|
|
||||||
|
# Verify server was created
|
||||||
|
assert server is not None
|
||||||
|
assert server.config == mock_config
|
||||||
|
|
||||||
|
# Verify components were initialized
|
||||||
|
assert hasattr(server, 'audit_logger')
|
||||||
|
assert hasattr(server, 'operation_validator')
|
||||||
|
assert hasattr(server, 'podman_client')
|
||||||
|
assert hasattr(server, 'container_manager')
|
||||||
|
assert hasattr(server, 'client_manager')
|
||||||
|
assert hasattr(server, 'bridge_server')
|
||||||
|
assert hasattr(server, 'simple_backend')
|
||||||
|
assert hasattr(server, 'jupyter_backend')
|
||||||
|
assert hasattr(server, 'environment_builder')
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_tools_registration(mock_config):
|
||||||
|
"""Test that tools are registered with the server."""
|
||||||
|
with patch('mcp_forge.server.server.AuditLogger'), \
|
||||||
|
patch('mcp_forge.server.server.OperationValidator'), \
|
||||||
|
patch('mcp_forge.server.server.PodmanClient'), \
|
||||||
|
patch('mcp_forge.server.server.SecureContainerManager'), \
|
||||||
|
patch('mcp_forge.server.server.MCPClientManager'), \
|
||||||
|
patch('mcp_forge.server.server.ToolBridgeServer'), \
|
||||||
|
patch('mcp_forge.server.server.ToolInjectionGenerator'), \
|
||||||
|
patch('mcp_forge.server.server.SimpleBackend'), \
|
||||||
|
patch('mcp_forge.server.server.JupyterKernelManager'), \
|
||||||
|
patch('mcp_forge.server.server.JupyterBackend'), \
|
||||||
|
patch('mcp_forge.server.server.SessionManager'), \
|
||||||
|
patch('mcp_forge.server.server.EnvironmentBuilder'):
|
||||||
|
|
||||||
|
server = ForgeServer(config=mock_config)
|
||||||
|
|
||||||
|
# Verify tool instances were created
|
||||||
|
assert hasattr(server, 'execute_python_tool')
|
||||||
|
assert hasattr(server, 'document_state_tool')
|
||||||
|
assert hasattr(server, 'build_environment_tool')
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_resources_registration(mock_config):
|
||||||
|
"""Test that resources are registered with the server."""
|
||||||
|
with patch('mcp_forge.server.server.AuditLogger'), \
|
||||||
|
patch('mcp_forge.server.server.OperationValidator'), \
|
||||||
|
patch('mcp_forge.server.server.PodmanClient'), \
|
||||||
|
patch('mcp_forge.server.server.SecureContainerManager'), \
|
||||||
|
patch('mcp_forge.server.server.MCPClientManager'), \
|
||||||
|
patch('mcp_forge.server.server.ToolBridgeServer'), \
|
||||||
|
patch('mcp_forge.server.server.ToolInjectionGenerator'), \
|
||||||
|
patch('mcp_forge.server.server.SimpleBackend'), \
|
||||||
|
patch('mcp_forge.server.server.JupyterKernelManager'), \
|
||||||
|
patch('mcp_forge.server.server.JupyterBackend'), \
|
||||||
|
patch('mcp_forge.server.server.SessionManager'), \
|
||||||
|
patch('mcp_forge.server.server.EnvironmentBuilder'):
|
||||||
|
|
||||||
|
server = ForgeServer(config=mock_config)
|
||||||
|
|
||||||
|
# Verify resource handler was created
|
||||||
|
assert hasattr(server, 'resource_handler')
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_server_shutdown(mock_config):
|
||||||
|
"""Test that server shuts down gracefully."""
|
||||||
|
with patch('mcp_forge.server.server.AuditLogger') as mock_audit, \
|
||||||
|
patch('mcp_forge.server.server.OperationValidator'), \
|
||||||
|
patch('mcp_forge.server.server.PodmanClient'), \
|
||||||
|
patch('mcp_forge.server.server.SecureContainerManager'), \
|
||||||
|
patch('mcp_forge.server.server.MCPClientManager') as mock_client_mgr, \
|
||||||
|
patch('mcp_forge.server.server.ToolBridgeServer') as mock_bridge, \
|
||||||
|
patch('mcp_forge.server.server.ToolInjectionGenerator'), \
|
||||||
|
patch('mcp_forge.server.server.SimpleBackend'), \
|
||||||
|
patch('mcp_forge.server.server.JupyterKernelManager'), \
|
||||||
|
patch('mcp_forge.server.server.JupyterBackend'), \
|
||||||
|
patch('mcp_forge.server.server.SessionManager'), \
|
||||||
|
patch('mcp_forge.server.server.EnvironmentBuilder'):
|
||||||
|
|
||||||
|
# Setup mocks
|
||||||
|
mock_client_mgr.return_value.shutdown = AsyncMock()
|
||||||
|
mock_bridge_instance = Mock()
|
||||||
|
mock_bridge_instance.stop = Mock() # Not async
|
||||||
|
mock_bridge.return_value = mock_bridge_instance
|
||||||
|
mock_audit_instance = Mock()
|
||||||
|
mock_audit.return_value = mock_audit_instance
|
||||||
|
|
||||||
|
server = ForgeServer(config=mock_config)
|
||||||
|
|
||||||
|
# Shutdown server
|
||||||
|
await server.shutdown()
|
||||||
|
|
||||||
|
# Verify cleanup was called
|
||||||
|
server.client_manager.shutdown.assert_called_once()
|
||||||
|
server.bridge_server.stop.assert_called_once()
|
||||||
|
server.audit_logger.log.assert_called()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_server_component_initialization_order(mock_config):
|
||||||
|
"""Test that components are initialized in correct order."""
|
||||||
|
init_order = []
|
||||||
|
|
||||||
|
def track_init(name):
|
||||||
|
def decorator(cls):
|
||||||
|
original_init = cls.__init__
|
||||||
|
def new_init(self, *args, **kwargs):
|
||||||
|
init_order.append(name)
|
||||||
|
return original_init(self, *args, **kwargs)
|
||||||
|
cls.__init__ = new_init
|
||||||
|
return cls
|
||||||
|
return decorator
|
||||||
|
|
||||||
|
with patch('mcp_forge.server.server.AuditLogger') as mock_audit, \
|
||||||
|
patch('mcp_forge.server.server.OperationValidator') as mock_validator, \
|
||||||
|
patch('mcp_forge.server.server.PodmanClient') as mock_podman, \
|
||||||
|
patch('mcp_forge.server.server.SecureContainerManager') as mock_container, \
|
||||||
|
patch('mcp_forge.server.server.MCPClientManager'), \
|
||||||
|
patch('mcp_forge.server.server.ToolBridgeServer'), \
|
||||||
|
patch('mcp_forge.server.server.ToolInjectionGenerator'), \
|
||||||
|
patch('mcp_forge.server.server.SimpleBackend'), \
|
||||||
|
patch('mcp_forge.server.server.JupyterKernelManager'), \
|
||||||
|
patch('mcp_forge.server.server.JupyterBackend'), \
|
||||||
|
patch('mcp_forge.server.server.SessionManager'), \
|
||||||
|
patch('mcp_forge.server.server.EnvironmentBuilder'):
|
||||||
|
|
||||||
|
mock_audit.side_effect = lambda *args, **kwargs: init_order.append('audit_logger') or Mock()
|
||||||
|
mock_validator.side_effect = lambda *args, **kwargs: init_order.append('operation_validator') or Mock()
|
||||||
|
mock_podman.side_effect = lambda *args, **kwargs: init_order.append('podman_client') or Mock()
|
||||||
|
mock_container.side_effect = lambda *args, **kwargs: init_order.append('container_manager') or Mock()
|
||||||
|
|
||||||
|
server = ForgeServer(config=mock_config)
|
||||||
|
|
||||||
|
# Verify security components are initialized first
|
||||||
|
assert init_order.index('audit_logger') < init_order.index('podman_client')
|
||||||
|
assert init_order.index('operation_validator') < init_order.index('podman_client')
|
||||||
|
assert init_order.index('podman_client') < init_order.index('container_manager')
|
||||||
86
tests/server/test_server_integration.py
Normal file
86
tests/server/test_server_integration.py
Normal file
|
|
@ -0,0 +1,86 @@
|
||||||
|
"""Integration tests for MCP Forge Server - tests real component initialization."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from pathlib import Path
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
from mcp_forge.server.server import ForgeServer
|
||||||
|
from mcp_forge.config.schema import (
|
||||||
|
ForgeConfig, ServerConfig, SecurityConfig, ExecutionConfig, SessionConfig,
|
||||||
|
ImageConfig, VolumeConfig, EnvironmentBuilderConfig, PackageValidationConfig
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def real_config(tmp_path):
|
||||||
|
"""Real configuration with all required fields."""
|
||||||
|
# Create required files
|
||||||
|
(tmp_path / "allowlist.txt").write_text("requests\npandas\nnumpy\n")
|
||||||
|
(tmp_path / "blocklist.txt").write_text("")
|
||||||
|
|
||||||
|
config = ForgeConfig(
|
||||||
|
server=ServerConfig(
|
||||||
|
host="localhost",
|
||||||
|
port=3000,
|
||||||
|
podman_socket=Path("/run/user/1000/podman/podman.sock")
|
||||||
|
),
|
||||||
|
security=SecurityConfig(
|
||||||
|
audit_log=tmp_path / "audit.log",
|
||||||
|
enforce_resource_limits=True,
|
||||||
|
allow_network=False
|
||||||
|
),
|
||||||
|
execution=ExecutionConfig(
|
||||||
|
default_backend="simple",
|
||||||
|
default_timeout=300,
|
||||||
|
max_timeout=1800,
|
||||||
|
default_memory="512m",
|
||||||
|
max_memory="2g"
|
||||||
|
),
|
||||||
|
images=ImageConfig(
|
||||||
|
allowed_python_versions=["3.11", "3.12"],
|
||||||
|
default_base_image="python:3.11"
|
||||||
|
),
|
||||||
|
sessions=SessionConfig(
|
||||||
|
max_concurrent=10,
|
||||||
|
idle_timeout=3600
|
||||||
|
),
|
||||||
|
volumes=VolumeConfig(
|
||||||
|
base_path=tmp_path / "volumes"
|
||||||
|
),
|
||||||
|
environment_builder=EnvironmentBuilderConfig(
|
||||||
|
uv_cache_path=tmp_path / "cache",
|
||||||
|
build_rate_limit={"requests": 5, "period": 60},
|
||||||
|
package_validation=PackageValidationConfig(
|
||||||
|
allowlist_path=tmp_path / "allowlist.txt",
|
||||||
|
blocklist_path=tmp_path / "blocklist.txt"
|
||||||
|
)
|
||||||
|
),
|
||||||
|
mcp_tools={}
|
||||||
|
)
|
||||||
|
return config
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_server_can_initialize_with_minimal_mocking(real_config):
|
||||||
|
"""Test that ForgeServer can initialize with real component instances."""
|
||||||
|
# Only mock the podman library since we don't have a real Podman socket
|
||||||
|
with patch('mcp_forge.podman.client.BasePodmanClient'):
|
||||||
|
# This should succeed if all parameter mismatches are fixed
|
||||||
|
server = ForgeServer(config=real_config)
|
||||||
|
|
||||||
|
# Verify all components were created
|
||||||
|
assert server.audit_logger is not None
|
||||||
|
assert server.operation_validator is not None
|
||||||
|
assert server.podman_client is not None
|
||||||
|
assert server.container_manager is not None
|
||||||
|
assert server.client_manager is not None
|
||||||
|
assert server.bridge_server is not None
|
||||||
|
assert server.simple_backend is not None
|
||||||
|
assert server.kernel_manager is not None
|
||||||
|
assert server.session_manager is not None
|
||||||
|
assert server.jupyter_backend is not None
|
||||||
|
assert server.environment_builder is not None
|
||||||
|
# Sub-components created by EnvironmentBuilder
|
||||||
|
assert server.package_validator is not None
|
||||||
|
assert server.uv_installer is not None
|
||||||
|
assert server.image_builder is not None
|
||||||
250
tests/server/tools/test_build_environment.py
Normal file
250
tests/server/tools/test_build_environment.py
Normal file
|
|
@ -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"]
|
||||||
188
tests/server/tools/test_document_state.py
Normal file
188
tests/server/tools/test_document_state.py
Normal file
|
|
@ -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)
|
||||||
389
tests/server/tools/test_execute_python.py
Normal file
389
tests/server/tools/test_execute_python.py
Normal file
|
|
@ -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)
|
||||||
Loading…
Add table
Add a link
Reference in a new issue