agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

state: pass state db handles through consumers (#20561)

Browse Source 
## Why

SQLite state was still being opened from consumer paths, including lazy
`OnceCell`-backed thread-store call sites. That let one process
construct multiple state DB connections for the same Codex home, which
makes SQLite lock contention and `database is locked` failures much
easier to hit.

State DB lifetime should be chosen by main-like entrypoints and tests,
then passed through explicitly. Consumers should use the supplied
`Option<StateDbHandle>` or `StateDbHandle` and keep their existing
filesystem fallback or error behavior when no handle is available.

The startup path also needs to keep the rollout crate in charge of
SQLite state initialization. Opening `codex_state::StateRuntime`
directly bypasses rollout metadata backfill, so entrypoints should
initialize through `codex_rollout::state_db` and receive a handle only
after required rollout backfills have completed.

## What Changed

- Initialize the state DB in main-like entrypoints for CLI, TUI,
app-server, exec, MCP server, and the thread-manager sample.
- Pass `Option<StateDbHandle>` through `ThreadManager`,
`LocalThreadStore`, app-server processors, TUI app wiring, rollout
listing/recording, personality migration, shell snapshot cleanup,
session-name lookup, and memory/device-key consumers.
- Remove the lazy local state DB wrapper from the thread store so
non-test consumers use only the supplied handle or their existing
fallback path.
- Make `codex_rollout::state_db::init` the local state startup path: it
opens/migrates SQLite, runs rollout metadata backfill when needed, waits
for concurrent backfill workers up to a bounded timeout, verifies
completion, and then returns the initialized handle.
- Keep optional/non-owning SQLite helpers, such as remote TUI local
reads, as open-only paths that do not run startup backfill.
- Switch app-server startup from direct
`codex_state::StateRuntime::init` to the rollout state initializer so
app-server cannot skip rollout backfill.
- Collapse split rollout lookup/list APIs so callers use the normal
methods with an optional state handle instead of `_with_state_db`
variants.
- Restore `getConversationSummary(ThreadId)` to delegate through
`ThreadStore::read_thread` instead of a LocalThreadStore-specific
rollout path special case.
- Keep DB-backed rollout path lookup keyed on the DB row and file
existence, without imposing the filesystem filename convention on
existing DB rows.
- Verify readable DB-backed rollout paths against `session_meta.id`
before returning them, so a stale SQLite row that points at another
thread's JSONL falls back to filesystem search and read-repairs the DB
row.
- Keep `debug prompt-input` filesystem-only so a one-off debug command
does not initialize or backfill SQLite state just to print prompt input.
- Keep goal-session test Codex homes alive only in the goal-specific
helper, rather than leaking tempdirs from the shared session test
helper.
- Update tests and call sites to pass explicit state handles where DB
behavior is expected and explicit `None` where filesystem-only behavior
is intended.

## Validation

- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo check -p
codex-rollout -p codex-thread-store -p codex-app-server -p codex-core -p
codex-tui -p codex-exec -p codex-cli --tests`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-rollout state_db_`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-rollout find_thread_path`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-rollout find_thread_path -- --nocapture`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-rollout try_init_ -- --nocapture`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-rollout`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo clippy -p
codex-rollout --lib -- -D warnings`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-thread-store
read_thread_falls_back_when_sqlite_path_points_to_another_thread --
--nocapture`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-thread-store`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p codex-core
shell_snapshot`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p codex-core
--test all personality_migration`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p codex-core
--test all rollout_list_find`
- `RUST_MIN_STACK=8388608 CODEX_SKIP_VENDORED_BWRAP=1
CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p codex-core
--test all rollout_list_find::find_prefers_sqlite_path_by_id --
--nocapture`
- `RUST_MIN_STACK=8388608 CODEX_SKIP_VENDORED_BWRAP=1
CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p codex-core
--test all rollout_list_find -- --nocapture`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p codex-core
interrupt_accounts_active_goal_before_pausing`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-app-server get_auth_status -- --test-threads=1`
- `CODEX_SKIP_VENDORED_BWRAP=1
CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo test -p
codex-app-server --lib`
- `CODEX_SKIP_VENDORED_BWRAP=1
CARGO_TARGET_DIR=/tmp/codex-target-state-db cargo check -p codex-rollout
-p codex-app-server --tests`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db just fix -p codex-rollout
-p codex-thread-store -p codex-core -p codex-app-server -p codex-tui -p
codex-exec -p codex-cli`
- `CODEX_SKIP_VENDORED_BWRAP=1
CARGO_TARGET_DIR=/tmp/codex-target-state-db just fix -p codex-rollout -p
codex-app-server`
- `CARGO_TARGET_DIR=/tmp/codex-target-state-db just fix -p
codex-rollout`
- `CODEX_SKIP_VENDORED_BWRAP=1
CARGO_TARGET_DIR=/tmp/codex-target-state-db just fix -p codex-core`
- `just argument-comment-lint -p codex-core`
- `just argument-comment-lint -p codex-rollout`

Focused coverage added in `codex-rollout`:

- `recorder::tests::state_db_init_backfills_before_returning` verifies
the rollout metadata row exists before startup init returns.
- `state_db::tests::try_init_waits_for_concurrent_startup_backfill`
verifies startup waits for another worker to finish backfill instead of
disabling the handle for the process.
-
`state_db::tests::try_init_times_out_waiting_for_stuck_startup_backfill`
verifies startup does not hang indefinitely on a stuck backfill lease.
-
`tests::find_thread_path_accepts_existing_state_db_path_without_canonical_filename`
verifies DB-backed lookup accepts valid existing rollout paths even when
the filename does not include the thread UUID.
-
`tests::find_thread_path_falls_back_when_db_path_points_to_another_thread`
verifies DB-backed lookup ignores a stale row whose existing path
belongs to another thread and read-repairs the row after filesystem
fallback.

Focused coverage updated in `codex-core`:

- `rollout_list_find::find_prefers_sqlite_path_by_id` now uses a
DB-preferred rollout file with matching `session_meta.id`, so it still
verifies that valid SQLite paths win without depending on stale/empty
rollout contents.

`cargo test -p codex-app-server thread_list_respects_search_term_filter
-- --test-threads=1 --nocapture` was attempted locally but timed out
waiting for the app-server test harness `initialize` response before
reaching the changed thread-list code path.

`bazel test //codex-rs/thread-store:thread-store-unit-tests
--test_output=errors` was attempted locally after the thread-store fix,
but this container failed before target analysis while fetching `v8+`
through BuildBuddy/direct GitHub. The equivalent local crate coverage,
including `cargo test -p codex-thread-store`, passes.

A plain local `cargo check -p codex-rollout -p codex-app-server --tests`
also requires system `libcap.pc` for `codex-linux-sandbox`; the
follow-up app-server check above used `CODEX_SKIP_VENDORED_BWRAP=1` in
this container.
This commit is contained in:
Ruslan Nigmatullin
2026-05-04 11:46:03 -07:00
committed by GitHub
Unverified
parent 0035d7bd18
commit 4d201e340e
74 changed files with 1286 additions and 517 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/agent/control.rs
+2 -2
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_rollout::state_db;
use codex_state::DirectionalThreadSpawnEdgeStatus;
use codex_thread_store::ReadThreadParams;
use serde::Serialize;
@@ -526,6 +525,7 @@ 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 {
@@ -536,7 +536,7 @@ impl AgentControl {
agent_nickname: _,
}) => {
let (resumed_agent_nickname, resumed_agent_role) =
if let Some(state_db_ctx) = state_db::get_state_db(&config).await {
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),
codex-rs/core/src/agent/control_tests.rs
+15 -3
View File
@@ -1,5 +1,6 @@
use super::*;
use crate::CodexThread;
use crate::StateDbHandle;
use crate::ThreadManager;
use crate::agent::agent_status_from_event;
use crate::config::AgentRoleConfig;
@@ -7,6 +8,7 @@ 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;
@@ -84,6 +86,7 @@ fn spawn_agent_call(call_id: &str) -> ResponseItem {
struct AgentControlHarness {
_home: TempDir,
config: Config,
state_db: Option<StateDbHandle>,
manager: ThreadManager,
control: AgentControl,
}
@@ -91,16 +94,19 @@ struct AgentControlHarness {
impl AgentControlHarness {
async fn new() -> Self {
let (home, config) = test_config().await;
let manager = ThreadManager::with_models_provider_and_home_for_tests(
let state_db = init_state_db(&config).await;
let manager = ThreadManager::with_models_provider_home_and_state_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(),
);
let control = manager.agent_control();
Self {
_home: home,
config,
state_db,
manager,
control,
}
@@ -1537,16 +1543,19 @@ async fn resume_thread_subagent_restores_stored_nickname_and_role() {
.features
.enable(Feature::Sqlite)
.expect("test config should allow sqlite");
let manager = ThreadManager::with_models_provider_and_home_for_tests(
let state_db = init_state_db(&config).await;
let manager = ThreadManager::with_models_provider_home_and_state_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(),
);
let control = manager.agent_control();
let harness = AgentControlHarness {
_home: home,
config,
state_db,
manager,
control,
};
@@ -1695,7 +1704,10 @@ async fn resume_agent_from_rollout_reads_archived_rollout_path() {
.shutdown_live_agent(child_thread_id)
.await
.expect("child shutdown should succeed");
let store = LocalThreadStore::new(LocalThreadStoreConfig::from_config(&harness.config));
let store = LocalThreadStore::new(
LocalThreadStoreConfig::from_config(&harness.config),
harness.state_db.clone(),
);
store
.archive_thread(ArchiveThreadParams {
thread_id: child_thread_id,
codex-rs/core/src/lib.rs
+1 -1
View File
@@ -145,7 +145,7 @@ pub(crate) mod shell_snapshot;
pub mod spawn;
pub(crate) mod state_db_bridge;
pub use state_db_bridge::StateDbHandle;
pub use state_db_bridge::get_state_db;
pub use state_db_bridge::init_state_db;
mod thread_rollout_truncation;
mod tools;
pub(crate) mod turn_diff_tracker;
codex-rs/core/src/personality_migration.rs
+16 -7
View File
@@ -1,6 +1,7 @@
use crate::config::edit::ConfigEditsBuilder;
use codex_config::config_toml::ConfigToml;
use codex_protocol::config_types::Personality;
use codex_rollout::state_db::StateDbHandle;
use codex_thread_store::ListThreadsParams;
use codex_thread_store::LocalThreadStore;
use codex_thread_store::LocalThreadStoreConfig;
@@ -24,6 +25,7 @@ pub enum PersonalityMigrationStatus {
pub async fn maybe_migrate_personality(
codex_home: &Path,
config_toml: &ConfigToml,
state_db: Option<StateDbHandle>,
) -> io::Result<PersonalityMigrationStatus> {
let marker_path = codex_home.join(PERSONALITY_MIGRATION_FILENAME);
if tokio::fs::try_exists(&marker_path).await? {
@@ -43,7 +45,7 @@ pub async fn maybe_migrate_personality(
.or_else(|| config_toml.model_provider.clone())
.unwrap_or_else(|| "openai".to_string());
if !has_recorded_sessions(codex_home, model_provider_id.as_str()).await? {
if !has_recorded_sessions(codex_home, model_provider_id.as_str(), state_db).await? {
create_marker(&marker_path).await?;
return Ok(PersonalityMigrationStatus::SkippedNoSessions);
}
@@ -60,12 +62,19 @@ pub async fn maybe_migrate_personality(
Ok(PersonalityMigrationStatus::Applied)
}
async fn has_recorded_sessions(codex_home: &Path, default_provider: &str) -> 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(),
});
async fn has_recorded_sessions(
codex_home: &Path,
default_provider: &str,
state_db: Option<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,
);
if has_threads(&store, /*archived*/ false).await? {
return Ok(true);
}
codex-rs/core/src/personality_migration_tests.rs
+5 -5
View File
@@ -87,7 +87,7 @@ 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).await?;
let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?;
assert_eq!(status, PersonalityMigrationStatus::Applied);
assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists());
@@ -103,7 +103,7 @@ 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).await?;
let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?;
assert_eq!(status, PersonalityMigrationStatus::Applied);
assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists());
@@ -119,7 +119,7 @@ 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).await?;
let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?;
assert_eq!(status, PersonalityMigrationStatus::SkippedMarker);
assert!(!temp.path().join("config.toml").exists());
@@ -136,7 +136,7 @@ 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).await?;
let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?;
assert_eq!(
status,
@@ -153,7 +153,7 @@ 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).await?;
let status = maybe_migrate_personality(temp.path(), &config_toml, /*state_db*/ None).await?;
assert_eq!(status, PersonalityMigrationStatus::SkippedNoSessions);
assert!(temp.path().join(PERSONALITY_MIGRATION_FILENAME).exists());
codex-rs/core/src/prompt_debug.rs
+5 -1
View File
@@ -16,6 +16,7 @@ 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::thread_store_from_config;
@@ -24,6 +25,7 @@ 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;
@@ -35,13 +37,15 @@ pub async fn build_prompt_input(
config.codex_linux_sandbox_exe.clone(),
)?;
let thread_store = thread_store_from_config(&config, 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,
thread_store_from_config(&config),
thread_store,
state_db.clone(),
);
let thread = thread_manager.start_thread(config).await?;
codex-rs/core/src/session/mod.rs
+10 -1
View File
@@ -556,7 +556,15 @@ impl Codex {
};
match thread_id {
Some(thread_id) => {
let state_db_ctx = state_db::get_state_db(&config).await;
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
};
state_db::get_dynamic_tools(state_db_ctx.as_deref(), thread_id, "codex_spawn")
.await
}
@@ -1298,6 +1306,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(),
);
}
codex-rs/core/src/session/session.rs
+1
View File
@@ -726,6 +726,7 @@ impl Session {
session_configuration.cwd.clone(),
&mut default_shell,
session_telemetry.clone(),
state_db_ctx.clone(),
)
}
} else {
codex-rs/core/src/session/tests.rs
+59 -15
View File
@@ -3492,6 +3492,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_zsh_path() {
/*analytics_events_client*/ None,
Arc::new(codex_thread_store::LocalThreadStore::new(
codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()),
/*state_db*/ None,
)),
codex_rollout_trace::ThreadTraceContext::disabled(),
)
@@ -3639,6 +3640,7 @@ 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,
)),
model_client: ModelClient::new(
Some(auth_manager.clone()),
@@ -3810,6 +3812,7 @@ async fn make_session_with_config_and_rx(
/*analytics_events_client*/ None,
Arc::new(codex_thread_store::LocalThreadStore::new(
codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()),
/*state_db*/ None,
)),
codex_rollout_trace::ThreadTraceContext::disabled(),
)
@@ -4998,10 +5001,44 @@ async fn make_session_and_context_with_auth_and_config_and_rx<F>(
where
F: FnOnce(&mut Config),
{
let (tx_event, rx_event) = async_channel::unbounded();
let codex_home = tempfile::tempdir().expect("create temp dir");
let mut config = build_test_config(codex_home.path()).await;
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;
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);
@@ -5127,10 +5164,11 @@ where
agent_control,
network_proxy: None,
network_approval: Arc::clone(&network_approval),
state_db: None,
state_db: state_db.clone(),
live_thread: None,
thread_store: Arc::new(codex_thread_store::LocalThreadStore::new(
codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()),
state_db,
)),
model_client: ModelClient::new(
Some(Arc::clone(&auth_manager)),
@@ -5225,10 +5263,13 @@ async fn make_goal_session_and_context_with_rx() -> (
Arc<Session>,
Arc<TurnContext>,
async_channel::Receiver<Event>,
tempfile::TempDir,
) {
let (session, turn_context, rx) = make_session_and_context_with_auth_and_config_and_rx(
let codex_home = tempfile::tempdir().expect("create temp dir");
let (session, turn_context, rx) = make_session_and_context_with_auth_config_home_and_rx(
CodexAuth::from_api_key("Test API Key"),
Vec::new(),
codex_home.path(),
|config| {
config
.features
@@ -5238,14 +5279,14 @@ async fn make_goal_session_and_context_with_rx() -> (
)
.await;
upsert_goal_test_thread(session.as_ref()).await;
(session, turn_context, rx)
(session, turn_context, rx, codex_home)
}
async fn upsert_goal_test_thread(session: &Session) {
let config = session.get_config().await;
let state_db = goal_test_state_db(session)
.await
.expect("goal test state db should initialize");
let state_db = session
.state_db()
.expect("goal test session should have a state db");
let mut builder = codex_state::ThreadMetadataBuilder::new(
session.conversation_id,
config
@@ -7062,7 +7103,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) = make_goal_session_and_context_with_rx().await;
let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
sess.set_thread_goal(
tc.as_ref(),
SetGoalRequest {
@@ -7316,6 +7357,9 @@ fn post_goal_token_usage() -> TokenUsage {
}
async fn goal_test_state_db(sess: &Session) -> anyhow::Result<crate::StateDbHandle> {
if let Some(state_db) = sess.state_db() {
return Ok(state_db);
}
let config = sess.get_config().await;
codex_state::StateRuntime::init(config.sqlite_home.clone(), config.model_provider_id.clone())
.await
@@ -7323,7 +7367,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) = make_goal_session_and_context_with_rx().await;
let (sess, tc, rx, _codex_home) = make_goal_session_and_context_with_rx().await;
sess.set_thread_goal(
tc.as_ref(),
SetGoalRequest {
@@ -7423,7 +7467,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) = make_goal_session_and_context_with_rx().await;
let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
sess.set_thread_goal(
tc.as_ref(),
SetGoalRequest {
@@ -7485,7 +7529,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) = make_goal_session_and_context_with_rx().await;
let (sess, tc, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
sess.spawn_task(
Arc::clone(&tc),
Vec::new(),
@@ -8149,7 +8193,7 @@ async fn sample_rollout(
#[tokio::test]
async fn create_goal_tool_rejects_existing_goal() {
let (session, turn_context, _rx) = make_goal_session_and_context_with_rx().await;
let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
let handler = GoalHandler;
@@ -8211,7 +8255,7 @@ async fn create_goal_tool_rejects_existing_goal() {
#[tokio::test]
async fn update_goal_tool_rejects_pausing_goal() {
let (session, turn_context, _rx) = make_goal_session_and_context_with_rx().await;
let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
let handler = GoalHandler;
@@ -8271,7 +8315,7 @@ async fn update_goal_tool_rejects_pausing_goal() {
#[tokio::test]
async fn update_goal_tool_marks_goal_complete() {
let (session, turn_context, _rx) = make_goal_session_and_context_with_rx().await;
let (session, turn_context, _rx, _codex_home) = make_goal_session_and_context_with_rx().await;
let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
let handler = GoalHandler;
codex-rs/core/src/session/tests/guardian_tests.rs
+1
View File
@@ -731,6 +731,7 @@ 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,
));
let CodexSpawnOk { codex, .. } = Codex::spawn(CodexSpawnArgs {
codex-rs/core/src/shell_snapshot.rs
+22 -6
View File
@@ -5,6 +5,7 @@ use std::sync::Arc;
use std::time::Duration;
use std::time::SystemTime;
use crate::StateDbHandle;
use crate::rollout::list::find_thread_path_by_id_str;
use crate::shell::Shell;
use crate::shell::ShellType;
@@ -41,6 +42,7 @@ impl ShellSnapshot {
session_cwd: AbsolutePathBuf,
shell: &mut Shell,
session_telemetry: SessionTelemetry,
state_db: Option<StateDbHandle>,
) -> watch::Sender<Option<Arc<ShellSnapshot>>> {
let (shell_snapshot_tx, shell_snapshot_rx) = watch::channel(None);
shell.shell_snapshot = shell_snapshot_rx;
@@ -52,6 +54,7 @@ impl ShellSnapshot {
shell.clone(),
shell_snapshot_tx.clone(),
session_telemetry,
state_db,
);
shell_snapshot_tx
@@ -64,6 +67,7 @@ impl ShellSnapshot {
shell: Shell,
shell_snapshot_tx: watch::Sender<Option<Arc<ShellSnapshot>>>,
session_telemetry: SessionTelemetry,
state_db: Option<StateDbHandle>,
) {
Self::spawn_snapshot_task(
codex_home,
@@ -72,6 +76,7 @@ impl ShellSnapshot {
shell,
shell_snapshot_tx,
session_telemetry,
state_db,
);
}
@@ -82,15 +87,21 @@ impl ShellSnapshot {
snapshot_shell: Shell,
shell_snapshot_tx: watch::Sender<Option<Arc<ShellSnapshot>>>,
session_telemetry: SessionTelemetry,
state_db: Option<StateDbHandle>,
) {
let snapshot_span = info_span!("shell_snapshot", thread_id = %session_id);
tokio::spawn(
async move {
let timer = session_telemetry.start_timer("codex.shell_snapshot.duration_ms", &[]);
let snapshot =
ShellSnapshot::try_new(&codex_home, session_id, &session_cwd, &snapshot_shell)
.await
.map(Arc::new);
let snapshot = ShellSnapshot::try_new(
&codex_home,
session_id,
&session_cwd,
&snapshot_shell,
state_db,
)
.await
.map(Arc::new);
let success = snapshot.is_ok();
let success_tag = if success { "true" } else { "false" };
let _ = timer.map(|timer| timer.record(&[("success", success_tag)]));
@@ -110,6 +121,7 @@ impl ShellSnapshot {
session_id: ThreadId,
session_cwd: &AbsolutePathBuf,
shell: &Shell,
state_db: Option<StateDbHandle>,
) -> std::result::Result<Self, &'static str> {
// File to store the snapshot
let extension = match shell.shell_type {
@@ -131,7 +143,9 @@ impl ShellSnapshot {
let codex_home = codex_home.clone();
let cleanup_session_id = session_id;
tokio::spawn(async move {
if let Err(err) = cleanup_stale_snapshots(&codex_home, cleanup_session_id).await {
if let Err(err) =
cleanup_stale_snapshots(&codex_home, cleanup_session_id, state_db).await
{
tracing::warn!("Failed to clean up shell snapshots: {err:?}");
}
});
@@ -486,6 +500,7 @@ $envVars | ForEach-Object {
pub async fn cleanup_stale_snapshots(
codex_home: &AbsolutePathBuf,
active_session_id: ThreadId,
state_db: Option<StateDbHandle>,
) -> Result<()> {
let snapshot_dir = codex_home.join(SNAPSHOT_DIR);
@@ -515,7 +530,8 @@ pub async fn cleanup_stale_snapshots(
continue;
}
let rollout_path = find_thread_path_by_id_str(codex_home, session_id).await?;
let rollout_path =
find_thread_path_by_id_str(codex_home, session_id, state_db.as_deref()).await?;
let Some(rollout_path) = rollout_path else {
remove_snapshot_file(&path).await;
continue;
codex-rs/core/src/shell_snapshot_tests.rs
+22 -11
View File
@@ -202,6 +202,7 @@ async fn try_new_creates_and_deletes_snapshot_file() -> Result<()> {
ThreadId::new(),
&dir.path().abs(),
&shell,
/*state_db*/ None,
)
.await
.expect("snapshot should be created");
@@ -227,14 +228,24 @@ async fn try_new_uses_distinct_generation_paths() -> Result<()> {
shell_snapshot: crate::shell::empty_shell_snapshot_receiver(),
};
let initial_snapshot =
ShellSnapshot::try_new(&dir.path().abs(), session_id, &dir.path().abs(), &shell)
.await
.expect("initial snapshot should be created");
let refreshed_snapshot =
ShellSnapshot::try_new(&dir.path().abs(), session_id, &dir.path().abs(), &shell)
.await
.expect("refreshed snapshot should be created");
let initial_snapshot = ShellSnapshot::try_new(
&dir.path().abs(),
session_id,
&dir.path().abs(),
&shell,
/*state_db*/ None,
)
.await
.expect("initial snapshot should be created");
let refreshed_snapshot = ShellSnapshot::try_new(
&dir.path().abs(),
session_id,
&dir.path().abs(),
&shell,
/*state_db*/ None,
)
.await
.expect("refreshed snapshot should be created");
let initial_path = initial_snapshot.path.clone();
let refreshed_path = refreshed_snapshot.path.clone();
@@ -428,7 +439,7 @@ async fn cleanup_stale_snapshots_removes_orphans_and_keeps_live() -> Result<()>
fs::write(&orphan_snapshot, "orphan").await?;
fs::write(&invalid_snapshot, "invalid").await?;
cleanup_stale_snapshots(&codex_home, ThreadId::new()).await?;
cleanup_stale_snapshots(&codex_home, ThreadId::new(), /*state_db*/ None).await?;
assert_eq!(live_snapshot.exists(), true);
assert_eq!(orphan_snapshot.exists(), false);
@@ -451,7 +462,7 @@ async fn cleanup_stale_snapshots_removes_stale_rollouts() -> Result<()> {
set_file_mtime(&rollout_path, SNAPSHOT_RETENTION + Duration::from_secs(60))?;
cleanup_stale_snapshots(&codex_home, ThreadId::new()).await?;
cleanup_stale_snapshots(&codex_home, ThreadId::new(), /*state_db*/ None).await?;
assert_eq!(stale_snapshot.exists(), false);
Ok(())
@@ -472,7 +483,7 @@ async fn cleanup_stale_snapshots_skips_active_session() -> Result<()> {
set_file_mtime(&rollout_path, SNAPSHOT_RETENTION + Duration::from_secs(60))?;
cleanup_stale_snapshots(&codex_home, active_session).await?;
cleanup_stale_snapshots(&codex_home, active_session, /*state_db*/ None).await?;
assert_eq!(active_snapshot.exists(), true);
Ok(())
codex-rs/core/src/state_db_bridge.rs
+2 -2
View File
@@ -3,6 +3,6 @@ pub use codex_rollout::state_db::StateDbHandle;
use crate::config::Config;
pub async fn get_state_db(config: &Config) -> Option<StateDbHandle> {
rollout_state_db::get_state_db(config).await
pub async fn init_state_db(config: &Config) -> Option<StateDbHandle> {
rollout_state_db::init(config).await
}
codex-rs/core/src/stream_events_utils.rs
+7 -4
View File
@@ -138,8 +138,11 @@ 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(turn_context, item).await;
let has_memory_citation = record_stage1_output_usage_and_detect_memory_citation(
sess.services.state_db.as_ref(),
item,
)
.await;
if has_memory_citation {
sess.record_memory_citation_for_turn(&turn_context.sub_id)
.await;
@@ -174,7 +177,7 @@ pub(crate) async fn mark_thread_memory_mode_polluted_if_external_context(
}
async fn record_stage1_output_usage_and_detect_memory_citation(
turn_context: &TurnContext,
state_db_ctx: Option<&state_db::StateDbHandle>,
item: &ResponseItem,
) -> bool {
let Some(raw_text) = raw_assistant_output_text_from_item(item) else {
@@ -190,7 +193,7 @@ async fn record_stage1_output_usage_and_detect_memory_citation(
return true;
}
if let Some(db) = state_db::get_state_db(turn_context.config.as_ref()).await {
if let Some(db) = state_db_ctx {
let _ = db.record_stage1_output_usage(&thread_ids).await;
}
true
codex-rs/core/src/test_support.rs
+16
View File
@@ -73,6 +73,22 @@ pub fn thread_manager_with_models_provider_and_home(
)
}
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,
)
}
pub async fn start_thread_with_user_shell_override(
thread_manager: &ThreadManager,
config: Config,
codex-rs/core/src/thread_manager.rs
+35 -4
View File
@@ -50,6 +50,7 @@ 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::StateDbHandle;
use codex_state::DirectionalThreadSpawnEdgeStatus;
use codex_thread_store::InMemoryThreadStore;
use codex_thread_store::LocalThreadStore;
@@ -248,6 +249,7 @@ pub(crate) struct ThreadManagerState {
thread_store: Arc<dyn ThreadStore>,
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>,
}
@@ -263,10 +265,14 @@ pub fn build_models_manager(
)
}
pub fn thread_store_from_config(config: &Config) -> Arc<dyn ThreadStore> {
pub fn thread_store_from_config(
config: &Config,
state_db: Option<StateDbHandle>,
) -> Arc<dyn ThreadStore> {
match &config.experimental_thread_store {
ThreadStoreConfig::Local => Arc::new(LocalThreadStore::new(
LocalThreadStoreConfig::from_config(config),
state_db,
)),
ThreadStoreConfig::Remote { endpoint } => Arc::new(RemoteThreadStore::new(endpoint)),
ThreadStoreConfig::InMemory { id } => InMemoryThreadStore::for_id(id),
@@ -281,6 +287,7 @@ impl ThreadManager {
environment_manager: Arc<EnvironmentManager>,
analytics_events_client: Option<AnalyticsEventsClient>,
thread_store: Arc<dyn ThreadStore>,
state_db: Option<StateDbHandle>,
) -> Self {
let codex_home = config.codex_home.clone();
let restriction_product = session_source.restriction_product();
@@ -310,6 +317,7 @@ impl ThreadManager {
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()))),
}),
@@ -347,6 +355,22 @@ impl ThreadManager {
provider: ModelProviderInfo,
codex_home: PathBuf,
environment_manager: Arc<EnvironmentManager>,
) -> Self {
Self::with_models_provider_home_and_state_for_tests(
auth,
provider,
codex_home,
environment_manager,
/*state_db*/ None,
)
}
pub(crate) fn with_models_provider_home_and_state_for_tests(
auth: CodexAuth,
provider: ModelProviderInfo,
codex_home: PathBuf,
environment_manager: Arc<EnvironmentManager>,
state_db: Option<StateDbHandle>,
) -> Self {
set_thread_manager_test_mode_for_tests(/*enabled*/ true);
let auth_manager = AuthManager::from_auth_for_testing(auth);
@@ -369,12 +393,14 @@ impl ThreadManager {
let skills_watcher = build_skills_watcher(Arc::clone(&skills_manager));
// This test constructor has no Config input. Tests that need a non-local
// process store should construct ThreadManager::new with an explicit store.
let thread_store: Arc<dyn ThreadStore> =
Arc::new(LocalThreadStore::new(LocalThreadStoreConfig {
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(),
));
Self {
state: Arc::new(ThreadManagerState {
threads: Arc::new(RwLock::new(HashMap::new())),
@@ -390,6 +416,7 @@ impl ThreadManager {
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()))),
}),
@@ -837,6 +864,10 @@ impl ThreadManager {
}
impl ThreadManagerState {
pub(crate) fn state_db(&self) -> Option<StateDbHandle> {
self.state_db.clone()
}
pub(crate) async fn list_thread_ids(&self) -> Vec<ThreadId> {
self.threads
.read()
codex-rs/core/src/thread_manager_tests.rs
+21 -8
View File
@@ -1,5 +1,6 @@
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;
@@ -389,7 +390,8 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() {
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
thread_store_from_config(&config),
thread_store_from_config(&config, /*state_db*/ None),
/*state_db*/ None,
);
let selected_cwd =
AbsolutePathBuf::try_from(config.cwd.as_path().join("selected")).expect("absolute path");
@@ -498,7 +500,8 @@ async fn resume_active_thread_from_rollout_returns_running_thread() {
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
thread_store_from_config(&config),
thread_store_from_config(&config, /*state_db*/ None),
/*state_db*/ None,
);
let source = manager
@@ -551,7 +554,8 @@ async fn resume_stopped_thread_from_rollout_spawns_new_thread() {
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
thread_store_from_config(&config),
thread_store_from_config(&config, /*state_db*/ None),
/*state_db*/ None,
);
let source = manager
@@ -614,7 +618,8 @@ async fn new_uses_active_provider_for_model_refresh() {
SessionSource::Exec,
Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
/*analytics_events_client*/ None,
thread_store_from_config(&config),
thread_store_from_config(&config, /*state_db*/ None),
/*state_db*/ None,
);
let _ = manager.list_models(RefreshStrategy::Online).await;
@@ -819,13 +824,15 @@ 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 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),
thread_store_from_config(&config, state_db.clone()),
state_db.clone(),
);
let source = manager
@@ -921,13 +928,15 @@ 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 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),
thread_store_from_config(&config, state_db.clone()),
state_db.clone(),
);
let source = manager
@@ -1012,13 +1021,15 @@ 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 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),
thread_store_from_config(&config, state_db.clone()),
state_db.clone(),
);
let source = manager
@@ -1148,13 +1159,15 @@ 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 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),
thread_store_from_config(&config, state_db.clone()),
state_db.clone(),
);
let source = manager
codex-rs/core/src/tools/handlers/multi_agents_tests.rs
+12 -1
View File
@@ -3,8 +3,10 @@ use crate::ThreadManager;
use crate::config::AgentRoleConfig;
use crate::config::DEFAULT_AGENT_MAX_DEPTH;
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::thread_store_from_config;
use crate::tools::context::ToolOutput;
use crate::tools::handlers::multi_agents_v2::CloseAgentHandler as CloseAgentHandlerV2;
use crate::tools::handlers::multi_agents_v2::FollowupTaskHandler as FollowupTaskHandlerV2;
@@ -3149,13 +3151,22 @@ async fn close_agent_submits_shutdown_and_returns_previous_status() {
#[tokio::test]
async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtrees_closed() {
let (_session, turn) = make_session_and_context().await;
let manager = thread_manager();
let mut config = turn.config.as_ref().clone();
config.agent_max_depth = 3;
config
.features
.enable(Feature::Sqlite)
.expect("test config should allow sqlite");
let state_db = init_state_db(&config).await;
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(),
);
let parent = manager
.start_thread(config.clone())