aljaz/goose
1
0
Fork 0
mirror of https://github.com/aljazceru/goose.git synced 2025-12-18 22:54:24 +01:00
Code Issues Packages Projects Releases Wiki Activity

Add pending request cleanup and config

Browse Source
This commit is contained in:
Aljaz
2025-05-29 14:53:48 +02:00
parent 493ba53c08
commit 89130c2067
15 changed files with 242 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

47
crates/goose-api/src/handlers.rs
View File

@@ -72,6 +72,10 @@ pub enum ExtensionConfigRequest {
#[serde(default)]
env_keys: Vec<String>,
timeout: Option<u64>,
#[serde(default)]
max_pending_requests: Option<usize>,
#[serde(default)]
pending_request_timeout: Option<u64>,
},
#[serde(rename = "stdio")]
Stdio {
@@ -84,12 +88,20 @@ pub enum ExtensionConfigRequest {
#[serde(default)]
env_keys: Vec<String>,
timeout: Option<u64>,
#[serde(default)]
max_pending_requests: Option<usize>,
#[serde(default)]
pending_request_timeout: Option<u64>,
},
#[serde(rename = "builtin")]
Builtin {
name: String,
display_name: Option<String>,
timeout: Option<u64>,
#[serde(default)]
max_pending_requests: Option<usize>,
#[serde(default)]
pending_request_timeout: Option<u64>,
},
#[serde(rename = "frontend")]
Frontend {
@@ -370,7 +382,15 @@ pub async fn add_extension_handler(
}
let extension = match req {
ExtensionConfigRequest::Sse { name, uri, envs, env_keys, timeout } => {
ExtensionConfigRequest::Sse {
name,
uri,
envs,
env_keys,
timeout,
max_pending_requests,
pending_request_timeout,
} => {
ExtensionConfig::Sse {
name,
uri,
@@ -378,10 +398,21 @@ pub async fn add_extension_handler(
env_keys,
description: None,
timeout,
max_pending_requests,
pending_request_timeout,
bundled: None,
}
}
ExtensionConfigRequest::Stdio { name, cmd, args, envs, env_keys, timeout } => {
ExtensionConfigRequest::Stdio {
name,
cmd,
args,
envs,
env_keys,
timeout,
max_pending_requests,
pending_request_timeout,
} => {
ExtensionConfig::Stdio {
name,
cmd,
@@ -389,15 +420,25 @@ pub async fn add_extension_handler(
envs,
env_keys,
timeout,
max_pending_requests,
pending_request_timeout,
description: None,
bundled: None,
}
}
ExtensionConfigRequest::Builtin { name, display_name, timeout } => {
ExtensionConfigRequest::Builtin {
name,
display_name,
timeout,
max_pending_requests,
pending_request_timeout,
} => {
ExtensionConfig::Builtin {
name,
display_name,
timeout,
max_pending_requests,
pending_request_timeout,
bundled: None,
}
}

4
crates/goose-cli/src/commands/configure.rs
View File

@@ -79,6 +79,8 @@ pub async fn handle_configure() -> Result<(), Box<dyn Error>> {
name: "developer".to_string(),
display_name: Some(goose::config::DEFAULT_DISPLAY_NAME.to_string()),
timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT),
max_pending_requests: Some(goose::config::DEFAULT_MAX_PENDING_REQUESTS),
pending_request_timeout: Some(goose::config::DEFAULT_PENDING_REQUEST_TIMEOUT),
bundled: Some(true),
},
})?;
@@ -548,6 +550,8 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
name: extension.clone(),
display_name: Some(display_name),
timeout: Some(timeout),
max_pending_requests: Some(goose::config::DEFAULT_MAX_PENDING_REQUESTS),
pending_request_timeout: Some(goose::config::DEFAULT_PENDING_REQUEST_TIMEOUT),
bundled: Some(true),
},
})?;

6
crates/goose-cli/src/session/mod.rs
View File

@@ -177,6 +177,8 @@ impl Session {
description: Some(goose::config::DEFAULT_EXTENSION_DESCRIPTION.to_string()),
// TODO: should set timeout
timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT),
max_pending_requests: Some(goose::config::DEFAULT_MAX_PENDING_REQUESTS),
pending_request_timeout: Some(goose::config::DEFAULT_PENDING_REQUEST_TIMEOUT),
bundled: None,
};
@@ -210,6 +212,8 @@ impl Session {
description: Some(goose::config::DEFAULT_EXTENSION_DESCRIPTION.to_string()),
// TODO: should set timeout
timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT),
max_pending_requests: Some(goose::config::DEFAULT_MAX_PENDING_REQUESTS),
pending_request_timeout: Some(goose::config::DEFAULT_PENDING_REQUEST_TIMEOUT),
bundled: None,
};
@@ -235,6 +239,8 @@ impl Session {
display_name: None,
// TODO: should set a timeout
timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT),
max_pending_requests: Some(goose::config::DEFAULT_MAX_PENDING_REQUESTS),
pending_request_timeout: Some(goose::config::DEFAULT_PENDING_REQUEST_TIMEOUT),
bundled: None,
};
self.agent

