VFS: cap concurrent scripts and enforce execution deadline

backend/routers/vfs.py

@@ -11,7 +11,7 @@
 
 from ..auth import get_current_user
 from ..services import security_audit_service, vfs_service
-from ..services.vfs_service import VfsBudgetExceeded
+from ..services.vfs_service import VfsBudgetExceeded, VfsBusy
 
 router = APIRouter(prefix="/api/v1/me/vfs", tags=["vfs"])
 
@@ -53,6 +53,8 @@ async def run_vfs(
         raise HTTPException(status_code=400, detail=str(e)) from e
     except VfsBudgetExceeded as e:
         raise HTTPException(status_code=413, detail=str(e)) from e
+    except VfsBusy as e:
+        raise HTTPException(status_code=429, detail=str(e), headers={"Retry-After": "2"}) from e
 
 
 @router.post("/searches", status_code=204)

backend/services/vfs_service.py

@@ -17,6 +17,7 @@
 import asyncio
 import functools
 import threading
+import time
 from contextlib import contextmanager
 
 import anyio
@@ -32,16 +33,29 @@
 # on an open connection until the client's own timeout fires.
 MAX_DOCUMENT_READS = 400
 
+# Each active script retains a full VFS model and the documents it reads. This
+# cap keeps one caller's burst below the API process's memory limit.
+MAX_CONCURRENT_SCRIPTS = 4
+
+# Nested ASGI requests must finish inside the script's remaining budget. The
+# worker itself stays attached so a timed-out request never frees a slot while
+# memory-heavy work is still running.
+MAX_SCRIPT_SECONDS = 60
+
 SOURCE_ENTRIES_PAGE = 1000
 
 
 class VfsBudgetExceeded(Exception):
-    """More direct document reads than one shell invocation is allowed.
+    """More document reads or time than one shell invocation is allowed.
     Deliberately not a VfsClientError: the shell downgrades those to per-file
     warnings, and this must abort the whole command. Reads inside a grep sweep
     raise VfsScanBudget instead, which the shell turns into a partial result."""
 
 
+class VfsBusy(Exception):
+    """This API process is already running its safe number of VFS scripts."""
+
+
 class InProcessVfsClient:
     """`VfsClient` served by the running app over nested ASGI calls.
 
@@ -56,6 +70,24 @@ def __init__(self, http: httpx.AsyncClient, loop: asyncio.AbstractEventLoop) ->
         self._scan = False
         self._document_reads = 0
         self._reads_lock = threading.Lock()
+        self._started = time.monotonic()
+
+    def _result_before_deadline(self, future):
+        seconds_left = MAX_SCRIPT_SECONDS - (time.monotonic() - self._started)
+        if seconds_left <= 0:
+            future.cancel()
+            raise VfsBudgetExceeded(
+                f"command ran longer than {MAX_SCRIPT_SECONDS}s; "
+                "scope it to a smaller subtree or use search"
+            )
+        try:
+            return future.result(timeout=seconds_left)
+        except TimeoutError:
+            future.cancel()
+            raise VfsBudgetExceeded(
+                f"command ran longer than {MAX_SCRIPT_SECONDS}s; "
+                "scope it to a smaller subtree or use search"
+            ) from None
 
     @contextmanager
     def internal_calls(self):
@@ -91,7 +123,7 @@ def record_search(self, pattern: str, roots: list[str], docs_scanned: int) -> No
             ),
             self._loop,
         )
-        response = future.result()
+        response = self._result_before_deadline(future)
         if response.status_code >= 400:
             raise VfsClientError(_error_detail(response))
 
@@ -110,7 +142,7 @@ def _request(self, method: str, endpoint: str, **params) -> httpx.Response:
             self._http.request(method, endpoint, params=params or None, headers=headers),
             self._loop,
         )
-        response = future.result()
+        response = self._result_before_deadline(future)
         if response.status_code >= 400:
             raise VfsClientError(_error_detail(response))
         return response
@@ -224,25 +256,37 @@ def _run_script(
     }
 
 
+_scripts_running = 0
+
+
 async def run_vfs_script(app, authorization: str, script: str, cwd: str) -> dict:
     """Execute one read-only shell script against the caller's Stash.
 
     `authorization` is forwarded verbatim onto every nested request, so the VFS
     sees precisely what that credential sees anywhere else in the API.
     """
-    loop = asyncio.get_running_loop()
-    transport = httpx.ASGITransport(app=app)
-    async with httpx.AsyncClient(
-        transport=transport,
-        base_url="http://vfs.internal",
-        # X-Stash-Via tags nested reads as ask-the-stash traffic in the audit
-        # trail (see auth._set_request_via).
-        headers={"Authorization": authorization, "X-Stash-Via": "ask"},
-        timeout=None,
-    ) as http:
-        return await anyio.to_thread.run_sync(
-            functools.partial(_run_script, http, loop, script, cwd)
+    global _scripts_running
+    if _scripts_running >= MAX_CONCURRENT_SCRIPTS:
+        raise VfsBusy(
+            f"{MAX_CONCURRENT_SCRIPTS} VFS commands are already running; retry in a moment"
         )
+    _scripts_running += 1
+    try:
+        loop = asyncio.get_running_loop()
+        transport = httpx.ASGITransport(app=app)
+        async with httpx.AsyncClient(
+            transport=transport,
+            base_url="http://vfs.internal",
+            # X-Stash-Via tags nested reads as ask-the-stash traffic in the audit
+            # trail (see auth._set_request_via).
+            headers={"Authorization": authorization, "X-Stash-Via": "ask"},
+            timeout=None,
+        ) as http:
+            return await anyio.to_thread.run_sync(
+                functools.partial(_run_script, http, loop, script, cwd)
+            )
+    finally:
+        _scripts_running -= 1
 
 
 def _resolve_node(http: httpx.AsyncClient, loop: asyncio.AbstractEventLoop, path: str) -> dict:

backend/tests/test_vfs.py

@@ -7,12 +7,14 @@
 not part of the tree.
 """
 
+import asyncio
+import threading
 from uuid import UUID
 
 import pytest
 from httpx import AsyncClient
 
-from backend.services import source_service
+from backend.services import source_service, vfs_service
 
 from .conftest import unique_name
 
@@ -215,6 +217,78 @@ async def test_document_read_budget_still_aborts_direct_reads(client: AsyncClien
     assert resp.status_code == 413
 
 
+async def test_concurrency_cap_rejects_an_overlapping_request(client: AsyncClient, monkeypatch):
+    """A burst must reject excess work before it builds another VFS model."""
+    monkeypatch.setattr("backend.services.vfs_service.MAX_CONCURRENT_SCRIPTS", 1)
+    started = threading.Event()
+    release = threading.Event()
+
+    def run_script(*_args):
+        started.set()
+        release.wait()
+        return {"stdout": "", "stderr": "", "exit_code": 0, "cwd": "/"}
+
+    monkeypatch.setattr("backend.services.vfs_service._run_script", run_script)
+    api_key, _ = await _register(client)
+    first = asyncio.create_task(_vfs(client, api_key, "ls /"))
+
+    try:
+        assert await asyncio.to_thread(started.wait, 1)
+        second = await _vfs(client, api_key, "ls /")
+
+        assert second.status_code == 429
+        assert second.headers["retry-after"] == "2"
+    finally:
+        release.set()
+
+    assert (await first).status_code == 200
+
+
+async def test_concurrency_slot_is_released_after_failure(client: AsyncClient, monkeypatch):
+    """A crashed command must not permanently reduce the process capacity."""
+    monkeypatch.setattr("backend.services.vfs_service.MAX_CONCURRENT_SCRIPTS", 1)
+    calls = 0
+
+    def run_script(*_args):
+        nonlocal calls
+        calls += 1
+        if calls == 1:
+            raise RuntimeError("script crashed")
+        return {"stdout": "", "stderr": "", "exit_code": 0, "cwd": "/"}
+
+    monkeypatch.setattr("backend.services.vfs_service._run_script", run_script)
+    api_key, _ = await _register(client)
+
+    first = await _vfs(client, api_key, "ls /")
+
+    assert first.status_code == 500
+    assert (await _vfs(client, api_key, "ls /")).status_code == 200
+
+
+async def test_wall_clock_budget_aborts_the_command(client: AsyncClient, monkeypatch):
+    monkeypatch.setattr("backend.services.vfs_service.MAX_SCRIPT_SECONDS", 0)
+    api_key, _ = await _register(client)
+
+    resp = await _vfs(client, api_key, "ls /")
+
+    assert resp.status_code == 413
+    assert "longer than" in resp.json()["detail"]
+
+
+async def test_wall_clock_budget_cancels_a_slow_nested_request(monkeypatch):
+    """A stalled content route must not hold a VFS slot past the deadline."""
+    monkeypatch.setattr("backend.services.vfs_service.MAX_SCRIPT_SECONDS", 0.05)
+
+    class SlowHttp:
+        async def request(self, *_args, **_kwargs):
+            await asyncio.sleep(10)
+
+    client = vfs_service.InProcessVfsClient(SlowHttp(), asyncio.get_running_loop())
+
+    with pytest.raises(vfs_service.VfsBudgetExceeded, match="longer than"):
+        await asyncio.to_thread(client._request, "GET", "/slow")
+
+
 async def test_machine_fs_404s_without_provisioned_computer(client: AsyncClient, monkeypatch):
     """Browsing must never conjure a VM: a user who never ran a cloud agent
     gets a 404 from the machine fs, not a freshly provisioned sprite."""