From 175167e091b74fc435beaefde77bd30c88499272 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= <erik@bjareho.lt>
Date: Fri, 1 Nov 2024 13:45:40 +0100
Subject: [PATCH] feat: implement anthropic-style computer tool (#225)

* feat: started working on anthropic-style computer tool

* Apply suggestions from code review

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

* fix: progress on computer use

* fix: added Dockerfile.server

* fix: fixed vnc in computer use webui

* docs: fixed docs for computer use

* fix: rewrote computer_action function to not be a generator

* docs: fixed server docs for computer use

* docs: refactored computer use warning into seperate file

* fix: optimized Dockerfile.computer for faster rebuilds

* fix: refactor and misc fixes to computer use

* fix: enable select tools in computer use context

* fix: multiple fixes to computer use and web ui

* fix: disable computer tool unless explicitly enabled

* fix: removed deleted file from .dockerignore

* docs: minor fix to computer use docs

---------

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>
---
 .dockerignore                                 |   1 +
 Makefile                                      |   5 +
 docs/computer-use-warning.rst                 |   4 +
 docs/conf.py                                  |   1 +
 docs/contributing.rst                         |   4 +-
 docs/examples.rst                             |  27 +++
 docs/server.rst                               |  47 +++-
 docs/tools.rst                                |  36 +++
 gptme/server/api.py                           |  12 +
 gptme/server/cli.py                           |  28 ++-
 gptme/server/static/computer.html             |  92 ++++++++
 gptme/server/static/index.html                |   3 +-
 gptme/server/static/main.js                   |  67 ++++--
 gptme/tools/__init__.py                       |  13 +-
 gptme/tools/computer.py                       | 219 ++++++++++++++++++
 gptme/tools/python.py                         |   4 +-
 gptme/tools/screenshot.py                     |   2 +-
 gptme/tools/vision.py                         |   2 +
 poetry.lock                                   |  19 +-
 pyproject.toml                                |   4 +
 scripts/Dockerfile                            |  13 --
 scripts/Dockerfile.computer                   | 102 ++++++++
 scripts/Dockerfile.server                     |  23 ++
 .../tint2/applications/firefox-custom.desktop |   8 +
 .../.config/tint2/applications/gedit.desktop  |   8 +
 .../tint2/applications/terminal.desktop       |   8 +
 scripts/computer_home/.config/tint2/tint2rc   | 100 ++++++++
 scripts/computer_home/entrypoint.sh           |  11 +
 scripts/computer_home/mutter_startup.sh       |  22 ++
 scripts/computer_home/novnc_startup.sh        |  29 +++
 scripts/computer_home/start_all.sh            |   8 +
 scripts/computer_home/tint2_startup.sh        |  22 ++
 scripts/computer_home/x11vnc_startup.sh       |  45 ++++
 scripts/computer_home/xvfb_startup.sh         |  23 ++
 34 files changed, 965 insertions(+), 47 deletions(-)
 create mode 100644 docs/computer-use-warning.rst
 create mode 100644 gptme/server/static/computer.html
 create mode 100644 gptme/tools/computer.py
 create mode 100644 scripts/Dockerfile.computer
 create mode 100644 scripts/Dockerfile.server
 create mode 100755 scripts/computer_home/.config/tint2/applications/firefox-custom.desktop
 create mode 100755 scripts/computer_home/.config/tint2/applications/gedit.desktop
 create mode 100644 scripts/computer_home/.config/tint2/applications/terminal.desktop
 create mode 100644 scripts/computer_home/.config/tint2/tint2rc
 create mode 100755 scripts/computer_home/entrypoint.sh
 create mode 100755 scripts/computer_home/mutter_startup.sh
 create mode 100755 scripts/computer_home/novnc_startup.sh
 create mode 100755 scripts/computer_home/start_all.sh
 create mode 100755 scripts/computer_home/tint2_startup.sh
 create mode 100755 scripts/computer_home/x11vnc_startup.sh
 create mode 100755 scripts/computer_home/xvfb_startup.sh

diff --git a/.dockerignore b/.dockerignore
index f3aff7dc..3fa7309b 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -10,6 +10,7 @@ gptme.toml
 
 # Build scripts
 scripts
+!scripts/computer_home
 .github
 
 # Build/test/coverage/docs/prof directories
diff --git a/Makefile b/Makefile
index 57464325..23a3ff5d 100644
--- a/Makefile
+++ b/Makefile
@@ -16,7 +16,12 @@ build:
 
 build-docker:
 	docker build . -t gptme:latest -f scripts/Dockerfile
+	docker build . -t gptme-server:latest -f scripts/Dockerfile.server
 	docker build . -t gptme-eval:latest -f scripts/Dockerfile.eval
+	# docker build . -t gptme-eval:latest -f scripts/Dockerfile.eval --build-arg RUST=yes --build-arg BROWSER=yes
+
+build-docker-computer:
+	docker build . -t gptme-computer:latest -f scripts/Dockerfile.computer
 
 build-docker-dev:
 	docker build . -t gptme-dev:latest -f scripts/Dockerfile.dev
diff --git a/docs/computer-use-warning.rst b/docs/computer-use-warning.rst
new file mode 100644
index 00000000..d90f162d
--- /dev/null
+++ b/docs/computer-use-warning.rst
@@ -0,0 +1,4 @@
+.. warning::
+
+   The computer use interface is experimental and has serious security implications.
+   Please use with caution and see Anthropic's documentation on `computer use <https://docs.anthropic.com/en/docs/build-with-claude/computer-use>`_ for additional guidance.
diff --git a/docs/conf.py b/docs/conf.py
index 8dbed0c1..f8617bea 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -118,6 +118,7 @@ def setup(app):
     ("py:class", "flask.app.Flask"),
     ("py:class", "gptme.tools.python.T"),
     ("py:class", "threading.Thread"),
+    ("py:class", "gptme.tools.computer.ScalingSource"),
 ]
 
 # -- Options for HTML output -------------------------------------------------
diff --git a/docs/contributing.rst b/docs/contributing.rst
index 03977a1a..abb48644 100644
--- a/docs/contributing.rst
+++ b/docs/contributing.rst
@@ -17,10 +17,10 @@ Install
    # checkout the code and navigate to the root of the project
    git clone https://github.com/ErikBjare/gptme.git
    cd gptme
