indent 0.1.13__py3-none-any.whl → 0.1.28__py3-none-any.whl

exponent/core/remote_execution/client.py

@@ -5,7 +5,7 @@ import json
 import logging
 import time
 import uuid
-from collections.abc import AsyncGenerator, Callable, Generator
+from collections.abc import AsyncGenerator, Callable, Coroutine, Generator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from typing import Any, TypeVar, cast
@@ -30,13 +30,24 @@ from exponent.core.remote_execution.cli_rpc_types import (
     CliRpcRequest,
     CliRpcResponse,
     ErrorResponse,
+    ErrorToolResult,
+    GenerateUploadUrlRequest,
+    GenerateUploadUrlResponse,
     GetAllFilesRequest,
     GetAllFilesResponse,
     HttpRequest,
     KeepAliveCliChatRequest,
     KeepAliveCliChatResponse,
+    StartTerminalRequest,
+    StartTerminalResponse,
+    StopTerminalRequest,
+    StopTerminalResponse,
     SwitchCLIChatRequest,
     SwitchCLIChatResponse,
+    TerminalInputRequest,
+    TerminalInputResponse,
+    TerminalResizeRequest,
+    TerminalResizeResponse,
     TerminateRequest,
     TerminateResponse,
     ToolExecutionRequest,
@@ -51,7 +62,10 @@ 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,
@@ -61,12 +75,9 @@ from exponent.core.remote_execution.types import (
     ChatSource,
     CLIConnectedState,
     CreateChatResponse,
-    GitInfo,
     HeartbeatInfo,
-    PrReviewWorkflowInput,
-    RemoteExecutionResponseType,
     RunWorkflowRequest,
-    StreamingCodeExecutionRequest,
+    WorkflowInput,
     WorkflowTriggerRequest,
     WorkflowTriggerResponse,
 )
@@ -114,6 +125,13 @@ class RemoteExecutionClient:
         # 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
@@ -173,23 +191,40 @@ class RemoteExecutionClient:
             # Handle cancellation gracefully
             return None
 
-    async def _handle_websocket_message(
+    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"] != "request":
+        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 as e:
+        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:
@@ -291,6 +326,101 @@ class RemoteExecutionClient:
                 )
             )
             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
@@ -337,33 +467,64 @@ class RemoteExecutionClient:
                     request = await requests.get()
 
                 try:
-                    # if isinstance(request, StreamingCodeExecutionRequest):
-                    #     async for streaming_response in self.handle_streaming_request(
-                    #         request
-                    #     ):
-                    #         async with results_lock:
-                    #             await results.put(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)
+                    # 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:
-                        await results.put(
-                            CliRpcResponse(
-                                request_id=request.request_id,
-                                response=ErrorResponse(
-                                    error_message=str(e),
-                                ),
-                            )
+                        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
 
@@ -381,13 +542,17 @@ class RemoteExecutionClient:
         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())
-            pending = {recv, get_beat, get_result}
+            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(
@@ -397,7 +562,7 @@ class RemoteExecutionClient:
                 if recv in done:
                     msg = str(recv.result())
                     exit_info = await self._handle_websocket_message(
-                        msg, websocket, requests
+                        msg, websocket, requests, terminal_session_manager
                     )
                     if exit_info is not None:
                         return exit_info
@@ -416,12 +581,24 @@ class RemoteExecutionClient:
 
                 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()
@@ -432,21 +609,27 @@ class RemoteExecutionClient:
         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)
 
-        beats: asyncio.Queue[HeartbeatInfo] = asyncio.Queue()
-        requests: asyncio.Queue[CliRpcRequest] = asyncio.Queue()
-        results: asyncio.Queue[CliRpcResponse] = asyncio.Queue()
-
-        tasks = await self._setup_tasks(beats, requests, results)
+        self._websocket = websocket
 
         try:
             return await self._process_websocket_messages(
-                websocket, beats, requests, results
+                websocket,
+                beats,
+                requests,
+                results,
+                terminal_output_queue,
+                terminal_session_manager,
             )
         except websockets.exceptions.ConnectionClosed as e:
             if e.rcvd is not None:
