feat(routing): session-routing plugin layer for per-session router affinity

docs/benchmark-datasets.md

@@ -67,7 +67,9 @@ This document describes datasets that AIPerf can use to generate stimulus. Addit
 
 ## Exgentic Agent Trace Replay
 
-The Exgentic loaders stream recorded agent sessions directly from Hugging Face. `exgentic` is pinned to v1 revision `70036b93a04e61b0ea2706a68b962f4f26774587`; `exgentic_v2` is independently pinned to v2 revision `4b8ad4ab198438e5a170f9171c19c6a2cf7c1814`. Each replays successful, positive-token chat call snapshots. Recorded messages, system instructions, tool definitions, output-token limits, request controls, and call start times are preserved. Tools are not executed, and live responses are not added to later requests. Every request carries the source session as `x-dynamo-session-id` for Dynamo agentic tracing while AIPerf retains its own request correlation ID.
+The Exgentic loaders stream recorded agent sessions directly from Hugging Face. `exgentic` is pinned to v1 revision `70036b93a04e61b0ea2706a68b962f4f26774587`; `exgentic_v2` is independently pinned to v2 revision `4b8ad4ab198438e5a170f9171c19c6a2cf7c1814`. Each replays successful, positive-token chat call snapshots. Recorded messages, system instructions, tool definitions, output-token limits, request controls, and call start times are preserved. Tools are not executed, and live responses are not added to later requests. The loaders no longer stamp a session header themselves; AIPerf keeps a stable per-conversation correlation ID, and pairing the run with `--session-routing dynamo_headers` sends `X-Dynamo-Session-ID` for Dynamo session affinity and agentic tracing.
+
+This changes fixed-schedule affinity grouping. `--session-routing dynamo_headers` keys affinity on the live per-conversation correlation ID, and in `--fixed-schedule` replay the loader explodes one recorded agent session into many single-turn span-conversations — so each span now gets its own key and routes independently, and Dynamo prefix-affinity no longer groups a recorded session's spans onto one worker. The removed stamp instead put the shared recorded session ID on every span, grouping them (but that recorded ID also collided across dataset-recycled replays, binding unrelated live conversations to a single router entry). The two behaviors serve different fidelity goals; restoring source-session grouping without the recycling collision would need a dataset-provided affinity key (possible future work).
 
 Provide a finite materialization bound through `--num-conversations`, `--num-dataset-entries`, or `--request-count`. `--benchmark-duration` limits request issuance, not dataset setup.
 

docs/cli-options.md

@@ -379,6 +379,15 @@ Content type for request body serialization. By default, requests are sent as 'a
 
 HTTP header name used to carry the per-session affinity identifier. When set, replaces the default `X-Correlation-ID` header with the provided name (e.g., `--session-header X-Session-ID`).
 
+#### `--session-routing` `<str>`
+
+Session-aware routing mode: stamps per-session identity on every request for router affinity. Built-ins: dynamo_headers (X-Dynamo-Session-ID + parent header), dynamo_nvext (nvext.session_control bind/close request-body metadata), smg_routing_key (X-SMG-Routing-Key for the SGLang Model Gateway manual policy), session_id_header (custom additive header). Parameterize with --session-routing-opt.
+<br/>_Choices: [`dynamo_headers`, `dynamo_nvext`, `smg_routing_key`, `session_id_header`]_
+
+#### `--session-routing-opt` `<list>`
+
+Repeatable key=value option for the selected --session-routing mode (e.g. --session-routing-opt timeout_seconds=600), validated against the plugin's Options model.
+
 ### Tokenizer
 
 #### `--tokenizer` `<str>`
@@ -1614,7 +1623,7 @@ Explore AIPerf plugins: aiperf plugins [category] [type]
 #### `--category` `<str>`
 
 Category to explore.
-<br/>_Choices: [`accuracy_benchmark`, `accuracy_grader`, `api_router`, `arrival_pattern`, `communication`, `communication_client`, `console_exporter`, `convergence_criterion`, `custom_dataset_loader`, `data_exporter`, `dataset_backing_store`, `dataset_client_store`, `dataset_composer`, `dataset_sampler`, `endpoint`, `gpu_telemetry_collector`, `gpu_telemetry_processor`, `plot`, `public_dataset_loader`, `ramp`, `record_processor`, `results_processor`, `search_planner`, `search_recipe`, `search_recipe_post_process`, `server_metrics_processor`, `service`, `service_manager`, `timing_strategy`, `transport`, `ui`, `url_selection_strategy`, `zmq_proxy`]_
+<br/>_Choices: [`accuracy_benchmark`, `accuracy_grader`, `api_router`, `arrival_pattern`, `communication`, `communication_client`, `console_exporter`, `convergence_criterion`, `custom_dataset_loader`, `data_exporter`, `dataset_backing_store`, `dataset_client_store`, `dataset_composer`, `dataset_sampler`, `endpoint`, `gpu_telemetry_collector`, `gpu_telemetry_processor`, `plot`, `public_dataset_loader`, `ramp`, `record_processor`, `results_processor`, `search_planner`, `search_recipe`, `search_recipe_post_process`, `server_metrics_processor`, `service`, `service_manager`, `session_routing`, `timing_strategy`, `transport`, `ui`, `url_selection_strategy`, `zmq_proxy`]_
 
 #### `--name` `<str>`
 
