From 9958f4785202844730e490c8b59a81f6c87b8f4d Mon Sep 17 00:00:00 2001
From: shuofengzhang <sofuncheung16@gmail.com>
Date: Mon, 27 Apr 2026 13:42:19 -0700
Subject: [PATCH] fix(async): preserve typed errors in research stream mode

---
 tavily/async_tavily.py |  5 ++++-
 tests/test_research.py | 20 ++++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/tavily/async_tavily.py b/tavily/async_tavily.py
index 663d228..99fd19f 100644
--- a/tavily/async_tavily.py
+++ b/tavily/async_tavily.py
@@ -723,8 +723,11 @@ async def stream_generator() -> AsyncGenerator[bytes, None]:
                                 yield chunk
                 except httpx.TimeoutException:
                     raise TimeoutError(timeout)
+                except (UsageLimitExceededError, ForbiddenError, InvalidAPIKeyError, BadRequestError, TimeoutError):
+                    # Preserve Tavily-specific error types for callers in streaming mode.
+                    raise
                 except Exception as e:
-                    raise Exception(f"Error during research stream: {str(e)}")
+                    raise Exception(f"Error during research stream: {str(e)}") from e
 
             return stream_generator()
         else:
diff --git a/tests/test_research.py b/tests/test_research.py
index 8c2747f..f1794f4 100644
--- a/tests/test_research.py
+++ b/tests/test_research.py
@@ -1,5 +1,9 @@
 import asyncio
 
+import pytest
+
+from tavily.errors import UsageLimitExceededError
+
 BASE_URL = "https://api.tavily.com"
 
 dummy_queued_response = {
@@ -147,3 +151,19 @@ def test_async_get_research(async_interceptor, async_client):
     request = async_interceptor.get_request()
     validate_get_research(request, response)
 
+
+def test_async_research_stream_preserves_usage_limit_error(async_interceptor, async_client):
+    async_interceptor.set_response(429, json={"detail": {"error": "quota exceeded"}})
+
+    async def run_test():
+        stream = await async_client.research(
+            input="Research the latest developments in AI",
+            stream=True,
+        )
+
+        with pytest.raises(UsageLimitExceededError, match="quota exceeded"):
+            async for _ in stream:
+                pass
+
+    asyncio.run(run_test())
+