acp: Fix startup error race condition test flake (#56654)

Self-Review Checklist:

- [x] I've reviewed my own diff for quality, security, and reliability
- [x] Unsafe blocks (if any) have justifying comments
- [x] The content is consistent with the [UI/UX
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
- [ ] Tests cover the new/changed behavior
- [x] Performance impact has been considered and is acceptable

Closes #ISSUE

Release Notes:

- N/A or Added/Fixed/Improved ...

crates/agent_servers/src/acp.rs

@@ -23,6 +23,7 @@ use std::path::PathBuf;
 use std::process::{ExitStatus, Stdio};
 use std::rc::Rc;
 use std::sync::{Arc, Mutex};
+use std::time::Duration;
 use std::{any::Any, cell::RefCell, collections::VecDeque};
 use task::{Shell, ShellBuilder, SpawnInTerminal};
 use thiserror::Error;
@@ -41,6 +42,8 @@ use crate::GEMINI_ID;
 
 pub const GEMINI_TERMINAL_AUTH_METHOD_ID: &str = "spawn-gemini-cli";
 const MAX_DEBUG_BACKLOG_MESSAGES: usize = 2000;
+const ACP_RESPONSE_CHANNEL_CANCELLED: &str =
+    "response channel cancelled — connection may have dropped";
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum AcpDebugMessageDirection {
@@ -249,10 +252,8 @@ fn into_foreground_future<T: JsonRpcResponse>(
     });
     async move {
         spawn_result?;
-        rx.await.map_err(|_| {
+        rx.await
-            acp::Error::internal_error()
+            .map_err(|_| acp::Error::internal_error().data(ACP_RESPONSE_CHANNEL_CANCELLED))?
-                .data("response channel cancelled — connection may have dropped")
-        })?
     }
 }
 
@@ -821,17 +822,23 @@ impl AcpConnection {
                 .context("Failed to receive ACP connection handle")
         }
         .boxed_local();
-        let status_fut = child.status().boxed_local();
+        let status_fut = child
+            .status()
+            .map({
+                let debug_log = debug_log.clone();
+                move |status| match status {
+                    Ok(status) => Ok(exited_load_error_with_stderr(status, &debug_log)),
+                    Err(err) => Err(anyhow!("failed to wait for agent server exit: {err}")),
+                }
+            })
+            .boxed_local();
         let (connection, status_fut) = match futures::future::select(connection_rx, status_fut)
             .await
         {
             futures::future::Either::Left((connection, status_fut)) => (connection?, status_fut),
-            futures::future::Either::Right((status, _connection_rx)) => match status {
+            futures::future::Either::Right((load_error, _connection_rx)) => {
-                Ok(status) => return Err(exited_load_error_with_stderr(status, &debug_log).into()),
+                return Err(load_error?.into());
-                Err(err) => {
+            }
-                    return Err(anyhow!("agent server exited before initialization: {err}"));
-                }
-            },
         };
 
         // Set up the foreground dispatch loop to process work items from handlers.
@@ -869,19 +876,34 @@ impl AcpConnection {
                     ),
             ),
         )
-        .map(|response| response.map_err(anyhow::Error::from))
         .boxed_local();
-        let (response, status_fut) = match futures::future::select(initialize_response, status_fut)
+        let (response, status_fut) =
-            .await
+            match futures::future::select(initialize_response, status_fut).await {
-        {
+                futures::future::Either::Left((Ok(response), status_fut)) => (response, status_fut),
-            futures::future::Either::Left((response, status_fut)) => (response?, status_fut),
+                futures::future::Either::Left((Err(error), status_fut)) => {
-            futures::future::Either::Right((status, _initialize_response)) => match status {
+                    let response_channel_cancelled = error.code == ErrorCode::InternalError
-                Ok(status) => return Err(exited_load_error_with_stderr(status, &debug_log).into()),
+                        && error.data.as_ref().and_then(|data| data.as_str())
-                Err(err) => {
+                            == Some(ACP_RESPONSE_CHANNEL_CANCELLED);
-                    return Err(anyhow!("agent server exited before initialization: {err}"));
+                    if !response_channel_cancelled {
+                        return Err(error.into());
+                    }
+
+                    let timer = cx
+                        .background_executor()
+                        .timer(Duration::from_millis(250))
+                        .boxed_local();
+                    if let futures::future::Either::Left((load_error, _timer)) =
+                        futures::future::select(status_fut, timer).await
+                    {
+                        return Err(load_error?.into());
+                    }
+
+                    return Err(error.into());
                 }
-            },
+                futures::future::Either::Right((load_error, _initialize_response)) => {
-        };
+                    return Err(load_error?.into());
+                }
+            };
 
         if response.protocol_version < MINIMUM_SUPPORTED_VERSION {
             return Err(UnsupportedVersion.into());
@@ -889,14 +911,9 @@ impl AcpConnection {
 
         let wait_task = cx.spawn({
             let sessions = sessions.clone();
-            let debug_log = debug_log.clone();
             async move |cx| {
-                let status = status_fut.await?;
+                let load_error = status_fut.await?;
-                emit_load_error_to_all_sessions(
+                emit_load_error_to_all_sessions(&sessions, load_error, cx);
-                    &sessions,
-                    exited_load_error_with_stderr(status, &debug_log),
-                    cx,
-                );
                 anyhow::Ok(())
             }
         });