Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -30,15 +30,18 @@ type queryServiceGetTracesInterface interface {
// This tool fetches full OTLP details (attributes, events, links, status) for specific spans
// within a trace, allowing deep inspection of individual span data for debugging and analysis.
type getSpanDetailsHandler struct {
queryService queryServiceGetTracesInterface
queryService queryServiceGetTracesInterface
maxSpanDetailsPerRequest int
}

// NewGetSpanDetailsHandler creates a new get_span_details handler and returns the handler function.
func NewGetSpanDetailsHandler(
queryService *querysvc.QueryService,
maxSpanDetailsPerRequest int,
) mcp.ToolHandlerFor[types.GetSpanDetailsInput, types.GetSpanDetailsOutput] {
h := &getSpanDetailsHandler{
queryService: queryService,
queryService: queryService,
maxSpanDetailsPerRequest: maxSpanDetailsPerRequest,
}
Comment thread
yurishkuro marked this conversation as resolved.
return h.handle
}
Expand Down Expand Up @@ -115,7 +118,7 @@ func (h *getSpanDetailsHandler) handle(
}

// buildQuery converts GetSpanDetailsInput to querysvc.GetTraceParams.
func (*getSpanDetailsHandler) buildQuery(input types.GetSpanDetailsInput) (querysvc.GetTraceParams, error) {
func (h *getSpanDetailsHandler) buildQuery(input types.GetSpanDetailsInput) (querysvc.GetTraceParams, error) {
// Validate input
if input.TraceID == "" {
return querysvc.GetTraceParams{}, errors.New("trace_id is required")
Expand All @@ -125,6 +128,15 @@ func (*getSpanDetailsHandler) buildQuery(input types.GetSpanDetailsInput) (query
return querysvc.GetTraceParams{}, errors.New("span_ids is required and must not be empty")
}

// Validate span count against configured limit
if len(input.SpanIDs) > h.maxSpanDetailsPerRequest {
return querysvc.GetTraceParams{}, fmt.Errorf(
"span_ids exceeds maximum limit: requested %d, max allowed %d",
len(input.SpanIDs),
h.maxSpanDetailsPerRequest,
)
}

traceID, err := parseTraceID(input.TraceID)
if err != nil {
return querysvc.GetTraceParams{}, fmt.Errorf("invalid trace_id: %w", err)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ func TestGetSpanDetailsHandler_Handle_Success(t *testing.T) {

mock := newMockYieldingTraces(testTrace)

handler := &getSpanDetailsHandler{queryService: mock}
handler := &getSpanDetailsHandler{queryService: mock, maxSpanDetailsPerRequest: 50}

input := types.GetSpanDetailsInput{
TraceID: traceID,
Expand Down Expand Up @@ -115,7 +115,7 @@ func TestGetSpanDetailsHandler_Handle_SingleSpan(t *testing.T) {

mock := newMockYieldingTraces(testTrace)

handler := &getSpanDetailsHandler{queryService: mock}
handler := &getSpanDetailsHandler{queryService: mock, maxSpanDetailsPerRequest: 50}

input := types.GetSpanDetailsInput{
TraceID: traceID,
Expand Down Expand Up @@ -143,7 +143,7 @@ func TestGetSpanDetailsHandler_Handle_FiltersBySpanIDs(t *testing.T) {

mock := newMockYieldingTraces(testTrace)

handler := &getSpanDetailsHandler{queryService: mock}
handler := &getSpanDetailsHandler{queryService: mock, maxSpanDetailsPerRequest: 50}

input := types.GetSpanDetailsInput{
TraceID: traceID,
Expand All @@ -166,7 +166,7 @@ func TestGetSpanDetailsHandler_Handle_FiltersBySpanIDs(t *testing.T) {
}

func TestGetSpanDetailsHandler_Handle_MissingTraceID(t *testing.T) {
handler := NewGetSpanDetailsHandler(nil)
handler := NewGetSpanDetailsHandler(nil, 50)

input := types.GetSpanDetailsInput{
SpanIDs: []string{spanIDToHex("span001")},
Expand All @@ -179,7 +179,7 @@ func TestGetSpanDetailsHandler_Handle_MissingTraceID(t *testing.T) {
}

func TestGetSpanDetailsHandler_Handle_MissingSpanIDs(t *testing.T) {
handler := NewGetSpanDetailsHandler(nil)
handler := NewGetSpanDetailsHandler(nil, 50)

input := types.GetSpanDetailsInput{
TraceID: testTraceID,
Expand All @@ -193,7 +193,7 @@ func TestGetSpanDetailsHandler_Handle_MissingSpanIDs(t *testing.T) {
}

func TestGetSpanDetailsHandler_Handle_InvalidTraceID(t *testing.T) {
handler := NewGetSpanDetailsHandler(nil)
handler := NewGetSpanDetailsHandler(nil, 50)

input := types.GetSpanDetailsInput{
TraceID: "invalid-trace-id",
Expand All @@ -209,7 +209,7 @@ func TestGetSpanDetailsHandler_Handle_InvalidTraceID(t *testing.T) {
func TestGetSpanDetailsHandler_Handle_TraceNotFound(t *testing.T) {
mock := newMockYieldingEmpty()

handler := &getSpanDetailsHandler{queryService: mock}
handler := &getSpanDetailsHandler{queryService: mock, maxSpanDetailsPerRequest: 50}

input := types.GetSpanDetailsInput{
TraceID: testTraceID,
Expand All @@ -225,7 +225,7 @@ func TestGetSpanDetailsHandler_Handle_TraceNotFound(t *testing.T) {
func TestGetSpanDetailsHandler_Handle_QueryError(t *testing.T) {
mock := newMockYieldingError(errors.New("database connection failed"))

handler := &getSpanDetailsHandler{queryService: mock}
handler := &getSpanDetailsHandler{queryService: mock, maxSpanDetailsPerRequest: 50}

input := types.GetSpanDetailsInput{
TraceID: testTraceID,
Expand Down Expand Up @@ -262,7 +262,7 @@ func TestGetSpanDetailsHandler_Handle_PartialResults(t *testing.T) {
},
}

handler := &getSpanDetailsHandler{queryService: mock}
handler := &getSpanDetailsHandler{queryService: mock, maxSpanDetailsPerRequest: 50}

input := types.GetSpanDetailsInput{
TraceID: traceID,
Expand Down Expand Up @@ -300,7 +300,7 @@ func TestGetSpanDetailsHandler_Handle_MultipleIterations(t *testing.T) {
},
}

handler := &getSpanDetailsHandler{queryService: mock}
handler := &getSpanDetailsHandler{queryService: mock, maxSpanDetailsPerRequest: 50}

input := types.GetSpanDetailsInput{
TraceID: traceID,
Expand Down Expand Up @@ -335,7 +335,7 @@ func TestGetSpanDetailsHandler_Handle_WithParentSpanID(t *testing.T) {

mock := newMockYieldingTraces(testTrace)

handler := &getSpanDetailsHandler{queryService: mock}
handler := &getSpanDetailsHandler{queryService: mock, maxSpanDetailsPerRequest: 50}

input := types.GetSpanDetailsInput{
TraceID: traceID,
Expand All @@ -360,7 +360,7 @@ func TestGetSpanDetailsHandler_Handle_NoMatchingSpans(t *testing.T) {

mock := newMockYieldingTraces(testTrace)

handler := &getSpanDetailsHandler{queryService: mock}
handler := &getSpanDetailsHandler{queryService: mock, maxSpanDetailsPerRequest: 50}

input := types.GetSpanDetailsInput{
TraceID: traceID,
Expand Down Expand Up @@ -389,7 +389,7 @@ func TestGetSpanDetailsHandler_Handle_PartialMissingSpans(t *testing.T) {

mock := newMockYieldingTraces(testTrace)

handler := &getSpanDetailsHandler{queryService: mock}
handler := &getSpanDetailsHandler{queryService: mock, maxSpanDetailsPerRequest: 50}

input := types.GetSpanDetailsInput{
TraceID: traceID,
Expand Down Expand Up @@ -497,6 +497,24 @@ func TestConvertAttributeValue(t *testing.T) {
}
}

func TestGetSpanDetailsHandler_ExceedsLimit(t *testing.T) {
// Create handler with a limit of 2 spans per request
handler := NewGetSpanDetailsHandler(nil, 2)

input := types.GetSpanDetailsInput{
TraceID: testTraceID,
// Request 3 spans but limit is 2 - should fail before querying
SpanIDs: []string{spanIDToHex("span001"), spanIDToHex("span002"), spanIDToHex("span003")},
}

_, _, err := handler(context.Background(), &mcp.CallToolRequest{}, input)

require.Error(t, err)
assert.Contains(t, err.Error(), "exceeds maximum limit")
assert.Contains(t, err.Error(), "requested 3")
assert.Contains(t, err.Error(), "max allowed 2")
}

func TestParseTraceID(t *testing.T) {
tests := []struct {
name string
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,11 +21,6 @@ import (
"github.com/jaegertracing/jaeger/internal/storage/v2/api/tracestore"
)

const (
defaultSearchLimit = 10
maxSearchLimit = 100
)

// queryServiceInterface defines the interface we need from QueryService for testing
type queryServiceInterface interface {
FindTraces(ctx context.Context, query querysvc.TraceQueryParams) iter.Seq2[[]ptrace.Traces, error]
Expand All @@ -37,14 +32,17 @@ type queryServiceInterface interface {
// browsing and filtering large result sets.
type searchTracesHandler struct {
queryService queryServiceInterface
maxResults int
}

// NewSearchTracesHandler creates a new search_traces handler and returns the handler function.
func NewSearchTracesHandler(
queryService *querysvc.QueryService,
maxResults int,
) mcp.ToolHandlerFor[types.SearchTracesInput, types.SearchTracesOutput] {
h := &searchTracesHandler{
queryService: queryService,
maxResults: maxResults,
}
Comment thread
yurishkuro marked this conversation as resolved.
return h.handle
}
Expand Down Expand Up @@ -93,7 +91,7 @@ func (h *searchTracesHandler) handle(
}

// buildQuery converts SearchTracesInput to querysvc.TraceQueryParams.
func (*searchTracesHandler) buildQuery(input types.SearchTracesInput) (querysvc.TraceQueryParams, error) {
func (h *searchTracesHandler) buildQuery(input types.SearchTracesInput) (querysvc.TraceQueryParams, error) {
// Use default start time if not provided
startTimeMinInput := input.StartTimeMin
if startTimeMinInput == "" {
Expand Down Expand Up @@ -141,12 +139,13 @@ func (*searchTracesHandler) buildQuery(input types.SearchTracesInput) (querysvc.
}

// Set default and max search depth
const defaultSearchDepth = 10
searchDepth := input.SearchDepth
if searchDepth <= 0 {
searchDepth = defaultSearchLimit
searchDepth = defaultSearchDepth
}
if searchDepth > maxSearchLimit {
searchDepth = maxSearchLimit
if searchDepth > h.maxResults {
searchDepth = h.maxResults
}
Comment thread
yurishkuro marked this conversation as resolved.

// Convert attributes map to pcommon.Map
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ func TestSearchTracesHandler_Handle_FullWorkflow(t *testing.T) {
},
}

handler := &searchTracesHandler{queryService: mock}
handler := &searchTracesHandler{queryService: mock, maxResults: 100}

input := types.SearchTracesInput{
StartTimeMin: "-1h",
Expand All @@ -81,7 +81,7 @@ func TestSearchTracesHandler_Handle_WithStartTimeMax(t *testing.T) {

mock := newMockFindTraces(testTrace)

handler := &searchTracesHandler{queryService: mock}
handler := &searchTracesHandler{queryService: mock, maxResults: 100}

input := types.SearchTracesInput{
StartTimeMin: "-2h",
Expand All @@ -108,7 +108,7 @@ func TestSearchTracesHandler_Handle_WithDurations(t *testing.T) {
},
}

handler := &searchTracesHandler{queryService: mock}
handler := &searchTracesHandler{queryService: mock, maxResults: 100}

input := types.SearchTracesInput{
StartTimeMin: "-1h",
Expand Down Expand Up @@ -138,7 +138,7 @@ func TestSearchTracesHandler_Handle_WithAttributes(t *testing.T) {
},
}

handler := &searchTracesHandler{queryService: mock}
handler := &searchTracesHandler{queryService: mock, maxResults: 100}

input := types.SearchTracesInput{
StartTimeMin: "-1h",
Expand Down Expand Up @@ -171,7 +171,7 @@ func TestSearchTracesHandler_Handle_WithErrorsFilter(t *testing.T) {
},
}

handler := &searchTracesHandler{queryService: mock}
handler := &searchTracesHandler{queryService: mock, maxResults: 100}

input := types.SearchTracesInput{
StartTimeMin: "-1h",
Expand Down Expand Up @@ -199,7 +199,7 @@ func TestSearchTracesHandler_Handle_SearchDepthDefault(t *testing.T) {
},
}

handler := &searchTracesHandler{queryService: mock}
handler := &searchTracesHandler{queryService: mock, maxResults: 100}

input := types.SearchTracesInput{
StartTimeMin: "-1h",
Expand All @@ -223,7 +223,7 @@ func TestSearchTracesHandler_Handle_SearchDepthMax(t *testing.T) {
},
}

handler := &searchTracesHandler{queryService: mock}
handler := &searchTracesHandler{queryService: mock, maxResults: 100}

input := types.SearchTracesInput{
StartTimeMin: "-1h",
Expand All @@ -249,7 +249,7 @@ func TestSearchTracesHandler_Handle_QueryError(t *testing.T) {
},
}

handler := &searchTracesHandler{queryService: mock}
handler := &searchTracesHandler{queryService: mock, maxResults: 100}

input := types.SearchTracesInput{
StartTimeMin: "-1h",
Expand Down Expand Up @@ -286,7 +286,7 @@ func TestSearchTracesHandler_Handle_PartialResults(t *testing.T) {
},
}

handler := &searchTracesHandler{queryService: mock}
handler := &searchTracesHandler{queryService: mock, maxResults: 100}

input := types.SearchTracesInput{
StartTimeMin: "-1h",
Expand All @@ -305,7 +305,7 @@ func TestSearchTracesHandler_Handle_PartialResults(t *testing.T) {
}

func TestSearchTracesHandler_Handle_MissingServiceName(t *testing.T) {
handler := NewSearchTracesHandler(nil)
handler := NewSearchTracesHandler(nil, 100)

input := types.SearchTracesInput{
StartTimeMin: "-1h",
Expand All @@ -319,7 +319,7 @@ func TestSearchTracesHandler_Handle_MissingServiceName(t *testing.T) {
}

func TestSearchTracesHandler_Handle_InvalidTimeFormat(t *testing.T) {
handler := NewSearchTracesHandler(nil)
handler := NewSearchTracesHandler(nil, 100)

input := types.SearchTracesInput{
StartTimeMin: "invalid-time",
Expand All @@ -333,7 +333,7 @@ func TestSearchTracesHandler_Handle_InvalidTimeFormat(t *testing.T) {
}

func TestSearchTracesHandler_Handle_InvalidStartTimeMax(t *testing.T) {
handler := NewSearchTracesHandler(nil)
handler := NewSearchTracesHandler(nil, 100)

input := types.SearchTracesInput{
StartTimeMin: "-1h",
Expand All @@ -348,7 +348,7 @@ func TestSearchTracesHandler_Handle_InvalidStartTimeMax(t *testing.T) {
}

func TestSearchTracesHandler_Handle_InvalidDurationMin(t *testing.T) {
handler := NewSearchTracesHandler(nil)
handler := NewSearchTracesHandler(nil, 100)

input := types.SearchTracesInput{
StartTimeMin: "-1h",
Expand All @@ -363,7 +363,7 @@ func TestSearchTracesHandler_Handle_InvalidDurationMin(t *testing.T) {
}

func TestSearchTracesHandler_Handle_InvalidDurationMax(t *testing.T) {
handler := NewSearchTracesHandler(nil)
handler := NewSearchTracesHandler(nil, 100)

input := types.SearchTracesInput{
StartTimeMin: "-1h",
Expand All @@ -378,7 +378,7 @@ func TestSearchTracesHandler_Handle_InvalidDurationMax(t *testing.T) {
}

func TestSearchTracesHandler_Handle_DurationMaxLessThanMin(t *testing.T) {
handler := NewSearchTracesHandler(nil)
handler := NewSearchTracesHandler(nil, 100)

input := types.SearchTracesInput{
StartTimeMin: "-1h",
Expand Down Expand Up @@ -464,7 +464,7 @@ func TestSearchTracesHandler_Handle_DefaultStartTime(t *testing.T) {
},
}

handler := &searchTracesHandler{queryService: mock}
handler := &searchTracesHandler{queryService: mock, maxResults: 100}

// Omit StartTimeMin to trigger default logic
input := types.SearchTracesInput{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,8 +35,8 @@ type SearchTracesInput struct {
// SearchDepth defines the maximum search depth. Depending on the backend storage implementation,
// this may behave like an SQL LIMIT clause. However, some implementations might not support
// precise limits, and a larger value generally results in more traces being returned.
// Default: 10, max: 100.
SearchDepth int `json:"search_depth,omitempty" jsonschema:"Maximum search depth (default: 10 max: 100)"`
// Default: 10, maximum is controlled by server configuration (MaxSearchResults).
SearchDepth int `json:"search_depth,omitempty" jsonschema:"Maximum search depth (default: 10, max controlled by server config)"`
}

// SearchTracesOutput defines the output of the search_traces MCP tool.
Expand Down
Loading
Loading