mirrors/zed
0
0 Fork 0
mirror of https://github.com/zed-industries/zed.git synced 2026-10-28 08:55:27 +08:00
Code Activity

Fix projects leaking via agents panel (#63765)

Browse Source 
When working with multiple windows, one agent worked in the external
hard drive's directory.
After finishing working on that external dir, I've closed the entire
window (the other one, unrelated, remained) only to discover I cannot
unmount the external hard drive.

Investigated, and seems that that manual `ref_count` is buggy — replaced
that with `WeakEntry` for both agent and ACP parts.
Also added leaked project logging into reliability.

Release Notes:

- Fixed projects leaking via agents panel
This commit is contained in:
Kirill Bulatov authored and GitHub committed 2026-09-04 18:14:56 +00:00
1 parent 9bda6b4e03
commit 5313741ffa
5 files changed
+413 -338
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

crates/agent/src/agent.rs
+237 -134
View File
@@ -228,7 +228,8 @@ struct Session {
/// The internal thread that processes messages
thread: Entity<Thread>,
/// The ACP thread that handles protocol communication
acp_thread: Entity<acp_thread::AcpThread>,
acp_thread: WeakEntity<acp_thread::AcpThread>,
subagents: Vec<Entity<acp_thread::AcpThread>>,
project_id: EntityId,
/// Latest snapshot to persist. Overwritten in place on every save request;
/// the single save worker drains it, coalescing bursts into one write.
@@ -236,12 +237,19 @@ struct Session {
save_wake: watch::Sender<()>,
save_worker: Task<Result<()>>,
_subscriptions: Vec<Subscription>,
ref_count: usize,
}
impl Session {
fn draft_prompt(&self, cx: &App) -> Option<Vec<acp::ContentBlock>> {
match self.acp_thread.upgrade() {
Some(acp_thread) => acp_thread.read(cx).draft_prompt().map(Vec::from),
None => self.thread.read(cx).draft_prompt().map(Vec::from),
}
}
}
struct PendingSession {
task: Shared<Task<Result<Entity<AcpThread>, Arc<anyhow::Error>>>>,
ref_count: usize,
}
pub struct LanguageModels {
@@ -776,14 +784,13 @@ impl NativeAgent {
)
});
self.register_session(thread, project_id, 1, cx)
self.register_session(thread, project_id, cx)
}
fn register_session(
&mut self,
thread_handle: Entity<Thread>,
project_id: EntityId,
ref_count: usize,
cx: &mut Context<Self>,
) -> Entity<AcpThread> {
let connection = Rc::new(NativeAgentConnection(cx.entity()));
@@ -848,6 +855,14 @@ impl NativeAgent {
cx.observe(&thread_handle, move |this, thread, cx| {
this.save_thread(thread, cx)
}),
cx.observe_release(&acp_thread, {
let session_id = session_id.clone();
let acp_thread_id = acp_thread.entity_id();
move |this, released_acp_thread, cx| {
let draft_prompt = released_acp_thread.draft_prompt().map(Vec::from);
this.release_session(&session_id, acp_thread_id, draft_prompt, cx);
}
}),
];
let (save_wake, save_wake_rx) = watch::channel(());
@@ -874,13 +889,13 @@ impl NativeAgent {
session_id,
Session {
thread: thread_handle,
acp_thread: acp_thread.clone(),
acp_thread: acp_thread.downgrade(),
subagents: Vec::new(),
project_id,
pending_save,
save_wake,
save_worker,
_subscriptions: subscriptions,
ref_count,
},
);
@@ -1359,7 +1374,7 @@ impl NativeAgent {
};
let thread = thread.downgrade();
let acp_thread = session.acp_thread.downgrade();
let acp_thread = session.acp_thread.clone();
cx.spawn(async move |_, cx| {
let title = thread.read_with(cx, |thread, _| thread.title())?;
if let Some(title) = title {
@@ -1381,9 +1396,12 @@ impl NativeAgent {
let Some(session) = self.sessions.get(thread.read(cx).id()) else {
return;
};
session.acp_thread.update(cx, |acp_thread, cx| {
acp_thread.update_token_usage(usage.0.clone(), cx);
});
session
.acp_thread
.update(cx, |acp_thread, cx| {
acp_thread.update_token_usage(usage.0.clone(), cx);
})
.ok();
}
fn handle_project_event(
@@ -1536,16 +1554,19 @@ impl NativeAgent {
if session.project_id != project_id {
continue;
}
session.acp_thread.update(cx, |thread, cx| {
thread
.handle_session_update(
acp::SessionUpdate::AvailableCommandsUpdate(
acp::AvailableCommandsUpdate::new(available_commands.clone()),
),
cx,
)
.log_err();
});
session
.acp_thread
.update(cx, |thread, cx| {
thread
.handle_session_update(
acp::SessionUpdate::AvailableCommandsUpdate(
acp::AvailableCommandsUpdate::new(available_commands.clone()),
),
cx,
)
.log_err();
})
.ok();
}
}
@@ -1661,13 +1682,16 @@ impl NativeAgent {
project: Entity<Project>,
cx: &mut Context<Self>,
) -> Task<Result<Entity<AcpThread>>> {
if let Some(session) = self.sessions.get_mut(&id) {
session.ref_count += 1;
return Task::ready(Ok(session.acp_thread.clone()));
if let Some(session) = self.sessions.get(&id) {
if let Some(acp_thread) = session.acp_thread.upgrade() {
return Task::ready(Ok(acp_thread));
}
let acp_thread_id = session.acp_thread.entity_id();
let draft_prompt = session.draft_prompt(cx);
self.release_session(&id, acp_thread_id, draft_prompt, cx);
}
if let Some(pending) = self.pending_sessions.get_mut(&id) {
pending.ref_count += 1;
if let Some(pending) = self.pending_sessions.get(&id) {
let task = pending.task.clone();
return cx.background_spawn(async move { task.await.map_err(|err| anyhow!(err)) });
}
@@ -1690,11 +1714,8 @@ impl NativeAgent {
let acp_thread = this
.update(cx, |this, cx| {
let project_id = this.get_or_create_project_state(&project, cx);
let ref_count = this
.pending_sessions
.remove(&id)
.map_or(1, |pending| pending.ref_count);
this.register_session(thread.clone(), project_id, ref_count, cx)
this.pending_sessions.remove(&id);
this.register_session(thread.clone(), project_id, cx)
})
.map_err(Arc::new)?;
let events = thread.update(cx, |thread, cx| thread.replay(cx));
@@ -1719,7 +1740,6 @@ impl NativeAgent {
id,
PendingSession {
task: shared_task.clone(),
ref_count: 1,
},
);
@@ -1739,48 +1759,44 @@ impl NativeAgent {
.update(cx, |this, cx| {
this.sessions
.get(&id)
.unwrap()
.context("session released before summary")?
.thread
.update(cx, |thread, cx| thread.summary(cx))
})?
.update(cx, |thread, cx| anyhow::Ok(thread.summary(cx)))
})??
.await
.context("Failed to generate summary")?;
this.update(cx, |this, cx| this.close_session(&id, cx))?
.await?;
drop(acp_thread);
Ok(result)
})
}
fn close_session(
fn release_session(
&mut self,
session_id: &acp::SessionId,
acp_thread_id: EntityId,
draft_prompt: Option<Vec<acp::ContentBlock>>,
cx: &mut Context<Self>,
) -> Task<Result<()>> {
let Some(session) = self.sessions.get_mut(session_id) else {
return Task::ready(Ok(()));
) {
let Some(session) = self.sessions.get(session_id) else {
return;
};
session.ref_count -= 1;
if session.ref_count > 0 {
return Task::ready(Ok(()));
if session.acp_thread.entity_id() != acp_thread_id {
return;
}
let thread = session.thread.clone();
self.save_thread(thread, cx);
self.enqueue_save(session_id, draft_prompt, cx);
let Some(session) = self.sessions.remove(session_id) else {
return Task::ready(Ok(()));
return;
};
let project_id = session.project_id;
session.save_worker.detach_and_log_err(cx);
let has_remaining = self.sessions.values().any(|s| s.project_id == project_id);
if !has_remaining {
self.projects.remove(&project_id);
self.publish_skill_index(cx);
}
session.save_worker
}
fn save_thread(&mut self, thread: Entity<Thread>, cx: &mut Context<Self>) {
@@ -1788,7 +1804,22 @@ impl NativeAgent {
let Some(session) = self.sessions.get(&id) else {
return;
};
let Some((id, folder_paths, db_thread)) = self.thread_save_payload(session, cx) else {
let draft_prompt = session.draft_prompt(cx);
self.enqueue_save(&id, draft_prompt, cx);
}
fn enqueue_save(
&mut self,
id: &acp::SessionId,
draft_prompt: Option<Vec<acp::ContentBlock>>,
cx: &mut Context<Self>,
) {
let Some(session) = self.sessions.get(id) else {
return;
};
let Some((id, folder_paths, db_thread)) =
self.thread_save_payload(session, draft_prompt, cx)
else {
return;
};
@@ -1843,6 +1874,7 @@ impl NativeAgent {
fn thread_save_payload(
&self,
session: &Session,
draft_prompt: Option<Vec<acp::ContentBlock>>,
cx: &mut App,
) -> Option<(acp::SessionId, PathList, Task<DbThread>)> {
if session.thread.read(cx).is_empty() {
@@ -1857,7 +1889,6 @@ impl NativeAgent {
.map(|worktree| worktree.read(cx).abs_path().to_path_buf())
.collect::<Vec<_>>(),
);
let draft_prompt = session.acp_thread.read(cx).draft_prompt().map(Vec::from);
let id = session.thread.read(cx).id().clone();
let db_thread = session.thread.update(cx, |thread, cx| {
thread.set_draft_prompt(draft_prompt);
@@ -1876,7 +1907,8 @@ impl NativeAgent {
let mut saves = Vec::new();
for session in self.sessions.values() {
saves.extend(self.thread_save_payload(session, cx));
let draft_prompt = session.draft_prompt(cx);
saves.extend(self.thread_save_payload(session, draft_prompt, cx));
}
async move {
@@ -1928,7 +1960,11 @@ impl NativeAgent {
.sessions
.get(&session_id)
.context("Failed to get session")?;
anyhow::Ok((session.acp_thread.clone(), session.thread.clone()))
let acp_thread = session
.acp_thread
.upgrade()
.context("Session was released")?;
anyhow::Ok((acp_thread, session.thread.clone()))
})??;
let mut last_is_user = true;
@@ -2018,7 +2054,11 @@ impl NativeAgent {
.sessions
.get(&session_id)
.context("Failed to get session")?;
anyhow::Ok((session.acp_thread.clone(), session.thread.clone()))
let acp_thread = session
.acp_thread
.upgrade()
.context("Session was released")?;
anyhow::Ok((acp_thread, session.thread.clone()))
})??;
let response_stream =
@@ -2069,7 +2109,11 @@ impl NativeAgent {
.sessions
.get(&session_id)
.context("Failed to get session")?;
anyhow::Ok((session.acp_thread.clone(), session.thread.clone()))
let acp_thread = session
.acp_thread
.upgrade()
.context("Session was released")?;
anyhow::Ok((acp_thread, session.thread.clone()))
})??;
// Build the model-context message: skill envelope first, then
@@ -2214,10 +2258,8 @@ impl NativeAgentConnection {
+ FnOnce(Entity<Thread>, &mut App) -> Result<mpsc::UnboundedReceiver<Result<ThreadEvent>>>,
) -> Task<Result<acp::PromptResponse>> {
let Some((thread, acp_thread)) = self.0.update(cx, |agent, _cx| {
agent
.sessions
.get_mut(&session_id)
.map(|s| (s.thread.clone(), s.acp_thread.clone()))
let session = agent.sessions.get(&session_id)?;
Some((session.thread.clone(), session.acp_thread.clone()))
}) else {
log::error!("Session not found in run_turn: {}", session_id);
return Task::ready(Err(anyhow!("Session not found")));
@@ -2228,12 +2270,7 @@ impl NativeAgentConnection {
Ok(stream) => stream,
Err(err) => return Task::ready(Err(err)),
};
Self::handle_thread_events(
response_stream,
acp_thread.downgrade(),
Some(self.clone()),
cx,
)
Self::handle_thread_events(response_stream, acp_thread, Some(self.clone()), cx)
}
fn handle_thread_events(
@@ -2723,19 +2760,6 @@ impl acp_thread::AgentConnection for NativeAgentConnection {
.update(cx, |agent, cx| agent.open_thread(session_id, project, cx))
}
fn supports_close_session(&self) -> bool {
true
}
fn close_session(
self: Rc<Self>,
session_id: &acp::SessionId,
cx: &mut App,
) -> Task<Result<()>> {
self.0
.update(cx, |agent, cx| agent.close_session(session_id, cx))
}
fn auth_methods(&self) -> &[acp::AuthMethod] {
&[] // No auth for in-process
}
@@ -2804,7 +2828,7 @@ impl acp_thread::AgentConnection for NativeAgentConnection {
agent.sessions.get(session_id).map(|session| {
Rc::new(NativeAgentSessionTruncate {
thread: session.thread.clone(),
acp_thread: session.acp_thread.downgrade(),
acp_thread: session.acp_thread.clone(),
}) as _
})
})
@@ -3192,7 +3216,13 @@ impl NativeThreadEnvironment {
.get(&parent_session_id)
.map(|s| s.project_id)
.context("parent session not found")?;
Ok(agent.register_session(subagent_thread.clone(), project_id, 1, cx))
let acp_thread = agent.register_session(subagent_thread.clone(), project_id, cx);
let parent_session = agent
.sessions
.get_mut(&parent_session_id)
.context("parent session not found")?;
parent_session.subagents.push(acp_thread.clone());
Ok(acp_thread)
})??;
let depth = current_depth + 1;
@@ -3218,7 +3248,11 @@ impl NativeThreadEnvironment {
.sessions
.get(&session_id)
.ok_or_else(|| anyhow!("No subagent session found with id {session_id}"))?;
anyhow::Ok((session.thread.clone(), session.acp_thread.clone()))
let acp_thread = session
.acp_thread
.upgrade()
.ok_or_else(|| anyhow!("Subagent session {session_id} was released"))?;
anyhow::Ok((session.thread.clone(), acp_thread))
})??;
let depth = subagent_thread.read(cx).depth();
@@ -3832,6 +3866,7 @@ mod internal_tests {
use std::path::Path;
use super::*;
use crate::tests::release_dropped_entities;
use acp_thread::{AgentConnection, AgentModelGroupName, AgentModelInfo, MentionUri};
use agent_settings::COMPACTION_PROMPT;
use fs::FakeFs;
@@ -5249,7 +5284,7 @@ mod internal_tests {
// Run the subagent through the production registration path.
// This is what installs the `SkillTool` on the thread.
let _subagent_acp = agent.update(cx, |agent, cx| {
agent.register_session(subagent_thread.clone(), parent_project_id, 1, cx)
agent.register_session(subagent_thread.clone(), parent_project_id, cx)
});
// Verify the subagent thread has the `SkillTool` installed —
@@ -5494,7 +5529,12 @@ mod internal_tests {
cx: &mut TestAppContext,
fs: Arc<FakeFs>,
root: &str,
) -> (Entity<NativeAgent>, Entity<Project>, WorktreeId) {
) -> (
Entity<NativeAgent>,
Entity<Project>,
WorktreeId,
Entity<AcpThread>,
) {
use collections::{HashMap, HashSet};
use project::trusted_worktrees::{self, PathTrust, TrustedWorktrees};
@@ -5508,7 +5548,7 @@ mod internal_tests {
cx.update(|cx| NativeAgent::new(thread_store, Templates::new(), fs.clone(), cx));
let connection = NativeAgentConnection(agent.clone());
let _acp_thread = cx
let acp_thread = cx
.update(|cx| {
Rc::new(connection).new_session(
project.clone(),
@@ -5536,7 +5576,7 @@ mod internal_tests {
});
cx.run_until_parked();
(agent, project, worktree_id)
(agent, project, worktree_id, acp_thread)
}
/// The body resolver for a project-local skill must read the file
@@ -5564,7 +5604,7 @@ mod internal_tests {
)
.await;
let (agent, project, worktree_id) =
let (agent, project, worktree_id, _acp_thread) =
open_trusted_project_skills(cx, fs.clone(), "/project").await;
let project_id = project.entity_id();
@@ -5631,7 +5671,7 @@ mod internal_tests {
)
.await;
let (agent, project, _worktree_id) =
let (agent, project, _worktree_id, _acp_thread) =
open_trusted_project_skills(cx, fs.clone(), "/project").await;
let project_id = project.entity_id();
@@ -5681,7 +5721,7 @@ mod internal_tests {
)
.await;
let (agent, project, _worktree_id) =
let (agent, project, _worktree_id, _acp_thread) =
open_trusted_project_skills(cx, fs.clone(), "/project").await;
let project_id = project.entity_id();
@@ -5729,7 +5769,7 @@ mod internal_tests {
)
.await;
let (agent, project, worktree_id) =
let (agent, project, worktree_id, _acp_thread) =
open_trusted_project_skills(cx, fs.clone(), "/project").await;
let project_id = project.entity_id();
@@ -6164,11 +6204,9 @@ mod internal_tests {
cx.run_until_parked();
// Close the session so it can be reloaded from disk.
cx.update(|cx| connection.clone().close_session(&session_id, cx))
.await
.unwrap();
drop(thread);
drop(acp_thread);
release_dropped_entities(cx);
agent.read_with(cx, |agent, _| {
assert!(agent.sessions.is_empty());
});
@@ -6267,11 +6305,9 @@ mod internal_tests {
cx.run_until_parked();
// Close the session so it can be reloaded from disk.
cx.update(|cx| connection.clone().close_session(&session_id, cx))
.await
.unwrap();
drop(thread);
drop(acp_thread);
release_dropped_entities(cx);
agent.read_with(cx, |agent, _| {
assert!(agent.sessions.is_empty());
});
@@ -6362,10 +6398,8 @@ mod internal_tests {
send.await.unwrap();
cx.run_until_parked();
cx.update(|cx| connection.clone().close_session(&session_id, cx))
.await
.unwrap();
drop(acp_thread);
release_dropped_entities(cx);
(agent, connection, project, session_id, provider)
}
@@ -6624,11 +6658,9 @@ mod internal_tests {
});
// Close the session so it can be reloaded from disk.
cx.update(|cx| connection.clone().close_session(&session_id, cx))
.await
.unwrap();
drop(thread);
drop(acp_thread);
release_dropped_entities(cx);
agent.read_with(cx, |agent, _| {
assert_eq!(agent.sessions.keys().cloned().collect::<Vec<_>>(), []);
});
@@ -6688,7 +6720,7 @@ mod internal_tests {
}
#[gpui::test]
async fn test_close_session_saves_thread(cx: &mut TestAppContext) {
async fn test_releasing_session_saves_thread(cx: &mut TestAppContext) {
init_test(cx);
let fs = FakeFs::new(cx.executor());
fs.insert_tree(
@@ -6746,10 +6778,12 @@ mod internal_tests {
});
// Close the session immediately — no run_until_parked in between.
cx.update(|cx| connection.clone().close_session(&session_id, cx))
.await
.unwrap();
cx.run_until_parked();
drop(thread);
drop(acp_thread);
release_dropped_entities(cx);
agent.read_with(cx, |agent, _| {
assert!(agent.sessions.is_empty());
});
// Reopen and verify the draft prompt was saved.
let reloaded = agent
@@ -6762,11 +6796,93 @@ mod internal_tests {
assert_eq!(
thread.draft_prompt(),
Some(draft_blocks.as_slice()),
"close_session must save the thread; draft prompt was lost"
"releasing the session must save the thread; draft prompt was lost"
);
});
}
#[gpui::test]
async fn test_releasing_thread_releases_subagent_sessions_and_project(cx: &mut TestAppContext) {
init_test(cx);
let fs = FakeFs::new(cx.executor());
fs.insert_tree(
"/",
json!({
"a": {
"file.txt": "hello"
}
}),
)
.await;
let thread_store = cx.new(|cx| ThreadStore::new(cx));
let agent = cx
.update(|cx| NativeAgent::new(thread_store.clone(), Templates::new(), fs.clone(), cx));
let connection = Rc::new(NativeAgentConnection(agent.clone()));
let warmup_project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
drop(warmup_project);
release_dropped_entities(cx);
let leak_snapshot = cx.update(|cx| cx.leak_detector_snapshot());
let project = Project::test(fs.clone(), [path!("/a").as_ref()], cx).await;
let weak_project = project.downgrade();
let acp_thread = cx
.update(|cx| {
connection
.clone()
.new_session(project.clone(), PathList::new(&[Path::new("")]), cx)
})
.await
.unwrap();
let session_id = acp_thread.read_with(cx, |thread, _| thread.session_id().clone());
let thread = cx.update(|cx| native_thread_for_session(&agent, &session_id, cx));
let environment = NativeThreadEnvironment {
agent: agent.downgrade(),
thread: thread.downgrade(),
acp_thread: acp_thread.downgrade(),
};
let first_subagent = cx
.update(|cx| environment.create_subagent_thread("first".to_string(), cx))
.unwrap();
let second_subagent = cx
.update(|cx| environment.create_subagent_thread("second".to_string(), cx))
.unwrap();
cx.run_until_parked();
let session_ids = agent.read_with(cx, |agent, _| {
let mut ids = agent
.sessions
.keys()
.map(ToString::to_string)
.collect::<Vec<_>>();
ids.sort();
ids
});
let mut expected_ids = vec![
session_id.to_string(),
first_subagent.id().to_string(),
second_subagent.id().to_string(),
];
expected_ids.sort();
assert_eq!(session_ids, expected_ids);
drop(first_subagent);
drop(second_subagent);
drop(environment);
drop(thread);
drop(acp_thread);
drop(project);
release_dropped_entities(cx);
cx.update(|cx| cx.assert_no_new_leaks(&leak_snapshot));
let (session_count, project_count) =
agent.read_with(cx, |agent, _| (agent.sessions.len(), agent.projects.len()));
assert_eq!(session_count, 0);
assert_eq!(project_count, 0);
assert!(weak_project.upgrade().is_none());
}
#[gpui::test]
async fn test_save_burst_preserves_in_flight_write(cx: &mut TestAppContext) {
init_test(cx);
@@ -6918,25 +7034,20 @@ mod internal_tests {
cx.run_until_parked();
agent.read_with(cx, |agent, _| {
let session = agent
.sessions
.get(&session_id)
.expect("thread_summary should not close the active session");
assert_eq!(
session.ref_count, 1,
"thread_summary should release its temporary session reference"
assert!(
agent.sessions.contains_key(&session_id),
"thread_summary should not close the active session"
);
});
cx.update(|cx| connection.clone().close_session(&session_id, cx))
.await
.unwrap();
cx.run_until_parked();
drop(thread);
drop(acp_thread);
release_dropped_entities(cx);
agent.read_with(cx, |agent, _| {
assert!(
agent.sessions.is_empty(),
"closing the active session after thread_summary should unload it"
"dropping the active session after thread_summary should unload it"
);
});
}
@@ -6993,11 +7104,9 @@ mod internal_tests {
send.await.unwrap();
cx.run_until_parked();
cx.update(|cx| connection.clone().close_session(&session_id, cx))
.await
.unwrap();
drop(thread);
drop(acp_thread);
release_dropped_entities(cx);
agent.read_with(cx, |agent, _| {
assert!(agent.sessions.is_empty());
});
@@ -7032,14 +7141,13 @@ mod internal_tests {
"concurrent loads for the same session should share one AcpThread"
);
cx.update(|cx| connection.clone().close_session(&session_id, cx))
.await
.unwrap();
drop(first_loaded_thread);
release_dropped_entities(cx);
agent.read_with(cx, |agent, _| {
assert!(
agent.sessions.contains_key(&session_id),
"closing one loaded session should not drop shared session state"
"dropping one loaded handle should not drop shared session state"
);
});
@@ -7078,14 +7186,9 @@ mod internal_tests {
);
});
cx.update(|cx| connection.clone().close_session(&session_id, cx))
.await
.unwrap();
cx.run_until_parked();
drop(first_loaded_thread);
drop(second_loaded_thread);
release_dropped_entities(cx);
agent.read_with(cx, |agent, _| {
assert!(agent.sessions.is_empty());
});
crates/agent/src/tests/mod.rs
+14 -7
View File
@@ -64,6 +64,11 @@ pub(crate) fn init_test(cx: &mut TestAppContext) {
});
}
pub(crate) fn release_dropped_entities(cx: &mut TestAppContext) {
cx.update(|_| ());
cx.run_until_parked();
}
pub(crate) struct FakeTerminalHandle {
killed: Arc<AtomicBool>,
stopped_by_user: Arc<AtomicBool>,
@@ -4267,10 +4272,8 @@ async fn test_agent_connection(cx: &mut TestAppContext) {
request.await.expect("prompt should fail gracefully");
// Explicitly close the session and drop the ACP thread.
cx.update(|cx| Rc::new(connection.clone()).close_session(&session_id, cx))
.await
.unwrap();
drop(acp_thread);
release_dropped_entities(cx);
let result = cx
.update(|cx| {
acp_thread::AgentSessionClientUserMessageIds::prompt(
@@ -5596,7 +5599,8 @@ async fn test_subagent_tool_call_end_to_end(cx: &mut TestAppContext) {
.get(&subagent_session_id)
.expect("subagent session should exist")
.acp_thread
.clone()
.upgrade()
.expect("subagent thread should be alive")
});
model.send_last_completion_stream_text_chunk("subagent task response");
@@ -5732,7 +5736,8 @@ async fn test_subagent_tool_output_does_not_include_thinking(cx: &mut TestAppCon
.get(&subagent_session_id)
.expect("subagent session should exist")
.acp_thread
.clone()
.upgrade()
.expect("subagent thread should be alive")
});
model.send_last_completion_stream_text_chunk("subagent task response 1");
@@ -5880,7 +5885,8 @@ async fn test_subagent_tool_call_cancellation_during_task_prompt(cx: &mut TestAp
.get(&subagent_session_id)
.expect("subagent session should exist")
.acp_thread
.clone()
.upgrade()
.expect("subagent thread should be alive")
});
// model.send_last_completion_stream_text_chunk("subagent task response");
@@ -6012,7 +6018,8 @@ async fn test_subagent_tool_resume_session(cx: &mut TestAppContext) {
.get(&subagent_session_id)
.expect("subagent session should exist")
.acp_thread
.clone()
.upgrade()
.expect("subagent thread should be alive")
});
// Subagent responds
crates/agent_servers/src/acp.rs
+116 -185
View File
@@ -491,7 +491,6 @@ impl AcpConnectionDefaults {
struct PendingAcpSession {
task: Shared<Task<Result<Entity<AcpThread>, Arc<anyhow::Error>>>>,
ref_count: usize,
}
struct SessionConfigResponse {
@@ -522,7 +521,7 @@ pub struct AcpSession {
suppress_abort_err: bool,
session_modes: Option<Rc<RefCell<acp::SessionModeState>>>,
config_options: Option<ConfigOptions>,
ref_count: usize,
_release_subscription: Subscription,
}
pub struct AcpSessionList {
@@ -1163,6 +1162,60 @@ impl AcpConnection {
session_directories_from_work_dirs(work_dirs, supports_additional_directories)
}
fn register_session(
&self,
session_id: acp::SessionId,
thread: &Entity<AcpThread>,
session_modes: Option<Rc<RefCell<acp::SessionModeState>>>,
config_options: Option<ConfigOptions>,
cx: &mut App,
) {
let release_subscription = cx.observe_release(thread, {
let session_id = session_id.clone();
let thread_id = thread.entity_id();
let sessions = self.sessions.clone();
let connection = self.connection.clone();
let supports_close = self.agent_supports_session_close();
move |_, cx| {
let removed = {
let mut sessions = sessions.borrow_mut();
match sessions.get(&session_id) {
Some(session) if session.thread.entity_id() == thread_id => {
sessions.remove(&session_id)
}
_ => None,
}
};
if removed.is_none() || !supports_close {
return;
}
cx.foreground_executor()
.spawn(async move {
connection
.send_request(acp::CloseSessionRequest::new(session_id))
.block_task()
.await
.log_err();
})
.detach();
}
});
self.sessions.borrow_mut().insert(
session_id,
AcpSession {
thread: thread.downgrade(),
suppress_abort_err: false,
session_modes,
config_options,
_release_subscription: release_subscription,
},
);
}
fn agent_supports_session_close(&self) -> bool {
self.agent_capabilities.session_capabilities.close.is_some()
}
fn open_or_create_session(
self: Rc<Self>,
session_id: acp::SessionId,
@@ -1184,19 +1237,20 @@ impl AcpConnection {
// Concurrent loads should still wait for the in-flight task so that
// ref-counting happens in one place and the caller sees a fully loaded
// session.
if let Some(pending) = self.pending_sessions.borrow_mut().get_mut(&session_id) {
pending.ref_count += 1;
if let Some(pending) = self.pending_sessions.borrow().get(&session_id) {
let task = pending.task.clone();
return cx
.foreground_executor()
.spawn(async move { task.await.map_err(|err| anyhow!(err)) });
}
if let Some(session) = self.sessions.borrow_mut().get_mut(&session_id) {
session.ref_count += 1;
if let Some(thread) = session.thread.upgrade() {
return Task::ready(Ok(thread));
}
if let Some(thread) = self
.sessions
.borrow()
.get(&session_id)
.and_then(|session| session.thread.upgrade())
{
return Task::ready(Ok(thread));
}
let directories = match self.session_directories_from_work_dirs(&work_dirs) {
@@ -1230,16 +1284,9 @@ impl AcpConnection {
// `session/update` notifications that arrive during the call
// (e.g. history replay during `session/load`) can find the thread.
// Modes/config are filled in once the response arrives.
this.sessions.borrow_mut().insert(
session_id.clone(),
AcpSession {
thread: thread.downgrade(),
suppress_abort_err: false,
session_modes: None,
config_options: None,
ref_count: 1,
},
);
cx.update(|cx| {
this.register_session(session_id.clone(), &thread, None, None, cx)
});
let response =
match rpc_call(this.connection.clone(), session_id.clone(), directories)
@@ -1260,17 +1307,8 @@ impl AcpConnection {
this.apply_default_config_options(&session_id, config_opts, cx);
}
let ref_count = this
.pending_sessions
.borrow_mut()
.remove(&session_id)
.map_or(1, |pending| pending.ref_count);
this.pending_sessions.borrow_mut().remove(&session_id);
// If `close_session` ran to completion while the load RPC was in
// flight, it will have removed both the pending entry and the
// sessions entry (and dispatched the ACP close RPC). In that case
// the thread has no live session to attach to, so fail the load
// instead of handing back an orphaned thread.
{
let mut sessions = this.sessions.borrow_mut();
let Some(session) = sessions.get_mut(&session_id) else {
@@ -1280,7 +1318,6 @@ impl AcpConnection {
};
session.session_modes = modes;
session.config_options = config_options.map(ConfigOptions::new);
session.ref_count = ref_count;
}
Ok(thread)
@@ -1292,7 +1329,6 @@ impl AcpConnection {
session_id,
PendingAcpSession {
task: shared_task.clone(),
ref_count: 1,
},
);
@@ -1693,16 +1729,15 @@ impl AgentConnection for AcpConnection {
)
});
self.sessions.borrow_mut().insert(
response.session_id,
AcpSession {
thread: thread.downgrade(),
suppress_abort_err: false,
session_modes: modes,
config_options: config_options.map(ConfigOptions::new),
ref_count: 1,
},
);
cx.update(|cx| {
self.register_session(
response.session_id,
&thread,
modes,
config_options.map(ConfigOptions::new),
cx,
)
});
Ok(thread)
})
@@ -1810,76 +1845,6 @@ impl AgentConnection for AcpConnection {
)
}
fn supports_close_session(&self) -> bool {
self.agent_capabilities.session_capabilities.close.is_some()
}
fn close_session(
self: Rc<Self>,
session_id: &acp::SessionId,
cx: &mut App,
) -> Task<Result<()>> {
if !self.supports_close_session() {
return Task::ready(Err(anyhow!(LoadError::Other(
"Closing sessions is not supported by this agent.".into()
))));
}
// If a load is still in flight, decrement its ref count. The pending
// entry is the source of truth for how many handles exist during a
// load, so we must tick it down here as well as the `sessions` entry
// that was pre-registered to receive history-replay notifications.
// Only once the pending ref count hits zero do we actually close the
// session; the load task will observe the missing sessions entry and
// fail with "session was closed before load completed".
let pending_ref_count = {
let mut pending_sessions = self.pending_sessions.borrow_mut();
pending_sessions.get_mut(session_id).map(|pending| {
pending.ref_count = pending.ref_count.saturating_sub(1);
pending.ref_count
})
};
match pending_ref_count {
Some(0) => {
self.pending_sessions.borrow_mut().remove(session_id);
self.sessions.borrow_mut().remove(session_id);
let conn = self.connection.clone();
let session_id = session_id.clone();
return cx.foreground_executor().spawn(async move {
conn.send_request(acp::CloseSessionRequest::new(session_id))
.block_task()
.await?;
Ok(())
});
}
Some(_) => return Task::ready(Ok(())),
None => {}
}
let mut sessions = self.sessions.borrow_mut();
let Some(session) = sessions.get_mut(session_id) else {
return Task::ready(Ok(()));
};
session.ref_count = session.ref_count.saturating_sub(1);
if session.ref_count > 0 {
return Task::ready(Ok(()));
}
sessions.remove(session_id);
drop(sessions);
let conn = self.connection.clone();
let session_id = session_id.clone();
cx.foreground_executor().spawn(async move {
conn.send_request(acp::CloseSessionRequest::new(session_id.clone()))
.block_task()
.await?;
Ok(())
})
}
fn auth_methods(&self) -> &[acp::AuthMethod] {
&self.auth_methods
}
@@ -2274,18 +2239,6 @@ pub mod test_support {
.load_session(session_id, project, work_dirs, title, cx)
}
fn supports_close_session(&self) -> bool {
self.inner.supports_close_session()
}
fn close_session(
self: Rc<Self>,
session_id: &acp::SessionId,
cx: &mut App,
) -> Task<Result<()>> {
self.inner.clone().close_session(session_id, cx)
}
fn supports_resume_session(&self) -> bool {
self.inner.supports_resume_session()
}
@@ -4035,7 +3988,9 @@ mod tests {
}
#[gpui::test]
async fn test_loaded_sessions_keep_state_until_last_close(cx: &mut gpui::TestAppContext) {
async fn test_loaded_sessions_keep_state_until_last_handle_drops(
cx: &mut gpui::TestAppContext,
) {
let (
connection,
project,
@@ -4085,36 +4040,30 @@ mod tests {
"underlying ACP load_session should be called exactly once for concurrent loads"
);
// The session has ref_count 2. The first close should not send the ACP
// close_session RPC — the session is still referenced.
cx.update(|cx| connection.clone().close_session(&session_id, cx))
.await
.expect("first close failed");
drop(first_thread);
release_dropped_entities(cx);
assert_eq!(
close_count.load(Ordering::SeqCst),
0,
"ACP close_session should not be sent while ref_count > 0"
"ACP close_session should not be sent while a thread handle is alive"
);
assert!(
connection.sessions.borrow().contains_key(&session_id),
"session should still be tracked after first close"
"session should still be tracked while a thread handle is alive"
);
// The second close drops ref_count to 0 — now the ACP RPC must be sent.
cx.update(|cx| connection.clone().close_session(&session_id, cx))
.await
.expect("second close failed");
cx.run_until_parked();
drop(second_thread);
release_dropped_entities(cx);
assert_eq!(
close_count.load(Ordering::SeqCst),
1,
"ACP close_session should be sent exactly once when ref_count reaches 0"
"ACP close_session should be sent exactly once when the last handle drops"
);
assert!(
!connection.sessions.borrow().contains_key(&session_id),
"session should be removed after final close"
"session should be removed after the last handle drops"
);
}
@@ -4189,13 +4138,10 @@ mod tests {
);
}
// Regression test: if `close_session` is issued while a `load_session`
// RPC is still in flight, the close must take effect cleanly — the load
// must fail with a recognizable error (not return an orphaned thread),
// no entry must remain in `sessions` or `pending_sessions`, and the ACP
// `close_session` RPC must be dispatched.
#[gpui::test]
async fn test_close_session_during_in_flight_load(cx: &mut gpui::TestAppContext) {
async fn test_dropping_loader_during_in_flight_load_closes_session(
cx: &mut gpui::TestAppContext,
) {
let (
connection,
project,
@@ -4246,24 +4192,19 @@ mod tests {
"sessions entry should be pre-registered to receive replay notifications"
);
// Close the session while the load is still parked. This should take
// the pending path and dispatch the ACP close RPC.
let close_task = cx.update(|cx| connection.clone().close_session(&session_id, cx));
drop(load_task);
cx.run_until_parked();
assert_eq!(
close_count.load(Ordering::SeqCst),
0,
"ACP close_session must not be sent while the load RPC is in flight"
);
// Release the gate so the load RPC can finally respond.
gate_tx.send(()).await.expect("gate send failed");
drop(gate_tx);
let load_result = load_task.await;
close_task.await.expect("close failed");
cx.run_until_parked();
let err = load_result.expect_err("load should fail after close-during-load");
assert!(
err.to_string()
.contains("session was closed before load completed"),
"expected close-during-load error, got: {err}"
);
release_dropped_entities(cx);
assert_eq!(
close_count.load(Ordering::SeqCst),
@@ -4272,26 +4213,19 @@ mod tests {
);
assert!(
!connection.sessions.borrow().contains_key(&session_id),
"sessions entry should be removed after close-during-load"
"sessions entry should be removed once the loaded thread has no holders"
);
assert!(
!connection
.pending_sessions
.borrow()
.contains_key(&session_id),
"pending_sessions entry should be removed after close-during-load"
"pending_sessions entry should be removed once the load resolves"
);
}
// Regression test: when two concurrent `load_session` calls share a pending
// task and one of them issues `close_session` before the load RPC
// resolves, the remaining load must still succeed and the session must
// stay live. If `close_session` incorrectly short-circuits via the
// `sessions` path (removing the entry while a load is still in flight),
// the pending task will fail and both concurrent loaders will lose
// their handle.
#[gpui::test]
async fn test_close_during_load_preserves_other_concurrent_loader(
async fn test_dropping_one_loader_during_load_preserves_other_concurrent_loader(
cx: &mut gpui::TestAppContext,
) {
let (
@@ -4340,34 +4274,29 @@ mod tests {
"load_session RPC should only be dispatched once for concurrent loads"
);
// Close one of the two handles while the shared load is still parked.
// Because a second loader still holds a pending ref, this should be a
// no-op on the wire.
cx.update(|cx| connection.clone().close_session(&session_id, cx))
.await
.expect("close during load failed");
drop(first_load);
cx.run_until_parked();
assert_eq!(
close_count.load(Ordering::SeqCst),
0,
"close_session RPC must not be dispatched while another load handle remains"
"close_session RPC must not be dispatched while another loader remains"
);
// Release the gate so the load RPC can finally respond.
gate_tx.send(()).await.expect("gate send failed");
drop(gate_tx);
let first_thread = first_load.await.expect("first load should still succeed");
let second_thread = second_load.await.expect("second load should still succeed");
cx.run_until_parked();
assert_eq!(
first_thread.entity_id(),
second_thread.entity_id(),
"concurrent loads should share one AcpThread"
close_count.load(Ordering::SeqCst),
0,
"close_session RPC must not be dispatched while a thread handle is alive"
);
assert!(
connection.sessions.borrow().contains_key(&session_id),
"session must remain tracked while a load handle is still outstanding"
"session must remain tracked while a thread handle is alive"
);
assert!(
!connection
@@ -4377,11 +4306,8 @@ mod tests {
"pending_sessions entry should be cleared once the load resolves"
);
// Final close drops ref_count to 0 and dispatches the ACP close RPC.
cx.update(|cx| connection.clone().close_session(&session_id, cx))
.await
.expect("final close failed");
cx.run_until_parked();
drop(second_thread);
release_dropped_entities(cx);
assert_eq!(
close_count.load(Ordering::SeqCst),
1,
@@ -4389,9 +4315,14 @@ mod tests {
);
assert!(
!connection.sessions.borrow().contains_key(&session_id),
"session should be removed after final close"
"session should be removed after the last handle is released"
);
}
fn release_dropped_entities(cx: &mut gpui::TestAppContext) {
cx.update(|_| ());
cx.run_until_parked();
}
}
fn mcp_servers_for_project(project: &Entity<Project>, cx: &App) -> Vec<acp::McpServer> {
crates/agent_ui/src/conversation_view.rs
+8 -1
View File
@@ -4323,10 +4323,12 @@ pub(crate) mod tests {
"Conversation should transition to LoadError when an ACP thread exits"
);
});
release_dropped_entities(cx);
assert_eq!(
close_session_count.load(std::sync::atomic::Ordering::SeqCst),
1,
"ConversationView should close the ACP session after a thread exit"
"dropping the thread views after a thread exit should close the ACP session"
);
}
@@ -6671,6 +6673,11 @@ pub(crate) mod tests {
cx.read(|cx| thread.read(cx).message_editor.clone())
}
fn release_dropped_entities(cx: &mut VisualTestContext) {
cx.update(|_, _| ());
cx.run_until_parked();
}
#[gpui::test]
async fn test_rewind_views(cx: &mut TestAppContext) {
init_test(cx);
crates/zed/src/reliability.rs
+38 -11
View File
@@ -2,7 +2,7 @@ use anyhow::{Context as _, Result};
use client::{Client, telemetry::MINIDUMP_ENDPOINT};
use feature_flags::FeatureFlagAppExt;
use futures::{AsyncReadExt, TryStreamExt};
use gpui::{App, AppContext, Entity, TaskExt};
use gpui::{App, AppContext, Entity, TaskExt, WeakEntity};
use http_client::{AsyncBody, HttpClient, Request};
use project::{Project, worktree_store::WorktreeStoreDiagnostics};
use proto::{CrashReport, GetCrashFilesResponse};
@@ -13,9 +13,11 @@ use reqwest::{
use serde::Deserialize;
use smol::stream::StreamExt;
use std::{
cell::RefCell,
collections::HashSet,
ffi::OsStr,
fs,
rc::Rc,
sync::Arc,
time::{Duration, Instant},
};
@@ -25,9 +27,12 @@ use workspace::WorkspaceStore;
mod hang_detection;
type ProjectRegistry = Rc<RefCell<Vec<WeakEntity<Project>>>>;
pub fn init(client: Arc<Client>, workspace_store: Entity<WorkspaceStore>, cx: &mut App) {
hang_detection::start(client.clone(), cx);
start_memory_usage_logging(workspace_store, cx);
let projects = ProjectRegistry::default();
start_memory_usage_logging(workspace_store, projects.clone(), cx);
cx.on_flags_ready({
let client = client.clone();
@@ -52,6 +57,7 @@ pub fn init(client: Arc<Client>, workspace_store: Entity<WorkspaceStore>, cx: &m
}
cx.observe_new(move |project: &mut Project, _, cx| {
projects.borrow_mut().push(cx.weak_entity());
let client = client.clone();
let Some(remote_client) = project.remote_client() else {
@@ -99,11 +105,15 @@ const MEMORY_USAGE_MINIMUM_LOGGED_DELTA: u64 = 64 * 1024 * 1024;
/// Logs on a fixed heartbeat, and additionally whenever resident memory changed
/// significantly since the last logged value, so that bursts of growth are timestamped
/// against the surrounding log entries.
fn start_memory_usage_logging(workspace_store: Entity<WorkspaceStore>, cx: &App) {
fn start_memory_usage_logging(
workspace_store: Entity<WorkspaceStore>,
projects: ProjectRegistry,
cx: &App,
) {
let (diagnostics_sender, mut diagnostics_receiver) = futures::channel::mpsc::unbounded();
cx.spawn(async move |cx| {
while diagnostics_receiver.next().await.is_some() {
cx.update(|cx| log_worktree_diagnostics(&workspace_store, cx));
cx.update(|cx| log_worktree_diagnostics(&workspace_store, &projects, cx));
}
})
.detach();
@@ -157,18 +167,35 @@ fn start_memory_usage_logging(workspace_store: Entity<WorkspaceStore>, cx: &App)
.detach();
}
fn log_worktree_diagnostics(workspace_store: &Entity<WorkspaceStore>, cx: &App) {
let workspaces = workspace_store
fn log_worktree_diagnostics(
workspace_store: &Entity<WorkspaceStore>,
projects: &ProjectRegistry,
cx: &App,
) {
let workspace_project_ids = workspace_store
.read(cx)
.workspaces()
.filter_map(|workspace| workspace.upgrade())
.collect::<Vec<_>>();
.map(|workspace| workspace.read(cx).project().entity_id())
.collect::<HashSet<_>>();
let live_projects = {
let mut projects = projects.borrow_mut();
projects.retain(|project| project.upgrade().is_some());
projects
.iter()
.filter_map(|project| project.upgrade())
.collect::<Vec<_>>()
};
let project_count = live_projects.len();
let orphaned_project_count = live_projects
.iter()
.filter(|project| !workspace_project_ids.contains(&project.entity_id()))
.count();
let mut worktree_store_ids = HashSet::new();
let mut store_count = 0;
let mut aggregate = WorktreeStoreDiagnostics::default();
for workspace in workspaces {
let project = workspace.read(cx).project().clone();
for project in live_projects {
let worktree_store = project.read(cx).worktree_store();
if !worktree_store_ids.insert(worktree_store.entity_id()) {
continue;
@@ -218,13 +245,13 @@ fn log_worktree_diagnostics(workspace_store: &Entity<WorkspaceStore>, cx: &App)
} = aggregate;
match largest_worktree {
Some(largest_worktree) => log::info!(
"worktree diagnostics: stores {store_count}, slots {worktree_slots}, live {live_worktrees}, visible {visible_worktrees}, strong {strong_handles}, dead weak {dead_weak_handles}, loading {loading_worktrees}, entries {total_entries}, visible entries {visible_entries}, largest {} ({} entries, {} visible)",
"worktree diagnostics: projects {project_count}, orphaned projects {orphaned_project_count}, stores {store_count}, slots {worktree_slots}, live {live_worktrees}, visible {visible_worktrees}, strong {strong_handles}, dead weak {dead_weak_handles}, loading {loading_worktrees}, entries {total_entries}, visible entries {visible_entries}, largest {} ({} entries, {} visible)",
largest_worktree.path.display(),
largest_worktree.entries,
largest_worktree.visible_entries,
),
None => log::info!(
"worktree diagnostics: stores {store_count}, slots {worktree_slots}, live {live_worktrees}, visible {visible_worktrees}, strong {strong_handles}, dead weak {dead_weak_handles}, loading {loading_worktrees}, entries {total_entries}, visible entries {visible_entries}, largest none",
"worktree diagnostics: projects {project_count}, orphaned projects {orphaned_project_count}, stores {store_count}, slots {worktree_slots}, live {live_worktrees}, visible {visible_worktrees}, strong {strong_handles}, dead weak {dead_weak_handles}, loading {loading_worktrees}, entries {total_entries}, visible entries {visible_entries}, largest none",
),
}
}