indent 0.1.26__py3-none-any.whl

exponent/core/remote_execution/client.py

@@ -0,0 +1,999 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import time
+import uuid
+from collections.abc import AsyncGenerator, Callable, Coroutine, Generator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import Any, TypeVar, cast
+
+import msgspec
+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
+from exponent.core.remote_execution import files, system_context
+from exponent.core.remote_execution.cli_rpc_types import (
+    BashToolInput,
+    BatchToolExecutionRequest,
+    BatchToolExecutionResponse,
+    CliRpcRequest,
+    CliRpcResponse,
+    ErrorResponse,
+    ErrorToolResult,
+    GenerateUploadUrlRequest,
+    GenerateUploadUrlResponse,
+    GetAllFilesRequest,
+    GetAllFilesResponse,
+    HttpRequest,
+    KeepAliveCliChatRequest,
+    KeepAliveCliChatResponse,
+    StartTerminalRequest,
+    StartTerminalResponse,
+    StopTerminalRequest,
+    StopTerminalResponse,
+    SwitchCLIChatRequest,
+    SwitchCLIChatResponse,
+    TerminalInputRequest,
+    TerminalInputResponse,
+    TerminalResizeRequest,
+    TerminalResizeResponse,
+    TerminateRequest,
+    TerminateResponse,
+    ToolExecutionRequest,
+    ToolExecutionResponse,
+    ToolResultType,
+)
+from exponent.core.remote_execution.code_execution import (
+    execute_code_streaming,
+)
+from exponent.core.remote_execution.files import file_walk
+from exponent.core.remote_execution.http_fetch import fetch_http_content
+from exponent.core.remote_execution.session import (
+    RemoteExecutionClientSession,
+    get_session,
+    send_exception_log,
+)
+from exponent.core.remote_execution.terminal_session import TerminalSessionManager
+from exponent.core.remote_execution.terminal_types import TerminalMessage
+from exponent.core.remote_execution.tool_execution import (
+    execute_bash_tool,
+    execute_tool,
+    truncate_result,
+)
+from exponent.core.remote_execution.types import (
+    ChatSource,
+    CLIConnectedState,
+    CreateChatResponse,
+    HeartbeatInfo,
+    RunWorkflowRequest,
+    WorkflowInput,
+    WorkflowTriggerRequest,
+    WorkflowTriggerResponse,
+)
+from exponent.core.remote_execution.utils import (
+    deserialize_api_response,
+)
+from exponent.utils.version import get_installed_version
+
+logger = logging.getLogger(__name__)
+
+
+TModel = TypeVar("TModel", bound=BaseModel)
+
+
+@dataclass
+class WSDisconnected:
+    error_message: str | None = None
+
+
+@dataclass
+class SwitchCLIChat:
+    new_chat_uuid: str
+
+
+REMOTE_EXECUTION_CLIENT_EXIT_INFO = WSDisconnected | SwitchCLIChat
+
+# UUID for a single run of the CLI
+cli_uuid = uuid.uuid4()
+
+
+class RemoteExecutionClient:
+    def __init__(
+        self,
+        session: RemoteExecutionClientSession,
+        file_cache: files.FileCache | None = None,
+    ):
+        self.current_session = session
+        self.file_cache = file_cache or files.FileCache(session.working_directory)
+
+        # for active code executions, track whether they should be halted
+        # correlation_id -> should_halt
+        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
+
+        # Track pending upload URL requests
+        self._pending_upload_requests: dict[
+            str, asyncio.Future[GenerateUploadUrlResponse]
+        ] = {}
+        self._upload_request_lock = asyncio.Lock()
+        self._websocket: ClientConnection | None = None
+
+    @property
+    def working_directory(self) -> str:
+        return self.current_session.working_directory
+
+    @property
+    def api_client(self) -> AsyncClient:
+        return self.current_session.api_client
+
+    @property
+    def ws_client(self) -> AsyncClient:
+        return self.current_session.ws_client
+
+    async def add_code_execution_to_halt_states(self, correlation_id: str) -> None:
+        async with self._halt_lock:
+            self._halt_states[correlation_id] = False
+
+    async def halt_all_code_executions(self) -> None:
+        logger.info(f"Halting all code executions: {self._halt_states}")
+        async with self._halt_lock:
+            self._halt_states = {
+                correlation_id: True for correlation_id in self._halt_states.keys()
+            }
+
+    async def clear_halt_state(self, correlation_id: str) -> None:
+        async with self._halt_lock:
+            self._halt_states.pop(correlation_id, None)
+
+    def get_halt_check(self, correlation_id: str) -> Callable[[], bool]:
+        def should_halt() -> bool:
+            # Don't need to lock here, since just reading from dict
+            return self._halt_states.get(correlation_id, False)
+
+        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(  # noqa: PLR0911, PLR0915
+        self,
+        msg: str,
+        websocket: ClientConnection,
+        requests: asyncio.Queue[CliRpcRequest],
+        terminal_session_manager: TerminalSessionManager,
+    ) -> REMOTE_EXECUTION_CLIENT_EXIT_INFO | None:
+        """Handle an incoming websocket message.
+        Returns None to continue processing, or a REMOTE_EXECUTION_CLIENT_EXIT_INFO to exit."""
+
+        self._last_request_time = time.time()
+
+        msg_data = json.loads(msg)
+        if msg_data["type"] == "result":
+            data = json.dumps(msg_data["data"])
+            try:
+                response = msgspec.json.decode(data, type=CliRpcResponse)
+                if isinstance(response.response, GenerateUploadUrlResponse):
+                    async with self._upload_request_lock:
+                        if response.request_id in self._pending_upload_requests:
+                            future = self._pending_upload_requests.pop(
+                                response.request_id
+                            )
+                            future.set_result(response.response)
+            except Exception as e:
+                logger.error(f"Error handling upload URL response: {e}")
+            return None
+        elif msg_data["type"] != "request":
+            return None
+
+        data = json.dumps(msg_data["data"])
+        try:
+            request = msgspec.json.decode(data, type=CliRpcRequest)
+        except (msgspec.DecodeError, msgspec.ValidationError) as e:
+            # Try and decode to get request_id if possible
+            request = msgspec.json.decode(data)
+            if isinstance(request, dict) and "request_id" in request:
+                request_id = request["request_id"]
+                if (
+                    request.get("request", {}).get("type", {}) == "tool_execution"
+                ) and (
+                    "tool_input" in request["request"]
+                    and "tool_name" in request["request"]["tool_input"]
+                ):
+                    tool_name = request["request"]["tool_input"]["tool_name"]
+                    logger.error(
+                        f"Error tool {tool_name} received in a request."
+                        "Please ensure you are running the latest version of Indent. If this issue persists, please contact support."
+                    )
+                    await websocket.send(
+                        json.dumps(
+                            {
+                                "type": "result",
+                                "data": msgspec.to_builtins(
+                                    CliRpcResponse(
+                                        request_id=request_id,
+                                        response=ErrorResponse(
+                                            error_message=f"Unknown tool: {tool_name}. If you are running an older version of Indent, please upgrade to the latest version to ensure compatibility."
+                                        ),
+                                    )
+                                ),
+                            }
+                        )
+                    )
+                else:
+                    logger.error(
+                        "Error decoding cli rpc request. Please ensure you are running the latest version of Indent."
+                    )
+                    await websocket.send(
+                        json.dumps(
+                            {
+                                "type": "result",
+                                "data": msgspec.to_builtins(
+                                    CliRpcResponse(
+                                        request_id=request_id,
+                                        response=ErrorResponse(
+                                            error_message=f"Unknown cli rpc request type: {request}",
+                                        ),
+                                    )
+                                ),
+                            }
+                        )
+                    )
+
+                return None
+            else:
+                # If we couldn't get a request_id, re-raise and fail noisily
+                raise e
+
+        if isinstance(request.request, TerminateRequest):
+            await self.halt_all_code_executions()
+            await websocket.send(
+                json.dumps(
+                    {
+                        "type": "result",
+                        "data": msgspec.to_builtins(
+                            CliRpcResponse(
+                                request_id=request.request_id,
+                                response=TerminateResponse(),
+                            )
+                        ),
+                    }
+                )
+            )
+            return None
+        elif isinstance(request.request, SwitchCLIChatRequest):
+            await websocket.send(
+                json.dumps(
+                    {
+                        "type": "result",
+                        "data": msgspec.to_builtins(
+                            CliRpcResponse(
+                                request_id=request.request_id,
+                                response=SwitchCLIChatResponse(),
+                            )
+                        ),
+                    }
+                )
+            )
+            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
+        elif isinstance(request.request, StartTerminalRequest):
+            # Start a new terminal session
+            session_id = await terminal_session_manager.start_session(
+                websocket=websocket,
+                session_id=request.request.session_id,
+                command=request.request.command,
+                cols=request.request.cols,
+                rows=request.request.rows,
+                env=request.request.env,
+            )
+            await websocket.send(
+                json.dumps(
+                    {
+                        "type": "result",
+                        "data": msgspec.to_builtins(
+                            CliRpcResponse(
+                                request_id=request.request_id,
+                                response=StartTerminalResponse(
+                                    session_id=session_id, success=True
+                                ),
+                            )
+                        ),
+                    }
+                )
+            )
+            return None
+        elif isinstance(request.request, TerminalInputRequest):
+            # Send input to terminal session
+            success = await terminal_session_manager.send_input(
+                session_id=request.request.session_id,
+                data=request.request.data,
+            )
+            await websocket.send(
+                json.dumps(
+                    {
+                        "type": "result",
+                        "data": msgspec.to_builtins(
+                            CliRpcResponse(
+                                request_id=request.request_id,
+                                response=TerminalInputResponse(
+                                    session_id=request.request.session_id,
+                                    success=success,
+                                ),
+                            )
+                        ),
+                    }
+                )
+            )
+            return None
+        elif isinstance(request.request, TerminalResizeRequest):
+            # Resize terminal session
+            success = await terminal_session_manager.resize_terminal(
+                session_id=request.request.session_id,
+                rows=request.request.rows,
+                cols=request.request.cols,
+            )
+            await websocket.send(
+                json.dumps(
+                    {
+                        "type": "result",
+                        "data": msgspec.to_builtins(
+                            CliRpcResponse(
+                                request_id=request.request_id,
+                                response=TerminalResizeResponse(
+                                    session_id=request.request.session_id,
+                                    success=success,
+                                ),
+                            )
+                        ),
+                    }
+                )
+            )
+            return None
+        elif isinstance(request.request, StopTerminalRequest):
+            # Stop terminal session
+            success = await terminal_session_manager.stop_session(
+                session_id=request.request.session_id
+            )
+            await websocket.send(
+                json.dumps(
+                    {
+                        "type": "result",
+                        "data": msgspec.to_builtins(
+                            CliRpcResponse(
+                                request_id=request.request_id,
+                                response=StopTerminalResponse(
+                                    session_id=request.request.session_id,
+                                    success=success,
+                                ),
+                            )
+                        ),
+                    }
+                )
+            )
+            return None
+        else:
+            if isinstance(request.request, ToolExecutionRequest) and isinstance(
+                request.request.tool_input, BashToolInput
+            ):
+                await self.add_code_execution_to_halt_states(request.request_id)
+            elif isinstance(request.request, BatchToolExecutionRequest):
+                # Add halt state if any of the batch tools are bash commands
+                if any(
+                    isinstance(tool_input, BashToolInput)
+                    for tool_input in request.request.tool_inputs
+                ):
+                    await self.add_code_execution_to_halt_states(request.request_id)
+
+            await requests.put(request)
+            return None
+
+    async def _setup_tasks(
+        self,
+        beats: asyncio.Queue[HeartbeatInfo],
+        requests: asyncio.Queue[CliRpcRequest],
+        results: asyncio.Queue[CliRpcResponse],
+    ) -> list[asyncio.Task[None]]:
+        """Setup heartbeat and executor tasks."""
+
+        async def beat() -> None:
+            while True:
+                info = await self.get_heartbeat_info()
+                await beats.put(info)
+                await asyncio.sleep(3)
+
+        # Lock to ensure that only one executor can grab a
+        # request at a time.
+        requests_lock = asyncio.Lock()
+
+        # Lock to ensure that only one executor can put a
+        # result in the results queue at a time.
+        results_lock = asyncio.Lock()
+
+        async def executor() -> None:
+            # We use locks here to protect the request/result
+            # queues from being accessed by multiple executors.
+            while True:
+                async with requests_lock:
+                    request = await requests.get()
+
+                try:
+                    # Check if this is a streaming request
+                    from exponent.core.remote_execution.cli_rpc_types import (
+                        StreamingCodeExecutionRequest,
+                    )
+
+                    if isinstance(request.request, StreamingCodeExecutionRequest):
+                        async for streaming_response in self.handle_streaming_request(
+                            request.request
+                        ):
+                            async with results_lock:
+                                await results.put(
+                                    CliRpcResponse(
+                                        request_id=request.request_id,
+                                        response=streaming_response,
+                                    )
+                                )
+                    else:
+                        # Note that we don't want to hold the lock here
+                        # because we want other executors to be able to
+                        # grab requests while we're handling a request.
+                        logger.info(f"Handling request {request}")
+                        response = await self.handle_request(request)
+                        async with results_lock:
+                            logger.info(f"Putting response {response}")
+                            await results.put(response)
+                except Exception as e:
+                    logger.info(f"Error handling request {request}:\n\n{e}")
+                    try:
+                        await send_exception_log(e, session=self.current_session)
+                    except Exception:
+                        pass
+                    async with results_lock:
+                        from exponent.core.remote_execution.cli_rpc_types import (
+                            StreamingCodeExecutionRequest,
+                            StreamingErrorResponse,
+                        )
+
+                        if isinstance(request.request, StreamingCodeExecutionRequest):
+                            # For streaming requests, send a streaming error response
+                            await results.put(
+                                CliRpcResponse(
+                                    request_id=request.request_id,
+                                    response=StreamingErrorResponse(
+                                        correlation_id=request.request.correlation_id,
+                                        error_message=str(e),
+                                    ),
+                                )
+                            )
+                        else:
+                            await results.put(
+                                CliRpcResponse(
+                                    request_id=request.request_id,
+                                    response=ErrorResponse(
+                                        error_message=str(e),
+                                    ),
+                                )
+                            )
+
+        beat_task = asyncio.create_task(beat())
+        # Three parallel executors to handle requests
+
+        executor_tasks = [
+            asyncio.create_task(executor()),
+            asyncio.create_task(executor()),
+            asyncio.create_task(executor()),
+        ]
+
+        return [beat_task, *executor_tasks]
+
+    async def _process_websocket_messages(
+        self,
+        websocket: ClientConnection,
+        beats: asyncio.Queue[HeartbeatInfo],
+        requests: asyncio.Queue[CliRpcRequest],
+        results: asyncio.Queue[CliRpcResponse],
+        terminal_output_queue: asyncio.Queue[TerminalMessage],
+        terminal_session_manager: TerminalSessionManager,
+    ) -> REMOTE_EXECUTION_CLIENT_EXIT_INFO:
+        """Process messages from the websocket connection."""
+        pending: set[asyncio.Task[object]] = set()
+        try:
+            recv = asyncio.create_task(websocket.recv())
+            get_beat = asyncio.create_task(beats.get())
+            get_result = asyncio.create_task(results.get())
+            get_terminal_output = asyncio.create_task(terminal_output_queue.get())
+            pending = {recv, get_beat, get_result, get_terminal_output}
+
+            while True:
+                done, pending = await asyncio.wait(
+                    pending, return_when=asyncio.FIRST_COMPLETED
+                )
+
+                if recv in done:
+                    msg = str(recv.result())
+                    exit_info = await self._handle_websocket_message(
+                        msg, websocket, requests, terminal_session_manager
+                    )
+                    if exit_info is not None:
+                        return exit_info
+
+                    recv = asyncio.create_task(websocket.recv())
+                    pending.add(recv)
+
+                if get_beat in done:
+                    info = get_beat.result()
+                    data = json.loads(info.model_dump_json())
+                    msg = json.dumps({"type": "heartbeat", "data": data})
+                    await websocket.send(msg)
+
+                    get_beat = asyncio.create_task(beats.get())
+                    pending.add(get_beat)
+
+                if get_result in done:
+                    response = get_result.result()
+                    # All responses are now CliRpcResponse with msgspec
+                    data = msgspec.to_builtins(response)
+                    msg = json.dumps({"type": "result", "data": data})
+                    await websocket.send(msg)
+
+                    get_result = asyncio.create_task(results.get())
+                    pending.add(get_result)
+
+                if get_terminal_output in done:
+                    terminal_message = get_terminal_output.result()
+                    data = msgspec.to_builtins(terminal_message)
+                    msg = json.dumps({"type": "terminal_message", "data": data})
+                    await websocket.send(msg)
+
+                    get_terminal_output = asyncio.create_task(
+                        terminal_output_queue.get()
+                    )
+                    pending.add(get_terminal_output)
+        finally:
+            for task in pending:
+                task.cancel()
+
+            await asyncio.gather(*pending, return_exceptions=True)
+
+    async def _handle_websocket_connection(
+        self,
+        websocket: ClientConnection,
+        connection_tracker: ConnectionTracker | None,
+        beats: asyncio.Queue[HeartbeatInfo],
+        requests: asyncio.Queue[CliRpcRequest],
+        results: asyncio.Queue[CliRpcResponse],
+        terminal_output_queue: asyncio.Queue[TerminalMessage],
+        terminal_session_manager: TerminalSessionManager,
+    ) -> REMOTE_EXECUTION_CLIENT_EXIT_INFO | None:
+        """Handle a single websocket connection.
+        Returns None to continue with reconnection attempts, or an exit info to terminate."""
+        if connection_tracker is not None:
+            await connection_tracker.set_connected(True)
+
+        self._websocket = websocket
+
+        try:
+            return await self._process_websocket_messages(
+                websocket,
+                beats,
+                requests,
+                results,
+                terminal_output_queue,
+                terminal_session_manager,
+            )
+        except websockets.exceptions.ConnectionClosed as e:
+            if e.rcvd is not None:
+                if e.rcvd.code == 1000:
+                    # Normal closure, exit completely
+                    return WSDisconnected()
+                elif e.rcvd.code == 1008:
+                    error_message = (
+                        "Error connecting to websocket"
+                        if e.rcvd.reason is None
+                        else e.rcvd.reason
+                    )
+                    return WSDisconnected(error_message=error_message)
+            # Otherwise, allow reconnection attempt
+            logger.debug("Websocket connection closed by remote.")
+            return None
+        except TimeoutError:
+            # Timeout, allow reconnection attempt
+            # TODO: investgate if this is needed, possibly scope it down
+            return None
+        finally:
+            if connection_tracker is not None:
+                await connection_tracker.set_connected(False)
+
+    async def run_connection(
+        self,
+        chat_uuid: str,
+        connection_tracker: ConnectionTracker | None = None,
+        timeout_seconds: int | None = None,
+    ) -> REMOTE_EXECUTION_CLIENT_EXIT_INFO:
+        """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()
+
+        # Create queues ONCE - persist across reconnections
+        beats: asyncio.Queue[HeartbeatInfo] = asyncio.Queue()
+        requests: asyncio.Queue[CliRpcRequest] = asyncio.Queue()
+        results: asyncio.Queue[CliRpcResponse] = asyncio.Queue()
+        terminal_output_queue: asyncio.Queue[TerminalMessage] = asyncio.Queue()
+
+        # Create terminal session manager ONCE - persist across reconnections
+        terminal_session_manager = TerminalSessionManager(terminal_output_queue)
+
+        # Create tasks ONCE - persist across reconnections
+        executors = await self._setup_tasks(beats, requests, results)
+
+        try:
+            async for websocket in self.ws_connect(f"/api/ws/chat/{chat_uuid}"):
+                # 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,
+                                beats,
+                                requests,
+                                results,
+                                terminal_output_queue,
+                                terminal_session_manager,
+                            )
+                        ),
+                        asyncio.create_task(self._timeout_monitor(timeout_seconds)),
+                    ],
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+
+                # 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
+            return WSDisconnected(
+                error_message="Could not establish websocket connection"
+            )
+        finally:
+            # Stop all terminal sessions to clean up PTY processes
+            await terminal_session_manager.stop_all_sessions()
+
+            # Cancel all background tasks when exiting
+            for task in executors:
+                task.cancel()
+            await asyncio.gather(*executors, return_exceptions=True)
+
+    async def create_chat(self, chat_source: ChatSource) -> CreateChatResponse:
+        response = await self.api_client.post(
+            "/api/remote_execution/create_chat",
+            params={"chat_source": chat_source.value},
+        )
+        return await deserialize_api_response(response, CreateChatResponse)
+
+    # deprecated
+    async def run_workflow(self, chat_uuid: str, workflow_id: str) -> dict[str, Any]:
+        response = await self.api_client.post(
+            "/api/remote_execution/run_workflow",
+            json=RunWorkflowRequest(
+                chat_uuid=chat_uuid,
+                workflow_id=workflow_id,
+            ).model_dump(),
+            timeout=60,
+        )
+        if response.status_code != http_status.OK:
+            raise Exception(
+                f"Failed to run workflow with status code {response.status_code} and response {response.text}"
+            )
+        return cast(dict[str, Any], response.json())
+
+    async def trigger_workflow(
+        self, workflow_name: str, workflow_input: WorkflowInput
+    ) -> WorkflowTriggerResponse:
+        response = await self.api_client.post(
+            "/api/remote_execution/trigger_workflow",
+            json=WorkflowTriggerRequest(
+                workflow_name=workflow_name,
+                workflow_input=workflow_input,
+            ).model_dump(),
+        )
+        return await deserialize_api_response(response, WorkflowTriggerResponse)
+
+    async def get_heartbeat_info(self) -> HeartbeatInfo:
+        return HeartbeatInfo(
+            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:
+        logger.info(f"Sending heartbeat for chat_uuid {chat_uuid}")
+        heartbeat_info = await self.get_heartbeat_info()
+        response = await self.api_client.post(
+            f"/api/remote_execution/{chat_uuid}/heartbeat",
+            content=heartbeat_info.model_dump_json(),
+            timeout=60,
+        )
+        if response.status_code != http_status.OK:
+            raise Exception(
+                f"Heartbeat failed with status code {response.status_code} and response {response.text}"
+            )
+        connected_state = await deserialize_api_response(response, CLIConnectedState)
+        logger.info(f"Heartbeat response: {connected_state}")
+        return connected_state
+
+    async def request_upload_url(
+        self, s3_key: str, content_type: str
+    ) -> GenerateUploadUrlResponse:
+        if self._websocket is None:
+            raise RuntimeError("No active websocket connection")
+
+        request_id = str(uuid.uuid4())
+        request = CliRpcRequest(
+            request_id=request_id,
+            request=GenerateUploadUrlRequest(s3_key=s3_key, content_type=content_type),
+        )
+
+        future: asyncio.Future[GenerateUploadUrlResponse] = asyncio.Future()
+        async with self._upload_request_lock:
+            self._pending_upload_requests[request_id] = future
+
+        try:
+            await self._websocket.send(
+                json.dumps({"type": "request", "data": msgspec.to_builtins(request)})
+            )
+
+            response = await asyncio.wait_for(future, timeout=30)
+            return response
+        except TimeoutError:
+            async with self._upload_request_lock:
+                self._pending_upload_requests.pop(request_id, None)
+            raise RuntimeError("Timeout waiting for upload URL response")
+        except Exception as e:
+            async with self._upload_request_lock:
+                self._pending_upload_requests.pop(request_id, None)
+            raise e
+
+    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):
+                    raw_result = await execute_bash_tool(
+                        request.request.tool_input,
+                        self.working_directory,
+                        should_halt=self.get_halt_check(request.request_id),
+                    )
+                else:
+                    raw_result = await execute_tool(  # type: ignore[assignment]
+                        request.request.tool_input, self.working_directory, self
+                    )
+                tool_result = truncate_result(raw_result)
+                return CliRpcResponse(
+                    request_id=request.request_id,
+                    response=ToolExecutionResponse(
+                        tool_result=tool_result,
+                    ),
+                )
+            elif isinstance(request.request, GetAllFilesRequest):
+                files = await file_walk(self.working_directory)
+                return CliRpcResponse(
+                    request_id=request.request_id,
+                    response=GetAllFilesResponse(files=files),
+                )
+            elif isinstance(request.request, BatchToolExecutionRequest):
+                coros: list[Coroutine[Any, Any, ToolResultType]] = []
+                for tool_input in request.request.tool_inputs:
+                    if isinstance(tool_input, BashToolInput):
+                        coros.append(
+                            execute_bash_tool(
+                                tool_input,
+                                self.working_directory,
+                                should_halt=self.get_halt_check(request.request_id),
+                            )
+                        )
+                    else:
+                        coros.append(
+                            execute_tool(tool_input, self.working_directory, self)
+                        )
+
+                results: list[ToolResultType | BaseException] = await asyncio.gather(
+                    *coros, return_exceptions=True
+                )
+
+                processed_results: list[ToolResultType] = []
+                for result in results:
+                    if not isinstance(result, BaseException):
+                        processed_results.append(truncate_result(result))
+                    else:
+                        processed_results.append(
+                            ErrorToolResult(error_message=str(result))
+                        )
+
+                return CliRpcResponse(
+                    request_id=request.request_id,
+                    response=BatchToolExecutionResponse(
+                        tool_results=processed_results,
+                    ),
+                )
+            elif isinstance(request.request, HttpRequest):
+                http_response = await fetch_http_content(request.request)
+                return CliRpcResponse(
+                    request_id=request.request_id,
+                    response=http_response,
+                )
+            elif isinstance(request.request, TerminateRequest):
+                raise ValueError(
+                    "TerminateRequest should not be handled by handle_request"
+                )
+
+            elif isinstance(request.request, SwitchCLIChatRequest):
+                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"
+                )
+            elif isinstance(request.request, StartTerminalRequest):
+                raise ValueError(
+                    "StartTerminalRequest should not be handled by handle_request"
+                )
+            elif isinstance(request.request, TerminalInputRequest):
+                raise ValueError(
+                    "TerminalInputRequest should not be handled by handle_request"
+                )
+            elif isinstance(request.request, TerminalResizeRequest):
+                raise ValueError(
+                    "TerminalResizeRequest should not be handled by handle_request"
+                )
+            elif isinstance(request.request, StopTerminalRequest):
+                raise ValueError(
+                    "StopTerminalRequest should not be handled by handle_request"
+                )
+
+            raise ValueError(f"Unhandled request type: {type(request)}")
+
+        except Exception as e:
+            logger.error(f"Error handling request {request}:\n\n{e}")
+            raise e
+        finally:
+            # Clean up halt state after request is complete
+            if isinstance(request.request, ToolExecutionRequest) and isinstance(
+                request.request.tool_input, BashToolInput
+            ):
+                await self.clear_halt_state(request.request_id)
+            elif isinstance(request.request, BatchToolExecutionRequest):
+                # Clear halt state if any of the batch tools were bash commands
+                if any(
+                    isinstance(tool_input, BashToolInput)
+                    for tool_input in request.request.tool_inputs
+                ):
+                    await self.clear_halt_state(request.request_id)
+
+    async def handle_streaming_request(
+        self,
+        request: Any,
+    ) -> AsyncGenerator[Any, None]:
+        from exponent.core.remote_execution.cli_rpc_types import (
+            StreamingCodeExecutionRequest,
+        )
+
+        if not isinstance(request, StreamingCodeExecutionRequest):
+            assert False, f"{type(request)} should be sent to handle_streaming_request"
+        async for output in execute_code_streaming(
+            request,
+            self.current_session,
+            working_directory=self.working_directory,
+            should_halt=self.get_halt_check(request.correlation_id),
+        ):
+            yield output
+
+    def ws_connect(self, path: str) -> connect:
+        base_url = (
+            str(self.ws_client.base_url)
+            .replace("http://", "ws://")
+            .replace("https://", "wss://")
+        )
+
+        url = f"{base_url}{path}"
+        headers = {"api-key": self.api_client.headers["api-key"]}
+
+        def custom_backoff() -> Generator[float, None, None]:
+            yield 0.1  # short initial delay
+
+            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
+
+    @staticmethod
+    @asynccontextmanager
+    async def session(
+        api_key: str,
+        base_url: str,
+        base_ws_url: str,
+        working_directory: str,
+        file_cache: files.FileCache | None = None,
+    ) -> AsyncGenerator[RemoteExecutionClient, None]:
+        async with get_session(
+            working_directory, base_url, base_ws_url, api_key
+        ) as session:
+            yield RemoteExecutionClient(session, file_cache)