nvidia-nat 1.4.0a20251120__py3-none-any.whl → 1.4.0a20260113__py3-none-any.whl

nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -39,6 +39,7 @@ from pydantic import BaseModel
 from pydantic import Field
 from starlette.websockets import WebSocket
 
+from nat.builder.context import Context
 from nat.builder.eval_builder import WorkflowEvalBuilder
 from nat.builder.evaluator import EvaluatorInfo
 from nat.builder.function import Function
@@ -54,6 +55,7 @@ from nat.eval.config import EvaluationRunOutput
 from nat.eval.evaluate import EvaluationRun
 from nat.eval.evaluate import EvaluationRunConfig
 from nat.eval.evaluator.evaluator_model import EvalInput
+from nat.front_ends.fastapi.async_job import run_generation
 from nat.front_ends.fastapi.auth_flow_handlers.http_flow_handler import HTTPAuthenticationFlowHandler
 from nat.front_ends.fastapi.auth_flow_handlers.websocket_flow_handler import FlowState
 from nat.front_ends.fastapi.auth_flow_handlers.websocket_flow_handler import WebSocketAuthenticationFlowHandler
@@ -74,6 +76,7 @@ from nat.front_ends.fastapi.utils import get_config_file_path
 from nat.object_store.models import ObjectStoreItem
 from nat.runtime.loader import load_workflow
 from nat.runtime.session import SessionManager
+from nat.utils.log_utils import setup_logging
 
 logger = logging.getLogger(__name__)
 
@@ -105,6 +108,9 @@ class FastApiFrontEndPluginWorkerBase(ABC):
         self._scheduler_address = os.environ.get("NAT_DASK_SCHEDULER_ADDRESS")
         self._db_url = os.environ.get("NAT_JOB_STORE_DB_URL")
         self._config_file_path = get_config_file_path()
+        self._use_dask_threads = os.environ.get("NAT_USE_DASK_THREADS", "0") == "1"
+        self._log_level = int(os.environ.get("NAT_FASTAPI_LOG_LEVEL", logging.INFO))
+        setup_logging(self._log_level)
 
         if self._scheduler_address is not None:
             if not _DASK_AVAILABLE:
@@ -232,6 +238,9 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
         self._outstanding_flows: dict[str, FlowState] = {}
         self._outstanding_flows_lock = asyncio.Lock()
 
+        # Track session managers for each route
+        self._session_managers: list[SessionManager] = []
+
         # Evaluator storage for single-item evaluation
         self._evaluators: dict[str, EvaluatorInfo] = {}
         self._eval_builder: WorkflowEvalBuilder | None = None
@@ -268,6 +277,27 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
             # Don't fail startup, just log the error
             self._evaluators = {}
 
+    async def _create_session_manager(self,
+                                      builder: WorkflowBuilder,
+                                      entry_function: str | None = None) -> SessionManager:
+        """Create and register a SessionManager."""
+
+        sm = await SessionManager.create(config=self._config, shared_builder=builder, entry_function=entry_function)
+        self._session_managers.append(sm)
+
+        return sm
+
+    async def cleanup_session_managers(self):
+        """Clean up all SessionManager resources on shutdown."""
+        for sm in self._session_managers:
+            try:
+                await sm.shutdown()
+            except Exception as e:
+                logger.error(f"Error cleaning up SessionManager: {e}")
+
+        self._session_managers.clear()
+        logger.info("All SessionManagers cleaned up")
+
     async def cleanup_evaluators(self):
         """Clean up evaluator resources on shutdown."""
         if self._eval_builder:
@@ -293,6 +323,9 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
         # TODO: we need config control over this as it's not always needed
         await self.initialize_evaluators(self._config)
 
+        # Ensure session manager resources are cleaned up when the app shuts down
+        app.add_event_handler("shutdown", self.cleanup_session_managers)
+
         # Ensure evaluator resources are cleaned up when the app shuts down
         app.add_event_handler("shutdown", self.cleanup_evaluators)
 
@@ -300,18 +333,20 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
     async def add_routes(self, app: FastAPI, builder: WorkflowBuilder):
 
