Use middleware instead of hooks for tool calls

Author: jboolean

contrib/mark3labs/mcp-go/example_test.go

@@ -6,7 +6,7 @@
 package mcpgo_test
 
 import (
-	mcpgo "github.com/DataDog/dd-trace-go/contrib/mark3labs/mcp-go/v2"
+	mcpgotrace "github.com/DataDog/dd-trace-go/contrib/mark3labs/mcp-go/v2"
 	"github.com/DataDog/dd-trace-go/v2/ddtrace/tracer"
 	"github.com/mark3labs/mcp-go/server"
 )
@@ -15,16 +15,11 @@ func Example() {
 	tracer.Start()
 	defer tracer.Stop()
 
-	// Create Datadog hooks
-	ddHooks := mcpgo.NewHooks()
-	defer ddHooks.Stop()
-
-	// Create server hooks and add Datadog tracing
-	hooks := &server.Hooks{}
-	ddHooks.AddHooks(hooks)
+	// Create Datadog middleware
+	mcpgotrace.NewToolHandlerMiddleware()
 
 	// Use hooks with your MCP server
 	srv := server.NewMCPServer("my-server", "1.0.0",
-		server.WithHooks(hooks))
+		server.WithToolHandlerMiddleware(mcpgotrace.NewToolHandlerMiddleware()))
 	_ = srv
 }

contrib/mark3labs/mcp-go/hooks.go

@@ -0,0 +1,50 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016 Datadog, Inc.
+
+package mcpgo // import "github.com/DataDog/dd-trace-go/contrib/mark3labs/mcp-go/v2"
+
+import (
+	"context"
+	"encoding/json"
+
+	"github.com/DataDog/dd-trace-go/v2/instrumentation"
+	"github.com/DataDog/dd-trace-go/v2/llmobs"
+
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/mark3labs/mcp-go/server"
+)
+
+var instr *instrumentation.Instrumentation
+
+func init() {
+	instr = instrumentation.Load(instrumentation.PackageMark3LabsMcpGo)
+}
+
+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)
+
+			result, err := next(ctx, request)
+
+			inputJSON, _ := json.Marshal(request)
+			var outputText string
+			if result != nil {
+				resultJSON, _ := json.Marshal(result)
+				outputText = string(resultJSON)
+			}
+
+			toolSpan.AnnotateTextIO(string(inputJSON), outputText)
+
+			if err != nil {
+				toolSpan.Finish(llmobs.WithError(err))
+			} else {
+				toolSpan.Finish()
+			}
+
+			return result, err
+		}
+	}
+}

contrib/mark3labs/mcp-go/mcpgo.go

@@ -21,76 +21,20 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestNewHooks(t *testing.T) {
+func TestNewToolHandlerMiddleware(t *testing.T) {
 	mt := mocktracer.Start()
 	defer mt.Stop()
 
-	hooks := NewHooks()
-	assert.NotNil(t, hooks)
-	assert.NotNil(t, hooks.toolCache)
-
-	hooks.Stop()
-}
-
-func TestAddHooks(t *testing.T) {
-	mt := mocktracer.Start()
-	defer mt.Stop()
-
-	ddHooks := NewHooks()
-	defer ddHooks.Stop()
-
-	serverHooks := &server.Hooks{}
-	ddHooks.AddHooks(serverHooks)
-
-	assert.Len(t, serverHooks.OnBeforeCallTool, 1)
-	assert.Len(t, serverHooks.OnAfterCallTool, 1)
-	assert.Len(t, serverHooks.OnError, 1)
-}
-
-// Integration Tests
-
-func TestIntegrationSessionInitialize(t *testing.T) {
-	tt := testTracer(t)
-	defer tt.Stop()
-
-	ddHooks := NewHooks()
-	defer ddHooks.Stop()
-
-	hooks := &server.Hooks{}
-	ddHooks.AddHooks(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"])
+	middleware := NewToolHandlerMiddleware()
+	assert.NotNil(t, middleware)
 }
 
 func TestIntegrationToolCallSuccess(t *testing.T) {
 	tt := testTracer(t)
 	defer tt.Stop()
 
-	ddHooks := NewHooks()
-	defer ddHooks.Stop()
-
-	hooks := &server.Hooks{}
-	ddHooks.AddHooks(hooks)
-
 	srv := server.NewMCPServer("test-server", "1.0.0",
-		server.WithHooks(hooks))
+		server.WithToolHandlerMiddleware(NewToolHandlerMiddleware()))
 
 	calcTool := mcp.NewTool("calculator",
 		mcp.WithDescription("A simple calculator"))
@@ -167,14 +111,8 @@ func TestIntegrationToolCallError(t *testing.T) {
 	tt := testTracer(t)
 	defer tt.Stop()
 
-	ddHooks := NewHooks()
-	defer ddHooks.Stop()
-
-	hooks := &server.Hooks{}
-	ddHooks.AddHooks(hooks)
-
 	srv := server.NewMCPServer("test-server", "1.0.0",
-		server.WithHooks(hooks))
+		server.WithToolHandlerMiddleware(NewToolHandlerMiddleware()))
 
 	errorTool := mcp.NewTool("error_tool",
 		mcp.WithDescription("A tool that always errors"))