feat(k8s): add a limit parameter to k8s_get_resources

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",

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) {