Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions crates/goose-cli/src/scenario_tests/mock_client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,7 @@ impl McpClientTrait for MockClient {
_session_id: &str,
name: &str,
arguments: Option<serde_json::Map<String, Value>>,
_working_dir: Option<&str>,
_cancel_token: CancellationToken,
) -> Result<CallToolResult, Error> {
if let Some(handler) = self.handlers.get(name) {
Expand Down
24 changes: 18 additions & 6 deletions crates/goose-mcp/src/developer/rmcp_developer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,13 +10,26 @@ use rmcp::{
model::{
CallToolResult, CancelledNotificationParam, Content, ErrorCode, ErrorData,
GetPromptRequestParams, GetPromptResult, Implementation, ListPromptsResult, LoggingLevel,
LoggingMessageNotificationParam, PaginatedRequestParams, Prompt, PromptArgument,
LoggingMessageNotificationParam, Meta, PaginatedRequestParams, Prompt, PromptArgument,
PromptMessage, PromptMessageRole, Role, ServerCapabilities, ServerInfo,
},
schemars::JsonSchema,
service::{NotificationContext, RequestContext},
tool, tool_handler, tool_router, RoleServer, ServerHandler,
};

/// Header name for passing working directory through MCP request metadata
const WORKING_DIR_HEADER: &str = "agent-working-dir";

/// Extract working directory from MCP request metadata
fn extract_working_dir_from_meta(meta: &Meta) -> Option<PathBuf> {
meta.0
.get(WORKING_DIR_HEADER)
.and_then(|v| v.as_str())
.filter(|s| !s.is_empty())
.map(PathBuf::from)
}

use serde::{Deserialize, Serialize};
use std::{
collections::HashMap,
Expand Down Expand Up @@ -873,6 +886,8 @@ impl DeveloperServer {
let peer = context.peer;
let request_id = context.id;

let working_dir = extract_working_dir_from_meta(&context.meta);

// Validate the shell command
self.validate_shell_command(command)?;

Expand All @@ -886,7 +901,7 @@ impl DeveloperServer {

// Execute the command and capture output
let output_result = self
.execute_shell_command(command, &peer, cancellation_token.clone())
.execute_shell_command(command, &peer, cancellation_token.clone(), working_dir)
.await;

// Clean up the process from tracking
Expand Down Expand Up @@ -970,17 +985,14 @@ impl DeveloperServer {
command: &str,
peer: &rmcp::service::Peer<RoleServer>,
cancellation_token: CancellationToken,
working_dir: Option<PathBuf>,
) -> Result<String, ErrorData> {
let mut shell_config = ShellConfig::default();
let shell_name = std::path::Path::new(&shell_config.executable)
.file_name()
.and_then(|s| s.to_str())
.unwrap_or("bash");

let working_dir = std::env::var("GOOSE_WORKING_DIR")
.ok()
.map(std::path::PathBuf::from);

if let Some(ref env_file) = self.bash_env_file {
if shell_name == "bash" {
shell_config.envs.push((
Expand Down
7 changes: 6 additions & 1 deletion crates/goose-server/src/routes/agent.rs
Original file line number Diff line number Diff line change
Expand Up @@ -947,7 +947,12 @@ async fn call_tool(

let tool_result = agent
.extension_manager
.dispatch_tool_call(&payload.session_id, tool_call, CancellationToken::default())
.dispatch_tool_call(
&payload.session_id,
tool_call,
None,
CancellationToken::default(),
)
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

Expand Down
1 change: 1 addition & 0 deletions crates/goose/src/agents/agent.rs
Original file line number Diff line number Diff line change
Expand Up @@ -581,6 +581,7 @@ impl Agent {
.dispatch_tool_call(
&session.id,
tool_call.clone(),
Some(session.working_dir.as_path()),
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe dispatch_tool_call should just take a reference to the session itself? now that we're passing two fields from it

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I looked into this and its a pretty large refactor because not all call sites have access to a full session and it would add additional overhead to create a session and/or mock a fake one etc. Gonna punt on this for now and we can come back to it later..

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

that makes sense - thanks for looking into it

cancellation_token.unwrap_or_default(),
)
.await;
Expand Down
1 change: 1 addition & 0 deletions crates/goose/src/agents/apps_extension.rs
Original file line number Diff line number Diff line change
Expand Up @@ -531,6 +531,7 @@ impl McpClientTrait for AppsManagerClient {
session_id: &str,
name: &str,
arguments: Option<JsonObject>,
_working_dir: Option<&str>,
_cancel_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let result = match name {
Expand Down
1 change: 1 addition & 0 deletions crates/goose/src/agents/chatrecall_extension.rs
Original file line number Diff line number Diff line change
Expand Up @@ -297,6 +297,7 @@ impl McpClientTrait for ChatRecallClient {
session_id: &str,
name: &str,
arguments: Option<JsonObject>,
_working_dir: Option<&str>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let content = match name {
Expand Down
3 changes: 2 additions & 1 deletion crates/goose/src/agents/code_execution_extension.rs
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,7 @@ fn create_tool_callback(
arguments: args.and_then(|v| v.as_object().cloned()),
};
match manager
.dispatch_tool_call(&session_id, tool_call, CancellationToken::new())
.dispatch_tool_call(&session_id, tool_call, None, CancellationToken::new())
.await
{
Ok(dispatch_result) => match dispatch_result.result.await {
Expand Down Expand Up @@ -422,6 +422,7 @@ impl McpClientTrait for CodeExecutionClient {
session_id: &str,
name: &str,
arguments: Option<JsonObject>,
_working_dir: Option<&str>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let result = match name {
Expand Down
65 changes: 47 additions & 18 deletions crates/goose/src/agents/extension_manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -599,16 +599,6 @@ impl ExtensionManager {
.await?;
Box::new(client)
} else {
// Set GOOSE_WORKING_DIR in the current process for builtin extensions
// since they run in-process and read from std::env::var
if effective_working_dir.exists() && effective_working_dir.is_dir() {
std::env::set_var("GOOSE_WORKING_DIR", &effective_working_dir);
tracing::info!(
"Set GOOSE_WORKING_DIR for builtin extension: {:?}",
effective_working_dir
);
}

let (server_read, client_write) = tokio::io::duplex(65536);
let (client_read, server_write) = tokio::io::duplex(65536);
extension_fn(server_read, server_write);
Expand Down Expand Up @@ -1189,6 +1179,7 @@ impl ExtensionManager {
&self,
session_id: &str,
tool_call: CallToolRequestParams,
working_dir: Option<&std::path::Path>,
cancellation_token: CancellationToken,
) -> Result<ToolCallResult> {
// Some models strip the tool prefix, so auto-add it for known code_execution tools
Expand Down Expand Up @@ -1248,16 +1239,24 @@ impl ExtensionManager {
let client = client.clone();
let notifications_receiver = client.lock().await.subscribe().await;
let session_id = session_id.to_string();
let working_dir_str = working_dir.map(|p| p.to_string_lossy().to_string());

let fut = async move {
tracing::debug!(
"dispatch_tool_call fut: calling client.call_tool tool={} session_id={}",
"dispatch_tool_call fut: calling client.call_tool tool={} session_id={} working_dir={:?}",
tool_name,
session_id
session_id,
working_dir_str
);
let client_guard = client.lock().await;
client_guard
.call_tool(&session_id, &tool_name, arguments, cancellation_token)
.call_tool(
&session_id,
&tool_name,
arguments,
working_dir_str.as_deref(),
cancellation_token,
)
.await
.map_err(|e| match e {
ServiceError::McpError(error_data) => error_data,
Expand Down Expand Up @@ -1591,6 +1590,7 @@ mod tests {
_session_id: &str,
name: &str,
_arguments: Option<JsonObject>,
_working_dir: Option<&str>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
match name {
Expand Down Expand Up @@ -1727,7 +1727,12 @@ mod tests {
};

let result = extension_manager
.dispatch_tool_call("test-session-id", tool_call, CancellationToken::default())
.dispatch_tool_call(
"test-session-id",
tool_call,
None,
CancellationToken::default(),
)
.await;
assert!(result.is_ok());

Expand All @@ -1739,7 +1744,12 @@ mod tests {
};

let result = extension_manager
.dispatch_tool_call("test-session-id", tool_call, CancellationToken::default())
.dispatch_tool_call(
"test-session-id",
tool_call,
None,
CancellationToken::default(),
)
.await;
assert!(result.is_ok());

Expand All @@ -1752,7 +1762,12 @@ mod tests {
};

let result = extension_manager
.dispatch_tool_call("test-session-id", tool_call, CancellationToken::default())
.dispatch_tool_call(
"test-session-id",
tool_call,
None,
CancellationToken::default(),
)
.await;
assert!(result.is_ok());

Expand All @@ -1765,7 +1780,12 @@ mod tests {
};

let result = extension_manager
.dispatch_tool_call("test-session-id", tool_call, CancellationToken::default())
.dispatch_tool_call(
"test-session-id",
tool_call,
None,
CancellationToken::default(),
)
.await;
assert!(result.is_ok());

Expand All @@ -1777,7 +1797,12 @@ mod tests {
};

let result = extension_manager
.dispatch_tool_call("test-session-id", tool_call, CancellationToken::default())
.dispatch_tool_call(
"test-session-id",
tool_call,
None,
CancellationToken::default(),
)
.await;
assert!(result.is_ok());

Expand All @@ -1793,6 +1818,7 @@ mod tests {
.dispatch_tool_call(
"test-session-id",
invalid_tool_call,
None,
CancellationToken::default(),
)
.await
Expand Down Expand Up @@ -1820,6 +1846,7 @@ mod tests {
.dispatch_tool_call(
"test-session-id",
invalid_tool_call,
None,
CancellationToken::default(),
)
.await;
Expand Down Expand Up @@ -1922,6 +1949,7 @@ mod tests {
.dispatch_tool_call(
"test-session-id",
unavailable_tool_call,
None,
CancellationToken::default(),
)
.await;
Expand All @@ -1947,6 +1975,7 @@ mod tests {
.dispatch_tool_call(
"test-session-id",
available_tool_call,
None,
CancellationToken::default(),
)
.await;
Expand Down
1 change: 1 addition & 0 deletions crates/goose/src/agents/extension_manager_extension.rs
Original file line number Diff line number Diff line change
Expand Up @@ -435,6 +435,7 @@ impl McpClientTrait for ExtensionManagerClient {
session_id: &str,
name: &str,
arguments: Option<JsonObject>,
_working_dir: Option<&str>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let result = match name {
Expand Down
Loading
Loading