fix(jaegermcp): Enforce response limits

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

@@ -21,15 +21,18 @@ import (
 // This tool retrieves all spans with error status from a specific trace, returning full
 // OTLP span details including attributes, events, and links for error analysis.
 type getTraceErrorsHandler struct {
-	queryService queryServiceGetTracesInterface
+	queryService             queryServiceGetTracesInterface
+	maxSpanDetailsPerRequest int
 }
 
 // NewGetTraceErrorsHandler creates a new get_trace_errors handler and returns the handler function.
 func NewGetTraceErrorsHandler(
 	queryService *querysvc.QueryService,
+	maxSpanDetailsPerRequest int,
 ) mcp.ToolHandlerFor[types.GetTraceErrorsInput, types.GetTraceErrorsOutput] {
 	h := &getTraceErrorsHandler{
-		queryService: queryService,
+		queryService:             queryService,
+		maxSpanDetailsPerRequest: maxSpanDetailsPerRequest,
 	}
 	return h.handle
 }
@@ -48,11 +51,13 @@ func (h *getTraceErrorsHandler) handle(
 
 	tracesIter := h.queryService.GetTraces(ctx, params)
 
-	// Wrap with AggregateTraces to ensure each ptrace.Traces contains a complete trace
-	aggregatedIter := jptrace.AggregateTraces(tracesIter)
+	// AggregateTracesWithLimit ensures a complete trace view while bounding server-side
+	// memory to maxSpanDetailsPerRequest spans, preventing unbounded work on large traces.
+	aggregatedIter := jptrace.AggregateTracesWithLimit(tracesIter, h.maxSpanDetailsPerRequest)
 
 	// Collect spans with error status
 	var errorSpans []types.SpanDetail
+	totalErrors := 0
 	traceFound := false
 
 	for trace, err := range aggregatedIter {
@@ -66,8 +71,12 @@ func (h *getTraceErrorsHandler) handle(
 		for pos, span := range jptrace.SpanIter(trace) {
 			// Check if span has error status
 			if span.Status().Code() == ptrace.StatusCodeError {
-				detail := buildSpanDetail(pos, span)
-				errorSpans = append(errorSpans, detail)
+				totalErrors++
+				// Only build and collect detail up to the limit
+				if h.maxSpanDetailsPerRequest == 0 || len(errorSpans) < h.maxSpanDetailsPerRequest {
+					detail := buildSpanDetail(pos, span)
+					errorSpans = append(errorSpans, detail)
+				}
 			}
 		}
 	}
@@ -78,7 +87,7 @@ func (h *getTraceErrorsHandler) handle(
 
 	output := types.GetTraceErrorsOutput{
 		TraceID:    input.TraceID,
-		ErrorCount: len(errorSpans),
+		ErrorCount: totalErrors,
 		Spans:      errorSpans,
 	}
 

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

@@ -153,7 +153,7 @@ func TestGetTraceErrorsHandler_Handle_SingleError(t *testing.T) {
 }
 
 func TestGetTraceErrorsHandler_Handle_MissingTraceID(t *testing.T) {
-	handler := NewGetTraceErrorsHandler(nil)
+	handler := NewGetTraceErrorsHandler(nil, 100)
 
 	input := types.GetTraceErrorsInput{
 		TraceID: "",
@@ -166,7 +166,7 @@ func TestGetTraceErrorsHandler_Handle_MissingTraceID(t *testing.T) {
 }
 
 func TestGetTraceErrorsHandler_Handle_InvalidTraceID(t *testing.T) {
-	handler := NewGetTraceErrorsHandler(nil)
+	handler := NewGetTraceErrorsHandler(nil, 100)
 
 	input := types.GetTraceErrorsInput{
 		TraceID: "invalid-trace-id",
@@ -372,3 +372,35 @@ func TestGetTraceErrorsHandler_Handle_ErrorSpanWithEvents(t *testing.T) {
 	assert.Equal(t, "RuntimeError", span.Events[0].Attributes["exception.type"])
 	assert.Equal(t, "Something went wrong", span.Events[0].Attributes["exception.message"])
 }
+
+func TestGetTraceErrorsHandler_Handle_LimitEnforced(t *testing.T) {
+	traceID := testTraceID
+
+	// Create 5 error spans
+	spanConfigs := []spanConfig{
+		{spanID: "span001", operation: "/api/error1", hasError: true, errorMessage: "err1"},
+		{spanID: "span002", operation: "/api/error2", hasError: true, errorMessage: "err2"},
+		{spanID: "span003", operation: "/api/error3", hasError: true, errorMessage: "err3"},
+		{spanID: "span004", operation: "/api/error4", hasError: true, errorMessage: "err4"},
+		{spanID: "span005", operation: "/api/error5", hasError: true, errorMessage: "err5"},
+	}
+
+	testTrace := createTestTraceWithSpans(traceID, spanConfigs)
+	mock := newMockYieldingTraces(testTrace)
+
+	// Set limit to 3 — should return at most 3 spans
+	handler := &getTraceErrorsHandler{
+		queryService:             mock,
+		maxSpanDetailsPerRequest: 3,
+	}
+
+	input := types.GetTraceErrorsInput{TraceID: traceID}
+	_, output, err := handler.handle(context.Background(), &mcp.CallToolRequest{}, input)
+
+	require.NoError(t, err)
+	// AggregateTracesWithLimit caps the trace at limit spans before error counting,
+	// so ErrorCount and len(Spans) are both bounded by maxSpanDetailsPerRequest.
+	assert.Equal(t, 3, output.ErrorCount)
+	// Returned spans are capped at exactly the limit (5 errors, limit=3 → exactly 3 spans)
+	assert.Len(t, output.Spans, 3)
+}

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

@@ -24,15 +24,18 @@ import (
 // This tool returns the structural tree of a trace showing parent-child relationships,
 // timing, and error locations WITHOUT returning attributes or logs to keep the response compact.
 type getTraceTopologyHandler struct {
-	queryService queryServiceGetTracesInterface
+	queryService             queryServiceGetTracesInterface
+	maxSpanDetailsPerRequest int
 }
 
 // NewGetTraceTopologyHandler creates a new get_trace_topology handler and returns the handler function.
 func NewGetTraceTopologyHandler(
 	queryService *querysvc.QueryService,
+	maxSpanDetailsPerRequest int,
 ) mcp.ToolHandlerFor[types.GetTraceTopologyInput, types.GetTraceTopologyOutput] {
 	h := &getTraceTopologyHandler{
-		queryService: queryService,
+		queryService:             queryService,
+		maxSpanDetailsPerRequest: maxSpanDetailsPerRequest,
 	}
 	return h.handle
 }
@@ -63,8 +66,9 @@ func (h *getTraceTopologyHandler) handle(
 
 	tracesIter := h.queryService.GetTraces(ctx, params)
 
-	// Wrap with AggregateTraces to ensure each ptrace.Traces contains a complete trace
-	aggregatedIter := jptrace.AggregateTraces(tracesIter)
+	// AggregateTracesWithLimit ensures a complete trace view while bounding server-side
+	// memory to maxSpanDetailsPerRequest spans, preventing unbounded work on large traces.
+	aggregatedIter := jptrace.AggregateTracesWithLimit(tracesIter, h.maxSpanDetailsPerRequest)
 
 	// Collect all spans from the trace
 	var spans []rawSpan
@@ -80,6 +84,9 @@ func (h *getTraceTopologyHandler) handle(
 		// Iterate through all spans in the trace and collect them
 		for pos, span := range jptrace.SpanIter(trace) {
 			spans = append(spans, extractRawSpan(pos, span))
+			if h.maxSpanDetailsPerRequest > 0 && len(spans) >= h.maxSpanDetailsPerRequest {
+				break
+			}
 		}
 	}
 

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

@@ -338,7 +338,7 @@ func TestGetTraceTopologyHandler_Handle_NoAttributes(t *testing.T) {
 }
 
 func TestGetTraceTopologyHandler_Handle_MissingTraceID(t *testing.T) {
-	handler := NewGetTraceTopologyHandler(nil)
+	handler := NewGetTraceTopologyHandler(nil, 100)
 
 	_, _, err := handler(context.Background(), &mcp.CallToolRequest{}, types.GetTraceTopologyInput{})
 
@@ -347,7 +347,7 @@ func TestGetTraceTopologyHandler_Handle_MissingTraceID(t *testing.T) {
 }
 
 func TestGetTraceTopologyHandler_Handle_InvalidTraceID(t *testing.T) {
-	handler := NewGetTraceTopologyHandler(nil)
+	handler := NewGetTraceTopologyHandler(nil, 100)
 
 	input := types.GetTraceTopologyInput{TraceID: "invalid-trace-id"}
 	_, _, err := handler(context.Background(), &mcp.CallToolRequest{}, input)
@@ -569,3 +569,33 @@ func TestGetTraceTopologyHandler_Handle_DFSOrder(t *testing.T) {
 	assert.Less(t, indexOf["C"], indexOf["B"])
 	assert.Less(t, indexOf["D"], indexOf["B"])
 }
+
+func TestGetTraceTopologyHandler_Handle_LimitEnforced(t *testing.T) {
+	traceID := testTraceID
+
+	// Create a trace with 6 spans
+	spanConfigs := []spanConfig{
+		{spanID: "root001", operation: "root"},
+		{spanID: "span002", parentSpanID: "root001", operation: "child1"},
+		{spanID: "span003", parentSpanID: "root001", operation: "child2"},
+		{spanID: "span004", parentSpanID: "span002", operation: "grandchild1"},
+		{spanID: "span005", parentSpanID: "span002", operation: "grandchild2"},
+		{spanID: "span006", parentSpanID: "span003", operation: "grandchild3"},
+	}
+
+	testTrace := createTestTraceWithSpans(traceID, spanConfigs)
+	mock := newMockYieldingTraces(testTrace)
+
+	// Set limit to 3 — should collect at most 3 spans before building topology
+	handler := &getTraceTopologyHandler{
+		queryService:             mock,
+		maxSpanDetailsPerRequest: 3,
+	}
+
+	input := types.GetTraceTopologyInput{TraceID: traceID}
+	_, output, err := handler.handle(context.Background(), &mcp.CallToolRequest{}, input)
+
+	require.NoError(t, err)
+	// Exactly 3 spans returned — 6-span trace with limit=3 must truncate to exactly 3
+	assert.Len(t, output.Spans, 3)
+}

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

@@ -78,6 +78,9 @@ func (h *searchTracesHandler) handle(
 
 		summary := buildTraceSummary(trace)
 		summaries = append(summaries, summary)
+		if h.maxResults > 0 && len(summaries) >= h.maxResults {
+			break
+		}
 	}
 
 	output := types.SearchTracesOutput{Traces: summaries}
@@ -144,7 +147,7 @@ func (h *searchTracesHandler) buildQuery(input types.SearchTracesInput) (querysv
 	if searchDepth <= 0 {
 		searchDepth = defaultSearchDepth
 	}
-	if searchDepth > h.maxResults {
+	if h.maxResults > 0 && searchDepth > h.maxResults {
 		searchDepth = h.maxResults
 	}
 

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

@@ -476,3 +476,40 @@ func TestSearchTracesHandler_Handle_DefaultStartTime(t *testing.T) {
 	require.NoError(t, err)
 	require.Len(t, output.Traces, 1)
 }
+
+func TestSearchTracesHandler_Handle_LimitEnforced(t *testing.T) {
+	// Create 5 distinct traces
+	traces := []ptrace.Traces{
+		createTestTrace("trace001", "svc", "/op1", false),
+		createTestTrace("trace002", "svc", "/op2", false),
+		createTestTrace("trace003", "svc", "/op3", false),
+		createTestTrace("trace004", "svc", "/op4", false),
+		createTestTrace("trace005", "svc", "/op5", false),
+	}
+
+	mock := &mockQueryService{
+		findTracesFunc: func(_ context.Context, _ querysvc.TraceQueryParams) iter.Seq2[[]ptrace.Traces, error] {
+			return func(yield func([]ptrace.Traces, error) bool) {
+				for _, tr := range traces {
+					if !yield([]ptrace.Traces{tr}, nil) {
+						return
+					}
+				}
+			}
+		},
+	}
+
+	// Set maxResults to 3 — should return at most 3 traces
+	handler := &searchTracesHandler{queryService: mock, maxResults: 3}
+
+	input := types.SearchTracesInput{
+		StartTimeMin: "-1h",
+		ServiceName:  "svc",
+	}
+
+	_, output, err := handler.handle(context.Background(), &mcp.CallToolRequest{}, input)
+
+	require.NoError(t, err)
+	// Returned traces are capped at exactly the limit (5 traces, limit=3 → exactly 3 traces)
+	assert.Len(t, output.Traces, 3)
+}

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

@@ -155,12 +155,12 @@ func (s *server) registerTools() {
 	mcp.AddTool(s.mcpServer, &mcp.Tool{
 		Name:        "get_trace_errors",
 		Description: "Get full details for all spans with error status.",
-	}, handlers.NewGetTraceErrorsHandler(s.queryAPI))
+	}, handlers.NewGetTraceErrorsHandler(s.queryAPI, s.config.MaxSpanDetailsPerRequest))
 
 	mcp.AddTool(s.mcpServer, &mcp.Tool{
 		Name:        "get_trace_topology",
 		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.",
-	}, handlers.NewGetTraceTopologyHandler(s.queryAPI))
+	}, handlers.NewGetTraceTopologyHandler(s.queryAPI, s.config.MaxSpanDetailsPerRequest))
 
 	mcp.AddTool(s.mcpServer, &mcp.Tool{
 		Name:        "get_critical_path",