Persist User Message ID For HTTP Connections (#696)

Fixes a bug where conversation_id was not being properly set for WebSocket connections and adds support for persisting user_message_id from HTTP connections. 
Closes: [Issue 658](#658)
## By Submitting this PR I confirm:
- I am familiar with the [Contributing Guidelines](https://github.com/NVIDIA/NeMo-Agent-Toolkit/blob/develop/docs/source/resources/contributing.md).
- We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  - Any contribution which contains commits that are not Signed-Off will not be accepted.
- When the PR is ready for review, new or existing tests cover these changes.
- When the PR is ready for review, the documentation is up to date with these changes.

Authors:
  - Eric Evans II (https://github.com/ericevans-nv)

Approvers:
  - Will Killian (https://github.com/willkill07)

URL: #696

Author: ericevans-nv

external/nat-ui

@@ -65,6 +65,7 @@ class ContextState(metaclass=Singleton):
 
     def __init__(self):
         self.conversation_id: ContextVar[str | None] = ContextVar("conversation_id", default=None)
+        self.user_message_id: ContextVar[str | None] = ContextVar("user_message_id", default=None)
         self.input_message: ContextVar[typing.Any] = ContextVar("input_message", default=None)
         self.user_manager: ContextVar[typing.Any] = ContextVar("user_manager", default=None)
         self.metadata: ContextVar[RequestAttributes] = ContextVar("request_attributes", default=RequestAttributes())
@@ -165,6 +166,13 @@ def conversation_id(self) -> str | None:
         """
         return self._context_state.conversation_id.get()
 
+    @property
+    def user_message_id(self) -> str | None:
+        """
+        This property retrieves the user message ID which is the unique identifier for the current user message.
+        """
+        return self._context_state.user_message_id.get()
+
     @contextmanager
     def push_active_function(self, function_name: str, input_data: typing.Any | None):
         """

src/nat/builder/context.py

@@ -307,7 +307,7 @@ async def run_evaluation(job_id: str, config_file: str, reps: int, session_manag
         async def start_evaluation(request: EvaluateRequest, background_tasks: BackgroundTasks, http_request: Request):
             """Handle evaluation requests."""
 
-            async with session_manager.session(request=http_request):
+            async with session_manager.session(http_connection=http_request):
 
                 # if job_id is present and already exists return the job info
                 if request.job_id:
@@ -336,7 +336,7 @@ async def get_job_status(job_id: str, http_request: Request) -> EvaluateStatusRe
             """Get the status of an evaluation job."""
             logger.info("Getting status for job %s", job_id)
 
-            async with session_manager.session(request=http_request):
+            async with session_manager.session(http_connection=http_request):
 
                 job = job_store.get_job(job_id)
                 if not job:
@@ -349,7 +349,7 @@ async def get_last_job_status(http_request: Request) -> EvaluateStatusResponse:
             """Get the status of the last created evaluation job."""
             logger.info("Getting last job status")
 
-            async with session_manager.session(request=http_request):
+            async with session_manager.session(http_connection=http_request):
 
                 job = job_store.get_last_job()
                 if not job:
@@ -361,7 +361,7 @@ async def get_last_job_status(http_request: Request) -> EvaluateStatusResponse:
         async def get_jobs(http_request: Request, status: str | None = None) -> list[EvaluateStatusResponse]:
             """Get all jobs, optionally filtered by status."""
 
-            async with session_manager.session(request=http_request):
+            async with session_manager.session(http_connection=http_request):
 
                 if status is None:
                     logger.info("Getting all jobs")
@@ -572,7 +572,7 @@ async def get_single(response: Response, request: Request):
 
                 response.headers["Content-Type"] = "application/json"
 
-                async with session_manager.session(request=request,
+                async with session_manager.session(http_connection=request,
                                                    user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return await generate_single_response(None, session_manager, result_type=result_type)
@@ -583,7 +583,7 @@ def get_streaming_endpoint(streaming: bool, result_type: type | None, output_typ
 
             async def get_stream(request: Request):
 
-                async with session_manager.session(request=request,
+                async with session_manager.session(http_connection=request,
                                                    user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
@@ -618,7 +618,7 @@ async def post_single(response: Response, request: Request, payload: request_typ
 
                 response.headers["Content-Type"] = "application/json"
 
-                async with session_manager.session(request=request,
+                async with session_manager.session(http_connection=request,
                                                    user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return await generate_single_response(payload, session_manager, result_type=result_type)
@@ -632,7 +632,7 @@ def post_streaming_endpoint(request_type: type,
 
             async def post_stream(request: Request, payload: request_type):
 
-                async with session_manager.session(request=request,
+                async with session_manager.session(http_connection=request,
                                                    user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
@@ -677,7 +677,7 @@ async def post_openai_api_compatible(response: Response, request: Request, paylo
                 # Check if streaming is requested
                 stream_requested = getattr(payload, 'stream', False)
 
-                async with session_manager.session(request=request):
+                async with session_manager.session(http_connection=request):
                     if stream_requested:
                         # Return streaming response
                         return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
@@ -688,42 +688,41 @@ async def post_openai_api_compatible(response: Response, request: Request, paylo
                                                      step_adaptor=self.get_step_adaptor(),
                                                      result_type=ChatResponseChunk,
                                                      output_type=ChatResponseChunk))
-                    else:
-                        # Return single response - check if workflow supports non-streaming
-                        try:
+
+                    # Return single response - check if workflow supports non-streaming
+                    try:
+                        response.headers["Content-Type"] = "application/json"
+                        return await generate_single_response(payload, session_manager, result_type=ChatResponse)
+                    except ValueError as e:
+                        if "Cannot get a single output value for streaming workflows" in str(e):
+                            # Workflow only supports streaming, but client requested non-streaming
+                            # Fall back to streaming and collect the result
+                            chunks = []
+                            async for chunk_str in generate_streaming_response_as_str(
+                                    payload,
+                                    session_manager=session_manager,
+                                    streaming=True,
+                                    step_adaptor=self.get_step_adaptor(),
+                                    result_type=ChatResponseChunk,
+                                    output_type=ChatResponseChunk):
+                                if chunk_str.startswith("data: ") and not chunk_str.startswith("data: [DONE]"):
+                                    chunk_data = chunk_str[6:].strip()  # Remove "data: " prefix
+                                    if chunk_data:
+                                        try:
+                                            chunk_json = ChatResponseChunk.model_validate_json(chunk_data)
+                                            if (chunk_json.choices and len(chunk_json.choices) > 0
+                                                    and chunk_json.choices[0].delta
+                                                    and chunk_json.choices[0].delta.content is not None):
+                                                chunks.append(chunk_json.choices[0].delta.content)
+                                        except Exception:
+                                            continue
+
+                            # Create a single response from collected chunks
+                            content = "".join(chunks)
+                            single_response = ChatResponse.from_string(content)
                             response.headers["Content-Type"] = "application/json"
-                            return await generate_single_response(payload, session_manager, result_type=ChatResponse)
-                        except ValueError as e:
-                            if "Cannot get a single output value for streaming workflows" in str(e):
-                                # Workflow only supports streaming, but client requested non-streaming
-                                # Fall back to streaming and collect the result
-                                chunks = []
-                                async for chunk_str in generate_streaming_response_as_str(
-                                        payload,
-                                        session_manager=session_manager,
-                                        streaming=True,
-                                        step_adaptor=self.get_step_adaptor(),
-                                        result_type=ChatResponseChunk,
-                                        output_type=ChatResponseChunk):
-                                    if chunk_str.startswith("data: ") and not chunk_str.startswith("data: [DONE]"):
-                                        chunk_data = chunk_str[6:].strip()  # Remove "data: " prefix
-                                        if chunk_data:
-                                            try:
-                                                chunk_json = ChatResponseChunk.model_validate_json(chunk_data)
-                                                if (chunk_json.choices and len(chunk_json.choices) > 0
-                                                        and chunk_json.choices[0].delta
-                                                        and chunk_json.choices[0].delta.content is not None):
-                                                    chunks.append(chunk_json.choices[0].delta.content)
-                                            except Exception:
-                                                continue
-
-                                # Create a single response from collected chunks
-                                content = "".join(chunks)
-                                single_response = ChatResponse.from_string(content)
-                                response.headers["Content-Type"] = "application/json"
-                                return single_response
-                            else:
-                                raise
+                            return single_response
+                        raise
 
             return post_openai_api_compatible
 
@@ -758,7 +757,7 @@ async def start_async_generation(
                     http_request: Request) -> AsyncGenerateResponse | AsyncGenerationStatusResponse:
                 """Handle async generation requests."""
 
-                async with session_manager.session(request=http_request):
+                async with session_manager.session(http_connection=http_request):
 
                     # if job_id is present and already exists return the job info
                     if request.job_id:
@@ -804,7 +803,7 @@ async def get_async_job_status(job_id: str, http_request: Request) -> AsyncGener
             """Get the status of an async job."""
             logger.info("Getting status for job %s", job_id)
 
-            async with session_manager.session(request=http_request):
+            async with session_manager.session(http_connection=http_request):
 
                 job = job_store.get_job(job_id)
                 if not job:

src/nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py

@@ -166,8 +166,9 @@ def _done_callback(task: asyncio.Task):  # pylint: disable=unused-argument
                     self._running_workflow_task = None
 
                 self._running_workflow_task = asyncio.create_task(
-                    self._run_workflow(content.text,
-                                       self._conversation_id,
+                    self._run_workflow(payload=content.text,
+                                       user_message_id=self._message_parent_id,
+                                       conversation_id=self._conversation_id,
                                        result_type=self._schema_output_mapping[self._workflow_schema_type],
                                        output_type=self._schema_output_mapping[
                                            self._workflow_schema_type])).add_done_callback(_done_callback)
@@ -290,14 +291,16 @@ async def human_interaction_callback(self, prompt: InteractionPrompt) -> HumanRe
 
     async def _run_workflow(self,
                             payload: typing.Any,
+                            user_message_id: str | None = None,
                             conversation_id: str | None = None,
                             result_type: type | None = None,
                             output_type: type | None = None) -> None:
 
         try:
             async with self._session_manager.session(
+                    user_message_id=user_message_id,
                     conversation_id=conversation_id,
-                    request=self._socket,
+                    http_connection=self._socket,
                     user_input_callback=self.human_interaction_callback,
                     user_authentication_callback=(self._flow_handler.authenticate
                                                   if self._flow_handler else None)) as session:

src/nat/front_ends/fastapi/message_handler.py

@@ -21,7 +21,9 @@
 from contextlib import asynccontextmanager
 from contextlib import nullcontext
 
+from fastapi import WebSocket
 from starlette.requests import HTTPConnection
+from starlette.requests import Request
 
 from nat.builder.context import Context
 from nat.builder.context import ContextState
@@ -89,7 +91,8 @@ def context(self) -> Context:
     @asynccontextmanager
     async def session(self,
                       user_manager=None,
-                      request: HTTPConnection | None = None,
+                      http_connection: HTTPConnection | None = None,
+                      user_message_id: str | None = None,
                       conversation_id: str | None = None,
                       user_input_callback: Callable[[InteractionPrompt], Awaitable[HumanResponse]] = None,
                       user_authentication_callback: Callable[[AuthProviderBaseConfig, AuthFlowType],
@@ -107,10 +110,11 @@ async def session(self,
         if user_authentication_callback is not None:
             token_user_authentication = self._context_state.user_auth_callback.set(user_authentication_callback)
 
-        if conversation_id is not None and request is None:
-            self._context_state.conversation_id.set(conversation_id)
+        if isinstance(http_connection, WebSocket):
+            self.set_metadata_from_websocket(user_message_id, conversation_id)
 
-        self.set_metadata_from_http_request(request)
+        if isinstance(http_connection, Request):
+            self.set_metadata_from_http_request(http_connection)
 
         try:
             yield self
@@ -135,14 +139,11 @@ async def run(self, message):
             async with self._workflow.run(message) as runner:
                 yield runner
 
-    def set_metadata_from_http_request(self, request: HTTPConnection | None) -> None:
+    def set_metadata_from_http_request(self, request: Request) -> None:
         """
         Extracts and sets user metadata request attributes from a HTTP request.
         If request is None, no attributes are set.
         """
-        if request is None:
-            return
-
         self._context.metadata._request.method = getattr(request, "method", None)
         self._context.metadata._request.url_path = request.url.path
         self._context.metadata._request.url_port = request.url.port
@@ -157,6 +158,20 @@ def set_metadata_from_http_request(self, request: HTTPConnection | None) -> None
         if request.headers.get("conversation-id"):
             self._context_state.conversation_id.set(request.headers["conversation-id"])
 
+        if request.headers.get("user-message-id"):
+            self._context_state.user_message_id.set(request.headers["user-message-id"])
+
+    def set_metadata_from_websocket(self, user_message_id: str | None, conversation_id: str | None) -> None:
+        """
+        Extracts and sets user metadata for Websocket connections.
+        """
+
+        if conversation_id is not None:
+            self._context_state.conversation_id.set(conversation_id)
+
+        if user_message_id is not None:
+            self._context_state.user_message_id.set(user_message_id)
+
 
 # Compatibility aliases with previous releases
 AIQSessionManager = SessionManager