diff --git a/docs/cli-options.md b/docs/cli-options.md
index 719398f2a7..bc115d6d1c 100644
--- a/docs/cli-options.md
+++ b/docs/cli-options.md
@@ -324,6 +324,22 @@ 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`).
+#### `--use-dynamo-conv-aware-routing`, `--use-dynamo-session-control`
+
+Emit Dynamo nvext.session_control in OpenAI-compatible request bodies so Dynamo can bind all turns from the same replayed conversation lineage to the same backend worker. This is only intended for Dynamo frontends that implement session_control.
+
_Flag (no value required)_
+
+#### `--use-legacy-dynamo-session-control`
+
+Emit the legacy Dynamo nvext.session_control lifecycle that released Dynamo (v1.2.x) understands: action 'open' on the first turn, session_id only on intermediate turns, and action 'close' on the final turn. Use this when the target Dynamo predates the 'bind' action (added in v1.3.0-dev); otherwise 'bind' is rejected with an HTTP 400. Requires --use-dynamo-conv-aware-routing, and the Dynamo deployment must expose a worker session_control endpoint for 'open' to take effect.
+
_Flag (no value required)_
+
+#### `--dynamo-session-timeout-seconds` ``
+
+Dynamo nvext.session_control timeout in seconds when --use-dynamo-conv-aware-routing is enabled.
+
_Constraints: ≥ 1_
+
_Default: `300`_
+
### Tokenizer
#### `--tokenizer` ``
@@ -1721,6 +1737,22 @@ 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`).
+#### `--use-dynamo-conv-aware-routing`, `--use-dynamo-session-control`
+
+Emit Dynamo nvext.session_control in OpenAI-compatible request bodies so Dynamo can bind all turns from the same replayed conversation lineage to the same backend worker. This is only intended for Dynamo frontends that implement session_control.
+
_Flag (no value required)_
+
+#### `--use-legacy-dynamo-session-control`
+
+Emit the legacy Dynamo nvext.session_control lifecycle that released Dynamo (v1.2.x) understands: action 'open' on the first turn, session_id only on intermediate turns, and action 'close' on the final turn. Use this when the target Dynamo predates the 'bind' action (added in v1.3.0-dev); otherwise 'bind' is rejected with an HTTP 400. Requires --use-dynamo-conv-aware-routing, and the Dynamo deployment must expose a worker session_control endpoint for 'open' to take effect.
+
_Flag (no value required)_
+
+#### `--dynamo-session-timeout-seconds` ``
+
+Dynamo nvext.session_control timeout in seconds when --use-dynamo-conv-aware-routing is enabled.
+
_Constraints: ≥ 1_
+
_Default: `300`_
+
### Tokenizer
#### `--tokenizer` ``
diff --git a/src/aiperf/common/models/model_endpoint_info.py b/src/aiperf/common/models/model_endpoint_info.py
index e5e98bb62b..6a96d1b064 100644
--- a/src/aiperf/common/models/model_endpoint_info.py
+++ b/src/aiperf/common/models/model_endpoint_info.py
@@ -139,6 +139,20 @@ def _redact_headers(self, value: list[tuple[str, str]]) -> list[tuple[str, str]]
description="HTTP header name to use for the per-session affinity identifier. "
"When set, replaces the default `X-Correlation-ID` header name with this value.",
)
+ use_dynamo_conv_aware_routing: bool = Field(
+ default=EndpointDefaults.USE_DYNAMO_CONV_AWARE_ROUTING,
+ description="Emit Dynamo nvext.session_control for conversation-aware routing.",
+ )
+ use_legacy_dynamo_session_control: bool = Field(
+ default=EndpointDefaults.USE_LEGACY_DYNAMO_SESSION_CONTROL,
+ description="Emit the v1.2.x-compatible open/close session_control lifecycle "
+ "instead of the 'bind' action (which only exists in Dynamo >= v1.3.0-dev).",
+ )
+ dynamo_session_timeout_seconds: int = Field(
+ default=EndpointDefaults.DYNAMO_SESSION_TIMEOUT_SECONDS,
+ ge=1,
+ description="Timeout in seconds for Dynamo nvext.session_control sessions.",
+ )
collect_trace_chunks: bool = Field(
default=False,
description="Collect per-chunk trace data (timestamps and sizes) for HTTP trace export. "
@@ -213,6 +227,15 @@ def from_run(cls, run: BenchmarkRun) -> ModelEndpointInfo:
collect_trace_chunks=False,
template=getattr(ep, "template", None),
session_header=getattr(ep, "session_header", None),
+ use_dynamo_conv_aware_routing=getattr(
+ ep, "use_dynamo_conv_aware_routing", False
+ ),
+ use_legacy_dynamo_session_control=getattr(
+ ep, "use_legacy_dynamo_session_control", False
+ ),
+ dynamo_session_timeout_seconds=getattr(
+ ep, "dynamo_session_timeout_seconds", 300
+ ),
),
transport=ep.transport,
)
diff --git a/src/aiperf/config/endpoint.py b/src/aiperf/config/endpoint.py
index 3af9882992..decd5a684b 100644
--- a/src/aiperf/config/endpoint.py
+++ b/src/aiperf/config/endpoint.py
@@ -57,6 +57,9 @@ class EndpointDefaults:
CONNECTION_REUSE_STRATEGY = ConnectionReuseStrategy.POOLED
DOWNLOAD_VIDEO_CONTENT = False
REQUEST_CONTENT_TYPE = None
+ USE_DYNAMO_CONV_AWARE_ROUTING = False
+ USE_LEGACY_DYNAMO_SESSION_CONTROL = False
+ DYNAMO_SESSION_TIMEOUT_SECONDS = 300
# Readiness probe defaults. Timeout 0 disables the probe (the default);
# any positive value enables it. Interval is only consulted when the
# probe is enabled but is validated positive so mis-configuration
@@ -317,6 +320,48 @@ def _redact_headers(self, value: dict[str, str]) -> dict[str, str]:
),
]
+ use_dynamo_conv_aware_routing: Annotated[
+ bool,
+ Field(
+ default=EndpointDefaults.USE_DYNAMO_CONV_AWARE_ROUTING,
+ description=(
+ "Emit Dynamo nvext.session_control in OpenAI-compatible request "
+ "bodies so Dynamo can bind all turns from the same replayed "
+ "conversation lineage to the same backend worker. This is only "
+ "intended for Dynamo frontends that implement session_control."
+ ),
+ ),
+ ]
+
+ use_legacy_dynamo_session_control: Annotated[
+ bool,
+ Field(
+ default=EndpointDefaults.USE_LEGACY_DYNAMO_SESSION_CONTROL,
+ description=(
+ "Emit the legacy Dynamo nvext.session_control lifecycle that "
+ "released Dynamo (v1.2.x) understands: action 'open' on the first "
+ "turn, session_id only on intermediate turns, and action 'close' "
+ "on the final turn. Use this when the target Dynamo predates the "
+ "'bind' action (added in v1.3.0-dev); otherwise 'bind' is rejected "
+ "with an HTTP 400. Requires use_dynamo_conv_aware_routing, and "
+ "the Dynamo deployment must expose a worker session_control "
+ "endpoint for 'open' to take effect."
+ ),
+ ),
+ ]
+
+ dynamo_session_timeout_seconds: Annotated[
+ int,
+ Field(
+ default=EndpointDefaults.DYNAMO_SESSION_TIMEOUT_SECONDS,
+ ge=1,
+ description=(
+ "Dynamo nvext.session_control timeout in seconds when "
+ "use_dynamo_conv_aware_routing is enabled."
+ ),
+ ),
+ ]
+
wait_for_model_timeout: Annotated[
float,
Field(
@@ -466,6 +511,23 @@ def _validate_template_required(self) -> Self:
raise ValueError("template is required when endpoint type is 'template'")
return self
+ @model_validator(mode="after")
+ def validate_dynamo_session_control_coherent(self) -> Self:
+ """Reject legacy Dynamo session control unless conversation-aware routing
+ is enabled, since the legacy flag only selects the wire contract for the
+ session_control that use_dynamo_conv_aware_routing emits.
+ """
+ if (
+ self.use_legacy_dynamo_session_control
+ and not self.use_dynamo_conv_aware_routing
+ ):
+ raise ValueError(
+ "--use-legacy-dynamo-session-control has no effect unless "
+ "--use-dynamo-conv-aware-routing is enabled. Enable conversation-"
+ "aware routing, or drop the legacy flag."
+ )
+ return self
+
@model_validator(mode="after")
def _validate_wait_for_model_coherent(self) -> Self:
"""Reject configurations where probe sub-options are set to non-default
diff --git a/src/aiperf/config/flags/_converter_endpoint.py b/src/aiperf/config/flags/_converter_endpoint.py
index afe4773c61..8a8cc73bf3 100644
--- a/src/aiperf/config/flags/_converter_endpoint.py
+++ b/src/aiperf/config/flags/_converter_endpoint.py
@@ -72,6 +72,9 @@ 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",
+ "use_dynamo_conv_aware_routing": "use_dynamo_conv_aware_routing",
+ "use_legacy_dynamo_session_control": "use_legacy_dynamo_session_control",
+ "dynamo_session_timeout_seconds": "dynamo_session_timeout_seconds",
}
diff --git a/src/aiperf/config/flags/_section_fields.py b/src/aiperf/config/flags/_section_fields.py
index 15acafabd1..40f0755c11 100644
--- a/src/aiperf/config/flags/_section_fields.py
+++ b/src/aiperf/config/flags/_section_fields.py
@@ -20,6 +20,7 @@
"connection_reuse_strategy",
"custom_endpoint",
"download_video_content",
+ "dynamo_session_timeout_seconds",
"model_names",
"model_selection_strategy",
"request_content_type",
@@ -30,6 +31,8 @@
"endpoint_type",
"url_selection_strategy",
"urls",
+ "use_dynamo_conv_aware_routing",
+ "use_legacy_dynamo_session_control",
"use_legacy_max_tokens",
"use_server_token_count",
"wait_for_model_interval",
diff --git a/src/aiperf/config/flags/cli_config.py b/src/aiperf/config/flags/cli_config.py
index 22fb5157ba..8eaea6b6f7 100644
--- a/src/aiperf/config/flags/cli_config.py
+++ b/src/aiperf/config/flags/cli_config.py
@@ -389,6 +389,60 @@ class CLIConfig(BaseConfig):
),
] = None
+ use_dynamo_conv_aware_routing: Annotated[
+ bool,
+ Field(
+ description=(
+ "Emit Dynamo nvext.session_control in OpenAI-compatible request "
+ "bodies so Dynamo can bind all turns from the same replayed "
+ "conversation lineage to the same backend worker. This is only "
+ "intended for Dynamo frontends that implement session_control."
+ ),
+ ),
+ CLIParameter(
+ name=(
+ "--use-dynamo-conv-aware-routing",
+ "--use-dynamo-session-control",
+ ),
+ group=Groups.ENDPOINT,
+ ),
+ ] = EndpointDefaults.USE_DYNAMO_CONV_AWARE_ROUTING
+
+ use_legacy_dynamo_session_control: Annotated[
+ bool,
+ Field(
+ description=(
+ "Emit the legacy Dynamo nvext.session_control lifecycle that "
+ "released Dynamo (v1.2.x) understands: action 'open' on the first "
+ "turn, session_id only on intermediate turns, and action 'close' "
+ "on the final turn. Use this when the target Dynamo predates the "
+ "'bind' action (added in v1.3.0-dev); otherwise 'bind' is rejected "
+ "with an HTTP 400. Requires --use-dynamo-conv-aware-routing, and "
+ "the Dynamo deployment must expose a worker session_control "
+ "endpoint for 'open' to take effect."
+ ),
+ ),
+ CLIParameter(
+ name=("--use-legacy-dynamo-session-control",),
+ group=Groups.ENDPOINT,
+ ),
+ ] = EndpointDefaults.USE_LEGACY_DYNAMO_SESSION_CONTROL
+
+ dynamo_session_timeout_seconds: Annotated[
+ int,
+ Field(
+ ge=1,
+ description=(
+ "Dynamo nvext.session_control timeout in seconds when "
+ "--use-dynamo-conv-aware-routing is enabled."
+ ),
+ ),
+ CLIParameter(
+ name=("--dynamo-session-timeout-seconds",),
+ group=Groups.ENDPOINT,
+ ),
+ ] = EndpointDefaults.DYNAMO_SESSION_TIMEOUT_SECONDS
+
@property
def url(self) -> str:
"""Return the first URL for backward compatibility."""
diff --git a/src/aiperf/config/schema/aiperf-config.schema.json b/src/aiperf/config/schema/aiperf-config.schema.json
index f5fd78efa0..448f1f18b5 100644
--- a/src/aiperf/config/schema/aiperf-config.schema.json
+++ b/src/aiperf/config/schema/aiperf-config.schema.json
@@ -7214,6 +7214,43 @@
"description": "HTTP header name used to carry the per-session affinity identifier. When set, replaces the default `X-Correlation-ID` header. Useful when the inference server expects a custom session-affinity header (e.g. `--session-header X-Session-ID`).",
"title": "Sessionheader"
},
+ "useDynamoConvAwareRouting": {
+ "default": false,
+ "description": "Emit Dynamo nvext.session_control in OpenAI-compatible request bodies so Dynamo can bind all turns from the same replayed conversation lineage to the same backend worker. This is only intended for Dynamo frontends that implement session_control.",
+ "title": "Usedynamoconvawarerouting",
+ "type": "boolean"
+ },
+ "useLegacyDynamoSessionControl": {
+ "default": false,
+ "description": "Emit the legacy Dynamo nvext.session_control lifecycle that released Dynamo (v1.2.x) understands: action 'open' on the first turn, session_id only on intermediate turns, and action 'close' on the final turn. Use this when the target Dynamo predates the 'bind' action (added in v1.3.0-dev); otherwise 'bind' is rejected with an HTTP 400. Requires use_dynamo_conv_aware_routing, and the Dynamo deployment must expose a worker session_control endpoint for 'open' to take effect.",
+ "title": "Uselegacydynamosessioncontrol",
+ "type": "boolean"
+ },
+ "dynamoSessionTimeoutSeconds": {
+ "default": 300,
+ "description": "Dynamo nvext.session_control timeout in seconds when use_dynamo_conv_aware_routing is enabled.",
+ "title": "Dynamosessiontimeoutseconds",
+ "oneOf": [
+ {
+ "default": 300,
+ "description": "Dynamo nvext.session_control timeout in seconds when use_dynamo_conv_aware_routing is enabled.",
+ "minimum": 1,
+ "title": "Dynamosessiontimeoutseconds",
+ "type": "integer"
+ },
+ {
+ "type": "string",
+ "pattern": ".*\\{\\{.*\\}\\}.*",
+ "description": "Jinja2 template (e.g., '{{ variable }}')."
+ },
+ {
+ "type": "string",
+ "pattern": ".*\\$\\{[A-Za-z_][A-Za-z0-9_]*(?::[^}]*)?\\}.*",
+ "description": "Environment variable (e.g., '${VAR}' or '${VAR:default}')."
+ }
+ ],
+ "x-jinja2-supported": true
+ },
"waitForModelTimeout": {
"default": 0.0,
"description": "Enable a pre-flight readiness probe by setting this to a positive value (seconds). aiperf applies this timeout to each URL/model probe before starting the benchmark, aborting with a non-zero exit if any probe times out. For multiple URLs or models, worst-case wall-clock time can be roughly this timeout multiplied by the number of URL/model probes. The probe strategy is controlled by `--wait-for-model-mode`, which defaults to sending a 1-token inference request. 0 (default) disables the probe. Eliminates the need for external shell-based readiness loops in containers and Kubernetes recipes.",
diff --git a/src/aiperf/exporters/console_api_error_exporter.py b/src/aiperf/exporters/console_api_error_exporter.py
index f789e46719..2a323f14c2 100644
--- a/src/aiperf/exporters/console_api_error_exporter.py
+++ b/src/aiperf/exporters/console_api_error_exporter.py
@@ -77,11 +77,64 @@ def detect(error_summary: list[ErrorDetailsCount]) -> ErrorInsight | None:
return None
+class DynamoSessionControlDetector:
+ @staticmethod
+ def detect(error_summary: list[ErrorDetailsCount]) -> ErrorInsight | None:
+ if not error_summary or not isinstance(error_summary, list):
+ return None
+
+ for item in error_summary:
+ err = getattr(item, "error_details", None)
+ if err is None:
+ continue
+
+ raw_msg = err.message or ""
+ parsed = None
+ with contextlib.suppress(Exception):
+ parsed = orjson.loads(raw_msg)
+
+ backend_msg = None
+ if isinstance(parsed, dict):
+ backend_msg = parsed.get("message")
+
+ error_blob = str(backend_msg or raw_msg).lower()
+
+ # serde rejects the unknown enum value with e.g.
+ # `unknown variant `bind`, expected `open` or `close``.
+ if "unknown variant" in error_blob and "bind" in error_blob:
+ return ErrorInsight(
+ title="Unsupported Dynamo session_control action: bind",
+ problem=(
+ "The Dynamo frontend rejected nvext.session_control with "
+ "action='bind'. This Dynamo build's SessionAction only "
+ "accepts 'open' and 'close' -- the 'bind' action was added "
+ "after the v1.2.x release line (first available in "
+ "v1.3.0-dev / upstream commit d97c889ba)."
+ ),
+ causes=[
+ "--use-dynamo-conv-aware-routing emits action='bind' on every non-final turn.",
+ "The target Dynamo server predates the 'bind' action (e.g. v1.2.1).",
+ ],
+ investigation=[
+ "Check the Dynamo frontend version and its supported SessionAction values.",
+ "Inspect request payloads in profile_export.jsonl -> nvext.session_control.",
+ ],
+ fixes=[
+ "Upgrade Dynamo to a build that supports action='bind' (>= v1.3.0-dev, upstream commit d97c889ba).",
+ "Or run with --use-legacy-dynamo-session-control to emit the v1.2.x-compatible open/close lifecycle (requires the worker to expose a session_control endpoint).",
+ "Or disable --use-dynamo-conv-aware-routing.",
+ ],
+ )
+
+ return None
+
+
class ConsoleApiErrorExporter(AIPerfLoggerMixin):
"""Displays helpful diagnostic panels for known API error patterns."""
DETECTORS: ClassVar[list] = [
MaxCompletionTokensDetector,
+ DynamoSessionControlDetector,
]
def __init__(self, exporter_config: ExporterConfig, **kwargs):
diff --git a/src/aiperf/workers/dynamo_session_control.py b/src/aiperf/workers/dynamo_session_control.py
new file mode 100644
index 0000000000..d3a45429c8
--- /dev/null
+++ b/src/aiperf/workers/dynamo_session_control.py
@@ -0,0 +1,113 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Dynamo conversation-aware routing via ``nvext.session_control``.
+
+Opt-in helpers that shape the ``nvext.session_control`` block Dynamo frontends
+read to pin every turn of a replayed conversation to the same backend worker
+(reusing its prefix KV cache). Kept separate from request building so the
+worker stays unaware of wire-body internals: the policy (which lifecycle action
+to emit) lives in :func:`build_session_control`, and the mechanism (overlaying
+it onto the structured request body) in :func:`merge_session_control`. The
+single caller is the serialization chokepoint
+:meth:`aiperf.workers.inference_client.InferenceClient._send_request_to_transport`,
+right after the endpoint formats the request dict.
+
+The verbatim PAYLOAD_BYTES mmap fast path (raw_payload / inputs_json /
+mooncake-with-payload datasets) is refused against this feature at dataset load
+(:meth:`aiperf.dataset.dataset_manager.DatasetManager._select_mmap_format`), so
+this module only ever sees a freshly built request dict -- mirroring how
+``--cache-bust`` is handled.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def build_session_control(
+ *,
+ session_id: str,
+ is_final_turn: bool,
+ timeout_seconds: int,
+ legacy: bool = False,
+ already_opened: bool = False,
+) -> dict[str, Any]:
+ """Build the ``nvext.session_control`` block for a single request.
+
+ Each conversation instance is its own sticky session keyed by its
+ X-Correlation-ID (``session_id``). Dynamo's router co-locates turns purely
+ by ``session_id``. Two wire contracts are produced depending on ``legacy``.
+
+ Modern (``legacy=False``, the default) -- targets Dynamo builds that
+ implement the ``bind`` action (>= v1.3.0-dev / upstream commit d97c889ba):
+
+ - non-final turn -> ``action: "bind"`` (router-only sticky affinity).
+ Re-bind is idempotent on the router and refreshes the inactivity TTL, so
+ long inter-turn replay delays cannot silently expire affinity
+ mid-conversation. That is why every non-final turn re-binds rather than
+ only the first -- it also removes any need for per-worker bind tracking.
+ - final turn -> ``action: "close"`` to release the router affinity (and any
+ worker-side session) immediately instead of leaking it until the TTL
+ reaper fires -- material on high-cardinality runs (millions of sessions).
+
+ Legacy (``legacy=True``) -- targets released Dynamo (v1.2.x), whose
+ ``SessionAction`` enum only accepts ``open`` and ``close`` (``bind`` does not
+ exist there and is rejected with an HTTP 400 ``unknown variant`` error):
+
+ - first request the worker sends for a session (``already_opened`` False)
+ -> ``action: "open"`` to create the worker session and bind router
+ affinity. ``open`` is NOT idempotent, so the caller tracks which sessions
+ it has opened and passes ``already_opened`` accordingly. The trigger is
+ "first request the worker sends", NOT ``turn_index == 0``: under agentic
+ replay the first request is the WARMUP turn (k_i), and profiling resumes
+ mid-trace at k_i+1, so no profiling request ever carries ``turn_index 0``.
+ ``open`` also requires the Dynamo deployment to expose a worker
+ ``session_control`` endpoint; without one Dynamo silently skips session
+ lifecycle (no 400, but no affinity).
+ - subsequent turns (``already_opened`` True) -> ``session_id`` only, which
+ keeps affinity on the sticky-router path.
+ - final turn -> ``action: "close"`` (same as modern).
+
+ Fork/spawn children use their own ``session_id`` rather than a shared
+ lineage key: ``close`` is keyed on ``session_id``, so a shared key would let
+ a child's ``close`` tear down the parent's still-live affinity. Cross-branch
+ KV reuse survives anyway because the child's first (still-unbound) request is
+ routed by Dynamo's prefix-overlap router onto the worker already holding the
+ shared prefix.
+ """
+ session_control: dict[str, Any] = {"session_id": session_id}
+ if is_final_turn:
+ session_control["action"] = "close"
+ elif legacy:
+ if not already_opened:
+ session_control["action"] = "open"
+ session_control["timeout"] = timeout_seconds
+ else:
+ session_control["action"] = "bind"
+ session_control["timeout"] = timeout_seconds
+ return session_control
+
+
+def merge_session_control(
+ payload: dict[str, Any],
+ session_control: dict[str, Any],
+) -> dict[str, Any]:
+ """Return a copy of ``payload`` with ``session_control`` overlaid under ``nvext``.
+
+ Never mutates ``payload`` nor its nested ``nvext`` / ``session_control``
+ dicts, so it is safe on a cached ``Turn.raw_payload`` or on an ``extra_body``
+ reference shared with the dataset. Existing ``nvext`` keys and any
+ pre-existing ``session_control`` fields are preserved; the computed fields
+ overlay them.
+ """
+ merged = dict(payload)
+ raw_nvext = merged.get("nvext")
+ nvext = dict(raw_nvext) if isinstance(raw_nvext, dict) else {}
+ raw_session_control = nvext.get("session_control")
+ merged_session_control = (
+ dict(raw_session_control) if isinstance(raw_session_control, dict) else {}
+ )
+ merged_session_control.update(session_control)
+ nvext["session_control"] = merged_session_control
+ merged["nvext"] = nvext
+ return merged
diff --git a/src/aiperf/workers/inference_client.py b/src/aiperf/workers/inference_client.py
index b376631526..78031357aa 100644
--- a/src/aiperf/workers/inference_client.py
+++ b/src/aiperf/workers/inference_client.py
@@ -20,6 +20,10 @@
from aiperf.common.redact import redact_headers
from aiperf.plugin import plugins
from aiperf.plugin.enums import PluginType, TransportType
+from aiperf.workers.dynamo_session_control import (
+ build_session_control,
+ merge_session_control,
+)
if TYPE_CHECKING:
from aiperf.transports.base_transports import FirstTokenCallback
@@ -61,6 +65,12 @@ def __init__(self, model_endpoint: ModelEndpointInfo, service_id: str, **kwargs)
self.model_endpoint = model_endpoint
self.service_id = service_id
+ # Legacy Dynamo session_control only: session_ids this worker has already
+ # sent an 'open' for. 'open' is not idempotent and must be sent exactly
+ # once on the first request the worker makes for a session. Entries are
+ # dropped on 'close' to bound the set to in-flight sessions.
+ self._dynamo_opened_sessions: set[str] = set()
+
# Detect and set transport type if not explicitly set
if not model_endpoint.transport:
model_endpoint.transport = TransportType(
@@ -108,6 +118,29 @@ async def _send_request_to_transport(
if raw_payload is not None
else self.endpoint.format_payload(request_info)
)
+ # Dynamo conversation-aware routing (opt-in): overlay
+ # nvext.session_control onto the structured request body. Done here,
+ # after the endpoint built the dict, so it is endpoint-agnostic and
+ # never mutates a cached Turn.
+ endpoint = self.model_endpoint.endpoint
+ if endpoint.use_dynamo_conv_aware_routing and isinstance(payload, dict):
+ session_id = request_info.x_correlation_id
+ legacy = endpoint.use_legacy_dynamo_session_control
+ session_control = build_session_control(
+ session_id=session_id,
+ is_final_turn=request_info.is_final_turn,
+ timeout_seconds=endpoint.dynamo_session_timeout_seconds,
+ legacy=legacy,
+ already_opened=session_id in self._dynamo_opened_sessions,
+ )
+ # Track the open/close lifecycle so legacy 'open' is sent exactly
+ # once per session (modern 'bind' is stateless and ignores this).
+ if legacy:
+ if session_control.get("action") == "open":
+ self._dynamo_opened_sessions.add(session_id)
+ elif request_info.is_final_turn:
+ self._dynamo_opened_sessions.discard(session_id)
+ payload = merge_session_control(payload, session_control)
request_info.payload_bytes = orjson.dumps(payload)
return await self.transport.send_request(
request_info,
@@ -115,6 +148,17 @@ async def _send_request_to_transport(
first_token_callback=first_token_callback,
)
+ def discard_dynamo_session(self, session_id: str) -> None:
+ """Drop a session from the legacy Dynamo 'open' tracking set.
+
+ The inline final-turn discard in ``_send_request_to_transport`` only
+ fires when a final turn is actually dispatched; callers that abandon a
+ session before its final turn (e.g. cancellation) use this to keep the
+ set bounded. Idempotent and a no-op for the modern (non-legacy) path,
+ which never populates this set.
+ """
+ self._dynamo_opened_sessions.discard(session_id)
+
async def _send_request_internal(
self,
request_info: RequestInfo,
diff --git a/src/aiperf/workers/worker.py b/src/aiperf/workers/worker.py
index 79cb7bb40b..347c6d5373 100644
--- a/src/aiperf/workers/worker.py
+++ b/src/aiperf/workers/worker.py
@@ -604,6 +604,11 @@ def _release_and_evict_for_terminal(
Non-FORK and non-parent sessions evict immediately.
"""
+ # Cancelled sessions never dispatch their final turn, so the inline
+ # legacy-'open' discard in the inference client cannot fire; discard
+ # here to keep the tracking set bounded (idempotent, no-op for the
+ # modern path).
+ self.inference_client.discard_dynamo_session(x_correlation_id)
if (
credit.parent_correlation_id is not None
and credit.branch_mode == ConversationBranchMode.FORK
diff --git a/tests/unit/config/test_endpoint.py b/tests/unit/config/test_endpoint.py
index 1a0b17fb6b..bd4a129b0a 100644
--- a/tests/unit/config/test_endpoint.py
+++ b/tests/unit/config/test_endpoint.py
@@ -19,6 +19,7 @@
import warnings
import pytest
+from pydantic import ValidationError
from pytest import param
from aiperf.common.enums import ModelSelectionStrategy
@@ -42,6 +43,48 @@ def test_endpoint_config_timeout_uses_endpoint_default() -> None:
assert endpoint.timeout == EndpointDefaults.TIMEOUT
+def test_endpoint_config_dynamo_session_control_defaults_disabled() -> None:
+ endpoint = EndpointConfig(urls=["http://localhost:8000"])
+ assert endpoint.use_dynamo_conv_aware_routing is False
+ assert endpoint.use_legacy_dynamo_session_control is False
+ assert (
+ endpoint.dynamo_session_timeout_seconds
+ == EndpointDefaults.DYNAMO_SESSION_TIMEOUT_SECONDS
+ )
+
+
+def test_endpoint_config_legacy_dynamo_without_routing_raises_error() -> None:
+ """Legacy session control only selects the wire contract for the routing
+ feature, so enabling it alone is incoherent and must be rejected."""
+ with pytest.raises(
+ ValidationError, match="use-legacy-dynamo-session-control has no effect"
+ ):
+ EndpointConfig(
+ urls=["http://localhost:8000"],
+ use_legacy_dynamo_session_control=True,
+ )
+
+
+@pytest.mark.parametrize(
+ ("use_routing", "use_legacy"),
+ [
+ param(True, False, id="modern-routing-only"),
+ param(True, True, id="legacy-with-routing"),
+ param(False, False, id="both-disabled"),
+ ],
+) # fmt: skip
+def test_endpoint_config_dynamo_session_control_coherent_combos_accepted(
+ use_routing: bool, use_legacy: bool
+) -> None:
+ endpoint = EndpointConfig(
+ urls=["http://localhost:8000"],
+ use_dynamo_conv_aware_routing=use_routing,
+ use_legacy_dynamo_session_control=use_legacy,
+ )
+ assert endpoint.use_dynamo_conv_aware_routing is use_routing
+ assert endpoint.use_legacy_dynamo_session_control is use_legacy
+
+
def _make_model_endpoint(
base_url: str, transport: TransportType | None = None
) -> ModelEndpointInfo:
diff --git a/tests/unit/exporters/test_console_api_error_exporter.py b/tests/unit/exporters/test_console_api_error_exporter.py
index 09e11224e2..d43612b87e 100644
--- a/tests/unit/exporters/test_console_api_error_exporter.py
+++ b/tests/unit/exporters/test_console_api_error_exporter.py
@@ -6,10 +6,12 @@
from unittest.mock import MagicMock
import pytest
+from pytest import param
from rich.console import Console
from aiperf.exporters.console_api_error_exporter import (
ConsoleApiErrorExporter,
+ DynamoSessionControlDetector,
MaxCompletionTokensDetector,
)
from aiperf.exporters.exporter_config import ExporterConfig
@@ -117,3 +119,100 @@ async def test_exporter_skips_when_no_insight(self):
await exporter.export(mock_console)
assert mock_console.print.call_count == 0
+
+
+class TestDynamoSessionControlDetector:
+ """Unit tests for the Dynamo session_control 'bind' rejection detector."""
+
+ @pytest.mark.parametrize(
+ "message",
+ [
+ param(
+ json.dumps(
+ {
+ "message": "Failed to deserialize the JSON body into the target type: "
+ "nvext.session_control.action: unknown variant `bind`, "
+ "expected `open` or `close` at line 1 column 100"
+ }
+ ),
+ id="json_wrapped_backend_message",
+ ),
+ param(
+ "unknown variant `bind`, expected `open` or `close` at line 1 column 100",
+ id="raw_non_json_message",
+ ),
+ param(
+ json.dumps("unknown variant `bind`, expected `open` or `close`"),
+ id="json_non_dict_falls_back_to_raw",
+ ),
+ ],
+ ) # fmt: skip
+ def test_detect_unknown_bind_variant_returns_insight(self, message):
+ """serde-style 'unknown variant `bind`' errors should map to the Dynamo insight."""
+ summary = make_summary(MockErrorDetails(message=message))
+
+ insight = DynamoSessionControlDetector.detect(summary)
+
+ assert insight is not None
+ assert "bind" in insight.title
+ assert "session_control" in insight.problem
+ assert any("--use-dynamo-conv-aware-routing" in c for c in insight.causes)
+ assert any("--use-legacy-dynamo-session-control" in f for f in insight.fixes)
+
+ @pytest.mark.parametrize(
+ "message",
+ [
+ param('{"message": "context_length_exceeded"}', id="unrelated_json_error"),
+ param(
+ "unknown variant `frobnicate`, expected `open` or `close`",
+ id="unknown_variant_without_bind",
+ ),
+ param(
+ '{"message": "session bind failed: worker unavailable"}',
+ id="bind_without_unknown_variant",
+ ),
+ param("", id="empty_message"),
+ ],
+ ) # fmt: skip
+ def test_detect_unrelated_error_returns_none(self, message):
+ summary = make_summary(MockErrorDetails(message=message))
+
+ assert DynamoSessionControlDetector.detect(summary) is None
+
+ def test_detect_no_errors_returns_none(self):
+ assert DynamoSessionControlDetector.detect(None) is None
+ assert DynamoSessionControlDetector.detect([]) is None
+
+ def test_detect_item_without_error_details_skipped_returns_none(self):
+ summary = [MockErrorDetailsCount(None, 1)]
+
+ assert DynamoSessionControlDetector.detect(summary) is None
+
+ def test_detect_none_message_returns_none(self):
+ summary = make_summary(MockErrorDetails(message=None))
+
+ assert DynamoSessionControlDetector.detect(summary) is None
+
+ @pytest.mark.asyncio
+ async def test_exporter_prints_panel_for_bind_rejection(self):
+ """The registered detector should surface a Rich panel via the exporter."""
+ mock_console = MagicMock(spec=Console)
+ err = MockErrorDetails(
+ message=json.dumps(
+ {"message": "unknown variant `bind`, expected `open` or `close`"}
+ )
+ )
+ exporter_config = MagicMock(spec=ExporterConfig)
+ exporter_config.results = MagicMock()
+ exporter_config.results.error_summary = make_summary(err)
+
+ exporter = ConsoleApiErrorExporter(exporter_config)
+
+ await exporter.export(mock_console)
+
+ assert mock_console.print.call_count >= 2
+ _, args, _ = mock_console.print.mock_calls[1]
+ panel = args[0]
+ assert "Unsupported Dynamo session_control action: bind" in str(panel.title)
+ panel_text = str(panel.renderable)
+ assert "--use-legacy-dynamo-session-control" in panel_text
diff --git a/tests/unit/workers/test_dynamo_session_control.py b/tests/unit/workers/test_dynamo_session_control.py
new file mode 100644
index 0000000000..b72d80b4ce
--- /dev/null
+++ b/tests/unit/workers/test_dynamo_session_control.py
@@ -0,0 +1,164 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Unit tests for the Dynamo nvext.session_control helpers."""
+
+from aiperf.workers.dynamo_session_control import (
+ build_session_control,
+ merge_session_control,
+)
+
+
+class TestBuildSessionControl:
+ """Policy: which lifecycle action each turn emits (modern / bind contract)."""
+
+ def test_non_final_turn_binds_with_timeout(self):
+ """Non-final turns re-bind (idempotent, refreshes the router TTL)."""
+ sc = build_session_control(
+ session_id="conv-1", is_final_turn=False, timeout_seconds=300
+ )
+ assert sc == {"session_id": "conv-1", "action": "bind", "timeout": 300}
+
+ def test_modern_ignores_already_opened(self):
+ """Modern mode always re-binds; already_opened only affects legacy."""
+ sc = build_session_control(
+ session_id="conv-1",
+ is_final_turn=False,
+ timeout_seconds=300,
+ already_opened=True,
+ )
+ assert sc == {"session_id": "conv-1", "action": "bind", "timeout": 300}
+
+ def test_final_turn_closes_without_timeout(self):
+ """Final turn closes the session; timeout is irrelevant to close."""
+ sc = build_session_control(
+ session_id="conv-1", is_final_turn=True, timeout_seconds=300
+ )
+ assert sc == {"session_id": "conv-1", "action": "close"}
+
+ def test_session_id_and_timeout_are_passed_through(self):
+ """session_id is the routing key; timeout flows through on bind."""
+ sc = build_session_control(
+ session_id="x-corr-abc",
+ is_final_turn=False,
+ timeout_seconds=42,
+ )
+ assert sc["session_id"] == "x-corr-abc"
+ assert sc["timeout"] == 42
+
+
+class TestBuildSessionControlLegacy:
+ """Legacy (v1.2.x) contract: open on first request, close on last, else bare.
+
+ The 'first request' is signalled by ``already_opened=False`` (the caller
+ tracks it per worker), NOT by turn_index -- agentic replay's first request
+ is the warmup turn k_i, never turn 0.
+ """
+
+ def test_first_request_opens_with_timeout(self):
+ """The first request for a session (not yet opened) emits open."""
+ sc = build_session_control(
+ session_id="conv-1",
+ is_final_turn=False,
+ timeout_seconds=300,
+ legacy=True,
+ already_opened=False,
+ )
+ assert sc == {"session_id": "conv-1", "action": "open", "timeout": 300}
+
+ def test_already_opened_sends_session_id_only(self):
+ """Once opened, subsequent turns carry only session_id (sticky routing)."""
+ sc = build_session_control(
+ session_id="conv-1",
+ is_final_turn=False,
+ timeout_seconds=300,
+ legacy=True,
+ already_opened=True,
+ )
+ assert sc == {"session_id": "conv-1"}
+
+ def test_final_turn_closes(self):
+ """Final turn closes, same as modern mode (even if already opened)."""
+ sc = build_session_control(
+ session_id="conv-1",
+ is_final_turn=True,
+ timeout_seconds=300,
+ legacy=True,
+ already_opened=True,
+ )
+ assert sc == {"session_id": "conv-1", "action": "close"}
+
+ def test_single_turn_closes_rather_than_opens(self):
+ """is_final_turn wins: a single-turn (never-opened) session emits close."""
+ sc = build_session_control(
+ session_id="conv-1",
+ is_final_turn=True,
+ timeout_seconds=300,
+ legacy=True,
+ already_opened=False,
+ )
+ assert sc == {"session_id": "conv-1", "action": "close"}
+
+ def test_never_emits_bind(self):
+ """Legacy mode must never emit the 'bind' action (rejected by v1.2.x)."""
+ actions = {
+ build_session_control(
+ session_id="c",
+ is_final_turn=False,
+ timeout_seconds=300,
+ legacy=True,
+ already_opened=opened,
+ ).get("action")
+ for opened in (False, True)
+ }
+ assert "bind" not in actions
+
+
+class TestMergeSessionControl:
+ """Mechanism: overlay session_control under nvext without mutating input."""
+
+ def test_adds_nvext_session_control_to_bare_payload(self):
+ payload = {"messages": [], "model": "m"}
+ sc = {"session_id": "c1", "action": "bind", "timeout": 300}
+
+ merged = merge_session_control(payload, sc)
+
+ assert merged["nvext"]["session_control"] == sc
+ assert merged["messages"] == []
+ assert merged["model"] == "m"
+
+ def test_preserves_existing_nvext_and_session_control_fields(self):
+ payload = {"nvext": {"trace": "keep", "session_control": {"existing": "keep"}}}
+ sc = {"session_id": "c1", "action": "bind", "timeout": 300}
+
+ merged = merge_session_control(payload, sc)
+
+ assert merged["nvext"] == {
+ "trace": "keep",
+ "session_control": {
+ "existing": "keep",
+ "session_id": "c1",
+ "action": "bind",
+ "timeout": 300,
+ },
+ }
+
+ def test_does_not_mutate_input_payload_or_nested_dicts(self):
+ """Safe to call on a cached Turn.raw_payload / shared extra_body."""
+ nested_sc = {"existing": "keep"}
+ nvext = {"session_control": nested_sc}
+ payload = {"nvext": nvext}
+ sc = {"session_id": "c1", "action": "close"}
+
+ merged = merge_session_control(payload, sc)
+
+ # Inputs are left pristine at every level.
+ assert payload == {"nvext": {"session_control": {"existing": "keep"}}}
+ assert nvext == {"session_control": {"existing": "keep"}}
+ assert nested_sc == {"existing": "keep"}
+ assert merged is not payload
+ # The returned copy carries the overlay.
+ assert merged["nvext"]["session_control"] == {
+ "existing": "keep",
+ "session_id": "c1",
+ "action": "close",
+ }
diff --git a/tests/unit/workers/test_inference_client.py b/tests/unit/workers/test_inference_client.py
index cee9cb695e..87b40a511a 100644
--- a/tests/unit/workers/test_inference_client.py
+++ b/tests/unit/workers/test_inference_client.py
@@ -5,6 +5,7 @@
import warnings
from unittest.mock import AsyncMock, MagicMock, patch
+import orjson
import pytest
from pytest import param
@@ -381,6 +382,211 @@ def test_enrich_request_record_uses_last_turn_model(self, inference_client):
assert result.model_name == "standalone-model"
+ def _make_routing_request_info(
+ self,
+ inference_client,
+ x_correlation_id: str = "conv-1",
+ is_final_turn: bool = False,
+ ) -> RequestInfo:
+ """Build a RequestInfo whose payload goes through endpoint formatting."""
+ return RequestInfo(
+ model_endpoint=inference_client.model_endpoint,
+ turns=[Turn(role="user", texts=[Text(contents=["hello"])])],
+ turn_index=0,
+ credit_num=0,
+ credit_phase=CreditPhase.PROFILING,
+ x_request_id="req-1",
+ x_correlation_id=x_correlation_id,
+ conversation_id="conv-template",
+ is_final_turn=is_final_turn,
+ )
+
+ def _sent_payload(self, inference_client) -> dict:
+ return inference_client.transport.send_request.call_args.kwargs["payload"]
+
+ @pytest.fixture
+ def routing_client(self, inference_client):
+ """Inference client with Dynamo conversation-aware routing enabled."""
+ inference_client.model_endpoint.endpoint.use_dynamo_conv_aware_routing = True
+ inference_client.endpoint.format_payload.return_value = {
+ "model": "test-model",
+ "messages": [{"role": "user", "content": "hello"}],
+ }
+ inference_client.transport.send_request = AsyncMock(
+ return_value=RequestRecord()
+ )
+ return inference_client
+
+ @pytest.mark.asyncio
+ async def test_send_request_to_transport_modern_non_final_turn_binds(
+ self, routing_client
+ ):
+ """Modern mode emits action=bind with timeout on non-final turns."""
+ routing_client.model_endpoint.endpoint.dynamo_session_timeout_seconds = 77
+ request_info = self._make_routing_request_info(
+ routing_client, x_correlation_id="conv-bind", is_final_turn=False
+ )
+
+ await routing_client._send_request_to_transport(request_info)
+
+ payload = self._sent_payload(routing_client)
+ assert payload["nvext"]["session_control"] == {
+ "session_id": "conv-bind",
+ "action": "bind",
+ "timeout": 77,
+ }
+ # Modern 'bind' is stateless: no open-tracking entries are created.
+ assert routing_client._dynamo_opened_sessions == set()
+ # The serialized wire bytes must carry the same overlaid payload.
+ assert orjson.loads(request_info.payload_bytes) == payload
+
+ @pytest.mark.asyncio
+ async def test_send_request_to_transport_modern_final_turn_closes(
+ self, routing_client
+ ):
+ """Modern mode emits action=close (no timeout) on the final turn."""
+ request_info = self._make_routing_request_info(
+ routing_client, x_correlation_id="conv-close", is_final_turn=True
+ )
+
+ await routing_client._send_request_to_transport(request_info)
+
+ payload = self._sent_payload(routing_client)
+ assert payload["nvext"]["session_control"] == {
+ "session_id": "conv-close",
+ "action": "close",
+ }
+ assert routing_client._dynamo_opened_sessions == set()
+
+ @pytest.mark.asyncio
+ async def test_send_request_to_transport_legacy_lifecycle_open_then_bare_then_close(
+ self, routing_client
+ ):
+ """Legacy mode sends open exactly once, then bare session_id, then close."""
+ endpoint = routing_client.model_endpoint.endpoint
+ endpoint.use_legacy_dynamo_session_control = True
+ endpoint.dynamo_session_timeout_seconds = 300
+
+ first = self._make_routing_request_info(
+ routing_client, x_correlation_id="conv-legacy", is_final_turn=False
+ )
+ await routing_client._send_request_to_transport(first)
+ assert self._sent_payload(routing_client)["nvext"]["session_control"] == {
+ "session_id": "conv-legacy",
+ "action": "open",
+ "timeout": 300,
+ }
+ assert routing_client._dynamo_opened_sessions == {"conv-legacy"}
+
+ middle = self._make_routing_request_info(
+ routing_client, x_correlation_id="conv-legacy", is_final_turn=False
+ )
+ await routing_client._send_request_to_transport(middle)
+ assert self._sent_payload(routing_client)["nvext"]["session_control"] == {
+ "session_id": "conv-legacy",
+ }
+ assert routing_client._dynamo_opened_sessions == {"conv-legacy"}
+
+ final = self._make_routing_request_info(
+ routing_client, x_correlation_id="conv-legacy", is_final_turn=True
+ )
+ await routing_client._send_request_to_transport(final)
+ assert self._sent_payload(routing_client)["nvext"]["session_control"] == {
+ "session_id": "conv-legacy",
+ "action": "close",
+ }
+ assert routing_client._dynamo_opened_sessions == set()
+
+ @pytest.mark.asyncio
+ async def test_send_request_to_transport_legacy_tracks_sessions_independently(
+ self, routing_client
+ ):
+ """Each session_id gets its own legacy open, even when interleaved."""
+ endpoint = routing_client.model_endpoint.endpoint
+ endpoint.use_legacy_dynamo_session_control = True
+
+ for session_id in ("conv-a", "conv-b"):
+ request_info = self._make_routing_request_info(
+ routing_client, x_correlation_id=session_id, is_final_turn=False
+ )
+ await routing_client._send_request_to_transport(request_info)
+ assert (
+ self._sent_payload(routing_client)["nvext"]["session_control"]["action"]
+ == "open"
+ )
+
+ assert routing_client._dynamo_opened_sessions == {"conv-a", "conv-b"}
+
+ @pytest.mark.asyncio
+ async def test_send_request_to_transport_preserves_existing_nvext_keys(
+ self, routing_client
+ ):
+ """The overlay must not clobber unrelated nvext fields from the endpoint."""
+ formatted = {
+ "model": "test-model",
+ "nvext": {"ignore_eos": True},
+ }
+ routing_client.endpoint.format_payload.return_value = formatted
+ request_info = self._make_routing_request_info(
+ routing_client, x_correlation_id="conv-nvext", is_final_turn=False
+ )
+
+ await routing_client._send_request_to_transport(request_info)
+
+ payload = self._sent_payload(routing_client)
+ assert payload["nvext"]["ignore_eos"] is True
+ assert payload["nvext"]["session_control"]["session_id"] == "conv-nvext"
+ # merge_session_control copies: the endpoint's dict is never mutated.
+ assert formatted["nvext"] == {"ignore_eos": True}
+
+ @pytest.mark.asyncio
+ async def test_send_request_to_transport_routing_disabled_leaves_payload_untouched(
+ self, inference_client
+ ):
+ """Without the opt-in flag, no nvext.session_control is injected."""
+ inference_client.endpoint.format_payload.return_value = {"model": "test-model"}
+ inference_client.transport.send_request = AsyncMock(
+ return_value=RequestRecord()
+ )
+ request_info = self._make_routing_request_info(
+ inference_client, is_final_turn=True
+ )
+
+ await inference_client._send_request_to_transport(request_info)
+
+ payload = self._sent_payload(inference_client)
+ assert payload == {"model": "test-model"}
+ assert "nvext" not in payload
+
+ @pytest.mark.asyncio
+ async def test_send_request_to_transport_non_dict_payload_skips_overlay(
+ self, routing_client
+ ):
+ """Non-dict payloads (defensive isinstance guard) are passed through as-is."""
+ routing_client.endpoint.format_payload.return_value = "raw-string-body"
+ request_info = self._make_routing_request_info(
+ routing_client, is_final_turn=False
+ )
+
+ await routing_client._send_request_to_transport(request_info)
+
+ assert self._sent_payload(routing_client) == "raw-string-body"
+ assert routing_client._dynamo_opened_sessions == set()
+
+ def test_discard_dynamo_session_removes_tracked_session_and_is_idempotent(
+ self, inference_client
+ ):
+ """Abandoned sessions can be dropped from the legacy open-tracking set."""
+ inference_client._dynamo_opened_sessions.add("conv-x")
+
+ inference_client.discard_dynamo_session("conv-x")
+ assert inference_client._dynamo_opened_sessions == set()
+
+ # Idempotent: discarding an unknown/already-removed session is a no-op.
+ inference_client.discard_dynamo_session("conv-x")
+ inference_client.discard_dynamo_session("never-opened")
+ assert inference_client._dynamo_opened_sessions == set()
+
@pytest.mark.parametrize(
"base_url",
[
diff --git a/tests/unit/workers/test_worker.py b/tests/unit/workers/test_worker.py
index 244443d2f5..144b6ee96a 100644
--- a/tests/unit/workers/test_worker.py
+++ b/tests/unit/workers/test_worker.py
@@ -288,3 +288,26 @@ async def test_falls_back_to_dataset_manager_when_no_client_and_not_stopping(
assert result == expected_conversation
mock_fallback.assert_called_once_with("test-conv-123", sample_credit_context)
+
+
+# --- Terminal Eviction / Dynamo Session Tracking Tests ---
+
+
+@pytest.mark.asyncio
+class TestReleaseAndEvictForTerminal:
+ """Test suite for Worker's _release_and_evict_for_terminal method."""
+
+ async def test_release_and_evict_for_terminal_cancelled_session_discards_dynamo_open_tracking(
+ self, mock_worker, sample_credit_context
+ ):
+ """A session abandoned before its final turn (e.g. cancellation) must be
+ dropped from the legacy Dynamo 'open' tracking set on terminal eviction,
+ since the inline final-turn discard never fires for it."""
+ credit = sample_credit_context.credit
+ mock_worker.inference_client._dynamo_opened_sessions.add(
+ credit.x_correlation_id
+ )
+
+ mock_worker._release_and_evict_for_terminal(credit, credit.x_correlation_id)
+
+ assert mock_worker.inference_client._dynamo_opened_sessions == set()