mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-27 07:50:43 +02:00
Merge branch 'main' into responsecachingmiddleware
This commit is contained in:
William Easton
GitHub
commit
28370827dc
187 changed files with 8368 additions and 2798 deletions
|
|
@ -11,6 +11,7 @@ from fastmcp.server.middleware.error_handling import (
|
|||
RetryMiddleware,
|
||||
)
|
||||
from fastmcp.server.middleware.middleware import MiddlewareContext
|
||||
from fastmcp.utilities.tests import caplog_for_fastmcp
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -60,8 +61,9 @@ class TestErrorHandlingMiddleware:
|
|||
middleware = ErrorHandlingMiddleware()
|
||||
error = ValueError("test error")
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
middleware._log_error(error, mock_context)
|
||||
with caplog_for_fastmcp(caplog):
|
||||
with caplog.at_level(logging.ERROR):
|
||||
middleware._log_error(error, mock_context)
|
||||
|
||||
assert "Error in test_method: ValueError: test error" in caplog.text
|
||||
assert "ValueError:test_method" in middleware.error_counts
|
||||
|
|
@ -72,8 +74,9 @@ class TestErrorHandlingMiddleware:
|
|||
middleware = ErrorHandlingMiddleware(include_traceback=True)
|
||||
error = ValueError("test error")
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
middleware._log_error(error, mock_context)
|
||||
with caplog_for_fastmcp(caplog):
|
||||
with caplog.at_level(logging.ERROR):
|
||||
middleware._log_error(error, mock_context)
|
||||
|
||||
assert "Error in test_method: ValueError: test error" in caplog.text
|
||||
# The traceback is added to the log message
|
||||
|
|
@ -95,8 +98,9 @@ class TestErrorHandlingMiddleware:
|
|||
middleware = ErrorHandlingMiddleware(error_callback=callback)
|
||||
error = ValueError("test error")
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
middleware._log_error(error, mock_context)
|
||||
with caplog_for_fastmcp(caplog):
|
||||
with caplog.at_level(logging.ERROR):
|
||||
middleware._log_error(error, mock_context)
|
||||
|
||||
assert "Error in error callback: callback error" in caplog.text
|
||||
|
||||
|
|
@ -189,9 +193,10 @@ class TestErrorHandlingMiddleware:
|
|||
middleware = ErrorHandlingMiddleware()
|
||||
mock_call_next = AsyncMock(side_effect=ValueError("test error"))
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await middleware.on_message(mock_context, mock_call_next)
|
||||
with caplog_for_fastmcp(caplog):
|
||||
with caplog.at_level(logging.ERROR):
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await middleware.on_message(mock_context, mock_call_next)
|
||||
|
||||
assert isinstance(exc_info.value, McpError)
|
||||
assert exc_info.value.error.code == -32602
|
||||
|
|
@ -293,8 +298,9 @@ class TestRetryMiddleware:
|
|||
]
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
result = await middleware.on_request(mock_context, mock_call_next)
|
||||
with caplog_for_fastmcp(caplog):
|
||||
with caplog.at_level(logging.WARNING):
|
||||
result = await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
assert result == "test_result"
|
||||
assert mock_call_next.call_count == 3
|
||||
|
|
@ -307,9 +313,10 @@ class TestRetryMiddleware:
|
|||
# Fail all attempts
|
||||
mock_call_next = AsyncMock(side_effect=ConnectionError("connection failed"))
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
with pytest.raises(ConnectionError):
|
||||
await middleware.on_request(mock_context, mock_call_next)
|
||||
with caplog_for_fastmcp(caplog):
|
||||
with caplog.at_level(logging.WARNING):
|
||||
with pytest.raises(ConnectionError):
|
||||
await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
assert mock_call_next.call_count == 3 # initial + 2 retries
|
||||
assert "Retrying in" in caplog.text
|
||||
|
|
@ -385,14 +392,19 @@ class TestErrorHandlingMiddlewareIntegration:
|
|||
|
||||
error_handling_server.add_middleware(ErrorHandlingMiddleware())
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
async with Client(error_handling_server) as client:
|
||||
# Test different types of errors
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool("failing_operation", {"error_type": "value"})
|
||||
with caplog_for_fastmcp(caplog):
|
||||
with caplog.at_level(logging.ERROR):
|
||||
async with Client(error_handling_server) as client:
|
||||
# Test different types of errors
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool(
|
||||
"failing_operation", {"error_type": "value"}
|
||||
)
|
||||
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool("failing_operation", {"error_type": "file"})
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool(
|
||||
"failing_operation", {"error_type": "file"}
|
||||
)
|
||||
|
||||
log_text = caplog.text
|
||||
|
||||
|
|
@ -443,17 +455,20 @@ class TestErrorHandlingMiddlewareIntegration:
|
|||
|
||||
error_handling_server.add_middleware(ErrorHandlingMiddleware())
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
async with Client(error_handling_server) as client:
|
||||
# Successful operation (should not generate error logs)
|
||||
await client.call_tool("reliable_operation", {"data": "test"})
|
||||
with caplog_for_fastmcp(caplog):
|
||||
with caplog.at_level(logging.ERROR):
|
||||
async with Client(error_handling_server) as client:
|
||||
# Successful operation (should not generate error logs)
|
||||
await client.call_tool("reliable_operation", {"data": "test"})
|
||||
|
||||
# Failed operation (should generate error log)
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool("failing_operation", {"error_type": "value"})
|
||||
# Failed operation (should generate error log)
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool(
|
||||
"failing_operation", {"error_type": "value"}
|
||||
)
|
||||
|
||||
# Another successful operation
|
||||
await client.call_tool("reliable_operation", {"data": "test2"})
|
||||
# Another successful operation
|
||||
await client.call_tool("reliable_operation", {"data": "test2"})
|
||||
|
||||
log_text = caplog.text
|
||||
|
||||
|
|
@ -533,18 +548,19 @@ class TestRetryMiddlewareIntegration:
|
|||
)
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
async with Client(error_handling_server) as client:
|
||||
# This operation fails intermittently - try several times
|
||||
success_count = 0
|
||||
for _ in range(5):
|
||||
try:
|
||||
await client.call_tool(
|
||||
"intermittent_operation", {"fail_rate": 0.7}
|
||||
)
|
||||
success_count += 1
|
||||
except Exception:
|
||||
pass # Some failures expected even with retries
|
||||
with caplog_for_fastmcp(caplog):
|
||||
with caplog.at_level(logging.WARNING):
|
||||
async with Client(error_handling_server) as client:
|
||||
# This operation fails intermittently - try several times
|
||||
success_count = 0
|
||||
for _ in range(5):
|
||||
try:
|
||||
await client.call_tool(
|
||||
"intermittent_operation", {"fail_rate": 0.7}
|
||||
)
|
||||
success_count += 1
|
||||
except Exception:
|
||||
pass # Some failures expected even with retries
|
||||
|
||||
# Should have some retry log messages
|
||||
# Note: Retry logs might not appear if the underlying errors are wrapped by FastMCP
|
||||
|
|
@ -584,17 +600,22 @@ class TestRetryMiddlewareIntegration:
|
|||
)
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
async with Client(error_handling_server) as client:
|
||||
# Try intermittent operation
|
||||
try:
|
||||
await client.call_tool("intermittent_operation", {"fail_rate": 0.9})
|
||||
except Exception:
|
||||
pass # May still fail even with retries
|
||||
with caplog_for_fastmcp(caplog):
|
||||
with caplog.at_level(logging.ERROR):
|
||||
async with Client(error_handling_server) as client:
|
||||
# Try intermittent operation
|
||||
try:
|
||||
await client.call_tool(
|
||||
"intermittent_operation", {"fail_rate": 0.9}
|
||||
)
|
||||
except Exception:
|
||||
pass # May still fail even with retries
|
||||
|
||||
# Try permanent failure
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool("failing_operation", {"error_type": "value"})
|
||||
# Try permanent failure
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool(
|
||||
"failing_operation", {"error_type": "value"}
|
||||
)
|
||||
|
||||
log_text = caplog.text
|
||||
|
||||
|
|
|
|||
251
tests/server/middleware/test_initialization_middleware.py
Normal file
251
tests/server/middleware/test_initialization_middleware.py
Normal file
|
|
@ -0,0 +1,251 @@
|
|||
"""Tests for middleware support during initialization."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import mcp.types as mt
|
||||
|
||||
from fastmcp import Client, FastMCP
|
||||
from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
|
||||
|
||||
|
||||
class InitializationMiddleware(Middleware):
|
||||
"""Middleware that captures initialization details."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.initialized = False
|
||||
self.client_info = None
|
||||
self.session_data = {}
|
||||
|
||||
async def on_initialize(
|
||||
self,
|
||||
context: MiddlewareContext[mt.InitializeRequest],
|
||||
call_next: CallNext[mt.InitializeRequest, None],
|
||||
) -> None:
|
||||
"""Capture initialization details and store session data."""
|
||||
self.initialized = True
|
||||
|
||||
# Extract client info from the initialize params
|
||||
if hasattr(context.message, "params") and hasattr(
|
||||
context.message.params, "clientInfo"
|
||||
):
|
||||
self.client_info = context.message.params.clientInfo
|
||||
|
||||
# Store data in the context state for cross-request access
|
||||
if context.fastmcp_context:
|
||||
context.fastmcp_context.set_state("client_initialized", True)
|
||||
if self.client_info:
|
||||
context.fastmcp_context.set_state(
|
||||
"client_name", getattr(self.client_info, "name", "unknown")
|
||||
)
|
||||
|
||||
return await call_next(context)
|
||||
|
||||
|
||||
class ClientDetectionMiddleware(Middleware):
|
||||
"""Middleware that detects specific clients and modifies behavior.
|
||||
|
||||
This demonstrates storing data in the middleware instance itself
|
||||
for cross-request access, since context state is request-scoped.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.is_test_client = False
|
||||
self.tools_modified = False
|
||||
self.initialization_called = False
|
||||
|
||||
async def on_initialize(
|
||||
self,
|
||||
context: MiddlewareContext[mt.InitializeRequest],
|
||||
call_next: CallNext[mt.InitializeRequest, None],
|
||||
) -> None:
|
||||
"""Detect test client during initialization."""
|
||||
self.initialization_called = True
|
||||
|
||||
# For testing purposes, always set it to true
|
||||
# Store in instance variable for cross-request access
|
||||
self.is_test_client = True
|
||||
|
||||
return await call_next(context)
|
||||
|
||||
async def on_list_tools(
|
||||
self,
|
||||
context: MiddlewareContext[mt.ListToolsRequest],
|
||||
call_next: CallNext[mt.ListToolsRequest, list],
|
||||
) -> list:
|
||||
"""Modify tools based on client detection."""
|
||||
tools = await call_next(context)
|
||||
|
||||
# Use the instance variable set during initialization
|
||||
if self.is_test_client:
|
||||
# Add a special annotation to tools for test clients
|
||||
for tool in tools:
|
||||
if not hasattr(tool, "annotations"):
|
||||
tool.annotations = mt.ToolAnnotations()
|
||||
if tool.annotations is None:
|
||||
tool.annotations = mt.ToolAnnotations()
|
||||
# Mark as read-only for test clients
|
||||
tool.annotations.readOnlyHint = True
|
||||
self.tools_modified = True
|
||||
|
||||
return tools
|
||||
|
||||
|
||||
async def test_simple_initialization_hook():
|
||||
"""Test that the on_initialize hook is called."""
|
||||
server = FastMCP("TestServer")
|
||||
|
||||
class SimpleInitMiddleware(Middleware):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.called = False
|
||||
|
||||
async def on_initialize(
|
||||
self,
|
||||
context: MiddlewareContext[mt.InitializeRequest],
|
||||
call_next: CallNext[mt.InitializeRequest, None],
|
||||
) -> None:
|
||||
self.called = True
|
||||
return await call_next(context)
|
||||
|
||||
middleware = SimpleInitMiddleware()
|
||||
server.add_middleware(middleware)
|
||||
|
||||
# Connect client
|
||||
async with Client(server):
|
||||
# Middleware should have been called
|
||||
assert middleware.called is True, "on_initialize was not called"
|
||||
|
||||
|
||||
async def test_middleware_receives_initialization():
|
||||
"""Test that middleware can intercept initialization requests."""
|
||||
server = FastMCP("TestServer")
|
||||
middleware = InitializationMiddleware()
|
||||
server.add_middleware(middleware)
|
||||
|
||||
@server.tool
|
||||
def test_tool(x: int) -> str:
|
||||
return f"Result: {x}"
|
||||
|
||||
# Connect client
|
||||
async with Client(server) as client:
|
||||
# Middleware should have been called during initialization
|
||||
assert middleware.initialized is True
|
||||
|
||||
# Test that the tool still works
|
||||
result = await client.call_tool("test_tool", {"x": 42})
|
||||
assert result.content[0].text == "Result: 42" # type: ignore[attr-defined]
|
||||
|
||||
|
||||
async def test_client_detection_middleware():
|
||||
"""Test middleware that detects specific clients and modifies behavior."""
|
||||
server = FastMCP("TestServer")
|
||||
middleware = ClientDetectionMiddleware()
|
||||
server.add_middleware(middleware)
|
||||
|
||||
@server.tool
|
||||
def example_tool() -> str:
|
||||
return "example"
|
||||
|
||||
# Connect with a client
|
||||
async with Client(server) as client:
|
||||
# Middleware should have been called during initialization
|
||||
assert middleware.initialization_called is True
|
||||
assert middleware.is_test_client is True
|
||||
|
||||
# List tools to trigger modification
|
||||
tools = await client.list_tools()
|
||||
assert len(tools) == 1
|
||||
assert middleware.tools_modified is True
|
||||
|
||||
# Check that the tool has the modified annotation
|
||||
tool = tools[0]
|
||||
assert tool.annotations is not None
|
||||
assert tool.annotations.readOnlyHint is True
|
||||
|
||||
|
||||
async def test_multiple_middleware_initialization():
|
||||
"""Test that multiple middleware can handle initialization."""
|
||||
server = FastMCP("TestServer")
|
||||
|
||||
init_mw = InitializationMiddleware()
|
||||
detect_mw = ClientDetectionMiddleware()
|
||||
|
||||
server.add_middleware(init_mw)
|
||||
server.add_middleware(detect_mw)
|
||||
|
||||
@server.tool
|
||||
def test_tool() -> str:
|
||||
return "test"
|
||||
|
||||
async with Client(server) as client:
|
||||
# Both middleware should have processed initialization
|
||||
assert init_mw.initialized is True
|
||||
assert detect_mw.initialization_called is True
|
||||
assert detect_mw.is_test_client is True
|
||||
|
||||
# List tools to check detection worked
|
||||
await client.list_tools()
|
||||
assert detect_mw.tools_modified is True
|
||||
|
||||
|
||||
async def test_initialization_middleware_with_state_sharing():
|
||||
"""Test that state set during initialization is available in later requests."""
|
||||
server = FastMCP("TestServer")
|
||||
|
||||
class StateTrackingMiddleware(Middleware):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.init_state = {}
|
||||
self.tool_state = {}
|
||||
|
||||
async def on_initialize(
|
||||
self,
|
||||
context: MiddlewareContext[mt.InitializeRequest],
|
||||
call_next: CallNext[mt.InitializeRequest, None],
|
||||
) -> None:
|
||||
# Store some state during initialization
|
||||
if context.fastmcp_context:
|
||||
context.fastmcp_context.set_state("init_timestamp", "2024-01-01")
|
||||
context.fastmcp_context.set_state("client_id", "test-123")
|
||||
self.init_state["timestamp"] = "2024-01-01"
|
||||
self.init_state["client_id"] = "test-123"
|
||||
|
||||
return await call_next(context)
|
||||
|
||||
async def on_call_tool(
|
||||
self,
|
||||
context: MiddlewareContext[mt.CallToolRequestParams],
|
||||
call_next: CallNext[mt.CallToolRequestParams, Any],
|
||||
) -> Any:
|
||||
# Try to access state from initialization
|
||||
if context.fastmcp_context:
|
||||
timestamp = context.fastmcp_context.get_state("init_timestamp")
|
||||
client_id = context.fastmcp_context.get_state("client_id")
|
||||
self.tool_state["timestamp"] = timestamp
|
||||
self.tool_state["client_id"] = client_id
|
||||
|
||||
return await call_next(context)
|
||||
|
||||
middleware = StateTrackingMiddleware()
|
||||
server.add_middleware(middleware)
|
||||
|
||||
@server.tool
|
||||
def test_tool() -> str:
|
||||
return "success"
|
||||
|
||||
async with Client(server) as client:
|
||||
# Initialization should have set state
|
||||
assert middleware.init_state["timestamp"] == "2024-01-01"
|
||||
assert middleware.init_state["client_id"] == "test-123"
|
||||
|
||||
# Call a tool - state should be accessible
|
||||
result = await client.call_tool("test_tool", {})
|
||||
assert result.content[0].text == "success" # type: ignore[attr-defined]
|
||||
|
||||
# State should have been accessible during tool call
|
||||
# Note: State is request-scoped, so it won't persist across requests
|
||||
# This test shows the pattern, but actual cross-request state would need
|
||||
# external storage (Redis, DB, etc.)
|
||||
# The middleware.tool_state might be None if state doesn't persist
|
||||
|
|
@ -2,37 +2,40 @@
|
|||
|
||||
import datetime
|
||||
import logging
|
||||
import re
|
||||
from collections.abc import Generator
|
||||
from typing import Any, Literal, TypeVar
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import mcp
|
||||
import mcp.types
|
||||
import pytest
|
||||
from inline_snapshot import snapshot
|
||||
from pydantic import AnyUrl
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.client import Client
|
||||
from fastmcp.resources.template import ResourceTemplate
|
||||
from fastmcp.server.middleware.logging import (
|
||||
LoggingMiddleware,
|
||||
StructuredLoggingMiddleware,
|
||||
)
|
||||
from fastmcp.server.middleware.middleware import MiddlewareContext
|
||||
from fastmcp.server.middleware.middleware import CallNext, MiddlewareContext
|
||||
from fastmcp.utilities.tests import caplog_for_fastmcp
|
||||
|
||||
FIXED_DATE = datetime.datetime(2023, 1, 1, tzinfo=datetime.timezone.utc)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
def remove_line_numbers(logs: str) -> str:
|
||||
"""Remove line numbers from log messages."""
|
||||
trimmed_logs = ""
|
||||
lines = logs.split("\n")
|
||||
for line in lines:
|
||||
# Match only the first `:\d+ `
|
||||
line = re.sub(pattern=r":\d+ ", repl=":LINE_NUMBER ", string=line, count=1)
|
||||
trimmed_logs += line + "\n"
|
||||
return trimmed_logs
|
||||
def get_log_lines(
|
||||
caplog: pytest.LogCaptureFixture, module: str | None = None
|
||||
) -> list[str]:
|
||||
"""Get log lines from a caplog fixture."""
|
||||
return [
|
||||
record.message
|
||||
for record in caplog.records
|
||||
if (module or "logging") in record.name
|
||||
]
|
||||
|
||||
|
||||
def new_mock_context(
|
||||
|
|
@ -51,6 +54,17 @@ def new_mock_context(
|
|||
return context
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_duration_ms() -> Generator[float, None]:
|
||||
"""Mock duration_ms."""
|
||||
patched = patch(
|
||||
"fastmcp.server.middleware.logging._get_duration_ms", return_value=0.02
|
||||
)
|
||||
patched.start()
|
||||
yield
|
||||
patched.stop()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_context():
|
||||
"""Create a mock middleware context."""
|
||||
|
|
@ -77,15 +91,14 @@ class TestStructuredLoggingMiddleware:
|
|||
|
||||
def test_init_default(self):
|
||||
"""Test default initialization."""
|
||||
middleware = LoggingMiddleware()
|
||||
middleware = StructuredLoggingMiddleware()
|
||||
|
||||
assert middleware.logger.name == "fastmcp.requests"
|
||||
assert middleware.logger.name == "fastmcp.middleware.structured_logging"
|
||||
assert middleware.log_level == logging.INFO
|
||||
assert middleware.include_payloads is False
|
||||
assert middleware.max_payload_length == 1000
|
||||
assert middleware.include_payload_length is False
|
||||
assert middleware.estimate_payload_tokens is False
|
||||
assert middleware.structured_logging is False
|
||||
assert middleware.structured_logging is True
|
||||
|
||||
def test_init_custom(self):
|
||||
"""Test custom initialization."""
|
||||
|
|
@ -108,14 +121,12 @@ class TestStructuredLoggingMiddleware:
|
|||
"""Test message formatting without payloads."""
|
||||
middleware = StructuredLoggingMiddleware()
|
||||
|
||||
message = middleware._create_before_message(mock_context, "test_event")
|
||||
message = middleware._create_before_message(mock_context)
|
||||
|
||||
assert message == snapshot(
|
||||
{
|
||||
"event": "test_event",
|
||||
"timestamp": "2023-01-01T00:00:00+00:00",
|
||||
"event": "request_start",
|
||||
"source": "client",
|
||||
"type": "request",
|
||||
"method": "test_method",
|
||||
}
|
||||
)
|
||||
|
|
@ -126,14 +137,12 @@ class TestStructuredLoggingMiddleware:
|
|||
"""Test message formatting with payloads."""
|
||||
middleware = StructuredLoggingMiddleware(include_payloads=True)
|
||||
|
||||
message = middleware._create_before_message(mock_context, "test_event")
|
||||
message = middleware._create_before_message(mock_context)
|
||||
|
||||
assert message == snapshot(
|
||||
{
|
||||
"event": "test_event",
|
||||
"timestamp": "2023-01-01T00:00:00+00:00",
|
||||
"event": "request_start",
|
||||
"source": "client",
|
||||
"type": "request",
|
||||
"method": "test_method",
|
||||
"payload": '{"method":"tools/call","params":{"_meta":null,"name":"test_method","arguments":{"param":"value"}}}',
|
||||
"payload_type": "CallToolRequest",
|
||||
|
|
@ -143,14 +152,12 @@ class TestStructuredLoggingMiddleware:
|
|||
def test_calculate_response_size(self, mock_context: MiddlewareContext[Any]):
|
||||
"""Test response size calculation."""
|
||||
middleware = StructuredLoggingMiddleware(include_payload_length=True)
|
||||
message = middleware._create_before_message(mock_context, "test_event")
|
||||
message = middleware._create_before_message(mock_context)
|
||||
|
||||
assert message == snapshot(
|
||||
{
|
||||
"event": "test_event",
|
||||
"timestamp": "2023-01-01T00:00:00+00:00",
|
||||
"event": "request_start",
|
||||
"source": "client",
|
||||
"type": "request",
|
||||
"method": "test_method",
|
||||
"payload_length": 98,
|
||||
}
|
||||
|
|
@ -163,14 +170,12 @@ class TestStructuredLoggingMiddleware:
|
|||
middleware = StructuredLoggingMiddleware(
|
||||
include_payload_length=True, estimate_payload_tokens=True
|
||||
)
|
||||
message = middleware._create_before_message(mock_context, "test_event")
|
||||
message = middleware._create_before_message(mock_context)
|
||||
|
||||
assert message == snapshot(
|
||||
{
|
||||
"event": "test_event",
|
||||
"timestamp": "2023-01-01T00:00:00+00:00",
|
||||
"event": "request_start",
|
||||
"source": "client",
|
||||
"type": "request",
|
||||
"method": "test_method",
|
||||
"payload_tokens": 24,
|
||||
"payload_length": 98,
|
||||
|
|
@ -186,16 +191,18 @@ class TestStructuredLoggingMiddleware:
|
|||
middleware = StructuredLoggingMiddleware()
|
||||
mock_call_next = AsyncMock(return_value="test_result")
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
result = await middleware.on_message(mock_context, mock_call_next)
|
||||
|
||||
assert result == "test_result"
|
||||
assert mock_call_next.called
|
||||
assert remove_line_numbers(caplog.text) == snapshot("""\
|
||||
INFO fastmcp.structured:logging.py:LINE_NUMBER Processing message: {"event": "request_start", "timestamp": "2023-01-01T00:00:00+00:00", "method": "test_method", "type": "request", "source": "client"}
|
||||
INFO fastmcp.structured:logging.py:LINE_NUMBER Completed message: {"event": "request_success", "timestamp": "2023-01-01T00:00:00+00:00", "method": "test_method", "type": "request", "source": "client"}
|
||||
|
||||
""")
|
||||
assert get_log_lines(caplog) == snapshot(
|
||||
[
|
||||
'{"event": "request_start", "method": "test_method", "source": "client"}',
|
||||
'{"event": "request_success", "method": "test_method", "source": "client", "duration_ms": 0.02}',
|
||||
]
|
||||
)
|
||||
|
||||
async def test_on_message_failure(
|
||||
self, mock_context: MiddlewareContext[Any], caplog: pytest.LogCaptureFixture
|
||||
|
|
@ -204,12 +211,16 @@ INFO fastmcp.structured:logging.py:LINE_NUMBER Completed message: {"event":
|
|||
middleware = StructuredLoggingMiddleware()
|
||||
mock_call_next = AsyncMock(side_effect=ValueError("test error"))
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
with pytest.raises(ValueError):
|
||||
await middleware.on_message(mock_context, mock_call_next)
|
||||
|
||||
assert "Processing message:" in caplog.text
|
||||
assert "Failed message: test_method - test error" in caplog.text
|
||||
assert get_log_lines(caplog) == snapshot(
|
||||
[
|
||||
'{"event": "request_start", "method": "test_method", "source": "client"}',
|
||||
'{"event": "request_error", "method": "test_method", "source": "client", "duration_ms": 0.02, "error": "test error"}',
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class TestLoggingMiddleware:
|
||||
|
|
@ -218,7 +229,7 @@ class TestLoggingMiddleware:
|
|||
def test_init_default(self):
|
||||
"""Test default initialization."""
|
||||
middleware = LoggingMiddleware()
|
||||
assert middleware.logger.name == "fastmcp.requests"
|
||||
assert middleware.logger.name == "fastmcp.middleware.logging"
|
||||
assert middleware.log_level == logging.INFO
|
||||
assert middleware.include_payloads is False
|
||||
assert middleware.include_payload_length is False
|
||||
|
|
@ -227,11 +238,11 @@ class TestLoggingMiddleware:
|
|||
def test_format_message(self, mock_context: MiddlewareContext[Any]):
|
||||
"""Test message formatting."""
|
||||
middleware = LoggingMiddleware()
|
||||
message = middleware._create_before_message(mock_context, "test_event")
|
||||
message = middleware._create_before_message(mock_context)
|
||||
formatted = middleware._format_message(message)
|
||||
|
||||
assert formatted == snapshot(
|
||||
"event=test_event timestamp=2023-01-01T00:00:00+00:00 method=test_method type=request source=client"
|
||||
"event=request_start method=test_method source=client"
|
||||
)
|
||||
|
||||
def test_create_before_message_long_payload(
|
||||
|
|
@ -240,12 +251,196 @@ class TestLoggingMiddleware:
|
|||
"""Test message formatting with long payload truncation."""
|
||||
middleware = LoggingMiddleware(include_payloads=True, max_payload_length=10)
|
||||
|
||||
message = middleware._create_before_message(mock_context, "test_event")
|
||||
message = middleware._create_before_message(mock_context)
|
||||
|
||||
formatted = middleware._format_message(message)
|
||||
|
||||
assert "payload=" in formatted
|
||||
assert "..." in formatted
|
||||
assert formatted == snapshot(
|
||||
'event=request_start method=test_method source=client payload={"method":... payload_type=CallToolRequest'
|
||||
)
|
||||
|
||||
async def test_on_message_failure(
|
||||
self, mock_context: MiddlewareContext[Any], caplog: pytest.LogCaptureFixture
|
||||
):
|
||||
"""Test structured logging of failed messages."""
|
||||
middleware = StructuredLoggingMiddleware()
|
||||
mock_call_next = AsyncMock(side_effect=ValueError("test error"))
|
||||
|
||||
with caplog_for_fastmcp(caplog):
|
||||
with pytest.raises(ValueError):
|
||||
await middleware.on_message(mock_context, mock_call_next)
|
||||
|
||||
# Check that we have structured JSON logs
|
||||
assert get_log_lines(caplog) == snapshot(
|
||||
[
|
||||
'{"event": "request_start", "method": "test_method", "source": "client"}',
|
||||
'{"event": "request_error", "method": "test_method", "source": "client", "duration_ms": 0.02, "error": "test error"}',
|
||||
]
|
||||
)
|
||||
|
||||
async def test_on_message_with_pydantic_types_in_payload(
|
||||
self,
|
||||
mock_call_next: CallNext[Any, Any],
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
):
|
||||
"""Ensure Pydantic AnyUrl in payload serializes correctly when include_payloads=True."""
|
||||
|
||||
mock_context = new_mock_context(
|
||||
message=mcp.types.ReadResourceRequest(
|
||||
method="resources/read",
|
||||
params=mcp.types.ReadResourceRequestParams(
|
||||
uri=AnyUrl("test://example/1"),
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
middleware = StructuredLoggingMiddleware(include_payloads=True)
|
||||
|
||||
with caplog_for_fastmcp(caplog):
|
||||
result = await middleware.on_message(mock_context, mock_call_next)
|
||||
|
||||
assert result == "test_result"
|
||||
|
||||
assert get_log_lines(caplog) == snapshot(
|
||||
[
|
||||
'{"event": "request_start", "method": "test_method", "source": "client", "payload": "{\\"method\\":\\"resources/read\\",\\"params\\":{\\"_meta\\":null,\\"uri\\":\\"test://example/1\\"}}", "payload_type": "ReadResourceRequest"}',
|
||||
'{"event": "request_success", "method": "test_method", "source": "client", "duration_ms": 0.02}',
|
||||
]
|
||||
)
|
||||
|
||||
async def test_on_message_with_resource_template_in_payload(
|
||||
self,
|
||||
mock_call_next: CallNext[Any, Any],
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
):
|
||||
"""Ensure ResourceTemplate in payload serializes via pydantic conversion without errors."""
|
||||
|
||||
mock_context = new_mock_context(
|
||||
message=ResourceTemplate(
|
||||
name="tmpl",
|
||||
uri_template="tmpl://{id}",
|
||||
parameters={"id": {"type": "string"}},
|
||||
)
|
||||
)
|
||||
|
||||
middleware = StructuredLoggingMiddleware(include_payloads=True)
|
||||
|
||||
with caplog_for_fastmcp(caplog):
|
||||
result = await middleware.on_message(mock_context, mock_call_next)
|
||||
|
||||
assert result == "test_result"
|
||||
|
||||
assert get_log_lines(caplog) == snapshot(
|
||||
[
|
||||
'{"event": "request_start", "method": "test_method", "source": "client", "payload": "{\\"name\\":\\"tmpl\\",\\"title\\":null,\\"description\\":null,\\"tags\\":[],\\"meta\\":null,\\"enabled\\":true,\\"uri_template\\":\\"tmpl://{id}\\",\\"mime_type\\":\\"text/plain\\",\\"parameters\\":{\\"id\\":{\\"type\\":\\"string\\"}},\\"annotations\\":null}", "payload_type": "ResourceTemplate"}',
|
||||
'{"event": "request_success", "method": "test_method", "source": "client", "duration_ms": 0.02}',
|
||||
]
|
||||
)
|
||||
|
||||
async def test_on_message_with_nonserializable_payload_falls_back_to_str(
|
||||
self, mock_call_next: CallNext[Any, Any], caplog: pytest.LogCaptureFixture
|
||||
):
|
||||
"""Ensure non-JSONable objects fall back to string serialization in payload."""
|
||||
|
||||
class NonSerializable:
|
||||
def __str__(self) -> str:
|
||||
return "NON_SERIALIZABLE"
|
||||
|
||||
mock_context = new_mock_context(
|
||||
message=mcp.types.CallToolRequest(
|
||||
method="tools/call",
|
||||
params=mcp.types.CallToolRequestParams(
|
||||
name="test_method",
|
||||
arguments={"obj": NonSerializable()},
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
middleware = StructuredLoggingMiddleware(include_payloads=True)
|
||||
|
||||
with caplog_for_fastmcp(caplog):
|
||||
result = await middleware.on_message(mock_context, mock_call_next)
|
||||
|
||||
assert result == "test_result"
|
||||
|
||||
assert get_log_lines(caplog) == snapshot(
|
||||
[
|
||||
'{"event": "request_start", "method": "test_method", "source": "client", "payload": "{\\"method\\":\\"tools/call\\",\\"params\\":{\\"_meta\\":null,\\"name\\":\\"test_method\\",\\"arguments\\":{\\"obj\\":\\"NON_SERIALIZABLE\\"}}}", "payload_type": "CallToolRequest"}',
|
||||
'{"event": "request_success", "method": "test_method", "source": "client", "duration_ms": 0.02}',
|
||||
]
|
||||
)
|
||||
|
||||
async def test_on_message_with_custom_serializer_applied(
|
||||
self, mock_call_next: CallNext[Any, Any], caplog: pytest.LogCaptureFixture
|
||||
):
|
||||
"""Ensure a custom serializer is used for non-JSONable payloads."""
|
||||
|
||||
# Provide a serializer that replaces entire payload with a fixed string
|
||||
def custom_serializer(_: Any) -> str:
|
||||
return "CUSTOM_PAYLOAD"
|
||||
|
||||
mock_context = new_mock_context(
|
||||
message=mcp.types.CallToolRequest(
|
||||
method="tools/call",
|
||||
params=mcp.types.CallToolRequestParams(
|
||||
name="test_method",
|
||||
arguments={"obj": "OBJECT"},
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
middleware = StructuredLoggingMiddleware(
|
||||
include_payloads=True, payload_serializer=custom_serializer
|
||||
)
|
||||
|
||||
with caplog_for_fastmcp(caplog):
|
||||
result = await middleware.on_message(mock_context, mock_call_next)
|
||||
|
||||
assert result == "test_result"
|
||||
|
||||
assert get_log_lines(caplog) == snapshot(
|
||||
[
|
||||
'{"event": "request_start", "method": "test_method", "source": "client", "payload": "CUSTOM_PAYLOAD", "payload_type": "CallToolRequest"}',
|
||||
'{"event": "request_success", "method": "test_method", "source": "client", "duration_ms": 0.02}',
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def logging_server():
|
||||
"""Create a FastMCP server specifically for logging middleware tests."""
|
||||
from fastmcp import FastMCP
|
||||
|
||||
mcp = FastMCP("LoggingTestServer")
|
||||
|
||||
@mcp.tool
|
||||
def simple_operation(data: str) -> str:
|
||||
"""A simple operation for testing logging."""
|
||||
return f"Processed: {data}"
|
||||
|
||||
@mcp.tool
|
||||
def complex_operation(items: list[str], mode: str = "default") -> dict:
|
||||
"""A complex operation with structured data."""
|
||||
return {"processed_items": len(items), "mode": mode, "result": "success"}
|
||||
|
||||
@mcp.tool
|
||||
def operation_with_error(should_fail: bool = False) -> str:
|
||||
"""An operation that can be made to fail."""
|
||||
if should_fail:
|
||||
raise ValueError("Operation failed intentionally")
|
||||
return "Operation completed successfully"
|
||||
|
||||
@mcp.resource("log://test")
|
||||
def test_resource() -> str:
|
||||
"""A test resource for logging."""
|
||||
return "Test resource content"
|
||||
|
||||
@mcp.prompt
|
||||
def test_prompt() -> str:
|
||||
"""A test prompt for logging."""
|
||||
return "Test prompt content"
|
||||
|
||||
return mcp
|
||||
|
||||
|
||||
class TestLoggingMiddlewareIntegration:
|
||||
|
|
@ -290,33 +485,29 @@ class TestLoggingMiddlewareIntegration:
|
|||
):
|
||||
"""Test that logging middleware captures successful operations."""
|
||||
logging_middleware = LoggingMiddleware(methods=["tools/call"])
|
||||
logging_middleware._get_timestamp_from_context = ( # ty: ignore[invalid-assignment]
|
||||
lambda _: FIXED_DATE.isoformat()
|
||||
)
|
||||
|
||||
logging_server.add_middleware(logging_middleware)
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
async with Client(logging_server) as client:
|
||||
await client.call_tool(
|
||||
name="simple_operation", arguments={"data": "test_data"}
|
||||
)
|
||||
await client.call_tool(
|
||||
name="complex_operation",
|
||||
arguments={"items": ["a", "b", "c"], "mode": "batch"},
|
||||
)
|
||||
with caplog_for_fastmcp(caplog):
|
||||
with caplog.at_level(logging.INFO):
|
||||
async with Client(logging_server) as client:
|
||||
await client.call_tool(
|
||||
name="simple_operation", arguments={"data": "test_data"}
|
||||
)
|
||||
await client.call_tool(
|
||||
name="complex_operation",
|
||||
arguments={"items": ["a", "b", "c"], "mode": "batch"},
|
||||
)
|
||||
|
||||
# Should have processing and completion logs for both operations
|
||||
assert remove_line_numbers(caplog.text) == snapshot("""\
|
||||
INFO mcp.server.lowlevel.server:server.py:LINE_NUMBER Processing request of type CallToolRequest
|
||||
INFO fastmcp.requests:logging.py:LINE_NUMBER Processing message: event=request_start timestamp=2023-01-01T00:00:00+00:00 method=tools/call type=request source=client
|
||||
INFO fastmcp.requests:logging.py:LINE_NUMBER Completed message: event=request_success timestamp=2023-01-01T00:00:00+00:00 method=tools/call type=request source=client
|
||||
INFO mcp.server.lowlevel.server:server.py:LINE_NUMBER Processing request of type ListToolsRequest
|
||||
INFO mcp.server.lowlevel.server:server.py:LINE_NUMBER Processing request of type CallToolRequest
|
||||
INFO fastmcp.requests:logging.py:LINE_NUMBER Processing message: event=request_start timestamp=2023-01-01T00:00:00+00:00 method=tools/call type=request source=client
|
||||
INFO fastmcp.requests:logging.py:LINE_NUMBER Completed message: event=request_success timestamp=2023-01-01T00:00:00+00:00 method=tools/call type=request source=client
|
||||
|
||||
""")
|
||||
assert get_log_lines(caplog) == snapshot(
|
||||
[
|
||||
"event=request_start method=tools/call source=client",
|
||||
"event=request_success method=tools/call source=client duration_ms=0.02",
|
||||
"event=request_start method=tools/call source=client",
|
||||
"event=request_success method=tools/call source=client duration_ms=0.02",
|
||||
]
|
||||
)
|
||||
|
||||
async def test_logging_middleware_logs_failures(
|
||||
self, logging_server: FastMCP, caplog: pytest.LogCaptureFixture
|
||||
|
|
@ -324,7 +515,7 @@ INFO fastmcp.requests:logging.py:LINE_NUMBER Completed message: event=reques
|
|||
"""Test that logging middleware captures failed operations."""
|
||||
logging_server.add_middleware(LoggingMiddleware(methods=["tools/call"]))
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
async with Client(logging_server) as client:
|
||||
# This should fail and be logged
|
||||
with pytest.raises(Exception):
|
||||
|
|
@ -335,8 +526,9 @@ INFO fastmcp.requests:logging.py:LINE_NUMBER Completed message: event=reques
|
|||
log_text = caplog.text
|
||||
|
||||
# Should have processing and failure logs
|
||||
assert "Processing message:" in log_text
|
||||
assert "Failed message: tools/call" in log_text
|
||||
assert log_text.splitlines()[-1] == snapshot(
|
||||
"ERROR fastmcp.middleware.logging:logging.py:122 event=request_error method=tools/call source=client duration_ms=0.02 error=Error calling tool 'operation_with_error': Operation failed intentionally"
|
||||
)
|
||||
|
||||
async def test_logging_middleware_with_payloads(
|
||||
self, logging_server: FastMCP, caplog: pytest.LogCaptureFixture
|
||||
|
|
@ -346,24 +538,18 @@ INFO fastmcp.requests:logging.py:LINE_NUMBER Completed message: event=reques
|
|||
middleware = LoggingMiddleware(
|
||||
include_payloads=True, max_payload_length=500, methods=["tools/call"]
|
||||
)
|
||||
middleware._get_timestamp_from_context = ( # ty: ignore[invalid-assignment]
|
||||
lambda _: FIXED_DATE.isoformat()
|
||||
)
|
||||
logging_server.add_middleware(middleware)
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
async with Client(logging_server) as client:
|
||||
await client.call_tool("simple_operation", {"data": "payload_test"})
|
||||
|
||||
log_text = caplog.text
|
||||
|
||||
assert remove_line_numbers(log_text) == snapshot("""\
|
||||
INFO mcp.server.lowlevel.server:server.py:LINE_NUMBER Processing request of type CallToolRequest
|
||||
INFO fastmcp.requests:logging.py:LINE_NUMBER Processing message: event=request_start timestamp=2023-01-01T00:00:00+00:00 method=tools/call type=request source=client payload={"_meta":null,"name":"simple_operation","arguments":{"data":"payload_test"}} payload_type=CallToolRequestParams
|
||||
INFO fastmcp.requests:logging.py:LINE_NUMBER Completed message: event=request_success timestamp=2023-01-01T00:00:00+00:00 method=tools/call type=request source=client
|
||||
INFO mcp.server.lowlevel.server:server.py:LINE_NUMBER Processing request of type ListToolsRequest
|
||||
|
||||
""")
|
||||
assert get_log_lines(caplog) == snapshot(
|
||||
[
|
||||
'event=request_start method=tools/call source=client payload={"_meta":null,"name":"simple_operation","arguments":{"data":"payload_test"}} payload_type=CallToolRequestParams',
|
||||
"event=request_success method=tools/call source=client duration_ms=0.02",
|
||||
]
|
||||
)
|
||||
|
||||
async def test_structured_logging_middleware_produces_json(
|
||||
self, logging_server: FastMCP, caplog: pytest.LogCaptureFixture
|
||||
|
|
@ -373,34 +559,21 @@ INFO mcp.server.lowlevel.server:server.py:LINE_NUMBER Processing request of
|
|||
logging_middleware = StructuredLoggingMiddleware(
|
||||
include_payloads=True, methods=["tools/call"]
|
||||
)
|
||||
logging_middleware._get_timestamp_from_context = ( # ty: ignore[invalid-assignment]
|
||||
lambda _: FIXED_DATE.isoformat()
|
||||
)
|
||||
|
||||
logging_server.add_middleware(logging_middleware)
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
async with Client(logging_server) as client:
|
||||
await client.call_tool(
|
||||
name="simple_operation", arguments={"data": "json_test"}
|
||||
)
|
||||
|
||||
# Extract JSON log entries
|
||||
log_lines = [
|
||||
record.message
|
||||
for record in caplog.records
|
||||
if record.name == "fastmcp.structured"
|
||||
]
|
||||
|
||||
assert len(log_lines) >= 2 # Should have start and success entries
|
||||
|
||||
assert remove_line_numbers(caplog.text) == snapshot("""\
|
||||
INFO mcp.server.lowlevel.server:server.py:LINE_NUMBER Processing request of type CallToolRequest
|
||||
INFO fastmcp.structured:logging.py:LINE_NUMBER Processing message: {"event": "request_start", "timestamp": "2023-01-01T00:00:00+00:00", "method": "tools/call", "type": "request", "source": "client", "payload": "{\\"_meta\\":null,\\"name\\":\\"simple_operation\\",\\"arguments\\":{\\"data\\":\\"json_test\\"}}", "payload_type": "CallToolRequestParams"}
|
||||
INFO fastmcp.structured:logging.py:LINE_NUMBER Completed message: {"event": "request_success", "timestamp": "2023-01-01T00:00:00+00:00", "method": "tools/call", "type": "request", "source": "client"}
|
||||
INFO mcp.server.lowlevel.server:server.py:LINE_NUMBER Processing request of type ListToolsRequest
|
||||
|
||||
""")
|
||||
assert get_log_lines(caplog) == snapshot(
|
||||
[
|
||||
'{"event": "request_start", "method": "tools/call", "source": "client", "payload": "{\\"_meta\\":null,\\"name\\":\\"simple_operation\\",\\"arguments\\":{\\"data\\":\\"json_test\\"}}", "payload_type": "CallToolRequestParams"}',
|
||||
'{"event": "request_success", "method": "tools/call", "source": "client", "duration_ms": 0.02}',
|
||||
]
|
||||
)
|
||||
|
||||
async def test_structured_logging_middleware_handles_errors(
|
||||
self, logging_server: FastMCP, caplog: pytest.LogCaptureFixture
|
||||
|
|
@ -408,25 +581,23 @@ INFO mcp.server.lowlevel.server:server.py:LINE_NUMBER Processing request of
|
|||
"""Test structured logging of errors with JSON format."""
|
||||
|
||||
logging_middleware = StructuredLoggingMiddleware(methods=["tools/call"])
|
||||
logging_middleware._get_timestamp_from_context = ( # ty: ignore[invalid-assignment]
|
||||
lambda _: FIXED_DATE.isoformat()
|
||||
)
|
||||
|
||||
logging_server.add_middleware(logging_middleware)
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
async with Client(logging_server) as client:
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool(
|
||||
"operation_with_error", {"should_fail": True}
|
||||
)
|
||||
with caplog_for_fastmcp(caplog):
|
||||
with caplog.at_level(logging.INFO):
|
||||
async with Client(logging_server) as client:
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool(
|
||||
"operation_with_error", {"should_fail": True}
|
||||
)
|
||||
|
||||
assert remove_line_numbers(caplog.text) == snapshot("""\
|
||||
INFO mcp.server.lowlevel.server:server.py:LINE_NUMBER Processing request of type CallToolRequest
|
||||
INFO fastmcp.structured:logging.py:LINE_NUMBER Processing message: {"event": "request_start", "timestamp": "2023-01-01T00:00:00+00:00", "method": "tools/call", "type": "request", "source": "client"}
|
||||
ERROR fastmcp.structured:logging.py:LINE_NUMBER Failed message: tools/call - Error calling tool 'operation_with_error': Operation failed intentionally
|
||||
|
||||
""")
|
||||
assert get_log_lines(caplog) == snapshot(
|
||||
[
|
||||
'{"event": "request_start", "method": "tools/call", "source": "client"}',
|
||||
'{"event": "request_error", "method": "tools/call", "source": "client", "duration_ms": 0.02, "error": "Error calling tool \'operation_with_error\': Operation failed intentionally"}',
|
||||
]
|
||||
)
|
||||
|
||||
async def test_logging_middleware_with_different_operations(
|
||||
self, logging_server: FastMCP, caplog: pytest.LogCaptureFixture
|
||||
|
|
@ -444,7 +615,7 @@ ERROR fastmcp.structured:logging.py:LINE_NUMBER Failed message: tools/call -
|
|||
)
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
async with Client(logging_server) as client:
|
||||
# Test different operation types
|
||||
await client.call_tool("simple_operation", {"data": "test"})
|
||||
|
|
@ -452,16 +623,18 @@ ERROR fastmcp.structured:logging.py:LINE_NUMBER Failed message: tools/call -
|
|||
await client.get_prompt("test_prompt")
|
||||
await client.list_resources()
|
||||
|
||||
log_text = caplog.text
|
||||
|
||||
# Should have logs for all different operation types
|
||||
# Note: Different operations may have different method names
|
||||
processing_count = log_text.count("Processing message:")
|
||||
completion_count = log_text.count("Completed message:")
|
||||
|
||||
# Should have processed all 4 operations
|
||||
assert processing_count == 4
|
||||
assert completion_count == 4
|
||||
assert get_log_lines(caplog) == snapshot(
|
||||
[
|
||||
"event=request_start method=tools/call source=client",
|
||||
"event=request_success method=tools/call source=client duration_ms=0.02",
|
||||
"event=request_start method=resources/read source=client",
|
||||
"event=request_success method=resources/read source=client duration_ms=0.02",
|
||||
"event=request_start method=prompts/get source=client",
|
||||
"event=request_success method=prompts/get source=client duration_ms=0.02",
|
||||
"event=request_start method=resources/list source=client",
|
||||
"event=request_success method=resources/list source=client duration_ms=0.02",
|
||||
]
|
||||
)
|
||||
|
||||
async def test_logging_middleware_custom_configuration(
|
||||
self, logging_server: FastMCP
|
||||
|
|
@ -491,5 +664,7 @@ ERROR fastmcp.structured:logging.py:LINE_NUMBER Failed message: tools/call -
|
|||
|
||||
# Check that our custom logger captured the logs
|
||||
log_output = log_buffer.getvalue()
|
||||
assert "Processing message:" in log_output
|
||||
assert "payload=" in log_output
|
||||
assert log_output == snapshot("""\
|
||||
event=request_start method=tools/call source=client payload={"_meta":null,"name":"simple_operation","arguments":{"data":"custom_test"}} payload_type=CallToolRequestParams
|
||||
event=request_success method=tools/call source=client duration_ms=0.02
|
||||
""")
|
||||
|
|
|
|||
|
|
@ -293,6 +293,17 @@ class TestMiddlewareHooks:
|
|||
result = list_prompts_calls[0].result
|
||||
assert isinstance(result, list)
|
||||
|
||||
async def test_initialize(
|
||||
self, mcp_server: FastMCP, recording_middleware: RecordingMiddleware
|
||||
):
|
||||
async with Client(mcp_server) as client:
|
||||
await client.ping()
|
||||
|
||||
assert recording_middleware.assert_called(at_least=1)
|
||||
assert recording_middleware.assert_called(hook="on_message", at_least=1)
|
||||
assert recording_middleware.assert_called(hook="on_request", at_least=1)
|
||||
assert recording_middleware.assert_called(hook="on_initialize", at_least=1)
|
||||
|
||||
async def test_list_tools_filtering_middleware(self):
|
||||
"""Test that middleware can filter tools."""
|
||||
|
||||
|
|
|
|||
|
|
@ -306,9 +306,10 @@ class TestRateLimitingMiddlewareIntegration:
|
|||
|
||||
async def test_rate_limiting_blocks_rapid_requests(self, rate_limit_server):
|
||||
"""Test that rate limiting blocks rapid successive requests."""
|
||||
# Very restrictive rate limit (accounting for extra list_tools calls per tool call)
|
||||
# Very restrictive rate limit (accounting for initialization and list_tools calls)
|
||||
# Requests: 1 initialize + 1 list_tools + 4 call_tools = 6 total before limit
|
||||
rate_limit_server.add_middleware(
|
||||
RateLimitingMiddleware(max_requests_per_second=10.0, burst_capacity=5)
|
||||
RateLimitingMiddleware(max_requests_per_second=10.0, burst_capacity=6)
|
||||
)
|
||||
|
||||
async with Client(rate_limit_server) as client:
|
||||
|
|
@ -356,7 +357,7 @@ class TestRateLimitingMiddlewareIntegration:
|
|||
"""Test sliding window rate limiting implementation."""
|
||||
rate_limit_server.add_middleware(
|
||||
SlidingWindowRateLimitingMiddleware(
|
||||
max_requests=5, # Accounting for extra list_tools calls
|
||||
max_requests=6, # 1 init + 1 list_tools + 3 calls + 1 to fail
|
||||
window_minutes=1, # 1-minute window
|
||||
)
|
||||
)
|
||||
|
|
@ -374,7 +375,7 @@ class TestRateLimitingMiddlewareIntegration:
|
|||
async def test_rate_limiting_with_different_operations(self, rate_limit_server):
|
||||
"""Test that rate limiting applies to all types of operations."""
|
||||
rate_limit_server.add_middleware(
|
||||
RateLimitingMiddleware(max_requests_per_second=9.0, burst_capacity=4)
|
||||
RateLimitingMiddleware(max_requests_per_second=9.0, burst_capacity=5)
|
||||
)
|
||||
|
||||
async with Client(rate_limit_server) as client:
|
||||
|
|
@ -395,8 +396,8 @@ class TestRateLimitingMiddlewareIntegration:
|
|||
|
||||
rate_limit_server.add_middleware(
|
||||
RateLimitingMiddleware(
|
||||
max_requests_per_second=6.0, # Accounting for extra list_tools calls
|
||||
burst_capacity=3,
|
||||
max_requests_per_second=6.0, # Accounting for initialization and list_tools calls
|
||||
burst_capacity=4,
|
||||
get_client_id=get_client_id,
|
||||
)
|
||||
)
|
||||
|
|
@ -416,8 +417,8 @@ class TestRateLimitingMiddlewareIntegration:
|
|||
rate_limit_server.add_middleware(
|
||||
RateLimitingMiddleware(
|
||||
max_requests_per_second=6.0,
|
||||
burst_capacity=4,
|
||||
global_limit=True, # Accounting for extra list_tools calls
|
||||
burst_capacity=5, # 1 init + 2 list_tools + 2 calls before limit
|
||||
global_limit=True, # Accounting for initialization and list_tools calls
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -435,7 +436,7 @@ class TestRateLimitingMiddlewareIntegration:
|
|||
rate_limit_server.add_middleware(
|
||||
RateLimitingMiddleware(
|
||||
max_requests_per_second=10.0, # 10 per second = 1 every 100ms
|
||||
burst_capacity=3,
|
||||
burst_capacity=4,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from fastmcp import FastMCP
|
|||
from fastmcp.client import Client
|
||||
from fastmcp.server.middleware.middleware import MiddlewareContext
|
||||
from fastmcp.server.middleware.timing import DetailedTimingMiddleware, TimingMiddleware
|
||||
from fastmcp.utilities.tests import caplog_for_fastmcp
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -47,7 +48,7 @@ class TestTimingMiddleware:
|
|||
"""Test timing successful requests."""
|
||||
middleware = TimingMiddleware()
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
result = await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
assert result == "test_result"
|
||||
|
|
@ -60,7 +61,7 @@ class TestTimingMiddleware:
|
|||
middleware = TimingMiddleware()
|
||||
mock_call_next = AsyncMock(side_effect=ValueError("test error"))
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
with pytest.raises(ValueError):
|
||||
await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
|
|
@ -84,7 +85,7 @@ class TestDetailedTimingMiddleware:
|
|||
context.message.name = "test_tool"
|
||||
mock_call_next = AsyncMock(return_value="tool_result")
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
result = await middleware.on_call_tool(context, mock_call_next)
|
||||
|
||||
assert result == "tool_result"
|
||||
|
|
@ -97,7 +98,7 @@ class TestDetailedTimingMiddleware:
|
|||
context.message.uri = "test://resource"
|
||||
mock_call_next = AsyncMock(return_value="resource_result")
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
result = await middleware.on_read_resource(context, mock_call_next)
|
||||
|
||||
assert result == "resource_result"
|
||||
|
|
@ -110,7 +111,7 @@ class TestDetailedTimingMiddleware:
|
|||
context.message.name = "test_prompt"
|
||||
mock_call_next = AsyncMock(return_value="prompt_result")
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
result = await middleware.on_get_prompt(context, mock_call_next)
|
||||
|
||||
assert result == "prompt_result"
|
||||
|
|
@ -122,7 +123,7 @@ class TestDetailedTimingMiddleware:
|
|||
context = MagicMock()
|
||||
mock_call_next = AsyncMock(return_value="tools_result")
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
result = await middleware.on_list_tools(context, mock_call_next)
|
||||
|
||||
assert result == "tools_result"
|
||||
|
|
@ -135,7 +136,7 @@ class TestDetailedTimingMiddleware:
|
|||
context.message.name = "failing_tool"
|
||||
mock_call_next = AsyncMock(side_effect=RuntimeError("operation failed"))
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
with pytest.raises(RuntimeError):
|
||||
await middleware.on_call_tool(context, mock_call_next)
|
||||
|
||||
|
|
@ -155,15 +156,15 @@ def timing_server():
|
|||
|
||||
@mcp.tool
|
||||
def short_task() -> str:
|
||||
"""A task that takes 0.1 seconds."""
|
||||
time.sleep(0.1)
|
||||
return "Done after 0.1s"
|
||||
"""A task that takes 0.01 seconds."""
|
||||
time.sleep(0.01)
|
||||
return "Done after 0.01 seconds"
|
||||
|
||||
@mcp.tool
|
||||
def medium_task() -> str:
|
||||
"""A task that takes 0.15 seconds."""
|
||||
time.sleep(0.15)
|
||||
return "Done after 0.15s"
|
||||
"""A task that takes 0.02 seconds."""
|
||||
time.sleep(0.02)
|
||||
return "Done after 0.02 seconds"
|
||||
|
||||
@mcp.tool
|
||||
def failing_task() -> str:
|
||||
|
|
@ -173,14 +174,14 @@ def timing_server():
|
|||
@mcp.resource("timer://test")
|
||||
def test_resource() -> str:
|
||||
"""A resource that takes time to read."""
|
||||
time.sleep(0.05)
|
||||
return "Resource content after 0.05s"
|
||||
time.sleep(0.005)
|
||||
return "Resource content after 0.005 seconds"
|
||||
|
||||
@mcp.prompt
|
||||
def test_prompt() -> str:
|
||||
"""A prompt that takes time to generate."""
|
||||
time.sleep(0.08)
|
||||
return "Prompt content after 0.08s"
|
||||
time.sleep(0.008)
|
||||
return "Prompt content after 0.008 seconds"
|
||||
|
||||
return mcp
|
||||
|
||||
|
|
@ -194,7 +195,7 @@ class TestTimingMiddlewareIntegration:
|
|||
"""Test that timing middleware accurately measures tool execution times."""
|
||||
timing_server.add_middleware(TimingMiddleware())
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
async with Client(timing_server) as client:
|
||||
# Test instant task
|
||||
await client.call_tool("instant_task")
|
||||
|
|
@ -225,7 +226,7 @@ class TestTimingMiddlewareIntegration:
|
|||
"""Test that timing middleware measures time even for failed operations."""
|
||||
timing_server.add_middleware(TimingMiddleware())
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
async with Client(timing_server) as client:
|
||||
# This should fail but still be timed
|
||||
with pytest.raises(Exception):
|
||||
|
|
@ -241,7 +242,7 @@ class TestTimingMiddlewareIntegration:
|
|||
"""Test that detailed timing middleware provides operation-specific timing."""
|
||||
timing_server.add_middleware(DetailedTimingMiddleware())
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
async with Client(timing_server) as client:
|
||||
# Test tool call
|
||||
await client.call_tool("short_task")
|
||||
|
|
@ -271,7 +272,7 @@ class TestTimingMiddlewareIntegration:
|
|||
"""Test timing middleware with concurrent operations."""
|
||||
timing_server.add_middleware(TimingMiddleware())
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with caplog_for_fastmcp(caplog):
|
||||
async with Client(timing_server) as client:
|
||||
# Run multiple operations concurrently
|
||||
tasks = [
|
||||
|
|
@ -290,7 +291,7 @@ class TestTimingMiddlewareIntegration:
|
|||
len(timing_logs) >= 3
|
||||
) # At least 3 tool calls, may have additional list_tools calls
|
||||
|
||||
async def test_timing_middleware_custom_logger(self, timing_server):
|
||||
async def test_timing_middleware_custom_logger(self, timing_server, caplog):
|
||||
"""Test timing middleware with custom logger configuration."""
|
||||
import io
|
||||
import logging
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue