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

exponent/commands/run_commands.py

@@ -61,12 +61,19 @@ def run_cli() -> None:
     hidden=True,
     required=False,
 )
+@click.option(
+    "--timeout-seconds",
+    type=int,
+    help="Number of seconds without receiving a request before shutting down",
+    envvar="INDENT_TIMEOUT_SECONDS",
+)
 @use_settings
 def run(
     settings: Settings,
     chat_id: str | None = None,
     prompt: str | None = None,
     workflow_id: str | None = None,
+    timeout_seconds: int | None = None,
 ) -> None:
     """[default] Start or reconnect to an Indent session."""
     check_exponent_version_and_upgrade(settings)
@@ -104,7 +111,9 @@ def run(
         launch_exponent_browser(settings.environment, base_url, chat_uuid)
 
     while True:
-        result = run_chat(loop, api_key, chat_uuid, settings, prompt, workflow_id)
+        result = run_chat(
+            loop, api_key, chat_uuid, settings, prompt, workflow_id, timeout_seconds
+        )
         if result is None or isinstance(result, WSDisconnected):
             # NOTE: None here means that handle_connection_changes exited
             # first. We should likely have a different message for this.
@@ -112,11 +121,11 @@ def run(
                 click.secho(f"Error: {result.error_message}", fg="red")
                 sys.exit(10)
             else:
-                print("Disconnected upon user request, shutting down...")
+                click.echo("Disconnected upon user request, shutting down...")
                 break
         elif isinstance(result, SwitchCLIChat):
             chat_uuid = result.new_chat_uuid
-            print("\nSwitching chats...")
+            click.echo("\nSwitching chats...")
         else:
             assert_unreachable(result)
 
@@ -128,6 +137,7 @@ def run_chat(
     settings: Settings,
     prompt: str | None,
     workflow_id: str | None,
+    timeout_seconds: int | None,
 ) -> REMOTE_EXECUTION_CLIENT_EXIT_INFO | None:
     start_ts = time.time()
     base_url = settings.base_url
@@ -135,7 +145,7 @@ def run_chat(
     base_ws_url = settings.get_base_ws_url()
 
     print_exponent_message(base_url, chat_uuid)
-    print()
+    click.echo()
 
     connection_tracker = ConnectionTracker()
 
@@ -149,6 +159,7 @@ def run_chat(
             prompt=prompt,
             workflow_id=workflow_id,
             connection_tracker=connection_tracker,
+            timeout_seconds=timeout_seconds,
         )
     )
 
@@ -179,27 +190,27 @@ async def handle_connection_changes(
     try:
         async with timeout(5):
             assert await connection_tracker.next_change()
-            print(ready_message(start_ts))
+            click.echo(ready_message(start_ts))
     except TimeoutError:
         spinner = Spinner("Connecting...")
         spinner.show()
         assert await connection_tracker.next_change()
         spinner.hide()
-        print(ready_message(start_ts))
+        click.echo(ready_message(start_ts))
 
     while True:
         assert not await connection_tracker.next_change()
 
-        print("Disconnected...", end="")
+        click.echo("Disconnected...", nl=False)
         await asyncio.sleep(1)
         spinner = Spinner("Reconnecting...")
         spinner.show()
         assert await connection_tracker.next_change()
         spinner.hide()
-        print("\x1b[1;32m✓ Reconnected", end="")
+        click.echo("\x1b[1;32m✓ Reconnected", nl=False)
         sys.stdout.flush()
         await asyncio.sleep(1)
-        print("\r\x1b[0m\x1b[2K", end="")
+        click.echo("\r\x1b[0m\x1b[2K", nl=False)
         sys.stdout.flush()
 
 

exponent/commands/workflow_commands.py

@@ -47,7 +47,7 @@ def trigger(settings: Settings, workflow_type: str) -> None:
 
     while True:
         result = run_chat(
-            loop, settings.api_key, response.chat_uuid, settings, None, None
+            loop, settings.api_key, response.chat_uuid, settings, None, None, None
         )
         if result is None or isinstance(result, WSDisconnected):
             # NOTE: None here means that handle_connection_changes exited

exponent/core/graphql/client.py

@@ -1,7 +1,7 @@
 from collections.abc import AsyncGenerator
 from typing import Any
 
-from gql import Client, gql
+from gql import Client, GraphQLRequest, gql
 from gql.transport.httpx import HTTPXAsyncTransport
 from gql.transport.websockets import WebsocketsTransport
 
@@ -41,8 +41,10 @@ class GraphQLClient:
             execute_timeout=timeout,
         ) as session:
             # Execute single query
-            query = gql(query_str)
-            result = await session.execute(query, vars, op_name)
+            query = GraphQLRequest(
+                query_str, variable_values=vars, operation_name=op_name
+            )
+            result = await session.execute(query)
             return result
 
     async def subscribe(

exponent/core/graphql/mutations.py

@@ -73,3 +73,117 @@ mutation CreateCloudChat($configId: String!) {
   }
 }
 """
+
+
+CREATE_CLOUD_CHAT_FROM_REPOSITORY_MUTATION = """
+mutation CreateCloudChatFromRepository($repositoryId: String!) {
+  createCloudChat(repositoryId: $repositoryId) {
+    __typename
+    ...on Chat {
+      chatUuid
+    }
+    ...on UnauthenticatedError {
+      message
+    }
+    ...on ChatNotFoundError {
+      message
+    }
+    ...on CloudConfigNotFoundError {
+      message
+    }
+    ...on GithubConfigNotFoundError {
+      message
+    }
+    ...on CloudSessionError {
+      message
+    }
+  }
+}
+"""
+
+
+ENABLE_CLOUD_REPOSITORY_MUTATION = """
+mutation EnableCloudRepository($orgName: String!, $repoName: String!) {
+  enableCloudRepository(orgName: $orgName, repoName: $repoName) {
+    __typename
+    ...on ContainerImage {
+      buildRef
+      createdAt
+      updatedAt
+    }
+    ...on UnauthenticatedError {
+      message
+    }
+    ...on CloudConfigNotFoundError {
+      message
+    }
+    ...on GithubConfigNotFoundError {
+      message
+    }
+    ...on CloudSessionError {
+      message
+    }
+    ...on Error {
+      message
+    }
+  }
+}
+"""
+
+
+INCREMENTAL_BUILD_CLOUD_REPOSITORY_MUTATION = """
+mutation IncrementalBuildCloudRepository($orgName: String!, $repoName: String!) {
+  incrementalBuildCloudRepository(orgName: $orgName, repoName: $repoName) {
+    __typename
+    ...on ContainerImage {
+      buildRef
+      createdAt
+      updatedAt
+    }
+    ...on UnauthenticatedError {
+      message
+    }
+    ...on CloudConfigNotFoundError {
+      message
+    }
+    ...on GithubConfigNotFoundError {
+      message
+    }
+    ...on CloudSessionError {
+      message
+    }
+    ...on Error {
+      message
+    }
+  }
+}
+"""
+
+
+REBUILD_CLOUD_REPOSITORY_MUTATION = """
+mutation RebuildCloudRepository($orgName: String!, $repoName: String!) {
+  rebuildCloudRepository(orgName: $orgName, repoName: $repoName) {
+    __typename
+    ...on ContainerImage {
+      buildRef
+      createdAt
+      updatedAt
+    }
+    ...on UnauthenticatedError {
+      message
+    }
+    ...on CloudConfigNotFoundError {
+      message
+    }
+    ...on GithubConfigNotFoundError {
+      message
+    }
+    ...on CloudSessionError {
+      message
+    }
+    ...on Error {
+      message
+    }
+  }
+}
+"""

exponent/core/graphql/queries.py

@@ -1,3 +1,26 @@
+GITHUB_REPOSITORIES_QUERY: str = """
+query GithubRepositories {
+  githubRepositories {
+    __typename
+    ... on GithubRepositories {
+      repositories {
+        id
+        githubOrgName
+        githubRepoName
+        baseHost
+        containerImageId
+        createdAt
+        updatedAt
+      }
+    }
+    ... on Error {
+      message
+    }
+  }
+}
+"""
+
+
 EVENTS_FOR_CHAT_QUERY: str = """query EventsForChat($chatUuid: UUID!) {
   eventsForChat(chatUuid: $chatUuid) {
     ... on EventHistory {

exponent/core/remote_execution/cli_rpc_types.py

@@ -27,6 +27,11 @@ class ToolInput(msgspec.Struct, tag_field="tool_name", omit_defaults=True):
         """
         return msgspec.to_builtins(self)  # type: ignore[no-any-return]
 
+    def to_text(self) -> str:
+        d = msgspec.to_builtins(self)
+        del d["tool_name"]
+        return yaml.dump(d)
+
 
 class ToolResult(msgspec.Struct, tag_field="tool_name", omit_defaults=True):
     """Concrete subclasses return data from a tool execution."""
@@ -236,6 +241,12 @@ class SwitchCLIChatRequest(msgspec.Struct, tag="switch_cli_chat"):
     new_chat_uuid: str
 
 
+# This message is sent periodically from the client to keep the connection alive while the chat is turning.
+# This synergizes with CLI-side timeouts to avoid disconnecting during long operations.
+class KeepAliveCliChatRequest(msgspec.Struct, tag="keep_alive_cli_chat"):
+    pass
+
+
 class BatchToolExecutionRequest(msgspec.Struct, tag="batch_tool_execution"):
     tool_inputs: list[ToolInputType]
 
@@ -256,6 +267,10 @@ class SwitchCLIChatResponse(msgspec.Struct, tag="switch_cli_chat"):
     pass
 
 
+class KeepAliveCliChatResponse(msgspec.Struct, tag="keep_alive_cli_chat"):
+    pass
+
+
 class CliRpcRequest(msgspec.Struct):
     request_id: str
     request: (
@@ -265,6 +280,7 @@ class CliRpcRequest(msgspec.Struct):
         | HttpRequest
         | BatchToolExecutionRequest
         | SwitchCLIChatRequest
+        | KeepAliveCliChatRequest
     )
 
 
@@ -286,4 +302,5 @@ class CliRpcResponse(msgspec.Struct):
         | BatchToolExecutionResponse
         | HttpResponse
         | SwitchCLIChatResponse
+        | KeepAliveCliChatResponse
     )

exponent/core/remote_execution/client.py

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