diff --git a/cli/Cargo.toml b/cli/Cargo.toml index e939680aae6..5cadf7ae6c4 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -67,7 +67,7 @@ winresource = "0.1" [target.'cfg(windows)'.dependencies] winreg = "0.56" winapi = "0.3.9" -windows-sys = { version = "0.61", features = ["Win32_System_Console", "Win32_UI_Input_KeyboardAndMouse"] } +windows-sys = { version = "0.61", features = ["Win32_Foundation", "Win32_Security", "Win32_Security_Authorization", "Win32_System_Console", "Win32_System_Threading", "Win32_UI_Input_KeyboardAndMouse"] } [target.'cfg(target_os = "macos")'.dependencies] core-foundation = "0.9.3" diff --git a/cli/src/bin/code/main.rs b/cli/src/bin/code/main.rs index 8a9c38f3afd..572d6f57de1 100644 --- a/cli/src/bin/code/main.rs +++ b/cli/src/bin/code/main.rs @@ -114,7 +114,9 @@ async fn main() -> Result<(), std::convert::Infallible> { Some(args::AgentSubcommand::Stop(stop_args)) => { agent_stop::agent_stop(context!(), stop_args).await } - Some(args::AgentSubcommand::Kill) => agent_kill::agent_kill(context!()).await, + Some(args::AgentSubcommand::Kill(kill_args)) => { + agent_kill::agent_kill(context!(), kill_args).await + } Some(args::AgentSubcommand::Logs(logs_args)) => { agent_logs::agent_logs(context!(), logs_args).await } diff --git a/cli/src/commands.rs b/cli/src/commands.rs index eeb8fc53336..14fce7cded6 100644 --- a/cli/src/commands.rs +++ b/cli/src/commands.rs @@ -6,6 +6,7 @@ mod context; pub mod agent; +pub mod agent_discovery; pub mod agent_host; pub mod agent_kill; pub mod agent_logs; diff --git a/cli/src/commands/agent.rs b/cli/src/commands/agent.rs index 701c56b3f56..5725fbf8c51 100644 --- a/cli/src/commands/agent.rs +++ b/cli/src/commands/agent.rs @@ -3,37 +3,46 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -use std::fs; +use std::sync::LazyLock; use ahp::{Client, Transport, TransportError, TransportMessage}; use ahp_types::commands::{AuthenticateParams, AuthenticateResult}; use ahp_types::errors::ahp_error_codes; use ahp_types::state::ProtectedResourceMetadata; -use ahp_types::{ROOT_RESOURCE_URI, PROTOCOL_VERSION}; +use ahp_types::{PROTOCOL_VERSION, ROOT_RESOURCE_URI}; use futures::{SinkExt, StreamExt}; use tokio::io::{AsyncRead, AsyncWrite}; +use tokio::sync::Mutex as AsyncMutex; use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::{connect_async, WebSocketStream}; +use crate::async_pipe::get_socket_rw_stream; use crate::auth::{Auth, AuthProvider}; use crate::constants::AGENT_HOST_PORT; use crate::log; +use crate::tunnels::agent_host_registry::{AgentHostEndpointAddress, AgentHostEndpointMetadata}; use crate::tunnels::dev_tunnels::DevTunnels; -use crate::util::errors::{wrap, AnyError, CodeError}; -use crate::util::machine::process_exists; +use crate::util::errors::{wrap, AnyError}; use super::CommandContext; -use crate::tunnels::agent_host_metadata::AgentHostMetadata; -/// Connects to an agent host, initializes the AHP session, and returns -/// the ready-to-use client. If an explicit `address` is given it is used -/// directly; if `tunnel_name` is given, the tunnel is looked up via the -/// dev tunnels API; otherwise the lockfile written by `code agent host` -/// is read to discover the local instance. +/// Connects to an agent host at an explicit `--address` or `--tunnel` +/// target, initializes the AHP session, and returns the ready-to-use +/// client. If `address` is given it is used directly; otherwise +/// `tunnel_name` is looked up via the dev tunnels API. +/// +/// This is deliberately explicit-target-only: automatic discovery of a +/// local standalone/editor instance lives in +/// [`super::agent_discovery`] instead (see +/// [`super::agent_discovery::discover_live_endpoints`] and +/// [`super::agent_discovery::connect_to_session_host`]). Every call site +/// already branches to one of those before falling through here, so +/// passing neither `address` nor `tunnel_name` is a caller bug, not a +/// "no target given" case to recover from silently. /// /// The returned client has been initialized but **not** authenticated. /// Use [`request_with_auth`] to issue commands that may require auth. -pub async fn connect( +pub async fn connect_explicit( ctx: &CommandContext, address: Option<&str>, tunnel_name: Option<&str>, @@ -41,18 +50,57 @@ pub async fn connect( let client = match (address, tunnel_name) { (Some(addr), _) => connect_ws(addr).await?, (None, Some(name)) => connect_via_tunnel(ctx, name).await?, - (None, None) => { - let addr = resolve_address_from_lockfile(ctx)?; - connect_ws(&addr).await? + (None, None) => unreachable!( + "connect_explicit requires --address or --tunnel; callers must route \ + the no-target case through agent_discovery instead" + ), + }; + + initialize_client(&client).await?; + Ok(client) +} + +/// Connects to a specific registry-discovered endpoint (used by +/// multi-host auto-discovery in `ps`/`logs`/`stop`/`kill`). Dispatches to +/// a `tcp` WebSocket connection or a `socket`/named-pipe raw-stream +/// WebSocket handshake depending on the endpoint's address kind, and +/// always includes the endpoint's connection token as the `tkn` query +/// parameter, matching the wire convention documented in +/// `LOCAL_ENDPOINT.md`. +/// +/// Like [`connect_explicit`], the returned client is initialized but not yet +/// authenticated; use [`request_with_auth`] for calls that may need it. +pub async fn connect_to_endpoint(endpoint: &AgentHostEndpointMetadata) -> Result { + let client = match &endpoint.endpoint { + AgentHostEndpointAddress::Tcp { host, port } => { + let dial_host = crate::commands::agent_host::dial_host(Some(host)); + let mut url = format!("ws://{dial_host}:{port}/"); + if !endpoint.connection_token.is_empty() { + url.push_str(&format!("?tkn={}", endpoint.connection_token)); + } + connect_ws(&url).await? + } + AgentHostEndpointAddress::Socket { path } => { + connect_ws_over_socket(path, &endpoint.connection_token).await? } }; + initialize_client(&client).await?; + Ok(client) +} + +/// Shared final step of establishing an AHP session: performs the +/// protocol `initialize` handshake common to every connection kind. +async fn initialize_client(client: &Client) -> Result<(), AnyError> { client - .initialize("code-cli".into(), vec![PROTOCOL_VERSION.to_string()], vec![]) + .initialize( + "code-cli".into(), + vec![PROTOCOL_VERSION.to_string()], + vec![], + ) .await .map_err(|e| wrap(e, "AHP initialize failed"))?; - - Ok(client) + Ok(()) } /// Opens a WebSocket connection and creates an AHP client. @@ -68,6 +116,37 @@ async fn connect_ws(address: &str) -> Result { .map_err(|e| wrap(e, "Failed to establish AHP session").into()) } +/// Opens a raw-stream WebSocket connection over a Unix domain socket +/// (Unix) or named pipe (Windows) using the existing [`async_pipe`] +/// abstraction, mirroring the tunnel raw-stream handshake in +/// [`connect_via_tunnel`]. This is how the CLI reaches `editor`-owned +/// registry endpoints, which are never TCP. +async fn connect_ws_over_socket(path: &str, connection_token: &str) -> Result { + let pipe = get_socket_rw_stream(std::path::Path::new(path)) + .await + .map_err(|e| { + wrap( + e, + format!("Failed to connect to agent host socket at {path}"), + ) + })?; + + let mut url = "ws://localhost/".to_string(); + if !connection_token.is_empty() { + url.push_str(&format!("?tkn={connection_token}")); + } + + let (ws_stream, _) = tokio_tungstenite::client_async(url, pipe) + .await + .map_err(|e| wrap(e, format!("WebSocket handshake over socket {path} failed")))?; + + let transport = WsTransport::new(ws_stream, ()); + + Client::connect(transport, ahp::ClientConfig::default()) + .await + .map_err(|e| wrap(e, "Failed to establish AHP session over socket").into()) +} + /// A [`Transport`] backed by a `tokio-tungstenite` WebSocket stream. /// /// `_guard` keeps an auxiliary resource alive for the lifetime of the @@ -81,7 +160,10 @@ struct WsTransport { impl WsTransport { fn new(inner: WebSocketStream, guard: G) -> Self { - Self { inner, _guard: guard } + Self { + inner, + _guard: guard, + } } } @@ -153,6 +235,34 @@ async fn connect_via_tunnel(ctx: &CommandContext, name: &str) -> Result> = LazyLock::new(|| AsyncMutex::new(())); + +/// Runs `authenticate` while holding [`AUTH_SERIALIZE`], releasing the +/// guard as soon as `authenticate` completes — *before* returning to the +/// caller. This is split out from [`request_with_auth`] specifically so +/// the guard is never accidentally held across the retried RPC that +/// follows: that RPC can be slow (or the host can be hung/unreachable), +/// and holding the lock across it would stall unrelated hosts' own +/// authentication attempts for no reason. Also split out so this +/// locking behavior can be exercised directly in tests without needing +/// a live AHP client. +async fn serialize_auth(authenticate: F) -> T +where + F: FnOnce() -> Fut, + Fut: std::future::Future, +{ + let _guard = AUTH_SERIALIZE.lock().await; + authenticate().await +} + /// Issues a JSON-RPC request, automatically handling `-32007` auth errors /// by running the device-flow login and retrying once. pub async fn request_with_auth( @@ -172,7 +282,10 @@ where ctx.log, "Server requires authentication, starting login flow..." ); - authenticate_from_error(ctx, client, e).await?; + serialize_auth(|| authenticate_from_error(ctx, client, e)).await?; + // Deliberately outside `serialize_auth`'s guard: this retry + // can be slow or hang if the host is unreachable, and must + // not block other hosts' concurrent authentication. client .request::(method, params) .await @@ -266,31 +379,75 @@ async fn authenticate_from_error( Ok(()) } -fn resolve_address_from_lockfile(ctx: &CommandContext) -> Result { - let lockfile_path = ctx.paths.agent_host_lockfile(); +#[cfg(test)] +mod tests { + use super::*; + use crate::async_pipe::{get_socket_name, listen_socket_rw_stream}; - let data = fs::read_to_string(&lockfile_path).map_err(|e| { - wrap( - e, - "No running agent host found. Start one with `code agent host` or specify --address", - ) - })?; + /// Exercises the real Unix-socket/named-pipe raw-stream WebSocket + /// handshake used to reach `editor`-owned (and socket-based + /// standalone) registry endpoints, using the same cross-platform + /// [`async_pipe`] listener/accept pattern already established for + /// this abstraction elsewhere in the crate (see + /// `tunnels::agent_host` tests), rather than mocking the transport. + #[tokio::test] + async fn connect_ws_over_socket_completes_handshake_over_a_real_socket() { + let path = get_socket_name(); + let mut listener = listen_socket_rw_stream(&path).await.unwrap(); - let metadata: AgentHostMetadata = serde_json::from_str(&data).map_err(|e| { - wrap( - e, - format!("Corrupt agent host lockfile at {}", lockfile_path.display()), - ) - })?; + let server = tokio::spawn(async move { + let pipe = listener.accept().await.unwrap(); + // Accepting the WS handshake server-side is enough to prove + // the client's raw-stream `client_async` handshake (with the + // `?tkn=` connection token in the URL) completes correctly + // against a real listener; `Client::connect` itself never + // blocks on server behavior past that. + let _ws = tokio_tungstenite::accept_async(pipe).await.unwrap(); + }); - if !process_exists(metadata.pid) { - let _ = fs::remove_file(&lockfile_path); - return Err(CodeError::NoRunningAgentHost.into()); + let client = connect_ws_over_socket(path.to_str().unwrap(), "test-token").await; + assert!(client.is_ok(), "expected Ok, got {:?}", client.err()); + + server.await.unwrap(); + + #[cfg(unix)] + let _ = std::fs::remove_file(&path); } - let mut url = format!("ws://127.0.0.1:{}/", metadata.port); - if let Some(token) = &metadata.connection_token { - url.push_str(&format!("?tkn={token}")); + /// Regression test for the `AUTH_SERIALIZE` scoping fix: the guard + /// must be released as soon as `authenticate` finishes, *not* held + /// across whatever the caller does afterwards (in production, the + /// retried RPC). Otherwise a slow/hung host's post-auth retry could + /// stall an unrelated host's own authentication attempt. We can't + /// easily stand up two real AHP `Client`s here, so this exercises + /// `serialize_auth` directly: task A's simulated "retry" (a long + /// sleep performed *after* `serialize_auth` returns) must not block + /// task B's concurrent `serialize_auth` call. + #[tokio::test] + async fn serialize_auth_releases_guard_before_caller_retries() { + let host_a = tokio::spawn(async { + serialize_auth(|| async {}).await; + // Simulates a slow/hung host's retried RPC, which happens + // *after* the guard should already have been released. + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + }); + + // Give host A a head start so it acquires the guard first. + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + + // If the guard were (incorrectly) held across host A's simulated + // retry, this would time out; with the fix it resolves almost + // immediately since the guard was already released. + let host_b = tokio::time::timeout(std::time::Duration::from_millis(150), async { + serialize_auth(|| async {}).await; + }) + .await; + + assert!( + host_b.is_ok(), + "host B's authenticate must not wait for host A's retry" + ); + + host_a.await.unwrap(); } - Ok(url) } diff --git a/cli/src/commands/agent_discovery.rs b/cli/src/commands/agent_discovery.rs new file mode 100644 index 00000000000..fb4db42847a --- /dev/null +++ b/cli/src/commands/agent_discovery.rs @@ -0,0 +1,315 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +//! Shared multi-host discovery helpers for `code agent ps|logs|stop`. +//! +//! `code agent host` supervisors (standalone) and running VS Code windows +//! (editor) each publish a live entry to the shared local agent-host +//! endpoint registry (schema v2; see +//! [`crate::tunnels::agent_host_registry`]). When the user doesn't pin a +//! specific `--address`/`--tunnel`, these commands should consider *every* +//! live entry a candidate host, not just one arbitrarily selected one — +//! that's what this module coordinates. + +use ahp::Client; +use ahp_types::commands::{ListSessionsParams, ListSessionsResult}; +use ahp_types::ROOT_RESOURCE_URI; +use futures::stream::FuturesUnordered; +use futures::StreamExt; + +use crate::log; +use crate::tunnels::agent_host_registry::{self, AgentHostEndpointMetadata}; +use crate::tunnels::user_data_path::resolve_user_data_path; +use crate::util::errors::{AnyError, CodeError}; + +use super::agent; +use super::CommandContext; + +/// Enumerates every live schema-v2 registry endpoint (both `editor` and +/// `standalone`, both socket/pipe and TCP addresses) as auto-discovery +/// candidates. The shared registry is the sole source of truth for +/// automatic discovery; callers should treat an empty `Vec` as "no agent +/// host is currently running" ([`CodeError::NoRunningAgentHost`]) rather +/// than falling back to any other discovery mechanism. +pub fn discover_live_endpoints( + ctx: &CommandContext, + user_data_dir: Option<&str>, +) -> Vec { + let user_data_path = resolve_user_data_path(user_data_dir); + trace!( + ctx.log, + "discovering live agent hosts in {}", + user_data_path.display() + ); + agent_host_registry::list_live_endpoints(&ctx.log, &user_data_path) +} + +/// Outcome of probing a single host for a given session, as consumed by +/// [`resolve_first_match`]. Generic over the "found" payload `T` (rather +/// than hard-coding a live [`Client`]) purely so the early-return/ +/// cancellation control flow in [`resolve_first_match`] can be exercised +/// directly in tests with cheap, deterministic fake futures instead of a +/// real AHP connection. +enum ProbeResult { + /// Connected and queried successfully; the session was present. The + /// payload is whatever the caller wants back for the matching host + /// (a ready-to-use [`Client`] in production). + Found(T), + /// Connected and queried successfully; the session was not present. + NotFound, + /// The host could not be connected to or queried at all (network + /// error, auth failure, protocol error, ...). Nothing can be + /// concluded about whether it owns the session; its label is kept + /// for diagnostics. + Unreachable(String), +} + +/// Failure produced by [`resolve_first_match`] when no host matched. +#[derive(Debug, Clone, PartialEq, Eq)] +enum SearchFailure { + /// Every host was searched successfully and none owns the session: + /// safe to report "session not found". + NotFound, + /// At least one host could not be searched (its label recorded + /// here), and no host that *could* be searched owned the session. We + /// must not claim the session doesn't exist — it might live on the + /// host(s) we couldn't reach. + Incomplete(Vec), +} + +/// Consumes host-probe results as they complete, returning the payload +/// from the first match *immediately*: dropping `probes` (which happens +/// as soon as this function returns) cancels any probes still in flight, +/// so a hung or unreachable host can never delay `logs`/`stop` once a +/// match is found on a faster host. A `Found` anywhere takes precedence +/// over `Unreachable` elsewhere: once we have a definitive match we +/// don't need to care that some other host was unreachable. When +/// nothing matches, returns [`SearchFailure::NotFound`] if every host +/// was searched successfully, or [`SearchFailure::Incomplete`] (carrying +/// the unreachable hosts' labels) if at least one host could not be +/// searched. +async fn resolve_first_match(mut probes: FuturesUnordered) -> Result +where + F: std::future::Future>, +{ + let mut unreachable_labels = Vec::new(); + + while let Some(probe) = probes.next().await { + match probe { + ProbeResult::Found(payload) => return Ok(payload), + ProbeResult::NotFound => {} + ProbeResult::Unreachable(label) => unreachable_labels.push(label), + } + } + + if unreachable_labels.is_empty() { + Err(SearchFailure::NotFound) + } else { + Err(SearchFailure::Incomplete(unreachable_labels)) + } +} + +/// Connects to `endpoint` and checks whether `session` appears in its +/// `listSessions` catalog — membership, not error-text sniffing, is the +/// source of truth per requirement. On success but no match, or on any +/// failure, the client (if any) is shut down before returning so callers +/// never need to remember to clean up non-matching connections. +async fn probe_host( + ctx: &CommandContext, + endpoint: AgentHostEndpointMetadata, + session: &str, +) -> ProbeResult { + let label = endpoint.label(); + + let client = match agent::connect_to_endpoint(&endpoint).await { + Ok(client) => client, + Err(e) => { + warning!(ctx.log, "Could not connect to agent host {}: {}", label, e); + return ProbeResult::Unreachable(label); + } + }; + + let listed: Result = agent::request_with_auth( + ctx, + &client, + "listSessions", + ListSessionsParams { + channel: ROOT_RESOURCE_URI.to_string(), + filter: None, + }, + ) + .await; + + match listed { + Ok(r) if r.items.iter().any(|s| s.resource == session) => ProbeResult::Found(client), + Ok(_) => { + client.shutdown().await; + ProbeResult::NotFound + } + Err(e) => { + warning!( + ctx.log, + "Could not search agent host {} for session {}: {}", + label, + session, + e + ); + client.shutdown().await; + ProbeResult::Unreachable(label) + } + } +} + +/// Shared auto-discovery entry point for session-targeted commands +/// (`logs`, `stop`). Searches every live registry endpoint concurrently +/// for `session`, preferring `listSessions` catalog membership over +/// error-text sniffing, and returns a ready-to-use client for whichever +/// host owns it as soon as that host responds — it does not wait for +/// slower or hung hosts once a match is found. Unused/non-matching +/// clients are shut down before this returns; any probes still in +/// flight when a match is found are cancelled (never leaked) by simply +/// dropping them. +/// +/// Returns [`CodeError::NoRunningAgentHost`] when the shared registry has +/// no live endpoints at all — the sole source of truth for automatic +/// discovery. +/// +/// Callers with an explicit `--address`/`--tunnel` should call +/// [`agent::connect_explicit`] directly instead of this helper; this +/// path is only for auto-discovery. +pub async fn connect_to_session_host( + ctx: &CommandContext, + session: &str, + user_data_dir: Option<&str>, +) -> Result { + let endpoints = discover_live_endpoints(ctx, user_data_dir); + + if endpoints.is_empty() { + return Err(CodeError::NoRunningAgentHost.into()); + } + + let probes = FuturesUnordered::new(); + for endpoint in endpoints { + probes.push(probe_host(ctx, endpoint, session)); + } + + match resolve_first_match(probes).await { + Ok(client) => Ok(client), + Err(SearchFailure::NotFound) => { + Err(CodeError::SessionNotFoundOnAnyHost(session.to_string()).into()) + } + Err(SearchFailure::Incomplete(unreachable_labels)) => Err( + CodeError::IncompleteSessionSearch(session.to_string(), unreachable_labels.join(", ")) + .into(), + ), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A fake host probe used by [`resolve_first_match`] tests: same + /// concrete future type regardless of arguments (required to share a + /// single [`FuturesUnordered`]), optionally sleeping first to + /// simulate a slow/hung host, and optionally flipping `completed` + /// so a test can assert it was never driven to completion. + async fn fake_probe( + delay: Option, + result: ProbeResult, + completed: Option>, + ) -> ProbeResult { + if let Some(delay) = delay { + tokio::time::sleep(delay).await; + } + if let Some(completed) = completed { + completed.store(true, std::sync::atomic::Ordering::SeqCst); + } + result + } + + #[tokio::test] + async fn resolve_first_match_returns_immediately_and_cancels_remaining_probes() { + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::Arc; + use std::time::Duration; + + let hung_probe_completed = Arc::new(AtomicBool::new(false)); + + let probes = FuturesUnordered::new(); + probes.push(fake_probe(None, ProbeResult::Found(42), None)); + probes.push(fake_probe( + // Long enough that if it were ever awaited to completion the + // test would time out; the fix must never wait on it. + Some(Duration::from_secs(3600)), + ProbeResult::Unreachable("slow-host".to_string()), + Some(hung_probe_completed.clone()), + )); + + let outcome = tokio::time::timeout(Duration::from_millis(200), resolve_first_match(probes)) + .await + .expect("resolve_first_match must not wait on the hung probe") + .expect("expected the fast host's match"); + + assert_eq!(outcome, 42); + assert!( + !hung_probe_completed.load(Ordering::SeqCst), + "the hung probe's future must be dropped/cancelled, not run to completion" + ); + } + + #[tokio::test] + async fn resolve_first_match_prefers_found_even_when_seen_after_unreachable() { + use std::time::Duration; + + // The first probe to complete is Unreachable; a slightly slower + // probe later resolves with a match. A `Found` anywhere must + // still win over an already-observed `Unreachable`. + let probes = FuturesUnordered::new(); + probes.push(fake_probe( + None, + ProbeResult::Unreachable("host-a".to_string()), + None, + )); + probes.push(fake_probe( + Some(Duration::from_millis(10)), + ProbeResult::Found(7), + None, + )); + + let outcome = resolve_first_match(probes) + .await + .expect("expected the delayed match to win"); + assert_eq!(outcome, 7); + } + + #[tokio::test] + async fn resolve_first_match_is_not_found_when_every_host_was_searched() { + let probes = FuturesUnordered::new(); + probes.push(fake_probe(None, ProbeResult::NotFound, None)); + probes.push(fake_probe(None, ProbeResult::NotFound, None)); + + let err = resolve_first_match::(probes) + .await + .expect_err("no host matched, so this should be an Err"); + assert_eq!(err, SearchFailure::NotFound); + } + + #[tokio::test] + async fn resolve_first_match_is_incomplete_when_a_host_could_not_be_searched() { + let probes = FuturesUnordered::new(); + probes.push(fake_probe(None, ProbeResult::NotFound, None)); + probes.push(fake_probe( + None, + ProbeResult::Unreachable("host-b".to_string()), + None, + )); + + let err = resolve_first_match::(probes) + .await + .expect_err("no host matched, so this should be an Err"); + assert_eq!(err, SearchFailure::Incomplete(vec!["host-b".to_string()])); + } +} diff --git a/cli/src/commands/agent_host.rs b/cli/src/commands/agent_host.rs index 362277834e5..deb9e67522f 100644 --- a/cli/src/commands/agent_host.rs +++ b/cli/src/commands/agent_host.rs @@ -17,13 +17,14 @@ use crate::constants::{self, AGENT_HOST_PORT}; use crate::log; use crate::state::LauncherPaths; use crate::tunnels::agent_host::{ - classify_agent_host_lockfile, AgentHostConfig, AgentHostLockfileDecision, AgentHostManager, + classify_agent_host, AgentHostConfig, AgentHostManager, AgentHostReuseDecision, AgentHostSidecar, LoopbackAuth, }; -use crate::tunnels::agent_host_metadata::remove_agent_host_metadata; +use crate::tunnels::agent_host_registry::{self, AgentHostEndpointIdentity, AgentHostServerType}; use crate::tunnels::code_server::CodeServerArgs; use crate::tunnels::dev_tunnels::DevTunnels; use crate::tunnels::shutdown_signal::ShutdownRequest; +use crate::tunnels::user_data_path::resolve_user_data_path; use crate::update_service::Platform; use crate::util::command::{kill_tree, DetachFromParent}; use crate::util::errors::{wrap, AnyError, CodeError}; @@ -36,13 +37,14 @@ use super::tunnels::fulfill_existing_tunnel_args; use super::CommandContext; /// Internal env var that flips `code agent host` into supervisor mode: -/// the body that actually binds the TCP listener, writes the lockfile, -/// owns the proxy sidecar, and manages the AH backend's lifecycle. The -/// foreground `code agent host` invocation re-execs itself detached with -/// this variable set so the supervisor outlives the user's terminal. +/// the body that actually binds the TCP listener, publishes the +/// supervisor's registry entry, owns the proxy sidecar, and manages the AH +/// backend's lifecycle. The foreground `code agent host` invocation +/// re-execs itself detached with this variable set so the supervisor +/// outlives the user's terminal. const SUPERVISOR_ENV: &str = "VSCODE_AGENT_HOST_SUPERVISOR"; /// Single-line sentinel the supervisor prints once the listener is bound, -/// the lockfile is written, and the banner has been flushed. The +/// its registry entry is published, and the banner has been flushed. The /// foreground process watches for this on the supervisor's stdout, then /// either exits (`--detach`) or starts forwarding output. const SUPERVISOR_READY_LINE: &str = "__VSCODE_AGENT_HOST_READY__"; @@ -52,16 +54,17 @@ const SUPERVISOR_READY_TIMEOUT: Duration = Duration::from_secs(5 * 60); /// Runs the `code agent host` command. Acts in one of two modes: /// -/// * **Foreground** (the default): classifies the canonical lockfile and -/// either prints info about the live supervisor (`Reuse`), or daemonizes -/// a new supervisor child (`SpawnFresh`) and either exits (`--detach`) -/// or follows the supervisor's stdout until Ctrl-C. +/// * **Foreground** (the default): consults the shared local agent-host +/// endpoint registry and either prints info about the live supervisor +/// (`Reuse`), or daemonizes a new supervisor child (`SpawnFresh`) and +/// either exits (`--detach`) or follows the supervisor's stdout until +/// Ctrl-C. /// /// * **Supervisor** (when [`SUPERVISOR_ENV`] is set): binds the public TCP -/// listener, writes the lockfile recording this process's PID + port, -/// runs the proxy accept loop, and manages the underlying VS Code server -/// as a regular child process so the supervisor can kill+respawn it on -/// update. +/// listener, publishes a registry entry recording this process's PID + +/// port, runs the proxy accept loop, and manages the underlying VS Code +/// server as a regular child process so the supervisor can kill+respawn +/// it on update. pub async fn agent_host(ctx: CommandContext, args: AgentHostArgs) -> Result { if std::env::var_os(SUPERVISOR_ENV).is_some() { return run_supervisor(ctx, args).await; @@ -74,16 +77,17 @@ pub async fn agent_host(ctx: CommandContext, args: AgentHostArgs) -> Result Result { let started = Instant::now(); - let lockfile_path = ctx.paths.agent_host_lockfile(); + let user_data_path = resolve_user_data_path(args.user_data_dir.as_deref()); - let decision = classify_agent_host_lockfile(&ctx.log, &lockfile_path); + let decision = classify_agent_host(&ctx.log, &user_data_path); - if let AgentHostLockfileDecision::Reuse { + if let AgentHostReuseDecision::Reuse { pid, host, port, token, tunnel_name, + instance_id, } = &decision { // User asked to replace explicitly: kill + spawn fresh, regardless @@ -95,7 +99,7 @@ async fn run_foreground(ctx: CommandContext, args: AgentHostArgs) -> Result Result Result { let started = Instant::now(); + let user_data_path = resolve_user_data_path(args.user_data_dir.as_deref()); + let instance_id = uuid::Uuid::new_v4().to_string(); // Attach a file log sink before anything else, so download progress, // AH child crashes, update-loop errors, and post-handoff diagnostics @@ -257,7 +263,8 @@ async fn run_supervisor(mut ctx: CommandContext, mut args: AgentHostArgs) -> Res args.host.clone(), loopback_auth, tunnel_name.clone(), - ctx.paths.agent_host_lockfile(), + user_data_path.clone(), + instance_id.clone(), ) .await?; let bound_port = sidecar.bound_addr().port(); @@ -359,10 +366,10 @@ fn print_reuse_banner( if let (Some(base), Some(name)) = (constants::EDITOR_WEB_URL, tunnel_name) { output::print_banner_line("Tunnel", &format!("{base}/agents/tunnel/{name}")); } - // Surface the host the supervisor was actually bound to (older - // lockfiles omit it; fall back to loopback). This lets the network - // hint correctly say "use --host to expose" only when the supervisor - // really is loopback-only. + // Surface the host the supervisor was actually bound to (falling back + // to loopback if unknown). This lets the network hint correctly say + // "use --host to expose" only when the supervisor really is + // loopback-only. let banner_listen_ip = host .and_then(|h| h.parse::().ok()) .unwrap_or(std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST)); @@ -377,43 +384,44 @@ fn print_reuse_banner( } /// Compare the user's requested supervisor configuration with what's -/// recorded in the lockfile. Returns a short human description of the -/// first conflict found (e.g. `"--host 0.0.0.0 conflicts with the -/// running supervisor (bound to 127.0.0.1)"`), or `None` when the -/// requested config is compatible with sharing the existing supervisor. +/// recorded for the running supervisor's registry entry. Returns a short +/// human description of the first conflict found (e.g. `"--host 0.0.0.0 +/// conflicts with the running supervisor (bound to 127.0.0.1)"`), or +/// `None` when the requested config is compatible with sharing the +/// existing supervisor. /// -/// `lockfile_host` may be `None` when the lockfile was written by an -/// older CLI; in that case we conservatively skip the host comparison -/// (the supervisor is most likely loopback, which is the default). +/// `running_host` may be `None` when the registry entry doesn't record a +/// host; in that case we conservatively skip the host comparison (the +/// supervisor is most likely loopback, which is the default). fn detect_config_conflict( args: &AgentHostArgs, - lockfile_host: Option<&str>, - lockfile_port: u16, - lockfile_token: Option<&str>, - lockfile_tunnel: Option<&str>, + running_host: Option<&str>, + running_port: u16, + running_token: Option<&str>, + running_tunnel: Option<&str>, ) -> Option { - if let (Some(requested), Some(running)) = (args.host.as_deref(), lockfile_host) { + if let (Some(requested), Some(running)) = (args.host.as_deref(), running_host) { if requested != running { return Some(format!( "--host {requested} conflicts with the running supervisor (bound to {running})" )); } } - if args.port != 0 && args.port != lockfile_port { + if args.port != 0 && args.port != running_port { return Some(format!( "--port {requested} conflicts with the running supervisor (bound to {running})", requested = args.port, - running = lockfile_port, + running = running_port, )); } - if args.without_connection_token && lockfile_token.is_some() { + if args.without_connection_token && running_token.is_some() { return Some( "--without-connection-token conflicts with the running supervisor (uses a token)" .to_string(), ); } if let Some(requested) = args.connection_token.as_deref() { - match lockfile_token { + match running_token { None => { return Some( "--connection-token conflicts with the running supervisor (no token configured)" @@ -428,7 +436,7 @@ fn detect_config_conflict( Some(_) => {} } } - if args.tunnel && lockfile_tunnel.is_none() { + if args.tunnel && running_tunnel.is_none() { return Some( "--tunnel conflicts with the running supervisor (not exposed via a tunnel)".to_string(), ); @@ -436,9 +444,16 @@ fn detect_config_conflict( None } -/// Kill the existing supervisor process tree and drop the lockfile so the -/// subsequent supervisor start writes a clean record. -async fn replace_existing(log: &log::Logger, lockfile: &Path, pid: u32) -> Result<(), AnyError> { +/// Kill the existing supervisor process tree and remove its exact +/// `(standalone, pid, instanceId)` entry from the shared local agent-host +/// endpoint registry, so the subsequent supervisor start publishes a +/// clean one. +async fn replace_existing( + log: &log::Logger, + user_data_path: &Path, + pid: u32, + instance_id: String, +) -> Result<(), AnyError> { if let Err(e) = kill_tree(pid).await { warning!( log, @@ -447,7 +462,12 @@ async fn replace_existing(log: &log::Logger, lockfile: &Path, pid: u32) -> Resul e ); } - let _ = remove_agent_host_metadata(lockfile); + let identity = AgentHostEndpointIdentity { + server_type: AgentHostServerType::Standalone, + pid, + instance_id, + }; + agent_host_registry::remove_agent_host_endpoint(log, user_data_path, &identity); Ok(()) } @@ -522,14 +542,14 @@ pub async fn ensure_supervisor_running( launcher_paths: &LauncherPaths, log: &log::Logger, ) -> Result { - let lockfile_path = launcher_paths.agent_host_lockfile(); - if let AgentHostLockfileDecision::Reuse { + let user_data_path = resolve_user_data_path(None); + if let AgentHostReuseDecision::Reuse { pid, host, port, token, .. - } = classify_agent_host_lockfile(log, &lockfile_path) + } = classify_agent_host(log, &user_data_path) { return Ok(ActiveAgentHost { pid, @@ -594,8 +614,8 @@ pub async fn ensure_supervisor_running( } } - match classify_agent_host_lockfile(log, &lockfile_path) { - AgentHostLockfileDecision::Reuse { + match classify_agent_host(log, &user_data_path) { + AgentHostReuseDecision::Reuse { pid, host, port, @@ -607,22 +627,25 @@ pub async fn ensure_supervisor_running( port, token, }), - AgentHostLockfileDecision::SpawnFresh => Err(CodeError::CouldNotListenOnInterface( - std::io::Error::other("agent host supervisor signalled ready but lockfile is missing"), - ) - .into()), + AgentHostReuseDecision::SpawnFresh => { + Err(CodeError::CouldNotListenOnInterface(std::io::Error::other( + "agent host supervisor signalled ready but its registry entry is missing", + )) + .into()) + } } } -/// Endpoint of a running agent host supervisor, as recorded in the -/// lockfile and consumed by tunnel + bridge callers. +/// Endpoint of a running agent host supervisor, as recorded in the shared +/// local agent-host endpoint registry and consumed by tunnel + bridge +/// callers. pub struct ActiveAgentHost { pub pid: u32, /// Host the supervisor was bound to (e.g. `"0.0.0.0"`, `"::1"`, - /// `"localhost"`, a specific IP). `None` for lockfiles written by - /// older CLIs. Consumers should pair this with [`dial_host`] to - /// pick the right loopback target when the supervisor was bound to - /// a wildcard. + /// `"localhost"`, a specific IP). `None` when the registry entry + /// doesn't record a host. Consumers should pair this with + /// [`dial_host`] to pick the right loopback target when the + /// supervisor was bound to a wildcard. pub host: Option, pub port: u16, pub token: Option, @@ -632,8 +655,8 @@ impl ActiveAgentHost { /// Loopback address callers should dial to reach this supervisor. /// Maps IPv4/IPv6 wildcards (`0.0.0.0` / `::`) to the corresponding /// loopback; passes specific hosts (e.g. `::1`, `localhost`, - /// `10.0.0.5`) through unchanged. Missing host (older lockfile) - /// falls back to IPv4 loopback to preserve the prior behaviour. + /// `10.0.0.5`) through unchanged. A missing host falls back to IPv4 + /// loopback to preserve the prior behaviour. pub fn dial_host(&self) -> &str { dial_host(self.host.as_deref()) } diff --git a/cli/src/commands/agent_kill.rs b/cli/src/commands/agent_kill.rs index 9fbefa3f5ac..2ff764301c4 100644 --- a/cli/src/commands/agent_kill.rs +++ b/cli/src/commands/agent_kill.rs @@ -3,54 +3,165 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -use std::fs; - +use crate::constants::IS_INTERACTIVE_CLI; use crate::log; -use crate::tunnels::agent_host_metadata::AgentHostMetadata; +use crate::tunnels::agent_host_registry::{self, AgentHostEndpointMetadata}; +use crate::tunnels::user_data_path::resolve_user_data_path; use crate::util::command::kill_tree; -use crate::util::errors::{wrap, AnyError}; -use crate::util::machine::process_exists; +use crate::util::errors::{wrap, AnyError, CodeError}; +use crate::util::input::prompt_index; +use super::args::AgentKillArgs; use super::CommandContext; -/// Forcefully kills the running agent host process tree and cleans up. -pub async fn agent_kill(ctx: CommandContext) -> Result { - let lockfile_path = ctx.paths.agent_host_lockfile(); +/// Forcefully kills a running **standalone** agent host process tree and +/// cleans up its registry entry. +/// +/// Only ever considers `standalone` entries (see +/// [`agent_host_registry::list_live_standalone_endpoints`]): `editor` +/// entries are owned by running VS Code windows and must never be killed +/// by this command. If exactly one live standalone host is registered it +/// is killed directly, matching prior single-instance behavior. If more +/// than one exists, `--instance-id` selects one non-interactively; +/// otherwise, in an interactive terminal, the user is prompted to choose +/// by pid/instance/address/quality. In a non-interactive context with +/// more than one candidate and no `--instance-id`, this fails rather than +/// guessing which host to kill. Returns +/// [`CodeError::NoRunningAgentHost`] when the shared registry has no live +/// standalone entry at all. +pub async fn agent_kill(ctx: CommandContext, args: AgentKillArgs) -> Result { + let user_data_path = resolve_user_data_path(args.user_data_dir.as_deref()); + let candidates = agent_host_registry::list_live_standalone_endpoints(&ctx.log, &user_data_path); - let data = fs::read_to_string(&lockfile_path).map_err(|e| { - wrap( - e, - "No running agent host found. Start one with `code agent host`", - ) - })?; - - let metadata: AgentHostMetadata = serde_json::from_str(&data).map_err(|e| { - wrap( - e, - format!("Corrupt agent host lockfile at {}", lockfile_path.display()), - ) - })?; - - if !process_exists(metadata.pid) { - let _ = fs::remove_file(&lockfile_path); - ctx.log - .result("Agent host is not running (stale lockfile cleaned up)."); - return Ok(0); + if candidates.is_empty() { + return Err(CodeError::NoRunningAgentHost.into()); } + let selected = select_candidate(&candidates, args.instance_id.as_deref())?; + kill_standalone_endpoint(&ctx, &user_data_path, selected).await +} + +/// Picks the standalone endpoint to kill out of `candidates`, which must +/// be non-empty. Behavior: +/// - `instance_id` given: the matching candidate, or an error if none +/// matches (even when there is only one candidate — an explicit +/// mismatch should not silently kill the wrong host). +/// - Exactly one candidate, no `instance_id`: that one, directly. +/// - Multiple candidates, no `instance_id`, interactive terminal: prompt. +/// - Multiple candidates, no `instance_id`, non-interactive: error asking +/// for `--instance-id`. +fn select_candidate<'a>( + candidates: &'a [AgentHostEndpointMetadata], + instance_id: Option<&str>, +) -> Result<&'a AgentHostEndpointMetadata, AnyError> { + if let Some(instance_id) = instance_id { + return candidates + .iter() + .find(|e| e.instance_id == instance_id) + .ok_or_else(|| CodeError::UnknownAgentHostInstance(instance_id.to_string()).into()); + } + + if candidates.len() == 1 { + return Ok(&candidates[0]); + } + + if !*IS_INTERACTIVE_CLI { + let ids = candidates + .iter() + .map(|e| e.instance_id.clone()) + .collect::>() + .join(", "); + return Err(CodeError::AmbiguousAgentHostInstance(ids).into()); + } + + let labels: Vec = candidates.iter().map(|e| e.label()).collect(); + let chosen = prompt_index( + "Multiple standalone agent hosts are running; pick one to kill", + &labels, + ) + .map_err(|e| wrap(e, "Failed to read agent host selection"))?; + Ok(&candidates[chosen]) +} + +async fn kill_standalone_endpoint( + ctx: &CommandContext, + user_data_path: &std::path::Path, + selected: &AgentHostEndpointMetadata, +) -> Result { debug!( ctx.log, - "Killing agent host process tree (pid {})", metadata.pid + "Killing agent host process tree (pid {})", selected.pid ); - kill_tree(metadata.pid) + kill_tree(selected.pid) .await .map_err(|e| wrap(e, "Failed to kill agent host process tree"))?; - let _ = fs::remove_file(&lockfile_path); + agent_host_registry::remove_agent_host_endpoint(&ctx.log, user_data_path, &selected.identity()); ctx.log - .result(format!("Killed agent host (pid {}).", metadata.pid)); - + .result(format!("Killed agent host (pid {}).", selected.pid)); Ok(0) } + +#[cfg(test)] +mod tests { + use super::*; + + fn standalone(pid: u32, instance_id: &str) -> AgentHostEndpointMetadata { + AgentHostEndpointMetadata::new_standalone( + pid, + instance_id.to_string(), + "127.0.0.1".to_string(), + 9000, + "tok".to_string(), + "1".to_string(), + None, + None, + ) + } + + #[test] + fn select_candidate_returns_single_candidate_directly() { + let candidates = vec![standalone(1, "only")]; + let selected = select_candidate(&candidates, None).unwrap(); + assert_eq!(selected.instance_id, "only"); + } + + #[test] + fn select_candidate_with_instance_id_matches_even_when_only_one_candidate() { + let candidates = vec![standalone(1, "only")]; + let selected = select_candidate(&candidates, Some("only")).unwrap(); + assert_eq!(selected.instance_id, "only"); + } + + #[test] + fn select_candidate_with_mismatched_instance_id_errors() { + let candidates = vec![standalone(1, "only")]; + let err = select_candidate(&candidates, Some("nope")).unwrap_err(); + assert!(matches!( + err, + AnyError::CodeError(CodeError::UnknownAgentHostInstance(ref id)) if id == "nope" + )); + } + + #[test] + fn select_candidate_with_instance_id_picks_matching_one_of_several() { + let candidates = vec![standalone(1, "a"), standalone(2, "b")]; + let selected = select_candidate(&candidates, Some("b")).unwrap(); + assert_eq!(selected.pid, 2); + } + + #[test] + fn select_candidate_multiple_without_instance_id_in_non_interactive_test_env_errors() { + // `cargo test` runs with no attached TTY, so `IS_INTERACTIVE_CLI` is + // false here; this exercises the non-interactive-disambiguation + // error path without needing to fake a terminal/prompt. + let candidates = vec![standalone(1, "a"), standalone(2, "b")]; + let err = select_candidate(&candidates, None).unwrap_err(); + assert!(matches!( + err, + AnyError::CodeError(CodeError::AmbiguousAgentHostInstance(_)) + )); + } +} diff --git a/cli/src/commands/agent_logs.rs b/cli/src/commands/agent_logs.rs index 528598e9e8c..7b88f7cf530 100644 --- a/cli/src/commands/agent_logs.rs +++ b/cli/src/commands/agent_logs.rs @@ -13,13 +13,27 @@ use crate::tunnels::shutdown_signal::ShutdownRequest; use crate::util::errors::AnyError; use super::agent; +use super::agent_discovery; use super::args::AgentLogsArgs; use super::output::Styles; use super::CommandContext; /// Subscribes to a session and streams actions/notifications in real time. pub async fn agent_logs(ctx: CommandContext, args: AgentLogsArgs) -> Result { - let client = agent::connect(&ctx, args.address.as_deref(), args.tunnel.as_deref()).await?; + let client = match ( + args.discovery.address.as_deref(), + args.discovery.tunnel.as_deref(), + ) { + (None, None) => { + agent_discovery::connect_to_session_host( + &ctx, + &args.session, + args.discovery.user_data_dir.as_deref(), + ) + .await? + } + (address, tunnel) => agent::connect_explicit(&ctx, address, tunnel).await?, + }; let (result, mut sub): (SubscribeResult, _) = { let r: SubscribeResult = agent::request_with_auth( diff --git a/cli/src/commands/agent_ps.rs b/cli/src/commands/agent_ps.rs index 32afc2be451..afbfae15d06 100644 --- a/cli/src/commands/agent_ps.rs +++ b/cli/src/commands/agent_ps.rs @@ -6,20 +6,63 @@ use ahp_types::commands::{ListSessionsParams, ListSessionsResult}; use ahp_types::state::{SessionStatus, SessionSummary}; use ahp_types::ROOT_RESOURCE_URI; +use futures::stream::FuturesUnordered; +use futures::StreamExt; +use serde::Serialize; -use crate::util::errors::AnyError; +use crate::log; +use crate::tunnels::agent_host_registry::AgentHostEndpointMetadata; +use crate::util::errors::{wrap, AnyError, CodeError}; use super::agent; +use super::agent_discovery; use super::args::AgentPsArgs; use super::output::{self, Styles}; use super::CommandContext; /// Lists active sessions on a running agent host. +/// +/// With an explicit `--address`/`--tunnel`, this queries exactly one host +/// and prints its sessions with no host-provenance header, same as before +/// multi-host discovery existed. Otherwise every live registry endpoint +/// (both `editor` and `standalone`, both socket/pipe and TCP) is queried +/// concurrently; human output is printed as each host completes rather +/// than waiting for the slowest one, and `--json` collects every host's +/// outcome into a single valid JSON document tagged with host +/// provenance. Returns [`CodeError::NoRunningAgentHost`] when the shared +/// registry has no live entries at all — the sole source of truth for +/// automatic discovery. pub async fn agent_ps(ctx: CommandContext, args: AgentPsArgs) -> Result { - let client = agent::connect(&ctx, args.address.as_deref(), args.tunnel.as_deref()).await?; + if args.discovery.address.is_some() || args.discovery.tunnel.is_some() { + return agent_ps_single( + &ctx, + &args, + args.discovery.address.as_deref(), + args.discovery.tunnel.as_deref(), + ) + .await; + } + + let endpoints = + agent_discovery::discover_live_endpoints(&ctx, args.discovery.user_data_dir.as_deref()); + if endpoints.is_empty() { + return Err(CodeError::NoRunningAgentHost.into()); + } + + agent_ps_multi(&ctx, &args, endpoints).await +} + +/// Single-host path used for an explicit `--address`/`--tunnel`. +async fn agent_ps_single( + ctx: &CommandContext, + args: &AgentPsArgs, + address: Option<&str>, + tunnel: Option<&str>, +) -> Result { + let client = agent::connect_explicit(ctx, address, tunnel).await?; let result: ListSessionsResult = agent::request_with_auth( - &ctx, + ctx, &client, "listSessions", ListSessionsParams { @@ -31,22 +74,11 @@ pub async fn agent_ps(ctx: CommandContext, args: AgentPsArgs) -> Result = if args.all { - result.items.iter().collect() - } else { - result - .items - .iter() - .filter(|s| is_active(s.status)) - .collect() - }; - - // Most-recently-modified first. - items.sort_by_key(|b| std::cmp::Reverse(b.modified_at)); + let items = select_and_sort(&result.items, args.all); if args.json { let json = serde_json::to_string_pretty(&items) - .map_err(|e| crate::util::errors::wrap(e, "Failed to serialize sessions"))?; + .map_err(|e| wrap(e, "Failed to serialize sessions"))?; output::print_paged(&json); } else if items.is_empty() { ctx.log.result("No active sessions."); @@ -58,6 +90,198 @@ pub async fn agent_ps(ctx: CommandContext, args: AgentPsArgs) -> Result, AnyError>, +} + +async fn query_host(ctx: &CommandContext, endpoint: AgentHostEndpointMetadata) -> HostSessions { + let sessions = query_host_sessions(ctx, &endpoint).await; + HostSessions { endpoint, sessions } +} + +async fn query_host_sessions( + ctx: &CommandContext, + endpoint: &AgentHostEndpointMetadata, +) -> Result, AnyError> { + let client = agent::connect_to_endpoint(endpoint).await?; + + let result: ListSessionsResult = agent::request_with_auth( + ctx, + &client, + "listSessions", + ListSessionsParams { + channel: ROOT_RESOURCE_URI.to_string(), + filter: None, + }, + ) + .await?; + + client.shutdown().await; + Ok(result.items) +} + +/// Multi-host auto-discovery path (requirement 3): connects to and +/// queries every discovered host concurrently via `FuturesUnordered`, +/// printing each host's results (human mode) as soon as that host +/// completes rather than waiting for the slowest one. Per-host failures +/// are reported as warnings and don't stop other hosts; the command only +/// returns an error if *no* host could be queried at all. `--json` +/// collects every host's outcome (success or error) before printing +/// once, so the emitted JSON is always a single valid document — never +/// corrupted by interleaving multiple hosts' output — tagged with host +/// provenance. +async fn agent_ps_multi( + ctx: &CommandContext, + args: &AgentPsArgs, + endpoints: Vec, +) -> Result { + let mut tasks = FuturesUnordered::new(); + for endpoint in endpoints { + tasks.push(query_host(ctx, endpoint)); + } + + let mut any_succeeded = false; + let mut collected: Vec = Vec::new(); + + while let Some(outcome) = tasks.next().await { + match &outcome.sessions { + Ok(_) => any_succeeded = true, + Err(e) => { + warning!( + ctx.log, + "Could not query agent host {}: {}", + outcome.endpoint.label(), + e + ); + } + } + + if args.json { + collected.push(outcome); + } else { + print_host_outcome_human(&outcome, args.all); + } + } + + if !any_succeeded { + return Err(CodeError::NoAgentHostReachable( + "no discovered agent host could be queried".to_string(), + ) + .into()); + } + + if args.json { + let json = build_json_output(&collected, args.all)?; + output::print_paged(&json); + } + + Ok(0) +} + +/// Prints one host's outcome immediately (bypasses the pager used by the +/// single-host path: paging would buffer output until the whole command +/// finishes, defeating the point of streaming results as hosts +/// complete). +fn print_host_outcome_human(outcome: &HostSessions, all: bool) { + let header = Styles::title(); + println!( + "\n{}", + header.apply_to(format!("── {} ──", outcome.endpoint.label())) + ); + + match &outcome.sessions { + Ok(sessions) => { + let items = select_and_sort(sessions, all); + if items.is_empty() { + println!(" {}", Styles::muted().apply_to("No active sessions.")); + } else { + print!("{}", format_sessions_list(&items)); + } + } + Err(e) => { + println!(" {}", Styles::error().apply_to(format!("⚠ {e}"))); + } + } +} + +/// JSON-serializable host provenance tag, kept intentionally small and +/// stable (not just a re-export of [`AgentHostEndpointMetadata`], whose +/// shape is an internal implementation detail of the registry file). +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +struct HostJson { + #[serde(rename = "type")] + server_type: &'static str, + pid: u32, + instance_id: String, + address: String, + #[serde(skip_serializing_if = "Option::is_none")] + quality: Option, + #[serde(skip_serializing_if = "Option::is_none")] + tunnel_name: Option, +} + +impl From<&AgentHostEndpointMetadata> for HostJson { + fn from(e: &AgentHostEndpointMetadata) -> Self { + HostJson { + server_type: match e.server_type { + crate::tunnels::agent_host_registry::AgentHostServerType::Editor => "editor", + crate::tunnels::agent_host_registry::AgentHostServerType::Standalone => { + "standalone" + } + }, + pid: e.pid, + instance_id: e.instance_id.clone(), + address: e.address_label(), + quality: e.quality.clone(), + tunnel_name: e.tunnel_name.clone(), + } + } +} + +#[derive(Serialize)] +#[serde(rename_all = "camelCase")] +struct HostSessionsJson<'a> { + host: HostJson, + #[serde(skip_serializing_if = "Option::is_none")] + sessions: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, +} + +fn build_json_output(collected: &[HostSessions], all: bool) -> Result { + let hosts: Vec = collected + .iter() + .map(|outcome| HostSessionsJson { + host: HostJson::from(&outcome.endpoint), + sessions: outcome + .sessions + .as_ref() + .ok() + .map(|s| select_and_sort(s, all)), + error: outcome.sessions.as_ref().err().map(|e| e.to_string()), + }) + .collect(); + + serde_json::to_string_pretty(&hosts).map_err(|e| wrap(e, "Failed to serialize sessions").into()) +} + +/// Applies the `--all` filter (active-only unless set) and sorts +/// most-recently-modified first, shared by every output path so human +/// and JSON output (single- or multi-host) always agree on ordering. +fn select_and_sort(sessions: &[SessionSummary], all: bool) -> Vec<&SessionSummary> { + let mut items: Vec<&SessionSummary> = if all { + sessions.iter().collect() + } else { + sessions.iter().filter(|s| is_active(s.status)).collect() + }; + + items.sort_by_key(|b| std::cmp::Reverse(b.modified_at)); + items +} + /// A session is "active" if it is in-progress, needs input, or errored /// (i.e. not just idle/archived). fn is_active(status: u32) -> bool { @@ -131,3 +355,72 @@ fn status_styled(status: u32) -> console::StyledObject { Styles::muted().apply_to(format!("? unknown ({})", status.bits())) } } + +#[cfg(test)] +mod tests { + use super::*; + + fn session(resource: &str, status: u32, modified_at: i64) -> SessionSummary { + SessionSummary { + resource: resource.to_string(), + provider: "test".to_string(), + title: String::new(), + status, + activity: None, + created_at: 0, + modified_at, + project: None, + model: None, + agent: None, + working_directory: None, + changes: None, + annotations: None, + } + } + + #[test] + fn select_and_sort_filters_idle_unless_all() { + let idle = session("a", SessionStatus::Idle.bits(), 1); + let active = session("b", SessionStatus::InProgress.bits(), 2); + let sessions = vec![idle.clone(), active.clone()]; + + let filtered = select_and_sort(&sessions, false); + assert_eq!(filtered.len(), 1); + assert_eq!(filtered[0].resource, "b"); + + let all = select_and_sort(&sessions, true); + assert_eq!(all.len(), 2); + } + + #[test] + fn select_and_sort_orders_most_recently_modified_first() { + let older = session("a", SessionStatus::InProgress.bits(), 1); + let newer = session("b", SessionStatus::InProgress.bits(), 2); + let sessions = vec![older, newer]; + + let sorted = select_and_sort(&sessions, true); + assert_eq!(sorted[0].resource, "b"); + assert_eq!(sorted[1].resource, "a"); + } + + #[test] + fn host_json_carries_provenance_fields() { + let entry = AgentHostEndpointMetadata::new_standalone( + 42, + "instance-a".to_string(), + "127.0.0.1".to_string(), + 8080, + "tok".to_string(), + "0.1.0".to_string(), + Some("insider".to_string()), + None, + ); + + let json = HostJson::from(&entry); + assert_eq!(json.server_type, "standalone"); + assert_eq!(json.pid, 42); + assert_eq!(json.instance_id, "instance-a"); + assert_eq!(json.address, "127.0.0.1:8080"); + assert_eq!(json.quality.as_deref(), Some("insider")); + } +} diff --git a/cli/src/commands/agent_stop.rs b/cli/src/commands/agent_stop.rs index 58e5915ead3..cb8bad3fe18 100644 --- a/cli/src/commands/agent_stop.rs +++ b/cli/src/commands/agent_stop.rs @@ -11,13 +11,27 @@ use crate::log; use crate::util::errors::{wrap, AnyError}; use super::agent; +use super::agent_discovery; use super::args::AgentStopArgs; use super::CommandContext; /// Cancels the active turn of every in-progress chat in a session on a running /// agent host. pub async fn agent_stop(ctx: CommandContext, args: AgentStopArgs) -> Result { - let client = agent::connect(&ctx, args.address.as_deref(), args.tunnel.as_deref()).await?; + let client = match ( + args.discovery.address.as_deref(), + args.discovery.tunnel.as_deref(), + ) { + (None, None) => { + agent_discovery::connect_to_session_host( + &ctx, + &args.session, + args.discovery.user_data_dir.as_deref(), + ) + .await? + } + (address, tunnel) => agent::connect_explicit(&ctx, address, tunnel).await?, + }; // Subscribe to the session to get its catalog of chats. let result: SubscribeResult = agent::request_with_auth( diff --git a/cli/src/commands/args.rs b/cli/src/commands/args.rs index ef6daa0742b..2095d31ca1a 100644 --- a/cli/src/commands/args.rs +++ b/cli/src/commands/args.rs @@ -243,7 +243,7 @@ pub struct AgentHostArgs { pub host: Option, /// Port the agent host should bind on. If 0 (the default) the OS /// picks a free ephemeral port; the chosen port is recorded in the - /// agent host lockfile. + /// shared agent-host endpoint registry. #[clap(long, default_value_t = 0)] pub port: u16, /// A secret that must be included with all requests. @@ -259,6 +259,15 @@ pub struct AgentHostArgs { #[clap(long)] pub server_data_dir: Option, + /// Overrides the resolved user data directory used to home the shared + /// local agent-host endpoint registry + /// (`/agent-host/local-endpoint/metadata.json`), the same + /// file editor windows publish to. Defaults to the platform user data + /// directory (honoring `VSCODE_PORTABLE` / `VSCODE_APPDATA` when set), + /// matching the editor's own resolution rules. + #[clap(long)] + pub user_data_dir: Option, + /// Stop any agent host already running on this machine and start a /// fresh one. Without this flag, the command reuses an existing live /// supervisor when its configuration is compatible, and errors out @@ -304,14 +313,29 @@ pub enum AgentSubcommand { Stop(AgentStopArgs), /// Forcefully kill the running agent host process tree. - Kill, + Kill(AgentKillArgs), /// Stream live session events. Logs(AgentLogsArgs), } +/// Discovery/connection target shared by every agent-host command that +/// can either auto-discover a local instance or target one explicitly: +/// `code agent ps|stop|logs`. `--user-data-dir` scopes automatic +/// discovery to a specific registry (see +/// [`crate::commands::agent_discovery::discover_live_endpoints`]); +/// `--address`/`--tunnel` bypass discovery entirely and connect to +/// exactly one explicit target (see +/// [`crate::commands::agent::connect_explicit`]). Passing neither +/// `--address` nor `--tunnel` means "discover automatically", not +/// "connect nowhere" — every consumer of this struct must branch on that +/// itself. #[derive(Args, Debug, Clone)] -pub struct AgentPsArgs { +pub struct AgentDiscoveryArgs { + /// Directory containing the shared agent host registry used for automatic discovery. + #[clap(long)] + pub user_data_dir: Option, + /// WebSocket address of a running agent host (e.g. ws://127.0.0.1:1234?tkn=secret). /// If omitted, the CLI discovers a locally running agent host automatically. #[clap(long)] @@ -320,6 +344,13 @@ pub struct AgentPsArgs { /// Connect via a named dev tunnel instead of the local address. #[clap(long)] pub tunnel: Option, +} + +#[derive(Args, Debug, Clone)] +pub struct AgentPsArgs { + /// Discovery/connection target; see [`AgentDiscoveryArgs`]. + #[clap(flatten)] + pub discovery: AgentDiscoveryArgs, /// Output results as JSON instead of a human-readable table. #[clap(long)] @@ -335,14 +366,9 @@ pub struct AgentStopArgs { /// Session URI to cancel the active turn of (e.g. copilot:/). pub session: String, - /// WebSocket address of a running agent host. - /// If omitted, the CLI discovers a locally running agent host automatically. - #[clap(long)] - pub address: Option, - - /// Connect via a named dev tunnel instead of the local address. - #[clap(long)] - pub tunnel: Option, + /// Discovery/connection target; see [`AgentDiscoveryArgs`]. + #[clap(flatten)] + pub discovery: AgentDiscoveryArgs, } #[derive(Args, Debug, Clone)] @@ -350,14 +376,22 @@ pub struct AgentLogsArgs { /// Session URI to stream events for (e.g. copilot:/). pub session: String, - /// WebSocket address of a running agent host. - /// If omitted, the CLI discovers a locally running agent host automatically. - #[clap(long)] - pub address: Option, + /// Discovery/connection target; see [`AgentDiscoveryArgs`]. + #[clap(flatten)] + pub discovery: AgentDiscoveryArgs, +} - /// Connect via a named dev tunnel instead of the local address. +#[derive(Args, Debug, Clone)] +pub struct AgentKillArgs { + /// Directory containing the shared agent host registry used for automatic discovery. #[clap(long)] - pub tunnel: Option, + pub user_data_dir: Option, + + /// Instance ID of the standalone agent host to kill, as shown when + /// multiple are running. Required to select non-interactively when + /// more than one live standalone agent host is registered. + #[clap(long)] + pub instance_id: Option, } #[derive(Args, Debug, Clone)] diff --git a/cli/src/constants.rs b/cli/src/constants.rs index 938f43432d7..02c4b00aa6e 100644 --- a/cli/src/constants.rs +++ b/cli/src/constants.rs @@ -67,6 +67,16 @@ pub const QUALITYLESS_PRODUCT_NAME: &str = match option_env!("VSCODE_CLI_QUALITY None => "Code", }; +/// Short product name, mirroring `product.json`'s `nameShort` (e.g. `Code - +/// OSS`, `Visual Studio Code`). Used as the leaf directory name when +/// resolving the platform user data directory, matching the TypeScript +/// resolver in `src/vs/platform/environment/node/userDataPath.ts` (which is +/// passed `product.nameShort`). +pub const PRODUCT_NAME_SHORT: &str = match option_env!("VSCODE_CLI_NAME_SHORT") { + Some(n) => n, + None => "Code - OSS", +}; + /// Name of the application without quality information. pub const QUALITYLESS_SERVER_NAME: &str = concatcp!(QUALITYLESS_PRODUCT_NAME, " Server"); @@ -94,9 +104,10 @@ pub const DEFAULT_DATA_PARENT_DIR: &str = match option_env!("VSCODE_CLI_DATA_FOL /// Canonical, machine-wide parent directory used to coordinate the agent /// host across CLI invocations. Mirrors the `serverDataFolderName` in -/// `product.json` so the lockfile/log written by `code agent host` lines -/// up with the directory the SSH `command-shell` entry point already uses -/// (otherwise local + remote would race on different lockfiles). +/// `product.json` so the supervisor log written by `code agent host` +/// lines up with the directory the SSH `command-shell` entry point +/// already uses (otherwise local + remote would race on different +/// directories). pub const SERVER_DATA_PARENT_DIR: &str = match option_env!("VSCODE_CLI_SERVER_DATA_FOLDER_NAME") { Some(n) => n, None => ".vscode-server-oss", diff --git a/cli/src/state.rs b/cli/src/state.rs index 19334123db0..f0797f856d3 100644 --- a/cli/src/state.rs +++ b/cli/src/state.rs @@ -237,18 +237,6 @@ impl LauncherPaths { }) } - /// Lockfile for the running agent host. Pinned to the canonical - /// server data dir (see [`agent_host_root`]) so a `code agent host` - /// invoked locally and the supervisor spawned by the SSH - /// `command-shell` path agree on the same lockfile regardless of - /// `--cli-data-dir`. - pub fn agent_host_lockfile(&self) -> PathBuf { - self.agent_host_root().join(format!( - "agent-host-{}.lock", - VSCODE_CLI_QUALITY.unwrap_or("oss") - )) - } - /// Suggested path for the detached `code agent host` supervisor's log /// file. The supervisor severs its inherited stdio after signalling /// readiness, so a file log is the only way to debug post-handoff @@ -260,10 +248,10 @@ impl LauncherPaths { )) } - /// Canonical machine-wide directory holding agent host coordination - /// files (lockfile + supervisor log). Anchored on `serverDataFolderName` + /// Canonical machine-wide directory holding the detached `code agent + /// host` supervisor's log file. Anchored on `serverDataFolderName` /// rather than `self.root` so different `--cli-data-dir` values still - /// land on the same lockfile. Falls back to `self.root` when no home + /// land on the same log file. Falls back to `self.root` when no home /// directory is available. fn agent_host_root(&self) -> PathBuf { match dirs::home_dir() { diff --git a/cli/src/tunnels.rs b/cli/src/tunnels.rs index 76b27899635..a92ba998361 100644 --- a/cli/src/tunnels.rs +++ b/cli/src/tunnels.rs @@ -14,7 +14,9 @@ pub mod singleton_client; pub mod singleton_server; pub mod agent_host; -pub mod agent_host_metadata; +pub mod agent_host_registry; +#[cfg(windows)] +mod agent_host_registry_acl_windows; mod challenge; mod control_server; mod nosleep; @@ -35,6 +37,7 @@ mod service_macos; #[cfg(target_os = "windows")] mod service_windows; mod socket_signal; +pub mod user_data_path; mod wsl_detect; pub use control_server::{ diff --git a/cli/src/tunnels/agent_host.rs b/cli/src/tunnels/agent_host.rs index 93651587620..ab05a88fda2 100644 --- a/cli/src/tunnels/agent_host.rs +++ b/cli/src/tunnels/agent_host.rs @@ -21,7 +21,9 @@ use tokio::io::{AsyncBufReadExt, BufReader}; use tokio::net::TcpListener; use tokio::sync::Mutex; -use crate::async_pipe::{get_socket_name, get_socket_rw_stream, listen_socket_rw_stream, AsyncPipeListener}; +use crate::async_pipe::{ + get_socket_name, get_socket_rw_stream, listen_socket_rw_stream, AsyncPipeListener, +}; use crate::constants::VSCODE_CLI_QUALITY; use crate::download_cache::DownloadCache; use crate::log; @@ -36,8 +38,9 @@ use crate::util::http::{empty_body, full_body, HyperBody}; use crate::util::io::SilentCopyProgress; use crate::util::sync::{new_barrier, Barrier, BarrierOpener}; -use super::agent_host_metadata::{ - remove_agent_host_metadata_for_pid, write_agent_host_metadata, AgentHostMetadata, +use super::agent_host_registry::{ + self, AgentHostEndpointIdentity, AgentHostEndpointMetadata, AgentHostServerType, + AGENT_HOST_PROTOCOL_VERSION, }; use super::paths::{get_server_folder_name, SERVER_FOLDER_NAME}; use super::shutdown_signal::ShutdownSignal; @@ -634,13 +637,21 @@ impl AgentHostManager { let mut listener = match listen_socket_rw_stream(path).await { Ok(l) => l, Err(e) => { - warning!(self.log, "Failed to bind management socket {:?}: {}", path, e); + warning!( + self.log, + "Failed to bind management socket {:?}: {}", + path, + e + ); self.management_listener_started .store(false, Ordering::SeqCst); return; } }; - debug!(self.log, "Listening for agent host management requests on {:?}", path); + debug!( + self.log, + "Listening for agent host management requests on {:?}", path + ); self.run_management_accept_loop(&mut listener).await; } @@ -704,7 +715,11 @@ impl AgentHostManager { let new_release = match self.get_latest_release().await { Ok(r) => r, Err(e) => { - warning!(self.log, "Upgrade request: latest release lookup failed: {}", e); + warning!( + self.log, + "Upgrade request: latest release lookup failed: {}", + e + ); return json_response( 503, &UpgradeResponse { @@ -762,7 +777,12 @@ impl AgentHostManager { // `upgradeStarted`. The background update loop usually pre-fetches // this, so the common path is a no-op. if let Err(e) = self.ensure_downloaded(&new_release).await { - warning!(self.log, "Failed to download upgrade {}: {}", new_release, e); + warning!( + self.log, + "Failed to download upgrade {}: {}", + new_release, + e + ); self.upgrade_in_progress.store(false, Ordering::SeqCst); return json_response( 503, @@ -791,9 +811,15 @@ impl AgentHostManager { // ready endpoint instead of paying for startup again. match self_clone.start_server().await { Ok(_) => info!(self_clone.log, "Restarted agent host on {}", release_commit), - Err(e) => warning!(self_clone.log, "Failed to restart agent host after upgrade: {}", e), + Err(e) => warning!( + self_clone.log, + "Failed to restart agent host after upgrade: {}", + e + ), } - self_clone.upgrade_in_progress.store(false, Ordering::SeqCst); + self_clone + .upgrade_in_progress + .store(false, Ordering::SeqCst); }); json_response( @@ -1000,21 +1026,32 @@ pub struct AgentHostSidecar { listener: TcpListener, bound_addr: SocketAddr, public_token: Option, - lockfile_path: PathBuf, + user_data_path: PathBuf, + instance_id: String, pid: u32, + /// Set once registry cleanup for this instance's identity has been + /// performed (successfully or not — a best-effort attempt counts), so + /// `Drop` never redundantly repeats it after an explicit [`Self::shutdown`]. + registry_cleaned_up: AtomicBool, } impl AgentHostSidecar { - /// Binds a TCP listener at `addr`, writes the canonical agent host - /// lockfile pointing at the bound port, and returns a sidecar ready to - /// [`serve`](Self::serve) connections. The agent host backend is *not* - /// started here — the wrapped [`AgentHostManager`] starts it on demand - /// when the first request arrives. + /// Binds a TCP listener at `addr`, publishes a `standalone` entry to the + /// shared local agent-host endpoint registry (schema v2, see + /// [`agent_host_registry`]) pointing at the bound port, and returns a + /// sidecar ready to [`serve`](Self::serve) connections. The agent host + /// backend is *not* started here — the wrapped [`AgentHostManager`] + /// starts it on demand when the first request arrives. /// /// `loopback_auth` decides whether the local TCP accept loop enforces a /// connection token. The caller MUST make this choice deliberately: /// loopback is reachable from any local process, so binding without a /// token must be a conscious user opt-in (e.g. `--without-connection-token`). + /// + /// `user_data_path` is the resolved user data directory that homes the + /// registry (see [`super::user_data_path`]); `instance_id` is this + /// process's stable identity within the registry, used to disambiguate + /// PID reuse and to scope `--replace`/removal to exactly this entry. pub async fn bind_tcp( log: log::Logger, manager: Arc, @@ -1022,7 +1059,8 @@ impl AgentHostSidecar { host_label: Option, loopback_auth: LoopbackAuth, tunnel_name: Option, - lockfile_path: PathBuf, + user_data_path: PathBuf, + instance_id: String, ) -> Result, AnyError> { let public_token = loopback_auth.into_token(); let listener = TcpListener::bind(addr) @@ -1039,15 +1077,44 @@ impl AgentHostSidecar { // character-equal without spuriously flagging hostname-vs-IP // equivalents as a config conflict. let host = host_label.unwrap_or_else(|| bound_addr.ip().to_string()); - let metadata = build_metadata( + let entry = AgentHostEndpointMetadata::new_standalone( pid, + instance_id.clone(), host, bound_addr.port(), - public_token.clone(), - tunnel_name.as_deref(), + public_token.clone().unwrap_or_default(), + AGENT_HOST_PROTOCOL_VERSION.to_string(), + VSCODE_CLI_QUALITY.map(str::to_string), + tunnel_name, ); - if let Err(e) = write_agent_host_metadata(&lockfile_path, &metadata) { - warning!(log, "Failed to write agent host lockfile: {}", e); + + // Registry publish does blocking filesystem I/O and may briefly + // spin-retry on a contended lock; run it on a blocking-safe thread + // so it never stalls the tokio runtime. + { + let publish_log = log.clone(); + let publish_path = user_data_path.clone(); + match tokio::task::spawn_blocking(move || { + agent_host_registry::publish_agent_host_endpoint( + &publish_log, + &publish_path, + &entry, + ) + }) + .await + { + Ok(Ok(())) => {} + Ok(Err(e)) => warning!( + log, + "Failed to publish agent host endpoint registry entry: {}", + e + ), + Err(e) => warning!( + log, + "Agent host endpoint registry publish task failed: {}", + e + ), + } } Ok(Arc::new(Self { @@ -1056,8 +1123,10 @@ impl AgentHostSidecar { listener, bound_addr, public_token, - lockfile_path, + user_data_path, + instance_id, pid, + registry_cleaned_up: AtomicBool::new(false), })) } @@ -1131,37 +1200,78 @@ impl AgentHostSidecar { } } - /// Stops the agent host backend and removes the lockfile if it still - /// belongs to this process. Safe to call multiple times. + /// Stops the agent host backend and removes this instance's entry from + /// the shared local agent-host endpoint registry. Safe to call multiple + /// times: only the first call performs registry cleanup, and `Drop` + /// will not repeat it afterwards (see `registry_cleaned_up`). pub async fn shutdown(&self) { self.manager.kill_running_server().await; - if let Err(e) = remove_agent_host_metadata_for_pid(&self.lockfile_path, self.pid) { - warning!(self.log, "Failed to clean up agent host lockfile: {}", e); + if self.registry_cleaned_up.swap(true, Ordering::SeqCst) { + return; + } + let identity = AgentHostEndpointIdentity { + server_type: AgentHostServerType::Standalone, + pid: self.pid, + instance_id: self.instance_id.clone(), + }; + let log = self.log.clone(); + let user_data_path = self.user_data_path.clone(); + if let Err(e) = tokio::task::spawn_blocking(move || { + agent_host_registry::remove_agent_host_endpoint(&log, &user_data_path, &identity); + }) + .await + { + warning!( + self.log, + "Agent host endpoint registry cleanup task failed: {}", + e + ); } } } impl Drop for AgentHostSidecar { fn drop(&mut self) { - // Best-effort cleanup if the caller forgot to call `shutdown`. Only - // removes the lockfile when the recorded PID still matches us. - let _ = remove_agent_host_metadata_for_pid(&self.lockfile_path, self.pid); - } -} + // If `shutdown` already performed (or is performing) registry + // cleanup, don't repeat it here. + if self.registry_cleaned_up.swap(true, Ordering::SeqCst) { + return; + } -fn build_metadata( - pid: u32, - host: String, - port: u16, - connection_token: Option, - tunnel_name: Option<&str>, -) -> AgentHostMetadata { - let mut metadata = AgentHostMetadata::new(pid, port); - metadata.host = Some(host); - metadata.connection_token = connection_token; - metadata.quality = VSCODE_CLI_QUALITY.map(str::to_string); - metadata.tunnel_name = tunnel_name.map(str::to_string); - metadata + // Best-effort cleanup for the case where the caller forgot to call + // `shutdown`. `remove_agent_host_endpoint` only removes the entry + // that exactly matches our own `(type, pid, instanceId)` identity. + let identity = AgentHostEndpointIdentity { + server_type: AgentHostServerType::Standalone, + pid: self.pid, + instance_id: self.instance_id.clone(), + }; + let log = self.log.clone(); + let user_data_path = self.user_data_path.clone(); + + // `drop` is synchronous and must not block a Tokio worker thread + // with this call's blocking filesystem I/O (lock acquisition, + // reads/writes). If a runtime is reachable from here, hand the + // cleanup off to a blocking-safe thread and don't wait for it — + // this is already a best-effort fallback, so fire-and-forget is + // acceptable. If no runtime is available (e.g. this sidecar + // outlived it), there is no worker thread left to protect, so it's + // safe to just do the blocking removal inline. + match tokio::runtime::Handle::try_current() { + Ok(handle) => { + handle.spawn_blocking(move || { + agent_host_registry::remove_agent_host_endpoint( + &log, + &user_data_path, + &identity, + ); + }); + } + Err(_) => { + agent_host_registry::remove_agent_host_endpoint(&log, &user_data_path, &identity); + } + } + } } /// How the loopback TCP accept loop authenticates incoming connections. @@ -1210,76 +1320,68 @@ async fn handle_request_with_auth( handle_request(manager, req).await } -// ---- Lockfile-aware reuse --------------------------------------------------- +// ---- Registry-based reuse --------------------------------------------------- -/// Decision derived from inspecting `agent-host-.lock`. Used by -/// CLI entry points (e.g. `code tunnel`, `code agent host`) to decide -/// whether they may safely own the agent host lockfile or should forward -/// to / share the existing one. +/// Decision derived from consulting the shared local agent-host endpoint +/// registry (schema v2; see [`agent_host_registry`]). Used by CLI entry +/// points (e.g. `code tunnel`, `code agent host`) to decide whether they +/// may safely start their own supervisor or should forward to / share an +/// existing one. /// /// The agent host server is downloaded on demand and may speak a newer /// protocol than the CLI itself is built with, so we deliberately do NOT /// check the protocol version: any live registered supervisor is always /// considered reusable. #[derive(Debug, Clone, PartialEq, Eq)] -pub enum AgentHostLockfileDecision { - /// No live agent host registered; the caller may start its own sidecar. +pub enum AgentHostReuseDecision { + /// No live standalone agent host registered; the caller may start its + /// own sidecar. SpawnFresh, - /// A live agent host supervisor owns the lockfile. Tunnel callers - /// should forward to `127.0.0.1:port` instead of binding a second - /// listener / clobbering the lockfile. `host` and `tunnel_name` - /// expose the supervisor's effective config so foreground callers - /// can detect a configuration conflict and refuse to silently reuse. + /// A live standalone agent host supervisor owns a registry entry. + /// Tunnel callers should forward to `127.0.0.1:port` instead of + /// binding a second listener / publishing a conflicting entry. `host` + /// and `tunnel_name` expose the supervisor's effective config so + /// foreground callers can detect a configuration conflict and refuse + /// to silently reuse. Reuse { pid: u32, host: Option, port: u16, token: Option, tunnel_name: Option, + /// This entry's stable identity within the registry, used by + /// `--replace` to scope removal to exactly this instance. + instance_id: String, }, } -/// Inspect the agent host lockfile at `path` and decide whether the caller -/// should spawn a fresh sidecar or reuse an existing live one. Missing / -/// unreadable / stale (dead-PID) lockfiles all map to -/// [`AgentHostLockfileDecision::SpawnFresh`]. -pub fn classify_agent_host_lockfile( +/// Preferred entry point for CLI commands that need to discover a live +/// standalone agent host: consults the shared local agent-host endpoint +/// registry (schema v2), the sole source of truth for automatic +/// discovery. +/// +/// `editor` entries are never selected here — they are owned by running VS +/// Code windows and must remain invisible to (and unkillable by) the +/// standalone CLI's discovery/`--replace` path. See +/// [`agent_host_registry::select_live_standalone_endpoint`]. +pub fn classify_agent_host( log: &log::Logger, - path: &std::path::Path, -) -> AgentHostLockfileDecision { - use super::agent_host_metadata::read_agent_host_metadata; - use crate::util::machine::process_exists; - - let metadata = match read_agent_host_metadata(path) { - Ok(Some(m)) => m, - Ok(None) => return AgentHostLockfileDecision::SpawnFresh, - Err(e) => { - debug!( - log, - "Could not read agent host lockfile {}: {}", - path.display(), - e - ); - return AgentHostLockfileDecision::SpawnFresh; - } - }; - - if !process_exists(metadata.pid) { - debug!( - log, - "Agent host lockfile {} references dead PID {}; treating as stale", - path.display(), - metadata.pid - ); - return AgentHostLockfileDecision::SpawnFresh; - } - - AgentHostLockfileDecision::Reuse { - pid: metadata.pid, - host: metadata.host, - port: metadata.port, - token: metadata.connection_token, - tunnel_name: metadata.tunnel_name, + user_data_path: &std::path::Path, +) -> AgentHostReuseDecision { + match agent_host_registry::select_live_standalone_endpoint(log, user_data_path) { + Some(selected) => AgentHostReuseDecision::Reuse { + pid: selected.pid, + host: Some(selected.host), + port: selected.port, + token: if selected.connection_token.is_empty() { + None + } else { + Some(selected.connection_token) + }, + tunnel_name: selected.tunnel_name, + instance_id: selected.instance_id, + }, + None => AgentHostReuseDecision::SpawnFresh, } } @@ -1382,9 +1484,7 @@ fn inject_connection_token(uri: &::http::Uri, token: &str) -> ::http::Uri { #[cfg(test)] mod tests { use super::*; - use crate::tunnels::agent_host_metadata::{ - read_agent_host_metadata, AGENT_HOST_PROTOCOL_VERSION, - }; + use crate::tunnels::agent_host_registry::AgentHostEndpointAddress; use crate::util::http::ReqwestSimpleHttp; use std::path::Path; @@ -1492,36 +1592,10 @@ mod tests { assert!(json.contains(r#""restartDelayMs":3000"#), "got: {}", json); } - #[test] - fn metadata_includes_quality_and_no_tunnel() { - let metadata = build_metadata( - 42, - "127.0.0.1".to_string(), - 8080, - Some("tok".to_string()), - None, - ); - - assert_eq!(metadata.pid, 42); - assert_eq!(metadata.host.as_deref(), Some("127.0.0.1")); - assert_eq!(metadata.port, 8080); - assert_eq!(metadata.connection_token.as_deref(), Some("tok")); - assert_eq!(metadata.protocol_version, AGENT_HOST_PROTOCOL_VERSION); - assert_eq!(metadata.tunnel_name, None); - } - - #[test] - fn metadata_records_tunnel_name() { - let metadata = build_metadata(42, "0.0.0.0".to_string(), 8080, None, Some("my-tunnel")); - - assert_eq!(metadata.host.as_deref(), Some("0.0.0.0")); - assert_eq!(metadata.tunnel_name.as_deref(), Some("my-tunnel")); - } - #[tokio::test] - async fn bind_tcp_writes_lockfile_with_bound_port_and_pid() { + async fn bind_tcp_publishes_registry_entry_with_bound_port_and_pid() { let dir = tempfile::tempdir().unwrap(); - let lockfile = dir.path().join("agent-host.lock"); + let user_data_path = dir.path().join("user-data"); let manager = make_test_manager(dir.path()); let sidecar = AgentHostSidecar::bind_tcp( @@ -1531,24 +1605,36 @@ mod tests { Some("localhost".to_string()), LoopbackAuth::Token("tok".to_string()), Some("my-tunnel".to_string()), - lockfile.clone(), + user_data_path.clone(), + "instance-a".to_string(), ) .await .unwrap(); - let metadata = read_agent_host_metadata(&lockfile).unwrap().unwrap(); - assert_eq!(metadata.pid, std::process::id()); - assert_eq!(metadata.host.as_deref(), Some("localhost")); - assert_eq!(metadata.port, sidecar.bound_addr().port()); - assert_ne!(metadata.port, 0); - assert_eq!(metadata.connection_token.as_deref(), Some("tok")); - assert_eq!(metadata.tunnel_name.as_deref(), Some("my-tunnel")); + let entries = + agent_host_registry::read_registry(&log::Logger::test(), &user_data_path).unwrap(); + assert_eq!(entries.len(), 1); + let entry = &entries[0]; + assert_eq!(entry.server_type, AgentHostServerType::Standalone); + assert_eq!(entry.pid, std::process::id()); + assert_eq!(entry.instance_id, "instance-a"); + assert_eq!(entry.connection_token, "tok"); + assert_eq!(entry.tunnel_name.as_deref(), Some("my-tunnel")); + assert_eq!(entry.protocol_version, AGENT_HOST_PROTOCOL_VERSION); + match &entry.endpoint { + AgentHostEndpointAddress::Tcp { host, port } => { + assert_eq!(host, "localhost"); + assert_eq!(*port, sidecar.bound_addr().port()); + assert_ne!(*port, 0); + } + other => panic!("expected a tcp endpoint, got {:?}", other), + } } #[tokio::test] - async fn drop_removes_lockfile_when_pid_matches() { + async fn drop_removes_registry_entry_matching_our_identity_without_blocking_worker() { let dir = tempfile::tempdir().unwrap(); - let lockfile = dir.path().join("agent-host.lock"); + let user_data_path = dir.path().join("user-data"); let manager = make_test_manager(dir.path()); { @@ -1559,20 +1645,43 @@ mod tests { None, LoopbackAuth::Disabled, None, - lockfile.clone(), + user_data_path.clone(), + "instance-fallback".to_string(), ) .await .unwrap(); - assert!(lockfile.exists()); + assert_eq!( + agent_host_registry::read_registry(&log::Logger::test(), &user_data_path) + .unwrap() + .len(), + 1 + ); } - assert!(!lockfile.exists()); + // `shutdown` was never called, so `Drop`'s fallback cleanup is + // responsible here. It dispatches the blocking removal to a + // separate blocking-safe thread rather than doing it inline on + // this async task's worker, so poll briefly for it to land instead + // of asserting immediately. + let deadline = Instant::now() + Duration::from_secs(2); + loop { + if agent_host_registry::read_registry(&log::Logger::test(), &user_data_path) + .unwrap() + .is_empty() + { + break; + } + if Instant::now() >= deadline { + panic!("drop's fallback registry cleanup did not complete in time"); + } + tokio::time::sleep(Duration::from_millis(20)).await; + } } #[tokio::test] - async fn shutdown_leaves_lockfile_when_pid_was_overwritten() { + async fn shutdown_leaves_registry_entry_owned_by_a_different_instance() { let dir = tempfile::tempdir().unwrap(); - let lockfile = dir.path().join("agent-host.lock"); + let user_data_path = dir.path().join("user-data"); let manager = make_test_manager(dir.path()); let sidecar = AgentHostSidecar::bind_tcp( @@ -1582,96 +1691,180 @@ mod tests { None, LoopbackAuth::Disabled, None, - lockfile.clone(), + user_data_path.clone(), + "instance-c".to_string(), ) .await .unwrap(); - // Simulate another process taking over the same lockfile path. - let foreign_pid = std::process::id().wrapping_add(1); - write_agent_host_metadata(&lockfile, &AgentHostMetadata::new(foreign_pid, 9999)).unwrap(); + // Simulate another live process taking over with a distinct + // instance ID; `shutdown`/`Drop` must only ever remove the entry + // exactly matching our own `(type, pid, instanceId)` identity. + let foreign = AgentHostEndpointMetadata::new_standalone( + std::process::id(), + "instance-foreign".to_string(), + "127.0.0.1".to_string(), + 9999, + String::new(), + AGENT_HOST_PROTOCOL_VERSION.to_string(), + None, + None, + ); + agent_host_registry::publish_agent_host_endpoint( + &log::Logger::test(), + &user_data_path, + &foreign, + ) + .unwrap(); sidecar.shutdown().await; - let preserved = read_agent_host_metadata(&lockfile).unwrap().unwrap(); - assert_eq!(preserved.pid, foreign_pid); - assert_eq!(preserved.port, 9999); + let entries = + agent_host_registry::read_registry(&log::Logger::test(), &user_data_path).unwrap(); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].instance_id, "instance-foreign"); + } + + #[tokio::test] + async fn drop_does_not_redundantly_clean_up_after_shutdown() { + let dir = tempfile::tempdir().unwrap(); + let user_data_path = dir.path().join("user-data"); + let manager = make_test_manager(dir.path()); + + let sidecar = AgentHostSidecar::bind_tcp( + log::Logger::test(), + manager, + SocketAddr::from(([127, 0, 0, 1], 0)), + None, + LoopbackAuth::Disabled, + None, + user_data_path.clone(), + "instance-shutdown-then-drop".to_string(), + ) + .await + .unwrap(); + + sidecar.shutdown().await; + assert!( + agent_host_registry::read_registry(&log::Logger::test(), &user_data_path) + .unwrap() + .is_empty() + ); + + // Republish an entry reusing our own identity, simulating a case + // where some other writer took over that exact (type, pid, + // instanceId) slot right after `shutdown` removed it. If `Drop` + // were to redundantly repeat cleanup after `shutdown` already + // claimed it, it would incorrectly remove this entry too. + let republished = AgentHostEndpointMetadata::new_standalone( + std::process::id(), + "instance-shutdown-then-drop".to_string(), + "127.0.0.1".to_string(), + 9999, + String::new(), + AGENT_HOST_PROTOCOL_VERSION.to_string(), + None, + None, + ); + agent_host_registry::publish_agent_host_endpoint( + &log::Logger::test(), + &user_data_path, + &republished, + ) + .unwrap(); + + drop(sidecar); + + // Give any (incorrectly) dispatched fallback cleanup a moment to + // run before asserting it left the republished entry untouched. + tokio::time::sleep(Duration::from_millis(100)).await; + + let entries = + agent_host_registry::read_registry(&log::Logger::test(), &user_data_path).unwrap(); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].instance_id, "instance-shutdown-then-drop"); } #[test] - fn classify_returns_spawn_fresh_when_lockfile_missing() { + fn classify_agent_host_returns_spawn_fresh_when_registry_empty() { let dir = tempfile::tempdir().unwrap(); - let lockfile = dir.path().join("missing.lock"); + let user_data_path = dir.path().join("user-data"); - let decision = classify_agent_host_lockfile(&log::Logger::test(), &lockfile); + let decision = classify_agent_host(&log::Logger::test(), &user_data_path); - assert_eq!(decision, AgentHostLockfileDecision::SpawnFresh); + assert_eq!(decision, AgentHostReuseDecision::SpawnFresh); } #[test] - fn classify_returns_spawn_fresh_for_stale_pid() { + fn classify_agent_host_prefers_live_registry_standalone_entry() { let dir = tempfile::tempdir().unwrap(); - let lockfile = dir.path().join("agent-host.lock"); - // Use a PID that is extremely unlikely to exist (max u32 - 1). - // `process_exists` returns false for unknown PIDs. - let mut metadata = AgentHostMetadata::new(u32::MAX - 1, 1234); - metadata.connection_token = Some("ignored".to_string()); - write_agent_host_metadata(&lockfile, &metadata).unwrap(); - - let decision = classify_agent_host_lockfile(&log::Logger::test(), &lockfile); - - assert_eq!(decision, AgentHostLockfileDecision::SpawnFresh); - } - - #[test] - fn classify_returns_reuse_for_live_compatible_lockfile() { - let dir = tempfile::tempdir().unwrap(); - let lockfile = dir.path().join("agent-host.lock"); + let user_data_path = dir.path().join("user-data"); let pid = std::process::id(); - let mut metadata = AgentHostMetadata::new(pid, 4321); - metadata.host = Some("127.0.0.1".to_string()); - metadata.connection_token = Some("tok".to_string()); - write_agent_host_metadata(&lockfile, &metadata).unwrap(); - let decision = classify_agent_host_lockfile(&log::Logger::test(), &lockfile); + let entry = AgentHostEndpointMetadata::new_standalone( + pid, + "instance-registry".to_string(), + "127.0.0.1".to_string(), + 4321, + "registry-tok".to_string(), + AGENT_HOST_PROTOCOL_VERSION.to_string(), + None, + None, + ); + agent_host_registry::publish_agent_host_endpoint( + &log::Logger::test(), + &user_data_path, + &entry, + ) + .unwrap(); + + let decision = classify_agent_host(&log::Logger::test(), &user_data_path); assert_eq!( decision, - AgentHostLockfileDecision::Reuse { + AgentHostReuseDecision::Reuse { pid, host: Some("127.0.0.1".to_string()), port: 4321, - token: Some("tok".to_string()), + token: Some("registry-tok".to_string()), tunnel_name: None, + instance_id: "instance-registry".to_string(), } ); } #[test] - fn classify_returns_reuse_for_live_newer_protocol() { - // The agent host server is downloaded on demand and may speak a - // newer protocol than the CLI is built with; we must still treat - // it as reusable rather than refusing to share it. + fn classify_agent_host_never_selects_an_editor_registry_entry() { let dir = tempfile::tempdir().unwrap(); - let lockfile = dir.path().join("agent-host.lock"); - let pid = std::process::id(); - let mut metadata = AgentHostMetadata::new(pid, 4321); - metadata.protocol_version = "9.9.9".to_string(); - metadata.connection_token = Some("tok".to_string()); - write_agent_host_metadata(&lockfile, &metadata).unwrap(); + let user_data_path = dir.path().join("user-data"); - let decision = classify_agent_host_lockfile(&log::Logger::test(), &lockfile); + let editor = AgentHostEndpointMetadata { + schema_version: + crate::tunnels::agent_host_registry::AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION, + server_type: AgentHostServerType::Editor, + pid: std::process::id(), + instance_id: "editor-instance".to_string(), + protocol_version: AGENT_HOST_PROTOCOL_VERSION.to_string(), + connection_token: "editor-tok".to_string(), + endpoint: AgentHostEndpointAddress::Socket { + path: "/tmp/editor.sock".to_string(), + }, + quality: None, + tunnel_name: None, + }; + agent_host_registry::publish_agent_host_endpoint( + &log::Logger::test(), + &user_data_path, + &editor, + ) + .unwrap(); - assert_eq!( - decision, - AgentHostLockfileDecision::Reuse { - pid, - host: None, - port: 4321, - token: Some("tok".to_string()), - tunnel_name: None, - } - ); + // With only an (ignored) editor entry present, the caller must be + // told to spawn a fresh standalone supervisor rather than ever + // touching the editor entry. + let decision = classify_agent_host(&log::Logger::test(), &user_data_path); + + assert_eq!(decision, AgentHostReuseDecision::SpawnFresh); } #[test] diff --git a/cli/src/tunnels/agent_host_metadata.rs b/cli/src/tunnels/agent_host_metadata.rs deleted file mode 100644 index 855081a4fc8..00000000000 --- a/cli/src/tunnels/agent_host_metadata.rs +++ /dev/null @@ -1,215 +0,0 @@ -/*--------------------------------------------------------------------------------------------- - * Copyright (c) Microsoft Corporation. All rights reserved. - * Licensed under the MIT License. See License.txt in the project root for license information. - *--------------------------------------------------------------------------------------------*/ - -use std::fs; -use std::io::{self, Write}; -use std::path::Path; - -use serde::{Deserialize, Serialize}; - -pub const AGENT_HOST_METADATA_SCHEMA_VERSION: u32 = 1; -pub const AGENT_HOST_PROTOCOL_VERSION: &str = "0.1.0"; - -/// Persisted record describing a running `code agent host` proxy, written to -/// the per-quality lockfile (`/agent-host-.lock`). -/// -/// This schema is shared with the TypeScript SSH client in -/// `src/vs/platform/agentHost/common/remoteAgentHostMetadata.ts`; field -/// renames or removals MUST be coordinated across both languages. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct AgentHostMetadata { - pub schema_version: u32, - pub pid: u32, - pub port: u16, - /// Host the supervisor's TCP listener was bound to (e.g. `127.0.0.1`, - /// `0.0.0.0`). Optional so older lockfiles still parse; consumers - /// fall back to loopback when absent. Used by the foreground - /// `code agent host` command to detect when a caller's requested - /// `--host` differs from what's already running. - #[serde(skip_serializing_if = "Option::is_none")] - pub host: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub connection_token: Option, - pub protocol_version: String, - #[serde(skip_serializing_if = "Option::is_none")] - pub quality: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub tunnel_name: Option, -} - -impl AgentHostMetadata { - pub fn new(pid: u32, port: u16) -> Self { - Self { - schema_version: AGENT_HOST_METADATA_SCHEMA_VERSION, - pid, - port, - host: None, - connection_token: None, - protocol_version: AGENT_HOST_PROTOCOL_VERSION.to_string(), - quality: None, - tunnel_name: None, - } - } -} - -pub fn read_agent_host_metadata(path: &Path) -> io::Result> { - let text = match fs::read_to_string(path) { - Ok(text) => text, - Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(None), - Err(err) => return Err(err), - }; - - serde_json::from_str(&text) - .map(Some) - .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err)) -} - -pub fn write_agent_host_metadata(path: &Path, metadata: &AgentHostMetadata) -> io::Result<()> { - #[cfg(not(windows))] - use std::os::unix::fs::PermissionsExt; - - let parent = path.parent().unwrap_or_else(|| Path::new(".")); - fs::create_dir_all(parent)?; - #[cfg(not(windows))] - fs::set_permissions(parent, fs::Permissions::from_mode(0o700))?; - - let mut temp = tempfile::NamedTempFile::new_in(parent)?; - #[cfg(not(windows))] - temp.as_file() - .set_permissions(fs::Permissions::from_mode(0o600))?; - temp.write_all(serde_json::to_string(metadata)?.as_bytes())?; - temp.flush()?; - temp.persist(path).map_err(|err| err.error)?; - #[cfg(not(windows))] - fs::set_permissions(path, fs::Permissions::from_mode(0o600))?; - Ok(()) -} - -pub fn remove_agent_host_metadata(path: &Path) -> io::Result<()> { - match fs::remove_file(path) { - Ok(()) => Ok(()), - Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(()), - Err(err) => Err(err), - } -} - -pub fn remove_agent_host_metadata_for_pid(path: &Path, pid: u32) -> io::Result<()> { - match read_agent_host_metadata(path)? { - Some(metadata) if metadata.pid == pid => remove_agent_host_metadata(path), - _ => Ok(()), - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::path::PathBuf; - - fn metadata_path(dir: &Path) -> PathBuf { - dir.join("agent-host.lock") - } - - #[test] - fn serializes_with_camel_case_fields() { - let mut metadata = AgentHostMetadata::new(1234, 8080); - metadata.host = Some("0.0.0.0".to_string()); - metadata.connection_token = Some("tok".to_string()); - metadata.quality = Some("insider".to_string()); - metadata.tunnel_name = Some("my-tunnel".to_string()); - - let value = serde_json::to_value(&metadata).unwrap(); - assert_eq!(value["schemaVersion"], 1); - assert_eq!(value["pid"], 1234); - assert_eq!(value["port"], 8080); - assert_eq!(value["host"], "0.0.0.0"); - assert_eq!(value["connectionToken"], "tok"); - assert_eq!(value["protocolVersion"], "0.1.0"); - assert_eq!(value["quality"], "insider"); - assert_eq!(value["tunnelName"], "my-tunnel"); - } - - #[test] - fn omits_optional_fields_when_unset() { - let metadata = AgentHostMetadata::new(1234, 8080); - - let value = serde_json::to_value(&metadata).unwrap(); - assert!(value.get("host").is_none()); - assert!(value.get("connectionToken").is_none()); - assert!(value.get("quality").is_none()); - assert!(value.get("tunnelName").is_none()); - } - - #[test] - fn round_trips_metadata() { - let dir = tempfile::tempdir().unwrap(); - let path = metadata_path(dir.path()); - let mut metadata = AgentHostMetadata::new(1234, 8080); - metadata.connection_token = Some("tok".to_string()); - - write_agent_host_metadata(&path, &metadata).unwrap(); - assert_eq!(read_agent_host_metadata(&path).unwrap(), Some(metadata)); - } - - #[test] - fn missing_metadata_returns_none() { - let dir = tempfile::tempdir().unwrap(); - let path = metadata_path(dir.path()); - - assert_eq!(read_agent_host_metadata(&path).unwrap(), None); - } - - #[test] - fn invalid_metadata_returns_invalid_data() { - let dir = tempfile::tempdir().unwrap(); - let path = metadata_path(dir.path()); - fs::write(&path, "not json").unwrap(); - - let err = read_agent_host_metadata(&path).unwrap_err(); - assert_eq!(err.kind(), io::ErrorKind::InvalidData); - } - - #[test] - fn remove_metadata_ignores_missing_file() { - let dir = tempfile::tempdir().unwrap(); - let path = metadata_path(dir.path()); - - remove_agent_host_metadata(&path).unwrap(); - } - - #[test] - fn remove_metadata_for_pid_only_removes_matching_process() { - let dir = tempfile::tempdir().unwrap(); - let path = metadata_path(dir.path()); - let metadata = AgentHostMetadata::new(1234, 8080); - - write_agent_host_metadata(&path, &metadata).unwrap(); - remove_agent_host_metadata_for_pid(&path, 4321).unwrap(); - assert_eq!(read_agent_host_metadata(&path).unwrap(), Some(metadata)); - - remove_agent_host_metadata_for_pid(&path, 1234).unwrap(); - assert_eq!(read_agent_host_metadata(&path).unwrap(), None); - } - - #[cfg(not(windows))] - #[test] - fn writes_owner_only_permissions() { - use std::os::unix::fs::PermissionsExt; - - let dir = tempfile::tempdir().unwrap(); - let path = metadata_path(dir.path()); - - write_agent_host_metadata(&path, &AgentHostMetadata::new(1234, 8080)).unwrap(); - - assert_eq!( - fs::metadata(dir.path()).unwrap().permissions().mode() & 0o777, - 0o700 - ); - assert_eq!( - fs::metadata(&path).unwrap().permissions().mode() & 0o777, - 0o600 - ); - } -} diff --git a/cli/src/tunnels/agent_host_registry.rs b/cli/src/tunnels/agent_host_registry.rs new file mode 100644 index 00000000000..38de8a82f53 --- /dev/null +++ b/cli/src/tunnels/agent_host_registry.rs @@ -0,0 +1,1348 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +//! Rust CLI writer/reader for the shared local agent-host endpoint registry +//! at `/agent-host/local-endpoint/metadata.json`. +//! +//! This schema and the multi-writer locking/upsert/removal protocol are +//! shared with, and MUST stay in lock-step with, the TypeScript +//! implementation in: +//! - `src/vs/platform/agentHost/common/agentHostEndpointRegistry.ts` (schema +//! + pure array helpers) +//! - `src/vs/platform/agentHost/node/localAgentHostMetadata.ts` (lock, +//! atomic write, directory security) +//! - `src/vs/platform/agentHost/LOCAL_ENDPOINT.md` (protocol document) +//! +//! The standalone `code agent host` CLI only ever publishes a `standalone` +//! entry with a `tcp` endpoint; it never publishes (and must never select +//! for reuse/`--replace`) an `editor` entry, since those are owned by +//! running VS Code windows. + +use std::collections::HashMap; +use std::fs::{self, File, OpenOptions}; +use std::io::{self, Write}; +use std::path::{Path, PathBuf}; +use std::time::{Duration, Instant}; + +use serde::{Deserialize, Serialize}; + +use crate::log; +use crate::util::machine::process_exists; + +/// Schema version for the shared registry. See module docs: version 1 was +/// the editor-only, socket-path-only shape; version 2 generalizes the +/// registry to hold both editor (socket/pipe) and standalone CLI (TCP) +/// endpoints. Entries with any other `schemaVersion` are ignored on read +/// (not rejected wholesale) so one writer's unsupported entry can never hide +/// another live writer's endpoint. +pub const AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION: u32 = 2; + +/// AHP protocol version this CLI build implements, recorded on every +/// registry entry it publishes. +pub const AGENT_HOST_PROTOCOL_VERSION: &str = "0.1.0"; + +const METADATA_DIRECTORY_NAME: &str = "agent-host"; +const ENDPOINT_DIRECTORY_NAME: &str = "local-endpoint"; +const METADATA_FILE_NAME: &str = "metadata.json"; + +/// How long [`publish_agent_host_endpoint`] waits to acquire the write lock +/// before giving up. Mirrors the TS `asyncLockAcquireTimeoutMs`, used on the +/// (infrequent, startup-time) publish path. +const PUBLISH_LOCK_ACQUIRE_TIMEOUT: Duration = Duration::from_millis(3000); +const PUBLISH_LOCK_RETRY_DELAY: Duration = Duration::from_millis(40); + +/// How long [`remove_agent_host_endpoint`] waits to acquire the write lock. +/// Mirrors the TS `syncLockAcquireTimeoutMs`, kept short since this runs on +/// the shutdown path and must not noticeably delay process exit. +const CLEANUP_LOCK_ACQUIRE_TIMEOUT: Duration = Duration::from_millis(500); +const CLEANUP_LOCK_RETRY_DELAY: Duration = Duration::from_millis(10); + +/// Grace period for a lock directory whose owner file has not appeared yet, +/// to avoid racing a concurrent acquirer that is mid-write. +const LOCK_OWNER_GRACE: Duration = Duration::from_millis(2000); + +// ---- Schema ----------------------------------------------------------------- + +/// Kind of process that owns an agent host endpoint. Controls +/// ownership/default-selection policy on the client; it is not by itself a +/// measure of trust or of registry identity (see [`AgentHostEndpointIdentity`]). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum AgentHostServerType { + Editor, + Standalone, +} + +/// How to physically connect to an endpoint. Editor endpoints are always a +/// Unix domain socket or Windows named pipe; the standalone Rust CLI +/// currently only ever publishes a `tcp` endpoint. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "lowercase")] +pub enum AgentHostEndpointAddress { + Socket { path: String }, + Tcp { host: String, port: u16 }, +} + +/// One entry of the shared local agent-host endpoint registry. The registry +/// file itself is a JSON array of these entries. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentHostEndpointMetadata { + pub schema_version: u32, + #[serde(rename = "type")] + pub server_type: AgentHostServerType, + pub pid: u32, + pub instance_id: String, + pub protocol_version: String, + pub connection_token: String, + pub endpoint: AgentHostEndpointAddress, + #[serde(skip_serializing_if = "Option::is_none")] + pub quality: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub tunnel_name: Option, +} + +/// The subset of [`AgentHostEndpointMetadata`] that identifies a unique +/// registry entry/owner: `(type, pid, instanceId)`. `instance_id` makes +/// identity safe across PID reuse; `pid` alone is not a safe key because +/// operating systems recycle PIDs. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct AgentHostEndpointIdentity { + pub server_type: AgentHostServerType, + pub pid: u32, + pub instance_id: String, +} + +impl AgentHostEndpointMetadata { + /// Builds a `standalone` entry, as published by `code agent host`. + #[allow(clippy::too_many_arguments)] + pub fn new_standalone( + pid: u32, + instance_id: String, + host: String, + port: u16, + connection_token: String, + protocol_version: String, + quality: Option, + tunnel_name: Option, + ) -> Self { + Self { + schema_version: AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION, + server_type: AgentHostServerType::Standalone, + pid, + instance_id, + protocol_version, + connection_token, + endpoint: AgentHostEndpointAddress::Tcp { host, port }, + quality, + tunnel_name, + } + } + + pub fn identity(&self) -> AgentHostEndpointIdentity { + AgentHostEndpointIdentity { + server_type: self.server_type, + pid: self.pid, + instance_id: self.instance_id.clone(), + } + } + + /// The connectable address as a short human string: `host:port` for a + /// `tcp` endpoint, or the raw socket/pipe path for a `socket` endpoint. + pub fn address_label(&self) -> String { + match &self.endpoint { + AgentHostEndpointAddress::Tcp { host, port } => format!("{host}:{port}"), + AgentHostEndpointAddress::Socket { path } => path.clone(), + } + } + + /// A stable, human-readable identifier for disambiguating this entry + /// from other discovered hosts in `ps`/`logs`/`stop`/`kill` output. Not + /// meant to be machine-parsed; see [`Self::identity`] for that. + pub fn label(&self) -> String { + let kind = match self.server_type { + AgentHostServerType::Editor => "editor", + AgentHostServerType::Standalone => "standalone", + }; + let mut label = format!("{kind} (pid {}, {})", self.pid, self.address_label()); + if let Some(quality) = &self.quality { + label.push_str(&format!(" [{quality}]")); + } + if let Some(tunnel_name) = &self.tunnel_name { + label.push_str(&format!(" (tunnel {tunnel_name})")); + } + label + } +} + +/// Deterministic ordering key for [`AgentHostServerType`] used when sorting +/// discovered endpoints, so iteration/print order is stable across runs +/// regardless of registry file write order. Standalone sorts first since it +/// is the more likely single-instance case users expect to see first. +fn server_type_sort_rank(server_type: AgentHostServerType) -> u8 { + match server_type { + AgentHostServerType::Standalone => 0, + AgentHostServerType::Editor => 1, + } +} + +fn is_same_identity(a: &AgentHostEndpointMetadata, identity: &AgentHostEndpointIdentity) -> bool { + a.server_type == identity.server_type + && a.pid == identity.pid + && a.instance_id == identity.instance_id +} + +/// Structurally validates one raw registry entry. Unsupported schema versions +/// are ignored; malformed entries return an error for the caller to log. +fn parse_registry_entry( + raw: &serde_json::Value, +) -> Result, String> { + let schema_version = raw + .get("schemaVersion") + .and_then(serde_json::Value::as_u64) + .ok_or_else(|| "schemaVersion must be a positive integer".to_string())?; + if schema_version != AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION as u64 { + return Ok(None); + } + + let entry: AgentHostEndpointMetadata = + serde_json::from_value(raw.clone()).map_err(|error| error.to_string())?; + if entry.pid == 0 { + return Err("pid must be greater than zero".to_string()); + } + if let AgentHostEndpointAddress::Tcp { port, .. } = entry.endpoint { + if port == 0 { + return Err("TCP endpoint port must be greater than zero".to_string()); + } + } + + Ok(Some(entry)) +} + +/// Parses the raw contents of the registry file (expected to be +/// `AgentHostEndpointMetadata[]`). Every entry is validated independently. +fn parse_registry(log: &log::Logger, raw: &[serde_json::Value]) -> Vec { + raw.iter() + .enumerate() + .filter_map(|(index, raw)| match parse_registry_entry(raw) { + Ok(entry) => entry, + Err(error) => { + warning!( + log, + "Ignoring malformed agent host registry entry at index {}: {}", + index, + error + ); + None + } + }) + .collect() +} + +/// Deduplicates `entries` by `(server_type, pid, instance_id)`. If a +/// duplicate identity appears more than once (for example a crashed writer +/// left a stale copy behind before another writer's cleanup ran), the entry +/// encountered later in `entries` wins, since it is presumed to be the more +/// recently written copy. Mirrors `dedupeAgentHostEndpointMetadata` in the TS +/// reference (a `Map` keyed by identity: inserting an already-seen key +/// updates its value but keeps its original iteration position), so the +/// surviving entry's *position* is that of its first occurrence, but its +/// *value* is that of its last occurrence. +fn dedupe_entries(entries: Vec) -> Vec { + let mut result: Vec = Vec::with_capacity(entries.len()); + let mut index_by_identity: HashMap = HashMap::new(); + for entry in entries { + let identity = entry.identity(); + if let Some(&index) = index_by_identity.get(&identity) { + result[index] = entry; + } else { + index_by_identity.insert(identity, result.len()); + result.push(entry); + } + } + result +} + +/// Returns `entries` with any existing entry sharing `metadata`'s identity +/// replaced by `metadata`. +fn upsert_entry( + entries: Vec, + metadata: AgentHostEndpointMetadata, +) -> Vec { + let identity = metadata.identity(); + let mut remaining: Vec<_> = entries + .into_iter() + .filter(|e| !is_same_identity(e, &identity)) + .collect(); + remaining.push(metadata); + remaining +} + +/// Returns `entries` with the exact-identity-matching entry removed, if any. +/// Used on shutdown so a writer only ever removes its own entry, never a +/// newer process's entry that happens to share its PID. +fn remove_entry( + entries: &[AgentHostEndpointMetadata], + identity: &AgentHostEndpointIdentity, +) -> Vec { + entries + .iter() + .filter(|e| !is_same_identity(e, identity)) + .cloned() + .collect() +} + +/// Drops entries whose PID is confirmed dead. Entries are only ever pruned +/// here (i.e. when death is certain via a PID liveness check); a live PID is +/// always kept. +fn prune_dead_entries( + log: &log::Logger, + entries: Vec, +) -> Vec { + entries + .into_iter() + .filter(|e| { + if process_exists(e.pid) { + true + } else { + info!( + log, + "Pruning stale local endpoint registry entry: {:?} PID {} (instance {}) is no longer running", + e.server_type, + e.pid, + e.instance_id + ); + false + } + }) + .collect() +} + +// ---- Paths -------------------------------------------------------------------- + +fn metadata_directory(user_data_path: &Path) -> PathBuf { + user_data_path + .join(METADATA_DIRECTORY_NAME) + .join(ENDPOINT_DIRECTORY_NAME) +} + +/// Path to the shared registry file for `user_data_path`. +fn metadata_path(user_data_path: &Path) -> PathBuf { + metadata_directory(user_data_path).join(METADATA_FILE_NAME) +} + +fn lock_directory_path(metadata_path: &Path) -> PathBuf { + let mut os_string = metadata_path.as_os_str().to_owned(); + os_string.push(".lock"); + PathBuf::from(os_string) +} + +fn lock_owner_file_path(lock_dir: &Path) -> PathBuf { + lock_dir.join("owner.json") +} + +fn temp_write_path(metadata_path: &Path, unique_suffix: &str) -> PathBuf { + let mut os_string = metadata_path.as_os_str().to_owned(); + os_string.push("."); + os_string.push(unique_suffix); + os_string.push(".tmp"); + PathBuf::from(os_string) +} + +// ---- Directory security ------------------------------------------------------- + +/// Creates the metadata directory (if needed) and restricts it to the +/// current user (Unix: `0700`; Windows: an ACL granting only the current +/// user, `SYSTEM`, and `Administrators`), mirroring +/// `prepareLocalAgentHostEndpointMetadataDirectory` in +/// `node/localAgentHostMetadata.ts`. +fn prepare_metadata_directory(user_data_path: &Path) -> io::Result<()> { + let dir = metadata_directory(user_data_path); + fs::create_dir_all(&dir)?; + + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + fs::set_permissions(&dir, fs::Permissions::from_mode(0o700))?; + } + + #[cfg(windows)] + super::agent_host_registry_acl_windows::apply_owner_only_acl(&dir)?; + + Ok(()) +} + +// ---- Read --------------------------------------------------------------------- + +/// Reads and validates every entry in the shared registry, without taking +/// the write lock. Safe to call frequently: the registry file is only ever +/// observed in a fully-written state (via atomic rename). +pub fn read_registry( + log: &log::Logger, + user_data_path: &Path, +) -> io::Result> { + read_registry_at(log, &metadata_path(user_data_path)) +} + +fn read_registry_at(log: &log::Logger, path: &Path) -> io::Result> { + let metadata = match fs::symlink_metadata(path) { + Ok(m) => m, + Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(e) => return Err(e), + }; + if !metadata.is_file() { + // Missing, a directory, or a symlink (defends against a symlink + // swap attack; genuine entries are only ever written via rename). + return Ok(Vec::new()); + } + + let raw = match fs::read_to_string(path) { + Ok(s) => s, + Err(e) if e.kind() == io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(e) => return Err(e), + }; + + let values: Vec = match serde_json::from_str(&raw) { + Ok(v) => v, + Err(error) => { + warning!( + log, + "Ignoring malformed agent host endpoint registry at {}: {}", + path.display(), + error + ); + return Ok(Vec::new()); + } + }; + + Ok(parse_registry(log, &values)) +} + +// ---- Atomic write --------------------------------------------------------------- + +fn write_registry_atomic( + path: &Path, + unique_suffix: &str, + entries: &[AgentHostEndpointMetadata], +) -> io::Result<()> { + let temp_path = temp_write_path(path, unique_suffix); + let json = serde_json::to_vec(entries)?; + + let mut open_options = OpenOptions::new(); + open_options.write(true).create_new(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + open_options.mode(0o600); + } + + { + let mut file = open_options.open(&temp_path)?; + file.write_all(&json)?; + file.sync_all()?; + } + + let rename_result = fs::rename(&temp_path, path); + // Best-effort cleanup, mirroring the TS `finally { rm force }`: a no-op + // once the rename above succeeded. + let _ = fs::remove_file(&temp_path); + rename_result?; + + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + fs::set_permissions(path, fs::Permissions::from_mode(0o600))?; + } + + Ok(()) +} + +// ---- Multi-writer lock ---------------------------------------------------------- +// +// The lock is a sibling directory to metadata.json (metadata.json.lock). +// Directory creation without `recursive` is used as the exclusive-acquire +// primitive because it is atomic on every platform we support and needs no +// native/optional dependency. The lock holder's `(pid, instanceId)` is +// written into an owner file inside the directory so a contending process +// can recognize and reclaim an abandoned lock: if the recorded PID is +// confirmed dead, the lock is stale and is reclaimed immediately; otherwise +// acquisition is retried until a bounded timeout elapses, after which the +// caller must log and continue running undiscoverable rather than silently +// bypassing the lock. Mirrors `node/localAgentHostMetadata.ts`. + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +struct LockOwner { + pid: u32, + instance_id: String, +} + +/// Holds the sibling lock directory for the duration of a registry +/// read-modify-write. Releasing (removing the lock directory) happens on +/// drop, but only if the owner recorded on disk still matches this guard +/// (i.e. nobody else has reclaimed it as stale in the meantime). +struct RegistryLock { + lock_dir: PathBuf, + owner: LockOwner, +} + +impl Drop for RegistryLock { + fn drop(&mut self) { + if let Some(current) = read_lock_owner(&self.lock_dir) { + if current != self.owner { + // Another process already reclaimed this lock as stale; it + // now owns this lock's lifecycle, so leave it alone. + return; + } + } + let _ = fs::remove_dir_all(&self.lock_dir); + } +} + +fn read_lock_owner(lock_dir: &Path) -> Option { + let raw = fs::read_to_string(lock_owner_file_path(lock_dir)).ok()?; + serde_json::from_str(&raw).ok() +} + +fn is_lock_directory_stale_without_owner(lock_dir: &Path) -> bool { + match fs::metadata(lock_dir) { + Ok(metadata) => match metadata.modified() { + Ok(modified) => match modified.elapsed() { + Ok(elapsed) => elapsed > LOCK_OWNER_GRACE, + Err(_) => false, + }, + Err(_) => false, + }, + // The directory disappeared already (another process reclaimed + // it); let the caller retry acquisition. + Err(_) => true, + } +} + +fn try_reclaim_stale_lock(lock_dir: &Path, log: &log::Logger) -> bool { + let owner = read_lock_owner(lock_dir); + match &owner { + Some(owner) if process_exists(owner.pid) => return false, + Some(_) => {} + None => { + if !is_lock_directory_stale_without_owner(lock_dir) { + return false; + } + } + } + + if fs::remove_dir_all(lock_dir).is_err() { + return false; + } + + warning!( + log, + "Reclaimed a stale local agent host endpoint registry lock{}", + match owner { + Some(o) => format!(" from PID {}", o.pid), + None => String::new(), + } + ); + true +} + +fn acquire_registry_lock( + lock_dir: &Path, + owner: &LockOwner, + timeout: Duration, + retry_delay: Duration, + log: &log::Logger, +) -> io::Result> { + let deadline = Instant::now() + timeout; + loop { + match fs::create_dir(lock_dir) { + Ok(()) => { + write_lock_owner(lock_dir, owner)?; + return Ok(Some(RegistryLock { + lock_dir: lock_dir.to_path_buf(), + owner: owner.clone(), + })); + } + Err(e) if e.kind() == io::ErrorKind::AlreadyExists => { + if try_reclaim_stale_lock(lock_dir, log) { + continue; + } + if Instant::now() >= deadline { + return Ok(None); + } + std::thread::sleep(retry_delay); + } + Err(e) => return Err(e), + } + } +} + +fn write_lock_owner(lock_dir: &Path, owner: &LockOwner) -> io::Result<()> { + let mut open_options = OpenOptions::new(); + open_options.write(true).create(true).truncate(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + open_options.mode(0o600); + } + let mut file: File = open_options.open(lock_owner_file_path(lock_dir))?; + file.write_all(serde_json::to_string(owner)?.as_bytes()) +} + +// ---- Publish / remove ----------------------------------------------------------- + +/// Upserts `metadata` into the shared local agent host endpoint registry. +/// +/// Multiple processes (editor windows and the standalone `code agent host` +/// CLI) can publish to the same registry file concurrently, so this +/// acquires the sibling lock first, serializing the read-prune-upsert-write +/// sequence across all writers. Readers remain lock-free because the final +/// write is an atomic rename. +/// +/// Returns an error if the lock cannot be acquired within a bounded timeout, +/// or if any filesystem operation fails; callers MUST treat that as +/// "continue running, but undiscoverable" and must not fall back to a +/// non-atomic write. +pub fn publish_agent_host_endpoint( + log: &log::Logger, + user_data_path: &Path, + metadata: &AgentHostEndpointMetadata, +) -> io::Result<()> { + prepare_metadata_directory(user_data_path)?; + let path = metadata_path(user_data_path); + let lock_dir = lock_directory_path(&path); + let owner = LockOwner { + pid: metadata.pid, + instance_id: metadata.instance_id.clone(), + }; + + let _lock = acquire_registry_lock( + &lock_dir, + &owner, + PUBLISH_LOCK_ACQUIRE_TIMEOUT, + PUBLISH_LOCK_RETRY_DELAY, + log, + )? + .ok_or_else(|| { + io::Error::new( + io::ErrorKind::TimedOut, + format!( + "Timed out acquiring the local agent host endpoint registry lock at {}", + lock_dir.display() + ), + ) + })?; + + let current = read_registry_at(log, &path)?; + let live = prune_dead_entries(log, current); + let deduped = dedupe_entries(live); + let next = upsert_entry(deduped, metadata.clone()); + write_registry_atomic(&path, &metadata.instance_id, &next) +} + +/// Removes exactly `identity`'s `(type, pid, instanceId)` entry from the +/// registry, reacquiring the write lock first. Deletes the file entirely +/// only when the resulting registry is empty. Best-effort: failures are +/// logged, never returned as fatal, so process shutdown is never blocked by +/// cleanup. +pub fn remove_agent_host_endpoint( + log: &log::Logger, + user_data_path: &Path, + identity: &AgentHostEndpointIdentity, +) { + let path = metadata_path(user_data_path); + let lock_dir = lock_directory_path(&path); + let owner = LockOwner { + pid: identity.pid, + instance_id: identity.instance_id.clone(), + }; + + let lock = match acquire_registry_lock( + &lock_dir, + &owner, + CLEANUP_LOCK_ACQUIRE_TIMEOUT, + CLEANUP_LOCK_RETRY_DELAY, + log, + ) { + Ok(Some(lock)) => lock, + Ok(None) => { + warning!( + log, + "Timed out acquiring the local agent host endpoint registry lock while removing our entry from {}", + path.display() + ); + return; + } + Err(e) => { + warning!( + log, + "Failed to acquire the local agent host endpoint registry lock: {}", + e + ); + return; + } + }; + + let result = (|| -> io::Result<()> { + let current = read_registry_at(log, &path)?; + let remaining = remove_entry(¤t, identity); + if remaining.len() == current.len() { + return Ok(()); + } + if remaining.is_empty() { + match fs::remove_file(&path) { + Ok(()) => Ok(()), + Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()), + Err(e) => Err(e), + } + } else { + write_registry_atomic(&path, &identity.instance_id, &remaining) + } + })(); + + if let Err(e) = result { + warning!( + log, + "Failed to remove our entry from the local agent host endpoint registry: {}", + e + ); + } + + drop(lock); +} + +// ---- Endpoint enumeration --------------------------------------------------------- + +/// Reads the registry and returns every live endpoint (both `editor` and +/// `standalone`, both `socket`/pipe and `tcp` addresses), deduped by +/// identity and pruned of dead-process entries, in a stable deterministic +/// order (standalone before editor, then by `instanceId`). +/// +/// This is the general-purpose discovery primitive backing `code agent +/// ps|logs|stop`'s auto-discovery: every live entry here is a candidate +/// host to query, not just the one this process would reuse for `code +/// agent host`. +pub fn list_live_endpoints( + log: &log::Logger, + user_data_path: &Path, +) -> Vec { + let entries = match read_registry(log, user_data_path) { + Ok(entries) => entries, + Err(e) => { + debug!( + log, + "Could not read the local agent host endpoint registry at {}: {}", + metadata_path(user_data_path).display(), + e + ); + return Vec::new(); + } + }; + + let live = prune_dead_entries(log, entries); + let mut deduped = dedupe_entries(live); + deduped.sort_by(|a, b| { + server_type_sort_rank(a.server_type) + .cmp(&server_type_sort_rank(b.server_type)) + .then_with(|| a.instance_id.cmp(&b.instance_id)) + }); + deduped +} + +/// Like [`list_live_endpoints`], but restricted to `standalone` entries. +/// `editor` entries are never included here: they are owned by running VS +/// Code windows and must never be selected, replaced, or killed by the +/// standalone CLI. Used by `code agent kill`'s multi-instance +/// disambiguation as well as by [`select_live_standalone_endpoint`]. +pub fn list_live_standalone_endpoints( + log: &log::Logger, + user_data_path: &Path, +) -> Vec { + list_live_endpoints(log, user_data_path) + .into_iter() + .filter(|e| e.server_type == AgentHostServerType::Standalone) + .collect() +} + +/// A live `standalone` registry entry selected for reuse by `code agent +/// host` / `code agent ps|stop|logs|kill`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LiveStandaloneEndpoint { + pub pid: u32, + pub instance_id: String, + pub host: String, + pub port: u16, + /// Empty when the standalone host was started with + /// `--without-connection-token`. + pub connection_token: String, + pub quality: Option, + pub tunnel_name: Option, +} + +/// Reads the registry and returns a live `standalone` entry to reuse, if +/// any. `editor` entries are never considered: they are owned by running VS +/// Code windows and must never be selected or replaced by the standalone +/// CLI. Only `tcp` entries are considered, since this helper backs `code +/// agent host`'s single-target TCP reuse path. If more than one live +/// standalone entry exists, selection is deterministic (lowest +/// `instanceId`) and a warning is logged recommending `--address` to +/// disambiguate, since there is currently no dedicated "target this +/// instance" flag. +/// +/// Callers that want *every* live standalone entry (e.g. `code agent +/// kill`'s multi-instance disambiguation, which must offer socket/pipe +/// entries too) should use [`list_live_standalone_endpoints`] instead. +pub fn select_live_standalone_endpoint( + log: &log::Logger, + user_data_path: &Path, +) -> Option { + let mut live: Vec = list_live_standalone_endpoints(log, user_data_path) + .into_iter() + .filter_map(|e| match e.endpoint { + AgentHostEndpointAddress::Tcp { host, port } => Some(LiveStandaloneEndpoint { + pid: e.pid, + instance_id: e.instance_id, + host, + port, + connection_token: e.connection_token, + quality: e.quality, + tunnel_name: e.tunnel_name, + }), + AgentHostEndpointAddress::Socket { .. } => None, + }) + .collect(); + + if live.is_empty() { + return None; + } + + if live.len() > 1 { + live.sort_by(|a, b| a.instance_id.cmp(&b.instance_id)); + warning!( + log, + "Multiple live standalone agent hosts are registered; selecting instance {} deterministically. Pass --address to target a specific one.", + live[0].instance_id + ); + } + + live.into_iter().next() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn standalone(pid: u32, instance_id: &str, port: u16) -> AgentHostEndpointMetadata { + AgentHostEndpointMetadata::new_standalone( + pid, + instance_id.to_string(), + "127.0.0.1".to_string(), + port, + "tok".to_string(), + "0.1.0".to_string(), + None, + None, + ) + } + + #[test] + fn serializes_with_camel_case_and_type_tag() { + let mut metadata = standalone(42, "instance-a", 8080); + metadata.quality = Some("insider".to_string()); + metadata.tunnel_name = Some("my-tunnel".to_string()); + + let value = serde_json::to_value(&metadata).unwrap(); + assert_eq!(value["schemaVersion"], 2); + assert_eq!(value["type"], "standalone"); + assert_eq!(value["pid"], 42); + assert_eq!(value["instanceId"], "instance-a"); + assert_eq!(value["protocolVersion"], "0.1.0"); + assert_eq!(value["connectionToken"], "tok"); + assert_eq!(value["endpoint"]["type"], "tcp"); + assert_eq!(value["endpoint"]["host"], "127.0.0.1"); + assert_eq!(value["endpoint"]["port"], 8080); + assert_eq!(value["quality"], "insider"); + assert_eq!(value["tunnelName"], "my-tunnel"); + } + + #[test] + fn omits_optional_fields_when_unset() { + let metadata = standalone(42, "instance-a", 8080); + let value = serde_json::to_value(&metadata).unwrap(); + assert!(value.get("quality").is_none()); + assert!(value.get("tunnelName").is_none()); + } + + #[test] + fn parse_entry_ignores_unsupported_schema_version() { + let raw = serde_json::json!({ + "schemaVersion": 1, + "type": "editor", + "pid": 42, + "instanceId": "a", + "endpointPath": "/tmp/foo.sock", + "connectionToken": "tok", + "protocolVersion": "0.1.0", + }); + assert!(parse_registry_entry(&raw).unwrap().is_none()); + } + + #[test] + fn parse_registry_ignores_unknown_server_type_without_dropping_known_entries() { + let unknown = serde_json::json!({ + "schemaVersion": 2, + "type": "future-host", + "pid": 42, + "instanceId": "unknown", + "protocolVersion": "1", + "connectionToken": "tok", + "endpoint": { "type": "tcp", "host": "127.0.0.1", "port": 8080 } + }); + let known = standalone(43, "known", 8081); + let raw = vec![unknown, serde_json::to_value(&known).unwrap()]; + + assert_eq!(parse_registry(&log::Logger::test(), &raw), vec![known]); + } + + #[test] + fn parse_entry_ignores_zero_pid_and_zero_port() { + let mut raw = serde_json::to_value(standalone(0, "a", 8080)).unwrap(); + assert!(parse_registry_entry(&raw).is_err()); + + raw = serde_json::to_value(standalone(42, "a", 0)).unwrap(); + assert!(parse_registry_entry(&raw).is_err()); + } + + #[test] + fn parse_entry_accepts_well_formed_editor_socket_entry() { + let raw = serde_json::json!({ + "schemaVersion": 2, + "type": "editor", + "pid": 42, + "instanceId": "a", + "protocolVersion": "0.1.0", + "connectionToken": "tok", + "endpoint": { "type": "socket", "path": "/tmp/foo.sock" }, + }); + let entry = parse_registry_entry(&raw).unwrap().unwrap(); + assert_eq!(entry.server_type, AgentHostServerType::Editor); + assert_eq!( + entry.endpoint, + AgentHostEndpointAddress::Socket { + path: "/tmp/foo.sock".to_string() + } + ); + } + + #[test] + fn upsert_replaces_only_matching_identity() { + let entries = vec![standalone(1, "a", 100), standalone(2, "b", 200)]; + let next = upsert_entry(entries, standalone(1, "a", 999)); + + assert_eq!(next.len(), 2); + let a = next.iter().find(|e| e.instance_id == "a").unwrap(); + assert_eq!( + a.endpoint, + AgentHostEndpointAddress::Tcp { + host: "127.0.0.1".to_string(), + port: 999 + } + ); + } + + #[test] + fn upsert_preserves_other_writers_entries() { + let entries = vec![standalone(1, "a", 100)]; + let next = upsert_entry(entries, standalone(2, "b", 200)); + + assert_eq!(next.len(), 2); + assert!(next.iter().any(|e| e.instance_id == "a")); + assert!(next.iter().any(|e| e.instance_id == "b")); + } + + #[test] + fn remove_entry_only_removes_exact_identity() { + let entries = vec![standalone(1, "a", 100), standalone(1, "a-newer", 200)]; + let identity = AgentHostEndpointIdentity { + server_type: AgentHostServerType::Standalone, + pid: 1, + instance_id: "a".to_string(), + }; + let remaining = remove_entry(&entries, &identity); + + assert_eq!(remaining.len(), 1); + assert_eq!(remaining[0].instance_id, "a-newer"); + } + + #[test] + fn dedupe_entries_keeps_first_position_but_last_value_for_duplicate_identity() { + // A crashed writer can leave a stale copy of its own identity behind + // (e.g. an earlier publish that raced a later one before this + // writer's own cleanup ran). The later occurrence's value must win, + // but — matching the TS `Map`-based reference — the surviving + // entry's position is that of the *first* occurrence, not the last. + let stale = standalone(1, "a", 100); + let other = standalone(2, "b", 200); + let fresh = standalone(1, "a", 999); + let entries = vec![stale, other, fresh]; + + let deduped = dedupe_entries(entries); + + assert_eq!(deduped.len(), 2); + assert_eq!(deduped[0].instance_id, "a"); + assert_eq!( + deduped[0].endpoint, + AgentHostEndpointAddress::Tcp { + host: "127.0.0.1".to_string(), + port: 999 + } + ); + assert_eq!(deduped[1].instance_id, "b"); + } + + #[test] + fn prune_dead_entries_keeps_live_and_drops_dead() { + let log = log::Logger::test(); + let live_pid = std::process::id(); + let dead_pid = u32::MAX - 1; + let entries = vec![ + standalone(live_pid, "live", 100), + standalone(dead_pid, "dead", 200), + ]; + + let pruned = prune_dead_entries(&log, entries); + + assert_eq!(pruned.len(), 1); + assert_eq!(pruned[0].instance_id, "live"); + } + + #[test] + fn publish_then_read_round_trips() { + let dir = tempfile::tempdir().unwrap(); + let log = log::Logger::test(); + let metadata = standalone(std::process::id(), "instance-a", 8080); + + publish_agent_host_endpoint(&log, dir.path(), &metadata).unwrap(); + + let entries = read_registry(&log::Logger::test(), dir.path()).unwrap(); + assert_eq!(entries, vec![metadata]); + } + + #[test] + fn publish_preserves_concurrent_writers_entries() { + let dir = tempfile::tempdir().unwrap(); + let log = log::Logger::test(); + let editor = AgentHostEndpointMetadata { + schema_version: AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION, + server_type: AgentHostServerType::Editor, + pid: std::process::id(), + instance_id: "editor-a".to_string(), + protocol_version: "0.1.0".to_string(), + connection_token: "editor-tok".to_string(), + endpoint: AgentHostEndpointAddress::Socket { + path: "/tmp/editor.sock".to_string(), + }, + quality: None, + tunnel_name: None, + }; + publish_agent_host_endpoint(&log, dir.path(), &editor).unwrap(); + + let standalone_entry = standalone(std::process::id(), "standalone-a", 9090); + publish_agent_host_endpoint(&log, dir.path(), &standalone_entry).unwrap(); + + let entries = read_registry(&log::Logger::test(), dir.path()).unwrap(); + assert_eq!(entries.len(), 2); + assert!(entries.contains(&editor)); + assert!(entries.contains(&standalone_entry)); + } + + #[test] + fn publish_prunes_dead_entries_from_other_writers() { + let dir = tempfile::tempdir().unwrap(); + let log = log::Logger::test(); + let dead = standalone(u32::MAX - 1, "dead", 100); + publish_agent_host_endpoint(&log, dir.path(), &dead).unwrap(); + + let live = standalone(std::process::id(), "live", 200); + publish_agent_host_endpoint(&log, dir.path(), &live).unwrap(); + + let entries = read_registry(&log::Logger::test(), dir.path()).unwrap(); + assert_eq!(entries, vec![live]); + } + + #[test] + fn remove_deletes_file_when_registry_becomes_empty() { + let dir = tempfile::tempdir().unwrap(); + let log = log::Logger::test(); + let metadata = standalone(std::process::id(), "instance-a", 8080); + publish_agent_host_endpoint(&log, dir.path(), &metadata).unwrap(); + + remove_agent_host_endpoint(&log, dir.path(), &metadata.identity()); + + assert!(!metadata_path(dir.path()).exists()); + } + + #[test] + fn remove_only_removes_exact_owner_leaving_other_entries() { + let dir = tempfile::tempdir().unwrap(); + let log = log::Logger::test(); + let mine = standalone(std::process::id(), "mine", 8080); + let theirs = standalone(std::process::id(), "theirs", 9090); + publish_agent_host_endpoint(&log, dir.path(), &mine).unwrap(); + publish_agent_host_endpoint(&log, dir.path(), &theirs).unwrap(); + + remove_agent_host_endpoint(&log, dir.path(), &mine.identity()); + + let entries = read_registry(&log::Logger::test(), dir.path()).unwrap(); + assert_eq!(entries, vec![theirs]); + } + + #[test] + fn remove_does_not_delete_newer_process_entry_with_same_pid() { + // Simulates PID reuse: our entry was already overwritten by a + // different instanceId sharing our old PID. Removal must target + // the exact (type, pid, instanceId) tuple, never the PID alone. + let dir = tempfile::tempdir().unwrap(); + let log = log::Logger::test(); + let ours = standalone(std::process::id(), "ours", 8080); + publish_agent_host_endpoint(&log, dir.path(), &ours).unwrap(); + + let newer = standalone(std::process::id(), "newer-same-pid", 9090); + publish_agent_host_endpoint(&log, dir.path(), &newer).unwrap(); + + remove_agent_host_endpoint(&log, dir.path(), &ours.identity()); + + let entries = read_registry(&log::Logger::test(), dir.path()).unwrap(); + assert_eq!(entries, vec![newer]); + } + + #[test] + fn select_live_standalone_ignores_editor_entries() { + let dir = tempfile::tempdir().unwrap(); + let log = log::Logger::test(); + let editor = AgentHostEndpointMetadata { + schema_version: AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION, + server_type: AgentHostServerType::Editor, + pid: std::process::id(), + instance_id: "editor-a".to_string(), + protocol_version: "0.1.0".to_string(), + connection_token: "editor-tok".to_string(), + endpoint: AgentHostEndpointAddress::Socket { + path: "/tmp/editor.sock".to_string(), + }, + quality: None, + tunnel_name: None, + }; + publish_agent_host_endpoint(&log, dir.path(), &editor).unwrap(); + + assert_eq!(select_live_standalone_endpoint(&log, dir.path()), None); + } + + #[test] + fn select_live_standalone_returns_live_entry() { + let dir = tempfile::tempdir().unwrap(); + let log = log::Logger::test(); + let entry = standalone(std::process::id(), "instance-a", 8080); + publish_agent_host_endpoint(&log, dir.path(), &entry).unwrap(); + + let selected = select_live_standalone_endpoint(&log, dir.path()).unwrap(); + assert_eq!(selected.pid, std::process::id()); + assert_eq!(selected.instance_id, "instance-a"); + assert_eq!(selected.host, "127.0.0.1"); + assert_eq!(selected.port, 8080); + } + + #[test] + fn select_live_standalone_is_deterministic_with_multiple_live_entries() { + let dir = tempfile::tempdir().unwrap(); + let log = log::Logger::test(); + let first = standalone(std::process::id(), "b-instance", 8080); + let second = standalone(std::process::id(), "a-instance", 9090); + publish_agent_host_endpoint(&log, dir.path(), &first).unwrap(); + publish_agent_host_endpoint(&log, dir.path(), &second).unwrap(); + + let selected = select_live_standalone_endpoint(&log, dir.path()).unwrap(); + assert_eq!(selected.instance_id, "a-instance"); + } + + #[test] + fn acquire_registry_lock_reclaims_stale_lock_from_dead_pid() { + let dir = tempfile::tempdir().unwrap(); + let log = log::Logger::test(); + let path = metadata_path(dir.path()); + let lock_dir = lock_directory_path(&path); + fs::create_dir_all(lock_dir.parent().unwrap()).unwrap(); + fs::create_dir(&lock_dir).unwrap(); + write_lock_owner( + &lock_dir, + &LockOwner { + pid: u32::MAX - 1, + instance_id: "dead-owner".to_string(), + }, + ) + .unwrap(); + + let owner = LockOwner { + pid: std::process::id(), + instance_id: "new-owner".to_string(), + }; + let lock = acquire_registry_lock( + &lock_dir, + &owner, + Duration::from_millis(500), + Duration::from_millis(10), + &log, + ) + .unwrap(); + + assert!(lock.is_some()); + } + + #[test] + fn acquire_registry_lock_times_out_when_holder_is_alive() { + let dir = tempfile::tempdir().unwrap(); + let log = log::Logger::test(); + let path = metadata_path(dir.path()); + let lock_dir = lock_directory_path(&path); + fs::create_dir_all(lock_dir.parent().unwrap()).unwrap(); + fs::create_dir(&lock_dir).unwrap(); + write_lock_owner( + &lock_dir, + &LockOwner { + pid: std::process::id(), + instance_id: "alive-owner".to_string(), + }, + ) + .unwrap(); + + let owner = LockOwner { + pid: std::process::id(), + instance_id: "contender".to_string(), + }; + let lock = acquire_registry_lock( + &lock_dir, + &owner, + Duration::from_millis(60), + Duration::from_millis(20), + &log, + ) + .unwrap(); + + assert!(lock.is_none()); + } + + fn editor(pid: u32, instance_id: &str, socket_path: &str) -> AgentHostEndpointMetadata { + AgentHostEndpointMetadata { + schema_version: AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION, + server_type: AgentHostServerType::Editor, + pid, + instance_id: instance_id.to_string(), + protocol_version: "0.1.0".to_string(), + connection_token: "editor-tok".to_string(), + endpoint: AgentHostEndpointAddress::Socket { + path: socket_path.to_string(), + }, + quality: None, + tunnel_name: None, + } + } + + #[test] + fn list_live_endpoints_includes_editor_and_standalone_sorted_stably() { + let dir = tempfile::tempdir().unwrap(); + let log = log::Logger::test(); + let pid = std::process::id(); + let ed = editor(pid, "editor-b", "/tmp/editor-b.sock"); + let sa = standalone(pid, "standalone-a", 8080); + // Publish editor first so we can assert sort order isn't just + // insertion order. + publish_agent_host_endpoint(&log, dir.path(), &ed).unwrap(); + publish_agent_host_endpoint(&log, dir.path(), &sa).unwrap(); + + let live = list_live_endpoints(&log, dir.path()); + + assert_eq!(live.len(), 2); + // Standalone sorts before editor per `server_type_sort_rank`. + assert_eq!(live[0].instance_id, "standalone-a"); + assert_eq!(live[1].instance_id, "editor-b"); + } + + #[test] + fn list_live_endpoints_excludes_dead_and_dedupes() { + let dir = tempfile::tempdir().unwrap(); + let log = log::Logger::test(); + let dead_pid = u32::MAX - 1; + let dead = standalone(dead_pid, "dead", 100); + let live_entry = standalone(std::process::id(), "live", 200); + publish_agent_host_endpoint(&log, dir.path(), &dead).unwrap(); + publish_agent_host_endpoint(&log, dir.path(), &live_entry).unwrap(); + + let live = list_live_endpoints(&log, dir.path()); + + assert_eq!(live, vec![live_entry]); + } + + #[test] + fn list_live_endpoints_returns_empty_when_registry_missing() { + let dir = tempfile::tempdir().unwrap(); + let log = log::Logger::test(); + + assert_eq!(list_live_endpoints(&log, dir.path()), Vec::new()); + } + + #[test] + fn list_live_standalone_endpoints_excludes_editor_but_includes_socket_standalones() { + let dir = tempfile::tempdir().unwrap(); + let log = log::Logger::test(); + let pid = std::process::id(); + let ed = editor(pid, "editor-a", "/tmp/editor-a.sock"); + let sa_tcp = standalone(pid, "standalone-tcp", 8080); + let mut sa_socket = standalone(pid, "standalone-socket", 0); + sa_socket.endpoint = AgentHostEndpointAddress::Socket { + path: "/tmp/standalone.sock".to_string(), + }; + publish_agent_host_endpoint(&log, dir.path(), &ed).unwrap(); + publish_agent_host_endpoint(&log, dir.path(), &sa_tcp).unwrap(); + publish_agent_host_endpoint(&log, dir.path(), &sa_socket).unwrap(); + + let standalones = list_live_standalone_endpoints(&log, dir.path()); + + assert_eq!(standalones.len(), 2); + assert!(standalones + .iter() + .all(|e| e.server_type == AgentHostServerType::Standalone)); + assert!(standalones + .iter() + .any(|e| e.instance_id == "standalone-tcp")); + assert!(standalones + .iter() + .any(|e| e.instance_id == "standalone-socket")); + } + + #[test] + fn address_label_formats_tcp_and_socket_endpoints() { + let tcp = standalone(1, "a", 8080); + assert_eq!(tcp.address_label(), "127.0.0.1:8080"); + + let socket = editor(1, "b", "/tmp/editor.sock"); + assert_eq!(socket.address_label(), "/tmp/editor.sock"); + } + + #[test] + fn label_includes_kind_pid_address_quality_and_tunnel() { + let mut entry = standalone(42, "a", 8080); + assert_eq!(entry.label(), "standalone (pid 42, 127.0.0.1:8080)"); + + entry.quality = Some("insider".to_string()); + entry.tunnel_name = Some("my-tunnel".to_string()); + assert_eq!( + entry.label(), + "standalone (pid 42, 127.0.0.1:8080) [insider] (tunnel my-tunnel)" + ); + + let ed = editor(7, "editor-a", "/tmp/editor.sock"); + assert_eq!(ed.label(), "editor (pid 7, /tmp/editor.sock)"); + } +} diff --git a/cli/src/tunnels/agent_host_registry_acl_windows.rs b/cli/src/tunnels/agent_host_registry_acl_windows.rs new file mode 100644 index 00000000000..641f39788ee --- /dev/null +++ b/cli/src/tunnels/agent_host_registry_acl_windows.rs @@ -0,0 +1,350 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +//! Windows-only owner-only ACL for the shared agent host endpoint registry +//! directory, via native Win32 security APIs (no `whoami.exe`/`icacls.exe` +//! subprocesses). The current user is identified by SID rather than by a +//! locale-sensitive account name. + +use std::ffi::OsStr; +use std::io; +use std::iter::once; +use std::os::windows::ffi::OsStrExt; +use std::path::Path; +use std::ptr; + +use windows_sys::core::PWSTR; +use windows_sys::Win32::Foundation::{CloseHandle, LocalFree, HANDLE, HLOCAL}; +use windows_sys::Win32::Security::Authorization::{ + ConvertSidToStringSidW, ConvertStringSecurityDescriptorToSecurityDescriptorW, + SetNamedSecurityInfoW, SDDL_REVISION_1, SE_FILE_OBJECT, +}; +use windows_sys::Win32::Security::{ + GetSecurityDescriptorDacl, GetTokenInformation, TokenUser, ACL, DACL_SECURITY_INFORMATION, + PROTECTED_DACL_SECURITY_INFORMATION, PSECURITY_DESCRIPTOR, PSID, TOKEN_QUERY, TOKEN_USER, +}; +use windows_sys::Win32::System::Threading::{GetCurrentProcess, OpenProcessToken}; + +/// Applies a protected DACL to `path` granting full control, with +/// object/container inheritance (so children inherit the same grants), to +/// only the current user, `LOCAL SYSTEM`, and `BUILTIN\Administrators`. +/// Mirrors `prepareLocalAgentHostEndpointMetadataDirectory`'s Windows branch +/// in `node/localAgentHostMetadata.ts`. +pub(super) fn apply_owner_only_acl(path: &Path) -> io::Result<()> { + let sid = current_user_sid_string()?; + // "D:P" = protected DACL, i.e. no inherited ACEs from the parent + // (matches `icacls /inheritance:r`). Each `(A;OICI;FA;;;)` grants + // Full-Access to , with Object-Inherit + Container-Inherit so + // files/subdirectories created underneath inherit the same grant. `SY` + // and `BA` are the SDDL well-known aliases for LOCAL SYSTEM and + // BUILTIN\Administrators. + let sddl = format!("D:P(A;OICI;FA;;;{sid})(A;OICI;FA;;;SY)(A;OICI;FA;;;BA)"); + + let descriptor = SecurityDescriptor::from_sddl(&sddl)?; + let dacl = descriptor.dacl()?; + let path_wide = to_wide_null(path.as_os_str()); + + // SAFETY: `path_wide` is a valid null-terminated wide string live for + // the call; `dacl` points into `descriptor`, which outlives this call. + // Null owner/group/sacl leave those unchanged. + let result = unsafe { + SetNamedSecurityInfoW( + path_wide.as_ptr(), + SE_FILE_OBJECT, + DACL_SECURITY_INFORMATION | PROTECTED_DACL_SECURITY_INFORMATION, + ptr::null_mut(), + ptr::null_mut(), + dacl, + ptr::null(), + ) + }; + if result != 0 { + return Err(io::Error::from_raw_os_error(result as i32)); + } + Ok(()) +} + +/// Returns the exact SID (never a locale-sensitive account name) of the +/// current process's user, in `S-1-5-...` string form. +fn current_user_sid_string() -> io::Result { + // SAFETY: `GetCurrentProcess` returns a pseudo-handle that must not be + // closed. `OpenProcessToken` with `TOKEN_QUERY` only opens a read + // handle to the process's own token. + let mut token: HANDLE = ptr::null_mut(); + if unsafe { OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &mut token) } == 0 { + return Err(io::Error::last_os_error()); + } + let token = OwnedHandle(token); + + let mut needed: u32 = 0; + // SAFETY: passing a null/zero-length buffer is the documented way to + // query the required size; it is written to `needed` regardless of + // the (expected) failure return. + unsafe { GetTokenInformation(token.0, TokenUser, ptr::null_mut(), 0, &mut needed) }; + if needed == 0 { + return Err(io::Error::last_os_error()); + } + + let mut buffer = vec![0u8; needed as usize]; + // SAFETY: `buffer` is exactly `needed` bytes, the size this same call + // reported above, and is only read as a `TOKEN_USER` once this + // succeeds. + let ok = unsafe { + GetTokenInformation( + token.0, + TokenUser, + buffer.as_mut_ptr() as *mut _, + needed, + &mut needed, + ) + }; + if ok == 0 { + return Err(io::Error::last_os_error()); + } + + // SAFETY: `buffer` was just filled with a valid `TOKEN_USER` above; + // `User.Sid` points inside `buffer` and stays valid for `buffer`'s + // lifetime, which outlives this call to `sid_to_string`. + let sid = unsafe { (*(buffer.as_ptr() as *const TOKEN_USER)).User.Sid }; + sid_to_string(sid) +} + +/// Converts a SID to its `S-1-5-...` string form. +fn sid_to_string(sid: PSID) -> io::Result { + let mut string_sid: PWSTR = ptr::null_mut(); + // SAFETY: `sid` is a valid SID for the duration of this call (see + // callers). `ConvertSidToStringSidW` allocates `string_sid` via + // `LocalAlloc`; ownership is transferred to `LocalWideString`, which + // frees it via `LocalFree` on drop. + if unsafe { ConvertSidToStringSidW(sid, &mut string_sid) } == 0 { + return Err(io::Error::last_os_error()); + } + Ok(LocalWideString(string_sid).to_string_lossy()) +} + +fn to_wide_null(value: &OsStr) -> Vec { + value.encode_wide().chain(once(0)).collect() +} + +/// An owned process token handle, closed on drop. +struct OwnedHandle(HANDLE); + +impl Drop for OwnedHandle { + fn drop(&mut self) { + // SAFETY: `self.0` is a valid handle opened by `OpenProcessToken` + // and not otherwise closed. + unsafe { + CloseHandle(self.0); + } + } +} + +/// An owned `LocalAlloc`-backed wide string, e.g. as returned by +/// `ConvertSidToStringSidW`, freed via `LocalFree` on drop. +struct LocalWideString(PWSTR); + +impl LocalWideString { + fn to_string_lossy(&self) -> String { + // SAFETY: `self.0` is a valid null-terminated wide string for the + // lifetime of `self`. + let len = unsafe { (0..).take_while(|&i| *self.0.add(i) != 0).count() }; + // SAFETY: `[self.0, self.0 + len)` was just measured as the + // null-terminated extent of a valid wide string above. + let slice = unsafe { std::slice::from_raw_parts(self.0, len) }; + String::from_utf16_lossy(slice) + } +} + +impl Drop for LocalWideString { + fn drop(&mut self) { + if !self.0.is_null() { + // SAFETY: `self.0` was allocated by `ConvertSidToStringSidW`, + // which documents `LocalFree` as the correct release for it. + unsafe { + LocalFree(self.0 as HLOCAL); + } + } + } +} + +/// An owned `LocalAlloc`-backed security descriptor, e.g. as returned by +/// `ConvertStringSecurityDescriptorToSecurityDescriptorW` or +/// `GetNamedSecurityInfoW`, freed via `LocalFree` on drop. +pub(super) struct SecurityDescriptor(PSECURITY_DESCRIPTOR); + +impl SecurityDescriptor { + fn from_sddl(sddl: &str) -> io::Result { + let wide = to_wide_null(OsStr::new(sddl)); + let mut descriptor: PSECURITY_DESCRIPTOR = ptr::null_mut(); + // SAFETY: `wide` is a valid null-terminated wide string live for + // the call. The resulting descriptor is allocated via + // `LocalAlloc` and is owned by the returned `Self`. + let ok = unsafe { + ConvertStringSecurityDescriptorToSecurityDescriptorW( + wide.as_ptr(), + SDDL_REVISION_1, + &mut descriptor, + ptr::null_mut(), + ) + }; + if ok == 0 { + return Err(io::Error::last_os_error()); + } + Ok(Self(descriptor)) + } + + /// Takes ownership of a security descriptor already allocated by a + /// Win32 API that documents `LocalFree` as its release (e.g. + /// `GetNamedSecurityInfoW`), used by tests to inspect an applied ACL. + #[cfg(test)] + pub(super) fn from_raw(descriptor: PSECURITY_DESCRIPTOR) -> Self { + Self(descriptor) + } + + /// Used by tests to independently verify security-descriptor control + /// flags (e.g. `SE_DACL_PROTECTED`) on the applied ACL. + #[cfg(test)] + pub(super) fn raw(&self) -> PSECURITY_DESCRIPTOR { + self.0 + } + + /// Returns the DACL embedded in this security descriptor. The + /// returned pointer's lifetime is tied to `self`. + pub(super) fn dacl(&self) -> io::Result<*const ACL> { + let mut present = 0; + let mut dacl: *mut ACL = ptr::null_mut(); + let mut defaulted = 0; + // SAFETY: `self.0` is a valid security descriptor kept alive for + // at least as long as the returned pointer is used. + let ok = + unsafe { GetSecurityDescriptorDacl(self.0, &mut present, &mut dacl, &mut defaulted) }; + if ok == 0 { + return Err(io::Error::last_os_error()); + } + if present == 0 { + return Err(io::Error::other( + "security descriptor unexpectedly has no DACL", + )); + } + Ok(dacl) + } +} + +impl Drop for SecurityDescriptor { + fn drop(&mut self) { + if !self.0.is_null() { + // SAFETY: `self.0` was allocated by either + // `ConvertStringSecurityDescriptorToSecurityDescriptorW` or + // `GetNamedSecurityInfoW`, both of which document `LocalFree` + // as the correct release for it. + unsafe { + LocalFree(self.0 as HLOCAL); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use windows_sys::Win32::Security::Authorization::GetNamedSecurityInfoW; + use windows_sys::Win32::Security::{ + GetAce, GetSecurityDescriptorControl, ACCESS_ALLOWED_ACE, CONTAINER_INHERIT_ACE, + OBJECT_INHERIT_ACE, SE_DACL_PROTECTED, + }; + + // winnt.h `ACCESS_ALLOWED_ACE_TYPE`; not worth an extra windows-sys + // feature for a single test-only constant. + const ACCESS_ALLOWED_ACE_TYPE: u8 = 0; + // winnt.h `FILE_ALL_ACCESS`; likewise kept local to this test. + const FILE_ALL_ACCESS: u32 = 0x001F_01FF; + + /// Reads back the DACL Win32 actually stored for `path` (as opposed to + /// the DACL passed to `SetNamedSecurityInfoW`), so the assertions below + /// exercise the real, applied ACL rather than our own construction of + /// it. + fn read_dacl(path: &Path) -> (SecurityDescriptor, *const ACL) { + let path_wide = to_wide_null(path.as_os_str()); + let mut dacl: *mut ACL = ptr::null_mut(); + let mut descriptor: PSECURITY_DESCRIPTOR = ptr::null_mut(); + // SAFETY: `path_wide` is a valid null-terminated wide string live + // for the call. The returned descriptor is `LocalAlloc`-backed and + // is wrapped below so it is freed on drop; `dacl` points inside it. + let result = unsafe { + GetNamedSecurityInfoW( + path_wide.as_ptr(), + SE_FILE_OBJECT, + DACL_SECURITY_INFORMATION | PROTECTED_DACL_SECURITY_INFORMATION, + ptr::null_mut(), + ptr::null_mut(), + &mut dacl, + ptr::null_mut(), + &mut descriptor, + ) + }; + assert_eq!(result, 0, "GetNamedSecurityInfoW failed: {result}"); + (SecurityDescriptor::from_raw(descriptor), dacl) + } + + #[test] + fn apply_owner_only_acl_grants_exactly_user_system_and_admins() { + let dir = tempfile::tempdir().unwrap(); + apply_owner_only_acl(dir.path()).unwrap(); + + let (descriptor, dacl) = read_dacl(dir.path()); + + let mut control: u16 = 0; + let mut revision: u32 = 0; + // SAFETY: `descriptor.raw()` is a valid security descriptor from + // `read_dacl` above, kept alive by `descriptor`. + let ok = + unsafe { GetSecurityDescriptorControl(descriptor.raw(), &mut control, &mut revision) }; + assert_ne!(ok, 0, "GetSecurityDescriptorControl failed"); + assert_ne!( + control & SE_DACL_PROTECTED, + 0, + "DACL must be protected (no inheritance from the parent directory)" + ); + + // SAFETY: `dacl` points into `descriptor`, which is still alive. + let ace_count = unsafe { (*dacl).AceCount }; + assert_eq!(ace_count, 3, "expected exactly one ACE per granted SID"); + + let mut granted_sids = Vec::new(); + for index in 0..ace_count as u32 { + let mut ace_ptr: *mut core::ffi::c_void = ptr::null_mut(); + // SAFETY: `index` is within `[0, ace_count)`; `dacl` is valid. + let ok = unsafe { GetAce(dacl, index, &mut ace_ptr) }; + assert_ne!(ok, 0, "GetAce failed for index {index}"); + + // SAFETY: every ACE in an SDDL-built `(A;...)` DACL is an + // `ACCESS_ALLOWED_ACE`, and `ace_ptr` was just validated above. + let ace = unsafe { &*(ace_ptr as *const ACCESS_ALLOWED_ACE) }; + assert_eq!(ace.Header.AceType, ACCESS_ALLOWED_ACE_TYPE); + let inherit_flags = OBJECT_INHERIT_ACE | CONTAINER_INHERIT_ACE; + assert_eq!( + ace.Header.AceFlags as u32 & inherit_flags, + inherit_flags, + "ACE must carry object + container inherit flags for index {index}" + ); + assert_eq!(ace.Mask, FILE_ALL_ACCESS, "ACE must grant full control"); + + // SAFETY: `SidStart` is the address of the variable-length SID + // data that immediately follows the fixed `ACCESS_ALLOWED_ACE` + // fields, per its documented layout. + let sid = &ace.SidStart as *const u32 as PSID; + granted_sids.push(sid_to_string(sid).unwrap()); + } + granted_sids.sort(); + + let mut expected = vec![ + current_user_sid_string().unwrap(), + "S-1-5-18".to_string(), // LOCAL SYSTEM + "S-1-5-32-544".to_string(), // BUILTIN\Administrators + ]; + expected.sort(); + assert_eq!(granted_sids, expected); + } +} diff --git a/cli/src/tunnels/control_server.rs b/cli/src/tunnels/control_server.rs index 41419274101..6cb0cb16d44 100644 --- a/cli/src/tunnels/control_server.rs +++ b/cli/src/tunnels/control_server.rs @@ -201,8 +201,9 @@ pub async fn serve( let (exit_barrier, signal_exit) = new_barrier(); // Kick off the agent host supervisor in the background. The supervisor - // is the only process that binds the user-facing TCP listener and owns - // the canonical lockfile; we never spawn an in-process sidecar here. + // is the only process that binds the user-facing TCP listener and + // publishes the canonical registry entry; we never spawn an + // in-process sidecar here. // We deliberately do NOT await this here — the tunnel needs to start // accepting connections immediately. Consumers that need the // supervisor's endpoint (currently `handle_serve` for the diff --git a/cli/src/tunnels/user_data_path.rs b/cli/src/tunnels/user_data_path.rs new file mode 100644 index 00000000000..7595b64837f --- /dev/null +++ b/cli/src/tunnels/user_data_path.rs @@ -0,0 +1,326 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +//! Resolves the platform "user data" directory used to home the shared +//! agent-host discovery registry (`/agent-host/local-endpoint/metadata.json`). +//! +//! Mirrors the precedence and per-platform rules implemented by the +//! TypeScript resolver in `src/vs/platform/environment/node/userDataPath.ts` +//! (which is passed `product.nameShort`), with one deliberate ordering +//! change: an explicit `--user-data-dir` always wins here. In the Electron +//! main process, `VSCODE_PORTABLE`/`VSCODE_APPDATA` are checked *before* the +//! CLI argument only to work around Electron implicitly re-injecting +//! `--user-data-dir` into argv; that quirk does not apply to this +//! standalone-CLI-only flag, so we can use the simpler, more predictable +//! "explicit flag always wins" order. + +use std::path::PathBuf; + +use crate::constants::PRODUCT_NAME_SHORT; + +/// The platform family used to select default user-data directory rules. +/// Kept distinct from `std::env::consts::OS` so unit tests can exercise all +/// three branches regardless of the host running the tests. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum UserDataOs { + Windows, + MacOs, + Linux, +} + +impl UserDataOs { + pub fn current() -> Self { + if cfg!(target_os = "windows") { + UserDataOs::Windows + } else if cfg!(target_os = "macos") { + UserDataOs::MacOs + } else { + UserDataOs::Linux + } + } +} + +/// Snapshot of the environment variables (and home directory) consulted by +/// the resolver. Injectable so tests never need to mutate real process-wide +/// environment variables (which would make tests order-dependent / racy). +#[derive(Debug, Clone, Default)] +pub struct UserDataPathEnv { + pub vscode_portable: Option, + pub vscode_appdata: Option, + pub appdata: Option, + pub userprofile: Option, + pub xdg_config_home: Option, + pub home_dir: Option, +} + +impl UserDataPathEnv { + /// Reads the real process environment / home directory. + pub fn from_process() -> Self { + Self { + vscode_portable: non_empty_env("VSCODE_PORTABLE"), + vscode_appdata: non_empty_env("VSCODE_APPDATA"), + appdata: non_empty_env("APPDATA"), + userprofile: non_empty_env("USERPROFILE"), + xdg_config_home: non_empty_env("XDG_CONFIG_HOME"), + home_dir: dirs::home_dir(), + } + } +} + +fn non_empty_env(name: &str) -> Option { + match std::env::var(name) { + Ok(v) if !v.is_empty() => Some(v), + _ => None, + } +} + +/// Resolves the user data directory using the real process environment and +/// the built-in [`PRODUCT_NAME_SHORT`]. +pub fn resolve_user_data_path(explicit: Option<&str>) -> PathBuf { + resolve_user_data_path_with( + explicit, + PRODUCT_NAME_SHORT, + UserDataOs::current(), + &UserDataPathEnv::from_process(), + ) +} + +/// Core, dependency-injected resolver. Precedence: +/// 1. Explicit `--user-data-dir`, if provided. +/// 2. `VSCODE_PORTABLE` -> `/user-data`. +/// 3. `VSCODE_APPDATA` -> `/`. +/// 4. The platform default (see [`default_user_data_path_with`]). +/// +/// Relative paths (explicit or portable) are resolved against the current +/// working directory, matching the TypeScript resolver's behavior of +/// resolving non-absolute paths against `process.cwd()`. +pub fn resolve_user_data_path_with( + explicit: Option<&str>, + product_name: &str, + os: UserDataOs, + env: &UserDataPathEnv, +) -> PathBuf { + if let Some(explicit) = explicit { + return resolve_relative_to_cwd(explicit, os); + } + + if let Some(portable) = &env.vscode_portable { + return join_component(resolve_relative_to_cwd(portable, os), "user-data", os); + } + + if let Some(appdata) = &env.vscode_appdata { + return join_component(resolve_relative_to_cwd(appdata, os), product_name, os); + } + + default_user_data_path_with(product_name, os, env) +} + +/// Resolves the platform default user data directory using the real process +/// environment and the built-in [`PRODUCT_NAME_SHORT`]. +pub fn default_user_data_path() -> PathBuf { + default_user_data_path_with( + PRODUCT_NAME_SHORT, + UserDataOs::current(), + &UserDataPathEnv::from_process(), + ) +} + +/// Core, dependency-injected platform-default resolver, mirroring +/// `getDefaultUserDataPath` in the TypeScript resolver: +/// - Windows: `%APPDATA%\`, falling back to +/// `%USERPROFILE%\AppData\Roaming\`. +/// - macOS: `~/Library/Application Support/`. +/// - Linux: `${XDG_CONFIG_HOME:-~/.config}/`. +pub fn default_user_data_path_with( + product_name: &str, + os: UserDataOs, + env: &UserDataPathEnv, +) -> PathBuf { + match os { + UserDataOs::Windows => { + let base = if let Some(appdata) = &env.appdata { + PathBuf::from(appdata) + } else if let Some(userprofile) = &env.userprofile { + PathBuf::from(userprofile).join("AppData").join("Roaming") + } else { + home_dir_or_empty(env) + }; + base.join(product_name) + } + UserDataOs::MacOs => home_dir_or_empty(env) + .join("Library") + .join("Application Support") + .join(product_name), + UserDataOs::Linux => { + let base = if let Some(xdg) = &env.xdg_config_home { + PathBuf::from(xdg) + } else { + home_dir_or_empty(env).join(".config") + }; + base.join(product_name) + } + } +} + +fn home_dir_or_empty(env: &UserDataPathEnv) -> PathBuf { + env.home_dir.clone().unwrap_or_else(|| PathBuf::from("")) +} + +/// Returns whether `path` is absolute under the *given* platform's rules, +/// independent of the host the code is actually compiled/running on. Plain +/// `Path::is_absolute` can't be used here because on Windows it returns +/// `false` for POSIX-style rooted paths like `/mnt/portable` (they're +/// "drive-relative", not absolute), which would otherwise let a supposedly +/// absolute override silently get re-rooted under the current directory. +/// This only matters for exercising all three `UserDataOs` branches from a +/// single test host; in production `os` always matches the real host. +fn is_absolute_for_os(path: &str, os: UserDataOs) -> bool { + match os { + UserDataOs::Windows => { + let bytes = path.as_bytes(); + let has_drive_root = bytes.len() >= 3 + && bytes[0].is_ascii_alphabetic() + && bytes[1] == b':' + && (bytes[2] == b'\\' || bytes[2] == b'/'); + has_drive_root || path.starts_with("\\\\") || path.starts_with("//") + } + UserDataOs::MacOs | UserDataOs::Linux => path.starts_with('/'), + } +} + +/// Joins `component` onto `base` using the path separator for `os`, without +/// going through `PathBuf::join`'s host-native (and, for a rooted-but-no-prefix +/// `base` on Windows, surprising) semantics. +fn join_component(base: PathBuf, component: &str, os: UserDataOs) -> PathBuf { + let sep = match os { + UserDataOs::Windows => '\\', + UserDataOs::MacOs | UserDataOs::Linux => '/', + }; + let mut s = base.to_string_lossy().into_owned(); + if !s.ends_with(['/', '\\']) { + s.push(sep); + } + s.push_str(component); + PathBuf::from(s) +} + +fn resolve_relative_to_cwd(path: &str, os: UserDataOs) -> PathBuf { + if is_absolute_for_os(path, os) { + return PathBuf::from(path); + } + + match std::env::current_dir() { + Ok(cwd) => cwd.join(path), + Err(_) => PathBuf::from(path), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn env_with_home(home: &str) -> UserDataPathEnv { + UserDataPathEnv { + home_dir: Some(PathBuf::from(home)), + ..Default::default() + } + } + + #[test] + fn default_windows_path_uses_appdata() { + let mut env = env_with_home(r"C:\Users\test"); + env.appdata = Some(r"C:\Users\test\AppData\Roaming".to_string()); + let path = default_user_data_path_with("Code - OSS", UserDataOs::Windows, &env); + assert_eq!( + path, + PathBuf::from(r"C:\Users\test\AppData\Roaming\Code - OSS") + ); + } + + #[test] + fn default_windows_path_falls_back_to_userprofile() { + let mut env = env_with_home(r"C:\Users\test"); + env.userprofile = Some(r"C:\Users\test".to_string()); + let path = default_user_data_path_with("Code - OSS", UserDataOs::Windows, &env); + assert_eq!( + path, + PathBuf::from(r"C:\Users\test\AppData\Roaming\Code - OSS") + ); + } + + #[test] + fn default_macos_path_uses_application_support() { + let env = env_with_home("/Users/test"); + let path = default_user_data_path_with("Code - OSS", UserDataOs::MacOs, &env); + assert_eq!( + path, + PathBuf::from("/Users/test/Library/Application Support/Code - OSS") + ); + } + + #[test] + fn default_linux_path_uses_xdg_config_home() { + let mut env = env_with_home("/home/test"); + env.xdg_config_home = Some("/home/test/.config".to_string()); + let path = default_user_data_path_with("Code - OSS", UserDataOs::Linux, &env); + assert_eq!(path, PathBuf::from("/home/test/.config/Code - OSS")); + } + + #[test] + fn default_linux_path_falls_back_to_home_dot_config() { + let env = env_with_home("/home/test"); + let path = default_user_data_path_with("Code - OSS", UserDataOs::Linux, &env); + assert_eq!(path, PathBuf::from("/home/test/.config/Code - OSS")); + } + + #[test] + fn vscode_portable_overrides_default() { + let mut env = env_with_home("/home/test"); + env.vscode_portable = Some("/mnt/portable".to_string()); + let path = resolve_user_data_path_with(None, "Code - OSS", UserDataOs::Linux, &env); + assert_eq!(path, PathBuf::from("/mnt/portable/user-data")); + } + + #[test] + fn vscode_appdata_overrides_default() { + let mut env = env_with_home("/home/test"); + env.vscode_appdata = Some("/mnt/appdata".to_string()); + let path = resolve_user_data_path_with(None, "Code - OSS", UserDataOs::Linux, &env); + assert_eq!(path, PathBuf::from("/mnt/appdata/Code - OSS")); + } + + #[test] + fn explicit_dir_takes_precedence_over_portable_and_appdata() { + let mut env = env_with_home("/home/test"); + env.vscode_portable = Some("/mnt/portable".to_string()); + env.vscode_appdata = Some("/mnt/appdata".to_string()); + let path = resolve_user_data_path_with( + Some("/explicit/dir"), + "Code - OSS", + UserDataOs::Linux, + &env, + ); + assert_eq!(path, PathBuf::from("/explicit/dir")); + } + + #[test] + fn relative_explicit_dir_resolves_against_cwd() { + let env = env_with_home("/home/test"); + let path = resolve_user_data_path_with( + Some("relative-dir"), + "Code - OSS", + UserDataOs::Linux, + &env, + ); + assert_eq!(path, std::env::current_dir().unwrap().join("relative-dir")); + } + + #[test] + fn falls_back_to_platform_default_when_nothing_set() { + let env = env_with_home("/home/test"); + let path = resolve_user_data_path_with(None, "Code - OSS", UserDataOs::Linux, &env); + assert_eq!(path, PathBuf::from("/home/test/.config/Code - OSS")); + } +} diff --git a/cli/src/util/errors.rs b/cli/src/util/errors.rs index 8439b7cff1a..155cbdf77ca 100644 --- a/cli/src/util/errors.rs +++ b/cli/src/util/errors.rs @@ -523,6 +523,16 @@ pub enum CodeError { ServerUnexpectedExit(String), #[error("Server binary is not executable: {0}")] ServerNotExecutable(String), + #[error("no agent host could be reached: {0}")] + NoAgentHostReachable(String), + #[error("no session matching \"{0}\" was found on any discovered agent host")] + SessionNotFoundOnAnyHost(String), + #[error("could not confirm whether session \"{0}\" exists: {1} could not be searched")] + IncompleteSessionSearch(String, String), + #[error("multiple live standalone agent hosts are registered ({0}); pass --instance-id to select one")] + AmbiguousAgentHostInstance(String), + #[error("no live standalone agent host with instance id \"{0}\" was found")] + UnknownAgentHostInstance(String), } makeAnyError!( diff --git a/cli/src/util/input.rs b/cli/src/util/input.rs index d0fec6832a9..c14157b2655 100644 --- a/cli/src/util/input.rs +++ b/cli/src/util/input.rs @@ -67,3 +67,17 @@ pub fn prompt_placeholder(question: &str, placeholder: &str) -> Result, options: &[String]) -> Result { + Select::with_theme(&ColorfulTheme::default()) + .with_prompt(text) + .items(options) + .default(0) + .interact() + .map_err(|e| wrap(e, "Failed to read select input")) +} diff --git a/src/vs/platform/agentHost/LOCAL_ENDPOINT.md b/src/vs/platform/agentHost/LOCAL_ENDPOINT.md index f35461a24f9..f936f67d38d 100644 --- a/src/vs/platform/agentHost/LOCAL_ENDPOINT.md +++ b/src/vs/platform/agentHost/LOCAL_ENDPOINT.md @@ -20,44 +20,75 @@ the product quality and any `--user-data-dir` argument. Implementations should resolve the active user data directory rather than assuming the default Stable or Insiders location. +The file is a **shared registry**: every locally running agent host process +(this editor's own utility process, other editor windows, and the standalone +`code agent host` CLI) upserts its own entry into the same file, so any one +process can discover every other live local agent host. + The file is optional. If VS Code cannot prepare or publish the external endpoint, it logs the error and continues running the agent host over its internal MessagePort transport. ## File format -The current schema version is `1`: +The current schema version is `2`: ```json [ { + "schemaVersion": 2, "type": "editor", - "schemaVersion": 1, "pid": 12345, "instanceId": "base64url-instance-id", - "endpointPath": "\\\\.\\pipe\\vscode-agent-host-...", + "protocolVersion": "0.7.0", "connectionToken": "base64url-bearer-token", - "protocolVersion": "0.7.0" + "endpoint": { + "type": "socket", + "path": "\\\\.\\pipe\\vscode-agent-host-..." + } } ] ``` | Property | Description | |---|---| -| `type` | Kind of local server. Currently always `editor`. | -| `schemaVersion` | Metadata schema version. Clients should reject unsupported versions. | -| `pid` | PID of the agent host utility process that owns the endpoint. | -| `instanceId` | Random identity used to distinguish successive endpoint owners. | -| `endpointPath` | Windows named pipe or Unix domain socket path. | -| `connectionToken` | Random bearer token required during the WebSocket upgrade. | +| `schemaVersion` | Metadata schema version. Clients must ignore entries whose version they do not understand rather than rejecting the whole file. | +| `type` | Kind of process that owns the endpoint: `editor` (a VS Code utility process) or `standalone` (the `code agent host` CLI). This controls ownership/default-selection policy on the client; it is not a measure of trust. | +| `pid` | PID of the process that owns the endpoint. | +| `instanceId` | Random identity used to distinguish successive endpoint owners. Combined with `type` and `pid`, this forms the entry's identity for dedupe/upsert/removal, since PIDs can be reused after a process exits. | | `protocolVersion` | AHP version spoken by the host. Clients must still perform the normal AHP `initialize` negotiation. | +| `connectionToken` | Random bearer token required during the WebSocket upgrade. | +| `endpoint` | Discriminated union describing how to connect: `{ "type": "socket", "path": string }` for a Windows named pipe or Unix domain socket (used by the editor today), or `{ "type": "tcp", "host": string, "port": number }` for a TCP listener (used by the standalone CLI). | +| `quality` | Optional. Product quality of a standalone CLI endpoint. Not part of entry identity. | +| `tunnelName` | Optional. Tunnel name associated with a standalone CLI endpoint. Not part of entry identity. | -Readers should treat every entry and field as untrusted input. +Readers must treat every entry and field as untrusted input: + +- structurally validate every entry, dropping malformed ones individually + rather than failing the whole read; +- ignore entries with an unsupported `schemaVersion`; +- ignore entries whose PID is confirmed dead, when a PID liveness check is + possible; +- deduplicate entries by `(type, pid, instanceId)`; +- never reconstruct `endpoint` values — always use the published address as-is; +- perform the normal AHP `initialize` negotiation after connecting, regardless + of the advertised `protocolVersion`. + +The shared parser/model lives in +[`common/agentHostEndpointRegistry.ts`](common/agentHostEndpointRegistry.ts) so +every local reader and writer (the editor publisher, a future registry +watcher, SSH discovery, and tests) validates entries identically. + +Schema version `1` (a flat `{ endpointPath: string }` shape, `type` always +`"editor"`) is the format previously written by the editor alone. Version-`2` +readers ignore version-`1` entries rather than attempting to interpret them, +consistent with the "ignore unsupported schema versions" rule above. ## Connecting -Connect to `endpointPath` using WebSocket framing and provide -`connectionToken` in the standard VS Code connection-token query parameter: +Connect to `endpoint.path` (socket endpoints) using WebSocket framing and +provide `connectionToken` in the standard VS Code connection-token query +parameter: ```text ?tkn= @@ -72,7 +103,7 @@ Connections without the token, or with the wrong token, are rejected with HTTP ## Endpoint paths -On Windows, `endpointPath` is a named pipe: +On Windows, the editor's `endpoint.path` is a named pipe: ```text \\.\pipe\vscode-agent-host-- @@ -87,6 +118,34 @@ user-data-specific directory to stay within Unix socket path-length limits: Clients must use the path from the metadata file rather than reconstructing it. +## Multi-writer safety + +Because more than one process can publish to `metadata.json` at once, an +atomic rename by itself is not sufficient: two writers could read the same +array concurrently and each overwrite the other's addition. Every writer +therefore: + +1. Acquires an exclusive lock: an atomically-created sibling lock directory + (`metadata.json.lock`) containing an `owner.json` file recording the + lock holder's `(pid, instanceId)`. +2. Reads the current array. +3. Drops entries whose PID is confirmed dead. +4. Upserts its own `(type, pid, instanceId)` entry. +5. Writes a mode-`0600` temporary file and atomically renames it over + `metadata.json`. +6. Releases the lock. + +Lock acquisition is bounded: if the lock directory already exists, a +contender inspects `owner.json`. If the recorded PID is no longer alive, the +lock is stale and is reclaimed immediately; otherwise acquisition is retried +until a short timeout elapses. On timeout, the writer does **not** silently +bypass the lock and write anyway — it logs the failure and continues running +undiscoverable, exactly as when the endpoint cannot be published at all. + +Readers never take the lock: because the registry file is only ever observed +in a fully-written state (via atomic rename), reads are always safe without +coordination. + ## Security and lifecycle - The metadata directory and file are restricted to the current user. On @@ -95,14 +154,17 @@ Clients must use the path from the metadata file rather than reconstructing it. metadata token is required to complete the WebSocket upgrade. - Metadata is written atomically only after the endpoint is listening and the protocol handler is installed. -- On shutdown, VS Code removes the metadata only if its PID and `instanceId` - still match. This prevents an older process from deleting a newer process's - endpoint record. +- On shutdown, VS Code reacquires the write lock and removes only the entry + whose `(type, pid, instanceId)` exactly matches its own. This prevents an + older process from deleting a newer process's endpoint record, and prevents + a writer from ever deleting another live writer's entry. The file itself is + deleted only when the resulting array is empty. - Clients should handle a missing file, a stale PID, endpoint closure, and the metadata being replaced while reconnecting. The implementation and lifecycle wiring live in: +- [`common/agentHostEndpointRegistry.ts`](common/agentHostEndpointRegistry.ts) - [`node/localAgentHostMetadata.ts`](node/localAgentHostMetadata.ts) - [`node/agentHostMain.ts`](node/agentHostMain.ts) - [`node/webSocketTransport.ts`](node/webSocketTransport.ts) diff --git a/src/vs/platform/agentHost/common/agentHostEndpointRegistry.ts b/src/vs/platform/agentHost/common/agentHostEndpointRegistry.ts new file mode 100644 index 00000000000..017e0054012 --- /dev/null +++ b/src/vs/platform/agentHost/common/agentHostEndpointRegistry.ts @@ -0,0 +1,182 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { vEnum, vLiteral, vNumber, vObj, vOptionalProp, vString, vUnion } from '../../../base/common/validation.js'; + +/** + * Schema version for the shared local agent-host endpoint registry at + * `/agent-host/local-endpoint/metadata.json`. This schema is + * shared with the Rust CLI (`cli/src/tunnels/agent_host_metadata.rs`); field + * renames or removals MUST be coordinated across both languages. + * + * Version 1 was the editor-only, socket-path-only schema + * (`ILocalAgentHostEndpointMetadata` in `node/localAgentHostMetadata.ts`). + * Version 2 generalizes the registry to hold both editor (socket/pipe) and + * standalone CLI (TCP) endpoints in the same file so every locally running + * agent host is discoverable from one place. + */ +export const AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION = 2; + +/** + * Kind of process that owns an agent host endpoint. Controls + * ownership/default-selection policy on the client; it is not by itself a + * measure of trust or of registry identity (see {@link IAgentHostEndpointIdentity}). + */ +export type AgentHostServerType = 'editor' | 'standalone'; + +/** + * How to physically connect to an endpoint. Editor endpoints are always a + * Unix domain socket or Windows named pipe; the standalone Rust CLI + * currently only ever publishes a TCP endpoint. + */ +export type AgentHostEndpointAddress = + | { readonly type: 'socket'; readonly path: string } + | { readonly type: 'tcp'; readonly host: string; readonly port: number }; + +/** + * One entry of the shared local agent-host endpoint registry. The registry + * file itself is a JSON array of these entries + * (`AgentHostEndpointMetadata[]`). + */ +export interface IAgentHostEndpointMetadata { + readonly schemaVersion: typeof AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION; + readonly type: AgentHostServerType; + readonly pid: number; + readonly instanceId: string; + readonly protocolVersion: string; + readonly connectionToken: string; + readonly endpoint: AgentHostEndpointAddress; + readonly quality?: string; + readonly tunnelName?: string; +} + +/** + * The subset of {@link IAgentHostEndpointMetadata} that identifies a unique + * registry entry/owner: `(type, pid, instanceId)`. `instanceId` makes + * identity safe across PID reuse and rapid process replacement; `pid` alone + * is not a safe key because operating systems recycle PIDs. + */ +export interface IAgentHostEndpointIdentity { + readonly type: AgentHostServerType; + readonly pid: number; + readonly instanceId: string; +} + +const endpointAddressValidator = vUnion( + vObj({ type: vLiteral('socket'), path: vString() }), + vObj({ type: vLiteral('tcp'), host: vString(), port: vNumber() }), +); + +const entryValidator = vObj({ + schemaVersion: vNumber(), + type: vEnum('editor', 'standalone'), + pid: vNumber(), + instanceId: vString(), + protocolVersion: vString(), + connectionToken: vString(), + endpoint: endpointAddressValidator, + quality: vOptionalProp(vString()), + tunnelName: vOptionalProp(vString()), +}); + +/** + * Structurally validates one raw registry entry and returns it typed, or + * `undefined` if it is malformed or its `schemaVersion` is not the one this + * build understands. Every field is treated as untrusted input. + */ +export function parseAgentHostEndpointMetadataEntry(raw: unknown): IAgentHostEndpointMetadata | undefined { + const { content, error } = entryValidator.validate(raw); + if (error) { + return undefined; + } + if (content.schemaVersion !== AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION) { + // Unsupported (older/newer) schema version: ignore rather than reject + // the whole registry, so a foreign writer's entry cannot hide every + // other live writer's endpoint. + return undefined; + } + if (!Number.isSafeInteger(content.pid) || content.pid <= 0) { + return undefined; + } + if (content.endpoint.type === 'tcp' && (!Number.isSafeInteger(content.endpoint.port) || content.endpoint.port <= 0 || content.endpoint.port > 65535)) { + return undefined; + } + + return { + schemaVersion: AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION, + type: content.type, + pid: content.pid, + instanceId: content.instanceId, + protocolVersion: content.protocolVersion, + connectionToken: content.connectionToken, + endpoint: content.endpoint, + quality: content.quality, + tunnelName: content.tunnelName, + }; +} + +/** + * Parses the raw contents of the registry file (expected to be + * `AgentHostEndpointMetadata[]`). Every entry is validated independently: + * malformed entries and entries with an unsupported `schemaVersion` are + * dropped rather than failing the entire read. A non-array top-level value + * yields an empty registry. + */ +export function parseAgentHostEndpointRegistry(raw: unknown): IAgentHostEndpointMetadata[] { + if (!Array.isArray(raw)) { + return []; + } + const entries: IAgentHostEndpointMetadata[] = []; + for (const item of raw) { + const entry = parseAgentHostEndpointMetadataEntry(item); + if (entry) { + entries.push(entry); + } + } + return entries; +} + +/** Stable string key for `(type, pid, instanceId)`, suitable for use as a Map key. */ +export function getAgentHostEndpointIdentityKey(identity: IAgentHostEndpointIdentity): string { + return `${identity.type}:${identity.pid}:${identity.instanceId}`; +} + +export function isSameAgentHostEndpointIdentity(a: IAgentHostEndpointIdentity, b: IAgentHostEndpointIdentity): boolean { + return a.type === b.type && a.pid === b.pid && a.instanceId === b.instanceId; +} + +/** + * Deduplicates entries by `(type, pid, instanceId)`. If a duplicate identity + * appears more than once (for example a crashed writer left a stale copy + * before another writer's cleanup ran), the entry encountered later in + * `entries` wins, since it is presumed to be the most recently written copy. + */ +export function dedupeAgentHostEndpointMetadata(entries: readonly IAgentHostEndpointMetadata[]): IAgentHostEndpointMetadata[] { + const byIdentity = new Map(); + for (const entry of entries) { + byIdentity.set(getAgentHostEndpointIdentityKey(entry), entry); + } + return [...byIdentity.values()]; +} + +/** + * Returns `entries` with any existing entry sharing `metadata`'s identity + * replaced by `metadata`. Used by a writer to upsert its own registry entry + * without disturbing other writers' entries. + */ +export function upsertAgentHostEndpointMetadata(entries: readonly IAgentHostEndpointMetadata[], metadata: IAgentHostEndpointMetadata): IAgentHostEndpointMetadata[] { + const remaining = entries.filter(entry => !isSameAgentHostEndpointIdentity(entry, metadata)); + remaining.push(metadata); + return remaining; +} + +/** + * Returns `entries` with the exact-identity-matching entry removed, if any. + * Used on shutdown so a writer only ever removes its own entry, never a + * newer process's entry that happens to share its PID. + */ +export function removeAgentHostEndpointMetadata(entries: readonly IAgentHostEndpointMetadata[], owner: IAgentHostEndpointIdentity): IAgentHostEndpointMetadata[] { + return entries.filter(entry => !isSameAgentHostEndpointIdentity(entry, owner)); +} diff --git a/src/vs/platform/agentHost/node/agentHostMain.ts b/src/vs/platform/agentHost/node/agentHostMain.ts index 07df867e2f3..36b6fcc4791 100644 --- a/src/vs/platform/agentHost/node/agentHostMain.ts +++ b/src/vs/platform/agentHost/node/agentHostMain.ts @@ -401,7 +401,7 @@ async function startAgentHost(): Promise { server.registerChannel(AgentHostIpcChannels.Protocol, messagePortProtocolServer); if (localEndpoint) { try { - await publishLocalAgentHostEndpointMetadata(environmentService.userDataPath, localEndpoint.metadata); + await publishLocalAgentHostEndpointMetadata(environmentService.userDataPath, localEndpoint.metadata, logService); localDataPlaneDisposables.add(toDisposable(() => { cleanupLocalAgentHostEndpoint(environmentService.userDataPath, localEndpoint.metadata, logService); })); @@ -557,7 +557,7 @@ async function startLocalAgentHostEndpoint( } server = await WebSocketProtocolServer.create( { - socketPath: endpointMetadata.endpointPath, + socketPath: endpointMetadata.endpoint.path, connectionTokenValidate: token => token === endpointMetadata.connectionToken, }, logService, @@ -585,12 +585,12 @@ function cleanupLocalAgentHostEndpoint( logService: ILogService, ): void { try { - cleanupLocalAgentHostEndpointMetadataSync(userDataPath, metadata); + cleanupLocalAgentHostEndpointMetadataSync(userDataPath, metadata, logService); } catch (error) { logService.error('[AgentHost] Failed to clean up local protocol metadata', error); } try { - cleanupLocalAgentHostEndpointSocketSync(metadata.endpointPath); + cleanupLocalAgentHostEndpointSocketSync(metadata.endpoint.path); } catch (error) { logService.error('[AgentHost] Failed to clean up local protocol socket', error); } diff --git a/src/vs/platform/agentHost/node/localAgentHostMetadata.ts b/src/vs/platform/agentHost/node/localAgentHostMetadata.ts index aed1b9b14d3..062060e603f 100644 --- a/src/vs/platform/agentHost/node/localAgentHostMetadata.ts +++ b/src/vs/platform/agentHost/node/localAgentHostMetadata.ts @@ -8,42 +8,37 @@ import { createHash, randomBytes } from 'crypto'; import * as fs from 'fs'; import * as os from 'os'; import { join } from '../../../base/common/path.js'; -import { vArray, vLiteral, vNumber, vObj, vString } from '../../../base/common/validation.js'; +import { ILogService } from '../../log/common/log.js'; +import { + AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION, + IAgentHostEndpointIdentity, + IAgentHostEndpointMetadata, + dedupeAgentHostEndpointMetadata, + parseAgentHostEndpointRegistry, + removeAgentHostEndpointMetadata, + upsertAgentHostEndpointMetadata, +} from '../common/agentHostEndpointRegistry.js'; import { PROTOCOL_VERSION } from '../common/state/protocol/version/registry.js'; +import { isPidAlive } from './agentHostLockfile.js'; -const metadataSchemaVersion = 1; const metadataDirectoryName = 'agent-host'; const endpointDirectoryName = 'local-endpoint'; const metadataFileName = 'metadata.json'; -export interface ILocalAgentHostEndpointMetadata { +/** The editor's own entry in the shared local agent host endpoint registry. */ +export type ILocalAgentHostEndpointMetadata = IAgentHostEndpointMetadata & { readonly type: 'editor'; - readonly schemaVersion: typeof metadataSchemaVersion; - readonly pid: number; - readonly instanceId: string; - readonly endpointPath: string; - readonly connectionToken: string; - readonly protocolVersion: string; -} - -const metadataValidator = vArray(vObj({ - type: vLiteral('editor'), - schemaVersion: vNumber(), - pid: vNumber(), - instanceId: vString(), - endpointPath: vString(), - connectionToken: vString(), - protocolVersion: vString(), -})); + readonly endpoint: { readonly type: 'socket'; readonly path: string }; +}; export function createLocalAgentHostEndpointMetadata(userDataPath: string): ILocalAgentHostEndpointMetadata { const instanceId = randomBytes(16).toString('base64url'); return { + schemaVersion: AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION, type: 'editor', - schemaVersion: metadataSchemaVersion, pid: process.pid, instanceId, - endpointPath: getEndpointPath(userDataPath, instanceId), + endpoint: { type: 'socket', path: getEndpointPath(userDataPath, instanceId) }, connectionToken: randomBytes(32).toString('base64url'), protocolVersion: PROTOCOL_VERSION, }; @@ -82,37 +77,64 @@ export async function prepareLocalAgentHostEndpointSocketDirectory(userDataPath: } } -export async function publishLocalAgentHostEndpointMetadata(userDataPath: string, metadata: ILocalAgentHostEndpointMetadata): Promise { +/** + * Upserts `metadata` into the shared local agent host endpoint registry. + * + * Multiple processes (this editor, other editor windows, and the standalone + * `code agent host` CLI) can publish to the same registry file concurrently, + * so an atomic rename alone is not sufficient: two writers could otherwise + * read the same array and overwrite each other's addition. This acquires a + * sibling exclusive lock first, so the read-prune-upsert-write sequence + * below is serialized across all writers. Readers remain lock-free because + * the final write is an atomic rename. + * + * Throws if the lock cannot be acquired within a bounded timeout, or if any + * filesystem operation fails; callers must treat that as "continue running, + * but undiscoverable" and must not fall back to a non-atomic write. + */ +export async function publishLocalAgentHostEndpointMetadata(userDataPath: string, metadata: ILocalAgentHostEndpointMetadata, logService?: ILogService): Promise { const metadataPath = getMetadataPath(userDataPath); - const temporaryPath = `${metadataPath}.${metadata.instanceId}.tmp`; - const entries = readMetadata(metadataPath).filter(entry => entry.pid !== metadata.pid || entry.type !== metadata.type); - entries.push(metadata); - const handle = await fs.promises.open(temporaryPath, 'wx', 0o600); - try { - await handle.writeFile(JSON.stringify(entries), 'utf8'); - await handle.sync(); - } finally { - await handle.close(); + const release = await acquireRegistryLockAsync(userDataPath, metadata, logService); + if (!release) { + throw new Error(`Timed out acquiring the local agent host endpoint registry lock at ${getLockDirectoryPath(userDataPath)}`); } - try { - await fs.promises.rename(temporaryPath, metadataPath); + const current = await readRegistryAsync(metadataPath); + const live = pruneDeadAgentHostEndpointMetadata(current, logService); + const next = upsertAgentHostEndpointMetadata(dedupeAgentHostEndpointMetadata(live), metadata); + await writeRegistryAtomicAsync(metadataPath, metadata.instanceId, next); } finally { - await fs.promises.rm(temporaryPath, { force: true }); + await release(); } } -export function cleanupLocalAgentHostEndpointMetadataSync(userDataPath: string, owner: ILocalAgentHostEndpointMetadata): void { +/** + * Removes exactly `owner`'s `(type, pid, instanceId)` entry from the + * registry, reacquiring the write lock first. Deletes the file entirely + * only when the resulting registry is empty. This is a best-effort shutdown + * operation: failures are logged, never thrown, so process exit is never + * blocked by cleanup. + */ +export function cleanupLocalAgentHostEndpointMetadataSync(userDataPath: string, owner: ILocalAgentHostEndpointMetadata, logService?: ILogService): void { const metadataPath = getMetadataPath(userDataPath); - const entries = readMetadata(metadataPath); - const remaining = entries.filter(entry => entry.pid !== owner.pid || entry.instanceId !== owner.instanceId || entry.type !== owner.type); - if (remaining.length === entries.length) { + const release = acquireRegistryLockSync(userDataPath, owner, logService); + if (!release) { + logService?.error(`[AgentHost] Timed out acquiring the local agent host endpoint registry lock while removing our entry from ${metadataPath}`); return; } - if (remaining.length === 0) { - fs.rmSync(metadataPath, { force: true }); - } else { - fs.writeFileSync(metadataPath, JSON.stringify(remaining), { encoding: 'utf8', mode: 0o600 }); + try { + const current = readRegistrySync(metadataPath); + const remaining = removeAgentHostEndpointMetadata(current, owner); + if (remaining.length === current.length) { + return; + } + if (remaining.length === 0) { + fs.rmSync(metadataPath, { force: true }); + } else { + writeRegistryAtomicSync(metadataPath, owner.instanceId, remaining); + } + } finally { + release(); } } @@ -122,6 +144,16 @@ export function cleanupLocalAgentHostEndpointSocketSync(endpointPath: string): v } } +/** + * Reads and validates every live entry in the shared local agent host + * endpoint registry, without taking the write lock. Safe to call frequently + * (e.g. from a file watcher) because the registry file is only ever + * observed in a fully-written state via atomic rename. + */ +export async function readLocalAgentHostEndpointRegistry(userDataPath: string): Promise { + return readRegistryAsync(getMetadataPath(userDataPath)); +} + function getMetadataDirectory(userDataPath: string): string { return join(userDataPath, metadataDirectoryName, endpointDirectoryName); } @@ -130,6 +162,14 @@ function getMetadataPath(userDataPath: string): string { return join(getMetadataDirectory(userDataPath), metadataFileName); } +function getLockDirectoryPath(userDataPath: string): string { + return `${getMetadataPath(userDataPath)}.lock`; +} + +function getLockOwnerFilePath(lockDirectoryPath: string): string { + return join(lockDirectoryPath, 'owner.json'); +} + function getSocketDirectory(userDataPath: string): string { const owner = process.getuid?.().toString() ?? ''; const hash = createHash('sha256').update(`${owner}:${userDataPath}`).digest('hex').slice(0, 12); @@ -144,25 +184,299 @@ function getEndpointPath(userDataPath: string, instanceId: string): string { return join(getSocketDirectory(userDataPath), `${instanceId}.sock`); } -function readMetadata(path: string): ILocalAgentHostEndpointMetadata[] { +async function readRegistryAsync(metadataPath: string): Promise { + let raw: string; try { - const stat = fs.lstatSync(path); + const stat = await fs.promises.lstat(metadataPath); if (!stat.isFile() || stat.isSymbolicLink()) { return []; } - const result = metadataValidator.validate(JSON.parse(fs.readFileSync(path, 'utf8'))); - if (result.error) { - return []; - } - return result.content - .filter(entry => entry.schemaVersion === metadataSchemaVersion) - .map(entry => ({ ...entry, schemaVersion: metadataSchemaVersion })); + raw = await fs.promises.readFile(metadataPath, 'utf8'); } catch (error) { - if (isNotFound(error) || error instanceof SyntaxError) { + if (isNotFound(error)) { return []; } throw error; } + return parseRegistryJson(raw); +} + +function readRegistrySync(metadataPath: string): IAgentHostEndpointMetadata[] { + let raw: string; + try { + const stat = fs.lstatSync(metadataPath); + if (!stat.isFile() || stat.isSymbolicLink()) { + return []; + } + raw = fs.readFileSync(metadataPath, 'utf8'); + } catch (error) { + if (isNotFound(error)) { + return []; + } + throw error; + } + return parseRegistryJson(raw); +} + +function parseRegistryJson(raw: string): IAgentHostEndpointMetadata[] { + try { + return parseAgentHostEndpointRegistry(JSON.parse(raw)); + } catch (error) { + if (error instanceof SyntaxError) { + return []; + } + throw error; + } +} + +/** + * Drops entries whose PID is confirmed dead. Entries are only ever pruned + * here (i.e. when death is certain via a PID liveness check); a live PID, or + * a PID we cannot check, is always kept. + */ +function pruneDeadAgentHostEndpointMetadata(entries: readonly IAgentHostEndpointMetadata[], logService?: ILogService): IAgentHostEndpointMetadata[] { + return entries.filter(entry => { + if (isPidAlive(entry.pid)) { + return true; + } + logService?.info(`[AgentHost] Pruning stale local endpoint registry entry: ${entry.type} PID ${entry.pid} (instance ${entry.instanceId}) is no longer running`); + return false; + }); +} + +async function writeRegistryAtomicAsync(metadataPath: string, uniqueSuffix: string, entries: readonly IAgentHostEndpointMetadata[]): Promise { + const temporaryPath = `${metadataPath}.${uniqueSuffix}.tmp`; + const handle = await fs.promises.open(temporaryPath, 'wx', 0o600); + try { + await handle.writeFile(JSON.stringify(entries), 'utf8'); + await handle.sync(); + } finally { + await handle.close(); + } + try { + await fs.promises.rename(temporaryPath, metadataPath); + } finally { + await fs.promises.rm(temporaryPath, { force: true }); + } +} + +function writeRegistryAtomicSync(metadataPath: string, uniqueSuffix: string, entries: readonly IAgentHostEndpointMetadata[]): void { + const temporaryPath = `${metadataPath}.${uniqueSuffix}.tmp`; + const fd = fs.openSync(temporaryPath, 'wx', 0o600); + try { + fs.writeFileSync(fd, JSON.stringify(entries), 'utf8'); + fs.fsyncSync(fd); + } finally { + fs.closeSync(fd); + } + try { + fs.renameSync(temporaryPath, metadataPath); + } finally { + fs.rmSync(temporaryPath, { force: true }); + } +} + +// #region Multi-writer lock +// +// The lock is a sibling directory to metadata.json (metadata.json.lock). +// `mkdir` without `recursive` is used as the exclusive-acquire primitive +// because directory creation is atomic on every platform we support and +// requires no native/optional dependency. The lock holder's `(pid, +// instanceId)` is written into an owner file inside the directory so a +// contending process can recognize and reclaim an abandoned lock: if the +// recorded PID is confirmed dead, the lock is stale and is reclaimed +// immediately; otherwise acquisition is retried until a bounded timeout +// elapses, after which the caller is told to log and continue +// undiscoverable rather than silently bypassing the lock. + +interface ILockOwner { + readonly pid: number; + readonly instanceId: string; +} + +const asyncLockAcquireTimeoutMs = 3000; +const asyncLockRetryDelayMs = 40; +const syncLockAcquireTimeoutMs = 500; +const syncLockRetryDelayMs = 10; +/** Grace period for a lock directory whose owner file has not appeared yet, to avoid racing a concurrent acquirer that is mid-write. */ +const lockOwnerGraceMs = 2000; + +async function acquireRegistryLockAsync(userDataPath: string, owner: ILockOwner, logService?: ILogService): Promise<(() => Promise) | undefined> { + const lockDirectoryPath = getLockDirectoryPath(userDataPath); + const deadline = Date.now() + asyncLockAcquireTimeoutMs; + for (; ;) { + try { + await fs.promises.mkdir(lockDirectoryPath); + await fs.promises.writeFile(getLockOwnerFilePath(lockDirectoryPath), JSON.stringify(owner), { encoding: 'utf8', mode: 0o600 }); + return () => releaseRegistryLockAsync(lockDirectoryPath, owner, logService); + } catch (error) { + if (!isAlreadyExists(error)) { + throw error; + } + if (await tryReclaimStaleLockAsync(lockDirectoryPath, logService)) { + continue; + } + if (Date.now() >= deadline) { + return undefined; + } + await delay(asyncLockRetryDelayMs); + } + } +} + +async function releaseRegistryLockAsync(lockDirectoryPath: string, owner: ILockOwner, logService?: ILogService): Promise { + try { + const current = await readLockOwnerAsync(lockDirectoryPath); + if (current && !isSameLockOwner(current, owner)) { + // Another process already reclaimed this lock as stale; it now owns + // this lock's lifecycle, so leave it alone. + return; + } + await fs.promises.rm(lockDirectoryPath, { recursive: true, force: true }); + } catch (error) { + logService?.error('[AgentHost] Failed to release the local agent host endpoint registry lock', error); + } +} + +async function tryReclaimStaleLockAsync(lockDirectoryPath: string, logService?: ILogService): Promise { + const owner = await readLockOwnerAsync(lockDirectoryPath); + if (owner) { + if (isPidAlive(owner.pid)) { + return false; + } + } else if (!(await isLockDirectoryStaleWithoutOwnerAsync(lockDirectoryPath))) { + return false; + } + try { + await fs.promises.rm(lockDirectoryPath, { recursive: true, force: true }); + } catch { + return false; + } + logService?.warn(`[AgentHost] Reclaimed a stale local agent host endpoint registry lock${owner ? ` from PID ${owner.pid}` : ''}`); + return true; +} + +async function readLockOwnerAsync(lockDirectoryPath: string): Promise { + try { + return parseLockOwner(JSON.parse(await fs.promises.readFile(getLockOwnerFilePath(lockDirectoryPath), 'utf8'))); + } catch { + return undefined; + } +} + +async function isLockDirectoryStaleWithoutOwnerAsync(lockDirectoryPath: string): Promise { + try { + const stat = await fs.promises.stat(lockDirectoryPath); + return Date.now() - stat.mtimeMs > lockOwnerGraceMs; + } catch { + // The directory disappeared already (another process reclaimed it); + // let the caller retry acquisition. + return true; + } +} + +function acquireRegistryLockSync(userDataPath: string, owner: ILockOwner, logService?: ILogService): (() => void) | undefined { + const lockDirectoryPath = getLockDirectoryPath(userDataPath); + const deadline = Date.now() + syncLockAcquireTimeoutMs; + for (; ;) { + try { + fs.mkdirSync(lockDirectoryPath); + fs.writeFileSync(getLockOwnerFilePath(lockDirectoryPath), JSON.stringify(owner), { encoding: 'utf8', mode: 0o600 }); + return () => releaseRegistryLockSync(lockDirectoryPath, owner, logService); + } catch (error) { + if (!isAlreadyExists(error)) { + throw error; + } + if (tryReclaimStaleLockSync(lockDirectoryPath, logService)) { + continue; + } + if (Date.now() >= deadline) { + return undefined; + } + sleepSync(syncLockRetryDelayMs); + } + } +} + +function releaseRegistryLockSync(lockDirectoryPath: string, owner: ILockOwner, logService?: ILogService): void { + try { + const current = readLockOwnerSync(lockDirectoryPath); + if (current && !isSameLockOwner(current, owner)) { + return; + } + fs.rmSync(lockDirectoryPath, { recursive: true, force: true }); + } catch (error) { + logService?.error('[AgentHost] Failed to release the local agent host endpoint registry lock', error); + } +} + +function tryReclaimStaleLockSync(lockDirectoryPath: string, logService?: ILogService): boolean { + const owner = readLockOwnerSync(lockDirectoryPath); + if (owner) { + if (isPidAlive(owner.pid)) { + return false; + } + } else if (!isLockDirectoryStaleWithoutOwnerSync(lockDirectoryPath)) { + return false; + } + try { + fs.rmSync(lockDirectoryPath, { recursive: true, force: true }); + } catch { + return false; + } + logService?.warn(`[AgentHost] Reclaimed a stale local agent host endpoint registry lock${owner ? ` from PID ${owner.pid}` : ''}`); + return true; +} + +function readLockOwnerSync(lockDirectoryPath: string): ILockOwner | undefined { + try { + return parseLockOwner(JSON.parse(fs.readFileSync(getLockOwnerFilePath(lockDirectoryPath), 'utf8'))); + } catch { + return undefined; + } +} + +function isLockDirectoryStaleWithoutOwnerSync(lockDirectoryPath: string): boolean { + try { + const stat = fs.statSync(lockDirectoryPath); + return Date.now() - stat.mtimeMs > lockOwnerGraceMs; + } catch { + return true; + } +} + +function parseLockOwner(raw: unknown): ILockOwner | undefined { + if (typeof raw !== 'object' || raw === null) { + return undefined; + } + const obj = raw as Record; + if (typeof obj.pid !== 'number' || typeof obj.instanceId !== 'string') { + return undefined; + } + return { pid: obj.pid, instanceId: obj.instanceId }; +} + +function isSameLockOwner(a: IAgentHostEndpointIdentity | ILockOwner, b: IAgentHostEndpointIdentity | ILockOwner): boolean { + return a.pid === b.pid && a.instanceId === b.instanceId; +} + +function isAlreadyExists(error: unknown): boolean { + return (error as NodeJS.ErrnoException | undefined)?.code === 'EEXIST'; +} + +function delay(ms: number): Promise { + return new Promise(resolve => setTimeout(resolve, ms)); +} + +/** Synchronous bounded sleep, only used on the shutdown/cleanup path where an `async` wait is not usable (dispose() is synchronous). */ +function sleepSync(ms: number): void { + Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, ms); +} + +// #endregion + +function isNotFound(error: unknown): boolean { + return (error as NodeJS.ErrnoException | undefined)?.code === 'ENOENT'; } async function applyWindowsOwnerOnlyAcl(path: string): Promise { @@ -193,7 +507,3 @@ function runWindowsCommand(command: string, args: readonly string[]): Promise error ? reject(error) : resolve(String(stdout))); }); } - -function isNotFound(error: unknown): boolean { - return (error as NodeJS.ErrnoException | undefined)?.code === 'ENOENT'; -} diff --git a/src/vs/platform/agentHost/test/node/agentHostEndpointRegistry.test.ts b/src/vs/platform/agentHost/test/node/agentHostEndpointRegistry.test.ts new file mode 100644 index 00000000000..e09c23a0d66 --- /dev/null +++ b/src/vs/platform/agentHost/test/node/agentHostEndpointRegistry.test.ts @@ -0,0 +1,110 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'assert'; +import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js'; +import { + AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION, + IAgentHostEndpointMetadata, + dedupeAgentHostEndpointMetadata, + getAgentHostEndpointIdentityKey, + parseAgentHostEndpointMetadataEntry, + parseAgentHostEndpointRegistry, + removeAgentHostEndpointMetadata, + upsertAgentHostEndpointMetadata, +} from '../../common/agentHostEndpointRegistry.js'; + +function createEntry(overrides: Partial = {}): IAgentHostEndpointMetadata { + return { + schemaVersion: AGENT_HOST_ENDPOINT_REGISTRY_SCHEMA_VERSION, + type: 'editor', + pid: 100, + instanceId: 'a', + protocolVersion: '0.7.0', + connectionToken: 'token', + endpoint: { type: 'socket', path: '/tmp/vscode-ah/a.sock' }, + ...overrides, + }; +} + +suite('Agent Host Endpoint Registry (schema v2)', () => { + ensureNoDisposablesAreLeakedInTestSuite(); + + test('structurally validates every entry, discriminating the endpoint union', () => { + const valid = createEntry(); + const validTcp = createEntry({ type: 'standalone', endpoint: { type: 'tcp', host: '127.0.0.1', port: 12345 } }); + const cases: unknown[] = [ + valid, + validTcp, + null, + 42, + 'not an object', + { ...valid, pid: 'not-a-number' }, + { ...valid, pid: -1 }, + { ...valid, type: 'bogus-type' }, + { ...valid, endpoint: { type: 'bogus-kind', path: '/x' } }, + { ...valid, endpoint: { type: 'tcp', host: '127.0.0.1', port: 'not-a-number' } }, + { ...valid, endpoint: { type: 'tcp', host: '127.0.0.1', port: 70000 } }, + { ...valid, connectionToken: undefined }, + ]; + + assert.deepStrictEqual( + cases.map(entry => parseAgentHostEndpointMetadataEntry(entry) !== undefined), + [true, true, false, false, false, false, false, false, false, false, false, false], + ); + }); + + test('ignores unsupported schema versions without failing the whole array', () => { + const current = createEntry({ instanceId: 'current' }); + const raw = [ + current, + { ...current, instanceId: 'v1', schemaVersion: 1 }, + { ...current, instanceId: 'v3', schemaVersion: 3 }, + { not: 'a valid entry' }, + 'garbage', + ]; + + assert.deepStrictEqual( + parseAgentHostEndpointRegistry(raw).map(entry => entry.instanceId), + ['current'], + ); + }); + + test('a non-array top-level value yields an empty registry', () => { + assert.deepStrictEqual(parseAgentHostEndpointRegistry({ not: 'an array' }), []); + assert.deepStrictEqual(parseAgentHostEndpointRegistry(null), []); + }); + + test('dedupes by (type, pid, instanceId), the later entry winning', () => { + const stale = createEntry({ instanceId: 'dup', connectionToken: 'stale-token' }); + const fresh = createEntry({ instanceId: 'dup', connectionToken: 'fresh-token' }); + const other = createEntry({ pid: 200, instanceId: 'other' }); + const sameIdsDifferentType = createEntry({ type: 'standalone', instanceId: 'dup', endpoint: { type: 'tcp', host: '127.0.0.1', port: 1 } }); + + const deduped = dedupeAgentHostEndpointMetadata([stale, other, sameIdsDifferentType, fresh]); + + assert.deepStrictEqual( + deduped.map(getAgentHostEndpointIdentityKey).sort(), + [other, sameIdsDifferentType, fresh].map(getAgentHostEndpointIdentityKey).sort(), + ); + assert.strictEqual(deduped.find(entry => entry.type === 'editor' && entry.instanceId === 'dup')?.connectionToken, 'fresh-token'); + }); + + test('upsert replaces only the matching identity, leaving other writers untouched', () => { + const a = createEntry({ instanceId: 'a' }); + const b = createEntry({ pid: 200, instanceId: 'b', type: 'standalone', endpoint: { type: 'tcp', host: '127.0.0.1', port: 1 } }); + const updatedA = createEntry({ instanceId: 'a', connectionToken: 'updated-token' }); + + assert.deepStrictEqual(upsertAgentHostEndpointMetadata([a, b], updatedA), [b, updatedA]); + }); + + test('remove takes only the exact-identity entry (same PID is not enough)', () => { + const owner = createEntry({ instanceId: 'a' }); + const impostor = createEntry({ instanceId: 'a-impostor' }); // same (type, pid), different instanceId + const other = createEntry({ pid: 200, instanceId: 'b' }); + + assert.deepStrictEqual(removeAgentHostEndpointMetadata([owner, impostor, other], owner), [impostor, other]); + }); +}); diff --git a/src/vs/platform/agentHost/test/node/localAgentHostMetadata.test.ts b/src/vs/platform/agentHost/test/node/localAgentHostMetadata.test.ts index e338983f7be..3ccf2eacb8e 100644 --- a/src/vs/platform/agentHost/test/node/localAgentHostMetadata.test.ts +++ b/src/vs/platform/agentHost/test/node/localAgentHostMetadata.test.ts @@ -4,6 +4,7 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; +import { spawnSync } from 'child_process'; import * as fs from 'fs'; import * as os from 'os'; import { dirname, join } from '../../../../base/common/path.js'; @@ -37,10 +38,10 @@ suite('Local Agent Host Endpoint Metadata', () => { pid: metadata.pid, protocolVersion: metadata.protocolVersion, tokenLength: metadata.connectionToken.length, - isScoped: metadata.endpointPath !== other.endpointPath, + isScoped: metadata.endpoint.path !== other.endpoint.path, }, { type: 'editor', - schemaVersion: 1, + schemaVersion: 2, pid: process.pid, protocolVersion: metadata.protocolVersion, tokenLength: 43, @@ -54,8 +55,8 @@ suite('Local Agent Host Endpoint Metadata', () => { const metadata = createLocalAgentHostEndpointMetadata(deeplyNested); assert.deepStrictEqual({ - isUnderTemp: dirname(dirname(metadata.endpointPath)) === os.tmpdir(), - isShort: Buffer.byteLength(metadata.endpointPath) < 104, + isUnderTemp: dirname(dirname(metadata.endpoint.path)) === os.tmpdir(), + isShort: Buffer.byteLength(metadata.endpoint.path) < 104, }, { isUnderTemp: true, isShort: true, @@ -68,27 +69,90 @@ suite('Local Agent Host Endpoint Metadata', () => { }); } - test('atomically replaces and owner-checks metadata', async () => { + test('preserves distinct writers and removes only the exact owner', async () => { const first = createLocalAgentHostEndpointMetadata(userDataPath); const second = createLocalAgentHostEndpointMetadata(userDataPath); await publishLocalAgentHostEndpointMetadata(userDataPath, first); await publishLocalAgentHostEndpointMetadata(userDataPath, second); + const publishedBoth = JSON.parse(await fs.promises.readFile(metadataPath, 'utf8')); + cleanupLocalAgentHostEndpointMetadataSync(userDataPath, first); - const published = JSON.parse(await fs.promises.readFile(metadataPath, 'utf8')); + const publishedAfterFirstRemoved = JSON.parse(await fs.promises.readFile(metadataPath, 'utf8')); cleanupLocalAgentHostEndpointMetadataSync(userDataPath, second); assert.deepStrictEqual({ - published, + publishedBoth, + publishedAfterFirstRemoved, removed: !fs.existsSync(metadataPath), files: await fs.promises.readdir(dirname(metadataPath)), }, { - published: [second], + publishedBoth: [first, second], + publishedAfterFirstRemoved: [second], removed: true, files: [], }); }); + test('concurrent writers preserve every entry (no lost updates)', async () => { + const writers = Array.from({ length: 5 }, () => createLocalAgentHostEndpointMetadata(userDataPath)); + + await Promise.all(writers.map(metadata => publishLocalAgentHostEndpointMetadata(userDataPath, metadata))); + + const published: Array<{ instanceId: string }> = JSON.parse(await fs.promises.readFile(metadataPath, 'utf8')); + assert.deepStrictEqual( + new Set(published.map(entry => entry.instanceId)), + new Set(writers.map(writer => writer.instanceId)), + ); + }); + + test('reclaims a lock abandoned by a dead process', async () => { + // A process that has already exited by the time spawnSync returns, so + // its PID is guaranteed to no longer be alive. `process.execPath` is + // Electron under the unit test runner, so it must be told to run as + // plain Node (matching the pattern used elsewhere in this codebase, + // e.g. node/claude/claudeSdkOptions.ts) rather than launch the full app. + const deadPid = spawnSync(process.execPath, ['-e', '0'], { env: { ...process.env, ELECTRON_RUN_AS_NODE: '1' } }).pid!; + const lockDirectoryPath = `${metadataPath}.lock`; + await fs.promises.mkdir(lockDirectoryPath); + await fs.promises.writeFile(join(lockDirectoryPath, 'owner.json'), JSON.stringify({ pid: deadPid, instanceId: 'stale-owner' }), 'utf8'); + + const metadata = createLocalAgentHostEndpointMetadata(userDataPath); + await publishLocalAgentHostEndpointMetadata(userDataPath, metadata); + + assert.deepStrictEqual({ + published: JSON.parse(await fs.promises.readFile(metadataPath, 'utf8')).map((entry: { instanceId: string }) => entry.instanceId), + lockRemoved: !fs.existsSync(lockDirectoryPath), + }, { + published: [metadata.instanceId], + lockRemoved: true, + }); + }); + + test('fails closed (throws) rather than bypassing a lock a live process holds', async function () { + this.timeout(10_000); + + // Our own PID is alive by definition, so recording it (with a different + // instanceId) as the lock owner simulates another live writer holding + // the lock for the entire bounded acquisition window. + const lockDirectoryPath = `${metadataPath}.lock`; + await fs.promises.mkdir(lockDirectoryPath); + await fs.promises.writeFile(join(lockDirectoryPath, 'owner.json'), JSON.stringify({ pid: process.pid, instanceId: 'other-live-writer' }), 'utf8'); + + const metadata = createLocalAgentHostEndpointMetadata(userDataPath); + await assert.rejects(() => publishLocalAgentHostEndpointMetadata(userDataPath, metadata)); + + assert.deepStrictEqual({ + metadataFileWritten: fs.existsSync(metadataPath), + lockStillHeld: fs.existsSync(lockDirectoryPath), + }, { + metadataFileWritten: false, + lockStillHeld: true, + }); + + await fs.promises.rm(lockDirectoryPath, { recursive: true, force: true }); + }); + if (process.platform !== 'win32') { test('writes owner-only metadata permissions', async () => { const metadata = createLocalAgentHostEndpointMetadata(userDataPath);