wcgw 2.8.10__py3-none-any.whl → 3.0.1__py3-none-any.whl

wcgw/client/encoder/__init__.py

@@ -0,0 +1,47 @@
+import threading
+from typing import Callable, Protocol, TypeVar, cast
+
+import tokenizers  # type: ignore[import-untyped]
+
+T = TypeVar("T")
+
+
+class EncoderDecoder(Protocol[T]):
+    def encoder(self, text: str) -> list[T]: ...
+
+    def decoder(self, tokens: list[T]) -> str: ...
+
+
+class LazyEncoder:
+    def __init__(self) -> None:
+        self._tokenizer: tokenizers.Tokenizer | None = None
+        self._init_lock = threading.Lock()
+        self._init_thread = threading.Thread(target=self._initialize, daemon=True)
+        self._init_thread.start()
+
+    def _initialize(self) -> None:
+        with self._init_lock:
+            if self._tokenizer is None:
+                self._tokenizer = tokenizers.Tokenizer.from_pretrained(
+                    "Xenova/claude-tokenizer"
+                )
+
+    def _ensure_initialized(self) -> None:
+        if self._tokenizer is None:
+            with self._init_lock:
+                if self._tokenizer is None:
+                    self._init_thread.join()
+
+    def encoder(self, text: str) -> list[int]:
+        self._ensure_initialized()
+        assert self._tokenizer is not None, "Couldn't initialize tokenizer"
+        return cast(list[int], self._tokenizer.encode(text).ids)
+
+    def decoder(self, tokens: list[int]) -> str:
+        self._ensure_initialized()
+        assert self._tokenizer is not None, "Couldn't initialize tokenizer"
+        return cast(str, self._tokenizer.decode(tokens))
+
+
+def get_default_encoder() -> EncoderDecoder[int]:
+    return LazyEncoder()

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/client/mcp_server/Readme.md

@@ -1,89 +1,3 @@
-# Claude desktop support
+# The doc has moved to main Readme.md
 
