indent 0.1.10__py3-none-any.whl → 0.1.12__py3-none-any.whl

exponent/core/remote_execution/client.py

@@ -3,19 +3,22 @@ from __future__ import annotations
 import asyncio
 import json
 import logging
-from collections.abc import AsyncGenerator, Callable
+import time
+import uuid
+from collections.abc import AsyncGenerator, Callable, Generator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
-from typing import Any, TypeVar, Union, cast
+from typing import Any, TypeVar, cast
 
 import msgspec
-import websockets.client
 import websockets.exceptions
 from httpx import (
     AsyncClient,
     codes as http_status,
 )
 from pydantic import BaseModel
+from websockets.asyncio import client as asyncio_websockets_client
+from websockets.asyncio.client import ClientConnection, connect
 
 from exponent.commands.utils import ConnectionTracker
 from exponent.core.config import is_editable_install
@@ -29,15 +32,16 @@ from exponent.core.remote_execution.cli_rpc_types import (
     ErrorResponse,
     GetAllFilesRequest,
     GetAllFilesResponse,
-    HttpResponse,
     HttpRequest,
+    KeepAliveCliChatRequest,
+    KeepAliveCliChatResponse,
+    SwitchCLIChatRequest,
+    SwitchCLIChatResponse,
     TerminateRequest,
     TerminateResponse,
     ToolExecutionRequest,
     ToolExecutionResponse,
     ToolResultType,
-    SwitchCLIChatRequest,
-    SwitchCLIChatResponse,
 )
 from exponent.core.remote_execution.code_execution import (
     execute_code_streaming,
@@ -87,7 +91,10 @@ class SwitchCLIChat:
     new_chat_uuid: str
 
 
-REMOTE_EXECUTION_CLIENT_EXIT_INFO = Union[WSDisconnected, SwitchCLIChat]
+REMOTE_EXECUTION_CLIENT_EXIT_INFO = WSDisconnected | SwitchCLIChat
+
+# UUID for a single run of the CLI
+cli_uuid = uuid.uuid4()
 
 
 class RemoteExecutionClient:
@@ -104,6 +111,9 @@ class RemoteExecutionClient:
         self._halt_states: dict[str, bool] = {}
         self._halt_lock = asyncio.Lock()
 
+        # Track last request time for timeout functionality
+        self._last_request_time: float | None = None
+
     @property
     def working_directory(self) -> str:
         return self.current_session.working_directory
@@ -138,10 +148,35 @@ class RemoteExecutionClient:
 
         return should_halt
 
+    async def _timeout_monitor(
+        self, timeout_seconds: int | None
+    ) -> WSDisconnected | None:
+        """Monitor for inactivity timeout and return WSDisconnected if timeout occurs.
+
+        If timeout_seconds is None, keeps looping indefinitely until cancelled.
+        """
+        try:
+            while True:
+                await asyncio.sleep(1)
+                if (
+                    timeout_seconds is not None
+                    and self._last_request_time is not None
+                    and time.time() - self._last_request_time > timeout_seconds
+                ):
+                    logger.info(
+                        f"No requests received for {timeout_seconds} seconds. Shutting down..."
+                    )
+                    return WSDisconnected(
+                        error_message=f"Timeout after {timeout_seconds} seconds of inactivity"
+                    )
+        except asyncio.CancelledError:
+            # Handle cancellation gracefully
+            return None
+
     async def _handle_websocket_message(
         self,
         msg: str,
-        websocket: websockets.client.WebSocketClientProtocol,
+        websocket: ClientConnection,
         requests: asyncio.Queue[CliRpcRequest],
     ) -> REMOTE_EXECUTION_CLIENT_EXIT_INFO | None:
         """Handle an incoming websocket message.
@@ -241,6 +276,21 @@ class RemoteExecutionClient:
                 )
             )
             return SwitchCLIChat(new_chat_uuid=request.request.new_chat_uuid)
+        elif isinstance(request.request, KeepAliveCliChatRequest):
+            await websocket.send(
+                json.dumps(
+                    {
+                        "type": "result",
+                        "data": msgspec.to_builtins(
+                            CliRpcResponse(
+                                request_id=request.request_id,
+                                response=KeepAliveCliChatResponse(),
+                            )
+                        ),
+                    }
+                )
+            )
+            return None
         else:
             if isinstance(request.request, ToolExecutionRequest) and isinstance(
                 request.request.tool_input, BashToolInput
@@ -302,7 +352,7 @@ class RemoteExecutionClient:
                     async with results_lock:
                         logger.info(f"Putting response {response}")
                         await results.put(response)
-                except Exception as e:  # noqa: BLE001
+                except Exception as e:
                     logger.info(f"Error handling request {request}:\n\n{e}")
                     async with results_lock:
                         await results.put(
@@ -327,7 +377,7 @@ class RemoteExecutionClient:
 
     async def _process_websocket_messages(
         self,
-        websocket: websockets.client.WebSocketClientProtocol,
+        websocket: ClientConnection,
         beats: asyncio.Queue[HeartbeatInfo],
         requests: asyncio.Queue[CliRpcRequest],
         results: asyncio.Queue[CliRpcResponse],
@@ -380,7 +430,7 @@ class RemoteExecutionClient:
 
     async def _handle_websocket_connection(
         self,
-        websocket: websockets.client.WebSocketClientProtocol,
+        websocket: ClientConnection,
         connection_tracker: ConnectionTracker | None,
     ) -> REMOTE_EXECUTION_CLIENT_EXIT_INFO | None:
         """Handle a single websocket connection.
@@ -427,17 +477,37 @@ class RemoteExecutionClient:
         self,
         chat_uuid: str,
         connection_tracker: ConnectionTracker | None = None,
+        timeout_seconds: int | None = None,
     ) -> REMOTE_EXECUTION_CLIENT_EXIT_INFO:
