diff --git a/tests/rust-python-harness/strategies/trace_parity/fixtures.py b/tests/rust-python-harness/strategies/trace_parity/fixtures.py index 84ac525efb6..a3d084c97de 100644 --- a/tests/rust-python-harness/strategies/trace_parity/fixtures.py +++ b/tests/rust-python-harness/strategies/trace_parity/fixtures.py @@ -70,16 +70,17 @@ def aws_event_stream_frame(payload: Mapping[str, object]) -> bytes: def aws_event_stream_response( events: Iterable[Mapping[str, object]], *, corrupt_last_frame: bool = False ) -> RecordedHttpStreamResponse: - frames: Final = [aws_event_stream_frame(event) for event in events] - if corrupt_last_frame: - corrupted: Final = bytearray(frames[-1]) - corrupted[-1] ^= 0xFF - frames[-1] = bytes(corrupted) + frames: Final = tuple(aws_event_stream_frame(event) for event in events) + body: Final = ( + b"".join((*frames[:-1], frames[-1][:-1] + bytes((frames[-1][-1] ^ 0xFF,)))) + if corrupt_last_frame + else b"".join(frames) + ) return RecordedHttpStreamResponse( kind="http_stream", status_code=200, headers=AWS_EVENT_STREAM_HEADERS, - chunks=(RecordedStreamChunk.from_bytes(b"".join(frames)),), + chunks=(RecordedStreamChunk.from_bytes(body),), ) diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py b/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py index 522d2d0259c..3dc6d731b4b 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/chat_completions/case.py @@ -9,8 +9,6 @@ from ...models import GatewayRouteSpec, RouteFixture, TraceScenario, TraceSuite MAPPINGS: Final = ( mapping(span="python_chat_gateway_route", python_frame=r"proxy_server\.py:\d+ chat_completion$"), mapping(span="python_gateway_service", python_frame=r"ProxyBaseLLMRequestProcessing\.base_process_llm_request$"), - mapping(rust_span="chat_completions_gateway_route"), - mapping(rust_span="chat_completions"), mapping(span="python_chat_entrypoint", python_frame=r"main\.py:\d+ a?completion$"), mapping(span="python_provider_config", python_frame=r"ProviderConfigManager\.get_provider_chat_config$"), mapping(rust_span="validate_environment", python_frame=r"(? RouteFixture: TRACE_SUITE: Final = TraceSuite( - route=GatewayRouteSpec("chat_completions"), + route=GatewayRouteSpec("chat_completions", rust_supported=False), scenarios=( TraceScenario(name="async-anthropic", fixture=_fixture, mappings=MAPPINGS, asynchronous=True), TraceScenario( 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 5f24760dfe3..94d6be7cebf 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/execution.py @@ -176,8 +176,9 @@ def execute_gateway_trace( scenario: TraceScenario, engine: TraceEngine = "both", ) -> TraceArtifact: - python_trace: Final = _collect(route, scenario, "python") if engine != "rust" else () - rust_trace: Final = _collect(route, scenario, "rust") if engine != "python" else () + effective_engine: Final[TraceEngine] = "python" if engine == "both" and not route.rust_supported else engine + python_trace: Final = _collect(route, scenario, "python") if effective_engine != "rust" else () + rust_trace: Final = _collect(route, scenario, "rust") if effective_engine != "python" else () collection_python_error: Final = None if isinstance(python_trace, tuple) else f"python: {python_trace.message}" 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 () @@ -185,7 +186,7 @@ def execute_gateway_trace( python, rust, projection_error = _projections(python_events, rust_events) python_error: Final = projection_error or collection_python_error return TraceArtifact.from_traces( - engine=engine, + engine=effective_engine, surface="gateway", sdk_function=route.route, scenario=scenario.name, diff --git a/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py b/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py index 52c1ebb391f..7d805a0603c 100644 --- a/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/gateway/responses/case.py @@ -11,7 +11,6 @@ MAPPINGS: Final = ( span="python_responses_gateway_route", python_frame=r"response_api_endpoints/endpoints\.py:\d+ responses_api$" ), mapping(span="python_gateway_service", python_frame=r"ProxyBaseLLMRequestProcessing\.base_process_llm_request$"), - mapping(rust_span="responses_gateway_route"), mapping(span="python_responses", python_frame=r"responses/main\.py:\d+ a?responses$"), mapping(span="python_provider_config", python_frame=r"ProviderConfigManager\.get_provider_responses_api_config$"), mapping(rust_span="validate_environment", python_frame=r"OpenAIResponsesAPIConfig\.validate_environment$"), @@ -50,7 +49,7 @@ def _stream_fixture(engine: Engine, base_url: str) -> RouteFixture: TRACE_SUITE: Final = TraceSuite( - route=GatewayRouteSpec("responses"), + route=GatewayRouteSpec("responses", rust_supported=False), scenarios=( TraceScenario(name="async-openai", fixture=_fixture, mappings=MAPPINGS, asynchronous=True), TraceScenario( diff --git a/tests/rust-python-harness/strategies/trace_parity/models.py b/tests/rust-python-harness/strategies/trace_parity/models.py index 3701fb41e14..d6ed42250c4 100644 --- a/tests/rust-python-harness/strategies/trace_parity/models.py +++ b/tests/rust-python-harness/strategies/trace_parity/models.py @@ -48,13 +48,14 @@ class RouteFixture: class RouteSpec: route: SdkFunction python_entrypoints: tuple[str, str] - rust_entrypoints: tuple[str, str] + rust_entrypoints: tuple[str, str] | None fixture: Callable[[Engine, str], RouteFixture] @dataclass(frozen=True, slots=True) class GatewayRouteSpec: route: SdkFunction + rust_supported: bool = True TraceRouteSpec: TypeAlias = RouteSpec | GatewayRouteSpec 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 ffdbe842e85..783c22a0dc0 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/execution.py @@ -62,6 +62,8 @@ def _entrypoint(spec: RouteSpec, engine: Engine, *, asynchronous: bool) -> SdkCa from litellm.rust_bridge import get_native_bridge if engine == "rust": + if spec.rust_entrypoints is None: + return TraceExecutionFailure("rust", f"{spec.route} has no native Rust trace entrypoint") bridge: Final = cast(object | None, get_native_bridge()) if bridge is None: return TraceExecutionFailure("rust", "native Rust bridge is required for trace parity") @@ -153,6 +155,7 @@ def execute_trace( surface: Surface, engine: TraceEngine = "both", ) -> TraceArtifact: + effective_engine: Final[TraceEngine] = "python" if engine == "both" and route.rust_entrypoints is None else engine scenario_route: Final = RouteSpec( route=route.route, python_entrypoints=route.python_entrypoints, @@ -165,11 +168,13 @@ def execute_trace( "python", asynchronous=scenario.asynchronous, ) - if engine != "rust" + if effective_engine != "rust" else () ) rust_trace: Final = ( - collect_trace(scenario_route, "rust", asynchronous=scenario.asynchronous) if engine != "python" else () + collect_trace(scenario_route, "rust", asynchronous=scenario.asynchronous) + if effective_engine != "python" + else () ) python_error: Final = _failure_message(python_trace) rust_error: Final = _failure_message(rust_trace) @@ -180,7 +185,7 @@ def execute_trace( rust: Final = pipeline_projection("rust", rust_events) except ValueError as error: return TraceArtifact.from_traces( - engine=engine, + engine=effective_engine, surface=surface, sdk_function=route.route, scenario=scenario.name, @@ -189,7 +194,7 @@ def execute_trace( python_error=f"harness: {error}", ) return TraceArtifact.from_traces( - engine=engine, + engine=effective_engine, surface=surface, sdk_function=route.route, scenario=scenario.name, diff --git a/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py b/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py index 68d2eaba9bf..3f6e540efc5 100644 --- a/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py +++ b/tests/rust-python-harness/strategies/trace_parity/sdk/responses/case.py @@ -161,7 +161,7 @@ def _anthropic_bridge_stream_fixture(engine: Engine, _base_url: str) -> RouteFix ) -SPEC: Final = RouteSpec("responses", ("responses", "aresponses"), ("responses", "aresponses"), _openai_fixture) +SPEC: Final = RouteSpec("responses", ("responses", "aresponses"), None, _openai_fixture) TRACE_SUITE: Final = TraceSuite( route=SPEC, scenarios=( 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 57364057788..be25dd53b02 100644 --- a/tests/rust-python-harness/strategies/trace_parity/test_runner.py +++ b/tests/rust-python-harness/strategies/trace_parity/test_runner.py @@ -203,6 +203,44 @@ def test_gateway_trace_keeps_calls_outside_scenario_mappings(monkeypatch: pytest ) +def test_default_trace_skips_unavailable_rust_sdk_entrypoint(monkeypatch: pytest.MonkeyPatch) -> None: + execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.sdk.execution") + route: Final = RouteSpec("responses", ("responses", "aresponses"), None, _fixture) + scenario: Final = TraceScenario("sync-openai", _fixture, (), asynchronous=False) + engines: list[Engine] = [] + + def collect(_route: RouteSpec, engine: Engine, *, asynchronous: bool) -> tuple[FunctionTraceEvent, ...]: + engines.append(engine) + return (FunctionTraceEvent(0, None, "responses"),) + + monkeypatch.setattr(execution, "collect_trace", collect) + + trace: Final = execution.execute_trace(route, scenario, "sdk") + + assert engines == ["python"] + assert trace.engine == "python" + assert trace.rust_error is None + + +def test_default_trace_skips_unavailable_rust_gateway_route(monkeypatch: pytest.MonkeyPatch) -> None: + execution: Final = importlib.import_module("tests.rust-python-harness.strategies.trace_parity.gateway.execution") + route: Final = GatewayRouteSpec("responses", rust_supported=False) + scenario: Final = TraceScenario("async-openai", _fixture, (), asynchronous=True) + engines: list[Engine] = [] + + def collect(_route: GatewayRouteSpec, _scenario: TraceScenario, engine: Engine) -> tuple[FunctionTraceEvent, ...]: + engines.append(engine) + return (FunctionTraceEvent(0, None, "responses"),) + + monkeypatch.setattr(execution, "_collect", collect) + + trace: Final = execution.execute_gateway_trace(route, scenario) + + assert engines == ["python"] + assert trace.engine == "python" + assert trace.rust_error is None + + def test_scenario_validation_rejects_duplicate_and_unsafe_names() -> None: route: Final = RouteSpec("ocr", ("ocr", "aocr"), ("ocr", "aocr"), _fixture) duplicate: Final = TraceSuite(