wcgw 5.5.4__py3-none-any.whl

wcgw/client/file_ops/extensions.py

@@ -0,0 +1,137 @@
+"""
+File with definitions of known source code file extensions.
+Used to determine the appropriate context length for files.
+Supports selecting between coding_max_tokens and noncoding_max_tokens
+based on file extensions.
+"""
+from typing import Dict, Optional, Set
+
+# Set of file extensions considered to be source code
+# Each extension should be listed without the dot (e.g., 'py' not '.py')
+SOURCE_CODE_EXTENSIONS: Set[str] = {
+    # Python
+    'py', 'pyx', 'pyi', 'pyw',
+    
+    # JavaScript and TypeScript
+    'js', 'jsx', 'ts', 'tsx', 'mjs', 'cjs',
+    
+    # Web
+    'html', 'htm', 'xhtml', 'css', 'scss', 'sass', 'less',
+    
+    # C and C++
+    'c', 'h', 'cpp', 'cxx', 'cc', 'hpp', 'hxx', 'hh', 'inl',
+    
+    # C#
+    'cs', 'csx',
+    
+    # Java
+    'java', 'scala', 'kt', 'kts', 'groovy',
+    
+    # Go
+    'go', 'mod',
+    
+    # Rust
+    'rs', 'rlib',
+    
+    # Swift
+    'swift',
+    
+    # Ruby
+    'rb', 'rake', 'gemspec',
+    
+    # PHP
+    'php', 'phtml', 'phar', 'phps',
+    
+    # Shell
+    'sh', 'bash', 'zsh', 'fish',
+    
+    # PowerShell
+    'ps1', 'psm1', 'psd1',
+    
+    # SQL
+    'sql', 'ddl', 'dml',
+    
+    # Markup and config
+    'xml', 'json', 'yaml', 'yml', 'toml', 'ini', 'cfg', 'conf',
+    
+    # Documentation
+    'md', 'markdown', 'rst', 'adoc', 'tex',
+    
+    # Build and dependency files
+    'Makefile', 'Dockerfile', 'Jenkinsfile',
+    
+    # Haskell
+    'hs', 'lhs',
+    
+    # Lisp family
+    'lisp', 'cl', 'el', 'clj', 'cljs', 'edn', 'scm',
+    
+    # Erlang and Elixir
+    'erl', 'hrl', 'ex', 'exs',
+    
+    # Dart and Flutter
+    'dart',
+    
+    # Objective-C
+    'm', 'mm',
+}
+
+# Context length limits based on file type (in tokens)
+CONTEXT_LENGTH_LIMITS: Dict[str, int] = {
+    'source_code': 24000,  # For known source code files
+    'default': 8000,       # For all other files
+}
+
+def is_source_code_file(filename: str) -> bool:
+    """
+    Determine if a file is a source code file based on its extension.
+    
+    Args:
+        filename: The name of the file to check
+        
+    Returns:
+        True if the file has a recognized source code extension, False otherwise
+    """
+    # Extract extension (without the dot)
+    parts = filename.split('.')
+    if len(parts) > 1:
+        ext = parts[-1].lower()
+        return ext in SOURCE_CODE_EXTENSIONS
+    
+    # Files without extensions (like 'Makefile', 'Dockerfile')
+    # Case-insensitive match for files without extensions
+    return filename.lower() in {ext.lower() for ext in SOURCE_CODE_EXTENSIONS}
+
+def get_context_length_for_file(filename: str) -> int:
+    """
+    Get the appropriate context length limit for a file based on its extension.
+    
+    Args:
+        filename: The name of the file to check
+        
+    Returns:
+        The context length limit in tokens
+    """
+    if is_source_code_file(filename):
+        return CONTEXT_LENGTH_LIMITS['source_code']
+    return CONTEXT_LENGTH_LIMITS['default']
+
+
+def select_max_tokens(filename: str, coding_max_tokens: Optional[int], noncoding_max_tokens: Optional[int]) -> Optional[int]:
+    """
+    Select the appropriate max_tokens limit based on file type.
+    
+    Args:
+        filename: The name of the file to check
+        coding_max_tokens: Maximum tokens for source code files
+        noncoding_max_tokens: Maximum tokens for non-source code files
+        
+    Returns:
+        The appropriate max_tokens limit for the file
+    """
+    if coding_max_tokens is None and noncoding_max_tokens is None:
+        return None
+        
+    if is_source_code_file(filename):
+        return coding_max_tokens
+    return noncoding_max_tokens

