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_exposure; 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::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::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; /// Takes initial turn input and runs a loop where, at each sampling request, /// the model replies with either: /// /// - requested function calls /// - an assistant message /// /// While it is possible for the model to return multiple of these items in a /// single sampling request, in practice, we generally one item per sampling request: /// /// - If the model requests a function call, we execute it and send the output /// back to the model in the next sampling request. /// - If the model sends only an assistant message, we record it in the /// conversation history and consider the turn complete. /// pub(crate) async fn run_turn( sess: Arc, turn_context: Arc, turn_extension_data: Arc, input: Vec, prewarmed_client_session: Option, cancellation_token: CancellationToken, ) -> CodexResult> { let mut client_session = prewarmed_client_session.unwrap_or_else(|| sess.services.model_client.new_session()); // TODO(ccunningham): Pre-turn compaction runs before context updates and the // new user message are recorded. Estimate pending incoming items (context // diffs/full reinjection + user input) and trigger compaction preemptively // when they would push the thread over the compaction threshold. 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); } // run_turn owns the step used to seed context and make the first sampling request. let first_step_context = sess.capture_step_context(Arc::clone(&turn_context)).await; // Keep the exact model-visible state used by this turn and its inline compactions. let mut world_state = sess .record_context_updates_and_set_reference_context_item(first_step_context.as_ref()) .await; 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 = None; let mut stop_hook_active = false; // Although from the perspective of codex.rs, TurnDiffTracker has the lifecycle of a Task which contains // many turns, from the perspective of the user, it is a single turn. let display_roots = turn_diff_display_roots(turn_context.as_ref()).await; let turn_diff_tracker = Arc::new(tokio::sync::Mutex::new( TurnDiffTracker::with_environment_display_roots(display_roots), )); // `ModelClientSession` is turn-scoped and caches WebSocket + sticky routing state, so we reuse // one instance across retries within this turn. // Pending input is drained into history before building the next model request. // However, we defer that drain until after sampling in two cases: // 1. At the start of a turn, so the fresh turn input in `input` gets sampled first. // 2. After auto-compact, when model/tool continuation needs to resume before any steer. let mut next_step_context = Some(first_step_context); loop { // Note that pending_input would be something like a message the user // submitted through the UI while the model was running. Though the UI // may support this, the model might not. 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; // Capture once so context, advertised tools, and tool calls share one request view. 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?; if turn_context .config .features .enabled(Feature::DeferredExecutor) { world_state = sess .record_step_world_state_if_changed(&world_state, step_context.as_ref()) .await; } // Construct the input that we will send to the model. let sampling_request_input: Vec = 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; 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" ); super::token_budget::maybe_record( sess.as_ref(), turn_context.as_ref(), token_status.tokens_until_compaction, ) .await; // as long as compaction works well in getting us way below the token limit, we shouldn't worry about being in an infinite loop. if needs_follow_up && (sess.take_new_context_window_request().await || token_limit_reached) { if let Err(err) = run_auto_compact( &sess, Arc::clone(&step_context), &mut client_session, InitialContextInjection::BeforeLastUserMessage(Arc::clone(&world_state)), 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); } 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_conversation_items( &turn_context, std::slice::from_ref(&hook_prompt_message), ) .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()) => { { let mut state = sess.state.lock().await; error_or_panic( "Invalid image detected; sanitizing tool output to prevent poisoning", ); if state.history.replace_last_turn_images("Invalid image") { continue; } } 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(/*message_prefix*/ None)); sess.send_event(&turn_context, event).await; // let the user continue the conversation 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 { // TODO(anp): Migrate git-root discovery and diff display roots to PathUri so foreign // environment roots can participate without host-native conversion. 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, turn_context: &Arc, 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, step_context: &StepContext, input: &[TurnInput], cancellation_token: &CancellationToken, ) -> Option<(Vec, HashSet)> { let turn_context = step_context.turn.as_ref(); // Guardian input embeds the parent transcript as untrusted evidence. Do not interpret skill or // plugin mentions from that generated prompt as requests to inject additional instructions. 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::>(); 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; // Structured plugin:// mentions are resolved from the current session's // enabled plugins, then converted into turn-scoped guidance below. 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() { // Plugin mentions need raw MCP/app inventory even when app tools // are normally hidden so we can describe the plugin's currently // usable capabilities for this turn. 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, turn_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::(); 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 = 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::>(); 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::>(); 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 = 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, turn_context: &TurnContext, user_input: &[UserInput], cancellation_token: &CancellationToken, ) -> Option> { 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 .iter() .enumerate() .filter_map(|(index, environment)| { // TODO(anp): Migrate extension turn-input environments to PathUri so foreign cwd // values are not omitted from extension context. Some(TurnInputEnvironment { environment_id: environment.environment_id.clone(), cwd: environment.cwd().to_abs_path().ok()?.into_path_buf(), is_primary: index == 0, }) }) .collect::>(); 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(), ) .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.collaboration_mode.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, turn_context: &Arc, 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; // Compact if the configured auto-compaction budget or usable context window is exhausted. if token_status.token_limit_reached { // Pre-turn compaction runs before run_turn creates the normal sampling step. let step_context = sess.capture_step_context(Arc::clone(turn_context)).await; run_auto_compact( sess, step_context, client_session, InitialContextInjection::DoNotInject, CompactionReason::ContextLimit, CompactionPhase::PreTurn, ) .await?; } Ok(()) } /// Returns true only when both turns declare compaction compatibility hashes and they differ. /// A missing hash does not provide enough information to trigger compaction. fn comp_hash_changed(previous: Option<&str>, current: Option<&str>) -> bool { previous .zip(current) .is_some_and(|(previous, current)| previous != current) } /// Runs pre-sampling compaction against the previous model when its compaction compatibility /// hash changed or when switching to a smaller context-window model. /// /// Returns `Err(_)` only when compaction was attempted and failed. async fn maybe_run_previous_model_inline_compact( sess: &Arc, turn_context: &Arc, 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_turn_context = Arc::new( turn_context .with_model(previous_turn_settings.model, &sess.services.models_manager) .await, ); if should_compact_for_comp_hash_change { // This pre-turn request needs a step built from the previous model's turn context. let step_context = sess .capture_step_context(Arc::clone(&previous_model_turn_context)) .await; run_auto_compact( sess, 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 { // This pre-turn request needs a step built from the previous model's turn context. let step_context = sess .capture_step_context(Arc::clone(&previous_model_turn_context)) .await; run_auto_compact( sess, 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, step_context: Arc, 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) { // Compaction is the reset request, so force a new context window // instead of consuming a pending `new_context` tool request. 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", /*manual*/ false, ); run_inline_remote_auto_compact_task_v2( Arc::clone(sess), step_context, client_session, initial_context_injection, reason, phase, ) .await?; return Ok(()); } emit_compact_metric( &sess.services.session_telemetry, "remote", /*manual*/ false, ); run_inline_remote_auto_compact_task( Arc::clone(sess), step_context, client_session.turn_state(), initial_context_injection, reason, phase, ) .await?; } else { emit_compact_metric( &sess.services.session_telemetry, "local", /*manual*/ 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, ) -> HashSet { 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::>(); 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::>(); 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::>(); 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, 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, step_context: Arc, turn_store: Arc, turn_diff_tracker: SharedTurnDiffTracker, client_session: &mut ModelClientSession, responses_metadata: &CodexResponsesMetadata, input: Vec, cancellation_token: CancellationToken, ) -> CodexResult<(SamplingRequestResult, Vec)> { 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> { let turn_context = step_context.turn.as_ref(); let mcp_connection_manager = step_context.mcp.manager(); let has_mcp_servers = mcp_connection_manager.has_servers(); let all_mcp_tools = mcp_connection_manager .list_all_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::>(); 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_exposure = build_mcp_tool_exposure( &all_mcp_tools, connectors.as_deref(), &turn_context.config, search_tool_enabled(turn_context), ); let mcp_tools = has_mcp_servers.then_some(mcp_tool_exposure.direct_tools); let deferred_mcp_tools = mcp_tool_exposure.deferred_tools; Ok(Arc::new(ToolRouter::from_context( step_context, ToolRouterParams { mcp_tools, deferred_mcp_tools, tool_suggest_candidates, extension_tool_executors: extension_tool_executors(sess), 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, } /// Ephemeral per-response state for streaming a single proposed plan. /// This is intentionally not persisted or stored in session/state since it /// only exists while a response is actively streaming. The final plan text /// is extracted from the completed assistant message. /// Tracks a single proposed plan item across a streaming response. struct ProposedPlanItemState { item_id: String, started: bool, completed: bool, } /// Aggregated state used only while streaming a plan-mode response. /// Includes per-item parsers, deferred agent message bookkeeping, and the plan item lifecycle. struct PlanModeStreamState { /// Agent message items started by the model but deferred until we see non-plan text. pending_agent_message_items: HashMap, /// Agent message items whose start notification has been emitted. started_agent_message_items: HashSet, /// Leading whitespace buffered until we see non-whitespace text for an item. leading_whitespace_by_item: HashMap, /// Tracks plan item lifecycle while streaming plan output. 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, } 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; } } /// In plan mode we defer agent message starts until the parser emits non-plan /// text. The parser buffers each line until it can rule out a tag prefix, so /// plan-only outputs never show up as empty assistant messages. 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()); } } /// Agent messages are text-only today; concatenate all text entries. 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)> { 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::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::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, } } /// Split the stream into normal assistant text vs. proposed plan content. /// Normal text becomes AgentMessage deltas; plan content becomes PlanDelta + /// TurnItem::Plan. async fn handle_plan_segments( sess: &Session, turn_context: &TurnContext, state: &mut PlanModeStreamState, item_id: &str, segments: Vec, ) { 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() { // Citation extraction is intentionally local for now; we strip citations from display text // but do not yet surface them in protocol events. 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; } /// Flush buffered assistant text parser state when an assistant message item ends. 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; } /// Flush any remaining buffered assistant text parser state at response completion. 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; } } /// Emit completion for plan items by parsing the finalized assistant message. 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; } } } /// Emit a completed agent message in plan mode, respecting deferred starts. 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); } /// Emit completion for a plan-mode turn item, handling agent messages specially. 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; } } } /// Handle a completed assistant response item in plan mode, returning true if handled. 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, ) -> 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, turn_context, TurnItemContributorPolicy::Run(turn_store), item, /*plan_mode*/ 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>>, sess: Arc, turn_context: Arc, ) -> 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(()) } #[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, turn_context: Arc, turn_store: Arc, client_session: &mut ModelClientSession, responses_metadata: &CodexResponsesMetadata, turn_diff_tracker: SharedTurnDiffTracker, prompt: &Prompt, cancellation_token: CancellationToken, ) -> CodexResult { 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 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>> = FuturesOrdered::new(); let mut needs_follow_up = false; let mut last_agent_message: Option = None; let mut active_item: Option = None; let mut active_tool_argument_diff_consumer: Option<( String, Box, )> = 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.collaboration_mode.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 = 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.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(item) => { 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; // todo: remove before stabilizing multi-agent v2 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(item) => { if let ResponseItem::CustomToolCall { call_id, name, .. } = &item { let tool_name = ToolName::plain(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(), turn_context.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 = None; let mut seeded_item_id: Option = 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) => { // Update internal state with latest rate limits, but defer sending until // token usage is available to avoid duplicate TokenCount events. sess.record_rate_limits_info(snapshot).await; should_emit_token_count = true; } ResponseEvent::ModelsEtag(etag) => { // Update internal state with latest models etag sess.services.models_manager.refresh_if_new_etag(etag).await; } ResponseEvent::Completed { token_usage, end_turn, .. } => { flush_assistant_text_segments_all( &sess, &turn_context, plan_mode_state.as_mut(), &mut assistant_message_stream_parsers, ) .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) => { // In review child threads, suppress assistant text deltas; the // UI will show a selection popup from the final ReviewOutput. 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 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 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::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 { // A tool call such as request_user_input can intentionally pause the turn. Emit token // counts only after pending tools resolve so clients do not see progress events while the // turn is waiting on the user. This also needs to happen before returning cancellation so // token usage already recorded from the completed response is still persisted. 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 { for item in responses.iter().rev() { if let Some(message) = last_assistant_message_from_item(item, /*plan_mode*/ false) { return Some(message); } } None } #[cfg(test)] #[path = "turn_tests.rs"] mod tests;