use anyhow::Result;
use codex_config::types::McpServerConfig;
use codex_config::types::McpServerTransportConfig;
use codex_core::config::Config;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_features::Feature;
use codex_login::CodexAuth;
use codex_protocol::ThreadId;
use codex_protocol::config_types::WebSearchMode;
use codex_protocol::dynamic_tools::DynamicToolFunctionSpec;
use codex_protocol::dynamic_tools::DynamicToolNamespaceSpec;
use codex_protocol::dynamic_tools::DynamicToolNamespaceTool;
use codex_protocol::dynamic_tools::DynamicToolSpec;
use codex_protocol::models::PermissionProfile;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use codex_protocol::protocol::RolloutItem;
use codex_protocol::protocol::RolloutLine;
use codex_protocol::protocol::SessionMeta;
use codex_protocol::protocol::SessionMetaLine;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::UserMessageEvent;
use codex_protocol::user_input::UserInput;
use codex_web_search_extension::install as install_web_search_extension;
use core_test_support::responses;
use core_test_support::responses::ev_completed;
use core_test_support::responses::ev_function_call;
use core_test_support::responses::ev_response_created;
use core_test_support::responses::ev_web_search_call_done;
use core_test_support::responses::mount_sse_once;
use core_test_support::responses::mount_sse_sequence;
use core_test_support::responses::start_mock_server;
use core_test_support::skip_if_no_network;
use core_test_support::stdio_server_bin;
use core_test_support::test_codex::local_selections;
use core_test_support::test_codex::test_codex;
use core_test_support::test_codex::turn_permission_fields;
use core_test_support::wait_for_event;
use core_test_support::wait_for_event_match;
use core_test_support::wait_for_mcp_server;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::collections::HashMap;
use std::fs;
use std::sync::Arc;
use tokio::time::Duration;
use tracing_subscriber::prelude::*;
use uuid::Uuid;
use wiremock::Mock;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn new_thread_is_recorded_in_state_db() -> Result<()> {
let server = start_mock_server().await;
let mut builder = test_codex().with_config(|config| {
config
.features
.enable(Feature::Sqlite)
.expect("test config should allow feature update");
});
let test = builder.build(&server).await?;
let thread_id = test.session_configured.thread_id;
let rollout_path = test.codex.rollout_path().expect("rollout path");
let db_path = codex_state::state_db_path(test.config.sqlite_home.as_path());
for _ in 0..100 {
if tokio::fs::try_exists(&db_path).await.unwrap_or(false) {
break;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
let db = test.codex.state_db().expect("state db enabled");
assert!(
!rollout_path.exists(),
"fresh thread rollout should not be materialized before first user message"
);
let initial_metadata = db.get_thread(thread_id).await?;
assert!(
initial_metadata.is_none(),
"fresh thread should not be recorded in state db before first user message"
);
test.submit_turn("materialize rollout").await?;
let mut metadata = None;
for _ in 0..100 {
metadata = db.get_thread(thread_id).await?;
if metadata.is_some() {
break;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
let metadata = metadata.expect("thread should exist in state db");
assert_eq!(metadata.id, thread_id);
assert_eq!(metadata.rollout_path, rollout_path.to_path_buf());
assert!(
rollout_path.exists(),
"rollout should be materialized after first user message"
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn resume_restores_dynamic_tools_from_rollout_with_sqlite_enabled() -> Result<()> {
let server = start_mock_server().await;
let mock = mount_sse_sequence(
&server,
vec![
responses::sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]),
responses::sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]),
],
)
.await;
let namespace = "resume_tools";
let namespace_description = "Tools available after resume.";
let tool_name = "resume_lookup";
let tool_description = "Look up a value after resume.";
let input_schema = json!({
"type": "object",
"properties": { "query": { "type": "string" } },
"required": ["query"],
"additionalProperties": false,
});
let dynamic_tool = DynamicToolSpec::Namespace(DynamicToolNamespaceSpec {
name: namespace.to_string(),
description: namespace_description.to_string(),
tools: vec![DynamicToolNamespaceTool::Function(
DynamicToolFunctionSpec {
name: tool_name.to_string(),
description: tool_description.to_string(),
input_schema: input_schema.clone(),
defer_loading: false,
},
)],
});
let mut builder = test_codex().with_config(|config| {
config
.features
.enable(Feature::Sqlite)
.expect("test config should allow feature update");
});
let base_test = builder.build(&server).await?;
let started = base_test
.thread_manager
.start_thread_with_tools(base_test.config.clone(), vec![dynamic_tool])
.await?;
let rollout_path = started
.session_configured
.rollout_path
.clone()
.expect("rollout path");
started
.thread
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "persist this thread".to_string(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
responsesapi_client_metadata: None,
additional_context: Default::default(),
thread_settings: Default::default(),
})
.await?;
wait_for_event(&started.thread, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
let mut resume_builder = test_codex().with_config(|config| {
config
.features
.enable(Feature::Sqlite)
.expect("test config should allow feature update");
});
let resumed = resume_builder
.resume(&server, base_test.home.clone(), rollout_path)
.await?;
resumed.submit_turn("use the restored tool").await?;
let requests = mock.requests();
assert_eq!(requests.len(), 2);
let resumed_body = requests[1].body_json();
let tools = resumed_body
.get("tools")
.and_then(serde_json::Value::as_array)
.expect("resumed request tools");
let restored_namespace = tools
.iter()
.find(|tool| tool.get("name") == Some(&json!(namespace)))
.expect("dynamic tool namespace should be restored from rollout metadata");
assert_eq!(
restored_namespace,
&json!({
"type": "namespace",
"name": namespace,
"description": namespace_description,
"tools": [{
"type": "function",
"name": tool_name,
"description": tool_description,
"strict": false,
"parameters": input_schema,
}],
})
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn resume_restores_legacy_dynamic_tools_from_rollout_with_sqlite_enabled() -> Result<()> {
let server = start_mock_server().await;
let mock = mount_sse_sequence(
&server,
vec![
responses::sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]),
responses::sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]),
],
)
.await;
let namespace = "resume_tools";
let tool_name = "resume_lookup";
let tool_description = "Look up a value after resume.";
let input_schema = json!({
"type": "object",
"properties": { "query": { "type": "string" } },
"required": ["query"],
"additionalProperties": false,
});
let mut builder = test_codex().with_config(|config| {
config
.features
.enable(Feature::Sqlite)
.expect("test config should allow feature update");
});
let base_test = builder.build(&server).await?;
let started = base_test
.thread_manager
.start_thread_with_tools(base_test.config.clone(), Vec::new())
.await?;
let rollout_path = started
.session_configured
.rollout_path
.clone()
.expect("rollout path");
started
.thread
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "persist this thread".to_string(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
responsesapi_client_metadata: None,
additional_context: Default::default(),
thread_settings: Default::default(),
})
.await?;
wait_for_event(&started.thread, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
started.thread.submit(Op::Shutdown).await?;
wait_for_event(&started.thread, |event| {
matches!(event, EventMsg::ShutdownComplete)
})
.await;
let mut rollout_lines = fs::read_to_string(&rollout_path)?
.lines()
.map(serde_json::from_str::<serde_json::Value>)
.collect::<serde_json::Result<Vec<_>>>()?;
rollout_lines.first_mut().expect("session metadata line")["payload"]["dynamic_tools"] = json!([{
"namespace": namespace,
"name": tool_name,
"description": tool_description,
"inputSchema": input_schema,
"exposeToContext": true,
}]);
let rollout = rollout_lines
.iter()
.map(serde_json::to_string)
.collect::<serde_json::Result<Vec<_>>>()?
.join("\n");
fs::write(&rollout_path, format!("{rollout}\n"))?;
let mut resume_builder = test_codex().with_config(|config| {
config
.features
.enable(Feature::Sqlite)
.expect("test config should allow feature update");
});
let resumed = resume_builder
.resume(&server, base_test.home.clone(), rollout_path)
.await?;
resumed.submit_turn("use the restored tool").await?;
let requests = mock.requests();
assert_eq!(requests.len(), 2);
let resumed_body = requests[1].body_json();
let tools = resumed_body
.get("tools")
.and_then(serde_json::Value::as_array)
.expect("resumed request tools");
let restored_namespace = tools
.iter()
.find(|tool| tool.get("name") == Some(&json!(namespace)))
.expect("dynamic tool namespace should be restored from rollout metadata");
assert_eq!(
restored_namespace,
&json!({
"type": "namespace",
"name": namespace,
"description": "Tools in the resume_tools namespace.",
"tools": [{
"type": "function",
"name": tool_name,
"description": tool_description,
"strict": false,
"parameters": input_schema,
}],
})
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn backfill_scans_existing_rollouts() -> Result<()> {
let server = start_mock_server().await;
let uuid = Uuid::now_v7();
let thread_id = ThreadId::from_string(&uuid.to_string())?;
let rollout_rel_path = format!("sessions/2026/01/27/rollout-2026-01-27T12-00-00-{uuid}.jsonl");
let rollout_rel_path_for_hook = rollout_rel_path.clone();
let mut builder = test_codex()
.with_pre_build_hook(move |codex_home| {
let rollout_path = codex_home.join(&rollout_rel_path_for_hook);
let parent = rollout_path
.parent()
.expect("rollout path should have parent");
fs::create_dir_all(parent).expect("should create rollout directory");
let session_meta_line = SessionMetaLine {
meta: SessionMeta {
session_id: thread_id.into(),
id: thread_id,
forked_from_id: None,
parent_thread_id: None,
timestamp: "2026-01-27T12:00:00Z".to_string(),
cwd: codex_home.to_path_buf(),
originator: "test".to_string(),
cli_version: "test".to_string(),
source: SessionSource::default(),
thread_source: None,
agent_path: None,
agent_nickname: None,
agent_role: None,
model_provider: None,
base_instructions: None,
dynamic_tools: None,
selected_capability_roots: Vec::new(),
memory_mode: None,
history_mode: Default::default(),
history_base: None,
subagent_history_start_ordinal: None,
multi_agent_version: None,
context_window: None,
},
git: None,
};
let lines = [
RolloutLine {
timestamp: "2026-01-27T12:00:00Z".to_string(),
ordinal: None,
item: RolloutItem::SessionMeta(session_meta_line),
},
RolloutLine {
timestamp: "2026-01-27T12:00:01Z".to_string(),
ordinal: None,
item: RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent {
client_id: None,
message: "hello from backfill".to_string(),
images: None,
local_images: Vec::new(),
text_elements: Vec::new(),
..Default::default()
})),
},
];
let jsonl = lines
.iter()
.map(|line| serde_json::to_string(line).expect("rollout line should serialize"))
.collect::<Vec<_>>()
.join("\n");
fs::write(&rollout_path, format!("{jsonl}\n")).expect("should write rollout file");
})
.with_config(|config| {
config
.features
.enable(Feature::Sqlite)
.expect("test config should allow feature update");
});
let test = builder.build(&server).await?;
let db_path = codex_state::state_db_path(test.config.sqlite_home.as_path());
let rollout_path = test.config.codex_home.join(&rollout_rel_path);
let default_provider = test.config.model_provider_id.clone();
for _ in 0..20 {
if tokio::fs::try_exists(&db_path).await.unwrap_or(false) {
break;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
let db = test.codex.state_db().expect("state db enabled");
let mut metadata = None;
for _ in 0..40 {
metadata = db.get_thread(thread_id).await?;
if metadata.is_some() {
break;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
let metadata = metadata.expect("backfilled thread should exist in state db");
assert_eq!(metadata.id, thread_id);
assert_eq!(metadata.rollout_path, rollout_path.to_path_buf());
assert_eq!(metadata.model_provider, default_provider);
assert!(metadata.first_user_message.is_some());
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn user_messages_persist_in_state_db() -> Result<()> {
let server = start_mock_server().await;
mount_sse_sequence(
&server,
vec![
responses::sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]),
responses::sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]),
],
)
.await;
let mut builder = test_codex().with_config(|config| {
config
.features
.enable(Feature::Sqlite)
.expect("test config should allow feature update");
});
let test = builder.build(&server).await?;
let db_path = codex_state::state_db_path(test.config.sqlite_home.as_path());
for _ in 0..100 {
if tokio::fs::try_exists(&db_path).await.unwrap_or(false) {
break;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
test.submit_turn("hello from sqlite").await?;
test.submit_turn("another message").await?;
let db = test.codex.state_db().expect("state db enabled");
let thread_id = test.session_configured.thread_id;
let mut metadata = None;
for _ in 0..100 {
metadata = db.get_thread(thread_id).await?;
if metadata
.as_ref()
.map(|entry| entry.first_user_message.is_some())
.unwrap_or(false)
{
break;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
let metadata = metadata.expect("thread should exist in state db");
assert!(metadata.first_user_message.is_some());
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn web_search_marks_thread_memory_mode_polluted_when_configured() -> Result<()> {
let server = start_mock_server().await;
mount_sse_sequence(
&server,
vec![responses::sse(vec![
ev_response_created("resp-1"),
ev_web_search_call_done("ws-1", "completed", "weather seattle"),
ev_completed("resp-1"),
])],
)
.await;
let mut builder = test_codex().with_config(|config| {
config
.features
.enable(Feature::Sqlite)
.expect("test config should allow feature update");
config.memories.disable_on_external_context = true;
});
let test = builder.build(&server).await?;
let db = test.codex.state_db().expect("state db enabled");
let thread_id = test.session_configured.thread_id;
test.submit_turn("search the web").await?;
let mut memory_mode = None;
for _ in 0..100 {
memory_mode = db.get_thread_memory_mode(thread_id).await?;
if memory_mode.as_deref() == Some("polluted") {
break;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
assert_eq!(memory_mode.as_deref(), Some("polluted"));
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn standalone_web_search_marks_thread_memory_mode_polluted_when_configured() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
Mock::given(method("POST"))
.and(path("/v1/alpha/search"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"output": "Search result",
})))
.expect(1)
.mount(&server)
.await;
mount_sse_sequence(
&server,
vec![
responses::sse(vec![
ev_response_created("resp-1"),
responses::ev_function_call_with_namespace(
"web-run-1",
"web",
"run",
&json!({
"search_query": [{"q": "standalone web search"}],
})
.to_string(),
),
ev_completed("resp-1"),
]),
responses::sse(vec![
responses::ev_assistant_message("msg-1", "done"),
ev_completed("resp-2"),
]),
],
)
.await;
let auth = CodexAuth::from_api_key("dummy");
let auth_manager = codex_core::test_support::auth_manager_from_auth(auth.clone());
let mut extension_builder = ExtensionRegistryBuilder::<Config>::new();
install_web_search_extension(&mut extension_builder, auth_manager);
let mut builder = test_codex()
.with_auth(auth)
.with_extensions(Arc::new(extension_builder.build()))
.with_config(|config| {
config
.features
.enable(Feature::Sqlite)
.expect("test config should allow feature update");
config
.features
.enable(Feature::StandaloneWebSearch)
.expect("standalone web search should be enabled");
config.memories.disable_on_external_context = true;
config
.web_search_mode
.set(WebSearchMode::Live)
.expect("web search mode should be accepted");
});
let test = builder.build(&server).await?;
let db = test.codex.state_db().expect("state db enabled");
let thread_id = test.session_configured.thread_id;
test.submit_turn("search the web").await?;
let mut memory_mode = None;
for _ in 0..100 {
memory_mode = db.get_thread_memory_mode(thread_id).await?;
if memory_mode.as_deref() == Some("polluted") {
break;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
assert_eq!(memory_mode.as_deref(), Some("polluted"));
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn mcp_call_marks_thread_memory_mode_polluted_when_configured() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let call_id = "call-123";
let server_name = "rmcp";
let namespace = format!("mcp__{server_name}");
mount_sse_once(
&server,
responses::sse(vec![
ev_response_created("resp-1"),
responses::ev_function_call_with_namespace(
call_id,
&namespace,
"echo",
"{\"message\":\"ping\"}",
),
ev_completed("resp-1"),
]),
)
.await;
mount_sse_once(
&server,
responses::sse(vec![
responses::ev_assistant_message("msg-1", "rmcp echo tool completed."),
ev_completed("resp-2"),
]),
)
.await;
let rmcp_test_server_bin = stdio_server_bin()?;
let mut builder = test_codex().with_config(move |config| {
config
.features
.enable(Feature::Sqlite)
.expect("test config should allow feature update");
config.memories.disable_on_external_context = true;
let mut servers = config.mcp_servers.get().clone();
servers.insert(
server_name.to_string(),
McpServerConfig {
auth: Default::default(),
transport: McpServerTransportConfig::Stdio {
command: rmcp_test_server_bin,
args: Vec::new(),
env: Some(HashMap::from([(
"MCP_TEST_VALUE".to_string(),
"propagated-env".to_string(),
)])),
env_vars: Vec::new(),
cwd: None,
},
environment_id: "local".to_string(),
enabled: true,
required: false,
supports_parallel_tool_calls: false,
disabled_reason: None,
startup_timeout_sec: Some(Duration::from_secs(10)),
tool_timeout_sec: None,
default_tools_approval_mode: None,
enabled_tools: None,
disabled_tools: None,
scopes: None,
oauth: None,
oauth_resource: None,
tools: HashMap::new(),
},
);
config
.mcp_servers
.set(servers)
.expect("test mcp servers should accept any configuration");
});
let test = builder.build(&server).await?;
wait_for_mcp_server(&test.codex, server_name).await?;
let db = test.codex.state_db().expect("state db enabled");
let thread_id = test.session_configured.thread_id;
let cwd = test.config.cwd.clone();
let (sandbox_policy, permission_profile) =
turn_permission_fields(PermissionProfile::read_only(), cwd.as_path());
test.codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "call the rmcp echo tool".to_string(),
text_elements: Vec::new(),
}],
final_output_json_schema: None,
responsesapi_client_metadata: None,
additional_context: Default::default(),
thread_settings: codex_protocol::protocol::ThreadSettingsOverrides {
environments: Some(local_selections(cwd)),
approval_policy: Some(AskForApproval::Never),
sandbox_policy: Some(sandbox_policy),
permission_profile,
collaboration_mode: Some(codex_protocol::config_types::CollaborationMode {
mode: codex_protocol::config_types::ModeKind::Default,
settings: codex_protocol::config_types::Settings {
model: test.session_configured.model.clone(),
reasoning_effort: None,
developer_instructions: None,
},
}),
..Default::default()
},
})
.await?;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::McpToolCallEnd(_))
})
.await;
wait_for_event_match(&test.codex, |event| match event {
EventMsg::Error(err) => Some(Err(anyhow::anyhow!(err.message.clone()))),
EventMsg::TurnComplete(_) => Some(Ok(())),
_ => None,
})
.await?;
let mut memory_mode = None;
for _ in 0..100 {
memory_mode = db.get_thread_memory_mode(thread_id).await?;
if memory_mode.as_deref() == Some("polluted") {
break;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
assert_eq!(memory_mode.as_deref(), Some("polluted"));
Ok(())
}
#[tokio::test(flavor = "current_thread")]
async fn tool_call_logs_include_thread_id() -> Result<()> {
let server = start_mock_server().await;
let call_id = "call-1";
let args = json!({
"command": "echo hello",
"timeout_ms": 1_000,
"login": false,
});
let args_json = serde_json::to_string(&args)?;
mount_sse_sequence(
&server,
vec![
responses::sse(vec![
ev_response_created("resp-1"),
ev_function_call(call_id, "shell_command", &args_json),
ev_completed("resp-1"),
]),
responses::sse(vec![ev_completed("resp-2")]),
],
)
.await;
let mut builder = test_codex().with_config(|config| {
config
.features
.enable(Feature::Sqlite)
.expect("test config should allow feature update");
});
let test = builder.build(&server).await?;
let db = test.codex.state_db().expect("state db enabled");
let expected_thread_id = test.session_configured.thread_id.to_string();
test.submit_turn("run a shell command").await?;
let log_db_layer = codex_state::log_db::start(db.clone());
let subscriber = tracing_subscriber::registry().with(log_db_layer.clone());
let dispatch = tracing::Dispatch::new(subscriber);
tracing::dispatcher::with_default(&dispatch, || {
let span = tracing::info_span!("test_log_span", thread_id = %expected_thread_id);
let _entered = span.enter();
tracing::info!("ToolCall: shell_command {{\"command\":\"echo hello\"}}");
});
log_db_layer.flush().await;
let mut found = None;
for _ in 0..80 {
let query = codex_state::LogQuery {
descending: true,
limit: Some(20),
..Default::default()
};
let rows = db.query_logs(&query).await?;
if let Some(row) = rows.into_iter().find(|row| {
row.message
.as_deref()
.is_some_and(|m| m.contains("ToolCall:"))
}) {
let thread_id = row.thread_id;
let message = row.message;
found = Some((thread_id, message));
break;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
let (thread_id, message) = found.expect("expected ToolCall log row");
assert_eq!(thread_id, Some(expected_thread_id));
assert!(
message
.as_deref()
.is_some_and(|text| text.contains("ToolCall:")),
"expected ToolCall message, got {message:?}"
);
Ok(())
}