wcgw/client/file_ops/search_replace.py

@@ -0,0 +1,212 @@
+import re
+from typing import Callable, Optional
+
+from .diff_edit import FileEditInput, FileEditOutput, SearchReplaceMatchError
+
+# Global regex patterns
+SEARCH_MARKER = re.compile(r"^<<<<<<+\s*SEARCH>?\s*$")
+DIVIDER_MARKER = re.compile(r"^======*\s*$")
+REPLACE_MARKER = re.compile(r"^>>>>>>+\s*REPLACE\s*$")
+
+
+class SearchReplaceSyntaxError(Exception):
+    def __init__(self, message: str):
+        message = f"""Got syntax error while parsing search replace blocks:
+{message}
+---
+
+Make sure blocks are in correct sequence, and the markers are in separate lines:
+
+<{"<<<<<< SEARCH"}
+    example old
+=======
+    example new
+>{">>>>>> REPLACE"}
+ 
+"""
+        super().__init__(message)
+
+
+def search_replace_edit(
+    lines: list[str], original_content: str, logger: Callable[[str], object]
+) -> tuple[str, str]:
+    if not lines:
+        raise SearchReplaceSyntaxError("Error: No input to search replace edit")
+
+    original_lines = original_content.split("\n")
+    n_lines = len(lines)
+    i = 0
+    search_replace_blocks = list[tuple[list[str], list[str]]]()
+
+    while i < n_lines:
+        if SEARCH_MARKER.match(lines[i]):
+            line_num = i + 1
+            search_block = []
+            i += 1
+
+            while i < n_lines and not DIVIDER_MARKER.match(lines[i]):
+                if SEARCH_MARKER.match(lines[i]) or REPLACE_MARKER.match(lines[i]):
+                    raise SearchReplaceSyntaxError(
+                        f"Line {i + 1}: Found stray marker in SEARCH block: {lines[i]}"
+                    )
+                search_block.append(lines[i])
+                i += 1
+
+            if i >= n_lines:
+                raise SearchReplaceSyntaxError(
+                    f"Line {line_num}: Unclosed SEARCH block - missing ======= marker"
+                )
+
+            if not search_block:
+                raise SearchReplaceSyntaxError(
+                    f"Line {line_num}: SEARCH block cannot be empty"
+                )
+
+            i += 1
+            replace_block = []
+
+            while i < n_lines and not REPLACE_MARKER.match(lines[i]):
+                if SEARCH_MARKER.match(lines[i]) or DIVIDER_MARKER.match(lines[i]):
+                    raise SearchReplaceSyntaxError(
+                        f"Line {i + 1}: Found stray marker in REPLACE block: {lines[i]}"
+                    )
+                replace_block.append(lines[i])
+                i += 1
+
+            if i >= n_lines:
+                raise SearchReplaceSyntaxError(
+                    f"Line {line_num}: Unclosed block - missing REPLACE marker"
+                )
+
+            i += 1
+
+            for line in search_block:
+                logger("> " + line)
+            logger("=======")
+            for line in replace_block:
+                logger("< " + line)
+            logger("\n\n\n\n")
+
+            search_replace_blocks.append((search_block, replace_block))
+        else:
+            if REPLACE_MARKER.match(lines[i]) or DIVIDER_MARKER.match(lines[i]):
+                raise SearchReplaceSyntaxError(
+                    f"Line {i + 1}: Found stray marker outside block: {lines[i]}"
+                )
+            i += 1
+
+    if not search_replace_blocks:
+        raise SearchReplaceSyntaxError(
+            "No valid search replace blocks found, ensure your SEARCH/REPLACE blocks are formatted correctly"
+        )
+
+    edited_content, comments_ = edit_with_individual_fallback(
+        original_lines, search_replace_blocks, False
+    )
+
+    edited_file = "\n".join(edited_content)
+    if not comments_:
+        comments = "File edited successfully."
+    else:
+        comments = (
+            "File edited successfully. However, following warnings were generated while matching search blocks.\n"
+            + "\n".join(comments_)
+        )
+    return edited_file, comments
+
+
+def identify_first_differing_block(
+    best_matches: list[FileEditOutput],
+) -> Optional[list[str]]:
+    """
+    Identify the first search block that differs across multiple best matches.
+    Returns the search block content that first shows different matches.
+    """
+    if not best_matches or len(best_matches) <= 1:
+        return None
+
+    # First, check if the number of blocks differs (shouldn't happen, but let's be safe)
+    block_counts = [len(match.edited_with_tolerances) for match in best_matches]
+    if not all(count == block_counts[0] for count in block_counts):
+        # If block counts differ, just return the first search block as problematic
+        return (
+            best_matches[0].orig_search_blocks[0]
+            if best_matches[0].orig_search_blocks
+            else None
+        )
+
+    # Go through each block position and see if the slices differ
+    for i in range(min(block_counts)):
+        slices = [match.edited_with_tolerances[i][0] for match in best_matches]
+
+        # Check if we have different slices for this block across matches
+        if any(s.start != slices[0].start or s.stop != slices[0].stop for s in slices):
+            # We found our differing block - return the search block content
+            if i < len(best_matches[0].orig_search_blocks):
+                return best_matches[0].orig_search_blocks[i]
+            else:
+                return None
+
+    # If we get here, we couldn't identify a specific differing block
+    return None
+
+
+def edit_with_individual_fallback(
+    original_lines: list[str],
+    search_replace_blocks: list[tuple[list[str], list[str]]],
+    replace_all: bool,
+) -> tuple[list[str], set[str]]:
+    outputs = FileEditInput(original_lines, 0, search_replace_blocks, 0).edit_file()
+    best_matches = FileEditOutput.get_best_match(outputs)
+    try:
+        edited_content, comments_ = best_matches[0].replace_or_throw(3)
+    except SearchReplaceMatchError:
+        if len(search_replace_blocks) > 1:
+            try:
+                # Try one at a time
+                all_comments = set[str]()
+                running_lines = list(original_lines)
+                for block in search_replace_blocks:
+                    running_lines, comments_ = edit_with_individual_fallback(
+                        running_lines, [block], replace_all
+                    )
+                    all_comments |= comments_
+                return running_lines, all_comments
+            except SearchReplaceMatchError:
+                # Raise the outer error instead
+                # Otherwise the suggested search block will be
+                # after applying previous N search blocks and that
+                # would signal to LLM that we've updated the file
+                pass
+        raise
+
+    if replace_all and len(best_matches) > 1 and len(search_replace_blocks) == 1:
+        # For only one search/replace block only replace all
+        try:
+            edited_content, comments__ = edit_with_individual_fallback(
+                edited_content, search_replace_blocks, replace_all
+            )
+            comments_ |= comments__
+        except SearchReplaceMatchError:
+            # Will not happen ideally, but still no use of throwing error here
+            pass
+    elif len(best_matches) > 1:
+        # Find the first block that differs across matches
+        first_diff_block = identify_first_differing_block(best_matches)
+        if first_diff_block is not None:
+            block_content = "\n".join(first_diff_block)
+            raise SearchReplaceMatchError(f"""
+The following block matched more than once:
+```
+{block_content}
+```
+Consider adding more context before and after this block to make the match unique.
+    """)
+        else:
+            raise SearchReplaceMatchError("""
+One of the blocks matched more than once
+
+Consider adding more context before and after all the blocks to make the match unique.
+    """)
+
+    return edited_content, comments_

