From d26a97318953defac364b65ffd8ad3358c04e213 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 24 Jun 2026 13:34:56 +0100 Subject: [PATCH] Harden sandbox sessions and artifact quotas Add runtime governance for the sandbox and artifact store: a per-process concurrent-session cap with least-recently-used eviction of idle sessions and a periodic idle reaper (Celery beat), so sandbox kernels do not accumulate. The session manager performs all backend start/stop outside its lock and tears down the captured handle, so eviction never closes a concurrently re-opened session. Add per-user artifact quotas (count, total bytes, and per-file size) enforced at persistence time as a soft cap, and best-effort cleanup of per-render scratch directories in the sandbox workspace. --- .../agents/tools/artifact_generator.py | 62 +-- .../agents/tools/document_extractor.py | 20 +- application/api/user/tasks.py | 28 ++ application/core/settings.py | 10 + application/sandbox/artifacts_capture.py | 60 ++- application/sandbox/daytona.py | 70 ++++ application/sandbox/jupyter_gateway.py | 22 +- application/sandbox/manager.py | 280 ++++++++++++-- application/sandbox/sandbox_creator.py | 1 + .../storage/db/repositories/artifacts.py | 21 + .../tools/test_document_extractor_unit.py | 42 ++ tests/sandbox/test_daytona.py | 83 ++++ tests/sandbox/test_sandbox_manager.py | 365 +++++++++++++++++- .../storage/db/repositories/test_artifacts.py | 98 +++++ 14 files changed, 1075 insertions(+), 87 deletions(-) diff --git a/application/agents/tools/artifact_generator.py b/application/agents/tools/artifact_generator.py index 2ed15620..451f6f37 100644 --- a/application/agents/tools/artifact_generator.py +++ b/application/agents/tools/artifact_generator.py @@ -19,7 +19,11 @@ from typing import Any, Dict, List, Optional from application.agents.tools.base import Tool from application.core.settings import settings -from application.sandbox.artifacts_capture import append_artifact_version, persist_new_artifact +from application.sandbox.artifacts_capture import ( + QuotaExceeded, + append_artifact_version, + persist_new_artifact, +) from application.sandbox.sandbox_creator import SandboxCreator from application.storage.db.repositories.artifacts import ArtifactsRepository from application.storage.db.session import db_readonly @@ -360,18 +364,21 @@ class ArtifactGeneratorTool(Tool): info = _KIND_INFO[kind] filename = self._filename(title, info["ext"]) - ref = persist_new_artifact( - user_id=self.user_id, - kind=info["kind"], - data=rendered["data"], - filename=filename, - mime_type=info["mime"], - title=title, - conversation_id=self.conversation_id, - workflow_run_id=self.workflow_run_id, - spec=spec, - produced_by=self._produced_by("create_artifact", kind), - ) + try: + ref = persist_new_artifact( + user_id=self.user_id, + kind=info["kind"], + data=rendered["data"], + filename=filename, + mime_type=info["mime"], + title=title, + conversation_id=self.conversation_id, + workflow_run_id=self.workflow_run_id, + spec=spec, + produced_by=self._produced_by("create_artifact", kind), + ) + except QuotaExceeded as exc: + return {"status": "error", "error": str(exc)} if ref is None: return {"status": "error", "error": "failed to persist artifact."} self._last_artifact_id = ref["artifact_id"] @@ -409,15 +416,18 @@ class ArtifactGeneratorTool(Tool): return {"status": "error", "error": rendered["error"]} info = _KIND_INFO[kind] filename = self._filename(None, info["ext"]) - ref = append_artifact_version( - user_id=self.user_id, - artifact_id=artifact_id, - data=rendered["data"], - filename=filename, - mime_type=info["mime"], - spec=spec, - produced_by=self._produced_by(action, kind), - ) + try: + ref = append_artifact_version( + user_id=self.user_id, + artifact_id=artifact_id, + data=rendered["data"], + filename=filename, + mime_type=info["mime"], + spec=spec, + produced_by=self._produced_by(action, kind), + ) + except QuotaExceeded as exc: + return {"status": "error", "error": str(exc)} if ref is None: return {"status": "error", "error": "failed to persist artifact version."} self._last_artifact_id = ref["artifact_id"] @@ -482,8 +492,9 @@ class ArtifactGeneratorTool(Tool): return {"error": "artifact_generator requires a conversation_id or workflow_run_id."} token = uuid.uuid4().hex - spec_path = f"artifacts/{token}/spec.json" - out_path = f"artifacts/{token}/out.{_KIND_INFO[kind]['ext']}" + token_dir = f"artifacts/{token}" + spec_path = f"{token_dir}/spec.json" + out_path = f"{token_dir}/out.{_KIND_INFO[kind]['ext']}" program = _RENDERERS[kind].format(spec_path=spec_path, out_path=out_path) timeout = float(getattr(settings, "SANDBOX_EXEC_TIMEOUT", 60)) @@ -510,6 +521,9 @@ class ArtifactGeneratorTool(Tool): logger.exception("artifact_generator: render failed") return {"error": f"render failed: {type(exc).__name__}: {exc}"} finally: + # Drop this render's scratch dir before tearing the session down so a + # warm/reused session doesn't accumulate per-render files on disk. + manager.remove_path(session_id, token_dir) try: manager.close(session_id) except Exception: diff --git a/application/agents/tools/document_extractor.py b/application/agents/tools/document_extractor.py index 94f76959..4f489e9f 100644 --- a/application/agents/tools/document_extractor.py +++ b/application/agents/tools/document_extractor.py @@ -24,7 +24,7 @@ from application.core.json_schema_utils import ( normalize_json_schema_payload, ) from application.core.settings import settings -from application.sandbox.artifacts_capture import persist_new_artifact +from application.sandbox.artifacts_capture import QuotaExceeded, persist_new_artifact from application.sandbox.sandbox_creator import SandboxCreator from application.storage.db.repositories.artifacts import ArtifactsRepository from application.storage.db.session import db_readonly @@ -212,9 +212,10 @@ class DocumentExtractorTool(Tool): } token = uuid.uuid4().hex - input_path = f"extract/{token}/inputs/{loaded['filename']}" - params_path = f"extract/{token}/params.json" - result_path = f"extract/{token}/result.json" + token_dir = f"extract/{token}" + input_path = f"{token_dir}/inputs/{loaded['filename']}" + params_path = f"{token_dir}/params.json" + result_path = f"{token_dir}/result.json" params = { "input_path": input_path, "markdown_max_bytes": _MARKDOWN_MAX_BYTES, @@ -248,6 +249,9 @@ class DocumentExtractorTool(Tool): logger.exception("document_extractor: extraction failed") return {"status": "error", "error": f"extraction failed: {type(exc).__name__}: {exc}"} finally: + # Drop this extraction's scratch dir (staged input + params + result) + # before close so a warm/reused session doesn't accumulate on disk. + manager.remove_path(session_id, token_dir) try: manager.close(session_id) except Exception: @@ -283,7 +287,13 @@ class DocumentExtractorTool(Tool): compact = self._compact_payload(extracted) payload: Dict[str, Any] = {"status": "ok", "structured": compact} if should_persist: - ref = self._persist(extracted, loaded["title"]) + try: + ref = self._persist(extracted, loaded["title"]) + except QuotaExceeded as exc: + # Extraction itself succeeded; surface the quota error alongside + # the in-context structured result rather than failing the call. + payload["artifact_error"] = str(exc) + ref = None if ref is not None: self._last_artifact_id = ref["artifact_id"] payload["artifact"] = ref diff --git a/application/api/user/tasks.py b/application/api/user/tasks.py index fa079fab..2415bc54 100644 --- a/application/api/user/tasks.py +++ b/application/api/user/tasks.py @@ -1,3 +1,4 @@ +import logging from datetime import timedelta from application.api.user.idempotency import with_idempotency @@ -342,6 +343,25 @@ def cleanup_schedule_runs(self): return {"deleted": deleted, "ttl_days": ttl_days} +@celery.task(bind=True, acks_late=False) +def reap_sandbox_sessions(self): + """Close sandbox sessions idle past their TTL in this worker process. + + The SandboxManager registry is per-process, so this reaps only sessions + bound in THIS worker; the API processes reap their own opportunistically on + ``open``. Artifacts are persisted eagerly, so reaping only closes idle + kernels and never loses a user-facing artifact. + """ + try: + from application.sandbox.sandbox_creator import SandboxCreator + + reaped = SandboxCreator.get_manager().reap_expired() + except Exception: # noqa: BLE001 - housekeeping must never crash the beat loop + logging.getLogger(__name__).exception("reap_sandbox_sessions failed") + return {"reaped": 0, "error": True} + return {"reaped": len(reaped)} + + @celery.on_after_configure.connect def setup_periodic_tasks(sender, **kwargs): from application.core.settings import settings @@ -407,6 +427,14 @@ def setup_periodic_tasks(sender, **kwargs): cleanup_schedule_runs.s(), name="cleanup-schedule-runs", ) + # Close idle-past-TTL sandbox sessions roughly every minute. The on-open + # opportunistic reap still runs in the API processes; this covers worker + # processes (and quiet periods where no new session is opened). + sender.add_periodic_task( + timedelta(seconds=60), + reap_sandbox_sessions.s(), + name="reap-sandbox-sessions", + ) # Bound time limits so a hung OAuth discovery (user never finishes the diff --git a/application/core/settings.py b/application/core/settings.py index 5ebb0cad..a7375cc7 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -333,6 +333,11 @@ class Settings(BaseSettings): SANDBOX_GATEWAY_AUTH_TOKEN: Optional[str] = None # gateway auth token, if set SANDBOX_KERNEL_NAME: str = "python3" # kernelspec name to launch per session SANDBOX_MAX_TTL: int = 1200 # hard cap (s) on agent-selectable keep-alive TTL + # Per-process/worker cap on concurrent live sandbox sessions. Backend-agnostic + # (complements DAYTONA_MAX_SANDBOXES); when reached, an LRU-idle session is + # evicted to make room. This bound is local to each app/worker process. + # 0 (or any non-positive value) disables the cap (unlimited sessions). + SANDBOX_MAX_SESSIONS: int = 32 SANDBOX_EXEC_TIMEOUT: int = 60 # default wall-clock cap (s) per exec call SANDBOX_HTTP_TIMEOUT: int = 10 # fixed cap (s) for REST control calls (create/delete/alive/interrupt) SANDBOX_MAX_OUTPUT_BYTES: int = 8 * 1024 * 1024 # cap on buffered stdout+stderr per exec @@ -354,6 +359,11 @@ class Settings(BaseSettings): DAYTONA_AUTO_STOP_INTERVAL: int = 15 # minutes idle before Daytona auto-stops a sandbox (0 disables) DAYTONA_AUTO_DELETE_INTERVAL: int = 60 # minutes after stop before Daytona auto-deletes (-1 disables) DAYTONA_MAX_SANDBOXES: int = 50 # cap on concurrent live Daytona sandboxes (cost-DoS guard) + # Per-user artifact quotas (generous defaults; enforced at persistence time). + # For all three, 0 (or any non-positive value) disables that quota (unlimited). + ARTIFACT_MAX_BYTES: int = 50 * 1024 * 1024 # cap on a single stored artifact version's bytes + ARTIFACT_MAX_COUNT_PER_USER: int = 5000 # cap on artifacts a user may own + ARTIFACT_MAX_TOTAL_BYTES_PER_USER: int = 5 * 1024 * 1024 * 1024 # cap on a user's total stored bytes @field_validator("POSTGRES_URI", mode="before") @classmethod diff --git a/application/sandbox/artifacts_capture.py b/application/sandbox/artifacts_capture.py index e779278a..f3f27006 100644 --- a/application/sandbox/artifacts_capture.py +++ b/application/sandbox/artifacts_capture.py @@ -16,6 +16,7 @@ from typing import Any, Dict, List, Optional, Tuple from sqlalchemy import text +from application.core.settings import settings from application.storage.db.repositories.artifacts import ArtifactsRepository from application.storage.db.session import db_session from application.storage.storage_creator import StorageCreator @@ -23,6 +24,10 @@ from application.utils import safe_filename logger = logging.getLogger(__name__) + +class QuotaExceeded(Exception): + """Raised when persisting an artifact would breach a per-user size/count quota.""" + # Cap the per-run capture work so a workspace full of pre-existing files can't turn # one exec into an unbounded read+persist sweep. MAX_CAPTURED_FILES = 64 @@ -115,14 +120,20 @@ def capture_artifacts( signature = (len(data), hashlib.sha256(data).hexdigest()) if pre_signatures.get(rel_path) == signature: continue - ref = persist_artifact( - rel_path, - data, - user_id=user_id, - conversation_id=conversation_id, - workflow_run_id=workflow_run_id, - produced_by=produced_by, - ) + try: + ref = persist_artifact( + rel_path, + data, + user_id=user_id, + conversation_id=conversation_id, + workflow_run_id=workflow_run_id, + produced_by=produced_by, + ) + except QuotaExceeded: + # Out of quota: stop capturing the remaining files rather than retry + # each one. Files already captured this run are returned as-is. + logger.warning("artifacts_capture: per-user quota reached; remaining files not captured") + break if ref is not None: captured.append(ref) return captured @@ -184,6 +195,31 @@ def _cleanup_orphan(storage: Any, saved_key: Optional[str]) -> None: logger.exception("artifacts_capture: orphaned-key cleanup failed for %s", saved_key) +def _check_single_artifact_size(size: int) -> None: + """Reject a single artifact version whose byte size exceeds ``ARTIFACT_MAX_BYTES``.""" + max_bytes = int(getattr(settings, "ARTIFACT_MAX_BYTES", 0) or 0) + if max_bytes > 0 and size > max_bytes: + raise QuotaExceeded(f"artifact is too large: {size} bytes exceeds the {max_bytes}-byte per-file cap") + + +def _enforce_user_quota(repo: ArtifactsRepository, user_id: str, added_bytes: int, *, new_artifact: bool) -> None: + """Reject the write when ``user_id`` is at/over their count or total-bytes quota. + + Best-effort SOFT cap (not hard under concurrency): the count/total-bytes reads run + on the same connection as the pending insert, but under READ COMMITTED two concurrent + persists can each read a pre-insert total and both pass, briefly overshooting the cap. + A NEW artifact also consumes one count slot; appending a version only adds bytes to an + existing identity. + """ + _check_single_artifact_size(added_bytes) + max_count = int(getattr(settings, "ARTIFACT_MAX_COUNT_PER_USER", 0) or 0) + if new_artifact and max_count > 0 and repo.count_for_user(user_id) >= max_count: + raise QuotaExceeded(f"artifact count quota reached ({max_count}); delete artifacts to free space") + max_total = int(getattr(settings, "ARTIFACT_MAX_TOTAL_BYTES_PER_USER", 0) or 0) + if max_total > 0 and repo.total_bytes_for_user(user_id) + added_bytes > max_total: + raise QuotaExceeded(f"artifact storage quota reached ({max_total} bytes); delete artifacts to free space") + + def persist_new_artifact( *, user_id: str, @@ -212,6 +248,7 @@ def persist_new_artifact( try: with db_session() as conn: repo = ArtifactsRepository(conn) + _enforce_user_quota(repo, user_id, size, new_artifact=True) artifact = repo.create_artifact( user_id, kind, @@ -232,6 +269,10 @@ def persist_new_artifact( _set_version_storage_path(conn, artifact_id, 1, storage_path) storage.save_file(io.BytesIO(data), storage_path) saved_key = storage_path + except QuotaExceeded: + # Quota check runs before any storage write, so nothing to clean up; + # surface a clean error the caller can render. + raise except Exception: logger.exception("artifacts_capture: failed to persist new artifact") _cleanup_orphan(storage, saved_key) @@ -265,6 +306,7 @@ def append_artifact_version( try: with db_session() as conn: repo = ArtifactsRepository(conn) + _enforce_user_quota(repo, user_id, size, new_artifact=False) version = repo.append_version( artifact_id, mime_type=mime_type, @@ -281,6 +323,8 @@ def append_artifact_version( _set_version_storage_path(conn, artifact_id, version_number, storage_path) storage.save_file(io.BytesIO(data), storage_path) saved_key = storage_path + except QuotaExceeded: + raise except Exception: logger.exception("artifacts_capture: failed to append artifact version") _cleanup_orphan(storage, saved_key) diff --git a/application/sandbox/daytona.py b/application/sandbox/daytona.py index ea95262f..524838a6 100644 --- a/application/sandbox/daytona.py +++ b/application/sandbox/daytona.py @@ -243,11 +243,81 @@ class DaytonaSandbox(CodeSandbox): handle = self._handles.pop(session_id, None) if handle is None: return + self._delete_sandbox(handle) + + def close_handle(self, session_id: str, sandbox_id: str) -> None: + """Delete the SPECIFIC sandbox captured at eviction time, never a re-opened one. + + Used when the manager evicts a session and a concurrent ``open`` of the same id + may have already created a fresh sandbox: this tears down only the sandbox whose + id was captured, leaving any newer registered handle untouched. Falls back to the + registry only when nothing matches the captured id. + """ + with self._lock: + current = self._handles.get(session_id) + if current is not None and current.sandbox_id == sandbox_id: + # Captured handle is still the registered one: pop and delete it. + self._handles.pop(session_id, None) + victim = current + else: + # A concurrent re-open replaced the handle (or it is already gone); do + # NOT touch the registry. Delete the captured sandbox by id directly so + # the freshly created one survives. + victim = None + if victim is not None: + self._delete_sandbox(victim) + else: + self._delete_sandbox_by_id(sandbox_id) + + def _delete_sandbox(self, handle: "_Handle") -> None: + """Best-effort delete of the sandbox behind ``handle`` (never raises).""" try: self._client.delete(handle.sandbox) except Exception as exc: # noqa: BLE001 - teardown is best-effort, never raise logger.warning("Failed to delete Daytona sandbox %s: %s", handle.sandbox_id, exc) + def _delete_sandbox_by_id(self, sandbox_id: str) -> None: + """Best-effort delete a sandbox by its id when no live handle is held (never raises). + + ``client.delete`` needs a ``Sandbox`` object, so fetch it by id first; this is the + evict-then-concurrent-reopen path where only the captured id is still known. + """ + try: + sandbox = self._client.get(sandbox_id) + self._client.delete(sandbox) + except Exception as exc: # noqa: BLE001 - teardown is best-effort, never raise + logger.warning("Failed to delete Daytona sandbox by id %s: %s", sandbox_id, exc) + + def remove_path(self, session_id: str, path: str) -> None: + """Best-effort delete a workspace-relative path (per-render scratch dir); never raises. + + The tools pass a token DIRECTORY (e.g. ``artifacts/{token}``), so the delete is + recursive — ``fs.delete_file`` with ``recursive=True`` removes a folder and its + contents (daytona==0.190.1). + """ + try: + handle = self._get_handle(session_id) + remote = self._remote_path(handle.workspace, path) + if remote == posixpath.normpath(handle.workspace): + return # refuse to delete the workspace root itself + self._delete_remote(handle, remote) + except Exception as exc: # noqa: BLE001 - cleanup is best-effort, never raise + logger.debug("remove_path best-effort delete returned for %r: %s", path, exc) + + def _delete_remote(self, handle: "_Handle", remote: str) -> None: + """Recursively delete ``remote`` inside the sandbox, with a contained-command fallback.""" + try: + handle.sandbox.fs.delete_file(remote, recursive=True) + return + except TypeError: + # Older SDK without the recursive kwarg: fall back to a contained rm -rf. + pass + try: + handle.sandbox.fs.delete_file(remote) + except Exception: # noqa: BLE001 - directory delete may need a recursive remove + quoted = "'" + remote.replace("'", "'\\''") + "'" + handle.sandbox.process.exec(f"rm -rf {quoted}", timeout=int(self._default_timeout)) + def _prime(self, handle: _Handle) -> None: """Create the per-session workspace directory inside the sandbox.""" try: diff --git a/application/sandbox/jupyter_gateway.py b/application/sandbox/jupyter_gateway.py index b9960986..1cec1d3d 100644 --- a/application/sandbox/jupyter_gateway.py +++ b/application/sandbox/jupyter_gateway.py @@ -155,14 +155,32 @@ class JupyterKernelGatewaySandbox(CodeSandbox): kernel = self._kernels.pop(session_id, None) if kernel is None: return + self._delete_kernel(kernel.kernel_id) + + def close_handle(self, session_id: str, kernel_id: str) -> None: + """Delete the SPECIFIC kernel captured at eviction time, never a re-opened one. + + When the manager evicts a session, a concurrent ``open`` of the same id may have + already started a fresh kernel; this deletes only the kernel whose id was captured + and pops the registry entry only when it still points at that same kernel, so the + new kernel survives. + """ + with self._lock: + current = self._kernels.get(session_id) + if current is not None and current.kernel_id == kernel_id: + self._kernels.pop(session_id, None) + self._delete_kernel(kernel_id) + + def _delete_kernel(self, kernel_id: str) -> None: + """Best-effort DELETE of a gateway kernel by id (teardown never raises).""" try: requests.delete( - f"{self._base_url}/api/kernels/{kernel.kernel_id}", + f"{self._base_url}/api/kernels/{kernel_id}", headers=self._headers(), timeout=self._http_timeout, ) except requests.RequestException as exc: # teardown is best-effort - logger.warning("Failed to delete kernel %s: %s", kernel.kernel_id, exc) + logger.warning("Failed to delete kernel %s: %s", kernel_id, exc) def _kernel_alive(self, kernel_id: str) -> bool: try: diff --git a/application/sandbox/manager.py b/application/sandbox/manager.py index 183a08a6..edd4c623 100644 --- a/application/sandbox/manager.py +++ b/application/sandbox/manager.py @@ -1,21 +1,37 @@ """Process-wide registry binding session ids to sandbox handles with idle expiry.""" +import logging import threading import time from dataclasses import dataclass, field -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Tuple from application.sandbox.base import CodeSandbox, ExecResult +logger = logging.getLogger(__name__) + + +class SandboxCapacityError(RuntimeError): + """Raised when the concurrent-session cap is reached and no idle session can be freed.""" + @dataclass class _Session: - """Bookkeeping for one bound sandbox session: its TTL and access timestamps.""" + """Bookkeeping for one bound sandbox session: its TTL, access timestamps, and backend handle. + + ``handle`` is the backend handle id returned by ``backend.open``; ``None`` while a + slot is RESERVED (a placeholder occupying a cap slot during a cold backend open that + runs outside the lock). ``ready`` is False for such a placeholder so reuse/reap/evict + skip it until the backend open finalizes it. + """ session_id: str ttl: float created_at: float last_access: float = field(default=0.0) + in_use: int = field(default=0) + handle: Optional[str] = field(default=None) + ready: bool = field(default=False) def is_expired(self, now: float) -> bool: """True when the session has been idle longer than its (clamped) TTL.""" @@ -23,13 +39,25 @@ class _Session: class SandboxManager: - """Binds session ids to a shared backend, clamps TTLs, and reaps idle sessions.""" + """Binds session ids to a shared backend, clamps TTLs, caps concurrency, and reaps idle sessions.""" - def __init__(self, backend: CodeSandbox, max_ttl: float, default_ttl: Optional[float] = None) -> None: - """Wrap ``backend`` with a registry clamped to ``max_ttl`` seconds per session.""" + def __init__( + self, + backend: CodeSandbox, + max_ttl: float, + default_ttl: Optional[float] = None, + max_sessions: Optional[int] = None, + ) -> None: + """Wrap ``backend`` with a registry clamped to ``max_ttl`` and bounded to ``max_sessions``. + + The session cap is per-process/worker: each app or Celery process keeps its + own registry, so the effective fleet-wide ceiling is ``max_sessions`` times + the number of live processes. + """ self._backend = backend self._max_ttl = max_ttl self._default_ttl = default_ttl if default_ttl is not None else max_ttl + self._max_sessions = max_sessions if max_sessions and max_sessions > 0 else None self._sessions: Dict[str, _Session] = {} self._lock = threading.Lock() @@ -42,94 +70,264 @@ class SandboxManager: return min(ttl, self._max_ttl) def open(self, session_id: str, ttl: Optional[float] = None) -> str: - """Open (or reuse) the sandbox for ``session_id`` with a clamped TTL.""" - self.reap_expired() + """Open (or reuse) the sandbox for ``session_id`` with a clamped TTL. + + The lock guards ONLY the in-memory registry and refcounts; the (potentially + ~60s) cold backend open runs WITHOUT the lock held so it can never serialize + other lock-taking methods. Flow: + + 1. Under the lock: if the session already exists and is ready, refresh it and + return its cached handle (reuse never does backend I/O). Otherwise reap + idle-expired sessions, evict an LRU-idle victim if at capacity, then RESERVE + a placeholder slot for ``session_id`` so concurrent opens can't overshoot + the cap. + 2. Outside the lock: call ``backend.open`` (cold start) and close any reaped / + evicted victims' captured backend resources. + 3. Under the lock: finalize the placeholder into a ready session. On failure, + free the reserved slot and re-raise. + """ with self._lock: session = self._sessions.get(session_id) now = time.monotonic() - if session is not None: + if session is not None and session.ready: session.last_access = now - handle = self._backend.attach(session_id) - return handle - handle = self._backend.open(session_id) + return session.handle + # A placeholder for this id means another thread is mid cold-open; refresh + # it and re-reserve below. Both threads call the (idempotent) backend.open, + # which dedupes by session id, so no cap overshoot and no orphaned runtime. + reaped = self._reap_locked(now) + evicted = self._make_room_locked() self._sessions[session_id] = _Session( session_id=session_id, ttl=self._clamp_ttl(ttl), created_at=now, last_access=now, ) - return handle + + # Cold backend open and victim teardown run OUTSIDE the lock. + for sid, handle in reaped: + self._close_backend(sid, handle) + if evicted is not None: + self._close_backend(evicted[0], evicted[1]) + try: + handle = self._backend.open(session_id) + except Exception: + # Free the reserved slot so a failed cold open never leaks a cap slot. + with self._lock: + placeholder = self._sessions.get(session_id) + if placeholder is not None and not placeholder.ready: + self._sessions.pop(session_id, None) + self._close_backend(session_id, None) # best-effort: tear down anything created + raise + + with self._lock: + session = self._sessions.get(session_id) + if session is None: + # Slot was reclaimed (e.g. closed) while we opened; drop the new runtime. + stale = handle + else: + session.handle = handle + session.ready = True + session.last_access = time.monotonic() + stale = None + if stale is not None: + self._close_backend(session_id, stale) + return handle + + def _make_room_locked(self) -> Optional[Tuple[str, Optional[str]]]: + """Evict the LRU-idle ready session when at capacity; return (id, handle) to close, else None. + + Caller holds the lock. The victim is removed from the registry here so its slot + is freed for the reserving caller; its backend ``close`` is deferred to outside + the lock (keyed by the captured handle so a concurrent re-open of the same id + is never torn down). Raises ``SandboxCapacityError`` when the registry is full + and every session is busy or still opening. + """ + if self._max_sessions is None or len(self._sessions) < self._max_sessions: + return None + idle = [s for s in self._sessions.values() if s.ready and s.in_use == 0] + if not idle: + raise SandboxCapacityError( + f"sandbox session cap reached ({self._max_sessions} live, all busy); cannot open another" + ) + victim = min(idle, key=lambda s: s.last_access) + self._sessions.pop(victim.session_id, None) + logger.info("SandboxManager evicting LRU-idle session %s to honor cap", victim.session_id) + return (victim.session_id, victim.handle) def attach(self, session_id: str) -> str: - """Reattach to an existing session, refreshing its idle clock.""" + """Reattach to an existing session, refreshing its idle clock. + + Returns the cached handle without backend I/O for a ready session; falls back + to ``backend.attach`` (outside the lock) only when no handle is cached yet. + """ with self._lock: session = self._sessions.get(session_id) if session is None: raise KeyError(f"No sandbox session bound for {session_id!r}") session.last_access = time.monotonic() + cached = session.handle + if cached is not None: + return cached return self._backend.attach(session_id) def exec(self, session_id: str, code: str, timeout: Optional[float] = None) -> ExecResult: - """Execute ``code`` in the bound session, touching its idle clock first.""" - with self._lock: - session = self._sessions.get(session_id) - if session is None: - raise KeyError(f"No sandbox session bound for {session_id!r}") - session.last_access = time.monotonic() - return self._backend.exec(session_id, code, timeout) + """Execute ``code`` in the bound session, holding it in-use so a reap/evict can't pull it.""" + self._enter(session_id) + try: + return self._backend.exec(session_id, code, timeout) + finally: + self._leave(session_id) def put_file(self, session_id: str, dest_path: str, data: bytes) -> None: """Write ``data`` into the bound session's workspace.""" - self._touch(session_id) - self._backend.put_file(session_id, dest_path, data) + self._enter(session_id) + try: + self._backend.put_file(session_id, dest_path, data) + finally: + self._leave(session_id) def get_file(self, session_id: str, path: str) -> bytes: """Read ``path`` from the bound session's workspace.""" - self._touch(session_id) - return self._backend.get_file(session_id, path) + self._enter(session_id) + try: + return self._backend.get_file(session_id, path) + finally: + self._leave(session_id) def list_files(self, session_id: str) -> List[str]: """List files in the bound session's workspace.""" - self._touch(session_id) - return self._backend.list_files(session_id) + self._enter(session_id) + try: + return self._backend.list_files(session_id) + finally: + self._leave(session_id) + + def remove_path(self, session_id: str, path: str) -> None: + """Best-effort delete a workspace-relative path; never raises (cleanup must not fail an op).""" + try: + self._enter(session_id) + except KeyError: + return + try: + remover = getattr(self._backend, "remove_path", None) + if callable(remover): + remover(session_id, path) + else: + self._remove_via_exec(session_id, path) + except Exception: + logger.exception("SandboxManager: best-effort remove_path failed for %r", path) + finally: + self._leave(session_id) + + def _remove_via_exec(self, session_id: str, path: str) -> None: + """Fallback workspace cleanup: run a contained shutil.rmtree of the relative path.""" + self._backend.exec(session_id, self._build_remove_program(path), None) + + @staticmethod + def _build_remove_program(path: str) -> str: + """Build the contained shutil.rmtree program for ``path`` with a workspace-root guard.""" + # ``path`` is a server-controlled token dir literal; it is passed as a + # repr so its contents are never interpreted as code, and the backend + # already chdirs into the per-session workspace before running. The + # guard rejects absolute, traversal, and empty/'.'/'./' paths so a caller + # can never rmtree the whole workspace root. + return ( + "import os, shutil\n" + f"_p = {path!r}\n" + "if _p and not _p.startswith('/') and '..' not in _p.split('/') and os.path.normpath(_p) != '.':\n" + " shutil.rmtree(_p, ignore_errors=True)\n" + ) def close(self, session_id: str) -> None: - """Tear down the backend runtime and drop the session from the registry.""" + """Tear down the backend runtime and drop the session from the registry. + + The backend ``close`` runs OUTSIDE the lock (it may be slow network I/O) and is + keyed by the captured handle so it tears down only the resource this session + owned, never one a concurrent re-open created. + """ with self._lock: - self._sessions.pop(session_id, None) - self._backend.close(session_id) + session = self._sessions.pop(session_id, None) + handle = session.handle if session is not None else None + self._close_backend(session_id, handle) def has_session(self, session_id: str) -> bool: - """True when ``session_id`` is currently bound in the registry.""" + """True when ``session_id`` is currently bound (ready or opening) in the registry.""" with self._lock: return session_id in self._sessions + def session_count(self) -> int: + """Return the number of sessions currently bound (ready or reserved) in the registry.""" + with self._lock: + return len(self._sessions) + def ttl_for(self, session_id: str) -> Optional[float]: """Return the clamped TTL bound to ``session_id``, or None if unbound.""" with self._lock: session = self._sessions.get(session_id) return session.ttl if session else None - def _touch(self, session_id: str) -> None: + def _enter(self, session_id: str) -> None: + """Touch the idle clock and mark the session in-use so a concurrent reap/evict skips it.""" with self._lock: session = self._sessions.get(session_id) - if session is None: + if session is None or not session.ready: raise KeyError(f"No sandbox session bound for {session_id!r}") session.last_access = time.monotonic() + session.in_use += 1 + + def _leave(self, session_id: str) -> None: + """Release an in-use hold taken by ``_enter`` (idempotent if the session was closed).""" + with self._lock: + session = self._sessions.get(session_id) + if session is not None and session.in_use > 0: + session.in_use -= 1 + session.last_access = time.monotonic() + + def _close_backend(self, session_id: str, handle: Optional[str]) -> None: + """Close the SPECIFIC backend resource captured for this session, best-effort. + + When the backend exposes ``close_handle``, the captured handle id is passed so a + concurrent re-open of the same ``session_id`` (which created a new backend handle) + is never the one torn down. Backends without that hook fall back to ``close`` by + id, which is correct for callers (``close``/reap/failed-open) where no concurrent + re-open of a still-registered session can be in flight. + """ + try: + closer = getattr(self._backend, "close_handle", None) + if handle is not None and callable(closer): + closer(session_id, handle) + else: + self._backend.close(session_id) + except Exception: + logger.exception("SandboxManager: backend close failed for %s", session_id) def reap_expired(self) -> List[str]: """Close sessions idle past their TTL and return the reaped session ids. - Extension point: a later hardening slice runs this on a background - reaper and, before close, flushes workspace files to the artifact - store (persist-on-reap) so a fresh kernel can re-mount them by - reference on next access. For now it is a plain idle close. + Artifacts are persisted eagerly by the tools/code node right after each + exec, so a session's workspace is scratch: reaping only closes the kernel + and never loses a user-facing artifact. Busy (in-use) sessions are left + alone so a reap can't pull a workspace out from under a running exec. """ now = time.monotonic() with self._lock: - expired = [sid for sid, s in self._sessions.items() if s.is_expired(now)] - for sid in expired: - self._sessions.pop(sid, None) - for sid in expired: - self._backend.close(sid) + expired = self._reap_locked(now) + for sid, handle in expired: + self._close_backend(sid, handle) + return [sid for sid, _ in expired] + + def _reap_locked(self, now: float) -> List[Tuple[str, Optional[str]]]: + """Pop idle-expired ready sessions from the registry; return (id, handle) pairs to close. + + Caller holds the lock. Placeholders (not yet ready) and busy sessions are left + alone. The backend close of each popped session is deferred to outside the lock. + """ + expired = [ + (sid, s.handle) + for sid, s in self._sessions.items() + if s.ready and s.in_use == 0 and s.is_expired(now) + ] + for sid, _ in expired: + self._sessions.pop(sid, None) return expired diff --git a/application/sandbox/sandbox_creator.py b/application/sandbox/sandbox_creator.py index b3917c71..425a9e62 100644 --- a/application/sandbox/sandbox_creator.py +++ b/application/sandbox/sandbox_creator.py @@ -64,6 +64,7 @@ class SandboxCreator: cls._instance = SandboxManager( backend=backend, max_ttl=float(settings.SANDBOX_MAX_TTL), + max_sessions=int(settings.SANDBOX_MAX_SESSIONS), ) return cls._instance diff --git a/application/storage/db/repositories/artifacts.py b/application/storage/db/repositories/artifacts.py index 86c81863..49ec3392 100644 --- a/application/storage/db/repositories/artifacts.py +++ b/application/storage/db/repositories/artifacts.py @@ -163,6 +163,27 @@ class ArtifactsRepository: ) return [_artifact_to_dict(r) for r in result.fetchall()] + def count_for_user(self, user_id: str) -> int: + """Return how many artifacts ``user_id`` currently owns (quota accounting).""" + row = self._conn.execute( + text("SELECT count(*) FROM artifacts WHERE user_id = :user_id"), + {"user_id": user_id}, + ).fetchone() + return int(row[0]) if row is not None else 0 + + def total_bytes_for_user(self, user_id: str) -> int: + """Return the summed byte size of every version a user owns (quota accounting).""" + row = self._conn.execute( + text( + "SELECT COALESCE(SUM(v.size), 0) " + "FROM artifact_versions v " + "JOIN artifacts a ON a.id = v.artifact_id " + "WHERE a.user_id = :user_id" + ), + {"user_id": user_id}, + ).fetchone() + return int(row[0]) if row is not None else 0 + def get_version(self, artifact_id: str, version: int) -> Optional[dict]: """Fetch a single version row by artifact id and version number.""" result = self._conn.execute( diff --git a/tests/agents/tools/test_document_extractor_unit.py b/tests/agents/tools/test_document_extractor_unit.py index 60f84dba..e8524c1c 100644 --- a/tests/agents/tools/test_document_extractor_unit.py +++ b/tests/agents/tools/test_document_extractor_unit.py @@ -53,11 +53,15 @@ class _FakeManager: self.exec_programs: List[str] = [] self.opened: List[Any] = [] self.closed: List[str] = [] + self.removed: List[str] = [] def open(self, session_id, ttl=None): self.opened.append((session_id, ttl)) return session_id + def remove_path(self, session_id, path): + self.removed.append(path) + def put_file(self, session_id, dest_path, data): self.put_files[dest_path] = data @@ -203,6 +207,24 @@ def test_extract_stages_input_and_shapes_compact_payload(monkeypatch): assert manager.closed == ["conv-1"] +def test_extract_removes_its_scratch_token_dir(monkeypatch): + _stub_repo(monkeypatch, found=True, conv="conv-1", run=None) + _patch_storage(monkeypatch) + _patch_no_persist(monkeypatch) + manager = _manager_with_extract(_CANNED_EXTRACT) + _patch_manager(monkeypatch, manager) + + _tool().execute_action("extract_document", input="art-1", persist=False) + + # The per-render token dir (extract/) is removed before/at teardown so + # a reused session doesn't accumulate staged inputs + results on disk. + assert manager.removed, "extractor should remove its scratch token dir" + removed_dir = manager.removed[0] + assert removed_dir.startswith("extract/") + # The dir prefix must cover the files that were staged under it. + assert all(p.startswith(removed_dir + "/") for p in manager.put_files) + + def test_extract_persists_data_artifact_by_reference(monkeypatch): _stub_repo(monkeypatch, found=True, conv="conv-1", run=None) _patch_storage(monkeypatch) @@ -229,6 +251,26 @@ def test_extract_persists_data_artifact_by_reference(monkeypatch): assert json.loads(captured["data"].decode("utf-8")) == _CANNED_EXTRACT +def test_quota_exceeded_on_persist_surfaces_cleanly(monkeypatch): + from application.sandbox.artifacts_capture import QuotaExceeded + + _stub_repo(monkeypatch, found=True, conv="conv-1", run=None) + _patch_storage(monkeypatch) + manager = _manager_with_extract(_CANNED_EXTRACT) + _patch_manager(monkeypatch, manager) + + def _quota_blocked(**kwargs): + raise QuotaExceeded("artifact storage quota reached") + + monkeypatch.setattr(de, "persist_new_artifact", _quota_blocked) + + out = _tool().execute_action("extract_document", input="art-1") + # Extraction still succeeds; the quota failure surfaces as a non-fatal note. + assert out["status"] == "ok" + assert "artifact" not in out + assert "quota" in out["artifact_error"].lower() + + # --------------------------------------------------------------------------- # Cross-tenant scope # --------------------------------------------------------------------------- diff --git a/tests/sandbox/test_daytona.py b/tests/sandbox/test_daytona.py index edc9cbe5..a4675523 100644 --- a/tests/sandbox/test_daytona.py +++ b/tests/sandbox/test_daytona.py @@ -47,6 +47,7 @@ class _FakeFS: self.create_folder = mock.Mock() self.upload_file = mock.Mock() self.download_file = mock.Mock(return_value=b"payload") + self.delete_file = mock.Mock() self.get_file_info = mock.Mock(return_value=_FakeFileInfo("a.txt", size=7)) # Path-aware tree: maps an absolute dir to the entries it contains. The # real Daytona fs.list_files is single-level and returns basenames, so the @@ -81,6 +82,16 @@ class _FakeDaytonaClient: self.delete = mock.Mock(side_effect=self._delete) self.list = mock.Mock(side_effect=self._list) self.start = mock.Mock(side_effect=self._start) + self.get = mock.Mock(side_effect=self._get) + + def _get(self, sandbox_id_or_name): + for _, sandbox in self.created: + if sandbox.id == sandbox_id_or_name: + return sandbox + for sandbox in self.existing: + if sandbox.id == sandbox_id_or_name: + return sandbox + raise KeyError(sandbox_id_or_name) def _create(self, params=None, timeout=60): labels = getattr(params, "labels", None) or {} @@ -218,6 +229,78 @@ def test_close_swallows_delete_errors(sandbox): assert not sandbox._handles +def test_remove_path_deletes_directory_recursively(sandbox): + """A token DIRECTORY is removed with a recursive delete (daytona delete_file recursive=True).""" + sandbox.open("conv-1") + _, created = sandbox._client.created[0] + sandbox.remove_path("conv-1", "artifacts/tok") + remote, kwargs = created.fs.delete_file.call_args.args, created.fs.delete_file.call_args.kwargs + assert remote[0].endswith("/artifacts/tok") + assert kwargs.get("recursive") is True + + +def test_remove_path_refuses_workspace_root(sandbox): + """remove_path must never delete the workspace root itself.""" + sandbox.open("conv-1") + _, created = sandbox._client.created[0] + sandbox.remove_path("conv-1", ".") + created.fs.delete_file.assert_not_called() + + +def test_remove_path_falls_back_to_rm_when_no_recursive_kwarg(sandbox): + """When delete_file lacks the recursive kwarg, fall back to a contained rm -rf.""" + sandbox.open("conv-1") + _, created = sandbox._client.created[0] + + def _no_recursive(path, recursive=None): + if recursive: + raise TypeError("delete_file() got an unexpected keyword argument 'recursive'") + raise RuntimeError("cannot delete a directory") + + created.fs.delete_file.side_effect = _no_recursive + created.process.exec = mock.Mock() + sandbox.remove_path("conv-1", "artifacts/tok") + cmd = created.process.exec.call_args.args[0] + assert cmd.startswith("rm -rf ") and cmd.endswith("/artifacts/tok'") + + +def test_remove_path_never_raises(sandbox): + """Any backend error during cleanup is swallowed (cleanup must not fail an op).""" + sandbox.open("conv-1") + _, created = sandbox._client.created[0] + created.fs.delete_file.side_effect = RuntimeError("cloud down") + created.process.exec = mock.Mock(side_effect=RuntimeError("also down")) + sandbox.remove_path("conv-1", "artifacts/tok") # must not raise + + +def test_close_handle_deletes_captured_sandbox(sandbox): + """close_handle tears down the captured sandbox and drops the registry entry.""" + sandbox.open("conv-1") + _, created = sandbox._client.created[0] + sandbox.close_handle("conv-1", created.id) + sandbox._client.delete.assert_called_once_with(created) + assert not sandbox._handles + + +def test_close_handle_leaves_reopened_sandbox_intact(sandbox): + """A concurrent re-open replaced the handle: close_handle deletes the OLD one by id only.""" + sandbox.open("conv-1") + _, old = sandbox._client.created[0] + # Simulate a concurrent re-open: a new sandbox is registered for the same session. + new = _FakeSandbox(sandbox_id="sbx-new", labels={"docsgpt_session_id": "conv-1"}) + from application.sandbox.daytona import _Handle, _WORKSPACE_ROOT + + sandbox._handles["conv-1"] = _Handle(new, new.id, _WORKSPACE_ROOT) + sandbox._client.created.append((None, new)) # so client.get(new.id) could resolve + + sandbox.close_handle("conv-1", old.id) # close the OLD captured id + + # The OLD sandbox was deleted (fetched by id); the NEW handle stays registered. + assert old in sandbox._client.deleted + assert new not in sandbox._client.deleted + assert sandbox._handles["conv-1"].sandbox_id == "sbx-new" + + def test_invalid_session_id_rejected(sandbox): with pytest.raises(ValueError): sandbox.open("../etc/passwd") diff --git a/tests/sandbox/test_sandbox_manager.py b/tests/sandbox/test_sandbox_manager.py index c6736cbb..577227d9 100644 --- a/tests/sandbox/test_sandbox_manager.py +++ b/tests/sandbox/test_sandbox_manager.py @@ -1,11 +1,12 @@ """Unit tests for SandboxManager and SandboxCreator using an in-memory backend.""" +import threading from typing import Dict, List import pytest from application.sandbox.base import CodeSandbox, ExecResult -from application.sandbox.manager import SandboxManager +from application.sandbox.manager import SandboxCapacityError, SandboxManager class FakeBackend(CodeSandbox): @@ -15,12 +16,16 @@ class FakeBackend(CodeSandbox): self.open_calls: List[str] = [] self.attach_calls: List[str] = [] self.closed: List[str] = [] + self.closed_handles: List[tuple] = [] + self.removed: List[tuple] = [] self.files: Dict[str, Dict[str, bytes]] = {} self._handles: Dict[str, str] = {} def open(self, session_id: str) -> str: self.open_calls.append(session_id) self.files.setdefault(session_id, {}) + # Distinct handle id per open call so an evict + concurrent re-open of the + # same id produce different handles (the test asserts the old one is closed). handle = f"handle-{session_id}-{len(self.open_calls)}" self._handles[session_id] = handle return handle @@ -36,6 +41,14 @@ class FakeBackend(CodeSandbox): self.files.pop(session_id, None) self._handles.pop(session_id, None) + def close_handle(self, session_id: str, handle: str) -> None: + # Close the SPECIFIC captured handle; only drop the registry entry when it + # still points at that handle (a concurrent re-open may have replaced it). + self.closed_handles.append((session_id, handle)) + if self._handles.get(session_id) == handle: + self._handles.pop(session_id, None) + self.files.pop(session_id, None) + def exec(self, session_id, code, timeout=None) -> ExecResult: return ExecResult(status="ok", stdout=f"ran:{code}") @@ -48,6 +61,17 @@ class FakeBackend(CodeSandbox): def list_files(self, session_id) -> List[str]: return list(self.files.get(session_id, {}).keys()) + def remove_path(self, session_id, path) -> None: + self.removed.append((session_id, path)) + files = self.files.get(session_id, {}) + for key in [k for k in files if k == path or k.startswith(path + "/")]: + files.pop(key, None) + + @property + def torn_down(self) -> List[str]: + """Session ids torn down via either close path, in teardown order.""" + return self.closed + [sid for sid, _ in self.closed_handles] + @pytest.fixture() def backend() -> FakeBackend: @@ -62,12 +86,15 @@ def test_open_registers_session_and_opens_backend(backend): assert mgr.has_session("conv-1") -def test_open_twice_reuses_via_attach(backend): +def test_open_twice_reuses_cached_handle_without_backend_io(backend): mgr = SandboxManager(backend, max_ttl=600) - mgr.open("conv-1") - mgr.open("conv-1") + h1 = mgr.open("conv-1") + h2 = mgr.open("conv-1") assert backend.open_calls == ["conv-1"] # opened once - assert backend.attach_calls == ["conv-1"] # second open reused via attach + # Reuse returns the cached handle WITHOUT any backend I/O (no second open/attach): + # backend calls must never run while the manager lock is held. + assert backend.attach_calls == [] + assert h1 == h2 def test_attach_reuse_requires_existing_session(backend): @@ -75,7 +102,9 @@ def test_attach_reuse_requires_existing_session(backend): with pytest.raises(KeyError): mgr.attach("missing") mgr.open("conv-1") + # attach returns the cached handle without backend I/O. assert mgr.attach("conv-1").startswith("handle-conv-1") + assert backend.attach_calls == [] def test_ttl_clamped_to_max(backend): @@ -103,7 +132,7 @@ def test_close_drops_registry_and_backend(backend): mgr.open("conv-1") mgr.close("conv-1") assert not mgr.has_session("conv-1") - assert backend.closed == ["conv-1"] + assert backend.torn_down == ["conv-1"] def test_reap_expired_closes_idle_sessions(backend, monkeypatch): @@ -114,7 +143,7 @@ def test_reap_expired_closes_idle_sessions(backend, monkeypatch): clock["t"] = 1051.0 # 51s idle > 50s ttl reaped = mgr.reap_expired() assert reaped == ["conv-1"] - assert backend.closed == ["conv-1"] + assert backend.torn_down == ["conv-1"] assert not mgr.has_session("conv-1") @@ -159,3 +188,325 @@ def test_sandbox_creator_manager_is_singleton(): m2 = SandboxCreator.get_manager() assert m1 is m2 SandboxCreator.reset() + + +# --------------------------------------------------------------------------- +# Concurrent-session cap +# --------------------------------------------------------------------------- + + +def test_open_under_cap_does_not_evict(backend): + mgr = SandboxManager(backend, max_ttl=600, max_sessions=3) + mgr.open("a") + mgr.open("b") + assert backend.torn_down == [] + assert mgr.session_count() == 2 + + +def test_cap_evicts_lru_idle_session(backend, monkeypatch): + clock = {"t": 1000.0} + monkeypatch.setattr("application.sandbox.manager.time.monotonic", lambda: clock["t"]) + mgr = SandboxManager(backend, max_ttl=600, max_sessions=2) + mgr.open("a") # last_access 1000 + clock["t"] = 1001.0 + mgr.open("b") # last_access 1001 + clock["t"] = 1002.0 + mgr.open("c") # at cap -> evict LRU-idle = "a" + assert backend.torn_down == ["a"] + assert not mgr.has_session("a") + assert mgr.has_session("b") and mgr.has_session("c") + assert mgr.session_count() == 2 + + +def test_cap_eviction_picks_least_recently_used(backend, monkeypatch): + clock = {"t": 1000.0} + monkeypatch.setattr("application.sandbox.manager.time.monotonic", lambda: clock["t"]) + mgr = SandboxManager(backend, max_ttl=600, max_sessions=2) + mgr.open("a") + clock["t"] = 1001.0 + mgr.open("b") + clock["t"] = 1002.0 + mgr.exec("a", "noop") # refresh "a" so "b" is now the LRU + clock["t"] = 1003.0 + mgr.open("c") + assert backend.torn_down == ["b"] + assert mgr.has_session("a") and mgr.has_session("c") + + +def test_cap_rejects_when_all_sessions_busy(backend): + mgr = SandboxManager(backend, max_ttl=600, max_sessions=1) + mgr.open("a") + # Hold "a" in-use; a concurrent open of "b" then cannot free a slot. + mgr._enter("a") + try: + with pytest.raises(SandboxCapacityError): + mgr.open("b") + finally: + mgr._leave("a") + # Once released, opening succeeds (evicting the now-idle "a"). + mgr.open("b") + assert mgr.has_session("b") + + +def test_reuse_existing_session_never_evicts_at_cap(backend): + mgr = SandboxManager(backend, max_ttl=600, max_sessions=1) + mgr.open("a") + mgr.open("a") # reuse, not a new session -> no eviction + assert backend.torn_down == [] + assert mgr.session_count() == 1 + + +# --------------------------------------------------------------------------- +# Idle reaper +# --------------------------------------------------------------------------- + + +def test_reap_closes_idle_past_ttl_and_keeps_fresh(backend, monkeypatch): + clock = {"t": 1000.0} + monkeypatch.setattr("application.sandbox.manager.time.monotonic", lambda: clock["t"]) + mgr = SandboxManager(backend, max_ttl=600) + mgr.open("stale", ttl=50) + clock["t"] = 1040.0 + mgr.open("fresh", ttl=50) # last_access 1040 + clock["t"] = 1051.0 # stale idle 51s > 50; fresh idle 11s < 50 + reaped = mgr.reap_expired() + assert reaped == ["stale"] + assert backend.torn_down == ["stale"] + assert mgr.has_session("fresh") + + +def test_reap_leaves_busy_session_even_if_expired(backend, monkeypatch): + clock = {"t": 1000.0} + monkeypatch.setattr("application.sandbox.manager.time.monotonic", lambda: clock["t"]) + mgr = SandboxManager(backend, max_ttl=600) + mgr.open("busy", ttl=10) + mgr._enter("busy") # mark in-use (e.g. a long exec in flight) + clock["t"] = 1100.0 # well past TTL + try: + assert mgr.reap_expired() == [] + assert mgr.has_session("busy") + assert backend.torn_down == [] + finally: + mgr._leave("busy") + + +# --------------------------------------------------------------------------- +# Workspace cleanup +# --------------------------------------------------------------------------- + + +def test_remove_path_delegates_to_backend(backend): + mgr = SandboxManager(backend, max_ttl=600) + mgr.open("a") + mgr.put_file("a", "artifacts/tok/out.pptx", b"x") + mgr.remove_path("a", "artifacts/tok") + assert backend.removed == [("a", "artifacts/tok")] + assert mgr.list_files("a") == [] + + +def test_remove_path_unknown_session_is_silent(backend): + mgr = SandboxManager(backend, max_ttl=600) + mgr.remove_path("missing", "artifacts/tok") # must not raise + assert backend.removed == [] + + +def test_remove_path_via_exec_when_backend_lacks_helper(): + class _NoRemoveBackend(FakeBackend): + remove_path = None # type: ignore[assignment] + + def __init__(self): + super().__init__() + self.exec_programs: List[str] = [] + + def exec(self, session_id, code, timeout=None): + self.exec_programs.append(code) + return ExecResult(status="ok") + + backend = _NoRemoveBackend() + mgr = SandboxManager(backend, max_ttl=600) + mgr.open("a") + mgr.remove_path("a", "artifacts/tok") + assert any("shutil.rmtree" in code for code in backend.exec_programs) + + +def test_remove_via_exec_program_guards_workspace_root(): + """The exec-fallback program must refuse to rmtree the workspace root ('', '.', './').""" + + class _NoRemoveBackend(FakeBackend): + remove_path = None # type: ignore[assignment] + + def __init__(self): + super().__init__() + self.exec_programs: List[str] = [] + + def exec(self, session_id, code, timeout=None): + self.exec_programs.append(code) + return ExecResult(status="ok") + + backend = _NoRemoveBackend() + mgr = SandboxManager(backend, max_ttl=600) + mgr.open("a") + program = mgr._build_remove_program(".") + # Run the generated guard program in-process against a sandbox dir; a '.' path + # must be a no-op (root guard), never rmtree the cwd. + import os + import tempfile + + with tempfile.TemporaryDirectory() as tmp: + keep = os.path.join(tmp, "keep.txt") + with open(keep, "w") as fh: + fh.write("x") + cwd = os.getcwd() + os.chdir(tmp) + try: + exec(compile(program, "", "exec"), {}) + finally: + os.chdir(cwd) + assert os.path.exists(keep), "root-guard let '.' delete the workspace" + # And empty/'./' are likewise refused. + for bad in ("", "./", "."): + assert "os.path.normpath(_p) != '.'" in mgr._build_remove_program(bad) + + +# --------------------------------------------------------------------------- +# Thread-safety smoke +# --------------------------------------------------------------------------- + + +def test_concurrent_open_close_stays_consistent(backend): + mgr = SandboxManager(backend, max_ttl=600, max_sessions=8) + errors: List[Exception] = [] + + def churn(i: int) -> None: + try: + for n in range(50): + sid = f"s-{i}-{n % 4}" + try: + mgr.open(sid) + mgr.exec(sid, "x") + except SandboxCapacityError: + pass + mgr.reap_expired() + mgr.close(sid) + except Exception as exc: # noqa: BLE001 - surface any thread error to the assertion + errors.append(exc) + + threads = [threading.Thread(target=churn, args=(i,)) for i in range(6)] + for t in threads: + t.start() + for t in threads: + t.join() + assert errors == [] + assert mgr.session_count() <= 8 + + +# --------------------------------------------------------------------------- +# Concurrency: lock must never be held across a blocking backend call +# --------------------------------------------------------------------------- + + +def test_lock_not_held_across_blocking_open(): + """While a cold backend.open blocks, other lock-taking methods return promptly.""" + + class _BlockingOpenBackend(FakeBackend): + def __init__(self) -> None: + super().__init__() + self.entered_open = threading.Event() + self.release_open = threading.Event() + + def open(self, session_id: str) -> str: + if session_id == "slow": + self.entered_open.set() + # Block here as if cold-starting a kernel/sandbox (network I/O ~60s). + assert self.release_open.wait(timeout=5) + return super().open(session_id) + + backend = _BlockingOpenBackend() + mgr = SandboxManager(backend, max_ttl=600, max_sessions=8) + + opener = threading.Thread(target=lambda: mgr.open("slow")) + opener.start() + try: + # Wait until the backend open is in flight (lock already released by design). + assert backend.entered_open.wait(timeout=5) + + # These take the manager lock; if the lock were held across backend.open they + # would block for the full 5s. They must return effectively immediately. + results: dict = {} + + def probe() -> None: + results["count"] = mgr.session_count() + results["has_other"] = mgr.has_session("other") + # A second open for a DIFFERENT id must also make progress (reserve + open). + results["other_handle"] = mgr.open("other") + + prober = threading.Thread(target=probe) + prober.start() + prober.join(timeout=3) + assert not prober.is_alive(), "lock was held across backend.open (probe blocked)" + assert results["count"] >= 1 # the reserved "slow" placeholder counts + assert results["has_other"] is False + assert results["other_handle"].startswith("handle-other") + finally: + backend.release_open.set() + opener.join(timeout=5) + assert mgr.has_session("slow") + + +# --------------------------------------------------------------------------- +# Concurrency: eviction closes the captured resource, never a re-opened one +# --------------------------------------------------------------------------- + + +def test_evict_then_concurrent_reopen_closes_old_handle_not_new(): + """Evicting V and concurrently re-opening V must tear down V's OLD handle, not the new one.""" + + class _SlowCloseBackend(FakeBackend): + def __init__(self) -> None: + super().__init__() + self.block_handle = None # only this exact handle's close blocks + self.close_started = threading.Event() + self.release_close = threading.Event() + + def close_handle(self, session_id: str, handle: str) -> None: + if handle == self.block_handle: + self.close_started.set() + # Hold the deferred close open so a concurrent open(V) lands first. + assert self.release_close.wait(timeout=5) + super().close_handle(session_id, handle) + + backend = _SlowCloseBackend() + # cap=2: one slot pinned by a busy filler, one free slot the eviction/reopen contend over. + mgr = SandboxManager(backend, max_ttl=600, max_sessions=2) + + mgr.open("U") + mgr._enter("U") # pin U busy so it is never the eviction victim + old_handle = mgr.open("V") # the victim's original handle (free slot) + backend.block_handle = old_handle # only V's OLD handle's close will block + + # open("W") is at cap -> evicts idle "V"; its deferred close_handle(V, old) blocks. + evictor = threading.Thread(target=lambda: mgr.open("W")) + evictor.start() + assert backend.close_started.wait(timeout=5) # eviction close of V's old handle is in flight + + # At this point W holds the free slot as a placeholder. Free it for the reopen by + # letting the evicting open finish its own backend.open; W then occupies the slot. + # To re-open V we must first release U so the cap has room. + mgr._leave("U") + + # Re-open V WHILE its old handle's close is still blocked: a brand-new handle. + new_handle = mgr.open("V") + assert new_handle != old_handle + + # Now let the deferred close of the OLD handle complete. + backend.release_close.set() + evictor.join(timeout=5) + + # The OLD handle was the one closed; the NEW handle for V survives intact. + assert (("V", old_handle) in backend.closed_handles) + assert (("V", new_handle) not in backend.closed_handles) + assert backend._handles.get("V") == new_handle # registry still points at the new one + assert mgr.has_session("V") + # The new V is usable (its workspace was not torn down by the stale close). + mgr.put_file("V", "f.txt", b"data") + assert mgr.get_file("V", "f.txt") == b"data" diff --git a/tests/storage/db/repositories/test_artifacts.py b/tests/storage/db/repositories/test_artifacts.py index 15bb9aca..2d2f15a4 100644 --- a/tests/storage/db/repositories/test_artifacts.py +++ b/tests/storage/db/repositories/test_artifacts.py @@ -224,6 +224,104 @@ class TestListArtifacts: repo.list_artifacts() +class TestQuotaAccounting: + def test_count_for_user_counts_only_that_user(self, pg_conn): + repo = _repo(pg_conn) + conv = _conversation_id() + repo.create_artifact("alice", "document", conversation_id=conv) + repo.create_artifact("alice", "document", conversation_id=conv) + repo.create_artifact("bob", "document", conversation_id=conv) + assert repo.count_for_user("alice") == 2 + assert repo.count_for_user("bob") == 1 + assert repo.count_for_user("nobody") == 0 + + def test_total_bytes_sums_all_versions(self, pg_conn): + repo = _repo(pg_conn) + conv = _conversation_id() + a = repo.create_artifact("alice", "document", conversation_id=conv, size=100) + repo.append_version(a["id"], size=250) # +250 -> 350 for the same artifact + repo.create_artifact("alice", "document", conversation_id=conv, size=50) + repo.create_artifact("bob", "document", conversation_id=conv, size=999) + assert repo.total_bytes_for_user("alice") == 400 + assert repo.total_bytes_for_user("bob") == 999 + + def test_total_bytes_zero_for_unknown_user(self, pg_conn): + repo = _repo(pg_conn) + assert repo.total_bytes_for_user("nobody") == 0 + + def test_total_bytes_treats_null_size_as_zero(self, pg_conn): + repo = _repo(pg_conn) + repo.create_artifact("alice", "document", conversation_id=_conversation_id()) + assert repo.total_bytes_for_user("alice") == 0 + + +class TestQuotaEnforcement: + """Exercises the shared enforcement helper against a real repo + connection.""" + + def _seed(self, pg_conn, *, count: int, size: int) -> None: + repo = _repo(pg_conn) + for _ in range(count): + repo.create_artifact("alice", "document", conversation_id=_conversation_id(), size=size) + + def test_under_quota_passes(self, pg_conn, monkeypatch): + from application.sandbox import artifacts_capture as ac + + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_COUNT_PER_USER", 10, raising=False) + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_TOTAL_BYTES_PER_USER", 10_000, raising=False) + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_BYTES", 10_000, raising=False) + self._seed(pg_conn, count=2, size=100) + # 2 existing + 1 new = 3 <= 10; 200 + 100 = 300 <= 10000. Must not raise. + ac._enforce_user_quota(_repo(pg_conn), "alice", 100, new_artifact=True) + + def test_over_count_quota_raises(self, pg_conn, monkeypatch): + from application.sandbox import artifacts_capture as ac + + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_COUNT_PER_USER", 2, raising=False) + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_TOTAL_BYTES_PER_USER", 10_000_000, raising=False) + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_BYTES", 10_000_000, raising=False) + self._seed(pg_conn, count=2, size=10) + with pytest.raises(ac.QuotaExceeded): + ac._enforce_user_quota(_repo(pg_conn), "alice", 10, new_artifact=True) + + def test_count_quota_ignored_when_appending_version(self, pg_conn, monkeypatch): + from application.sandbox import artifacts_capture as ac + + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_COUNT_PER_USER", 2, raising=False) + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_TOTAL_BYTES_PER_USER", 10_000_000, raising=False) + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_BYTES", 10_000_000, raising=False) + self._seed(pg_conn, count=2, size=10) + # Appending a version to an existing identity must not trip the count cap. + ac._enforce_user_quota(_repo(pg_conn), "alice", 10, new_artifact=False) + + def test_over_total_bytes_quota_raises(self, pg_conn, monkeypatch): + from application.sandbox import artifacts_capture as ac + + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_COUNT_PER_USER", 10_000, raising=False) + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_TOTAL_BYTES_PER_USER", 500, raising=False) + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_BYTES", 10_000, raising=False) + self._seed(pg_conn, count=1, size=450) + with pytest.raises(ac.QuotaExceeded): + ac._enforce_user_quota(_repo(pg_conn), "alice", 100, new_artifact=True) # 450 + 100 > 500 + + def test_single_artifact_too_large_raises(self, pg_conn, monkeypatch): + from application.sandbox import artifacts_capture as ac + + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_COUNT_PER_USER", 10_000, raising=False) + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_TOTAL_BYTES_PER_USER", 10_000_000, raising=False) + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_BYTES", 1000, raising=False) + with pytest.raises(ac.QuotaExceeded): + ac._enforce_user_quota(_repo(pg_conn), "alice", 1001, new_artifact=True) + + def test_zero_settings_disable_enforcement(self, pg_conn, monkeypatch): + from application.sandbox import artifacts_capture as ac + + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_COUNT_PER_USER", 0, raising=False) + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_TOTAL_BYTES_PER_USER", 0, raising=False) + monkeypatch.setattr(ac.settings, "ARTIFACT_MAX_BYTES", 0, raising=False) + self._seed(pg_conn, count=5, size=10_000) + ac._enforce_user_quota(_repo(pg_conn), "alice", 10_000_000, new_artifact=True) # disabled -> no raise + + class TestCascadeDelete: def test_deleting_artifact_removes_versions(self, pg_conn): repo = _repo(pg_conn)