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

wcgw/__init__.py

@@ -1,2 +1,3 @@
 from .client.cli import app
 from .client.tools import run as listen
+from .client.mcp_server import main as mcp_server

wcgw/client/anthropic_client.py

@@ -29,7 +29,6 @@ from ..types_ import (
     FileEdit,
     ReadFile,
     ReadImage,
-    Writefile,
     ResetShell,
 )
 
@@ -165,6 +164,7 @@ def loop(
 - Optionally `exit shell has restarted` is the output, in which case environment resets, you can run fresh commands.
 - 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.
 """,
         ),
         ToolParam(
@@ -193,7 +193,7 @@ def loop(
 - 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.
+- For editing existing files, use FileEdit instead of this tool.
 """,
         ),
         ToolParam(
@@ -219,7 +219,7 @@ def loop(
     uname_machine = os.uname().machine
 
     system = f"""
-You're a cli assistant.
+You're an expert software engineer with shell and code knowledge.
 
 Instructions:
 
@@ -227,10 +227,11 @@ Instructions:
     - 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.
-    
+
 System information:
     - System: {uname_sysname}
     - Machine: {uname_machine}
+    - Current directory: {os.getcwd()}
 """
 
     with open(os.path.join(os.path.dirname(__file__), "diff-instructions.txt")) as f:
@@ -361,7 +362,7 @@ System information:
                                     enc,
                                     limit - cost,
                                     loop,
-                                    max_tokens=8096,
+                                    max_tokens=8000,
                                 )
                             except Exception as e:
                                 output_or_done = (

wcgw/client/diff-instructions.txt

@@ -4,6 +4,7 @@ Instructions for editing files.
 
 Only edit the files using the following SEARCH/REPLACE blocks.
 ```
+file_edit_using_search_replace_blocks="""
 <<<<<<< SEARCH
 def hello():
     "print a greeting"
@@ -32,6 +33,7 @@ def call_hello_renamed():
     hello_renamed()
     impl2()
 >>>>>>> REPLACE
+"""
 ```
 
 # *SEARCH/REPLACE block* Rules:
@@ -43,7 +45,7 @@ Every *SEARCH/REPLACE block* must use this format:
 4. The lines to replace into the source code
 5. The end of the replace block: >>>>>>> REPLACE
 
-Every "<<<<<<< SEARCH" section must *EXACTLY MATCH* the existing file content, character for character, including all comments, docstrings, etc.
+Every "<<<<<<< SEARCH" section must *EXACTLY MATCH* the existing file content, character for character, including all comments, docstrings, whitespaces, etc.
 
 *SEARCH/REPLACE* blocks will *only* replace the first match occurrence.
 Including multiple unique *SEARCH/REPLACE* blocks if needed.
@@ -53,3 +55,5 @@ Keep *SEARCH/REPLACE* blocks concise.
 Break large *SEARCH/REPLACE* blocks into a series of smaller blocks that each change a small portion of the file.
 Include just the changing lines, and a few surrounding lines if needed for uniqueness.
 Do not include long runs of unchanging lines in *SEARCH/REPLACE* blocks.
+
+Preserve leading spaces and indentations in both SEARCH and REPLACE blocks.

wcgw/client/mcp_server/Readme.md

@@ -0,0 +1,26 @@
+# Claude desktop support
+
+## Setup
+
+Update `claude_desktop_config.json`
+
+```json
+{
+  "mcpServers": {
+    "wcgw": {
+      "command": "uvx",
+      "args": ["--from", "wcgw", "wcgw_mcp"]
+    }
+  }
+}
+```
+
+Then restart claude app.
+
+## Usage
+
+You should be able to see this icon if everything goes right.
+
+![mcp icon](https://github.com/rusiaaman/wcgw/blob/main/static/rocket-icon.png?raw=true)
+
+Then ask claude to execute shell commands, read files, edit files, run your code, etc.

wcgw/client/mcp_server/__init__.py

@@ -0,0 +1,11 @@
+from wcgw.client.mcp_server import server
+import asyncio
+
+
+def main():
+    """Main entry point for the package."""
+    asyncio.run(server.main())
+
+
+# Optionally expose other important items at package level
+__all__ = ["main", "server"]

wcgw/client/mcp_server/server.py

@@ -0,0 +1,222 @@
+import asyncio
+import importlib
+import json
+import os
+import traceback
+from typing import Any
+
+from mcp.server.models import InitializationOptions
+import mcp.types as types
+from mcp.types import Tool as ToolParam
+from mcp.server import NotificationOptions, Server
+from pydantic import AnyUrl, ValidationError
+import mcp.server.stdio
+from ..tools import DoneFlag, get_tool_output, which_tool_name, default_enc
+from ...types_ import (
+    BashCommand,
+    BashInteraction,
+    CreateFileNew,
+    FileEdit,
+    ReadFile,
+    ReadImage,
+    ResetShell,
+    Initialize,
+)
+
+
+server = Server("wcgw")
+
+
+@server.list_resources()
+async def handle_list_resources() -> list[types.Resource]:
+    return []
+
+
+@server.read_resource()
+async def handle_read_resource(uri: AnyUrl) -> str:
+    raise ValueError("No resources available")
+
+
+@server.list_prompts()
+async def handle_list_prompts() -> list[types.Prompt]:
+    return []
+
+
+@server.get_prompt()
+async def handle_get_prompt(
+    name: str, arguments: dict[str, str] | None
+) -> types.GetPromptResult:
+    types.GetPromptResult(messages=[])
+
+
+@server.list_tools()
+async def handle_list_tools() -> list[types.Tool]:
+    """
+    List available tools.
+    Each tool specifies its arguments using JSON Schema validation.
+    """
+
+    with open(
+        os.path.join(
+            os.path.dirname(os.path.dirname(__file__)), "diff-instructions.txt"
+        )
+    ) as f:
+        diffinstructions = f.read()
+    return [
+        ToolParam(
+            inputSchema=Initialize.model_json_schema(),
+            name="Initialize",
+            description="""
+- Always call this at the start of the conversation before anything else.
+""",
+        ),
+        ToolParam(
+            inputSchema=BashCommand.model_json_schema(),
+            name="BashCommand",
+            description="""
+- Execute a bash command. This is stateful (beware with subsequent calls).
+- Do not use interactive commands like nano. Prefer writing simpler commands.
+- Status of the command and the current working directory will always be returned at the end.
+- Optionally `exit shell has restarted` is the output, in which case environment resets, you can run fresh commands.
+- 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.
+""",
+        ),
+        ToolParam(
+            inputSchema=BashInteraction.model_json_schema(),
+            name="BashInteraction",
+            description="""
+- Interact with running program using this tool
+- Special keys like arrows, interrupts, enter, etc.
+- Send text input to the running program.
+- Send send_specials=["Enter"] to recheck status of a running program.
+- Only one of send_text, send_specials, send_ascii should be provided.
+""",
+        ),
+        ToolParam(
+            inputSchema=ReadFile.model_json_schema(),
+            name="ReadFile",
+            description="""
+- Read full file content
+- Provide absolute file path only
+""",
+        ),
+        ToolParam(
+            inputSchema=CreateFileNew.model_json_schema(),
+            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.
+""",
+        ),
+        ToolParam(
+            inputSchema=ReadImage.model_json_schema(),
+            name="ReadImage",
+            description="Read an image from the shell.",
+        ),
+        ToolParam(
+            inputSchema=ResetShell.model_json_schema(),
+            name="ResetShell",
+            description="Resets the shell. Use only if all interrupts and prompt reset attempts have failed repeatedly.",
+        ),
+        ToolParam(
+            inputSchema=FileEdit.model_json_schema(),
+            name="FileEdit",
+            description="""
+- Use absolute file path only.
+- Use SEARCH/REPLACE blocks to edit the file.
+"""
+            + diffinstructions,
+        ),
+        ToolParam(
+            inputSchema=ReadImage.model_json_schema(),
+            name="ReadImage",
+            description="""
+- Read an image from the shell.
+""",
+        ),
+    ]
+
+
+@server.call_tool()
+async def handle_call_tool(
+    name: str, arguments: dict | None
+) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
+    if not arguments:
+        raise ValueError("Missing arguments")
+
+    tool_type = which_tool_name(name)
+
+    try:
+        tool_call = tool_type(**arguments)
+    except ValidationError:
+
+        def try_json(x: str) -> Any:
+            try:
+                return json.loads(x)
+            except json.JSONDecodeError:
+                return x
+
+        tool_call = tool_type(**{k: try_json(v) for k, v in arguments.items()})
+
+    try:
+        output_or_done, _ = 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}"
+        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.
+
+    
+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.
+
+    Always write production ready, syntactically correct code.
+"""
+
+        return [types.TextContent(type="text", text=output_or_done)]
+
+    return [
+        types.ImageContent(
+            type="image",
+            data=output_or_done.data,
+            mimeType=output_or_done.media_type,
+        )
+    ]
+
+
+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={},
+                ),
+            ),
+        )

wcgw/client/openai_client.py

@@ -27,7 +27,6 @@ from ..types_ import (
     FileEdit,
     ReadImage,
     ReadFile,
-    Writefile,
     ResetShell,
 )
 
@@ -148,8 +147,7 @@ def loop(
     my_dir = os.path.dirname(__file__)
 
     config = Config(
-        model=cast(
-            Models, os.getenv("OPENAI_MODEL", "gpt-4o-2024-08-06").lower()),
+        model=cast(Models, os.getenv("OPENAI_MODEL", "gpt-4o-2024-08-06").lower()),
         cost_limit=0.1,
         cost_unit="$",
         cost_file={
@@ -177,6 +175,7 @@ def loop(
 - Optionally `exit shell has restarted` is the output, in which case environment resets, you can run fresh commands.
 - 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.
 """,
         ),
         openai.pydantic_function_tool(
@@ -201,7 +200,7 @@ def loop(
 - 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.""",
+- For editing existing files, use FileEdit instead of this tool.""",
         ),
         openai.pydantic_function_tool(
             FileEdit,
@@ -223,7 +222,7 @@ def loop(
     uname_machine = os.uname().machine
 
     system = f"""
-You're a cli assistant.
+You're an expert software engineer with shell and code knowledge.
 
 Instructions:
 
@@ -235,6 +234,7 @@ Instructions:
 System information:
     - System: {uname_sysname}
     - Machine: {uname_machine}
+    - Current directory: {os.getcwd()}
 
 """
 
@@ -341,7 +341,7 @@ System information:
                                     enc,
                                     limit - cost,
                                     loop,
-                                    max_tokens=2048,
+                                    max_tokens=8000,
                                 )
                             except Exception as e:
                                 output_or_done = (

wcgw/client/tools.py

@@ -45,7 +45,7 @@ from openai.types.chat import (
     ChatCompletionMessage,
     ParsedChatCompletionMessage,
 )
-from nltk.metrics.distance import edit_distance
+from nltk.metrics.distance import edit_distance  # type: ignore[import-untyped]
 
 from ..types_ import (
     BashCommand,
@@ -53,10 +53,10 @@ from ..types_ import (
     CreateFileNew,
     FileEditFindReplace,
     FileEdit,
+    Initialize,
     ReadFile,
     ReadImage,
     ResetShell,
-    Writefile,
 )
 
 from .common import CostData, Models, discard_input
@@ -95,7 +95,8 @@ def ask_confirmation(prompt: Confirmation) -> str:
     return "Yes" if response.lower() == "y" else "No"
 
 
-PROMPT = "#@@"
+PROMPT_CONST = "#@wcgw@#"
+PROMPT = PROMPT_CONST
 
 
 def start_shell() -> pexpect.spawn:  # type: ignore
@@ -124,7 +125,7 @@ def _is_int(mystr: str) -> bool:
 
 
 def _get_exit_code() -> int:
-    if PROMPT != "#@@":
+    if PROMPT != PROMPT_CONST:
         return 0
     # First reset the prompt in case venv was sourced or other reasons.
     SHELL.sendline(f"export PS1={PROMPT}")
@@ -155,6 +156,16 @@ BASH_STATE: BASH_CLF_OUTPUT = "repl"
 CWD = os.getcwd()
 
 
+def initial_info() -> str:
+    uname_sysname = os.uname().sysname
+    uname_machine = os.uname().machine
+    return f"""
+System: {uname_sysname}
+Machine: {uname_machine}
+Current working directory: {CWD}
+"""
+
+
 def reset_shell() -> str:
     global SHELL, BASH_STATE, CWD
     SHELL.close(True)
@@ -306,9 +317,7 @@ def execute_bash(
                 updated_repl_mode = update_repl_prompt(bash_arg.send_text)
                 if updated_repl_mode:
                     BASH_STATE = "repl"
-                    response = (
-                        "Prompt updated, you can execute REPL lines using BashCommand now"
-                    )
+                    response = "Prompt updated, you can execute REPL lines using BashCommand now"
                     console.print(response)
                     return (
                         response,
@@ -334,7 +343,7 @@ def execute_bash(
         tokens = enc.encode(text)
 
         if max_tokens and len(tokens) >= max_tokens:
-            text = "...(truncated)\n" + enc.decode(tokens[-(max_tokens - 1):])
+            text = "...(truncated)\n" + enc.decode(tokens[-(max_tokens - 1) :])
 
         if is_interrupt:
             text = (
@@ -361,7 +370,7 @@ Otherwise, you may want to try Ctrl-c again or program specific exit interactive
 
     tokens = enc.encode(output)
     if max_tokens and len(tokens) >= max_tokens:
-        output = "...(truncated)\n" + enc.decode(tokens[-(max_tokens - 1):])
+        output = "...(truncated)\n" + enc.decode(tokens[-(max_tokens - 1) :])
 
     try:
         exit_status = get_status()
@@ -436,10 +445,10 @@ def read_image_from_shell(file_path: str) -> ImageData:
         image_bytes = image_file.read()
         image_b64 = base64.b64encode(image_bytes).decode("utf-8")
         image_type = mimetypes.guess_type(file_path)[0]
-        return ImageData(media_type=image_type, data=image_b64)
+        return ImageData(media_type=image_type, data=image_b64)  # type: ignore
 
 
-def write_file(writefile: Writefile | CreateFileNew, error_on_exist: bool) -> str:
+def write_file(writefile: CreateFileNew, error_on_exist: bool) -> str:
     if not os.path.isabs(writefile.file_path):
         return f"Failure: file_path should be absolute path, current working directory is {CWD}"
     else:
@@ -488,15 +497,21 @@ def find_least_edit_distance_substring(
         edit_distance_sum = 0
         for j in range(len(find_lines)):
             if (i + j) < len(content_lines):
-                edit_distance_sum += edit_distance(
-                    content_lines[i + j], find_lines[j])
+                edit_distance_sum += edit_distance(content_lines[i + j], find_lines[j])
             else:
                 edit_distance_sum += len(find_lines[j])
         if edit_distance_sum < min_edit_distance:
             min_edit_distance = edit_distance_sum
             orig_start_index = new_to_original_indices[i]
-            orig_end_index = new_to_original_indices.get(i + len(find_lines) - 1, len(orig_content_lines) - 1) + 1
-            min_edit_distance_lines = orig_content_lines[orig_start_index:orig_end_index]
+            orig_end_index = (
+                new_to_original_indices.get(
+                    i + len(find_lines) - 1, len(orig_content_lines) - 1
+                )
+                + 1
+            )
+            min_edit_distance_lines = orig_content_lines[
+                orig_start_index:orig_end_index
+            ]
     return "\n".join(min_edit_distance_lines), min_edit_distance
 
 
@@ -525,7 +540,8 @@ def do_diff_edit(fedit: FileEdit) -> str:
 
     if not os.path.isabs(fedit.file_path):
         raise Exception(
-            f"Failure: file_path should be absolute path, current working directory is {CWD}")
+            f"Failure: file_path should be absolute path, current working directory is {CWD}"
+        )
     else:
         path_ = fedit.file_path
 
@@ -535,13 +551,16 @@ def do_diff_edit(fedit: FileEdit) -> str:
     with open(path_) as f:
         apply_diff_to = f.read()
 
+    fedit.file_edit_using_search_replace_blocks = (
+        fedit.file_edit_using_search_replace_blocks.strip()
+    )
     lines = fedit.file_edit_using_search_replace_blocks.split("\n")
 
     if not lines or not re.match(r"^<<<<<<+\s*SEARCH\s*$", lines[0]):
         raise Exception(
             "Error: first line should be `<<<<<< SEARCH` to start a search-replace block"
         )
-    
+
     n_lines = len(lines)
     i = 0
     replacement_count = 0
@@ -561,15 +580,14 @@ def do_diff_edit(fedit: FileEdit) -> str:
 
             for line in search_block:
                 console.log("> " + line)
-            console.log("---")
+            console.log("=======")
             for line in replace_block:
                 console.log("< " + line)
-
+            console.log("\n\n\n\n")
             search_block_ = "\n".join(search_block)
             replace_block_ = "\n".join(replace_block)
 
-            apply_diff_to = edit_content(
-                apply_diff_to, search_block_, replace_block_)
+            apply_diff_to = edit_content(apply_diff_to, search_block_, replace_block_)
             replacement_count += 1
         else:
             i += 1
@@ -588,7 +606,8 @@ def do_diff_edit(fedit: FileEdit) -> str:
 def file_edit(fedit: FileEditFindReplace) -> str:
     if not os.path.isabs(fedit.file_path):
         raise Exception(
-            f"Failure: file_path should be absolute path, current working directory is {CWD}")
+            f"Failure: file_path should be absolute path, current working directory is {CWD}"
+        )
     else:
         path_ = fedit.file_path
 
@@ -598,17 +617,14 @@ def file_edit(fedit: FileEditFindReplace) -> str:
     if not fedit.find_lines:
         raise Exception("Error: `find_lines` cannot be empty")
 
-    out_string = "\n".join(
-        "> " + line for line in fedit.find_lines.split("\n"))
-    in_string = "\n".join(
-        "< " + line for line in fedit.replace_with_lines.split("\n"))
+    out_string = "\n".join("> " + line for line in fedit.find_lines.split("\n"))
+    in_string = "\n".join("< " + line for line in fedit.replace_with_lines.split("\n"))
     console.log(f"Editing file: {path_}\n---\n{out_string}\n---\n{in_string}\n---")
     try:
         with open(path_) as f:
             content = f.read()
 
-        content = edit_content(content, fedit.find_lines,
-                               fedit.replace_with_lines)
+        content = edit_content(content, fedit.find_lines, fedit.replace_with_lines)
 
         with open(path_, "w") as f:
             f.write(content)
@@ -645,7 +661,6 @@ TOOLS = (
     | BashCommand
     | BashInteraction
     | ResetShell
-    | Writefile
     | CreateFileNew
     | FileEditFindReplace
     | FileEdit
@@ -653,6 +668,7 @@ TOOLS = (
     | DoneFlag
     | ReadImage
     | ReadFile
+    | Initialize
 )
 
 
@@ -670,8 +686,6 @@ def which_tool_name(name: str) -> Type[TOOLS]:
         return BashInteraction
     elif name == "ResetShell":
         return ResetShell
-    elif name == "Writefile":
-        return Writefile
     elif name == "CreateFileNew":
         return CreateFileNew
     elif name == "FileEditFindReplace":
@@ -686,6 +700,8 @@ def which_tool_name(name: str) -> Type[TOOLS]:
         return ReadImage
     elif name == "ReadFile":
         return ReadFile
+    elif name == "Initialize":
+        return Initialize
     else:
         raise ValueError(f"Unknown tool name: {name}")
 
@@ -696,13 +712,13 @@ def get_tool_output(
     | BashCommand
     | BashInteraction
     | ResetShell
-    | Writefile
     | CreateFileNew
     | FileEditFindReplace
     | FileEdit
     | AIAssistant
     | DoneFlag
     | ReadImage
+    | Initialize
     | ReadFile,
     enc: tiktoken.Encoding,
     limit: float,
@@ -715,7 +731,6 @@ def get_tool_output(
             | BashCommand
             | BashInteraction
             | ResetShell
-            | Writefile
             | CreateFileNew
             | FileEditFindReplace
             | FileEdit
@@ -723,12 +738,12 @@ def get_tool_output(
             | DoneFlag
             | ReadImage
             | ReadFile
+            | Initialize
         ](
             Confirmation
             | BashCommand
             | BashInteraction
             | ResetShell
-            | Writefile
             | CreateFileNew
             | FileEditFindReplace
             | FileEdit
@@ -736,6 +751,7 @@ def get_tool_output(
             | DoneFlag
             | ReadImage
             | ReadFile
+            | Initialize
         )
         arg = adapter.validate_python(args)
     else:
@@ -747,9 +763,6 @@ def get_tool_output(
     elif isinstance(arg, (BashCommand | BashInteraction)):
         console.print("Calling execute bash tool")
         output = execute_bash(enc, arg, max_tokens)
-    elif isinstance(arg, Writefile):
-        console.print("Calling write file tool")
-        output = write_file(arg, False), 0
     elif isinstance(arg, CreateFileNew):
         console.print("Calling write file tool")
         output = write_file(arg, True), 0
@@ -770,10 +783,13 @@ def get_tool_output(
         output = read_image_from_shell(arg.file_path), 0.0
     elif isinstance(arg, ReadFile):
         console.print("Calling read file tool")
-        output = read_file(arg), 0.0
+        output = read_file(arg, max_tokens), 0.0
     elif isinstance(arg, ResetShell):
         console.print("Calling reset shell tool")
         output = reset_shell(), 0.0
+    elif isinstance(arg, Initialize):
+        console.print("Calling initial info tool")
+        output = initial_info(), 0.0
     else:
         raise ValueError(f"Unknown tool: {arg}")
 
@@ -785,8 +801,7 @@ History = list[ChatCompletionMessageParam]
 
 default_enc = tiktoken.encoding_for_model("gpt-4o")
 default_model: Models = "gpt-4o-2024-08-06"
-default_cost = CostData(cost_per_1m_input_tokens=0.15,
-                        cost_per_1m_output_tokens=0.6)
+default_cost = CostData(cost_per_1m_input_tokens=0.15, cost_per_1m_output_tokens=0.6)
 curr_cost = 0.0
 
 
@@ -794,13 +809,13 @@ class Mdata(BaseModel):
     data: (
         BashCommand
         | BashInteraction
-        | Writefile
         | CreateFileNew
         | ResetShell
         | FileEditFindReplace
         | FileEdit
         | str
         | ReadFile
+        | Initialize
     )
 
 
@@ -829,8 +844,7 @@ def register_client(server_url: str, client_uuid: str = "") -> None:
                     raise Exception(mdata)
                 try:
                     output, cost = get_tool_output(
-                        mdata.data, default_enc, 0.0, lambda x, y: (
-                            "", 0), None
+                        mdata.data, default_enc, 0.0, lambda x, y: ("", 0), 8000
                     )
                     curr_cost += cost
                     print(f"{curr_cost=}")
@@ -863,8 +877,7 @@ def app(
     register_client(server_url, client_uuid or "")
 
 
-def read_file(readfile: ReadFile) -> str:
-
+def read_file(readfile: ReadFile, max_tokens: Optional[int]) -> str:
     console.print(f"Reading file: {readfile.file_path}")
 
     if not os.path.isabs(readfile.file_path):
@@ -876,4 +889,11 @@ def read_file(readfile: ReadFile) -> str:
 
     with path.open("r") as f:
         content = f.read()
+
+    if max_tokens is not None:
+        tokens = default_enc.encode(content)
+        if len(tokens) > max_tokens:
+            content = default_enc.decode(tokens[: max_tokens - 5])
+            content += "\n...(truncated)"
+
     return content

wcgw/relay/serve.py

@@ -20,9 +20,9 @@ from ..types_ import (
     CreateFileNew,
     FileEditFindReplace,
     FileEdit,
+    Initialize,
     ReadFile,
     ResetShell,
-    Writefile,
     Specials,
 )
 
@@ -31,12 +31,12 @@ class Mdata(BaseModel):
     data: (
         BashCommand
         | BashInteraction
-        | Writefile
         | CreateFileNew
         | ResetShell
         | FileEditFindReplace
         | FileEdit
         | ReadFile
+        | Initialize
         | str
     )
     user_id: UUID
@@ -51,7 +51,7 @@ gpts: dict[UUID, Callable[[str], None]] = {}
 images: DefaultDict[UUID, dict[str, dict[str, Any]]] = DefaultDict(dict)
 
 
-CLIENT_VERSION_MINIMUM = "1.2.0"
+CLIENT_VERSION_MINIMUM = "1.3.0"
 
 
 @app.websocket("/v1/register/{uuid}")
@@ -65,14 +65,12 @@ async def register_websocket(websocket: WebSocket, uuid: UUID) -> None:
     # receive client version
     client_version = await websocket.receive_text()
     sem_version_client = semantic_version.Version.coerce(client_version)
-    sem_version_server = semantic_version.Version.coerce(
-        CLIENT_VERSION_MINIMUM)
+    sem_version_server = semantic_version.Version.coerce(CLIENT_VERSION_MINIMUM)
     if sem_version_client < sem_version_server:
         await websocket.send_text(
             Mdata(
                 user_id=uuid,
-                data=f"Client version {client_version} is outdated. Please upgrade to {
-                    CLIENT_VERSION_MINIMUM} or higher.",
+                data=f"Client version {client_version} is outdated. Please upgrade to {CLIENT_VERSION_MINIMUM} or higher.",
             ).model_dump_json()
         )
         await websocket.close(
@@ -92,8 +90,7 @@ async def register_websocket(websocket: WebSocket, uuid: UUID) -> None:
         while True:
             received_data = await websocket.receive_text()
             if uuid not in gpts:
-                raise fastapi.HTTPException(
-                    status_code=400, detail="No call made")
+                raise fastapi.HTTPException(status_code=400, detail="No call made")
             gpts[uuid](received_data)
     except WebSocketDisconnect:
         # Remove the client if the WebSocket is disconnected
@@ -280,11 +277,36 @@ async def read_file_endpoint(read_file_data: ReadFileWithUUID) -> str:
 
     gpts[user_id] = put_results
 
-    await clients[user_id](
-        Mdata(data=ReadFile(file_path=read_file_data.file_path,
-                            type=read_file_data.type
-                            ), user_id=user_id)
-    )
+    await clients[user_id](Mdata(data=read_file_data, user_id=user_id))
+
+    start_time = time.time()
+    while time.time() - start_time < 30:
+        if results is not None:
+            return results
+        await asyncio.sleep(0.1)
+
+    raise fastapi.HTTPException(status_code=500, detail="Timeout error")
+
+
+class InitializeWithUUID(Initialize):
+    user_id: UUID
+
+
+@app.post("/v1/initialize")
+async def initialize(initialize_data: InitializeWithUUID) -> str:
+    user_id = initialize_data.user_id
+    if user_id not in clients:
+        return "Failure: id not found, ask the user to check it."
+
+    results: Optional[str] = None
+
+    def put_results(result: str) -> None:
+        nonlocal results
+        results = result
+
+    gpts[user_id] = put_results
+
+    await clients[user_id](Mdata(data=initialize_data, user_id=user_id))
 
     start_time = time.time()
     while time.time() - start_time < 30:

wcgw/types_.py

@@ -24,11 +24,6 @@ class ReadImage(BaseModel):
     type: Literal["ReadImage"]
 
 
-class Writefile(BaseModel):
-    file_path: str
-    file_content: str
-
-
 class CreateFileNew(BaseModel):
     file_path: str
     file_content: str
@@ -52,3 +47,7 @@ class ResetShell(BaseModel):
 class FileEdit(BaseModel):
     file_path: str
     file_edit_using_search_replace_blocks: str
+
+
+class Initialize(BaseModel):
+    type: Literal["Initialize"]

{wcgw-1.2.2.dist-info → wcgw-1.4.0.dist-info}/METADATA

@@ -1,12 +1,13 @@
 Metadata-Version: 2.3
 Name: wcgw
-Version: 1.2.2
+Version: 1.4.0
 Summary: What could go wrong giving full shell access to chatgpt?
 Project-URL: Homepage, https://github.com/rusiaaman/wcgw
 Author-email: Aman Rusia <gapypi@arcfu.com>
 Requires-Python: <3.13,>=3.10
 Requires-Dist: anthropic>=0.39.0
 Requires-Dist: fastapi>=0.115.0
+Requires-Dist: mcp>=1.0.0
 Requires-Dist: mypy>=1.11.2
 Requires-Dist: nltk>=3.9.1
 Requires-Dist: openai>=1.46.0
@@ -50,7 +51,7 @@ You need to keep running this client for GPT to access your shell. Run it in a v
 ### Option 1: using uv [Recommended]
 ```sh
 $ curl -LsSf https://astral.sh/uv/install.sh | sh
-$ uv tool run --python 3.12 wcgw@latest
+$ uvx wcgw@latest
 ```
 
 ### Option 2: using pip
@@ -109,6 +110,9 @@ If you don't have public ip and domain name, you can use `ngrok` or similar serv
 The specify the server url in the `wcgw` command like so
 `wcgw --server-url https://your-url/v1/register`
 
+# Claude Support
+WCGW now supports Claude Desktop through the MCP protocol, allowing you to use Claude's capabilities directly from your desktop environment. This integration enables seamless interaction between Claude and your local shell.
+
 # [Optional] Local shell access with openai API key
 
 Add `OPENAI_API_KEY` and `OPENAI_ORG_ID` env variables.

wcgw-1.4.0.dist-info/RECORD

@@ -0,0 +1,20 @@
+wcgw/__init__.py,sha256=9K2QW7QuSLhMTVbKbBYd9UUp-ZyrfBrxcjuD_xk458k,118
+wcgw/types_.py,sha256=AWQUtzv7S_YQ2_36RMQOAA2gJ26ZQZcOA6l5EvFR6hk,1054
+wcgw/client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+wcgw/client/__main__.py,sha256=ngI_vBcLAv7fJgmS4w4U7tuWtalGB8c7W5qebuT6Z6o,30
+wcgw/client/anthropic_client.py,sha256=G7uCtUd5pt7ZQO7k2TJI3FtgK4QtzdIvd4Z1S18xGNU,15595
+wcgw/client/cli.py,sha256=Oja42CHkVO8puqOXflko9NeephYCMa85aBmQTEjBZtI,932
+wcgw/client/common.py,sha256=grH-yV_4tnTQZ29xExn4YicGLxEq98z-HkEZwH0ReSg,1410
+wcgw/client/diff-instructions.txt,sha256=s5AJKG23JsjwRYhFZFQVvwDpF67vElawrmdXwvukR1A,1683
+wcgw/client/openai_client.py,sha256=kQCtQX0x-jsnw-tHBxi41XqnOkCaoTUqY8K4gHeI4xM,17703
+wcgw/client/openai_utils.py,sha256=YNwCsA-Wqq7jWrxP0rfQmBTb1dI0s7dWXzQqyTzOZT4,2629
+wcgw/client/tools.py,sha256=jjQ8eFQvcE8xfoGR5WkoOgWYE0dOIThxFrrppY4gQDQ,27954
+wcgw/client/mcp_server/Readme.md,sha256=GH59caNySMMphqiimsVKrteU-jmNC7kcdBu4MPxtZZ8,477
+wcgw/client/mcp_server/__init__.py,sha256=cQ7PUrEmXUpio8x0SEoGWP5hCRPd7z2bAkNCbYbtTys,236
+wcgw/client/mcp_server/server.py,sha256=G8fyXQzVurM8rzekiXU62YaseDIT2mthc0lj64FDJKs,7089
+wcgw/relay/serve.py,sha256=RUcUeyL4Xt0EEo12Ul6VQjb4tRle4uIdsa85v7XXxEw,8771
+wcgw/relay/static/privacy.txt,sha256=s9qBdbx2SexCpC_z33sg16TptmAwDEehMCLz4L50JLc,529
+wcgw-1.4.0.dist-info/METADATA,sha256=I31G-r47C-9acBusZsKFt2BjbqgP1S83amCnQJuL_XQ,5501
+wcgw-1.4.0.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
+wcgw-1.4.0.dist-info/entry_points.txt,sha256=eKo1omwbAggWlQ0l7GKoR7uV1-j16nk9tK0BhC2Oz_E,120
+wcgw-1.4.0.dist-info/RECORD,,

{wcgw-1.2.2.dist-info → wcgw-1.4.0.dist-info}/entry_points.txt

@@ -1,4 +1,5 @@
 [console_scripts]
 wcgw = wcgw:listen
 wcgw_local = wcgw:app
+wcgw_mcp = wcgw:mcp_server
 wcgw_relay = wcgw.relay.serve:run

wcgw-1.2.2.dist-info/RECORD

@@ -1,17 +0,0 @@
-wcgw/__init__.py,sha256=PNWvBvjUKA3aj4bHOtIqBKCAtOW88pr0hAXZ7RylVr8,68
-wcgw/types_.py,sha256=HND4V3KruqkAhzzEf2ve7HhF8y_o5w11Fu784U8xN_I,1062
-wcgw/client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-wcgw/client/__main__.py,sha256=ngI_vBcLAv7fJgmS4w4U7tuWtalGB8c7W5qebuT6Z6o,30
-wcgw/client/anthropic_client.py,sha256=YQ2puVKlrgbI6QvV_DQjWkvq2qw6KhPPrivm0OGfRsc,15353
-wcgw/client/cli.py,sha256=Oja42CHkVO8puqOXflko9NeephYCMa85aBmQTEjBZtI,932
-wcgw/client/common.py,sha256=grH-yV_4tnTQZ29xExn4YicGLxEq98z-HkEZwH0ReSg,1410
-wcgw/client/diff-instructions.txt,sha256=Fy6op3wkr5gYgMHCRcitG4804zLrTxMbs-VSMB1oIzA,1548
-wcgw/client/openai_client.py,sha256=v6ZBW4oh1CJpyWLTt16S-7ew8Gmq6YrJzaTux2ffywI,17470
-wcgw/client/openai_utils.py,sha256=YNwCsA-Wqq7jWrxP0rfQmBTb1dI0s7dWXzQqyTzOZT4,2629
-wcgw/client/tools.py,sha256=zeUOemmFwReY1jEJnnAkhPbBwkL4pjZvQ6G5BWPux4c,27301
-wcgw/relay/serve.py,sha256=sMbERQTM1GhpnPWLzdtpcb23TIsgbP-ZuqVj-vjp5Rw,8186
-wcgw/relay/static/privacy.txt,sha256=s9qBdbx2SexCpC_z33sg16TptmAwDEehMCLz4L50JLc,529
-wcgw-1.2.2.dist-info/METADATA,sha256=0zQpOcQD8eRiQMVibiksflPhrngU49FSgpEHAmXQ_DA,5255
-wcgw-1.2.2.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
-wcgw-1.2.2.dist-info/entry_points.txt,sha256=WlIB825-Vm9ZtNzgENQsbHj4DRMkbpVR7uSkQyBlaPA,93
-wcgw-1.2.2.dist-info/RECORD,,