@@ -1776,6 +1785,15 @@ Content type for request body serialization. By default, requests are sent as 'a
 
 HTTP header name used to carry the per-session affinity identifier. When set, replaces the default `X-Correlation-ID` header with the provided name (e.g., `--session-header X-Session-ID`).
 
+#### `--session-routing` `<str>`
+
+Session-aware routing mode: stamps per-session identity on every request for router affinity. Built-ins: dynamo_headers (X-Dynamo-Session-ID + parent header), dynamo_nvext (nvext.session_control bind/close request-body metadata), smg_routing_key (X-SMG-Routing-Key for the SGLang Model Gateway manual policy), session_id_header (custom additive header). Parameterize with --session-routing-opt.
+<br/>_Choices: [`dynamo_headers`, `dynamo_nvext`, `smg_routing_key`, `session_id_header`]_
+
+#### `--session-routing-opt` `<list>`
+
+Repeatable key=value option for the selected --session-routing mode (e.g. --session-routing-opt timeout_seconds=600), validated against the plugin's Options model.
+
 ### Tokenizer
 
 #### `--tokenizer` `<str>`

docs/plugins/plugin-system.md

@@ -100,7 +100,7 @@ for entry, cls in plugins.iter_all(PluginType.ENDPOINT):
 
 ## Plugin Categories
 
-AIPerf supports 33 plugin categories organized by function, including `api_router` and `public_dataset_loader`:
+AIPerf supports 34 plugin categories organized by function, including `api_router` and `public_dataset_loader`:
 
 ### Timing Categories
 