-        """Run the websocket connection loop."""
+        """Run the websocket connection loop with optional inactivity timeout."""
         self.current_session.set_chat_uuid(chat_uuid)
 
+        # Initialize last request time for timeout monitoring
+        self._last_request_time = time.time()
+
         async for websocket in self.ws_connect(f"/api/ws/chat/{chat_uuid}"):
-            result = await self._handle_websocket_connection(
-                websocket, connection_tracker
+            # Always run connection and timeout monitor concurrently
+            # If timeout_seconds is None, timeout monitor will loop indefinitely
+            done, pending = await asyncio.wait(
+                [
+                    asyncio.create_task(
+                        self._handle_websocket_connection(websocket, connection_tracker)
+                    ),
+                    asyncio.create_task(self._timeout_monitor(timeout_seconds)),
+                ],
+                return_when=asyncio.FIRST_COMPLETED,
             )
-            if result is not None:
-                return result
-            # If we get None, we'll try to reconnect
+
+            # Cancel pending tasks
+            for task in pending:
+                task.cancel()
+
+            # Return result from completed task
+            for task in done:
+                result = await task
+                # If we get None, we'll try to reconnect
+                if result is not None:
+                    return result
 
         # If we exit the websocket connection loop without returning,
         # it means we couldn't establish a connection
@@ -490,6 +560,7 @@ class RemoteExecutionClient:
             system_info=await system_context.get_system_info(self.working_directory),
             exponent_version=get_installed_version(),
             editable_installation=is_editable_install(),
+            cli_uuid=str(cli_uuid),
         )
 
     async def send_heartbeat(self, chat_uuid: str) -> CLIConnectedState:
@@ -509,6 +580,9 @@ class RemoteExecutionClient:
         return connected_state
 
     async def handle_request(self, request: CliRpcRequest) -> CliRpcResponse:
+        # Update last request time for timeout functionality
+        self._last_request_time = time.time()
+
         try:
             if isinstance(request.request, ToolExecutionRequest):
                 if isinstance(request.request.tool_input, BashToolInput):
@@ -550,7 +624,7 @@ class RemoteExecutionClient:
                             )
                         tool_result = truncate_result(raw_result)
                         results.append(tool_result)
-                    except Exception as e:  # noqa: BLE001
+                    except Exception as e:
                         logger.error(f"Error executing tool {tool_input}: {e}")
                         from exponent.core.remote_execution.cli_rpc_types import (
                             ErrorToolResult,
@@ -579,6 +653,10 @@ class RemoteExecutionClient:
                 raise ValueError(
                     "SwitchCLIChatRequest should not be handled by handle_request"
                 )
+            elif isinstance(request.request, KeepAliveCliChatRequest):
+                raise ValueError(
+                    "KeepAliveCliChatRequest should not be handled by handle_request"
+                )
 
             raise ValueError(f"Unhandled request type: {type(request)}")
 
@@ -612,7 +690,7 @@ class RemoteExecutionClient:
         ):
             yield output
 