24
crates/goose-server/src/routes/extension.rs
View File

@@ -29,6 +29,10 @@ enum ExtensionConfigRequest {
#[serde(default)]
env_keys: Vec<String>,
timeout: Option<u64>,
#[serde(default)]
max_pending_requests: Option<usize>,
#[serde(default)]
pending_request_timeout: Option<u64>,
},
/// Standard I/O (stdio) extension.
#[serde(rename = "stdio")]
@@ -47,6 +51,10 @@ enum ExtensionConfigRequest {
#[serde(default)]
env_keys: Vec<String>,
timeout: Option<u64>,
#[serde(default)]
max_pending_requests: Option<usize>,
#[serde(default)]
pending_request_timeout: Option<u64>,
},
/// Built-in extension that is part of the goose binary.
#[serde(rename = "builtin")]
@@ -55,6 +63,10 @@ enum ExtensionConfigRequest {
name: String,
display_name: Option<String>,
timeout: Option<u64>,
#[serde(default)]
max_pending_requests: Option<usize>,
#[serde(default)]
pending_request_timeout: Option<u64>,
},
/// Frontend extension that provides tools to be executed by the frontend.
#[serde(rename = "frontend")]
@@ -167,6 +179,8 @@ async fn add_extension(
envs,
env_keys,
timeout,
max_pending_requests,
pending_request_timeout,
} => ExtensionConfig::Sse {
name,
uri,
@@ -174,6 +188,8 @@ async fn add_extension(
env_keys,
description: None,
timeout,
max_pending_requests,
pending_request_timeout,
bundled: None,
},
ExtensionConfigRequest::Stdio {
@@ -183,6 +199,8 @@ async fn add_extension(
envs,
env_keys,
timeout,
max_pending_requests,
pending_request_timeout,
} => {
// TODO: We can uncomment once bugs are fixed. Check allowlist for Stdio extensions
// if !is_command_allowed(&cmd, &args) {
@@ -204,6 +222,8 @@ async fn add_extension(
envs,
env_keys,
timeout,
max_pending_requests,
pending_request_timeout,
bundled: None,
}
}
@@ -211,10 +231,14 @@ async fn add_extension(
name,
display_name,
timeout,
max_pending_requests,
pending_request_timeout,
} => ExtensionConfig::Builtin {
name,
display_name,
timeout,
max_pending_requests,
pending_request_timeout,
bundled: None,
},
ExtensionConfigRequest::Frontend {

22
crates/goose/src/agents/extension.rs
View File

@@ -136,6 +136,10 @@ pub enum ExtensionConfig {
// NOTE: set timeout to be optional for compatibility.
// However, new configurations should include this field.
timeout: Option<u64>,
#[serde(default)]
max_pending_requests: Option<usize>,
#[serde(default)]
pending_request_timeout: Option<u64>,
/// Whether this extension is bundled with Goose
#[serde(default)]
bundled: Option<bool>,
@@ -153,6 +157,10 @@ pub enum ExtensionConfig {
env_keys: Vec<String>,
timeout: Option<u64>,
description: Option<String>,
#[serde(default)]
max_pending_requests: Option<usize>,
#[serde(default)]
pending_request_timeout: Option<u64>,
/// Whether this extension is bundled with Goose
#[serde(default)]
bundled: Option<bool>,
@@ -164,6 +172,10 @@ pub enum ExtensionConfig {
name: String,
display_name: Option<String>, // needed for the UI
timeout: Option<u64>,
#[serde(default)]
max_pending_requests: Option<usize>,
#[serde(default)]
pending_request_timeout: Option<u64>,
/// Whether this extension is bundled with Goose
#[serde(default)]
bundled: Option<bool>,
@@ -189,6 +201,8 @@ impl Default for ExtensionConfig {
name: config::DEFAULT_EXTENSION.to_string(),
display_name: Some(config::DEFAULT_DISPLAY_NAME.to_string()),
timeout: Some(config::DEFAULT_EXTENSION_TIMEOUT),
max_pending_requests: Some(config::DEFAULT_MAX_PENDING_REQUESTS),
pending_request_timeout: Some(config::DEFAULT_PENDING_REQUEST_TIMEOUT),
bundled: Some(true),
}
}
@@ -203,6 +217,8 @@ impl ExtensionConfig {
env_keys: Vec::new(),
description: Some(description.into()),
timeout: Some(timeout.into()),
max_pending_requests: Some(config::DEFAULT_MAX_PENDING_REQUESTS),
pending_request_timeout: Some(config::DEFAULT_PENDING_REQUEST_TIMEOUT),
bundled: None,
}
}
@@ -221,6 +237,8 @@ impl ExtensionConfig {
env_keys: Vec::new(),
description: Some(description.into()),
timeout: Some(timeout.into()),
max_pending_requests: Some(config::DEFAULT_MAX_PENDING_REQUESTS),
pending_request_timeout: Some(config::DEFAULT_PENDING_REQUEST_TIMEOUT),
bundled: None,
}
}
@@ -237,6 +255,8 @@ impl ExtensionConfig {
envs,
env_keys,
timeout,
max_pending_requests,
pending_request_timeout,
description,
bundled,
..
@@ -248,6 +268,8 @@ impl ExtensionConfig {
args: args.into_iter().map(Into::into).collect(),
description,
timeout,
max_pending_requests,
pending_request_timeout,
bundled,
},
other => other,

23
crates/goose/src/agents/extension_manager.rs
View File

@@ -178,10 +178,17 @@ impl ExtensionManager {
envs,
env_keys,
timeout,
max_pending_requests,
pending_request_timeout,
..
} => {
let all_envs = merge_environments(envs, env_keys, &sanitized_name).await?;
let transport = SseTransport::new(uri, all_envs);
let transport = SseTransport::new(
uri,
all_envs,
*max_pending_requests,
pending_request_timeout.map(Duration::from_secs),
);
let handle = transport.start().await?;
let service = McpService::with_timeout(
handle,
@@ -197,10 +204,18 @@ impl ExtensionManager {
envs,
env_keys,
timeout,
max_pending_requests,
pending_request_timeout,
..
} => {
let all_envs = merge_environments(envs, env_keys, &sanitized_name).await?;
let transport = StdioTransport::new(cmd, args.to_vec(), all_envs);
let transport = StdioTransport::new(
cmd,
args.to_vec(),
all_envs,
*max_pending_requests,
pending_request_timeout.map(Duration::from_secs),
);
let handle = transport.start().await?;
let service = McpService::with_timeout(
handle,
@@ -214,6 +229,8 @@ impl ExtensionManager {
name,
display_name: _,
timeout,
max_pending_requests,
pending_request_timeout,
bundled: _,
} => {
let cmd = std::env::current_exe()
@@ -225,6 +242,8 @@ impl ExtensionManager {
&cmd,
vec!["mcp".to_string(), name.clone()],
HashMap::new(),
*max_pending_requests,
pending_request_timeout.map(Duration::from_secs),
);
let handle = transport.start().await?;
let service = McpService::with_timeout(

4
crates/goose/src/config/extensions.rs
View File

@@ -7,6 +7,8 @@ use utoipa::ToSchema;
pub const DEFAULT_EXTENSION: &str = "developer";
pub const DEFAULT_EXTENSION_TIMEOUT: u64 = 300;
pub const DEFAULT_MAX_PENDING_REQUESTS: usize = 128;
pub const DEFAULT_PENDING_REQUEST_TIMEOUT: u64 = 30;
pub const DEFAULT_EXTENSION_DESCRIPTION: &str = "";
pub const DEFAULT_DISPLAY_NAME: &str = "Developer";
@@ -45,6 +47,8 @@ impl ExtensionConfigManager {
name: DEFAULT_EXTENSION.to_string(),
display_name: Some(DEFAULT_DISPLAY_NAME.to_string()),
timeout: Some(DEFAULT_EXTENSION_TIMEOUT),
max_pending_requests: Some(DEFAULT_MAX_PENDING_REQUESTS),
pending_request_timeout: Some(DEFAULT_PENDING_REQUEST_TIMEOUT),
bundled: Some(true),
},
},

2
crates/goose/src/config/mod.rs
View File

@@ -13,3 +13,5 @@ pub use extensions::DEFAULT_DISPLAY_NAME;
pub use extensions::DEFAULT_EXTENSION;
pub use extensions::DEFAULT_EXTENSION_DESCRIPTION;
pub use extensions::DEFAULT_EXTENSION_TIMEOUT;
pub use extensions::DEFAULT_MAX_PENDING_REQUESTS;
pub use extensions::DEFAULT_PENDING_REQUEST_TIMEOUT;

23
crates/mcp-client/examples/clients.rs
View File

@@ -18,17 +18,34 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
)
.init();
let transport1 = StdioTransport::new("uvx", vec!["mcp-server-git".to_string()], HashMap::new());
let transport1 = StdioTransport::new(
"uvx",
vec!["mcp-server-git".to_string()],
HashMap::new(),
None,
None,
);
let handle1 = transport1.start().await?;
let service1 = McpService::with_timeout(handle1, Duration::from_secs(30));
let client1 = McpClient::new(service1);
let transport2 = StdioTransport::new("uvx", vec!["mcp-server-git".to_string()], HashMap::new());
let transport2 = StdioTransport::new(
"uvx",
vec!["mcp-server-git".to_string()],
HashMap::new(),
None,
None,
);
let handle2 = transport2.start().await?;
let service2 = McpService::with_timeout(handle2, Duration::from_secs(30));
let client2 = McpClient::new(service2);
let transport3 = SseTransport::new("http://localhost:8000/sse", HashMap::new());
let transport3 = SseTransport::new(
"http://localhost:8000/sse",
HashMap::new(),
None,
None,
);
let handle3 = transport3.start().await?;
let service3 = McpService::with_timeout(handle3, Duration::from_secs(10));
let client3 = McpClient::new(service3);

7
crates/mcp-client/examples/sse.rs
View File

@@ -18,7 +18,12 @@ async fn main() -> Result<()> {
.init();
// Create the base transport
let transport = SseTransport::new("http://localhost:8000/sse", HashMap::new());
let transport = SseTransport::new(
"http://localhost:8000/sse",
HashMap::new(),
None,
None,
);
// Start transport
let handle = transport.start().await?;

8
crates/mcp-client/examples/stdio.rs
View File

@@ -20,7 +20,13 @@ async fn main() -> Result<(), ClientError> {
.init();
// 1) Create the transport
let transport = StdioTransport::new("uvx", vec!["mcp-server-git".to_string()], HashMap::new());
let transport = StdioTransport::new(
"uvx",
vec!["mcp-server-git".to_string()],
HashMap::new(),
None,
None,
);
// 2) Start the transport to get a handle
let transport_handle = transport.start().await?;

2
crates/mcp-client/examples/stdio_integration.rs
View File

@@ -29,6 +29,8 @@ async fn main() -> Result<(), ClientError> {
.map(|s| s.to_string())
.collect(),
HashMap::new(),
None,
None,
);
// Start the transport to get a handle

43
crates/mcp-client/src/transport/mod.rs
View File

@@ -1,6 +1,7 @@
use async_trait::async_trait;
use mcp_core::protocol::JsonRpcMessage;
use std::collections::HashMap;
use std::time::{Duration, Instant};
use thiserror::Error;
use tokio::sync::{mpsc, oneshot, RwLock};
@@ -31,6 +32,9 @@ pub enum Error {
#[error("HTTP error: {status} - {message}")]
HttpError { status: u16, message: String },
#[error("Too many pending requests")]
PendingRequestsFull,
}
/// A message that can be sent through the transport
@@ -89,7 +93,10 @@ pub async fn send_message(
// A data structure to store pending requests and their response channels
pub struct PendingRequests {
requests: RwLock<HashMap<String, oneshot::Sender<Result<JsonRpcMessage, Error>>>>,
requests:
RwLock<HashMap<String, (oneshot::Sender<Result<JsonRpcMessage, Error>>, Instant)>>,
max_size: Option<usize>,
timeout: Option<Duration>,
}
impl Default for PendingRequests {
@@ -100,21 +107,49 @@ impl Default for PendingRequests {
impl PendingRequests {
pub fn new() -> Self {
Self::with_limits(None, None)
}
pub fn with_limits(max_size: Option<usize>, timeout: Option<Duration>) -> Self {
Self {
requests: RwLock::new(HashMap::new()),
max_size,
timeout,
}
}
pub async fn insert(&self, id: String, sender: oneshot::Sender<Result<JsonRpcMessage, Error>>) {
self.requests.write().await.insert(id, sender);
pub async fn insert(
&self,
id: String,
sender: oneshot::Sender<Result<JsonRpcMessage, Error>>,
) -> Result<(), Error> {
self.cleanup().await;
let mut map = self.requests.write().await;
if let Some(max) = self.max_size {
if map.len() >= max {
return Err(Error::PendingRequestsFull);
}
}
map.insert(id, (sender, Instant::now()));
Ok(())
}
pub async fn respond(&self, id: &str, response: Result<JsonRpcMessage, Error>) {
if let Some(tx) = self.requests.write().await.remove(id) {
if let Some((tx, _)) = self.requests.write().await.remove(id) {
let _ = tx.send(response);
}
}
pub async fn cleanup(&self) {
if let Some(timeout) = self.timeout {
let now = Instant::now();
self.requests
.write()
.await
.retain(|_, (_, ts)| now.duration_since(*ts) <= timeout);
}
}
pub async fn clear(&self) {
self.requests.write().await.clear();
}

25
crates/mcp-client/src/transport/sse.rs
View File

@@ -179,7 +179,14 @@ impl SseActor {
if let JsonRpcMessage::Request(JsonRpcRequest { id: Some(id), .. }) =
&transport_msg.message
{
pending_requests.insert(id.to_string(), response_tx).await;
pending_requests.cleanup().await;
if let Err(e) = pending_requests
.insert(id.to_string(), response_tx)
.await
{
let _ = response_tx.send(Err(e));
continue;
}
}
}
@@ -236,14 +243,23 @@ impl TransportHandle for SseTransportHandle {
pub struct SseTransport {
sse_url: String,
env: HashMap<String, String>,
max_pending: Option<usize>,
pending_timeout: Option<Duration>,
}
/// The SSE transport spawns an `SseActor` on `start()`.
impl SseTransport {
pub fn new<S: Into<String>>(sse_url: S, env: HashMap<String, String>) -> Self {
pub fn new<S: Into<String>>(
sse_url: S,
env: HashMap<String, String>,
max_pending: Option<usize>,
pending_timeout: Option<Duration>,
) -> Self {
Self {
sse_url: sse_url.into(),
env,
max_pending,
pending_timeout,
}
}
@@ -286,7 +302,10 @@ impl Transport for SseTransport {
// Build the actor
let actor = SseActor::new(
rx,
Arc::new(PendingRequests::new()),
Arc::new(PendingRequests::with_limits(
self.max_pending,
self.pending_timeout,
)),
self.sse_url.clone(),
post_endpoint,
);

21
crates/mcp-client/src/transport/stdio.rs
View File

@@ -1,4 +1,5 @@
use std::collections::HashMap;
use std::time::Duration;
use std::sync::atomic::{AtomicI32, Ordering};
use std::sync::Arc;
use tokio::process::{Child, ChildStderr, ChildStdin, ChildStdout, Command};
@@ -162,7 +163,14 @@ impl StdioActor {
if let Some(response_tx) = transport_msg.response_tx.take() {
if let JsonRpcMessage::Request(request) = &transport_msg.message {
if let Some(id) = &request.id {
pending_requests.insert(id.to_string(), response_tx).await;
pending_requests.cleanup().await;
if let Err(e) = pending_requests
.insert(id.to_string(), response_tx)
.await
{
let _ = response_tx.send(Err(e));
continue;
}
}
}
}
@@ -218,6 +226,8 @@ pub struct StdioTransport {
command: String,
args: Vec<String>,
env: HashMap<String, String>,
max_pending: Option<usize>,
pending_timeout: Option<Duration>,
}
impl StdioTransport {
@@ -225,11 +235,15 @@ impl StdioTransport {
command: S,
args: Vec<String>,
env: HashMap<String, String>,
max_pending: Option<usize>,
pending_timeout: Option<Duration>,
) -> Self {
Self {
command: command.into(),
args,
env,
max_pending,
pending_timeout,
}
}
@@ -294,7 +308,10 @@ impl Transport for StdioTransport {
let actor = StdioActor {
receiver: Some(message_rx),
pending_requests: Arc::new(PendingRequests::new()),
pending_requests: Arc::new(PendingRequests::with_limits(
self.max_pending,
self.pending_timeout,
)),
process,
error_sender: error_tx,
stdin: Some(stdin),