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
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
//! Coalescing dispatch implementation for [`super::EventCoalescer`].
use super::{
AgentEvent, CoalescedPayload, EventCoalescer, ExtensionEventName, ExtensionManager,
extension_event_name_from_agent, is_coalescable_event, is_lifecycle_event,
};
use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex};
impl EventCoalescer {
/// Create a new coalescer backed by the given extension manager.
pub fn new(manager: ExtensionManager) -> Self {
Self {
manager,
pending: Arc::new(Mutex::new(HashMap::new())),
in_flight: Arc::new(Mutex::new(HashSet::new())),
batch_buffer: Arc::new(Mutex::new(Vec::new())),
batch_drain_scheduled: Arc::new(std::sync::atomic::AtomicBool::new(false)),
}
}
/// Dispatch a fire-and-forget event, coalescing if applicable.
///
/// For coalescable events (`MessageUpdate`, `ToolExecutionUpdate`):
/// - If no dispatch is in-flight for this event type, spawns one immediately.
/// - If a dispatch is already in-flight, replaces the pending payload so
/// the in-flight task will dispatch the latest version on completion.
///
/// For non-coalescable events, buffers the event and schedules a batch
/// drain task that dispatches all buffered events in a single JS bridge
/// call. This saves ~21µs of fixed overhead per additional event in the
/// batch.
#[allow(clippy::too_many_lines)]
pub(super) fn dispatch_fire_and_forget(
&self,
event: ExtensionEventName,
data: CoalescedPayload,
runtime_handle: &asupersync::runtime::RuntimeHandle,
) {
// Fast path: skip entirely if no hooks registered. Asked through
// `as_str`, before the owned name exists, so an event nothing
// subscribes to costs no allocation at all (bd-82331).
if !self.manager.has_hook_for(event.as_str()) {
return;
}
let event_name_str = event.to_string();
if !is_coalescable_event(&event) {
// Non-coalescable: buffer for batch dispatch.
{
let mut buf = self
.batch_buffer
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
buf.push((event, data));
}
// Schedule a drain task if one isn't already pending.
if !self
.batch_drain_scheduled
.swap(true, std::sync::atomic::Ordering::AcqRel)
{
let manager = self.manager.clone();
let buffer = self.batch_buffer.clone();
let flag = self.batch_drain_scheduled.clone();
runtime_handle.spawn(async move {
loop {
// Drain the buffer; events that arrived between scheduling
// and execution are included in this batch.
let raw = {
let mut buf = buffer
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
std::mem::take(&mut *buf)
};
if raw.is_empty() {
// No work left. Release the scheduled flag, but guard against
// a race where producers appended while the flag was still true.
flag.store(false, std::sync::atomic::Ordering::Release);
let should_continue = {
if buffer
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.is_empty()
{
false
} else {
!flag.swap(true, std::sync::atomic::Ordering::AcqRel)
}
};
if should_continue {
continue;
}
break;
}
// Resolve lazy payloads off the main thread.
let events = raw
.into_iter()
.map(|(evt, payload)| (evt, payload.resolve()))
.collect::<Vec<_>>();
let _ = manager.dispatch_event_batch(events).await;
}
});
}
return;
}
// Coalescable path: check if a dispatch is already in-flight.
{
let mut in_flight = self
.in_flight
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if in_flight.contains(&event_name_str) {
// Replace pending payload; the in-flight task will pick it up.
self.pending
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(event_name_str, data);
return;
}
in_flight.insert(event_name_str.clone());
}
let manager = self.manager.clone();
let pending = self.pending.clone();
let in_flight = self.in_flight.clone();
let event_name_owned = event_name_str;
runtime_handle.spawn(async move {
let mut next_payload = Some(data);
loop {
let Some(payload) = next_payload.take() else {
break;
};
// Re-parse the event name back.
let dispatch_event = match event_name_owned.as_str() {
"message_update" => ExtensionEventName::MessageUpdate,
"tool_execution_update" => ExtensionEventName::ToolExecutionUpdate,
_ => {
in_flight
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.remove(&event_name_owned);
break;
}
};
let _ = manager
.dispatch_event(dispatch_event, payload.resolve())
.await;
// Fast path: drain pending replacement payload if present.
if let Some(new_data) = {
let mut p = pending
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
p.remove(&event_name_owned)
} {
next_payload = Some(new_data);
continue;
}
// Hand off atomically with writers (which lock in_flight then pending)
// so we don't strand a payload that arrives right before completion.
let maybe_new_data = {
let mut f = in_flight
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let mut p = pending
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
p.remove(&event_name_owned).or_else(|| {
f.remove(&event_name_owned);
None
})
};
if let Some(new_data) = maybe_new_data {
next_payload = Some(new_data);
continue;
}
break;
}
});
}
/// Like [`dispatch_fire_and_forget`](Self::dispatch_fire_and_forget) but
/// takes the raw [`AgentEvent`] and defers serialization until after
/// verifying that a hook is actually registered. This avoids the
/// `serde_json::to_value()` cost (~2-5µs) for events that no extension
/// listens to.
pub fn dispatch_agent_event_lazy(
&self,
event: &AgentEvent,
runtime_handle: &asupersync::runtime::RuntimeHandle,
) {
let Some(event_name) = extension_event_name_from_agent(event) else {
return;
};
if is_lifecycle_event(&event_name) {
return;
}
if !self.manager.has_hook_for(event_name.as_str()) {
return;
}
// Hook exists — defer serialization to the async task.
let event_clone = event.clone();
let lazy = Box::new(move || serde_json::to_value(&event_clone).ok());
self.dispatch_fire_and_forget(event_name, CoalescedPayload::Lazy(lazy), runtime_handle);
}
}
#[cfg(test)]
mod tests {
use super::{
AgentEvent, ExtensionEventName, extension_event_name_from_agent, is_lifecycle_event,
};
use crate::model::{
AssistantMessage, AssistantMessageEvent, Message, UserContent, UserMessage,
};
use crate::tools::ToolOutput;
use std::sync::Arc;
/// The two-route split this file depends on, pinned (bd-82331).
///
/// Lifecycle events are dispatched to extensions from inside the agent
/// loop, so the coalescer must skip them or every surface that installs
/// one would deliver them twice. Observation events are the coalescer's
/// job, and a surface that does not install one delivers none of them —
/// which is the whole of bd-82331. If this classification ever changes,
/// both halves change with it and this test is where that shows up.
#[test]
fn lifecycle_events_are_skipped_and_observation_events_are_not() {
let lifecycle = [
ExtensionEventName::AgentStart,
ExtensionEventName::AgentEnd,
ExtensionEventName::TurnStart,
ExtensionEventName::TurnEnd,
];
for name in lifecycle {
assert!(
is_lifecycle_event(&name),
"{name:?} must stay a lifecycle event: it is dispatched in-loop, and treating it \
as an observation event would double-deliver it on every surface that installs a \
coalescer"
);
}
let observation = [
ExtensionEventName::MessageStart,
ExtensionEventName::MessageUpdate,
ExtensionEventName::MessageEnd,
ExtensionEventName::ToolExecutionStart,
ExtensionEventName::ToolExecutionUpdate,
ExtensionEventName::ToolExecutionEnd,
];
for name in observation {
assert!(
!is_lifecycle_event(&name),
"{name:?} must stay an observation event: it reaches extensions only through a \
surface-installed coalescer (bd-82331)"
);
}
}
/// Every observation event an agent can emit maps to the extension event
/// name extensions subscribe to (bd-82331).
///
/// `dispatch_agent_event_lazy` returns early when
/// `extension_event_name_from_agent` yields `None`, which is correct for
/// the events extensions have no concept of (`auto_retry_*`,
/// `failover_*`, `provider_error`, …). The match in that function is
/// exhaustive, so a *new* `AgentEvent` variant cannot be forgotten — the
/// compiler demands an arm. What the compiler cannot catch is an existing
/// observation variant being folded into the trailing `=> None` group, or
/// wired to the wrong name: both compile, both are silent, and both stop
/// extensions receiving an event they are registered for with no error
/// anywhere. So assert the exact name, for all six, rather than merely
/// `is_some()`.
#[test]
fn every_observation_event_maps_to_an_extension_event_name() {
let message = || {
Message::User(UserMessage {
content: UserContent::Text("probe".to_string()),
timestamp: 1_700_000_000,
})
};
let tool_output = || ToolOutput {
content: Vec::new(),
details: None,
is_error: false,
};
let events = [
(
AgentEvent::MessageStart { message: message() },
ExtensionEventName::MessageStart,
),
(
AgentEvent::MessageUpdate {
message: message(),
assistant_message_event: AssistantMessageEvent::Start {
partial: Arc::new(AssistantMessage::default()),
},
},
ExtensionEventName::MessageUpdate,
),
(
AgentEvent::MessageEnd { message: message() },
ExtensionEventName::MessageEnd,
),
(
AgentEvent::ToolExecutionStart {
tool_call_id: "call-1".to_string(),
tool_name: "probe".to_string(),
args: serde_json::Value::Null,
},
ExtensionEventName::ToolExecutionStart,
),
(
AgentEvent::ToolExecutionUpdate {
tool_call_id: "call-1".to_string(),
tool_name: "probe".to_string(),
args: serde_json::Value::Null,
partial_result: tool_output(),
},
ExtensionEventName::ToolExecutionUpdate,
),
(
AgentEvent::ToolExecutionEnd {
tool_call_id: "call-1".to_string(),
tool_name: "probe".to_string(),
result: tool_output(),
is_error: false,
},
ExtensionEventName::ToolExecutionEnd,
),
];
for (event, expected) in events {
let actual = extension_event_name_from_agent(&event);
assert_eq!(
actual,
Some(expected),
"this observation event must map to {expected:?}; a None means extensions \
silently stop seeing it and a different name means it is delivered to the \
wrong subscribers — neither reports the loss (bd-82331)"
);
}
}
}