wcgw 5.0.2__py3-none-any.whl → 5.1.1__py3-none-any.whl

wcgw/client/bash_state/bash_state.py

@@ -996,7 +996,7 @@ def execute_bash(
     bash_state: BashState,
     enc: EncoderDecoder[int],
     bash_arg: BashCommand,
-    max_tokens: Optional[int],
+    max_tokens: Optional[int],  # This will be noncoding_max_tokens
     timeout_s: Optional[float],
 ) -> tuple[str, float]:
     try:
@@ -1026,7 +1026,7 @@ def _execute_bash(
     bash_state: BashState,
     enc: EncoderDecoder[int],
     bash_arg: BashCommand,
-    max_tokens: Optional[int],
+    max_tokens: Optional[int],  # This will be noncoding_max_tokens
     timeout_s: Optional[float],
 ) -> tuple[str, float]:
     try:

wcgw/client/file_ops/diff_edit.py

@@ -46,10 +46,13 @@ class FileEditOutput:
         last_idx = 0
         errors = []
         warnings = set[str]()
+        info = set[str]()
+        score = 0.0
         for (span, tolerances, replace_with), search_ in zip(
             self.edited_with_tolerances, self.orig_search_blocks
         ):
             for tol in tolerances:
+                score += tol.count * tol.score_multiplier
                 if tol.count > 0:
                     if tol.severity_cat == "WARNING":
                         warnings.add(tol.error_name)
@@ -66,6 +69,8 @@ Error:
 {tol.error_name}
 ---
                                   """)
+                    else:
+                        info.add(tol.error_name)
                     if len(errors) >= max_errors:
                         raise SearchReplaceMatchError("\n".join(errors))
             if last_idx < span.start:
@@ -80,12 +85,19 @@ Error:
         if errors:
             raise SearchReplaceMatchError("\n".join(errors))
 
+        if score > 1000:
+            display = (list(warnings) + list(info))[:max_errors]
+            raise SearchReplaceMatchError(
+                "Too many warnings generated, not apply the edits\n"
+                + "\n".join(display)
+            )
+
         return new_lines, set(warnings)
 
     @staticmethod
     def get_best_match(
         outputs: list["FileEditOutput"],
-    ) -> tuple[list["FileEditOutput"], bool]:
+    ) -> list["FileEditOutput"]:
         best_hits: list[FileEditOutput] = []
         best_score = float("-inf")
         assert outputs
@@ -103,7 +115,7 @@ Error:
                     best_score = hit_score
                 elif abs(hit_score - best_score) < 1e-3:
                     best_hits.append(output)
-        return best_hits, best_score > 1000
+        return best_hits
 
 
 def line_process_max_space_tolerance(line: str) -> str:

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

@@ -155,7 +155,7 @@ def edit_with_individual_fallback(
     original_lines: list[str], search_replace_blocks: list[tuple[list[str], list[str]]]
 ) -> tuple[list[str], set[str]]:
     outputs = FileEditInput(original_lines, 0, search_replace_blocks, 0).edit_file()
-    best_matches, is_error = FileEditOutput.get_best_match(outputs)
+    best_matches = FileEditOutput.get_best_match(outputs)
 
     try:
         edited_content, comments_ = best_matches[0].replace_or_throw(3)
@@ -171,7 +171,6 @@ def edit_with_individual_fallback(
                 all_comments |= comments_
             return running_lines, all_comments
         raise
-    assert not is_error
 
     if len(best_matches) > 1:
         # Find the first block that differs across matches

wcgw/client/mcp_server/server.py

@@ -1,13 +1,12 @@
 import importlib
 import logging
 import os
-from typing import Any
+from typing import Any, Optional
 
-import mcp_wcgw.server.stdio
-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
+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
@@ -25,7 +24,7 @@ from ..tools import (
     which_tool_name,
 )
 
-server = Server("wcgw")
+server: Server[Any] = Server("wcgw")
 
 # Log only time stamp
 logging.basicConfig(level=logging.INFO, format="%(asctime)s: %(message)s")
@@ -89,15 +88,7 @@ async def handle_list_tools() -> list[types.Tool]:
     Each tool specifies its arguments using JSON Schema validation.
     """
 
-    tools_ = [
-        ToolParam(
-            inputSchema=tool.inputSchema,
-            name=tool.name,
-            description=tool.description,
-        )
-        for tool in TOOL_PROMPTS
-    ]
-    return tools_
+    return TOOL_PROMPTS
 
 
 @server.call_tool()  # type: ignore
