Skip to content

MPT-22959 Buffer streamed bodies before 401 retry in extension framework auth #366

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
svazquezco wants to merge 3 commits into
base: main
Choose a base branch
from
+123 −58
Open

MPT-22959 Buffer streamed bodies before 401 retry in extension framework auth #366

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
770013f
fix(auth): buffer streamed bodies before 401 retry in extension frame…
svazquezco Jul 9, 2026
fed5e13
test(auth): use fixtures and result naming in streamed-body retry tests
svazquezco Jul 9, 2026
d0ad4ff
test(auth): reuse extension client fixtures across the test module
svazquezco Jul 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions docs/usage.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,8 +46,9 @@ implementations are available:
- `ExtensionFrameworkAuthentication` — a short-lived installation or account-scoped token
fetched from an extension secret via `POST /installations/-/token`. It refreshes
proactively once the token nears its JWT `exp` (default leeway 60s) and reactively on
`401`. Pass `account_id` to request a token scoped to a specific account
(`?account.id=<id>`); use one provider instance per account scope.
`401`. Request bodies are buffered in memory before sending so the `401` retry can
replay one-shot streamed bodies intact. Pass `account_id` to request a token scoped to
a specific account (`?account.id=<id>`); use one provider instance per account scope.

## Instantiate The Client

