Inject state DB, agent graph store (#20689)

## Why

We want the agent graph store to be passed down the stack as a real
dependency, the same way we already treat the thread store.

This will let us inject the agent graph store as a real dependency and
support implementations other than the local SQLite-backed one. Right
now most code instantiates a state DB and an agent graph store
just-in-time. Ideally, we would not depend on the state DB directly but
only read through the higher-level interfaces.

This change makes the dependency boundaries explicit and moves state DB
initialization to process bootstrap instead of hiding it inside local
store implementations.

## What changed

- `ThreadManager` now requires a `StateDbHandle` and an
`AgentGraphStore` at construction time instead of treating them as
optional internals.
- The local store constructors no longer lazily initialize SQLite.
Callers now initialize the state DB once per process and use that shared
handle to build:
  - `LocalThreadStore`
  - `LocalAgentGraphStore`
- App bootstraps (`app-server`, `mcp-server`, `prompt_debug`, and the
thread-manager sample) now initialize the state DB up front and inject
the resulting handle down the stack.
- `app-server` now consistently uses its process-scoped state DB handle
instead of reopening SQLite or trying to recover it from loaded threads.
- Device-key storage now reuses the shared state DB handle instead of
maintaining its own lazy opener.
- The thread archive / descendant traversal paths now use the injected
`AgentGraphStore` instead of reaching through local
thread-store-specific state.

## Verification

- `cargo check -p codex-core -p codex-thread-store -p codex-app-server
-p codex-mcp-server -p codex-thread-manager-sample --tests`
- `cargo test -p codex-thread-store`
- `cargo test -p codex-core
thread_manager_accepts_separate_agent_graph_store_and_thread_store --
--nocapture`
- `cargo test -p codex-app-server
thread_archive_archives_spawned_descendants -- --nocapture`
This commit is contained in:
Rasmus Rygaard
2026-05-05 14:45:29 -07:00
committed by GitHub
Unverified
parent 36460387ec
commit 7e310bc7f3
53 changed files with 768 additions and 745 deletions
+1
View File
@@ -2411,6 +2411,7 @@ dependencies = [
"bm25",
"chrono",
"clap",
"codex-agent-graph-store",
"codex-analytics",
"codex-api",
"codex-app-server-protocol",
+1 -1
View File
@@ -29,7 +29,6 @@ pub use codex_app_server::in_process::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY;
pub use codex_app_server::in_process::InProcessServerEvent;
use codex_app_server::in_process::InProcessStartArgs;
use codex_app_server::in_process::LogDbLayer;
pub use codex_app_server::in_process::StateDbHandle;
use codex_app_server_protocol::ClientInfo;
use codex_app_server_protocol::ClientNotification;
use codex_app_server_protocol::ClientRequest;
@@ -47,6 +46,7 @@ use codex_config::LoaderOverrides;
use codex_config::NoopThreadConfigLoader;
use codex_config::RemoteThreadConfigLoader;
use codex_config::ThreadConfigLoader;
pub use codex_core::StateDbHandle;
use codex_core::config::Config;
pub use codex_exec_server::EnvironmentManager;
pub use codex_exec_server::EnvironmentManagerArgs;
@@ -2629,7 +2629,8 @@ mod tests {
config.model_provider.clone(),
config.codex_home.to_path_buf(),
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
),
)
.await,
);
let codex_core::NewThread {
thread_id: conversation_id,
@@ -3214,7 +3215,8 @@ mod tests {
config.model_provider.clone(),
config.codex_home.to_path_buf(),
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
),
)
.await,
);
let codex_core::NewThread {
thread_id: conversation_id,
+18 -7
View File
@@ -82,11 +82,12 @@ use codex_config::CloudRequirementsLoader;
use codex_config::LoaderOverrides;
use codex_config::ThreadConfigLoader;
use codex_core::config::Config;
use codex_core::init_state_db_from_config;
use codex_exec_server::EnvironmentManager;
use codex_feedback::CodexFeedback;
use codex_login::AuthManager;
use codex_protocol::protocol::SessionSource;
pub use codex_rollout::StateDbHandle;
use codex_rollout::state_db::StateDbHandle;
pub use codex_state::log_db::LogDbLayer;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
@@ -127,7 +128,7 @@ pub struct InProcessStartArgs {
pub feedback: CodexFeedback,
/// SQLite tracing layer used to flush recently emitted logs before feedback upload.
pub log_db: Option<LogDbLayer>,
/// Process-wide SQLite state handle shared with embedded app-server consumers.
/// Optional state DB handle to use for the in-process runtime.
pub state_db: Option<StateDbHandle>,
/// Environment manager used by core execution and filesystem operations.
pub environment_manager: Arc<EnvironmentManager>,
@@ -344,7 +345,7 @@ impl InProcessClientHandle {
/// the runtime is shut down and an `InvalidData` error is returned.
pub async fn start(args: InProcessStartArgs) -> IoResult<InProcessClientHandle> {
let initialize = args.initialize.clone();
let client = start_uninitialized(args);
let client = start_uninitialized(args).await;
let initialize_response = client
.request(ClientRequest::Initialize {
@@ -364,8 +365,12 @@ pub async fn start(args: InProcessStartArgs) -> IoResult<InProcessClientHandle>
Ok(client)
}
fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
async fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
let channel_capacity = args.channel_capacity.max(1);
let state_db = match args.state_db.clone() {
Some(state_db) => Some(state_db),
None => init_state_db_from_config(args.config.as_ref()).await,
};
let (client_tx, mut client_rx) = mpsc::channel::<InProcessClientMessage>(channel_capacity);
let (event_tx, event_rx) = mpsc::channel::<InProcessServerEvent>(channel_capacity);
@@ -414,6 +419,12 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
);
let (processor_tx, mut processor_rx) = mpsc::channel::<ProcessorCommand>(channel_capacity);
let mut processor_handle = tokio::spawn(async move {
let Some(state_db) = state_db else {
warn!(
"in-process app-server state db initialization failed; shutting down processor task"
);
return;
};
let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs {
outgoing: Arc::clone(&processor_outgoing),
analytics_events_client,
@@ -423,7 +434,7 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
environment_manager: args.environment_manager,
feedback: args.feedback,
log_db: args.log_db,
state_db: args.state_db,
state_db,
config_warnings: args.config_warnings,
session_source: args.session_source,
auth_manager,
@@ -761,7 +772,7 @@ mod tests {
) -> InProcessClientHandle {
let codex_home = TempDir::new().expect("temp dir");
let config = Arc::new(build_test_config(codex_home.path()).await);
let state_db = codex_rollout::state_db::try_init(config.as_ref())
let state_db = init_state_db_from_config(config.as_ref())
.await
.expect("state db should initialize for in-process test");
let args = InProcessStartArgs {
@@ -819,7 +830,7 @@ mod tests {
}
#[tokio::test]
async fn in_process_allows_device_key_requests_to_reach_device_key_processor() {
async fn in_process_allows_device_key_requests_to_reach_device_key_api() {
let client = start_test_client(SessionSource::Cli).await;
const MALFORMED_KEY_ID_MESSAGE: &str = concat!(
"invalid device key payload: keyId must be dk_hse_, dk_tpm_, or dk_osn_ ",
+14 -24
View File
@@ -50,11 +50,11 @@ use codex_config::TextRange as CoreTextRange;
use codex_core::ExecPolicyError;
use codex_core::check_execpolicy_for_warnings;
use codex_core::config::find_codex_home;
use codex_core::init_state_db_from_config;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::ExecServerRuntimePaths;
use codex_feedback::CodexFeedback;
use codex_protocol::protocol::SessionSource;
use codex_rollout::state_db as rollout_state_db;
use codex_state::log_db;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
@@ -487,9 +487,9 @@ pub async fn run_main_with_transport_options(
}
};
let state_db_result = rollout_state_db::try_init(&config).await;
let state_db_init_error = state_db_result.as_ref().err().map(ToString::to_string);
let state_db = state_db_result.ok();
let state_db = init_state_db_from_config(&config)
.await
.ok_or_else(|| std::io::Error::other("failed to initialize sqlite state db"))?;
if should_run_personality_migration {
let effective_toml = config.config_layer_stack.effective_config();
@@ -598,10 +598,12 @@ pub async fn run_main_with_transport_options(
let feedback_layer = feedback.logger_layer();
let feedback_metadata_layer = feedback.metadata_layer();
let log_db = state_db.clone().map(log_db::start);
let log_db_layer = log_db
.clone()
.map(|layer| layer.with_filter(Targets::new().with_default(Level::TRACE)));
let log_db = log_db::start(state_db.clone());
let log_db_layer = Some(
log_db
.clone()
.with_filter(Targets::new().with_default(Level::TRACE)),
);
let otel_logger_layer = otel.as_ref().and_then(|o| o.logger_layer());
let otel_tracing_layer = otel.as_ref().and_then(|o| o.tracing_layer());
let _ = tracing_subscriber::registry()
@@ -618,10 +620,6 @@ pub async fn run_main_with_transport_options(
None => error!("{}", warning.summary),
}
}
if let Some(err) = &state_db_init_error {
error!("failed to initialize sqlite state db: {err}");
}
let transport_shutdown_token = CancellationToken::new();
let mut transport_accept_handles = Vec::<JoinHandle<()>>::new();
@@ -666,25 +664,17 @@ pub async fn run_main_with_transport_options(
let auth_manager =
AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false).await;
let remote_control_config_enabled = config.features.enabled(Feature::RemoteControl);
let remote_control_enabled = remote_control_config_enabled && state_db.is_some();
if remote_control_config_enabled && state_db.is_none() {
error!("remote control disabled because sqlite state db is unavailable");
}
let remote_control_enabled = config.features.enabled(Feature::RemoteControl);
if transport_accept_handles.is_empty() && !remote_control_enabled {
return Err(std::io::Error::new(
ErrorKind::InvalidInput,
if remote_control_config_enabled && state_db.is_none() {
"no transport configured; remote control disabled because sqlite state db is unavailable"
} else {
"no transport configured; use --listen or enable remote control"
},
"no transport configured; use --listen or enable remote control",
));
}
let (remote_control_accept_handle, remote_control_handle) = start_remote_control(
config.chatgpt_base_url.clone(),
state_db.clone(),
Some(state_db.clone()),
auth_manager.clone(),
transport_event_tx.clone(),
transport_shutdown_token.clone(),
@@ -768,7 +758,7 @@ pub async fn run_main_with_transport_options(
config_manager,
environment_manager,
feedback: feedback.clone(),
log_db,
log_db: Some(log_db),
state_db: state_db.clone(),
config_warnings,
session_source,
+7 -4
View File
@@ -61,6 +61,7 @@ use codex_app_server_protocol::experimental_required_message;
use codex_arg0::Arg0DispatchPaths;
use codex_chatgpt::workspace_settings;
use codex_core::ThreadManager;
use codex_core::agent_graph_store_from_state_db;
use codex_core::config::Config;
use codex_core::thread_store_from_config;
use codex_exec_server::EnvironmentManager;
@@ -254,7 +255,7 @@ pub(crate) struct MessageProcessorArgs {
pub(crate) environment_manager: Arc<EnvironmentManager>,
pub(crate) feedback: CodexFeedback,
pub(crate) log_db: Option<LogDbLayer>,
pub(crate) state_db: Option<StateDbHandle>,
pub(crate) state_db: StateDbHandle,
pub(crate) config_warnings: Vec<ConfigWarningNotification>,
pub(crate) session_source: SessionSource,
pub(crate) auth_manager: Arc<AuthManager>,
@@ -291,14 +292,16 @@ impl MessageProcessor {
// affect per-thread behavior, but they must not move newly started,
// resumed, or forked threads to a different persistence backend/root.
let thread_store = thread_store_from_config(config.as_ref(), state_db.clone());
let agent_graph_store = agent_graph_store_from_state_db(state_db.clone());
let thread_manager = Arc::new(ThreadManager::new(
config.as_ref(),
auth_manager.clone(),
session_source,
environment_manager,
Some(analytics_events_client.clone()),
Arc::clone(&thread_store),
state_db.clone(),
Arc::clone(&thread_store),
agent_graph_store.clone(),
));
thread_manager
.plugins_manager()
@@ -344,7 +347,7 @@ impl MessageProcessor {
Arc::clone(&config),
feedback,
log_db,
state_db.clone(),
Some(state_db.clone()),
);
let git_processor = GitRequestProcessor::new();
let initialize_processor = InitializeRequestProcessor::new(
@@ -395,7 +398,7 @@ impl MessageProcessor {
thread_watch_manager.clone(),
Arc::clone(&thread_list_state_permit),
thread_goal_processor.clone(),
state_db.clone(),
Some(state_db.clone()),
);
let turn_processor = TurnRequestProcessor::new(
auth_manager.clone(),
@@ -32,6 +32,7 @@ use codex_config::CloudRequirementsLoader;
use codex_config::LoaderOverrides;
use codex_core::config::Config;
use codex_core::config::ConfigBuilder;
use codex_core::init_state_db_from_config;
use codex_exec_server::EnvironmentManager;
use codex_feedback::CodexFeedback;
use codex_login::AuthManager;
@@ -281,6 +282,9 @@ async fn build_test_processor(
outgoing_tx,
analytics_events_client.clone(),
));
let state_db = init_state_db_from_config(config.as_ref())
.await
.expect("tracing test processor requires state db");
let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs {
outgoing,
analytics_events_client,
@@ -290,7 +294,7 @@ async fn build_test_processor(
environment_manager: Arc::new(EnvironmentManager::default_for_tests()),
feedback: CodexFeedback::new(),
log_db: None,
state_db: None,
state_db,
config_warnings: Vec::new(),
session_source: SessionSource::VSCode,
auth_manager,
@@ -33,8 +33,8 @@ use codex_device_key::RemoteControlClientConnectionAudience;
use codex_device_key::RemoteControlClientConnectionSignPayload;
use codex_device_key::RemoteControlClientEnrollmentAudience;
use codex_device_key::RemoteControlClientEnrollmentSignPayload;
use codex_rollout::state_db::StateDbHandle;
use codex_state::DeviceKeyBindingRecord;
use codex_state::StateRuntime;
#[derive(Clone)]
pub(crate) struct DeviceKeyRequestProcessor {
@@ -43,10 +43,7 @@ pub(crate) struct DeviceKeyRequestProcessor {
}
impl DeviceKeyRequestProcessor {
pub(crate) fn new(
outgoing: Arc<OutgoingMessageSender>,
state_db: Option<Arc<StateRuntime>>,
) -> Self {
pub(crate) fn new(outgoing: Arc<OutgoingMessageSender>, state_db: StateDbHandle) -> Self {
Self {
outgoing,
store: DeviceKeyStore::new(Arc::new(StateDeviceKeyBindingStore::new(state_db))),
@@ -170,25 +167,18 @@ async fn sign_device_key(
}
struct StateDeviceKeyBindingStore {
state_db: Option<Arc<StateRuntime>>,
state_db: StateDbHandle,
}
impl StateDeviceKeyBindingStore {
fn new(state_db: Option<Arc<StateRuntime>>) -> Self {
fn new(state_db: StateDbHandle) -> Self {
Self { state_db }
}
async fn state_db(&self) -> Result<Arc<StateRuntime>, DeviceKeyError> {
self.state_db
.clone()
.ok_or_else(|| DeviceKeyError::Platform("sqlite state db unavailable".to_string()))
}
}
impl fmt::Debug for StateDeviceKeyBindingStore {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("StateDeviceKeyBindingStore")
.field("has_state_db", &self.state_db.is_some())
.finish_non_exhaustive()
}
}
@@ -196,7 +186,7 @@ impl fmt::Debug for StateDeviceKeyBindingStore {
#[async_trait]
impl DeviceKeyBindingStore for StateDeviceKeyBindingStore {
async fn get_binding(&self, key_id: &str) -> Result<Option<DeviceKeyBinding>, DeviceKeyError> {
let state_db = self.state_db().await?;
let state_db = self.state_db.clone();
state_db
.get_device_key_binding(key_id)
.await
@@ -214,7 +204,7 @@ impl DeviceKeyBindingStore for StateDeviceKeyBindingStore {
key_id: &str,
binding: &DeviceKeyBinding,
) -> Result<(), DeviceKeyError> {
let state_db = self.state_db().await?;
let state_db = self.state_db.clone();
state_db
.upsert_device_key_binding(&DeviceKeyBindingRecord {
key_id: key_id.to_string(),
@@ -7,7 +7,7 @@ pub(crate) struct ThreadGoalRequestProcessor {
outgoing: Arc<OutgoingMessageSender>,
config: Arc<Config>,
thread_state_manager: ThreadStateManager,
state_db: Option<StateDbHandle>,
state_db: StateDbHandle,
}
impl ThreadGoalRequestProcessor {
@@ -16,7 +16,7 @@ impl ThreadGoalRequestProcessor {
outgoing: Arc<OutgoingMessageSender>,
config: Arc<Config>,
thread_state_manager: ThreadStateManager,
state_db: Option<StateDbHandle>,
state_db: StateDbHandle,
) -> Self {
Self {
thread_manager,
@@ -72,23 +72,6 @@ impl ThreadGoalRequestProcessor {
}
}
pub(crate) async fn pending_resume_goal_state(
&self,
thread: &CodexThread,
) -> (bool, Option<StateDbHandle>) {
let emit_thread_goal_update = self.config.features.enabled(Feature::Goals);
let thread_goal_state_db = if emit_thread_goal_update {
if let Some(state_db) = thread.state_db() {
Some(state_db)
} else {
self.state_db.clone()
}
} else {
None
};
(emit_thread_goal_update, thread_goal_state_db)
}
async fn thread_goal_set_inner(
&self,
request_id: ConnectionRequestId,
@@ -110,7 +93,7 @@ impl ThreadGoalRequestProcessor {
None => find_thread_path_by_id_str(
&self.config.codex_home,
&thread_id.to_string(),
self.state_db.as_deref(),
Some(self.state_db.as_ref()),
)
.await
.map_err(|err| {
@@ -275,7 +258,7 @@ impl ThreadGoalRequestProcessor {
None => find_thread_path_by_id_str(
&self.config.codex_home,
&thread_id.to_string(),
self.state_db.as_deref(),
Some(self.state_db.as_ref()),
)
.await
.map_err(|err| {
@@ -339,7 +322,7 @@ impl ThreadGoalRequestProcessor {
find_thread_path_by_id_str(
&self.config.codex_home,
&thread_id.to_string(),
self.state_db.as_deref(),
Some(self.state_db.as_ref()),
)
.await
.map_err(|err| {
@@ -348,9 +331,7 @@ impl ThreadGoalRequestProcessor {
.ok_or_else(|| invalid_request(format!("thread not found: {thread_id}")))?;
}
self.state_db
.clone()
.ok_or_else(|| internal_error("sqlite state db unavailable for thread goals"))
Ok(self.state_db.clone())
}
async fn emit_thread_goal_snapshot(&self, thread_id: ThreadId) {
@@ -2611,10 +2611,10 @@ impl ThreadRequestProcessor {
)));
};
let (emit_thread_goal_update, thread_goal_state_db) = self
.thread_goal_processor
.pending_resume_goal_state(existing_thread.as_ref())
.await;
let emit_thread_goal_update = self.config.features.enabled(Feature::Goals);
let thread_goal_state_db = emit_thread_goal_update
.then(|| self.state_db.clone())
.flatten();
let command = crate::thread_state::ThreadListenerCommand::SendThreadResumeResponse(
Box::new(crate::thread_state::PendingThreadResumeRequest {
@@ -42,6 +42,7 @@ use codex_core::config::ConfigBuilder;
use codex_exec_server::EnvironmentManager;
use codex_feedback::CodexFeedback;
use codex_protocol::protocol::SessionSource;
use codex_state::StateRuntime;
use codex_thread_store::InMemoryThreadStore;
use pretty_assertions::assert_eq;
use tempfile::TempDir;
@@ -67,6 +68,13 @@ async fn thread_start_with_non_local_thread_store_does_not_create_local_persiste
.loader_overrides(loader_overrides.clone())
.build()
.await?;
let sqlite_home = TempDir::new()?;
let state_db = StateRuntime::init(
sqlite_home.path().to_path_buf(),
config.model_provider_id.clone(),
)
.await
.expect("remote thread store regression test should initialize state db");
let thread_store = InMemoryThreadStore::for_id(store_id.clone());
let _in_memory_store = InMemoryThreadStoreId { store_id };
@@ -80,7 +88,7 @@ async fn thread_start_with_non_local_thread_store_does_not_create_local_persiste
thread_config_loader: Arc::new(NoopThreadConfigLoader),
feedback: CodexFeedback::new(),
log_db: None,
state_db: None,
state_db: Some(state_db),
environment_manager: Arc::new(EnvironmentManager::default_for_tests()),
config_warnings: Vec::new(),
session_source: SessionSource::Cli,
+1 -1
View File
@@ -1387,7 +1387,7 @@ async fn run_debug_prompt_input_command(
});
}
let prompt_input = codex_core::build_prompt_input(config, input, /*state_db*/ None).await?;
let prompt_input = codex_core::build_prompt_input(config, input).await?;
println!("{}", serde_json::to_string_pretty(&prompt_input)?);
Ok(())
+2
View File
@@ -31,6 +31,7 @@ pub use codex_core::StartThreadOptions;
pub use codex_core::StateDbHandle;
pub use codex_core::ThreadManager;
pub use codex_core::ThreadShutdownReport;
pub use codex_core::agent_graph_store_from_state_db;
pub use codex_core::config::Config;
pub use codex_core::config::Constrained;
pub use codex_core::config::GhostSnapshotConfig;
@@ -40,6 +41,7 @@ pub use codex_core::config::TerminalResizeReflowConfig;
pub use codex_core::config::ThreadStoreConfig;
pub use codex_core::config::find_codex_home;
pub use codex_core::init_state_db;
pub use codex_core::init_state_db_from_config;
pub use codex_core::skills::SkillsManager;
pub use codex_core::thread_store_from_config;
pub use codex_exec_server::EnvironmentManager;
+1
View File
@@ -26,6 +26,7 @@ bm25 = { workspace = true }
chrono = { workspace = true, features = ["serde"] }
clap = { workspace = true, features = ["derive"] }
codex-analytics = { workspace = true }
codex-agent-graph-store = { workspace = true }
codex-api = { workspace = true }
codex-app-server-protocol = { workspace = true }
codex-apply-patch = { workspace = true }
+19 -29
View File
@@ -29,7 +29,6 @@ use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_protocol::user_input::UserInput;
use codex_state::DirectionalThreadSpawnEdgeStatus;
use codex_thread_store::ReadThreadParams;
use serde::Serialize;
use std::collections::HashMap;
@@ -296,7 +295,6 @@ impl AgentControl {
state.notify_thread_created(new_thread.thread_id);
self.persist_thread_spawn_edge_for_source(
new_thread.thread.as_ref(),
new_thread.thread_id,
notification_source.as_ref(),
)
@@ -449,19 +447,14 @@ impl AgentControl {
))
.await?;
let state = self.upgrade()?;
let Ok(resumed_thread) = state.get_thread(resumed_thread_id).await else {
return Ok(resumed_thread_id);
};
let Some(state_db_ctx) = resumed_thread.state_db() else {
return Ok(resumed_thread_id);
};
let agent_graph_store = state.agent_graph_store();
let mut resume_queue = VecDeque::from([(thread_id, root_depth)]);
while let Some((parent_thread_id, parent_depth)) = resume_queue.pop_front() {
let child_ids = match state_db_ctx
.list_thread_spawn_children_with_status(
let child_ids = match agent_graph_store
.list_thread_spawn_children(
parent_thread_id,
DirectionalThreadSpawnEdgeStatus::Open,
Some(codex_agent_graph_store::ThreadSpawnEdgeStatus::Open),
)
.await
{
@@ -525,7 +518,6 @@ impl AgentControl {
let _ = config.features.disable(Feature::Collab);
}
let state = self.upgrade()?;
let state_db_ctx = state.state_db();
let mut reservation = self.state.reserve_spawn_slot(config.agent_max_threads)?;
let (session_source, agent_metadata) = match session_source {
SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
@@ -535,14 +527,11 @@ impl AgentControl {
agent_role: _,
agent_nickname: _,
}) => {
let state_db_ctx = state.state_db();
let (resumed_agent_nickname, resumed_agent_role) =
if let Some(state_db_ctx) = state_db_ctx.as_ref() {
match state_db_ctx.get_thread(thread_id).await {
Ok(Some(metadata)) => (metadata.agent_nickname, metadata.agent_role),
Ok(None) | Err(_) => (None, None),
}
} else {
(None, None)
match state_db_ctx.get_thread(thread_id).await {
Ok(Some(metadata)) => (metadata.agent_nickname, metadata.agent_role),
Ok(None) | Err(_) => (None, None),
};
self.prepare_thread_spawn(
&mut reservation,
@@ -609,7 +598,6 @@ impl AgentControl {
);
}
self.persist_thread_spawn_edge_for_source(
resumed_thread.thread.as_ref(),
resumed_thread.thread_id,
Some(&notification_source),
)
@@ -722,11 +710,13 @@ impl AgentControl {
/// agent and any live descendants reached from the in-memory tree.
pub(crate) async fn close_agent(&self, agent_id: ThreadId) -> CodexResult<String> {
let state = self.upgrade()?;
if let Ok(thread) = state.get_thread(agent_id).await
&& let Some(state_db_ctx) = thread.state_db()
&& let Err(err) = state_db_ctx
.set_thread_spawn_edge_status(agent_id, DirectionalThreadSpawnEdgeStatus::Closed)
.await
if let Err(err) = state
.agent_graph_store()
.set_thread_spawn_edge_status(
agent_id,
codex_agent_graph_store::ThreadSpawnEdgeStatus::Closed,
)
.await
{
warn!("failed to persist thread-spawn edge status for {agent_id}: {err}");
}
@@ -1142,21 +1132,21 @@ impl AgentControl {
async fn persist_thread_spawn_edge_for_source(
&self,
thread: &crate::CodexThread,
child_thread_id: ThreadId,
session_source: Option<&SessionSource>,
) {
let Some(parent_thread_id) = session_source.and_then(thread_spawn_parent_thread_id) else {
return;
};
let Some(state_db_ctx) = thread.state_db() else {
let Ok(state) = self.upgrade() else {
return;
};
if let Err(err) = state_db_ctx
if let Err(err) = state
.agent_graph_store()
.upsert_thread_spawn_edge(
parent_thread_id,
child_thread_id,
DirectionalThreadSpawnEdgeStatus::Open,
codex_agent_graph_store::ThreadSpawnEdgeStatus::Open,
)
.await
{
+22 -19
View File
@@ -1,6 +1,5 @@
use super::*;
use crate::CodexThread;
use crate::StateDbHandle;
use crate::ThreadManager;
use crate::agent::agent_status_from_event;
use crate::config::AgentRoleConfig;
@@ -8,7 +7,6 @@ use crate::config::Config;
use crate::config::ConfigBuilder;
use crate::context::ContextualUserFragment;
use crate::context::SubagentNotification;
use crate::init_state_db;
use assert_matches::assert_matches;
use codex_features::Feature;
use codex_login::CodexAuth;
@@ -86,7 +84,6 @@ fn spawn_agent_call(call_id: &str) -> ResponseItem {
struct AgentControlHarness {
_home: TempDir,
config: Config,
state_db: Option<StateDbHandle>,
manager: ThreadManager,
control: AgentControl,
}
@@ -94,19 +91,17 @@ struct AgentControlHarness {
impl AgentControlHarness {
async fn new() -> Self {
let (home, config) = test_config().await;
let state_db = init_state_db(&config).await;
let manager = ThreadManager::with_models_provider_home_and_state_for_tests(
let manager = ThreadManager::with_models_provider_and_home_for_tests(
CodexAuth::from_api_key("dummy"),
config.model_provider.clone(),
config.codex_home.to_path_buf(),
std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
state_db.clone(),
);
)
.await;
let control = manager.agent_control();
Self {
_home: home,
config,
state_db,
manager,
control,
}
@@ -955,7 +950,8 @@ async fn spawn_agent_respects_max_threads_limit() {
config.model_provider.clone(),
config.codex_home.to_path_buf(),
std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
);
)
.await;
let control = manager.agent_control();
let _ = manager
@@ -1007,7 +1003,8 @@ async fn spawn_agent_releases_slot_after_shutdown() {
config.model_provider.clone(),
config.codex_home.to_path_buf(),
std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
);
)
.await;
let control = manager.agent_control();
let first_agent_id = control
@@ -1050,7 +1047,8 @@ async fn spawn_agent_limit_shared_across_clones() {
config.model_provider.clone(),
config.codex_home.to_path_buf(),
std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
);
)
.await;
let control = manager.agent_control();
let cloned = control.clone();
@@ -1095,7 +1093,8 @@ async fn resume_agent_respects_max_threads_limit() {
config.model_provider.clone(),
config.codex_home.to_path_buf(),
std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
);
)
.await;
let control = manager.agent_control();
let resumable_id = control
@@ -1151,7 +1150,8 @@ async fn resume_agent_releases_slot_after_resume_failure() {
config.model_provider.clone(),
config.codex_home.to_path_buf(),
std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
);
)
.await;
let control = manager.agent_control();
let _ = control
@@ -1543,19 +1543,17 @@ async fn resume_thread_subagent_restores_stored_nickname_and_role() {
.features
.enable(Feature::Sqlite)
.expect("test config should allow sqlite");
let state_db = init_state_db(&config).await;
let manager = ThreadManager::with_models_provider_home_and_state_for_tests(
let manager = ThreadManager::with_models_provider_and_home_for_tests(
CodexAuth::from_api_key("dummy"),
config.model_provider.clone(),
config.codex_home.to_path_buf(),
std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
state_db.clone(),
);
)
.await;
let control = manager.agent_control();
let harness = AgentControlHarness {
_home: home,
config,
state_db,
manager,
control,
};
@@ -1706,7 +1704,12 @@ async fn resume_agent_from_rollout_reads_archived_rollout_path() {
.expect("child shutdown should succeed");
let store = LocalThreadStore::new(
LocalThreadStoreConfig::from_config(&harness.config),
harness.state_db.clone(),
codex_state::StateRuntime::init(
harness.config.sqlite_home.clone(),
harness.config.model_provider_id.clone(),
)
.await
.expect("state db should initialize"),
);
store
.archive_thread(ArchiveThreadParams {
+1
View File
@@ -95,6 +95,7 @@ pub(crate) async fn run_codex_thread_interactive(
parent_trace: None,
environment_selections: parent_ctx.environments.clone(),
analytics_events_client: Some(parent_session.services.analytics_events_client.clone()),
state_db: parent_session.services.state_db.clone(),
thread_store: Arc::clone(&parent_session.services.thread_store),
}))
.or_cancel(&cancel_token)
-12
View File
@@ -29,7 +29,6 @@ use codex_protocol::protocol::TokenUsage;
use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::protocol::validate_thread_goal_objective;
use codex_rollout::state_db::reconcile_rollout;
use codex_thread_store::LocalThreadStore;
use codex_utils_template::Template;
use futures::future::BoxFuture;
use std::sync::Arc;
@@ -1338,17 +1337,6 @@ impl Session {
state_db
} else if let Some(state_db) = self.goal_runtime.state_db.lock().await.clone() {
state_db
} else if let Some(local_store) = self
.services
.thread_store
.as_any()
.downcast_ref::<LocalThreadStore>()
{
local_store.state_db().await.ok_or_else(|| {
anyhow::anyhow!(
"thread goals require a local persisted thread with a state database"
)
})?
} else {
anyhow::bail!("thread goals require a local persisted thread with a state database");
};
+2
View File
@@ -124,7 +124,9 @@ pub use thread_manager::NewThread;
pub use thread_manager::StartThreadOptions;
pub use thread_manager::ThreadManager;
pub use thread_manager::ThreadShutdownReport;
pub use thread_manager::agent_graph_store_from_state_db;
pub use thread_manager::build_models_manager;
pub use thread_manager::init_state_db_from_config;
pub use thread_manager::thread_store_from_config;
pub use web_search::web_search_action_detail;
pub use web_search::web_search_detail;
+7 -10
View File
@@ -25,7 +25,7 @@ pub enum PersonalityMigrationStatus {
pub async fn maybe_migrate_personality(
codex_home: &Path,
config_toml: &ConfigToml,
state_db: Option<StateDbHandle>,
state_db: StateDbHandle,
) -> io::Result<PersonalityMigrationStatus> {
let marker_path = codex_home.join(PERSONALITY_MIGRATION_FILENAME);
if tokio::fs::try_exists(&marker_path).await? {
@@ -65,16 +65,13 @@ pub async fn maybe_migrate_personality(
async fn has_recorded_sessions(
codex_home: &Path,
default_provider: &str,
state_db: Option<StateDbHandle>,
state_db: StateDbHandle,
) -> io::Result<bool> {
let store = LocalThreadStore::new(
LocalThreadStoreConfig {
codex_home: codex_home.to_path_buf(),
sqlite_home: codex_home.to_path_buf(),
default_model_provider_id: default_provider.to_string(),
},
state_db,
);
let config = LocalThreadStoreConfig {
codex_home: codex_home.to_path_buf(),
default_model_provider_id: default_provider.to_string(),
};
let store = LocalThreadStore::new(config, state_db);
if has_threads(&store, /*archived*/ false).await? {
return Ok(true);
}
@@ -8,7 +8,10 @@ use codex_protocol::protocol::SessionMetaLine;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::UserMessageEvent;
use codex_rollout::ARCHIVED_SESSIONS_SUBDIR;
use codex_rollout::RolloutConfig;
use codex_rollout::SESSIONS_SUBDIR;
use codex_rollout::state_db::StateDbHandle;
use codex_state::state_db_path;
use pretty_assertions::assert_eq;
use tempfile::TempDir;
use tokio::io::AsyncWriteExt;
@@ -20,6 +23,26 @@ async fn read_config_toml(codex_home: &Path) -> io::Result<ConfigToml> {
toml::from_str(&contents).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))
}
async fn state_db_for_test(codex_home: &Path) -> io::Result<StateDbHandle> {
state_db_for_test_with_sqlite_home(codex_home, codex_home).await
}
async fn state_db_for_test_with_sqlite_home(
codex_home: &Path,
sqlite_home: &Path,
) -> io::Result<StateDbHandle> {
let config = RolloutConfig {
codex_home: codex_home.to_path_buf(),
sqlite_home: sqlite_home.to_path_buf(),
cwd: codex_home.to_path_buf(),
model_provider_id: "openai".to_string(),
generate_memories: false,
};
codex_rollout::state_db::try_init(&config)
.await
.map_err(io::Error::other)
}
async fn write_session_with_user_event(codex_home: &Path) -> io::Result<()> {
let thread_id = ThreadId::new();
let dir = codex_home
@@ -87,7 +110,8 @@ async fn applies_when_sessions_exist_and_no_personality() -> io::Result<()> {
write_session_with_user_event(temp.path()).await?;
let config_toml = ConfigToml::default();
let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?;
let state_db = state_db_for_test(temp.path()).await?;
let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?;
assert_eq!(status, PersonalityMigrationStatus::Applied);
assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists());
@@ -103,7 +127,8 @@ async fn applies_when_only_archived_sessions_exist_and_no_personality() -> io::R
write_archived_session_with_user_event(temp.path()).await?;
let config_toml = ConfigToml::default();
let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?;
let state_db = state_db_for_test(temp.path()).await?;
let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?;
assert_eq!(status, PersonalityMigrationStatus::Applied);
assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists());
@@ -119,7 +144,8 @@ async fn skips_when_marker_exists() -> io::Result<()> {
create_marker(&temp.path().join(PERSONALITY_MIGRATION_FILENAME)).await?;
let config_toml = ConfigToml::default();
let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?;
let state_db = state_db_for_test(temp.path()).await?;
let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?;
assert_eq!(status, PersonalityMigrationStatus::SkippedMarker);
assert!(!temp.path().join("config.toml").exists());
@@ -136,7 +162,8 @@ async fn skips_when_personality_explicit() -> io::Result<()> {
.map_err(|err| io::Error::other(format!("failed to write config: {err}")))?;
let config_toml = read_config_toml(temp.path()).await?;
let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?;
let state_db = state_db_for_test(temp.path()).await?;
let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?;
assert_eq!(
status,
@@ -153,10 +180,37 @@ async fn skips_when_personality_explicit() -> io::Result<()> {
async fn skips_when_no_sessions() -> io::Result<()> {
let temp = TempDir::new()?;
let config_toml = ConfigToml::default();
let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?;
let state_db = state_db_for_test(temp.path()).await?;
let status = maybe_migrate_personality(temp.path(), &config_toml, state_db).await?;
assert_eq!(status, PersonalityMigrationStatus::SkippedNoSessions);
assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists());
assert!(!temp.path().join("config.toml").exists());
Ok(())
}
#[tokio::test]
async fn uses_configured_sqlite_home_when_checking_for_sessions() -> io::Result<()> {
let codex_home = TempDir::new()?;
let sqlite_home = TempDir::new()?;
write_session_with_user_event(codex_home.path()).await?;
let config_toml = ConfigToml::default();
let state_db =
state_db_for_test_with_sqlite_home(codex_home.path(), sqlite_home.path()).await?;
let status = maybe_migrate_personality(codex_home.path(), &config_toml, state_db).await?;
assert_eq!(status, PersonalityMigrationStatus::Applied);
assert!(
codex_home
.path()
.join(PERSONALITY_MIGRATION_FILENAME)
.exists()
);
let persisted = read_config_toml(codex_home.path()).await?;
assert_eq!(persisted.personality, Some(Personality::Pragmatic));
assert!(!state_db_path(codex_home.path()).exists());
assert!(state_db_path(sqlite_home.path()).exists());
Ok(())
}
+8 -3
View File
@@ -16,8 +16,9 @@ use crate::config::Config;
use crate::session::session::Session;
use crate::session::turn::build_prompt;
use crate::session::turn::built_tools;
use crate::state_db_bridge::StateDbHandle;
use crate::thread_manager::ThreadManager;
use crate::thread_manager::agent_graph_store_from_state_db;
use crate::thread_manager::init_state_db_from_config;
use crate::thread_manager::thread_store_from_config;
/// Build the model-visible `input` list for a single debug turn.
@@ -25,7 +26,6 @@ use crate::thread_manager::thread_store_from_config;
pub async fn build_prompt_input(
mut config: Config,
input: Vec<UserInput>,
state_db: Option<StateDbHandle>,
) -> CodexResult<Vec<ResponseItem>> {
config.ephemeral = true;
@@ -37,15 +37,20 @@ pub async fn build_prompt_input(
config.codex_linux_sandbox_exe.clone(),
)?;
let state_db = init_state_db_from_config(&config)
.await
.ok_or_else(|| std::io::Error::other("prompt debug requires state db"))?;
let thread_store = thread_store_from_config(&config, state_db.clone());
let agent_graph_store = agent_graph_store_from_state_db(state_db.clone());
let thread_manager = ThreadManager::new(
&config,
Arc::clone(&auth_manager),
SessionSource::Exec,
Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await),
/*analytics_events_client*/ None,
state_db,
thread_store,
state_db.clone(),
agent_graph_store,
);
let thread = thread_manager.start_thread(config).await?;
+5 -12
View File
@@ -132,7 +132,6 @@ use codex_terminal_detection::user_agent;
use codex_thread_store::CreateThreadParams;
use codex_thread_store::LiveThread;
use codex_thread_store::LiveThreadInitGuard;
use codex_thread_store::LocalThreadStore;
use codex_thread_store::ResumeThreadParams;
use codex_thread_store::ThreadEventPersistenceMode;
use codex_thread_store::ThreadPersistenceMetadata;
@@ -356,7 +355,6 @@ use codex_protocol::protocol::TokenUsage;
use codex_protocol::protocol::TokenUsageInfo;
use codex_protocol::protocol::WarningEvent;
use codex_protocol::user_input::UserInput;
use codex_tools::ToolEnvironmentMode;
use codex_tools::ToolsConfig;
use codex_tools::ToolsConfigParams;
use codex_utils_absolute_path::AbsolutePathBuf;
@@ -412,6 +410,7 @@ pub(crate) struct CodexSpawnArgs {
pub(crate) parent_trace: Option<W3cTraceContext>,
pub(crate) environment_selections: ResolvedTurnEnvironments,
pub(crate) analytics_events_client: Option<AnalyticsEventsClient>,
pub(crate) state_db: Option<state_db::StateDbHandle>,
pub(crate) thread_store: Arc<dyn ThreadStore>,
}
@@ -469,6 +468,7 @@ impl Codex {
parent_trace: _,
environment_selections,
analytics_events_client,
state_db,
thread_store,
} = args;
let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY);
@@ -557,15 +557,7 @@ impl Codex {
};
match thread_id {
Some(thread_id) => {
let state_db_ctx = if config.ephemeral {
None
} else if let Some(local_store) =
thread_store.as_any().downcast_ref::<LocalThreadStore>()
{
local_store.state_db().await
} else {
None
};
let state_db_ctx = state_db.clone();
state_db::get_dynamic_tools(state_db_ctx.as_deref(), thread_id, "codex_spawn")
.await
}
@@ -651,6 +643,7 @@ impl Codex {
agent_control,
environment_manager,
analytics_events_client,
state_db,
thread_store,
parent_rollout_thread_trace,
)
@@ -1311,7 +1304,7 @@ impl Session {
self.services.user_shell.as_ref().clone(),
self.services.shell_snapshot_tx.clone(),
self.services.session_telemetry.clone(),
self.services.state_db.clone(),
self.state_db(),
);
}
+3 -23
View File
@@ -343,6 +343,7 @@ impl Session {
agent_control: AgentControl,
environment_manager: Arc<EnvironmentManager>,
analytics_events_client: Option<AnalyticsEventsClient>,
state_db: Option<state_db::StateDbHandle>,
thread_store: Arc<dyn ThreadStore>,
parent_rollout_thread_trace: ThreadTraceContext,
) -> anyhow::Result<Arc<Self>> {
@@ -441,22 +442,7 @@ impl Session {
otel.name = "session_init.thread_persistence",
session_init.ephemeral = config.ephemeral,
));
let state_db_fut = async {
if config.ephemeral {
None
} else if let Some(local_store) =
thread_store.as_any().downcast_ref::<LocalThreadStore>()
{
local_store.state_db().await
} else {
None
}
}
.instrument(info_span!(
"session_init.state_db",
otel.name = "session_init.state_db",
session_init.ephemeral = config.ephemeral,
));
let state_db_ctx = if config.ephemeral { None } else { state_db };
let is_subagent = session_configuration.session_source.is_non_root_agent();
let history_meta_fut = async {
@@ -495,15 +481,9 @@ impl Session {
// Join all independent futures.
let (
thread_persistence_result,
state_db_ctx,
(history_log_id, history_entry_count),
(auth, mcp_servers, auth_statuses),
) = tokio::join!(
thread_persistence_fut,
state_db_fut,
history_meta_fut,
auth_and_mcp_fut
);
) = tokio::join!(thread_persistence_fut, history_meta_fut, auth_and_mcp_fut);
let mut live_thread_init =
LiveThreadInitGuard::new(thread_persistence_result.map_err(|e| {
+67 -79
View File
@@ -867,7 +867,7 @@ async fn danger_full_access_turns_do_not_expose_managed_network_proxy() -> anyho
&permission_profile_for_sandbox_policy(&SandboxPolicy::DangerFullAccess),
)?;
let session = make_session_with_config(move |config| {
let (session, _codex_home) = make_session_with_config(move |config| {
let cwd = config.cwd.clone();
config
.permissions
@@ -933,7 +933,7 @@ async fn danger_full_access_tool_attempts_do_not_enforce_managed_network() -> an
&permission_profile_for_sandbox_policy(&SandboxPolicy::DangerFullAccess),
)?;
let session = make_session_with_config(move |config| {
let (session, _codex_home) = make_session_with_config(move |config| {
let cwd = config.cwd.clone();
config
.permissions
@@ -1008,7 +1008,7 @@ async fn workspace_write_turns_continue_to_expose_managed_network_proxy() -> any
&permission_profile_for_sandbox_policy(&sandbox_policy),
)?;
let session = make_session_with_config(move |config| {
let (session, _codex_home) = make_session_with_config(move |config| {
let cwd = config.cwd.clone();
config
.permissions
@@ -1035,7 +1035,7 @@ async fn user_shell_commands_do_not_inherit_managed_network_proxy() -> anyhow::R
&permission_profile_for_sandbox_policy(&sandbox_policy),
)?;
let (session, rx) = make_session_with_config_and_rx(move |config| {
let (session, rx, _codex_home) = make_session_with_config_and_rx(move |config| {
let cwd = config.cwd.clone();
config
.permissions
@@ -1165,7 +1165,7 @@ async fn reload_user_config_layer_updates_effective_apps_config() {
#[tokio::test]
async fn reload_user_config_layer_refreshes_hooks() -> anyhow::Result<()> {
let session = make_session_with_config(|config| {
let (session, _codex_home) = make_session_with_config(|config| {
config
.features
.enable(Feature::CodexHooks)
@@ -2982,15 +2982,13 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati
fn turn_environments_for_tests(
environment: &Arc<codex_exec_server::Environment>,
cwd: &codex_utils_absolute_path::AbsolutePathBuf,
) -> crate::environment_selection::ResolvedTurnEnvironments {
crate::environment_selection::ResolvedTurnEnvironments {
turn_environments: vec![TurnEnvironment {
environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(),
environment: Arc::clone(environment),
cwd: cwd.clone(),
shell: "bash".to_string(),
}],
}
) -> Vec<TurnEnvironment> {
vec![TurnEnvironment {
environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(),
environment: Arc::clone(environment),
cwd: cwd.clone(),
shell: "bash".to_string(),
}]
}
#[tokio::test]
@@ -3528,9 +3526,15 @@ async fn session_new_fails_when_zsh_fork_enabled_without_zsh_path() {
AgentControl::default(),
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
/*state_db*/ None,
Arc::new(codex_thread_store::LocalThreadStore::new(
codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()),
/*state_db*/ None,
codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.model_provider_id.clone(),
)
.await
.expect("state db should initialize"),
)),
codex_rollout_trace::ThreadTraceContext::disabled(),
)
@@ -3678,7 +3682,12 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
live_thread: None,
thread_store: Arc::new(codex_thread_store::LocalThreadStore::new(
codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()),
/*state_db*/ None,
codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.model_provider_id.clone(),
)
.await
.expect("state db should initialize"),
)),
model_client: ModelClient::new(
Some(auth_manager.clone()),
@@ -3723,7 +3732,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
model_info,
&models_manager,
/*network*/ None,
turn_environments,
crate::environment_selection::ResolvedTurnEnvironments { turn_environments },
session_configuration.cwd.clone(),
"turn_id".to_string(),
skills_outcome,
@@ -3756,14 +3765,18 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
async fn make_session_with_config(
mutator: impl FnOnce(&mut Config),
) -> anyhow::Result<Arc<Session>> {
let (session, _rx_event) = make_session_with_config_and_rx(mutator).await?;
Ok(session)
) -> anyhow::Result<(Arc<Session>, tempfile::TempDir)> {
let (session, _rx_event, codex_home) = make_session_with_config_and_rx(mutator).await?;
Ok((session, codex_home))
}
async fn make_session_with_config_and_rx(
mutator: impl FnOnce(&mut Config),
) -> anyhow::Result<(Arc<Session>, async_channel::Receiver<Event>)> {
) -> anyhow::Result<(
Arc<Session>,
async_channel::Receiver<Event>,
tempfile::TempDir,
)> {
let codex_home = tempfile::tempdir().expect("create temp dir");
let mut config = build_test_config(codex_home.path()).await;
mutator(&mut config);
@@ -3848,15 +3861,21 @@ async fn make_session_with_config_and_rx(
AgentControl::default(),
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
/*state_db*/ None,
Arc::new(codex_thread_store::LocalThreadStore::new(
codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()),
/*state_db*/ None,
codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.model_provider_id.clone(),
)
.await
.expect("state db should initialize"),
)),
codex_rollout_trace::ThreadTraceContext::disabled(),
)
.await?;
Ok((session, rx_event))
Ok((session, rx_event, codex_home))
}
#[tokio::test]
@@ -4518,7 +4537,7 @@ async fn default_turn_honors_empty_stored_thread_environments() {
let turn_context = session.new_default_turn().await;
assert!(turn_context.environments.primary().is_none());
assert!(turn_context.environments.primary_environment().is_none());
assert!(turn_context.environments.turn_environments.is_empty());
assert_eq!(turn_context.cwd, session_cwd);
assert_eq!(turn_context.config.cwd, session_cwd);
@@ -4537,7 +4556,7 @@ async fn primary_environment_uses_first_turn_environment() {
environment_id: "second".to_string(),
environment: Arc::clone(&first_environment.environment),
cwd: second_cwd.clone(),
shell: first_environment.shell.clone(),
shell: "bash".to_string(),
});
assert_eq!(
@@ -4580,7 +4599,7 @@ async fn empty_turn_environments_clear_primary_environment() {
.await
.expect("turn should start");
assert!(turn_context.environments.primary().is_none());
assert!(turn_context.environments.primary_environment().is_none());
assert!(turn_context.environments.turn_environments.is_empty());
assert_eq!(turn_context.cwd, session.get_config().await.cwd);
assert_eq!(turn_context.config.cwd, session.get_config().await.cwd);
@@ -5036,47 +5055,13 @@ async fn make_session_and_context_with_auth_and_config_and_rx<F>(
Arc<TurnContext>,
async_channel::Receiver<Event>,
)
where
F: FnOnce(&mut Config),
{
let codex_home = tempfile::tempdir().expect("create temp dir");
make_session_and_context_with_auth_config_home_and_rx(
auth,
dynamic_tools,
codex_home.path(),
configure_config,
)
.await
}
async fn make_session_and_context_with_auth_config_home_and_rx<F>(
auth: CodexAuth,
dynamic_tools: Vec<DynamicToolSpec>,
codex_home: &Path,
configure_config: F,
) -> (
Arc<Session>,
Arc<TurnContext>,
async_channel::Receiver<Event>,
)
where
F: FnOnce(&mut Config),
{
let (tx_event, rx_event) = async_channel::unbounded();
let mut config = build_test_config(codex_home).await;
let codex_home = tempfile::tempdir().expect("create temp dir").keep();
let mut config = build_test_config(codex_home.as_path()).await;
configure_config(&mut config);
let state_db = if config.features.enabled(Feature::Goals) {
Some(
codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.model_provider_id.clone(),
)
.await
.expect("goal tests should initialize sqlite state db"),
)
} else {
None
};
let config = Arc::new(config);
let conversation_id = ThreadId::default();
let auth_manager = AuthManager::from_auth_for_testing(auth);
@@ -5162,6 +5147,12 @@ where
.expect("create environment"),
);
let state_db = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let skills_watcher = Arc::new(SkillsWatcher::noop());
let services = SessionServices {
mcp_connection_manager: Arc::new(RwLock::new(McpConnectionManager::new_uninitialized(
@@ -5202,7 +5193,7 @@ where
agent_control,
network_proxy: None,
network_approval: Arc::clone(&network_approval),
state_db: state_db.clone(),
state_db: Some(state_db.clone()),
live_thread: None,
thread_store: Arc::new(codex_thread_store::LocalThreadStore::new(
codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()),
@@ -5251,7 +5242,7 @@ where
model_info,
&models_manager,
/*network*/ None,
turn_environments,
crate::environment_selection::ResolvedTurnEnvironments { turn_environments },
session_configuration.cwd.clone(),
"turn_id".to_string(),
skills_outcome,
@@ -5301,13 +5292,10 @@ async fn make_goal_session_and_context_with_rx() -> (
Arc<Session>,
Arc<TurnContext>,
async_channel::Receiver<Event>,
tempfile::TempDir,
) {
let codex_home = tempfile::tempdir().expect("create temp dir");
let (session, turn_context, rx) = make_session_and_context_with_auth_config_home_and_rx(
let (session, turn_context, rx) = make_session_and_context_with_auth_and_config_and_rx(
CodexAuth::from_api_key("Test API Key"),
Vec::new(),
codex_home.path(),
|config| {
config
.features
@@ -5317,14 +5305,14 @@ async fn make_goal_session_and_context_with_rx() -> (
)
.await;
upsert_goal_test_thread(session.as_ref()).await;
(session, turn_context, rx, codex_home)
(session, turn_context, rx)
}
async fn upsert_goal_test_thread(session: &Session) {
let config = session.get_config().await;
let state_db = session
.state_db()
.expect("goal test session should have a state db");
let state_db = goal_test_state_db(session)
.await
.expect("goal test state db should initialize");
let mut builder = codex_state::ThreadMetadataBuilder::new(
session.conversation_id,
config
@@ -7148,7 +7136,7 @@ async fn abort_empty_active_turn_preserves_pending_input() {
#[tokio::test]
async fn interrupt_accounts_active_goal_before_pausing() -> anyhow::Result<()> {
let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
let (sess, tc, _rx) = make_goal_session_and_context_with_rx().await;
sess.set_thread_goal(
tc.as_ref(),
SetGoalRequest {
@@ -7412,7 +7400,7 @@ async fn goal_test_state_db(sess: &Session) -> anyhow::Result<crate::StateDbHand
#[tokio::test]
async fn budget_limited_accounting_steers_active_turn_without_aborting() -> anyhow::Result<()> {
let (sess, tc, rx, _codex_home) = make_goal_session_and_context_with_rx().await;
let (sess, tc, rx) = make_goal_session_and_context_with_rx().await;
sess.set_thread_goal(
tc.as_ref(),
SetGoalRequest {
@@ -7512,7 +7500,7 @@ async fn budget_limited_accounting_steers_active_turn_without_aborting() -> anyh
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn external_goal_mutation_accounts_active_turn_before_status_change() -> anyhow::Result<()> {
let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
let (sess, tc, _rx) = make_goal_session_and_context_with_rx().await;
sess.set_thread_goal(
tc.as_ref(),
SetGoalRequest {
@@ -7579,7 +7567,7 @@ async fn external_goal_mutation_accounts_active_turn_before_status_change() -> a
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn external_active_goal_set_marks_current_turn_for_accounting() -> anyhow::Result<()> {
let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
let (sess, tc, _rx) = make_goal_session_and_context_with_rx().await;
sess.spawn_task(
Arc::clone(&tc),
Vec::new(),
@@ -8246,7 +8234,7 @@ async fn sample_rollout(
#[tokio::test]
async fn create_goal_tool_rejects_existing_goal() {
let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
let (session, turn_context, _rx) = make_goal_session_and_context_with_rx().await;
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
let handler = CreateGoalHandler;
@@ -8308,7 +8296,7 @@ async fn create_goal_tool_rejects_existing_goal() {
#[tokio::test]
async fn update_goal_tool_rejects_pausing_goal() {
let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
let (session, turn_context, _rx) = make_goal_session_and_context_with_rx().await;
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
let create_handler = CreateGoalHandler;
let update_handler = UpdateGoalHandler;
@@ -8369,7 +8357,7 @@ async fn update_goal_tool_rejects_pausing_goal() {
#[tokio::test]
async fn update_goal_tool_marks_goal_complete() {
let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
let (session, turn_context, _rx) = make_goal_session_and_context_with_rx().await;
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
let create_handler = CreateGoalHandler;
let update_handler = UpdateGoalHandler;
@@ -731,7 +731,12 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() {
let skills_watcher = Arc::new(SkillsWatcher::noop());
let thread_store = Arc::new(codex_thread_store::LocalThreadStore::new(
codex_thread_store::LocalThreadStoreConfig::from_config(&config),
/*state_db*/ None,
codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.model_provider_id.clone(),
)
.await
.expect("state db should initialize"),
));
let CodexSpawnOk { codex, .. } = Codex::spawn(CodexSpawnArgs {
@@ -760,6 +765,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() {
turn_environments: Vec::new(),
},
analytics_events_client: None,
state_db: None,
thread_store,
})
.await
@@ -8,6 +8,7 @@ use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_sandboxing::compatibility_sandbox_policy_for_permission_profile;
use codex_sandboxing::policy_transforms::effective_file_system_sandbox_policy;
use codex_sandboxing::policy_transforms::effective_network_sandbox_policy;
use codex_tools::ToolEnvironmentMode;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
+3 -6
View File
@@ -138,11 +138,8 @@ pub(crate) async fn record_completed_response_item(
.await;
}
mark_thread_memory_mode_polluted_if_external_context(sess, turn_context, item).await;
let has_memory_citation = record_stage1_output_usage_and_detect_memory_citation(
sess.services.state_db.as_ref(),
item,
)
.await;
let has_memory_citation =
record_stage1_output_usage_and_detect_memory_citation(sess.state_db(), item).await;
if has_memory_citation {
sess.record_memory_citation_for_turn(&turn_context.sub_id)
.await;
@@ -177,7 +174,7 @@ pub(crate) async fn mark_thread_memory_mode_polluted_if_external_context(
}
async fn record_stage1_output_usage_and_detect_memory_citation(
state_db_ctx: Option<&state_db::StateDbHandle>,
state_db_ctx: Option<state_db::StateDbHandle>,
item: &ResponseItem,
) -> bool {
let Some(raw_text) = raw_assistant_output_text_from_item(item) else {
+4 -19
View File
@@ -52,14 +52,14 @@ pub fn auth_manager_from_auth_with_home(auth: CodexAuth, codex_home: PathBuf) ->
AuthManager::from_auth_for_testing_with_home(auth, codex_home)
}
pub fn thread_manager_with_models_provider(
pub async fn thread_manager_with_models_provider(
auth: CodexAuth,
provider: ModelProviderInfo,
) -> ThreadManager {
ThreadManager::with_models_provider_for_tests(auth, provider)
ThreadManager::with_models_provider_for_tests(auth, provider).await
}
pub fn thread_manager_with_models_provider_and_home(
pub async fn thread_manager_with_models_provider_and_home(
auth: CodexAuth,
provider: ModelProviderInfo,
codex_home: PathBuf,
@@ -71,22 +71,7 @@ pub fn thread_manager_with_models_provider_and_home(
codex_home,
environment_manager,
)
}
pub fn thread_manager_with_models_provider_home_and_state(
auth: CodexAuth,
provider: ModelProviderInfo,
codex_home: PathBuf,
environment_manager: Arc<EnvironmentManager>,
state_db: Option<crate::StateDbHandle>,
) -> ThreadManager {
ThreadManager::with_models_provider_home_and_state_for_tests(
auth,
provider,
codex_home,
environment_manager,
state_db,
)
.await
}
pub async fn start_thread_with_user_shell_override(
+76 -34
View File
@@ -18,6 +18,8 @@ use crate::skills_watcher::SkillsWatcher;
use crate::skills_watcher::SkillsWatcherEvent;
use crate::tasks::InterruptedTurnHistoryMarker;
use crate::tasks::interrupted_turn_history_marker;
use codex_agent_graph_store::AgentGraphStore;
use codex_agent_graph_store::LocalAgentGraphStore;
use codex_analytics::AnalyticsEventsClient;
use codex_app_server_protocol::ThreadHistoryBuilder;
use codex_app_server_protocol::TurnStatus;
@@ -50,8 +52,8 @@ use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::protocol::TurnAbortedEvent;
use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_protocol::protocol::W3cTraceContext;
use codex_rollout::state_db;
use codex_rollout::state_db::StateDbHandle;
use codex_state::DirectionalThreadSpawnEdgeStatus;
use codex_thread_store::InMemoryThreadStore;
use codex_thread_store::LocalThreadStore;
use codex_thread_store::LocalThreadStoreConfig;
@@ -247,9 +249,10 @@ pub(crate) struct ThreadManagerState {
mcp_manager: Arc<McpManager>,
skills_watcher: Arc<SkillsWatcher>,
thread_store: Arc<dyn ThreadStore>,
state_db: StateDbHandle,
agent_graph_store: Arc<dyn AgentGraphStore>,
session_source: SessionSource,
analytics_events_client: Option<AnalyticsEventsClient>,
state_db: Option<StateDbHandle>,
// Captures submitted ops for testing purpose when test mode is enabled.
ops_log: Option<SharedCapturedOps>,
}
@@ -265,10 +268,11 @@ pub fn build_models_manager(
)
}
pub fn thread_store_from_config(
config: &Config,
state_db: Option<StateDbHandle>,
) -> Arc<dyn ThreadStore> {
pub async fn init_state_db_from_config(config: &Config) -> Option<StateDbHandle> {
state_db::init(config).await
}
pub fn thread_store_from_config(config: &Config, state_db: StateDbHandle) -> Arc<dyn ThreadStore> {
match &config.experimental_thread_store {
ThreadStoreConfig::Local => Arc::new(LocalThreadStore::new(
LocalThreadStoreConfig::from_config(config),
@@ -279,15 +283,38 @@ pub fn thread_store_from_config(
}
}
pub fn agent_graph_store_from_state_db(state_db: StateDbHandle) -> Arc<dyn AgentGraphStore> {
Arc::new(LocalAgentGraphStore::new(state_db))
}
async fn state_db_from_roots_for_tests(
codex_home: PathBuf,
sqlite_home: PathBuf,
default_model_provider_id: String,
) -> StateDbHandle {
let config = codex_rollout::RolloutConfig {
codex_home: codex_home.clone(),
sqlite_home,
cwd: codex_home,
model_provider_id: default_model_provider_id,
generate_memories: false,
};
state_db::try_init(&config)
.await
.unwrap_or_else(|err| panic!("test state db should initialize: {err}"))
}
impl ThreadManager {
#[allow(clippy::too_many_arguments)]
pub fn new(
config: &Config,
auth_manager: Arc<AuthManager>,
session_source: SessionSource,
environment_manager: Arc<EnvironmentManager>,
analytics_events_client: Option<AnalyticsEventsClient>,
state_db: StateDbHandle,
thread_store: Arc<dyn ThreadStore>,
state_db: Option<StateDbHandle>,
agent_graph_store: Arc<dyn AgentGraphStore>,
) -> Self {
let codex_home = config.codex_home.clone();
let restriction_product = session_source.restriction_product();
@@ -314,10 +341,11 @@ impl ThreadManager {
mcp_manager,
skills_watcher,
thread_store,
state_db,
agent_graph_store,
auth_manager,
session_source,
analytics_events_client,
state_db,
ops_log: should_use_test_thread_manager_behavior()
.then(|| Arc::new(std::sync::Mutex::new(Vec::new()))),
}),
@@ -327,7 +355,7 @@ impl ThreadManager {
/// Construct with a dummy AuthManager containing the provided CodexAuth.
/// Used for integration tests: should not be used by ordinary business logic.
pub(crate) fn with_models_provider_for_tests(
pub(crate) async fn with_models_provider_for_tests(
auth: CodexAuth,
provider: ModelProviderInfo,
) -> Self {
@@ -338,11 +366,18 @@ impl ThreadManager {
));
std::fs::create_dir_all(&codex_home)
.unwrap_or_else(|err| panic!("temp codex home dir create failed: {err}"));
let mut manager = Self::with_models_provider_and_home_for_tests(
let state_db = state_db_from_roots_for_tests(
codex_home.clone(),
codex_home.clone(),
OPENAI_PROVIDER_ID.to_string(),
)
.await;
let mut manager = Self::with_models_provider_and_home_and_state_db_for_tests(
auth,
provider,
codex_home.clone(),
Arc::new(EnvironmentManager::default_for_tests()),
state_db,
);
manager._test_codex_home_guard = Some(TempCodexHomeGuard { path: codex_home });
manager
@@ -350,27 +385,33 @@ impl ThreadManager {
/// Construct with a dummy AuthManager containing the provided CodexAuth and codex home.
/// Used for integration tests: should not be used by ordinary business logic.
pub(crate) fn with_models_provider_and_home_for_tests(
pub(crate) async fn with_models_provider_and_home_for_tests(
auth: CodexAuth,
provider: ModelProviderInfo,
codex_home: PathBuf,
environment_manager: Arc<EnvironmentManager>,
) -> Self {
Self::with_models_provider_home_and_state_for_tests(
let state_db = state_db_from_roots_for_tests(
codex_home.clone(),
codex_home.clone(),
OPENAI_PROVIDER_ID.to_string(),
)
.await;
Self::with_models_provider_and_home_and_state_db_for_tests(
auth,
provider,
codex_home,
environment_manager,
/*state_db*/ None,
state_db,
)
}
pub(crate) fn with_models_provider_home_and_state_for_tests(
fn with_models_provider_and_home_and_state_db_for_tests(
auth: CodexAuth,
provider: ModelProviderInfo,
codex_home: PathBuf,
environment_manager: Arc<EnvironmentManager>,
state_db: Option<StateDbHandle>,
state_db: StateDbHandle,
) -> Self {
set_thread_manager_test_mode_for_tests(/*enabled*/ true);
let auth_manager = AuthManager::from_auth_for_testing(auth);
@@ -396,11 +437,11 @@ impl ThreadManager {
let thread_store: Arc<dyn ThreadStore> = Arc::new(LocalThreadStore::new(
LocalThreadStoreConfig {
codex_home: codex_home.clone(),
sqlite_home: codex_home.clone(),
default_model_provider_id: OPENAI_PROVIDER_ID.to_string(),
},
state_db.clone(),
));
let agent_graph_store = agent_graph_store_from_state_db(state_db.clone());
Self {
state: Arc::new(ThreadManagerState {
threads: Arc::new(RwLock::new(HashMap::new())),
@@ -413,10 +454,11 @@ impl ThreadManager {
mcp_manager,
skills_watcher,
thread_store,
state_db,
agent_graph_store,
auth_manager,
session_source: SessionSource::Exec,
analytics_events_client: None,
state_db,
ops_log: should_use_test_thread_manager_behavior()
.then(|| Arc::new(std::sync::Mutex::new(Vec::new()))),
}),
@@ -523,22 +565,17 @@ impl ThreadManager {
subtree_thread_ids.push(thread_id);
seen_thread_ids.insert(thread_id);
if let Some(state_db_ctx) = thread.state_db() {
for status in [
DirectionalThreadSpawnEdgeStatus::Open,
DirectionalThreadSpawnEdgeStatus::Closed,
] {
for descendant_id in state_db_ctx
.list_thread_spawn_descendants_with_status(thread_id, status)
.await
.map_err(|err| {
CodexErr::Fatal(format!("failed to load thread-spawn descendants: {err}"))
})?
{
if seen_thread_ids.insert(descendant_id) {
subtree_thread_ids.push(descendant_id);
}
}
for descendant_id in self
.state
.agent_graph_store
.list_thread_spawn_descendants(thread_id, /*status_filter*/ None)
.await
.map_err(|err| {
CodexErr::Fatal(format!("failed to load thread-spawn descendants: {err}"))
})?
{
if seen_thread_ids.insert(descendant_id) {
subtree_thread_ids.push(descendant_id);
}
}
@@ -864,10 +901,14 @@ impl ThreadManager {
}
impl ThreadManagerState {
pub(crate) fn state_db(&self) -> Option<StateDbHandle> {
pub(crate) fn state_db(&self) -> StateDbHandle {
self.state_db.clone()
}
pub(crate) fn agent_graph_store(&self) -> Arc<dyn AgentGraphStore> {
self.agent_graph_store.clone()
}
pub(crate) async fn list_thread_ids(&self) -> Vec<ThreadId> {
self.threads
.read()
@@ -1172,6 +1213,7 @@ impl ThreadManagerState {
parent_trace,
environment_selections,
analytics_events_client: self.analytics_events_client.clone(),
state_db: Some(self.state_db.clone()),
thread_store: Arc::clone(&self.thread_store),
})
.await?;
+58 -24
View File
@@ -1,6 +1,5 @@
use super::*;
use crate::config::test_config;
use crate::init_state_db;
use crate::rollout::RolloutRecorder;
use crate::session::session::SessionSettingsUpdate;
use crate::session::tests::make_session_and_context;
@@ -47,6 +46,21 @@ fn assistant_msg(text: &str) -> ResponseItem {
}
}
async fn state_backed_stores(
config: &Config,
) -> (
StateDbHandle,
Arc<dyn ThreadStore>,
Arc<dyn AgentGraphStore>,
) {
let state_db = init_state_db_from_config(config)
.await
.expect("thread manager test requires state db");
let thread_store = thread_store_from_config(config, state_db.clone());
let agent_graph_store = agent_graph_store_from_state_db(state_db.clone());
(state_db, thread_store, agent_graph_store)
}
fn contextual_user_interrupted_marker() -> ResponseItem {
interrupted_turn_history_marker(InterruptedTurnHistoryMarker::ContextualUser)
.expect("contextual-user interrupted marker should be enabled")
@@ -261,7 +275,8 @@ async fn shutdown_all_threads_bounded_submits_shutdown_to_every_thread() {
config.model_provider.clone(),
config.codex_home.to_path_buf(),
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
);
)
.await;
let thread_1 = manager
.start_thread(config.clone())
.await
@@ -310,7 +325,8 @@ async fn start_thread_accepts_explicit_environment_when_default_environment_is_d
config.model_provider.clone(),
config.codex_home.to_path_buf(),
environment_manager,
);
)
.await;
let thread = manager
.start_thread_with_options(StartThreadOptions {
@@ -345,7 +361,8 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() {
config.model_provider.clone(),
config.codex_home.to_path_buf(),
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
);
)
.await;
let thread = manager
.start_thread_with_options(StartThreadOptions {
config,
@@ -384,14 +401,16 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() {
let auth_manager =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await;
let manager = ThreadManager::new(
&config,
auth_manager.clone(),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
thread_store_from_config(&config, /*state_db*/ None),
/*state_db*/ None,
state_db,
thread_store,
agent_graph_store,
);
let selected_cwd =
AbsolutePathBuf::try_from(config.cwd.as_path().join("selected")).expect("absolute path");
@@ -494,14 +513,16 @@ async fn resume_active_thread_from_rollout_returns_running_thread() {
let auth_manager =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await;
let manager = ThreadManager::new(
&config,
auth_manager.clone(),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
thread_store_from_config(&config, /*state_db*/ None),
/*state_db*/ None,
state_db,
thread_store,
agent_graph_store,
);
let source = manager
@@ -548,14 +569,16 @@ async fn resume_stopped_thread_from_rollout_spawns_new_thread() {
let auth_manager =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await;
let manager = ThreadManager::new(
&config,
auth_manager.clone(),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
thread_store_from_config(&config, /*state_db*/ None),
/*state_db*/ None,
state_db,
thread_store,
agent_graph_store,
);
let source = manager
@@ -612,14 +635,16 @@ async fn new_uses_active_provider_for_model_refresh() {
let auth_manager =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await;
let manager = ThreadManager::new(
&config,
auth_manager,
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
thread_store_from_config(&config, /*state_db*/ None),
/*state_db*/ None,
state_db,
thread_store,
agent_graph_store,
);
let _ = manager.list_models(RefreshStrategy::Online).await;
@@ -824,15 +849,16 @@ async fn interrupted_fork_snapshot_does_not_synthesize_turn_id_for_legacy_histor
let auth_manager =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let state_db = init_state_db(&config).await;
let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await;
let manager = ThreadManager::new(
&config,
auth_manager.clone(),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
thread_store_from_config(&config, state_db.clone()),
state_db.clone(),
state_db,
thread_store,
agent_graph_store,
);
let source = manager
@@ -928,15 +954,16 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() {
let auth_manager =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let state_db = init_state_db(&config).await;
let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await;
let manager = ThreadManager::new(
&config,
auth_manager.clone(),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
thread_store_from_config(&config, state_db.clone()),
state_db.clone(),
state_db,
thread_store,
agent_graph_store,
);
let source = manager
@@ -1021,15 +1048,16 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_
let auth_manager =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let state_db = init_state_db(&config).await;
let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await;
let manager = ThreadManager::new(
&config,
auth_manager.clone(),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
thread_store_from_config(&config, state_db.clone()),
state_db.clone(),
state_db,
thread_store,
agent_graph_store,
);
let source = manager
@@ -1159,15 +1187,16 @@ async fn resumed_thread_keeps_paused_goal_paused() -> anyhow::Result<()> {
let auth_manager =
AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
let state_db = init_state_db(&config).await;
let (state_db, thread_store, agent_graph_store) = state_backed_stores(&config).await;
let manager = ThreadManager::new(
&config,
auth_manager.clone(),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
thread_store_from_config(&config, state_db.clone()),
state_db.clone(),
state_db,
thread_store,
agent_graph_store,
);
let source = manager
@@ -1236,6 +1265,11 @@ async fn resumed_thread_keeps_paused_goal_paused() -> anyhow::Result<()> {
.await
.is_none()
);
let goal = state_db
.get_thread_goal(resumed.thread_id)
.await?
.expect("goal should still exist after resume");
assert_eq!(codex_state::ThreadGoalStatus::Paused, goal.status);
resumed.thread.shutdown_and_wait().await?;
Ok(())
@@ -6,6 +6,7 @@ use crate::function_tool::FunctionCallError;
use crate::init_state_db;
use crate::session::tests::make_session_and_context;
use crate::session_prefix::format_subagent_notification_message;
use crate::thread_manager::agent_graph_store_from_state_db;
use crate::thread_manager::thread_store_from_config;
use crate::tools::context::ToolOutput;
use crate::tools::handlers::multi_agents_v2::CloseAgentHandler as CloseAgentHandlerV2;
@@ -90,11 +91,12 @@ fn parse_agent_id(id: &str) -> ThreadId {
ThreadId::from_string(id).expect("agent id should be valid")
}
fn thread_manager() -> ThreadManager {
async fn thread_manager() -> ThreadManager {
ThreadManager::with_models_provider_for_tests(
CodexAuth::from_api_key("dummy"),
built_in_model_providers(/* openai_base_url */ /*openai_base_url*/ None)["openai"].clone(),
)
.await
}
async fn install_role_with_model_override(turn: &mut TurnContext) -> String {
@@ -241,7 +243,7 @@ async fn spawn_agent_uses_explorer_role_and_preserves_approval_policy() {
}
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let mut config = (*turn.config).clone();
let provider_info =
@@ -296,7 +298,7 @@ async fn spawn_agent_uses_explorer_role_and_preserves_approval_policy() {
async fn spawn_agent_fork_context_rejects_agent_type_override() {
let (mut session, mut turn) = make_session_and_context().await;
let role_name = install_role_with_model_override(&mut turn).await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -328,7 +330,7 @@ async fn spawn_agent_fork_context_rejects_agent_type_override() {
#[tokio::test]
async fn spawn_agent_fork_context_rejects_child_model_overrides() {
let (mut session, turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -363,7 +365,7 @@ async fn spawn_agent_fork_context_rejects_child_model_overrides() {
async fn multi_agent_v2_spawn_fork_turns_all_rejects_agent_type_override() {
let (mut session, mut turn) = make_session_and_context().await;
let role_name = install_role_with_model_override(&mut turn).await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -406,7 +408,7 @@ async fn multi_agent_v2_spawn_fork_turns_all_rejects_agent_type_override() {
#[tokio::test]
async fn multi_agent_v2_spawn_defaults_to_full_fork_and_rejects_child_model_overrides() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -447,7 +449,7 @@ async fn multi_agent_v2_spawn_defaults_to_full_fork_and_rejects_child_model_over
async fn multi_agent_v2_spawn_partial_fork_turns_allows_agent_type_override() {
let (mut session, mut turn) = make_session_and_context().await;
let role_name = install_role_with_model_override(&mut turn).await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -503,7 +505,7 @@ async fn multi_agent_v2_spawn_partial_fork_turns_allows_agent_type_override() {
#[tokio::test]
async fn spawn_agent_returns_agent_id_without_task_name() {
let (mut session, turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let output = SpawnAgentHandler
@@ -530,7 +532,7 @@ async fn spawn_agent_returns_agent_id_without_task_name() {
#[tokio::test]
async fn multi_agent_v2_spawn_requires_task_name() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -564,7 +566,7 @@ async fn multi_agent_v2_spawn_requires_task_name() {
#[tokio::test]
async fn multi_agent_v2_spawn_rejects_legacy_items_field() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -624,7 +626,7 @@ async fn multi_agent_v2_spawn_returns_path_and_send_message_accepts_relative_pat
}
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -721,7 +723,7 @@ async fn multi_agent_v2_spawn_returns_path_and_send_message_accepts_relative_pat
#[tokio::test]
async fn multi_agent_v2_spawn_rejects_legacy_fork_context() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -760,7 +762,7 @@ async fn multi_agent_v2_spawn_rejects_legacy_fork_context() {
#[tokio::test]
async fn multi_agent_v2_spawn_rejects_invalid_fork_turns_string() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -799,7 +801,7 @@ async fn multi_agent_v2_spawn_rejects_invalid_fork_turns_string() {
#[tokio::test]
async fn multi_agent_v2_spawn_rejects_zero_fork_turns() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -838,7 +840,7 @@ async fn multi_agent_v2_spawn_rejects_zero_fork_turns() {
#[tokio::test]
async fn multi_agent_v2_send_message_accepts_root_target_from_child() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -914,7 +916,7 @@ async fn multi_agent_v2_send_message_accepts_root_target_from_child() {
#[tokio::test]
async fn multi_agent_v2_followup_task_rejects_root_target_from_child() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -995,7 +997,7 @@ async fn multi_agent_v2_followup_task_rejects_root_target_from_child() {
#[tokio::test]
async fn multi_agent_v2_list_agents_returns_completed_status_and_last_task_message() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -1089,7 +1091,7 @@ async fn multi_agent_v2_list_agents_returns_completed_status_and_last_task_messa
#[tokio::test]
async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -1176,7 +1178,7 @@ async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() {
#[tokio::test]
async fn multi_agent_v2_list_agents_omits_closed_agents() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -1240,7 +1242,7 @@ async fn multi_agent_v2_list_agents_omits_closed_agents() {
#[tokio::test]
async fn multi_agent_v2_send_message_rejects_legacy_items_field() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -1296,7 +1298,7 @@ async fn multi_agent_v2_send_message_rejects_legacy_items_field() {
#[tokio::test]
async fn multi_agent_v2_send_message_rejects_interrupt_parameter() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -1369,7 +1371,7 @@ async fn multi_agent_v2_send_message_rejects_interrupt_parameter() {
#[tokio::test]
async fn multi_agent_v2_followup_task_completion_notifies_parent_on_every_turn() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -1504,7 +1506,7 @@ async fn multi_agent_v2_followup_task_completion_notifies_parent_on_every_turn()
#[tokio::test]
async fn multi_agent_v2_followup_task_rejects_legacy_items_field() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -1557,7 +1559,7 @@ async fn multi_agent_v2_followup_task_rejects_legacy_items_field() {
#[tokio::test]
async fn multi_agent_v2_interrupted_turn_does_not_notify_parent() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -1634,7 +1636,7 @@ async fn multi_agent_v2_interrupted_turn_does_not_notify_parent() {
#[tokio::test]
async fn multi_agent_v2_spawn_omits_agent_id_when_named() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -1673,7 +1675,7 @@ async fn multi_agent_v2_spawn_omits_agent_id_when_named() {
#[tokio::test]
async fn multi_agent_v2_spawn_surfaces_task_name_validation_errors() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -1716,7 +1718,7 @@ async fn spawn_agent_reapplies_runtime_sandbox_after_role_config() {
}
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let expected_sandbox = turn.config.legacy_sandbox_policy();
let mut expected_file_system_sandbox_policy =
@@ -1797,7 +1799,7 @@ async fn spawn_agent_reapplies_runtime_sandbox_after_role_config() {
#[tokio::test]
async fn spawn_agent_rejects_when_depth_limit_exceeded() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let max_depth = turn.config.agent_max_depth;
@@ -1835,7 +1837,7 @@ async fn spawn_agent_allows_depth_up_to_configured_max_depth() {
}
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let mut config = (*turn.config).clone();
@@ -1881,7 +1883,7 @@ async fn multi_agent_v2_spawn_agent_ignores_configured_max_depth() {
}
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let mut config = (*turn.config).clone();
config.agent_max_depth = 1;
config
@@ -1989,7 +1991,7 @@ async fn send_input_rejects_invalid_id() {
#[tokio::test]
async fn send_input_reports_missing_agent() {
let (mut session, turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let agent_id = ThreadId::new();
let invocation = invocation(
@@ -2010,7 +2012,7 @@ async fn send_input_reports_missing_agent() {
#[tokio::test]
async fn send_input_interrupts_before_prompt() {
let (mut session, turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let config = turn.config.as_ref().clone();
let thread = manager
@@ -2052,7 +2054,7 @@ async fn send_input_interrupts_before_prompt() {
#[tokio::test]
async fn send_input_accepts_structured_items() {
let (mut session, turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let config = turn.config.as_ref().clone();
let thread = manager
@@ -2126,7 +2128,7 @@ async fn resume_agent_rejects_invalid_id() {
#[tokio::test]
async fn resume_agent_reports_missing_agent() {
let (mut session, turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let agent_id = ThreadId::new();
let invocation = invocation(
@@ -2147,7 +2149,7 @@ async fn resume_agent_reports_missing_agent() {
#[tokio::test]
async fn resume_agent_noops_for_active_agent() {
let (mut session, turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let config = turn.config.as_ref().clone();
let thread = manager
@@ -2186,7 +2188,7 @@ async fn resume_agent_noops_for_active_agent() {
#[tokio::test]
async fn resume_agent_restores_closed_agent_and_accepts_send_input() {
let (mut session, turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let config = turn.config.as_ref().clone();
let thread = manager
@@ -2265,7 +2267,7 @@ async fn resume_agent_restores_closed_agent_and_accepts_send_input() {
#[tokio::test]
async fn resume_agent_rejects_when_depth_limit_exceeded() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let max_depth = turn.config.agent_max_depth;
@@ -2354,7 +2356,7 @@ async fn wait_agent_rejects_empty_targets() {
#[tokio::test]
async fn multi_agent_v2_wait_agent_accepts_timeout_only_argument() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -2493,7 +2495,7 @@ async fn multi_agent_v2_wait_agent_uses_configured_min_timeout() {
#[tokio::test]
async fn wait_agent_returns_not_found_for_missing_agents() {
let (mut session, turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let id_a = ThreadId::new();
let id_b = ThreadId::new();
@@ -2529,7 +2531,7 @@ async fn wait_agent_returns_not_found_for_missing_agents() {
#[tokio::test]
async fn wait_agent_times_out_when_status_is_not_final() {
let (mut session, turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let config = turn.config.as_ref().clone();
let thread = manager
@@ -2572,7 +2574,7 @@ async fn wait_agent_times_out_when_status_is_not_final() {
#[tokio::test]
async fn wait_agent_clamps_short_timeouts_to_minimum() {
let (mut session, turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let config = turn.config.as_ref().clone();
let thread = manager
@@ -2610,7 +2612,7 @@ async fn wait_agent_clamps_short_timeouts_to_minimum() {
#[tokio::test]
async fn wait_agent_returns_final_status_without_timeout() {
let (mut session, turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let config = turn.config.as_ref().clone();
let thread = manager
@@ -2662,7 +2664,7 @@ async fn wait_agent_returns_final_status_without_timeout() {
#[tokio::test]
async fn multi_agent_v2_wait_agent_returns_summary_for_mailbox_activity() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -2753,7 +2755,7 @@ async fn multi_agent_v2_wait_agent_returns_summary_for_mailbox_activity() {
#[tokio::test]
async fn multi_agent_v2_wait_agent_returns_for_already_queued_mail() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -2831,7 +2833,7 @@ async fn multi_agent_v2_wait_agent_returns_for_already_queued_mail() {
#[tokio::test]
async fn multi_agent_v2_wait_agent_wakes_on_any_mailbox_notification() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -2919,7 +2921,7 @@ async fn multi_agent_v2_wait_agent_wakes_on_any_mailbox_notification() {
#[tokio::test]
async fn multi_agent_v2_wait_agent_does_not_return_completed_content() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -3005,7 +3007,7 @@ async fn multi_agent_v2_wait_agent_does_not_return_completed_content() {
#[tokio::test]
async fn multi_agent_v2_close_agent_accepts_task_name_target() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -3064,7 +3066,7 @@ async fn multi_agent_v2_close_agent_accepts_task_name_target() {
#[tokio::test]
async fn multi_agent_v2_close_agent_rejects_root_target_and_id() {
let (mut session, mut turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
let root = manager
.start_thread((*turn.config).clone())
.await
@@ -3112,7 +3114,7 @@ async fn multi_agent_v2_close_agent_rejects_root_target_and_id() {
#[tokio::test]
async fn close_agent_submits_shutdown_and_returns_previous_status() {
let (mut session, turn) = make_session_and_context().await;
let manager = thread_manager();
let manager = thread_manager().await;
session.services.agent_control = manager.agent_control();
let config = turn.config.as_ref().clone();
let thread = manager
@@ -3157,15 +3159,18 @@ async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtr
.features
.enable(Feature::Sqlite)
.expect("test config should allow sqlite");
let state_db = init_state_db(&config).await;
let state_db = init_state_db(&config)
.await
.expect("test config should initialize state db");
let manager = ThreadManager::new(
&config,
AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")),
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
thread_store_from_config(&config, state_db.clone()),
state_db.clone(),
thread_store_from_config(&config, state_db.clone()),
agent_graph_store_from_state_db(state_db.clone()),
);
let parent = manager
+12 -5
View File
@@ -15,7 +15,9 @@ use anyhow::anyhow;
use codex_config::CloudRequirementsLoader;
use codex_core::CodexThread;
use codex_core::ThreadManager;
use codex_core::agent_graph_store_from_state_db;
use codex_core::config::Config;
use codex_core::init_state_db_from_config;
use codex_core::shell::Shell;
use codex_core::shell::get_shell_by_model_provided_path;
use codex_core::thread_store_from_config;
@@ -423,25 +425,30 @@ impl TestCodexBuilder {
environment_manager: Arc<codex_exec_server::EnvironmentManager>,
) -> anyhow::Result<TestCodex> {
let auth = self.auth.clone();
let state_db = codex_core::init_state_db(&config).await;
let thread_manager = if config.model_catalog.is_some() {
let state_db = init_state_db_from_config(&config)
.await
.expect("test codex requires state db");
let thread_store = thread_store_from_config(&config, state_db.clone());
let agent_graph_store = agent_graph_store_from_state_db(state_db.clone());
ThreadManager::new(
&config,
codex_core::test_support::auth_manager_from_auth(auth.clone()),
SessionSource::Exec,
Arc::clone(&environment_manager),
/*analytics_events_client*/ None,
thread_store_from_config(&config, state_db.clone()),
state_db.clone(),
state_db,
thread_store,
agent_graph_store,
)
} else {
codex_core::test_support::thread_manager_with_models_provider_home_and_state(
codex_core::test_support::thread_manager_with_models_provider_and_home(
auth.clone(),
config.model_provider.clone(),
config.codex_home.to_path_buf(),
Arc::clone(&environment_manager),
state_db.clone(),
)
.await
};
let thread_manager = Arc::new(thread_manager);
let user_shell_override = self.user_shell_override.clone();
+11 -8
View File
@@ -5,6 +5,8 @@ use codex_core::NewThread;
use codex_core::Prompt;
use codex_core::ResponseEvent;
use codex_core::ThreadManager;
use codex_core::agent_graph_store_from_state_db;
use codex_core::init_state_db_from_config;
use codex_core::thread_store_from_config;
use codex_features::Feature;
use codex_login::AuthManager;
@@ -770,15 +772,10 @@ async fn includes_conversation_id_and_model_headers_in_request() {
let installation_id =
std::fs::read_to_string(test.codex_home_path().join(INSTALLATION_ID_FILENAME))
.expect("read installation id");
let session_id_string = session_id.to_string();
assert_eq!(request_session_id, session_id_string);
assert_eq!(request_session_id, session_id.to_string());
assert_eq!(request_originator, originator().value);
assert_eq!(request_authorization, "Bearer Test API Key");
assert_eq!(
request_body["prompt_cache_key"].as_str(),
Some(session_id_string.as_str())
);
assert_eq!(
request_body["client_metadata"]["x-codex-installation-id"].as_str(),
Some(installation_id.as_str())
@@ -1106,14 +1103,20 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() {
Ok(None) => panic!("No CodexAuth found in codex_home"),
Err(e) => panic!("Failed to load CodexAuth: {e}"),
};
let state_db = init_state_db_from_config(&config)
.await
.expect("client test requires state db");
let thread_store = thread_store_from_config(&config, state_db.clone());
let agent_graph_store = agent_graph_store_from_state_db(state_db.clone());
let thread_manager = ThreadManager::new(
&config,
auth_manager,
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
thread_store_from_config(&config, /*state_db*/ None),
/*state_db*/ None,
state_db,
thread_store,
agent_graph_store,
);
let NewThread { thread: codex, .. } = thread_manager
.start_thread(config.clone())
@@ -13,6 +13,8 @@ use codex_protocol::protocol::SessionMeta;
use codex_protocol::protocol::SessionMetaLine;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::UserMessageEvent;
use codex_rollout::RolloutConfig;
use codex_rollout::state_db::StateDbHandle;
use pretty_assertions::assert_eq;
use std::io;
use std::path::Path;
@@ -26,6 +28,27 @@ async fn read_config_toml(codex_home: &Path) -> io::Result<ConfigToml> {
toml::from_str(&contents).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))
}
async fn state_db_for_test(codex_home: &Path) -> io::Result<StateDbHandle> {
let config = RolloutConfig {
codex_home: codex_home.to_path_buf(),
sqlite_home: codex_home.to_path_buf(),
cwd: codex_home.to_path_buf(),
model_provider_id: "openai".to_string(),
generate_memories: false,
};
codex_rollout::state_db::try_init(&config)
.await
.map_err(io::Error::other)
}
async fn run_migration(
codex_home: &Path,
config_toml: &ConfigToml,
) -> io::Result<PersonalityMigrationStatus> {
let state_db = state_db_for_test(codex_home).await?;
maybe_migrate_personality(codex_home, config_toml, state_db).await
}
async fn write_session_with_user_event(codex_home: &Path) -> io::Result<()> {
let thread_id = ThreadId::new();
let dir = codex_home
@@ -141,8 +164,7 @@ async fn migration_marker_exists_no_sessions_no_change() -> io::Result<()> {
let marker_path = temp.path().join(PERSONALITY_MIGRATION_FILENAME);
tokio::fs::write(&marker_path, "v1\n").await?;
let status =
maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?;
let status = run_migration(temp.path(), &ConfigToml::default()).await?;
assert_eq!(status, PersonalityMigrationStatus::SkippedMarker);
assert_eq!(
@@ -156,8 +178,7 @@ async fn migration_marker_exists_no_sessions_no_change() -> io::Result<()> {
async fn no_marker_no_sessions_no_change() -> io::Result<()> {
let temp = TempDir::new()?;
let status =
maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?;
let status = run_migration(temp.path(), &ConfigToml::default()).await?;
assert_eq!(status, PersonalityMigrationStatus::SkippedNoSessions);
assert_eq!(
@@ -176,8 +197,7 @@ async fn no_marker_sessions_sets_personality() -> io::Result<()> {
let temp = TempDir::new()?;
write_session_with_user_event(temp.path()).await?;
let status =
maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?;
let status = run_migration(temp.path(), &ConfigToml::default()).await?;
assert_eq!(status, PersonalityMigrationStatus::Applied);
assert_eq!(
@@ -197,7 +217,7 @@ async fn no_marker_sessions_preserves_existing_config_fields() -> io::Result<()>
tokio::fs::write(temp.path().join("config.toml"), "model = \"gpt-5.4\"\n").await?;
let config_toml = read_config_toml(temp.path()).await?;
let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?;
let status = run_migration(temp.path(), &config_toml).await?;
assert_eq!(status, PersonalityMigrationStatus::Applied);
let persisted = read_config_toml(temp.path()).await?;
@@ -211,8 +231,7 @@ async fn no_marker_meta_only_rollout_is_treated_as_no_sessions() -> io::Result<(
let temp = TempDir::new()?;
write_session_with_meta_only(temp.path()).await?;
let status =
maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?;
let status = run_migration(temp.path(), &ConfigToml::default()).await?;
assert_eq!(status, PersonalityMigrationStatus::SkippedNoSessions);
assert_eq!(
@@ -232,7 +251,7 @@ async fn no_marker_explicit_global_personality_skips_migration() -> io::Result<(
write_session_with_user_event(temp.path()).await?;
let config_toml = parse_config_toml("personality = \"friendly\"\n")?;
let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?;
let status = run_migration(temp.path(), &config_toml).await?;
assert_eq!(
status,
@@ -262,7 +281,7 @@ personality = "friendly"
"#,
)?;
let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?;
let status = run_migration(temp.path(), &config_toml).await?;
assert_eq!(
status,
@@ -285,7 +304,7 @@ async fn marker_short_circuits_invalid_profile_resolution() -> io::Result<()> {
tokio::fs::write(temp.path().join(PERSONALITY_MIGRATION_FILENAME), "v1\n").await?;
let config_toml = parse_config_toml("profile = \"missing\"\n")?;
let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?;
let status = run_migration(temp.path(), &config_toml).await?;
assert_eq!(status, PersonalityMigrationStatus::SkippedMarker);
Ok(())
@@ -296,7 +315,7 @@ async fn invalid_selected_profile_returns_error_and_does_not_write_marker() -> i
let temp = TempDir::new()?;
let config_toml = parse_config_toml("profile = \"missing\"\n")?;
let err = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None)
let err = run_migration(temp.path(), &config_toml)
.await
.expect_err("missing profile should fail");
@@ -313,10 +332,8 @@ async fn applied_migration_is_idempotent_on_second_run() -> io::Result<()> {
let temp = TempDir::new()?;
write_session_with_user_event(temp.path()).await?;
let first_status =
maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?;
let second_status =
maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?;
let first_status = run_migration(temp.path(), &ConfigToml::default()).await?;
let second_status = run_migration(temp.path(), &ConfigToml::default()).await?;
assert_eq!(first_status, PersonalityMigrationStatus::Applied);
assert_eq!(second_status, PersonalityMigrationStatus::SkippedMarker);
@@ -330,8 +347,7 @@ async fn no_marker_archived_sessions_sets_personality() -> io::Result<()> {
let temp = TempDir::new()?;
write_archived_session_with_user_event(temp.path()).await?;
let status =
maybe_migrate_personality(temp.path(), &ConfigToml::default(), /*state_db*/ None).await?;
let status = run_migration(temp.path(), &ConfigToml::default()).await?;
assert_eq!(status, PersonalityMigrationStatus::Applied);
assert_eq!(
@@ -29,7 +29,6 @@ async fn build_prompt_input_includes_context_and_user_message() -> Result<()> {
text: "hello from debug prompt".to_string(),
text_elements: Vec::new(),
}],
/*state_db*/ None,
)
.await?;
+2 -1
View File
@@ -95,7 +95,8 @@ async fn emits_warning_when_resumed_model_differs() {
let thread_manager = codex_core::test_support::thread_manager_with_models_provider(
CodexAuth::from_api_key("test"),
config.model_provider.clone(),
);
)
.await;
let auth_manager =
codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("test"));
+10 -2
View File
@@ -3,6 +3,8 @@
use anyhow::Result;
use codex_core::ThreadManager;
use codex_core::agent_graph_store_from_state_db;
use codex_core::init_state_db_from_config;
use codex_core::thread_store_from_config;
use codex_exec_server::CreateDirectoryOptions;
use codex_exec_server::EnvironmentManager;
@@ -235,6 +237,11 @@ async fn list_skills_skips_cwd_roots_when_environment_disabled() -> Result<()> {
let mut config = load_default_config_for_test(&codex_home).await;
config.cwd = AbsolutePathBuf::from_absolute_path_checked(cwd.path())?;
let state_db = init_state_db_from_config(&config)
.await
.expect("skills test requires state db");
let thread_store = thread_store_from_config(&config, state_db.clone());
let agent_graph_store = agent_graph_store_from_state_db(state_db.clone());
let thread_manager = ThreadManager::new(
&config,
codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("dummy")),
@@ -246,8 +253,9 @@ async fn list_skills_skips_cwd_roots_when_environment_disabled() -> Result<()> {
)?,
)),
/*analytics_events_client*/ None,
thread_store_from_config(&config, /*state_db*/ None),
/*state_db*/ None,
state_db,
thread_store,
agent_graph_store,
);
let new_thread = thread_manager.start_thread(config.clone()).await?;
let cwd = config.cwd.to_path_buf();
@@ -34,7 +34,8 @@ async fn emits_warning_when_unstable_features_enabled_via_config() {
let thread_manager = codex_core::test_support::thread_manager_with_models_provider(
CodexAuth::from_api_key("test"),
config.model_provider.clone(),
);
)
.await;
let auth_manager =
codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("test"));
@@ -81,7 +82,8 @@ async fn suppresses_warning_when_configured() {
let thread_manager = codex_core::test_support::thread_manager_with_models_provider(
CodexAuth::from_api_key("test"),
config.model_provider.clone(),
);
)
.await;
let auth_manager =
codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("test"));
+5 -3
View File
@@ -83,7 +83,6 @@ pub async fn run_main(
std::io::Error::new(ErrorKind::InvalidData, format!("error loading config: {e}"))
})?;
set_default_client_residency_requirement(config.enforce_residency.value());
let state_db = codex_core::init_state_db(&config).await;
let otel = codex_core::otel_init::build_provider(
&config,
@@ -140,15 +139,18 @@ pub async fn run_main(
// Task: process incoming messages.
let processor_handle = tokio::spawn({
let outgoing_message_sender = OutgoingMessageSender::new(outgoing_tx);
let mut processor = MessageProcessor::new(
let processor = MessageProcessor::new(
outgoing_message_sender,
arg0_paths,
Arc::new(config),
environment_manager,
state_db,
)
.await;
async move {
let Some(mut processor) = processor else {
error!("failed to initialize MCP processor");
return;
};
while let Some(msg) = incoming_rx.recv().await {
match msg {
JsonRpcMessage::Request(r) => processor.process_request(r).await,
+11 -7
View File
@@ -2,9 +2,10 @@ use std::collections::HashMap;
use std::sync::Arc;
use codex_arg0::Arg0DispatchPaths;
use codex_core::StateDbHandle;
use codex_core::ThreadManager;
use codex_core::agent_graph_store_from_state_db;
use codex_core::config::Config;
use codex_core::init_state_db_from_config;
use codex_core::thread_store_from_config;
use codex_exec_server::EnvironmentManager;
use codex_login::AuthManager;
@@ -54,30 +55,33 @@ impl MessageProcessor {
arg0_paths: Arg0DispatchPaths,
config: Arc<Config>,
environment_manager: Arc<EnvironmentManager>,
state_db: Option<StateDbHandle>,
) -> Self {
) -> Option<Self> {
let outgoing = Arc::new(outgoing);
let auth_manager = AuthManager::shared_from_config(
config.as_ref(),
/*enable_codex_api_key_env*/ false,
)
.await;
let state_db = init_state_db_from_config(config.as_ref()).await?;
let thread_store = thread_store_from_config(config.as_ref(), state_db.clone());
let agent_graph_store = agent_graph_store_from_state_db(state_db.clone());
let thread_manager = Arc::new(ThreadManager::new(
config.as_ref(),
auth_manager,
SessionSource::Mcp,
environment_manager,
/*analytics_events_client*/ None,
thread_store_from_config(config.as_ref(), state_db.clone()),
state_db.clone(),
state_db,
thread_store,
agent_graph_store,
));
Self {
Some(Self {
outgoing,
initialized: false,
arg0_paths,
thread_manager,
running_requests_id_to_codex_uuid: Arc::new(Mutex::new(HashMap::new())),
}
})
}
pub(crate) async fn process_request(&mut self, request: JsonRpcRequest<ClientRequest>) {
+8 -3
View File
@@ -52,10 +52,11 @@ use codex_core_api::TuiNotificationSettings;
use codex_core_api::UriBasedFileOpener;
use codex_core_api::UserInput;
use codex_core_api::WebSearchMode;
use codex_core_api::agent_graph_store_from_state_db;
use codex_core_api::arg0_dispatch_or_else;
use codex_core_api::built_in_model_providers;
use codex_core_api::find_codex_home;
use codex_core_api::init_state_db;
use codex_core_api::init_state_db_from_config;
use codex_core_api::item_event_to_server_notification;
use codex_core_api::set_default_originator;
use codex_core_api::thread_store_from_config;
@@ -104,7 +105,6 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> {
};
let config = new_config(args.model, arg0_paths)?;
let state_db = init_state_db(&config).await;
let auth_manager =
AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false).await;
@@ -112,7 +112,11 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> {
config.codex_self_exe.clone(),
config.codex_linux_sandbox_exe.clone(),
)?;
let Some(state_db) = init_state_db_from_config(&config).await else {
bail!("thread manager sample requires state db");
};
let thread_store = thread_store_from_config(&config, state_db.clone());
let agent_graph_store = agent_graph_store_from_state_db(state_db.clone());
let environment_manager =
Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await);
let thread_manager = ThreadManager::new(
@@ -121,8 +125,9 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> {
SessionSource::Exec,
environment_manager,
/*analytics_events_client*/ None,
Arc::clone(&thread_store),
state_db,
Arc::clone(&thread_store),
agent_graph_store,
);
let NewThread {
@@ -13,11 +13,11 @@ pub(super) async fn archive_thread(
params: ArchiveThreadParams,
) -> ThreadStoreResult<()> {
let thread_id = params.thread_id;
let state_db_ctx = store.state_db().await;
let state_db = store.state_db();
let rollout_path = find_thread_path_by_id_str(
store.config.codex_home.as_path(),
&thread_id.to_string(),
state_db_ctx.as_deref(),
Some(state_db.as_ref()),
)
.await
.map_err(|err| ThreadStoreError::InvalidRequest {
@@ -52,11 +52,10 @@ pub(super) async fn archive_thread(
}
})?;
if let Some(ctx) = state_db_ctx {
let _ = ctx
.mark_archived(thread_id, archived_path.as_path(), Utc::now())
.await;
}
let _ = store
.state_db()
.mark_archived(thread_id, archived_path.as_path(), Utc::now())
.await;
Ok(())
}
@@ -75,13 +74,15 @@ mod tests {
use crate::ThreadSortKey;
use crate::ThreadStore;
use crate::local::LocalThreadStore;
use crate::local::test_support::init_test_state_db;
use crate::local::test_support::test_config;
use crate::local::test_support::test_store;
use crate::local::test_support::write_session_file;
#[tokio::test]
async fn archive_thread_moves_rollout_to_archived_collection() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(201);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let active_path =
@@ -127,21 +128,12 @@ mod tests {
async fn archive_thread_updates_sqlite_metadata_when_present() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(202);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let active_path =
write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
runtime
.mark_backfill_complete(/*last_watermark*/ None)
.await
.expect("backfill should be complete");
let mut builder = codex_state::ThreadMetadataBuilder::new(
thread_id,
active_path.clone(),
@@ -22,12 +22,12 @@ pub(super) async fn create_thread(
})?;
let config = RolloutConfig {
codex_home: store.config.codex_home.clone(),
sqlite_home: store.config.sqlite_home.clone(),
sqlite_home: store.sqlite_home(),
cwd,
model_provider_id: params.metadata.model_provider.clone(),
generate_memories: matches!(params.metadata.memory_mode, ThreadMemoryMode::Enabled),
};
let state_db_ctx = store.state_db().await;
let state_db_ctx = Some(store.state_db());
let recorder = RolloutRecorder::new(
&config,
RolloutRecorderParams::new(
+32 -36
View File
@@ -39,16 +39,16 @@ pub(super) async fn list_threads(
SortDirection::Asc => codex_rollout::SortDirection::Asc,
SortDirection::Desc => codex_rollout::SortDirection::Desc,
};
let state_db = store.state_db().await;
let rollout_config = RolloutConfig {
codex_home: store.config.codex_home.clone(),
sqlite_home: store.config.sqlite_home.clone(),
sqlite_home: store.sqlite_home(),
cwd: store.config.codex_home.clone(),
model_provider_id: store.config.default_model_provider_id.clone(),
generate_memories: false,
};
let state_db_ctx = Some(store.state_db());
let page = list_rollout_threads(
state_db,
state_db_ctx,
&rollout_config,
store.config.default_model_provider_id.as_str(),
&params,
@@ -80,14 +80,13 @@ pub(super) async fn list_threads(
.map(|thread| thread.thread_id)
.collect::<HashSet<_>>();
let mut names = HashMap::<ThreadId, String>::with_capacity(thread_ids.len());
if let Some(state_db_ctx) = store.state_db().await {
for &thread_id in &thread_ids {
let Ok(Some(metadata)) = state_db_ctx.get_thread(thread_id).await else {
continue;
};
if let Some(title) = distinct_thread_metadata_title(&metadata) {
names.insert(thread_id, title);
}
let state_db_ctx = store.state_db();
for &thread_id in &thread_ids {
let Ok(Some(metadata)) = state_db_ctx.get_thread(thread_id).await else {
continue;
};
if let Some(title) = distinct_thread_metadata_title(&metadata) {
names.insert(thread_id, title);
}
}
if names.len() < thread_ids.len()
@@ -108,9 +107,9 @@ pub(super) async fn list_threads(
}
async fn list_rollout_threads(
state_db: Option<codex_rollout::StateDbHandle>,
state_db_ctx: Option<codex_rollout::StateDbHandle>,
config: &RolloutConfig,
default_model_provider_id: &str,
default_model_provider: &str,
params: &ListThreadsParams,
cursor: Option<&codex_rollout::Cursor>,
sort_key: codex_rollout::ThreadSortKey,
@@ -118,7 +117,7 @@ async fn list_rollout_threads(
) -> ThreadStoreResult<codex_rollout::ThreadsPage> {
let page = if params.use_state_db_only && params.archived {
RolloutRecorder::list_archived_threads_from_state_db(
state_db,
state_db_ctx.clone(),
config,
params.page_size,
cursor,
@@ -127,13 +126,13 @@ async fn list_rollout_threads(
params.allowed_sources.as_slice(),
params.model_providers.as_deref(),
params.cwd_filters.as_deref(),
default_model_provider_id,
default_model_provider,
params.search_term.as_deref(),
)
.await
} else if params.use_state_db_only {
RolloutRecorder::list_threads_from_state_db(
state_db,
state_db_ctx.clone(),
config,
params.page_size,
cursor,
@@ -142,13 +141,13 @@ async fn list_rollout_threads(
params.allowed_sources.as_slice(),
params.model_providers.as_deref(),
params.cwd_filters.as_deref(),
default_model_provider_id,
default_model_provider,
params.search_term.as_deref(),
)
.await
} else if params.archived {
RolloutRecorder::list_archived_threads(
state_db,
state_db_ctx.clone(),
config,
params.page_size,
cursor,
@@ -157,13 +156,13 @@ async fn list_rollout_threads(
params.allowed_sources.as_slice(),
params.model_providers.as_deref(),
params.cwd_filters.as_deref(),
default_model_provider_id,
default_model_provider,
params.search_term.as_deref(),
)
.await
} else {
RolloutRecorder::list_threads(
state_db,
state_db_ctx,
config,
params.page_size,
cursor,
@@ -172,7 +171,7 @@ async fn list_rollout_threads(
params.allowed_sources.as_slice(),
params.model_providers.as_deref(),
params.cwd_filters.as_deref(),
default_model_provider_id,
default_model_provider,
params.search_term.as_deref(),
)
.await
@@ -195,7 +194,9 @@ mod tests {
use super::*;
use crate::ThreadStore;
use crate::local::LocalThreadStore;
use crate::local::test_support::init_test_state_db;
use crate::local::test_support::test_config;
use crate::local::test_support::test_store;
use crate::local::test_support::write_archived_session_file;
use crate::local::test_support::write_session_file;
use crate::local::test_support::write_session_file_with;
@@ -203,7 +204,7 @@ mod tests {
#[tokio::test]
async fn list_threads_uses_default_provider_when_rollout_omits_provider() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
write_session_file_with(
home.path(),
home.path().join("sessions/2025/01/03"),
@@ -238,22 +239,13 @@ mod tests {
async fn list_threads_preserves_sqlite_title_search_results() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(103);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path = home.path().join("rollout-title-search.jsonl");
fs::write(&rollout_path, "").expect("placeholder rollout file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
runtime
.mark_backfill_complete(/*last_watermark*/ None)
.await
.expect("backfill should be complete");
let created_at = Utc::now();
let mut builder = codex_state::ThreadMetadataBuilder::new(
thread_id,
@@ -267,6 +259,10 @@ mod tests {
let mut metadata = builder.build(config.default_model_provider_id.as_str());
metadata.title = "needle title".to_string();
metadata.first_user_message = Some("plain preview".to_string());
runtime
.mark_backfill_complete(/*last_watermark*/ None)
.await
.expect("backfill should be complete");
runtime
.upsert_thread(&metadata)
.await
@@ -303,7 +299,7 @@ mod tests {
#[tokio::test]
async fn list_threads_selects_active_or_archived_collection() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let active_uuid = Uuid::from_u128(105);
let archived_uuid = Uuid::from_u128(106);
write_session_file(home.path(), "2025-01-03T12-00-00", active_uuid)
@@ -372,7 +368,7 @@ mod tests {
async fn list_threads_returns_local_rollout_summary() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let store = LocalThreadStore::new(config, /*state_db*/ None);
let store = LocalThreadStore::new(config.clone(), init_test_state_db(&config).await);
let uuid = Uuid::from_u128(101);
let path =
write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file");
@@ -411,7 +407,7 @@ mod tests {
#[tokio::test]
async fn list_threads_rejects_invalid_cursor() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let err = store
.list_threads(ListThreadsParams {
@@ -66,12 +66,12 @@ pub(super) async fn resume_thread(
})?;
let config = RolloutConfig {
codex_home: store.config.codex_home.clone(),
sqlite_home: store.config.sqlite_home.clone(),
sqlite_home: store.sqlite_home(),
cwd,
model_provider_id: params.metadata.model_provider.clone(),
generate_memories: matches!(params.metadata.memory_mode, ThreadMemoryMode::Enabled),
};
let state_db_ctx = store.state_db().await;
let state_db_ctx = Some(store.state_db());
let recorder = RolloutRecorder::new(
&config,
RolloutRecorderParams::resume(
+24 -18
View File
@@ -41,7 +41,7 @@ use crate::UpdateThreadMetadataParams;
pub struct LocalThreadStore {
pub(super) config: LocalThreadStoreConfig,
live_recorders: Arc<Mutex<HashMap<ThreadId, RolloutRecorder>>>,
state_db: Option<StateDbHandle>,
state_db: StateDbHandle,
}
/// Process-scoped configuration for local thread storage.
@@ -51,7 +51,6 @@ pub struct LocalThreadStore {
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct LocalThreadStoreConfig {
pub codex_home: PathBuf,
pub sqlite_home: PathBuf,
/// Provider used only when older local metadata does not contain one.
pub default_model_provider_id: String,
}
@@ -60,7 +59,6 @@ impl LocalThreadStoreConfig {
pub fn from_config(config: &impl codex_rollout::RolloutConfigView) -> Self {
Self {
codex_home: config.codex_home().to_path_buf(),
sqlite_home: config.sqlite_home().to_path_buf(),
default_model_provider_id: config.model_provider_id().to_string(),
}
}
@@ -75,8 +73,9 @@ impl std::fmt::Debug for LocalThreadStore {
}
impl LocalThreadStore {
/// Create a local store using an already initialized state DB handle.
pub fn new(config: LocalThreadStoreConfig, state_db: Option<StateDbHandle>) -> Self {
/// Create a local store from process-scoped local storage configuration and
/// the caller-provided shared state DB handle.
pub fn new(config: LocalThreadStoreConfig, state_db: StateDbHandle) -> Self {
Self {
config,
live_recorders: Arc::new(Mutex::new(HashMap::new())),
@@ -85,10 +84,14 @@ impl LocalThreadStore {
}
/// Return the state DB handle used by local rollout writers.
pub async fn state_db(&self) -> Option<StateDbHandle> {
pub fn state_db(&self) -> StateDbHandle {
self.state_db.clone()
}
pub(super) fn sqlite_home(&self) -> PathBuf {
self.state_db.codex_home().to_path_buf()
}
/// Read a local rollout-backed thread by path.
pub async fn read_thread_by_rollout_path(
&self,
@@ -282,14 +285,16 @@ mod tests {
use super::*;
use crate::ThreadEventPersistenceMode;
use crate::ThreadPersistenceMetadata;
use crate::local::test_support::init_test_state_db;
use crate::local::test_support::test_config;
use crate::local::test_support::test_store;
use crate::local::test_support::write_archived_session_file;
use crate::local::test_support::write_session_file;
#[tokio::test]
async fn live_writer_lifecycle_writes_and_closes() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let thread_id = ThreadId::default();
store
@@ -338,7 +343,7 @@ mod tests {
#[tokio::test]
async fn create_thread_rejects_missing_cwd() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let thread_id = ThreadId::default();
let mut params = create_thread_params(thread_id);
params.metadata.cwd = None;
@@ -358,7 +363,7 @@ mod tests {
#[tokio::test]
async fn discard_thread_drops_unmaterialized_live_writer() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let thread_id = ThreadId::default();
store
@@ -396,8 +401,9 @@ mod tests {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let thread_id = ThreadId::default();
let state_db = init_test_state_db(&config).await;
let first_store = LocalThreadStore::new(config.clone(), /*state_db*/ None);
let first_store = LocalThreadStore::new(config.clone(), state_db.clone());
first_store
.create_thread(create_thread_params(thread_id))
.await
@@ -426,7 +432,7 @@ mod tests {
.await
.expect("shutdown initial writer");
let resumed_store = LocalThreadStore::new(config, /*state_db*/ None);
let resumed_store = LocalThreadStore::new(config, state_db);
resumed_store
.resume_thread(ResumeThreadParams {
thread_id,
@@ -457,7 +463,7 @@ mod tests {
#[tokio::test]
async fn create_thread_rejects_duplicate_live_writer() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let thread_id = ThreadId::default();
store
@@ -477,7 +483,7 @@ mod tests {
#[tokio::test]
async fn resume_thread_rejects_duplicate_live_writer() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let thread_id = ThreadId::default();
store
@@ -506,7 +512,7 @@ mod tests {
#[tokio::test]
async fn resume_thread_rejects_missing_cwd() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = uuid::Uuid::from_u128(407);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path =
@@ -535,7 +541,7 @@ mod tests {
async fn load_history_uses_live_writer_rollout_path() {
let home = TempDir::new().expect("temp dir");
let external_home = TempDir::new().expect("external temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = uuid::Uuid::from_u128(404);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path = write_session_file(external_home.path(), "2025-01-04T10-00-00", uuid)
@@ -584,7 +590,7 @@ mod tests {
async fn read_thread_uses_live_writer_rollout_path_for_external_resume() {
let home = TempDir::new().expect("temp dir");
let external_home = TempDir::new().expect("external temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = uuid::Uuid::from_u128(406);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path = write_session_file(external_home.path(), "2025-01-04T11-00-00", uuid)
@@ -623,7 +629,7 @@ mod tests {
#[tokio::test]
async fn load_history_uses_live_writer_rollout_path_for_archived_source() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = uuid::Uuid::from_u128(405);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path = write_archived_session_file(home.path(), "2025-01-04T10-30-00", uuid)
@@ -691,7 +697,7 @@ mod tests {
#[tokio::test]
async fn read_thread_by_rollout_path_includes_history() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let thread_id = ThreadId::default();
store
+33 -77
View File
@@ -176,12 +176,12 @@ async fn resolve_rollout_path(
return Ok(Some(path));
}
let state_db_ctx = store.state_db().await;
let state_db = store.state_db();
if include_archived {
match find_thread_path_by_id_str(
store.config.codex_home.as_path(),
&thread_id.to_string(),
state_db_ctx.as_deref(),
Some(state_db.as_ref()),
)
.await
.map_err(|err| ThreadStoreError::InvalidRequest {
@@ -191,7 +191,7 @@ async fn resolve_rollout_path(
None => find_archived_thread_path_by_id_str(
store.config.codex_home.as_path(),
&thread_id.to_string(),
state_db_ctx.as_deref(),
Some(state_db.as_ref()),
)
.await
.map_err(|err| ThreadStoreError::InvalidRequest {
@@ -202,7 +202,7 @@ async fn resolve_rollout_path(
find_thread_path_by_id_str(
store.config.codex_home.as_path(),
&thread_id.to_string(),
state_db_ctx.as_deref(),
Some(state_db.as_ref()),
)
.await
.map_err(|err| ThreadStoreError::InvalidRequest {
@@ -260,8 +260,7 @@ async fn read_sqlite_metadata(
store: &LocalThreadStore,
thread_id: codex_protocol::ThreadId,
) -> Option<ThreadMetadata> {
let runtime = store.state_db().await?;
runtime.get_thread(thread_id).await.ok().flatten()
store.state_db().get_thread(thread_id).await.ok().flatten()
}
async fn stored_thread_from_sqlite_metadata(
@@ -412,7 +411,9 @@ mod tests {
use super::*;
use crate::ThreadStore;
use crate::local::LocalThreadStore;
use crate::local::test_support::init_test_state_db;
use crate::local::test_support::test_config;
use crate::local::test_support::test_store;
use crate::local::test_support::write_archived_session_file;
use crate::local::test_support::write_session_file;
use crate::local::test_support::write_session_file_with_fork;
@@ -420,7 +421,7 @@ mod tests {
#[tokio::test]
async fn read_thread_returns_active_rollout_summary() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(205);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let active_path =
@@ -448,7 +449,7 @@ mod tests {
#[tokio::test]
async fn read_thread_returns_rollout_path_summary() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(211);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let active_path =
@@ -479,17 +480,12 @@ mod tests {
async fn read_thread_by_rollout_path_prefers_sqlite_git_info() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(223);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let active_path =
write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let mut builder = ThreadMetadataBuilder::new(
thread_id,
active_path.clone(),
@@ -527,7 +523,7 @@ mod tests {
#[tokio::test]
async fn read_thread_returns_archived_rollout_when_requested() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(207);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let archived_path = write_archived_session_file(home.path(), "2025-01-03T12-00-00", uuid)
@@ -568,7 +564,7 @@ mod tests {
#[tokio::test]
async fn read_thread_prefers_active_rollout_over_archived() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(208);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let active_path =
@@ -593,7 +589,7 @@ mod tests {
#[tokio::test]
async fn read_thread_returns_forked_from_id() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(209);
let parent_uuid = Uuid::from_u128(210);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
@@ -626,17 +622,12 @@ mod tests {
async fn read_thread_applies_sqlite_thread_name() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(212);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path =
write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let mut builder =
ThreadMetadataBuilder::new(thread_id, rollout_path, Utc::now(), SessionSource::Cli);
builder.model_provider = Some(config.default_model_provider_id.clone());
@@ -666,13 +657,8 @@ mod tests {
async fn read_thread_preserves_rollout_cwd_when_sqlite_metadata_exists() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(224);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let day_dir = home.path().join("sessions/2025/01/03");
@@ -741,7 +727,7 @@ mod tests {
#[tokio::test]
async fn read_thread_uses_legacy_thread_name_when_sqlite_title_is_missing() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(213);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file");
@@ -765,6 +751,8 @@ mod tests {
async fn read_thread_uses_sqlite_metadata_for_rollout_without_user_preview() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(217);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let day_dir = home.path().join("sessions/2025/01/03");
@@ -786,13 +774,6 @@ mod tests {
});
writeln!(file, "{meta}").expect("write session meta");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let mut builder = ThreadMetadataBuilder::new(
thread_id,
rollout_path.clone(),
@@ -835,18 +816,13 @@ mod tests {
let home = TempDir::new().expect("temp dir");
let external = TempDir::new().expect("external temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(220);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path =
write_session_file(home.path(), "2025-01-03T12-00-00", uuid).expect("session file");
let stale_path = external.path().join("missing-rollout.jsonl");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let mut builder = ThreadMetadataBuilder::new(
thread_id,
stale_path.clone(),
@@ -884,6 +860,8 @@ mod tests {
let home = TempDir::new().expect("temp dir");
let external = TempDir::new().expect("external temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(221);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path =
@@ -891,13 +869,6 @@ mod tests {
let other_uuid = Uuid::from_u128(222);
let stale_path = write_session_file(external.path(), "2025-01-04T12-00-00", other_uuid)
.expect("other session file");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let mut builder =
ThreadMetadataBuilder::new(thread_id, stale_path, Utc::now(), SessionSource::Cli);
builder.model_provider = Some("wrong-sqlite-provider".to_string());
@@ -929,7 +900,7 @@ mod tests {
#[tokio::test]
async fn read_thread_uses_session_meta_for_rollout_without_user_preview_or_sqlite_metadata() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(218);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let day_dir = home.path().join("sessions/2025/01/03");
@@ -984,18 +955,13 @@ mod tests {
let home = TempDir::new().expect("temp dir");
let external = TempDir::new().expect("external temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(214);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path = external
.path()
.join(format!("rollout-2025-01-03T12-00-00-{uuid}.jsonl"));
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let mut builder = ThreadMetadataBuilder::new(
thread_id,
rollout_path.clone(),
@@ -1042,20 +1008,15 @@ mod tests {
let home = TempDir::new().expect("temp dir");
let external = TempDir::new().expect("external temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(216);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let rollout_path = external
.path()
.join(format!("rollout-2025-01-03T12-00-00-{uuid}.jsonl"));
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let mut builder =
ThreadMetadataBuilder::new(thread_id, rollout_path, Utc::now(), SessionSource::Cli);
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
builder.archived_at = Some(Utc::now());
let mut metadata = builder.build(config.default_model_provider_id.as_str());
metadata.first_user_message = Some("Archived SQLite preview".to_string());
@@ -1098,17 +1059,12 @@ mod tests {
async fn read_thread_sqlite_fallback_loads_archived_history() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(219);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let archived_path = write_archived_session_file(home.path(), "2025-01-03T12-00-00", uuid)
.expect("archived session file");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let mut builder = ThreadMetadataBuilder::new(
thread_id,
archived_path.clone(),
@@ -1144,7 +1100,7 @@ mod tests {
#[tokio::test]
async fn read_thread_fails_without_rollout() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(206);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
@@ -4,18 +4,34 @@ use std::path::Path;
use std::path::PathBuf;
use codex_rollout::ARCHIVED_SESSIONS_SUBDIR;
use codex_rollout::StateDbHandle;
use uuid::Uuid;
use super::LocalThreadStore;
use super::LocalThreadStoreConfig;
pub(super) fn test_config(codex_home: &Path) -> LocalThreadStoreConfig {
LocalThreadStoreConfig {
codex_home: codex_home.to_path_buf(),
sqlite_home: codex_home.to_path_buf(),
default_model_provider_id: "test-provider".to_string(),
}
}
pub(super) async fn init_test_state_db(config: &LocalThreadStoreConfig) -> StateDbHandle {
codex_state::StateRuntime::init(
config.codex_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize")
}
pub(super) async fn test_store(codex_home: &Path) -> LocalThreadStore {
let config = test_config(codex_home);
let state_db = init_test_state_db(&config).await;
LocalThreadStore::new(config, state_db)
}
pub(super) fn write_session_file(root: &Path, ts: &str, uuid: Uuid) -> std::io::Result<PathBuf> {
write_session_file_with(
root,
@@ -17,11 +17,11 @@ pub(super) async fn unarchive_thread(
params: ArchiveThreadParams,
) -> ThreadStoreResult<StoredThread> {
let thread_id = params.thread_id;
let state_db_ctx = store.state_db().await;
let state_db = store.state_db();
let archived_path = find_archived_thread_path_by_id_str(
store.config.codex_home.as_path(),
&thread_id.to_string(),
state_db_ctx.as_deref(),
Some(state_db.as_ref()),
)
.await
.map_err(|err| ThreadStoreError::InvalidRequest {
@@ -73,11 +73,10 @@ pub(super) async fn unarchive_thread(
message: format!("failed to update unarchived thread timestamp: {err}"),
})?;
if let Some(ctx) = state_db_ctx {
let _ = ctx
.mark_unarchived(thread_id, restored_path.as_path())
.await;
}
let _ = store
.state_db()
.mark_unarchived(thread_id, restored_path.as_path())
.await;
let item = read_thread_item_from_rollout(restored_path.clone())
.await
@@ -112,13 +111,15 @@ mod tests {
use super::*;
use crate::ThreadStore;
use crate::local::LocalThreadStore;
use crate::local::test_support::init_test_state_db;
use crate::local::test_support::test_config;
use crate::local::test_support::test_store;
use crate::local::test_support::write_archived_session_file;
#[tokio::test]
async fn unarchive_thread_restores_rollout_and_returns_updated_thread() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(203);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let archived_path = write_archived_session_file(home.path(), "2025-01-03T13-00-00", uuid)
@@ -149,21 +150,12 @@ mod tests {
async fn unarchive_thread_updates_sqlite_metadata_when_present() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(204);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let archived_path = write_archived_session_file(home.path(), "2025-01-03T13-00-00", uuid)
.expect("archived session file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
runtime
.mark_backfill_complete(/*last_watermark*/ None)
.await
.expect("backfill should be complete");
let mut builder = codex_state::ThreadMetadataBuilder::new(
thread_id,
archived_path.clone(),
@@ -59,9 +59,8 @@ pub(super) async fn update_thread_metadata(
.await?;
}
let state_db_ctx = store.state_db().await;
codex_rollout::state_db::reconcile_rollout(
state_db_ctx.as_deref(),
Some(store.state_db()).as_deref(),
resolved_rollout_path.path.as_path(),
store.config.default_model_provider_id.as_str(),
/*builder*/ None,
@@ -73,11 +72,7 @@ pub(super) async fn update_thread_metadata(
let resolved_git_info = match git_info {
Some(git_info) => {
let Some(state_db) = store.state_db().await else {
return Err(ThreadStoreError::Internal {
message: format!("sqlite state db unavailable for thread {thread_id}"),
});
};
let state_db = store.state_db();
let metadata =
state_db
.get_thread(thread_id)
@@ -157,11 +152,7 @@ async fn apply_thread_git_info(
branch: &Option<String>,
origin_url: &Option<String>,
) -> ThreadStoreResult<()> {
let Some(state_db) = store.state_db().await else {
return Err(ThreadStoreError::Internal {
message: format!("sqlite state db unavailable for thread {thread_id}"),
});
};
let state_db = store.state_db();
let updated = state_db
.update_thread_git_info(
thread_id,
@@ -307,11 +298,11 @@ async fn resolve_rollout_path(
return Ok(ResolvedRolloutPath { path, archived });
}
let state_db_ctx = store.state_db().await;
let state_db = store.state_db();
let active_path = find_thread_path_by_id_str(
store.config.codex_home.as_path(),
&thread_id.to_string(),
state_db_ctx.as_deref(),
Some(state_db.as_ref()),
)
.await
.map_err(|err| ThreadStoreError::InvalidRequest {
@@ -331,7 +322,7 @@ async fn resolve_rollout_path(
find_archived_thread_path_by_id_str(
store.config.codex_home.as_path(),
&thread_id.to_string(),
state_db_ctx.as_deref(),
Some(state_db.as_ref()),
)
.await
.map_err(|err| ThreadStoreError::InvalidRequest {
@@ -366,14 +357,16 @@ mod tests {
use crate::ThreadPersistenceMetadata;
use crate::ThreadStore;
use crate::local::LocalThreadStore;
use crate::local::test_support::init_test_state_db;
use crate::local::test_support::test_config;
use crate::local::test_support::test_store;
use crate::local::test_support::write_archived_session_file;
use crate::local::test_support::write_session_file;
#[tokio::test]
async fn update_thread_metadata_sets_name_on_active_rollout_and_indexes_name() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(301);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let path =
@@ -408,18 +401,12 @@ mod tests {
async fn update_thread_metadata_sets_memory_mode_on_active_rollout() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(302);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let path =
write_session_file(home.path(), "2025-01-03T14-30-00", uuid).expect("session file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let thread = store
.update_thread_metadata(UpdateThreadMetadataParams {
thread_id,
@@ -452,13 +439,8 @@ mod tests {
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let path =
write_session_file(home.path(), "2025-01-03T18-30-00", uuid).expect("session file");
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
store
.update_thread_metadata(UpdateThreadMetadataParams {
@@ -517,7 +499,7 @@ mod tests {
async fn update_thread_metadata_uses_live_rollout_path_for_external_resume() {
let home = TempDir::new().expect("temp dir");
let external_home = TempDir::new().expect("external temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(307);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let path = write_session_file(external_home.path(), "2025-01-03T14-45-00", uuid)
@@ -558,13 +540,8 @@ mod tests {
async fn update_thread_metadata_sets_git_info() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config, Some(runtime));
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config, runtime);
let uuid = Uuid::from_u128(309);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
write_session_file(home.path(), "2025-01-03T17-00-00", uuid).expect("session file");
@@ -601,13 +578,8 @@ mod tests {
async fn update_thread_metadata_partially_updates_git_info() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config, Some(runtime));
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config, runtime);
let uuid = Uuid::from_u128(310);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
write_session_file(home.path(), "2025-01-03T17-30-00", uuid).expect("session file");
@@ -659,13 +631,8 @@ mod tests {
async fn update_thread_metadata_clears_git_info_fields() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = codex_state::StateRuntime::init(
config.sqlite_home.clone(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(311);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let path =
@@ -829,7 +796,7 @@ mod tests {
#[tokio::test]
async fn update_thread_metadata_rejects_mismatched_session_meta_id() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let filename_uuid = Uuid::from_u128(303);
let metadata_uuid = Uuid::from_u128(304);
let thread_id = ThreadId::from_string(&filename_uuid.to_string()).expect("valid thread id");
@@ -861,7 +828,7 @@ mod tests {
#[tokio::test]
async fn update_thread_metadata_rejects_multi_field_patch_without_partial_write() {
let home = TempDir::new().expect("temp dir");
let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None);
let store = test_store(home.path()).await;
let uuid = Uuid::from_u128(305);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let path =
@@ -896,21 +863,12 @@ mod tests {
async fn update_thread_metadata_keeps_archived_thread_archived_in_sqlite() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(306);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let archived_path = write_archived_session_file(home.path(), "2025-01-03T16-00-00", uuid)
.expect("archived session file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
runtime
.mark_backfill_complete(/*last_watermark*/ None)
.await
.expect("backfill should be complete");
codex_rollout::state_db::reconcile_rollout(
Some(runtime.as_ref()),
archived_path.as_path(),
@@ -959,21 +917,12 @@ mod tests {
async fn update_thread_metadata_keeps_live_archived_thread_archived_in_sqlite() {
let home = TempDir::new().expect("temp dir");
let config = test_config(home.path());
let runtime = init_test_state_db(&config).await;
let store = LocalThreadStore::new(config.clone(), runtime.clone());
let uuid = Uuid::from_u128(308);
let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id");
let archived_path = write_archived_session_file(home.path(), "2025-01-03T16-30-00", uuid)
.expect("archived session file");
let runtime = codex_state::StateRuntime::init(
home.path().to_path_buf(),
config.default_model_provider_id.clone(),
)
.await
.expect("state db should initialize");
let store = LocalThreadStore::new(config.clone(), Some(runtime.clone()));
runtime
.mark_backfill_complete(/*last_watermark*/ None)
.await
.expect("backfill should be complete");
codex_rollout::state_db::reconcile_rollout(
Some(runtime.as_ref()),
archived_path.as_path(),
+32 -30
View File
@@ -881,38 +881,40 @@ pub async fn run_main(
AppServerTarget::Remote { .. } => state_db::get_state_db(&config).await,
};
let effective_toml = config.config_layer_stack.effective_config();
match effective_toml.try_into() {
Ok(config_toml) => {
match crate::legacy_core::personality_migration::maybe_migrate_personality(
&config.codex_home,
&config_toml,
state_db.clone(),
)
.await
{
Ok(
crate::legacy_core::personality_migration::PersonalityMigrationStatus::Applied,
) => {
config = load_config_or_exit(
cli_kv_overrides.clone(),
overrides.clone(),
cloud_requirements.clone(),
)
.await;
}
Ok(
crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedMarker
| crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedExplicitPersonality
| crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedNoSessions,
) => {}
Err(err) => {
tracing::warn!(error = %err, "failed to run personality migration");
if let Some(state_db) = state_db.clone() {
let effective_toml = config.config_layer_stack.effective_config();
match effective_toml.try_into() {
Ok(config_toml) => {
match crate::legacy_core::personality_migration::maybe_migrate_personality(
&config.codex_home,
&config_toml,
state_db,
)
.await
{
Ok(
crate::legacy_core::personality_migration::PersonalityMigrationStatus::Applied,
) => {
config = load_config_or_exit(
cli_kv_overrides.clone(),
overrides.clone(),
cloud_requirements.clone(),
)
.await;
}
Ok(
crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedMarker
| crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedExplicitPersonality
| crate::legacy_core::personality_migration::PersonalityMigrationStatus::SkippedNoSessions,
) => {}
Err(err) => {
tracing::warn!(error = %err, "failed to run personality migration");
}
}
}
}
Err(err) => {
tracing::warn!(error = %err, "failed to deserialize config for personality migration");
Err(err) => {
tracing::warn!(error = %err, "failed to deserialize config for personality migration");
}
}
}
+3
View File
@@ -5733,6 +5733,7 @@ session_picker_view = "dense"
name: None,
turns: vec![codex_app_server_protocol::Turn {
id: String::from("turn-1"),
items_view: codex_app_server_protocol::TurnItemsView::Full,
items: vec![
ThreadItem::UserMessage {
id: String::from("user-1"),
@@ -5798,6 +5799,7 @@ session_picker_view = "dense"
name: None,
turns: vec![codex_app_server_protocol::Turn {
id: String::from("turn-1"),
items_view: codex_app_server_protocol::TurnItemsView::Full,
items: vec![ThreadItem::Reasoning {
id: String::from("reasoning-1"),
summary: Vec::new(),
@@ -5853,6 +5855,7 @@ session_picker_view = "dense"
name: None,
turns: vec![codex_app_server_protocol::Turn {
id: String::from("turn-1"),
items_view: codex_app_server_protocol::TurnItemsView::Full,
items: vec![ThreadItem::Reasoning {
id: String::from("reasoning-1"),
summary: vec![String::from("public summary")],