Persist subagent relationships before creating child storage (#58129)

## Why

Creating child history before persisting its parent-child relationship can leave history that parent subtree deletion cannot discover, especially when fork startup copies history. A relationship write failure must reject the spawn before child storage is created.

## What changed

- Reserve the child thread ID and persist its spawn edge before starting the child. Propagate persistence failures to the caller.
- Keep pending relationship writes owned through caller cancellation, and wait for them to finish before closing the edge during cleanup, including when a write commits but reports an error.
- Skip edge cleanup when publication was skipped, and reject IDs belonging to running threads before publication.

## Testing

Add regression coverage for cancellation before commit or acknowledgement, errors after commit, skipped publication, and duplicate child IDs. Integration tests cover fresh children and full-history forks, verifying that failed publication leaves no child history, execution, or creation notification and preserves unrelated history during parent deletion.

GitOrigin-RevId: 6e2e3f32d9558091312651ab47aa845d7d13b4ca
This commit is contained in:
radai-oai
2026-10-29 17:25:51 +00:00
committed by copyberry
parent 468b71f7a6
commit 875a1a78f0
10 changed files with 667 additions and 56 deletions
+1
View File
@@ -5894,6 +5894,7 @@ dependencies = [
"anyhow",
"assert_cmd",
"base64 0.22.1",
"codex-agent-graph-store",
"codex-analytics",
"codex-arg0",
"codex-attachment-store",
+2
View File
@@ -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,
+12 -9
View File
@@ -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<bool> {
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)
}
}
+36 -33
View File
@@ -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(&notification_source),
)
.await;
let _ = self
.persist_thread_spawn_edge_for_source(
resumed_thread.thread.config_snapshot().await.ephemeral,
resumed_thread.thread_id,
Some(&notification_source),
)
.await;
Ok((resumed_thread.thread_id, multi_agent_version))
}
+22 -12
View File
@@ -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<ThreadManagerState>,
pending: Option<(ThreadId, AgentTreeMembership)>,
edge_write: Option<JoinHandle<()>>,
edge_write: Option<JoinHandle<CodexResult<bool>>>,
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<CodexResult<bool>>) {
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)
@@ -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<ThreadSpawnEdgeStatus>,
) -> AgentGraphStoreFuture<'_, Vec<ThreadId>> {
self.inner.list_thread_spawn_children(parent, status)
}
fn list_thread_spawn_descendants(
&self,
root: ThreadId,
status: Option<ThreadSpawnEdgeStatus>,
) -> AgentGraphStoreFuture<'_, Vec<ThreadId>> {
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<FaultingGraphStore>) {
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<ThreadId> {
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::<ThreadId>::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::<ThreadId>::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");
}
+35 -1
View File
@@ -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]
+1
View File
@@ -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 }
+11 -1
View File
@@ -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<ThreadHistoryMode>,
models_manager: Option<SharedModelsManager>,
thread_store: Option<Arc<dyn ThreadStore>>,
agent_graph_store: Option<Arc<dyn AgentGraphStore>>,
image_store: Arc<dyn AttachmentStore>,
}
@@ -438,6 +440,11 @@ impl TestCodexBuilder {
self
}
pub fn with_agent_graph_store(mut self, agent_graph_store: Arc<dyn AgentGraphStore>) -> 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(),
}
}
@@ -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<Option<ThreadId>>,
}
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<ThreadSpawnEdgeStatus>,
) -> AgentGraphStoreFuture<'_, Vec<ThreadId>> {
Box::pin(async { Ok(Vec::new()) })
}
fn list_thread_spawn_descendants(
&self,
_root: ThreadId,
_status: Option<ThreadSpawnEdgeStatus>,
) -> AgentGraphStoreFuture<'_, Vec<ThreadId>> {
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(()));