From eb2ffaae79654883228fd6a51b85039f69016882 Mon Sep 17 00:00:00 2001
From: Yujong Lee <yujong@berri.ai>
Date: Mon, 14 Sep 2026 11:37:40 -0700
Subject: [PATCH] fix(harness): rebuild bridge for trace parity

---
 .../shared/native_build.py                    | 20 ++++++++++++++++++-
 .../shared/test_native_build.py               | 11 +++++-----
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/tests/rust-python-harness/shared/native_build.py b/tests/rust-python-harness/shared/native_build.py
index b50def98e3a..f67488cecb4 100644
--- a/tests/rust-python-harness/shared/native_build.py
+++ b/tests/rust-python-harness/shared/native_build.py
@@ -16,6 +16,11 @@ _RUST_ROOT: Final = "litellm-rust"
 _LOCKFILE: Final = "Cargo.lock"
 _SOURCE_SUFFIXES: Final = frozenset({".rs", ".toml"})
 _FAILURE_OUTPUT_LINES: Final = 15
+_TRACE_CHECK: Final = (
+    "from litellm.rust_bridge import get_native_bridge; "
+    "bridge = get_native_bridge(); "
+    "raise SystemExit(0 if bridge is not None and getattr(bridge, '_trace', None) is not None else 1)"
+)
 
 
 def needs_rebuild(native_mtime: float | None, newest_source_mtime: float | None) -> bool:
@@ -73,6 +78,17 @@ def _rebuild(repo_root: Path) -> tuple[bool, str]:
     return completed.returncode == 0, "\n".join(lines[-_FAILURE_OUTPUT_LINES:])
 
 
+def _installed_bridge_has_trace(repo_root: Path) -> bool:
+    completed: Final = subprocess.run(
+        (sys.executable, "-c", _TRACE_CHECK),
+        cwd=repo_root,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        check=False,
+    )
+    return completed.returncode == 0
+
+
 def trace_bridge_error() -> str | None:
     bridge: Final = get_native_bridge()
     if bridge is None:
@@ -85,7 +101,9 @@ def trace_bridge_error() -> str | None:
 def ensure_trace_bridge(repo_root: Path) -> str | None:
     native_path: Final = _native_module_path()
     native_mtime: Final = native_path.stat().st_mtime if native_path is not None and native_path.exists() else None
-    rebuild_required: Final = needs_rebuild(native_mtime, _newest_source_mtime(repo_root)) or trace_bridge_error() is not None
+    rebuild_required: Final = needs_rebuild(
+        native_mtime, _newest_source_mtime(repo_root)
+    ) or not _installed_bridge_has_trace(repo_root)
     if rebuild_required:
         print(f"Rebuilding native Rust bridge ({BRIDGE_FEATURE} feature)...", flush=True)
         succeeded: Final
diff --git a/tests/rust-python-harness/shared/test_native_build.py b/tests/rust-python-harness/shared/test_native_build.py
index 982f787aefa..dc08bc1a2b6 100644
--- a/tests/rust-python-harness/shared/test_native_build.py
+++ b/tests/rust-python-harness/shared/test_native_build.py
@@ -105,13 +105,14 @@ def test_ensure_trace_bridge_rebuilds_when_trace_feature_is_missing(
         state.rebuilt = True
         return True, ""
 
+    def fake_get_native_bridge() -> SimpleNamespace:
+        assert state.rebuilt
+        return SimpleNamespace(_trace=object())
+
     monkeypatch.setattr(native_build, "_native_module_path", lambda: native)
     monkeypatch.setattr(native_build, "_rebuild", fake_rebuild)
-    monkeypatch.setattr(
-        native_build,
-        "get_native_bridge",
-        lambda: SimpleNamespace(_trace=object() if state.rebuilt else None),
-    )
+    monkeypatch.setattr(native_build, "_installed_bridge_has_trace", lambda repo_root: False)
+    monkeypatch.setattr(native_build, "get_native_bridge", fake_get_native_bridge)
 
     message: Final = native_build.ensure_trace_bridge(tmp_path)