feat: token-level capture for RL training (Switchyard ↔ NeMo Gym)

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

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,14 +215,15 @@ 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).
 
     ``subprocess.run`` inherits stdin/stdout/stderr so piped use works.
     ``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"),
     )

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"),
     )

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: