mirror of
https://github.com/openai/codex.git
synced 2026-11-06 08:31:30 +00:00
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 authored and GitHub committed
1 parent
36460387ec
commit
7e310bc7f3
53 files changed
+768
-745
No files matched your search
@@ -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,
|
||||
|
||||
@@ -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_ ",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in new issue
Block a user