Skip to content

Commit 5d28623

Browse files
committed
write remaining list tools, make flatten recursive with per-tool depth available
1 parent f5291c4 commit 5d28623

File tree

6 files changed

+104
-70
lines changed

6 files changed

+104
-70
lines changed

pkg/github/issues.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313
ghErrors "github.com/github/github-mcp-server/pkg/errors"
1414
"github.com/github/github-mcp-server/pkg/inventory"
1515
"github.com/github/github-mcp-server/pkg/octicons"
16+
"github.com/github/github-mcp-server/pkg/response"
1617
"github.com/github/github-mcp-server/pkg/sanitize"
1718
"github.com/github/github-mcp-server/pkg/scopes"
1819
"github.com/github/github-mcp-server/pkg/translations"
@@ -610,7 +611,7 @@ func ListIssueTypes(t translations.TranslationHelperFunc) inventory.ServerTool {
610611
return ghErrors.NewGitHubAPIStatusErrorResponse(ctx, "failed to list issue types", resp, body), nil, nil
611612
}
612613

613-
r, err := json.Marshal(issueTypes)
614+
r, err := response.MarshalItems(issueTypes)
614615
if err != nil {
615616
return utils.NewToolResultErrorFromErr("failed to marshal issue types", err), nil, nil
616617
}
@@ -1605,7 +1606,7 @@ func ListIssues(t translations.TranslationHelperFunc) inventory.ServerTool {
16051606
}
16061607

