indent 0.1.26__py3-none-any.whl

exponent/core/remote_execution/cli_rpc_types.py

@@ -0,0 +1,499 @@
+from typing import TYPE_CHECKING, Any
+
+import msgspec
+import yaml
+from pydantic_ai.format_as_xml import format_as_xml
+
+if TYPE_CHECKING:
+    from exponent_server.core.tools.edit_tool import (
+        EditToolInput as ServerSideEditToolInput,
+    )
+
+
+class PartialToolResult(msgspec.Struct, tag_field="tool_name", omit_defaults=True):
+    pass
+
+
+class ToolInput(msgspec.Struct, tag_field="tool_name", omit_defaults=True):
+    """Concrete subclasses describe the full input schema for a tool."""
+
+    def to_llm(self) -> dict[str, Any]:
+        """Convert ToolInput to LLM-friendly typed dict format.
+
+        Returns a dictionary with the tool parameters, excluding the tool_name
+        which is handled separately by the LLM integration layer.
+
+        Returns:
+            self by default, which msgspec will serialize appropriately.
+        """
+        return msgspec.to_builtins(self)  # type: ignore[no-any-return]
+
+    def to_text(self) -> str:
+        """This text is purely used for human debugging, it's not fed to the llm"""
+        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."""
+
+    def to_text(self) -> str:
+        """
+        This provides a default textual representation of the tool result. Override it as needed for your tool."""
+        d = msgspec.to_builtins(self)
+        del d["tool_name"]
+        return format_as_xml(d, include_root_tag=False, item_tag="item")
+
+
+class ErrorToolResult(ToolResult, tag="error"):
+    error_message: str
+    is_assistant_error: bool = False
+
+
+READ_TOOL_NAME = "read"
+READ_TOOL_ARTIFACT_NAME = "read_tool_artifact"
+
+
+class ReadToolInput(ToolInput, tag=READ_TOOL_NAME):
+    file_path: str
+    offset: int | None = None
+    limit: int | None = None
+
+
+class FileMetadata(msgspec.Struct):
+    modified_timestamp: float
+    file_mode: str
+
+
+class ReadToolArtifactResult(ToolResult, tag=READ_TOOL_ARTIFACT_NAME):
+    s3_uri: str
+    file_path: str
+    media_type: str
+
+    def to_text(self) -> str:
+        return f"[Image artifact uploaded to {self.s3_uri}]"
+
+
+class ReadToolResult(ToolResult, tag=READ_TOOL_NAME):
+    content: str | None = None
+    num_lines: int | None = None
+    start_line: int | None = None
+    total_lines: int | None = None
+    metadata: FileMetadata | None = None
+    artifact: ReadToolArtifactResult | None = None
+
+    def to_text(self) -> str:
+        if self.artifact:
+            return self.artifact.to_text()
+        assert self.content is not None
+        assert self.start_line is not None
+        lines = self.content.splitlines()
+        lines = [
+            f"{str(i).rjust(6)}→{line}"
+            for i, line in enumerate(lines, start=self.start_line + 1)
+        ]
+        return "\n".join(lines)
+
+
+LIST_TOOL_NAME = "ls"
+
+
+class ListToolInput(ToolInput, tag=LIST_TOOL_NAME):
+    path: str
+    ignore: list[str] | None = None
+
+
+class ListToolResult(ToolResult, tag=LIST_TOOL_NAME):
+    files: list[str]
+    truncated: bool = False
+
+
+GLOB_TOOL_NAME = "glob"
+
+
+class GlobToolInput(ToolInput, tag=GLOB_TOOL_NAME):
+    pattern: str
+    path: str | None = None
+
+
+class GlobToolResult(ToolResult, tag=GLOB_TOOL_NAME):
+    filenames: list[str]
+    duration_ms: int
+    num_files: int
+    truncated: bool
+
+
+WRITE_TOOL_NAME = "write"
+
+
+class WriteToolInput(ToolInput, tag=WRITE_TOOL_NAME):
+    file_path: str
+    content: str
+
+
+class WriteToolResult(ToolResult, tag=WRITE_TOOL_NAME):
+    message: str
+    metadata: FileMetadata | None = None
+
+
+GREP_TOOL_NAME = "grep"
+
+
+class GrepToolInput(ToolInput, tag=GREP_TOOL_NAME):
+    pattern: str
+    path: str | None = None
+    include: str | None = None
+    multiline: bool | None = None
+
+
+class GrepToolResult(ToolResult, tag=GREP_TOOL_NAME):
+    matches: list[str]
+    truncated: bool = False
+
+
+EDIT_TOOL_NAME = "edit"
+
+
+# This is only used in the CLI. The server side type is edit_tool.py
+class EditToolInput(ToolInput, tag=EDIT_TOOL_NAME):
+    file_path: str
+    old_string: str
+    new_string: str
+    replace_all: bool = False
+    # The last retrieved modified timestamp. This is used to check if the file has been modified since the last read/write that we are aware of (in which case we should re-read the file before editing).
+    last_known_modified_timestamp: float | None = None
+
+    @classmethod
+    def from_server_side_input(
+        cls,
+        server_side_input: "ServerSideEditToolInput",
+        last_known_modified_timestamp: float | None,
+    ) -> "EditToolInput":
+        return cls(
+            file_path=server_side_input.file_path,
+            old_string=server_side_input.old_string,
+            new_string=server_side_input.new_string,
+            replace_all=server_side_input.replace_all,
+            last_known_modified_timestamp=last_known_modified_timestamp,
+        )
+
+
+class EditToolResult(ToolResult, tag=EDIT_TOOL_NAME):
+    message: str
+    metadata: FileMetadata | None = None
+
+
+BASH_TOOL_NAME = "bash"
+
+
+class BashToolInput(ToolInput, tag=BASH_TOOL_NAME):
+    command: str
+    timeout: int | None = None
+    description: str | None = None
+
+
+class PartialBashToolResult(PartialToolResult, tag=BASH_TOOL_NAME):
+    shell_output: str | None = None
+
+
+class BashToolResult(ToolResult, tag=BASH_TOOL_NAME):
+    shell_output: str
+    duration_ms: int
+    exit_code: int | None
+    timed_out: bool
+    stopped_by_user: bool
+
+
+DOWNLOAD_ARTIFACT_TOOL_NAME = "download_artifact"
+
+
+class DownloadArtifactToolInput(ToolInput, tag=DOWNLOAD_ARTIFACT_TOOL_NAME):
+    presigned_url: str
+    file_path: str
+    artifact_id: str
+    overwrite: bool = True
+
+
+class DownloadArtifactToolResult(ToolResult, tag=DOWNLOAD_ARTIFACT_TOOL_NAME):
+    file_path: str
+    artifact_id: str
+    file_size_bytes: int
+    content_preview: str | None = None
+    num_lines: int | None = None
+    total_lines: int | None = None
+    truncated: bool = False
+
+
+UPLOAD_ARTIFACT_TOOL_NAME = "upload_artifact"
+
+
+class UploadArtifactToolInput(ToolInput, tag=UPLOAD_ARTIFACT_TOOL_NAME):
+    presigned_url: str
+    file_path: str
+    artifact_id: str
+    content_type: str
+
+
+class UploadArtifactToolResult(ToolResult, tag=UPLOAD_ARTIFACT_TOOL_NAME):
+    artifact_id: str
+    file_size_bytes: int
+    content_type: str
+
+
+class HttpRequest(msgspec.Struct, tag="http_fetch_cli"):
+    url: str
+    method: str = "GET"
+    headers: dict[str, str] | None = None
+    timeout: int | None = None
+
+
+class HttpResponse(msgspec.Struct, tag="http_fetch_cli"):
+    status_code: int | None = None
+    content: str | None = None
+    error_message: str | None = None
+    duration_ms: int | None = None
+    headers: dict[str, str] | None = None
+
+
+ToolInputType = (
+    ReadToolInput
+    | WriteToolInput
+    | ListToolInput
+    | GlobToolInput
+    | GrepToolInput
+    | EditToolInput
+    | BashToolInput
+    | DownloadArtifactToolInput
+    | UploadArtifactToolInput
+)
+PartialToolResultType = PartialBashToolResult
+
+ToolResultType = (
+    ReadToolResult
+    | ReadToolArtifactResult
+    | WriteToolResult
+    | ListToolResult
+    | GlobToolResult
+    | GrepToolResult
+    | EditToolResult
+    | BashToolResult
+    | DownloadArtifactToolResult
+    | UploadArtifactToolResult
+    | ErrorToolResult
+)
+
+
+class ToolExecutionRequest(msgspec.Struct, tag="tool_execution"):
+    tool_input: ToolInputType
+
+
+class GetAllFilesRequest(msgspec.Struct, tag="get_all_files"):
+    pass
+
+
+class TerminateRequest(msgspec.Struct, tag="terminate"):
+    pass
+
+
+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]
+
+
+class GetAllFilesResponse(msgspec.Struct, tag="get_all_files"):
+    files: list[str]
+
+
+class TerminateResponse(msgspec.Struct, tag="terminate"):
+    pass
+
+
+class TimeoutResponse(msgspec.Struct, tag="timeout"):
+    pass
+
+
+class BatchToolExecutionResponse(msgspec.Struct, tag="batch_tool_execution"):
+    tool_results: list[ToolResultType]
+
+
+class SwitchCLIChatResponse(msgspec.Struct, tag="switch_cli_chat"):
+    pass
+
+
+class KeepAliveCliChatResponse(msgspec.Struct, tag="keep_alive_cli_chat"):
+    pass
+
+
+class GenerateUploadUrlRequest(msgspec.Struct, tag="generate_upload_url"):
+    s3_key: str
+    content_type: str
+
+
+class GenerateUploadUrlResponse(msgspec.Struct, tag="generate_upload_url"):
+    upload_url: str
+    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: (
+        ToolExecutionRequest
+        | GetAllFilesRequest
+        | TerminateRequest
+        | HttpRequest
+        | BatchToolExecutionRequest
+        | SwitchCLIChatRequest
+        | KeepAliveCliChatRequest
+        | GenerateUploadUrlRequest
+        | StartTerminalRequest
+        | TerminalInputRequest
+        | TerminalResizeRequest
+        | StopTerminalRequest
+        | StreamingCodeExecutionRequest
+    )
+
+
+class ToolExecutionResponse(msgspec.Struct, tag="tool_execution"):
+    tool_result: ToolResultType
+
+
+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: (
+        ToolExecutionResponse
+        | GetAllFilesResponse
+        | ErrorResponse
+        | TerminateResponse
+        | TimeoutResponse
+        | BatchToolExecutionResponse
+        | HttpResponse
+        | SwitchCLIChatResponse
+        | KeepAliveCliChatResponse
+        | GenerateUploadUrlResponse
+        | StartTerminalResponse
+        | TerminalInputResponse
+        | TerminalResizeResponse
+        | StopTerminalResponse
+        | StreamingCodeExecutionResponseChunk
+        | StreamingCodeExecutionResponse
+        | StreamingErrorResponse
+    )