link telemetry to ocaml server

Author: liukatkat

src/semgrep_mcp/semgrep.py

@@ -175,7 +175,9 @@ async def send_request(self, request: str, **kwargs: Any) -> str:
 ################################################################################
 
 
-async def run_semgrep_process(args: list[str]) -> asyncio.subprocess.Process:
+async def run_semgrep_process(
+    top_level_span: trace.Span | None, args: list[str]
+) -> asyncio.subprocess.Process:
     """
     Runs semgrep with the given arguments as a subprocess, without waiting for it to finish.
     """
@@ -185,7 +187,14 @@ async def run_semgrep_process(args: list[str]) -> asyncio.subprocess.Process:
 
     # Just so we get the debug logs for the MCP server
     env = os.environ.copy()
-    env["SEMGREP_LOG_SRCS"] = "mcp"
+    env["SEMGREP_LOG_SRCS"] = "mcp,commons"
+    if top_level_span:
+        env["SEMGREP_TRACE_PARENT_SPAN_ID"] = trace.format_span_id(
+            top_level_span.get_span_context().span_id
+        )
+        env["SEMGREP_TRACE_PARENT_TRACE_ID"] = trace.format_trace_id(
+            top_level_span.get_span_context().trace_id
+        )
 
     # Execute semgrep command
     process = await asyncio.create_subprocess_exec(
@@ -202,7 +211,7 @@ async def run_semgrep_process(args: list[str]) -> asyncio.subprocess.Process:
     return process
 
 
-async def run_semgrep(args: list[str]) -> str:
+async def run_semgrep(top_level_span: trace.Span | None, args: list[str]) -> str:
     """
     Runs semgrep with the given arguments
 
@@ -213,7 +222,7 @@ async def run_semgrep(args: list[str]) -> str:
         Output of semgrep command
     """
 
-    process = await run_semgrep_process(args)
+    process = await run_semgrep_process(top_level_span, args)
 
     stdout, stderr = await process.communicate()
 

src/semgrep_mcp/server.py

@@ -284,7 +284,7 @@ async def server_lifespan(_server: Server) -> AsyncIterator[SemgrepContext]:
     # Initialize resources on startup with tracing
     # MCP requires Pro Engine
     with start_tracing("mcp-python-server") as span:
-        process = await run_semgrep_process(["mcp", "--pro", "--debug"])
+        process = await run_semgrep_process(top_level_span=span, args=["mcp", "--pro", "--debug"])
 
         try:
             yield SemgrepContext(process=process, top_level_span=span)
@@ -350,7 +350,7 @@ async def get_supported_languages() -> list[str]:
     args = ["show", "supported-languages", "--experimental"]
 
     # Parse output and return list of languages
-    languages = await run_semgrep(args)
+    languages = await run_semgrep(top_level_span=None, args=args)
     return [lang.strip() for lang in languages.strip().split("\n") if lang.strip()]
 
 
@@ -586,7 +586,7 @@ async def semgrep_scan_with_custom_rule(
 
         # Run semgrep scan with custom rule
         args = get_semgrep_scan_args(temp_dir, rule_file_path)
-        output = await run_semgrep(args)
+        output = await run_semgrep(top_level_span=None, args=args)
         results: SemgrepScanResult = SemgrepScanResult.model_validate_json(output)
         remove_temp_dir_from_results(results, temp_dir)
         return results
@@ -631,7 +631,7 @@ async def semgrep_scan(
         # Create temporary files from code content
         temp_dir = create_temp_files_from_code_content(code_files)
         args = get_semgrep_scan_args(temp_dir, config)
-        output = await run_semgrep(args)
+        output = await run_semgrep(top_level_span=None, args=args)
         results: SemgrepScanResult = SemgrepScanResult.model_validate_json(output)
         remove_temp_dir_from_results(results, temp_dir)
         return results
@@ -731,7 +731,7 @@ async def semgrep_scan_local(
         results = []
         for cf in code_files:
             args = get_semgrep_scan_args(cf.path, config)
-            output = await run_semgrep(args)
+            output = await run_semgrep(top_level_span=None, args=args)
             result: SemgrepScanResult = SemgrepScanResult.model_validate_json(output)
             results.append(result)
         return results
@@ -787,7 +787,7 @@ async def security_check(
         # Create temporary files from code content
         temp_dir = create_temp_files_from_code_content(code_files)
         args = get_semgrep_scan_args(temp_dir)
-        output = await run_semgrep(args)
+        output = await run_semgrep(top_level_span=None, args=args)
         results: SemgrepScanResult = SemgrepScanResult.model_validate_json(output)
         remove_temp_dir_from_results(results, temp_dir)
 
@@ -850,7 +850,7 @@ async def get_abstract_syntax_tree(
             "--json",
             temp_file_path,
         ]
-        return await run_semgrep(args)
+        return await run_semgrep(top_level_span=None, args=args)
 
     except McpError as e:
         raise e