16071608
// Create response with issues
1608-
response := map[string]any{
1609+
issueResponse := map[string]any{
16091610
"issues": issues,
16101611
"pageInfo": map[string]any{
16111612
"hasNextPage": pageInfo.HasNextPage,
@@ -1615,7 +1616,7 @@ func ListIssues(t translations.TranslationHelperFunc) inventory.ServerTool {
16151616
},
16161617
"totalCount": totalCount,
16171618
}
1618-
out, err := json.Marshal(response)
1619+
out, err := response.MarshalItems(issueResponse)
16191620
if err != nil {
16201621
return nil, nil, fmt.Errorf("failed to marshal issues: %w", err)
16211622
}

pkg/github/issues_test.go

Lines changed: 17 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1188,7 +1188,6 @@ func Test_ListIssues(t *testing.T) {
11881188
expectError bool
11891189
errContains string
11901190
expectedCount int
1191-
verifyOrder func(t *testing.T, issues []*github.Issue)
11921191
}{
11931192
{
11941193
name: "list all issues",
@@ -1296,32 +1295,29 @@ func Test_ListIssues(t *testing.T) {
12961295
}
12971296
require.NoError(t, err)
12981297

1299-
// Parse the structured response with pagination info
1300-
var response struct {
1301-
Issues []*github.Issue `json:"issues"`
1302-
PageInfo struct {
1303-
HasNextPage bool `json:"hasNextPage"`
1304-
HasPreviousPage bool `json:"hasPreviousPage"`
1305-
StartCursor string `json:"startCursor"`
1306-
EndCursor string `json:"endCursor"`
1307-
} `json:"pageInfo"`
1308-
TotalCount int `json:"totalCount"`
1309-
}
1298+
// Parse the response
1299+
var response map[string]any
13101300
err = json.Unmarshal([]byte(text), &response)
13111301
require.NoError(t, err)
13121302

1313-
assert.Len(t, response.Issues, tc.expectedCount, "Expected %d issues, got %d", tc.expectedCount, len(response.Issues))
1303+
// Metadata should be preserved
1304+
assert.NotNil(t, response["totalCount"])
1305+
pageInfo, ok := response["pageInfo"].(map[string]any)
1306+
require.True(t, ok, "pageInfo should be a map")
1307+
assert.Contains(t, pageInfo, "hasNextPage")
1308+
assert.Contains(t, pageInfo, "endCursor")
13141309

1315-
// Verify order if verifyOrder function is provided
1316-
if tc.verifyOrder != nil {
1317-
tc.verifyOrder(t, response.Issues)
1318-
}
1310+
issues, ok := response["issues"].([]any)
1311+
require.True(t, ok)
1312+
assert.Len(t, issues, tc.expectedCount, "Expected %d issues, got %d", tc.expectedCount, len(issues))
13191313

13201314
// Verify that returned issues have expected structure
1321-
for _, issue := range response.Issues {
1322-
assert.NotNil(t, issue.Number, "Issue should have number")
1323-
assert.NotNil(t, issue.Title, "Issue should have title")
1324-
assert.NotNil(t, issue.State, "Issue should have state")
1315+
for _, issue := range issues {
1316+
m, ok := issue.(map[string]any)
1317+
require.True(t, ok)
1318+
assert.NotNil(t, m["number"], "Issue should have number")
1319+
assert.NotNil(t, m["title"], "Issue should have title")
1320+
assert.NotNil(t, m["state"], "Issue should have state")
13251321
}
13261322
})
13271323
}

pkg/github/repositories.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313
ghErrors "github.com/github/github-mcp-server/pkg/errors"
1414
"github.com/github/github-mcp-server/pkg/inventory"
1515
"github.com/github/github-mcp-server/pkg/octicons"
16+
"github.com/github/github-mcp-server/pkg/response"
1617
"github.com/github/github-mcp-server/pkg/scopes"
1718
"github.com/github/github-mcp-server/pkg/translations"
1819
"github.com/github/github-mcp-server/pkg/utils"
@@ -216,7 +217,7 @@ func ListCommits(t translations.TranslationHelperFunc) inventory.ServerTool {
216217
minimalCommits[i] = convertToMinimalCommit(commit, false)
217218
}
218219

219-
r, err := json.Marshal(minimalCommits)
220+
r, err := response.MarshalItems(minimalCommits, 3)
220221
if err != nil {
221222
return nil, nil, fmt.Errorf("failed to marshal response: %w", err)
222223
}
@@ -1496,7 +1497,7 @@ func ListTags(t translations.TranslationHelperFunc) inventory.ServerTool {
14961497
return ghErrors.NewGitHubAPIStatusErrorResponse(ctx, "failed to list tags", resp, body), nil, nil
14971498
}
14981499

1499-
r, err := json.Marshal(tags)
1500+
r, err := response.MarshalItems(tags)
15001501
if err != nil {
15011502
return nil, nil, fmt.Errorf("failed to marshal response: %w", err)
15021503
}
@@ -1669,7 +1670,7 @@ func ListReleases(t translations.TranslationHelperFunc) inventory.ServerTool {
16691670
return ghErrors.NewGitHubAPIStatusErrorResponse(ctx, "failed to list releases", resp, body), nil, nil
16701671
}
16711672

1672-
r, err := json.Marshal(releases)
1673+
r, err := response.MarshalItems(releases)
16731674
if err != nil {
16741675
return nil, nil, fmt.Errorf("failed to marshal response: %w", err)
16751676
}

pkg/github/repositories_test.go

Lines changed: 17 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1053,23 +1053,30 @@ func Test_ListCommits(t *testing.T) {
10531053
textContent := getTextResult(t, result)
10541054

10551055
// Unmarshal and verify the result
1056-
var returnedCommits []MinimalCommit
1056+
var returnedCommits []map[string]any
10571057
err = json.Unmarshal([]byte(textContent.Text), &returnedCommits)
10581058
require.NoError(t, err)
10591059
assert.Len(t, returnedCommits, len(tc.expectedCommits))
10601060
for i, commit := range returnedCommits {
1061-
assert.Equal(t, tc.expectedCommits[i].GetSHA(), commit.SHA)
1062-
assert.Equal(t, tc.expectedCommits[i].GetHTMLURL(), commit.HTMLURL)
1061+
assert.Equal(t, tc.expectedCommits[i].GetSHA(), commit["sha"])
1062+
assert.Equal(t, tc.expectedCommits[i].GetHTMLURL(), commit["html_url"])
10631063
if tc.expectedCommits[i].Commit != nil {
1064-
assert.Equal(t, tc.expectedCommits[i].Commit.GetMessage(), commit.Commit.Message)
1064+
assert.Equal(t, tc.expectedCommits[i].Commit.GetMessage(), commit["commit.message"])
1065+
if tc.expectedCommits[i].Commit.Author != nil {
1066+
assert.Equal(t, tc.expectedCommits[i].Commit.Author.GetName(), commit["commit.author.name"])
1067+
assert.Equal(t, tc.expectedCommits[i].Commit.Author.GetEmail(), commit["commit.author.email"])
1068+
if tc.expectedCommits[i].Commit.Author.Date != nil {
1069+
assert.NotEmpty(t, commit["commit.author.date"], "commit.author.date should be present")
1070+
}
1071+
}
10651072
}
10661073
if tc.expectedCommits[i].Author != nil {
1067-
assert.Equal(t, tc.expectedCommits[i].Author.GetLogin(), commit.Author.Login)
1074+
assert.Equal(t, tc.expectedCommits[i].Author.GetLogin(), commit["author.login"])
10681075
}
10691076

10701077
// Files and stats are never included in list_commits
1071-
assert.Nil(t, commit.Files)
1072-
assert.Nil(t, commit.Stats)
1078+
assert.Nil(t, commit["files"])
1079+
assert.Nil(t, commit["stats"])
10731080
}
10741081
})
10751082
}
@@ -2782,15 +2789,15 @@ func Test_ListTags(t *testing.T) {
27822789
textContent := getTextResult(t, result)
27832790

27842791
// Parse and verify the result
2785-
var returnedTags []*github.RepositoryTag
2792+
var returnedTags []map[string]any
27862793
err = json.Unmarshal([]byte(textContent.Text), &returnedTags)
27872794
require.NoError(t, err)
27882795

27892796
// Verify each tag
27902797
require.Equal(t, len(tc.expectedTags), len(returnedTags))
27912798
for i, expectedTag := range tc.expectedTags {
2792-
assert.Equal(t, *expectedTag.Name, *returnedTags[i].Name)
2793-
assert.Equal(t, *expectedTag.Commit.SHA, *returnedTags[i].Commit.SHA)
2799+
assert.Equal(t, *expectedTag.Name, returnedTags[i]["name"])
2800+
assert.Equal(t, *expectedTag.Commit.SHA, returnedTags[i]["commit.sha"])
27942801
}
27952802
})
27962803
}

