krisyotam/zed

crates/client/src/client.rs+29−0

@@ -24,6 +24,7 @@ use futures::{ · 23 unmodified lines
24	AsyncReadExt, FutureExt, SinkExt, Stream, StreamExt, TryFutureExt as _, TryStreamExt,	24	AsyncReadExt, FutureExt, SinkExt, Stream, StreamExt, TryFutureExt as _, TryStreamExt,
25	channel::{mpsc, oneshot},	25	channel::{mpsc, oneshot},
26	future::BoxFuture,	26	future::BoxFuture,
		27	stream::BoxStream,
27	};	28	};
28	use gpui::{App, AsyncApp, Entity, Global, Task, WeakEntity, actions};	29	use gpui::{App, AsyncApp, Entity, Global, Task, WeakEntity, actions};
29	use http_client::{HttpClient, HttpClientWithUrl, http, read_proxy_from_env};	30	use http_client::{HttpClient, HttpClientWithUrl, http, read_proxy_from_env};
@@ -1789,6 +1790,34 @@ impl ProtoClient for Client { · 1759 unmodified lines
1789	self.request_dynamic(envelope, request_type).boxed()	1790	self.request_dynamic(envelope, request_type).boxed()
1790	}	1791	}
1791		1792
		1793	fn request_stream(
		1794	&self,
		1795	envelope: proto::Envelope,
		1796	request_type: &'static str,
		1797	) -> BoxFuture<'static, Result<BoxStream<'static, Result<proto::Envelope>>>> {
		1798	let client_id = self.id();
		1799	let response = self.connection_id().map(\|connection_id\| {
		1800	self.peer
		1801	.request_stream_dynamic(connection_id, envelope, request_type)
		1802	});
		1803
		1804	async move {
		1805	log::debug!(
		1806	"rpc stream request start. client_id:{}. name:{}",
		1807	client_id,
		1808	request_type
		1809	);
		1810	let response = response?.await;
		1811	log::debug!(
		1812	"rpc stream request opened. client_id:{}. name:{}",
		1813	client_id,
		1814	request_type
		1815	);
		1816	response
		1817	}
		1818	.boxed()
		1819	}
		1820