-    def ws_connect(self, path: str) -> websockets.client.connect:
+    def ws_connect(self, path: str) -> connect:
         base_url = (
             str(self.ws_client.base_url)
             .replace("http://", "ws://")
@@ -622,14 +700,24 @@ class RemoteExecutionClient:
         url = f"{base_url}{path}"
         headers = {"api-key": self.api_client.headers["api-key"]}
 
-        conn = websockets.client.connect(
-            url, extra_headers=headers, timeout=10, ping_timeout=10
-        )
+        def custom_backoff() -> Generator[float, None, None]:
+            yield 0.1  # short initial delay
 
-        # Stop exponential backoff from blowing up
-        # the wait time between connection attempts
-        conn.BACKOFF_MAX = 2
-        conn.BACKOFF_INITIAL = 1  # pyright: ignore
+            delay = 0.5
+            while True:
+                if delay < 2.0:
+                    yield delay
+                    delay *= 1.5
+                else:
+                    yield 2.0
+
+        # Can remove if this is added to public API
+        # https://github.com/python-websockets/websockets/issues/1395#issuecomment-3225670409
+        asyncio_websockets_client.backoff = custom_backoff  # type: ignore[attr-defined, assignment]
+
+        conn = connect(
+            url, additional_headers=headers, open_timeout=10, ping_timeout=10
+        )
 
         return conn
 

exponent/core/remote_execution/file_write.py

@@ -1,85 +1,13 @@
-import logging
 import os
-import re
-import subprocess
-from collections.abc import Callable
-from textwrap import dedent, indent
 
 from anyio import Path as AsyncPath
-from diff_match_patch import diff_match_patch
-from pydantic import BaseModel
 
 from exponent.core.remote_execution.types import (
     FilePath,
-    FileWriteRequest,
-    FileWriteResponse,
 )
 from exponent.core.remote_execution.utils import (
-    assert_unreachable,
-    safe_read_file,
     safe_write_file,
 )
-from exponent.core.types.command_data import (
-    WRITE_STRATEGY_FULL_FILE_REWRITE,
-    WRITE_STRATEGY_NATURAL_EDIT,
-    WRITE_STRATEGY_SEARCH_REPLACE,
-    WRITE_STRATEGY_UDIFF,
-)
-
-logger = logging.getLogger(__name__)
-
-
-class FileEditResult(BaseModel):
-    content: str | None
-    failed_edits: list[tuple[str, str]]
-
-
-async def execute_file_write(
-    event: FileWriteRequest, working_directory: str
-) -> FileWriteResponse:
-    write_strategy = event.write_strategy
-    content = event.content
-
-    if write_strategy == WRITE_STRATEGY_FULL_FILE_REWRITE:
-        result = await execute_full_file_rewrite(
-            event.file_path, content, working_directory
-        )
-    elif write_strategy == WRITE_STRATEGY_UDIFF:
-        result = await execute_udiff_edit(event.file_path, content, working_directory)
-    elif write_strategy == WRITE_STRATEGY_SEARCH_REPLACE:
-        result = await execute_search_replace_edit(
-            event.file_path, content, working_directory
-        )
-    elif write_strategy == WRITE_STRATEGY_NATURAL_EDIT:
-        result = await execute_full_file_rewrite(
-            event.file_path, content, working_directory
-        )
-    else:
-        assert_unreachable(write_strategy)
-    return FileWriteResponse(
-        content=result,
-        correlation_id=event.correlation_id,
-    )
-
-
-def lint_file(file_path: str, working_directory: str) -> str:
-    try:
-        # Construct the absolute path
-        full_file_path = os.path.join(working_directory, file_path)
-
-        # Run ruff check --fix on the file
-        result = subprocess.run(
-            ["ruff", "check", "--fix", full_file_path],
-            capture_output=True,
-            text=True,
-            check=True,
-        )
-
-        # If the subprocess ran successfully, return a success message
-        return f"Lint results:\n\n{result.stdout}\n\n{result.stderr}"
-    except Exception as e:  # noqa: BLE001
-        # For any other errors, return a generic error message
-        return f"An error occurred while linting: {e!s}"
 
 
 async def execute_full_file_rewrite(
@@ -103,308 +31,5 @@ async def execute_full_file_rewrite(
 
         return result
 
-    except Exception as e:  # noqa: BLE001
-        return f"An error occurred: {e!s}"
-
-
-async def execute_udiff_edit(
-    file_path: str, content: str, working_directory: str
-) -> str:
-    return await execute_partial_edit(
-        file_path, content, working_directory, apply_udiff
-    )
-
-
-async def execute_search_replace_edit(
-    file_path: str, content: str, working_directory: str
-) -> str:
-    return await execute_partial_edit(
-        file_path, content, working_directory, apply_all_search_replace
-    )
-
-
-async def execute_partial_edit(
-    file_path: str,
-    edit_content: str,
-    working_directory: str,
-    edit_function: Callable[[str, str], FileEditResult],
-) -> str:
-    try:
-        # Construct the absolute path
-        full_file_path = AsyncPath(os.path.join(working_directory, file_path))
-
-        # Check if the directory exists, if not, create it
-        await full_file_path.parent.mkdir(parents=True, exist_ok=True)
-
-        # Determine if the file exists and write the new content
-        file_content, created = await read_or_init_file(full_file_path)
-
-        success = await open_file_and_apply_edit(
-            file_path=full_file_path,
-            file_content=file_content,
-            edit_content=edit_content,
-            edit_function=edit_function,
-        )
-
-        if success:
-            verb = "Created" if created else "Modified"
-            return f"{verb} file {file_path}"
-        else:
-            verb = "create" if created else "modify"
-            return f"Failed to {verb} file {file_path}"
-
     except Exception as e:
-        raise e
-
-
-async def read_or_init_file(file_path: FilePath) -> tuple[str, bool]:
-    path = AsyncPath(file_path)
-
-    if not (await path.exists()):
-        await path.touch()
-        return "", True
-
-    content = await safe_read_file(path)
-    return content, False
-
-
-async def open_file_and_apply_edit(
-    file_path: FilePath,
-    file_content: str,
-    edit_content: str,
-    edit_function: Callable[[str, str], FileEditResult],
-) -> bool:
-    result = edit_function(file_content, edit_content)
-
-    if not result.content:
-        return False
-
-    await safe_write_file(file_path, result.content)
-
-    return True
-
-
-def find_leading_whitespace(existing_content: str, search: str) -> str | None:
-    existing_lines = existing_content.splitlines()
-
-    search_line_count = len(search.splitlines())
-    dedented_search = dedent(search)
-
-    for i in range(len(existing_lines)):
-        existing_window_content = "\n".join(existing_lines[i : i + search_line_count])
-        dedented_existing_window = dedent(existing_window_content)
-
-        leading_ws_len = len(existing_window_content) - len(
-            existing_window_content.lstrip()
-        )
-        leading_ws = existing_window_content[:leading_ws_len]
-
-        if dedented_existing_window == dedented_search:
-            return leading_ws
-
-    return None
-
-
-def try_fix_whitespace(
-    existing_content: str, search: str, replace: str
-) -> tuple[str, str] | None:
-    # Try to fix the whitespace of the search and replace
-    # to make the edit more likely to apply
-    leading_ws = find_leading_whitespace(existing_content, search)
-    if leading_ws is None:
-        return None
-
-    dedented_search = dedent(search)
-    dedented_replace = dedent(replace)
-
-    return indent(dedented_search, leading_ws), indent(dedented_replace, leading_ws)
-
-
-def try_search_replace(existing_content: str, search: str, replace: str) -> str | None:
-    # Try simple search and replace first
-    new_content = simple_search_and_replace(existing_content, search, replace)
-    if new_content:
-        return new_content
-
-    fixed_ws = try_fix_whitespace(existing_content, search, replace)
-    if not fixed_ws:
-        return None
-
-    search, replace = fixed_ws
-
-    new_content = simple_search_and_replace(existing_content, search, replace)
-    if new_content:
-        return new_content
-
-    return None
-
-
-def try_diff_patch(existing_content: str, search: str, replace: str) -> str | None:
-    new_content = diff_patch_search_and_replace(existing_content, search, replace)
-    if new_content:
-        print("Applied diff patch search and replace")
-        return new_content
-
-    return None
-
-
-def apply_udiff(existing_content: str, diff_content: str) -> FileEditResult:
-    hunks = get_raw_udiff_hunks(diff_content)
-
-    for hunk in hunks:
-        if not hunk:
-            continue
-
-        search, replace = split_hunk_for_search_and_replace(hunk)
-
-        # Exact match
-        new_content = try_search_replace(existing_content, search, replace)
-        if new_content is not None:
-            print("Applied successfully!")
-            return FileEditResult(content=new_content, failed_edits=[])
-
-        # Fuzzy match
-        new_content = try_diff_patch(existing_content, search, replace)
-        if new_content is not None:
-            print("Applied successfully!")
-            return FileEditResult(content=new_content, failed_edits=[])
-
-        print("Failed to apply hunk, exiting!")
-        return FileEditResult(content=None, failed_edits=[(search, replace)])
-
-    return FileEditResult(content=existing_content, failed_edits=[])
-
-
-def get_raw_udiff_hunks(content: str) -> list[list[str]]:
-    lines = content.splitlines(keepends=True)
-    hunks: list[list[str]] = []
-    current_hunk: list[str] = []
-    for line in lines:
-        if line.startswith("@@"):
-            if current_hunk:
-                hunks.append(current_hunk)
-                current_hunk = []
-        else:
-            current_hunk.append(line)
-    if current_hunk:
-        hunks.append(current_hunk)
-    return hunks
-
-
-def split_hunk_for_search_and_replace(hunk: list[str]) -> tuple[str, str]:
-    search_lines = []
-    replace_lines = []
-
-    search_prefixes = ["-", " "]
-    replace_prefixes = ["+", " "]
-    for line in hunk:
-        if not line:
-            continue
-        prefix, content = line[0], line[1:]
-        if not content:
-            continue
-        if prefix in search_prefixes:
-            search_lines.append(content)
-        if prefix in replace_prefixes:
-            replace_lines.append(content)
-    return "".join(search_lines), "".join(replace_lines)
-
-
-def simple_search_and_replace(content: str, search: str, replace: str) -> str | None:
-    if content.count(search) >= 1:
-        return content.replace(search, replace)
-    return None
-
-
-def diff_patch_search_and_replace(
-    content: str, search: str, replace: str
-) -> str | None:
-    patcher = diff_match_patch()
-    # 3 second tieout for computing diffs
-    patcher.Diff_Timeout = 3
-    patcher.Match_Threshold = 0.95
-    patcher.Match_Distance = 500
-    patcher.Match_MaxBits = 128
-    patcher.Patch_Margin = 32
-    search_vs_replace_diff = patcher.diff_main(search, replace, False)
-
-    # Simplify the diff as much as possible
-    patcher.diff_cleanupEfficiency(search_vs_replace_diff)
-    patcher.diff_cleanupSemantic(search_vs_replace_diff)
-
-    original_vs_search_diff = patcher.diff_main(search, content)
-    new_diffs = patcher.patch_make(search, search_vs_replace_diff)
-    # Offset the search vs. replace diffs with the offset
-    # of the search diff within the original content.
-    for new_diff in new_diffs:
-        new_diff.start1 = patcher.diff_xIndex(original_vs_search_diff, new_diff.start1)
-        new_diff.start2 = patcher.diff_xIndex(original_vs_search_diff, new_diff.start2)
-
-    new_content, successes = patcher.patch_apply(new_diffs, content)
-    if not all(successes):
-        return None
-
-    return str(new_content)
-
-
-SEARCH_REPLACE_RE = re.compile(
-    r"[^<>]*<<<+\s*SEARCH\n((?P<search>.*?)\n)??===+\n((?P<replace>.*?)\n)??>>>+\s*?REPLACE\s*?[^<>]*",
-    re.DOTALL,
-)
-
-TAGGED_SEARCH_REPLACE_RE = re.compile(
-    r"<search>(?P<search>.*?)??</search>\s*?<replace>(?P<replace>.*?)??</replace>",
-    re.DOTALL,
-)
-
-
-def apply_search_replace(result: str, search: str, replace: str) -> str | None:
-    if not search and not replace:
-        # Nonsense
-        return None
-
-    if not search and not result:
-        # New file, just return replace
-        return replace
-
-    if not search.strip():
-        # Search on just whitespace,
-        # too dangerous to apply
-        return None
-
-    return try_search_replace(result, search, replace)
-
-
-def apply_all_search_replace(
-    existing_content: str,
-    sr_content: str,
-    match_re: re.Pattern[str] = SEARCH_REPLACE_RE,
-) -> FileEditResult:
-    # Same as apply_search_replace, but applies all search and replace pairs
-    # in the sr_content to the existing_content
-
-    result = existing_content
-    failed_edits: list[tuple[str, str]] = []
-
-    for match in match_re.finditer(sr_content):
-        match_dict = match.groupdict()
-        search, replace = match_dict.get("search"), match_dict.get("replace")
-        search = search or ""
-        replace = replace or ""
-
-        new_result = apply_search_replace(result, search, replace)
-        if new_result is None:
-            failed_edits.append((search, replace))
-            continue
-
-        result = new_result
-
-    return FileEditResult(content=result, failed_edits=failed_edits)
-
-
-def apply_all_tagged_search_replace(
-    existing_content: str, sr_content: str
-) -> FileEditResult:
-    return apply_all_search_replace(
-        existing_content, sr_content, TAGGED_SEARCH_REPLACE_RE
-    )
+        return f"An error occurred: {e!s}"