Merge pull request #126 from golf-mcp/aschlean/fixing-otel-instum

Aschlean/fixing otel instumentation

Author: aschlean

src/golf/auth/factory.py

@@ -264,6 +264,7 @@ def _create_oauth_proxy_provider(config: OAuthProxyConfig) -> "AuthProvider":
             env_value = env_value.strip()
             try:
                 from urllib.parse import urlparse
+
                 parsed = urlparse(env_value)
                 if not parsed.scheme or not parsed.netloc:
                     raise ValueError(
@@ -292,6 +293,7 @@ def _create_oauth_proxy_provider(config: OAuthProxyConfig) -> "AuthProvider":
             env_value = env_value.strip()
             try:
                 from urllib.parse import urlparse
+
                 parsed = urlparse(env_value)
                 if not parsed.scheme or not parsed.netloc:
                     raise ValueError(
@@ -300,26 +302,21 @@ def _create_oauth_proxy_provider(config: OAuthProxyConfig) -> "AuthProvider":
                     )
                 if parsed.scheme not in ("http", "https"):
                     raise ValueError(
-                        f"Token endpoint from {config.token_endpoint_env_var} "
-                        f"must use http or https: '{env_value}'"
+                        f"Token endpoint from {config.token_endpoint_env_var} must use http or https: '{env_value}'"
                     )
                 token_endpoint = env_value
             except Exception as e:
                 if isinstance(e, ValueError):
                     raise
-                raise ValueError(
-                    f"Invalid token_endpoint from {config.token_endpoint_env_var}: {e}"
-                ) from e
+                raise ValueError(f"Invalid token_endpoint from {config.token_endpoint_env_var}: {e}") from e
 
     client_id = config.client_id
     if config.client_id_env_var:
         env_value = os.environ.get(config.client_id_env_var)
         if env_value:
             client_id = env_value.strip()
             if not client_id:
-                raise ValueError(
-                    f"Client ID from environment variable {config.client_id_env_var} cannot be empty"
-                )
+                raise ValueError(f"Client ID from environment variable {config.client_id_env_var} cannot be empty")
 
     client_secret = config.client_secret
     if config.client_secret_env_var:
@@ -339,24 +336,19 @@ def _create_oauth_proxy_provider(config: OAuthProxyConfig) -> "AuthProvider":
             env_value = env_value.strip()
             try:
                 from urllib.parse import urlparse
+
                 parsed = urlparse(env_value)
                 if not parsed.scheme or not parsed.netloc:
                     raise ValueError(
-                        f"Invalid base_url from environment variable "
-                        f"{config.base_url_env_var}: '{env_value}'"
+                        f"Invalid base_url from environment variable {config.base_url_env_var}: '{env_value}'"
                     )
                 if parsed.scheme not in ("http", "https"):
-                    raise ValueError(
-                        f"Base URL from {config.base_url_env_var} "
-                        f"must use http or https: '{env_value}'"
-                    )
+                    raise ValueError(f"Base URL from {config.base_url_env_var} must use http or https: '{env_value}'")
                 base_url = env_value
             except Exception as e:
                 if isinstance(e, ValueError):
                     raise
-                raise ValueError(
-                    f"Invalid base_url from {config.base_url_env_var}: {e}"
-                ) from e
+                raise ValueError(f"Invalid base_url from {config.base_url_env_var}: {e}") from e
 
     revocation_endpoint = config.revocation_endpoint
     if config.revocation_endpoint_env_var:
@@ -367,6 +359,7 @@ def _create_oauth_proxy_provider(config: OAuthProxyConfig) -> "AuthProvider":
             if env_value:  # Only validate if not empty
                 try:
                     from urllib.parse import urlparse
+
                     parsed = urlparse(env_value)
                     if not parsed.scheme or not parsed.netloc:
                         raise ValueError(
@@ -388,28 +381,37 @@ def _create_oauth_proxy_provider(config: OAuthProxyConfig) -> "AuthProvider":
 
     # Final validation: ensure all required fields have values after env resolution
     if not authorization_endpoint:
-        env_var_hint = f" (environment variable {config.authorization_endpoint_env_var} is not set)" \
-                      if config.authorization_endpoint_env_var else ""
+        env_var_hint = (
+            f" (environment variable {config.authorization_endpoint_env_var} is not set)"
+            if config.authorization_endpoint_env_var
+            else ""
+        )
         raise ValueError(f"Authorization endpoint is required but not provided{env_var_hint}")
 
     if not token_endpoint:
-        env_var_hint = f" (environment variable {config.token_endpoint_env_var} is not set)" \
-                      if config.token_endpoint_env_var else ""
+        env_var_hint = (
+            f" (environment variable {config.token_endpoint_env_var} is not set)"
+            if config.token_endpoint_env_var
+            else ""
+        )
         raise ValueError(f"Token endpoint is required but not provided{env_var_hint}")
 
     if not client_id:
-        env_var_hint = f" (environment variable {config.client_id_env_var} is not set)" \
-                      if config.client_id_env_var else ""
+        env_var_hint = (
+            f" (environment variable {config.client_id_env_var} is not set)" if config.client_id_env_var else ""
+        )
         raise ValueError(f"Client ID is required but not provided{env_var_hint}")
 
     if not client_secret:
-        env_var_hint = f" (environment variable {config.client_secret_env_var} is not set)" \
-                      if config.client_secret_env_var else ""
+        env_var_hint = (
+            f" (environment variable {config.client_secret_env_var} is not set)" if config.client_secret_env_var else ""
+        )
         raise ValueError(f"Client secret is required but not provided{env_var_hint}")
 
     if not base_url:
-        env_var_hint = f" (environment variable {config.base_url_env_var} is not set)" \
-                      if config.base_url_env_var else ""
+        env_var_hint = (
+            f" (environment variable {config.base_url_env_var} is not set)" if config.base_url_env_var else ""
+        )
         raise ValueError(f"Base URL is required but not provided{env_var_hint}")
 
     # Production security checks
@@ -430,9 +432,7 @@ def _create_oauth_proxy_provider(config: OAuthProxyConfig) -> "AuthProvider":
         ]:
             parsed = urlparse(url_value)
             if parsed.scheme == "http":
-                raise ValueError(
-                    f"OAuth proxy {url_name} must use HTTPS in production environment: '{url_value}'"
-                )
+                raise ValueError(f"OAuth proxy {url_name} must use HTTPS in production environment: '{url_value}'")
 
         # Check for localhost in production
         parsed_base = urlparse(base_url)

src/golf/core/builder.py

@@ -990,23 +990,20 @@ def _generate_server(self) -> None:
                 imports.append(f"import {full_module_path}")
 
                 # Add code to register this component
+                # Note: When opentelemetry_enabled, we use OpenTelemetryMiddleware for span creation
+                # instead of wrapping individual functions, to ensure proper context propagation
                 if self.settings.opentelemetry_enabled:
-                    # Use telemetry instrumentation
-                    registration = f"# Register the {component_type.value} '{component.name}' with telemetry"
+                    # Register components without function wrapping - middleware handles tracing
                     entry_func = (
                         component.entry_function
                         if hasattr(component, "entry_function") and component.entry_function
                         else "export"
                     )
 
-                    registration += (
-                        f"\n_wrapped_func = instrument_{component_type.value}("
-                        f"{full_module_path}.{entry_func}, '{component.name}')"
-                    )
-
                     if component_type == ComponentType.TOOL:
+                        registration = f"# Register the tool '{component.name}'"
                         registration += (
-                            f"\n_tool = Tool.from_function(_wrapped_func, "
+                            f"\n_tool = Tool.from_function({full_module_path}.{entry_func}, "
                             f'name="{component.name}", '
                             f"description={repr(component.docstring or '')})"
                         )
@@ -1015,23 +1012,25 @@ def _generate_server(self) -> None:
                             registration += f".with_annotations({component.annotations})"
                         registration += "\nmcp.add_tool(_tool)"
                     elif component_type == ComponentType.RESOURCE:
+                        registration = f"# Register the resource '{component.name}'"
                         if self._is_resource_template(component):
                             registration += (
-                                f"\n_template = ResourceTemplate.from_function(_wrapped_func, "
+                                f"\n_template = ResourceTemplate.from_function({full_module_path}.{entry_func}, "
                                 f'uri_template="{component.uri_template}", name="{component.name}", '
                                 f"description={repr(component.docstring or '')})\n"
                                 f"mcp.add_template(_template)"
                             )
                         else:
                             registration += (
-                                f"\n_resource = Resource.from_function(_wrapped_func, "
+                                f"\n_resource = Resource.from_function({full_module_path}.{entry_func}, "
                                 f'uri="{component.uri_template}", name="{component.name}", '
                                 f"description={repr(component.docstring or '')})\n"
                                 f"mcp.add_resource(_resource)"
                             )
                     else:  # PROMPT
+                        registration = f"# Register the prompt '{component.name}'"
                         registration += (
-                            f"\n_prompt = Prompt.from_function(_wrapped_func, "
+                            f"\n_prompt = Prompt.from_function({full_module_path}.{entry_func}, "
                             f'name="{component.name}", '
                             f"description={repr(component.docstring or '')})\n"
                             f"mcp.add_prompt(_prompt)"
@@ -1198,6 +1197,13 @@ def _generate_server(self) -> None:
             imports.append("")
             component_registrations.append("")
 
+        # Add OpenTelemetry FastMCP middleware if enabled (must be added before other middleware)
+        if self.settings.opentelemetry_enabled:
+            component_registrations.append("# Register OpenTelemetry middleware for proper span context propagation")
+            component_registrations.append("from golf.telemetry import OpenTelemetryMiddleware")
+            component_registrations.append("mcp.add_middleware(OpenTelemetryMiddleware())")
+            component_registrations.append("")
+
         # Check for custom middleware.py file and register middleware classes
         middleware_classes = self._discover_middleware_classes(self.project_path)
         if middleware_classes:
@@ -1308,10 +1314,11 @@ def _generate_server(self) -> None:
                 middleware_list.append("Middleware(MetricsMiddleware)")
 
             # Add OpenTelemetry middleware if enabled
+            # Note: We intentionally do NOT use opentelemetry.instrumentation.asgi.OpenTelemetryMiddleware
+            # because it creates noisy low-level ASGI spans (http receive/send). Golf's FastMCP middleware
+            # (OpenTelemetryMiddleware) creates cleaner, more meaningful MCP-level spans.
             if self.settings.opentelemetry_enabled:
-                middleware_setup.append("    from opentelemetry.instrumentation.asgi import OpenTelemetryMiddleware")
-                middleware_setup.append("    from starlette.middleware import Middleware")
-                middleware_list.append("Middleware(OpenTelemetryMiddleware)")
+                pass  # OpenTelemetry middleware is added via mcp.add_middleware() earlier in the code
 
             if middleware_setup:
                 main_code.extend(middleware_setup)
@@ -1368,10 +1375,11 @@ def _generate_server(self) -> None:
                 middleware_list.append("Middleware(MetricsMiddleware)")
 
             # Add OpenTelemetry middleware if enabled
+            # Note: We intentionally do NOT use opentelemetry.instrumentation.asgi.OpenTelemetryMiddleware
+            # because it creates noisy low-level ASGI spans (http receive/send). Golf's FastMCP middleware
+            # (OpenTelemetryMiddleware) creates cleaner, more meaningful MCP-level spans.
             if self.settings.opentelemetry_enabled:
-                middleware_setup.append("    from opentelemetry.instrumentation.asgi import OpenTelemetryMiddleware")
-                middleware_setup.append("    from starlette.middleware import Middleware")
-                middleware_list.append("Middleware(OpenTelemetryMiddleware)")
+                pass  # OpenTelemetry middleware is added via mcp.add_middleware() earlier in the code
 
             if middleware_setup:
                 main_code.extend(middleware_setup)

src/golf/telemetry/__init__.py

@@ -1,6 +1,7 @@
 """Golf telemetry module for OpenTelemetry instrumentation."""
 
 from golf.telemetry.instrumentation import (
+    get_provider,
     get_tracer,
     init_telemetry,
     instrument_elicitation,
@@ -9,6 +10,8 @@
     instrument_sampling,
     instrument_tool,
     telemetry_lifespan,
+    OpenTelemetryMiddleware,
+    OTelContextCapturingMiddleware,
 )
 
 __all__ = [
@@ -19,5 +22,8 @@
     "instrument_sampling",
     "telemetry_lifespan",
     "init_telemetry",
+    "get_provider",
     "get_tracer",
+    "OpenTelemetryMiddleware",
+    "OTelContextCapturingMiddleware",
 ]