mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-04-25 08:26:21 +02:00
MCP auth for the simple case (#557)
* MCP auth token header * Mention limitations * Fix AgentStep schema error by converting argument values to strings. * Added tests for MCP auth and agent step parsing
This commit is contained in:
cybermaggedon
GitHub
parent
d9d4c91363
commit
4c3db4dbbe
8 changed files with 1361 additions and 56 deletions
233
tests/unit/test_agent/test_mcp_tool_auth.py
Normal file
233
tests/unit/test_agent/test_mcp_tool_auth.py
Normal file
|
|
@ -0,0 +1,233 @@
|
|||
"""
|
||||
Unit tests for MCP tool bearer token authentication
|
||||
|
||||
Tests the authentication feature added to MCP tool service that allows
|
||||
configuring optional bearer tokens for MCP server connections.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock, AsyncMock, patch, MagicMock
|
||||
import json
|
||||
|
||||
|
||||
class TestMcpToolAuthentication:
|
||||
"""Test cases for MCP tool bearer token authentication"""
|
||||
|
||||
def test_mcp_tool_with_auth_token_header_building(self):
|
||||
"""Test that auth token is correctly formatted in headers"""
|
||||
# Arrange
|
||||
mcp_config = {
|
||||
"url": "https://secure.example.com/mcp",
|
||||
"remote-name": "secure-tool",
|
||||
"auth-token": "test-token-12345"
|
||||
}
|
||||
|
||||
# Act - simulate header building logic from service.py
|
||||
headers = {}
|
||||
if "auth-token" in mcp_config and mcp_config["auth-token"]:
|
||||
token = mcp_config["auth-token"]
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
|
||||
# Assert
|
||||
assert "Authorization" in headers
|
||||
assert headers["Authorization"] == "Bearer test-token-12345"
|
||||
assert headers["Authorization"].startswith("Bearer ")
|
||||
|
||||
def test_mcp_tool_without_auth_token_header_building(self):
|
||||
"""Test that no auth header is added when token is not present (backward compatibility)"""
|
||||
# Arrange
|
||||
mcp_config = {
|
||||
"url": "http://public.example.com/mcp",
|
||||
"remote-name": "public-tool"
|
||||
# No auth-token field
|
||||
}
|
||||
|
||||
# Act - simulate header building logic from service.py
|
||||
headers = {}
|
||||
if "auth-token" in mcp_config and mcp_config["auth-token"]:
|
||||
token = mcp_config["auth-token"]
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
|
||||
# Assert
|
||||
assert headers == {}
|
||||
assert "Authorization" not in headers
|
||||
|
||||
def test_mcp_config_with_auth_token(self):
|
||||
"""Test MCP configuration parsing with auth-token"""
|
||||
# Arrange
|
||||
config = {
|
||||
"mcp": {
|
||||
"secure-tool": json.dumps({
|
||||
"url": "https://secure.example.com/mcp",
|
||||
"remote-name": "secure-tool",
|
||||
"auth-token": "test-token-xyz"
|
||||
}),
|
||||
"public-tool": json.dumps({
|
||||
"url": "http://public.example.com/mcp",
|
||||
"remote-name": "public-tool"
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
# Act - simulate on_mcp_config
|
||||
mcp_services = {
|
||||
k: json.loads(v)
|
||||
for k, v in config["mcp"].items()
|
||||
}
|
||||
|
||||
# Assert
|
||||
assert "secure-tool" in mcp_services
|
||||
assert mcp_services["secure-tool"]["auth-token"] == "test-token-xyz"
|
||||
assert mcp_services["secure-tool"]["url"] == "https://secure.example.com/mcp"
|
||||
|
||||
assert "public-tool" in mcp_services
|
||||
assert "auth-token" not in mcp_services["public-tool"]
|
||||
assert mcp_services["public-tool"]["url"] == "http://public.example.com/mcp"
|
||||
|
||||
def test_auth_token_with_empty_string(self):
|
||||
"""Test that empty auth-token string is treated as no auth"""
|
||||
# Arrange
|
||||
config_data = {
|
||||
"url": "https://example.com/mcp",
|
||||
"remote-name": "test-tool",
|
||||
"auth-token": ""
|
||||
}
|
||||
|
||||
# Act - simulate header building logic
|
||||
headers = {}
|
||||
if "auth-token" in config_data and config_data["auth-token"]:
|
||||
headers["Authorization"] = f"Bearer {config_data['auth-token']}"
|
||||
|
||||
# Assert
|
||||
assert headers == {}, "Empty auth-token should not add Authorization header"
|
||||
|
||||
def test_auth_token_with_special_characters(self):
|
||||
"""Test auth token with special characters (JWT-like)"""
|
||||
# Arrange
|
||||
jwt_token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"
|
||||
|
||||
config_data = {
|
||||
"url": "https://example.com/mcp",
|
||||
"auth-token": jwt_token
|
||||
}
|
||||
|
||||
# Act - simulate header building
|
||||
headers = {}
|
||||
if "auth-token" in config_data and config_data["auth-token"]:
|
||||
token = config_data["auth-token"]
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
|
||||
# Assert
|
||||
assert headers["Authorization"] == f"Bearer {jwt_token}"
|
||||
assert "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9" in headers["Authorization"]
|
||||
|
||||
def test_multiple_tools_with_different_auth_configs(self):
|
||||
"""Test handling multiple MCP tools with mixed auth configurations"""
|
||||
# Arrange
|
||||
mcp_services = {
|
||||
"tool-a": {
|
||||
"url": "https://a.example.com/mcp",
|
||||
"auth-token": "token-a"
|
||||
},
|
||||
"tool-b": {
|
||||
"url": "https://b.example.com/mcp",
|
||||
"auth-token": "token-b"
|
||||
},
|
||||
"tool-c": {
|
||||
"url": "http://c.example.com/mcp"
|
||||
# No auth-token
|
||||
}
|
||||
}
|
||||
|
||||
# Act - simulate header building for each tool
|
||||
headers_a = {}
|
||||
if "auth-token" in mcp_services["tool-a"] and mcp_services["tool-a"]["auth-token"]:
|
||||
headers_a["Authorization"] = f"Bearer {mcp_services['tool-a']['auth-token']}"
|
||||
|
||||
headers_b = {}
|
||||
if "auth-token" in mcp_services["tool-b"] and mcp_services["tool-b"]["auth-token"]:
|
||||
headers_b["Authorization"] = f"Bearer {mcp_services['tool-b']['auth-token']}"
|
||||
|
||||
headers_c = {}
|
||||
if "auth-token" in mcp_services["tool-c"] and mcp_services["tool-c"]["auth-token"]:
|
||||
headers_c["Authorization"] = f"Bearer {mcp_services['tool-c']['auth-token']}"
|
||||
|
||||
# Assert
|
||||
assert headers_a == {"Authorization": "Bearer token-a"}
|
||||
assert headers_b == {"Authorization": "Bearer token-b"}
|
||||
assert headers_c == {}
|
||||
|
||||
def test_auth_token_not_logged(self):
|
||||
"""Test that auth tokens are not exposed in logs"""
|
||||
# This is more of a guideline test - in real implementation,
|
||||
# we should ensure tokens are never logged
|
||||
|
||||
# Arrange
|
||||
auth_token = "super-secret-token-123"
|
||||
config = {
|
||||
"url": "https://secure.example.com/mcp",
|
||||
"auth-token": auth_token
|
||||
}
|
||||
|
||||
# Act - simulate log-safe representation
|
||||
def get_log_safe_config(cfg):
|
||||
"""Return config with sensitive data masked"""
|
||||
safe_config = cfg.copy()
|
||||
if "auth-token" in safe_config and safe_config["auth-token"]:
|
||||
safe_config["auth-token"] = "****"
|
||||
return safe_config
|
||||
|
||||
log_safe = get_log_safe_config(config)
|
||||
|
||||
# Assert
|
||||
assert log_safe["auth-token"] == "****"
|
||||
assert auth_token not in str(log_safe)
|
||||
assert "url" in log_safe
|
||||
assert log_safe["url"] == "https://secure.example.com/mcp"
|
||||
|
||||
def test_auth_token_with_remote_name_configuration(self):
|
||||
"""Test MCP tool configuration with both auth-token and remote-name"""
|
||||
# Arrange
|
||||
mcp_config = {
|
||||
"url": "https://server.example.com/mcp",
|
||||
"remote-name": "actual_tool_name",
|
||||
"auth-token": "my-token-456"
|
||||
}
|
||||
|
||||
# Act - simulate header building and remote name extraction
|
||||
headers = {}
|
||||
if "auth-token" in mcp_config and mcp_config["auth-token"]:
|
||||
token = mcp_config["auth-token"]
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
|
||||
remote_name = mcp_config.get("remote-name", "default-name")
|
||||
|
||||
# Assert
|
||||
assert headers["Authorization"] == "Bearer my-token-456"
|
||||
assert remote_name == "actual_tool_name"
|
||||
assert "url" in mcp_config
|
||||
assert mcp_config["url"] == "https://server.example.com/mcp"
|
||||
|
||||
def test_bearer_token_format(self):
|
||||
"""Test that Bearer token format is correct"""
|
||||
# Arrange
|
||||
tokens = [
|
||||
"simple-token",
|
||||
"token_with_underscore",
|
||||
"token-with-dash",
|
||||
"TokenWithMixedCase123",
|
||||
"eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.payload.signature"
|
||||
]
|
||||
|
||||
# Act & Assert
|
||||
for token in tokens:
|
||||
headers = {}
|
||||
if token:
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
|
||||
# Verify format is "Bearer <token>" with single space
|
||||
assert headers["Authorization"].startswith("Bearer ")
|
||||
assert headers["Authorization"] == f"Bearer {token}"
|
||||
# Verify no extra spaces
|
||||
assert headers["Authorization"].count("Bearer") == 1
|
||||
assert headers["Authorization"].split("Bearer ")[1] == token
|
||||
Loading…
Add table
Add a link
Reference in a new issue