-        await self.add_default_route(app, SessionManager(await builder.build()))
-        await self.add_evaluate_route(app, SessionManager(await builder.build()))
-        await self.add_evaluate_item_route(app, SessionManager(await builder.build()))
+        await self.add_default_route(app, await self._create_session_manager(builder))
+        await self.add_evaluate_route(app, await self._create_session_manager(builder))
+        await self.add_evaluate_item_route(app, await self._create_session_manager(builder))
+
         await self.add_static_files_route(app, builder)
         await self.add_authorization_route(app)
         await self.add_mcp_client_tool_list_route(app, builder)
+        await self.add_monitor_route(app)
 
         for ep in self.front_end_config.endpoints:
 
-            entry_workflow = await builder.build(entry_function=ep.function_name)
-
-            await self.add_route(app, endpoint=ep, session_manager=SessionManager(entry_workflow))
+            await self.add_route(app,
+                                 endpoint=ep,
+                                 session_manager=await self._create_session_manager(builder, ep.function_name))
 
     async def add_default_route(self, app: FastAPI, session_manager: SessionManager):
 
@@ -662,11 +697,15 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                         endpoint: FastApiFrontEndConfig.EndpointBase,
                         session_manager: SessionManager):
 
-        workflow = session_manager.workflow
+        GenerateBodyType = session_manager.get_workflow_input_schema()
+        GenerateStreamResponseType = session_manager.get_workflow_streaming_output_schema()
+        GenerateSingleResponseType = session_manager.get_workflow_single_output_schema()
 
-        GenerateBodyType = workflow.input_schema
-        GenerateStreamResponseType = workflow.streaming_output_schema
-        GenerateSingleResponseType = workflow.single_output_schema
+        def add_context_headers_to_response(response: Response) -> None:
+            """Add context-based headers to response if available."""
+            observability_trace_id = Context.get().observability_trace_id
+            if observability_trace_id:
+                response.headers["Observability-Trace-Id"] = observability_trace_id
 
         # Skip async generation for custom routes (those with function_name)
         if self._dask_available and not hasattr(endpoint, 'function_name'):
@@ -715,10 +754,13 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
                 response.headers["Content-Type"] = "application/json"
 
-                async with session_manager.session(http_connection=request,
-                                                   user_authentication_callback=self._http_flow_handler.authenticate):
+                async with session_manager.session(
+                        http_connection=request,
+                        user_authentication_callback=self._http_flow_handler.authenticate) as session:
 
-                    return await generate_single_response(None, session_manager, result_type=result_type)
+                    result = await generate_single_response(None, session, result_type=result_type)
+                    add_context_headers_to_response(response)
+                    return result
 
             return get_single
 
@@ -726,13 +768,14 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             async def get_stream(request: Request):
 
-                async with session_manager.session(http_connection=request,
-                                                   user_authentication_callback=self._http_flow_handler.authenticate):
+                async with session_manager.session(
+                        http_connection=request,
+                        user_authentication_callback=self._http_flow_handler.authenticate) as session:
 
                     return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
                                              content=generate_streaming_response_as_str(
                                                  None,
-                                                 session_manager=session_manager,
+                                                 session=session,
                                                  streaming=streaming,
                                                  step_adaptor=self.get_step_adaptor(),
                                                  result_type=result_type,
@@ -744,14 +787,14 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             async def get_stream(filter_steps: str | None = None):
 
-                return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
-                                         content=generate_streaming_response_full_as_str(
-                                             None,
-                                             session_manager=session_manager,
-                                             streaming=streaming,
-                                             result_type=result_type,
-                                             output_type=output_type,
-                                             filter_steps=filter_steps))
+                async with session_manager.session(http_connection=None) as session:
+                    return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
+                                             content=generate_streaming_response_full_as_str(None,
+                                                                                             session=session,
+                                                                                             streaming=streaming,
+                                                                                             result_type=result_type,
+                                                                                             output_type=output_type,
+                                                                                             filter_steps=filter_steps))
 
             return get_stream
 
