diff --git a/crates/goose/src/acp/fs.rs b/crates/goose/src/acp/fs.rs index 311514bb595d..13755da1243a 100644 --- a/crates/goose/src/acp/fs.rs +++ b/crates/goose/src/acp/fs.rs @@ -6,10 +6,10 @@ use crate::agents::platform_extensions::developer::edit::{ use crate::agents::platform_extensions::developer::shell::{ShellParams, OUTPUT_LIMIT_BYTES}; use crate::agents::platform_extensions::developer::DeveloperClient; use agent_client_protocol::schema::v1::{ - CreateTerminalRequest, Diff, KillTerminalRequest, ReadTextFileRequest, ReleaseTerminalRequest, - SessionId, SessionNotification, SessionUpdate, Terminal, TerminalOutputRequest, - ToolCallContent, ToolCallId, ToolCallLocation, ToolCallUpdate, ToolCallUpdateFields, ToolKind, - WaitForTerminalExitRequest, WriteTextFileRequest, + CreateTerminalRequest, Diff, EnvVariable, KillTerminalRequest, ReadTextFileRequest, + ReleaseTerminalRequest, SessionId, SessionNotification, SessionUpdate, Terminal, + TerminalOutputRequest, ToolCallContent, ToolCallId, ToolCallLocation, ToolCallUpdate, + ToolCallUpdateFields, ToolKind, WaitForTerminalExitRequest, WriteTextFileRequest, }; use agent_client_protocol::{Client, ConnectionTo}; use agent_client_protocol_schema::v1::TerminalId; @@ -69,6 +69,17 @@ pub(crate) struct AcpTools { pub(crate) terminal: bool, } +fn create_terminal_request( + session_id: &SessionId, + params: &ShellParams, + ctx: &crate::agents::ToolCallContext, +) -> CreateTerminalRequest { + CreateTerminalRequest::new(session_id.clone(), ¶ms.command) + .env(vec![EnvVariable::new("AGENT_SESSION_ID", &ctx.session_id)]) + .cwd(ctx.working_dir.clone()) + .output_byte_limit(OUTPUT_LIMIT_BYTES as u64) +} + fn error_result(msg: impl std::fmt::Display) -> CallToolResult { CallToolResult::error(vec![RmcpContent::text(msg.to_string()).with_priority(0.0)]) } @@ -250,11 +261,7 @@ impl AcpTools { let create_res = self .cx - .send_request( - CreateTerminalRequest::new(self.session_id.clone(), ¶ms.command) - .cwd(ctx.working_dir.clone()) - .output_byte_limit(OUTPUT_LIMIT_BYTES as u64), - ) + .send_request(create_terminal_request(&self.session_id, ¶ms, ctx)) .block_task() .await .map_err(|e| { @@ -438,3 +445,31 @@ impl McpClientTrait for AcpTools { self.inner.get_info() } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::agents::ToolCallContext; + + #[test] + fn terminal_request_includes_agent_session_id() { + let session_id = SessionId::new("acp-session"); + let params = ShellParams { + command: "echo test".to_string(), + timeout_secs: None, + }; + let ctx = ToolCallContext::new( + "agent-session".to_string(), + Some(std::path::PathBuf::from("/tmp/worktree")), + None, + ); + + let request = create_terminal_request(&session_id, ¶ms, &ctx); + + assert_eq!( + request.env, + vec![EnvVariable::new("AGENT_SESSION_ID", "agent-session")] + ); + assert_eq!(request.cwd, Some(std::path::PathBuf::from("/tmp/worktree"))); + } +} diff --git a/crates/goose/src/agents/platform_extensions/developer/mod.rs b/crates/goose/src/agents/platform_extensions/developer/mod.rs index 4b272b161255..095bf2c8d358 100644 --- a/crates/goose/src/agents/platform_extensions/developer/mod.rs +++ b/crates/goose/src/agents/platform_extensions/developer/mod.rs @@ -199,7 +199,7 @@ impl McpClientTrait for DeveloperClient { "shell" => match Self::parse_args::(arguments) { Ok(params) => Ok(self .shell_tool - .shell_with_cwd(params, working_dir, cancel_token) + .shell_with_cwd(params, working_dir, Some(&ctx.session_id), cancel_token) .await), Err(error) => Ok(ShellTool::error_result(&format!("Error: {error}"), None)), }, @@ -338,6 +338,29 @@ mod tests { ); } + #[cfg(not(windows))] + #[tokio::test] + async fn developer_client_passes_session_id_to_shell_tool() { + let temp = tempfile::tempdir().unwrap(); + let client = DeveloperClient::new(test_context(temp.path().join("sessions"))).unwrap(); + let ctx = ToolCallContext::new("session-789".to_owned(), None, None); + + let result = client + .call_tool( + &ctx, + "shell", + Some(object!({ + "command": "printenv AGENT_SESSION_ID" + })), + CancellationToken::new(), + ) + .await + .unwrap(); + + assert_eq!(result.is_error, Some(false)); + assert_eq!(first_text(&result), "session-789"); + } + #[cfg(not(windows))] #[tokio::test] async fn developer_client_uses_working_dir_for_shell_tool() { diff --git a/crates/goose/src/agents/platform_extensions/developer/shell.rs b/crates/goose/src/agents/platform_extensions/developer/shell.rs index fdc48ab27284..7c3bbaff1588 100644 --- a/crates/goose/src/agents/platform_extensions/developer/shell.rs +++ b/crates/goose/src/agents/platform_extensions/developer/shell.rs @@ -348,7 +348,7 @@ impl ShellTool { } pub async fn shell(&self, params: ShellParams) -> CallToolResult { - self.shell_with_cwd(params, None, CancellationToken::new()) + self.shell_with_cwd(params, None, None, CancellationToken::new()) .await } @@ -356,6 +356,7 @@ impl ShellTool { &self, params: ShellParams, working_dir: Option<&std::path::Path>, + session_id: Option<&str>, cancellation_token: CancellationToken, ) -> CallToolResult { if params.command.trim().is_empty() { @@ -374,6 +375,7 @@ impl ShellTool { params.timeout_secs, working_dir, login_path_ref, + session_id, cancellation_token, ) .await @@ -520,11 +522,12 @@ async fn run_command( timeout_secs: Option, working_dir: Option<&std::path::Path>, login_path: Option<&str>, + session_id: Option<&str>, cancellation_token: CancellationToken, ) -> Result { let timeout_secs = Some(resolve_shell_timeout(timeout_secs)); - let mut command = build_shell_command(command_line, working_dir, login_path); + let mut command = build_shell_command(command_line, working_dir, login_path, session_id); command.stdout(Stdio::piped()); command.stderr(Stdio::piped()); @@ -599,8 +602,8 @@ async fn run_command( } Err(_) => { tracing::debug!( - "output drain timed out after {OUTPUT_DRAIN_TIMEOUT_MILLIS}ms (backgrounded process?)" - ); + "output drain timed out after {OUTPUT_DRAIN_TIMEOUT_MILLIS}ms (backgrounded process?)" + ); abort_handle.abort(); true } @@ -625,6 +628,7 @@ fn build_shell_command( command_line: &str, working_dir: Option<&std::path::Path>, login_path: Option<&str>, + session_id: Option<&str>, ) -> tokio::process::Command { #[cfg(windows)] let mut command = { @@ -664,6 +668,7 @@ fn build_shell_command( if let Some(path) = login_path { command.arg(format!("--env=PATH={}", path)); } + apply_flatpak_session_environment(&mut command, session_id); command .arg(&shell) .args(unix_shell_command_args(command_line)); @@ -677,14 +682,37 @@ fn build_shell_command( if let Some(path) = login_path { command.env("PATH", path); } + apply_session_environment(&mut command, session_id); command } }; + #[cfg(windows)] + apply_session_environment(&mut command, session_id); command.set_no_window(); command } +fn apply_session_environment(command: &mut tokio::process::Command, session_id: Option<&str>) { + if let Some(session_id) = session_id.filter(|id| !id.is_empty()) { + command.env("AGENT_SESSION_ID", session_id); + } else { + command.env_remove("AGENT_SESSION_ID"); + } +} + +#[cfg(not(windows))] +fn apply_flatpak_session_environment( + command: &mut tokio::process::Command, + session_id: Option<&str>, +) { + if let Some(session_id) = session_id.filter(|id| !id.is_empty()) { + command.arg(format!("--env=AGENT_SESSION_ID={session_id}")); + } else { + command.arg("--unset-env=AGENT_SESSION_ID"); + } +} + /// Split tagged lines into (stdout, stderr, interleaved) strings. fn split_lines(lines: &[(bool, String)]) -> (String, String, String) { let mut stdout = String::new(); @@ -872,6 +900,7 @@ mod tests { timeout_secs: None, }, Some(dir.path()), + None, CancellationToken::new(), ) .await; @@ -882,6 +911,52 @@ mod tests { assert_eq!(observed, expected); } + #[test] + fn session_environment_is_set_or_removed() { + for (session_id, expected) in [ + ( + Some("session-123"), + Some(Some(std::ffi::OsStr::new("session-123"))), + ), + (None, Some(None)), + ] { + let mut command = tokio::process::Command::new("ignored"); + command.env("AGENT_SESSION_ID", "stale-session"); + + apply_session_environment(&mut command, session_id); + + assert_eq!( + command + .as_std() + .get_envs() + .find_map(|(key, value)| (key == "AGENT_SESSION_ID").then_some(value)), + expected + ); + } + } + + #[cfg(not(windows))] + #[test] + fn flatpak_session_environment_is_set_or_unset() { + for (session_id, expected) in [ + (Some("session-123"), "--env=AGENT_SESSION_ID=session-123"), + (None, "--unset-env=AGENT_SESSION_ID"), + ] { + let mut command = tokio::process::Command::new("flatpak-spawn"); + + apply_flatpak_session_environment(&mut command, session_id); + + assert_eq!( + command + .as_std() + .get_args() + .map(|arg| arg.to_string_lossy().into_owned()) + .collect::>(), + vec![expected] + ); + } + } + #[cfg(not(windows))] #[tokio::test] async fn shell_kills_child_on_cancellation() { @@ -902,6 +977,7 @@ mod tests { timeout_secs: None, }, None, + None, token, ) .await;