Added input/output types, further code cleanup, got rid of pylint disable self use

Johnnyl202 · Johnnyl202 · commit 3f576c29b250 · 2025-07-25T00:49:34.000-07:00
diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/mcpinstrumentor/README.md b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/mcpinstrumentor/README.md
@@ -1,11 +1,11 @@
 # MCP Instrumentor
 
-OpenTelemetry MCP instrumentation package.
+OpenTelemetry MCP instrumentation package for AWS Distro.
 
 ## Installation
 
 ```bash
-pip install mcpinstrumentor
+pip install amazon-opentelemetry-distro-mcpinstrumentor
 ```
 
 ## Usage
@@ -14,4 +14,20 @@ pip install mcpinstrumentor
 from mcpinstrumentor import MCPInstrumentor
 
 MCPInstrumentor().instrument()
-```
+```
+
+## Configuration
+
+### Environment Variables
+
+- `MCP_SERVICE_NAME`: Sets the service name for MCP client spans. Defaults to "Generic MCP Server" if not set.
+
+```bash
+export MCP_SERVICE_NAME="My Custom MCP Server"
+```
+
+## Features
+
+- Automatic instrumentation of MCP client and server requests
+- Distributed tracing support with trace context propagation
+- Configurable service naming via environment variables
diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/mcpinstrumentor/mcpinstrumentor.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/mcpinstrumentor/mcpinstrumentor.py
@@ -1,8 +1,9 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 import logging
-from typing import Any, Collection
+from typing import Any, Callable, Collection, Dict, Tuple
 
+from mcp import ClientRequest
 from wrapt import register_post_import_hook, wrap_function_wrapper
 
 from opentelemetry import trace
@@ -29,12 +30,46 @@ class MCPInstrumentor(BaseInstrumentor):
     An instrumenter for MCP.
     """
 
-    def __init__(self):  # pylint: disable=no-self-use
+    def __init__(self):
         super().__init__()
         self.tracer = None
 
+    @staticmethod
+    def instrumentation_dependencies() -> Collection[str]:
+        return _instruments
+
+    def _instrument(self, **kwargs: Any) -> None:
+        tracer_provider = kwargs.get("tracer_provider")
+        if tracer_provider:
+            self.tracer = tracer_provider.get_tracer("mcp")
+        else:
+            self.tracer = trace.get_tracer("mcp")
+        register_post_import_hook(
+            lambda _: wrap_function_wrapper(
+                "mcp.shared.session",
+                "BaseSession.send_request",
+                self._wrap_send_request,
+            ),
+            "mcp.shared.session",
+        )
+        register_post_import_hook(
+            lambda _: wrap_function_wrapper(
+                "mcp.server.lowlevel.server",
+                "Server._handle_request",
+                self._wrap_handle_request,
+            ),
+            "mcp.server.lowlevel.server",
+        )
+
+    @staticmethod
+    def _uninstrument(**kwargs: Any) -> None:
+        unwrap("mcp.shared.session", "BaseSession.send_request")
+        unwrap("mcp.server.lowlevel.server", "Server._handle_request")
+
     # Send Request Wrapper
-    def _wrap_send_request(self, wrapped, instance, args, kwargs):  # pylint: disable=no-self-use
+    def _wrap_send_request(
+        self, wrapped: Callable, instance: Any, args: Tuple[Any, ...], kwargs: Dict[str, Any]
+    ) -> Callable:
         """
         Changes made:
             The wrapper intercepts the request before sending, injects distributed tracing context into the
@@ -43,14 +78,14 @@ def _wrap_send_request(self, wrapped, instance, args, kwargs):  # pylint: disabl
             type and calling the original function with identical parameters.
         """
 
-        async def async_wrapper():  # pylint: disable=no-self-use
+        async def async_wrapper():
             with self.tracer.start_as_current_span("client.send_request", kind=trace.SpanKind.CLIENT) as span:
                 span_ctx = span.get_span_context()
                 request = args[0] if len(args) > 0 else kwargs.get("request")
                 if request:
                     req_root = request.root if hasattr(request, "root") else request
 
-                    self.handle_attributes(span, req_root, True)
+                    self._generate_mcp_attributes(span, req_root, is_client=True)
                     request_data = request.model_dump(by_alias=True, mode="json", exclude_none=True)
                     self._inject_trace_context(request_data, span_ctx)
                     # Reconstruct request object with injected trace context
@@ -68,7 +103,9 @@ async def async_wrapper():  # pylint: disable=no-self-use
         return async_wrapper()
 
     # Handle Request Wrapper
-    async def _wrap_handle_request(self, wrapped, instance, args, kwargs):  # pylint: disable=no-self-use
+    async def _wrap_handle_request(
+        self, wrapped: Callable, instance: Any, args: Tuple[Any, ...], kwargs: Dict[str, Any]
+    ) -> Any:
         """
         Changes made:
         This wrapper intercepts requests before processing, extracts distributed tracing context from
@@ -87,19 +124,35 @@ async def _wrap_handle_request(self, wrapped, instance, args, kwargs):  # pylint
             traceparent = getattr(req.params.meta, "traceparent", None)
         span_context = self._extract_span_context_from_traceparent(traceparent) if traceparent else None
         if span_context:
-            span_name = self._get_span_name(req)
+            span_name = self._get_mcp_operation(req)
             with self.tracer.start_as_current_span(
                 span_name,
                 kind=trace.SpanKind.SERVER,
                 context=trace.set_span_in_context(trace.NonRecordingSpan(span_context)),
             ) as span:
-                self.handle_attributes(span, req, False)
+                self._generate_mcp_attributes(span, req, False)
                 result = await wrapped(*args, **kwargs)
                 return result
         else:
             return await wrapped(*args, **kwargs)
 
