Reclaim ephemeral agent homes left behind by abnormal exits

Moving ephemeral homes out of the system temp directory and under Studio's auth
tree removed the only thing that ever cleaned them up: the OS. When the wrapper
is killed by SIGKILL, the console closes or the machine crashes, the context
manager's finally never runs, and nothing prunes
<STUDIO_HOME>/auth/agents/.tmp, so interrupted sessions accumulate there
indefinitely. Only the Windows codex path had reclamation.

`_temporary_agent_config` now goes through the same locked session helper that
path already used, so every agent gets the scavenge on launch, the advisory
lock that keeps a live session from being swept, and the heartbeat that anchors
the stale window to wrapper death rather than session start.
`_reclaim_stale_ephemeral_sessions` and `_short_ephemeral_session` take the
prefix to glob and create, which is the only part that was codex specific.

An age-only sweep was not enough on its own: without the live marker a session
still running after the stale window would be deleted underneath itself.

unsloth_cli/tests/test_start.py: 382 passed.
This commit is contained in:
danielhanchen 2026-07-29 08:24:51 +00:00
commit eec7c67d1e
2 changed files with 36 additions and 12 deletions

View file

@ -2799,13 +2799,13 @@ def _agents_config_root() -> Path:
@contextlib.contextmanager
def _temporary_agent_config(prefix: str):
# These homes live under Studio's auth tree, which nothing else prunes, so reuse the
# locked session helper: a wrapper killed before its finally runs leaves a home that
# the next launch reclaims, and the lock keeps a live session from being swept.
temp_root = _agents_config_root() / ".tmp"
temp_root.mkdir(parents = True, exist_ok = True, mode = 0o700)
path = Path(tempfile.mkdtemp(prefix = prefix, dir = temp_root))
try:
with _short_ephemeral_session(temp_root, prefix) as path:
yield path
finally:
shutil.rmtree(path, ignore_errors = True)
def _ephemeral_session_parent(agent: str) -> Optional[Path]:
@ -2874,9 +2874,9 @@ def _locked_file(path: Path, blocking: bool = True):
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-*"):
def _reclaim_stale_ephemeral_sessions(parent: Path, prefix: str) -> None:
"""Remove abandoned session homes while preserving locked live sessions."""
for path in parent.glob(f"{prefix}*"):
if not path.is_dir():
continue
active_lock = path / ".active.lock"
@ -2907,8 +2907,8 @@ def _refresh_ephemeral_session_marker(path: Path, stop: threading.Event) -> None
@contextlib.contextmanager
def _short_ephemeral_session(parent: Path):
"""Create a short Codex home whose lock makes crash cleanup concurrency-safe."""
def _short_ephemeral_session(parent: Path, prefix: str = "u-codex-"):
"""Create a session home whose lock makes crash cleanup concurrency-safe."""
path = None
active_lock = contextlib.ExitStack()
heartbeat_stop = None
@ -2917,8 +2917,8 @@ def _short_ephemeral_session(parent: Path):
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))
_reclaim_stale_ephemeral_sessions(parent, prefix)
path = Path(tempfile.mkdtemp(prefix = prefix, 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}")
@ -2926,7 +2926,7 @@ def _short_ephemeral_session(parent: Path):
heartbeat = threading.Thread(
target = _refresh_ephemeral_session_marker,
args = (path / ".active.lock", heartbeat_stop),
name = "unsloth-codex-home-heartbeat",
name = "unsloth-agent-home-heartbeat",
daemon = True,
)
heartbeat.start()

View file

@ -5199,6 +5199,30 @@ def test_session_config_reclaims_old_short_homes_but_keeps_recent_and_live(monke
assert not first.exists()
def test_session_config_reclaims_abandoned_homes_for_non_codex_agents(monkeypatch, tmp_path):
# These homes sit under Studio's auth tree, which nothing else prunes, so a wrapper
# killed before its finally runs must be reclaimed by the next launch.
agents_root = tmp_path / "agents"
temp_root = agents_root / ".tmp"
temp_root.mkdir(parents = True)
monkeypatch.setattr(start, "_agents_config_root", lambda: agents_root)
abandoned = temp_root / "unsloth-claude-abandoned"
abandoned.mkdir()
(abandoned / ".active.lock").write_bytes(b"\0")
(abandoned / "state.json").write_text("left behind")
old = time.time() - start._CODEX_EPHEMERAL_STALE_SECONDS - 1
os.utime(abandoned / ".active.lock", (old, old))
recent = temp_root / "unsloth-claude-still-running"
recent.mkdir()
(recent / ".active.lock").write_bytes(b"\0")
with start._session_config("claude", launch = True) as home:
assert not abandoned.exists()
assert recent.exists()
assert home.parent == temp_root
assert not home.exists()
def test_session_config_serializes_normal_short_home_deletion(monkeypatch, tmp_path):
short_parent = tmp_path / "u"
short_parent.mkdir()