from __future__ import annotations import hashlib import json import os import threading import time from dataclasses import dataclass from yuxi import config as conf from yuxi.utils.logging_config import logger from .provisioner_client import ProvisionerClient, SandboxRecord def sandbox_id_for_thread(thread_id: str, skills_thread_id: str | None = None, *, uid: str | None = None) -> str: file_thread_id = str(thread_id or "").strip() skills_id = str(skills_thread_id or file_thread_id).strip() uid_id = str(uid or "").strip() scope = file_thread_id if skills_id == file_thread_id else f"{file_thread_id}:{skills_id}" identity = f"{uid_id}:{scope}" if uid_id else scope digest = hashlib.sha256(identity.encode("utf-8")).hexdigest() return digest[:12] def _sandbox_key(uid: str, file_thread_id: str, skills_thread_id: str) -> str: return f"{uid}::{file_thread_id}::{skills_thread_id}" def normalize_env(env: dict | None) -> dict[str, str]: if not isinstance(env, dict): return {} return {str(key): "" if value is None else str(value) for key, value in env.items() if str(key)} def postgres_conninfo() -> str: db_url = os.getenv("POSTGRES_URL", "").strip() return db_url.replace("+asyncpg", "").replace("+psycopg", "") def load_user_agent_env(uid: str) -> dict[str, str]: conninfo = postgres_conninfo() if not conninfo: return {} try: import psycopg with psycopg.connect(conninfo, connect_timeout=3) as conn: with conn.cursor() as cursor: cursor.execute("SELECT env FROM agent_envs WHERE uid = %s", (uid,)) row = cursor.fetchone() except Exception as exc: raise RuntimeError(f"failed to load agent env for uid {uid}: {exc}") from exc if not row: return {} value = row[0] if isinstance(value, str): try: value = json.loads(value) except json.JSONDecodeError as exc: raise RuntimeError(f"stored agent env for uid {uid} is not valid JSON") from exc return normalize_env(value) @dataclass(slots=True) class SandboxConnection: cache_key: str thread_id: str file_thread_id: str skills_thread_id: str uid: str sandbox_id: str sandbox_url: str class ProvisionerSandboxProvider: def __init__(self): provider_name = str(getattr(conf, "sandbox_provider", "provisioner")).strip().lower() if provider_name != "provisioner": raise RuntimeError("only sandbox_provider=provisioner is supported") provisioner_url = str(getattr(conf, "sandbox_provisioner_url", "") or "").strip() if not provisioner_url: raise RuntimeError("sandbox_provisioner_url is required") self._client = ProvisionerClient(provisioner_url) self._lock = threading.Lock() self._thread_locks: dict[str, threading.Lock] = {} self._connections: dict[str, SandboxConnection] = {} self._last_touch_at: dict[str, float] = {} self._touch_interval_seconds = int(getattr(conf, "sandbox_keepalive_interval_seconds", 30)) def _thread_lock(self, cache_key: str) -> threading.Lock: with self._lock: lock = self._thread_locks.get(cache_key) if lock is None: lock = threading.Lock() self._thread_locks[cache_key] = lock return lock def _record_to_connection( self, *, cache_key: str, thread_id: str, file_thread_id: str, skills_thread_id: str, uid: str, record: SandboxRecord, ) -> SandboxConnection: connection = SandboxConnection( cache_key=cache_key, thread_id=thread_id, file_thread_id=file_thread_id, skills_thread_id=skills_thread_id, uid=uid, sandbox_id=record.sandbox_id, sandbox_url=record.sandbox_url, ) self._connections[cache_key] = connection self._last_touch_at[cache_key] = time.time() return connection def _should_touch(self, cache_key: str) -> bool: if self._touch_interval_seconds <= 0: return False last_touch = self._last_touch_at.get(cache_key) if last_touch is None: return True return (time.time() - last_touch) >= self._touch_interval_seconds def _touch_if_needed(self, connection: SandboxConnection) -> bool: if not self._should_touch(connection.cache_key): return True is_alive = self._client.touch(connection.sandbox_id) self._last_touch_at[connection.cache_key] = time.time() return is_alive def acquire( self, thread_id: str, *, uid: str, file_thread_id: str | None = None, skills_thread_id: str | None = None, ) -> str: file_id = str(file_thread_id or thread_id).strip() skills_id = str(skills_thread_id or thread_id).strip() cache_key = _sandbox_key(uid, file_id, skills_id) lock = self._thread_lock(cache_key) with lock: current = self._connections.get(cache_key) if current: if current.uid != uid: raise RuntimeError(f"sandbox scope {cache_key} belongs to uid {current.uid}, not {uid}") try: if self._touch_if_needed(current): return current.sandbox_id self._connections.pop(cache_key, None) self._last_touch_at.pop(cache_key, None) except Exception as exc: # noqa: BLE001 logger.warning(f"Failed to touch sandbox {current.sandbox_id} for {cache_key}: {exc}") return current.sandbox_id sandbox_id = sandbox_id_for_thread(file_id, skills_id, uid=uid) logger.info(f"Ensuring sandbox {sandbox_id} for file thread {file_id} and skills thread {skills_id}") record = self._client.create( sandbox_id, thread_id, uid, load_user_agent_env(uid), file_thread_id=file_id, skills_thread_id=skills_id, ) connection = self._record_to_connection( cache_key=cache_key, thread_id=thread_id, file_thread_id=file_id, skills_thread_id=skills_id, uid=uid, record=record, ) return connection.sandbox_id def get( self, thread_id: str, *, uid: str, create_if_missing: bool = False, file_thread_id: str | None = None, skills_thread_id: str | None = None, ) -> SandboxConnection | None: file_id = str(file_thread_id or thread_id).strip() skills_id = str(skills_thread_id or thread_id).strip() cache_key = _sandbox_key(uid, file_id, skills_id) lock = self._thread_lock(cache_key) with lock: current = self._connections.get(cache_key) if current: if current.uid != uid: raise RuntimeError(f"sandbox scope {cache_key} belongs to uid {current.uid}, not {uid}") try: if self._touch_if_needed(current): return current self._connections.pop(cache_key, None) self._last_touch_at.pop(cache_key, None) except Exception as exc: # noqa: BLE001 logger.warning(f"Failed to touch sandbox {current.sandbox_id} for {cache_key}: {exc}") return current sandbox_id = sandbox_id_for_thread(file_id, skills_id, uid=uid) if create_if_missing: record = self._client.create( sandbox_id, thread_id, uid, load_user_agent_env(uid), file_thread_id=file_id, skills_thread_id=skills_id, ) else: record = self._client.discover(sandbox_id) if record is None: return None return self._record_to_connection( cache_key=cache_key, thread_id=thread_id, file_thread_id=file_id, skills_thread_id=skills_id, uid=uid, record=record, ) def shutdown(self) -> None: with self._lock: connections = list(self._connections.values()) self._connections.clear() self._last_touch_at.clear() for connection in connections: try: self._client.delete(connection.sandbox_id) except Exception as exc: # noqa: BLE001 logger.warning(f"Failed to release sandbox {connection.sandbox_id} for {connection.cache_key}: {exc}") _sandbox_provider: ProvisionerSandboxProvider | None = None _sandbox_provider_lock = threading.Lock() def init_sandbox_provider() -> ProvisionerSandboxProvider: global _sandbox_provider with _sandbox_provider_lock: if _sandbox_provider is None: _sandbox_provider = ProvisionerSandboxProvider() return _sandbox_provider def get_sandbox_provider() -> ProvisionerSandboxProvider: provider = _sandbox_provider if provider is not None: return provider return init_sandbox_provider() def shutdown_sandbox_provider() -> None: global _sandbox_provider with _sandbox_provider_lock: provider = _sandbox_provider _sandbox_provider = None if provider is not None: provider.shutdown()