fix: ensure child processes are killed when terminating sandbox

src/deno_sandbox/sandbox.py

@@ -98,7 +98,7 @@ def create(self, options: Optional[SandboxCreateOptions] = None):
         rpc = RpcClient(async_handle._rpc, self._bridge)
 
         try:
-            yield Sandbox(self._client, rpc, async_handle.id)
+            yield Sandbox(self._client, rpc, async_handle)
         except Exception:
             import sys
 
@@ -115,7 +115,7 @@ def connect(self, options: SandboxConnectOptions):
         rpc = RpcClient(async_handle._rpc, self._bridge)
 
         try:
-            yield Sandbox(self._client, rpc, async_handle.id)
+            yield Sandbox(self._client, rpc, async_handle)
         except Exception:
             import sys
 
@@ -253,9 +253,15 @@ class VsCodeOptions(TypedDict):
 
 
 class AsyncSandboxDeno:
-    def __init__(self, client: AsyncConsoleClient, rpc: AsyncRpcClient):
+    def __init__(
+        self,
+        client: AsyncConsoleClient,
+        rpc: AsyncRpcClient,
+        processes: list[AsyncChildProcess],
+    ):
         self._client = client
         self._rpc = rpc
+        self._processes = processes
 
     async def run(self, options: DenoRunOptions) -> AsyncDenoProcess:
         """Create a new Deno process from the specified entrypoint file or code. The runtime will execute the given code to completion, and then exit."""
@@ -285,7 +291,11 @@ async def run(self, options: DenoRunOptions) -> AsyncDenoProcess:
 
         result = await self._rpc.call("spawnDeno", params)
 
-        return await AsyncDenoProcess.create(result, self._rpc, opts)
+        process = await AsyncDenoProcess.create(
+            result, self._rpc, opts, self._processes
+        )
+        self._processes.append(process)
+        return process
 
     async def eval(self, code: str) -> Any:
         repl = await self.repl()
@@ -312,15 +322,24 @@ async def repl(self, options: Optional[DenoReplOptions] = None) -> AsyncDenoRepl
 
         result: ProcessSpawnResult = await self._rpc.call("spawnDenoRepl", params)
 
-        return await AsyncDenoRepl.create(result, self._rpc, opts)
+        process = await AsyncDenoRepl.create(result, self._rpc, opts, self._processes)
+        self._processes.append(process)
+        return process
 
 
 class SandboxDeno:
-    def __init__(self, client: ConsoleClient, rpc: RpcClient):
+    def __init__(
+        self,
+        client: ConsoleClient,
+        rpc: RpcClient,
+        processes: list[AsyncChildProcess],
+    ):
         self._client = client
         self._rpc = rpc
 