@@ -761,10 +804,13 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
                 response.headers["Content-Type"] = "application/json"
 
-                async with session_manager.session(http_connection=request,
-                                                   user_authentication_callback=self._http_flow_handler.authenticate):
+                async with session_manager.session(
+                        http_connection=request,
+                        user_authentication_callback=self._http_flow_handler.authenticate) as session:
 
-                    return await generate_single_response(payload, session_manager, result_type=result_type)
+                    result = await generate_single_response(payload, session, result_type=result_type)
+                    add_context_headers_to_response(response)
+                    return result
 
             return post_single
 
@@ -775,13 +821,14 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             async def post_stream(request: Request, payload: request_type):
 
-                async with session_manager.session(http_connection=request,
-                                                   user_authentication_callback=self._http_flow_handler.authenticate):
+                async with session_manager.session(
+                        http_connection=request,
+                        user_authentication_callback=self._http_flow_handler.authenticate) as session:
 
                     return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
                                              content=generate_streaming_response_as_str(
                                                  payload,
-                                                 session_manager=session_manager,
+                                                 session=session,
                                                  streaming=streaming,
                                                  step_adaptor=self.get_step_adaptor(),
                                                  result_type=result_type,
@@ -799,14 +846,14 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             async def post_stream(payload: request_type, filter_steps: str | None = None):
 
-                return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
-                                         content=generate_streaming_response_full_as_str(
-                                             payload,
-                                             session_manager=session_manager,
-                                             streaming=streaming,
-                                             result_type=result_type,
-                                             output_type=output_type,
-                                             filter_steps=filter_steps))
+                async with session_manager.session(http_connection=None) as session:
+                    return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
+                                             content=generate_streaming_response_full_as_str(payload,
+                                                                                             session=session,
+                                                                                             streaming=streaming,
+                                                                                             result_type=result_type,
+                                                                                             output_type=output_type,
+                                                                                             filter_steps=filter_steps))
 
             return post_stream
 
@@ -822,20 +869,22 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                 response.headers["Content-Type"] = "application/json"
                 stream_requested = getattr(payload, 'stream', False)
 
-                async with session_manager.session(http_connection=request):
+                async with session_manager.session(http_connection=request) as session:
                     if stream_requested:
 
                         # Return streaming response
                         return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
                                                  content=generate_streaming_response_as_str(
                                                      payload,
-                                                     session_manager=session_manager,
+                                                     session=session,
                                                      streaming=True,
                                                      step_adaptor=self.get_step_adaptor(),
                                                      result_type=ChatResponseChunk,
                                                      output_type=ChatResponseChunk))
 
-                    return await generate_single_response(payload, session_manager, result_type=ChatResponse)
+                    result = await generate_single_response(payload, session, result_type=ChatResponse)
+                    add_context_headers_to_response(response)
+                    return result
 
             return post_openai_api_compatible
 
@@ -856,23 +905,6 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                                                  updated_at=job.updated_at,
                                                  expires_at=self._job_store.get_expires_at(job))
 
-        async def run_generation(scheduler_address: str,
-                                 db_url: str,
-                                 config_file_path: str,
-                                 job_id: str,
-                                 payload: typing.Any):
-            """Background task to run the workflow."""
-            job_store = JobStore(scheduler_address=scheduler_address, db_url=db_url)
-            try:
-                async with load_workflow(config_file_path) as local_session_manager:
-                    result = await generate_single_response(
-                        payload, local_session_manager, result_type=local_session_manager.workflow.single_output_schema)
-
-                await job_store.update_status(job_id, JobStatus.SUCCESS, output=result)
-            except Exception as e:
-                logger.exception("Error in async job %s", job_id)
-                await job_store.update_status(job_id, JobStatus.FAILURE, error=str(e))
-
         def post_async_generation(request_type: type):
 
             async def start_async_generation(
@@ -895,6 +927,8 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                         job_fn=run_generation,
                         sync_timeout=request.sync_timeout,
                         job_args=[
+                            not self._use_dask_threads,
+                            self._log_level,
                             self._scheduler_address,
                             self._db_url,
                             self._config_file_path,
@@ -1228,7 +1262,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                     if configured_group.config.type != "mcp_client":
                         continue
 
-                    from nat.plugins.mcp.client_config import MCPClientConfig
+                    from nat.plugins.mcp.client.client_config import MCPClientConfig
 
                     config = configured_group.config
                     assert isinstance(config, MCPClientConfig)
@@ -1377,6 +1411,107 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                 }
             })
 
