Merge lab after verified process lifecycle

routes/shell_routes.py

@@ -26,6 +26,7 @@ from src.host_docker_access import (
 )
 from src.optional_deps import prepare_optional_dependency_import
 from src.auth_helpers import _auth_disabled
+from src import process_lifecycle
 
 # POSIX-only: `pty`/`fcntl` transitively import `termios`, which does NOT exist
 # on Windows, so importing them unconditionally crashed app startup there
@@ -683,18 +684,18 @@ def _session_pgid(pid: int) -> int | None:
     """Process-group id of the session ``pid`` leads, or None if unavailable.
 
     Read this *before* the leader is reaped: once it is, ``getpgid`` fails and
-    the group id can no longer be recovered from the pid.
+    the group id can no longer be recovered from the pid. If the group is the
+    server's own — ``setsid`` did not take effect — there is no session group
+    to signal, and None makes teardown reach the child alone instead of the
+    whole server.
     """
-    getpgid = getattr(os, "getpgid", None)
-    if getpgid is None:  # no process groups (native Windows)
-        return None
-    try:
-        return getpgid(pid)
-    except OSError:
+    pgid = process_lifecycle.pgid_of(pid)
+    if pgid is None or pgid == process_lifecycle.own_pgid():
         return None
+    return pgid
 
 
-def _signal_session(pgid: int | None, pid: int, sig: int) -> bool:
+def _signal_session(pgid: int | None, pid: int | None, sig: int) -> bool:
     """Send ``sig`` to the whole process group, or to the lone process.
 
     Returns whether anything was signalled, so a caller can tell "the session
@@ -702,20 +703,7 @@ def _signal_session(pgid: int | None, pid: int, sig: int) -> bool:
     matters: if ``setsid`` did not take effect, or the platform has no process
     groups, teardown must still reach the child rather than do nothing.
     """
-    killpg = getattr(os, "killpg", None)
-    if pgid is not None and killpg is not None:
-        try:
-            killpg(pgid, sig)
-            return True
-        except ProcessLookupError:
-            return False
-        except OSError:
-            pass  # group signalling refused — fall through to the single pid
-    try:
-        os.kill(pid, sig)
-        return True
-    except OSError:
-        return False
+    return process_lifecycle.signal_group(pid, pgid, sig)
 
 
 def _session_alive(pgid: int | None, pid: int) -> bool:
@@ -723,41 +711,35 @@ def _session_alive(pgid: int | None, pid: int) -> bool:
 
     An unreaped zombie is still signallable, so a True here can also mean the
     leader has exited but not yet been collected. Without a group id this can
