diff --git a/.gitignore b/.gitignore index 5cf9451b..21bb18ab 100644 --- a/.gitignore +++ b/.gitignore @@ -180,6 +180,9 @@ openai_api_key.txt # Dev stuff dev/ +# Mock server certificates (auto-generated) +dev-tools/mcp-mock-server/.certs/ + # VSCode .vscode/ diff --git a/Makefile b/Makefile index 5ca07bed..1651398b 100644 --- a/Makefile +++ b/Makefile @@ -31,10 +31,10 @@ test-e2e-local: ## Run end to end tests for the service check-types: ## Checks type hints in sources - uv run mypy --explicit-package-bases --disallow-untyped-calls --disallow-untyped-defs --disallow-incomplete-defs --ignore-missing-imports --disable-error-code attr-defined src/ tests/unit tests/integration tests/e2e/ + uv run mypy --explicit-package-bases --disallow-untyped-calls --disallow-untyped-defs --disallow-incomplete-defs --ignore-missing-imports --disable-error-code attr-defined src/ tests/unit tests/integration tests/e2e/ dev-tools/ security-check: ## Check the project for security issues - bandit -c pyproject.toml -r src tests + uv run bandit -c pyproject.toml -r src tests dev-tools format: ## Format the code into unified format uv run black . @@ -84,13 +84,13 @@ black: ## Check source code using Black code formatter uv run black --check . pylint: ## Check source code using Pylint static code analyser - uv run pylint src tests + uv run pylint src tests dev-tools pyright: ## Check source code using Pyright static type checker - uv run pyright src + uv run pyright src dev-tools docstyle: ## Check the docstring style using Docstyle checker - uv run pydocstyle -v src + uv run pydocstyle -v src dev-tools ruff: ## Check source code using Ruff linter uv run ruff check . --per-file-ignores=tests/*:S101 --per-file-ignores=scripts/*:S101 diff --git a/README.md b/README.md index 8006af76..3bb7a188 100644 --- a/README.md +++ b/README.md @@ -297,36 +297,145 @@ user_data_collection: **Note**: The `run.yaml` configuration is currently an implementation detail. In the future, all configuration will be available directly from the lightspeed-core config. +**Important**: Only MCP servers defined in the `lightspeed-stack.yaml` configuration are available to the agents. Tools configured in the llama-stack `run.yaml` are not accessible to lightspeed-core agents. + #### Configuring MCP Servers -MCP (Model Context Protocol) servers provide tools and capabilities to the AI agents. These are configured in the `mcp_servers` section of your `lightspeed-stack.yaml`: +MCP (Model Context Protocol) servers provide tools and capabilities to the AI agents. These are configured in the `mcp_servers` section of your `lightspeed-stack.yaml`. + +**Basic Configuration Structure:** + +Each MCP server requires two fields: +- `name`: Unique identifier for the MCP server +- `url`: The endpoint where the MCP server is running + +And one optional field: +- `provider_id`: MCP provider identification (defaults to `"model-context-protocol"`) + +**Minimal Example:** ```yaml mcp_servers: - name: "filesystem-tools" - provider_id: "model-context-protocol" - url: "http://localhost:3000" + url: "http://localhost:9000" - name: "git-tools" - provider_id: "model-context-protocol" - url: "http://localhost:3001" - - name: "database-tools" - provider_id: "model-context-protocol" - url: "http://localhost:3002" + url: "http://localhost:9001" ``` -**Important**: Only MCP servers defined in the `lightspeed-stack.yaml` configuration are available to the agents. Tools configured in the llama-stack `run.yaml` are not accessible to lightspeed-core agents. +In addition to the basic configuration above, you can configure authentication headers for your MCP servers to securely communicate with services that require credentials. + +#### Configuring MCP Server Authentication -#### Configuring MCP Headers +Lightspeed Core Stack supports three methods for authenticating with MCP servers, each suited for different use cases: -MCP headers allow you to pass authentication tokens, API keys, or other metadata to MCP servers. These are configured **per request** via the `MCP-HEADERS` HTTP header: +##### 1. Static Tokens from Files (Recommended for Service Credentials) + +Store authentication tokens in secret files and reference them in your configuration. This is ideal for API keys, service tokens, or any credentials that don't change per-user: + +```yaml +mcp_servers: + - name: "api-service" + url: "http://api-service:8080" + authorization_headers: + Authorization: "/var/secrets/api-token" # Path to file containing token + X-API-Key: "/var/secrets/api-key" # Multiple headers supported +``` + +The secret files should contain only the header value (tokens are automatically stripped of whitespace): + +```bash +# /var/secrets/api-token +Bearer sk-abc123def456... + +# /var/secrets/api-key +my-api-key-value +``` + +##### 2. Kubernetes Service Account Tokens (For K8s Deployments) + +Use the special `"kubernetes"` keyword to automatically use the authenticated user's Kubernetes token. This is perfect for MCP servers running in the same Kubernetes cluster: + +```yaml +mcp_servers: + - name: "k8s-internal-service" + url: "http://internal-mcp.default.svc.cluster.local:8080" + authorization_headers: + Authorization: "kubernetes" # Uses user's k8s token from request auth +``` + +The user's Kubernetes token is extracted from the incoming request's `Authorization` header and forwarded to the MCP server. + +##### 3. Client-Provided Tokens (For Per-User Authentication) + +Use the special `"client"` keyword to allow clients to provide custom tokens per-request. This enables user-specific authentication: + +```yaml +mcp_servers: + - name: "user-specific-service" + url: "http://user-service:8080" + authorization_headers: + Authorization: "client" # Token provided via MCP-HEADERS + X-User-Token: "client" # Multiple client headers supported +``` + +Clients then provide tokens via the `MCP-HEADERS` HTTP header: ```bash curl -X POST "http://localhost:8080/v1/query" \ -H "Content-Type: application/json" \ - -H "MCP-HEADERS: {\"filesystem-tools\": {\"Authorization\": \"Bearer token123\"}}" \ - -d '{"query": "List files in /tmp"}' + -H "MCP-HEADERS: {\"user-specific-service\": {\"Authorization\": \"Bearer user-token-123\", \"X-User-Token\": \"custom-value\"}}" \ + -d '{"query": "Get my data"}' +``` + +**Note**: `MCP-HEADERS` is an **HTTP request header** containing a JSON-encoded dictionary. The dictionary is keyed by **server name** (not URL), matching the `name` field in your MCP server configuration. Each server name maps to another dictionary containing the HTTP headers to forward to that specific MCP server. + +**Structure**: `MCP-HEADERS: {"": {"": "", ...}, ...}` + +##### Combining Authentication Methods + +You can mix and match authentication methods across different MCP servers, and even combine multiple methods for a single server: + +```yaml +mcp_servers: + # Static credentials for public API + - name: "weather-api" + url: "http://weather-api:8080" + authorization_headers: + X-API-Key: "/var/secrets/weather-api-key" + + # Kubernetes auth for internal services + - name: "internal-db" + url: "http://db-mcp.cluster.local:8080" + authorization_headers: + Authorization: "kubernetes" + + # Mixed: static API key + per-user token + - name: "multi-tenant-service" + url: "http://multi-tenant:8080" + authorization_headers: + X-Service-Key: "/var/secrets/service-key" # Static service credential + Authorization: "client" # User-specific token ``` +##### Authentication Method Comparison + +| Method | Use Case | Configuration | Token Scope | Example | +|--------|----------|---------------|-------------|---------| +| **Static File** | Service tokens, API keys | File path in config | Global (all users) | `"/var/secrets/token"` | +| **Kubernetes** | K8s service accounts | `"kubernetes"` keyword | Per-user (from auth) | `"kubernetes"` | +| **Client** | User-specific tokens | `"client"` keyword + HTTP header | Per-request | `"client"` | + +##### Important: Automatic Server Skipping + +**If an MCP server has `authorization_headers` configured but the required tokens cannot be resolved at runtime, the server will be automatically skipped for that request.** This prevents failed authentication attempts to MCP servers. + +**Examples:** +- A server with `Authorization: "kubernetes"` will be skipped if the user's request doesn't include a Kubernetes token +- A server with `Authorization: "client"` will be skipped if no `MCP-HEADERS` are provided in the request +- A server with multiple headers will be skipped if **any** required header cannot be resolved + +Skipped servers are logged as warnings. Check Lightspeed Core logs to see which servers were skipped and why. + ### Llama Stack project and configuration @@ -995,6 +1104,36 @@ The version X.Y.Z indicates: * Y is the minor version (backward-compatible), and * Z is the patch version (backward-compatible bug fix). +# Development Tools + +Lightspeed Core Stack includes development utilities to help with local testing and debugging. These tools are located in the `dev-tools/` directory. + +## MCP Mock Server + +A lightweight mock MCP server for testing MCP integrations locally without requiring real MCP infrastructure. + +**Quick Start:** +```bash +# Start the mock server +python dev-tools/mcp-mock-server/server.py + +# Configure Lightspeed Core Stack to use it +# Add to lightspeed-stack.yaml: +mcp_servers: + - name: "mock-test" + url: "http://localhost:9000" + authorization_headers: + Authorization: "/tmp/test-token" +``` + +**Features:** +- Test authorization header configuration +- Debug MCP connectivity issues +- Inspect captured headers via debug endpoints +- No external dependencies (pure Python stdlib) + +For detailed usage instructions, see [`dev-tools/mcp-mock-server/README.md`](dev-tools/mcp-mock-server/README.md). + # Konflux The official image of Lightspeed Core Stack is built on [Konflux](https://konflux-ui.apps.kflux-prd-rh02.0fk9.p1.openshiftapps.com/ns/lightspeed-core-tenant/applications/lightspeed-stack). diff --git a/dev-tools/MANUAL_TESTING.md b/dev-tools/MANUAL_TESTING.md new file mode 100644 index 00000000..2301c3d1 --- /dev/null +++ b/dev-tools/MANUAL_TESTING.md @@ -0,0 +1,187 @@ +# Manual Testing Guide for MCP Authorization + +This guide walks through testing the MCP server authorization feature using the mock MCP server. + +## Prerequisites + +### 1. Create Test Secret File + +```bash +echo "Bearer test-secret-token-12345" > /tmp/lightspeed-mcp-test-token +``` + +### 2. Start Mock MCP Server + +```bash +# Terminal 1: Start mock MCP server on HTTP port 9000 +python3 dev-tools/mcp-mock-server/server.py 9000 +``` + +Verify the server starts and shows the HTTP endpoint. + +### 3. Install Library Mode Dependencies + +The test configuration uses Llama Stack in library mode, which requires additional dependencies: + +```bash +uv pip install emoji langdetect aiosqlite pythainlp asyncpg nltk 'mcp>=1.23.0' matplotlib 'sqlalchemy[asyncio]' chardet scikit-learn faiss-cpu pillow 'datasets>=4.0.0' psycopg2-binary pandas pypdf pymongo redis tree_sitter requests +``` + +**Note:** These dependencies are needed for Llama Stack's inline providers (agents, vector_io, eval, etc.) to work in library mode. This is a one-time installation. + +### 4. Start Lightspeed Core + +```bash +# Terminal 2: Start Lightspeed Core with test config (Llama Stack runs as library) +# Make sure OPENAI_API_KEY is set first! +export OPENAI_API_KEY="your-api-key-here" +uv run src/lightspeed_stack.py --config dev-tools/test-configs/mcp-mock-test.yaml +``` + +Wait for Lightspeed Core to start (you should see "Application startup complete"). + +**Note:** The test configuration uses Llama Stack in library mode with a dedicated test config (`dev-tools/test-configs/llama-stack-mcp-test.yaml`), so you don't need to start it separately! + +--- + +## Test: All Three Authorization Types in One Request + +The test configuration defines **3 MCP servers**, which means **every query will contact all 3 servers** to discover their tools. This allows us to test all three authorization types in a single request. + +### Step 1: Make a Query Request + +```bash +# Terminal 3: Make test query with all required headers +curl -X POST http://localhost:8080/v1/streaming_query \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer my-k8s-token" \ + -H 'MCP-HEADERS: {"mock-client-auth": {"Authorization": "Bearer my-client-token"}}' \ + -d '{"query": "Test all MCP auth types"}' +``` + +
+Optional: Using Real Tokens (for production testing) + +If you want to test with actual tokens instead of mock values: + +```bash +# Extract real Kubernetes token (requires oc or kubectl) +K8S_TOKEN=$(oc whoami -t 2>/dev/null || kubectl get secret -o jsonpath='{.data.token}' | base64 -d) + +# Set your client token +CLIENT_TOKEN="your-actual-client-token" + +# Make request with real tokens +curl -X POST http://localhost:8080/v1/streaming_query \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer ${K8S_TOKEN}" \ + -H "MCP-HEADERS: {\"mock-client-auth\": {\"Authorization\": \"Bearer ${CLIENT_TOKEN}\"}}" \ + -d '{"query": "Test with real tokens"}' +``` + +**Note:** The mock MCP server doesn't validate tokens, so the simple example above is sufficient for local testing. + +
+ +**What This Tests:** +- **`mock-file-auth`**: Uses static token from `/tmp/lightspeed-mcp-test-token` +- **`mock-k8s-auth`**: Forwards the Kubernetes token from your `Authorization` header +- **`mock-client-auth`**: Uses the client-provided token from `MCP-HEADERS` + +### Step 2: Verify All Three Auth Types Worked + +Check the mock server terminal output. You should see **6 requests total** (2 per server: initialize + list_tools): + +1. **First pair (mock-file-auth)**: + - `Authorization: Bearer test-secret-token-12345` + +2. **Second pair (mock-k8s-auth)**: + - `Authorization: Bearer my-k8s-token` + +3. **Third pair (mock-client-auth)**: + - `Authorization: Bearer my-client-token` + +Or check the debug endpoint: + +```bash +curl http://localhost:9000/debug/requests | jq +``` + +### Expected Result + +The mock server should return unique tool names for each auth type: +- `mock_tool_file` - from `mock-file-auth` +- `mock_tool_k8s` - from `mock-k8s-auth` +- `mock_tool_client` - from `mock-client-auth` + +Check the Lightspeed Core logs, you should see: +```text +DEBUG Configured 3 MCP tools: ['mock-file-auth', 'mock-k8s-auth', 'mock-client-auth'] +``` + +### Step 3: View Request History + +```bash +curl http://localhost:9000/debug/requests | jq +``` + +This will show all 6 requests with their respective Authorization headers. + +--- + +## Additional Test Scenarios + +### Test Without Client Headers + +Test what happens when you don't provide `MCP-HEADERS`: + +```bash +curl -X POST http://localhost:8080/v1/streaming_query \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer my-k8s-token" \ + -d '{"query": "Test without client headers"}' +``` + +**Expected Result:** +- `mock-file-auth`: ✅ Works (static token) +- `mock-k8s-auth`: ✅ Works (uses your k8s token) +- `mock-client-auth`: ⚠️ **Skipped** (required auth header not available - see warning in logs) + +## Cleanup + +Stop all services (Ctrl+C in each terminal): +1. Terminal 1: Mock MCP server +2. Terminal 2: Lightspeed Core (includes Llama Stack) + +Remove test secret: +```bash +rm /tmp/lightspeed-mcp-test-token +``` + +## Troubleshooting + +### Missing OPENAI_API_KEY +- Error: "API key not found" or authentication errors +- Solution: Set the environment variable before starting Lightspeed Core + ```bash + export OPENAI_API_KEY="your-api-key-here" + ``` + +### Lightspeed Core startup errors +- Check that `dev-tools/test-configs/llama-stack-mcp-test.yaml` exists and is valid +- Verify OPENAI_API_KEY is set +- Check the logs for specific error messages + +### Mock server not receiving requests +- Verify mock server is running: `curl http://localhost:9000/debug/headers` +- Check Lightspeed Core logs for errors +- Ensure config file path is correct + +### Headers not captured +- Check mock server output for incoming requests +- Verify the secret file exists and is readable +- Check that the MCP server name matches in config and MCP-HEADERS + +### Connection refused +- Ensure all services are running on expected ports (9000, 8080) +- Check firewall settings diff --git a/dev-tools/README.md b/dev-tools/README.md new file mode 100644 index 00000000..173a1c8d --- /dev/null +++ b/dev-tools/README.md @@ -0,0 +1,32 @@ +# Development Tools + +This directory contains utilities and tools for local development and testing of Lightspeed Core Stack. + +## Available Tools + +### MCP Mock Server + +A lightweight mock Model Context Protocol (MCP) server for testing MCP integrations and authorization headers locally. + +**Location:** `dev-tools/mcp-mock-server/` + +**Use Cases:** +- Test MCP server authentication headers without real MCP infrastructure +- Debug authorization header configuration +- Local development of MCP-related features +- Validate MCP server connectivity + +See [`mcp-mock-server/README.md`](mcp-mock-server/README.md) for usage instructions. + +## Testing MCP Integration with Lightspeed Core + +For comprehensive step-by-step instructions on manually testing MCP authorization, see [`MANUAL_TESTING.md`](MANUAL_TESTING.md). + +## Adding New Tools + +When adding new development tools to this directory: +1. Create a subdirectory for the tool +2. Include a README.md explaining what it does and how to use it +3. Update this file to list the new tool +4. Keep tools self-contained with their own dependencies (if any) + diff --git a/dev-tools/mcp-mock-server/README.md b/dev-tools/mcp-mock-server/README.md new file mode 100644 index 00000000..4d112a03 --- /dev/null +++ b/dev-tools/mcp-mock-server/README.md @@ -0,0 +1,277 @@ +# MCP Mock Server + +A lightweight mock Model Context Protocol (MCP) server for local development and testing. + +## Purpose + +This mock server helps developers: +- Test MCP server authentication without real infrastructure +- Verify authorization headers are correctly sent from Lightspeed Core Stack +- Debug MCP configuration issues locally +- Develop and test MCP-related features +- Test both HTTP and HTTPS connections + +**⚠️ Testing Only:** This server is single-threaded and handles requests sequentially. It is designed purely for development and testing purposes, not for production or high-load scenarios. + +## Features + +- ✅ **Pure Python** - No external dependencies (uses stdlib only) +- ✅ **HTTP & HTTPS** - Runs both protocols simultaneously for comprehensive testing +- ✅ **Header Capture** - Captures and displays all request headers +- ✅ **Debug Endpoints** - Inspect captured headers and request history +- ✅ **MCP Protocol** - Implements basic MCP endpoints for testing +- ✅ **Request Logging** - Tracks recent requests with timestamps +- ✅ **Self-Signed Certs** - Auto-generates certificates for HTTPS testing + +## Quick Start + +### 1. Start the Mock Server + +```bash +# Default ports (HTTP: 3000, HTTPS: 3001) +python dev-tools/mcp-mock-server/server.py + +# Custom ports (HTTP: 8080, HTTPS: 8081) +python dev-tools/mcp-mock-server/server.py 8080 +``` + +You should see: +```text +====================================================================== +MCP Mock Server starting with HTTP and HTTPS +====================================================================== +HTTP: http://localhost:3000 +HTTPS: https://localhost:3001 +====================================================================== +Debug endpoints: + • /debug/headers - View captured headers + • /debug/requests - View request log +MCP endpoint: + • POST /mcp/v1/list_tools +====================================================================== +Note: HTTPS uses a self-signed certificate (for testing only) +``` + +**Note:** The server will automatically generate a self-signed certificate in `dev-tools/mcp-mock-server/.certs/` on first run. + +### 2. Configure Lightspeed Core Stack + +Create a test secret file: +```bash +echo "Bearer test-secret-token-123" > /tmp/mcp-test-token +``` + +Add MCP server to your `lightspeed-stack.yaml`: + +**For HTTP testing:** +```yaml +mcp_servers: + - name: "mock-mcp-test-http" + provider_id: "model-context-protocol" + url: "http://localhost:3000" + authorization_headers: + Authorization: "/tmp/mcp-test-token" +``` + +**For HTTPS testing:** +```yaml +mcp_servers: + - name: "mock-mcp-test-https" + provider_id: "model-context-protocol" + url: "https://localhost:3001" + authorization_headers: + Authorization: "/tmp/mcp-test-token" +``` + +**Note:** For HTTPS with self-signed certificates, you may need to disable SSL verification in your test environment. + +### 3. Test It + +Start Lightspeed Core Stack and make a query: +```bash +# In another terminal +uv run make run + +# Make a query +curl -X POST http://localhost:8080/v1/query \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer user-token" \ + -d '{"query": "Test MCP tools"}' +``` + +### 4. Verify Headers Were Sent + +Check the mock server terminal output or visit the debug endpoints: + +**For HTTP:** +```bash +curl http://localhost:3000/debug/headers +``` + +**For HTTPS (with self-signed cert warning):** +```bash +curl -k https://localhost:3001/debug/headers +``` + +You should see all headers that were sent from the request. + +## Debug Endpoints + +Both HTTP and HTTPS servers expose the same debug endpoints. + +### View All Captured Headers + +**HTTP:** +```bash +curl http://localhost:3000/debug/headers +``` + +**HTTPS:** +```bash +curl -k https://localhost:3001/debug/headers +``` + +Response: +```json +{ + "last_headers": { + "Authorization": "Bearer test-secret-token-123", + "Host": "localhost:3000", + "User-Agent": "curl/7.64.1", + "Accept": "*/*", + "Content-Type": "application/json" + }, + "request_count": 5 +} +``` + +### View Request History + +**HTTP:** +```bash +curl http://localhost:3000/debug/requests +``` + +**HTTPS:** +```bash +curl -k https://localhost:3001/debug/requests +``` + +Response: +```json +[ + { + "timestamp": "2026-01-08T10:30:45.123456", + "method": "POST", + "path": "/mcp/v1/list_tools", + "headers": { + "Authorization": "Bearer test-secret-token-123" + } + } +] +``` + +## Testing Different Authentication Methods + +### Static Token from File +```yaml +mcp_servers: + - name: "file-auth-test" + url: "http://localhost:3000" + authorization_headers: + Authorization: "/tmp/test-token" + X-API-Key: "/tmp/api-key" +``` + +### Kubernetes Token +```yaml +mcp_servers: + - name: "k8s-auth-test" + url: "http://localhost:3000" + authorization_headers: + Authorization: "kubernetes" +``` + +The mock server will receive the user's Kubernetes token from the request. + +### Client-Provided Token +```yaml +mcp_servers: + - name: "client-auth-test" + url: "http://localhost:3000" + authorization_headers: + Authorization: "client" +``` + +Send request with `MCP-HEADERS`: +```bash +curl -X POST http://localhost:8080/v1/query \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer user-k8s-token" \ + -H 'MCP-HEADERS: {"client-auth-test": {"Authorization": "Bearer client-custom-token"}}' \ + -d '{"query": "Test"}' +``` + +## Usage in E2E Tests + +You can use this mock server in automated e2e tests: + +```bash +# Start mock server in background +python dev-tools/mcp-mock-server/server.py 3000 & +MCP_PID=$! + +# Run tests +make test-e2e + +# Cleanup +kill $MCP_PID +``` + +## Running Tests + +The mock server includes its own pytest test suite to verify functionality: + +```bash +# Run tests with pytest +uv run pytest dev-tools/mcp-mock-server/test_mock_mcp_server.py -v + +# Run tests without verbose output +uv run pytest dev-tools/mcp-mock-server/test_mock_mcp_server.py +``` + +The test suite automatically: +- Starts the mock server on ports 9000/9001 +- Tests HTTP and HTTPS endpoints +- Verifies header capture functionality +- Tests request logging +- Cleans up the server after tests complete + +## Troubleshooting + +### Mock server not receiving requests +- Check that Lightspeed Core Stack is configured with the correct URL +- Verify the mock server is running on the expected port +- Check firewall/network settings + +### Headers not captured +- Ensure the header name matches what's configured +- Check mock server logs for incoming requests +- Use `/debug/requests` endpoint to see all recent requests + +### Port already in use +```bash +# Use a different port +python dev-tools/mcp-mock-server/server.py 8080 +``` + +## Limitations + +This is a **development/testing tool only**: +- ❌ Not for production use +- ❌ No authentication/security +- ❌ Limited MCP protocol implementation +- ❌ Single-threaded (one request at a time) + +For production, use real MCP servers. + diff --git a/dev-tools/mcp-mock-server/server.py b/dev-tools/mcp-mock-server/server.py new file mode 100644 index 00000000..c2604e18 --- /dev/null +++ b/dev-tools/mcp-mock-server/server.py @@ -0,0 +1,330 @@ +#!/usr/bin/env python3 +"""Minimal MCP mock server for testing authorization headers. + +This is a simple HTTP/HTTPS server that implements basic MCP protocol endpoints +for testing purposes. It captures and logs authorization headers, making it +useful for validating that Lightspeed Core Stack correctly sends auth headers +to MCP servers. + +The server runs both HTTP and HTTPS simultaneously on consecutive ports. + +Usage: + python server.py [http_port] + +Example: + python server.py 3000 # HTTP on 3000, HTTPS on 3001 +""" + +import json +import ssl +import subprocess +import sys +import threading +from http.server import HTTPServer, BaseHTTPRequestHandler +from datetime import datetime +from pathlib import Path +from typing import Any + + +# Global storage for captured headers (last request) +last_headers: dict[str, str] = {} +request_log: list = [] + + +class MCPMockHandler(BaseHTTPRequestHandler): + """HTTP request handler for mock MCP server.""" + + def log_message(self, format: str, *args: Any) -> None: + """Log requests with timestamp.""" # pylint: disable=redefined-builtin + timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + print(f"[{timestamp}] {format % args}") + + def _capture_headers(self) -> None: + """Capture all headers from the request.""" + last_headers.clear() + + # Capture all headers for debugging + for header_name, value in self.headers.items(): + last_headers[header_name] = value + + # Log the request + request_log.append( + { + "timestamp": datetime.now().isoformat(), + "method": self.command, + "path": self.path, + "headers": dict(last_headers), + } + ) + + # Keep only last 10 requests + if len(request_log) > 10: + request_log.pop(0) + + def do_POST(self) -> None: # pylint: disable=invalid-name + """Handle POST requests (MCP protocol endpoints).""" + self._capture_headers() + + # Read request body to get JSON-RPC request + content_length = int(self.headers.get("Content-Length", 0)) + body = self.rfile.read(content_length) if content_length > 0 else b"{}" + + try: + request_data = json.loads(body.decode("utf-8")) + request_id = request_data.get("id", 1) + method = request_data.get("method", "unknown") + except (json.JSONDecodeError, UnicodeDecodeError): + request_id = 1 + method = "unknown" + + # Determine tool name based on authorization header to avoid collisions + auth_header = self.headers.get("Authorization", "") + + # Match based on token content + match auth_header: + case _ if "test-secret-token" in auth_header: + tool_name = "mock_tool_file" + tool_desc = "Mock tool with file-based auth" + case _ if "my-k8s-token" in auth_header: + tool_name = "mock_tool_k8s" + tool_desc = "Mock tool with Kubernetes token" + case _ if "my-client-token" in auth_header: + tool_name = "mock_tool_client" + tool_desc = "Mock tool with client-provided token" + case _: + # No auth header or unrecognized token + tool_name = "mock_tool_no_auth" + tool_desc = "Mock tool with no authorization" + + # Handle MCP protocol methods + if method == "initialize": + # Return MCP initialize response + response = { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "protocolVersion": "2024-11-05", + "capabilities": { + "tools": {}, + }, + "serverInfo": { + "name": "mock-mcp-server", + "version": "1.0.0", + }, + }, + } + elif method == "tools/list": + # Return list of tools with unique name + response = { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "tools": [ + { + "name": tool_name, + "description": tool_desc, + "inputSchema": { + "type": "object", + "properties": { + "message": { + "type": "string", + "description": "Test message", + } + }, + }, + } + ] + }, + } + else: + # Generic success response for other methods + response = { + "jsonrpc": "2.0", + "id": request_id, + "result": {"status": "ok"}, + } + + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(json.dumps(response).encode()) + + print(f" → Captured headers: {last_headers}") + + def do_GET(self) -> None: # pylint: disable=invalid-name + """Handle GET requests (debug endpoints).""" + # Handle different GET endpoints + match self.path: + case "/debug/headers": + self._send_json_response( + {"last_headers": last_headers, "request_count": len(request_log)} + ) + case "/debug/requests": + self._send_json_response(request_log) + case "/": + self._send_help_page() + case _: + self.send_response(404) + self.end_headers() + + def _send_json_response(self, data: dict | list) -> None: + """Send a JSON response.""" + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(json.dumps(data, indent=2).encode()) + + def _send_help_page(self) -> None: + """Send HTML help page for root endpoint.""" + self.send_response(200) + self.send_header("Content-Type", "text/html") + self.end_headers() + help_html = """ + + MCP Mock Server + +

