fix(closes OPEN-10343): handle Google ADK callback exceptions cleanly

When a user-defined `before_model_callback` raises (e.g. a guardrail
short-circuit), the Openlayer tracer used to surface a chained
`ValueError("<Token> was created in a different Context")` from its
contextvar cleanup paths and overwrite span output with `"Error: ..."`,
polluting tracebacks and traces. Make exception unwinding boring:

- Add `_safe_reset_contextvar`: wraps `ContextVar.reset(token)` with the
  same `try/except ValueError` pattern used by `tracer.create_step`, so
  cross-context tokens during async unwind no longer raise a chained
  exception. Apply it to the three bare reset sites in the agent
  transfer-token path, async callback wrapper, and sync callback wrapper.
- Add `_record_step_error`: stores `metadata["error"] = {type, message}`
  on the step instead of overwriting `step.output` with `"Error: ..."`.
  Apply it to all five wrappers (agent, LLM, tool, async/sync callback).
- Demote the three pass-through `logger.error("Error in ...")` lines to
  `logger.debug` — the exception already shows up in the user's
  traceback, our duplicate log only added confusion.
- Add `tests/test_google_adk_integration.py` with end-to-end coverage
  asserting the user's exception is the sole exception in the chain,
  that metadata records the error, and that `step.output` is not
  overwritten — plus unit tests for the new helpers.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: viniciusdsmello

src/openlayer/lib/integrations/google_adk_tracer.py

@@ -47,6 +47,7 @@
 
 from ..tracing import tracer, steps, enums
 from ..tracing.tracer import _current_step as _tracer_current_step
+from ..tracing.tracer import _safe_reset_contextvar
 
 logger = logging.getLogger(__name__)
 
@@ -284,6 +285,21 @@ def _sort_steps_by_time(step: Any, recursive: bool = True) -> None:
             _sort_steps_by_time(child_step, recursive=True)
 
 
