Files
codex/codex-rs/core/src/prompt_debug.rs
T
Rasmus Rygaard 7e310bc7f3 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`
2026-05-05 21:45:29 +00:00

103 lines
3.4 KiB
Rust

use std::collections::HashSet;
use std::sync::Arc;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::EnvironmentManagerArgs;
use codex_exec_server::ExecServerRuntimePaths;
use codex_login::AuthManager;
use codex_protocol::error::Result as CodexResult;
use codex_protocol::models::ResponseInputItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::SessionSource;
use codex_protocol::user_input::UserInput;
use tokio_util::sync::CancellationToken;
use crate::config::Config;
use crate::session::session::Session;
use crate::session::turn::build_prompt;
use crate::session::turn::built_tools;
use crate::thread_manager::ThreadManager;
use crate::thread_manager::agent_graph_store_from_state_db;
use crate::thread_manager::init_state_db_from_config;
use crate::thread_manager::thread_store_from_config;
/// Build the model-visible `input` list for a single debug turn.
#[doc(hidden)]
pub async fn build_prompt_input(
mut config: Config,
input: Vec<UserInput>,
) -> CodexResult<Vec<ResponseItem>> {
config.ephemeral = true;
let auth_manager =
AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false).await;
let local_runtime_paths = ExecServerRuntimePaths::from_optional_paths(
config.codex_self_exe.clone(),
config.codex_linux_sandbox_exe.clone(),
)?;
let state_db = init_state_db_from_config(&config)
.await
.ok_or_else(|| std::io::Error::other("prompt debug requires state db"))?;
let thread_store = thread_store_from_config(&config, state_db.clone());
let agent_graph_store = agent_graph_store_from_state_db(state_db.clone());
let thread_manager = ThreadManager::new(
&config,
Arc::clone(&auth_manager),
SessionSource::Exec,
Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::new(local_runtime_paths)).await),
/*analytics_events_client*/ None,
state_db,
thread_store,
agent_graph_store,
);
let thread = thread_manager.start_thread(config).await?;
let output = build_prompt_input_from_session(thread.thread.codex.session.as_ref(), input).await;
let shutdown = thread.thread.shutdown_and_wait().await;
let _removed = thread_manager.remove_thread(&thread.thread_id).await;
shutdown?;
output
}
pub(crate) async fn build_prompt_input_from_session(
sess: &Session,
input: Vec<UserInput>,
) -> CodexResult<Vec<ResponseItem>> {
let turn_context = sess.new_default_turn().await;
sess.record_context_updates_and_set_reference_context_item(turn_context.as_ref())
.await;
if !input.is_empty() {
let input_item = ResponseInputItem::from(input);
let response_item = ResponseItem::from(input_item);
sess.record_conversation_items(turn_context.as_ref(), std::slice::from_ref(&response_item))
.await;
}
let prompt_input = sess
.clone_history()
.await
.for_prompt(&turn_context.model_info.input_modalities);
let router = built_tools(
sess,
turn_context.as_ref(),
&prompt_input,
&HashSet::new(),
Some(turn_context.turn_skills.outcome.as_ref()),
&CancellationToken::new(),
)
.await?;
let base_instructions = sess.get_base_instructions().await;
let prompt = build_prompt(
prompt_input,
router.as_ref(),
turn_context.as_ref(),
base_instructions,
);
Ok(prompt.get_formatted_input())
}