indent 0.1.21__py3-none-any.whl → 0.1.23__py3-none-any.whl

exponent/__init__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.21'
-__version_tuple__ = version_tuple = (0, 1, 21)
+__version__ = version = '0.1.23'
+__version_tuple__ = version_tuple = (0, 1, 23)
 
 __commit_id__ = commit_id = None

exponent/commands/cloud_commands.py

@@ -111,6 +111,7 @@ async def create_cloud_chat_from_repository(
     base_api_url: str,
     base_ws_url: str,
     repository_id: str,
+    provider: str | None = None,
 ) -> dict[str, Any]:
     graphql_client = GraphQLClient(
         api_key=api_key, base_api_url=base_api_url, base_ws_url=base_ws_url
@@ -118,6 +119,7 @@ async def create_cloud_chat_from_repository(
 
     variables = {
         "repositoryId": repository_id,
+        "provider": provider,
     }
 
     result = await graphql_client.execute(

exponent/core/graphql/mutations.py

@@ -76,8 +76,8 @@ mutation CreateCloudChat($configId: String!) {
 
 
 CREATE_CLOUD_CHAT_FROM_REPOSITORY_MUTATION = """
-mutation CreateCloudChatFromRepository($repositoryId: String!) {
-  createCloudChat(repositoryId: $repositoryId) {
+mutation CreateCloudChatFromRepository($repositoryId: String!, $provider: String) {
+  createCloudChat(repositoryId: $repositoryId, provider: $provider) {
     __typename
     ...on Chat {
       chatUuid

exponent/core/remote_execution/cli_rpc_types.py

@@ -304,6 +304,77 @@ class GenerateUploadUrlResponse(msgspec.Struct, tag="generate_upload_url"):
     s3_uri: str
 
 
+# Terminal session management
+class StartTerminalRequest(msgspec.Struct, tag="start_terminal"):
+    """Start a new terminal session with PTY"""
+
+    session_id: str
+    command: list[str] | None = None  # None = default shell, or specific command
+    cols: int = 80
+    rows: int = 24
+    env: dict[str, str] | None = None  # Additional environment variables
+
+
+class StartTerminalResponse(msgspec.Struct, tag="start_terminal"):
+    """Response after starting terminal"""
+
+    session_id: str
+    success: bool
+    error_message: str | None = None
+
+
+# Terminal input (user typing)
+class TerminalInputRequest(msgspec.Struct, tag="terminal_input"):
+    """Send user input to terminal"""
+
+    session_id: str
+    data: str  # Raw input data from user
+
+
+class TerminalInputResponse(msgspec.Struct, tag="terminal_input"):
+    """Acknowledge input received"""
+
+    session_id: str
+    success: bool
+
+
+# Terminal resize
+class TerminalResizeRequest(msgspec.Struct, tag="terminal_resize"):
+    """Resize terminal dimensions"""
+
+    session_id: str
+    cols: int
+    rows: int
+
+
+class TerminalResizeResponse(msgspec.Struct, tag="terminal_resize"):
+    """Acknowledge resize"""
+
+    session_id: str
+    success: bool
+
+
+# Terminal stop
+class StopTerminalRequest(msgspec.Struct, tag="stop_terminal"):
+    """Stop a terminal session"""
+
+    session_id: str
+
+
+class StopTerminalResponse(msgspec.Struct, tag="stop_terminal"):
+    """Acknowledge terminal stopped"""
+
+    session_id: str
+    success: bool
+
+
+class StreamingCodeExecutionRequest(msgspec.Struct, tag="streaming_code_execution"):
+    correlation_id: str
+    language: str  # "python" or "shell"
+    content: str
+    timeout: int
+
+
 class CliRpcRequest(msgspec.Struct):
     request_id: str
     request: (
@@ -315,6 +386,11 @@ class CliRpcRequest(msgspec.Struct):
         | SwitchCLIChatRequest
         | KeepAliveCliChatRequest
         | GenerateUploadUrlRequest
+        | StartTerminalRequest
+        | TerminalInputRequest
+        | TerminalResizeRequest
+        | StopTerminalRequest
+        | StreamingCodeExecutionRequest
     )
 
 
@@ -326,6 +402,40 @@ class ErrorResponse(msgspec.Struct, tag="error"):
     error_message: str
 
 
+class StreamingCodeExecutionResponseChunk(
+    msgspec.Struct, tag="streaming_code_execution_chunk"
+):
+    correlation_id: str
+    content: str
+    truncated: bool = False
+
+    def add(
+        self, new_chunk: "StreamingCodeExecutionResponseChunk"
+    ) -> "StreamingCodeExecutionResponseChunk":
+        """Aggregates content of this and a new chunk."""
+        assert self.correlation_id == new_chunk.correlation_id
+        return StreamingCodeExecutionResponseChunk(
+            correlation_id=self.correlation_id,
+            content=self.content + new_chunk.content,
+            truncated=self.truncated or new_chunk.truncated,
+        )
+
+
+class StreamingCodeExecutionResponse(msgspec.Struct, tag="streaming_code_execution"):
+    correlation_id: str
+    content: str
+    truncated: bool = False
+    # Only present for shell code execution
+    cancelled_for_timeout: bool = False
+    exit_code: int | None = None
+    halted: bool = False
+
+
+class StreamingErrorResponse(msgspec.Struct, tag="streaming_error"):
+    correlation_id: str
+    error_message: str
+
+
 class CliRpcResponse(msgspec.Struct):
     request_id: str
     response: (
@@ -339,4 +449,11 @@ class CliRpcResponse(msgspec.Struct):
         | SwitchCLIChatResponse
         | KeepAliveCliChatResponse
         | GenerateUploadUrlResponse
+        | StartTerminalResponse
+        | TerminalInputResponse
+        | TerminalResizeResponse
+        | StopTerminalResponse
+        | StreamingCodeExecutionResponseChunk
+        | StreamingCodeExecutionResponse
+        | StreamingErrorResponse
     )

exponent/core/remote_execution/client.py

@@ -38,8 +38,16 @@ from exponent.core.remote_execution.cli_rpc_types import (
     HttpRequest,
     KeepAliveCliChatRequest,
     KeepAliveCliChatResponse,
+    StartTerminalRequest,
+    StartTerminalResponse,
+    StopTerminalRequest,
+    StopTerminalResponse,
     SwitchCLIChatRequest,
     SwitchCLIChatResponse,
+    TerminalInputRequest,
+    TerminalInputResponse,
+    TerminalResizeRequest,
+    TerminalResizeResponse,
     TerminateRequest,
     TerminateResponse,
     ToolExecutionRequest,
@@ -56,6 +64,8 @@ from exponent.core.remote_execution.session import (
     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,
@@ -66,9 +76,7 @@ from exponent.core.remote_execution.types import (
     CLIConnectedState,
     CreateChatResponse,
     HeartbeatInfo,
-    RemoteExecutionResponseType,
     RunWorkflowRequest,
-    StreamingCodeExecutionRequest,
     WorkflowInput,
     WorkflowTriggerRequest,
     WorkflowTriggerResponse,
@@ -183,11 +191,12 @@ 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."""
@@ -215,7 +224,7 @@ class RemoteExecutionClient:
         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:
@@ -317,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
@@ -363,21 +467,31 @@ 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:
@@ -385,15 +499,32 @@ class RemoteExecutionClient:
                     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
 
@@ -411,6 +542,8 @@ 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()
@@ -418,7 +551,8 @@ class RemoteExecutionClient:
             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(
@@ -428,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
@@ -447,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()
@@ -466,6 +612,8 @@ 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 | None:
         """Handle a single websocket connection.
         Returns None to continue with reconnection attempts, or an exit info to terminate."""
@@ -476,7 +624,12 @@ class RemoteExecutionClient:
 
         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:
@@ -517,6 +670,10 @@ class RemoteExecutionClient:
         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)
@@ -529,7 +686,13 @@ class RemoteExecutionClient:
                     [
                         asyncio.create_task(
                             self._handle_websocket_connection(
-                                websocket, connection_tracker, beats, requests, results
+                                websocket,
+                                connection_tracker,
+                                beats,
+                                requests,
+                                results,
+                                terminal_output_queue,
+                                terminal_session_manager,
                             )
                         ),
                         asyncio.create_task(self._timeout_monitor(timeout_seconds)),
@@ -554,6 +717,9 @@ class RemoteExecutionClient:
                 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()
@@ -733,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)}")
 
@@ -754,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(

exponent/core/remote_execution/code_execution.py

@@ -1,5 +1,10 @@
 from collections.abc import AsyncGenerator, Callable
 
+from exponent.core.remote_execution.cli_rpc_types import (
+    StreamingCodeExecutionRequest,
+    StreamingCodeExecutionResponse,
+    StreamingCodeExecutionResponseChunk,
+)
 from exponent.core.remote_execution.languages.python_execution import (
     execute_python_streaming,
 )
@@ -8,13 +13,9 @@ from exponent.core.remote_execution.languages.shell_streaming import (
 )
 from exponent.core.remote_execution.languages.types import StreamedOutputPiece
 from exponent.core.remote_execution.session import RemoteExecutionClientSession
-from exponent.core.remote_execution.types import (
-    StreamingCodeExecutionRequest,
-    StreamingCodeExecutionResponse,
-    StreamingCodeExecutionResponseChunk,
-)
 
 EMPTY_OUTPUT_STRING = "(No output)"
+MAX_OUTPUT_LENGTH = 50000  # Maximum characters to keep in final output
 
 
 async def execute_code_streaming(
@@ -34,9 +35,18 @@ async def execute_code_streaming(
                     content=output.content, correlation_id=request.correlation_id
                 )
             else:
+                final_output = output.output or EMPTY_OUTPUT_STRING
+                truncated = len(final_output) > MAX_OUTPUT_LENGTH
+                if truncated:
+                    final_output = final_output[-MAX_OUTPUT_LENGTH:]
+                    final_output = (
+                        f"[Truncated to last {MAX_OUTPUT_LENGTH} characters]\n\n"
+                        + final_output
+                    )
                 yield StreamingCodeExecutionResponse(
                     correlation_id=request.correlation_id,
-                    content=output.output or EMPTY_OUTPUT_STRING,
+                    content=final_output,
+                    truncated=truncated,
                     halted=output.halted,
                 )
 
@@ -49,9 +59,18 @@ async def execute_code_streaming(
                     content=shell_output.content, correlation_id=request.correlation_id
                 )
             else:
+                final_output = shell_output.output or EMPTY_OUTPUT_STRING
+                truncated = len(final_output) > MAX_OUTPUT_LENGTH
+                if truncated:
+                    final_output = final_output[-MAX_OUTPUT_LENGTH:]
+                    final_output = (
+                        f"[Truncated to last {MAX_OUTPUT_LENGTH} characters]\n\n"
+                        + final_output
+                    )
                 yield StreamingCodeExecutionResponse(
                     correlation_id=request.correlation_id,
-                    content=shell_output.output or EMPTY_OUTPUT_STRING,
+                    content=final_output,
+                    truncated=truncated,
                     halted=shell_output.halted,
                     exit_code=shell_output.exit_code,
                     cancelled_for_timeout=shell_output.cancelled_for_timeout,

exponent/core/remote_execution/languages/shell_streaming.py

@@ -50,7 +50,7 @@ async def read_stream(
 
     while True:
         try:
-            data = await stream.read(4096)
+            data = await stream.read(50_000)
             if not data:
                 break
             chunk = data.decode(encoding=encoding, errors="replace")
@@ -182,10 +182,8 @@ async def execute_shell_streaming(  # noqa: PLR0915
                     piece = await task
                     yield piece
 
-                    if process.returncode is not None:
-                        break
-
                     # Schedule next read from the same stream
+                    # Don't check process.returncode here - we need to drain all buffered output
                     if task is stdout_task and not process.stdout.at_eof():
                         stdout_task = asyncio.create_task(stdout_gen.__anext__())
                         pending.add(stdout_task)