mirror of
https://github.com/openai/codex.git
synced 2026-10-20 17:38:45 +00:00
Add experimental thread/sample requests to the app server (#55339)
## What changed Add `thread/sample` to collect one model response on an ephemeral fork that inherits a loaded thread's context and configuration. Return the source turn ID, response ID, assistant text, and available token usage without modifying the source conversation. Require an idle, persistent root thread and an `expectedTurnId` matching its latest completed turn. Reject results if the source history changes during sampling. Keep the parent responsive while sampling, hide sample forks from `thread/loaded/list`, and reclaim them after completion or timeout. Intercept responses before tool execution or Stop hooks can run, returning null text for tool calls or responses without text. Apply optional reasoning-effort updates only when the provider and model support cache-preserving updates. ## Testing Add integration coverage for inherited context and reasoning effort, source preservation, tool and Stop hook suppression, invalid requests, stream errors, parent responsiveness, and timeout cleanup. GitOrigin-RevId: 08f4867e53ec7f6c8f4428d5542eeea4e6dadd0b
This commit is contained in:
Generated
+1
@@ -2168,6 +2168,7 @@ dependencies = [
|
||||
"futures",
|
||||
"hmac 0.12.1",
|
||||
"http 1.4.0",
|
||||
"insta",
|
||||
"keyring",
|
||||
"libc",
|
||||
"opentelemetry",
|
||||
|
||||
BIN
Binary file not shown.
@@ -567,6 +567,13 @@ client_request_definitions! {
|
||||
serialization: thread_or_path(params.thread_id, params.path),
|
||||
response: v2::ThreadForkResponse,
|
||||
},
|
||||
#[experimental("thread/sample")]
|
||||
ThreadSample => "thread/sample" {
|
||||
params: v2::ThreadSampleParams,
|
||||
// Waiting for a sample must not block requests on its parent thread.
|
||||
serialization: None,
|
||||
response: v2::ThreadSampleResponse,
|
||||
},
|
||||
ThreadArchive => "thread/archive" {
|
||||
params: v2::ThreadArchiveParams,
|
||||
serialization: thread_id(params.thread_id),
|
||||
|
||||
@@ -35,6 +35,7 @@ mod thread;
|
||||
mod thread_attachment;
|
||||
mod thread_data;
|
||||
mod thread_prediction;
|
||||
mod thread_sample;
|
||||
mod thread_usage;
|
||||
mod turn;
|
||||
mod user_verification;
|
||||
@@ -76,6 +77,7 @@ pub use thread::*;
|
||||
pub use thread_attachment::*;
|
||||
pub use thread_data::*;
|
||||
pub use thread_prediction::*;
|
||||
pub use thread_sample::*;
|
||||
pub use thread_usage::*;
|
||||
pub use turn::*;
|
||||
pub use user_verification::*;
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
//! One model response on an ephemeral fork that inherits the source's configuration and context.
|
||||
|
||||
use crate::JsonSchema;
|
||||
use crate::TS;
|
||||
use codex_protocol::openai_models::ReasoningEffort;
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug, Default, Clone, PartialEq, JsonSchema, TS)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[ts(export_to = "v2/")]
|
||||
pub struct ThreadSampleParams {
|
||||
pub thread_id: String,
|
||||
/// The loaded source's latest completed turn. Stale or active sources are rejected.
|
||||
pub expected_turn_id: String,
|
||||
pub prompt: String,
|
||||
/// Applied only when the provider and model support cache-preserving reasoning-effort updates.
|
||||
#[ts(optional = nullable)]
|
||||
pub effort: Option<ReasoningEffort>,
|
||||
#[ts(optional = nullable)]
|
||||
pub turn_trigger: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[ts(export_to = "v2/")]
|
||||
pub struct ThreadSampleResponse {
|
||||
pub source_turn_id: String,
|
||||
pub response_id: String,
|
||||
/// Assistant text from the response, or null if it attempted a tool call or completed without text.
|
||||
pub text: Option<String>,
|
||||
pub input_tokens: Option<i64>,
|
||||
pub cached_input_tokens: Option<i64>,
|
||||
pub output_tokens: Option<i64>,
|
||||
}
|
||||
@@ -144,6 +144,7 @@ core_test_support = { workspace = true }
|
||||
flate2 = { workspace = true }
|
||||
hmac = { workspace = true }
|
||||
http = { workspace = true }
|
||||
insta = { workspace = true }
|
||||
keyring = { workspace = true }
|
||||
opentelemetry = { workspace = true }
|
||||
opentelemetry_sdk = { workspace = true }
|
||||
|
||||
@@ -32,6 +32,7 @@ use crate::thread_state::ThreadListenerCommand;
|
||||
use crate::thread_state::ThreadStateManager;
|
||||
|
||||
pub(crate) struct ThreadExtensionDependencies {
|
||||
pub(crate) sample_responses: Arc<crate::thread_sample::SampleResponseRegistry>,
|
||||
pub(crate) event_sink: Arc<dyn ExtensionEventSink>,
|
||||
pub(crate) auth_manager: Arc<AuthManager>,
|
||||
pub(crate) state_db: Option<StateDbHandle>,
|
||||
@@ -52,6 +53,7 @@ pub(crate) fn thread_extensions(
|
||||
dependencies: ThreadExtensionDependencies,
|
||||
) -> Arc<ExtensionRegistry<Config>> {
|
||||
let ThreadExtensionDependencies {
|
||||
sample_responses,
|
||||
event_sink,
|
||||
auth_manager,
|
||||
state_db,
|
||||
@@ -125,6 +127,7 @@ pub(crate) fn thread_extensions(
|
||||
.enabled(codex_features::Feature::SkillSearch),
|
||||
},
|
||||
);
|
||||
builder.model_request_contributor(sample_responses);
|
||||
Arc::new(builder.build())
|
||||
}
|
||||
|
||||
|
||||
@@ -142,6 +142,7 @@ mod request_processors;
|
||||
mod request_serialization;
|
||||
mod server_request_error;
|
||||
mod skills_watcher;
|
||||
mod thread_sample;
|
||||
mod thread_state;
|
||||
mod thread_status;
|
||||
mod transport;
|
||||
|
||||
@@ -522,6 +522,9 @@ enabled = false
|
||||
thread_extensions(
|
||||
&good_config,
|
||||
ThreadExtensionDependencies {
|
||||
sample_responses: Arc::new(
|
||||
crate::thread_sample::SampleResponseRegistry::default(),
|
||||
),
|
||||
event_sink: Arc::new(NoopExtensionEventSink),
|
||||
auth_manager: auth_manager.clone(),
|
||||
state_db: Some(state_db.clone()),
|
||||
|
||||
@@ -329,6 +329,7 @@ impl MessageProcessor {
|
||||
let extension_event_sink =
|
||||
app_server_extension_event_sink(outgoing.clone(), thread_state_manager.clone());
|
||||
let mut queue_service = None;
|
||||
let sample_responses = Arc::new(crate::thread_sample::SampleResponseRegistry::default());
|
||||
let thread_manager = Arc::new_cyclic(|thread_manager| {
|
||||
queue_service = queue_store.map(|queue| {
|
||||
Arc::new(QueuedItemService::new(
|
||||
@@ -347,6 +348,7 @@ impl MessageProcessor {
|
||||
thread_extensions(
|
||||
config.as_ref(),
|
||||
ThreadExtensionDependencies {
|
||||
sample_responses: Arc::clone(&sample_responses),
|
||||
event_sink: Arc::clone(&extension_event_sink),
|
||||
auth_manager: auth_manager.clone(),
|
||||
state_db: state_db.clone(),
|
||||
@@ -535,6 +537,7 @@ impl MessageProcessor {
|
||||
Arc::clone(&skills_watcher),
|
||||
turn_cost_worker.as_ref().map(TurnCostWorker::handle),
|
||||
config_warnings,
|
||||
sample_responses,
|
||||
);
|
||||
let turn_processor = TurnRequestProcessor::new(
|
||||
auth_manager,
|
||||
@@ -1005,7 +1008,8 @@ impl MessageProcessor {
|
||||
| ClientRequest::TurnSettingsUpdate { .. }
|
||||
| ClientRequest::ThreadDelete { .. }
|
||||
| ClientRequest::ThreadArchive { .. } => (Some(self.turn_admission.admit()?), false),
|
||||
ClientRequest::TurnStart { .. }
|
||||
ClientRequest::ThreadSample { .. }
|
||||
| ClientRequest::TurnStart { .. }
|
||||
| ClientRequest::TurnSteer { .. }
|
||||
| ClientRequest::ReviewStart { .. }
|
||||
| ClientRequest::ThreadCompactStart { .. }
|
||||
@@ -1335,6 +1339,11 @@ impl MessageProcessor {
|
||||
)
|
||||
.await
|
||||
}
|
||||
ClientRequest::ThreadSample { params, .. } => self
|
||||
.thread_processor
|
||||
.thread_sample(params, Arc::clone(&session.rpc_gate), request_context)
|
||||
.await
|
||||
.map(|response| Some(response.into())),
|
||||
ClientRequest::ThreadArchive { params, .. } => {
|
||||
self.thread_processor
|
||||
.thread_archive(request_id.clone(), params)
|
||||
|
||||
@@ -11,7 +11,11 @@ use super::thread_fork_goal::inherit_thread_goal_snapshot;
|
||||
use super::thread_input::can_accept_direct_input;
|
||||
use super::thread_input::ensure_direct_input_allowed;
|
||||
use super::*;
|
||||
use crate::connection_rpc_gate::ConnectionRpcGate;
|
||||
use crate::error_code::method_not_found;
|
||||
use crate::thread_sample::SampleTask;
|
||||
use crate::thread_sample::SampleThread;
|
||||
use crate::thread_sample::generate_sample;
|
||||
use codex_app_server_protocol::SelectedCapabilityRoot;
|
||||
use codex_app_server_protocol::ThreadHistoryMode as ApiThreadHistoryMode;
|
||||
use codex_app_server_protocol::ThreadItemsListAnchor;
|
||||
@@ -19,6 +23,8 @@ use codex_app_server_protocol::ThreadItemsListCursor;
|
||||
use codex_app_server_protocol::ThreadRevertParams;
|
||||
use codex_app_server_protocol::ThreadRevertResponse;
|
||||
use codex_app_server_protocol::ThreadRevertedNotification;
|
||||
use codex_app_server_protocol::ThreadSampleParams;
|
||||
use codex_app_server_protocol::ThreadSampleResponse;
|
||||
use codex_app_server_protocol::ThreadSection;
|
||||
use codex_app_server_protocol::ThreadSectionAppearance;
|
||||
use codex_app_server_protocol::ThreadSectionMoveParams;
|
||||
@@ -463,6 +469,7 @@ pub(crate) struct ThreadRequestProcessor {
|
||||
pub(super) state_db: Option<StateDbHandle>,
|
||||
pub(super) log_db: Option<LogDbLayer>,
|
||||
pub(super) background_tasks: TaskTracker,
|
||||
pub(super) sample_responses: Arc<crate::thread_sample::SampleResponseRegistry>,
|
||||
pub(super) skills_watcher: Arc<SkillsWatcher>,
|
||||
pub(super) turn_cost_worker: Option<crate::turn_cost_worker::TurnCostWorkerHandle>,
|
||||
pub(super) initial_config_warnings: Arc<Vec<ConfigWarningNotification>>,
|
||||
@@ -505,6 +512,7 @@ impl ThreadRequestProcessor {
|
||||
skills_watcher: Arc<SkillsWatcher>,
|
||||
turn_cost_worker: Option<crate::turn_cost_worker::TurnCostWorkerHandle>,
|
||||
initial_config_warnings: Vec<ConfigWarningNotification>,
|
||||
sample_responses: Arc<crate::thread_sample::SampleResponseRegistry>,
|
||||
) -> Self {
|
||||
Self {
|
||||
auth_manager,
|
||||
@@ -522,6 +530,7 @@ impl ThreadRequestProcessor {
|
||||
state_db,
|
||||
log_db,
|
||||
background_tasks: TaskTracker::new(),
|
||||
sample_responses,
|
||||
skills_watcher,
|
||||
turn_cost_worker,
|
||||
initial_config_warnings: Arc::new(initial_config_warnings),
|
||||
@@ -603,6 +612,38 @@ impl ThreadRequestProcessor {
|
||||
.map(|()| None)
|
||||
}
|
||||
|
||||
pub(crate) async fn thread_sample(
|
||||
&self,
|
||||
params: ThreadSampleParams,
|
||||
rpc_gate: Arc<ConnectionRpcGate>,
|
||||
context: RequestContext,
|
||||
) -> Result<ThreadSampleResponse, JSONRPCErrorError> {
|
||||
let id = ThreadId::from_string(¶ms.thread_id)
|
||||
.map_err(|error| invalid_request(error.to_string()))?;
|
||||
let parent = self
|
||||
.thread_manager
|
||||
.get_thread(id)
|
||||
.await
|
||||
.map_err(|_| invalid_request("thread/sample requires a loaded source"))?;
|
||||
self.config_manager
|
||||
.check_thread_model_provider(parent.config().await.as_ref())
|
||||
.await
|
||||
.map_err(|error| config_load_error(&error))?;
|
||||
generate_sample(
|
||||
SampleTask {
|
||||
manager: Arc::clone(&self.thread_manager),
|
||||
store: Arc::clone(&self.thread_store),
|
||||
parent,
|
||||
samples: Arc::clone(&self.sample_responses),
|
||||
params,
|
||||
trace: context.request_trace(),
|
||||
},
|
||||
&self.background_tasks,
|
||||
rpc_gate,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn thread_revert(
|
||||
&self,
|
||||
request_id: ConnectionRequestId,
|
||||
@@ -2887,13 +2928,17 @@ impl ThreadRequestProcessor {
|
||||
params: ThreadLoadedListParams,
|
||||
) -> Result<ThreadLoadedListResponse, JSONRPCErrorError> {
|
||||
let ThreadLoadedListParams { cursor, limit } = params;
|
||||
let mut data: Vec<String> = self
|
||||
.thread_manager
|
||||
.list_thread_ids()
|
||||
.await
|
||||
.into_iter()
|
||||
.map(|thread_id| thread_id.to_string())
|
||||
.collect();
|
||||
let mut data = Vec::new();
|
||||
for thread_id in self.thread_manager.list_thread_ids().await {
|
||||
if let Ok(thread) = self.thread_manager.get_thread(thread_id).await
|
||||
&& thread
|
||||
.thread_extension_data()
|
||||
.get::<SampleThread>()
|
||||
.is_none()
|
||||
{
|
||||
data.push(thread_id.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
if data.is_empty() {
|
||||
return Ok(ThreadLoadedListResponse {
|
||||
|
||||
@@ -0,0 +1,405 @@
|
||||
//! Captures one model response on an inherited ephemeral fork before Core can act on it.
|
||||
//! Tool calls produce a null sample; an owned task retains startup and cleanup.
|
||||
|
||||
use crate::connection_rpc_gate::ConnectionRpcGate;
|
||||
use crate::error_code::internal_error;
|
||||
use crate::error_code::invalid_request;
|
||||
use crate::request_processors::build_legacy_api_turns_from_rollout_items;
|
||||
use codex_app_server_protocol::JSONRPCErrorError;
|
||||
use codex_app_server_protocol::ThreadSampleParams;
|
||||
use codex_app_server_protocol::ThreadSampleResponse;
|
||||
use codex_app_server_protocol::TurnStatus;
|
||||
use codex_core::CodexThread;
|
||||
use codex_core::ForkSnapshot;
|
||||
use codex_core::ThreadManager;
|
||||
use codex_extension_api::ModelRequestContributor;
|
||||
use codex_extension_api::ModelRequestInput;
|
||||
use codex_extension_api::ModelRequestKind;
|
||||
use codex_extension_api::ModelResponseInterceptor;
|
||||
use codex_extension_api::ModelResponseStream;
|
||||
use codex_extension_api::ResponseEvent;
|
||||
use codex_features::Feature;
|
||||
use codex_otel::auth_storage::AuthStorageOriginator;
|
||||
use codex_protocol::models::ContentItem;
|
||||
use codex_protocol::models::MessagePhase;
|
||||
use codex_protocol::models::ResponseItem;
|
||||
use codex_protocol::protocol::AgentStatus;
|
||||
use codex_protocol::protocol::EventMsg;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use codex_protocol::protocol::ThreadHistoryMode;
|
||||
use codex_protocol::protocol::ThreadSettingsOverrides;
|
||||
use codex_protocol::protocol::ThreadSource;
|
||||
use codex_protocol::protocol::W3cTraceContext;
|
||||
use codex_protocol::turn_input::StartIfIdleSubmission;
|
||||
use codex_protocol::turn_input::TurnInputRequest;
|
||||
use codex_protocol::turn_input::TurnStartOptions;
|
||||
use codex_protocol::user_input::UserInput;
|
||||
use codex_rollout::InitialHistory;
|
||||
use codex_rollout::ResumedHistory;
|
||||
use codex_thread_store::ListTurnsParams;
|
||||
use codex_thread_store::LoadThreadHistoryParams;
|
||||
use codex_thread_store::SortDirection;
|
||||
use codex_thread_store::StoredTurnItemsView;
|
||||
use codex_thread_store::StoredTurnStatus;
|
||||
use codex_thread_store::ThreadStore;
|
||||
use futures::StreamExt;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio::time::Instant;
|
||||
use tokio_util::task::TaskTracker;
|
||||
|
||||
const SAMPLE_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
// Bytes conservatively bound tokens, including for byte-fallback tokenizers.
|
||||
const MAX_PROMPT_BYTES: usize = 4_000;
|
||||
const MAX_OUTPUT_BYTES: usize = 64 * 1024;
|
||||
type SampleResult = Result<ThreadSampleResponse, JSONRPCErrorError>;
|
||||
type SampleSender = Arc<Mutex<Option<oneshot::Sender<SampleResult>>>>;
|
||||
|
||||
/// Marks private sample children from startup through cleanup.
|
||||
pub(crate) struct SampleThread;
|
||||
|
||||
pub(crate) struct SampleTask {
|
||||
pub(crate) manager: Arc<ThreadManager>,
|
||||
pub(crate) store: Arc<dyn ThreadStore>,
|
||||
pub(crate) parent: Arc<CodexThread>,
|
||||
pub(crate) samples: Arc<SampleResponseRegistry>,
|
||||
pub(crate) params: ThreadSampleParams,
|
||||
pub(crate) trace: Option<W3cTraceContext>,
|
||||
}
|
||||
|
||||
pub(crate) async fn generate_sample(
|
||||
task: SampleTask,
|
||||
tasks: &TaskTracker,
|
||||
rpc_gate: Arc<ConnectionRpcGate>,
|
||||
) -> SampleResult {
|
||||
let deadline = Instant::now() + SAMPLE_TIMEOUT;
|
||||
if task.params.expected_turn_id.trim().is_empty()
|
||||
|| task.params.prompt.trim().is_empty()
|
||||
|| task.params.prompt.len() > MAX_PROMPT_BYTES
|
||||
{
|
||||
return Err(invalid_request(format!(
|
||||
"thread/sample requires expectedTurnId and a prompt of at most {MAX_PROMPT_BYTES} bytes",
|
||||
)));
|
||||
}
|
||||
let (sender, receiver) = oneshot::channel();
|
||||
let originator = AuthStorageOriginator::current();
|
||||
tasks.spawn(async move {
|
||||
rpc_gate
|
||||
.run(originator.scope(async move {
|
||||
let result = Box::pin(task.run(deadline)).await;
|
||||
let _ = sender.send(result);
|
||||
}))
|
||||
.await;
|
||||
});
|
||||
tokio::time::timeout_at(deadline, receiver)
|
||||
.await
|
||||
.map_err(|_| internal_error("sample timed out"))?
|
||||
.map_err(|_| internal_error("sample task stopped"))?
|
||||
}
|
||||
|
||||
impl SampleTask {
|
||||
async fn run(self, deadline: Instant) -> SampleResult {
|
||||
let Self {
|
||||
manager,
|
||||
store,
|
||||
parent,
|
||||
samples,
|
||||
params,
|
||||
trace,
|
||||
} = self;
|
||||
let id = parent.startup_metadata().thread_id;
|
||||
if parent.agent_status().await == AgentStatus::Running {
|
||||
return Err(invalid_request("sample source must be idle"));
|
||||
}
|
||||
let snapshot = parent.config_snapshot().await;
|
||||
if snapshot.ephemeral || snapshot.session_source.is_non_root_agent() {
|
||||
return Err(invalid_request(
|
||||
"sample source must be a persistent root thread",
|
||||
));
|
||||
}
|
||||
let revision = parent
|
||||
.conversation_history_snapshot()
|
||||
.await
|
||||
.history_version();
|
||||
parent
|
||||
.flush_rollout()
|
||||
.await
|
||||
.map_err(|error| internal_error(error.to_string()))?;
|
||||
let history = store
|
||||
.load_latest_model_context(LoadThreadHistoryParams {
|
||||
thread_id: id,
|
||||
include_archived: false,
|
||||
})
|
||||
.await
|
||||
.map_err(|error| invalid_request(error.to_string()))?;
|
||||
let matches_source = match snapshot.history_mode {
|
||||
ThreadHistoryMode::Legacy => build_legacy_api_turns_from_rollout_items(&history.items)
|
||||
.last()
|
||||
.is_some_and(|turn| {
|
||||
turn.id == params.expected_turn_id && turn.status == TurnStatus::Completed
|
||||
}),
|
||||
ThreadHistoryMode::Paginated => store
|
||||
.list_turns(ListTurnsParams {
|
||||
thread_id: id,
|
||||
include_archived: false,
|
||||
cursor: None,
|
||||
page_size: 1,
|
||||
sort_direction: SortDirection::Desc,
|
||||
items_view: StoredTurnItemsView::NotLoaded,
|
||||
})
|
||||
.await
|
||||
.map_err(|error| invalid_request(error.to_string()))?
|
||||
.turns
|
||||
.first()
|
||||
.is_some_and(|turn| {
|
||||
turn.turn_id == params.expected_turn_id
|
||||
&& turn.status == StoredTurnStatus::Completed
|
||||
}),
|
||||
};
|
||||
if !matches_source {
|
||||
return Err(invalid_request(
|
||||
"sample source turn is no longer the latest completed turn",
|
||||
));
|
||||
}
|
||||
let (mut options, settings) = Box::pin(async {
|
||||
manager
|
||||
.fork_options_from_parent(id)
|
||||
.await
|
||||
.map(|(options, settings)| (Box::new(options), Box::new(settings)))
|
||||
})
|
||||
.await
|
||||
.map_err(|error| internal_error(error.to_string()))?;
|
||||
let model = manager
|
||||
.get_models_manager()
|
||||
.get_model_info(
|
||||
options.config.model.as_deref().unwrap_or(&snapshot.model),
|
||||
&options.config.to_models_manager_config(),
|
||||
)
|
||||
.await;
|
||||
options.config.ephemeral = true;
|
||||
options
|
||||
.config
|
||||
.features
|
||||
.enable(Feature::ReasoningEffortOverride)
|
||||
.map_err(|error| invalid_request(error.to_string()))?;
|
||||
options
|
||||
.thread_extension_init
|
||||
.insert(codex_core::compact::AutoCompactionDisabled);
|
||||
options.thread_extension_init.insert(SampleThread);
|
||||
let effort = params.effort.filter(|_| {
|
||||
model.supports_reasoning_effort_updates && options.config.model_provider.is_openai()
|
||||
});
|
||||
options.session_source = Some(SessionSource::Custom("thread_sample".to_string()));
|
||||
options.thread_source = Some(ThreadSource::Feature("thread_sample".to_string()));
|
||||
options.parent_trace = trace.clone();
|
||||
if Instant::now() >= deadline {
|
||||
return Err(internal_error("sample timed out"));
|
||||
}
|
||||
// Do not cancel startup by dropping it: the owner must reclaim a partially initialized fork.
|
||||
let fork = Box::pin(manager.fork_thread_from_history(
|
||||
ForkSnapshot::Interrupted,
|
||||
*options,
|
||||
InitialHistory::Resumed(ResumedHistory {
|
||||
history_revision: history.revision,
|
||||
conversation_id: id,
|
||||
history: Arc::new(history.items),
|
||||
rollout_path: None,
|
||||
}),
|
||||
))
|
||||
.await
|
||||
.map_err(|error| internal_error(format!("sample fork failed: {error}")))?;
|
||||
let (sender, mut receiver) = oneshot::channel();
|
||||
samples
|
||||
.captures
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.insert(
|
||||
fork.thread_id.to_string(),
|
||||
Arc::new(Mutex::new(Some(sender))),
|
||||
);
|
||||
let sample = async {
|
||||
Box::pin(fork.thread.restore_thread_settings(*settings))
|
||||
.await
|
||||
.map_err(|error| internal_error(error.to_string()))?;
|
||||
let input = TurnInputRequest::user_input(vec![UserInput::Text {
|
||||
text: params.prompt,
|
||||
text_elements: Vec::new(),
|
||||
}])
|
||||
.with_thread_settings(ThreadSettingsOverrides {
|
||||
effort: effort.map(Some),
|
||||
..Default::default()
|
||||
})
|
||||
.on_start(TurnStartOptions {
|
||||
turn_trigger: params.turn_trigger,
|
||||
..Default::default()
|
||||
})
|
||||
.with_trace(trace);
|
||||
if !matches!(
|
||||
fork.thread
|
||||
.start_turn_if_idle(input)
|
||||
.await
|
||||
.map_err(|error| internal_error(error.to_string()))?,
|
||||
StartIfIdleSubmission::Started { .. }
|
||||
) {
|
||||
return Err(internal_error("sample turn was not admitted"));
|
||||
}
|
||||
loop {
|
||||
let event = tokio::select! {
|
||||
biased;
|
||||
result = &mut receiver => {
|
||||
return result.map_err(|_| internal_error("sample request stopped"))?;
|
||||
}
|
||||
event = fork.thread.next_event() => {
|
||||
event.map_err(|error| internal_error(error.to_string()))?
|
||||
}
|
||||
};
|
||||
match event.msg {
|
||||
EventMsg::TurnComplete(turn) => {
|
||||
return Err(internal_error(turn.error.map_or_else(
|
||||
|| "sample completed without a response".to_string(),
|
||||
|error| error.message,
|
||||
)));
|
||||
}
|
||||
EventMsg::TurnAborted(_) | EventMsg::ShutdownComplete => {
|
||||
return Err(internal_error("sample turn stopped"));
|
||||
}
|
||||
EventMsg::ExecApprovalRequest(_)
|
||||
| EventMsg::ApplyPatchApprovalRequest(_)
|
||||
| EventMsg::RequestPermissions(_)
|
||||
| EventMsg::RequestUserInput(_)
|
||||
| EventMsg::DynamicToolCallRequest(_)
|
||||
| EventMsg::ElicitationRequest(_) => {
|
||||
return Err(internal_error("sample turn requires client interaction"));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
};
|
||||
let result = tokio::select! {
|
||||
biased;
|
||||
_ = tokio::time::sleep_until(deadline) => Err(internal_error("sample timed out")),
|
||||
result = sample => result,
|
||||
};
|
||||
let cleanup = fork.thread.shutdown_and_wait().await;
|
||||
manager
|
||||
.remove_thread_if_matches(&fork.thread_id, &fork.thread)
|
||||
.await;
|
||||
samples
|
||||
.captures
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.remove(&fork.thread_id.to_string());
|
||||
let mut response = result?;
|
||||
cleanup.map_err(|error| internal_error(error.to_string()))?;
|
||||
if parent
|
||||
.conversation_history_snapshot()
|
||||
.await
|
||||
.history_version()
|
||||
!= revision
|
||||
{
|
||||
return Err(invalid_request("sample source changed"));
|
||||
}
|
||||
response.source_turn_id = params.expected_turn_id;
|
||||
Ok(response)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub(crate) struct SampleResponseRegistry {
|
||||
captures: Mutex<HashMap<String, SampleSender>>,
|
||||
}
|
||||
|
||||
impl ModelRequestContributor for SampleResponseRegistry {
|
||||
fn request(&self, input: ModelRequestInput<'_>) -> Option<Box<dyn ModelResponseInterceptor>> {
|
||||
if input.kind != ModelRequestKind::Generation {
|
||||
return None;
|
||||
}
|
||||
self.captures
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.get(input.thread_id)
|
||||
.cloned()
|
||||
.map(|sender| {
|
||||
Box::new(SampleInterceptor { sender }) as Box<dyn ModelResponseInterceptor>
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct SampleInterceptor {
|
||||
sender: SampleSender,
|
||||
}
|
||||
|
||||
impl ModelResponseInterceptor for SampleInterceptor {
|
||||
fn intercept(self: Box<Self>, upstream: ModelResponseStream) -> ModelResponseStream {
|
||||
Box::pin(futures::stream::once(async move {
|
||||
let result = collect_response(upstream).await;
|
||||
if let Some(sender) = self
|
||||
.sender
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.take()
|
||||
{
|
||||
let _ = sender.send(result);
|
||||
}
|
||||
// Completion permits Stop hooks; EOF triggers stream recovery. Wait for owner cleanup.
|
||||
std::future::pending().await
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
async fn collect_response(mut upstream: ModelResponseStream) -> SampleResult {
|
||||
let mut text = String::new();
|
||||
let mut text_only = true;
|
||||
while let Some(event) = upstream.next().await {
|
||||
match event.map_err(|error| internal_error(format!("sample failed: {error}")))? {
|
||||
ResponseEvent::OutputItemDone(ResponseItem::Message {
|
||||
phase: Some(MessagePhase::Commentary),
|
||||
..
|
||||
}) => {}
|
||||
ResponseEvent::OutputItemDone(ResponseItem::Message { role, content, .. })
|
||||
if role == "assistant" =>
|
||||
{
|
||||
for item in content {
|
||||
match item {
|
||||
ContentItem::OutputText { text: part, .. } => {
|
||||
if text.len().saturating_add(part.len()) > MAX_OUTPUT_BYTES {
|
||||
return Err(internal_error("sample output exceeds limit"));
|
||||
}
|
||||
text.push_str(&part);
|
||||
}
|
||||
_ => text_only = false,
|
||||
}
|
||||
}
|
||||
}
|
||||
ResponseEvent::OutputItemDone(ResponseItem::Reasoning { .. }) => {}
|
||||
ResponseEvent::OutputItemDone(_)
|
||||
| ResponseEvent::OutputItemAdded(
|
||||
ResponseItem::FunctionCall { .. }
|
||||
| ResponseItem::CustomToolCall { .. }
|
||||
| ResponseItem::LocalShellCall { .. }
|
||||
| ResponseItem::ToolSearchCall { .. },
|
||||
) => text_only = false,
|
||||
ResponseEvent::Completed {
|
||||
response_id,
|
||||
token_usage,
|
||||
..
|
||||
} => {
|
||||
return Ok(ThreadSampleResponse {
|
||||
source_turn_id: String::new(),
|
||||
response_id,
|
||||
text: (text_only && !text.trim().is_empty()).then_some(text),
|
||||
input_tokens: token_usage.as_ref().map(|usage| usage.input_tokens),
|
||||
cached_input_tokens: token_usage
|
||||
.as_ref()
|
||||
.map(|usage| usage.cached_input_tokens),
|
||||
output_tokens: token_usage.as_ref().map(|usage| usage.output_tokens),
|
||||
});
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Err(internal_error("sample stream ended before completion"))
|
||||
}
|
||||
@@ -154,6 +154,7 @@ mod thread_read;
|
||||
mod thread_read_state;
|
||||
mod thread_resume;
|
||||
mod thread_revert;
|
||||
mod thread_sample;
|
||||
mod thread_sections;
|
||||
mod thread_settings_update;
|
||||
mod thread_shell_command;
|
||||
|
||||
+195
@@ -0,0 +1,195 @@
|
||||
---
|
||||
source: app-server/tests/suite/v2/thread_sample.rs
|
||||
expression: "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.\",\n&requests, &ContextSnapshotOptions::default().include_request_settings(),)"
|
||||
---
|
||||
Scenario: 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.
|
||||
|
||||
## Window 1
|
||||
Settings:
|
||||
include: ["reasoning.encrypted_content"]
|
||||
model: "gpt-5.4"
|
||||
parallel_tool_calls: false
|
||||
prompt_cache_key: "<PROMPT_CACHE_KEY 1>"
|
||||
reasoning: {"context":"all_turns","effort":"high","summary":"auto"}
|
||||
service_tier: "priority"
|
||||
store: false
|
||||
stream: true
|
||||
tool_choice: "auto"
|
||||
-- request 1 (turn) --
|
||||
00:additional_tools/developer (3; hash=bf13940dee826ebd):
|
||||
- namespace/functions; hash=028af59ad835936d
|
||||
- function/exec_command
|
||||
- function/write_stdin
|
||||
- function/request_user_input
|
||||
- function/request_permissions
|
||||
- function/view_image
|
||||
- function/get_goal
|
||||
- function/create_goal
|
||||
- function/update_goal
|
||||
- function/parent_tool
|
||||
- namespace/multi_agent_v1: Tools for spawning and managing sub-agents.; hash=c826e5b07fd2820f
|
||||
- function/close_agent
|
||||
- function/resume_agent
|
||||
- function/send_input
|
||||
- function/spawn_agent
|
||||
- function/wait_agent
|
||||
- namespace/web: Tools in the web namespace.; hash=0037b1ed8d50fab5
|
||||
- function/run
|
||||
01:message/developer:
|
||||
Retain these base instructions.
|
||||
02:message/developer[4]:
|
||||
[01] Retain this system policy context.
|
||||
[02] <skills_instructions>
|
||||
## Skills
|
||||
A skill is a set of local instructions to follow that is stored in a `SKILL.md` file. Below is the list ...ed into an absolute path using the skill roots table. [hash=ca27c3957f42502b]
|
||||
### Skill roots
|
||||
- `r0` = `<SKILLS_ROOT>/.system`
|
||||
### Available skills
|
||||
- imagegen: Generate or edit raster images when the task benefits from AI-created bitmap visuals such as...ctly in HTML/CSS/canvas. (file: r0/imagegen/SKILL.md) [hash=dcbada3f7f876b5c]
|
||||
- openai-docs: Use for Codex models/pricing, scheduled tasks, skills, settings, setup, troubleshooting, ...merely mention Codex. (file: r0/openai-docs/SKILL.md) [hash=daf7d1f4fd2d2abe]
|
||||
- skill-creator: Create or update a Codex skill with appropriately scoped instructions and any needed supporting resources. (file: r0/skill-creator/SKILL.md)
|
||||
- skill-installer: Install Codex skills into $CODEX_HOME/skills from a curated list or a GitHub repo pat...g private repos). (file: r0/skill-installer/SKILL.md) [hash=d3eb4395bd8e52c5]
|
||||
</skills_instructions>
|
||||
[03] <permissions instructions>
|
||||
Filesystem sandboxing defines which files can be read or written. `sandbox_mode` is `read-only`: The san... permits reading files. Network access is restricted. [hash=be7da4425f89053c]
|
||||
Approval policy is currently never. Do not provide the `sandbox_permissions` for any reason, commands will be rejected.
|
||||
</permissions instructions>
|
||||
[04] <collaboration_mode>Keep the parent's planning instructions.</collaboration_mode>
|
||||
03:message/user:
|
||||
<environment_context>
|
||||
<cwd><CWD></cwd>
|
||||
<shell><HOST_SHELL></shell>
|
||||
<current_date><CURRENT_DATE></current_date>
|
||||
<timezone><HOST_TIMEZONE></timezone>
|
||||
<filesystem><workspace_roots><root><CWD></root></workspace_roots><permission_profile type="managed"><f...ntry></file_system></permission_profile></filesystem> [hash=e8c1de60307d0a3d]
|
||||
</environment_context>
|
||||
04:message/user:
|
||||
Initial turn
|
||||
05:configuration_update:reasoning={"effort":"high"}
|
||||
-- request 2 (turn) --
|
||||
06:reasoning:summary=Check the release.:encrypted=true
|
||||
07:message/assistant:
|
||||
Done
|
||||
08:message/user:
|
||||
Continue identically
|
||||
09:configuration_update:reasoning={"effort":"medium"}
|
||||
|
||||
## Window 2 (after request 2: input repeated)
|
||||
Settings: same as window 1
|
||||
-- request 3 (turn) --
|
||||
00:additional_tools/developer (3; hash=bf13940dee826ebd):
|
||||
- namespace/functions; hash=028af59ad835936d
|
||||
- function/exec_command
|
||||
- function/write_stdin
|
||||
- function/request_user_input
|
||||
- function/request_permissions
|
||||
- function/view_image
|
||||
- function/get_goal
|
||||
- function/create_goal
|
||||
- function/update_goal
|
||||
- function/parent_tool
|
||||
- namespace/multi_agent_v1: Tools for spawning and managing sub-agents.; hash=c826e5b07fd2820f
|
||||
- function/close_agent
|
||||
- function/resume_agent
|
||||
- function/send_input
|
||||
- function/spawn_agent
|
||||
- function/wait_agent
|
||||
- namespace/web: Tools in the web namespace.; hash=0037b1ed8d50fab5
|
||||
- function/run
|
||||
01:message/developer:
|
||||
Retain these base instructions.
|
||||
02:message/developer[4]:
|
||||
[01] Retain this system policy context.
|
||||
[02] <skills_instructions>
|
||||
## Skills
|
||||
A skill is a set of local instructions to follow that is stored in a `SKILL.md` file. Below is the list ...ed into an absolute path using the skill roots table. [hash=ca27c3957f42502b]
|
||||
### Skill roots
|
||||
- `r0` = `<SKILLS_ROOT>/.system`
|
||||
### Available skills
|
||||
- imagegen: Generate or edit raster images when the task benefits from AI-created bitmap visuals such as...ctly in HTML/CSS/canvas. (file: r0/imagegen/SKILL.md) [hash=dcbada3f7f876b5c]
|
||||
- openai-docs: Use for Codex models/pricing, scheduled tasks, skills, settings, setup, troubleshooting, ...merely mention Codex. (file: r0/openai-docs/SKILL.md) [hash=daf7d1f4fd2d2abe]
|
||||
- skill-creator: Create or update a Codex skill with appropriately scoped instructions and any needed supporting resources. (file: r0/skill-creator/SKILL.md)
|
||||
- skill-installer: Install Codex skills into $CODEX_HOME/skills from a curated list or a GitHub repo pat...g private repos). (file: r0/skill-installer/SKILL.md) [hash=d3eb4395bd8e52c5]
|
||||
</skills_instructions>
|
||||
[03] <permissions instructions>
|
||||
Filesystem sandboxing defines which files can be read or written. `sandbox_mode` is `read-only`: The san... permits reading files. Network access is restricted. [hash=be7da4425f89053c]
|
||||
Approval policy is currently never. Do not provide the `sandbox_permissions` for any reason, commands will be rejected.
|
||||
</permissions instructions>
|
||||
[04] <collaboration_mode>Keep the parent's planning instructions.</collaboration_mode>
|
||||
03:message/user:
|
||||
<environment_context>
|
||||
<cwd><CWD></cwd>
|
||||
<shell><HOST_SHELL></shell>
|
||||
<current_date><CURRENT_DATE></current_date>
|
||||
<timezone><HOST_TIMEZONE></timezone>
|
||||
<filesystem><workspace_roots><root><CWD></root></workspace_roots><permission_profile type="managed"><f...ntry></file_system></permission_profile></filesystem> [hash=e8c1de60307d0a3d]
|
||||
</environment_context>
|
||||
04:message/user:
|
||||
Initial turn
|
||||
05:configuration_update:reasoning={"effort":"high"}
|
||||
06:reasoning:summary=Check the release.:encrypted=true
|
||||
07:message/assistant:
|
||||
Done
|
||||
08:message/user:
|
||||
Continue identically
|
||||
09:configuration_update:reasoning={"effort":"medium"}
|
||||
|
||||
## Window 3 (after request 3: input repeated)
|
||||
Settings: same as window 1
|
||||
-- request 4 (turn) --
|
||||
00:additional_tools/developer (3; hash=bf13940dee826ebd):
|
||||
- namespace/functions; hash=028af59ad835936d
|
||||
- function/exec_command
|
||||
- function/write_stdin
|
||||
- function/request_user_input
|
||||
- function/request_permissions
|
||||
- function/view_image
|
||||
- function/get_goal
|
||||
- function/create_goal
|
||||
- function/update_goal
|
||||
- function/parent_tool
|
||||
- namespace/multi_agent_v1: Tools for spawning and managing sub-agents.; hash=c826e5b07fd2820f
|
||||
- function/close_agent
|
||||
- function/resume_agent
|
||||
- function/send_input
|
||||
- function/spawn_agent
|
||||
- function/wait_agent
|
||||
- namespace/web: Tools in the web namespace.; hash=0037b1ed8d50fab5
|
||||
- function/run
|
||||
01:message/developer:
|
||||
Retain these base instructions.
|
||||
02:message/developer[4]:
|
||||
[01] Retain this system policy context.
|
||||
[02] <skills_instructions>
|
||||
## Skills
|
||||
A skill is a set of local instructions to follow that is stored in a `SKILL.md` file. Below is the list ...ed into an absolute path using the skill roots table. [hash=ca27c3957f42502b]
|
||||
### Skill roots
|
||||
- `r0` = `<SKILLS_ROOT>/.system`
|
||||
### Available skills
|
||||
- imagegen: Generate or edit raster images when the task benefits from AI-created bitmap visuals such as...ctly in HTML/CSS/canvas. (file: r0/imagegen/SKILL.md) [hash=dcbada3f7f876b5c]
|
||||
- openai-docs: Use for Codex models/pricing, scheduled tasks, skills, settings, setup, troubleshooting, ...merely mention Codex. (file: r0/openai-docs/SKILL.md) [hash=daf7d1f4fd2d2abe]
|
||||
- skill-creator: Create or update a Codex skill with appropriately scoped instructions and any needed supporting resources. (file: r0/skill-creator/SKILL.md)
|
||||
- skill-installer: Install Codex skills into $CODEX_HOME/skills from a curated list or a GitHub repo pat...g private repos). (file: r0/skill-installer/SKILL.md) [hash=d3eb4395bd8e52c5]
|
||||
</skills_instructions>
|
||||
[03] <permissions instructions>
|
||||
Filesystem sandboxing defines which files can be read or written. `sandbox_mode` is `read-only`: The san... permits reading files. Network access is restricted. [hash=be7da4425f89053c]
|
||||
Approval policy is currently never. Do not provide the `sandbox_permissions` for any reason, commands will be rejected.
|
||||
</permissions instructions>
|
||||
[04] <collaboration_mode>Keep the parent's planning instructions.</collaboration_mode>
|
||||
03:message/user:
|
||||
<environment_context>
|
||||
<cwd><CWD></cwd>
|
||||
<shell><HOST_SHELL></shell>
|
||||
<current_date><CURRENT_DATE></current_date>
|
||||
<timezone><HOST_TIMEZONE></timezone>
|
||||
<filesystem><workspace_roots><root><CWD></root></workspace_roots><permission_profile type="managed"><f...ntry></file_system></permission_profile></filesystem> [hash=e8c1de60307d0a3d]
|
||||
</environment_context>
|
||||
04:message/user:
|
||||
Initial turn
|
||||
05:configuration_update:reasoning={"effort":"high"}
|
||||
06:reasoning:summary=Check the release.:encrypted=true
|
||||
07:message/assistant:
|
||||
Done
|
||||
08:message/user:
|
||||
Continue identically
|
||||
09:configuration_update:reasoning={"effort":"medium"}
|
||||
@@ -0,0 +1,957 @@
|
||||
//! 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::ThreadForkParams;
|
||||
use codex_app_server_protocol::ThreadForkResponse;
|
||||
use codex_app_server_protocol::ThreadHistoryMode;
|
||||
use codex_app_server_protocol::ThreadLoadedListResponse;
|
||||
use codex_app_server_protocol::ThreadSampleParams;
|
||||
use codex_app_server_protocol::ThreadSampleResponse;
|
||||
use codex_app_server_protocol::ThreadStartParams;
|
||||
use codex_app_server_protocol::TurnStartParams;
|
||||
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::responses;
|
||||
use core_test_support::skip_if_host_windows;
|
||||
use core_test_support::skip_if_remote;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::Value;
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tempfile::TempDir;
|
||||
use test_case::test_case;
|
||||
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_case(ThreadHistoryMode::Legacy; "legacy")]
|
||||
#[test_case(ThreadHistoryMode::Paginated; "paginated")]
|
||||
#[tokio::test]
|
||||
async fn sample_preserves_context_and_keeps_source_unchanged(
|
||||
history_mode: ThreadHistoryMode,
|
||||
) -> Result<()> {
|
||||
let server = responses::start_mock_server().await;
|
||||
let mut parent_responses = vec![responses::sse(vec![
|
||||
responses::ev_reasoning_item(
|
||||
"parent-reasoning",
|
||||
&["Check the release."],
|
||||
&["Keep the existing context."],
|
||||
),
|
||||
responses::ev_assistant_message("answer", "Done"),
|
||||
responses::ev_completed_with_tokens(
|
||||
"parent",
|
||||
/*total_tokens*/
|
||||
if history_mode == ThreadHistoryMode::Paginated {
|
||||
330_000
|
||||
} else {
|
||||
100
|
||||
},
|
||||
),
|
||||
])];
|
||||
if history_mode == ThreadHistoryMode::Paginated {
|
||||
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 history_mode == ThreadHistoryMode::Paginated {
|
||||
75
|
||||
} else {
|
||||
0
|
||||
}),
|
||||
),
|
||||
])),
|
||||
..Default::default()
|
||||
})
|
||||
.await?;
|
||||
let completed = client
|
||||
.start_turn_and_wait_for_completion(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 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(),
|
||||
prompt: "Continue identically".to_string(),
|
||||
effort: Some(ReasoningEffort::Medium),
|
||||
turn_trigger: Some("composer_predictions".to_string()),
|
||||
},
|
||||
)
|
||||
.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 {
|
||||
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 history_mode == ThreadHistoryMode::Paginated {
|
||||
assert_eq!(
|
||||
requests[sample_index].inputs_of_type("compaction")[0]["encrypted_content"],
|
||||
"sample-test-compacted-parent",
|
||||
);
|
||||
}
|
||||
let prediction = &requests[sample_index];
|
||||
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["turn_trigger"], "composer_predictions");
|
||||
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!(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()],
|
||||
})
|
||||
.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(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sample_blocks_mcp_execution_without_changing_parent_tools() -> 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::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(),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
assert_eq!(
|
||||
(response.text, response.response_id),
|
||||
(None, "sample-tool".to_string())
|
||||
);
|
||||
let fork = prediction_fork(&mut client, &parent_id).await?;
|
||||
client
|
||||
.start_turn_and_wait_for_completion(TurnStartParams {
|
||||
thread_id: fork,
|
||||
input: input("Call the tool again"),
|
||||
..Default::default()
|
||||
})
|
||||
.await?;
|
||||
let requests = mock.requests();
|
||||
assert_eq!(requests.len(), 4);
|
||||
assert_eq!(
|
||||
comparable_request(requests[2].body_json()),
|
||||
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(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sample_wait_keeps_parent_responsive_and_timeout_reclaims_child() -> 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 error = timeout(
|
||||
WAIT,
|
||||
client.read_stream_until_error_message(RequestId::Integer(sample_id)),
|
||||
)
|
||||
.await??;
|
||||
assert!(
|
||||
error.error.message.contains("timed out"),
|
||||
"{}",
|
||||
error.error.message
|
||||
);
|
||||
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 completed 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(())
|
||||
}
|
||||
|
||||
#[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(())
|
||||
}
|
||||
Reference in New Issue
Block a user