wcgw 3.0.1rc1__tar.gz → 3.0.1rc3__tar.gz

{wcgw-3.0.1rc1 → wcgw-3.0.1rc3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wcgw
-Version: 3.0.1rc1
+Version: 3.0.1rc3
 Summary: Shell and coding agent on claude and chatgpt
 Project-URL: Homepage, https://github.com/rusiaaman/wcgw
 Author-email: Aman Rusia <gapypi@arcfu.com>

{wcgw-3.0.1rc1 → wcgw-3.0.1rc3}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 authors = [{ name = "Aman Rusia", email = "gapypi@arcfu.com" }]
 name = "wcgw"
-version = "3.0.1rc1"
+version = "3.0.1rc3"
 description = "Shell and coding agent on claude and chatgpt"
 readme = "README.md"
 requires-python = ">=3.11"

{wcgw-3.0.1rc1 → wcgw-3.0.1rc3}/src/wcgw/client/bash_state/bash_state.py

@@ -6,7 +6,15 @@ import threading
 import time
 import traceback
 from dataclasses import dataclass
-from typing import Any, Literal, Optional
+from typing import (
+    Any,
+    Callable,
+    Concatenate,
+    Literal,
+    Optional,
+    ParamSpec,
+    TypeVar,
+)
 
 import pexpect
 import pyte
@@ -24,7 +32,8 @@ from ...types_ import (
 from ..encoder import EncoderDecoder
 from ..modes import BashCommandMode, FileEditMode, WriteIfEmptyMode
 
-PROMPT_CONST = "#" + "@wcgw@#"
+PROMPT_CONST = "wcgw→" + " "
+PROMPT_STATEMENT = "export GIT_PAGER=cat PAGER=cat PROMPT_COMMAND= PS1='wcgw→'' '"
 BASH_CLF_OUTPUT = Literal["repl", "pending"]
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
@@ -63,7 +72,7 @@ def get_tmpdir() -> str:
             timeout=CONFIG.timeout,
         ).strip()
         return result
-    except subprocess.CalledProcessError:
+    except (subprocess.CalledProcessError, FileNotFoundError):
         return "//tmp"
     except Exception:
         return ""
@@ -71,9 +80,7 @@ def get_tmpdir() -> str:
 
 def check_if_screen_command_available() -> bool:
     try:
-        subprocess.run(
-            ["screen", "-v"], capture_output=True, check=True, timeout=CONFIG.timeout
-        )
+        subprocess.run(["screen", "-v"], capture_output=True, check=True, timeout=0.2)
         return True
     except (subprocess.CalledProcessError, FileNotFoundError):
         return False
@@ -92,12 +99,14 @@ def cleanup_all_screens_with_name(name: str, console: Console) -> None:
             capture_output=True,
             text=True,
             check=True,
-            timeout=CONFIG.timeout,
+            timeout=0.2,
         )
         output = result.stdout
     except subprocess.CalledProcessError as e:
         # When no screens exist, screen may return a non-zero exit code.
         output = (e.stdout or "") + (e.stderr or "")
+    except FileNotFoundError:
+        return
 
     sessions_to_kill = []
 
@@ -121,7 +130,7 @@ def cleanup_all_screens_with_name(name: str, console: Console) -> None:
                 check=True,
                 timeout=CONFIG.timeout,
             )
-        except subprocess.CalledProcessError:
+        except (subprocess.CalledProcessError, FileNotFoundError):
             console.log(f"Failed to kill screen session: {session}")
 
 
@@ -136,23 +145,21 @@ def start_shell(
         **os.environ,
         "PS1": PROMPT_CONST,
         "TMPDIR": get_tmpdir(),
-        "TERM": "vt100",
+        "TERM": "xterm-256color",
     }
     try:
         shell = pexpect.spawn(
             cmd,
             env=overrideenv,  # type: ignore[arg-type]
-            echo=False,
+            echo=True,
             encoding="utf-8",
             timeout=CONFIG.timeout,
             cwd=initial_dir,
             codec_errors="backslashreplace",
             dimensions=(500, 160),
         )
