mirror of
https://github.com/openai/codex.git
synced 2026-10-30 19:18:49 +00:00
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:
Generated
+1
@@ -5894,6 +5894,7 @@ dependencies = [
|
||||
"anyhow",
|
||||
"assert_cmd",
|
||||
"base64 0.22.1",
|
||||
"codex-agent-graph-store",
|
||||
"codex-analytics",
|
||||
"codex-arg0",
|
||||
"codex-attachment-store",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
@@ -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]
|
||||
|
||||
@@ -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 }
|
||||
|
||||
@@ -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(()));
|
||||
|
||||
Reference in New Issue
Block a user