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..736b54200179 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. @@ -39,21 +38,6 @@ 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. - - 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 - - def engine_conversation_affinity_enabled(llm: Any) -> bool: """Whether the resolved engine config enables conversation-affinity ADP routing. 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 = [ 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: