Websocket context solution - draft

Author: marktayl1

src/backend/app_kernel.py

@@ -667,8 +667,7 @@ async def get_plans(
             "UserIdNotFound", {"status_code": 400, "detail": "no user"}
         )
         raise HTTPException(status_code=400, detail="no user")
-    
-    await connection_config.send_status_update_async("Test message from get_plans", user_id)
+  
 
     #### <To do: Francia> Replace the following with code to get plan run history from the database
 

src/backend/common/models/messages_kernel.py

@@ -178,6 +178,10 @@ class StartingTask(KernelBaseModel):
     creator: str
     logo: str
 
+class TeamSelectionRequest(KernelBaseModel):
+    """Request model for team selection."""
+    team_id: str
+    session_id: Optional[str] = None
 
 class TeamConfiguration(BaseDataModel):
     """Represents a team configuration stored in the database."""

src/backend/v3/api/router.py

@@ -1,4 +1,5 @@
 import asyncio
+import contextvars
 import json
 import logging
 import uuid
@@ -9,28 +10,21 @@
 from common.config.app_config import config
 from common.database.database_factory import DatabaseFactory
 from common.models.messages_kernel import (GeneratePlanRequest, InputTask,
-                                           Plan, PlanStatus)
+                                           TeamSelectionRequest)
 from common.utils.event_utils import track_event_if_configured
 from common.utils.utils_kernel import rai_success, rai_validate_team_config
 from fastapi import (APIRouter, BackgroundTasks, Depends, FastAPI, File,
                      HTTPException, Request, UploadFile, WebSocket,
                      WebSocketDisconnect)
 from kernel_agents.agent_factory import AgentFactory
-from pydantic import BaseModel
 from semantic_kernel.agents.runtime import InProcessRuntime
 from v3.common.services.team_service import TeamService
-from v3.config.settings import connection_config, team_config
+from v3.config.settings import connection_config, current_user_id, team_config
 from v3.orchestration.orchestration_manager import OrchestrationManager
 
 router = APIRouter()
 logger = logging.getLogger(__name__)
 
