Clean up usage.

Author: milesj

CHANGELOG.md

@@ -33,8 +33,8 @@
   - Connecting to the daemon and starting it are now a single operation, so a command that needs the
     daemon will start one itself if the background pre-warm hasn't yet, instead of silently running
     without it. Concurrent starts still coordinate so only one daemon is spawned.
-  - When connecting, the client now checks the running daemon's moon and protocol version against its
-    own and, on a mismatch, restarts it — so a daemon left over from before a `moon upgrade` is
+  - When connecting, the client now checks the running daemon's moon and protocol version against
+    its own and, on a mismatch, restarts it — so a daemon left over from before a `moon upgrade` is
     replaced instead of serving the old binary indefinitely.
 - **Processes**
   - Improved our "stream and capture output" child process handling to operate on bytes instead of

crates/app/src/session.rs

@@ -53,7 +53,7 @@ pub struct MoonSession {
 
     // Lazy components
     pub(crate) cache_engine: OnceLock<Arc<CacheEngine>>,
-    // pub(crate) daemon_client: OnceCell<Option<DaemonClient>>,
+    pub(crate) daemon_client: OnceLock<DaemonClient>,
     pub(crate) extension_registry: OnceCell<Arc<ExtensionRegistry>>,
     pub(crate) project_graph: OnceLock<Arc<ProjectGraph>>,
     pub(crate) task_graph: OnceLock<Arc<TaskGraph>>,
@@ -83,7 +83,7 @@ impl MoonSession {
             config_dir: PathBuf::new(),
             config_loader: ConfigLoader::default(),
             console: Console::new(cli.quiet || is_formatted_output()),
-            // daemon_client: OnceCell::new(),
+            daemon_client: OnceLock::new(),
             extensions_config: Arc::new(ExtensionsConfig::default()),
             extension_registry: OnceCell::new(),
             moon_env: Arc::new(MoonEnvironment::default()),
@@ -139,16 +139,19 @@ impl MoonSession {
             return Ok(None);
         }
 
-        // let client = self
-        //     .daemon_client
-        //     .get_or_try_init(async move || self.get_daemon_connector()?.acquire().await)
-        //     .await?;
+        if let Some(client) = self.daemon_client.get() {
+            return Ok(Some(client.to_owned()));
+        }
+
+        let client = self.get_daemon_connector()?.acquire().await?;
 
-        // Ok(client.clone())
+        // Only cache the client if we successfully connected to a daemon.
+        // If we failed to connect, we don't want to cache so that we try again.
+        if let Some(client) = &client {
+            let _ = self.daemon_client.set(client.to_owned());
+        }
 
-        // `acquire` connects to a running daemon or starts one, degrading to
-        // `None` on failure, so there's nothing to handle here.
-        self.get_daemon_connector()?.acquire().await
+        Ok(client)
     }
 
     pub async fn create_workspace_graph_context(&self) -> miette::Result<WorkspaceBuilderContext> {
@@ -476,7 +479,9 @@ impl AppSession for MoonSession {
         // `connect_to_daemon`, so the two coordinate and at most one spawns —
         // and a spawn failure degrades to `None` instead of failing the run.
         if self.is_daemon_allowed() {
-            self.get_daemon_connector()?.acquire().await?;
+            if let Some(client) = self.get_daemon_connector()?.acquire().await? {
+                let _ = self.daemon_client.set(client);
+            }
         }
 
         Ok(None)

crates/daemon-client/src/daemon_client.rs

@@ -103,14 +103,26 @@ impl DaemonClient {
         debug!(endpoint = endpoint, "Connecting to daemon");
 
         let channel = match timeout(CONNECT_TIMEOUT, connect_channel(&endpoint)).await {
-            Ok(Ok(channel)) => channel,
+            Ok(Ok(channel)) => {
+                debug!(endpoint = endpoint, "Connected to daemon");
+
+                channel
+            }
             Ok(Err(error)) => {
+                debug!(
+                    endpoint = endpoint,
+                    error = error.to_string(),
+                    "Failed to connect to daemon"
+                );
+
                 return Err(DaemonClientError::ConnectFailed {
                     endpoint,
                     error: Box::new(error),
                 });
             }
             Err(_) => {
+                debug!(endpoint = endpoint, "Timed out connecting to daemon");
+
                 return Err(DaemonClientError::ConnectTimedOut {
                     endpoint,
                     timeout_secs: CONNECT_TIMEOUT.as_secs(),
@@ -147,6 +159,8 @@ impl DaemonClient {
     /// A failure to read status keeps the daemon rather than churning on a
     /// transient error — a real RPC will surface any genuine problem.
     pub async fn handshake(&mut self, client_version: &str) -> HandshakeOutcome {
+        debug!("Initiating handshake with daemon");
+
         match self.status().await {
             Ok(status) => {
                 let outcome = if status.protocol_version == PROTOCOL_VERSION
@@ -159,9 +173,9 @@ impl DaemonClient {
 
                 if outcome == HandshakeOutcome::Restart {
                     debug!(
-                        daemon_version = status.moon_version,
+                        server_version = status.moon_version,
                         client_version,
-                        daemon_protocol = status.protocol_version,
+                        server_protocol = status.protocol_version,
                         client_protocol = PROTOCOL_VERSION,
                         "Daemon version handshake mismatch"
                     );