-`wcgw` enables Claude desktop app on Mac to access shell and file system in order to automate tasks, run code, etc.
-
-It also has a computer use feature to connect to linux running on docker. Claude can fully control it including mouse and keyboard.
-
-## Setup
-
-Update `claude_desktop_config.json` (~/Library/Application Support/Claude/claude_desktop_config.json)
-
-```json
-{
-  "mcpServers": {
-    "wcgw": {
-      "command": "uv",
-      "args": [
-        "tool",
-        "run",
-        "--from",
-        "wcgw@latest",
-        "--python",
-        "3.12",
-        "wcgw_mcp"
-      ]
-    }
-  }
-}
-```
-
-Then restart claude app.
-
-### [Optional] Computer use support using desktop on docker
-
-Computer use is disabled by default. Add `--computer-use` to enable it. This will add necessary tools to Claude including ScreenShot, Mouse and Keyboard control.
-
-```json
-{
-  "mcpServers": {
-    "wcgw": {
-      "command": "uv",
-      "args": [
-        "tool",
-        "run",
-        "--from",
-        "wcgw@latest",
-        "--python",
-        "3.12",
-        "wcgw_mcp",
-        "--computer-use"
-      ]
-    }
-  }
-}
-```
-
-Claude will be able to connect to any docker container with linux environment. Native system control isn't supported outside docker.
-
-You'll need to run a docker image with desktop and optional VNC connection. Here's a demo image:
-
-```sh
-docker run -p 6080:6080 ghcr.io/anthropics/anthropic-quickstarts:computer-use-demo-latest
-```
-
-Then ask claude desktop app to control the docker os. It'll connect to the docker container and control it.
-
-Connect to `http://localhost:6080/vnc.html` for desktop view (VNC) of the system running in the docker.
-
-## Usage
-
-Wait for a few seconds. 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)
-over here
-
-![mcp icon](https://github.com/rusiaaman/wcgw/blob/main/static/claude-ss.jpg?raw=true)
-
-Then ask claude to execute shell commands, read files, edit files, run your code, etc.
-
-If you've run the docker for LLM to access, you can ask it to control the "docker os". If you don't provide the docker container id to it, it'll try to search for available docker using `docker ps` command.
-
-## Example
-
-### Computer use example
-
-![computer-use](https://github.com/rusiaaman/wcgw/blob/main/static/computer-use.jpg?raw=true)
-
-### Shell example
-
-![example](https://github.com/rusiaaman/wcgw/blob/main/static/example.jpg?raw=true)
+![main readme](https://github.com/rusiaaman/wcgw/blob/main/README.md)

wcgw/client/mcp_server/__init__.py

@@ -7,9 +7,9 @@ main = Typer()
 
 
 @main.command()
-def app(computer_use: bool = False) -> None:
+def app() -> None:
     """Main entry point for the package."""
-    asyncio.run(server.main(computer_use))
+    asyncio.run(server.main())
 
 
 # Optionally expose other important items at package level

wcgw/client/mcp_server/server.py

@@ -1,5 +1,4 @@
 import importlib
-import json
 import logging
 import os
 from typing import Any
@@ -9,29 +8,22 @@ 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
 
 from ...types_ import (
-    BashCommand,
-    BashInteraction,
-    ContextSave,
-    FileEdit,
-    GetScreenInfo,
     Initialize,
-    Keyboard,
-    Mouse,
-    ReadFiles,
-    ReadImage,
-    ResetShell,
-    ScreenShot,
-    WriteIfEmpty,
 )
-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
+from ..bash_state.bash_state import CONFIG, BashState
+from ..tools import (
+    Context,
+    default_enc,
+    get_tool_output,
+    parse_tool_by_name,
+    which_tool_name,
+)
 
 server = Server("wcgw")
 
@@ -64,7 +56,7 @@ PROMPTS = {
             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,
+        KTS,
     )
 }
 
@@ -78,12 +70,12 @@ async def handle_list_prompts() -> list[types.Prompt]:
 async def handle_get_prompt(
     name: str, arguments: dict[str, str] | None
 ) -> types.GetPromptResult:
+    assert BASH_STATE
     messages = [
         types.PromptMessage(
             role="user",
             content=types.TextContent(
-                type="text",
-                text=PROMPTS[name][1](),
+                type="text", text=PROMPTS[name][1][BASH_STATE.mode]
             ),
         )
     ]
@@ -97,187 +89,37 @@ async def handle_list_tools() -> list[types.Tool]:
     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()
-
-    tools = [
-        ToolParam(
-            inputSchema=Initialize.model_json_schema(),
-            name="Initialize",
-            description="""
-- Always call this at the start of the conversation before using any of the shell tools from wcgw.
-- This will reset the shell.
-- Use `any_workspace_path` to initialize the shell in the appropriate project directory.
-- If the user has mentioned a workspace or project root, use it to set `any_workspace_path`.
-- If the user has mentioned a folder or file with unclear project root, use the file or folder as `any_workspace_path`.
-- 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(
-            inputSchema=BashCommand.model_json_schema(),
-            name="BashCommand",
-            description=f"""
-- 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 or the last 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 {SLEEP_TIME_MAX_S} 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(
-            inputSchema=BashInteraction.model_json_schema(),
-            name="BashInteraction",
-            description=f"""
-- 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.
-- This returns within {SLEEP_TIME_MAX_S} seconds, for heavy programs keep checking status for upto 10 turns before asking user to continue checking again.
-- Programs don't hang easily, so most likely explanation for no output is usually that the program is still running, and you need to check status again using ["Enter"].
-- Do not send Ctrl-c before checking for status till 10 minutes or whatever is appropriate for the program to finish.
-- Set longer wait_for_seconds when program is expected to run for a long time.
-""",
-        ),
-        ToolParam(
-            inputSchema=ReadFiles.model_json_schema(),
-            name="ReadFiles",
-            description="""
-- Read full file content of one or more files.
-- Provide absolute file paths only
-""",
-        ),
-        ToolParam(
-            inputSchema=WriteIfEmpty.model_json_schema(),
-            name="WriteIfEmpty",
-            description="""
-- Write content to an empty or non-existent file. Provide file path and content. Use this instead of BashCommand for writing new files.
-- 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.\nAlso exits the docker environment.\nYou need to call GetScreenInfo again.",
-        ),
-        ToolParam(
-            inputSchema=FileEdit.model_json_schema(),
-            name="FileEdit",
-            description="""
-- Use absolute file path only.
-- Use SEARCH/REPLACE blocks to edit the file.
-- If the edit fails due to block not matching, please retry with correct block till it matches. Re-read the file to ensure you've all the lines correct.
-"""
-            + diffinstructions,
-        ),
+    tools_ = [
         ToolParam(
-            inputSchema=ContextSave.model_json_schema(),
-            name="ContextSave",
-            description="""
-Saves provided description and file contents of all the relevant file paths or globs in a single text file.
-- Provide random unqiue id or whatever user provided.
-- Leave project path as empty string if no project path""",
-        ),
+            inputSchema=tool.inputSchema,
+            name=tool.name,
+            description=tool.description,
+        )
+        for tool in TOOL_PROMPTS
     ]
-    if COMPUTER_USE_ON_DOCKER_ENABLED:
-        tools += [
-            ToolParam(
-                inputSchema=GetScreenInfo.model_json_schema(),
-                name="GetScreenInfo",
-                description="""
-- Important: call this first in the conversation before ScreenShot, Mouse, and Keyboard tools.
-- Get display information of a linux os running on docker using image "ghcr.io/anthropics/anthropic-quickstarts:computer-use-demo-latest"
-- If user hasn't provided docker image id, check using `docker ps` and provide the id.
-- If the docker is not running, run using `docker run -d -p 6080:6080 ghcr.io/anthropics/anthropic-quickstarts:computer-use-demo-latest`
-- Connects shell to the docker environment.
-- Note: once this is called, the shell enters the docker environment. All bash commands will run over there.
-""",
-            ),
-            ToolParam(
-                inputSchema=ScreenShot.model_json_schema(),
-                name="ScreenShot",
-                description="""
-- Capture screenshot of the linux os on docker.
-- All actions on UI using mouse and keyboard return within 0.5 seconds.
-    * So if you're doing something that takes longer for UI to update like heavy page loading, keep checking UI for update using ScreenShot upto 10 turns. 
-    * Notice for smallest of the loading icons to check if your action worked.
-    * After 10 turns of no change, ask user for permission to keep checking.
-    * If you don't notice even slightest of the change, it's likely you clicked on the wrong place.
-""",
-            ),
-            ToolParam(
-                inputSchema=Mouse.model_json_schema(),
-                name="Mouse",
-                description="""
-- Interact with the linux os on docker using mouse.
-- Uses xdotool
-- About left_click_drag: the current mouse position will be used as the starting point, click and drag to the given x, y coordinates. Useful in things like sliders, moving things around, etc.
-- The output of this command has the screenshot after doing this action. Use this to verify if the action was successful.
-""",
-            ),
-            ToolParam(
-                inputSchema=Keyboard.model_json_schema(),
-                name="Keyboard",
-                description="""
-- Interact with the linux os on docker using keyboard.
-- Emulate keyboard input to the screen
-- Uses xdootool to send keyboard input, keys like Return, BackSpace, Escape, Page_Up, etc. can be used.
-- Do not use it to interact with Bash tool.
-- Make sure you've selected a text area or an editable element before sending text.
-- The output of this command has the screenshot after doing this action. Use this to verify if the action was successful.
-""",
-            ),
-        ]
-
-    return tools
+    return tools_
 
 
 @server.call_tool()  # type: ignore
 async def handle_call_tool(
     name: str, arguments: dict[str, Any] | None
 ) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
+    global BASH_STATE
     if not arguments:
         raise ValueError("Missing arguments")
 
     tool_type = which_tool_name(name)
+    tool_call = parse_tool_by_name(name, arguments)
 
     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()})
-
-    try:
+        assert BASH_STATE
         output_or_dones, _ = get_tool_output(
-            tool_call, default_enc, 0.0, lambda x, y: ("", 0), 8000
+            Context(BASH_STATE, BASH_STATE.console),
+            tool_call,
+            default_enc,
+            0.0,
+            lambda x, y: ("", 0),
+            8000,
         )
 
     except Exception as e:
@@ -285,7 +127,6 @@ async def handle_call_tool(
 
     content: list[types.TextContent | types.ImageContent | types.EmbeddedResource] = []
     for output_or_done in output_or_dones:
-        assert not isinstance(output_or_done, DoneFlag)
         if isinstance(output_or_done, str):
             if issubclass(tool_type, Initialize):
                 output_or_done += """
@@ -307,31 +148,30 @@ Initialize call done.
     return content
 
 
-async def main(computer_use: bool) -> None:
-    global COMPUTER_USE_ON_DOCKER_ENABLED
+BASH_STATE = None
 
-    tools.TIMEOUT = SLEEP_TIME_MAX_S
-    tools.TIMEOUT_WHILE_OUTPUT = 55
-    tools.OUTPUT_WAIT_PATIENCE = 5
-    tools.console = Console()
-
-    if computer_use:
-        COMPUTER_USE_ON_DOCKER_ENABLED = True
 
+async def main() -> None:
+    global BASH_STATE
+    CONFIG.update(3, 55, 5)
     version = str(importlib.metadata.version("wcgw"))
-    tools.console.log("wcgw version: " + version)
-    # Run the server using stdin/stdout streams
-    async with mcp_wcgw.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={},
+    home_dir = os.path.expanduser("~")
+    with BashState(
+        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
+        async with mcp_wcgw.server.stdio.stdio_server() as (read_stream, write_stream):
+            await server.run(
+                read_stream,
+                write_stream,
+                InitializationOptions(
+                    server_name="wcgw",
+                    server_version=version,
+                    capabilities=server.get_capabilities(
+                        notification_options=NotificationOptions(),
+                        experimental_capabilities={},
+                    ),
                 ),
-            ),
-            raise_exceptions=False,
-        )
+                raise_exceptions=False,
+            )

wcgw/client/modes.py

@@ -61,7 +61,7 @@ You are now running in "code_writer" mode.
     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)}
+    - You are allowed to run FileEdit for files matching only the following globs: {", ".join(allowed_file_edit_globs)}
 """
         else:
             path_prompt = """
@@ -76,7 +76,7 @@ You are now running in "code_writer" mode.
     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)}
+    - You are allowed to run WriteIfEmpty files matching only the following globs: {", ".join(allowed_file_edit_globs)}
 """
         else:
             path_prompt = """
@@ -85,7 +85,7 @@ You are now running in "code_writer" mode.
     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 Ctrl-c 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.
@@ -101,7 +101,7 @@ You are now running in "code_writer" mode.
     if allowed_commands != "all":
         if allowed_commands:
             command_prompt = f"""
-    - You are only allowed to run the following commands: {', '.join(allowed_commands)}
+    - You are only allowed to run the following commands: {", ".join(allowed_commands)}
 {run_command_common}
 """
         else:
@@ -124,7 +124,7 @@ Instructions:
     - 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 Ctrl-c 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:
@@ -138,10 +138,11 @@ ARCHITECT_PROMPT = """You are now running 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.
+- Do not use Ctrl-c 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)
+- Share only snippets when any implementation is requested.
 
-Your response should be in self-critique and brainstorm style.
+Respond only after doing the following:
 - 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.)
@@ -149,17 +150,17 @@ Your response should be in self-critique and brainstorm style.
 
 
 DEFAULT_MODES: dict[Modes, ModeImpl] = {
-    Modes.wcgw: ModeImpl(
+    "wcgw": ModeImpl(
         bash_command_mode=BashCommandMode("normal_mode", "all"),
         write_if_empty_mode=WriteIfEmptyMode("all"),
         file_edit_mode=FileEditMode("all"),
     ),
-    Modes.architect: ModeImpl(
+    "architect": ModeImpl(
         bash_command_mode=BashCommandMode("restricted_mode", "all"),
         write_if_empty_mode=WriteIfEmptyMode([]),
         file_edit_mode=FileEditMode([]),
     ),
-    Modes.code_writer: ModeImpl(
+    "code_writer": ModeImpl(
         bash_command_mode=BashCommandMode("normal_mode", "all"),
         write_if_empty_mode=WriteIfEmptyMode("all"),
         file_edit_mode=FileEditMode("all"),
@@ -172,11 +173,11 @@ def modes_to_state(
 ) -> 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]
+        mode_impl = DEFAULT_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]
+        mode_impl = DEFAULT_MODES["code_writer"]
         # Override with custom settings from CodeWriterMode
         mode_impl = ModeImpl(
             bash_command_mode=BashCommandMode(
@@ -186,7 +187,7 @@ def modes_to_state(
             file_edit_mode=FileEditMode(mode.allowed_globs),
             write_if_empty_mode=WriteIfEmptyMode(mode.allowed_globs),
         )
-        mode_name = Modes.code_writer
+        mode_name = "code_writer"
     return (
         mode_impl.bash_command_mode,
         mode_impl.file_edit_mode,
@@ -233,10 +234,4 @@ Provide all relevant file paths in order to understand and solve the the task. E
 (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]
+KTS = {"wcgw": WCGW_KT, "architect": ARCHITECT_KT, "code_writer": WCGW_KT}