feat(contrib/mcp-go): Simplified way to add tracing

[jboolean]

What does this PR do?

Provides a simpler, more extensible, and less leaky way to add tracing to mcp-go.

Motivation

Previously I proposed a version which used reflection, but changed approaches: #4115

Reviewer's Checklist

• Changed code has unit tests for its functionality at or near 100% coverage.
• System-Tests covering this feature have been added and enabled with the va.b.c-dev version tag.
• There is a benchmark for any new code, or changes to existing code.
• If this interacts with the agent in a new way, a system test has been added.
• New code is free of linting errors. You can check this by running ./scripts/lint.sh locally.
• Add an appropriate team label so this PR gets put in the right place for the release notes.
• Non-trivial go.mod changes, e.g. adding new modules, are reviewed by @DataDog/dd-trace-go-guild.

Unsure? Have a question? Request a review!

[jboolean]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• feat(contrib/mcp-go): Intent capture #4123
• feat(contrib/mcp-go): Simplified way to add tracing #4122 👈 (View in Graphite)
• feat(contrib/mcp-go): Trace MCP session initializations with MLObs spans #4101 : 2 other dependent PRs (#4106 , #4115 )
• feat(contrib/mcp-go): Initial mcp-go tracer implementation #4100
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[jboolean]

In the future (next PR in stack) a boolean option to enable intent capture will be added to this config.

[pr-commenter]

Benchmarks

Benchmark execution time: 2025-11-14 15:38:38

Comparing candidate commit 58df1a2 in PR branch jb/simplified-way-to-add-tracing with baseline commit 7f50759 in branch jb/contrib-mcp-go-init-span.

Found 0 performance improvements and 0 performance regressions! Performance is the same for 6 metrics, 0 unstable metrics.

[datadog-official]

✅ Tests

🎉 All green!

❄️ No new flaky tests detected
🧪 All tests passed

_{This comment will be updated automatically if new data arrives.
🔗 Commit SHA: 58df1a2 | Docs | Datadog PR Page | Was this helpful? Give us feedback!}

[genesor]

Suggested change

	hooks := options.Hooks
	// Append hooks (hooks is a private field)
	if hooks == nil {
	hooks = &server.Hooks{}
	}
	AddServerHooks(hooks)
	server.WithHooks(hooks)(s)
	server.WithToolHandlerMiddleware(NewToolHandlerMiddleware())(s)
	if options == nil {
	option = new(TracingConfig)
	}
	// Add tracing hooks to server
	AddServerHooks(options.Hooks)
	server.WithHooks(options.Hooks)(s)
	server.WithToolHandlerMiddleware(NewToolHandlerMiddleware())(s)

This would be cleaner, you could also rename AddServerHooks to AddTracingHooks so we know what is happening here and handle the hooks == nil in the func too.

This still does not protect us from having our hooks overwritten by another server.WithHooks if the customer has one in their code.

[jboolean]

Should add AddServerHooks be made private?

[jboolean]

I cannot move creation of the empty hooks object into AddServerHooks because it must mutate its input in order for that input to later be passed to WithHooks by the caller. Unless I misunderstood something.

[genesor]

Should add AddServerHooks be made private?

I'd say yes but still provide a way to have hooks added separately in case the user wants to.

My comment #4115 (comment) would still apply. Make it simpler for the user and us to understand. No need to have NewXX funcs, I'd go for clarity and transparent use whenever I could.

WithTracing would become a wrapper around WithHooks and WithToolHandlerMiddleware

func WithHooks(hooks *server.Hooks) server.ServerOption {
	addServerHooks(hooks) // handle `nil` hooks inside

	return server.WithHooks(hooks)
}

func WithToolHandlerMiddleware() server.ServerOption {
	return server.WithToolHandlerMiddleware(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, 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, marshalErr := json.Marshal(result)
				if marshalErr != nil {
					instr.Logger().Warn("mcp-go: failed to marshal tool result: %v", marshalErr)
				}
				outputText = string(resultJSON)
			}

			toolSpan.AnnotateTextIO(string(inputJSON), outputText)

			if err != nil {
				toolSpan.Finish(llmobs.WithError(err))
			} else {
				toolSpan.Finish()
			}

			return result, err
		}
	})
}

func WithTracing(cfg *TracingConfig) server.ServerOption {
	return func(s *server.MCPServer) {
		if cfg == nil {
			cfg = new(TracingConfig)
		}

		// Add tracing hooks & middleware to server
		WithHooks(cfg.Hooks)(s)
		WithToolHandlerMiddleware()(s)
	}
}

[jboolean]

I think I still like having things organized into separate files with the responsibility of defining the ServerOption in its own file. But I renamed the public factory method to a private variable and renamed the other file to tracing to clarify its purpose.