fix(langgraph): return an awaited @task's value while streaming

With stream=True, arun_with_retry drains the task's astream() and returned
nothing, so a @task awaited inside a node got None whenever the node ran in
streaming mode: under astream_events (v1 and v2) and anything built on it.
Return the value the task wrote instead, on both stream paths.

Same idea as #9295 by AaronG18 and the lookup amishaguptaberk suggested.

libs/langgraph/langgraph/pregel/_retry.py

@@ -34,6 +34,7 @@ from langgraph._internal._constants import (
     CONFIG_KEY_THREAD_ID,
     CONFIG_KEY_TIMED_ATTEMPT_OBSERVER,
     NS_SEP,
+    RETURN,
 )
 from langgraph._internal._runnable import create_task_in_config_context
 from langgraph._internal._timeout import sync_timeout_unsupported
@@ -733,15 +734,14 @@ async def arun_with_retry(
                 if stream:
                     async for _ in task.proc.astream(task.input, config):
                         pass
-                    break
+                    return _return_value(task)
                 return await task.proc.ainvoke(task.input, config)
             result = await _arun_with_timeout(
                 task, config, resolved_timeout, attempt_ctx, stream=stream
             )
             _finish_timed_attempt(config, attempt_ctx)
             if stream:
-                # if successful, end
-                break
+                return _return_value(task)
             return result
         except ParentCommand as exc:
             ns: str = config[CONF][CONFIG_KEY_CHECKPOINT_NS]
@@ -831,6 +831,12 @@ async def arun_with_retry(
             config = patch_configurable(config, {CONFIG_KEY_RESUMING: True})
 
 
+def _return_value(task: PregelExecutableTask) -> Any:
+    # Streaming drains the task instead of returning its result, so a caller
+    # awaiting a `@task` needs the value the task wrote.
+    return next((value for channel, value in task.writes if channel == RETURN), None)
+
+
 def _should_retry_on(retry_policy: RetryPolicy, exc: Exception) -> bool:
     """Check if the given exception should be retried based on the retry policy."""
     if isinstance(retry_policy.retry_on, Sequence):

libs/langgraph/tests/test_pregel_async.py

@@ -7495,6 +7495,39 @@ async def test_entrypoint_from_async_generator() -> None:
             yield "b"
 
 
+@NEEDS_CONTEXTVARS
+@pytest.mark.parametrize("version", ["v1", "v2"])
+async def test_awaited_task_returns_its_value_under_astream_events(
+    version: Literal["v1", "v2"],
+) -> None:
+    class State(TypedDict):
+        doubled: int
+        zeroed: int
+
+    @task
+    async def double(x: int) -> int:
+        return x * 2
+
+    @task
+    async def zero() -> int:
+        return 0
+
+    async def node(state: State) -> State:
+        return {"doubled": await double(21), "zeroed": await zero()}
+
+    graph = StateGraph(State).add_node("node", node).add_edge(START, "node").compile()
+
+    outputs = [
+        event["data"]["output"]
+        async for event in graph.astream_events(
+            {"doubled": 0, "zeroed": 1}, version=version
+        )
+        if event["event"] == "on_chain_end" and event["name"] == "node"
+    ]
+
+    assert outputs == [{"doubled": 42, "zeroed": 0}]
+
+
 @NEEDS_CONTEXTVARS
 async def test_named_tasks_functional() -> None:
     class Foo:

libs/sdk-py/langgraph_sdk/_async/http.py

@@ -217,7 +217,7 @@ class HttpClient:
             current_headers = dict(
                 request_headers if reconnect_path is None else reconnect_headers
             )
-            if last_event_id:
+            if last_event_id is not None:
                 current_headers["Last-Event-ID"] = last_event_id
 
             current_method = method if reconnect_path is None else "GET"

libs/sdk-py/langgraph_sdk/_sync/http.py

@@ -219,7 +219,7 @@ class SyncHttpClient:
             current_headers = dict(
                 request_headers if reconnect_path is None else reconnect_headers
             )
-            if last_event_id:
+            if last_event_id is not None:
                 current_headers["Last-Event-ID"] = last_event_id
 
             current_method = method if reconnect_path is None else "GET"

libs/sdk-py/langgraph_sdk/sse.py

@@ -79,27 +79,18 @@ class SSEDecoder:
     def __init__(self) -> None:
         self._event = ""
         self._data = bytearray()
-        self._last_event_id: str | None = None
+        self._last_event_id = ""
         self._retry: int | None = None
-        self._at_stream_start = True
 
     @property
     def last_event_id(self) -> str | None:
-        """Return the last event identifier that was seen.
+        """Return the last event identifier that was seen."""
 
-        `None` until the stream sends an `id` field, and `""` after an empty
-        one, which resets the reconnect cursor.
-        """
-
-        return self._last_event_id
+        return self._last_event_id or None
 
     def decode(self, line: bytes) -> StreamPart | None:
         # See: https://html.spec.whatwg.org/multipage/server-sent-events.html#event-stream-interpretation
 
-        if self._at_stream_start:
-            self._at_stream_start = False
-            line = line.removeprefix(b"\xef\xbb\xbf")
-
         if not line:
             if (
                 not self._event
@@ -112,7 +103,7 @@ class SSEDecoder:
             sse = StreamPart(
                 event=self._event,
                 data=orjson.loads(self._data) if self._data else None,  # ty: ignore[invalid-argument-type]
-                id=self._last_event_id or None,
+                id=self.last_event_id,
             )
 
             # NOTE: as per the SSE spec, do not reset last_event_id.
@@ -133,8 +124,6 @@ class SSEDecoder:
         if fieldname == b"event":
             self._event = value.decode()
         elif fieldname == b"data":
-            if self._data:
-                self._data.extend(b"\n")
             self._data.extend(value)
         elif fieldname == b"id":
             if b"\0" in value:

libs/sdk-py/tests/test_client_stream.py

@@ -5,7 +5,6 @@ from pathlib import Path
 from typing import Any
 
 import httpx
-import orjson
 import pytest
 from typing_extensions import assert_type
 
@@ -111,25 +110,6 @@ def test_stream_sse():
         assert len(parts) == 79
 
 
-def test_sse_decoder_joins_data_lines_with_a_newline():
-    decoder = SSEDecoder()
-    for line in (b"data: [1,", b"data: 2]"):
-        decoder.decode(line)
-    assert decoder.decode(b"") == StreamPart(event="", data=[1, 2], id=None)  # ty: ignore[invalid-argument-type]
-
-    for line in (b"data: [1", b"data: 0]"):
-        decoder.decode(line)
-    with pytest.raises(orjson.JSONDecodeError):
-        decoder.decode(b"")
-
-
-def test_sse_decoder_strips_a_leading_bom():
-    decoder = SSEDecoder()
-    for line in (b"\xef\xbb\xbfevent: values", b"data: 1"):
-        decoder.decode(line)
-    assert decoder.decode(b"") == StreamPart(event="values", data=1, id=None)  # ty: ignore[invalid-argument-type]
-
-
 # --- HTTP client streaming ---
 
 
@@ -312,52 +292,6 @@ async def test_http_client_stream_recovers_after_disconnect():
     ]
 
 
-def _drop_after_an_empty_id(stream_cls: type, reconnects: list[httpx.Request]):
-    def handler(request: httpx.Request) -> httpx.Response:
-        if request.method == "POST":
-            return httpx.Response(
-                200,
-                headers={"Content-Type": "text/event-stream", "Location": "/reconnect"},
-                stream=stream_cls(
-                    [
-                        b"id: 1\nevent: values\ndata: 1\n\n",
-                        b"id:\nevent: values\ndata: 2\n\n",
-                    ],
-                    httpx.RemoteProtocolError("incomplete chunked read"),
-                ),
-            )
-        reconnects.append(request)
-        return httpx.Response(
-            200,
-            headers={"Content-Type": "text/event-stream"},
-            stream=stream_cls([b"event: end\ndata: null\n\n"]),
-        )
-
-    return handler
-
-
-def test_sync_stream_reconnects_without_a_cursor_after_an_empty_id():
-    reconnects: list[httpx.Request] = []
-    transport = httpx.MockTransport(_drop_after_an_empty_id(ListByteStream, reconnects))
-    with httpx.Client(transport=transport, base_url="https://example.com") as client:
-        list(SyncHttpClient(client).stream("/stream", "POST", json={}))
-
-    assert "last-event-id" not in reconnects[0].headers
-
-
-async def test_async_stream_reconnects_without_a_cursor_after_an_empty_id():
-    reconnects: list[httpx.Request] = []
-    transport = httpx.MockTransport(
-        _drop_after_an_empty_id(AsyncListByteStream, reconnects)
-    )
-    async with httpx.AsyncClient(
-        transport=transport, base_url="https://example.com"
-    ) as client:
-        [part async for part in HttpClient(client).stream("/stream", "POST", json={})]
-
-    assert "last-event-id" not in reconnects[0].headers
-
-
 # --- _sse_to_v2_dict conversion ---