Skip to content

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
awslabs:mainfrom
finklen-amazon:bcm-pagination-1-core
Open

finklen-amazon wants to merge 4 commits into
awslabs:mainfrom
finklen-amazon:bcm-pagination-1-core

Conversation

@finklen-amazon

@finklen-amazon finklen-amazon commented Oct 16, 2026 •

Copy link
Copy Markdown
Contributor

Part 1 of 3. Supersedes part of #4785, which is split into three stacked PRs so each can be reviewed on its own:

  1. feat(billing-cost-management-mcp-server)!: add paginate_aws_items and move three tools to it #4786 - shared helper (paginate_aws_items, PageSpec, envelope, error contract, SQL offload) and three worked examples: cost-category, cost-allocation-tags, cost-anomaly.
  2. feat(billing-cost-management-mcp-server)!: move Cost Explorer and list tools to max_items #4787 - Cost Explorer and simple list tools: cost-explorer, cost-comparison, ri-performance, sp-performance, sp-recommendation, sp-purchase-analyzer, sp-explorer, free-tier-usage, invoicing, invoice-units, billing views, get-billing-preferences, procurement-preferences, enterprise-support.
  3. feat(billing-cost-management-mcp-server)!: move remaining tools to max_items and drop old helper #4788 - tools with special handling (Billing Conductor, budgets, Compute Optimizer, Compute Optimizer Automation, Cost Optimization Hub, AWS Pricing, Pricing Calculator, credits), then removal of the old paginate_aws_response and the SQL offload's support for the old pagination shapes.

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_results was a page size in some tools and a total in others, and max_pages counted 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_items is the total number of items to return. next_token continues a previous call.
  • Without max_items, every page is fetched. With it, the operation stops before a page that could take the result past max_items, so it never returns more than max_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_token in a response is the service's own token, passed back unchanged.
  • Time limit: one call stops paging, on a page boundary, after BCM_MCP_PAGINATION_MAX_SECONDS (default 300) seconds. The result is partial and continues with next_token.
  • Partial results lead with complete_dataset: false and a warning (telling the caller that totals and counts undercount, and how to continue), followed by a lowercase pagination block 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"}
    The same fields, in the same order, are carried into the SQL-offload result (plus 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.
  • A failure after the first page keeps the pages already fetched and returns them as partial, with the token that requested the failed page so a continuation retries exactly that page. A failure on the first page is returned as an error as before.
  • A rejected or expired token (InvalidNextTokenException, ExpiredNextTokenException, RequestChangedException, BillExpirationException, plus per-operation codes such as Compute Optimizer's InvalidParameterValueException when the call carried a token) returns error_type: invalid_next_token with retry_strategy: modify_request and a hint to call again without next_token.
  • Operations that are not paginated (forecasts, single-object gets, start_*) take neither parameter.

Implementation (this PR):

  • utilities/aws_service_base.py: new paginate_aws_items, built around a per-operation PageSpec (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_error maps InvalidNextTokenError. The old paginate_aws_response is kept unchanged for the tools that have not moved yet.
  • utilities/error_contract.py: InvalidNextTokenError and the invalid_next_token guidance.
  • utilities/sql_utils.py: _derive_pagination_envelope reads the lowercase pagination block and carries the partial signal into the offload result. Until every tool has moved, it still reads the two older shapes unmigrated tools emit (the Pagination wrapper built after the old paginate_aws_response, and a raw single-page continuation token named by pagination_token_key), with the same output as before; tests cover both formats.
  • utilities/constants.py: PAGINATION_MAX_SECONDS.
  • paginate_aws_items, PageSpec and with_pagination are exported from utilities and the package root, next to paginate_aws_response.
  • Worked examples: cost-anomaly (GetAnomalies, replacing its hand-written loop), list-cost-category-definitions, list-cost-allocation-tags and list-cost-allocation-tag-backfill-history. Each PageSpec is defined next to its handler, and each tool description states what max_items counts, that omitting it fetches every page, and what to do with a partial result.
  • README: new "Pagination" section (scoped to the migrated tools until feat(billing-cost-management-mcp-server)!: move remaining tools to max_items and drop old helper #4788), BCM_MCP_PAGINATION_MAX_SECONDS and the invalid_next_token row in the error table. CHANGELOG entries under Changed and Fixed.
  • tests/tools/test_cost_explorer_operations.py: the shared mock context gains an async debug, because the offload now logs its size check outside the try that 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-tags and list-cost-category-definitions paged to the end with no limit whenever a next_token was passed without max_pages.

Pagination envelope

Every paginated tool response carries a lowercase pagination block:

Field Meaning
complete_dataset true when the last page was reached, so the result holds all the data the request matches. false when more data remains or paging stopped early.
has_more true only when next_token can continue the result. A result can be partial with has_more: false (a page failed in a way that invalidates the token, or one region of a multi-region call failed), so complete_dataset alone reports partiality.
next_token The service's own continuation token, passed through unchanged, or null when 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_returned The number of items in this response (counted by the operation's PageSpec, see below). Not a total count of the data.

Key order. When the result is partial, the response leads with top-level complete_dataset: false and warning, then pagination, then the data. A complete result starts with pagination, 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 warning text 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 with next_token or 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, warning and rows_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 same pagination block, then data_stored, table_name, schema, row_count, sample_queries and preview. On a partial result each sample query description is prefixed PARTIAL 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_token was 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/:

Helper Module Purpose
PageSpec aws_service_base.py Frozen dataclass describing how one AWS operation pages. Defined once per operation, next to its handler.
paginate_aws_items aws_service_base.py Fetches pages for one operation according to its PageSpec, max_items and next_token, and builds the pagination block.
with_pagination aws_service_base.py Attaches the pagination block to a tool's response data in the documented key order (partial lead, then pagination, then data). Moves warning out of the block to the top level.
first_page_fields aws_service_base.py Returns fields only the first page carries (totals, group definitions) so a tool can merge them into a multi-page result.
partial_data_warning aws_service_base.py Warning text for a result that stopped before the last page.
failure_warning aws_service_base.py Warning text for a result cut short by a failed page; tells the caller to continue with next_token or to restart.
paging_limit_reached aws_service_base.py Returns the name of the time limit (e.g. 300-second) once a call has paged for PAGINATION_MAX_SECONDS, otherwise None.
paging_limit_warning aws_service_base.py Warning text for a result stopped by the time limit.
TOKEN_ERROR_CODES aws_service_base.py Default error codes meaning a token can no longer be used: InvalidNextTokenException, ExpiredNextTokenException, RequestChangedException, BillExpirationException.
PAGINATION_MAX_SECONDS constants.py Seconds one call may spend paging. Default 300, set with BCM_MCP_PAGINATION_MAX_SECONDS (positive integer).
InvalidNextTokenError error_contract.py Raised when the caller's token is rejected on the first request. handle_aws_error maps it to error_type: invalid_next_token, which the error contract maps to retry_strategy: modify_request and the restart hint.
_derive_pagination_envelope sql_utils.py Reads the pagination block (and a partial result's warning) from a response so convert_response_if_needed carries it into the offload result.

PageSpec fields:

Field Default Controls
result_key required Response key holding the page's items, or a tuple path to a nested list, e.g. ('SavingsPlansPurchaseRecommendation', 'Details').
token_param required Request field that carries the continuation token (e.g. NextPageToken, nextToken).
token_key required Response field that returns the continuation token.
page_max required Largest number of items one page can hold. Bounds the worst case when deciding whether the next page could exceed max_items. A larger page is logged and kept.
page_size_param None Request field that sets the page size, or None when the service chooses it.
page_min 1 Smallest page size the service accepts.
fixed_page_size False Always request the largest page. For services whose token is bound to the page size, so a continuation must repeat it.
items_per_page_unit 1 Items one unit of page size can return, for APIs whose page size counts a parent object (e.g. DescribeBudgetNotificationsForAccount, where each budget holds up to 10 notifications).
count_items len Counts the items in one page's result list, for results whose items are nested (e.g. notifications per budget).
token_error_codes TOKEN_ERROR_CODES Codes meaning the token can no longer be used, on any page. On the first request of a continuation they raise InvalidNextTokenError; on a later page they end the call as partial with next_token: null.
caller_token_error_codes empty Codes that mean a rejected token only on the first request of a continuation, for services that report a bad token with a generic code (e.g. Compute Optimizer's InvalidParameterValueException). On a later page they are an ordinary failure.

paginate_aws_items:

async def paginate_aws_items(
    ctx: Context,
    operation: str,                  # operation name, for logging
    api_call: Callable[..., Any],    # called with keyword arguments; may return an awaitable
    request_params: Dict[str, Any],  # request parameters, without token or page size
    spec: PageSpec,
    *,
    max_items: Optional[int] = None,
    next_token: Optional[str] = None,
) -> Tuple[List[Any], List[Dict[str, Any]], Dict[str, Any]]:

It returns (items, pages, pagination): every kept item, the raw page responses (for merging fields outside the result list, e.g. with first_page_fields), and the pagination block, which carries warning when partial for with_pagination to lift to the top level.

The algorithm:

  • No max_items: fetch every page until the service returns no token.
  • With max_items: request the exact remaining page size where the API allows it ((max_items - collected) // items_per_page_unit, clamped to page_min and the largest page; the largest page when fixed_page_size). Before each page after the first, stop if the page could take the result past max_items (page size * items_per_page_unit, or page_max when the service chooses the size). The first page is always returned, even when it alone exceeds max_items (e.g. a fixed 300-row page).
  • Tokens: the service token is passed through unchanged; the server never builds or rewrites one. Every request resends request_params unchanged apart from the token and the page size.
  • Repeated token: if the service returns a token already used in this call, paging stops as partial with next_token: null, because continuing would repeat the same pages.
  • Failure after page 1: the pages already fetched are kept and returned as partial, with the token that requested the failed page, so a continuation retries exactly that page. If the failure was a token error, the token is dropped and the warning says to restart.
  • Failure on page 1: if the caller sent next_token and the code is in token_error_codes or caller_token_error_codes, InvalidNextTokenError is raised; any other error is re-raised unchanged and the tool returns it through handle_aws_error as before.
  • max_items < 1: raises ValueError ("max_items must be at least 1. Omit it to fetch every page.").
  • Time limit: after PAGINATION_MAX_SECONDS (default 300, env BCM_MCP_PAGINATION_MAX_SECONDS), paging stops on a page boundary and the result is partial with the service's next_token.

How a tool uses it (from tools/cost_anomaly_tools.py):

# 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',
    token_key='NextPageToken',
    page_max=10000,
    page_size_param='MaxResults',
)

...
        all_anomalies, _, pagination = await paginate_aws_items(
            ctx,
            'GetAnomalies',
            ce_client.get_anomalies,
            request_params,
            GET_ANOMALIES,
            max_items=max_items,
            next_token=next_token,
        )
        formatted_response = {'anomalies': [...]}  # built from all_anomalies
        return format_response('success', with_pagination(formatted_response, pagination))

    except Exception as e:
        return await handle_aws_error(ctx, e, 'get_anomalies', 'Cost Explorer')

User experience

With a limit, e.g. cost-anomaly with max_items: 1000 on 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

  • Tool parameters: max_results and max_pages are removed from cost-anomaly, list-cost-category-definitions, list-cost-allocation-tags and list-cost-allocation-tag-backfill-history, which take max_items / next_token. MCP clients reject unknown parameters, so callers that pass the old names must switch.
  • Default behavior: omitting max_items fetches every page (bounded by the 300 s time limit).
  • Response shape: these tools' responses carry the lowercase pagination block (complete_dataset, has_more, next_token, items_returned) instead of the uppercase Pagination block. Partial results add top-level complete_dataset: false and warning as the first keys.
  • Python API: additive in this PR. paginate_aws_items(ctx, operation, api_call, request_params, spec, *, max_items=None, next_token=None) returns (items, pages, pagination); paginate_aws_response and convert_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.

  • I have reviewed the contributing guidelines
  • I have performed a self-review of this change
  • Changes have been tested
  • Changes are documented

Is this a breaking change? (Y/N) Y

RFC issue number:

Checklist:

  • Migration process documented
  • Implement warnings (if it can live side by side)

Testing: uv run --frozen pytest --cov --cov-branch (1974 passed, 97% branch coverage), ruff check, ruff format --check and pyright clean in src/billing-cost-management-mcp-server on this branch alone. New helper tests in tests/utilities/test_pagination.py; each migrated operation has tests for fetching every page without max_items, stopping at max_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.


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

codecov Bot commented Oct 16, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 93.56%. Comparing base (d534d0f) to head (f4ebf36).

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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

… 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.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: To triage

Development

Successfully merging this pull request may close these issues.

2 participants