From 1256d2188a08fae4e5138aa32927411af8136fc7 Mon Sep 17 00:00:00 2001
From: Evan Kravitz <evakravi@amazon.com>
Date: Fri, 11 Sep 2026 17:40:26 +0000
Subject: [PATCH] fix(jumpstart): send NextToken when Hub.list_models pages
 through hub contents

`Hub._list_and_paginate_models` reads `NextToken` from each `list_hub_contents` response but never sends it, so every call requests the first page. On a hub with more than one page the loop never ends. `Hub.list_models()` on `SageMakerPublicHub` (743 models) ran for 5 minutes without a result.

Pass the token as `next_token` on each call after the first. `Session.list_hub_contents` maps it to `NextToken`.

---
X-AI-Prompt: Why does Hub.list_models never return on SageMakerPublicHub, and how do we fix the pagination?
X-AI-Tool: claude-code
---
 src/sagemaker/jumpstart/hub/hub.py            |  2 +
 .../unit/sagemaker/jumpstart/hub/test_hub.py  | 39 +++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/src/sagemaker/jumpstart/hub/hub.py b/src/sagemaker/jumpstart/hub/hub.py
index 692966cee4..825aef7afc 100644
--- a/src/sagemaker/jumpstart/hub/hub.py
+++ b/src/sagemaker/jumpstart/hub/hub.py
@@ -126,6 +126,8 @@ def _list_and_paginate_models(self, **kwargs) -> List[Dict[str, Any]]:
 
         while first_iteration or next_token:
             first_iteration = False
+            if next_token:
+                kwargs["next_token"] = next_token
             list_hub_content_response = self._sagemaker_session.list_hub_contents(**kwargs)
             hub_model_summaries.extend(list_hub_content_response.get("HubContentSummaries", []))
             next_token = list_hub_content_response.get("NextToken")
diff --git a/tests/unit/sagemaker/jumpstart/hub/test_hub.py b/tests/unit/sagemaker/jumpstart/hub/test_hub.py
index 29efb6b31f..0c10f33401 100644
--- a/tests/unit/sagemaker/jumpstart/hub/test_hub.py
+++ b/tests/unit/sagemaker/jumpstart/hub/test_hub.py
@@ -212,6 +212,45 @@ def test_describe_model_one_thrown_error(mock_describe_hub_content_response, sag
         )
 
 
+def test_list_models_sends_next_token_to_the_following_page(sagemaker_session):
+    pages = {
+        ("ModelReference", None): {
+            "HubContentSummaries": [{"HubContentName": "reference-a"}],
+            "NextToken": "page-2",
+        },
+        ("ModelReference", "page-2"): {"HubContentSummaries": [{"HubContentName": "reference-b"}]},
+        ("Model", None): {
+            "HubContentSummaries": [{"HubContentName": "model-a"}],
+            "NextToken": "page-2",
+        },
+        ("Model", "page-2"): {
+            "HubContentSummaries": [{"HubContentName": "model-b"}],
+            "NextToken": "page-3",
+        },
+        ("Model", "page-3"): {"HubContentSummaries": [{"HubContentName": "model-c"}]},
+    }
+
+    def list_hub_contents(**kwargs):
+        assert kwargs["hub_name"] == HUB_NAME
+        assert kwargs["max_results"] == 2
+        # Each page is served once. A second request for a page raises KeyError, not a loop.
+        return pages.pop((kwargs["hub_content_type"], kwargs.get("next_token")))
+
+    sagemaker_session.list_hub_contents = Mock(side_effect=list_hub_contents)
+    hub = Hub(hub_name=HUB_NAME, sagemaker_session=sagemaker_session)
+
+    response = hub.list_models(max_results=2)
+
+    assert [summary["HubContentName"] for summary in response["hub_content_summaries"]] == [
+        "reference-a",
+        "reference-b",
+        "model-a",
+        "model-b",
+        "model-c",
+    ]
+    assert pages == {}
+
+
 def test_create_hub_content_reference(sagemaker_session):
     hub = Hub(hub_name=HUB_NAME, sagemaker_session=sagemaker_session)
     model_name = "mock-model-one-huggingface"