-        self._async = AsyncSandboxDeno(self._client._async, rpc._async_client)
+        self._async = AsyncSandboxDeno(
+            self._client._async, rpc._async_client, processes
+        )
 
     def run(self, options: DenoRunOptions) -> DenoProcess:
         """
@@ -343,12 +362,13 @@ def __init__(
     ):
         self._client = client
         self._rpc = rpc
+        self._processes: list[AsyncChildProcess] = []
 
         self.url: str | None = None
         self.ssh: None = None
         self.id = sandbox_id
         self.fs = AsyncSandboxFs(client, rpc)
-        self.deno = AsyncSandboxDeno(client, rpc)
+        self.deno = AsyncSandboxDeno(client, rpc, self._processes)
         self.env = AsyncSandboxEnv(client, rpc)
 
     @property
@@ -380,7 +400,11 @@ async def spawn(
             params["stderr"] = "piped"
 
         result: ProcessSpawnResult = await self._rpc.call("spawn", params)
-        return await AsyncChildProcess.create(result, self._rpc, opts)
+        process = await AsyncChildProcess.create(
+            result, self._rpc, opts, self._processes
+        )
+        self._processes.append(process)
+        return process
 
     async def fetch(
         self,
@@ -392,6 +416,10 @@ async def fetch(
         return await self._rpc.fetch(url, method, headers, redirect)
 
     async def close(self) -> None:
+        # Kill all tracked processes
+        for process in self._processes:
+            await process.kill()
+        self._processes.clear()
         await self._rpc.close()
 
     async def kill(self) -> None:
@@ -458,16 +486,18 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
 
 
 class Sandbox:
-    def __init__(self, client: ConsoleClient, rpc: RpcClient, sandbox_id: str):
+    def __init__(
+        self, client: ConsoleClient, rpc: RpcClient, async_sandbox: AsyncSandbox
+    ):
         self._client = client
         self._rpc = rpc
-        self._async = AsyncSandbox(self._client._async, rpc._async_client, sandbox_id)
+        self._async = async_sandbox
 
         self.url: str | None = None
         self.ssh: None = None
-        self.id = sandbox_id
+        self.id = async_sandbox.id
         self.fs = SandboxFs(client, rpc)
-        self.deno = SandboxDeno(client, rpc)
+        self.deno = SandboxDeno(client, rpc, self._async._processes)
         self.env = SandboxEnv(client, rpc)
 
     @property

src/deno_sandbox/wrappers.py

@@ -226,6 +226,7 @@ def __init__(
         rpc: AsyncRpcClient,
         _stdout_task: Optional[asyncio.Task] = None,
         _stderr_task: Optional[asyncio.Task] = None,
+        _process_list: Optional[list["AsyncChildProcess"]] = None,
     ):
         self.pid = pid
         self.stdout = stdout
@@ -235,19 +236,27 @@ def __init__(
         self._rpc = rpc
         self._stdout_task = _stdout_task
         self._stderr_task = _stderr_task
+        self._process_list = _process_list
+
+    def _remove_from_list(self) -> None:
+        """Remove this process from the tracking list."""
+        if self._process_list is not None and self in self._process_list:
+            self._process_list.remove(self)
 
     @classmethod
     async def create(
         cls: type["AsyncChildProcess"],
         res: ProcessSpawnResult,
         rpc: AsyncRpcClient,
         options: RemoteProcessOptions,
+        process_list: Optional[list[AsyncChildProcess]] = None,
     ) -> AsyncChildProcess:
-        return create_process_like(cls, res, rpc, options)
+        return create_process_like(cls, res, rpc, options, process_list)
 
     async def wait(self) -> ChildProcessStatus:
         raw = await self._wait_task
         result = cast(ProcessWaitResult, raw)
+        self._remove_from_list()
         return ChildProcessStatus(
             success=result["success"], code=result["code"], signal=result["signal"]
         )
@@ -268,6 +277,8 @@ async def kill(self) -> None:
         if self._stderr_task is not None:
             self._stderr_task.cancel()
 
+        self._remove_from_list()
+
     async def __aenter__(self):
         return self
 
@@ -300,12 +311,14 @@ def create_process_like[T](
             AsyncRpcClient,
             Optional[asyncio.Task],
             Optional[asyncio.Task],
+            Optional[list[AsyncChildProcess]],
         ],
         T,
     ],
     res: ProcessSpawnResult,
     rpc: AsyncRpcClient,
     options: RemoteProcessOptions,
+    process_list: Optional[list[AsyncChildProcess]] = None,
 ) -> T:
     pid = res["pid"]
 
@@ -327,7 +340,9 @@ def create_process_like[T](
         stderr_coro = _pipe_stream(stderr, sys.stderr.buffer)
         stderr_task = rpc._loop.create_task(stderr_coro)
 
-    instance = cls(pid, stdout, stderr, wait_task, rpc, stdout_task, stderr_task)
+    instance = cls(
+        pid, stdout, stderr, wait_task, rpc, stdout_task, stderr_task, process_list
+    )
 
     return instance
 
@@ -386,8 +401,11 @@ def __init__(
         rpc: AsyncRpcClient,
         stdout_task: Optional[asyncio.Task],
         stderr_task: Optional[asyncio.Task],
+        process_list: Optional[list[AsyncChildProcess]] = None,
     ):
-        super().__init__(pid, stdout, stderr, wait_task, rpc, stdout_task, stderr_task)
+        super().__init__(
+            pid, stdout, stderr, wait_task, rpc, stdout_task, stderr_task, process_list
+        )
         self._listening_task: asyncio.Task | None = None
 
     @classmethod
@@ -396,8 +414,9 @@ async def create(
         res: ProcessSpawnResult,
         rpc: AsyncRpcClient,
         options: RemoteProcessOptions,
+        process_list: Optional[list[AsyncChildProcess]] = None,
     ) -> AsyncDenoProcess:
-        p = create_process_like(cls, res, rpc, options)
+        p = create_process_like(cls, res, rpc, options, process_list)
 
         p._listening_task = rpc._loop.create_task(
             rpc.call("denoHttpWait", {"pid": p.pid})
@@ -467,8 +486,9 @@ async def create(
         res: ProcessSpawnResult,
         rpc: AsyncRpcClient,
         options: RemoteProcessOptions,
+        process_list: Optional[list[AsyncChildProcess]] = None,
     ) -> AsyncDenoRepl:
-        return create_process_like(cls, res, rpc, options)
+        return create_process_like(cls, res, rpc, options, process_list)
 
     async def eval(self, code: str) -> str:
         """Evaluate code in the REPL and return the output."""
@@ -495,6 +515,7 @@ async def close(self) -> None:
             self._stdout_task.cancel()
         if self._stderr_task is not None:
             self._stderr_task.cancel()
+        self._remove_from_list()
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         await self.close()