1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
//! Event-stream bridge for direct host tool invocations.
use super::*;
use tokio::sync::{broadcast, oneshot};
impl DirectToolRuntime {
pub(in crate::agent_api) fn spawn_call_with_agent_events(
self,
name: String,
args: serde_json::Value,
) -> (
mpsc::Receiver<AgentEvent>,
JoinHandle<Result<ToolCallResult>>,
) {
if let Err(error) = self.ensure_open() {
let (tx, rx) = mpsc::channel(1);
drop(tx);
return (rx, tokio::spawn(async move { Err(error) }));
}
let (agent_tx, agent_rx) = mpsc::channel::<AgentEvent>(256);
let (runtime_tx, mut runtime_rx) = mpsc::channel::<AgentEvent>(256);
let (broadcast_tx, mut broadcast_rx) = broadcast::channel::<AgentEvent>(256);
let (broadcast_shutdown_tx, mut broadcast_shutdown_rx) = oneshot::channel::<()>();
let security_provider = self.security_provider.clone();
let forward_tx = agent_tx.clone();
let forwarder = tokio::spawn(async move {
let mut sanitizer = crate::security::AgentEventStreamSanitizer::new(security_provider);
let mut runtime_open = true;
let mut broadcast_open = true;
let mut broadcast_drain_remaining = None;
while runtime_open || broadcast_open {
let event = tokio::select! {
biased;
_ = &mut broadcast_shutdown_rx,
if broadcast_open && broadcast_drain_remaining.is_none() =>
{
// A detached producer (notably `task` with
// `background: true`) may retain its sender long after
// the direct tool invocation returns. Snapshot the
// already-accepted backlog at the invocation boundary
// and drain only that prefix instead of waiting for all
// producer clones to disappear.
let pending = broadcast_rx.len();
if pending == 0 {
broadcast_open = false;
} else {
broadcast_drain_remaining = Some(pending);
}
None
}
event = runtime_rx.recv(), if runtime_open => {
match event {
Some(event) => Some(event),
None => {
runtime_open = false;
None
}
}
}
event = broadcast_rx.recv(), if broadcast_open => {
match event {
Ok(event) => {
if let Some(remaining) = &mut broadcast_drain_remaining {
*remaining = remaining.saturating_sub(1);
if *remaining == 0 {
broadcast_open = false;
}
}
Some(event)
}
Err(broadcast::error::RecvError::Lagged(skipped)) => {
tracing::warn!(skipped, "direct tool event bridge lagged");
if let Some(remaining) = &mut broadcast_drain_remaining {
let skipped = usize::try_from(skipped).unwrap_or(usize::MAX);
*remaining = remaining.saturating_sub(skipped);
if *remaining == 0 {
broadcast_open = false;
}
}
None
}
Err(broadcast::error::RecvError::Closed) => {
broadcast_open = false;
None
}
}
}
};
let Some(event) = event else {
continue;
};
for event in sanitizer.push(event) {
if forward_tx.send(event).await.is_err() {
return;
}
}
}
for event in sanitizer.finish() {
if forward_tx.send(event).await.is_err() {
return;
}
}
});
let agent_loop = self.agent_loop;
let session_id = self.session_id;
let cancel = self.session_cancel.child_token();
let mut ctx = self.tool_context;
ctx.agent_event_tx = Some(broadcast_tx);
let tool_name = name.clone();
let tool_id = format!("host-{tool_name}-{}", uuid::Uuid::new_v4());
let event_args = args.clone();
let event_tx = Some(runtime_tx);
let security_provider = self.security_provider;
let handle = tokio::spawn(async move {
let result = agent_loop
.invoke_host_tool(
ToolInvocation::host_direct(tool_id.clone(), tool_name.clone(), args),
&session_id,
&event_tx,
&cancel,
&ctx,
)
.await;
// Runtime deltas have an owned producer lifecycle, so closing that
// path still guarantees all tool-stream deltas precede ToolEnd.
// The broadcast path uses an explicit invocation boundary because
// detached background tasks intentionally outlive this call.
broadcast_shutdown_tx.send(()).ok();
drop(event_tx);
drop(ctx);
if let Err(error) = forwarder.await {
tracing::warn!(%error, "direct tool event bridge failed");
}
let end = AgentEvent::ToolEnd {
id: tool_id,
name: tool_name,
args: Some(event_args),
output: result.output.clone(),
exit_code: result.exit_code,
metadata: result.metadata.clone(),
error_kind: result.error_kind.clone(),
};
let end = security_provider
.as_deref()
.map(|provider| crate::security::sanitize_agent_event(provider, &end))
.unwrap_or(end);
agent_tx.send(end).await.ok();
Ok(project_tool_result(result))
});
(agent_rx, handle)
}
}