wcgw/client/mcp_server/Readme.md

@@ -0,0 +1,3 @@
+# The doc has moved to main Readme.md
+
+![main readme](https://github.com/rusiaaman/wcgw/blob/main/README.md)

wcgw/client/mcp_server/__init__.py

@@ -0,0 +1,32 @@
+# mypy: disable-error-code="import-untyped"
+import asyncio
+import importlib
+
+import typer
+from typer import Typer
+
+from wcgw.client.mcp_server import server
+
+main = Typer()
+
+
+@main.command()
+def app(
+    version: bool = typer.Option(
+        False, "--version", "-v", help="Show version and exit"
+    ),
+    shell: str = typer.Option(
+        "", "--shell", help="Path to shell executable (defaults to $SHELL or /bin/bash)"
+    ),
+) -> None:
+    """Main entry point for the package."""
+    if version:
+        version_ = importlib.metadata.version("wcgw")
+        print(f"wcgw version: {version_}")
+        raise typer.Exit()
+
+    asyncio.run(server.main(shell))
+
+
+# Optionally expose other important items at package level
+__all__ = ["main", "server"]

wcgw/client/mcp_server/server.py

@@ -0,0 +1,184 @@
+import importlib
+import logging
+import os
+from typing import Any, Optional
+
+import mcp.server.stdio
+import mcp.types as types
+from mcp.server import NotificationOptions, Server
+from mcp.server.models import InitializationOptions
+from pydantic import AnyUrl
+
+from wcgw.client.modes import KTS
+from wcgw.client.tool_prompts import TOOL_PROMPTS
+
+from ...types_ import (
+    Initialize,
+)
+from ..bash_state.bash_state import CONFIG, BashState, get_tmpdir
+from ..tools import (
+    Context,
+    default_enc,
+    get_tool_output,
+    parse_tool_by_name,
+    which_tool_name,
+)
+
+server: Server[Any] = Server("wcgw")
+
+# Log only time stamp
+logging.basicConfig(level=logging.INFO, format="%(asctime)s: %(message)s")
+logger = logging.getLogger("wcgw")
+
+
+class Console:
+    def print(self, msg: str, *args: Any, **kwargs: Any) -> None:
+        logger.info(msg)
+
+    def log(self, msg: str, *args: Any, **kwargs: Any) -> None:
+        logger.info(msg)
+
+
+@server.list_resources()  # type: ignore
+async def handle_list_resources() -> list[types.Resource]:
+    return []
+
+
+@server.read_resource()  # type: ignore
+async def handle_read_resource(uri: AnyUrl) -> str:
+    raise ValueError("No resources available")
+
+
+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.",
+        ),
+        KTS,
+    )
+}
+
+
+@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:
+    assert BASH_STATE
+    messages = [
+        types.PromptMessage(
+            role="user",
+            content=types.TextContent(
+                type="text", text=PROMPTS[name][1][BASH_STATE.mode]
+            ),
+        )
+    ]
+    return types.GetPromptResult(messages=messages)
+
+
+@server.list_tools()  # type: ignore
+async def handle_list_tools() -> list[types.Tool]:
+    """
+    List available tools.
+    Each tool specifies its arguments using JSON Schema validation.
+    """
+
+    return TOOL_PROMPTS
+
+
+@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:
+        assert BASH_STATE
+        output_or_dones, _ = get_tool_output(
+            Context(BASH_STATE, BASH_STATE.console),
+            tool_call,
+            default_enc,
+            0.0,
+            lambda x, y: ("", 0),
+            24000,  # coding_max_tokens
+            8000,   # noncoding_max_tokens
+        )
+
+    except Exception as e:
+        output_or_dones = [f"GOT EXCEPTION while calling tool. Error: {e}"]
+
+    content: list[types.TextContent | types.ImageContent | types.EmbeddedResource] = []
+    for output_or_done in output_or_dones:
+        if isinstance(output_or_done, str):
+            if issubclass(tool_type, Initialize):
+                # Prepare the original hardcoded message
+                original_message = """
+- 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.
+    """
+                
+                # If custom instructions exist, prepend them to the original message
+                if CUSTOM_INSTRUCTIONS:
+                    output_or_done += f"\n{CUSTOM_INSTRUCTIONS}\n{original_message}"
+                else:
+                    output_or_done += original_message
+
+            content.append(types.TextContent(type="text", text=output_or_done))
+        else:
+            content.append(
+                types.ImageContent(
+                    type="image",
+                    data=output_or_done.data,
+                    mimeType=output_or_done.media_type,
+                )
+            )
+
+    return content
+
+
+BASH_STATE = None
+CUSTOM_INSTRUCTIONS = None
+
+
+async def main(shell_path: str = "") -> None:
+    global BASH_STATE, CUSTOM_INSTRUCTIONS
+    CONFIG.update(3, 55, 5)
+    version = str(importlib.metadata.version("wcgw"))
+    
+    # Read custom instructions from environment variable
+    CUSTOM_INSTRUCTIONS = os.getenv("WCGW_SERVER_INSTRUCTIONS")
+
+    # starting_dir is inside tmp dir
+    tmp_dir = get_tmpdir()
+    starting_dir = os.path.join(tmp_dir, "claude_playground")
+
+    with BashState(
+        Console(), starting_dir, None, None, None, None, True, None, None, shell_path or None
+    ) as BASH_STATE:
+        BASH_STATE.console.log("wcgw version: " + version)
+        # Run the server using stdin/stdout streams
+        async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
+            await server.run(
+                read_stream,
+                write_stream,
+                InitializationOptions(
+                    server_name="wcgw",
+                    server_version=version,
+                    capabilities=server.get_capabilities(
+                        notification_options=NotificationOptions(),
+                        experimental_capabilities={},
+                    ),
+                ),
+                raise_exceptions=False,
+            )

