Skip to content

agent: Support streaming tool input #50099

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bennetbo merged 2 commits into from
Feb 25, 2026
Merged

agent: Support streaming tool input #50099

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
01a3bfa
agent: Support streaming tool input
bennetbo Feb 25, 2026
6c5d458
Remove unrelated changes
bennetbo Feb 25, 2026
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
66 changes: 33 additions & 33 deletions crates/agent/src/tests/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -310,11 +310,11 @@ async fn test_terminal_tool_timeout_kills_handle(cx: &mut TestAppContext) {

let task = cx.update(|cx| {
tool.run(
crate::TerminalToolInput {
ToolInput::resolved(crate::TerminalToolInput {
command: "sleep 1000".to_string(),
cd: ".".to_string(),
timeout_ms: Some(5),
},
}),
event_stream,
cx,
)
Expand Down Expand Up @@ -377,11 +377,11 @@ async fn test_terminal_tool_without_timeout_does_not_kill_handle(cx: &mut TestAp

let _task = cx.update(|cx| {
tool.run(
crate::TerminalToolInput {
ToolInput::resolved(crate::TerminalToolInput {
command: "sleep 1000".to_string(),
cd: ".".to_string(),
timeout_ms: None,
},
}),
event_stream,
cx,
)
Expand Down Expand Up @@ -3991,11 +3991,11 @@ async fn test_terminal_tool_permission_rules(cx: &mut TestAppContext) {

let task = cx.update(|cx| {
tool.run(
crate::TerminalToolInput {
ToolInput::resolved(crate::TerminalToolInput {
command: "rm -rf /".to_string(),
cd: ".".to_string(),
timeout_ms: None,
},
}),
event_stream,
cx,
)
Expand Down Expand Up @@ -4043,11 +4043,11 @@ async fn test_terminal_tool_permission_rules(cx: &mut TestAppContext) {

let task = cx.update(|cx| {
tool.run(
crate::TerminalToolInput {
ToolInput::resolved(crate::TerminalToolInput {
command: "echo hello".to_string(),
cd: ".".to_string(),
timeout_ms: None,
},
}),
event_stream,
cx,
)
Expand Down Expand Up @@ -4101,11 +4101,11 @@ async fn test_terminal_tool_permission_rules(cx: &mut TestAppContext) {

let _task = cx.update(|cx| {
tool.run(
crate::TerminalToolInput {
ToolInput::resolved(crate::TerminalToolInput {
command: "sudo rm file".to_string(),
cd: ".".to_string(),
timeout_ms: None,
},
}),
event_stream,
cx,
)
Expand Down Expand Up @@ -4148,11 +4148,11 @@ async fn test_terminal_tool_permission_rules(cx: &mut TestAppContext) {

let task = cx.update(|cx| {
tool.run(
crate::TerminalToolInput {
ToolInput::resolved(crate::TerminalToolInput {
command: "echo hello".to_string(),
cd: ".".to_string(),
timeout_ms: None,
},
}),
event_stream,
cx,
)
Expand Down Expand Up @@ -5309,11 +5309,11 @@ async fn test_edit_file_tool_deny_rule_blocks_edit(cx: &mut TestAppContext) {

let task = cx.update(|cx| {
tool.run(
crate::EditFileToolInput {
ToolInput::resolved(crate::EditFileToolInput {
display_description: "Edit sensitive file".to_string(),
path: "root/sensitive_config.txt".into(),
mode: crate::EditFileMode::Edit,
},
}),
event_stream,
cx,
)
Expand Down Expand Up @@ -5359,9 +5359,9 @@ async fn test_delete_path_tool_deny_rule_blocks_deletion(cx: &mut TestAppContext

let task = cx.update(|cx| {
tool.run(
crate::DeletePathToolInput {
ToolInput::resolved(crate::DeletePathToolInput {
path: "root/important_data.txt".to_string(),
},
}),
event_stream,
cx,
)
Expand Down Expand Up @@ -5411,10 +5411,10 @@ async fn test_move_path_tool_denies_if_destination_denied(cx: &mut TestAppContex

let task = cx.update(|cx| {
tool.run(
crate::MovePathToolInput {
ToolInput::resolved(crate::MovePathToolInput {
source_path: "root/safe.txt".to_string(),
destination_path: "root/protected/safe.txt".to_string(),
},
}),
event_stream,
cx,
)
Expand Down Expand Up @@ -5467,10 +5467,10 @@ async fn test_move_path_tool_denies_if_source_denied(cx: &mut TestAppContext) {

let task = cx.update(|cx| {
tool.run(
crate::MovePathToolInput {
ToolInput::resolved(crate::MovePathToolInput {
source_path: "root/secret.txt".to_string(),
destination_path: "root/public/not_secret.txt".to_string(),
},
}),
event_stream,
cx,
)
Expand Down Expand Up @@ -5525,10 +5525,10 @@ async fn test_copy_path_tool_deny_rule_blocks_copy(cx: &mut TestAppContext) {

let task = cx.update(|cx| {
tool.run(
crate::CopyPathToolInput {
ToolInput::resolved(crate::CopyPathToolInput {
source_path: "root/confidential.txt".to_string(),
destination_path: "root/dest/copy.txt".to_string(),
},
}),
event_stream,
cx,
)
Expand Down Expand Up @@ -5580,12 +5580,12 @@ async fn test_save_file_tool_denies_if_any_path_denied(cx: &mut TestAppContext)

let task = cx.update(|cx| {
tool.run(
crate::SaveFileToolInput {
ToolInput::resolved(crate::SaveFileToolInput {
paths: vec![
std::path::PathBuf::from("root/normal.txt"),
std::path::PathBuf::from("root/readonly/config.txt"),
],
},
}),
event_stream,
cx,
)
Expand Down Expand Up @@ -5632,9 +5632,9 @@ async fn test_save_file_tool_respects_deny_rules(cx: &mut TestAppContext) {

let task = cx.update(|cx| {
tool.run(
crate::SaveFileToolInput {
ToolInput::resolved(crate::SaveFileToolInput {
paths: vec![std::path::PathBuf::from("root/config.secret")],
},
}),
event_stream,
cx,
)
Expand Down Expand Up @@ -5676,7 +5676,7 @@ async fn test_web_search_tool_deny_rule_blocks_search(cx: &mut TestAppContext) {
let input: crate::WebSearchToolInput =
serde_json::from_value(json!({"query": "internal.company.com secrets"})).unwrap();

let task = cx.update(|cx| tool.run(input, event_stream, cx));
let task = cx.update(|cx| tool.run(ToolInput::resolved(input), event_stream, cx));

let result = task.await;
assert!(result.is_err(), "expected search to be blocked");
Expand Down Expand Up @@ -5741,11 +5741,11 @@ async fn test_edit_file_tool_allow_rule_skips_confirmation(cx: &mut TestAppConte

let _task = cx.update(|cx| {
tool.run(
crate::EditFileToolInput {
ToolInput::resolved(crate::EditFileToolInput {
display_description: "Edit README".to_string(),
path: "root/README.md".into(),
mode: crate::EditFileMode::Edit,
},
}),
event_stream,
cx,
)
Expand Down Expand Up @@ -5811,11 +5811,11 @@ async fn test_edit_file_tool_allow_still_prompts_for_local_settings(cx: &mut Tes
let (event_stream, mut rx) = crate::ToolCallEventStream::test();
let _task = cx.update(|cx| {
tool.run(
crate::EditFileToolInput {
ToolInput::resolved(crate::EditFileToolInput {
display_description: "Edit local settings".to_string(),
path: "root/.zed/settings.json".into(),
mode: crate::EditFileMode::Edit,
},
}),
event_stream,
cx,
)
Expand Down Expand Up @@ -5855,7 +5855,7 @@ async fn test_fetch_tool_deny_rule_blocks_url(cx: &mut TestAppContext) {
let input: crate::FetchToolInput =
serde_json::from_value(json!({"url": "https://internal.company.com/api"})).unwrap();

let task = cx.update(|cx| tool.run(input, event_stream, cx));
let task = cx.update(|cx| tool.run(ToolInput::resolved(input), event_stream, cx));

let result = task.await;
assert!(result.is_err(), "expected fetch to be blocked");
Expand Down Expand Up @@ -5893,7 +5893,7 @@ async fn test_fetch_tool_allow_rule_skips_confirmation(cx: &mut TestAppContext)
let input: crate::FetchToolInput =
serde_json::from_value(json!({"url": "https://docs.rs/some-crate"})).unwrap();

let _task = cx.update(|cx| tool.run(input, event_stream, cx));
let _task = cx.update(|cx| tool.run(ToolInput::resolved(input), event_stream, cx));

cx.run_until_parked();

Expand Down
89 changes: 62 additions & 27 deletions crates/agent/src/tests/test_tools.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ use agent_settings::AgentSettings;
use gpui::{App, SharedString, Task};
use std::future;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;

/// A tool that echoes its input
#[derive(JsonSchema, Serialize, Deserialize)]
Expand Down Expand Up @@ -33,11 +34,17 @@ impl AgentTool for EchoTool {

fn run(
self: Arc<Self>,
input: Self::Input,
input: ToolInput<Self::Input>,
_event_stream: ToolCallEventStream,
_cx: &mut App,
cx: &mut App,
) -> Task<Result<String, String>> {
Task::ready(Ok(input.text))
cx.spawn(async move |_cx| {
let input = input
.recv()
.await
.map_err(|e| format!("Failed to receive tool input: {e}"))?;
Ok(input.text)
})
}
}

Expand Down Expand Up @@ -74,7 +81,7 @@ impl AgentTool for DelayTool {

fn run(
self: Arc<Self>,
input: Self::Input,
input: ToolInput<Self::Input>,
_event_stream: ToolCallEventStream,
cx: &mut App,
) -> Task<Result<String, String>>
Expand All @@ -83,6 +90,10 @@ impl AgentTool for DelayTool {
{
let executor = cx.background_executor().clone();
cx.foreground_executor().spawn(async move {
let input = input
.recv()
.await
.map_err(|e| format!("Failed to receive tool input: {e}"))?;
executor.timer(Duration::from_millis(input.ms)).await;
Ok("Ding".to_string())
})
Expand Down Expand Up @@ -114,28 +125,38 @@ impl AgentTool for ToolRequiringPermission {

fn run(
self: Arc<Self>,
_input: Self::Input,
input: ToolInput<Self::Input>,
event_stream: ToolCallEventStream,
cx: &mut App,
) -> Task<Result<String, String>> {
let settings = AgentSettings::get_global(cx);
let decision = decide_permission_from_settings(Self::NAME, &[String::new()], settings);

let authorize = match decision {
ToolPermissionDecision::Allow => None,
ToolPermissionDecision::Deny(reason) => {
return Task::ready(Err(reason));
}
ToolPermissionDecision::Confirm => {
let context = crate::ToolPermissionContext::new(
"tool_requiring_permission",
vec![String::new()],
);
Some(event_stream.authorize("Authorize?", context, cx))
}
};
cx.spawn(async move |cx| {
let _input = input
.recv()
.await
.map_err(|e| format!("Failed to receive tool input: {e}"))?;

let decision = cx.update(|cx| {
decide_permission_from_settings(
Self::NAME,
&[String::new()],
AgentSettings::get_global(cx),
)
});

let authorize = match decision {
ToolPermissionDecision::Allow => None,
ToolPermissionDecision::Deny(reason) => {
return Err(reason);
}
ToolPermissionDecision::Confirm => Some(cx.update(|cx| {
let context = crate::ToolPermissionContext::new(
"tool_requiring_permission",
vec![String::new()],
);
event_stream.authorize("Authorize?", context, cx)
})),
};

cx.foreground_executor().spawn(async move {
if let Some(authorize) = authorize {
authorize.await.map_err(|e| e.to_string())?;
}
Expand Down Expand Up @@ -169,11 +190,15 @@ impl AgentTool for InfiniteTool {

fn run(
self: Arc<Self>,
_input: Self::Input,
input: ToolInput<Self::Input>,
_event_stream: ToolCallEventStream,
cx: &mut App,
) -> Task<Result<String, String>> {
cx.foreground_executor().spawn(async move {
let _input = input
.recv()
.await
.map_err(|e| format!("Failed to receive tool input: {e}"))?;
future::pending::<()>().await;
unreachable!()
})
Expand Down Expand Up @@ -221,11 +246,15 @@ impl AgentTool for CancellationAwareTool {

fn run(
self: Arc<Self>,
_input: Self::Input,
input: ToolInput<Self::Input>,
event_stream: ToolCallEventStream,
cx: &mut App,
) -> Task<Result<String, String>> {
cx.foreground_executor().spawn(async move {
let _input = input
.recv()
.await
.map_err(|e| format!("Failed to receive tool input: {e}"))?;
// Wait for cancellation - this tool does nothing but wait to be cancelled
event_stream.cancelled_by_user().await;
self.was_cancelled.store(true, Ordering::SeqCst);
Expand Down Expand Up @@ -276,10 +305,16 @@ impl AgentTool for WordListTool {

fn run(
self: Arc<Self>,
_input: Self::Input,
input: ToolInput<Self::Input>,
_event_stream: ToolCallEventStream,
_cx: &mut App,
cx: &mut App,
) -> Task<Result<String, String>> {
Task::ready(Ok("ok".to_string()))
cx.spawn(async move |_cx| {
let _input = input
.recv()
.await
.map_err(|e| format!("Failed to receive tool input: {e}"))?;
Ok("ok".to_string())
})
}
}
Loading