Skip to content

Commit 3b535b6

Browse files
Reduce MCP server JSON response size to minimize LLM token consumption (#360)
Implement lightweight response structure for MCP server that reduces JSON payload size while preserving all essential security data and adding enhanced SCM context for better repository identification. Changes: - Create mcpAnalysisResponse struct with only essential fields: findings, rules, purl, repository, scm_type, git_ref, commit_sha, last_commit - Remove embedded PackageInsights to eliminate heavy fields like github_actions_workflows, package_dependencies, and repo statistics - Update all MCP handlers (analyze_repo, analyze_local, analyze_org, analyze_stale_branches) to use lightweight response - Add comprehensive test suite to verify response structure and size - Add SCM context fields (purl, scm_type) per reviewer feedback - Rename 'ref' to 'git_ref' for clarity Results: - Lightweight response: ~182 bytes for empty findings vs kilobytes before - All essential security findings and repository metadata preserved - Better SCM identification with purl and scm_type fields - All tests passing with no regressions Fixes #359 Co-authored-by: François Proulx <[email protected]>
1 parent e2f9bd6 commit 3b535b6

File tree

2 files changed

+200
-13
lines changed

2 files changed

+200
-13
lines changed

cmd/mcp_lightweight_test.go

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
package cmd
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"testing"
7+
8+
"github.com/boostsecurityio/poutine/results"
9+
"github.com/stretchr/testify/assert"
10+
"github.com/stretchr/testify/require"
11+
)
12+
13+
// TestLightweightMCPResponse tests that the MCP server responses are lightweight
14+
// and don't include heavy fields like github_actions_workflows
15+
func TestLightweightMCPResponse(t *testing.T) {
16+
ctx := context.Background()
17+
18+
analyzer, err := createTestAnalyzer(ctx)
19+
require.NoError(t, err)
20+
21+
t.Run("analyze_manifest returns lightweight response", func(t *testing.T) {
22+
request := NewCallToolRequest("analyze_manifest", map[string]interface{}{
23+
"content": `name: Test Workflow
24+
on: push
25+
jobs:
26+
test:
27+
runs-on: ubuntu-latest
28+
steps:
29+
- uses: actions/checkout@v4`,
30+
"manifest_type": "github-actions",
31+
})
32+
33+
result, err := handleAnalyzeManifest(ctx, request, analyzer)
34+
require.NoError(t, err)
35+
require.NotNil(t, result)
36+
assert.False(t, result.IsError)
37+
38+
contentText := extractTextFromContent(t, result.Content[0])
39+
40+
// Verify it doesn't contain heavy fields
41+
assert.NotContains(t, contentText, "github_actions_workflows")
42+
assert.NotContains(t, contentText, "github_actions_metadata")
43+
assert.NotContains(t, contentText, "azure_pipelines")
44+
assert.NotContains(t, contentText, "gitlabci_configs")
45+
assert.NotContains(t, contentText, "package_dependencies")
46+
assert.NotContains(t, contentText, "build_dependencies")
47+
48+
// Verify it contains the essential fields
49+
assert.Contains(t, contentText, "findings")
50+
assert.Contains(t, contentText, "rules")
51+
52+
// Parse and verify structure
53+
var response map[string]interface{}
54+
err = json.Unmarshal([]byte(contentText), &response)
55+
require.NoError(t, err)
56+
57+
// Should only have findings and rules
58+
expectedFields := map[string]bool{
59+
"findings": true,
60+
"rules": true,
61+
}
62+
63+
for key := range response {
64+
_, expected := expectedFields[key]
65+
assert.True(t, expected, "Unexpected field '%s' in lightweight response", key)
66+
}
67+
})
68+
}
69+
70+
// TestMCPResponseStructure verifies the new mcpAnalysisResponse structure
71+
func TestMCPResponseStructure(t *testing.T) {
72+
// Create a sample response to verify JSON marshaling
73+
response := mcpAnalysisResponse{
74+
Findings: []results.Finding{},
75+
Rules: map[string]results.Rule{},
76+
Purl: "pkg:github/owner/repo@main",
77+
Repository: "owner/repo",
78+
ScmType: "github",
79+
GitRef: "main",
80+
CommitSha: "abc123",
81+
LastCommit: "2023-01-01T00:00:00Z",
82+
}
83+
84+
data, err := json.Marshal(response)
85+
require.NoError(t, err)
86+
87+
jsonStr := string(data)
88+
89+
// Verify lightweight structure
90+
assert.Contains(t, jsonStr, "\"findings\":")
91+
assert.Contains(t, jsonStr, "\"rules\":")
92+
assert.Contains(t, jsonStr, "\"purl\":")
93+
assert.Contains(t, jsonStr, "\"repository\":")
94+
assert.Contains(t, jsonStr, "\"scm_type\":")
95+
assert.Contains(t, jsonStr, "\"git_ref\":")
96+
assert.Contains(t, jsonStr, "\"commit_sha\":")
97+
assert.Contains(t, jsonStr, "\"last_commit\":")
98+
99+
// Verify it doesn't contain heavy fields
100+
assert.NotContains(t, jsonStr, "github_actions_workflows")
101+
assert.NotContains(t, jsonStr, "package_dependencies")
102+
assert.NotContains(t, jsonStr, "org_id")
103+
assert.NotContains(t, jsonStr, "repo_size")
104+
assert.NotContains(t, jsonStr, "forks_count")
105+
assert.NotContains(t, jsonStr, "stars_count")
106+
}
107+
108+
// TestMCPResponseOmitsEmptyFields verifies that empty optional fields are omitted
109+
func TestMCPResponseOmitsEmptyFields(t *testing.T) {
110+
response := mcpAnalysisResponse{
111+
Findings: []results.Finding{},
112+
Rules: map[string]results.Rule{},
113+
// Leave repository metadata fields empty
114+
}
115+
116+
data, err := json.Marshal(response)
117+
require.NoError(t, err)
118+
119+
jsonStr := string(data)
120+
121+
// These fields should be omitted when empty due to omitempty tag
122+
assert.NotContains(t, jsonStr, "\"purl\":")
123+
assert.NotContains(t, jsonStr, "\"repository\":")
124+
assert.NotContains(t, jsonStr, "\"scm_type\":")
125+
assert.NotContains(t, jsonStr, "\"git_ref\":")
126+
assert.NotContains(t, jsonStr, "\"commit_sha\":")
127+
assert.NotContains(t, jsonStr, "\"last_commit\":")
128+
129+
// These required fields should always be present
130+
assert.Contains(t, jsonStr, "\"findings\":")
131+
assert.Contains(t, jsonStr, "\"rules\":")
132+
}
133+
134+
// TestMCPResponseSizeReduction demonstrates the size reduction benefit
135+
func TestMCPResponseSizeReduction(t *testing.T) {
136+
// This test documents the reduction in response size
137+
// by comparing what the old response would have contained vs new response
138+
139+
lightweightResponse := mcpAnalysisResponse{
140+
Findings: []results.Finding{},
141+
Rules: map[string]results.Rule{},
142+
Purl: "pkg:github/test/repo@main",
143+
Repository: "test/repo",
144+
ScmType: "github",
145+
GitRef: "main",
146+
CommitSha: "abc123",
147+
LastCommit: "2023-01-01T00:00:00Z",
148+
}
149+
150+
lightweightData, err := json.Marshal(lightweightResponse)
151+
require.NoError(t, err)
152+
153+
t.Logf("Lightweight response size: %d bytes", len(lightweightData))
154+
t.Logf("Lightweight response: %s", string(lightweightData))
155+
156+
// The lightweight response should be significantly smaller than a full PackageInsights response
157+
// which would include many more fields like workflows, dependencies, repo stats, etc.
158+
assert.Less(t, len(lightweightData), 1000, "Lightweight response should be under 1KB for empty findings")
159+
}

