tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

fix(developer): expose AGENT_SESSION_ID to shell commands (#10428)

Browse Source 
Co-authored-by: John Tennant <jtennant@block.xyz>
This commit is contained in:
John Matthew Tennant
2026-07-17 12:22:04 -04:00
committed by GitHub
parent ceb94b24dc
commit 8d5bc5d497
3 changed files with 148 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/acp/fs.rs
+44 -9
View File
@@ -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(), &params.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(), &params.command)
.cwd(ctx.working_dir.clone())
.output_byte_limit(OUTPUT_LIMIT_BYTES as u64),
)
.send_request(create_terminal_request(&self.session_id, &params, 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, &params, &ctx);
assert_eq!(
request.env,
vec![EnvVariable::new("AGENT_SESSION_ID", "agent-session")]
);
assert_eq!(request.cwd, Some(std::path::PathBuf::from("/tmp/worktree")));
}
}
crates/goose/src/agents/platform_extensions/developer/mod.rs
+24 -1
View File
@@ -199,7 +199,7 @@ impl McpClientTrait for DeveloperClient {
"shell" => match Self::parse_args::<ShellParams>(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() {
crates/goose/src/agents/platform_extensions/developer/shell.rs
+80 -4
View File
@@ -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<u64>,
working_dir: Option<&std::path::Path>,
login_path: Option<&str>,
session_id: Option<&str>,
cancellation_token: CancellationToken,
) -> Result<ExecutionOutput, String> {
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<_>>(),
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;