diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index a0500e41fd..c09b51b7cf 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -5894,6 +5894,7 @@ dependencies = [ "anyhow", "assert_cmd", "base64 0.22.1", + "codex-agent-graph-store", "codex-analytics", "codex-arg0", "codex-attachment-store", diff --git a/codex-rs/agent-graph-store/src/store.rs b/codex-rs/agent-graph-store/src/store.rs index 0760cb15d4..a37a3e958c 100644 --- a/codex-rs/agent-graph-store/src/store.rs +++ b/codex-rs/agent-graph-store/src/store.rs @@ -19,6 +19,8 @@ pub trait AgentGraphStore: Send + Sync { /// /// `child_thread_id` has at most one persisted parent. Re-inserting the same child should /// update both the parent and status to match the supplied values. + /// The edge may be published before the child's metadata or history exists. A successful + /// write must make the child discoverable through the list methods. fn upsert_thread_spawn_edge( &self, parent_thread_id: ThreadId, diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 541eedbb35..579d6016f3 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -634,24 +634,23 @@ impl LocalAgentControl { Some(parent_thread.session.inherited_environment_snapshot().await) } + /// Returns true after publication, or false when the source or configuration skips it. async fn persist_thread_spawn_edge_for_source( &self, - child_thread: &crate::CodexThread, + ephemeral: bool, child_thread_id: ThreadId, session_source: Option<&SessionSource>, - ) { + ) -> CodexResult { let Some(parent_thread_id) = session_source.and_then(SessionSource::parent_thread_id) else { - return; + return Ok(false); }; - if child_thread.config_snapshot().await.ephemeral { - return; + if ephemeral { + return Ok(false); } - let Ok(state) = self.runtime.upgrade() else { - return; - }; + let state = self.runtime.upgrade()?; let Some(agent_graph_store) = state.agent_graph_store() else { - return; + return Ok(false); }; if let Err(err) = agent_graph_store .upsert_thread_spawn_edge( @@ -662,7 +661,11 @@ impl LocalAgentControl { .await { warn!("failed to persist thread-spawn edge: {err}"); + return Err(CodexErr::Fatal( + "failed to persist spawned agent relationship".to_string(), + )); } + Ok(true) } } diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index 9f1660cf61..5587e230f0 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -753,6 +753,33 @@ impl LocalAgentControl { } (None, _) => None, }; + // Fork startup can persist copied history before returning a live child. + // Publish its deletion relationship before any child storage is created. + let child_thread_id = self.runtime.generate_thread_id(); + if state.threads.read().await.contains_key(&child_thread_id) { + return Err(CodexErr::InvalidRequest(format!( + "thread {child_thread_id} is already running", + )) + .with_agent_context(AgentErrorContext::ChildStartup)); + } + start_options.reserved_thread_id = Some(child_thread_id); + let mut pending_spawn = PendingSpawn::new( + Arc::clone(&state), + child_thread_id, + membership, + SpawnOwner::Local, + ); + let edge_write_started_at = Instant::now(); + let control = self.clone(); + let source = notification_source.clone(); + let ephemeral = config.ephemeral; + pending_spawn.set_edge_write(tokio::spawn(async move { + control + .persist_thread_spawn_edge_for_source(ephemeral, child_thread_id, source.as_ref()) + .await + })); + pending_spawn.wait_for_edge().await?; + let edge_write = edge_write_started_at.elapsed(); let child_create_started_at = Instant::now(); let new_thread = state .spawn_child_thread( @@ -765,13 +792,6 @@ impl LocalAgentControl { .map_err(|err| err.with_agent_context(AgentErrorContext::ChildStartup))?; let child_create = child_create_started_at.elapsed(); agent_metadata.agent_id = Some(new_thread.thread_id); - let mut pending_spawn = PendingSpawn::new( - Arc::clone(&state), - new_thread.thread_id, - membership, - SpawnOwner::Local, - ); - if let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id, .. })) = notification_source.as_ref() @@ -796,29 +816,11 @@ impl LocalAgentControl { .await; } - let control = self.clone(); - let child = Arc::clone(&new_thread.thread); - let child_thread_id = new_thread.thread_id; - let source = notification_source.clone(); - pending_spawn.set_edge_write(tokio::spawn(async move { - control - .persist_thread_spawn_edge_for_source( - child.as_ref(), - child_thread_id, - source.as_ref(), - ) - .await; - })); let durability_wait_started_at = Instant::now(); if options.fork_mode.is_some() { - tokio::join!( - new_thread.thread.ensure_rollout_materialized(), - pending_spawn.wait_for_edge(), - ); - } else { - pending_spawn.wait_for_edge().await; + new_thread.thread.ensure_rollout_materialized().await; } - let durability_wait = durability_wait_started_at.elapsed(); + let durability_wait = edge_write + durability_wait_started_at.elapsed(); // Capture startup settings before the child can publish metadata for a later turn. let spawn_config = if multi_agent_version == MultiAgentVersion::V2 { @@ -1093,12 +1095,13 @@ impl LocalAgentControl { agent_metadata.agent_path.clone(), ); } - self.persist_thread_spawn_edge_for_source( - resumed_thread.thread.as_ref(), - resumed_thread.thread_id, - Some(¬ification_source), - ) - .await; + let _ = self + .persist_thread_spawn_edge_for_source( + resumed_thread.thread.config_snapshot().await.ephemeral, + resumed_thread.thread_id, + Some(¬ification_source), + ) + .await; Ok((resumed_thread.thread_id, multi_agent_version)) } diff --git a/codex-rs/core/src/agent/control/spawn_guard.rs b/codex-rs/core/src/agent/control/spawn_guard.rs index 1918af66dc..2c0a8baaa4 100644 --- a/codex-rs/core/src/agent/control/spawn_guard.rs +++ b/codex-rs/core/src/agent/control/spawn_guard.rs @@ -1,4 +1,5 @@ //! Owns a provisional child until the spawn result can be returned without yielding. +//! Graph publication stays owned through cancellation and settles before cleanup closes its edge. use super::AgentTreeMembership; use crate::thread_manager::ThreadManagerState; @@ -7,6 +8,7 @@ use codex_agent_graph_store::AgentGraphStoreError; use codex_agent_graph_store::ThreadSpawnEdgeStatus; use codex_protocol::ThreadId; use codex_protocol::error::CodexErrKind; +use codex_protocol::error::Result as CodexResult; use codex_thread_store::ThreadStoreError; use std::sync::Arc; use tokio::task::JoinHandle; @@ -20,7 +22,8 @@ pub(in crate::agent) enum SpawnOwner { pub(in crate::agent) struct PendingSpawn { state: Arc, pending: Option<(ThreadId, AgentTreeMembership)>, - edge_write: Option>, + edge_write: Option>>, + close_spawn_edge: bool, owner: SpawnOwner, } @@ -35,22 +38,26 @@ impl PendingSpawn { state, pending: Some((child, membership)), edge_write: None, + close_spawn_edge: false, owner, } } - pub(super) fn set_edge_write(&mut self, edge_write: JoinHandle<()>) { + pub(super) fn set_edge_write(&mut self, edge_write: JoinHandle>) { self.edge_write = Some(edge_write); } - pub(super) async fn wait_for_edge(&mut self) { - if let Some(edge_write) = self.edge_write.as_mut() { - assert!( - edge_write.await.is_ok(), - "spawn edge write task should complete" - ); - } + pub(super) async fn wait_for_edge(&mut self) -> CodexResult<()> { + let Some(edge_write) = self.edge_write.as_mut() else { + return Ok(()); + }; + // Keep ownership while waiting, then clear the completed handle before returning errors. + let result = edge_write.await; self.edge_write = None; + // An error may follow a committed write; only an explicit skip rules it out. + self.close_spawn_edge = !matches!(&result, Ok(Ok(false))); + result??; + Ok(()) } pub(in crate::agent) fn disarm(mut self) -> AgentTreeMembership { @@ -69,6 +76,7 @@ impl Drop for PendingSpawn { let state = Arc::clone(&self.state); let teardown = membership.into_teardown_guard("child_spawn_cleanup", Some(child)); let edge_write = self.edge_write.take(); + let close_spawn_edge = self.close_spawn_edge; let local_owner = matches!(self.owner, SpawnOwner::Local); drop(tokio::spawn(async move { if let Some(thread) = state.remove_thread(&child).await { @@ -91,10 +99,12 @@ impl Drop for PendingSpawn { } } // A pending Open write must finish before cleanup writes Closed. - if let Some(edge_write) = edge_write { - let _ = edge_write.await; - } + let close_spawn_edge = match edge_write { + Some(edge_write) => !matches!(edge_write.await, Ok(Ok(false))), + None => close_spawn_edge, + }; if local_owner + && close_spawn_edge && let Some(store) = state.agent_graph_store() && let Err(error) = store .set_thread_spawn_edge_status(child, ThreadSpawnEdgeStatus::Closed) diff --git a/codex-rs/core/src/agent/control/spawn_publication_tests.rs b/codex-rs/core/src/agent/control/spawn_publication_tests.rs new file mode 100644 index 0000000000..3df67dbbdc --- /dev/null +++ b/codex-rs/core/src/agent/control/spawn_publication_tests.rs @@ -0,0 +1,356 @@ +#![cfg(test)] + +//! Exercises spawn cleanup against real SQLite graph writes with controlled acknowledgements. +//! Publication must survive caller cancellation and finish before cleanup closes the edge. + +use super::*; +use codex_agent_graph_store::AgentGraphStore; +use codex_agent_graph_store::AgentGraphStoreError; +use codex_agent_graph_store::AgentGraphStoreFuture; +use codex_agent_graph_store::LocalAgentGraphStore; +use codex_agent_graph_store::ThreadSpawnEdgeStatus; +use pretty_assertions::assert_eq; + +#[derive(Clone, Copy)] +enum SpawnEdgeFailure { + CancelledBeforeCommit, + CancelledBeforeAck, + ErrorAfterCommit, +} + +enum SkippedPublication { + Ephemeral, + NoParent, +} + +struct FaultingGraphStore { + inner: LocalAgentGraphStore, + failure: SpawnEdgeFailure, + started: CancellationToken, + held: CancellationToken, + release: CancellationToken, + cancelled: CancellationToken, + close_attempted: CancellationToken, +} + +impl FaultingGraphStore { + async fn hold_write(&self) { + self.held.cancel(); + self.release.cancelled().await; + } +} + +impl AgentGraphStore for FaultingGraphStore { + fn upsert_thread_spawn_edge( + &self, + parent: ThreadId, + child: ThreadId, + status: ThreadSpawnEdgeStatus, + ) -> AgentGraphStoreFuture<'_, ()> { + Box::pin(async move { + let cancellation_guard = self.cancelled.clone().drop_guard(); + self.started.cancel(); + if matches!(self.failure, SpawnEdgeFailure::CancelledBeforeCommit) { + self.hold_write().await; + } + self.inner + .upsert_thread_spawn_edge(parent, child, status) + .await?; + if matches!(self.failure, SpawnEdgeFailure::CancelledBeforeAck) { + self.hold_write().await; + } + cancellation_guard.disarm(); + match self.failure { + SpawnEdgeFailure::CancelledBeforeCommit | SpawnEdgeFailure::CancelledBeforeAck => { + Ok(()) + } + SpawnEdgeFailure::ErrorAfterCommit => Err(AgentGraphStoreError::Internal { + message: "injected response loss after successful commit".to_string(), + }), + } + }) + } + + fn set_thread_spawn_edge_status( + &self, + child: ThreadId, + status: ThreadSpawnEdgeStatus, + ) -> AgentGraphStoreFuture<'_, ()> { + self.close_attempted.cancel(); + self.inner.set_thread_spawn_edge_status(child, status) + } + + fn list_thread_spawn_children( + &self, + parent: ThreadId, + status: Option, + ) -> AgentGraphStoreFuture<'_, Vec> { + self.inner.list_thread_spawn_children(parent, status) + } + + fn list_thread_spawn_descendants( + &self, + root: ThreadId, + status: Option, + ) -> AgentGraphStoreFuture<'_, Vec> { + self.inner.list_thread_spawn_descendants(root, status) + } +} + +async fn publication_harness( + failure: SpawnEdgeFailure, + generate_thread_id: impl Fn() -> ThreadId + Send + Sync + 'static, +) -> (AgentControlHarness, Arc) { + let mut harness = AgentControlHarness::new().await; + let graph = Arc::new(FaultingGraphStore { + inner: LocalAgentGraphStore::new(harness.state_db.clone().expect("state db")), + failure, + started: CancellationToken::new(), + held: CancellationToken::new(), + release: CancellationToken::new(), + cancelled: CancellationToken::new(), + close_attempted: CancellationToken::new(), + }); + harness.manager = ThreadManager::new( + &harness.config, + harness.manager.auth_manager(), + harness.manager.get_models_manager(), + crate::CodexAppsToolsCache::default(), + SessionSource::Exec, + harness.manager.environment_manager(), + empty_extension_registry(), + Arc::new(crate::test_support::EmptyUserInstructionsProvider), + /*analytics_events_client*/ None, + harness.manager.image_store(), + crate::thread_manager::thread_store_from_config(&harness.config, harness.state_db.clone()), + Some(graph.clone()), + "spawn-publication-cleanup-test".to_string(), + /*attestation_provider*/ None, + /*external_time_provider*/ None, + ) + .with_thread_id_generator(generate_thread_id); + harness.control = harness.manager.agent_control(); + (harness, graph) +} + +async fn spawn_publication_child( + harness: &AgentControlHarness, + parent: ThreadId, +) -> CodexResult { + harness + .control + .spawn_agent( + harness.config.clone(), + text_input("child task"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: parent, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + })), + ) + .await +} + +async fn assert_publication_closed( + harness: &AgentControlHarness, + graph: &FaultingGraphStore, + parent: ThreadId, +) { + harness + .control + .runtime + .request_shutdown() + .wait() + .await + .expect("cleanup completes"); + let all = graph + .list_thread_spawn_children(parent, /*status_filter*/ None) + .await + .expect("graph readable"); + assert_eq!( + all.len(), + 1, + "the store durably accepted the prepublication" + ); + assert_thread_not_loaded(&harness.manager, all[0]).await; + assert_eq!( + graph + .list_thread_spawn_children(parent, Some(ThreadSpawnEdgeStatus::Open)) + .await + .expect("open graph"), + Vec::::new(), + "rejected or cancelled spawn must close its committed edge", + ); + assert_eq!( + graph + .list_thread_spawn_children(parent, Some(ThreadSpawnEdgeStatus::Closed)) + .await + .expect("closed graph"), + all, + ); +} + +#[test_case::test_case(SpawnEdgeFailure::CancelledBeforeCommit, 0; "before_commit")] +#[test_case::test_case(SpawnEdgeFailure::CancelledBeforeAck, 1; "before_acknowledgement")] +#[tokio::test] +async fn cancelled_spawn_edge_publication_cleanup( + failure: SpawnEdgeFailure, + committed_children: usize, +) { + let (harness, graph) = publication_harness(failure, ThreadId::new).await; + let parent = ThreadId::new(); + let mut spawn = Box::pin(spawn_publication_child(&harness, parent)); + tokio::select! { + biased; + () = graph.held.cancelled() => {}, + result = &mut spawn => panic!("spawn returned while publication was held: {result:?}"), + } + assert_eq!( + graph + .list_thread_spawn_children(parent, /*status_filter*/ None) + .await + .expect("graph before cancellation") + .len(), + committed_children + ); + drop(spawn); + // Dropping a directly-awaited store future cancels this token synchronously. + // Keeping only the child ID in the guard cannot satisfy this ownership check. + assert!( + !graph.cancelled.is_cancelled(), + "caller cancellation must retain the submitted write" + ); + // Paused time advances only after runnable cleanup has been polled. No real + // storage operation is pending while the graph-store future is held. + tokio::time::pause(); + let early_close = timeout(Duration::from_secs(1), graph.close_attempted.cancelled()).await; + tokio::time::resume(); + assert!(early_close.is_err(), "Closed must wait for Open"); + graph.release.cancel(); + assert_publication_closed(&harness, &graph, parent).await; +} + +#[tokio::test] +async fn failed_spawn_edge_publication_cleanup() { + let (harness, graph) = + publication_harness(SpawnEdgeFailure::ErrorAfterCommit, ThreadId::new).await; + let parent = ThreadId::new(); + let error = spawn_publication_child(&harness, parent) + .await + .expect_err("lost commit acknowledgement rejects spawn"); + assert_eq!( + error.to_string(), + "Fatal error: failed to persist spawned agent relationship" + ); + assert_publication_closed(&harness, &graph, parent).await; +} + +#[test_case::test_case(SkippedPublication::Ephemeral; "ephemeral")] +#[test_case::test_case(SkippedPublication::NoParent; "no_parent")] +#[tokio::test] +async fn failed_spawn_without_publication_does_not_close_graph(skipped: SkippedPublication) { + let (mut harness, graph) = + publication_harness(SpawnEdgeFailure::ErrorAfterCommit, ThreadId::new).await; + let parent = ThreadId::new(); + let source = match skipped { + SkippedPublication::Ephemeral => { + harness.config.ephemeral = true; + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: parent, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + })) + } + SkippedPublication::NoParent => None, + }; + let state = harness.control.runtime.upgrade().expect("live manager"); + let spawn = + harness + .control + .spawn_agent(harness.config.clone(), text_input("child task"), source); + tokio::pin!(spawn); + let shutdown = { + let _threads = state.threads.write().await; + assert!( + spawn + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())) + .is_pending() + ); + harness.control.runtime.request_shutdown() + }; + assert_eq!( + CodexErrKind::from(&spawn.await.expect_err("child startup rejects shutdown")), + CodexErrKind::InvalidRequest, + ); + shutdown + .wait() + .await + .expect("cleanup completes without graph errors"); + assert!(!graph.started.is_cancelled(), "no Open write was issued"); + assert!( + !graph.close_attempted.is_cancelled(), + "skipped publication needs no Closed write" + ); + assert_eq!( + graph + .list_thread_spawn_children(parent, /*status_filter*/ None) + .await + .expect("graph readable"), + Vec::::new() + ); +} + +#[tokio::test] +async fn repeated_child_id_preserves_existing_thread_and_graph() { + let repeated_id = ThreadId::new(); + let (harness, graph) = + publication_harness(SpawnEdgeFailure::ErrorAfterCommit, move || repeated_id).await; + let (parent, existing_thread) = harness.start_thread().await; + assert_eq!(parent, repeated_id); + let children_before = graph + .list_thread_spawn_children(parent, /*status_filter*/ None) + .await + .expect("initial graph"); + let error = spawn_publication_child(&harness, parent) + .await + .expect_err("loaded IDs are rejected"); + assert_eq!( + (error.agent_context(), CodexErrKind::from(&error)), + ( + Some(AgentErrorContext::ChildStartup), + CodexErrKind::InvalidRequest + ), + ); + assert!(Arc::ptr_eq( + &existing_thread, + &harness + .manager + .get_thread(parent) + .await + .expect("existing thread survives") + )); + assert!( + !graph.started.is_cancelled(), + "duplicate ID must be rejected before publication" + ); + assert!( + !graph.close_attempted.is_cancelled(), + "existing thread is not provisional cleanup" + ); + assert_eq!( + graph + .list_thread_spawn_children(parent, /*status_filter*/ None) + .await + .expect("graph after rejection"), + children_before + ); + existing_thread + .shutdown_and_wait() + .await + .expect("stop existing thread"); +} diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 86eaf38a4e..87247661d2 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -124,6 +124,9 @@ use tokio::time::timeout; use tokio_util::sync::CancellationToken; use toml::Value as TomlValue; +#[path = "control/spawn_publication_tests.rs"] +mod spawn_publication; + impl LocalAgentControl { /// Create a child fixture through the production spawn entry point. pub(crate) async fn spawn_agent_with_metadata( @@ -784,11 +787,12 @@ async fn spawn_failure_context_survives_manager_and_fork_errors() { async fn spawn_failure_context_preserves_shutdown_during_child_startup() { let harness = AgentControlHarness::new().await; let state = harness.control.runtime.upgrade().expect("live manager"); + let parent_id = ThreadId::new(); let spawn = harness.control.spawn_agent( harness.config.clone(), text_input("hello"), Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id: ThreadId::new(), + parent_thread_id: parent_id, depth: 1, agent_path: None, agent_nickname: None, @@ -827,6 +831,22 @@ async fn spawn_failure_context_preserves_shutdown_during_child_startup() { .wait() .await .expect("spawn releases its membership"); + let closed_children = harness + .state_db + .as_ref() + .expect("manager should retain state db") + .list_thread_spawn_children_with_status( + parent_id, + codex_state::DirectionalThreadSpawnEdgeStatus::Closed, + ) + .await + .expect("closed child edges should load"); + assert_eq!( + closed_children.len(), + 1, + "failed startup closes its published edge" + ); + assert_thread_not_loaded(&harness.manager, closed_children[0]).await; } #[tokio::test] @@ -890,6 +910,20 @@ async fn cancelled_spawn_after_input_admission_discards_child() { .wait() .await .expect("provisional cleanup releases tree membership"); + assert_eq!( + harness + .state_db + .as_ref() + .expect("manager should retain state db") + .list_thread_spawn_children_with_status( + parent_id, + codex_state::DirectionalThreadSpawnEdgeStatus::Closed, + ) + .await + .expect("closed child edge should load"), + vec![child_id], + "cancelled startup must close its published edge" + ); } #[tokio::test] diff --git a/codex-rs/core/tests/common/Cargo.toml b/codex-rs/core/tests/common/Cargo.toml index 5644a06e0a..cb5c222039 100644 --- a/codex-rs/core/tests/common/Cargo.toml +++ b/codex-rs/core/tests/common/Cargo.toml @@ -15,6 +15,7 @@ workspace = true anyhow = { workspace = true } assert_cmd = { workspace = true } base64 = { workspace = true } +codex-agent-graph-store = { workspace = true } codex-analytics = { workspace = true } codex-arg0 = { workspace = true } codex-attachment-store = { workspace = true } diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 77e9a03ece..89081ae5ca 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -15,6 +15,7 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; use anyhow::anyhow; +use codex_agent_graph_store::AgentGraphStore; use codex_analytics::AnalyticsEventsClient; use codex_attachment_store::AttachmentStore; use codex_config::CloudConfigBundleLoader; @@ -424,6 +425,7 @@ pub struct TestCodexBuilder { history_mode: Option, models_manager: Option, thread_store: Option>, + agent_graph_store: Option>, image_store: Arc, } @@ -438,6 +440,11 @@ impl TestCodexBuilder { self } + pub fn with_agent_graph_store(mut self, agent_graph_store: Arc) -> Self { + self.agent_graph_store = Some(agent_graph_store); + self + } + pub fn with_thread_manager( mut self, configure: impl FnOnce(ThreadManager) -> ThreadManager + Send + 'static, @@ -877,7 +884,9 @@ impl TestCodexBuilder { self.analytics_events_client.clone(), Arc::clone(&self.image_store), Arc::clone(&thread_store), - codex_core::local_agent_graph_store_from_state_db(state_db.as_ref()), + self.agent_graph_store.clone().or_else(|| { + codex_core::local_agent_graph_store_from_state_db(state_db.as_ref()) + }), installation_id, /*attestation_provider*/ None, /*external_time_provider*/ self.external_time_provider.clone(), @@ -1558,6 +1567,7 @@ pub fn test_codex() -> TestCodexBuilder { history_mode: Some(ThreadHistoryMode::Legacy), models_manager: None, thread_store: None, + agent_graph_store: None, image_store: codex_core::passthrough_image_store(), } } diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index c75dd859e1..ec628fd960 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -1,4 +1,9 @@ use anyhow::Result; +use assert_matches::assert_matches; +use codex_agent_graph_store::AgentGraphStore; +use codex_agent_graph_store::AgentGraphStoreError; +use codex_agent_graph_store::AgentGraphStoreFuture; +use codex_agent_graph_store::ThreadSpawnEdgeStatus; use codex_config::ConfigLayerEntry; use codex_config::ConfigLayerSource; use codex_config::ConfigLayerStack; @@ -35,7 +40,11 @@ use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::user_input::UserInput; use codex_skills_extension::SkillsExtensionConfig; +use codex_thread_store::DeleteThreadsParams; use codex_thread_store::LoadThreadHistoryParams; +use codex_thread_store::LocalThreadStore; +use codex_thread_store::LocalThreadStoreConfig; +use codex_thread_store::ThreadStoreError; use core_test_support::hooks::trust_discovered_hooks; use core_test_support::responses::ResponsesRequest; use core_test_support::responses::assert_parent_turn; @@ -3452,6 +3461,188 @@ async fn spawn_agent_preserves_configured_defaults_through_unrelated_role() -> R Ok(()) } +#[derive(Default)] +struct FailingSpawnEdgeStore { + child: Mutex>, +} + +impl AgentGraphStore for FailingSpawnEdgeStore { + fn upsert_thread_spawn_edge( + &self, + _parent: ThreadId, + child: ThreadId, + _status: ThreadSpawnEdgeStatus, + ) -> AgentGraphStoreFuture<'_, ()> { + *self.child.lock().expect("lock attempted child ID") = Some(child); + Box::pin(async { + Err(AgentGraphStoreError::Internal { + message: "injected edge publication failure".to_string(), + }) + }) + } + + fn set_thread_spawn_edge_status( + &self, + _child: ThreadId, + _status: ThreadSpawnEdgeStatus, + ) -> AgentGraphStoreFuture<'_, ()> { + Box::pin(async { Ok(()) }) + } + + fn list_thread_spawn_children( + &self, + _parent: ThreadId, + _status: Option, + ) -> AgentGraphStoreFuture<'_, Vec> { + Box::pin(async { Ok(Vec::new()) }) + } + + fn list_thread_spawn_descendants( + &self, + _root: ThreadId, + _status: Option, + ) -> AgentGraphStoreFuture<'_, Vec> { + Box::pin(async { Ok(Vec::new()) }) + } +} + +#[test_case(false; "fresh child")] +#[test_case(true; "full-history fork")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn failed_spawn_edge_leaves_no_child_history(fork_context: bool) -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let spawn_args = json!({ "message": CHILD_PROMPT, "fork_context": fork_context }).to_string(); + mount_sse_once_match( + &server, + |request: &wiremock::Request| body_contains(request, TURN_1_PROMPT), + sse(vec![ + ev_response_created("spawn"), + ev_function_call_with_namespace( + SPAWN_CALL_ID, + MULTI_AGENT_V1_NAMESPACE, + "spawn_agent", + &spawn_args, + ), + ev_completed("spawn"), + ]), + ) + .await; + let tool_output = mount_sse_once_match( + &server, + |request: &wiremock::Request| body_contains(request, SPAWN_CALL_ID), + sse(vec![ + ev_response_created("parent-done"), + ev_completed("parent-done"), + ]), + ) + .await; + let child_requests = mount_sse_once_match( + &server, + |request: &wiremock::Request| { + body_contains(request, CHILD_PROMPT) && !body_contains(request, SPAWN_CALL_ID) + }, + sse(vec![ + ev_response_created("child-done"), + ev_completed("child-done"), + ]), + ) + .await; + + let home = Arc::new(tempfile::tempdir()?); + // Keep history independent of topology; SQLite metadata can repair missing edges. + let store = Arc::new(LocalThreadStore::new( + LocalThreadStoreConfig { + codex_home: home.path().to_path_buf(), + sqlite: codex_state::SqliteConfig::new_for_testing(home.path().try_into()?), + default_model_provider_id: "openai".to_string(), + }, + /*state_db*/ None, + )); + let graph = Arc::new(FailingSpawnEdgeStore::default()); + let test = test_codex() + .with_home(home) + .with_thread_store(store) + .with_agent_graph_store(graph.clone()) + .with_history_mode(ThreadHistoryMode::Paginated) + .with_config(|config| { + config + .features + .enable(Feature::Collab) + .expect("enable native spawn"); + }) + .build_with_auto_env(&server) + .await?; + let unrelated = test + .thread_manager + .start_thread(test.start_thread_options().await) + .await?; + unrelated.thread.ensure_rollout_materialized().await; + unrelated.thread.flush_rollout().await?; + let mut created = test.thread_manager.subscribe_thread_created(); + test.submit_turn(TURN_1_PROMPT).await?; + let child_id = graph + .child + .lock() + .expect("lock attempted child ID") + .expect("edge was attempted"); + test.thread_manager + .shutdown_all_threads_bounded(Duration::from_secs(10)) + .await; + + let child_before_delete = test + .thread_store + .load_latest_model_context(LoadThreadHistoryParams { + thread_id: child_id, + include_archived: true, + }) + .await; + let parent_id = test.session_configured.thread_id; + let deletion_ids = test + .thread_manager + .list_agent_subtree_thread_ids(parent_id) + .await?; + assert_eq!(deletion_ids, vec![parent_id]); + test.thread_store + .delete_threads(DeleteThreadsParams { + thread_ids: deletion_ids, + }) + .await?; + test.thread_store + .load_latest_model_context(LoadThreadHistoryParams { + thread_id: unrelated.thread_id, + include_archived: true, + }) + .await + .expect("unrelated history survives parent deletion"); + for thread_id in [parent_id, child_id] { + assert_matches!(test.thread_store.load_latest_model_context(LoadThreadHistoryParams { + thread_id, + include_archived: true, + }).await, Err(ThreadStoreError::InvalidRequest { message }) + if message == format!("no rollout found for thread id {thread_id}")); + } + assert_matches!(child_before_delete, Err(ThreadStoreError::InvalidRequest { message }) + if message == format!("no rollout found for thread id {child_id}"), + "rejected child must not leave durable history"); + assert_eq!( + tool_output + .function_call_output_text(SPAWN_CALL_ID) + .as_deref(), + Some("collab spawn failed: Fatal error: failed to persist spawned agent relationship") + ); + assert!( + child_requests.requests().is_empty(), + "rejected child must not execute" + ); + assert!( + created.try_recv().is_err(), + "rejected child must not be published" + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn spawn_agent_rejects_reasoning_effort_unsupported_by_role_model() -> Result<()> { skip_if_no_network!(Ok(()));