feat(contrib/mcp-go): Trace MCP session initializations with MLObs spans (#4101)

Co-authored-by: julian.boilen <[email protected]>

Author: jboolean

contrib/mark3labs/mcp-go/README.md

@@ -2,6 +2,8 @@
 
 This integration provides Datadog tracing for the [mark3labs/mcp-go](https://github.com/mark3labs/mcp-go) library.
 
+Both hooks and middleware are used.
+
 ## Usage
 
 ```go
@@ -15,13 +17,18 @@ func main() {
 	tracer.Start()
 	defer tracer.Stop()
 
-	srv := server.NewMCPServer("my-server", "1.0.0",
-		server.WithToolHandlerMiddleware(mcpgotrace.NewToolHandlerMiddleware()))
-	_ = srv
+    // Add tracing to your server hooks
+    hooks := &server.Hooks{}
+    mcpgotrace.AddServerHooks(hooks)
+
+    srv := server.NewMCPServer("my-server", "1.0.0",
+        server.WithHooks(hooks),
+        server.WithToolHandlerMiddleware(mcpgotrace.NewToolHandlerMiddleware()))
 }
 ```
 
 ## Features
 
 The integration automatically traces:
 - **Tool calls**: Creates LLMObs tool spans with input/output annotation for all tool invocations
+- **Session initialization**: Create LLMObs task spans for session initialization, including client information.

contrib/mark3labs/mcp-go/example_test.go

@@ -15,7 +15,12 @@ func Example() {
 	tracer.Start()
 	defer tracer.Stop()
 
+	// Create server hooks and add Datadog tracing
+	hooks := &server.Hooks{}
+	mcpgotrace.AddServerHooks(hooks)
+
 	srv := server.NewMCPServer("my-server", "1.0.0",
+		server.WithHooks(hooks),
 		server.WithToolHandlerMiddleware(mcpgotrace.NewToolHandlerMiddleware()))
 	_ = srv
 }

contrib/mark3labs/mcp-go/mcpgo.go

@@ -8,6 +8,7 @@ package mcpgo // import "github.com/DataDog/dd-trace-go/contrib/mark3labs/mcp-go
 import (
 	"context"
 	"encoding/json"
+	"sync"
 
 	"github.com/DataDog/dd-trace-go/v2/instrumentation"
 	"github.com/DataDog/dd-trace-go/v2/llmobs"
@@ -22,17 +23,35 @@ func init() {
 	instr = instrumentation.Load(instrumentation.PackageMark3LabsMCPGo)
 }
 
+type hooks struct {
+	spanCache *sync.Map
+}
+
+// AddServerHooks appends Datadog tracing hooks to an existing server.Hooks object.
+func AddServerHooks(hooks *server.Hooks) {
+	ddHooks := newHooks()
+	hooks.AddBeforeInitialize(ddHooks.onBeforeInitialize)
+	hooks.AddAfterInitialize(ddHooks.onAfterInitialize)
+	hooks.AddOnError(ddHooks.onError)
+}
+
 func NewToolHandlerMiddleware() server.ToolHandlerMiddleware {
 	return func(next server.ToolHandlerFunc) server.ToolHandlerFunc {
 		return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
 			toolSpan, ctx := llmobs.StartToolSpan(ctx, request.Params.Name, llmobs.WithIntegration(string(instrumentation.PackageMark3LabsMCPGo)))
 
 			result, err := next(ctx, request)
 
-			inputJSON, _ := json.Marshal(request)
+			inputJSON, marshalErr := json.Marshal(request)
+			if marshalErr != nil {
+				instr.Logger().Warn("mcp-go: failed to marshal tool request: %v", marshalErr)
+			}
 			var outputText string
 			if result != nil {
-				resultJSON, _ := json.Marshal(result)
+				resultJSON, marshalErr := json.Marshal(result)
+				if marshalErr != nil {
+					instr.Logger().Warn("mcp-go: failed to marshal tool result: %v", marshalErr)
+				}
 				outputText = string(resultJSON)
 			}
 
@@ -48,3 +67,71 @@ func NewToolHandlerMiddleware() server.ToolHandlerMiddleware {
 		}
 	}
 }
+
+func newHooks() *hooks {
+	return &hooks{
+		spanCache: &sync.Map{},
+	}
+}
+
+func (h *hooks) onBeforeInitialize(ctx context.Context, id any, request *mcp.InitializeRequest) {
+	taskSpan, _ := llmobs.StartTaskSpan(ctx, "mcp.initialize", llmobs.WithIntegration("mark3labs/mcp-go"))
+
+	clientName := request.Params.ClientInfo.Name
+	clientVersion := request.Params.ClientInfo.Version
+	taskSpan.Annotate(llmobs.WithAnnotatedTags(map[string]string{"client_name": clientName, "client_version": clientName + "_" + clientVersion}))
+
+	h.spanCache.Store(id, taskSpan)
+}
+
+func (h *hooks) onAfterInitialize(ctx context.Context, id any, request *mcp.InitializeRequest, result *mcp.InitializeResult) {
+	finishSpanWithIO(h, id, request, result)
+}
+
+func (h *hooks) onError(ctx context.Context, id any, method mcp.MCPMethod, message any, err error) {
+	if method != mcp.MethodInitialize {
+		return
+	}
+	value, ok := h.spanCache.LoadAndDelete(id)
+	if !ok {
+		return
+	}
+
+	span, ok := value.(*llmobs.TaskSpan)
+	if !ok {
+		return
+	}
+
+	defer span.Finish(llmobs.WithError(err))
+
+	inputJSON, marshalErr := json.Marshal(message)
+	if marshalErr != nil {
+		instr.Logger().Warn("mcp-go: failed to marshal error message: %v", marshalErr)
+	}
+	span.AnnotateTextIO(string(inputJSON), err.Error())
+
+}
+
+func finishSpanWithIO[Req any, Res any](h *hooks, id any, request Req, result Res) {
+	value, ok := h.spanCache.LoadAndDelete(id)
+	if !ok {
+		return
+	}
+	span, ok := value.(*llmobs.TaskSpan)
+	if !ok {
+		return
+	}
+
+	defer span.Finish()
+
+	inputJSON, marshalErr := json.Marshal(request)
+	if marshalErr != nil {
+		instr.Logger().Warn("mcp-go: failed to marshal request: %v", marshalErr)
+	}
+	resultJSON, marshalErr := json.Marshal(result)
+	if marshalErr != nil {
+		instr.Logger().Warn("mcp-go: failed to marshal result: %v", marshalErr)
+	}
+
+	span.AnnotateTextIO(string(inputJSON), string(resultJSON))
+}

contrib/mark3labs/mcp-go/mcpgo_test.go

@@ -12,6 +12,7 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/DataDog/dd-trace-go/v2/ddtrace/mocktracer"
 	"github.com/DataDog/dd-trace-go/v2/ddtrace/tracer"
 	"github.com/DataDog/dd-trace-go/v2/instrumentation/testutils/testtracer"
 	"github.com/mark3labs/mcp-go/mcp"
@@ -20,6 +21,81 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestNewToolHandlerMiddleware(t *testing.T) {
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	middleware := NewToolHandlerMiddleware()
+	assert.NotNil(t, middleware)
+}
+
+func TestAddServerHooks(t *testing.T) {
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	serverHooks := &server.Hooks{}
+	AddServerHooks(serverHooks)
+
+	assert.Len(t, serverHooks.OnBeforeInitialize, 1)
+	assert.Len(t, serverHooks.OnAfterInitialize, 1)
+	assert.Len(t, serverHooks.OnError, 1)
+}
+
+func TestIntegrationSessionInitialize(t *testing.T) {
+	tt := testTracer(t)
+	defer tt.Stop()
+
+	hooks := &server.Hooks{}
+	AddServerHooks(hooks)
+
+	srv := server.NewMCPServer("test-server", "1.0.0",
+		server.WithHooks(hooks))
+
+	ctx := context.Background()
+	initRequest := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"test-client","version":"1.0.0"}}}`
+
+	response := srv.HandleMessage(ctx, []byte(initRequest))
+	assert.NotNil(t, response)
+
+	responseBytes, err := json.Marshal(response)
+	require.NoError(t, err)
+
+	var resp map[string]interface{}
+	err = json.Unmarshal(responseBytes, &resp)
+	require.NoError(t, err)
+	assert.Equal(t, "2.0", resp["jsonrpc"])
+	assert.Equal(t, float64(1), resp["id"])
+	assert.NotNil(t, resp["result"])
+
+	spans := tt.WaitForLLMObsSpans(t, 1)
+	require.Len(t, spans, 1)
+
+	taskSpan := spans[0]
+	assert.Equal(t, "mcp.initialize", taskSpan.Name)
+	assert.Equal(t, "task", taskSpan.Meta["span.kind"])
+
+	assert.Contains(t, taskSpan.Tags, "client_name:test-client")
+	assert.Contains(t, taskSpan.Tags, "client_version:test-client_1.0.0")
+
+	assert.Contains(t, taskSpan.Meta, "input")
+	assert.Contains(t, taskSpan.Meta, "output")
+
+	inputMeta := taskSpan.Meta["input"]
+	assert.NotNil(t, inputMeta)
+	inputJSON, err := json.Marshal(inputMeta)
+	require.NoError(t, err)
+	inputStr := string(inputJSON)
+	assert.Contains(t, inputStr, "2024-11-05")
+	assert.Contains(t, inputStr, "test-client")
+
+	outputMeta := taskSpan.Meta["output"]
+	assert.NotNil(t, outputMeta)
+	outputJSON, err := json.Marshal(outputMeta)
+	require.NoError(t, err)
+	outputStr := string(outputJSON)
+	assert.Contains(t, outputStr, "serverInfo")
+}
+
 // Test tool spans are recorded on a successful tool call
 func TestIntegrationToolCallSuccess(t *testing.T) {
 	tt := testTracer(t)