fix(cloakserve): add idle cleanup for seeded profiles (#352)

* fix(cloakserve): add idle cleanup for seeded profiles

* docs(cloakserve): document idle process cleanup

README.md

@@ -892,6 +892,10 @@ docker run -d --name cloak -p 127.0.0.1:9222:9222 cloakhq/cloakbrowser \
 # Headed mode (renders to Xvfb inside container)
 docker run -d --name cloak -p 127.0.0.1:9222:9222 cloakhq/cloakbrowser \
   cloakserve --headless=false
+
+# Reap disconnected per-seed browser processes after 5 minutes
+docker run -d --name cloak -p 127.0.0.1:9222:9222 cloakhq/cloakbrowser \
+  cloakserve --idle-timeout=300
 ```
 
 Stop the server:
@@ -943,7 +947,9 @@ b4 = pw.chromium.connect_over_cdp(
 )
 ```
 
-Supported query params: `fingerprint`, `timezone`, `locale`, `platform`, `platform-version`, `brand`, `brand-version`, `gpu-vendor`, `gpu-renderer`, `hardware-concurrency`, `device-memory`, `screen-width`, `screen-height`, `proxy`, `geoip`. Same seed reuses the same process (first connection's params win). No seed = shared default process (backward compatible). Check active processes at `GET /` (returns JSON with PIDs, ports, and connection counts).
+Supported query params: `fingerprint`, `timezone`, `locale`, `platform`, `platform-version`, `brand`, `brand-version`, `gpu-vendor`, `gpu-renderer`, `hardware-concurrency`, `device-memory`, `screen-width`, `screen-height`, `proxy`, `geoip`. Same seed reuses the same process (first connection's params win). No seed = shared default process (backward compatible).
+
+By default, per-seed processes stay alive until `cloakserve` exits. If clients create many unique seeds, set `--idle-timeout=SECONDS` or `CLOAKSERVE_IDLE_TIMEOUT=SECONDS` to automatically terminate a seed's Chrome process after its last CDP WebSocket disconnects. `0`, `off`, `false`, `none`, or `disabled` disable idle cleanup. When cleanup runs, the seed's temporary profile directory under `--data-dir` is removed too. Check active processes at `GET /` (returns JSON with PIDs, ports, connection counts, idle timeout, and pending cleanup status).
 
 **Persistent profiles** — mount a volume to keep cookies and sessions across container restarts:
 

bin/cloakserve

@@ -181,6 +181,7 @@ class ChromePool:
         default_seed: str | None = None,
         default_locale: str | None = None,
         default_timezone: str | None = None,
+        idle_timeout: float = 0.0,
     ):
         self._binary = binary
         self._global_args = global_args
@@ -189,12 +190,14 @@ class ChromePool:
         self._default_seed = default_seed
         self._default_locale = default_locale
         self._default_timezone = default_timezone
+        self._idle_timeout = idle_timeout
         self._processes: dict[str, ChromeProcess] = {}
         self._default: ChromeProcess | None = None
         self._locks: dict[str, asyncio.Lock] = {}
         self._next_port = BASE_CDP_PORT
         # Connection refcounting for status reporting
         self._connections: dict[str, int] = {}
+        self._idle_tasks: dict[str, asyncio.Task] = {}
 
     def _get_lock(self, seed: str) -> asyncio.Lock:
         if seed not in self._locks:
@@ -224,6 +227,7 @@ class ChromePool:
 
     def connect(self, seed_key: str) -> None:
         """Increment connection refcount for a seed."""
+        self._cancel_idle_cleanup(seed_key)
         self._connections[seed_key] = self._connections.get(seed_key, 0) + 1
 
     def disconnect(self, seed_key: str) -> None:
@@ -231,9 +235,54 @@ class ChromePool:
         count = self._connections.get(seed_key, 0) - 1
         if count <= 0:
             self._connections.pop(seed_key, None)
+            self._schedule_idle_cleanup(seed_key)
         else:
             self._connections[seed_key] = count
 
+    def _cancel_idle_cleanup(self, seed_key: str) -> None:
+        task = self._idle_tasks.pop(seed_key, None)
+        if task is None or task.done():
+            return
+        try:
+            current_task = asyncio.current_task()
+        except RuntimeError:
+            current_task = None
+        if task is not current_task:
+            task.cancel()
+
+    def _discard_idle_task(self, seed_key: str, task: asyncio.Task) -> None:
+        if self._idle_tasks.get(seed_key) is task:
+            self._idle_tasks.pop(seed_key, None)
+
+    def _schedule_idle_cleanup(self, seed_key: str) -> None:
+        if self._idle_timeout <= 0 or seed_key not in self._processes:
+            return
+
+        self._cancel_idle_cleanup(seed_key)
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            return
+
+        task = loop.create_task(
+            self._cleanup_after_idle(seed_key, self._idle_timeout),
+            name=f"cloakserve-idle-cleanup-{seed_key}",
+        )
+        self._idle_tasks[seed_key] = task
+        task.add_done_callback(lambda done_task: self._discard_idle_task(seed_key, done_task))
+
+    async def _cleanup_after_idle(self, seed_key: str, timeout: float) -> None:
+        try:
+            await asyncio.sleep(timeout)
+            if self._connections.get(seed_key, 0) > 0 or seed_key not in self._processes:
+                return
+            logger.info("Cleaning up idle Chrome process (seed=%s)", seed_key)
+            await self._cleanup_process(seed_key)
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            logger.exception("Idle cleanup failed for seed=%s", seed_key)
+
     async def get_or_launch(
         self,
         seed: str | None,
@@ -271,6 +320,8 @@ class ChromePool:
             if seed_key in self._processes:
                 proc = self._processes[seed_key]
                 if proc.process.poll() is None:
+                    if seed_key in self._idle_tasks:
+                        self._schedule_idle_cleanup(seed_key)
                     if any([extra_args, timezone, locale, proxy, geoip]):
                         logger.warning(
                             "Seed %s already running (port %d, tz=%s, locale=%s, proxy=%s) — "
@@ -360,6 +411,7 @@ class ChromePool:
 
     async def _cleanup_process(self, key: str) -> None:
         """Terminate a Chrome process and clean up."""
+        self._cancel_idle_cleanup(key)
         proc = self._processes.pop(key, None)
         if not proc:
             return
@@ -377,6 +429,13 @@ class ChromePool:
 
     async def shutdown(self) -> None:
         """Terminate all Chrome processes."""
+        idle_tasks = list(self._idle_tasks.values())
+        self._idle_tasks.clear()
+        for task in idle_tasks:
+            if not task.done():
+                task.cancel()
+        if idle_tasks:
+            await asyncio.gather(*idle_tasks, return_exceptions=True)
         for key in list(self._processes.keys()):
             await self._cleanup_process(key)
         logger.info("All Chrome processes terminated")
@@ -479,6 +538,7 @@ async def handle_root(request: web.Request) -> web.Response:
                 "port": proc.cdp_port,
                 "seed": proc.seed,
                 "connections": pool._connections.get(key, 0),
+                "idle_cleanup_pending": key in pool._idle_tasks,
                 "timezone": proc.timezone,
                 "locale": proc.locale,
                 "proxy": proc.proxy,
@@ -486,6 +546,7 @@ async def handle_root(request: web.Request) -> web.Response:
     return web.json_response({
         "status": "ok",
         "active": len(processes),
+        "idle_timeout": pool._idle_timeout,
         "processes": processes,
     })
 
@@ -687,6 +748,23 @@ def _default_data_dir() -> str:
     return str(Path.home() / ".cloakbrowser" / "cloakserve")
 
 
+def _parse_idle_timeout(value: str) -> float:
+    value = value.strip()
+    if value.lower() in {"0", "false", "off", "none", "disabled"}:
+        return 0.0
+    timeout = float(value)
+    if timeout < 0:
+        raise ValueError("--idle-timeout must be greater than or equal to 0")
+    return timeout
+
+
+def _default_idle_timeout() -> float:
+    value = os.environ.get("CLOAKSERVE_IDLE_TIMEOUT")
+    if value is None:
+        return 0.0
+    return _parse_idle_timeout(value)
+
+
 def parse_cli_args(argv: list[str]) -> tuple[dict, list[str]]:
     """Parse cloakserve-specific args, return (config, passthrough_args).
 
@@ -702,12 +780,14 @@ def parse_cli_args(argv: list[str]) -> tuple[dict, list[str]]:
         "default_seed": None,
         "default_locale": None,
         "default_timezone": None,
+        "idle_timeout": _default_idle_timeout(),
     }
     passthrough = []
     # Flags consumed by cloakserve (not passed to Chrome)
     consumed_prefixes = (
         "--port=",
         "--data-dir=",
+        "--idle-timeout=",
         "--remote-debugging-port=",
         "--remote-debugging-address=",
     )
@@ -717,6 +797,8 @@ def parse_cli_args(argv: list[str]) -> tuple[dict, list[str]]:
             config["port"] = int(arg.split("=", 1)[1])
         elif arg.startswith("--data-dir="):
             config["data_dir"] = arg.split("=", 1)[1]
+        elif arg.startswith("--idle-timeout="):
+            config["idle_timeout"] = _parse_idle_timeout(arg.split("=", 1)[1])
         elif arg == "--headless=false" or arg == "--headless=False":
             config["headless"] = False
             passthrough.append(arg)
@@ -761,6 +843,7 @@ def main() -> None:
         default_seed=config["default_seed"],
         default_locale=config["default_locale"],
         default_timezone=config["default_timezone"],
+        idle_timeout=config["idle_timeout"],
     )
 
     app = web.Application()

tests/test_cloakserve.py

@@ -93,6 +93,7 @@ class TestParseCliArgs:
         assert config["port"] == 9222
         assert config["headless"] is True
         assert config["data_dir"] is not None
+        assert config["idle_timeout"] == 0.0
         assert passthrough == []
 
     def test_custom_port(self):
@@ -131,6 +132,31 @@ class TestParseCliArgs:
         _, passthrough = parse_cli_args(["--data-dir=/tmp/test"])
         assert not any(a.startswith("--data-dir=") for a in passthrough)
 
+    def test_idle_timeout_not_in_passthrough(self):
+        config, passthrough = parse_cli_args(["--idle-timeout=30", "--no-sandbox"])
+        assert config["idle_timeout"] == 30.0
+        assert "--idle-timeout=30" not in passthrough
+        assert "--no-sandbox" in passthrough
+
+    @pytest.mark.parametrize("value", ["0", "off", "false", "none", "disabled"])
+    def test_idle_timeout_disabled_values(self, value):
+        config, _ = parse_cli_args([f"--idle-timeout={value}"])
+        assert config["idle_timeout"] == 0.0
+
+    def test_idle_timeout_env_default(self, monkeypatch):
+        monkeypatch.setenv("CLOAKSERVE_IDLE_TIMEOUT", "2.5")
+        config, _ = parse_cli_args([])
+        assert config["idle_timeout"] == 2.5
+
+    def test_idle_timeout_cli_overrides_env(self, monkeypatch):
+        monkeypatch.setenv("CLOAKSERVE_IDLE_TIMEOUT", "2.5")
+        config, _ = parse_cli_args(["--idle-timeout=9"])
+        assert config["idle_timeout"] == 9.0
+
+    def test_idle_timeout_rejects_negative_values(self):
+        with pytest.raises(ValueError):
+            parse_cli_args(["--idle-timeout=-1"])
+
     @patch("os.path.exists", return_value=True)
     def test_default_data_dir_docker(self, _mock):
         assert _default_data_dir() == "/tmp/cloakserve"
@@ -431,12 +457,21 @@ class TestHandlerURLRewriting:
 class TestConnectionTracking:
     """Test ChromePool.connect() / disconnect() without real Chrome."""
 
-    def _make_pool(self):
+    def _make_pool(self, idle_timeout: float = 0.0):
         return ChromePool(
             binary="/fake/chrome",
             global_args=[],
             headless=True,
             data_dir="/tmp/test-cloakserve",
+            idle_timeout=idle_timeout,
+        )
+
+    def _track_process(self, pool, seed="seed1"):
+        pool._processes[seed] = SimpleNamespace()
+
+    def _track_live_process(self, pool, seed="seed1"):
+        pool._processes[seed] = SimpleNamespace(
+            process=SimpleNamespace(poll=lambda: None),
         )
 
     def test_connect_increments(self):
@@ -473,6 +508,83 @@ class TestConnectionTracking:
         assert pool._connections["a"] == 1
         assert pool._connections["b"] == 1
 
+    def test_idle_cleanup_disabled_by_default(self):
+        async def run():
+            pool = self._make_pool()
+            self._track_process(pool)
+
+            pool.connect("seed1")
+            pool.disconnect("seed1")
+
+            await asyncio.sleep(0)
+            assert pool._idle_tasks == {}
+
+        asyncio.run(run())
+
+    def test_disconnect_to_zero_schedules_idle_cleanup(self):
+        async def run():
+            pool = self._make_pool(idle_timeout=0.01)
+            self._track_process(pool)
+            cleaned = []
+
+            async def fake_cleanup(seed):
+                cleaned.append(seed)
+                pool._processes.pop(seed, None)
+
+            pool._cleanup_process = fake_cleanup
+            pool.connect("seed1")
+            pool.disconnect("seed1")
+
+            assert "seed1" in pool._idle_tasks
+            await asyncio.sleep(0.05)
+            assert cleaned == ["seed1"]
+            assert "seed1" not in pool._idle_tasks
+
+        asyncio.run(run())
+
+    def test_reconnect_cancels_pending_idle_cleanup(self):
+        async def run():
+            pool = self._make_pool(idle_timeout=0.03)
+            self._track_process(pool)
+            cleaned = []
+
+            async def fake_cleanup(seed):
+                cleaned.append(seed)
+                pool._processes.pop(seed, None)
+
+            pool._cleanup_process = fake_cleanup
+            pool.connect("seed1")
+            pool.disconnect("seed1")
+            assert "seed1" in pool._idle_tasks
+
+            pool.connect("seed1")
+            await asyncio.sleep(0.06)
+
+            assert cleaned == []
+            assert pool._connections["seed1"] == 1
+            assert "seed1" not in pool._idle_tasks
+
+        asyncio.run(run())
+
+    def test_discovery_refreshes_pending_idle_cleanup(self):
+        async def run():
+            pool = self._make_pool(idle_timeout=1.0)
+            self._track_live_process(pool)
+
+            pool.connect("seed1")
+            pool.disconnect("seed1")
+            first_task = pool._idle_tasks["seed1"]
+
+            await pool.get_or_launch("seed1")
+            second_task = pool._idle_tasks["seed1"]
+
+            assert second_task is not first_task
+            pool._cancel_idle_cleanup("seed1")
+            await asyncio.sleep(0)
+            assert "seed1" not in pool._idle_tasks
+
+        asyncio.run(run())
+
 
 # ---------------------------------------------------------------------------
 # Seed validation (CVE fix — path traversal via fingerprint param)