mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
state: pass state db handles through consumers (#20561)
## 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:
committed by
GitHub
Unverified
parent
0035d7bd18
commit
4d201e340e
@@ -86,6 +86,7 @@ use codex_exec_server::EnvironmentManager;
|
||||
use codex_feedback::CodexFeedback;
|
||||
use codex_login::AuthManager;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
pub use codex_rollout::StateDbHandle;
|
||||
pub use codex_state::log_db::LogDbLayer;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
@@ -126,6 +127,8 @@ 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.
|
||||
pub state_db: Option<StateDbHandle>,
|
||||
/// Environment manager used by core execution and filesystem operations.
|
||||
pub environment_manager: Arc<EnvironmentManager>,
|
||||
/// Startup warnings emitted after initialize succeeds.
|
||||
@@ -251,6 +254,8 @@ pub struct InProcessClientHandle {
|
||||
client: InProcessClientSender,
|
||||
event_rx: mpsc::Receiver<InProcessServerEvent>,
|
||||
runtime_handle: tokio::task::JoinHandle<()>,
|
||||
#[cfg(test)]
|
||||
_test_codex_home: Option<tempfile::TempDir>,
|
||||
}
|
||||
|
||||
impl InProcessClientHandle {
|
||||
@@ -418,6 +423,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,
|
||||
config_warnings: args.config_warnings,
|
||||
session_source: args.session_source,
|
||||
auth_manager,
|
||||
@@ -717,6 +723,8 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
|
||||
client: InProcessClientSender { client_tx },
|
||||
event_rx,
|
||||
runtime_handle,
|
||||
#[cfg(test)]
|
||||
_test_codex_home: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -738,13 +746,22 @@ mod tests {
|
||||
use codex_app_server_protocol::TurnStatus;
|
||||
use codex_core::config::ConfigBuilder;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::path::Path;
|
||||
use tempfile::TempDir;
|
||||
|
||||
async fn build_test_config() -> Config {
|
||||
match ConfigBuilder::default().build().await {
|
||||
async fn build_test_config(codex_home: &Path) -> Config {
|
||||
match ConfigBuilder::default()
|
||||
.codex_home(codex_home.to_path_buf())
|
||||
.build()
|
||||
.await
|
||||
{
|
||||
Ok(config) => config,
|
||||
Err(_) => Config::load_default_with_cli_overrides(Vec::new())
|
||||
.await
|
||||
.expect("default config should load"),
|
||||
Err(_) => Config::load_default_with_cli_overrides_for_codex_home(
|
||||
codex_home.to_path_buf(),
|
||||
Vec::new(),
|
||||
)
|
||||
.await
|
||||
.expect("default config should load"),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -752,15 +769,21 @@ mod tests {
|
||||
session_source: SessionSource,
|
||||
channel_capacity: usize,
|
||||
) -> 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())
|
||||
.await
|
||||
.expect("state db should initialize for in-process test");
|
||||
let args = InProcessStartArgs {
|
||||
arg0_paths: Arg0DispatchPaths::default(),
|
||||
config: Arc::new(build_test_config().await),
|
||||
config,
|
||||
cli_overrides: Vec::new(),
|
||||
loader_overrides: LoaderOverrides::default(),
|
||||
cloud_requirements: CloudRequirementsLoader::default(),
|
||||
thread_config_loader: Arc::new(codex_config::NoopThreadConfigLoader),
|
||||
feedback: CodexFeedback::new(),
|
||||
log_db: None,
|
||||
state_db: Some(state_db),
|
||||
environment_manager: Arc::new(EnvironmentManager::default_for_tests()),
|
||||
config_warnings: Vec::new(),
|
||||
session_source,
|
||||
@@ -775,7 +798,9 @@ mod tests {
|
||||
},
|
||||
channel_capacity,
|
||||
};
|
||||
start(args).await.expect("in-process runtime should start")
|
||||
let mut client = start(args).await.expect("in-process runtime should start");
|
||||
client._test_codex_home = Some(codex_home);
|
||||
client
|
||||
}
|
||||
|
||||
async fn start_test_client(session_source: SessionSource) -> InProcessClientHandle {
|
||||
|
||||
@@ -54,6 +54,7 @@ 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;
|
||||
@@ -453,23 +454,6 @@ pub async fn run_main_with_transport_options(
|
||||
.await
|
||||
{
|
||||
Ok(config) => {
|
||||
let effective_toml = config.config_layer_stack.effective_config();
|
||||
match effective_toml.try_into() {
|
||||
Ok(config_toml) => {
|
||||
if let Err(err) = codex_core::personality_migration::maybe_migrate_personality(
|
||||
&config.codex_home,
|
||||
&config_toml,
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!(error = %err, "Failed to run personality migration");
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(error = %err, "Failed to deserialize config for personality migration");
|
||||
}
|
||||
}
|
||||
|
||||
let discovered_thread_config_loader = configured_thread_config_loader(&config);
|
||||
config_manager
|
||||
.replace_thread_config_loader(Arc::clone(&discovered_thread_config_loader));
|
||||
@@ -483,23 +467,70 @@ pub async fn run_main_with_transport_options(
|
||||
}
|
||||
};
|
||||
let mut config_warnings = Vec::new();
|
||||
let config = match config_manager
|
||||
let (mut config, should_run_personality_migration) = match config_manager
|
||||
.load_latest_config(/*fallback_cwd*/ None)
|
||||
.await
|
||||
{
|
||||
Ok(config) => config,
|
||||
Ok(config) => (config, true),
|
||||
Err(err) => {
|
||||
let message = config_warning_from_error("Invalid configuration; using defaults.", &err);
|
||||
config_warnings.push(message);
|
||||
config_manager.load_default_config().await.map_err(|e| {
|
||||
std::io::Error::new(
|
||||
ErrorKind::InvalidData,
|
||||
format!("error loading default config after config error: {e}"),
|
||||
)
|
||||
})?
|
||||
(
|
||||
config_manager.load_default_config().await.map_err(|e| {
|
||||
std::io::Error::new(
|
||||
ErrorKind::InvalidData,
|
||||
format!("error loading default config after config error: {e}"),
|
||||
)
|
||||
})?,
|
||||
false,
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
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();
|
||||
|
||||
if should_run_personality_migration {
|
||||
let effective_toml = config.config_layer_stack.effective_config();
|
||||
match effective_toml.try_into() {
|
||||
Ok(config_toml) => {
|
||||
match codex_core::personality_migration::maybe_migrate_personality(
|
||||
&config.codex_home,
|
||||
&config_toml,
|
||||
state_db.clone(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(codex_core::personality_migration::PersonalityMigrationStatus::Applied) => {
|
||||
config = config_manager
|
||||
.load_latest_config(/*fallback_cwd*/ None)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
std::io::Error::new(
|
||||
ErrorKind::InvalidData,
|
||||
format!(
|
||||
"error reloading config after personality migration: {err}"
|
||||
),
|
||||
)
|
||||
})?;
|
||||
}
|
||||
Ok(
|
||||
codex_core::personality_migration::PersonalityMigrationStatus::SkippedMarker
|
||||
| codex_core::personality_migration::PersonalityMigrationStatus::SkippedExplicitPersonality
|
||||
| codex_core::personality_migration::PersonalityMigrationStatus::SkippedNoSessions,
|
||||
) => {}
|
||||
Err(err) => {
|
||||
warn!(error = %err, "Failed to run personality migration");
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(error = %err, "Failed to deserialize config for personality migration");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Ok(Some(err)) = check_execpolicy_for_warnings(&config.config_layer_stack).await {
|
||||
let (path, range) = exec_policy_warning_location(&err);
|
||||
let message = ConfigWarningNotification {
|
||||
@@ -567,13 +598,6 @@ pub async fn run_main_with_transport_options(
|
||||
|
||||
let feedback_layer = feedback.logger_layer();
|
||||
let feedback_metadata_layer = feedback.metadata_layer();
|
||||
let state_db_result = codex_state::StateRuntime::init(
|
||||
config.sqlite_home.clone(),
|
||||
config.model_provider_id.clone(),
|
||||
)
|
||||
.await;
|
||||
let state_db_init_error = state_db_result.as_ref().err().map(ToString::to_string);
|
||||
let state_db = state_db_result.ok();
|
||||
let log_db = state_db.clone().map(log_db::start);
|
||||
let log_db_layer = log_db
|
||||
.clone()
|
||||
@@ -745,6 +769,7 @@ pub async fn run_main_with_transport_options(
|
||||
environment_manager,
|
||||
feedback: feedback.clone(),
|
||||
log_db,
|
||||
state_db: state_db.clone(),
|
||||
config_warnings,
|
||||
session_source,
|
||||
auth_manager,
|
||||
|
||||
@@ -72,6 +72,7 @@ use codex_login::auth::ExternalAuthTokens;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use codex_protocol::protocol::W3cTraceContext;
|
||||
use codex_rollout::StateDbHandle;
|
||||
use codex_state::log_db::LogDbLayer;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::sync::Semaphore;
|
||||
@@ -251,6 +252,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) config_warnings: Vec<ConfigWarningNotification>,
|
||||
pub(crate) session_source: SessionSource,
|
||||
pub(crate) auth_manager: Arc<AuthManager>,
|
||||
@@ -272,6 +274,7 @@ impl MessageProcessor {
|
||||
environment_manager,
|
||||
feedback,
|
||||
log_db,
|
||||
state_db,
|
||||
config_warnings,
|
||||
session_source,
|
||||
auth_manager,
|
||||
@@ -285,7 +288,7 @@ impl MessageProcessor {
|
||||
// The thread store is intentionally process-scoped. Config reloads can
|
||||
// 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());
|
||||
let thread_store = thread_store_from_config(config.as_ref(), state_db.clone());
|
||||
let thread_manager = Arc::new(ThreadManager::new(
|
||||
config.as_ref(),
|
||||
auth_manager.clone(),
|
||||
@@ -293,6 +296,7 @@ impl MessageProcessor {
|
||||
environment_manager,
|
||||
Some(analytics_events_client.clone()),
|
||||
Arc::clone(&thread_store),
|
||||
state_db.clone(),
|
||||
));
|
||||
thread_manager
|
||||
.plugins_manager()
|
||||
@@ -337,6 +341,7 @@ impl MessageProcessor {
|
||||
Arc::clone(&config),
|
||||
feedback,
|
||||
log_db,
|
||||
state_db.clone(),
|
||||
);
|
||||
let git_processor = GitRequestProcessor::new();
|
||||
let initialize_processor = InitializeRequestProcessor::new(
|
||||
@@ -371,6 +376,7 @@ impl MessageProcessor {
|
||||
outgoing.clone(),
|
||||
Arc::clone(&config),
|
||||
thread_state_manager.clone(),
|
||||
state_db.clone(),
|
||||
);
|
||||
let thread_processor = ThreadRequestProcessor::new(
|
||||
auth_manager.clone(),
|
||||
@@ -386,6 +392,7 @@ impl MessageProcessor {
|
||||
thread_watch_manager.clone(),
|
||||
Arc::clone(&thread_list_state_permit),
|
||||
thread_goal_processor.clone(),
|
||||
state_db.clone(),
|
||||
);
|
||||
let turn_processor = TurnRequestProcessor::new(
|
||||
auth_manager.clone(),
|
||||
@@ -399,6 +406,7 @@ impl MessageProcessor {
|
||||
thread_state_manager,
|
||||
thread_watch_manager,
|
||||
thread_list_state_permit,
|
||||
state_db.clone(),
|
||||
);
|
||||
if matches!(plugin_startup_tasks, crate::PluginStartupTasks::Start) {
|
||||
// Keep plugin startup warmups aligned at app-server startup.
|
||||
@@ -429,11 +437,7 @@ impl MessageProcessor {
|
||||
arg0_paths,
|
||||
config.codex_home.to_path_buf(),
|
||||
);
|
||||
let device_key_processor = DeviceKeyRequestProcessor::new(
|
||||
outgoing.clone(),
|
||||
config.sqlite_home.clone(),
|
||||
config.model_provider_id.clone(),
|
||||
);
|
||||
let device_key_processor = DeviceKeyRequestProcessor::new(outgoing.clone(), state_db);
|
||||
let fs_processor = FsRequestProcessor::new(
|
||||
thread_manager
|
||||
.environment_manager()
|
||||
|
||||
@@ -290,6 +290,7 @@ async fn build_test_processor(
|
||||
environment_manager: Arc::new(EnvironmentManager::default_for_tests()),
|
||||
feedback: CodexFeedback::new(),
|
||||
log_db: None,
|
||||
state_db: None,
|
||||
config_warnings: Vec::new(),
|
||||
session_source: SessionSource::VSCode,
|
||||
auth_manager,
|
||||
|
||||
@@ -369,7 +369,6 @@ use codex_rmcp_client::perform_oauth_login_return_url;
|
||||
use codex_rollout::EventPersistenceMode;
|
||||
use codex_rollout::is_persisted_rollout_item;
|
||||
use codex_rollout::state_db::StateDbHandle;
|
||||
use codex_rollout::state_db::get_state_db;
|
||||
use codex_rollout::state_db::reconcile_rollout;
|
||||
use codex_state::StateRuntime;
|
||||
use codex_state::ThreadMetadata;
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
use std::fmt;
|
||||
use std::future::Future;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::error_code::internal_error;
|
||||
@@ -36,7 +35,6 @@ use codex_device_key::RemoteControlClientEnrollmentAudience;
|
||||
use codex_device_key::RemoteControlClientEnrollmentSignPayload;
|
||||
use codex_state::DeviceKeyBindingRecord;
|
||||
use codex_state::StateRuntime;
|
||||
use tokio::sync::OnceCell;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct DeviceKeyRequestProcessor {
|
||||
@@ -47,15 +45,11 @@ pub(crate) struct DeviceKeyRequestProcessor {
|
||||
impl DeviceKeyRequestProcessor {
|
||||
pub(crate) fn new(
|
||||
outgoing: Arc<OutgoingMessageSender>,
|
||||
sqlite_home: PathBuf,
|
||||
default_provider: String,
|
||||
state_db: Option<Arc<StateRuntime>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
outgoing,
|
||||
store: DeviceKeyStore::new(Arc::new(StateDeviceKeyBindingStore::new(
|
||||
sqlite_home,
|
||||
default_provider,
|
||||
))),
|
||||
store: DeviceKeyStore::new(Arc::new(StateDeviceKeyBindingStore::new(state_db))),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -176,39 +170,25 @@ async fn sign_device_key(
|
||||
}
|
||||
|
||||
struct StateDeviceKeyBindingStore {
|
||||
sqlite_home: PathBuf,
|
||||
default_provider: String,
|
||||
state_db: OnceCell<Arc<StateRuntime>>,
|
||||
state_db: Option<Arc<StateRuntime>>,
|
||||
}
|
||||
|
||||
impl StateDeviceKeyBindingStore {
|
||||
fn new(sqlite_home: PathBuf, default_provider: String) -> Self {
|
||||
Self {
|
||||
sqlite_home,
|
||||
default_provider,
|
||||
state_db: OnceCell::new(),
|
||||
}
|
||||
fn new(state_db: Option<Arc<StateRuntime>>) -> Self {
|
||||
Self { state_db }
|
||||
}
|
||||
|
||||
async fn state_db(&self) -> Result<Arc<StateRuntime>, DeviceKeyError> {
|
||||
let sqlite_home = self.sqlite_home.clone();
|
||||
let default_provider = self.default_provider.clone();
|
||||
self.state_db
|
||||
.get_or_try_init(|| async move {
|
||||
StateRuntime::init(sqlite_home, default_provider)
|
||||
.await
|
||||
.map_err(|err| DeviceKeyError::Platform(err.to_string()))
|
||||
})
|
||||
.await
|
||||
.cloned()
|
||||
.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("sqlite_home", &self.sqlite_home)
|
||||
.field("default_provider", &self.default_provider)
|
||||
.field("has_state_db", &self.state_db.is_some())
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,6 +7,7 @@ pub(crate) struct FeedbackRequestProcessor {
|
||||
config: Arc<Config>,
|
||||
feedback: CodexFeedback,
|
||||
log_db: Option<LogDbLayer>,
|
||||
state_db: Option<StateDbHandle>,
|
||||
}
|
||||
|
||||
impl FeedbackRequestProcessor {
|
||||
@@ -16,6 +17,7 @@ impl FeedbackRequestProcessor {
|
||||
config: Arc<Config>,
|
||||
feedback: CodexFeedback,
|
||||
log_db: Option<LogDbLayer>,
|
||||
state_db: Option<StateDbHandle>,
|
||||
) -> Self {
|
||||
Self {
|
||||
auth_manager,
|
||||
@@ -23,6 +25,7 @@ impl FeedbackRequestProcessor {
|
||||
config,
|
||||
feedback,
|
||||
log_db,
|
||||
state_db,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -75,7 +78,7 @@ impl FeedbackRequestProcessor {
|
||||
if let Some(log_db) = self.log_db.as_ref() {
|
||||
log_db.flush().await;
|
||||
}
|
||||
let state_db_ctx = get_state_db(&self.config).await;
|
||||
let state_db_ctx = self.state_db.clone();
|
||||
let feedback_thread_ids = match conversation_id {
|
||||
Some(conversation_id) => match self
|
||||
.thread_manager
|
||||
|
||||
@@ -7,6 +7,7 @@ pub(crate) struct ThreadGoalRequestProcessor {
|
||||
outgoing: Arc<OutgoingMessageSender>,
|
||||
config: Arc<Config>,
|
||||
thread_state_manager: ThreadStateManager,
|
||||
state_db: Option<StateDbHandle>,
|
||||
}
|
||||
|
||||
impl ThreadGoalRequestProcessor {
|
||||
@@ -15,12 +16,14 @@ impl ThreadGoalRequestProcessor {
|
||||
outgoing: Arc<OutgoingMessageSender>,
|
||||
config: Arc<Config>,
|
||||
thread_state_manager: ThreadStateManager,
|
||||
state_db: Option<StateDbHandle>,
|
||||
) -> Self {
|
||||
Self {
|
||||
thread_manager,
|
||||
outgoing,
|
||||
config,
|
||||
thread_state_manager,
|
||||
state_db,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -78,7 +81,7 @@ impl ThreadGoalRequestProcessor {
|
||||
if let Some(state_db) = thread.state_db() {
|
||||
Some(state_db)
|
||||
} else {
|
||||
open_state_db_for_direct_thread_lookup(&self.config).await
|
||||
self.state_db.clone()
|
||||
}
|
||||
} else {
|
||||
None
|
||||
@@ -104,12 +107,16 @@ impl ThreadGoalRequestProcessor {
|
||||
"ephemeral thread does not support goals: {thread_id}"
|
||||
))
|
||||
})?,
|
||||
None => find_thread_path_by_id_str(&self.config.codex_home, &thread_id.to_string())
|
||||
.await
|
||||
.map_err(|err| {
|
||||
internal_error(format!("failed to locate thread id {thread_id}: {err}"))
|
||||
})?
|
||||
.ok_or_else(|| invalid_request(format!("thread not found: {thread_id}")))?,
|
||||
None => find_thread_path_by_id_str(
|
||||
&self.config.codex_home,
|
||||
&thread_id.to_string(),
|
||||
self.state_db.as_deref(),
|
||||
)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
internal_error(format!("failed to locate thread id {thread_id}: {err}"))
|
||||
})?
|
||||
.ok_or_else(|| invalid_request(format!("thread not found: {thread_id}")))?,
|
||||
};
|
||||
reconcile_rollout(
|
||||
Some(&state_db),
|
||||
@@ -247,12 +254,16 @@ impl ThreadGoalRequestProcessor {
|
||||
"ephemeral thread does not support goals: {thread_id}"
|
||||
))
|
||||
})?,
|
||||
None => find_thread_path_by_id_str(&self.config.codex_home, &thread_id.to_string())
|
||||
.await
|
||||
.map_err(|err| {
|
||||
internal_error(format!("failed to locate thread id {thread_id}: {err}"))
|
||||
})?
|
||||
.ok_or_else(|| invalid_request(format!("thread not found: {thread_id}")))?,
|
||||
None => find_thread_path_by_id_str(
|
||||
&self.config.codex_home,
|
||||
&thread_id.to_string(),
|
||||
self.state_db.as_deref(),
|
||||
)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
internal_error(format!("failed to locate thread id {thread_id}: {err}"))
|
||||
})?
|
||||
.ok_or_else(|| invalid_request(format!("thread not found: {thread_id}")))?,
|
||||
};
|
||||
reconcile_rollout(
|
||||
Some(&state_db),
|
||||
@@ -307,16 +318,20 @@ impl ThreadGoalRequestProcessor {
|
||||
return Ok(state_db);
|
||||
}
|
||||
} else {
|
||||
find_thread_path_by_id_str(&self.config.codex_home, &thread_id.to_string())
|
||||
.await
|
||||
.map_err(|err| {
|
||||
internal_error(format!("failed to locate thread id {thread_id}: {err}"))
|
||||
})?
|
||||
.ok_or_else(|| invalid_request(format!("thread not found: {thread_id}")))?;
|
||||
find_thread_path_by_id_str(
|
||||
&self.config.codex_home,
|
||||
&thread_id.to_string(),
|
||||
self.state_db.as_deref(),
|
||||
)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
internal_error(format!("failed to locate thread id {thread_id}: {err}"))
|
||||
})?
|
||||
.ok_or_else(|| invalid_request(format!("thread not found: {thread_id}")))?;
|
||||
}
|
||||
|
||||
open_state_db_for_direct_thread_lookup(&self.config)
|
||||
.await
|
||||
self.state_db
|
||||
.clone()
|
||||
.ok_or_else(|| internal_error("sqlite state db unavailable for thread goals"))
|
||||
}
|
||||
|
||||
|
||||
@@ -259,6 +259,7 @@ pub(crate) struct ThreadRequestProcessor {
|
||||
pub(super) thread_watch_manager: ThreadWatchManager,
|
||||
pub(super) thread_list_state_permit: Arc<Semaphore>,
|
||||
pub(super) thread_goal_processor: ThreadGoalRequestProcessor,
|
||||
pub(super) state_db: Option<StateDbHandle>,
|
||||
pub(super) background_tasks: TaskTracker,
|
||||
}
|
||||
|
||||
@@ -278,6 +279,7 @@ impl ThreadRequestProcessor {
|
||||
thread_watch_manager: ThreadWatchManager,
|
||||
thread_list_state_permit: Arc<Semaphore>,
|
||||
thread_goal_processor: ThreadGoalRequestProcessor,
|
||||
state_db: Option<StateDbHandle>,
|
||||
) -> Self {
|
||||
Self {
|
||||
auth_manager,
|
||||
@@ -293,6 +295,7 @@ impl ThreadRequestProcessor {
|
||||
thread_watch_manager,
|
||||
thread_list_state_permit,
|
||||
thread_goal_processor,
|
||||
state_db,
|
||||
background_tasks: TaskTracker::new(),
|
||||
}
|
||||
}
|
||||
@@ -1170,7 +1173,7 @@ impl ThreadRequestProcessor {
|
||||
.map_err(|err| invalid_request(format!("invalid thread id: {err}")))?;
|
||||
|
||||
let mut thread_ids = vec![thread_id];
|
||||
if let Some(state_db_ctx) = get_state_db(&self.config).await {
|
||||
if let Some(state_db_ctx) = self.state_db.as_ref() {
|
||||
let descendants = state_db_ctx
|
||||
.list_thread_spawn_descendants(thread_id)
|
||||
.await
|
||||
@@ -1391,14 +1394,10 @@ impl ThreadRequestProcessor {
|
||||
}
|
||||
|
||||
async fn memory_reset_response_inner(&self) -> Result<MemoryResetResponse, JSONRPCErrorError> {
|
||||
let state_db = StateRuntime::init(
|
||||
self.config.sqlite_home.clone(),
|
||||
self.config.model_provider_id.clone(),
|
||||
)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
internal_error(format!("failed to open state db for memory reset: {err}"))
|
||||
})?;
|
||||
let state_db = self
|
||||
.state_db
|
||||
.clone()
|
||||
.ok_or_else(|| internal_error("sqlite state db unavailable for memory reset"))?;
|
||||
|
||||
state_db.clear_memory_data().await.map_err(|err| {
|
||||
internal_error(format!("failed to clear memory rows in state db: {err}"))
|
||||
@@ -1445,7 +1444,7 @@ impl ThreadRequestProcessor {
|
||||
let loaded_thread = self.thread_manager.get_thread(thread_uuid).await.ok();
|
||||
let mut state_db_ctx = loaded_thread.as_ref().and_then(|thread| thread.state_db());
|
||||
if state_db_ctx.is_none() {
|
||||
state_db_ctx = get_state_db(&self.config).await;
|
||||
state_db_ctx = self.state_db.clone();
|
||||
}
|
||||
let Some(state_db_ctx) = state_db_ctx else {
|
||||
return Err(internal_error(format!(
|
||||
@@ -1583,33 +1582,37 @@ impl ThreadRequestProcessor {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let rollout_path =
|
||||
match find_thread_path_by_id_str(&self.config.codex_home, &thread_uuid.to_string())
|
||||
.await
|
||||
let rollout_path = match find_thread_path_by_id_str(
|
||||
&self.config.codex_home,
|
||||
&thread_uuid.to_string(),
|
||||
self.state_db.as_deref(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(path)) => path,
|
||||
Ok(None) => match find_archived_thread_path_by_id_str(
|
||||
&self.config.codex_home,
|
||||
&thread_uuid.to_string(),
|
||||
self.state_db.as_deref(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(path)) => path,
|
||||
Ok(None) => match find_archived_thread_path_by_id_str(
|
||||
&self.config.codex_home,
|
||||
&thread_uuid.to_string(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(Some(path)) => path,
|
||||
Ok(None) => {
|
||||
return Err(invalid_request(format!("thread not found: {thread_uuid}")));
|
||||
}
|
||||
Err(err) => {
|
||||
return Err(internal_error(format!(
|
||||
"failed to locate archived thread id {thread_uuid}: {err}"
|
||||
)));
|
||||
}
|
||||
},
|
||||
Ok(None) => {
|
||||
return Err(invalid_request(format!("thread not found: {thread_uuid}")));
|
||||
}
|
||||
Err(err) => {
|
||||
return Err(internal_error(format!(
|
||||
"failed to locate thread id {thread_uuid}: {err}"
|
||||
"failed to locate archived thread id {thread_uuid}: {err}"
|
||||
)));
|
||||
}
|
||||
};
|
||||
},
|
||||
Err(err) => {
|
||||
return Err(internal_error(format!(
|
||||
"failed to locate thread id {thread_uuid}: {err}"
|
||||
)));
|
||||
}
|
||||
};
|
||||
|
||||
reconcile_rollout(
|
||||
Some(state_db_ctx),
|
||||
@@ -2555,7 +2558,7 @@ impl ThreadRequestProcessor {
|
||||
let InitialHistory::Resumed(resumed_history) = thread_history else {
|
||||
return None;
|
||||
};
|
||||
let state_db_ctx = get_state_db(&self.config).await?;
|
||||
let state_db_ctx = self.state_db.clone()?;
|
||||
let persisted_metadata = state_db_ctx
|
||||
.get_thread(resumed_history.conversation_id)
|
||||
.await
|
||||
@@ -2922,7 +2925,9 @@ impl ThreadRequestProcessor {
|
||||
}
|
||||
|
||||
async fn attach_thread_name(&self, thread_id: ThreadId, thread: &mut Thread) {
|
||||
if let Some(title) = title_from_state_db(&self.config, thread_id).await {
|
||||
if let Some(title) =
|
||||
title_from_state_db(&self.config, self.state_db.as_ref(), thread_id).await
|
||||
{
|
||||
set_thread_name_from_title(thread, title);
|
||||
}
|
||||
}
|
||||
@@ -3683,8 +3688,12 @@ async fn read_summary_from_state_db_context_by_thread_id(
|
||||
Some(summary_from_thread_metadata(&metadata))
|
||||
}
|
||||
|
||||
async fn title_from_state_db(config: &Config, thread_id: ThreadId) -> Option<String> {
|
||||
if let Some(state_db_ctx) = open_state_db_for_direct_thread_lookup(config).await
|
||||
async fn title_from_state_db(
|
||||
config: &Config,
|
||||
state_db_ctx: Option<&StateDbHandle>,
|
||||
thread_id: ThreadId,
|
||||
) -> Option<String> {
|
||||
if let Some(state_db_ctx) = state_db_ctx
|
||||
&& let Some(metadata) = state_db_ctx.get_thread(thread_id).await.ok().flatten()
|
||||
&& let Some(title) = distinct_title(&metadata)
|
||||
{
|
||||
|
||||
@@ -1,13 +1,5 @@
|
||||
use super::*;
|
||||
|
||||
pub(super) async fn open_state_db_for_direct_thread_lookup(
|
||||
config: &Config,
|
||||
) -> Option<StateDbHandle> {
|
||||
StateRuntime::init(config.sqlite_home.clone(), config.model_provider_id.clone())
|
||||
.await
|
||||
.ok()
|
||||
}
|
||||
|
||||
pub(crate) async fn read_summary_from_rollout(
|
||||
path: &Path,
|
||||
fallback_provider: &str,
|
||||
|
||||
@@ -13,6 +13,7 @@ pub(crate) struct TurnRequestProcessor {
|
||||
thread_state_manager: ThreadStateManager,
|
||||
thread_watch_manager: ThreadWatchManager,
|
||||
thread_list_state_permit: Arc<Semaphore>,
|
||||
state_db: Option<StateDbHandle>,
|
||||
}
|
||||
|
||||
impl TurnRequestProcessor {
|
||||
@@ -29,6 +30,7 @@ impl TurnRequestProcessor {
|
||||
thread_state_manager: ThreadStateManager,
|
||||
thread_watch_manager: ThreadWatchManager,
|
||||
thread_list_state_permit: Arc<Semaphore>,
|
||||
state_db: Option<StateDbHandle>,
|
||||
) -> Self {
|
||||
Self {
|
||||
auth_manager,
|
||||
@@ -42,6 +44,7 @@ impl TurnRequestProcessor {
|
||||
thread_state_manager,
|
||||
thread_watch_manager,
|
||||
thread_list_state_permit,
|
||||
state_db,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -891,16 +894,20 @@ impl TurnRequestProcessor {
|
||||
let rollout_path = if let Some(path) = parent_thread.rollout_path() {
|
||||
path
|
||||
} else {
|
||||
find_thread_path_by_id_str(&self.config.codex_home, &parent_thread_id.to_string())
|
||||
.await
|
||||
.map_err(|err| {
|
||||
internal_error(format!(
|
||||
"failed to locate thread id {parent_thread_id}: {err}"
|
||||
))
|
||||
})?
|
||||
.ok_or_else(|| {
|
||||
invalid_request(format!("no rollout found for thread id {parent_thread_id}"))
|
||||
})?
|
||||
find_thread_path_by_id_str(
|
||||
&self.config.codex_home,
|
||||
&parent_thread_id.to_string(),
|
||||
self.state_db.as_deref(),
|
||||
)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
internal_error(format!(
|
||||
"failed to locate thread id {parent_thread_id}: {err}"
|
||||
))
|
||||
})?
|
||||
.ok_or_else(|| {
|
||||
invalid_request(format!("no rollout found for thread id {parent_thread_id}"))
|
||||
})?
|
||||
};
|
||||
|
||||
let mut config = self.config.as_ref().clone();
|
||||
|
||||
Reference in New Issue
Block a user