diff --git a/agent/chat_completion_helpers.py b/agent/chat_completion_helpers.py
index 67425f6647..f4c24e9b6d 100644
--- a/agent/chat_completion_helpers.py
+++ b/agent/chat_completion_helpers.py
@@ -2335,7 +2335,7 @@ def cleanup_task_resources(agent, task_id: str) -> None:
     and ``cleanup_browser`` in headed mode (the inactivity reaper handles idle sessions)."""
     def _headed() -> bool:
         try:
-            from tools.browser_tool import _is_headed_mode
+            from tools.browser_tool_cloud import _is_headed_mode
             return _is_headed_mode()
         except Exception:
             return bool(os.environ.get("AGENT_BROWSER_HEADED"))
diff --git a/run_agent.py b/run_agent.py
index 1cc624a452..f6d31bc8e2 100644
--- a/run_agent.py
+++ b/run_agent.py
@@ -102,7 +102,7 @@ if not _loaded_env_paths:
 from model_tools import get_toolset_for_tool
 from tools.terminal_tool_lifecycle import cleanup_vm, get_active_env
 from tools.interrupt import set_interrupt as _set_interrupt
-from tools.browser_tool import cleanup_browser
+from tools.browser_tool_lifecycle import cleanup_browser
 
 from agent.memory_provider import is_trivial_prompt
 from agent.client_lifecycle import ClientLifecycleMixin