From ac2d09737fb00b137214a3c85041a1769c9812c9 Mon Sep 17 00:00:00 2001
From: Jimmy Byrd <jimmybyrd@gmail.com>
Date: Mon, 20 Feb 2023 22:51:56 -0500
Subject: [PATCH] adds JsonRpc tracing

---
 src/FsAutoComplete.Core/Utils.fs              | 27 +++++++++++++++++++
 .../LspServers/AdaptiveFSharpLspServer.fs     |  5 ++--
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/FsAutoComplete.Core/Utils.fs b/src/FsAutoComplete.Core/Utils.fs
index 750f6de69..184bbd4c0 100644
--- a/src/FsAutoComplete.Core/Utils.fs
+++ b/src/FsAutoComplete.Core/Utils.fs
@@ -827,6 +827,7 @@ module Tracing =
 
   open System.Diagnostics
   open FsOpenTelemetry
+  open StreamJsonRpc
 
   module SemanticConventions =
     [<Literal>]
@@ -845,3 +846,29 @@ module Tracing =
   let serviceName = "FsAutoComplete"
 
   let fsacActivitySource = new ActivitySource(serviceName, Version.info().Version)
+
+  type StreamJsonRpcTracingStrategy (activitySource : ActivitySource) =
+    interface IActivityTracingStrategy with
+        member this.ApplyInboundActivity(request: Protocol.JsonRpcRequest): IDisposable =
+
+          let a = activitySource.StartActivity(request.Method)
+          a
+            .SetTagSafe("rpc.system", "jsonrpc")
+            .SetTagSafe("rpc.jsonrpc.argumentNames",  String.Join(',', request.ArgumentNames))
+            .SetTagSafe("rpc.jsonrpc.isNotification", request.IsNotification)
+            .SetTagSafe("rpc.jsonrpc.IsResponseExpected", request.IsResponseExpected)
+            .SetTagSafe("rpc.jsonrpc.version", request.Version)
+            .SetTagSafe("rpc.jsonrpc.request_id", request.RequestId)
+            .SetTagSafe("rpc.method", request.Method)
+            |> ignore
+
+          a.TraceStateString <- request.TraceState
+          if request.TraceParent <> null then
+            a.SetParentId(request.TraceParent) |> ignore
+          a
+
+        member this.ApplyOutboundActivity(request: Protocol.JsonRpcRequest): unit =
+            if Activity.Current <> null && Activity.Current.IdFormat = ActivityIdFormat.W3C then
+              request.TraceParent <- Activity.Current.Id
+              request.TraceState <- Activity.Current.TraceStateString
+            ()
diff --git a/src/FsAutoComplete/LspServers/AdaptiveFSharpLspServer.fs b/src/FsAutoComplete/LspServers/AdaptiveFSharpLspServer.fs
index ec09dce3a..7d48faf93 100644
--- a/src/FsAutoComplete/LspServers/AdaptiveFSharpLspServer.fs
+++ b/src/FsAutoComplete/LspServers/AdaptiveFSharpLspServer.fs
@@ -4426,12 +4426,13 @@ module AdaptiveFSharpLspServer =
           None
       | _ -> None
 
-    { new JsonRpc(handler) with
+    let strategy = StreamJsonRpcTracingStrategy(Tracing.fsacActivitySource)
+
+    { new JsonRpc(handler, ActivityTracingStrategy = strategy) with
         member this.IsFatalException(ex: Exception) =
           match ex with
           | HandleableException -> false
           | _ -> true }
-
   let startCore toolsPath workspaceLoaderFactory =
     use input = Console.OpenStandardInput()
     use output = Console.OpenStandardOutput()