wcgw/client/memory.py

@@ -0,0 +1,103 @@
+import json
+import os
+import re
+import shlex
+from typing import Any, Callable, Optional, TypeVar
+
+from ..types_ import ContextSave
+
+
+def get_app_dir_xdg() -> str:
+    xdg_data_dir = os.environ.get("XDG_DATA_HOME", os.path.expanduser("~/.local/share"))
+    return os.path.join(xdg_data_dir, "wcgw")
+
+
+def format_memory(task_memory: ContextSave, relevant_files: str) -> str:
+    memory_data = ""
+    if task_memory.project_root_path:
+        memory_data += (
+            f"# PROJECT ROOT = {shlex.quote(task_memory.project_root_path)}\n"
+        )
+    memory_data += task_memory.description
+
+    memory_data += (
+        "\n\n"
+        + "# Relevant file paths\n"
+        + ", ".join(map(shlex.quote, task_memory.relevant_file_globs))
+    )
+
+    memory_data += "\n\n# Relevant Files:\n" + relevant_files
+
+    return memory_data
+
+
+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)
+
+    task_id = task_memory.id
+    if not task_id:
+        raise Exception("Task id can not be empty")
+    memory_data = format_memory(task_memory, relevant_files)
+
+    memory_file_full = os.path.join(memory_dir, f"{task_id}.txt")
+
+    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
+
+
+T = TypeVar("T")
+
+
+def load_memory(
+    task_id: str,
+    coding_max_tokens: Optional[int],
+    noncoding_max_tokens: Optional[int],
+    encoder: Callable[[str], list[T]],
+    decoder: Callable[[list[T]], 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")
+
+    with open(memory_file, "r") as f:
+        data = f.read()
+
+    # Memory files are considered non-code files for token limits
+    max_tokens = noncoding_max_tokens
+    if max_tokens:
+        toks = encoder(data)
+        if len(toks) > max_tokens:
+            toks = toks[: max(0, max_tokens - 10)]
+            data = decoder(toks)
+            data += "\n(... truncated)"
+
+    project_root_match = re.search(r"# PROJECT ROOT = \s*(.*?)\s*$", data, re.MULTILINE)
+    project_root_path = ""
+    if project_root_match:
+        matched_path = project_root_match.group(1)
+        parsed_ = shlex.split(matched_path)
+        if parsed_ and len(parsed_) == 1:
+            project_root_path = parsed_[0]
+
+    # 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