mirror of
https://gitee.com/wanwujie/deer-flow
synced 2026-04-21 13:24:44 +08:00
fix(subagents): cleanup background tasks after completion to prevent memory leak (#1030)
* fix(subagents): cleanup background tasks after completion to prevent memory leak Added cleanup_background_task() function to remove completed subagent results from the global _background_tasks dict. Found a small issue: completed tasks were never removed, causing memory to grow indefinitely with each subagent execution. Alternative approaches considered: - Future + SubagentHandle pattern: Not chosen due to requiring refactoring Chose the simple cleanup approach for minimal code changes while effectively resolving the memory leak. Changes: - Add cleanup_background_task() in executor.py - Call cleanup in all task_tool return paths (completed, failed, timed out) * fix(subagents): prevent race condition in background task cleanup Address Copilot review feedback on memory leak fix: - Add terminal state check in cleanup_background_task() to only remove tasks that are COMPLETED/FAILED/TIMED_OUT or have completed_at set - Remove cleanup call from polling safety-timeout branch in task_tool since the task may still be running - Add comprehensive tests for cleanup behavior including: - Verification that cleanup is called on terminal states - Verification that cleanup is NOT called on polling timeout - Tests for terminal state check logic in executor This prevents KeyError when the background executor tries to update a task that was prematurely removed from _background_tasks. --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
@@ -452,3 +452,40 @@ def list_background_tasks() -> list[SubagentResult]:
|
||||
"""
|
||||
with _background_tasks_lock:
|
||||
return list(_background_tasks.values())
|
||||
|
||||
|
||||
def cleanup_background_task(task_id: str) -> None:
|
||||
"""Remove a completed task from background tasks.
|
||||
|
||||
Should be called by task_tool after it finishes polling and returns the result.
|
||||
This prevents memory leaks from accumulated completed tasks.
|
||||
|
||||
Only removes tasks that are in a terminal state (COMPLETED/FAILED/TIMED_OUT)
|
||||
to avoid race conditions with the background executor still updating the task entry.
|
||||
|
||||
Args:
|
||||
task_id: The task ID to remove.
|
||||
"""
|
||||
with _background_tasks_lock:
|
||||
result = _background_tasks.get(task_id)
|
||||
if result is None:
|
||||
# Nothing to clean up; may have been removed already.
|
||||
logger.debug("Requested cleanup for unknown background task %s", task_id)
|
||||
return
|
||||
|
||||
# Only clean up tasks that are in a terminal state to avoid races with
|
||||
# the background executor still updating the task entry.
|
||||
is_terminal_status = result.status in {
|
||||
SubagentStatus.COMPLETED,
|
||||
SubagentStatus.FAILED,
|
||||
SubagentStatus.TIMED_OUT,
|
||||
}
|
||||
if is_terminal_status or result.completed_at is not None:
|
||||
del _background_tasks[task_id]
|
||||
logger.debug("Cleaned up background task: %s", task_id)
|
||||
else:
|
||||
logger.debug(
|
||||
"Skipping cleanup for non-terminal background task %s (status=%s)",
|
||||
task_id,
|
||||
result.status.value if hasattr(result.status, "value") else result.status,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user