Initial public release of Warp.
Repo-Sync-Origin: warpdotdev/warp-internal@12af1d983b
This commit is contained in:
@@ -0,0 +1,37 @@
|
||||
[package]
|
||||
name = "remote_server"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
authors = ["Warp Team <dev@warp.dev>"]
|
||||
publish.workspace = true
|
||||
license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
anyhow.workspace = true
|
||||
async-channel.workspace = true
|
||||
command.workspace = true
|
||||
dashmap.workspace = true
|
||||
futures.workspace = true
|
||||
futures-lite.workspace = true
|
||||
log.workspace = true
|
||||
prost.workspace = true
|
||||
repo_metadata.workspace = true
|
||||
serde.workspace = true
|
||||
thiserror.workspace = true
|
||||
uuid.workspace = true
|
||||
warp_core.workspace = true
|
||||
warp_util.workspace = true
|
||||
warpui.workspace = true
|
||||
|
||||
[target.'cfg(not(target_family = "wasm"))'.dependencies]
|
||||
async-process.workspace = true
|
||||
|
||||
[target.'cfg(target_family = "wasm")'.dependencies]
|
||||
getrandom = { version = "0.2", features = ["js"] }
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { workspace = true, features = ["io-util", "macros", "rt-multi-thread"] }
|
||||
tokio-util = { workspace = true, features = ["compat"] }
|
||||
|
||||
[build-dependencies]
|
||||
prost-build.workspace = true
|
||||
@@ -0,0 +1,5 @@
|
||||
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
println!("cargo:rerun-if-changed=proto/remote_server.proto");
|
||||
prost_build::compile_protos(&["proto/remote_server.proto"], &["proto/"])?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,297 @@
|
||||
// A protobuf-based API for communication between
|
||||
// the Warp client and the remote server binary.
|
||||
//
|
||||
// Messages are length-prefixed: [4-byte little-endian length][protobuf bytes].
|
||||
|
||||
syntax = "proto3";
|
||||
|
||||
package remote_server;
|
||||
|
||||
// ── Top-level envelopes ───────────────────────────────────────────
|
||||
|
||||
// Top-level envelope for all client → server messages.
|
||||
message ClientMessage {
|
||||
string request_id = 1;
|
||||
oneof message {
|
||||
Initialize initialize = 2;
|
||||
Abort abort = 3;
|
||||
NavigatedToDirectory navigated_to_directory = 4;
|
||||
LoadRepoMetadataDirectory load_repo_metadata_directory = 5;
|
||||
SessionBootstrapped session_bootstrapped = 6;
|
||||
WriteFile write_file = 7;
|
||||
DeleteFile delete_file = 8;
|
||||
RunCommandRequest run_command = 9;
|
||||
ReadFileContextRequest read_file_context = 10;
|
||||
}
|
||||
}
|
||||
// Top-level envelope for all server → client messages.
|
||||
// Push messages use an empty request_id to distinguish them from
|
||||
// request/response pairs.
|
||||
message ServerMessage {
|
||||
string request_id = 1;
|
||||
oneof message {
|
||||
InitializeResponse initialize_response = 2;
|
||||
ErrorResponse error = 3;
|
||||
NavigatedToDirectoryResponse navigated_to_directory_response = 4;
|
||||
RepoMetadataSnapshot repo_metadata_snapshot = 5;
|
||||
RepoMetadataUpdatePush repo_metadata_update = 6;
|
||||
LoadRepoMetadataDirectoryResponse load_repo_metadata_directory_response = 7;
|
||||
WriteFileResponse write_file_response = 8;
|
||||
DeleteFileResponse delete_file_response = 9;
|
||||
RunCommandResponse run_command_response = 10;
|
||||
ReadFileContextResponse read_file_context_response = 11;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Initialize handshake
|
||||
|
||||
// Sent by the client immediately after connecting to negotiate the protocol.
|
||||
message Initialize {}
|
||||
|
||||
// Sent by the client to cancel an in-progress request.
|
||||
// This is a notification (fire-and-forget) — the server does not send a response.
|
||||
message Abort {
|
||||
string request_id_to_abort = 1;
|
||||
}
|
||||
|
||||
// Returned by the server in response to Initialize.
|
||||
message InitializeResponse {
|
||||
string server_version = 1;
|
||||
string host_id = 2;
|
||||
}
|
||||
|
||||
// Sent by the client after the SSH session has been bootstrapped.
|
||||
// This is a notification (fire-and-forget) — the server does not send a
|
||||
// response. The server uses this to create a per-session
|
||||
// LocalCommandExecutor matching the bootstrapped shell.
|
||||
message SessionBootstrapped {
|
||||
uint64 session_id = 1;
|
||||
string shell_type = 2;
|
||||
// The full path to the shell binary (e.g. "/usr/bin/zsh").
|
||||
// When present, the server uses this directly instead of doing a PATH lookup.
|
||||
optional string shell_path = 3;
|
||||
}
|
||||
|
||||
// Sent by the client to execute a shell command on the remote host.
|
||||
message RunCommandRequest {
|
||||
string command = 1;
|
||||
// Working directory for the command. If empty, uses the server's default.
|
||||
optional string working_directory = 2;
|
||||
// Environment variables to set for the command.
|
||||
map<string, string> environment_variables = 3;
|
||||
// The session whose shell executor should run this command.
|
||||
uint64 session_id = 4;
|
||||
}
|
||||
|
||||
// Specific error codes for RunCommandRequest failures.
|
||||
enum RunCommandErrorCode {
|
||||
RUN_COMMAND_ERROR_CODE_UNSPECIFIED = 0;
|
||||
// The session ID has no associated executor (session was never bootstrapped).
|
||||
SESSION_NOT_FOUND = 1;
|
||||
// The command could not be executed (e.g. the shell process failed to spawn).
|
||||
EXECUTION_FAILED = 2;
|
||||
}
|
||||
|
||||
message RunCommandError {
|
||||
RunCommandErrorCode code = 1;
|
||||
string message = 2;
|
||||
}
|
||||
|
||||
message RunCommandSuccess {
|
||||
bytes stdout = 1;
|
||||
bytes stderr = 2;
|
||||
// Absent when the process was killed by a signal (Unix).
|
||||
optional int32 exit_code = 3;
|
||||
}
|
||||
|
||||
// Returned by the server with the result of a RunCommandRequest.
|
||||
message RunCommandResponse {
|
||||
oneof result {
|
||||
RunCommandSuccess success = 1;
|
||||
RunCommandError error = 2;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Shared error response ─────────────────────────────────────────
|
||||
|
||||
enum ErrorCode {
|
||||
ERROR_CODE_UNSPECIFIED = 0;
|
||||
// The request was malformed (e.g. missing oneof variant).
|
||||
INVALID_REQUEST = 1;
|
||||
// An unexpected server-side failure.
|
||||
INTERNAL = 2;
|
||||
}
|
||||
|
||||
message ErrorResponse {
|
||||
ErrorCode code = 1;
|
||||
string message = 2;
|
||||
}
|
||||
|
||||
// ── Shared repo metadata sub-messages ─────────────────────────────
|
||||
// Mirror the Rust types in repo_metadata/src/file_tree_update.rs.
|
||||
|
||||
// Mirrors RepoNodeMetadata in Rust.
|
||||
message RepoNodeMetadata {
|
||||
oneof node {
|
||||
DirectoryNodeMetadata directory = 1;
|
||||
FileNodeMetadata file = 2;
|
||||
}
|
||||
}
|
||||
|
||||
// Mirrors DirectoryNodeMetadata in Rust.
|
||||
message DirectoryNodeMetadata {
|
||||
string path = 1;
|
||||
bool ignored = 2;
|
||||
bool loaded = 3;
|
||||
}
|
||||
|
||||
// Mirrors FileNodeMetadata in Rust.
|
||||
message FileNodeMetadata {
|
||||
string path = 1;
|
||||
optional string extension = 2;
|
||||
bool ignored = 3;
|
||||
}
|
||||
|
||||
// Mirrors FileTreeEntryUpdate in Rust. Describes a subtree patch rooted
|
||||
// at a parent directory.
|
||||
message RepoMetadataEntryUpdate {
|
||||
string parent_path_to_replace = 1;
|
||||
repeated RepoNodeMetadata subtree_metadata = 2;
|
||||
}
|
||||
|
||||
// ── NavigatedToDirectory ──────────────────────────────────────────
|
||||
|
||||
// Client → server: "I navigated to this directory, please index it."
|
||||
message NavigatedToDirectory {
|
||||
string path = 1;
|
||||
}
|
||||
|
||||
// Response after the server has run git detection on the requested path.
|
||||
// When is_git is true, indexed_path is the git repo root and full indexing
|
||||
// runs in the background. A RepoMetadataSnapshot push will follow.
|
||||
// When is_git is false, the directory was lazily indexed at first level.
|
||||
// A RepoMetadataSnapshot push with the lazy tree data will follow.
|
||||
message NavigatedToDirectoryResponse {
|
||||
string indexed_path = 1;
|
||||
bool is_git = 2;
|
||||
}
|
||||
|
||||
// ── LoadRepoMetadataDirectory ─────────────────────────────────────
|
||||
|
||||
// Client → server: load the next level of a subdirectory within an
|
||||
// already-tracked repo (lazy expand).
|
||||
message LoadRepoMetadataDirectory {
|
||||
string repo_path = 1;
|
||||
string dir_path = 2;
|
||||
}
|
||||
|
||||
// Response with the loaded subtree entries for the requested directory.
|
||||
message LoadRepoMetadataDirectoryResponse {
|
||||
string repo_path = 1;
|
||||
string dir_path = 2;
|
||||
repeated RepoMetadataEntryUpdate entries = 3;
|
||||
}
|
||||
|
||||
// ── File write/delete operations ──────────────────────────────────
|
||||
|
||||
// Shared error type for file operations (read/write/delete).
|
||||
message FileOperationError {
|
||||
string message = 1;
|
||||
}
|
||||
|
||||
// Client → server: write content to a file, creating parent dirs if needed.
|
||||
message WriteFile {
|
||||
string path = 1;
|
||||
string content = 2;
|
||||
}
|
||||
|
||||
// Server → client: result of a WriteFile request.
|
||||
message WriteFileResponse {
|
||||
oneof result {
|
||||
WriteFileSuccess success = 1;
|
||||
FileOperationError error = 2;
|
||||
}
|
||||
}
|
||||
|
||||
message WriteFileSuccess {}
|
||||
|
||||
// Client → server: delete a file.
|
||||
message DeleteFile {
|
||||
string path = 1;
|
||||
}
|
||||
|
||||
// Server → client: result of a DeleteFile request.
|
||||
message DeleteFileResponse {
|
||||
oneof result {
|
||||
DeleteFileSuccess success = 1;
|
||||
FileOperationError error = 2;
|
||||
}
|
||||
}
|
||||
|
||||
message DeleteFileSuccess {}
|
||||
|
||||
// ── Read file context (batch) ─────────────────────────────────────
|
||||
|
||||
// A single file to read, with optional line ranges.
|
||||
message ReadFileContextFile {
|
||||
string path = 1;
|
||||
// 1-indexed line ranges (start..end). Empty = read entire file.
|
||||
repeated LineRange line_ranges = 2;
|
||||
}
|
||||
|
||||
message LineRange {
|
||||
uint32 start = 1;
|
||||
uint32 end = 2;
|
||||
}
|
||||
|
||||
// Client → server: batch read multiple files with full context.
|
||||
message ReadFileContextRequest {
|
||||
repeated ReadFileContextFile files = 1;
|
||||
// Per-file byte limit. Absent = use server default.
|
||||
optional uint32 max_file_bytes = 2;
|
||||
// Cumulative byte budget across all files. Absent = no batch limit.
|
||||
optional uint32 max_batch_bytes = 3;
|
||||
}
|
||||
|
||||
// Server → client: result of a ReadFileContextRequest.
|
||||
// Per-file failures are reported in `failed_files`, not as a top-level error.
|
||||
// Catastrophic server errors (malformed request, etc.) use the generic ErrorResponse.
|
||||
message ReadFileContextResponse {
|
||||
repeated FileContextProto file_contexts = 1;
|
||||
repeated FailedFileRead failed_files = 2;
|
||||
}
|
||||
|
||||
message FailedFileRead {
|
||||
string path = 1;
|
||||
FileOperationError error = 2;
|
||||
}
|
||||
|
||||
message FileContextProto {
|
||||
string file_name = 1;
|
||||
oneof content {
|
||||
string text_content = 2;
|
||||
bytes binary_content = 3;
|
||||
}
|
||||
// Optional 1-indexed line range this segment covers.
|
||||
optional uint32 line_range_start = 4;
|
||||
optional uint32 line_range_end = 5;
|
||||
optional uint64 last_modified_epoch_millis = 6;
|
||||
uint32 line_count = 7;
|
||||
}
|
||||
|
||||
// ── Server → client push messages ─────────────────────────────────
|
||||
|
||||
// Full or lazy-loaded repo metadata snapshot. Pushed by the server after
|
||||
// NavigatedToDirectory completes indexing (either lazy or full git).
|
||||
message RepoMetadataSnapshot {
|
||||
string repo_path = 1;
|
||||
repeated RepoMetadataEntryUpdate entries = 2;
|
||||
bool sync_complete = 3;
|
||||
}
|
||||
|
||||
// Incremental repo metadata update. Mirrors RepoMetadataUpdate in Rust.
|
||||
message RepoMetadataUpdatePush {
|
||||
string repo_path = 1;
|
||||
repeated string remove_entries = 2;
|
||||
repeated RepoMetadataEntryUpdate update_entries = 3;
|
||||
}
|
||||
@@ -0,0 +1,602 @@
|
||||
use std::collections::HashMap;
|
||||
use std::fmt;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use dashmap::DashMap;
|
||||
use futures::channel::oneshot;
|
||||
use futures::io::{AsyncRead, AsyncWrite};
|
||||
use warpui::r#async::{executor, FutureExt as _};
|
||||
|
||||
use crate::proto::{
|
||||
client_message, server_message, Abort, ClientMessage, DeleteFile, ErrorCode, Initialize,
|
||||
InitializeResponse, LoadRepoMetadataDirectoryResponse, NavigatedToDirectoryResponse,
|
||||
ReadFileContextRequest, ReadFileContextResponse, RunCommandRequest, RunCommandResponse,
|
||||
ServerMessage, SessionBootstrapped, WriteFile,
|
||||
};
|
||||
|
||||
use crate::protocol::{self, ProtocolError, RequestId};
|
||||
|
||||
use warp_core::SessionId;
|
||||
use warpui::r#async::TransportStream;
|
||||
|
||||
/// Default request timeout (2 minutes).
|
||||
const REQUEST_TIMEOUT: Duration = Duration::from_secs(120);
|
||||
|
||||
/// Errors from the `RemoteServerClient`.
|
||||
#[derive(thiserror::Error, Debug)]
|
||||
pub enum ClientError {
|
||||
#[error("Connection was dropped")]
|
||||
Disconnected,
|
||||
|
||||
#[error("Protocol error: {0}")]
|
||||
Protocol(#[from] ProtocolError),
|
||||
|
||||
#[error("Response channel closed before receiving a reply")]
|
||||
ResponseChannelClosed,
|
||||
|
||||
#[error("Unexpected response from server")]
|
||||
UnexpectedResponse,
|
||||
|
||||
#[error("Server error ({code:?}): {message}")]
|
||||
ServerError { code: ErrorCode, message: String },
|
||||
|
||||
#[error("Request timed out after {0:?}")]
|
||||
Timeout(Duration),
|
||||
|
||||
#[error("File operation failed: {0}")]
|
||||
FileOperationFailed(String),
|
||||
}
|
||||
|
||||
/// Events received from the remote server, delivered through the event
|
||||
/// channel returned by [`RemoteServerClient::new`].
|
||||
///
|
||||
/// The consumer (typically `RemoteServerManager`) drains this channel to
|
||||
/// react to connection lifecycle changes and server-pushed data.
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum ClientEvent {
|
||||
/// The reader task detected EOF or a fatal error. The connection is gone.
|
||||
/// This is always the last event sent on the channel.
|
||||
Disconnected,
|
||||
/// A full or lazy-loaded repo metadata snapshot was pushed by the server.
|
||||
RepoMetadataSnapshotReceived {
|
||||
update: repo_metadata::RepoMetadataUpdate,
|
||||
},
|
||||
/// An incremental repo metadata update was pushed by the server.
|
||||
RepoMetadataUpdated {
|
||||
update: repo_metadata::RepoMetadataUpdate,
|
||||
},
|
||||
/// A server message could not be decoded and had no parseable request_id.
|
||||
MessageDecodingError,
|
||||
}
|
||||
/// Client for communicating with a `remote_server` process over the remote server protocol.
|
||||
///
|
||||
/// Exposes async request/response APIs over generic I/O streams (child-process pipes,
|
||||
/// SSH channels, or in-memory streams for testing).
|
||||
///
|
||||
/// Designed to be wrapped in `Arc` for sharing across threads. Construction
|
||||
/// returns an event receiver that delivers push events and a final
|
||||
/// `Disconnected` event when the connection drops.
|
||||
///
|
||||
/// This type does **not** own the child subprocess whose stdio backs it.
|
||||
/// For transports that spawn a subprocess (e.g. SSH), the caller is
|
||||
/// responsible for holding the `Child` for the lifetime of the session
|
||||
/// so that `kill_on_drop` fires when teardown occurs. In Warp this is
|
||||
/// the `RemoteServerManager`, which stores the child in
|
||||
/// `RemoteSessionState` alongside the `Arc<RemoteServerClient>`. That
|
||||
/// way the child's lifetime is gated by the manager's session map
|
||||
/// rather than by `Arc` refcount -- cloning `Arc<RemoteServerClient>`
|
||||
/// into other owners (e.g. the command executor) no longer keeps the
|
||||
/// child alive.
|
||||
pub struct RemoteServerClient {
|
||||
/// Channel for queuing ClientMessages to send to the remote server.
|
||||
outbound_tx: async_channel::Sender<ClientMessage>,
|
||||
|
||||
/// Maps `request_id` → oneshot sender for the correlated response from the remote server.
|
||||
pending_requests: Arc<DashMap<RequestId, oneshot::Sender<Result<ServerMessage, ClientError>>>>,
|
||||
|
||||
/// Set to `true` by the reader task when the connection is lost. Checked by
|
||||
/// `send_request` after inserting into `pending_requests` to avoid hanging
|
||||
/// on a dead connection.
|
||||
disconnected: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl fmt::Debug for RemoteServerClient {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("RemoteServerClient").finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
impl RemoteServerClient {
|
||||
/// Creates a client from a child process's stdin, stdout, and stderr.
|
||||
///
|
||||
/// The caller retains ownership of the `Child` itself. Typically the
|
||||
/// caller spawns the `Command` with `kill_on_drop(true)` and stashes
|
||||
/// the returned `Child` somewhere whose lifetime matches the
|
||||
/// session's (in Warp, on the `RemoteServerManager`'s
|
||||
/// `RemoteSessionState`). Dropping the `Child` there triggers
|
||||
/// SIGKILL on the subprocess, regardless of how many
|
||||
/// `Arc<RemoteServerClient>` clones are still alive.
|
||||
///
|
||||
/// Internally forwards stderr lines to local logging via
|
||||
/// [`spawn_stderr_forwarder`], then delegates to [`Self::new`] for the
|
||||
/// protocol reader/writer setup.
|
||||
///
|
||||
/// Returns the client and an event receiver that delivers push events
|
||||
/// and a final `Disconnected` event when the connection drops.
|
||||
pub fn from_child_streams(
|
||||
stdin: async_process::ChildStdin,
|
||||
stdout: async_process::ChildStdout,
|
||||
stderr: async_process::ChildStderr,
|
||||
executor: &executor::Background,
|
||||
) -> (Self, async_channel::Receiver<ClientEvent>) {
|
||||
spawn_stderr_forwarder(stderr, executor);
|
||||
Self::new(stdout, stdin, executor)
|
||||
}
|
||||
}
|
||||
|
||||
impl RemoteServerClient {
|
||||
/// Creates a new client, spawning background reader and writer tasks on the
|
||||
/// provided executor.
|
||||
///
|
||||
/// Returns the client and an event receiver that delivers push events
|
||||
/// and a final `Disconnected` event when the connection drops.
|
||||
pub fn new(
|
||||
reader: impl AsyncRead + TransportStream,
|
||||
writer: impl AsyncWrite + TransportStream,
|
||||
executor: &executor::Background,
|
||||
) -> (Self, async_channel::Receiver<ClientEvent>) {
|
||||
let pending_requests: Arc<
|
||||
DashMap<RequestId, oneshot::Sender<Result<ServerMessage, ClientError>>>,
|
||||
> = Arc::new(DashMap::new());
|
||||
let (outbound_tx, outbound_rx) = async_channel::unbounded::<ClientMessage>();
|
||||
let (event_tx, event_rx) = async_channel::unbounded::<ClientEvent>();
|
||||
let disconnected = Arc::new(AtomicBool::new(false));
|
||||
|
||||
executor
|
||||
.spawn(Self::writer_task(
|
||||
writer,
|
||||
outbound_rx,
|
||||
Arc::clone(&pending_requests),
|
||||
))
|
||||
.detach();
|
||||
executor
|
||||
.spawn(Self::reader_task(
|
||||
reader,
|
||||
Arc::clone(&pending_requests),
|
||||
event_tx,
|
||||
Arc::clone(&disconnected),
|
||||
))
|
||||
.detach();
|
||||
|
||||
(
|
||||
Self {
|
||||
outbound_tx,
|
||||
pending_requests,
|
||||
disconnected,
|
||||
},
|
||||
event_rx,
|
||||
)
|
||||
}
|
||||
|
||||
/// Sends an `Initialize` request and awaits the `InitializeResponse`.
|
||||
pub async fn initialize(&self) -> Result<InitializeResponse, ClientError> {
|
||||
let request_id = RequestId::new();
|
||||
let msg = ClientMessage {
|
||||
request_id: request_id.to_string(),
|
||||
message: Some(client_message::Message::Initialize(Initialize {})),
|
||||
};
|
||||
|
||||
let response = self.send_request(request_id, msg).await?;
|
||||
|
||||
match response.message {
|
||||
Some(server_message::Message::InitializeResponse(resp)) => Ok(resp),
|
||||
other => {
|
||||
log::error!("Unexpected response variant for Initialize: {other:?}");
|
||||
Err(ClientError::UnexpectedResponse)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Sends a `SessionBootstrapped` notification (fire-and-forget) so the
|
||||
/// server can create a `LocalCommandExecutor` for the session.
|
||||
pub fn notify_session_bootstrapped(
|
||||
&self,
|
||||
session_id: SessionId,
|
||||
shell_type: &str,
|
||||
shell_path: Option<&str>,
|
||||
) {
|
||||
let msg = ClientMessage {
|
||||
request_id: String::new(),
|
||||
message: Some(client_message::Message::SessionBootstrapped(
|
||||
SessionBootstrapped {
|
||||
session_id: session_id.as_u64(),
|
||||
shell_type: shell_type.to_owned(),
|
||||
shell_path: shell_path.map(ToOwned::to_owned),
|
||||
},
|
||||
)),
|
||||
};
|
||||
self.send_notification(msg);
|
||||
}
|
||||
|
||||
/// Sends a `NavigatedToDirectory` request and awaits the response.
|
||||
pub async fn navigate_to_directory(
|
||||
&self,
|
||||
path: String,
|
||||
) -> Result<NavigatedToDirectoryResponse, ClientError> {
|
||||
let request_id = RequestId::new();
|
||||
let msg = ClientMessage {
|
||||
request_id: request_id.to_string(),
|
||||
message: Some(client_message::Message::NavigatedToDirectory(
|
||||
crate::proto::NavigatedToDirectory { path },
|
||||
)),
|
||||
};
|
||||
|
||||
let response = self.send_request(request_id, msg).await?;
|
||||
|
||||
match response.message {
|
||||
Some(server_message::Message::NavigatedToDirectoryResponse(resp)) => Ok(resp),
|
||||
other => {
|
||||
log::error!("Unexpected response variant for NavigatedToDirectory: {other:?}");
|
||||
Err(ClientError::UnexpectedResponse)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Sends a `LoadRepoMetadataDirectory` request and awaits the response.
|
||||
pub async fn load_repo_metadata_directory(
|
||||
&self,
|
||||
repo_path: String,
|
||||
dir_path: String,
|
||||
) -> Result<LoadRepoMetadataDirectoryResponse, ClientError> {
|
||||
let request_id = RequestId::new();
|
||||
let msg = ClientMessage {
|
||||
request_id: request_id.to_string(),
|
||||
message: Some(client_message::Message::LoadRepoMetadataDirectory(
|
||||
crate::proto::LoadRepoMetadataDirectory {
|
||||
repo_path,
|
||||
dir_path,
|
||||
},
|
||||
)),
|
||||
};
|
||||
|
||||
let response = self.send_request(request_id, msg).await?;
|
||||
|
||||
match response.message {
|
||||
Some(server_message::Message::LoadRepoMetadataDirectoryResponse(resp)) => Ok(resp),
|
||||
other => {
|
||||
log::error!("Unexpected response variant for LoadRepoMetadataDirectory: {other:?}");
|
||||
Err(ClientError::UnexpectedResponse)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Writes content to a file on the remote host.
|
||||
/// Creates parent directories if they don't exist.
|
||||
pub async fn write_file(&self, path: String, content: String) -> Result<(), ClientError> {
|
||||
let request_id = RequestId::new();
|
||||
let msg = ClientMessage {
|
||||
request_id: request_id.to_string(),
|
||||
message: Some(client_message::Message::WriteFile(WriteFile {
|
||||
path,
|
||||
content,
|
||||
})),
|
||||
};
|
||||
let response = self.send_request(request_id, msg).await?;
|
||||
match response.message {
|
||||
Some(server_message::Message::WriteFileResponse(resp)) => match resp.result {
|
||||
Some(crate::proto::write_file_response::Result::Success(_)) | None => Ok(()),
|
||||
Some(crate::proto::write_file_response::Result::Error(e)) => {
|
||||
Err(ClientError::FileOperationFailed(e.message))
|
||||
}
|
||||
},
|
||||
other => {
|
||||
log::error!("Unexpected response variant for WriteFile: {other:?}");
|
||||
Err(ClientError::UnexpectedResponse)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Batch-reads one or more files from the remote host with full context
|
||||
/// (line ranges, binary/image support, metadata, size limits).
|
||||
///
|
||||
/// Per-file failures are reported in `ReadFileContextResponse::failed_files`
|
||||
/// rather than as a top-level error. The method only returns `Err` for
|
||||
/// transport-level failures (disconnect, timeout, etc.).
|
||||
pub async fn read_file_context(
|
||||
&self,
|
||||
request: ReadFileContextRequest,
|
||||
) -> Result<ReadFileContextResponse, ClientError> {
|
||||
let request_id = RequestId::new();
|
||||
let msg = ClientMessage {
|
||||
request_id: request_id.to_string(),
|
||||
message: Some(client_message::Message::ReadFileContext(request)),
|
||||
};
|
||||
let response = self.send_request(request_id, msg).await?;
|
||||
match response.message {
|
||||
Some(server_message::Message::ReadFileContextResponse(resp)) => Ok(resp),
|
||||
other => {
|
||||
log::error!("Unexpected response variant for ReadFileContext: {other:?}");
|
||||
Err(ClientError::UnexpectedResponse)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Deletes a file on the remote host.
|
||||
pub async fn delete_file(&self, path: String) -> Result<(), ClientError> {
|
||||
let request_id = RequestId::new();
|
||||
let msg = ClientMessage {
|
||||
request_id: request_id.to_string(),
|
||||
message: Some(client_message::Message::DeleteFile(DeleteFile { path })),
|
||||
};
|
||||
let response = self.send_request(request_id, msg).await?;
|
||||
match response.message {
|
||||
Some(server_message::Message::DeleteFileResponse(resp)) => match resp.result {
|
||||
Some(crate::proto::delete_file_response::Result::Success(_)) | None => Ok(()),
|
||||
Some(crate::proto::delete_file_response::Result::Error(e)) => {
|
||||
Err(ClientError::FileOperationFailed(e.message))
|
||||
}
|
||||
},
|
||||
other => {
|
||||
log::error!("Unexpected response variant for DeleteFile: {other:?}");
|
||||
Err(ClientError::UnexpectedResponse)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Converts a server push message (empty request_id) into a domain event.
|
||||
fn push_message_to_event(msg: ServerMessage) -> Option<ClientEvent> {
|
||||
match msg.message? {
|
||||
server_message::Message::RepoMetadataSnapshot(snapshot) => {
|
||||
let update = crate::repo_metadata_proto::proto_snapshot_to_update(&snapshot)?;
|
||||
Some(ClientEvent::RepoMetadataSnapshotReceived { update })
|
||||
}
|
||||
server_message::Message::RepoMetadataUpdate(push) => {
|
||||
let update = crate::repo_metadata_proto::proto_to_repo_metadata_update(&push)?;
|
||||
Some(ClientEvent::RepoMetadataUpdated { update })
|
||||
}
|
||||
other => {
|
||||
log::warn!("Unhandled push message variant: {other:?}");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Sends a `RunCommand` request
|
||||
pub async fn run_command(
|
||||
&self,
|
||||
session_id: SessionId,
|
||||
command: String,
|
||||
working_directory: Option<String>,
|
||||
environment_variables: HashMap<String, String>,
|
||||
) -> Result<RunCommandResponse, ClientError> {
|
||||
let request_id = RequestId::new();
|
||||
let msg = ClientMessage {
|
||||
request_id: request_id.to_string(),
|
||||
message: Some(client_message::Message::RunCommand(RunCommandRequest {
|
||||
command,
|
||||
working_directory,
|
||||
environment_variables,
|
||||
session_id: session_id.as_u64(),
|
||||
})),
|
||||
};
|
||||
|
||||
let response = self.send_request(request_id, msg).await?;
|
||||
|
||||
match response.message {
|
||||
Some(server_message::Message::RunCommandResponse(resp)) => Ok(resp),
|
||||
other => {
|
||||
log::error!("Unexpected response variant for RunCommand: {other:?}");
|
||||
Err(ClientError::UnexpectedResponse)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Generic request/response correlation.
|
||||
///
|
||||
/// Registers a oneshot channel keyed by `request_id`, sends the message
|
||||
/// through the outbound channel, and awaits the correlated response.
|
||||
/// Times out after `REQUEST_TIMEOUT` and sends an `Abort` to the server.
|
||||
async fn send_request(
|
||||
&self,
|
||||
request_id: RequestId,
|
||||
msg: ClientMessage,
|
||||
) -> Result<ServerMessage, ClientError> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
self.pending_requests.insert(request_id.clone(), tx);
|
||||
|
||||
// Check if the reader task has already marked the connection as dead.
|
||||
// The DashMap lock from `insert` above synchronizes with the lock from
|
||||
// `clear` in `reader_task`, so if `clear` ran before our insert the
|
||||
// flag is guaranteed to be visible here.
|
||||
if self.disconnected.load(Ordering::Acquire) {
|
||||
self.pending_requests.clear();
|
||||
return Err(ClientError::Disconnected);
|
||||
}
|
||||
|
||||
if self.outbound_tx.send(msg).await.is_err() {
|
||||
self.pending_requests.remove(&request_id);
|
||||
return Err(ClientError::Disconnected);
|
||||
}
|
||||
|
||||
let result = match rx.with_timeout(REQUEST_TIMEOUT).await {
|
||||
Ok(Ok(inner)) => inner,
|
||||
Ok(Err(_)) => return Err(ClientError::ResponseChannelClosed),
|
||||
Err(_) => {
|
||||
// Timed out — clean up and send abort.
|
||||
self.pending_requests.remove(&request_id);
|
||||
self.send_abort(&request_id);
|
||||
return Err(ClientError::Timeout(REQUEST_TIMEOUT));
|
||||
}
|
||||
};
|
||||
|
||||
// Unwrap the inner Result (reader task may send Err for decode failures).
|
||||
let response = result?;
|
||||
|
||||
// Convert server-reported ErrorResponse into ClientError so callers
|
||||
// only need to match on success variants.
|
||||
if let Some(server_message::Message::Error(ref e)) = response.message {
|
||||
return Err(ClientError::ServerError {
|
||||
code: e.code(),
|
||||
message: e.message.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
/// Sends an `Abort` notification for the given request ID.
|
||||
fn send_abort(&self, request_id_to_abort: &RequestId) {
|
||||
let msg = ClientMessage {
|
||||
request_id: RequestId::new().to_string(),
|
||||
message: Some(client_message::Message::Abort(Abort {
|
||||
request_id_to_abort: request_id_to_abort.to_string(),
|
||||
})),
|
||||
};
|
||||
self.send_notification(msg);
|
||||
}
|
||||
|
||||
/// Sends a message without registering a pending request (fire-and-forget).
|
||||
fn send_notification(&self, msg: ClientMessage) {
|
||||
// Use try_send to avoid blocking; if the channel is full or closed,
|
||||
// the notification is best-effort.
|
||||
if let Err(e) = self.outbound_tx.try_send(msg) {
|
||||
log::debug!("Failed to send notification (best-effort): {e}");
|
||||
}
|
||||
}
|
||||
|
||||
/// Background task that writes `ClientMessage`s to the underlying stream.
|
||||
async fn writer_task(
|
||||
writer: impl AsyncWrite + TransportStream,
|
||||
outbound_rx: async_channel::Receiver<ClientMessage>,
|
||||
pending_requests: Arc<
|
||||
DashMap<RequestId, oneshot::Sender<Result<ServerMessage, ClientError>>>,
|
||||
>,
|
||||
) {
|
||||
let mut writer = futures::io::BufWriter::new(writer);
|
||||
while let Ok(msg) = outbound_rx.recv().await {
|
||||
if let Err(e) = protocol::write_client_message(&mut writer, &msg).await {
|
||||
let request_id = RequestId::from(msg.request_id);
|
||||
if !e.is_write_recoverable() {
|
||||
log::error!("Writer task fatal error: request_id={request_id}: {e}");
|
||||
pending_requests.clear();
|
||||
break;
|
||||
}
|
||||
log::error!("Writer task: request_id={request_id}: {e}");
|
||||
// Drop the sender so the caller receives ResponseChannelClosed.
|
||||
pending_requests.remove(&request_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Background task that reads `ServerMessage`s and resolves pending
|
||||
/// requests by `request_id`, or converts push messages to events.
|
||||
///
|
||||
/// Sends `ClientEvent::Disconnected` as the final event when the
|
||||
/// connection is lost.
|
||||
async fn reader_task(
|
||||
reader: impl AsyncRead + TransportStream,
|
||||
pending_requests: Arc<
|
||||
DashMap<RequestId, oneshot::Sender<Result<ServerMessage, ClientError>>>,
|
||||
>,
|
||||
event_tx: async_channel::Sender<ClientEvent>,
|
||||
disconnected: Arc<AtomicBool>,
|
||||
) {
|
||||
let mut reader = futures::io::BufReader::new(reader);
|
||||
loop {
|
||||
match protocol::read_server_message(&mut reader).await {
|
||||
Ok(msg) => {
|
||||
let request_id = RequestId::from(msg.request_id.clone());
|
||||
if request_id.is_empty() {
|
||||
// Push message — convert to a domain event and forward.
|
||||
if let Some(event) = Self::push_message_to_event(msg) {
|
||||
if event_tx.send(event).await.is_err() {
|
||||
log::warn!("Event channel closed, dropping push message");
|
||||
}
|
||||
}
|
||||
} else if let Some((_, tx)) = pending_requests.remove(&request_id) {
|
||||
// Ignore send failure — the caller may have dropped the receiver.
|
||||
let _ = tx.send(Ok(msg));
|
||||
} else {
|
||||
log::warn!("Received unexpected response with request_id={request_id}");
|
||||
}
|
||||
}
|
||||
Err(ProtocolError::Decode(ref err, Some(ref request_id))) => {
|
||||
if let Some((_, tx)) = pending_requests.remove(request_id) {
|
||||
log::warn!(
|
||||
"Reader task: malformed response \
|
||||
(request_id={request_id}): {err}"
|
||||
);
|
||||
let _ = tx.send(Err(ClientError::Protocol(ProtocolError::Decode(
|
||||
err.clone(),
|
||||
Some(request_id.clone()),
|
||||
))));
|
||||
} else {
|
||||
log::warn!(
|
||||
"Reader task: malformed response for \
|
||||
unknown request (request_id={request_id}): {err}"
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(ProtocolError::Decode(ref err, None)) => {
|
||||
log::warn!(
|
||||
"Reader task: skipping malformed response \
|
||||
(no parseable request_id): {err}"
|
||||
);
|
||||
let _ = event_tx.send(ClientEvent::MessageDecodingError).await;
|
||||
}
|
||||
Err(e) if e.is_read_recoverable() => {
|
||||
log::warn!("Reader task: skipping message: {e}");
|
||||
}
|
||||
Err(e) => {
|
||||
match e {
|
||||
ProtocolError::UnexpectedEof => {
|
||||
log::info!("Reader task: server disconnected (EOF)");
|
||||
}
|
||||
_ => log::error!("Reader task fatal error: {e}"),
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Mark the connection as dead so that any new `send_request` calls
|
||||
// fail immediately rather than hanging forever. This prevents a race
|
||||
// where `pending_requests.clear()` runs before `send_request` has
|
||||
// inserted its oneshot entry.
|
||||
disconnected.store(true, Ordering::Release);
|
||||
|
||||
// Notify all pending requests that the connection is gone.
|
||||
pending_requests.clear();
|
||||
|
||||
// Signal disconnection as the final event.
|
||||
let _ = event_tx.send(ClientEvent::Disconnected).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawns a background task that reads lines from the server's stderr and
|
||||
/// forwards them to the client's logging.
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
pub fn spawn_stderr_forwarder(
|
||||
stderr: impl AsyncRead + TransportStream,
|
||||
executor: &executor::Background,
|
||||
) {
|
||||
use futures::io::AsyncBufReadExt;
|
||||
use futures::StreamExt;
|
||||
|
||||
executor
|
||||
.spawn(async move {
|
||||
let reader = futures::io::BufReader::new(stderr);
|
||||
let mut lines = reader.lines();
|
||||
while let Some(Ok(line)) = lines.next().await {
|
||||
log::info!("[remote_server] {line}");
|
||||
}
|
||||
})
|
||||
.detach();
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "../client_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,250 @@
|
||||
use futures::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
|
||||
use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
|
||||
|
||||
use crate::proto::{
|
||||
client_message, run_command_response, server_message, ClientMessage, ErrorCode,
|
||||
InitializeResponse, RunCommandResponse, RunCommandSuccess, ServerMessage,
|
||||
};
|
||||
use crate::protocol;
|
||||
use warp_core::SessionId;
|
||||
use warpui::r#async::executor;
|
||||
|
||||
use super::*;
|
||||
|
||||
/// Generic mock server: loops reading ClientMessages and responds using the
|
||||
/// provided closure. Exits cleanly on EOF.
|
||||
async fn mock_server_with<F>(
|
||||
mut reader: impl AsyncRead + Unpin,
|
||||
mut writer: impl AsyncWrite + Unpin,
|
||||
responder: F,
|
||||
) where
|
||||
F: Fn(&ClientMessage) -> server_message::Message,
|
||||
{
|
||||
loop {
|
||||
match protocol::read_client_message(&mut reader).await {
|
||||
Ok(msg) => {
|
||||
let response = ServerMessage {
|
||||
request_id: msg.request_id.clone(),
|
||||
message: Some(responder(&msg)),
|
||||
};
|
||||
protocol::write_server_message(&mut writer, &response)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
Err(protocol::ProtocolError::UnexpectedEof) => break,
|
||||
Err(e) => panic!("mock server error: {e}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Sets up a duplex stream, spawns `mock_server_with` with the given responder,
|
||||
/// and returns a connected `RemoteServerClient`, its event receiver, and the
|
||||
/// background executor (which must be kept alive for the test duration).
|
||||
fn setup_mock_client<F>(
|
||||
responder: F,
|
||||
) -> (
|
||||
RemoteServerClient,
|
||||
async_channel::Receiver<ClientEvent>,
|
||||
executor::Background,
|
||||
)
|
||||
where
|
||||
F: Fn(&ClientMessage) -> server_message::Message + Send + 'static,
|
||||
{
|
||||
let (client_stream, server_stream) = tokio::io::duplex(4096);
|
||||
let (server_read, server_write) = tokio::io::split(server_stream);
|
||||
let (client_read, client_write) = tokio::io::split(client_stream);
|
||||
|
||||
tokio::spawn(mock_server_with(
|
||||
server_read.compat(),
|
||||
server_write.compat_write(),
|
||||
responder,
|
||||
));
|
||||
|
||||
let executor = executor::Background::default();
|
||||
let (client, event_rx) =
|
||||
RemoteServerClient::new(client_read.compat(), client_write.compat_write(), &executor);
|
||||
(client, event_rx, executor)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn initialize_round_trip() {
|
||||
let (client, _disconnect_rx, _executor) = setup_mock_client(|_| {
|
||||
server_message::Message::InitializeResponse(InitializeResponse {
|
||||
server_version: "test-0.1.0".to_string(),
|
||||
host_id: "test-host-id".to_string(),
|
||||
})
|
||||
});
|
||||
|
||||
let resp = client.initialize().await.unwrap();
|
||||
assert_eq!(resp.server_version, "test-0.1.0");
|
||||
assert_eq!(resp.host_id, "test-host-id");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disconnected_on_closed_stream() {
|
||||
let (client_stream, server_stream) = tokio::io::duplex(4096);
|
||||
// Drop the server side immediately.
|
||||
drop(server_stream);
|
||||
|
||||
let (client_read, client_write) = tokio::io::split(client_stream);
|
||||
let executor = executor::Background::default();
|
||||
let (client, disconnect_rx) =
|
||||
RemoteServerClient::new(client_read.compat(), client_write.compat_write(), &executor);
|
||||
|
||||
// An initialize call on a dead stream must complete with an error rather than hang.
|
||||
let result = client.initialize().await;
|
||||
assert!(result.is_err());
|
||||
|
||||
// The reader task should detect EOF and emit a Disconnected event.
|
||||
let event = disconnect_rx.recv().await.unwrap();
|
||||
assert!(matches!(event, ClientEvent::Disconnected));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn run_command_round_trip() {
|
||||
let (client, _disconnect_rx, _executor) = setup_mock_client(|msg| {
|
||||
let command = match &msg.message {
|
||||
Some(client_message::Message::RunCommand(req)) => req.command.clone(),
|
||||
other => panic!("Expected RunCommand, got {other:?}"),
|
||||
};
|
||||
server_message::Message::RunCommandResponse(RunCommandResponse {
|
||||
result: Some(run_command_response::Result::Success(RunCommandSuccess {
|
||||
stdout: format!("output of: {command}").into_bytes(),
|
||||
stderr: Vec::new(),
|
||||
exit_code: Some(0),
|
||||
})),
|
||||
})
|
||||
});
|
||||
|
||||
let resp = client
|
||||
.run_command(
|
||||
SessionId::from(42u64),
|
||||
"echo hello".to_string(),
|
||||
None,
|
||||
Default::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let success = match resp.result {
|
||||
Some(run_command_response::Result::Success(s)) => s,
|
||||
other => panic!("Expected RunCommandSuccess, got {other:?}"),
|
||||
};
|
||||
assert_eq!(success.stdout, b"output of: echo hello");
|
||||
assert!(success.stderr.is_empty());
|
||||
assert_eq!(success.exit_code, Some(0));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn concurrent_in_flight_requests() {
|
||||
let (client, _disconnect_rx, _executor) = setup_mock_client(|_| {
|
||||
server_message::Message::InitializeResponse(InitializeResponse {
|
||||
server_version: "test-0.1.0".to_string(),
|
||||
host_id: "test-host-id".to_string(),
|
||||
})
|
||||
});
|
||||
let client = std::sync::Arc::new(client);
|
||||
|
||||
let mut handles = Vec::new();
|
||||
for _ in 0..10 {
|
||||
let c = std::sync::Arc::clone(&client);
|
||||
handles.push(tokio::spawn(async move {
|
||||
c.initialize().await.expect("concurrent initialize failed")
|
||||
}));
|
||||
}
|
||||
|
||||
for h in handles {
|
||||
let resp = h.await.unwrap();
|
||||
assert_eq!(resp.server_version, "test-0.1.0");
|
||||
assert_eq!(resp.host_id, "test-host-id");
|
||||
}
|
||||
}
|
||||
|
||||
/// Simulates a server that reads raw bytes, sends an error response for
|
||||
/// malformed messages where the request_id is parseable, then continues
|
||||
/// processing valid messages.
|
||||
async fn mock_server_with_error_handling(
|
||||
mut reader: impl AsyncRead + Unpin,
|
||||
mut writer: impl AsyncWrite + Unpin,
|
||||
) {
|
||||
loop {
|
||||
match protocol::read_client_message(&mut reader).await {
|
||||
Ok(msg) => {
|
||||
let response = ServerMessage {
|
||||
request_id: msg.request_id,
|
||||
message: Some(server_message::Message::InitializeResponse(
|
||||
InitializeResponse {
|
||||
server_version: "test-0.1.0".to_string(),
|
||||
host_id: "test-host-id".to_string(),
|
||||
},
|
||||
)),
|
||||
};
|
||||
protocol::write_server_message(&mut writer, &response)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
Err(protocol::ProtocolError::Decode(_, Some(ref id))) => {
|
||||
let error_response = ServerMessage {
|
||||
request_id: id.to_string(),
|
||||
message: Some(server_message::Message::Error(
|
||||
crate::proto::ErrorResponse {
|
||||
code: ErrorCode::InvalidRequest.into(),
|
||||
message: "malformed message".to_string(),
|
||||
},
|
||||
)),
|
||||
};
|
||||
protocol::write_server_message(&mut writer, &error_response)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
Err(protocol::ProtocolError::Decode(_, None)) => {}
|
||||
Err(protocol::ProtocolError::UnexpectedEof) => break,
|
||||
Err(e) => panic!("mock server error: {e}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Sends a corrupted protobuf with a valid request_id to the server,
|
||||
/// verifying the server responds with an ErrorResponse for that request_id.
|
||||
#[tokio::test]
|
||||
async fn server_returns_error_for_malformed_message_with_parseable_id() {
|
||||
let (client_stream, server_stream) = tokio::io::duplex(4096);
|
||||
let (server_read, server_write) = tokio::io::split(server_stream);
|
||||
let (client_read, client_write) = tokio::io::split(client_stream);
|
||||
|
||||
tokio::spawn(mock_server_with_error_handling(
|
||||
server_read.compat(),
|
||||
server_write.compat_write(),
|
||||
));
|
||||
|
||||
// Manually construct a corrupted message with a valid request_id field
|
||||
// followed by bytes that cause a prost decode failure.
|
||||
let mut payload = Vec::new();
|
||||
// Field 1 (string): tag=0x0a, length=15, "malformed-req-1"
|
||||
payload.push(0x0a);
|
||||
payload.push(15);
|
||||
payload.extend_from_slice(b"malformed-req-1");
|
||||
// Invalid trailing bytes: field tag with reserved wire type 7 causes
|
||||
// prost to fail, but our try_extract_request_id stops after field 1.
|
||||
payload.extend_from_slice(&[0x0F, 0x01]); // field 1, wire type 7 (invalid)
|
||||
|
||||
// Write the corrupted message with length prefix.
|
||||
let mut client_write = client_write.compat_write();
|
||||
let len = payload.len() as u32;
|
||||
client_write.write_all(&len.to_le_bytes()).await.unwrap();
|
||||
client_write.write_all(&payload).await.unwrap();
|
||||
client_write.flush().await.unwrap();
|
||||
|
||||
// Read the error response from the server.
|
||||
let mut client_reader = futures::io::BufReader::new(client_read.compat());
|
||||
let response: ServerMessage = protocol::read_server_message(&mut client_reader)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.request_id, "malformed-req-1");
|
||||
match response.message {
|
||||
Some(server_message::Message::Error(e)) => {
|
||||
assert_eq!(e.code(), ErrorCode::InvalidRequest);
|
||||
}
|
||||
other => panic!("expected ErrorResponse, got: {other:?}"),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
// Re-export from warp_core so existing `remote_server::HostId` imports
|
||||
// continue to work.
|
||||
pub use warp_core::HostId;
|
||||
@@ -0,0 +1,39 @@
|
||||
#!/usr/bin/env bash
|
||||
# Installs the Warp remote server binary on a remote host.
|
||||
#
|
||||
# Placeholders (substituted at runtime by setup.rs):
|
||||
# {download_base_url} — e.g. https://app.warp.dev/download/cli
|
||||
# {channel} — stable | preview | dev
|
||||
# {install_dir} — e.g. ~/.warp/remote-server
|
||||
# {binary_name} — e.g. oz | oz-dev | oz-preview
|
||||
set -e
|
||||
|
||||
arch=$(uname -m)
|
||||
case "$arch" in
|
||||
x86_64) arch_name=x86_64 ;;
|
||||
aarch64|arm64) arch_name=aarch64 ;;
|
||||
*) echo "unsupported arch: $arch" >&2; exit 2 ;;
|
||||
esac
|
||||
|
||||
os_kernel=$(uname -s)
|
||||
case "$os_kernel" in
|
||||
Darwin) os_name=macos ;;
|
||||
Linux) os_name=linux ;;
|
||||
*) echo "unsupported OS: $os_kernel" >&2; exit 2 ;;
|
||||
esac
|
||||
|
||||
install_dir="{install_dir}"
|
||||
install_dir="${install_dir/#\~/"$HOME"}"
|
||||
mkdir -p "$install_dir"
|
||||
|
||||
tmpdir=$(mktemp -d "$install_dir/.install.XXXXXX")
|
||||
trap 'rm -rf "$tmpdir"' EXIT
|
||||
|
||||
curl -fSL "{download_base_url}?package=tar&os=$os_name&arch=$arch_name&channel={channel}" \
|
||||
-o "$tmpdir/oz.tar.gz"
|
||||
tar -xzf "$tmpdir/oz.tar.gz" -C "$tmpdir"
|
||||
|
||||
bin=$(find "$tmpdir" -type f -name 'oz*' ! -name '*.tar.gz' | head -n1)
|
||||
if [ -z "$bin" ]; then echo "no binary found in tarball" >&2; exit 1; fi
|
||||
chmod +x "$bin"
|
||||
mv "$bin" "$install_dir/{binary_name}"
|
||||
@@ -0,0 +1,15 @@
|
||||
pub mod client;
|
||||
pub mod host_id;
|
||||
pub mod manager;
|
||||
pub mod protocol;
|
||||
pub mod repo_metadata_proto;
|
||||
pub mod setup;
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
pub mod ssh;
|
||||
pub mod transport;
|
||||
|
||||
pub use host_id::HostId;
|
||||
|
||||
pub mod proto {
|
||||
include!(concat!(env!("OUT_DIR"), "/remote_server.rs"));
|
||||
}
|
||||
@@ -0,0 +1,966 @@
|
||||
use std::collections::{HashMap, HashSet};
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use crate::client::ClientEvent;
|
||||
use crate::client::RemoteServerClient;
|
||||
use crate::setup::RemotePlatform;
|
||||
use crate::setup::RemoteServerSetupState;
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use crate::transport::Connection;
|
||||
use crate::transport::RemoteTransport;
|
||||
use crate::HostId;
|
||||
use repo_metadata::RepoMetadataUpdate;
|
||||
use serde::Serialize;
|
||||
use warp_core::SessionId;
|
||||
use warpui::{Entity, ModelContext, ModelSpawner, SingletonEntity};
|
||||
|
||||
/// Which phase of the remote server connection flow failed.
|
||||
#[derive(Clone, Copy, Debug, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum RemoteServerInitPhase {
|
||||
/// `transport.connect()` failed (SSH/process spawn level).
|
||||
Connect,
|
||||
/// `client.initialize()` failed (protocol handshake level).
|
||||
Initialize,
|
||||
}
|
||||
|
||||
/// The remote server client operation that failed.
|
||||
#[derive(Clone, Copy, Debug, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum RemoteServerOperation {
|
||||
NavigateToDirectory,
|
||||
LoadRepoMetadataDirectory,
|
||||
}
|
||||
|
||||
/// Classification of a remote server client error for telemetry.
|
||||
#[derive(Clone, Copy, Debug, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum RemoteServerErrorKind {
|
||||
Timeout,
|
||||
Disconnected,
|
||||
ServerError,
|
||||
Other,
|
||||
}
|
||||
|
||||
impl RemoteServerErrorKind {
|
||||
/// Classify a [`ClientError`] into a telemetry error kind.
|
||||
pub fn from_client_error(error: &crate::client::ClientError) -> Self {
|
||||
use crate::client::ClientError;
|
||||
match error {
|
||||
ClientError::Timeout(_) => Self::Timeout,
|
||||
ClientError::Disconnected | ClientError::ResponseChannelClosed => Self::Disconnected,
|
||||
ClientError::ServerError { .. } => Self::ServerError,
|
||||
ClientError::Protocol(_)
|
||||
| ClientError::UnexpectedResponse
|
||||
| ClientError::FileOperationFailed(_) => Self::Other,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-session connection state. Encodes which data is available at each
|
||||
/// lifecycle stage so the compiler prevents invalid combinations.
|
||||
///
|
||||
/// For subprocess-backed transports (SSH), the `Initializing` and
|
||||
/// `Connected` variants also own the transport's `Child`. Dropping or
|
||||
/// replacing the state sends SIGKILL to the subprocess via
|
||||
/// `kill_on_drop`, which is the authoritative teardown path -- it fires
|
||||
/// on both explicit deregistration and spontaneous disconnect, and is
|
||||
/// unaffected by lingering `Arc<RemoteServerClient>` clones held
|
||||
/// elsewhere (e.g. the per-session command executor).
|
||||
///
|
||||
/// They also optionally carry a `control_path` pointing at the SSH
|
||||
/// `ControlMaster` socket for this session. On explicit teardown
|
||||
/// (after the user's shell exits), `deregister_session` uses this to
|
||||
/// run `ssh -O exit`, forcing the master to terminate without waiting
|
||||
/// for half-closed multiplexed channels to finish cleanup on the
|
||||
/// remote side.
|
||||
#[derive(Debug)]
|
||||
pub enum RemoteSessionState {
|
||||
/// `connect_session` has been called; background task is starting the
|
||||
/// server process over SSH.
|
||||
Connecting,
|
||||
/// Server process spawned, client exists, initialize handshake in progress.
|
||||
Initializing {
|
||||
client: Arc<RemoteServerClient>,
|
||||
/// The transport's owning `Child`. Dropped when the state is
|
||||
/// replaced or removed, killing the subprocess via
|
||||
/// `kill_on_drop`.
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
_child: async_process::Child,
|
||||
/// See type-level doc.
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
control_path: Option<PathBuf>,
|
||||
},
|
||||
/// Initialize handshake succeeded. Client is ready for requests.
|
||||
Connected {
|
||||
client: Arc<RemoteServerClient>,
|
||||
host_id: HostId,
|
||||
/// The transport's owning `Child`. See `Initializing::_child`.
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
_child: async_process::Child,
|
||||
/// See type-level doc.
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
control_path: Option<PathBuf>,
|
||||
},
|
||||
/// Connection dropped (EOF/error from the reader task).
|
||||
Disconnected,
|
||||
}
|
||||
|
||||
/// Events emitted by [`RemoteServerManager`].
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum RemoteServerManagerEvent {
|
||||
// --- Session-scoped events ---
|
||||
/// A connection flow has started for this session.
|
||||
SessionConnecting { session_id: SessionId },
|
||||
/// This session's server is connected and ready. Includes the `HostId`
|
||||
/// received from the initialize handshake, for model deduplication.
|
||||
SessionConnected {
|
||||
session_id: SessionId,
|
||||
host_id: HostId,
|
||||
},
|
||||
/// The remote server launch or handshake failed.
|
||||
SessionConnectionFailed {
|
||||
session_id: SessionId,
|
||||
/// Which phase of the connection flow failed.
|
||||
phase: RemoteServerInitPhase,
|
||||
/// The error message from the failed phase.
|
||||
error: String,
|
||||
},
|
||||
/// This session's connection dropped. Carries `host_id` so consumers
|
||||
/// don't need to look it up from the already-transitioned state.
|
||||
/// This session's underlying connection is no longer usable: the
|
||||
/// stream closed (EOF/error), the initialize handshake failed, or the
|
||||
/// session was explicitly deregistered while `Connected`. Signals to
|
||||
/// subscribers that they should drop any `Arc<RemoteServerClient>` they
|
||||
/// hold for this session. Carries `host_id` so consumers don't need to
|
||||
/// look it up from the already-transitioned state.
|
||||
///
|
||||
/// Note this is about *transport* state, not manager tracking: after
|
||||
/// this event fires the session may still be present in the manager
|
||||
/// in the `Disconnected` state (e.g. when the stream dropped on its
|
||||
/// own). Use `SessionDeregistered` to observe removal from the manager.
|
||||
SessionDisconnected {
|
||||
session_id: SessionId,
|
||||
host_id: HostId,
|
||||
},
|
||||
/// The manager is no longer tracking this session -- it has been
|
||||
/// removed from the `sessions` map via `deregister_session`. Fires
|
||||
/// exactly once per session, and only on explicit teardown (never as
|
||||
/// a result of a spontaneous connection drop).
|
||||
///
|
||||
/// If the session was `Connected` at the point of deregistration, a
|
||||
/// `SessionDisconnected` event is emitted first so transport-level
|
||||
/// subscribers can release their client references.
|
||||
SessionDeregistered { session_id: SessionId },
|
||||
|
||||
// --- Host-scoped events ---
|
||||
/// The first session for this host reached `Connected`. Downstream
|
||||
/// features should create per-host models (e.g. `RepoMetadataModel`).
|
||||
HostConnected { host_id: HostId },
|
||||
/// The last session for this host was disconnected or deregistered.
|
||||
/// Downstream features should tear down per-host models.
|
||||
HostDisconnected { host_id: HostId },
|
||||
|
||||
// --- Repo metadata events (forwarded from ClientEvent push channel) ---
|
||||
/// Response to a `navigate_to_directory` request.
|
||||
NavigatedToDirectory {
|
||||
session_id: SessionId,
|
||||
host_id: HostId,
|
||||
indexed_path: String,
|
||||
is_git: bool,
|
||||
},
|
||||
/// A full or lazy-loaded repo metadata snapshot was pushed by the server.
|
||||
RepoMetadataSnapshot {
|
||||
host_id: HostId,
|
||||
update: RepoMetadataUpdate,
|
||||
},
|
||||
/// An incremental repo metadata update was pushed by the server.
|
||||
RepoMetadataUpdated {
|
||||
host_id: HostId,
|
||||
update: RepoMetadataUpdate,
|
||||
},
|
||||
/// A `LoadRepoMetadataDirectory` response was received from the server.
|
||||
RepoMetadataDirectoryLoaded {
|
||||
host_id: HostId,
|
||||
update: RepoMetadataUpdate,
|
||||
},
|
||||
|
||||
// --- Setup events ---
|
||||
/// Intermediate state change during the binary check/install flow.
|
||||
SetupStateChanged {
|
||||
session_id: SessionId,
|
||||
state: RemoteServerSetupState,
|
||||
},
|
||||
/// Result of [`RemoteServerManager::check_binary`]. Returns a result where:
|
||||
/// - `Ok(true)` means the binary is installed and executable,
|
||||
/// - `Ok(false)` means it is definitively not installed, and
|
||||
/// - `Err(_)` means the check itself failed (e.g. SSH error or timeout).
|
||||
BinaryCheckComplete {
|
||||
session_id: SessionId,
|
||||
result: Result<bool, String>,
|
||||
/// The detected remote platform (OS + arch) from `uname -sm`.
|
||||
/// `None` if detection failed or was not attempted.
|
||||
remote_platform: Option<RemotePlatform>,
|
||||
},
|
||||
/// Result of [`RemoteServerManager::install_binary`]. Returns a result where:
|
||||
/// - `Ok(())` means the install succeeded, and
|
||||
/// - `Err(_)` means the install failed and carries the failure reason (SSH error, timeout, script error, etc.).
|
||||
BinaryInstallComplete {
|
||||
session_id: SessionId,
|
||||
result: Result<(), String>,
|
||||
},
|
||||
|
||||
// --- Telemetry events ---
|
||||
/// A client request to the remote server failed.
|
||||
ClientRequestFailed {
|
||||
session_id: SessionId,
|
||||
operation: RemoteServerOperation,
|
||||
error_kind: RemoteServerErrorKind,
|
||||
},
|
||||
/// A server message could not be decoded (no parseable request_id).
|
||||
ServerMessageDecodingError { session_id: SessionId },
|
||||
}
|
||||
|
||||
/// Shell info stashed by [`RemoteServerManager::notify_session_bootstrapped`]
|
||||
/// when the session is not yet in `Connected` state. Flushed automatically
|
||||
/// when [`RemoteServerManager::mark_session_connected`] fires.
|
||||
#[cfg_attr(target_family = "wasm", allow(dead_code))]
|
||||
struct PendingSessionBootstrappedNotification {
|
||||
shell_type: String,
|
||||
shell_path: Option<String>,
|
||||
}
|
||||
|
||||
/// Singleton model that manages connections to `remote_server` processes on
|
||||
/// remote hosts.
|
||||
///
|
||||
/// Each SSH session gets its own `RemoteServerClient` and SSH connection.
|
||||
/// Deduplication of the underlying long-lived server process happens on the
|
||||
/// remote host. The `HostId` returned by the server's `InitializeResponse`
|
||||
/// is used on the client to deduplicate host-scoped models (e.g.
|
||||
/// `RepoMetadataModel`), not connections.
|
||||
pub struct RemoteServerManager {
|
||||
/// Per-session connection state. Each SSH session gets its own dedicated
|
||||
/// connection to the remote server.
|
||||
sessions: HashMap<SessionId, RemoteSessionState>,
|
||||
/// Reverse index: host → sessions for O(1) lookup by `HostId`.
|
||||
host_to_sessions: HashMap<HostId, HashSet<SessionId>>,
|
||||
/// Spawner for running closures back on the main thread.
|
||||
spawner: ModelSpawner<Self>,
|
||||
/// Last path requested per session for dedup. Avoids redundant
|
||||
/// `navigate_to_directory` calls when `update_active_session` fires
|
||||
/// repeatedly for the same CWD.
|
||||
last_navigated_path: HashMap<SessionId, String>,
|
||||
/// Per-session `SessionBootstrapped` notifications that arrived before the
|
||||
/// session reached `Connected`. Flushed in `mark_session_connected`.
|
||||
pending_bootstrapped_notifications: HashMap<SessionId, PendingSessionBootstrappedNotification>,
|
||||
/// Detected remote platform per session, populated during the binary check
|
||||
/// phase via `detect_platform()`. Used for telemetry.
|
||||
session_platforms: HashMap<SessionId, RemotePlatform>,
|
||||
}
|
||||
|
||||
impl Entity for RemoteServerManager {
|
||||
type Event = RemoteServerManagerEvent;
|
||||
}
|
||||
|
||||
impl SingletonEntity for RemoteServerManager {}
|
||||
|
||||
impl RemoteServerManager {
|
||||
pub fn new(ctx: &mut ModelContext<Self>) -> Self {
|
||||
Self {
|
||||
sessions: HashMap::new(),
|
||||
host_to_sessions: HashMap::new(),
|
||||
spawner: ctx.spawner(),
|
||||
last_navigated_path: HashMap::new(),
|
||||
pending_bootstrapped_notifications: HashMap::new(),
|
||||
session_platforms: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a connected client for the given host by picking an arbitrary
|
||||
/// session from the host's session pool.
|
||||
pub fn client_for_host(&self, host_id: &HostId) -> Option<&Arc<RemoteServerClient>> {
|
||||
let sessions = self.host_to_sessions.get(host_id)?;
|
||||
sessions
|
||||
.iter()
|
||||
.find_map(|session_id| self.client_for_session(*session_id))
|
||||
}
|
||||
|
||||
/// Checks if the remote server binary is installed and executable.
|
||||
/// Emits `BinaryCheckComplete { result }`.
|
||||
///
|
||||
/// Returns Ok(true) if the binary is installed and executable,
|
||||
/// Ok(false) if it is definitively not installed, and
|
||||
/// Err(_) if the check failed (e.g. SSH timeout/unreachable).
|
||||
#[cfg_attr(target_family = "wasm", allow(unused_variables))]
|
||||
pub fn check_binary<T>(
|
||||
&mut self,
|
||||
session_id: SessionId,
|
||||
transport: T,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) where
|
||||
T: RemoteTransport + 'static,
|
||||
{
|
||||
#[cfg(target_family = "wasm")]
|
||||
{
|
||||
log::warn!("check_binary is a no-op on WASM");
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
{
|
||||
ctx.emit(RemoteServerManagerEvent::SetupStateChanged {
|
||||
session_id,
|
||||
state: RemoteServerSetupState::Checking,
|
||||
});
|
||||
let spawner = self.spawner.clone();
|
||||
ctx.background_executor()
|
||||
.spawn(async move {
|
||||
// Run platform detection and binary check concurrently.
|
||||
let (platform_result, check_result) =
|
||||
futures::join!(transport.detect_platform(), transport.check_binary(),);
|
||||
let platform = match platform_result {
|
||||
Ok(p) => Some(p),
|
||||
Err(e) => {
|
||||
log::warn!("Platform detection failed for session {session_id:?}: {e}");
|
||||
None
|
||||
}
|
||||
};
|
||||
let _ = spawner
|
||||
.spawn(move |me, ctx| {
|
||||
if let Some(ref p) = platform {
|
||||
me.session_platforms.insert(session_id, p.clone());
|
||||
}
|
||||
ctx.emit(RemoteServerManagerEvent::BinaryCheckComplete {
|
||||
session_id,
|
||||
result: check_result,
|
||||
remote_platform: platform,
|
||||
});
|
||||
})
|
||||
.await;
|
||||
})
|
||||
.detach();
|
||||
}
|
||||
}
|
||||
|
||||
/// Installs the remote server binary.
|
||||
/// Emits `BinaryInstallComplete { result }`.
|
||||
///
|
||||
/// Returns Ok(()) if the install succeeded, and
|
||||
/// Err(_) if the install failed (e.g. SSH timeout/unreachable).
|
||||
#[cfg_attr(target_family = "wasm", allow(unused_variables))]
|
||||
pub fn install_binary<T>(
|
||||
&mut self,
|
||||
session_id: SessionId,
|
||||
transport: T,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) where
|
||||
T: RemoteTransport + 'static,
|
||||
{
|
||||
#[cfg(target_family = "wasm")]
|
||||
{
|
||||
log::warn!("install_binary is a no-op on WASM");
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
{
|
||||
ctx.emit(RemoteServerManagerEvent::SetupStateChanged {
|
||||
session_id,
|
||||
state: RemoteServerSetupState::Installing {
|
||||
progress_percent: None,
|
||||
},
|
||||
});
|
||||
let spawner = self.spawner.clone();
|
||||
ctx.background_executor()
|
||||
.spawn(async move {
|
||||
let result = transport.install_binary().await;
|
||||
let _ = spawner
|
||||
.spawn(move |_me, ctx| {
|
||||
ctx.emit(RemoteServerManagerEvent::BinaryInstallComplete {
|
||||
session_id,
|
||||
result,
|
||||
});
|
||||
})
|
||||
.await;
|
||||
})
|
||||
.detach();
|
||||
}
|
||||
}
|
||||
|
||||
/// Entry point for establishing a remote server connection for a session.
|
||||
/// This assumes the binary is already installed and executable.
|
||||
/// Callers should first call `check_binary` and `install_binary` to ensure the binary is present.
|
||||
///
|
||||
/// The full flow is:
|
||||
/// 1. **Connect** — `transport.connect()` establishes the I/O streams and
|
||||
/// creates the `RemoteServerClient`.
|
||||
/// 2. **Handshake** — perform the initialize handshake (which returns the
|
||||
/// `HostId`) and transition to `Connected`.
|
||||
///
|
||||
/// No-op on WASM (remote server connections use a different transport).
|
||||
#[cfg_attr(target_family = "wasm", allow(unused_variables, unused_mut))]
|
||||
pub fn connect_session<T>(
|
||||
&mut self,
|
||||
session_id: SessionId,
|
||||
transport: T,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) where
|
||||
T: RemoteTransport + 'static,
|
||||
{
|
||||
#[cfg(target_family = "wasm")]
|
||||
{
|
||||
log::warn!("connect_session is a no-op on WASM");
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
{
|
||||
log::info!("Starting remote server connection for session {session_id:?}");
|
||||
|
||||
// Advance the user-visible setup pipeline. Both callers (binary
|
||||
// already installed, and binary just installed) enter this
|
||||
// method right when the Initializing phase begins, so we emit
|
||||
// the state change from one place.
|
||||
ctx.emit(RemoteServerManagerEvent::SetupStateChanged {
|
||||
session_id,
|
||||
state: RemoteServerSetupState::Initializing,
|
||||
});
|
||||
|
||||
self.sessions
|
||||
.insert(session_id, RemoteSessionState::Connecting);
|
||||
ctx.emit(RemoteServerManagerEvent::SessionConnecting { session_id });
|
||||
|
||||
let spawner = self.spawner.clone();
|
||||
let executor = ctx.background_executor().clone();
|
||||
|
||||
ctx.background_executor()
|
||||
.spawn(async move {
|
||||
// ---- Phase 1: Connect (establish streams, create client) ----
|
||||
match transport.connect(&executor).await {
|
||||
Ok(Connection {
|
||||
client,
|
||||
event_rx,
|
||||
child,
|
||||
control_path,
|
||||
}) => {
|
||||
let client = Arc::new(client);
|
||||
|
||||
// Transition to Initializing and start draining
|
||||
// the event channel for push events and disconnect.
|
||||
// The `Child` is stashed on the session state so
|
||||
// its lifetime is controlled by the manager -- on
|
||||
// teardown the state is dropped, which runs the
|
||||
// `Child`'s destructor and SIGKILLs the subprocess
|
||||
// via `kill_on_drop`. `control_path` is stashed
|
||||
// for explicit teardown's `ssh -O exit` call.
|
||||
let client_for_state = Arc::clone(&client);
|
||||
let _ = spawner
|
||||
.spawn(move |me, ctx| {
|
||||
me.sessions.insert(
|
||||
session_id,
|
||||
RemoteSessionState::Initializing {
|
||||
client: client_for_state,
|
||||
_child: child,
|
||||
control_path,
|
||||
},
|
||||
);
|
||||
|
||||
// Drain the event channel on the main thread.
|
||||
// Each push event is forwarded as a manager
|
||||
// event in real-time. When the stream closes
|
||||
// (after Disconnected or channel drop), we
|
||||
// transition the session to Disconnected.
|
||||
ctx.spawn_stream_local(
|
||||
event_rx,
|
||||
move |me, event, ctx| {
|
||||
me.forward_client_event(session_id, event, ctx);
|
||||
},
|
||||
move |me, ctx| {
|
||||
me.mark_session_disconnected(session_id, ctx);
|
||||
},
|
||||
);
|
||||
})
|
||||
.await;
|
||||
|
||||
// ---- Phase 2: Initialize handshake ----
|
||||
match client.initialize().await {
|
||||
Ok(resp) => {
|
||||
let host_id = HostId::new(resp.host_id);
|
||||
let _ = spawner
|
||||
.spawn(move |me, ctx| {
|
||||
me.mark_session_connected(session_id, host_id, ctx);
|
||||
})
|
||||
.await;
|
||||
}
|
||||
Err(e) => {
|
||||
log::error!(
|
||||
"Initialize handshake failed for session {session_id:?}: {e}"
|
||||
);
|
||||
let error = format!("{e:#}");
|
||||
let _ = spawner
|
||||
.spawn(move |me, ctx| {
|
||||
ctx.emit(
|
||||
RemoteServerManagerEvent::SetupStateChanged {
|
||||
session_id,
|
||||
state: RemoteServerSetupState::Failed {
|
||||
error: error.clone(),
|
||||
},
|
||||
},
|
||||
);
|
||||
ctx.emit(
|
||||
RemoteServerManagerEvent::SessionConnectionFailed {
|
||||
session_id,
|
||||
phase: RemoteServerInitPhase::Initialize,
|
||||
error,
|
||||
},
|
||||
);
|
||||
me.mark_session_disconnected(session_id, ctx);
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
log::error!(
|
||||
"Failed to connect remote server for session {session_id:?}: {e:#}"
|
||||
);
|
||||
let error = format!("{e:#}");
|
||||
let _ = spawner
|
||||
.spawn(move |me, ctx| {
|
||||
ctx.emit(RemoteServerManagerEvent::SetupStateChanged {
|
||||
session_id,
|
||||
state: RemoteServerSetupState::Failed {
|
||||
error: error.clone(),
|
||||
},
|
||||
});
|
||||
ctx.emit(RemoteServerManagerEvent::SessionConnectionFailed {
|
||||
session_id,
|
||||
phase: RemoteServerInitPhase::Connect,
|
||||
error,
|
||||
});
|
||||
me.mark_session_disconnected(session_id, ctx);
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
})
|
||||
.detach();
|
||||
}
|
||||
}
|
||||
|
||||
/// Removes a session from the manager and tears down its connection.
|
||||
///
|
||||
/// Assumes the caller has already observed that the user's shell
|
||||
/// has exited (in practice this is only invoked from the
|
||||
/// `ExitShell` teardown path). Under that assumption we also force
|
||||
/// the local SSH `ControlMaster` to exit immediately via
|
||||
/// `ssh -O exit`, which is required because the master is the
|
||||
/// user's interactive ssh process and, without the explicit
|
||||
/// `-O exit`, it hangs waiting for remote-side cleanup of
|
||||
/// multiplexed channels (see [`crate::ssh::stop_control_master`]).
|
||||
///
|
||||
/// Mechanically:
|
||||
/// 1. Remove the session entry. Dropping the `RemoteSessionState`
|
||||
/// drops the transport's owned `Child`, which SIGKILLs the
|
||||
/// `ssh … remote-server-proxy` subprocess via `kill_on_drop`.
|
||||
/// 2. If the session had a ControlMaster `control_path`, spawn a
|
||||
/// background task that runs `ssh -O exit` against it.
|
||||
///
|
||||
/// The `Child` is owned by the manager's state, *not* by
|
||||
/// `Arc<RemoteServerClient>`. Lingering `Arc` clones held elsewhere
|
||||
/// (e.g. by the per-session command executor) do *not* keep the
|
||||
/// subprocess alive -- removing the state here always SIGKILLs the
|
||||
/// child, regardless of client refcount.
|
||||
///
|
||||
/// Two separate events can fire here, and they mean different things:
|
||||
///
|
||||
/// * `SessionDisconnected` -- the *transport* went away. Emitted only
|
||||
/// when the session was `Connected` at the time of deregistration.
|
||||
/// Subscribers can use this to drop their
|
||||
/// `Arc<RemoteServerClient>` references and cancel in-flight
|
||||
/// requests. The same event also fires independently from
|
||||
/// `mark_session_disconnected` when the stream drops on its own.
|
||||
/// * `SessionDeregistered` -- the manager is no longer *tracking* this
|
||||
/// session. Always emitted, regardless of which state the session
|
||||
/// was in, because the entry is being removed from `sessions`
|
||||
/// outright. Unlike `SessionDisconnected`, this one never fires for
|
||||
/// spontaneous drops -- only for explicit teardown.
|
||||
pub fn deregister_session(&mut self, session_id: SessionId, ctx: &mut ModelContext<Self>) {
|
||||
self.last_navigated_path.remove(&session_id);
|
||||
self.pending_bootstrapped_notifications.remove(&session_id);
|
||||
self.session_platforms.remove(&session_id);
|
||||
|
||||
// Remove the session entry. Dropping the `RemoteSessionState`
|
||||
// here drops the transport's owned `Child` (if any), which
|
||||
// SIGKILLs the `ssh … remote-server-proxy` subprocess via
|
||||
// `kill_on_drop`.
|
||||
let prev = self.sessions.remove(&session_id);
|
||||
|
||||
// Extract the ControlMaster socket path (if any) so we can
|
||||
// force the master to exit below. Safe to do under the
|
||||
// "caller already observed ExitShell" assumption documented
|
||||
// above.
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
let control_path = match &prev {
|
||||
Some(RemoteSessionState::Connected { control_path, .. })
|
||||
| Some(RemoteSessionState::Initializing { control_path, .. }) => control_path.clone(),
|
||||
_ => None,
|
||||
};
|
||||
|
||||
if let Some(RemoteSessionState::Connected { host_id, .. }) = prev {
|
||||
self.remove_from_host_index(&host_id, session_id);
|
||||
ctx.emit(RemoteServerManagerEvent::SessionDisconnected {
|
||||
session_id,
|
||||
host_id: host_id.clone(),
|
||||
});
|
||||
if !self.host_to_sessions.contains_key(&host_id) {
|
||||
ctx.emit(RemoteServerManagerEvent::HostDisconnected {
|
||||
host_id: host_id.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
ctx.emit(RemoteServerManagerEvent::SessionDeregistered { session_id });
|
||||
|
||||
// Force the local SSH ControlMaster to exit after teardown.
|
||||
// Spawned detached because the ssh subcommand may take a moment
|
||||
// to complete and we don't want to block the main thread on it.
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
if let Some(control_path) = control_path {
|
||||
ctx.background_executor()
|
||||
.spawn(async move {
|
||||
crate::ssh::stop_control_master(&control_path).await;
|
||||
})
|
||||
.detach();
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the client for this session, if connected.
|
||||
pub fn client_for_session(&self, session_id: SessionId) -> Option<&Arc<RemoteServerClient>> {
|
||||
match self.sessions.get(&session_id) {
|
||||
Some(RemoteSessionState::Connected { client, .. }) => Some(client),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the connection state for this session.
|
||||
pub fn session(&self, session_id: SessionId) -> Option<&RemoteSessionState> {
|
||||
self.sessions.get(&session_id)
|
||||
}
|
||||
|
||||
/// Returns the detected remote platform for this session, if available.
|
||||
pub fn platform_for_session(&self, session_id: SessionId) -> Option<&RemotePlatform> {
|
||||
self.session_platforms.get(&session_id)
|
||||
}
|
||||
|
||||
/// Returns the `HostId` for this session, if the initialize handshake
|
||||
/// has completed. Downstream features use this to deduplicate
|
||||
/// host-scoped models (e.g. `RepoMetadataModel`).
|
||||
pub fn host_id_for_session(&self, session_id: SessionId) -> Option<&HostId> {
|
||||
match self.sessions.get(&session_id) {
|
||||
Some(RemoteSessionState::Connected { host_id, .. }) => Some(host_id),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns all session IDs connected to a given host. O(1) via the
|
||||
/// reverse index.
|
||||
pub fn sessions_for_host(&self, host_id: &HostId) -> Option<&HashSet<SessionId>> {
|
||||
self.host_to_sessions.get(host_id)
|
||||
}
|
||||
|
||||
/// Sends a `NavigatedToDirectory` request to the remote server for
|
||||
/// the given session and emits the response as a manager event.
|
||||
///
|
||||
/// Deduplicates: if the same `(session_id, path)` was already requested,
|
||||
/// the call is a no-op.
|
||||
pub fn navigate_to_directory(
|
||||
&mut self,
|
||||
session_id: SessionId,
|
||||
path: String,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
// Dedup: skip if this session already navigated to the same path.
|
||||
if self.last_navigated_path.get(&session_id) == Some(&path) {
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(client) = self.client_for_session(session_id).cloned() else {
|
||||
log::warn!("navigate_to_directory: no connected client for session {session_id:?}");
|
||||
return;
|
||||
};
|
||||
let Some(host_id) = self.host_id_for_session(session_id).cloned() else {
|
||||
log::warn!("navigate_to_directory: no host_id for session {session_id:?}");
|
||||
return;
|
||||
};
|
||||
|
||||
// Record only after confirming the client is connected, so that a
|
||||
// retry after SessionConnected is not incorrectly deduplicated.
|
||||
self.last_navigated_path.insert(session_id, path.clone());
|
||||
|
||||
let spawner = self.spawner.clone();
|
||||
ctx.background_executor()
|
||||
.spawn(async move {
|
||||
match client.navigate_to_directory(path).await {
|
||||
Ok(resp) => {
|
||||
let _ = spawner
|
||||
.spawn(move |_me, ctx| {
|
||||
ctx.emit(RemoteServerManagerEvent::NavigatedToDirectory {
|
||||
session_id,
|
||||
host_id,
|
||||
indexed_path: resp.indexed_path,
|
||||
is_git: resp.is_git,
|
||||
});
|
||||
})
|
||||
.await;
|
||||
}
|
||||
Err(e) => {
|
||||
log::error!("navigate_to_directory failed for session {session_id:?}: {e}");
|
||||
let error_kind = RemoteServerErrorKind::from_client_error(&e);
|
||||
let _ = spawner
|
||||
.spawn(move |_me, ctx| {
|
||||
ctx.emit(RemoteServerManagerEvent::ClientRequestFailed {
|
||||
session_id,
|
||||
operation: RemoteServerOperation::NavigateToDirectory,
|
||||
error_kind,
|
||||
});
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
})
|
||||
.detach();
|
||||
}
|
||||
|
||||
/// Sends a `SessionBootstrapped` notification to the remote server.
|
||||
///
|
||||
/// If the session is already in `Connected` state the notification is sent
|
||||
/// immediately. Otherwise it is stashed and automatically flushed when
|
||||
/// `mark_session_connected` transitions the session to `Connected`.
|
||||
pub fn notify_session_bootstrapped(
|
||||
&mut self,
|
||||
session_id: SessionId,
|
||||
shell_type: &str,
|
||||
shell_path: Option<&str>,
|
||||
) {
|
||||
if let Some(client) = self.client_for_session(session_id) {
|
||||
client.notify_session_bootstrapped(session_id, shell_type, shell_path);
|
||||
} else {
|
||||
log::info!(
|
||||
"notify_session_bootstrapped: session {session_id:?} not yet connected, \
|
||||
stashing notification"
|
||||
);
|
||||
self.pending_bootstrapped_notifications.insert(
|
||||
session_id,
|
||||
PendingSessionBootstrappedNotification {
|
||||
shell_type: shell_type.to_owned(),
|
||||
shell_path: shell_path.map(ToOwned::to_owned),
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Sends a `LoadRepoMetadataDirectory` request to the remote server for
|
||||
/// the given session and emits the response as a manager event.
|
||||
pub fn load_remote_repo_metadata_directory(
|
||||
&mut self,
|
||||
session_id: SessionId,
|
||||
repo_path: String,
|
||||
dir_path: String,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
let Some(client) = self.client_for_session(session_id).cloned() else {
|
||||
log::warn!(
|
||||
"load_remote_repo_metadata_directory: no connected client for session {session_id:?}"
|
||||
);
|
||||
return;
|
||||
};
|
||||
let Some(host_id) = self.host_id_for_session(session_id).cloned() else {
|
||||
log::warn!(
|
||||
"load_remote_repo_metadata_directory: no host_id for session {session_id:?}"
|
||||
);
|
||||
return;
|
||||
};
|
||||
|
||||
let spawner = self.spawner.clone();
|
||||
ctx.background_executor()
|
||||
.spawn(async move {
|
||||
match client
|
||||
.load_repo_metadata_directory(repo_path, dir_path)
|
||||
.await
|
||||
{
|
||||
Ok(resp) => {
|
||||
if let Some(update) =
|
||||
crate::repo_metadata_proto::proto_load_repo_metadata_directory_response_to_update(&resp)
|
||||
{
|
||||
let _ = spawner
|
||||
.spawn(move |_me, ctx| {
|
||||
ctx.emit(
|
||||
RemoteServerManagerEvent::RepoMetadataDirectoryLoaded {
|
||||
host_id,
|
||||
update,
|
||||
},
|
||||
);
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
log::error!(
|
||||
"load_repo_metadata_directory failed for session {session_id:?}: {e}"
|
||||
);
|
||||
let error_kind = RemoteServerErrorKind::from_client_error(&e);
|
||||
let _ = spawner
|
||||
.spawn(move |_me, ctx| {
|
||||
ctx.emit(RemoteServerManagerEvent::ClientRequestFailed {
|
||||
session_id,
|
||||
operation: RemoteServerOperation::LoadRepoMetadataDirectory,
|
||||
error_kind,
|
||||
});
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
})
|
||||
.detach();
|
||||
}
|
||||
|
||||
/// Forwards a push event from the client event channel as a manager event.
|
||||
/// No-ops if the session is not in `Connected` state (i.e. `host_id` not
|
||||
/// yet available).
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
fn forward_client_event(
|
||||
&self,
|
||||
session_id: SessionId,
|
||||
event: ClientEvent,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
let Some(host_id) = self.host_id_for_session(session_id) else {
|
||||
log::debug!("Dropping push event for session {session_id:?}: not connected yet");
|
||||
return;
|
||||
};
|
||||
let host_id = host_id.clone();
|
||||
|
||||
match event {
|
||||
ClientEvent::RepoMetadataSnapshotReceived { update } => {
|
||||
ctx.emit(RemoteServerManagerEvent::RepoMetadataSnapshot { host_id, update });
|
||||
}
|
||||
ClientEvent::RepoMetadataUpdated { update } => {
|
||||
ctx.emit(RemoteServerManagerEvent::RepoMetadataUpdated { host_id, update });
|
||||
}
|
||||
ClientEvent::MessageDecodingError => {
|
||||
ctx.emit(RemoteServerManagerEvent::ServerMessageDecodingError { session_id });
|
||||
}
|
||||
ClientEvent::Disconnected => {
|
||||
// Handled by the drain loop's completion callback.
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
fn mark_session_connected(
|
||||
&mut self,
|
||||
session_id: SessionId,
|
||||
host_id: HostId,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
log::info!("Remote server connected for session {session_id:?}, host {host_id}");
|
||||
|
||||
// Only transition if the session is still in Initializing state.
|
||||
// Remove first so we can move the client handle (and owned `Child`)
|
||||
// out.
|
||||
let Some(RemoteSessionState::Initializing {
|
||||
client,
|
||||
_child,
|
||||
control_path,
|
||||
}) = self.sessions.remove(&session_id)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
|
||||
let is_first_session = !self.host_to_sessions.contains_key(&host_id);
|
||||
self.sessions.insert(
|
||||
session_id,
|
||||
RemoteSessionState::Connected {
|
||||
client,
|
||||
host_id: host_id.clone(),
|
||||
_child,
|
||||
control_path,
|
||||
},
|
||||
);
|
||||
self.host_to_sessions
|
||||
.entry(host_id.clone())
|
||||
.or_default()
|
||||
.insert(session_id);
|
||||
if is_first_session {
|
||||
ctx.emit(RemoteServerManagerEvent::HostConnected {
|
||||
host_id: host_id.clone(),
|
||||
});
|
||||
}
|
||||
ctx.emit(RemoteServerManagerEvent::SetupStateChanged {
|
||||
session_id,
|
||||
state: RemoteServerSetupState::Ready,
|
||||
});
|
||||
ctx.emit(RemoteServerManagerEvent::SessionConnected {
|
||||
session_id,
|
||||
host_id,
|
||||
});
|
||||
|
||||
// Flush any SessionBootstrapped notification that was stashed before
|
||||
// the session reached Connected.
|
||||
if let Some(notif) = self.pending_bootstrapped_notifications.remove(&session_id) {
|
||||
if let Some(client) = self.client_for_session(session_id) {
|
||||
log::info!(
|
||||
"Flushing stashed SessionBootstrapped notification for session \
|
||||
{session_id:?}"
|
||||
);
|
||||
client.notify_session_bootstrapped(
|
||||
session_id,
|
||||
¬if.shell_type,
|
||||
notif.shell_path.as_deref(),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
pub(crate) fn mark_session_disconnected(
|
||||
&mut self,
|
||||
session_id: SessionId,
|
||||
ctx: &mut ModelContext<Self>,
|
||||
) {
|
||||
self.pending_bootstrapped_notifications.remove(&session_id);
|
||||
let Some(prev) = self.sessions.remove(&session_id) else {
|
||||
return;
|
||||
};
|
||||
self.sessions
|
||||
.insert(session_id, RemoteSessionState::Disconnected);
|
||||
|
||||
if let RemoteSessionState::Connected { host_id, .. } = prev {
|
||||
self.remove_from_host_index(&host_id, session_id);
|
||||
// Emit `SessionDisconnected` before `HostDisconnected` so that
|
||||
// subscribers (e.g. the command executor) drop their
|
||||
// `Arc<RemoteServerClient>` reference before any host-scoped
|
||||
// teardown runs. This matches the ordering in
|
||||
// `deregister_session` so both teardown paths look identical
|
||||
// to subscribers.
|
||||
ctx.emit(RemoteServerManagerEvent::SessionDisconnected {
|
||||
session_id,
|
||||
host_id: host_id.clone(),
|
||||
});
|
||||
if !self.host_to_sessions.contains_key(&host_id) {
|
||||
ctx.emit(RemoteServerManagerEvent::HostDisconnected { host_id });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Removes a session from the host → sessions reverse index.
|
||||
/// Cleans up the entry entirely if the set becomes empty.
|
||||
fn remove_from_host_index(&mut self, host_id: &HostId, session_id: SessionId) {
|
||||
if let Some(set) = self.host_to_sessions.get_mut(host_id) {
|
||||
set.remove(&session_id);
|
||||
if set.is_empty() {
|
||||
self.host_to_sessions.remove(host_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,238 @@
|
||||
use std::fmt;
|
||||
|
||||
use futures::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
|
||||
use prost::Message;
|
||||
|
||||
use crate::proto::{ClientMessage, ServerMessage};
|
||||
|
||||
/// Maximum allowed message payload size (64 MB).
|
||||
///
|
||||
/// `read_message` rejects payloads exceeding this limit after decoding the
|
||||
/// length prefix but before allocating the payload buffer, preventing OOM from
|
||||
/// corrupted or adversarial length prefixes.
|
||||
pub const MAX_MESSAGE_SIZE: usize = 64 * 1024 * 1024;
|
||||
|
||||
/// Errors that can occur during protocol-level read/write operations.
|
||||
#[derive(thiserror::Error, Debug)]
|
||||
pub enum ProtocolError {
|
||||
#[error("I/O error: {0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
|
||||
/// When full protobuf decode fails, the protocol layer attempts to extract
|
||||
/// the `request_id` from the raw bytes so callers can correlate the error.
|
||||
#[error("Failed to decode protobuf message: {0}")]
|
||||
Decode(prost::DecodeError, Option<RequestId>),
|
||||
|
||||
#[error("Unexpected EOF while reading message")]
|
||||
UnexpectedEof,
|
||||
|
||||
#[error("Message too large: {size} bytes exceeds limit of {max} bytes")]
|
||||
MessageTooLarge { size: usize, max: usize },
|
||||
}
|
||||
|
||||
/// A typed wrapper around the proto `string request_id` field.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
pub struct RequestId(String);
|
||||
|
||||
impl RequestId {
|
||||
#[allow(clippy::new_without_default)]
|
||||
pub fn new() -> Self {
|
||||
Self(uuid::Uuid::new_v4().to_string())
|
||||
}
|
||||
|
||||
/// Returns true if this is an empty request ID, indicating a push message
|
||||
/// from the server (not correlated to any client request).
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.0.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<String> for RequestId {
|
||||
fn from(s: String) -> Self {
|
||||
Self(s)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<RequestId> for String {
|
||||
fn from(id: RequestId) -> Self {
|
||||
id.0
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for RequestId {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.write_str(&self.0)
|
||||
}
|
||||
}
|
||||
|
||||
/// Reads a length-delimited protobuf message from `reader`.
|
||||
///
|
||||
/// Wire format: `[4-byte little-endian length][protobuf bytes]`.
|
||||
pub async fn read_message<M: Message + Default>(
|
||||
reader: &mut (impl AsyncRead + Unpin),
|
||||
) -> Result<M, ProtocolError> {
|
||||
let mut len_buf = [0u8; 4];
|
||||
match reader.read_exact(&mut len_buf).await {
|
||||
Ok(_) => {}
|
||||
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
|
||||
return Err(ProtocolError::UnexpectedEof);
|
||||
}
|
||||
Err(e) => return Err(ProtocolError::Io(e)),
|
||||
}
|
||||
let len = u32::from_le_bytes(len_buf) as usize;
|
||||
|
||||
if len > MAX_MESSAGE_SIZE {
|
||||
return Err(ProtocolError::MessageTooLarge {
|
||||
size: len,
|
||||
max: MAX_MESSAGE_SIZE,
|
||||
});
|
||||
}
|
||||
|
||||
let mut buf = vec![0u8; len];
|
||||
reader.read_exact(&mut buf).await.map_err(|e| {
|
||||
if e.kind() == std::io::ErrorKind::UnexpectedEof {
|
||||
ProtocolError::UnexpectedEof
|
||||
} else {
|
||||
ProtocolError::Io(e)
|
||||
}
|
||||
})?;
|
||||
|
||||
M::decode(&buf[..]).map_err(|e| {
|
||||
let request_id = try_extract_request_id(&buf).map(RequestId::from);
|
||||
ProtocolError::Decode(e, request_id)
|
||||
})
|
||||
}
|
||||
|
||||
/// Writes a length-delimited protobuf message to `writer`.
|
||||
///
|
||||
/// Wire format: `[4-byte little-endian length][protobuf bytes]`.
|
||||
pub async fn write_message<M: Message>(
|
||||
writer: &mut (impl AsyncWrite + Unpin),
|
||||
msg: &M,
|
||||
) -> Result<(), ProtocolError> {
|
||||
let encoded = msg.encode_to_vec();
|
||||
if encoded.len() > MAX_MESSAGE_SIZE {
|
||||
return Err(ProtocolError::MessageTooLarge {
|
||||
size: encoded.len(),
|
||||
max: MAX_MESSAGE_SIZE,
|
||||
});
|
||||
}
|
||||
let len = encoded.len() as u32;
|
||||
writer.write_all(&len.to_le_bytes()).await?;
|
||||
writer.write_all(&encoded).await?;
|
||||
writer.flush().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Reads a `ClientMessage` from the given reader.
|
||||
pub async fn read_client_message(
|
||||
reader: &mut (impl AsyncRead + Unpin),
|
||||
) -> Result<ClientMessage, ProtocolError> {
|
||||
read_message(reader).await
|
||||
}
|
||||
|
||||
/// Writes a `ClientMessage` to the given writer.
|
||||
pub async fn write_client_message(
|
||||
writer: &mut (impl AsyncWrite + Unpin),
|
||||
msg: &ClientMessage,
|
||||
) -> Result<(), ProtocolError> {
|
||||
write_message(writer, msg).await
|
||||
}
|
||||
|
||||
/// Reads a `ServerMessage` from the given reader.
|
||||
pub async fn read_server_message(
|
||||
reader: &mut (impl AsyncRead + Unpin),
|
||||
) -> Result<ServerMessage, ProtocolError> {
|
||||
read_message(reader).await
|
||||
}
|
||||
|
||||
/// Writes a `ServerMessage` to the given writer.
|
||||
pub async fn write_server_message(
|
||||
writer: &mut (impl AsyncWrite + Unpin),
|
||||
msg: &ServerMessage,
|
||||
) -> Result<(), ProtocolError> {
|
||||
write_message(writer, msg).await
|
||||
}
|
||||
|
||||
impl ProtocolError {
|
||||
/// Whether a read loop can safely continue after this error.
|
||||
///
|
||||
/// True only when the payload was fully consumed, keeping the stream aligned
|
||||
/// at the next length prefix.
|
||||
pub fn is_read_recoverable(&self) -> bool {
|
||||
match self {
|
||||
ProtocolError::Decode(..) => true,
|
||||
ProtocolError::Io(_) => false,
|
||||
ProtocolError::UnexpectedEof => false,
|
||||
ProtocolError::MessageTooLarge { .. } => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether a write loop can safely continue after this error.
|
||||
///
|
||||
/// True only when nothing was written to the stream, keeping it aligned.
|
||||
pub fn is_write_recoverable(&self) -> bool {
|
||||
match self {
|
||||
ProtocolError::MessageTooLarge { .. } => true,
|
||||
ProtocolError::Io(_) => false,
|
||||
ProtocolError::Decode(..) => false,
|
||||
ProtocolError::UnexpectedEof => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Attempts to extract the `request_id` from raw protobuf bytes by parsing
|
||||
/// only field 1 (string) and ignoring the rest of the buffer.
|
||||
///
|
||||
/// This uses manual wire-format parsing: field 1 of type string has tag byte
|
||||
/// `0x0a` (field_number=1, wire_type=2) followed by a varint length and UTF-8
|
||||
/// bytes. We stop as soon as field 1 is extracted, so corruption in later
|
||||
/// bytes does not affect extraction.
|
||||
///
|
||||
/// **Note**: This assumes `request_id` is always field 1 in the message schema.
|
||||
/// If the protobuf schema changes, update this accordingly.
|
||||
///
|
||||
/// Returns `None` if the buffer doesn't start with a valid field 1 string,
|
||||
/// or if the extracted string is empty.
|
||||
fn try_extract_request_id(buf: &[u8]) -> Option<String> {
|
||||
// Field 1 (string) wire tag: field_number=1, wire_type=2 (length-delimited).
|
||||
if buf.first() != Some(&0x0a) {
|
||||
return None;
|
||||
}
|
||||
let buf = &buf[1..];
|
||||
|
||||
// Decode varint-encoded string length.
|
||||
let (len, consumed) = decode_varint(buf)?;
|
||||
let buf = &buf[consumed..];
|
||||
|
||||
if buf.len() < len {
|
||||
return None;
|
||||
}
|
||||
|
||||
let s = std::str::from_utf8(&buf[..len]).ok()?;
|
||||
if s.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(s.to_string())
|
||||
}
|
||||
|
||||
/// Decodes a protobuf varint from the start of `buf`.
|
||||
/// Returns `(value, bytes_consumed)` or `None` if the varint is malformed.
|
||||
fn decode_varint(buf: &[u8]) -> Option<(usize, usize)> {
|
||||
let mut result: u64 = 0;
|
||||
for (i, &byte) in buf.iter().enumerate() {
|
||||
if i >= 10 {
|
||||
// Varint too long.
|
||||
return None;
|
||||
}
|
||||
result |= ((byte & 0x7F) as u64) << (i * 7);
|
||||
if byte & 0x80 == 0 {
|
||||
return Some((result as usize, i + 1));
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "protocol_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,200 @@
|
||||
use prost::Message;
|
||||
|
||||
use crate::proto::{
|
||||
client_message, server_message, ClientMessage, Initialize, InitializeResponse, ServerMessage,
|
||||
};
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn round_trip_client_message() {
|
||||
let msg = ClientMessage {
|
||||
request_id: "test-123".to_string(),
|
||||
message: Some(client_message::Message::Initialize(Initialize {})),
|
||||
};
|
||||
|
||||
let mut buf = Vec::new();
|
||||
write_client_message(&mut buf, &msg).await.unwrap();
|
||||
|
||||
let mut cursor = &buf[..];
|
||||
let decoded: ClientMessage = read_client_message(&mut cursor).await.unwrap();
|
||||
|
||||
assert_eq!(decoded.request_id, "test-123");
|
||||
match decoded.message {
|
||||
Some(client_message::Message::Initialize(_)) => {}
|
||||
other => panic!("unexpected message variant: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn round_trip_server_message() {
|
||||
let msg = ServerMessage {
|
||||
request_id: "resp-456".to_string(),
|
||||
message: Some(server_message::Message::InitializeResponse(
|
||||
InitializeResponse {
|
||||
server_version: "0.1.0".to_string(),
|
||||
host_id: "test-host".to_string(),
|
||||
},
|
||||
)),
|
||||
};
|
||||
|
||||
let mut buf = Vec::new();
|
||||
write_server_message(&mut buf, &msg).await.unwrap();
|
||||
|
||||
let mut cursor = &buf[..];
|
||||
let decoded: ServerMessage = read_server_message(&mut cursor).await.unwrap();
|
||||
|
||||
assert_eq!(decoded.request_id, "resp-456");
|
||||
match decoded.message {
|
||||
Some(server_message::Message::InitializeResponse(resp)) => {
|
||||
assert_eq!(resp.server_version, "0.1.0");
|
||||
}
|
||||
other => panic!("unexpected message variant: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_unexpected_eof_on_empty_input() {
|
||||
let mut cursor: &[u8] = &[];
|
||||
let result = read_client_message(&mut cursor).await;
|
||||
assert!(matches!(result, Err(ProtocolError::UnexpectedEof)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_truncated_payload() {
|
||||
// Write a length prefix claiming 100 bytes, but only provide 4.
|
||||
let mut buf = Vec::new();
|
||||
buf.extend_from_slice(&100u32.to_le_bytes());
|
||||
buf.extend_from_slice(&[0u8; 4]);
|
||||
|
||||
let mut cursor = &buf[..];
|
||||
let result = read_client_message(&mut cursor).await;
|
||||
assert!(matches!(result, Err(ProtocolError::UnexpectedEof)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn round_trip_zero_length_message() {
|
||||
// A default ClientMessage with no fields set encodes to zero bytes.
|
||||
let msg = ClientMessage::default();
|
||||
|
||||
let mut buf = Vec::new();
|
||||
write_client_message(&mut buf, &msg).await.unwrap();
|
||||
|
||||
// The first 4 bytes should be the length (0).
|
||||
assert_eq!(&buf[..4], &0u32.to_le_bytes());
|
||||
|
||||
let mut cursor = &buf[..];
|
||||
let decoded: ClientMessage = read_client_message(&mut cursor).await.unwrap();
|
||||
assert_eq!(decoded.request_id, "");
|
||||
assert!(decoded.message.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_message_too_large() {
|
||||
// Write a length prefix exceeding MAX_MESSAGE_SIZE.
|
||||
let oversized_len = (MAX_MESSAGE_SIZE as u32) + 1;
|
||||
let buf = oversized_len.to_le_bytes();
|
||||
|
||||
let mut cursor = &buf[..];
|
||||
let result = read_client_message(&mut cursor).await;
|
||||
assert!(matches!(result, Err(ProtocolError::MessageTooLarge { .. })));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_message_too_large() {
|
||||
// Build a ClientMessage whose encoded size exceeds MAX_MESSAGE_SIZE.
|
||||
let msg = ClientMessage {
|
||||
request_id: "x".repeat(MAX_MESSAGE_SIZE + 1),
|
||||
message: None,
|
||||
};
|
||||
|
||||
let mut buf = Vec::new();
|
||||
let result = write_client_message(&mut buf, &msg).await;
|
||||
assert!(matches!(result, Err(ProtocolError::MessageTooLarge { .. })));
|
||||
// Nothing should have been written to the stream.
|
||||
assert!(buf.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn try_extract_request_id_from_valid_message() {
|
||||
let msg = ClientMessage {
|
||||
request_id: "abc-123".to_string(),
|
||||
message: Some(client_message::Message::Initialize(Initialize {})),
|
||||
};
|
||||
let buf = msg.encode_to_vec();
|
||||
assert_eq!(try_extract_request_id(&buf), Some("abc-123".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn try_extract_request_id_from_corrupted_payload_with_valid_id() {
|
||||
// Manually construct bytes: valid request_id field followed by
|
||||
// corrupt trailing bytes (unterminated varint that would crash
|
||||
// a full prost decode but doesn't affect our field-1 extraction).
|
||||
let mut buf = Vec::new();
|
||||
// Field 1 (string): tag=0x0a, length=7, "req-456"
|
||||
buf.push(0x0a);
|
||||
buf.push(7);
|
||||
buf.extend_from_slice(b"req-456");
|
||||
// Corrupt trailing bytes: unterminated varint (all continuation bits set).
|
||||
buf.extend_from_slice(&[0xFF, 0xFF, 0xFF, 0xFF]);
|
||||
|
||||
// request_id should still be extractable despite trailing corruption.
|
||||
assert_eq!(try_extract_request_id(&buf), Some("req-456".to_string()));
|
||||
}
|
||||
#[test]
|
||||
fn try_extract_request_id_from_empty_bytes() {
|
||||
assert_eq!(try_extract_request_id(&[]), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn try_extract_request_id_from_garbage_bytes() {
|
||||
// Completely random bytes that don't form a valid protobuf.
|
||||
// This may or may not decode depending on what prost makes of it,
|
||||
// but should not panic. If it decodes to an empty request_id, we
|
||||
// return None.
|
||||
let result = try_extract_request_id(&[0xFF, 0xFF, 0xFF, 0xFF]);
|
||||
// We don't assert a specific value — just that it doesn't panic.
|
||||
// If prost happens to decode something, it'll be empty or garbage.
|
||||
let _ = result;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decode_error_extracts_request_id() {
|
||||
// Construct a corrupted message with a valid request_id field.
|
||||
let mut payload = Vec::new();
|
||||
// Field 1 (string): tag=0x0a, length=6, "req-42"
|
||||
payload.push(0x0a);
|
||||
payload.push(6);
|
||||
payload.extend_from_slice(b"req-42");
|
||||
// Invalid trailing bytes that cause prost decode failure.
|
||||
payload.extend_from_slice(&[0x0F, 0x01]);
|
||||
|
||||
let mut buf = Vec::new();
|
||||
buf.extend_from_slice(&(payload.len() as u32).to_le_bytes());
|
||||
buf.extend_from_slice(&payload);
|
||||
|
||||
let mut cursor = &buf[..];
|
||||
let result = read_client_message(&mut cursor).await;
|
||||
match result {
|
||||
Err(ProtocolError::Decode(_, Some(id))) => {
|
||||
assert_eq!(id.to_string(), "req-42");
|
||||
}
|
||||
other => panic!("expected Decode error with request_id, got: {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decode_error_none_when_no_request_id() {
|
||||
// Completely invalid protobuf bytes with no valid field 1.
|
||||
let garbage = vec![0xFF, 0xFE, 0xFD, 0xFC];
|
||||
let mut buf = Vec::new();
|
||||
buf.extend_from_slice(&(garbage.len() as u32).to_le_bytes());
|
||||
buf.extend_from_slice(&garbage);
|
||||
|
||||
let mut cursor = &buf[..];
|
||||
let result = read_client_message(&mut cursor).await;
|
||||
match result {
|
||||
Err(ProtocolError::Decode(_, None)) => {}
|
||||
other => panic!("expected Decode error with None request_id, got: {other:?}"),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,291 @@
|
||||
//! Conversion between `repo_metadata` Rust types and proto-generated types.
|
||||
//!
|
||||
//! The Rust types in `repo_metadata::file_tree_update` were designed to mirror the
|
||||
//! proto schema 1:1, so these conversions are straightforward field mappings.
|
||||
|
||||
use repo_metadata::file_tree_store::{FileTreeEntry, FileTreeEntryState};
|
||||
use repo_metadata::file_tree_update::{
|
||||
DirectoryNodeMetadata, FileNodeMetadata, FileTreeEntryUpdate, RepoMetadataUpdate,
|
||||
RepoNodeMetadata,
|
||||
};
|
||||
use warp_util::standardized_path::StandardizedPath;
|
||||
|
||||
use crate::proto;
|
||||
|
||||
// ── Rust → Proto ────────────────────────────────────────────
|
||||
|
||||
impl From<&RepoMetadataUpdate> for proto::RepoMetadataUpdatePush {
|
||||
fn from(update: &RepoMetadataUpdate) -> Self {
|
||||
Self {
|
||||
repo_path: update.repo_path.to_string(),
|
||||
remove_entries: update
|
||||
.remove_entries
|
||||
.iter()
|
||||
.map(|p| p.to_string())
|
||||
.collect(),
|
||||
update_entries: update
|
||||
.update_entries
|
||||
.iter()
|
||||
.map(proto::RepoMetadataEntryUpdate::from)
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&FileTreeEntryUpdate> for proto::RepoMetadataEntryUpdate {
|
||||
fn from(update: &FileTreeEntryUpdate) -> Self {
|
||||
Self {
|
||||
parent_path_to_replace: update.parent_path_to_replace.to_string(),
|
||||
subtree_metadata: update
|
||||
.subtree_metadata
|
||||
.iter()
|
||||
.map(proto::RepoNodeMetadata::from)
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&RepoNodeMetadata> for proto::RepoNodeMetadata {
|
||||
fn from(node: &RepoNodeMetadata) -> Self {
|
||||
let node_oneof = match node {
|
||||
RepoNodeMetadata::Directory(dir) => {
|
||||
proto::repo_node_metadata::Node::Directory(proto::DirectoryNodeMetadata {
|
||||
path: dir.path.to_string(),
|
||||
ignored: dir.ignored,
|
||||
loaded: dir.loaded,
|
||||
})
|
||||
}
|
||||
RepoNodeMetadata::File(file) => {
|
||||
proto::repo_node_metadata::Node::File(proto::FileNodeMetadata {
|
||||
path: file.path.to_string(),
|
||||
extension: file.extension.clone(),
|
||||
ignored: file.ignored,
|
||||
})
|
||||
}
|
||||
};
|
||||
Self {
|
||||
node: Some(node_oneof),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Serializes a full `FileTreeEntry`
|
||||
/// messages suitable for a `RepoMetadataSnapshot`.
|
||||
///
|
||||
/// Walks the tree breadth-first from the root, producing one `RepoMetadataEntryUpdate`
|
||||
/// per parent directory containing its immediate children as `RepoNodeMetadata`.
|
||||
pub fn file_tree_entry_to_snapshot_proto(
|
||||
entry: &FileTreeEntry,
|
||||
) -> Vec<proto::RepoMetadataEntryUpdate> {
|
||||
let mut result = Vec::new();
|
||||
let mut queue = std::collections::VecDeque::new();
|
||||
queue.push_back(entry.root_directory().clone());
|
||||
|
||||
while let Some(current_path) = queue.pop_front() {
|
||||
let children: Vec<_> = entry.child_paths(¤t_path).cloned().collect();
|
||||
if children.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let mut subtree_metadata = Vec::with_capacity(children.len());
|
||||
for child_path in &children {
|
||||
match entry.get(child_path) {
|
||||
Some(FileTreeEntryState::Directory(dir)) => {
|
||||
subtree_metadata.push(proto::RepoNodeMetadata {
|
||||
node: Some(proto::repo_node_metadata::Node::Directory(
|
||||
proto::DirectoryNodeMetadata {
|
||||
path: dir.path.to_string(),
|
||||
ignored: dir.ignored,
|
||||
loaded: dir.loaded,
|
||||
},
|
||||
)),
|
||||
});
|
||||
// Enqueue for breadth-first traversal.
|
||||
queue.push_back(child_path.clone());
|
||||
}
|
||||
Some(FileTreeEntryState::File(file)) => {
|
||||
subtree_metadata.push(proto::RepoNodeMetadata {
|
||||
node: Some(proto::repo_node_metadata::Node::File(
|
||||
proto::FileNodeMetadata {
|
||||
path: file.path.to_string(),
|
||||
extension: file.extension.clone(),
|
||||
ignored: file.ignored,
|
||||
},
|
||||
)),
|
||||
});
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
}
|
||||
|
||||
if !subtree_metadata.is_empty() {
|
||||
result.push(proto::RepoMetadataEntryUpdate {
|
||||
parent_path_to_replace: current_path.to_string(),
|
||||
subtree_metadata,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
// ── Proto → Rust ──────────────────────────────────────────────────
|
||||
|
||||
/// Converts a `RepoMetadataUpdatePush` proto message into a `RepoMetadataUpdate`.
|
||||
pub fn proto_to_repo_metadata_update(
|
||||
push: &proto::RepoMetadataUpdatePush,
|
||||
) -> Option<RepoMetadataUpdate> {
|
||||
let repo_path = StandardizedPath::try_new(&push.repo_path).ok()?;
|
||||
|
||||
let remove_entries: Vec<StandardizedPath> = push
|
||||
.remove_entries
|
||||
.iter()
|
||||
.filter_map(|p| match StandardizedPath::try_new(p) {
|
||||
Ok(path) => Some(path),
|
||||
Err(e) => {
|
||||
log::warn!("Skipping invalid remove_entry path {p:?}: {e}");
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
let update_entries: Vec<FileTreeEntryUpdate> = push
|
||||
.update_entries
|
||||
.iter()
|
||||
.filter_map(proto_to_entry_update)
|
||||
.collect();
|
||||
|
||||
Some(RepoMetadataUpdate {
|
||||
repo_path,
|
||||
remove_entries,
|
||||
update_entries,
|
||||
})
|
||||
}
|
||||
|
||||
/// Converts a `RepoMetadataSnapshot` proto into a `RepoMetadataUpdate`
|
||||
/// (with no removals) that can be applied to a `RemoteRepoMetadataModel`.
|
||||
pub fn proto_snapshot_to_update(
|
||||
snapshot: &proto::RepoMetadataSnapshot,
|
||||
) -> Option<RepoMetadataUpdate> {
|
||||
let repo_path = StandardizedPath::try_new(&snapshot.repo_path).ok()?;
|
||||
|
||||
let update_entries: Vec<FileTreeEntryUpdate> = snapshot
|
||||
.entries
|
||||
.iter()
|
||||
.filter_map(proto_to_entry_update)
|
||||
.collect();
|
||||
|
||||
Some(RepoMetadataUpdate {
|
||||
repo_path,
|
||||
remove_entries: Vec::new(),
|
||||
update_entries,
|
||||
})
|
||||
}
|
||||
|
||||
fn proto_to_entry_update(
|
||||
proto_update: &proto::RepoMetadataEntryUpdate,
|
||||
) -> Option<FileTreeEntryUpdate> {
|
||||
let parent_path = StandardizedPath::try_new(&proto_update.parent_path_to_replace).ok()?;
|
||||
|
||||
let subtree_metadata: Vec<RepoNodeMetadata> = proto_update
|
||||
.subtree_metadata
|
||||
.iter()
|
||||
.filter_map(proto_to_repo_node_metadata)
|
||||
.collect();
|
||||
|
||||
Some(FileTreeEntryUpdate {
|
||||
parent_path_to_replace: parent_path,
|
||||
subtree_metadata,
|
||||
})
|
||||
}
|
||||
|
||||
/// Converts a `LoadRepoMetadataDirectoryResponse` proto into a `RepoMetadataUpdate`
|
||||
/// (with no removals) that can be applied to a `RemoteRepoMetadataModel`.
|
||||
pub fn proto_load_repo_metadata_directory_response_to_update(
|
||||
resp: &proto::LoadRepoMetadataDirectoryResponse,
|
||||
) -> Option<RepoMetadataUpdate> {
|
||||
let repo_path = StandardizedPath::try_new(&resp.repo_path).ok()?;
|
||||
|
||||
let update_entries: Vec<FileTreeEntryUpdate> = resp
|
||||
.entries
|
||||
.iter()
|
||||
.filter_map(proto_to_entry_update)
|
||||
.collect();
|
||||
|
||||
Some(RepoMetadataUpdate {
|
||||
repo_path,
|
||||
remove_entries: Vec::new(),
|
||||
update_entries,
|
||||
})
|
||||
}
|
||||
|
||||
/// Serializes the immediate children of a directory in a `FileTreeEntry` as
|
||||
/// `RepoMetadataEntryUpdate` protos. Used to build a `LoadRepoMetadataDirectoryResponse`.
|
||||
pub fn file_tree_children_to_proto_entries(
|
||||
entry: &FileTreeEntry,
|
||||
dir_path: &StandardizedPath,
|
||||
) -> Vec<proto::RepoMetadataEntryUpdate> {
|
||||
let children: Vec<_> = entry.child_paths(dir_path).cloned().collect();
|
||||
if children.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let mut subtree_metadata = Vec::with_capacity(children.len());
|
||||
for child_path in &children {
|
||||
match entry.get(child_path) {
|
||||
Some(FileTreeEntryState::Directory(dir)) => {
|
||||
subtree_metadata.push(proto::RepoNodeMetadata {
|
||||
node: Some(proto::repo_node_metadata::Node::Directory(
|
||||
proto::DirectoryNodeMetadata {
|
||||
path: dir.path.to_string(),
|
||||
ignored: dir.ignored,
|
||||
loaded: dir.loaded,
|
||||
},
|
||||
)),
|
||||
});
|
||||
}
|
||||
Some(FileTreeEntryState::File(file)) => {
|
||||
subtree_metadata.push(proto::RepoNodeMetadata {
|
||||
node: Some(proto::repo_node_metadata::Node::File(
|
||||
proto::FileNodeMetadata {
|
||||
path: file.path.to_string(),
|
||||
extension: file.extension.clone(),
|
||||
ignored: file.ignored,
|
||||
},
|
||||
)),
|
||||
});
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
}
|
||||
|
||||
if subtree_metadata.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
vec![proto::RepoMetadataEntryUpdate {
|
||||
parent_path_to_replace: dir_path.to_string(),
|
||||
subtree_metadata,
|
||||
}]
|
||||
}
|
||||
|
||||
fn proto_to_repo_node_metadata(proto_node: &proto::RepoNodeMetadata) -> Option<RepoNodeMetadata> {
|
||||
match proto_node.node.as_ref()? {
|
||||
proto::repo_node_metadata::Node::Directory(dir) => {
|
||||
let path = StandardizedPath::try_new(&dir.path).ok()?;
|
||||
Some(RepoNodeMetadata::Directory(DirectoryNodeMetadata {
|
||||
path,
|
||||
ignored: dir.ignored,
|
||||
loaded: dir.loaded,
|
||||
}))
|
||||
}
|
||||
proto::repo_node_metadata::Node::File(file) => {
|
||||
let path = StandardizedPath::try_new(&file.path).ok()?;
|
||||
Some(RepoNodeMetadata::File(FileNodeMetadata {
|
||||
path,
|
||||
extension: file.extension.clone(),
|
||||
ignored: file.ignored,
|
||||
}))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,215 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use warp_core::channel::{Channel, ChannelState};
|
||||
|
||||
/// State machine for the remote server install → launch → initialize flow.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum RemoteServerSetupState {
|
||||
/// Checking if the binary exists on remote.
|
||||
Checking,
|
||||
/// Downloading and installing the binary.
|
||||
Installing { progress_percent: Option<u8> },
|
||||
/// Binary is launched, waiting for InitializeResponse.
|
||||
Initializing,
|
||||
/// Handshake complete. Ready.
|
||||
Ready,
|
||||
/// Something failed. Fall back to ControlMaster.
|
||||
Failed { error: String },
|
||||
}
|
||||
|
||||
impl RemoteServerSetupState {
|
||||
pub fn is_ready(&self) -> bool {
|
||||
matches!(self, Self::Ready)
|
||||
}
|
||||
|
||||
pub fn is_failed(&self) -> bool {
|
||||
matches!(self, Self::Failed { .. })
|
||||
}
|
||||
|
||||
pub fn is_terminal(&self) -> bool {
|
||||
self.is_ready() || self.is_failed()
|
||||
}
|
||||
|
||||
pub fn is_in_progress(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::Checking | Self::Installing { .. } | Self::Initializing
|
||||
)
|
||||
}
|
||||
|
||||
pub fn is_connecting(&self) -> bool {
|
||||
matches!(self, Self::Installing { .. } | Self::Initializing)
|
||||
}
|
||||
}
|
||||
|
||||
/// Detected remote platform from `uname -sm` output.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct RemotePlatform {
|
||||
pub os: RemoteOs,
|
||||
pub arch: RemoteArch,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum RemoteOs {
|
||||
Linux,
|
||||
MacOs,
|
||||
}
|
||||
|
||||
impl RemoteOs {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::Linux => "linux",
|
||||
Self::MacOs => "macos",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum RemoteArch {
|
||||
X86_64,
|
||||
Aarch64,
|
||||
}
|
||||
|
||||
impl RemoteArch {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::X86_64 => "x86_64",
|
||||
Self::Aarch64 => "aarch64",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse `uname -sm` output into a `RemotePlatform`.
|
||||
///
|
||||
/// The expected format is `<os> <arch>`, e.g. `Linux x86_64` or `Darwin arm64`.
|
||||
/// Takes the last line to skip any shell initialization output.
|
||||
pub fn parse_uname_output(output: &str) -> Result<RemotePlatform> {
|
||||
let line = output
|
||||
.lines()
|
||||
.last()
|
||||
.ok_or_else(|| anyhow!("empty uname output"))?
|
||||
.trim();
|
||||
|
||||
let mut parts = line.split_whitespace();
|
||||
let os_str = parts
|
||||
.next()
|
||||
.ok_or_else(|| anyhow!("missing OS in uname output: {line}"))?;
|
||||
let arch_str = parts
|
||||
.next()
|
||||
.ok_or_else(|| anyhow!("missing arch in uname output: {line}"))?;
|
||||
|
||||
let os = match os_str {
|
||||
"Linux" => RemoteOs::Linux,
|
||||
"Darwin" => RemoteOs::MacOs,
|
||||
other => return Err(anyhow!("unsupported OS: {other}")),
|
||||
};
|
||||
|
||||
let arch = match arch_str {
|
||||
"x86_64" => RemoteArch::X86_64,
|
||||
"aarch64" | "arm64" | "armv8l" => RemoteArch::Aarch64,
|
||||
other => return Err(anyhow!("unsupported arch: {other}")),
|
||||
};
|
||||
|
||||
Ok(RemotePlatform { os, arch })
|
||||
}
|
||||
|
||||
/// Returns the remote directory where the binary is installed, keyed by channel.
|
||||
///
|
||||
/// - stable: `~/.warp/remote-server`
|
||||
/// - preview: `~/.warp-preview/remote-server`
|
||||
/// - dev: `~/.warp-dev/remote-server`
|
||||
/// - local: `~/.warp-local/remote-server`
|
||||
/// - integration: `~/.warp-dev/remote-server`
|
||||
/// - warp-oss: `~/.warp-oss/remote-server`
|
||||
pub fn remote_server_dir() -> String {
|
||||
let warp_dir = match ChannelState::channel() {
|
||||
Channel::Stable => ".warp",
|
||||
Channel::Preview => ".warp-preview",
|
||||
Channel::Dev | Channel::Integration => ".warp-dev",
|
||||
Channel::Local => ".warp-local",
|
||||
Channel::Oss => {
|
||||
// TODO(alokedesai): need to figure out how remote server works with warp-oss
|
||||
// For now, return what Dev returns.
|
||||
".warp-dev"
|
||||
}
|
||||
};
|
||||
format!("~/{warp_dir}/remote-server")
|
||||
}
|
||||
|
||||
/// Returns the binary name, keyed by channel.
|
||||
///
|
||||
/// Matches the CLI command names: `oz` (stable), `oz-preview`, `oz-dev`.
|
||||
pub fn binary_name() -> &'static str {
|
||||
ChannelState::channel().cli_command_name()
|
||||
}
|
||||
|
||||
/// Returns the full remote binary path.
|
||||
pub fn remote_server_binary() -> String {
|
||||
format!("{}/{}", remote_server_dir(), binary_name())
|
||||
}
|
||||
|
||||
/// Returns the shell command to check if the remote server binary exists and
|
||||
/// is executable.
|
||||
pub fn binary_check_command() -> String {
|
||||
let bin = remote_server_binary();
|
||||
format!("test -x {bin}")
|
||||
}
|
||||
|
||||
/// The install script template, loaded from a standalone `.sh` file for
|
||||
/// readability. Placeholders like `{download_base_url}` are substituted by
|
||||
/// [`install_script`].
|
||||
const INSTALL_SCRIPT_TEMPLATE: &str = include_str!("install_remote_server.sh");
|
||||
|
||||
/// Returns the install script that downloads and installs the CLI binary.
|
||||
///
|
||||
/// The script detects the remote architecture via `uname -m`, downloads the
|
||||
/// correct Oz CLI tarball from the download URL (with os, arch, package, and
|
||||
/// channel query params), and extracts it to the install directory.
|
||||
///
|
||||
/// All parameters (URL, channel, directory, binary name) are derived
|
||||
/// internally from the current channel configuration.
|
||||
pub fn install_script() -> String {
|
||||
INSTALL_SCRIPT_TEMPLATE
|
||||
.replace("{download_base_url}", &download_url())
|
||||
.replace("{channel}", download_channel())
|
||||
.replace("{install_dir}", &remote_server_dir())
|
||||
.replace("{binary_name}", binary_name())
|
||||
}
|
||||
|
||||
/// Construct the download URL from the server root URL.
|
||||
///
|
||||
/// For example, given `https://app.warp.dev`, returns
|
||||
/// `https://app.warp.dev/download/cli`.
|
||||
fn download_url() -> String {
|
||||
let base = ChannelState::server_root_url();
|
||||
let base = base.trim_end_matches('/');
|
||||
format!("{base}/download/cli")
|
||||
}
|
||||
|
||||
/// Maps the client's [`Channel`] to the server's download channel parameter.
|
||||
///
|
||||
/// The server recognises `"stable"`, `"preview"`, and `"dev"`. Local and
|
||||
/// Integration builds map to `"dev"` so they fetch dogfood artifacts.
|
||||
fn download_channel() -> &'static str {
|
||||
match ChannelState::channel() {
|
||||
Channel::Stable => "stable",
|
||||
Channel::Preview => "preview",
|
||||
Channel::Dev | Channel::Local | Channel::Integration => "dev",
|
||||
Channel::Oss => {
|
||||
// TODO(alokedesai): need to figure out how remote server works with warp-oss
|
||||
// For now, return what Dev returns.
|
||||
"dev"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Timeout for the binary existence check.
|
||||
pub const CHECK_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
|
||||
/// Timeout for the install script.
|
||||
pub const INSTALL_TIMEOUT: Duration = Duration::from_secs(60);
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "setup_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,108 @@
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parse_uname_linux_x86_64() {
|
||||
let platform = parse_uname_output("Linux x86_64").unwrap();
|
||||
assert_eq!(platform.os, RemoteOs::Linux);
|
||||
assert_eq!(platform.arch, RemoteArch::X86_64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_uname_linux_aarch64() {
|
||||
let platform = parse_uname_output("Linux aarch64").unwrap();
|
||||
assert_eq!(platform.os, RemoteOs::Linux);
|
||||
assert_eq!(platform.arch, RemoteArch::Aarch64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_uname_darwin_arm64() {
|
||||
let platform = parse_uname_output("Darwin arm64").unwrap();
|
||||
assert_eq!(platform.os, RemoteOs::MacOs);
|
||||
assert_eq!(platform.arch, RemoteArch::Aarch64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_uname_darwin_x86_64() {
|
||||
let platform = parse_uname_output("Darwin x86_64").unwrap();
|
||||
assert_eq!(platform.os, RemoteOs::MacOs);
|
||||
assert_eq!(platform.arch, RemoteArch::X86_64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_uname_linux_armv8l() {
|
||||
let platform = parse_uname_output("Linux armv8l").unwrap();
|
||||
assert_eq!(platform.os, RemoteOs::Linux);
|
||||
assert_eq!(platform.arch, RemoteArch::Aarch64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_uname_skips_shell_initialization_output() {
|
||||
let output = "Last login: Mon Apr 7 10:00:00 2025\nWelcome to Ubuntu\nLinux x86_64";
|
||||
let platform = parse_uname_output(output).unwrap();
|
||||
assert_eq!(platform.os, RemoteOs::Linux);
|
||||
assert_eq!(platform.arch, RemoteArch::X86_64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_uname_trims_whitespace() {
|
||||
let platform = parse_uname_output(" Linux x86_64 \n").unwrap();
|
||||
assert_eq!(platform.os, RemoteOs::Linux);
|
||||
assert_eq!(platform.arch, RemoteArch::X86_64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_uname_unsupported_os() {
|
||||
let result = parse_uname_output("Windows x86_64");
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().to_string().contains("unsupported OS"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_uname_unsupported_arch() {
|
||||
let result = parse_uname_output("Linux mips");
|
||||
assert!(result.is_err());
|
||||
assert!(result.unwrap_err().to_string().contains("unsupported arch"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_uname_empty_output() {
|
||||
let result = parse_uname_output("");
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_uname_missing_arch() {
|
||||
let result = parse_uname_output("Linux");
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn state_is_ready() {
|
||||
assert!(RemoteServerSetupState::Ready.is_ready());
|
||||
assert!(!RemoteServerSetupState::Checking.is_ready());
|
||||
assert!(!RemoteServerSetupState::Initializing.is_ready());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn state_is_failed() {
|
||||
assert!(RemoteServerSetupState::Failed {
|
||||
error: "test".into()
|
||||
}
|
||||
.is_failed());
|
||||
assert!(!RemoteServerSetupState::Ready.is_failed());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn state_is_terminal() {
|
||||
assert!(RemoteServerSetupState::Ready.is_terminal());
|
||||
assert!(RemoteServerSetupState::Failed {
|
||||
error: "test".into()
|
||||
}
|
||||
.is_terminal());
|
||||
assert!(!RemoteServerSetupState::Checking.is_terminal());
|
||||
assert!(!RemoteServerSetupState::Installing {
|
||||
progress_percent: None
|
||||
}
|
||||
.is_terminal());
|
||||
assert!(!RemoteServerSetupState::Initializing.is_terminal());
|
||||
}
|
||||
@@ -0,0 +1,155 @@
|
||||
use std::path::Path;
|
||||
use std::process::Output;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use command::r#async::Command;
|
||||
use warpui::r#async::FutureExt as _;
|
||||
|
||||
/// Timeout for `ssh -O exit`. The command only talks to the local
|
||||
/// ControlMaster over a Unix socket, so it should return almost
|
||||
/// immediately; if it doesn't, we'd rather give up than block
|
||||
/// teardown.
|
||||
const STOP_CONTROL_MASTER_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
|
||||
/// Builds the common SSH argument list for multiplexed connections through
|
||||
/// an existing ControlMaster socket.
|
||||
pub fn ssh_args(socket_path: &Path) -> Vec<String> {
|
||||
vec![
|
||||
"-q".to_string(),
|
||||
"-o".to_string(),
|
||||
"PasswordAuthentication=no".to_string(),
|
||||
"-o".to_string(),
|
||||
"ForwardX11=no".to_string(),
|
||||
"-o".to_string(),
|
||||
format!("ControlPath={}", socket_path.display()),
|
||||
"placeholder@placeholder".to_string(),
|
||||
]
|
||||
}
|
||||
|
||||
/// Runs `ssh -O exit -o ControlPath=<socket_path>` to force the local
|
||||
/// SSH `ControlMaster` managing `socket_path` to exit immediately,
|
||||
/// without waiting for multiplexed channels to finish draining.
|
||||
///
|
||||
/// The user's interactive ssh is spawned with `-o ControlMaster=yes` by
|
||||
/// `warp_ssh_helper`, so it is both the interactive session and the
|
||||
/// multiplex master. When the user's remote shell exits, that ssh can
|
||||
/// hang waiting for half-closed slave channels (e.g. from
|
||||
/// `ssh ... remote-server-proxy`) to finish cleanup on the remote
|
||||
/// side. Sending `-O exit` bypasses that wait.
|
||||
///
|
||||
/// **Only safe to call once the user's shell has already exited** --
|
||||
/// this tears down the interactive ssh outright. In practice it is
|
||||
/// invoked from the `ExitShell` teardown path on the client.
|
||||
///
|
||||
/// Fire-and-forget. Errors are logged but not propagated: at teardown
|
||||
/// time there is nothing useful to do with them.
|
||||
pub async fn stop_control_master(socket_path: &Path) {
|
||||
let args = ssh_args(socket_path);
|
||||
let result = async {
|
||||
Command::new("ssh")
|
||||
.arg("-O")
|
||||
.arg("exit")
|
||||
.args(&args)
|
||||
.kill_on_drop(true)
|
||||
.output()
|
||||
.await
|
||||
}
|
||||
.with_timeout(STOP_CONTROL_MASTER_TIMEOUT)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(Ok(output)) if output.status.success() => {
|
||||
log::info!(
|
||||
"stop_control_master: `ssh -O exit` succeeded for {}",
|
||||
socket_path.display()
|
||||
);
|
||||
}
|
||||
Ok(Ok(output)) => {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
log::info!(
|
||||
"stop_control_master: `ssh -O exit` for {} exited with {:?}: {stderr}",
|
||||
socket_path.display(),
|
||||
output.status.code(),
|
||||
);
|
||||
}
|
||||
Ok(Err(e)) => {
|
||||
log::info!(
|
||||
"stop_control_master: failed to spawn `ssh -O exit` for {}: {e}",
|
||||
socket_path.display()
|
||||
);
|
||||
}
|
||||
Err(_) => {
|
||||
log::warn!(
|
||||
"stop_control_master: `ssh -O exit` for {} timed out after {:?}",
|
||||
socket_path.display(),
|
||||
STOP_CONTROL_MASTER_TIMEOUT,
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Run a single SSH command through the ControlMaster socket and return a result where:
|
||||
/// - `Err` for transport-level failures (e.g. couldn't spawn `ssh`, or timeout).
|
||||
/// - `Ok(output)` callers should check `output.status` to distinguish a successful remote command from a non-zero remote exit.
|
||||
pub async fn run_ssh_command(
|
||||
socket_path: &Path,
|
||||
remote_command: &str,
|
||||
timeout: Duration,
|
||||
) -> Result<Output> {
|
||||
async {
|
||||
Command::new("ssh")
|
||||
.args(ssh_args(socket_path))
|
||||
.arg(remote_command)
|
||||
.kill_on_drop(true)
|
||||
.output()
|
||||
.await
|
||||
}
|
||||
.with_timeout(timeout)
|
||||
.await
|
||||
.map_err(|_| anyhow!("SSH command timed out after {timeout:?}"))?
|
||||
.map_err(|e| anyhow!("SSH command failed to execute: {e}"))
|
||||
}
|
||||
|
||||
/// Pipe a script into `bash -s` on the remote host via the ControlMaster
|
||||
/// socket. Returns a result where:
|
||||
/// - `Err` for transport-level failures (e.g. couldn't spawn `ssh`, or timeout).
|
||||
/// - `Ok(output)` callers should check `output.status` to distinguish a successful remote script from a non-zero remote exit.
|
||||
///
|
||||
/// We pipe via stdin rather than passing the script as an SSH command-line
|
||||
/// argument because the install script is multi-line and contains shell
|
||||
/// constructs (case statements, variable expansions, single/double quotes)
|
||||
/// that would require complex, fragile escaping if passed as an argument.
|
||||
/// The `bash -s` + stdin approach avoids all escaping issues and has no
|
||||
/// argument length limits.
|
||||
pub async fn run_ssh_script(socket_path: &Path, script: &str, timeout: Duration) -> Result<Output> {
|
||||
use std::process::Stdio;
|
||||
|
||||
let mut child = Command::new("ssh")
|
||||
.args(ssh_args(socket_path))
|
||||
.arg("bash -s")
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.kill_on_drop(true)
|
||||
.spawn()
|
||||
.map_err(|e| anyhow!("Failed to spawn SSH for script: {e}"))?;
|
||||
|
||||
// Write the script to stdin.
|
||||
if let Some(mut stdin) = child.stdin.take() {
|
||||
use futures_lite::io::AsyncWriteExt;
|
||||
stdin
|
||||
.write_all(script.as_bytes())
|
||||
.await
|
||||
.map_err(|e| anyhow!("Failed to write script to stdin: {e}"))?;
|
||||
// Close stdin so the remote bash exits after reading the script.
|
||||
drop(stdin);
|
||||
}
|
||||
|
||||
child
|
||||
.output()
|
||||
.with_timeout(timeout)
|
||||
.await
|
||||
.map_err(|_| anyhow!("Script timed out after {timeout:?}"))?
|
||||
.map_err(|e| anyhow!("Script failed: {e}"))
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
//! Transport abstraction for [`RemoteServerManager`].
|
||||
//!
|
||||
//! Separates SSH-specific concerns (ControlMaster sockets, binary install,
|
||||
//! process spawning) from the transport-agnostic session lifecycle managed
|
||||
//! by [`RemoteServerManager`]. Alternative transports (Docker exec,
|
||||
//! in-process for tests) implement the same trait without touching the
|
||||
//! manager.
|
||||
//!
|
||||
//! Methods are async. Callers use the trait via generics
|
||||
//! (`T: RemoteTransport`) rather than `dyn` dispatch.
|
||||
//!
|
||||
//! [`RemoteServerManager`]: crate::manager::RemoteServerManager
|
||||
use std::future::Future;
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
use std::path::PathBuf;
|
||||
|
||||
use async_channel::Receiver;
|
||||
use warpui::r#async::executor;
|
||||
|
||||
use crate::client::{ClientEvent, RemoteServerClient};
|
||||
use crate::setup::RemotePlatform;
|
||||
|
||||
/// A successful return from [`RemoteTransport::connect`].
|
||||
///
|
||||
/// Bundles the live [`RemoteServerClient`] and its [`ClientEvent`]
|
||||
/// receiver together with any transport-specific resources whose
|
||||
/// lifetime must match the session (notably an owning `Child` for
|
||||
/// subprocess-backed transports). The caller -- typically
|
||||
/// [`RemoteServerManager`] -- stashes the whole `Connection` on its
|
||||
/// per-session state so that dropping the state cleans everything up at
|
||||
/// once.
|
||||
///
|
||||
/// [`RemoteServerManager`]: crate::manager::RemoteServerManager
|
||||
#[cfg_attr(target_family = "wasm", allow(dead_code))]
|
||||
pub struct Connection {
|
||||
pub client: RemoteServerClient,
|
||||
pub event_rx: Receiver<ClientEvent>,
|
||||
/// The subprocess whose stdio backs the client (e.g.
|
||||
/// `ssh … remote-server-proxy`). Spawned with `kill_on_drop(true)`
|
||||
/// by the transport, so dropping this `Child` sends SIGKILL to the
|
||||
/// subprocess. The [`RemoteServerManager`] holds it for the
|
||||
/// lifetime of the session and drops it on teardown.
|
||||
///
|
||||
/// [`RemoteServerManager`]: crate::manager::RemoteServerManager
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
pub child: async_process::Child,
|
||||
/// For transports that multiplex through a local SSH
|
||||
/// `ControlMaster` socket: the path to that socket, used on
|
||||
/// explicit teardown (after the user's shell exits) to run
|
||||
/// `ssh -O exit` and force the master to terminate without
|
||||
/// waiting for half-closed channels. `None` for transports with
|
||||
/// no separate master process (in-process tests, etc.).
|
||||
///
|
||||
/// See [`crate::ssh::stop_control_master`] for the exact command.
|
||||
#[cfg(not(target_family = "wasm"))]
|
||||
pub control_path: Option<PathBuf>,
|
||||
}
|
||||
|
||||
pub trait RemoteTransport: Send + Sync {
|
||||
/// Detects the remote host's OS and architecture by running `uname -sm`.
|
||||
///
|
||||
/// Returns the parsed [`RemotePlatform`] on success, or an error string
|
||||
/// if the command fails or the output cannot be parsed.
|
||||
fn detect_platform(&self) -> impl Future<Output = Result<RemotePlatform, String>> + Send;
|
||||
|
||||
/// Checks whether the remote server binary is present on the remote host.
|
||||
///
|
||||
/// Pure I/O — does not emit any events. The caller
|
||||
/// ([`RemoteServerManager::check_binary`]) is responsible for emitting
|
||||
/// [`SetupStateChanged`] and [`BinaryCheckComplete`].
|
||||
///
|
||||
/// Returns `Ok(true)` if the binary is installed and executable,
|
||||
/// `Ok(false)` if it is definitively not installed, and
|
||||
/// `Err(_)` if the check failed (e.g. SSH timeout/unreachable).
|
||||
fn check_binary(&self) -> impl Future<Output = Result<bool, String>> + Send;
|
||||
|
||||
/// Installs the remote server binary on the remote host.
|
||||
///
|
||||
/// Pure I/O — does not emit any events. The caller
|
||||
/// ([`RemoteServerManager::install_binary`]) is responsible for emitting
|
||||
/// [`SetupStateChanged`] and [`BinaryInstallComplete`].
|
||||
///
|
||||
/// Returns `Ok(())` if the install succeeded, and
|
||||
/// `Err(_)` if the install failed (e.g. SSH timeout, script error).
|
||||
fn install_binary(&self) -> impl Future<Output = Result<(), String>> + Send;
|
||||
|
||||
/// Establish a new connection to the remote server.
|
||||
///
|
||||
/// Called on both the initial connect and every subsequent reconnect
|
||||
/// attempt. Returns a [`Connection`] carrying the live client, its
|
||||
/// event channel, and any transport-specific resources (e.g. an
|
||||
/// owning `Child`) whose lifetime must match the session.
|
||||
///
|
||||
/// The implementation is responsible for any transport-specific setup
|
||||
/// required before messages can flow (e.g. spawning a process, connecting
|
||||
/// a socket). Stderr forwarding to local logging should also happen here.
|
||||
fn connect(
|
||||
&self,
|
||||
executor: &executor::Background,
|
||||
) -> impl Future<Output = anyhow::Result<Connection>> + Send;
|
||||
}
|
||||
Reference in New Issue
Block a user