feat: sdk-py: Reconnect to long-lived responses on wait/join/cancel endpoints (#6168)

- When connection is dropped while waiting, reconnect up to 5 times if a
Location header is present
This commit is contained in:
Nuno Campos
2025-09-20 19:44:07 +01:00
committed by GitHub
parent 11834512db
commit 5ba02d5b46
2 changed files with 141 additions and 15 deletions
+1 -1
View File
@@ -728,7 +728,7 @@ def is_studio_user(
return (
isinstance(user, types.StudioUser)
or isinstance(user, dict)
and user.get("kind") == "StudioUser"
and user.get("kind") == "StudioUser" # ty: ignore[invalid-argument-type]
)
+140 -14
View File
@@ -437,6 +437,54 @@ class HttpClient:
logger.error(f"Error from langgraph-api: {body}", exc_info=e)
raise e
async def request_reconnect(
self,
path: str,
method: str,
*,
json: dict[str, Any] | None = None,
params: QueryParamTypes | None = None,
headers: Mapping[str, str] | None = None,
on_response: Callable[[httpx.Response], None] | None = None,
reconnect_limit: int = 5,
) -> Any:
"""Send a request that automatically reconnects to Location header."""
request_headers, content = await _aencode_json(json)
if headers:
request_headers.update(headers)
async with self.client.stream(
method, path, headers=request_headers, content=content, params=params
) as r:
if on_response:
on_response(r)
try:
r.raise_for_status()
except httpx.HTTPStatusError as e:
body = (await r.aread()).decode()
if sys.version_info >= (3, 11):
e.add_note(body)
else:
logger.error(f"Error from langgraph-api: {body}", exc_info=e)
raise e
loc = r.headers.get("location")
if reconnect_limit <= 0 or not loc:
return await _adecode_json(r)
try:
return await _adecode_json(r)
except httpx.HTTPError:
warnings.warn(
f"Request failed, attempting reconnect to Location: {loc}",
stacklevel=2,
)
await r.aclose()
return await self.request_reconnect(
loc,
"GET",
headers=request_headers,
# don't pass on_response so it's only called once
reconnect_limit=reconnect_limit - 1,
)
async def stream(
self,
path: str,
@@ -2533,8 +2581,9 @@ class RunsClient:
if on_run_created and (metadata := _get_run_metadata_from_response(res)):
on_run_created(metadata)
response = await self.http.post(
response = await self.http.request_reconnect(
endpoint,
"POST",
json={k: v for k, v in payload.items() if v is not None},
params=params,
headers=headers,
@@ -2679,12 +2728,20 @@ class RunsClient:
}
if params:
query_params.update(params)
return await self.http.post(
f"/threads/{thread_id}/runs/{run_id}/cancel",
json=None,
params=query_params,
headers=headers,
)
if wait:
return await self.http.request_reconnect(
f"/threads/{thread_id}/runs/{run_id}/cancel",
"POST",
params=query_params,
headers=headers,
)
else:
return await self.http.post(
f"/threads/{thread_id}/runs/{run_id}/cancel",
json=None,
params=query_params,
headers=headers,
)
async def join(
self,
@@ -2716,8 +2773,11 @@ class RunsClient:
```
""" # noqa: E501
return await self.http.get(
f"/threads/{thread_id}/runs/{run_id}/join", headers=headers, params=params
return await self.http.request_reconnect(
f"/threads/{thread_id}/runs/{run_id}/join",
"GET",
headers=headers,
params=params,
)
def join_stream(
@@ -3689,6 +3749,54 @@ class SyncHttpClient:
logger.error(f"Error from langgraph-api: {body}", exc_info=e)
raise e
def request_reconnect(
self,
path: str,
method: str,
*,
json: dict[str, Any] | None = None,
params: QueryParamTypes | None = None,
headers: Mapping[str, str] | None = None,
on_response: Callable[[httpx.Response], None] | None = None,
reconnect_limit: int = 5,
) -> Any:
"""Send a request that automatically reconnects to Location header."""
request_headers, content = _encode_json(json)
if headers:
request_headers.update(headers)
with self.client.stream(
method, path, headers=request_headers, content=content, params=params
) as r:
if on_response:
on_response(r)
try:
r.raise_for_status()
except httpx.HTTPStatusError as e:
body = r.read().decode()
if sys.version_info >= (3, 11):
e.add_note(body)
else:
logger.error(f"Error from langgraph-api: {body}", exc_info=e)
raise e
loc = r.headers.get("location")
if reconnect_limit <= 0 or not loc:
return _decode_json(r)
try:
return _decode_json(r)
except httpx.HTTPError:
warnings.warn(
f"Request failed, attempting reconnect to Location: {loc}",
stacklevel=2,
)
r.close()
return self.request_reconnect(
loc,
"GET",
headers=request_headers,
# don't pass on_response so it's only called once
reconnect_limit=reconnect_limit - 1,
)
def stream(
self,
path: str,
@@ -5754,8 +5862,9 @@ class SyncRunsClient:
endpoint = (
f"/threads/{thread_id}/runs/wait" if thread_id is not None else "/runs/wait"
)
return self.http.post(
return self.http.request_reconnect(
endpoint,
"POST",
json={k: v for k, v in payload.items() if v is not None},
params=params,
headers=headers,
@@ -5878,11 +5987,25 @@ class SyncRunsClient:
```
""" # noqa: E501
query_params = {
"wait": 1 if wait else 0,
"action": action,
}
if params:
query_params.update(params)
if wait:
return self.http.request_reconnect(
f"/threads/{thread_id}/runs/{run_id}/cancel",
"POST",
json=None,
params=query_params,
headers=headers,
)
return self.http.post(
f"/threads/{thread_id}/runs/{run_id}/cancel?wait={1 if wait else 0}&action={action}",
f"/threads/{thread_id}/runs/{run_id}/cancel",
json=None,
params=query_params,
headers=headers,
params=params,
)
def join(
@@ -5915,8 +6038,11 @@ class SyncRunsClient:
```
""" # noqa: E501
return self.http.get(
f"/threads/{thread_id}/runs/{run_id}/join", headers=headers, params=params
return self.http.request_reconnect(
f"/threads/{thread_id}/runs/{run_id}/join",
"GET",
headers=headers,
params=params,
)
def join_stream(