diff --git a/switchyard/cli/launch_command.py b/switchyard/cli/launch_command.py index abb15665..196d711f 100644 --- a/switchyard/cli/launch_command.py +++ b/switchyard/cli/launch_command.py @@ -39,7 +39,10 @@ from switchyard.lib.profiles.random_routing import ( RandomRoutingConfig, ) -from switchyard.server.server_util import load_secrets, resolve_rl_log_dir +from switchyard.server.server_util import ( + load_secrets, + resolve_rl_log_dir, +) logger = logging.getLogger(__name__) diff --git a/switchyard/cli/launchers/claude_code_launcher.py b/switchyard/cli/launchers/claude_code_launcher.py index a7c28409..755cc171 100644 --- a/switchyard/cli/launchers/claude_code_launcher.py +++ b/switchyard/cli/launchers/claude_code_launcher.py @@ -46,6 +46,7 @@ from switchyard.cli.launchers.launch_intake_config import ( LaunchIntakeConfig, build_launch_capture_processors, + capture_session_headers, print_intake_warning, ) from switchyard.cli.launchers.launcher_runtime import ( @@ -68,6 +69,7 @@ from switchyard.cli.launchers.session_summary import print_session_summary from switchyard.cli.route_bundle import ( load_route_bundle_table, + route_bundle_declares_token_capture, ) from switchyard.lib.backends.llm_target import ( BackendFormat, @@ -172,6 +174,7 @@ def _claude_env( port: int, model: str, intake: LaunchIntakeConfig | None = None, + capture_headers: dict[str, str] | None = None, ) -> dict[str, str]: """Build the env-var overrides that route Claude Code through our proxy. @@ -200,6 +203,9 @@ def _claude_env( intake.opt_in_headers(), ) env["SWITCHYARD_SESSION_ID"] = intake.session_id + elif capture_headers: + env["ANTHROPIC_CUSTOM_HEADERS"] = _format_anthropic_custom_headers(capture_headers) + env["SWITCHYARD_SESSION_ID"] = capture_headers["proxy_x_session_id"] return env @@ -209,6 +215,7 @@ def _supervise_claude_plain( port: int, model: str, intake: LaunchIntakeConfig | None = None, + capture_headers: dict[str, str] | None = None, ) -> int: """Run ``claude`` via plain subprocess (non-TTY / headless fallback). @@ -216,7 +223,7 @@ def _supervise_claude_plain( ``KeyboardInterrupt`` is translated to exit code 130. """ env = os.environ.copy() - env.update(_claude_env(port, model, intake=intake)) + env.update(_claude_env(port, model, intake=intake, capture_headers=capture_headers)) try: result = subprocess.run([claude_bin, *claude_args], env=env, check=False) return result.returncode @@ -292,6 +299,7 @@ def _run_claude_with_switchyard( stats: StatsAccumulator, intake: LaunchIntakeConfig | None = None, strategy_summary: str | None = None, + capture_headers: dict[str, str] | None = None, ) -> int: """Chain-agnostic supervisor: host ``switchyard`` then spawn claude. @@ -377,6 +385,7 @@ def _run_claude_with_switchyard( resolved_port, display_model, intake=intake, + capture_headers=capture_headers, ) logger.debug( "claude env ANTHROPIC_BASE_URL=%s ANTHROPIC_MODEL=%s " @@ -401,6 +410,7 @@ def _run_claude_with_switchyard( return _supervise_claude_plain( claude_bin, claude_args, resolved_port, display_model, intake=intake, + capture_headers=capture_headers, ) finally: print_session_summary(stats) @@ -438,7 +448,12 @@ def launch_claude( """ _quiet_launch_loggers() stats = StatsAccumulator() - intake_request, intake_response = build_launch_capture_processors(intake, rl_log_dir) + # Token capture activates when RL logging is on AND the bundle declares + # token_capture_engine on a route; the capture pair replaces the rl-logging pair. + capture = rl_log_dir is not None and route_bundle_declares_token_capture(routing_profiles) + intake_request, intake_response = build_launch_capture_processors( + intake, rl_log_dir, token_capture=capture, + ) switchyard = _build_claude_switchyard( model=model, api_key=api_key, @@ -461,6 +476,7 @@ def launch_claude( stats_accumulator=stats, pre_routing_request_processors=intake_request, extra_response_processors=intake_response, + token_capture_enabled=rl_log_dir is not None, ) for sub_model, sub_chain, sub_metadata in yaml_table.items(): table.register(sub_model, sub_chain, metadata=sub_metadata) @@ -480,6 +496,7 @@ def launch_claude( stats=stats, intake=intake, strategy_summary=strategy_summary, + capture_headers=capture_session_headers(intake, capture, "claude"), ) @@ -513,7 +530,9 @@ def _discovery_fn(base_url: str, api_key: str) -> list[str]: _quiet_launch_loggers() stats = StatsAccumulator() - intake_request, intake_response = build_launch_capture_processors(intake, rl_log_dir) + intake_request, intake_response = build_launch_capture_processors( + intake, rl_log_dir, + ) switchyard = build_deterministic_routing_switchyard( config, stats, @@ -539,4 +558,5 @@ def _discovery_fn(base_url: str, api_key: str) -> list[str]: stats=stats, intake=intake, strategy_summary=deterministic_strategy_summary(config), + capture_headers=capture_session_headers(intake, False, "claude"), ) diff --git a/switchyard/cli/launchers/codex_cli_launcher.py b/switchyard/cli/launchers/codex_cli_launcher.py index 63605c90..ce78eb9c 100644 --- a/switchyard/cli/launchers/codex_cli_launcher.py +++ b/switchyard/cli/launchers/codex_cli_launcher.py @@ -47,6 +47,7 @@ from switchyard.cli.launchers.launch_intake_config import ( LaunchIntakeConfig, build_launch_capture_processors, + capture_session_headers, print_intake_warning, ) from switchyard.cli.launchers.launcher_runtime import ( @@ -69,6 +70,7 @@ from switchyard.cli.launchers.session_summary import print_session_summary from switchyard.cli.route_bundle import ( load_route_bundle_table, + route_bundle_declares_token_capture, ) from switchyard.lib.backends.llm_target import BackendFormat, LlmTarget from switchyard.lib.processors.model_rewrite_request_processor import ( @@ -222,6 +224,7 @@ def _provider_overrides( port: int, *, intake: LaunchIntakeConfig | None = None, model_catalog_json: str | None = None, + capture_headers: dict[str, str] | None = None, ) -> list[str]: """Build the ``-c key=value`` argv pairs that point codex at the proxy. @@ -264,6 +267,11 @@ def _provider_overrides( overrides.extend([ "-c", f"model_providers.{_PROVIDER_ID}.http_headers={headers_toml}", ]) + elif capture_headers: + headers_toml = _format_codex_http_headers(capture_headers) + overrides.extend([ + "-c", f"model_providers.{_PROVIDER_ID}.http_headers={headers_toml}", + ]) return overrides @@ -283,11 +291,17 @@ def _codex_command( model: str, intake: LaunchIntakeConfig | None = None, model_catalog_json: str | None = None, + capture_headers: dict[str, str] | None = None, ) -> list[str]: """Build the exact Codex argv for the transient Switchyard provider.""" return [ codex_bin, - *_provider_overrides(port, intake=intake, model_catalog_json=model_catalog_json), + *_provider_overrides( + port, + intake=intake, + model_catalog_json=model_catalog_json, + capture_headers=capture_headers, + ), "-m", model, *codex_args, @@ -301,6 +315,7 @@ def _supervise_codex( model: str, intake: LaunchIntakeConfig | None = None, model_catalog_json: str | None = None, + capture_headers: dict[str, str] | None = None, ) -> int: """Run ``codex`` with proxy provider injected; return its exit code. @@ -335,6 +350,7 @@ def _supervise_codex( model, intake=intake, model_catalog_json=model_catalog_json, + capture_headers=capture_headers, ), env=_codex_env(intake=intake), check=False, @@ -353,6 +369,7 @@ def _run_codex_with_switchyard( intake: LaunchIntakeConfig | None = None, codex_model_catalog: Sequence[CodexModelCatalogEntry] = (), strategy_summary: str | None = None, + capture_headers: dict[str, str] | None = None, ) -> int: """Chain-agnostic supervisor: host ``switchyard`` then spawn codex.""" codex_bin = _find_codex_binary() @@ -413,6 +430,7 @@ def _run_codex_with_switchyard( display_model, intake=intake, model_catalog_json=model_catalog_json, + capture_headers=capture_headers, ), footer_fn=footer.as_footer_fn(), footer_height=lambda: footer.height, @@ -422,6 +440,7 @@ def _run_codex_with_switchyard( return _supervise_codex( codex_bin, codex_args, resolved_port, display_model, intake=intake, model_catalog_json=model_catalog_json, + capture_headers=capture_headers, ) finally: print_session_summary(stats) @@ -461,7 +480,12 @@ def launch_codex( binary wasn't found, ``130`` on Ctrl-C). """ stats = StatsAccumulator() - intake_request, intake_response = build_launch_capture_processors(intake, rl_log_dir) + # Token capture activates when RL logging is on AND the bundle declares + # token_capture_engine on a route; the capture pair replaces the rl-logging pair. + capture = rl_log_dir is not None and route_bundle_declares_token_capture(routing_profiles) + intake_request, intake_response = build_launch_capture_processors( + intake, rl_log_dir, token_capture=capture, + ) switchyard = _build_switchyard( model, api_key, @@ -493,6 +517,7 @@ def launch_codex( stats_accumulator=stats, pre_routing_request_processors=intake_request, extra_response_processors=intake_response, + token_capture_enabled=rl_log_dir is not None, ) for sub_model, sub_chain, sub_metadata in yaml_table.items(): table.register(sub_model, sub_chain, metadata=sub_metadata) @@ -524,6 +549,7 @@ def launch_codex( intake=intake, codex_model_catalog=codex_model_catalog, strategy_summary=strategy_summary, + capture_headers=capture_session_headers(intake, capture, "codex"), ) @@ -589,7 +615,9 @@ def _discovery_fn(base_url: str, api_key: str) -> list[str]: return fetch_model_ids(base_url, api_key) stats = StatsAccumulator() - intake_request, intake_response = build_launch_capture_processors(intake, rl_log_dir) + intake_request, intake_response = build_launch_capture_processors( + intake, rl_log_dir, + ) switchyard = build_deterministic_routing_switchyard( config, stats, @@ -639,4 +667,5 @@ def _discovery_fn(base_url: str, api_key: str) -> list[str]: intake=intake, codex_model_catalog=codex_model_catalog, strategy_summary=deterministic_strategy_summary(config), + capture_headers=capture_session_headers(intake, False, "codex"), ) diff --git a/switchyard/cli/launchers/launch_intake_config.py b/switchyard/cli/launchers/launch_intake_config.py index dbefbff9..bcd1cb15 100644 --- a/switchyard/cli/launchers/launch_intake_config.py +++ b/switchyard/cli/launchers/launch_intake_config.py @@ -156,20 +156,70 @@ def build_intake_processors( def build_launch_capture_processors( intake: LaunchIntakeConfig | None, rl_log_dir: Path | None, + token_capture: bool = False, ) -> tuple[list[Any], list[Any]]: - """Combine intake + RL-logging request/response processor lists for launchers. + """Combine intake + RL-logging / token-capture processor lists for launchers. The intake sink (``intake``) and the local RL trace logger (``rl_log_dir``) - are independent: either, both, or neither may be active. Returns - ``([], [])`` when neither is. + are independent: any subset may be active. When ``token_capture`` is set + (the route bundle declares ``token_capture_engine`` on a route and RL logging is + on), the capture pair replaces the rl-logging pair — the unified record + carries the trace fields, so running both would double-write. Returns + ``([], [])`` when nothing is active. """ from switchyard.lib.processors.rl_logging_response_processor import ( build_rl_logging_processors, ) + from switchyard.lib.processors.token_capture_response_processor import ( + build_token_capture_processors, + ) intake_request, intake_response = build_intake_processors(intake) - rl_request, rl_response = build_rl_logging_processors(rl_log_dir) - return [*intake_request, *rl_request], [*intake_response, *rl_response] + if token_capture: + log_request, log_response = build_token_capture_processors(rl_log_dir) + else: + log_request, log_response = build_rl_logging_processors(rl_log_dir) + return ( + [*intake_request, *log_request], + [*intake_response, *log_response], + ) + + +def capture_session_headers( + intake: LaunchIntakeConfig | None, + capture_enabled: bool, + target: str, +) -> dict[str, str]: + """Per-launch session header for token capture when intake is off. + + Intake's opt-in headers already carry a per-launch session id; this covers + token capture without intake, so captured records group per launch instead + of going uncaptured for lack of a session. Returns ``{}`` when intake is + active (its headers win) or capture is disabled. + """ + if intake is not None or not capture_enabled: + return {} + return {"proxy_x_session_id": _default_session_id(target)} + + +def capture_session_id_for_api_key( + intake: LaunchIntakeConfig | None, + capture_enabled: bool, + target: str, +) -> str | None: + """Per-launch session id for harnesses that ride it on the API-key field. + + OpenClaw has no custom-header surface, so intake's opt-in headers never + reach it — unlike :func:`capture_session_headers`, intake being active + does not cover session delivery here. With capture on, the session id + (intake's when present, else a fresh one) is substituted for the harness's + opaque API-key placeholder and resolved proxy-side from the caller key. + """ + if not capture_enabled: + return None + if intake is not None: + return intake.session_id + return _default_session_id(target) def _default_session_id(target: str) -> str: diff --git a/switchyard/cli/launchers/openclaw_launcher.py b/switchyard/cli/launchers/openclaw_launcher.py index 1eb70610..7bf5fafb 100644 --- a/switchyard/cli/launchers/openclaw_launcher.py +++ b/switchyard/cli/launchers/openclaw_launcher.py @@ -56,6 +56,7 @@ from switchyard.cli.launchers.launch_intake_config import ( LaunchIntakeConfig, build_launch_capture_processors, + capture_session_id_for_api_key, print_intake_warning, ) from switchyard.cli.launchers.launcher_runtime import ( @@ -77,11 +78,13 @@ from switchyard.cli.launchers.proxy_health_monitor import ProxyHealthMonitor from switchyard.cli.route_bundle import ( load_route_bundle_table, + route_bundle_declares_token_capture, ) from switchyard.lib.backends.llm_target import BackendFormat, LlmTarget from switchyard.lib.processors.model_rewrite_request_processor import ( ModelRewriteRequestProcessor, ) +from switchyard.lib.processors.token_capture_request_processor import SESSION_API_KEY_MARKER from switchyard.lib.profiles import ( DeterministicRoutingConfig, ) @@ -318,6 +321,7 @@ def _remove_openclaw_workspace(workspace: str | None) -> None: def _openclaw_env( workspace: str, intake: LaunchIntakeConfig | None = None, + capture_session_id: str | None = None, ) -> dict[str, str]: """Environment that points OpenClaw at the transient workspace. @@ -341,7 +345,15 @@ def _openclaw_env( env["OPENCLAW_HOME"] = workspace env["OPENCLAW_CONFIG_PATH"] = str(Path(workspace) / "openclaw.json") env["OPENCLAW_HIDE_BANNER"] = "1" - env[_API_KEY_ENV] = _API_KEY_PLACEHOLDER + # OpenClaw has no custom-header surface; when token capture needs a + # session id, ride the apiKey field instead — OpenClaw echoes it as the + # Authorization bearer, and the capture processor recognizes the marker + # prefix to recover the session id (any id, not only a fixed shape). + if capture_session_id: + env[_API_KEY_ENV] = SESSION_API_KEY_MARKER + capture_session_id + env["SWITCHYARD_SESSION_ID"] = capture_session_id + else: + env[_API_KEY_ENV] = _API_KEY_PLACEHOLDER if intake is not None: env["SWITCHYARD_SESSION_ID"] = intake.session_id return env @@ -373,6 +385,7 @@ def _supervise_openclaw( openclaw_args: list[str], workspace: str, intake: LaunchIntakeConfig | None = None, + capture_session_id: str | None = None, ) -> int: """Run ``openclaw chat`` with the transient workspace; return exit code. @@ -383,7 +396,10 @@ def _supervise_openclaw( try: result = subprocess.run( _openclaw_command(openclaw_bin, openclaw_args), - env=_openclaw_env(workspace=workspace, intake=intake), + env=_openclaw_env( + workspace=workspace, intake=intake, + capture_session_id=capture_session_id, + ), check=False, ) return result.returncode @@ -400,6 +416,7 @@ def _run_openclaw_with_switchyard( catalog_entries: Sequence[OpenClawModelCatalogEntry], intake: LaunchIntakeConfig | None = None, strategy_summary: str | None = None, + capture_session_id: str | None = None, ) -> int: """Chain-agnostic supervisor: host ``switchyard`` then spawn openclaw.""" openclaw_bin = _find_openclaw_binary() @@ -461,12 +478,16 @@ def _run_openclaw_with_switchyard( command=_openclaw_command(openclaw_bin, openclaw_args), footer_fn=footer.as_footer_fn(), footer_height=lambda: footer.height, - env=_openclaw_env(workspace=workspace, intake=intake), + env=_openclaw_env( + workspace=workspace, intake=intake, + capture_session_id=capture_session_id, + ), ).run() return _supervise_openclaw( openclaw_bin, openclaw_args, workspace=workspace, intake=intake, + capture_session_id=capture_session_id, ) finally: if server is not None: @@ -524,7 +545,12 @@ def launch_openclaw( binary wasn't found, ``130`` on Ctrl-C). """ stats = StatsAccumulator() - intake_request, intake_response = build_launch_capture_processors(intake, rl_log_dir) + # Token capture activates when RL logging is on AND the bundle declares + # token_capture_engine on a route; the capture pair replaces the rl-logging pair. + capture = rl_log_dir is not None and route_bundle_declares_token_capture(routing_profiles) + intake_request, intake_response = build_launch_capture_processors( + intake, rl_log_dir, token_capture=capture, + ) switchyard = _build_switchyard( model, api_key, @@ -556,6 +582,7 @@ def launch_openclaw( stats_accumulator=stats, pre_routing_request_processors=intake_request, extra_response_processors=intake_response, + token_capture_enabled=rl_log_dir is not None, ) for sub_model, sub_chain, sub_metadata in yaml_table.items(): table.register(sub_model, sub_chain, metadata=sub_metadata) @@ -585,6 +612,7 @@ def launch_openclaw( catalog_entries=catalog_entries, intake=intake, strategy_summary=strategy_summary, + capture_session_id=capture_session_id_for_api_key(intake, capture, "openclaw"), ) @@ -649,7 +677,9 @@ def _discovery_fn(base_url: str, api_key: str) -> list[str]: return fetch_model_ids(base_url, api_key) stats = StatsAccumulator() - intake_request, intake_response = build_launch_capture_processors(intake, rl_log_dir) + intake_request, intake_response = build_launch_capture_processors( + intake, rl_log_dir, + ) switchyard = build_deterministic_routing_switchyard( config, stats, @@ -699,4 +729,5 @@ def _discovery_fn(base_url: str, api_key: str) -> list[str]: catalog_entries=catalog_entries, intake=intake, strategy_summary=deterministic_strategy_summary(config), + capture_session_id=capture_session_id_for_api_key(intake, False, "openclaw"), ) diff --git a/switchyard/cli/route_bundle.py b/switchyard/cli/route_bundle.py index c41a13a8..fdf00a9a 100644 --- a/switchyard/cli/route_bundle.py +++ b/switchyard/cli/route_bundle.py @@ -22,6 +22,7 @@ argparse-driven and YAML-driven assembly share one builder. """ +import logging import os import re from collections.abc import Mapping, Sequence @@ -58,6 +59,8 @@ ) from switchyard.lib.stats_accumulator import StatsAccumulator +logger = logging.getLogger(__name__) + def _default_discovery_fn(base_url: str, api_key: str) -> list[str]: """Wrap provider catalog fetching for the lib-level callable shape. @@ -164,6 +167,9 @@ def llm_target_to_route_dict(target: LlmTarget) -> dict[str, Any]: _MODEL_ROUTE_KEYS = _ROUTE_METADATA_KEYS | _TARGET_DEFAULT_ROUTE_KEYS | frozenset({ "target", "model", + # Opts this route's single target into token capture. Applied to the target + # before coercion. + "token_capture_engine", }) _RANDOM_ROUTING_ROUTE_KEYS = _ROUTE_METADATA_KEYS | _TARGET_DEFAULT_ROUTE_KEYS | frozenset({ "strong", @@ -191,7 +197,7 @@ def llm_target_to_route_dict(target: LlmTarget) -> dict[str, Any]: }) _PASSTHROUGH_ROUTE_KEYS = ( _ROUTE_METADATA_KEYS - | frozenset({"defaults", "enable_stats"}) + | frozenset({"defaults", "enable_stats", "token_capture_engine"}) | _PASSTHROUGH_SETTING_KEYS ) _LATENCY_ENDPOINT_DEFAULT_KEYS = frozenset({ @@ -418,6 +424,7 @@ def load_route_bundle_table( stats_accumulator: StatsAccumulator | None = None, pre_routing_request_processors: Sequence[Any] = (), extra_response_processors: Sequence[Any] = (), + token_capture_enabled: bool = False, ) -> RouteTable: """Load *path* and return a table keyed by route model id. @@ -430,12 +437,17 @@ def load_route_bundle_table( ``pre_routing_request_processors`` / ``extra_response_processors`` let callers attach process-level components such as Intake telemetry to every YAML-declared route. + + ``token_capture_enabled`` honors routes' ``token_capture_engine`` declarations + (deriving token-capture request params); when ``False`` the key is ignored + with a warning so clients never receive token fields nobody will capture. """ return build_route_bundle_table( parse_routing_profiles_file(path), stats_accumulator=stats_accumulator, pre_routing_request_processors=pre_routing_request_processors, extra_response_processors=extra_response_processors, + token_capture_enabled=token_capture_enabled, ) @@ -444,6 +456,7 @@ def build_route_bundle_table( stats_accumulator: StatsAccumulator | None = None, pre_routing_request_processors: Sequence[Any] = (), extra_response_processors: Sequence[Any] = (), + token_capture_enabled: bool = False, ) -> RouteTable: """Parse *raw* dict and build a :class:`RouteTable`. @@ -452,6 +465,8 @@ def build_route_bundle_table( then delegates to :func:`build_table_from_bundle`. Optional ``stats_accumulator`` overrides the parsed bundle's default ``None``. """ + if not token_capture_enabled: + raw = _strip_token_capture_engine(raw) bundle = _parse_route_bundle_dict(raw) if stats_accumulator is not None: bundle = replace(bundle, stats_accumulator=stats_accumulator) @@ -462,6 +477,71 @@ def build_route_bundle_table( ) +def route_bundle_declares_token_capture(raw_or_path: object) -> bool: + """True when the bundle declares ``token_capture_engine`` on any route. + + Accepts either a path to a routing-profiles YAML file or an + already-parsed bundle dict. Token capture activates under + ``--enable-rl-logging`` when at least one route opts in via + ``token_capture_engine``. A bundle that cannot be parsed reports ``False`` — + the real table load surfaces the parse error. + """ + if raw_or_path is None: + return False + raw: object + if isinstance(raw_or_path, str | Path): + try: + raw = parse_routing_profiles_file(raw_or_path) + except RouteBundleConfigError: + return False + else: + raw = raw_or_path + + def _walk(node: object) -> bool: + if isinstance(node, Mapping): + return "token_capture_engine" in node or any(_walk(value) for value in node.values()) + if isinstance(node, list): + return any(_walk(item) for item in node) + return False + + return _walk(raw) + + +def _strip_token_capture_engine(raw: object) -> object: + """Drop ``token_capture_engine`` keys when token capture is disabled. + + The field derives token-capture request params into the route target's + ``extra_body`` (see ``llm_target_with_token_capture``); honoring it while + RL logging is off would make clients receive token fields that nothing + captures. Warns once so the config mismatch is visible. + """ + + stripped = 0 + + def _walk(node: object) -> object: + nonlocal stripped + if isinstance(node, Mapping): + out: dict[str, object] = {} + for key, value in node.items(): + if key == "token_capture_engine": + stripped += 1 + continue + out[key] = _walk(value) + return out + if isinstance(node, list): + return [_walk(item) for item in node] + return node + + result = _walk(raw) + if stripped: + logger.warning( + "route bundle declares token_capture_engine on %d route(s) but token " + "capture is disabled; ignoring (enable with --enable-rl-logging)", + stripped, + ) + return result + + def _parse_route_bundle_dict(raw: object) -> RouteBundle: """Validate *raw* and return a :class:`RouteBundle`. @@ -1129,10 +1209,12 @@ def _passthrough_target( if route_type == "model": raise RouteBundleConfigError(f"route {model_id!r} requires target or model") target_raw = model_id - return coerce_llm_target( - _target_mapping(target_raw, target_defaults, default_id=model_id, where="target"), - default_id=model_id, - ) + target = _target_mapping(target_raw, target_defaults, default_id=model_id, where="target") + # Apply the route's token_capture_engine to its single target. + route_engine = route.get("token_capture_engine") + if route_engine is not None: + target["token_capture_engine"] = route_engine + return coerce_llm_target(target, default_id=model_id) def _latency_service_switchyard( @@ -1506,5 +1588,6 @@ def _optional_bool(value: object, default: bool) -> bool: "build_route_bundle_table", "load_route_bundle_table", "parse_routing_profiles_file", + "route_bundle_declares_token_capture", "routing_profile_model_ids", ] diff --git a/switchyard/cli/switchyard_cli.py b/switchyard/cli/switchyard_cli.py index 82e7acbe..7e681f95 100644 --- a/switchyard/cli/switchyard_cli.py +++ b/switchyard/cli/switchyard_cli.py @@ -88,11 +88,15 @@ from switchyard.cli.route_bundle import ( RouteBundleConfigError, load_route_bundle_table, + route_bundle_declares_token_capture, ) from switchyard.lib.config import IntakeSinkConfig from switchyard.lib.processors.intake_request_processor import IntakeRequestProcessor from switchyard.lib.processors.intake_response_processor import IntakeResponseProcessor from switchyard.lib.processors.rl_logging_response_processor import build_rl_logging_processors +from switchyard.lib.processors.token_capture_response_processor import ( + build_token_capture_processors, +) from switchyard.server.server_util import ( DEFAULT_SECRETS_FILE, add_transport_args, @@ -413,27 +417,40 @@ def _cmd_serve(args: argparse.Namespace) -> None: _cmd_serve_profile_config(args) return - # Intake sink + local RL trace logging both attach as chain processors; - # combine them so a single serve invocation can run either or both. + # Intake sink and local RL trace logging attach as chain processors. + # Token capture activates when RL logging is on AND the bundle declares + # token_capture_engine on a route; the capture pair then replaces the + # rl-logging pair (the unified record carries the trace fields — running + # both would double-write). + saved = None + if not routing_profiles: + saved = load_user_config().routing_profiles + if not saved: + raise SystemExit( + "serve: no routing-profiles given. Pass --routing-profiles " + "PATH or run `switchyard configure --routing-profiles PATH` " + "to save one." + ) intake_request, intake_response = _resolve_intake_processors(args) - rl_request, rl_response = build_rl_logging_processors(resolve_rl_log_dir(args)) - request_processors = [*intake_request, *rl_request] - response_processors = [*intake_response, *rl_response] + rl_log_dir = resolve_rl_log_dir(args) + capture = rl_log_dir is not None and route_bundle_declares_token_capture( + routing_profiles or saved + ) + if capture: + log_request, log_response = build_token_capture_processors(rl_log_dir) + else: + log_request, log_response = build_rl_logging_processors(rl_log_dir) + request_processors = [*intake_request, *log_request] + response_processors = [*intake_response, *log_response] if routing_profiles: table = load_route_bundle_table( routing_profiles, pre_routing_request_processors=request_processors, extra_response_processors=response_processors, + token_capture_enabled=rl_log_dir is not None, ) source = routing_profiles else: - saved = load_user_config().routing_profiles - if not saved: - raise SystemExit( - "serve: no routing-profiles given. Pass --routing-profiles " - "PATH or run `switchyard configure --routing-profiles PATH` " - "to save one." - ) # Saved bundles are stored as parsed dicts, so we can skip the # YAML parse step and feed them straight into the dict-driven # entrypoint. Env-var references inside the dict expand inside @@ -443,6 +460,7 @@ def _cmd_serve(args: argparse.Namespace) -> None: saved, pre_routing_request_processors=request_processors, extra_response_processors=response_processors, + token_capture_enabled=rl_log_dir is not None, ) source = f"" logger.info( @@ -889,9 +907,11 @@ def _build_parser() -> argparse.ArgumentParser: action="store_true", help=( "Write per-turn RL training traces (message_history JSON, one " - "file per request/response pair) for `launch` sessions. Global " - "flag — place it before the subcommand, e.g. " - "switchyard --enable-rl-logging launch claude." + "file per request/response pair) for `launch` and `serve` sessions. " + "When the route bundle declares `token_capture_engine` on a route, " + "traces are token-level capture records grouped by session and " + "served via GET /v1/sessions. Global flag — place it before the " + "subcommand, e.g. switchyard --enable-rl-logging launch claude." ), ) parser.add_argument( diff --git a/switchyard/lib/backends/llm_target.py b/switchyard/lib/backends/llm_target.py index ad22d073..a0ab29fa 100644 --- a/switchyard/lib/backends/llm_target.py +++ b/switchyard/lib/backends/llm_target.py @@ -45,12 +45,13 @@ def coerce_llm_target(value: object, *, default_id: str) -> LlmTarget: timeout_secs = data.pop("timeout_secs", data.pop("timeout", None)) extra_body = data.pop("extra_body", None) extra_headers = data.pop("extra_headers", None) + token_capture_engine = data.pop("token_capture_engine", None) data.pop("tuning", None) if data: unknown = ", ".join(sorted(data)) raise ValueError(f"unknown LlmTarget fields: {unknown}") - return LlmTarget( + target = LlmTarget( id=target_id, model=model, format=target_format, @@ -61,6 +62,11 @@ def coerce_llm_target(value: object, *, default_id: str) -> LlmTarget: extra_body=extra_body, extra_headers=extra_headers, ) + if token_capture_engine is not None: + # Declaring the serving engine opts the target into token-capture + # request params (pairs with `--enable-rl-logging` on the server). + target = llm_target_with_token_capture(target, str(token_capture_engine)) + return target def llm_target_with_format(target: LlmTarget, target_format: BackendFormat) -> LlmTarget: @@ -96,6 +102,39 @@ def _should_disable_thinking(model: str) -> bool: return any(fragment in normalized for fragment in _DISABLE_THINKING_MODEL_FRAGMENTS) +_TOKEN_CAPTURE_ENGINE_PARAMS: dict[str, dict[str, Any]] = { + # vLLM: `return_token_ids` emits `response.prompt_token_ids` + `choice.token_ids`; + # `top_logprobs` must be non-None for `logprobs.content[]` to populate given + # `logprobs=true` (0 returns just the sampled token's logprob). + "vllm": {"logprobs": True, "return_token_ids": True, "top_logprobs": 0}, +} + + +def llm_target_with_token_capture(target: LlmTarget, token_capture_engine: str) -> LlmTarget: + """Return ``target`` with the engine's token-capture request params in ``extra_body``. + + These params must ride ``extra_body`` (merged after request translation in the + backend) — request-side injection is dropped by the translation allowlist for + cross-format traffic. Explicit ``extra_body`` keys on the target win over derived + params; harness-side conflicts are resolved caller-wins at request time. + """ + params = _TOKEN_CAPTURE_ENGINE_PARAMS.get(token_capture_engine) + if params is None: + supported = ", ".join(sorted(_TOKEN_CAPTURE_ENGINE_PARAMS)) + raise ValueError( + f"unknown token_capture_engine {token_capture_engine!r} for token " + f"capture; supported: {supported}" + ) + return LlmTarget( + id=target.id, + model=target.model, + format=target.format, + endpoint=target.endpoint, + extra_body={**params, **(target.extra_body or {})}, + extra_headers=target.extra_headers, + ) + + __all__ = [ "BackendFormat", "EndpointConfig", @@ -103,4 +142,5 @@ def _should_disable_thinking(model: str) -> bool: "coerce_llm_target", "llm_target_with_format", "llm_target_with_runtime_defaults", + "llm_target_with_token_capture", ] diff --git a/switchyard/lib/endpoints/token_capture_endpoint.py b/switchyard/lib/endpoints/token_capture_endpoint.py new file mode 100644 index 00000000..26172aeb --- /dev/null +++ b/switchyard/lib/endpoints/token_capture_endpoint.py @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""FastAPI endpoint exposing captured token-level completion records. + +Serves the on-disk record store written by +:class:`~switchyard.lib.processors.token_capture_response_processor.TokenCaptureResponseProcessor` +via ``GET /v1/sessions/{session_id}/completions`` — one session's records in +capture order. + +Reads go straight to disk (no in-memory registry), so retrieval works across +uvicorn workers and after a proxy restart. +""" + +from __future__ import annotations + +import json +import logging +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from fastapi import APIRouter, HTTPException + +from switchyard.lib.endpoints.base import Endpoint as NemoSwitchyardEndpoint +from switchyard.lib.processors.token_capture_response_processor import ( + SCHEMA_VERSION, + session_dir_name, + sessions_root, +) + +if TYPE_CHECKING: + from fastapi import FastAPI + +logger = logging.getLogger(__name__) + + +def _session_head(session_dir: Path) -> dict[str, Any] | None: + """``{session_id, parent_session_id}`` from the first readable record in *session_dir*. + + ``None`` when the directory holds no readable, session-tagged record. + """ + for record_path in sorted(session_dir.glob("*.json")): + try: + record = json.loads(record_path.read_text()) + except (OSError, json.JSONDecodeError): + continue + session_id = record.get("session_id") + if isinstance(session_id, str) and session_id: + return { + "session_id": session_id, + "parent_session_id": record.get("parent_session_id"), + } + return None + + +class TokenCaptureSessionsEndpoint(NemoSwitchyardEndpoint): + """Exposes captured completion records for retrieval by session id. + + Contributed automatically by + :meth:`TokenCaptureResponseProcessor.get_endpoint` — no manual wiring + required. + """ + + register_once = True + + def __init__(self, capture_dir: Path | str) -> None: + self._capture_dir = Path(capture_dir) + + def register(self, app: FastAPI) -> None: + routes = APIRouter() + capture_dir = self._capture_dir + + async def get_session_completions(session_id: str) -> dict[str, Any]: + """All captured records for one session, in capture order.""" + # The directory name is a hash of the session id (path-safe by + # construction, no traversal possible). + session_dir = sessions_root(capture_dir) / session_dir_name(session_id) + if not session_dir.is_dir(): + raise HTTPException(status_code=404, detail="unknown session") + + completions: list[dict[str, Any]] = [] + for record_path in sorted(session_dir.glob("*.json")): + try: + record = json.loads(record_path.read_text()) + except (OSError, json.JSONDecodeError) as exc: + # A torn record must not hide the rest of the session. + logger.warning( + "token capture: skipping unreadable record %s: %s", + record_path, + exc, + ) + continue + # Defense in depth: only surface records that actually belong to + # the requested session, never a neighbor sharing the directory. + if record.get("session_id") == session_id: + completions.append(record) + completions.sort( + key=lambda record: (record.get("captured_at", ""), record.get("uuid", "")) + ) + return { + "schema_version": SCHEMA_VERSION, + "session_id": session_id, + "completions": completions, + } + + async def list_sessions() -> dict[str, Any]: + """Every session id captured under this log dir, with its parent link. + + Lets a caller enumerate a run's sessions when the harness mints its + own ids (e.g. OpenCode's ``X-Session-Id``), without reading any + harness logs. Directory names are session-id hashes, so the ids come + from the records. Scope one run to one ``--rl-log-dir`` so the list is + not mixed across rollouts. + """ + root = sessions_root(capture_dir) + sessions: list[dict[str, Any]] = [] + if root.is_dir(): + for session_dir in sorted(root.iterdir()): + entry = _session_head(session_dir) if session_dir.is_dir() else None + if entry is not None: + sessions.append(entry) + return {"schema_version": SCHEMA_VERSION, "sessions": sessions} + + routes.get("/v1/sessions")(list_sessions) + routes.get("/v1/sessions/{session_id}/completions")(get_session_completions) + app.include_router(routes, tags=["TokenCapture"]) diff --git a/switchyard/lib/processors/rl_logging_response_processor.py b/switchyard/lib/processors/rl_logging_response_processor.py index 529a8c1c..f5cb246a 100644 --- a/switchyard/lib/processors/rl_logging_response_processor.py +++ b/switchyard/lib/processors/rl_logging_response_processor.py @@ -98,28 +98,12 @@ def _build_entry(self, request: JsonObject, response: ChatResponse) -> JsonObjec if not isinstance(message, dict): return None - messages = [dict(m) for m in request.get("messages", []) if isinstance(m, dict)] - assistant: JsonObject = {"role": "assistant"} - content = message.get("content") - if content is not None: - assistant["content"] = content - tool_calls = message.get("tool_calls") - if tool_calls: - assistant["tool_calls"] = tool_calls - messages.append(assistant) - - usage = body.get("usage") - usage = usage if isinstance(usage, dict) else {} return { "uuid": str(uuid_lib.uuid4()), - "messages": messages, - "tools": _format_tools(request.get("tools", [])), - "tool_choice": _format_tool_choice(request.get("tool_choice")), - "token_count": { - "prompt_tokens": usage.get("prompt_tokens", 0), - "completion_tokens": usage.get("completion_tokens", 0), - "total_tokens": usage.get("total_tokens", 0), - }, + "messages": build_trace_messages(request, message), + "tools": format_trace_tools(request.get("tools", [])), + "tool_choice": format_trace_tool_choice(request.get("tool_choice")), + "token_count": build_trace_token_count(body.get("usage")), "is_valid": True, } @@ -151,7 +135,35 @@ def _trace_filename() -> str: return f"{timestamp}_trace_{trace_id}_{suffix_id}.json" -def _format_tools(raw_tools: object) -> list[JsonObject]: +def build_trace_messages(request: JsonObject, message: JsonObject) -> list[JsonObject]: + """Request-snapshot history plus the appended assistant turn. + + Shared by the RL trace logger and the token-capture record writer so the + text-trace shape stays identical across both outputs. + """ + messages = [dict(m) for m in request.get("messages", []) if isinstance(m, dict)] + assistant: JsonObject = {"role": "assistant"} + content = message.get("content") + if content is not None: + assistant["content"] = content + tool_calls = message.get("tool_calls") + if tool_calls: + assistant["tool_calls"] = tool_calls + messages.append(assistant) + return messages + + +def build_trace_token_count(usage: object) -> JsonObject: + """Normalize a response ``usage`` block into the trace ``token_count`` shape.""" + usage = usage if isinstance(usage, dict) else {} + return { + "prompt_tokens": usage.get("prompt_tokens", 0), + "completion_tokens": usage.get("completion_tokens", 0), + "total_tokens": usage.get("total_tokens", 0), + } + + +def format_trace_tools(raw_tools: object) -> list[JsonObject]: """Port the V1 message_history tool shape: ``{id, description, inputSchema}``.""" if not isinstance(raw_tools, list): return [] @@ -177,7 +189,8 @@ def _format_tools(raw_tools: object) -> list[JsonObject]: return tools -def _format_tool_choice(tool_choice: object) -> str: +def format_trace_tool_choice(tool_choice: object) -> str: + """Extract the tool-choice ``type`` string, defaulting to ``"auto"``.""" if isinstance(tool_choice, str): return tool_choice if isinstance(tool_choice, dict): diff --git a/switchyard/lib/processors/token_capture_request_processor.py b/switchyard/lib/processors/token_capture_request_processor.py new file mode 100644 index 00000000..5ad5cc7f --- /dev/null +++ b/switchyard/lib/processors/token_capture_request_processor.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Request-side processor that resolves the capture session and preps the upstream call.""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping + +from switchyard.lib.proxy_context import CTX_CALLER_API_KEY, ProxyContext +from switchyard.lib.request_metadata import CTX_PROFILE_REQUEST_HEADERS, CTX_REQUEST_METADATA +from switchyard_rust.core import ChatRequest + +logger = logging.getLogger(__name__) + +#: Context metadata key holding the resolved capture session id. Written by +#: :class:`TokenCaptureRequestProcessor`; read by +#: :class:`~switchyard.lib.processors.token_capture_response_processor.TokenCaptureResponseProcessor` +#: — a call without it is forwarded uncaptured. +CTX_TOKEN_CAPTURE_SESSION = "_token_capture_session" + +#: Context metadata key holding the parent session id for a subagent call, from +#: the harness's native ``x-parent-session-id`` header. Written by +#: :class:`TokenCaptureRequestProcessor`, stored on the record so the session +#: tree is recoverable from Switchyard alone (no harness logs). Absent for +#: root/single-agent calls. +CTX_TOKEN_CAPTURE_PARENT_SESSION = "_token_capture_parent_session" + +#: Context metadata key set when the inbound request asked for streaming and +#: the request was flipped to non-streaming for faithful token capture. Read +#: by the response processor, which synthesizes the client-facing stream. +CTX_TOKEN_CAPTURE_ORIGINAL_STREAM = "_token_capture_original_stream" + +#: Header-less harnesses (OpenClaw) ride the session id on their API key. The +#: launcher prefixes it with this marker so the server can recognize a session +#: id explicitly, rather than guessing from the string shape — a real client +#: credential (no marker) is therefore never mistaken for a session id and +#: leaked into record files / directory names. Shared with the launcher, which +#: applies the prefix (see ``launch_intake_config``). The colon makes an +#: accidental collision with a real API key effectively impossible. +SESSION_API_KEY_MARKER = "sycap-session:" + +#: Top-level request-body key carrying the capture session for harness clients +#: with no custom-header or API-key surface (e.g. clients whose only reachable +#: knob is extra JSON body fields). Always stripped before the request is +#: forwarded upstream, captured or not. +BODY_SESSION_KEY = "proxy_x_session_id" + +#: Harness-native session headers read as the capture session id when the +#: explicit ``proxy_x_session_id`` (header/body) is absent. Lets a stock, +#: unmodified harness be captured through its own correlation header — the whole +#: point of capture-and-reconstruct. OpenCode (run stock, no fork) emits +#: ``X-Session-Id`` on every request to a custom OpenAI-compatible provider, and +#: a distinct id per subagent session. Matched case-insensitively against the +#: retained request header map. Kept as an explicit allowlist (not any header) +#: so an arbitrary caller header is never mistaken for a session id. +NATIVE_SESSION_HEADERS: tuple[str, ...] = ("x-session-id",) + +#: Harness-native header naming the parent session of a subagent call (OpenCode +#: sends it on every subagent request). Captured onto the record so the session +#: tree is durable in Switchyard. Matched case-insensitively. +NATIVE_PARENT_SESSION_HEADER = "x-parent-session-id" + +#: Caller-supplied sampling params stripped so the target's derived +#: token-capture ``extra_body`` params (see ``llm_target_with_token_capture``) +#: always win. +_CALLER_PARAM_KEYS = ("logprobs", "top_logprobs", "return_token_ids") + + +class TokenCaptureRequestProcessor: + """Resolve the capture session and force the upstream call non-streaming. + + Runs after :class:`~switchyard.lib.processors.rl_logging_request_processor.RlLoggingRequestProcessor`, + whose translated snapshot supplies the record's ``messages``. Calls with no + resolvable session are forwarded untouched and go uncaptured. + + For captured calls, caller-supplied ``logprobs`` / ``top_logprobs`` / + ``return_token_ids`` are stripped so the target's derived params win, and + streaming requests are flipped to ``stream: false`` upstream — vLLM only + returns token IDs on buffered completions. The original intent is recorded + so the response processor can synthesize an equivalent stream for the + client (harnesses like Claude Code cannot disable streaming). + """ + + async def process(self, ctx: ProxyContext, request: ChatRequest) -> ChatRequest: + """Resolve the session; strip caller params and flip ``stream`` off.""" + body = dict(request.body) + mutated = BODY_SESSION_KEY in body + body_session = body.pop(BODY_SESSION_KEY, None) + + session_id = _resolve_session_id(ctx, body_session) + if session_id is None: + logger.debug("token capture: no session id on request; forwarding uncaptured") + if mutated: + request.replace_body(body) + return request + ctx.metadata[CTX_TOKEN_CAPTURE_SESSION] = session_id + parent_session = _native_parent_session_id(ctx) + if parent_session is not None: + ctx.metadata[CTX_TOKEN_CAPTURE_PARENT_SESSION] = parent_session + + for key in _CALLER_PARAM_KEYS: + if key in body: + del body[key] + mutated = True + if body.get("stream"): + ctx.metadata[CTX_TOKEN_CAPTURE_ORIGINAL_STREAM] = True + body["stream"] = False + # stream_options is only valid alongside stream=true. + body.pop("stream_options", None) + mutated = True + if mutated: + request.replace_body(body) + return request + + +def _resolve_session_id(ctx: ProxyContext, body_session: object) -> str | None: + metadata = ctx.metadata.get(CTX_REQUEST_METADATA) + session_id = getattr(metadata, "session_id", None) + if isinstance(session_id, str) and session_id: + return session_id + # Clients whose only reachable knob is extra JSON body fields ride the + # session id on a top-level body key (stripped by the caller above). + if isinstance(body_session, str) and body_session: + return body_session + # Stock harnesses that emit their own correlation header (e.g. OpenCode's + # ``X-Session-Id``) are captured through it — no proxy-specific header, body + # field, or fork required. + native = _native_session_id(ctx) + if native is not None: + return native + # Harnesses with no custom-header surface (OpenClaw) ride the session id on + # their API key, marker-prefixed by the launcher. Strip the marker to get + # the session id; a caller key without the marker is a real credential and + # is never treated as a session. + caller_key = ctx.metadata.get(CTX_CALLER_API_KEY) + if isinstance(caller_key, str) and caller_key.startswith(SESSION_API_KEY_MARKER): + session = caller_key[len(SESSION_API_KEY_MARKER):] + return session or None + return None + + +def _lowercased_headers(ctx: ProxyContext) -> dict[str, str] | None: + """The retained request header map with lowercased names, or ``None`` if absent.""" + headers = ctx.metadata.get(CTX_PROFILE_REQUEST_HEADERS) + if not isinstance(headers, Mapping): + return None + return {str(name).lower(): value for name, value in headers.items()} + + +def _native_session_id(ctx: ProxyContext) -> str | None: + """First non-empty value among :data:`NATIVE_SESSION_HEADERS` in the request headers. + + Matched case-insensitively (``NATIVE_SESSION_HEADERS`` are already lowercased). + """ + lowered = _lowercased_headers(ctx) + if lowered is None: + return None + for header in NATIVE_SESSION_HEADERS: + value = lowered.get(header) + if isinstance(value, str) and value: + return value + return None + + +def _native_parent_session_id(ctx: ProxyContext) -> str | None: + """The :data:`NATIVE_PARENT_SESSION_HEADER` value, or ``None`` (root/single-agent call).""" + lowered = _lowercased_headers(ctx) + if lowered is None: + return None + value = lowered.get(NATIVE_PARENT_SESSION_HEADER) + return value if isinstance(value, str) and value else None diff --git a/switchyard/lib/processors/token_capture_response_processor.py b/switchyard/lib/processors/token_capture_response_processor.py new file mode 100644 index 00000000..2616f261 --- /dev/null +++ b/switchyard/lib/processors/token_capture_response_processor.py @@ -0,0 +1,354 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Response-side processor that writes per-call token-level completion records.""" + +from __future__ import annotations + +import hashlib +import json +import logging +import math +import os +import uuid as uuid_lib +from collections.abc import AsyncIterator +from datetime import UTC, datetime +from pathlib import Path +from typing import Any, TypeGuard, cast + +from switchyard.lib.processors.rl_logging_request_processor import ( + CTX_RL_LOGGING_REQUEST, + RlLoggingRequestProcessor, +) +from switchyard.lib.processors.rl_logging_response_processor import ( + build_trace_messages, + build_trace_token_count, + format_trace_tool_choice, + format_trace_tools, +) +from switchyard.lib.processors.token_capture_request_processor import ( + CTX_TOKEN_CAPTURE_ORIGINAL_STREAM, + CTX_TOKEN_CAPTURE_PARENT_SESSION, + CTX_TOKEN_CAPTURE_SESSION, + TokenCaptureRequestProcessor, +) +from switchyard.lib.proxy_context import ProxyContext +from switchyard_rust.core import ChatResponse, ChatResponseType, response_type_matches + +logger = logging.getLogger(__name__) + +JsonObject = dict[str, Any] + +#: Records live under ``/sessions//.json`` so token +#: capture never collides with flat text traces in the same log directory. +_SESSIONS_SUBDIR = "sessions" + +#: Schema version stamped on every stored record and retrieval envelope. +SCHEMA_VERSION = 1 + + +class TokenCaptureResponseProcessor: + """Write one token-level completion record per captured turn to ``capture_dir``. + + Runs before the terminal ``TranslationEngine``, so the raw backend body + still carries the engine token fields (``prompt_token_ids``, + ``choice.token_ids``, ``logprobs.content[]``) that translation would drop. + Each record unifies the RL text trace (messages/tools built from the + ``CTX_RL_LOGGING_REQUEST`` snapshot) with those token-level fields. The + live response is returned unchanged — this processor only observes. + + Calls without a resolved capture session (``CTX_TOKEN_CAPTURE_SESSION``) + are skipped. Records that fail token-field validation are still stored, + flagged ``is_valid: false`` — capture never affects the harness response. + """ + + def __init__(self, capture_dir: Path | str) -> None: + self._capture_dir = Path(capture_dir) + self._capture_dir.mkdir(parents=True, exist_ok=True) + self._warned_streaming = False + + async def process(self, ctx: ProxyContext, response: ChatResponse) -> ChatResponse: + """Record the completed turn's token-level data. + + Returns the response unchanged, except when the request processor + flipped a streaming request to non-streaming: the buffered completion + is then re-emitted as a synthetic OpenAI chunk stream so the terminal + ``TranslationEngine`` can stream to the client in its native format. + """ + session_id = ctx.metadata.get(CTX_TOKEN_CAPTURE_SESSION) + if not isinstance(session_id, str) or not session_id: + logger.debug("token capture: no capture session on context; skipping record") + return response + if not response_type_matches(response, ChatResponseType.OPENAI_COMPLETION): + self._note_skip(response) + return response + + record = self._build_record(ctx, session_id, response) + try: + self._write_record(session_id, record) + except OSError as exc: + logger.warning( + "token capture: failed to write record to %s: %s", + self._capture_dir, + exc, + ) + if ctx.metadata.get(CTX_TOKEN_CAPTURE_ORIGINAL_STREAM): + return _synthesize_stream(dict(response.body)) + return response + + def _note_skip(self, response: ChatResponse) -> None: + streaming = response_type_matches(response, ChatResponseType.OPENAI_STREAM) + if streaming and not self._warned_streaming: + self._warned_streaming = True + logger.warning( + "token capture: a streaming response reached the capture " + "processor; record skipped (unexpected — capture forces " + "non-streaming upstream)", + ) + + def _build_record( + self, ctx: ProxyContext, session_id: str, response: ChatResponse + ) -> JsonObject: + body = dict(response.body) + request = ctx.metadata.get(CTX_RL_LOGGING_REQUEST) + request = request if isinstance(request, dict) else {} + choices = body.get("choices") + choices = choices if isinstance(choices, list) else [] + choice = choices[0] if choices and isinstance(choices[0], dict) else {} + message = choice.get("message") + message = message if isinstance(message, dict) else {} + + request_id = body.get("id") + model = body.get("model") + prompt_token_ids = body.get("prompt_token_ids") + generation_token_ids = choice.get("token_ids") + generation_log_probs = _extract_log_probs(choice.get("logprobs")) + problem = _validate_token_fields( + request_id, model, prompt_token_ids, generation_token_ids, generation_log_probs, len(choices) + ) + if problem is not None: + logger.warning("token capture: storing record with is_valid=false: %s", problem) + + return { + "schema_version": SCHEMA_VERSION, + "uuid": str(uuid_lib.uuid4()), + "session_id": session_id, + "parent_session_id": ctx.metadata.get(CTX_TOKEN_CAPTURE_PARENT_SESSION), + "captured_at": datetime.now(UTC).isoformat(), + "request_id": request_id, + "model": model, + "messages": build_trace_messages(request, message), + "tools": format_trace_tools(request.get("tools", [])), + "tool_choice": format_trace_tool_choice(request.get("tool_choice")), + "token_count": build_trace_token_count(body.get("usage")), + "prompt_token_ids": prompt_token_ids, + "generation_token_ids": generation_token_ids, + "generation_log_probs": generation_log_probs, + "finish_reason": choice.get("finish_reason"), + "is_valid": problem is None, + } + + def _write_record(self, session_id: str, record: JsonObject) -> None: + session_dir = sessions_root(self._capture_dir) / session_dir_name(session_id) + session_dir.mkdir(parents=True, exist_ok=True) + path = session_dir / f"{record['uuid']}.json" + # Write-then-rename so the retrieval endpoint never reads a torn record + # (the endpoint globs *.json; the .tmp suffix keeps partials invisible). + tmp_path = path.with_name(path.name + ".tmp") + with open(tmp_path, "w") as handle: + json.dump(record, handle, indent=2) + os.replace(tmp_path, path) + # Records carry full conversation text; keep them owner-only. + os.chmod(path, 0o600) + + def get_endpoint(self) -> Any: + """Contribute ``GET /v1/sessions`` record retrieval to the server.""" + from switchyard.lib.endpoints.token_capture_endpoint import ( + TokenCaptureSessionsEndpoint, + ) + + return TokenCaptureSessionsEndpoint(self._capture_dir) + + +def build_token_capture_processors( + capture_dir: Path | None, +) -> tuple[list[Any], list[Any]]: + """Request/response processor lists for token-level capture. + + The request side pairs :class:`RlLoggingRequestProcessor` — whose + translated snapshot supplies the record's ``messages``/``tools`` — with + :class:`TokenCaptureRequestProcessor`. Returns ``([], [])`` when + ``capture_dir`` is ``None`` (capture disabled). Shared by the ``launch`` + and ``serve`` wiring. + """ + if capture_dir is None: + return [], [] + return ( + [RlLoggingRequestProcessor(), TokenCaptureRequestProcessor()], + [TokenCaptureResponseProcessor(capture_dir)], + ) + + +def _extract_log_probs(logprobs: object) -> list[Any] | None: + """``[entry.logprob for entry in logprobs.content]``, or ``None`` off-shape.""" + if not isinstance(logprobs, dict): + return None + content = logprobs.get("content") + if not isinstance(content, list): + return None + return [entry.get("logprob") if isinstance(entry, dict) else None for entry in content] + + +def _validate_token_fields( + request_id: object, + model: object, + prompt_token_ids: object, + generation_token_ids: object, + generation_log_probs: object, + choice_count: int, +) -> str | None: + """Reason the record is unusable for training, or ``None`` if valid.""" + if choice_count != 1: + return f"response has {choice_count} choices; multiple-choice capture is unsupported" + # Provenance: the design requires request id and model identity. ``model`` + # is load-bearing — token ids are meaningless without the tokenizer that + # produced them — and real vLLM always emits both, so absence is anomalous. + if not isinstance(request_id, str) or not request_id: + return "request_id is missing or not a non-empty string" + if not isinstance(model, str) or not model: + return "model is missing or not a non-empty string" + if not _is_token_id_list(prompt_token_ids): + return "prompt_token_ids is not a non-empty list of ints" + if not _is_token_id_list(generation_token_ids): + return "generation_token_ids is not a non-empty list of ints" + if not isinstance(generation_log_probs, list) or not all( + _is_finite_number(value) for value in generation_log_probs + ): + return "generation_log_probs is not a list of finite floats" + if len(generation_token_ids) != len(generation_log_probs): + return ( + f"generation_token_ids ({len(generation_token_ids)}) and " + f"generation_log_probs ({len(generation_log_probs)}) are misaligned" + ) + return None + + +def _is_token_id_list(value: object) -> TypeGuard[list[int]]: + return ( + isinstance(value, list) + and bool(value) + and all(isinstance(item, int) and not isinstance(item, bool) for item in value) + ) + + +def _is_finite_number(value: object) -> bool: + return isinstance(value, int | float) and not isinstance(value, bool) and math.isfinite(value) + + +def _synthesize_stream(body: JsonObject) -> ChatResponse: + """Re-emit a buffered completion as a synthetic OpenAI chunk stream. + + The harness asked for streaming, but the upstream call was forced + non-streaming for token capture. Two chunks reproduce the completion: + the first carries the assistant delta (content + tool calls), the second + carries ``finish_reason`` and usage. First choice only — training + harnesses run ``n=1``. Token-level fields stay in the stored record; + chunks carry only the standard wire fields. + """ + from openai.types import CompletionUsage + from openai.types.chat import ChatCompletionChunk + from openai.types.chat.chat_completion_chunk import ( + Choice as ChunkChoice, + ) + from openai.types.chat.chat_completion_chunk import ( + ChoiceDelta, + ChoiceDeltaToolCall, + ChoiceDeltaToolCallFunction, + ) + from pydantic import ValidationError + + from switchyard.lib.chat_response.openai_chat import ResponseStream + + completion_id = str(body.get("id") or "chatcmpl-token-capture") + created = int(body.get("created") or 0) + model = str(body.get("model") or "unknown") + + choices = body.get("choices") + choice: JsonObject = {} + if isinstance(choices, list) and choices and isinstance(choices[0], dict): + choice = choices[0] + message_obj = choice.get("message") + message: JsonObject = message_obj if isinstance(message_obj, dict) else {} + + delta_tool_calls: list[ChoiceDeltaToolCall] | None = None + raw_tool_calls = message.get("tool_calls") + if isinstance(raw_tool_calls, list) and raw_tool_calls: + delta_tool_calls = [ + ChoiceDeltaToolCall( + index=index, + id=call.get("id"), + type="function", + function=ChoiceDeltaToolCallFunction( + name=(call.get("function") or {}).get("name"), + arguments=(call.get("function") or {}).get("arguments"), + ), + ) + for index, call in enumerate(raw_tool_calls) + if isinstance(call, dict) + ] + + usage = body.get("usage") + final_usage = None + if isinstance(usage, dict): + # A malformed usage block must not abort the client-facing stream — + # drop it rather than raise from model_validate. + try: + final_usage = CompletionUsage.model_validate(usage) + except ValidationError as exc: + logger.warning("token capture: dropping malformed usage in synthesized stream: %s", exc) + finish_reason = cast(Any, choice.get("finish_reason") or "stop") + + async def _chunks() -> AsyncIterator[ChatCompletionChunk]: + yield ChatCompletionChunk( + id=completion_id, + object="chat.completion.chunk", + created=created, + model=model, + choices=[ + ChunkChoice( + index=0, + delta=ChoiceDelta( + role="assistant", + content=message.get("content"), + tool_calls=delta_tool_calls, + ), + finish_reason=None, + ) + ], + ) + yield ChatCompletionChunk( + id=completion_id, + object="chat.completion.chunk", + created=created, + model=model, + choices=[ChunkChoice(index=0, delta=ChoiceDelta(), finish_reason=finish_reason)], + usage=final_usage, + ) + + return ChatResponse.openai_stream(ResponseStream(_chunks())) + + +def sessions_root(capture_dir: Path) -> Path: + """Root under which per-session record directories live (writer + endpoint).""" + return capture_dir / _SESSIONS_SUBDIR + + +def session_dir_name(session_id: str) -> str: + """Collision-resistant directory name for a session id (writer + endpoint). + + A hash, not a sanitized form: sanitizing distinct opaque ids to the same + safe string (e.g. ``tenant:a`` and ``tenant*a`` → ``tenant_a``) would let + one session's records surface under another. The fixed-length hex output is + also inherently path-safe, so no traversal guard is needed. + """ + return hashlib.sha256(session_id.encode()).hexdigest() diff --git a/tests/test_token_capture.py b/tests/test_token_capture.py new file mode 100644 index 00000000..badbbc78 --- /dev/null +++ b/tests/test_token_capture.py @@ -0,0 +1,974 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for token-level capture (RL logging + ``token_capture_engine`` routes).""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import pytest + +from switchyard.cli.route_bundle import route_bundle_declares_token_capture +from switchyard.cli.switchyard_cli import _build_parser +from switchyard.lib.processors.rl_logging_request_processor import RlLoggingRequestProcessor +from switchyard.lib.processors.token_capture_request_processor import ( + CTX_TOKEN_CAPTURE_ORIGINAL_STREAM, + CTX_TOKEN_CAPTURE_SESSION, + TokenCaptureRequestProcessor, +) +from switchyard.lib.processors.token_capture_response_processor import ( + TokenCaptureResponseProcessor, + build_token_capture_processors, +) +from switchyard.lib.request_metadata import CTX_PROFILE_REQUEST_HEADERS, CTX_REQUEST_METADATA +from switchyard.server.server_util import resolve_rl_log_dir +from switchyard_rust.components import RequestMetadata +from switchyard_rust.core import ChatRequest, ChatResponse, ProxyContext + +_SESSION = "claude-1700000000000-abc12345" + + +def _anthropic_request() -> ChatRequest: + return ChatRequest.anthropic( + { + "model": "claude-test", + "max_tokens": 32, + "messages": [{"role": "user", "content": "hello"}], + } + ) + + +def _vllm_completion(*, choices: list | None = None) -> ChatResponse: + """Backend body as vLLM emits it with token-capture params enabled.""" + if choices is None: + choices = [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + "token_ids": [7, 8], + "logprobs": { + "content": [ + {"token": "hi", "logprob": -0.1}, + {"token": "!", "logprob": -0.2}, + ] + }, + } + ] + return ChatResponse.openai_completion( + { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1700000000, + "model": "Qwen/Qwen3-0.6B", + "prompt_token_ids": [1, 2, 3], + "choices": choices, + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + } + ) + + +def _read_only_record(capture_dir: Path) -> dict: + files = list(capture_dir.rglob("*.json")) + assert len(files) == 1, f"expected exactly one record file, got {files}" + return json.loads(files[0].read_text()) + + +def _session_ctx(session_id: str | None = _SESSION) -> ProxyContext: + ctx = ProxyContext() + if session_id is not None: + ctx.metadata[CTX_REQUEST_METADATA] = RequestMetadata.from_headers( + {"proxy_x_session_id": session_id} + ) + return ctx + + +async def _run( + capture_dir: Path, + response: ChatResponse, + *, + session_id: str | None = None, + request: ChatRequest | None = None, + ctx: ProxyContext | None = None, +) -> ProxyContext: + """Run the capture pair the way the wiring installs it (rl snapshot first).""" + if ctx is None: + ctx = _session_ctx(session_id) + if request is None: + request = _anthropic_request() + await RlLoggingRequestProcessor().process(ctx, request) + await TokenCaptureRequestProcessor().process(ctx, request) + await TokenCaptureResponseProcessor(capture_dir).process(ctx, response) + return ctx + + +# --------------------------------------------------------------------------- +# Activation: --enable-rl-logging + token_capture_engine route +# --------------------------------------------------------------------------- + + +def test_capture_gates_on_rl_logging_flag(tmp_path: Path) -> None: + parser = _build_parser() + + args = parser.parse_args(["serve"]) + assert resolve_rl_log_dir(args) is None + + args = parser.parse_args(["--enable-rl-logging", "--rl-log-dir", str(tmp_path), "serve"]) + assert resolve_rl_log_dir(args) == tmp_path + + +def test_route_bundle_declares_token_capture_dict() -> None: + declaring = { + "routes": { + "m": { + "type": "model", + "target": {"model": "a", "base_url": "http://x/v1"}, + "token_capture_engine": "vllm", + }, + }, + } + plain = { + "routes": { + "m": {"type": "model", "target": {"model": "a", "base_url": "http://x/v1"}}, + }, + } + assert route_bundle_declares_token_capture(declaring) is True + assert route_bundle_declares_token_capture(plain) is False + assert route_bundle_declares_token_capture(None) is False + + +def test_route_bundle_declares_token_capture_path(tmp_path: Path) -> None: + profile = tmp_path / "profiles.yaml" + profile.write_text( + "routes:\n" + " m:\n" + " type: model\n" + " target:\n" + " model: a\n" + " base_url: http://x/v1\n" + " token_capture_engine: vllm\n" + ) + assert route_bundle_declares_token_capture(str(profile)) is True + # An unreadable bundle reports False; the real table load raises. + assert route_bundle_declares_token_capture(str(tmp_path / "missing.yaml")) is False + + +def test_route_level_token_capture_engine_parses() -> None: + """token_capture_engine sits at route level, beside a string target; it must + parse and derive the vLLM capture params for the route's single target.""" + from switchyard.cli.route_bundle import build_route_bundle_table + + raw = { + "defaults": {"api_key": "dummy", "base_url": "http://vllm:8000/v1", "format": "openai"}, + "routes": { + "policy-model": { + "type": "model", + "target": "served-model", + "token_capture_engine": "vllm", + } + }, + } + table = build_route_bundle_table(raw, token_capture_enabled=True) + assert "policy-model" in table.registered_models() + assert route_bundle_declares_token_capture(raw) is True + + +def test_serve_installs_capture_pair_for_declaring_bundle( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """rl-logging on + token_capture_engine route -> capture pair replaces rl pair.""" + import switchyard.cli.switchyard_cli as cli + + profile = tmp_path / "profiles.yaml" + profile.write_text( + "routes:\n" + " m:\n" + " type: model\n" + " target:\n" + " model: a\n" + " base_url: http://x/v1\n" + " token_capture_engine: vllm\n" + ) + captured: dict[str, list] = {} + + class _FakeTable: + def registered_models(self) -> list[str]: + return ["m"] + + def default_model(self) -> str | None: + return None + + def _fake_load( + routing_profiles, + *, + pre_routing_request_processors=(), + extra_response_processors=(), + **_kwargs, + ): + captured["request"] = list(pre_routing_request_processors) + captured["response"] = list(extra_response_processors) + return _FakeTable() + + monkeypatch.setattr(cli, "load_route_bundle_table", _fake_load) + monkeypatch.setattr(cli, "build_and_serve", lambda *a, **k: None) + + args = argparse.Namespace( + config=None, + routing_profiles=str(profile), + enable_rl_logging=True, + rl_log_dir=str(tmp_path / "rl_data"), + intake_enabled=False, + intake_base_url=None, + intake_workspace=None, + intake_api_key=None, + intake_nvdataflow_project=None, + ) + cli._cmd_serve(args) + + assert [type(p).__name__ for p in captured["request"]] == [ + "RlLoggingRequestProcessor", + "TokenCaptureRequestProcessor", + ] + assert [type(p).__name__ for p in captured["response"]] == [ + "TokenCaptureResponseProcessor", + ] + + +def test_builder_disabled_returns_empty_lists() -> None: + assert build_token_capture_processors(None) == ([], []) + + +def test_builder_returns_capture_pair(tmp_path: Path) -> None: + request, response = build_token_capture_processors(tmp_path) + assert [type(p).__name__ for p in request] == [ + "RlLoggingRequestProcessor", + "TokenCaptureRequestProcessor", + ] + assert [type(p).__name__ for p in response] == ["TokenCaptureResponseProcessor"] + + +def test_build_launch_capture_processors_token_capture(tmp_path: Path) -> None: + from switchyard.cli.launchers.launch_intake_config import build_launch_capture_processors + + request, response = build_launch_capture_processors(None, tmp_path, token_capture=True) + assert [type(p).__name__ for p in request] == [ + "RlLoggingRequestProcessor", + "TokenCaptureRequestProcessor", + ] + assert [type(p).__name__ for p in response] == ["TokenCaptureResponseProcessor"] + + +# --------------------------------------------------------------------------- +# Request processor: session resolution, caller-param strip, stream flip +# --------------------------------------------------------------------------- + + +def _streaming_request() -> ChatRequest: + return ChatRequest.openai_chat( + { + "model": "m", + "messages": [{"role": "user", "content": "hello"}], + "stream": True, + "stream_options": {"include_usage": True}, + } + ) + + +async def test_streaming_request_flipped_when_session_present() -> None: + ctx = _session_ctx() + request = _streaming_request() + await TokenCaptureRequestProcessor().process(ctx, request) + + body = dict(request.body) + assert body["stream"] is False + assert "stream_options" not in body + assert ctx.metadata[CTX_TOKEN_CAPTURE_ORIGINAL_STREAM] is True + assert ctx.metadata[CTX_TOKEN_CAPTURE_SESSION] == _SESSION + + +async def test_no_session_leaves_request_untouched() -> None: + ctx = ProxyContext() + request = _streaming_request() + await TokenCaptureRequestProcessor().process(ctx, request) + + body = dict(request.body) + assert body["stream"] is True + assert body["stream_options"] == {"include_usage": True} + assert CTX_TOKEN_CAPTURE_ORIGINAL_STREAM not in ctx.metadata + assert CTX_TOKEN_CAPTURE_SESSION not in ctx.metadata + + +async def test_caller_token_params_stripped_when_session_present() -> None: + ctx = _session_ctx() + request = ChatRequest.openai_chat( + { + "model": "m", + "messages": [], + "logprobs": False, + "top_logprobs": None, + "return_token_ids": False, + } + ) + await TokenCaptureRequestProcessor().process(ctx, request) + + body = dict(request.body) + # The target's derived extra_body params must win over caller-supplied ones. + assert "logprobs" not in body + assert "top_logprobs" not in body + assert "return_token_ids" not in body + + +async def test_caller_token_params_kept_without_session() -> None: + ctx = ProxyContext() + request = ChatRequest.openai_chat( + {"model": "m", "messages": [], "logprobs": False, "top_logprobs": None} + ) + await TokenCaptureRequestProcessor().process(ctx, request) + + body = dict(request.body) + assert body["logprobs"] is False + assert body["top_logprobs"] is None + + +async def test_non_streaming_request_left_untouched() -> None: + ctx = _session_ctx() + request = ChatRequest.openai_chat({"model": "m", "messages": []}) + await TokenCaptureRequestProcessor().process(ctx, request) + + assert "stream" not in dict(request.body) + assert CTX_TOKEN_CAPTURE_ORIGINAL_STREAM not in ctx.metadata + + +async def test_session_id_from_marked_caller_key(tmp_path: Path) -> None: + from switchyard.lib.processors.token_capture_request_processor import SESSION_API_KEY_MARKER + from switchyard.lib.proxy_context import CTX_CALLER_API_KEY + + # OpenClaw rides the session id on the marker-prefixed API key. A custom + # (non-launcher-shaped) session id must still resolve. + ctx = ProxyContext() + ctx.metadata[CTX_CALLER_API_KEY] = SESSION_API_KEY_MARKER + "my-custom-rerun-42" + await _run(tmp_path, _vllm_completion(), ctx=ctx) + + record = _read_only_record(tmp_path) + assert record["session_id"] == "my-custom-rerun-42" + + +async def test_unmarked_api_key_never_becomes_session_id(tmp_path: Path) -> None: + from switchyard.lib.proxy_context import CTX_CALLER_API_KEY + + ctx = ProxyContext() + # A real credential (no marker) — even one shaped like an old-style id — + # must not become a session or leak to disk. + ctx.metadata[CTX_CALLER_API_KEY] = "openclaw-1783600121000-a1b2c3d4" + await _run(tmp_path, _vllm_completion(), ctx=ctx) + + assert list(tmp_path.rglob("*.json")) == [] + + +async def test_session_id_from_body_field(tmp_path: Path) -> None: + ctx = ProxyContext() + request = ChatRequest.openai_chat({"model": "m", "messages": [], "proxy_x_session_id": "abc123"}) + await _run(tmp_path, _vllm_completion(), ctx=ctx, request=request) + + # The session field never reaches the upstream body. + assert "proxy_x_session_id" not in dict(request.body) + record = _read_only_record(tmp_path) + assert record["session_id"] == "abc123" + + +async def test_header_session_wins_over_body_field(tmp_path: Path) -> None: + request = ChatRequest.openai_chat({"model": "m", "messages": [], "proxy_x_session_id": "from-body"}) + await _run(tmp_path, _vllm_completion(), session_id="from-header", request=request) + + assert "proxy_x_session_id" not in dict(request.body) + record = _read_only_record(tmp_path) + assert record["session_id"] == "from-header" + + +async def test_non_string_body_session_ignored_but_stripped(tmp_path: Path) -> None: + ctx = ProxyContext() + request = ChatRequest.openai_chat({"model": "m", "messages": [], "proxy_x_session_id": 123}) + await _run(tmp_path, _vllm_completion(), ctx=ctx, request=request) + + # Ignored as a session, but still never forwarded upstream. + assert "proxy_x_session_id" not in dict(request.body) + assert list(tmp_path.rglob("*.json")) == [] + + +async def test_session_id_from_native_header(tmp_path: Path) -> None: + # OpenCode (stock, no fork) rides its session id on the native X-Session-Id + # header; capture keys on it without a proxy header, body field, or api key. + ctx = ProxyContext() + ctx.metadata[CTX_PROFILE_REQUEST_HEADERS] = {"X-Session-Id": "oc-sess-1"} + await _run(tmp_path, _vllm_completion(), ctx=ctx) + + assert _read_only_record(tmp_path)["session_id"] == "oc-sess-1" + + +async def test_native_header_matched_case_insensitively(tmp_path: Path) -> None: + ctx = ProxyContext() + ctx.metadata[CTX_PROFILE_REQUEST_HEADERS] = {"x-session-id": "oc-sess-2"} + await _run(tmp_path, _vllm_completion(), ctx=ctx) + + assert _read_only_record(tmp_path)["session_id"] == "oc-sess-2" + + +async def test_proxy_session_wins_over_native_header(tmp_path: Path) -> None: + ctx = _session_ctx("from-proxy-header") + ctx.metadata[CTX_PROFILE_REQUEST_HEADERS] = {"X-Session-Id": "from-native"} + await _run(tmp_path, _vllm_completion(), ctx=ctx) + + assert _read_only_record(tmp_path)["session_id"] == "from-proxy-header" + + +async def test_unlisted_header_never_becomes_session(tmp_path: Path) -> None: + # Only allowlisted native headers are honored; an arbitrary caller header is + # never mistaken for a session id. + ctx = ProxyContext() + ctx.metadata[CTX_PROFILE_REQUEST_HEADERS] = {"x-random-id": "nope"} + await _run(tmp_path, _vllm_completion(), ctx=ctx) + + assert list(tmp_path.rglob("*.json")) == [] + + +async def test_parent_session_captured_from_native_header(tmp_path: Path) -> None: + # OpenCode sends x-parent-session-id on every subagent call; capture it so the + # session tree is recoverable from Switchyard alone. + ctx = ProxyContext() + ctx.metadata[CTX_PROFILE_REQUEST_HEADERS] = { + "X-Session-Id": "sub-1", + "x-parent-session-id": "root-1", + } + await _run(tmp_path, _vllm_completion(), ctx=ctx) + + record = _read_only_record(tmp_path) + assert record["session_id"] == "sub-1" + assert record["parent_session_id"] == "root-1" + + +async def test_parent_session_absent_is_none(tmp_path: Path) -> None: + ctx = ProxyContext() + ctx.metadata[CTX_PROFILE_REQUEST_HEADERS] = {"X-Session-Id": "root-1"} + await _run(tmp_path, _vllm_completion(), ctx=ctx) + + assert _read_only_record(tmp_path)["parent_session_id"] is None + + +# --------------------------------------------------------------------------- +# Response processor: unified record +# --------------------------------------------------------------------------- + + +async def test_unified_record_schema(tmp_path: Path) -> None: + await _run(tmp_path, _vllm_completion(), session_id=_SESSION) + record = _read_only_record(tmp_path) + + import uuid as uuid_lib + from datetime import datetime + + assert record["schema_version"] == 1 + uuid_lib.UUID(record["uuid"]) # must parse as a valid UUID + assert record["session_id"] == _SESSION + # captured_at must parse as an ISO timestamp (retrieval sorts on it). + datetime.fromisoformat(record["captured_at"]) + assert record["request_id"] == "chatcmpl-test" + assert record["model"] == "Qwen/Qwen3-0.6B" + # Text trace from the translated request snapshot + raw-body assistant turn. + assert [m["role"] for m in record["messages"]] == ["user", "assistant"] + assert record["messages"][0]["content"] == "hello" + assert record["messages"][-1]["content"] == "hi" + assert record["tools"] == [] + assert record["tool_choice"] == "auto" + assert record["token_count"] == { + "prompt_tokens": 3, + "completion_tokens": 2, + "total_tokens": 5, + } + # Token-level fields from the raw vLLM body. + assert record["prompt_token_ids"] == [1, 2, 3] + assert record["generation_token_ids"] == [7, 8] + assert record["generation_log_probs"] == [-0.1, -0.2] + assert record["finish_reason"] == "stop" + assert record["is_valid"] is True + + +async def test_record_file_is_owner_only(tmp_path: Path) -> None: + await _run(tmp_path, _vllm_completion(), session_id=_SESSION) + files = list(tmp_path.rglob("*.json")) + assert len(files) == 1 + assert files[0].stat().st_mode & 0o777 == 0o600 + + +async def test_no_session_is_not_captured(tmp_path: Path) -> None: + await _run(tmp_path, _vllm_completion()) + assert list(tmp_path.rglob("*.json")) == [] + + +async def test_misaligned_logprobs_marks_record_invalid(tmp_path: Path) -> None: + choices = [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + "token_ids": [7, 8], + "logprobs": {"content": [{"token": "hi", "logprob": -0.1}]}, + } + ] + await _run(tmp_path, _vllm_completion(choices=choices), session_id=_SESSION) + record = _read_only_record(tmp_path) + + assert record["is_valid"] is False + assert record["generation_token_ids"] == [7, 8] + assert record["generation_log_probs"] == [-0.1] + + +async def test_multiple_choices_marks_record_invalid(tmp_path: Path) -> None: + choice = { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + "token_ids": [7, 8], + "logprobs": { + "content": [ + {"token": "hi", "logprob": -0.1}, + {"token": "!", "logprob": -0.2}, + ] + }, + } + await _run( + tmp_path, + _vllm_completion(choices=[choice, {**choice, "index": 1}]), + session_id=_SESSION, + ) + record = _read_only_record(tmp_path) + assert record["is_valid"] is False + + +async def test_missing_token_fields_mark_record_invalid(tmp_path: Path) -> None: + """A non-engine backend (no token fields) still yields a stored, invalid record.""" + choices = [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ] + await _run(tmp_path, _vllm_completion(choices=choices), session_id=_SESSION) + record = _read_only_record(tmp_path) + assert record["is_valid"] is False + assert record["messages"][-1]["content"] == "hi" + + +async def test_missing_model_marks_record_invalid(tmp_path: Path) -> None: + # Token ids are meaningless without the tokenizer, so a record with no model + # is untrainable — stored, flagged invalid. + body = dict(_vllm_completion().body) + body.pop("model", None) + await _run(tmp_path, ChatResponse.openai_completion(body), session_id=_SESSION) + record = _read_only_record(tmp_path) + assert record["is_valid"] is False + assert record["model"] is None + + +async def test_missing_request_id_marks_record_invalid(tmp_path: Path) -> None: + body = dict(_vllm_completion().body) + body.pop("id", None) + await _run(tmp_path, ChatResponse.openai_completion(body), session_id=_SESSION) + record = _read_only_record(tmp_path) + assert record["is_valid"] is False + assert record["request_id"] is None + + +async def test_non_finite_logprobs_mark_record_invalid(tmp_path: Path) -> None: + choices = [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + "token_ids": [7, 8], + "logprobs": { + "content": [ + {"token": "hi", "logprob": float("nan")}, + {"token": "!", "logprob": -0.2}, + ] + }, + } + ] + await _run(tmp_path, _vllm_completion(choices=choices), session_id=_SESSION) + record = _read_only_record(tmp_path) + assert record["is_valid"] is False + + +async def test_non_int_token_ids_mark_record_invalid(tmp_path: Path) -> None: + choices = [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + "token_ids": [7.5, "8"], + "logprobs": { + "content": [ + {"token": "hi", "logprob": -0.1}, + {"token": "!", "logprob": -0.2}, + ] + }, + } + ] + await _run(tmp_path, _vllm_completion(choices=choices), session_id=_SESSION) + record = _read_only_record(tmp_path) + assert record["is_valid"] is False + + +async def test_synthetic_stream_preserves_tool_calls(tmp_path: Path) -> None: + from switchyard_rust.core import ChatResponseType, response_type_matches + + choices = [ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": {"name": "write_file", "arguments": '{"path": "hello.py"}'}, + } + ], + }, + "finish_reason": "tool_calls", + "token_ids": [7, 8], + "logprobs": { + "content": [ + {"token": "a", "logprob": -0.1}, + {"token": "b", "logprob": -0.2}, + ] + }, + } + ] + ctx = _session_ctx() + request = _streaming_request() + await RlLoggingRequestProcessor().process(ctx, request) + await TokenCaptureRequestProcessor().process(ctx, request) + out = await TokenCaptureResponseProcessor(tmp_path).process( + ctx, _vllm_completion(choices=choices) + ) + + assert response_type_matches(out, ChatResponseType.OPENAI_STREAM) + chunks = [chunk async for chunk in out.stream] + delta_calls = chunks[0].choices[0].delta.tool_calls + assert delta_calls is not None and len(delta_calls) == 1 + assert delta_calls[0].id == "call_abc123" + assert delta_calls[0].index == 0 + assert delta_calls[0].function.name == "write_file" + assert delta_calls[0].function.arguments == '{"path": "hello.py"}' + assert chunks[-1].choices[0].finish_reason == "tool_calls" + + +async def test_streaming_response_is_skipped(tmp_path: Path) -> None: + async def _iter(): + return + yield # pragma: no cover + + from switchyard.lib.chat_response import ResponseStream + + ctx = ProxyContext() + ctx.metadata[CTX_TOKEN_CAPTURE_SESSION] = _SESSION + response = ChatResponse.openai_stream(ResponseStream(_iter())) + out = await TokenCaptureResponseProcessor(tmp_path).process(ctx, response) + + assert out is response + assert list(tmp_path.rglob("*.json")) == [] + + +async def test_synthetic_stream_replaces_buffered_response(tmp_path: Path) -> None: + ctx = await _run( + tmp_path, + _vllm_completion(), + session_id=_SESSION, + request=_streaming_request(), + ) + assert ctx.metadata[CTX_TOKEN_CAPTURE_ORIGINAL_STREAM] is True + + # The record was still captured from the buffered body. + record = _read_only_record(tmp_path) + assert record["generation_token_ids"] == [7, 8] + + +async def test_synthetic_stream_chunks(tmp_path: Path) -> None: + from switchyard_rust.core import ChatResponseType, response_type_matches + + ctx = _session_ctx() + request = _streaming_request() + await RlLoggingRequestProcessor().process(ctx, request) + await TokenCaptureRequestProcessor().process(ctx, request) + out = await TokenCaptureResponseProcessor(tmp_path).process(ctx, _vllm_completion()) + + # The client gets an OpenAI chunk stream reproducing the completion. + assert response_type_matches(out, ChatResponseType.OPENAI_STREAM) + chunks = [chunk async for chunk in out.stream] + assert len(chunks) == 2 + assert chunks[0].choices[0].delta.content == "hi" + assert chunks[0].choices[0].delta.role == "assistant" + assert chunks[1].choices[0].finish_reason == "stop" + assert chunks[1].usage.total_tokens == 5 + + +async def test_buffered_request_gets_buffered_response(tmp_path: Path) -> None: + from switchyard_rust.core import ChatResponseType, response_type_matches + + response = _vllm_completion() + ctx = _session_ctx() + await RlLoggingRequestProcessor().process(ctx, _anthropic_request()) + await TokenCaptureRequestProcessor().process(ctx, _anthropic_request()) + result = await TokenCaptureResponseProcessor(tmp_path).process(ctx, response) + + assert result is response + assert response_type_matches(result, ChatResponseType.OPENAI_COMPLETION) + + +# --------------------------------------------------------------------------- +# Retrieval endpoint +# --------------------------------------------------------------------------- + + +def _retrieval_client(capture_dir: Path): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + processor = TokenCaptureResponseProcessor(capture_dir) + app = FastAPI() + processor.get_endpoint().register(app) + return TestClient(app, raise_server_exceptions=False) + + +async def test_retrieval_endpoint_serves_session_records(tmp_path: Path) -> None: + await _run(tmp_path, _vllm_completion(), session_id=_SESSION) + await _run(tmp_path, _vllm_completion(), session_id=_SESSION) + client = _retrieval_client(tmp_path) + + result = client.get(f"/v1/sessions/{_SESSION}/completions") + assert result.status_code == 200 + payload = result.json() + assert payload["schema_version"] == 1 + assert payload["session_id"] == _SESSION + assert len(payload["completions"]) == 2 + assert payload["completions"][0]["generation_token_ids"] == [7, 8] + + +async def _capture_native(tmp_path: Path, session_id: str, parent: str | None = None) -> None: + """Capture one call for a stock-harness session via its native headers.""" + ctx = ProxyContext() + headers = {"X-Session-Id": session_id} + if parent is not None: + headers["x-parent-session-id"] = parent + ctx.metadata[CTX_PROFILE_REQUEST_HEADERS] = headers + await _run(tmp_path, _vllm_completion(), ctx=ctx) + + +async def test_list_sessions_endpoint_returns_tree(tmp_path: Path) -> None: + await _capture_native(tmp_path, "root-1") + await _capture_native(tmp_path, "sub-1", parent="root-1") + await _capture_native(tmp_path, "sub-2", parent="root-1") + + payload = _retrieval_client(tmp_path).get("/v1/sessions").json() + assert payload["schema_version"] == 1 + tree = {s["session_id"]: s["parent_session_id"] for s in payload["sessions"]} + assert tree == {"root-1": None, "sub-1": "root-1", "sub-2": "root-1"} + + +def test_list_sessions_endpoint_empty(tmp_path: Path) -> None: + payload = _retrieval_client(tmp_path).get("/v1/sessions").json() + assert payload == {"schema_version": 1, "sessions": []} + + +def _session_dir(tmp_path: Path, session_id: str) -> Path: + from switchyard.lib.processors.token_capture_response_processor import ( + session_dir_name, + sessions_root, + ) + + return sessions_root(tmp_path) / session_dir_name(session_id) + + +def test_retrieval_completions_sorted_by_captured_at_then_uuid(tmp_path: Path) -> None: + session_dir = _session_dir(tmp_path, _SESSION) + session_dir.mkdir(parents=True) + + def _rec(captured_at: str, uuid: str) -> str: + return json.dumps({"session_id": _SESSION, "captured_at": captured_at, "uuid": uuid}) + + # Filenames deliberately reverse the capture order. + (session_dir / "z.json").write_text(_rec("2026-01-01T00:00:00+00:00", "aaa")) + (session_dir / "a.json").write_text(_rec("2026-01-02T00:00:00+00:00", "bbb")) + (session_dir / "m.json").write_text(_rec("2026-01-01T00:00:00+00:00", "zzz")) + client = _retrieval_client(tmp_path) + + payload = client.get(f"/v1/sessions/{_SESSION}/completions").json() + assert [(r["captured_at"], r["uuid"]) for r in payload["completions"]] == [ + ("2026-01-01T00:00:00+00:00", "aaa"), + ("2026-01-01T00:00:00+00:00", "zzz"), + ("2026-01-02T00:00:00+00:00", "bbb"), + ] + + +async def test_retrieval_skips_torn_record_and_serves_the_rest(tmp_path: Path) -> None: + await _run(tmp_path, _vllm_completion(), session_id=_SESSION) + # Simulate a torn/corrupt record alongside the good one. + (_session_dir(tmp_path, _SESSION) / "0000-torn.json").write_text('{"schema_version": 1, "uuid": ') + client = _retrieval_client(tmp_path) + + payload = client.get(f"/v1/sessions/{_SESSION}/completions").json() + assert len(payload["completions"]) == 1 + assert payload["completions"][0]["generation_token_ids"] == [7, 8] + + +async def test_sessions_with_sanitizer_colliding_ids_stay_isolated(tmp_path: Path) -> None: + # These two distinct ids collapse to the same name under a naive sanitizer + # (both -> "tenant_a"); the hashed dir keeps them in separate directories so + # one session's records never surface under the other. + await _run(tmp_path, _vllm_completion(), session_id="tenant:a") + await _run(tmp_path, _vllm_completion(choices=[ + {"index": 0, "message": {"role": "assistant", "content": "other"}, + "finish_reason": "stop", "token_ids": [9], + "logprobs": {"content": [{"token": "other", "logprob": -0.5}]}}, + ]), session_id="tenant*a") + client = _retrieval_client(tmp_path) + + a = client.get("/v1/sessions/tenant:a/completions").json() + b = client.get("/v1/sessions/tenant*a/completions").json() + assert len(a["completions"]) == 1 and a["completions"][0]["session_id"] == "tenant:a" + assert len(b["completions"]) == 1 and b["completions"][0]["session_id"] == "tenant*a" + assert a["completions"][0]["generation_token_ids"] != b["completions"][0]["generation_token_ids"] + + +def test_retrieval_endpoint_unknown_session_is_404(tmp_path: Path) -> None: + client = _retrieval_client(tmp_path) + assert client.get("/v1/sessions/nope/completions").status_code == 404 + + +def test_retrieval_endpoint_rejects_path_traversal(tmp_path: Path) -> None: + # A sentinel record one level above sessions/ — a vulnerable join would + # surface it. Encoded ".." reaches the handler un-normalized (plain "../" + # is collapsed by the client before routing, so it only proves routing). + (tmp_path / "secret.json").write_text('{"leaked": true}') + client = _retrieval_client(tmp_path) + + resp = client.get("/v1/sessions/%2E%2E/completions") + assert resp.status_code == 404 + assert "leaked" not in resp.text + + +def test_retrieval_endpoint_registers_once() -> None: + from switchyard.lib.endpoints.token_capture_endpoint import ( + TokenCaptureSessionsEndpoint, + ) + + assert TokenCaptureSessionsEndpoint.register_once is True + + +# --------------------------------------------------------------------------- +# Route bundle + launcher wiring +# --------------------------------------------------------------------------- + + +def test_token_capture_engine_stripped_when_capture_disabled() -> None: + from switchyard.cli.route_bundle import _strip_token_capture_engine + + raw = { + "routes": { + "m1": { + "type": "model", + "target": {"model": "a", "base_url": "http://x/v1"}, + "token_capture_engine": "vllm", + }, + }, + } + stripped = _strip_token_capture_engine(raw) + assert "token_capture_engine" not in stripped["routes"]["m1"] + # Original mapping untouched; other keys intact. + assert raw["routes"]["m1"]["token_capture_engine"] == "vllm" + assert stripped["routes"]["m1"]["target"]["model"] == "a" + + +def test_capture_session_headers() -> None: + from switchyard.cli.launchers.launch_intake_config import ( + LaunchIntakeConfig, + capture_session_headers, + ) + + # Capture off -> no headers. + assert capture_session_headers(None, False, "claude") == {} + # Intake on -> its own headers already carry the session id. + intake = LaunchIntakeConfig.from_resolved( + base_url=None, + workspace=None, + api_key=None, + app="a", + task="t", + session_id="s1", + target="claude", + ) + assert capture_session_headers(intake, True, "claude") == {} + # Capture on, intake off -> per-launch session header. + headers = capture_session_headers(None, True, "claude") + assert set(headers) == {"proxy_x_session_id"} + assert headers["proxy_x_session_id"].startswith("claude-") + + +def test_capture_session_id_for_api_key_covers_intake_on() -> None: + from switchyard.cli.launchers.launch_intake_config import ( + LaunchIntakeConfig, + capture_session_id_for_api_key, + ) + + intake = LaunchIntakeConfig.from_resolved( + base_url=None, + workspace=None, + api_key=None, + app="a", + task="t", + session_id="openclaw-1700000000000-abcd1234", + target="openclaw", + ) + # Capture off -> no session id regardless of intake. + assert capture_session_id_for_api_key(intake, False, "openclaw") is None + # Capture on + intake on -> intake's session id rides the key (headers + # can't reach OpenClaw). + assert ( + capture_session_id_for_api_key(intake, True, "openclaw") + == "openclaw-1700000000000-abcd1234" + ) + # Capture on, intake off -> fresh launcher-shaped id. + generated = capture_session_id_for_api_key(None, True, "openclaw") + assert generated is not None and generated.startswith("openclaw-") + + +def test_openclaw_env_carries_marked_capture_session_id_as_api_key() -> None: + from switchyard.cli.launchers.openclaw_launcher import _openclaw_env + from switchyard.lib.processors.token_capture_request_processor import SESSION_API_KEY_MARKER + + env = _openclaw_env( + workspace="/tmp/ws", + capture_session_id="openclaw-1783600121000-a1b2c3d4", + ) + # The session id rides the API key behind the marker the server strips. + assert env["SWITCHYARD_API_KEY"] == SESSION_API_KEY_MARKER + "openclaw-1783600121000-a1b2c3d4" + + # Without capture, the opaque placeholder is unchanged. + assert _openclaw_env(workspace="/tmp/ws")["SWITCHYARD_API_KEY"] == "switchyard" + + +def test_claude_env_carries_capture_session_header() -> None: + from switchyard.cli.launchers.claude_code_launcher import _claude_env + + env = _claude_env(1234, "m", capture_headers={"proxy_x_session_id": "claude-1-ab"}) + assert env["ANTHROPIC_CUSTOM_HEADERS"] == "proxy_x_session_id: claude-1-ab" + assert env["SWITCHYARD_SESSION_ID"] == "claude-1-ab" + + # Without capture headers the env is unchanged from the intake-off default. + assert "ANTHROPIC_CUSTOM_HEADERS" not in _claude_env(1234, "m") diff --git a/tests/test_token_capture_target.py b/tests/test_token_capture_target.py new file mode 100644 index 00000000..abf12333 --- /dev/null +++ b/tests/test_token_capture_target.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Token-capture engine params derived into ``LlmTarget.extra_body``.""" + +from __future__ import annotations + +import pytest + +from switchyard.lib.backends.llm_target import ( + BackendFormat, + LlmTarget, + coerce_llm_target, + llm_target_with_token_capture, +) + + +def _target(**overrides: object) -> LlmTarget: + params: dict[str, object] = { + "id": "policy", + "model": "Qwen/Qwen3-0.6B", + "format": BackendFormat.OPENAI, + "base_url": "https://example.invalid/v1", + "api_key": "sk-test", + } + params.update(overrides) + return LlmTarget(**params) + + +class TestLlmTargetWithTokenCapture: + def test_vllm_params_derived_into_extra_body(self) -> None: + target = llm_target_with_token_capture(_target(), "vllm") + + assert target.extra_body == { + "logprobs": True, + "return_token_ids": True, + "top_logprobs": 0, + } + + def test_explicit_extra_body_key_wins_over_derived(self) -> None: + target = llm_target_with_token_capture(_target(extra_body={"top_logprobs": 5}), "vllm") + + assert target.extra_body == { + "logprobs": True, + "return_token_ids": True, + "top_logprobs": 5, + } + + def test_unrelated_extra_body_keys_preserved(self) -> None: + target = llm_target_with_token_capture( + _target(extra_body={"chat_template_kwargs": {"enable_thinking": False}}), + "vllm", + ) + + assert target.extra_body["chat_template_kwargs"] == {"enable_thinking": False} + assert target.extra_body["return_token_ids"] is True + + def test_target_identity_and_connection_fields_preserved(self) -> None: + source = _target() + target = llm_target_with_token_capture(source, "vllm") + + assert target.id == source.id + assert target.model == source.model + assert target.format == source.format + # Connection settings live on ``endpoint`` and must survive the rewrap. + assert target.endpoint.base_url == source.endpoint.base_url + assert target.endpoint.api_key == source.endpoint.api_key + assert target.endpoint.timeout_secs == source.endpoint.timeout_secs + + def test_unknown_engine_raises(self) -> None: + with pytest.raises(ValueError, match="unknown token_capture_engine 'tgi'"): + llm_target_with_token_capture(_target(), "tgi") + + +class TestCoerceLlmTargetTokenCaptureEngine: + def test_token_capture_engine_field_derives_params(self) -> None: + target = coerce_llm_target( + { + "model": "Qwen/Qwen3-0.6B", + "base_url": "https://example.invalid/v1", + "api_key": "sk-test", + "token_capture_engine": "vllm", + }, + default_id="policy", + ) + + assert target.extra_body == { + "logprobs": True, + "return_token_ids": True, + "top_logprobs": 0, + } + + def test_without_token_capture_engine_no_params(self) -> None: + target = coerce_llm_target( + {"model": "m", "base_url": "https://example.invalid/v1", "api_key": "k"}, + default_id="policy", + ) + + assert not target.extra_body diff --git a/tests/test_token_capture_translation.py b/tests/test_token_capture_translation.py new file mode 100644 index 00000000..d4befa07 --- /dev/null +++ b/tests/test_token_capture_translation.py @@ -0,0 +1,336 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Golden translation tests for the token-capture flow. + +Fixed request/response fixtures, asserted structurally against their validated +translated forms — covering the two translation boundaries token capture +relies on: + +* inbound harness request (Anthropic / Responses) → the OpenAI Chat body the + vLLM backend receives, and +* the synthesized OpenAI chunk stream (built from a captured buffered vLLM + completion) → the client's native stream events. + +The fixtures are plain dicts at module top so more validated pairs can be +appended without touching the test logic. Assertions are structural (event +types, block shapes, exact field values) rather than substring checks, so a +translation that emits the right strings in the wrong shape still fails. +""" + +from __future__ import annotations + +import json + +from switchyard.lib.processors.rl_logging_request_processor import RlLoggingRequestProcessor +from switchyard.lib.processors.token_capture_request_processor import ( + TokenCaptureRequestProcessor, +) +from switchyard.lib.processors.token_capture_response_processor import ( + TokenCaptureResponseProcessor, +) +from switchyard.lib.request_metadata import CTX_REQUEST_METADATA +from switchyard_rust.components import RequestMetadata +from switchyard_rust.core import ChatRequest, ChatRequestType, ChatResponse, ProxyContext +from switchyard_rust.translation import TranslationEngine + +_SESSION = "claude-1700000000000-abc12345" + +# --- Fixed inbound requests (what harnesses send) ---------------------------- + +ANTHROPIC_REQUEST = { + "model": "tito-model", + "max_tokens": 64, + "stream": True, + "system": "be brief", + "messages": [{"role": "user", "content": "say hello"}], +} + +# Turn 2 of a tool-using conversation: the shape capture sees on every call +# after the first (assistant tool_use + tool_result history). +ANTHROPIC_TOOL_HISTORY_REQUEST = { + "model": "tito-model", + "max_tokens": 64, + "messages": [ + {"role": "user", "content": "create hello.py"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Creating the file now."}, + { + "type": "tool_use", + "id": "call_abc123", + "name": "write_file", + "input": {"path": "hello.py"}, + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_abc123", + "content": "file written", + } + ], + }, + ], +} + +RESPONSES_REQUEST = { + "model": "tito-model", + "stream": True, + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "say hello"}], + } + ], +} + +RESPONSES_TOOL_HISTORY_REQUEST = { + "model": "tito-model", + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "create hello.py"}], + }, + { + "type": "function_call", + "call_id": "call_abc123", + "name": "write_file", + "arguments": '{"path": "hello.py"}', + }, + { + "type": "function_call_output", + "call_id": "call_abc123", + "output": "file written", + }, + ], +} + +# --- Fixed captured vLLM completion (what the backend returns) --------------- + +VLLM_COMPLETION_WITH_TOOL_CALL = { + "id": "chatcmpl-golden", + "object": "chat.completion", + "created": 1700000000, + "model": "Qwen/Qwen3-0.6B", + "prompt_token_ids": [1, 2, 3], + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Creating the file now.", + "tool_calls": [ + { + "id": "call_abc123", + "type": "function", + "function": { + "name": "write_file", + "arguments": '{"path": "hello.py"}', + }, + } + ], + }, + "finish_reason": "tool_calls", + "token_ids": [7, 8, 9], + "logprobs": { + "content": [ + {"token": "a", "logprob": -0.1}, + {"token": "b", "logprob": -0.2}, + {"token": "c", "logprob": -0.3}, + ] + }, + } + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 3, "total_tokens": 6}, +} + + +def _capture_ctx() -> ProxyContext: + ctx = ProxyContext() + ctx.metadata[CTX_REQUEST_METADATA] = RequestMetadata.from_headers( + {"proxy_x_session_id": _SESSION} + ) + return ctx + + +async def _synthesized_stream(tmp_path, inbound: ChatRequest) -> ChatResponse: + """Run the capture pair the way the chain does and return the synthetic stream.""" + ctx = _capture_ctx() + await RlLoggingRequestProcessor().process(ctx, inbound) + await TokenCaptureRequestProcessor().process(ctx, inbound) + return await TokenCaptureResponseProcessor(tmp_path).process( + ctx, ChatResponse.openai_completion(dict(VLLM_COMPLETION_WITH_TOOL_CALL)) + ) + + +# --- Request direction: harness format → upstream OpenAI Chat body ----------- + + +def test_anthropic_request_translates_to_openai_chat() -> None: + engine = TranslationEngine() + translated = engine.request_to( + ChatRequestType.OPENAI_CHAT, ChatRequest.anthropic(dict(ANTHROPIC_REQUEST)) + ) + body = dict(translated.body) + + assert body["model"] == "tito-model" + assert body["max_completion_tokens"] == 64 + roles = [m["role"] for m in body["messages"]] + assert roles == ["system", "user"] + assert body["messages"][0]["content"] == "be brief" + assert body["messages"][1]["content"] == "say hello" + + +def test_anthropic_tool_history_translates_to_openai_chat() -> None: + engine = TranslationEngine() + translated = engine.request_to( + ChatRequestType.OPENAI_CHAT, + ChatRequest.anthropic(dict(ANTHROPIC_TOOL_HISTORY_REQUEST)), + ) + messages = dict(translated.body)["messages"] + + roles = [m["role"] for m in messages] + assert roles == ["user", "assistant", "tool"] + + assistant = messages[1] + calls = assistant["tool_calls"] + assert len(calls) == 1 + assert calls[0]["id"] == "call_abc123" + assert calls[0]["function"]["name"] == "write_file" + assert json.loads(calls[0]["function"]["arguments"]) == {"path": "hello.py"} + + tool = messages[2] + # The tool result must round-trip keyed by the SAME call id. + assert tool["tool_call_id"] == "call_abc123" + assert "file written" in json.dumps(tool["content"]) + + +def test_responses_request_translates_to_openai_chat() -> None: + engine = TranslationEngine() + translated = engine.request_to( + ChatRequestType.OPENAI_CHAT, ChatRequest.openai_responses(dict(RESPONSES_REQUEST)) + ) + body = dict(translated.body) + + assert body["model"] == "tito-model" + user_messages = [m for m in body["messages"] if m["role"] == "user"] + assert len(user_messages) == 1 + content = user_messages[0]["content"] + # Typed input_text parts must arrive as clean text, not serialized JSON. + text = content if isinstance(content, str) else "".join( + part.get("text", "") for part in content + ) + assert text == "say hello" + assert "{" not in text + + +def test_responses_tool_history_translates_to_openai_chat() -> None: + engine = TranslationEngine() + translated = engine.request_to( + ChatRequestType.OPENAI_CHAT, + ChatRequest.openai_responses(dict(RESPONSES_TOOL_HISTORY_REQUEST)), + ) + messages = dict(translated.body)["messages"] + + assistant = next(m for m in messages if m["role"] == "assistant") + calls = assistant["tool_calls"] + assert calls[0]["id"] == "call_abc123" + assert calls[0]["function"]["name"] == "write_file" + assert json.loads(calls[0]["function"]["arguments"]) == {"path": "hello.py"} + + tool = next(m for m in messages if m["role"] == "tool") + assert tool["tool_call_id"] == "call_abc123" + assert "file written" in json.dumps(tool["content"]) + + +# --- Stream direction: synthetic OpenAI chunks → client-native events -------- + + +async def test_synthetic_stream_translates_to_anthropic_events(tmp_path) -> None: + inbound = ChatRequest.anthropic(dict(ANTHROPIC_REQUEST)) + out = await _synthesized_stream(tmp_path, inbound) + + events = [event async for event in TranslationEngine().stream_for_request(inbound, out)] + types = [e["type"] for e in events] + + # Anthropic stream grammar: proper frame, in order. + assert types[0] == "message_start" + assert types[-1] == "message_stop" + + # Text arrives as a text block with the exact delta — and ONLY the text: + # a tool call leaking into a text delta (the JSON-blob failure class) + # must fail here. + text_deltas = [ + e["delta"]["text"] + for e in events + if e["type"] == "content_block_delta" and e["delta"]["type"] == "text_delta" + ] + assert text_deltas == ["Creating the file now."] + + # The tool call arrives as a structured tool_use block with the same id. + tool_starts = [ + e["content_block"] + for e in events + if e["type"] == "content_block_start" and e["content_block"]["type"] == "tool_use" + ] + assert len(tool_starts) == 1 + assert tool_starts[0]["id"] == "call_abc123" + assert tool_starts[0]["name"] == "write_file" + json_deltas = [ + e["delta"]["partial_json"] + for e in events + if e["type"] == "content_block_delta" and e["delta"]["type"] == "input_json_delta" + ] + assert json.loads("".join(json_deltas)) == {"path": "hello.py"} + + # Terminal frame: tool_use stop reason and usage from the captured body. + message_delta = next(e for e in events if e["type"] == "message_delta") + assert message_delta["delta"]["stop_reason"] == "tool_use" + assert message_delta["usage"]["output_tokens"] == 6 - 3 + + +def _parse_sse(events: list[str]) -> list[dict]: + parsed = [] + for raw in events: + for line in raw.splitlines(): + if line.startswith("data: "): + parsed.append(json.loads(line[len("data: "):])) + return parsed + + +async def test_synthetic_stream_translates_to_responses_events(tmp_path) -> None: + inbound = ChatRequest.openai_responses(dict(RESPONSES_REQUEST)) + out = await _synthesized_stream(tmp_path, inbound) + + raw = [event async for event in TranslationEngine().stream_for_request(inbound, out)] + events = _parse_sse([str(e) for e in raw]) + types = [e["type"] for e in events] + + assert types[0] == "response.created" + assert types[-1] == "response.completed" + + # Text arrives only through output_text deltas, with the exact payload. + text_deltas = [e["delta"] for e in events if e["type"] == "response.output_text.delta"] + assert text_deltas == ["Creating the file now."] + + # The tool call arrives as a function_call output item with the same + # call id, plus its arguments through the dedicated delta channel. + fc_items = [ + e["item"] + for e in events + if e["type"] == "response.output_item.added" and e["item"]["type"] == "function_call" + ] + assert len(fc_items) == 1 + assert fc_items[0]["call_id"] == "call_abc123" + assert fc_items[0]["name"] == "write_file" + args_done = next( + e for e in events if e["type"] == "response.function_call_arguments.done" + ) + assert json.loads(args_done["arguments"]) == {"path": "hello.py"}