Files
codex/codex-rs/core/tests/suite/prompt_caching.rs
pakrym-oai 30ca89424c Always fallback to real shell (#6953)
Either cmd.exe or `/bin/sh`.
2025-11-20 10:58:46 -08:00

745 lines
25 KiB
Rust

#![allow(clippy::unwrap_used)]
use codex_core::features::Feature;
use codex_core::model_family::find_family_for_model;
use codex_core::protocol::AskForApproval;
use codex_core::protocol::ENVIRONMENT_CONTEXT_OPEN_TAG;
use codex_core::protocol::EventMsg;
use codex_core::protocol::Op;
use codex_core::protocol::SandboxPolicy;
use codex_core::protocol_config_types::ReasoningEffort;
use codex_core::protocol_config_types::ReasoningSummary;
use codex_core::shell::Shell;
use codex_core::shell::default_user_shell;
use codex_protocol::user_input::UserInput;
use core_test_support::load_sse_fixture_with_id;
use core_test_support::responses::mount_sse_once;
use core_test_support::responses::start_mock_server;
use core_test_support::skip_if_no_network;
use core_test_support::test_codex::TestCodex;
use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use tempfile::TempDir;
fn text_user_input(text: String) -> serde_json::Value {
serde_json::json!({
"type": "message",
"role": "user",
"content": [ { "type": "input_text", "text": text } ]
})
}
fn default_env_context_str(cwd: &str, shell: &Shell) -> String {
let shell_name = shell.name();
format!(
r#"<environment_context>
<cwd>{cwd}</cwd>
<approval_policy>on-request</approval_policy>
<sandbox_mode>read-only</sandbox_mode>
<network_access>restricted</network_access>
<shell>{shell_name}</shell>
</environment_context>"#
)
}
/// Build minimal SSE stream with completed marker using the JSON fixture.
fn sse_completed(id: &str) -> String {
load_sse_fixture_with_id("tests/fixtures/completed_template.json", id)
}
fn assert_tool_names(body: &serde_json::Value, expected_names: &[&str]) {
assert_eq!(
body["tools"]
.as_array()
.unwrap()
.iter()
.map(|t| t["name"].as_str().unwrap().to_string())
.collect::<Vec<_>>(),
expected_names
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn codex_mini_latest_tools() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
use pretty_assertions::assert_eq;
let server = start_mock_server().await;
let req1 = mount_sse_once(&server, sse_completed("resp-1")).await;
let req2 = mount_sse_once(&server, sse_completed("resp-2")).await;
let TestCodex { codex, .. } = test_codex()
.with_config(|config| {
config.user_instructions = Some("be consistent and helpful".to_string());
config.features.disable(Feature::ApplyPatchFreeform);
config.model = "codex-mini-latest".to_string();
config.model_family = find_family_for_model("codex-mini-latest")
.expect("model family for codex-mini-latest");
})
.build(&server)
.await?;
codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "hello 1".into(),
}],
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "hello 2".into(),
}],
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
let expected_instructions = [
include_str!("../../prompt.md"),
include_str!("../../../apply-patch/apply_patch_tool_instructions.md"),
]
.join("\n");
let body0 = req1.single_request().body_json();
assert_eq!(
body0["instructions"],
serde_json::json!(expected_instructions),
);
let body1 = req2.single_request().body_json();
assert_eq!(
body1["instructions"],
serde_json::json!(expected_instructions),
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn prompt_tools_are_consistent_across_requests() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
use pretty_assertions::assert_eq;
let server = start_mock_server().await;
let req1 = mount_sse_once(&server, sse_completed("resp-1")).await;
let req2 = mount_sse_once(&server, sse_completed("resp-2")).await;
let TestCodex { codex, config, .. } = test_codex()
.with_config(|config| {
config.user_instructions = Some("be consistent and helpful".to_string());
})
.build(&server)
.await?;
let base_instructions = config.model_family.base_instructions.clone();
codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "hello 1".into(),
}],
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "hello 2".into(),
}],
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
let expected_tools_names = vec![
"shell_command",
"list_mcp_resources",
"list_mcp_resource_templates",
"read_mcp_resource",
"update_plan",
"apply_patch",
"view_image",
];
let body0 = req1.single_request().body_json();
let expected_instructions = if expected_tools_names.contains(&"apply_patch") {
base_instructions
} else {
[
base_instructions.clone(),
include_str!("../../../apply-patch/apply_patch_tool_instructions.md").to_string(),
]
.join("\n")
};
assert_eq!(
body0["instructions"],
serde_json::json!(expected_instructions),
);
assert_tool_names(&body0, &expected_tools_names);
let body1 = req2.single_request().body_json();
assert_eq!(
body1["instructions"],
serde_json::json!(expected_instructions),
);
assert_tool_names(&body1, &expected_tools_names);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn prefixes_context_and_instructions_once_and_consistently_across_requests()
-> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
use pretty_assertions::assert_eq;
let server = start_mock_server().await;
let req1 = mount_sse_once(&server, sse_completed("resp-1")).await;
let req2 = mount_sse_once(&server, sse_completed("resp-2")).await;
let TestCodex { codex, config, .. } = test_codex()
.with_config(|config| {
config.user_instructions = Some("be consistent and helpful".to_string());
})
.build(&server)
.await?;
codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "hello 1".into(),
}],
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "hello 2".into(),
}],
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
let shell = default_user_shell();
let cwd_str = config.cwd.to_string_lossy();
let expected_env_text = default_env_context_str(&cwd_str, &shell);
let expected_ui_text = format!(
"# AGENTS.md instructions for {cwd_str}\n\n<INSTRUCTIONS>\nbe consistent and helpful\n</INSTRUCTIONS>"
);
let expected_env_msg = serde_json::json!({
"type": "message",
"role": "user",
"content": [ { "type": "input_text", "text": expected_env_text } ]
});
let expected_ui_msg = serde_json::json!({
"type": "message",
"role": "user",
"content": [ { "type": "input_text", "text": expected_ui_text } ]
});
let expected_user_message_1 = serde_json::json!({
"type": "message",
"role": "user",
"content": [ { "type": "input_text", "text": "hello 1" } ]
});
let body1 = req1.single_request().body_json();
assert_eq!(
body1["input"],
serde_json::json!([expected_ui_msg, expected_env_msg, expected_user_message_1])
);
let expected_user_message_2 = serde_json::json!({
"type": "message",
"role": "user",
"content": [ { "type": "input_text", "text": "hello 2" } ]
});
let body2 = req2.single_request().body_json();
let expected_body2 = serde_json::json!(
[
body1["input"].as_array().unwrap().as_slice(),
[expected_user_message_2].as_slice(),
]
.concat()
);
assert_eq!(body2["input"], expected_body2);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn overrides_turn_context_but_keeps_cached_prefix_and_key_constant() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
use pretty_assertions::assert_eq;
let server = start_mock_server().await;
let req1 = mount_sse_once(&server, sse_completed("resp-1")).await;
let req2 = mount_sse_once(&server, sse_completed("resp-2")).await;
let TestCodex { codex, .. } = test_codex()
.with_config(|config| {
config.user_instructions = Some("be consistent and helpful".to_string());
})
.build(&server)
.await?;
// First turn
codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "hello 1".into(),
}],
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
let writable = TempDir::new().unwrap();
codex
.submit(Op::OverrideTurnContext {
cwd: None,
approval_policy: Some(AskForApproval::Never),
sandbox_policy: Some(SandboxPolicy::WorkspaceWrite {
writable_roots: vec![writable.path().to_path_buf()],
network_access: true,
exclude_tmpdir_env_var: true,
exclude_slash_tmp: true,
}),
model: Some("o3".to_string()),
effort: Some(Some(ReasoningEffort::High)),
summary: Some(ReasoningSummary::Detailed),
})
.await?;
// Second turn after overrides
codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "hello 2".into(),
}],
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
let body1 = req1.single_request().body_json();
let body2 = req2.single_request().body_json();
// prompt_cache_key should remain constant across overrides
assert_eq!(
body1["prompt_cache_key"], body2["prompt_cache_key"],
"prompt_cache_key should not change across overrides"
);
// The entire prefix from the first request should be identical and reused
// as the prefix of the second request, ensuring cache hit potential.
let expected_user_message_2 = serde_json::json!({
"type": "message",
"role": "user",
"content": [ { "type": "input_text", "text": "hello 2" } ]
});
// After overriding the turn context, the environment context should be emitted again
// reflecting the new approval policy and sandbox settings. Omit cwd because it did
// not change.
let shell = default_user_shell();
let expected_env_text_2 = format!(
r#"<environment_context>
<approval_policy>never</approval_policy>
<sandbox_mode>workspace-write</sandbox_mode>
<network_access>enabled</network_access>
<writable_roots>
<root>{}</root>
</writable_roots>
<shell>{}</shell>
</environment_context>"#,
writable.path().display(),
shell.name()
);
let expected_env_msg_2 = serde_json::json!({
"type": "message",
"role": "user",
"content": [ { "type": "input_text", "text": expected_env_text_2 } ]
});
let expected_body2 = serde_json::json!(
[
body1["input"].as_array().unwrap().as_slice(),
[expected_env_msg_2, expected_user_message_2].as_slice(),
]
.concat()
);
assert_eq!(body2["input"], expected_body2);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn override_before_first_turn_emits_environment_context() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let req = mount_sse_once(&server, sse_completed("resp-1")).await;
let TestCodex { codex, .. } = test_codex().build(&server).await?;
codex
.submit(Op::OverrideTurnContext {
cwd: None,
approval_policy: Some(AskForApproval::Never),
sandbox_policy: None,
model: None,
effort: None,
summary: None,
})
.await?;
codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "first message".into(),
}],
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
let body = req.single_request().body_json();
let input = body["input"]
.as_array()
.expect("input array must be present");
assert!(
!input.is_empty(),
"expected at least environment context and user message"
);
let env_msg = &input[1];
let env_text = env_msg["content"][0]["text"]
.as_str()
.expect("environment context text");
assert!(
env_text.starts_with(ENVIRONMENT_CONTEXT_OPEN_TAG),
"second entry should be environment context, got: {env_text}"
);
assert!(
env_text.contains("<approval_policy>never</approval_policy>"),
"environment context should reflect overridden approval policy: {env_text}"
);
let env_count = input
.iter()
.filter(|msg| {
msg["content"]
.as_array()
.and_then(|content| {
content.iter().find(|item| {
item["type"].as_str() == Some("input_text")
&& item["text"]
.as_str()
.map(|text| text.starts_with(ENVIRONMENT_CONTEXT_OPEN_TAG))
.unwrap_or(false)
})
})
.is_some()
})
.count();
assert_eq!(
env_count, 2,
"environment context should appear exactly twice, found {env_count}"
);
let user_msg = &input[2];
let user_text = user_msg["content"][0]["text"]
.as_str()
.expect("user message text");
assert_eq!(user_text, "first message");
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn per_turn_overrides_keep_cached_prefix_and_key_constant() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
use pretty_assertions::assert_eq;
let server = start_mock_server().await;
let req1 = mount_sse_once(&server, sse_completed("resp-1")).await;
let req2 = mount_sse_once(&server, sse_completed("resp-2")).await;
let TestCodex { codex, .. } = test_codex()
.with_config(|config| {
config.user_instructions = Some("be consistent and helpful".to_string());
})
.build(&server)
.await?;
// First turn
codex
.submit(Op::UserInput {
items: vec![UserInput::Text {
text: "hello 1".into(),
}],
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
// Second turn using per-turn overrides via UserTurn
let new_cwd = TempDir::new().unwrap();
let writable = TempDir::new().unwrap();
codex
.submit(Op::UserTurn {
items: vec![UserInput::Text {
text: "hello 2".into(),
}],
cwd: new_cwd.path().to_path_buf(),
approval_policy: AskForApproval::Never,
sandbox_policy: SandboxPolicy::WorkspaceWrite {
writable_roots: vec![writable.path().to_path_buf()],
network_access: true,
exclude_tmpdir_env_var: true,
exclude_slash_tmp: true,
},
model: "o3".to_string(),
effort: Some(ReasoningEffort::High),
summary: ReasoningSummary::Detailed,
final_output_json_schema: None,
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
let body1 = req1.single_request().body_json();
let body2 = req2.single_request().body_json();
// prompt_cache_key should remain constant across per-turn overrides
assert_eq!(
body1["prompt_cache_key"], body2["prompt_cache_key"],
"prompt_cache_key should not change across per-turn overrides"
);
// The entire prefix from the first request should be identical and reused
// as the prefix of the second request.
let expected_user_message_2 = serde_json::json!({
"type": "message",
"role": "user",
"content": [ { "type": "input_text", "text": "hello 2" } ]
});
let shell = default_user_shell();
let expected_env_text_2 = format!(
r#"<environment_context>
<cwd>{}</cwd>
<approval_policy>never</approval_policy>
<sandbox_mode>workspace-write</sandbox_mode>
<network_access>enabled</network_access>
<writable_roots>
<root>{}</root>
</writable_roots>
<shell>{}</shell>
</environment_context>"#,
new_cwd.path().display(),
writable.path().display(),
shell.name(),
);
let expected_env_msg_2 = serde_json::json!({
"type": "message",
"role": "user",
"content": [ { "type": "input_text", "text": expected_env_text_2 } ]
});
let expected_body2 = serde_json::json!(
[
body1["input"].as_array().unwrap().as_slice(),
[expected_env_msg_2, expected_user_message_2].as_slice(),
]
.concat()
);
assert_eq!(body2["input"], expected_body2);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn send_user_turn_with_no_changes_does_not_send_environment_context() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
use pretty_assertions::assert_eq;
let server = start_mock_server().await;
let req1 = mount_sse_once(&server, sse_completed("resp-1")).await;
let req2 = mount_sse_once(&server, sse_completed("resp-2")).await;
let TestCodex { codex, config, .. } = test_codex()
.with_config(|config| {
config.user_instructions = Some("be consistent and helpful".to_string());
})
.build(&server)
.await?;
let default_cwd = config.cwd.clone();
let default_approval_policy = config.approval_policy;
let default_sandbox_policy = config.sandbox_policy.clone();
let default_model = config.model.clone();
let default_effort = config.model_reasoning_effort;
let default_summary = config.model_reasoning_summary;
codex
.submit(Op::UserTurn {
items: vec![UserInput::Text {
text: "hello 1".into(),
}],
cwd: default_cwd.clone(),
approval_policy: default_approval_policy,
sandbox_policy: default_sandbox_policy.clone(),
model: default_model.clone(),
effort: default_effort,
summary: default_summary,
final_output_json_schema: None,
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
codex
.submit(Op::UserTurn {
items: vec![UserInput::Text {
text: "hello 2".into(),
}],
cwd: default_cwd.clone(),
approval_policy: default_approval_policy,
sandbox_policy: default_sandbox_policy.clone(),
model: default_model.clone(),
effort: default_effort,
summary: default_summary,
final_output_json_schema: None,
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
let body1 = req1.single_request().body_json();
let body2 = req2.single_request().body_json();
let shell = default_user_shell();
let default_cwd_lossy = default_cwd.to_string_lossy();
let expected_ui_text = format!(
"# AGENTS.md instructions for {default_cwd_lossy}\n\n<INSTRUCTIONS>\nbe consistent and helpful\n</INSTRUCTIONS>"
);
let expected_ui_msg = text_user_input(expected_ui_text);
let expected_env_msg_1 = text_user_input(default_env_context_str(&default_cwd_lossy, &shell));
let expected_user_message_1 = text_user_input("hello 1".to_string());
let expected_input_1 = serde_json::Value::Array(vec![
expected_ui_msg.clone(),
expected_env_msg_1.clone(),
expected_user_message_1.clone(),
]);
assert_eq!(body1["input"], expected_input_1);
let expected_user_message_2 = text_user_input("hello 2".to_string());
let expected_input_2 = serde_json::Value::Array(vec![
expected_ui_msg,
expected_env_msg_1,
expected_user_message_1,
expected_user_message_2,
]);
assert_eq!(body2["input"], expected_input_2);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn send_user_turn_with_changes_sends_environment_context() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
use pretty_assertions::assert_eq;
let server = start_mock_server().await;
let req1 = mount_sse_once(&server, sse_completed("resp-1")).await;
let req2 = mount_sse_once(&server, sse_completed("resp-2")).await;
let TestCodex { codex, config, .. } = test_codex()
.with_config(|config| {
config.user_instructions = Some("be consistent and helpful".to_string());
})
.build(&server)
.await?;
let default_cwd = config.cwd.clone();
let default_approval_policy = config.approval_policy;
let default_sandbox_policy = config.sandbox_policy.clone();
let default_model = config.model.clone();
let default_effort = config.model_reasoning_effort;
let default_summary = config.model_reasoning_summary;
codex
.submit(Op::UserTurn {
items: vec![UserInput::Text {
text: "hello 1".into(),
}],
cwd: default_cwd.clone(),
approval_policy: default_approval_policy,
sandbox_policy: default_sandbox_policy.clone(),
model: default_model,
effort: default_effort,
summary: default_summary,
final_output_json_schema: None,
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
codex
.submit(Op::UserTurn {
items: vec![UserInput::Text {
text: "hello 2".into(),
}],
cwd: default_cwd.clone(),
approval_policy: AskForApproval::Never,
sandbox_policy: SandboxPolicy::DangerFullAccess,
model: "o3".to_string(),
effort: Some(ReasoningEffort::High),
summary: ReasoningSummary::Detailed,
final_output_json_schema: None,
})
.await?;
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
let body1 = req1.single_request().body_json();
let body2 = req2.single_request().body_json();
let shell = default_user_shell();
let expected_ui_text = format!(
"# AGENTS.md instructions for {}\n\n<INSTRUCTIONS>\nbe consistent and helpful\n</INSTRUCTIONS>",
default_cwd.to_string_lossy()
);
let expected_ui_msg = serde_json::json!({
"type": "message",
"role": "user",
"content": [ { "type": "input_text", "text": expected_ui_text } ]
});
let expected_env_text_1 = default_env_context_str(&default_cwd.to_string_lossy(), &shell);
let expected_env_msg_1 = text_user_input(expected_env_text_1);
let expected_user_message_1 = text_user_input("hello 1".to_string());
let expected_input_1 = serde_json::Value::Array(vec![
expected_ui_msg.clone(),
expected_env_msg_1.clone(),
expected_user_message_1.clone(),
]);
assert_eq!(body1["input"], expected_input_1);
let shell_name = shell.name();
let expected_env_msg_2 = text_user_input(format!(
r#"<environment_context>
<approval_policy>never</approval_policy>
<sandbox_mode>danger-full-access</sandbox_mode>
<network_access>enabled</network_access>
<shell>{shell_name}</shell>
</environment_context>"#
));
let expected_user_message_2 = text_user_input("hello 2".to_string());
let expected_input_2 = serde_json::Value::Array(vec![
expected_ui_msg,
expected_env_msg_1,
expected_user_message_1,
expected_env_msg_2,
expected_user_message_2,
]);
assert_eq!(body2["input"], expected_input_2);
Ok(())
}