fix(harness): rebuild bridge for trace parity

This commit is contained in:
Yujong Lee 2026-09-14 11:37:40 -07:00
parent 561533c596
commit eb2ffaae79
2 changed files with 25 additions and 6 deletions

View file

@ -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

View file

@ -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)