diff --git a/crates/agentos-sidecar-core/src/engine.rs b/crates/agentos-sidecar-core/src/engine.rs index 508cad50e1..2388102aa3 100644 --- a/crates/agentos-sidecar-core/src/engine.rs +++ b/crates/agentos-sidecar-core/src/engine.rs @@ -14,8 +14,8 @@ use std::collections::BTreeMap; use agentos_protocol::generated::v1::{ AcpCloseSessionRequest, AcpCreateSessionRequest, AcpDeliverAgentOutputRequest, - AcpGetSessionStateRequest, AcpPendingResponse, AcpRequest, AcpResponse, AcpRuntimeKind, - AcpResumeSessionRequest, AcpSessionClosedResponse, AcpSessionRequest, + AcpGetSessionStateRequest, AcpPendingResponse, AcpRequest, AcpResponse, + AcpResumeSessionRequest, AcpRuntimeKind, AcpSessionClosedResponse, AcpSessionRequest, AcpSessionResumedResponse, AcpSessionRpcResponse, }; use serde_json::{json, Map, Value}; diff --git a/crates/agentos-sidecar/src/acp/mod.rs b/crates/agentos-sidecar/src/acp/mod.rs index ed30ebe3e5..3d3c343ecc 100644 --- a/crates/agentos-sidecar/src/acp/mod.rs +++ b/crates/agentos-sidecar/src/acp/mod.rs @@ -16,8 +16,8 @@ use agentos_native_sidecar::wire::{ StreamChannel, WriteStdinRequest, }; use agentos_native_sidecar::{ - Extension, ExtensionContext, ExtensionFuture, ExtensionInterruptRequest, - ExtensionInterruptResponse, ExtensionResponse, SidecarError, + Extension, ExtensionContext, ExtensionFuture, ExtensionOrderingPolicy, ExtensionRequestClass, + ExtensionResponse, SidecarError, }; use agentos_protocol::generated::v1::*; use agentos_protocol::ACP_EXTENSION_NAMESPACE; @@ -160,7 +160,7 @@ impl InactivityWarnings { pub struct AcpExtension { next_process_id: AtomicUsize, sessions: Mutex>, - prompt_cancellations: StdMutex>>, + routes: StdMutex>>, pending_permission_responses: Arc>>, } @@ -177,6 +177,439 @@ enum PendingPermissionSignal { Terminal(String), } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum AcpRoutePhase { + Idle, + StartingOrRestoring, + AdapterRpcRunning, + PromptRunning, + Stopping, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct AcpRouteSnapshot { + phase: AcpRoutePhase, + prompt_active: bool, +} + +impl AcpRouteSnapshot { + const IDLE: Self = Self { + phase: AcpRoutePhase::Idle, + prompt_active: false, + }; +} + +#[derive(Debug)] +enum AcpRouteState { + Idle, + StartingOrRestoring, + AdapterRpcRunning { + operation: String, + }, + PromptRunning { + cancellation: tokio::sync::watch::Sender, + }, + Stopping { + prompt_active: bool, + cancellation: Option>, + teardown_active: bool, + }, +} + +impl AcpRouteState { + fn snapshot(&self) -> AcpRouteSnapshot { + match self { + Self::Idle => AcpRouteSnapshot::IDLE, + Self::StartingOrRestoring => AcpRouteSnapshot { + phase: AcpRoutePhase::StartingOrRestoring, + prompt_active: false, + }, + Self::AdapterRpcRunning { .. } => AcpRouteSnapshot { + phase: AcpRoutePhase::AdapterRpcRunning, + prompt_active: false, + }, + Self::PromptRunning { .. } => AcpRouteSnapshot { + phase: AcpRoutePhase::PromptRunning, + prompt_active: true, + }, + Self::Stopping { prompt_active, .. } => AcpRouteSnapshot { + phase: AcpRoutePhase::Stopping, + prompt_active: *prompt_active, + }, + } + } +} + +#[derive(Debug)] +struct AcpRouteEntry { + owner_connection_id: String, + state: StdMutex, + changes: tokio::sync::watch::Sender, +} + +impl AcpRouteEntry { + fn new(owner_connection_id: String) -> Self { + let (changes, _receiver) = tokio::sync::watch::channel(AcpRouteSnapshot::IDLE); + Self { + owner_connection_id, + state: StdMutex::new(AcpRouteState::Idle), + changes, + } + } + + fn publish(&self, state: &AcpRouteState) { + self.changes.send_replace(state.snapshot()); + } + + fn begin_start(self: &Arc) -> Result { + let mut state = self.lock_state()?; + match &*state { + AcpRouteState::Idle => { + *state = AcpRouteState::StartingOrRestoring; + self.publish(&state); + Ok(AcpRouteStart::Leader(AcpRouteStartGuard { + route: Arc::clone(self), + armed: true, + })) + } + AcpRouteState::StartingOrRestoring => Ok(AcpRouteStart::Wait(self.changes.subscribe())), + AcpRouteState::PromptRunning { .. } => Err(route_busy_error( + "cannot start or restore an adapter while a prompt is running", + )), + AcpRouteState::AdapterRpcRunning { operation } => Err(route_busy_error(&format!( + "cannot start or restore an adapter while {operation} is running" + ))), + AcpRouteState::Stopping { .. } => Err(route_busy_error( + "cannot start or restore an adapter while the session is stopping", + )), + } + } + + fn begin_prompt( + self: &Arc, + ) -> Result<(AcpPromptRouteGuard, tokio::sync::watch::Receiver), SidecarError> { + let (cancellation, receiver) = tokio::sync::watch::channel(false); + let mut state = self.lock_state()?; + match &*state { + AcpRouteState::Idle => { + *state = AcpRouteState::PromptRunning { cancellation }; + self.publish(&state); + Ok(( + AcpPromptRouteGuard { + route: Arc::clone(self), + armed: true, + }, + receiver, + )) + } + AcpRouteState::StartingOrRestoring => Err(route_busy_error( + "cannot prompt while the adapter is starting or restoring", + )), + AcpRouteState::PromptRunning { .. } => { + Err(route_busy_error("session already has an active prompt")) + } + AcpRouteState::AdapterRpcRunning { operation } => Err(route_busy_error(&format!( + "cannot prompt while {operation} is running" + ))), + AcpRouteState::Stopping { .. } => Err(route_busy_error( + "cannot prompt while the session is stopping", + )), + } + } + + fn begin_stopping(self: &Arc) -> Result { + let mut state = self.lock_state()?; + match &*state { + AcpRouteState::Idle => { + *state = AcpRouteState::Stopping { + prompt_active: false, + cancellation: None, + teardown_active: true, + }; + } + AcpRouteState::PromptRunning { cancellation } => { + let cancellation = cancellation.clone(); + let _ = cancellation.send(true); + *state = AcpRouteState::Stopping { + prompt_active: true, + cancellation: Some(cancellation), + teardown_active: true, + }; + } + AcpRouteState::StartingOrRestoring => { + return Err(route_busy_error( + "cannot stop while the adapter is starting or restoring", + )); + } + AcpRouteState::AdapterRpcRunning { operation } => { + return Err(route_busy_error(&format!( + "cannot stop while {operation} is running" + ))); + } + AcpRouteState::Stopping { + teardown_active: true, + .. + } => { + return Err(route_busy_error("session teardown is already active")); + } + AcpRouteState::Stopping { + cancellation, + teardown_active: false, + .. + } => { + if let Some(cancellation) = cancellation { + let _ = cancellation.send(true); + } + let AcpRouteState::Stopping { + teardown_active, .. + } = &mut *state + else { + unreachable!("matched stopping route") + }; + *teardown_active = true; + } + } + self.publish(&state); + Ok(AcpRouteStopGuard { + route: Arc::clone(self), + armed: true, + }) + } + + fn signal_prompt_cancellation(&self) -> Result { + let state = self.lock_state()?; + let cancellation = match &*state { + AcpRouteState::PromptRunning { cancellation } => Some(cancellation), + AcpRouteState::Stopping { + cancellation: Some(cancellation), + .. + } => Some(cancellation), + AcpRouteState::Idle + | AcpRouteState::StartingOrRestoring + | AcpRouteState::AdapterRpcRunning { .. } + | AcpRouteState::Stopping { + cancellation: None, .. + } => None, + }; + Ok(cancellation.is_some_and(|sender| sender.send(true).is_ok())) + } + + fn begin_adapter_rpc( + self: &Arc, + operation: &str, + ) -> Result { + let mut state = self.lock_state()?; + match &*state { + AcpRouteState::Idle => { + *state = AcpRouteState::AdapterRpcRunning { + operation: operation.to_owned(), + }; + self.publish(&state); + Ok(AcpRouteAdapterRpcGuard { + route: Arc::clone(self), + armed: true, + }) + } + AcpRouteState::StartingOrRestoring => Err(route_busy_error(&format!( + "cannot {operation} while the adapter is starting or restoring" + ))), + AcpRouteState::AdapterRpcRunning { operation: active } => Err(route_busy_error( + &format!("cannot {operation} while {active} is running"), + )), + AcpRouteState::PromptRunning { .. } => Err(route_busy_error(&format!( + "cannot {operation} while a prompt is running" + ))), + AcpRouteState::Stopping { .. } => Err(route_busy_error(&format!( + "cannot {operation} while the session is stopping" + ))), + } + } + + async fn wait_until_not_starting( + mut receiver: tokio::sync::watch::Receiver, + ) -> Result<(), SidecarError> { + while receiver.borrow().phase == AcpRoutePhase::StartingOrRestoring { + receiver.changed().await.map_err(|_| { + SidecarError::InvalidState(String::from( + "ACP route startup waiter closed before startup completed", + )) + })?; + } + Ok(()) + } + + async fn wait_for_prompt_completion(&self) -> Result<(), SidecarError> { + let mut receiver = self.changes.subscribe(); + while receiver.borrow().prompt_active { + receiver.changed().await.map_err(|_| { + SidecarError::InvalidState(String::from( + "ACP prompt completion waiter closed before the prompt completed", + )) + })?; + } + Ok(()) + } + + fn finish_start(&self) { + let Ok(mut state) = self.state.lock() else { + eprintln!("ERR_AGENTOS_ACP_ROUTE_STATE: startup route lock is poisoned"); + return; + }; + if matches!(*state, AcpRouteState::StartingOrRestoring) { + *state = AcpRouteState::Idle; + self.publish(&state); + } + } + + fn finish_prompt(&self) { + let Ok(mut state) = self.state.lock() else { + eprintln!("ERR_AGENTOS_ACP_ROUTE_STATE: prompt route lock is poisoned"); + return; + }; + match &*state { + AcpRouteState::PromptRunning { .. } => { + *state = AcpRouteState::Idle; + } + AcpRouteState::Stopping { .. } => { + let teardown_active = match &*state { + AcpRouteState::Stopping { + teardown_active, .. + } => *teardown_active, + _ => unreachable!("matched stopping route"), + }; + *state = AcpRouteState::Stopping { + prompt_active: false, + cancellation: None, + teardown_active, + }; + } + AcpRouteState::Idle + | AcpRouteState::StartingOrRestoring + | AcpRouteState::AdapterRpcRunning { .. } => return, + } + self.publish(&state); + } + + fn complete_stopping(&self) { + let Ok(mut state) = self.state.lock() else { + eprintln!("ERR_AGENTOS_ACP_ROUTE_STATE: stopping route lock is poisoned"); + return; + }; + if matches!(*state, AcpRouteState::Stopping { .. }) { + *state = AcpRouteState::Idle; + self.publish(&state); + } + } + + fn release_stopping_attempt(&self) { + let Ok(mut state) = self.state.lock() else { + eprintln!("ERR_AGENTOS_ACP_ROUTE_STATE: stopping route lock is poisoned"); + return; + }; + if let AcpRouteState::Stopping { + teardown_active, .. + } = &mut *state + { + *teardown_active = false; + self.publish(&state); + } + } + + fn finish_adapter_rpc(&self) { + let Ok(mut state) = self.state.lock() else { + eprintln!("ERR_AGENTOS_ACP_ROUTE_STATE: adapter RPC route lock is poisoned"); + return; + }; + if matches!(*state, AcpRouteState::AdapterRpcRunning { .. }) { + *state = AcpRouteState::Idle; + self.publish(&state); + } + } + + fn lock_state(&self) -> Result, SidecarError> { + self.state.lock().map_err(|_| { + SidecarError::InvalidState(String::from("ACP route state lock is poisoned")) + }) + } +} + +#[derive(Debug)] +enum AcpRouteStart { + Leader(AcpRouteStartGuard), + Wait(tokio::sync::watch::Receiver), +} + +#[derive(Debug)] +struct AcpRouteStartGuard { + route: Arc, + armed: bool, +} + +impl Drop for AcpRouteStartGuard { + fn drop(&mut self) { + if self.armed { + self.route.finish_start(); + } + } +} + +#[derive(Debug)] +struct AcpPromptRouteGuard { + route: Arc, + armed: bool, +} + +impl Drop for AcpPromptRouteGuard { + fn drop(&mut self) { + if self.armed { + self.route.finish_prompt(); + } + } +} + +#[derive(Debug)] +struct AcpRouteStopGuard { + route: Arc, + armed: bool, +} + +#[derive(Debug)] +struct AcpRouteAdapterRpcGuard { + route: Arc, + armed: bool, +} + +impl Drop for AcpRouteAdapterRpcGuard { + fn drop(&mut self) { + if self.armed { + self.route.finish_adapter_rpc(); + } + } +} + +impl AcpRouteStopGuard { + async fn wait_for_prompt_completion(&self) -> Result<(), SidecarError> { + self.route.wait_for_prompt_completion().await + } + + fn complete(mut self) { + self.route.complete_stopping(); + self.armed = false; + } +} + +impl Drop for AcpRouteStopGuard { + fn drop(&mut self) { + if self.armed { + self.route.release_stopping_attempt(); + } + } +} + +fn route_busy_error(message: &str) -> SidecarError { + SidecarError::InvalidState(format!("session_busy: {message}")) +} + #[derive(Debug, Clone)] struct LiveAcpRuntime { acp_session_id: String, @@ -209,9 +642,138 @@ impl AcpExtension { Self::default() } + fn route_entry( + &self, + ownership: &OwnershipScope, + route_key: &str, + ) -> Result, SidecarError> { + let owner_connection_id = ownership_connection_id(ownership); + let mut routes = self.routes.lock().map_err(|_| { + SidecarError::InvalidState(String::from("ACP route registry is poisoned")) + })?; + if let Some(route) = routes.get(route_key) { + if route.owner_connection_id != owner_connection_id { + return Err(SidecarError::InvalidState(format!( + "unknown ACP session {route_key}" + ))); + } + return Ok(Arc::clone(route)); + } + let route = Arc::new(AcpRouteEntry::new(owner_connection_id)); + routes.insert(route_key.to_owned(), Arc::clone(&route)); + Ok(route) + } + + fn begin_route_start( + &self, + ownership: &OwnershipScope, + route_key: &str, + ) -> Result { + self.route_entry(ownership, route_key)?.begin_start() + } + + fn begin_route_prompt( + &self, + ownership: &OwnershipScope, + route_key: &str, + ) -> Result<(AcpPromptRouteGuard, tokio::sync::watch::Receiver), SidecarError> { + self.route_entry(ownership, route_key)?.begin_prompt() + } + + fn begin_route_stopping( + &self, + ownership: &OwnershipScope, + route_key: &str, + ) -> Result { + self.route_entry(ownership, route_key)?.begin_stopping() + } + + fn begin_route_adapter_rpc( + &self, + ownership: &OwnershipScope, + route_key: &str, + operation: &str, + ) -> Result { + self.route_entry(ownership, route_key)? + .begin_adapter_rpc(operation) + } + + async fn drain_prompt_for_teardown( + &self, + ctx: &mut ExtensionContext, + route_key: &str, + stop_guard: &AcpRouteStopGuard, + ) -> Result { + let grace = ctx.acp_termination_grace().await?; + match tokio::time::timeout(grace, stop_guard.wait_for_prompt_completion()).await { + Ok(result) => { + result?; + return Ok(false); + } + Err(_) => tracing::warn!( + target: "agentos_sidecar::acp_extension", + route_key, + grace_ms = grace.as_millis() as u64, + "ACP prompt did not finish after cancellation; escalating adapter teardown", + ), + } + + self.force_kill_acp_runtime(ctx, route_key).await?; + match tokio::time::timeout(grace, stop_guard.wait_for_prompt_completion()).await { + Ok(result) => { + result?; + Ok(true) + } + Err(_) => Err(SidecarError::InvalidState(format!( + "acp_teardown_timeout: route {route_key} did not commit its terminal prompt state within {}ms after SIGKILL; raise NativeSidecarConfig.acp_termination_grace", + grace.as_millis() + ))), + } + } + + fn cancel_routes_for_connection(&self, connection_id: &str, reason: &str) { + let route_entries = match self.routes.lock() { + Ok(routes) => routes + .iter() + .filter(|(_, route)| route.owner_connection_id == connection_id) + .map(|(key, route)| (key.clone(), Arc::clone(route))) + .collect::>(), + Err(_) => { + eprintln!("ERR_AGENTOS_ACP_ROUTE_STATE: route registry is poisoned"); + return; + } + }; + for (route_key, route) in route_entries { + if let Err(error) = route.signal_prompt_cancellation() { + eprintln!( + "ERR_AGENTOS_PROMPT_CANCELLATION_REGISTRY: failed to signal {route_key}: {error}" + ); + } + self.cancel_pending_permissions(&route_key, reason); + } + } + + fn clear_routes_for_connection(&self, connection_id: &str) { + match self.routes.lock() { + Ok(mut routes) => { + routes.retain(|_, route| route.owner_connection_id != connection_id); + } + Err(_) => eprintln!("ERR_AGENTOS_ACP_ROUTE_STATE: route registry is poisoned"), + } + } + + fn remove_route(&self, route_key: &str) { + match self.routes.lock() { + Ok(mut routes) => { + routes.remove(route_key); + } + Err(_) => eprintln!("ERR_AGENTOS_ACP_ROUTE_STATE: route registry is poisoned"), + } + } + async fn handle_payload( &self, - mut ctx: ExtensionContext<'_>, + mut ctx: ExtensionContext, payload: &[u8], ) -> Result { use tracing::Instrument as _; @@ -324,7 +886,7 @@ impl AcpExtension { async fn session_store( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, ) -> Result { let limits = ctx.vm_acp_limits().await?; let database = ctx.vm_database().await?.ok_or_else(|| { @@ -337,7 +899,7 @@ impl AcpExtension { async fn open_session( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: AcpOpenSessionRequest, ) -> AcpHandlerOutput { let store = match self.session_store(ctx).await { @@ -351,6 +913,7 @@ impl AcpExtension { if let Err(error) = validate_user_session_id(&session_id) { return AcpHandlerOutput::response(Err(error)); } + let route_key = durable_route_key(ctx.ownership(), &session_id); let cwd = request .cwd .clone() @@ -440,7 +1003,42 @@ impl AcpExtension { skip_os_instructions: request.skip_os_instructions.unwrap_or(false), additional_instructions: request.additional_instructions.clone(), }; - let route_key = durable_route_key(ctx.ownership(), &session_id); + let _start_guard = loop { + match self.begin_route_start(ctx.ownership(), &route_key) { + Ok(AcpRouteStart::Leader(guard)) => break guard, + Ok(AcpRouteStart::Wait(receiver)) => { + if let Err(error) = AcpRouteEntry::wait_until_not_starting(receiver).await { + return AcpHandlerOutput::response(Err(error)); + } + match store.get(&session_id).await { + Ok(Some(existing)) => { + if existing.agent != request.agent + || existing.creation_options_json != creation_options + { + return AcpHandlerOutput::response(Err( + SidecarError::InvalidState(format!( + "session_conflict: session {session_id} already exists with different immutable creation options" + )), + )); + } + return match self.ensure_durable_runtime(ctx, &store, existing).await { + Ok(_) => AcpHandlerOutput::response(Ok( + AcpResponse::AcpOpenSessionResponse(AcpOpenSessionResponse { + reserved: false, + }), + )), + Err(error) => AcpHandlerOutput::response(Err(error)), + }; + } + Ok(None) => continue, + Err(error) => { + return AcpHandlerOutput::response(Err(session_store_error(error))); + } + } + } + Err(error) => return AcpHandlerOutput::response(Err(error)), + } + }; let created = self .start_acp_runtime(ctx, create, &session_id, &route_key, additional_directories) .await; @@ -485,7 +1083,7 @@ impl AcpExtension { async fn get_durable_session( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: AcpGetDurableSessionRequest, ) -> Result { let session_id = default_session_id(request.session_id)?; @@ -497,7 +1095,7 @@ impl AcpExtension { async fn list_durable_sessions( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: AcpListDurableSessionsRequest, ) -> Result { const DEFAULT_LIMIT: usize = 50; @@ -541,7 +1139,7 @@ impl AcpExtension { async fn delete_durable_session( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: AcpDeleteSessionRequest, ) -> AcpHandlerOutput { let session_id = match default_session_id(request.session_id) { @@ -557,6 +1155,17 @@ impl AcpExtension { Err(error) => return AcpHandlerOutput::response(Err(session_store_error(error))), }; let route_key = durable_route_key(ctx.ownership(), &session_id); + let stop_guard = match self.begin_route_stopping(ctx.ownership(), &route_key) { + Ok(guard) => guard, + Err(error) => return AcpHandlerOutput::response(Err(error)), + }; + self.cancel_pending_permissions(&route_key, "session_deleted"); + if let Err(error) = self + .drain_prompt_for_teardown(ctx, &route_key, &stop_guard) + .await + { + return AcpHandlerOutput::response(Err(error)); + } if self.sessions.lock().await.contains_key(&route_key) { if let Err(error) = self.stop_acp_runtime(ctx, &route_key).await { return AcpHandlerOutput::response(Err(SidecarError::InvalidState(format!( @@ -564,19 +1173,28 @@ impl AcpExtension { session_id )))); } + } else if let Err(error) = ctx.dispose_session_resources_wire(&route_key).await { + return AcpHandlerOutput::response(Err(SidecarError::InvalidState(format!( + "session_delete_cleanup_failed: session {} was retained because its resources could not be disposed: {error}", + session_id + )))); } let _ = stored; match store.delete(&session_id).await { - Ok(()) => AcpHandlerOutput::response(Ok(AcpResponse::AcpDeleteSessionResponse( - AcpDeleteSessionResponse { reserved: false }, - ))), + Ok(()) => { + self.remove_route(&route_key); + stop_guard.complete(); + AcpHandlerOutput::response(Ok(AcpResponse::AcpDeleteSessionResponse( + AcpDeleteSessionResponse { reserved: false }, + ))) + } Err(error) => AcpHandlerOutput::response(Err(session_store_error(error))), } } async fn unload_durable_session( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: AcpUnloadSessionRequest, ) -> AcpHandlerOutput { let session_id = match default_session_id(request.session_id) { @@ -592,11 +1210,28 @@ impl AcpExtension { Err(error) => return AcpHandlerOutput::response(Err(error)), }; let route_key = durable_route_key(ctx.ownership(), &session_id); + let stop_guard = match self.begin_route_stopping(ctx.ownership(), &route_key) { + Ok(guard) => guard, + Err(error) => return AcpHandlerOutput::response(Err(error)), + }; + self.cancel_pending_permissions(&route_key, "prompt_cancelled"); + if let Err(error) = self + .drain_prompt_for_teardown(ctx, &route_key, &stop_guard) + .await + { + return AcpHandlerOutput::response(Err(error)); + } if self.sessions.lock().await.contains_key(&route_key) { if let Err(error) = self.stop_acp_runtime(ctx, &route_key).await { return AcpHandlerOutput::response(Err(error)); } + } else if let Err(error) = ctx.dispose_session_resources_wire(&route_key).await { + return AcpHandlerOutput::response(Err(SidecarError::InvalidState(format!( + "session_unload_cleanup_failed: session {session_id} resources could not be disposed: {error}" + )))); } + self.remove_route(&route_key); + stop_guard.complete(); let _ = stored; AcpHandlerOutput::response(Ok(AcpResponse::AcpUnloadSessionResponse( AcpUnloadSessionResponse { reserved: false }, @@ -605,7 +1240,7 @@ impl AcpExtension { async fn read_history( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: AcpReadHistoryRequest, ) -> Result { const DEFAULT_LIMIT: usize = 100; @@ -670,7 +1305,7 @@ impl AcpExtension { async fn get_durable_config( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: AcpGetSessionConfigRequest, ) -> Result { let id = default_session_id(request.session_id)?; @@ -685,7 +1320,7 @@ impl AcpExtension { async fn get_durable_capabilities( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: AcpGetSessionCapabilitiesRequest, ) -> Result { let id = default_session_id(request.session_id)?; @@ -699,7 +1334,7 @@ impl AcpExtension { async fn get_durable_agent_info( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: AcpGetSessionAgentInfoRequest, ) -> Result { let id = default_session_id(request.session_id)?; @@ -716,9 +1351,25 @@ impl Extension for AcpExtension { ACP_EXTENSION_NAMESPACE } + fn request_ordering_key(&self, ownership: &OwnershipScope, payload: &[u8]) -> Option> { + let request = decode_request(payload).ok()?; + let session_id = durable_request_session_id(&request)?; + Some(durable_route_key(ownership, session_id).into_bytes()) + } + + fn request_ordering_policy( + &self, + _ownership: &OwnershipScope, + _payload: &[u8], + ) -> ExtensionOrderingPolicy { + // AcpRouteEntry performs the same bounded exclusion and must return + // ACP's established typed `session_busy` response. + ExtensionOrderingPolicy::ExtensionManaged + } + fn handle_request<'a>( &'a self, - ctx: ExtensionContext<'a>, + ctx: ExtensionContext, payload: Vec, ) -> ExtensionFuture<'a, ExtensionResponse> { Box::pin(async move { @@ -727,6 +1378,15 @@ impl Extension for AcpExtension { }) } + fn request_class(&self, payload: &[u8]) -> ExtensionRequestClass { + match decode_request(payload) { + Ok( + AcpRequest::AcpCancelPromptRequest(_) | AcpRequest::AcpRespondPermissionRequest(_), + ) => ExtensionRequestClass::Progress, + _ => ExtensionRequestClass::Ordinary, + } + } + fn bootstrap_vm_database<'a>( &'a self, database: agentos_native_sidecar::vm_sqlite::SharedVmSqliteDatabase, @@ -742,20 +1402,31 @@ impl Extension for AcpExtension { }) } - fn is_blocking_request(&self, payload: &[u8]) -> bool { - matches!( - decode_request(payload), - Ok(AcpRequest::AcpSessionRequest(request)) if request.method == "session/prompt" - ) || matches!(decode_request(payload), Ok(AcpRequest::AcpPromptRequest(_))) - } - fn on_dispose<'a>(&'a self) -> ExtensionFuture<'a, ()> { Box::pin(async move { + let connection_ids = match self.routes.lock() { + Ok(routes) => routes + .values() + .map(|route| route.owner_connection_id.clone()) + .collect::>(), + Err(_) => { + eprintln!("ERR_AGENTOS_ACP_ROUTE_STATE: route registry is poisoned"); + BTreeSet::new() + } + }; + for connection_id in connection_ids { + self.cancel_routes_for_connection(&connection_id, "vm_shutdown"); + } // Extension/sidecar teardown: drop every remaining session record so // no `stdout_buffer` survives the host process. The adapter processes // themselves are reaped by the host's own session/VM dispose; this // only frees the wrapper-side tracking map. self.sessions.lock().await.clear(); + if let Ok(mut routes) = self.routes.lock() { + routes.clear(); + } else { + eprintln!("ERR_AGENTOS_ACP_ROUTE_STATE: route registry is poisoned"); + } Ok(()) }) } @@ -769,156 +1440,12 @@ impl Extension for AcpExtension { // connection. This closes the disconnect path of H4 (the per-request // process-exit eviction and `on_dispose` cover the other paths). let connection_id = ownership_connection_id(ctx.ownership()); + self.cancel_routes_for_connection(&connection_id, "vm_shutdown"); self.cleanup_sessions_for_connection(&connection_id).await; + self.clear_routes_for_connection(&connection_id); Ok(()) }) } - - fn interrupt_blocking_request( - &self, - blocking_payload: &[u8], - interrupt: ExtensionInterruptRequest<'_>, - ) -> Option { - let blocking = decode_request(blocking_payload).ok()?; - if let AcpRequest::AcpPromptRequest(blocking_request) = blocking { - let user_session_id = blocking_request - .session_id - .unwrap_or_else(|| String::from("main")); - return match interrupt { - ExtensionInterruptRequest::KillProcess => Some(ExtensionInterruptResponse { - interrupt_active: true, - interrupted_response_payload: encode_durable_interrupted_prompt( - &user_session_id, - )?, - interrupting_response_payload: None, - }), - ExtensionInterruptRequest::ExtensionPayload { payload, ownership } => { - match decode_request(payload).ok()? { - AcpRequest::AcpCancelPromptRequest(cancel) => { - let cancel_session_id = - cancel.session_id.unwrap_or_else(|| String::from("main")); - if cancel_session_id != user_session_id { - return None; - } - let key = durable_route_key(ownership, &user_session_id); - let signalled = self.signal_prompt_cancellation(&key); - // Permission response and prompt cancellation share - // the same registry lock, so exactly one wins. A - // cancelled permission waiter receives an explicit - // ACP cancelled outcome immediately. - self.cancel_pending_permissions(&key, "prompt_cancelled"); - Some(ExtensionInterruptResponse { - interrupt_active: false, - interrupted_response_payload: encode_durable_interrupted_prompt( - &user_session_id, - )?, - interrupting_response_payload: Some( - encode_durable_cancel_response(signalled)?, - ), - }) - } - AcpRequest::AcpRespondPermissionRequest(response) => { - let response_session_id = response.session_id; - if response_session_id != user_session_id { - return None; - } - let key = format!( - "{}:{}", - durable_route_key(ownership, &user_session_id), - response.request_id - ); - let selected_option_id = response.option_id.clone(); - // Validate before acknowledging the interrupt. An invalid - // option must not consume the live request and prevent the - // caller from correcting its response. - let (accepted, valid_options) = match self - .pending_permission_responses - .lock() - { - Ok(mut pending) => { - if let Some(entry) = pending.get(&key) { - if entry.offered_option_ids.contains(&response.option_id) { - (pending.remove(&key), None) - } else { - ( - None, - Some( - entry - .offered_option_ids - .iter() - .cloned() - .collect::>(), - ), - ) - } - } else { - (None, None) - } - } - Err(_) => { - eprintln!("ERR_AGENTOS_PERMISSION_RESPONSE: permission response registry is poisoned"); - (None, None) - } - }; - let accepted = accepted.is_some_and(|entry| { - entry - .sender - .send(PendingPermissionSignal::Selected(selected_option_id)) - .is_ok() - }); - Some(ExtensionInterruptResponse { - interrupt_active: false, - interrupted_response_payload: encode_durable_interrupted_prompt( - &user_session_id, - )?, - interrupting_response_payload: Some( - if let Some(valid_options) = valid_options { - encode_response(AcpResponse::AcpErrorResponse(AcpErrorResponse { - code: String::from("invalid_permission_option"), - message: format!( - "invalid_permission_option: request {} does not offer {}; valid option IDs: {}", - response.request_id, - response.option_id, - valid_options.join(", ") - ), - })).ok()? - } else { - encode_durable_permission_response(accepted)? - }, - ), - }) - } - AcpRequest::AcpUnloadSessionRequest(unload) - if unload.session_id.as_deref().unwrap_or("main") - == user_session_id => - { - let key = durable_route_key(ownership, &user_session_id); - self.cancel_pending_permissions(&key, "prompt_cancelled"); - self.signal_prompt_cancellation(&key); - // Returning None deliberately queues the unload behind - // the active prompt. The signal makes that prompt commit - // its terminal cancellation first; normal dispatch then - // tears down the adapter and answers unload. - None - } - AcpRequest::AcpDeleteSessionRequest(delete) - if delete.session_id.as_deref().unwrap_or("main") - == user_session_id => - { - let key = durable_route_key(ownership, &user_session_id); - self.cancel_pending_permissions(&key, "session_deleted"); - self.signal_prompt_cancellation(&key); - // As with unload, queue deletion until the prompt has - // durably recorded its terminal cancellation. - None - } - _ => None, - } - } - }; - } - None - } } struct AcpHandlerOutput { @@ -976,17 +1503,68 @@ fn durable_route_key(ownership: &OwnershipScope, session_id: &str) -> String { key } -fn session_store_error(error: agentos_native_sidecar::vm_sqlite::VmSqliteError) -> SidecarError { - match error { - error @ (agentos_native_sidecar::vm_sqlite::VmSqliteError::ResultTooLarge { .. } - | agentos_native_sidecar::vm_sqlite::VmSqliteError::HistoryEventBatchTooLarge { - .. +/// Return the public durable session identity targeted by a native ACP +/// request. Global queries and legacy browser-only messages have no durable +/// session conflict domain. Validation remains in the request handler; this +/// ingress classification hook is deliberately total and side-effect free. +fn durable_request_session_id(request: &AcpRequest) -> Option<&str> { + match request { + AcpRequest::AcpOpenSessionRequest(request) => { + Some(request.session_id.as_deref().unwrap_or("main")) } - | agentos_native_sidecar::vm_sqlite::VmSqliteError::HistoryByteBatchTooLarge { - .. + AcpRequest::AcpGetDurableSessionRequest(request) => { + Some(request.session_id.as_deref().unwrap_or("main")) } - | agentos_native_sidecar::vm_sqlite::VmSqliteError::DurableCollectionLimit { - .. + AcpRequest::AcpDeleteSessionRequest(request) => { + Some(request.session_id.as_deref().unwrap_or("main")) + } + AcpRequest::AcpUnloadSessionRequest(request) => { + Some(request.session_id.as_deref().unwrap_or("main")) + } + AcpRequest::AcpPromptRequest(request) => { + Some(request.session_id.as_deref().unwrap_or("main")) + } + AcpRequest::AcpCancelPromptRequest(request) => { + Some(request.session_id.as_deref().unwrap_or("main")) + } + AcpRequest::AcpRespondPermissionRequest(request) => Some(&request.session_id), + AcpRequest::AcpReadHistoryRequest(request) => { + Some(request.session_id.as_deref().unwrap_or("main")) + } + AcpRequest::AcpGetSessionConfigRequest(request) => { + Some(request.session_id.as_deref().unwrap_or("main")) + } + AcpRequest::AcpSetSessionConfigOptionRequest(request) => { + Some(request.session_id.as_deref().unwrap_or("main")) + } + AcpRequest::AcpGetSessionCapabilitiesRequest(request) => { + Some(request.session_id.as_deref().unwrap_or("main")) + } + AcpRequest::AcpGetSessionAgentInfoRequest(request) => { + Some(request.session_id.as_deref().unwrap_or("main")) + } + AcpRequest::AcpListDurableSessionsRequest(_) + | AcpRequest::AcpCreateSessionRequest(_) + | AcpRequest::AcpSessionRequest(_) + | AcpRequest::AcpGetSessionStateRequest(_) + | AcpRequest::AcpCloseSessionRequest(_) + | AcpRequest::AcpResumeSessionRequest(_) + | AcpRequest::AcpDeliverAgentOutputRequest(_) + | AcpRequest::AcpListAgentsRequest(_) => None, + } +} + +fn session_store_error(error: agentos_native_sidecar::vm_sqlite::VmSqliteError) -> SidecarError { + match error { + error @ (agentos_native_sidecar::vm_sqlite::VmSqliteError::ResultTooLarge { .. } + | agentos_native_sidecar::vm_sqlite::VmSqliteError::HistoryEventBatchTooLarge { + .. + } + | agentos_native_sidecar::vm_sqlite::VmSqliteError::HistoryByteBatchTooLarge { + .. + } + | agentos_native_sidecar::vm_sqlite::VmSqliteError::DurableCollectionLimit { + .. }) => SidecarError::InvalidState(error.to_string()), error => SidecarError::InvalidState(format!("session_storage_error: {error}")), } @@ -1227,6 +1805,1074 @@ mod tests { assert_eq!(AcpExtension::new().namespace(), ACP_EXTENSION_NAMESPACE); } + #[test] + fn cancel_and_permission_are_progress_classified_by_acp() { + let extension = AcpExtension::new(); + let cancel = serde_bare::to_vec(&AcpRequest::AcpCancelPromptRequest( + AcpCancelPromptRequest { + session_id: Some(String::from("main")), + }, + )) + .expect("encode cancel"); + let permission = serde_bare::to_vec(&AcpRequest::AcpRespondPermissionRequest( + AcpRespondPermissionRequest { + session_id: String::from("main"), + request_id: String::from("permission-1"), + option_id: String::from("allow"), + }, + )) + .expect("encode permission response"); + + assert_eq!( + extension.request_class(&cancel), + ExtensionRequestClass::Progress + ); + assert_eq!( + extension.request_class(&permission), + ExtensionRequestClass::Progress + ); + assert_eq!( + extension.request_class(b"not an ACP request"), + ExtensionRequestClass::Ordinary + ); + } + + #[test] + fn durable_acp_requests_expose_an_opaque_owned_session_ordering_key() { + use agentos_native_sidecar::wire::VmOwnership; + + let extension = AcpExtension::new(); + let ownership = OwnershipScope::VmOwnership(VmOwnership { + connection_id: String::from("conn-1"), + session_id: String::from("owner-session"), + vm_id: String::from("vm-1"), + }); + let other_vm = OwnershipScope::VmOwnership(VmOwnership { + connection_id: String::from("conn-1"), + session_id: String::from("owner-session"), + vm_id: String::from("vm-2"), + }); + let default_prompt = serde_bare::to_vec(&AcpRequest::AcpPromptRequest(AcpPromptRequest { + session_id: None, + idempotency_key: None, + content: String::from("[]"), + })) + .expect("encode default prompt"); + let explicit_main_cancel = serde_bare::to_vec(&AcpRequest::AcpCancelPromptRequest( + AcpCancelPromptRequest { + session_id: Some(String::from("main")), + }, + )) + .expect("encode explicit main cancellation"); + let other_prompt = serde_bare::to_vec(&AcpRequest::AcpPromptRequest(AcpPromptRequest { + session_id: Some(String::from("other")), + idempotency_key: None, + content: String::from("[]"), + })) + .expect("encode other prompt"); + + let main_key = durable_route_key(&ownership, "main").into_bytes(); + assert_eq!( + extension.request_ordering_key(&ownership, &default_prompt), + Some(main_key.clone()), + "omitted durable session IDs must order as main" + ); + assert_eq!( + extension.request_ordering_key(&ownership, &explicit_main_cancel), + Some(main_key), + "progress and ordinary messages for one durable session must expose the same target key" + ); + assert_ne!( + extension.request_ordering_key(&ownership, &default_prompt), + extension.request_ordering_key(&ownership, &other_prompt), + "different durable ACP sessions must remain independent" + ); + assert_ne!( + extension.request_ordering_key(&ownership, &default_prompt), + extension.request_ordering_key(&other_vm, &default_prompt), + "identically named durable sessions in different VMs must remain independent" + ); + assert_eq!( + extension.request_ordering_policy(&ownership, &default_prompt), + ExtensionOrderingPolicy::ExtensionManaged, + "ACP must retain its protocol-specific session_busy rejection" + ); + } + + #[test] + fn global_and_invalid_acp_requests_have_no_session_ordering_key() { + use agentos_native_sidecar::wire::ConnectionOwnership; + + let extension = AcpExtension::new(); + let ownership = OwnershipScope::ConnectionOwnership(ConnectionOwnership { + connection_id: String::from("conn-1"), + }); + let list = serde_bare::to_vec(&AcpRequest::AcpListDurableSessionsRequest( + AcpListDurableSessionsRequest { + cursor: None, + limit: None, + }, + )) + .expect("encode list request"); + + assert_eq!(extension.request_ordering_key(&ownership, &list), None); + assert_eq!( + extension.request_ordering_key(&ownership, b"not an ACP request"), + None + ); + } + + #[test] + fn route_start_is_single_flight_and_releases_waiters_without_sleeping() { + let route = Arc::new(AcpRouteEntry::new(String::from("conn-1"))); + let leader = match route.begin_start().expect("first startup is admitted") { + AcpRouteStart::Leader(guard) => guard, + AcpRouteStart::Wait(_) => panic!("first startup unexpectedly waited"), + }; + let waiter = match route.begin_start().expect("second startup waits") { + AcpRouteStart::Wait(receiver) => receiver, + AcpRouteStart::Leader(_) => panic!("second startup became another leader"), + }; + assert_eq!(waiter.borrow().phase, AcpRoutePhase::StartingOrRestoring); + + drop(leader); + poll_uncontended(AcpRouteEntry::wait_until_not_starting(waiter)) + .expect("startup waiter observes leader completion"); + + assert!(matches!( + route.begin_start().expect("route becomes reusable"), + AcpRouteStart::Leader(_) + )); + } + + #[tokio::test(flavor = "current_thread")] + async fn concurrent_open_and_restore_attempts_launch_one_adapter_each() { + use std::sync::atomic::AtomicBool; + + async fn attempt_start( + route: Arc, + live: Arc, + launches: Arc, + mut release: tokio::sync::watch::Receiver, + role: tokio::sync::mpsc::UnboundedSender<&'static str>, + ) { + loop { + if live.load(Ordering::Acquire) { + return; + } + match route.begin_start().expect("route startup") { + AcpRouteStart::Leader(guard) => { + launches.fetch_add(1, Ordering::AcqRel); + role.send("leader").expect("record leader"); + while !*release.borrow() { + release.changed().await.expect("release startup leader"); + } + live.store(true, Ordering::Release); + drop(guard); + return; + } + AcpRouteStart::Wait(waiter) => { + role.send("waiter").expect("record waiter"); + AcpRouteEntry::wait_until_not_starting(waiter) + .await + .expect("startup leader completes"); + } + } + } + } + + for operation in ["open", "restore"] { + let route = Arc::new(AcpRouteEntry::new(format!("conn-{operation}"))); + let live = Arc::new(AtomicBool::new(false)); + let launches = Arc::new(AtomicUsize::new(0)); + let (release, release_rx) = tokio::sync::watch::channel(false); + let (role_tx, mut role_rx) = tokio::sync::mpsc::unbounded_channel(); + let first = tokio::spawn(attempt_start( + Arc::clone(&route), + Arc::clone(&live), + Arc::clone(&launches), + release_rx.clone(), + role_tx.clone(), + )); + let second = tokio::spawn(attempt_start( + route, + Arc::clone(&live), + Arc::clone(&launches), + release_rx, + role_tx, + )); + + let mut roles = [ + role_rx.recv().await.expect("first startup role"), + role_rx.recv().await.expect("second startup role"), + ]; + roles.sort_unstable(); + assert_eq!(roles, ["leader", "waiter"]); + assert_eq!( + launches.load(Ordering::Acquire), + 1, + "concurrent {operation} must launch one adapter" + ); + + release.send(true).expect("release startup"); + first.await.expect("first startup task"); + second.await.expect("second startup task"); + assert!(live.load(Ordering::Acquire)); + assert_eq!(launches.load(Ordering::Acquire), 1); + } + } + + #[test] + fn prompt_registration_is_early_busy_and_directly_cancellable() { + let route = Arc::new(AcpRouteEntry::new(String::from("conn-1"))); + let (prompt, cancellation) = route.begin_prompt().expect("prompt is admitted"); + let busy = route + .begin_prompt() + .expect_err("same-route second prompt must be rejected"); + assert!(busy.to_string().contains("session_busy")); + assert!(!*cancellation.borrow()); + + assert!(route + .signal_prompt_cancellation() + .expect("cancellation state is readable")); + assert!(*cancellation.borrow()); + + drop(prompt); + let (_next_prompt, next_cancellation) = route + .begin_prompt() + .expect("completed prompt releases route"); + assert!(!*next_cancellation.borrow()); + } + + #[test] + fn cancel_before_prompt_reservation_does_not_poison_a_future_prompt() { + let route = Arc::new(AcpRouteEntry::new(String::from("conn-1"))); + + assert!(!route + .signal_prompt_cancellation() + .expect("idle route cancellation is readable")); + + let (_prompt, cancellation) = route + .begin_prompt() + .expect("a later prompt still reserves the route"); + assert!( + !*cancellation.borrow(), + "a no-active-prompt cancellation must not leak into future work" + ); + } + + #[tokio::test(flavor = "current_thread")] + async fn cancel_after_reservation_is_visible_before_durable_acceptance() { + let (_dir, store) = durable_test_store().await; + create_durable_test_session(&store, "cancel-before-accept").await; + let route = Arc::new(AcpRouteEntry::new(String::from("conn-1"))); + let (prompt, cancellation) = route.begin_prompt().expect("prompt reserves route"); + + assert!(route + .signal_prompt_cancellation() + .expect("cancel reaches reserved prompt")); + assert!(*cancellation.borrow()); + + store + .accept_prompt( + "cancel-before-accept", + "prompt-1", + Some("key-1"), + vec![1; 32], + &[], + ) + .await + .expect("durable acceptance remains atomic"); + assert!( + *cancellation.borrow(), + "durable acceptance must not replace the installed cancellation token" + ); + finish_prompt_failure( + &store, + "cancel-before-accept", + "prompt-1", + None, + "prompt_interrupted", + SidecarError::InvalidState(String::from("cancelled before adapter write")), + ) + .await; + drop(prompt); + } + + #[tokio::test(flavor = "current_thread")] + async fn cancel_before_adapter_write_wins_without_polling_the_write() { + use std::future::poll_fn; + use std::sync::atomic::AtomicBool; + use std::task::Poll; + + let (sender, mut cancellation) = tokio::sync::watch::channel(false); + sender.send(true).expect("signal cancellation"); + let write_polled = AtomicBool::new(false); + + let outcome = race_prompt_cancellation( + &mut cancellation, + poll_fn(|_| { + write_polled.store(true, Ordering::Release); + Poll::<()>::Pending + }), + ) + .await; + + assert_eq!(outcome, PromptCancellationRace::Cancelled); + assert!( + !write_polled.load(Ordering::Acquire), + "a cancellation already visible at the write boundary must prevent the adapter write" + ); + } + + #[tokio::test(flavor = "current_thread")] + async fn cancel_during_output_wait_wakes_without_polling_or_time_delays() { + let (sender, mut cancellation) = tokio::sync::watch::channel(false); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (_release_tx, release_rx) = tokio::sync::oneshot::channel::<()>(); + + let wait = async { + race_prompt_cancellation(&mut cancellation, async { + started_tx.send(()).expect("mark output wait started"); + let _ = release_rx.await; + }) + .await + }; + let cancel = async { + started_rx.await.expect("output wait starts"); + sender.send(true).expect("cancel output wait"); + }; + let (outcome, ()) = tokio::join!(wait, cancel); + + assert_eq!(outcome, PromptCancellationRace::Cancelled); + } + + #[tokio::test(flavor = "current_thread")] + async fn cancel_after_terminal_commit_cannot_change_the_durable_outcome() { + let (_dir, store) = durable_test_store().await; + create_durable_test_session(&store, "cancel-after-terminal").await; + store + .accept_prompt( + "cancel-after-terminal", + "prompt-1", + Some("key-1"), + vec![2; 32], + &[], + ) + .await + .expect("accept prompt"); + let route = Arc::new(AcpRouteEntry::new(String::from("conn-1"))); + let (prompt, _) = route.begin_prompt().expect("prompt route starts"); + let result_json = + r#"{"sessionId":"cancel-after-terminal","message":null,"stopReason":"end_turn"}"#; + + store + .finish_prompt( + "cancel-after-terminal", + "prompt-1", + &[], + None, + Some(result_json), + None, + ) + .await + .expect("natural completion commits"); + drop(prompt); + + assert!( + !route + .signal_prompt_cancellation() + .expect("completed route cancellation is readable"), + "a cancel after terminal commit must not target completed work" + ); + assert!(store + .finish_prompt( + "cancel-after-terminal", + "prompt-1", + &[], + None, + None, + Some(r#"{"code":"prompt_interrupted"}"#), + ) + .await + .is_err()); + let stored = store + .prompt_by_idempotency_key("cancel-after-terminal", "key-1") + .await + .expect("read completed prompt") + .expect("completed prompt exists"); + assert_eq!(stored.state, "completed"); + assert_eq!(stored.result_json.as_deref(), Some(result_json)); + assert!(stored.error_json.is_none()); + } + + #[tokio::test(flavor = "current_thread")] + async fn natural_completion_and_interruption_race_commit_exactly_once() { + let (_dir, store) = durable_test_store().await; + create_durable_test_session(&store, "terminal-race").await; + store + .accept_prompt("terminal-race", "prompt-1", Some("key-1"), vec![3; 32], &[]) + .await + .expect("accept prompt"); + let barrier = Arc::new(tokio::sync::Barrier::new(3)); + let completed_store = store.clone(); + let completed_barrier = Arc::clone(&barrier); + let completed = tokio::spawn(async move { + completed_barrier.wait().await; + completed_store + .finish_prompt( + "terminal-race", + "prompt-1", + &[], + None, + Some(r#"{"sessionId":"terminal-race","message":null,"stopReason":"end_turn"}"#), + None, + ) + .await + }); + let interrupted_store = store.clone(); + let interrupted_barrier = Arc::clone(&barrier); + let interrupted = tokio::spawn(async move { + interrupted_barrier.wait().await; + interrupted_store + .finish_prompt( + "terminal-race", + "prompt-1", + &[], + None, + None, + Some(r#"{"code":"prompt_interrupted"}"#), + ) + .await + }); + barrier.wait().await; + let (completed, interrupted) = tokio::join!(completed, interrupted); + let completed = completed.expect("completion task"); + let interrupted = interrupted.expect("interruption task"); + + assert_ne!( + completed.is_ok(), + interrupted.is_ok(), + "the durable compare-and-set must admit exactly one terminal outcome" + ); + let stored = store + .prompt_by_idempotency_key("terminal-race", "key-1") + .await + .expect("read terminal prompt") + .expect("terminal prompt exists"); + assert!(matches!(stored.state.as_str(), "completed" | "failed")); + assert_ne!(stored.result_json.is_some(), stored.error_json.is_some()); + } + + #[test] + fn adapter_response_loop_is_single_consumer_per_route() { + let route = Arc::new(AcpRouteEntry::new(String::from("conn-1"))); + let rpc = route + .begin_adapter_rpc("set configuration") + .expect("first adapter RPC is admitted"); + let prompt_busy = route + .begin_prompt() + .expect_err("prompt cannot compete for the same adapter stdout"); + assert!(prompt_busy.to_string().contains("session_busy")); + let rpc_busy = route + .begin_adapter_rpc("another RPC") + .expect_err("second adapter response loop is rejected"); + assert!(rpc_busy.to_string().contains("session_busy")); + + drop(rpc); + assert!(route.begin_prompt().is_ok()); + + let other_route = Arc::new(AcpRouteEntry::new(String::from("conn-1"))); + assert!(other_route.begin_adapter_rpc("different route RPC").is_ok()); + } + + #[test] + fn adapter_waits_use_process_targeted_event_broker_only() { + let runtime = include_str!("runtime.rs"); + let turn = include_str!("turn.rs"); + assert!( + runtime + .matches("poll_process_event_wire(process_id") + .count() + >= 2, + "response and exit waits must target the owned adapter process" + ); + assert!( + turn.contains("poll_process_event_wire(process_id"), + "permission waits must target the owned adapter process" + ); + assert!( + !runtime.contains(".poll_event_wire(") && !turn.contains(".poll_event_wire("), + "ACP must not consume from the VM-wide process event queue" + ); + } + + #[test] + fn different_routes_run_and_complete_prompts_independently() { + let route_a = Arc::new(AcpRouteEntry::new(String::from("conn-1"))); + let route_b = Arc::new(AcpRouteEntry::new(String::from("conn-1"))); + let (prompt_a, _) = route_a.begin_prompt().expect("route A prompt starts"); + let (prompt_b, _) = route_b.begin_prompt().expect("route B prompt starts"); + assert_eq!( + route_a + .lock_state() + .expect("route A state") + .snapshot() + .phase, + AcpRoutePhase::PromptRunning + ); + assert_eq!( + route_b + .lock_state() + .expect("route B state") + .snapshot() + .phase, + AcpRoutePhase::PromptRunning + ); + + drop(prompt_a); + assert!(route_a.begin_adapter_rpc("route A follow-up").is_ok()); + assert!(route_b + .begin_adapter_rpc("route B competing RPC") + .expect_err("route B prompt remains independently active") + .to_string() + .contains("session_busy")); + + drop(prompt_b); + assert!(route_b.begin_prompt().is_ok()); + } + + #[test] + fn stopping_cancels_then_waits_for_prompt_guard_to_finish() { + use std::future::Future as _; + use std::task::{Context, Poll}; + + let route = Arc::new(AcpRouteEntry::new(String::from("conn-1"))); + let (prompt, cancellation) = route.begin_prompt().expect("prompt is admitted"); + let stop = route + .begin_stopping() + .expect("stop transitions a running prompt"); + assert!(*cancellation.borrow()); + + { + let mut wait = std::pin::pin!(stop.wait_for_prompt_completion()); + let waker = std::task::Waker::noop(); + let mut cx = Context::from_waker(waker); + assert!(matches!(wait.as_mut().poll(&mut cx), Poll::Pending)); + + drop(prompt); + assert!(matches!(wait.as_mut().poll(&mut cx), Poll::Ready(Ok(())))); + } + let busy = route + .begin_prompt() + .expect_err("stopping route rejects new prompts"); + assert!(busy.to_string().contains("session_busy")); + stop.complete(); + assert!(route.begin_prompt().is_ok()); + } + + #[test] + fn failed_teardown_keeps_route_stopping_and_allows_bounded_retry() { + let route = Arc::new(AcpRouteEntry::new(String::from("conn-1"))); + let (prompt, cancellation) = route.begin_prompt().expect("prompt is admitted"); + let first_stop = route.begin_stopping().expect("first teardown starts"); + assert!(*cancellation.borrow()); + + drop(first_stop); + assert!(route + .begin_prompt() + .expect_err("failed teardown must not reopen the route") + .to_string() + .contains("session_busy")); + + let retry = route + .begin_stopping() + .expect("a later unload/delete can retry the retained stopping route"); + assert!(route + .begin_stopping() + .expect_err("only one teardown attempt may own the stopping route") + .to_string() + .contains("session_busy")); + drop(prompt); + retry.complete(); + assert!(route.begin_prompt().is_ok()); + } + + async fn durable_test_store() -> (tempfile::TempDir, SessionStore) { + use agentos_native_sidecar::limits::DEFAULT_SQLITE_MAX_RESULT_BYTES; + use agentos_vm_config::VmSqliteDescriptor; + + let dir = tempfile::tempdir().expect("durable ACP test directory"); + let database = agentos_native_sidecar::vm_sqlite::resolve_vm_sqlite( + &VmSqliteDescriptor::SqliteFile { + path: dir.path().join("acp.sqlite").display().to_string(), + }, + agentos_runtime::SidecarRuntime::process(&agentos_runtime::RuntimeConfig::default()) + .expect("sidecar runtime") + .context(), + DEFAULT_SQLITE_MAX_RESULT_BYTES, + ) + .await + .expect("local durable ACP database"); + let store = SessionStore::open(database) + .await + .expect("migrate durable ACP store"); + (dir, store) + } + + async fn create_durable_test_session(store: &SessionStore, session_id: &str) { + store + .create( + session_id, + "pi", + &format!("native-{session_id}"), + "/workspace", + "{}", + None, + None, + "[]", + ) + .await + .expect("create durable ACP session"); + } + + #[tokio::test(flavor = "current_thread")] + async fn unload_and_delete_wait_for_durable_prompt_terminal_before_teardown() { + use std::future::Future as _; + use std::task::{Context, Poll}; + + let (_dir, store) = durable_test_store().await; + for action in ["unload", "delete"] { + let session_id = format!("{action}-session"); + let prompt_id = format!("{action}-prompt"); + let idempotency_key = format!("{action}-key"); + create_durable_test_session(&store, &session_id).await; + store + .accept_prompt( + &session_id, + &prompt_id, + Some(&idempotency_key), + vec![1; 32], + &[json!({ + "sessionUpdate": "user_message_chunk", + "content": { "type": "text", "text": action }, + "messageId": format!("{action}-message"), + })], + ) + .await + .expect("accept durable prompt"); + + let route = Arc::new(AcpRouteEntry::new(String::from("conn-1"))); + let (prompt, cancellation) = route.begin_prompt().expect("prompt route starts"); + let stop = route.begin_stopping().expect("teardown enters stopping"); + assert!(*cancellation.borrow()); + assert!(route.begin_prompt().is_err()); + assert!(route.begin_adapter_rpc("set configuration").is_err()); + + let mut wait = Box::pin(stop.wait_for_prompt_completion()); + let waker = std::task::Waker::noop(); + let mut cx = Context::from_waker(waker); + assert!(matches!(wait.as_mut().poll(&mut cx), Poll::Pending)); + + finish_prompt_failure( + &store, + &session_id, + &prompt_id, + None, + "prompt_interrupted", + SidecarError::InvalidState(format!("{action} cancelled the prompt")), + ) + .await; + let stored_prompt = store + .prompt_by_idempotency_key(&session_id, &idempotency_key) + .await + .expect("read durable prompt") + .expect("durable prompt exists"); + assert_eq!(stored_prompt.state, "failed"); + assert_eq!( + store + .get(&session_id) + .await + .expect("read durable session") + .expect("durable session exists") + .state, + "failed" + ); + assert!( + matches!(wait.as_mut().poll(&mut cx), Poll::Pending), + "teardown must wait until the prompt future releases its route guard" + ); + + drop(prompt); + wait.as_mut() + .await + .expect("teardown observes prompt completion"); + drop(wait); + stop.complete(); + } + } + + #[tokio::test(flavor = "current_thread")] + async fn adapter_kill_is_observed_by_prompt_and_commits_real_terminal_failure() { + let (_dir, store) = durable_test_store().await; + create_durable_test_session(&store, "kill-session").await; + store + .accept_prompt( + "kill-session", + "kill-prompt", + Some("kill-key"), + vec![4; 32], + &[], + ) + .await + .expect("accept killed prompt"); + let route = Arc::new(AcpRouteEntry::new(String::from("conn-1"))); + let (prompt, _) = route.begin_prompt().expect("prompt route starts"); + let event = + EventPayload::ProcessExitedEvent(agentos_native_sidecar::wire::ProcessExitedEvent { + process_id: String::from("adapter-killed"), + exit_code: 137, + }); + let (_, error) = matching_adapter_exit_error( + &event, + "adapter-killed", + 42, + &[String::from("sent request session/prompt id=42")], + "forced termination", + ) + .expect("prompt observes its adapter ProcessExited event"); + assert!(is_adapter_exited_error(&error)); + assert_eq!(adapter_exit_code_from_error(&error), Some(137)); + assert!(matching_adapter_exit_error(&event, "different-adapter", 42, &[], "").is_none()); + + let returned = finish_prompt_failure( + &store, + "kill-session", + "kill-prompt", + None, + "prompt_interrupted", + error, + ) + .await; + assert!(returned.to_string().contains("exited with code 137")); + let stored = store + .prompt_by_idempotency_key("kill-session", "kill-key") + .await + .expect("read killed prompt") + .expect("killed prompt exists"); + assert_eq!(stored.state, "failed"); + assert!(stored + .error_json + .as_deref() + .is_some_and(|error| error.contains("exited with code 137"))); + + drop(prompt); + assert!(route.begin_prompt().is_ok()); + } + + #[tokio::test(flavor = "current_thread")] + async fn adapter_kill_progresses_independently_without_dropping_prompt_future() { + let (_dir, store) = durable_test_store().await; + create_durable_test_session(&store, "independent-kill").await; + store + .accept_prompt( + "independent-kill", + "kill-prompt", + Some("kill-key"), + vec![5; 32], + &[], + ) + .await + .expect("accept prompt before adapter exit"); + let route = Arc::new(AcpRouteEntry::new(String::from("conn-1"))); + let (prompt_guard, _) = route.begin_prompt().expect("prompt route starts"); + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (event_tx, mut event_rx) = tokio::sync::mpsc::unbounded_channel(); + let prompt_store = store.clone(); + let prompt_task = tokio::spawn(async move { + started_tx.send(()).expect("mark prompt future started"); + let event = event_rx.recv().await.expect("adapter exit is delivered"); + let (_, error) = matching_adapter_exit_error( + &event, + "adapter-killed", + 42, + &[String::from("sent request session/prompt id=42")], + "forced termination", + ) + .expect("prompt future observes its adapter exit"); + let error = finish_prompt_failure( + &prompt_store, + "independent-kill", + "kill-prompt", + None, + "prompt_interrupted", + error, + ) + .await; + drop(prompt_guard); + error + }); + started_rx.await.expect("prompt reaches its event wait"); + + let (unrelated_tx, unrelated_rx) = tokio::sync::oneshot::channel(); + let unrelated = tokio::spawn(async move { + unrelated_rx.await.expect("unrelated request is released"); + "unrelated-complete" + }); + unrelated_tx.send(()).expect("release unrelated request"); + assert_eq!( + unrelated.await.expect("unrelated request task"), + "unrelated-complete" + ); + assert!( + !prompt_task.is_finished(), + "independent work must complete without dropping or replacing the prompt future" + ); + + event_tx + .send(EventPayload::ProcessExitedEvent( + agentos_native_sidecar::wire::ProcessExitedEvent { + process_id: String::from("adapter-killed"), + exit_code: 137, + }, + )) + .expect("deliver ProcessExited to prompt"); + let error = prompt_task.await.expect("prompt task remains supervised"); + assert!(error.to_string().contains("exited with code 137")); + let stored = store + .prompt_by_idempotency_key("independent-kill", "kill-key") + .await + .expect("read terminal prompt") + .expect("terminal prompt exists"); + assert_eq!(stored.state, "failed"); + assert!(route.begin_prompt().is_ok()); + } + + #[test] + fn permission_response_resumes_only_its_target_while_unrelated_route_stays_active() { + use agentos_native_sidecar::wire::VmOwnership; + use tokio::sync::oneshot::error::TryRecvError; + + let extension = AcpExtension::new(); + let ownership = OwnershipScope::VmOwnership(VmOwnership { + connection_id: String::from("conn-1"), + session_id: String::from("session-1"), + vm_id: String::from("vm-1"), + }); + let target_route_key = durable_route_key(&ownership, "target"); + let unrelated_route_key = durable_route_key(&ownership, "unrelated"); + let target_route = extension + .route_entry(&ownership, &target_route_key) + .expect("target route"); + let unrelated_route = extension + .route_entry(&ownership, &unrelated_route_key) + .expect("unrelated route"); + let (target_prompt, _) = target_route.begin_prompt().expect("target prompt starts"); + let (unrelated_prompt, _) = unrelated_route + .begin_prompt() + .expect("unrelated prompt starts concurrently"); + let (target_sender, mut target_receiver) = tokio::sync::oneshot::channel(); + let (unrelated_sender, mut unrelated_receiver) = tokio::sync::oneshot::channel(); + { + let mut pending = extension + .pending_permission_responses + .lock() + .expect("permission registry"); + pending.insert( + format!("{target_route_key}:permission-target"), + PendingPermissionResponse { + offered_option_ids: BTreeSet::from([String::from("allow")]), + acp_request_id: json!(7), + sender: target_sender, + }, + ); + pending.insert( + format!("{unrelated_route_key}:permission-unrelated"), + PendingPermissionResponse { + offered_option_ids: BTreeSet::from([String::from("allow")]), + acp_request_id: json!(8), + sender: unrelated_sender, + }, + ); + } + + let response = extension + .deliver_live_permission_response( + &ownership, + &AcpRespondPermissionRequest { + session_id: String::from("target"), + request_id: String::from("permission-target"), + option_id: String::from("allow"), + }, + ) + .expect("target permission response") + .expect("target waiter is live"); + assert_eq!(response.status, "accepted"); + assert!(matches!( + target_receiver.try_recv(), + Ok(PendingPermissionSignal::Selected(option)) if option == "allow" + )); + assert!(matches!( + unrelated_receiver.try_recv(), + Err(TryRecvError::Empty) + )); + assert_eq!( + unrelated_route + .lock_state() + .expect("unrelated route state") + .snapshot() + .phase, + AcpRoutePhase::PromptRunning + ); + + extension.cancel_pending_permissions(&unrelated_route_key, "prompt_cancelled"); + assert!(matches!( + unrelated_receiver.try_recv(), + Ok(PendingPermissionSignal::Terminal(reason)) if reason == "prompt_cancelled" + )); + drop(target_prompt); + drop(unrelated_prompt); + } + + #[test] + fn cancel_during_permission_wait_selects_exactly_one_terminal_signal() { + use agentos_native_sidecar::wire::VmOwnership; + + let extension = AcpExtension::new(); + let ownership = OwnershipScope::VmOwnership(VmOwnership { + connection_id: String::from("conn-1"), + session_id: String::from("session-1"), + vm_id: String::from("vm-1"), + }); + let route_key = durable_route_key(&ownership, "main"); + let request_key = format!("{route_key}:permission-race"); + let (sender, mut receiver) = tokio::sync::oneshot::channel(); + extension + .pending_permission_responses + .lock() + .expect("permission registry") + .insert( + request_key.clone(), + PendingPermissionResponse { + offered_option_ids: BTreeSet::from([String::from("allow")]), + acp_request_id: json!(9), + sender, + }, + ); + let barrier = Arc::new(std::sync::Barrier::new(3)); + + std::thread::scope(|scope| { + let response_barrier = Arc::clone(&barrier); + let response_extension = &extension; + let response_ownership = &ownership; + let response = scope.spawn(move || { + response_barrier.wait(); + response_extension.deliver_live_permission_response( + response_ownership, + &AcpRespondPermissionRequest { + session_id: String::from("main"), + request_id: String::from("permission-race"), + option_id: String::from("allow"), + }, + ) + }); + let cancel_barrier = Arc::clone(&barrier); + let cancel_extension = &extension; + let cancel_route_key = &route_key; + let cancel = scope.spawn(move || { + cancel_barrier.wait(); + cancel_extension.cancel_pending_permissions(cancel_route_key, "prompt_cancelled"); + }); + barrier.wait(); + response.join().expect("permission response thread").expect( + "a racing permission response is either delivered live or observes no waiter", + ); + cancel.join().expect("permission cancel thread"); + }); + + let delivered = receiver.try_recv().expect("one race winner signals waiter"); + assert!(match delivered { + PendingPermissionSignal::Selected(option) => option == "allow", + PendingPermissionSignal::Terminal(reason) => reason == "prompt_cancelled", + }); + assert!(extension + .pending_permission_responses + .lock() + .expect("permission registry") + .get(&request_key) + .is_none()); + assert!(matches!( + receiver.try_recv(), + Err(tokio::sync::oneshot::error::TryRecvError::Closed) + )); + } + + #[test] + fn invalid_permission_option_does_not_consume_live_waiter() { + use agentos_native_sidecar::wire::VmOwnership; + + let extension = AcpExtension::new(); + let ownership = OwnershipScope::VmOwnership(VmOwnership { + connection_id: String::from("conn-1"), + session_id: String::from("session-1"), + vm_id: String::from("vm-1"), + }); + let route_key = durable_route_key(&ownership, "main"); + let key = format!("{route_key}:permission-1"); + let (sender, mut receiver) = tokio::sync::oneshot::channel(); + extension + .pending_permission_responses + .lock() + .expect("permission registry") + .insert( + key.clone(), + PendingPermissionResponse { + offered_option_ids: BTreeSet::from([String::from("allow")]), + acp_request_id: json!(7), + sender, + }, + ); + + let invalid = extension + .deliver_live_permission_response( + &ownership, + &AcpRespondPermissionRequest { + session_id: String::from("main"), + request_id: String::from("permission-1"), + option_id: String::from("deny"), + }, + ) + .expect_err("invalid option is rejected"); + assert!(invalid.to_string().contains("invalid_permission_option")); + assert!(extension + .pending_permission_responses + .lock() + .expect("permission registry") + .contains_key(&key)); + + let accepted = extension + .deliver_live_permission_response( + &ownership, + &AcpRespondPermissionRequest { + session_id: String::from("main"), + request_id: String::from("permission-1"), + option_id: String::from("allow"), + }, + ) + .expect("valid option is accepted") + .expect("live waiter exists"); + assert_eq!(accepted.status, "accepted"); + assert!(matches!( + receiver.try_recv(), + Ok(PendingPermissionSignal::Selected(option)) if option == "allow" + )); + assert!(!extension + .pending_permission_responses + .lock() + .expect("permission registry") + .contains_key(&key)); + } + #[test] fn omitted_session_permission_policy_defaults_to_allow_all() { let request = AcpOpenSessionRequest { @@ -1522,6 +3168,30 @@ mod tests { assert!(buffer.is_empty()); } + #[test] + fn sequential_same_process_response_lines_preserve_ids_across_exchanges() { + let mut stdout = String::new(); + for (id, method) in [(1, "initialize"), (2, "session/new")] { + let chunk = format!( + "{{\"jsonrpc\":\"2.0\",\"id\":{id},\"result\":{{\"method\":\"{method}\"}}}}\n" + ); + let lines = append_stdout_chunk(&mut stdout, chunk.as_bytes(), 4096) + .expect("sequential response line"); + assert_eq!(lines.len(), 1); + let response: Value = + serde_json::from_str(&lines[0]).expect("sequential JSON-RPC response"); + assert_eq!(response.get("id").and_then(Value::as_i64), Some(id)); + assert_eq!( + response.pointer("/result/method").and_then(Value::as_str), + Some(method) + ); + assert!( + stdout.is_empty(), + "completed response leaves no stale bytes" + ); + } + } + #[test] fn bounded_stdout_lines_reject_complete_overlong_line() { let mut buffer = String::new(); diff --git a/crates/agentos-sidecar/src/acp/restore.rs b/crates/agentos-sidecar/src/acp/restore.rs index be709a57b4..13aeeadb79 100644 --- a/crates/agentos-sidecar/src/acp/restore.rs +++ b/crates/agentos-sidecar/src/acp/restore.rs @@ -3,14 +3,26 @@ use super::*; impl AcpExtension { pub(super) async fn ensure_durable_runtime( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, store: &SessionStore, - session: StoredSession, + mut session: StoredSession, ) -> Result { let route_key = durable_route_key(ctx.ownership(), &session.session_id); - if self.sessions.lock().await.contains_key(&route_key) { - return Ok(session); - } + let _start_guard = loop { + if self.sessions.lock().await.contains_key(&route_key) { + return Ok(session); + } + match self.begin_route_start(ctx.ownership(), &route_key)? { + AcpRouteStart::Leader(guard) => break guard, + AcpRouteStart::Wait(receiver) => { + AcpRouteEntry::wait_until_not_starting(receiver).await?; + if self.sessions.lock().await.contains_key(&route_key) { + return required_stored_session(store, &session.session_id).await; + } + session = required_stored_session(store, &session.session_id).await?; + } + } + }; let env = serde_json::from_str::>(&session.env_json) .map_err(|error| SidecarError::InvalidState(format!("invalid stored env: {error}")))?; let additional_directories = serde_json::from_str::>( @@ -119,7 +131,7 @@ impl AcpExtension { pub(super) async fn reapply_stored_config( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, route_key: &str, config_options_json: &str, ) -> Result<(), SidecarError> { @@ -174,7 +186,7 @@ impl AcpExtension { pub(super) async fn restore_acp_runtime( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: RestoreRuntimeRequest, user_session_id: &str, route_key: &str, @@ -318,7 +330,7 @@ impl AcpExtension { /// bootstrap state plus the chosen `mode` and any armed preamble. pub(super) async fn restore_acp_runtime_inner( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: &RestoreRuntimeRequest, create_like: &AcpCreateSessionRequest, process_id: &str, diff --git a/crates/agentos-sidecar/src/acp/runtime.rs b/crates/agentos-sidecar/src/acp/runtime.rs index afbfb7e005..40f23369e2 100644 --- a/crates/agentos-sidecar/src/acp/runtime.rs +++ b/crates/agentos-sidecar/src/acp/runtime.rs @@ -3,7 +3,7 @@ use super::*; impl AcpExtension { pub(super) async fn start_acp_runtime( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: AcpCreateSessionRequest, user_session_id: &str, route_key: &str, @@ -127,7 +127,7 @@ impl AcpExtension { /// whose manifest carries a non-empty `agent.acpEntrypoint` is an agent. The /// client parses no manifests — the sidecar owns agent enumeration too. Sorted /// by id. - pub(super) async fn list_agents(&self, mut ctx: ExtensionContext<'_>) -> AcpHandlerOutput { + pub(super) async fn list_agents(&self, mut ctx: ExtensionContext) -> AcpHandlerOutput { // The sidecar-owned projected-agent state is the SOURCE OF TRUTH for // installed agents (it reflects `ConfigureVm` and live `linkSoftware` // updates). Packed `.aospkg` packages ship no `agentos-package.json` in @@ -151,7 +151,7 @@ impl AcpExtension { pub(super) async fn start_acp_runtime_inner( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: &AcpCreateSessionRequest, process_id: &str, additional_directories: Vec, @@ -266,7 +266,7 @@ impl AcpExtension { pub(super) async fn stop_acp_runtime( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, route_key: &str, ) -> Result<(), SidecarError> { // Enforce per-connection ownership before tearing anything down: only the @@ -410,10 +410,47 @@ impl AcpExtension { Ok(()) } + /// Escalate a stopping route without starting another adapter response + /// loop. The active prompt remains the sole stdout/event consumer and + /// observes `ProcessExited`, so it can commit its real durable terminal + /// result before unload/delete continues. + pub(super) async fn force_kill_acp_runtime( + &self, + ctx: &mut ExtensionContext, + route_key: &str, + ) -> Result { + let caller_connection_id = ownership_connection_id(ctx.ownership()); + let process_id = { + let sessions = self.sessions.lock().await; + let Some(session) = sessions.get(route_key) else { + return Ok(false); + }; + if session.owner_connection_id != caller_connection_id { + return Err(SidecarError::InvalidState(format!( + "unknown ACP session {route_key}" + ))); + } + session.process_id.clone() + }; + match ctx + .kill_process_wire(KillProcessRequest { + process_id: process_id.clone(), + signal: String::from("SIGKILL"), + }) + .await + { + Ok(_) => Ok(true), + Err(error) if is_process_already_gone_error(&error) => Ok(false), + Err(error) => Err(SidecarError::InvalidState(format!( + "ACP adapter {process_id} could not be killed during bounded teardown: {error}" + ))), + } + } + #[allow(clippy::needless_option_as_deref)] pub(super) async fn send_runtime_request_with_sink( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: AcpSessionRequest, mut durable_sink: Option<&mut DurableUpdateSink>, cancellation: Option<&mut tokio::sync::watch::Receiver>, @@ -687,7 +724,7 @@ impl AcpExtension { /// caller retry cannot accidentally target the dead process. pub(super) async fn handle_adapter_exit( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, session_id: &str, exit_code: Option, error: SidecarError, @@ -710,6 +747,11 @@ impl AcpExtension { exit_code = ?exit_code, "ACP adapter process exited unexpectedly; live session route evicted", ); + if let Err(dispose_error) = ctx.dispose_session_resources_wire(session_id).await { + eprintln!( + "ERR_AGENTOS_ACP_EXIT_CLEANUP: failed to dispose resources for exited ACP route {session_id}: {dispose_error}" + ); + } let frame = encode_event(AcpEvent::AcpAgentExitedEvent(AcpAgentExitedEvent { session_id: session @@ -896,7 +938,7 @@ impl LiveAcpRuntime { /// is what makes `session/update`s arrive mid-turn instead of all arriving at /// once when the `session/prompt` dispatch finally resolves. pub(super) fn deliver_event( - ctx: &ExtensionContext<'_>, + ctx: &ExtensionContext, events: &mut Vec, frame: agentos_native_sidecar::wire::EventFrame, ) -> Result<(), SidecarError> { @@ -906,9 +948,59 @@ pub(super) fn deliver_event( Ok(()) } +#[derive(Debug, PartialEq, Eq)] +pub(super) enum PromptCancellationRace { + Cancelled, + Completed(T), +} + +/// Race one prompt-owned operation against its cancellation token. +/// +/// The current token value is checked before polling `operation`, and a +/// simultaneously-ready cancellation notification wins. This closes the two +/// narrow windows where cancellation could otherwise arrive after a caller's +/// manual check but before an adapter write or while its output wait was +/// parked. +pub(super) async fn race_prompt_cancellation( + cancellation: &mut tokio::sync::watch::Receiver, + operation: impl std::future::Future, +) -> PromptCancellationRace { + if *cancellation.borrow() { + return PromptCancellationRace::Cancelled; + } + tokio::pin!(operation); + loop { + tokio::select! { + biased; + changed = cancellation.changed() => match changed { + Ok(()) if *cancellation.borrow() => { + return PromptCancellationRace::Cancelled; + } + Ok(()) => continue, + Err(_) => { + return PromptCancellationRace::Completed(operation.await); + } + }, + result = &mut operation => return PromptCancellationRace::Completed(result), + } + } +} + +fn cancelled_prompt_exchange(response_id: i64) -> JsonRpcExchange { + JsonRpcExchange { + response: json!({ + "jsonrpc": "2.0", + "id": response_id, + "result": { "stopReason": "cancelled" }, + }), + events: Vec::new(), + notifications: Vec::new(), + } +} + #[allow(clippy::too_many_arguments)] pub(super) async fn send_json_rpc_request( - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, process_id: &str, agent_type: &str, request: Value, @@ -920,24 +1012,41 @@ pub(super) async fn send_json_rpc_request( mut cancellation: Option<&mut tokio::sync::watch::Receiver>, ) -> Result { let max_read_line_bytes = ctx.vm_acp_limits().await?.max_read_line_bytes; + let method = request + .get("method") + .and_then(Value::as_str) + .unwrap_or("unknown") + .to_string(); + if cancellation + .as_ref() + .is_some_and(|receiver| *receiver.borrow()) + { + return Ok(cancelled_prompt_exchange(response_id)); + } let mut line = serde_json::to_vec(&request).map_err(|error| { SidecarError::InvalidState(format!("failed to serialize ACP request: {error}")) })?; line.push(b'\n'); - ctx.write_stdin_wire(WriteStdinRequest { + let write = ctx.write_stdin_wire(WriteStdinRequest { process_id: process_id.to_string(), chunk: line, - }) - .await?; + }); + if let Some(receiver) = cancellation.as_deref_mut() { + match race_prompt_cancellation(receiver, write).await { + PromptCancellationRace::Cancelled => { + return Ok(cancelled_prompt_exchange(response_id)); + } + PromptCancellationRace::Completed(result) => { + result?; + } + } + } else { + write.await?; + } let deadline = timeout.map(|timeout| Instant::now() + timeout); let mut events = Vec::new(); let mut notifications = Vec::new(); - let method = request - .get("method") - .and_then(Value::as_str) - .unwrap_or("unknown") - .to_string(); let mut recent_activity = Vec::new(); let mut adapter_stderr = String::new(); record_recent_activity( @@ -1021,26 +1130,10 @@ pub(super) async fn send_json_rpc_request( if let Some(inactivity) = inactivity.as_ref() { remaining = remaining.min(inactivity.wait_duration(now)); } - // `poll_event_wire` already waits on the execution event receiver. Use - // the real request deadline so output/exit wakes this task directly; + // The process-targeted broker wait is durable and ownership checked. + // Use the real request deadline so output/exit wakes this task directly; // a sub-millisecond timeout loop only burns runtime turns while idle. - let event = if let Some(receiver) = cancellation.as_deref_mut() { - tokio::select! { - biased; - changed = receiver.changed() => { - if changed.is_ok() && *receiver.borrow() { - if let Some(session_id) = event_session_id { - write_session_cancel_notification(ctx, process_id, session_id).await?; - } - } - cancellation = None; - continue; - } - event = ctx.poll_event_wire(remaining) => event?, - } - } else { - ctx.poll_event_wire(remaining).await? - }; + let event = ctx.poll_process_event_wire(process_id, remaining).await?; let Some(event) = event else { if response_drain_deadline.is_some() { return Ok(JsonRpcExchange { @@ -1052,6 +1145,26 @@ pub(super) async fn send_json_rpc_request( continue; }; + if let Some((exit_code, error)) = matching_adapter_exit_error( + &event.payload, + process_id, + response_id, + &recent_activity, + &adapter_stderr, + ) { + let stderr_tail = adapter_stderr_tail(&adapter_stderr); + tracing::warn!( + target: "agentos_sidecar::acp_extension", + process_id, + agent_type, + session_id = ?event_session_id, + exit_code, + stderr_tail = %stderr_tail, + "ACP adapter process exited before answering request id={response_id}", + ); + return Err(error); + } + match event.payload { EventPayload::ProcessOutputEvent(output) if output.process_id == process_id && output.channel == StreamChannel::Stdout => @@ -1201,32 +1314,6 @@ pub(super) async fn send_json_rpc_request( events.push(frame); } } - EventPayload::ProcessExitedEvent(exited) if exited.process_id == process_id => { - // Embed ADAPTER_EXITED_ERROR_MARKER directly so is_adapter_exited_error() - // stays coupled to this producer: changing the wording can't silently - // disable session eviction (the H4 leak fix) without touching the const. - let stderr_tail: String = adapter_stderr - .chars() - .rev() - .take(4000) - .collect::() - .chars() - .rev() - .collect(); - tracing::warn!( - target: "agentos_sidecar::acp_extension", - process_id, - agent_type, - session_id = ?event_session_id, - exit_code = exited.exit_code, - stderr_tail = %stderr_tail, - "ACP adapter process exited before answering request id={response_id}", - ); - return Err(SidecarError::InvalidState(format!( - "ACP adapter process {process_id} {ADAPTER_EXITED_ERROR_MARKER} {} before response id={response_id}; recent_activity={:?}; adapter_stderr={:?}", - exited.exit_code, recent_activity, stderr_tail - ))); - } EventPayload::ProcessOutputEvent(_) | EventPayload::ProcessExitedEvent(_) | EventPayload::ExecutionOutputEvent(_) @@ -1238,6 +1325,43 @@ pub(super) async fn send_json_rpc_request( } } +pub(super) fn matching_adapter_exit_error( + payload: &EventPayload, + process_id: &str, + response_id: i64, + recent_activity: &[String], + adapter_stderr: &str, +) -> Option<(i32, SidecarError)> { + let EventPayload::ProcessExitedEvent(exited) = payload else { + return None; + }; + if exited.process_id != process_id { + return None; + } + let stderr_tail = adapter_stderr_tail(adapter_stderr); + // Embed ADAPTER_EXITED_ERROR_MARKER directly so is_adapter_exited_error() + // stays coupled to this producer: changing the wording cannot silently + // disable session eviction without updating its regression coverage. + Some(( + exited.exit_code, + SidecarError::InvalidState(format!( + "ACP adapter process {process_id} {ADAPTER_EXITED_ERROR_MARKER} {} before response id={response_id}; recent_activity={recent_activity:?}; adapter_stderr={stderr_tail:?}", + exited.exit_code + )), + )) +} + +fn adapter_stderr_tail(adapter_stderr: &str) -> String { + adapter_stderr + .chars() + .rev() + .take(4000) + .collect::() + .chars() + .rev() + .collect() +} + pub(super) fn record_recent_activity(recent_activity: &mut Vec, entry: String) { if recent_activity.len() == 16 { recent_activity.remove(0); @@ -1309,7 +1433,7 @@ pub(super) fn timeout_error_message( } pub(super) async fn wait_for_process_exit( - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, process_id: &str, timeout: Duration, ) -> bool { @@ -1320,7 +1444,7 @@ pub(super) async fn wait_for_process_exit( return false; } let Ok(event) = ctx - .poll_event_wire(deadline.saturating_duration_since(now)) + .poll_process_event_wire(process_id, deadline.saturating_duration_since(now)) .await else { return false; @@ -1337,13 +1461,13 @@ pub(super) async fn wait_for_process_exit( } pub(super) async fn handle_inbound_request( - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, process_id: &str, session_id: &str, message: &Value, events: &mut Vec, durable_sink: Option<&mut DurableUpdateSink>, - cancellation: Option<&mut tokio::sync::watch::Receiver>, + mut cancellation: Option<&mut tokio::sync::watch::Receiver>, ) -> Result<(), SidecarError> { let id = message.get("id").cloned().ok_or_else(|| { SidecarError::InvalidState(String::from("ACP inbound request missing id")) @@ -1359,11 +1483,10 @@ pub(super) async fn handle_inbound_request( .handle_permission_request( ctx, process_id, - session_id, message.get("id"), ¶ms, events, - cancellation, + cancellation.as_deref_mut(), ) .await?; json!({ @@ -1382,7 +1505,17 @@ pub(super) async fn handle_inbound_request( }) } } - _ => forward_inbound_host_request(ctx, session_id, message, &id, method)?, + _ => { + forward_inbound_host_request( + ctx, + session_id, + message, + &id, + method, + cancellation.as_deref_mut(), + ) + .await? + } }; let mut line = serde_json::to_vec(&response).map_err(|error| { SidecarError::InvalidState(format!("failed to serialize ACP inbound response: {error}")) @@ -1396,12 +1529,13 @@ pub(super) async fn handle_inbound_request( Ok(()) } -pub(super) fn forward_inbound_host_request( - ctx: &ExtensionContext<'_>, +pub(super) async fn forward_inbound_host_request( + ctx: &ExtensionContext, session_id: &str, message: &Value, id: &Value, method: &str, + mut cancellation: Option<&mut tokio::sync::watch::Receiver>, ) -> Result { let callback = AcpCallback::AcpHostRequestCallback(AcpHostRequestCallback { session_id: session_id.to_string(), @@ -1412,10 +1546,30 @@ pub(super) fn forward_inbound_host_request( // This path contains only noninteractive filesystem/terminal/internal host // RPCs. Human permission requests are handled durably above and deliberately // have no deadline. - let response = ctx.invoke_callback( + let callback = ctx.invoke_callback_async( encode_callback(callback)?, ACP_MACHINE_HOST_CALLBACK_TIMEOUT, - )?; + ); + tokio::pin!(callback); + let response = if let Some(cancellation) = cancellation.as_deref_mut() { + if *cancellation.borrow() { + return Ok(callback_cancelled_response(id.clone(), method)); + } + tokio::select! { + biased; + changed = cancellation.changed() => { + if changed.is_err() || *cancellation.borrow() { + return Ok(callback_cancelled_response(id.clone(), method)); + } + return Err(SidecarError::InvalidState(format!( + "ACP host request {method} cancellation channel changed without cancellation" + ))); + } + response = &mut callback => response?, + } + } else { + callback.await? + }; let response: AcpCallbackResponse = serde_bare::from_slice(&response).map_err(|error| { SidecarError::InvalidState(format!("invalid ACP host request response: {error}")) })?; @@ -1434,6 +1588,17 @@ pub(super) fn forward_inbound_host_request( Ok(response) } +fn callback_cancelled_response(id: Value, method: &str) -> Value { + json!({ + "jsonrpc": "2.0", + "id": id, + "error": { + "code": -32800, + "message": format!("request cancelled: {method}"), + }, + }) +} + pub(super) fn method_not_found_response(id: Value, method: &str) -> Value { json!({ "jsonrpc": "2.0", @@ -1540,7 +1705,7 @@ pub(super) struct AgentPackageAgentBlock { /// packages ship no `agentos-package.json` in the guest filesystem). A package /// without an agent block yields `None`. pub(super) async fn read_projected_agent_block( - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, agent_type: &str, ) -> Option { let launches = ctx.projected_agents().await.ok()?; @@ -1561,7 +1726,7 @@ pub(super) async fn read_projected_agent_block( /// missing file, a missing `agent` block, or an empty `agent.acpEntrypoint` all map /// to a single typed "unknown agent" error naming the agent and how to fix it. pub(super) async fn resolve_agent( - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, agent_type: &str, ) -> Result { match read_projected_agent_block(ctx, agent_type).await { @@ -1704,7 +1869,7 @@ pub(super) fn prepend_prompt_preamble(params: &mut Map, preamble: } } -pub(super) async fn kill_process_best_effort(ctx: &mut ExtensionContext<'_>, process_id: &str) { +pub(super) async fn kill_process_best_effort(ctx: &mut ExtensionContext, process_id: &str) { if let Err(error) = ctx .kill_process_wire(KillProcessRequest { process_id: process_id.to_owned(), diff --git a/crates/agentos-sidecar/src/acp/turn.rs b/crates/agentos-sidecar/src/acp/turn.rs index 89476c093a..ecd687bef3 100644 --- a/crates/agentos-sidecar/src/acp/turn.rs +++ b/crates/agentos-sidecar/src/acp/turn.rs @@ -3,7 +3,7 @@ use super::*; impl AcpExtension { pub(super) async fn prompt_durable_session( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: AcpPromptRequest, ) -> AcpHandlerOutput { let session_id = match default_session_id(request.session_id) { @@ -98,6 +98,12 @@ impl AcpExtension { "session_restore_failed: session {session_id} has no private ACP id" )))); } + let cancellation_key = durable_route_key(ctx.ownership(), &session_id); + let (_prompt_route_guard, mut cancellation_receiver) = + match self.begin_route_prompt(ctx.ownership(), &cancellation_key) { + Ok(prompt) => prompt, + Err(error) => return AcpHandlerOutput::response(Err(error)), + }; let prompt_id = uuid::Uuid::new_v4().to_string(); let user_message_id = uuid::Uuid::new_v4().to_string(); let user_updates = content @@ -121,6 +127,11 @@ impl AcpExtension { .await { Ok(events) => events, + Err(agentos_native_sidecar::vm_sqlite::VmSqliteError::UnexpectedChanges { .. }) => { + return AcpHandlerOutput::response(Err(route_busy_error(&format!( + "session {session_id} already has an active prompt" + )))); + } Err(error) => return AcpHandlerOutput::response(Err(session_store_error(error))), }; let mut events = Vec::new(); @@ -185,27 +196,6 @@ impl AcpExtension { }; } }; - let cancellation_key = durable_route_key(ctx.ownership(), &session_id); - let (cancellation_sender, mut cancellation_receiver) = tokio::sync::watch::channel(false); - if let Ok(mut cancellations) = self.prompt_cancellations.lock() { - cancellations.insert(cancellation_key.clone(), cancellation_sender); - } else { - let error = finish_prompt_failure( - &store, - &session_id, - &prompt_id, - sink.last_output_sequence, - "prompt_cancellation_registry_failed", - SidecarError::InvalidState(String::from( - "prompt cancellation registry is poisoned", - )), - ) - .await; - return AcpHandlerOutput { - response: Err(error), - events, - }; - } let raw = self .send_runtime_request_with_sink( ctx, @@ -218,13 +208,6 @@ impl AcpExtension { Some(&mut cancellation_receiver), ) .await; - if let Ok(mut cancellations) = self.prompt_cancellations.lock() { - cancellations.remove(&cancellation_key); - } else { - eprintln!( - "ERR_AGENTOS_PROMPT_CANCELLATION_REGISTRY: failed to remove completed prompt {session_id}" - ); - } events.extend(raw.events); let rpc = match raw.response { Ok(AcpResponse::AcpSessionRpcResponse(response)) => response, @@ -435,21 +418,59 @@ impl AcpExtension { pub(super) async fn cancel_durable_prompt( &self, - _ctx: &mut ExtensionContext<'_>, - _request: AcpCancelPromptRequest, + ctx: &mut ExtensionContext, + request: AcpCancelPromptRequest, ) -> AcpHandlerOutput { + let session_id = match default_session_id(request.session_id) { + Ok(session_id) => session_id, + Err(error) => return AcpHandlerOutput::response(Err(error)), + }; + let route_key = durable_route_key(ctx.ownership(), &session_id); + let signalled = self.signal_prompt_cancellation(&route_key); + self.cancel_pending_permissions(&route_key, "prompt_cancelled"); + if signalled { + let target = self + .sessions + .lock() + .await + .get(&route_key) + .map(|runtime| (runtime.process_id.clone(), runtime.acp_session_id.clone())); + let Some((process_id, acp_session_id)) = target else { + return AcpHandlerOutput::response(Err(SidecarError::InvalidState(format!( + "prompt_cancellation_target_missing: active prompt {session_id} has no live adapter runtime" + )))); + }; + if let Err(error) = + write_session_cancel_notification(ctx, &process_id, &acp_session_id).await + { + return AcpHandlerOutput::response(Err(error)); + } + } AcpHandlerOutput::response(Ok(AcpResponse::AcpCancelPromptResponse( AcpCancelPromptResponse { - status: String::from("no_active_prompt"), + status: if signalled { + String::from("cancelled") + } else { + String::from("no_active_prompt") + }, }, ))) } pub(super) async fn respond_durable_permission( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: AcpRespondPermissionRequest, ) -> AcpHandlerOutput { + match self.deliver_live_permission_response(ctx.ownership(), &request) { + Ok(Some(response)) => { + return AcpHandlerOutput::response(Ok(AcpResponse::AcpRespondPermissionResponse( + response, + ))); + } + Ok(None) => {} + Err(error) => return AcpHandlerOutput::response(Err(error)), + } let store = match self.session_store(ctx).await { Ok(store) => store, Err(error) => return AcpHandlerOutput::response(Err(error)), @@ -467,7 +488,7 @@ impl AcpExtension { pub(super) async fn set_durable_config( &self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, request: AcpSetSessionConfigOptionRequest, ) -> AcpHandlerOutput { let session_id = match default_session_id(request.session_id) { @@ -509,6 +530,12 @@ impl AcpExtension { "session_restore_failed: missing private ACP id", )))); } + let route_key = durable_route_key(ctx.ownership(), &session_id); + let _adapter_rpc_guard = + match self.begin_route_adapter_rpc(ctx.ownership(), &route_key, "set configuration") { + Ok(guard) => guard, + Err(error) => return AcpHandlerOutput::response(Err(error)), + }; let mut params = Map::from_iter([ (String::from("configId"), Value::String(request.config_id)), (String::from("value"), value.clone()), @@ -542,7 +569,7 @@ impl AcpExtension { .send_runtime_request_with_sink( ctx, AcpSessionRequest { - session_id: durable_route_key(ctx.ownership(), &session_id), + session_id: route_key.clone(), method: String::from("session/set_config_option"), params: Some(params), }, @@ -589,7 +616,6 @@ impl AcpExtension { events: output.events, }; } - let route_key = durable_route_key(ctx.ownership(), &session_id); let runtime = match self.sessions.lock().await.get(&route_key).cloned() { Some(runtime) => runtime, None => { @@ -631,20 +657,75 @@ impl AcpExtension { } pub(super) fn signal_prompt_cancellation(&self, route_key: &str) -> bool { - match self.prompt_cancellations.lock() { - Ok(cancellations) => cancellations - .get(route_key) - .cloned() - .is_some_and(|sender| sender.send(true).is_ok()), + let route = match self.routes.lock() { + Ok(routes) => routes.get(route_key).cloned(), Err(_) => { - eprintln!( - "ERR_AGENTOS_PROMPT_CANCELLATION_REGISTRY: cancellation registry is poisoned" - ); - false + eprintln!("ERR_AGENTOS_PROMPT_CANCELLATION_REGISTRY: route registry is poisoned"); + return false; } + }; + match route { + Some(route) => match route.signal_prompt_cancellation() { + Ok(signalled) => signalled, + Err(error) => { + eprintln!( + "ERR_AGENTOS_PROMPT_CANCELLATION_REGISTRY: failed to signal {route_key}: {error}" + ); + false + } + }, + None => false, } } + pub(super) fn deliver_live_permission_response( + &self, + ownership: &OwnershipScope, + response: &AcpRespondPermissionRequest, + ) -> Result, SidecarError> { + let key = format!( + "{}:{}", + durable_route_key(ownership, &response.session_id), + response.request_id + ); + let selected_option_id = response.option_id.clone(); + let pending = { + let mut pending = self.pending_permission_responses.lock().map_err(|_| { + SidecarError::InvalidState(String::from("permission response registry is poisoned")) + })?; + let Some(entry) = pending.get(&key) else { + return Ok(None); + }; + if !entry.offered_option_ids.contains(&response.option_id) { + let valid_options = entry + .offered_option_ids + .iter() + .cloned() + .collect::>() + .join(", "); + return Err(SidecarError::InvalidState(format!( + "invalid_permission_option: request {} does not offer {}; valid option IDs: {}", + response.request_id, response.option_id, valid_options + ))); + } + pending.remove(&key) + }; + let accepted = pending.is_some_and(|entry| { + entry + .sender + .send(PendingPermissionSignal::Selected(selected_option_id)) + .is_ok() + }); + Ok(Some(AcpRespondPermissionResponse { + status: if accepted { + String::from("accepted") + } else { + String::from("not_pending") + }, + reason: (!accepted).then(|| String::from("already_resolved")), + })) + } + pub(super) fn cancel_pending_permissions(&self, route_key: &str, reason: &str) { let prefix = format!("{route_key}:"); let mut pending = match self.pending_permission_responses.lock() { @@ -756,9 +837,8 @@ impl DurableUpdateSink { #[allow(clippy::too_many_arguments)] pub(super) async fn handle_permission_request( &mut self, - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, process_id: &str, - acp_session_id: &str, rpc_id: Option<&Value>, params: &Value, events: &mut Vec, @@ -882,7 +962,6 @@ impl DurableUpdateSink { { self.emit_stored(ctx, events, std::slice::from_ref(&event))?; } - write_session_cancel_notification(ctx, process_id, acp_session_id).await?; return Ok(json!({ "outcome": { "outcome": "cancelled" } })); } }; @@ -921,7 +1000,7 @@ impl DurableUpdateSink { pub(super) async fn handle_notification( &mut self, - ctx: &ExtensionContext<'_>, + ctx: &ExtensionContext, notification: &Value, events: &mut Vec, ) -> Result { @@ -1036,7 +1115,7 @@ impl DurableUpdateSink { pub(super) async fn flush( &mut self, - ctx: &ExtensionContext<'_>, + ctx: &ExtensionContext, events: &mut Vec, ) -> Result<(), SidecarError> { if self.buffered.is_empty() { @@ -1056,7 +1135,7 @@ impl DurableUpdateSink { pub(super) async fn persist( &mut self, - ctx: &ExtensionContext<'_>, + ctx: &ExtensionContext, events: &mut Vec, updates: Vec, ) -> Result<(), SidecarError> { @@ -1079,7 +1158,7 @@ impl DurableUpdateSink { pub(super) fn emit_stored( &self, - ctx: &ExtensionContext<'_>, + ctx: &ExtensionContext, events: &mut Vec, stored: &[StoredEvent], ) -> Result<(), SidecarError> { @@ -1189,7 +1268,7 @@ pub(super) fn decode_durable_event(event_json: &str) -> Result, + ctx: &mut ExtensionContext, process_id: &str, request_id: &str, key: &str, @@ -1230,7 +1309,10 @@ async fn wait_for_permission_signal( "prompt_cancelled", ))); } - let polled = async { ctx.poll_event_wire(Duration::from_secs(1)).await }; + let polled = async { + ctx.poll_process_event_wire(process_id, Duration::from_secs(1)) + .await + }; let signal = if let Some(cancellation) = cancellation.as_deref_mut() { tokio::select! { biased; @@ -1271,7 +1353,7 @@ async fn wait_for_permission_signal( } fn handle_permission_wait_event( - ctx: &ExtensionContext<'_>, + ctx: &ExtensionContext, process_id: &str, key: &str, pending: &Arc>>, @@ -1302,7 +1384,7 @@ fn handle_permission_wait_event( } pub(super) async fn write_session_cancel_notification( - ctx: &mut ExtensionContext<'_>, + ctx: &mut ExtensionContext, process_id: &str, session_id: &str, ) -> Result<(), SidecarError> { @@ -1343,42 +1425,6 @@ pub(super) fn cancel_notification_fallback_response(id: Value) -> Value { }) } -pub(super) fn encode_durable_interrupted_prompt(session_id: &str) -> Option> { - encode_response(AcpResponse::AcpPromptResponse(AcpPromptResponse { - session_id: session_id.to_owned(), - message: None, - stop_reason: String::from("cancelled"), - })) - .ok() -} - -pub(super) fn encode_durable_cancel_response(signalled: bool) -> Option> { - encode_response(AcpResponse::AcpCancelPromptResponse( - AcpCancelPromptResponse { - status: if signalled { - String::from("cancelled") - } else { - String::from("no_active_prompt") - }, - }, - )) - .ok() -} - -pub(super) fn encode_durable_permission_response(accepted: bool) -> Option> { - encode_response(AcpResponse::AcpRespondPermissionResponse( - AcpRespondPermissionResponse { - status: if accepted { - String::from("accepted") - } else { - String::from("not_pending") - }, - reason: (!accepted).then(|| String::from("already_resolved")), - }, - )) - .ok() -} - pub(super) fn synthetic_mode_update(mode_id: &str) -> Value { json!({ "jsonrpc": "2.0", diff --git a/crates/agentos-sidecar/src/session_store.rs b/crates/agentos-sidecar/src/session_store.rs index 980d435a22..6638783903 100644 --- a/crates/agentos-sidecar/src/session_store.rs +++ b/crates/agentos-sidecar/src/session_store.rs @@ -614,7 +614,7 @@ impl SessionStore { .map(|event| event.payload_bytes) .sum::(); let mut statements = vec![SqlStatement::new( - "UPDATE agentos_core_sessions SET state = ?, state_prompt_id = CASE WHEN ? = 'idle' THEN NULL ELSE state_prompt_id END, state_started_at_ms = CASE WHEN ? = 'idle' THEN NULL ELSE state_started_at_ms END, latest_sequence = latest_sequence + ?, retained_event_count = retained_event_count + ?, retained_event_bytes = retained_event_bytes + ?, updated_at_ms = ? WHERE session_id = ? AND state_prompt_id = ? AND latest_sequence <= ?", + "UPDATE agentos_core_sessions SET state = ?, state_prompt_id = CASE WHEN ? = 'idle' THEN NULL ELSE state_prompt_id END, state_started_at_ms = CASE WHEN ? = 'idle' THEN NULL ELSE state_started_at_ms END, latest_sequence = latest_sequence + ?, retained_event_count = retained_event_count + ?, retained_event_bytes = retained_event_bytes + ?, updated_at_ms = ? WHERE session_id = ? AND state_prompt_id = ? AND latest_sequence <= ? AND EXISTS (SELECT 1 FROM agentos_core_prompts WHERE session_id = ? AND prompt_id = ? AND state = 'accepted')", vec![ text(state), text(state), @@ -626,6 +626,8 @@ impl SessionStore { text(session_id), text(prompt_id), SqlValue::SqlInteger(MAX_SAFE_SEQUENCE - count), + text(session_id), + text(prompt_id), ], ) .expect_changes(1)]; @@ -647,7 +649,7 @@ impl SessionStore { )); } statements.push(SqlStatement::new( - "UPDATE agentos_core_prompts SET state = ?, result_json = ?, error_json = ?, last_output_sequence = COALESCE(?, CASE WHEN ? = 0 THEN last_output_sequence ELSE (SELECT latest_sequence FROM agentos_core_sessions WHERE session_id = ?) END), updated_at_ms = ? WHERE session_id = ? AND prompt_id = ?", + "UPDATE agentos_core_prompts SET state = ?, result_json = ?, error_json = ?, last_output_sequence = COALESCE(?, CASE WHEN ? = 0 THEN last_output_sequence ELSE (SELECT latest_sequence FROM agentos_core_sessions WHERE session_id = ?) END), updated_at_ms = ? WHERE session_id = ? AND prompt_id = ? AND state = 'accepted'", vec![ text(prompt_state), optional_text(result_json),