+    async def add_monitor_route(self, app: FastAPI):
+        """Add the per-user monitoring endpoint to the FastAPI app.
+
+        Security Warning:
+            This endpoint exposes per-user identifiers and usage metrics. It should be
+            protected by deploying behind an internal network, a reverse proxy with
+            authentication, or similar access controls to prevent exposure to untrusted callers.
+        """
+        # Check if monitoring is enabled in config
+        if not self._config.general.enable_per_user_monitoring:
+            logger.debug("Per-user monitoring disabled, skipping /monitor/users endpoint")
+            return
+
+        from nat.runtime.metrics import PerUserMetricsCollector
+        from nat.runtime.metrics import PerUserMonitorResponse
+        from nat.runtime.metrics import PerUserResourceUsage
+
+        async def get_per_user_metrics(user_id: str | None = None) -> PerUserMonitorResponse:
+            """
+            Get resource usage metrics for per-user workflows.
+
+            Args:
+                user_id: Optional user ID to filter metrics for a specific user
+
+            Returns:
+                PerUserMonitorResponse with metrics for all or specified users
+            """
+            # Collect metrics from all session managers that have per-user workflows
+            all_users: list[PerUserResourceUsage] = []
+
+            for session_manager in self._session_managers:
+                if not session_manager.is_workflow_per_user:
+                    continue
+
+                collector = PerUserMetricsCollector(session_manager)
+
+                if user_id is not None:
+                    # Filter for specific user
+                    user_metrics = await collector.collect_user_metrics(user_id)
+                    if user_metrics:
+                        all_users.append(user_metrics)
+                else:
+                    # Get all users
+                    response = await collector.collect_all_metrics()
+                    all_users.extend(response.users)
+
+            from datetime import datetime
+            return PerUserMonitorResponse(
+                timestamp=datetime.now(),
+                total_active_users=len(all_users),
+                users=all_users,
+            )
+
+        # Register the monitoring endpoint
+        app.add_api_route(path="/monitor/users",
+                          endpoint=get_per_user_metrics,
+                          methods=["GET"],
+                          response_model=PerUserMonitorResponse,
+                          description="Get resource usage metrics for per-user workflows",
+                          tags=["Monitoring"],
+                          responses={
+                              200: {
+                                  "description": "Successfully retrieved per-user metrics",
+                                  "content": {
+                                      "application/json": {
+                                          "example": {
+                                              "timestamp":
+                                                  "2025-12-16T10:30:00Z",
+                                              "total_active_users":
+                                                  2,
+                                              "users": [{
+                                                  "user_id": "alice",
+                                                  "session": {
+                                                      "created_at": "2025-12-16T09:00:00Z",
+                                                      "last_activity": "2025-12-16T10:29:55Z",
+                                                      "ref_count": 1,
+                                                      "is_active": True
+                                                  },
+                                                  "requests": {
+                                                      "total_requests": 42,
+                                                      "active_requests": 1,
+                                                      "avg_latency_ms": 1250.5,
+                                                      "error_count": 2
+                                                  },
+                                                  "memory": {
+                                                      "per_user_functions_count": 2,
+                                                      "per_user_function_groups_count": 1,
+                                                      "exit_stack_size": 3
+                                                  }
+                                              }]
+                                          }
+                                      }
+                                  }
+                              },
+                              500: {
+                                  "description": "Internal Server Error"
+                              }
+                          })
+
+        logger.info("Added per-user monitoring endpoint at /monitor/users")
+
     async def _add_flow(self, state: str, flow_state: FlowState):
         async with self._outstanding_flows_lock:
             self._outstanding_flows[state] = flow_state