-    def _inject_trace_context(self, request_data, span_ctx):  # pylint: disable=no-self-use
+    def _generate_mcp_attributes(self, span: trace.Span, request: ClientRequest, is_client: bool) -> None:
+        import mcp.types as types  # pylint: disable=import-outside-toplevel,consider-using-from-import
+
+        operation = "UnknownOperation"
+        if isinstance(request, types.ListToolsRequest):
+            operation = "ListTool"
+            span.set_attribute("mcp.list_tools", True)
+        elif isinstance(request, types.CallToolRequest):
+            operation = request.params.name
+            span.set_attribute("mcp.call_tool", True)
+        if is_client:
+            self._add_client_attributes(span, operation, request)
+        else:
+            self._add_server_attributes(span, operation, request)
+
+    @staticmethod
+    def _inject_trace_context(request_data: Dict[str, Any], span_ctx) -> None:
         if "params" not in request_data:
             request_data["params"] = {}
         if "_meta" not in request_data["params"]:
@@ -110,7 +163,8 @@ def _inject_trace_context(self, request_data, span_ctx):  # pylint: disable=no-s
         traceparent = f"00-{trace_id_hex}-{span_id_hex}-{trace_flags}"
         request_data["params"]["_meta"]["traceparent"] = traceparent
 
-    def _extract_span_context_from_traceparent(self, traceparent):  # pylint: disable=no-self-use
+    @staticmethod
+    def _extract_span_context_from_traceparent(traceparent: str):
         parts = traceparent.split("-")
         if len(parts) == 4:
             try:
@@ -127,72 +181,29 @@ def _extract_span_context_from_traceparent(self, traceparent):  # pylint: disabl
                 return None
         return None
 
-    def _get_span_name(self, req):  # pylint: disable=no-self-use
-        span_name = "unknown"
+    @staticmethod
+    def _get_mcp_operation(req: ClientRequest) -> str:
         import mcp.types as types  # pylint: disable=import-outside-toplevel,consider-using-from-import
 
+        span_name = "unknown"
+
         if isinstance(req, types.ListToolsRequest):
             span_name = "tools/list"
         elif isinstance(req, types.CallToolRequest):
-            if hasattr(req, "params") and hasattr(req.params, "name"):
-                span_name = f"tools/{req.params.name}"
-            else:
-                span_name = "unknown"
+            span_name = f"tools/{req.params.name}"
         return span_name
 
-    def handle_attributes(self, span, request, is_client=True):  # pylint: disable=no-self-use
-        import mcp.types as types  # pylint: disable=import-outside-toplevel,consider-using-from-import
-
-        operation = self._get_span_name(request)
-        if isinstance(request, types.ListToolsRequest):
-            operation = "ListTool"
-            span.set_attribute("mcp.list_tools", True)
-        elif isinstance(request, types.CallToolRequest):
-            if hasattr(request, "params") and hasattr(request.params, "name"):
-                operation = request.params.name
-            span.set_attribute("mcp.call_tool", True)
-        if is_client:
-            self._add_client_attributes(span, operation, request)
-        else:
-            self._add_server_attributes(span, operation, request)
+    @staticmethod
+    def _add_client_attributes(span: trace.Span, operation: str, request: ClientRequest) -> None:
+        import os  # pylint: disable=import-outside-toplevel
 
-    def _add_client_attributes(self, span, operation, request):  # pylint: disable=no-self-use
-        span.set_attribute("aws.remote.service", "Appsignals MCP Server")
+        service_name = os.environ.get("MCP_SERVICE_NAME", "Generic MCP Server")
+        span.set_attribute("aws.remote.service", service_name)
         span.set_attribute("aws.remote.operation", operation)
-        if hasattr(request, "params") and hasattr(request.params, "name"):
+        if hasattr(request, "params") and request.params and hasattr(request.params, "name"):
             span.set_attribute("tool.name", request.params.name)
 
-    def _add_server_attributes(self, span, operation, request):  # pylint: disable=no-self-use
-        span.set_attribute("server_side", True)
-        if hasattr(request, "params") and hasattr(request.params, "name"):
+    @staticmethod
+    def _add_server_attributes(span: trace.Span, operation: str, request: ClientRequest) -> None:
+        if hasattr(request, "params") and request.params and hasattr(request.params, "name"):
             span.set_attribute("tool.name", request.params.name)
-
-    def instrumentation_dependencies(self) -> Collection[str]:  # pylint: disable=no-self-use
-        return _instruments
-
-    def _instrument(self, **kwargs: Any) -> None:  # pylint: disable=no-self-use
-        tracer_provider = kwargs.get("tracer_provider")
-        if tracer_provider:
-            self.tracer = tracer_provider.get_tracer("mcp")
-        else:
-            self.tracer = trace.get_tracer("mcp")
-        register_post_import_hook(
-            lambda _: wrap_function_wrapper(
-                "mcp.shared.session",
-                "BaseSession.send_request",
-                self._wrap_send_request,
-            ),
-            "mcp.shared.session",
-        )
-        register_post_import_hook(
-            lambda _: wrap_function_wrapper(
-                "mcp.server.lowlevel.server",
-                "Server._handle_request",
-                self._wrap_handle_request,
-            ),
-            "mcp.server.lowlevel.server",
-        )
-
-    def _uninstrument(self, **kwargs: Any) -> None:  # pylint: disable=no-self-use
-        unwrap("mcp.shared.session", "BaseSession.send_request")
-        unwrap("mcp.server.lowlevel.server", "Server._handle_request")
diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_mcpinstrumentor.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_mcpinstrumentor.py
