Codex Handbook
core/src/tasks/regular.rs 89 lines
use std::sync::Arc;use tokio_util::sync::CancellationToken;use crate::session::TurnInput;use crate::session::turn::run_turn;use crate::session::turn_context::TurnContext;use crate::session_startup_prewarm::SessionStartupPrewarmResolution;use crate::state::TaskKind;use codex_protocol::protocol::EventMsg;use codex_protocol::protocol::TurnStartedEvent;use tracing::Instrument;use tracing::trace_span;use super::SessionTask;use super::SessionTaskContext;#[derive(Default)]pub(crate) struct RegularTask;impl RegularTask {    pub(crate) fn new() -> Self {        Self    }}impl SessionTask for RegularTask {    fn kind(&self) -> TaskKind {        TaskKind::Regular    }    fn span_name(&self) -> &'static str {        "session_task.turn"    }    async fn run(        self: Arc<Self>,        session: Arc<SessionTaskContext>,        ctx: Arc<TurnContext>,        input: Vec<TurnInput>,        cancellation_token: CancellationToken,    ) -> Option<String> {        let sess = session.clone_session();        let turn_extension_data = session.turn_extension_data();        let run_turn_span = trace_span!("run_turn");        // Regular turns emit `TurnStarted` inline so first-turn lifecycle does        // not wait on startup prewarm resolution.        let prewarmed_client_session = async {            let event = EventMsg::TurnStarted(TurnStartedEvent {                turn_id: ctx.sub_id.clone(),                trace_id: ctx.trace_id.clone(),                started_at: ctx.turn_timing_state.started_at_unix_secs().await,                model_context_window: ctx.model_context_window(),                collaboration_mode_kind: ctx.collaboration_mode.mode,            });            sess.send_event(ctx.as_ref(), event).await;            sess.set_server_reasoning_included(/*included*/ false).await;            sess.consume_startup_prewarm_for_regular_turn(&cancellation_token)                .await        }        .instrument(trace_span!("regular_task.prepare_run_turn"))        .await;        let prewarmed_client_session = match prewarmed_client_session {            SessionStartupPrewarmResolution::Cancelled => return None,            SessionStartupPrewarmResolution::Unavailable { .. } => None,            SessionStartupPrewarmResolution::Ready(prewarmed_client_session) => {                Some(*prewarmed_client_session)            }        };        let mut next_input = input;        let mut prewarmed_client_session = prewarmed_client_session;        loop {            let last_agent_message = run_turn(                Arc::clone(&sess),                Arc::clone(&ctx),                Arc::clone(&turn_extension_data),                next_input,                prewarmed_client_session.take(),                cancellation_token.child_token(),            )            .instrument(run_turn_span.clone())            .await;            if !sess.input_queue.has_pending_input(&sess.active_turn).await {                return last_agent_message;            }            next_input = Vec::new();        }    }}