From 338180f383500b54106ea542b48419e98a7d1389 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Fri, 14 Jun 2024 12:54:20 -0700
Subject: [PATCH] core: in astream_events v2 always await task even if already
 finished (#22916)

- this ensures exceptions propagate to the caller
---
 libs/core/langchain_core/tracers/event_stream.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/libs/core/langchain_core/tracers/event_stream.py b/libs/core/langchain_core/tracers/event_stream.py
index ad2d4eb6924..b7c6fa0359c 100644
--- a/libs/core/langchain_core/tracers/event_stream.py
+++ b/libs/core/langchain_core/tracers/event_stream.py
@@ -940,9 +940,10 @@ async def _astream_events_implementation_v2(
 
             yield event
     finally:
-        # Wait for the runnable to finish, if not cancelled (eg. by break)
-        if task.cancel():
-            try:
-                await task
-            except asyncio.CancelledError:
-                pass
+        # Cancel the task if it's still running
+        task.cancel()
+        # Await it anyway, to run any cleanup code, and propagate any exceptions
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass