@team-agent/installer 0.3.2 → 0.3.4

package/crates/team-agent/src/lifecycle/restart/rebuild.rs

@@ -1,6 +1,6 @@
-use super::*;
 use super::common::*;
 use super::selection::classify_restart_plan;
+use super::*;
 
 // ── lifecycle::restart —— 整队 Route B resume-or-fresh 重建 ──────────────────
 
@@ -12,13 +12,23 @@ pub fn restart(
     workspace: &Path,
     allow_fresh: bool,
     team: Option<&str>,
+) -> Result<RestartReport, LifecycleError> {
+    restart_with_session_convergence_deadline(workspace, allow_fresh, team, None)
+}
+
+pub fn restart_with_session_convergence_deadline(
+    workspace: &Path,
+    allow_fresh: bool,
+    team: Option<&str>,
+    session_converge_deadline_ms: Option<u64>,
 ) -> Result<RestartReport, LifecycleError> {
     let run_ws = lifecycle_run_workspace(workspace)?;
-    restart_with_transport(
+    restart_with_transport_with_session_convergence_deadline(
         workspace,
         allow_fresh,
         team,
         &crate::tmux_backend::TmuxBackend::for_workspace(&run_ws),
+        session_converge_deadline_ms,
     )
 }
 
@@ -30,6 +40,42 @@ pub fn restart_with_transport(
     allow_fresh: bool,
     team: Option<&str>,
     transport: &dyn crate::transport::Transport,
+) -> Result<RestartReport, LifecycleError> {
+    match restart_with_transport_with_session_convergence_deadline(
+        workspace,
+        allow_fresh,
+        team,
+        transport,
+        None,
+    )? {
+        RestartReport::RefusedResumeNotReady {
+            missing,
+            allow_fresh,
+            error,
+            ..
+        } => Ok(RestartReport::RefusedResumeAtomicity {
+            unresumable: missing
+                .into_iter()
+                .map(|agent_id| UnresumableWorker {
+                    agent_id,
+                    reason: "session_capture_incomplete".to_string(),
+                    session_id: None,
+                    first_send_at: None,
+                })
+                .collect(),
+            allow_fresh,
+            error,
+        }),
+        report => Ok(report),
+    }
+}
+
+pub fn restart_with_transport_with_session_convergence_deadline(
+    workspace: &Path,
+    allow_fresh: bool,
+    team: Option<&str>,
+    transport: &dyn crate::transport::Transport,
+    session_converge_deadline_ms: Option<u64>,
 ) -> Result<RestartReport, LifecycleError> {
     if crate::lifecycle::restart::input_has_no_local_team_context(workspace) {
         return Err(LifecycleError::TeamSelect(format!(
@@ -55,18 +101,62 @@ pub fn restart_with_transport(
     .map_err(|e| LifecycleError::TeamSelect(e.to_string()))?;
     let mut state = selected.state;
     crate::lifecycle::launch::ensure_owner_allowed_for_state(&state, None)?;
-    let spec_workspace = selected
-        .spec_workspace
-        .as_ref()
-        .ok_or_else(|| LifecycleError::TeamSelect("active team spec workspace not found".to_string()))?;
+    let spec_workspace = selected.spec_workspace.as_ref().ok_or_else(|| {
+        LifecycleError::TeamSelect("active team spec workspace not found".to_string())
+    })?;
     let spec = load_team_spec(spec_workspace)?;
     let safety = crate::lifecycle::launch::effective_runtime_config(&spec)?;
-    if refresh_missing_provider_sessions(&mut state)? {
+    let mut convergence = converge_missing_provider_sessions(
+        &mut state,
+        session_convergence_deadline(session_converge_deadline_ms),
+        session_convergence_poll_interval(),
+        &selected.run_workspace,
+        allow_fresh,
+    )?;
+    if convergence.converged && convergence.changed {
+        crate::state::projection::save_team_scoped_state(&selected.run_workspace, &state)
+            .map_err(|e| LifecycleError::StatePersist(e.to_string()))?;
+    }
+    if repair_resume_sessions_from_event_log(&selected.run_workspace, &mut state)? {
+        crate::state::projection::save_team_scoped_state(&selected.run_workspace, &state)
+            .map_err(|e| LifecycleError::StatePersist(e.to_string()))?;
+        let missing_after_repair = restart_required_missing_session_agent_ids(&state);
+        convergence.changed = true;
+        convergence.converged = missing_after_repair.is_empty();
+        convergence.missing = missing_after_repair;
+    }
+    if !convergence.converged && !allow_fresh {
+        return Ok(RestartReport::RefusedResumeNotReady {
+            missing: convergence
+                .missing
+                .iter()
+                .map(|agent_id| AgentId::new(agent_id.clone()))
+                .collect(),
+            allow_fresh,
+            deadline: convergence.deadline,
+            elapsed: convergence.elapsed,
+            error: "resume_not_ready: session_capture_incomplete".to_string(),
+        });
+    }
+    if !convergence.converged && convergence.changed {
         crate::state::projection::save_team_scoped_state(&selected.run_workspace, &state)
             .map_err(|e| LifecycleError::StatePersist(e.to_string()))?;
     }
+    let forced_fresh_missing = if convergence.converged {
+        std::collections::BTreeSet::new()
+    } else {
+        convergence.missing.iter().cloned().collect()
+    };
+    let forced_fresh_convergence = (!convergence.converged).then_some(convergence.clone());
     let plan = classify_restart_plan(&state, allow_fresh)?;
-    write_restart_resume_decision_events(&selected.run_workspace, &state, allow_fresh, &plan.decisions)?;
+    write_restart_resume_decision_events(
+        &selected.run_workspace,
+        &state,
+        allow_fresh,
+        &plan.decisions,
+        &forced_fresh_missing,
+        forced_fresh_convergence.as_ref(),
+    )?;
     if !plan.corrupt_entries.is_empty() {
         return Ok(RestartReport::RefusedInvalidFirstSendAt {
             invalid: plan.corrupt_entries,
@@ -86,8 +176,13 @@ pub fn restart_with_transport(
         transport
             .kill_session(&session_name)
             .map_err(|e| LifecycleError::Transport(e.to_string()))?;
+        mark_leader_receiver_rebind_required(&mut state, &session_name);
+        mark_restart_targets_stopped_after_teardown(&mut state, &plan.decisions);
+        crate::state::projection::save_team_scoped_state(&selected.run_workspace, &state)
+            .map_err(|e| LifecycleError::StatePersist(e.to_string()))?;
     }
-    for (idx, decision) in plan.decisions.iter().enumerate() {
+    let mut last_spawned: Option<AgentId> = None;
+    for decision in &plan.decisions {
         let agent = state
             .get("agents")
             .and_then(|v| v.get(decision.agent_id.as_str()))
@@ -96,36 +191,355 @@ pub fn restart_with_transport(
                     "agent {} not found for restart",
                     decision.agent_id
                 ))
-            })?;
+            })?
+            .clone();
         let session_id = if matches!(decision.restart_mode, StartMode::Resumed) {
             decision.session_id.as_ref()
         } else {
             None
         };
-        let _ = spawn_agent_window(
+        let session_live = session_live_or_default(transport, &session_name, false);
+        if !session_live {
+            if let Some(previous) = &last_spawned {
+                return Err(LifecycleError::Transport(format!(
+                    "session_disappeared_after_spawn: provider_resume_exited for {}; session {} disappeared before spawning {}",
+                    previous,
+                    session_name.as_str(),
+                    decision.agent_id
+                )));
+            }
+        }
+        let spawn = spawn_agent_window(
             &selected.run_workspace,
             &session_name,
             &decision.agent_id,
-            agent,
+            &agent,
             session_id,
-            idx > 0,
+            session_live,
             transport,
             Some(&safety),
+            Some(spec_workspace),
         )?;
+        verify_spawned_agent_live(&decision.agent_id, &spawn, transport)?;
+        mark_agent_respawned(&mut state, &decision.agent_id, &spawn, transport, &safety)?;
+        last_spawned = Some(decision.agent_id.clone());
+        if let Some(agent) = state
+            .get_mut("agents")
+            .and_then(serde_json::Value::as_object_mut)
+            .and_then(|agents| agents.get_mut(decision.agent_id.as_str()))
+            .and_then(serde_json::Value::as_object_mut)
+        {
+            persist_effective_approval_policy_for_restart(agent, &safety);
+        }
     }
+    crate::state::projection::save_team_scoped_state(&selected.run_workspace, &state)
+        .map_err(|e| LifecycleError::StatePersist(e.to_string()))?;
     let coordinator_started = start_coordinator_for_workspace(&selected.run_workspace)?;
+    let attach_commands = crate::tmux_backend::attach_commands_for_windows(
+        &selected.run_workspace,
+        &session_name,
+        plan.decisions
+            .iter()
+            .map(|decision| decision.agent_id.as_str()),
+    );
+    let next_actions = attach_commands.clone();
     Ok(RestartReport::Restarted {
         session_name,
         agents: plan.decisions,
         coordinator_started,
+        next_actions,
+        attach_commands,
     })
 }
 
+fn repair_resume_sessions_from_event_log(
+    workspace: &Path,
+    state: &mut serde_json::Value,
+) -> Result<bool, LifecycleError> {
+    let agent_ids = state
+        .get("agents")
+        .and_then(serde_json::Value::as_object)
+        .map(|agents| agents.keys().cloned().collect::<Vec<_>>())
+        .unwrap_or_default();
+    let mut changed = false;
+    for agent_id in agent_ids {
+        let previous = state
+            .get("agents")
+            .and_then(|agents| agents.get(&agent_id))
+            .cloned()
+            .unwrap_or(serde_json::Value::Null);
+        if previous
+            .get("session_id")
+            .and_then(serde_json::Value::as_str)
+            .is_some_and(|session| !session.is_empty())
+        {
+            continue;
+        }
+        let Some(provider) = previous
+            .get("provider")
+            .and_then(serde_json::Value::as_str)
+            .and_then(parse_provider)
+        else {
+            continue;
+        };
+        let auth_mode = previous
+            .get("auth_mode")
+            .and_then(serde_json::Value::as_str)
+            .and_then(parse_auth_mode)
+            .unwrap_or(AuthMode::Subscription);
+        let exclude_session_ids = claimed_session_ids_except(state, &agent_id);
+        let adapter = crate::provider::get_adapter(provider);
+        let repaired = crate::session_capture::recover_resume_session_from_events(
+            workspace,
+            &agent_id,
+            &previous,
+            adapter.as_ref(),
+            auth_mode,
+            &exclude_session_ids,
+        )
+        .map_err(|e| LifecycleError::Provider(e.to_string()))?;
+        let Some(repaired) = repaired else {
+            continue;
+        };
+        let old_session_id = previous
+            .get("session_id")
+            .and_then(serde_json::Value::as_str)
+            .filter(|session| !session.is_empty())
+            .map(str::to_string);
+        let session_id = repaired
+            .get("session_id")
+            .and_then(serde_json::Value::as_str)
+            .filter(|session| !session.is_empty())
+            .map(str::to_string);
+        let rollout_path = repaired
+            .get("rollout_path")
+            .and_then(serde_json::Value::as_str)
+            .filter(|path| !path.is_empty())
+            .map(str::to_string);
+        if let Some(agent) = state
+            .get_mut("agents")
+            .and_then(serde_json::Value::as_object_mut)
+            .and_then(|agents| agents.get_mut(&agent_id))
+        {
+            *agent = repaired.clone();
+        }
+        crate::event_log::EventLog::new(workspace)
+            .write(
+                "resume.session_repaired",
+                serde_json::json!({
+                    "agent_id": agent_id,
+                    "provider": provider_wire(provider),
+                    "old_session_id": old_session_id,
+                    "session_id": session_id,
+                    "rollout_path": rollout_path,
+                    "captured_via": "event_log_repair",
+                    "attribution_confidence": repaired.get("attribution_confidence").cloned().unwrap_or(serde_json::Value::Null),
+                }),
+            )
+            .map_err(|e| LifecycleError::StatePersist(e.to_string()))?;
+        changed = true;
+    }
+    Ok(changed)
+}
+
+fn claimed_session_ids_except(
+    state: &serde_json::Value,
+    current_agent_id: &str,
+) -> std::collections::BTreeSet<String> {
+    state
+        .get("agents")
+        .and_then(serde_json::Value::as_object)
+        .map(|agents| {
+            agents
+                .iter()
+                .filter(|(agent_id, _)| agent_id.as_str() != current_agent_id)
+                .filter_map(|(_, agent)| {
+                    agent
+                        .get("session_id")
+                        .and_then(serde_json::Value::as_str)
+                        .filter(|session| !session.is_empty())
+                        .map(str::to_string)
+                })
+                .collect()
+        })
+        .unwrap_or_default()
+}
+
+fn session_convergence_deadline(requested_ms: Option<u64>) -> std::time::Duration {
+    if let Some(ms) = requested_ms {
+        return std::time::Duration::from_millis(ms);
+    }
+    env_duration_ms(
+        &[
+            "TEAM_AGENT_RESTART_SESSION_CAPTURE_DEADLINE_MS",
+            "TEAM_AGENT_RESTART_SESSION_CONVERGENCE_DEADLINE_MS",
+            "TEAM_AGENT_RESTART_CAPTURE_DEADLINE_MS",
+            "TEAM_AGENT_RESTART_CAPTURE_TIMEOUT_MS",
+            "TEAM_AGENT_RESTART_SESSION_CAPTURE_TIMEOUT_MS",
+            "TEAM_AGENT_RESTART_SESSION_CONVERGENCE_TIMEOUT_MS",
+            "TEAM_AGENT_SESSION_CAPTURE_DEADLINE_MS",
+            "TEAM_AGENT_SESSION_CAPTURE_CONVERGENCE_DEADLINE_MS",
+            "TEAM_AGENT_SESSION_CAPTURE_TIMEOUT_MS",
+            "TEAM_AGENT_SESSION_CAPTURE_CONVERGENCE_TIMEOUT_MS",
+            "TEAM_AGENT_SESSION_CONVERGENCE_DEADLINE_MS",
+            "TEAM_AGENT_SESSION_CONVERGENCE_TIMEOUT_MS",
+            "TEAM_AGENT_PROVIDER_SESSION_CONVERGENCE_DEADLINE_MS",
+            "TEAM_AGENT_PROVIDER_SESSION_CONVERGENCE_TIMEOUT_MS",
+        ],
+        crate::session_capture::RESTART_SESSION_CONVERGENCE_DEADLINE_MS,
+    )
+}
+
+fn session_convergence_poll_interval() -> std::time::Duration {
+    env_duration_ms(
+        &[
+            "TEAM_AGENT_RESTART_SESSION_CAPTURE_POLL_MS",
+            "TEAM_AGENT_RESTART_SESSION_CONVERGENCE_POLL_MS",
+            "TEAM_AGENT_RESTART_CAPTURE_POLL_MS",
+            "TEAM_AGENT_SESSION_CAPTURE_POLL_MS",
+            "TEAM_AGENT_SESSION_CAPTURE_CONVERGENCE_POLL_MS",
+            "TEAM_AGENT_SESSION_CONVERGENCE_POLL_MS",
+            "TEAM_AGENT_PROVIDER_SESSION_CONVERGENCE_POLL_MS",
+        ],
+        crate::session_capture::RESTART_SESSION_CONVERGENCE_POLL_MS,
+    )
+}
+
+fn env_duration_ms(names: &[&str], default_ms: u64) -> std::time::Duration {
+    let ms = names
+        .iter()
+        .find_map(|name| {
+            std::env::var(name)
+                .ok()
+                .and_then(|value| parse_duration_value_ms(&value))
+                .or_else(|| {
+                    name.strip_suffix("_MS").and_then(|prefix| {
+                        std::env::var(prefix)
+                            .ok()
+                            .and_then(|value| parse_duration_value_seconds_ms(&value))
+                    })
+                })
+        })
+        .unwrap_or(default_ms);
+    std::time::Duration::from_millis(ms)
+}
+
+fn parse_duration_value_ms(value: &str) -> Option<u64> {
+    value.parse::<u64>().ok()
+}
+
+fn parse_duration_value_seconds_ms(value: &str) -> Option<u64> {
+    let seconds = value.parse::<f64>().ok()?;
+    if seconds.is_finite() && seconds >= 0.0 {
+        Some((seconds * 1000.0).round() as u64)
+    } else {
+        None
+    }
+}
+
+fn verify_spawned_agent_live(
+    _agent_id: &AgentId,
+    _spawn: &SpawnedAgentWindow,
+    _transport: &dyn crate::transport::Transport,
+) -> Result<(), LifecycleError> {
+    Ok(())
+}
+
+fn mark_leader_receiver_rebind_required(state: &mut serde_json::Value, session_name: &SessionName) {
+    let Some(receiver) = state
+        .get_mut("leader_receiver")
+        .and_then(serde_json::Value::as_object_mut)
+    else {
+        return;
+    };
+    let same_session = receiver
+        .get("session_name")
+        .and_then(|v| v.as_str())
+        .map(|session| session == session_name.as_str())
+        .unwrap_or(true);
+    if !same_session {
+        return;
+    }
+    if receiver
+        .get("status")
+        .and_then(|v| v.as_str())
+        .is_some_and(|status| status == "attached")
+    {
+        receiver.insert("status".to_string(), serde_json::json!("rebind_required"));
+    }
+}
+
+fn mark_restart_targets_stopped_after_teardown(
+    state: &mut serde_json::Value,
+    decisions: &[RestartedAgent],
+) {
+    let Some(agents) = state
+        .get_mut("agents")
+        .and_then(serde_json::Value::as_object_mut)
+    else {
+        return;
+    };
+    for decision in decisions {
+        let Some(agent) = agents
+            .get_mut(decision.agent_id.as_str())
+            .and_then(serde_json::Value::as_object_mut)
+        else {
+            continue;
+        };
+        agent.insert("status".to_string(), serde_json::json!("stopped"));
+        agent.remove("pane_id");
+        agent.remove("pane_pid");
+    }
+}
+
+fn mark_agent_respawned(
+    state: &mut serde_json::Value,
+    agent_id: &AgentId,
+    spawn: &SpawnedAgentWindow,
+    transport: &dyn crate::transport::Transport,
+    safety: &DangerousApproval,
+) -> Result<(), LifecycleError> {
+    let Some(agent) = state
+        .get_mut("agents")
+        .and_then(serde_json::Value::as_object_mut)
+        .and_then(|agents| agents.get_mut(agent_id.as_str()))
+        .and_then(serde_json::Value::as_object_mut)
+    else {
+        return Err(LifecycleError::StatePersist(format!(
+            "agent {} state is not an object",
+            agent_id
+        )));
+    };
+    agent.insert("status".to_string(), serde_json::json!("running"));
+    agent.insert(
+        "pane_id".to_string(),
+        serde_json::json!(spawn.spawn.pane_id.as_str()),
+    );
+    let pane_pid = spawn.spawn.child_pid.or_else(|| {
+        transport
+            .list_targets()
+            .unwrap_or_default()
+            .into_iter()
+            .find(|pane| pane.pane_id == spawn.spawn.pane_id)
+            .and_then(|pane| pane.pane_pid)
+    });
+    if let Some(pane_pid) = pane_pid {
+        agent.insert("pane_pid".to_string(), serde_json::json!(pane_pid));
+    }
+    crate::lifecycle::launch::persist_command_plan_state(agent, &spawn.plan, &spawn.profile_launch);
+    persist_effective_approval_policy_for_restart(agent, safety);
+    agent.remove("startup_prompts");
+    agent.remove("startup_prompt_status");
+    Ok(())
+}
+
 fn write_restart_resume_decision_events(
     workspace: &Path,
     state: &serde_json::Value,
     allow_fresh: bool,
     decisions: &[RestartedAgent],
+    forced_fresh_missing: &std::collections::BTreeSet<String>,
+    forced_fresh_convergence: Option<&crate::session_capture::SessionConvergence>,
 ) -> Result<(), LifecycleError> {
     for decision in decisions {
         let agent = state
@@ -150,6 +564,8 @@ fn write_restart_resume_decision_events(
             session_id,
             allow_fresh,
             decision_wire,
+            forced_fresh_missing.contains(decision.agent_id.as_str()),
+            forced_fresh_convergence,
         )?;
     }
     Ok(())
@@ -162,15 +578,16 @@ fn write_restart_resume_decision_event(
     session_id: Option<String>,
     allow_fresh: bool,
     decision: &str,
+    forced_fresh: bool,
+    forced_fresh_convergence: Option<&crate::session_capture::SessionConvergence>,
 ) -> Result<(), LifecycleError> {
     use std::io::Write as _;
 
     let path = workspace.join(".team").join("logs").join("events.jsonl");
     if let Some(parent) = path.parent() {
-        std::fs::create_dir_all(parent)
-            .map_err(|e| LifecycleError::StatePersist(e.to_string()))?;
+        std::fs::create_dir_all(parent).map_err(|e| LifecycleError::StatePersist(e.to_string()))?;
     }
-    let event = serde_json::json!({
+    let mut event = serde_json::json!({
         "ts": chrono::Utc::now().to_rfc3339(),
         "event": crate::lifecycle::types::event_names::RESTART_RESUME_DECISION,
         "worker_id": worker_id,
@@ -181,8 +598,26 @@ fn write_restart_resume_decision_event(
         "first_send_at": first_send_at,
         "session_id": session_id,
     });
-    let line = serde_json::to_string(&event)
-        .map_err(|e| LifecycleError::StatePersist(e.to_string()))?;
+    if forced_fresh {
+        if let Some(event) = event.as_object_mut() {
+            event.insert("forced_fresh".to_string(), serde_json::json!(true));
+            event.insert("reason".to_string(), serde_json::json!("resume_not_ready"));
+            if let Some(convergence) = forced_fresh_convergence {
+                event.insert(
+                    "session_convergence".to_string(),
+                    serde_json::json!({
+                        "complete": false,
+                        "deadline_s": convergence.deadline.as_secs_f64(),
+                        "deadline_ms": convergence.deadline.as_millis(),
+                        "elapsed_ms": convergence.elapsed.as_millis(),
+                        "pending_agent_ids": convergence.missing.clone(),
+                    }),
+                );
+            }
+        }
+    }
+    let line =
+        serde_json::to_string(&event).map_err(|e| LifecycleError::StatePersist(e.to_string()))?;
     let mut file = std::fs::OpenOptions::new()
         .create(true)
         .append(true)
@@ -233,7 +668,10 @@ pub fn select_restart_state(
         .get("active_team_key")
         .and_then(serde_json::Value::as_str)
         .filter(|s| !s.is_empty())
-        .map_or_else(|| crate::state::projection::team_state_key(&selected), str::to_string);
+        .map_or_else(
+            || crate::state::projection::team_state_key(&selected),
+            str::to_string,
+        );
     Ok(restart_candidate_from_state(workspace, &key, &selected))
 }
 
@@ -299,12 +737,14 @@ fn restart_candidate_has_context(state: &serde_json::Value) -> bool {
         .and_then(serde_json::Value::as_object)
         .is_some_and(|agents| {
             agents.values().any(|agent| {
-                ["session_id", "rollout_path", "first_send_at"].iter().any(|key| {
-                    agent
-                        .get(*key)
-                        .and_then(serde_json::Value::as_str)
-                        .is_some_and(|s| !s.is_empty())
-                })
+                ["session_id", "rollout_path", "first_send_at"]
+                    .iter()
+                    .any(|key| {
+                        agent
+                            .get(*key)
+                            .and_then(serde_json::Value::as_str)
+                            .is_some_and(|s| !s.is_empty())
+                    })
             })
         })
 }

package/crates/team-agent/src/lifecycle/restart/remove.rs

@@ -230,6 +230,7 @@ fn remove_agent_inner(
         "agent_health",
         None,
     )?;
+    maybe_fail_remove_after_agent_health_delete()?;
     Ok(RemoveSuccess {
         outcome: RemoveAgentOutcome::Removed {
             agent_id: agent_id.clone(),
@@ -585,15 +586,16 @@ fn select_agent_health(
         .map_err(|e| LifecycleError::StatePersist(e.to_string()))?;
     let row = conn
         .query_row(
-            "select status, last_output_at, context_usage_pct, current_task_id \
+            "select owner_team_id, status, last_output_at, context_usage_pct, current_task_id \
              from agent_health where agent_id = ?1",
             [agent_id.as_str()],
             |r| {
                 Ok(CapturedHealth {
-                    status: r.get::<_, Option<String>>(0)?,
-                    last_output_at: r.get::<_, Option<String>>(1)?,
-                    context_usage_pct: r.get::<_, Option<i64>>(2)?,
-                    current_task_id: r.get::<_, Option<String>>(3)?,
+                    owner_team_id: r.get::<_, Option<String>>(0)?,
+                    status: r.get::<_, Option<String>>(1)?,
+                    last_output_at: r.get::<_, Option<String>>(2)?,
+                    context_usage_pct: r.get::<_, Option<i64>>(3)?,
+                    current_task_id: r.get::<_, Option<String>>(4)?,
                 })
             },
         )
@@ -622,8 +624,9 @@ fn restore_agent_health(
     // health (golden _restore_agent_health re-upserts status||"IDLE" + the captured columns).
     conn.execute(
         "insert into agent_health (owner_team_id, agent_id, status, last_output_at, context_usage_pct, current_task_id, updated_at) \
-         values (null, ?1, ?2, ?3, ?4, ?5, ?6)",
+         values (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
         rusqlite::params![
+            row.owner_team_id,
             agent_id.as_str(),
             status,
             row.last_output_at,
@@ -638,12 +641,25 @@ fn restore_agent_health(
 
 #[derive(Clone)]
 struct CapturedHealth {
+    owner_team_id: Option<String>,
     status: Option<String>,
     last_output_at: Option<String>,
     context_usage_pct: Option<i64>,
     current_task_id: Option<String>,
 }
 
+fn maybe_fail_remove_after_agent_health_delete() -> Result<(), LifecycleError> {
+    let Ok(reason) = std::env::var("TEAM_AGENT_TEST_FAIL_REMOVE_AFTER_AGENT_HEALTH_DELETE") else {
+        return Ok(());
+    };
+    if reason.is_empty() {
+        return Ok(());
+    }
+    Err(LifecycleError::StatePersist(format!(
+        "injected remove failure after agent_health delete: {reason}"
+    )))
+}
+
 struct RemoveRollback {
     agent_id: AgentId,
     spec_text: Option<String>,