fix(desktop): stabilize agent identity restore (#1831)

desktop/src-tauri/src/app_state.rs

@@ -21,6 +21,15 @@ pub struct AppState {
     /// Workspace-provided relay URL override. Set by `apply_workspace` on app
     /// init and takes priority over env vars and compile-time defaults.
     pub relay_url_override: Mutex<Option<String>>,
+    /// Set during backend setup when managed agents are eligible for launch
+    /// restore. `apply_workspace` consumes it after installing the workspace
+    /// relay and identity, so agents never start against the fallback relay.
+    pub managed_agent_restore_pending: AtomicBool,
+    /// Shared shutdown signal checked by launch-time agent restoration.
+    pub shutdown_started: AtomicBool,
+    /// Serializes the restore spawn/register transition with shutdown cleanup,
+    /// preventing an agent from spawning after shutdown has swept processes.
+    pub managed_agent_restore_transition: Mutex<()>,
     pub managed_agents_store_lock: Mutex<()>,
     pub channel_templates_store_lock: Mutex<()>,
     pub managed_agent_processes: Mutex<HashMap<String, ManagedAgentProcess>>,
@@ -143,6 +152,9 @@ pub fn build_app_state() -> AppState {
             .build()
             .unwrap_or_else(|_| reqwest::Client::new()),
         relay_url_override: Mutex::new(None),
+        managed_agent_restore_pending: AtomicBool::new(false),
+        shutdown_started: AtomicBool::new(false),
+        managed_agent_restore_transition: Mutex::new(()),
         identity_mutation: Mutex::new(()),
         managed_agents_store_lock: Mutex::new(()),
         channel_templates_store_lock: Mutex::new(()),

desktop/src-tauri/src/commands/workspace.rs

@@ -1,11 +1,12 @@
 use nostr::Keys;
 use serde::{Deserialize, Serialize};
+use std::sync::atomic::Ordering;
 use tauri::{AppHandle, Emitter, Manager, State};
 
 use crate::app_state::AppState;
 use crate::managed_agents::{
-    effective_repos_dir, ensure_repos_symlink, nest_dir, try_regenerate_nest,
-    write_persisted_repos_dir,
+    effective_repos_dir, ensure_repos_symlink, nest_dir, restore_managed_agents_on_launch,
+    try_regenerate_nest, write_persisted_repos_dir,
 };
 use crate::relay;
 
@@ -104,6 +105,7 @@ pub async fn apply_workspace(
     repos_dir: Option<String>,
     app: AppHandle,
 ) -> Result<(), String> {
+    let restore_app = app.clone();
     tokio::task::spawn_blocking(move || {
         let state = app.state::<AppState>();
 
@@ -168,8 +170,26 @@ pub async fn apply_workspace(
 
         try_regenerate_nest(&app);
 
-        Ok(())
+        Ok::<(), String>(())
     })
     .await
-    .map_err(|e| format!("spawn_blocking failed: {e}"))?
+    .map_err(|e| format!("spawn_blocking failed: {e}"))??;
+
+    let state = restore_app.state::<AppState>();
+    if state
+        .managed_agent_restore_pending
+        .swap(false, Ordering::AcqRel)
+    {
+        let app = restore_app.clone();
+        tauri::async_runtime::spawn(async move {
+            let state = app.state::<AppState>();
+            if let Err(error) =
+                restore_managed_agents_on_launch(&app, &state.shutdown_started).await
+            {
+                eprintln!("buzz-desktop: failed to restore managed agents: {error}");
+            }
+        });
+    }
+
+    Ok(())
 }

desktop/src-tauri/src/lib.rs

@@ -40,9 +40,7 @@ use huddle::{
     join_huddle, leave_huddle, push_audio_pcm, set_huddle_transcription_enabled, set_tts_enabled,
     set_voice_input_mode, speak_agent_message, start_huddle, start_stt_pipeline,
 };
-use managed_agents::{
-    backfill_persona_snapshots, ensure_nest, restore_managed_agents_on_launch, try_regenerate_nest,
-};
+use managed_agents::{backfill_persona_snapshots, ensure_nest, try_regenerate_nest};
 use shutdown::shutdown_managed_agents;
 use std::sync::{
     atomic::{AtomicBool, Ordering},
@@ -404,8 +402,6 @@ pub fn run() {
     #[cfg(not(buzz_updater_enabled))]
     let builder = builder;
 
-    let shutdown_started = Arc::new(AtomicBool::new(false));
-    let restore_shutdown_started = Arc::clone(&shutdown_started);
     let app = builder
         .register_asynchronous_uri_scheme_protocol("buzz-media", |ctx, request, responder| {
             let app = ctx.app_handle().clone();
@@ -418,7 +414,6 @@ pub fn run() {
         .manage(commands::pairing::PairingHandle::new())
         .setup(move |app| {
             let app_handle = app.handle().clone();
-            let shutdown_started = Arc::clone(&restore_shutdown_started);
 
             // Run all pre-identity data migrations before state loads from disk.
             migration::run_boot_migrations(&app_handle);
@@ -577,22 +572,16 @@ pub fn run() {
                 });
             }
 
-            // Keep launch-time agent restoration off the synchronous setup path
-            // so the frontend can mount and reveal the window promptly. Gated on
-            // the boot-time repos symlink result (see restore_agents above):
-            // skip when a configured repos_dir could not be resolved, so no
-            // agent clones into a REPOS that isn't the user's target.
-            // Also skipped in recovery mode — agents must not be spawned
-            // under an ephemeral owner key.
+            // Defer launch-time agent restoration until `apply_workspace` has
+            // installed the active workspace relay and identity. Starting here
+            // would race React initialization and send agents whose saved record
+            // has no relay override to the localhost fallback. Preserve the
+            // boot-time repos and identity recovery safety gates by only marking
+            // restoration pending when both allow it.
             if restore_agents && !recovery_mode {
-                tauri::async_runtime::spawn(async move {
-                    if let Err(error) =
-                        restore_managed_agents_on_launch(&app_handle, shutdown_started.as_ref())
-                            .await
-                    {
-                        eprintln!("buzz-desktop: failed to restore managed agents: {error}");
-                    }
-                });
+                state
+                    .managed_agent_restore_pending
+                    .store(true, Ordering::Release);
             }
 
             // Periodic sweep: reap orphaned agents from dead instances every 60s.
@@ -900,9 +889,11 @@ pub fn run() {
     {
         let signal_app = app.handle().clone();
         let signal_shutdown_done = Arc::clone(&shutdown_done);
-        let signal_shutdown_started = Arc::clone(&shutdown_started);
         if let Err(e) = ctrlc::set_handler(move || {
-            signal_shutdown_started.store(true, Ordering::SeqCst);
+            signal_app
+                .state::<AppState>()
+                .shutdown_started
+                .store(true, Ordering::SeqCst);
             if !signal_shutdown_done.swap(true, Ordering::SeqCst) {
                 let _ = shutdown_managed_agents(&signal_app);
             }
@@ -915,7 +906,10 @@ pub fn run() {
     let run_shutdown_done = Arc::clone(&shutdown_done);
     app.run(move |app_handle, event| match event {
         RunEvent::ExitRequested { .. } | RunEvent::Exit => {
-            shutdown_started.store(true, Ordering::SeqCst);
+            app_handle
+                .state::<AppState>()
+                .shutdown_started
+                .store(true, Ordering::SeqCst);
             if !run_shutdown_done.swap(true, Ordering::SeqCst) {
                 prevent_sleep::release(&app_handle.state::<AppState>().prevent_sleep);
                 if let Err(error) = shutdown_managed_agents(app_handle) {

desktop/src-tauri/src/managed_agents/restore.rs

@@ -236,6 +236,18 @@ pub async fn restore_managed_agents_on_launch(
         return Ok(());
     }
 
+    // Serialize spawning and runtime registration with shutdown cleanup. The
+    // shutdown flag is rechecked after taking the lock so shutdown either
+    // prevents this transition or waits until every child is tracked and can
+    // be terminated.
+    let restore_transition = state
+        .managed_agent_restore_transition
+        .lock()
+        .map_err(|error| error.to_string())?;
+    if shutdown_started.load(Ordering::SeqCst) {
+        return Ok(());
+    }
+
     // ── Phase B (no locks): resolve commands and spawn processes in parallel ──
     let spawn_results: Vec<AgentSpawnResult> = std::thread::scope(|scope| {
         let owner_hex_ref = owner_hex.as_deref();
@@ -328,6 +340,9 @@ pub async fn restore_managed_agents_on_launch(
             .collect();
 
     save_managed_agents(app, &records)?;
+    drop(runtimes);
+    drop(_store_guard);
+    drop(restore_transition);
 
     // ── Profile reconciliation (fire-and-forget) ────────────────────────────
     // Spawn background tasks to ensure each restored agent's kind:0 profile is

desktop/src-tauri/src/shutdown.rs

@@ -9,6 +9,10 @@ use crate::util;
 
 pub(crate) fn shutdown_managed_agents(app: &tauri::AppHandle) -> Result<(), String> {
     let state = app.state::<AppState>();
+    let _restore_transition = state
+        .managed_agent_restore_transition
+        .lock()
+        .map_err(|error| error.to_string())?;
     let _store_guard = state
         .managed_agents_store_lock
         .lock()

desktop/src/features/messages/lib/useMentions.ts

@@ -326,6 +326,9 @@ export function useMentions(
 
     for (const member of members ?? []) {
       const pubkey = normalizePubkey(member.pubkey);
+      const linkedPersonaId = activePersonaById.has(pubkey)
+        ? pubkey
+        : undefined;
       const agentName =
         managedAgentNamesByPubkey.get(pubkey) ??
         relayAgentNamesByPubkey.get(pubkey) ??
@@ -342,7 +345,8 @@ export function useMentions(
           null,
         avatarUrl: profile?.avatarUrl ?? null,
         isMember: true,
-        personaId: managedAgentPersonaIdsByPubkey.get(pubkey),
+        personaId:
+          managedAgentPersonaIdsByPubkey.get(pubkey) ?? linkedPersonaId,
         isAgent:
           member.isAgent === true ||
           profile?.isAgent === true ||
@@ -360,11 +364,15 @@ export function useMentions(
     }
 
     for (const agent of relayAgentsQuery.data ?? []) {
+      const pubkey = normalizePubkey(agent.pubkey);
       addCandidate({
         kind: "identity",
-        pubkey: agent.pubkey,
+        pubkey,
         displayName: agent.name,
         isMember: false,
+        personaId:
+          managedAgentPersonaIdsByPubkey.get(pubkey) ??
+          (activePersonaById.has(pubkey) ? pubkey : undefined),
         ownerPubkey: null,
         isAgent: true,
       });
@@ -387,27 +395,25 @@ export function useMentions(
 
     if (canSearchGlobalUsers) {
       for (const user of userSearchResults) {
+        const pubkey = normalizePubkey(user.pubkey);
         addCandidate({
           kind: "identity",
-          pubkey: user.pubkey,
+          pubkey,
           displayName: formatSearchUserDisplayName(user),
           avatarUrl: user.avatarUrl ?? null,
-          personaId: managedAgentPersonaIdsByPubkey.get(
-            normalizePubkey(user.pubkey),
-          ),
+          personaId:
+            managedAgentPersonaIdsByPubkey.get(pubkey) ??
+            (activePersonaById.has(pubkey) ? pubkey : undefined),
           isMember: false,
           isAgent:
             user.isAgent ||
-            managedAgentNamesByPubkey.has(normalizePubkey(user.pubkey)) ||
-            relayAgentNamesByPubkey.has(normalizePubkey(user.pubkey)),
-          personaName:
-            personaNameByPubkey.get(normalizePubkey(user.pubkey)) ?? null,
+            managedAgentNamesByPubkey.has(pubkey) ||
+            relayAgentNamesByPubkey.has(pubkey),
+          personaName: personaNameByPubkey.get(pubkey) ?? null,
           secondaryLabel: formatSearchUserSecondaryLabel(user),
           ownerPubkey: user.ownerPubkey ?? null,
           isGlobalSearchResult: true,
-          isManagedAgent: managedAgentNamesByPubkey.has(
-            normalizePubkey(user.pubkey),
-          ),
+          isManagedAgent: managedAgentNamesByPubkey.has(pubkey),
         });
       }
     }
@@ -436,6 +442,7 @@ export function useMentions(
       },
     );
   }, [
+    activePersonaById,
     activePersonas,
     userSearchResults,
     canSearchGlobalUsers,