+def _record_step_error(step: Any, error: BaseException) -> None:
+    """Record exception info on a step's metadata without overwriting output."""
+    if step is None:
+        return
+    try:
+        if getattr(step, "metadata", None) is None:
+            step.metadata = {}
+        step.metadata["error"] = {
+            "type": type(error).__name__,
+            "message": str(error),
+        }
+    except Exception:  # pragma: no cover - defensive: never break unwinding
+        logger.debug("Failed to record error on step metadata", exc_info=True)
+
+
 def _build_llm_request_for_trace(llm_request: Any) -> Dict[str, Any]:
     """Build a dictionary representation of the LLM request for tracing.
 
@@ -676,8 +692,8 @@ async def new_function():
                         step.output = "Agent execution completed"
 
                 except Exception as e:
-                    step.output = f"Error: {str(e)}"
-                    logger.error(f"Error in agent execution: {e}")
+                    _record_step_error(step, e)
+                    logger.debug("Agent execution raised; propagating: %s", e)
                     raise
                 finally:
                     # Sort all nested steps recursively by start_time to ensure chronological order
@@ -691,8 +707,7 @@ async def new_function():
 
             # Restore the current step context if we changed it for transfer
             # This must be done AFTER the with block exits
-            if transfer_token is not None:
-                _tracer_current_step.reset(transfer_token)
+            _safe_reset_contextvar(_tracer_current_step, transfer_token)
 
         return new_function()
 
@@ -923,8 +938,8 @@ async def new_function():
                             pass
 
                 except Exception as e:
-                    step.output = f"Error: {str(e)}"
-                    logger.error(f"Error in LLM call: {e}")
+                    _record_step_error(step, e)
+                    logger.debug("LLM call raised; propagating: %s", e)
                     raise
                 finally:
                     # Sort nested steps by start_time for correct chronological order
@@ -1025,8 +1040,8 @@ async def new_function():
                     return result
 
                 except Exception as e:
-                    step.output = f"Error: {str(e)}"
-                    logger.error(f"Error in tool execution: {e}")
+                    _record_step_error(step, e)
+                    logger.debug("Tool execution raised; propagating: %s", e)
                     raise
                 finally:
                     # Sort nested steps by start_time for correct chronological order
@@ -1296,12 +1311,14 @@ async def async_traced_callback(*args, **kwargs):
 
                             return result
                         except Exception as e:
-                            step.output = f"Error: {str(e)}"
+                            _record_step_error(step, e)
                             raise
                 finally:
-                    # Restore the previous current step
-                    if saved_token is not None:
-                        _tracer_current_step.reset(saved_token)
+                    # Restore the previous current step. Wrapped defensively
+                    # because the await chain can cross asyncio Contexts and
+                    # would otherwise raise a chained ValueError from this
+                    # cleanup path during exception unwinding (OPEN-10343).
+                    _safe_reset_contextvar(_tracer_current_step, saved_token)
 
             return async_traced_callback
         else:
@@ -1363,12 +1380,12 @@ def sync_traced_callback(*args, **kwargs):
 
                             return result
                         except Exception as e:
-                            step.output = f"Error: {str(e)}"
+                            _record_step_error(step, e)
                             raise
                 finally:
-                    # Restore the previous current step
-                    if saved_token is not None:
-                        _tracer_current_step.reset(saved_token)
+                    # Restore the previous current step. Wrapped defensively
+                    # for symmetry with the async path; see OPEN-10343.
+                    _safe_reset_contextvar(_tracer_current_step, saved_token)
 
             return sync_traced_callback
 

src/openlayer/lib/tracing/tracer.py

@@ -466,6 +466,30 @@ def get_rag_question() -> Optional[str]:
     return _rag_question.get(None)
 
 
+_CROSS_CONTEXT_TOKEN_MESSAGE = "was created in a different Context"
+
+
+def _safe_reset_contextvar(var: "contextvars.ContextVar[Any]", token: Any) -> None:
+    """Reset a ContextVar token, swallowing the cross-context ``ValueError``.
+
+    Async generators driven by integrations (Google ADK, etc.) can be entered
+    and exited in different asyncio Contexts. The token from ``set()`` then
+    becomes invalid and ``reset()`` raises
+    ``ValueError("<Token ...> was created in a different Context")``. During
+    exception unwinding that secondary failure becomes a chained exception
+    polluting the user's traceback (OPEN-10343). Other ``ValueError``s are
+    re-raised.
+    """
+    if token is None:
+        return
+    try:
+        var.reset(token)
+    except ValueError as exc:
+        if _CROSS_CONTEXT_TOKEN_MESSAGE not in str(exc):
+            raise
+        logger.debug("Skipping cross-context ContextVar reset: %s", exc)
+
+
 @contextmanager
 def create_step(
     name: str,
@@ -493,13 +517,7 @@ def create_step(
             latency = (new_step.end_time - new_step.start_time) * 1000  # in ms
             new_step.latency = latency
 
-        try:
-            _current_step.reset(token)
-        except ValueError as e:
-            # Handle context variable mismatch gracefully
-            # This can occur when async generators cross context boundaries
-            if "was created in a different Context" not in str(e):
-                raise
+        _safe_reset_contextvar(_current_step, token)
 
         _handle_trace_completion(
             is_root_step=is_root_step,

tests/test_google_adk_integration.py

@@ -0,0 +1,224 @@
+"""Tests for the Google ADK tracer integration.
+
+Regression coverage for OPEN-10343: when a user-defined ``before_model_callback``
+raises, the Openlayer tracer must let the original exception propagate
+unchanged (no chained ``ValueError`` from contextvar cleanup) and must record
+the error on the span's metadata without overwriting ``step.output``.
+"""
+
+import asyncio
+from typing import Any, Dict, List, Optional, Tuple
+from unittest.mock import patch
+
+import pytest
+
+pytest.importorskip("google.adk")
+pytest.importorskip("wrapt")
+
+from openlayer.lib.integrations.google_adk_tracer import (
+    _record_step_error,
+    _safe_reset_contextvar,
+)
+
+
+class ErrorHandling(Exception):
+    """User-defined exception raised by a guardrail-style callback."""
+
+
+@pytest.fixture(autouse=True)
+def _disable_publish(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Keep all tracer publish paths off during the test."""
+    monkeypatch.setenv("OPENLAYER_DISABLE_PUBLISH", "true")
+    monkeypatch.setenv("OPENLAYER_API_KEY", "fake")
+    monkeypatch.setenv("GOOGLE_API_KEY", "fake")
+
+    from openlayer.lib.tracing import tracer as _tracer
+
+    monkeypatch.setattr(_tracer, "_publish", False, raising=False)
+
+
+def _collect_exception_chain(exc: BaseException) -> List[BaseException]:
+    """Walk ``__cause__`` / ``__context__`` and return the unique chain."""
+    seen: List[BaseException] = []
+    current: Optional[BaseException] = exc
+    while current is not None and current not in seen:
+        seen.append(current)
+        nxt = current.__cause__ if current.__cause__ is not None else current.__context__
+        current = nxt
+    return seen
+
+
+class TestGoogleADKCallbackExceptions:
+    """OPEN-10343: chained / noisy exceptions from before_model_callback."""
+
+    APP_NAME = "openlayer_open_10343"
+    USER_ID = "u"
+    SESSION_ID = "s"
+
+    def _build_runner_and_agent(self) -> Tuple[Any, Any]:
+        from google.adk.agents import LlmAgent
+        from google.adk.runners import Runner
+        from google.adk.sessions import InMemorySessionService
+
+        from openlayer.lib.integrations import trace_google_adk
+
+        trace_google_adk()
+
+        async def before_model(callback_context: Any, llm_request: Any, **_: Any) -> None:
+            raise ErrorHandling("input blocked by guardrail (pi_and_jailbreak)")
+
+        agent = LlmAgent(
+            model="gemini-2.5-flash",
+            name="OpenLayerTestAgent",
+            instruction="test",
+            before_model_callback=before_model,
+        )
+        session_service = InMemorySessionService()
+        runner = Runner(agent=agent, app_name=self.APP_NAME, session_service=session_service)
+        return runner, session_service
+
+    async def _drive_runner(self, runner: Any, session_service: Any) -> None:
+        from google.genai import types
+
+        await session_service.create_session(
+            app_name=self.APP_NAME, user_id=self.USER_ID, session_id=self.SESSION_ID
+        )
+        content = types.Content(role="user", parts=[types.Part(text="hi")])
+        async for _event in runner.run_async(
+            user_id=self.USER_ID, session_id=self.SESSION_ID, new_message=content
+        ):
+            pass
+
+    @pytest.mark.filterwarnings("ignore::DeprecationWarning")
+    def test_user_exception_propagates_without_chaining(self) -> None:
+        """The user's ``ErrorHandling`` must be the only exception surfaced."""
+        runner, session_service = self._build_runner_and_agent()
+
+        with pytest.raises(ErrorHandling) as exc_info:
+            asyncio.run(self._drive_runner(runner, session_service))
+
+        chain = _collect_exception_chain(exc_info.value)
+
+        # No chained ValueError from contextvars / token-reset cleanup.
+        assert not any(
+            isinstance(item, ValueError)
+            and "was created in a different Context" in str(item)
+            for item in chain
+        ), f"chained ValueError appeared in exception chain: {chain}"
+
+        # Anything in the chain that isn't the user's exception would itself
+        # be a tracer-introduced failure — fail loudly.
+        for item in chain:
+            assert isinstance(item, ErrorHandling), (
+                "unexpected non-ErrorHandling exception in chain: "
+                f"{type(item).__name__}: {item}"
+            )
+
+    @pytest.mark.filterwarnings("ignore::DeprecationWarning")
+    def test_callback_error_recorded_on_step_metadata(self) -> None:
+        """The callback step's metadata['error'] should describe the failure."""
+        from openlayer.lib.tracing import tracer as _tracer
+        from openlayer.lib.tracing.traces import Trace
+
+        captured: Dict[str, Any] = {}
+        original_handle = _tracer._handle_trace_completion
+
+        def _capture_handle(*args: Any, **kwargs: Any) -> None:
+            current_trace = _tracer.get_current_trace()
+            if current_trace is not None and "trace" not in captured:
+                captured["trace"] = current_trace
+            return original_handle(*args, **kwargs)
+
+        runner, session_service = self._build_runner_and_agent()
+
+        with patch.object(_tracer, "_handle_trace_completion", _capture_handle):
+            with pytest.raises(ErrorHandling):
+                asyncio.run(self._drive_runner(runner, session_service))
+
+        trace: Trace = captured["trace"]
+
+        callback_steps = self._find_callback_steps(trace)
+        assert callback_steps, "no callback steps were recorded on the trace"
+
+        before_model_step = next(
+            (s for s in callback_steps if s.metadata.get("callback_type") == "before_model"),
+            None,
+        )
+        assert before_model_step is not None, "before_model callback step missing"
+
+        recorded = before_model_step.metadata.get("error")
+        assert recorded is not None, "error metadata not recorded on callback step"
+        assert recorded["type"] == "ErrorHandling"
+        assert "pi_and_jailbreak" in recorded["message"]
+
+        # And critically: step.output was *not* overwritten with "Error: ...".
+        if before_model_step.output is not None:
+            assert not (
+                isinstance(before_model_step.output, str)
+                and before_model_step.output.startswith("Error:")
+            ), (
+                "step.output should not be overwritten with 'Error: ...'; "
+                f"got {before_model_step.output!r}"
+            )
+
+    @staticmethod
+    def _find_callback_steps(trace: Any) -> List[Any]:
+        """Walk the trace tree and return all callback steps (USER_CALL)."""
+        from openlayer.lib.tracing import enums
+
+        out: List[Any] = []
+
+        def _walk(step: Any) -> None:
+            if getattr(step, "step_type", None) == enums.StepType.USER_CALL:
+                if step.metadata.get("is_callback"):
+                    out.append(step)
+            for child in getattr(step, "steps", []) or []:
+                _walk(child)
+
+        for step in getattr(trace, "steps", []) or []:
+            _walk(step)
+        return out
+
+
+class TestGoogleADKHelpers:
+    """Unit tests for the small helper functions added in OPEN-10343."""
+
+    def test_safe_reset_swallows_cross_context_value_error(self) -> None:
+        class FakeVar:
+            def reset(self, _token: Any) -> None:
+                raise ValueError("<Token at 0x1> was created in a different Context")
+
+        # Should not raise.
+        _safe_reset_contextvar(FakeVar(), object())  # type: ignore[arg-type]
+
+    def test_safe_reset_reraises_other_value_errors(self) -> None:
+        class FakeVar:
+            def reset(self, _token: Any) -> None:
+                raise ValueError("some other reason")
+
+        with pytest.raises(ValueError, match="some other reason"):
+            _safe_reset_contextvar(FakeVar(), object())  # type: ignore[arg-type]
+
+    def test_safe_reset_no_op_for_none_token(self) -> None:
+        class FakeVar:
+            def reset(self, _token: Any) -> None:  # pragma: no cover
+                raise AssertionError("should not be called when token is None")
+
+        _safe_reset_contextvar(FakeVar(), None)  # type: ignore[arg-type]
+
+    def test_record_step_error_writes_metadata(self) -> None:
+        class FakeStep:
+            metadata: Optional[Dict[str, Any]] = None
+            output: Optional[Any] = "preserved"
+
+        step = FakeStep()
+        _record_step_error(step, ErrorHandling("boom"))
+
+        assert step.metadata is not None
+        assert step.metadata["error"] == {"type": "ErrorHandling", "message": "boom"}
+        # Output left untouched.
+        assert step.output == "preserved"
+
+    def test_record_step_error_handles_none_step(self) -> None:
+        # Should not raise.
+        _record_step_error(None, ErrorHandling("boom"))