wcgw 2.7.2__py3-none-any.whl → 2.8.0__py3-none-any.whl

wcgw/client/anthropic_client.py

@@ -130,7 +130,7 @@ def loop(
     memory = None
     if resume:
         try:
-            _, memory = load_memory(
+            _, memory, _ = load_memory(
                 resume,
                 8000,
                 lambda x: default_enc.encode(x).ids,
@@ -166,6 +166,7 @@ def loop(
 - 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 "&".
 - Use longer wait_for_seconds if the command is expected to run for a long time.
+- Do not use 'cat' to read files, use ReadFiles tool instead.
 """,
         ),
         ToolParam(
@@ -281,22 +282,13 @@ Saves provided description and file contents of all the relevant file paths or g
             ),
         ]
 
-    initial_info = initialize(
-        os.getcwd(), [], resume if (memory and resume) else "", 8000
+    system = initialize(
+        os.getcwd(),
+        [],
+        resume if (memory and resume) else "",
+        max_tokens=8000,
+        mode="wcgw",
     )
-    system = f"""
-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 add/edit the code.
-    - Do not install new tools/packages before ensuring no such tools/package or an alternative already exists.
-
-{initial_info}
-"""
 
     with open(os.path.join(os.path.dirname(__file__), "diff-instructions.txt")) as f:
         system += f.read()

wcgw/client/mcp_server/server.py

@@ -28,6 +28,7 @@ from ...types_ import (
 )
 from .. import tools
 from ..computer_use import SLEEP_TIME_MAX_S
+from ..modes import get_kt_prompt
 from ..tools import DoneFlag, default_enc, get_tool_output, which_tool_name
 
 COMPUTER_USE_ON_DOCKER_ENABLED = False
@@ -45,48 +46,35 @@ async def handle_read_resource(uri: AnyUrl) -> str:
     raise ValueError("No resources available")
 
 
-@server.list_prompts()  # type: ignore
-async def handle_list_prompts() -> list[types.Prompt]:
-    return [
+PROMPTS = {
+    "KnowledgeTransfer": (
         types.Prompt(
             name="KnowledgeTransfer",
             description="Prompt for invoking ContextSave tool in order to do a comprehensive knowledge transfer of a coding task. Prompts to save detailed error log and instructions.",
-        )
-    ]
+        ),
+        get_kt_prompt,
+    )
+}
+
+
+@server.list_prompts()  # type: ignore
+async def handle_list_prompts() -> list[types.Prompt]:
+    return [x[0] for x in PROMPTS.values()]
 
 
 @server.get_prompt()  # type: ignore
 async def handle_get_prompt(
     name: str, arguments: dict[str, str] | None
 ) -> types.GetPromptResult:
-    messages = []
-    if name == "KnowledgeTransfer":
-        messages = [
-            types.PromptMessage(
-                role="user",
-                content=types.TextContent(
-                    type="text",
-                    text="""Use `ContextSave` tool to do a knowledge transfer of the task in hand.
-Write detailed description in order to do a KT.
-Save all information necessary for a person to understand the task and the problems.
-
-Format the `description` field using Markdown with the following sections.
-- "# Objective" section containing project and task objective.
-- "# All user instructions" section should be provided containing all instructions user shared in the conversation.
-- "# Current status of the task" should be provided containing only what is already achieved, not what's remaining.
-- "# All issues with snippets" section containing snippets of error, traceback, file snippets, commands, etc. But no comments or solutions.
-- Be very verbose in the all issues with snippets section providing as much error context as possible.
-- "# Build and development instructions" section containing instructions to build or run project or run tests, or envrionment related information. Only include what's known. Leave empty if unknown.
-- After the tool completes succesfully, tell me the task id and the file path the tool generated (important!)
-- This tool marks end of your conversation, do not run any further tools after calling this.
-
-Provide all relevant file paths in order to understand and solve the the task. Err towards providing more file paths than fewer.
-
-(Note to self: this conversation can then be resumed later asking "Resume `<generated id>`" which should call Initialize tool)
-""",
-                ),
-            )
-        ]
+    messages = [
+        types.PromptMessage(
+            role="user",
+            content=types.TextContent(
+                type="text",
+                text=PROMPTS[name][1](),
+            ),
+        )
+    ]
     return types.GetPromptResult(messages=messages)
 
 
@@ -117,6 +105,9 @@ async def handle_list_tools() -> list[types.Tool]:
 - If user has mentioned any files use `initial_files_to_read` to read, use absolute paths only.
 - If `any_workspace_path` is provided, a tree structure of the workspace will be shown.
 - Leave `any_workspace_path` as empty if no file or folder is mentioned.
+- By default use mode "wcgw"
+- In "code-writer" mode, set the commands and globs which user asked to set, otherwise use 'all'.
+- In order to change the mode later, call this tool again but be sure to not provide any other argument like task_id_to_resume unnecessarily.
 """,
         ),
         ToolParam(
@@ -263,6 +254,8 @@ async def handle_call_tool(
     except ValidationError:
 
         def try_json(x: str) -> Any:
+            if not isinstance(x, str):
+                return x
             try:
                 return json.loads(x)
             except json.JSONDecodeError:
@@ -284,24 +277,9 @@ async def handle_call_tool(
         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 add/edit the code.
-    - Do not install new tools/packages before ensuring no such tools/package or an alternative already exists.
-    - Do not use artifacts if you have access to the repository and not asked by the user to provide artifacts/snippets. Directly create/update using shell tools.
-    - Do not use Ctrl-c or Ctrl-z or interrupt commands without asking the user, because often the program don't show any update but they still are running.
-    - Do not use echo to write multi-line files, always use FileEdit tool to update a 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.
+- Additional important note: as soon as you encounter "The user has chosen to disallow the tool call.", immediately stop doing everything and ask user for the reason.
+
+Initialize call done.
     """
 
             content.append(types.TextContent(type="text", text=output_or_done))

wcgw/client/memory.py

@@ -1,7 +1,8 @@
+import json
 import os
 import re
 import shlex
-from typing import Callable, Optional
+from typing import Any, Callable, Optional
 
 from ..types_ import ContextSave
 
@@ -30,7 +31,11 @@ def format_memory(task_memory: ContextSave, relevant_files: str) -> str:
     return memory_data
 
 
-def save_memory(task_memory: ContextSave, relevant_files: str) -> str:
+def save_memory(
+    task_memory: ContextSave,
+    relevant_files: str,
+    bash_state_dict: Optional[dict[str, Any]] = None,
+) -> str:
     app_dir = get_app_dir_xdg()
     memory_dir = os.path.join(app_dir, "memory")
     os.makedirs(memory_dir, exist_ok=True)
@@ -45,6 +50,12 @@ def save_memory(task_memory: ContextSave, relevant_files: str) -> str:
     with open(memory_file_full, "w") as f:
         f.write(memory_data)
 
+    # Save bash state if provided
+    if bash_state_dict is not None:
+        state_file = os.path.join(memory_dir, f"{task_id}_bash_state.json")
+        with open(state_file, "w") as f:
+            json.dump(bash_state_dict, f, indent=2)
+
     return memory_file_full
 
 
@@ -53,7 +64,7 @@ def load_memory[T](
     max_tokens: Optional[int],
     encoder: Callable[[str], list[T]],
     decoder: Callable[[list[T]], str],
-) -> tuple[str, str]:
+) -> tuple[str, str, Optional[dict[str, Any]]]:
     app_dir = get_app_dir_xdg()
     memory_dir = os.path.join(app_dir, "memory")
     memory_file = os.path.join(memory_dir, f"{task_id}.txt")
@@ -75,4 +86,12 @@ def load_memory[T](
         parsed_ = shlex.split(matched_path)
         if parsed_ and len(parsed_) == 1:
             project_root_path = parsed_[0]
-    return project_root_path, data
+
+    # Try to load bash state if exists
+    state_file = os.path.join(memory_dir, f"{task_id}_bash_state.json")
+    bash_state: Optional[dict[str, Any]] = None
+    if os.path.exists(state_file):
+        with open(state_file) as f:
+            bash_state = json.load(f)
+
+    return project_root_path, data, bash_state

wcgw/client/modes.py

@@ -0,0 +1,242 @@
+from dataclasses import dataclass
+from typing import Any, Literal, NamedTuple
+
+from ..types_ import Modes, ModesConfig
+
+
+class BashCommandMode(NamedTuple):
+    bash_mode: Literal["normal_mode", "restricted_mode"]
+    allowed_commands: Literal["all", "none"]
+
+    def serialize(self) -> dict[str, Any]:
+        return {"bash_mode": self.bash_mode, "allowed_commands": self.allowed_commands}
+
+    @classmethod
+    def deserialize(cls, data: dict[str, Any]) -> "BashCommandMode":
+        return cls(data["bash_mode"], data["allowed_commands"])
+
+
+class FileEditMode(NamedTuple):
+    allowed_globs: Literal["all"] | list[str]
+
+    def serialize(self) -> dict[str, Any]:
+        return {"allowed_globs": self.allowed_globs}
+
+    @classmethod
+    def deserialize(cls, data: dict[str, Any]) -> "FileEditMode":
+        return cls(data["allowed_globs"])
+
+
+class WriteIfEmptyMode(NamedTuple):
+    allowed_globs: Literal["all"] | list[str]
+
+    def serialize(self) -> dict[str, Any]:
+        return {"allowed_globs": self.allowed_globs}
+
+    @classmethod
+    def deserialize(cls, data: dict[str, Any]) -> "WriteIfEmptyMode":
+        return cls(data["allowed_globs"])
+
+
+@dataclass
+class ModeImpl:
+    bash_command_mode: BashCommandMode
+    file_edit_mode: FileEditMode
+    write_if_empty_mode: WriteIfEmptyMode
+
+
+def code_writer_prompt(
+    allowed_file_edit_globs: Literal["all"] | list[str],
+    all_write_new_globs: Literal["all"] | list[str],
+    allowed_commands: Literal["all"] | list[str],
+) -> str:
+    base = """
+You have to run in "code_writer" mode.
+"""
+
+    path_prompt = """
+    - You are allowed to run FileEdit in the provided repository only.
+    """
+
+    if allowed_file_edit_globs != "all":
+        if allowed_file_edit_globs:
+            path_prompt = f"""
+    - You are allowed to run FileEdit for files matching only the following globs: {', '.join(allowed_file_edit_globs)}
+"""
+        else:
+            path_prompt = """
+    - You are not allowed to run FileEdit.
+"""
+    base += path_prompt
+
+    path_prompt = """
+    - You are allowed to run WriteIfEmpty in the provided repository only.
+    """
+
+    if all_write_new_globs != "all":
+        if all_write_new_globs:
+            path_prompt = f"""
+    - You are allowed to run WriteIfEmpty files matching only the following globs: {', '.join(allowed_file_edit_globs)}
+"""
+        else:
+            path_prompt = """
+    - You are not allowed to run WriteIfEmpty.
+"""
+    base += path_prompt
+
+    run_command_common = """
+    - Do not use Ctrl-c or Ctrl-z or interrupt commands without asking the user, because often the programs don't show any update but they still are running.
+    - Do not use echo to write multi-line files, always use FileEdit tool to update a code.
+    - Do not provide code snippets unless asked by the user, instead directly add/edit the code.
+    - 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.)
+    - Do not use artifacts if you have access to the repository and not asked by the user to provide artifacts/snippets. Directly create/update using wcgw tools.
+"""
+
+    command_prompt = f"""
+    - You are only allowed to run commands for project setup, code writing, editing, updating, testing, running and debugging related to the project.
+    - Do not run anything that adds or removes packages, changes system configuration or environment.
+{run_command_common}
+"""
+    if allowed_commands != "all":
+        if allowed_commands:
+            command_prompt = f"""
+    - You are only allowed to run the following commands: {', '.join(allowed_commands)}
+{run_command_common}
+"""
+        else:
+            command_prompt = """
+    - You are not allowed to run any commands.
+"""
+
+    base += command_prompt
+    return base
+
+
+WCGW_PROMPT = """
+---
+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.)
+    - Do not provide code snippets unless asked by the user, instead directly add/edit the code.
+    - Do not install new tools/packages before ensuring no such tools/package or an alternative already exists.
+    - Do not use artifacts if you have access to the repository and not asked by the user to provide artifacts/snippets. Directly create/update using wcgw tools
+    - Do not use Ctrl-c or Ctrl-z or interrupt commands without asking the user, because often the programs don't show any update but they still are running.
+    - Do not use echo to write multi-line files, always use FileEdit tool to update a 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.
+
+
+"""
+ARCHITECT_PROMPT = """You have to run in "architect" mode. This means
+- You are not allowed to edit or update any file. You are not allowed to create any file. 
+- You are not allowed to run any commands that may change disk, system configuration, packages or environment. Only read-only commands are allowed.
+- Only run commands that allows you to explore the repository, understand the system or read anything of relevance. 
+- Do not use Ctrl-c or Ctrl-z or interrupt commands without asking the user, because often the programs don't show any update but they still are running.
+- You are not allowed to change directory (bash will run in -r mode)
+
+Your response should be in self-critique and brainstorm style.
+- Read as many relevant files as possible. 
+- Be comprehensive in your understanding and search of relevant files.
+- First understand about the project by getting the folder structure (ignoring .git, node_modules, venv, etc.)
+"""
+
+
+DEFAULT_MODES: dict[Modes, ModeImpl] = {
+    Modes.wcgw: ModeImpl(
+        bash_command_mode=BashCommandMode("normal_mode", "all"),
+        write_if_empty_mode=WriteIfEmptyMode("all"),
+        file_edit_mode=FileEditMode("all"),
+    ),
+    Modes.architect: ModeImpl(
+        bash_command_mode=BashCommandMode("restricted_mode", "all"),
+        write_if_empty_mode=WriteIfEmptyMode([]),
+        file_edit_mode=FileEditMode([]),
+    ),
+    Modes.code_writer: ModeImpl(
+        bash_command_mode=BashCommandMode("normal_mode", "all"),
+        write_if_empty_mode=WriteIfEmptyMode("all"),
+        file_edit_mode=FileEditMode("all"),
+    ),
+}
+
+
+def modes_to_state(
+    mode: ModesConfig,
+) -> tuple[BashCommandMode, FileEditMode, WriteIfEmptyMode, Modes]:
+    # First get default mode config
+    if isinstance(mode, str):
+        mode_impl = DEFAULT_MODES[Modes[mode]]  # converts str to Modes enum
+        mode_name = Modes[mode]
+    else:
+        # For CodeWriterMode, use code_writer as base and override
+        mode_impl = DEFAULT_MODES[Modes.code_writer]
+        # Override with custom settings from CodeWriterMode
+        mode_impl = ModeImpl(
+            bash_command_mode=BashCommandMode(
+                mode_impl.bash_command_mode.bash_mode,
+                "all" if mode.allowed_commands else "none",
+            ),
+            file_edit_mode=FileEditMode(mode.allowed_globs),
+            write_if_empty_mode=WriteIfEmptyMode(mode.allowed_globs),
+        )
+        mode_name = Modes.code_writer
+    return (
+        mode_impl.bash_command_mode,
+        mode_impl.file_edit_mode,
+        mode_impl.write_if_empty_mode,
+        mode_name,
+    )
+
+
+WCGW_KT = """Use `ContextSave` tool to do a knowledge transfer of the task in hand.
+Write detailed description in order to do a KT.
+Save all information necessary for a person to understand the task and the problems.
+
+Format the `description` field using Markdown with the following sections.
+- "# Objective" section containing project and task objective.
+- "# All user instructions" section should be provided containing all instructions user shared in the conversation.
+- "# Current status of the task" should be provided containing only what is already achieved, not what's remaining.
+- "# All issues with snippets" section containing snippets of error, traceback, file snippets, commands, etc. But no comments or solutions.
+- Be very verbose in the all issues with snippets section providing as much error context as possible.
+- "# Build and development instructions" section containing instructions to build or run project or run tests, or envrionment related information. Only include what's known. Leave empty if unknown.
+- Any other relevant sections following the above.
+- After the tool completes succesfully, tell me the task id and the file path the tool generated (important!)
+- This tool marks end of your conversation, do not run any further tools after calling this.
+
+Provide all relevant file paths in order to understand and solve the the task. Err towards providing more file paths than fewer.
+
+(Note to self: this conversation can then be resumed later asking "Resume wcgw task `<generated id>`" which should call Initialize tool)
+"""
+
+
+ARCHITECT_KT = """Use `ContextSave` tool to do a knowledge transfer of the task in hand.
+Write detailed description in order to do a KT.
+Save all information necessary for a person to understand the task and the problems.
+
+Format the `description` field using Markdown with the following sections.
+- "# Objective" section containing project and task objective.
+- "# All user instructions" section should be provided containing all instructions user shared in the conversation.
+- "# Designed plan" should be provided containing the designed plan as discussed.
+- Any other relevant sections following the above.
+- After the tool completes succesfully, tell me the task id and the file path the tool generated (important!)
+- This tool marks end of your conversation, do not run any further tools after calling this.
+
+Provide all relevant file paths in order to understand and solve the the task. Err towards providing more file paths than fewer.
+
+(Note to self: this conversation can then be resumed later asking "Resume wcgw task `<generated id>`" which should call Initialize tool)
+"""
+
+KTS = {Modes.wcgw: WCGW_KT, Modes.architect: ARCHITECT_KT, Modes.code_writer: WCGW_KT}
+
+
+def get_kt_prompt() -> str:
+    from .tools import BASH_STATE
+
+    return KTS[BASH_STATE.mode]

wcgw/client/openai_client.py

@@ -125,7 +125,7 @@ def loop(
     memory = None
     if resume:
         try:
-            _, memory = load_memory(
+            _, memory, _ = load_memory(
                 resume,
                 8000,
                 lambda x: default_enc.encode(x).ids,
@@ -176,6 +176,7 @@ def loop(
 - 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 "&".
+- Do not use 'cat' to read files, use ReadFiles tool instead.
 """,
         ),
         openai.pydantic_function_tool(
@@ -226,23 +227,13 @@ Saves provided description and file contents of all the relevant file paths or g
         ),
     ]
 
-    initial_info = initialize(
-        os.getcwd(), [], resume if (memory and resume) else "", 8000
+    system = initialize(
+        os.getcwd(),
+        [],
+        resume if (memory and resume) else "",
+        max_tokens=8000,
+        mode="wcgw",
     )
-    system = f"""
-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 add/edit the code.
-    - Do not install new tools/packages before ensuring no such tools/package or an alternative already exists.
-
-{initial_info}
-
-"""
 
     with open(os.path.join(os.path.dirname(__file__), "diff-instructions.txt")) as f:
         system += f.read()

wcgw/client/tools.py

@@ -1,5 +1,6 @@
 import base64
 import datetime
+import fnmatch
 import glob
 import importlib.metadata
 import json
@@ -14,6 +15,7 @@ from os.path import expanduser
 from pathlib import Path
 from tempfile import NamedTemporaryFile, TemporaryDirectory
 from typing import (
+    Any,
     Callable,
     Literal,
     Optional,
@@ -39,12 +41,15 @@ from websockets.sync.client import connect as syncconnect
 from ..types_ import (
     BashCommand,
     BashInteraction,
+    CodeWriterMode,
     ContextSave,
     FileEdit,
     FileEditFindReplace,
     GetScreenInfo,
     Initialize,
     Keyboard,
+    Modes,
+    ModesConfig,
     Mouse,
     ReadFiles,
     ReadImage,
@@ -55,6 +60,15 @@ from ..types_ import (
 from .computer_use import run_computer_tool
 from .file_ops.search_replace import search_replace_edit
 from .memory import load_memory, save_memory
+from .modes import (
+    ARCHITECT_PROMPT,
+    WCGW_PROMPT,
+    BashCommandMode,
+    FileEditMode,
+    WriteIfEmptyMode,
+    code_writer_prompt,
+    modes_to_state,
+)
 from .repo_ops.repo_context import get_repo_context
 from .sys_utils import command_run
 
@@ -123,14 +137,19 @@ PROMPT_CONST = "#@wcgw@#"
 PROMPT = PROMPT_CONST
 
 
-def start_shell() -> pexpect.spawn:  # type: ignore
+def start_shell(is_restricted_mode: bool, initial_dir: str) -> pexpect.spawn:  # type: ignore
     try:
+        cmd = "/bin/bash"
+        if is_restricted_mode:
+            cmd += " -r"
+
         shell = pexpect.spawn(
-            "/bin/bash",
+            cmd,
             env={**os.environ, **{"PS1": PROMPT}},  # type: ignore[arg-type]
             echo=False,
             encoding="utf-8",
             timeout=TIMEOUT,
+            cwd=initial_dir,
         )
         shell.sendline(f"export PS1={PROMPT}")
     except Exception as e:
@@ -203,15 +222,54 @@ BASH_CLF_OUTPUT = Literal["repl", "pending"]
 
 
 class BashState:
-    def __init__(self) -> None:
-        self._init()
+    def __init__(
+        self,
+        working_dir: str,
+        bash_command_mode: Optional[BashCommandMode],
+        file_edit_mode: Optional[FileEditMode],
+        write_if_empty_mode: Optional[WriteIfEmptyMode],
+        mode: Optional[Modes],
+        whitelist_for_overwrite: Optional[set[str]] = None,
+    ) -> None:
+        self._cwd = working_dir or os.getcwd()
+        self._bash_command_mode: BashCommandMode = bash_command_mode or BashCommandMode(
+            "normal_mode", "all"
+        )
+        self._file_edit_mode: FileEditMode = file_edit_mode or FileEditMode("all")
+        self._write_if_empty_mode: WriteIfEmptyMode = (
+            write_if_empty_mode or WriteIfEmptyMode("all")
+        )
+        self._mode = mode or Modes.wcgw
+        self._whitelist_for_overwrite: set[str] = whitelist_for_overwrite or set()
+
+        self._init_shell()
+
+    @property
+    def mode(self) -> Modes:
+        return self._mode
 
-    def _init(self) -> None:
+    @property
+    def bash_command_mode(self) -> BashCommandMode:
+        return self._bash_command_mode
+
+    @property
+    def file_edit_mode(self) -> FileEditMode:
+        return self._file_edit_mode
+
+    @property
+    def write_if_empty_mode(self) -> WriteIfEmptyMode:
+        return self._write_if_empty_mode
+
+    def _init_shell(self) -> None:
         self._state: Literal["repl"] | datetime.datetime = "repl"
         self._is_in_docker: Optional[str] = ""
-        self._cwd: str = os.getcwd()
-        self._shell = start_shell()
-        self._whitelist_for_overwrite: set[str] = set()
+        # Ensure self._cwd exists
+        os.makedirs(self._cwd, exist_ok=True)
+        self._shell = start_shell(
+            self._bash_command_mode.bash_mode == "restricted_mode",
+            self._cwd,
+        )
+
         self._pending_output = ""
 
         # Get exit info to ensure shell is ready
@@ -258,9 +316,49 @@ class BashState:
         self._cwd = current_dir
         return current_dir
 
-    def reset(self) -> None:
+    def reset_shell(self) -> None:
         self.shell.close(True)
-        self._init()
+        self._init_shell()
+
+    def serialize(self) -> dict[str, Any]:
+        """Serialize BashState to a dictionary for saving"""
+        return {
+            "bash_command_mode": self._bash_command_mode.serialize(),
+            "file_edit_mode": self._file_edit_mode.serialize(),
+            "write_if_empty_mode": self._write_if_empty_mode.serialize(),
+            "whitelist_for_overwrite": list(self._whitelist_for_overwrite),
+            "mode": self._mode,
+        }
+
+    @staticmethod
+    def parse_state(
+        state: dict[str, Any],
+    ) -> tuple[BashCommandMode, FileEditMode, WriteIfEmptyMode, Modes, list[str]]:
+        return (
+            BashCommandMode.deserialize(state["bash_command_mode"]),
+            FileEditMode.deserialize(state["file_edit_mode"]),
+            WriteIfEmptyMode.deserialize(state["write_if_empty_mode"]),
+            Modes[str(state["mode"])],
+            state["whitelist_for_overwrite"],
+        )
+
+    def load_state(
+        self,
+        bash_command_mode: BashCommandMode,
+        file_edit_mode: FileEditMode,
+        write_if_empty_mode: WriteIfEmptyMode,
+        mode: Modes,
+        whitelist_for_overwrite: list[str],
+        cwd: str,
+    ) -> None:
+        """Create a new BashState instance from a serialized state dictionary"""
+        self._bash_command_mode = bash_command_mode
+        self._cwd = cwd or self._cwd
+        self._file_edit_mode = file_edit_mode
+        self._write_if_empty_mode = write_if_empty_mode
+        self._whitelist_for_overwrite = set(whitelist_for_overwrite)
+        self._mode = mode
+        self.reset_shell()
 
     def get_pending_for(self) -> str:
         if isinstance(self._state, datetime.datetime):
@@ -290,7 +388,8 @@ class BashState:
         return self._pending_output
 
 
-BASH_STATE = BashState()
+BASH_STATE = BashState(os.getcwd(), None, None, None, None)
+INITIALIZED = False
 
 
 def initialize(
@@ -298,57 +397,122 @@ def initialize(
     read_files_: list[str],
     task_id_to_resume: str,
     max_tokens: Optional[int],
+    mode: ModesConfig,
 ) -> str:
-    reset_shell()
+    global BASH_STATE
 
     # Expand the workspace path
-    any_workspace_path = expand_user(any_workspace_path, BASH_STATE.is_in_docker)
+    any_workspace_path = expand_user(any_workspace_path, None)
     repo_context = ""
 
     memory = ""
+    bash_state = None
     if task_id_to_resume:
         try:
-            project_root_path, task_mem = load_memory(
+            project_root_path, task_mem, bash_state = load_memory(
                 task_id_to_resume,
                 max_tokens,
                 lambda x: default_enc.encode(x).ids,
                 lambda x: default_enc.decode(x),
             )
             memory = "Following is the retrieved task:\n" + task_mem
-            if (
-                not any_workspace_path or not os.path.exists(any_workspace_path)
-            ) and os.path.exists(project_root_path):
+            if os.path.exists(project_root_path):
                 any_workspace_path = project_root_path
+
         except Exception:
             memory = f'Error: Unable to load task with ID "{task_id_to_resume}" '
 
+    folder_to_start = None
     if any_workspace_path:
         if os.path.exists(any_workspace_path):
             repo_context, folder_to_start = get_repo_context(any_workspace_path, 200)
 
-            BASH_STATE.shell.sendline(f"cd {shlex.quote(str(folder_to_start))}")
-            BASH_STATE.shell.expect(PROMPT, timeout=0.2)
-            BASH_STATE.update_cwd()
-
             repo_context = f"---\n# Workspace structure\n{repo_context}\n---\n"
+
+            # update modes if they're relative
+            if isinstance(mode, CodeWriterMode):
+                mode.update_relative_globs(any_workspace_path)
+            else:
+                assert isinstance(mode, str)
         else:
-            repo_context = (
-                f"\nInfo: Workspace path {any_workspace_path} does not exist\n"
-            )
+            if os.path.abspath(any_workspace_path):
+                os.makedirs(any_workspace_path, exist_ok=True)
+                repo_context = f"\nInfo: Workspace path {any_workspace_path} did not exist. I've created it for you.\n"
+            else:
+                repo_context = (
+                    f"\nInfo: Workspace path {any_workspace_path} does not exist."
+                )
+    # Restore bash state if available
+    if bash_state is not None:
+        try:
+            parsed_state = BashState.parse_state(bash_state)
+            if mode == "wcgw":
+                BASH_STATE.load_state(
+                    parsed_state[0],
+                    parsed_state[1],
+                    parsed_state[2],
+                    parsed_state[3],
+                    parsed_state[4] + list(BASH_STATE.whitelist_for_overwrite),
+                    str(folder_to_start) if folder_to_start else "",
+                )
+            else:
+                state = modes_to_state(mode)
+                BASH_STATE.load_state(
+                    state[0],
+                    state[1],
+                    state[2],
+                    state[3],
+                    parsed_state[4] + list(BASH_STATE.whitelist_for_overwrite),
+                    str(folder_to_start) if folder_to_start else "",
+                )
+        except ValueError:
+            console.print(traceback.format_exc())
+            console.print("Error: couldn't load bash state")
+            pass
+    else:
+        state = modes_to_state(mode)
+        BASH_STATE.load_state(
+            state[0],
+            state[1],
+            state[2],
+            state[3],
+            list(BASH_STATE.whitelist_for_overwrite),
+            str(folder_to_start) if folder_to_start else "",
+        )
+    del mode
 
     initial_files_context = ""
     if read_files_:
+        if folder_to_start:
+            read_files_ = [
+                os.path.join(folder_to_start, f) if not os.path.isabs(f) else f
+                for f in read_files_
+            ]
         initial_files = read_files(read_files_, max_tokens)
         initial_files_context = f"---\n# Requested files\n{initial_files}\n---\n"
 
     uname_sysname = os.uname().sysname
     uname_machine = os.uname().machine
 
+    mode_prompt = ""
+    if BASH_STATE.mode == Modes.code_writer:
+        mode_prompt = code_writer_prompt(
+            BASH_STATE.file_edit_mode.allowed_globs,
+            BASH_STATE.write_if_empty_mode.allowed_globs,
+            "all" if BASH_STATE.bash_command_mode.allowed_commands else [],
+        )
+    elif BASH_STATE.mode == Modes.architect:
+        mode_prompt = ARCHITECT_PROMPT
+    else:
+        mode_prompt = WCGW_PROMPT
+
     output = f"""
+{mode_prompt}
+
 # Environment
 System: {uname_sysname}
 Machine: {uname_machine}
-Current working directory: {BASH_STATE.cwd}
+Initialized in directory (also cwd): {BASH_STATE.cwd}
 
 {repo_context}
 
@@ -359,11 +523,14 @@ Current working directory: {BASH_STATE.cwd}
 {memory}
 """
 
+    global INITIALIZED
+    INITIALIZED = True
+
     return output
 
 
 def reset_shell() -> str:
-    BASH_STATE.reset()
+    BASH_STATE.reset_shell()
     return "Reset successful" + get_status()
 
 
@@ -472,6 +639,8 @@ def execute_bash(
     try:
         is_interrupt = False
         if isinstance(bash_arg, BashCommand):
+            if BASH_STATE.bash_command_mode.allowed_commands == "none":
+                return "Error: BashCommand not allowed in current mode", 0.0
             updated_repl_mode = update_repl_prompt(bash_arg.command)
             if updated_repl_mode:
                 BASH_STATE.set_repl()
@@ -647,31 +816,11 @@ def execute_bash(
         console.print(traceback.format_exc())
         console.print("Malformed output, restarting shell", style="red")
         # Malformed output, restart shell
-        BASH_STATE.reset()
+        BASH_STATE.reset_shell()
         output = "(exit shell has restarted)"
     return output, 0
 
 
-def serve_image_in_bg(file_path: str, client_uuid: str, name: str) -> None:
-    if not client_uuid:
-        client_uuid = str(uuid.uuid4())
-
-    server_url = "wss://wcgw.arcfu.com/register_serve_image"
-
-    with open(file_path, "rb") as image_file:
-        image_bytes = image_file.read()
-        media_type = mimetypes.guess_type(file_path)[0]
-        image_b64 = base64.b64encode(image_bytes).decode("utf-8")
-        uu = {"name": name, "image_b64": image_b64, "media_type": media_type}
-
-    with syncconnect(f"{server_url}/{client_uuid}") as websocket:
-        try:
-            websocket.send(json.dumps(uu))
-        except websockets.ConnectionClosed:
-            console.print(f"Connection closed for UUID: {client_uuid}, retrying")
-            serve_image_in_bg(file_path, client_uuid, name)
-
-
 MEDIA_TYPES = Literal["image/jpeg", "image/png", "image/gif", "image/webp"]
 
 
@@ -766,6 +915,13 @@ def write_file(
         path_ = expand_user(writefile.file_path, BASH_STATE.is_in_docker)
 
     error_on_exist_ = error_on_exist and path_ not in BASH_STATE.whitelist_for_overwrite
+
+    # Validate using write_if_empty_mode after checking whitelist
+    allowed_globs = BASH_STATE.write_if_empty_mode.allowed_globs
+    if allowed_globs != "all" and not any(
+        fnmatch.fnmatch(path_, pattern) for pattern in allowed_globs
+    ):
+        return f"Error: updating file {path_} not allowed in current mode. Doesn't match allowed globs: {allowed_globs}"
     add_overwrite_warning = ""
     if not BASH_STATE.is_in_docker:
         if (error_on_exist or error_on_exist_) and os.path.exists(path_):
@@ -897,6 +1053,15 @@ def _do_diff_edit(fedit: FileEdit, max_tokens: Optional[int]) -> str:
     else:
         path_ = expand_user(fedit.file_path, BASH_STATE.is_in_docker)
 
+    # Validate using file_edit_mode
+    allowed_globs = BASH_STATE.file_edit_mode.allowed_globs
+    if allowed_globs != "all" and not any(
+        fnmatch.fnmatch(path_, pattern) for pattern in allowed_globs
+    ):
+        raise Exception(
+            f"Error: updating file {path_} not allowed in current mode. Doesn't match allowed globs: {allowed_globs}"
+        )
+
     # The LLM is now aware that the file exists
     BASH_STATE.add_to_whitelist_for_overwrite(path_)
 
@@ -1063,7 +1228,7 @@ def get_tool_output(
     loop_call: Callable[[str, float], tuple[str, float]],
     max_tokens: Optional[int],
 ) -> tuple[list[str | ImageData | DoneFlag], float]:
-    global IS_IN_DOCKER, TOOL_CALLS
+    global IS_IN_DOCKER, TOOL_CALLS, INITIALIZED
     if isinstance(args, dict):
         adapter = TypeAdapter[TOOLS](TOOLS, config={"extra": "forbid"})
         arg = adapter.validate_python(args)
@@ -1071,17 +1236,27 @@ def get_tool_output(
         arg = args
     output: tuple[str | DoneFlag | ImageData, float]
     TOOL_CALLS.append(arg)
+
     if isinstance(arg, Confirmation):
         console.print("Calling ask confirmation tool")
         output = ask_confirmation(arg), 0.0
     elif isinstance(arg, (BashCommand | BashInteraction)):
         console.print("Calling execute bash tool")
+        if not INITIALIZED:
+            raise Exception("Initialize tool not called yet.")
+
         output = execute_bash(enc, arg, max_tokens, arg.wait_for_seconds)
     elif isinstance(arg, WriteIfEmpty):
         console.print("Calling write file tool")
+        if not INITIALIZED:
+            raise Exception("Initialize tool not called yet.")
+
         output = write_file(arg, True, max_tokens), 0
     elif isinstance(arg, FileEdit):
         console.print("Calling full file edit tool")
+        if not INITIALIZED:
+            raise Exception("Initialize tool not called yet.")
+
         output = do_diff_edit(arg, max_tokens), 0.0
     elif isinstance(arg, DoneFlag):
         console.print("Calling mark finish tool")
@@ -1106,6 +1281,7 @@ def get_tool_output(
                 arg.initial_files_to_read,
                 arg.task_id_to_resume,
                 max_tokens,
+                arg.mode,
             ),
             0.0,
         )
@@ -1158,7 +1334,7 @@ def get_tool_output(
             if not globs:
                 warnings += f"Warning: No files found for the glob: {fglob}\n"
         relevant_files_data = read_files(relevant_files[:10_000], None)
-        output_ = save_memory(arg, relevant_files_data)
+        output_ = save_memory(arg, relevant_files_data, BASH_STATE.serialize())
         if not relevant_files and arg.relevant_file_globs:
             output_ = f'Error: No files found for the given globs. Context file successfully saved at "{output_}", but please fix the error.'
         elif warnings:

wcgw/types_.py

@@ -1,4 +1,6 @@
-from typing import Literal, Optional, Sequence
+import os
+from enum import Enum
+from typing import Any, Literal, Optional, Sequence, Union
 
 from pydantic import BaseModel as PydanticBaseModel
 
@@ -11,6 +13,54 @@ class NoExtraArgs(PydanticBaseModel):
 BaseModel = NoExtraArgs
 
 
+class Modes(str, Enum):
+    wcgw = "wcgw"
+    architect = "architect"
+    code_writer = "code_writer"
+
+
+class CodeWriterMode(BaseModel):
+    allowed_globs: Literal["all"] | list[str]
+    allowed_commands: Literal["all"] | list[str]
+
+    def update_relative_globs(self, workspace_root: str) -> None:
+        """Update globs if they're relative paths"""
+        if self.allowed_globs != "all":
+            self.allowed_globs = [
+                glob if os.path.isabs(glob) else os.path.join(workspace_root, glob)
+                for glob in self.allowed_globs
+            ]
+
+
+ModesConfig = Union[Literal["wcgw", "architect"], CodeWriterMode]
+
+
+class Initialize(BaseModel):
+    any_workspace_path: str
+    initial_files_to_read: list[str]
+    task_id_to_resume: str
+    mode_name: Literal["wcgw", "architect", "code_writer"]
+    code_writer_config: Optional[CodeWriterMode] = None
+
+    def model_post_init(self, __context: Any) -> None:
+        if self.mode_name == "code_writer":
+            assert (
+                self.code_writer_config is not None
+            ), "code_writer_config can't be null when the mode is code_writer"
+        return super().model_post_init(__context)
+
+    @property
+    def mode(self) -> ModesConfig:
+        if self.mode_name == "wcgw":
+            return "wcgw"
+        if self.mode_name == "architect":
+            return "architect"
+        assert (
+            self.code_writer_config is not None
+        ), "code_writer_config can't be null when the mode is code_writer"
+        return self.code_writer_config
+
+
 class BashCommand(BaseModel):
     command: str
     wait_for_seconds: Optional[int] = None
@@ -56,12 +106,6 @@ class FileEdit(BaseModel):
     file_edit_using_search_replace_blocks: str
 
 
-class Initialize(BaseModel):
-    any_workspace_path: str
-    initial_files_to_read: list[str]
-    task_id_to_resume: str
-
-
 class GetScreenInfo(BaseModel):
     docker_image_id: str
 

{wcgw-2.7.2.dist-info → wcgw-2.8.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wcgw
-Version: 2.7.2
+Version: 2.8.0
 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-2.7.2.dist-info → wcgw-2.8.0.dist-info}/RECORD

@@ -1,22 +1,23 @@
 wcgw/__init__.py,sha256=9K2QW7QuSLhMTVbKbBYd9UUp-ZyrfBrxcjuD_xk458k,118
-wcgw/types_.py,sha256=6gqdwC0F_rAQ80uPQpTWnw7uUByAD9qPa6-hKFgxqUU,1946
+wcgw/types_.py,sha256=D3518a2azSKeW3D-ACYWJwsaqo7Oj-8BRRR2IhCUtNU,3414
 wcgw/client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 wcgw/client/__main__.py,sha256=wcCrL4PjG51r5wVKqJhcoJPTLfHW0wNbD31DrUN0MWI,28
-wcgw/client/anthropic_client.py,sha256=p-QsqjTl5LJxCFL5GfPNlQCRa4Bh4XM6MWfkaYJj4qc,21439
+wcgw/client/anthropic_client.py,sha256=TLPQCc96cxLGGMHg1sYqX3_IxTm_-WQ1e-XPQD6i6xU,20958
 wcgw/client/cli.py,sha256=-z0kpDAW3mzfQrQeZfaVJhBCAQY3HXnt9GdgQ8s-u0Y,1003
 wcgw/client/common.py,sha256=OCH7Tx64jojz3M3iONUrGMadE07W21DiZs5sOxWX1Qc,1456
 wcgw/client/computer_use.py,sha256=35NKAlMrxwD0TBlMMRnbCwz4g8TBRGOlcy-cmS-yJ_A,15247
 wcgw/client/diff-instructions.txt,sha256=tmJ9Fu9XdO_72lYXQQNY9RZyx91bjxrXJf9d_KBz57k,1611
-wcgw/client/memory.py,sha256=x7AsZo2GwYyEmnxtTvmQJBlUxD_26LuENQbCbNhTGyY,2272
-wcgw/client/openai_client.py,sha256=UCiamF27iVhYh5fI5Kpa1IIMMXI6R5uPtRDjdKyHhOQ,18340
+wcgw/client/memory.py,sha256=8LdYsOhvCOoC1kfvDr85kNy07WnhPMvE6B2FRM2w85Y,2902
+wcgw/client/modes.py,sha256=ZhCY-Orbk2JtbxbrPz-YNqLO6PA_zNvE1uNCD3JmnYA,10463
+wcgw/client/openai_client.py,sha256=WvCnTij5QOPhe0VCdsk5Qo1Ob6yBdxFSvsCmnF0Jozw,17858
 wcgw/client/openai_utils.py,sha256=KfMB1-p2zDiA7pPWwAVarochf7-qeL1UMgtlDV9DtKA,2662
 wcgw/client/sys_utils.py,sha256=GajPntKhaTUMn6EOmopENWZNR2G_BJyuVbuot0x6veI,1376
-wcgw/client/tools.py,sha256=H9_OPkQOB8QMi6BumRow1JTGCffRVNTEmeagUTF3HE4,44930
+wcgw/client/tools.py,sha256=9_pbB1K83teY9kiYs1i_bUsKIYSv7kGhi2-VxLZOFOI,51034
 wcgw/client/file_ops/diff_edit.py,sha256=o0ucArVwn5p6QTDgYsjLfMy4TJXxUG3qcppFBNF3bbQ,16751
 wcgw/client/file_ops/search_replace.py,sha256=89ieDC9fTsIKPDx7CJwnwpX32dRdSlMKoBtKVXc7VWI,3971
 wcgw/client/mcp_server/Readme.md,sha256=I8N4dHkTUVGNQ63BQkBMBhCCBTgqGOSF_pUR6iOEiUk,2495
 wcgw/client/mcp_server/__init__.py,sha256=hyPPwO9cabAJsOMWhKyat9yl7OlSmIobaoAZKHu3DMc,381
-wcgw/client/mcp_server/server.py,sha256=UWUHBlu7UllR4q3xgyHQLeoUSLKmtkGc9nNrZaoKYe4,14905
+wcgw/client/mcp_server/server.py,sha256=C-wqaoS7sGWhQQW0pKS3HLjsUwjo2zzDJZm8d6znozU,12812
 wcgw/client/repo_ops/display_tree.py,sha256=5FD4hfMkM2cIZnXlu7WfJswJLthj0SkuHlkGH6dpWQU,4632
 wcgw/client/repo_ops/path_prob.py,sha256=SWf0CDn37rtlsYRQ51ufSxay-heaQoVIhr1alB9tZ4M,2144
 wcgw/client/repo_ops/paths_model.vocab,sha256=M1pXycYDQehMXtpp-qAgU7rtzeBbCOiJo4qcYFY0kqk,315087
@@ -46,8 +47,8 @@ mcp_wcgw/shared/memory.py,sha256=dBsOghxHz8-tycdSVo9kSujbsC8xb_tYsGmuJobuZnw,281
 mcp_wcgw/shared/progress.py,sha256=ymxOsb8XO5Mhlop7fRfdbmvPodANj7oq6O4dD0iUcnw,1048
 mcp_wcgw/shared/session.py,sha256=e44a0LQOW8gwdLs9_DE9oDsxqW2U8mXG3d5KT95bn5o,10393
 mcp_wcgw/shared/version.py,sha256=d2LZii-mgsPIxpshjkXnOTUmk98i0DT4ff8VpA_kAvE,111
-wcgw-2.7.2.dist-info/METADATA,sha256=Cn7oNhfpRHJ2rMyFvcyJgxKe2Xa_iFHJrF1CE_fcvkY,9408
-wcgw-2.7.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-wcgw-2.7.2.dist-info/entry_points.txt,sha256=eKo1omwbAggWlQ0l7GKoR7uV1-j16nk9tK0BhC2Oz_E,120
-wcgw-2.7.2.dist-info/licenses/LICENSE,sha256=BvY8xqjOfc3X2qZpGpX3MZEmF-4Dp0LqgKBbT6L_8oI,11142
-wcgw-2.7.2.dist-info/RECORD,,
+wcgw-2.8.0.dist-info/METADATA,sha256=3Mr_eVjHoBtbPcswgwoQ8PCrsHJWhr4WptUxrhhU27A,9408
+wcgw-2.8.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+wcgw-2.8.0.dist-info/entry_points.txt,sha256=eKo1omwbAggWlQ0l7GKoR7uV1-j16nk9tK0BhC2Oz_E,120
+wcgw-2.8.0.dist-info/licenses/LICENSE,sha256=BvY8xqjOfc3X2qZpGpX3MZEmF-4Dp0LqgKBbT6L_8oI,11142
+wcgw-2.8.0.dist-info/RECORD,,