1792	fn send(&self, envelope: proto::Envelope, message_type: &'static str) -> Result<()> {	1821	fn send(&self, envelope: proto::Envelope, message_type: &'static str) -> Result<()> {
1793	log::debug!("rpc send. client_id:{}, name:{}", self.id(), message_type);	1822	log::debug!("rpc send. client_id:{}, name:{}", self.id(), message_type);
1794	let connection_id = self.connection_id()?;	1823	let connection_id = self.connection_id()?;

crates/fs/src/fake_git_repo.rs+36−3

@@ -1435,10 +1435,43 @@ impl GitRepository for FakeGitRepository { · 1434 unmodified lines
1435	fn search_commits(	1435	fn search_commits(
1436	&self,	1436	&self,
1437	_log_source: LogSource,	1437	_log_source: LogSource,
1438	_search_args: SearchCommitArgs,	1438	search_args: SearchCommitArgs,
1439	_request_tx: Sender<Oid>,	1439	request_tx: Sender<Oid>,
1440	) -> BoxFuture<'_, Result<()>> {	1440	) -> BoxFuture<'_, Result<()>> {
1441	async { bail!("search_commits not supported for FakeGitRepository") }.boxed()	1441	async move {
		1442	let query = if search_args.case_sensitive {
		1443	search_args.query.to_string()
		1444	} else {
		1445	search_args.query.to_lowercase()
		1446	};
		1447
		1448	let matching_shas = self.fs.with_git_state(&self.dot_git_path, false, \|state\| {
		1449	state
		1450	.commit_data
		1451	.iter()
		1452	.filter_map(\|(sha, entry)\| {
		1453	let FakeCommitDataEntry::Success(commit_data) = entry else {
		1454	return None;
		1455	};
		1456	let message = if search_args.case_sensitive {
		1457	commit_data.message.to_string()
		1458	} else {
		1459	commit_data.message.to_lowercase()
		1460	};
		1461	message.contains(&query).then_some(*sha)
		1462	})
		1463	.collect::<Vec<_>>()
		1464	})?;
		1465
		1466	for sha in matching_shas {
		1467	if request_tx.send(sha).await.is_err() {
		1468	break;
		1469	}
		1470	}
		1471
		1472	Ok(())
		1473	}
		1474	.boxed()
1442	}	1475	}
1443		1476
1444	fn commit_data_reader(&self) -> Result<CommitDataReader> {	1477	fn commit_data_reader(&self) -> Result<CommitDataReader> {

crates/project/src/git_store.rs+129−3

@@ -20,7 +20,7 @@ use collections::HashMap; · 19 unmodified lines
20	pub use conflict_set::{ConflictRegion, ConflictSet, ConflictSetSnapshot, ConflictSetUpdate};	20	pub use conflict_set::{ConflictRegion, ConflictSet, ConflictSetSnapshot, ConflictSetUpdate};
21	use fs::{Fs, RemoveOptions};	21	use fs::{Fs, RemoveOptions};
22	use futures::{	22	use futures::{
23	FutureExt, StreamExt,	23	FutureExt, SinkExt, Stream, StreamExt,
24	channel::{	24	channel::{
25	mpsc,	25	mpsc,
26	oneshot::{self, Canceled},	26	oneshot::{self, Canceled},
@@ -680,6 +680,7 @@ impl GitStore { · 653 unmodified lines
680	client.add_entity_request_handler(Self::handle_edit_ref);	680	client.add_entity_request_handler(Self::handle_edit_ref);
681	client.add_entity_request_handler(Self::handle_repair_worktrees);	681	client.add_entity_request_handler(Self::handle_repair_worktrees);
682	client.add_entity_request_handler(Self::handle_get_commit_data);	682	client.add_entity_request_handler(Self::handle_get_commit_data);
		683	client.add_entity_stream_request_handler(Self::handle_search_commits);
683	}	684	}
684		685
685	pub fn is_local(&self) -> bool {	686	pub fn is_local(&self) -> bool {
@@ -2669,6 +2670,63 @@ impl GitStore { · 1983 unmodified lines
2669	Ok(proto::GetCommitDataResponse { commits })	2670	Ok(proto::GetCommitDataResponse { commits })
2670	}	2671	}
2671		2672
		2673	async fn handle_search_commits(
		2674	this: Entity<Self>,
		2675	envelope: TypedEnvelope<proto::SearchCommits>,
		2676	mut cx: AsyncApp,
		2677	) -> Result<impl Stream<Item = Result<proto::SearchCommitsResponse>>> {
		2678	const CHUNK_SIZE: usize = 100;
		2679
		2680	let repository_id = RepositoryId::from_proto(envelope.payload.repository_id);
		2681	let repository_handle = Self::repository_for_request(&this, repository_id, &mut cx)?;
		2682	let log_source = log_source_from_proto(
		2683	envelope
		2684	.payload
		2685	.log_source
		2686	.context("missing search commit log source")?,
		2687	)?;
		2688	let search_args = SearchCommitArgs {
		2689	query: SharedString::from(envelope.payload.query),
		2690	case_sensitive: envelope.payload.case_sensitive,
		2691	};
		2692
		2693	let (request_tx, request_rx) = async_channel::unbounded();
		2694	repository_handle.update(&mut cx, \|repository, cx\| {
		2695	repository.search_commits(log_source, search_args, request_tx, cx);
		2696	});
		2697
		2698	let (mut response_tx, response_rx) = mpsc::unbounded();
		2699	cx.background_spawn(async move {
		2700	let mut shas = Vec::new();
		2701
		2702	while let Ok(sha) = request_rx.recv().await {
		2703	shas.push(sha.to_string());
		2704
		2705	if shas.len() >= CHUNK_SIZE {
		2706	if response_tx
		2707	.send(Ok(proto::SearchCommitsResponse {
		2708	shas: mem::take(&mut shas),
		2709	}))
		2710	.await
		2711	.is_err()
		2712	{
		2713	return;
		2714	}
		2715	}
		2716	}
		2717
		2718	if !shas.is_empty() {
		2719	response_tx
		2720	.send(Ok(proto::SearchCommitsResponse { shas }))
		2721	.await
		2722	.ok();
		2723	}
		2724	})
		2725	.detach();
		2726
		2727	Ok(response_rx)
		2728	}
		2729
2672	async fn handle_edit_ref(	2730	async fn handle_edit_ref(
2673	this: Entity<Self>,	2731	this: Entity<Self>,
2674	envelope: TypedEnvelope<proto::GitEditRef>,	2732	envelope: TypedEnvelope<proto::GitEditRef>,
@@ -4974,6 +5032,7 @@ impl Repository { · 2299 unmodified lines
4974	cx: &mut Context<Self>,	5032	cx: &mut Context<Self>,
4975	) {	5033	) {
4976	let repository_state = self.repository_state.clone();	5034	let repository_state = self.repository_state.clone();
		5035	let repository_id = self.id;
4977		5036
4978	cx.background_spawn(async move {	5037	cx.background_spawn(async move {
4979	let repo_state = repository_state.await;	5038	let repo_state = repository_state.await;
@@ -4985,8 +5044,50 @@ impl Repository { · 5 unmodified lines
4985	.await	5044	.await
4986	.log_err();	5045	.log_err();
4987	}	5046	}
4988	Ok(RepositoryState::Remote(_)) => {}	5047
4989	Err(_) => {}	5048	Ok(RepositoryState::Remote(RemoteRepositoryState { client, project_id })) => {
		5049	let result = client
		5050	.request_stream(proto::SearchCommits {
		5051	project_id: project_id.to_proto(),
		5052	repository_id: repository_id.to_proto(),
		5053	log_source: Some(log_source_to_proto(&log_source)),
		5054	query: search_args.query.to_string(),
		5055	case_sensitive: search_args.case_sensitive,
		5056	})
		5057	.await;
		5058
		5059	let mut stream = match result {
		5060	Ok(stream) => stream,
		5061	Err(error) => {
		5062	log::error!("failed to search commits remotely: {error:?}");
		5063	return;
		5064	}
		5065	};
		5066
		5067	while let Some(response) = stream.next().await {
		5068	let response = match response {
		5069	Ok(response) => response,
		5070	Err(error) => {
		5071	log::error!(
		5072	"failed to receive remote commit search results: {error:?}"
		5073	);
		5074	return;
		5075	}
		5076	};
		5077
		5078	for sha in &response.shas {
		5079	let Ok(oid) = Oid::from_str(sha) else {
		5080	return;
		5081	};
		5082	if request_tx.send(oid).await.is_err() {
		5083	return;
		5084	}
		5085	}
		5086	}
		5087	}
		5088	Err(error) => {
		5089	log::error!("failed to get repository state for commit search: {error}");
		5090	}
4990	};	5091	};
4991	})	5092	})
4992	.detach();	5093	.detach();
@@ -8119,6 +8220,31 @@ fn deserialize_blame_buffer_response( · 3126 unmodified lines
8119	Some(Blame { entries, messages })	8220	Some(Blame { entries, messages })
8120	}	8221	}
8121		8222
		8223	fn log_source_to_proto(log_source: &LogSource) -> proto::GitLogSource {
		8224	proto::GitLogSource {
		8225	source: Some(match log_source {
		8226	LogSource::All => proto::git_log_source::Source::All(proto::GitLogSourceAll {}),
		8227	LogSource::Branch(branch) => proto::git_log_source::Source::Branch(branch.to_string()),
		8228	LogSource::Sha(sha) => proto::git_log_source::Source::Sha(sha.to_string()),
		8229	LogSource::Path(path) => proto::git_log_source::Source::Path(path.to_proto()),
		8230	}),
		8231	}
		8232	}
		8233
		8234	fn log_source_from_proto(log_source: proto::GitLogSource) -> Result<LogSource> {
		8235	match log_source
		8236	.source
		8237	.context("git log source is missing source")?
		8238	{
		8239	proto::git_log_source::Source::All(_) => Ok(LogSource::All),
		8240	proto::git_log_source::Source::Branch(branch) => Ok(LogSource::Branch(branch.into())),
		8241	proto::git_log_source::Source::Sha(sha) => Ok(LogSource::Sha(Oid::from_str(&sha)?)),
		8242	proto::git_log_source::Source::Path(path) => {
		8243	Ok(LogSource::Path(RepoPath::from_proto(&path)?))
		8244	}
		8245	}
		8246	}
		8247
8122	fn commit_data_to_proto(commit: &CommitData) -> proto::CommitData {	8248	fn commit_data_to_proto(commit: &CommitData) -> proto::CommitData {
8123	proto::CommitData {	8249	proto::CommitData {
8124	sha: commit.sha.to_string(),	8250	sha: commit.sha.to_string(),

crates/proto/proto/git.proto+23−0

@@ -693,3 +693,26 @@ message CommitData { · 692 unmodified lines
693	message GetCommitDataResponse {	693	message GetCommitDataResponse {
694	repeated CommitData commits = 1;	694	repeated CommitData commits = 1;
695	}	695	}
		696
		697	message GitLogSourceAll {}
		698
		699	message GitLogSource {
		700	oneof source {
		701	GitLogSourceAll all = 1;
		702	string branch = 2;
		703	string sha = 3;
		704	string path = 4;
		705	}
		706	}
		707
		708	message SearchCommits {
		709	uint64 project_id = 1;
		710	uint64 repository_id = 2;
		711	GitLogSource log_source = 3;
		712	string query = 4;
		713	bool case_sensitive = 5;
		714	}
		715
		716	message SearchCommitsResponse {
		717	repeated string shas = 1;
		718	}

crates/proto/proto/zed.proto+3−1

     GitCreateArchiveCheckpointResponse git_create_archive_checkpoint_response = 445;
     GitRestoreArchiveCheckpoint git_restore_archive_checkpoint = 446;
     GetCommitData get_commit_data = 447;
-    GetCommitDataResponse get_commit_data_response = 448; // current max
+    GetCommitDataResponse get_commit_data_response = 448;
+    SearchCommits search_commits = 449;
+    SearchCommitsResponse search_commits_response = 450; // current max
   }
   reserved 87 to 88;

crates/proto/src/proto.rs+4−0

     (GitRepairWorktrees, Background),
     (GetCommitData, Background),
     (GetCommitDataResponse, Background),
+    (SearchCommits, Background),
+    (SearchCommitsResponse, Background),
     (GitWorktreesResponse, Background),
     (GitCreateWorktree, Background),
     (GitRemoveWorktree, Background),
     (GitEditRef, Ack),
     (GitRepairWorktrees, Ack),
     (GetCommitData, GetCommitDataResponse),
+    (SearchCommits, SearchCommitsResponse),
     (GitCreateWorktree, Ack),
     (GitRemoveWorktree, Ack),
     (GitRenameWorktree, Ack),
     GitEditRef,
     GitRepairWorktrees,
     GetCommitData,
+    SearchCommits,
     GitCreateArchiveCheckpoint,
     GitRestoreArchiveCheckpoint,
     GitCreateWorktree,

crates/remote/src/remote_client.rs+225−3

@@ -22,6 +22,7 @@ use futures::{ · 21 unmodified lines
22	},	22	},
23	future::{BoxFuture, Shared, WeakShared},	23	future::{BoxFuture, Shared, WeakShared},
24	select, select_biased,	24	select, select_biased,
		25	stream::BoxStream,
25	};	26	};
26	use gpui::{	27	use gpui::{
27	App, AppContext as _, AsyncApp, BackgroundExecutor, BorrowAppContext, Context, Entity,	28	App, AppContext as _, AsyncApp, BackgroundExecutor, BorrowAppContext, Context, Entity,
@@ -1320,6 +1321,8 @@ impl RemoteConnectionOptions { · 1292 unmodified lines
1320	#[cfg(test)]	1321	#[cfg(test)]
1321	mod tests {	1322	mod tests {
1322	use super::*;	1323	use super::*;
		1324	use gpui::TestAppContext;
		1325	use rpc::{ErrorCodeExt, proto::ErrorCode};
1323		1326
1324	#[test]	1327	#[test]
1325	fn test_ssh_display_name_prefers_nickname() {	1328	fn test_ssh_display_name_prefers_nickname() {
@@ -1341,6 +1344,137 @@ mod tests { · 15 unmodified lines
1341		1344
1342	assert_eq!(options.display_name(), "1.2.3.4");	1345	assert_eq!(options.display_name(), "1.2.3.4");
1343	}	1346	}
		1347
		1348	#[gpui::test]
		1349	async fn test_channel_client_request_stream_terminates_on_error(cx: &mut TestAppContext) {
		1350	let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
		1351	let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded::<Envelope>();
		1352
		1353	let client =
		1354	cx.update(\|cx\| ChannelClient::new(incoming_rx, outgoing_tx, cx, "test-client", false));
		1355
		1356	// The client sends RemoteStarted on startup; drain the outgoing channel
		1357	// so it doesn't block.
		1358	let _drain_outgoing = cx
		1359	.executor()
		1360	.spawn(async move { while outgoing_rx.next().await.is_some() {} });
		1361
		1362	let mut stream = client
		1363	.request_stream_dynamic(proto::Test { id: 0 }.into_envelope(0, None, None), "Test")
		1364	.await
		1365	.unwrap();
		1366
		1367	let request_id = 0;
		1368
		1369	incoming_tx
		1370	.unbounded_send(proto::Test { id: 1 }.into_envelope(100, Some(request_id), None))
		1371	.unwrap();
		1372
		1373	let first = stream.next().await.unwrap().unwrap();
		1374	assert_eq!(
		1375	proto::Test::from_envelope(first).unwrap(),
		1376	proto::Test { id: 1 }
		1377	);
		1378
		1379	// Send an Error without a trailing EndStream. The Error alone should
		1380	// terminate the stream.
		1381	incoming_tx
		1382	.unbounded_send(
		1383	ErrorCode::Internal
		1384	.message("boom".to_string())
		1385	.to_proto()
		1386	.into_envelope(101, Some(request_id), None),
		1387	)
		1388	.unwrap();
		1389
		1390	let second = stream.next().await.unwrap();
		1391	let error = second.unwrap_err();
		1392	assert!(
		1393	format!("{error}").contains("boom"),
		1394	"expected error to surface server message, got: {error}"
		1395	);
		1396
		1397	assert!(stream.next().await.is_none());
		1398	assert_eq!(client.stream_response_channels.lock().len(), 0);
		1399	}
		1400
		1401	#[gpui::test]
		1402	async fn test_channel_client_dropping_stream_request_before_response_cleans_up_channel(
		1403	cx: &mut TestAppContext,
		1404	) {
		1405	let (_incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
		1406	let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded::<Envelope>();
		1407
		1408	let client =
		1409	cx.update(\|cx\| ChannelClient::new(incoming_rx, outgoing_tx, cx, "test-client", false));
		1410
		1411	let _drain_outgoing = cx
		1412	.executor()
		1413	.spawn(async move { while outgoing_rx.next().await.is_some() {} });
		1414
		1415	let stream = client
		1416	.request_stream_dynamic(proto::Test { id: 0 }.into_envelope(0, None, None), "Test")
		1417	.await
		1418	.unwrap();
		1419
		1420	assert_eq!(client.stream_response_channels.lock().len(), 1);
		1421
		1422	drop(stream);
		1423	cx.run_until_parked();
		1424
		1425	assert_eq!(
		1426	client.stream_response_channels.lock().len(),
		1427	0,
		1428	"dropping a stream before any responses arrive should remove response channel bookkeeping"
		1429	);
		1430	}
		1431
		1432	#[gpui::test]
		1433	async fn test_channel_client_dropping_stream_request_before_completion(
		1434	cx: &mut TestAppContext,
		1435	) {
		1436	let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
		1437	let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded::<Envelope>();
		1438
		1439	let client =
		1440	cx.update(\|cx\| ChannelClient::new(incoming_rx, outgoing_tx, cx, "test-client", false));
		1441
		1442	let _drain_outgoing = cx
		1443	.executor()
		1444	.spawn(async move { while outgoing_rx.next().await.is_some() {} });
		1445
		1446	let mut stream = client
		1447	.request_stream_dynamic(proto::Test { id: 0 }.into_envelope(0, None, None), "Test")
		1448	.await
		1449	.unwrap();
		1450
		1451	let request_id = 0;
		1452
		1453	incoming_tx
		1454	.unbounded_send(proto::Test { id: 1 }.into_envelope(100, Some(request_id), None))
		1455	.unwrap();
		1456	let _ = stream.next().await.unwrap().unwrap();
		1457
		1458	assert_eq!(client.stream_response_channels.lock().len(), 1);
		1459
		1460	drop(stream);
		1461
		1462	// Inject an orphaned non-terminal response. The read loop should detect
		1463	// that the consumer has been dropped and clean up its bookkeeping (no
		1464	// EndStream sent here on purpose, otherwise the cleanup would happen
		1465	// via the terminal-response path and mask the bug under test).
		1466	incoming_tx
		1467	.unbounded_send(proto::Test { id: 2 }.into_envelope(101, Some(request_id), None))
		1468	.unwrap();
		1469
		1470	cx.run_until_parked();
		1471
		1472	assert_eq!(
		1473	client.stream_response_channels.lock().len(),
		1474	0,
		1475	"stream channel should be removed once the consumer has dropped the stream"
		1476	);
		1477	}
1344	}	1478	}
1345		1479
1346	impl From<SshConnectionOptions> for RemoteConnectionOptions {	1480	impl From<SshConnectionOptions> for RemoteConnectionOptions {
@@ -1418,6 +1552,8 @@ pub trait RemoteConnection: Send + Sync { · 71 unmodified lines
1418	}	1552	}
1419		1553
1420	type ResponseChannels = Mutex<HashMap<MessageId, oneshot::Sender<(Envelope, oneshot::Sender<()>)>>>;	1554	type ResponseChannels = Mutex<HashMap<MessageId, oneshot::Sender<(Envelope, oneshot::Sender<()>)>>>;
		1555	type StreamResponseChannels =
		1556	Arc<Mutex<HashMap<MessageId, UnboundedSender<(Result<Envelope>, oneshot::Sender<()>)>>>>;
1421		1557
1422	struct Signal<T> {	1558	struct Signal<T> {
1423	tx: Mutex<Option<oneshot::Sender<T>>>,	1559	tx: Mutex<Option<oneshot::Sender<T>>>,
@@ -1455,6 +1591,7 @@ pub(crate) struct ChannelClient { · 31 unmodified lines
1455	outgoing_tx: Mutex<mpsc::UnboundedSender<Envelope>>,	1591	outgoing_tx: Mutex<mpsc::UnboundedSender<Envelope>>,
1456	buffer: Mutex<VecDeque<Envelope>>,	1592	buffer: Mutex<VecDeque<Envelope>>,
1457	response_channels: ResponseChannels,	1593	response_channels: ResponseChannels,
		1594	stream_response_channels: StreamResponseChannels,
1458	message_handlers: Mutex<ProtoMessageHandlerSet>,	1595	message_handlers: Mutex<ProtoMessageHandlerSet>,
1459	max_received: AtomicU32,	1596	max_received: AtomicU32,
1460	name: &'static str,	1597	name: &'static str,
@@ -1477,6 +1614,7 @@ impl ChannelClient { · 16 unmodified lines
1477	next_message_id: AtomicU32::new(0),	1614	next_message_id: AtomicU32::new(0),
1478	max_received: AtomicU32::new(0),	1615	max_received: AtomicU32::new(0),
1479	response_channels: ResponseChannels::default(),	1616	response_channels: ResponseChannels::default(),
		1617	stream_response_channels: StreamResponseChannels::default(),
1480	message_handlers: Default::default(),	1618	message_handlers: Default::default(),
1481	buffer: Mutex::new(VecDeque::new()),	1619	buffer: Mutex::new(VecDeque::new()),
1482	name,	1620	name,
@@ -1550,13 +1688,40 @@ impl ChannelClient { · 67 unmodified lines
1550		1688
1551	if let Some(request_id) = incoming.responding_to {	1689	if let Some(request_id) = incoming.responding_to {
1552	let request_id = MessageId(request_id);	1690	let request_id = MessageId(request_id);
		1691	// An incoming response with no payload is malformed; drop
		1692	// it. The request future and any stream consumers will
		1693	// remain pending until either a real response arrives or
		1694	// the connection is torn down.
		1695	if incoming.payload.is_none() {
		1696	continue;
		1697	}
1553	let sender = this.response_channels.lock().remove(&request_id);	1698	let sender = this.response_channels.lock().remove(&request_id);
1554	if let Some(sender) = sender {	1699	if let Some(sender) = sender {
1555	let (tx, rx) = oneshot::channel();	1700	let (tx, rx) = oneshot::channel();
1556	if incoming.payload.is_some() {	1701	sender.send((incoming, tx)).ok();
1557	sender.send((incoming, tx)).ok();
1558	}
1559	rx.await.ok();	1702	rx.await.ok();
		1703	} else {
		1704	let terminal_stream_response = matches!(
		1705	&incoming.payload,
		1706	Some(proto::envelope::Payload::Error(_))
		1707	\| Some(proto::envelope::Payload::EndStream(_))
		1708	);
		1709	let sender = if terminal_stream_response {
		1710	this.stream_response_channels.lock().remove(&request_id)
		1711	} else {
		1712	this.stream_response_channels
		1713	.lock()
		1714	.get(&request_id)
		1715	.cloned()
		1716	};
		1717	if let Some(sender) = sender {
		1718	let (tx, rx) = oneshot::channel();
		1719	if sender.unbounded_send((Ok(incoming), tx)).is_err() {
		1720	this.stream_response_channels.lock().remove(&request_id);
		1721	continue;
		1722	}
		1723	rx.await.ok();
		1724	}
1560	}	1725	}
1561	} else if let Some(envelope) =	1726	} else if let Some(envelope) =
1562	build_typed_envelope(peer_id, Instant::now(), incoming)	1727	build_typed_envelope(peer_id, Instant::now(), incoming)
@@ -1721,6 +1886,55 @@ impl ChannelClient { · 158 unmodified lines
1721	}	1886	}
1722	}	1887	}
1723		1888
		1889	fn request_stream_dynamic(
		1890	&self,
		1891	mut envelope: proto::Envelope,
		1892	type_name: &'static str,
		1893	) -> impl 'static + Future<Output = Result<BoxStream<'static, Result<proto::Envelope>>>> {
		1894	envelope.id = self.next_message_id.fetch_add(1, SeqCst);
		1895	let message_id = MessageId(envelope.id);
		1896	let (tx, rx) = mpsc::unbounded();
		1897	let stream_response_channels = self.stream_response_channels.clone();
		1898	stream_response_channels.lock().insert(message_id, tx);
		1899
		1900	let result = self.send_buffered(envelope);
		1901	async move {
		1902	if let Err(error) = &result {
		1903	log::error!("failed to send message: {error}");
		1904	anyhow::bail!("failed to send message: {error}");
		1905	}
		1906
		1907	let cleanup_stream_response_channel = util::defer({
		1908	let stream_response_channels = stream_response_channels.clone();
		1909	move \|\| {
		1910	stream_response_channels.lock().remove(&message_id);
		1911	}
		1912	});
		1913
		1914	Ok(rx
		1915	.filter_map(move \|(response, _barrier)\| {
		1916	// Keep the cleanup guard alive until the returned stream is dropped.
		1917	let _keep_cleanup_guard_alive = &cleanup_stream_response_channel;
		1918	futures::future::ready(match response {
		1919	Ok(response) => {
		1920	if let Some(proto::envelope::Payload::Error(error)) = &response.payload
		1921	{
		1922	Some(Err(RpcError::from_proto(error, type_name)))
		1923	} else if let Some(proto::envelope::Payload::EndStream(_)) =
		1924	&response.payload
		1925	{
		1926	None
		1927	} else {
		1928	Some(Ok(response))
		1929	}
		1930	}
		1931	Err(error) => Some(Err(error)),
		1932	})
		1933	})
		1934	.boxed())
		1935	}
		1936	}
		1937
1724	pub fn send_dynamic(&self, mut envelope: proto::Envelope) -> Result<()> {	1938	pub fn send_dynamic(&self, mut envelope: proto::Envelope) -> Result<()> {
1725	envelope.id = self.next_message_id.fetch_add(1, SeqCst);	1939	envelope.id = self.next_message_id.fetch_add(1, SeqCst);
1726	self.send_buffered(envelope)	1940	self.send_buffered(envelope)
@@ -1751,6 +1965,14 @@ impl ProtoClient for ChannelClient { · 24 unmodified lines
1751	self.request_dynamic(envelope, request_type, true).boxed()	1965	self.request_dynamic(envelope, request_type, true).boxed()
1752	}	1966	}
1753		1967
		1968	fn request_stream(
		1969	&self,
		1970	envelope: proto::Envelope,
		1971	request_type: &'static str,
		1972	) -> BoxFuture<'static, Result<BoxStream<'static, Result<proto::Envelope>>>> {
		1973	self.request_stream_dynamic(envelope, request_type).boxed()
		1974	}
		1975
1754	fn send(&self, envelope: proto::Envelope, _message_type: &'static str) -> Result<()> {	1976	fn send(&self, envelope: proto::Envelope, _message_type: &'static str) -> Result<()> {
1755	self.send_dynamic(envelope)	1977	self.send_dynamic(envelope)
1756	}	1978	}

crates/remote_server/src/remote_editing_tests.rs+108−1

@@ -11,7 +11,10 @@ use languages::rust_lang; · 10 unmodified lines
11		11
12	use extension::ExtensionHostProxy;	12	use extension::ExtensionHostProxy;
13	use fs::{FakeFs, Fs};	13	use fs::{FakeFs, Fs};
14	use git::repository::Worktree as GitWorktree;	14	use git::{
		15	Oid,
		16	repository::{CommitData, Worktree as GitWorktree},
		17	};
15	use gpui::{AppContext as _, Entity, SharedString, TestAppContext};	18	use gpui::{AppContext as _, Entity, SharedString, TestAppContext};
16	use http_client::{BlockedHttpClient, FakeHttpClient};	19	use http_client::{BlockedHttpClient, FakeHttpClient};
17	use language::{	20	use language::{
@@ -29,11 +32,13 @@ use project::{ · 11 unmodified lines
29	search::{SearchQuery, SearchResult},	32	search::{SearchQuery, SearchResult},
30	};	33	};
31	use remote::RemoteClient;	34	use remote::RemoteClient;
		35	use rpc::proto;
32	use serde_json::json;	36	use serde_json::json;
33	use settings::{Settings, SettingsLocation, SettingsStore, initial_server_settings_content};	37	use settings::{Settings, SettingsLocation, SettingsStore, initial_server_settings_content};
34	use smol::stream::StreamExt;	38	use smol::stream::StreamExt;
35	use std::{	39	use std::{
36	path::{Path, PathBuf},	40	path::{Path, PathBuf},
		41	str::FromStr,
37	sync::Arc,	42	sync::Arc,
38	};	43	};
39	use unindent::Unindent as _;	44	use unindent::Unindent as _;
@@ -1626,6 +1631,108 @@ async fn test_remote_root_repo_common_dir(cx: &mut TestAppContext, server_cx: &m · 1586 unmodified lines
1626	assert_eq!(common_dir, None);	1631	assert_eq!(common_dir, None);
1627	}	1632	}
1628		1633
		1634	#[gpui::test]
		1635	async fn test_remote_search_commits_streams_proto_chunks(
		1636	cx: &mut TestAppContext,
		1637	server_cx: &mut TestAppContext,
		1638	) {
		1639	const COMMIT_COUNT: usize = 900;
		1640	const RESPONSE_MAX_SIZE: usize = 100;
		1641
		1642	let fs = FakeFs::new(server_cx.executor());
		1643	fs.insert_tree(
		1644	path!("/code"),
		1645	json!({
		1646	"project1": {
		1647	".git": {},
		1648	"file.txt": "content",
		1649	},
		1650	}),
		1651	)
		1652	.await;
		1653
		1654	let commit_data = (0..COMMIT_COUNT)
		1655	.map(\|index\| {
		1656	let sha = Oid::from_str(&format!("{:040x}", index + 1)).unwrap();
		1657	(
		1658	CommitData {
		1659	sha,
		1660	parents: Default::default(),
		1661	author_name: SharedString::from("Author"),
		1662	author_email: SharedString::from("author@example.com"),
		1663	commit_timestamp: index as i64,
		1664	subject: SharedString::from(format!("Subject {index}")),
		1665	message: SharedString::from(format!("needle commit {index}")),
		1666	},
		1667	false,
		1668	)
		1669	})
		1670	.collect::<Vec<_>>();
		1671	let expected_shas = commit_data
		1672	.iter()
		1673	.map(\|(commit_data, _)\| commit_data.sha.to_string())
		1674	.collect::<HashSet<_>>();
		1675	fs.set_commit_data(Path::new(path!("/code/project1/.git")), commit_data);
		1676
		1677	let (project, _headless) = init_test(&fs, cx, server_cx).await;
		1678	project
		1679	.update(cx, \|project, cx\| {
		1680	project.find_or_create_worktree(path!("/code/project1"), true, cx)
		1681	})
		1682	.await
		1683	.expect("should open remote worktree");
		1684	server_cx.run_until_parked();
		1685	cx.run_until_parked();
		1686	project
		1687	.update(cx, \|project, cx\| project.git_scans_complete(cx))
		1688	.await;
		1689
		1690	let (remote_client, repository_id) = project.read_with(cx, \|project, cx\| {
		1691	let repository = project
		1692	.active_repository(cx)
		1693	.expect("remote project should have an active repository");
		1694	let repository_id = repository.read(cx).snapshot().id;
		1695	let remote_client = project
		1696	.remote_client()
		1697	.expect("project should have a remote client");
		1698	(remote_client, repository_id)
		1699	});
		1700	let proto_client = remote_client.read_with(cx, \|remote_client, _\| remote_client.proto_client());
		1701	let mut stream = proto_client
		1702	.request_stream(proto::SearchCommits {
		1703	project_id: proto::REMOTE_SERVER_PROJECT_ID,
		1704	repository_id: repository_id.to_proto(),
		1705	log_source: Some(proto::GitLogSource {
		1706	source: Some(proto::git_log_source::Source::All(
		1707	proto::GitLogSourceAll {},
		1708	)),
		1709	}),
		1710	query: "needle".to_string(),
		1711	case_sensitive: true,
		1712	})
		1713	.await
		1714	.expect("search commits stream should start");
		1715
		1716	let mut chunks = Vec::new();
		1717	while let Some(response) = futures::StreamExt::next(&mut stream).await {
		1718	chunks.push(response.expect("search commits chunk should succeed").shas);
		1719	}
		1720
		1721	assert!(
		1722	chunks.len() > 1,
		1723	"expected search results to stream in multiple chunks"
		1724	);
		1725	for chunk in chunks.iter().take(chunks.len() - 1) {
		1726	assert!(
		1727	chunk.len() <= RESPONSE_MAX_SIZE,
		1728	"non-final chunks should meet the target byte size"
		1729	);
		1730	}
		1731
		1732	let actual_shas = chunks.into_iter().flatten().collect::<HashSet<_>>();
		1733	assert_eq!(actual_shas, expected_shas);
		1734	}
		1735
1629	#[gpui::test]	1736	#[gpui::test]
1630	async fn test_remote_archive_git_operations_are_supported(	1737	async fn test_remote_archive_git_operations_are_supported(
1631	cx: &mut TestAppContext,	1738	cx: &mut TestAppContext,

crates/rpc/src/peer.rs+364−48

@@ -8,7 +8,7 @@ use super::{ · 7 unmodified lines
8	use anyhow::{Context as _, Result, anyhow};	8	use anyhow::{Context as _, Result, anyhow};
9	use collections::HashMap;	9	use collections::HashMap;
10	use futures::{	10	use futures::{
11	FutureExt, SinkExt, Stream, StreamExt, TryFutureExt,	11	FutureExt, SinkExt, StreamExt, TryFutureExt,
12	channel::{mpsc, oneshot},	12	channel::{mpsc, oneshot},
13	stream::BoxStream,	13	stream::BoxStream,
14	};	14	};
@@ -278,11 +278,23 @@ impl Peer { · 263 unmodified lines
278	);	278	);
279	let response_channel =	279	let response_channel =
280	response_channels.lock().as_mut()?.remove(&responding_to);	280	response_channels.lock().as_mut()?.remove(&responding_to);
281	let stream_response_channel = stream_response_channels	281	let terminal_stream_response = matches!(
282	.lock()	282	&incoming.payload,
283	.as_ref()?	283	Some(proto::envelope::Payload::Error(_))
284	.get(&responding_to)	284	\| Some(proto::envelope::Payload::EndStream(_))
285	.cloned();	285	);
		286	let stream_response_channel = if terminal_stream_response {
		287	stream_response_channels
		288	.lock()
		289	.as_mut()?
		290	.remove(&responding_to)
		291	} else {
		292	stream_response_channels
		293	.lock()
		294	.as_ref()?
		295	.get(&responding_to)
		296	.cloned()
		297	};
286		298
287	if let Some(tx) = response_channel {	299	if let Some(tx) = response_channel {
288	let requester_resumed = oneshot::channel();	300	let requester_resumed = oneshot::channel();
@@ -319,21 +331,15 @@ impl Peer { · 30 unmodified lines
319	?error,	331	?error,
320	"incoming stream response: request future dropped",	332	"incoming stream response: request future dropped",
321	);	333	);
		334	// The consumer has gone away, so drop the bookkeeping
		335	// for this stream rather than letting it accumulate
		336	// every subsequent message until a terminal frame.
		337	if let Some(channels) = stream_response_channels.lock().as_mut() {
		338	channels.remove(&responding_to);
		339	}
		340	} else {
		341	let _ = requester_resumed.1.await;
322	}	342	}
323
324	tracing::debug!(
325	%connection_id,
326	message_id,
327	responding_to,
328	"incoming stream response: waiting to resume requester"
329	);
330	let _ = requester_resumed.1.await;
331	tracing::debug!(
332	%connection_id,
333	message_id,
334	responding_to,
335	"incoming stream response: requester resumed"
336	);
337	} else {	343	} else {
338	let message_type = proto::build_typed_envelope(	344	let message_type = proto::build_typed_envelope(
339	connection_id.into(),	345	connection_id.into(),
@@ -484,55 +490,96 @@ impl Peer { · 144 unmodified lines
484	&self,	490	&self,
485	receiver_id: ConnectionId,	491	receiver_id: ConnectionId,
486	request: T,	492	request: T,
487	) -> impl Future<Output = Result<impl Unpin + Stream<Item = Result<T::Response>>>> {	493	) -> impl Future<Output = Result<BoxStream<'static, Result<T::Response>>>> {
		494	let stream =
		495	self.request_stream_dynamic(receiver_id, request.into_envelope(0, None, None), T::NAME);
		496
		497	async move {
		498	Ok(stream
		499	.await?
		500	.map(\|response\| {
		501	T::Response::from_envelope(response?)
		502	.context("received response of the wrong type")
		503	})
		504	.boxed())
		505	}
		506	}
		507
		508	pub fn request_stream_dynamic(
		509	&self,
		510	receiver_id: ConnectionId,
		511	mut envelope: proto::Envelope,
		512	request_type: &'static str,
		513	) -> impl Future<Output = Result<BoxStream<'static, Result<proto::Envelope>>>> + use<> {
488	let (tx, rx) = mpsc::unbounded();	514	let (tx, rx) = mpsc::unbounded();
489	let send = self.connection_state(receiver_id).and_then(\|connection\| {	515	let send = self.connection_state(receiver_id).and_then(\|connection\| {
490	let message_id = connection.next_message_id.fetch_add(1, SeqCst);	516	let message_id = connection.next_message_id.fetch_add(1, SeqCst);
		517	envelope.id = message_id;
491	let stream_response_channels = connection.stream_response_channels.clone();	518	let stream_response_channels = connection.stream_response_channels.clone();
492	stream_response_channels	519	stream_response_channels
493	.lock()	520	.lock()
494	.as_mut()	521	.as_mut()
495	.context("connection was closed")?	522	.context("connection was closed")?
496	.insert(message_id, tx);	523	.insert(message_id, tx);
497	connection	524	if let Err(error) = connection
498	.outgoing_tx	525	.outgoing_tx
499	.unbounded_send(Message::Envelope(	526	.unbounded_send(Message::Envelope(envelope))
500	request.into_envelope(message_id, None, None),	527	{
501	))	528	if let Some(channels) = stream_response_channels.lock().as_mut() {
502	.context("connection was closed")?;	529	channels.remove(&message_id);
		530	}
		531	return Err(error).context("connection was closed");
		532	}
503	Ok((message_id, stream_response_channels))	533	Ok((message_id, stream_response_channels))
504	});	534	});
505		535
506	async move {	536	async move {
507	let (message_id, stream_response_channels) = send?;	537	let (message_id, stream_response_channels) = send?;
508	let stream_response_channels = Arc::downgrade(&stream_response_channels);	538	let stream_response_channels = Arc::downgrade(&stream_response_channels);
509		539	let cleanup_stream_response_channel = util::defer({
510	Ok(rx.filter_map(move \|(response, _barrier)\| {
511	let stream_response_channels = stream_response_channels.clone();	540	let stream_response_channels = stream_response_channels.clone();
512	future::ready(match response {	541	move \|\| {
513	Ok(response) => {	542	if let Some(channels) = stream_response_channels.upgrade()
514	if let Some(proto::envelope::Payload::Error(error)) = &response.payload {	543	&& let Some(channels) = channels.lock().as_mut()
515	Some(Err(RpcError::from_proto(error, T::NAME)))	544	{
516	} else if let Some(proto::envelope::Payload::EndStream(_)) =	545	channels.remove(&message_id);
517	&response.payload	546	}
518	{	547	}
519	// Remove the transmitting end of the response channel to end the stream.	548	});
520	if let Some(channels) = stream_response_channels.upgrade()	549
521	&& let Some(channels) = channels.lock().as_mut()	550	Ok(rx
		551	.filter_map(move \|(response, _barrier)\| {
		552	let _keep_cleanup_guard_alive = &cleanup_stream_response_channel;
		553	let stream_response_channels = stream_response_channels.clone();
		554	future::ready(match response {
		555	Ok(response) => {
		556	if let Some(proto::envelope::Payload::Error(error)) = &response.payload
522	{	557	{
523	channels.remove(&message_id);	558	// Remove the transmitting end of the response channel to end the stream.
		559	if let Some(channels) = stream_response_channels.upgrade()
		560	&& let Some(channels) = channels.lock().as_mut()
		561	{
		562	channels.remove(&message_id);
		563	}
		564	Some(Err(RpcError::from_proto(error, request_type)))
		565	} else if let Some(proto::envelope::Payload::EndStream(_)) =
		566	&response.payload
		567	{
		568	// Remove the transmitting end of the response channel to end the stream.
		569	if let Some(channels) = stream_response_channels.upgrade()
		570	&& let Some(channels) = channels.lock().as_mut()
		571	{
		572	channels.remove(&message_id);
		573	}
		574	None
		575	} else {
		576	Some(Ok(response))
524	}	577	}
525	None
526	} else {
527	Some(
528	T::Response::from_envelope(response)
529	.context("received response of the wrong type"),
530	)
531	}	578	}
532	}	579	Err(error) => Some(Err(error)),
533	Err(error) => Some(Err(error)),	580	})
534	})	581	})
535	}))	582	.boxed())
536	}	583	}
537	}	584	}
538		585
@@ -661,6 +708,13 @@ impl Peer { · 122 unmodified lines
661	.with_context(\|\| format!("no such connection: {connection_id}"))?;	708	.with_context(\|\| format!("no such connection: {connection_id}"))?;
662	Ok(connection.clone())	709	Ok(connection.clone())
663	}	710	}
		711
		712	#[cfg(any(test, feature = "test-support"))]
		713	pub fn pending_stream_request_count(&self, connection_id: ConnectionId) -> Option<usize> {
		714	let connection = self.connection_state(connection_id).ok()?;
		715	let channels = connection.stream_response_channels.lock();
		716	Some(channels.as_ref()?.len())
		717	}
664	}	718	}
665		719
666	impl Serialize for Peer {	720	impl Serialize for Peer {
@@ -992,6 +1046,268 @@ mod tests { · 325 unmodified lines
992	);	1046	);
993	}	1047	}
994		1048
		1049	#[gpui::test(iterations = 50)]
		1050	async fn test_request_stream(cx: &mut TestAppContext) {
		1051	init_logger();
		1052
		1053	let executor = cx.executor();
		1054	let server = Peer::new(0);
		1055	let client = Peer::new(0);
		1056
		1057	let (client_to_server_conn, server_to_client_conn, _kill) =
		1058	Connection::in_memory(executor.clone());
		1059	let (client_to_server_conn_id, io_task1, mut client_incoming) =
		1060	client.add_test_connection(client_to_server_conn, executor.clone());
		1061	let (_, io_task2, mut server_incoming) =
		1062	server.add_test_connection(server_to_client_conn, executor.clone());
		1063
		1064	executor.spawn(io_task1).detach();
		1065	executor.spawn(io_task2).detach();
		1066	executor
		1067	.spawn(async move { while client_incoming.next().await.is_some() {} })
		1068	.detach();
		1069
		1070	executor
		1071	.spawn({
		1072	let server = server.clone();
		1073	async move {
		1074	let request = server_incoming
		1075	.next()
		1076	.await
		1077	.unwrap()
		1078	.into_any()
		1079	.downcast::<TypedEnvelope<proto::Test>>()
		1080	.unwrap();
		1081	let receipt = request.receipt();
		1082	server.respond(receipt, proto::Test { id: 1 }).unwrap();
		1083	server.respond(receipt, proto::Test { id: 2 }).unwrap();
		1084	server.respond(receipt, proto::Test { id: 3 }).unwrap();
		1085	server.end_stream(receipt).unwrap();
		1086
		1087	// Prevent the connection from being dropped.
		1088	server_incoming.next().await;
		1089	}
		1090	})
		1091	.detach();
		1092
		1093	let mut stream = client
		1094	.request_stream(client_to_server_conn_id, proto::Test { id: 0 })
		1095	.await
		1096	.unwrap();
		1097
		1098	let mut received = Vec::new();
		1099	while let Some(item) = stream.next().await {
		1100	received.push(item.unwrap());
		1101	}
		1102
		1103	assert_eq!(
		1104	received,
		1105	vec![
		1106	proto::Test { id: 1 },
		1107	proto::Test { id: 2 },
		1108	proto::Test { id: 3 },
		1109	]
		1110	);
		1111	assert_eq!(
		1112	client.pending_stream_request_count(client_to_server_conn_id),
		1113	Some(0)
		1114	);
		1115	}
		1116
		1117	#[gpui::test]
		1118	async fn test_request_stream_send_failure_cleans_up_response_channel(cx: &mut TestAppContext) {
		1119	init_logger();
		1120
		1121	let executor = cx.executor();
		1122	let client = Peer::new(0);
		1123
		1124	let (client_to_server_conn, _server_to_client_conn, _kill) =
		1125	Connection::in_memory(executor.clone());
		1126	let (client_to_server_conn_id, io_task, _client_incoming) =
		1127	client.add_test_connection(client_to_server_conn, executor.clone());
		1128
		1129	drop(io_task);
		1130
		1131	let result = client
		1132	.request_stream(client_to_server_conn_id, proto::Test { id: 0 })
		1133	.await;
		1134
		1135	assert!(
		1136	result.is_err(),
		1137	"stream request should fail when the connection write task has gone away"
		1138	);
		1139	assert_eq!(
		1140	client.pending_stream_request_count(client_to_server_conn_id),
		1141	Some(0),
		1142	"failed stream request should not leave response channel bookkeeping behind"
		1143	);
		1144	}
		1145
		1146	#[gpui::test(iterations = 50)]
		1147	async fn test_request_stream_terminates_on_error(cx: &mut TestAppContext) {
		1148	init_logger();
		1149
		1150	let executor = cx.executor();
		1151	let server = Peer::new(0);
		1152	let client = Peer::new(0);
		1153
		1154	let (client_to_server_conn, server_to_client_conn, _kill) =
		1155	Connection::in_memory(executor.clone());
		1156	let (client_to_server_conn_id, io_task1, mut client_incoming) =
		1157	client.add_test_connection(client_to_server_conn, executor.clone());
		1158	let (_, io_task2, mut server_incoming) =
		1159	server.add_test_connection(server_to_client_conn, executor.clone());
		1160
		1161	executor.spawn(io_task1).detach();
		1162	executor.spawn(io_task2).detach();
		1163	executor
		1164	.spawn(async move { while client_incoming.next().await.is_some() {} })
		1165	.detach();
		1166
		1167	executor
		1168	.spawn({
		1169	let server = server.clone();
		1170	async move {
		1171	let request = server_incoming
		1172	.next()
		1173	.await
		1174	.unwrap()
		1175	.into_any()
		1176	.downcast::<TypedEnvelope<proto::Test>>()
		1177	.unwrap();
		1178	let receipt = request.receipt();
		1179	server.respond(receipt, proto::Test { id: 1 }).unwrap();
		1180	// Send an Error without a trailing EndStream. The Error alone
		1181	// should be treated as a terminal stream response.
		1182	server
		1183	.respond_with_error(
		1184	receipt,
		1185	ErrorCode::Internal.message("boom".to_string()).to_proto(),
		1186	)
		1187	.unwrap();
		1188
		1189	// Prevent the connection from being dropped.
		1190	server_incoming.next().await;
		1191	}
		1192	})
		1193	.detach();
		1194
		1195	let mut stream = client
		1196	.request_stream(client_to_server_conn_id, proto::Test { id: 0 })
		1197	.await
		1198	.unwrap();
		1199
		1200	assert_eq!(stream.next().await.unwrap().unwrap(), proto::Test { id: 1 });
		1201
		1202	let error = stream.next().await.unwrap().unwrap_err();
		1203	assert!(
		1204	format!("{error}").contains("boom"),
		1205	"expected error to surface server message, got: {error}"
		1206	);
		1207
		1208	// The error alone (without an EndStream) should terminate the stream.
		1209	assert!(stream.next().await.is_none());
		1210	assert_eq!(
		1211	client.pending_stream_request_count(client_to_server_conn_id),
		1212	Some(0)
		1213	);
		1214	}
		1215
		1216	#[gpui::test(iterations = 50)]
		1217	async fn test_dropping_stream_request_before_completion(cx: &mut TestAppContext) {
		1218	init_logger();
		1219
		1220	let executor = cx.executor();
		1221	let server = Peer::new(0);
		1222	let client = Peer::new(0);
		1223
		1224	let (client_to_server_conn, server_to_client_conn, _kill) =
		1225	Connection::in_memory(executor.clone());
		1226	let (client_to_server_conn_id, io_task1, mut client_incoming) =
		1227	client.add_test_connection(client_to_server_conn, executor.clone());
		1228	let (_, io_task2, mut server_incoming) =
		1229	server.add_test_connection(server_to_client_conn, executor.clone());
		1230
		1231	executor.spawn(io_task1).detach();
		1232	executor.spawn(io_task2).detach();
		1233	executor
		1234	.spawn(async move { while client_incoming.next().await.is_some() {} })
		1235	.detach();
		1236
		1237	let (drop_signal_tx, drop_signal_rx) = oneshot::channel::<()>();
		1238	let server_task = executor.spawn({
		1239	let server = server.clone();
		1240	async move {
		1241	let request = server_incoming
		1242	.next()
		1243	.await
		1244	.unwrap()
		1245	.into_any()
		1246	.downcast::<TypedEnvelope<proto::Test>>()
		1247	.unwrap();
		1248	let receipt = request.receipt();
		1249	server.respond(receipt, proto::Test { id: 1 }).unwrap();
		1250
		1251	// Wait until the consumer has dropped the stream.
		1252	drop_signal_rx.await.ok();
		1253
		1254	// Send a non-terminal response after the consumer is gone. The
		1255	// peer should detect that the receiver has been dropped and clean
		1256	// up its bookkeeping. Crucially, we do NOT send EndStream here
		1257	// because that would clean up via the terminal-response path and
		1258	// mask the bug.
		1259	server.respond(receipt, proto::Test { id: 2 }).unwrap();
		1260
		1261	// A Ping/Ack round-trip after the response acts as a sync
		1262	// barrier: because messages over the in-memory connection are
		1263	// delivered in order, by the time the client observes the Ack,
		1264	// it has already processed the dropped response above.
		1265	let ping = server_incoming
		1266	.next()
		1267	.await
		1268	.unwrap()
		1269	.into_any()
		1270	.downcast::<TypedEnvelope<proto::Ping>>()
		1271	.unwrap();
		1272	server.respond(ping.receipt(), proto::Ack {}).unwrap();
		1273
		1274	// Prevent the connection from being dropped.
		1275	server_incoming.next().await;
		1276	}
		1277	});
		1278
		1279	let mut stream = client
		1280	.request_stream(client_to_server_conn_id, proto::Test { id: 0 })
		1281	.await
		1282	.unwrap();
		1283
		1284	assert_eq!(stream.next().await.unwrap().unwrap(), proto::Test { id: 1 });
		1285
		1286	// The stream is mid-flight, so the channel should be tracked.
		1287	assert_eq!(
		1288	client.pending_stream_request_count(client_to_server_conn_id),
		1289	Some(1)
		1290	);
		1291
		1292	drop(stream);
		1293	drop_signal_tx.send(()).ok();
		1294
		1295	// Synchronization barrier: once this Ack arrives, the read loop has
		1296	// already processed the orphaned stream response that came before it.
		1297	client
		1298	.request(client_to_server_conn_id, proto::Ping {})
		1299	.await
		1300	.unwrap();
		1301
		1302	assert_eq!(
		1303	client.pending_stream_request_count(client_to_server_conn_id),
		1304	Some(0),
		1305	"stream channel should be removed once the consumer has dropped the stream"
		1306	);
		1307
		1308	drop(server_task);
		1309	}
		1310
995	#[gpui::test(iterations = 50)]	1311	#[gpui::test(iterations = 50)]
996	async fn test_disconnect(cx: &mut TestAppContext) {	1312	async fn test_disconnect(cx: &mut TestAppContext) {
997	let executor = cx.executor();	1313	let executor = cx.executor();

crates/rpc/src/proto_client.rs+95−1

@@ -1,9 +1,10 @@
1	use anyhow::{Context, Result};	1	use anyhow::{Context, Result};
2	use collections::HashMap;	2	use collections::HashMap;
3	use futures::{	3	use futures::{
4	Future, FutureExt as _,	4	Future, FutureExt as _, Stream, StreamExt as _,
5	channel::oneshot,	5	channel::oneshot,
6	future::{BoxFuture, LocalBoxFuture},	6	future::{BoxFuture, LocalBoxFuture},
		7	stream::BoxStream,
7	};	8	};
8	use gpui::{AnyEntity, AnyWeakEntity, AsyncApp, BackgroundExecutor, Entity, FutureExt as _};	9	use gpui::{AnyEntity, AnyWeakEntity, AsyncApp, BackgroundExecutor, Entity, FutureExt as _};
9	use parking_lot::Mutex;	10	use parking_lot::Mutex;
@@ -61,6 +62,20 @@ pub trait ProtoClient: Send + Sync { · 51 unmodified lines
61	request_type: &'static str,	62	request_type: &'static str,
62	) -> BoxFuture<'static, Result<Envelope>>;	63	) -> BoxFuture<'static, Result<Envelope>>;
63		64
		65	fn request_stream(
		66	&self,
		67	envelope: Envelope,
		68	request_type: &'static str,
		69	) -> BoxFuture<'static, Result<BoxStream<'static, Result<Envelope>>>> {
		70	async move {
		71	anyhow::bail!(
		72	"stream requests are not supported for {request_type}: {:?}",
		73	envelope.payload
		74	)
		75	}
		76	.boxed()
		77	}
		78
64	fn send(&self, envelope: Envelope, message_type: &'static str) -> Result<()>;	79	fn send(&self, envelope: Envelope, message_type: &'static str) -> Result<()>;
65		80
66	fn send_response(&self, envelope: Envelope, message_type: &'static str) -> Result<()>;	81	fn send_response(&self, envelope: Envelope, message_type: &'static str) -> Result<()>;
@@ -223,6 +238,23 @@ impl AnyProtoClient { · 156 unmodified lines
223	}	238	}
224	}	239	}
225		240
		241	pub fn request_stream<T: RequestMessage>(
		242	&self,
		243	request: T,
		244	) -> impl Future<Output = Result<BoxStream<'static, Result<T::Response>>>> + use<T> {
		245	let envelope = request.into_envelope(0, None, None);
		246	let response_stream = self.0.client.request_stream(envelope, T::NAME);
		247	async move {
		248	Ok(response_stream
		249	.await?
		250	.map(\|response\| {
		251	T::Response::from_envelope(response?)
		252	.context("received response of the wrong type")
		253	})
		254	.boxed())
		255	}
		256	}
		257
226	pub fn send<T: EnvelopedMessage>(&self, request: T) -> Result<()> {	258	pub fn send<T: EnvelopedMessage>(&self, request: T) -> Result<()> {
227	let envelope = request.into_envelope(0, None, None);	259	let envelope = request.into_envelope(0, None, None);
228	self.0.client.send(envelope, T::NAME)	260	self.0.client.send(envelope, T::NAME)
@@ -479,6 +511,68 @@ impl AnyProtoClient { · 250 unmodified lines
479	);	511	);
480	}	512	}
481		513
		514	pub fn add_entity_stream_request_handler<M, E, H, F, S>(&self, handler: H)
		515	where
		516	M: EnvelopedMessage + RequestMessage + EntityMessage,
		517	E: 'static,
		518	H: 'static + Sync + Send + Fn(gpui::Entity<E>, TypedEnvelope<M>, AsyncApp) -> F,
		519	F: 'static + Future<Output = Result<S>>,
		520	S: 'static + Stream<Item = Result<M::Response>>,
		521	{
		522	let message_type_id = TypeId::of::<M>();
		523	let entity_type_id = TypeId::of::<E>();
		524	let entity_id_extractor = \|envelope: &dyn AnyTypedEnvelope\| {
		525	(envelope as &dyn Any)
		526	.downcast_ref::<TypedEnvelope<M>>()
		527	.unwrap()
		528	.payload
		529	.remote_entity_id()
		530	};
		531	self.0
		532	.client
		533	.message_handler_set()
		534	.lock()
		535	.add_entity_message_handler(
		536	message_type_id,
		537	entity_type_id,
		538	entity_id_extractor,
		539	Arc::new(move \|entity, envelope, client, cx\| {
		540	let entity = entity.downcast::<E>().unwrap();
		541	let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
		542	let request_id = envelope.message_id();
		543	let stream = handler(entity, *envelope, cx);
		544	async move {
		545	// An Error response is itself a terminal stream frame on
		546	// both transports (Peer and ChannelClient), so we don't
		547	// need to follow it with an EndStream.
		548	match stream.await {
		549	Ok(stream) => {
		550	futures::pin_mut!(stream);
		551	while let Some(result) = stream.next().await {
		552	match result {
		553	Ok(response) => {
		554	client.send_response(request_id, response)?
		555	}
		556	Err(error) => {
		557	client.send_response(request_id, error.to_proto())?;
		558	return Err(error);
		559	}
		560	}
		561	}
		562	client.send_response(request_id, proto::EndStream {})?;
		563	Ok(())
		564	}
		565	Err(error) => {
		566	client.send_response(request_id, error.to_proto())?;
		567	Err(error)
		568	}
		569	}
		570	}
		571	.boxed_local()
		572	}),
		573	);
		574	}
		575
482	pub fn add_entity_message_handler<M, E, H, F>(&self, handler: H)	576	pub fn add_entity_message_handler<M, E, H, F>(&self, handler: H)
483	where	577	where
484	M: EnvelopedMessage + EntityMessage,	578	M: EnvelopedMessage + EntityMessage,