Use custom OTelMiddleware with either Azure or Logfire

Author: pamelafox

pyproject.toml

@@ -18,6 +18,7 @@ dependencies = [
     "azure-cosmos>=4.9.0",
     "azure-monitor-opentelemetry>=1.6.4",
     "logfire>=3.11.0",
+    "azure-core-tracing-opentelemetry>=1.0.0b12"
 ]
 
 [dependency-groups]

servers/__init__.py

@@ -6,43 +6,31 @@
 from typing import Annotated
 
 import logfire
+from azure.core.settings import settings
 from azure.cosmos.aio import CosmosClient
 from azure.identity.aio import DefaultAzureCredential, ManagedIdentityCredential
 from azure.monitor.opentelemetry import configure_azure_monitor
-from azure.monitor.opentelemetry.exporter import AzureMonitorTraceExporter
 from dotenv import load_dotenv
 from fastmcp import FastMCP
-from opentelemetry import trace
-from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry_middleware import OpenTelemetryMiddleware
 
 RUNNING_IN_PRODUCTION = os.getenv("RUNNING_IN_PRODUCTION", "false").lower() == "true"
 
 if not RUNNING_IN_PRODUCTION:
     load_dotenv(override=True)
 
-logging.basicConfig(level=logging.WARNING, format="%(asctime)s - %(message)s")
+logging.basicConfig(level=logging.DEBUG, format="%(asctime)s - %(message)s")
 logger = logging.getLogger("ExpensesMCP")
 logger.setLevel(logging.INFO)
 
 # Configure OpenTelemetry tracing
-APPLICATIONINSIGHTS_CONNECTION_STRING = os.getenv("APPLICATIONINSIGHTS_CONNECTION_STRING")
-if APPLICATIONINSIGHTS_CONNECTION_STRING:
+if os.getenv("APPLICATIONINSIGHTS_CONNECTION_STRING"):
     logger.info("Setting up Azure Monitor instrumentation")
     configure_azure_monitor()
-
-# Use Logfire to instrument MCP tool calls
-logfire.configure(
-    service_name="expenses-mcp",
-    send_to_logfire=False,  # Send spans to Application Insights, not Logfire backend
-)
-logfire.instrument_mcp()
-
-if APPLICATIONINSIGHTS_CONNECTION_STRING:
-    logger.info("Adding Azure Monitor Trace Exporter to TracerProvider")
-    tracer_provider = trace.get_tracer_provider()
-    exporter = AzureMonitorTraceExporter(connection_string=APPLICATIONINSIGHTS_CONNECTION_STRING)
-    span_processor = BatchSpanProcessor(exporter)
-    tracer_provider.add_span_processor(span_processor)
+if os.getenv("LOGFIRE_PROJECT_NAME"):
+    logger.info("Setting up Logfire instrumentation")
+    settings.tracing_implementation = "opentelemetry"  # Send Azure Monitor traces via OpenTelemetry
+    logfire.configure(service_name="expenses-mcp", send_to_logfire=True)
 
 # Cosmos DB configuration from environment variables
 AZURE_COSMOSDB_ACCOUNT = os.environ["AZURE_COSMOSDB_ACCOUNT"]
@@ -65,6 +53,7 @@
 logger.info(f"Connected to Cosmos DB: {AZURE_COSMOSDB_ACCOUNT}")
 
 mcp = FastMCP("Expenses Tracker")
+mcp.add_middleware(OpenTelemetryMiddleware("ExpensesMCP"))
 
 
 class PaymentMethod(Enum):

servers/deployed_mcp.py

@@ -0,0 +1,83 @@
+from fastmcp.server.middleware import Middleware, MiddlewareContext
+from opentelemetry import trace
+from opentelemetry.trace import Status, StatusCode
+
+
+class OpenTelemetryMiddleware(Middleware):
+    """Middleware that creates OpenTelemetry spans for MCP operations."""
+
+    def __init__(self, tracer_name: str):
+        self.tracer = trace.get_tracer(tracer_name)
+
+    async def on_call_tool(self, context: MiddlewareContext, call_next):
+        """Create a span for each tool call with detailed attributes."""
+        tool_name = context.message.name
+
+        with self.tracer.start_as_current_span(
+            f"tool.{tool_name}",
+            attributes={
+                "mcp.method": context.method,
+                "mcp.source": context.source,
+                "mcp.tool.name": tool_name,
+                "mcp.tool.arguments": str(context.message.arguments),
+            },
+        ) as span:
+            try:
+                result = await call_next(context)
+                span.set_attribute("mcp.tool.success", True)
+                span.set_status(Status(StatusCode.OK))
+                return result
+            except Exception as e:
+                span.set_attribute("mcp.tool.success", False)
+                span.set_attribute("mcp.tool.error", str(e))
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                raise
+
+    async def on_read_resource(self, context: MiddlewareContext, call_next):
+        """Create a span for each resource read."""
+        resource_uri = str(getattr(context.message, "uri", "unknown"))
+
+        with self.tracer.start_as_current_span(
+            f"resource.{resource_uri}",
+            attributes={
+                "mcp.method": context.method,
+                "mcp.source": context.source,
+                "mcp.resource.uri": resource_uri,
+            },
+        ) as span:
+            try:
+                result = await call_next(context)
+                span.set_attribute("mcp.resource.success", True)
+                span.set_status(Status(StatusCode.OK))
+                return result
+            except Exception as e:
+                span.set_attribute("mcp.resource.success", False)
+                span.set_attribute("mcp.resource.error", str(e))
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                raise
+
+    async def on_get_prompt(self, context: MiddlewareContext, call_next):
+        """Create a span for each prompt retrieval."""
+        prompt_name = getattr(context.message, "name", "unknown")
+
+        with self.tracer.start_as_current_span(
+            f"prompt.{prompt_name}",
+            attributes={
+                "mcp.method": context.method,
+                "mcp.source": context.source,
+                "mcp.prompt.name": prompt_name,
+            },
+        ) as span:
+            try:
+                result = await call_next(context)
+                span.set_attribute("mcp.prompt.success", True)
+                span.set_status(Status(StatusCode.OK))
+                return result
+            except Exception as e:
+                span.set_attribute("mcp.prompt.success", False)
+                span.set_attribute("mcp.prompt.error", str(e))
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                span.record_exception(e)
+                raise