From ff7db8b4d0a325e1bcc65cacb2ad1b722af89a61 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Fri, 11 Sep 2026 17:43:39 -0400 Subject: [PATCH 1/4] enhancement: improve daemon worker and gateway concurrency Signed-off-by: Will Killian --- crates/cli/src/daemon/broker/registry.rs | 48 ++ crates/cli/src/daemon/broker/server/socket.rs | 13 +- crates/cli/src/daemon/worker/managed.rs | 77 ++- crates/cli/src/daemon/worker/runtime.rs | 26 +- crates/cli/src/gateway/mod.rs | 32 +- crates/cli/src/gateway/request.rs | 17 + crates/cli/src/sessions/idle.rs | 10 +- crates/cli/src/sessions/mod.rs | 477 ++++++++++++++---- crates/cli/src/sessions/routing.rs | 204 +++++--- .../tests/coverage/daemon/registry_tests.rs | 25 + .../tests/coverage/shared/session_tests.rs | 84 +++ crates/core/src/codec/streaming.rs | 39 +- crates/core/src/observability/atof.rs | 3 + .../core/tests/unit/codec/streaming_tests.rs | 15 + 14 files changed, 858 insertions(+), 212 deletions(-) diff --git a/crates/cli/src/daemon/broker/registry.rs b/crates/cli/src/daemon/broker/registry.rs index ca6913d54..5fbe92484 100644 --- a/crates/cli/src/daemon/broker/registry.rs +++ b/crates/cli/src/daemon/broker/registry.rs @@ -226,6 +226,22 @@ impl Registry { fingerprint: Fingerprint, session_id: &McpSessionId, ) -> Result { + // Most directive polls observe stable ready/pass-through state. Keep + // those polls on the shared read lock; recovering-with-target remains + // intentionally mutable because it transitions to Ready below. + let inner = self.read(); + let route = inner + .routes + .get(&fingerprint) + .ok_or(RegistryError::UnknownRoute)?; + if !route.refs.contains_key(session_id) { + return Err(RegistryError::UnknownMcpSession); + } + if let Some(directive) = route.stable_current_directive(session_id, self.retry_after_ms) { + return Ok(directive); + } + drop(inner); + let mut inner = self.write(); let route = inner .routes @@ -886,6 +902,38 @@ impl RouteEntry { } } + // Return a directive without mutating the route. `Recovering` with an + // existing target deliberately returns None because the write-path + // promotes it back to Ready as part of serving the directive. + fn stable_current_directive( + &self, + session_id: &McpSessionId, + retry_after_ms: u64, + ) -> Option { + Some(match &self.state { + RouteState::Empty => BrokerDirective::WaitForWorker { retry_after_ms }, + RouteState::Activating { + owner, + launch: active_launch, + } if owner == session_id => active_launch.clone().into_directive(), + RouteState::Activating { .. } + | RouteState::Draining { .. } + | RouteState::Recovering { target: None, .. } => { + BrokerDirective::WaitForWorker { retry_after_ms } + } + RouteState::Ready { target } if !target.control_available() => { + BrokerDirective::WaitForWorker { retry_after_ms } + } + RouteState::Ready { target } => BrokerDirective::ReuseWorker { + endpoint: target.endpoint().to_owned(), + }, + RouteState::PassThrough { .. } => BrokerDirective::UsePassThrough, + RouteState::Recovering { + target: Some(_), .. + } => return None, + }) + } + fn after_reference_removed( &mut self, removed_session: &McpSessionId, diff --git a/crates/cli/src/daemon/broker/server/socket.rs b/crates/cli/src/daemon/broker/server/socket.rs index 586bf31fa..9c7f74ff5 100644 --- a/crates/cli/src/daemon/broker/server/socket.rs +++ b/crates/cli/src/daemon/broker/server/socket.rs @@ -310,6 +310,13 @@ async fn run(state: Arc, role: ComponentRole, local: bool, socket: if send.try_send(Message::Text(text.into())).is_err() { break; } continue; } + // ACKs affect only the peer's bounded directive queue. Challenge + // issuance has no routing effect either, so neither needs to + // wake every established control connection. + let local_control_command = matches!( + &request.command, + Command::Acknowledge { .. } | Command::Challenge(_) + ); let readiness = if let Command::Ready(payload) = &request.command { Some(socket_ready(&state, role, id.as_deref(), &generation, payload).await) } else { @@ -328,7 +335,7 @@ async fn run(state: Arc, role: ComponentRole, local: bool, socket: None => dispatch(&state, role, &mut id, &mut challenge, request.command).await, }; let success = response.status().is_success(); - if success && let Some(id) = &id { + if success && !local_control_command && let Some(id) = &id { let mut peers = lock(&state.sockets.peers); let entry = peers.entry(key(role, id)); use std::collections::hash_map::Entry; @@ -362,7 +369,9 @@ async fn run(state: Arc, role: ComponentRole, local: bool, socket: last_reply = Some((request.request_id, text.to_string(), event.clone())); let Ok(text) = serde_json::to_string(&event) else { break }; if send.try_send(Message::Text(text.into())).is_err() { break; } - state.sockets.changed.notify_waiters(); + if !local_control_command { + state.sockets.changed.notify_waiters(); + } } } } diff --git a/crates/cli/src/daemon/worker/managed.rs b/crates/cli/src/daemon/worker/managed.rs index 8ca31c888..754d8e378 100644 --- a/crates/cli/src/daemon/worker/managed.rs +++ b/crates/cli/src/daemon/worker/managed.rs @@ -106,6 +106,11 @@ const INTERNAL_DISPATCH_ROUTE_HEADER: &str = "x-nemo-relay-internal-dispatch-rou const INTERNAL_DISPATCH_BACKEND_HEADER: &str = "x-nemo-relay-internal-dispatch-backend"; const INTERNAL_RETRY_AWARE_HEADER: &str = "x-nemo-relay-internal-retry-aware"; +#[derive(Clone, Copy)] +pub(super) struct ProviderMiddlewareRequirements { + request_body_decode_required: bool, +} + /// Runtime-owned plugin activation, hook sessions, and response observation. pub(super) struct ManagedRuntime { config: GatewayConfig, @@ -125,7 +130,7 @@ impl ManagedRuntime { let activation = crate::server::initialize_plugin_host(config.plugin_config.clone(), dynamic_plugins) .await?; - if let Err(error) = reject_incompatible_execution_middleware() { + if let Err(error) = provider_middleware_requirements() { if let Some(activation) = activation { let _ = activation.clear(); } @@ -144,8 +149,18 @@ impl ManagedRuntime { /// Rechecks the transport contract before a provider body is polled. Plugin activation is /// normally static, but this also fails closed if a component installs middleware later. + #[cfg(test)] pub(super) fn ensure_streaming_transport_compatible(&self) -> Result<(), CliError> { - reject_incompatible_execution_middleware() + self.provider_middleware_requirements().map(|_| ()) + } + + /// Snapshots the middleware contract once for an incoming provider request. + /// This keeps late registration fail-closed while avoiding separate global + /// registry enumeration for raw-delivery compatibility and body decoding. + pub(super) fn provider_middleware_requirements( + &self, + ) -> Result { + provider_middleware_requirements() } pub(super) async fn close(&self) -> Result<(), CliError> { @@ -256,16 +271,33 @@ impl ManagedRuntime { } } + #[cfg(test)] pub(super) async fn proxy_provider( + &self, + upstream: PooledClient, + request: Request, + route: ProviderRoute, + ) -> Result, CliError> { + self.proxy_provider_with_requirements( + upstream, + request, + route, + self.provider_middleware_requirements()?, + ) + .await + } + + pub(super) async fn proxy_provider_with_requirements( &self, upstream: PooledClient, mut request: Request, route: ProviderRoute, + middleware: ProviderMiddlewareRequirements, ) -> Result, CliError> { let Some(surface) = provider_surface(request.uri().path()) else { return dispatch_unmanaged(upstream, request, route, &self.config).await; }; - if !request_body_decode_required()? { + if !middleware.request_body_decode_required { strip_worker_headers(request.headers_mut()); strip_untrusted_dispatch_headers(request.headers_mut()); let streaming_hint = request_streaming_hint(request.headers()); @@ -1015,14 +1047,30 @@ fn prepared_streaming(request: &LlmRequest) -> bool { stream_mode(request) } -fn request_body_decode_required() -> Result { +fn provider_middleware_requirements() -> Result { let kinds = BTreeSet::from([ RuntimeRegistrationKind::LlmSanitizeRequestGuardrail, RuntimeRegistrationKind::LlmConditionalExecutionGuardrail, RuntimeRegistrationKind::LlmRequestIntercept, + RuntimeRegistrationKind::LlmExecutionIntercept, + RuntimeRegistrationKind::LlmStreamExecutionIntercept, ]); let registrations = list_runtime_registrations(Some(&kinds)).map_err(CliError::from)?; - Ok(registrations.iter().any(registration_reads_request_body)) + let incompatible = incompatible_registration_names(®istrations); + if !incompatible.is_empty() { + return Err(CliError::Config(format!( + "daemon worker raw delivery is incompatible with LLM execution middleware: {}", + incompatible.join(", ") + ))); + } + Ok(ProviderMiddlewareRequirements { + request_body_decode_required: registrations.iter().any(registration_reads_request_body), + }) +} + +#[cfg(test)] +fn request_body_decode_required() -> Result { + provider_middleware_requirements().map(|requirements| requirements.request_body_decode_required) } fn registration_reads_request_body(registration: &RuntimeRegistrationIdentity) -> bool { @@ -1034,24 +1082,9 @@ fn registration_reads_request_body(registration: &RuntimeRegistrationIdentity) - ) } +#[cfg(test)] fn reject_incompatible_execution_middleware() -> Result<(), CliError> { - // Execution intercepts own the provider callback and may replace, suppress, retry, or mutate - // its result. The raw worker transport cannot safely invoke that contract while also returning - // the provider's response head and frames unchanged. Request intercepts and conditional - // execution guardrails remain supported above the transport boundary. - let kinds = BTreeSet::from([ - RuntimeRegistrationKind::LlmExecutionIntercept, - RuntimeRegistrationKind::LlmStreamExecutionIntercept, - ]); - let registrations = list_runtime_registrations(Some(&kinds)).map_err(CliError::from)?; - let incompatible = incompatible_registration_names(®istrations); - if incompatible.is_empty() { - return Ok(()); - } - Err(CliError::Config(format!( - "daemon worker raw delivery is incompatible with LLM execution middleware: {}", - incompatible.join(", ") - ))) + provider_middleware_requirements().map(|_| ()) } fn incompatible_registration_names(registrations: &[RuntimeRegistrationIdentity]) -> Vec { diff --git a/crates/cli/src/daemon/worker/runtime.rs b/crates/cli/src/daemon/worker/runtime.rs index 543687d95..7bf4bfb97 100644 --- a/crates/cli/src/daemon/worker/runtime.rs +++ b/crates/cli/src/daemon/worker/runtime.rs @@ -565,12 +565,19 @@ async fn proxy(State(state): State>, request: Request) -> let Some(route) = PublicRoute::from_path(request.uri().path()) else { return StatusCode::NOT_FOUND.into_response(); }; - if matches!(route, PublicRoute::Provider(_)) - && let Some(managed) = state.managed.as_ref() - && let Err(error) = managed.ensure_streaming_transport_compatible() - { - return route_failure_response(error); - } + let middleware = if matches!(route, PublicRoute::Provider(_)) { + match state + .managed + .as_ref() + .map(|managed| managed.provider_middleware_requirements()) + { + Some(Ok(requirements)) => Some(requirements), + Some(Err(error)) => return route_failure_response(error), + None => None, + } + } else { + None + }; let Some(in_flight) = state.admit() else { let mut response = message( StatusCode::SERVICE_UNAVAILABLE, @@ -601,7 +608,12 @@ async fn proxy(State(state): State>, request: Request) -> PublicRoute::Provider(provider) => { if let Some(managed) = state.managed.as_ref() { let response = managed - .proxy_provider(state.upstream.clone(), request, provider) + .proxy_provider_with_requirements( + state.upstream.clone(), + request, + provider, + middleware.expect("managed provider requests have middleware requirements"), + ) .await; return match response { Ok(response) => { diff --git a/crates/cli/src/gateway/mod.rs b/crates/cli/src/gateway/mod.rs index e1268db41..89077560d 100644 --- a/crates/cli/src/gateway/mod.rs +++ b/crates/cli/src/gateway/mod.rs @@ -161,10 +161,11 @@ pub(crate) async fn passthrough( ) -> Result, CliError> { state.touch(); let authorization = state.authorize_provider_request(request.headers_mut())?; - let prepared = prepare_gateway_request(&state.config, request, authorization).await?; + let mut prepared = prepare_gateway_request(&state.config, request, authorization).await?; + let start = take_llm_gateway_start(&mut prepared); let prep = state .sessions - .prepare_gateway_call(&prepared.headers, build_llm_gateway_start(&prepared)) + .prepare_gateway_call(&prepared.headers, start) .await?; run_managed_gateway(state, prepared, prep).await } @@ -747,8 +748,7 @@ fn client_sse_body( while let Some(item) = json_stream.next().await { match item { Ok(event_json) => { - let frame = encode_sse_frame(&event_json, route); - yield Ok::(Bytes::from(frame)); + yield Ok::(encode_sse_frame(&event_json, route)); } Err(error) => { guard.finish().await; @@ -864,19 +864,25 @@ impl Drop for GatewayCallGuard { // Formats one SSE frame from a parsed event payload. Anthropic and OpenAI Responses events carry // the event name in the `type` field, so it is mirrored back onto the `event:` line; OpenAI Chat // chunks have no event name and emit only `data:`. -fn encode_sse_frame(event_json: &Value, route: ProviderRoute) -> String { - let serialized = serde_json::to_string(event_json).unwrap_or_else(|_| "null".to_string()); +fn encode_sse_frame(event_json: &Value, route: ProviderRoute) -> Bytes { let event_name = match route { - ProviderRoute::AnthropicMessages | ProviderRoute::OpenAiResponses => event_json - .get("type") - .and_then(Value::as_str) - .map(ToOwned::to_owned), + ProviderRoute::AnthropicMessages | ProviderRoute::OpenAiResponses => { + event_json.get("type").and_then(Value::as_str) + } _ => None, }; - match event_name { - Some(name) => format!("event: {name}\ndata: {serialized}\n\n"), - None => format!("data: {serialized}\n\n"), + let mut frame = Vec::with_capacity(64); + if let Some(name) = event_name { + frame.extend_from_slice(b"event: "); + frame.extend_from_slice(name.as_bytes()); + frame.push(b'\n'); + } + frame.extend_from_slice(b"data: "); + if serde_json::to_writer(&mut frame, event_json).is_err() { + frame.extend_from_slice(b"null"); } + frame.extend_from_slice(b"\n\n"); + Bytes::from(frame) } // Forwards the buffered request to the upstream provider with only the safe request headers. This diff --git a/crates/cli/src/gateway/request.rs b/crates/cli/src/gateway/request.rs index 08dfb7c25..fbdb38cfc 100644 --- a/crates/cli/src/gateway/request.rs +++ b/crates/cli/src/gateway/request.rs @@ -211,6 +211,7 @@ fn passthrough_body_error(error: axum::Error) -> CliError { } } +#[cfg(test)] pub(super) fn build_llm_gateway_start(request: &PreparedGatewayRequest) -> LlmGatewayStart { build_llm_gateway_start_from_parts( &request.headers, @@ -221,6 +222,22 @@ pub(super) fn build_llm_gateway_start(request: &PreparedGatewayRequest) -> LlmGa ) } +/// Transfers the already-parsed request JSON into the session start event. +/// +/// The gateway only needs this representation once after request preparation, so moving it avoids +/// a second full copy for large provider prompts. The borrowed builder above remains available for +/// callers that must retain a prepared request. +pub(super) fn take_llm_gateway_start(request: &mut PreparedGatewayRequest) -> LlmGatewayStart { + let request_json = std::mem::take(&mut request.request_json); + build_llm_gateway_start_from_parts( + &request.headers, + &request.path, + request.provider, + request_json, + request.streaming, + ) +} + pub(super) fn build_llm_gateway_start_from_parts( headers: &HeaderMap, path: &str, diff --git a/crates/cli/src/sessions/idle.rs b/crates/cli/src/sessions/idle.rs index 6f90bf5dd..8ec60a00e 100644 --- a/crates/cli/src/sessions/idle.rs +++ b/crates/cli/src/sessions/idle.rs @@ -74,10 +74,16 @@ pub(super) async fn release_closed_owner_ids( if released_owner_ids.is_empty() { return; } - let mut owners = authenticated_owners.lock().await; let sessions = inner.lock().await; + let retained = released_owner_ids + .iter() + .filter(|session_id| sessions.contains_key(session_id.as_str())) + .cloned() + .collect::>(); + drop(sessions); + let mut owners = authenticated_owners.lock().await; owners.retain(|session_id, _| { - !released_owner_ids.contains(session_id) || sessions.contains_key(session_id) + !released_owner_ids.contains(session_id) || retained.contains(session_id) }); } diff --git a/crates/cli/src/sessions/mod.rs b/crates/cli/src/sessions/mod.rs index bc87aef24..541633d77 100644 --- a/crates/cli/src/sessions/mod.rs +++ b/crates/cli/src/sessions/mod.rs @@ -3,6 +3,7 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::time::{Duration, Instant}; use axum::http::HeaderMap; @@ -21,7 +22,7 @@ use nemo_relay::api::tool::{ tool_conditional_execution, tool_request_intercepts, }; use serde_json::{Map, Value, json}; -use tokio::sync::Mutex; +use tokio::sync::{Mutex, Notify}; use crate::agents::shared::adapters::{SKILL_LOAD_SOURCE_KEY, SKILL_LOAD_SOURCE_PROMPT_EXPANSION}; use crate::agents::shared::alignment::{ @@ -81,10 +82,43 @@ pub(crate) struct HookEffects { pub(crate) tool_argument_transform: Option, } +struct HookEventApplication { + released_owner_ids: HashSet, + subscriber_delivery: Option, + tool_argument_transform: Option, +} + +impl HookEventApplication { + fn empty() -> Self { + Self { + released_owner_ids: HashSet::new(), + subscriber_delivery: None, + tool_argument_transform: None, + } + } + + fn release(session_id: String) -> Self { + let mut application = Self::empty(); + application.released_owner_ids.insert(session_id); + application + } +} + #[derive(Clone)] pub(crate) struct SessionManager { inner: Arc>>, - authenticated_owners: Arc>>, + // The session map is only a directory. Work for one session is serialized by a + // dedicated gate so middleware awaiting inside that session does not hold the + // directory lock and block unrelated sessions. + session_gates: SessionGates, + // Applying a hook temporarily takes its session out of the directory while middleware runs. + // Track those operations so an idle shutdown or teardown cannot mistake them for no work. + session_activity: SessionActivity, + authenticated_owners: AuthenticatedOwners, + // Ownership is committed only after the hook batch succeeds. Reservations + // prevent another client from claiming the same session while that batch is + // running without holding the owner map across middleware awaits. + authenticated_reservations: Arc>>, // Cross-session alignment state owns child-session aliases and child-first SessionStart hooks. // Applies to Codex child threads today; the generic state lives in `alignment` so session code // only orchestrates when promotion is safe. @@ -92,6 +126,112 @@ pub(crate) struct SessionManager { default_config: GatewayConfig, } +pub(super) type SessionGates = Arc>>>>; +pub(super) type AuthenticatedOwners = Arc>>; +pub(super) type AuthenticatedReservations = Arc>>; + +#[derive(Clone)] +pub(super) struct SessionActivity { + active: Arc, + closing: Arc, + changed: Arc, +} + +impl SessionActivity { + fn new() -> Self { + Self { + active: Arc::new(AtomicUsize::new(0)), + closing: Arc::new(AtomicBool::new(false)), + changed: Arc::new(Notify::new()), + } + } + + pub(super) fn begin(&self) -> SessionActivityGuard { + self.active.fetch_add(1, Ordering::AcqRel); + SessionActivityGuard(self.clone()) + } + + fn has_active(&self) -> bool { + self.active.load(Ordering::Acquire) != 0 + } + + fn begin_closing(&self) { + self.closing.store(true, Ordering::Release); + } + + pub(super) fn is_closing(&self) -> bool { + self.closing.load(Ordering::Acquire) + } + + async fn wait_for_idle(&self) { + loop { + let changed = self.changed.notified(); + if !self.has_active() { + return; + } + changed.await; + } + } +} + +pub(super) struct SessionActivityGuard(SessionActivity); + +impl Drop for SessionActivityGuard { + fn drop(&mut self) { + if self.0.active.fetch_sub(1, Ordering::AcqRel) == 1 { + self.0.changed.notify_waiters(); + } + } +} + +pub(super) async fn session_gate(gates: &SessionGates, session_id: &str) -> Arc> { + let mut gates = gates.lock().await; + gates + .entry(session_id.to_string()) + .or_insert_with(|| Arc::new(Mutex::new(()))) + .clone() +} + +async fn owner_matches( + owners: &AuthenticatedOwners, + reservations: &Arc>>, + session_id: &str, + owner: &str, +) -> bool { + if owners + .lock() + .await + .get(session_id) + .is_some_and(|existing| existing == owner) + { + return true; + } + reservations + .lock() + .await + .get(session_id) + .is_some_and(|(reserved_owner, _)| reserved_owner == owner) +} + +fn release_reservations( + reservations: &mut HashMap, + session_ids: &HashSet, + owner: &str, +) { + for session_id in session_ids { + let remove = match reservations.get_mut(session_id) { + Some((reserved_owner, pending)) if reserved_owner == owner => { + *pending = pending.saturating_sub(1); + *pending == 0 + } + _ => false, + }; + if remove { + reservations.remove(session_id); + } + } +} + struct RoutingIdentityHeaderContext<'a> { session_id: &'a str, agent_kind: AgentKind, @@ -304,7 +444,10 @@ impl SessionManager { pub(crate) fn new(default_config: GatewayConfig) -> Self { Self { inner: Arc::new(Mutex::new(HashMap::new())), + session_gates: Arc::new(Mutex::new(HashMap::new())), + session_activity: SessionActivity::new(), authenticated_owners: Arc::new(Mutex::new(HashMap::new())), + authenticated_reservations: Arc::new(Mutex::new(HashMap::new())), alignment: Arc::new(Mutex::new(SessionAlignmentState::default())), default_config, } @@ -317,13 +460,17 @@ impl SessionManager { events: Vec, owner: &str, ) -> Result<(), CliError> { - let mut owners = self.authenticated_owners.lock().await; - let mut prospective_owners = owners.clone(); + let owners = self.authenticated_owners.lock().await; + let mut reservations = self.authenticated_reservations.lock().await; + let mut reserved_ids = HashSet::new(); for event in &events { let session_id = event.session_id(); - if prospective_owners + if owners .get(session_id) .is_some_and(|existing| existing != owner) + || reservations + .get(session_id) + .is_some_and(|(existing, _)| existing != owner) { return Err(CliError::Unauthorized(format!( "Relay hook client does not own session '{session_id}'" @@ -331,18 +478,45 @@ impl SessionManager { } } for event in &events { - prospective_owners - .entry(event.session_id().to_string()) - .or_insert_with(|| owner.to_string()); + let session_id = event.session_id(); + if owners.contains_key(session_id) || !reserved_ids.insert(session_id.to_string()) { + continue; + } + reservations + .entry(session_id.to_string()) + .and_modify(|(_, pending)| *pending += 1) + .or_insert_with(|| (owner.to_string(), 1)); } - let (released_owner_ids, _effects) = self - .apply_events_inner(headers, events, Some(&prospective_owners), Some(owner)) - .await?; - *owners = prospective_owners; drop(owners); - release_closed_owner_ids(&self.inner, &self.authenticated_owners, &released_owner_ids) - .await; - Ok(()) + drop(reservations); + + let result = self.apply_events_inner(headers, events, Some(owner)).await; + + let mut owners = self.authenticated_owners.lock().await; + let mut reservations = self.authenticated_reservations.lock().await; + match result { + Ok((released_owner_ids, _effects)) => { + for session_id in &reserved_ids { + owners + .entry(session_id.clone()) + .or_insert_with(|| owner.to_string()); + } + release_reservations(&mut reservations, &reserved_ids, owner); + drop(reservations); + drop(owners); + release_closed_owner_ids( + &self.inner, + &self.authenticated_owners, + &released_owner_ids, + ) + .await; + Ok(()) + } + Err(error) => { + release_reservations(&mut reservations, &reserved_ids, owner); + Err(error) + } + } } /// Evaluates a final host permission request inside its existing session scope. @@ -368,6 +542,8 @@ impl SessionManager { } } drop(owners); + let gate = session_gate(&self.session_gates, &event.session_id).await; + let _gate = gate.lock().await; let sessions = self.inner.lock().await; let session = sessions.get(&event.session_id).ok_or_else(|| { CliError::InvalidPayload(format!( @@ -462,7 +638,7 @@ impl SessionManager { headers: &HeaderMap, events: Vec, ) -> Result { - self.apply_events_inner(headers, events, None, None) + self.apply_events_inner(headers, events, None) .await .map(|(_, effects)| effects) } @@ -473,98 +649,174 @@ impl SessionManager { &self, headers: &HeaderMap, events: Vec, - authenticated_owners: Option<&HashMap>, authenticated_owner: Option<&str>, ) -> Result<(HashSet, HookEffects), CliError> { let mut effects = HookEffects::default(); let mut subscriber_deliveries = Vec::new(); let mut released_owner_ids = HashSet::new(); - let mut alignment_state = self.alignment.lock().await; - let mut sessions = self.inner.lock().await; + let config = self.default_config.session_config_from_headers(headers); + let authenticated = AuthenticatedRouting::new( + authenticated_owner, + &self.authenticated_owners, + &self.authenticated_reservations, + ); for event in events { - let original_session_id = event.session_id().to_string(); - let original_was_terminal = event.is_terminal(); - let mut event = event; - let config = self.default_config.session_config_from_headers(headers); - if queue_or_promote_child_start( - &mut event, - &mut sessions, - &mut alignment_state, - config.clone(), - authenticated_owners, - authenticated_owner, - ) - .await? - { - continue; - } - - if let Some(owners) = authenticated_owners - && let Some(alias) = alignment_state.alias_for_session(&original_session_id) - && let Some(alias_owner) = alias.authenticated_owner() - && owners - .get(&alias.parent_session_id) - .is_none_or(|existing| existing != alias_owner) - { - return Err(CliError::Unauthorized(format!( - "Relay hook client does not own session '{}'", - alias.parent_session_id - ))); - } - - let Some((event, session_id, is_agent_started)) = - route_event_for_session(event, &mut sessions, &mut alignment_state) - else { - if original_was_terminal { - released_owner_ids.insert(original_session_id); - } - continue; - }; - if original_was_terminal && original_session_id != session_id { - released_owner_ids.insert(original_session_id); - } - let event_kind = event_agent_kind(&event); - let (should_remove_session, subscriber_delivery, tool_argument_transform) = - apply_event_to_session( - &mut sessions, - &session_id, - event, - event_kind, - config.clone(), - is_agent_started, - ) + let application = self + .apply_hook_event(event, config.clone(), authenticated) .await?; - if let Some(subscriber_delivery) = subscriber_delivery { + released_owner_ids.extend(application.released_owner_ids); + if let Some(subscriber_delivery) = application.subscriber_delivery { subscriber_deliveries.push(subscriber_delivery); } - if tool_argument_transform.is_some() { - effects.tool_argument_transform = tool_argument_transform; - } - 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. - promote_pending_subagents_for_parent( - &mut sessions, - &mut alignment_state, - &session_id, - config.clone(), - authenticated_owners, - ) - .await?; - } - if should_remove_session { - sessions.remove(&session_id); - released_owner_ids.insert(session_id); + if application.tool_argument_transform.is_some() { + effects.tool_argument_transform = application.tool_argument_transform; } } - drop(sessions); - drop(alignment_state); for subscriber_delivery in subscriber_deliveries { subscriber_delivery.wait().await?; } Ok((released_owner_ids, effects)) } + async fn apply_hook_event( + &self, + event: NormalizedEvent, + config: SessionConfig, + authenticated: AuthenticatedRouting<'_>, + ) -> Result { + let original_session_id = event.session_id().to_string(); + let original_was_terminal = event.is_terminal(); + let mut event = event; + // Provider-specific child discovery can await harness state. Do it before taking the + // shared directory or alignment locks; the remaining routing work is map-only. + let pending_child = alignment::pending_subagent_start(&mut event).await; + let Some((event, session_id, is_agent_started)) = self + .route_hook_event( + pending_child, + event, + &original_session_id, + config.clone(), + authenticated, + ) + .await? + else { + return Ok(if original_was_terminal { + HookEventApplication::release(original_session_id) + } else { + HookEventApplication::empty() + }); + }; + let mut application = HookEventApplication::empty(); + if original_was_terminal && original_session_id != session_id { + application + .released_owner_ids + .insert(original_session_id.clone()); + } + let event_kind = event_agent_kind(&event); + let applied = SessionEventApplier::new( + &self.inner, + &self.session_gates, + &self.session_activity, + config.clone(), + ) + .apply(&session_id, event, event_kind, is_agent_started) + .await?; + let Some((should_remove_session, subscriber_delivery, tool_argument_transform)) = applied + else { + if original_was_terminal { + application.released_owner_ids.insert(original_session_id); + } + return Ok(application); + }; + application.subscriber_delivery = subscriber_delivery; + application.tool_argument_transform = tool_argument_transform; + if is_agent_started { + self.promote_pending_children(&session_id, config, authenticated) + .await?; + } + if should_remove_session { + application.released_owner_ids.insert(session_id); + } + Ok(application) + } + + async fn route_hook_event( + &self, + pending_child: Option<(String, alignment::PendingSubagentStart)>, + event: NormalizedEvent, + original_session_id: &str, + config: SessionConfig, + authenticated: AuthenticatedRouting<'_>, + ) -> Result, CliError> { + let (routed, alias) = { + let mut alignment_state = self.alignment.lock().await; + let mut sessions = self.inner.lock().await; + if queue_or_promote_child_start( + pending_child, + &mut sessions, + &mut alignment_state, + config, + authenticated, + ) + .await? + { + return Ok(None); + } + let alias = alignment_state.alias_for_session(original_session_id); + (route_event_for_session(event, &mut alignment_state), alias) + }; + self.authorize_alias_owner(alias.as_ref(), authenticated) + .await?; + Ok(routed) + } + + async fn authorize_alias_owner( + &self, + alias: Option<&SessionAlias>, + authenticated: AuthenticatedRouting<'_>, + ) -> Result<(), CliError> { + let (Some(owner), Some(alias), Some(alias_owner)) = ( + authenticated.owner, + alias, + alias.and_then(SessionAlias::authenticated_owner), + ) else { + return Ok(()); + }; + if alias_owner == owner + && owner_matches( + &self.authenticated_owners, + &self.authenticated_reservations, + &alias.parent_session_id, + alias_owner, + ) + .await + { + return Ok(()); + } + Err(CliError::Unauthorized(format!( + "Relay hook client does not own session '{}'", + alias.parent_session_id + ))) + } + + async fn promote_pending_children( + &self, + parent_session_id: &str, + config: SessionConfig, + authenticated: AuthenticatedRouting<'_>, + ) -> Result<(), CliError> { + let mut alignment_state = self.alignment.lock().await; + let mut sessions = self.inner.lock().await; + promote_pending_subagents_for_parent( + &mut sessions, + &mut alignment_state, + parent_session_id, + config, + authenticated, + ) + .await + } + /// Legacy manual-lifecycle entry point retained for tests that drive correlation behavior /// directly. Production gateway traffic uses [`Self::prepare_gateway_call`] + /// `llm_call_execute` / `llm_stream_call_execute` so the runtime owns start/end events. @@ -587,12 +839,17 @@ impl SessionManager { let mut start = start; let config = self.default_config.session_config_from_headers(headers); let alias = self.resolve_start_alias(&mut start, config.clone()).await?; + let session_id = { + let sessions = self.inner.lock().await; + start + .session_id + .clone() + .or_else(|| single_active_session_id(&sessions)) + .unwrap_or_else(|| format!("{}-gateway", AgentKind::Gateway.as_str())) + }; + let gate = session_gate(&self.session_gates, &session_id).await; + let _gate = gate.lock().await; let mut sessions = self.inner.lock().await; - let session_id = start - .session_id - .clone() - .or_else(|| single_active_session_id(&sessions)) - .unwrap_or_else(|| format!("{}-gateway", AgentKind::Gateway.as_str())); let inferred_agent_kind = alignment::agent_kind_for_gateway_provider(&start.provider); let session = sessions .entry(session_id.clone()) @@ -623,8 +880,13 @@ impl SessionManager { let mut start = start; let config = self.default_config.session_config_from_headers(headers); self.resolve_start_alias(&mut start, config.clone()).await?; + let (session_id, session_finish) = { + let sessions = self.inner.lock().await; + gateway_session_for_call(&start, &sessions) + }; + let gate = session_gate(&self.session_gates, &session_id).await; + let _gate = gate.lock().await; let mut sessions = self.inner.lock().await; - let (session_id, session_finish) = gateway_session_for_call(&start, &sessions); // Match `start_llm`: when this path creates a brand-new session (real agent's gateway // request beats its SessionStart hook), label the session by the provider so ATIF and // Phoenix scopes carry the agent identity instead of freezing on "gateway". @@ -666,6 +928,8 @@ impl SessionManager { /// in-flight counter to prevent the idle sweeper from closing a turn while an upstream /// provider request or streaming response is still active. pub(crate) async fn finish_gateway_call(&self, session_id: &str, finish: GatewaySessionFinish) { + let gate = session_gate(&self.session_gates, session_id).await; + let _gate = gate.lock().await; let mut sessions = self.inner.lock().await; if let Some(session) = sessions.get_mut(session_id) { session.finish_gateway_call(); @@ -723,6 +987,9 @@ impl SessionManager { /// gateway calls still block idle shutdown; [`Self::close_all`] balances the dormant agent scope /// when the gateway exits. pub(crate) async fn has_open_sessions(&self) -> bool { + if self.session_activity.has_active() { + return true; + } self.inner .lock() .await @@ -746,6 +1013,8 @@ impl SessionManager { let session_id = active.session_id.clone(); let llm_id = active.handle.uuid.to_string(); let owner_subagent_id = active.owner_subagent_id.clone(); + let gate = session_gate(&self.session_gates, &session_id).await; + let _gate = gate.lock().await; { let mut sessions = self.inner.lock().await; let Some(session) = sessions.get_mut(&session_id) else { @@ -801,6 +1070,8 @@ impl SessionManager { ), None => (session_id.to_string(), owner_subagent_id), }; + let gate = session_gate(&self.session_gates, &session_id).await; + let _gate = gate.lock().await; let mut sessions = self.inner.lock().await; if let Some(session) = sessions.get_mut(&session_id) { session.record_completed_llm_response(response, owner_subagent_id); @@ -813,6 +1084,8 @@ impl SessionManager { /// deterministic lifecycle boundary for those sessions, so close open scopes while /// observability plugins are still active. Applies to Codex transparent runs today. pub(crate) async fn close_all(&self, reason: &str) -> Result<(), CliError> { + self.session_activity.begin_closing(); + self.session_activity.wait_for_idle().await; self.alignment.lock().await.clear(); self.authenticated_owners.lock().await.clear(); let mut sessions = { @@ -888,6 +1161,10 @@ impl SessionManager { } } } + // The parent ownership decision is complete. Do not carry the owner + // map lock into session promotion; routing acquires the directory first + // before checking ownership. + drop(owners); let mut sessions = self.inner.lock().await; let alias = promote_pending_subagent( &mut sessions, @@ -895,7 +1172,11 @@ impl SessionManager { session_id, pending, config, - Some(&owners), + AuthenticatedRouting::new( + None, + &self.authenticated_owners, + &self.authenticated_reservations, + ), ) .await?; if let Some(alias) = alias.as_ref() { diff --git a/crates/cli/src/sessions/routing.rs b/crates/cli/src/sessions/routing.rs index b82ac4327..245f7d4fe 100644 --- a/crates/cli/src/sessions/routing.rs +++ b/crates/cli/src/sessions/routing.rs @@ -4,18 +4,129 @@ //! Child-session aliasing and lifecycle-event routing. use std::collections::HashMap; +use std::sync::Arc; use nemo_relay::api::runtime::SubscriberDelivery; use serde_json::Value; +use tokio::sync::Mutex; use crate::agents::shared::alignment::{ - self, PendingSubagentStart, SessionAlias, SessionAlignmentState, merge_metadata, + PendingSubagentStart, SessionAlias, SessionAlignmentState, merge_metadata, }; use crate::configuration::SessionConfig; use crate::error::CliError; use crate::events::{AgentKind, NormalizedEvent, SessionEvent}; -use super::{LlmGatewayStart, Session, ToolArgumentTransform}; +use super::{ + AuthenticatedOwners, AuthenticatedReservations, LlmGatewayStart, Session, SessionActivity, + SessionGates, ToolArgumentTransform, session_gate, +}; + +#[derive(Clone, Copy)] +pub(super) struct AuthenticatedRouting<'a> { + pub(super) owner: Option<&'a str>, + owners: Option<&'a AuthenticatedOwners>, + reservations: Option<&'a AuthenticatedReservations>, +} + +impl<'a> AuthenticatedRouting<'a> { + pub(super) fn new( + owner: Option<&'a str>, + owners: &'a AuthenticatedOwners, + reservations: &'a AuthenticatedReservations, + ) -> Self { + Self { + owner, + owners: owner.map(|_| owners), + reservations: owner.map(|_| reservations), + } + } +} + +pub(super) struct SessionEventApplier<'a> { + sessions: &'a Arc>>, + gates: &'a SessionGates, + activity: &'a SessionActivity, + config: SessionConfig, +} + +impl<'a> SessionEventApplier<'a> { + pub(super) fn new( + sessions: &'a Arc>>, + gates: &'a SessionGates, + activity: &'a SessionActivity, + config: SessionConfig, + ) -> Self { + Self { + sessions, + gates, + activity, + config, + } + } + + pub(super) async fn apply( + &self, + session_id: &str, + event: NormalizedEvent, + event_kind: AgentKind, + is_agent_started: bool, + ) -> Result< + Option<( + bool, + Option, + Option, + )>, + CliError, + > { + let _activity = self.activity.begin(); + let gate = session_gate(self.gates, session_id).await; + let _gate = gate.lock().await; + if self.activity.is_closing() { + return Ok(None); + } + let session = { + let mut sessions = self.sessions.lock().await; + sessions.remove(session_id) + }; + if session.is_none() && event.is_terminal() { + return Ok(None); + } + let mut session = session.unwrap_or_else(|| { + Session::new(session_id.to_string(), event_kind, self.config.clone()) + }); + if is_agent_started + && session.agent_kind == AgentKind::Gateway + && event_kind != AgentKind::Gateway + { + session.agent_kind = event_kind; + } + match session.apply(event).await { + Ok(subscriber_delivery) => { + let is_empty = session.is_empty(); + let tool_argument_transform = session.take_tool_argument_transform(); + if !is_empty { + self.sessions + .lock() + .await + .insert(session_id.to_string(), session); + } + Ok(Some(( + is_empty, + subscriber_delivery, + tool_argument_transform, + ))) + } + Err(error) => { + self.sessions + .lock() + .await + .insert(session_id.to_string(), session); + Err(error) + } + } + } +} pub(super) fn apply_start_alias(start: &mut LlmGatewayStart, alias: &SessionAlias) { start.session_id = Some(alias.parent_session_id.clone()); @@ -24,18 +135,16 @@ pub(super) fn apply_start_alias(start: &mut LlmGatewayStart, alias: &SessionAlia } pub(super) async fn queue_or_promote_child_start( - event: &mut NormalizedEvent, + pending_child: Option<(String, PendingSubagentStart)>, sessions: &mut HashMap, alignment_state: &mut SessionAlignmentState, config: SessionConfig, - authenticated_owners: Option<&HashMap>, - authenticated_owner: Option<&str>, + authenticated: AuthenticatedRouting<'_>, ) -> Result { - let Some((child_session_id, mut pending)) = alignment::pending_subagent_start(event).await - else { + let Some((child_session_id, mut pending)) = pending_child else { return Ok(false); }; - pending.set_authenticated_owner(authenticated_owner.map(ToOwned::to_owned)); + pending.set_authenticated_owner(authenticated.owner.map(ToOwned::to_owned)); if sessions .get(&child_session_id) .is_some_and(|session| !session.can_reparent_as_subagent_alias()) @@ -44,10 +153,12 @@ pub(super) async fn queue_or_promote_child_start( } if sessions.contains_key(pending.parent_session_id()) { if !parent_owner_matches( - authenticated_owners, + authenticated, pending.parent_session_id(), pending.authenticated_owner(), - ) { + ) + .await + { return Err(CliError::Unauthorized(format!( "Relay hook client does not own session '{}'", pending.parent_session_id() @@ -60,7 +171,7 @@ pub(super) async fn queue_or_promote_child_start( child_session_id, pending, config, - authenticated_owners, + authenticated, ) .await?; } else { @@ -70,52 +181,21 @@ pub(super) async fn queue_or_promote_child_start( Ok(true) } -pub(super) async fn apply_event_to_session( - sessions: &mut HashMap, - session_id: &str, - event: NormalizedEvent, - event_kind: AgentKind, - config: SessionConfig, - is_agent_started: bool, -) -> Result< - ( - bool, - Option, - Option, - ), - CliError, -> { - let session = sessions - .entry(session_id.to_string()) - .or_insert_with(|| Session::new(session_id.to_string(), event_kind, config)); - if is_agent_started - && session.agent_kind == AgentKind::Gateway - && event_kind != AgentKind::Gateway - { - session.agent_kind = event_kind; - } - let subscriber_delivery = session.apply(event).await?; - let tool_argument_transform = session.take_tool_argument_transform(); - Ok(( - session.is_empty(), - subscriber_delivery, - tool_argument_transform, - )) -} - pub(super) async fn promote_pending_subagents_for_parent( sessions: &mut HashMap, alignment_state: &mut SessionAlignmentState, parent_session_id: &str, config: SessionConfig, - authenticated_owners: Option<&HashMap>, + authenticated: AuthenticatedRouting<'_>, ) -> Result<(), CliError> { for (child_session_id, pending) in alignment_state.pending_for_parent(parent_session_id) { if !parent_owner_matches( - authenticated_owners, + authenticated, parent_session_id, pending.authenticated_owner(), - ) { + ) + .await + { continue; } promote_pending_subagent( @@ -124,7 +204,7 @@ pub(super) async fn promote_pending_subagents_for_parent( child_session_id, pending, config.clone(), - authenticated_owners, + authenticated, ) .await?; } @@ -137,7 +217,7 @@ pub(super) async fn promote_pending_subagent( child_session_id: String, pending: PendingSubagentStart, config: SessionConfig, - authenticated_owners: Option<&HashMap>, + authenticated: AuthenticatedRouting<'_>, ) -> Result, CliError> { if sessions .get(&child_session_id) @@ -148,10 +228,12 @@ pub(super) async fn promote_pending_subagent( sessions.remove(&child_session_id); let parent_session_id = pending.parent_session_id().to_string(); if !parent_owner_matches( - authenticated_owners, + authenticated, &parent_session_id, pending.authenticated_owner(), - ) { + ) + .await + { return Ok(None); } let parent_session = sessions @@ -181,30 +263,30 @@ pub(super) async fn promote_pending_subagent( Ok(Some(alias)) } -fn parent_owner_matches( - authenticated_owners: Option<&HashMap>, +async fn parent_owner_matches( + authenticated: AuthenticatedRouting<'_>, parent_session_id: &str, pending_owner: Option<&str>, ) -> bool { - match (authenticated_owners, pending_owner) { - (Some(owners), Some(owner)) => owners - .get(parent_session_id) - .is_some_and(|existing| existing == owner), + match ( + authenticated.owners, + authenticated.reservations, + pending_owner, + ) { + (Some(owners), Some(reservations), Some(owner)) => { + super::owner_matches(owners, reservations, parent_session_id, owner).await + } _ => true, } } pub(super) fn route_event_for_session( event: NormalizedEvent, - sessions: &mut HashMap, alignment_state: &mut SessionAlignmentState, ) -> Option<(NormalizedEvent, String, bool)> { let event = alignment_state.route_event(event); let session_id = event.session_id().to_string(); let is_agent_started = matches!(&event, NormalizedEvent::AgentStarted(_)); - if event.is_terminal() && !sessions.contains_key(&session_id) { - return None; - } Some((event, session_id, is_agent_started)) } diff --git a/crates/cli/tests/coverage/daemon/registry_tests.rs b/crates/cli/tests/coverage/daemon/registry_tests.rs index 145b3b54a..2b901ed00 100644 --- a/crates/cli/tests/coverage/daemon/registry_tests.rs +++ b/crates/cli/tests/coverage/daemon/registry_tests.rs @@ -888,6 +888,31 @@ fn recovery_completion_covers_live_empty_and_draining_routes() { ); } +#[test] +fn current_directive_promotes_a_recovered_target_before_reusing_it() { + let registry = Registry::new(false); + let fingerprint = fingerprint(33); + let token = TokenDigest::from_token(b"token-33"); + let mcp = session("mcp"); + registry + .register_mcp(registration(fingerprint, token, "mcp"), launch("launch")) + .unwrap(); + registry + .begin_recovery(fingerprint, Some(worker("survivor")), 100) + .unwrap(); + + assert_eq!( + registry.current_directive(fingerprint, &mcp).unwrap(), + BrokerDirective::ReuseWorker { + endpoint: "http://127.0.0.1:41000".to_owned(), + } + ); + assert_eq!( + registry.snapshot(fingerprint).unwrap().state, + RouteStateKind::Ready + ); +} + #[test] fn registry_rejects_stale_worker_generations_and_invalid_state_transitions() { let registry = Registry::new(false); diff --git a/crates/cli/tests/coverage/shared/session_tests.rs b/crates/cli/tests/coverage/shared/session_tests.rs index 6c5da80a8..6c8f22a31 100644 --- a/crates/cli/tests/coverage/shared/session_tests.rs +++ b/crates/cli/tests/coverage/shared/session_tests.rs @@ -4,6 +4,9 @@ use axum::http::HeaderMap; use nemo_relay::api::event::{Event, ScopeCategory}; use nemo_relay::api::llm::{LlmCallExecuteParams, llm_call_execute}; +use nemo_relay::api::registry::{ + deregister_tool_conditional_execution_guardrail, register_tool_conditional_execution_guardrail, +}; use nemo_relay::api::runtime::EventSubscriberFn; use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; use nemo_relay::codec::resolve::{ @@ -1776,6 +1779,87 @@ async fn terminal_subscriber_wait_releases_session_manager_locks() { .unwrap(); } +#[tokio::test] +async fn slow_tool_guardrail_does_not_block_another_session() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + const GUARDRAIL: &str = "cli-session-isolation-slow-tool"; + const TOOL: &str = "session-isolation-slow-tool"; + let _ = deregister_tool_conditional_execution_guardrail(GUARDRAIL); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let started_tx = Arc::new(StdMutex::new(Some(started_tx))); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let release_rx = Arc::new(StdMutex::new(Some(release_rx))); + register_tool_conditional_execution_guardrail( + GUARDRAIL, + 1, + Arc::new(move |name, _| { + let started_tx = Arc::clone(&started_tx); + let release_rx = Arc::clone(&release_rx); + let blocks = name == TOOL; + Box::pin(async move { + if blocks { + if let Some(started) = started_tx.lock().unwrap().take() { + let _ = started.send(()); + } + let release = { release_rx.lock().unwrap().take() }; + if let Some(release) = release { + let _ = release.await; + } + } + Ok(None) + }) + }), + ) + .unwrap(); + + let manager = SessionManager::new(session_test_config()); + let blocked_manager = manager.clone(); + let blocked = tokio::spawn(async move { + blocked_manager + .apply_events( + &HeaderMap::new(), + vec![NormalizedEvent::ToolStarted(ToolEvent { + session_id: "blocked-tool-session".into(), + agent_kind: AgentKind::Codex, + event_name: "PreToolUse".into(), + tool_call_id: "blocked-tool".into(), + tool_name: TOOL.into(), + subagent_id: None, + arguments: json!({}), + result: Value::Null, + status: None, + payload: json!({}), + metadata: json!({}), + })], + ) + .await + }); + tokio::time::timeout(Duration::from_secs(1), started_rx) + .await + .expect("slow guardrail should start") + .unwrap(); + + let parallel = tokio::time::timeout( + Duration::from_millis(250), + manager.apply_events( + &HeaderMap::new(), + vec![NormalizedEvent::Notification(codex_session_event( + "parallel-tool-session", + "notification", + json!({ "session_id": "parallel-tool-session" }), + ))], + ), + ) + .await; + + release_tx.send(()).unwrap(); + blocked.await.unwrap().unwrap(); + deregister_tool_conditional_execution_guardrail(GUARDRAIL).unwrap(); + parallel + .expect("another session must progress while a tool guardrail is waiting") + .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/codec/streaming.rs b/crates/core/src/codec/streaming.rs index e1bbbd777..d9740679c 100644 --- a/crates/core/src/codec/streaming.rs +++ b/crates/core/src/codec/streaming.rs @@ -99,6 +99,11 @@ pub trait StreamingCodec: Send + Sync { #[derive(Default)] pub struct SseEventDecoder { buffer: String, + // Start the next delimiter search near the only position where a newly + // appended chunk can complete an unfinished `\n\n` terminator. Completed + // frames are compacted once after each push instead of shifting the buffer + // for every frame. + scan_from: usize, } /// One decoded SSE frame, paired with the parsed `data:` payload. @@ -143,15 +148,25 @@ impl SseEventDecoder { // and remove it only when the next byte completes the sequence. if self.buffer.ends_with('\r') && bytes.first() == Some(&b'\n') { self.buffer.pop(); + // Removing the trailing CR means the preceding LF can now pair with the incoming LF. + self.scan_from = self.scan_from.min(self.buffer.len().saturating_sub(1)); + } + let chunk = String::from_utf8_lossy(bytes); + if chunk.contains("\r\n") { + self.buffer.push_str(&chunk.replace("\r\n", "\n")); + } else { + self.buffer.push_str(&chunk); } - let chunk = String::from_utf8_lossy(bytes).replace("\r\n", "\n"); - self.buffer.push_str(&chunk); let mut results = Vec::new(); - while let Some(cut) = self.buffer.find("\n\n") { - let frame: String = self.buffer.drain(..cut).collect(); - // Drop the `\n\n` terminator itself. - self.buffer.drain(..2); - match parse_sse_frame(&frame) { + let mut consumed = 0; + let mut search_from = self.scan_from.min(self.buffer.len()); + while let Some(relative_cut) = self.buffer[search_from..].find("\n\n") { + let cut = search_from + relative_cut; + // `consumed` marks the start of the unparsed frame. Borrowing the + // frame avoids allocating a string for every SSE event. + let frame = &self.buffer[consumed..cut]; + consumed = cut + 2; + match parse_sse_frame(frame) { Ok(Some(event)) => results.push(Ok(event)), Ok(None) => {} Err(error) => { @@ -159,6 +174,16 @@ impl SseEventDecoder { break; } } + search_from = consumed; + } + if consumed > 0 { + self.buffer.drain(..consumed); + self.scan_from = 0; + } + if results.last().is_none_or(Result::is_ok) { + // The next terminator can only start at the prior final byte; do + // not rescan an incomplete frame from its beginning on every chunk. + self.scan_from = self.buffer.len().saturating_sub(1); } results } diff --git a/crates/core/src/observability/atof.rs b/crates/core/src/observability/atof.rs index 03fc022a9..76520ae16 100644 --- a/crates/core/src/observability/atof.rs +++ b/crates/core/src/observability/atof.rs @@ -440,6 +440,9 @@ impl AtofExporter { state.last_error = Some(error); return; } + if state.endpoints.is_empty() { + return; + } let Ok(raw_json) = serde_json::to_string(&value) else { state.last_error = Some("failed to serialize ATOF event".to_string()); return; diff --git a/crates/core/tests/unit/codec/streaming_tests.rs b/crates/core/tests/unit/codec/streaming_tests.rs index 3a9856429..9b0718be7 100644 --- a/crates/core/tests/unit/codec/streaming_tests.rs +++ b/crates/core/tests/unit/codec/streaming_tests.rs @@ -112,3 +112,18 @@ fn preserves_successes_before_a_later_parse_error() { assert!(error.contains("not valid json"), "{error}"); assert!(results.next().is_none()); } + +#[test] +fn resumes_with_frames_after_a_malformed_frame() { + let mut decoder = SseEventDecoder::new(); + let results = decoder.push_bytes_results( + b"data: {\"chunk\":\"first\"}\n\ndata: {not valid json}\n\ndata: {\"chunk\":\"later\"}\n\n", + ); + assert_eq!(results.len(), 2); + assert_eq!(results[0].as_ref().unwrap().data, json!({"chunk": "first"})); + assert!(results[1].is_err()); + + let resumed = decoder.push_bytes_results(b""); + assert_eq!(resumed.len(), 1); + assert_eq!(resumed[0].as_ref().unwrap().data, json!({"chunk": "later"})); +} From ee59ebcaa6e78ed890fb3da8e3c20c3f8ebac672 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Fri, 11 Sep 2026 18:15:09 -0400 Subject: [PATCH 2/4] fix: harden session concurrency cleanup Signed-off-by: Will Killian --- crates/cli/src/sessions/idle.rs | 33 ++-- crates/cli/src/sessions/mod.rs | 32 +++- .../tests/coverage/shared/session_tests.rs | 167 +++++++++++++++++- crates/core/src/codec/streaming.rs | 4 +- .../core/tests/unit/codec/streaming_tests.rs | 17 ++ 5 files changed, 233 insertions(+), 20 deletions(-) diff --git a/crates/cli/src/sessions/idle.rs b/crates/cli/src/sessions/idle.rs index 8ec60a00e..3df885406 100644 --- a/crates/cli/src/sessions/idle.rs +++ b/crates/cli/src/sessions/idle.rs @@ -13,7 +13,7 @@ use tokio::sync::Mutex; use crate::agents::shared::alignment::SessionAlignmentState; use crate::error::CliError; -use super::Session; +use super::{Session, SessionGates, session_gate}; pub(super) const AGENT_IDLE_TIMEOUT: Duration = Duration::from_secs(30); pub(super) const AGENT_IDLE_SWEEP_INTERVAL: Duration = Duration::from_secs(5); @@ -35,6 +35,7 @@ pub(super) async fn close_sessions_for_shutdown( pub(super) async fn close_idle_sessions_from_parts( inner: &Arc>>, + session_gates: &SessionGates, authenticated_owners: &Arc>>, alignment: &Arc>, now: Instant, @@ -62,29 +63,31 @@ pub(super) async fn close_idle_sessions_from_parts( let cleanup_sessions = restore_retained_sessions(inner, retained_sessions, &closed_subagents).await; clear_closed_subagents(alignment, closed_subagents, &cleanup_sessions).await; - release_closed_owner_ids(inner, authenticated_owners, &released_owner_ids).await; + release_closed_owner_ids( + inner, + session_gates, + authenticated_owners, + &released_owner_ids, + ) + .await; first_error.map_or(Ok(closed_turns), Err) } pub(super) async fn release_closed_owner_ids( inner: &Arc>>, + session_gates: &SessionGates, authenticated_owners: &Arc>>, released_owner_ids: &HashSet, ) { - if released_owner_ids.is_empty() { - return; + for session_id in released_owner_ids { + // The gate protects this presence check and deletion from a hook recreating the same + // session while it is applying middleware. + let gate = session_gate(session_gates, session_id).await; + let _gate = gate.lock().await; + if !inner.lock().await.contains_key(session_id) { + authenticated_owners.lock().await.remove(session_id); + } } - let sessions = inner.lock().await; - let retained = released_owner_ids - .iter() - .filter(|session_id| sessions.contains_key(session_id.as_str())) - .cloned() - .collect::>(); - drop(sessions); - let mut owners = authenticated_owners.lock().await; - owners.retain(|session_id, _| { - !released_owner_ids.contains(session_id) || retained.contains(session_id) - }); } async fn take_idle_sessions( diff --git a/crates/cli/src/sessions/mod.rs b/crates/cli/src/sessions/mod.rs index 541633d77..35a654669 100644 --- a/crates/cli/src/sessions/mod.rs +++ b/crates/cli/src/sessions/mod.rs @@ -123,6 +123,9 @@ pub(crate) struct SessionManager { // Applies to Codex child threads today; the generic state lives in `alignment` so session code // only orchestrates when promotion is safe. alignment: Arc>, + // Alignment routes coordinate ownership checks and session-directory updates. This lock is + // never held while middleware or subscriber delivery runs. + alignment_routing: Arc>, default_config: GatewayConfig, } @@ -449,6 +452,7 @@ impl SessionManager { authenticated_owners: Arc::new(Mutex::new(HashMap::new())), authenticated_reservations: Arc::new(Mutex::new(HashMap::new())), alignment: Arc::new(Mutex::new(SessionAlignmentState::default())), + alignment_routing: Arc::new(Mutex::new(())), default_config, } } @@ -506,6 +510,7 @@ impl SessionManager { drop(owners); release_closed_owner_ids( &self.inner, + &self.session_gates, &self.authenticated_owners, &released_owner_ids, ) @@ -578,6 +583,7 @@ impl SessionManager { /// shutdown paths. pub(crate) fn start_idle_sweeper(&self) { let inner = Arc::downgrade(&self.inner); + let session_gates = Arc::downgrade(&self.session_gates); let authenticated_owners = Arc::downgrade(&self.authenticated_owners); let alignment = Arc::downgrade(&self.alignment); tokio::spawn(async move { @@ -585,8 +591,9 @@ impl SessionManager { interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); loop { interval.tick().await; - let (Some(inner), Some(authenticated_owners), Some(alignment)) = ( + let (Some(inner), Some(session_gates), Some(authenticated_owners), Some(alignment)) = ( inner.upgrade(), + session_gates.upgrade(), authenticated_owners.upgrade(), alignment.upgrade(), ) else { @@ -594,6 +601,7 @@ impl SessionManager { }; if let Err(error) = close_idle_sessions_from_parts( &inner, + &session_gates, &authenticated_owners, &alignment, Instant::now(), @@ -748,6 +756,17 @@ impl SessionManager { config: SessionConfig, authenticated: AuthenticatedRouting<'_>, ) -> Result, CliError> { + let _routing = self.alignment_routing.lock().await; + let parent_gate = match pending_child.as_ref() { + Some((_, pending)) => { + Some(session_gate(&self.session_gates, pending.parent_session_id()).await) + } + None => None, + }; + let _parent_gate = match parent_gate.as_ref() { + Some(gate) => Some(gate.lock().await), + None => None, + }; let (routed, alias) = { let mut alignment_state = self.alignment.lock().await; let mut sessions = self.inner.lock().await; @@ -805,6 +824,9 @@ impl SessionManager { config: SessionConfig, authenticated: AuthenticatedRouting<'_>, ) -> Result<(), CliError> { + let _routing = self.alignment_routing.lock().await; + let gate = session_gate(&self.session_gates, parent_session_id).await; + let _gate = gate.lock().await; let mut alignment_state = self.alignment.lock().await; let mut sessions = self.inner.lock().await; promote_pending_subagents_for_parent( @@ -929,7 +951,7 @@ impl SessionManager { /// provider request or streaming response is still active. pub(crate) async fn finish_gateway_call(&self, session_id: &str, finish: GatewaySessionFinish) { let gate = session_gate(&self.session_gates, session_id).await; - let _gate = gate.lock().await; + let gate_guard = gate.lock().await; let mut sessions = self.inner.lock().await; if let Some(session) = sessions.get_mut(session_id) { session.finish_gateway_call(); @@ -944,10 +966,12 @@ impl SessionManager { }); let mut closing = completed.then(|| sessions.remove(session_id)).flatten(); drop(sessions); + drop(gate_guard); if completed { release_closed_owner_ids( &self.inner, + &self.session_gates, &self.authenticated_owners, &HashSet::from([session_id.to_string()]), ) @@ -1107,6 +1131,7 @@ impl SessionManager { ) -> Result { close_idle_sessions_from_parts( &self.inner, + &self.session_gates, &self.authenticated_owners, &self.alignment, now, @@ -1128,6 +1153,7 @@ impl SessionManager { let Some(session_id) = start.session_id.clone() else { return Ok(None); }; + let _routing = self.alignment_routing.lock().await; let mut owners = self.authenticated_owners.lock().await; let mut alignment_state = self.alignment.lock().await; if let Some(alias) = alignment_state.alias_for_session(&session_id) { @@ -1165,6 +1191,8 @@ impl SessionManager { // map lock into session promotion; routing acquires the directory first // before checking ownership. drop(owners); + let gate = session_gate(&self.session_gates, pending.parent_session_id()).await; + let _gate = gate.lock().await; let mut sessions = self.inner.lock().await; let alias = promote_pending_subagent( &mut sessions, diff --git a/crates/cli/tests/coverage/shared/session_tests.rs b/crates/cli/tests/coverage/shared/session_tests.rs index 6c8f22a31..36e2e01f1 100644 --- a/crates/cli/tests/coverage/shared/session_tests.rs +++ b/crates/cli/tests/coverage/shared/session_tests.rs @@ -12,6 +12,7 @@ use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, regi use nemo_relay::codec::resolve::{ ProviderSurface, request_codec as build_request_codec, response_codec as build_response_codec, }; +use nemo_relay::error::FlowError; use nemo_relay::observability::OpenTelemetryType; use nemo_relay::observability::atof::{AtofExporter, AtofExporterConfig, AtofExporterMode}; use nemo_relay::observability::otel::OpenTelemetrySubscriber; @@ -1840,7 +1841,7 @@ async fn slow_tool_guardrail_does_not_block_another_session() { .unwrap(); let parallel = tokio::time::timeout( - Duration::from_millis(250), + Duration::from_secs(5), manager.apply_events( &HeaderMap::new(), vec![NormalizedEvent::Notification(codex_session_event( @@ -1860,6 +1861,170 @@ async fn slow_tool_guardrail_does_not_block_another_session() { .unwrap(); } +#[tokio::test] +async fn failed_authenticated_batch_releases_its_session_reservation() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + const GUARDRAIL: &str = "cli-session-reservation-release"; + const TOOL: &str = "session-reservation-release"; + let _ = deregister_tool_conditional_execution_guardrail(GUARDRAIL); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let started_tx = Arc::new(StdMutex::new(Some(started_tx))); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let release_rx = Arc::new(StdMutex::new(Some(release_rx))); + register_tool_conditional_execution_guardrail( + GUARDRAIL, + 1, + Arc::new(move |name, _| { + let started_tx = Arc::clone(&started_tx); + let release_rx = Arc::clone(&release_rx); + Box::pin(async move { + if name == TOOL { + if let Some(started) = started_tx.lock().unwrap().take() { + let _ = started.send(()); + } + let release = { release_rx.lock().unwrap().take() }; + if let Some(release) = release { + let _ = release.await; + } + return Err(FlowError::Internal("expected reservation failure".into())); + } + Ok(None) + }) + }), + ) + .unwrap(); + + let manager = SessionManager::new(session_test_config()); + let event = || { + NormalizedEvent::ToolStarted(ToolEvent { + session_id: "reserved-session".into(), + agent_kind: AgentKind::Codex, + event_name: "PreToolUse".into(), + tool_call_id: "reserved-tool".into(), + tool_name: TOOL.into(), + subagent_id: None, + arguments: json!({}), + result: Value::Null, + status: None, + payload: json!({}), + metadata: json!({}), + }) + }; + let pending_manager = manager.clone(); + let pending = tokio::spawn(async move { + pending_manager + .apply_authenticated_events(&HeaderMap::new(), vec![event()], "client-a") + .await + }); + tokio::time::timeout(Duration::from_secs(1), started_rx) + .await + .expect("guardrail should start") + .unwrap(); + + let competing = manager + .apply_authenticated_events( + &HeaderMap::new(), + vec![NormalizedEvent::AgentStarted(session_event( + "reserved-session", + "SessionStart", + ))], + "client-b", + ) + .await + .unwrap_err(); + assert!(matches!(competing, CliError::Unauthorized(_))); + + release_tx.send(()).unwrap(); + assert!(pending.await.unwrap().is_err()); + manager + .apply_authenticated_events( + &HeaderMap::new(), + vec![NormalizedEvent::AgentStarted(session_event( + "reserved-session", + "SessionStart", + ))], + "client-b", + ) + .await + .unwrap(); + deregister_tool_conditional_execution_guardrail(GUARDRAIL).unwrap(); +} + +#[tokio::test] +async fn close_all_waits_for_an_in_flight_session_application() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + const GUARDRAIL: &str = "cli-session-close-all-drain"; + const TOOL: &str = "session-close-all-drain"; + let _ = deregister_tool_conditional_execution_guardrail(GUARDRAIL); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let started_tx = Arc::new(StdMutex::new(Some(started_tx))); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let release_rx = Arc::new(StdMutex::new(Some(release_rx))); + register_tool_conditional_execution_guardrail( + GUARDRAIL, + 1, + Arc::new(move |name, _| { + let started_tx = Arc::clone(&started_tx); + let release_rx = Arc::clone(&release_rx); + Box::pin(async move { + if name == TOOL { + if let Some(started) = started_tx.lock().unwrap().take() { + let _ = started.send(()); + } + let release = { release_rx.lock().unwrap().take() }; + if let Some(release) = release { + let _ = release.await; + } + } + Ok(None) + }) + }), + ) + .unwrap(); + + let manager = SessionManager::new(session_test_config()); + let applying_manager = manager.clone(); + let applying = tokio::spawn(async move { + applying_manager + .apply_events( + &HeaderMap::new(), + vec![NormalizedEvent::ToolStarted(ToolEvent { + session_id: "close-all-session".into(), + agent_kind: AgentKind::Codex, + event_name: "PreToolUse".into(), + tool_call_id: "close-all-tool".into(), + tool_name: TOOL.into(), + subagent_id: None, + arguments: json!({}), + result: Value::Null, + status: None, + payload: json!({}), + metadata: json!({}), + })], + ) + .await + }); + tokio::time::timeout(Duration::from_secs(1), started_rx) + .await + .expect("guardrail should start") + .unwrap(); + + let closing_manager = manager.clone(); + let mut closing = tokio::spawn(async move { closing_manager.close_all("test_shutdown").await }); + assert!( + tokio::time::timeout(Duration::from_millis(100), &mut closing) + .await + .is_err(), + "close_all must wait until the in-flight event releases its activity" + ); + + release_tx.send(()).unwrap(); + applying.await.unwrap().unwrap(); + closing.await.unwrap().unwrap(); + assert!(manager.inner.lock().await.is_empty()); + deregister_tool_conditional_execution_guardrail(GUARDRAIL).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/codec/streaming.rs b/crates/core/src/codec/streaming.rs index d9740679c..0b8560792 100644 --- a/crates/core/src/codec/streaming.rs +++ b/crates/core/src/codec/streaming.rs @@ -149,7 +149,7 @@ impl SseEventDecoder { if self.buffer.ends_with('\r') && bytes.first() == Some(&b'\n') { self.buffer.pop(); // Removing the trailing CR means the preceding LF can now pair with the incoming LF. - self.scan_from = self.scan_from.min(self.buffer.len().saturating_sub(1)); + self.scan_from = self.buffer.len() - usize::from(self.buffer.ends_with('\n')); } let chunk = String::from_utf8_lossy(bytes); if chunk.contains("\r\n") { @@ -183,7 +183,7 @@ impl SseEventDecoder { if results.last().is_none_or(Result::is_ok) { // The next terminator can only start at the prior final byte; do // not rescan an incomplete frame from its beginning on every chunk. - self.scan_from = self.buffer.len().saturating_sub(1); + self.scan_from = self.buffer.len() - usize::from(self.buffer.ends_with('\n')); } results } diff --git a/crates/core/tests/unit/codec/streaming_tests.rs b/crates/core/tests/unit/codec/streaming_tests.rs index 9b0718be7..061b84b1a 100644 --- a/crates/core/tests/unit/codec/streaming_tests.rs +++ b/crates/core/tests/unit/codec/streaming_tests.rs @@ -127,3 +127,20 @@ fn resumes_with_frames_after_a_malformed_frame() { assert_eq!(resumed.len(), 1); assert_eq!(resumed[0].as_ref().unwrap().data, json!({"chunk": "later"})); } + +#[test] +fn resumes_after_a_multibyte_character_in_an_incomplete_frame() { + let mut decoder = SseEventDecoder::new(); + assert!( + decoder + .push_bytes_results("data: {\"chunk\":\"café".as_bytes()) + .is_empty() + ); + + let completed = decoder.push_bytes_results(b"\"}\n\n"); + assert_eq!(completed.len(), 1); + assert_eq!( + completed[0].as_ref().unwrap().data, + json!({"chunk": "café"}) + ); +} From 2ec881b34aef6f3fb578c74f5612938a48639dad Mon Sep 17 00:00:00 2001 From: Will Killian Date: Fri, 11 Sep 2026 18:35:36 -0400 Subject: [PATCH 3/4] fix: preserve session ownership during cleanup Signed-off-by: Will Killian --- crates/cli/src/sessions/idle.rs | 94 ++++++++----------- crates/cli/src/sessions/mod.rs | 26 +++++ .../tests/coverage/shared/session_tests.rs | 83 +++++++++++++++- 3 files changed, 145 insertions(+), 58 deletions(-) diff --git a/crates/cli/src/sessions/idle.rs b/crates/cli/src/sessions/idle.rs index 3df885406..542c4c0f6 100644 --- a/crates/cli/src/sessions/idle.rs +++ b/crates/cli/src/sessions/idle.rs @@ -3,7 +3,7 @@ //! Idle-session sweeping and shutdown closure. -use std::collections::{HashMap, HashSet, hash_map::Entry}; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -42,26 +42,16 @@ pub(super) async fn close_idle_sessions_from_parts( timeout: Duration, reason: &str, ) -> Result { - let idle_sessions = take_idle_sessions(inner, now, timeout).await; - if idle_sessions.is_empty() { + let candidate_ids = idle_session_ids(inner, now, timeout).await; + if candidate_ids.is_empty() { return Ok(0); } - let idle_session_ids = idle_sessions + let (closed_turns, closed_subagents, released_owner_ids, first_error) = + close_idle_turns(inner, session_gates, candidate_ids, now, timeout, reason).await; + let cleanup_sessions = closed_subagents .iter() .map(|(session_id, _)| session_id.clone()) - .collect::>(); - let (closed_turns, closed_subagents, retained_sessions, first_error) = - close_idle_turns(idle_sessions, reason).await; - let retained_session_ids = retained_sessions - .iter() - .map(|(session_id, _)| session_id.clone()) - .collect::>(); - let released_owner_ids = idle_session_ids - .difference(&retained_session_ids) - .cloned() - .collect::>(); - let cleanup_sessions = - restore_retained_sessions(inner, retained_sessions, &closed_subagents).await; + .collect(); clear_closed_subagents(alignment, closed_subagents, &cleanup_sessions).await; release_closed_owner_ids( inner, @@ -90,42 +80,55 @@ pub(super) async fn release_closed_owner_ids( } } -async fn take_idle_sessions( +async fn idle_session_ids( inner: &Arc>>, now: Instant, timeout: Duration, -) -> Vec<(String, Session)> { - let mut sessions = inner.lock().await; - let ids = sessions +) -> Vec { + inner + .lock() + .await .iter() .filter_map(|(session_id, session)| { session .is_idle_for(now, timeout) .then_some(session_id.clone()) }) - .collect::>(); - ids.into_iter() - .filter_map(|session_id| { - sessions - .remove(&session_id) - .map(|session| (session_id, session)) - }) .collect() } type ClosedIdleTurns = ( usize, Vec<(String, String)>, - Vec<(String, Session)>, + HashSet, Option, ); -async fn close_idle_turns(idle_sessions: Vec<(String, Session)>, reason: &str) -> ClosedIdleTurns { +async fn close_idle_turns( + inner: &Arc>>, + session_gates: &SessionGates, + candidate_ids: Vec, + now: Instant, + timeout: Duration, + reason: &str, +) -> ClosedIdleTurns { let mut closed_turns = 0; let mut closed_subagents = Vec::new(); - let mut retained_sessions = Vec::new(); + let mut released_owner_ids = HashSet::new(); let mut first_error = None; - for (session_id, mut session) in idle_sessions { + for session_id in candidate_ids { + let gate = session_gate(session_gates, &session_id).await; + let _gate = gate.lock().await; + let Some(mut session) = ({ + let mut sessions = inner.lock().await; + sessions + .get(&session_id) + .is_some_and(|session| session.is_idle_for(now, timeout)) + .then(|| sessions.remove(&session_id)) + .flatten() + }) else { + continue; + }; let stack = session.scope_stack.clone(); match TASK_SCOPE_STACK .scope(stack, async { @@ -151,38 +154,19 @@ async fn close_idle_turns(idle_sessions: Vec<(String, Session)>, reason: &str) - Err(_) => {} } if !session.is_empty() { - retained_sessions.push((session_id, session)); + inner.lock().await.insert(session_id, session); + } else { + released_owner_ids.insert(session_id); } } ( closed_turns, closed_subagents, - retained_sessions, + released_owner_ids, first_error, ) } -async fn restore_retained_sessions( - inner: &Arc>>, - retained_sessions: Vec<(String, Session)>, - closed_subagents: &[(String, String)], -) -> HashSet { - let mut cleanup_sessions = HashSet::new(); - let mut sessions = inner.lock().await; - for (session_id, session) in retained_sessions { - if let Entry::Vacant(entry) = sessions.entry(session_id.clone()) { - entry.insert(session); - cleanup_sessions.insert(session_id); - } - } - for (session_id, _) in closed_subagents { - if !sessions.contains_key(session_id) { - cleanup_sessions.insert(session_id.clone()); - } - } - cleanup_sessions -} - async fn clear_closed_subagents( alignment: &Arc>, closed_subagents: Vec<(String, String)>, diff --git a/crates/cli/src/sessions/mod.rs b/crates/cli/src/sessions/mod.rs index 35a654669..2dbd95f76 100644 --- a/crates/cli/src/sessions/mod.rs +++ b/crates/cli/src/sessions/mod.rs @@ -518,12 +518,38 @@ impl SessionManager { Ok(()) } Err(error) => { + drop(reservations); + drop(owners); + self.bind_retained_session_owners(&reserved_ids, owner) + .await; + let mut reservations = self.authenticated_reservations.lock().await; release_reservations(&mut reservations, &reserved_ids, owner); Err(error) } } } + // A hook batch can have applied an earlier event before a later event fails. The session is + // retained in that case, so bind the reserving client before releasing its reservation. The + // session gate makes the directory check and owner insertion one atomic lifecycle decision. + async fn bind_retained_session_owners(&self, session_ids: &HashSet, owner: &str) { + for session_id in session_ids { + let gate = session_gate(&self.session_gates, session_id).await; + let _gate = gate.lock().await; + let sessions = self.inner.lock().await; + if sessions + .get(session_id) + .is_some_and(|session| !session.is_empty()) + { + self.authenticated_owners + .lock() + .await + .entry(session_id.clone()) + .or_insert_with(|| owner.to_string()); + } + } + } + /// Evaluates a final host permission request inside its existing session scope. pub(crate) async fn authorize_tool_permission( &self, diff --git a/crates/cli/tests/coverage/shared/session_tests.rs b/crates/cli/tests/coverage/shared/session_tests.rs index 36e2e01f1..19ca28c44 100644 --- a/crates/cli/tests/coverage/shared/session_tests.rs +++ b/crates/cli/tests/coverage/shared/session_tests.rs @@ -146,6 +146,20 @@ async fn authenticated_child_cannot_promote_into_another_clients_parent() { .subagents .contains_key("child-thread") ); + + // The rejected child never created session state, so its reservation is released for a + // future independent session with the same client-provided ID. + manager + .apply_authenticated_events( + &HeaderMap::new(), + vec![NormalizedEvent::AgentStarted(session_event( + "child-thread", + "SessionStart", + ))], + "client-c", + ) + .await + .unwrap(); } #[tokio::test] @@ -1862,7 +1876,7 @@ async fn slow_tool_guardrail_does_not_block_another_session() { } #[tokio::test] -async fn failed_authenticated_batch_releases_its_session_reservation() { +async fn failed_guardrail_batch_keeps_its_partial_session_owner() { let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; const GUARDRAIL: &str = "cli-session-reservation-release"; const TOOL: &str = "session-reservation-release"; @@ -1936,7 +1950,7 @@ async fn failed_authenticated_batch_releases_its_session_reservation() { release_tx.send(()).unwrap(); assert!(pending.await.unwrap().is_err()); - manager + let competing = manager .apply_authenticated_events( &HeaderMap::new(), vec![NormalizedEvent::AgentStarted(session_event( @@ -1946,7 +1960,70 @@ async fn failed_authenticated_batch_releases_its_session_reservation() { "client-b", ) .await - .unwrap(); + .unwrap_err(); + assert!(matches!(competing, CliError::Unauthorized(_))); + deregister_tool_conditional_execution_guardrail(GUARDRAIL).unwrap(); +} + +#[tokio::test] +async fn partially_applied_authenticated_batch_keeps_its_owner() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + const GUARDRAIL: &str = "cli-session-partial-batch-owner"; + const TOOL: &str = "session-partial-batch-owner"; + let _ = deregister_tool_conditional_execution_guardrail(GUARDRAIL); + register_tool_conditional_execution_guardrail( + GUARDRAIL, + 1, + Arc::new(|name, _| { + Box::pin(async move { + (name == TOOL) + .then(|| FlowError::Internal("expected partial batch failure".into())) + .map_or(Ok(None), Err) + }) + }), + ) + .unwrap(); + + let manager = SessionManager::new(session_test_config()); + let result = manager + .apply_authenticated_events( + &HeaderMap::new(), + vec![ + NormalizedEvent::AgentStarted(session_event( + "partially-applied-session", + "SessionStart", + )), + NormalizedEvent::ToolStarted(ToolEvent { + session_id: "partially-applied-session".into(), + agent_kind: AgentKind::Codex, + event_name: "PreToolUse".into(), + tool_call_id: "partially-applied-tool".into(), + tool_name: TOOL.into(), + subagent_id: None, + arguments: json!({}), + result: Value::Null, + status: None, + payload: json!({}), + metadata: json!({}), + }), + ], + "client-a", + ) + .await; + assert!(result.is_err()); + + let competing = manager + .apply_authenticated_events( + &HeaderMap::new(), + vec![NormalizedEvent::AgentStarted(session_event( + "partially-applied-session", + "SessionStart", + ))], + "client-b", + ) + .await + .unwrap_err(); + assert!(matches!(competing, CliError::Unauthorized(_))); deregister_tool_conditional_execution_guardrail(GUARDRAIL).unwrap(); } From beaedad4534fdf8ab63f7ce1a231ee807d7f7aa6 Mon Sep 17 00:00:00 2001 From: Will Killian Date: Fri, 11 Sep 2026 18:48:59 -0400 Subject: [PATCH 4/4] test: assert retained session ownership Signed-off-by: Will Killian --- .../tests/coverage/shared/session_tests.rs | 43 ++++++++----------- 1 file changed, 18 insertions(+), 25 deletions(-) diff --git a/crates/cli/tests/coverage/shared/session_tests.rs b/crates/cli/tests/coverage/shared/session_tests.rs index 19ca28c44..17f37e9e2 100644 --- a/crates/cli/tests/coverage/shared/session_tests.rs +++ b/crates/cli/tests/coverage/shared/session_tests.rs @@ -1950,18 +1950,15 @@ async fn failed_guardrail_batch_keeps_its_partial_session_owner() { release_tx.send(()).unwrap(); assert!(pending.await.unwrap().is_err()); - let competing = manager - .apply_authenticated_events( - &HeaderMap::new(), - vec![NormalizedEvent::AgentStarted(session_event( - "reserved-session", - "SessionStart", - ))], - "client-b", - ) - .await - .unwrap_err(); - assert!(matches!(competing, CliError::Unauthorized(_))); + assert_eq!( + manager + .authenticated_owners + .lock() + .await + .get("reserved-session") + .map(String::as_str), + Some("client-a") + ); deregister_tool_conditional_execution_guardrail(GUARDRAIL).unwrap(); } @@ -2011,19 +2008,15 @@ async fn partially_applied_authenticated_batch_keeps_its_owner() { ) .await; assert!(result.is_err()); - - let competing = manager - .apply_authenticated_events( - &HeaderMap::new(), - vec![NormalizedEvent::AgentStarted(session_event( - "partially-applied-session", - "SessionStart", - ))], - "client-b", - ) - .await - .unwrap_err(); - assert!(matches!(competing, CliError::Unauthorized(_))); + assert_eq!( + manager + .authenticated_owners + .lock() + .await + .get("partially-applied-session") + .map(String::as_str), + Some("client-a") + ); deregister_tool_conditional_execution_guardrail(GUARDRAIL).unwrap(); }