-   
+
    # install poetry (if not installed)
    pipx install poetry
-   
+
    # activate the virtualenv
    poetry shell
 
diff --git a/docs/examples.rst b/docs/examples.rst
index 34bb0ad6..72afdd71 100644
--- a/docs/examples.rst
+++ b/docs/examples.rst
@@ -64,3 +64,30 @@ Generate docstrings for all functions in a file:
    gptme --non-interactive "Patch these files to include concise docstrings for all functions, skip functions that already have docstrings. Include: brief description, parameters." $@
 
 These examples demonstrate how gptme can be used to create simple yet powerful automation tools. Each script can be easily customized and expanded to fit specific project needs.
+
+.. rubric:: Computer Use Examples
+
+Using the computer tool for GUI automation and desktop interaction (requires running the server with computer use support):
+
+.. code-block:: bash
+
+    # Start server with computer use support
+    docker run -p 5000:5000 -p 8080:8080 -p 6080:6080 ghcr.io/erikbjare/gptme:latest-server
+
+    # Then in another terminal:
+
+    # Open and interact with an application
+    gptme 'open firefox and navigate to example.com'
+
+    # GUI automation with visual feedback
+    gptme 'create a simple drawing in xpaint'
+
+    # Desktop automation with keyboard/mouse
+    gptme 'open calculator and compute 15 * 23'
+
+The computer use interface at http://localhost:8080 provides a split view with:
+- Chat interface on the left
+- Desktop view on the right
+- Controls for toggling interaction mode
+
+This enables complex GUI automation tasks with visual feedback and confirmation.
diff --git a/docs/server.rst b/docs/server.rst
index 8e33665a..76c3c1cb 100644
--- a/docs/server.rst
+++ b/docs/server.rst
@@ -16,12 +16,49 @@ It can be started by running the following command:
 Web UI
 ------
 
-.. code-block:: bash
+The server provides two interfaces:
 
-    gptme-server
+1. Basic Chat Interface
+
+   .. code-block:: bash
+
+       gptme-server
+
+   Access the basic chat interface at http://localhost:5000
+
+   For more usage, see :ref:`the CLI documentation <cli:gptme-server>`.
+
+2. Computer Use Interface
+
+   Requires Docker.
+
+   .. code-block:: bash
+
+       # Clone the repository
+       git clone https://github.com/ErikBjare/gptme.git
+       cd gptme
+       # Build container
+       make build-docker-computer
+       # Run container
+       docker run -v ~/.config/gptme:/home/computeruse/.config/gptme -p 6080:6080 -p 8080:8080 gptme-computer:latest
+
+   The computer use interface provides:
+
+   - Combined view at http://localhost:8080/computer
+   - Chat view at http://localhost:8080
+   - Desktop view at http://localhost:6080/vnc.html
+
+   Features:
+
+   - Split view with chat on the left, desktop on the right
+   - Toggle for view-only/interactive desktop mode
+   - Fullscreen support
+   - Automatic screen scaling for optimal LLM vision
 
-This should let you view your chats in a web browser and make basic requests.
+   Requirements:
 
-You can then access the web UI by visiting http://localhost:5000 in your browser.
+   - Docker for running the server with X11 support
+   - Browser with WebSocket support for VNC
+   - Network ports 6080 (VNC) and 8080 (web UI) available
 
-For more usage, see :ref:`the CLI documentation <cli:gptme-server>`.
+.. include:: computer-use-warning.rst
diff --git a/docs/tools.rst b/docs/tools.rst
index 36f55a9b..05599798 100644
--- a/docs/tools.rst
+++ b/docs/tools.rst
@@ -26,6 +26,7 @@ The tools can be grouped into the following categories:
 
   - `Screenshot`_
   - `Vision`_
+  - `Computer`_
 
 - Chat management
 
@@ -107,3 +108,38 @@ Chats
 .. automodule:: gptme.tools.chats
     :members:
     :noindex:
+
+Computer
+--------
+
+.. automodule:: gptme.tools.computer
+    :members:
+    :noindex:
+
+The computer tool provides direct interaction with the desktop environment through X11, allowing for:
+
+- Keyboard input simulation
+- Mouse control (movement, clicks, dragging)
+- Screen capture with automatic scaling
+- Cursor position tracking
+
+To use the computer tool, see the instructions for :doc:`server`.
+
+Example usage::
+
+    # Type text
+    computer(action="type", text="Hello, World!")
+
+    # Move mouse and click
+    computer(action="mouse_move", coordinate=(100, 100))
+    computer(action="left_click")
+
+    # Take screenshot
+    computer(action="screenshot")
+
+    # Send keyboard shortcuts
+    computer(action="key", text="Control_L+c")
+
+The tool automatically handles screen resolution scaling to ensure optimal performance with LLM vision capabilities.
+
+.. include:: computer-use-warning.rst
diff --git a/gptme/server/api.py b/gptme/server/api.py
index fda5538e..72bb9eef 100644
--- a/gptme/server/api.py
+++ b/gptme/server/api.py
@@ -148,6 +148,18 @@ def root():
     return current_app.send_static_file("index.html")
 
 
+# serve computer interface
+@api.route("/computer")
+def computer():
+    return current_app.send_static_file("computer.html")
+
+
+# serve chat interface (for embedding in computer view)
+@api.route("/chat")
+def chat():
+    return current_app.send_static_file("index.html")
+
+
 @api.route("/favicon.png")
 def favicon():
     return flask.send_from_directory(media_path, "logo.png")
diff --git a/gptme/server/cli.py b/gptme/server/cli.py
index 08e1474d..f0decfdf 100644
--- a/gptme/server/cli.py
+++ b/gptme/server/cli.py
@@ -16,14 +16,36 @@
     default=None,
     help="Model to use by default, can be overridden in each request.",
 )
-def main(debug: bool, verbose: bool, model: str | None):  # pragma: no cover
+@click.option(
+    "--host",
+    default="127.0.0.1",
+    help="Host to bind the server to.",
+)
+@click.option(
+    "--port",
+    default="5000",
+    help="Port to run the server on.",
+)
+@click.option("--tools", default=None, help="Tools to enable, comma separated.")
+def main(
+    debug: bool,
+    verbose: bool,
+    model: str | None,
+    host: str,
+    port: str,
+    tools: str | None,
+):  # pragma: no cover
     """
     Starts a server and web UI for gptme.
 
     Note that this is very much a work in progress, and is not yet ready for normal use.
     """
     init_logging(verbose)