nat/front_ends/fastapi/html_snippets/__init__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/front_ends/fastapi/html_snippets/auth_code_grant_success.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/front_ends/fastapi/intermediate_steps_subscriber.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -32,6 +32,7 @@ async def pull_intermediate(_q, adapter):
     intermediate_done = asyncio.Event()
     context = Context.get()
     loop = asyncio.get_running_loop()
+    trace_id_emitted = False
 
     async def set_intermediate_done():
         intermediate_done.set()
@@ -43,6 +44,16 @@ async def pull_intermediate(_q, adapter):
         If adapter is None, convert the raw IntermediateStep into the complete
         ResponseIntermediateStep and place it into the queue.
         """
+        nonlocal trace_id_emitted
+
+        # Check if trace ID is now available and emit it once
+        if not trace_id_emitted:
+            observability_trace_id = context.observability_trace_id
+            if observability_trace_id:
+                from nat.data_models.api_server import ResponseObservabilityTrace
+                loop.create_task(_q.put(ResponseObservabilityTrace(observability_trace_id=observability_trace_id)))
+                trace_id_emitted = True
+
         if adapter is None:
             adapted = ResponseIntermediateStep(id=item.UUID,
                                                type=item.event_type,

nat/front_ends/fastapi/job_store.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -208,13 +208,14 @@ class JobStore(DaskClientMixin):
         AsyncSession
             An active SQLAlchemy async session with an open transaction.
         """
-        async with self._session() as session:
-            async with session.begin():
-                yield session
-
-        # Removes the current task key from the session registry, preventing
-        # potential memory leaks
-        await self._session.remove()
+        try:
+            async with self._session() as session:
+                async with session.begin():
+                    yield session
+        finally:
+            # Removes the current task key from the session registry, preventing
+            # potential memory leaks
+            await self._session.remove()
 
     def ensure_job_id(self, job_id: str | None) -> str:
         """
@@ -502,13 +503,14 @@ class JobStore(DaskClientMixin):
 
         return updated_at + timedelta(seconds=job.expiry_seconds)
 
-    async def cleanup_expired_jobs(self):
+    async def cleanup_expired_jobs(self) -> int:
         """
         Cleanup expired jobs, keeping the most recent one.
 
         Updated_at is used instead of created_at to determine the most recent job. This is because jobs may not be
         processed in the order they are created.
         """
+        logger.info("Starting cleanup of expired jobs")
         now = datetime.now(UTC)
 
         stmt = select(JobInfo).where(
@@ -536,9 +538,11 @@ class JobStore(DaskClientMixin):
                         elif os.path.isdir(job.output_path):
                             shutil.rmtree(job.output_path)
 
-            if len(expired_ids) > 0:
+            num_expired = len(expired_ids)
+            if num_expired > 0:
                 successfully_expired = []
                 for job_id in expired_ids:
+                    var = None
                     try:
                         var = Variable(name=job_id, client=client)
                         try:
@@ -549,14 +553,22 @@ class JobStore(DaskClientMixin):
                         except TimeoutError:
                             pass
 
-                        var.delete()
                         successfully_expired.append(job_id)
                     except Exception:
                         logger.exception("Failed to expire %s", job_id)
 
+                    finally:
+                        if var is not None:
+                            try:
+                                var.delete()
+                            except Exception:
+                                logger.exception("Failed to delete variable %s", job_id)
+
                 await session.execute(
                     update(JobInfo).where(JobInfo.job_id.in_(successfully_expired)).values(is_expired=True))
 
+            return num_expired
+
 
 def get_db_engine(db_url: str | None = None, echo: bool = False, use_async: bool = True) -> "Engine | AsyncEngine":
     """

nat/front_ends/fastapi/main.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/front_ends/fastapi/message_handler.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -30,6 +30,7 @@ from nat.data_models.api_server import ChatResponse
 from nat.data_models.api_server import ChatResponseChunk
 from nat.data_models.api_server import Error
 from nat.data_models.api_server import ErrorTypes
+from nat.data_models.api_server import ResponseObservabilityTrace
 from nat.data_models.api_server import ResponsePayloadOutput
 from nat.data_models.api_server import ResponseSerializable
 from nat.data_models.api_server import SystemResponseContent
