Fimeg/letta-server
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

fix: handle faulty schemas from bad mcp servers better

Browse Source 
Co-authored-by: jnjpng <jin@letta.com>
Co-authored-by: Jin Peng <jinjpeng@Jins-MacBook-Pro.local>
This commit is contained in:
Charles Packer
2025-08-15 16:08:52 -07:00
committed by GitHub
parent 1ff2434598
commit 773a6452d1
10 changed files with 858 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

77
tests/mcp_tests/test_mcp.py
View File

@@ -210,6 +210,83 @@ def test_stdio_mcp_server(client, agent_state):
assert len(ret.tool_return.strip()) >= 10, f"Expected at least 10 characters in tool_return, got {len(ret.tool_return.strip())}"
# Optional OpenAI validation test for MCP-normalized schema
# Skips unless OPENAI_API_KEY is set to avoid network flakiness in CI
EXAMPLE_BAD_SCHEMA = {
"type": "object",
"properties": {
"conversation_type": {
"type": "string",
"const": "Group",
"description": "Specifies the type of conversation to be created. Must be 'Group' for this action.",
},
"message": {
"type": "object",
"additionalProperties": {}, # invalid for OpenAI: missing "type"
"description": "Initial message payload",
},
"participant_ids": {
"type": "array",
"items": {"type": "string"},
"description": "Participant IDs",
},
},
"required": ["conversation_type", "message", "participant_ids"],
"additionalProperties": False,
"$schema": "http://json-schema.org/draft-07/schema#",
}
@pytest.mark.skipif(
not os.getenv("OPENAI_API_KEY"),
reason="Requires OPENAI_API_KEY to call OpenAI for schema validation",
)
def test_openai_rejects_untyped_additional_properties_and_accepts_normalized_schema():
"""Test written to check if our extra schema validation works.
Some MCP servers will return faulty schemas that require correction, or they will brick the LLM client calls.
"""
import copy
try:
from openai import OpenAI
except Exception as e: # pragma: no cover
pytest.skip(f"openai package not available: {e}")
client = OpenAI()
def run_request_with_schema(schema: dict):
tools = [
{
"type": "function",
"function": {
"name": "TWITTER_CREATE_A_NEW_DM_CONVERSATION",
"description": "Create a DM conversation",
"parameters": schema,
"strict": True,
},
}
]
return client.chat.completions.create(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hello"}],
tools=tools,
)
# Bad schema should raise
with pytest.raises(Exception):
run_request_with_schema(EXAMPLE_BAD_SCHEMA)
# Normalized should succeed
normalized = copy.deepcopy(EXAMPLE_BAD_SCHEMA)
normalized["properties"]["message"]["additionalProperties"] = False
normalized["properties"]["message"]["properties"] = {"text": {"type": "string"}}
normalized["properties"]["message"]["required"] = ["text"]
resp = run_request_with_schema(normalized)
assert getattr(resp, "id", None)
@pytest.mark.asyncio
async def test_streamable_http_mcp_server_update_schema_no_docstring_required(client, agent_state, server_url):
"""

185
tests/mcp_tests/test_mcp_schema_validation.py Normal file
View File

@@ -0,0 +1,185 @@
"""
Test MCP tool schema validation integration.
"""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from letta.functions.mcp_client.types import MCPTool, MCPToolHealth
from letta.functions.schema_validator import SchemaHealth, validate_complete_json_schema
@pytest.mark.asyncio
async def test_mcp_tools_get_health_status():
"""Test that MCP tools receive health status when listed."""
from letta.server.server import SyncServer
# Create mock tools with different schema types
mock_tools = [
# Strict compliant tool
MCPTool(
name="strict_tool",
inputSchema={"type": "object", "properties": {"text": {"type": "string"}}, "required": ["text"], "additionalProperties": False},
),
# Non-strict tool (free-form object)
MCPTool(
name="non_strict_tool",
inputSchema={
"type": "object",
"properties": {"message": {"type": "object", "additionalProperties": {}}}, # Free-form object
"required": ["message"],
"additionalProperties": False,
},
),
# Invalid tool (missing type)
MCPTool(name="invalid_tool", inputSchema={"properties": {"data": {"type": "string"}}, "required": ["data"]}),
]
# Mock the server and client
mock_client = AsyncMock()
mock_client.list_tools = AsyncMock(return_value=mock_tools)
# Call the method directly
actual_server = SyncServer.__new__(SyncServer)
actual_server.mcp_clients = {"test_server": mock_client}
tools = await actual_server.get_tools_from_mcp_server("test_server")
# Verify health status was added
assert len(tools) == 3
# Check strict tool
strict_tool = tools[0]
assert strict_tool.name == "strict_tool"
assert strict_tool.health is not None
assert strict_tool.health.status == SchemaHealth.STRICT_COMPLIANT.value
assert strict_tool.health.reasons == []
# Check non-strict tool
non_strict_tool = tools[1]
assert non_strict_tool.name == "non_strict_tool"
assert non_strict_tool.health is not None
assert non_strict_tool.health.status == SchemaHealth.NON_STRICT_ONLY.value
assert len(non_strict_tool.health.reasons) > 0
assert any("additionalProperties" in reason for reason in non_strict_tool.health.reasons)
# Check invalid tool
invalid_tool = tools[2]
assert invalid_tool.name == "invalid_tool"
assert invalid_tool.health is not None
assert invalid_tool.health.status == SchemaHealth.INVALID.value
assert len(invalid_tool.health.reasons) > 0
assert any("type" in reason for reason in invalid_tool.health.reasons)
def test_composio_like_schema_marked_non_strict():
"""Test that Composio-like schemas are correctly marked as NON_STRICT_ONLY."""
# Example schema from Composio with free-form message object
composio_schema = {
"type": "object",
"properties": {
"message": {"type": "object", "additionalProperties": {}, "description": "Message to send"} # Free-form, missing "type"
},
"required": ["message"],
"additionalProperties": False,
}
status, reasons = validate_complete_json_schema(composio_schema)
assert status == SchemaHealth.NON_STRICT_ONLY
assert len(reasons) > 0
assert any("additionalProperties" in reason for reason in reasons)
def test_empty_object_in_required_marked_invalid():
"""Test that required properties allowing empty objects are marked INVALID."""
schema = {
"type": "object",
"properties": {
"config": {"type": "object", "properties": {}, "required": [], "additionalProperties": False} # Empty object schema
},
"required": ["config"], # Required but allows empty object
"additionalProperties": False,
}
status, reasons = validate_complete_json_schema(schema)
assert status == SchemaHealth.INVALID
assert any("empty object" in reason for reason in reasons)
assert any("config" in reason for reason in reasons)
@pytest.mark.asyncio
async def test_add_mcp_tool_rejects_non_strict_schemas():
"""Test that adding MCP tools with non-strict schemas is rejected."""
from fastapi import HTTPException
from letta.server.rest_api.routers.v1.tools import add_mcp_tool
from letta.settings import tool_settings
# Mock a non-strict tool
non_strict_tool = MCPTool(
name="test_tool",
inputSchema={
"type": "object",
"properties": {"message": {"type": "object"}}, # Missing additionalProperties: false
"required": ["message"],
"additionalProperties": False,
},
)
non_strict_tool.health = MCPToolHealth(status=SchemaHealth.NON_STRICT_ONLY.value, reasons=["Missing additionalProperties for message"])
# Mock server response
with patch("letta.server.rest_api.routers.v1.tools.get_letta_server") as mock_get_server:
with patch.object(tool_settings, "mcp_read_from_config", True): # Ensure we're using config path
mock_server = AsyncMock()
mock_server.get_tools_from_mcp_server = AsyncMock(return_value=[non_strict_tool])
mock_server.user_manager.get_user_or_default = MagicMock()
mock_get_server.return_value = mock_server
# Should raise HTTPException for non-strict schema
with pytest.raises(HTTPException) as exc_info:
await add_mcp_tool(mcp_server_name="test_server", mcp_tool_name="test_tool", server=mock_server, actor_id=None)
assert exc_info.value.status_code == 400
assert "non-strict schema" in exc_info.value.detail["message"].lower()
assert exc_info.value.detail["health_status"] == SchemaHealth.NON_STRICT_ONLY.value
@pytest.mark.asyncio
async def test_add_mcp_tool_rejects_invalid_schemas():
"""Test that adding MCP tools with invalid schemas is rejected."""
from fastapi import HTTPException
from letta.server.rest_api.routers.v1.tools import add_mcp_tool
from letta.settings import tool_settings
# Mock an invalid tool
invalid_tool = MCPTool(
name="test_tool",
inputSchema={
"properties": {"data": {"type": "string"}},
"required": ["data"],
# Missing "type": "object"
},
)
invalid_tool.health = MCPToolHealth(status=SchemaHealth.INVALID.value, reasons=["Missing 'type' at root level"])
# Mock server response
with patch("letta.server.rest_api.routers.v1.tools.get_letta_server") as mock_get_server:
with patch.object(tool_settings, "mcp_read_from_config", True): # Ensure we're using config path
mock_server = AsyncMock()
mock_server.get_tools_from_mcp_server = AsyncMock(return_value=[invalid_tool])
mock_server.user_manager.get_user_or_default = MagicMock()
mock_get_server.return_value = mock_server
# Should raise HTTPException for invalid schema
with pytest.raises(HTTPException) as exc_info:
await add_mcp_tool(mcp_server_name="test_server", mcp_tool_name="test_tool", server=mock_server, actor_id=None)
assert exc_info.value.status_code == 400
assert "invalid schema" in exc_info.value.detail["message"].lower()
assert exc_info.value.detail["health_status"] == SchemaHealth.INVALID.value

314
tests/mcp_tests/test_schema_validator.py Normal file
View File

@@ -0,0 +1,314 @@
"""
Unit tests for the JSON Schema validator for OpenAI strict mode compliance.
"""
from letta.functions.schema_validator import SchemaHealth, validate_complete_json_schema
class TestSchemaValidator:
"""Test cases for the schema validator."""
def test_valid_strict_compliant_schema(self):
"""Test a fully strict-compliant schema."""
schema = {
"type": "object",
"properties": {
"name": {"type": "string", "description": "The name of the user"},
"age": {"type": "integer", "description": "The age of the user"},
"address": {
"type": "object",
"properties": {"street": {"type": "string"}, "city": {"type": "string"}},
"required": ["street", "city"],
"additionalProperties": False,
},
},
"required": ["name", "age"],
"additionalProperties": False,
}
status, reasons = validate_complete_json_schema(schema)
assert status == SchemaHealth.STRICT_COMPLIANT
assert reasons == []
def test_free_form_object_non_strict(self):
"""Test that free-form objects (like Composio message) are marked as NON_STRICT_ONLY."""
schema = {
"type": "object",
"properties": {
"message": {
"type": "object",
"description": "A message object",
# Missing additionalProperties: false makes this free-form
}
},
"required": ["message"],
"additionalProperties": False,
}
status, reasons = validate_complete_json_schema(schema)
assert status == SchemaHealth.NON_STRICT_ONLY
assert any("additionalProperties" in reason for reason in reasons)
def test_empty_object_in_required_invalid(self):
"""Test that required properties allowing empty objects are marked INVALID."""
schema = {
"type": "object",
"properties": {
"config": {"type": "object", "properties": {}, "required": [], "additionalProperties": False} # Empty object schema
},
"required": ["config"], # Required but allows empty object
"additionalProperties": False,
}
status, reasons = validate_complete_json_schema(schema)
assert status == SchemaHealth.INVALID
assert any("empty object" in reason for reason in reasons)
def test_missing_type_invalid(self):
"""Test that schemas missing type are marked INVALID."""
schema = {
# Missing "type": "object"
"properties": {"name": {"type": "string"}},
"required": ["name"],
}
status, reasons = validate_complete_json_schema(schema)
assert status == SchemaHealth.INVALID
assert any("type" in reason.lower() for reason in reasons)
def test_missing_items_in_array_invalid(self):
"""Test that arrays without items definition are marked INVALID."""
schema = {
"type": "object",
"properties": {
"tags": {
"type": "array"
# Missing "items" definition
}
},
"required": ["tags"],
"additionalProperties": False,
}
status, reasons = validate_complete_json_schema(schema)
assert status == SchemaHealth.INVALID
assert any("items" in reason for reason in reasons)
def test_required_property_not_in_properties_invalid(self):
"""Test that required properties not defined in properties are marked INVALID."""
schema = {
"type": "object",
"properties": {"name": {"type": "string"}},
"required": ["name", "email"], # "email" not in properties
"additionalProperties": False,
}
status, reasons = validate_complete_json_schema(schema)
assert status == SchemaHealth.INVALID
assert any("email" in reason and "not found" in reason for reason in reasons)
def test_nested_object_validation(self):
"""Test that nested objects are properly validated."""
schema = {
"type": "object",
"properties": {
"user": {
"type": "object",
"properties": {
"profile": {
"type": "object",
"properties": {"bio": {"type": "string"}},
# Missing additionalProperties and required
}
},
"required": ["profile"],
"additionalProperties": False,
}
},
"required": ["user"],
"additionalProperties": False,
}
status, reasons = validate_complete_json_schema(schema)
assert status == SchemaHealth.NON_STRICT_ONLY
# Should have warnings about nested profile object
assert any("profile" in reason.lower() or "properties.profile" in reason for reason in reasons)
def test_union_types_with_anyof(self):
"""Test schemas with anyOf union types."""
schema = {
"type": "object",
"properties": {"value": {"anyOf": [{"type": "string"}, {"type": "number"}]}},
"required": ["value"],
"additionalProperties": False,
}
status, reasons = validate_complete_json_schema(schema)
assert status == SchemaHealth.STRICT_COMPLIANT
assert reasons == []
def test_array_with_proper_items(self):
"""Test arrays with properly defined items."""
schema = {
"type": "object",
"properties": {
"items": {
"type": "array",
"items": {
"type": "object",
"properties": {"id": {"type": "string"}, "value": {"type": "number"}},
"required": ["id", "value"],
"additionalProperties": False,
},
}
},
"required": ["items"],
"additionalProperties": False,
}
status, reasons = validate_complete_json_schema(schema)
assert status == SchemaHealth.STRICT_COMPLIANT
assert reasons == []
def test_empty_array_in_required_invalid(self):
"""Test that required properties allowing empty arrays are marked INVALID."""
schema = {
"type": "object",
"properties": {
"tags": {
"type": "array",
"items": {"type": "string"},
# No minItems constraint, allows empty array
}
},
"required": ["tags"],
"additionalProperties": False,
}
# This should actually be STRICT_COMPLIANT since empty arrays with defined items are OK
status, reasons = validate_complete_json_schema(schema)
assert status == SchemaHealth.STRICT_COMPLIANT
def test_array_without_constraints_invalid(self):
"""Test that arrays without any constraints in required props are invalid."""
schema = {
"type": "object",
"properties": {
"data": {
"type": "array"
# No items defined at all - completely unconstrained
}
},
"required": ["data"],
"additionalProperties": False,
}
status, reasons = validate_complete_json_schema(schema)
assert status == SchemaHealth.INVALID
assert any("items" in reason for reason in reasons)
def test_composio_like_schema(self):
"""Test a schema similar to Composio's free-form message structure."""
schema = {
"type": "object",
"properties": {
"message": {
"type": "object",
"description": "Message to send",
# No properties defined, no additionalProperties: false
# This is a free-form object
}
},
"required": ["message"],
"additionalProperties": False,
}
status, reasons = validate_complete_json_schema(schema)
assert status == SchemaHealth.NON_STRICT_ONLY
assert any("additionalProperties" in reason for reason in reasons)
def test_non_dict_schema(self):
"""Test that non-dict schemas are marked INVALID."""
schema = "not a dict"
status, reasons = validate_complete_json_schema(schema)
assert status == SchemaHealth.INVALID
assert any("dict" in reason for reason in reasons)
def test_schema_with_defaults_strict_compliant(self):
"""Test that root-level schemas without required field are STRICT_COMPLIANT."""
schema = {
"type": "object",
"properties": {"name": {"type": "string"}, "optional": {"type": "string"}},
# Missing "required" field at root level is OK
"additionalProperties": False,
}
status, reasons = validate_complete_json_schema(schema)
# After fix, root level without required should be STRICT_COMPLIANT
assert status == SchemaHealth.STRICT_COMPLIANT
assert reasons == []
def test_composio_schema_with_optional_root_properties_strict_compliant(self):
"""Test that Composio-like schemas with optional root properties are STRICT_COMPLIANT."""
schema = {
"type": "object",
"properties": {
"thinking": {"type": "string", "description": "Deep inner monologue"},
"connected_account_id": {"type": "string", "description": "Specific connected account ID"},
"toolkit": {"type": "string", "description": "Name of the toolkit"},
"request_heartbeat": {"type": "boolean", "description": "Request immediate heartbeat"},
},
"required": ["thinking", "request_heartbeat"], # Not all properties are required
"additionalProperties": False,
}
status, reasons = validate_complete_json_schema(schema)
assert status == SchemaHealth.STRICT_COMPLIANT
assert reasons == []
def test_root_level_without_required_strict_compliant(self):
"""Test that root-level objects without 'required' field are STRICT_COMPLIANT."""
schema = {
"type": "object",
"properties": {
"name": {"type": "string"},
"age": {"type": "integer"},
},
# No "required" field at root level
"additionalProperties": False,
}
status, reasons = validate_complete_json_schema(schema)
# Root level without required should be STRICT_COMPLIANT
assert status == SchemaHealth.STRICT_COMPLIANT
assert reasons == []
def test_nested_object_without_required_non_strict(self):
"""Test that nested objects without 'required' remain NON_STRICT_ONLY."""
schema = {
"type": "object",
"properties": {
"user": {
"type": "object",
"properties": {
"preferences": {
"type": "object",
"properties": {"theme": {"type": "string"}, "language": {"type": "string"}},
# Missing "required" field in nested object
"additionalProperties": False,
},
"name": {"type": "string"},
},
"required": ["name"], # Don't require preferences so it's not marked INVALID
"additionalProperties": False,
}
},
"required": ["user"],
"additionalProperties": False,
}
status, reasons = validate_complete_json_schema(schema)
assert status == SchemaHealth.NON_STRICT_ONLY
# Should have warning about nested preferences object missing 'required'
assert any("required" in reason and "preferences" in reason for reason in reasons)