Skip to content

Commit 383182c

Browse files
MikeHsu0618MikeHsu0618
authored
feat(pyroscope): add series query and unified query tool (#672)
1 parent c7a0c61 commit 383182c

3 files changed

Lines changed: 372 additions & 81 deletions

File tree

‎tools/pyroscope.go‎

Lines changed: 187 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,10 @@ package tools
22

33
import (
44
"context"
5+
"encoding/json"
56
"fmt"
67
"io"
8+
"math"
79
"net/http"
810
"net/url"
911
"regexp"
@@ -23,7 +25,7 @@ func AddPyroscopeTools(mcp *server.MCPServer) {
2325
ListPyroscopeLabelNames.Register(mcp)
2426
ListPyroscopeLabelValues.Register(mcp)
2527
ListPyroscopeProfileTypes.Register(mcp)
26-
FetchPyroscopeProfile.Register(mcp)
28+
QueryPyroscope.Register(mcp)
2729
}
2830

2931
const listPyroscopeLabelNamesToolPrompt = `
@@ -210,79 +212,6 @@ func listPyroscopeProfileTypes(ctx context.Context, args ListPyroscopeProfileTyp
210212
return profileTypes, nil
211213
}
212214

213-
const fetchPyroscopeProfileToolPrompt = `
214-
Fetches a profile from a Pyroscope data source for a given time range. By default, the time range is tha past 1 hour.
215-
The profile type is required, available profile types can be fetched via the list_pyroscope_profile_types tool. Not all
216-
profile types are available for every service. Expect some queries to return empty result sets, this indicates the
217-
profile type does not exist for that query. In such a case, consider trying a related profile type or giving up.
218-
Matchers are not required, but highly recommended, they are generally used to select an application by the service_name
219-
label (e.g. {service_name="foo"}). Use the list_pyroscope_label_names tool to fetch available label names, and the
220-
list_pyroscope_label_values tool to fetch available label values. The returned profile is in DOT format.
221-
`
222-
223-
var FetchPyroscopeProfile = mcpgrafana.MustTool(
224-
"fetch_pyroscope_profile",
225-
fetchPyroscopeProfileToolPrompt,
226-
fetchPyroscopeProfile,
227-
mcp.WithTitleAnnotation("Fetch Pyroscope profile"),
228-
mcp.WithIdempotentHintAnnotation(true),
229-
mcp.WithReadOnlyHintAnnotation(true),
230-
)
231-
232-
type FetchPyroscopeProfileParams struct {
233-
DataSourceUID string `json:"data_source_uid" jsonschema:"required,description=The UID of the datasource to query"`
234-
ProfileType string `json:"profile_type" jsonschema:"required,description=Type profile type\\, use the list_pyroscope_profile_types tool to fetch available profile types"`
235-
Matchers string `json:"matchers,omitempty" jsonschema:"description=Optionally\\, Prometheus style matchers used to filter the result set (defaults to: {})"`
236-
MaxNodeDepth int `json:"max_node_depth,omitempty" jsonschema:"description=Optionally\\, the maximum depth of nodes in the resulting profile. Less depth results in smaller profiles that execute faster\\, more depth result in larger profiles that have more detail. A value of -1 indicates to use an unbounded node depth (default: 100). Reducing max node depth from the default will negatively impact the accuracy of the profile"`
237-
StartRFC3339 string `json:"start_rfc_3339,omitempty" jsonschema:"description=Optionally\\, the start time of the query in RFC3339 format (defaults to 1 hour ago)"`
238-
EndRFC3339 string `json:"end_rfc_3339,omitempty" jsonschema:"description=Optionally\\, the end time of the query in RFC3339 format (defaults to now)"`
239-
}
240-
241-
func fetchPyroscopeProfile(ctx context.Context, args FetchPyroscopeProfileParams) (string, error) {
242-
args.Matchers = stringOrDefault(args.Matchers, "{}")
243-
matchersRegex := regexp.MustCompile(`^\{.*\}$`)
244-
if !matchersRegex.MatchString(args.Matchers) {
245-
args.Matchers = fmt.Sprintf("{%s}", args.Matchers)
246-
}
247-
248-
args.MaxNodeDepth = intOrDefault(args.MaxNodeDepth, 100)
249-
250-
start, err := rfc3339OrDefault(args.StartRFC3339, time.Time{})
251-
if err != nil {
252-
return "", fmt.Errorf("failed to parse start timestamp %q: %w", args.StartRFC3339, err)
253-
}
254-
255-
end, err := rfc3339OrDefault(args.EndRFC3339, time.Time{})
256-
if err != nil {
257-
return "", fmt.Errorf("failed to parse end timestamp %q: %w", args.EndRFC3339, err)
258-
}
259-
260-
start, end, err = validateTimeRange(start, end)
261-
if err != nil {
262-
return "", err
263-
}
264-
265-
client, err := newPyroscopeClient(ctx, args.DataSourceUID)
266-
if err != nil {
267-
return "", fmt.Errorf("failed to create Pyroscope client: %w", err)
268-
}
269-
270-
req := &renderRequest{
271-
ProfileType: args.ProfileType,
272-
Matcher: args.Matchers,
273-
Start: start,
274-
End: end,
275-
Format: "dot",
276-
MaxNodes: args.MaxNodeDepth,
277-
}
278-
res, err := client.Render(ctx, req)
279-
if err != nil {
280-
return "", fmt.Errorf("failed to call Pyroscope API: %w", err)
281-
}
282-
283-
res = cleanupDotProfile(res)
284-
return res, nil
285-
}
286215

287216
func newPyroscopeClient(ctx context.Context, uid string) (*pyroscopeClient, error) {
288217
cfg := mcpgrafana.GrafanaConfigFromContext(ctx)
@@ -458,3 +387,187 @@ func cleanupDotProfile(profile string) string {
458387
return ""
459388
})
460389
}
390+
391+
392+
var matchersRegex = regexp.MustCompile(`^\{.*\}$`)
393+
394+
// rawSeries is the JSON structure returned for a single time-series.
395+
type rawSeries struct {
396+
Labels map[string]string `json:"labels"`
397+
Points [][2]float64 `json:"points"` // [[timestamp_ms, value], ...]
398+
}
399+
400+
// seriesResponse is the structured metrics response embedded in the query_pyroscope result.
401+
type seriesResponse struct {
402+
Series []rawSeries `json:"series"`
403+
TimeRange map[string]string `json:"time_range"`
404+
StepSecs float64 `json:"step_seconds"`
405+
}
406+
407+
func buildSeriesResponse(series []*typesv1.Series, start, end time.Time, step float64) *seriesResponse {
408+
raw := make([]rawSeries, 0, len(series))
409+
for _, s := range series {
410+
labels := make(map[string]string, len(s.Labels))
411+
for _, lp := range s.Labels {
412+
labels[lp.Name] = lp.Value
413+
}
414+
415+
points := make([][2]float64, 0, len(s.Points))
416+
for _, p := range s.Points {
417+
points = append(points, [2]float64{float64(p.Timestamp), p.Value})
418+
}
419+
420+
if len(points) == 0 {
421+
continue
422+
}
423+
424+
raw = append(raw, rawSeries{
425+
Labels: labels,
426+
Points: points,
427+
})
428+
}
429+
430+
return &seriesResponse{
431+
Series: raw,
432+
TimeRange: map[string]string{"from": start.Format(time.RFC3339), "to": end.Format(time.RFC3339)},
433+
StepSecs: step,
434+
}
435+
}
436+
437+
// ---------------------------------------------------------------------------
438+
// query_pyroscope — unified tool: profile + metrics + both
439+
// ---------------------------------------------------------------------------
440+
441+
const queryPyroscopeToolPrompt = `
442+
Unified Pyroscope query tool for fetching profiles or metrics from Pyroscope. Profile data shows WHICH functions consume resources; metrics data
443+
shows WHEN consumption spiked. Use query_type="both" for complete analysis in one call.
444+
445+
query_type options (extends Grafana's PyroscopeQueryType):
446+
- "profile": returns DOT-format call graph
447+
- "metrics": returns time-series data points
448+
- "both" (default): returns both profile and metrics in one response
449+
`
450+
451+
var QueryPyroscope = mcpgrafana.MustTool(
452+
"query_pyroscope",
453+
queryPyroscopeToolPrompt,
454+
queryPyroscope,
455+
mcp.WithTitleAnnotation("Query Pyroscope"),
456+
mcp.WithIdempotentHintAnnotation(true),
457+
mcp.WithReadOnlyHintAnnotation(true),
458+
)
459+
460+
type QueryPyroscopeParams struct {
461+
DataSourceUID string `json:"data_source_uid" jsonschema:"required,description=The UID of the datasource to query"`
462+
ProfileType string `json:"profile_type" jsonschema:"required,description=The profile type\\, use list_pyroscope_profile_types to discover available types"`
463+
QueryType string `json:"query_type,omitempty" jsonschema:"description=Query type: \"profile\" (flamegraph)\\, \"metrics\" (time-series)\\, or \"both\" (default). Use \"both\" for complete analysis"`
464+
Matchers string `json:"matchers,omitempty" jsonschema:"description=Prometheus style matchers (defaults to: {})"`
465+
GroupBy []string `json:"group_by,omitempty" jsonschema:"description=Labels to group metrics series by"`
466+
Step float64 `json:"step,omitempty" jsonschema:"description=Seconds between metrics data points (default: auto)"`
467+
MaxNodeDepth int `json:"max_node_depth,omitempty" jsonschema:"description=Max depth for profile call graph (default: 100)"`
468+
StartRFC3339 string `json:"start_rfc_3339,omitempty" jsonschema:"description=Start time in RFC3339 (defaults to 1 hour ago)"`
469+
EndRFC3339 string `json:"end_rfc_3339,omitempty" jsonschema:"description=End time in RFC3339 (defaults to now)"`
470+
}
471+
472+
func queryPyroscope(ctx context.Context, args QueryPyroscopeParams) (string, error) {
473+
queryType := strings.ToLower(strings.TrimSpace(args.QueryType))
474+
if queryType == "" {
475+
queryType = "both"
476+
}
477+
if queryType != "profile" && queryType != "metrics" && queryType != "both" {
478+
return "", fmt.Errorf("invalid query_type %q: must be \"profile\", \"metrics\", or \"both\"", args.QueryType)
479+
}
480+
481+
// Common setup
482+
matchers := stringOrDefault(args.Matchers, "{}")
483+
if !matchersRegex.MatchString(matchers) {
484+
matchers = fmt.Sprintf("{%s}", matchers)
485+
}
486+
487+
start, err := rfc3339OrDefault(args.StartRFC3339, time.Time{})
488+
if err != nil {
489+
return "", fmt.Errorf("failed to parse start timestamp %q: %w", args.StartRFC3339, err)
490+
}
491+
492+
end, err := rfc3339OrDefault(args.EndRFC3339, time.Time{})
493+
if err != nil {
494+
return "", fmt.Errorf("failed to parse end timestamp %q: %w", args.EndRFC3339, err)
495+
}
496+
497+
start, end, err = validateTimeRange(start, end)
498+
if err != nil {
499+
return "", err
500+
}
501+
502+
client, err := newPyroscopeClient(ctx, args.DataSourceUID)
503+
if err != nil {
504+
return "", fmt.Errorf("failed to create Pyroscope client: %w", err)
505+
}
506+
507+
wantProfile := queryType == "profile" || queryType == "both"
508+
wantMetrics := queryType == "metrics" || queryType == "both"
509+
510+
result := make(map[string]any)
511+
result["query_type"] = queryType
512+
513+
if wantProfile {
514+
maxNodes := intOrDefault(args.MaxNodeDepth, 100)
515+
res, profileErr := client.Render(ctx, &renderRequest{
516+
ProfileType: args.ProfileType,
517+
Matcher: matchers,
518+
Start: start,
519+
End: end,
520+
Format: "dot",
521+
MaxNodes: maxNodes,
522+
})
523+
if profileErr != nil {
524+
// Single-type query: propagate error so MCP framework sets IsError=true.
525+
// "both" mode: embed error for partial results.
526+
if queryType == "profile" {
527+
return "", fmt.Errorf("failed to fetch profile: %w", profileErr)
528+
}
529+
result["profile"] = map[string]string{"error": profileErr.Error()}
530+
} else {
531+
result["profile"] = cleanupDotProfile(res)
532+
}
533+
}
534+
535+
if wantMetrics {
536+
step := args.Step
537+
if step <= 0 {
538+
step = math.Max(end.Sub(start).Seconds()/50.0, 15.0)
539+
}
540+
541+
seriesRes, metricsErr := client.SelectSeries(ctx, connect.NewRequest(&querierv1.SelectSeriesRequest{
542+
ProfileTypeID: args.ProfileType,
543+
LabelSelector: matchers,
544+
Start: start.UnixMilli(),
545+
End: end.UnixMilli(),
546+
GroupBy: args.GroupBy,
547+
Step: step,
548+
}))
549+
if metricsErr != nil {
550+
if queryType == "metrics" {
551+
return "", fmt.Errorf("failed to fetch metrics: %w", metricsErr)
552+
}
553+
result["metrics"] = map[string]string{"error": metricsErr.Error()}
554+
} else {
555+
result["metrics"] = buildSeriesResponse(seriesRes.Msg.Series, start, end, step)
556+
}
557+
}
558+
559+
// If both queries were attempted and both failed, propagate error.
560+
_, profileFailed := result["profile"].(map[string]string)
561+
_, metricsFailed := result["metrics"].(map[string]string)
562+
if queryType == "both" && profileFailed && metricsFailed {
563+
return "", fmt.Errorf("both queries failed — profile: %s; metrics: %s",
564+
result["profile"].(map[string]string)["error"],
565+
result["metrics"].(map[string]string)["error"])
566+
}
567+
568+
out, err := json.Marshal(result)
569+
if err != nil {
570+
return "", fmt.Errorf("failed to marshal response: %w", err)
571+
}
572+
return string(out), nil
573+
}

‎tools/pyroscope_test.go‎

Lines changed: 40 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,10 @@
33
package tools
44

55
import (
6+
"encoding/json"
67
"testing"
78

9+
"github.com/stretchr/testify/assert"
810
"github.com/stretchr/testify/require"
911
)
1012

@@ -65,24 +67,55 @@ func TestPyroscopeTools(t *testing.T) {
6567
})
6668
})
6769

68-
t.Run("fetch Pyroscope profile", func(t *testing.T) {
70+
t.Run("query Pyroscope both", func(t *testing.T) {
6971
ctx := newTestContext()
70-
profile, err := fetchPyroscopeProfile(ctx, FetchPyroscopeProfileParams{
72+
result, err := queryPyroscope(ctx, QueryPyroscopeParams{
7173
DataSourceUID: "pyroscope",
7274
ProfileType: "process_cpu:cpu:nanoseconds:cpu:nanoseconds",
7375
Matchers: `{service_name="pyroscope"}`,
76+
QueryType: "both",
7477
})
7578
require.NoError(t, err)
76-
require.NotEmpty(t, profile)
79+
require.NotEmpty(t, result)
80+
81+
var parsed map[string]any
82+
require.NoError(t, json.Unmarshal([]byte(result), &parsed))
83+
assert.Equal(t, "both", parsed["query_type"])
84+
assert.NotNil(t, parsed["profile"], "profile should be present")
85+
assert.NotNil(t, parsed["metrics"], "metrics should be present")
7786
})
7887

79-
t.Run("fetch empty Pyroscope profile", func(t *testing.T) {
88+
t.Run("query Pyroscope profile only", func(t *testing.T) {
8089
ctx := newTestContext()
81-
_, err := fetchPyroscopeProfile(ctx, FetchPyroscopeProfileParams{
90+
result, err := queryPyroscope(ctx, QueryPyroscopeParams{
8291
DataSourceUID: "pyroscope",
8392
ProfileType: "process_cpu:cpu:nanoseconds:cpu:nanoseconds",
84-
Matchers: `{service_name="pyroscope", label_does_not_exit="missing"}`,
93+
Matchers: `{service_name="pyroscope"}`,
94+
QueryType: "profile",
8595
})
86-
require.EqualError(t, err, "failed to call Pyroscope API: pyroscope API returned an empty response")
96+
require.NoError(t, err)
97+
98+
var parsed map[string]any
99+
require.NoError(t, json.Unmarshal([]byte(result), &parsed))
100+
assert.Equal(t, "profile", parsed["query_type"])
101+
assert.NotNil(t, parsed["profile"])
102+
assert.Nil(t, parsed["metrics"], "metrics should not be present for profile-only")
103+
})
104+
105+
t.Run("query Pyroscope metrics only", func(t *testing.T) {
106+
ctx := newTestContext()
107+
result, err := queryPyroscope(ctx, QueryPyroscopeParams{
108+
DataSourceUID: "pyroscope",
109+
ProfileType: "process_cpu:cpu:nanoseconds:cpu:nanoseconds",
110+
Matchers: `{service_name="pyroscope"}`,
111+
QueryType: "metrics",
112+
})
113+
require.NoError(t, err)
114+
115+
var parsed map[string]any
116+
require.NoError(t, json.Unmarshal([]byte(result), &parsed))
117+
assert.Equal(t, "metrics", parsed["query_type"])
118+
assert.Nil(t, parsed["profile"], "profile should not be present for metrics-only")
119+
assert.NotNil(t, parsed["metrics"])
87120
})
88121
}

0 commit comments

Comments
 (0)