pkg/response/optimize.go

Lines changed: 37 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,16 @@ const defaultFillRateThreshold = 0.1
1212
// minFillRateRows is the minimum number of items required to apply fill-rate filtering
1313
const minFillRateRows = 3
1414

15+
// maxFlattenDepth is the maximum nesting depth that flatten will recurse into.
16+
// Deeper nested maps are silently dropped.
17+
const maxFlattenDepth = 2
18+
1519
// preservedFields is a set of keys that are exempt from all destructive strategies except whitespace normalization.
1620
// Keys are matched against post-flatten map keys, so for nested fields like "user.html_url", the dotted key must be
1721
// added explicitly. Empty collections are still dropped. Wins over collectionFieldExtractors.
1822
var preservedFields = map[string]bool{
1923
"html_url": true,
24+
"draft": true,
2025
}
2126

2227
// collectionFieldExtractors controls how array fields are handled instead of being summarized as "[N items]".
@@ -32,17 +37,24 @@ var collectionFieldExtractors = map[string][]string{
3237

3338
// MarshalItems is the single entry point for response optimization.
3439
// Handles two shapes: plain JSON arrays and wrapped objects with metadata.
35-
func MarshalItems(data any) ([]byte, error) {
40+
// An optional maxDepth controls how many nesting levels flatten will recurse
41+
// into; it defaults to maxFlattenDepth when omitted.
42+
func MarshalItems(data any, maxDepth ...int) ([]byte, error) {
43+
depth := maxFlattenDepth
44+
if len(maxDepth) > 0 {
45+
depth = maxDepth[0]
46+
}
47+
3648
raw, err := json.Marshal(data)
3749
if err != nil {
3850
return nil, fmt.Errorf("failed to marshal data: %w", err)
3951
}
4052

4153
switch raw[0] {
4254
case '[':
43-
return optimizeArray(raw)
55+
return optimizeArray(raw, depth)
4456
case '{':
45-
return optimizeObject(raw)
57+
return optimizeObject(raw, depth)
4658
default:
4759
return raw, nil
4860
}
@@ -51,13 +63,13 @@ func MarshalItems(data any) ([]byte, error) {
5163
// OptimizeItems runs the full optimization pipeline on a slice of items:
5264
// flatten, remove URLs, remove zero-values, normalize whitespace,
5365
// summarize collections, and fill-rate filtering.
54-
func OptimizeItems(items []map[string]any) []map[string]any {
66+
func OptimizeItems(items []map[string]any, depth int) []map[string]any {
5567
if len(items) == 0 {
5668
return items
5769
}
5870

5971
for i, item := range items {
60-
flattenedItem := flatten(item)
72+
flattenedItem := flattenTo(item, depth)
6173
items[i] = optimizeItem(flattenedItem)
6274
}
6375

@@ -68,29 +80,26 @@ func OptimizeItems(items []map[string]any) []map[string]any {
6880
return items
6981
}
7082

71-
// flatten promotes values from one-level-deep nested maps into the parent
72-
// using dot-notation keys ("user.login", "user.id"). Nested maps and arrays
73-
// within those nested maps are dropped.
74-
func flatten(item map[string]any) map[string]any {
83+
// flattenTo recursively promotes values from nested maps into the parent
84+
// using dot-notation keys ("user.login", "commit.author.date"). Arrays
85+
// within nested maps are preserved at their dotted key position.
86+
// Recursion stops at the given maxDepth; deeper nested maps are dropped.
87+
func flattenTo(item map[string]any, maxDepth int) map[string]any {
7588
result := make(map[string]any, len(item))
76-
for key, value := range item {
77-
nested, ok := value.(map[string]any)
78-
if !ok {
79-
result[key] = value
80-
continue
81-
}
89+
flattenInto(item, "", result, 1, maxDepth)
90+
return result
91+
}
8292

83-
for nk, nv := range nested {
84-
switch nv.(type) {
85-
case map[string]any, []any:
86-
// skip complex nested values
87-
default:
88-
result[key+"."+nk] = nv
89-
}
93+
// flattenInto is the recursive worker for flattenTo.
94+
func flattenInto(item map[string]any, prefix string, result map[string]any, depth int, maxDepth int) {
95+
for key, value := range item {
96+
fullKey := prefix + key
97+
if nested, ok := value.(map[string]any); ok && depth < maxDepth {
98+
flattenInto(nested, fullKey+".", result, depth+1, maxDepth)
99+
} else if !ok {
100+
result[fullKey] = value
90101
}
91102
}
92-
93-
return result
94103
}
95104

96105
// filterByFillRate drops keys that appear on less than the threshold proportion of items.
@@ -126,16 +135,16 @@ func filterByFillRate(items []map[string]any, threshold float64) []map[string]an
126135
}
127136

128137
// optimizeArray is the entry point for optimizing a raw JSON array.
129-
func optimizeArray(raw []byte) ([]byte, error) {
138+
func optimizeArray(raw []byte, depth int) ([]byte, error) {
130139
var items []map[string]any
131140
if err := json.Unmarshal(raw, &items); err != nil {
132141
return nil, fmt.Errorf("failed to unmarshal JSON: %w", err)
133142
}
134-
return json.Marshal(OptimizeItems(items))
143+
return json.Marshal(OptimizeItems(items, depth))
135144
}
136145

137146
// optimizeObject is the entry point for optimizing a raw JSON object.
138-
func optimizeObject(raw []byte) ([]byte, error) {
147+
func optimizeObject(raw []byte, depth int) ([]byte, error) {
139148
var wrapper map[string]any
140149
if err := json.Unmarshal(raw, &wrapper); err != nil {
141150
return nil, fmt.Errorf("failed to unmarshal JSON: %w", err)
@@ -161,7 +170,7 @@ func optimizeObject(raw []byte) ([]byte, error) {
161170
items = append(items, m)
162171
}
163172
}
164-
wrapper[dataKey] = OptimizeItems(items)
173+
wrapper[dataKey] = OptimizeItems(items, depth)
165174

166175
return json.Marshal(wrapper)
167176
}

0 commit comments

Comments
 (0)