wcgw 1.4.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

wcgw/client/mcp_server/server.py

@@ -2,6 +2,7 @@ import asyncio
 import importlib
 import json
 import os
+import sys
 import traceback
 from typing import Any
 
@@ -17,39 +18,44 @@ from ...types_ import (
     BashInteraction,
     CreateFileNew,
     FileEdit,
+    Keyboard,
+    Mouse,
     ReadFile,
     ReadImage,
     ResetShell,
     Initialize,
+    ScreenShot,
+    GetScreenInfo,
 )
+from ..computer_use import Computer
 
 
 server = Server("wcgw")
 
 
-@server.list_resources()
+@server.list_resources()  # type: ignore
 async def handle_list_resources() -> list[types.Resource]:
     return []
 
 
-@server.read_resource()
+@server.read_resource()  # type: ignore
 async def handle_read_resource(uri: AnyUrl) -> str:
     raise ValueError("No resources available")
 
 
-@server.list_prompts()
+@server.list_prompts()  # type: ignore
 async def handle_list_prompts() -> list[types.Prompt]:
     return []
 
 
-@server.get_prompt()
+@server.get_prompt()  # type: ignore
 async def handle_get_prompt(
     name: str, arguments: dict[str, str] | None
 ) -> types.GetPromptResult:
-    types.GetPromptResult(messages=[])
+    return types.GetPromptResult(messages=[])
 
 
-@server.list_tools()
+@server.list_tools()  # type: ignore
 async def handle_list_tools() -> list[types.Tool]:
     """
     List available tools.
@@ -62,6 +68,7 @@ async def handle_list_tools() -> list[types.Tool]:
         )
     ) as f:
         diffinstructions = f.read()
+
     return [
         ToolParam(
             inputSchema=Initialize.model_json_schema(),
@@ -81,6 +88,7 @@ async def handle_list_tools() -> list[types.Tool]:
 - The first line might be `(...truncated)` if the output is too long.
 - Always run `pwd` if you get any file or directory not found error to make sure you're not lost.
 - The control will return to you in 5 seconds regardless of the status. For heavy commands, keep checking status using BashInteraction till they are finished.
+- Run long running commands in background using screen instead of "&".
 """,
         ),
         ToolParam(
@@ -107,7 +115,6 @@ async def handle_list_tools() -> list[types.Tool]:
             name="CreateFileNew",
             description="""
 - Write content to a new file. Provide file path and content. Use this instead of BashCommand for writing new files.
-- This doesn't create any directories, please create directories using `mkdir -p` BashCommand.
 - Provide absolute file path only.
 - For editing existing files, use FileEdit instead of this tool.
 """,
@@ -120,7 +127,7 @@ async def handle_list_tools() -> list[types.Tool]:
         ToolParam(
             inputSchema=ResetShell.model_json_schema(),
             name="ResetShell",
-            description="Resets the shell. Use only if all interrupts and prompt reset attempts have failed repeatedly.",
+            description="Resets the shell. Use only if all interrupts and prompt reset attempts have failed repeatedly.\nAlso exits the docker environment.\nYou need to call GetScreenInfo again.",
         ),
         ToolParam(
             inputSchema=FileEdit.model_json_schema(),
@@ -136,14 +143,54 @@ async def handle_list_tools() -> list[types.Tool]:
             name="ReadImage",
             description="""
 - Read an image from the shell.
+""",
+        ),
+        ToolParam(
+            inputSchema=GetScreenInfo.model_json_schema(),
+            name="GetScreenInfo",
+            description="""
+- Get display information of an OS running on docker using image "ghcr.io/anthropics/anthropic-quickstarts:computer-use-demo-latest"
+- If user hasn't provided docker image id, check using `docker ps` and provide the id.
+- Important: call this first in the conversation before ScreenShot, Mouse, and Keyboard tools.
+- Connects shell to the docker environment.
+- Note: once this is called, the shell enters the docker environment. All bash commands will run over there.
+""",
+        ),
+        ToolParam(
+            inputSchema=ScreenShot.model_json_schema(),
+            name="ScreenShot",
+            description="""
+- Capture screenshot of an OS running on docker using image "ghcr.io/anthropics/anthropic-quickstarts:computer-use-demo-latest"
+- If user hasn't provided docker image id, check using `docker ps` and provide the id.
+- Capture ScreenShot of the current screen for automation.
+""",
+        ),
+        ToolParam(
+            inputSchema=Mouse.model_json_schema(),
+            name="Mouse",
+            description="""
+- Interact with docker container running image "ghcr.io/anthropics/anthropic-quickstarts:computer-use-demo-latest"
+- If user hasn't provided docker image id, check using `docker ps` and provide the id.
+- Interact with the screen using mouse
+""",
+        ),
+        ToolParam(
+            inputSchema=Keyboard.model_json_schema(),
+            name="Keyboard",
+            description="""
+- Interact with docker container running image "ghcr.io/anthropics/anthropic-quickstarts:computer-use-demo-latest"
+- If user hasn't provided docker image id, check using `docker ps` and provide the id.
+- Emulate keyboard input to the screen
+- Uses xdootool to send keyboard input, keys like Return, BackSpace, Escape, Page_Up, etc. can be used.
+- Do not use it to interact with Bash tool.
 """,
         ),
     ]
 
 
