fix(vertexai): make dont_throw work on async functions

packages/opentelemetry-instrumentation-vertexai/opentelemetry/instrumentation/vertexai/utils.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 import os
 import traceback
@@ -17,26 +18,43 @@ def should_send_prompts():
 def dont_throw(func):
     """
     A decorator that wraps the passed in function and logs exceptions instead of throwing them.
+    Works for both synchronous and asynchronous functions.
 
     @param func: The function to wrap
     @return: The wrapper function
     """
     # Obtain a logger specific to the function's module
     logger = logging.getLogger(func.__module__)
 
-    def wrapper(*args, **kwargs):
+    def _handle_exception(e):
+        logger.debug(
+            "OpenLLMetry failed to trace in %s, error: %s",
+            func.__name__,
+            traceback.format_exc(),
+        )
+        if Config.exception_logger:
+            try:
+                Config.exception_logger(e)
+            except Exception:
+                logger.debug(
+                    "OpenLLMetry exception logger failed in %s",
+                    func.__name__,
+                    exc_info=True,
+                )
+
+    async def async_wrapper(*args, **kwargs):
+        try:
+            return await func(*args, **kwargs)
+        except Exception as e:
+            _handle_exception(e)
+
+    def sync_wrapper(*args, **kwargs):
         try:
             return func(*args, **kwargs)
         except Exception as e:
-            logger.debug(
-                "OpenLLMetry failed to trace in %s, error: %s",
-                func.__name__,
-                traceback.format_exc(),
-            )
-            if Config.exception_logger:
-                Config.exception_logger(e)
+            _handle_exception(e)
 
-    return wrapper
+    return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
 
 
 def should_emit_events():

packages/opentelemetry-instrumentation-vertexai/tests/test_dont_throw.py

@@ -0,0 +1,94 @@
+import asyncio
+
+import pytest
+from opentelemetry.instrumentation.vertexai.utils import dont_throw
+
+
+def test_dont_throw_swallows_sync_exceptions():
+    @dont_throw
+    def boom():
+        raise RuntimeError("instrumentation failed")
+
+    assert boom() is None, "sync instrumentation errors must not reach the caller"
+
+
+@pytest.mark.asyncio
+async def test_dont_throw_swallows_async_exceptions():
+    """An async function returns a coroutine immediately, so a sync-only
+    wrapper exits its try block before the body runs and the caller awaits
+    outside the guard. span_utils.set_input_attributes and _handle_request are
+    both async and both decorated, and _handle_request is awaited unguarded."""
+
+    @dont_throw
+    async def boom():
+        raise RuntimeError("instrumentation failed")
+
+    assert await boom() is None, "async instrumentation errors must not reach the caller"
+
+
+@pytest.mark.asyncio
+async def test_dont_throw_returns_async_values():
+    @dont_throw
+    async def fine():
+        return "ok"
+
+    assert await fine() == "ok", "decorator must not swallow the return value"
+
+
+def test_dont_throw_preserves_sync_return():
+    @dont_throw
+    def fine():
+        return "ok"
+
+    assert fine() == "ok"
+
+
+def test_dont_throw_picks_wrapper_by_function_kind():
+    @dont_throw
+    async def coro():
+        return None
+
+    @dont_throw
+    def plain():
+        return None
+
+    assert asyncio.iscoroutinefunction(coro), "async functions need the async wrapper"
+    assert not asyncio.iscoroutinefunction(plain)
+
+
+def test_dont_throw_survives_a_failing_exception_logger():
+    """A user-supplied exception_logger that itself raises must not reach the caller."""
+    from opentelemetry.instrumentation.vertexai.config import Config
+
+    def angry_logger(_e):
+        raise ValueError("exception logger is broken")
+
+    previous = Config.exception_logger
+    Config.exception_logger = angry_logger
+    try:
+        @dont_throw
+        def boom():
+            raise RuntimeError("instrumentation failed")
+
+        assert boom() is None, "a failing exception_logger must not surface to the caller"
+    finally:
+        Config.exception_logger = previous
+
+
+@pytest.mark.asyncio
+async def test_dont_throw_async_survives_a_failing_exception_logger():
+    from opentelemetry.instrumentation.vertexai.config import Config
+
+    def angry_logger(_e):
+        raise ValueError("exception logger is broken")
+
+    previous = Config.exception_logger
+    Config.exception_logger = angry_logger
+    try:
+        @dont_throw
+        async def boom():
+            raise RuntimeError("instrumentation failed")
+
+        assert await boom() is None
+    finally:
+        Config.exception_logger = previous