-    init(model, interactive=False, tool_allowlist=None)
+    init(
+        model,
+        interactive=False,
+        tool_allowlist=None if tools is None else tools.split(","),
+    )
 
     # if flask not installed, ask the user to install `server` extras
     try:
@@ -37,4 +59,4 @@ def main(debug: bool, verbose: bool, model: str | None):  # pragma: no cover
     click.echo("Initialization complete, starting server")
 
     app = create_app()
-    app.run(debug=debug)
+    app.run(debug=debug, host=host, port=int(port))
diff --git a/gptme/server/static/computer.html b/gptme/server/static/computer.html
new file mode 100644
index 00000000..33994147
--- /dev/null
+++ b/gptme/server/static/computer.html
@@ -0,0 +1,92 @@
+<!DOCTYPE html>
+<html>
+<head>
+    <title>gptme - Computer Use</title>
+    <meta name="permissions-policy" content="fullscreen=*" />
+    <style>
+        body {
+            margin: 0;
+            padding: 0;
+            overflow: hidden;
+            font-family: system-ui, -apple-system, sans-serif;
+        }
+        .container {
+            display: flex;
+            height: 100vh;
+            width: 100vw;
+        }
+        .chat {
+            flex: 1;
+            border: none;
+            height: 100vh;
+            background: #f5f5f5;
+        }
+        .desktop {
+            flex: 2;
+            border: none;
+            height: 100vh;
+        }
+        .controls {
+            position: absolute;
+            top: 10px;
+            right: 10px;
+            z-index: 1000;
+            display: flex;
+            gap: 10px;
+        }
+        button {
+            padding: 8px 16px;
+            border-radius: 4px;
+            border: 1px solid #ccc;
+            background: white;
+            cursor: pointer;
+            font-size: 14px;
+        }
+        button:hover {
+            background: #f0f0f0;
+        }
+    </style>
+</head>
+<body>
+    <div class="container">
+        <div class="chat">
+            <!-- Will be replaced with gptme chat interface -->
+            <iframe src="/chat" style="width: 100%; height: 100%; border: none;"></iframe>
+        </div>
+        <iframe
+            id="vnc"
+            class="desktop"
+            src="http://127.0.0.1:6080/vnc.html?&resize=scale&autoconnect=1&view_only=1&reconnect=1&reconnect_delay=2000"
+            allow="fullscreen"
+        ></iframe>
+    </div>
+    <div class="controls">
+        <button id="toggleViewOnly">Toggle Screen Control (Off)</button>
+        <button id="toggleFullscreen">Toggle Fullscreen</button>
+    </div>
+    <script>
+        // Toggle view-only mode
+        document.getElementById("toggleViewOnly").addEventListener("click", function() {
+            var vncIframe = document.getElementById("vnc");
+            var button = document.getElementById("toggleViewOnly");
+            var currentSrc = vncIframe.src;
+            if (currentSrc.includes("view_only=1")) {
+                vncIframe.src = currentSrc.replace("view_only=1", "view_only=0");
+                button.innerText = "Toggle Screen Control (On)";
+            } else {
+                vncIframe.src = currentSrc.replace("view_only=0", "view_only=1");
+                button.innerText = "Toggle Screen Control (Off)";
+            }
+        });
+
+        // Toggle fullscreen
+        document.getElementById("toggleFullscreen").addEventListener("click", function() {
+            if (!document.fullscreenElement) {
+                document.documentElement.requestFullscreen();
+            } else {
+                document.exitFullscreen();
+            }
+        });
+    </script>
+</body>
+</html>
diff --git a/gptme/server/static/index.html b/gptme/server/static/index.html
index 9e1263a3..8df9f719 100644
--- a/gptme/server/static/index.html
+++ b/gptme/server/static/index.html
@@ -151,8 +151,9 @@ <h1 class="text-2xl font-bold text-gray-800 mb-2 md:mb-0">{{ selectedConversatio
                   <textarea
                     class="flex-grow border rounded-lg p-3 mb-2 md:mb-0 md:mr-2 focus:outline-none focus:ring-2 focus:ring-blue-500 resize-none"
                     v-model="newMessage"
-                    placeholder="Type your message"
+                    placeholder="Type your message (Enter to send, Shift+Enter for newline)"
                     rows="3"
+                    @keydown="handleKeyDown"
                     ></textarea>
                   <button
                     type="submit"
diff --git a/gptme/server/static/main.js b/gptme/server/static/main.js
index b9396826..91bddada 100644
--- a/gptme/server/static/main.js
+++ b/gptme/server/static/main.js
@@ -190,29 +190,52 @@ new Vue({
       this.selectConversation(name);
     },
     async sendMessage() {
+      const messageContent = this.newMessage;
+      // Clear input immediately
+      this.newMessage = "";
+
+      // Add message to chat log immediately
+      const tempMessage = {
+        role: "user",
+        content: messageContent,
+        timestamp: new Date().toISOString(),
+        html: this.mdToHtml(messageContent)
+      };
+      this.chatLog.push(tempMessage);
+      this.scrollToBottom();
+
+      // Send to server
       const payload = JSON.stringify({
         role: "user",
-        content: this.newMessage,
+        content: messageContent,
         branch: this.branch,
       });
-      const req = await fetch(`${apiRoot}/${this.selectedConversation}`, {
-        method: "POST",
-        headers: {
-          "Content-Type": "application/json",
-        },
-        body: payload,
-      });
-      if (!req.ok) {
-        this.error = req.statusText;
-        return;
+
+      try {
+        const req = await fetch(`${apiRoot}/${this.selectedConversation}`, {
+          method: "POST",
+          headers: {
+            "Content-Type": "application/json",
+          },
+          body: payload,
+        });
+
+        if (!req.ok) {
+            throw new Error(req.statusText);
+        }
+
+        await req.json();
+        // Reload conversation to get server-side state
+        await this.selectConversation(this.selectedConversation, this.branch);
+        // Generate response
+        this.generate();
+      } catch (error) {
+        this.error = error.toString();
+        // Remove temporary message on error
+        this.chatLog.pop();
+        // Refill input
+        this.newMessage = messageContent;
       }
-      console.log(req);
-      console.log(await req.json());
-      this.newMessage = "";
-      // reload conversation
-      await this.selectConversation(this.selectedConversation, this.branch);
-      // generate
-      this.generate();
     },
     async generate() {
       this.generating = true;
@@ -309,5 +332,13 @@ new Vue({
       this.conversationsLimit += 100;
       await this.getConversations();
     },
+    handleKeyDown(e) {
+      // If Enter is pressed without Shift
+      if (e.key === 'Enter' && !e.shiftKey) {
+        e.preventDefault();  // Prevent default newline
+        this.sendMessage();  // Send the message
+      }
+      // If Shift+Enter, let the default behavior happen (create newline)
+    },
   },
 });
diff --git a/gptme/tools/__init__.py b/gptme/tools/__init__.py
index 8912ed6d..351adb84 100644
--- a/gptme/tools/__init__.py
+++ b/gptme/tools/__init__.py
@@ -6,6 +6,7 @@
 from .base import ConfirmFunc, ToolSpec, ToolUse
 from .browser import tool as browser_tool
 from .chats import tool as chats_tool
+from .computer import tool as computer_tool
 from .gh import tool as gh_tool
 from .patch import tool as patch_tool
 from .python import register_function
@@ -43,11 +44,18 @@
     youtube_tool,
     screenshot_tool,
     vision_tool,
+    computer_tool,
     # python tool is loaded last to ensure all functions are registered
     python_tool,
 ]
 loaded_tools: list[ToolSpec] = []
 
+# Tools that are disabled by default, unless explicitly enabled
+# TODO: find a better way to handle this
+tools_default_disabled = [
+    "computer",
+]
+
 
 def init_tools(allowlist=None) -> None:
     """Runs initialization logic for tools."""
@@ -60,11 +68,14 @@ def init_tools(allowlist=None) -> None:
             continue
         if tool in loaded_tools:
             continue
+        if tool.name in tools_default_disabled:
+            if not allowlist or tool.name not in allowlist:
+                continue
         load_tool(tool)
 
     for tool_name in allowlist or []:
         if not has_tool(tool_name):
-            logger.warning(f"Tool '{tool_name}' not found")
+            raise ValueError(f"Tool '{tool_name}' not found")
 
 
 def load_tool(tool: ToolSpec) -> None:
diff --git a/gptme/tools/computer.py b/gptme/tools/computer.py
new file mode 100644
index 00000000..2775e2ab
--- /dev/null
+++ b/gptme/tools/computer.py
@@ -0,0 +1,219 @@
+"""
+Tool for computer interaction through X11, including screen capture, keyboard, and mouse control.
+Similar to Anthropic's computer use demo, but integrated with gptme's architecture.
+"""
+
+import os
+import shlex
+import shutil
+import subprocess
+from enum import Enum
+from pathlib import Path
+from typing import Literal, TypedDict
+
+from ..message import Message
+from .base import ToolSpec, ToolUse
+from .screenshot import _screenshot
+from .vision import view_image
+
+# Constants from Anthropic's implementation
+TYPING_DELAY_MS = 12
+TYPING_GROUP_SIZE = 50
+OUTPUT_DIR = "/tmp/outputs"
+
+Action = Literal[
+    "key",
+    "type",
+    "mouse_move",
+    "left_click",
+    "left_click_drag",
+    "right_click",
+    "middle_click",
+    "double_click",
+    "screenshot",
+    "cursor_position",
+]
+
+
+class Resolution(TypedDict):
+    width: int
+    height: int
+
+
+# Recommended maximum resolutions for LLM vision
+MAX_SCALING_TARGETS: dict[str, Resolution] = {
+    "XGA": Resolution(width=1024, height=768),  # 4:3
+    "WXGA": Resolution(width=1280, height=800),  # 16:10
+    "FWXGA": Resolution(width=1366, height=768),  # ~16:9
+}
+
+
+class ScalingSource(Enum):
+    COMPUTER = "computer"
+    API = "api"
+
+
+def chunks(s: str, chunk_size: int) -> list[str]:
+    """Split string into chunks for typing simulation."""
+    return [s[i : i + chunk_size] for i in range(0, len(s), chunk_size)]
+
+
+def scale_coordinates(
+    source: ScalingSource, x: int, y: int, current_width: int, current_height: int
+) -> tuple[int, int]:
+    """Scale coordinates to/from recommended resolutions."""
+    ratio = current_width / current_height
+    target_dimension = None
+
+    for dimension in MAX_SCALING_TARGETS.values():
+        if abs(dimension["width"] / dimension["height"] - ratio) < 0.02:
+            if dimension["width"] < current_width:
+                target_dimension = dimension
+                break
+
+    if target_dimension is None:
+        return x, y
+
+    x_scaling_factor = target_dimension["width"] / current_width
+    y_scaling_factor = target_dimension["height"] / current_height
+
+    if source == ScalingSource.API:
+        if x > current_width or y > current_height:
+            raise ValueError(f"Coordinates {x}, {y} are out of bounds")
+        # Scale up
+        return round(x / x_scaling_factor), round(y / y_scaling_factor)
+    # Scale down
+    return round(x * x_scaling_factor), round(y * y_scaling_factor)
+
+
+def run_xdotool(cmd: str, display: str | None = None) -> str:
+    """Run an xdotool command with optional display setting and wait for completion."""
+    env = os.environ.copy()
+    if display:
+        env["DISPLAY"] = display
+    try:
+        result = subprocess.run(
+            f"xdotool {cmd}",
+            shell=True,
+            env=env,
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+        return result.stdout
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(f"xdotool command failed: {e.stderr}") from e
+
+
+def computer(
+    action: Action, text: str | None = None, coordinate: tuple[int, int] | None = None
+) -> Message | None:
+    """
+    Perform computer interactions through X11.
+
+    Args:
+        action: The type of action to perform
+        text: Text to type or key sequence to send
+        coordinate: X,Y coordinates for mouse actions
+    """
+    display = os.getenv("DISPLAY", ":1")
+    width = int(os.getenv("WIDTH", "1024"))
+    height = int(os.getenv("HEIGHT", "768"))
+
+    if action in ("mouse_move", "left_click_drag"):
+        if not coordinate:
+            raise ValueError(f"coordinate is required for {action}")
+        x, y = scale_coordinates(
+            ScalingSource.API, coordinate[0], coordinate[1], width, height
+        )
+
+        if action == "mouse_move":
+            run_xdotool(f"mousemove --sync {x} {y}", display)
+        else:  # left_click_drag
+            run_xdotool(f"mousedown 1 mousemove --sync {x} {y} mouseup 1", display)
+
+        print(f"Moved mouse to {x},{y}")
+    elif action in ("key", "type"):
+        if not text:
+            raise ValueError(f"text is required for {action}")
+
+        if action == "key":
+            run_xdotool(f"key -- {text}", display)
+            print(f"Sent key sequence: {text}")
+        else:  # type
+            for chunk in chunks(text, TYPING_GROUP_SIZE):
+                run_xdotool(
+                    f"type --delay {TYPING_DELAY_MS} -- {shlex.quote(chunk)}",
+                    display,
+                )
+            print(f"Typed text: {text}")
+    elif action in ("left_click", "right_click", "middle_click", "double_click"):
+        click_arg = {
+            "left_click": "1",
+            "right_click": "3",
+            "middle_click": "2",
+            "double_click": "--repeat 2 --delay 500 1",
+        }[action]
+        run_xdotool(f"click {click_arg}", display)
+        print(f"Performed {action}")
+    elif action == "screenshot":
+        # Use X11-specific screenshot if available, fall back to native
+        output_dir = Path(OUTPUT_DIR)
+        output_dir.mkdir(parents=True, exist_ok=True)
+        path = output_dir / "screenshot.png"
+
+        if shutil.which("gnome-screenshot"):
+            run_xdotool(f"gnome-screenshot -f {path} -p", display)
+        elif os.name == "posix":
+            path = _screenshot(path)  # Use existing screenshot function
+        else:
+            raise NotImplementedError("Screenshot not supported on this platform")
+
+        # Scale if needed
+        if path.exists():
+            x, y = scale_coordinates(
+                ScalingSource.COMPUTER, width, height, width, height
+            )
+            subprocess.run(
+                f"convert {path} -resize {x}x{y}! {path}", shell=True, check=True
+            )
+            return view_image(path)
+        else:
+            print("Error: Screenshot failed")
+    elif action == "cursor_position":
+        output = run_xdotool("getmouselocation --shell", display)
+        x = int(output.split("X=")[1].split("\n")[0])
+        y = int(output.split("Y=")[1].split("\n")[0])
+        x, y = scale_coordinates(ScalingSource.COMPUTER, x, y, width, height)
+        print(f"Cursor position: X={x},Y={y}")
+
+    raise ValueError(f"Invalid action: {action}")
+
+
+instructions = """
+Use this tool to interact with the computer through X11.
+Available actions:
+- key: Send key sequence (e.g., "Return", "Control_L+c")
+- type: Type text with realistic delays
+- mouse_move: Move mouse to coordinates
+- left_click, right_click, middle_click, double_click: Mouse clicks
+- left_click_drag: Click and drag to coordinates
+- screenshot: Take and view a screenshot
+- cursor_position: Get current mouse position
+"""
+
+examples = f"""
+#### View a screenshot
+> User: What do you see on the screen?
+> Assistant:
+{ToolUse("ipython", [], 'computer("screenshot")').to_output()}
+> System: Viewing image...
+"""
+
+tool = ToolSpec(
+    name="computer",
+    desc="Control the computer through X11 (keyboard, mouse, screen)",
+    instructions=instructions,
+    examples=examples,
+    functions=[computer],
+)
diff --git a/gptme/tools/python.py b/gptme/tools/python.py
index 2866939f..da0c8c70 100644
--- a/gptme/tools/python.py
+++ b/gptme/tools/python.py
@@ -115,16 +115,18 @@ def execute_python(
         result = _ipython.run_cell(code, silent=False, store_history=False)
 
     output = ""
+    # TODO: should we include captured stdout with messages like these?
+    # used by vision tool
     if isinstance(result.result, Message):
         yield result.result
         return
 
     if result.result is not None:
         output += f"Result:\n```\n{result.result}\n```\n\n"
+
     # only show stdout if there is no result
     elif captured.stdout:
         output += f"```stdout\n{captured.stdout.rstrip()}\n```\n\n"
-
     if captured.stderr:
         output += f"```stderr\n{captured.stderr.rstrip()}\n```\n\n"
     if result.error_in_exec:
diff --git a/gptme/tools/screenshot.py b/gptme/tools/screenshot.py
index 7510f376..1a26ef02 100644
--- a/gptme/tools/screenshot.py
+++ b/gptme/tools/screenshot.py
@@ -18,7 +18,7 @@ def _screenshot(path: Path) -> Path:
             subprocess.run(["screencapture", str(path)], check=True)
         else:  # Linux
             # TODO: add support for specifying window/fullscreen?
-            subprocess.run(["scrot", str(path)], check=True)  # --focused
+            subprocess.run(["scrot", "--overwrite", str(path)], check=True)  # --focused
     else:
         raise NotImplementedError(
             "Screenshot functionality is only available on macOS and Linux."
diff --git a/gptme/tools/vision.py b/gptme/tools/vision.py
index 56b4afb8..5f7f2476 100644
--- a/gptme/tools/vision.py
+++ b/gptme/tools/vision.py
@@ -14,6 +14,8 @@ def view_image(image_path: Path | str) -> Message:
     """View an image."""
     if isinstance(image_path, str):
         image_path = Path(image_path)
+    if not image_path.exists():
+        return Message("system", f"Image not found at {image_path}")
     return Message("system", f"Viewing image at {image_path}", files=[image_path])
 
 
diff --git a/poetry.lock b/poetry.lock
index 96020bd6..5bcbc787 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -2498,6 +2498,20 @@ files = [
 [package.extras]
 cli = ["click (>=5.0)"]
 
+[[package]]
+name = "python-xlib"
+version = "0.33"
+description = "Python X Library"
+optional = true
+python-versions = "*"
+files = [
+    {file = "python-xlib-0.33.tar.gz", hash = "sha256:55af7906a2c75ce6cb280a584776080602444f75815a7aff4d287bb2d7018b32"},
+    {file = "python_xlib-0.33-py2.py3-none-any.whl", hash = "sha256:c3534038d42e0df2f1392a1b30a15a4ff5fdc2b86cfa94f072bf11b10a164398"},
+]
+
+[package.dependencies]
+six = ">=1.10.0"
+
 [[package]]
 name = "pytz"
 version = "2024.2"
@@ -3411,12 +3425,13 @@ files = [
 requests = "*"
 
 [extras]
-all = ["flask", "matplotlib", "numpy", "pandas", "pillow", "playwright"]
+all = ["flask", "matplotlib", "numpy", "pandas", "pillow", "playwright", "python-xlib"]
 browser = ["playwright"]
+computer = ["pillow", "python-xlib"]
 datascience = ["matplotlib", "numpy", "pandas", "pillow"]
 server = ["flask"]
 
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "043894e5aad0f1d3bfb95134b3469d5eba53ce7f8e8f5e83bef97ff90ee483cb"
+content-hash = "1a818642af5d1853c1f8bbc628218757ec088e756a1098e5caca4b4403268258"
diff --git a/pyproject.toml b/pyproject.toml
index fae3bbf2..7018a53e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -36,6 +36,7 @@ ipython = "^8.17.2"
 bashlex = "^0.18"
 playwright = {version = "1.47.*", optional=true}  # version constrained due to annoying to have to run `playwright install` on every update
 youtube_transcript_api = {version = "^0.6.1", optional = true}
+python-xlib = {version = "^0.33", optional = true}  # for X11 interaction
 
 # providers
 openai = "^1.0"
@@ -85,6 +86,7 @@ types-lxml = "*"
 server = ["flask"]
 browser = ["playwright"]
 datascience = ["matplotlib", "pandas", "numpy", "pillow"]
+computer = ["python-xlib", "pillow"]  # pillow already in datascience but listed for clarity
 all = [
     # server
     "flask",
@@ -92,6 +94,8 @@ all = [
     "playwright",
     # datascience
     "matplotlib", "pandas", "numpy", "pillow",
+    # computer
+    "python-xlib",
 ]
 
 [tool.ruff]
diff --git a/scripts/Dockerfile b/scripts/Dockerfile
index 59b5bbd6..82c52eac 100644
--- a/scripts/Dockerfile
+++ b/scripts/Dockerfile
@@ -52,18 +52,5 @@ ENV PYTHONPATH=/app
 # Set the working directory
 WORKDIR /workspace
 
-# Expose port 5000
-EXPOSE 5000
-
-# Healthcheck
-# TODO: only relevant for server
-#HEALTHCHECK --interval=30s --timeout=30s --start-period=5s --retries=3 \
-#    CMD curl -f http://localhost:5000/ || exit 1
-
-# TODO: make into separate image
-# CMD ["poetry", "run", "python", "-m", "gptme.server"]
-
-RUN poetry config virtualenvs.create false
-
 # Entrypoint if prompt/args given, run the CLI
 ENTRYPOINT ["python", "-m", "gptme"]
diff --git a/scripts/Dockerfile.computer b/scripts/Dockerfile.computer
new file mode 100644
index 00000000..292dc7b5
--- /dev/null
+++ b/scripts/Dockerfile.computer
@@ -0,0 +1,102 @@
+FROM ubuntu:22.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV DEBIAN_PRIORITY=high
+
+# Install system dependencies
+RUN apt-get update && \
+    apt-get -y upgrade && \
+    apt-get -y install \
+    build-essential \
+    python3.10 \
+    python3.10-dev \
+    python3-pip \
+    # UI Requirements
+    xvfb \
+    xterm \
+    xdotool \
+    scrot \
+    imagemagick \
+    sudo \
+    mutter \
+    x11vnc \
+    tint2 \
+    x11-apps \
+    # Tools
+    make \
+    git \
+    tmux \
+    curl \
+    pandoc \
+    netcat-openbsd \
+    net-tools \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install noVNC for web access
+RUN git clone --branch v1.5.0 https://github.com/novnc/noVNC.git /opt/noVNC && \
+    git clone --branch v0.12.0 https://github.com/novnc/websockify /opt/noVNC/utils/websockify && \
+    ln -s /opt/noVNC/vnc.html /opt/noVNC/index.html
+
+# Create user
+ENV USERNAME=computeruse
+ENV HOME=/home/$USERNAME
+RUN useradd -m -s /bin/bash -d $HOME $USERNAME && \
+    echo "${USERNAME} ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers && \
+    mkdir -p /workspace && \
+    mkdir -p /tmp/.X11-unix && \
+    chmod 1777 /tmp/.X11-unix && \
+    chown -R $USERNAME:$USERNAME /workspace
+
+# Set up environment
+ENV PYTHONPATH=/app
+ENV PATH="/usr/local/bin:$PATH"
+ENV DISPLAY_NUM=1
+ENV WIDTH=1024
+ENV HEIGHT=768
+ENV DISPLAY=:${DISPLAY_NUM}
+
+# Install Poetry and dependencies
+RUN python3.10 -m pip install --upgrade pip && \
+    python3.10 -m pip install poetry
+
+# Set up project
+WORKDIR /app
+COPY --chown=$USERNAME:$USERNAME pyproject.toml poetry.lock /app/
+RUN poetry config virtualenvs.create false
+RUN poetry install --no-interaction --no-ansi -E server --no-root
+
+COPY --chown=$USERNAME:$USERNAME . /app
+RUN poetry install --no-interaction --no-ansi -E server
+
+# Copy desktop environment configs (at the end for faster rebuilds)
+RUN mkdir -p $HOME && chmod 777 $HOME
+COPY scripts/computer_home $HOME
+RUN chown -R $USERNAME:$USERNAME $HOME
+
+# Switch to non-root user
+USER $USERNAME
+WORKDIR $HOME
+
+# Configure git
+RUN git config --global user.name "gptme" && \
+    git config --global user.email "gptme@superuserlabs.org" && \
+    git config --global init.defaultBranch main
+
+# Copy startup scripts
+COPY --chown=$USERNAME:$USERNAME scripts/computer_home/* $HOME/
+RUN chmod +x $HOME/*.sh
+
+# Set working directory
+WORKDIR /workspace
+
+# Expose ports
+# 6080: noVNC web interface
+# 8080: Chat view
+EXPOSE 6080 8080
+
+HEALTHCHECK --interval=30s --timeout=30s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:8080 || exit 1
+
+# Start services
+WORKDIR $HOME
+ENTRYPOINT ["./entrypoint.sh"]
diff --git a/scripts/Dockerfile.server b/scripts/Dockerfile.server
new file mode 100644
index 00000000..9ebf7f74
--- /dev/null
+++ b/scripts/Dockerfile.server
@@ -0,0 +1,23 @@
+# Use the main Dockerfile as the base image
+FROM gptme:latest
+
+# Install server dependencies
+USER root
+WORKDIR /app
+RUN poetry install -E server --without=dev
+
+# Switch back to non-root user
+USER appuser
+
+# Set the working directory
+WORKDIR /workspace
+
+# Expose the server port
+EXPOSE 5000
+
+# Add healthcheck
+HEALTHCHECK --interval=30s --timeout=30s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:5000/ || exit 1
+
+# Set the entrypoint to run the server
+ENTRYPOINT ["python", "-m", "gptme.server", "--host", "0.0.0.0"]
diff --git a/scripts/computer_home/.config/tint2/applications/firefox-custom.desktop b/scripts/computer_home/.config/tint2/applications/firefox-custom.desktop
new file mode 100755
index 00000000..94802126
--- /dev/null
+++ b/scripts/computer_home/.config/tint2/applications/firefox-custom.desktop
@@ -0,0 +1,8 @@
+[Desktop Entry]
+Name=Firefox Custom
+Comment=Open Firefox with custom URL
+Exec=firefox-esr -new-window
+Icon=firefox-esr
+Terminal=false
+Type=Application
+Categories=Network;WebBrowser;
diff --git a/scripts/computer_home/.config/tint2/applications/gedit.desktop b/scripts/computer_home/.config/tint2/applications/gedit.desktop
new file mode 100755
index 00000000..d5af03f4
--- /dev/null
+++ b/scripts/computer_home/.config/tint2/applications/gedit.desktop
@@ -0,0 +1,8 @@
+[Desktop Entry]
+Name=Gedit
+Comment=Open gedit
+Exec=gedit
+Icon=text-editor-symbolic
+Terminal=false
+Type=Application
+Categories=TextEditor;
diff --git a/scripts/computer_home/.config/tint2/applications/terminal.desktop b/scripts/computer_home/.config/tint2/applications/terminal.desktop
new file mode 100644
index 00000000..0c2d45d4
--- /dev/null
+++ b/scripts/computer_home/.config/tint2/applications/terminal.desktop
@@ -0,0 +1,8 @@
+[Desktop Entry]
+Name=Terminal
+Comment=Open Terminal
+Exec=xterm
+Icon=utilities-terminal
+Terminal=false
+Type=Application
+Categories=System;TerminalEmulator;
diff --git a/scripts/computer_home/.config/tint2/tint2rc b/scripts/computer_home/.config/tint2/tint2rc
new file mode 100644
index 00000000..5db6d312
--- /dev/null
+++ b/scripts/computer_home/.config/tint2/tint2rc
@@ -0,0 +1,100 @@
+#-------------------------------------
+# Panel
+panel_items = TL
+panel_size = 100% 60
+panel_margin = 0 0
+panel_padding = 2 0 2
+panel_background_id = 1
+wm_menu = 0
+panel_dock = 0
+panel_position = bottom center horizontal
+panel_layer = top
+panel_monitor = all
+panel_shrink = 0
+autohide = 0
+autohide_show_timeout = 0
+autohide_hide_timeout = 0.5
+autohide_height = 2
+strut_policy = follow_size
+panel_window_name = tint2
+disable_transparency = 1
+mouse_effects = 1
+font_shadow = 0
+mouse_hover_icon_asb = 100 0 10
+mouse_pressed_icon_asb = 100 0 0
+scale_relative_to_dpi = 0
+scale_relative_to_screen_height = 0
+
+#-------------------------------------
+# Taskbar
+taskbar_mode = single_desktop
+taskbar_hide_if_empty = 0
+taskbar_padding = 0 0 2
+taskbar_background_id = 0
+taskbar_active_background_id = 0
+taskbar_name = 1
+taskbar_hide_inactive_tasks = 0
+taskbar_hide_different_monitor = 0
+taskbar_hide_different_desktop = 0
+taskbar_always_show_all_desktop_tasks = 0
+taskbar_name_padding = 4 2
+taskbar_name_background_id = 0
+taskbar_name_active_background_id = 0
+taskbar_name_font_color = #e3e3e3 100
+taskbar_name_active_font_color = #ffffff 100
+taskbar_distribute_size = 0
+taskbar_sort_order = none
+task_align = left
+
+#-------------------------------------
+# Launcher
+launcher_padding = 4 8 4
+launcher_background_id = 0
+launcher_icon_background_id = 0
+launcher_icon_size = 48
+launcher_icon_asb = 100 0 0
+launcher_icon_theme_override = 0
+startup_notifications = 1
+launcher_tooltip = 1
+
+#-------------------------------------
+# Launcher icon
+launcher_item_app = /usr/share/applications/libreoffice-calc.desktop
+launcher_item_app = /home/computeruse/.config/tint2/applications/terminal.desktop
+launcher_item_app = /home/computeruse/.config/tint2/applications/firefox-custom.desktop
+launcher_item_app = /usr/share/applications/xpaint.desktop
+launcher_item_app = /usr/share/applications/xpdf.desktop
+launcher_item_app = /home/computeruse/.config/tint2/applications/gedit.desktop
+launcher_item_app = /usr/share/applications/galculator.desktop
+
+#-------------------------------------
+# Background definitions
+# ID 1
+rounded = 0
+border_width = 0
+background_color = #000000 60
+border_color = #000000 30
+
+# ID 2
+rounded = 4
+border_width = 1
+background_color = #777777 20
+border_color = #777777 30
+
+# ID 3
+rounded = 4
+border_width = 1
+background_color = #777777 20
+border_color = #ffffff 40
+
+# ID 4
+rounded = 4
+border_width = 1
+background_color = #aa4400 100
+border_color = #aa7733 100
+
+# ID 5
+rounded = 4
+border_width = 1
+background_color = #aaaa00 100
+border_color = #aaaa00 100
diff --git a/scripts/computer_home/entrypoint.sh b/scripts/computer_home/entrypoint.sh
new file mode 100755
index 00000000..8102f465
--- /dev/null
+++ b/scripts/computer_home/entrypoint.sh
@@ -0,0 +1,11 @@
+#!/bin/bash
+set -e
+
+./start_all.sh
+./novnc_startup.sh
+
+# Start gptme server
+python3 -m gptme.server --host 0.0.0.0 --port 8080 --tools python,computer
+
+# Keep the container running
+tail -f /dev/null
diff --git a/scripts/computer_home/mutter_startup.sh b/scripts/computer_home/mutter_startup.sh
new file mode 100755
index 00000000..7c5404f2
--- /dev/null
+++ b/scripts/computer_home/mutter_startup.sh
@@ -0,0 +1,22 @@
+#!/bin/bash
+set -e
+
+echo "Starting mutter..."
+XDG_SESSION_TYPE=x11 mutter --replace --sm-disable 2>/tmp/mutter_stderr.log &
+
+timeout=30
+while [ $timeout -gt 0 ]; do
+    if xdotool search --class "mutter" >/dev/null 2>&1; then
+        break
+    fi
+    sleep 1
+    ((timeout--))
+done
+
+if [ $timeout -eq 0 ]; then
+    echo "mutter stderr output:" >&2
+    cat /tmp/mutter_stderr.log >&2
+    exit 1
+fi
+
+rm /tmp/mutter_stderr.log
diff --git a/scripts/computer_home/novnc_startup.sh b/scripts/computer_home/novnc_startup.sh
new file mode 100755
index 00000000..ca2d5c0b
--- /dev/null
+++ b/scripts/computer_home/novnc_startup.sh
@@ -0,0 +1,29 @@
+#!/bin/bash
+echo "starting noVNC"
+
+# Start noVNC with optimized settings
+/opt/noVNC/utils/novnc_proxy \
+    --vnc localhost:5900 \
+    --listen 6080 \
+    --web /opt/noVNC \
+    --heartbeat 15 \
+    --idle-timeout 0 \
+    > /tmp/novnc.log 2>&1 &
+
+# Wait for noVNC to start
+timeout=10
+while [ $timeout -gt 0 ]; do
+    if netstat -tuln | grep -q ":6080 "; then
+        break
+    fi
+    sleep 1
+    ((timeout--))
+done
+
+if [ $timeout -eq 0 ]; then
+    echo "noVNC failed to start, log output:" >&2
+    cat /tmp/novnc.log >&2
+    exit 1
+fi
+
+echo "noVNC started successfully"
diff --git a/scripts/computer_home/start_all.sh b/scripts/computer_home/start_all.sh
new file mode 100755
index 00000000..96e4675b
--- /dev/null
+++ b/scripts/computer_home/start_all.sh
@@ -0,0 +1,8 @@
+#!/bin/bash
+set -e
+
+export DISPLAY=:${DISPLAY_NUM}
+./xvfb_startup.sh
+./tint2_startup.sh
+./mutter_startup.sh
+./x11vnc_startup.sh
diff --git a/scripts/computer_home/tint2_startup.sh b/scripts/computer_home/tint2_startup.sh
new file mode 100755
index 00000000..0227aabb
--- /dev/null
+++ b/scripts/computer_home/tint2_startup.sh
@@ -0,0 +1,22 @@
+#!/bin/bash
+set -e
+
+echo "Starting tint2..."
+tint2 2>/tmp/tint2_stderr.log &
+
+timeout=30
+while [ $timeout -gt 0 ]; do
+    if xdotool search --class "tint2" >/dev/null 2>&1; then
+        break
+    fi
+    sleep 1
+    ((timeout--))
+done
+
+if [ $timeout -eq 0 ]; then
+    echo "tint2 stderr output:" >&2
+    cat /tmp/tint2_stderr.log >&2
+    exit 1
+fi
+
+rm /tmp/tint2_stderr.log
diff --git a/scripts/computer_home/x11vnc_startup.sh b/scripts/computer_home/x11vnc_startup.sh
new file mode 100755
index 00000000..0253dbc4
--- /dev/null
+++ b/scripts/computer_home/x11vnc_startup.sh
@@ -0,0 +1,45 @@
+#!/bin/bash
+echo "starting vnc"
+
+(x11vnc -display $DISPLAY \
+    -forever \
+    -shared \
+    -rfbport 5900 \
+    -nopw \
+    2>/tmp/x11vnc_stderr.log) &
+
+x11vnc_pid=$!
+
+# Wait for x11vnc to start
+timeout=10
+while [ $timeout -gt 0 ]; do
+    if netstat -tuln | grep -q ":5900 "; then
+        break
+    fi
+    sleep 1
+    ((timeout--))
+done
+
+if [ $timeout -eq 0 ]; then
+    echo "x11vnc failed to start, stderr output:" >&2
+    cat /tmp/x11vnc_stderr.log >&2
+    exit 1
+fi
+
+: > /tmp/x11vnc_stderr.log
+
+# Monitor x11vnc process in the background
+(
+    while true; do
+        if ! kill -0 $x11vnc_pid 2>/dev/null; then
+            echo "x11vnc process crashed, restarting..." >&2
+            if [ -f /tmp/x11vnc_stderr.log ]; then
+                echo "x11vnc stderr output:" >&2
+                cat /tmp/x11vnc_stderr.log >&2
+                rm /tmp/x11vnc_stderr.log
+            fi
+            exec "$0"
+        fi
+        sleep 5
+    done
+) &
diff --git a/scripts/computer_home/xvfb_startup.sh b/scripts/computer_home/xvfb_startup.sh
new file mode 100755
index 00000000..92d8db35
--- /dev/null
+++ b/scripts/computer_home/xvfb_startup.sh
@@ -0,0 +1,23 @@
+#!/bin/bash
+set -e
+
+DPI=96
+RES_AND_DEPTH=${WIDTH}x${HEIGHT}x24
+
+echo "Starting Xvfb..."
+Xvfb $DISPLAY -ac -screen 0 $RES_AND_DEPTH -retro -dpi $DPI -nolisten tcp -nolisten unix &
+XVFB_PID=$!
+
+# Wait for Xvfb to be ready
+timeout=10
+start_time=$(date +%s)
+while ! xdpyinfo >/dev/null 2>&1; do
+    if [ $(($(date +%s) - start_time)) -gt $timeout ]; then
+        echo "Xvfb failed to start within $timeout seconds" >&2
+        exit 1
+    fi
+    sleep 0.1
+done
+
+echo "Xvfb started successfully on display ${DISPLAY}"
+echo "Xvfb PID: $XVFB_PID"