feat(routing): lineage-scoped affinity for dynamo_nvext

Adds scope={conversation,lineage} to the dynamo_nvext mode
(--session-routing-opt scope=lineage; default conversation, unchanged
behavior). Under lineage, every session in an agent tree binds
nvext.session_control with the tree ROOT's correlation ID instead of its
own, so fork/spawn children co-locate on the worker (and data-parallel
rank) already holding the shared parent prefix -- an explicit client-side
affinity for Dynamo deployments without KV-event prefix indexing, where
an unbound child request is otherwise placed arbitrarily and re-prefills
the shared parent prefix once per child.

Close discipline: a shared affinity key must never be torn down while
sibling sessions may still run (a straggler's next bind would re-place it
arbitrarily and lose co-location), so under lineage the close fires only
on a request the credit issuer stamped is_tree_final -- provably the last
request of the whole tree, available under agentic replay, where the
has_branches gate guarantees a branch-declaring final turn never reads as
tree-final while children are pending. Under indeterminate modes
(reactive dag_jsonl) nothing closes and the session_control TTL reclaims
the key.

Lineage-scope design and motivation from #14
(YAMY1234), rebased onto the session-routing plugin layer: the flag
surface there (--dynamo-session-affinity-scope + EndpointDefaults +
EndpointInfo plumbing) collapses into this one Options field, and the
root-final-turn close is replaced by the is_tree_final discipline.

Verified with back-to-back --export-level raw mock-server runs of
--scenario inferencex-agentx-mvp against the PR #14 implementation:
child binds are wire-identical (keyed on the tree-root ID in both), and
the lineage close fired exactly once -- on the provably-last wire request
of a fully-drained tree.

Co-Authored-By: YAMY1234 <yamy1234@users.noreply.github.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com>

Author: 3 people

docs/plugins/plugin-system.md

@@ -146,7 +146,7 @@ the request-serialization chokepoint. The base class is `SessionRoutingBase`
 | 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). |
+| `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`. Options: `timeout_seconds` (default 300); `scope` (default `conversation`; `lineage` binds every session in an agent tree with the tree root's correlation ID for whole-lineage co-location on deployments without KV-event prefix indexing, closing the shared key only on a request stamped provably-last for the whole tree). |
 | `smg_routing_key` | `SmgRoutingKeyRouting` | SGLang Model Gateway `manual`-policy stickiness via `X-SMG-Routing-Key`. No options. |
 | `session_id_header` | `SessionIdHeaderRouting` | Preset: additive `X-Session-ID` header carrying the session's correlation ID. No options. |
 | `identity_headers` | `IdentityHeadersRouting` | Fully generic tiered identity headers: emits exactly what you configure, nothing by default. Options (each a comma-separable name list): `session` (this session's ID), `parent` (immediate parent's ID; omitted for roots), `root` (session-tree root's ID, for whole-tree affinity). Requires at least one name across tiers; names must be unique (case-insensitive) across all tiers combined, including within a single tier. |

src/aiperf/plugin/plugins.yaml

@@ -449,7 +449,12 @@ session_routing:
       Only for Dynamo builds that implement session_control (current
       upstream main does not; use dynamo_headers there). Mutates request
       bodies, so it is incompatible with the PAYLOAD_BYTES mmap fast path.
-      Option: timeout_seconds (default 300).
+      Options: timeout_seconds (default 300); scope (default
+      conversation) -- scope=lineage binds every session in an agent
+      tree with the tree ROOT's correlation ID so the whole lineage
+      co-locates with the shared parent prefix (for deployments without
+      KV-event prefix indexing), closing the shared key only on a
+      request stamped provably-last for the whole tree.
 
   smg_routing_key:
     class: aiperf.workers.session_routing:SmgRoutingKeyRouting

src/aiperf/workers/session_routing.py

@@ -27,7 +27,7 @@
 
 import re
 from abc import ABC
-from typing import Annotated, Any, ClassVar, Generic, TypeVar
+from typing import Annotated, Any, ClassVar, Generic, Literal, TypeVar
 
 from msgspec import Struct
 from pydantic import BeforeValidator, ConfigDict, Field, model_validator
@@ -141,6 +141,17 @@ class DynamoNvextOptions(AIPerfBaseModel):
         ge=1,
         description="Dynamo session_control inactivity timeout carried on every bind.",
     )
+    scope: Literal["conversation", "lineage"] = Field(
+        default="conversation",
+        description="Affinity-key scope. 'conversation' (default) binds each "
+        "session with its own correlation ID and closes on its final turn. "
+        "'lineage' binds every session in an agent tree with the tree ROOT's "
+        "correlation ID so the whole lineage co-locates on the worker holding "
+        "the shared parent prefix (for Dynamo deployments without KV-event "
+        "prefix indexing); the shared key is closed only on a request stamped "
+        "provably-last for the whole tree (is_tree_final, agentic replay) -- "
+        "otherwise the session_control TTL reclaims it.",
+    )
 
 
 class DynamoNvextRouting(SessionRoutingBase[DynamoNvextOptions]):
@@ -150,6 +161,14 @@ class DynamoNvextRouting(SessionRoutingBase[DynamoNvextOptions]):
     router, refreshes the TTL), 'close' on the final turn. Targets Dynamo
     builds that implement session_control; current upstream Dynamo main does
     not (use dynamo_headers there).
