link parent span

Author: liukatkat

src/semgrep_mcp/semgrep.py

@@ -6,6 +6,7 @@
 
 from mcp.shared.exceptions import McpError
 from mcp.types import INTERNAL_ERROR, ErrorData
+from opentelemetry import trace
 
 from semgrep_mcp.models import CodeFile
 from semgrep_mcp.semgrep_interfaces.semgrep_output_v1 import CliOutput
@@ -139,9 +140,11 @@ class SemgrepContext:
     process: asyncio.subprocess.Process
     stdin: asyncio.StreamWriter
     stdout: asyncio.StreamReader
+    top_level_span: trace.Span
 
-    def __init__(self, process: asyncio.subprocess.Process) -> None:
+    def __init__(self, process: asyncio.subprocess.Process, top_level_span: trace.Span) -> None:
         self.process = process
+        self.top_level_span = top_level_span
 
         if process.stdin is not None and process.stdout is not None:
             self.stdin = process.stdin

src/semgrep_mcp/server.py

@@ -30,7 +30,7 @@
     set_semgrep_executable,
 )
 from semgrep_mcp.semgrep_interfaces.semgrep_output_v1 import CliOutput
-from utilities.tracing import initialize_tracing
+from utilities.tracing import start_tracing, with_span
 
 # ---------------------------------------------------------------------------------
 # Constants
@@ -282,18 +282,18 @@ def remove_temp_dir_from_results(results: SemgrepScanResult, temp_dir: str) -> N
 async def server_lifespan(_server: Server) -> AsyncIterator[SemgrepContext]:
     """Manage server startup and shutdown lifecycle."""
     # Initialize resources on startup with tracing
-    # MCP requires Pro Engine  
-    with initialize_tracing("mcp-python-server") as span:
+    # MCP requires Pro Engine
+    with start_tracing("mcp-python-server") as span:
         process = await run_semgrep_process(["mcp", "--pro", "--debug"])
 
-    try:
-        yield SemgrepContext(process=process)
-    finally:
-        if process.returncode is None:
-            # Clean up on shutdown
-            process.terminate()
-        else:
-            print(f"`semgrep mcp` process exited with code {process.returncode} already")
+        try:
+            yield SemgrepContext(process=process, top_level_span=span)
+        finally:
+            if process.returncode is None:
+                # Clean up on shutdown
+                process.terminate()
+            else:
+                print(f"`semgrep mcp` process exited with code {process.returncode} already")
 
 
 # Create a fast MCP server
@@ -678,7 +678,8 @@ async def semgrep_scan_rpc(
     temp_dir = None
     try:
         # TODO: perhaps should return more interpretable results?
-        cli_output = await run_semgrep_via_rpc(context, code_files)
+        with with_span(context.top_level_span, "semgrep_scan_rpc"):
+            cli_output = await run_semgrep_via_rpc(context, code_files)
         return cli_output
     except McpError as e:
         raise e

src/utilities/tracing.py

@@ -1,12 +1,12 @@
 #!/usr/bin/env python3
 
 import os
+from collections.abc import Generator
 from contextlib import contextmanager
-from typing import Generator
 
 from opentelemetry import trace
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
-from opentelemetry.sdk.resources import SERVICE_NAME, DEPLOYMENT_ENVIRONMENT, Resource
+from opentelemetry.sdk.resources import DEPLOYMENT_ENVIRONMENT, SERVICE_NAME, Resource
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 
@@ -17,62 +17,63 @@
 
 MCP_SERVICE_NAME = "mcp"
 
-top_level_span : trace.Span | None = None
 
-
-def get_trace_endpoint() -> (str, str):
+def get_trace_endpoint() -> tuple[str, str]:
     """Get the appropriate trace endpoint based on environment."""
     env = os.environ.get("ENVIRONMENT", "dev").lower()
-    
+
     if env == "prod":
         return (DEFAULT_TRACE_ENDPOINT, "prod")
     elif env == "local":
         return (DEFAULT_LOCAL_ENDPOINT, "local")
     else:
         return (DEFAULT_DEV_ENDPOINT, "dev")
 
+
 @contextmanager
-def initialize_tracing(name: str) -> Generator[trace.Span, None, None]:
+def start_tracing(name: str) -> Generator[trace.Span, None, None]:
     """Initialize OpenTelemetry tracing."""
-
     (endpoint, env) = get_trace_endpoint()
-    
+
     # Create resource with basic attributes
-    resource = Resource.create({
-        SERVICE_NAME: MCP_SERVICE_NAME,
-        DEPLOYMENT_ENVIRONMENT: env,
-    })
+    resource = Resource.create(
+        {
+            SERVICE_NAME: MCP_SERVICE_NAME,
+            DEPLOYMENT_ENVIRONMENT: env,
+        }
+    )
     # Create tracer provider
     provider = TracerProvider(resource=resource)
-    
+
     # Create OTLP exporter
     exporter = OTLPSpanExporter(endpoint=endpoint)
-    
+
     # Create span processor
     processor = BatchSpanProcessor(exporter)
     provider.add_span_processor(processor)
-    
+
     # Set the global tracer provider
     trace.set_tracer_provider(provider)
-    
+
     # Get tracer instance
     tracer = trace.get_tracer(MCP_SERVICE_NAME)
 
     with tracer.start_as_current_span(name) as span:
-        top_level_span = span
-        trace_id = trace.format_trace_id(top_level_span.get_span_context().trace_id)
+        trace_id = trace.format_trace_id(span.get_span_context().trace_id)
         # TODO: use logging
         print("Tracing initialized")
         print(f"Tracing initialized with trace ID: {trace_id}")
-        
+
         yield span
 
 
 @contextmanager
-def trace_span(
-    name: str, 
+def with_span(
+    parent_span: trace.Span,
+    name: str,
 ) -> Generator[trace.Span, None, None]:
     tracer = trace.get_tracer(MCP_SERVICE_NAME)
-    
-    with tracer.start_as_current_span(name) as span:
-        yield span
+
+    context = trace.set_span_in_context(parent_span)
+    with tracer.start_span(name, context=context) as span:
+        yield span