From db378d896349b7ee0ca02cf48d0cb266f4b95736 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 14 Sep 2026 12:08:31 -0700 Subject: [PATCH] fix(harness): preserve unmapped calls in execution traces --- tests/rust-python-harness/AGENTS.md | 2 +- .../shared/tracing/steps.py | 4 +- .../shared/tracing/test_steps.py | 17 ++++ .../strategies/trace_parity/AGENTS.md | 2 +- .../trace_parity/gateway/execution.py | 8 +- .../strategies/trace_parity/models.py | 2 + .../strategies/trace_parity/sdk/execution.py | 10 +-- .../strategies/trace_parity/sdk/ocr/case.py | 46 +++++++++++ .../strategies/trace_parity/test_runner.py | 80 ++++++++++++++++++- 9 files changed, 156 insertions(+), 15 deletions(-) diff --git a/tests/rust-python-harness/AGENTS.md b/tests/rust-python-harness/AGENTS.md index ec973035785..71e17541fd2 100644 --- a/tests/rust-python-harness/AGENTS.md +++ b/tests/rust-python-harness/AGENTS.md @@ -63,7 +63,7 @@ tests/rust-python-harness/ - Examples: `run e2e_parity --surface sdk --function ocr`, `run unit_tests_parity --function ocr --pytest-arg=-x`, or `run all --function ocr` - `cli/catalog.py` discovers strategies, validates their Python definitions, and orders them; `cli/__init__.py` builds the Click command tree; `cli/commands.py` runs selected cases - `e2e_parity/` compares SDK objects, exceptions, callbacks, and streams, or gateway HTTP responses -- `trace_parity/` prints filtered Python and Rust execution traces without comparing them; before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`) +- `trace_parity/` prints every collected Python call under `litellm/` and every Rust span without comparing them; mappings only filter the separate unit-test mapping strategy. Before running it rebuilds the native bridge with the `trace-parity` feature whenever `litellm-rust` sources are newer than the installed extension (`shared/native_build.py`) - E2E and trace strategies load their registered module cases and run surface-specific execution from their folders - `unit_tests_mapping/contracts.py` owns typed harness-side mapping contracts, per-function contracts live below `cases/`, and `mappings.py` exports the registry; live test discovery derives unmapped Python and Rust-only tests without an exhaustive manifest - `unit_tests_mapping/runner.py` validates confirmed mappings against the live Python and Rust inventories and attaches the derived status report diff --git a/tests/rust-python-harness/shared/tracing/steps.py b/tests/rust-python-harness/shared/tracing/steps.py index 152f753a536..492ffab64e5 100644 --- a/tests/rust-python-harness/shared/tracing/steps.py +++ b/tests/rust-python-harness/shared/tracing/steps.py @@ -76,7 +76,7 @@ def _span_for(engine: Engine, function: str, mappings: Sequence[TraceMapping]) - def pipeline_projection( - engine: Engine, events: Sequence[FunctionTraceEvent], mappings: Sequence[TraceMapping] + engine: Engine, events: Sequence[FunctionTraceEvent], mappings: Sequence[TraceMapping] | None = None ) -> PipelineProjection: raw_parents: dict[int, int | None] = {} projected_ids: set[int] = set() @@ -88,7 +88,7 @@ def pipeline_projection( if event.parent_id is not None and event.parent_id not in raw_parents: raise ValueError(f"trace event {event.id} references unknown or later parent {event.parent_id}") raw_parents[event.id] = event.parent_id - span = _span_for(engine, event.function, mappings) + span = event.function if mappings is None else _span_for(engine, event.function, mappings) if span is None: unmatched += 1 continue diff --git a/tests/rust-python-harness/shared/tracing/test_steps.py b/tests/rust-python-harness/shared/tracing/test_steps.py index 1b2f02a4ed7..ee5e0bafd28 100644 --- a/tests/rust-python-harness/shared/tracing/test_steps.py +++ b/tests/rust-python-harness/shared/tracing/test_steps.py @@ -40,6 +40,23 @@ def test_python_projection_collapses_unmapped_parents_and_counts_noise() -> None ] +@pytest.mark.parametrize("engine", ("python", "rust")) +def test_projection_without_mappings_keeps_every_call_and_parent(engine: Engine) -> None: + events: Final = ( + event(0, "module.py:1 entry"), + event(1, "module.py:2 internal_helper", 0), + event(2, "module.py:3 nested", 1), + event(3, "module.py:2 internal_helper", 0), + ) + + projection: Final = pipeline_projection(engine, events) + + assert projection.unmatched == 0 + assert tuple((step.id, step.parent_id, step.span, step.raw) for step in projection.steps) == tuple( + (item.id, item.parent_id, item.function, item.raw) for item in events + ) + + def test_rust_projection_keeps_unknown_spans() -> None: projection: Final = pipeline_projection("rust", (event(0, "route"), event(1, "new_span", 0)), MAPPINGS) assert [(step.span, step.parent_id) for step in projection.steps] == [("route", None), ("new_span", 0)] diff --git a/tests/rust-python-harness/strategies/trace_parity/AGENTS.md b/tests/rust-python-harness/strategies/trace_parity/AGENTS.md index 5db17974d07..19861aea2b4 100644 --- a/tests/rust-python-harness/strategies/trace_parity/AGENTS.md +++ b/tests/rust-python-harness/strategies/trace_parity/AGENTS.md @@ -1 +1 @@ -Prints filtered Python profiler frames and feature-gated Rust spans from live traces against a replayed provider response. The two traces are independent and are not compared. +Prints every collected Python call under litellm/ and every feature-gated Rust span from live traces against replayed HTTP responses. The two traces are independent and are not compared. API-key and Vertex credentials scenarios exercise separate authentication paths; credentials scenarios replay the token exchange locally. diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py index 88130dab808..5f24760dfe3 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py @@ -160,13 +160,11 @@ def _collect( def _projections( python_events: tuple[FunctionTraceEvent, ...], rust_events: tuple[FunctionTraceEvent, ...], - scenario: TraceScenario, ) -> tuple[PipelineProjection, PipelineProjection, str | None]: - mappings: Final = scenario.mappings try: return ( - pipeline_projection("python", python_events, mappings), - pipeline_projection("rust", rust_events, mappings), + pipeline_projection("python", python_events), + pipeline_projection("rust", rust_events), None, ) except ValueError as error: @@ -184,7 +182,7 @@ def execute_gateway_trace( rust_error: Final = None if isinstance(rust_trace, tuple) else f"rust: {rust_trace.message}" python_events: Final = python_trace if isinstance(python_trace, tuple) else () rust_events: Final = rust_trace if isinstance(rust_trace, tuple) else () - python, rust, projection_error = _projections(python_events, rust_events, scenario) + python, rust, projection_error = _projections(python_events, rust_events) python_error: Final = projection_error or collection_python_error return TraceArtifact.from_traces( engine=engine, diff --git a/tests/rust-python-harness/strategies/trace_parity/models.py b/tests/rust-python-harness/strategies/trace_parity/models.py index be7375dc604..fd4585d85c1 100644 --- a/tests/rust-python-harness/strategies/trace_parity/models.py +++ b/tests/rust-python-harness/strategies/trace_parity/models.py @@ -18,6 +18,7 @@ class RouteFixture: provider_responses: tuple[RecordedResponse, ...] expected_failure: bool = False consume_stream: bool = False + environment: tuple[tuple[str, str], ...] = () def derive( self, @@ -32,6 +33,7 @@ class RouteFixture: provider_responses=self.provider_responses if provider_responses is None else provider_responses, expected_failure=self.expected_failure if expected_failure is None else expected_failure, consume_stream=self.consume_stream if consume_stream is None else consume_stream, + environment=self.environment, ) def with_body(self, **updates: object) -> RouteFixture: diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py index f8a674451b8..2ef8a0584bc 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py @@ -120,21 +120,22 @@ def collect_trace( provider.enqueue_response(response) fixture: Final = RouteFixture( kwargs={ - **base_fixture.kwargs, "api_key": "test-key", + **base_fixture.kwargs, "api_base": provider.url, **({"timeout_seconds": 5} if engine == "rust" else {"timeout": 5}), }, provider_responses=base_fixture.provider_responses, expected_failure=base_fixture.expected_failure, consume_stream=base_fixture.consume_stream, + environment=base_fixture.environment, ) environment: Final = ( patch.dict(os.environ, {"LITELLM_RUST": "1" if python_rust_enabled else "0"}) if engine == "python" else nullcontext() ) - with environment: + with environment, patch.dict(os.environ, fixture.environment): collected: Final = _collect(function, fixture, engine, asynchronous=asynchronous) provider.take_requests(len(fixture.provider_responses)) except Exception as error: @@ -160,7 +161,6 @@ def execute_trace( surface: Surface, engine: TraceEngine = "both", ) -> TraceArtifact: - mappings: Final = scenario.mappings scenario_route: Final = RouteSpec( route=route.route, python_entrypoints=route.python_entrypoints, @@ -185,8 +185,8 @@ def execute_trace( python_events: Final = python_trace if isinstance(python_trace, tuple) else () rust_events: Final = rust_trace if isinstance(rust_trace, tuple) else () try: - python: Final = pipeline_projection("python", python_events, mappings) - rust: Final = pipeline_projection("rust", rust_events, mappings) + python: Final = pipeline_projection("python", python_events) + rust: Final = pipeline_projection("rust", rust_events) except ValueError as error: return TraceArtifact.from_traces( engine=engine, diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py index b2c0e4c7acc..58b16c68383 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/ocr/case.py @@ -204,6 +204,40 @@ def _vertex_deepseek_fixture(engine: Engine, _base_url: str) -> RouteFixture: ) +def _vertex_deepseek_credentials_fixture(engine: Engine, base_url: str) -> RouteFixture: + from cryptography.hazmat.primitives import serialization + from cryptography.hazmat.primitives.asymmetric import rsa + + fixture: Final = _vertex_deepseek_fixture(engine, base_url) + private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + credentials: Final = json.dumps( + { + "type": "service_account", + "project_id": "trace-project", + "private_key_id": "trace-key", + "private_key": private_key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ).decode(), + "client_email": "trace@trace-project.iam.gserviceaccount.com", + "token_uri": f"{base_url}/token", + } + ) + return RouteFixture( + kwargs={**fixture.kwargs, "api_key": None}, + environment=(("VERTEXAI_CREDENTIALS", credentials), ("VERTEX_AI_API_KEY", "")), + provider_responses=( + RecordedHttpResponse.from_bytes( + 200, + (HttpHeader(name="content-type", value="application/json"),), + b'{"access_token":"trace-token","token_type":"Bearer","expires_in":3600}', + ), + *fixture.provider_responses, + ), + ) + + def _cohere_fixture(engine: Engine, _base_url: str) -> RouteFixture: return RouteFixture( kwargs={ @@ -441,6 +475,18 @@ TRACE_SUITE: Final = TraceSuite( mappings=(*DEEPSEEK_ASYNC_MAPPINGS, IGNORED_SUCCESS_CALLBACK_MAPPING), asynchronous=True, ), + TraceScenario( + name="sync-vertex-deepseek-credentials", + fixture=_vertex_deepseek_credentials_fixture, + mappings=DEEPSEEK_SYNC_MAPPINGS, + asynchronous=False, + ), + TraceScenario( + name="async-vertex-deepseek-credentials", + fixture=_vertex_deepseek_credentials_fixture, + mappings=DEEPSEEK_ASYNC_MAPPINGS, + asynchronous=True, + ), TraceScenario( name="async-cohere", fixture=_cohere_fixture, diff --git a/tests/rust-python-harness/strategies/trace_parity/test_runner.py b/tests/rust-python-harness/strategies/trace_parity/test_runner.py index 9dfd0d51c89..b940221bb7a 100644 --- a/tests/rust-python-harness/strategies/trace_parity/test_runner.py +++ b/tests/rust-python-harness/strategies/trace_parity/test_runner.py @@ -13,7 +13,7 @@ import litellm from ...shared.reporting.models import Coverage, HarnessCase, HarnessRun, RunStatus, SdkFunction, Surface from ...shared.reporting.strategy import ModuleCaseSpec from ...shared.tracing.profiler import FunctionTraceEvent -from ...shared.tracing.steps import Engine, PipelineStep +from ...shared.tracing.steps import Engine, PipelineStep, mapping from .models import GatewayRouteSpec, RouteFixture, RouteSpec, TraceScenario, TraceSuite from .reporting import TraceArtifact from .runner import run_trace_cases, run_trace_scenario, runner_selection, scenario_nodeids, validate_trace_suite @@ -124,6 +124,84 @@ def test_expected_provider_failure_omits_feedback_banner( assert litellm.suppress_debug_info is False +@pytest.mark.parametrize("asynchronous", (False, True)) +def test_vertex_trace_keeps_unmapped_helpers_and_parents(asynchronous: bool) -> None: + loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case") + suite: Final = cast(TraceSuite, loaded.TRACE_SUITE) + name: Final = f"{'async' if asynchronous else 'sync'}-vertex-deepseek" + scenario: Final = next(item for item in suite.scenarios if item.name == name) + assert isinstance(suite.route, RouteSpec) + + trace: Final = execute_trace(suite.route, scenario, "sdk", engine="python") + + assert trace.python_error is None + url: Final = next( + event for event in trace.python if event.raw.endswith(" VertexAIDeepSeekOCRConfig.get_complete_url") + ) + project: Final = next( + event for event in trace.python if event.raw.endswith(" VertexBase.safe_get_vertex_ai_project") + ) + location: Final = next( + event for event in trace.python if event.raw.endswith(" VertexBase.safe_get_vertex_ai_location") + ) + assert project.parent_id == location.parent_id == url.id + assert not any(event.raw.endswith(" VertexBase.get_access_token") for event in trace.python) + + +@pytest.mark.parametrize("asynchronous", (False, True)) +def test_vertex_credentials_trace_runs_real_auth_helpers(asynchronous: bool, monkeypatch: pytest.MonkeyPatch) -> None: + loaded: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.ocr.case") + suite: Final = cast(TraceSuite, loaded.TRACE_SUITE) + name: Final = f"{'async' if asynchronous else 'sync'}-vertex-deepseek-credentials" + scenario: Final = next(item for item in suite.scenarios if item.name == name) + monkeypatch.setenv("VERTEXAI_CREDENTIALS", "original-credentials") + monkeypatch.setenv("VERTEX_AI_API_KEY", "original-api-key") + assert isinstance(suite.route, RouteSpec) + + trace: Final = execute_trace(suite.route, scenario, "sdk", engine="python") + + assert trace.python_error is None + validate: Final = next( + event for event in trace.python if event.raw.endswith(" VertexAIDeepSeekOCRConfig.validate_environment") + ) + helpers: Final = ( + "VertexBase.safe_get_vertex_ai_project", + "VertexBase.safe_get_vertex_ai_credentials", + "VertexBase.get_access_token", + ) + assert tuple(event.raw.split(" ", 1)[1] for event in trace.python if event.parent_id == validate.id) == helpers + token: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.get_access_token")) + load: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.load_auth")) + refresh: Final = next(event for event in trace.python if event.raw.endswith(" VertexBase.refresh_auth")) + assert load.parent_id == token.id + assert refresh.parent_id == load.id + assert os.environ["VERTEXAI_CREDENTIALS"] == "original-credentials" + assert os.environ["VERTEX_AI_API_KEY"] == "original-api-key" + + +def test_gateway_trace_keeps_calls_outside_scenario_mappings(monkeypatch: pytest.MonkeyPatch) -> None: + execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.gateway.execution") + events: Final = ( + FunctionTraceEvent(0, None, "route.py:1 entry"), + FunctionTraceEvent(1, 0, "auth.py:2 authenticate"), + FunctionTraceEvent(2, 1, "auth.py:3 credentials"), + ) + scenario: Final = TraceScenario( + "async-gateway", + _fixture, + (mapping(rust_span="entry", python_frame=r" entry$"),), + asynchronous=True, + ) + monkeypatch.setattr(execution, "_collect", lambda *_args: events) + + trace: Final = execution.execute_gateway_trace(GatewayRouteSpec("messages"), scenario, engine="python") + + assert trace.python_error is None + assert tuple((event.id, event.parent_id, event.raw) for event in trace.python) == tuple( + (event.id, event.parent_id, event.raw) for event in events + ) + + def test_scenario_validation_rejects_duplicate_and_unsafe_names() -> None: route: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture) duplicate: Final = TraceSuite(