From 51cc7aa33c97d740d5f284f72b3bd64a401277eb Mon Sep 17 00:00:00 2001 From: Lucas Anjos Date: Wed, 16 Sep 2026 21:32:06 +0100 Subject: [PATCH] feat(k8s): add a limit parameter to k8s_get_resources k8s_get_resources had no way to bound its result size, so a call with all_namespaces=true returned every row on the cluster into the model's context and could exhaust the context window. The limit is applied per output format, because slicing lines out of a structured document would return an invalid one: tabular output is cut by row while preserving its header, json is cut at its items array, and yaml is left untouched since it has no equally safe cut point. A truncated result carries a notice naming the totals, so a shortened listing is not mistaken for a complete one. limit=0 returns everything. Closes #82 Signed-off-by: Lucas Anjos --- pkg/k8s/k8s.go | 93 ++++++++++++++++++++++++- pkg/k8s/k8s_test.go | 164 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 256 insertions(+), 1 deletion(-) diff --git a/pkg/k8s/k8s.go b/pkg/k8s/k8s.go index 6def2f2..6d324ae 100644 --- a/pkg/k8s/k8s.go +++ b/pkg/k8s/k8s.go @@ -3,6 +3,7 @@ package k8s import ( "context" _ "embed" + "encoding/json" "fmt" "maps" "math/rand" @@ -84,7 +85,96 @@ func (k *K8sTool) handleKubectlGetEnhanced(ctx context.Context, request mcp.Call args = append(args, "-o", "json") } - return k.runKubectlCommand(ctx, request.Header, args...) + result, err := k.runKubectlCommand(ctx, request.Header, args...) + if err != nil || result == nil || result.IsError { + return result, err + } + + return truncateGetResult(result, output, mcp.ParseInt(request, "limit", defaultGetResourcesLimit)), nil +} + +// defaultGetResourcesLimit bounds how many resources k8s_get_resources returns. +// A listing is fed straight into a model's context, so an unbounded one on a large +// cluster can exhaust the context window and end the session rather than returning +// a large answer. Callers that need everything can pass limit=0. +const defaultGetResourcesLimit = 200 + +// truncationNotice tells the caller the listing was shortened. It matters as much +// as the cap itself: a silently shortened listing reads as complete, so a model +// would draw conclusions from partial data. +const truncationNotice = "... truncated: showing %d of %d resources. Narrow with namespace or resource_name, or pass limit=0 for all." + +// truncateGetResult caps how many resources a kubectl get result carries. The +// output format decides how, because slicing lines out of a structured document +// would yield an invalid one: json is cut at its items array, and yaml is left +// alone since it has no equally safe cut point. +func truncateGetResult(result *mcp.CallToolResult, output string, limit int) *mcp.CallToolResult { + if limit <= 0 || result == nil || len(result.Content) == 0 { + return result + } + textContent, ok := result.Content[0].(mcp.TextContent) + if !ok || textContent.Text == "" { + return result + } + + switch { + case output == "json": + return truncateJSONItems(result, textContent.Text, limit) + case output == "" || output == "wide" || output == "name" || + strings.HasPrefix(output, "custom-columns"): + return truncateRows(result, textContent.Text, limit) + default: + // yaml, jsonpath and go-template have no row structure to cut on. + return result + } +} + +// truncateRows caps a line-oriented listing, preserving any header line. +func truncateRows(result *mcp.CallToolResult, text string, limit int) *mcp.CallToolResult { + lines := strings.Split(strings.TrimSuffix(text, "\n"), "\n") + + // A table carries a header line that is not itself a resource. + header := 0 + if strings.HasPrefix(lines[0], "NAME") { + header = 1 + } + + total := len(lines) - header + if total <= limit { + return result + } + + return mcp.NewToolResultText(fmt.Sprintf("%s\n"+truncationNotice, + strings.Join(lines[:header+limit], "\n"), limit, total)) +} + +// truncateJSONItems caps a kubectl List document at its items array. Anything +// that is not such a document (a single object, or unparseable output) is +// returned untouched rather than guessed at. +func truncateJSONItems(result *mcp.CallToolResult, text string, limit int) *mcp.CallToolResult { + var list struct { + APIVersion string `json:"apiVersion"` + Kind string `json:"kind"` + Metadata map[string]any `json:"metadata"` + Items []json.RawMessage `json:"items"` + } + if err := json.Unmarshal([]byte(text), &list); err != nil || list.Items == nil { + return result + } + + total := len(list.Items) + if total <= limit { + return result + } + + list.Items = list.Items[:limit] + truncated, err := json.MarshalIndent(list, "", " ") + if err != nil { + return result + } + + return mcp.NewToolResultText(fmt.Sprintf("%s\n"+truncationNotice, + truncated, limit, total)) } // Get pod logs @@ -653,6 +743,7 @@ func RegisterTools(s *server.MCPServer, llm llms.Model, kubeconfig string, readO mcp.WithString("namespace", mcp.Description("Namespace to query (optional)")), mcp.WithString("all_namespaces", mcp.Description("Query all namespaces (true/false)")), mcp.WithString("output", mcp.Description("Output format (json, yaml, wide)"), mcp.DefaultString("wide")), + mcp.WithNumber("limit", mcp.Description("Maximum number of resources to return; 0 returns all"), mcp.DefaultNumber(defaultGetResourcesLimit)), ), telemetry.AdaptToolHandler(telemetry.WithTracing("k8s_get_resources", k8sTool.handleKubectlGetEnhanced))) s.AddTool(mcp.NewTool("k8s_get_pod_logs", diff --git a/pkg/k8s/k8s_test.go b/pkg/k8s/k8s_test.go index 6c008a2..68f1652 100644 --- a/pkg/k8s/k8s_test.go +++ b/pkg/k8s/k8s_test.go @@ -2,7 +2,10 @@ package k8s import ( "context" + "encoding/json" + "fmt" "net/http" + "strings" "testing" "github.com/kagent-dev/tools/internal/cmd" @@ -561,6 +564,167 @@ func TestHandleKubectlGetEnhanced(t *testing.T) { assert.NotNil(t, result) assert.False(t, result.IsError) }) + + t.Run("truncates a listing longer than the limit", func(t *testing.T) { + mock := cmd.NewMockShellExecutor() + lines := []string{"NAME READY STATUS RESTARTS AGE"} + for i := 0; i < 5; i++ { + lines = append(lines, fmt.Sprintf("pod-%d 1/1 Running 0 1d", i)) + } + mock.AddCommandString("kubectl", []string{"get", "pods", "-o", "wide"}, strings.Join(lines, "\n"), nil) + ctx := cmd.WithShellExecutor(ctx, mock) + + k8sTool := newTestK8sTool() + req := mcp.CallToolRequest{} + req.Params.Arguments = map[string]interface{}{"resource_type": "pods", "limit": 2} + result, err := k8sTool.handleKubectlGetEnhanced(ctx, req) + assert.NoError(t, err) + assert.False(t, result.IsError) + + resultText := getResultText(result) + assert.Contains(t, resultText, "NAME") + assert.Contains(t, resultText, "pod-0") + assert.Contains(t, resultText, "pod-1") + // A truncated listing must not read as complete, or the model draws + // conclusions from partial data. + assert.NotContains(t, resultText, "pod-2") + assert.Contains(t, resultText, "showing 2 of 5 resources") + }) + + t.Run("leaves a listing shorter than the limit untouched", func(t *testing.T) { + mock := cmd.NewMockShellExecutor() + expectedOutput := "NAME READY STATUS RESTARTS AGE\npod-0 1/1 Running 0 1d" + mock.AddCommandString("kubectl", []string{"get", "pods", "-o", "wide"}, expectedOutput, nil) + ctx := cmd.WithShellExecutor(ctx, mock) + + k8sTool := newTestK8sTool() + req := mcp.CallToolRequest{} + req.Params.Arguments = map[string]interface{}{"resource_type": "pods", "limit": 10} + result, err := k8sTool.handleKubectlGetEnhanced(ctx, req) + assert.NoError(t, err) + assert.False(t, result.IsError) + + resultText := getResultText(result) + assert.Equal(t, expectedOutput, resultText) + assert.NotContains(t, resultText, "truncated") + }) + + t.Run("limit=0 returns every row", func(t *testing.T) { + mock := cmd.NewMockShellExecutor() + lines := []string{"NAME READY STATUS RESTARTS AGE"} + for i := 0; i < 3; i++ { + lines = append(lines, fmt.Sprintf("pod-%d 1/1 Running 0 1d", i)) + } + expectedOutput := strings.Join(lines, "\n") + mock.AddCommandString("kubectl", []string{"get", "pods", "-o", "wide"}, expectedOutput, nil) + ctx := cmd.WithShellExecutor(ctx, mock) + + k8sTool := newTestK8sTool() + req := mcp.CallToolRequest{} + req.Params.Arguments = map[string]interface{}{"resource_type": "pods", "limit": 0} + result, err := k8sTool.handleKubectlGetEnhanced(ctx, req) + assert.NoError(t, err) + assert.False(t, result.IsError) + + resultText := getResultText(result) + assert.Equal(t, expectedOutput, resultText) + assert.NotContains(t, resultText, "truncated") + }) + + t.Run("applies the default limit when none is given", func(t *testing.T) { + mock := cmd.NewMockShellExecutor() + lines := []string{"NAME READY STATUS RESTARTS AGE"} + for i := 0; i < defaultGetResourcesLimit+5; i++ { + lines = append(lines, fmt.Sprintf("pod-%d 1/1 Running 0 1d", i)) + } + mock.AddCommandString("kubectl", []string{"get", "pods", "-o", "wide"}, strings.Join(lines, "\n"), nil) + ctx := cmd.WithShellExecutor(ctx, mock) + + k8sTool := newTestK8sTool() + req := mcp.CallToolRequest{} + req.Params.Arguments = map[string]interface{}{"resource_type": "pods"} + result, err := k8sTool.handleKubectlGetEnhanced(ctx, req) + assert.NoError(t, err) + assert.False(t, result.IsError) + + resultText := getResultText(result) + assert.Contains(t, resultText, fmt.Sprintf("showing %d of %d resources", defaultGetResourcesLimit, defaultGetResourcesLimit+5)) + }) + + t.Run("caps json output at the items array, keeping it parseable", func(t *testing.T) { + mock := cmd.NewMockShellExecutor() + items := make([]string, 0, 5) + for i := 0; i < 5; i++ { + items = append(items, fmt.Sprintf(`{"metadata":{"name":"pod-%d"}}`, i)) + } + listJSON := fmt.Sprintf(`{"apiVersion":"v1","kind":"List","items":[%s]}`, strings.Join(items, ",")) + mock.AddCommandString("kubectl", []string{"get", "pods", "-o", "json"}, listJSON, nil) + ctx := cmd.WithShellExecutor(ctx, mock) + + k8sTool := newTestK8sTool() + req := mcp.CallToolRequest{} + req.Params.Arguments = map[string]interface{}{"resource_type": "pods", "output": "json", "limit": 2} + result, err := k8sTool.handleKubectlGetEnhanced(ctx, req) + assert.NoError(t, err) + assert.False(t, result.IsError) + + resultText := getResultText(result) + assert.Contains(t, resultText, "showing 2 of 5 resources") + + // Cutting a structured document by line would leave it unparseable. + jsonPart := resultText[:strings.LastIndex(resultText, "\n... truncated")] + var decoded struct { + Kind string `json:"kind"` + Items []struct { + Metadata struct { + Name string `json:"name"` + } `json:"metadata"` + } `json:"items"` + } + require.NoError(t, json.Unmarshal([]byte(jsonPart), &decoded)) + assert.Equal(t, "List", decoded.Kind) + require.Len(t, decoded.Items, 2) + assert.Equal(t, "pod-0", decoded.Items[0].Metadata.Name) + assert.Equal(t, "pod-1", decoded.Items[1].Metadata.Name) + }) + + t.Run("leaves yaml untouched since it has no safe cut point", func(t *testing.T) { + mock := cmd.NewMockShellExecutor() + lines := make([]string, 0, defaultGetResourcesLimit+10) + for i := 0; i < defaultGetResourcesLimit+10; i++ { + lines = append(lines, fmt.Sprintf(" key%d: value%d", i, i)) + } + expectedOutput := "apiVersion: v1\n" + strings.Join(lines, "\n") + mock.AddCommandString("kubectl", []string{"get", "pod", "my-pod", "-o", "yaml"}, expectedOutput, nil) + ctx := cmd.WithShellExecutor(ctx, mock) + + k8sTool := newTestK8sTool() + req := mcp.CallToolRequest{} + req.Params.Arguments = map[string]interface{}{ + "resource_type": "pod", "resource_name": "my-pod", "output": "yaml", + } + result, err := k8sTool.handleKubectlGetEnhanced(ctx, req) + assert.NoError(t, err) + assert.False(t, result.IsError) + assert.Equal(t, expectedOutput, getResultText(result)) + }) + + t.Run("leaves a single json object untouched", func(t *testing.T) { + mock := cmd.NewMockShellExecutor() + expectedOutput := `{"apiVersion":"v1","kind":"Pod","metadata":{"name":"my-pod"}}` + mock.AddCommandString("kubectl", []string{"get", "pod", "my-pod", "-o", "json"}, expectedOutput, nil) + ctx := cmd.WithShellExecutor(ctx, mock) + + k8sTool := newTestK8sTool() + req := mcp.CallToolRequest{} + req.Params.Arguments = map[string]interface{}{ + "resource_type": "pod", "resource_name": "my-pod", "output": "json", "limit": 1, + } + result, err := k8sTool.handleKubectlGetEnhanced(ctx, req) + assert.NoError(t, err) + assert.False(t, result.IsError) + assert.Equal(t, expectedOutput, getResultText(result)) + }) } func TestHandleKubectlLogsEnhanced(t *testing.T) {