From da1702e48788ccccb1c89e18d55877bc68cbd7ac Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Fri, 21 Aug 2026 13:11:45 +1000 Subject: [PATCH 1/2] fix(acp): defer authentication until required Signed-off-by: Matt Toohey --- crates/acp-client/src/driver.rs | 294 ++++++++++++++++++++++++++++---- 1 file changed, 262 insertions(+), 32 deletions(-) diff --git a/crates/acp-client/src/driver.rs b/crates/acp-client/src/driver.rs index a251f059..e2bad532 100644 --- a/crates/acp-client/src/driver.rs +++ b/crates/acp-client/src/driver.rs @@ -19,7 +19,7 @@ use agent_client_protocol::{ schema::{ v1::{ AgentCapabilities, AuthMethod, AuthenticateRequest, CancelNotification, - ContentBlock as AcpContentBlock, ContentChunk, ImageContent, Implementation, + ContentBlock as AcpContentBlock, ContentChunk, ErrorCode, ImageContent, Implementation, InitializeRequest, InitializeResponse, LoadSessionRequest, McpCapabilities, McpServer, NewSessionRequest, PermissionOption as SchemaPermissionOption, PermissionOptionId, PermissionOptionKind as SchemaPermissionOptionKind, PromptRequest, PromptResponse, @@ -31,7 +31,7 @@ use agent_client_protocol::{ }, ProtocolVersion, }, - Agent, ByteStreams, Client, ConnectionTo, + Agent, ByteStreams, Client, ConnectionTo, JsonRpcRequest, }; use async_trait::async_trait; use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; @@ -2343,7 +2343,9 @@ async fn setup_acp_session(context: AcpSessionSetupContext<'_>) -> Result) -> Result) -> Result { - let new_session_request = - NewSessionRequest::new(working_dir.to_path_buf()).mcp_servers(mcp_servers.to_vec()); - let session_response = connection - .send_request(new_session_request) - .block_task() - .await - .map_err(|e| format!("Failed to create ACP session: {e:?}"))?; + let new_session_request = || { + NewSessionRequest::new(working_dir.to_path_buf()).mcp_servers(mcp_servers.to_vec()) + }; + let session_response = send_session_setup_request( + connection, + new_session_request, + &init_response.auth_methods, + "create ACP session", + ) + .await?; let new_id = session_response.session_id.to_string(); store @@ -2486,16 +2494,41 @@ fn describe_auth_methods(auth_methods: &[AuthMethod]) -> String { .join(", ") } -async fn authenticate_if_advertised( +async fn send_session_setup_request( + connection: &ConnectionTo, + make_request: MakeRequest, + auth_methods: &[AuthMethod], + operation: &str, +) -> Result +where + Req: JsonRpcRequest, + MakeRequest: Fn() -> Req, +{ + match connection.send_request(make_request()).block_task().await { + Ok(response) => Ok(response), + Err(error) if error.code == ErrorCode::AuthRequired => { + authenticate_with_usable_method(connection, auth_methods).await?; + connection + .send_request(make_request()) + .block_task() + .await + .map_err(|error| format!("Failed to {operation} after authentication: {error:?}")) + } + Err(error) => Err(format!("Failed to {operation}: {error:?}")), + } +} + +async fn authenticate_with_usable_method( connection: &ConnectionTo, auth_methods: &[AuthMethod], ) -> Result<(), String> { - let Some(method) = auth_methods.first() else { - return Ok(()); - }; + let method = select_auth_method(auth_methods).ok_or_else(|| { + "ACP authentication is required, but the agent advertised no usable authentication method" + .to_string() + })?; log::debug!( - "ACP agent advertised authentication methods; selecting {} ({})", + "ACP authentication required; selecting {} ({})", method.name(), method.id() ); @@ -2504,9 +2537,9 @@ async fn authenticate_if_advertised( .send_request(AuthenticateRequest::new(method.id().clone())) .block_task() .await - .map_err(|e| { + .map_err(|error| { format!( - "ACP authentication failed with method {} ({}): {e:?}", + "ACP authentication failed with method {} ({}): {error:?}", method.name(), method.id() ) @@ -2515,6 +2548,46 @@ async fn authenticate_if_advertised( Ok(()) } +fn select_auth_method(auth_methods: &[AuthMethod]) -> Option<&AuthMethod> { + // Agent-managed methods do not require the client to collect credentials. + // Prefer a non-API-key method so providers such as Codex use an existing + // browser login instead of an unavailable key merely because `api-key` was + // advertised first. + auth_methods + .iter() + .find(|method| matches!(method, AuthMethod::Agent(_)) && !looks_like_api_key(method)) + .or_else(|| { + auth_methods + .iter() + .find(|method| auth_method_is_usable(method)) + }) +} + +fn auth_method_is_usable(method: &AuthMethod) -> bool { + match method { + // Agent-managed methods own their credential lookup. If no better + // method is available, let the agent report any missing credential. + AuthMethod::Agent(_) => true, + AuthMethod::EnvVar(method) => method.vars.iter().all(|var| { + var.optional || std::env::var_os(&var.name).is_some_and(|value| !value.is_empty()) + }), + // Terminal methods require a separate interactive client flow that the + // driver does not currently implement. Unknown future methods are also + // unusable until the client explicitly supports their flow. + AuthMethod::Terminal(_) => false, + _ => false, + } +} + +fn looks_like_api_key(method: &AuthMethod) -> bool { + [method.id().to_string(), method.name().to_string()] + .iter() + .any(|value| { + let value = value.to_ascii_lowercase(); + value.contains("api") && value.contains("key") + }) +} + fn build_prompt_content_blocks( prompt: &str, images: &[(String, String)], @@ -2668,17 +2741,20 @@ mod tests { is_config_selection_unavailable_error, is_missing_mcp_transport_error, mcp_server_transport_supported, permission_response_for_options, remote_acp_segments, resolve_acp_working_dir, resolve_session_config_option_selection, - resolve_spawn_working_dir, sanitize_remote_acp_chunk, shell_exec_line, shell_quote, - AcpDriver, AcpEventMetadata, AcpNotificationHandler, AcpPermissionOption, - AcpPermissionOptionKind, AcpPermissionRequest, AcpSessionConfigOptionSelection, + resolve_spawn_working_dir, sanitize_remote_acp_chunk, select_auth_method, + send_session_setup_request, setup_acp_session, shell_exec_line, shell_quote, AcpDriver, + AcpEventMetadata, AcpNotificationHandler, AcpPermissionOption, AcpPermissionOptionKind, + AcpPermissionRequest, AcpSessionConfigOptionSelection, AcpSessionSetupContext, AgentRunOutcome, BasicMessageWriter, MessageWriter, RemoteLineOutcome, ReplayBoundary, - ReplayBuffer, ReplayEvent, + ReplayBuffer, ReplayEvent, Store, }; use agent_client_protocol::schema::v1::{ + AuthMethod, AuthMethodAgent, AuthenticateRequest, AuthenticateResponse, ContentBlock as AcpContentBlock, McpCapabilities, McpServer, McpServerHttp, McpServerSse, - McpServerStdio, PermissionOption, PermissionOptionKind, Plan, PlanEntry, PlanEntryPriority, - PlanEntryStatus, RequestPermissionOutcome, SessionConfigOption, - SessionConfigOptionCategory, SessionConfigSelectOption, SessionNotification, SessionUpdate, + McpServerStdio, NewSessionRequest, NewSessionResponse, PermissionOption, + PermissionOptionKind, Plan, PlanEntry, PlanEntryPriority, PlanEntryStatus, + RequestPermissionOutcome, SessionConfigOption, SessionConfigOptionCategory, + SessionConfigSelectOption, SessionNotification, SessionUpdate, SetSessionConfigOptionRequest, SetSessionConfigOptionResponse, StopReason, }; use std::ffi::OsString; @@ -2687,6 +2763,160 @@ mod tests { use std::time::{SystemTime, UNIX_EPOCH}; use tokio_util::sync::CancellationToken; + #[test] + fn auth_selection_prefers_chat_login_over_first_advertised_api_key() { + let methods = vec![ + AuthMethod::Agent(AuthMethodAgent::new("api-key", "API Key")), + AuthMethod::Agent(AuthMethodAgent::new("chat-gpt", "ChatGPT")), + ]; + + let selected = select_auth_method(&methods).expect("chat login should be usable"); + + assert_eq!(selected.id().to_string(), "chat-gpt"); + } + + #[tokio::test(flavor = "current_thread")] + async fn session_setup_authenticates_with_usable_method_and_retries_on_auth_required() { + let calls = Arc::new(Mutex::new(Vec::::new())); + let calls_for_auth = Arc::clone(&calls); + let calls_for_session = Arc::clone(&calls); + let session_attempts = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let session_attempts_for_handler = Arc::clone(&session_attempts); + let agent = agent_client_protocol::Agent + .builder() + .on_receive_request( + async move |request: AuthenticateRequest, responder, _cx| { + calls_for_auth + .lock() + .unwrap() + .push(format!("authenticate:{}", request.method_id)); + responder.respond(AuthenticateResponse::new()) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |_request: NewSessionRequest, responder, _cx| { + calls_for_session.lock().unwrap().push("session/new".into()); + if session_attempts_for_handler + .fetch_add(1, std::sync::atomic::Ordering::SeqCst) + == 0 + { + responder.respond_with_error(agent_client_protocol::Error::auth_required()) + } else { + responder.respond(NewSessionResponse::new("session-1")) + } + }, + agent_client_protocol::on_receive_request!(), + ); + let methods = vec![ + AuthMethod::Agent(AuthMethodAgent::new("api-key", "API Key")), + AuthMethod::Agent(AuthMethodAgent::new("chat-gpt", "ChatGPT")), + ]; + + agent_client_protocol::Client + .connect_with(agent, async |connection| { + send_session_setup_request( + &connection, + || NewSessionRequest::new(PathBuf::from("/tmp")), + &methods, + "create ACP session", + ) + .await + .map(|_| ()) + .map_err(agent_client_protocol::util::internal_error) + }) + .await + .expect("protocol should succeed"); + + assert_eq!( + calls.lock().unwrap().as_slice(), + &["session/new", "authenticate:chat-gpt", "session/new"] + ); + } + + #[derive(Default)] + struct RecordingStore { + agent_session_ids: Mutex>, + } + + #[async_trait::async_trait] + impl Store for RecordingStore { + fn set_agent_session_id( + &self, + session_id: &str, + agent_session_id: &str, + ) -> Result<(), String> { + self.agent_session_ids + .lock() + .unwrap() + .push((session_id.to_string(), agent_session_id.to_string())); + Ok(()) + } + } + + #[tokio::test(flavor = "current_thread")] + async fn full_session_setup_uses_existing_login_without_eager_authentication() { + use agent_client_protocol::schema::v1::{InitializeRequest, InitializeResponse}; + + let calls = Arc::new(Mutex::new(Vec::::new())); + let calls_for_auth = Arc::clone(&calls); + let calls_for_session = Arc::clone(&calls); + let agent = agent_client_protocol::Agent + .builder() + .on_receive_request( + async |request: InitializeRequest, responder, _cx| { + responder.respond( + InitializeResponse::new(request.protocol_version).auth_methods(vec![ + AuthMethod::Agent(AuthMethodAgent::new("api-key", "API Key")), + AuthMethod::Agent(AuthMethodAgent::new("chat-gpt", "ChatGPT")), + ]), + ) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: AuthenticateRequest, responder, _cx| { + calls_for_auth + .lock() + .unwrap() + .push(format!("authenticate:{}", request.method_id)); + responder.respond_with_error(agent_client_protocol::Error::internal_error()) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |_request: NewSessionRequest, responder, _cx| { + calls_for_session.lock().unwrap().push("session/new".into()); + responder.respond(NewSessionResponse::new("session-1")) + }, + agent_client_protocol::on_receive_request!(), + ); + let store: Arc = Arc::new(RecordingStore::default()); + let writer: Arc = Arc::new(BasicMessageWriter::new()); + + agent_client_protocol::Client + .connect_with(agent, async |connection| { + setup_acp_session(AcpSessionSetupContext { + connection: &connection, + working_dir: Path::new("/tmp"), + store: &store, + writer: &writer, + our_session_id: "local-session", + acp_session_id: None, + config_options: &[], + mcp_servers: &[], + agent_label: "Codex", + }) + .await + .map(|_| ()) + .map_err(agent_client_protocol::util::internal_error) + }) + .await + .expect("existing login should create a session without authenticate"); + + assert_eq!(calls.lock().unwrap().as_slice(), &["session/new"]); + } + fn unique_test_dir(prefix: &str) -> PathBuf { let nonce = SystemTime::now() .duration_since(UNIX_EPOCH) From bf76d48ffbb81a110f7f3af9865c91c88d5c30c8 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Fri, 21 Aug 2026 14:29:32 +1000 Subject: [PATCH 2/2] fix(acp): require explicit authentication handling Signed-off-by: Matt Toohey --- apps/staged/src-tauri/src/doctor.rs | 117 ++- apps/staged/src-tauri/src/lib.rs | 2 + apps/staged/src-tauri/src/pikchr_mcp.rs | 1 + .../staged/src-tauri/src/pikchr_subsession.rs | 4 + apps/staged/src-tauri/src/session_runner.rs | 1 + apps/staged/src-tauri/src/web_server.rs | 11 + apps/staged/src/lib/commands.ts | 17 + .../features/sessions/SessionChatPane.svelte | 111 +++ .../features/sessions/authRecovery.test.ts | 68 ++ .../src/lib/features/sessions/authRecovery.ts | 48 ++ crates/acp-client/src/driver.rs | 480 +++++++++-- crates/acp-client/src/lib.rs | 9 +- crates/acp-client/src/simple.rs | 3 + crates/doctor/src/command.rs | 6 +- crates/doctor/src/lib.rs | 753 +++++++++++++++++- 15 files changed, 1500 insertions(+), 131 deletions(-) create mode 100644 apps/staged/src/lib/features/sessions/authRecovery.test.ts create mode 100644 apps/staged/src/lib/features/sessions/authRecovery.ts diff --git a/apps/staged/src-tauri/src/doctor.rs b/apps/staged/src-tauri/src/doctor.rs index 6feae12c..2b64339d 100644 --- a/apps/staged/src-tauri/src/doctor.rs +++ b/apps/staged/src-tauri/src/doctor.rs @@ -1,13 +1,35 @@ //! Tauri command wrappers for the doctor health-check system. +use std::collections::HashMap; use std::path::{Path, PathBuf}; +use std::sync::{Mutex, OnceLock}; + +use serde::Serialize; pub use doctor::types::{AuthStatus, InstallSource}; pub use doctor::{ - AgentVersionInfo, CheckStatus, DoctorCheck, DoctorReport, ExecuteFixOptions, FixType, - RunChecksOptions, + AgentVersionInfo, CheckStatus, DoctorCheck, DoctorReport, ExecuteFixOptions, FixStdin, + FixStdinWriter, FixType, RunChecksOptions, }; +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct DoctorLoginOutput { + pub check_id: String, + pub line: Option, + pub done: bool, + pub error: Option, +} + +/// Writers for login fixes currently owned by the UI. This is intentionally +/// only a lifetime map for active subprocesses, not a cache of authentication +/// state; doctor remains the source of truth for whether login is available. +static ACTIVE_LOGINS: OnceLock>> = OnceLock::new(); + +fn active_logins() -> &'static Mutex> { + ACTIVE_LOGINS.get_or_init(|| Mutex::new(HashMap::new())) +} + /// Environment snapshot for doctor checks and fixes. Shaped through /// `apply_managed_tools_env` so checks resolve binaries from the same PATH /// the agent spawn path uses — a bridge Staged manages must never be @@ -55,10 +77,17 @@ fn execute_fix_options( command_override: Option, env_vars: Vec<(String, String)>, ) -> ExecuteFixOptions { + // Everything else stays at doctor's defaults: Staged's fixes are + // non-interactive, so nothing here feeds a prompt and the child keeps + // inheriting stdin rather than getting a piped one; the standard fix + // timeout is far above any install or login this runs. Spelled with + // `..Default::default()` so a new doctor option doesn't break this + // workspace-excluded crate, which `cargo check` under `crates/` never + // compiles but `staged-ci.yml` does. ExecuteFixOptions { command_override, npm_registry: crate::managed_acp_tools::npm_registry().map(str::to_string), - env: None, + ..Default::default() } .with_env_snapshot(env_vars) } @@ -106,17 +135,77 @@ async fn run_doctor_report(check_freshness: bool) -> DoctorReport { report } -/// Run a fix for a doctor check, identified by check ID and fix type. -/// -/// The actual shell command is looked up from the static check definitions — -/// the caller never sends a raw command string. Two families of fixes are -/// native rather than shell commands: the node-runtime fix (re)installs the -/// pinned managed runtime, and install fixes for the managed ACP bridges run -/// the floating managed installer so the bridge lands in -/// `~/.staged/packages/tools` with an absolute-path shim instead of the -/// crate's `npm install -g`. Remaining npm-backed fixes install the managed -/// runtime first, since they run npm from it into the private prefix (the -/// existing "Running…" spinner covers the one-time download). +/// Start an interactive login fix and stream its output to the frontend. +#[tauri::command] +pub async fn start_doctor_login( + app_handle: tauri::AppHandle, + check_id: String, +) -> Result<(), String> { + doctor::agents::lookup_fix_command(&check_id, &FixType::Auth) + .ok_or_else(|| format!("No login fix available for {check_id}"))?; + let env_vars = doctor_env_vars().await; + let (writer, stdin) = FixStdin::pipe(); + { + let mut logins = active_logins().lock().unwrap_or_else(|e| e.into_inner()); + if logins.contains_key(&check_id) { + return Err(format!("A login is already running for {check_id}")); + } + logins.insert(check_id.clone(), writer); + } + + let event_check_id = check_id.clone(); + let event_app = app_handle.clone(); + tokio::spawn(async move { + let result = doctor::execute_fix_streaming_with_env_options( + check_id.clone(), + FixType::Auth, + ExecuteFixOptions::default() + .with_env_snapshot(env_vars) + .with_stdin(stdin), + move |line| { + crate::web_server::emit_to_all( + &event_app, + "doctor-login-output", + DoctorLoginOutput { + check_id: event_check_id.clone(), + line: Some(line.to_string()), + done: false, + error: None, + }, + ); + }, + ) + .await; + + active_logins() + .lock() + .unwrap_or_else(|e| e.into_inner()) + .remove(&check_id); + crate::web_server::emit_to_all( + &app_handle, + "doctor-login-output", + DoctorLoginOutput { + check_id, + line: None, + done: true, + error: result.err(), + }, + ); + }); + Ok(()) +} + +#[tauri::command] +pub fn send_doctor_login_code(check_id: String, code: String) -> Result<(), String> { + let writer = active_logins() + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(&check_id) + .cloned() + .ok_or_else(|| format!("No active login for {check_id}"))?; + writer.send_line(code) +} + #[tauri::command] pub async fn run_doctor_fix(check_id: String, fix_type: FixType) -> Result<(), String> { if check_id == NODE_RUNTIME_CHECK_ID { diff --git a/apps/staged/src-tauri/src/lib.rs b/apps/staged/src-tauri/src/lib.rs index 5d1f18ad..6605cfde 100644 --- a/apps/staged/src-tauri/src/lib.rs +++ b/apps/staged/src-tauri/src/lib.rs @@ -2375,6 +2375,8 @@ pub fn run() { doctor::run_doctor, doctor::run_doctor_freshness, doctor::run_doctor_fix, + doctor::start_doctor_login, + doctor::send_doctor_login_code, doctor::run_doctor_update, ]) .build(tauri::generate_context!()) diff --git a/apps/staged/src-tauri/src/pikchr_mcp.rs b/apps/staged/src-tauri/src/pikchr_mcp.rs index 48fbe98f..0629c2db 100644 --- a/apps/staged/src-tauri/src/pikchr_mcp.rs +++ b/apps/staged/src-tauri/src/pikchr_mcp.rs @@ -1522,6 +1522,7 @@ arrow from COLL.e to SNOW.w"#; cancel_token: &CancellationToken, _agent_session_id: Option<&str>, _config_options: &[acp_client::AcpSessionConfigOptionSelection], + _auth_selection: Option<&acp_client::AcpAuthenticationSelection>, ) -> Result { assert!( self.registry.cancel(session_id), diff --git a/apps/staged/src-tauri/src/pikchr_subsession.rs b/apps/staged/src-tauri/src/pikchr_subsession.rs index 5d58fab2..23cde89c 100644 --- a/apps/staged/src-tauri/src/pikchr_subsession.rs +++ b/apps/staged/src-tauri/src/pikchr_subsession.rs @@ -328,6 +328,7 @@ async fn generate_pikchr_source_inner( cancel_token, agent_session_id.as_deref(), config_options, + None, ) .await; writer_dyn.finalize().await; @@ -559,6 +560,7 @@ mod tests { _cancel_token: &CancellationToken, agent_session_id: Option<&str>, config_options: &[acp_client::AcpSessionConfigOptionSelection], + _auth_selection: Option<&acp_client::AcpAuthenticationSelection>, ) -> Result { *self.calls.lock().unwrap() += 1; self.seen_session_ids @@ -1162,6 +1164,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE cancel_token: &CancellationToken, agent_session_id: Option<&str>, config_options: &[acp_client::AcpSessionConfigOptionSelection], + _auth_selection: Option<&acp_client::AcpAuthenticationSelection>, ) -> Result { self.store .update_session_status( @@ -1182,6 +1185,7 @@ agent: http=false, sse=false). Select a provider that supports MCP over HTTP/SSE cancel_token, agent_session_id, config_options, + None, ) .await } diff --git a/apps/staged/src-tauri/src/session_runner.rs b/apps/staged/src-tauri/src/session_runner.rs index 7e1cfdc9..3324b58a 100644 --- a/apps/staged/src-tauri/src/session_runner.rs +++ b/apps/staged/src-tauri/src/session_runner.rs @@ -723,6 +723,7 @@ pub fn start_session( &cancel_token, agent_session_id.as_deref(), &selected_acp_config_options, + None, ) .await; diff --git a/apps/staged/src-tauri/src/web_server.rs b/apps/staged/src-tauri/src/web_server.rs index 4c98df4d..ffd00739 100644 --- a/apps/staged/src-tauri/src/web_server.rs +++ b/apps/staged/src-tauri/src/web_server.rs @@ -3760,6 +3760,17 @@ async fn dispatch(command: &str, args: Value, state: &WebAppState) -> Result { + let check_id: String = arg(&args, "checkId")?; + crate::doctor::start_doctor_login(app_handle.clone(), check_id).await?; + Ok(Value::Null) + } + "send_doctor_login_code" => { + let check_id: String = arg(&args, "checkId")?; + let code: String = arg(&args, "code")?; + crate::doctor::send_doctor_login_code(check_id, code)?; + Ok(Value::Null) + } "run_doctor_fix" => { let check_id: String = arg(&args, "checkId")?; let fix_type: doctor::FixType = arg(&args, "fixType")?; diff --git a/apps/staged/src/lib/commands.ts b/apps/staged/src/lib/commands.ts index eba9ef37..084d0063 100644 --- a/apps/staged/src/lib/commands.ts +++ b/apps/staged/src/lib/commands.ts @@ -1493,6 +1493,23 @@ export function runDoctorFix( return invokeCommand('run_doctor_fix', { checkId, fixType }); } +/** Start an interactive login fix. Output is delivered through doctor-login-output events. */ +export function startDoctorLogin(checkId: string): Promise { + return invokeCommand('start_doctor_login', { checkId }); +} + +/** Submit a line to an interactive doctor login started by startDoctorLogin. */ +export function sendDoctorLoginCode(checkId: string, code: string): Promise { + return invokeCommand('send_doctor_login_code', { checkId, code }); +} + +export interface DoctorLoginOutput { + checkId: string; + line: string | null; + done: boolean; + error: string | null; +} + /** * Run a source-aware update for a single readout (main CLI or ACP bridge). * diff --git a/apps/staged/src/lib/features/sessions/SessionChatPane.svelte b/apps/staged/src/lib/features/sessions/SessionChatPane.svelte index f1cc61e5..9ce62ac3 100644 --- a/apps/staged/src/lib/features/sessions/SessionChatPane.svelte +++ b/apps/staged/src/lib/features/sessions/SessionChatPane.svelte @@ -75,8 +75,17 @@ sendQueuedSessionMessage, type AcpConfigDiscovery, type AcpConfigSelector, + type DoctorLoginOutput, } from '../../api/commands'; import { listenToEvent, type UnlistenFn } from '../../transport'; + import { openSettings } from '../layout/navigation.svelte'; + import { doctorState, runChecks } from '../doctor/doctor.svelte'; + import { + canOfferLogin, + doctorCheckForProvider, + isAuthCodePrompt, + isAuthenticationError, + } from './authRecovery'; import AcpFixedConfigPicker from '../agents/AcpFixedConfigPicker.svelte'; import { agentState } from '../agents/agent.svelte'; import { @@ -200,6 +209,13 @@ let unlistenStatus: UnlistenFn | null = null; let statusEventVersion = 0; let closed = false; + let loginRunning = $state(false); + let loginError = $state(null); + let loginCodePrompt = $state(false); + let loginCode = $state(''); + let loginOutputUnlisten: UnlistenFn | null = null; + let loginCheck = $derived(doctorCheckForProvider(session?.provider, doctorState.report)); + let canLogin = $derived(canOfferLogin(loginCheck)); let inputText = $state(''); let queuedMessages = $state([]); @@ -565,6 +581,7 @@ closed = true; stopPolling(); unlistenStatus?.(); + loginOutputUnlisten?.(); }); // This pane can be mounted once and reused across opens (the `active` prop toggles @@ -635,6 +652,45 @@ }; }); + async function startLogin() { + if (!session?.provider || !canLogin || loginRunning) return; + loginRunning = true; + loginError = null; + loginCodePrompt = false; + try { + const { startDoctorLogin } = await import('../../api/commands'); + loginOutputUnlisten?.(); + const unlisten = listenToEvent('doctor-login-output', (output) => { + if (output.checkId !== `ai-agent-${session?.provider}`) return; + if (output.line && isAuthCodePrompt(output.line)) loginCodePrompt = true; + if (output.done) { + loginRunning = false; + unlisten(); + loginOutputUnlisten = null; + if (output.error) loginError = output.error; + else void runChecks(); + } + }); + loginOutputUnlisten = unlisten; + await startDoctorLogin(`ai-agent-${session.provider}`); + } catch (e) { + loginRunning = false; + loginError = e instanceof Error ? e.message : String(e); + } + } + + async function submitLoginCode() { + if (!session?.provider || !loginCode.trim()) return; + try { + const { sendDoctorLoginCode } = await import('../../api/commands'); + await sendDoctorLoginCode(`ai-agent-${session.provider}`, loginCode.trim()); + loginCode = ''; + loginCodePrompt = false; + } catch (e) { + loginError = e instanceof Error ? e.message : String(e); + } + } + function isComposerFocused(): boolean { return document.activeElement === inputEl; } @@ -2051,10 +2107,47 @@ was killed from outside with a recorded reason (e.g. a Pikchr child session whose generate_pikchr call timed out) and reads as an error. --> {#if (session?.status === 'error' || session?.status === 'cancelled') && session.errorMessage} + {@const authError = isAuthenticationError(session.errorMessage)} {session.errorMessage} + {#if authError} + +
+ + {#if canLogin} + + {/if} +
+ + {/if} + {#if loginError} +

{loginError}

+ {/if} + {#if loginCodePrompt} + + {/if} {:else if session && session.status !== 'running' && session.status !== 'queued'} {#if isResumableReason(session.completionReason)} {@const isWarning = @@ -2664,6 +2757,24 @@ /* ----- Input wrapper + queue popover ----------------------------------- */ + .auth-actions, + .login-code-row { + display: flex; + align-items: center; + gap: 8px; + flex-wrap: wrap; + } + + .login-code-input { + min-width: 180px; + border: 1px solid var(--border-subtle); + border-radius: 6px; + background: var(--bg-primary); + color: var(--text-primary); + padding: 4px 8px; + font-size: var(--size-xs); + } + .input-wrapper { flex-shrink: 0; } diff --git a/apps/staged/src/lib/features/sessions/authRecovery.test.ts b/apps/staged/src/lib/features/sessions/authRecovery.test.ts new file mode 100644 index 00000000..ee524488 --- /dev/null +++ b/apps/staged/src/lib/features/sessions/authRecovery.test.ts @@ -0,0 +1,68 @@ +import { describe, expect, it } from 'vitest'; +import type { DoctorCheck } from '../../api/commands'; +import { + canOfferLogin, + doctorCheckForProvider, + isAuthCodePrompt, + isAuthenticationError, +} from './authRecovery'; + +function check(overrides: Partial = {}): DoctorCheck { + return { + id: 'ai-agent-claude', + label: 'Claude Code', + status: 'warn', + message: 'Installed, not authenticated', + fixUrl: null, + fixCommand: 'claude-agent-acp --cli auth login', + fixType: 'auth', + path: '/usr/local/bin/claude-agent-acp', + bridgePath: null, + rawOutput: null, + authStatus: 'notAuthenticated', + installedVersion: null, + latestVersion: null, + updateAvailable: null, + installSource: null, + selfUpdating: null, + main: null, + bridge: null, + ...overrides, + }; +} + +describe('authentication recovery helpers', () => { + it.each([ + 'ACP protocol failed: OAuth token has expired; authentication required', + 'Error: missing CODEX_API_KEY (or OPENAI_API_KEY)', + 'nested ACP error: Unauthorized (401)', + ])('recognizes authentication error: %s', (message) => { + expect(isAuthenticationError(message)).toBe(true); + }); + + it('does not turn unrelated failures into authentication actions', () => { + expect(isAuthenticationError('ACP protocol failed: connection refused')).toBe(false); + expect(isAuthenticationError('npm install failed with exit code 1')).toBe(false); + }); + + it('only offers login for a positively detected signed-out agent', () => { + expect(canOfferLogin(check())).toBe(true); + expect(canOfferLogin(check({ authStatus: 'unknown' }))).toBe(false); + expect(canOfferLogin(check({ authStatus: 'authenticated' }))).toBe(false); + expect(canOfferLogin(check({ fixType: null }))).toBe(false); + }); + + it('matches a session provider to the existing doctor report', () => { + const report = { checks: [check(), check({ id: 'ai-agent-codex', label: 'Codex' })] }; + expect(doctorCheckForProvider('codex', report)?.label).toBe('Codex'); + expect(doctorCheckForProvider('pi', report)).toBeNull(); + expect(doctorCheckForProvider(null, report)).toBeNull(); + }); + + it.each(['Enter authentication code:', 'Paste the code here', 'input your token'])( + 'recognizes a code prompt: %s', + (line) => { + expect(isAuthCodePrompt(line)).toBe(true); + } + ); +}); diff --git a/apps/staged/src/lib/features/sessions/authRecovery.ts b/apps/staged/src/lib/features/sessions/authRecovery.ts new file mode 100644 index 00000000..0c077b57 --- /dev/null +++ b/apps/staged/src/lib/features/sessions/authRecovery.ts @@ -0,0 +1,48 @@ +import type { DoctorCheck, DoctorReport } from '../../api/commands'; + +/** Authentication failures commonly arrive wrapped in one or more ACP errors. */ +export function isAuthenticationError(message: string | null | undefined): boolean { + if (!message) return false; + const text = message.toLowerCase(); + + return ( + /authenticat(?:e|ion|ed|ing)/.test(text) || + /auth[_ -]?required/.test(text) || + /unauthori[sz]ed/.test(text) || + /oauth/.test(text) || + /(?:api[_ ]?key|access token|refresh token|credential).*(?:missing|invalid|expired|required)/.test( + text + ) || + /(?:missing|invalid|expired|required).*(?:api[_ ]?key|access token|refresh token|credential)/.test( + text + ) || + /\b(?:codex_api_key|openai_api_key)\b/.test(text) + ); +} + +/** Find the doctor check for the provider recorded on a session. */ +export function doctorCheckForProvider( + provider: string | null | undefined, + report: DoctorReport | null | undefined +): DoctorCheck | null { + if (!provider || !report) return null; + return report.checks.find((check) => check.id === `ai-agent-${provider}`) ?? null; +} + +/** + * A login action is offered only when doctor's existing auth probe positively + * identified a signed-out agent. Unknown status is deliberately not guessed: + * a login command can be ineffective when the binary or its credentials could + * not be detected. + */ +export function canOfferLogin(check: DoctorCheck | null | undefined): boolean { + return check?.authStatus === 'notAuthenticated' && check.fixType === 'auth' && !!check.fixCommand; +} + +export function isAuthCodePrompt(line: string): boolean { + const text = line.toLowerCase(); + return ( + /\b(?:enter|paste|提供|input|type|write|submit)\b.{0,40}\b(?:code|token)\b/.test(text) || + /\b(?:code|token)\b.{0,40}\b(?:enter|paste|input|type|write|submit)\b/.test(text) + ); +} diff --git a/crates/acp-client/src/driver.rs b/crates/acp-client/src/driver.rs index e2bad532..81ebbe58 100644 --- a/crates/acp-client/src/driver.rs +++ b/crates/acp-client/src/driver.rs @@ -267,6 +267,152 @@ pub struct AcpSessionConfigOptionSelection { pub value_id: String, } +/// Provider-specific opt-in to authenticate with a known ACP method. +/// +/// Generic ACP setup never derives this from advertised method IDs, display +/// names, or ordering. Callers may populate it only after their integration has +/// explicit knowledge that the method is safe and currently usable. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AcpAuthenticationSelection { + pub method_id: String, +} + +impl AcpAuthenticationSelection { + pub fn new(method_id: impl Into) -> Self { + Self { + method_id: method_id.into(), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub enum AcpAuthenticationMethodCategory { + AgentManaged, + EnvironmentBacked, + Terminal, + Unsupported, +} + +impl AcpAuthenticationMethodCategory { + fn label(self) -> &'static str { + match self { + Self::AgentManaged => "agent-managed", + Self::EnvironmentBacked => "environment-backed", + Self::Terminal => "terminal", + Self::Unsupported => "unsupported", + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AcpAuthenticationMethod { + pub id: String, + pub display_name: String, + pub description: Option, + pub category: AcpAuthenticationMethodCategory, + pub can_handle: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub struct AcpAuthenticationRequired { + pub methods: Vec, + pub attempted_method_id: Option, +} + +impl AcpAuthenticationRequired { + fn from_auth_methods(auth_methods: &[AuthMethod]) -> Self { + Self { + methods: auth_methods.iter().map(auth_method_details).collect(), + attempted_method_id: None, + } + } + + fn after_authentication_attempt(mut self, method_id: impl Into) -> Self { + self.attempted_method_id = Some(method_id.into()); + self + } + + fn describe(&self, operation: &str) -> String { + let retry = self + .attempted_method_id + .as_deref() + .map(|method_id| { + format!(" after authenticating with explicitly selected method '{method_id}'") + }) + .unwrap_or_default(); + let methods = if self.methods.is_empty() { + "no advertised authentication methods".to_string() + } else { + self.methods + .iter() + .map(|method| { + let support = if method.can_handle { + "client-supported" + } else { + "unsupported by this client" + }; + let description = method + .description + .as_deref() + .filter(|description| !description.trim().is_empty()) + .map(|description| format!(", description: {description}")) + .unwrap_or_default(); + format!( + "{} (id: {}, category: {}, {support}{description})", + method.display_name, + method.id, + method.category.label(), + ) + }) + .collect::>() + .join("; ") + }; + + format!( + "ACP authentication is required to {operation}{retry}; not retrying automatically without an explicit supported authentication method. Advertised methods: {methods}" + ) + } +} + +fn auth_method_details(method: &AuthMethod) -> AcpAuthenticationMethod { + let category = auth_method_category(method); + AcpAuthenticationMethod { + id: method.id().to_string(), + display_name: method.name().to_string(), + description: method.description().map(str::to_string), + category, + can_handle: auth_method_can_be_handled(category), + } +} + +fn auth_method_category(method: &AuthMethod) -> AcpAuthenticationMethodCategory { + match method { + AuthMethod::Agent(_) => AcpAuthenticationMethodCategory::AgentManaged, + AuthMethod::EnvVar(_) => AcpAuthenticationMethodCategory::EnvironmentBacked, + AuthMethod::Terminal(_) => AcpAuthenticationMethodCategory::Terminal, + _ => AcpAuthenticationMethodCategory::Unsupported, + } +} + +fn auth_method_can_be_handled(category: AcpAuthenticationMethodCategory) -> bool { + match category { + // Staged can pass an explicitly selected agent-managed method ID to + // `authenticate`, but the generic driver still must not choose one by + // guessing from provider-defined IDs, names, or list order. + AcpAuthenticationMethodCategory::AgentManaged => true, + // Environment-backed methods need provider-specific confirmation that + // credentials are available in the agent process environment. + AcpAuthenticationMethodCategory::EnvironmentBacked => false, + // Terminal authentication requires a complete interactive terminal + // flow. Until that exists, never send it through `authenticate`. + AcpAuthenticationMethodCategory::Terminal => false, + AcpAuthenticationMethodCategory::Unsupported => false, + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct ReplayBoundary { pub role: String, @@ -382,6 +528,7 @@ pub trait AgentDriver { cancel_token: &CancellationToken, agent_session_id: Option<&str>, config_options: &[AcpSessionConfigOptionSelection], + auth_selection: Option<&AcpAuthenticationSelection>, ) -> Result; } @@ -718,6 +865,7 @@ impl AgentDriver for AcpDriver { cancel_token: &CancellationToken, agent_session_id: Option<&str>, config_options: &[AcpSessionConfigOptionSelection], + auth_selection: Option<&AcpAuthenticationSelection>, ) -> Result { let spawn_working_dir = resolve_spawn_working_dir(working_dir, self.is_remote); let acp_working_dir = resolve_acp_working_dir( @@ -983,6 +1131,7 @@ impl AgentDriver for AcpDriver { &self.mcp_servers, &self.agent_label, cancel_token, + auth_selection, ) .await .map_err(agent_client_protocol::util::internal_error) @@ -1981,6 +2130,7 @@ async fn run_acp_protocol( mcp_servers: &[McpServer], agent_label: &str, cancel_token: &CancellationToken, + auth_selection: Option<&AcpAuthenticationSelection>, ) -> Result { let setup_task = tokio::time::timeout( ACP_SETUP_TIMEOUT, @@ -1994,6 +2144,7 @@ async fn run_acp_protocol( config_options, mcp_servers, agent_label, + auth_selection, }), ); let setup = tokio::select! { @@ -2306,6 +2457,7 @@ struct AcpSessionSetupContext<'a> { config_options: &'a [AcpSessionConfigOptionSelection], mcp_servers: &'a [McpServer], agent_label: &'a str, + auth_selection: Option<&'a AcpAuthenticationSelection>, } async fn setup_acp_session(context: AcpSessionSetupContext<'_>) -> Result { @@ -2319,6 +2471,7 @@ async fn setup_acp_session(context: AcpSessionSetupContext<'_>) -> Result) -> Result) -> Result( connection: &ConnectionTo, make_request: MakeRequest, auth_methods: &[AuthMethod], + auth_selection: Option<&AcpAuthenticationSelection>, operation: &str, ) -> Result where @@ -2507,28 +2663,57 @@ where match connection.send_request(make_request()).block_task().await { Ok(response) => Ok(response), Err(error) if error.code == ErrorCode::AuthRequired => { - authenticate_with_usable_method(connection, auth_methods).await?; - connection - .send_request(make_request()) - .block_task() - .await - .map_err(|error| format!("Failed to {operation} after authentication: {error:?}")) + let required = AcpAuthenticationRequired::from_auth_methods(auth_methods); + let Some(selection) = auth_selection else { + return Err(required.describe(operation)); + }; + + authenticate_with_explicit_method(connection, auth_methods, selection).await?; + let attempted = selection.method_id.clone(); + match connection.send_request(make_request()).block_task().await { + Ok(response) => Ok(response), + Err(error) if error.code == ErrorCode::AuthRequired => Err(required + .after_authentication_attempt(attempted) + .describe(operation)), + Err(error) => Err(format!( + "Failed to {operation} after authentication with explicitly selected method '{}': {error:?}", + selection.method_id + )), + } } Err(error) => Err(format!("Failed to {operation}: {error:?}")), } } -async fn authenticate_with_usable_method( +async fn authenticate_with_explicit_method( connection: &ConnectionTo, auth_methods: &[AuthMethod], + selection: &AcpAuthenticationSelection, ) -> Result<(), String> { - let method = select_auth_method(auth_methods).ok_or_else(|| { - "ACP authentication is required, but the agent advertised no usable authentication method" - .to_string() - })?; + let method = auth_methods + .iter() + .find(|method| method.id().to_string() == selection.method_id) + .ok_or_else(|| { + let required = AcpAuthenticationRequired::from_auth_methods(auth_methods); + format!( + "ACP authentication method '{}' was selected explicitly, but the agent did not advertise it. {}", + selection.method_id, + required.describe("authenticate") + ) + })?; + let details = auth_method_details(method); + if !details.can_handle { + return Err(format!( + "ACP authentication method '{}' ({}) cannot be handled by this client because it is {}. Advertised methods: {}", + details.display_name, + details.id, + details.category.label(), + AcpAuthenticationRequired::from_auth_methods(auth_methods).describe("authenticate") + )); + } log::debug!( - "ACP authentication required; selecting {} ({})", + "ACP authentication required; using explicitly selected method {} ({})", method.name(), method.id() ); @@ -2539,7 +2724,7 @@ async fn authenticate_with_usable_method( .await .map_err(|error| { format!( - "ACP authentication failed with method {} ({}): {error:?}", + "ACP authentication failed with explicitly selected method {} ({}): {error:?}", method.name(), method.id() ) @@ -2548,46 +2733,6 @@ async fn authenticate_with_usable_method( Ok(()) } -fn select_auth_method(auth_methods: &[AuthMethod]) -> Option<&AuthMethod> { - // Agent-managed methods do not require the client to collect credentials. - // Prefer a non-API-key method so providers such as Codex use an existing - // browser login instead of an unavailable key merely because `api-key` was - // advertised first. - auth_methods - .iter() - .find(|method| matches!(method, AuthMethod::Agent(_)) && !looks_like_api_key(method)) - .or_else(|| { - auth_methods - .iter() - .find(|method| auth_method_is_usable(method)) - }) -} - -fn auth_method_is_usable(method: &AuthMethod) -> bool { - match method { - // Agent-managed methods own their credential lookup. If no better - // method is available, let the agent report any missing credential. - AuthMethod::Agent(_) => true, - AuthMethod::EnvVar(method) => method.vars.iter().all(|var| { - var.optional || std::env::var_os(&var.name).is_some_and(|value| !value.is_empty()) - }), - // Terminal methods require a separate interactive client flow that the - // driver does not currently implement. Unknown future methods are also - // unusable until the client explicitly supports their flow. - AuthMethod::Terminal(_) => false, - _ => false, - } -} - -fn looks_like_api_key(method: &AuthMethod) -> bool { - [method.id().to_string(), method.name().to_string()] - .iter() - .any(|value| { - let value = value.to_ascii_lowercase(); - value.contains("api") && value.contains("key") - }) -} - fn build_prompt_content_blocks( prompt: &str, images: &[(String, String)], @@ -2741,12 +2886,13 @@ mod tests { is_config_selection_unavailable_error, is_missing_mcp_transport_error, mcp_server_transport_supported, permission_response_for_options, remote_acp_segments, resolve_acp_working_dir, resolve_session_config_option_selection, - resolve_spawn_working_dir, sanitize_remote_acp_chunk, select_auth_method, - send_session_setup_request, setup_acp_session, shell_exec_line, shell_quote, AcpDriver, - AcpEventMetadata, AcpNotificationHandler, AcpPermissionOption, AcpPermissionOptionKind, - AcpPermissionRequest, AcpSessionConfigOptionSelection, AcpSessionSetupContext, - AgentRunOutcome, BasicMessageWriter, MessageWriter, RemoteLineOutcome, ReplayBoundary, - ReplayBuffer, ReplayEvent, Store, + resolve_spawn_working_dir, sanitize_remote_acp_chunk, send_session_setup_request, + setup_acp_session, shell_exec_line, shell_quote, AcpAuthenticationMethodCategory, + AcpAuthenticationRequired, AcpAuthenticationSelection, AcpDriver, AcpEventMetadata, + AcpNotificationHandler, AcpPermissionOption, AcpPermissionOptionKind, AcpPermissionRequest, + AcpSessionConfigOptionSelection, AcpSessionSetupContext, AgentRunOutcome, + BasicMessageWriter, MessageWriter, RemoteLineOutcome, ReplayBoundary, ReplayBuffer, + ReplayEvent, Store, }; use agent_client_protocol::schema::v1::{ AuthMethod, AuthMethodAgent, AuthenticateRequest, AuthenticateResponse, @@ -2763,20 +2909,96 @@ mod tests { use std::time::{SystemTime, UNIX_EPOCH}; use tokio_util::sync::CancellationToken; + #[derive(Default)] + struct RecordingStore { + agent_session_ids: Mutex>, + } + + #[async_trait::async_trait] + impl Store for RecordingStore { + fn set_agent_session_id( + &self, + session_id: &str, + agent_session_id: &str, + ) -> Result<(), String> { + self.agent_session_ids + .lock() + .unwrap() + .push((session_id.to_string(), agent_session_id.to_string())); + Ok(()) + } + } + #[test] - fn auth_selection_prefers_chat_login_over_first_advertised_api_key() { + fn auth_method_metadata_does_not_guess_from_provider_names() { + let methods = vec![ + AuthMethod::Agent(AuthMethodAgent::new("api-key", "API Key")), + AuthMethod::Agent(AuthMethodAgent::new("chat-gpt", "ChatGPT")), + AuthMethod::Agent(AuthMethodAgent::new("arbitrary", "Arbitrary")), + ]; + + let required = AcpAuthenticationRequired::from_auth_methods(&methods); + + assert_eq!(required.methods.len(), 3); + assert!(required.methods.iter().all(|method| method.can_handle)); + assert!(required + .methods + .iter() + .all(|method| method.category == AcpAuthenticationMethodCategory::AgentManaged)); + assert_eq!(required.methods[0].id, "api-key"); + } + + #[tokio::test(flavor = "current_thread")] + async fn session_setup_returns_auth_required_without_guessing_a_method() { + let calls = Arc::new(Mutex::new(Vec::::new())); + let calls_for_auth = Arc::clone(&calls); + let calls_for_session = Arc::clone(&calls); + let agent = agent_client_protocol::Agent + .builder() + .on_receive_request( + async move |request: AuthenticateRequest, responder, _cx| { + calls_for_auth + .lock() + .unwrap() + .push(format!("authenticate:{}", request.method_id)); + responder.respond(AuthenticateResponse::new()) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |_request: NewSessionRequest, responder, _cx| { + calls_for_session.lock().unwrap().push("session/new".into()); + responder.respond_with_error(agent_client_protocol::Error::auth_required()) + }, + agent_client_protocol::on_receive_request!(), + ); let methods = vec![ AuthMethod::Agent(AuthMethodAgent::new("api-key", "API Key")), AuthMethod::Agent(AuthMethodAgent::new("chat-gpt", "ChatGPT")), ]; - let selected = select_auth_method(&methods).expect("chat login should be usable"); + let error = agent_client_protocol::Client + .connect_with(agent, async |connection| { + send_session_setup_request( + &connection, + || NewSessionRequest::new(PathBuf::from("/tmp")), + &methods, + None, + "create ACP session", + ) + .await + .map(|_| ()) + .map_err(agent_client_protocol::util::internal_error) + }) + .await + .expect_err("auth_required should be surfaced without authenticate"); - assert_eq!(selected.id().to_string(), "chat-gpt"); + assert!(format!("{error:?}").contains("ACP authentication is required")); + assert_eq!(calls.lock().unwrap().as_slice(), &["session/new"]); } #[tokio::test(flavor = "current_thread")] - async fn session_setup_authenticates_with_usable_method_and_retries_on_auth_required() { + async fn session_setup_authenticates_explicit_method_and_retries_once() { let calls = Arc::new(Mutex::new(Vec::::new())); let calls_for_auth = Arc::clone(&calls); let calls_for_session = Arc::clone(&calls); @@ -2812,6 +3034,7 @@ mod tests { AuthMethod::Agent(AuthMethodAgent::new("api-key", "API Key")), AuthMethod::Agent(AuthMethodAgent::new("chat-gpt", "ChatGPT")), ]; + let selection = AcpAuthenticationSelection::new("chat-gpt"); agent_client_protocol::Client .connect_with(agent, async |connection| { @@ -2819,6 +3042,7 @@ mod tests { &connection, || NewSessionRequest::new(PathBuf::from("/tmp")), &methods, + Some(&selection), "create ACP session", ) .await @@ -2834,24 +3058,121 @@ mod tests { ); } - #[derive(Default)] - struct RecordingStore { - agent_session_ids: Mutex>, + #[tokio::test(flavor = "current_thread")] + async fn session_setup_stops_after_second_auth_required() { + let calls = Arc::new(Mutex::new(Vec::::new())); + let calls_for_auth = Arc::clone(&calls); + let calls_for_session = Arc::clone(&calls); + let agent = agent_client_protocol::Agent + .builder() + .on_receive_request( + async move |request: AuthenticateRequest, responder, _cx| { + calls_for_auth + .lock() + .unwrap() + .push(format!("authenticate:{}", request.method_id)); + responder.respond(AuthenticateResponse::new()) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |_request: NewSessionRequest, responder, _cx| { + calls_for_session.lock().unwrap().push("session/new".into()); + responder.respond_with_error(agent_client_protocol::Error::auth_required()) + }, + agent_client_protocol::on_receive_request!(), + ); + let methods = vec![AuthMethod::Agent(AuthMethodAgent::new( + "chat-gpt", "ChatGPT", + ))]; + let selection = AcpAuthenticationSelection::new("chat-gpt"); + + let error = agent_client_protocol::Client + .connect_with(agent, async |connection| { + send_session_setup_request( + &connection, + || NewSessionRequest::new(PathBuf::from("/tmp")), + &methods, + Some(&selection), + "create ACP session", + ) + .await + .map(|_| ()) + .map_err(agent_client_protocol::util::internal_error) + }) + .await + .expect_err("a second auth_required should terminate the setup"); + + let error = format!("{error:?}"); + assert!(error.contains("after authenticating with explicitly selected method 'chat-gpt'")); + assert_eq!( + calls.lock().unwrap().as_slice(), + &["session/new", "authenticate:chat-gpt", "session/new"] + ); } - #[async_trait::async_trait] - impl Store for RecordingStore { - fn set_agent_session_id( - &self, - session_id: &str, - agent_session_id: &str, - ) -> Result<(), String> { - self.agent_session_ids - .lock() - .unwrap() - .push((session_id.to_string(), agent_session_id.to_string())); - Ok(()) - } + #[test] + fn terminal_methods_are_reported_unsupported() { + use agent_client_protocol::schema::v1::AuthMethodTerminal; + + let methods = vec![AuthMethod::Terminal(AuthMethodTerminal::new( + "terminal-login", + "Terminal Login", + ))]; + let required = AcpAuthenticationRequired::from_auth_methods(&methods); + + assert_eq!( + required.methods[0].category, + AcpAuthenticationMethodCategory::Terminal + ); + assert!(!required.methods[0].can_handle); + } + + #[tokio::test(flavor = "current_thread")] + async fn terminal_methods_are_never_sent_to_authenticate() { + use agent_client_protocol::schema::v1::AuthMethodTerminal; + + let auth_called = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let auth_called_for_handler = Arc::clone(&auth_called); + let agent = agent_client_protocol::Agent + .builder() + .on_receive_request( + async move |_request: AuthenticateRequest, responder, _cx| { + auth_called_for_handler.store(true, std::sync::atomic::Ordering::SeqCst); + responder.respond(AuthenticateResponse::new()) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |_request: NewSessionRequest, responder, _cx| { + responder.respond_with_error(agent_client_protocol::Error::auth_required()) + }, + agent_client_protocol::on_receive_request!(), + ); + let methods = vec![AuthMethod::Terminal(AuthMethodTerminal::new( + "terminal-login", + "Terminal Login", + ))]; + let selection = AcpAuthenticationSelection::new("terminal-login"); + + let error = agent_client_protocol::Client + .connect_with(agent, async |connection| { + send_session_setup_request( + &connection, + || NewSessionRequest::new(PathBuf::from("/tmp")), + &methods, + Some(&selection), + "create ACP session", + ) + .await + .map(|_| ()) + .map_err(agent_client_protocol::util::internal_error) + }) + .await + .expect_err("terminal auth is unsupported"); + + assert!(format!("{error:?}").contains("terminal")); + assert!(!auth_called.load(std::sync::atomic::Ordering::SeqCst)); } #[tokio::test(flavor = "current_thread")] @@ -2906,6 +3227,7 @@ mod tests { config_options: &[], mcp_servers: &[], agent_label: "Codex", + auth_selection: None, }) .await .map(|_| ()) diff --git a/crates/acp-client/src/lib.rs b/crates/acp-client/src/lib.rs index 0405c627..36817c1d 100644 --- a/crates/acp-client/src/lib.rs +++ b/crates/acp-client/src/lib.rs @@ -28,10 +28,11 @@ pub use agent_client_protocol::schema::v1::{ }; pub use driver::{ autoapprove_permission_decision, is_config_selection_unavailable_error, - is_missing_mcp_transport_error, strip_code_fences, AcpDriver, AcpEventMetadata, - AcpInitializeMetadata, AcpPermissionDecision, AcpPermissionOption, AcpPermissionOptionKind, - AcpPermissionRequest, AcpSessionConfigOptionSelection, AcpToolCallMetadata, AgentDriver, - AgentRunOutcome, BasicMessageWriter, MessageWriter, ReplayBoundary, Store, + is_missing_mcp_transport_error, strip_code_fences, AcpAuthenticationSelection, AcpDriver, + AcpEventMetadata, AcpInitializeMetadata, AcpPermissionDecision, AcpPermissionOption, + AcpPermissionOptionKind, AcpPermissionRequest, AcpSessionConfigOptionSelection, + AcpToolCallMetadata, AgentDriver, AgentRunOutcome, BasicMessageWriter, MessageWriter, + ReplayBoundary, Store, }; pub use simple::{run_acp_prompt, run_acp_prompt_with_interpreter_env_snapshot}; pub use types::{ diff --git a/crates/acp-client/src/simple.rs b/crates/acp-client/src/simple.rs index 901a02e9..17675f75 100644 --- a/crates/acp-client/src/simple.rs +++ b/crates/acp-client/src/simple.rs @@ -58,6 +58,7 @@ impl AgentDriver for SimpleDriverWrapper { cancel_token: &CancellationToken, agent_session_id: Option<&str>, config_options: &[crate::driver::AcpSessionConfigOptionSelection], + auth_selection: Option<&crate::driver::AcpAuthenticationSelection>, ) -> Result { if !images.is_empty() { log::debug!( @@ -77,6 +78,7 @@ impl AgentDriver for SimpleDriverWrapper { cancel_token, agent_session_id, config_options, + auth_selection, ) .await } @@ -160,6 +162,7 @@ async fn run_acp_prompt_with_options( &cancel_token, None, &[], + None, ) .await .map_err(|e| anyhow::anyhow!("ACP driver error: {e}"))?; diff --git a/crates/doctor/src/command.rs b/crates/doctor/src/command.rs index 4ccc3a2a..c2eb1b4a 100644 --- a/crates/doctor/src/command.rs +++ b/crates/doctor/src/command.rs @@ -193,7 +193,11 @@ fn clean_up_after_incomplete_wait(child: &mut Child) { let _ = child.wait(); } -fn kill_child_process_group_or_child(child: &mut Child) { +/// Best-effort kill: target the child's process group first so a shell's whole +/// command tree goes with it, falling back to the direct child when the group +/// lookup fails (the child wasn't spawned with `process_group(0)`, or it isn't +/// Unix). Callers must still reap afterwards. +pub(crate) fn kill_child_process_group_or_child(child: &mut Child) { if kill_child_process_group(child) { return; } diff --git a/crates/doctor/src/lib.rs b/crates/doctor/src/lib.rs index abdd2325..572d47b2 100644 --- a/crates/doctor/src/lib.rs +++ b/crates/doctor/src/lib.rs @@ -20,6 +20,7 @@ pub use types::{AgentVersionInfo, CheckStatus, DoctorCheck, DoctorReport, FixTyp use std::collections::{HashMap, HashSet}; use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; use agents::{ bundled_version_probe_args, check_single_ai_agent, derive_update_command, lookup_fix_command, @@ -673,6 +674,253 @@ struct FreshnessTarget { version_args: Option<&'static [&'static str]>, } +/// Opt-in piped stdin for a fix subprocess. Create with [`FixStdin::pipe`]; +/// keep the [`FixStdinWriter`], put the `FixStdin` in +/// [`ExecuteFixOptions::stdin`]. +/// +/// Single-use: the first execution claims the pipe, and any later execution +/// handed the same `FixStdin` — or a clone of it, including one carried along by +/// a cloned [`ExecuteFixOptions`] — fails with an error instead of spawning. +/// Retrying a fix needs a fresh pipe. +#[derive(Debug, Clone)] +pub struct FixStdin { + state: Arc>, +} + +/// The pipe's whole life cycle: `Buffered` until the fix spawns, `Live` while it +/// runs, then `Closed` — terminal, and reached when the fix ends, when the last +/// writer drops, or when a write finds the read end gone. Holding the child's +/// stdin handle here rather than in a thread of its own is what lets +/// [`FixStdinWriter::send_line`] write through and report the real outcome. +#[derive(Debug)] +enum FixStdinState { + /// Before the fix spawns: lines the host queued, replayed at spawn. + /// `claimed` marks the execution that reserved this pipe, so a second one + /// is rejected before it spawns. `eof` records that every writer dropped + /// pre-spawn, so the replay is followed immediately by closing the pipe. + Buffered { + lines: Vec, + eof: bool, + claimed: bool, + }, + /// Fix running: writes go straight into the child's stdin. + Live(std::process::ChildStdin), + /// Fix finished, every writer gone, or a write hit a dead pipe. + Closed, +} + +/// Rejection for an execution handed a `FixStdin` another one already claimed. +const FIX_STDIN_REUSED: &str = "FixStdin already consumed by a previous fix execution; \ + create a fresh pipe with FixStdin::pipe() for each run"; + +/// Rejection for a line the pipe cannot deliver because it is closed. +const FIX_STDIN_CLOSED: &str = "Fix is no longer accepting input"; + +/// Locking the pipe state recovers from poisoning instead of propagating it: no +/// invariant spans the lock (the state is a plain enum, and the only work done +/// under it is a `Vec` push or a pipe write), while treating a poisoned lock as +/// a failure would cost `send_line` its delivery guarantee and leak the child's +/// stdin handle for the lifetime of the writer. +fn lock_fix_stdin_state(state: &Mutex) -> std::sync::MutexGuard<'_, FixStdinState> { + state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) +} + +impl FixStdinState { + /// Queue or write `line` — with the trailing newline the caller doesn't + /// supply — according to the current state. A failed write latches `Closed` + /// so later sends fail without re-discovering the dead pipe. + fn send_line(&mut self, line: String) -> Result<(), String> { + match self { + FixStdinState::Buffered { lines, .. } => { + lines.push(line); + Ok(()) + } + FixStdinState::Live(pipe) => { + use std::io::Write; + match pipe + .write_all(format!("{line}\n").as_bytes()) + .and_then(|()| pipe.flush()) + { + Ok(()) => Ok(()), + Err(e) => { + *self = FixStdinState::Closed; + Err(format!("{FIX_STDIN_CLOSED}: {e}")) + } + } + } + FixStdinState::Closed => Err(FIX_STDIN_CLOSED.to_string()), + } + } +} + +impl FixStdin { + /// Create a connected pair: a cloneable writer for the caller to keep and + /// the `FixStdin` to place in [`ExecuteFixOptions::stdin`]. Lines sent + /// before the fix subprocess spawns are queued and replayed once it does; + /// dropping every writer clone closes the child's stdin (EOF). + /// + /// Dropping the writers is the only way to say "no more input", and a fix + /// that reads *to EOF* rather than a fixed number of lines will not exit + /// until that happens — a host that leaves its input UI open pins the fix + /// until [`ExecuteFixOptions::timeout`] fires. Nothing else is at stake in + /// dropping them: the child's stdin handle lives with the fix and is + /// reclaimed when it ends, held writer or not. + pub fn pipe() -> (FixStdinWriter, FixStdin) { + let state = Arc::new(Mutex::new(FixStdinState::Buffered { + lines: Vec::new(), + eof: false, + claimed: false, + })); + ( + FixStdinWriter { + inner: Arc::new(FixStdinWriterInner { + state: state.clone(), + }), + }, + FixStdin { state }, + ) + } + + /// Reserve this pipe for a child about to be spawned. First caller wins; + /// `Err` on every later call (a clone already fed an execution), which the + /// caller surfaces instead of spawning a fix whose stdin is already dead. + fn claim(&self) -> Result<(), String> { + match &mut *lock_fix_stdin_state(&self.state) { + FixStdinState::Buffered { claimed, .. } if !*claimed => { + *claimed = true; + Ok(()) + } + _ => Err(FIX_STDIN_REUSED.to_string()), + } + } + + /// Hand the spawned child's stdin to the pipe, replay whatever the host + /// queued before the spawn, and go live. + /// + /// Only ever reached after a successful [`FixStdin::claim`], which is what + /// guarantees the state is still `Buffered`; any other state means another + /// execution owns the pipe, and dropping the handle — an immediate EOF for + /// this child — is the only safe reading of that. A replay write that fails + /// is not the fix's failure (a command is free to exit successfully without + /// reading its stdin), so it only latches `Closed`; the host hears about it + /// from its next `send_line`. + fn attach(&self, child_stdin: std::process::ChildStdin) { + let mut state = lock_fix_stdin_state(&self.state); + let FixStdinState::Buffered { lines, eof, .. } = &mut *state else { + return; + }; + let queued = std::mem::take(lines); + let eof = *eof; + *state = FixStdinState::Live(child_stdin); + for line in queued { + if state.send_line(line).is_err() { + break; + } + } + if eof { + // Every writer was dropped before the spawn, so the queued lines + // above are all the input there will ever be and closing now is the + // EOF the fix is waiting for. + *state = FixStdinState::Closed; + } + } + + /// The fix is over: close the pipe so every later send fails immediately. + /// A write hitting `EPIPE` cannot be the signal on its own — a backgrounded + /// grandchild that inherited the child's stdin keeps the read end open, and + /// writes into it go on succeeding long after the fix is gone. + fn close(&self) { + *lock_fix_stdin_state(&self.state) = FixStdinState::Closed; + } +} + +/// Cloneable handle for feeding lines to a fix subprocess's stdin. Dropping +/// every clone closes the fix's stdin (EOF). +#[derive(Debug, Clone)] +pub struct FixStdinWriter { + inner: Arc, +} + +/// Shared by every [`FixStdinWriter`] clone so EOF is delivered exactly when +/// the last one drops, which is what keeps the writer `Clone`. +#[derive(Debug)] +struct FixStdinWriterInner { + state: Arc>, +} + +impl Drop for FixStdinWriterInner { + fn drop(&mut self) { + match &mut *lock_fix_stdin_state(&self.state) { + // Pre-spawn the queued lines still have to reach the child first, so + // record the EOF for `attach` to deliver after the replay. + FixStdinState::Buffered { eof, .. } => *eof = true, + // Otherwise dropping the state's `ChildStdin` *is* the EOF. + state => *state = FixStdinState::Closed, + } + } +} + +impl FixStdinWriter { + /// Write one line to the fix's stdin; a trailing `\n` is appended and the + /// pipe is flushed. + /// + /// `Ok` means the bytes were handed to the child's stdin pipe — not that the + /// fix read them, since a fix can exit with bytes still buffered. `Err` + /// means the line was *not* delivered: the fix has finished, its stdin is + /// closed, or this pipe was never attached to a spawned fix. + /// + /// Lines sent before the fix spawns are queued and replayed at spawn, so + /// they return `Ok` before any pipe exists; if the fix never spawns they are + /// dropped. + /// + /// Completion is signalled by the fix's own `Result`, never by `send_line`. + /// May block if the fix isn't reading and the pipe buffer fills, so a host + /// sending anything bulkier than a pasted code should call this off its + /// async runtime. + pub fn send_line(&self, line: impl Into) -> Result<(), String> { + lock_fix_stdin_state(&self.inner.state).send_line(line.into()) + } +} + +/// Wall-clock bound on a single fix execution. +/// +/// Fixes are install/auth/update actions, so the bound has to clear a +/// cold-cache `npm install -g` behind a corporate proxy and a human doing SSO +/// in a browser — orders of magnitude above the probe timeouts in +/// [`crate::command`]. This is an enum rather than `Option` because +/// `None` reads as both "use the default" and "no timeout"; here every literal +/// has to say which it means, and `Unbounded` stays reachable for a caller +/// that genuinely wants the old forever-wait. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum FixTimeout { + /// [`DEFAULT_FIX_TIMEOUT`]. + #[default] + Standard, + /// A caller-chosen bound. + After(Duration), + /// No bound at all: the fix runs until it exits on its own. + Unbounded, +} + +impl FixTimeout { + /// The wall-clock bound, or `None` for [`FixTimeout::Unbounded`]. + fn duration(self) -> Option { + match self { + FixTimeout::Standard => Some(DEFAULT_FIX_TIMEOUT), + FixTimeout::After(duration) => Some(duration), + FixTimeout::Unbounded => None, + } + } +} + +/// Deadline applied by [`FixTimeout::Standard`]. Deliberately generous: it +/// exists to stop a wedged fix from pinning a blocking worker and a process +/// tree for the lifetime of the host, not to police slow-but-honest installs +/// or a leisurely browser login. +pub const DEFAULT_FIX_TIMEOUT: Duration = Duration::from_secs(600); + /// Options for executing a doctor fix command. #[derive(Debug, Clone, Default)] pub struct ExecuteFixOptions { @@ -682,6 +930,16 @@ pub struct ExecuteFixOptions { pub npm_registry: Option, /// Optional caller-provided environment snapshot for the fix subprocess. pub env: Option, + /// Opt-in piped stdin for the fix subprocess (see [`FixStdin::pipe`]). + /// `None` keeps the child inheriting the host process's stdin, so + /// terminal hosts can still run interactive fixes directly. + /// + /// A `FixStdin` feeds exactly one execution, so a cached options struct + /// must have this field refreshed (or be rebuilt) before a fix is retried; + /// reusing it fails the run. + pub stdin: Option, + /// Wall-clock bound on the fix. Defaults to [`FixTimeout::Standard`]. + pub timeout: FixTimeout, } impl ExecuteFixOptions { @@ -689,6 +947,20 @@ impl ExecuteFixOptions { self.env = Some(DoctorEnv::new(vars)); self } + + /// Attach an opt-in stdin pipe (see [`FixStdin::pipe`]). The `FixStdin` + /// feeds exactly one execution: call this again with a fresh pipe for + /// every retry rather than reusing a built options struct. + pub fn with_stdin(mut self, stdin: FixStdin) -> Self { + self.stdin = Some(stdin); + self + } + + /// Override the wall-clock bound on the fix (see [`FixTimeout`]). + pub fn with_timeout(mut self, timeout: FixTimeout) -> Self { + self.timeout = timeout; + self + } } /// Run a fix command for a doctor check, identified by check ID and fix type. @@ -722,7 +994,7 @@ pub async fn execute_fix_with_options( ExecuteFixOptions { command_override, npm_registry: npm_registry.map(str::to_string), - env: None, + ..Default::default() }, ) .await @@ -779,7 +1051,7 @@ where ExecuteFixOptions { command_override, npm_registry: npm_registry.map(str::to_string), - env: None, + ..Default::default() }, on_line, ) @@ -813,21 +1085,24 @@ where // Fixes are intentionally not routed through the bounded probe runner: // these are user-triggered install/auth/update actions and can reasonably - // be interactive or long-running. - run_command_streaming(command, opts.env, on_line).await + // be interactive or long-running, so they get the far more generous + // `FixTimeout` bound instead of a probe timeout. + run_command_streaming(command, opts.env, opts.stdin, opts.timeout, on_line).await } /// Async wrapper that runs `run_command_streaming_blocking` on the blocking pool. pub(crate) async fn run_command_streaming( command: String, env: Option, + stdin: Option, + timeout: FixTimeout, on_line: F, ) -> Result<(), String> where F: FnMut(&str) + Send + 'static, { tokio::task::spawn_blocking(move || { - run_command_streaming_blocking(&command, env.as_ref(), on_line) + run_command_streaming_blocking(&command, env.as_ref(), stdin, timeout, on_line) }) .await .unwrap_or_else(|e| Err(format!("Task failed: {e}"))) @@ -1011,31 +1286,98 @@ pub(crate) fn execute_command_with_path_prefix_with_env( } } +/// Closes the fix's stdin pipe when `run_command_streaming_blocking` leaves its +/// body — normal return, error return, timeout, spawn failure, or a panic in +/// `on_line`. Every path has to close it: a host that still holds a +/// [`FixStdinWriter`] would otherwise keep getting `Ok` from `send_line` for a +/// fix that is already over, and the child's stdin handle would live as long as +/// that writer. +struct FixStdinCloser<'a>(&'a FixStdin); + +impl Drop for FixStdinCloser<'_> { + fn drop(&mut self) { + self.0.close(); + } +} + /// Spawn `command` through a login shell, stream stdout/stderr lines to -/// `on_line`, and return based on the process exit status. This path is -/// deliberately unbounded: fix commands are user-triggered install/auth/update -/// actions and may prompt or run package managers. Stderr lines are also -/// accumulated so a non-zero exit can surface a useful error message (matching -/// the non-streaming behavior of the previous `execute_command`). +/// `on_line`, and return based on the process exit status. Bounded by +/// `timeout`, which is generous rather than tight: fix commands are +/// user-triggered install/auth/update actions and may prompt or run package +/// managers. Stderr lines are also accumulated so a non-zero exit can surface a +/// useful error message (matching the non-streaming behavior of the previous +/// `execute_command`). fn run_command_streaming_blocking( command: &str, env: Option<&DoctorEnv>, + stdin: Option, + timeout: FixTimeout, mut on_line: F, ) -> Result<(), String> where F: FnMut(&str), { use std::io::{BufRead, BufReader}; + use std::sync::mpsc::RecvTimeoutError; + + use wait_timeout::ChildExt; + + fn consume(msg: StreamLine, on_line: &mut F, stderr_accum: &mut String) { + match msg { + StreamLine::Stdout(s) => { + on_line(&s); + } + StreamLine::Stderr(s) => { + on_line(&s); + if !stderr_accum.is_empty() { + stderr_accum.push('\n'); + } + stderr_accum.push_str(&s); + } + } + } + + // Claim the pipe before anything is launched: a `FixStdin` another execution + // already consumed can never deliver a line, so the child would block + // forever on a pipe nobody writes — the exact hang this option exists to + // fix. Always a caller bug, so surface it at the call site rather than + // spawning a doomed subprocess. + if let Some(fix_stdin) = &stdin { + fix_stdin.claim()?; + } - let mut command = build_shell_command(command, &[], env); - command + let mut shell_command = build_shell_command(command, &[], env); + shell_command .stdout(std::process::Stdio::piped()) .stderr(std::process::Stdio::piped()); - command::configure_command(&mut command); - let mut child = command + // Opt-in only: without a `FixStdin` the child keeps inheriting the host + // process's stdin, so interactive fixes in terminal hosts are untouched. + if stdin.is_some() { + shell_command.stdin(std::process::Stdio::piped()); + // Own the whole tree so a timeout can kill more than the login shell: + // `kill(-pid)` only reaches an `npm install` under `zsh -lc` if the + // shell leads its own group. Gated on piped stdin because a child in + // its own group that reads the controlling terminal gets SIGTTIN and + // stops — impossible here precisely because doctor owns its stdin, but + // a real regression for a terminal host on the inherited-stdin path. + #[cfg(unix)] + { + use std::os::unix::process::CommandExt; + shell_command.process_group(0); + } + } + command::configure_command(&mut shell_command); + + // Declared ahead of the spawn so a spawn failure closes the pipe too: the + // claim above is already spent, so the host must not keep getting `Ok` for a + // fix that never started. + let _stdin_closer = stdin.as_ref().map(FixStdinCloser); + + let mut child = shell_command .spawn() .map_err(|e| format!("Failed to run command: {e}"))?; + let child_stdin = child.stdin.take(); let stdout = child.stdout.take().expect("stdout was piped"); let stderr = child.stderr.take().expect("stderr was piped"); @@ -1059,28 +1401,82 @@ where } }); + // Deliberately after the readers are running: the replay of pre-spawn lines + // writes inline on this thread, so a queue larger than the pipe buffer would + // deadlock against a child whose output nobody is draining yet. + if let (Some(fix_stdin), Some(child_stdin)) = (&stdin, child_stdin) { + fix_stdin.attach(child_stdin); + } + + let limit = timeout.duration(); + let deadline = limit.map(|limit| Instant::now() + limit); let mut stderr_accum = String::new(); - for msg in rx.iter() { - match msg { - StreamLine::Stdout(s) => { - on_line(&s); - } - StreamLine::Stderr(s) => { - on_line(&s); - if !stderr_accum.is_empty() { - stderr_accum.push('\n'); + let mut expired = false; + + loop { + let msg = match deadline { + Some(deadline) => { + match rx.recv_timeout(deadline.saturating_duration_since(Instant::now())) { + Ok(msg) => msg, + Err(RecvTimeoutError::Timeout) => { + expired = true; + break; + } + Err(RecvTimeoutError::Disconnected) => break, } - stderr_accum.push_str(&s); } - } + // `recv_timeout(Duration::MAX)` overflows instantly, so the + // unbounded case keeps the plain blocking receive. + None => match rx.recv() { + Ok(msg) => msg, + Err(_) => break, + }, + }; + consume(msg, &mut on_line, &mut stderr_accum); } - let _ = stdout_thread.join(); - let _ = stderr_thread.join(); + let status = if expired { + None + } else { + // Both pipes hit EOF, so the readers are already done and joining is + // immediate. The process can still outlive its pipes, though, so the + // reap is bounded by the same deadline. + let _ = stdout_thread.join(); + let _ = stderr_thread.join(); + match deadline { + Some(deadline) => child + .wait_timeout(deadline.saturating_duration_since(Instant::now())) + .map_err(|e| format!("Failed to wait for command: {e}"))?, + None => Some( + child + .wait() + .map_err(|e| format!("Failed to wait for command: {e}"))?, + ), + } + }; - let status = child - .wait() - .map_err(|e| format!("Failed to wait for command: {e}"))?; + let Some(status) = status else { + let limit = limit.expect("a deadline only exists when the fix is bounded"); + // Anything the readers already queued is real output the user should + // see before the notice explaining why it stopped. + while let Ok(msg) = rx.try_recv() { + consume(msg, &mut on_line, &mut stderr_accum); + } + on_line(&format!( + "doctor: fix timed out after {} — terminating", + format_duration(limit) + )); + command::kill_child_process_group_or_child(&mut child); + let _ = child.wait(); + // The reader threads are deliberately not joined: a descendant that + // escaped the process group can hold the inherited stdout open long + // after the fix is dead, and waiting on that is the hang this timeout + // exists to end. Dropping `rx` retires them at their next send. + return Err(format!( + "Fix timed out after {} without finishing: {command}", + format_duration(limit) + )); + }; if status.success() { Ok(()) @@ -1100,7 +1496,7 @@ mod tests { use std::path::Path; use std::sync::{Arc, Mutex}; - use std::time::Duration; + use std::time::{Duration, Instant}; fn timeout(label: &str, command: &str) -> CommandTimeout { CommandTimeout::new(label, command, Duration::from_secs(15)) @@ -1198,6 +1594,8 @@ mod tests { let result = run_command_streaming( "echo doctor-streaming-marker-hello && echo doctor-streaming-marker-world".to_string(), None, + None, + FixTimeout::Standard, move |line| { lines_clone.lock().unwrap().push(line.to_string()); }, @@ -1220,6 +1618,295 @@ mod tests { ); } + /// A line sent through the `FixStdin` pipe must reach the child's stdin + /// and dropping the last writer must deliver EOF: `cat` echoes the line + /// and exits 0 only when its stdin closes. Sending before the child + /// spawns also exercises the pre-spawn buffering guarantee. + #[tokio::test] + async fn run_command_streaming_piped_stdin_round_trips_through_cat() { + let lines: Arc>> = Arc::new(Mutex::new(Vec::new())); + let lines_clone = lines.clone(); + let (writer, stdin) = FixStdin::pipe(); + + writer.send_line("doctor-stdin-marker-echo").unwrap(); + drop(writer); + + let result = run_command_streaming( + "cat".to_string(), + None, + Some(stdin), + FixTimeout::Standard, + move |line| { + lines_clone.lock().unwrap().push(line.to_string()); + }, + ) + .await; + + assert!(result.is_ok(), "cat should exit 0 on EOF; got {result:?}"); + let captured = lines.lock().unwrap().clone(); + assert!( + captured.iter().any(|l| l == "doctor-stdin-marker-echo"), + "cat should echo the line written to its piped stdin; captured: {captured:?}", + ); + } + + /// The paste-an-auth-code shape: the command prompts by blocking on a line + /// read, and the caller feeds the answer through the writer while the fix is + /// running. Sending from inside `on_line` — on the fix's own thread, in + /// response to the prompt the fix printed — pins the send to a moment when + /// the pipe is provably live, so the `Ok` asserted here is the delivery + /// guarantee and not the pre-spawn queueing one. + #[tokio::test] + async fn run_command_streaming_piped_stdin_feeds_prompt_style_read() { + let lines: Arc>> = Arc::new(Mutex::new(Vec::new())); + let lines_clone = lines.clone(); + let live_send: Arc>>> = Arc::new(Mutex::new(None)); + let live_send_clone = live_send.clone(); + let (writer, stdin) = FixStdin::pipe(); + + let result = run_command_streaming( + "echo doctor-stdin-prompt; read -r line && echo \"got-$line\"".to_string(), + None, + Some(stdin), + FixTimeout::Standard, + move |line| { + lines_clone.lock().unwrap().push(line.to_string()); + if line == "doctor-stdin-prompt" { + *live_send_clone.lock().unwrap() = + Some(writer.send_line("doctor-stdin-auth-code")); + } + }, + ) + .await; + + assert!(result.is_ok(), "read/echo should exit 0; got {result:?}"); + let captured = lines.lock().unwrap().clone(); + let sent = live_send + .lock() + .unwrap() + .take() + .expect("the fix's prompt line should have reached on_line"); + assert!( + sent.is_ok(), + "a send while the fix is live should report delivery; got {sent:?}", + ); + assert!( + captured.iter().any(|l| l == "got-doctor-stdin-auth-code"), + "prompt-style read should see the sent line; captured: {captured:?}", + ); + } + + /// A writer held across the fix's completion must not hang the run, and the + /// *first* send after it must fail: the runner closes the pipe as it returns, + /// so `Ok` never means "queued for a fix that is already over". That is the + /// berd#99 shape — the login subprocess dies, the user pastes the auth code + /// a beat later — and a host keying off `Ok` would otherwise wait forever + /// with nothing in the log to explain it. + #[tokio::test] + async fn run_command_streaming_piped_stdin_rejects_sends_once_the_fix_finishes() { + let (writer, stdin) = FixStdin::pipe(); + + let result = run_command_streaming( + "echo doctor-stdin-done".to_string(), + None, + Some(stdin), + FixTimeout::Standard, + |_| {}, + ) + .await; + + assert!(result.is_ok(), "echo fix should complete; got {result:?}"); + let err = writer + .send_line("late-line") + .expect_err("the first send after the fix finished should fail"); + assert!( + err.contains("no longer accepting input"), + "error should say the input is closed; got {err:?}", + ); + } + + /// `EPIPE` alone can't carry "the fix is over": a backgrounded grandchild + /// inherits the child's stdin and keeps the read end open, so a write into a + /// finished fix's pipe still succeeds. Only the runner's explicit close on + /// the way out makes this send fail. The grandchild's stdout is redirected so + /// it doesn't also hold the reader threads open — this test is about stdin. + #[cfg(unix)] + #[tokio::test] + async fn run_command_streaming_piped_stdin_rejects_sends_when_a_grandchild_holds_the_pipe() { + let (writer, stdin) = FixStdin::pipe(); + + let result = run_command_streaming( + "sleep 2 >/dev/null 2>&1 & echo doctor-stdin-done".to_string(), + None, + Some(stdin), + FixTimeout::Standard, + |_| {}, + ) + .await; + + assert!(result.is_ok(), "echo fix should complete; got {result:?}"); + assert!( + writer.send_line("late-line").is_err(), + "a grandchild holding the read end must not make a dead fix look writable", + ); + } + + /// Reusing a `FixStdin` (or a clone) for a second execution must fail + /// loudly rather than hand the child an immediately-EOF'd stdin — the + /// receiver lives with the first run, so a second could only hang. The + /// second run must also never spawn: nothing reaches `on_line`. + #[tokio::test] + async fn run_command_streaming_piped_stdin_errors_when_reused() { + let (writer, stdin) = FixStdin::pipe(); + let reused = stdin.clone(); + writer.send_line("doctor-stdin-reuse-first").unwrap(); + drop(writer); + + let first = run_command_streaming( + "cat".to_string(), + None, + Some(stdin), + FixTimeout::Standard, + |_| {}, + ) + .await; + assert!(first.is_ok(), "first run should succeed; got {first:?}"); + + let lines: Arc>> = Arc::new(Mutex::new(Vec::new())); + let lines_clone = lines.clone(); + let second = run_command_streaming( + "echo doctor-stdin-reuse-second".to_string(), + None, + Some(reused), + FixTimeout::Standard, + move |line| lines_clone.lock().unwrap().push(line.to_string()), + ) + .await; + + let err = second.expect_err("reusing a consumed FixStdin should fail"); + let captured = lines.lock().unwrap().clone(); + assert!( + err.contains("already consumed"), + "error should name the reuse; got {err:?}", + ); + assert!( + captured.is_empty(), + "second run must not spawn; captured: {captured:?}", + ); + } + + /// The default bound must stay at fix scale, not probe scale. A fix is an + /// `npm install -g` behind a corporate proxy or a human doing SSO in a + /// browser; retuning this toward `DEFAULT_PROBE_TIMEOUT` would kill honest + /// work mid-flight. + #[test] + fn default_fix_timeout_stays_at_fix_scale() { + assert_eq!(DEFAULT_FIX_TIMEOUT, Duration::from_secs(600)); + assert_eq!(ExecuteFixOptions::default().timeout, FixTimeout::Standard); + assert_eq!(FixTimeout::Standard.duration(), Some(DEFAULT_FIX_TIMEOUT)); + assert_eq!(FixTimeout::Unbounded.duration(), None); + assert!( + DEFAULT_FIX_TIMEOUT >= DEFAULT_PROBE_TIMEOUT * 30, + "fix timeout must stay far above probe scale", + ); + } + + /// A fix that never finishes must return on its deadline instead of + /// pinning the blocking worker forever — the whole point of the bound. + #[tokio::test] + async fn run_command_streaming_returns_when_the_fix_outlives_its_timeout() { + let lines: Arc>> = Arc::new(Mutex::new(Vec::new())); + let lines_clone = lines.clone(); + let started = Instant::now(); + + let result = run_command_streaming( + "sleep 60".to_string(), + None, + None, + FixTimeout::After(Duration::from_millis(100)), + move |line| lines_clone.lock().unwrap().push(line.to_string()), + ) + .await; + + let err = result.expect_err("a fix past its deadline should fail"); + assert!( + err.contains("timed out") && err.contains("sleep 60"), + "error should name the timeout and the command; got {err:?}", + ); + assert!( + started.elapsed() < Duration::from_secs(2), + "timeout path waited for the fix instead of its deadline", + ); + let captured = lines.lock().unwrap().clone(); + assert!( + captured + .iter() + .any(|l| l.starts_with("doctor: fix timed out")), + "callers should see a notice line explaining the stop; captured: {captured:?}", + ); + } + + /// With piped stdin the shell leads its own process group, so the timeout + /// kill must take the whole tree — not just the login shell, leaving a + /// backgrounded installer running. + #[cfg(unix)] + #[tokio::test] + async fn run_command_streaming_timeout_kills_the_whole_process_tree() { + let tmp = unique_tmp_dir("fix-timeout-tree"); + let marker = tmp.join("grandchild-ran"); + let (_writer, stdin) = FixStdin::pipe(); + + let result = run_command_streaming( + format!("(sleep 2; touch {}) & sleep 60", marker.display()), + None, + Some(stdin), + FixTimeout::After(Duration::from_millis(300)), + |_| {}, + ) + .await; + + assert!(result.is_err(), "timed-out fix should fail; got {result:?}"); + // Past when the backgrounded grandchild would have written its marker + // had it survived the group kill. + tokio::time::sleep(Duration::from_secs(3)).await; + let survived = marker.exists(); + let _ = std::fs::remove_dir_all(&tmp); + assert!( + !survived, + "backgrounded grandchild outlived the timeout kill", + ); + } + + /// A descendant that escaped the process group keeps the inherited + /// stdout/stderr open, so the reader threads never see EOF. The timeout + /// path must not join them — it must return on the deadline regardless + /// (the streaming twin of `command_runner_returns_when_escaped_descendant_ + /// keeps_pipes_open`). + #[cfg(unix)] + #[tokio::test] + async fn run_command_streaming_timeout_returns_when_escaped_descendant_keeps_pipes_open() { + let started = Instant::now(); + + let result = run_command_streaming( + "perl -MPOSIX=setsid -e 'setsid(); sleep 5' & wait".to_string(), + None, + None, + FixTimeout::After(Duration::from_millis(250)), + |_| {}, + ) + .await; + + let err = result.expect_err("a fix past its deadline should fail"); + assert!( + err.contains("timed out"), + "error should name the timeout; got {err:?}", + ); + assert!( + started.elapsed() < Duration::from_secs(2), + "timeout path waited for the escaped descendant to close the pipes", + ); + } + /// `execute_fix(|_| {})` and `execute_fix_streaming(.., |_| {})` must /// produce identical results for the same fix lookup — `execute_fix` is /// supposed to be a thin delegate. @@ -1620,8 +2307,8 @@ mod tests { FixType::UpdateMain, ExecuteFixOptions { command_override: Some(script_name.to_string()), - npm_registry: None, env: Some(env), + ..Default::default() }, move |line| { lines_clone.lock().unwrap().push(line.to_string()); @@ -1674,8 +2361,8 @@ mod tests { FixType::UpdateMain, ExecuteFixOptions { command_override: Some(command.to_string()), - npm_registry: None, env: Some(env), + ..Default::default() }, move |line| { lines_clone.lock().unwrap().push(line.to_string());