@@ -461,15 +644,13 @@ class RemoteExecutionClient:
                     )
                     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:
-            for task in tasks:
-                task.cancel()
-            await asyncio.gather(*tasks, return_exceptions=True)
             if connection_tracker is not None:
                 await connection_tracker.set_connected(False)
 
@@ -485,33 +666,64 @@ class RemoteExecutionClient:
         # 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}"):
-            # 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,
-            )
+        # 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()
 
-            # Cancel pending tasks
-            for task in pending:
-                task.cancel()
+        # Create terminal session manager ONCE - persist across reconnections
+        terminal_session_manager = TerminalSessionManager(terminal_output_queue)
 
-            # 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
+        # Create tasks ONCE - persist across reconnections
+        executors = await self._setup_tasks(beats, requests, results)
 
-        # 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")
+        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(
@@ -520,13 +732,6 @@ class RemoteExecutionClient:
         )
         return await deserialize_api_response(response, CreateChatResponse)
 
-    async def get_gh_installation_token(self, git_info: GitInfo) -> dict[str, Any]:
-        response = await self.api_client.post(
-            "/github_app/exchange_token",
-            json=git_info.model_dump(),
-        )
-        return cast(dict[str, Any], response.json())
-
     # deprecated
     async def run_workflow(self, chat_uuid: str, workflow_id: str) -> dict[str, Any]:
         response = await self.api_client.post(
@@ -544,7 +749,7 @@ class RemoteExecutionClient:
         return cast(dict[str, Any], response.json())
 
     async def trigger_workflow(
-        self, workflow_name: str, workflow_input: PrReviewWorkflowInput
+        self, workflow_name: str, workflow_input: WorkflowInput
     ) -> WorkflowTriggerResponse:
         response = await self.api_client.post(
             "/api/remote_execution/trigger_workflow",
@@ -579,6 +784,38 @@ class RemoteExecutionClient:
         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()
@@ -593,7 +830,7 @@ class RemoteExecutionClient:
                     )
                 else:
                     raw_result = await execute_tool(  # type: ignore[assignment]
-                        request.request.tool_input, self.working_directory
+                        request.request.tool_input, self.working_directory, self
                     )
                 tool_result = truncate_result(raw_result)
                 return CliRpcResponse(
@@ -609,33 +846,38 @@ class RemoteExecutionClient:
                     response=GetAllFilesResponse(files=files),
                 )
             elif isinstance(request.request, BatchToolExecutionRequest):
-                results: list[ToolResultType] = []
+                coros: list[Coroutine[Any, Any, ToolResultType]] = []
                 for tool_input in request.request.tool_inputs:
-                    try:
-                        if isinstance(tool_input, BashToolInput):
-                            raw_result = await execute_bash_tool(
+                    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:
-                            raw_result = await execute_tool(  # type: ignore[assignment]
-                                tool_input, self.working_directory
-                            )
-                        tool_result = truncate_result(raw_result)
-                        results.append(tool_result)
-                    except Exception as e:
-                        logger.error(f"Error executing tool {tool_input}: {e}")
-                        from exponent.core.remote_execution.cli_rpc_types import (
-                            ErrorToolResult,
+                        )
+                    else:
+                        coros.append(
+                            execute_tool(tool_input, self.working_directory, self)
                         )
 
-                        results.append(ErrorToolResult(error_message=str(e)))
+                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=results,
+                        tool_results=processed_results,
                     ),
                 )
             elif isinstance(request.request, HttpRequest):
@@ -657,6 +899,22 @@ class RemoteExecutionClient:
                 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)}")
 
@@ -678,8 +936,13 @@ class RemoteExecutionClient:
                     await self.clear_halt_state(request.request_id)
 
     async def handle_streaming_request(
-        self, request: StreamingCodeExecutionRequest
-    ) -> AsyncGenerator[RemoteExecutionResponseType, None]:
+        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(