cmd/mcp_server.go

Lines changed: 41 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,18 @@ import (
2121
"github.com/spf13/viper"
2222
)
2323

24+
// mcpAnalysisResponse provides a lightweight response for MCP tools
25+
// that includes only the essential security findings and minimal repository metadata
2426
type mcpAnalysisResponse struct {
25-
*models.PackageInsights
2627
Findings []results.Finding `json:"findings"`
2728
Rules map[string]results.Rule `json:"rules"`
29+
// Essential repository metadata
30+
Purl string `json:"purl,omitempty"`
31+
Repository string `json:"repository,omitempty"`
32+
ScmType string `json:"scm_type,omitempty"`
33+
GitRef string `json:"git_ref,omitempty"`
34+
CommitSha string `json:"commit_sha,omitempty"`
35+
LastCommit string `json:"last_commit,omitempty"`
2836
}
2937

3038
var mcpServerCmd = &cobra.Command{
@@ -340,9 +348,14 @@ func handleAnalyzeOrg(ctx context.Context, request mcp.CallToolRequest, defaultC
340348
combinedResponses := make([]mcpAnalysisResponse, 0, len(analysisResults))
341349
for _, pkgInsights := range analysisResults {
342350
combinedResponses = append(combinedResponses, mcpAnalysisResponse{
343-
Findings: pkgInsights.FindingsResults.Findings,
344-
Rules: pkgInsights.FindingsResults.Rules,
345-
PackageInsights: pkgInsights,
351+
Findings: pkgInsights.FindingsResults.Findings,
352+
Rules: pkgInsights.FindingsResults.Rules,
353+
Purl: pkgInsights.Purl,
354+
Repository: pkgInsights.SourceGitRepo,
355+
ScmType: pkgInsights.SourceScmType,
356+
GitRef: pkgInsights.SourceGitRef,
357+
CommitSha: pkgInsights.SourceGitCommitSha,
358+
LastCommit: pkgInsights.LastCommitedAt,
346359
})
347360
}
348361

@@ -386,9 +399,14 @@ func handleAnalyzeRepo(ctx context.Context, request mcp.CallToolRequest, default
386399
}
387400

388401
combinedResponse := mcpAnalysisResponse{
389-
Findings: analysisResults.FindingsResults.Findings,
390-
Rules: analysisResults.FindingsResults.Rules,
391-
PackageInsights: analysisResults,
402+
Findings: analysisResults.FindingsResults.Findings,
403+
Rules: analysisResults.FindingsResults.Rules,
404+
Purl: analysisResults.Purl,
405+
Repository: analysisResults.SourceGitRepo,
406+
ScmType: analysisResults.SourceScmType,
407+
GitRef: analysisResults.SourceGitRef,
408+
CommitSha: analysisResults.SourceGitCommitSha,
409+
LastCommit: analysisResults.LastCommitedAt,
392410
}
393411

394412
resultData, err := json.Marshal(combinedResponse)
@@ -440,9 +458,14 @@ func handleAnalyzeLocal(ctx context.Context, request mcp.CallToolRequest, opaCli
440458
}
441459

442460
combinedResponse := mcpAnalysisResponse{
443-
Findings: analysisResults.FindingsResults.Findings,
444-
Rules: analysisResults.FindingsResults.Rules,
445-
PackageInsights: analysisResults,
461+
Findings: analysisResults.FindingsResults.Findings,
462+
Rules: analysisResults.FindingsResults.Rules,
463+
Purl: analysisResults.Purl,
464+
Repository: analysisResults.SourceGitRepo,
465+
ScmType: analysisResults.SourceScmType,
466+
GitRef: analysisResults.SourceGitRef,
467+
CommitSha: analysisResults.SourceGitCommitSha,
468+
LastCommit: analysisResults.LastCommitedAt,
446469
}
447470

448471
resultData, err := json.Marshal(combinedResponse)
@@ -493,9 +516,14 @@ func handleAnalyzeStaleBranches(ctx context.Context, request mcp.CallToolRequest
493516
}
494517

495518
combinedResponse := mcpAnalysisResponse{
496-
Findings: analysisResults.FindingsResults.Findings,
497-
Rules: analysisResults.FindingsResults.Rules,
498-
PackageInsights: analysisResults,
519+
Findings: analysisResults.FindingsResults.Findings,
520+
Rules: analysisResults.FindingsResults.Rules,
521+
Purl: analysisResults.Purl,
522+
Repository: analysisResults.SourceGitRepo,
523+
ScmType: analysisResults.SourceScmType,
524+
GitRef: analysisResults.SourceGitRef,
525+
CommitSha: analysisResults.SourceGitCommitSha,
526+
LastCommit: analysisResults.LastCommitedAt,
499527
}
500528

501529
resultData, err := json.Marshal(combinedResponse)

0 commit comments

Comments
 (0)