🔧 Fix: Resource Leaks in AgentV2 with Proper Cleanup

[ivosetyadi]

🔧 Fix: Resource Leaks in AgentV2 with Proper Cleanup

Summary

Implemented comprehensive resource management to prevent websocket handler memory leaks. Added del destructor, async context manager support, and explicit cleanup() method to ensure proper resource cleanup.

Problem Statement

Original Issue (MEDIUM Severity):
Websocket handlers registered in AgentV2.init were not reliably cleaned up, causing memory leaks when:

• Objects destroyed without proper flow
• Cleanup failed silently
• No way to track cleanup state

Solution Implemented

1. Added del Destructor

Ensures cleanup even if object destroyed abnormally

• Guaranteed cleanup on garbage collection
• Safe with hasattr() for partial initialization
• Tracks state with _handler_registered flag

2. Async Context Manager Support

Enable Pythonic resource management:

async with AgentV2(...) as agent:
    result = await agent.run()
# Automatic cleanup

3. Explicit cleanup() Method

Manual cleanup with proper logging

• Can be called explicitly when needed
• Idempotent - safe to call multiple times
• Proper success/failure logging

4. Updated Finally Blocks

Use centralized cleanup() method for consistency

5. State Tracking

_handler_registered flag prevents double cleanup

Benefits

Before	After
Only cleanup in try-finally	Multiple cleanup mechanisms
No cleanup on destruction	del cleanup
No state tracking	_handler_registered flag
No context manager	Full async context support
Silent failures	Proper logging

Usage Patterns

Recommended:

async with AgentV2(...) as agent:
    await agent.run()

Explicit:

agent = AgentV2(...)
try:
    await agent.run()
finally:
    await agent.cleanup()

Automatic:

agent = AgentV2(...)
await agent.run()
# __del__ ensures cleanup

Testing

• ✅ Python syntax validation passed
• ✅ Backward compatible
• ✅ No breaking changes
• ✅ Context manager optional

Code Changes

• File: backend/app/ai/agent_v2.py
• Changes: +40 insertions, -9 deletions
• Net: +31 lines

Checklist

• Added del destructor
• Added context manager support
• Added explicit cleanup() method
• Added state tracking flag
• Updated finally blocks
• Proper logging

[yochze]

thanks @ivosetyadi!
@claude review

[Copilot]

Pull request overview

This PR implements comprehensive resource management for AgentV2 to prevent websocket handler memory leaks. The changes add multiple cleanup mechanisms including a destructor, async context manager support, and an explicit cleanup method.

Key Changes:

• Added __del__ destructor for automatic cleanup on garbage collection
• Implemented async context manager (__aenter__/__aexit__) for Pythonic resource management
• Added explicit cleanup() method with idempotency checks and logging

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The __del__ destructor calls synchronous websocket_manager.remove_handler() which is safe since it's a synchronous method. However, there's a potential race condition: if the async cleanup() method is running in another thread/task when the object is being garbage collected, both could try to remove the handler simultaneously. While the idempotency check with _handler_registered helps, the flag check and update are not atomic operations.

Consider using a lock or atomic flag to ensure thread-safe cleanup, or document that cleanup should complete before object destruction.

[Copilot]

The _handler_registered flag is initialized to True at the end of __init__, but the websocket handler is registered earlier in the constructor (line 79). If an exception occurs in __init__ between line 79 (handler registration) and line 156 (flag initialization), the __del__ method won't clean up the handler because hasattr(self, '_handler_registered') will return False.

Move the _handler_registered = True line to immediately after the websocket_manager.add_handler() call on line 79 to ensure the flag accurately reflects the handler registration state.

Suggested change

	# Track if handler was registered for cleanup
	self._handler_registered = True
	def __del__(self):
	"""Cleanup websocket handler on object destruction to prevent memory leaks."""
	if hasattr(self, '_handler_registered') and self._handler_registered:
	try:
	websocket_manager.remove_handler(self._handle_completion_update)
	self._handler_registered = False
	except Exception:
	# Silently ignore cleanup errors during destruction
	pass
	def __del__(self):
	"""Cleanup websocket handler on object destruction to prevent memory leaks."""
	try:
	websocket_manager.remove_handler(self._handle_completion_update)
	except Exception:
	# Silently ignore cleanup errors during destruction
	pass

[Copilot]

The new resource management features (del, aenter, aexit, cleanup) lack test coverage. Consider adding tests that verify:

1. Cleanup is called when using async context manager
2. Handler is properly removed after cleanup
3. Idempotency of cleanup (calling it multiple times is safe)
4. del properly cleans up when object is destroyed
5. The _handler_registered flag correctly tracks state

[Copilot]

Overly complex 'del' method.

Suggested change

	"""Cleanup websocket handler on object destruction to prevent memory leaks."""
	if hasattr(self, '_handler_registered') and self._handler_registered:
	try:
	websocket_manager.remove_handler(self._handle_completion_update)
	self._handler_registered = False
	except Exception:
	# Silently ignore cleanup errors during destruction
	pass
	"""Best-effort notice on object destruction; rely on explicit cleanup instead."""
	try:
	if getattr(self, "_handler_registered", False):
	# Object is being garbage-collected while a handler may still be registered.
	# Actual cleanup should be done via `cleanup` or the async context manager.
	logger.debug("AgentV2 instance deleted without explicit cleanup.")
	except Exception:
	# Silently ignore any errors during interpreter shutdown or GC.
	pass