diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 3ae0276473..c702b90748 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -6,6 +6,7 @@ import atexit import base64 import contextlib +import errno import json import os import re @@ -102,6 +103,8 @@ _CODEX_SUBAGENT_MCP_SERVER = "unsloth_local_agent" _CODEX_SUBAGENT_MCP_TOOL = "spawn_local_agent" _CODEX_SUBAGENT_CONFIG_ENV = "UNSLOTH_CODEX_SUBAGENT_CONFIG" _CODEX_PARENT_OVERLAY_MANIFEST = ".unsloth-parent-overlay.json" +_CODEX_EPHEMERAL_STALE_SECONDS = 24 * 60 * 60 +_CODEX_EPHEMERAL_HEARTBEAT_SECONDS = 60 _CODEX_SUBAGENT_TOOL_DESCRIPTION = ( f"{_SUBAGENT_DESCRIPTION} Use this tool instead of the built-in spawn_agent tool for those " "requests. Other subagent requests may use the built-in tools normally." @@ -2680,6 +2683,147 @@ def _agents_config_root() -> Path: return auth_root() / "agents" +def _ephemeral_session_parent(agent: str) -> Optional[Path]: + """Return a non-system-temp parent when an agent needs one.""" + if os.name != "nt" or agent != "codex": + return None + # Codex creates a deeply nested curated-plugin checkout below CODEX_HOME. + # A normal %TEMP%\unsloth-codex-* home can exceed legacy Windows path + # limits during startup, and Codex also refuses to create its PATH helpers + # below the system temp directory. Keep the throwaway home short but still + # private to the current user; _session_config removes it on exit. + root = Path.home() / ".unsloth" / ".tmp" + root.mkdir(parents = True, exist_ok = True, mode = 0o700) + return root + + +def _ephemeral_session_prefix(agent: str, parent: Optional[Path]) -> str: + """Return the platform-specific prefix for an ephemeral agent home.""" + return "u-codex-" if agent == "codex" and parent is not None else f"unsloth-{agent}-" + + +@contextlib.contextmanager +def _locked_file(path: Path, blocking: bool = True): + """Yield whether an advisory lock was acquired for the first byte of path.""" + handle = path.open("a+b") + acquired = False + try: + if os.name == "nt": + import msvcrt + + handle.seek(0, os.SEEK_END) + if handle.tell() == 0: + handle.write(b"\0") + handle.flush() + handle.seek(0) + while True: + try: + msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1) + acquired = True + break + except OSError as exc: + if exc.errno not in (errno.EACCES, errno.EAGAIN, errno.EDEADLK): + raise + if not blocking: + break + # LK_LOCK gives up after roughly ten seconds. Poll LK_NBLCK + # instead so a large stale plugin checkout cannot make a + # concurrent launch fail just because cleanup takes longer. + time.sleep(0.05) + else: + import fcntl + mode = fcntl.LOCK_EX | (0 if blocking else fcntl.LOCK_NB) + try: + fcntl.flock(handle.fileno(), mode) + acquired = True + except BlockingIOError: + acquired = False + yield acquired + finally: + if acquired: + if os.name == "nt": + handle.seek(0) + msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1) + else: + fcntl.flock(handle.fileno(), fcntl.LOCK_UN) + handle.close() + + +def _reclaim_stale_ephemeral_sessions(parent: Path) -> None: + """Remove abandoned short Codex homes while preserving locked live sessions.""" + for path in parent.glob("u-codex-*"): + if not path.is_dir(): + continue + active_lock = path / ".active.lock" + try: + modified = active_lock.stat().st_mtime if active_lock.exists() else path.stat().st_mtime + except FileNotFoundError: + continue + # The wrapper owns the advisory lock, not the Codex child. If only the + # wrapper is killed, its child may still be using CODEX_HOME; give that + # process a full day to finish before treating the unlocked home as stale. + if time.time() - modified < _CODEX_EPHEMERAL_STALE_SECONDS: + continue + try: + with _locked_file(active_lock, blocking = False) as stale: + pass + except FileNotFoundError: + # A normally exiting session may have removed itself after the glob. + continue + if stale: + shutil.rmtree(path, ignore_errors = True) + + +def _refresh_ephemeral_session_marker(path: Path, stop: threading.Event) -> None: + """Keep the stale grace period relative to wrapper death, not session start.""" + while not stop.wait(_CODEX_EPHEMERAL_HEARTBEAT_SECONDS): + with contextlib.suppress(OSError): + os.utime(path, None) + + +@contextlib.contextmanager +def _short_ephemeral_session(parent: Path): + """Create a short Codex home whose lock makes crash cleanup concurrency-safe.""" + path = None + active_lock = contextlib.ExitStack() + heartbeat_stop = None + heartbeat = None + try: + with _locked_file(parent / ".cleanup.lock") as cleanup_lock: + if not cleanup_lock: # The blocking acquisition should always succeed. + raise RuntimeError(f"Could not lock ephemeral session root: {parent}") + _reclaim_stale_ephemeral_sessions(parent) + path = Path(tempfile.mkdtemp(prefix = "u-codex-", dir = parent)) + locked = active_lock.enter_context(_locked_file(path / ".active.lock")) + if not locked: + raise RuntimeError(f"Could not lock ephemeral session home: {path}") + heartbeat_stop = threading.Event() + heartbeat = threading.Thread( + target = _refresh_ephemeral_session_marker, + args = (path / ".active.lock", heartbeat_stop), + name = "unsloth-codex-home-heartbeat", + daemon = True, + ) + heartbeat.start() + yield path + finally: + if heartbeat_stop is not None: + heartbeat_stop.set() + if heartbeat is not None: + heartbeat.join(timeout = 1) + try: + with _locked_file(parent / ".cleanup.lock") as cleanup_lock: + if not cleanup_lock: # The blocking acquisition should always succeed. + raise RuntimeError(f"Could not lock ephemeral session root: {parent}") + # Release the live marker only after deletion is serialized with + # startup scavenging, so no scanner can race this rmtree. + active_lock.close() + if path is not None: + shutil.rmtree(path, ignore_errors = True) + finally: + active_lock.close() + + @contextlib.contextmanager def _session_config( agent: str, @@ -2695,11 +2839,16 @@ def _session_config( resumed next time. Either way the user's real ~/. config is left untouched. """ if launch and not persist: - path = Path(tempfile.mkdtemp(prefix = f"unsloth-{agent}-")) - try: - yield path - finally: - shutil.rmtree(path, ignore_errors = True) + parent = _ephemeral_session_parent(agent) + if parent is not None: + with _short_ephemeral_session(parent) as path: + yield path + else: + path = Path(tempfile.mkdtemp(prefix = _ephemeral_session_prefix(agent, parent))) + try: + yield path + finally: + shutil.rmtree(path, ignore_errors = True) else: # Never wipe this dir: a previously printed recipe may still be running # an agent whose sessions/state live here, and every config writer diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index 2d4be513fc..dd8a1daf2d 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -10,6 +10,7 @@ import os import re import shlex import sys +import time import urllib.error from pathlib import Path from types import SimpleNamespace @@ -1550,7 +1551,8 @@ def test_connect_codex_launch_uses_ephemeral_home(fake_studio, monkeypatch): assert result.exit_code == 0, result.output home = Path(captured["home"]) assert captured["config_present"] # config existed while codex ran - assert "unsloth-codex-" in home.name # an ephemeral temp dir, not ~/.codex + parent = start._ephemeral_session_parent("codex") + assert home.name.startswith(start._ephemeral_session_prefix("codex", parent)) assert not home.exists() # cleaned up after the agent exits @@ -4734,7 +4736,96 @@ def test_session_config_default_launch_is_ephemeral(): # Default launch (no --persist) still uses a throwaway temp dir wiped on exit. with start._session_config("codex", launch = True) as home: assert home.exists() - assert "unsloth-codex-" in home.name + parent = start._ephemeral_session_parent("codex") + assert home.name.startswith(start._ephemeral_session_prefix("codex", parent)) + assert not home.exists() + + +def test_session_config_codex_uses_short_ephemeral_parent(monkeypatch, tmp_path): + # Windows Codex checks out its curated plugins under CODEX_HOME/.tmp/plugins. + # Put its throwaway home outside the longer system temp path so that checkout + # stays below legacy MAX_PATH and Codex does not reject temp-dir PATH helpers. + short_parent = tmp_path / "u" + short_parent.mkdir() + monkeypatch.setattr( + start, + "_ephemeral_session_parent", + lambda agent: short_parent if agent == "codex" else None, + ) + + with start._session_config("codex", launch = True) as home: + assert home.parent == short_parent + assert home.name.startswith("u-codex-") + assert home.exists() + assert not home.exists() + + +def test_locked_file_windows_blocking_retries_until_acquired(monkeypatch, tmp_path): + attempts = [] + sleeps = [] + + def locking(_fd, mode, _length): + if mode == 1: + attempts.append(mode) + if len(attempts) < 3: + raise PermissionError(start.errno.EACCES, "busy") + + fake_msvcrt = SimpleNamespace(LK_NBLCK = 1, LK_UNLCK = 2, locking = locking) + monkeypatch.setitem(sys.modules, "msvcrt", fake_msvcrt) + _simulate_windows(monkeypatch) + monkeypatch.setattr(start.time, "sleep", sleeps.append) + + with start._locked_file(tmp_path / "lock") as acquired: + assert acquired + assert len(attempts) == 3 + assert sleeps == [0.05, 0.05] + + +def test_session_config_reclaims_old_short_homes_but_keeps_recent_and_live(monkeypatch, tmp_path): + short_parent = tmp_path / "u" + short_parent.mkdir() + stale = short_parent / "u-codex-abandoned" + stale.mkdir() + (stale / ".active.lock").write_bytes(b"\0") + (stale / "plugin-checkout").write_text("left behind") + old = time.time() - start._CODEX_EPHEMERAL_STALE_SECONDS - 1 + os.utime(stale / ".active.lock", (old, old)) + recent = short_parent / "u-codex-surviving-child" + recent.mkdir() + (recent / ".active.lock").write_bytes(b"\0") + monkeypatch.setattr( + start, + "_ephemeral_session_parent", + lambda agent: short_parent if agent == "codex" else None, + ) + + with start._session_config("codex", launch = True) as first: + assert not stale.exists() + assert recent.exists() + with start._session_config("codex", launch = True) as second: + assert first.exists() + assert second.exists() + assert first != second + assert first.exists() + assert not second.exists() + assert not first.exists() + + +def test_session_config_serializes_normal_short_home_deletion(monkeypatch, tmp_path): + short_parent = tmp_path / "u" + short_parent.mkdir() + monkeypatch.setattr(start, "_ephemeral_session_parent", lambda _agent: short_parent) + original_rmtree = start.shutil.rmtree + + def checked_rmtree(path, *args, **kwargs): + if path.parent == short_parent and path.name.startswith("u-codex-"): + with start._locked_file(short_parent / ".cleanup.lock", blocking = False) as unlocked: + assert not unlocked + return original_rmtree(path, *args, **kwargs) + + monkeypatch.setattr(start.shutil, "rmtree", checked_rmtree) + with start._session_config("codex", launch = True) as home: + assert home.exists() assert not home.exists() @@ -4782,7 +4873,8 @@ def test_default_launch_home_is_ephemeral(agent, fake_studio, tmp_path, monkeypa monkeypatch.setattr(start.shutil, "which", lambda _: f"/usr/local/bin/{agent}") captured = _capture_launch(monkeypatch, [agent]) home = captured["env"][_RESUME_ENV_VAR[agent]] - assert f"unsloth-{agent}-" in home + parent = start._ephemeral_session_parent(agent) + assert start._ephemeral_session_prefix(agent, parent) in home assert str(tmp_path / "agents") not in home