Skip to content

Commit 307e11f

Browse files
committed
uses event buffer for mcp events
1 parent dc7de59 commit 307e11f

File tree

1 file changed

+11
-8
lines changed
  • crates/agent/src/agent/mcp

1 file changed

+11
-8
lines changed

crates/agent/src/agent/mcp/mod.rs

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,7 @@ use tokio::sync::{
132132
use tracing::{
133133
debug,
134134
error,
135+
info,
135136
warn,
136137
};
137138
use types::Prompt;
@@ -265,6 +266,7 @@ pub struct McpManager {
265266

266267
initializing_servers: HashMap<String, McpServerActorHandle>,
267268
servers: HashMap<String, McpServerActorHandle>,
269+
event_buf: Vec<McpServerActorEvent>,
268270
}
269271

270272
impl McpManager {
@@ -280,6 +282,7 @@ impl McpManager {
280282
cred_path,
281283
initializing_servers: HashMap::new(),
282284
servers: HashMap::new(),
285+
event_buf: Vec::<McpServerActorEvent>::new(),
283286
}
284287
}
285288

@@ -302,6 +305,10 @@ impl McpManager {
302305

303306
async fn main_loop(mut self, mcp_main_loop_to_handle_server_event_tx: broadcast::Sender<McpServerActorEvent>) {
304307
loop {
308+
self.event_buf
309+
.drain(..)
310+
.for_each(|evt| _ = mcp_main_loop_to_handle_server_event_tx.send(evt));
311+
305312
tokio::select! {
306313
req = self.request_rx.recv() => {
307314
let Some(req) = req else {
@@ -313,7 +320,7 @@ impl McpManager {
313320
},
314321
res = self.server_event_rx.recv() => {
315322
if let Some(evt) = res {
316-
self.handle_mcp_actor_event(evt, &mcp_main_loop_to_handle_server_event_tx);
323+
self.handle_mcp_actor_event(evt);
317324
}
318325
}
319326
}
@@ -361,11 +368,7 @@ impl McpManager {
361368
}
362369
}
363370

364-
fn handle_mcp_actor_event(
365-
&mut self,
366-
evt: McpServerActorEvent,
367-
mcp_main_loop_to_handle_server_event_tx: &broadcast::Sender<McpServerActorEvent>,
368-
) {
371+
fn handle_mcp_actor_event(&mut self, evt: McpServerActorEvent) {
369372
// TODO: keep a record of all the different server events received in this layer?
370373
match &evt {
371374
McpServerActorEvent::Initialized {
@@ -387,10 +390,10 @@ impl McpManager {
387390
self.initializing_servers.remove(server_name);
388391
},
389392
McpServerActorEvent::OauthRequest { server_name, oauth_url } => {
390-
tracing::info!(?server_name, ?oauth_url, "received oauth request");
393+
info!(?server_name, ?oauth_url, "received oauth request");
391394
},
392395
}
393-
let _ = mcp_main_loop_to_handle_server_event_tx.send(evt);
396+
self.event_buf.push(evt);
394397
}
395398
}
396399

0 commit comments

Comments
 (0)