Merge pull request #479 from spa5k/feat/corrupt-branch-auto-repair

fix: auto-repair corrupt derived branch stores

Author: ScriptedAlchemy

.changeset/derived-branch-corruption-recovery.md

@@ -0,0 +1,5 @@
+---
+"tracedecay": patch
+---
+
+Preserve corrupt derived-branch database families under the project recovery directory and rebuild the active branch index from a healthy tracked ancestor, while keeping default stores fail-closed, correcting branch-specific doctor recovery paths, and reloading rotated daemon credentials when long-lived clients reconnect after a restart.

src/daemon.rs

@@ -917,9 +917,9 @@ fn is_transient_daemon_connect_error(kind: std::io::ErrorKind) -> bool {
     )
 }
 
-async fn connect_to_daemon_connection(connection: &DaemonConnection) -> Result<BrokerStream> {
-    connect_with_restart_grace(
-        connection,
+async fn connect_to_current_daemon(socket_path: &Path) -> Result<(DaemonConnection, BrokerStream)> {
+    connect_with_restart_grace_resolving(
+        || client_connection(socket_path),
         DAEMON_RESTART_GRACE,
         DAEMON_RESTART_POLL_INTERVAL,
     )
@@ -935,10 +935,27 @@ async fn connect_with_restart_grace(
     grace: Duration,
     poll_interval: Duration,
 ) -> Result<BrokerStream> {
+    let (_, stream) =
+        connect_with_restart_grace_resolving(|| Ok(connection.clone()), grace, poll_interval)
+            .await?;
+    Ok(stream)
+}
+
+/// Resolves the current authority record on every connect attempt.
+///
+/// A daemon restart replaces both its endpoint authority epoch and auth token.
+/// Resolving only before the restart window would connect to the new socket
+/// with stale credentials after it rebinds.
+async fn connect_with_restart_grace_resolving(
+    mut resolve: impl FnMut() -> Result<DaemonConnection>,
+    grace: Duration,
+    poll_interval: Duration,
+) -> Result<(DaemonConnection, BrokerStream)> {
     let deadline = tokio::time::Instant::now() + grace;
     loop {
+        let connection = resolve()?;
         match BrokerStream::connect(&connection.endpoint).await {
-            Ok(stream) => return Ok(stream),
+            Ok(stream) => return Ok((connection, stream)),
             Err(TraceDecayError::Io(err)) => {
                 if !is_transient_daemon_connect_error(err.kind())
                     || tokio::time::Instant::now() >= deadline
@@ -1291,8 +1308,7 @@ async fn send_daemon_request_line_with_liveness_poll(
     line: &str,
     liveness_poll_interval: Duration,
 ) -> Result<Vec<String>> {
-    let connection = client_connection(socket_path)?;
-    let stream = connect_to_daemon_connection(&connection).await?;
+    let (connection, stream) = connect_to_current_daemon(socket_path).await?;
     let (reader, mut writer) = stream.into_split();
 
     write_daemon_preamble(&mut writer, &connection, handshake).await?;
@@ -1490,8 +1506,7 @@ async fn call_tool_with_liveness_poll(
     arguments: serde_json::Value,
     liveness_poll_interval: Duration,
 ) -> Result<serde_json::Value> {
-    let connection = client_connection(socket_path)?;
-    let stream = connect_to_daemon_connection(&connection).await?;
+    let (connection, stream) = connect_to_current_daemon(socket_path).await?;
     let (reader, mut writer) = stream.into_split();
     let id = json!(1);
     let request = JsonRpcRequest {

src/daemon/service.rs

@@ -1078,9 +1078,9 @@ mod tests {
     #[cfg(unix)]
     use tempfile::TempDir;
 
-    use super::{
-        DaemonServiceSpec, DaemonServiceState, LaunchctlFailureMode, LaunchdCommand, ServiceRunner,
-    };
+    use super::{DaemonServiceSpec, LaunchctlFailureMode, LaunchdCommand};
+    #[cfg(target_os = "linux")]
+    use super::{DaemonServiceState, ServiceRunner};
     use crate::config::lock_user_data_dir_test_env;
 
     struct EnvVarGuard {

src/daemon/tests.rs

@@ -729,6 +729,55 @@ async fn answer_one_proxy_request(listener: tokio::net::UnixListener, generation
     writer.shutdown().await.expect("shutdown fake daemon");
 }
 
+#[cfg(unix)]
+async fn answer_one_authenticated_proxy_request(
+    listener: tokio::net::UnixListener,
+    expected_token: &str,
+    generation: u64,
+) {
+    let (stream, _addr) = listener.accept().await.expect("accept proxied client");
+    let (reader, mut writer) = stream.into_split();
+    let mut lines = tokio::io::BufReader::new(reader).lines();
+    let auth_line = lines
+        .next_line()
+        .await
+        .expect("read auth preface")
+        .expect("auth preface line");
+    let preface =
+        super::transport::DaemonAuthPreface::from_line(auth_line.trim()).expect("auth preface");
+    assert!(
+        preface.authenticate(expected_token),
+        "proxy must reload the current daemon authority token"
+    );
+    let handshake_line = lines
+        .next_line()
+        .await
+        .expect("read handshake")
+        .expect("handshake line");
+    DaemonHandshake::from_line(&handshake_line).expect("parse handshake");
+    let request_line = lines
+        .next_line()
+        .await
+        .expect("read request")
+        .expect("request line");
+    let request: Value = serde_json::from_str(&request_line).expect("request json");
+    let response = json!({
+        "jsonrpc": "2.0",
+        "id": request["id"],
+        "result": { "generation": generation }
+    });
+    writer
+        .write_all(
+            serde_json::to_string(&response)
+                .expect("response json")
+                .as_bytes(),
+        )
+        .await
+        .expect("write response");
+    writer.write_all(b"\n").await.expect("write newline");
+    writer.shutdown().await.expect("shutdown fake daemon");
+}
+
 #[cfg(unix)]
 async fn daemon_round_trip(
     engine: super::DaemonEngine,
@@ -1306,6 +1355,86 @@ async fn long_lived_proxy_reconnects_after_daemon_socket_rebind() {
     await_test_task(daemon, "daemon rebind task").await;
 }
 
+#[cfg(unix)]
+#[tokio::test]
+async fn long_lived_proxy_reloads_rotated_auth_after_daemon_restart() {
+    let dir = TempDir::new().expect("temp dir");
+    let profile = dir.path().canonicalize().expect("canonical profile");
+    let socket = profile.join("daemon.sock");
+    let endpoint = super::transport::DaemonEndpoint::Unix(socket.clone());
+    let first_listener = tokio::net::UnixListener::bind(&socket).expect("bind first daemon socket");
+    let first_authority = super::authority::DaemonAuthority::acquire(&profile, &endpoint, "first")
+        .expect("first daemon authority");
+    let first_token = first_authority.auth_token().to_string();
+    let rebound_socket = socket.clone();
+    let rebound_profile = profile.clone();
+    let rebound_endpoint = endpoint.clone();
+    let (unbound_tx, unbound_rx) = tokio::sync::oneshot::channel();
+    let daemon = tokio::spawn(async move {
+        answer_one_authenticated_proxy_request(first_listener, &first_token, 1).await;
+        drop(first_authority);
+        std::fs::remove_file(&rebound_socket).expect("unlink first daemon socket");
+        unbound_tx.send(()).expect("notify daemon outage");
+        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+
+        let second_listener =
+            tokio::net::UnixListener::bind(&rebound_socket).expect("bind second daemon socket");
+        let second_authority = super::authority::DaemonAuthority::acquire(
+            &rebound_profile,
+            &rebound_endpoint,
+            "second",
+        )
+        .expect("second daemon authority");
+        let second_token = second_authority.auth_token().to_string();
+        assert_ne!(first_token, second_token, "daemon restart must rotate auth");
+        answer_one_authenticated_proxy_request(second_listener, &second_token, 2).await;
+        drop(second_authority);
+    });
+
+    let (mut transport, sender, mut receiver) = crate::mcp::transport::ChannelTransport::new();
+    let proxy_socket = socket.clone();
+    let proxy = tokio::spawn(async move {
+        super::proxy_transport_to_daemon(
+            &proxy_socket,
+            &test_handshake_defaults(),
+            None,
+            &mut transport,
+        )
+        .await
+    });
+    let request = |id| {
+        serde_json::to_string(&json!({
+            "jsonrpc": "2.0",
+            "id": id,
+            "method": "tools/list"
+        }))
+        .expect("request json")
+    };
+
+    sender.send(request(1)).expect("send first request");
+    let first = tokio::time::timeout(std::time::Duration::from_secs(2), receiver.recv())
+        .await
+        .expect("first response timed out")
+        .expect("first response");
+    let first: Value = serde_json::from_str(first.trim()).expect("first response json");
+    assert_eq!(first["result"]["generation"], json!(1));
+
+    unbound_rx.await.expect("first daemon should unlink socket");
+    sender.send(request(2)).expect("send second request");
+    let second = tokio::time::timeout(std::time::Duration::from_secs(2), receiver.recv())
+        .await
+        .expect("second response timed out")
+        .expect("second response");
+    let second: Value = serde_json::from_str(second.trim()).expect("second response json");
+    assert_eq!(second["result"]["generation"], json!(2));
+
+    drop(sender);
+    await_test_task(proxy, "rotating-auth proxy task")
+        .await
+        .expect("proxy transport");
+    await_test_task(daemon, "rotating-auth daemon task").await;
+}
+
 #[cfg(unix)]
 #[tokio::test]
 async fn proxy_uses_daemon_initialize_route_without_registry_access() {

src/doctor.rs

@@ -375,7 +375,12 @@ fn fallback_database_path(project_path: &Path) -> Option<PathBuf> {
 fn database_recovery_guidance(db_path: &Path) -> String {
     let wal_path = db_path.with_extension("db-wal");
     let shm_path = db_path.with_extension("db-shm");
-    let data_root = db_path.parent().unwrap_or_else(|| Path::new("."));
+    let graph_parent = db_path.parent().unwrap_or_else(|| Path::new("."));
+    let data_root = if graph_parent.file_name() == Some(std::ffi::OsStr::new("branches")) {
+        graph_parent.parent().unwrap_or(graph_parent)
+    } else {
+        graph_parent
+    };
     let mut graph_dirty = db_path.as_os_str().to_os_string();
     graph_dirty.push(".dirty");
     let graph_dirty = PathBuf::from(graph_dirty);
@@ -391,7 +396,8 @@ fn database_recovery_guidance(db_path: &Path) -> String {
          graph dirty sentinel: {}\n\
          legacy dirty sentinel (if present): {}\n\
          `sessions.db` is separate and must not be removed: {}\n\
-         Facts are stored in the graph database; automatic rebuild is intentionally blocked because it cannot preserve them generically.\n\
+         Facts are stored in the graph database; automatic default-store rebuild is intentionally blocked because it cannot preserve them generically.\n\
+         Derived branch indexes are preserved under `recovery/` and rebuilt automatically from a healthy tracked ancestor.\n\
          Do not run `tracedecay init`, `tracedecay sync --force`, or `tracedecay wipe` until that recovery set is safely copied.\n\
          Report the preserved set at https://github.com/ScriptedAlchemy/tracedecay/issues for offline recovery.",
         db_path.display(),

src/doctor/tests.rs

@@ -236,7 +236,16 @@ fn database_recovery_guidance_names_the_preserved_recovery_set() {
     );
     assert!(guidance.contains("`sessions.db` is separate and must not be removed"));
     assert!(guidance.contains("Facts are stored in the graph database"));
-    assert!(guidance.contains("automatic rebuild is intentionally blocked"));
+    assert!(guidance.contains("automatic default-store rebuild is intentionally blocked"));
+    assert!(guidance.contains("Derived branch indexes are preserved"));
+
+    let branch_db = PathBuf::from("/profile/projects/proj_test/branches/feature.db");
+    let branch_guidance = database_recovery_guidance(&branch_db);
+    let branches_root = branch_db.parent().unwrap();
+    let data_root = branches_root.parent().unwrap();
+    assert!(branch_guidance.contains(&data_root.join("dirty").display().to_string()));
+    assert!(branch_guidance.contains(&data_root.join("sessions.db").display().to_string()));
+    assert!(!branch_guidance.contains(&branches_root.join("sessions.db").display().to_string()));
 }
 
 #[tokio::test]