langchain-ai/deepagents
1
0
Fork 0
mirror of https://github.com/langchain-ai/deepagents.git synced 2026-08-27 20:50:04 -04:00
Code Issues 169 Packages Projects Releases 16 Wiki Activity

fix(code): cancel server-side runs before re-trying interrupted-state writes (#3611)

Browse Source 
When a user hits `Esc` while connected to a remote LangGraph agent, the
local SSE stream closes but the server-side run keeps executing. The
follow-up `aupdate_state` POST then returned HTTP 409 and surfaced
*"Could not save interrupted state (ConflictError). Subsequent turns may
see stale state."* in the chat. Now the client cancels in-flight runs
server-side before retrying the state write, so the partial AI message +
cancellation marker land cleanly and the next turn starts with the right
context.
This commit is contained in:
Mason Daugherty
2026-05-26 21:57:48 -04:00
committed by GitHub
parent f74ebef1f2
commit 7d46357c54
2 changed files with 338 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/code/deepagents_code/remote_client.py
+130 -5
View File
@@ -8,6 +8,7 @@ state snapshots in the server's serialized form.
from __future__ import annotations
import asyncio
import logging
from typing import TYPE_CHECKING, Any
@@ -19,6 +20,15 @@ from deepagents_code._debug import configure_debug_logging
logger = logging.getLogger(__name__)
configure_debug_logging(logger)
_RUN_CANCEL_WAIT_SECONDS = 10.0
"""Per-run cancel wait. Picked so a stuck server-side run can't hang the UI on
Esc for more than ~10s, while leaving room for an actually-cancelling run to
finish its in-flight tool call.
Concurrent cancels keep aggregate wall time bounded by this value regardless of
how many runs are active.
"""
def _require_thread_id(config: dict[str, Any] | None) -> str:
"""Extract and validate that `thread_id` is present in config.
@@ -240,9 +250,18 @@ class RemoteAgent:
) -> None:
"""Update the state of a thread.
Exceptions from the underlying graph (server/network errors) are logged
at DEBUG level and then re-raised so callers can decide how to surface
them (callers typically log at WARNING with a friendlier message).
On HTTP 409 (`ConflictError`) the server still considers the thread
busy — typically because the client cancelled the SSE stream before
the server finished the run. In that case, cancel any pending/running
runs with `wait=True` and retry the state update once. Per-run cancel
waits are bounded by `_RUN_CANCEL_WAIT_SECONDS` and run concurrently,
so callers cannot block indefinitely regardless of how many runs were
active.
Other exceptions from the underlying graph (server/network errors) are
logged at DEBUG level and re-raised so callers can decide how to
surface them (callers typically log at WARNING with a friendlier
message).
Args:
config: Config with `configurable.thread_id`.
@@ -251,16 +270,35 @@ class RemoteAgent:
Raises:
ValueError: If `thread_id` is not present in `config`.
""" # noqa: DOC502 — raised by _require_thread_id
thread_id = _require_thread_id(config)
from langgraph_sdk.errors import ConflictError
thread_id = _require_thread_id(config)
prepared = _prepare_config(config)
graph = self._get_graph()
try:
await graph.aupdate_state(_prepare_config(config), values)
await graph.aupdate_state(prepared, values)
except ConflictError:
pass
except Exception:
logger.debug(
"Failed to update state for thread %s", thread_id, exc_info=True
)
raise
else:
return
await _cancel_active_runs(graph, thread_id)
try:
await graph.aupdate_state(prepared, values)
except Exception:
logger.debug(
"Retry of update_state still failed for thread %s",
thread_id,
exc_info=True,
)
raise
async def aensure_thread(self, config: dict[str, Any]) -> None:
"""Ensure the remote thread record exists before mutating state.
@@ -317,6 +355,93 @@ class RemoteAgent:
return self
async def _cancel_active_runs(graph: Any, thread_id: str) -> None: # noqa: ANN401
"""Cancel pending/running runs on a thread and wait for them to settle.
Best-effort: per-run cancellation failures are logged at DEBUG and
swallowed. Conditions that imply the retry will likely still 409 — failing
to obtain the SDK client, or failing to list runs in every status — are
logged at WARNING so they show up in default logs.
The SDK client is reached via `graph._validate_client()`, a private
attribute on `langgraph.pregel.remote.RemoteGraph`. If upstream renames
or removes it, this helper degrades to no-op and the caller's retry will
re-raise the original `ConflictError`.
Per-run cancels run concurrently and are bounded by
`_RUN_CANCEL_WAIT_SECONDS`, so aggregate wall time stays near that bound
regardless of how many runs are active.
Args:
graph: Underlying `RemoteGraph` instance.
thread_id: Server-side thread identifier.
"""
try:
client = graph._validate_client()
except Exception:
logger.warning(
"Could not obtain SDK client for thread %s; retry will likely "
"still see the conflict",
thread_id,
exc_info=True,
)
return
run_ids: list[str] = []
listed_any = False
for status in ("running", "pending"):
try:
runs = await client.runs.list(thread_id, status=status, limit=10)
except Exception:
logger.debug(
"Failed to list %s runs for thread %s",
status,
thread_id,
exc_info=True,
)
continue
listed_any = True
for run in runs:
run_id = run.get("run_id") if isinstance(run, dict) else None
if run_id:
run_ids.append(run_id)
if not listed_any:
logger.warning(
"Could not list active runs for thread %s; retry will likely "
"still see the conflict",
thread_id,
)
return
if not run_ids:
return
async def _cancel_one(run_id: str) -> None:
try:
await asyncio.wait_for(
client.runs.cancel(thread_id, run_id, wait=True, action="interrupt"),
timeout=_RUN_CANCEL_WAIT_SECONDS,
)
except TimeoutError:
logger.warning(
"Timed out after %.1fs waiting for run %s on thread %s to "
"cancel; retry may still see the conflict",
_RUN_CANCEL_WAIT_SECONDS,
run_id,
thread_id,
)
except Exception:
logger.debug(
"Failed to cancel run %s on thread %s",
run_id,
thread_id,
exc_info=True,
)
await asyncio.gather(*(_cancel_one(rid) for rid in run_ids))
# ---------------------------------------------------------------------------
# Config helpers
# ---------------------------------------------------------------------------
libs/code/tests/unit_tests/test_remote_client.py
+208
View File
@@ -513,6 +513,214 @@ class TestRemoteAgentUpdateState:
uuid.UUID(call_config["configurable"]["thread_id"])
def _conflict_error() -> Exception:
"""Build a `ConflictError` (HTTP 409) for tests."""
import httpx
from langgraph_sdk.errors import ConflictError
request = httpx.Request("POST", "http://localhost:8123/threads/x/state")
response = httpx.Response(409, request=request)
return ConflictError("Thread busy", response=response, body=None)
class TestRemoteAgentUpdateStateConflictRecovery:
"""`aupdate_state` cancels in-flight runs on 409 and retries once."""
def _agent_with_client(
self,
*,
runs_list: AsyncMock,
runs_cancel: AsyncMock,
update_side_effect: list[Any],
) -> tuple[RemoteAgent, MagicMock]:
agent = RemoteAgent(url="http://localhost:8123", graph_name="agent")
mock_graph = MagicMock()
mock_graph.aupdate_state = AsyncMock(side_effect=update_side_effect)
mock_runs = MagicMock()
mock_runs.list = runs_list
mock_runs.cancel = runs_cancel
mock_client = MagicMock()
mock_client.runs = mock_runs
mock_graph._validate_client.return_value = mock_client
agent._graph = mock_graph
return agent, mock_graph
async def test_cancels_all_active_runs_then_retries(self) -> None:
runs_list = AsyncMock(
side_effect=[
[{"run_id": "run-1"}, {"run_id": "run-2"}], # running
[{"run_id": "run-3"}], # pending
]
)
runs_cancel = AsyncMock()
agent, mock_graph = self._agent_with_client(
runs_list=runs_list,
runs_cancel=runs_cancel,
update_side_effect=[_conflict_error(), None],
)
await agent.aupdate_state(_config(), {"messages": []})
assert runs_list.await_count == 2
assert runs_cancel.await_count == 3
cancelled_ids = {call.args[1] for call in runs_cancel.await_args_list}
assert cancelled_ids == {"run-1", "run-2", "run-3"}
# wait=True + action="interrupt" are contractual — `wait` is what
# actually settles the thread before the retry.
for call in runs_cancel.await_args_list:
assert call.kwargs == {"wait": True, "action": "interrupt"}
assert mock_graph.aupdate_state.await_count == 2
async def test_no_active_runs_still_retries(self) -> None:
runs_list = AsyncMock(return_value=[])
runs_cancel = AsyncMock()
agent, mock_graph = self._agent_with_client(
runs_list=runs_list,
runs_cancel=runs_cancel,
update_side_effect=[_conflict_error(), None],
)
await agent.aupdate_state(_config(), {"messages": []})
assert runs_cancel.await_count == 0
assert mock_graph.aupdate_state.await_count == 2
async def test_retry_still_conflict_raises(self) -> None:
runs_list = AsyncMock(return_value=[])
runs_cancel = AsyncMock()
agent, mock_graph = self._agent_with_client(
runs_list=runs_list,
runs_cancel=runs_cancel,
update_side_effect=[_conflict_error(), _conflict_error()],
)
from langgraph_sdk.errors import ConflictError
with pytest.raises(ConflictError):
await agent.aupdate_state(_config(), {"messages": []})
assert mock_graph.aupdate_state.await_count == 2
async def test_cancel_timeout_still_retries(self) -> None:
import asyncio
async def slow_cancel(*_args: Any, **_kwargs: Any) -> None:
await asyncio.sleep(60) # exceeds wait_for timeout
runs_list = AsyncMock(return_value=[{"run_id": "run-1"}])
runs_cancel = AsyncMock(side_effect=slow_cancel)
agent, mock_graph = self._agent_with_client(
runs_list=runs_list,
runs_cancel=runs_cancel,
update_side_effect=[_conflict_error(), None],
)
with patch("deepagents_code.remote_client._RUN_CANCEL_WAIT_SECONDS", 0.01):
await agent.aupdate_state(_config(), {"messages": []})
assert mock_graph.aupdate_state.await_count == 2
async def test_cancel_non_timeout_exception_is_swallowed(self) -> None:
runs_list = AsyncMock(side_effect=[[{"run_id": "run-1"}], []])
runs_cancel = AsyncMock(side_effect=RuntimeError("server hiccup"))
agent, mock_graph = self._agent_with_client(
runs_list=runs_list,
runs_cancel=runs_cancel,
update_side_effect=[_conflict_error(), None],
)
await agent.aupdate_state(_config(), {"messages": []})
assert runs_cancel.await_count == 1
assert mock_graph.aupdate_state.await_count == 2
async def test_runs_list_partial_failure_still_retries(self) -> None:
# First status list raises; second returns runs. Recovery should still
# cancel what it can find and retry.
runs_list = AsyncMock(side_effect=[RuntimeError("boom"), [{"run_id": "run-2"}]])
runs_cancel = AsyncMock()
agent, mock_graph = self._agent_with_client(
runs_list=runs_list,
runs_cancel=runs_cancel,
update_side_effect=[_conflict_error(), None],
)
await agent.aupdate_state(_config(), {"messages": []})
assert runs_list.await_count == 2
assert runs_cancel.await_count == 1
assert runs_cancel.await_args_list[0].args[1] == "run-2"
assert mock_graph.aupdate_state.await_count == 2
async def test_runs_list_total_failure_skips_cancel(self) -> None:
# Both status calls raise. With nothing listed, no cancels happen and
# the retry surfaces the persistent conflict.
runs_list = AsyncMock(side_effect=[RuntimeError("boom"), RuntimeError("boom")])
runs_cancel = AsyncMock()
agent, mock_graph = self._agent_with_client(
runs_list=runs_list,
runs_cancel=runs_cancel,
update_side_effect=[_conflict_error(), _conflict_error()],
)
from langgraph_sdk.errors import ConflictError
with pytest.raises(ConflictError):
await agent.aupdate_state(_config(), {"messages": []})
runs_cancel.assert_not_called()
assert mock_graph.aupdate_state.await_count == 2
async def test_validate_client_raises_skips_cancel_and_retries(self) -> None:
agent = RemoteAgent(url="http://localhost:8123", graph_name="agent")
mock_graph = MagicMock()
mock_graph.aupdate_state = AsyncMock(
side_effect=[_conflict_error(), _conflict_error()]
)
mock_graph._validate_client.side_effect = RuntimeError("no client")
agent._graph = mock_graph
from langgraph_sdk.errors import ConflictError
with pytest.raises(ConflictError):
await agent.aupdate_state(_config(), {"messages": []})
assert mock_graph.aupdate_state.await_count == 2
async def test_runs_without_run_id_are_skipped(self) -> None:
runs_list = AsyncMock(
side_effect=[
# Mixed shapes: missing key, None id, non-dict — all skipped.
[{"run_id": "ok"}, {"run_id": None}, {"status": "running"}, "garbage"],
[],
]
)
runs_cancel = AsyncMock()
agent, mock_graph = self._agent_with_client(
runs_list=runs_list,
runs_cancel=runs_cancel,
update_side_effect=[_conflict_error(), None],
)
await agent.aupdate_state(_config(), {"messages": []})
assert runs_cancel.await_count == 1
assert runs_cancel.await_args_list[0].args[1] == "ok"
assert mock_graph.aupdate_state.await_count == 2
async def test_non_conflict_exception_does_not_retry(self) -> None:
runs_list = AsyncMock()
runs_cancel = AsyncMock()
agent, mock_graph = self._agent_with_client(
runs_list=runs_list,
runs_cancel=runs_cancel,
update_side_effect=[ConnectionError("down")],
)
with pytest.raises(ConnectionError, match="down"):
await agent.aupdate_state(_config(), {"messages": []})
assert mock_graph.aupdate_state.await_count == 1
runs_list.assert_not_called()
runs_cancel.assert_not_called()
class TestRemoteAgentEnsureThread:
"""Verify remote thread registration before state writes."""