fix(mcp): release the death supervisor once nothing is left to reap

Follow-up to the salvaged #93517. Two gaps found in review:

- After the last unregister the supervisor stayed resident for the life of
  the process (~15 MB + a pipe) in any gateway/cron that ever connected a
  stdio server; main's per-server watchdog exited with its server. Close
  our write end when the supervised set empties: EOF with nothing
  registered makes the supervisor exit without reaping, and the next
  register already respawns and replays.
- A child that raced and exited before os.getpgid dropped its group from
  coverage entirely. The SDK spawns stdio servers as session leaders
  (pgid == pid), so fall back to the pid instead; the prune forgets the
  group once nothing in it is alive.

Tests: fake-supervisor release/re-spawn sequence, and a real-process EOF
release (supervisor exits 0, unregistered child untouched). Mutation
checked: disabling the release branch fails both.
This commit is contained in:
kshitijk4poor
2026-09-03 01:42:04 +05:30
committed by kshitij
parent b5b2ab1bb0
commit 4a2b23d77e
2 changed files with 87 additions and 2 deletions

View File

@@ -350,11 +350,25 @@ class _FakeSupervisor:
self.stdin = io.StringIO()
self.pid = 4242
self._exited = exited
self._sent = ""
self.closed = False
_real_close = self.stdin.close
def _close():
# Mirror a real pipe: capture what was written before the write
# end goes away, so tests can still assert on the control stream.
self._sent = self.stdin.getvalue()
self.closed = True
_real_close()
self.stdin.close = _close
def poll(self):
return 1 if self._exited else None
def lines(self):
if self.closed:
return self._sent.splitlines()
return self.stdin.getvalue().splitlines()
@@ -405,6 +419,55 @@ def test_unregister_is_forwarded(monkeypatch, all_groups_alive):
assert mcp_tool._supervised_pgids == set()
def test_supervisor_is_released_once_nothing_is_left_to_reap(monkeypatch, all_groups_alive):
"""An empty registration set must not keep a supervisor resident.
A gateway that once connected a stdio server would otherwise carry a
~15 MB process and a live pipe for the rest of its life. Closing our
write end is the same EOF the supervisor treats as parent death; with
nothing registered it exits without reaping. The next register starts a
fresh one, exactly like the dead-supervisor replay path.
"""
spawned = []
def _spawn():
fake = _FakeSupervisor()
spawned.append(fake)
return fake
monkeypatch.setattr(mcp_tool, "_spawn_death_supervisor", _spawn)
mcp_tool._update_death_supervisor("register", [111, 222])
mcp_tool._update_death_supervisor("unregister", [111])
assert not spawned[0].closed, "released the supervisor while a group was still registered"
mcp_tool._update_death_supervisor("unregister", [222])
assert spawned[0].closed, "supervisor kept resident with nothing left to reap"
assert spawned[0].lines()[-1] == "unregister 222", "release happened before the last unregister was sent"
assert mcp_tool._death_supervisor is None
mcp_tool._update_death_supervisor("register", [333])
assert len(spawned) == 2 and spawned[1].lines() == ["register 333"]
def test_supervisor_survives_the_real_eof_release():
"""End to end: closing the control pipe with nothing registered exits cleanly."""
if os.name != "posix":
pytest.skip("POSIX-only supervisor")
child = subprocess.Popen(_VICTIM, start_new_session=True)
try:
mcp_tool._update_death_supervisor("register", [os.getpgid(child.pid)])
proc = mcp_tool._death_supervisor
assert proc is not None and proc.poll() is None
mcp_tool._update_death_supervisor("unregister", [os.getpgid(child.pid)])
assert mcp_tool._death_supervisor is None
assert proc.wait(timeout=10) == 0, "supervisor did not exit on the release EOF"
assert child.poll() is None, "release reaped a group that had been unregistered"
finally:
_kill(child.pid)
child.wait(timeout=10)
def test_unregister_alone_does_not_start_a_supervisor(monkeypatch):
spawned = []
monkeypatch.setattr(

View File

@@ -1210,6 +1210,20 @@ def _update_death_supervisor(verb: str, pgids) -> None:
# ``_supervised_pgids``. Nothing is lost in between because that
# set, not the pipe, is the record of what needs reaping.
_death_supervisor = None
return
if not _supervised_pgids:
# Nothing left to reap: release the supervisor instead of keeping
# a ~15 MB process and a pipe resident for the life of a gateway
# that once connected a stdio server. Closing our write end is
# the same EOF signal parent death sends; with an empty set the
# supervisor reaps nothing and exits. The next register respawns
# and replays from ``_supervised_pgids`` as it already does.
try:
proc.stdin.close()
except (BrokenPipeError, ValueError, OSError):
pass
_death_supervisor = None
# ---------------------------------------------------------------------------
@@ -3440,9 +3454,17 @@ class MCPServerTask:
for _pid in new_pids:
try:
new_pgids[_pid] = os.getpgid(_pid)
except (AttributeError, ProcessLookupError, OSError):
except ProcessLookupError:
# The child raced and already exited. The MCP SDK
# spawns stdio servers with start_new_session=True,
# so the child was its own group leader (pgid ==
# pid); keep that group covered rather than drop
# it -- any descendant it left behind still has
# to be reaped, and the prune forgets the group
# once nothing in it is alive.
new_pgids[_pid] = _pid
except (AttributeError, OSError):
# AttributeError: Windows (os.getpgid is POSIX-only)
# ProcessLookupError: child raced and already exited
pass
with _lock:
for _pid in new_pids: