From 8413987fcde4c5d58eb41e6b619f34afd509a08c Mon Sep 17 00:00:00 2001
From: 0xallam <ahmed39652003@gmail.com>
Date: Mon, 19 Jan 2026 18:21:34 -0800
Subject: [PATCH] feat: remove docker container on shutdown

Add automatic cleanup of Docker containers when the application exits.
Uses a singleton runtime pattern and spawns a detached subprocess for
cleanup to ensure fast exit without blocking the UI.
---
 strix/interface/cli.py          |  3 +++
 strix/interface/tui.py          |  3 +++
 strix/runtime/__init__.py       | 19 +++++++++++++++++--
 strix/runtime/docker_runtime.py | 19 +++++++++++++++++++
 strix/runtime/runtime.py        |  3 +++
 5 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/strix/interface/cli.py b/strix/interface/cli.py
index ef99f5e..4b5d109 100644
--- a/strix/interface/cli.py
+++ b/strix/interface/cli.py
@@ -106,7 +106,10 @@ async def run_cli(args: Any) -> None:  # noqa: PLR0915
     tracer.vulnerability_found_callback = display_vulnerability
 
     def cleanup_on_exit() -> None:
+        from strix.runtime import cleanup_runtime
+
         tracer.cleanup()
+        cleanup_runtime()
 
     def signal_handler(_signum: int, _frame: Any) -> None:
         tracer.cleanup()
diff --git a/strix/interface/tui.py b/strix/interface/tui.py
index 38aaad8..e360ab0 100644
--- a/strix/interface/tui.py
+++ b/strix/interface/tui.py
@@ -742,7 +742,10 @@ class StrixTUIApp(App):  # type: ignore[misc]
 
     def _setup_cleanup_handlers(self) -> None:
         def cleanup_on_exit() -> None:
+            from strix.runtime import cleanup_runtime
+
             self.tracer.cleanup()
+            cleanup_runtime()
 
         def signal_handler(_signum: int, _frame: Any) -> None:
             self.tracer.cleanup()
diff --git a/strix/runtime/__init__.py b/strix/runtime/__init__.py
index ae01f38..5d0cbda 100644
--- a/strix/runtime/__init__.py
+++ b/strix/runtime/__init__.py
@@ -12,17 +12,32 @@ class SandboxInitializationError(Exception):
         self.details = details
 
 
+_global_runtime: AbstractRuntime | None = None
+
+
 def get_runtime() -> AbstractRuntime:
+    global _global_runtime  # noqa: PLW0603
+
     runtime_backend = Config.get("strix_runtime_backend")
 
     if runtime_backend == "docker":
         from .docker_runtime import DockerRuntime
 
-        return DockerRuntime()
+        if _global_runtime is None:
+            _global_runtime = DockerRuntime()
+        return _global_runtime
 
     raise ValueError(
         f"Unsupported runtime backend: {runtime_backend}. Only 'docker' is supported for now."
     )
 
 
-__all__ = ["AbstractRuntime", "SandboxInitializationError", "get_runtime"]
+def cleanup_runtime() -> None:
+    global _global_runtime  # noqa: PLW0603
+
+    if _global_runtime is not None:
+        _global_runtime.cleanup()
+        _global_runtime = None
+
+
+__all__ = ["AbstractRuntime", "SandboxInitializationError", "cleanup_runtime", "get_runtime"]
diff --git a/strix/runtime/docker_runtime.py b/strix/runtime/docker_runtime.py
index e42d00b..b783dcc 100644
--- a/strix/runtime/docker_runtime.py
+++ b/strix/runtime/docker_runtime.py
@@ -316,3 +316,22 @@ class DockerRuntime(AbstractRuntime):
             self._tool_server_token = None
         except (NotFound, DockerException):
             pass
+
+    def cleanup(self) -> None:
+        if self._scan_container is not None:
+            container_name = self._scan_container.name
+            self._scan_container = None
+            self._tool_server_port = None
+            self._tool_server_token = None
+
+            if container_name is None:
+                return
+
+            import subprocess
+
+            subprocess.Popen(  # noqa: S603
+                ["docker", "rm", "-f", container_name],  # noqa: S607
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+                start_new_session=True,
+            )
diff --git a/strix/runtime/runtime.py b/strix/runtime/runtime.py
index 9de1afe..e33c08d 100644
--- a/strix/runtime/runtime.py
+++ b/strix/runtime/runtime.py
@@ -27,3 +27,6 @@ class AbstractRuntime(ABC):
     @abstractmethod
     async def destroy_sandbox(self, container_id: str) -> None:
         raise NotImplementedError
+
+    def cleanup(self) -> None:
+        raise NotImplementedError