Restore outer tracing contexts after nested predictions

great_ai/tracing/tracing_context.py

@@ -74,7 +74,7 @@ class TracingContext(Generic[T]):
         return _current_tracing_context.get()
 
     def __enter__(self) -> "TracingContext":
-        _current_tracing_context.set(self)
+        self._token = _current_tracing_context.set(self)
         return self
 
     def __exit__(
@@ -83,7 +83,7 @@ class TracingContext(Generic[T]):
         exception: Optional[BaseException],
         traceback: Optional[TracebackType],
     ) -> Literal[False]:
-        _current_tracing_context.set(None)
+        _current_tracing_context.reset(self._token)
 
         if exception is not None and type is not None:
             self.finalise(exception=exception)

tests/test_tracing_regressions.py

@@ -0,0 +1,27 @@
+import pytest
+from great_ai import GreatAI, log_metric
+from great_ai.tracing import TracingContext
+
+
+@pytest.mark.parametrize("fail", [False, True])
+def test_nested_predictions_restore_outer_trace(
+    isolated_context: None, fail: bool
+) -> None:
+    @GreatAI.create
+    def inner(value: int) -> int:
+        if fail:
+            raise ValueError("inner failed")
+        return value
+
+    @GreatAI.create
+    def outer(value: int) -> int:
+        try:
+            inner(value)
+        except ValueError:
+            pass
+        log_metric("after_inner", 7, disable_logging=True)
+        return value
+
+    trace = outer(1)
+    assert trace.logged_values["metric:outer:after_inner"] == 7
+    assert TracingContext.get_current_tracing_context() is None