ive harnessed the harness
Something went wrong. Try again.
21 kB · 552 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553//! The remote endpoint boundary: one trait for reaching a named target, a manager that//! dispatches the four remote `HostRequest` variants to the registered endpoint, and the//! relay endpoint that starts a real worker over SSH for each ensured namespace.//!//! Opaque call dispatch is refused until the target worker's dispatcher is present, and//! nothing here is rerouted to local execution.
use anyhow::{bail, ensure, Context, Result};use std::{collections::BTreeMap, time::Duration};
pub type RemoteFuture = std::pin::Pin<Box<dyn std::future::Future<Output = Result<serde_json::Value>> + Send>>;
pub trait RemoteEndpoint: Send + Sync { fn ensure(&self, session: String, workspace: Option<String>) -> RemoteFuture; fn call( &self, namespace_id: String, namespace_generation: String, operation_id: String, operation: String, payload: serde_json::Value, ) -> RemoteFuture; fn cancel( &self, session: String, execution_id: String, namespace_id: String, namespace_generation: String, operation_id: String, ) -> RemoteFuture; fn close( &self, session: String, execution_id: String, namespace_id: String, namespace_generation: String, ) -> RemoteFuture;
fn shutdown(&self) -> RemoteFuture { Box::pin(async { Ok(serde_json::json!({"state": "closed"})) }) }}
/// Dispatches the four remote HostRequest variants to the registered/// host-owned endpoint. It validates identity-shaped fields and JSON bounds,/// but never interprets the endpoint's operation vocabulary.type Endpoint = std::sync::Arc<dyn RemoteEndpoint>;type Registry<K, V> = std::sync::Arc<std::sync::RwLock<BTreeMap<K, V>>>;
#[derive(Clone, Default)]pub struct RemoteManager { endpoints: Registry<String, Endpoint>, namespaces: Registry<(String, String), (String, Endpoint)>,}
impl RemoteManager { pub fn new() -> Self { Self::default() }
/// Stops every namespace worker and clears all target/binding registries. pub async fn shutdown(&self) -> Result<()> { let endpoints = self .endpoints .read() .map_err(|_| anyhow::anyhow!("remote target registry lock poisoned"))? .values() .cloned() .collect::<Vec<_>>(); for endpoint in endpoints { endpoint.shutdown().await?; } self.namespaces .write() .map_err(|_| anyhow::anyhow!("remote namespace registry lock poisoned"))? .clear(); self.endpoints .write() .map_err(|_| anyhow::anyhow!("remote target registry lock poisoned"))? .clear(); Ok(()) }
pub fn register_target( &self, target: crate::target::TargetId, endpoint: std::sync::Arc<dyn RemoteEndpoint>, ) -> Result<()> { self.endpoints .write() .map_err(|_| anyhow::anyhow!("remote target registry lock poisoned"))? .insert(target.to_string(), endpoint); Ok(()) }
/// Registers the production relay endpoint for a configured target. /// Callers pass the already-selected host backend; a remote placement can /// therefore never be silently demoted to a local adapter. pub fn register_relay_target( &self, target: crate::target::RemoteTarget, backend: std::sync::Arc<dyn crate::worker::WorkerBackend>, startup_timeout: Duration, ) -> Result<std::sync::Arc<RelayEndpoint>> { let endpoint = std::sync::Arc::new(RelayEndpoint::new(target.clone(), backend, startup_timeout)); self.register_target(target.id, endpoint.clone())?; Ok(endpoint) }
pub fn unregister_target(&self, target: &crate::target::TargetId) -> Result<bool> { Ok(self .endpoints .write() .map_err(|_| anyhow::anyhow!("remote target registry lock poisoned"))? .remove(target.as_str()) .is_some()) }
fn endpoint(&self, target: &str) -> Result<std::sync::Arc<dyn RemoteEndpoint>> { self.endpoints .read() .map_err(|_| anyhow::anyhow!("remote target registry lock poisoned"))? .get(target) .cloned() .ok_or_else(|| anyhow::anyhow!("remote target {target:?} is unavailable")) }
/// Handles all remote requests in one production dispatch point. The /// outer session/execution identities are retained by the endpoint call; /// an unavailable target is returned as a typed HostReply error and is /// never rerouted locally. pub async fn handle( &self, session: &crate::protocol::SessionId, execution_id: &crate::protocol::OperationId, request: crate::protocol::HostRequest, ) -> crate::protocol::HostReply { let result = self.handle_result(session, execution_id, request).await; match result { Ok(value) => crate::protocol::HostReply::Ok { value }, Err(error) => { let message = format!("{error:#}"); let code = if message.contains("remote call payload") && message.contains("remote value bound") { "payload_too_large" } else if message.contains("remote operation result") && message.contains("remote value bound") { "result_too_large" } else if message.contains("remote error") && message.contains("remote value bound") { "error_too_large" } else if message.contains("unavailable") { "unknown" } else { "remote" }; crate::protocol::HostReply::error(code, message) } } }
async fn handle_result( &self, session: &crate::protocol::SessionId, execution_id: &crate::protocol::OperationId, request: crate::protocol::HostRequest, ) -> Result<serde_json::Value> { use crate::protocol::HostRequest; match request { HostRequest::RemoteEnsure { target, workspace } => { validate_target_name(&target)?; let endpoint = self.endpoint(&target)?; let value = endpoint.ensure(session.to_string(), workspace).await?; ensure_json_bound(&value, "remote ensure result")?; let object = value .as_object() .ok_or_else(|| anyhow::anyhow!("remote ensure result is not an object"))?; let namespace = object .get("namespace_id") .and_then(serde_json::Value::as_str) .ok_or_else(|| anyhow::anyhow!("remote ensure result has no namespace_id"))?; let generation = object .get("namespace_generation") .and_then(serde_json::Value::as_str) .ok_or_else(|| { anyhow::anyhow!("remote ensure result has no namespace_generation") })?; crate::target::NamespaceId::parse(namespace.to_owned())?; crate::target::NamespaceGeneration::parse(generation.to_owned())?; self.namespaces .write() .map_err(|_| anyhow::anyhow!("remote namespace registry lock poisoned"))? .insert( (namespace.to_owned(), generation.to_owned()), (target, endpoint), ); Ok(value) } HostRequest::RemoteCall { target, namespace_id, namespace_generation, operation_id, operation, payload, } => { validate_target_name(&target)?; validate_namespace_binding(&namespace_id, &namespace_generation, &operation_id)?; ensure_remote_value_bound(&payload, "remote call payload")?; let (bound_target, endpoint) = self.namespace(&namespace_id, &namespace_generation)?; ensure!( bound_target == target, "remote call target does not match namespace binding" ); let value = endpoint .call( namespace_id, namespace_generation, operation_id, operation, payload, ) .await?; ensure_json_bound(&value, "remote operation result")?; Ok(value) } HostRequest::RemoteCancel { namespace_id, namespace_generation, operation_id, } => { validate_namespace_binding(&namespace_id, &namespace_generation, &operation_id)?; let (_, endpoint) = self.namespace(&namespace_id, &namespace_generation)?; let value = endpoint .cancel( session.to_string(), execution_id.to_string(), namespace_id, namespace_generation, operation_id, ) .await?; ensure_json_bound(&value, "remote cancellation result")?; Ok(value) } HostRequest::RemoteClose { namespace_id, namespace_generation, } => { validate_namespace_binding( &namespace_id, &namespace_generation, "close-operation", )?; let key = (namespace_id.clone(), namespace_generation.clone()); let (_, endpoint) = self.namespace(&namespace_id, &namespace_generation)?; let value = endpoint .close( session.to_string(), execution_id.to_string(), namespace_id, namespace_generation, ) .await?; ensure_json_bound(&value, "remote close result")?; self.namespaces .write() .map_err(|_| anyhow::anyhow!("remote namespace registry lock poisoned"))? .remove(&key); Ok(value) } _ => bail!("request is not a remote namespace operation"), } }
fn namespace( &self, namespace: &str, generation: &str, ) -> Result<(String, std::sync::Arc<dyn RemoteEndpoint>)> { self.namespaces .read() .map_err(|_| anyhow::anyhow!("remote namespace registry lock poisoned"))? .get(&(namespace.to_owned(), generation.to_owned())) .cloned() .ok_or_else(|| anyhow::anyhow!("remote namespace binding is unavailable")) }}
/// A production remote namespace endpoint backed by the shared M1 worker/// launch path. The backend is normally `RemoteBackend`, so each ensure starts/// a real worker over SSH with the issued target binding and verified identity./// Opaque call dispatch is deliberately refused until the target worker's/// dispatcher is present; this is not rerouted to local execution.pub struct RelayEndpoint { target: crate::target::RemoteTarget, backend: std::sync::Arc<dyn crate::worker::WorkerBackend>, startup_timeout: Duration, workers: std::sync::Arc<tokio::sync::Mutex<BTreeMap<(String, String), crate::worker::Worker>>>,}
impl RelayEndpoint { pub fn new( target: crate::target::RemoteTarget, backend: std::sync::Arc<dyn crate::worker::WorkerBackend>, startup_timeout: Duration, ) -> Self { Self { target, backend, startup_timeout, workers: std::sync::Arc::new(tokio::sync::Mutex::new(BTreeMap::new())), } }
async fn ensure_inner( &self, session: String, workspace: Option<String>, ) -> Result<serde_json::Value> { let session = crate::protocol::SessionId::try_from(session) .map_err(|error| anyhow::anyhow!(error))?; let workspace = workspace.unwrap_or_else(|| self.target.workspace.path.clone()); ensure!( workspace == self.target.workspace.path, "remote workspace does not match the configured target workspace" ); let namespace_id = crate::target::NamespaceId::parse(format!("ns-{}", uuid::Uuid::new_v4().simple()))?; let namespace_generation = crate::target::NamespaceGeneration::fresh(); let binding = crate::protocol::TargetStamp { target_id: self.target.id.clone(), namespace_id: namespace_id.clone(), namespace_generation: namespace_generation.clone(), }; let worker = crate::worker::Worker::launch_with( self.backend.as_ref(), self.startup_timeout, session, crate::protocol::Role::Remote, None, Some(binding), ) .await .context("launch remote relay namespace")?; let identity = worker .target_identity() .context("remote worker did not retain verified target identity")?; let capabilities = worker .worker_identity() .map(|identity| identity.runtime.capabilities) .unwrap_or_default(); self.workers.lock().await.insert( (namespace_id.to_string(), namespace_generation.to_string()), worker, ); Ok(serde_json::json!({ "target": self.target.id, "workspace": workspace, "namespace_id": namespace_id, "namespace_generation": namespace_generation, "persistence": "connection-scoped", "identity": identity, "capabilities": capabilities, })) }
async fn cancel_inner( &self, namespace_id: String, namespace_generation: String, ) -> Result<serde_json::Value> { let workers = self.workers.lock().await; ensure!( workers.contains_key(&(namespace_id, namespace_generation)), "remote namespace binding is unavailable" ); // The current shared worker API has no per-operation cancellation; // never collapse that limitation into a false confirmation. Ok(serde_json::json!({ "state": "unknown", "reason": "remote worker has no per-operation cancellation boundary" })) }
async fn close_inner( &self, namespace_id: String, namespace_generation: String, ) -> Result<serde_json::Value> { let worker = self .workers .lock() .await .remove(&(namespace_id, namespace_generation)) .ok_or_else(|| anyhow::anyhow!("remote namespace binding is unavailable"))?; worker.stop().await; Ok(serde_json::json!({"state": "closed"})) }}
impl RemoteEndpoint for RelayEndpoint { fn ensure(&self, session: String, workspace: Option<String>) -> RemoteFuture { let this = self.clone_for_task(); Box::pin(async move { this.ensure_inner(session, workspace).await }) }
fn call( &self, namespace_id: String, namespace_generation: String, operation_id: String, operation: String, payload: serde_json::Value, ) -> RemoteFuture { let workers = self.workers.clone(); Box::pin(async move { ensure_remote_value_bound(&payload, "remote call payload")?; let operation_id = crate::protocol::OperationId::try_from(operation_id) .map_err(|error| anyhow::anyhow!(error))?; let worker = workers .lock() .await .get(&(namespace_id, namespace_generation)) .cloned() .ok_or_else(|| anyhow::anyhow!("remote namespace binding is unavailable"))?; let outcome = worker .call( crate::worker::WorkCommand::RemoteDispatch { id: operation_id, operation, payload, }, Duration::from_secs(60), crate::worker::no_host_operations(), ) .await; match outcome { Ok(crate::protocol::Outcome::Remote { value }) => { ensure_remote_value_bound(&value, "remote operation result")?; Ok(serde_json::json!({"status": "ok", "value": value})) } Ok(crate::protocol::Outcome::RemoteError { code, message }) => { Ok(serde_json::json!({ "status": "error", "error": {"kind": code, "message": message} })) } Ok(crate::protocol::Outcome::Unknown { reason }) => Ok(serde_json::json!({ "status": "error", "error": {"kind": "unknown", "reason": reason} })), Ok(other) => Err(anyhow::anyhow!( "remote worker returned non-remote outcome: {other:?}" )), Err(error) => Ok(serde_json::json!({ "status": "error", "error": { "kind": "unknown", "reason": format!("remote dispatch outcome is unknown: {error:#}") } })), } }) }
fn shutdown(&self) -> RemoteFuture { let workers = self.workers.clone(); Box::pin(async move { let workers = { let mut guard = workers.lock().await; std::mem::take(&mut *guard) .into_values() .collect::<Vec<_>>() }; for worker in workers { worker.stop().await; } Ok(serde_json::json!({"state": "closed"})) }) }
fn cancel( &self, _session: String, _execution_id: String, namespace_id: String, namespace_generation: String, _operation_id: String, ) -> RemoteFuture { let this = self.clone_for_task(); Box::pin(async move { this.cancel_inner(namespace_id, namespace_generation).await }) }
fn close( &self, _session: String, _execution_id: String, namespace_id: String, namespace_generation: String, ) -> RemoteFuture { let this = self.clone_for_task(); Box::pin(async move { this.close_inner(namespace_id, namespace_generation).await }) }}
impl RelayEndpoint { fn clone_for_task(&self) -> Self { Self { target: self.target.clone(), backend: self.backend.clone(), startup_timeout: self.startup_timeout, workers: self.workers.clone(), } }}
fn validate_target_name(target: &str) -> Result<()> { crate::target::TargetId::parse(target.to_owned())?; Ok(())}
fn validate_namespace_binding(namespace: &str, generation: &str, operation: &str) -> Result<()> { crate::target::NamespaceId::parse(namespace.to_owned())?; crate::target::NamespaceGeneration::parse(generation.to_owned())?; crate::protocol::OperationId::try_from(operation.to_owned()) .map_err(|error| anyhow::anyhow!(error))?; Ok(())}
fn ensure_json_bound(value: &serde_json::Value, label: &str) -> Result<()> { ensure!( serde_json::to_vec(value)?.len() <= crate::protocol::MAX_FRAME, "{label} exceeds the {}-byte frame bound", crate::protocol::MAX_FRAME ); Ok(())}
fn ensure_remote_value_bound(value: &serde_json::Value, label: &str) -> Result<()> { ensure!( serde_json::to_vec(value)?.len() <= crate::protocol::MAX_REMOTE_VALUE, "{label} exceeds the {}-byte remote value bound", crate::protocol::MAX_REMOTE_VALUE ); Ok(())}