MCP Mock Server

+

Development mock server for testing MCP integrations.

+

Debug Endpoints:

+ +

MCP Protocol:

+

POST requests to any path with JSON-RPC format:

+
    +
  • {"jsonrpc": "2.0", "id": 1, "method": "initialize"}
  • +
  • {"jsonrpc": "2.0", "id": 1, "method": "tools/list"}
  • +
+ + + """ + self.wfile.write(help_html.encode()) + + +def generate_self_signed_cert(cert_dir: Path) -> tuple[Path, Path]: + """Generate self-signed certificate for HTTPS testing. + + Args: + cert_dir: Directory to store certificate files + + Returns: + Tuple of (cert_file, key_file) paths + """ + cert_file = cert_dir / "cert.pem" + key_file = cert_dir / "key.pem" + + # Only generate if files don't exist + if cert_file.exists() and key_file.exists(): + return cert_file, key_file + + cert_dir.mkdir(parents=True, exist_ok=True) + + # Generate self-signed certificate using openssl + try: + subprocess.run( + [ + "openssl", + "req", + "-x509", + "-newkey", + "rsa:4096", + "-keyout", + str(key_file), + "-out", + str(cert_file), + "-days", + "365", + "-nodes", + "-subj", + "/CN=localhost", + ], + check=True, + capture_output=True, + ) + print(f"Generated self-signed certificate: {cert_file}") + except subprocess.CalledProcessError as e: + print(f"Failed to generate certificate: {e}") + raise + + return cert_file, key_file + + +def run_http_server(port: int, httpd: HTTPServer) -> None: + """Run HTTP server in a thread.""" + print(f"HTTP server started on http://localhost:{port}") + try: + httpd.serve_forever() + except Exception as e: # pylint: disable=broad-except + print(f"HTTP server error: {e}") + + +def run_https_server(port: int, httpd: HTTPServer) -> None: + """Run HTTPS server in a thread.""" + print(f"HTTPS server started on https://localhost:{port}") + try: + httpd.serve_forever() + except Exception as e: # pylint: disable=broad-except + print(f"HTTPS server error: {e}") + + +def main() -> None: + """Start the mock MCP server with both HTTP and HTTPS.""" + http_port = int(sys.argv[1]) if len(sys.argv) > 1 else 3000 + https_port = http_port + 1 + + # Create HTTP server + http_server = HTTPServer(("", http_port), MCPMockHandler) + + # Create HTTPS server with self-signed certificate + https_server = HTTPServer(("", https_port), MCPMockHandler) + + # Generate or load self-signed certificate + script_dir = Path(__file__).parent + cert_dir = script_dir / ".certs" + cert_file, key_file = generate_self_signed_cert(cert_dir) + + # Wrap socket with SSL + context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(cert_file, key_file) + https_server.socket = context.wrap_socket(https_server.socket, server_side=True) + + print("=" * 70) + print("MCP Mock Server starting with HTTP and HTTPS") + print("=" * 70) + print(f"HTTP: http://localhost:{http_port}") + print(f"HTTPS: https://localhost:{https_port}") + print("=" * 70) + print("Debug endpoints:") + print(" • /debug/headers - View captured headers") + print(" • /debug/requests - View request log") + print("MCP endpoint:") + print(" • POST to any path (e.g., / or /mcp/v1/list_tools)") + print("=" * 70) + print("Note: HTTPS uses a self-signed certificate (for testing only)") + print("Press Ctrl+C to stop") + print() + + # Start servers in separate threads + http_thread = threading.Thread( + target=run_http_server, args=(http_port, http_server), daemon=True + ) + https_thread = threading.Thread( + target=run_https_server, args=(https_port, https_server), daemon=True + ) + + http_thread.start() + https_thread.start() + + try: + # Keep main thread alive + http_thread.join() + https_thread.join() + except KeyboardInterrupt: + print("\nShutting down mock servers...") + http_server.shutdown() + https_server.shutdown() + + +if __name__ == "__main__": + main() diff --git a/dev-tools/mcp-mock-server/test_mock_mcp_server.py b/dev-tools/mcp-mock-server/test_mock_mcp_server.py new file mode 100755 index 00000000..0f7d0f5c --- /dev/null +++ b/dev-tools/mcp-mock-server/test_mock_mcp_server.py @@ -0,0 +1,300 @@ +""" +Pytest tests for the MCP Mock Server. + +This test suite verifies the mock server functionality without requiring +the full Lightspeed Stack infrastructure. +""" + +# pylint: disable=redefined-outer-name +# pyright: reportAttributeAccessIssue=false + +import json +import subprocess +import sys +import time +import urllib.request +import urllib.error +from typing import Any + +import pytest + + +@pytest.fixture(scope="module") +def mock_server() -> Any: + """Start mock server for testing and stop it after tests complete.""" + # Using fixed ports for simplicity. For parallel test execution, + # consider using dynamic port allocation (e.g., bind to port 0). + http_port = 9000 + https_port = 9001 + + print(f"\n🚀 Starting mock server on ports {http_port}/{https_port}...") + # Keep stdout/stderr as PIPE to capture errors if startup fails + with subprocess.Popen( + [sys.executable, "dev-tools/mcp-mock-server/server.py", str(http_port)], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) as process: + # Poll server health endpoint instead of blind sleep + max_attempts = 10 + server_url = f"http://localhost:{http_port}/" + + for attempt in range(max_attempts): + if process.poll() is not None: + # Server crashed during startup + _, stderr = process.communicate() + pytest.fail(f"Server failed to start: {stderr.decode('utf-8')}") + + # Try to connect to health endpoint + try: + with urllib.request.urlopen(server_url, timeout=1) as response: + if response.status == 200: + print(f"✅ Server ready after {attempt + 1} attempt(s)") + break + except (urllib.error.URLError, OSError): + # Server not ready yet + time.sleep(0.5) + else: + # Timeout waiting for server + process.terminate() + pytest.fail(f"Server did not respond after {max_attempts} attempts") + + yield { + "process": process, + "http_url": f"http://localhost:{http_port}", + "https_url": f"https://localhost:{https_port}", + } + + # Cleanup: stop server + print("\n🛑 Stopping mock server...") + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + + +def make_request( + url: str, + method: str = "GET", + data: dict[str, Any] | None = None, + headers: dict[str, str] | None = None, + use_https: bool = False, +) -> tuple[int, dict[str, Any] | str]: + """Make HTTP request and return status code and response.""" + req_headers = headers or {} + req_data = None + if data: + req_data = json.dumps(data).encode("utf-8") + req_headers["Content-Type"] = "application/json" + + request = urllib.request.Request( + url, data=req_data, headers=req_headers, method=method + ) + + # For HTTPS with self-signed certs, disable SSL verification + import ssl # pylint: disable=import-outside-toplevel + + context = ( + ssl._create_unverified_context() # pylint: disable=protected-access + if use_https + else None + ) + + try: + with urllib.request.urlopen(request, context=context, timeout=5) as response: + body = response.read().decode("utf-8") + try: + return response.status, json.loads(body) + except json.JSONDecodeError: + return response.status, body + + except urllib.error.HTTPError as e: + body = e.read().decode("utf-8") + try: + return e.code, json.loads(body) + except json.JSONDecodeError: + return e.code, body + except Exception as e: # pylint: disable=broad-except + pytest.fail(f"Request failed: {e}") + return 500, "" # Never reached, but makes pylint happy + + +def test_http_mcp_list_tools(mock_server: Any) -> None: + """Test the MCP list_tools endpoint over HTTP.""" + status, response = make_request( + f"{mock_server['http_url']}/", + method="POST", + data={"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, + headers={"Authorization": "Bearer test-token-123"}, + ) + + assert status == 200, f"Expected 200 OK, got {status}" + assert isinstance(response, dict), f"Expected dict, got {type(response)}" + assert "result" in response, "Response should contain 'result' key (JSON-RPC)" + assert "tools" in response["result"], "Result should contain 'tools' key" + assert isinstance(response["result"]["tools"], list), "Tools should be a list" + assert len(response["result"]["tools"]) == 1, "Should have 1 mock tool" + # Tool name varies based on auth header + assert response["result"]["tools"][0]["name"].startswith( + "mock_tool" + ), "Tool name should start with 'mock_tool'" + + +def test_https_mcp_list_tools(mock_server: Any) -> None: + """Test the MCP list_tools endpoint over HTTPS.""" + status, response = make_request( + f"{mock_server['https_url']}/", + method="POST", + data={"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, + headers={"Authorization": "Bearer test-https-token"}, + use_https=True, + ) + + assert status == 200, f"Expected 200 OK over HTTPS, got {status}" + assert isinstance(response, dict), f"Expected dict, got {type(response)}" + assert "result" in response, "Response should contain 'result' key (JSON-RPC)" + assert "tools" in response["result"], "Result should contain 'tools' key" + + +def test_debug_headers_endpoint(mock_server: Any) -> None: + """Test the debug headers endpoint captures authorization headers.""" + # First, make a request with custom headers + make_request( + f"{mock_server['http_url']}/mcp/v1/list_tools", + method="POST", + headers={ + "Authorization": "Bearer debug-test-token", + "X-Custom-Header": "custom-value-123", + }, + ) + + # Now check if headers were captured + status, response = make_request(f"{mock_server['http_url']}/debug/headers") + + assert status == 200, f"Expected 200 OK, got {status}" + assert isinstance(response, dict), f"Expected dict, got {type(response)}" + + last_headers = response.get("last_headers", {}) + assert "Authorization" in last_headers, "Authorization header not captured" + assert ( + last_headers["Authorization"] == "Bearer debug-test-token" + ), f"Wrong Authorization value: {last_headers.get('Authorization')}" + assert "X-Custom-Header" in last_headers, "Custom header not captured" + assert ( + last_headers["X-Custom-Header"] == "custom-value-123" + ), f"Wrong custom header value: {last_headers.get('X-Custom-Header')}" + assert "request_count" in response, "request_count not in response" + + +def test_debug_requests_endpoint(mock_server: Any) -> None: + """Test the debug requests endpoint logs request history.""" + # Make a request + make_request( + f"{mock_server['http_url']}/mcp/v1/list_tools", + method="POST", + headers={"Authorization": "Bearer request-log-test"}, + ) + + # Check request log + status, response = make_request(f"{mock_server['http_url']}/debug/requests") + + assert status == 200, f"Expected 200 OK, got {status}" + assert isinstance(response, list), f"Expected list, got {type(response)}" + assert len(response) > 0, "No requests logged" + + # Type narrowing for the last request + last_request: dict[str, Any] = response[-1] # type: ignore[assignment,index] + assert "timestamp" in last_request, "Request missing timestamp" + assert "method" in last_request, "Request missing method" + assert ( + last_request["method"] == "POST" + ), f"Wrong method: {last_request.get('method')}" + assert "path" in last_request, "Request missing path" + assert ( + "/mcp/v1/list_tools" in last_request["path"] + ), f"Wrong path: {last_request.get('path')}" + assert "headers" in last_request, "Request missing headers" + assert isinstance(last_request["headers"], dict), "Headers should be dict" + + +def test_multiple_authorization_headers(mock_server: Any) -> None: + """Test capturing multiple authorization headers simultaneously.""" + make_request( + f"{mock_server['http_url']}/mcp/v1/list_tools", + method="POST", + headers={ + "Authorization": "Bearer multi-test-token", + "X-Custom-Auth": "custom-auth-value", + }, + ) + + # Check headers were captured + status, response = make_request(f"{mock_server['http_url']}/debug/headers") + + assert status == 200, f"Expected 200 OK, got {status}" + assert isinstance(response, dict), f"Expected dict, got {type(response)}" + + last_headers = response.get("last_headers", {}) + assert "Authorization" in last_headers, "Authorization header not captured" + assert ( + last_headers["Authorization"] == "Bearer multi-test-token" + ), f"Wrong Authorization: {last_headers.get('Authorization')}" + assert "X-Custom-Auth" in last_headers, "X-Custom-Auth not captured" + assert ( + last_headers["X-Custom-Auth"] == "custom-auth-value" + ), f"Wrong X-Custom-Auth: {last_headers.get('X-Custom-Auth')}" + # Verify multiple headers captured + assert ( + len(last_headers) >= 3 + ), f"Should capture multiple headers, got {len(last_headers)}" + + +def test_all_headers_captured(mock_server: Any) -> None: + """Test that all request headers are captured, not just predefined ones.""" + make_request( + f"{mock_server['http_url']}/mcp/v1/list_tools", + method="POST", + headers={ + "Authorization": "Bearer test", + "X-Random-Header-Name": "random-value", + "X-Another-One": "another-value", + }, + ) + + status, response = make_request(f"{mock_server['http_url']}/debug/headers") + + assert status == 200 + assert isinstance(response, dict), f"Expected dict, got {type(response)}" + last_headers = response.get("last_headers", {}) + + # All custom headers should be captured + assert "X-Random-Header-Name" in last_headers, "Random header not captured" + assert "X-Another-One" in last_headers, "Another header not captured" + + +def test_request_count_increments(mock_server: Any) -> None: + """Test that request count increments with each request. + + Note: The mock server only logs POST requests in request_log, + so GET requests to /debug/headers do not increment the count. + """ + # Get initial count (this GET is not logged) + _, initial_response = make_request(f"{mock_server['http_url']}/debug/headers") + assert isinstance(initial_response, dict), "Expected dict response" + initial_count = initial_response.get("request_count", 0) + + # Make a POST request (this will be logged) + make_request( + f"{mock_server['http_url']}/mcp/v1/list_tools", + method="POST", + ) + + # Check count increased by exactly 1 (only POST requests are logged) + _, final_response = make_request(f"{mock_server['http_url']}/debug/headers") + assert isinstance(final_response, dict), "Expected dict response" + final_count = final_response.get("request_count", 0) + + # The count should have increased by 1 (only the POST to /mcp/v1/list_tools) + assert ( + final_count == initial_count + 1 + ), f"Expected count to increase by 1, but went from {initial_count} to {final_count}" diff --git a/dev-tools/test-configs/llama-stack-mcp-test.yaml b/dev-tools/test-configs/llama-stack-mcp-test.yaml new file mode 100644 index 00000000..24d65342 --- /dev/null +++ b/dev-tools/test-configs/llama-stack-mcp-test.yaml @@ -0,0 +1,98 @@ +# Llama Stack configuration for MCP testing +# Minimal configuration with only APIs needed for MCP server testing +# +# Prerequisites: +# - Set OPENAI_API_KEY environment variable +# +version: 2 +image_name: mcp-test-config + +apis: +- agents +- files +- inference +- safety +- tool_runtime +- vector_io + +models: +- model_id: gpt-4o-mini + provider_id: openai + model_type: llm + provider_model_id: gpt-4o-mini + +shields: +- shield_id: default-shield + provider_id: llama-guard + provider_shield_id: llama-guard + +providers: + agents: + - provider_id: meta-reference + provider_type: inline::meta-reference + config: + persistence: + agent_state: + namespace: agents_state + backend: kv_default + responses: + table_name: agents_responses + backend: sql_default + + files: + - provider_id: meta-reference-files + provider_type: inline::localfs + config: + metadata_store: + table_name: files_metadata + backend: sql_default + storage_dir: /tmp/llama_mcp_test_files + + inference: + - provider_id: openai + provider_type: remote::openai + config: + api_key: ${env.OPENAI_API_KEY} + + safety: + - provider_id: llama-guard + provider_type: inline::llama-guard + config: + excluded_categories: [] + + tool_runtime: + - provider_id: rag-runtime + provider_type: inline::rag-runtime + config: {} + - provider_id: model-context-protocol + provider_type: remote::model-context-protocol + config: {} + + vector_io: + - provider_id: faiss + provider_type: inline::faiss + config: + persistence: + namespace: faiss_store + backend: kv_default + +storage: + backends: + kv_default: + type: kv_sqlite + db_path: /tmp/llama_mcp_test_kv_store.db + sql_default: + type: sql_sqlite + db_path: /tmp/llama_mcp_test_sql_store.db + stores: + metadata: + namespace: registry + backend: kv_default + inference: + table_name: inference_store + backend: sql_default + max_write_queue_size: 10000 + num_writers: 4 + conversations: + table_name: openai_conversations + backend: sql_default \ No newline at end of file diff --git a/dev-tools/test-configs/mcp-mock-test.yaml b/dev-tools/test-configs/mcp-mock-test.yaml new file mode 100644 index 00000000..4760e03c --- /dev/null +++ b/dev-tools/test-configs/mcp-mock-test.yaml @@ -0,0 +1,38 @@ +name: Lightspeed Core Service - MCP Mock Server Test +service: + host: localhost + port: 8080 + auth_enabled: false + workers: 1 + color_log: true + access_log: true +llama_stack: + use_as_library_client: true + library_client_config_path: "dev-tools/test-configs/llama-stack-mcp-test.yaml" +user_data_collection: + feedback_enabled: false + transcripts_enabled: false +authentication: + module: "noop-with-token" +inference: + default_model: "gpt-4o-mini" + default_provider: "openai" +mcp_servers: + # Test 1: Static file-based authentication (HTTP) + - name: "mock-file-auth" + provider_id: "model-context-protocol" + url: "http://localhost:9000" + authorization_headers: + Authorization: "/tmp/lightspeed-mcp-test-token" + # Test 2: Kubernetes token forwarding (HTTP) + - name: "mock-k8s-auth" + provider_id: "model-context-protocol" + url: "http://localhost:9000" + authorization_headers: + Authorization: "kubernetes" + # Test 3: Client-provided token (HTTP - simplified for testing) + - name: "mock-client-auth" + provider_id: "model-context-protocol" + url: "http://localhost:9000" + authorization_headers: + Authorization: "client" diff --git a/examples/lightspeed-stack-mcp-servers.yaml b/examples/lightspeed-stack-mcp-servers.yaml index d44f20e7..e71d0eb2 100644 --- a/examples/lightspeed-stack-mcp-servers.yaml +++ b/examples/lightspeed-stack-mcp-servers.yaml @@ -27,3 +27,23 @@ mcp_servers: - name: "server3" provider_id: "provider3" url: "http://url.com:3" + # Example with authorization headers using secret file paths + - name: "server-with-secret-auth" + provider_id: "provider4" + url: "http://url.com:4" + authorization_headers: + Authorization: "/var/secrets/auth-token" # Path to file containing auth token + X-API-Key: "/var/secrets/api-key" # Path to file containing API key + # Example with kubernetes token (special case) + - name: "server-with-k8s-auth" + provider_id: "provider5" + url: "http://url.com:5" + authorization_headers: + Authorization: "kubernetes" # Special value to use kubernetes service account token + # Example with client-provided token (special case) + - name: "server-with-client-auth" + provider_id: "provider6" + url: "http://url.com:6" + authorization_headers: + Authorization: "client" # Special value to forward the client's token + timeout: 30 # Optional: timeout in seconds (future Llama Stack feature) \ No newline at end of file diff --git a/src/app/endpoints/query_v2.py b/src/app/endpoints/query_v2.py index 69422c42..e6b14522 100644 --- a/src/app/endpoints/query_v2.py +++ b/src/app/endpoints/query_v2.py @@ -22,7 +22,7 @@ from authorization.middleware import authorize from configuration import AppConfig, configuration from constants import DEFAULT_RAG_TOOL -from models.config import Action +from models.config import Action, ModelContextProtocolServer from models.requests import QueryRequest from models.responses import ( ForbiddenResponse, @@ -80,7 +80,7 @@ def _build_tool_call_summary( # pylint: disable=too-many-return-statements,too- The OpenAI ``response.output`` array may contain any ``OpenAIResponseOutput`` variant: ``message``, ``function_call``, ``file_search_call``, ``web_search_call``, ``mcp_call``, ``mcp_list_tools``, or ``mcp_approval_request``. The OpenAI Spec supports more types - but as llamastack does not support them yet they are not considered here. + but as llamastack does not support them, yet they are not considered here. """ item_type = getattr(output_item, "type", None) @@ -507,7 +507,8 @@ def parse_referenced_documents_from_responses_api( # Use a set to track unique documents by (doc_url, doc_title) tuple seen_docs: set[tuple[Optional[str], Optional[str]]] = set() - if not response.output: + # Handle None response (e.g., when agent fails) + if response is None or not response.output: return documents for output_item in response.output: @@ -716,9 +717,9 @@ def get_rag_tools(vector_store_ids: list[str]) -> Optional[list[dict[str, Any]]] def get_mcp_tools( - mcp_servers: list, - token: Optional[str] = None, - mcp_headers: Optional[dict[str, dict[str, str]]] = None, + mcp_servers: list[ModelContextProtocolServer], + token: str | None = None, + mcp_headers: dict[str, dict[str, str]] | None = None, ) -> list[dict[str, Any]]: """ Convert MCP servers to tools format for Responses API. @@ -731,9 +732,47 @@ def get_mcp_tools( Returns: list[dict[str, Any]]: List of MCP tool definitions with server details and optional auth headers + + The way it works is we go through all the defined mcp servers and + create a tool definitions for each of them. If MCP server definition + has a non-empty resolved_authorization_headers we create invocation + headers, following the algorithm: + 1. If the header value is 'kubernetes' the header value is a k8s token + 2. If the header value is 'client': + find the value for a given MCP server/header in mcp_headers. + if the value is not found omit this header, otherwise use found value + 3. otherwise use the value from resolved_authorization_headers directly + + This algorithm allows to: + 1. Use static global header values, provided by configuration + 2. Use user specific k8s token, which will work for the majority of kubernetes + based MCP servers + 3. Use user specific tokens (passed by the client) for user specific MCP headers """ + + def _get_token_value(original: str, header: str) -> str | None: + """Convert to header value.""" + match original: + case "kubernetes": + # use k8s token + if token is None or token == "": + return None + return f"Bearer {token}" + case "client": + # use client provided token + if mcp_headers is None: + return None + c_headers = mcp_headers.get(mcp_server.name, None) + if c_headers is None: + return None + return c_headers.get(header, None) + case _: + # use provided + return original + tools = [] for mcp_server in mcp_servers: + # Base tool definition tool_def = { "type": "mcp", "server_label": mcp_server.name, @@ -741,18 +780,31 @@ def get_mcp_tools( "require_approval": "never", } - # Build headers: start with token auth, then merge in per-server headers - if token or mcp_headers: - headers = {} - # Add token-based auth if available - if token: - headers["Authorization"] = f"Bearer {token}" - # Merge in per-server headers (can override Authorization if needed) - server_headers = (mcp_headers or {}).get(mcp_server.url) - if server_headers: - headers.update(server_headers) - tool_def["headers"] = headers + # Build headers + headers = {} + for name, value in mcp_server.resolved_authorization_headers.items(): + # for each defined header + h_value = _get_token_value(value, name) + # only add the header if we got value + if h_value is not None: + headers[name] = h_value + + # Skip server if auth headers were configured but not all could be resolved + if mcp_server.authorization_headers and len(headers) != len( + mcp_server.authorization_headers + ): + logger.warning( + "Skipping MCP server %s: required %d auth headers but only resolved %d", + mcp_server.name, + len(mcp_server.authorization_headers), + len(headers), + ) + continue + if len(headers) > 0: + # add headers to tool definition + tool_def["headers"] = headers # type: ignore[index] + # collect tools info tools.append(tool_def) return tools diff --git a/src/configuration.py b/src/configuration.py index 4c9fd48e..ef19c619 100644 --- a/src/configuration.py +++ b/src/configuration.py @@ -334,7 +334,7 @@ def token_usage_history(self) -> Optional[TokenUsageHistory]: raise LogicError("logic error: configuration is not loaded") if ( self._token_usage_history is None - and self._configuration.quota_handlers.enable_token_history + and self._configuration.quota_handlers.enable_token_history # pylint: disable=no-member ): self._token_usage_history = TokenUsageHistory( self._configuration.quota_handlers diff --git a/src/models/config.py b/src/models/config.py index ce4cf34c..1e8335d5 100644 --- a/src/models/config.py +++ b/src/models/config.py @@ -2,6 +2,7 @@ # pylint: disable=too-many-lines +import logging from pathlib import Path from typing import Optional, Any, Pattern from enum import Enum @@ -21,6 +22,7 @@ PositiveInt, NonNegativeInt, SecretStr, + PrivateAttr, ) from pydantic.dataclasses import dataclass @@ -29,6 +31,9 @@ import constants from utils import checks +from utils.mcp_auth_headers import resolve_authorization_headers + +logger = logging.getLogger(__name__) class ConfigurationBase(BaseModel): @@ -462,6 +467,56 @@ class ModelContextProtocolServer(ConfigurationBase): description="URL of the MCP server", ) + authorization_headers: dict[str, str] = Field( + default_factory=dict, + title="Authorization headers", + description=( + "Headers to send to the MCP server. " + "The map contains the header name and the path to a file containing " + "the header value (secret). " + "There are 2 special cases: " + "1. Usage of the kubernetes token in the header. " + "To specify this use a string 'kubernetes' instead of the file path. " + "2. Usage of the client provided token in the header. " + "To specify this use a string 'client' instead of the file path." + ), + ) + + _resolved_authorization_headers: dict[str, str] = PrivateAttr(default_factory=dict) + + @property + def resolved_authorization_headers(self) -> dict[str, str]: + """Resolved authorization headers (computed from authorization_headers).""" + return self._resolved_authorization_headers + + timeout: Optional[PositiveInt] = Field( + default=None, + title="Request timeout", + description=( + "Timeout in seconds for requests to the MCP server. " + "If not specified, the default timeout from Llama Stack will be used. " + "Note: This field is reserved for future use when Llama Stack adds timeout support." + ), + ) + + @model_validator(mode="after") + def resolve_auth_headers(self) -> Self: + """ + Resolve authorization headers by reading secret files. + + Automatically populates resolved_authorization_headers by reading + secret files specified in authorization_headers. Special values + 'kubernetes' and 'client' are preserved for later substitution. + + Returns: + Self: The model instance with resolved_authorization_headers populated. + """ + if self.authorization_headers: + self._resolved_authorization_headers = resolve_authorization_headers( + self.authorization_headers + ) + return self + class LlamaStackConfiguration(ConfigurationBase): """Llama stack configuration. @@ -1589,7 +1644,46 @@ class Configuration(ConfigurationBase): description="Quota handlers configuration", ) - def dump(self, filename: str = "configuration.json") -> None: + @model_validator(mode="after") + def validate_mcp_auth_headers(self) -> Self: + """ + Validate MCP server authorization headers against authentication module. + + Removes any MCP server with authorization_headers="kubernetes" when the + authentication module is not "k8s". This prevents sending wrong credential + types to MCP servers. + + Returns: + Self: The model instance after validation. + """ + # Get authentication module value (pyright: ignore attribute access on Field) + auth_module = getattr(self.authentication, "module", None) + + # Filter out misconfigured MCP servers + valid_mcp_servers = [] + for mcp_server in self.mcp_servers: + is_valid = True + if mcp_server.authorization_headers: + for value in mcp_server.authorization_headers.values(): + if value.strip() == "kubernetes" and auth_module != "k8s": + logger.warning( + "Removing MCP server '%s': has authorization_headers with " + "value 'kubernetes' but authentication module is '%s' " + "(not 'k8s'). Either change authentication.module to 'k8s' " + "or update the MCP server's authorization_headers to use a " + "file path or 'client'.", + mcp_server.name, + auth_module, + ) + is_valid = False + break + if is_valid: + valid_mcp_servers.append(mcp_server) + + self.mcp_servers = valid_mcp_servers + return self + + def dump(self, filename: str | Path = "configuration.json") -> None: """ Write the current Configuration model to a JSON file. @@ -1598,7 +1692,7 @@ def dump(self, filename: str = "configuration.json") -> None: file exists it will be overwritten. Parameters: - filename (str): Path to the output file (defaults to "configuration.json"). + filename (str | Path): Path to the output file (defaults to "configuration.json"). """ with open(filename, "w", encoding="utf-8") as fout: fout.write(self.model_dump_json(indent=4)) diff --git a/src/utils/mcp_auth_headers.py b/src/utils/mcp_auth_headers.py new file mode 100644 index 00000000..8c5b2452 --- /dev/null +++ b/src/utils/mcp_auth_headers.py @@ -0,0 +1,87 @@ +"""Utilities for resolving MCP server authorization headers.""" + +import logging +from pathlib import Path + +logger = logging.getLogger(__name__) + + +def resolve_authorization_headers( + authorization_headers: dict[str, str], +) -> dict[str, str]: + """ + Resolve authorization headers by reading secret files or preserving special values. + + Parameters: + authorization_headers: Map of header names to secret locations or special keywords. + - If value is "kubernetes": leave is unchanged. We substitute it during request. + - If value is "client": leave it unchanged. . We substitute it during request. + - Otherwise: Treat as file path and read the secret from that file + + Returns: + dict[str, str]: Map of header names to resolved header values or special keywords + + Examples: + >>> # With file paths + >>> resolve_authorization_headers({"Authorization": "/var/secrets/token"}) + {"Authorization": "secret-value-from-file"} + + >>> # With kubernetes special case (kept as-is) + >>> resolve_authorization_headers({"Authorization": "kubernetes"}) + {"Authorization": "kubernetes"} + + >>> # With client special case (kept as-is) + >>> resolve_authorization_headers({"Authorization": "client"}) + {"Authorization": "client"} + """ + resolved: dict[str, str] = {} + + for header_name, value in authorization_headers.items(): + value = value.strip() + try: + if value == "kubernetes": + # Special case: Keep kubernetes keyword for later substitution + resolved[header_name] = "kubernetes" + logger.debug( + "Header %s will use Kubernetes token (resolved at request time)", + header_name, + ) + elif value == "client": + # Special case: Keep client keyword for later substitution + resolved[header_name] = "client" + logger.debug( + "Header %s will use client-provided token (resolved at request time)", + header_name, + ) + else: + # Regular case: Read secret from file path + secret_path = Path(value).expanduser() + if secret_path.exists() and secret_path.is_file(): + secret_value = secret_path.read_text(encoding="utf-8").strip() + if not secret_value: + logger.warning( + "Secret file %s for header %s is empty", + secret_path, + header_name, + ) + else: + resolved[header_name] = secret_value + logger.debug( + "Resolved header %s from secret file %s", + header_name, + secret_path, + ) + else: + logger.warning( + "Secret file not found or not a file: %s for header %s", + secret_path, + header_name, + ) + except Exception: # pylint: disable=broad-except + # Don't log value: it might be a literal token. + logger.exception( + "Failed to resolve authorization header %s (value treated as path)", + header_name, + ) + + return resolved diff --git a/tests/unit/app/endpoints/test_query_v2.py b/tests/unit/app/endpoints/test_query_v2.py index 77c04dba..fd4ece75 100644 --- a/tests/unit/app/endpoints/test_query_v2.py +++ b/tests/unit/app/endpoints/test_query_v2.py @@ -1,6 +1,8 @@ -# pylint: disable=redefined-outer-name, import-error,too-many-locals +# pylint: disable=redefined-outer-name, import-error,too-many-locals,too-many-lines +# pyright: reportCallIssue=false """Unit tests for the /query (v2) REST API endpoint using Responses API.""" +from pathlib import Path from typing import Any import pytest @@ -54,63 +56,163 @@ def test_get_rag_tools() -> None: def test_get_mcp_tools_with_and_without_token() -> None: - """Test get_mcp_tools generates correct tool definitions with and without auth tokens.""" - servers = [ + """Test get_mcp_tools with resolved_authorization_headers.""" + # Servers without authorization headers + servers_no_auth = [ ModelContextProtocolServer(name="fs", url="http://localhost:3000"), ModelContextProtocolServer(name="git", url="https://git.example.com/mcp"), ] - tools_no_token = get_mcp_tools(servers, token=None) - assert len(tools_no_token) == 2 - assert tools_no_token[0]["type"] == "mcp" - assert tools_no_token[0]["server_label"] == "fs" - assert tools_no_token[0]["server_url"] == "http://localhost:3000" - assert "headers" not in tools_no_token[0] - - tools_with_token = get_mcp_tools(servers, token="abc") - assert len(tools_with_token) == 2 - assert tools_with_token[1]["type"] == "mcp" - assert tools_with_token[1]["server_label"] == "git" - assert tools_with_token[1]["server_url"] == "https://git.example.com/mcp" - assert tools_with_token[1]["headers"] == {"Authorization": "Bearer abc"} + tools_no_auth = get_mcp_tools(servers_no_auth, token=None) + assert len(tools_no_auth) == 2 + assert tools_no_auth[0]["type"] == "mcp" + assert tools_no_auth[0]["server_label"] == "fs" + assert tools_no_auth[0]["server_url"] == "http://localhost:3000" + assert "headers" not in tools_no_auth[0] + + # Servers with kubernetes auth + servers_k8s = [ + ModelContextProtocolServer( + name="k8s-server", + url="http://localhost:3000", + authorization_headers={"Authorization": "kubernetes"}, + ), + ] + tools_k8s = get_mcp_tools(servers_k8s, token="user-k8s-token") + assert len(tools_k8s) == 1 + assert tools_k8s[0]["headers"] == {"Authorization": "Bearer user-k8s-token"} def test_get_mcp_tools_with_mcp_headers() -> None: - """Test get_mcp_tools merges token auth and per-server headers correctly.""" + """Test get_mcp_tools with client-provided headers.""" + # Server with client auth servers = [ - ModelContextProtocolServer(name="fs", url="http://localhost:3000"), - ModelContextProtocolServer(name="git", url="https://git.example.com/mcp"), + ModelContextProtocolServer( + name="fs", + url="http://localhost:3000", + authorization_headers={"Authorization": "client", "X-Custom": "client"}, + ), ] - # Test with mcp_headers only (no token) + # Test with mcp_headers provided mcp_headers = { - "http://localhost:3000": {"X-Custom-Header": "value1"}, - "https://git.example.com/mcp": {"X-API-Key": "secret123"}, + "fs": { + "Authorization": "client-provided-token", + "X-Custom": "custom-value", + } } tools = get_mcp_tools(servers, token=None, mcp_headers=mcp_headers) - assert len(tools) == 2 - assert tools[0]["headers"] == {"X-Custom-Header": "value1"} - assert tools[1]["headers"] == {"X-API-Key": "secret123"} - - # Test with both token and mcp_headers (should merge) - tools_merged = get_mcp_tools(servers, token="abc", mcp_headers=mcp_headers) - assert len(tools_merged) == 2 - assert tools_merged[0]["headers"] == { - "Authorization": "Bearer abc", - "X-Custom-Header": "value1", + assert len(tools) == 1 + assert tools[0]["headers"] == { + "Authorization": "client-provided-token", + "X-Custom": "custom-value", } - assert tools_merged[1]["headers"] == { - "Authorization": "Bearer abc", - "X-API-Key": "secret123", + + # Test with mcp_headers=None (server should be skipped since auth is required but unavailable) + tools_no_headers = get_mcp_tools(servers, token=None, mcp_headers=None) + assert len(tools_no_headers) == 0 # Server skipped due to missing required auth + + +def test_get_mcp_tools_with_static_headers(tmp_path: Path) -> None: + """Test get_mcp_tools with static headers from config files.""" + # Create a secret file + secret_file = tmp_path / "token.txt" + secret_file.write_text("static-secret-token") + + servers = [ + ModelContextProtocolServer( + name="server1", + url="http://localhost:3000", + authorization_headers={"Authorization": str(secret_file)}, + ), + ] + + tools = get_mcp_tools(servers, token=None) + assert len(tools) == 1 + assert tools[0]["headers"] == {"Authorization": "static-secret-token"} + + +def test_get_mcp_tools_with_mixed_headers(tmp_path: Path) -> None: + """Test get_mcp_tools with mixed header types.""" + # Create a secret file + secret_file = tmp_path / "api-key.txt" + secret_file.write_text("secret-api-key") + + servers = [ + ModelContextProtocolServer( + name="mixed-server", + url="http://localhost:3000", + authorization_headers={ + "Authorization": "kubernetes", + "X-API-Key": str(secret_file), + "X-Custom": "client", + }, + ), + ] + + mcp_headers = { + "mixed-server": { + "X-Custom": "client-custom-value", + } } - # Test mcp_headers can override Authorization - override_headers = { - "http://localhost:3000": {"Authorization": "Custom auth"}, + tools = get_mcp_tools(servers, token="k8s-token", mcp_headers=mcp_headers) + assert len(tools) == 1 + assert tools[0]["headers"] == { + "Authorization": "Bearer k8s-token", + "X-API-Key": "secret-api-key", + "X-Custom": "client-custom-value", } - tools_override = get_mcp_tools(servers, token="abc", mcp_headers=override_headers) - assert tools_override[0]["headers"] == {"Authorization": "Custom auth"} - assert tools_override[1]["headers"] == {"Authorization": "Bearer abc"} + + +def test_get_mcp_tools_skips_server_with_missing_auth() -> None: + """Test that servers with required but unavailable auth headers are skipped.""" + servers = [ + # Server with kubernetes auth but no token provided + ModelContextProtocolServer( + name="missing-k8s-auth", + url="http://localhost:3001", + authorization_headers={"Authorization": "kubernetes"}, + ), + # Server with client auth but no MCP-HEADERS provided + ModelContextProtocolServer( + name="missing-client-auth", + url="http://localhost:3002", + authorization_headers={"X-Token": "client"}, + ), + # Server with partial auth (2 headers required, only 1 available) + ModelContextProtocolServer( + name="partial-auth", + url="http://localhost:3003", + authorization_headers={ + "Authorization": "kubernetes", + "X-Custom": "client", + }, + ), + ] + + # No token, no mcp_headers + tools = get_mcp_tools(servers, token=None, mcp_headers=None) + # All servers should be skipped + assert len(tools) == 0 + + +def test_get_mcp_tools_includes_server_without_auth() -> None: + """Test that servers without auth config are always included.""" + servers = [ + # Server with no auth requirements + ModelContextProtocolServer( + name="public-server", + url="http://localhost:3000", + authorization_headers={}, + ), + ] + + # Should work even without token or headers + tools = get_mcp_tools(servers, token=None, mcp_headers=None) + assert len(tools) == 1 + assert tools[0]["server_label"] == "public-server" + assert "headers" not in tools[0] @pytest.mark.asyncio @@ -181,7 +283,11 @@ async def test_retrieve_response_builds_rag_and_mcp_tools( # pylint: disable=to mocker.patch("app.endpoints.query_v2.get_system_prompt", return_value="PROMPT") mock_cfg = mocker.Mock() mock_cfg.mcp_servers = [ - ModelContextProtocolServer(name="fs", url="http://localhost:3000"), + ModelContextProtocolServer( + name="fs", + url="http://localhost:3000", + authorization_headers={"Authorization": "kubernetes"}, + ), ] mocker.patch("app.endpoints.query_v2.configuration", mock_cfg) @@ -554,7 +660,7 @@ def _raise(*_args: Any, **_kwargs: Any) -> Exception: assert exc.value.status_code == status.HTTP_503_SERVICE_UNAVAILABLE detail = exc.value.detail assert isinstance(detail, dict) - assert detail["response"] == "Unable to connect to Llama Stack" + assert detail["response"] == "Unable to connect to Llama Stack" # type: ignore[index] fail_metric.inc.assert_called_once() diff --git a/tests/unit/app/endpoints/test_streaming_query_v2.py b/tests/unit/app/endpoints/test_streaming_query_v2.py index a9220ab8..20a379b7 100644 --- a/tests/unit/app/endpoints/test_streaming_query_v2.py +++ b/tests/unit/app/endpoints/test_streaming_query_v2.py @@ -60,9 +60,14 @@ async def test_retrieve_response_builds_rag_and_mcp_tools( mock_cfg = mocker.Mock() mock_cfg.mcp_servers = [ - ModelContextProtocolServer(name="fs", url="http://localhost:3000"), + ModelContextProtocolServer( + name="fs", + url="http://localhost:3000", + authorization_headers={"Authorization": "kubernetes"}, + ), ] mocker.patch("app.endpoints.streaming_query_v2.configuration", mock_cfg) + mocker.patch("app.endpoints.query_v2.configuration", mock_cfg) qr = QueryRequest(query="hello") await retrieve_response(mock_client, "model-z", qr, token="tok") diff --git a/tests/unit/models/config/test_dump_configuration.py b/tests/unit/models/config/test_dump_configuration.py index 7ed81c72..286cae17 100644 --- a/tests/unit/models/config/test_dump_configuration.py +++ b/tests/unit/models/config/test_dump_configuration.py @@ -1,6 +1,7 @@ """Unit tests checking ability to dump configuration.""" # pylint: disable=too-many-lines +# pyright: reportCallIssue=false import json @@ -248,6 +249,8 @@ def test_dump_configuration_with_one_mcp_server(tmp_path: Path) -> None: "name": "test-server", "url": "http://localhost:8080", "provider_id": "model-context-protocol", + "authorization_headers": {}, + "timeout": None, } ] @@ -305,16 +308,22 @@ def test_dump_configuration_with_more_mcp_servers(tmp_path: Path) -> None: "name": "test-server-1", "provider_id": "model-context-protocol", "url": "http://localhost:8081", + "authorization_headers": {}, + "timeout": None, }, { "name": "test-server-2", "provider_id": "model-context-protocol", "url": "http://localhost:8082", + "authorization_headers": {}, + "timeout": None, }, { "name": "test-server-3", "provider_id": "model-context-protocol", "url": "http://localhost:8083", + "authorization_headers": {}, + "timeout": None, }, ] diff --git a/tests/unit/models/config/test_model_context_protocol_server.py b/tests/unit/models/config/test_model_context_protocol_server.py index fb6664cf..3bc2f7b2 100644 --- a/tests/unit/models/config/test_model_context_protocol_server.py +++ b/tests/unit/models/config/test_model_context_protocol_server.py @@ -1,15 +1,20 @@ """Unit tests for ModelContextProtocolServer model.""" +# pyright: reportCallIssue=false + +from pathlib import Path + import pytest from pydantic import ValidationError -from models.config import ( - ModelContextProtocolServer, - LlamaStackConfiguration, - UserDataCollection, +from models.config import ( # type: ignore[import-not-found] + AuthenticationConfiguration, Configuration, + LlamaStackConfiguration, + ModelContextProtocolServer, ServiceConfiguration, + UserDataCollection, ) @@ -20,6 +25,7 @@ def test_model_context_protocol_server_constructor() -> None: assert mcp.name == "test-server" assert mcp.provider_id == "model-context-protocol" assert mcp.url == "http://localhost:8080" + assert mcp.authorization_headers == {} # Default should be empty dict def test_model_context_protocol_server_custom_provider() -> None: @@ -132,3 +138,168 @@ def test_configuration_multiple_mcp_servers() -> None: assert cfg.mcp_servers[1].name == "server2" assert cfg.mcp_servers[1].provider_id == "custom-provider" assert cfg.mcp_servers[2].name == "server3" + + +def test_model_context_protocol_server_with_authorization_headers( + tmp_path: Path, +) -> None: + """Test ModelContextProtocolServer with authorization headers.""" + auth_file = tmp_path / "auth.txt" + auth_file.write_text("my-secret") + api_key_file = tmp_path / "api_key.txt" + api_key_file.write_text("api-key-secret") + + mcp = ModelContextProtocolServer( + name="auth-server", + url="http://localhost:8080", + authorization_headers={ + "Authorization": str(auth_file), + "X-API-Key": str(api_key_file), + }, + ) + assert mcp is not None + assert mcp.name == "auth-server" + assert mcp.url == "http://localhost:8080" + assert mcp.authorization_headers == { + "Authorization": str(auth_file), + "X-API-Key": str(api_key_file), + } + assert mcp.resolved_authorization_headers == { + "Authorization": "my-secret", + "X-API-Key": "api-key-secret", + } + + +def test_model_context_protocol_server_kubernetes_special_case() -> None: + """Test ModelContextProtocolServer with kubernetes special case.""" + mcp = ModelContextProtocolServer( + name="k8s-server", + url="http://localhost:8080", + authorization_headers={"Authorization": "kubernetes"}, + ) + assert mcp is not None + assert mcp.authorization_headers == {"Authorization": "kubernetes"} + + +def test_model_context_protocol_server_client_special_case() -> None: + """Test ModelContextProtocolServer with client special case.""" + mcp = ModelContextProtocolServer( + name="client-server", + url="http://localhost:8080", + authorization_headers={"Authorization": "client"}, + ) + assert mcp is not None + assert mcp.authorization_headers == {"Authorization": "client"} + + +def test_configuration_mcp_servers_with_mixed_auth_headers(tmp_path: Path) -> None: + """ + Test Configuration with MCP servers having mixed authorization headers. + + Verifies backward compatibility (servers without auth headers) and + new functionality (servers with auth headers) work together. + """ + secret_file = tmp_path / "secret.txt" + secret_file.write_text("my-secret") + + mcp_servers = [ + ModelContextProtocolServer(name="server-no-auth", url="http://localhost:8080"), + ModelContextProtocolServer( + name="server-with-secret", + url="http://localhost:8081", + authorization_headers={"Authorization": str(secret_file)}, + ), + ModelContextProtocolServer( + name="server-with-k8s", + url="http://localhost:8082", + authorization_headers={"Authorization": "kubernetes"}, + ), + ModelContextProtocolServer( + name="server-with-client", + url="http://localhost:8083", + authorization_headers={"Authorization": "client"}, + ), + ] + cfg = Configuration( + name="test_name", + service=ServiceConfiguration(), + llama_stack=LlamaStackConfiguration( + use_as_library_client=True, + library_client_config_path="tests/configuration/run.yaml", + ), + user_data_collection=UserDataCollection( + feedback_enabled=False, feedback_storage=None + ), + mcp_servers=mcp_servers, + authentication=AuthenticationConfiguration(module="k8s"), + customization=None, + ) + assert cfg is not None + assert len(cfg.mcp_servers) == 4 + + # Server without auth headers (backward compatibility) + assert cfg.mcp_servers[0].name == "server-no-auth" + assert cfg.mcp_servers[0].authorization_headers == {} + + # Server with secret reference + assert cfg.mcp_servers[1].name == "server-with-secret" + assert cfg.mcp_servers[1].authorization_headers == { + "Authorization": str(secret_file) + } + assert cfg.mcp_servers[1].resolved_authorization_headers == { + "Authorization": "my-secret" + } + + # Server with kubernetes special case + assert cfg.mcp_servers[2].name == "server-with-k8s" + assert cfg.mcp_servers[2].authorization_headers == {"Authorization": "kubernetes"} + + # Server with client special case + assert cfg.mcp_servers[3].name == "server-with-client" + assert cfg.mcp_servers[3].authorization_headers == {"Authorization": "client"} + + +def test_model_context_protocol_server_resolved_headers_with_special_values() -> None: + """Test that resolved_authorization_headers preserves special values.""" + mcp = ModelContextProtocolServer( + name="test-server", + url="http://localhost:8080", + authorization_headers={ + "Authorization": "kubernetes", + "X-Custom": "client", + }, + ) + assert mcp is not None + # Special values should be preserved in resolved headers + assert mcp.resolved_authorization_headers == { + "Authorization": "kubernetes", + "X-Custom": "client", + } + + +def test_model_context_protocol_server_resolved_headers_with_file( + tmp_path: Path, +) -> None: + """Test that resolved_authorization_headers reads from files.""" + # Create a temporary secret file + secret_file = tmp_path / "secret.txt" + secret_file.write_text("my-secret-value") + + mcp = ModelContextProtocolServer( + name="test-server", + url="http://localhost:8080", + authorization_headers={"Authorization": str(secret_file)}, + ) + assert mcp is not None + # File content should be read into resolved headers + assert mcp.resolved_authorization_headers == {"Authorization": "my-secret-value"} + + +def test_model_context_protocol_server_resolved_headers_empty() -> None: + """Test that resolved_authorization_headers is empty when no auth headers.""" + mcp = ModelContextProtocolServer( + name="test-server", + url="http://localhost:8080", + ) + assert mcp is not None + assert not mcp.resolved_authorization_headers diff --git a/tests/unit/utils/test_mcp_auth_headers.py b/tests/unit/utils/test_mcp_auth_headers.py new file mode 100644 index 00000000..fe70fe19 --- /dev/null +++ b/tests/unit/utils/test_mcp_auth_headers.py @@ -0,0 +1,116 @@ +"""Unit tests for MCP authorization headers utilities.""" + +from pathlib import Path + + +from utils.mcp_auth_headers import resolve_authorization_headers + + +def test_resolve_authorization_headers_empty() -> None: + """Test resolving empty authorization headers.""" + result = resolve_authorization_headers({}) + assert not result + + +def test_resolve_authorization_headers_with_file(tmp_path: Path) -> None: + """Test resolving authorization headers from file.""" + # Create a temporary secret file + secret_file = tmp_path / "secret.txt" + secret_file.write_text("my-secret-token") + + headers = {"Authorization": str(secret_file)} + result = resolve_authorization_headers(headers) + + assert result == {"Authorization": "my-secret-token"} + + +def test_resolve_authorization_headers_with_file_strips_whitespace( + tmp_path: Path, +) -> None: + """Test that resolving headers strips whitespace from file content.""" + secret_file = tmp_path / "secret.txt" + secret_file.write_text(" my-secret-token\n ") + + headers = {"Authorization": str(secret_file)} + result = resolve_authorization_headers(headers) + + assert result == {"Authorization": "my-secret-token"} + + +def test_resolve_authorization_headers_with_nonexistent_file() -> None: + """Test resolving headers with nonexistent file logs warning and skips.""" + headers = {"Authorization": "/nonexistent/path/to/secret.txt"} + result = resolve_authorization_headers(headers) + + # Should return empty dict when file doesn't exist + assert not result + + +def test_resolve_authorization_headers_client_token() -> None: + """Test that client token keyword is preserved.""" + headers = {"Authorization": "client"} + result = resolve_authorization_headers(headers) + + # Should keep "client" as-is for later substitution + assert result == {"Authorization": "client"} + + +def test_resolve_authorization_headers_kubernetes_token() -> None: + """Test that kubernetes keyword is preserved.""" + headers = {"Authorization": "kubernetes"} + result = resolve_authorization_headers(headers) + + # Should keep "kubernetes" as-is for later substitution + assert result == {"Authorization": "kubernetes"} + + +def test_resolve_authorization_headers_multiple_headers(tmp_path: Path) -> None: + """Test resolving multiple authorization headers.""" + # Create multiple secret files + auth_file = tmp_path / "auth.txt" + auth_file.write_text("auth-token") + api_key_file = tmp_path / "api_key.txt" + api_key_file.write_text("api-key-value") + + headers = { + "Authorization": str(auth_file), + "X-API-Key": str(api_key_file), + } + result = resolve_authorization_headers(headers) + + assert result == { + "Authorization": "auth-token", + "X-API-Key": "api-key-value", + } + + +def test_resolve_authorization_headers_mixed_types(tmp_path: Path) -> None: + """Test resolving mixed header types (file, client, kubernetes).""" + # Create a secret file + secret_file = tmp_path / "secret.txt" + secret_file.write_text("file-secret") + + headers = { + "Authorization": "client", + "X-API-Key": str(secret_file), + "X-K8s-Token": "kubernetes", + } + result = resolve_authorization_headers(headers) + + # Special keywords should be preserved, file should be resolved + assert result["Authorization"] == "client" + assert result["X-API-Key"] == "file-secret" + assert result["X-K8s-Token"] == "kubernetes" + + +def test_resolve_authorization_headers_file_read_error(tmp_path: Path) -> None: + """Test handling of file read errors.""" + # Create a directory instead of a file to cause an error + secret_dir = tmp_path / "secret_dir" + secret_dir.mkdir() + + headers = {"Authorization": str(secret_dir)} + result = resolve_authorization_headers(headers) + + # Should handle error gracefully and return empty dict + assert not result