indent 0.0.8__py3-none-any.whl

exponent/core/remote_execution/tool_execution.py

@@ -0,0 +1,289 @@
+import uuid
+from collections.abc import Callable
+from pathlib import Path
+from time import time
+
+from anyio import Path as AsyncPath
+
+from exponent.core.remote_execution import files
+from exponent.core.remote_execution.cli_rpc_types import (
+    BashToolInput,
+    BashToolResult,
+    ErrorToolResult,
+    GlobToolInput,
+    GlobToolResult,
+    GrepToolInput,
+    GrepToolResult,
+    ListToolInput,
+    ListToolResult,
+    ReadToolInput,
+    ReadToolResult,
+    ToolInputType,
+    ToolResultType,
+    WriteToolInput,
+    WriteToolResult,
+)
+from exponent.core.remote_execution.code_execution import (
+    execute_code_streaming,
+)
+from exponent.core.remote_execution.file_write import execute_full_file_rewrite
+from exponent.core.remote_execution.truncation import truncate_tool_result
+from exponent.core.remote_execution.types import (
+    StreamingCodeExecutionRequest,
+    StreamingCodeExecutionResponse,
+)
+from exponent.core.remote_execution.utils import (
+    assert_unreachable,
+    safe_read_file,
+)
+
+GREP_MAX_RESULTS = 100
+
+
+async def execute_tool(
+    tool_input: ToolInputType, working_directory: str
+) -> ToolResultType:
+    if isinstance(tool_input, ReadToolInput):
+        return await execute_read_file(tool_input, working_directory)
+    elif isinstance(tool_input, WriteToolInput):
+        return await execute_write_file(tool_input, working_directory)
+    elif isinstance(tool_input, ListToolInput):
+        return await execute_list_files(tool_input, working_directory)
+    elif isinstance(tool_input, GlobToolInput):
+        return await execute_glob_files(tool_input, working_directory)
+    elif isinstance(tool_input, GrepToolInput):
+        return await execute_grep_files(tool_input, working_directory)
+    elif isinstance(tool_input, BashToolInput):
+        raise ValueError("Bash tool input should be handled by execute_bash_tool")
+    else:
+        assert_unreachable(tool_input)
+
+
+def truncate_result[T: ToolResultType](tool_result: T) -> T:
+    return truncate_tool_result(tool_result)
+
+
+async def execute_read_file(  # noqa: PLR0911
+    tool_input: ReadToolInput, working_directory: str
+) -> ReadToolResult | ErrorToolResult:
+    # Validate absolute path requirement
+    if not tool_input.file_path.startswith("/"):
+        return ErrorToolResult(
+            error_message=f"File path must be absolute, got relative path: {tool_input.file_path}"
+        )
+
+    # Validate offset and limit
+    offset = tool_input.offset if tool_input.offset is not None else 0
+    limit = tool_input.limit if tool_input.limit is not None else 2000
+
+    if offset < 0:
+        return ErrorToolResult(
+            error_message=f"Offset must be non-negative, got: {offset}"
+        )
+
+    if limit <= 0:
+        return ErrorToolResult(error_message=f"Limit must be positive, got: {limit}")
+
+    file = AsyncPath(working_directory, tool_input.file_path)
+
+    try:
+        exists = await file.exists()
+    except (OSError, PermissionError) as e:
+        return ErrorToolResult(error_message=f"Cannot access file: {e!s}")
+
+    if not exists:
+        return ErrorToolResult(
+            error_message="File not found",
+        )
+
+    try:
+        if await file.is_dir():
+            return ErrorToolResult(
+                error_message=f"{await file.absolute()} is a directory",
+            )
+    except (OSError, PermissionError) as e:
+        return ErrorToolResult(error_message=f"Cannot check file type: {e!s}")
+
+    try:
+        content = await safe_read_file(file)
+    except PermissionError:
+        return ErrorToolResult(
+            error_message=f"Permission denied: cannot read {tool_input.file_path}"
+        )
+    except UnicodeDecodeError:
+        return ErrorToolResult(
+            error_message="File appears to be binary or has invalid text encoding"
+        )
+    except Exception as e:  # noqa: BLE001
+        return ErrorToolResult(error_message=f"Error reading file: {e!s}")
+
+    # Handle empty files
+    if not content:
+        return ReadToolResult(
+            content="",
+            num_lines=0,
+            start_line=0,
+            total_lines=0,
+        )
+
+    content_lines = content.splitlines(keepends=True)
+    total_lines = len(content_lines)
+
+    # Handle offset beyond file length
+    if offset >= total_lines:
+        return ReadToolResult(
+            content="",
+            num_lines=0,
+            start_line=offset,
+            total_lines=total_lines,
+        )
+
+    # Apply offset and limit
+    content_lines = content_lines[offset : offset + limit]
+
+    # Apply character-level truncation at line boundaries to ensure consistency
+    # This ensures the content field and num_lines field remain in sync
+    CHARACTER_LIMIT = 90_000  # Match the limit in truncation.py
+
+    # Join lines and check total size
+    final_content = "".join(content_lines)
+
+    if len(final_content) > CHARACTER_LIMIT:
+        # Truncate at line boundaries to stay under the limit
+        truncated_lines: list[str] = []
+        current_size = 0
+        truncation_message = "\n[Content truncated due to size limit]"
+        truncation_size = len(truncation_message)
+        lines_included = 0
+
+        for line in content_lines:
+            # Check if adding this line would exceed the limit (accounting for truncation message)
+            if current_size + len(line) + truncation_size > CHARACTER_LIMIT:
+                final_content = "".join(truncated_lines) + truncation_message
+                break
+            truncated_lines.append(line)
+            current_size += len(line)
+            lines_included += 1
+        else:
+            # All lines fit (shouldn't happen if we got here, but be safe)
+            final_content = "".join(truncated_lines)
+            lines_included = len(content_lines)
+
+        num_lines = lines_included
+    else:
+        num_lines = len(content_lines)
+
+    return ReadToolResult(
+        content=final_content,
+        num_lines=num_lines,
+        start_line=offset,
+        total_lines=total_lines,
+    )
+
+
+async def execute_write_file(
+    tool_input: WriteToolInput, working_directory: str
+) -> WriteToolResult:
+    file_path = tool_input.file_path
+    path = Path(working_directory, file_path)
+    result = await execute_full_file_rewrite(
+        path, tool_input.content, working_directory
+    )
+    return WriteToolResult(message=result)
+
+
+async def execute_list_files(
+    tool_input: ListToolInput, working_directory: str
+) -> ListToolResult | ErrorToolResult:
+    path = AsyncPath(tool_input.path)
+
+    try:
+        exists = await path.exists()
+    except (OSError, PermissionError) as e:
+        return ErrorToolResult(error_message=f"Cannot access path: {e!s}")
+
+    if not exists:
+        return ErrorToolResult(error_message=f"Directory not found: {tool_input.path}")
+
+    try:
+        is_dir = await path.is_dir()
+    except (OSError, PermissionError) as e:
+        return ErrorToolResult(
+            error_message=f"Cannot check if path is directory: {e!s}"
+        )
+
+    if not is_dir:
+        return ErrorToolResult(
+            error_message=f"Path is not a directory: {tool_input.path}"
+        )
+
+    try:
+        filenames = [entry.name async for entry in path.iterdir()]
+    except (OSError, PermissionError) as e:
+        return ErrorToolResult(error_message=f"Cannot list directory contents: {e!s}")
+
+    return ListToolResult(
+        files=[filename for filename in filenames],
+    )
+
+
+async def execute_glob_files(
+    tool_input: GlobToolInput, working_directory: str
+) -> GlobToolResult:
+    # async timer
+    start_time = time()
+    results = await files.glob(
+        path=working_directory if tool_input.path is None else tool_input.path,
+        glob_pattern=tool_input.pattern,
+    )
+    duration_ms = int((time() - start_time) * 1000)
+    return GlobToolResult(
+        filenames=results,
+        duration_ms=duration_ms,
+        num_files=len(results),
+        truncated=len(results) >= files.GLOB_MAX_COUNT,
+    )
+
+
+async def execute_grep_files(
+    tool_input: GrepToolInput, working_directory: str
+) -> GrepToolResult:
+    results = await files.search_files(
+        path_str=working_directory if tool_input.path is None else tool_input.path,
+        file_pattern=tool_input.include,
+        regex=tool_input.pattern,
+        working_directory=working_directory,
+    )
+    return GrepToolResult(
+        matches=results[:GREP_MAX_RESULTS],
+        truncated=bool(len(results) > GREP_MAX_RESULTS),
+    )
+
+
+async def execute_bash_tool(
+    tool_input: BashToolInput, working_directory: str, should_halt: Callable[[], bool]
+) -> BashToolResult:
+    start_time = time()
+    result = None
+    async for result in execute_code_streaming(
+        StreamingCodeExecutionRequest(
+            language="shell",
+            content=tool_input.command,
+            timeout=120 if tool_input.timeout is None else tool_input.timeout,
+            correlation_id=str(uuid.uuid4()),
+        ),
+        working_directory=working_directory,
+        session=None,  # type: ignore
+        should_halt=should_halt,
+    ):
+        pass
+
+    assert isinstance(result, StreamingCodeExecutionResponse)
+
+    return BashToolResult(
+        shell_output=result.content,
+        exit_code=result.exit_code,
+        duration_ms=int((time() - start_time) * 1000),
+        timed_out=result.cancelled_for_timeout,
+        stopped_by_user=result.halted,
+    )

