use std::collections::HashMap;
use std::collections::HashSet;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use crate::SkillInjections;
use crate::build_skill_injections;
use crate::client::ModelClientSession;
use crate::client_common::Prompt;
use crate::client_common::ResponseEvent;
use crate::collect_explicit_skill_mentions;
use crate::compact::InitialContextInjection;
use crate::compact::run_inline_auto_compact_task;
use crate::compact::should_use_remote_compact_task;
use crate::compact_remote::run_inline_remote_auto_compact_task;
use crate::compact_remote_v2::run_inline_remote_auto_compact_task as run_inline_remote_auto_compact_task_v2;
use crate::connectors;
use crate::context::ContextualUserFragment;
use crate::feedback_tags;
use crate::hook_runtime::inspect_pending_input;
use crate::hook_runtime::record_additional_contexts;
use crate::hook_runtime::record_pending_input;
use crate::hook_runtime::run_legacy_after_agent_hook;
use crate::hook_runtime::run_pending_session_start_hooks;
use crate::hook_runtime::run_turn_stop_hooks;
use crate::injection::ToolMentionKind;
use crate::injection::app_id_from_path;
use crate::injection::tool_kind_for_path;
use crate::mcp_skill_dependencies::maybe_prompt_and_install_mcp_dependencies;
use crate::mcp_tool_exposure::build_mcp_tool_runtimes;
use crate::mentions::build_connector_slug_counts;
use crate::mentions::build_skill_name_counts;
use crate::mentions::collect_explicit_app_ids;
use crate::mentions::collect_explicit_plugin_mentions;
use crate::mentions::collect_tool_mentions_from_messages;
use crate::plugins::build_plugin_injections;
use crate::responses_metadata::CodexResponsesMetadata;
use crate::responses_metadata::CodexResponsesRequestKind;
use crate::responses_retry::ResponsesStreamRequest;
use crate::responses_retry::handle_retryable_response_stream_error;
use crate::session::PreviousTurnSettings;
use crate::session::TurnInput;
use crate::session::session::Session;
use crate::session::step_context::StepContext;
use crate::session::turn_context::TurnContext;
use crate::stream_events_utils::HandleOutputCtx;
use crate::stream_events_utils::TurnItemContributorPolicy;
use crate::stream_events_utils::finalize_non_tool_response_item;
use crate::stream_events_utils::handle_non_tool_response_item;
use crate::stream_events_utils::handle_output_item_done;
use crate::stream_events_utils::last_assistant_message_from_item;
use crate::stream_events_utils::mark_thread_memory_mode_polluted_if_external_context;
use crate::stream_events_utils::raw_assistant_output_text_from_item;
use crate::stream_events_utils::record_completed_response_item_with_finalized_facts;
use crate::tasks::emit_compact_metric;
use crate::tools::ToolRouter;
use crate::tools::context::SharedTurnDiffTracker;
use crate::tools::parallel::ToolCallRuntime;
use crate::tools::registry::ToolArgumentDiffConsumer;
use crate::tools::router::ToolRouterParams;
use crate::tools::router::ToolSuggestCandidates;
use crate::tools::router::ToolSuggestPresentation;
use crate::tools::router::extension_tool_executors;
use crate::tools::spec_plan::search_tool_enabled;
use crate::tools::spec_plan::tool_suggest_enabled;
use crate::turn_diff_tracker::TurnDiffTracker;
use crate::turn_timing::record_turn_ttft_metric;
use crate::util::error_or_panic;
use codex_analytics::AppInvocation;
use codex_analytics::CompactionPhase;
use codex_analytics::CompactionReason;
use codex_analytics::InvocationType;
use codex_analytics::TurnResolvedConfigFact;
use codex_analytics::build_track_events_context;
use codex_async_utils::OrCancelExt;
use codex_core_plugins::RecommendedPluginCandidatesInput;
use codex_core_skills::injection::InjectedHostSkillPrompts;
use codex_extension_api::TurnInputContext;
use codex_extension_api::TurnInputEnvironment;
use codex_features::Feature;
use codex_git_utils::get_git_repo_root_with_fs;
use codex_protocol::ResponseItemId;
use codex_protocol::config_types::AutoCompactTokenLimitScope;
use codex_protocol::config_types::ModeKind;
use codex_protocol::config_types::ServiceTier;
use codex_protocol::error::CodexErr;
use codex_protocol::error::Result as CodexResult;
use codex_protocol::items::PlanItem;
use codex_protocol::items::TurnItem;
use codex_protocol::items::build_hook_prompt_message;
use codex_protocol::models::BaseInstructions;
use codex_protocol::models::ContentItem;
use codex_protocol::models::MessagePhase;
use codex_protocol::models::ResponseInputItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::AgentMessageContentDeltaEvent;
use codex_protocol::protocol::AgentReasoningSectionBreakEvent;
use codex_protocol::protocol::CodexErrorInfo;
use codex_protocol::protocol::ErrorEvent;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::PlanDeltaEvent;
use codex_protocol::protocol::RawResponseCompletedEvent;
use codex_protocol::protocol::ReasoningContentDeltaEvent;
use codex_protocol::protocol::ReasoningRawContentDeltaEvent;
use codex_protocol::protocol::SafetyBufferingEvent;
use codex_protocol::protocol::TurnDiffEvent;
use codex_protocol::protocol::WarningEvent;
use codex_protocol::user_input::UserInput;
use codex_tools::ToolName;
use codex_tools::filter_request_plugin_install_discoverable_tools_for_client;
use codex_utils_stream_parser::AssistantTextChunk;
use codex_utils_stream_parser::AssistantTextStreamParser;
use codex_utils_stream_parser::ProposedPlanSegment;
use codex_utils_stream_parser::extract_proposed_plan_text;
use codex_utils_stream_parser::strip_citations;
use futures::future::BoxFuture;
use futures::prelude::*;
use futures::stream::FuturesOrdered;
use tokio_util::sync::CancellationToken;
use tracing::Instrument;
use tracing::error;
use tracing::field;
use tracing::info;
use tracing::instrument;
use tracing::trace;
use tracing::trace_span;
use tracing::warn;
pub(crate) async fn run_turn(
sess: Arc<Session>,
turn_context: Arc<TurnContext>,
turn_extension_data: Arc<codex_extension_api::ExtensionData>,
input: Vec<TurnInput>,
prewarmed_client_session: Option<ModelClientSession>,
cancellation_token: CancellationToken,
) -> CodexResult<Option<String>> {
let mut client_session =
prewarmed_client_session.unwrap_or_else(|| sess.services.model_client.new_session());
if let Err(err) = run_pre_sampling_compact(&sess, &turn_context, &mut client_session).await {
if matches!(err, CodexErr::TurnAborted) {
return Err(err);
}
let error = err.to_codex_protocol_error();
sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone())
.await;
error!("Failed to run pre-sampling compact");
return Ok(None);
}
let first_step_context = sess.capture_step_context(Arc::clone(&turn_context)).await;
let (mut world_state, display_roots) = tokio::join!(
sess.record_context_updates_and_set_reference_context_item(first_step_context.as_ref()),
turn_diff_display_roots(turn_context.as_ref()),
);
let Some((injection_items, explicitly_enabled_connectors)) = build_skills_and_plugins(
&sess,
first_step_context.as_ref(),
&input,
&cancellation_token,
)
.await
else {
return Ok(None);
};
if run_pending_session_start_hooks(&sess, &turn_context).await {
return Ok(None);
}
let mut can_drain_pending_input = input.is_empty();
if run_hooks_and_record_inputs(&sess, &turn_context, &input).await {
return Ok(None);
}
sess.merge_connector_selection(explicitly_enabled_connectors.clone())
.await;
sess.set_previous_turn_settings(Some(PreviousTurnSettings {
model: turn_context.model_info.slug.clone(),
comp_hash: turn_context.model_info.comp_hash.clone(),
realtime_active: Some(turn_context.realtime_active),
}))
.await;
for response_item in injection_items {
sess.record_conversation_items(&turn_context, std::slice::from_ref(&response_item))
.await;
}
track_turn_resolved_config_analytics(&sess, &turn_context, &input).await;
let mut last_agent_message: Option<String> = None;
let mut stop_hook_active = false;
let turn_diff_tracker = Arc::new(tokio::sync::Mutex::new(
TurnDiffTracker::with_environment_display_roots(display_roots),
));
let mut next_step_context = Some(first_step_context);
loop {
let pending_input = if can_drain_pending_input {
sess.input_queue.get_pending_input(&sess.active_turn).await
} else {
Vec::new()
};
if run_hooks_and_record_inputs(&sess, &turn_context, &pending_input).await {
break;
}
let window_id = sess.current_window_id().await;
super::rollout_budget::maybe_record_reminder(
sess.as_ref(),
turn_context.as_ref(),
&window_id,
)
.await;
let step_context = match next_step_context.take() {
Some(step_context) => step_context,
None => sess.capture_step_context(Arc::clone(&turn_context)).await,
};
let sampling_request_result: CodexResult<_> = async {
super::time_reminder::maybe_record_current_time_reminder(
sess.as_ref(),
turn_context.as_ref(),
&window_id,
)
.await?;
world_state = sess
.record_step_world_state_if_changed(&world_state, step_context.as_ref())
.await;
let sampling_request_input: Vec<ResponseItem> = async {
sess.clone_history()
.await
.for_prompt(&turn_context.model_info.input_modalities)
}
.instrument(trace_span!("run_turn.prepare_sampling_request_input"))
.await;
let responses_metadata = turn_context.turn_metadata_state.to_responses_metadata(
sess.installation_id.clone(),
window_id,
CodexResponsesRequestKind::Turn,
);
run_sampling_request(
Arc::clone(&sess),
Arc::clone(&step_context),
Arc::clone(&turn_extension_data),
Arc::clone(&turn_diff_tracker),
&mut client_session,
&responses_metadata,
sampling_request_input,
cancellation_token.child_token(),
)
.await
}
.await;
match sampling_request_result {
Ok((sampling_request_output, sampling_request_input)) => {
let SamplingRequestResult {
needs_follow_up: model_needs_follow_up,
last_agent_message: sampling_request_last_agent_message,
} = sampling_request_output;
if model_needs_follow_up {
sess.input_queue
.accept_mailbox_delivery_for_current_turn(
&sess.active_turn,
&turn_context.sub_id,
)
.await;
}
can_drain_pending_input = true;
let (has_pending_input, token_status, estimated_token_count) = async {
let has_pending_input =
sess.input_queue.has_pending_input(&sess.active_turn).await;
let token_status = super::context_window::context_window_token_status(
sess.as_ref(),
turn_context.as_ref(),
)
.await;
let estimated_token_count =
sess.get_estimated_token_count(turn_context.as_ref()).await;
(has_pending_input, token_status, estimated_token_count)
}
.instrument(trace_span!("run_turn.collect_post_sampling_state"))
.await;
let needs_follow_up = model_needs_follow_up || has_pending_input;
let token_limit_reached = token_status.token_limit_reached;
trace!(
turn_id = %turn_context.sub_id,
total_usage_tokens = token_status.active_context_tokens,
auto_compact_scope_tokens = token_status.auto_compact_scope_tokens,
estimated_token_count = ?estimated_token_count,
auto_compact_scope_limit = ?token_status.auto_compact_scope_limit,
auto_compact_limit_scope = ?turn_context.config.model_auto_compact_token_limit_scope,
auto_compact_window_prefill_tokens = ?token_status.auto_compact_window_prefill_tokens,
full_context_window_limit = ?token_status.full_context_window_limit,
full_context_window_limit_reached = token_status.full_context_window_limit_reached,
token_limit_reached,
model_needs_follow_up,
has_pending_input,
needs_follow_up,
"post sampling token usage"
);
let should_roll_over = needs_follow_up
&& (sess.take_new_context_window_request().await || token_limit_reached);
let allow_auto_compact_fallback = !should_roll_over && !token_limit_reached;
super::token_budget::maybe_record(
sess.as_ref(),
turn_context.as_ref(),
token_status.base_window_tokens_remaining,
allow_auto_compact_fallback,
)
.await;
if should_roll_over {
if let Err(err) = run_auto_compact(
&sess,
Arc::clone(&step_context),
None,
&mut client_session,
InitialContextInjection::BeforeLastUserMessage {
world_state: Arc::clone(&world_state),
step_context: Arc::clone(&step_context),
},
CompactionReason::ContextLimit,
CompactionPhase::MidTurn,
)
.await
{
if matches!(err, CodexErr::TurnAborted) {
return Err(err);
}
let error = err.to_codex_protocol_error();
sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone())
.await;
return Ok(None);
}
if run_pending_session_start_hooks(&sess, &turn_context).await {
return Ok(None);
}
can_drain_pending_input = !model_needs_follow_up;
continue;
}
if !needs_follow_up {
last_agent_message = sampling_request_last_agent_message;
let stop_outcome = run_turn_stop_hooks(
&sess,
&turn_context,
stop_hook_active,
last_agent_message.clone(),
)
.await;
if stop_outcome.should_block {
if let Some(hook_prompt_message) =
build_hook_prompt_message(&stop_outcome.continuation_fragments)
{
sess.record_response_item_and_emit_turn_item(
&turn_context,
hook_prompt_message,
)
.await;
sess.input_queue
.accept_mailbox_delivery_for_current_turn(
&sess.active_turn,
&turn_context.sub_id,
)
.await;
stop_hook_active = true;
continue;
} else {
sess.send_event(
&turn_context,
EventMsg::Warning(WarningEvent {
message: "Stop hook requested continuation without a prompt; ignoring the block.".to_string(),
}),
)
.await;
}
}
if stop_outcome.should_stop {
break;
}
if run_legacy_after_agent_hook(
&sess,
&turn_context,
&sampling_request_input,
last_agent_message.clone(),
)
.await
{
return Ok(None);
}
break;
}
continue;
}
Err(err @ CodexErr::TurnAborted) => {
return Err(err);
}
Err(codex_error @ CodexErr::InvalidImageRequest()) => {
sess.track_turn_codex_error(turn_context.as_ref(), &codex_error);
let error = CodexErrorInfo::BadRequest;
sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone())
.await;
let event = EventMsg::Error(ErrorEvent {
message: "Invalid image in your last message. Please remove it and try again."
.to_string(),
codex_error_info: Some(error),
});
sess.send_event(&turn_context, event).await;
break;
}
Err(e) => {
info!("Turn error: {e:#}");
let error = e.to_codex_protocol_error();
sess.emit_turn_error_lifecycle(turn_context.as_ref(), error.clone())
.await;
sess.track_turn_codex_error(turn_context.as_ref(), &e);
let event = EventMsg::Error(e.to_error_event( None));
sess.send_event(&turn_context, event).await;
break;
}
}
}
Ok(last_agent_message)
}
#[instrument(level = "trace", skip_all)]
async fn turn_diff_display_roots(turn_context: &TurnContext) -> Vec<(String, PathBuf)> {
let mut display_roots = Vec::new();
for turn_environment in turn_context.environments.turn_environments() {
let Ok(cwd) = turn_environment.cwd().to_abs_path() else {
continue;
};
let root =
get_git_repo_root_with_fs(turn_environment.environment.get_filesystem().as_ref(), &cwd)
.await
.unwrap_or(cwd)
.into_path_buf();
display_roots.push((turn_environment.environment_id.clone(), root));
}
display_roots
}
#[instrument(level = "trace", skip_all)]
async fn run_hooks_and_record_inputs(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
input: &[TurnInput],
) -> bool {
let mut blocked_input = false;
let mut accepted_user_input = false;
for input_item in input {
let hook_outcome = inspect_pending_input(sess, turn_context, input_item).await;
if hook_outcome.should_stop {
blocked_input = true;
record_additional_contexts(sess, turn_context, hook_outcome.additional_contexts).await;
} else {
if matches!(input_item, TurnInput::UserInput { content, .. } if !content.is_empty()) {
accepted_user_input = true;
}
record_pending_input(
sess,
turn_context,
input_item.clone(),
hook_outcome.additional_contexts,
)
.await;
}
}
blocked_input && !accepted_user_input
}
#[instrument(level = "trace", skip_all)]
async fn build_skills_and_plugins(
sess: &Arc<Session>,
step_context: &StepContext,
input: &[TurnInput],
cancellation_token: &CancellationToken,
) -> Option<(Vec<ResponseItem>, HashSet<String>)> {
let turn_context = step_context.turn.as_ref();
if crate::guardian::is_guardian_reviewer_source(&turn_context.session_source) {
return Some((Vec::new(), HashSet::new()));
}
let user_input = input
.iter()
.filter_map(|item| match item {
TurnInput::UserInput { content, .. } => Some(content.as_slice()),
TurnInput::ResponseItem(_) | TurnInput::InterAgentCommunication(_) => None,
})
.flatten()
.cloned()
.collect::<Vec<_>>();
let tracking = build_track_events_context(
turn_context.model_info.slug.clone(),
sess.thread_id.to_string(),
turn_context.sub_id.clone(),
turn_context.originator.clone(),
);
let loaded_plugins = sess
.services
.plugins_manager
.plugins_for_config(&turn_context.config.plugins_config_input())
.await;
let mentioned_plugins =
collect_explicit_plugin_mentions(&user_input, loaded_plugins.capability_summaries());
let connector_snapshot = step_context.mcp.config().connector_snapshot.clone();
let mcp_tools = if turn_context.apps_enabled() || !mentioned_plugins.is_empty() {
match step_context
.mcp
.manager_arc()
.list_all_tools()
.or_cancel(cancellation_token)
.await
{
Ok(mcp_tools) => mcp_tools,
Err(_) if turn_context.apps_enabled() => return None,
Err(_) => Vec::new(),
}
} else {
Vec::new()
};
let available_connectors = if turn_context.apps_enabled() {
let connectors = codex_connectors::merge::merge_plugin_connectors_with_accessible(
connector_snapshot
.connector_ids()
.iter()
.map(|connector_id| connector_id.0.clone()),
connectors::accessible_connectors_from_mcp_tools(&mcp_tools),
);
connectors::with_app_enabled_state(connectors, &turn_context.config)
} else {
Vec::new()
};
let skills_outcome = turn_context.turn_skills.snapshot.outcome();
let connector_slug_counts = build_connector_slug_counts(&available_connectors);
let extension_injection_items =
build_extension_turn_input_items(sess, step_context, &user_input, cancellation_token)
.await?;
let skill_name_counts_lower =
build_skill_name_counts(&skills_outcome.skills, &skills_outcome.disabled_paths).1;
let mentioned_skills = collect_explicit_skill_mentions(
&user_input,
&skills_outcome.skills,
&skills_outcome.disabled_paths,
&connector_slug_counts,
);
maybe_prompt_and_install_mcp_dependencies(
sess,
turn_context,
cancellation_token,
&mentioned_skills,
Some(sess.mcp_elicitation_reviewer()),
)
.await;
let injected_host_skill_prompts = turn_context
.extension_data
.get::<InjectedHostSkillPrompts>();
let SkillInjections {
items: skill_injections,
warnings: skill_warnings,
} = build_skill_injections(
&mentioned_skills,
Some(skills_outcome),
Some(&turn_context.session_telemetry),
&sess.services.analytics_events_client,
tracking.clone(),
)
.await;
for message in skill_warnings {
sess.send_event(turn_context, EventMsg::Warning(WarningEvent { message }))
.await;
}
let skill_items: Vec<ResponseItem> = skill_injections
.iter()
.map(|skill| ContextualUserFragment::into(crate::context::SkillInstructions::from(skill)))
.collect();
let skill_connector_ids = collect_explicit_app_ids_from_skill_items(
&skill_items,
&available_connectors,
&skill_name_counts_lower,
);
let plugin_items =
build_plugin_injections(&mentioned_plugins, &mcp_tools, &available_connectors);
let mut explicitly_enabled_connectors = collect_explicit_app_ids(&user_input);
explicitly_enabled_connectors.extend(skill_connector_ids);
let connector_names_by_id = available_connectors
.iter()
.map(|connector| (connector.id.as_str(), connector.name.as_str()))
.collect::<HashMap<&str, &str>>();
let mentioned_app_invocations = explicitly_enabled_connectors
.iter()
.map(|connector_id| AppInvocation {
connector_id: Some(connector_id.clone()),
app_name: connector_names_by_id
.get(connector_id.as_str())
.map(|name| (*name).to_string()),
invocation_type: Some(InvocationType::Explicit),
})
.collect::<Vec<_>>();
sess.services
.analytics_events_client
.track_app_mentioned(tracking.clone(), mentioned_app_invocations);
for summary in &mentioned_plugins {
if let Some(plugin) = sess
.services
.plugins_manager
.telemetry_metadata_for_capability_summary(summary)
{
sess.services
.analytics_events_client
.track_plugin_used(tracking.clone(), plugin);
}
}
let mut injection_items: Vec<ResponseItem> = match injected_host_skill_prompts {
Some(injected_host_skill_prompts) => skill_injections
.iter()
.filter(|skill| !injected_host_skill_prompts.contains_path(&skill.path))
.map(|skill| {
ContextualUserFragment::into(crate::context::SkillInstructions::from(skill))
})
.collect(),
None => skill_items,
};
injection_items.extend(plugin_items);
injection_items.extend(extension_injection_items);
Some((injection_items, explicitly_enabled_connectors))
}
#[tracing::instrument(
level = "trace",
skip_all,
fields(user_input_count = user_input.len())
)]
async fn build_extension_turn_input_items(
sess: &Arc<Session>,
step_context: &StepContext,
user_input: &[UserInput],
cancellation_token: &CancellationToken,
) -> Option<Vec<ResponseItem>> {
let turn_context = step_context.turn.as_ref();
let contributors = sess.services.extensions.turn_input_contributors().to_vec();
if contributors.is_empty() {
return Some(Vec::new());
}
let environments = turn_context
.environments
.turn_environments()
.enumerate()
.filter_map(|(index, environment)| {
Some(TurnInputEnvironment {
environment_id: environment.environment_id.clone(),
cwd: environment.cwd().to_abs_path().ok()?.into_path_buf(),
is_primary: index == 0,
})
})
.collect::<Vec<_>>();
let input = TurnInputContext {
turn_id: turn_context.sub_id.to_string(),
user_input: user_input.to_vec(),
environments,
};
let mut items = Vec::new();
for contributor in contributors {
let contributed_fragments = contributor
.contribute(
input.clone(),
&sess.services.session_extension_data,
&sess.services.thread_extension_data,
turn_context.extension_data.as_ref(),
&step_context.extension_data,
)
.or_cancel(cancellation_token)
.await
.ok()?;
items.extend(
contributed_fragments
.into_iter()
.map(ContextualUserFragment::into_boxed_response_item),
);
}
Some(items)
}
#[tracing::instrument(
level = "trace",
skip_all,
fields(input_count = input.len())
)]
async fn track_turn_resolved_config_analytics(
sess: &Session,
turn_context: &TurnContext,
input: &[TurnInput],
) {
let thread_config = {
let state = sess.state.lock().await;
state.session_configuration.thread_config_snapshot()
};
let is_first_turn = {
let mut state = sess.state.lock().await;
state.take_next_turn_is_first()
};
sess.services
.analytics_events_client
.track_turn_resolved_config(TurnResolvedConfigFact {
turn_id: turn_context.sub_id.clone(),
thread_id: sess.thread_id.to_string(),
num_input_images: input
.iter()
.filter_map(|item| match item {
TurnInput::UserInput { content, .. } => Some(content.as_slice()),
TurnInput::ResponseItem(_) | TurnInput::InterAgentCommunication(_) => None,
})
.flatten()
.filter(|item| {
matches!(item, UserInput::Image { .. } | UserInput::LocalImage { .. })
})
.count(),
submission_type: None,
ephemeral: thread_config.ephemeral,
session_source: thread_config.session_source,
model: turn_context.model_info.slug.clone(),
model_provider: turn_context.config.model_provider_id.clone(),
permission_profile: turn_context.permission_profile(),
#[allow(deprecated)]
permission_profile_cwd: turn_context.cwd.to_path_buf(),
reasoning_effort: turn_context.reasoning_effort.clone(),
reasoning_summary: Some(turn_context.reasoning_summary),
service_tier: turn_context
.config
.service_tier
.as_deref()
.and_then(ServiceTier::from_request_value),
approval_policy: turn_context.approval_policy.value(),
approvals_reviewer: turn_context.config.approvals_reviewer,
sandbox_network_access: turn_context.network_sandbox_policy().is_enabled(),
collaboration_mode: turn_context.mode,
personality: turn_context.personality,
workspace_kind: turn_context.turn_metadata_state.workspace_kind(),
is_first_turn,
});
}
#[instrument(level = "trace", skip_all)]
async fn run_pre_sampling_compact(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
client_session: &mut ModelClientSession,
) -> CodexResult<()> {
maybe_run_previous_model_inline_compact(sess, turn_context, client_session).await?;
let token_status =
super::context_window::context_window_token_status(sess.as_ref(), turn_context.as_ref())
.await;
if token_status.token_limit_reached {
let step_context = sess.capture_step_context(Arc::clone(turn_context)).await;
run_auto_compact(
sess,
step_context,
None,
client_session,
InitialContextInjection::DoNotInject,
CompactionReason::ContextLimit,
CompactionPhase::PreTurn,
)
.await?;
}
Ok(())
}
fn comp_hash_changed(previous: Option<&str>, current: Option<&str>) -> bool {
previous
.zip(current)
.is_some_and(|(previous, current)| previous != current)
}
async fn capture_current_model_fallback_step_context(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
previous_model: &str,
) -> Option<Arc<StepContext>> {
let uses_codex_backend = turn_context
.auth_manager
.as_deref()
.is_some_and(codex_login::AuthManager::current_auth_uses_codex_backend);
if !uses_codex_backend
|| !turn_context.provider.info().is_openai()
|| previous_model == turn_context.model_info.slug
{
return None;
}
Some(sess.capture_step_context(Arc::clone(turn_context)).await)
}
async fn maybe_run_previous_model_inline_compact(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
client_session: &mut ModelClientSession,
) -> CodexResult<()> {
let Some(previous_turn_settings) = sess.previous_turn_settings().await else {
return Ok(());
};
let should_compact_for_comp_hash_change = comp_hash_changed(
previous_turn_settings.comp_hash.as_deref(),
turn_context.model_info.comp_hash.as_deref(),
);
let previous_model = previous_turn_settings.model;
let previous_model_turn_context = Arc::new(
turn_context
.with_model(previous_model.clone(), &sess.services.models_manager)
.await,
);
if should_compact_for_comp_hash_change {
let step_context = sess
.capture_step_context(Arc::clone(&previous_model_turn_context))
.await;
let fallback_step_context = capture_current_model_fallback_step_context(
sess,
turn_context,
previous_model.as_str(),
)
.await;
run_auto_compact(
sess,
step_context,
fallback_step_context,
client_session,
InitialContextInjection::DoNotInject,
CompactionReason::CompHashChanged,
CompactionPhase::PreTurn,
)
.await?;
return Ok(());
}
let Some(old_context_window) = previous_model_turn_context.model_context_window() else {
return Ok(());
};
let Some(new_context_window) = turn_context.model_context_window() else {
return Ok(());
};
let active_context_tokens = sess.get_total_token_usage().await;
let previous_model_limit_reached = match turn_context
.config
.model_auto_compact_token_limit_scope
{
AutoCompactTokenLimitScope::Total => {
let new_auto_compact_limit = turn_context
.model_info
.auto_compact_token_limit()
.unwrap_or(i64::MAX);
active_context_tokens > new_auto_compact_limit
|| active_context_tokens >= new_context_window
}
AutoCompactTokenLimitScope::BodyAfterPrefix => active_context_tokens >= new_context_window,
};
let should_run = previous_model_limit_reached
&& previous_model_turn_context.model_info.slug != turn_context.model_info.slug
&& old_context_window > new_context_window;
if should_run {
let step_context = sess
.capture_step_context(Arc::clone(&previous_model_turn_context))
.await;
let fallback_step_context = capture_current_model_fallback_step_context(
sess,
turn_context,
previous_model.as_str(),
)
.await;
run_auto_compact(
sess,
step_context,
fallback_step_context,
client_session,
InitialContextInjection::DoNotInject,
CompactionReason::ModelDownshift,
CompactionPhase::PreTurn,
)
.await?;
}
Ok(())
}
#[instrument(
level = "trace",
skip_all,
fields(reason = ?reason, phase = ?phase)
)]
async fn run_auto_compact(
sess: &Arc<Session>,
step_context: Arc<StepContext>,
fallback_step_context: Option<Arc<StepContext>>,
client_session: &mut ModelClientSession,
initial_context_injection: InitialContextInjection,
reason: CompactionReason,
phase: CompactionPhase,
) -> CodexResult<()> {
let turn_context = &step_context.turn;
if turn_context.config.features.enabled(Feature::TokenBudget) {
crate::compact_token_budget::run_inline_auto_compact_task(
Arc::clone(sess),
step_context,
initial_context_injection,
)
.await?;
return Ok(());
}
if should_use_remote_compact_task(turn_context.provider.info()) {
if turn_context
.config
.features
.enabled(Feature::RemoteCompactionV2)
{
emit_compact_metric(
&sess.services.session_telemetry,
"remote_v2",
false,
);
run_inline_remote_auto_compact_task_v2(
Arc::clone(sess),
step_context,
fallback_step_context,
client_session,
initial_context_injection,
reason,
phase,
)
.await?;
return Ok(());
}
emit_compact_metric(
&sess.services.session_telemetry,
"remote",
false,
);
run_inline_remote_auto_compact_task(
Arc::clone(sess),
step_context,
fallback_step_context,
client_session.turn_state(),
initial_context_injection,
reason,
phase,
)
.await?;
} else {
emit_compact_metric(
&sess.services.session_telemetry,
"local",
false,
);
run_inline_auto_compact_task(
Arc::clone(sess),
Arc::clone(turn_context),
initial_context_injection,
reason,
phase,
)
.await?;
}
Ok(())
}
pub(super) fn collect_explicit_app_ids_from_skill_items(
skill_items: &[ResponseItem],
connectors: &[connectors::AppInfo],
skill_name_counts_lower: &HashMap<String, usize>,
) -> HashSet<String> {
if skill_items.is_empty() || connectors.is_empty() {
return HashSet::new();
}
let skill_messages = skill_items
.iter()
.filter_map(|item| match item {
ResponseItem::Message { content, .. } => {
content.iter().find_map(|content_item| match content_item {
ContentItem::InputText { text } => Some(text.clone()),
_ => None,
})
}
_ => None,
})
.collect::<Vec<String>>();
if skill_messages.is_empty() {
return HashSet::new();
}
let mentions = collect_tool_mentions_from_messages(&skill_messages);
let mention_names_lower = mentions
.plain_names
.iter()
.map(|name| name.to_ascii_lowercase())
.collect::<HashSet<String>>();
let mut connector_ids = mentions
.paths
.iter()
.filter(|path| tool_kind_for_path(path) == ToolMentionKind::App)
.filter_map(|path| app_id_from_path(path).map(str::to_string))
.collect::<HashSet<String>>();
let connector_slug_counts = build_connector_slug_counts(connectors);
for connector in connectors {
let slug = codex_connectors::metadata::connector_mention_slug(connector);
let connector_count = connector_slug_counts.get(&slug).copied().unwrap_or(0);
let skill_count = skill_name_counts_lower.get(&slug).copied().unwrap_or(0);
if connector_count == 1 && skill_count == 0 && mention_names_lower.contains(&slug) {
connector_ids.insert(connector.id.clone());
}
}
connector_ids
}
#[instrument(level = "trace", skip_all)]
pub(crate) fn build_prompt(
input: Vec<ResponseItem>,
router: &ToolRouter,
turn_context: &TurnContext,
base_instructions: BaseInstructions,
) -> Prompt {
Prompt {
input,
tools: router.model_visible_specs(),
parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls,
base_instructions,
output_schema: turn_context.final_output_json_schema.clone(),
output_schema_strict: !crate::guardian::is_guardian_reviewer_source(
&turn_context.session_source,
),
}
}
#[allow(clippy::too_many_arguments)]
#[allow(deprecated)]
#[instrument(level = "trace",
skip_all,
fields(
turn_id = %step_context.turn.sub_id,
model = %step_context.turn.model_info.slug,
cwd = %step_context.turn.cwd.display()
)
)]
async fn run_sampling_request(
sess: Arc<Session>,
step_context: Arc<StepContext>,
turn_store: Arc<codex_extension_api::ExtensionData>,
turn_diff_tracker: SharedTurnDiffTracker,
client_session: &mut ModelClientSession,
responses_metadata: &CodexResponsesMetadata,
input: Vec<ResponseItem>,
cancellation_token: CancellationToken,
) -> CodexResult<(SamplingRequestResult, Vec<ResponseItem>)> {
let turn_context = Arc::clone(&step_context.turn);
let router = built_tools(sess.as_ref(), step_context.as_ref(), &cancellation_token).await?;
let base_instructions = sess.get_base_instructions().await;
let tool_runtime = ToolCallRuntime::new(
Arc::clone(&router),
Arc::clone(&sess),
Arc::clone(&step_context),
Arc::clone(&turn_diff_tracker),
);
let _code_mode_worker = sess.services.code_mode_service.start_turn_worker(
&sess,
Arc::clone(&step_context),
Arc::clone(&router),
Arc::clone(&turn_diff_tracker),
);
let max_retries = turn_context.provider.info().stream_max_retries();
let mut retries = 0;
let mut initial_input = Some(input);
let mut original_input = None;
loop {
let prompt_input = if let Some(input) = initial_input.take() {
input
} else {
sess.clone_history()
.await
.for_prompt(&turn_context.model_info.input_modalities)
};
let prompt = build_prompt(
prompt_input,
router.as_ref(),
turn_context.as_ref(),
base_instructions.clone(),
);
let err = match try_run_sampling_request(
tool_runtime.clone(),
Arc::clone(&sess),
Arc::clone(&turn_context),
Arc::clone(&turn_store),
client_session,
responses_metadata,
Arc::clone(&turn_diff_tracker),
&prompt,
cancellation_token.child_token(),
)
.await
{
Ok(output) => {
return Ok((output, original_input.unwrap_or(prompt.input)));
}
Err(CodexErr::ContextWindowExceeded) => {
sess.set_total_tokens_full(&turn_context).await;
return Err(CodexErr::ContextWindowExceeded);
}
Err(CodexErr::UsageLimitReached(e)) => {
let rate_limits = e.rate_limits.clone();
if let Some(rate_limits) = rate_limits {
sess.update_rate_limits(&turn_context, *rate_limits).await;
}
return Err(CodexErr::UsageLimitReached(e));
}
Err(err) => err,
};
if original_input.is_none() {
original_input = Some(prompt.input);
}
if !err.is_retryable() {
return Err(err);
}
handle_retryable_response_stream_error(
&mut retries,
max_retries,
err,
client_session,
&sess,
&turn_context,
ResponsesStreamRequest::Sampling,
)
.await?;
turn_context.turn_timing_state.record_sampling_retry();
}
}
#[instrument(level = "trace",
skip_all,
fields(
turn_id = %step_context.turn.sub_id,
model = %step_context.turn.model_info.slug,
apps_enabled = step_context.turn.apps_enabled()
)
)]
pub(crate) async fn built_tools(
sess: &Session,
step_context: &StepContext,
cancellation_token: &CancellationToken,
) -> CodexResult<Arc<ToolRouter>> {
let turn_context = step_context.turn.as_ref();
let all_mcp_tools = step_context
.mcp_tools()
.or_cancel(cancellation_token)
.await?;
let loaded_plugins = sess
.services
.plugins_manager
.plugins_for_config(&turn_context.config.plugins_config_input())
.instrument(trace_span!("built_tools.load_plugins"))
.await;
let connector_snapshot = step_context.mcp.config().connector_snapshot.clone();
let apps_enabled = turn_context.apps_enabled();
let accessible_connectors =
apps_enabled.then(|| connectors::accessible_connectors_from_mcp_tools(all_mcp_tools));
let accessible_connectors_with_enabled_state =
accessible_connectors.as_ref().map(|connectors| {
connectors::with_app_enabled_state(connectors.clone(), &turn_context.config)
});
let connectors = if apps_enabled {
let connectors = codex_connectors::merge::merge_plugin_connectors_with_accessible(
connector_snapshot
.connector_ids()
.iter()
.map(|connector_id| connector_id.0.clone()),
accessible_connectors.clone().unwrap_or_default(),
);
Some(connectors::with_app_enabled_state(
connectors,
&turn_context.config,
))
} else {
None
};
let tool_suggest_is_enabled = tool_suggest_enabled(turn_context);
let auth = if tool_suggest_is_enabled {
sess.services.auth_manager.auth().await
} else {
None
};
let endpoint_recommended_plugin_candidates = if tool_suggest_is_enabled {
let plugins_config = turn_context.config.plugins_config_input();
sess.services
.plugins_manager
.recommended_plugin_candidates_for_config(RecommendedPluginCandidatesInput {
plugins_config: &plugins_config,
loaded_plugins: &loaded_plugins,
auth: auth.as_ref(),
disabled_tools: &turn_context.config.tool_suggest.disabled_tools,
app_server_client_name: turn_context.app_server_client_name.as_deref(),
})
.await
} else {
None
};
let tool_suggest_candidates =
if let Some(recommended_plugin_candidates) = endpoint_recommended_plugin_candidates {
Some(ToolSuggestCandidates {
tools: recommended_plugin_candidates,
presentation: ToolSuggestPresentation::RecommendationContext,
})
} else {
let loaded_plugin_app_connector_ids = connector_snapshot
.connector_ids()
.iter()
.map(|connector_id| connector_id.0.clone())
.collect::<Vec<_>>();
async {
if apps_enabled && tool_suggest_is_enabled {
if let Some(accessible_connectors) =
accessible_connectors_with_enabled_state.as_ref()
{
match connectors::list_tool_suggest_discoverable_tools_with_auth(
&turn_context.config,
sess.services.plugins_manager.as_ref(),
auth.as_ref(),
accessible_connectors.as_slice(),
&loaded_plugin_app_connector_ids,
)
.await
.map(|discoverable_tools| {
filter_request_plugin_install_discoverable_tools_for_client(
discoverable_tools,
turn_context.app_server_client_name.as_deref(),
)
}) {
Ok(discoverable_tools) if discoverable_tools.is_empty() => None,
Ok(discoverable_tools) => Some(ToolSuggestCandidates {
tools: discoverable_tools,
presentation: ToolSuggestPresentation::ListTool,
}),
Err(err) => {
warn!("failed to load discoverable tool suggestions: {err:#}");
None
}
}
} else {
None
}
} else {
None
}
}
.instrument(trace_span!("built_tools.load_discoverable_tools"))
.await
};
let mcp_tool_runtimes = build_mcp_tool_runtimes(
all_mcp_tools,
connectors.as_deref(),
&turn_context.config,
search_tool_enabled(turn_context),
);
Ok(Arc::new(ToolRouter::from_context(
step_context,
ToolRouterParams {
tool_runtimes: mcp_tool_runtimes,
tool_suggest_candidates,
extension_tool_executors: extension_tool_executors(sess, step_context),
dynamic_tools: turn_context.dynamic_tools.as_slice(),
},
&sess.services.tool_search_handler_cache,
)))
}
#[derive(Debug)]
struct SamplingRequestResult {
needs_follow_up: bool,
last_agent_message: Option<String>,
}
struct ProposedPlanItemState {
item_id: String,
started: bool,
completed: bool,
}
struct PlanModeStreamState {
pending_agent_message_items: HashMap<String, TurnItem>,
started_agent_message_items: HashSet<String>,
leading_whitespace_by_item: HashMap<String, String>,
plan_item_state: ProposedPlanItemState,
}
impl PlanModeStreamState {
fn new(turn_id: &str) -> Self {
Self {
pending_agent_message_items: HashMap::new(),
started_agent_message_items: HashSet::new(),
leading_whitespace_by_item: HashMap::new(),
plan_item_state: ProposedPlanItemState::new(turn_id),
}
}
}
#[derive(Debug, Default)]
pub(super) struct AssistantMessageStreamParsers {
plan_mode: bool,
parsers_by_item: HashMap<String, AssistantTextStreamParser>,
}
type ParsedAssistantTextDelta = AssistantTextChunk;
impl AssistantMessageStreamParsers {
pub(super) fn new(plan_mode: bool) -> Self {
Self {
plan_mode,
parsers_by_item: HashMap::new(),
}
}
fn parser_mut(&mut self, item_id: &str) -> &mut AssistantTextStreamParser {
let plan_mode = self.plan_mode;
self.parsers_by_item
.entry(item_id.to_string())
.or_insert_with(|| AssistantTextStreamParser::new(plan_mode))
}
pub(super) fn seed_item_text(&mut self, item_id: &str, text: &str) -> ParsedAssistantTextDelta {
if text.is_empty() {
return ParsedAssistantTextDelta::default();
}
self.parser_mut(item_id).push_str(text)
}
pub(super) fn parse_delta(&mut self, item_id: &str, delta: &str) -> ParsedAssistantTextDelta {
self.parser_mut(item_id).push_str(delta)
}
pub(super) fn finish_item(&mut self, item_id: &str) -> ParsedAssistantTextDelta {
let Some(mut parser) = self.parsers_by_item.remove(item_id) else {
return ParsedAssistantTextDelta::default();
};
parser.finish()
}
fn drain_finished(&mut self) -> Vec<(String, ParsedAssistantTextDelta)> {
let parsers_by_item = std::mem::take(&mut self.parsers_by_item);
parsers_by_item
.into_iter()
.map(|(item_id, mut parser)| (item_id, parser.finish()))
.collect()
}
}
impl ProposedPlanItemState {
fn new(turn_id: &str) -> Self {
Self {
item_id: format!("{turn_id}-plan"),
started: false,
completed: false,
}
}
async fn start(&mut self, sess: &Session, turn_context: &TurnContext) {
if self.started || self.completed {
return;
}
self.started = true;
let item = TurnItem::Plan(PlanItem {
id: self.item_id.clone(),
text: String::new(),
});
sess.emit_turn_item_started(turn_context, &item).await;
}
async fn push_delta(&mut self, sess: &Session, turn_context: &TurnContext, delta: &str) {
if self.completed {
return;
}
if delta.is_empty() {
return;
}
let event = PlanDeltaEvent {
thread_id: sess.thread_id.to_string(),
turn_id: turn_context.sub_id.clone(),
item_id: self.item_id.clone(),
delta: delta.to_string(),
};
sess.send_event(turn_context, EventMsg::PlanDelta(event))
.await;
}
async fn complete_with_text(
&mut self,
sess: &Session,
turn_context: &TurnContext,
text: String,
) {
if self.completed || !self.started {
return;
}
self.completed = true;
let item = TurnItem::Plan(PlanItem {
id: self.item_id.clone(),
text,
});
sess.emit_turn_item_completed(turn_context, item).await;
}
}
async fn maybe_emit_pending_agent_message_start(
sess: &Session,
turn_context: &TurnContext,
state: &mut PlanModeStreamState,
item_id: &str,
) {
if state.started_agent_message_items.contains(item_id) {
return;
}
if let Some(item) = state.pending_agent_message_items.remove(item_id) {
sess.emit_turn_item_started(turn_context, &item).await;
state
.started_agent_message_items
.insert(item_id.to_string());
}
}
pub(super) fn agent_message_text(item: &codex_protocol::items::AgentMessageItem) -> String {
item.content
.iter()
.map(|entry| match entry {
codex_protocol::items::AgentMessageContent::Text { text } => text.as_str(),
})
.collect()
}
pub(super) fn realtime_text_for_event(msg: &EventMsg) -> Option<(String, Option<MessagePhase>)> {
match msg {
EventMsg::AgentMessage(event) => Some((event.message.clone(), event.phase.clone())),
EventMsg::ItemCompleted(event) => match &event.item {
TurnItem::AgentMessage(item) => Some((agent_message_text(item), item.phase.clone())),
_ => None,
},
EventMsg::Error(_)
| EventMsg::Warning(_)
| EventMsg::GuardianWarning(_)
| EventMsg::RealtimeConversationStarted(_)
| EventMsg::RealtimeConversationSdp(_)
| EventMsg::RealtimeConversationRealtime(_)
| EventMsg::RealtimeConversationClosed(_)
| EventMsg::ModelReroute(_)
| EventMsg::ModelVerification(_)
| EventMsg::TurnModerationMetadata(_)
| EventMsg::SafetyBuffering(_)
| EventMsg::ContextCompacted(_)
| EventMsg::ThreadRolledBack(_)
| EventMsg::TurnStarted(_)
| EventMsg::ThreadSettingsApplied(_)
| EventMsg::TurnComplete(_)
| EventMsg::TokenCount(_)
| EventMsg::UserMessage(_)
| EventMsg::AgentReasoning(_)
| EventMsg::AgentReasoningRawContent(_)
| EventMsg::AgentReasoningSectionBreak(_)
| EventMsg::SessionConfigured(_)
| EventMsg::EnvironmentConnected(_)
| EventMsg::EnvironmentDisconnected(_)
| EventMsg::ThreadGoalUpdated(_)
| EventMsg::McpStartupUpdate(_)
| EventMsg::McpStartupComplete(_)
| EventMsg::McpToolCallBegin(_)
| EventMsg::McpToolCallEnd(_)
| EventMsg::WebSearchBegin(_)
| EventMsg::WebSearchEnd(_)
| EventMsg::ExecCommandBegin(_)
| EventMsg::ExecCommandOutputDelta(_)
| EventMsg::TerminalInteraction(_)
| EventMsg::ExecCommandEnd(_)
| EventMsg::PatchApplyBegin(_)
| EventMsg::PatchApplyUpdated(_)
| EventMsg::PatchApplyEnd(_)
| EventMsg::ImageGenerationBegin(_)
| EventMsg::ImageGenerationEnd(_)
| EventMsg::ViewImageToolCall(_)
| EventMsg::ExecApprovalRequest(_)
| EventMsg::RequestPermissions(_)
| EventMsg::RequestUserInput(_)
| EventMsg::DynamicToolCallRequest(_)
| EventMsg::DynamicToolCallResponse(_)
| EventMsg::GuardianAssessment(_)
| EventMsg::ElicitationRequest(_)
| EventMsg::ApplyPatchApprovalRequest(_)
| EventMsg::DeprecationNotice(_)
| EventMsg::StreamError(_)
| EventMsg::TurnDiff(_)
| EventMsg::RealtimeConversationListVoicesResponse(_)
| EventMsg::PlanUpdate(_)
| EventMsg::TurnAborted(_)
| EventMsg::ShutdownComplete
| EventMsg::EnteredReviewMode(_)
| EventMsg::ExitedReviewMode(_)
| EventMsg::RawResponseItem(_)
| EventMsg::RawResponseCompleted(_)
| EventMsg::ItemStarted(_)
| EventMsg::HookStarted(_)
| EventMsg::HookCompleted(_)
| EventMsg::AgentMessageContentDelta(_)
| EventMsg::PlanDelta(_)
| EventMsg::ReasoningContentDelta(_)
| EventMsg::ReasoningRawContentDelta(_)
| EventMsg::CollabAgentSpawnBegin(_)
| EventMsg::CollabAgentSpawnEnd(_)
| EventMsg::CollabAgentInteractionBegin(_)
| EventMsg::CollabAgentInteractionEnd(_)
| EventMsg::CollabWaitingBegin(_)
| EventMsg::CollabWaitingEnd(_)
| EventMsg::CollabCloseBegin(_)
| EventMsg::CollabCloseEnd(_)
| EventMsg::CollabResumeBegin(_)
| EventMsg::CollabResumeEnd(_)
| EventMsg::SubAgentActivity(_) => None,
}
}
async fn handle_plan_segments(
sess: &Session,
turn_context: &TurnContext,
state: &mut PlanModeStreamState,
item_id: &str,
segments: Vec<ProposedPlanSegment>,
) {
for segment in segments {
match segment {
ProposedPlanSegment::Normal(delta) => {
if delta.is_empty() {
continue;
}
let has_non_whitespace = delta.chars().any(|ch| !ch.is_whitespace());
if !has_non_whitespace && !state.started_agent_message_items.contains(item_id) {
let entry = state
.leading_whitespace_by_item
.entry(item_id.to_string())
.or_default();
entry.push_str(&delta);
continue;
}
let delta = if !state.started_agent_message_items.contains(item_id) {
if let Some(prefix) = state.leading_whitespace_by_item.remove(item_id) {
format!("{prefix}{delta}")
} else {
delta
}
} else {
delta
};
maybe_emit_pending_agent_message_start(sess, turn_context, state, item_id).await;
let event = AgentMessageContentDeltaEvent {
thread_id: sess.thread_id.to_string(),
turn_id: turn_context.sub_id.clone(),
item_id: item_id.to_string(),
delta,
};
sess.send_event(turn_context, EventMsg::AgentMessageContentDelta(event))
.await;
}
ProposedPlanSegment::ProposedPlanStart => {
if !state.plan_item_state.completed {
state.plan_item_state.start(sess, turn_context).await;
}
}
ProposedPlanSegment::ProposedPlanDelta(delta) => {
if !state.plan_item_state.completed {
if !state.plan_item_state.started {
state.plan_item_state.start(sess, turn_context).await;
}
state
.plan_item_state
.push_delta(sess, turn_context, &delta)
.await;
}
}
ProposedPlanSegment::ProposedPlanEnd => {}
}
}
}
async fn emit_streamed_assistant_text_delta(
sess: &Session,
turn_context: &TurnContext,
plan_mode_state: Option<&mut PlanModeStreamState>,
item_id: &str,
parsed: ParsedAssistantTextDelta,
) {
if parsed.is_empty() {
return;
}
if !parsed.citations.is_empty() {
let _citations = parsed.citations;
}
if let Some(state) = plan_mode_state {
if !parsed.plan_segments.is_empty() {
handle_plan_segments(sess, turn_context, state, item_id, parsed.plan_segments).await;
}
return;
}
if parsed.visible_text.is_empty() {
return;
}
let event = AgentMessageContentDeltaEvent {
thread_id: sess.thread_id.to_string(),
turn_id: turn_context.sub_id.clone(),
item_id: item_id.to_string(),
delta: parsed.visible_text,
};
sess.send_event(turn_context, EventMsg::AgentMessageContentDelta(event))
.await;
}
async fn flush_assistant_text_segments_for_item(
sess: &Session,
turn_context: &TurnContext,
plan_mode_state: Option<&mut PlanModeStreamState>,
parsers: &mut AssistantMessageStreamParsers,
item_id: &str,
) {
let parsed = parsers.finish_item(item_id);
emit_streamed_assistant_text_delta(sess, turn_context, plan_mode_state, item_id, parsed).await;
}
async fn flush_assistant_text_segments_all(
sess: &Session,
turn_context: &TurnContext,
mut plan_mode_state: Option<&mut PlanModeStreamState>,
parsers: &mut AssistantMessageStreamParsers,
) {
for (item_id, parsed) in parsers.drain_finished() {
emit_streamed_assistant_text_delta(
sess,
turn_context,
plan_mode_state.as_deref_mut(),
&item_id,
parsed,
)
.await;
}
}
async fn maybe_complete_plan_item_from_message(
sess: &Session,
turn_context: &TurnContext,
state: &mut PlanModeStreamState,
item: &ResponseItem,
) {
if let ResponseItem::Message { role, content, .. } = item
&& role == "assistant"
{
let mut text = String::new();
for entry in content {
if let ContentItem::OutputText { text: chunk } = entry {
text.push_str(chunk);
}
}
if let Some(plan_text) = extract_proposed_plan_text(&text) {
let (plan_text, _citations) = strip_citations(&plan_text);
if !state.plan_item_state.started {
state.plan_item_state.start(sess, turn_context).await;
}
state
.plan_item_state
.complete_with_text(sess, turn_context, plan_text)
.await;
}
}
}
async fn emit_agent_message_in_plan_mode(
sess: &Session,
turn_context: &TurnContext,
agent_message: codex_protocol::items::AgentMessageItem,
state: &mut PlanModeStreamState,
) {
let agent_message_id = agent_message.id.clone();
let text = agent_message_text(&agent_message);
if text.trim().is_empty() {
state.pending_agent_message_items.remove(&agent_message_id);
state.started_agent_message_items.remove(&agent_message_id);
return;
}
maybe_emit_pending_agent_message_start(sess, turn_context, state, &agent_message_id).await;
if !state
.started_agent_message_items
.contains(&agent_message_id)
{
let start_item = state
.pending_agent_message_items
.remove(&agent_message_id)
.unwrap_or_else(|| {
TurnItem::AgentMessage(codex_protocol::items::AgentMessageItem {
id: agent_message_id.clone(),
content: Vec::new(),
phase: None,
memory_citation: None,
})
});
sess.emit_turn_item_started(turn_context, &start_item).await;
state
.started_agent_message_items
.insert(agent_message_id.clone());
}
sess.emit_turn_item_completed(turn_context, TurnItem::AgentMessage(agent_message))
.await;
state.started_agent_message_items.remove(&agent_message_id);
}
async fn emit_turn_item_in_plan_mode(
sess: &Session,
turn_context: &TurnContext,
turn_item: TurnItem,
previously_active_item: Option<&TurnItem>,
state: &mut PlanModeStreamState,
) {
match turn_item {
TurnItem::AgentMessage(agent_message) => {
emit_agent_message_in_plan_mode(sess, turn_context, agent_message, state).await;
}
_ => {
if previously_active_item.is_none() {
sess.emit_turn_item_started(turn_context, &turn_item).await;
}
sess.emit_turn_item_completed(turn_context, turn_item).await;
}
}
}
async fn handle_assistant_item_done_in_plan_mode(
sess: &Session,
turn_context: &TurnContext,
turn_store: &codex_extension_api::ExtensionData,
item: &ResponseItem,
state: &mut PlanModeStreamState,
previously_active_item: Option<&TurnItem>,
last_agent_message: &mut Option<String>,
) -> bool {
if let ResponseItem::Message { role, .. } = item
&& role == "assistant"
{
maybe_complete_plan_item_from_message(sess, turn_context, state, item).await;
let mut finalized_facts = None;
if let Some(finalized_turn_item) = finalize_non_tool_response_item(
sess,
TurnItemContributorPolicy::Run(turn_store),
item,
true,
)
.await
{
finalized_facts = Some(finalized_turn_item.facts.clone());
emit_turn_item_in_plan_mode(
sess,
turn_context,
finalized_turn_item.turn_item,
previously_active_item,
state,
)
.await;
}
let final_last_agent_message = finalized_facts
.as_ref()
.and_then(|facts| facts.last_agent_message.clone());
record_completed_response_item_with_finalized_facts(
sess,
turn_context,
item,
finalized_facts.as_ref(),
)
.await;
if let Some(agent_message) = final_last_agent_message {
*last_agent_message = Some(agent_message);
}
return true;
}
false
}
#[instrument(level = "trace", skip_all)]
async fn drain_in_flight(
in_flight: &mut FuturesOrdered<BoxFuture<'static, CodexResult<ResponseInputItem>>>,
sess: Arc<Session>,
turn_context: Arc<TurnContext>,
) -> CodexResult<()> {
while let Some(res) = in_flight.next().await {
match res {
Ok(response_input) => {
let response_item = response_input.into();
sess.record_conversation_items(&turn_context, std::slice::from_ref(&response_item))
.await;
mark_thread_memory_mode_polluted_if_external_context(
sess.as_ref(),
turn_context.as_ref(),
&response_item,
)
.await;
}
Err(err) => {
error_or_panic(format!("in-flight tool future failed during drain: {err}"));
}
}
}
Ok(())
}
fn assign_missing_streamed_response_item_id(
item: &mut ResponseItem,
active_item: Option<&TurnItem>,
) {
if item.id().is_some_and(|id| !id.is_empty()) {
return;
}
let active_item_id = active_item
.map(|item| ResponseItemId::from_server(item.id()))
.filter(|item_id| !item_id.is_empty());
item.set_id(active_item_id);
Session::assign_missing_response_item_id(item);
}
#[allow(clippy::too_many_arguments)]
#[instrument(level = "trace",
skip_all,
fields(
turn_id = %turn_context.sub_id,
model = %turn_context.model_info.slug
)
)]
async fn try_run_sampling_request(
tool_runtime: ToolCallRuntime,
sess: Arc<Session>,
turn_context: Arc<TurnContext>,
turn_store: Arc<codex_extension_api::ExtensionData>,
client_session: &mut ModelClientSession,
responses_metadata: &CodexResponsesMetadata,
turn_diff_tracker: SharedTurnDiffTracker,
prompt: &Prompt,
cancellation_token: CancellationToken,
) -> CodexResult<SamplingRequestResult> {
feedback_tags!(
model = turn_context.model_info.slug.clone(),
approval_policy = turn_context.approval_policy.value(),
sandbox_policy = &turn_context.sandbox_policy(),
effort = turn_context.reasoning_effort,
auth_mode = sess.services.auth_manager.auth_mode(),
features = sess.features.enabled_features(),
);
let inference_trace = sess.services.rollout_thread_trace.inference_trace_context(
turn_context.sub_id.as_str(),
turn_context.model_info.slug.as_str(),
turn_context.provider.info().name.as_str(),
);
let sampling_timing_guard = turn_context.turn_timing_state.begin_sampling();
let uses_sequential_cutoff_reasoning_summaries = turn_context
.config
.features
.enabled(Feature::ConcurrentReasoningSummaries)
&& turn_context.provider.info().is_openai();
let mut stream = client_session
.stream(
prompt,
&turn_context.model_info,
&turn_context.session_telemetry,
turn_context.reasoning_effort.clone(),
turn_context.reasoning_summary,
turn_context.config.service_tier.clone(),
responses_metadata,
&inference_trace,
)
.instrument(trace_span!("stream_request"))
.or_cancel(&cancellation_token)
.await??;
let mut in_flight: FuturesOrdered<BoxFuture<'static, CodexResult<ResponseInputItem>>> =
FuturesOrdered::new();
let mut needs_follow_up = false;
let mut last_agent_message: Option<String> = None;
let mut active_item: Option<TurnItem> = None;
let mut active_tool_argument_diff_consumer: Option<(
String,
Box<dyn ToolArgumentDiffConsumer>,
)> = None;
let mut should_emit_turn_diff = false;
let mut should_emit_token_count = false;
let reasoning_effort = turn_context.effective_reasoning_effort_for_tracing();
let plan_mode = turn_context.mode == ModeKind::Plan;
let mut assistant_message_stream_parsers = AssistantMessageStreamParsers::new(plan_mode);
let mut plan_mode_state = plan_mode.then(|| PlanModeStreamState::new(&turn_context.sub_id));
let defer_streamed_turn_items_for_contributors =
!sess.services.extensions.turn_item_contributors().is_empty();
let mut active_item_is_streaming_to_client = false;
let receiving_span = trace_span!("receiving_stream");
let outcome: CodexResult<SamplingRequestResult> = loop {
let handle_responses = trace_span!(
parent: &receiving_span,
"handle_responses",
otel.name = field::Empty,
tool_name = field::Empty,
from = field::Empty,
codex.request.reasoning_effort = %reasoning_effort,
gen_ai.usage.input_tokens = field::Empty,
gen_ai.usage.cache_read.input_tokens = field::Empty,
gen_ai.usage.cache_write.input_tokens = field::Empty,
gen_ai.usage.output_tokens = field::Empty,
codex.usage.reasoning_output_tokens = field::Empty,
codex.usage.total_tokens = field::Empty,
);
let event = match stream
.next()
.instrument(trace_span!(parent: &handle_responses, "receiving"))
.or_cancel(&cancellation_token)
.await
{
Ok(event) => event,
Err(codex_async_utils::CancelErr::Cancelled) => break Err(CodexErr::TurnAborted),
};
let event = match event {
Some(Ok(event)) => event,
Some(Err(err)) => break Err(err),
None => {
break Err(CodexErr::Stream(
"stream closed before response.completed".into(),
None,
));
}
};
sess.services
.session_telemetry
.record_responses(&handle_responses, &event);
record_turn_ttft_metric(&turn_context, &event).await;
match event {
ResponseEvent::Created => {}
ResponseEvent::OutputItemDone(mut item) => {
assign_missing_streamed_response_item_id(&mut item, active_item.as_ref());
if let Some((_, mut consumer)) = active_tool_argument_diff_consumer.take()
&& let Ok(Some(event)) = consumer.finish()
{
sess.send_event(&turn_context, event).await;
}
let previously_active_item = active_item.take();
let previously_streamed_item = if active_item_is_streaming_to_client {
previously_active_item
} else {
None
};
active_item_is_streaming_to_client = false;
if let Some(previous) = previously_streamed_item.as_ref()
&& matches!(previous, TurnItem::AgentMessage(_))
{
let item_id = previous.id();
flush_assistant_text_segments_for_item(
&sess,
&turn_context,
plan_mode_state.as_mut(),
&mut assistant_message_stream_parsers,
&item_id,
)
.await;
}
if let Some(state) = plan_mode_state.as_mut()
&& handle_assistant_item_done_in_plan_mode(
&sess,
&turn_context,
turn_store.as_ref(),
&item,
state,
previously_streamed_item.as_ref(),
&mut last_agent_message,
)
.await
{
continue;
}
let mut ctx = HandleOutputCtx {
sess: sess.clone(),
turn_context: turn_context.clone(),
turn_store: Arc::clone(&turn_store),
tool_runtime: tool_runtime.clone(),
cancellation_token: cancellation_token.child_token(),
};
let preempt_for_mailbox_mail = match &item {
ResponseItem::Message { role, phase, .. } => {
role == "assistant" && matches!(phase, Some(MessagePhase::Commentary))
}
ResponseItem::Reasoning { .. } => true,
ResponseItem::AgentMessage { .. } => false,
ResponseItem::AdditionalTools { .. }
| ResponseItem::LocalShellCall { .. }
| ResponseItem::FunctionCall { .. }
| ResponseItem::ToolSearchCall { .. }
| ResponseItem::FunctionCallOutput { .. }
| ResponseItem::CustomToolCall { .. }
| ResponseItem::CustomToolCallOutput { .. }
| ResponseItem::ToolSearchOutput { .. }
| ResponseItem::WebSearchCall { .. }
| ResponseItem::ImageGenerationCall { .. }
| ResponseItem::Compaction { .. }
| ResponseItem::CompactionTrigger { .. }
| ResponseItem::ContextCompaction { .. }
| ResponseItem::Other => false,
};
let output_result =
match handle_output_item_done(&mut ctx, item, previously_streamed_item)
.instrument(handle_responses)
.await
{
Ok(output_result) => output_result,
Err(err) => break Err(err),
};
if let Some(tool_future) = output_result.tool_future {
in_flight.push_back(tool_future);
}
if let Some(agent_message) = output_result.last_agent_message {
last_agent_message = Some(agent_message);
}
needs_follow_up |= output_result.needs_follow_up;
if preempt_for_mailbox_mail && sess.input_queue.has_pending_mailbox_items().await {
break Ok(SamplingRequestResult {
needs_follow_up: true,
last_agent_message,
});
}
}
ResponseEvent::OutputItemAdded(mut item) => {
assign_missing_streamed_response_item_id(&mut item, None);
if let ResponseItem::CustomToolCall {
call_id,
name,
namespace,
..
} = &item
{
let tool_name = ToolName::new(namespace.clone(), name.as_str());
active_tool_argument_diff_consumer = tool_runtime
.create_diff_consumer(&tool_name)
.map(|consumer| (call_id.clone(), consumer));
} else if matches!(&item, ResponseItem::FunctionCall { .. }) {
active_tool_argument_diff_consumer = None;
}
if let Some(turn_item) = handle_non_tool_response_item(
sess.as_ref(),
TurnItemContributorPolicy::Skip,
&item,
plan_mode,
)
.await
{
let mut turn_item = turn_item;
let stream_item_to_client = !defer_streamed_turn_items_for_contributors;
let mut seeded_parsed: Option<ParsedAssistantTextDelta> = None;
let mut seeded_item_id: Option<String> = None;
if stream_item_to_client
&& matches!(turn_item, TurnItem::AgentMessage(_))
&& let Some(raw_text) = raw_assistant_output_text_from_item(&item)
{
let item_id = turn_item.id();
let mut seeded =
assistant_message_stream_parsers.seed_item_text(&item_id, &raw_text);
if let TurnItem::AgentMessage(agent_message) = &mut turn_item {
agent_message.content =
vec![codex_protocol::items::AgentMessageContent::Text {
text: if plan_mode {
String::new()
} else {
std::mem::take(&mut seeded.visible_text)
},
}];
}
seeded_parsed = plan_mode.then_some(seeded);
seeded_item_id = Some(item_id);
}
if stream_item_to_client {
if let Some(state) = plan_mode_state.as_mut()
&& matches!(turn_item, TurnItem::AgentMessage(_))
{
let item_id = turn_item.id();
state
.pending_agent_message_items
.insert(item_id, turn_item.clone());
} else {
sess.emit_turn_item_started(&turn_context, &turn_item).await;
}
if let (Some(state), Some(item_id), Some(parsed)) = (
plan_mode_state.as_mut(),
seeded_item_id.as_deref(),
seeded_parsed,
) {
emit_streamed_assistant_text_delta(
&sess,
&turn_context,
Some(state),
item_id,
parsed,
)
.await;
}
}
active_item = Some(turn_item);
active_item_is_streaming_to_client = stream_item_to_client;
}
}
ResponseEvent::ServerModel(server_model) => {
if !turn_context
.server_model_warning_emitted
.load(Ordering::Relaxed)
&& sess
.maybe_warn_on_server_model_mismatch(&turn_context, server_model)
.await
{
turn_context
.server_model_warning_emitted
.store(true, Ordering::Relaxed);
}
}
ResponseEvent::ModelVerifications(verifications) => {
if !turn_context
.model_verification_emitted
.swap(true, Ordering::Relaxed)
{
sess.emit_model_verification(&turn_context, verifications)
.await;
}
}
ResponseEvent::TurnModerationMetadata(metadata) => {
sess.emit_turn_moderation_metadata(&turn_context, metadata)
.await;
}
ResponseEvent::SafetyBuffering(buffering) => {
sess.send_event(
&turn_context,
EventMsg::SafetyBuffering(SafetyBufferingEvent {
model: turn_context.model_info.slug.clone(),
use_cases: buffering.use_cases,
reasons: buffering.reasons,
show_buffering_ui: buffering.show_buffering_ui,
faster_model: buffering.faster_model,
}),
)
.await;
}
ResponseEvent::ServerReasoningIncluded(included) => {
sess.set_server_reasoning_included(included).await;
}
ResponseEvent::RateLimits(snapshot) => {
sess.record_rate_limits_info(snapshot).await;
should_emit_token_count = true;
}
ResponseEvent::ModelsEtag(etag) => {
sess.services
.models_manager
.refresh_if_new_etag(etag, turn_context.config.http_client_factory())
.await;
}
ResponseEvent::Completed {
response_id,
token_usage,
end_turn,
} => {
flush_assistant_text_segments_all(
&sess,
&turn_context,
plan_mode_state.as_mut(),
&mut assistant_message_stream_parsers,
)
.await;
sess.send_event(
&turn_context,
EventMsg::RawResponseCompleted(RawResponseCompletedEvent {
response_id,
token_usage: token_usage.clone(),
}),
)
.await;
let budget_result = sess
.record_token_usage_info(&turn_context, token_usage.as_ref())
.await;
should_emit_token_count = true;
should_emit_turn_diff = true;
if let Err(err) = budget_result {
break Err(err);
}
if let Some(false) = end_turn {
needs_follow_up = true;
}
break Ok(SamplingRequestResult {
needs_follow_up,
last_agent_message,
});
}
ResponseEvent::OutputTextDelta(delta) => {
if let Some(active) = active_item.as_ref() {
if !active_item_is_streaming_to_client {
continue;
}
let item_id = active.id();
if matches!(active, TurnItem::AgentMessage(_)) {
let parsed = assistant_message_stream_parsers.parse_delta(&item_id, &delta);
emit_streamed_assistant_text_delta(
&sess,
&turn_context,
plan_mode_state.as_mut(),
&item_id,
parsed,
)
.await;
} else {
let event = AgentMessageContentDeltaEvent {
thread_id: sess.thread_id.to_string(),
turn_id: turn_context.sub_id.clone(),
item_id,
delta,
};
sess.send_event(&turn_context, EventMsg::AgentMessageContentDelta(event))
.await;
}
} else {
error_or_panic("OutputTextDelta without active item".to_string());
}
}
ResponseEvent::ToolCallInputDelta {
item_id: _,
call_id,
delta,
} => {
let Some((active_call_id, consumer)) = active_tool_argument_diff_consumer.as_mut()
else {
continue;
};
let call_id = match call_id {
Some(call_id) if call_id.as_str() != active_call_id.as_str() => continue,
Some(call_id) => call_id,
None => active_call_id.clone(),
};
if let Some(event) = consumer.consume_diff(turn_context.as_ref(), call_id, &delta) {
sess.send_event(&turn_context, event).await;
}
}
ResponseEvent::ReasoningSummaryDelta {
delta,
summary_index,
} => {
if uses_sequential_cutoff_reasoning_summaries {
continue;
}
if let Some(active) = active_item.as_ref() {
if !active_item_is_streaming_to_client {
continue;
}
let event = ReasoningContentDeltaEvent {
thread_id: sess.thread_id.to_string(),
turn_id: turn_context.sub_id.clone(),
item_id: active.id(),
delta,
summary_index,
};
sess.send_event(&turn_context, EventMsg::ReasoningContentDelta(event))
.await;
} else {
error_or_panic("ReasoningSummaryDelta without active item".to_string());
}
}
ResponseEvent::ReasoningSummaryPartAdded { summary_index } => {
if uses_sequential_cutoff_reasoning_summaries {
continue;
}
if let Some(active) = active_item.as_ref() {
if !active_item_is_streaming_to_client {
continue;
}
let event =
EventMsg::AgentReasoningSectionBreak(AgentReasoningSectionBreakEvent {
item_id: active.id(),
summary_index,
});
sess.send_event(&turn_context, event).await;
} else {
error_or_panic("ReasoningSummaryPartAdded without active item".to_string());
}
}
ResponseEvent::ReasoningSummaryDone {
item_id,
text,
summary_index,
} => {
if !uses_sequential_cutoff_reasoning_summaries {
continue;
}
let Some(active) = active_item.as_ref() else {
continue;
};
if !active_item_is_streaming_to_client || active.id() != item_id {
continue;
}
if summary_index > 0 {
sess.send_event(
&turn_context,
EventMsg::AgentReasoningSectionBreak(AgentReasoningSectionBreakEvent {
item_id: item_id.clone(),
summary_index,
}),
)
.await;
}
let event = ReasoningContentDeltaEvent {
thread_id: sess.thread_id.to_string(),
turn_id: turn_context.sub_id.clone(),
item_id,
delta: text,
summary_index,
};
sess.send_event(&turn_context, EventMsg::ReasoningContentDelta(event))
.await;
}
ResponseEvent::ReasoningContentDelta {
delta,
content_index,
} => {
if let Some(active) = active_item.as_ref() {
if !active_item_is_streaming_to_client {
continue;
}
let event = ReasoningRawContentDeltaEvent {
thread_id: sess.thread_id.to_string(),
turn_id: turn_context.sub_id.clone(),
item_id: active.id(),
delta,
content_index,
};
sess.send_event(&turn_context, EventMsg::ReasoningRawContentDelta(event))
.await;
} else {
error_or_panic("ReasoningRawContentDelta without active item".to_string());
}
}
}
};
drop(sampling_timing_guard);
flush_assistant_text_segments_all(
&sess,
&turn_context,
plan_mode_state.as_mut(),
&mut assistant_message_stream_parsers,
)
.await;
let tool_blocking_timing_guard = if in_flight.is_empty() {
None
} else {
Some(turn_context.turn_timing_state.begin_tool_blocking())
};
drain_in_flight(&mut in_flight, sess.clone(), turn_context.clone()).await?;
drop(tool_blocking_timing_guard);
if should_emit_token_count {
sess.send_token_count_event(&turn_context).await;
}
if cancellation_token.is_cancelled() {
return Err(CodexErr::TurnAborted);
}
if should_emit_turn_diff {
let unified_diff = {
let tracker = turn_diff_tracker.lock().await;
tracker.get_unified_diff()
};
if let Some(unified_diff) = unified_diff {
let msg = EventMsg::TurnDiff(TurnDiffEvent { unified_diff });
sess.clone().send_event(&turn_context, msg).await;
}
}
outcome
}
pub(crate) fn get_last_assistant_message_from_turn(responses: &[ResponseItem]) -> Option<String> {
for item in responses.iter().rev() {
if let Some(message) = last_assistant_message_from_item(item, false) {
return Some(message);
}
}
None
}
#[cfg(test)]
#[path = "turn_tests.rs"]
mod tests;