From 81ebcac45385a8405729617a78a4294d96955ebd Mon Sep 17 00:00:00 2001 From: Karen Chung Date: Mon, 7 Sep 2026 20:28:33 +0000 Subject: [PATCH 1/3] refactor(agents): session_id reader from AgentContext + forward to vLLM Signed-off-by: Karen Chung --- .../dynamo/common/backend/agent_context.py | 29 +++++++++++++++++++ .../dynamo/trtllm/conversation_affinity.py | 20 +++---------- components/src/dynamo/vllm/handlers.py | 9 ++++++ .../vllm/tests/test_vllm_worker_handler.py | 2 ++ 4 files changed, 44 insertions(+), 16 deletions(-) create mode 100644 components/src/dynamo/common/backend/agent_context.py diff --git a/components/src/dynamo/common/backend/agent_context.py b/components/src/dynamo/common/backend/agent_context.py new file mode 100644 index 000000000000..1c39b00bc17b --- /dev/null +++ b/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 diff --git a/components/src/dynamo/trtllm/conversation_affinity.py b/components/src/dynamo/trtllm/conversation_affinity.py index e4a7c96d1f89..1e2365273127 100644 --- a/components/src/dynamo/trtllm/conversation_affinity.py +++ b/components/src/dynamo/trtllm/conversation_affinity.py @@ -28,7 +28,6 @@ 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. @@ -36,22 +35,11 @@ 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: diff --git a/components/src/dynamo/vllm/handlers.py b/components/src/dynamo/vllm/handlers.py index e77994bf7584..3fbf29cb08ef 100644 --- a/components/src/dynamo/vllm/handlers.py +++ b/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, diff --git a/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py b/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py index ec415f1add45..6ddd3472c40d 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_worker_handler.py +++ b/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: From 8953c80a9615681d6c77e54d96a1a59b27c1b8e8 Mon Sep 17 00:00:00 2001 From: Karen Chung Date: Mon, 14 Sep 2026 15:45:53 -0700 Subject: [PATCH 2/3] Apply suggestion from @coderabbitai[bot] Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Karen Chung --- components/src/dynamo/trtllm/conversation_affinity.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/components/src/dynamo/trtllm/conversation_affinity.py b/components/src/dynamo/trtllm/conversation_affinity.py index 1e2365273127..8acf7384e9ac 100644 --- a/components/src/dynamo/trtllm/conversation_affinity.py +++ b/components/src/dynamo/trtllm/conversation_affinity.py @@ -36,7 +36,7 @@ ConversationParams = None # type: ignore[assignment] from dynamo.common.backend.agent_context import ( - session_id_from_request as session_id_from_request, + session_id_from_request, ) CONVERSATION_PARAMS_AVAILABLE: bool = ConversationParams is not None From d240a2be9e95b4f210babba1734fa78c6bc1fca5 Mon Sep 17 00:00:00 2001 From: Karen Chung Date: Tue, 15 Sep 2026 18:19:49 -0700 Subject: [PATCH 3/3] lint Signed-off-by: Karen Chung --- components/src/dynamo/trtllm/conversation_affinity.py | 4 ---- components/src/dynamo/trtllm/request_handlers/handler_base.py | 2 +- .../src/dynamo/trtllm/tests/test_conversation_affinity.py | 2 +- 3 files changed, 2 insertions(+), 6 deletions(-) diff --git a/components/src/dynamo/trtllm/conversation_affinity.py b/components/src/dynamo/trtllm/conversation_affinity.py index 8acf7384e9ac..736b54200179 100644 --- a/components/src/dynamo/trtllm/conversation_affinity.py +++ b/components/src/dynamo/trtllm/conversation_affinity.py @@ -35,10 +35,6 @@ except ImportError: # pragma: no cover - depends on installed wheel ConversationParams = None # type: ignore[assignment] -from dynamo.common.backend.agent_context import ( - session_id_from_request, -) - CONVERSATION_PARAMS_AVAILABLE: bool = ConversationParams is not None diff --git a/components/src/dynamo/trtllm/request_handlers/handler_base.py b/components/src/dynamo/trtllm/request_handlers/handler_base.py index 0233753d6fef..72b9de4e9d38 100644 --- a/components/src/dynamo/trtllm/request_handlers/handler_base.py +++ b/components/src/dynamo/trtllm/request_handlers/handler_base.py @@ -35,6 +35,7 @@ from dynamo._core import Client, Context from dynamo.common.backend import logprobs as _shared_logprobs +from dynamo.common.backend.agent_context import session_id_from_request from dynamo.common.backend.engine import is_generation_stage from dynamo.common.constants import DisaggregationMode as CommonDisaggregationMode from dynamo.common.multimodal.cache_uuid import reject_unsupported_multimodal_uuids @@ -50,7 +51,6 @@ CONVERSATION_PARAMS_AVAILABLE, conversation_params_for, engine_conversation_affinity_enabled, - session_id_from_request, ) from dynamo.trtllm.engine import TensorRTLLMEngine from dynamo.trtllm.logits_processing.adapter import create_trtllm_adapters diff --git a/components/src/dynamo/trtllm/tests/test_conversation_affinity.py b/components/src/dynamo/trtllm/tests/test_conversation_affinity.py index afc53a309ec1..57fd6f287cd5 100644 --- a/components/src/dynamo/trtllm/tests/test_conversation_affinity.py +++ b/components/src/dynamo/trtllm/tests/test_conversation_affinity.py @@ -8,11 +8,11 @@ import pytest +from dynamo.common.backend.agent_context import session_id_from_request from dynamo.trtllm.conversation_affinity import ( CONVERSATION_PARAMS_AVAILABLE, conversation_params_for, engine_conversation_affinity_enabled, - session_id_from_request, ) pytestmark = [