Skip to content

Commit 8a4c314

Browse files
abhay1999claude
andauthored
fix(jaegermcp): Enforce response limits (#8242)
## What this PR does `MaxSpanDetailsPerRequest` and `MaxSearchResults` were defined in the config but not consistently enforced across all MCP handlers. This PR wires the limits all the way through. ## Changes **`get_trace_errors.go`** - Continues iterating all spans to count total errors, but caps the `Spans` slice at `maxSpanDetailsPerRequest` - `ErrorCount` now reflects the **true total** error count, not the truncated slice length **`get_trace_topology.go`** - Breaks collection loop once `maxSpanDetailsPerRequest` is reached **`search_traces.go`** - Breaks aggregation loop once `maxResults` is reached - Fixed `searchDepth` clamping: guards with `maxResults > 0` so unlimited (`0`) is not incorrectly clamped to `0` **`server.go`** - Wires `MaxSpanDetailsPerRequest` into `NewGetTraceErrorsHandler` and `NewGetTraceTopologyHandler` constructors ## Tests added - `TestGetTraceErrorsHandler_Handle_LimitEnforced` — 5 error spans, limit 3, asserts `ErrorCount==5` and `len(Spans)<=3` - `TestGetTraceTopologyHandler_Handle_LimitEnforced` — 6-span trace, limit 3, asserts `len(Spans)<=3` - `TestSearchTracesHandler_Handle_LimitEnforced` — 5 traces, limit 3, asserts `len(Traces)<=3` Fixes #8118 --------- Signed-off-by: abhay1999 <abhaychaurasiya19@gmail.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent e75b845 commit 8a4c314

File tree

7 files changed

+133
-16
lines changed

7 files changed

+133
-16
lines changed

cmd/jaeger/internal/extension/jaegermcp/internal/handlers/get_trace_errors.go

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,15 +21,18 @@ import (
2121
// This tool retrieves all spans with error status from a specific trace, returning full
2222
// OTLP span details including attributes, events, and links for error analysis.
2323
type getTraceErrorsHandler struct {
24-
queryService queryServiceGetTracesInterface
24+
queryService queryServiceGetTracesInterface
25+
maxSpanDetailsPerRequest int
2526
}
2627

2728
// NewGetTraceErrorsHandler creates a new get_trace_errors handler and returns the handler function.
2829
func NewGetTraceErrorsHandler(
2930
queryService *querysvc.QueryService,
31+
maxSpanDetailsPerRequest int,
3032
) mcp.ToolHandlerFor[types.GetTraceErrorsInput, types.GetTraceErrorsOutput] {
3133
h := &getTraceErrorsHandler{
32-
queryService: queryService,
34+
queryService: queryService,
35+
maxSpanDetailsPerRequest: maxSpanDetailsPerRequest,
3336
}
3437
return h.handle
3538
}
@@ -48,11 +51,13 @@ func (h *getTraceErrorsHandler) handle(
4851

4952
tracesIter := h.queryService.GetTraces(ctx, params)
5053

51-
// Wrap with AggregateTraces to ensure each ptrace.Traces contains a complete trace
54+
// AggregateTraces reassembles the full trace so ErrorCount reflects every error span.
55+
// The Spans slice is capped inside the iteration loop below.
5256
aggregatedIter := jptrace.AggregateTraces(tracesIter)
5357

5458
// Collect spans with error status
5559
var errorSpans []types.SpanDetail
60+
totalErrors := 0
5661
traceFound := false
5762

5863
for trace, err := range aggregatedIter {
@@ -66,8 +71,12 @@ func (h *getTraceErrorsHandler) handle(
6671
for pos, span := range jptrace.SpanIter(trace) {
6772
// Check if span has error status
6873
if span.Status().Code() == ptrace.StatusCodeError {
69-
detail := buildSpanDetail(pos, span)
70-
errorSpans = append(errorSpans, detail)
74+
totalErrors++
75+
// Only build and collect detail up to the limit
76+
if h.maxSpanDetailsPerRequest == 0 || len(errorSpans) < h.maxSpanDetailsPerRequest {
77+
detail := buildSpanDetail(pos, span)
78+
errorSpans = append(errorSpans, detail)
79+
}
7180
}
7281
}
7382
}
@@ -78,7 +87,7 @@ func (h *getTraceErrorsHandler) handle(
7887

7988
output := types.GetTraceErrorsOutput{
8089
TraceID: input.TraceID,
81-
ErrorCount: len(errorSpans),
90+
ErrorCount: totalErrors,
8291
Spans: errorSpans,
8392
}
8493

cmd/jaeger/internal/extension/jaegermcp/internal/handlers/get_trace_errors_test.go

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,7 @@ func TestGetTraceErrorsHandler_Handle_SingleError(t *testing.T) {
153153
}
154154

155155
func TestGetTraceErrorsHandler_Handle_MissingTraceID(t *testing.T) {
156-
handler := NewGetTraceErrorsHandler(nil)
156+
handler := NewGetTraceErrorsHandler(nil, 100)
157157

158158
input := types.GetTraceErrorsInput{
159159
TraceID: "",
@@ -166,7 +166,7 @@ func TestGetTraceErrorsHandler_Handle_MissingTraceID(t *testing.T) {
166166
}
167167

168168
func TestGetTraceErrorsHandler_Handle_InvalidTraceID(t *testing.T) {
169-
handler := NewGetTraceErrorsHandler(nil)
169+
handler := NewGetTraceErrorsHandler(nil, 100)
170170

171171
input := types.GetTraceErrorsInput{
172172
TraceID: "invalid-trace-id",
@@ -372,3 +372,34 @@ func TestGetTraceErrorsHandler_Handle_ErrorSpanWithEvents(t *testing.T) {
372372
assert.Equal(t, "RuntimeError", span.Events[0].Attributes["exception.type"])
373373
assert.Equal(t, "Something went wrong", span.Events[0].Attributes["exception.message"])
374374
}
375+
376+
func TestGetTraceErrorsHandler_Handle_LimitEnforced(t *testing.T) {
377+
traceID := testTraceID
378+
379+
// Create 5 error spans
380+
spanConfigs := []spanConfig{
381+
{spanID: "span001", operation: "/api/error1", hasError: true, errorMessage: "err1"},
382+
{spanID: "span002", operation: "/api/error2", hasError: true, errorMessage: "err2"},
383+
{spanID: "span003", operation: "/api/error3", hasError: true, errorMessage: "err3"},
384+
{spanID: "span004", operation: "/api/error4", hasError: true, errorMessage: "err4"},
385+
{spanID: "span005", operation: "/api/error5", hasError: true, errorMessage: "err5"},
386+
}
387+
388+
testTrace := createTestTraceWithSpans(traceID, spanConfigs)
389+
mock := newMockYieldingTraces(testTrace)
390+
391+
// Set limit to 3 — should return at most 3 spans
392+
handler := &getTraceErrorsHandler{
393+
queryService: mock,
394+
maxSpanDetailsPerRequest: 3,
395+
}
396+
397+
input := types.GetTraceErrorsInput{TraceID: traceID}
398+
_, output, err := handler.handle(context.Background(), &mcp.CallToolRequest{}, input)
399+
400+
require.NoError(t, err)
401+
// ErrorCount reflects all error spans in the full trace (unbounded aggregation).
402+
assert.Equal(t, 5, output.ErrorCount)
403+
// Returned Spans are capped at exactly the limit (5 errors, limit=3 → exactly 3 spans).
404+
assert.Len(t, output.Spans, 3)
405+
}

cmd/jaeger/internal/extension/jaegermcp/internal/handlers/get_trace_topology.go

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,15 +24,18 @@ import (
2424
// This tool returns the structural tree of a trace showing parent-child relationships,
2525
// timing, and error locations WITHOUT returning attributes or logs to keep the response compact.
2626
type getTraceTopologyHandler struct {
27-
queryService queryServiceGetTracesInterface
27+
queryService queryServiceGetTracesInterface
28+
maxSpanDetailsPerRequest int
2829
}
2930

3031
// NewGetTraceTopologyHandler creates a new get_trace_topology handler and returns the handler function.
3132
func NewGetTraceTopologyHandler(
3233
queryService *querysvc.QueryService,
34+
maxSpanDetailsPerRequest int,
3335
) mcp.ToolHandlerFor[types.GetTraceTopologyInput, types.GetTraceTopologyOutput] {
3436
h := &getTraceTopologyHandler{
35-
queryService: queryService,
37+
queryService: queryService,
38+
maxSpanDetailsPerRequest: maxSpanDetailsPerRequest,
3639
}
3740
return h.handle
3841
}
@@ -63,8 +66,9 @@ func (h *getTraceTopologyHandler) handle(
6366

6467
tracesIter := h.queryService.GetTraces(ctx, params)
6568

66-
// Wrap with AggregateTraces to ensure each ptrace.Traces contains a complete trace
67-
aggregatedIter := jptrace.AggregateTraces(tracesIter)
69+
// AggregateTracesWithLimit ensures a complete trace view while bounding server-side
70+
// memory to maxSpanDetailsPerRequest spans, preventing unbounded work on large traces.
71+
aggregatedIter := jptrace.AggregateTracesWithLimit(tracesIter, h.maxSpanDetailsPerRequest)
6872

6973
// Collect all spans from the trace
7074
var spans []rawSpan
@@ -80,6 +84,9 @@ func (h *getTraceTopologyHandler) handle(
8084
// Iterate through all spans in the trace and collect them
8185
for pos, span := range jptrace.SpanIter(trace) {
8286
spans = append(spans, extractRawSpan(pos, span))
87+
if h.maxSpanDetailsPerRequest > 0 && len(spans) >= h.maxSpanDetailsPerRequest {
88+
break
89+
}
8390
}
8491
}
8592

cmd/jaeger/internal/extension/jaegermcp/internal/handlers/get_trace_topology_test.go

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -338,7 +338,7 @@ func TestGetTraceTopologyHandler_Handle_NoAttributes(t *testing.T) {
338338
}
339339

340340
func TestGetTraceTopologyHandler_Handle_MissingTraceID(t *testing.T) {
341-
handler := NewGetTraceTopologyHandler(nil)
341+
handler := NewGetTraceTopologyHandler(nil, 100)
342342

343343
_, _, err := handler(context.Background(), &mcp.CallToolRequest{}, types.GetTraceTopologyInput{})
344344

@@ -347,7 +347,7 @@ func TestGetTraceTopologyHandler_Handle_MissingTraceID(t *testing.T) {
347347
}
348348

349349
func TestGetTraceTopologyHandler_Handle_InvalidTraceID(t *testing.T) {
350-
handler := NewGetTraceTopologyHandler(nil)
350+
handler := NewGetTraceTopologyHandler(nil, 100)
351351

352352
input := types.GetTraceTopologyInput{TraceID: "invalid-trace-id"}
353353
_, _, err := handler(context.Background(), &mcp.CallToolRequest{}, input)
@@ -569,3 +569,33 @@ func TestGetTraceTopologyHandler_Handle_DFSOrder(t *testing.T) {
569569
assert.Less(t, indexOf["C"], indexOf["B"])
570570
assert.Less(t, indexOf["D"], indexOf["B"])
571571
}
572+
573+
func TestGetTraceTopologyHandler_Handle_LimitEnforced(t *testing.T) {
574+
traceID := testTraceID
575+
576+
// Create a trace with 6 spans
577+
spanConfigs := []spanConfig{
578+
{spanID: "root001", operation: "root"},
579+
{spanID: "span002", parentSpanID: "root001", operation: "child1"},
580+
{spanID: "span003", parentSpanID: "root001", operation: "child2"},
581+
{spanID: "span004", parentSpanID: "span002", operation: "grandchild1"},
582+
{spanID: "span005", parentSpanID: "span002", operation: "grandchild2"},
583+
{spanID: "span006", parentSpanID: "span003", operation: "grandchild3"},
584+
}
585+
586+
testTrace := createTestTraceWithSpans(traceID, spanConfigs)
587+
mock := newMockYieldingTraces(testTrace)
588+
589+
// Set limit to 3 — should collect at most 3 spans before building topology
590+
handler := &getTraceTopologyHandler{
591+
queryService: mock,
592+
maxSpanDetailsPerRequest: 3,
593+
}
594+
595+
input := types.GetTraceTopologyInput{TraceID: traceID}
596+
_, output, err := handler.handle(context.Background(), &mcp.CallToolRequest{}, input)
597+
598+
require.NoError(t, err)
599+
// Exactly 3 spans returned — 6-span trace with limit=3 must truncate to exactly 3
600+
assert.Len(t, output.Spans, 3)
601+
}

cmd/jaeger/internal/extension/jaegermcp/internal/handlers/search_traces.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,9 @@ func (h *searchTracesHandler) handle(
7878

7979
summary := buildTraceSummary(trace)
8080
summaries = append(summaries, summary)
81+
if h.maxResults > 0 && len(summaries) >= h.maxResults {
82+
break
83+
}
8184
}
8285

8386
output := types.SearchTracesOutput{Traces: summaries}

cmd/jaeger/internal/extension/jaegermcp/internal/handlers/search_traces_test.go

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -476,3 +476,40 @@ func TestSearchTracesHandler_Handle_DefaultStartTime(t *testing.T) {
476476
require.NoError(t, err)
477477
require.Len(t, output.Traces, 1)
478478
}
479+
480+
func TestSearchTracesHandler_Handle_LimitEnforced(t *testing.T) {
481+
// Create 5 distinct traces
482+
traces := []ptrace.Traces{
483+
createTestTrace("trace001", "svc", "/op1", false),
484+
createTestTrace("trace002", "svc", "/op2", false),
485+
createTestTrace("trace003", "svc", "/op3", false),
486+
createTestTrace("trace004", "svc", "/op4", false),
487+
createTestTrace("trace005", "svc", "/op5", false),
488+
}
489+
490+
mock := &mockQueryService{
491+
findTracesFunc: func(_ context.Context, _ querysvc.TraceQueryParams) iter.Seq2[[]ptrace.Traces, error] {
492+
return func(yield func([]ptrace.Traces, error) bool) {
493+
for _, tr := range traces {
494+
if !yield([]ptrace.Traces{tr}, nil) {
495+
return
496+
}
497+
}
498+
}
499+
},
500+
}
501+
502+
// Set maxResults to 3 — should return at most 3 traces
503+
handler := &searchTracesHandler{queryService: mock, maxResults: 3}
504+
505+
input := types.SearchTracesInput{
506+
StartTimeMin: "-1h",
507+
ServiceName: "svc",
508+
}
509+
510+
_, output, err := handler.handle(context.Background(), &mcp.CallToolRequest{}, input)
511+
512+
require.NoError(t, err)
513+
// Returned traces are capped at exactly the limit (5 traces, limit=3 → exactly 3 traces)
514+
assert.Len(t, output.Traces, 3)
515+
}

cmd/jaeger/internal/extension/jaegermcp/server.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -155,12 +155,12 @@ func (s *server) registerTools() {
155155
mcp.AddTool(s.mcpServer, &mcp.Tool{
156156
Name: "get_trace_errors",
157157
Description: "Get full details for all spans with error status.",
158-
}, handlers.NewGetTraceErrorsHandler(s.queryAPI))
158+
}, handlers.NewGetTraceErrorsHandler(s.queryAPI, s.config.MaxSpanDetailsPerRequest))
159159

160160
mcp.AddTool(s.mcpServer, &mcp.Tool{
161161
Name: "get_trace_topology",
162162
Description: "Get the structural topology of a trace as a flat, depth-first list of spans. Each span's 'path' field encodes ancestry as slash-delimited span IDs (e.g. rootID/parentID/spanID). Does NOT return attributes or logs.",
163-
}, handlers.NewGetTraceTopologyHandler(s.queryAPI))
163+
}, handlers.NewGetTraceTopologyHandler(s.queryAPI, s.config.MaxSpanDetailsPerRequest))
164164

165165
mcp.AddTool(s.mcpServer, &mcp.Tool{
166166
Name: "get_critical_path",

0 commit comments

Comments
 (0)