Expand Down
8 changes: 8 additions & 0 deletions mpt_api_client/auth/extension_framework.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,12 @@ class ExtensionFrameworkAuthentication(Authentication):
account scopes should construct one provider per scope.
"""

# The 401 retry re-sends the request, so one-shot streamed bodies (raw iterators,
# non-seekable files) must be buffered before the first send; otherwise the retry
# replays an already-consumed stream. httpx only honours this flag in the default
# ``auth_flow`` wrappers, so the overridden flows below also buffer explicitly.
requires_request_body = True

def __init__(
self,
secret: str,
Expand Down Expand Up @@ -76,6 +82,7 @@ def sync_auth_flow(
self, request: httpx.Request
) -> Generator[httpx.Request, httpx.Response, None]:
"""Attach a cached token, refreshing it proactively and on 401."""
request.read()

@codeant-ai codeant-ai Bot Jul 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion: The new unconditional pre-read buffers every request body into memory before the first send, which breaks true streaming semantics and can cause very large uploads to consume excessive RAM (or fail with out-of-memory) even when no 401 retry happens. Restrict buffering to bounded-size bodies or add a safe fallback path for large/non-bufferable streams instead of always materializing the full payload. [performance]

Severity Level: Major ⚠️ 
- ❌ All sync authenticated requests buffer bodies in memory.
- ⚠️ Large streamed uploads risk high memory consumption.
- ⚠️ Streaming semantics lost for iterator-based request bodies.
Steps of Reproduction ✅ 
1. In tests/unit/auth/test_extension_framework.py:24, construct an HTTPClient with
ExtensionFrameworkAuthentication by calling HTTPClient(base_url=API_URL,
authentication=ExtensionFrameworkAuthentication(SECRET)).

2. In the same test file at lines 26-28, call client.httpx_client.request("POST",
"/orders", content=iter([b"chunk-1", b"chunk-2"])) to send a streamed body backed by an
iterator.

3. This request triggers ExtensionFrameworkAuthentication.sync_auth_flow at
mpt_api_client/auth/extension_framework.py:81-93; at line 85 the code executes
request.read(), which causes httpx to fully buffer the entire request body into memory and
replace the one-shot stream with a ByteStream.

4. If the iterator in step 2 were producing a very large upload instead of two small
chunks, sync_auth_flow would still unconditionally materialize the whole payload before
sending, defeating streaming semantics and potentially consuming excessive memory or
causing out-of-memory errors even when the first request succeeds without a 401 retry.

Fix in Cursor Fix in VSCode Claude

(Use Cmd/Ctrl + Click for best experience)

Prompt for AI Agent 🤖 
This is a comment left during a code review.

**Path:** mpt_api_client/auth/extension_framework.py
**Line:** 85:85
**Comment:**
	*Performance: The new unconditional pre-read buffers every request body into memory before the first send, which breaks true streaming semantics and can cause very large uploads to consume excessive RAM (or fail with out-of-memory) even when no 401 retry happens. Restrict buffering to bounded-size bodies or add a safe fallback path for large/non-bufferable streams instead of always materializing the full payload.

Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
👍 | 👎

@svazquezco svazquezco Jul 9, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the review — this behaviour is intentional, so we're keeping the unconditional buffering:

  1. It is the deliberate trade-off of this fix. The goal of MPT-22959 is that the request body is read into memory before the auth flow so the 401 retry is always replayable. It is also exactly what httpx itself does for auth schemes that declare requires_request_body = True: the default auth-flow wrapper performs the same unconditional request.read(), with no size bound or fallback. Since this class overrides the flow wrappers, the explicit request.read() / await request.aread() just replicates the standard httpx convention.

  2. Restricting buffering to bounded-size bodies would reintroduce the bug. Large non-seekable streamed uploads are precisely the case where the 401 retry silently corrupts data (truncated multipart parts, httpx.StreamConsumed). A size-capped or fallback path would trade silent data corruption for memory savings on exactly the requests that matter most — the wrong trade for an API client.

  3. The practical impact is limited. Typical bodies in this client (json=, multipart with seekable files) are either already in memory or cheaply replayable; huge non-seekable streams are a marginal case, and the buffering behaviour is now documented in docs/usage.md so callers can account for it.

@codeant-ai codeant-ai Bot Jul 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ Customized review instruction saved!

Instruction:

Do not flag unconditional request-body buffering in auth-flow wrappers that require replayable requests for 401 retries and intentionally set requires_request_body = True.

Applied to:

  • mpt_api_client/auth/**

💡 To manage or update this instruction, visit: CodeAnt AI Settings

if self._token is None or self._is_expired():
self._refresh_sync()
request.headers["Authorization"] = f"Bearer {self._token}"
Expand All @@ -90,6 +97,7 @@ async def async_auth_flow(
self, request: httpx.Request
) -> AsyncGenerator[httpx.Request, httpx.Response]:
"""Attach a cached token, refreshing it proactively and on 401."""
await request.aread()

@codeant-ai codeant-ai Bot Jul 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion: The async flow now always drains the full request stream into memory before sending, which can block or exhaust memory for large/long-lived async streams and removes expected streaming behavior for normal successful requests. Add size-aware buffering or skip replay logic for oversized/non-replayable streams to avoid forcing full in-memory reads. [performance]

Severity Level: Major ⚠️ 
- ❌ Async extension client always buffers request bodies in memory.
- ⚠️ Large async streams can exhaust memory resources.
- ⚠️ Async streaming uploads lose progressive send behavior.
Steps of Reproduction ✅ 
1. In tests/unit/auth/test_extension_framework.py:154-156, construct an AsyncHTTPClient
with ExtensionFrameworkAuthentication by calling AsyncHTTPClient(base_url=API_URL,
authentication=ExtensionFrameworkAuthentication(SECRET)).

2. In the same file at lines 158-162, define an async generator stream_body that yields
multiple byte chunks, simulating a streamed request body produced over time.

3. At lines 163-165, call client.httpx_client.request("POST", "/orders",
content=stream_body()) to send the async streamed content through the authenticated
client.

4. This call triggers ExtensionFrameworkAuthentication.async_auth_flow at
mpt_api_client/auth/extension_framework.py:96-108; at line 100 the code executes await
request.aread(), which drains the entire async stream into memory before the first send,
so for large or long-lived streams the authentication layer forces full in-memory
buffering and removes true streaming behavior, increasing memory usage even when the
initial request does not receive a 401 response.

Fix in Cursor Fix in VSCode Claude

(Use Cmd/Ctrl + Click for best experience)

Prompt for AI Agent 🤖 
This is a comment left during a code review.

**Path:** mpt_api_client/auth/extension_framework.py
**Line:** 100:100
**Comment:**
	*Performance: The async flow now always drains the full request stream into memory before sending, which can block or exhaust memory for large/long-lived async streams and removes expected streaming behavior for normal successful requests. Add size-aware buffering or skip replay logic for oversized/non-replayable streams to avoid forcing full in-memory reads.

Validate the correctness of the flagged issue. If correct, How can I resolve this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask user if the user wants to fix the rest of the comments as well. if said yes, then fetch all the comments validate the correctness and implement a minimal fix
👍 | 👎

if self._token is None or self._is_expired():
await self._refresh_async()
request.headers["Authorization"] = f"Bearer {self._token}"
Expand Down
168 changes: 112 additions & 56 deletions tests/unit/auth/test_extension_framework.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import base64
import datetime as dt
import json
from collections.abc import AsyncGenerator

import httpx
import pytest
Expand All @@ -16,6 +17,18 @@
ORDERS_URL = f"{API_URL}/orders"


@pytest.fixture
def extension_http_client() -> HTTPClient:
return HTTPClient(base_url=API_URL, authentication=ExtensionFrameworkAuthentication(SECRET))


@pytest.fixture
def async_extension_http_client() -> AsyncHTTPClient:
return AsyncHTTPClient(
base_url=API_URL, authentication=ExtensionFrameworkAuthentication(SECRET)
)


def _jwt_with_exp(expires_at: dt.datetime, subject: str = "token") -> str:
def encode(payload: object) -> str:
raw = json.dumps(payload).encode("utf-8")
Expand All @@ -26,14 +39,13 @@ def encode(payload: object) -> str:


@respx.mock
def test_extension_framework_fetches_and_applies_token():
def test_extension_framework_fetches_and_applies_token(extension_http_client):
token_route = respx.post(TOKEN_URL).mock(
return_value=httpx.Response(200, json={"token": "installation-token"})
)
target_route = respx.get(ORDERS_URL).mock(return_value=httpx.Response(200, json={"data": []}))
client = HTTPClient(base_url=API_URL, authentication=ExtensionFrameworkAuthentication(SECRET))

client.request("GET", "/orders") # act
extension_http_client.request("GET", "/orders") # act

token_request = token_route.calls.last.request
target_request = target_route.calls.last.request
Expand All @@ -42,22 +54,21 @@ def test_extension_framework_fetches_and_applies_token():


@respx.mock
def test_extension_framework_caches_token_across_requests():
def test_extension_framework_caches_token_across_requests(extension_http_client):
token_route = respx.post(TOKEN_URL).mock(
return_value=httpx.Response(200, json={"token": "installation-token"})
)
respx.get(ORDERS_URL).mock(return_value=httpx.Response(200, json={"data": []}))
client = HTTPClient(base_url=API_URL, authentication=ExtensionFrameworkAuthentication(SECRET))

client.request("GET", "/orders") # first call populates the cache
extension_http_client.request("GET", "/orders") # first call populates the cache

client.request("GET", "/orders") # act
extension_http_client.request("GET", "/orders") # act

assert token_route.call_count == 1


@respx.mock
def test_extension_framework_refreshes_expired_token():
def test_extension_framework_refreshes_expired_token(extension_http_client):
Comment thread
coderabbitai[bot] marked this conversation as resolved.
token_route = respx.post(TOKEN_URL).mock(
side_effect=[
httpx.Response(200, json={"token": "stale-token"}),
Expand All @@ -70,9 +81,8 @@ def test_extension_framework_refreshes_expired_token():
httpx.Response(200, json={"data": []}),
]
)
client = HTTPClient(base_url=API_URL, authentication=ExtensionFrameworkAuthentication(SECRET))

response = client.request("GET", "/orders") # act
response = extension_http_client.request("GET", "/orders") # act

target_request = target_route.calls.last.request
assert response.status_code == httpx.codes.OK
Expand All @@ -81,7 +91,9 @@ def test_extension_framework_refreshes_expired_token():


@respx.mock
def test_extension_framework_retries_non_idempotent_request_on_unauthorized():
def test_extension_framework_retries_non_idempotent_request_on_unauthorized(
extension_http_client,
):
token_route = respx.post(TOKEN_URL).mock(
side_effect=[
httpx.Response(200, json={"token": "stale-token"}),
Expand All @@ -94,9 +106,8 @@ def test_extension_framework_retries_non_idempotent_request_on_unauthorized():
httpx.Response(201, json={"id": "ORD-1"}),
]
)
client = HTTPClient(base_url=API_URL, authentication=ExtensionFrameworkAuthentication(SECRET))

response = client.request("POST", "/orders", json={"x": 1}) # act
response = extension_http_client.request("POST", "/orders", json={"x": 1}) # act

target_request = target_route.calls.last.request
assert response.status_code == httpx.codes.CREATED
Expand All @@ -106,56 +117,78 @@ def test_extension_framework_retries_non_idempotent_request_on_unauthorized():


@respx.mock
def test_extension_framework_surfaces_repeated_unauthorized():
def test_extension_framework_retry_resends_streamed_body(extension_http_client):
token_route = respx.post(TOKEN_URL).mock(
side_effect=[
httpx.Response(200, json={"token": "stale-token"}),
httpx.Response(200, json={"token": "fresh-token"}),
]
)
target_route = respx.post(ORDERS_URL).mock(
side_effect=[
httpx.Response(401, json={"error": "expired"}),
httpx.Response(201, json={"id": "ORD-1"}),
]
)

result = extension_http_client.httpx_client.request(
"POST", "/orders", content=iter([b"chunk-1", b"chunk-2"])
)

target_request = target_route.calls.last.request
assert result.status_code == httpx.codes.CREATED
assert token_route.call_count == 2
assert target_route.call_count == 2
assert target_request.headers["Authorization"] == "Bearer fresh-token"
assert target_request.content == b"chunk-1chunk-2"


@respx.mock
def test_extension_framework_surfaces_repeated_unauthorized(extension_http_client):
respx.post(TOKEN_URL).mock(return_value=httpx.Response(200, json={"token": "any-token"}))
target_route = respx.get(ORDERS_URL).mock(
return_value=httpx.Response(401, json={"error": "nope"})
)
client = HTTPClient(base_url=API_URL, authentication=ExtensionFrameworkAuthentication(SECRET))

with pytest.raises(MPTAPIError) as exc_info: # act
client.request("GET", "/orders")
extension_http_client.request("GET", "/orders")

assert exc_info.value.status_code == httpx.codes.UNAUTHORIZED
assert target_route.call_count == 2 # original + exactly one retry, then surfaced


@respx.mock
async def test_extension_framework_works_with_async_client():
async def test_extension_framework_works_with_async_client(async_extension_http_client):
token_route = respx.post(TOKEN_URL).mock(
return_value=httpx.Response(200, json={"token": "installation-token"})
)
target_route = respx.get(ORDERS_URL).mock(return_value=httpx.Response(200, json={"data": []}))
client = AsyncHTTPClient(
base_url=API_URL, authentication=ExtensionFrameworkAuthentication(SECRET)
)

await client.request("GET", "/orders") # act
await async_extension_http_client.request("GET", "/orders") # act

target_request = target_route.calls.last.request
assert token_route.called
assert target_request.headers["Authorization"] == "Bearer installation-token"


@respx.mock
async def test_extension_framework_caches_token_across_requests_async():
async def test_extension_framework_caches_token_across_requests_async(
async_extension_http_client,
):
token_route = respx.post(TOKEN_URL).mock(
return_value=httpx.Response(200, json={"token": "installation-token"})
)
respx.get(ORDERS_URL).mock(return_value=httpx.Response(200, json={"data": []}))
client = AsyncHTTPClient(
base_url=API_URL, authentication=ExtensionFrameworkAuthentication(SECRET)
)

await client.request("GET", "/orders") # first call populates the cache
await async_extension_http_client.request("GET", "/orders") # first call populates the cache

await client.request("GET", "/orders") # act
await async_extension_http_client.request("GET", "/orders") # act

assert token_route.call_count == 1


@respx.mock
async def test_extension_framework_refreshes_expired_token_async():
async def test_extension_framework_refreshes_expired_token_async(async_extension_http_client):
token_route = respx.post(TOKEN_URL).mock(
side_effect=[
httpx.Response(200, json={"token": "stale-token"}),
Expand All @@ -168,11 +201,8 @@ async def test_extension_framework_refreshes_expired_token_async():
httpx.Response(200, json={"data": []}),
]
)
client = AsyncHTTPClient(
base_url=API_URL, authentication=ExtensionFrameworkAuthentication(SECRET)
)

response = await client.request("GET", "/orders") # act
response = await async_extension_http_client.request("GET", "/orders") # act

target_request = target_route.calls.last.request
assert response.status_code == httpx.codes.OK
Expand All @@ -181,7 +211,9 @@ async def test_extension_framework_refreshes_expired_token_async():


@respx.mock
async def test_extension_framework_retries_non_idempotent_request_on_unauthorized_async():
async def test_extension_framework_retries_non_idempotent_request_on_unauthorized_async(
async_extension_http_client,
):
token_route = respx.post(TOKEN_URL).mock(
side_effect=[
httpx.Response(200, json={"token": "stale-token"}),
Expand All @@ -194,11 +226,8 @@ async def test_extension_framework_retries_non_idempotent_request_on_unauthorize
httpx.Response(201, json={"id": "ORD-1"}),
]
)
client = AsyncHTTPClient(
base_url=API_URL, authentication=ExtensionFrameworkAuthentication(SECRET)
)

response = await client.request("POST", "/orders", json={"x": 1}) # act
response = await async_extension_http_client.request("POST", "/orders", json={"x": 1}) # act

target_request = target_route.calls.last.request
assert response.status_code == httpx.codes.CREATED
Expand All @@ -208,29 +237,59 @@ async def test_extension_framework_retries_non_idempotent_request_on_unauthorize


@respx.mock
async def test_extension_framework_surfaces_repeated_unauthorized_async():
async def test_extension_framework_retry_resends_streamed_body_async(async_extension_http_client):
token_route = respx.post(TOKEN_URL).mock(
side_effect=[
httpx.Response(200, json={"token": "stale-token"}),
httpx.Response(200, json={"token": "fresh-token"}),
]
)
target_route = respx.post(ORDERS_URL).mock(
side_effect=[
httpx.Response(401, json={"error": "expired"}),
httpx.Response(201, json={"id": "ORD-1"}),
]
)

# httpx AsyncClient requires an async iterable for streamed content
async def stream_body() -> AsyncGenerator[bytes]: # noqa: RUF029
yield b"chunk-1"
yield b"chunk-2"

result = await async_extension_http_client.httpx_client.request(
"POST", "/orders", content=stream_body()
)

target_request = target_route.calls.last.request
assert result.status_code == httpx.codes.CREATED
assert token_route.call_count == 2
assert target_route.call_count == 2
assert target_request.headers["Authorization"] == "Bearer fresh-token"
assert target_request.content == b"chunk-1chunk-2"


@respx.mock
async def test_extension_framework_surfaces_repeated_unauthorized_async(
async_extension_http_client,
):
respx.post(TOKEN_URL).mock(return_value=httpx.Response(200, json={"token": "any-token"}))
target_route = respx.get(ORDERS_URL).mock(
return_value=httpx.Response(401, json={"error": "nope"})
)
client = AsyncHTTPClient(
base_url=API_URL, authentication=ExtensionFrameworkAuthentication(SECRET)
)

with pytest.raises(MPTAPIError) as exc_info: # act
await client.request("GET", "/orders")
await async_extension_http_client.request("GET", "/orders")

assert exc_info.value.status_code == httpx.codes.UNAUTHORIZED
assert target_route.call_count == 2 # original + exactly one retry, then surfaced


@respx.mock
def test_extension_framework_token_error():
def test_extension_framework_token_error(extension_http_client):
respx.post(TOKEN_URL).mock(return_value=httpx.Response(403, json={"error": "forbidden"}))
client = HTTPClient(base_url=API_URL, authentication=ExtensionFrameworkAuthentication(SECRET))

with pytest.raises(MPTError): # act
client.request("GET", "/orders")
extension_http_client.request("GET", "/orders")


def test_extension_framework_requires_configuration():
Expand All @@ -242,12 +301,11 @@ def test_extension_framework_requires_configuration():


@respx.mock
def test_extension_framework_rejects_empty_token():
def test_extension_framework_rejects_empty_token(extension_http_client):
respx.post(TOKEN_URL).mock(return_value=httpx.Response(200, json={"token": None}))
client = HTTPClient(base_url=API_URL, authentication=ExtensionFrameworkAuthentication(SECRET))

with pytest.raises(MPTError): # act
client.request("GET", "/orders")
extension_http_client.request("GET", "/orders")


@respx.mock
Expand All @@ -268,7 +326,7 @@ def test_extension_framework_account_scoped_token_request():


@respx.mock
def test_extension_framework_refreshes_proactively_before_expiry():
def test_extension_framework_refreshes_proactively_before_expiry(extension_http_client):
near_expiry = dt.datetime.now(dt.UTC) + dt.timedelta(seconds=10)
far_expiry = dt.datetime.now(dt.UTC) + dt.timedelta(hours=1)
token_route = respx.post(TOKEN_URL).mock(
Expand All @@ -278,26 +336,24 @@ def test_extension_framework_refreshes_proactively_before_expiry():
]
)
respx.get(ORDERS_URL).mock(return_value=httpx.Response(200, json={"data": []}))
client = HTTPClient(base_url=API_URL, authentication=ExtensionFrameworkAuthentication(SECRET))

client.request("GET", "/orders") # first call caches a near-expiry token
extension_http_client.request("GET", "/orders") # first call caches a near-expiry token

client.request("GET", "/orders") # act: token within leeway -> proactive refresh
extension_http_client.request("GET", "/orders") # act: token within leeway -> refresh

assert token_route.call_count == 2


@respx.mock
def test_extension_framework_reuses_unexpired_token():
def test_extension_framework_reuses_unexpired_token(extension_http_client):
far_expiry = dt.datetime.now(dt.UTC) + dt.timedelta(hours=1)
token_route = respx.post(TOKEN_URL).mock(
return_value=httpx.Response(200, json={"token": _jwt_with_exp(far_expiry)})
)
respx.get(ORDERS_URL).mock(return_value=httpx.Response(200, json={"data": []}))
client = HTTPClient(base_url=API_URL, authentication=ExtensionFrameworkAuthentication(SECRET))

client.request("GET", "/orders") # caches a long-lived token
extension_http_client.request("GET", "/orders") # caches a long-lived token

client.request("GET", "/orders") # act
extension_http_client.request("GET", "/orders") # act

assert token_route.call_count == 1