-    only speak for the child itself, not for anything it spawned.
+    only speak for the child itself, not for anything it spawned. Only ESRCH
+    proves a group is gone; EPERM is a live group we may not signal, and
+    reporting a surviving session as contained is the one outcome teardown
+    must never produce.
     """
-    killpg = getattr(os, "killpg", None)
-    if pgid is not None and killpg is not None:
-        try:
-            killpg(pgid, 0)
-        except ProcessLookupError:
-            return False  # ESRCH — no member of the group is left
-        except OSError:
-            # Anything else (EPERM when the group holds a process we may not
-            # signal, EINVAL) answers the probe without proving the group is
-            # gone. Only ESRCH does that, so treat the rest as still running:
-            # reporting a surviving session as contained is the one outcome
-            # teardown must never produce.
-            return True
-        return True
+    if pgid is not None:
+        return process_lifecycle.group_present(pgid)
     return pid_alive(pid)
 
 
-async def _await_session_exit(proc, pgid: int | None, pid: int) -> bool:
-    """Wait up to the grace period for the leader and its group to go away."""
-    loop = asyncio.get_running_loop()
-    deadline = loop.time() + PTY_KILL_GRACE
-    while True:
-        remaining = deadline - loop.time()
-        if proc.returncode is None and remaining > 0:
-            # Reap the leader, otherwise its own zombie keeps the group alive
-            # and the liveness probe below can never come back clean.
-            with contextlib.suppress(asyncio.TimeoutError):
-                await asyncio.wait_for(proc.wait(), remaining)
-        if not _session_alive(pgid, pid):
-            return True
-        if loop.time() >= deadline:
-            return False
-        await asyncio.sleep(PTY_KILL_POLL_INTERVAL)
+def _bind_pty_spawn_identity(proc) -> None:
+    """Freeze the PTY leader's identity and its session group at spawn.
+
+    Called immediately after the spawn, while the pid is known to be the child
+    just created: we hold it unreaped, so the slot cannot have been reissued.
+    The group is recorded only when it is the leader's own (``setsid``
+    applied: pgid == pid) and the identity still verifies after reading it.
+    Teardown works from this record alone and never re-derives ownership
+    from ``proc.pid``, which outlives the process it named.
+    """
+    pid = getattr(proc, "pid", None)
+    if not pid:
+        return
+    identity = process_lifecycle.ProcessIdentity.capture(pid)
+    pgid = _session_pgid(pid)
+    if pgid != pid or identity.verdict() != process_lifecycle.OWNED:
+        pgid = None  # No safe session group: teardown reaches the child alone.
+    proc._ody_pty_identity = process_lifecycle.ProcessIdentity(
+        pid=identity.pid, start_token=identity.start_token, pgid=pgid)
 
 
 async def _terminate_pty_session(proc) -> bool:
@@ -770,18 +752,67 @@ async def _terminate_pty_session(proc) -> bool:
     command as terminated. Signal the group instead, escalate to SIGKILL if it
     outlives the grace period, and return whether the session is actually gone
     so the caller can say so rather than assume it.
+
+    Ownership is the identity frozen at spawn (:func:`_bind_pty_spawn_identity`),
+    re-verified before every signal:
+
+    * leader OWNED and still leading the recorded group → signal the group;
+    * leader GONE (exited and reaped) → the recorded group only, never the
+      pid: a group id is not reissued while the group lives, so a present
+      group with no process in its leader's slot is still ours;
+    * leader FOREIGN → the pid was reissued, which proves our group's
+      lifetime had already ended; nothing of ours is left to signal;
+    * leader UNVERIFIABLE, or no spawn identity at all → nothing is
+      signalled and the session is not reported gone.
+
+    The ladder itself is :func:`src.process_lifecycle.escalate_async`; the
+    leader is reaped through ``proc.wait()`` inside each window, otherwise its
+    own zombie keeps the group alive and the probe can never come back clean.
     """
     pid = getattr(proc, "pid", None)
     if pid is None:
         return True
-    pgid = _session_pgid(pid)
+    frozen = getattr(proc, "_ody_pty_identity", None)
+    if frozen is None:
+        logger.warning("PTY teardown for pid %s has no spawn identity; not signalling", pid)
+        return False
+    pgid = frozen.pgid
 
-    for sig in PTY_KILL_ESCALATION:
-        if not _signal_session(pgid, pid, sig):
-            break  # nothing left to signal
-        if await _await_session_exit(proc, pgid, pid):
+    def _gone() -> bool:
+        verdict = frozen.verdict()
+        if verdict == process_lifecycle.FOREIGN:
             return True
-    return not _session_alive(pgid, pid)
+        if verdict == process_lifecycle.UNVERIFIABLE:
+            return False
+        if pgid is not None:
+            return not _session_alive(pgid, frozen.pid)
+        return verdict == process_lifecycle.GONE or process_lifecycle.is_zombie(frozen.pid)
+
+    def _send(sig) -> bool:
+        verdict = frozen.verdict()
+        if verdict == process_lifecycle.OWNED:
+            if pgid is not None and process_lifecycle.pgid_of(frozen.pid) == pgid:
+                return _signal_session(pgid, frozen.pid, sig)
+            # Child-only: no safe group, or the leader no longer leads it.
+            return _signal_session(None, frozen.pid, sig)
+        if verdict == process_lifecycle.GONE and pgid is not None:
+            # Never fall back to the pid: it names no process of ours now.
+            return _signal_session(pgid, None, sig)
+        return False
+
+    async def _reap_leader():
+        if proc.returncode is None:
+            await proc.wait()
+
+    result = await process_lifecycle.escalate_async(
+        _gone,
+        _send,
+        steps=tuple((sig, PTY_KILL_GRACE) for sig in PTY_KILL_ESCALATION),
+        wait=_reap_leader,
+        poll_s=PTY_KILL_POLL_INTERVAL,
+        wait_floor_s=0.0,
+    )
+    return result.dead
 
 
 async def _terminate_pty_session_quietly(proc) -> None:
@@ -829,6 +860,7 @@ async def _generate_pty(cmd: str, timeout: int, request: Request):
         cwd=str(Path.home()),
         preexec_fn=os.setsid,
     )
+    _bind_pty_spawn_identity(proc)
     os.close(slave_fd)  # parent doesn't need the slave side
 
     deadline = (loop.time() + timeout) if timeout else None