wehub-resource-sync
fed8b2eed7
Backend release / release (push) Waiting to run
Bandit Security Scan / bandit_scan (push) Waiting to run
Build and push multi-arch DocsGPT Docker image / build (linux/amd64, ubuntu-latest, amd64) (push) Waiting to run
Build and push multi-arch DocsGPT Docker image / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Waiting to run
Build and push multi-arch DocsGPT Docker image / manifest (push) Blocked by required conditions
Build and push DocsGPT FE Docker image for development / build (linux/amd64, ubuntu-latest, amd64) (push) Waiting to run
Build and push DocsGPT FE Docker image for development / build (linux/arm64, ubuntu-24.04-arm, arm64) (push) Waiting to run
Build and push DocsGPT FE Docker image for development / manifest (push) Blocked by required conditions
Python linting / ruff (push) Waiting to run
Run python tests with pytest / Run tests and count coverage (3.12) (push) Waiting to run
React Widget Build / build (push) Waiting to run
392 lines
18 KiB
Python
392 lines
18 KiB
Python
"""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, 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, 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. ``pending_close`` marks a session whose
|
|
``close`` arrived while an op held it ``in_use``; the last ``_leave`` then tears it down.
|
|
"""
|
|
|
|
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)
|
|
pending_close: bool = field(default=False)
|
|
|
|
def is_expired(self, now: float) -> bool:
|
|
"""True when the session has been idle longer than its (clamped) TTL."""
|
|
return (now - self.last_access) > self.ttl
|
|
|
|
|
|
class SandboxManager:
|
|
"""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,
|
|
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()
|
|
|
|
def _clamp_ttl(self, ttl: Optional[float]) -> float:
|
|
"""Return an agent-requested TTL bounded to (0, ``max_ttl``]."""
|
|
if ttl is None:
|
|
ttl = self._default_ttl
|
|
if ttl <= 0:
|
|
return self._default_ttl
|
|
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.
|
|
|
|
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 and session.ready:
|
|
session.last_access = now
|
|
# Reuse must honor an explicit longer keep-alive: a run_code
|
|
# persist/ttl on a session first opened by another tool (e.g.
|
|
# artifact_generator at the 60s exec timeout) would otherwise be
|
|
# dropped, and the kernel + its background state reaped early.
|
|
# Extend only (never shrink) so a default reuse can't cut short a
|
|
# session another caller kept alive.
|
|
if ttl is not None:
|
|
session.ttl = max(session.ttl, self._clamp_ttl(ttl))
|
|
return session.handle
|
|
reaped = self._reap_locked(now)
|
|
if session is None:
|
|
# Genuinely new key: evict an LRU-idle victim if at capacity (may
|
|
# raise SandboxCapacityError), then RESERVE a placeholder slot so
|
|
# concurrent opens can't overshoot the cap.
|
|
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,
|
|
)
|
|
else:
|
|
# A not-yet-ready placeholder for THIS id already occupies a cap slot
|
|
# (another thread is mid cold-open). Overwriting the same key adds no
|
|
# slot, so do NOT call _make_room_locked here -- it would wrongly evict
|
|
# an innocent LRU-idle session or raise SandboxCapacityError. Refresh
|
|
# in place; both threads call the (idempotent) backend.open and the
|
|
# finalize step re-checks before binding the handle.
|
|
evicted = None
|
|
session.last_access = now
|
|
session.ttl = self._clamp_ttl(ttl)
|
|
|
|
# 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.
|
|
|
|
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, holding it in-use so a reap/evict can't pull it."""
|
|
session = self._enter(session_id)
|
|
try:
|
|
result = self._backend.exec(session_id, code, timeout)
|
|
if result.runtime_invalidated:
|
|
self._drop_invalidated_session(session_id, session)
|
|
return result
|
|
finally:
|
|
self._leave(session_id, expected=session)
|
|
|
|
def put_file(self, session_id: str, dest_path: str, data: bytes) -> None:
|
|
"""Write ``data`` into the bound session's workspace."""
|
|
session = self._enter(session_id)
|
|
try:
|
|
self._backend.put_file(session_id, dest_path, data)
|
|
finally:
|
|
self._leave(session_id, expected=session)
|
|
|
|
def get_file(self, session_id: str, path: str) -> bytes:
|
|
"""Read ``path`` from the bound session's workspace."""
|
|
session = self._enter(session_id)
|
|
try:
|
|
return self._backend.get_file(session_id, path)
|
|
finally:
|
|
self._leave(session_id, expected=session)
|
|
|
|
def list_files(self, session_id: str) -> List[str]:
|
|
"""List files in the bound session's workspace."""
|
|
session = self._enter(session_id)
|
|
try:
|
|
return self._backend.list_files(session_id)
|
|
finally:
|
|
self._leave(session_id, expected=session)
|
|
|
|
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:
|
|
session = 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, expected=session)
|
|
|
|
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.
|
|
|
|
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.
|
|
|
|
Teardown is DEFERRED when an op holds the session ``in_use`` (a concurrent
|
|
exec/put_file/get_file): the session is flagged ``pending_close`` and the last
|
|
``_leave`` performs the actual close, so a concurrent close never kills an
|
|
in-flight exec (which would lose its captured files). The common path -- the
|
|
caller's own exec has already ``_left``, so ``in_use == 0`` -- stays synchronous.
|
|
"""
|
|
with self._lock:
|
|
session = self._sessions.get(session_id)
|
|
if session is None:
|
|
return
|
|
if session.in_use > 0:
|
|
session.pending_close = True
|
|
return
|
|
self._sessions.pop(session_id, None)
|
|
handle = session.handle
|
|
self._close_backend(session_id, handle)
|
|
|
|
def has_session(self, session_id: str) -> bool:
|
|
"""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 _enter(self, session_id: str) -> _Session:
|
|
"""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 or not session.ready:
|
|
raise KeyError(f"No sandbox session bound for {session_id!r}")
|
|
session.last_access = time.monotonic()
|
|
session.in_use += 1
|
|
return session
|
|
|
|
def _drop_invalidated_session(self, session_id: str, expected: _Session) -> None:
|
|
"""Drop the exact manager entry whose backend runtime was already destroyed."""
|
|
with self._lock:
|
|
if self._sessions.get(session_id) is expected:
|
|
self._sessions.pop(session_id, None)
|
|
|
|
def _leave(self, session_id: str, expected: Optional[_Session] = None) -> None:
|
|
"""Release an in-use hold taken by ``_enter``; run a close deferred by ``close`` on the last release.
|
|
|
|
Idempotent if the session was already closed. If ``expected`` is supplied,
|
|
a newer entry under the same id is never modified (generation/ABA guard).
|
|
When the final hold is released and a ``close`` was deferred
|
|
(``pending_close``), the session is popped here and its backend torn down
|
|
OUTSIDE the lock, keyed by the captured handle.
|
|
"""
|
|
handle: Optional[str] = None
|
|
do_close = False
|
|
with self._lock:
|
|
session = self._sessions.get(session_id)
|
|
if session is not None and (expected is None or session is expected) and session.in_use > 0:
|
|
session.in_use -= 1
|
|
session.last_access = time.monotonic()
|
|
if session.in_use == 0 and session.pending_close:
|
|
self._sessions.pop(session_id, None)
|
|
handle = session.handle
|
|
do_close = True
|
|
if do_close:
|
|
self._close_backend(session_id, handle)
|
|
|
|
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.
|
|
|
|
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 = 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
|