1use std::cell::RefCell;
11use std::collections::BTreeMap;
12use std::sync::atomic::{AtomicU64, Ordering};
13use std::sync::{Arc, LazyLock, Mutex};
14
15use crate::agent_events::{
16 self, AgentEvent, AgentEventSink, ToolCallErrorCategory, ToolCallStatus, ToolMutationStatus,
17};
18use crate::mcp::VmMcpClientHandle;
19use crate::value::VmValue;
20
21const LOOP_EXIT_ABANDON_REASON: &str =
24 "The agent loop exited while this tool call was still in flight; it was \
25 never dispatched to a result. No terminal update was emitted by the \
26 dispatch path, so the loop resolved it as abandoned at loop exit.";
27
28pub type SessionEndHook = Arc<dyn Fn(&str) + Send + Sync>;
30
31thread_local! {
32 static CURRENT_HOST_BRIDGE: RefCell<Option<Arc<crate::bridge::HostBridge>>> =
33 const { RefCell::new(None) };
34 static CURRENT_LOOP_SINKS: RefCell<Vec<Arc<dyn AgentEventSink>>> =
45 const { RefCell::new(Vec::new()) };
46}
47
48static NEXT_SESSION_END_HOOK_ID: AtomicU64 = AtomicU64::new(1);
52static SESSION_END_HOOKS: LazyLock<Mutex<BTreeMap<u64, SessionEndHook>>> =
53 LazyLock::new(|| Mutex::new(BTreeMap::new()));
54
55#[must_use = "dropping the registration unregisters the session-end hook"]
61pub struct SessionEndHookRegistration {
62 id: u64,
63}
64
65impl Drop for SessionEndHookRegistration {
66 fn drop(&mut self) {
67 SESSION_END_HOOKS
68 .lock()
69 .unwrap_or_else(std::sync::PoisonError::into_inner)
70 .remove(&self.id);
71 }
72}
73
74static SESSION_MCP_CLIENTS: LazyLock<Mutex<BTreeMap<String, BTreeMap<String, VmMcpClientHandle>>>> =
75 LazyLock::new(|| Mutex::new(BTreeMap::new()));
76
77#[derive(Default)]
78struct ToolLifecycleStarts {
79 live: BTreeMap<(String, String), String>,
85}
86
87impl ToolLifecycleStarts {
88 fn observe(&mut self, event: &AgentEvent) -> bool {
89 match event {
90 AgentEvent::ToolCall {
91 session_id,
92 tool_call_id,
93 tool_name,
94 ..
95 } => {
96 if tool_call_id.trim().is_empty() {
97 return true;
98 }
99 self.live
100 .insert(
101 (session_id.clone(), tool_call_id.clone()),
102 tool_name.clone(),
103 )
104 .is_none()
105 }
106 AgentEvent::ToolCallUpdate {
107 session_id,
108 tool_call_id,
109 status: ToolCallStatus::Completed | ToolCallStatus::Failed,
110 ..
111 } => {
112 self.live
113 .remove(&(session_id.clone(), tool_call_id.clone()));
114 true
115 }
116 _ => true,
117 }
118 }
119
120 fn drain_session(&mut self, session_id: &str) -> Vec<(String, String)> {
125 let live = std::mem::take(&mut self.live);
126 let mut drained = Vec::new();
127 for ((active_session_id, tool_call_id), tool_name) in live {
128 if active_session_id == session_id {
129 drained.push((tool_call_id, tool_name));
130 } else {
131 self.live
132 .insert((active_session_id, tool_call_id), tool_name);
133 }
134 }
135 drained
136 }
137
138 fn clear_session(&mut self, session_id: &str) {
142 let _ = self.drain_session(session_id);
143 }
144}
145
146static TOOL_LIFECYCLE_STARTS: LazyLock<Mutex<ToolLifecycleStarts>> =
147 LazyLock::new(|| Mutex::new(ToolLifecycleStarts::default()));
148
149pub(crate) fn reset_session_state() {
163 TOOL_LIFECYCLE_STARTS
164 .lock()
165 .unwrap_or_else(std::sync::PoisonError::into_inner)
166 .live
167 .clear();
168 SESSION_MCP_CLIENTS
169 .lock()
170 .unwrap_or_else(std::sync::PoisonError::into_inner)
171 .clear();
172}
173
174fn observe_tool_lifecycle(event: &AgentEvent) -> bool {
175 TOOL_LIFECYCLE_STARTS
176 .lock()
177 .unwrap_or_else(std::sync::PoisonError::into_inner)
178 .observe(event)
179}
180
181pub(crate) struct LoopSinkGuard {
184 pushed: bool,
185}
186
187impl LoopSinkGuard {
188 pub(crate) fn install(sink: Option<Arc<dyn AgentEventSink>>) -> Self {
189 if let Some(sink) = sink {
190 CURRENT_LOOP_SINKS.with(|stack| stack.borrow_mut().push(sink));
191 Self { pushed: true }
192 } else {
193 Self { pushed: false }
194 }
195 }
196}
197
198impl Drop for LoopSinkGuard {
199 fn drop(&mut self) {
200 if self.pushed {
201 CURRENT_LOOP_SINKS.with(|stack| {
202 let _ = stack.borrow_mut().pop();
203 });
204 }
205 }
206}
207
208pub(crate) fn emit_agent_event_sync(event: &AgentEvent) {
223 if observe_tool_lifecycle(event) {
224 agent_events::emit_event(event);
225 let loop_sink = CURRENT_LOOP_SINKS.with(|stack| stack.borrow().last().cloned());
226 if let Some(sink) = loop_sink {
227 sink.handle_event(event);
228 }
229 }
230}
231
232pub async fn scope_agent_event_sink<F, T>(sink: Option<Arc<dyn AgentEventSink>>, future: F) -> T
240where
241 F: std::future::Future<Output = T>,
242{
243 let _guard = LoopSinkGuard::install(sink);
244 future.await
245}
246
247pub(crate) async fn emit_agent_event_with_ctx(
259 ctx: Option<&crate::vm::AsyncBuiltinCtx>,
260 event: &AgentEvent,
261) {
262 if observe_tool_lifecycle(event) {
263 agent_events::emit_event(event);
264
265 let loop_sink = CURRENT_LOOP_SINKS.with(|stack| stack.borrow().last().cloned());
266 if let Some(sink) = loop_sink {
267 sink.handle_event(event);
268 }
269 }
270
271 let subscribers = crate::agent_sessions::subscribers_for(event.session_id());
272 if subscribers.is_empty() {
273 return;
274 }
275 let payload = serde_json::to_value(event).unwrap_or(serde_json::Value::Null);
276 let arg = crate::stdlib::json_to_vm_value(&payload);
277 for closure in subscribers {
278 let VmValue::Closure(closure) = closure else {
279 continue;
280 };
281 let Some(ctx) = ctx else {
282 continue;
283 };
284 let mut vm = ctx.child_vm();
285 let result = vm.call_closure_pub(&closure, &[arg.clone()]).await;
288 ctx.forward_output(&vm.take_output());
289 if let Err(err) = result {
290 crate::events::log_warn(
291 "agent.subscriber",
292 &format!(
293 "session={} event={:?} subscriber error: {}",
294 event.session_id(),
295 std::mem::discriminant(event),
296 err
297 ),
298 );
299 }
300 }
301}
302
303pub fn register_session_end_hook(hook: SessionEndHook) -> SessionEndHookRegistration {
315 let id = NEXT_SESSION_END_HOOK_ID.fetch_add(1, Ordering::Relaxed);
316 SESSION_END_HOOKS
317 .lock()
318 .unwrap_or_else(std::sync::PoisonError::into_inner)
319 .insert(id, hook);
320 SessionEndHookRegistration { id }
321}
322
323pub(crate) fn fire_session_end_hooks(session_id: &str, abandon_in_flight: bool) {
343 let abandoned = {
347 let mut tracker = TOOL_LIFECYCLE_STARTS
348 .lock()
349 .unwrap_or_else(std::sync::PoisonError::into_inner);
350 if abandon_in_flight {
351 tracker.drain_session(session_id)
352 } else {
353 tracker.clear_session(session_id);
356 Vec::new()
357 }
358 };
359 for (tool_call_id, tool_name) in abandoned {
360 emit_agent_event_sync(&AgentEvent::ToolCallUpdate {
361 session_id: session_id.to_string(),
362 tool_call_id,
363 tool_name,
364 status: ToolCallStatus::Failed,
365 raw_output: None,
366 error: Some(LOOP_EXIT_ABANDON_REASON.to_string()),
367 duration_ms: None,
368 execution_duration_ms: None,
369 error_category: Some(ToolCallErrorCategory::AbandonedAtLoopExit),
370 mutation_status: ToolMutationStatus::Unknown,
371 changed_paths: None,
372 executor: None,
373 parsing: None,
374 raw_input: None,
375 raw_input_partial: None,
376 audit: None,
377 });
378 }
379 let hooks = SESSION_END_HOOKS
380 .lock()
381 .unwrap_or_else(std::sync::PoisonError::into_inner)
382 .values()
383 .cloned()
384 .collect::<Vec<_>>();
385 for hook in hooks {
386 hook(session_id);
387 }
388}
389
390pub(crate) fn install_current_host_bridge(bridge: Arc<crate::bridge::HostBridge>) {
391 CURRENT_HOST_BRIDGE.with(|slot| {
392 *slot.borrow_mut() = Some(bridge);
393 });
394}
395
396pub(crate) fn clear_current_host_bridge() {
397 CURRENT_HOST_BRIDGE.with(|slot| {
398 *slot.borrow_mut() = None;
399 });
400}
401
402pub(crate) fn swap_current_host_bridge(
403 bridge: Option<Arc<crate::bridge::HostBridge>>,
404) -> Option<Arc<crate::bridge::HostBridge>> {
405 CURRENT_HOST_BRIDGE.with(|slot| std::mem::replace(&mut *slot.borrow_mut(), bridge))
406}
407
408pub(crate) fn current_host_bridge() -> Option<Arc<crate::bridge::HostBridge>> {
409 CURRENT_HOST_BRIDGE.with(|slot| slot.borrow().clone())
410}
411
412pub fn current_agent_session_id() -> Option<String> {
416 crate::agent_sessions::current_session_id()
417}
418
419pub(crate) fn install_session_mcp_clients(
426 session_id: &str,
427 clients: BTreeMap<String, VmMcpClientHandle>,
428) {
429 if let Ok(mut map) = SESSION_MCP_CLIENTS.lock() {
430 let existing = map.entry(session_id.to_string()).or_default();
431 for (name, handle) in clients {
432 existing.insert(name, handle);
433 }
434 }
435}
436
437pub(crate) fn take_session_mcp_clients(
438 session_id: &str,
439) -> Option<BTreeMap<String, VmMcpClientHandle>> {
440 SESSION_MCP_CLIENTS
441 .lock()
442 .ok()
443 .and_then(|mut map| map.remove(session_id))
444}
445
446pub(crate) fn session_mcp_client(session_id: &str, server_name: &str) -> Option<VmMcpClientHandle> {
447 SESSION_MCP_CLIENTS.lock().ok().and_then(|map| {
448 map.get(session_id)
449 .and_then(|clients| clients.get(server_name))
450 .cloned()
451 })
452}
453
454#[cfg(test)]
455mod tests {
456 use super::*;
457 use crate::agent_events::{ToolCallErrorCategory, ToolCallStatus, ToolMutationStatus};
458 use serde_json::json;
459 use std::collections::BTreeSet;
460 use std::sync::atomic::AtomicUsize;
461
462 #[derive(Default)]
463 struct RecordingSink {
464 events: Mutex<Vec<AgentEvent>>,
465 }
466
467 impl AgentEventSink for RecordingSink {
468 fn handle_event(&self, event: &AgentEvent) {
469 self.events
470 .lock()
471 .expect("recording sink")
472 .push(event.clone());
473 }
474 }
475
476 fn start(session_id: &str, tool_call_id: &str) -> AgentEvent {
477 AgentEvent::ToolCall {
478 session_id: session_id.to_string(),
479 tool_call_id: tool_call_id.to_string(),
480 tool_name: "verify".to_string(),
481 kind: None,
482 status: ToolCallStatus::Pending,
483 raw_input: json!({}),
484 parsing: None,
485 audit: None,
486 }
487 }
488
489 fn finish(session_id: &str, tool_call_id: &str) -> AgentEvent {
490 AgentEvent::ToolCallUpdate {
491 session_id: session_id.to_string(),
492 tool_call_id: tool_call_id.to_string(),
493 tool_name: "verify".to_string(),
494 status: ToolCallStatus::Completed,
495 raw_output: None,
496 error: None,
497 duration_ms: None,
498 execution_duration_ms: None,
499 error_category: None,
500 mutation_status: ToolMutationStatus::Unknown,
501 changed_paths: None,
502 executor: None,
503 parsing: None,
504 raw_input: None,
505 raw_input_partial: None,
506 audit: None,
507 }
508 }
509
510 #[test]
511 fn session_end_hook_process_global_lifetime_is_guard_owned_across_resets() {
512 const SESSION_ID: &str = "hook-registration-lifetime";
513 let calls = Arc::new(AtomicUsize::new(0));
514 let observed = Arc::clone(&calls);
515 let registration = register_session_end_hook(Arc::new(move |session_id| {
516 if session_id == SESSION_ID {
517 observed.fetch_add(1, Ordering::Relaxed);
518 }
519 }));
520
521 crate::reset_thread_local_state();
522 fire_session_end_hooks(SESSION_ID, false);
523 assert_eq!(calls.load(Ordering::Relaxed), 1);
524
525 drop(registration);
526 fire_session_end_hooks(SESSION_ID, false);
527 assert_eq!(
528 calls.load(Ordering::Relaxed),
529 1,
530 "dropping the owner must unregister its process-global hook"
531 );
532 }
533
534 #[test]
535 fn lifecycle_start_is_single_writer_until_terminal_update() {
536 let mut starts = ToolLifecycleStarts::default();
537
538 assert!(starts.observe(&start("session-a", "tool-1")));
539 assert!(!starts.observe(&start("session-a", "tool-1")));
540 assert!(starts.observe(&start("session-a", "tool-2")));
541 assert!(starts.observe(&start("session-b", "tool-1")));
542 assert!(starts.observe(&finish("session-a", "tool-1")));
543 assert!(starts.observe(&start("session-a", "tool-1")));
544 }
545
546 #[test]
547 fn lifecycle_start_session_cleanup_releases_unfinished_ids() {
548 let mut starts = ToolLifecycleStarts::default();
549 assert!(starts.observe(&start("session-a", "tool-1")));
550
551 starts.clear_session("session-a");
552
553 assert!(starts.observe(&start("session-a", "tool-1")));
554 }
555
556 #[tokio::test(flavor = "current_thread")]
557 async fn observation_sink_receives_one_start_across_stream_and_dispatch() {
558 const SESSION_ID: &str = "single-start-observation-test";
559 const TOOL_CALL_ID: &str = "tool-1";
560 if let Ok(mut starts) = TOOL_LIFECYCLE_STARTS.lock() {
561 starts.clear_session(SESSION_ID);
562 }
563 let sink = Arc::new(RecordingSink::default());
564 let _guard = LoopSinkGuard::install(Some(sink.clone()));
565
566 emit_agent_event_sync(&start(SESSION_ID, TOOL_CALL_ID));
567 emit_agent_event_with_ctx(None, &start(SESSION_ID, TOOL_CALL_ID)).await;
568 emit_agent_event_with_ctx(None, &finish(SESSION_ID, TOOL_CALL_ID)).await;
569
570 let events = sink.events.lock().expect("recorded events");
571 assert_eq!(
572 events
573 .iter()
574 .filter(|event| matches!(event, AgentEvent::ToolCall { .. }))
575 .count(),
576 1,
577 "the streaming and dispatch paths share one observable start authority"
578 );
579 assert_eq!(events.len(), 2, "one start and one terminal update remain");
580 }
581
582 #[tokio::test(flavor = "current_thread")]
587 async fn loop_exit_resolves_every_in_flight_call_to_a_terminal_update() {
588 const SESSION_ID: &str = "loop-exit-abandon-test";
589 const OTHER_SESSION: &str = "loop-exit-abandon-other";
590 if let Ok(mut starts) = TOOL_LIFECYCLE_STARTS.lock() {
591 starts.clear_session(SESSION_ID);
592 starts.clear_session(OTHER_SESSION);
593 }
594 let sink = Arc::new(RecordingSink::default());
595 let _guard = LoopSinkGuard::install(Some(sink.clone()));
596
597 emit_agent_event_sync(&start(SESSION_ID, "call-a"));
600 emit_agent_event_sync(&start(SESSION_ID, "call-b"));
601 emit_agent_event_sync(&start(OTHER_SESSION, "call-c"));
602
603 fire_session_end_hooks(SESSION_ID, true);
604
605 let abandoned: Vec<(String, String, String, Option<String>)> = {
606 let events = sink.events.lock().expect("recorded events");
607 events
608 .iter()
609 .filter_map(|event| match event {
610 AgentEvent::ToolCallUpdate {
611 session_id,
612 tool_call_id,
613 tool_name,
614 status: ToolCallStatus::Failed,
615 error,
616 error_category: Some(ToolCallErrorCategory::AbandonedAtLoopExit),
617 ..
618 } => Some((
619 session_id.clone(),
620 tool_call_id.clone(),
621 tool_name.clone(),
622 error.clone(),
623 )),
624 _ => None,
625 })
626 .collect()
627 };
628
629 assert_eq!(
630 abandoned.len(),
631 2,
632 "both in-flight calls for the finalizing session get one terminal update each"
633 );
634 for (session_id, _id, tool_name, error) in &abandoned {
635 assert_eq!(session_id, SESSION_ID);
636 assert_eq!(
637 tool_name, "verify",
638 "the terminal update carries the tool name from the observed start"
639 );
640 assert!(
641 error.as_deref().is_some_and(|reason| !reason.is_empty()),
642 "the abandoned call carries a human-readable reason"
643 );
644 }
645 let resolved: BTreeSet<&str> = abandoned.iter().map(|(_, id, _, _)| id.as_str()).collect();
646 assert_eq!(
647 resolved,
648 BTreeSet::from(["call-a", "call-b"]),
649 "a different session's in-flight call is not swept"
650 );
651
652 let before = sink.events.lock().expect("recorded events").len();
654 fire_session_end_hooks(SESSION_ID, true);
655 let after = sink.events.lock().expect("recorded events").len();
656 assert_eq!(before, after, "re-finalizing a drained session is a no-op");
657
658 if let Ok(mut starts) = TOOL_LIFECYCLE_STARTS.lock() {
660 starts.clear_session(OTHER_SESSION);
661 }
662 }
663
664 #[tokio::test(flavor = "current_thread")]
668 async fn loop_exit_does_not_resurrect_a_completed_call() {
669 const SESSION_ID: &str = "loop-exit-completed-test";
670 if let Ok(mut starts) = TOOL_LIFECYCLE_STARTS.lock() {
671 starts.clear_session(SESSION_ID);
672 }
673 let sink = Arc::new(RecordingSink::default());
674 let _guard = LoopSinkGuard::install(Some(sink.clone()));
675
676 emit_agent_event_sync(&start(SESSION_ID, "call-done"));
677 emit_agent_event_sync(&finish(SESSION_ID, "call-done"));
678 fire_session_end_hooks(SESSION_ID, true);
679
680 let events = sink.events.lock().expect("recorded events");
681 assert!(
682 !events.iter().any(|event| matches!(
683 event,
684 AgentEvent::ToolCallUpdate {
685 error_category: Some(ToolCallErrorCategory::AbandonedAtLoopExit),
686 ..
687 }
688 )),
689 "a call that already reached a terminal update is not swept at loop exit"
690 );
691 }
692
693 #[tokio::test(flavor = "current_thread")]
700 async fn reset_drops_in_flight_calls_so_a_later_session_cannot_inherit_them() {
701 const SESSION_ID: &str = "loop-exit-reset-test";
702 emit_agent_event_sync(&start(SESSION_ID, "call-from-a-dead-run"));
703
704 crate::reset_thread_local_state();
705
706 let sink = Arc::new(RecordingSink::default());
707 let _guard = LoopSinkGuard::install(Some(sink.clone()));
708 fire_session_end_hooks(SESSION_ID, true);
710
711 let events = sink.events.lock().expect("recorded events");
712 assert!(
713 !events.iter().any(|event| matches!(
714 event,
715 AgentEvent::ToolCallUpdate {
716 error_category: Some(ToolCallErrorCategory::AbandonedAtLoopExit),
717 ..
718 }
719 )),
720 "a reset run's in-flight calls must not surface as the next run's abandoned calls"
721 );
722 }
723
724 #[tokio::test(flavor = "current_thread")]
728 async fn suspend_releases_in_flight_calls_without_a_terminal_update() {
729 const SESSION_ID: &str = "loop-exit-suspend-test";
730 if let Ok(mut starts) = TOOL_LIFECYCLE_STARTS.lock() {
731 starts.clear_session(SESSION_ID);
732 }
733 let sink = Arc::new(RecordingSink::default());
734 let _guard = LoopSinkGuard::install(Some(sink.clone()));
735
736 emit_agent_event_sync(&start(SESSION_ID, "call-suspended"));
737 fire_session_end_hooks(SESSION_ID, false);
738
739 let events = sink.events.lock().expect("recorded events");
740 assert!(
741 !events.iter().any(|event| matches!(
742 event,
743 AgentEvent::ToolCallUpdate {
744 error_category: Some(ToolCallErrorCategory::AbandonedAtLoopExit),
745 ..
746 }
747 )),
748 "a suspended session must not emit an abandoned terminal update"
749 );
750 drop(events);
752 if let Ok(mut starts) = TOOL_LIFECYCLE_STARTS.lock() {
753 assert!(
754 starts.observe(&start(SESSION_ID, "call-suspended")),
755 "suspend released the in-flight entry"
756 );
757 starts.clear_session(SESSION_ID);
758 }
759 }
760}