@@ -38,6 +39,7 @@ from nat.data_models.api_server import UserMessageContentRoleType
 from nat.data_models.api_server import UserMessages
 from nat.data_models.api_server import WebSocketMessageStatus
 from nat.data_models.api_server import WebSocketMessageType
+from nat.data_models.api_server import WebSocketObservabilityTraceMessage
 from nat.data_models.api_server import WebSocketSystemInteractionMessage
 from nat.data_models.api_server import WebSocketSystemIntermediateStepMessage
 from nat.data_models.api_server import WebSocketSystemResponseTokenMessage
@@ -69,14 +71,15 @@ class WebSocketMessageHandler:
         self._conversation_id: str | None = None
         self._workflow_schema_type: str | None = None
         self._user_interaction_response: asyncio.Future[TextContent] | None = None
+        self._pending_observability_trace: ResponseObservabilityTrace | None = None
 
         self._flow_handler: FlowHandlerBase | None = None
 
         self._schema_output_mapping: dict[str, type[BaseModel] | type[None]] = {
-            WorkflowSchemaType.GENERATE: self._session_manager.workflow.single_output_schema,
+            WorkflowSchemaType.GENERATE: self._session_manager.get_workflow_single_output_schema(),
             WorkflowSchemaType.CHAT: ChatResponse,
             WorkflowSchemaType.CHAT_STREAM: ChatResponseChunk,
-            WorkflowSchemaType.GENERATE_STREAM: self._session_manager.workflow.streaming_output_schema,
+            WorkflowSchemaType.GENERATE_STREAM: self._session_manager.get_workflow_streaming_output_schema(),
         }
 
     def set_flow_handler(self, flow_handler: FlowHandlerBase) -> None:
@@ -175,6 +178,7 @@ class WebSocketMessageHandler:
             self._message_parent_id = user_message_as_validated_type.id
             self._workflow_schema_type = user_message_as_validated_type.schema_type
             self._conversation_id = user_message_as_validated_type.conversation_id
+            self._pending_observability_trace = None
 
             message_content: typing.Any = await self._process_websocket_user_message(user_message_as_validated_type)
 
@@ -250,6 +254,13 @@ class WebSocketMessageHandler:
                     content=content,
                     status=status)
 
+            elif issubclass(message_schema, WebSocketObservabilityTraceMessage):
+                message = await self._message_validator.create_observability_trace_message(
+                    message_id=message_id,
+                    parent_id=self._message_parent_id,
+                    conversation_id=self._conversation_id,
+                    content=content)
+
             elif isinstance(content, Error):
                 raise ValidationError(f"Invalid input data creating websocket message. {data_model.model_dump_json()}")
 
@@ -327,12 +338,18 @@ class WebSocketMessageHandler:
                                                      user_authentication_callback=auth_callback) as session:
 
                 async for value in generate_streaming_response(payload,
-                                                               session_manager=session,
+                                                               session=session,
                                                                streaming=True,
                                                                step_adaptor=self._step_adaptor,
                                                                result_type=result_type,
                                                                output_type=output_type):
 
+                    # Store observability trace to send after completion message
+                    if isinstance(value, ResponseObservabilityTrace):
+                        if self._pending_observability_trace is None:
+                            self._pending_observability_trace = value
+                        continue
+
                     if not isinstance(value, ResponseSerializable):
                         value = ResponsePayloadOutput(payload=value)
 
@@ -342,3 +359,9 @@ class WebSocketMessageHandler:
             await self.create_websocket_message(data_model=SystemResponseContent(),
                                                 message_type=WebSocketMessageType.RESPONSE_MESSAGE,
                                                 status=WebSocketMessageStatus.COMPLETE)
+
+            # Send observability trace after completion message
+            if self._pending_observability_trace is not None:
+                await self.create_websocket_message(data_model=self._pending_observability_trace,
+                                                    message_type=WebSocketMessageType.OBSERVABILITY_TRACE_MESSAGE)
+                self._pending_observability_trace = None