Documentation of how to customize compute use

Author: chriscarrollsmith

README.md

@@ -14,7 +14,7 @@ Lower-level feature improvements include:
 - Support for calling tools via remote MCP servers
 - Support for more models
 
-Note that this template implements support for MCP servers, but does not yet support web search, image generation, or computer use.
+Note that this template implements does not yet support image generation.
 
 ## Quickstart Setup
 
@@ -111,3 +111,50 @@ After defining your function (and optionally a template), start the server and n
 ![Function Registration](./docs/functions.png)
 
 Don't forget to click "Regenerate tool.config.json" to save your changes. (This will regenerate the `tool.config.json` file with your new configuration.)
+
+## Customizing Computer Use
+
+The computer use tool lets the assistant control a virtual screen by issuing mouse and keyboard actions (click, type, scroll, etc.) and receiving screenshots in return.
+
+By default, this template uses a headless Playwright browser as the computer use backend. When the assistant starts a computer use session, it sees a landing page with a URL input so it can navigate to websites.
+
+To swap in a different backend (e.g., a full virtual desktop via VNC/xdotool, or a GUI automation library like pyautogui), implement the `ComputerSession` and `ComputerSessionManager` protocols defined in `utils/computer_use.py`:
+
+```python
+from utils.computer_use import ComputerSession, ComputerSessionManager, Action
+
+class MySession:
+    """Your custom session — must implement ComputerSession protocol."""
+
+    async def screenshot(self) -> str:
+        """Capture the screen and return a base64-encoded PNG string."""
+        ...
+
+    async def execute(self, action: Action) -> str:
+        """Perform an action (click, type, etc.) and return a screenshot."""
+        ...
+
+    async def close(self) -> None:
+        """Release resources."""
+        ...
+
+class MySessionManager:
+    """Your custom manager — must implement ComputerSessionManager protocol."""
+
+    def get_or_create(self, conversation_id: str, width: int = 1024, height: int = 768) -> MySession:
+        ...
+
+    async def close(self, conversation_id: str) -> None:
+        ...
+
+    async def close_all(self) -> None:
+        ...
+```
+
+Then replace the `session_manager` singleton in `utils/computer_use.py`:
+
+```python
+session_manager: ComputerSessionManager = MySessionManager()
+```
+
+No other changes are needed — the chat router and shutdown logic use `session_manager` through the protocol interface.

tests/test_computer_use.py

@@ -32,6 +32,8 @@
 from utils.computer_use import (
     BrowserSession,
     BrowserSessionManager,
+    ComputerSession,
+    ComputerSessionManager,
     build_computer_tool,
     describe_action,
     execute_computer_actions,
@@ -336,6 +338,16 @@ def test_only_type_key(self):
         assert tool == {"type": "computer"}
 
 
+class TestProtocolConformance:
+    """Tests that concrete classes satisfy the Protocol interfaces."""
+
+    def test_browser_session_is_computer_session(self):
+        assert isinstance(BrowserSession(), ComputerSession)
+
+    def test_browser_session_manager_is_computer_session_manager(self):
+        assert isinstance(BrowserSessionManager(), ComputerSessionManager)
+
+
 class TestExecuteComputerAction:
     """Integration tests for execute_computer_action()."""
 
@@ -361,18 +373,22 @@ async def test_returns_valid_png(self):
     async def test_reuses_session_for_same_conversation(self):
         from utils.computer_use import session_manager
 
+        manager = session_manager
+        assert isinstance(manager, BrowserSessionManager)
         action = Screenshot(type="screenshot")
         await execute_computer_actions([action],"test-conv")
         await execute_computer_actions([action],"test-conv")
-        assert "test-conv" in session_manager._sessions
+        assert "test-conv" in manager._sessions
 
     @pytest.mark.anyio
     async def test_different_conversations_different_sessions(self):
         from utils.computer_use import session_manager
 
+        manager = session_manager
+        assert isinstance(manager, BrowserSessionManager)
         action = Screenshot(type="screenshot")
         await execute_computer_actions([action],"conv-a")
         await execute_computer_actions([action],"conv-b")
-        assert "conv-a" in session_manager._sessions
-        assert "conv-b" in session_manager._sessions
-        assert session_manager._sessions["conv-a"] is not session_manager._sessions["conv-b"]
+        assert "conv-a" in manager._sessions
+        assert "conv-b" in manager._sessions
+        assert manager._sessions["conv-a"] is not manager._sessions["conv-b"]

utils/computer_use.py

@@ -1,7 +1,7 @@
 import asyncio
 import base64
 import logging
-from typing import Union
+from typing import Protocol, Union, runtime_checkable
 
 from playwright.async_api import async_playwright, Browser, Page, Playwright
 
@@ -19,6 +19,48 @@
 
 logger = logging.getLogger("uvicorn.error")
 
+Action = Union[
+    Click, DoubleClick, Drag, Keypress,
+    Move, Screenshot, Scroll, Type, Wait,
+]
+
+
+@runtime_checkable
+class ComputerSession(Protocol):
+    """Protocol for computer use sessions.
+
+    Implement this to provide a custom backend (e.g., VNC, xdotool, pyautogui)
+    instead of the default Playwright-based BrowserSession.
+
+    Methods:
+        screenshot: Capture the current screen as a base64-encoded PNG string.
+        execute: Perform an action and return a base64-encoded PNG screenshot.
+        close: Release resources held by this session.
+    """
+
+    async def screenshot(self) -> str: ...
+    async def execute(self, action: Action) -> str: ...
+    async def close(self) -> None: ...
+
+
+@runtime_checkable
+class ComputerSessionManager(Protocol):
+    """Protocol for managing computer use sessions by conversation ID.
+
+    Implement this to provide custom session lifecycle management.
+
+    Methods:
+        get_or_create: Return an existing session or create a new one.
+        close: Close and remove a specific session.
+        close_all: Close all sessions (called during app shutdown).
+    """
+
+    def get_or_create(
+        self, conversation_id: str, width: int = 1024, height: int = 768,
+    ) -> ComputerSession: ...
+    async def close(self, conversation_id: str) -> None: ...
+    async def close_all(self) -> None: ...
+
 _LANDING_PAGE_HTML = """\
 <!DOCTYPE html>
 <html>
@@ -45,11 +87,6 @@
 </html>
 """
 
-Action = Union[
-    Click, DoubleClick, Drag, Keypress,
-    Move, Screenshot, Scroll, Type, Wait,
-]
-
 # Map OpenAI key names to Playwright key names
 _KEY_MAP: dict[str, str] = {
     "ctrl": "Control",
@@ -227,8 +264,9 @@ async def close_all(self) -> None:
         self._sessions.clear()
 
 
-# Module-level singleton
-session_manager = BrowserSessionManager()
+# Module-level singleton — replace with your own ComputerSessionManager implementation
+# to use a different backend (e.g., VNC, xdotool, pyautogui).
+session_manager: ComputerSessionManager = BrowserSessionManager()
 
 
 def build_computer_tool(**kwargs: object) -> dict[str, str]: