@@ -4464,6 +4464,146 @@ def _get_tools_cache_key(self, tools):
44644464 @contextlib .contextmanager
44654465
44664466
4467+ # -------------------------------------------------------------------------
4468+ # Resource Lifecycle Management
4469+ # -------------------------------------------------------------------------
4470+
4471+ def close (self ) -> None :
4472+ """Synchronously close the agent and clean up resources."""
4473+ if getattr (self , '_closed' , False ):
4474+ return
4475+
4476+ # Memory cleanup
4477+ try :
4478+ memory = getattr (self , "_memory_instance" , None )
4479+ if memory and hasattr (memory , 'close_connections' ):
4480+ memory .close_connections ()
4481+ except Exception as e :
4482+ logger .warning (f"Memory cleanup failed: { e } " )
4483+
4484+ # MCP cleanup
4485+ try :
4486+ if hasattr (self , '_mcp_clients' ) and self ._mcp_clients :
4487+ for client_name , client in self ._mcp_clients .items ():
4488+ if hasattr (client , 'close' ):
4489+ client .close ()
4490+ self ._mcp_clients .clear ()
4491+ except Exception as e :
4492+ logger .warning (f"MCP cleanup failed: { e } " )
4493+
4494+ # Server registry cleanup
4495+ try :
4496+ self ._cleanup_server_registrations ()
4497+ except Exception as e :
4498+ logger .warning (f"Server cleanup failed: { e } " )
4499+
4500+ # Background tasks cleanup
4501+ try :
4502+ for task in getattr (self , '_background_tasks' , []):
4503+ if hasattr (task , 'cancel' ):
4504+ task .cancel ()
4505+ except Exception as e :
4506+ logger .warning (f"Task cleanup failed: { e } " )
4507+
4508+ # Always set closed flag
4509+ self ._closed = True
4510+
4511+ async def aclose (self ) -> None :
4512+ """Async version of close() for async context managers."""
4513+ try :
4514+ # Close memory connections asynchronously if supported
4515+ if hasattr (self , 'memory' ) and self .memory :
4516+ if hasattr (self .memory , 'aclose' ):
4517+ await self .memory .aclose ()
4518+ elif hasattr (self .memory , 'close_connections' ):
4519+ self .memory .close_connections ()
4520+
4521+ # Close MCP sessions asynchronously if supported
4522+ if hasattr (self , '_mcp_clients' ) and self ._mcp_clients :
4523+ for client in self ._mcp_clients .values ():
4524+ if hasattr (client , 'aclose' ):
4525+ await client .aclose ()
4526+ elif hasattr (client , 'close' ):
4527+ client .close ()
4528+ self ._mcp_clients .clear ()
4529+
4530+ # Clean up server registrations and tasks
4531+ self ._cleanup_server_registrations ()
4532+
4533+ if hasattr (self , '_background_tasks' ):
4534+ for task in self ._background_tasks :
4535+ if hasattr (task , 'cancel' ):
4536+ task .cancel ()
4537+ # Wait for cancellation to complete
4538+ try :
4539+ await task
4540+ except asyncio .CancelledError :
4541+ pass
4542+
4543+ self ._closed = True
4544+
4545+ except Exception as e :
4546+ logger .warning (f"Error during async agent cleanup: { e } " )
4547+
4548+ def _cleanup_server_registrations (self ) -> None :
4549+ """Clean up global server registry entries for this agent."""
4550+ try :
4551+ agent_id = self .agent_id
4552+ with _server_lock :
4553+ # Remove from _registered_agents
4554+ ports_to_clean = []
4555+ for port , path_dict in _registered_agents .items ():
4556+ paths_to_remove = []
4557+ for path , registered_id in path_dict .items ():
4558+ if registered_id == agent_id :
4559+ paths_to_remove .append (path )
4560+
4561+ for path in paths_to_remove :
4562+ del path_dict [path ]
4563+
4564+ # If no paths left for this port, mark port for cleanup
4565+ if not path_dict :
4566+ ports_to_clean .append (port )
4567+
4568+ # Clean up empty port entries
4569+ for port in ports_to_clean :
4570+ _registered_agents .pop (port , None )
4571+ _server_started .pop (port , None )
4572+ # Note: We don't clean up _shared_apps here as other agents might be using them
4573+
4574+ except Exception as e :
4575+ logger .warning (f"Error cleaning up server registrations: { e } " )
4576+
4577+ def __enter__ (self ):
4578+ """Context manager entry."""
4579+ return self
4580+
4581+ def __exit__ (self , exc_type , exc_val , exc_tb ):
4582+ """Context manager exit - clean up resources."""
4583+ self .close ()
4584+
4585+ async def __aenter__ (self ):
4586+ """Async context manager entry."""
4587+ return self
4588+
4589+ async def __aexit__ (self , exc_type , exc_val , exc_tb ):
4590+ """Async context manager exit - clean up resources."""
4591+ await self .aclose ()
4592+
4593+ def __del__ (self ):
4594+ """Destructor - ensure resources are cleaned up."""
4595+ try :
4596+ if not getattr (self , '_closed' , False ):
4597+ self .close ()
4598+ except Exception :
4599+ # Ignore errors in destructor to avoid issues during shutdown
4600+ pass
4601+
4602+ @property
4603+ def is_closed (self ) -> bool :
4604+ """Returns True if the agent has been closed."""
4605+ return getattr (self , '_closed' , False )
4606+
44674607 def __str__ (self ):
44684608 return f"Agent(name='{ self .name } ', role='{ self .role } ', goal='{ self .goal } ')"
44694609
0 commit comments