1010
1111class AgentRegistry :
1212 """Global registry for tracking and managing all agent instances across the application."""
13-
13+
1414 def __init__ (self ):
1515 self .logger = logging .getLogger (__name__ )
1616 self ._lock = threading .Lock ()
1717 self ._all_agents : WeakSet = WeakSet ()
1818 self ._agent_metadata : Dict [int , Dict [str , Any ]] = {}
19-
19+
2020 def register_agent (self , agent : Any , user_id : Optional [str ] = None ) -> None :
2121 """Register an agent instance for tracking and lifecycle management."""
2222 with self ._lock :
@@ -31,7 +31,7 @@ def register_agent(self, agent: Any, user_id: Optional[str] = None) -> None:
3131 self .logger .info (f"Registered agent: { type (agent ).__name__ } (ID: { agent_id } , User: { user_id } )" )
3232 except Exception as e :
3333 self .logger .error (f"Failed to register agent: { e } " )
34-
34+
3535 def unregister_agent (self , agent : Any ) -> None :
3636 """Unregister an agent instance."""
3737 with self ._lock :
@@ -43,34 +43,34 @@ def unregister_agent(self, agent: Any) -> None:
4343 self .logger .info (f"Unregistered agent: { metadata .get ('type' , 'Unknown' )} (ID: { agent_id } )" )
4444 except Exception as e :
4545 self .logger .error (f"Failed to unregister agent: { e } " )
46-
46+
4747 def get_all_agents (self ) -> List [Any ]:
4848 """Get all currently registered agents."""
4949 with self ._lock :
5050 return list (self ._all_agents )
51-
51+
5252 def get_agent_count (self ) -> int :
5353 """Get the total number of registered agents."""
5454 with self ._lock :
5555 return len (self ._all_agents )
56-
56+
5757 async def cleanup_all_agents (self ) -> None :
5858 """Clean up all registered agents across all users."""
5959 all_agents = self .get_all_agents ()
60-
60+
6161 if not all_agents :
6262 self .logger .info ("No agents to clean up" )
6363 return
64-
64+
6565 self .logger .info (f"🧹 Starting cleanup of { len (all_agents )} total agents" )
66-
66+
6767 # Log agent details for debugging
6868 for i , agent in enumerate (all_agents ):
6969 agent_name = getattr (agent , 'agent_name' , getattr (agent , 'name' , type (agent ).__name__ ))
7070 agent_type = type (agent ).__name__
7171 has_close = hasattr (agent , 'close' )
72- self .logger .info (f"Agent { i + 1 } : { agent_name } (Type: { agent_type } , Has close(): { has_close } )" )
73-
72+ self .logger .info (f"Agent { i + 1 } : { agent_name } (Type: { agent_type } , Has close(): { has_close } )" )
73+
7474 # Clean up agents concurrently
7575 cleanup_tasks = []
7676 for agent in all_agents :
@@ -80,59 +80,59 @@ async def cleanup_all_agents(self) -> None:
8080 agent_name = getattr (agent , 'agent_name' , getattr (agent , 'name' , type (agent ).__name__ ))
8181 self .logger .warning (f"⚠️ Agent { agent_name } has no close() method - just unregistering from registry" )
8282 self .unregister_agent (agent )
83-
83+
8484 if cleanup_tasks :
8585 self .logger .info (f"🔄 Executing { len (cleanup_tasks )} cleanup tasks..." )
8686 results = await asyncio .gather (* cleanup_tasks , return_exceptions = True )
87-
87+
8888 # Log any exceptions that occurred during cleanup
8989 success_count = 0
9090 for i , result in enumerate (results ):
9191 if isinstance (result , Exception ):
9292 self .logger .error (f"❌ Error cleaning up agent { i } : { result } " )
9393 else :
9494 success_count += 1
95-
95+
9696 self .logger .info (f"✅ Successfully cleaned up { success_count } /{ len (cleanup_tasks )} agents" )
97-
97+
9898 # Clear all tracking
9999 with self ._lock :
100100 self ._all_agents .clear ()
101101 self ._agent_metadata .clear ()
102-
102+
103103 self .logger .info ("🎉 Completed cleanup of all agents" )
104-
104+
105105 async def _safe_close_agent (self , agent : Any ) -> None :
106106 """Safely close an agent with error handling."""
107107 try :
108108 agent_name = getattr (agent , 'agent_name' , getattr (agent , 'name' , type (agent ).__name__ ))
109109 self .logger .info (f"Closing agent: { agent_name } " )
110-
110+
111111 # Call the agent's close method - it should handle Azure deletion and registry cleanup
112112 if asyncio .iscoroutinefunction (agent .close ):
113113 await agent .close ()
114114 else :
115115 agent .close ()
116-
116+
117117 self .logger .info (f"Successfully closed agent: { agent_name } " )
118-
118+
119119 except Exception as e :
120120 agent_name = getattr (agent , 'agent_name' , getattr (agent , 'name' , type (agent ).__name__ ))
121121 self .logger .error (f"Failed to close agent { agent_name } : { e } " )
122-
122+
123123 def get_registry_status (self ) -> Dict [str , Any ]:
124124 """Get current status of the agent registry for debugging and monitoring."""
125125 with self ._lock :
126126 status = {
127127 'total_agents' : len (self ._all_agents ),
128128 'agent_types' : {}
129129 }
130-
130+
131131 # Count agents by type
132132 for agent in self ._all_agents :
133133 agent_type = type (agent ).__name__
134134 status ['agent_types' ][agent_type ] = status ['agent_types' ].get (agent_type , 0 ) + 1
135-
135+
136136 return status
137137
138138
0 commit comments