-class TeamSelectionRequest(BaseModel):
-    """Request model for team selection."""
-    team_id: str
-    session_id: Optional[str] = None
-
-
 app_v3 = APIRouter(
     prefix="/api/v3",
     responses={404: {"description": "Not found"}},
@@ -50,32 +44,38 @@ async def start_comms(websocket: WebSocket, process_id: str):
         authenticated_user = get_authenticated_user_details(request_headers=headers)
         user_id = authenticated_user.get("user_principal_id")
         if not user_id:
-            user_id = f"anonymous_{process_id}"
+            user_id = "00000000-0000-0000-0000-000000000000"
     except Exception as e:
         logging.warning(f"Could not extract user from WebSocket headers: {e}")
-        user_id = f"anonymous_{user_id}"
+        user_id = "00000000-0000-0000-0000-000000000000"
 
-    # Add to the connection manager for backend updates
+    current_user_id.set(user_id)
 
-    connection_config.add_connection(user_id, websocket)
-    track_event_if_configured("WebSocketConnectionAccepted", {"process_id": "user_id"})
+    # Add to the connection manager for backend updates
+    connection_config.add_connection(process_id=process_id, connection=websocket, user_id=user_id)
+    track_event_if_configured("WebSocketConnectionAccepted", {"process_id": process_id, "user_id": user_id})
 
       # Keep the connection open - FastAPI will close the connection if this returns
-    while True:
-        # no expectation that we will receive anything from the client but this keeps
-        # the connection open and does not take cpu cycle
-        try:
-            await websocket.receive_text()
-        except asyncio.TimeoutError:
-            pass
-
-        except WebSocketDisconnect:
-            track_event_if_configured("WebSocketDisconnect", {"process_id": process_id})
-            logging.info(f"Client disconnected from batch {process_id}")
-            await connection_config.close_connection(user_id)
-        except Exception as e:
-            logging.error("Error in WebSocket connection", error=str(e))
-            await connection_config.close_connection(user_id)
+    try:
+        # Keep the connection open - FastAPI will close the connection if this returns
+        while True:
+            # no expectation that we will receive anything from the client but this keeps
+            # the connection open and does not take cpu cycle
+            try:
+                message = await websocket.receive_text()
+                logging.debug(f"Received WebSocket message from {user_id}: {message}")
+            except asyncio.TimeoutError:
+                pass
+            except WebSocketDisconnect:
+                track_event_if_configured("WebSocketDisconnect", {"process_id": process_id, "user_id": user_id})
+                logging.info(f"Client disconnected from batch {process_id}")
+                break
+    except Exception as e:
+        # Fixed logging syntax - removed the error= parameter
+        logging.error(f"Error in WebSocket connection: {str(e)}")
+    finally:
+        # Always clean up the connection
+        await connection_config.close_connection(user_id)
 
 @app_v3.get("/init_team")
 async def init_team(
@@ -112,7 +112,7 @@ async def init_team(
           )
 
       # Set as current team in memory
-      team_config.set_current_team(user_id=user_id, team_config=team_configuration)
+      team_config.set_current_team(user_id=user_id, team_configuration=team_configuration)
 
       # Initialize agent team for this user session
       await OrchestrationManager.get_current_or_new_orchestration(user_id=user_id, team_config=team_configuration)
@@ -227,8 +227,16 @@ async def process_request(background_tasks: BackgroundTasks, input_task: InputTa
         input_task.session_id = str(uuid.uuid4())
 
     try:
-        background_tasks.add_task(OrchestrationManager.run_orchestration, user_id, input_task)
-        #await connection_config.send_status_update_async("Test message from process_request", user_id)
+        current_user_id.set(user_id)  # Set context
+        current_context = contextvars.copy_context()  # Capture context
+        # background_tasks.add_task(
+        #     lambda: current_context.run(lambda:OrchestrationManager().run_orchestration, user_id, input_task)
+        # )
+
+        async def run_with_context():
+            return await current_context.run(OrchestrationManager().run_orchestration, user_id, input_task)
+
+        background_tasks.add_task(run_with_context)
 
         return {
             "status": "Request started successfully",
@@ -788,46 +796,4 @@ async def get_search_indexes_endpoint(request: Request):
         return {"search_summary": summary}
     except Exception as e:
         logging.error(f"Error retrieving search indexes: {str(e)}")
-        raise HTTPException(status_code=500, detail="Internal server error occurred")
-    
-
-# @app_v3.websocket("/socket/{process_id}")
-# async def process_outputs(websocket: WebSocket, process_id: str):
-#     """ Web-Socket endpoint for real-time process status updates. """
-        
-#     # Always accept the WebSocket connection first
-#     await websocket.accept()
-
-#     user_id = None
-#     try:
-#         # WebSocket headers are different, try to get user info
-#         headers = dict(websocket.headers)
-#         authenticated_user = get_authenticated_user_details(request_headers=headers)
-#         user_id = authenticated_user.get("user_principal_id")
-#         if not user_id:
-#             user_id = f"anonymous_{process_id}"
-#     except Exception as e:
-#         logger.warning(f"Could not extract user from WebSocket headers: {e}")
-#         # user_id = f"anonymous_{user_id}"
-
-#     # Add to the connection manager for backend updates
-
-#     connection_config.add_connection(user_id, websocket)
-#     track_event_if_configured("WebSocketConnectionAccepted", {"process_id": user_id})
-
-#       # Keep the connection open - FastAPI will close the connection if this returns
-#     while True:
-#         # no expectation that we will receive anything from the client but this keeps
-#         # the connection open and does not take cpu cycle
-#         try:
-#             await websocket.receive_text()
-#         except asyncio.TimeoutError:
-#             pass
-
-#         except WebSocketDisconnect:
-#             track_event_if_configured("WebSocketDisconnect", {"process_id": user_id})
-#             logger.info(f"Client disconnected from batch {user_id}")
-#             await connection_config.close_connection(user_id)
-#         except Exception as e:
-#             logger.error("Error in WebSocket connection", error=str(e))
-#             await connection_config.close_connection(user_id)
+        raise HTTPException(status_code=500, detail="Internal server error occurred")

src/backend/v3/callbacks/response_handlers.py

@@ -2,14 +2,23 @@
 Enhanced response callbacks for employee onboarding agent system.
 Provides detailed monitoring and response handling for different agent types.
 """
-
+import asyncio
 import sys
 
 from semantic_kernel.contents import (ChatMessageContent,
                                       StreamingChatMessageContent)
+from v3.config.settings import connection_config, current_user_id
 
 coderagent = False
 
+# Module-level variable to store current user_id
+_current_user_id: str = None
+
+def set_user_context(user_id: str):
+    """Set the user context for callbacks in this module."""
+    global _current_user_id
+    _current_user_id = user_id
+
 def agent_response_callback(message: ChatMessageContent) -> None:
     """Observer function to print detailed information about streaming messages."""
     global coderagent
@@ -31,8 +40,19 @@ def agent_response_callback(message: ChatMessageContent) -> None:
         return
     elif coderagent == True:
         coderagent = False
+        
     role = getattr(message, 'role', 'unknown')
 
+    # Send to WebSocket
+    if _current_user_id:
+        try:
+            asyncio.create_task(connection_config.send_status_update_async({
+                "type": "agent_message",
+                "data": {"agent_name": agent_name, "content": message.content, "role": role}
+            }, _current_user_id))
+        except Exception as e:
+            print(f"Error sending WebSocket message: {e}")
+
     print(f"\n🧠 **{agent_name}** [{message_type}] (role: {role})")
     print("-" * (len(agent_name) + len(message_type) + 10))
     if message.items[-1].content_type == 'function_call':