Keep a still-diverted fd claimed when its restore fails

A failed exit-time restore is still swallowed so it never masks what
ended the transport, but the fd now stays in the claimed set: a later
stdio_server() in that process is refused instead of claiming the
diversion and serving the null device as its wire.

Author: maxisbey

src/mcp/server/stdio.py

@@ -101,18 +101,23 @@ def restore() -> None:
         # Drain text buffered during the claim (a stray print) to the diversion.
         with suppress(OSError, ValueError):
             stream.flush()
-        _restore_fd(fd, private_fd)
-        unclaim()
+        # A failed restore leaves fd diverted; keep it claimed so later transports are refused.
+        if _restore_fd(fd, private_fd):
+            unclaim()
 
     # closefd=False: a blocked worker thread may still read this descriptor after exit.
     return os.fdopen(private_fd, mode, closefd=False), restore
 
 
-def _restore_fd(fd: int, private_fd: int) -> None:
-    with suppress(OSError):
+def _restore_fd(fd: int, private_fd: int) -> bool:
+    """Point fd back at the protocol pipe; a failure never masks the transport's exit."""
+    try:
         os.dup2(private_fd, fd)
         if sys.platform == "win32":  # pragma: no cover
             rebind_std_handle_to_fd(fd)
+    except OSError:
+        return False
+    return True
 
 
 @asynccontextmanager

tests/server/test_stdio.py

@@ -273,11 +273,14 @@ def failing_dup2(fd: int, fd2: int, inheritable: bool = True) -> int:
 async def test_stdio_server_exits_cleanly_when_the_stdin_restore_fails(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
-    """A failed fd 0 restore on exit is swallowed, not raised.
+    """A failed fd 0 restore on exit is swallowed, not raised, and the fd stays claimed.
 
-    SDK-defined behavior: the restore in the finally must never mask what ended the transport.
+    SDK-defined behavior: the restore must never mask what ended the transport, and a
+    still-diverted fd must refuse later transports rather than serve them the diversion.
     """
     request = JSONRPCRequest(jsonrpc="2.0", id=1, method="ping")
+    fresh_claims: set[int] = set()
+    monkeypatch.setattr("mcp.server.stdio._claimed", fresh_claims)  # this test leaves fd 0 claimed
     with _pipe_planted_on_fd0(monkeypatch) as (_, in_w):
         os.write(in_w, _frame(request))
         os.close(in_w)
@@ -311,6 +314,12 @@ def flaky_dup2(fd: int, fd2: int, inheritable: bool = True) -> int:
         finally:
             os.close(devnull_probe)
 
+        # The still-diverted fd stays claimed: a later transport is refused,
+        # not handed the null device as its wire.
+        with pytest.raises(RuntimeError, match="already claimed fd 0"):
+            async with stdio_server():
+                pytest.fail("unreachable")  # pragma: no cover
+
 
 @pytest.mark.anyio
 async def test_stdio_server_takes_stdout_off_the_descriptor_table_while_serving(