-        shell.sendline(
-            f"export PROMPT_COMMAND= PS1={PROMPT_CONST}"
-        )  # Unset prompt command to avoid interfering
-        shell.expect(PROMPT_CONST, timeout=CONFIG.timeout)
+        shell.sendline(PROMPT_STATEMENT)  # Unset prompt command to avoid interfering
+        shell.expect(PROMPT_CONST, timeout=0.2)
     except Exception as e:
         console.print(traceback.format_exc())
         console.log(f"Error starting shell: {e}. Retrying without rc ...")
@@ -160,13 +167,13 @@ def start_shell(
         shell = pexpect.spawn(
             "/bin/bash --noprofile --norc",
             env=overrideenv,  # type: ignore[arg-type]
-            echo=False,
+            echo=True,
             encoding="utf-8",
             timeout=CONFIG.timeout,
             codec_errors="backslashreplace",
         )
-        shell.sendline(f"export PS1={PROMPT_CONST}")
-        shell.expect(PROMPT_CONST, timeout=CONFIG.timeout)
+        shell.sendline(PROMPT_STATEMENT)
+        shell.expect(PROMPT_CONST, timeout=0.2)
 
     shellid = "wcgw." + time.strftime("%H%M%S")
     if over_screen:
@@ -174,22 +181,11 @@ def start_shell(
             raise ValueError("Screen command not available")
         # shellid is just hour, minute, second number
         shell.sendline(f"trap 'screen -X -S {shellid} quit' EXIT")
-        shell.expect(PROMPT_CONST, timeout=CONFIG.timeout)
+        shell.expect(PROMPT_CONST, timeout=0.2)
 
-        shell.sendline(f"screen -q -s /bin/bash -S {shellid}")
+        shell.sendline(f"screen -q -S {shellid} /bin/bash --noprofile --norc")
         shell.expect(PROMPT_CONST, timeout=CONFIG.timeout)
 
-        console.log(f"Entering screen session, name: {shellid}")
-
-    shell.sendline("stty -icanon -echo")
-    shell.expect(PROMPT_CONST, timeout=CONFIG.timeout)
-
-    shell.sendline("set +o pipefail")
-    shell.expect(PROMPT_CONST, timeout=CONFIG.timeout)
-
-    shell.sendline("export GIT_PAGER=cat PAGER=cat")
-    shell.expect(PROMPT_CONST, timeout=CONFIG.timeout)
-
     return shell, shellid
 
 
@@ -217,7 +213,39 @@ def render_terminal_output(text: str) -> list[str]:
     return lines
 
 
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+def requires_shell(
+    func: Callable[Concatenate["BashState", "pexpect.spawn[str]", P], R],
+) -> Callable[Concatenate["BashState", P], R]:
+    def wrapper(self: "BashState", /, *args: P.args, **kwargs: P.kwargs) -> R:
+        if not self._shell_loading.is_set():
+            if not self._shell_loading.wait(
+                timeout=CONFIG.timeout * 2
+            ):  # Twice in worst case if screen fails
+                raise RuntimeError("Shell initialization timeout")
+
+        if self._shell_error:
+            raise RuntimeError(f"Shell failed to initialize: {self._shell_error}.")
+
+        if not self._shell:
+            raise RuntimeError("Shell not initialized")
+
+        return func(self, self._shell, *args, **kwargs)
+
+    return wrapper
+
+
 class BashState:
+    _shell: Optional["pexpect.spawn[str]"]
+    _shell_id: Optional[str]
+    _shell_lock: threading.Lock
+    _shell_loading: threading.Event
+    _shell_error: Optional[Exception]
+    _use_screen: bool
+
     def __init__(
         self,
         console: Console,
@@ -240,35 +268,65 @@ class BashState:
         )
         self._mode = mode or Modes.wcgw
         self._whitelist_for_overwrite: set[str] = whitelist_for_overwrite or set()
-        self._prompt = PROMPT_CONST
         self._bg_expect_thread: Optional[threading.Thread] = None
         self._bg_expect_thread_stop_event = threading.Event()
-        self._init_shell(use_screen)
-
-    def expect(self, pattern: Any, timeout: Optional[float] = -1) -> int:
+        self._shell = None
+        self._shell_id = None
+        self._shell_lock = threading.Lock()
+        self._shell_loading = threading.Event()
+        self._shell_error = None
+        self._use_screen = use_screen
+        self._start_shell_loading()
+
+    def _start_shell_loading(self) -> None:
+        def load_shell() -> None:
+            try:
+                with self._shell_lock:
+                    if self._shell is not None:
+                        return
+                    self._init_shell()
+                self.run_bg_expect_thread()
+            except Exception as e:
+                self._shell_error = e
+            finally:
+                self._shell_loading.set()
+
+        threading.Thread(target=load_shell).start()
+
+    @requires_shell
+    def expect(
+        self, shell: "pexpect.spawn[str]", pattern: Any, timeout: Optional[float] = -1
+    ) -> int:
         self.close_bg_expect_thread()
-        return self._shell.expect(pattern, timeout)
+        output = shell.expect(pattern, timeout)
+        return output
 
-    def send(self, s: str | bytes) -> int:
-        output = self._shell.send(s)
-        self.run_bg_expect_thread()
+    @requires_shell
+    def send(self, shell: "pexpect.spawn[str]", s: str | bytes) -> int:
+        self.close_bg_expect_thread()
+        output = shell.send(s)
         return output
 
-    def sendline(self, s: str | bytes) -> int:
-        output = self._shell.sendline(s)
-        self.run_bg_expect_thread()
+    @requires_shell
+    def sendline(self, shell: "pexpect.spawn[str]", s: str | bytes) -> int:
+        self.close_bg_expect_thread()
+        output = shell.sendline(s)
         return output
 
     @property
-    def linesep(self) -> Any:
-        return self._shell.linesep
+    @requires_shell
+    def linesep(self, shell: "pexpect.spawn[str]") -> Any:
+        return shell.linesep
 
-    def sendintr(self) -> None:
-        self._shell.sendintr()
+    @requires_shell
+    def sendintr(self, shell: "pexpect.spawn[str]") -> None:
+        self.close_bg_expect_thread()
+        shell.sendintr()
 
     @property
-    def before(self) -> Optional[str]:
-        return self._shell.before
+    @requires_shell
+    def before(self, shell: "pexpect.spawn[str]") -> Optional[str]:
+        return shell.before
 
     def run_bg_expect_thread(self) -> None:
         """
@@ -279,11 +337,16 @@ class BashState:
             while True:
                 if self._bg_expect_thread_stop_event.is_set():
                     break
+                if self._shell is None:
+                    time.sleep(0.1)
+                    continue
                 output = self._shell.expect([pexpect.EOF, pexpect.TIMEOUT], timeout=0.1)
                 if output == 0:
                     break
 
-        if self._bg_expect_thread:
+        if self._bg_expect_thread and self._bg_expect_thread.is_alive():
+            if not self._bg_expect_thread_stop_event.is_set():
+                return
             self.close_bg_expect_thread()
 
         self._bg_expect_thread = threading.Thread(
@@ -300,8 +363,13 @@ class BashState:
 
     def cleanup(self) -> None:
         self.close_bg_expect_thread()
-        self._shell.close(True)
-        cleanup_all_screens_with_name(self._shell_id, self.console)
+        with self._shell_lock:
+            if self._shell:
+                self._shell.close(True)
+                if self._shell_id:
+                    cleanup_all_screens_with_name(self._shell_id, self.console)
+                self._shell = None
+                self._shell_id = None
 
     def __enter__(self) -> "BashState":
         return self
@@ -325,26 +393,26 @@ class BashState:
     def write_if_empty_mode(self) -> WriteIfEmptyMode:
         return self._write_if_empty_mode
 
-    def ensure_env_and_bg_jobs(self) -> Optional[int]:
-        if self._prompt != PROMPT_CONST:
-            return None
+    @requires_shell
+    def ensure_env_and_bg_jobs(self, _: "pexpect.spawn[str]") -> Optional[int]:
+        return self._ensure_env_and_bg_jobs()
+
+    def _ensure_env_and_bg_jobs(self) -> Optional[int]:
+        # Do not add @requires_shell decorator here, as it will cause deadlock
+        self.close_bg_expect_thread()
+        assert self._shell is not None, "Bad state, shell is not initialized"
         quick_timeout = 0.2 if not self.over_screen else 1
         # First reset the prompt in case venv was sourced or other reasons.
-        self.sendline(f"export PS1={self._prompt}")
-        self.expect(self._prompt, timeout=quick_timeout)
+        self._shell.sendline(PROMPT_STATEMENT)
+        self._shell.expect(PROMPT_CONST, timeout=quick_timeout)
         # Reset echo also if it was enabled
-        self.sendline("stty -icanon -echo")
-        self.expect(self._prompt, timeout=quick_timeout)
-        self.sendline("set +o pipefail")
-        self.expect(self._prompt, timeout=quick_timeout)
-        self.sendline("export GIT_PAGER=cat PAGER=cat")
-        self.expect(self._prompt, timeout=quick_timeout)
-        self.sendline("jobs | wc -l")
+        command = "jobs | wc -l"
+        self._shell.sendline(command)
         before = ""
         counts = 0
         while not _is_int(before):  # Consume all previous output
             try:
-                self.expect(self._prompt, timeout=quick_timeout)
+                self._shell.expect(PROMPT_CONST, timeout=quick_timeout)
             except pexpect.TIMEOUT:
                 self.console.print(f"Couldn't get exit code, before: {before}")
                 raise
@@ -354,20 +422,18 @@ class BashState:
                 before_val = str(before_val)
             assert isinstance(before_val, str)
             before_lines = render_terminal_output(before_val)
-            before = "\n".join(before_lines).strip()
+            before = "\n".join(before_lines).replace(command, "").strip()
             counts += 1
             if counts > 100:
                 raise ValueError(
                     "Error in understanding shell output. This shouldn't happen, likely shell is in a bad state, please reset it"
                 )
-
         try:
             return int(before)
         except ValueError:
             raise ValueError(f"Malformed output: {before}")
 
-    def _init_shell(self, use_screen: bool) -> None:
-        self._prompt = PROMPT_CONST
+    def _init_shell(self) -> None:
         self._state: Literal["repl"] | datetime.datetime = "repl"
         self._is_in_docker: Optional[str] = ""
         # Ensure self._cwd exists
@@ -377,9 +443,9 @@ class BashState:
                 self._bash_command_mode.bash_mode == "restricted_mode",
                 self._cwd,
                 self.console,
-                over_screen=use_screen,
+                over_screen=self._use_screen,
             )
-            self.over_screen = use_screen
+            self.over_screen = self._use_screen
         except Exception as e:
             if not isinstance(e, ValueError):
                 self.console.log(traceback.format_exc())
@@ -394,15 +460,15 @@ class BashState:
             self.over_screen = False
 
         self._pending_output = ""
-
-        # Get exit info to ensure shell is ready
-        self.ensure_env_and_bg_jobs()
+        try:
+            self._ensure_env_and_bg_jobs()
+        except ValueError as e:
+            self.console.log("Error while running _ensure_env_and_bg_jobs" + str(e))
 
     def set_pending(self, last_pending_output: str) -> None:
         if not isinstance(self._state, datetime.datetime):
             self._state = datetime.datetime.now()
         self._pending_output = last_pending_output
-        self.run_bg_expect_thread()
 
     def set_repl(self) -> None:
         self._state = "repl"
@@ -427,12 +493,13 @@ class BashState:
 
     @property
     def prompt(self) -> str:
-        return self._prompt
+        return PROMPT_CONST
 
-    def update_cwd(self) -> str:
-        self.sendline("pwd")
-        self.expect(self._prompt, timeout=0.2)
-        before_val = self._shell.before
+    @requires_shell
+    def update_cwd(self, shell: "pexpect.spawn[str]") -> str:
+        shell.sendline("pwd")
+        shell.expect(PROMPT_CONST, timeout=0.2)
+        before_val = shell.before
         if not isinstance(before_val, str):
             before_val = str(before_val)
         before_lines = render_terminal_output(before_val)
@@ -442,7 +509,9 @@ class BashState:
 
     def reset_shell(self) -> None:
         self.cleanup()
-        self._init_shell(True)
+        self._shell_loading.clear()
+        self._shell_error = None
+        self._start_shell_loading()
 
     def serialize(self) -> dict[str, Any]:
         """Serialize BashState to a dictionary for saving"""
@@ -476,6 +545,16 @@ class BashState:
         cwd: str,
     ) -> None:
         """Create a new BashState instance from a serialized state dictionary"""
+        if (
+            self._bash_command_mode == bash_command_mode
+            and ((self._cwd == cwd) or not cwd)
+            and (self._file_edit_mode == file_edit_mode)
+            and (self._write_if_empty_mode == write_if_empty_mode)
+            and (self._mode == mode)
+            and (self._whitelist_for_overwrite == set(whitelist_for_overwrite))
+        ):
+            # No need to reset shell if the state is the same
+            return
         self._bash_command_mode = bash_command_mode
         self._cwd = cwd or self._cwd
         self._file_edit_mode = file_edit_mode
@@ -546,9 +625,8 @@ def _incremental_text(text: str, last_pending_output: str) -> str:
     # text = render_terminal_output(text[-100_000:])
     text = text[-100_000:]
 
-    last_pending_output_rendered_lines = render_terminal_output(last_pending_output)
-    last_pending_output_rendered = "\n".join(last_pending_output_rendered_lines)
-    last_rendered_lines = last_pending_output_rendered.split("\n")
+    last_rendered_lines = render_terminal_output(last_pending_output)
+    last_pending_output_rendered = "\n".join(last_rendered_lines)
     if not last_rendered_lines:
         return rstrip(render_terminal_output(text))
 
@@ -599,6 +677,27 @@ def execute_bash(
     bash_arg: BashCommand,
     max_tokens: Optional[int],
     timeout_s: Optional[float],
+) -> tuple[str, float]:
+    try:
+        output, cost = _execute_bash(bash_state, enc, bash_arg, max_tokens, timeout_s)
+
+        # Remove echo if it's a command
+        if isinstance(bash_arg.action, Command):
+            command = bash_arg.action.command.strip()
+            if output.startswith(command):
+                output = output[len(command) :]
+
+    finally:
+        bash_state.run_bg_expect_thread()
+    return output, cost
+
+
+def _execute_bash(
+    bash_state: BashState,
+    enc: EncoderDecoder[int],
+    bash_arg: BashCommand,
+    max_tokens: Optional[int],
+    timeout_s: Optional[float],
 ) -> tuple[str, float]:
     try:
         is_interrupt = False

{wcgw-3.0.1rc1 → wcgw-3.0.1rc3}/src/wcgw/client/file_ops/diff_edit.py

@@ -102,6 +102,8 @@ def line_process_max_space_tolerance(line: str) -> str:
     return re.sub(r"\s", "", line)
 
 
+REMOVE_INDENTATION = "Warning: matching after removing all spaces in lines."
+
 DEFAULT_TOLERANCES = [
     Tolerance(
         line_process=str.rstrip,
@@ -119,11 +121,50 @@ DEFAULT_TOLERANCES = [
         line_process=line_process_max_space_tolerance,
         severity_cat="WARNING",
         score_multiplier=50,
-        error_name="Warning: matching after removing all spaces in lines.",
+        error_name=REMOVE_INDENTATION,
     ),
 ]
 
 
+def fix_indentation(
+    matched_lines: list[str], searched_lines: list[str], replaced_lines: list[str]
+) -> list[str]:
+    if not matched_lines or not searched_lines or not replaced_lines:
+        return replaced_lines
+
+    def get_indentation(line: str) -> str:
+        match = re.match(r"^(\s*)", line)
+        assert match
+        return match.group(0)
+
+    matched_indents = [get_indentation(line) for line in matched_lines if line.strip()]
+    searched_indents = [
+        get_indentation(line) for line in searched_lines if line.strip()
+    ]
+    if len(matched_indents) != len(searched_indents):
+        return replaced_lines
+    diffs: list[int] = [
+        len(searched) - len(matched)
+        for matched, searched in zip(matched_indents, searched_indents)
+    ]
+    if not all(diff == diffs[0] for diff in diffs):
+        return replaced_lines
+    if diffs[0] == 0:
+        return replaced_lines
+
+    # At this point we have same number of non-empty lines and the same indentation difference
+    # We can now adjust the indentation of the replaced lines
+    def adjust_indentation(line: str, diff: int) -> str:
+        if diff < 0:
+            return matched_indents[0][:-diff] + line
+        return line[diff:]
+
+    if diffs[0] > 0:
+        if not (all(not line[: diffs[0]].strip() for line in replaced_lines)):
+            return replaced_lines
+    return [adjust_indentation(line, diffs[0]) for line in replaced_lines]
+
+
 def remove_leading_trailing_empty_lines(lines: list[str]) -> list[str]:
     start = 0
     end = len(lines) - 1
@@ -247,6 +288,16 @@ class FileEditInput:
                         ]
 
             for match, tolerances in matches_with_tolerances:
+                if any(
+                    tolerance.error_name == REMOVE_INDENTATION
+                    for tolerance in tolerances
+                ):
+                    replace_by = fix_indentation(
+                        self.file_lines[match.start : match.stop],
+                        first_block[0],
+                        replace_by,
+                    )
+
                 file_edit_input = FileEditInput(
                     self.file_lines,
                     match.stop,

{wcgw-3.0.1rc1 → wcgw-3.0.1rc3}/src/wcgw/client/mcp_server/server.py

@@ -1,5 +1,4 @@
 import importlib
-import json
 import logging
 import os
 from typing import Any
@@ -9,7 +8,7 @@ import mcp_wcgw.types as types
 from mcp_wcgw.server import NotificationOptions, Server
 from mcp_wcgw.server.models import InitializationOptions
 from mcp_wcgw.types import Tool as ToolParam
-from pydantic import AnyUrl, ValidationError
+from pydantic import AnyUrl
 
 from wcgw.client.modes import KTS
 from wcgw.client.tool_prompts import TOOL_PROMPTS
@@ -18,7 +17,13 @@ from ...types_ import (
     Initialize,
 )
 from ..bash_state.bash_state import CONFIG, BashState
-from ..tools import Context, default_enc, get_tool_output, which_tool_name
+from ..tools import (
+    Context,
+    default_enc,
+    get_tool_output,
+    parse_tool_by_name,
+    which_tool_name,
+)
 
 server = Server("wcgw")
 
@@ -104,20 +109,7 @@ async def handle_call_tool(
         raise ValueError("Missing arguments")
 
     tool_type = which_tool_name(name)
-
-    try:
-        tool_call = tool_type(**arguments)
-    except ValidationError:
-
-        def try_json(x: str) -> Any:
-            if not isinstance(x, str):
-                return x
-            try:
-                return json.loads(x)
-            except json.JSONDecodeError:
-                return x
-
-        tool_call = tool_type(**{k: try_json(v) for k, v in arguments.items()})
+    tool_call = parse_tool_by_name(name, arguments)
 
     try:
         assert BASH_STATE
@@ -165,7 +157,7 @@ async def main() -> None:
     version = str(importlib.metadata.version("wcgw"))
     home_dir = os.path.expanduser("~")
     with BashState(
-        Console(), home_dir, None, None, None, None, False, None
+        Console(), home_dir, None, None, None, None, True, None
     ) as BASH_STATE:
         BASH_STATE.console.log("wcgw version: " + version)
         # Run the server using stdin/stdout streams