refactor

Author: felixweinberger

src/mcp/client/stdio/win32.py

@@ -6,7 +6,7 @@
 import subprocess
 import sys
 from pathlib import Path
-from typing import Any, BinaryIO, TextIO, Union, cast
+from typing import BinaryIO, TextIO, Union, cast
 
 import anyio
 from anyio import to_thread
@@ -21,10 +21,12 @@
     import win32job
 else:
     # Type stubs for non-Windows platforms
-    win32api = None  # type: ignore
-    win32con = None  # type: ignore
-    win32job = None  # type: ignore
-    pywintypes = None  # type: ignore
+    win32api = None
+    win32con = None
+    win32job = None
+    pywintypes = None
+
+JobHandle = int
 
 
 def get_windows_executable_command(command: str) -> str:
@@ -157,6 +159,7 @@ async def create_windows_process(
     """
     # Create a job object for this process
     job = _create_job_object()
+    process = None
 
     try:
         # First try using anyio with Windows-specific flags to hide console window
@@ -170,28 +173,9 @@ async def create_windows_process(
             stderr=errlog,
             cwd=cwd,
         )
-
-        # Assign process to job if we have one
-        if job and hasattr(process, "pid") and win32api and win32con:
-            try:
-                # Open the process with required permissions
-                process_handle = win32api.OpenProcess(
-                    win32con.PROCESS_SET_QUOTA | win32con.PROCESS_TERMINATE, False, process.pid
-                )
-                if process_handle:
-                    _assign_process_to_job(job, process_handle)
-                    win32api.CloseHandle(process_handle)
-                    # Store job on process for later cleanup
-                    process._job_object = job  # type: ignore
-            except Exception:
-                # If we can't assign to job, close it
-                if job and win32api:
-                    win32api.CloseHandle(job)
-
-        return process
     except NotImplementedError:
         # Windows often doesn't support async subprocess creation, use fallback
-        return await _create_windows_fallback_process(command, args, env, errlog, cwd, job)
+        process = await _create_windows_fallback_process(command, args, env, errlog, cwd)
     except Exception:
         # Try again without creation flags
         process = await anyio.open_process(
@@ -201,21 +185,8 @@ async def create_windows_process(
             cwd=cwd,
         )
 
-        # Assign process to job if we have one
-        if job and hasattr(process, "pid") and win32api and win32con:
-            try:
-                process_handle = win32api.OpenProcess(
-                    win32con.PROCESS_SET_QUOTA | win32con.PROCESS_TERMINATE, False, process.pid
-                )
-                if process_handle:
-                    _assign_process_to_job(job, process_handle)
-                    win32api.CloseHandle(process_handle)
-                    process._job_object = job  # type: ignore
-            except Exception:
-                if job and win32api:
-                    win32api.CloseHandle(job)
-
-        return process
+    _maybe_assign_process_to_job(process, job)
+    return process
 
 
 async def _create_windows_fallback_process(
@@ -224,7 +195,6 @@ async def _create_windows_fallback_process(
     env: dict[str, str] | None = None,
     errlog: TextIO | None = sys.stderr,
     cwd: Path | str | None = None,
-    job: Any = None,
 ) -> FallbackProcess:
     """
     Create a subprocess using subprocess.Popen as a fallback when anyio fails.
@@ -244,24 +214,6 @@ async def _create_windows_fallback_process(
             bufsize=0,  # Unbuffered output
             creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
         )
-        process = FallbackProcess(popen_obj)
-
-        # Assign to job if provided
-        if job and win32api and win32con:
-            try:
-                process_handle = win32api.OpenProcess(
-                    win32con.PROCESS_SET_QUOTA | win32con.PROCESS_TERMINATE, False, popen_obj.pid
-                )
-                if process_handle:
-                    _assign_process_to_job(job, process_handle)
-                    win32api.CloseHandle(process_handle)
-                    process._job_object = job  # type: ignore
-            except Exception:
-                if job:
-                    win32api.CloseHandle(job)
-
-        return process
-
     except Exception:
         # If creationflags failed, fallback without them
         popen_obj = subprocess.Popen(
@@ -273,26 +225,11 @@ async def _create_windows_fallback_process(
             cwd=cwd,
             bufsize=0,
         )
-        process = FallbackProcess(popen_obj)
-
-        # Assign to job if provided
-        if job and win32api and win32con:
-            try:
-                process_handle = win32api.OpenProcess(
-                    win32con.PROCESS_SET_QUOTA | win32con.PROCESS_TERMINATE, False, popen_obj.pid
-                )
-                if process_handle:
-                    _assign_process_to_job(job, process_handle)
-                    win32api.CloseHandle(process_handle)
-                    process._job_object = job  # type: ignore
-            except Exception:
-                if job:
-                    win32api.CloseHandle(job)
+    process = FallbackProcess(popen_obj)
+    return process
 
-        return process
 
-
-def _create_job_object() -> Any:
+def _create_job_object() -> int | None:
     """
     Create a Windows Job Object configured to terminate all processes when closed.
 
@@ -321,25 +258,41 @@ def _create_job_object() -> Any:
         return None
 
 
-def _assign_process_to_job(job: Any, process_handle: int) -> bool:
+def _maybe_assign_process_to_job(process: Union[Process, "FallbackProcess"], job: JobHandle | None) -> None:
     """
-    Assign a process to a job object.
+    Try to assign a process to a job object. If assignment fails
+    for any reason, the job handle is closed.
 
     Args:
-        job: The job object handle
-        process_handle: The process handle to assign
-
-    Returns:
-        True if successful, False otherwise
+        process: The process to assign to the job
+        job: The job object handle (may be None)
     """
-    if sys.platform != "win32" or not job or not win32job:
-        return False
+    if not job:
+        return
+
+    if sys.platform != "win32" or not win32api or not win32con or not win32job:
+        return
 
     try:
-        win32job.AssignProcessToJobObject(job, process_handle)
-        return True
+        # Open the process with required permissions
+        process_handle = win32api.OpenProcess(
+            win32con.PROCESS_SET_QUOTA | win32con.PROCESS_TERMINATE, False, process.pid
+        )
+        if not process_handle:
+            raise Exception("Failed to open process handle")
+
+        try:
+            # Assign process to job
+            win32job.AssignProcessToJobObject(job, process_handle)
+            # Store job on process for later cleanup
+            process._job_object = job  # type: ignore
+        finally:
+            # Always close the process handle
+            win32api.CloseHandle(process_handle)
     except Exception:
-        return False
+        # If we can't assign to job, close it
+        if win32api:
+            win32api.CloseHandle(job)
 
 
 async def terminate_windows_process_tree(process: Union[Process, "FallbackProcess"]) -> None: