From 5d6e7f2161cce8c3b06a002f827a39c0dcb69e65 Mon Sep 17 00:00:00 2001 From: Mason Date: Fri, 18 Sep 2026 03:57:36 +0000 Subject: [PATCH] sdk-py: add per-run context to thread stream run.start and run.respond Add optional keyword-only `context` to `run.start` and `run.respond` on both AsyncThreadStream and SyncThreadStream. The context is sent as `params.context` on the `run.start` / `input.respond` commands when non-None and omitted entirely otherwise, preserving the existing wire payload for callers that don't pass it. Co-authored-by: open-swe[bot] --- libs/sdk-py/langgraph_sdk/_async/stream.py | 19 +++++- libs/sdk-py/langgraph_sdk/_sync/stream.py | 19 +++++- .../streaming/test_sync_thread_stream.py | 58 +++++++++++++++++++ .../tests/streaming/test_thread_stream.py | 54 +++++++++++++++++ 4 files changed, 148 insertions(+), 2 deletions(-) diff --git a/libs/sdk-py/langgraph_sdk/_async/stream.py b/libs/sdk-py/langgraph_sdk/_async/stream.py index bc7e68ff6..3a27a9e1c 100644 --- a/libs/sdk-py/langgraph_sdk/_async/stream.py +++ b/libs/sdk-py/langgraph_sdk/_async/stream.py @@ -173,8 +173,18 @@ class RunModule: config: dict[str, Any] | None = None, metadata: dict[str, Any] | None = None, langsmith_tracing: LangSmithTracing | None = None, + context: dict[str, Any] | None = None, ) -> dict[str, Any]: - """Send `run.start` to the server. Returns the result (`{"run_id": ...}`).""" + """Send `run.start` to the server. Returns the result (`{"run_id": ...}`). + + Args: + input: the run input; omitted from the wire payload when None. + config: the run config; omitted when None. + metadata: run metadata; omitted when None. + langsmith_tracing: tracing options; omitted when None. + context: per-run static context; omitted from the wire payload + when None (server applies its default context behavior). + """ params: dict[str, Any] = {"assistant_id": self._owner.assistant_id} if input is not None: params["input"] = input @@ -184,6 +194,8 @@ class RunModule: params["metadata"] = metadata if langsmith_tracing is not None: params["langsmith_tracer"] = langsmith_tracing + if context is not None: + params["context"] = context loop = asyncio.get_running_loop() gate: asyncio.Future[None] = loop.create_future() self._owner._run_start_ready = gate @@ -216,6 +228,7 @@ class RunModule: response: Any, *, interrupt_id: str | None = None, + context: dict[str, Any] | None = None, ) -> dict[str, Any]: """Reply to a server-side interrupt and resume the run. @@ -224,6 +237,8 @@ class RunModule: wire (protocol field name). interrupt_id: optional explicit id. When omitted, requires exactly one outstanding interrupt and uses its id. + context: optional per-run static context for the resumed run; + forwarded with the `input.respond` command when non-None. Raises: RuntimeError: no outstanding interrupts; `interrupt_id` is None but @@ -266,6 +281,8 @@ class RunModule: "namespace": match["namespace"], "response": response, } + if context is not None: + params["context"] = context return await self._owner._send_command("input.respond", params) diff --git a/libs/sdk-py/langgraph_sdk/_sync/stream.py b/libs/sdk-py/langgraph_sdk/_sync/stream.py index fabbff24d..60dedeb06 100644 --- a/libs/sdk-py/langgraph_sdk/_sync/stream.py +++ b/libs/sdk-py/langgraph_sdk/_sync/stream.py @@ -216,8 +216,18 @@ class SyncRunModule: config: dict[str, Any] | None = None, metadata: dict[str, Any] | None = None, langsmith_tracing: LangSmithTracing | None = None, + context: dict[str, Any] | None = None, ) -> dict[str, Any]: - """Send `run.start` to the server. Returns the result (`{"run_id": ...}`).""" + """Send `run.start` to the server. Returns the result (`{"run_id": ...}`). + + Args: + input: the run input; omitted from the wire payload when None. + config: the run config; omitted when None. + metadata: run metadata; omitted when None. + langsmith_tracing: tracing options; omitted when None. + context: per-run static context; omitted from the wire payload + when None (server applies its default context behavior). + """ params: dict[str, Any] = {"assistant_id": self._owner.assistant_id} if input is not None: params["input"] = input @@ -227,6 +237,8 @@ class SyncRunModule: params["metadata"] = metadata if langsmith_tracing is not None: params["langsmith_tracer"] = langsmith_tracing + if context is not None: + params["context"] = context result = self._owner._send_command("run.start", params) self._owner._run_seen = True controller = self._owner._controller @@ -239,6 +251,7 @@ class SyncRunModule: response: Any, *, interrupt_id: str | None = None, + context: dict[str, Any] | None = None, ) -> dict[str, Any]: """Reply to a server-side interrupt and resume the run. @@ -246,6 +259,8 @@ class SyncRunModule: response: the response value forwarded as `params.response` on the wire. interrupt_id: optional explicit id. When omitted, requires exactly one outstanding interrupt. + context: optional per-run static context for the resumed run; + forwarded with the `input.respond` command when non-None. Raises: RuntimeError: no outstanding interrupts; `interrupt_id` is None but @@ -282,6 +297,8 @@ class SyncRunModule: "namespace": match["namespace"], "response": response, } + if context is not None: + params["context"] = context return self._owner._send_command("input.respond", params) diff --git a/libs/sdk-py/tests/streaming/test_sync_thread_stream.py b/libs/sdk-py/tests/streaming/test_sync_thread_stream.py index 5eecd74b8..38063a2fe 100644 --- a/libs/sdk-py/tests/streaming/test_sync_thread_stream.py +++ b/libs/sdk-py/tests/streaming/test_sync_thread_stream.py @@ -439,6 +439,64 @@ def test_sync_run_start_sends_command(): } +def test_sync_run_start_forwards_context(): + fake = SyncFakeServer() + fake.script([lifecycle_completed_event(seq=1)]) + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + thread.run.start(input={"x": 1}, context={"user_id": "u-1"}) + + assert fake.received_commands[0]["params"]["context"] == {"user_id": "u-1"} + + +def test_sync_run_start_omits_context_when_not_provided(): + fake = SyncFakeServer() + fake.script([lifecycle_completed_event(seq=1)]) + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + thread.run.start(input={"x": 1}) + + assert "context" not in fake.received_commands[0]["params"] + + +def test_sync_run_respond_forwards_context(): + fake = SyncFakeServer() + fake.script([lifecycle_completed_event(seq=1)]) + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + thread.run.start(input={}) + thread.interrupts.append( + {"interrupt_id": "i-1", "value": None, "namespace": []} + ) + thread.interrupted = True + thread.run.respond("yes", context={"user_id": "u-1"}) + + command = fake.received_commands[-1] + assert command["method"] == "input.respond" + assert command["params"]["context"] == {"user_id": "u-1"} + + +def test_sync_run_respond_omits_context_when_not_provided(): + fake = SyncFakeServer() + fake.script([lifecycle_completed_event(seq=1)]) + with httpx.Client(transport=fake.transport, base_url="http://test") as raw: + threads = SyncThreadsClient(SyncHttpClient(raw)) + with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + thread.run.start(input={}) + thread.interrupts.append( + {"interrupt_id": "i-1", "value": None, "namespace": []} + ) + thread.interrupted = True + thread.run.respond("yes") + + command = fake.received_commands[-1] + assert command["method"] == "input.respond" + assert "context" not in command["params"] + + def test_sync_events_iterates_raw_events(): fake = SyncFakeServer() diff --git a/libs/sdk-py/tests/streaming/test_thread_stream.py b/libs/sdk-py/tests/streaming/test_thread_stream.py index f683928be..56de789d5 100644 --- a/libs/sdk-py/tests/streaming/test_thread_stream.py +++ b/libs/sdk-py/tests/streaming/test_thread_stream.py @@ -311,6 +311,28 @@ async def test_run_start_forwards_config_metadata_and_langsmith_tracing(): } +async def test_run_start_forwards_context(): + fake = FakeServer() + transport = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={"x": 1}, context={"user_id": "u-1"}) + params = fake.received_commands[0]["params"] + assert params["context"] == {"user_id": "u-1"} + + +async def test_run_start_omits_context_when_not_provided(): + fake = FakeServer() + transport = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={"x": 1}) + params = fake.received_commands[0]["params"] + assert "context" not in params + + async def test_run_start_raises_outside_context_manager(): async with httpx.AsyncClient(base_url="http://test") as raw: @@ -616,6 +638,38 @@ async def test_run_respond_dispatches_input_respond_command(): assert command["params"]["namespace"] == [] +async def test_run_respond_forwards_context(): + fake = FakeServer() + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + thread.interrupts.append( + {"interrupt_id": "i-1", "value": None, "namespace": []} + ) + thread.interrupted = True + await thread.run.respond("yes", context={"user_id": "u-1"}) + params = fake.received_commands[-1]["params"] + assert params["context"] == {"user_id": "u-1"} + + +async def test_run_respond_omits_context_when_not_provided(): + fake = FakeServer() + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + await thread.run.start(input={}) + thread.interrupts.append( + {"interrupt_id": "i-1", "value": None, "namespace": []} + ) + thread.interrupted = True + await thread.run.respond("yes") + params = fake.received_commands[-1]["params"] + assert "context" not in params + + async def test_run_respond_with_explicit_interrupt_id(): fake = FakeServer() asgi = httpx.ASGITransport(app=fake.app)