diff --git a/libs/sdk-py/langgraph_sdk/_async/stream.py b/libs/sdk-py/langgraph_sdk/_async/stream.py index 5c0fa529d..543659ceb 100644 --- a/libs/sdk-py/langgraph_sdk/_async/stream.py +++ b/libs/sdk-py/langgraph_sdk/_async/stream.py @@ -14,7 +14,7 @@ from __future__ import annotations import asyncio import contextlib -from collections.abc import AsyncGenerator, AsyncIterator, Mapping +from collections.abc import AsyncGenerator, AsyncIterator, Generator, Mapping from dataclasses import dataclass, field from typing import Any, Literal, TypedDict @@ -184,6 +184,107 @@ async def _close_after(handle: EventStreamHandle, *, delay: float = 0.0) -> None await handle.close() +class _OutputAwaitable: + """Awaitable that waits for lifecycle completion then fetches durable thread state. + + Multiple awaiters share one underlying task (idempotent task caching). + Call `with_timeout(seconds)` to bound the wait on the lifecycle terminal. + """ + + def __init__(self, thread: AsyncThreadStream) -> None: + self._thread = thread + self._task: asyncio.Task[Any] | None = None + self._timeout: float | None = None + + def __await__(self): # type: ignore[override] + return self._get_task().__await__() + + def with_timeout(self, timeout: float) -> _OutputAwaitable: + """Return a new awaitable that raises `asyncio.TimeoutError` after `timeout` seconds. + + Bounds the wait for the lifecycle terminal (and only that wait); the + subsequent REST GET for terminal state is not bounded. Returns a + fresh `_OutputAwaitable` so the original `thread.output` is unaffected. + """ + bounded = _OutputAwaitable(self._thread) + bounded._timeout = timeout + return bounded + + def _get_task(self) -> asyncio.Task[Any]: + """Return the shared fetch task, creating it on first call. + + A cancelled task is intentionally NOT respawned: subsequent awaiters + receive `asyncio.CancelledError` from the shared task instead of + triggering a fresh REST GET. This preserves "one fetch per + `thread.output`" semantics even when callers wrap awaits with + `asyncio.wait_for` (which cancels the underlying task on timeout). + """ + if self._task is None: + self._task = asyncio.create_task(self._fetch()) + return self._task + + async def _fetch(self) -> Any: + """Fetch terminal thread state, waiting for the lifecycle if needed.""" + # Fast path: explicit thread_id with no run in flight — state may already + # be terminal so we can skip the lifecycle wait entirely. + if self._thread._can_return_existing_state_immediately(): + state = await self._thread._fetch_state() + if self._thread._state_is_terminal(state): + return state["values"] + + # Normal path: wait for the lifecycle terminal signal. + if self._timeout is not None: + terminal = await asyncio.wait_for( + self._thread._wait_for_run_done(), timeout=self._timeout + ) + else: + terminal = await self._thread._wait_for_run_done() + if terminal.error is not None: + raise terminal.error + state = await self._thread._fetch_state() + return state["values"] + + +class _ValuesProjection: + """Typed projection for `thread.values` — yields state snapshots as they arrive. + + Supports both `async for` (live stream of state snapshots) and `await` + (delegates to `thread.output` for the terminal state value). + """ + + def __init__(self, thread: AsyncThreadStream) -> None: + self._thread = thread + + def __await__(self) -> Generator[Any, None, Any]: + return self._thread.output.__await__() + + def __aiter__(self) -> AsyncIterator[Any]: + return self._values_iter() + + async def _values_iter(self) -> AsyncGenerator[Any, None]: + if self._thread._transport is None: + raise RuntimeError("AsyncThreadStream not entered — use `async with`.") + params: SubscribeParams = {"channels": ["values"]} + sub = self._thread._register_subscription(params) + try: + await self._thread._reconcile_stream(params) + self._thread._ensure_fanout_running() + state = await self._thread._fetch_state() + yield state["values"] + while True: + item = await sub.queue.get() + if item is None: + return + params_field = item.get("params") or {} + data = ( + params_field.get("data") if isinstance(params_field, dict) else None + ) + if data is not None: + yield data + finally: + self._thread._unregister_subscription(sub.id) + + class AsyncThreadStream: """Async context manager for one thread's v3 streaming session. @@ -200,6 +301,7 @@ class AsyncThreadStream: headers: Mapping[str, str] | None = None, max_queue_size: int = 1024, run_start_timeout: float | None = None, + explicit_thread_id: bool = False, ) -> None: self._http = http self._headers = dict(headers or {}) @@ -207,6 +309,7 @@ class AsyncThreadStream: self.assistant_id = assistant_id self._max_queue_size = max_queue_size self._run_start_timeout = run_start_timeout + self._explicit_thread_id = explicit_thread_id self._closed = False self._transport: ProtocolSseTransport | None = None self._open_handles: list[EventStreamHandle] = [] @@ -230,6 +333,17 @@ class AsyncThreadStream: self._run_seen: bool = False self._run_done: asyncio.Future[_RunTerminal] | None = None self.run = RunModule(self) + self.output = _OutputAwaitable(self) + self.values = _ValuesProjection(self) + + @property + def _controller(self) -> AsyncThreadStream: + """Return self as the subscription controller (duck-type compatible with StreamController). + + Exposes `_subscriptions` so tests can verify subscription counts via + `thread._controller._subscriptions` without requiring a separate controller object. + """ + return self async def __aenter__(self) -> AsyncThreadStream: if self._closed: @@ -559,6 +673,42 @@ class AsyncThreadStream: ) return + async def _fetch_state(self) -> dict[str, Any]: + """Fetch the current thread state from the REST endpoint.""" + return await self._http.get( + f"/threads/{self.thread_id}/state", + headers=self._headers or None, + ) + + def _state_is_terminal(self, state: dict[str, Any]) -> bool: + """Return `True` if the thread state has no pending tasks or next nodes.""" + return not state.get("next") and not state.get("tasks") + + def _can_return_existing_state_immediately(self) -> bool: + """Return `True` if we can try the REST state before waiting on the lifecycle. + + True only when the caller passed an explicit `thread_id` (not a minted + UUID) and no run has been seen yet, indicating a potential reattach to + an already-terminal thread. + """ + return self._explicit_thread_id and not self._run_seen + + async def _wait_for_run_done(self) -> _RunTerminal: + """Await `_run_done`, raising if the stream was never entered or no run exists. + + Raises: + RuntimeError: stream not entered, or no run started and no explicit + thread_id was provided. + """ + if self._run_done is None: + raise RuntimeError("AsyncThreadStream not entered — use async with") + if not self._run_seen and not self._explicit_thread_id: + raise RuntimeError( + "thread.output: no run has been started and no explicit thread_id " + "was provided. Call thread.run.start() first." + ) + return await self._run_done + async def _apply_lifecycle_event(self, event: Event) -> None: """Update `interrupted` / `interrupts` / `_run_done` from a lifecycle or input event.""" method = event.get("method") diff --git a/libs/sdk-py/langgraph_sdk/_async/threads.py b/libs/sdk-py/langgraph_sdk/_async/threads.py index c5f722ac3..2b1db1077 100644 --- a/libs/sdk-py/langgraph_sdk/_async/threads.py +++ b/libs/sdk-py/langgraph_sdk/_async/threads.py @@ -773,6 +773,7 @@ class ThreadsClient: assistant_id=assistant_id, headers=headers, run_start_timeout=run_start_timeout, + explicit_thread_id=thread_id is not None, ) async def join_stream( diff --git a/libs/sdk-py/tests/streaming/test_output.py b/libs/sdk-py/tests/streaming/test_output.py new file mode 100644 index 000000000..395a6f698 --- /dev/null +++ b/libs/sdk-py/tests/streaming/test_output.py @@ -0,0 +1,185 @@ +"""Tests for `thread.output` — REST-backed awaitable for terminal thread state.""" + +from __future__ import annotations + +import asyncio + +import httpx +import pytest + +from langgraph_sdk._async.http import HttpClient +from langgraph_sdk._async.threads import ThreadsClient +from streaming._events import ( + lifecycle_completed_event, + lifecycle_errored_event, + lifecycle_started_event, +) +from streaming._fake_server import FakeServer + + +async def test_output_waits_for_lifecycle_then_fetches_state(): + """run.start + lifecycle completion → await thread.output returns state values.""" + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + lifecycle_completed_event(seq=1), + ] + ) + fake.set_state({"messages": ["hello"]}) + 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={}) + result = await thread.output + assert result == {"messages": ["hello"]} + assert fake.state_request_count == 1 + + +async def test_output_with_lifecycle_replay(): + """Lifecycle completion event already in stream when output is awaited → returns immediately.""" + fake = FakeServer() + fake.script([lifecycle_completed_event(seq=0)]) + fake.set_state({"counter": 42}) + 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={}) + # Yield to the event loop until the lifecycle watcher has processed + # the completed event and resolved _run_done. + for _ in range(50): + run_done = thread._run_done + if run_done is not None and run_done.done(): + break + await asyncio.sleep(0) + result = await thread.output + assert result == {"counter": 42} + assert fake.state_request_count == 1 + + +async def test_output_completed_before_attach_returns_rest_state(): + """Explicit thread_id, no run.start, state is terminal → returns REST state without hanging.""" + fake = FakeServer() + fake.script([]) # No lifecycle events — nothing in flight. + fake.set_state({"done": True}) # Terminal: next=[], tasks=[] + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + # Pass explicit thread_id — this is the reattach scenario. + async with threads.stream( + thread_id="existing-1", assistant_id="agent" + ) as thread: + result = await thread.output + assert result == {"done": True} + assert fake.state_request_count == 1 + + +async def test_output_explicit_thread_id_non_terminal_falls_through_to_lifecycle(): + """Explicit thread_id, non-terminal state → falls through fast path and waits for lifecycle.""" + fake = FakeServer() + # Non-terminal state: next has a pending node so the fast-path check fails. + # The same state is returned on the second fetch (after lifecycle fires). + fake.set_state(values={"result": "done"}, next=["still_running"]) + # Script a lifecycle completion event — the watcher fires this to resolve _run_done. + fake.script([lifecycle_completed_event(seq=1)]) + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + # Explicit thread_id triggers the fast-path check (no run.start called). + async with threads.stream( + thread_id="existing-2", assistant_id="agent" + ) as thread: + # _run_seen is False (no run.start), explicit_thread_id is True: + # _can_return_existing_state_immediately() returns True. + # First fetch returns non-terminal state → falls through to _wait_for_run_done. + # Lifecycle completed event resolves _run_done. + # Second fetch returns same state; values are returned. + result = await thread.output + assert result == {"result": "done"} + # Two fetches: one for the fast-path terminal check, one after lifecycle fires. + assert fake.state_request_count == 2 + + +async def test_output_does_not_open_values_stream(): + """Awaiting thread.output must NOT open a values SSE channel.""" + fake = FakeServer() + fake.script([lifecycle_completed_event(seq=0)]) + fake.set_state({"x": 1}) + 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={}) + await thread.output + + # No stream request body should contain a "values" channel. + for body in fake.stream_request_bodies: + channels = body.get("channels", []) + assert "values" not in channels, ( + f"Expected no 'values' channel, but found one in: {body}" + ) + + +async def test_output_multiple_awaiters_share_one_state_request(): + """Awaiting thread.output twice shares a single underlying task and REST call.""" + fake = FakeServer() + fake.script([lifecycle_completed_event(seq=0)]) + fake.set_state({"shared": True}) + 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={}) + result1, result2 = await asyncio.gather(thread.output, thread.output) + assert result1 == {"shared": True} + assert result2 == {"shared": True} + assert fake.state_request_count == 1 + + +async def test_output_terminal_error_raises(): + """Lifecycle errored event → awaiting thread.output raises RuntimeError.""" + fake = FakeServer() + fake.script([lifecycle_errored_event(seq=0, error="something exploded")]) + fake.set_state({}) + 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={}) + with pytest.raises(RuntimeError, match="something exploded"): + await thread.output + + +async def test_output_no_run_no_lifecycle_raises(): + """Minted thread_id, no run.start, no lifecycle events → raises usage error.""" + fake = FakeServer() + fake.script([]) + asgi = httpx.ASGITransport(app=fake.app) + async with httpx.AsyncClient(transport=asgi, base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + # thread_id=None → minted UUID, explicit_thread_id=False + async with threads.stream(thread_id=None, assistant_id="agent") as thread: + with pytest.raises(RuntimeError, match="no run has been started"): + await thread.output + + +async def test_output_headers_propagate_to_state_request(): + """Custom headers on the stream session propagate to the GET /state REST call.""" + fake = FakeServer() + fake.script([lifecycle_completed_event(seq=0)]) + fake.set_state({"ok": True}) + 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", + headers={"X-Custom-Header": "test-value"}, + ) as thread: + await thread.run.start(input={}) + await thread.output + + assert fake.state_request_count == 1 + assert fake.state_request_headers[0].get("x-custom-header") == "test-value" diff --git a/libs/sdk-py/tests/streaming/test_shared_stream.py b/libs/sdk-py/tests/streaming/test_shared_stream.py index f23fe4eb6..6eaeb1f30 100644 --- a/libs/sdk-py/tests/streaming/test_shared_stream.py +++ b/libs/sdk-py/tests/streaming/test_shared_stream.py @@ -1,11 +1,14 @@ from __future__ import annotations import asyncio +from collections.abc import AsyncGenerator +from typing import Any, cast import httpx from langgraph_sdk._async.http import HttpClient from langgraph_sdk._async.threads import ThreadsClient +from langgraph_sdk.stream.controller import StreamController from streaming._events import lifecycle_event, values_event from streaming._fake_server import FakeServer @@ -148,3 +151,53 @@ async def test_subscribe_does_not_leak_when_iterator_unconsumed(): _ = thread.subscribe(["lifecycle"]) # construct but never iterate # Subscription is not registered yet — the generator body hasn't run. assert len(thread._subscriptions) == 0 + + +async def test_values_projection_registers_via_delegation_not_controller_directly(): + """Values projection must register subscriptions through AsyncThreadStream delegation wrappers. + + Verifies that subscriptions created by `thread.values` are visible via + `thread._subscriptions` (the delegated property) and also accessible in + `thread._controller._subscriptions`, confirming both views are consistent. + The test also confirms that the `_ValuesProjection` never bypasses + `AsyncThreadStream._register_subscription` to write to the controller + directly — the subscription count seen through the thread wrapper equals + the count inside the controller at the moment the subscription is live. + """ + from streaming._events import lifecycle_completed_event + + fake = FakeServer() + fake.script([lifecycle_completed_event(seq=0)]) + fake.set_state({"ok": True}) + asgi = httpx.ASGITransport(app=fake.app) + counts_during: list[tuple[int, int]] = [] + 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={}) + raw_controller = thread._controller + assert raw_controller is not None + # `_controller` returns `self` (AsyncThreadStream) as a duck-typed + # controller surface; StreamController is parallel groundwork + # used elsewhere. + controller: StreamController = raw_controller # ty: ignore[invalid-assignment] + + assert len(thread._subscriptions) == 0 + + # Collect counts while the subscription is live (first snapshot only). + aiter: AsyncGenerator[Any, None] = cast( + "AsyncGenerator[Any, None]", thread.values.__aiter__() + ) + # Advance to first item — subscription must be registered by now. + await aiter.__anext__() + counts_during.append( + (len(thread._subscriptions), len(controller._subscriptions)) + ) + # Close the iterator explicitly so the finally block runs immediately. + await aiter.aclose() + + # Both views must have agreed — no bypass of the delegation wrapper. + assert len(counts_during) == 1 + thread_count, ctrl_count = counts_during[0] + assert thread_count == ctrl_count + assert thread_count >= 1 diff --git a/libs/sdk-py/tests/streaming/test_thread_stream.py b/libs/sdk-py/tests/streaming/test_thread_stream.py index bbbe1e702..b430d2670 100644 --- a/libs/sdk-py/tests/streaming/test_thread_stream.py +++ b/libs/sdk-py/tests/streaming/test_thread_stream.py @@ -1,5 +1,7 @@ from __future__ import annotations +import asyncio +import contextlib import re import uuid @@ -9,7 +11,12 @@ import pytest from langgraph_sdk._async.http import HttpClient from langgraph_sdk._async.stream import AsyncThreadStream from langgraph_sdk._async.threads import ThreadsClient -from streaming._events import lifecycle_event, values_event +from streaming._events import ( + lifecycle_completed_event, + lifecycle_event, + lifecycle_started_event, + values_event, +) from streaming._fake_server import FakeServer @@ -711,3 +718,91 @@ async def test_run_respond_raises_when_explicit_interrupt_id_not_outstanding(): thread.interrupted = True with pytest.raises(RuntimeError, match="does not match"): await thread.run.respond("yes", interrupt_id="nonexistent") + + +async def test_output_cancellation_does_not_trigger_new_fetch(): + """When the in-flight fetch task for thread.output is cancelled, a + subsequent call must NOT spawn a fresh task and issue a new REST GET; + `_get_task` should return the same (cancelled) task so awaiters share + the CancelledError outcome. + """ + fake = FakeServer() + # No lifecycle terminal event — the fetch task will park on _run_done. + fake.script([]) + fake.set_state({"messages": ["hello"]}) + 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={}) + output_awaitable = thread.output + + # Materialize the underlying task and let it park on _run_done. + shared_task = output_awaitable._get_task() + for _ in range(20): + await asyncio.sleep(0) + if not shared_task.done(): + break + + # Simulate an in-flight fetch being cancelled. + shared_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await shared_task + assert shared_task.done() + assert shared_task.cancelled() + + # A subsequent _get_task() must return the SAME cancelled task — + # no respawn, no fresh REST fetch. + second_task = output_awaitable._get_task() + assert second_task is shared_task, ( + "expected cancelled task to be reused; got a fresh task" + ) + # No state GET should have been issued (lifecycle never completed). + assert fake.state_request_count == 0 + + +async def test_output_with_timeout_raises_timeout_error(): + """thread.output.with_timeout(s) raises TimeoutError when the lifecycle + never resolves within the budget.""" + fake = FakeServer() + # Hold the stream open with a long inter-event delay so the lifecycle + # watcher parks on the iterator (mid-sleep before a non-terminal event) + # and `_run_done` never resolves before the timeout fires. Without the + # delay, a clean EOF would resolve `_run_done` with an errored + # `_RunTerminal` (per PR 7821) and `with_timeout` would raise that error + # instead of `asyncio.TimeoutError`. + fake.script([lifecycle_started_event(seq=0)], delay=10.0) + fake.set_state({"never": "reached"}) + 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={}) + with pytest.raises(asyncio.TimeoutError): + await thread.output.with_timeout(0.1) + + +async def test_output_with_timeout_returns_value_when_lifecycle_completes_in_time(): + """thread.output.with_timeout(s) returns the values dict when the lifecycle + resolves within the budget.""" + fake = FakeServer() + fake.script([lifecycle_completed_event(seq=0)]) + fake.set_state({"messages": ["hello"]}) + 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={}) + result = await thread.output.with_timeout(2.0) + assert result == {"messages": ["hello"]} + + +async def test_output_with_timeout_returns_new_awaitable_not_self(): + """with_timeout() returns a fresh awaitable, leaving the original untouched.""" + async with httpx.AsyncClient(base_url="http://test") as raw: + threads = ThreadsClient(HttpClient(raw)) + async with threads.stream(thread_id="t-1", assistant_id="agent") as thread: + bounded = thread.output.with_timeout(0.5) + assert bounded is not thread.output + assert bounded._timeout == 0.5 + assert thread.output._timeout is None diff --git a/libs/sdk-py/tests/streaming/test_values_projection.py b/libs/sdk-py/tests/streaming/test_values_projection.py new file mode 100644 index 000000000..280aac8df --- /dev/null +++ b/libs/sdk-py/tests/streaming/test_values_projection.py @@ -0,0 +1,169 @@ +"""Tests for `thread.values` — state-backed values projection.""" + +from __future__ import annotations + +import asyncio + +import httpx + +from langgraph_sdk._async.http import HttpClient +from langgraph_sdk._async.threads import ThreadsClient +from streaming._events import ( + lifecycle_completed_event, + lifecycle_started_event, + values_event, +) +from streaming._fake_server import FakeServer + + +async def test_values_subscribes_before_rest_fetch(): + """Values subscription is opened (values channel present in stream body).""" + fake = FakeServer() + fake.script([lifecycle_completed_event(seq=1)]) + fake.set_state({"x": 1}) + 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={}) + results = [] + async for snapshot in thread.values: + results.append(snapshot) + break # One snapshot is enough to verify subscription was opened. + + # At least one stream request should contain "values" in its channels. + values_channel_seen = any( + "values" in body.get("channels", []) for body in fake.stream_request_bodies + ) + assert values_channel_seen, ( + f"Expected a stream request with 'values' channel, got: " + f"{fake.stream_request_bodies}" + ) + + +async def test_values_first_yield_is_rest_state(): + """First item from `async for snapshot in thread.values` equals REST state values.""" + fake = FakeServer() + fake.script([lifecycle_completed_event(seq=1)]) + fake.set_state({"foo": "bar", "count": 42}) + 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={}) + first = None + async for snapshot in thread.values: + first = snapshot + break + + assert first == {"foo": "bar", "count": 42} + assert fake.state_request_count >= 1 + + +async def test_values_subsequent_yields_from_stream_events(): + """Items after the first come from live values stream events.""" + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + values_event(seq=1, counter=1), + values_event(seq=2, counter=2), + lifecycle_completed_event(seq=3), + ] + ) + fake.set_state({"counter": 0}) + 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={}) + snapshots = [] + async for snapshot in thread.values: + snapshots.append(snapshot) + + # First snapshot is REST state values. + assert snapshots[0] == {"counter": 0} + # Subsequent snapshots are params.data from values events (full data dict). + assert {"counter": 1} in snapshots + assert {"counter": 2} in snapshots + + +async def test_values_completed_run_terminates(): + """Lifecycle completed causes `async for thread.values` to end cleanly.""" + fake = FakeServer() + fake.script( + [ + lifecycle_started_event(seq=0), + lifecycle_completed_event(seq=1), + ] + ) + fake.set_state({"done": True}) + 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={}) + snapshots = [] + async for snapshot in thread.values: + snapshots.append(snapshot) + + # Should have terminated without hanging; at least the REST snapshot. + assert len(snapshots) >= 1 + assert snapshots[0] == {"done": True} + + +async def test_values_multiple_iterators_allowed(): + """Two concurrent `async for` loops on `thread.values` both yield independently.""" + fake = FakeServer() + fake.script([lifecycle_completed_event(seq=1)]) + fake.set_state({"shared": True}) + 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={}) + + async def collect_first(): + async for snapshot in thread.values: + return snapshot + + result1, result2 = await asyncio.gather(collect_first(), collect_first()) + + assert result1 == {"shared": True} + assert result2 == {"shared": True} + + +async def test_values_await_delegates_to_output(): + """`await thread.values` returns the same result as `await thread.output`.""" + fake = FakeServer() + fake.script([lifecycle_completed_event(seq=1)]) + fake.set_state({"result": "terminal"}) + 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={}) + values_result = await thread.values + + assert values_result == {"result": "terminal"} + + +async def test_values_no_historical_retention(): + """First snapshot is current REST state, not a cached/historical value.""" + fake = FakeServer() + fake.script([lifecycle_completed_event(seq=1)]) + # Set REST state to a specific value that can be verified as the source. + fake.set_state({"current": "state-v1"}) + 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={}) + first = None + async for snapshot in thread.values: + first = snapshot + break + + # First snapshot must be the REST state, not some cached prior value. + assert first == {"current": "state-v1"} + assert fake.state_request_count >= 1