Skip to content

Commit 9ae69b2

Browse files
committed
Make middleware factory a private constant
1 parent cede691 commit 9ae69b2

File tree

3 files changed

+24
-26
lines changed

3 files changed

+24
-26
lines changed

contrib/mark3labs/mcp-go/option.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,6 @@ func WithTracing(options *TracingConfig) server.ServerOption {
3434

3535
server.WithHooks(hooks)(s)
3636

37-
server.WithToolHandlerMiddleware(NewToolHandlerMiddleware())(s)
37+
server.WithToolHandlerMiddleware(toolHandlerMiddleware)(s)
3838
}
3939
}

contrib/mark3labs/mcp-go/tracing.go

Lines changed: 21 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -39,36 +39,34 @@ func appendTracingHooks(hooks *server.Hooks) {
3939
hooks.AddOnError(tracingHooks.onError)
4040
}
4141

42-
func NewToolHandlerMiddleware() server.ToolHandlerMiddleware {
43-
return func(next server.ToolHandlerFunc) server.ToolHandlerFunc {
44-
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
45-
toolSpan, ctx := llmobs.StartToolSpan(ctx, request.Params.Name, llmobs.WithIntegration(string(instrumentation.PackageMark3LabsMCPGo)))
42+
var toolHandlerMiddleware = func(next server.ToolHandlerFunc) server.ToolHandlerFunc {
43+
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
44+
toolSpan, ctx := llmobs.StartToolSpan(ctx, request.Params.Name, llmobs.WithIntegration(string(instrumentation.PackageMark3LabsMCPGo)))
4645

47-
result, err := next(ctx, request)
46+
result, err := next(ctx, request)
4847

49-
inputJSON, marshalErr := json.Marshal(request)
48+
inputJSON, marshalErr := json.Marshal(request)
49+
if marshalErr != nil {
50+
instr.Logger().Warn("mcp-go: failed to marshal tool request: %v", marshalErr)
51+
}
52+
var outputText string
53+
if result != nil {
54+
resultJSON, marshalErr := json.Marshal(result)
5055
if marshalErr != nil {
51-
instr.Logger().Warn("mcp-go: failed to marshal tool request: %v", marshalErr)
52-
}
53-
var outputText string
54-
if result != nil {
55-
resultJSON, marshalErr := json.Marshal(result)
56-
if marshalErr != nil {
57-
instr.Logger().Warn("mcp-go: failed to marshal tool result: %v", marshalErr)
58-
}
59-
outputText = string(resultJSON)
56+
instr.Logger().Warn("mcp-go: failed to marshal tool result: %v", marshalErr)
6057
}
58+
outputText = string(resultJSON)
59+
}
6160

62-
toolSpan.AnnotateTextIO(string(inputJSON), outputText)
63-
64-
if err != nil {
65-
toolSpan.Finish(llmobs.WithError(err))
66-
} else {
67-
toolSpan.Finish()
68-
}
61+
toolSpan.AnnotateTextIO(string(inputJSON), outputText)
6962

70-
return result, err
63+
if err != nil {
64+
toolSpan.Finish(llmobs.WithError(err))
65+
} else {
66+
toolSpan.Finish()
7167
}
68+
69+
return result, err
7270
}
7371
}
7472

contrib/mark3labs/mcp-go/tracing_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,11 @@ import (
2121
"github.com/stretchr/testify/require"
2222
)
2323

24-
func TestNewToolHandlerMiddleware(t *testing.T) {
24+
func TestToolHandlerMiddleware(t *testing.T) {
2525
mt := mocktracer.Start()
2626
defer mt.Stop()
2727

28-
middleware := NewToolHandlerMiddleware()
28+
middleware := toolHandlerMiddleware
2929
assert.NotNil(t, middleware)
3030
}
3131

0 commit comments

Comments
 (0)