ai-space/goose
3
0
Fork 0
mirror of https://github.com/aaif-goose/goose.git synced 2026-07-03 14:10:03 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix: per session working dir isolation (#6920)

Browse Source
This commit is contained in:
Zane
2026-02-03 10:20:02 -08:00
committed by GitHub
parent 8631caa890
commit aa356bd460
16 changed files with 143 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/scenario_tests/mock_client.rs
+1
View File
@@ -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) {
crates/goose-mcp/src/developer/rmcp_developer.rs
+18 -6
View File
@@ -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,
@@ -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)?;
@@ -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
@@ -970,6 +985,7 @@ 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)
@@ -977,10 +993,6 @@ impl DeveloperServer {
.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((
crates/goose-server/src/routes/agent.rs
+6 -1
View File
@@ -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)?;
crates/goose/src/agents/agent.rs
+1
View File
@@ -581,6 +581,7 @@ impl Agent {
.dispatch_tool_call(
&session.id,
tool_call.clone(),
Some(session.working_dir.as_path()),
cancellation_token.unwrap_or_default(),
)
.await;
crates/goose/src/agents/apps_extension.rs
+1
View File
@@ -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 {
crates/goose/src/agents/chatrecall_extension.rs
+1
View File
@@ -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 {
crates/goose/src/agents/code_execution_extension.rs
+2 -1
View File
@@ -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 {
@@ -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 {
crates/goose/src/agents/extension_manager.rs
+47 -18
View File
@@ -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);
@@ -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
@@ -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,
@@ -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 {
@@ -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());
@@ -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());
@@ -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());
@@ -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());
@@ -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());
@@ -1793,6 +1818,7 @@ mod tests {
.dispatch_tool_call(
"test-session-id",
invalid_tool_call,
None,
CancellationToken::default(),
)
.await
@@ -1820,6 +1846,7 @@ mod tests {
.dispatch_tool_call(
"test-session-id",
invalid_tool_call,
None,
CancellationToken::default(),
)
.await;
@@ -1922,6 +1949,7 @@ mod tests {
.dispatch_tool_call(
"test-session-id",
unavailable_tool_call,
None,
CancellationToken::default(),
)
.await;
@@ -1947,6 +1975,7 @@ mod tests {
.dispatch_tool_call(
"test-session-id",
available_tool_call,
None,
CancellationToken::default(),
)
.await;
crates/goose/src/agents/extension_manager_extension.rs
+1
View File
@@ -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 {
crates/goose/src/agents/mcp_client.rs
+54 -30
View File
@@ -1,6 +1,6 @@
use crate::action_required_manager::ActionRequiredManager;
use crate::agents::types::SharedProvider;
use crate::session_context::SESSION_ID_HEADER;
use crate::session_context::{SESSION_ID_HEADER, WORKING_DIR_HEADER};
use rmcp::model::{
Content, CreateElicitationRequestParams, CreateElicitationResult, ElicitationAction, ErrorCode,
Extensions, JsonObject, Meta,
@@ -51,6 +51,7 @@ pub trait McpClientTrait: Send + Sync {
session_id: &str,
name: &str,
arguments: Option<JsonObject>,
working_dir: Option<&str>,
cancel_token: CancellationToken,
) -> Result<CallToolResult, Error>;
@@ -105,7 +106,7 @@ pub trait McpClientTrait: Send + Sync {
pub struct GooseClient {
notification_handlers: Arc<Mutex<Vec<Sender<ServerNotification>>>>,
provider: SharedProvider,
// Single-slot because calls are serialized per MCP client; see send_request_with_session.
// Single-slot because calls are serialized per MCP client.
current_session_id: Arc<Mutex<Option<String>>>,
}
@@ -386,13 +387,14 @@ impl McpClient {
self.docker_container.as_deref()
}
async fn send_request_with_session(
async fn send_request_with_context(
&self,
session_id: &str,
working_dir: Option<&str>,
request: ClientRequest,
cancel_token: CancellationToken,
) -> Result<ServerResult, Error> {
let request = inject_session_id_into_request(request, Some(session_id));
let request = inject_session_context_into_request(request, Some(session_id), working_dir);
// ExtensionManager serializes calls per MCP connection, so one current_session_id slot
// is sufficient for mapping callbacks to the active request session.
let handle = {
@@ -473,8 +475,9 @@ impl McpClientTrait for McpClient {
cancel_token: CancellationToken,
) -> Result<ListResourcesResult, Error> {
let res = self
.send_request_with_session(
.send_request_with_context(
session_id,
None,
ClientRequest::ListResourcesRequest(ListResourcesRequest {
params: Some(PaginatedRequestParams { meta: None, cursor }),
method: Default::default(),
@@ -497,8 +500,9 @@ impl McpClientTrait for McpClient {
cancel_token: CancellationToken,
) -> Result<ReadResourceResult, Error> {
let res = self
.send_request_with_session(
.send_request_with_context(
session_id,
None,
ClientRequest::ReadResourceRequest(ReadResourceRequest {
params: ReadResourceRequestParams {
meta: None,
@@ -524,8 +528,9 @@ impl McpClientTrait for McpClient {
cancel_token: CancellationToken,
) -> Result<ListToolsResult, Error> {
let res = self
.send_request_with_session(
.send_request_with_context(
session_id,
None,
ClientRequest::ListToolsRequest(ListToolsRequest {
params: Some(PaginatedRequestParams { meta: None, cursor }),
method: Default::default(),
@@ -546,6 +551,7 @@ impl McpClientTrait for McpClient {
session_id: &str,
name: &str,
arguments: Option<JsonObject>,
working_dir: Option<&str>,
cancel_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let request = ClientRequest::CallToolRequest(CallToolRequest {
@@ -560,7 +566,7 @@ impl McpClientTrait for McpClient {
});
let result = self
.send_request_with_session(session_id, request, cancel_token)
.send_request_with_context(session_id, working_dir, request, cancel_token)
.await;
match result? {
@@ -576,8 +582,9 @@ impl McpClientTrait for McpClient {
cancel_token: CancellationToken,
) -> Result<ListPromptsResult, Error> {
let res = self
.send_request_with_session(
.send_request_with_context(
session_id,
None,
ClientRequest::ListPromptsRequest(ListPromptsRequest {
params: Some(PaginatedRequestParams { meta: None, cursor }),
method: Default::default(),
@@ -605,8 +612,9 @@ impl McpClientTrait for McpClient {
_ => None,
};
let res = self
.send_request_with_session(
.send_request_with_context(
session_id,
None,
ClientRequest::GetPromptRequest(GetPromptRequest {
params: GetPromptRequestParams {
meta: None,
@@ -633,20 +641,24 @@ impl McpClientTrait for McpClient {
}
}
/// Injects the given session_id into Extensions._meta.
/// None (or empty) removes any existing session id.
fn inject_session_id_into_extensions(
/// Injects the given session_id and working_dir into Extensions._meta.
/// None (or empty) removes any existing values.
fn inject_session_context_into_extensions(
mut extensions: Extensions,
session_id: Option<&str>,
working_dir: Option<&str>,
) -> Extensions {
let session_id = session_id.filter(|id| !id.is_empty());
let working_dir = working_dir.filter(|dir| !dir.is_empty());
let mut meta_map = extensions
.get::<Meta>()
.map(|meta| meta.0.clone())
.unwrap_or_default();
// JsonObject is case-sensitive, so we use retain for case-insensitive removal
meta_map.retain(|k, _| !k.eq_ignore_ascii_case(SESSION_ID_HEADER));
meta_map.retain(|k, _| {
!k.eq_ignore_ascii_case(SESSION_ID_HEADER) && !k.eq_ignore_ascii_case(WORKING_DIR_HEADER)
});
if let Some(session_id) = session_id {
meta_map.insert(
@@ -655,37 +667,51 @@ fn inject_session_id_into_extensions(
);
}
if let Some(working_dir) = working_dir {
meta_map.insert(
WORKING_DIR_HEADER.to_string(),
Value::String(working_dir.to_string()),
);
}
extensions.insert(Meta(meta_map));
extensions
}
fn inject_session_id_into_request(
fn inject_session_context_into_request(
request: ClientRequest,
session_id: Option<&str>,
working_dir: Option<&str>,
) -> ClientRequest {
match request {
ClientRequest::ListResourcesRequest(mut req) => {
req.extensions = inject_session_id_into_extensions(req.extensions, session_id);
req.extensions =
inject_session_context_into_extensions(req.extensions, session_id, working_dir);
ClientRequest::ListResourcesRequest(req)
}
ClientRequest::ReadResourceRequest(mut req) => {
req.extensions = inject_session_id_into_extensions(req.extensions, session_id);
req.extensions =
inject_session_context_into_extensions(req.extensions, session_id, working_dir);
ClientRequest::ReadResourceRequest(req)
}
ClientRequest::ListToolsRequest(mut req) => {
req.extensions = inject_session_id_into_extensions(req.extensions, session_id);
req.extensions =
inject_session_context_into_extensions(req.extensions, session_id, working_dir);
ClientRequest::ListToolsRequest(req)
}
ClientRequest::CallToolRequest(mut req) => {
req.extensions = inject_session_id_into_extensions(req.extensions, session_id);
req.extensions =
inject_session_context_into_extensions(req.extensions, session_id, working_dir);
ClientRequest::CallToolRequest(req)
}
ClientRequest::ListPromptsRequest(mut req) => {
req.extensions = inject_session_id_into_extensions(req.extensions, session_id);
req.extensions =
inject_session_context_into_extensions(req.extensions, session_id, working_dir);
ClientRequest::ListPromptsRequest(req)
}
ClientRequest::GetPromptRequest(mut req) => {
req.extensions = inject_session_id_into_extensions(req.extensions, session_id);
req.extensions =
inject_session_context_into_extensions(req.extensions, session_id, working_dir);
ClientRequest::GetPromptRequest(req)
}
other => other,
@@ -814,7 +840,8 @@ mod tests {
*slot = Some(session_id.to_string());
}
let extensions = inject_session_id_into_extensions(Extensions::new(), ext_session);
let extensions =
inject_session_context_into_extensions(Extensions::new(), ext_session, None);
let resolved = client.resolve_session_id(&extensions).await;
@@ -830,8 +857,6 @@ mod tests {
#[test_case(list_prompts_request; "list_prompts")]
#[test_case(get_prompt_request; "get_prompt")]
fn test_request_injects_session(request_builder: fn(Extensions) -> ClientRequest) {
use serde_json::json;
let session_id = "test-session-id";
let mut extensions = Extensions::new();
extensions.insert(
@@ -843,7 +868,7 @@ mod tests {
);
let request = request_builder(extensions);
let request = inject_session_id_into_request(request, Some(session_id));
let request = inject_session_context_into_request(request, Some(session_id), None);
let extensions = request_extensions(&request).expect("request should have extensions");
let meta = extensions
.get::<Meta>()
@@ -861,10 +886,9 @@ mod tests {
#[test]
fn test_session_id_in_mcp_meta() {
use serde_json::json;
let session_id = "test-session-789";
let extensions = inject_session_id_into_extensions(Default::default(), Some(session_id));
let extensions =
inject_session_context_into_extensions(Default::default(), Some(session_id), None);
let mcp_meta = extensions.get::<Meta>().unwrap();
assert_eq!(
@@ -904,7 +928,7 @@ mod tests {
expected_meta: serde_json::Value,
) {
use rmcp::model::Extensions;
use serde_json::{from_value, json};
use serde_json::from_value;
let mut extensions = Extensions::new();
extensions.insert(
@@ -916,7 +940,7 @@ mod tests {
.unwrap(),
);
let extensions = inject_session_id_into_extensions(extensions, session_id);
let extensions = inject_session_context_into_extensions(extensions, session_id, None);
let mcp_meta = extensions.get::<Meta>().unwrap();
assert_eq!(&mcp_meta.0, expected_meta.as_object().unwrap());
crates/goose/src/agents/skills_extension.rs
+1
View File
@@ -306,6 +306,7 @@ impl McpClientTrait for SkillsClient {
_session_id: &str,
name: &str,
arguments: Option<JsonObject>,
_working_dir: Option<&str>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let content = match name {
crates/goose/src/agents/todo_extension.rs
+1
View File
@@ -174,6 +174,7 @@ impl McpClientTrait for TodoClient {
session_id: &str,
name: &str,
arguments: Option<JsonObject>,
_working_dir: Option<&str>,
_cancellation_token: CancellationToken,
) -> Result<CallToolResult, Error> {
let content = match name {
crates/goose/src/session_context.rs
+1
View File
@@ -1,6 +1,7 @@
use tokio::task_local;
pub const SESSION_ID_HEADER: &str = "agent-session-id";
pub const WORKING_DIR_HEADER: &str = "agent-working-dir";
task_local! {
pub static SESSION_ID: Option<String>;
crates/goose/tests/mcp_integration_test.rs
+6 -1
View File
@@ -267,7 +267,12 @@ async fn test_replayed_session(
arguments: tool_call.arguments,
};
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;
let tool_result = result?;
ui/desktop/src/components/bottom_menu/DirSwitcher.tsx
-2
View File
@@ -3,7 +3,6 @@ import { FolderDot } from 'lucide-react';
import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from '../ui/Tooltip';
import { updateWorkingDir } from '../../api';
import { toast } from 'react-toastify';
import { setCurrentWorkingDir } from '../../utils/workingDir';
interface DirSwitcherProps {
className: string;
@@ -43,7 +42,6 @@ export const DirSwitcher: React.FC<DirSwitcherProps> = ({
const newDir = result.filePaths[0];
window.electron.addRecentDir(newDir);
setCurrentWorkingDir(newDir);
if (sessionId) {
onWorkingDirChange?.(newDir);
ui/desktop/src/utils/workingDir.ts
+2 -9
View File
@@ -1,11 +1,4 @@
// Track the current working dir for this window (updated when user changes it)
let currentWorkingDir: string | null = null;
export const setCurrentWorkingDir = (dir: string): void => {
currentWorkingDir = dir;
};
export const getInitialWorkingDir = (): string => {
// Use the current dir if set, otherwise fall back to initial config
return currentWorkingDir ?? (window.appConfig?.get('GOOSE_WORKING_DIR') as string) ?? '';
// Fall back to initial config from app startup
return (window.appConfig?.get('GOOSE_WORKING_DIR') as string) ?? '';
};