+
+    Under ``scope=lineage`` the affinity key is the tree root's correlation ID
+    and the close discipline changes: a shared key must never be torn down
+    while sibling sessions may still run (a later bind would re-place the
+    straggler arbitrarily and lose co-location), so close fires only on a
+    request the issuer stamped ``is_tree_final`` -- conservative-exact under
+    agentic replay, never under indeterminate modes, where the TTL reaper
+    reclaims the key instead.
     """
 
     mutates_body: ClassVar[bool] = True
@@ -160,18 +179,25 @@ def __init__(self, options: DynamoNvextOptions) -> None:
         # Hot path: one attribute hop per request instead of two through the
         # pydantic options model.
         self._timeout = options.timeout_seconds
+        self._lineage = options.scope == "lineage"
 
     def transform_body(
         self, payload: dict[str, Any], ctx: RoutingContext
     ) -> dict[str, Any]:
-        if ctx.is_final_turn:
+        if self._lineage:
+            session_id = ctx.root_correlation_id or ctx.x_correlation_id
+            is_close = ctx.is_tree_final
+        else:
+            session_id = ctx.x_correlation_id
+            is_close = ctx.is_final_turn
+        if is_close:
             session_control: dict[str, Any] = {
-                "session_id": ctx.x_correlation_id,
+                "session_id": session_id,
                 "action": "close",
             }
         else:
             session_control = {
-                "session_id": ctx.x_correlation_id,
+                "session_id": session_id,
                 "action": "bind",
                 "timeout": self._timeout,
             }

tests/unit/workers/test_session_routing.py

@@ -129,6 +129,52 @@ def test_non_dict_nvext_replaced(self):
         merged = plugin.transform_body({"nvext": "bogus"}, _ctx())
         assert merged["nvext"]["session_control"]["session_id"] == "corr-1"
 
+    def test_lineage_scope_child_binds_with_root_id(self):
+        """Under scope=lineage, every session in the tree binds the ROOT's
+        correlation ID so the whole lineage co-locates."""
+        plugin = DynamoNvextRouting(
+            DynamoNvextOptions(scope="lineage", timeout_seconds=77)
+        )
+        ctx = _ctx(
+            parent_correlation_id="root-1",
+            root_correlation_id="root-1",
+            is_final_turn=True,  # child's own final turn must NOT close
+        )
+        sc = plugin.transform_body({}, ctx)["nvext"]["session_control"]
+        assert sc == {"session_id": "root-1", "action": "bind", "timeout": 77}
+
+    def test_lineage_scope_root_final_turn_binds_not_closes(self):
+        """A shared key must never be torn down while siblings may still run:
+        even the root's own final turn only binds unless the issuer stamped
+        the request provably-last for the whole tree."""
+        plugin = DynamoNvextRouting(DynamoNvextOptions(scope="lineage"))
+        sc = plugin.transform_body({}, _ctx(is_final_turn=True))["nvext"][
+            "session_control"
+        ]
+        assert sc["action"] == "bind"
+        assert sc["session_id"] == "corr-1"
+
+    def test_lineage_scope_closes_only_on_tree_final(self):
+        plugin = DynamoNvextRouting(DynamoNvextOptions(scope="lineage"))
+        ctx = _ctx(
+            root_correlation_id="root-1",
+            is_final_turn=True,
+            is_tree_final=True,
+        )
+        sc = plugin.transform_body({}, ctx)["nvext"]["session_control"]
+        assert sc == {"session_id": "root-1", "action": "close"}
+
+    def test_conversation_scope_default_unchanged(self):
+        """Default scope keeps per-conversation identity and final-turn close."""
+        plugin = DynamoNvextRouting(DynamoNvextOptions())
+        assert plugin.transform_body({}, _ctx(is_final_turn=True))["nvext"][
+            "session_control"
+        ] == {"session_id": "corr-1", "action": "close"}
+
+    def test_invalid_scope_rejected(self):
+        with pytest.raises(ValidationError):
+            DynamoNvextOptions(scope="tree")
+
 
 class TestSmgRoutingKeyRouting:
     def test_emits_routing_key(self):