From b2e4dfa3ae4041daacbfb99f96c8fd274a21c62e Mon Sep 17 00:00:00 2001 From: Benjamin Huo Date: Wed, 11 Nov 2026 23:05:08 +0800 Subject: [PATCH] fix: use object output schema for Kubescape vulnerability details Signed-off-by: Benjamin Huo --- cmd/tools_output_schema_test.go | 31 +++++++-- pkg/kubescape/kubescape.go | 34 +++++---- pkg/kubescape/kubescape_test.go | 118 +++++++++++++++++++++++++++++--- 3 files changed, 157 insertions(+), 26 deletions(-) diff --git a/cmd/tools_output_schema_test.go b/cmd/tools_output_schema_test.go index cadf9d7..d0e5817 100644 --- a/cmd/tools_output_schema_test.go +++ b/cmd/tools_output_schema_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "testing" + "time" sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/stretchr/testify/require" @@ -15,7 +16,8 @@ import ( // inference or resolution fails (for example a third-party k8s type with a // custom JSON marshaller), so a panic here means the server would not start. func TestEveryToolHasValidOutputSchema(t *testing.T) { - ctx := context.Background() + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() srv := sdkmcp.NewServer(&sdkmcp.Implementation{Name: "schema", Version: "test"}, nil) @@ -41,10 +43,29 @@ func TestEveryToolHasValidOutputSchema(t *testing.T) { require.NotNilf(t, tool.OutputSchema, "tool %q has no output schema; its handler still returns Out=any", tool.Name) - // The advertised schema must survive a JSON round-trip (it is sent on the - // wire as part of tools/list). - _, err := json.Marshal(tool.OutputSchema) - require.NoErrorf(t, err, "tool %q output schema is not JSON-serializable", tool.Name) + // Clients implementing the MCP 2025-06-18 Tool definition require + // object roots. A single array output rejects the whole tools/list + // response even when the SDK can infer and serialize that schema. + t.Run(tool.Name, func(t *testing.T) { + inputSchema, err := json.Marshal(tool.InputSchema) + require.NoError(t, err) + outputSchema, err := json.Marshal(tool.OutputSchema) + require.NoError(t, err) + for _, schema := range []struct { + name string + value []byte + }{ + {name: "inputSchema", value: inputSchema}, + {name: "outputSchema", value: outputSchema}, + } { + var root struct { + Type string `json:"type"` + } + require.NoErrorf(t, json.Unmarshal(schema.value, &root), + "%s must have a single object root: %s", schema.name, schema.value) + require.Equal(t, "object", root.Type, schema.name) + } + }) checked++ } require.NotEmpty(t, checked, "expected the server to advertise tools") diff --git a/pkg/kubescape/kubescape.go b/pkg/kubescape/kubescape.go index dcadfee..9e7207f 100644 --- a/pkg/kubescape/kubescape.go +++ b/pkg/kubescape/kubescape.go @@ -191,6 +191,12 @@ type listVulnerabilityManifestsOutput struct { TotalCount int `json:"total_count"` } +// getVulnerabilityDetailsOutput keeps the MCP output schema rooted at an +// object, as required by clients implementing the MCP 2025-06-18 Tool schema. +type getVulnerabilityDetailsOutput struct { + Matches []v1beta1.Match `json:"matches"` +} + type severitySummary struct { Critical int `json:"Critical"` High int `json:"High"` @@ -812,10 +818,10 @@ func (k *KubescapeTool) handleListVulnerabilitiesInManifest(ctx context.Context, } // handleGetVulnerabilityDetails gets detailed info about a specific CVE in a manifest -func (k *KubescapeTool) handleGetVulnerabilityDetails(ctx context.Context, request *sdkmcp.CallToolRequest, in getVulnerabilityDetailsInput) (*sdkmcp.CallToolResult, []v1beta1.Match, error) { +func (k *KubescapeTool) handleGetVulnerabilityDetails(ctx context.Context, request *sdkmcp.CallToolRequest, in getVulnerabilityDetailsInput) (*sdkmcp.CallToolResult, getVulnerabilityDetailsOutput, error) { if k.initError != nil { toolErr := errors.NewKubescapeError("get_vulnerability_details", k.initError) - return kubescapeErrResult(toolErr), nil, nil + return kubescapeErrResult(toolErr), getVulnerabilityDetailsOutput{Matches: []v1beta1.Match{}}, nil } namespace := in.Namespace @@ -826,10 +832,10 @@ func (k *KubescapeTool) handleGetVulnerabilityDetails(ctx context.Context, reque cveID := in.CveID if manifestName == "" { - return mcp.NewToolResultError("manifest_name parameter is required"), nil, nil + return mcp.NewToolResultError("manifest_name parameter is required"), getVulnerabilityDetailsOutput{Matches: []v1beta1.Match{}}, nil } if cveID == "" { - return mcp.NewToolResultError("cve_id parameter is required"), nil, nil + return mcp.NewToolResultError("cve_id parameter is required"), getVulnerabilityDetailsOutput{Matches: []v1beta1.Match{}}, nil } manifest, err := k.spdxClient.VulnerabilityManifests(namespace).Get(ctx, manifestName, metav1.GetOptions{}) @@ -837,7 +843,7 @@ func (k *KubescapeTool) handleGetVulnerabilityDetails(ctx context.Context, reque toolErr := errors.NewKubescapeError("get_vulnerability_manifest", err). WithContext("namespace", namespace). WithContext("manifest_name", manifestName) - return kubescapeErrResult(toolErr), nil, nil + return kubescapeErrResult(toolErr), getVulnerabilityDetailsOutput{Matches: []v1beta1.Match{}}, nil } // Find matching CVE entries @@ -849,14 +855,16 @@ func (k *KubescapeTool) handleGetVulnerabilityDetails(ctx context.Context, reque } if len(matches) == 0 { - return mcp.NewToolResultError(fmt.Sprintf("CVE %s not found in manifest %s", cveID, manifestName)), nil, nil + return mcp.NewToolResultError(fmt.Sprintf("CVE %s not found in manifest %s", cveID, manifestName)), getVulnerabilityDetailsOutput{Matches: []v1beta1.Match{}}, nil } - // Return the matches as the typed Out and leave Content unset. For an array - // (non-object) output the SDK serialises the value into a TextContent block - // itself; setting Content here as well would emit the same list twice for - // clients that read Content. - return nil, matches, nil + // Preserve the array-shaped text response for existing consumers, while + // exposing an object-shaped structured result and output schema. + content, err := json.Marshal(matches) + if err != nil { + return mcp.NewToolResultError(fmt.Sprintf("failed to marshal result: %v", err)), getVulnerabilityDetailsOutput{Matches: []v1beta1.Match{}}, nil + } + return mcp.NewToolResultText(string(content)), getVulnerabilityDetailsOutput{Matches: matches}, nil } // handleListConfigurationScans lists configuration security scan results @@ -1333,7 +1341,7 @@ type KubescapeToolInterface interface { HandleCheckHealth(ctx context.Context, in checkHealthInput) (*sdkmcp.CallToolResult, HealthCheckResult, error) HandleListVulnerabilityManifests(ctx context.Context, in listVulnerabilityManifestsInput) (*sdkmcp.CallToolResult, listVulnerabilityManifestsOutput, error) HandleListVulnerabilitiesInManifest(ctx context.Context, in listVulnerabilitiesInManifestInput) (*sdkmcp.CallToolResult, listVulnerabilitiesInManifestOutput, error) - HandleGetVulnerabilityDetails(ctx context.Context, in getVulnerabilityDetailsInput) (*sdkmcp.CallToolResult, []v1beta1.Match, error) + HandleGetVulnerabilityDetails(ctx context.Context, in getVulnerabilityDetailsInput) (*sdkmcp.CallToolResult, getVulnerabilityDetailsOutput, error) HandleListConfigurationScans(ctx context.Context, in listConfigurationScansInput) (*sdkmcp.CallToolResult, listConfigurationScansOutput, error) HandleGetConfigurationScan(ctx context.Context, in getConfigurationScanInput) (*sdkmcp.CallToolResult, mcp.TextOutput, error) HandleListApplicationProfiles(ctx context.Context, in listApplicationProfilesInput) (*sdkmcp.CallToolResult, listApplicationProfilesOutput, error) @@ -1358,7 +1366,7 @@ func (k *KubescapeTool) HandleListVulnerabilitiesInManifest(ctx context.Context, return k.handleListVulnerabilitiesInManifest(ctx, &sdkmcp.CallToolRequest{}, in) } -func (k *KubescapeTool) HandleGetVulnerabilityDetails(ctx context.Context, in getVulnerabilityDetailsInput) (*sdkmcp.CallToolResult, []v1beta1.Match, error) { +func (k *KubescapeTool) HandleGetVulnerabilityDetails(ctx context.Context, in getVulnerabilityDetailsInput) (*sdkmcp.CallToolResult, getVulnerabilityDetailsOutput, error) { return k.handleGetVulnerabilityDetails(ctx, &sdkmcp.CallToolRequest{}, in) } diff --git a/pkg/kubescape/kubescape_test.go b/pkg/kubescape/kubescape_test.go index fd9a672..601054f 100644 --- a/pkg/kubescape/kubescape_test.go +++ b/pkg/kubescape/kubescape_test.go @@ -5,6 +5,9 @@ import ( "encoding/json" "errors" "testing" + "time" + + toolmcp "github.com/kagent-dev/tools/internal/mcp" "github.com/kubescape/storage/pkg/apis/softwarecomposition/v1beta1" kubescapefake "github.com/kubescape/storage/pkg/generated/clientset/versioned/fake" @@ -559,19 +562,19 @@ func TestHandleGetVulnerabilityDetails_Success(t *testing.T) { tool := NewKubescapeToolWithClients(nil, nil, spdxClient.SpdxV1beta1()) - // The handler returns the matches as its typed Out and leaves the tool result - // for the SDK to synthesise. For an array Out the SDK serialises the value - // into a TextContent block itself, so setting Content in the handler as well - // would emit the list twice. - result, matches, err := tool.HandleGetVulnerabilityDetails(context.Background(), getVulnerabilityDetailsInput{ + result, output, err := tool.HandleGetVulnerabilityDetails(context.Background(), getVulnerabilityDetailsInput{ ManifestName: "test-manifest", CveID: "CVE-2021-1234", }) require.NoError(t, err) - assert.Nil(t, result, "handler should leave result to the SDK for an array Out") + require.NotNil(t, result) + require.False(t, result.IsError) - require.Len(t, matches, 1) - assert.Equal(t, "CVE-2021-1234", matches[0].Vulnerability.ID) + require.Len(t, output.Matches, 1) + assert.Equal(t, "CVE-2021-1234", output.Matches[0].Vulnerability.ID) + var textMatches []v1beta1.Match + require.NoError(t, json.Unmarshal([]byte(getResultText(result)), &textMatches)) + assert.Equal(t, output.Matches, textMatches) } func TestHandleGetVulnerabilityDetails_MissingManifestName(t *testing.T) { @@ -1180,3 +1183,102 @@ func TestHandleListVulnerabilityManifests_OmitsVulnerabilityCount(t *testing.T) // The useful metadata is still there. assert.Contains(t, getResultText(result), "nginx:1.14.0") } + +// TestGetVulnerabilityDetailsMCP exercises SDK schema inference and actual +// tools/call serialization, rather than only calling the handler directly. +func TestGetVulnerabilityDetailsMCP(t *testing.T) { + for _, testCase := range []struct { + name string + input getVulnerabilityDetailsInput + initError bool + wantError string + }{ + { + name: "success", + input: getVulnerabilityDetailsInput{ManifestName: "test-manifest", CveID: "CVE-2021-1234"}, + }, + { + name: "missing manifest name", + input: getVulnerabilityDetailsInput{CveID: "CVE-2021-1234"}, + wantError: "manifest_name parameter is required", + }, + { + name: "missing CVE id", + input: getVulnerabilityDetailsInput{ManifestName: "test-manifest"}, + wantError: "cve_id parameter is required", + }, + { + name: "CVE not found", + input: getVulnerabilityDetailsInput{ManifestName: "test-manifest", CveID: "CVE-2021-9999"}, + wantError: "CVE CVE-2021-9999 not found", + }, + { + name: "manifest not found", + input: getVulnerabilityDetailsInput{ManifestName: "missing-manifest", CveID: "CVE-2021-1234"}, + wantError: "not found", + }, + { + name: "initialization error", + input: getVulnerabilityDetailsInput{ManifestName: "test-manifest", CveID: "CVE-2021-1234"}, + initError: true, + wantError: "test initialization failure", + }, + } { + t.Run(testCase.name, func(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + match := v1beta1.Match{ + Vulnerability: v1beta1.Vulnerability{ + VulnerabilityMetadata: v1beta1.VulnerabilityMetadata{ID: "CVE-2021-1234"}, + }, + } + spdxClient := kubescapefake.NewClientset(&v1beta1.VulnerabilityManifest{ + ObjectMeta: metav1.ObjectMeta{Name: "test-manifest", Namespace: defaultKubescapeNamespace}, + Spec: v1beta1.VulnerabilityManifestSpec{ + Payload: v1beta1.GrypeDocument{Matches: []v1beta1.Match{match}}, + }, + }) + tool := NewKubescapeToolWithClients(nil, nil, spdxClient.SpdxV1beta1()) + if testCase.initError { + tool = NewKubescapeToolWithError(errors.New("test initialization failure")) + } + + server := sdkmcp.NewServer(&sdkmcp.Implementation{Name: "kubescape", Version: "test"}, nil) + toolmcp.AddTool(server, "kubescape", &sdkmcp.Tool{ + Name: "kubescape_get_vulnerability_details", + }, tool.handleGetVulnerabilityDetails) + serverTransport, clientTransport := sdkmcp.NewInMemoryTransports() + go func() { _ = server.Run(ctx, serverTransport) }() + + client := sdkmcp.NewClient(&sdkmcp.Implementation{Name: "test", Version: "test"}, nil) + session, err := client.Connect(ctx, clientTransport, nil) + require.NoError(t, err) + defer func() { _ = session.Close() }() + + result, err := session.CallTool(ctx, &sdkmcp.CallToolParams{ + Name: "kubescape_get_vulnerability_details", + Arguments: testCase.input, + }) + require.NoError(t, err, "SDK must accept the typed response, including tool errors") + require.NotNil(t, result) + require.Len(t, result.Content, 1, "text content must not be duplicated") + if testCase.wantError != "" { + require.True(t, result.IsError) + require.Contains(t, getResultText(result), testCase.wantError) + return + } + require.False(t, result.IsError) + + structured, err := json.Marshal(result.StructuredContent) + require.NoError(t, err) + var output getVulnerabilityDetailsOutput + require.NoError(t, json.Unmarshal(structured, &output)) + require.Equal(t, []v1beta1.Match{match}, output.Matches) + + var textMatches []v1beta1.Match + require.NoError(t, json.Unmarshal([]byte(getResultText(result)), &textMatches)) + require.Equal(t, output.Matches, textMatches) + }) + } +}