wcgw 3.0.1rc1__tar.gz → 3.0.1rc2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wcgw
-Version: 3.0.1rc1
+Version: 3.0.1rc2
 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.1rc2}/pyproject.toml

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

{wcgw-3.0.1rc1 → wcgw-3.0.1rc2}/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
@@ -63,7 +71,7 @@ def get_tmpdir() -> str:
             timeout=CONFIG.timeout,
         ).strip()
         return result
-    except subprocess.CalledProcessError:
+    except (subprocess.CalledProcessError, FileNotFoundError):
         return "//tmp"
     except Exception:
         return ""
@@ -98,6 +106,8 @@ def cleanup_all_screens_with_name(name: str, console: Console) -> None:
     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 +131,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}")
 
 
@@ -176,11 +186,9 @@ def start_shell(
         shell.sendline(f"trap 'screen -X -S {shellid} quit' EXIT")
         shell.expect(PROMPT_CONST, timeout=CONFIG.timeout)
 
-        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)
 
@@ -217,7 +225,37 @@ 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):
+                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,
@@ -243,34 +281,63 @@ class BashState:
         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()
+            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)
+        return shell.expect(pattern, timeout)
 
-    def send(self, s: str | bytes) -> int:
-        output = self._shell.send(s)
+    @requires_shell
+    def send(self, shell: "pexpect.spawn[str]", s: str | bytes) -> int:
+        output = shell.send(s)
         self.run_bg_expect_thread()
         return output
 
-    def sendline(self, s: str | bytes) -> int:
-        output = self._shell.sendline(s)
+    @requires_shell
+    def sendline(self, shell: "pexpect.spawn[str]", s: str | bytes) -> int:
+        output = shell.sendline(s)
         self.run_bg_expect_thread()
         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:
+        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:
+    @requires_shell
+    def run_bg_expect_thread(self, shell: "pexpect.spawn[str]") -> None:
         """
         Run background expect thread for handling shell interactions.
         """
@@ -279,7 +346,7 @@ class BashState:
             while True:
                 if self._bg_expect_thread_stop_event.is_set():
                     break
-                output = self._shell.expect([pexpect.EOF, pexpect.TIMEOUT], timeout=0.1)
+                output = shell.expect([pexpect.EOF, pexpect.TIMEOUT], timeout=0.1)
                 if output == 0:
                     break
 
@@ -300,8 +367,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 +397,34 @@ class BashState:
     def write_if_empty_mode(self) -> WriteIfEmptyMode:
         return self._write_if_empty_mode
 
-    def ensure_env_and_bg_jobs(self) -> Optional[int]:
+    @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"
         if self._prompt != PROMPT_CONST:
             return None
         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(f"export PS1={self._prompt}")
+        self._shell.expect(self._prompt, 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")
+        self._shell.sendline("stty -icanon -echo")
+        self._shell.expect(self._prompt, timeout=quick_timeout)
+        self._shell.sendline("set +o pipefail")
+        self._shell.expect(self._prompt, timeout=quick_timeout)
+        self._shell.sendline("export GIT_PAGER=cat PAGER=cat")
+        self._shell.expect(self._prompt, timeout=quick_timeout)
+        self._shell.sendline("jobs | wc -l")
         before = ""
         counts = 0
         while not _is_int(before):  # Consume all previous output
             try:
-                self.expect(self._prompt, timeout=quick_timeout)
+                self._shell.expect(self._prompt, timeout=quick_timeout)
             except pexpect.TIMEOUT:
                 self.console.print(f"Couldn't get exit code, before: {before}")
                 raise
@@ -366,7 +446,7 @@ class BashState:
         except ValueError:
             raise ValueError(f"Malformed output: {before}")
 
-    def _init_shell(self, use_screen: bool) -> None:
+    def _init_shell(self) -> None:
         self._prompt = PROMPT_CONST
         self._state: Literal["repl"] | datetime.datetime = "repl"
         self._is_in_docker: Optional[str] = ""
@@ -377,9 +457,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,9 +474,7 @@ class BashState:
             self.over_screen = False
 
         self._pending_output = ""
-
-        # Get exit info to ensure shell is ready
-        self.ensure_env_and_bg_jobs()
+        self._ensure_env_and_bg_jobs()
 
     def set_pending(self, last_pending_output: str) -> None:
         if not isinstance(self._state, datetime.datetime):
@@ -429,10 +507,11 @@ class BashState:
     def prompt(self) -> str:
         return self._prompt
 
-    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(self._prompt, 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 +521,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 +557,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

{wcgw-3.0.1rc1 → wcgw-3.0.1rc2}/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.1rc2}/src/wcgw/client/mcp_server/server.py

@@ -165,7 +165,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

{wcgw-3.0.1rc1 → wcgw-3.0.1rc2}/src/wcgw/client/tools.py

@@ -206,6 +206,8 @@ Initialized in directory (also cwd): {context.bash_state.cwd}
 {memory}
 """
 
+    global INITIALIZED
+    INITIALIZED = True
     return output, context
 
 
@@ -257,6 +259,8 @@ def reset_wcgw(context: Context, reset_wcgw: ResetWcgw) -> str:
             list(context.bash_state.whitelist_for_overwrite),
             reset_wcgw.starting_directory,
         )
+    global INITIALIZED
+    INITIALIZED = True
     return "Reset successful" + get_status(context.bash_state)
 
 
@@ -595,7 +599,6 @@ def get_tool_output(
         context.console.print("Calling reset wcgw tool")
         output = reset_wcgw(context, arg), 0.0
 
-        INITIALIZED = True
     elif isinstance(arg, Initialize):
         context.console.print("Calling initial info tool")
         output_, context = initialize(
@@ -608,7 +611,6 @@ def get_tool_output(
         )
         output = output_, 0.0
 
-        INITIALIZED = True
     elif isinstance(arg, ContextSave):
         context.console.print("Calling task memory tool")
         relevant_files = []

{wcgw-3.0.1rc1 → wcgw-3.0.1rc2}/src/wcgw/relay/client.py

@@ -25,7 +25,7 @@ def register_client(server_url: str, client_uuid: str = "") -> None:
     # Create the WebSocket connection and context
     the_console = rich.console.Console(style="magenta", highlight=False, markup=False)
     with BashState(
-        the_console, os.getcwd(), None, None, None, None, False, None
+        the_console, os.getcwd(), None, None, None, None, True, None
     ) as bash_state:
         context = Context(bash_state=bash_state, console=the_console)
 

{wcgw-3.0.1rc1 → wcgw-3.0.1rc2}/src/wcgw_cli/anthropic_client.py

@@ -214,10 +214,9 @@ def loop(
     )
 
     with BashState(
-        system_console, os.getcwd(), None, None, None, None, False, None
+        system_console, os.getcwd(), None, None, None, None, True, None
     ) as bash_state:
         context = Context(bash_state, system_console)
-
         system, context = initialize(
             context,
             os.getcwd(),

{wcgw-3.0.1rc1 → wcgw-3.0.1rc2}/src/wcgw_cli/openai_client.py

@@ -178,7 +178,7 @@ def loop(
     )
 
     with BashState(
-        system_console, os.getcwd(), None, None, None, None, False, None
+        system_console, os.getcwd(), None, None, None, None, True, None
     ) as bash_state:
         context = Context(bash_state, system_console)
         system, context = initialize(