clean up user safe websocket calling across contexts

Author: marktayl1

src/backend/v3/callbacks/response_handlers.py

@@ -9,61 +9,61 @@
                                       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:
+def agent_response_callback(message: ChatMessageContent, user_id: str = None) -> None:
     """Observer function to print detailed information about streaming messages."""
-    global coderagent
     # import sys
 
     # Get agent name to determine handling
     agent_name = message.name or "Unknown Agent"
 
+
     # Debug information about the message
-    message_type = type(message).__name__
-    metadata = getattr(message, 'metadata', {})
-    # when streaming code - list the coder info first once - 
-    if 'code' in metadata and metadata['code'] is True:
-        if coderagent == False:
-            print(f"\n🧠 **{agent_name}** [{message_type}]")
-            print("-" * (len(agent_name) + len(message_type) + 10))
-            coderagent = True
-        print(message.content, end='', flush=False)
-        return
-    elif coderagent == True:
-        coderagent = False
+    # message_type = type(message).__name__
+    # metadata = getattr(message, 'metadata', {})
+    # # when streaming code - list the coder info first once - 
+    # if 'code' in metadata and metadata['code'] is True:
+    #     if coderagent == False:
+    #         print(f"\n **{agent_name}** [{message_type}]")
+    #         print("-" * (len(agent_name) + len(message_type) + 10))
+    #         coderagent = True
+    #     print(message.content, end='', flush=False)
+    #     return
+    # elif coderagent == True:
+    #     coderagent = False
 
     role = getattr(message, 'role', 'unknown')
 
     # Send to WebSocket
-    if _current_user_id:
+    if 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))
+            }, 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))
+    print(f"\n **{agent_name}** (role: {role})")
+
     if message.items[-1].content_type == 'function_call':
         print(f"🔧 Function call: {message.items[-1].function_name}, Arguments: {message.items[-1].arguments}")
-    if metadata:
-        print(f"📋 Metadata: {metadata}")
+
 
 # Add this function after your agent_response_callback function
-async def streaming_agent_response_callback(streaming_message: StreamingChatMessageContent, is_final: bool) -> None:
+async def streaming_agent_response_callback(streaming_message: StreamingChatMessageContent, is_final: bool, user_id: str = None) -> None:
     """Simple streaming callback to show real-time agent responses."""
     if streaming_message.name != "CoderAgent":
         # Print streaming content as it arrives
         if hasattr(streaming_message, 'content') and streaming_message.content:
             print(streaming_message.content, end='', flush=False)
+
+            # Send to WebSocket
+            if user_id:
+                try:
+                    await connection_config.send_status_update_async({
+                        "type": "streaming_message",
+                        "data": {"agent_name": streaming_message.name or "Unknown Agent", "content": streaming_message.content, "is_final": is_final}
+                    }, user_id)
+                except Exception as e:
+                    print(f"Error sending streaming WebSocket message: {e}")

src/backend/v3/magentic_agents/proxy_agent.py

@@ -102,7 +102,7 @@ async def invoke(self, message: str,*, thread: AgentThread | None = None,**kwarg
             expected_type=DummyAgentThread,
         )
         # Replace with websocket call when available
-        print(f"\n🤔 ProxyAgent: Another agent is asking for clarification about:")
+        print(f"\nProxyAgent: Another agent is asking for clarification about:")
         print(f"   Request: {message}")
         print("-" * 60)
 

src/backend/v3/orchestration/orchestration_manager.py

@@ -14,8 +14,9 @@
 # Create custom execution settings to fix schema issues
 from semantic_kernel.connectors.ai.open_ai import (
     AzureChatCompletion, OpenAIChatPromptExecutionSettings)
+from semantic_kernel.contents import (ChatMessageContent,
+                                      StreamingChatMessageContent)
 from v3.callbacks.response_handlers import (agent_response_callback,
-                                            set_user_context,
                                             streaming_agent_response_callback)
 from v3.config.settings import config, current_user_id, orchestration_config
 from v3.magentic_agents.magentic_agent_factory import MagenticAgentFactory
@@ -32,7 +33,7 @@ def __init__(self):
         self.user_id: Optional[str] = None
 
     @classmethod
-    async def init_orchestration(cls, agents: List)-> MagenticOrchestration:
+    async def init_orchestration(cls, agents: List, user_id: str = None)-> MagenticOrchestration:
         """Main function to run the agents."""
 
         # Custom execution settings that should work with Azure OpenAI
@@ -59,28 +60,39 @@ def get_token():
                 ),
                 execution_settings=execution_settings
             ),
-            agent_response_callback=agent_response_callback,
-            streaming_agent_response_callback=streaming_agent_response_callback,  # Add streaming callback
+            agent_response_callback=cls._user_aware_agent_callback(user_id),
+            streaming_agent_response_callback=cls._user_aware_streaming_callback(user_id)
         )
         return magentic_orchestration
 
+    @staticmethod
+    def _user_aware_agent_callback(user_id: str):
+        """Factory method that creates a callback with captured user_id"""
+        def callback(message: ChatMessageContent):
+            return agent_response_callback(message, user_id)
+        return callback
+    
+    @staticmethod
+    def _user_aware_streaming_callback(user_id: str):
+        """Factory method that creates a streaming callback with captured user_id"""
+        async def callback(streaming_message: StreamingChatMessageContent, is_final: bool):
+            return await streaming_agent_response_callback(streaming_message, is_final, user_id)
+        return callback
+    
     @classmethod
     async def get_current_or_new_orchestration(self, user_id: str, team_config: TeamConfiguration) -> MagenticOrchestration:
         """get existing orchestration instance."""
         current_orchestration = orchestration_config.get_current_orchestration(user_id)
         if current_orchestration is None:
             factory = MagenticAgentFactory()
-            # to do: change to parsing teams from cosmos db
             agents = await factory.get_agents(team_config_input=team_config)
-            orchestration_config.orchestrations[user_id] = await self.init_orchestration(agents)
+            orchestration_config.orchestrations[user_id] = await self.init_orchestration(agents, user_id)
         return orchestration_config.get_current_orchestration(user_id)
 
     async def run_orchestration(self, user_id, input_task) -> None:
         """ Run the orchestration with user input loop."""
         token = current_user_id.set(user_id)
 
-        set_user_context(user_id)
-
         job_id = str(uuid.uuid4())
         orchestration_config.approvals[job_id] = None
 
@@ -90,9 +102,7 @@ async def run_orchestration(self, user_id, input_task) -> None:
             raise ValueError("Orchestration not initialized for user.")
 
         try:
-            # ADD THIS: Set user_id on the approval manager before invoke
             if hasattr(magentic_orchestration, '_manager') and hasattr(magentic_orchestration._manager, 'current_user_id'):
-                #object.__setattr__(magentic_orchestration._manager, 'current_user_id', user_id)
                 magentic_orchestration._manager.current_user_id = user_id
                 print(f"🔍 DEBUG: Set user_id on manager = {user_id}")
         except Exception as e: