File size: 3,353 Bytes
52a9af3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | use std::sync::Arc;
use tokio_util::sync::CancellationToken;
use crate::session::TurnInput;
use crate::session::session::Session;
use crate::session::turn::McpStartupRequirements;
use crate::session::turn::run_hooks_and_record_inputs;
use crate::session::turn::run_turn;
use crate::session::turn_context::TurnContext;
use crate::session_startup_prewarm::SessionStartupPrewarmResolution;
use crate::state::TaskKind;
use codex_thread_store::PersistContext;
use tracing::Instrument;
use tracing::trace_span;
use super::SessionTask;
use super::SessionTaskResult;
#[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>,
sess: Arc<Session>,
ctx: Arc<TurnContext>,
input: Vec<TurnInput>,
cancellation_token: CancellationToken,
) -> SessionTaskResult {
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 {
sess.emit_turn_started(&ctx).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 => {
run_hooks_and_record_inputs(
&sess,
&ctx,
&ctx.capture_current_model_info(),
&input,
PersistContext::Standard,
)
.await;
return Ok(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;
let mut mcp_startup_requirements = McpStartupRequirements::default();
loop {
let last_agent_message = run_turn(
Arc::clone(&sess),
Arc::clone(&ctx),
next_input,
&mut mcp_startup_requirements,
prewarmed_client_session.take(),
cancellation_token.child_token(),
)
.instrument(run_turn_span.clone())
.await?;
// Terminal errors are already reported. Let task completion preserve pending
// input instead of restarting the failed turn for that same input.
if ctx.terminal_error.lock().await.is_some() {
return Ok(last_agent_message);
}
if !sess.input_queue.has_pending_input(&sess.active_turn).await {
return Ok(last_agent_message);
}
next_input = Vec::new();
}
}
}
|