Repository navigation
feat(billing-cost-management-mcp-server)!: add paginate_aws_items and move three tools to it - #4786
Open
finklen-amazon wants to merge 4 commits into
Open
finklen-amazon wants to merge 4 commits into
finklen-amazon wants to merge 4 commits into
Conversation
|
|
||
| LIST_COST_ALLOCATION_TAGS = PageSpec( | ||
| result_key='CostAllocationTags', | ||
| token_param='NextToken', |
| # Pages can come back short and still carry a token, so paging continues until no token. | ||
| GET_ANOMALIES = PageSpec( | ||
| result_key='Anomalies', | ||
| token_param='NextPageToken', |
|
|
||
| LIST_COST_CATEGORY_DEFINITIONS = PageSpec( | ||
| result_key='CostCategoryReferences', | ||
| token_param='NextToken', |
| _HINT_INVALID = 'The request arguments are invalid. Correct them using the message, then retry.' | ||
| _HINT_NOT_FOUND = 'Nothing matched the request. Broaden or correct the filters, then retry.' | ||
| _HINT_INVALID_NEXT_TOKEN = ( | ||
| 'The next_token is expired or no longer matches the request. Call again without ' |
|
|
||
|
|
||
| def _derive_pagination_envelope( | ||
| def _derive_legacy_pagination_envelope( |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #4786 +/- ##
==========================================
+ Coverage 93.55% 93.56% +0.01%
==========================================
Files 1068 1068
Lines 92355 92504 +149
Branches 15027 15048 +21
==========================================
+ Hits 86400 86549 +149
Misses 3562 3562
Partials 2393 2393 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
finklen-amazon
marked this pull request as ready for review
October 16, 2026 14:45
finklen-amazon
requested review from
a team,
Comusus,
Rahul-1404,
anniewn,
bhatia-di,
callnmm,
chittev,
johnwangwyx,
mattkelsey,
shsrams,
somsubhro and
wangyuhere
as code owners
October 16, 2026 14:45
5 of 6 tasks
… move three tools to it List tools returned only the first page unless the caller asked for more, so a truncated result looked complete, and the paging parameters meant different things per tool. paginate_aws_items pages one AWS operation according to a per-operation PageSpec, with keyword-only max_items (the total number of items to return) and next_token. Without max_items it fetches every page; with it, it stops before a page that could take the result past max_items and returns the service's own token unchanged. One call stops paging, on a page boundary, after BCM_MCP_PAGINATION_MAX_SECONDS (default 300). A partial result leads with complete_dataset: false and a warning, and the response carries a lowercase pagination block, inline and in the SQL-offload result. has_more is true only when next_token can continue. A rejected or expired token maps to error_type invalid_next_token. cost-category, cost-allocation-tags and cost-anomaly move to the new helper as worked examples. The old paginate_aws_response, and the offload's reading of its Pagination wrapper and of raw single-page tokens, stay unchanged for the tools that still use them. Also fixes list-cost-allocation-tags and list-cost-category-definitions paging to the end with no limit when next_token was passed without max_pages. BREAKING CHANGE: list-cost-category-definitions, list-cost-allocation-tags, list-cost-allocation-tag-backfill-history and cost-anomaly take max_items and next_token; max_results and max_pages are removed from them, and their responses carry the lowercase pagination block instead of Pagination.
Cover three paths PR 1 adds that no test reached: a nested result_key whose parent object is missing from a page, a PriceList entry that is not valid JSON during SQL offload, and a partial offload through the generic key/value converter, which must not report rows_retrieved.
paginate_aws_items treated an empty or blank next_token as no token, so the call silently restarted from the first page while the caller believed it continued a previous result. Reject it as a validation error before any request is made, the same way as max_items below 1.
finklen-amazon
force-pushed
the
bcm-pagination-1-core
branch
from
October 16, 2026 19:45
abd4b62 to
f4ebf36
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes
Summary
Changes
Most list tools returned only the first page unless the caller asked for more, so on large accounts a truncated result looked complete (for example a 5,000-row page reported as "5,000 accounts"). The paging parameters also meant different things per tool:
max_resultswas a page size in some tools and a total in others, andmax_pagescounted pages whose size ranged from 10 to 10,000 items.This PR adds the shared paging helper and moves three small tools to it as worked examples. The other tools keep the old helper, unchanged, until #4787 and #4788 move them; #4788 then removes it.
New contract, for every paginated operation:
max_itemsis the total number of items to return.next_tokencontinues a previous call.max_items, every page is fetched. With it, the operation stops before a page that could take the result pastmax_items, so it never returns more thanmax_items(the first page is always returned, even when it alone is larger, e.g. a fixed 300-row page). The server never builds or rewrites a token:next_tokenin a response is the service's own token, passed back unchanged.BCM_MCP_PAGINATION_MAX_SECONDS(default 300) seconds. The result is partial and continues withnext_token.complete_dataset: falseand awarning(telling the caller that totals and counts undercount, and how to continue), followed by a lowercasepaginationblock ahead of the data:{"complete_dataset": false, "warning": "Partial data: ...", "pagination": {"complete_dataset": false, "has_more": true, "next_token": "<service token>", "items_returned": 5000}, "...": "data"}rows_retrieved), and its sample-query descriptions are prefixed "PARTIAL DATA, undercounts:". If the offload fails, the data is returned inline instead of an error.InvalidNextTokenException,ExpiredNextTokenException,RequestChangedException,BillExpirationException, plus per-operation codes such as Compute Optimizer'sInvalidParameterValueExceptionwhen the call carried a token) returnserror_type: invalid_next_tokenwithretry_strategy: modify_requestand a hint to call again withoutnext_token.start_*) take neither parameter.Implementation (this PR):
utilities/aws_service_base.py: newpaginate_aws_items, built around a per-operationPageSpec(result key or nested path, token request/response fields, largest page, page-size parameter and minimum, fixed page size for services that bind the token to it, items per page-size unit, item counting, token error codes). Every request resends the same parameters apart from the token and, where allowed, the page size. It detects a repeated token, and logs (and keeps) a page larger than the declared maximum. New helpers:with_pagination,first_page_fields,partial_data_warning,failure_warning,paging_limit_warning,paging_limit_reached,TOKEN_ERROR_CODES.handle_aws_errormapsInvalidNextTokenError. The oldpaginate_aws_responseis kept unchanged for the tools that have not moved yet.utilities/error_contract.py:InvalidNextTokenErrorand theinvalid_next_tokenguidance.utilities/sql_utils.py:_derive_pagination_envelopereads the lowercasepaginationblock and carries the partial signal into the offload result. Until every tool has moved, it still reads the two older shapes unmigrated tools emit (thePaginationwrapper built after the oldpaginate_aws_response, and a raw single-page continuation token named bypagination_token_key), with the same output as before; tests cover both formats.utilities/constants.py:PAGINATION_MAX_SECONDS.paginate_aws_items,PageSpecandwith_paginationare exported fromutilitiesand the package root, next topaginate_aws_response.cost-anomaly(GetAnomalies, replacing its hand-written loop),list-cost-category-definitions,list-cost-allocation-tagsandlist-cost-allocation-tag-backfill-history. EachPageSpecis defined next to its handler, and each tool description states whatmax_itemscounts, that omitting it fetches every page, and what to do with a partial result.BCM_MCP_PAGINATION_MAX_SECONDSand theinvalid_next_tokenrow in the error table. CHANGELOG entries under Changed and Fixed.tests/tools/test_cost_explorer_operations.py: the shared mock context gains an asyncdebug, because the offload now logs its size check outside thetrythat used to swallow a failed log call. The rest of that file's changes are in feat(billing-cost-management-mcp-server)!: move Cost Explorer and list tools to max_items #4787.Bug fixed along the way:
list-cost-allocation-tagsandlist-cost-category-definitionspaged to the end with no limit whenever anext_tokenwas passed withoutmax_pages.Pagination envelope
Every paginated tool response carries a lowercase
paginationblock:complete_datasettruewhen the last page was reached, so the result holds all the data the request matches.falsewhen more data remains or paging stopped early.has_moretrueonly whennext_tokencan continue the result. A result can be partial withhas_more: false(a page failed in a way that invalidates the token, or one region of a multi-region call failed), socomplete_datasetalone reports partiality.next_tokennullwhen there is nothing to continue (the data ended, or the token can no longer be used and the caller must restart). Pass it back with every other parameter unchanged.items_returnedPageSpec, see below). Not a total count of the data.Key order. When the result is partial, the response leads with top-level
complete_dataset: falseandwarning, thenpagination, then the data. A complete result starts withpagination, then the data. The order lets a reader, or a client that trims an oversized result, see whether the result is complete and how to continue before reading the data.Complete:
{"status": "success", "data": {"pagination": {"complete_dataset": true, "has_more": false, "next_token": null, "items_returned": 42}, "anomalies": ["..."]}}Partial (stopped at
max_items):{"status": "success", "data": {"complete_dataset": false, "warning": "Partial data: returned 1000 items and more remain. Totals, counts and top-N computed from this result undercount the real data. Call again with a larger max_items, continue with next_token, or state in the answer that the result is partial. Do not report the number of items returned as a total count.", "pagination": {"complete_dataset": false, "has_more": true, "next_token": "<service token>", "items_returned": 1000}, "anomalies": ["..."]}}The
warningtext depends on why paging stopped: more data remains (partial_data_warning), a later page failed (failure_warning, which names the error code and says whether to continue withnext_tokenor restart without it), or the time limit was reached (paging_limit_warning).SQL offload. When a response is large enough to be stored in the session SQL database, the offload result (returned as the tool's
data) carries the same signal in the same order:complete_dataset: false,warningandrows_retrieved(rows stored for this partial result; omitted for the generic key/value flatten, whose row count does not count items) when partial, then the samepaginationblock, thendata_stored,table_name,schema,row_count,sample_queriesandpreview. On a partial result each sample query description is prefixedPARTIAL DATA, undercounts:. If the offload fails, the data is returned inline, led by the same fields.{"status": "success", "complete_dataset": false, "warning": "Partial data: ...", "rows_retrieved": 5000, "pagination": {"complete_dataset": false, "has_more": true, "next_token": "<service token>", "items_returned": 5000}, "data_stored": true, "table_name": "...", "row_count": 5000, "sample_queries": ["..."]}Rejected token. When the first request of a continuation fails because the caller's
next_tokenwas rejected, the tool returns an error instead of data:{"status": "error", "error_type": "invalid_next_token", "message": "The next_token was rejected (ExpiredNextTokenException). It is expired or no longer matches the request.", "retry_strategy": "modify_request", "hint": "The next_token is expired or no longer matches the request. Call again without next_token to restart from the first page."}Shared utilities
Every list tool pages through the same helpers in
utilities/:PageSpecaws_service_base.pypaginate_aws_itemsaws_service_base.pyPageSpec,max_itemsandnext_token, and builds thepaginationblock.with_paginationaws_service_base.pypaginationblock to a tool's response data in the documented key order (partial lead, thenpagination, then data). Moveswarningout of the block to the top level.first_page_fieldsaws_service_base.pypartial_data_warningaws_service_base.pyfailure_warningaws_service_base.pynext_tokenor to restart.paging_limit_reachedaws_service_base.py300-second) once a call has paged forPAGINATION_MAX_SECONDS, otherwiseNone.paging_limit_warningaws_service_base.pyTOKEN_ERROR_CODESaws_service_base.pyInvalidNextTokenException,ExpiredNextTokenException,RequestChangedException,BillExpirationException.PAGINATION_MAX_SECONDSconstants.pyBCM_MCP_PAGINATION_MAX_SECONDS(positive integer).InvalidNextTokenErrorerror_contract.pyhandle_aws_errormaps it toerror_type: invalid_next_token, which the error contract maps toretry_strategy: modify_requestand the restart hint._derive_pagination_envelopesql_utils.pypaginationblock (and a partial result'swarning) from a response soconvert_response_if_neededcarries it into the offload result.PageSpecfields:result_key('SavingsPlansPurchaseRecommendation', 'Details').token_paramNextPageToken,nextToken).token_keypage_maxmax_items. A larger page is logged and kept.page_size_paramNoneNonewhen the service chooses it.page_min1fixed_page_sizeFalseitems_per_page_unit1DescribeBudgetNotificationsForAccount, where each budget holds up to 10 notifications).count_itemslentoken_error_codesTOKEN_ERROR_CODESInvalidNextTokenError; on a later page they end the call as partial withnext_token: null.caller_token_error_codesInvalidParameterValueException). On a later page they are an ordinary failure.paginate_aws_items:It returns
(items, pages, pagination): every kept item, the raw page responses (for merging fields outside the result list, e.g. withfirst_page_fields), and thepaginationblock, which carrieswarningwhen partial forwith_paginationto lift to the top level.The algorithm:
max_items: fetch every page until the service returns no token.max_items: request the exact remaining page size where the API allows it ((max_items - collected) // items_per_page_unit, clamped topage_minand the largest page; the largest page whenfixed_page_size). Before each page after the first, stop if the page could take the result pastmax_items(page size * items_per_page_unit, orpage_maxwhen the service chooses the size). The first page is always returned, even when it alone exceedsmax_items(e.g. a fixed 300-row page).request_paramsunchanged apart from the token and the page size.next_token: null, because continuing would repeat the same pages.next_tokenand the code is intoken_error_codesorcaller_token_error_codes,InvalidNextTokenErroris raised; any other error is re-raised unchanged and the tool returns it throughhandle_aws_erroras before.max_items < 1: raisesValueError("max_items must be at least 1. Omit it to fetch every page.").PAGINATION_MAX_SECONDS(default 300, envBCM_MCP_PAGINATION_MAX_SECONDS), paging stops on a page boundary and the result is partial with the service'snext_token.How a tool uses it (from
tools/cost_anomaly_tools.py):User experience
With a limit, e.g.
cost-anomalywithmax_items: 1000on an account with more anomalies, the response holds 1,000 anomalies and leads with:{"complete_dataset": false, "warning": "Partial data: returned 1000 items and more remain. ... Do not report the number of items returned as a total count.", "pagination": {"complete_dataset": false, "has_more": true, "next_token": "<service token>", "items_returned": 1000}}Without
max_items, the same call returns every anomaly.Breaking changes
max_resultsandmax_pagesare removed fromcost-anomaly,list-cost-category-definitions,list-cost-allocation-tagsandlist-cost-allocation-tag-backfill-history, which takemax_items/next_token. MCP clients reject unknown parameters, so callers that pass the old names must switch.max_itemsfetches every page (bounded by the 300 s time limit).paginationblock (complete_dataset,has_more,next_token,items_returned) instead of the uppercasePaginationblock. Partial results add top-levelcomplete_dataset: falseandwarningas the first keys.paginate_aws_items(ctx, operation, api_call, request_params, spec, *, max_items=None, next_token=None)returns(items, pages, pagination);paginate_aws_responseandconvert_response_if_needed(..., pagination_token_key=...)are unchanged until feat(billing-cost-management-mcp-server)!: move remaining tools to max_items and drop old helper #4788.Checklist
If your change doesn't seem to apply, please leave them unchecked.
Is this a breaking change? (Y/N) Y
RFC issue number:
Checklist:
Testing:
uv run --frozen pytest --cov --cov-branch(1974 passed, 97% branch coverage),ruff check,ruff format --checkandpyrightclean insrc/billing-cost-management-mcp-serveron this branch alone. New helper tests intests/utilities/test_pagination.py; each migrated operation has tests for fetching every page withoutmax_items, stopping atmax_items, and continuing with the caller's token and identical parameters.Acknowledgment
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of the project license.