refactor(agents): session_id reader from AgentContext + forward to vLLM

components/src/dynamo/common/backend/agent_context.py

@@ -0,0 +1,29 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Helpers for reading agent context from backend requests."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+
+def agent_context_from_request(
+    request: Mapping[str, Any],
+) -> Mapping[str, Any] | None:
+    """Return a request's serialized agent context when it is object-shaped."""
+    agent_context = request.get("agent_context")
+    return agent_context if isinstance(agent_context, Mapping) else None
+
+
+def session_id_from_request(request: Mapping[str, Any]) -> str | None:
+    """Return the request's non-empty ``agent_context.session_id``."""
+    agent_context = agent_context_from_request(request)
+    if agent_context is None:
+        return None
+    session_id = agent_context.get("session_id")
+    if not isinstance(session_id, str):
+        return None
+    session_id = session_id.strip()
+    return session_id or None

components/src/dynamo/trtllm/conversation_affinity.py

@@ -28,30 +28,18 @@
 
 from __future__ import annotations
 
-from collections.abc import Mapping
 from typing import Any, Optional
 
 try:  # Requires a TensorRT-LLM release newer than 1.3.0rc20.
     from tensorrt_llm.llmapi import ConversationParams
 except ImportError:  # pragma: no cover - depends on installed wheel
     ConversationParams = None  # type: ignore[assignment]
 
-CONVERSATION_PARAMS_AVAILABLE: bool = ConversationParams is not None
-
-
-def session_id_from_request(request: Mapping[str, Any]) -> Optional[str]:
-    """Return the stable conversation/session id the frontend forwards as
-    ``agent_context.session_id``, or ``None`` if absent, blank, or malformed.
+from dynamo.common.backend.agent_context import (
+    session_id_from_request as session_id_from_request,
+)
 
-    Uses ``session_id`` (the active reasoning/tool chain), not ``parent_session_id``.
-    """
-    agent_context = request.get("agent_context")
-    if not isinstance(agent_context, dict):
-        return None
-    session_id = agent_context.get("session_id")
-    if isinstance(session_id, str) and session_id.strip():
-        return session_id.strip()
-    return None
+CONVERSATION_PARAMS_AVAILABLE: bool = ConversationParams is not None
 
 
 def engine_conversation_affinity_enabled(llm: Any) -> bool:

components/src/dynamo/vllm/handlers.py

@@ -48,6 +48,7 @@
 
 from dynamo._core import Context
 from dynamo.common.backend import logprobs as _shared_logprobs
+from dynamo.common.backend.agent_context import session_id_from_request
 from dynamo.common.lora.manager import LoRAInfo, get_lora_manager
 from dynamo.common.memory.multimodal_embedding_cache_manager import (
     MultimodalEmbeddingCacheManager,
@@ -3234,6 +3235,7 @@ async def generate_tokens(
         priority=0,
         reasoning_ended=None,
         reasoning_parser_kwargs=None,
+        session_id=None,
     ):
         try:
             # Log LoRA usage for this generation (debug level to avoid log spam)
@@ -3252,6 +3254,7 @@ async def generate_tokens(
                     data_parallel_rank=data_parallel_rank,
                     trace_headers=trace_headers,
                     priority=priority,
+                    session_id=session_id,
                     **_engine_generate_reasoning_kwargs(
                         self.engine_client,
                         reasoning_ended,
@@ -3682,6 +3685,7 @@ async def _generate_token_mode(self, request, context, request_id):
 
         trace_headers = context.trace_headers()
         reasoning_ended, reasoning_parser_kwargs = _request_reasoning_metadata(request)
+        session_id = session_id_from_request(request)
 
         # In disagg decode mode, defer engine_client.abort() until the first
         # token so we don't abort while a NIXL KV transfer is still in flight
@@ -3730,6 +3734,7 @@ async def _generate_token_mode(self, request, context, request_id):
                         priority=priority,
                         reasoning_ended=reasoning_ended,
                         reasoning_parser_kwargs=reasoning_parser_kwargs,
+                        session_id=session_id,
                     ):
                         if abort_guard is not None:
                             abort_guard.signal_first_token()
@@ -3780,6 +3785,7 @@ async def _generate_text_mode(self, request, context, request_id):
         first_token_output_seen = False
 
         trace_headers = context.trace_headers()
+        session_id = session_id_from_request(request)
 
         is_decode_only = self.config.disaggregation_mode == DisaggregationMode.DECODE
         if is_decode_only and BYPASS_REMOTE_PREFILL_ANNOTATION in (
@@ -3813,6 +3819,7 @@ async def _generate_text_mode(self, request, context, request_id):
                     data_parallel_rank=dp_rank,
                     trace_headers=trace_headers,
                     priority=priority,
+                    session_id=session_id,
                 )
 
                 async for res in gen:
@@ -3997,6 +4004,7 @@ async def _generate_token_mode(self, request, context, request_id):
 
         trace_headers = context.trace_headers()
         reasoning_ended, reasoning_parser_kwargs = _request_reasoning_metadata(request)
+        session_id = session_id_from_request(request)
 
         async with self._abort_monitor(context, request_id, is_prefill=True):
             try:
@@ -4010,6 +4018,7 @@ async def _generate_token_mode(self, request, context, request_id):
                         lora_request=admitted_lora_request,
                         trace_headers=trace_headers,
                         priority=priority,
+                        session_id=session_id,
                         **_engine_generate_reasoning_kwargs(
                             self.engine_client,
                             reasoning_ended,

components/src/dynamo/vllm/tests/test_vllm_worker_handler.py

@@ -251,6 +251,7 @@ async def fake_generate(
             data_parallel_rank=None,
             trace_headers=None,
             priority=0,
+            session_id=None,
             reasoning_ended=None,
             reasoning_parser_kwargs=None,
         ):
@@ -298,6 +299,7 @@ async def fake_generate(
             data_parallel_rank=None,
             trace_headers=None,
             priority=0,
+            session_id=None,
         ):
             calls["called"] = True
             if False: