diff --git a/crates/cli/src/sessions/idle.rs b/crates/cli/src/sessions/idle.rs index 70fe45fb8..7833ef124 100644 --- a/crates/cli/src/sessions/idle.rs +++ b/crates/cli/src/sessions/idle.rs @@ -93,13 +93,19 @@ async fn close_idle_turns(idle_sessions: Vec<(String, Session)>, reason: &str) - .scope(stack, async { session.close_turn_for_reason(reason).await }) .await { - Ok(subagent_ids) => { + Ok((subagent_ids, subscriber_delivery)) => { closed_turns += 1; closed_subagents.extend( subagent_ids .into_iter() .map(|subagent_id| (session_id.clone(), subagent_id)), ); + if let Some(subscriber_delivery) = subscriber_delivery + && let Err(error) = subscriber_delivery.wait().await + && first_error.is_none() + { + first_error = Some(error.into()); + } } Err(error) if first_error.is_none() => first_error = Some(error), Err(_) => {} diff --git a/crates/cli/src/sessions/mod.rs b/crates/cli/src/sessions/mod.rs index 5f5510966..98993a15c 100644 --- a/crates/cli/src/sessions/mod.rs +++ b/crates/cli/src/sessions/mod.rs @@ -10,11 +10,11 @@ use nemo_relay::api::llm::{ LlmAttributes, LlmCallEndParams, LlmCallParams, LlmHandle, LlmRequest, llm_call, llm_call_end, }; use nemo_relay::api::runtime::{ - ScopeStackHandle, TASK_SCOPE_STACK, create_scope_stack, task_scope_push, + ScopeStackHandle, SubscriberDelivery, TASK_SCOPE_STACK, create_scope_stack, task_scope_push, }; use nemo_relay::api::scope::{ EmitMarkEventParams, PopScopeParams, PushScopeParams, ScopeHandle, ScopeType, - event as emit_mark_event, get_handle, pop_scope, push_scope, + event as emit_mark_event, get_handle, pop_scope_with_subscriber_delivery, push_scope, }; use nemo_relay::api::tool::{ ToolCallEndParams, ToolCallParams, ToolHandle, tool_call, tool_call_end, @@ -340,6 +340,7 @@ impl SessionManager { headers: &HeaderMap, events: Vec, ) -> Result<(), CliError> { + let mut subscriber_deliveries = Vec::new(); let mut alignment_state = self.alignment.lock().await; let mut sessions = self.inner.lock().await; for event in events { @@ -362,7 +363,7 @@ impl SessionManager { continue; }; let event_kind = event_agent_kind(&event); - let should_remove_session = apply_event_to_session( + let (should_remove_session, subscriber_delivery) = apply_event_to_session( &mut sessions, &session_id, event, @@ -371,6 +372,9 @@ impl SessionManager { is_agent_started, ) .await?; + if let Some(subscriber_delivery) = subscriber_delivery { + subscriber_deliveries.push(subscriber_delivery); + } if is_agent_started { // A just-opened parent may unlock one or more child SessionStart hooks that arrived // earlier in this batch or an earlier request. @@ -386,6 +390,11 @@ impl SessionManager { sessions.remove(&session_id); } } + drop(sessions); + drop(alignment_state); + for subscriber_delivery in subscriber_deliveries { + subscriber_delivery.wait().await?; + } Ok(()) } @@ -776,25 +785,40 @@ impl Session { // Runs one normalized hook event inside this session's scope stack. Dispatch stays synchronous // inside the scoped closure so lifecycle ordering from each hook request is preserved exactly. - async fn apply(&mut self, event: NormalizedEvent) -> Result<(), CliError> { + async fn apply( + &mut self, + event: NormalizedEvent, + ) -> Result, CliError> { self.touch_activity(); let stack = self.scope_stack.clone(); TASK_SCOPE_STACK .scope(stack, async move { match event { - NormalizedEvent::AgentStarted(event) => self.start_agent(event), + NormalizedEvent::AgentStarted(event) => self.start_agent(event).map(|()| None), NormalizedEvent::AgentEnded(event) => self.end_agent(event).await, NormalizedEvent::TurnEnded(event) => self.end_turn(event).await, - NormalizedEvent::SubagentStarted(event) => self.start_subagent(event).await, + NormalizedEvent::SubagentStarted(event) => { + self.start_subagent(event).await.map(|()| None) + } NormalizedEvent::SubagentEnded(event) => self.end_subagent(event).await, - NormalizedEvent::LlmHint(event) => self.add_llm_hint(event), - NormalizedEvent::LlmStarted(event) => self.start_hook_llm(event).await, - NormalizedEvent::LlmEnded(event) => self.end_hook_llm(event).await, - NormalizedEvent::ToolStarted(event) => self.start_tool(event).await, + NormalizedEvent::LlmHint(event) => self.add_llm_hint(event).map(|()| None), + NormalizedEvent::LlmStarted(event) => { + self.start_hook_llm(event).await.map(|()| None) + } + NormalizedEvent::LlmEnded(event) => { + self.end_hook_llm(event).await.map(|()| None) + } + NormalizedEvent::ToolStarted(event) => { + self.start_tool(event).await.map(|()| None) + } NormalizedEvent::ToolEnded(event) => self.end_tool(event).await, NormalizedEvent::PromptSubmitted(event) => self.start_turn(event).await, - NormalizedEvent::Compaction(event) => self.mark("compaction", event), - NormalizedEvent::Notification(event) => self.mark("notification", event), + NormalizedEvent::Compaction(event) => { + self.mark("compaction", event).map(|()| None) + } + NormalizedEvent::Notification(event) => { + self.mark("notification", event).map(|()| None) + } NormalizedEvent::HookMark(event) => { let name = if event .metadata @@ -806,7 +830,7 @@ impl Session { } else { "hook_mark" }; - self.mark(name, event) + self.mark(name, event).map(|()| None) } } }) @@ -991,20 +1015,29 @@ impl Session { // Opens a new Custom turn scope for a user prompt. If the previous turn never received a // terminal hook, close it first so each user input gets a bounded reviewable trace segment. - async fn start_turn(&mut self, event: SessionEvent) -> Result<(), CliError> { + async fn start_turn( + &mut self, + event: SessionEvent, + ) -> Result, CliError> { if alignment::aliased_turn_subagent_id(&event).is_some() { self.ensure_turn_started(event.metadata.clone())?; - return self.mark("prompt_submitted", event); + self.mark("prompt_submitted", event)?; + return Ok(None); } + let mut subscriber_delivery = None; if self.turn_scope.is_some() { if self.gateway_request_turn_open { self.gateway_request_turn_open = false; - return self.mark("prompt_submitted", event); + self.mark("prompt_submitted", event)?; + return Ok(None); } - self.close_turn_for_reason("superseded_by_next_turn") + let (_, delivery) = self + .close_turn_for_reason("superseded_by_next_turn") .await?; + subscriber_delivery = delivery; } - self.open_turn(event.metadata, event.payload, "user_prompt") + self.open_turn(event.metadata, event.payload, "user_prompt")?; + Ok(subscriber_delivery) } // Lazily creates an implicit turn when gateway/tool/LLM activity arrives before a prompt hook. @@ -1121,18 +1154,23 @@ impl Session { ) } - async fn end_turn(&mut self, event: SessionEvent) -> Result<(), CliError> { + async fn end_turn( + &mut self, + event: SessionEvent, + ) -> Result, CliError> { if let Some(subagent_id) = alignment::aliased_turn_subagent_id(&event) { - self.close_subagent_scope(&subagent_id, event.payload) - .await?; - return Ok(()); + return self.close_subagent_scope(&subagent_id, event.payload).await; } - self.close_turn(event.payload, Some(event.metadata), "closed_by_turn_end") + let (_, subscriber_delivery) = self + .close_turn(event.payload, Some(event.metadata), "closed_by_turn_end") .await?; - Ok(()) + Ok(subscriber_delivery) } - async fn close_turn_for_reason(&mut self, reason: &str) -> Result, CliError> { + async fn close_turn_for_reason( + &mut self, + reason: &str, + ) -> Result<(Vec, Option), CliError> { self.close_turn(json!({ "status": reason }), None, reason) .await } @@ -1142,30 +1180,35 @@ impl Session { output: Value, boundary_metadata: Option, reason: &str, - ) -> Result, CliError> { + ) -> Result<(Vec, Option), CliError> { if self.turn_scope.is_none() { - return Ok(Vec::new()); + return Ok((Vec::new(), None)); } self.close_active_llms(reason).await?; self.close_active_tools(reason).await?; let closed_subagents = self.close_active_subagents(reason).await?; let output = self.last_turn_llm_output.take().unwrap_or(output); self.clear_correlation_state(); - self.close_turn_scope(output, boundary_metadata)?; - Ok(closed_subagents) + let subscriber_delivery = self.close_turn_scope(output, boundary_metadata)?; + Ok((closed_subagents, subscriber_delivery)) } // Closes the session in a fail-safe order: active turn first, then the root agent scope when // the harness has one. Duplicate terminal hooks must not reopen scopes. - async fn end_agent(&mut self, event: SessionEvent) -> Result<(), CliError> { + async fn end_agent( + &mut self, + event: SessionEvent, + ) -> Result, CliError> { if !self.session_started && self.agent_scope.is_none() && self.turn_scope.is_none() { - return Ok(()); + return Ok(None); } - self.close_turn_for_reason("closed_by_agent_end").await?; + let (_, turn_delivery) = self.close_turn_for_reason("closed_by_agent_end").await?; self.clear_correlation_state(); - self.close_agent_scope(event.payload)?; + let agent_delivery = self.close_agent_scope(event.payload)?; self.session_started = false; - Ok(()) + // Agent end is queued after turn end on the serial dispatcher. Waiting for the later + // receipt therefore covers both terminal events without a process-wide flush. + Ok(agent_delivery.or(turn_delivery)) } async fn close_for_shutdown(&mut self, reason: &str) -> Result<(), CliError> { @@ -1176,9 +1219,9 @@ impl Session { if self.agent_scope.is_none() && self.turn_scope.is_none() { return Ok(()); } - self.close_turn_for_reason(reason).await?; + let _ = self.close_turn_for_reason(reason).await?; self.clear_correlation_state(); - self.close_agent_scope(payload)?; + let _ = self.close_agent_scope(payload)?; self.session_started = false; Ok(()) }) @@ -1226,7 +1269,10 @@ impl Session { async fn close_active_subagents(&mut self, reason: &str) -> Result, CliError> { let mut closed = Vec::new(); while let Some(subagent_id) = self.subagent_stack.pop() { - self.close_subagent_scope(&subagent_id, json!({ "status": reason })) + // Subagent ends precede the turn end on the serial dispatcher, so the turn receipt + // returned by `close_turn` also covers these deliveries. + let _ = self + .close_subagent_scope(&subagent_id, json!({ "status": reason })) .await?; closed.push(subagent_id); } @@ -1245,36 +1291,39 @@ impl Session { // Ends the root agent scope when present. Duplicate agent-end hooks can reach this path after the // scope is already gone, so absence is treated as a no-op. - fn close_agent_scope(&mut self, payload: Value) -> Result<(), CliError> { + fn close_agent_scope( + &mut self, + payload: Value, + ) -> Result, CliError> { let Some(scope) = self.agent_scope.take() else { - return Ok(()); + return Ok(None); }; - pop_scope( + let subscriber_delivery = pop_scope_with_subscriber_delivery( PopScopeParams::builder() .handle_uuid(&scope.uuid) .output(payload) .build(), )?; - Ok(()) + Ok(Some(subscriber_delivery)) } fn close_turn_scope( &mut self, output: Value, boundary_metadata: Option, - ) -> Result<(), CliError> { + ) -> Result, CliError> { let Some(scope) = self.turn_scope.take() else { - return Ok(()); + return Ok(None); }; self.gateway_request_turn_open = false; - pop_scope( + let subscriber_delivery = pop_scope_with_subscriber_delivery( PopScopeParams::builder() .handle_uuid(&scope.uuid) .output(output) .metadata_opt(boundary_metadata) .build(), )?; - Ok(()) + Ok(Some(subscriber_delivery)) } fn root_work_scope(&self) -> Option { @@ -1339,9 +1388,12 @@ impl Session { // a subagent already closed by another provider-specific completion signal are ignored. Claude // Code can also report late orphan stops after a turn has closed; those are logged and ignored // when there is no active turn so they cannot create lifecycle-only traces. - async fn end_subagent(&mut self, event: SubagentEvent) -> Result<(), CliError> { + async fn end_subagent( + &mut self, + event: SubagentEvent, + ) -> Result, CliError> { if self.completed_subagents.contains(&event.subagent_id) { - return Ok(()); + return Ok(None); } if !self.subagents.contains_key(&event.subagent_id) { log::warn!( @@ -1353,9 +1405,9 @@ impl Session { "Subagent lifecycle event had no matching start" ); if self.agent_kind == AgentKind::ClaudeCode && self.turn_scope.is_none() { - return Ok(()); + return Ok(None); } - return self.mark( + self.mark( "subagent_end_without_start", SessionEvent { session_id: event.session_id, @@ -1364,12 +1416,12 @@ impl Session { payload: event.payload, metadata: event.metadata, }, - ); + )?; + return Ok(None); }; self.ensure_turn_started(event.metadata.clone())?; self.close_subagent_scope(&event.subagent_id, event.payload) - .await?; - Ok(()) + .await } // Closes one subagent using that subagent's own scope stack. This is shared by explicit end @@ -1379,17 +1431,17 @@ impl Session { &mut self, subagent_id: &str, output: Value, - ) -> Result { + ) -> Result, CliError> { let Some(scope) = self.subagents.remove(subagent_id) else { - return Ok(false); + return Ok(None); }; let stack = self .subagent_stacks .remove(subagent_id) .unwrap_or_else(|| self.scope_stack.clone()); - TASK_SCOPE_STACK + let subscriber_delivery = TASK_SCOPE_STACK .scope(stack, async { - pop_scope( + pop_scope_with_subscriber_delivery( PopScopeParams::builder() .handle_uuid(&scope.uuid) .output(output) @@ -1411,7 +1463,7 @@ impl Session { { self.last_llm_owner = None; } - Ok(true) + Ok(Some(subscriber_delivery)) } // Stores an LLM correlation hint from hook activity after pruning expired hints. Hints do not @@ -1565,7 +1617,7 @@ impl Session { // Ends a tool call, synthesizing a start if no matching handle exists. This keeps post-only // hooks observable and preserves the final result/status instead of dropping orphaned endings. - async fn end_tool(&mut self, event: ToolEvent) -> Result<(), CliError> { + async fn end_tool(&mut self, event: ToolEvent) -> Result, CliError> { self.ensure_turn_started(event.metadata.clone())?; let event_metadata = self.event_identity_metadata(event.metadata.clone()); let completed_agent_subagent_id = alignment::completed_subagent_from_tool(&event); @@ -1616,11 +1668,10 @@ impl Session { .build(), )?; self.set_last_tool_owner(explicit_subagent_id); - if let Some(subagent_id) = completed_agent_subagent_id { - self.close_subagent_scope(&subagent_id, event.result) - .await?; + match completed_agent_subagent_id { + Some(subagent_id) => self.close_subagent_scope(&subagent_id, event.result).await, + None => Ok(None), } - Ok(()) } // Hermes pre/post tool hooks can disagree on call IDs: pre hooks may omit the provider id diff --git a/crates/cli/src/sessions/routing.rs b/crates/cli/src/sessions/routing.rs index 44de606e0..d8bfd10ce 100644 --- a/crates/cli/src/sessions/routing.rs +++ b/crates/cli/src/sessions/routing.rs @@ -5,6 +5,7 @@ use std::collections::HashMap; +use nemo_relay::api::runtime::SubscriberDelivery; use serde_json::Value; use crate::agents::shared::alignment::{ @@ -55,7 +56,7 @@ pub(super) async fn apply_event_to_session( event_kind: AgentKind, config: SessionConfig, is_agent_started: bool, -) -> Result { +) -> Result<(bool, Option), CliError> { let session = sessions .entry(session_id.to_string()) .or_insert_with(|| Session::new(session_id.to_string(), event_kind, config)); @@ -65,8 +66,8 @@ pub(super) async fn apply_event_to_session( { session.agent_kind = event_kind; } - session.apply(event).await?; - Ok(session.is_empty()) + let subscriber_delivery = session.apply(event).await?; + Ok((session.is_empty(), subscriber_delivery)) } pub(super) async fn promote_pending_subagents_for_parent( @@ -109,7 +110,7 @@ pub(super) async fn promote_pending_subagent( Session::new(parent_session_id.clone(), pending.event.agent_kind, config) }); if !parent_session.session_started && parent_session.agent_scope.is_none() { - parent_session + let _ = parent_session .apply(NormalizedEvent::AgentStarted(SessionEvent { session_id: parent_session_id, agent_kind: pending.event.agent_kind, @@ -119,7 +120,7 @@ pub(super) async fn promote_pending_subagent( })) .await?; } - parent_session + let _ = parent_session .apply(NormalizedEvent::SubagentStarted( pending.subagent_start_event(), )) diff --git a/crates/cli/tests/coverage/shared/server_tests.rs b/crates/cli/tests/coverage/shared/server_tests.rs index e25c048c7..084193e9c 100644 --- a/crates/cli/tests/coverage/shared/server_tests.rs +++ b/crates/cli/tests/coverage/shared/server_tests.rs @@ -18,9 +18,10 @@ use nemo_relay::api::event::ScopeCategory; use nemo_relay::api::llm::LlmRequestInterceptOutcome; use nemo_relay::api::registry::{ deregister_llm_execution_intercept, deregister_llm_request_intercept, - deregister_llm_stream_execution_intercept, deregister_tool_conditional_execution_guardrail, - register_llm_execution_intercept, register_llm_request_intercept, - register_llm_stream_execution_intercept, register_tool_conditional_execution_guardrail, + deregister_llm_stream_execution_intercept, deregister_scope_sanitize_end_guardrail, + deregister_tool_conditional_execution_guardrail, register_llm_execution_intercept, + register_llm_request_intercept, register_llm_stream_execution_intercept, + register_scope_sanitize_end_guardrail, register_tool_conditional_execution_guardrail, }; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; use nemo_relay::plugin::dynamic::DynamicPluginKind; @@ -126,6 +127,14 @@ impl Drop for SubscriberCleanup { } } +struct ScopeEndSanitizerCleanup(&'static str); + +impl Drop for ScopeEndSanitizerCleanup { + fn drop(&mut self) { + let _ = deregister_scope_sanitize_end_guardrail(self.0); + } +} + struct RequestInterceptCleanup(&'static str); impl Drop for RequestInterceptCleanup { @@ -1020,6 +1029,19 @@ async fn serve_listener_activates_plugin_config_and_clears_on_shutdown() { .unwrap(); assert_eq!(response.status(), StatusCode::OK); } + for hook_event_name in ["sessionStart", "UserPromptSubmit"] { + let response = client + .post(format!("{url}/hooks/codex")) + .json(&json!({ + "session_id": "plugin-shutdown-open-session", + "hook_event_name": hook_event_name, + "prompt": "Leave this turn open until Relay shuts down." + })) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + } shutdown_tx.send(()).unwrap(); handle.await.unwrap().unwrap(); @@ -1051,6 +1073,180 @@ async fn serve_listener_activates_plugin_config_and_clears_on_shutdown() { .as_array() .is_some_and(|events| events.len() >= 2) ); + assert!( + trajectories.iter().any(|trajectory| { + atif_matches_session(trajectory, "plugin-shutdown-open-session") + && trajectory["extra"]["observed_events"] + .as_array() + .is_some_and(|events| { + events.iter().any(|event| { + event["name"] == json!("codex-turn") + && event["scope_category"] == json!("end") + }) + }) + }), + "full server teardown must flush an open session's terminal ATIF snapshot before clearing plugins: {}", + serde_json::to_string_pretty(&trajectories).unwrap() + ); +} + +#[tokio::test] +async fn terminal_hook_responses_wait_for_their_atif_snapshot() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let _ = nemo_relay::plugin::clear_plugin_configuration(); + + let temp = tempfile::tempdir().unwrap(); + let atif_dir = temp.path().join("atif"); + std::fs::create_dir_all(&atif_dir).unwrap(); + let mut config = test_config(); + config.plugin_config = Some(json!({ + "version": 1, + "components": [{ + "kind": "observability", + "enabled": true, + "config": { + "version": 3, + "atif": { + "enabled": true, + "output_directory": atif_dir, + "filename_template": "trajectory-{session_id}.json" + } + } + }] + })); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let url = format!("http://{address}"); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let handle = + tokio::spawn(async move { serve_listener(listener, config, Some(shutdown_rx)).await }); + wait_for_gateway(&url).await; + let client = test_http_client(); + + for (path, session_id, turn_name, terminal_event, sanitizer_name) in [ + ( + "/hooks/codex", + "codex-atif-response-boundary", + "codex-turn", + "Stop", + "codex-atif-response-boundary-sanitizer", + ), + ( + "/hooks/claude-code", + "claude-atif-response-boundary", + "claude-code-turn", + "SessionEnd", + "claude-atif-response-boundary-sanitizer", + ), + ] { + for hook_event_name in ["sessionStart", "UserPromptSubmit"] { + let response = client + .post(format!("{url}{path}")) + .json(&json!({ + "session_id": session_id, + "hook_event_name": hook_event_name, + "prompt": "Return one short answer." + })) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + } + + let _ = deregister_scope_sanitize_end_guardrail(sanitizer_name); + let (started_tx, started_rx) = oneshot::channel(); + let (release_tx, release_rx) = oneshot::channel(); + let started_tx = Arc::new(Mutex::new(Some(started_tx))); + let release_rx = Arc::new(Mutex::new(Some(release_rx))); + let expected_session_id = session_id.to_string(); + let expected_turn_name = turn_name.to_string(); + register_scope_sanitize_end_guardrail( + sanitizer_name, + 0, + Arc::new(move |event, fields| { + let should_block = event.scope_category() == Some(ScopeCategory::End) + && event.name() == expected_turn_name + && event + .metadata() + .and_then(|metadata| metadata.get("session_id")) + .and_then(Value::as_str) + == Some(expected_session_id.as_str()); + let started = should_block + .then(|| started_tx.lock().unwrap().take()) + .flatten(); + let release = should_block + .then(|| release_rx.lock().unwrap().take()) + .flatten(); + Box::pin(async move { + if let Some(started) = started { + let _ = started.send(()); + if let Some(release) = release { + let _ = release.await; + } + } + Ok(fields) + }) + }), + ) + .unwrap(); + let sanitizer_cleanup = ScopeEndSanitizerCleanup(sanitizer_name); + + let terminal_client = client.clone(); + let terminal_url = format!("{url}{path}"); + let terminal_session_id = session_id.to_string(); + let mut terminal = tokio::spawn(async move { + terminal_client + .post(terminal_url) + .json(&json!({ + "session_id": terminal_session_id, + "hook_event_name": terminal_event, + "response": "Done." + })) + .send() + .await + .unwrap() + }); + tokio::time::timeout(std::time::Duration::from_secs(10), started_rx) + .await + .expect("terminal scope sanitizer should start") + .unwrap(); + + let early_response = + tokio::time::timeout(std::time::Duration::from_millis(100), &mut terminal) + .await + .ok(); + let returned_early = early_response.is_some(); + let _ = release_tx.send(()); + let response = match early_response { + Some(response) => response.unwrap(), + None => terminal.await.unwrap(), + }; + assert_eq!(response.status(), StatusCode::OK); + assert!( + !returned_early, + "{terminal_event} returned before its terminal subscribers completed" + ); + + let trajectories = std::fs::read_dir(temp.path().join("atif")) + .unwrap() + .filter_map(Result::ok) + .filter_map(|entry| { + serde_json::from_slice::(&std::fs::read(entry.path()).ok()?).ok() + }) + .collect::>(); + assert!( + trajectories + .iter() + .any(|trajectory| atif_matches_session(trajectory, session_id)), + "terminal hook response must not precede the ATIF snapshot for {session_id}: {}", + serde_json::to_string_pretty(&trajectories).unwrap() + ); + drop(sanitizer_cleanup); + } + + shutdown_tx.send(()).unwrap(); + handle.await.unwrap().unwrap(); } fn atif_matches_session(trajectory: &Value, session_id: &str) -> bool { diff --git a/crates/cli/tests/coverage/shared/session_tests.rs b/crates/cli/tests/coverage/shared/session_tests.rs index 4223e2da8..9b986002d 100644 --- a/crates/cli/tests/coverage/shared/session_tests.rs +++ b/crates/cli/tests/coverage/shared/session_tests.rs @@ -1329,6 +1329,97 @@ async fn turn_end_metadata_comes_only_from_the_real_turn_boundary() { deregister_subscriber(subscriber_name).unwrap(); } +#[tokio::test] +async fn terminal_subscriber_wait_releases_session_manager_locks() { + let subscriber_name = "cli-terminal-delivery-lock-release-test"; + let _ = deregister_subscriber(subscriber_name); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let started_tx = Arc::new(StdMutex::new(Some(started_tx))); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + let release_rx = Arc::new(StdMutex::new(release_rx)); + register_subscriber( + subscriber_name, + Arc::new(move |event| { + if event.scope_category() == Some(ScopeCategory::End) + && event.name() == "codex-turn" + && event + .metadata() + .and_then(|metadata| metadata.get("session_id")) + .and_then(Value::as_str) + == Some("blocked-terminal-session") + { + if let Some(started) = started_tx.lock().unwrap().take() { + let _ = started.send(()); + } + let _ = release_rx.lock().unwrap().recv(); + } + }), + ) + .unwrap(); + + let manager = SessionManager::new(session_test_config()); + for session_id in ["blocked-terminal-session", "parallel-session"] { + manager + .apply_events( + &HeaderMap::new(), + vec![ + NormalizedEvent::AgentStarted(codex_session_event( + session_id, + "SessionStart", + json!({ "session_id": session_id }), + )), + NormalizedEvent::PromptSubmitted(codex_session_event( + session_id, + "UserPromptSubmit", + json!({ "session_id": session_id }), + )), + ], + ) + .await + .unwrap(); + } + flush_subscribers().unwrap(); + + let terminal_manager = manager.clone(); + let terminal = tokio::spawn(async move { + terminal_manager + .apply_events( + &HeaderMap::new(), + vec![NormalizedEvent::TurnEnded(codex_session_event( + "blocked-terminal-session", + "Stop", + json!({ "session_id": "blocked-terminal-session" }), + ))], + ) + .await + }); + tokio::time::timeout(std::time::Duration::from_secs(1), started_rx) + .await + .expect("terminal subscriber should start") + .unwrap(); + + let parallel_result = tokio::time::timeout( + std::time::Duration::from_millis(250), + manager.apply_events( + &HeaderMap::new(), + vec![NormalizedEvent::Notification(codex_session_event( + "parallel-session", + "notification", + json!({ "session_id": "parallel-session" }), + ))], + ), + ) + .await; + release_tx.send(()).unwrap(); + terminal.await.unwrap().unwrap(); + flush_subscribers().unwrap(); + deregister_subscriber(subscriber_name).unwrap(); + + parallel_result + .expect("another session must remain writable while terminal subscribers are active") + .unwrap(); +} + #[tokio::test] async fn new_subagent_claims_first_unhinted_llm_when_siblings_active() { let manager = SessionManager::new(session_test_config()); diff --git a/crates/core/src/api/runtime.rs b/crates/core/src/api/runtime.rs index 12d6612c3..4b26c6d9e 100644 --- a/crates/core/src/api/runtime.rs +++ b/crates/core/src/api/runtime.rs @@ -32,4 +32,6 @@ pub use scope_stack::{ task_scope_top, with_active_event_uuid, with_scope_stack, }; pub use state::NemoRelayContextState; +#[doc(hidden)] +pub use subscriber_dispatcher::SubscriberDelivery; pub use subscriber_dispatcher::flush_subscribers; diff --git a/crates/core/src/api/runtime/subscriber_dispatcher.rs b/crates/core/src/api/runtime/subscriber_dispatcher.rs index e4255ffa7..bd08ca038 100644 --- a/crates/core/src/api/runtime/subscriber_dispatcher.rs +++ b/crates/core/src/api/runtime/subscriber_dispatcher.rs @@ -8,7 +8,7 @@ use crate::api::registry::Guardrail; use crate::api::runtime::{ EventSanitizeFn, EventSubscriberFn, NemoRelayContextState, ScopeStackHandle, }; -use crate::error::Result; +use crate::error::{FlowError, Result}; use std::any::Any; use std::cell::RefCell; use std::future::Future; @@ -82,6 +82,37 @@ pub(crate) type EventTransformFn = Box< dyn FnOnce(Event) -> Pin + Send + 'static>> + Send + 'static, >; +/// Completion receipt for one queued subscriber delivery. +/// +/// Unlike [`flush_subscribers`], this receipt waits only for sanitizer and +/// subscriber processing of the event that created it. Events queued later are +/// not part of the wait. +#[doc(hidden)] +pub struct SubscriberDelivery { + completion: tokio::sync::oneshot::Receiver<()>, +} + +impl SubscriberDelivery { + fn completed() -> Self { + let (completion_tx, completion) = tokio::sync::oneshot::channel(); + let _ = completion_tx.send(()); + Self { completion } + } + + /// Wait until this event's subscriber delivery is complete. + /// + /// Do not call this from a subscriber, event-sanitizer, guardrail, or + /// intercept callback. The dispatcher signals completion on its own + /// thread, so waiting there creates a wait cycle. + pub async fn wait(self) -> Result<()> { + self.completion.await.map_err(|error| { + FlowError::Internal(format!( + "subscriber delivery completion channel closed: {error}" + )) + }) + } +} + mod native { use std::cell::{Cell, RefCell}; use std::panic::{AssertUnwindSafe, catch_unwind}; @@ -108,6 +139,7 @@ mod native { scope_stack: ScopeStackHandle, publication_context: Option, lineage: Option, + completion: Option>, }, Flush { done: Sender<()>, @@ -486,6 +518,7 @@ mod native { scope_stack, publication_context: current_publication_context(), lineage: None, + completion: None, }; send_dispatch_message(message) } @@ -510,10 +543,44 @@ mod native { scope_stack, publication_context: current_publication_context(), lineage: None, + completion: None, }; enqueue_dispatch_message(message) } + pub(super) fn dispatch_sanitized_event_with_delivery( + event: Event, + sanitizers: Vec>, + subscribers: &[EventSubscriberFn], + scope_stack: ScopeStackHandle, + ) -> Result { + if subscribers.is_empty() { + return Ok(SubscriberDelivery::completed()); + } + let Some(scope_stack) = immutable_scope_stack(&scope_stack) else { + return Err(FlowError::Internal( + "failed to snapshot scope stack for subscriber delivery".into(), + )); + }; + let (completion_tx, completion) = tokio::sync::oneshot::channel(); + let message = DispatcherMessage::Deliver { + event: Box::new(event), + transform: None, + sanitizers, + subscribers: subscribers.to_vec(), + scope_stack, + publication_context: current_publication_context(), + lineage: None, + completion: Some(completion_tx), + }; + if !enqueue_dispatch_message(message) { + return Err(FlowError::Internal( + "failed to queue tracked subscriber delivery".into(), + )); + } + Ok(SubscriberDelivery { completion }) + } + pub(super) fn dispatch_reserved_sanitized_event( event: Event, sanitizers: Vec>, @@ -534,6 +601,7 @@ mod native { scope_stack, publication_context: current_publication_context(), lineage: None, + completion: None, }; enqueue_dispatch_message(message) } @@ -556,6 +624,7 @@ mod native { scope_stack, publication_context: current_publication_context(), lineage: None, + completion: None, }; enqueue_dispatch_message(message) } @@ -884,6 +953,7 @@ mod native { scope_stack, publication_context, lineage: permit, + completion, } => { let nested_publications = with_publication_lineage(Arc::clone(&lineage), || { deliver_event( @@ -899,6 +969,9 @@ mod native { for publication in nested_publications { handle_message(publication, state, Some(&lineage)); } + if let Some(completion) = completion { + let _ = completion.send(()); + } } DispatcherMessage::Barrier { publications, @@ -1182,6 +1255,15 @@ pub(crate) fn dispatch_sanitized_event( native::dispatch_sanitized_event(event, sanitizers, subscribers, scope_stack) } +pub(crate) fn dispatch_sanitized_event_with_delivery( + event: Event, + sanitizers: Vec>, + subscribers: &[EventSubscriberFn], + scope_stack: ScopeStackHandle, +) -> Result { + native::dispatch_sanitized_event_with_delivery(event, sanitizers, subscribers, scope_stack) +} + /// Publish a stream-finalization event at its reserved FIFO position. pub(crate) fn dispatch_reserved_sanitized_event( event: Event, diff --git a/crates/core/src/api/scope.rs b/crates/core/src/api/scope.rs index 03ecdf261..1bbe04d90 100644 --- a/crates/core/src/api/scope.rs +++ b/crates/core/src/api/scope.rs @@ -4,7 +4,7 @@ use crate::api::event::{BaseEvent, CategoryProfile, DataSchema, EventCategory, MarkEvent}; use crate::api::runtime::global_context; use crate::api::runtime::scope_stack::snapshot_scope_stack; -use crate::api::runtime::subscriber_dispatcher; +use crate::api::runtime::subscriber_dispatcher::{self, SubscriberDelivery}; use crate::api::runtime::{ current_scope_stack, task_scope_push, task_scope_remove, task_scope_top, }; @@ -296,6 +296,26 @@ pub fn push_scope(params: PushScopeParams<'_>) -> Result { /// snapshot, so cleanup does not change the middleware applied to the emitted /// event. pub fn pop_scope(params: PopScopeParams<'_>) -> Result<()> { + pop_scope_inner(params, false).map(|_| ()) +} + +/// Pop the current scope and return a receipt for its scope-end subscriber delivery. +/// +/// The receipt covers sanitizer and subscriber processing for the scope-end event. +/// It does not wait for unrelated events queued after that event. +#[doc(hidden)] +pub fn pop_scope_with_subscriber_delivery( + params: PopScopeParams<'_>, +) -> Result { + pop_scope_inner(params, true)?.ok_or_else(|| { + FlowError::Internal("tracked scope pop did not create a subscriber delivery receipt".into()) + }) +} + +fn pop_scope_inner( + params: PopScopeParams<'_>, + track_delivery: bool, +) -> Result> { ensure_runtime_owner()?; let scope_stack = current_scope_stack(); let (scope, event, subscribers, emission_scope_stack) = { @@ -335,13 +355,23 @@ pub fn pop_scope(params: PopScopeParams<'_>) -> Result<()> { let publication_scope_stack = snapshot_scope_stack(&emission_scope_stack)?; let removed = task_scope_remove(params.handle_uuid)?; debug_assert_eq!(removed.uuid, scope.uuid); - let _ = subscriber_dispatcher::dispatch_sanitized_event( - event, - sanitizers, - &subscribers, - publication_scope_stack, - ); - Ok(()) + if track_delivery { + subscriber_dispatcher::dispatch_sanitized_event_with_delivery( + event, + sanitizers, + &subscribers, + publication_scope_stack, + ) + .map(Some) + } else { + let _ = subscriber_dispatcher::dispatch_sanitized_event( + event, + sanitizers, + &subscribers, + publication_scope_stack, + ); + Ok(None) + } } /// Emit a standalone mark event under the current or provided scope. diff --git a/crates/core/tests/unit/subscriber_dispatcher_tests.rs b/crates/core/tests/unit/subscriber_dispatcher_tests.rs index 00bb3c9aa..d6ee7a45c 100644 --- a/crates/core/tests/unit/subscriber_dispatcher_tests.rs +++ b/crates/core/tests/unit/subscriber_dispatcher_tests.rs @@ -2,10 +2,10 @@ // SPDX-License-Identifier: Apache-2.0 use super::native::{ DispatcherLoopState, DispatcherMessage, PendingFlush, PublicationLineage, PublicationPermit, - dispatcher_sender, enqueue_dispatch_message, flush_queued_subscribers, flush_subscribers, - prepare_for_fork, register_async_publication, register_pending_publication, - resume_after_fork_parent, sanitize_event_snapshot, set_sanitizer_runtime_failure_for_test, - spawn_background_publication, + dispatch_sanitized_event_with_delivery, dispatcher_sender, enqueue_dispatch_message, + flush_queued_subscribers, flush_subscribers, prepare_for_fork, register_async_publication, + register_pending_publication, resume_after_fork_parent, sanitize_event_snapshot, + set_sanitizer_runtime_failure_for_test, spawn_background_publication, }; use super::{EventSubscriberFn, publication_context}; use crate::api::registry::RegistryRecord; @@ -59,6 +59,7 @@ fn flush_waits_for_active_but_not_later_publication_barriers() { scope_stack: current_scope_stack(), publication_context: None, lineage: None, + completion: None, }) .unwrap(); let (flush_tx, flush_rx) = mpsc::channel(); @@ -94,6 +95,7 @@ fn flush_waits_for_active_but_not_later_publication_barriers() { scope_stack: current_scope_stack(), publication_context: None, lineage: None, + completion: None, }]) .unwrap(); flush_rx @@ -135,6 +137,7 @@ fn pending_publication_defers_flush_without_blocking_unrelated_delivery() { scope_stack: current_scope_stack(), publication_context: None, lineage: None, + completion: None, }); delivered_rx .recv_timeout(std::time::Duration::from_secs(1)) @@ -202,6 +205,7 @@ fn flush_does_not_wait_for_later_delivery() { scope_stack: current_scope_stack(), publication_context: None, lineage: None, + completion: None, }) .unwrap(); barrier.sender.send(Vec::new()).unwrap(); @@ -215,6 +219,128 @@ fn flush_does_not_wait_for_later_delivery() { ); } +#[test] +fn subscriber_delivery_receipt_waits_for_its_event() { + let _lock = crate::shared_runtime::runtime_owner_test_mutex() + .lock() + .unwrap_or_else(|error| error.into_inner()); + flush_subscribers().unwrap(); + let (started_tx, started_rx) = mpsc::channel(); + let (release_tx, release_rx) = mpsc::channel(); + let release_rx = Arc::new(Mutex::new(release_rx)); + let subscriber: EventSubscriberFn = Arc::new(move |_event| { + started_tx.send(()).unwrap(); + release_rx + .lock() + .unwrap_or_else(|error| error.into_inner()) + .recv() + .unwrap(); + }); + let event = serde_json::from_value(serde_json::json!({ + "kind": "mark", + "atof_version": "0.1", + "uuid": "019c1df6-4a57-7000-8000-000000000017", + "timestamp": "2026-07-28T00:00:00Z", + "name": "tracked-delivery" + })) + .expect("valid event"); + let delivery = dispatch_sanitized_event_with_delivery( + event, + Vec::new(), + &[subscriber], + current_scope_stack(), + ) + .unwrap(); + started_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .expect("tracked subscriber should start"); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap(); + let mut wait = Box::pin(delivery.wait()); + assert!( + runtime + .block_on(async { + tokio::time::timeout(std::time::Duration::from_millis(50), wait.as_mut()).await + }) + .is_err(), + "delivery receipt must remain pending while its subscriber is active" + ); + release_tx.send(()).unwrap(); + runtime + .block_on(async { + tokio::time::timeout(std::time::Duration::from_secs(1), wait.as_mut()).await + }) + .expect("delivery receipt should complete after subscriber delivery") + .unwrap(); + flush_subscribers().unwrap(); +} + +#[test] +fn subscriber_delivery_receipt_does_not_capture_later_events() { + let _lock = crate::shared_runtime::runtime_owner_test_mutex() + .lock() + .unwrap_or_else(|error| error.into_inner()); + flush_subscribers().unwrap(); + let event = |uuid: &str, name: &str| { + serde_json::from_value(serde_json::json!({ + "kind": "mark", + "atof_version": "0.1", + "uuid": uuid, + "timestamp": "2026-07-28T00:00:00Z", + "name": name + })) + .expect("valid event") + }; + let delivery = dispatch_sanitized_event_with_delivery( + event("019c1df6-4a57-7000-8000-000000000018", "tracked"), + Vec::new(), + &[Arc::new(|_event| {})], + current_scope_stack(), + ) + .unwrap(); + + let (later_started_tx, later_started_rx) = mpsc::channel(); + let (release_later_tx, release_later_rx) = mpsc::channel(); + enqueue_dispatch_message(DispatcherMessage::Deliver { + event: Box::new(event( + "019c1df6-4a57-7000-8000-000000000019", + "later-blocked", + )), + transform: Some(Box::new(move |event| { + Box::pin(async move { + later_started_tx.send(()).unwrap(); + release_later_rx.recv().unwrap(); + event + }) + })), + sanitizers: Vec::new(), + subscribers: Vec::new(), + scope_stack: current_scope_stack(), + publication_context: None, + lineage: None, + completion: None, + }); + later_started_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .expect("later delivery should block the dispatcher"); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap(); + let result = runtime.block_on(async { + tokio::time::timeout(std::time::Duration::from_millis(100), delivery.wait()).await + }); + release_later_tx.send(()).unwrap(); + flush_subscribers().unwrap(); + result + .expect("tracked delivery must not wait for a later queued event") + .unwrap(); +} + #[test] fn pending_flushes_do_not_acknowledge_out_of_order() { let first_lineage = Arc::new(PublicationLineage::default()); @@ -311,6 +437,7 @@ fn nested_publication_barrier_precedes_already_queued_delivery() { scope_stack: nested_scope_stack.clone(), publication_context: None, lineage: None, + completion: None, })); let publication = register_async_publication().expect("nested publication barrier"); @@ -333,6 +460,7 @@ fn nested_publication_barrier_precedes_already_queued_delivery() { scope_stack: nested_scope_stack, publication_context: None, lineage: None, + completion: None, }]) .unwrap(); event @@ -343,6 +471,7 @@ fn nested_publication_barrier_precedes_already_queued_delivery() { scope_stack: current_scope_stack(), publication_context: None, lineage: None, + completion: None, }) .unwrap(); started_rx @@ -357,6 +486,7 @@ fn nested_publication_barrier_precedes_already_queued_delivery() { scope_stack: current_scope_stack(), publication_context: None, lineage: None, + completion: None, }) .unwrap(); release_tx.send(()).unwrap(); @@ -423,6 +553,7 @@ fn flush_waits_for_transitive_subscriber_publications_without_reordering() { scope_stack: current_scope_stack(), publication_context: None, lineage: None, + completion: None, })); }) }; @@ -460,6 +591,7 @@ fn flush_waits_for_transitive_subscriber_publications_without_reordering() { scope_stack: current_scope_stack(), publication_context: None, lineage: None, + completion: None, })); }) }; @@ -482,6 +614,7 @@ fn flush_waits_for_transitive_subscriber_publications_without_reordering() { scope_stack: current_scope_stack(), publication_context: None, lineage: None, + completion: None, }) .unwrap(); outer_started_rx @@ -496,6 +629,7 @@ fn flush_waits_for_transitive_subscriber_publications_without_reordering() { scope_stack: current_scope_stack(), publication_context: None, lineage: None, + completion: None, }) .unwrap(); release_outer_tx diff --git a/docs/nemo-relay-cli/basic-usage.mdx b/docs/nemo-relay-cli/basic-usage.mdx index 55d9b4e8a..a615a527f 100644 --- a/docs/nemo-relay-cli/basic-usage.mdx +++ b/docs/nemo-relay-cli/basic-usage.mdx @@ -29,6 +29,13 @@ the payload in a shared gateway envelope. - `POST /hooks/hermes` accepts Hermes shell hook JSON and returns the empty JSON object expected by Hermes hook commands. +When a hook closes a turn, subagent, or session scope, Relay returns the hook +response after subscribers finish processing that scope-end event. This makes +synchronous outputs such as the final ATIF trajectory visible at the hook +boundary without draining subscriber work queued later by another session. +Full gateway shutdown uses a process-wide flush after closing all remaining +sessions and before clearing plugins. + The adapters preserve vendor fields such as session IDs, working directories, transcript paths, model names, tool payloads, shell payloads, MCP payloads, file payloads, user identity, and subagent metadata in NeMo Relay event metadata.