mirror of
https://github.com/openai/codex.git
synced 2026-11-14 21:40:39 +00:00
Abhinav
4598f33790
## Why Model lookups through the startup catalog can use the wrong provider's capabilities or instructions after a default provider switch or when a thread has its own catalog. ## What changed - Resolve model metadata for sampling, prediction forks, and memory writes through the owning thread's catalog. - Make the thread's models manager available to extensions before startup callbacks and use it for guardian sampler configuration. - Import external sessions using the destination provider and originator's catalog, loading its offline cache only when no refreshed catalog is available. ## Testing Add regression coverage for sampling after a default provider switch, parent-specific prediction fork capabilities, thread-scoped model lookups, and imported sessions retaining destination provider instructions and originator. GitOrigin-RevId: c86081b9e55c5e2b42a7f082966485d7fb5139d4
1515 lines
56 KiB
Rust
1515 lines
56 KiB
Rust
//! Exercises sampling through JSON-RPC, including cache identity and cleanup.
|
||
|
||
use anyhow::Result;
|
||
use app_test_support::MockResponsesConfig;
|
||
use app_test_support::TestAppServer;
|
||
use codex_app_server_protocol::DynamicToolFunctionSpec;
|
||
use codex_app_server_protocol::DynamicToolSpec;
|
||
use codex_app_server_protocol::HooksListParams;
|
||
use codex_app_server_protocol::HooksListResponse;
|
||
use codex_app_server_protocol::RequestId;
|
||
use codex_app_server_protocol::ServerRequest;
|
||
use codex_app_server_protocol::ThreadForkParams;
|
||
use codex_app_server_protocol::ThreadForkResponse;
|
||
use codex_app_server_protocol::ThreadHistoryMode;
|
||
use codex_app_server_protocol::ThreadLoadedListResponse;
|
||
use codex_app_server_protocol::ThreadSampleCancelResponse;
|
||
use codex_app_server_protocol::ThreadSampleParams;
|
||
use codex_app_server_protocol::ThreadSampleResponse;
|
||
use codex_app_server_protocol::ThreadSource;
|
||
use codex_app_server_protocol::ThreadStartParams;
|
||
use codex_app_server_protocol::TurnCompletedNotification;
|
||
use codex_app_server_protocol::TurnInterruptParams;
|
||
use codex_app_server_protocol::TurnInterruptResponse;
|
||
use codex_app_server_protocol::TurnStartParams;
|
||
use codex_app_server_protocol::TurnStartResponse;
|
||
use codex_app_server_protocol::TurnStatus;
|
||
use codex_app_server_protocol::UserInput;
|
||
use codex_core::config::ConfigBuilder;
|
||
use codex_core::test_support::construct_model_info_offline;
|
||
use codex_features::Feature;
|
||
use codex_protocol::config_types::CollaborationMode;
|
||
use codex_protocol::config_types::ModeKind;
|
||
use codex_protocol::config_types::Settings;
|
||
use codex_protocol::openai_models::ModelServiceTier;
|
||
use codex_protocol::openai_models::ReasoningEffort;
|
||
use core_test_support::context_snapshot;
|
||
use core_test_support::context_snapshot::ContextSnapshotOptions;
|
||
use core_test_support::context_snapshot::SnapshotEntry;
|
||
use core_test_support::responses;
|
||
use core_test_support::skip_if_host_windows;
|
||
use core_test_support::skip_if_remote;
|
||
use core_test_support::streaming_sse::StreamingSseChunk;
|
||
use core_test_support::streaming_sse::start_streaming_sse_server;
|
||
use pretty_assertions::assert_eq;
|
||
use serde_json::Value;
|
||
use serde_json::json;
|
||
use std::collections::HashMap;
|
||
use std::sync::Arc;
|
||
use std::time::Duration;
|
||
use tempfile::TempDir;
|
||
use test_case::test_case;
|
||
use test_case::test_matrix;
|
||
use tokio::sync::Notify;
|
||
use tokio::time::timeout;
|
||
use wiremock::ResponseTemplate;
|
||
|
||
const WAIT: Duration = Duration::from_secs(40);
|
||
|
||
async fn client(
|
||
home: &TempDir,
|
||
server: &wiremock::MockServer,
|
||
overrides: bool,
|
||
) -> Result<TestAppServer> {
|
||
client_with_effort_updates(home, server, overrides, EffortUpdates::Supported, "").await
|
||
}
|
||
|
||
enum EffortUpdates {
|
||
Supported,
|
||
Unsupported,
|
||
}
|
||
|
||
async fn client_with_effort_updates(
|
||
home: &TempDir,
|
||
server: &wiremock::MockServer,
|
||
overrides: bool,
|
||
effort_updates: EffortUpdates,
|
||
extra_config: &str,
|
||
) -> Result<TestAppServer> {
|
||
let config = ConfigBuilder::default()
|
||
.codex_home(home.path().to_path_buf())
|
||
.loader_overrides(codex_config::LoaderOverrides::without_managed_config_for_tests())
|
||
.build()
|
||
.await?;
|
||
let mut model = construct_model_info_offline("gpt-5.4", &config);
|
||
model.supports_reasoning_effort_updates = matches!(effort_updates, EffortUpdates::Supported);
|
||
model.use_responses_lite = true;
|
||
model.service_tiers = vec![ModelServiceTier {
|
||
id: "priority".to_string(),
|
||
name: "Fast".to_string(),
|
||
description: "Priority processing.".to_string(),
|
||
}];
|
||
let catalog = home.path().join("catalog.json");
|
||
std::fs::write(
|
||
&catalog,
|
||
serde_json::to_string(&json!({"models": [model]}))?,
|
||
)?;
|
||
let mut config = MockResponsesConfig::new(&server.uri()).with_model("gpt-5.4").with_provider_name("OpenAI")
|
||
.with_root_config(&format!("model_catalog_json = {}\nmodel_reasoning_effort = \"high\"\nmodel_post_turn_compact_threshold_percent = 0", serde_json::to_string(&catalog)?));
|
||
config = if overrides {
|
||
config.enable_feature(Feature::ReasoningEffortOverride)
|
||
} else {
|
||
config.disable_feature(Feature::ReasoningEffortOverride)
|
||
};
|
||
config.with_extra_config(extra_config).write(home.path())?;
|
||
TestAppServer::builder()
|
||
.with_codex_home(home.path())
|
||
.build_initialized_with_timeout(WAIT)
|
||
.await
|
||
}
|
||
|
||
fn input(text: &str) -> Vec<UserInput> {
|
||
vec![UserInput::Text {
|
||
text: text.to_string(),
|
||
text_elements: Vec::new(),
|
||
}]
|
||
}
|
||
|
||
fn comparable_request(body: Value) -> Value {
|
||
let mut body = responses::strip_response_item_ids_from_json(body);
|
||
body.as_object_mut()
|
||
.expect("request object")
|
||
.remove("client_metadata");
|
||
for item in body["input"].as_array_mut().expect("request input") {
|
||
if let Some(metadata) = item
|
||
.get_mut("internal_chat_message_metadata_passthrough")
|
||
.and_then(Value::as_object_mut)
|
||
{
|
||
metadata.remove("create_time");
|
||
metadata.remove("turn_id");
|
||
}
|
||
}
|
||
body
|
||
}
|
||
|
||
async fn sample(
|
||
client: &mut TestAppServer,
|
||
params: ThreadSampleParams,
|
||
) -> Result<ThreadSampleResponse> {
|
||
let id = client
|
||
.send_raw_request("thread/sample", Some(serde_json::to_value(params)?))
|
||
.await?;
|
||
timeout(WAIT, client.read_response(id)).await?
|
||
}
|
||
|
||
async fn ready_parent(
|
||
client: &mut TestAppServer,
|
||
params: ThreadStartParams,
|
||
) -> Result<(String, String)> {
|
||
let parent = client.start_thread(params).await?.thread.id;
|
||
let turn = client
|
||
.start_turn_and_wait_for_completion(TurnStartParams {
|
||
thread_id: parent.clone(),
|
||
input: input("Initial turn"),
|
||
..Default::default()
|
||
})
|
||
.await?
|
||
.turn
|
||
.id;
|
||
Ok((parent, turn))
|
||
}
|
||
|
||
async fn prediction_fork(client: &mut TestAppServer, parent: &str) -> Result<String> {
|
||
let id = client
|
||
.send_thread_fork_request(ThreadForkParams {
|
||
thread_id: parent.to_string(),
|
||
experimental_prediction_mode: true,
|
||
ephemeral: true,
|
||
exclude_turns: true,
|
||
..Default::default()
|
||
})
|
||
.await?;
|
||
let fork: ThreadForkResponse = timeout(WAIT, client.read_response(id)).await??;
|
||
Ok(fork.thread.id)
|
||
}
|
||
|
||
#[test_matrix(
|
||
[ThreadHistoryMode::Legacy, ThreadHistoryMode::Paginated],
|
||
[TurnStatus::Completed, TurnStatus::Failed, TurnStatus::Interrupted],
|
||
[false, true]
|
||
)]
|
||
#[tokio::test]
|
||
async fn sample_preserves_context_and_keeps_source_unchanged(
|
||
history_mode: ThreadHistoryMode,
|
||
parent_status: TurnStatus,
|
||
allow_in_progress_turn: bool,
|
||
) -> Result<()> {
|
||
let server = responses::start_mock_server().await;
|
||
let compact_parent =
|
||
history_mode == ThreadHistoryMode::Paginated && parent_status == TurnStatus::Completed;
|
||
let mut parent_responses = vec![responses::sse(vec![
|
||
responses::ev_reasoning_item(
|
||
"parent-reasoning",
|
||
&["Check the release."],
|
||
&["Keep the existing context."],
|
||
),
|
||
if parent_status == TurnStatus::Interrupted {
|
||
responses::ev_function_call("parent-call", "parent_tool", "{}")
|
||
} else {
|
||
responses::ev_assistant_message("answer", "Done")
|
||
},
|
||
if parent_status == TurnStatus::Failed {
|
||
json!({"type":"response.failed","response":{"id":"parent","error":{"code":"insufficient_quota","message":"quota exceeded"}}})
|
||
} else {
|
||
responses::ev_completed_with_tokens(
|
||
"parent",
|
||
/*total_tokens*/ if compact_parent { 330_000 } else { 100 },
|
||
)
|
||
},
|
||
])];
|
||
if compact_parent {
|
||
parent_responses.push(responses::sse(vec![
|
||
json!({"type":"response.output_item.done","item":{
|
||
"type":"compaction", "encrypted_content":"sample-test-compacted-parent",
|
||
}}),
|
||
responses::ev_completed_with_tokens("compacted", /*total_tokens*/ 200),
|
||
]));
|
||
}
|
||
let sample_index = parent_responses.len();
|
||
parent_responses.extend([
|
||
responses::sse(vec![
|
||
json!({"type":"response.output_item.done","item":{"type":"message","id":"commentary","role":"assistant","phase":"commentary","content":[{"type":"output_text","text":"Checking the context"}]}}),
|
||
responses::ev_reasoning_item("reasoning", &["Summary"], &["Reasoning"]),
|
||
responses::ev_assistant_message("prediction", "Continue"),
|
||
json!({"type":"response.completed","response":{"id":"sample","usage":{"input_tokens":100,"input_tokens_details":{"cached_tokens":80},"output_tokens":5,"total_tokens":105}}}),
|
||
]),
|
||
responses::sse(vec![responses::ev_completed("prediction-fork")]),
|
||
responses::sse(vec![responses::ev_completed("follow-up")]),
|
||
]);
|
||
let mock = responses::mount_sse_sequence(&server, parent_responses).await;
|
||
let home = TempDir::new()?;
|
||
let mut client = client(&home, &server, /*overrides*/ true).await?;
|
||
let parent = client
|
||
.start_thread(ThreadStartParams {
|
||
history_mode: Some(history_mode),
|
||
base_instructions: Some("Retain these base instructions.".to_string()),
|
||
developer_instructions: Some("Retain this system policy context.".to_string()),
|
||
dynamic_tools: Some(vec![DynamicToolSpec::Function(DynamicToolFunctionSpec {
|
||
name: "parent_tool".to_string(),
|
||
description: "The parent's client-provided tool description.".to_string(),
|
||
input_schema: json!({"type":"object","properties":{}}),
|
||
defer_loading: false,
|
||
})]),
|
||
config: Some(std::collections::HashMap::from([
|
||
("features.request_permissions_tool".to_string(), json!(true)),
|
||
("model_auto_compact_token_limit".to_string(), json!(200_000)),
|
||
(
|
||
"model_post_turn_compact_threshold_percent".to_string(),
|
||
json!(if compact_parent { 75 } else { 0 }),
|
||
),
|
||
])),
|
||
..Default::default()
|
||
})
|
||
.await?;
|
||
let start_id = client
|
||
.send_turn_start_request(TurnStartParams {
|
||
thread_id: parent.thread.id.clone(),
|
||
input: input("Initial turn"),
|
||
turn_trigger: Some("parent-trigger".to_string()),
|
||
service_tier: Some(Some("priority".to_string())),
|
||
collaboration_mode: Some(CollaborationMode {
|
||
mode: ModeKind::Plan,
|
||
settings: Settings {
|
||
model: "gpt-5.4".to_string(),
|
||
reasoning_effort: Some(ReasoningEffort::High),
|
||
developer_instructions: Some(
|
||
"Keep the parent's planning instructions.".to_string(),
|
||
),
|
||
},
|
||
}),
|
||
disabled_plugin_ids: Some(vec!["unused@plugin".to_string()]),
|
||
..Default::default()
|
||
})
|
||
.await?;
|
||
let started: TurnStartResponse = timeout(WAIT, client.read_response(start_id)).await??;
|
||
if parent_status == TurnStatus::Interrupted {
|
||
assert!(matches!(
|
||
timeout(WAIT, client.read_stream_until_request_message()).await??,
|
||
ServerRequest::DynamicToolCall { .. }
|
||
));
|
||
let sample_id = client
|
||
.send_raw_request(
|
||
"thread/sample",
|
||
Some(json!({"threadId": parent.thread.id, "expectedTurnId": started.turn.id, "prompt": "Predict"})),
|
||
)
|
||
.await?;
|
||
let error = timeout(
|
||
WAIT,
|
||
client.read_stream_until_error_message(RequestId::Integer(sample_id)),
|
||
)
|
||
.await??;
|
||
assert_eq!(error.error.message, "sample source must be idle");
|
||
let interrupt_id = client
|
||
.send_turn_interrupt_request(TurnInterruptParams {
|
||
thread_id: parent.thread.id.clone(),
|
||
turn_id: started.turn.id,
|
||
})
|
||
.await?;
|
||
let _: TurnInterruptResponse = timeout(WAIT, client.read_response(interrupt_id)).await??;
|
||
}
|
||
let completed: TurnCompletedNotification =
|
||
timeout(WAIT, client.read_notification("turn/completed")).await??;
|
||
assert_eq!(completed.turn.status, parent_status);
|
||
let path = parent.thread.path.as_ref().expect("persistent source path");
|
||
let before = std::fs::read(path)?;
|
||
client.clear_message_buffer();
|
||
let result = sample(
|
||
&mut client,
|
||
ThreadSampleParams {
|
||
thread_id: parent.thread.id.clone(),
|
||
expected_turn_id: completed.turn.id.clone(),
|
||
allow_in_progress_turn,
|
||
prompt: "Continue identically".to_string(),
|
||
effort: Some(ReasoningEffort::Medium),
|
||
interrupt_on_cancel: Some(true),
|
||
thread_source: Some(ThreadSource::Feature("composer_predictions".to_string())),
|
||
turn_trigger: Some("composer_predictions".to_string()),
|
||
responsesapi_client_metadata: Some(HashMap::from([
|
||
(
|
||
"generation_trigger".to_string(),
|
||
"composer_prediction".to_string(),
|
||
),
|
||
("turn_trigger".to_string(), "client-supplied".to_string()),
|
||
])),
|
||
..Default::default()
|
||
},
|
||
)
|
||
.await?;
|
||
assert_eq!(
|
||
result,
|
||
ThreadSampleResponse {
|
||
source_turn_id: completed.turn.id,
|
||
response_id: "sample".to_string(),
|
||
text: Some("Continue".to_string()),
|
||
input_tokens: Some(100),
|
||
cached_input_tokens: Some(80),
|
||
output_tokens: Some(5),
|
||
}
|
||
);
|
||
assert_eq!(std::fs::read(path)?, before);
|
||
assert!(
|
||
!client
|
||
.pending_notification_methods()
|
||
.iter()
|
||
.any(|name| name == "thread/started")
|
||
);
|
||
let list_id = client
|
||
.send_raw_request("thread/loaded/list", Some(json!({})))
|
||
.await?;
|
||
let loaded: ThreadLoadedListResponse = timeout(WAIT, client.read_response(list_id)).await??;
|
||
assert_eq!(loaded.data, vec![parent.thread.id.clone()]);
|
||
let fork = prediction_fork(&mut client, &parent.thread.id).await?;
|
||
for thread_id in [fork, parent.thread.id] {
|
||
client
|
||
.start_turn_and_wait_for_completion(TurnStartParams {
|
||
thread_id,
|
||
input: input("Continue identically"),
|
||
effort: Some(ReasoningEffort::Medium),
|
||
..Default::default()
|
||
})
|
||
.await?;
|
||
}
|
||
let requests = mock.requests();
|
||
assert_eq!(requests.len(), sample_index + 3);
|
||
if history_mode == ThreadHistoryMode::Legacy && parent_status == TurnStatus::Completed {
|
||
insta::assert_snapshot!(
|
||
"sample_inherits_parent_context",
|
||
context_snapshot::format_request_history_snapshot(
|
||
"A parent uses high effort with custom instructions, client tools, permission tools, and priority service. A medium-effort sample, prediction fork, and subsequent parent turn preserve the same request context and high request-effort baseline.",
|
||
&requests,
|
||
&ContextSnapshotOptions::default().include_request_settings(),
|
||
)
|
||
);
|
||
}
|
||
if compact_parent {
|
||
assert_eq!(
|
||
requests[sample_index].inputs_of_type("compaction")[0]["encrypted_content"],
|
||
"sample-test-compacted-parent",
|
||
);
|
||
}
|
||
let prediction = &requests[sample_index];
|
||
let sample_input = prediction.input();
|
||
assert_eq!(marker_count(&sample_input, "<thread_snapshot>"), 0);
|
||
assert_eq!(
|
||
marker_count(&sample_input, "<turn_aborted>"),
|
||
usize::from(parent_status == TurnStatus::Interrupted),
|
||
"only the parent's real interruption marker should be inherited"
|
||
);
|
||
assert_eq!(
|
||
comparable_request(prediction.body_json()),
|
||
comparable_request(requests[sample_index + 1].body_json())
|
||
);
|
||
assert_eq!(
|
||
comparable_request(prediction.body_json()),
|
||
comparable_request(requests[sample_index + 2].body_json())
|
||
);
|
||
assert_eq!(prediction.body_json()["service_tier"], "priority");
|
||
assert!(prediction.body_json().to_string().contains("parent_tool"));
|
||
assert!(
|
||
prediction
|
||
.body_json()
|
||
.to_string()
|
||
.contains("request_permissions")
|
||
);
|
||
let metadata: Value = serde_json::from_str(
|
||
&prediction
|
||
.header("x-codex-turn-metadata")
|
||
.expect("turn metadata"),
|
||
)?;
|
||
assert_eq!(metadata["thread_source"], "composer_predictions");
|
||
assert_eq!(metadata["turn_trigger"], "composer_predictions");
|
||
assert_eq!(metadata["generation_trigger"], "composer_prediction");
|
||
let body_metadata: Value = serde_json::from_str(
|
||
prediction.body_json()["client_metadata"]["x-codex-turn-metadata"]
|
||
.as_str()
|
||
.expect("body turn metadata"),
|
||
)?;
|
||
assert_eq!(body_metadata, metadata);
|
||
let parent_metadata: Value = serde_json::from_str(
|
||
&requests[sample_index + 2]
|
||
.header("x-codex-turn-metadata")
|
||
.expect("parent turn metadata"),
|
||
)?;
|
||
assert!(parent_metadata.get("generation_trigger").is_none());
|
||
Ok(())
|
||
}
|
||
|
||
fn marker_count(input: &[Value], marker: &str) -> usize {
|
||
input
|
||
.iter()
|
||
.filter_map(|item| item["content"].as_array())
|
||
.flatten()
|
||
.filter_map(|part| part["text"].as_str())
|
||
.map(|text| text.matches(marker).count())
|
||
.sum()
|
||
}
|
||
|
||
fn assert_open_sample_input(input: &[Value]) {
|
||
assert_eq!(
|
||
(
|
||
marker_count(input, "<thread_snapshot>"),
|
||
marker_count(input, "<turn_aborted>"),
|
||
),
|
||
(1, 0),
|
||
"expected one snapshot note and no interruption marker"
|
||
);
|
||
let [.., note, prompt] = input else {
|
||
panic!("expected the snapshot note immediately before the sample prompt");
|
||
};
|
||
assert_eq!(note["role"], "developer");
|
||
let content = note["content"].as_array().expect("developer content");
|
||
assert_eq!(
|
||
content.last().expect("snapshot note")["text"],
|
||
concat!(
|
||
"<thread_snapshot>\n",
|
||
"This snapshot is a copy of the source thread’s history taken while its turn was still running. ",
|
||
"Making the copy did not stop the source turn. ",
|
||
"A tool call that was pending at capture may appear as \"aborted\" in this copy.\n",
|
||
"</thread_snapshot>"
|
||
)
|
||
);
|
||
assert_eq!(prompt["role"], "user");
|
||
assert_eq!(prompt["content"][0]["text"], "Predict");
|
||
}
|
||
|
||
#[test_case(ThreadHistoryMode::Legacy, false; "legacy")]
|
||
#[test_case(ThreadHistoryMode::Paginated, false; "paginated")]
|
||
#[test_case(ThreadHistoryMode::Paginated, true; "paginated_after_mid_turn_compaction")]
|
||
#[tokio::test]
|
||
async fn sample_in_progress_turn_uses_captured_history(
|
||
history_mode: ThreadHistoryMode,
|
||
compact_mid_turn: bool,
|
||
) -> Result<()> {
|
||
let (release_sample, sample_gate) = tokio::sync::oneshot::channel();
|
||
let mut streams = Vec::new();
|
||
if compact_mid_turn {
|
||
// A tool call keeps the turn open while its token usage triggers compaction.
|
||
streams.push(vec![StreamingSseChunk {
|
||
gate: None,
|
||
body: responses::sse(vec![
|
||
responses::ev_function_call("before-compaction", "unknown_tool", "{}"),
|
||
responses::ev_completed_with_tokens(
|
||
"before-compaction",
|
||
/*total_tokens*/ 330_000,
|
||
),
|
||
]),
|
||
}]);
|
||
streams.push(vec![StreamingSseChunk {
|
||
gate: None,
|
||
body: responses::sse(vec![
|
||
responses::ev_assistant_message("summary", "Mid-turn compaction summary"),
|
||
responses::ev_completed_with_tokens("compacted", /*total_tokens*/ 200),
|
||
]),
|
||
}]);
|
||
}
|
||
let child_request_index = streams.len() + 1;
|
||
streams.extend([
|
||
vec![StreamingSseChunk {
|
||
gate: None,
|
||
body: responses::sse(vec![
|
||
responses::ev_function_call("pending-call", "parent_tool", "{}"),
|
||
responses::ev_completed("parent"),
|
||
]),
|
||
}],
|
||
vec![StreamingSseChunk {
|
||
gate: Some(sample_gate),
|
||
body: responses::sse(vec![
|
||
responses::ev_assistant_message("prediction", "Continue"),
|
||
responses::ev_completed("sample"),
|
||
]),
|
||
}],
|
||
vec![StreamingSseChunk {
|
||
gate: None,
|
||
body: responses::sse(vec![responses::ev_completed("parent-done")]),
|
||
}],
|
||
]);
|
||
let (server, _) = start_streaming_sse_server(streams).await;
|
||
let home = TempDir::new()?;
|
||
MockResponsesConfig::new(server.uri())
|
||
.with_root_config("model_auto_compact_token_limit = 200000")
|
||
.write(home.path())?;
|
||
let mut client = TestAppServer::builder()
|
||
.with_codex_home(home.path())
|
||
.build_initialized_with_timeout(WAIT)
|
||
.await?;
|
||
let parent = client
|
||
.start_thread(ThreadStartParams {
|
||
history_mode: Some(history_mode),
|
||
base_instructions: compact_mid_turn
|
||
.then(|| "Help the user with their task.".to_string()),
|
||
dynamic_tools: Some(vec![DynamicToolSpec::Function(DynamicToolFunctionSpec {
|
||
name: "parent_tool".to_string(),
|
||
description: "Wait for the client.".to_string(),
|
||
input_schema: json!({"type":"object","properties":{}}),
|
||
defer_loading: false,
|
||
})]),
|
||
..Default::default()
|
||
})
|
||
.await?
|
||
.thread
|
||
.id;
|
||
let start = client
|
||
.send_turn_start_request(TurnStartParams {
|
||
thread_id: parent.clone(),
|
||
input: input("Initial turn"),
|
||
..Default::default()
|
||
})
|
||
.await?;
|
||
let started: TurnStartResponse = timeout(WAIT, client.read_response(start)).await??;
|
||
let ServerRequest::DynamicToolCall { request_id, .. } =
|
||
timeout(WAIT, client.read_stream_until_request_message()).await??
|
||
else {
|
||
anyhow::bail!("expected the parent's pending tool call");
|
||
};
|
||
let sample_id = client
|
||
.send_raw_request(
|
||
"thread/sample",
|
||
Some(json!({
|
||
"threadId": parent, "expectedTurnId": started.turn.id,
|
||
"allowInProgressTurn": true, "prompt": "Predict"
|
||
})),
|
||
)
|
||
.await?;
|
||
timeout(WAIT, server.wait_for_request_count(child_request_index + 1)).await?;
|
||
// Finish the parent after the child has captured the pending call.
|
||
client
|
||
.send_response(
|
||
request_id,
|
||
json!({
|
||
"contentItems": [{"type":"inputText","text":"parent tool result"}],
|
||
"success": true
|
||
}),
|
||
)
|
||
.await?;
|
||
let completed: TurnCompletedNotification =
|
||
timeout(WAIT, client.read_notification("turn/completed")).await??;
|
||
assert_eq!(completed.turn.id, started.turn.id);
|
||
assert_eq!(completed.turn.status, TurnStatus::Completed);
|
||
|
||
let requests = server
|
||
.requests()
|
||
.await
|
||
.iter()
|
||
.map(|request| serde_json::from_slice::<Value>(request))
|
||
.collect::<serde_json::Result<Vec<_>>>()?;
|
||
if compact_mid_turn {
|
||
let [initial, compaction, continuation, sample, follow_up] = requests.as_slice() else {
|
||
anyhow::bail!("expected five model requests, got {}", requests.len());
|
||
};
|
||
let entries = [
|
||
SnapshotEntry::body(initial).labeled("Parent starts"),
|
||
SnapshotEntry::body(compaction).labeled("Parent compacts"),
|
||
SnapshotEntry::body(continuation).labeled("Parent requests a tool"),
|
||
SnapshotEntry::body(sample).labeled("Child samples the open turn"),
|
||
SnapshotEntry::body(follow_up).labeled("Parent receives the tool result"),
|
||
];
|
||
insta::assert_snapshot!(
|
||
"sample_compacted_parent_with_pending_tool_call",
|
||
context_snapshot::format_context_snapshot(
|
||
"A child samples a compacted parent while a tool call is pending. The parent then receives the real tool result.",
|
||
&entries,
|
||
&ContextSnapshotOptions::default().rewrite_known_segments(),
|
||
)
|
||
);
|
||
}
|
||
let child = &requests[child_request_index];
|
||
let follow_up = &requests[child_request_index + 1];
|
||
let child_input = child["input"].as_array().expect("sample input");
|
||
let parent_input = follow_up["input"].as_array().expect("parent input");
|
||
assert_eq!(marker_count(parent_input, "<thread_snapshot>"), 0);
|
||
assert_eq!(marker_count(parent_input, "<turn_aborted>"), 0);
|
||
for (input, expected_output) in [
|
||
(parent_input, "parent tool result"),
|
||
(child_input, "aborted"),
|
||
] {
|
||
let call = input
|
||
.iter()
|
||
.position(|item| item["type"] == "function_call" && item["call_id"] == "pending-call")
|
||
.expect("captured call");
|
||
assert_eq!(input[call]["name"], "parent_tool");
|
||
if compact_mid_turn {
|
||
assert!(input[..call].iter().any(|item| {
|
||
item["content"][0]["text"]
|
||
.as_str()
|
||
.is_some_and(|text| text.contains("Mid-turn compaction summary"))
|
||
}));
|
||
}
|
||
let output = input
|
||
.iter()
|
||
.position(|item| {
|
||
item["type"] == "function_call_output" && item["call_id"] == "pending-call"
|
||
})
|
||
.expect("paired output");
|
||
assert!(call < output);
|
||
assert_eq!(input[output]["output"], expected_output);
|
||
}
|
||
assert_open_sample_input(child_input);
|
||
release_sample.send(()).expect("sample stream waiting");
|
||
let response: ThreadSampleResponse = timeout(WAIT, client.read_response(sample_id)).await??;
|
||
assert_eq!(response.source_turn_id, started.turn.id);
|
||
assert_eq!(response.text.as_deref(), Some("Continue"));
|
||
Ok(())
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn sample_after_mid_turn_compaction_without_a_pending_tool_call() -> Result<()> {
|
||
let (release_parent, parent_gate) = tokio::sync::oneshot::channel();
|
||
let mut first_response =
|
||
responses::ev_completed_with_tokens("before-compaction", /*total_tokens*/ 330_000);
|
||
first_response["response"]["end_turn"] = json!(false);
|
||
let stream = |gate, events| {
|
||
vec![StreamingSseChunk {
|
||
gate,
|
||
body: responses::sse(events),
|
||
}]
|
||
};
|
||
let (server, _) = start_streaming_sse_server(vec![
|
||
stream(None, vec![first_response]),
|
||
stream(
|
||
None,
|
||
vec![
|
||
responses::ev_assistant_message("summary", "Mid-turn compaction summary"),
|
||
responses::ev_completed_with_tokens("compacted", /*total_tokens*/ 200),
|
||
],
|
||
),
|
||
stream(
|
||
Some(parent_gate),
|
||
vec![responses::ev_completed("parent-done")],
|
||
),
|
||
stream(
|
||
None,
|
||
vec![
|
||
responses::ev_assistant_message("prediction", "Continue"),
|
||
responses::ev_completed("sample"),
|
||
],
|
||
),
|
||
])
|
||
.await;
|
||
let home = TempDir::new()?;
|
||
MockResponsesConfig::new(server.uri())
|
||
.with_root_config("model_auto_compact_token_limit = 200000")
|
||
.write(home.path())?;
|
||
let mut client = TestAppServer::builder()
|
||
.with_codex_home(home.path())
|
||
.build_initialized_with_timeout(WAIT)
|
||
.await?;
|
||
let parent = client
|
||
.start_thread(ThreadStartParams {
|
||
history_mode: Some(ThreadHistoryMode::Paginated),
|
||
..Default::default()
|
||
})
|
||
.await?
|
||
.thread;
|
||
let start = client
|
||
.send_turn_start_request(TurnStartParams {
|
||
thread_id: parent.id.clone(),
|
||
input: input("Initial turn"),
|
||
..Default::default()
|
||
})
|
||
.await?;
|
||
let started: TurnStartResponse = timeout(WAIT, client.read_response(start)).await??;
|
||
|
||
// The third request is the parent's first model step after compaction.
|
||
// Its response is gated, so the turn is still running when the child samples it.
|
||
timeout(WAIT, server.wait_for_request_count(/*count*/ 3)).await?;
|
||
let path = parent.path.as_ref().expect("persistent source path");
|
||
let before = std::fs::read(path)?;
|
||
let response = sample(
|
||
&mut client,
|
||
ThreadSampleParams {
|
||
thread_id: parent.id,
|
||
expected_turn_id: started.turn.id.clone(),
|
||
allow_in_progress_turn: true,
|
||
prompt: "Predict".to_string(),
|
||
..Default::default()
|
||
},
|
||
)
|
||
.await?;
|
||
assert_eq!(
|
||
std::fs::read(path)?,
|
||
before,
|
||
"sampling changed the parent rollout"
|
||
);
|
||
assert_eq!(response.source_turn_id, started.turn.id);
|
||
release_parent.send(()).expect("parent response waiting");
|
||
let completed: TurnCompletedNotification =
|
||
timeout(WAIT, client.read_notification("turn/completed")).await??;
|
||
assert_eq!(completed.turn.id, started.turn.id);
|
||
assert_eq!(completed.turn.status, TurnStatus::Completed);
|
||
|
||
let requests = server.requests().await;
|
||
assert_eq!(requests.len(), 4);
|
||
let child: Value = serde_json::from_slice(&requests[3])?;
|
||
let input = child["input"].as_array().expect("sample input");
|
||
assert!(
|
||
input.iter().any(|item| {
|
||
item["content"][0]["text"]
|
||
.as_str()
|
||
.is_some_and(|text| text.contains("Mid-turn compaction summary"))
|
||
}),
|
||
"expected the inherited compaction summary"
|
||
);
|
||
assert!(
|
||
input.iter().all(|item| item.get("call_id").is_none()),
|
||
"unexpected tool call or output"
|
||
);
|
||
assert_open_sample_input(input);
|
||
Ok(())
|
||
}
|
||
|
||
#[test_case(true, ReasoningEffort::High, ReasoningEffort::Low, ReasoningEffort::Medium; "pinned parent")]
|
||
#[test_case(false, ReasoningEffort::High, ReasoningEffort::Low, ReasoningEffort::Medium; "parent overrides disabled")]
|
||
#[test_case(true, ReasoningEffort::Low, ReasoningEffort::Medium, ReasoningEffort::High; "increasing pinned effort")]
|
||
#[test_case(false, ReasoningEffort::Low, ReasoningEffort::Medium, ReasoningEffort::High; "increasing effort without parent overrides")]
|
||
#[tokio::test]
|
||
async fn sample_preserves_effort_baseline(
|
||
overrides: bool,
|
||
initial_effort: ReasoningEffort,
|
||
selected_effort: ReasoningEffort,
|
||
sample_effort: ReasoningEffort,
|
||
) -> Result<()> {
|
||
let server = responses::start_mock_server().await;
|
||
let mock = responses::mount_sse_sequence(
|
||
&server,
|
||
["initial", "selected", "sample", "prediction-fork"]
|
||
.map(|id| {
|
||
responses::sse(vec![
|
||
responses::ev_reasoning_item(
|
||
&format!("reasoning-{id}"),
|
||
&["Preserve the context."],
|
||
&["Continue from the previous response."],
|
||
),
|
||
responses::ev_assistant_message(&format!("answer-{id}"), "Ready."),
|
||
responses::ev_completed(id),
|
||
])
|
||
})
|
||
.to_vec(),
|
||
)
|
||
.await;
|
||
let home = TempDir::new()?;
|
||
let mut client = client(&home, &server, overrides).await?;
|
||
let (parent_id, _) = ready_parent(
|
||
&mut client,
|
||
ThreadStartParams {
|
||
config: Some(std::collections::HashMap::from([(
|
||
"model_reasoning_effort".to_string(),
|
||
serde_json::to_value(&initial_effort)?,
|
||
)])),
|
||
..Default::default()
|
||
},
|
||
)
|
||
.await?;
|
||
let completed = client
|
||
.start_turn_and_wait_for_completion(TurnStartParams {
|
||
thread_id: parent_id.clone(),
|
||
input: input("Second turn"),
|
||
effort: Some(selected_effort.clone()),
|
||
turn_trigger: Some("parent-trigger".to_string()),
|
||
..Default::default()
|
||
})
|
||
.await?;
|
||
sample(
|
||
&mut client,
|
||
ThreadSampleParams {
|
||
thread_id: parent_id.clone(),
|
||
expected_turn_id: completed.turn.id,
|
||
prompt: "Predict".to_string(),
|
||
effort: Some(sample_effort.clone()),
|
||
..Default::default()
|
||
},
|
||
)
|
||
.await?;
|
||
let fork = prediction_fork(&mut client, &parent_id).await?;
|
||
client
|
||
.start_turn_and_wait_for_completion(TurnStartParams {
|
||
thread_id: fork,
|
||
input: input("Predict"),
|
||
effort: Some(sample_effort.clone()),
|
||
..Default::default()
|
||
})
|
||
.await?;
|
||
let requests = mock.requests();
|
||
assert_eq!(requests.len(), 4);
|
||
assert_eq!(
|
||
requests[0].body_json()["reasoning"]["effort"],
|
||
json!(initial_effort)
|
||
);
|
||
assert_eq!(
|
||
comparable_request(requests[2].body_json()),
|
||
comparable_request(requests[3].body_json())
|
||
);
|
||
let mut parent = comparable_request(requests[1].body_json());
|
||
let mut child = comparable_request(requests[2].body_json());
|
||
assert_eq!(
|
||
child["reasoning"]["effort"],
|
||
json!(if overrides {
|
||
initial_effort
|
||
} else {
|
||
selected_effort
|
||
})
|
||
);
|
||
let parent_input = parent["input"].take();
|
||
let child_input = child["input"].take();
|
||
assert_eq!(child, parent);
|
||
let parent_input = parent_input.as_array().expect("parent input");
|
||
let child_input = child_input.as_array().expect("sample input");
|
||
assert_eq!(&child_input[..parent_input.len()], parent_input.as_slice());
|
||
assert_eq!(
|
||
child_input.last(),
|
||
Some(&json!({"type":"configuration_update","reasoning":{"effort":sample_effort}}))
|
||
);
|
||
let metadata: Value = serde_json::from_str(
|
||
&requests[2]
|
||
.header("x-codex-turn-metadata")
|
||
.expect("turn metadata"),
|
||
)?;
|
||
assert_eq!(metadata["thread_source"], "thread_sample");
|
||
assert!(metadata.get("turn_trigger").is_none());
|
||
Ok(())
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn sample_skips_inherited_stop_hooks() -> Result<()> {
|
||
skip_if_host_windows!(Ok(()));
|
||
skip_if_remote!(Ok(()), "command hook uses host-local paths");
|
||
let server = responses::start_mock_server().await;
|
||
let mock = responses::mount_sse_sequence(
|
||
&server,
|
||
["parent", "sample"]
|
||
.map(|id| {
|
||
responses::sse(vec![
|
||
responses::ev_assistant_message(id, id),
|
||
responses::ev_completed(id),
|
||
])
|
||
})
|
||
.to_vec(),
|
||
)
|
||
.await;
|
||
let home = TempDir::new()?;
|
||
let log = home.path().join("stop.log");
|
||
let script = home.path().join("stop.sh");
|
||
std::fs::write(
|
||
&script,
|
||
format!(
|
||
"cat >/dev/null\nif [ -f '{log}' ]; then\n printf '%s' '{{\"decision\":\"block\",\"reason\":\"Continue working\"}}'\nfi\nprintf 'stop\\n' >> '{log}'\n",
|
||
log = log.display(),
|
||
),
|
||
)?;
|
||
let extra_config = format!(
|
||
"[[hooks.Stop]]\n[[hooks.Stop.hooks]]\ntype = \"command\"\ncommand = {}\n",
|
||
serde_json::to_string(&format!("sh '{}'", script.display()))?,
|
||
);
|
||
let mut client = client_with_effort_updates(
|
||
&home,
|
||
&server,
|
||
/*overrides*/ true,
|
||
EffortUpdates::Supported,
|
||
&extra_config,
|
||
)
|
||
.await?;
|
||
let list_id = client
|
||
.send_hooks_list_request(HooksListParams {
|
||
cwds: vec![home.path().to_path_buf()],
|
||
..Default::default()
|
||
})
|
||
.await?;
|
||
let hooks: HooksListResponse = client.read_response(list_id).await?;
|
||
let trust = hooks.data[0]
|
||
.hooks
|
||
.iter()
|
||
.map(|hook| (hook.key.clone(), json!({"trusted_hash": hook.current_hash})))
|
||
.collect::<serde_json::Map<_, _>>();
|
||
assert!(!trust.is_empty());
|
||
let (parent, turn) = ready_parent(
|
||
&mut client,
|
||
ThreadStartParams {
|
||
config: Some(std::collections::HashMap::from([(
|
||
"hooks.state".to_string(),
|
||
Value::Object(trust),
|
||
)])),
|
||
..Default::default()
|
||
},
|
||
)
|
||
.await?;
|
||
assert_eq!(std::fs::read_to_string(&log)?, "stop\n");
|
||
let result = sample(
|
||
&mut client,
|
||
ThreadSampleParams {
|
||
thread_id: parent,
|
||
expected_turn_id: turn,
|
||
prompt: "Predict".to_string(),
|
||
..Default::default()
|
||
},
|
||
)
|
||
.await?;
|
||
assert_eq!(result.text.as_deref(), Some("sample"));
|
||
assert_eq!(mock.requests().len(), 2);
|
||
assert_eq!(std::fs::read_to_string(log)?, "stop\n");
|
||
Ok(())
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn sample_returns_null_for_custom_tool_calls() -> Result<()> {
|
||
let server = responses::start_mock_server().await;
|
||
let mut tool = responses::ev_custom_tool_call(
|
||
"call",
|
||
"apply_patch",
|
||
"*** Begin Patch\n*** Add File: blocked.txt\n+blocked\n*** End Patch",
|
||
);
|
||
tool["type"] = json!("response.output_item.added");
|
||
let mock = responses::mount_sse_sequence(
|
||
&server,
|
||
vec![
|
||
responses::sse(vec![responses::ev_completed("parent")]),
|
||
responses::sse(vec![
|
||
responses::ev_assistant_message("partial", "Discard this partial answer"),
|
||
tool,
|
||
responses::ev_assistant_message("after-tool", "Discard this answer too"),
|
||
json!({"type":"response.completed","response":{"id":"blocked","end_turn":false,"usage":{"input_tokens":100,"input_tokens_details":{"cached_tokens":80},"output_tokens":5,"total_tokens":105}}}),
|
||
]),
|
||
],
|
||
)
|
||
.await;
|
||
let home = TempDir::new()?;
|
||
let mut client = client(&home, &server, /*overrides*/ true).await?;
|
||
let (parent_id, turn_id) = ready_parent(&mut client, ThreadStartParams::default()).await?;
|
||
let response = sample(
|
||
&mut client,
|
||
ThreadSampleParams {
|
||
thread_id: parent_id.clone(),
|
||
expected_turn_id: turn_id.clone(),
|
||
prompt: "Call the tool".to_string(),
|
||
..Default::default()
|
||
},
|
||
)
|
||
.await?;
|
||
assert_eq!(
|
||
response,
|
||
ThreadSampleResponse {
|
||
source_turn_id: turn_id,
|
||
response_id: "blocked".to_string(),
|
||
text: None,
|
||
input_tokens: Some(100),
|
||
cached_input_tokens: Some(80),
|
||
output_tokens: Some(5),
|
||
}
|
||
);
|
||
assert_eq!(mock.requests().len(), 2);
|
||
let id = client
|
||
.send_raw_request("thread/loaded/list", Some(json!({})))
|
||
.await?;
|
||
let loaded: ThreadLoadedListResponse = client.read_response(id).await?;
|
||
assert_eq!(loaded.data, vec![parent_id]);
|
||
Ok(())
|
||
}
|
||
|
||
#[test_case(false, true; "inherit MCPs by default")]
|
||
#[test_case(true, true; "disable MCPs with incremental tools")]
|
||
#[test_case(true, false; "retain MCPs without incremental tools")]
|
||
#[tokio::test]
|
||
async fn sample_blocks_mcp_execution_without_changing_parent_tools(
|
||
disable_mcp: bool,
|
||
incremental_tools: bool,
|
||
) -> Result<()> {
|
||
let server = responses::start_mock_server().await;
|
||
let mcp = wiremock::MockServer::start().await;
|
||
wiremock::Mock::given(wiremock::matchers::method("POST"))
|
||
.and(wiremock::matchers::path("/mcp"))
|
||
.respond_with(|request: &wiremock::Request| {
|
||
let request: Value = serde_json::from_slice(&request.body).expect("MCP request");
|
||
let result = match request["method"].as_str() {
|
||
Some("initialize") => json!({
|
||
"protocolVersion":"2025-11-25", "capabilities":{"tools":{}},
|
||
"serverInfo":{"name":"sample-test","version":"1.0"}
|
||
}),
|
||
Some("tools/list") => json!({"tools":[{
|
||
"name":"probe", "description":"Record a test call", "inputSchema":{"type":"object","properties":{}},
|
||
"annotations":{"readOnlyHint":true}
|
||
}]}),
|
||
Some("tools/call") => json!({"content":[{"type":"text","text":"Recorded"}]}),
|
||
_ => return ResponseTemplate::new(202),
|
||
};
|
||
ResponseTemplate::new(200).set_body_json(json!({"jsonrpc":"2.0","id":request["id"],"result":result}))
|
||
})
|
||
.mount(&mcp).await;
|
||
let mock = responses::mount_sse_sequence(
|
||
&server,
|
||
vec![
|
||
responses::sse(vec![
|
||
responses::ev_function_call_with_namespace(
|
||
"parent-call",
|
||
"mcp__sample_test",
|
||
"probe",
|
||
"{}",
|
||
),
|
||
responses::ev_completed("parent-tool"),
|
||
]),
|
||
responses::sse(vec![
|
||
responses::ev_assistant_message("answer", "Done"),
|
||
responses::ev_completed("parent"),
|
||
]),
|
||
responses::sse(vec![
|
||
responses::ev_function_call_with_namespace(
|
||
"sample-call",
|
||
"mcp__sample_test",
|
||
"probe",
|
||
"{}",
|
||
),
|
||
responses::ev_completed("sample-tool"),
|
||
]),
|
||
responses::sse(vec![responses::ev_completed("prediction-fork")]),
|
||
],
|
||
)
|
||
.await;
|
||
let home = TempDir::new()?;
|
||
let extra_config = format!(
|
||
"[mcp_servers.sample_test]\nurl = \"{}/mcp\"\ndefault_tools_approval_mode = \"approve\"",
|
||
mcp.uri()
|
||
);
|
||
let mut client = client_with_effort_updates(
|
||
&home,
|
||
&server,
|
||
/*overrides*/ true,
|
||
EffortUpdates::Supported,
|
||
&extra_config,
|
||
)
|
||
.await?;
|
||
let (parent_id, turn_id) = ready_parent(
|
||
&mut client,
|
||
ThreadStartParams {
|
||
config: Some(std::collections::HashMap::from([(
|
||
"features.incremental_tools".to_string(),
|
||
json!(incremental_tools),
|
||
)])),
|
||
..Default::default()
|
||
},
|
||
)
|
||
.await?;
|
||
let response = sample(
|
||
&mut client,
|
||
ThreadSampleParams {
|
||
thread_id: parent_id.clone(),
|
||
expected_turn_id: turn_id,
|
||
prompt: "Call the tool again".to_string(),
|
||
disable_mcp,
|
||
..Default::default()
|
||
},
|
||
)
|
||
.await?;
|
||
assert_eq!(
|
||
(response.text, response.response_id),
|
||
(None, "sample-tool".to_string())
|
||
);
|
||
let fork_id = client
|
||
.send_thread_fork_request(ThreadForkParams {
|
||
thread_id: parent_id,
|
||
experimental_prediction_mode: true,
|
||
experimental_disable_prediction_mcp: disable_mcp,
|
||
ephemeral: true,
|
||
exclude_turns: true,
|
||
..Default::default()
|
||
})
|
||
.await?;
|
||
let fork: ThreadForkResponse = timeout(WAIT, client.read_response(fork_id)).await??;
|
||
client
|
||
.start_turn_and_wait_for_completion(TurnStartParams {
|
||
thread_id: fork.thread.id,
|
||
input: input("Call the tool again"),
|
||
..Default::default()
|
||
})
|
||
.await?;
|
||
let requests = mock.requests();
|
||
assert_eq!(requests.len(), 4);
|
||
let parent = comparable_request(requests[1].body_json());
|
||
let sample = comparable_request(requests[2].body_json());
|
||
let prefix = parent["input"].as_array().expect("parent input");
|
||
let sample_input = sample["input"].as_array().expect("sample input");
|
||
assert_eq!(&sample_input[..prefix.len()], prefix);
|
||
assert!(!sample_input[prefix.len()..].iter().any(|item| {
|
||
item["type"] == "additional_tools" || item.to_string().contains("no longer available")
|
||
}));
|
||
assert_eq!(sample, comparable_request(requests[3].body_json()));
|
||
assert_eq!(
|
||
requests[2].body_json()["tools"],
|
||
requests[1].body_json()["tools"]
|
||
);
|
||
let calls = mcp
|
||
.received_requests()
|
||
.await
|
||
.expect("MCP requests")
|
||
.iter()
|
||
.filter(|request| {
|
||
serde_json::from_slice::<Value>(&request.body)
|
||
.is_ok_and(|body| body["method"] == "tools/call")
|
||
})
|
||
.count();
|
||
// The parent executes the tool; the sample never sends a second MCP call.
|
||
assert_eq!(calls, 1);
|
||
Ok(())
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn sample_does_not_hide_stream_errors_after_tool_calls() -> Result<()> {
|
||
let server = responses::start_mock_server().await;
|
||
let mock = responses::mount_sse_sequence(&server, vec![
|
||
responses::sse(vec![responses::ev_completed("parent")]),
|
||
responses::sse(vec![
|
||
responses::ev_function_call("call", "view_image", "{}"),
|
||
json!({"type":"response.failed","response":{"id":"failed","error":{"code":"insufficient_quota","message":"quota exceeded"}}}),
|
||
]),
|
||
]).await;
|
||
let home = TempDir::new()?;
|
||
let mut client = client(&home, &server, /*overrides*/ true).await?;
|
||
let (parent, turn) = ready_parent(&mut client, ThreadStartParams::default()).await?;
|
||
let id = client
|
||
.send_raw_request(
|
||
"thread/sample",
|
||
Some(json!({
|
||
"threadId":parent, "expectedTurnId":turn, "prompt":"Call a tool"
|
||
})),
|
||
)
|
||
.await?;
|
||
let error = timeout(
|
||
WAIT,
|
||
client.read_stream_until_error_message(RequestId::Integer(id)),
|
||
)
|
||
.await??;
|
||
assert_eq!(error.error.code, -32603);
|
||
assert_eq!(mock.requests().len(), 2);
|
||
Ok(())
|
||
}
|
||
|
||
enum SampleStop {
|
||
Timeout,
|
||
Cancel,
|
||
}
|
||
|
||
#[test_case(SampleStop::Timeout; "timeout")]
|
||
#[test_case(SampleStop::Cancel; "cancel")]
|
||
#[tokio::test]
|
||
async fn sample_wait_keeps_parent_responsive_and_reclaims_child(stop: SampleStop) -> Result<()> {
|
||
let server = responses::start_mock_server().await;
|
||
responses::mount_sse_once(
|
||
&server,
|
||
responses::sse(vec![responses::ev_completed("parent")]),
|
||
)
|
||
.await;
|
||
let home = TempDir::new()?;
|
||
let mut client = client(&home, &server, /*overrides*/ true).await?;
|
||
let (parent_id, turn_id) = ready_parent(&mut client, ThreadStartParams::default()).await?;
|
||
let arrived = Arc::new(Notify::new());
|
||
let notify = Arc::clone(&arrived);
|
||
let mock = responses::mount_response_once_match(
|
||
&server,
|
||
move |_: &wiremock::Request| {
|
||
notify.notify_one();
|
||
true
|
||
},
|
||
ResponseTemplate::new(200)
|
||
.insert_header("content-type", "text/event-stream")
|
||
.set_body_string(responses::sse(vec![responses::ev_completed("delayed")]))
|
||
.set_delay(Duration::from_secs(60)),
|
||
)
|
||
.await;
|
||
let params = ThreadSampleParams {
|
||
thread_id: parent_id.clone(),
|
||
expected_turn_id: turn_id,
|
||
prompt: "Predict".to_string(),
|
||
..Default::default()
|
||
};
|
||
let sample_id = client
|
||
.send_raw_request("thread/sample", Some(serde_json::to_value(¶ms)?))
|
||
.await?;
|
||
timeout(WAIT, arrived.notified()).await?;
|
||
responses::mount_sse_once(
|
||
&server,
|
||
responses::sse(vec![responses::ev_completed("next-parent-turn")]),
|
||
)
|
||
.await;
|
||
timeout(
|
||
Duration::from_secs(10),
|
||
client.start_turn_and_wait_for_completion(TurnStartParams {
|
||
thread_id: parent_id.clone(),
|
||
input: input("Continue the parent while sampling"),
|
||
..Default::default()
|
||
}),
|
||
)
|
||
.await??;
|
||
let (expected_error, wait) = match stop {
|
||
SampleStop::Timeout => ("sample timed out", WAIT),
|
||
SampleStop::Cancel => {
|
||
let cancel_id = client
|
||
.send_raw_request(
|
||
"thread/sample/cancel",
|
||
Some(json!({"requestId": sample_id})),
|
||
)
|
||
.await?;
|
||
let response: ThreadSampleCancelResponse =
|
||
timeout(Duration::from_secs(10), client.read_response(cancel_id)).await??;
|
||
assert_eq!(response, ThreadSampleCancelResponse {});
|
||
("sample canceled", Duration::from_secs(10))
|
||
}
|
||
};
|
||
let error = timeout(
|
||
wait,
|
||
client.read_stream_until_error_message(RequestId::Integer(sample_id)),
|
||
)
|
||
.await??;
|
||
assert_eq!(error.error.message, expected_error);
|
||
if matches!(stop, SampleStop::Cancel) {
|
||
for request_id in [json!(sample_id), json!("unknown")] {
|
||
let cancel_id = client
|
||
.send_raw_request(
|
||
"thread/sample/cancel",
|
||
Some(json!({"requestId": request_id})),
|
||
)
|
||
.await?;
|
||
let response: ThreadSampleCancelResponse =
|
||
timeout(Duration::from_secs(10), client.read_response(cancel_id)).await??;
|
||
assert_eq!(response, ThreadSampleCancelResponse {});
|
||
}
|
||
}
|
||
assert_eq!(mock.requests().len(), 1);
|
||
timeout(WAIT, async {
|
||
loop {
|
||
let list_id = client
|
||
.send_raw_request("thread/loaded/list", Some(json!({})))
|
||
.await?;
|
||
let loaded: ThreadLoadedListResponse = client.read_response(list_id).await?;
|
||
if loaded.data == vec![parent_id.clone()] {
|
||
return Ok::<(), anyhow::Error>(());
|
||
}
|
||
}
|
||
})
|
||
.await??;
|
||
Ok(())
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn sample_rejects_invalid_sources_and_prompts() -> Result<()> {
|
||
let server = responses::start_mock_server().await;
|
||
let mock = responses::mount_sse_sequence(
|
||
&server,
|
||
vec![responses::sse(vec![responses::ev_completed("parent")])],
|
||
)
|
||
.await;
|
||
let home = TempDir::new()?;
|
||
let mut client = client(&home, &server, /*overrides*/ true).await?;
|
||
let fresh = client
|
||
.start_thread(ThreadStartParams::default())
|
||
.await?
|
||
.thread
|
||
.id;
|
||
let ephemeral = client
|
||
.start_thread(ThreadStartParams {
|
||
ephemeral: Some(true),
|
||
..Default::default()
|
||
})
|
||
.await?
|
||
.thread
|
||
.id;
|
||
let (parent, turn) = ready_parent(&mut client, ThreadStartParams::default()).await?;
|
||
for (id, expected_turn, prompt, expected_error) in [
|
||
(
|
||
codex_protocol::ThreadId::new().to_string(),
|
||
turn.clone(),
|
||
"Predict".to_string(),
|
||
"loaded source",
|
||
),
|
||
(fresh, turn.clone(), "Predict".to_string(), "rollout"),
|
||
(
|
||
ephemeral,
|
||
turn.clone(),
|
||
"Predict".to_string(),
|
||
"persistent root",
|
||
),
|
||
(
|
||
parent.clone(),
|
||
"stale".to_string(),
|
||
"Predict".to_string(),
|
||
"latest finished turn",
|
||
),
|
||
(parent.clone(), turn.clone(), String::new(), "prompt"),
|
||
(
|
||
parent.clone(),
|
||
turn.clone(),
|
||
"x".repeat(4_001),
|
||
"at most 4000 bytes",
|
||
),
|
||
] {
|
||
let id = client
|
||
.send_raw_request(
|
||
"thread/sample",
|
||
Some(json!({"threadId": id, "expectedTurnId": expected_turn, "prompt": prompt})),
|
||
)
|
||
.await?;
|
||
let error = timeout(
|
||
WAIT,
|
||
client.read_stream_until_error_message(RequestId::Integer(id)),
|
||
)
|
||
.await??;
|
||
assert!(
|
||
error.error.message.contains(expected_error),
|
||
"{}",
|
||
error.error.message
|
||
);
|
||
}
|
||
assert_eq!(mock.requests().len(), 1);
|
||
Ok(())
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn sample_uses_the_parent_catalog_after_a_default_provider_switch() -> Result<()> {
|
||
let alpha = wiremock::MockServer::start().await;
|
||
let beta = wiremock::MockServer::start().await;
|
||
let mut model = codex_models_manager::bundled_models_response()?
|
||
.models
|
||
.into_iter()
|
||
.find(|model| model.slug == "gpt-5.5")
|
||
.expect("bundled reasoning model");
|
||
for (provider, supports_updates) in [(&alpha, false), (&beta, true)] {
|
||
model.supports_reasoning_effort_updates = supports_updates;
|
||
wiremock::Mock::given(wiremock::matchers::method("GET"))
|
||
.and(wiremock::matchers::path("/v1/models"))
|
||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"models": [&model]})))
|
||
.mount(provider)
|
||
.await;
|
||
}
|
||
let replies = ["parent", "sample"].map(|id| {
|
||
responses::sse(vec![
|
||
responses::ev_assistant_message(id, "Ready."),
|
||
responses::ev_completed(id),
|
||
])
|
||
});
|
||
let alpha_requests = responses::mount_sse_sequence(&alpha, replies.to_vec()).await;
|
||
let beta_requests = responses::mount_sse_sequence(&beta, replies.to_vec()).await;
|
||
let home = TempDir::new()?;
|
||
let alpha_url = alpha.uri();
|
||
let beta_url = beta.uri();
|
||
MockResponsesConfig::new(&alpha_url)
|
||
.with_model_provider("alpha")
|
||
.with_provider_name("OpenAI")
|
||
.with_model(&model.slug)
|
||
.with_root_config("model_reasoning_effort = 'medium'")
|
||
.enable_feature(Feature::ApiKeyModelDiscovery)
|
||
.enable_feature(Feature::ReasoningEffortOverride)
|
||
.with_provider_config(&format!(
|
||
"env_key = 'TEST_CATALOG_KEY'\nmodel_catalog_url = '{alpha_url}/v1/models'"
|
||
))
|
||
.with_extra_config(&format!(
|
||
r#"[model_providers.beta]
|
||
name = "OpenAI"
|
||
base_url = "{beta_url}/v1"
|
||
model_catalog_url = "{beta_url}/v1/models"
|
||
env_key = "TEST_CATALOG_KEY""#
|
||
))
|
||
.write(home.path())?;
|
||
let mut client = TestAppServer::builder()
|
||
.with_codex_home(home.path())
|
||
.with_env_overrides(&[("TEST_CATALOG_KEY", Some("test-key"))])
|
||
.build_initialized_with_timeout(WAIT)
|
||
.await?;
|
||
let alpha_parent = ready_parent(&mut client, ThreadStartParams::default()).await?;
|
||
let id = client
|
||
.send_raw_request(
|
||
"modelProvider/default/set",
|
||
Some(json!({"providerId": "beta"})),
|
||
)
|
||
.await?;
|
||
let write: codex_app_server_protocol::ConfigWriteResponse = client.read_response(id).await?;
|
||
assert_eq!(write.status, codex_app_server_protocol::WriteStatus::Ok);
|
||
let beta_parent = ready_parent(&mut client, ThreadStartParams::default()).await?;
|
||
for ((thread_id, expected_turn_id), requests, supports_updates) in [
|
||
(alpha_parent, alpha_requests, false),
|
||
(beta_parent, beta_requests, true),
|
||
] {
|
||
sample(
|
||
&mut client,
|
||
ThreadSampleParams {
|
||
thread_id,
|
||
expected_turn_id,
|
||
prompt: "Predict".into(),
|
||
effort: Some(ReasoningEffort::High),
|
||
..Default::default()
|
||
},
|
||
)
|
||
.await?;
|
||
let requests = requests.requests();
|
||
assert_eq!(requests.len(), 2);
|
||
let updates = requests[1].inputs_of_type("configuration_update");
|
||
assert_eq!(
|
||
updates
|
||
.iter()
|
||
.any(|item| item["reasoning"]["effort"] == "high"),
|
||
supports_updates
|
||
);
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
#[test_case(false; "unsupported model")]
|
||
#[test_case(true; "unsupported provider")]
|
||
#[tokio::test]
|
||
async fn unsupported_effort_updates_keep_inherited_effort(
|
||
unsupported_provider: bool,
|
||
) -> Result<()> {
|
||
let server = responses::start_mock_server().await;
|
||
let mock = responses::mount_sse_sequence(
|
||
&server,
|
||
["parent", "sample"]
|
||
.map(|id| {
|
||
responses::sse(vec![
|
||
responses::ev_assistant_message(id, "Completed without usage"),
|
||
json!({"type":"response.completed","response":{"id":id}}),
|
||
])
|
||
})
|
||
.to_vec(),
|
||
)
|
||
.await;
|
||
let home = TempDir::new()?;
|
||
let effort_updates = if unsupported_provider {
|
||
EffortUpdates::Supported
|
||
} else {
|
||
EffortUpdates::Unsupported
|
||
};
|
||
let mut client =
|
||
client_with_effort_updates(&home, &server, /*overrides*/ true, effort_updates, "").await?;
|
||
let overrides = if unsupported_provider {
|
||
json!({"model_providers.mock_provider.name":"Another provider"})
|
||
} else {
|
||
json!({})
|
||
};
|
||
let (parent_id, turn_id) = ready_parent(
|
||
&mut client,
|
||
ThreadStartParams {
|
||
config: Some(serde_json::from_value(overrides)?),
|
||
..Default::default()
|
||
},
|
||
)
|
||
.await?;
|
||
let response = sample(
|
||
&mut client,
|
||
ThreadSampleParams {
|
||
thread_id: parent_id,
|
||
expected_turn_id: turn_id.clone(),
|
||
prompt: "Predict".to_string(),
|
||
effort: Some(ReasoningEffort::Medium),
|
||
..Default::default()
|
||
},
|
||
)
|
||
.await?;
|
||
assert_eq!(
|
||
response,
|
||
ThreadSampleResponse {
|
||
source_turn_id: turn_id,
|
||
response_id: "sample".to_string(),
|
||
text: Some("Completed without usage".to_string()),
|
||
input_tokens: None,
|
||
cached_input_tokens: None,
|
||
output_tokens: None,
|
||
}
|
||
);
|
||
let requests = mock.requests();
|
||
assert_eq!(
|
||
requests[1].body_json()["reasoning"],
|
||
requests[0].body_json()["reasoning"]
|
||
);
|
||
assert!(
|
||
!requests[1]
|
||
.input()
|
||
.iter()
|
||
.any(|item| item["type"] == "configuration_update")
|
||
);
|
||
Ok(())
|
||
}
|