mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
6c874f9b34
Adds the app-server v2 goal API on top of the persisted goal state from PR 1. ## Why Clients need a stable app-server surface for reading and controlling materialized thread goals before the model tools and TUI can use them. Goal changes also need to be observable by app-server clients, including clients that resume an existing thread. ## What changed - Added v2 `thread/goal/get`, `thread/goal/set`, and `thread/goal/clear` RPCs for materialized threads. - Added `thread/goal/updated` and `thread/goal/cleared` notifications so clients can keep local goal state in sync. - Added resume/snapshot wiring so reconnecting clients see the current goal state for a thread. - Added app-server handlers that reconcile persisted rollout state before direct goal mutations. - Updated the app-server README plus generated JSON and TypeScript schema fixtures for the new API surface. ## Verification - Added app-server v2 coverage for goal get/set/clear behavior, notification emission, resume snapshots, and non-local thread-store interactions.
2286 lines
86 KiB
Rust
2286 lines
86 KiB
Rust
use std::collections::HashMap;
|
||
use std::collections::HashSet;
|
||
use std::sync::Arc;
|
||
use std::sync::atomic::Ordering;
|
||
|
||
use crate::SkillInjections;
|
||
use crate::SkillLoadOutcome;
|
||
use crate::build_skill_injections;
|
||
use crate::client::ModelClientSession;
|
||
use crate::client_common::Prompt;
|
||
use crate::client_common::ResponseEvent;
|
||
use crate::collect_env_var_dependencies;
|
||
use crate::collect_explicit_skill_mentions;
|
||
use crate::compact::InitialContextInjection;
|
||
use crate::compact::collect_user_messages;
|
||
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::connectors;
|
||
use crate::context::ContextualUserFragment;
|
||
use crate::feedback_tags;
|
||
use crate::hook_runtime::PendingInputHookDisposition;
|
||
use crate::hook_runtime::emit_hook_completed_events;
|
||
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_pending_session_start_hooks;
|
||
use crate::hook_runtime::run_user_prompt_submit_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::parse_turn_item;
|
||
use crate::plugins::build_plugin_injections;
|
||
use crate::resolve_skill_dependencies_for_turn;
|
||
use crate::session::PreviousTurnSettings;
|
||
use crate::session::session::Session;
|
||
use crate::session::turn_context::TurnContext;
|
||
use crate::stream_events_utils::HandleOutputCtx;
|
||
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;
|
||
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::turn_diff_tracker::TurnDiffTracker;
|
||
use crate::turn_timing::record_turn_ttft_metric;
|
||
use crate::unavailable_tool::collect_unavailable_called_tools;
|
||
use crate::util::backoff;
|
||
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_features::Feature;
|
||
use codex_hooks::HookEvent;
|
||
use codex_hooks::HookEventAfterAgent;
|
||
use codex_hooks::HookPayload;
|
||
use codex_hooks::HookResult;
|
||
use codex_protocol::config_types::ModeKind;
|
||
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::UserMessageItem;
|
||
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::AskForApproval;
|
||
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::TurnDiffEvent;
|
||
use codex_protocol::protocol::WarningEvent;
|
||
use codex_protocol::user_input::UserInput;
|
||
use codex_tools::ResponsesApiNamespaceTool;
|
||
use codex_tools::ToolName;
|
||
use codex_tools::ToolSpec;
|
||
use codex_tools::filter_tool_suggest_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 a user message as 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.
|
||
///
|
||
#[expect(
|
||
clippy::await_holding_invalid_type,
|
||
reason = "turn execution must keep active-turn state transitions atomic"
|
||
)]
|
||
pub(crate) async fn run_turn(
|
||
sess: Arc<Session>,
|
||
turn_context: Arc<TurnContext>,
|
||
input: Vec<UserInput>,
|
||
prewarmed_client_session: Option<ModelClientSession>,
|
||
cancellation_token: CancellationToken,
|
||
) -> Option<String> {
|
||
if input.is_empty() && !sess.has_pending_input().await {
|
||
return None;
|
||
}
|
||
|
||
let model_info = turn_context.model_info.clone();
|
||
let auto_compact_limit = model_info.auto_compact_token_limit().unwrap_or(i64::MAX);
|
||
let mut prewarmed_client_session = prewarmed_client_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.
|
||
let pre_sampling_compacted = match run_pre_sampling_compact(&sess, &turn_context).await {
|
||
Ok(pre_sampling_compacted) => pre_sampling_compacted,
|
||
Err(_) => {
|
||
error!("Failed to run pre-sampling compact");
|
||
return None;
|
||
}
|
||
};
|
||
if pre_sampling_compacted && let Some(mut client_session) = prewarmed_client_session.take() {
|
||
client_session.reset_websocket_session();
|
||
}
|
||
|
||
let skills_outcome = Some(turn_context.turn_skills.outcome.as_ref());
|
||
|
||
sess.record_context_updates_and_set_reference_context_item(turn_context.as_ref())
|
||
.await;
|
||
|
||
let loaded_plugins = sess
|
||
.services
|
||
.plugins_manager
|
||
.plugins_for_config(&turn_context.config)
|
||
.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(&input, loaded_plugins.capability_summaries());
|
||
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 sess
|
||
.services
|
||
.mcp_connection_manager
|
||
.read()
|
||
.await
|
||
.list_all_tools()
|
||
.or_cancel(&cancellation_token)
|
||
.await
|
||
{
|
||
Ok(mcp_tools) => mcp_tools,
|
||
Err(_) if turn_context.apps_enabled() => return None,
|
||
Err(_) => HashMap::new(),
|
||
}
|
||
} else {
|
||
HashMap::new()
|
||
};
|
||
let available_connectors = if turn_context.apps_enabled() {
|
||
let connectors = codex_connectors::merge::merge_plugin_connectors_with_accessible(
|
||
loaded_plugins
|
||
.effective_apps()
|
||
.into_iter()
|
||
.map(|connector_id| connector_id.0),
|
||
connectors::accessible_connectors_from_mcp_tools(&mcp_tools),
|
||
);
|
||
connectors::with_app_enabled_state(connectors, &turn_context.config)
|
||
} else {
|
||
Vec::new()
|
||
};
|
||
let connector_slug_counts = build_connector_slug_counts(&available_connectors);
|
||
let skill_name_counts_lower = skills_outcome
|
||
.as_ref()
|
||
.map_or_else(HashMap::new, |outcome| {
|
||
build_skill_name_counts(&outcome.skills, &outcome.disabled_paths).1
|
||
});
|
||
let mentioned_skills = skills_outcome.as_ref().map_or_else(Vec::new, |outcome| {
|
||
collect_explicit_skill_mentions(
|
||
&input,
|
||
&outcome.skills,
|
||
&outcome.disabled_paths,
|
||
&connector_slug_counts,
|
||
)
|
||
});
|
||
let config = turn_context.config.clone();
|
||
if config
|
||
.features
|
||
.enabled(Feature::SkillEnvVarDependencyPrompt)
|
||
{
|
||
let env_var_dependencies = collect_env_var_dependencies(&mentioned_skills);
|
||
resolve_skill_dependencies_for_turn(&sess, &turn_context, &env_var_dependencies).await;
|
||
}
|
||
|
||
maybe_prompt_and_install_mcp_dependencies(
|
||
sess.as_ref(),
|
||
turn_context.as_ref(),
|
||
&cancellation_token,
|
||
&mentioned_skills,
|
||
)
|
||
.await;
|
||
|
||
let session_telemetry = turn_context.session_telemetry.clone();
|
||
let thread_id = sess.conversation_id.to_string();
|
||
let tracking = build_track_events_context(
|
||
turn_context.model_info.slug.clone(),
|
||
thread_id,
|
||
turn_context.sub_id.clone(),
|
||
);
|
||
let SkillInjections {
|
||
items: skill_injections,
|
||
warnings: skill_warnings,
|
||
} = build_skill_injections(
|
||
&mentioned_skills,
|
||
skills_outcome,
|
||
Some(&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 plugin_items =
|
||
build_plugin_injections(&mentioned_plugins, &mcp_tools, &available_connectors);
|
||
let mentioned_plugin_metadata = mentioned_plugins
|
||
.iter()
|
||
.filter_map(crate::plugins::PluginCapabilitySummary::telemetry_metadata)
|
||
.collect::<Vec<_>>();
|
||
|
||
let mut explicitly_enabled_connectors = collect_explicit_app_ids(&input);
|
||
explicitly_enabled_connectors.extend(collect_explicit_app_ids_from_skill_items(
|
||
&skill_items,
|
||
&available_connectors,
|
||
&skill_name_counts_lower,
|
||
));
|
||
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<_>>();
|
||
|
||
if run_pending_session_start_hooks(&sess, &turn_context).await {
|
||
return None;
|
||
}
|
||
let additional_contexts = if input.is_empty() {
|
||
Vec::new()
|
||
} else {
|
||
let initial_input_for_turn: ResponseInputItem = ResponseInputItem::from(input.clone());
|
||
let response_item: ResponseItem = initial_input_for_turn.clone().into();
|
||
let user_prompt_submit_outcome = run_user_prompt_submit_hooks(
|
||
&sess,
|
||
&turn_context,
|
||
UserMessageItem::new(&input).message(),
|
||
)
|
||
.await;
|
||
if user_prompt_submit_outcome.should_stop {
|
||
record_additional_contexts(
|
||
&sess,
|
||
&turn_context,
|
||
user_prompt_submit_outcome.additional_contexts,
|
||
)
|
||
.await;
|
||
return None;
|
||
}
|
||
sess.record_user_prompt_and_emit_turn_item(turn_context.as_ref(), &input, response_item)
|
||
.await;
|
||
user_prompt_submit_outcome.additional_contexts
|
||
};
|
||
sess.services
|
||
.analytics_events_client
|
||
.track_app_mentioned(tracking.clone(), mentioned_app_invocations);
|
||
for plugin in mentioned_plugin_metadata {
|
||
sess.services
|
||
.analytics_events_client
|
||
.track_plugin_used(tracking.clone(), plugin);
|
||
}
|
||
sess.merge_connector_selection(explicitly_enabled_connectors.clone())
|
||
.await;
|
||
record_additional_contexts(&sess, &turn_context, additional_contexts).await;
|
||
if !input.is_empty() {
|
||
// Track the previous-turn baseline from the regular user-turn path only so
|
||
// standalone tasks (compact/shell/review/undo) cannot suppress future
|
||
// model/realtime injections.
|
||
sess.set_previous_turn_settings(Some(PreviousTurnSettings {
|
||
model: turn_context.model_info.slug.clone(),
|
||
realtime_active: Some(turn_context.realtime_active),
|
||
}))
|
||
.await;
|
||
}
|
||
if !skill_items.is_empty() {
|
||
sess.record_conversation_items(&turn_context, &skill_items)
|
||
.await;
|
||
}
|
||
if !plugin_items.is_empty() {
|
||
sess.record_conversation_items(&turn_context, &plugin_items)
|
||
.await;
|
||
}
|
||
|
||
track_turn_resolved_config_analytics(&sess, &turn_context, &input).await;
|
||
|
||
let skills_outcome = Some(turn_context.turn_skills.outcome.as_ref());
|
||
sess.maybe_start_ghost_snapshot(Arc::clone(&turn_context), cancellation_token.child_token())
|
||
.await;
|
||
let mut last_agent_message: Option<String> = 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 turn_diff_tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new()));
|
||
|
||
// `ModelClientSession` is turn-scoped and caches WebSocket + sticky routing state, so we reuse
|
||
// one instance across retries within this turn.
|
||
let mut client_session =
|
||
prewarmed_client_session.unwrap_or_else(|| sess.services.model_client.new_session());
|
||
// 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 user prompt in `input` gets sampled first.
|
||
// 2. After auto-compact, when model/tool continuation needs to resume before any steer.
|
||
let mut can_drain_pending_input = input.is_empty();
|
||
|
||
loop {
|
||
if run_pending_session_start_hooks(&sess, &turn_context).await {
|
||
break;
|
||
}
|
||
|
||
// 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.get_pending_input().await
|
||
} else {
|
||
Vec::new()
|
||
};
|
||
|
||
let mut blocked_pending_input = false;
|
||
let mut blocked_pending_input_contexts = Vec::new();
|
||
let mut requeued_pending_input = false;
|
||
let mut accepted_pending_input = Vec::new();
|
||
if !pending_input.is_empty() {
|
||
let mut pending_input_iter = pending_input.into_iter();
|
||
while let Some(pending_input_item) = pending_input_iter.next() {
|
||
match inspect_pending_input(&sess, &turn_context, pending_input_item).await {
|
||
PendingInputHookDisposition::Accepted(pending_input) => {
|
||
accepted_pending_input.push(*pending_input);
|
||
}
|
||
PendingInputHookDisposition::Blocked {
|
||
additional_contexts,
|
||
} => {
|
||
let remaining_pending_input = pending_input_iter.collect::<Vec<_>>();
|
||
if !remaining_pending_input.is_empty() {
|
||
let _ = sess.prepend_pending_input(remaining_pending_input).await;
|
||
requeued_pending_input = true;
|
||
}
|
||
blocked_pending_input_contexts = additional_contexts;
|
||
blocked_pending_input = true;
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
let has_accepted_pending_input = !accepted_pending_input.is_empty();
|
||
for pending_input in accepted_pending_input {
|
||
record_pending_input(&sess, &turn_context, pending_input).await;
|
||
}
|
||
record_additional_contexts(&sess, &turn_context, blocked_pending_input_contexts).await;
|
||
|
||
if blocked_pending_input && !has_accepted_pending_input {
|
||
if requeued_pending_input {
|
||
continue;
|
||
}
|
||
break;
|
||
}
|
||
|
||
// Construct the input that we will send to the model.
|
||
let sampling_request_input: Vec<ResponseItem> = {
|
||
sess.clone_history()
|
||
.await
|
||
.for_prompt(&turn_context.model_info.input_modalities)
|
||
};
|
||
|
||
let sampling_request_input_messages = sampling_request_input
|
||
.iter()
|
||
.filter_map(|item| match parse_turn_item(item) {
|
||
Some(TurnItem::UserMessage(user_message)) => Some(user_message),
|
||
_ => None,
|
||
})
|
||
.map(|user_message| user_message.message())
|
||
.collect::<Vec<String>>();
|
||
let turn_metadata_header = turn_context.turn_metadata_state.current_header_value();
|
||
match run_sampling_request(
|
||
Arc::clone(&sess),
|
||
Arc::clone(&turn_context),
|
||
Arc::clone(&turn_diff_tracker),
|
||
&mut client_session,
|
||
turn_metadata_header.as_deref(),
|
||
sampling_request_input,
|
||
&explicitly_enabled_connectors,
|
||
skills_outcome,
|
||
cancellation_token.child_token(),
|
||
)
|
||
.await
|
||
{
|
||
Ok(sampling_request_output) => {
|
||
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 = sess.has_pending_input().await;
|
||
let needs_follow_up = model_needs_follow_up || has_pending_input;
|
||
let total_usage_tokens = sess.get_total_token_usage().await;
|
||
let token_limit_reached = total_usage_tokens >= auto_compact_limit;
|
||
|
||
let estimated_token_count =
|
||
sess.get_estimated_token_count(turn_context.as_ref()).await;
|
||
|
||
trace!(
|
||
turn_id = %turn_context.sub_id,
|
||
total_usage_tokens,
|
||
estimated_token_count = ?estimated_token_count,
|
||
auto_compact_limit,
|
||
token_limit_reached,
|
||
model_needs_follow_up,
|
||
has_pending_input,
|
||
needs_follow_up,
|
||
"post sampling token usage"
|
||
);
|
||
|
||
// 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 token_limit_reached && needs_follow_up {
|
||
if run_auto_compact(
|
||
&sess,
|
||
&turn_context,
|
||
InitialContextInjection::BeforeLastUserMessage,
|
||
CompactionReason::ContextLimit,
|
||
CompactionPhase::MidTurn,
|
||
)
|
||
.await
|
||
.is_err()
|
||
{
|
||
return None;
|
||
}
|
||
client_session.reset_websocket_session();
|
||
can_drain_pending_input = !model_needs_follow_up;
|
||
continue;
|
||
}
|
||
|
||
if !needs_follow_up {
|
||
last_agent_message = sampling_request_last_agent_message;
|
||
let stop_hook_permission_mode = match turn_context.approval_policy.value() {
|
||
AskForApproval::Never => "bypassPermissions",
|
||
AskForApproval::UnlessTrusted
|
||
| AskForApproval::OnFailure
|
||
| AskForApproval::OnRequest
|
||
| AskForApproval::Granular(_) => "default",
|
||
}
|
||
.to_string();
|
||
let stop_request = codex_hooks::StopRequest {
|
||
session_id: sess.conversation_id,
|
||
turn_id: turn_context.sub_id.clone(),
|
||
cwd: turn_context.cwd.clone(),
|
||
transcript_path: sess.hook_transcript_path().await,
|
||
model: turn_context.model_info.slug.clone(),
|
||
permission_mode: stop_hook_permission_mode,
|
||
stop_hook_active,
|
||
last_assistant_message: last_agent_message.clone(),
|
||
};
|
||
for run in sess.hooks().preview_stop(&stop_request) {
|
||
sess.send_event(
|
||
&turn_context,
|
||
EventMsg::HookStarted(codex_protocol::protocol::HookStartedEvent {
|
||
turn_id: Some(turn_context.sub_id.clone()),
|
||
run,
|
||
}),
|
||
)
|
||
.await;
|
||
}
|
||
let stop_outcome = sess.hooks().run_stop(stop_request).await;
|
||
emit_hook_completed_events(&sess, &turn_context, stop_outcome.hook_events)
|
||
.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;
|
||
}
|
||
let hook_outcomes = sess
|
||
.hooks()
|
||
.dispatch(HookPayload {
|
||
session_id: sess.conversation_id,
|
||
cwd: turn_context.cwd.clone(),
|
||
client: turn_context.app_server_client_name.clone(),
|
||
triggered_at: chrono::Utc::now(),
|
||
hook_event: HookEvent::AfterAgent {
|
||
event: HookEventAfterAgent {
|
||
thread_id: sess.conversation_id,
|
||
turn_id: turn_context.sub_id.clone(),
|
||
input_messages: sampling_request_input_messages,
|
||
last_assistant_message: last_agent_message.clone(),
|
||
},
|
||
},
|
||
})
|
||
.await;
|
||
|
||
let mut abort_message = None;
|
||
for hook_outcome in hook_outcomes {
|
||
let hook_name = hook_outcome.hook_name;
|
||
match hook_outcome.result {
|
||
HookResult::Success => {}
|
||
HookResult::FailedContinue(error) => {
|
||
warn!(
|
||
turn_id = %turn_context.sub_id,
|
||
hook_name = %hook_name,
|
||
error = %error,
|
||
"after_agent hook failed; continuing"
|
||
);
|
||
}
|
||
HookResult::FailedAbort(error) => {
|
||
let message = format!(
|
||
"after_agent hook '{hook_name}' failed and aborted turn completion: {error}"
|
||
);
|
||
warn!(
|
||
turn_id = %turn_context.sub_id,
|
||
hook_name = %hook_name,
|
||
error = %error,
|
||
"after_agent hook failed; aborting operation"
|
||
);
|
||
if abort_message.is_none() {
|
||
abort_message = Some(message);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
if let Some(message) = abort_message {
|
||
sess.send_event(
|
||
&turn_context,
|
||
EventMsg::Error(ErrorEvent {
|
||
message,
|
||
codex_error_info: None,
|
||
}),
|
||
)
|
||
.await;
|
||
return None;
|
||
}
|
||
break;
|
||
}
|
||
continue;
|
||
}
|
||
Err(CodexErr::TurnAborted) => {
|
||
// Aborted turn is reported via a different event.
|
||
break;
|
||
}
|
||
Err(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;
|
||
}
|
||
}
|
||
|
||
let event = EventMsg::Error(ErrorEvent {
|
||
message: "Invalid image in your last message. Please remove it and try again."
|
||
.to_string(),
|
||
codex_error_info: Some(CodexErrorInfo::BadRequest),
|
||
});
|
||
sess.send_event(&turn_context, event).await;
|
||
break;
|
||
}
|
||
Err(e) => {
|
||
info!("Turn error: {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;
|
||
}
|
||
}
|
||
}
|
||
|
||
last_agent_message
|
||
}
|
||
|
||
async fn track_turn_resolved_config_analytics(
|
||
sess: &Session,
|
||
turn_context: &TurnContext,
|
||
input: &[UserInput],
|
||
) {
|
||
if !sess.enabled(Feature::GeneralAnalytics) {
|
||
return;
|
||
}
|
||
|
||
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.conversation_id.to_string(),
|
||
num_input_images: input
|
||
.iter()
|
||
.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(),
|
||
sandbox_policy: turn_context.sandbox_policy.get().clone(),
|
||
reasoning_effort: turn_context.reasoning_effort,
|
||
reasoning_summary: Some(turn_context.reasoning_summary),
|
||
service_tier: turn_context.config.service_tier,
|
||
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,
|
||
is_first_turn,
|
||
});
|
||
}
|
||
|
||
async fn run_pre_sampling_compact(
|
||
sess: &Arc<Session>,
|
||
turn_context: &Arc<TurnContext>,
|
||
) -> CodexResult<bool> {
|
||
let total_usage_tokens_before_compaction = sess.get_total_token_usage().await;
|
||
let mut pre_sampling_compacted = maybe_run_previous_model_inline_compact(
|
||
sess,
|
||
turn_context,
|
||
total_usage_tokens_before_compaction,
|
||
)
|
||
.await?;
|
||
let total_usage_tokens = sess.get_total_token_usage().await;
|
||
let auto_compact_limit = turn_context
|
||
.model_info
|
||
.auto_compact_token_limit()
|
||
.unwrap_or(i64::MAX);
|
||
// Compact if the total usage tokens are greater than the auto compact limit
|
||
if total_usage_tokens >= auto_compact_limit {
|
||
run_auto_compact(
|
||
sess,
|
||
turn_context,
|
||
InitialContextInjection::DoNotInject,
|
||
CompactionReason::ContextLimit,
|
||
CompactionPhase::PreTurn,
|
||
)
|
||
.await?;
|
||
pre_sampling_compacted = true;
|
||
}
|
||
Ok(pre_sampling_compacted)
|
||
}
|
||
|
||
/// Runs pre-sampling compaction against the previous model when switching to a smaller
|
||
/// context-window model.
|
||
///
|
||
/// Returns `Ok(true)` when compaction ran successfully, `Ok(false)` when compaction was skipped
|
||
/// because the model/context-window preconditions were not met, and `Err(_)` only when compaction
|
||
/// was attempted and failed.
|
||
async fn maybe_run_previous_model_inline_compact(
|
||
sess: &Arc<Session>,
|
||
turn_context: &Arc<TurnContext>,
|
||
total_usage_tokens: i64,
|
||
) -> CodexResult<bool> {
|
||
let Some(previous_turn_settings) = sess.previous_turn_settings().await else {
|
||
return Ok(false);
|
||
};
|
||
let previous_model_turn_context = Arc::new(
|
||
turn_context
|
||
.with_model(previous_turn_settings.model, &sess.services.models_manager)
|
||
.await,
|
||
);
|
||
|
||
let Some(old_context_window) = previous_model_turn_context.model_context_window() else {
|
||
return Ok(false);
|
||
};
|
||
let Some(new_context_window) = turn_context.model_context_window() else {
|
||
return Ok(false);
|
||
};
|
||
let new_auto_compact_limit = turn_context
|
||
.model_info
|
||
.auto_compact_token_limit()
|
||
.unwrap_or(i64::MAX);
|
||
let should_run = total_usage_tokens > new_auto_compact_limit
|
||
&& previous_model_turn_context.model_info.slug != turn_context.model_info.slug
|
||
&& old_context_window > new_context_window;
|
||
if should_run {
|
||
run_auto_compact(
|
||
sess,
|
||
&previous_model_turn_context,
|
||
InitialContextInjection::DoNotInject,
|
||
CompactionReason::ModelDownshift,
|
||
CompactionPhase::PreTurn,
|
||
)
|
||
.await?;
|
||
return Ok(true);
|
||
}
|
||
Ok(false)
|
||
}
|
||
|
||
async fn run_auto_compact(
|
||
sess: &Arc<Session>,
|
||
turn_context: &Arc<TurnContext>,
|
||
initial_context_injection: InitialContextInjection,
|
||
reason: CompactionReason,
|
||
phase: CompactionPhase,
|
||
) -> CodexResult<()> {
|
||
if should_use_remote_compact_task(turn_context.provider.info()) {
|
||
run_inline_remote_auto_compact_task(
|
||
Arc::clone(sess),
|
||
Arc::clone(turn_context),
|
||
initial_context_injection,
|
||
reason,
|
||
phase,
|
||
)
|
||
.await?;
|
||
} else {
|
||
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
|
||
}
|
||
|
||
pub(super) fn filter_connectors_for_input(
|
||
connectors: &[connectors::AppInfo],
|
||
input: &[ResponseItem],
|
||
explicitly_enabled_connectors: &HashSet<String>,
|
||
skill_name_counts_lower: &HashMap<String, usize>,
|
||
) -> Vec<connectors::AppInfo> {
|
||
let connectors: Vec<connectors::AppInfo> = connectors
|
||
.iter()
|
||
.filter(|connector| connector.is_enabled)
|
||
.cloned()
|
||
.collect::<Vec<_>>();
|
||
if connectors.is_empty() {
|
||
return Vec::new();
|
||
}
|
||
|
||
let user_messages = collect_user_messages(input);
|
||
if user_messages.is_empty() && explicitly_enabled_connectors.is_empty() {
|
||
return Vec::new();
|
||
}
|
||
|
||
let mentions = collect_tool_mentions_from_messages(&user_messages);
|
||
let mention_names_lower = mentions
|
||
.plain_names
|
||
.iter()
|
||
.map(|name| name.to_ascii_lowercase())
|
||
.collect::<HashSet<String>>();
|
||
|
||
let connector_slug_counts = build_connector_slug_counts(&connectors);
|
||
let mut allowed_connector_ids = explicitly_enabled_connectors.clone();
|
||
for path in mentions
|
||
.paths
|
||
.iter()
|
||
.filter(|path| tool_kind_for_path(path) == ToolMentionKind::App)
|
||
{
|
||
if let Some(connector_id) = app_id_from_path(path) {
|
||
allowed_connector_ids.insert(connector_id.to_string());
|
||
}
|
||
}
|
||
|
||
connectors
|
||
.into_iter()
|
||
.filter(|connector| {
|
||
connector_inserted_in_messages(
|
||
connector,
|
||
&mention_names_lower,
|
||
&allowed_connector_ids,
|
||
&connector_slug_counts,
|
||
skill_name_counts_lower,
|
||
)
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
fn connector_inserted_in_messages(
|
||
connector: &connectors::AppInfo,
|
||
mention_names_lower: &HashSet<String>,
|
||
allowed_connector_ids: &HashSet<String>,
|
||
connector_slug_counts: &HashMap<String, usize>,
|
||
skill_name_counts_lower: &HashMap<String, usize>,
|
||
) -> bool {
|
||
if allowed_connector_ids.contains(&connector.id) {
|
||
return true;
|
||
}
|
||
|
||
let mention_slug = codex_connectors::metadata::connector_mention_slug(connector);
|
||
let connector_count = connector_slug_counts
|
||
.get(&mention_slug)
|
||
.copied()
|
||
.unwrap_or(0);
|
||
let skill_count = skill_name_counts_lower
|
||
.get(&mention_slug)
|
||
.copied()
|
||
.unwrap_or(0);
|
||
connector_count == 1 && skill_count == 0 && mention_names_lower.contains(&mention_slug)
|
||
}
|
||
|
||
pub(crate) fn build_prompt(
|
||
input: Vec<ResponseItem>,
|
||
router: &ToolRouter,
|
||
turn_context: &TurnContext,
|
||
base_instructions: BaseInstructions,
|
||
) -> Prompt {
|
||
let deferred_dynamic_tools = turn_context
|
||
.dynamic_tools
|
||
.iter()
|
||
.filter(|tool| tool.defer_loading)
|
||
.map(|tool| ToolName::new(tool.namespace.clone(), tool.name.clone()))
|
||
.collect::<HashSet<_>>();
|
||
let tools = if deferred_dynamic_tools.is_empty() {
|
||
router.model_visible_specs()
|
||
} else {
|
||
router
|
||
.model_visible_specs()
|
||
.into_iter()
|
||
.filter_map(|spec| filter_deferred_dynamic_tool_spec(spec, &deferred_dynamic_tools))
|
||
.collect()
|
||
};
|
||
|
||
Prompt {
|
||
input,
|
||
tools,
|
||
parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls,
|
||
base_instructions,
|
||
personality: turn_context.personality,
|
||
output_schema: turn_context.final_output_json_schema.clone(),
|
||
output_schema_strict: !crate::guardian::is_guardian_reviewer_source(
|
||
&turn_context.session_source,
|
||
),
|
||
}
|
||
}
|
||
|
||
fn filter_deferred_dynamic_tool_spec(
|
||
spec: ToolSpec,
|
||
deferred_dynamic_tools: &HashSet<ToolName>,
|
||
) -> Option<ToolSpec> {
|
||
match spec {
|
||
ToolSpec::Function(tool) => {
|
||
if deferred_dynamic_tools.contains(&ToolName::plain(tool.name.as_str())) {
|
||
None
|
||
} else {
|
||
Some(ToolSpec::Function(tool))
|
||
}
|
||
}
|
||
ToolSpec::Namespace(mut namespace) => {
|
||
let namespace_name = namespace.name.clone();
|
||
namespace.tools.retain(|tool| match tool {
|
||
ResponsesApiNamespaceTool::Function(tool) => !deferred_dynamic_tools.contains(
|
||
&ToolName::namespaced(namespace_name.as_str(), tool.name.as_str()),
|
||
),
|
||
});
|
||
if namespace.tools.is_empty() {
|
||
None
|
||
} else {
|
||
Some(ToolSpec::Namespace(namespace))
|
||
}
|
||
}
|
||
spec => Some(spec),
|
||
}
|
||
}
|
||
|
||
#[allow(clippy::too_many_arguments)]
|
||
#[instrument(level = "trace",
|
||
skip_all,
|
||
fields(
|
||
turn_id = %turn_context.sub_id,
|
||
model = %turn_context.model_info.slug,
|
||
cwd = %turn_context.cwd.display()
|
||
)
|
||
)]
|
||
async fn run_sampling_request(
|
||
sess: Arc<Session>,
|
||
turn_context: Arc<TurnContext>,
|
||
turn_diff_tracker: SharedTurnDiffTracker,
|
||
client_session: &mut ModelClientSession,
|
||
turn_metadata_header: Option<&str>,
|
||
input: Vec<ResponseItem>,
|
||
explicitly_enabled_connectors: &HashSet<String>,
|
||
skills_outcome: Option<&SkillLoadOutcome>,
|
||
cancellation_token: CancellationToken,
|
||
) -> CodexResult<SamplingRequestResult> {
|
||
let router = built_tools(
|
||
sess.as_ref(),
|
||
turn_context.as_ref(),
|
||
&input,
|
||
explicitly_enabled_connectors,
|
||
skills_outcome,
|
||
&cancellation_token,
|
||
)
|
||
.await?;
|
||
|
||
let base_instructions = sess.get_base_instructions().await;
|
||
|
||
let tool_runtime = ToolCallRuntime::new(
|
||
Arc::clone(&router),
|
||
Arc::clone(&sess),
|
||
Arc::clone(&turn_context),
|
||
Arc::clone(&turn_diff_tracker),
|
||
);
|
||
let _code_mode_worker = sess
|
||
.services
|
||
.code_mode_service
|
||
.start_turn_worker(
|
||
&sess,
|
||
&turn_context,
|
||
Arc::clone(&router),
|
||
Arc::clone(&turn_diff_tracker),
|
||
)
|
||
.await;
|
||
let mut retries = 0;
|
||
let mut initial_input = Some(input);
|
||
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),
|
||
client_session,
|
||
turn_metadata_header,
|
||
Arc::clone(&turn_diff_tracker),
|
||
&prompt,
|
||
cancellation_token.child_token(),
|
||
)
|
||
.await
|
||
{
|
||
Ok(output) => {
|
||
return Ok(output);
|
||
}
|
||
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 !err.is_retryable() {
|
||
return Err(err);
|
||
}
|
||
|
||
// Use the configured provider-specific stream retry budget.
|
||
let max_retries = turn_context.provider.info().stream_max_retries();
|
||
if retries >= max_retries
|
||
&& client_session.try_switch_fallback_transport(
|
||
&turn_context.session_telemetry,
|
||
&turn_context.model_info,
|
||
)
|
||
{
|
||
sess.send_event(
|
||
&turn_context,
|
||
EventMsg::Warning(WarningEvent {
|
||
message: format!("Falling back from WebSockets to HTTPS transport. {err:#}"),
|
||
}),
|
||
)
|
||
.await;
|
||
retries = 0;
|
||
continue;
|
||
}
|
||
if retries < max_retries {
|
||
retries += 1;
|
||
let delay = match &err {
|
||
CodexErr::Stream(_, requested_delay) => {
|
||
requested_delay.unwrap_or_else(|| backoff(retries))
|
||
}
|
||
_ => backoff(retries),
|
||
};
|
||
warn!(
|
||
"stream disconnected - retrying sampling request ({retries}/{max_retries} in {delay:?})...",
|
||
);
|
||
|
||
// In release builds, hide the first websocket retry notification to reduce noisy
|
||
// transient reconnect messages. In debug builds, keep full visibility for diagnosis.
|
||
let report_error = retries > 1
|
||
|| cfg!(debug_assertions)
|
||
|| !sess.services.model_client.responses_websocket_enabled();
|
||
if report_error {
|
||
// Surface retry information to any UI/front‑end so the
|
||
// user understands what is happening instead of staring
|
||
// at a seemingly frozen screen.
|
||
sess.notify_stream_error(
|
||
&turn_context,
|
||
format!("Reconnecting... {retries}/{max_retries}"),
|
||
err,
|
||
)
|
||
.await;
|
||
}
|
||
tokio::time::sleep(delay).await;
|
||
} else {
|
||
return Err(err);
|
||
}
|
||
}
|
||
}
|
||
|
||
#[expect(
|
||
clippy::await_holding_invalid_type,
|
||
reason = "tool router construction reads through the session-owned manager guard"
|
||
)]
|
||
pub(crate) async fn built_tools(
|
||
sess: &Session,
|
||
turn_context: &TurnContext,
|
||
input: &[ResponseItem],
|
||
explicitly_enabled_connectors: &HashSet<String>,
|
||
skills_outcome: Option<&SkillLoadOutcome>,
|
||
cancellation_token: &CancellationToken,
|
||
) -> CodexResult<Arc<ToolRouter>> {
|
||
let mcp_connection_manager = sess.services.mcp_connection_manager.read().await;
|
||
let has_mcp_servers = mcp_connection_manager.has_servers();
|
||
let all_mcp_tools = mcp_connection_manager
|
||
.list_all_tools()
|
||
.or_cancel(cancellation_token)
|
||
.await?;
|
||
drop(mcp_connection_manager);
|
||
let loaded_plugins = sess
|
||
.services
|
||
.plugins_manager
|
||
.plugins_for_config(&turn_context.config)
|
||
.await;
|
||
|
||
let mut effective_explicitly_enabled_connectors = explicitly_enabled_connectors.clone();
|
||
effective_explicitly_enabled_connectors.extend(sess.get_connector_selection().await);
|
||
|
||
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(
|
||
loaded_plugins
|
||
.effective_apps()
|
||
.into_iter()
|
||
.map(|connector_id| connector_id.0),
|
||
accessible_connectors.clone().unwrap_or_default(),
|
||
);
|
||
Some(connectors::with_app_enabled_state(
|
||
connectors,
|
||
&turn_context.config,
|
||
))
|
||
} else {
|
||
None
|
||
};
|
||
let auth = sess.services.auth_manager.auth().await;
|
||
let discoverable_tools = if apps_enabled && turn_context.tools_config.tool_suggest {
|
||
if let Some(accessible_connectors) = accessible_connectors_with_enabled_state.as_ref() {
|
||
match connectors::list_tool_suggest_discoverable_tools_with_auth(
|
||
&turn_context.config,
|
||
auth.as_ref(),
|
||
accessible_connectors.as_slice(),
|
||
)
|
||
.await
|
||
.map(|discoverable_tools| {
|
||
filter_tool_suggest_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(discoverable_tools),
|
||
Err(err) => {
|
||
warn!("failed to load discoverable tool suggestions: {err:#}");
|
||
None
|
||
}
|
||
}
|
||
} else {
|
||
None
|
||
}
|
||
} else {
|
||
None
|
||
};
|
||
|
||
let explicitly_enabled = if let Some(connectors) = connectors.as_ref() {
|
||
let skill_name_counts_lower = skills_outcome.map_or_else(HashMap::new, |outcome| {
|
||
build_skill_name_counts(&outcome.skills, &outcome.disabled_paths).1
|
||
});
|
||
|
||
filter_connectors_for_input(
|
||
connectors,
|
||
input,
|
||
&effective_explicitly_enabled_connectors,
|
||
&skill_name_counts_lower,
|
||
)
|
||
} else {
|
||
Vec::new()
|
||
};
|
||
let mcp_tool_exposure = build_mcp_tool_exposure(
|
||
&all_mcp_tools,
|
||
connectors.as_deref(),
|
||
explicitly_enabled.as_slice(),
|
||
&turn_context.config,
|
||
&turn_context.tools_config,
|
||
);
|
||
let mcp_tools = has_mcp_servers.then_some(mcp_tool_exposure.direct_tools);
|
||
let deferred_mcp_tools = mcp_tool_exposure.deferred_tools;
|
||
let unavailable_called_tools = if turn_context
|
||
.config
|
||
.features
|
||
.enabled(Feature::UnavailableDummyTools)
|
||
{
|
||
let exposed_tool_names = mcp_tools
|
||
.iter()
|
||
.chain(deferred_mcp_tools.iter())
|
||
.flat_map(|tools| tools.keys().map(String::as_str))
|
||
.collect::<HashSet<_>>();
|
||
collect_unavailable_called_tools(input, &exposed_tool_names)
|
||
} else {
|
||
Vec::new()
|
||
};
|
||
|
||
let parallel_mcp_server_names = turn_context
|
||
.config
|
||
.mcp_servers
|
||
.get()
|
||
.iter()
|
||
.filter_map(|(server_name, server_config)| {
|
||
server_config
|
||
.supports_parallel_tool_calls
|
||
.then_some(server_name.clone())
|
||
})
|
||
.collect::<HashSet<_>>();
|
||
|
||
Ok(Arc::new(ToolRouter::from_config(
|
||
&turn_context.tools_config,
|
||
ToolRouterParams {
|
||
mcp_tools,
|
||
deferred_mcp_tools,
|
||
unavailable_called_tools,
|
||
parallel_mcp_server_names,
|
||
discoverable_tools,
|
||
dynamic_tools: turn_context.dynamic_tools.as_slice(),
|
||
},
|
||
)))
|
||
}
|
||
|
||
#[derive(Debug)]
|
||
struct SamplingRequestResult {
|
||
needs_follow_up: bool,
|
||
last_agent_message: Option<String>,
|
||
}
|
||
|
||
/// 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<String, TurnItem>,
|
||
/// Agent message items whose start notification has been emitted.
|
||
started_agent_message_items: HashSet<String>,
|
||
/// Leading whitespace buffered until we see non-whitespace text for an item.
|
||
leading_whitespace_by_item: HashMap<String, String>,
|
||
/// 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<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.conversation_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> {
|
||
match msg {
|
||
EventMsg::AgentMessage(event) => Some(event.message.clone()),
|
||
EventMsg::ItemCompleted(event) => match &event.item {
|
||
TurnItem::AgentMessage(item) => Some(agent_message_text(item)),
|
||
_ => None,
|
||
},
|
||
EventMsg::Error(_)
|
||
| EventMsg::Warning(_)
|
||
| EventMsg::GuardianWarning(_)
|
||
| EventMsg::RealtimeConversationStarted(_)
|
||
| EventMsg::RealtimeConversationSdp(_)
|
||
| EventMsg::RealtimeConversationRealtime(_)
|
||
| EventMsg::RealtimeConversationClosed(_)
|
||
| EventMsg::ModelReroute(_)
|
||
| EventMsg::ModelVerification(_)
|
||
| EventMsg::ContextCompacted(_)
|
||
| EventMsg::ThreadRolledBack(_)
|
||
| EventMsg::TurnStarted(_)
|
||
| EventMsg::TurnComplete(_)
|
||
| EventMsg::TokenCount(_)
|
||
| EventMsg::UserMessage(_)
|
||
| EventMsg::AgentMessageDelta(_)
|
||
| EventMsg::AgentReasoning(_)
|
||
| EventMsg::AgentReasoningDelta(_)
|
||
| EventMsg::AgentReasoningRawContent(_)
|
||
| EventMsg::AgentReasoningRawContentDelta(_)
|
||
| EventMsg::AgentReasoningSectionBreak(_)
|
||
| EventMsg::SessionConfigured(_)
|
||
| EventMsg::ThreadNameUpdated(_)
|
||
| 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::ViewImageToolCall(_)
|
||
| EventMsg::ImageGenerationBegin(_)
|
||
| EventMsg::ImageGenerationEnd(_)
|
||
| EventMsg::ExecApprovalRequest(_)
|
||
| EventMsg::RequestPermissions(_)
|
||
| EventMsg::RequestUserInput(_)
|
||
| EventMsg::DynamicToolCallRequest(_)
|
||
| EventMsg::DynamicToolCallResponse(_)
|
||
| EventMsg::GuardianAssessment(_)
|
||
| EventMsg::ElicitationRequest(_)
|
||
| EventMsg::ApplyPatchApprovalRequest(_)
|
||
| EventMsg::DeprecationNotice(_)
|
||
| EventMsg::BackgroundEvent(_)
|
||
| EventMsg::UndoStarted(_)
|
||
| EventMsg::UndoCompleted(_)
|
||
| EventMsg::StreamError(_)
|
||
| EventMsg::TurnDiff(_)
|
||
| EventMsg::GetHistoryEntryResponse(_)
|
||
| EventMsg::McpListToolsResponse(_)
|
||
| EventMsg::ListSkillsResponse(_)
|
||
| EventMsg::RealtimeConversationListVoicesResponse(_)
|
||
| EventMsg::SkillsUpdateAvailable
|
||
| 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(_) => 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<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.conversation_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.conversation_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,
|
||
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;
|
||
|
||
if let Some(turn_item) =
|
||
handle_non_tool_response_item(sess, turn_context, item, /*plan_mode*/ true).await
|
||
{
|
||
emit_turn_item_in_plan_mode(
|
||
sess,
|
||
turn_context,
|
||
turn_item,
|
||
previously_active_item,
|
||
state,
|
||
)
|
||
.await;
|
||
}
|
||
|
||
record_completed_response_item(sess, turn_context, item).await;
|
||
if let Some(agent_message) = last_assistant_message_from_item(item, /*plan_mode*/ true) {
|
||
*last_agent_message = Some(agent_message);
|
||
}
|
||
return true;
|
||
}
|
||
false
|
||
}
|
||
|
||
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(())
|
||
}
|
||
|
||
#[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>,
|
||
client_session: &mut ModelClientSession,
|
||
turn_metadata_header: Option<&str>,
|
||
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.get(),
|
||
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 mut stream = client_session
|
||
.stream(
|
||
prompt,
|
||
&turn_context.model_info,
|
||
&turn_context.session_telemetry,
|
||
turn_context.reasoning_effort,
|
||
turn_context.reasoning_summary,
|
||
turn_context.config.service_tier,
|
||
turn_metadata_header,
|
||
&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 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 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,
|
||
);
|
||
|
||
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 Some(event) = consumer.flush_on_complete()
|
||
{
|
||
sess.send_event(&turn_context, event).await;
|
||
}
|
||
let previously_active_item = active_item.take();
|
||
if let Some(previous) = previously_active_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,
|
||
&item,
|
||
state,
|
||
previously_active_item.as_ref(),
|
||
&mut last_agent_message,
|
||
)
|
||
.await
|
||
{
|
||
continue;
|
||
}
|
||
|
||
let mut ctx = HandleOutputCtx {
|
||
sess: sess.clone(),
|
||
turn_context: turn_context.clone(),
|
||
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::LocalShellCall { .. }
|
||
| ResponseItem::FunctionCall { .. }
|
||
| ResponseItem::ToolSearchCall { .. }
|
||
| ResponseItem::FunctionCallOutput { .. }
|
||
| ResponseItem::CustomToolCall { .. }
|
||
| ResponseItem::CustomToolCallOutput { .. }
|
||
| ResponseItem::ToolSearchOutput { .. }
|
||
| ResponseItem::WebSearchCall { .. }
|
||
| ResponseItem::ImageGenerationCall { .. }
|
||
| ResponseItem::GhostSnapshot { .. }
|
||
| ResponseItem::Compaction { .. }
|
||
| ResponseItem::Other => false,
|
||
};
|
||
|
||
let output_result =
|
||
match handle_output_item_done(&mut ctx, item, previously_active_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.mailbox_rx.lock().await.has_pending() {
|
||
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(),
|
||
&item,
|
||
plan_mode,
|
||
)
|
||
.await
|
||
{
|
||
let mut turn_item = turn_item;
|
||
let mut seeded_parsed: Option<ParsedAssistantTextDelta> = None;
|
||
let mut seeded_item_id: Option<String> = None;
|
||
if 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 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);
|
||
}
|
||
}
|
||
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::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.update_rate_limits(&turn_context, snapshot).await;
|
||
}
|
||
ResponseEvent::ModelsEtag(etag) => {
|
||
// Update internal state with latest models etag
|
||
sess.services.models_manager.refresh_if_new_etag(etag).await;
|
||
}
|
||
ResponseEvent::Completed {
|
||
response_id: _,
|
||
token_usage,
|
||
} => {
|
||
flush_assistant_text_segments_all(
|
||
&sess,
|
||
&turn_context,
|
||
plan_mode_state.as_mut(),
|
||
&mut assistant_message_stream_parsers,
|
||
)
|
||
.await;
|
||
sess.update_token_usage_info(&turn_context, token_usage.as_ref())
|
||
.await;
|
||
should_emit_turn_diff = 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() {
|
||
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.conversation_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() {
|
||
let event = ReasoningContentDeltaEvent {
|
||
thread_id: sess.conversation_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() {
|
||
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() {
|
||
let event = ReasoningRawContentDeltaEvent {
|
||
thread_id: sess.conversation_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());
|
||
}
|
||
}
|
||
}
|
||
};
|
||
|
||
flush_assistant_text_segments_all(
|
||
&sess,
|
||
&turn_context,
|
||
plan_mode_state.as_mut(),
|
||
&mut assistant_message_stream_parsers,
|
||
)
|
||
.await;
|
||
|
||
drain_in_flight(&mut in_flight, sess.clone(), turn_context.clone()).await?;
|
||
|
||
if cancellation_token.is_cancelled() {
|
||
return Err(CodexErr::TurnAborted);
|
||
}
|
||
|
||
if should_emit_turn_diff {
|
||
let unified_diff = {
|
||
let mut tracker = turn_diff_tracker.lock().await;
|
||
tracker.get_unified_diff()
|
||
};
|
||
if let Ok(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, /*plan_mode*/ false) {
|
||
return Some(message);
|
||
}
|
||
}
|
||
None
|
||
}
|