exponent/core/remote_execution/truncation.py

@@ -0,0 +1,284 @@
+"""Generalized truncation framework for tool results."""
+
+from abc import ABC, abstractmethod
+from typing import Any, TypeVar, cast
+
+from msgspec.structs import replace
+
+from exponent.core.remote_execution.cli_rpc_types import (
+    BashToolResult,
+    ErrorToolResult,
+    GlobToolResult,
+    GrepToolResult,
+    ListToolResult,
+    ReadToolResult,
+    ToolResult,
+    WriteToolResult,
+)
+from exponent.core.remote_execution.utils import truncate_output
+
+DEFAULT_CHARACTER_LIMIT = 90_000
+DEFAULT_LIST_ITEM_LIMIT = 1000
+DEFAULT_LIST_PREVIEW_ITEMS = 10
+
+
+class TruncationStrategy(ABC):
+    @abstractmethod
+    def should_truncate(self, result: ToolResult) -> bool:
+        pass
+
+    @abstractmethod
+    def truncate(self, result: ToolResult) -> ToolResult:
+        pass
+
+
+class StringFieldTruncation(TruncationStrategy):
+    def __init__(
+        self,
+        field_name: str,
+        character_limit: int = DEFAULT_CHARACTER_LIMIT,
+    ):
+        self.field_name = field_name
+        self.character_limit = character_limit
+
+    def should_truncate(self, result: ToolResult) -> bool:
+        if hasattr(result, self.field_name):
+            value = getattr(result, self.field_name)
+            if isinstance(value, str):
+                return len(value) > self.character_limit
+        return False
+
+    def truncate(self, result: ToolResult) -> ToolResult:
+        if not hasattr(result, self.field_name):
+            return result
+
+        value = getattr(result, self.field_name)
+        if not isinstance(value, str):
+            return result
+
+        truncated_value, was_truncated = truncate_output(value, self.character_limit)
+
+        updates: dict[str, Any] = {self.field_name: truncated_value}
+        if hasattr(result, "truncated") and was_truncated:
+            updates["truncated"] = True
+
+        return replace(result, **updates)
+
+
+class ListFieldTruncation(TruncationStrategy):
+    def __init__(
+        self,
+        field_name: str,
+        item_limit: int = DEFAULT_LIST_ITEM_LIMIT,
+        preview_items: int = DEFAULT_LIST_PREVIEW_ITEMS,
+    ):
+        self.field_name = field_name
+        self.item_limit = item_limit
+        self.preview_items = preview_items
+
+    def should_truncate(self, result: ToolResult) -> bool:
+        if hasattr(result, self.field_name):
+            value = getattr(result, self.field_name)
+            if isinstance(value, list):
+                return len(value) > self.item_limit
+        return False
+
+    def truncate(self, result: ToolResult) -> ToolResult:
+        if not hasattr(result, self.field_name):
+            return result
+
+        value = getattr(result, self.field_name)
+        if not isinstance(value, list):
+            return result
+
+        total_items = len(value)
+        if total_items <= self.item_limit:
+            return result
+
+        truncated_count = max(0, total_items - 2 * self.preview_items)
+        truncated_list = (
+            value[: self.preview_items]
+            + [f"... {truncated_count} items truncated ..."]
+            + value[-self.preview_items :]
+        )
+
+        updates: dict[str, Any] = {self.field_name: truncated_list}
+        if hasattr(result, "truncated"):
+            updates["truncated"] = True
+
+        return replace(result, **updates)
+
+
+class CompositeTruncation(TruncationStrategy):
+    def __init__(self, strategies: list[TruncationStrategy]):
+        self.strategies = strategies
+
+    def should_truncate(self, result: ToolResult) -> bool:
+        return any(strategy.should_truncate(result) for strategy in self.strategies)
+
+    def truncate(self, result: ToolResult) -> ToolResult:
+        for strategy in self.strategies:
+            if strategy.should_truncate(result):
+                result = strategy.truncate(result)
+        return result
+
+
+class StringListTruncation(TruncationStrategy):
+    """Truncation for lists of strings that limits both number of items and individual string length."""
+
+    def __init__(
+        self,
+        field_name: str,
+        max_items: int = DEFAULT_LIST_ITEM_LIMIT,
+        preview_items: int = DEFAULT_LIST_PREVIEW_ITEMS,
+        max_item_length: int = 1000,
+    ):
+        self.field_name = field_name
+        self.max_items = max_items
+        self.preview_items = preview_items
+        self.max_item_length = max_item_length
+
+    def should_truncate(self, result: ToolResult) -> bool:
+        if not hasattr(result, self.field_name):
+            return False
+
+        items = getattr(result, self.field_name)
+        if not isinstance(items, list):
+            return False
+
+        # Check if we need to truncate number of items
+        if len(items) > self.max_items:
+            return True
+
+        # Check if any individual item is too long
+        for item in items:
+            if isinstance(item, str) and len(item) > self.max_item_length:
+                return True
+            # Handle dict items (e.g., with metadata like file path and line number)
+            elif isinstance(item, dict) and "content" in item:
+                if len(item["content"]) > self.max_item_length:
+                    return True
+
+        return False
+
+    def _truncate_item_content(
+        self, item: str | dict[str, Any]
+    ) -> str | dict[str, Any]:
+        """Truncate an individual item's content."""
+        if isinstance(item, str):
+            if len(item) <= self.max_item_length:
+                return item
+            # Truncate string item
+            truncated, _ = truncate_output(item, self.max_item_length)
+            return truncated
+        elif isinstance(item, dict) and "content" in item:
+            # Handle dict-style items (e.g., with metadata like file path and line number)
+            if len(item["content"]) <= self.max_item_length:
+                return item
+            truncated_content, _ = truncate_output(
+                item["content"], self.max_item_length
+            )
+            return {**item, "content": truncated_content}
+        else:
+            return item
+
+    def truncate(self, result: ToolResult) -> ToolResult:
+        if not hasattr(result, self.field_name):
+            return result
+
+        items = getattr(result, self.field_name)
+        if not isinstance(items, list):
+            return result
+
+        # First, truncate individual item contents
+        truncated_items = [self._truncate_item_content(item) for item in items]
+
+        # Then, limit the number of items if needed
+        total_items = len(truncated_items)
+        if total_items > self.max_items:
+            truncated_count = max(0, total_items - 2 * self.preview_items)
+            final_items = (
+                truncated_items[: self.preview_items]
+                + [f"... {truncated_count} items truncated ..."]
+                + truncated_items[-self.preview_items :]
+            )
+        else:
+            final_items = truncated_items
+
+        updates: dict[str, Any] = {self.field_name: final_items}
+        if hasattr(result, "truncated"):
+            updates["truncated"] = True
+
+        return replace(result, **updates)
+
+
+TRUNCATION_REGISTRY: dict[type[ToolResult], TruncationStrategy] = {
+    ReadToolResult: StringFieldTruncation("content"),
+    WriteToolResult: StringFieldTruncation("message"),
+    BashToolResult: StringFieldTruncation("shell_output"),
+    GrepToolResult: StringListTruncation("matches"),
+    GlobToolResult: StringListTruncation("filenames", max_item_length=4096),
+    ListToolResult: StringListTruncation("files", max_item_length=4096),
+}
+
+
+T = TypeVar("T", bound=ToolResult)
+
+
+def truncate_tool_result(result: T) -> T:
+    if isinstance(result, ErrorToolResult):
+        return cast(T, result)
+
+    result_type = type(result)
+    if result_type in TRUNCATION_REGISTRY:
+        strategy = TRUNCATION_REGISTRY[result_type]
+        if strategy.should_truncate(result):
+            return cast(T, strategy.truncate(result))
+
+    return result
+
+
+def register_truncation_strategy(
+    result_type: type[ToolResult],
+    strategy: TruncationStrategy,
+) -> None:
+    TRUNCATION_REGISTRY[result_type] = strategy
+
+
+def configure_truncation_limits(
+    character_limit: int | None = None,
+    list_item_limit: int | None = None,
+    list_preview_items: int | None = None,
+    string_item_length: int | None = None,
+) -> None:
+    for strategy in TRUNCATION_REGISTRY.values():
+        if isinstance(strategy, StringFieldTruncation) and character_limit is not None:
+            strategy.character_limit = character_limit
+        elif isinstance(strategy, ListFieldTruncation):
+            if list_item_limit is not None:
+                strategy.item_limit = list_item_limit
+            if list_preview_items is not None:
+                strategy.preview_items = list_preview_items
+        elif isinstance(strategy, StringListTruncation):
+            if list_item_limit is not None:
+                strategy.max_items = list_item_limit
+            if list_preview_items is not None:
+                strategy.preview_items = list_preview_items
+            if string_item_length is not None:
+                strategy.max_item_length = string_item_length
+        elif isinstance(strategy, CompositeTruncation):
+            for sub_strategy in strategy.strategies:
+                if isinstance(sub_strategy, StringFieldTruncation) and character_limit:
+                    sub_strategy.character_limit = character_limit
+                elif isinstance(sub_strategy, ListFieldTruncation):
+                    if list_item_limit is not None:
+                        sub_strategy.item_limit = list_item_limit
+                    if list_preview_items is not None:
+                        sub_strategy.preview_items = list_preview_items
+                elif isinstance(sub_strategy, StringListTruncation):
+                    if list_item_limit is not None:
+                        sub_strategy.max_items = list_item_limit
+                    if list_preview_items is not None:
+                        sub_strategy.preview_items = list_preview_items
+                    if string_item_length is not None:
+                        sub_strategy.max_item_length = string_item_length