Skip to content

Commit 63d313a

Browse files
mattdhollowayCopilotSamMorrowDrums
authored
fix(ui_get): bound pagination to cap latency (#2766)
ui_get backs a synchronous UI picker (label/assignee/etc. dropdowns in the MCP App issue/PR write surfaces). Each handler paginated GitHub API results in an unbounded loop (PerPage 100, looping until NextPage==0 / HasNextPage==false). On very large repos/orgs this fans out into dozens of sequential round-trips, and a single slow page inflates the whole call — production telemetry showed a tail spiking to ~20 minutes. Bound every ui_get pagination loop to uiGetMaxPages (10 pages, ~1000 items) and add an additive "has_more" flag indicating results were truncated. Truncation is acceptable here because the picker pairs it with typeahead, so responsiveness matters more than completeness. Affected methods: labels, assignees, milestones, branches, reviewers (both the collaborators and teams loops). uiGetIssueTypes and issue_fields are single-request and left unchanged. has_more is response-only and backward-compatible: no existing keys are removed or renamed, and the tool input schema is unchanged. HTTP-client timeouts are intentionally a separate follow-up. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Sam Morrow <info@sam-morrow.com>
1 parent 3c4749d commit 63d313a

2 files changed

Lines changed: 191 additions & 6 deletions

File tree

pkg/github/ui_tools.go

Lines changed: 49 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,15 @@ import (
2020
"github.com/shurcooL/githubv4"
2121
)
2222

23+
// uiGetMaxPages bounds how many pages each ui_get pagination loop will fetch.
24+
// ui_get backs a synchronous UI picker (dropdowns for labels/assignees/etc. in
25+
// the MCP App issue/PR write surfaces), so responsiveness matters more than
26+
// completeness. At PerPage 100 this caps a call at ~1000 items and a bounded
27+
// number of API round-trips, keeping latency predictable on very large
28+
// repos/orgs. Results past the cap are truncated and surfaced via a "has_more"
29+
// flag, which is acceptable because the picker pairs truncation with typeahead.
30+
const uiGetMaxPages = 10
31+
2332
// UIGet creates a tool to fetch UI data for MCP Apps.
2433
func UIGet(t translations.TranslationHelperFunc) inventory.ServerTool {
2534
st := NewTool(
@@ -131,7 +140,8 @@ func uiGetLabels(ctx context.Context, deps ToolDependencies, args map[string]any
131140

132141
labels := make([]map[string]any, 0)
133142
var totalCount int
134-
for {
143+
hasMore := false
144+
for page := 1; ; page++ {
135145
if err := client.Query(ctx, &query, vars); err != nil {
136146
return ghErrors.NewGitHubGraphQLErrorResponse(ctx, "Failed to list labels", err), nil, nil
137147
}
@@ -147,12 +157,17 @@ func uiGetLabels(ctx context.Context, deps ToolDependencies, args map[string]any
147157
if !query.Repository.Labels.PageInfo.HasNextPage {
148158
break
149159
}
160+
if page >= uiGetMaxPages {
161+
hasMore = true
162+
break
163+
}
150164
vars["cursor"] = githubv4.NewString(query.Repository.Labels.PageInfo.EndCursor)
151165
}
152166

153167
response := map[string]any{
154168
"labels": labels,
155169
"totalCount": totalCount,
170+
"has_more": hasMore,
156171
}
157172

158173
out, err := json.Marshal(response)
@@ -176,8 +191,9 @@ func uiGetAssignees(ctx context.Context, deps ToolDependencies, args map[string]
176191

177192
opts := &github.ListOptions{PerPage: 100}
178193
var allAssignees []*github.User
194+
hasMore := false
179195

180-
for {
196+
for page := 1; ; page++ {
181197
assignees, resp, err := client.Issues.ListAssignees(ctx, owner, repo, opts)
182198
if err != nil {
183199
return ghErrors.NewGitHubAPIErrorResponse(ctx, "failed to list assignees", resp, err), nil, nil
@@ -189,6 +205,10 @@ func uiGetAssignees(ctx context.Context, deps ToolDependencies, args map[string]
189205
if resp.NextPage == 0 {
190206
break
191207
}
208+
if page >= uiGetMaxPages {
209+
hasMore = true
210+
break
211+
}
192212
opts.Page = resp.NextPage
193213
}
194214

@@ -203,6 +223,7 @@ func uiGetAssignees(ctx context.Context, deps ToolDependencies, args map[string]
203223
out, err := json.Marshal(map[string]any{
204224
"assignees": result,
205225
"totalCount": len(result),
226+
"has_more": hasMore,
206227
})
207228
if err != nil {
208229
return utils.NewToolResultErrorFromErr("failed to marshal assignees", err), nil, nil
@@ -228,7 +249,8 @@ func uiGetMilestones(ctx context.Context, deps ToolDependencies, args map[string
228249
}
229250

230251
var allMilestones []*github.Milestone
231-
for {
252+
hasMore := false
253+
for page := 1; ; page++ {
232254
milestones, resp, err := client.Issues.ListMilestones(ctx, owner, repo, opts)
233255
if err != nil {
234256
return ghErrors.NewGitHubAPIErrorResponse(ctx, "failed to list milestones", resp, err), nil, nil
@@ -240,6 +262,10 @@ func uiGetMilestones(ctx context.Context, deps ToolDependencies, args map[string
240262
if resp.NextPage == 0 {
241263
break
242264
}
265+
if page >= uiGetMaxPages {
266+
hasMore = true
267+
break
268+
}
243269
opts.Page = resp.NextPage
244270
}
245271

@@ -262,6 +288,7 @@ func uiGetMilestones(ctx context.Context, deps ToolDependencies, args map[string
262288
out, err := json.Marshal(map[string]any{
263289
"milestones": result,
264290
"totalCount": len(result),
291+
"has_more": hasMore,
265292
})
266293
if err != nil {
267294
return utils.NewToolResultErrorFromErr("failed to marshal milestones", err), nil, nil
@@ -314,7 +341,8 @@ func uiGetBranches(ctx context.Context, deps ToolDependencies, args map[string]a
314341
}
315342

316343
var allBranches []*github.Branch
317-
for {
344+
hasMore := false
345+
for page := 1; ; page++ {
318346
branches, resp, err := client.Repositories.ListBranches(ctx, owner, repo, opts)
319347
if err != nil {
320348
return ghErrors.NewGitHubAPIErrorResponse(ctx, "failed to list branches", resp, err), nil, nil
@@ -326,6 +354,10 @@ func uiGetBranches(ctx context.Context, deps ToolDependencies, args map[string]a
326354
if resp.NextPage == 0 {
327355
break
328356
}
357+
if page >= uiGetMaxPages {
358+
hasMore = true
359+
break
360+
}
329361
opts.Page = resp.NextPage
330362
}
331363

@@ -337,6 +369,7 @@ func uiGetBranches(ctx context.Context, deps ToolDependencies, args map[string]a
337369
r, err := json.Marshal(map[string]any{
338370
"branches": minimalBranches,
339371
"totalCount": len(minimalBranches),
372+
"has_more": hasMore,
340373
})
341374
if err != nil {
342375
return utils.NewToolResultErrorFromErr("failed to marshal response", err), nil, nil
@@ -446,7 +479,8 @@ func uiGetReviewers(ctx context.Context, deps ToolDependencies, args map[string]
446479
ListOptions: github.ListOptions{PerPage: 100},
447480
}
448481
var allCollaborators []*github.User
449-
for {
482+
hasMore := false
483+
for page := 1; ; page++ {
450484
collaborators, resp, err := client.Repositories.ListCollaborators(ctx, owner, repo, collaboratorOpts)
451485
if err != nil {
452486
return ghErrors.NewGitHubAPIErrorResponse(ctx, "failed to list reviewers", resp, err), nil, nil
@@ -458,12 +492,16 @@ func uiGetReviewers(ctx context.Context, deps ToolDependencies, args map[string]
458492
if resp.NextPage == 0 {
459493
break
460494
}
495+
if page >= uiGetMaxPages {
496+
hasMore = true
497+
break
498+
}
461499
collaboratorOpts.Page = resp.NextPage
462500
}
463501

464502
teamOpts := &github.ListOptions{PerPage: 100}
465503
var allTeams []*github.Team
466-
for {
504+
for page := 1; ; page++ {
467505
teams, resp, err := client.Repositories.ListTeams(ctx, owner, repo, teamOpts)
468506
if err != nil {
469507
return ghErrors.NewGitHubAPIErrorResponse(ctx, "failed to list reviewer teams", resp, err), nil, nil
@@ -475,6 +513,10 @@ func uiGetReviewers(ctx context.Context, deps ToolDependencies, args map[string]
475513
if resp.NextPage == 0 {
476514
break
477515
}
516+
if page >= uiGetMaxPages {
517+
hasMore = true
518+
break
519+
}
478520
teamOpts.Page = resp.NextPage
479521
}
480522

@@ -503,6 +545,7 @@ func uiGetReviewers(ctx context.Context, deps ToolDependencies, args map[string]
503545
"users": users,
504546
"teams": teams,
505547
"totalCount": len(users) + len(teams),
548+
"has_more": hasMore,
506549
})
507550
if err != nil {
508551
return utils.NewToolResultErrorFromErr("failed to marshal reviewers", err), nil, nil

pkg/github/ui_tools_test.go

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,10 @@ package github
33
import (
44
"context"
55
"encoding/json"
6+
"fmt"
67
"net/http"
8+
"net/http/httptest"
9+
"strconv"
710
"testing"
811
"time"
912

@@ -17,6 +20,78 @@ import (
1720
"github.com/stretchr/testify/require"
1821
)
1922

23+
// recorderTransport routes HTTP requests through an in-process handler, mirroring
24+
// internal/githubv4mock's own transport. We need it because githubv4mock keys its
25+
// matchers by query string, so it cannot model a multi-page labels query: every
26+
// page issues the identical query and differs only by the $cursor variable. This
27+
// transport lets a single handler answer each page dynamically.
28+
type recorderTransport struct{ handler http.Handler }
29+
30+
func (rt recorderTransport) RoundTrip(req *http.Request) (*http.Response, error) {
31+
rec := httptest.NewRecorder()
32+
rt.handler.ServeHTTP(rec, req)
33+
return rec.Result(), nil
34+
}
35+
36+
// alwaysHasNextPageLabelsClient returns a GraphQL client whose labels query always
37+
// reports another page, advancing the cursor on each call. It exercises uiGetLabels'
38+
// page cap: the loop fetches one label per page until it stops at uiGetMaxPages with
39+
// has_more=true. totalCount is reported as a large server-side count so the test can
40+
// confirm it stays the full repo count even when results are truncated.
41+
func alwaysHasNextPageLabelsClient(t *testing.T) *http.Client {
42+
t.Helper()
43+
var calls int
44+
mux := http.NewServeMux()
45+
mux.HandleFunc("/graphql", func(w http.ResponseWriter, _ *http.Request) {
46+
calls++
47+
resp := map[string]any{
48+
"data": map[string]any{
49+
"repository": map[string]any{
50+
"labels": map[string]any{
51+
"nodes": []any{
52+
map[string]any{
53+
"id": fmt.Sprintf("label-%d", calls),
54+
"name": fmt.Sprintf("label-%d", calls),
55+
"color": "ededed",
56+
"description": "",
57+
},
58+
},
59+
"totalCount": 9999,
60+
"pageInfo": map[string]any{
61+
"hasNextPage": true,
62+
"endCursor": fmt.Sprintf("cursor-%d", calls),
63+
},
64+
},
65+
},
66+
},
67+
}
68+
w.Header().Set("Content-Type", "application/json")
69+
_ = json.NewEncoder(w).Encode(resp)
70+
})
71+
return &http.Client{Transport: recorderTransport{handler: mux}}
72+
}
73+
74+
// alwaysNextPageHandler returns a REST handler that always advertises another page
75+
// via the Link header, regardless of the page requested. It drives a pagination loop
76+
// purely off the page cap so tests can assert ui_get stops at uiGetMaxPages and sets
77+
// has_more=true. The same body is returned for every page, so the number of items
78+
// collected equals the number of pages fetched.
79+
func alwaysNextPageHandler(t *testing.T, body any) http.HandlerFunc {
80+
t.Helper()
81+
return func(w http.ResponseWriter, r *http.Request) {
82+
page := 1
83+
if p := r.URL.Query().Get("page"); p != "" {
84+
if parsed, err := strconv.Atoi(p); err == nil {
85+
page = parsed
86+
}
87+
}
88+
w.Header().Set("Link", fmt.Sprintf(`<https://api.github.com/next?page=%d>; rel="next"`, page+1))
89+
w.Header().Set("Content-Type", "application/json")
90+
w.WriteHeader(http.StatusOK)
91+
_ = json.NewEncoder(w).Encode(body)
92+
}
93+
}
94+
2095
func Test_UIGet(t *testing.T) {
2196
// Verify tool definition
2297
serverTool := UIGet(translations.NullTranslationHelper)
@@ -95,6 +170,7 @@ func Test_UIGet(t *testing.T) {
95170
require.NoError(t, json.Unmarshal([]byte(responseText), &response))
96171
assert.Contains(t, response, "assignees")
97172
assert.Contains(t, response, "totalCount")
173+
assert.Equal(t, false, response["has_more"], "results within the page cap should not be truncated")
98174
},
99175
},
100176
{
@@ -113,6 +189,7 @@ func Test_UIGet(t *testing.T) {
113189
require.NoError(t, json.Unmarshal([]byte(responseText), &response))
114190
assert.Contains(t, response, "branches")
115191
assert.Contains(t, response, "totalCount")
192+
assert.Equal(t, false, response["has_more"], "results within the page cap should not be truncated")
116193
},
117194
},
118195
{
@@ -228,6 +305,7 @@ func Test_UIGet(t *testing.T) {
228305
require.Len(t, labels, 1)
229306
assert.Equal(t, "bug", labels[0].(map[string]any)["name"])
230307
assert.Equal(t, float64(1), response["totalCount"])
308+
assert.Equal(t, false, response["has_more"], "results within the page cap should not be truncated")
231309
},
232310
},
233311
{
@@ -300,6 +378,70 @@ func Test_UIGet(t *testing.T) {
300378
assert.Equal(t, "docs", teams[0].(map[string]any)["slug"])
301379
assert.Equal(t, "owner", teams[0].(map[string]any)["org"])
302380
assert.Equal(t, float64(2), response["totalCount"])
381+
assert.Equal(t, false, response["has_more"], "results within the page cap should not be truncated")
382+
},
383+
},
384+
{
385+
name: "branches pagination stops at the page cap",
386+
mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{
387+
"GET /repos/owner/repo/branches": alwaysNextPageHandler(t, []*github.Branch{{Name: github.Ptr("feature")}}),
388+
}),
389+
requestArgs: map[string]any{
390+
"method": "branches",
391+
"owner": "owner",
392+
"repo": "repo",
393+
},
394+
expectError: false,
395+
validateResult: func(t *testing.T, responseText string) {
396+
var response map[string]any
397+
require.NoError(t, json.Unmarshal([]byte(responseText), &response))
398+
branches, ok := response["branches"].([]any)
399+
require.True(t, ok, "branches should be a list")
400+
assert.Len(t, branches, uiGetMaxPages, "loop should stop at the page cap")
401+
assert.Equal(t, float64(uiGetMaxPages), response["totalCount"], "totalCount should be the bounded count")
402+
assert.Equal(t, true, response["has_more"], "truncated results should set has_more")
403+
},
404+
},
405+
{
406+
name: "labels pagination stops at the page cap",
407+
mockedGQLClient: alwaysHasNextPageLabelsClient(t),
408+
requestArgs: map[string]any{
409+
"method": "labels",
410+
"owner": "owner",
411+
"repo": "repo",
412+
},
413+
expectError: false,
414+
validateResult: func(t *testing.T, responseText string) {
415+
var response map[string]any
416+
require.NoError(t, json.Unmarshal([]byte(responseText), &response))
417+
labels, ok := response["labels"].([]any)
418+
require.True(t, ok, "labels should be a list")
419+
assert.Len(t, labels, uiGetMaxPages, "loop should stop at the page cap")
420+
assert.Equal(t, true, response["has_more"], "truncated results should set has_more")
421+
// totalCount stays the server-reported full count, so it can exceed
422+
// the number of labels returned once results are truncated.
423+
assert.Equal(t, float64(9999), response["totalCount"])
424+
},
425+
},
426+
{
427+
name: "reviewers pagination stops at the page cap",
428+
mockedClient: MockHTTPClientWithHandlers(map[string]http.HandlerFunc{
429+
"GET /repos/owner/repo/collaborators": alwaysNextPageHandler(t, []*github.User{{Login: github.Ptr("octocat")}}),
430+
"GET /repos/owner/repo/teams": mockResponse(t, http.StatusOK, mockReviewerTeams),
431+
}),
432+
requestArgs: map[string]any{
433+
"method": "reviewers",
434+
"owner": "owner",
435+
"repo": "repo",
436+
},
437+
expectError: false,
438+
validateResult: func(t *testing.T, responseText string) {
439+
var response map[string]any
440+
require.NoError(t, json.Unmarshal([]byte(responseText), &response))
441+
users, ok := response["users"].([]any)
442+
require.True(t, ok, "users should be a list")
443+
assert.Len(t, users, uiGetMaxPages, "collaborators loop should stop at the page cap")
444+
assert.Equal(t, true, response["has_more"], "truncating either loop should set has_more")
303445
},
304446
},
305447
{

0 commit comments

Comments
 (0)