@@ -119,7 +110,8 @@ async def handle_call_tool(
             default_enc,
             0.0,
             lambda x, y: ("", 0),
-            8000,
+            24000,  # coding_max_tokens
+            8000,   # noncoding_max_tokens
         )
 
     except Exception as e:
@@ -165,7 +157,7 @@ async def main() -> 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):
+        async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
             await server.run(
                 read_stream,
                 write_stream,

wcgw/client/memory.py

@@ -64,7 +64,8 @@ T = TypeVar("T")
 
 def load_memory(
     task_id: str,
-    max_tokens: Optional[int],
+    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]]]:
@@ -75,6 +76,8 @@ def load_memory(
     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:

wcgw/client/repo_ops/display_tree.py

@@ -15,7 +15,7 @@ class DirectoryTree:
         self.root = root
         self.max_files = max_files
         self.expanded_files: Set[Path] = set()
-        self.expanded_dirs = set[Path]()
+        self.expanded_dirs: Set[Path] = set()
 
         if not self.root.exists():
             raise ValueError(f"Root path {root} does not exist")
@@ -77,11 +77,11 @@ class DirectoryTree:
         def _display_recursive(
             current_path: Path, indent: int = 0, depth: int = 0
         ) -> None:
-            # Print current directory name
+            # Print current directory name with a trailing slash for directories
             if current_path == self.root:
-                writer.write(f"{current_path}\n")
+                writer.write(f"{current_path}/\n")
             else:
-                writer.write(f"{' ' * indent}{current_path.name}\n")
+                writer.write(f"{' ' * indent}{current_path.name}/\n")
 
             # Don't recurse beyond depth 1 unless path contains expanded files
             if depth > 0 and current_path not in self.expanded_dirs:

wcgw/client/tool_prompts.py

@@ -1,6 +1,6 @@
 import os
-from dataclasses import dataclass
-from typing import Any
+
+from mcp.types import Tool, ToolAnnotations
 
 from ..types_ import (
     BashCommand,
@@ -15,15 +15,8 @@ with open(os.path.join(os.path.dirname(__file__), "diff-instructions.txt")) as f
     diffinstructions = f.read()
 
 
-@dataclass
-class Prompts:
-    inputSchema: dict[str, Any]
-    name: str
-    description: str
-
-
 TOOL_PROMPTS = [
-    Prompts(
+    Tool(
         inputSchema=Initialize.model_json_schema(),
         name="Initialize",
         description="""
@@ -38,8 +31,9 @@ TOOL_PROMPTS = [
 - Use type="reset_shell" if in a conversation shell is not working after multiple tries.
 - Use type="user_asked_change_workspace" if in a conversation user asked to change workspace
 """,
+        annotations=ToolAnnotations(readOnlyHint=True, openWorldHint=False),
     ),
-    Prompts(
+    Tool(
         inputSchema=BashCommand.model_json_schema(),
         name="BashCommand",
         description="""
@@ -54,8 +48,9 @@ TOOL_PROMPTS = [
 - 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.
 - Do not send Ctrl-c before checking for status till 10 minutes or whatever is appropriate for the program to finish.
 """,
+        annotations=ToolAnnotations(destructiveHint=True, openWorldHint=True),
     ),
-    Prompts(
+    Tool(
         inputSchema=ReadFiles.model_json_schema(),
         name="ReadFiles",
         description="""
@@ -65,13 +60,15 @@ TOOL_PROMPTS = [
     - You may populate "show_line_numbers_reason" with your reason, by default null/empty means no line numbers are shown.
     - You may extract a range of lines. E.g., `/path/to/file:1-10` for lines 1-10. You can drop start or end like `/path/to/file:1-` or `/path/to/file:-10` 
 """,
+        annotations=ToolAnnotations(readOnlyHint=True, openWorldHint=False),
     ),
-    Prompts(
+    Tool(
         inputSchema=ReadImage.model_json_schema(),
         name="ReadImage",
         description="Read an image from the shell.",
+        annotations=ToolAnnotations(readOnlyHint=True, openWorldHint=False),
     ),
-    Prompts(
+    Tool(
         inputSchema=FileWriteOrEdit.model_json_schema(),
         name="FileWriteOrEdit",
         description="""
@@ -85,13 +82,17 @@ TOOL_PROMPTS = [
 
 """
         + diffinstructions,
+        annotations=ToolAnnotations(
+            destructiveHint=True, idempotentHint=True, openWorldHint=False
+        ),
     ),
-    Prompts(
+    Tool(
         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 3 word unqiue id or whatever user provided.
 - Leave project path as empty string if no project path""",
+        annotations=ToolAnnotations(readOnlyHint=True, openWorldHint=False),
     ),
 ]