-@server.call_tool()
+@server.call_tool()  # type: ignore
 async def handle_call_tool(
-    name: str, arguments: dict | None
+    name: str, arguments: dict[str, Any] | None
 ) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
     if not arguments:
         raise ValueError("Missing arguments")
@@ -163,60 +210,74 @@ async def handle_call_tool(
         tool_call = tool_type(**{k: try_json(v) for k, v in arguments.items()})
 
     try:
-        output_or_done, _ = get_tool_output(
+        output_or_dones, _ = get_tool_output(
             tool_call, default_enc, 0.0, lambda x, y: ("", 0), 8000
         )
 
     except Exception as e:
-        output_or_done = f"GOT EXCEPTION while calling tool. Error: {e}"
+        output_or_dones = [f"GOT EXCEPTION while calling tool. Error: {e}"]
         tb = traceback.format_exc()
-        print(output_or_done + "\n" + tb)
-
-    assert not isinstance(output_or_done, DoneFlag)
-    if isinstance(output_or_done, str):
-        if issubclass(tool_type, Initialize):
-            output_or_done += """
-            
-You're an expert software engineer with shell and code knowledge.
-
-Instructions:
-   
-    - You should use the provided bash execution, reading and writing file tools to complete objective.
-    - First understand about the project by getting the folder structure (ignoring .git, node_modules, venv, etc.)
-    - Always read relevant files before editing.
-    - Do not provide code snippets unless asked by the user, instead directly edit the code.
+        print(str(output_or_dones[0]) + "\n" + tb)
+
+    content: list[types.TextContent | types.ImageContent | types.EmbeddedResource] = []
+    for output_or_done in output_or_dones:
+        assert not isinstance(output_or_done, DoneFlag)
+        if isinstance(output_or_done, str):
+            if issubclass(tool_type, Initialize):
+                output_or_done += """
+                
+    You're an expert software engineer with shell and code knowledge.
 
+    Instructions:
     
-Additional instructions:
-    Always run `pwd` if you get any file or directory not found error to make sure you're not lost, or to get absolute cwd.
+        - You should use the provided bash execution, reading and writing file tools to complete objective.
+        - First understand about the project by getting the folder structure (ignoring .git, node_modules, venv, etc.)
+        - Always read relevant files before editing.
+        - Do not provide code snippets unless asked by the user, instead directly edit the code.
 
-    Always write production ready, syntactically correct code.
-"""
+        
+    Additional instructions:
+        Always run `pwd` if you get any file or directory not found error to make sure you're not lost, or to get absolute cwd.
 
-        return [types.TextContent(type="text", text=output_or_done)]
+        Always write production ready, syntactically correct code.
+    """
 
-    return [
-        types.ImageContent(
-            type="image",
-            data=output_or_done.data,
-            mimeType=output_or_done.media_type,
-        )
-    ]
+            content.append(types.TextContent(type="text", text=output_or_done))
+        else:
+            content.append(
+                types.ImageContent(
+                    type="image",
+                    data=output_or_done.data,
+                    mimeType=output_or_done.media_type,
+                )
+            )
+
+    return content
 
 
 async def main() -> None:
     version = importlib.metadata.version("wcgw")
-    # Run the server using stdin/stdout streams
-    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
-        await server.run(
-            read_stream,
-            write_stream,
-            InitializationOptions(
-                server_name="wcgw",
-                server_version=version,
-                capabilities=server.get_capabilities(
-                    notification_options=NotificationOptions(),
-                    experimental_capabilities={},
-                ),
-            ),
-        )
+    while True:
+        try:
+            # Run the server using stdin/stdout streams
+            async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
+                await server.run(
+                    read_stream,
+                    write_stream,
+                    InitializationOptions(
+                        server_name="wcgw",
+                        server_version=version,
+                        capabilities=server.get_capabilities(
+                            notification_options=NotificationOptions(),
+                            experimental_capabilities={},
+                        ),
+                    ),
+                    raise_exceptions=True,
+                )
+        except BaseException as e:
+            print(f"Server encountered an error: {e}", file=sys.stderr)
+            print("Stack trace:", file=sys.stderr)
+            traceback.print_exc(file=sys.stderr)
+            print("Restarting server in 5 seconds...", file=sys.stderr)
+            await asyncio.sleep(5)
+            continue

wcgw/client/openai_client.py

@@ -176,6 +176,7 @@ def loop(
 - The first line might be `(...truncated)` if the output is too long.
 - Always run `pwd` if you get any file or directory not found error to make sure you're not lost.
 - The control will return to you in 5 seconds regardless of the status. For heavy commands, keep checking status using BashInteraction till they are finished.
+- Run long running commands in background using screen instead of "&".
 """,
         ),
         openai.pydantic_function_tool(
@@ -198,7 +199,6 @@ def loop(
             CreateFileNew,
             description="""
 - Write content to a new file. Provide file path and content. Use this instead of BashCommand for writing new files.
-- This doesn't create any directories, please create directories using `mkdir -p` BashCommand.
 - Provide absolute file path only.
 - For editing existing files, use FileEdit instead of this tool.""",
         ),
@@ -336,13 +336,14 @@ System information:
                     for tool_call_id, toolcallargs in tool_call_args_by_id.items():
                         for toolindex, tool_args in toolcallargs.items():
                             try:
-                                output_or_done, cost_ = get_tool_output(
+                                output_or_dones, cost_ = get_tool_output(
                                     json.loads(tool_args),
                                     enc,
                                     limit - cost,
                                     loop,
                                     max_tokens=8000,
                                 )
+                                output_or_done = output_or_dones[0]
                             except Exception as e:
                                 output_or_done = (
                                     f"GOT EXCEPTION while calling tool. Error: {e}"

wcgw/client/sys_utils.py

@@ -0,0 +1,40 @@
+import subprocess
+
+MAX_RESPONSE_LEN: int = 16000
+
+
+def maybe_truncate(content: str, truncate_after: int | None = MAX_RESPONSE_LEN) -> str:
+    """Truncate content and append a notice if content exceeds the specified length."""
+    return (
+        content
+        if not truncate_after or len(content) <= truncate_after
+        else content[:truncate_after] + TRUNCATED_MESSAGE
+    )
+
+
+def command_run(
+    cmd: str,
+    timeout: float | None = 30.0,  # seconds
+    truncate_after: int | None = MAX_RESPONSE_LEN,
+    text: bool = True,
+) -> tuple[int, str, str]:
+    """Run a shell command synchronously with a timeout."""
+    try:
+        process = subprocess.Popen(
+            cmd,
+            shell=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=text,
+        )
+        stdout, stderr = process.communicate(timeout=timeout)
+        return (
+            process.returncode or 0,
+            maybe_truncate(stdout, truncate_after=truncate_after),
+            maybe_truncate(stderr, truncate_after=truncate_after),
+        )
+    except subprocess.TimeoutExpired as exc:
+        process.kill()
+        raise TimeoutError(
+            f"Command '{cmd}' timed out after {timeout} seconds"
+        ) from exc