Merge pull request #14 from tech-envision/codex/create-persistent-linux-vm-before-executing-commands

README.md

@@ -5,8 +5,10 @@ and demonstrates basic tool usage. Chat histories are stored in a local SQLite
 database using Peewee. Histories are persisted per user and session so
 conversations can be resumed with context. One example tool is included:
 
-* **execute_terminal** – Executes a shell command in a Linux VM with network
-  access. Output from ``stdout`` and ``stderr`` is captured and returned.
+* **execute_terminal** – Executes a shell command inside a persistent Linux VM
+  with network access. Output from ``stdout`` and ``stderr`` is captured and
+  returned. The VM is created when a chat session starts and reused for all
+  subsequent tool calls.
 
 The application now injects a system prompt that instructs the model to chain
 multiple tools when required. This prompt ensures the assistant can orchestrate

src/__init__.py

@@ -1,4 +1,5 @@
 from .chat import ChatSession
-from .tools import execute_terminal
+from .tools import execute_terminal, set_vm
+from .vm import LinuxVM
 
-__all__ = ["ChatSession", "execute_terminal"]
+__all__ = ["ChatSession", "execute_terminal", "set_vm", "LinuxVM"]

src/chat.py

@@ -15,7 +15,8 @@ from .config import (
 from .db import Conversation, Message as DBMessage, User, _db, init_db
 from .log import get_logger
 from .schema import Msg
-from .tools import execute_terminal
+from .tools import execute_terminal, set_vm
+from .vm import LinuxVM
 
 _LOG = get_logger(__name__)
 
@@ -35,13 +36,18 @@ class ChatSession:
         self._conversation, _ = Conversation.get_or_create(
             user=self._user, session_name=session
         )
+        self._vm = LinuxVM()
         self._messages: List[Msg] = self._load_history()
         self._ensure_system_prompt()
 
     async def __aenter__(self) -> "ChatSession":
+        self._vm.start()
+        set_vm(self._vm)
         return self
 
     async def __aexit__(self, exc_type, exc, tb) -> None:
+        set_vm(None)
+        self._vm.stop()
         if not _db.is_closed():
             _db.close()
 

src/tools.py

@@ -1,8 +1,20 @@
 from __future__ import annotations
 
-__all__ = ["execute_terminal"]
+__all__ = ["execute_terminal", "set_vm"]
 
 import subprocess
+from typing import Optional
+
+from .vm import LinuxVM
+
+_VM: Optional[LinuxVM] = None
+
+
+def set_vm(vm: LinuxVM | None) -> None:
+    """Register the VM instance used for command execution."""
+
+    global _VM
+    _VM = vm
 
 
 def execute_terminal(command: str, *, timeout: int = 3) -> str:
@@ -14,6 +26,12 @@ def execute_terminal(command: str, *, timeout: int = 3) -> str:
     ``stdout`` and ``stderr`` is captured and returned. Commands are killed if
     they exceed ``timeout`` seconds.
     """
+    if _VM:
+        try:
+            return _VM.execute(command, timeout=timeout)
+        except Exception as exc:  # pragma: no cover - unforeseen errors
+            return f"Failed to execute command in VM: {exc}"
+
     try:
         completed = subprocess.run(
             command,

src/vm.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+
+from typing import Optional
+import subprocess
+import uuid
+
+from .log import get_logger
+
+_LOG = get_logger(__name__)
+
+
+class LinuxVM:
+    """Manage a lightweight Linux VM using Docker."""
+
+    def __init__(self, image: str = "ubuntu:latest") -> None:
+        self._image = image
+        self._name = f"chat-vm-{uuid.uuid4().hex[:8]}"
+        self._running = False
+
+    def start(self) -> None:
+        """Start the VM if it is not already running."""
+        if self._running:
+            return
+
+        try:
+            subprocess.run(
+                ["docker", "pull", self._image],
+                check=False,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+            )
+            subprocess.run(
+                [
+                    "docker",
+                    "run",
+                    "-d",
+                    "--name",
+                    self._name,
+                    self._image,
+                    "sleep",
+                    "infinity",
+                ],
+                check=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+            )
+            self._running = True
+        except Exception as exc:  # pragma: no cover - runtime failures
+            _LOG.error("Failed to start VM: %s", exc)
+            raise RuntimeError(f"Failed to start VM: {exc}") from exc
+
+    def execute(self, command: str, *, timeout: int = 3) -> str:
+        """Execute a command inside the running VM."""
+        if not self._running:
+            raise RuntimeError("VM is not running")
+
+        try:
+            completed = subprocess.run(
+                [
+                    "docker",
+                    "exec",
+                    self._name,
+                    "bash",
+                    "-lc",
+                    command,
+                ],
+                capture_output=True,
+                text=True,
+                timeout=timeout,
+            )
+        except subprocess.TimeoutExpired as exc:
+            return f"Command timed out after {timeout}s: {exc.cmd}"
+        except Exception as exc:  # pragma: no cover - unforeseen errors
+            return f"Failed to execute command: {exc}"
+
+        output = completed.stdout
+        if completed.stderr:
+            output = f"{output}\n{completed.stderr}" if output else completed.stderr
+        return output.strip()
+
+    def stop(self) -> None:
+        """Terminate the VM if running."""
+        if not self._running:
+            return
+
+        subprocess.run(
+            ["docker", "rm", "-f", self._name],
+            check=False,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+        self._running = False
+
+    def __enter__(self) -> "LinuxVM":
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.stop()