diff --git a/.agents/skills/switchyard-lib-core/SKILL.md b/.agents/skills/switchyard-lib-core/SKILL.md index de8b3b05..57ddb0b7 100644 --- a/.agents/skills/switchyard-lib-core/SKILL.md +++ b/.agents/skills/switchyard-lib-core/SKILL.md @@ -41,6 +41,8 @@ right validation set. If the change is driven by a launcher need, also read | Route YAML / model dispatch | Use `switchyard/cli/route_bundle.py` and `switchyard/lib/route_table_builders.py`. They build `RouteTable` entries from profile-backed runtimes and keep launchers plus `switchyard serve --routing-profiles` on one path. | | Shared/persistent session-affinity pins across workers or pod churn | Configure the latency route with `session_affinity: true` + `affinity_store: redis` + `affinity_store_url` (optional `affinity_store_ttl_seconds`, `affinity_key_prefix`). `SessionAffinity` keeps the Rust `SessionCache` as L1 and reads/writes through the `AffinityPinStore` L2 (`switchyard/lib/redis_pin_store.py`), fail-open behind a 0.1s socket timeout and a 3-failure/10s-cooldown circuit breaker (`switchyard_affinity_l2_breaker_open` gauge). Requires the `switchyard[affinity-redis]` extra. | | Stats / telemetry | Reuse `StatsRequestProcessor`, `StatsResponseProcessor`, `StatsLlmBackend`, and `StatsAccumulator`. A profile config should thread one accumulator through all three when stats are enabled. Do not write a parallel collector. | +| RL token capture | `TokenCaptureRequestProcessor` / `TokenCaptureResponseProcessor` (`switchyard/lib/processors/token_capture_*.py`), activated by `--enable-rl-logging` + a route's `token_capture_engine: vllm`. One record per call under `/sessions/`, retrieval via `GET /v1/sessions[...]` (`token_capture_endpoint.py`). Session id resolution ladder lives in the request processor. Tests: `tests/test_token_capture*.py`. | +| Token-continuity injection | `TokenInjectionBackend` (`switchyard/lib/backends/token_injection_backend.py`) wraps the chat backend when a capture route sets `token_injection: true` (+ `tool_parser`/`reasoning_parser` mirroring the vLLM server flags). Per-session conversation chains, env-suffix split at the end-of-turn token, raw token-ID generation via `/v1/completions`, parsing via `vllm_parsers.py` (lazy version-gated vLLM import). Wrapped through `with_runtime_components(backend_wrapper=...)`. Any failure degrades that call to capture-only. Tests: `tests/test_token_injection.py`. | | A fixed-path endpoint contributed by per-route components | Set `Endpoint.register_once = True`; `build_switchyard_app(...)` mounts the first instance while still running every component's lifecycle. Leave the default `False` for configurable endpoint classes that may mount distinct instances. | | Per-endpoint attribution on `/metrics` for a Python backend that can't be wrapped by `StatsLlmBackend` | Set `ctx.selected_model = endpoint_id` before returning the response. Also set `ctx.backend_call_latency_ms = upstream_call_ms` so the response processor can compute routing overhead. `LatencyServiceLLMBackend.call` is the reference. | | State metrics on `/metrics` | Register a `PrometheusEmitter` via `switchyard.lib.endpoints.prometheus_emitter.register(...)` and unregister on `shutdown()`. This is for backend-owned state, not request-flow counters. | 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..d7b879a9 100644 --- a/switchyard/cli/route_bundle.py +++ b/switchyard/cli/route_bundle.py @@ -22,9 +22,10 @@ argparse-driven and YAML-driven assembly share one builder. """ +import logging import os import re -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass, field, replace from importlib import import_module from pathlib import Path @@ -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,14 @@ 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", + # Opts this route into token-continuity injection (requires + # token_capture_engine). Parser names mirror the vLLM server flags. + "token_injection", + "tool_parser", + "reasoning_parser", }) _RANDOM_ROUTING_ROUTE_KEYS = _ROUTE_METADATA_KEYS | _TARGET_DEFAULT_ROUTE_KEYS | frozenset({ "strong", @@ -191,7 +202,14 @@ 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", + "token_injection", + "tool_parser", + "reasoning_parser", + }) | _PASSTHROUGH_SETTING_KEYS ) _LATENCY_ENDPOINT_DEFAULT_KEYS = frozenset({ @@ -418,6 +436,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 +449,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 +468,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 +477,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 +489,75 @@ 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 + # Injection rides on capture; without capture it must not + # activate (its records would go unstored). + if key == "token_injection": + 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`. @@ -573,7 +669,10 @@ def build_table_from_bundle( table.register( model_id, switchyard, - metadata=_route_metadata(model_id, route, route_type), + metadata={ + **_route_metadata(model_id, route, route_type), + **_capture_route_metadata_extras(model_id, route, route_type, route_defaults), + }, ) return table @@ -786,6 +885,7 @@ def _build_switchyard_for_route( enable_stats=_optional_bool(route.get("enable_stats"), default=True), extra_request_processors=pre_routing_request_processors, extra_response_processors=extra_response_processors, + backend_wrapper=_token_injection_wrapper(model_id, route, target), ) if route_type == "routellm": @@ -1129,10 +1229,61 @@ 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 _token_injection_wrapper( + model_id: str, + route: Mapping[str, object], + target: LlmTarget, +) -> Callable[[Any], Any] | None: + """Backend wrapper enabling token-continuity injection, or ``None``. + + ``token_injection: true`` opts the route in; it requires + ``token_capture_engine`` (injection's records must be captured) and an + explicit target ``base_url`` for the direct ``/tokenize`` and + ``/v1/completions`` calls. Parser names are optional and mirror the vLLM + server's ``--tool-call-parser`` / ``--reasoning-parser`` flags. + """ + if not route.get("token_injection"): + return None + if route.get("token_capture_engine") is None: + raise RouteBundleConfigError( + f"route {model_id!r}: token_injection requires token_capture_engine" + ) + endpoint = target.endpoint + base_url = endpoint.base_url if endpoint is not None else None + if not base_url: + raise RouteBundleConfigError( + f"route {model_id!r}: token_injection requires an explicit base_url" + ) + tool_parser = _optional_str(route["tool_parser"]) if "tool_parser" in route else None + reasoning_parser = ( + _optional_str(route["reasoning_parser"]) if "reasoning_parser" in route else None ) + api_key = endpoint.api_key if endpoint is not None else None + timeout_secs = (endpoint.timeout_secs if endpoint is not None else None) or 600.0 + model = target.model + + def _wrap(backend: Any) -> Any: + from switchyard.lib.backends.token_injection_backend import TokenInjectionBackend + + return TokenInjectionBackend( + backend, + base_url=base_url, + model=model, + api_key=api_key, + tool_parser=tool_parser, + reasoning_parser=reasoning_parser, + timeout_secs=timeout_secs, + ) + + return _wrap def _latency_service_switchyard( @@ -1393,6 +1544,54 @@ def _validate_allowed_keys( ) +def _capture_route_metadata_extras( + model_id: str, + route: Mapping[str, object], + route_type: str, + target_defaults: Mapping[str, object], +) -> dict[str, object]: + """Upstream-model facts exposed on ``/v1/models`` for token-capture routes. + + Training consumers (e.g. NeMo Gym's opencode provider config) need the + served model's context length; the route target is the only component that + always knows the engine's address, so the fact is surfaced here. Absent on + fetch failure — consumers degrade gracefully. + """ + if route_type not in ("model", "passthrough") or route.get("token_capture_engine") is None: + return {} + try: + target = _passthrough_target(model_id, route, target_defaults, route_type) + except (RouteBundleConfigError, TypeError, ValueError): + return {} + endpoint = target.endpoint + base_url = endpoint.base_url if endpoint is not None else None + if not base_url: + return {} + max_model_len = _fetch_upstream_max_model_len(base_url, endpoint.api_key) + return {"max_model_len": max_model_len} if max_model_len is not None else {} + + +def _fetch_upstream_max_model_len(base_url: str, api_key: str | None) -> int | None: + """The engine's ``max_model_len`` from its ``/v1/models`` card, or ``None``.""" + import httpx + + headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} + try: + response = httpx.get( + f"{base_url.rstrip('/')}/models", headers=headers, timeout=5.0 + ) + response.raise_for_status() + data = response.json().get("data") + except Exception: # noqa: BLE001 - metadata is best-effort at startup + logger.debug("could not fetch max_model_len from %s", base_url, exc_info=True) + return None + for entry in data if isinstance(data, list) else []: + value = entry.get("max_model_len") if isinstance(entry, dict) else None + if isinstance(value, int) and not isinstance(value, bool): + return value + return None + + def _route_metadata( model_id: str, route: Mapping[str, object], @@ -1506,5 +1705,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..aa4a27b3 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,13 @@ 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; `token_injection: true` additionally " + "preserves multi-turn token continuity for reasoning models. " + "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/backends/token_injection_backend.py b/switchyard/lib/backends/token_injection_backend.py new file mode 100644 index 00000000..dd945948 --- /dev/null +++ b/switchyard/lib/backends/token_injection_backend.py @@ -0,0 +1,457 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Stateful vLLM backend wrapper that preserves token continuity by prompt injection.""" + +from __future__ import annotations + +import logging +import time +from dataclasses import dataclass +from typing import Any + +import httpx + +from switchyard.lib.backends.vllm_parsers import ( + VllmParserError, + parse_generation, + parse_guided_tool_array, + tool_choice_json_schema, +) +from switchyard.lib.processors.token_capture_request_processor import CTX_TOKEN_CAPTURE_SESSION +from switchyard.lib.proxy_context import ProxyContext +from switchyard.lib.roles import LLMBackend +from switchyard_rust.core import ChatRequest, ChatResponse + +logger = logging.getLogger(__name__) + +JsonObject = dict[str, Any] + +#: finish_reasons where the model emitted its natural end-of-turn token itself, +#: making the generation's final token usable as the EOT split marker. +_NATURAL_STOP_REASONS = frozenset({"stop", "tool_calls", "stop_sequence"}) + +#: Clamp headroom: the chat endpoint can render the prompt a few tokens longer +#: than /tokenize does (measured on vLLM 0.24.0 + Qwen3: the chat path's +#: disabled-thinking render pre-fills an empty think block, exactly 4 tokens). +#: The pre-fill size is a template property, so keep headroom beyond the +#: observed value rather than pinning to it. +_CONTEXT_MARGIN_TOKENS = 8 + +#: Chat request params forwarded verbatim onto the completions request. +_FORWARDED_SAMPLING_KEYS = ( + "temperature", + "top_p", + "stop", + "seed", + "frequency_penalty", + "presence_penalty", + "logit_bias", +) + + +@dataclass +class _Chain: + """Canonical token history of one conversation within a capture session. + + Harnesses interleave several conversations under one session id (e.g. + opencode's title-generator side-call rides the main agent's session), so a + session holds a list of chains and each call is routed to the chain it + extends. ``prefix`` is the no-generation-prompt template tokenization of + the chain's last call — the stable split offset for env-suffix extraction + (the generation-prompt suffix differs between ``/tokenize`` and chat + completions when a reasoning parser is active, so full-prompt offsets are + not comparable). + """ + + accumulated: list[int] + prefix: list[int] + #: The chain's last raw generation; its final token on a natural stop is + #: the end-of-turn marker used to slice the env suffix. Verified: vLLM + #: includes the stop token in token_ids on both 0.11 and 0.24. + last_generation: list[int] + #: End-of-turn token id, observed from a naturally stopped generation. + #: None (e.g. after a length-stop) means the chain cannot be extended and + #: is reseeded by the next matching call instead. + eot_token_id: int | None + + +class TokenInjectionBackend(LLMBackend): + """Guarantee per-session token continuity by injecting raw prompt token IDs. + + Wraps the normal chat backend. Each call is matched to the session chain + whose history it extends; matched calls bypass chat-template re-rendering + — which strips historical reasoning for reasoning models — by sending the + chain's accumulated history plus the newly tokenized environment suffix + directly to vLLM's ``/v1/completions`` as token IDs. The raw generation is + parsed with vLLM's own parsers and returned as a chat-shaped response + carrying the same engine token fields the capture processor already reads. + + A call that matches no chain is served via the wrapped backend and seeds a + new chain from its response, so interleaved side-calls and arrival order + never disable injection for the real conversation. Any failure (missing + token fields, parser errors, transport errors) also falls back to the + wrapped backend — behavior then equals capture-only mode and Gym's + validation masks the affected sample. Chain state is in-process only: one + proxy instance must own all calls of a session. + """ + + def __init__( + self, + inner: LLMBackend, + *, + base_url: str, + model: str, + api_key: str | None = None, + tool_parser: str | None = None, + reasoning_parser: str | None = None, + timeout_secs: float = 600.0, + ) -> None: + self._inner = inner + # /tokenize lives at the server root, not under /v1. + self._v1_url = base_url.rstrip("/") + self._root_url = self._v1_url.removesuffix("/v1") + self._model = model + self._api_key = api_key + self._tool_parser = tool_parser + self._reasoning_parser = reasoning_parser + self._timeout = timeout_secs + self._sessions: dict[str, list[_Chain]] = {} + # Model context length, learned from /tokenize responses; used to clamp + # the completion budget the way older vLLM chat endpoints did implicitly + # (newer versions reject over-budget requests outright, on both paths). + self._max_model_len: int | None = None + + async def call(self, ctx: ProxyContext, request: ChatRequest) -> ChatResponse: + """Serve one model call, injecting the matching chain's history when possible.""" + session_id = ctx.metadata.get(CTX_TOKEN_CAPTURE_SESSION) + if not isinstance(session_id, str) or not session_id: + return await self._inner.call(ctx, request) + body = dict(request.body) + if body.get("n", 1) not in (None, 1): + return await self._inner.call(ctx, request) + + chains = self._sessions.setdefault(session_id, []) + rendered: list[int] | None = None + try: + chain = None + if chains: + rendered = await self._tokenize(body, add_generation_prompt=True) + chain = _longest_matching_chain(chains, rendered) + if rendered is not None and chain is not None and chain.eot_token_id is not None: + return await self._injected_call(chain, body, rendered) + + # No chain extends this call (first call, side-call, rewritten + # history) or the matched chain has no end-of-turn marker: serve + # via the chat path and (re)seed a chain from the response. + prefix = await self._tokenize(body, add_generation_prompt=False) + response = await self._serve_inner(ctx, request, body, prompt_len=len(prefix)) + self._seed_chain(chains, chain, response, prefix) + return response + except (httpx.HTTPError, VllmParserError, KeyError, ValueError) as exc: + logger.warning( + "token injection: session %s: serving via the chat path: %s", session_id, exc + ) + if rendered is not None: + # The budget clamp must apply on this path too, or the call + # trades an injection error for a context-window rejection. + return await self._serve_inner(ctx, request, body, prompt_len=len(rendered)) + return await self._inner.call(ctx, request) + + async def _serve_inner( + self, ctx: ProxyContext, request: ChatRequest, body: JsonObject, prompt_len: int + ) -> ChatResponse: + """Delegate to the chat backend with the completion budget clamped. + + The clamp applies on this path too: newer vLLM rejects + ``prompt + max_tokens > max_model_len`` on chat completions as well. + """ + max_tokens = self._resolve_max_tokens(body, prompt_len) + if max_tokens is not None and max_tokens != body.get("max_tokens"): + body = dict(body) + body["max_tokens"] = max_tokens + body.pop("max_completion_tokens", None) + request.replace_body(body) + return await self._inner.call(ctx, request) + + def _seed_chain( + self, + chains: list[_Chain], + chain: _Chain | None, + response: ChatResponse, + prefix: list[int], + ) -> None: + """Start (or refresh) a chain from a chat-path response's token fields. + + Responses without engine token fields seed nothing — the call stays + capture-only and a later call may seed the chain instead. + """ + try: + body = dict(response.body) + except (TypeError, ValueError): + return + harvested = _harvest_token_fields(body) + if harvested is None: + return + prompt_ids, generation_ids, finish_reason = harvested + if prompt_ids[: len(prefix)] != prefix: + return + seeded = _Chain( + accumulated=prompt_ids + generation_ids, + prefix=prefix, + last_generation=generation_ids, + eot_token_id=( + generation_ids[-1] if finish_reason in _NATURAL_STOP_REASONS else None + ), + ) + if chain is not None: + chains[chains.index(chain)] = seeded + else: + chains.append(seeded) + + async def _injected_call( + self, chain: _Chain, body: JsonObject, rendered: list[int] + ) -> ChatResponse: + """Generate through ``/v1/completions`` with the chain's token history.""" + if chain.eot_token_id is None: # caller guarantees; narrows the type + raise ValueError("chain has no end-of-turn marker") + env_tokens = _slice_env_suffix( + rendered[len(chain.prefix):], chain.last_generation, chain.eot_token_id + ) + injected_prompt = chain.accumulated + env_tokens + + completion = await self._completions(body, injected_prompt) + choice = completion["choices"][0] + prompt_ids = choice.get("prompt_token_ids") or completion.get("prompt_token_ids") + generation_ids = choice.get("token_ids") + if prompt_ids != injected_prompt: + raise ValueError("completions endpoint did not echo the injected prompt") + if not isinstance(generation_ids, list) or not generation_ids: + raise ValueError("completions response carries no generation token ids") + + chat_body = self._synthesize_chat_body(body, completion, injected_prompt) + + # State advances only after a fully successful call, so a failed call + # falls back without corrupting the history. + chain.accumulated = injected_prompt + generation_ids + chain.prefix = await self._tokenize(body, add_generation_prompt=False) + chain.last_generation = generation_ids + if choice.get("finish_reason") in _NATURAL_STOP_REASONS: + chain.eot_token_id = generation_ids[-1] + return ChatResponse.openai_completion(chat_body) + + def _synthesize_chat_body( + self, request_body: JsonObject, completion: JsonObject, injected_prompt: list[int] + ) -> JsonObject: + """Chat-completions-shaped body carrying the engine token fields. + + The shape matches what vLLM's chat endpoint returns with the capture + params on, so the capture response processor and the stream + synthesizer consume it unchanged. + """ + choice = completion["choices"][0] + text = choice.get("text") or "" + tools = request_body.get("tools") if isinstance(request_body.get("tools"), list) else None + tool_choice = request_body.get("tool_choice") + + if tool_choice_json_schema(tools or [], tool_choice) is not None: + # Guided decoding: the generation is the schema-constrained JSON + # array itself; the tool parser never applies. + tool_calls = parse_guided_tool_array(text) + reasoning_content: str | None = None + content: str | None = None + else: + parsed = parse_generation( + text, + model=self._model, + tools=tools, + tool_parser=self._tool_parser, + reasoning_parser=self._reasoning_parser, + ) + tool_calls = parsed.tool_calls + reasoning_content = parsed.reasoning_content + content = parsed.content + + message: JsonObject = {"role": "assistant", "content": content} + if reasoning_content: + # Both spellings: vLLM chat responses renamed reasoning_content -> + # reasoning around 0.24; harnesses may read either. + message["reasoning_content"] = reasoning_content + message["reasoning"] = reasoning_content + if tool_calls: + message["tool_calls"] = tool_calls + + finish_reason = choice.get("finish_reason") or "stop" + if tool_calls and finish_reason == "stop": + finish_reason = "tool_calls" + + generation_ids = choice.get("token_ids") or [] + return { + "id": completion.get("id") or "chatcmpl-token-injection", + "object": "chat.completion", + "created": completion.get("created") or int(time.time()), + "model": completion.get("model") or self._model, + "prompt_token_ids": injected_prompt, + "choices": [ + { + "index": 0, + "message": message, + "finish_reason": finish_reason, + "token_ids": generation_ids, + "logprobs": _chat_shaped_logprobs(choice.get("logprobs"), generation_ids), + } + ], + "usage": completion.get("usage"), + } + + def _resolve_max_tokens(self, body: JsonObject, prompt_len: int) -> int | None: + """The effective completion budget for a call with a *prompt_len*-token prompt. + + Honors the smaller of ``max_tokens`` / ``max_completion_tokens`` when + both are present, clamped to the remaining model context when the + context length is known (learned from ``/tokenize`` responses). When + no budget fits at all, the requested value is returned unchanged — a + genuinely exhausted context should surface as the engine's own error. + """ + requested = [ + value + for value in (body.get("max_tokens"), body.get("max_completion_tokens")) + if isinstance(value, int) and not isinstance(value, bool) + ] + max_tokens = min(requested) if requested else None + if max_tokens is None or self._max_model_len is None: + return max_tokens + remaining = self._max_model_len - prompt_len - _CONTEXT_MARGIN_TOKENS + if remaining <= 0: + return max_tokens + return min(max_tokens, remaining) + + async def _tokenize(self, body: JsonObject, *, add_generation_prompt: bool) -> list[int]: + """Template-render and tokenize the request's messages without generating.""" + payload: JsonObject = { + "model": self._model, + "messages": body.get("messages") or [], + "add_generation_prompt": add_generation_prompt, + } + if isinstance(body.get("tools"), list): + payload["tools"] = body["tools"] + if isinstance(body.get("chat_template_kwargs"), dict): + payload["chat_template_kwargs"] = body["chat_template_kwargs"] + data = await self._post(f"{self._root_url}/tokenize", payload) + max_model_len = data.get("max_model_len") + if isinstance(max_model_len, int) and not isinstance(max_model_len, bool): + self._max_model_len = max_model_len + tokens = data.get("tokens") + if not isinstance(tokens, list) or not all(isinstance(t, int) for t in tokens): + raise ValueError("/tokenize response has no integer token list") + return tokens + + async def _completions(self, body: JsonObject, prompt: list[int]) -> JsonObject: + """One ``/v1/completions`` generation with the injected token-ID prompt.""" + payload: JsonObject = { + "model": self._model, + "prompt": prompt, + "stream": False, + "return_token_ids": True, + "logprobs": 0, + } + max_tokens = self._resolve_max_tokens(body, len(prompt)) + if max_tokens is not None: + payload["max_tokens"] = max_tokens + for key in _FORWARDED_SAMPLING_KEYS: + if body.get(key) is not None: + payload[key] = body[key] + tools = body.get("tools") if isinstance(body.get("tools"), list) else None + schema = tool_choice_json_schema(tools or [], body.get("tool_choice")) + if schema is not None: + payload["structured_outputs"] = {"json": schema} + return await self._post(f"{self._v1_url}/completions", payload) + + async def _post(self, url: str, payload: JsonObject) -> JsonObject: + headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {} + async with httpx.AsyncClient(timeout=self._timeout) as client: + response = await client.post(url, json=payload, headers=headers) + response.raise_for_status() + data = response.json() + if not isinstance(data, dict): + raise ValueError(f"{url} returned a non-object JSON body") + return data + + +def _longest_matching_chain(chains: list[_Chain], rendered: list[int]) -> _Chain | None: + """The chain whose prefix the rendered prompt extends; longest match wins.""" + best: _Chain | None = None + for chain in chains: + n = len(chain.prefix) + if (best is None or n > len(best.prefix)) and rendered[:n] == chain.prefix: + best = chain + return best + + +def _harvest_token_fields(body: JsonObject) -> tuple[list[int], list[int], Any] | None: + """``(prompt_token_ids, generation_token_ids, finish_reason)`` from a raw chat body.""" + choices = body.get("choices") + if not isinstance(choices, list) or len(choices) != 1 or not isinstance(choices[0], dict): + return None + choice = choices[0] + prompt_ids = body.get("prompt_token_ids") + generation_ids = choice.get("token_ids") + if not isinstance(prompt_ids, list) or not isinstance(generation_ids, list): + return None + if not _is_token_list(prompt_ids) or not _is_token_list(generation_ids): + return None + return list(prompt_ids), list(generation_ids), choice.get("finish_reason") + + +def _is_token_list(value: object) -> bool: + return ( + isinstance(value, list) + and bool(value) + and all(isinstance(item, int) and not isinstance(item, bool) for item in value) + ) + + +def _slice_env_suffix( + canonical_tail: list[int], last_generation: list[int], eot_token_id: int +) -> list[int]: + """Environment tokens added since the chain's last call (tool results, user turns, glue). + + ``canonical_tail`` re-renders the previous assistant turn (reasoning + stripped by the template) followed by the new environment messages; the + first end-of-turn token marks where the re-rendered turn ends. This never + renders a conversation ENDING with an assistant message, so thinking + templates' special-casing of a final assistant turn cannot affect it. + When the raw generation already ended with the end-of-turn token it is + skipped in the tail to avoid duplication; otherwise it is kept so the + stream still closes the assistant turn. + """ + try: + index = canonical_tail.index(eot_token_id) + except ValueError: + raise ValueError( + "end-of-turn token not found in the re-rendered prompt tail " + f"(eot={eot_token_id}, tail starts {canonical_tail[:8]})" + ) from None + if last_generation and last_generation[-1] == eot_token_id: + return canonical_tail[index + 1:] + return canonical_tail[index:] + + +def _chat_shaped_logprobs(logprobs: object, generation_ids: list[int]) -> JsonObject | None: + """Convert completions-endpoint logprobs to the chat shape capture reads. + + Completions returns ``{token_logprobs: [...], tokens: [...]}``; the capture + processor reads ``{content: [{token_id, logprob}, ...]}`` aligned with the + generation token ids. + """ + if not isinstance(logprobs, dict): + return None + token_logprobs = logprobs.get("token_logprobs") + if not isinstance(token_logprobs, list) or len(token_logprobs) != len(generation_ids): + return None + return { + "content": [ + {"token_id": token_id, "logprob": logprob} + for token_id, logprob in zip(generation_ids, token_logprobs, strict=True) + ] + } diff --git a/switchyard/lib/backends/vllm_parsers.py b/switchyard/lib/backends/vllm_parsers.py new file mode 100644 index 00000000..c3060047 --- /dev/null +++ b/switchyard/lib/backends/vllm_parsers.py @@ -0,0 +1,281 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lazy, version-gated access to vLLM's tool-call and reasoning parsers. + +The token-injection backend generates through vLLM's ``/v1/completions`` +endpoint, which returns raw text — the chat endpoint's tool-call and +reasoning parsing never runs. This module re-runs that parsing in the proxy +using vLLM's own parser registries, so the parsed output matches what the +chat endpoint would have produced for the same serving version. + +vLLM is an optional runtime dependency: it is imported on first use only, so +capture-only deployments and environments without vLLM are unaffected. The +import paths are gated on the installed vLLM version (verified boundaries: +tool parsers moved at 0.13.0, the OpenAI protocol module split at 0.15.0). +Patched vLLM builds may not match the upstream layout for their reported +version, so a failed import names the expected module path explicitly. +""" + +from __future__ import annotations + +import json +import logging +import uuid as uuid_lib +from dataclasses import dataclass, field +from typing import Any + +logger = logging.getLogger(__name__) + +JsonObject = dict[str, Any] + + +class VllmParserError(RuntimeError): + """vLLM's parsers are unavailable or failed for this configuration.""" + + +@dataclass +class ParsedGeneration: + """Structured view of one raw completions-endpoint generation.""" + + content: str | None + reasoning_content: str | None = None + tool_calls: list[JsonObject] = field(default_factory=list) + + +@dataclass +class _ParserRuntime: + """Resolved vLLM parser instances for one (model, config) pair.""" + + tool_parser: Any | None + reasoning_parser: Any | None + chat_request_cls: Any + + +_runtime_cache: dict[tuple[str, str | None, str | None], _ParserRuntime] = {} + + +def parse_generation( + text: str, + *, + model: str, + tools: list[JsonObject] | None, + tool_parser: str | None, + reasoning_parser: str | None, +) -> ParsedGeneration: + """Split raw generated text into reasoning, content, and structured tool calls. + + Parser names mirror the flags the vLLM server was launched with + (``--tool-call-parser``, ``--reasoning-parser``). A ``None`` parser name + skips that stage: content passes through unparsed. Raises + :class:`VllmParserError` when a named parser cannot be loaded. + """ + if tool_parser is None and reasoning_parser is None: + return ParsedGeneration(content=text) + + runtime = _resolve_runtime(model, tool_parser, reasoning_parser) + request = runtime.chat_request_cls( + messages=[{"role": "user", "content": ""}], + model=model, + tools=tools or None, + ) + + reasoning_content: str | None = None + content: str | None = text + if runtime.reasoning_parser is not None: + # Renamed upstream: extract_reasoning_content (<= 0.11.x) -> extract_reasoning. + extract = getattr(runtime.reasoning_parser, "extract_reasoning_content", None) + if extract is None: + extract = runtime.reasoning_parser.extract_reasoning + reasoning_content, content = extract(text, request=request) + + tool_calls: list[JsonObject] = [] + if runtime.tool_parser is not None and content: + extracted = runtime.tool_parser.extract_tool_calls(content, request=request) + if getattr(extracted, "tools_called", False): + tool_calls = [ + _tool_call_to_dict(call) for call in getattr(extracted, "tool_calls", []) + ] + content = getattr(extracted, "content", None) + return ParsedGeneration( + content=content, reasoning_content=reasoning_content, tool_calls=tool_calls + ) + + +def parse_guided_tool_array(text: str) -> list[JsonObject]: + """Parse a ``tool_choice: required``/named guided-decoding generation. + + Under guided decoding the generation is the schema-constrained JSON array + of ``{name, parameters}`` objects itself (see :func:`tool_choice_json_schema`) + — the tool parser never applies. Mirrors vLLM's chat-endpoint behavior. + """ + calls = json.loads(text) + if not isinstance(calls, list): + raise VllmParserError("guided tool generation is not a JSON array") + tool_calls: list[JsonObject] = [] + for call in calls: + if not isinstance(call, dict) or not isinstance(call.get("name"), str): + raise VllmParserError("guided tool generation entry is not {name, parameters}") + tool_calls.append({ + "id": f"chatcmpl-tool-{uuid_lib.uuid4().hex}", + "type": "function", + "function": { + "name": call["name"], + "arguments": json.dumps(call.get("parameters") or {}), + }, + }) + return tool_calls + + +def tool_choice_json_schema(tools: list[JsonObject], tool_choice: object) -> JsonObject | None: + """The guided-decoding JSON schema for ``required``/named tool choice, else ``None``. + + Ports vLLM's ``ChatCompletionRequest._get_json_schema_from_tool`` so the + completions-endpoint request constrains generation exactly as the chat + endpoint would. ``auto``/``none``/unset need no schema. + """ + if not tools or tool_choice in (None, "auto", "none"): + return None + + if isinstance(tool_choice, dict): + name = ((tool_choice.get("function") or {}) or {}).get("name") + by_name = { + (tool.get("function") or {}).get("name"): (tool.get("function") or {}) + for tool in tools + } + if name not in by_name: + raise VllmParserError(f"named tool {name!r} not present in tools") + return by_name[name].get("parameters") or {"type": "object", "properties": {}} + + if tool_choice != "required": + return None + + def tool_schema(function: JsonObject) -> JsonObject: + return { + "properties": { + "name": {"type": "string", "enum": [function.get("name")]}, + "parameters": function.get("parameters") + or {"type": "object", "properties": {}}, + }, + "required": ["name", "parameters"], + } + + return { + "type": "array", + "minItems": 1, + "items": { + "type": "object", + "anyOf": [tool_schema(tool.get("function") or {}) for tool in tools], + }, + } + + +def _resolve_runtime( + model: str, tool_parser: str | None, reasoning_parser: str | None +) -> _ParserRuntime: + key = (model, tool_parser, reasoning_parser) + cached = _runtime_cache.get(key) + if cached is not None: + return cached + + tool_parser_cls, reasoning_parser_cls, chat_request_cls = _import_vllm_parsers( + tool_parser, reasoning_parser + ) + tokenizer = _load_tokenizer(model) + runtime = _ParserRuntime( + tool_parser=tool_parser_cls(tokenizer) if tool_parser_cls is not None else None, + reasoning_parser=( + reasoning_parser_cls(tokenizer) if reasoning_parser_cls is not None else None + ), + chat_request_cls=chat_request_cls, + ) + _runtime_cache[key] = runtime + return runtime + + +def _import_vllm_parsers( + tool_parser: str | None, reasoning_parser: str | None +) -> tuple[Any | None, Any | None, Any]: + """Import parser classes from the installed vLLM, gated on its version. + + Verified layout boundaries: ``vllm.tool_parsers`` exists from 0.13.0 + (``vllm.entrypoints.openai.tool_parsers`` before), ``vllm.reasoning`` + never moved, and the monolithic ``vllm.entrypoints.openai.protocol`` + split into per-endpoint packages at 0.15.0. + """ + vllm_version = _installed_vllm_version() + + if vllm_version >= (0, 13, 0): + tool_parsers_module = "vllm.tool_parsers" + else: + tool_parsers_module = "vllm.entrypoints.openai.tool_parsers" + if vllm_version >= (0, 15, 0): + protocol_module = "vllm.entrypoints.openai.chat_completion.protocol" + else: + protocol_module = "vllm.entrypoints.openai.protocol" + + tool_parser_cls: Any | None = None + if tool_parser is not None: + manager = _import_module(tool_parsers_module, vllm_version).ToolParserManager + tool_parser_cls = manager.get_tool_parser(tool_parser) + reasoning_parser_cls: Any | None = None + if reasoning_parser is not None: + manager = _import_module("vllm.reasoning", vllm_version).ReasoningParserManager + reasoning_parser_cls = manager.get_reasoning_parser(reasoning_parser) + chat_request_cls = _import_module(protocol_module, vllm_version).ChatCompletionRequest + return tool_parser_cls, reasoning_parser_cls, chat_request_cls + + +def _import_module(name: str, vllm_version: tuple[int, ...]) -> Any: + import importlib + + try: + return importlib.import_module(name) + except ImportError as exc: + # Patched vLLM builds may not match the upstream layout for their + # reported version — name the expected path so the mismatch is obvious. + raise VllmParserError( + f"cannot import {name!r} from the installed vLLM " + f"(version {'.'.join(map(str, vllm_version))}); a patched build may " + f"not match the upstream module layout for its version" + ) from exc + + +def _installed_vllm_version() -> tuple[int, ...]: + from importlib.metadata import PackageNotFoundError, version + + try: + raw = version("vllm") + except PackageNotFoundError as exc: + raise VllmParserError( + "vLLM is not installed; token injection with configured parsers " + "requires the vllm package at runtime" + ) from exc + parts: list[int] = [] + for piece in raw.split(".")[:3]: + digits = "".join(ch for ch in piece if ch.isdigit()) + parts.append(int(digits) if digits else 0) + return tuple(parts) + + +def _load_tokenizer(model: str) -> Any: + """Fetch the model's tokenizer once; parser constructors require it.""" + try: + from transformers import AutoTokenizer + except ImportError as exc: + raise VllmParserError( + "transformers is required to construct vLLM parsers (tokenizer)" + ) from exc + return AutoTokenizer.from_pretrained(model) + + +def _tool_call_to_dict(call: Any) -> JsonObject: + """Normalize a vLLM ``ToolCall`` (pydantic model or dict) to a plain dict.""" + if isinstance(call, dict): + return call + dump = getattr(call, "model_dump", None) + if callable(dump): + result = dump() + if isinstance(result, dict): + return result + raise VllmParserError(f"unsupported tool call shape: {type(call).__name__}") 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/switchyard/lib/profiles/chain.py b/switchyard/lib/profiles/chain.py index 4cd1f140..bbabaa45 100644 --- a/switchyard/lib/profiles/chain.py +++ b/switchyard/lib/profiles/chain.py @@ -5,7 +5,7 @@ from __future__ import annotations -from collections.abc import Sequence +from collections.abc import Callable, Sequence from dataclasses import dataclass from typing import Any, cast @@ -108,12 +108,15 @@ def with_runtime_components( pre_request_processors: Sequence[Any] = (), post_request_processors: Sequence[Any] = (), response_processors: Sequence[Any] = (), + backend_wrapper: Callable[[LLMBackend], LLMBackend] | None = None, ) -> ComponentChainProfile: """Return a copy with serving-level stats and processor hooks applied. Profile configs stay user-facing and parseable; shared serving resources such as one route-table accumulator or Intake processors are attached by - the builder that hosts the profile. + the builder that hosts the profile. ``backend_wrapper`` wraps the + profile's backend (e.g. token injection) before stats attach, so stats + observe the wrapper's responses. """ from switchyard.lib.processors.stats_request_processor import ( StatsRequestProcessor, @@ -125,6 +128,8 @@ def with_runtime_components( request_chain: list[Any] = [] response_chain: list[Any] = list(self._response_processors) backend = self._backend + if backend_wrapper is not None: + backend = backend_wrapper(backend) stats: StatsAccumulator | None = None if enable_stats: diff --git a/switchyard/lib/route_table_builders.py b/switchyard/lib/route_table_builders.py index f4bf5c49..c0f5978f 100644 --- a/switchyard/lib/route_table_builders.py +++ b/switchyard/lib/route_table_builders.py @@ -97,6 +97,7 @@ def build_tier_passthrough_switchyard( enable_stats: bool = True, extra_request_processors: Sequence[Any] = (), extra_response_processors: Sequence[Any] = (), + backend_wrapper: Callable[[Any], Any] | None = None, ) -> ChainRuntime: """Build a single-tier chain for an explicitly selected model.""" from switchyard.lib.profiles import PassthroughProfileConfig, ProfileSwitchyard @@ -111,6 +112,7 @@ def build_tier_passthrough_switchyard( stats_accumulator=stats, pre_request_processors=extra_request_processors, response_processors=extra_response_processors, + backend_wrapper=backend_wrapper, ) ) 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"} diff --git a/tests/test_token_injection.py b/tests/test_token_injection.py new file mode 100644 index 00000000..2c7d906a --- /dev/null +++ b/tests/test_token_injection.py @@ -0,0 +1,635 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the token-continuity injection backend.""" + +from __future__ import annotations + +import json +from typing import Any + +import httpx +import respx + +from switchyard.lib.backends.token_injection_backend import ( + TokenInjectionBackend, + _slice_env_suffix, +) +from switchyard.lib.backends.vllm_parsers import ( + parse_generation, + tool_choice_json_schema, +) +from switchyard.lib.processors.token_capture_request_processor import CTX_TOKEN_CAPTURE_SESSION +from switchyard.lib.roles import LLMBackend +from switchyard_rust.core import ChatRequest, ChatResponse, ProxyContext + +_BASE_URL = "http://vllm.test/v1" +_MODEL = "Qwen/Qwen3-0.6B" +_EOT = 9 + +_TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Weather for a city.", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + }, + } +] + + +class _FakeInner(LLMBackend): + """Chat backend double returning preset vLLM-shaped completions. + + Accepts one body (returned for every call) or a list of bodies (consumed + in call order, last one repeating). Records each request body it serves. + """ + + def __init__(self, body: dict[str, Any] | list[dict[str, Any]] | None = None) -> None: + self.calls = 0 + self.seen_bodies: list[dict[str, Any]] = [] + bodies = body if body is not None else _first_turn_body() + self._bodies = bodies if isinstance(bodies, list) else [bodies] + + async def call(self, ctx: ProxyContext, request: ChatRequest) -> ChatResponse: + self.seen_bodies.append(dict(request.body)) + body = self._bodies[min(self.calls, len(self._bodies) - 1)] + self.calls += 1 + return ChatResponse.openai_completion(body) + + +def _first_turn_body() -> dict[str, Any]: + """vLLM chat body for the seed call: p1=[1,2,3], g1=[4,5,EOT], natural stop.""" + return { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1700000000, + "model": _MODEL, + "prompt_token_ids": [1, 2, 3], + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "step one"}, + "finish_reason": "stop", + "token_ids": [4, 5, _EOT], + "logprobs": {"content": [{"logprob": -0.1}, {"logprob": -0.2}, {"logprob": -0.3}]}, + } + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 3, "total_tokens": 6}, + } + + +def _backend(inner: _FakeInner | None = None, **kwargs: Any) -> TokenInjectionBackend: + return TokenInjectionBackend( + inner or _FakeInner(), + base_url=_BASE_URL, + model=_MODEL, + **kwargs, + ) + + +def _ctx(session: str | None = "sess-1") -> ProxyContext: + ctx = ProxyContext() + if session is not None: + ctx.metadata[CTX_TOKEN_CAPTURE_SESSION] = session + return ctx + + +def _chat_request(turns: int = 1) -> ChatRequest: + messages: list[dict[str, Any]] = [{"role": "user", "content": "hello"}] + if turns > 1: + messages += [ + {"role": "assistant", "content": "step one"}, + {"role": "user", "content": "continue"}, + ] + if turns > 2: + messages += [ + {"role": "assistant", "content": "step two"}, + {"role": "user", "content": "continue"}, + ] + return ChatRequest.openai_chat({"model": _MODEL, "messages": messages}) + + +def _mock_tokenize(respx_mock: respx.MockRouter, token_lists: list[list[int]]) -> None: + """Queue /tokenize responses in call order.""" + respx_mock.post("http://vllm.test/tokenize").mock( + side_effect=[httpx.Response(200, json={"tokens": tokens}) for tokens in token_lists] + ) + + +def _completion_response(prompt: list[int], generation: list[int], text: str) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "cmpl-2", + "object": "text_completion", + "created": 1700000100, + "model": _MODEL, + "choices": [ + { + "index": 0, + "text": text, + "finish_reason": "stop", + "prompt_token_ids": prompt, + "token_ids": generation, + "logprobs": { + "tokens": ["x"] * len(generation), + "token_logprobs": [-0.5] * len(generation), + }, + } + ], + "usage": { + "prompt_tokens": len(prompt), + "completion_tokens": len(generation), + "total_tokens": len(prompt) + len(generation), + }, + }, + ) + + +# --------------------------------------------------------------------------- +# Injection path +# --------------------------------------------------------------------------- + + +@respx.mock +async def test_second_call_injects_contiguous_prompt(respx_mock: respx.MockRouter) -> None: + """The injected prompt is exactly accumulated history + env suffix.""" + inner = _FakeInner() + backend = _backend(inner) + ctx = _ctx() + + # Seed: prefix (no gen prompt) = [1, 2]; p1 = [1, 2, 3] starts with it. + _mock_tokenize( + respx_mock, + [ + [1, 2], # seed prefix + [1, 2, _EOT, 20, 21, 30], # turn-2 rendered (gen prompt on) + [1, 2, _EOT, 20, 21], # turn-2 prefix (gen prompt off) + ], + ) + injected = [1, 2, 3, 4, 5, _EOT, 20, 21, 30] + completions = respx_mock.post(f"{_BASE_URL}/completions").mock( + return_value=_completion_response(injected, [40, 41], "The answer") + ) + + await backend.call(ctx, _chat_request()) + response = await backend.call(ctx, _chat_request(turns=2)) + + assert inner.calls == 1 # only the seed call hit the chat path + sent = json.loads(completions.calls[0].request.content) + assert sent["prompt"] == injected + assert sent["return_token_ids"] is True + assert sent["logprobs"] == 0 + + body = dict(response.body) + assert body["prompt_token_ids"] == injected + choice = body["choices"][0] + assert choice["token_ids"] == [40, 41] + assert choice["message"]["content"] == "The answer" + assert choice["logprobs"]["content"] == [ + {"token_id": 40, "logprob": -0.5}, + {"token_id": 41, "logprob": -0.5}, + ] + # Contiguity by construction: the injected prompt extends p1 + g1. + history = [1, 2, 3, 4, 5, _EOT] + assert body["prompt_token_ids"][: len(history)] == history + + +@respx.mock +async def test_third_call_extends_advanced_state(respx_mock: respx.MockRouter) -> None: + """State advances across injected calls: call 3 extends call 2's tokens.""" + backend = _backend() + ctx = _ctx() + injected_2 = [1, 2, 3, 4, 5, _EOT, 20, 21, 30] + injected_3 = injected_2 + [40, 41, 50, 51] + _mock_tokenize( + respx_mock, + [ + [1, 2], + [1, 2, _EOT, 20, 21, 30], + [1, 2, _EOT, 20, 21], + [1, 2, _EOT, 20, 21, 41, 50, 51], # turn-3 rendered; EOT 41 ends prior turn + [1, 2, _EOT, 20, 21, 41, 50], + ], + ) + completions = respx_mock.post(f"{_BASE_URL}/completions").mock( + side_effect=[ + _completion_response(injected_2, [40, 41], "step two"), + _completion_response(injected_3, [60], "done"), + ] + ) + + await backend.call(ctx, _chat_request()) + await backend.call(ctx, _chat_request(turns=2)) + await backend.call(ctx, _chat_request(turns=3)) + + sent = json.loads(completions.calls[1].request.content) + assert sent["prompt"] == injected_3 + + +@respx.mock +async def test_max_tokens_takes_smaller_of_both_keys(respx_mock: respx.MockRouter) -> None: + """Harnesses may send both max_tokens and max_completion_tokens; honor the smaller.""" + backend = _backend() + ctx = _ctx() + _mock_tokenize(respx_mock, [[1, 2], [1, 2, _EOT, 20], [1, 2, _EOT]]) + injected = [1, 2, 3, 4, 5, _EOT, 20] + completions = respx_mock.post(f"{_BASE_URL}/completions").mock( + return_value=_completion_response(injected, [40], "ok") + ) + + await backend.call(ctx, _chat_request()) + request = ChatRequest.openai_chat({ + "model": _MODEL, + "messages": dict(_chat_request(turns=2).body)["messages"], + "max_tokens": 32000, + "max_completion_tokens": 512, + }) + await backend.call(ctx, request) + assert json.loads(completions.calls[0].request.content)["max_tokens"] == 512 + + +@respx.mock +async def test_max_tokens_clamped_to_model_context(respx_mock: respx.MockRouter) -> None: + """The chat endpoint caps the budget to remaining context; the injected + completions call must do the same or vLLM rejects it with a 400.""" + backend = _backend() + ctx = _ctx() + respx_mock.post("http://vllm.test/tokenize").mock( + side_effect=[ + httpx.Response(200, json={"tokens": [1, 2], "max_model_len": 40}), + httpx.Response(200, json={"tokens": [1, 2, _EOT, 20], "max_model_len": 40}), + httpx.Response(200, json={"tokens": [1, 2, _EOT], "max_model_len": 40}), + ] + ) + injected = [1, 2, 3, 4, 5, _EOT, 20] # 7 tokens; remaining context = 33 + completions = respx_mock.post(f"{_BASE_URL}/completions").mock( + return_value=_completion_response(injected, [40], "ok") + ) + + await backend.call(ctx, _chat_request()) + request = ChatRequest.openai_chat({ + "model": _MODEL, + "messages": dict(_chat_request(turns=2).body)["messages"], + "max_tokens": 32000, + }) + await backend.call(ctx, request) + assert json.loads(completions.calls[0].request.content)["max_tokens"] == 25 # 40 - 7 - margin 8 + + +@respx.mock +async def test_sampling_params_forwarded(respx_mock: respx.MockRouter) -> None: + backend = _backend() + ctx = _ctx() + _mock_tokenize(respx_mock, [[1, 2], [1, 2, _EOT, 20], [1, 2, _EOT]]) + injected = [1, 2, 3, 4, 5, _EOT, 20] + completions = respx_mock.post(f"{_BASE_URL}/completions").mock( + return_value=_completion_response(injected, [40], "ok") + ) + + await backend.call(ctx, _chat_request()) + request = ChatRequest.openai_chat({ + "model": _MODEL, + "messages": dict(_chat_request(turns=2).body)["messages"], + "temperature": 0.7, + "max_completion_tokens": 128, + "stop": ["\n\n"], + }) + await backend.call(ctx, request) + + sent = json.loads(completions.calls[0].request.content) + assert sent["temperature"] == 0.7 + assert sent["max_tokens"] == 128 + assert sent["stop"] == ["\n\n"] + + +@respx.mock +async def test_required_tool_choice_uses_guided_decoding(respx_mock: respx.MockRouter) -> None: + """tool_choice: required rides structured_outputs and parses the JSON array.""" + backend = _backend() + ctx = _ctx() + _mock_tokenize(respx_mock, [[1, 2], [1, 2, _EOT, 20], [1, 2, _EOT]]) + injected = [1, 2, 3, 4, 5, _EOT, 20] + guided_text = json.dumps([{"name": "get_weather", "parameters": {"city": "Paris"}}]) + completions = respx_mock.post(f"{_BASE_URL}/completions").mock( + return_value=_completion_response(injected, [40, 41], guided_text) + ) + + await backend.call(ctx, _chat_request()) + request = ChatRequest.openai_chat({ + "model": _MODEL, + "messages": dict(_chat_request(turns=2).body)["messages"], + "tools": _TOOLS, + "tool_choice": "required", + }) + response = await backend.call(ctx, request) + + sent = json.loads(completions.calls[0].request.content) + assert sent["structured_outputs"]["json"]["minItems"] == 1 + + choice = dict(response.body)["choices"][0] + assert choice["finish_reason"] == "tool_calls" + call = choice["message"]["tool_calls"][0] + assert call["function"]["name"] == "get_weather" + assert json.loads(call["function"]["arguments"]) == {"city": "Paris"} + + +# --------------------------------------------------------------------------- +# Chain routing and fallback +# --------------------------------------------------------------------------- + + +async def test_no_session_delegates_without_state() -> None: + inner = _FakeInner() + backend = _backend(inner) + await backend.call(_ctx(session=None), _chat_request()) + assert inner.calls == 1 + assert backend._sessions == {} + + +@respx.mock +async def test_missing_token_fields_seeds_nothing_and_retries( + respx_mock: respx.MockRouter, +) -> None: + inner = _FakeInner(body={"id": "x", "model": _MODEL, "choices": [{"message": {}}]}) + backend = _backend(inner) + ctx = _ctx() + _mock_tokenize(respx_mock, [[1, 2], [1, 2, 30]]) + + await backend.call(ctx, _chat_request()) + await backend.call(ctx, _chat_request(turns=2)) + assert inner.calls == 2 # both served via the chat path + assert backend._sessions["sess-1"] == [] # nothing seeded, retried each call + + +@respx.mock +async def test_side_call_is_served_without_disturbing_the_chain( + respx_mock: respx.MockRouter, +) -> None: + """A non-extending side-call (opencode's title generator) is served via the + chat path; the main chain survives and the next extending call injects.""" + inner = _FakeInner() + backend = _backend(inner) + ctx = _ctx() + _mock_tokenize( + respx_mock, + [ + [1, 2], # c1 (main) seed prefix + [99, 98], # c2 (title-gen) rendered: matches nothing + [99], # c2 seed-attempt prefix (harvest prefix check fails -> no seed) + [1, 2, _EOT, 20, 21, 30], # c3 (main turn 2) rendered: matches the chain + [1, 2, _EOT, 20, 21], # c3 post-injection prefix + ], + ) + injected = [1, 2, 3, 4, 5, _EOT, 20, 21, 30] + completions = respx_mock.post(f"{_BASE_URL}/completions").mock( + return_value=_completion_response(injected, [40, 41], "back on chain") + ) + + await backend.call(ctx, _chat_request()) # seed main chain + await backend.call(ctx, _chat_request(turns=2)) # title-gen: chat path + await backend.call(ctx, _chat_request(turns=2)) # main turn 2: injected + + assert inner.calls == 2 + assert json.loads(completions.calls[0].request.content)["prompt"] == injected + + +@respx.mock +async def test_side_call_arriving_first_seeds_its_own_chain( + respx_mock: respx.MockRouter, +) -> None: + """Arrival order must not matter: a title-gen call arriving FIRST seeds its + own chain, and the main conversation still gets injected.""" + title_body = dict(_first_turn_body()) + title_body["prompt_token_ids"] = [9, 9, 3] + title_body["choices"] = [dict(_first_turn_body()["choices"][0])] + title_body["choices"][0]["token_ids"] = [7, _EOT] + inner = _FakeInner(body=[title_body, _first_turn_body()]) + backend = _backend(inner) + ctx = _ctx() + _mock_tokenize( + respx_mock, + [ + [9, 9], # c1 (title-gen) seed prefix -> chain B + [1, 2, 30], # c2 (main) rendered: matches nothing + [1, 2], # c2 seed prefix -> chain A + [1, 2, _EOT, 20, 21, 30], # c3 (main turn 2) rendered: matches chain A + [1, 2, _EOT, 20, 21], # c3 post-injection prefix + ], + ) + injected = [1, 2, 3, 4, 5, _EOT, 20, 21, 30] + completions = respx_mock.post(f"{_BASE_URL}/completions").mock( + return_value=_completion_response(injected, [40, 41], "main injects") + ) + + await backend.call(ctx, _chat_request()) # title-gen first + await backend.call(ctx, _chat_request(turns=2)) # main call 1: seeds chain A + await backend.call(ctx, _chat_request(turns=3)) # main turn 2: injected + + assert inner.calls == 2 + assert len(backend._sessions["sess-1"]) == 2 + assert json.loads(completions.calls[0].request.content)["prompt"] == injected + + +@respx.mock +async def test_chat_path_max_tokens_clamped(respx_mock: respx.MockRouter) -> None: + """Newer vLLM rejects over-budget chat calls too; the seed call is clamped.""" + inner = _FakeInner() + backend = _backend(inner) + ctx = _ctx() + respx_mock.post("http://vllm.test/tokenize").mock( + return_value=httpx.Response(200, json={"tokens": [1, 2], "max_model_len": 40}) + ) + + request = ChatRequest.openai_chat({ + "model": _MODEL, + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 32000, + }) + await backend.call(ctx, request) + assert inner.seen_bodies[0]["max_tokens"] == 30 # 40 - len([1, 2]) - margin 8 + + +@respx.mock +async def test_length_stop_reseeds_instead_of_injecting( + respx_mock: respx.MockRouter, +) -> None: + body = _first_turn_body() + body["choices"][0]["finish_reason"] = "length" + inner = _FakeInner(body=body) + backend = _backend(inner) + ctx = _ctx() + _mock_tokenize(respx_mock, [[1, 2], [1, 2, 30], [1, 2]]) + + await backend.call(ctx, _chat_request()) # seeds with no EOT marker + await backend.call(ctx, _chat_request(turns=2)) # matches, but cannot inject + assert inner.calls == 2 + assert len(backend._sessions["sess-1"]) == 1 # reseeded in place, not duplicated + + +@respx.mock +async def test_multi_choice_delegates(respx_mock: respx.MockRouter) -> None: + inner = _FakeInner() + backend = _backend(inner) + ctx = _ctx() + _mock_tokenize(respx_mock, [[1, 2]]) + + await backend.call(ctx, _chat_request()) + request = ChatRequest.openai_chat( + {"model": _MODEL, "messages": [{"role": "user", "content": "hi"}], "n": 2} + ) + await backend.call(ctx, request) + assert inner.calls == 2 + assert len(backend._sessions["sess-1"]) == 1 # chain untouched + + +@respx.mock +async def test_prompt_echo_mismatch_falls_back_keeping_chain( + respx_mock: respx.MockRouter, +) -> None: + """A completions response that does not echo the injected prompt is rejected; + the call is served via the chat path and the chain survives.""" + inner = _FakeInner() + backend = _backend(inner) + ctx = _ctx() + _mock_tokenize(respx_mock, [[1, 2], [1, 2, _EOT, 20]]) + respx_mock.post(f"{_BASE_URL}/completions").mock( + return_value=_completion_response([7, 7, 7], [40], "bad echo") + ) + + await backend.call(ctx, _chat_request()) + await backend.call(ctx, _chat_request(turns=2)) + assert inner.calls == 2 + assert len(backend._sessions["sess-1"]) == 1 + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def test_env_suffix_skips_eot_after_natural_stop() -> None: + assert _slice_env_suffix([_EOT, 20, 21], [4, 5, _EOT], _EOT) == [20, 21] + + +def test_env_suffix_keeps_eot_after_truncation() -> None: + assert _slice_env_suffix([_EOT, 20, 21], [4, 5], _EOT) == [_EOT, 20, 21] + + +def test_parse_generation_without_parsers_passes_through() -> None: + parsed = parse_generation( + "raw text", model=_MODEL, tools=None, tool_parser=None, reasoning_parser=None + ) + assert parsed.content == "raw text" + assert parsed.reasoning_content is None + assert parsed.tool_calls == [] + + +# --------------------------------------------------------------------------- +# Route bundle wiring +# --------------------------------------------------------------------------- + + +def _bundle(route: dict[str, Any]) -> dict[str, Any]: + return {"routes": {"policy-model": route}} + + +def _injection_route(**overrides: Any) -> dict[str, Any]: + route: dict[str, Any] = { + "type": "model", + "target": {"model": _MODEL, "base_url": _BASE_URL, "api_key": "k"}, + "token_capture_engine": "vllm", + "token_injection": True, + "tool_parser": "hermes", + "reasoning_parser": "qwen3", + } + route.update(overrides) + return route + + +def _route_backends(table: Any, model_id: str = "policy-model") -> list[Any]: + return list(table.lookup_switchyard(model_id).iter_components()) + + +def test_bundle_builds_injection_backend() -> None: + from switchyard.cli.route_bundle import build_route_bundle_table + + table = build_route_bundle_table(_bundle(_injection_route()), token_capture_enabled=True) + components = _route_backends(table) + injection = [c for c in components if isinstance(c, TokenInjectionBackend)] + assert len(injection) == 1 + backend = injection[0] + assert backend._model == _MODEL + assert backend._tool_parser == "hermes" + assert backend._reasoning_parser == "qwen3" + + +def test_bundle_injection_requires_capture_engine() -> None: + import pytest + + from switchyard.cli.route_bundle import RouteBundleConfigError, build_route_bundle_table + + route = _injection_route() + del route["token_capture_engine"] + with pytest.raises(RouteBundleConfigError, match="token_capture_engine"): + build_route_bundle_table(_bundle(route), token_capture_enabled=True) + + +def test_bundle_injection_stripped_when_capture_disabled() -> None: + from switchyard.cli.route_bundle import build_route_bundle_table + + table = build_route_bundle_table(_bundle(_injection_route()), token_capture_enabled=False) + components = _route_backends(table) + assert not any(isinstance(c, TokenInjectionBackend) for c in components) + + +@respx.mock +def test_bundle_capture_route_exposes_max_model_len(respx_mock: respx.MockRouter) -> None: + """Token-capture routes surface the engine's context length on /v1/models.""" + from switchyard.cli.route_bundle import build_route_bundle_table + + respx_mock.get(f"{_BASE_URL}/models").mock( + return_value=httpx.Response( + 200, json={"data": [{"id": _MODEL, "max_model_len": 32768}]} + ) + ) + table = build_route_bundle_table(_bundle(_injection_route()), token_capture_enabled=True) + entry = next(e for e in table.registered_model_entries() if e["id"] == "policy-model") + assert entry["max_model_len"] == 32768 + + +@respx.mock +def test_bundle_max_model_len_absent_on_fetch_failure(respx_mock: respx.MockRouter) -> None: + from switchyard.cli.route_bundle import build_route_bundle_table + + respx_mock.get(f"{_BASE_URL}/models").mock(return_value=httpx.Response(503)) + table = build_route_bundle_table(_bundle(_injection_route()), token_capture_enabled=True) + entry = next(e for e in table.registered_model_entries() if e["id"] == "policy-model") + assert "max_model_len" not in entry + + +def test_bundle_without_injection_key_is_unwrapped() -> None: + from switchyard.cli.route_bundle import build_route_bundle_table + + route = _injection_route() + del route["token_injection"] + table = build_route_bundle_table(_bundle(route), token_capture_enabled=True) + components = _route_backends(table) + assert not any(isinstance(c, TokenInjectionBackend) for c in components) + + +def test_tool_choice_schema_shapes() -> None: + assert tool_choice_json_schema(_TOOLS, "auto") is None + assert tool_choice_json_schema(_TOOLS, None) is None + assert tool_choice_json_schema([], "required") is None + + required = tool_choice_json_schema(_TOOLS, "required") + assert required is not None + assert required["items"]["anyOf"][0]["properties"]["name"]["enum"] == ["get_weather"] + + named = tool_choice_json_schema( + _TOOLS, {"type": "function", "function": {"name": "get_weather"}} + ) + assert named == _TOOLS[0]["function"]["parameters"]