@@ -129,6 +129,50 @@ AIPerf supports 33 plugin categories organized by function, including `api_route
 | `endpoint` | `EndpointType` | API endpoint implementations (chat, completions, embeddings, etc.) |
 | `transport` | `TransportType` | Network transport (HTTP via aiohttp) |
 
+### Session Routing Category
+
+| Category | Enum | Description |
+|----------|------|-------------|
+| `session_routing` | `SessionRoutingType` | Stamps per-session identity (headers or body metadata) onto outbound requests so an external router pins every turn of a session to one worker; selected via `--session-routing` |
+
+**Purpose.** A session-routing plugin gives an external router (SGLang Model Gateway, Dynamo, a
+generic session-affinity load balancer) the identity of the session a request belongs to, so all of
+a conversation's turns re-land on the replica holding its KV prefix. Exactly one mode runs per
+invocation; the selected plugin is instantiated once per worker by `InferenceClient` and invoked at
+the request-serialization chokepoint. The base class is `SessionRoutingBase`
+(`src/aiperf/workers/session_routing.py`); passthrough is the default (no headers, body unchanged).
+
+**Built-ins:**
+
+| Name | Class | Description |
+|------|-------|-------------|
+| `dynamo_headers` | `DynamoHeadersRouting` | Dynamo session affinity via `X-Dynamo-Session-ID`, plus `X-Dynamo-Parent-Session-ID` on subagent children. No options. |
+| `dynamo_nvext` | `DynamoNvextRouting` | Dynamo session affinity via `nvext.session_control` request-body metadata (bind on non-final turns, close on the final turn). Only for Dynamo builds that implement `session_control`. Option: `timeout_seconds` (default 300). |
+| `smg_routing_key` | `SmgRoutingKeyRouting` | SGLang Model Gateway `manual`-policy stickiness via `X-SMG-Routing-Key`. No options. |
+| `session_id_header` | `SessionIdHeaderRouting` | Generic additive session-affinity header carrying the session's correlation ID. Option: `header_name` (default `X-Session-ID`). |
+
+**Options (`--session-routing-opt key=value`).** Each plugin exposes an `Options` Pydantic model
+(`extra="forbid"`, so unknown keys are rejected at startup). Repeated `--session-routing-opt`
+pairs populate it; values are coerced to the model's field types and canonicalized at config
+resolution, so downstream code always sees typed values. `--session-routing-opt` without
+`--session-routing` is an error, and parameterless modes (`dynamo_headers`, `smg_routing_key`)
+reject every opt key.
+
+**`mutates_body`.** A plugin sets the `mutates_body` class var to `True` when `transform_body`
+changes the payload (only `dynamo_nvext` does); header-only modes leave the body untouched. It is
+protocol metadata that marks a mode as incompatible with any verbatim-bytes request path.
+`transform_body` must never mutate its input — the request path shares cached `Turn.raw_payload`
+dicts with the dataset — so it returns a new dict.
+
+**`on_session_end` contract.** Fires strictly after the session's last worker-side activity, on
+every terminal path (final turn, cancellation, cancel-before-start). It is post-session cleanup
+only and must be idempotent (default no-op).
+
+**Stateful-plugin rule.** A stateful plugin must key its instance state on `ctx.x_correlation_id`
+only. A session tree deliberately spans workers, so tree-keyed worker state fragments across
+processes. For tree-scoped behavior, use the stateless per-request context facts
+`root_correlation_id` and `is_tree_final` instead of accumulating state.
+
 ### Processing Categories
 
 | Category | Enum | Description |

src/aiperf/common/models/model_endpoint_info.py

@@ -122,6 +122,14 @@ def _redact_headers(self, value: list[tuple[str, str]]) -> list[tuple[str, str]]
         default=EndpointDefaults.USE_SERVER_TOKEN_COUNT,
         description="Use server-reported token counts from API usage fields instead of client-side tokenization.",
     )
+    session_routing: str | None = Field(
+        default=None,
+        description="Selected session-routing plugin name (None = off).",
+    )
+    session_routing_opts: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Raw opts for the routing plugin's Options model.",
+    )
     connection_reuse_strategy: ConnectionReuseStrategy = Field(
         default=EndpointDefaults.CONNECTION_REUSE_STRATEGY,
         description="Transport connection reuse strategy.",
@@ -207,6 +215,10 @@ def from_run(cls, run: BenchmarkRun) -> ModelEndpointInfo:
                 api_key=ep.api_key,
                 use_legacy_max_tokens=ep.use_legacy_max_tokens,
                 use_server_token_count=ep.use_server_token_count,
+                session_routing=(
+                    str(ep.session_routing) if ep.session_routing is not None else None
+                ),
+                session_routing_opts=dict(ep.session_routing_opts),
                 connection_reuse_strategy=ep.connection_reuse,
                 download_video_content=ep.download_video_content,
                 request_content_type=ep.request_content_type,

src/aiperf/common/models/record_models.py

@@ -653,6 +653,25 @@ class RequestInfo(RecordContext):
         description="Whether this is the final turn in the conversation. "
         "Used by per-conversation connection strategy to release the connection lease.",
     )
+    root_correlation_id: str | None = Field(
+        default=None,
+        description="The x_correlation_id of the depth-0 root of this request's session TREE. "
+        "Stable across the whole tree (root + every descendant subagent at any depth); equals "
+        "x_correlation_id for a root session. Sourced from the originating Credit. Per-request "
+        "transport context read by the session-routing transform to stamp tree-scoped identity "
+        "on the outbound request; it stays worker-side and is not carried onto the exported record.",
+    )
+    is_parent_final: bool | None = Field(
+        default=None,
+        description="Parent conversation had already returned its final turn at "
+        "credit-issue time; None for roots or when not determinable. Sourced from "
+        "the originating Credit.",
+    )
+    is_tree_final: bool = Field(
+        default=False,
+        description="Provably the last request this session tree will send "
+        "(conservative False when indeterminate). Sourced from the originating Credit.",
+    )
     url_index: int | None = Field(
         default=None,
         ge=0,

src/aiperf/config/endpoint.py

@@ -30,6 +30,7 @@
 from aiperf.config.loader.parsing import normalize_http_urls
 from aiperf.plugin.enums import (
     EndpointType,
+    SessionRoutingType,
     TransportType,
     URLSelectionStrategy,
 )
@@ -316,6 +317,34 @@ def _redact_headers(self, value: dict[str, str]) -> dict[str, str]:
         ),
     ]
 
+    session_routing: Annotated[
+        SessionRoutingType | None,
+        Field(
+            default=None,
+            description=(
+                "Session-routing transform stamping per-session identity on "
+                "every request so an external router (SGLang Model Gateway, Dynamo, "
+                "generic session-affinity LBs) can pin a session's turns to "
+                "one replica. One mode per run; parameterize with "
+                "--session-routing-opt key=value."
+            ),
+        ),
+    ]
+
+    session_routing_opts: Annotated[
+        dict[str, Any],
+        Field(
+            default_factory=dict,
+            description=(
+                "Mode-specific options for --session-routing, validated "
+                "against the selected plugin's Options model (unknown keys and "
+                "invalid values are rejected at startup). Values are "
+                "canonicalized/coerced to the plugin's Options model types at "
+                "validation, so downstream consumers see typed values."
+            ),
+        ),
+    ]
+
     wait_for_model_timeout: Annotated[
         float,
         Field(
@@ -525,3 +554,29 @@ def _validate_request_content_type(self) -> Self:
                 f"video_generation); endpoint type {self.type} does not."
             )
         return self
+
+    @model_validator(mode="after")
+    def validate_session_routing(self) -> Self:
+        """Fail fast: opts require a mode; opts must satisfy the plugin's Options.
+
+        Canonicalizes the opts to the plugin's Options model types so downstream
+        consumers (including the pickled BenchmarkRun that reaches workers) carry
+        coerced values (e.g. ``{"timeout_seconds": 600}``, not ``"600"``).
+        """
+        if self.session_routing is None:
+            if self.session_routing_opts:
+                raise ValueError(
+                    "--session-routing-opt requires --session-routing to select a mode."
+                )
+            return self
+        from aiperf.plugin import plugins
+        from aiperf.plugin.enums import PluginType
+
+        routing_cls = plugins.get_class(
+            PluginType.SESSION_ROUTING, str(self.session_routing)
+        )
+        options = routing_cls.Options(**self.session_routing_opts)
+        canonical = options.model_dump(mode="json", exclude_unset=True)
+        if canonical != self.session_routing_opts:
+            self.session_routing_opts = canonical
+        return self

src/aiperf/config/flags/_converter_endpoint.py

@@ -38,6 +38,21 @@ def _endpoint_template_from_extra(
     }
 
 
+def _parse_routing_opts(values: list[str]) -> dict[str, str]:
+    opts: dict[str, str] = {}
+    for item in values:
+        key, separator, value = item.partition("=")
+        key, value = key.strip(), value.strip()
+        if not separator or not key or not value:
+            raise ValueError(
+                f"Invalid --session-routing-opt {item!r}; expected non-empty key=value"
+            )
+        if key in opts:
+            raise ValueError(f"Duplicate --session-routing-opt key {key!r}")
+        opts[key] = value
+    return opts
+
+
 def _endpoint_template_fallback(endpoint: dict[str, Any]) -> None:
     from aiperf.plugin.enums import EndpointType
 
@@ -72,6 +87,7 @@ def _endpoint_template_fallback(endpoint: dict[str, Any]) -> None:
     "download_video_content": "download_video_content",
     "request_content_type": "request_content_type",
     "session_header": "session_header",
+    "session_routing": "session_routing",
 }
 
 
@@ -97,6 +113,8 @@ def build_endpoint(cli: CLIConfig) -> dict[str, Any]:
         extra = dict(cli.extra_inputs)
         _endpoint_template_from_extra(endpoint, extra)
         endpoint["extra"] = extra
+    if "session_routing_opt" in cli_set and cli.session_routing_opt:
+        endpoint["session_routing_opts"] = _parse_routing_opts(cli.session_routing_opt)
 
     _endpoint_template_fallback(endpoint)
     return endpoint

src/aiperf/config/flags/_section_fields.py

@@ -24,6 +24,8 @@
         "model_selection_strategy",
         "request_content_type",
         "session_header",
+        "session_routing",
+        "session_routing_opt",
         "streaming",
         "timeout_seconds",
         "transport",

src/aiperf/config/flags/cli_config.py

@@ -81,6 +81,7 @@
     GPUTelemetryCollectorType,
     PublicDatasetType,
     SearchPlannerType,
+    SessionRoutingType,
     TransportType,
     UIType,
     URLSelectionStrategy,
@@ -389,6 +390,42 @@ class CLIConfig(BaseConfig):
         ),
     ] = None
 
+    session_routing: Annotated[
+        SessionRoutingType | None,
+        Field(
+            description=(
+                "Session-aware routing mode: stamps per-session identity on "
+                "every request for router affinity. Built-ins: dynamo_headers "
+                "(X-Dynamo-Session-ID + parent header), dynamo_nvext "
+                "(nvext.session_control bind/close request-body metadata), "
+                "smg_routing_key (X-SMG-Routing-Key for the SGLang Model Gateway "
+                "manual policy), session_id_header (custom additive header). "
+                "Parameterize with --session-routing-opt."
+            ),
+        ),
+        CLIParameter(
+            name=("--session-routing",),
+            group=Groups.ENDPOINT,
+        ),
+    ] = None
+
+    session_routing_opt: Annotated[
+        list[str],
+        Field(
+            default_factory=list,
+            description=(
+                "Repeatable key=value option for the selected --session-routing "
+                "mode (e.g. --session-routing-opt timeout_seconds=600), "
+                "validated against the plugin's Options model."
+            ),
+        ),
+        CLIParameter(
+            name=("--session-routing-opt",),
+            consume_multiple=True,
+            group=Groups.ENDPOINT,
+        ),
+    ]
+
     @property
     def url(self) -> str:
         """Return the first URL for backward compatibility."""