codex / codex-rs /core /src /tasks /review.rs

Add files using upload-large-folder tool

52a9af3 verified about 1 month ago

10.1 kB

use std::sync::Arc;

use codex_prompts::render_review_exit_interrupted;
use codex_prompts::render_review_exit_success;
use codex_protocol::ResponseItemId;
use codex_protocol::config_types::WebSearchMode;
use codex_protocol::items::ExitedReviewModeItem;
use codex_protocol::items::TurnItem;
use codex_protocol::models::BaseInstructionsProvenance;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::AgentMessageContentDeltaEvent;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::ItemCompletedEvent;
use codex_protocol::protocol::ReviewOutputEvent;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::review_format::format_review_findings_block;
use codex_protocol::review_format::render_review_output_text;
use tokio_util::sync::CancellationToken;

use crate::codex_delegate::run_codex_thread_one_shot;
use crate::config::Constrained;
use crate::session::TurnInput;
use crate::session::session::Session;
use crate::session::turn_context::TurnContext;
use crate::state::TaskKind;
use codex_features::Feature;
use codex_protocol::user_input::UserInput;
use codex_thread_store::PersistContext;

use super::SessionTask;
use super::SessionTaskResult;

#[derive(Clone, Copy)]
pub(crate) struct ReviewTask;

impl ReviewTask {
    pub(crate) fn new() -> Self {
        Self
    }
}

impl SessionTask for ReviewTask {
    fn kind(&self) -> TaskKind {
        TaskKind::Review
    }

    fn span_name(&self) -> &'static str {
        "session_task.review"
    }

    async fn run(
        self: Arc<Self>,
        session: Arc<Session>,
        ctx: Arc<TurnContext>,
        input: Vec<TurnInput>,
        cancellation_token: CancellationToken,
    ) -> SessionTaskResult {
        session
            .services
            .session_telemetry
            .counter("codex.task.review", /*inc*/ 1, &[]);

        let mut user_input = Vec::new();
        for item in input {
            match item {
                TurnInput::UserInput { mut content, .. } => user_input.append(&mut content),
                TurnInput::ResponseItem(_)
                | TurnInput::FunctionCallOutput(_)
                | TurnInput::InterAgentCommunication(_) => {}
            }
        }

        // Start sub-codex conversation and get the receiver for events.
        let output = match start_review_conversation(
            session.clone(),
            ctx.clone(),
            user_input,
            cancellation_token.clone(),
        )
        .await
        {
            Some(receiver) => process_review_events(session.clone(), ctx.clone(), receiver).await,
            None => None,
        };
        if !cancellation_token.is_cancelled() {
            exit_review_mode(Arc::clone(&session), output.clone(), ctx.clone()).await;
        }
        Ok(None)
    }

    async fn abort(&self, session: Arc<Session>, ctx: Arc<TurnContext>) {
        exit_review_mode(session, /*review_output*/ None, ctx).await;
    }
}

async fn start_review_conversation(
    session: Arc<Session>,
    ctx: Arc<TurnContext>,
    input: Vec<UserInput>,
    cancellation_token: CancellationToken,
) -> Option<async_channel::Receiver<Event>> {
    let config = ctx.config.clone();
    let mut sub_agent_config = config.as_ref().clone();
    // Carry over review-only feature restrictions so the delegate cannot
    // re-enable blocked tools (web search, collab tools, view image).
    if let Err(err) = sub_agent_config
        .web_search_mode
        .set(WebSearchMode::Disabled)
    {
        panic!("by construction Constrained<WebSearchMode> must always support Disabled: {err}");
    }
    let _ = sub_agent_config.features.disable(Feature::Collab);
    let _ = sub_agent_config.features.disable(Feature::MultiAgentV2);

    // Set explicit review rubric for the sub-agent
    sub_agent_config.base_instructions = Some(crate::REVIEW_PROMPT.to_string());
    sub_agent_config.base_instructions_provenance = Some(BaseInstructionsProvenance::Custom);
    sub_agent_config.permissions.approval_policy = Constrained::allow_only(AskForApproval::Never);

    let model = config
        .review_model
        .clone()
        .unwrap_or_else(|| ctx.model_info().slug.clone());
    sub_agent_config.model = Some(model);
    (run_codex_thread_one_shot(
        sub_agent_config,
        Arc::clone(&session.services.auth_manager),
        Arc::clone(&session.services.models_manager),
        input,
        Arc::clone(&session),
        ctx.clone(),
        cancellation_token,
        SubAgentSource::Review,
        /*final_output_json_schema*/ None,
        /*initial_history*/ None,
    )
    .await)
        .ok()
        .map(|(_session, io)| io.rx_event)
}

async fn process_review_events(
    session: Arc<Session>,
    ctx: Arc<TurnContext>,
    receiver: async_channel::Receiver<Event>,
) -> Option<ReviewOutputEvent> {
    let mut prev_agent_message: Option<Event> = None;
    while let Ok(event) = receiver.recv().await {
        match event.clone().msg {
            EventMsg::AgentMessage(_) => {
                if let Some(prev) = prev_agent_message.take() {
                    session.send_event(ctx.as_ref(), prev.msg).await;
                }
                prev_agent_message = Some(event);
            }
            // Suppress ItemCompleted only for assistant messages: forwarding it
            // would trigger legacy AgentMessage via as_legacy_events(), which this
            // review flow intentionally hides in favor of structured output.
            EventMsg::ItemCompleted(ItemCompletedEvent {
                item: TurnItem::AgentMessage(_),
                ..
            })
            | EventMsg::AgentMessageContentDelta(AgentMessageContentDeltaEvent { .. }) => {}
            EventMsg::TurnComplete(task_complete) => {
                // Parse review output from the last agent message (if present).
                let out = task_complete
                    .last_agent_message
                    .as_deref()
                    .map(parse_review_output_event);
                return out;
            }
            EventMsg::TurnAborted(_) => {
                // Cancellation or abort: consumer will finalize with None.
                return None;
            }
            other => {
                session.send_event(ctx.as_ref(), other).await;
            }
        }
    }
    // Channel closed without TurnComplete: treat as interrupted.
    None
}

/// Parse a ReviewOutputEvent from a text blob returned by the reviewer model.
/// If the text is valid JSON matching ReviewOutputEvent, deserialize it.
/// Otherwise, attempt to extract the first JSON object substring and parse it.
/// If parsing still fails, return a structured fallback carrying the plain text
/// in `overall_explanation`.
fn parse_review_output_event(text: &str) -> ReviewOutputEvent {
    if let Ok(ev) = serde_json::from_str::<ReviewOutputEvent>(text) {
        return ev;
    }
    if let (Some(start), Some(end)) = (text.find('{'), text.rfind('}'))
        && start < end
        && let Some(slice) = text.get(start..=end)
        && let Ok(ev) = serde_json::from_str::<ReviewOutputEvent>(slice)
    {
        return ev;
    }
    ReviewOutputEvent {
        overall_explanation: text.to_string(),
        ..Default::default()
    }
}

/// Emits ExitedReviewMode item lifecycle with optional ReviewOutput,
/// and records the review output back into conversation history.
pub(crate) async fn exit_review_mode(
    session: Arc<Session>,
    review_output: Option<ReviewOutputEvent>,
    ctx: Arc<TurnContext>,
) {
    let (user_message, assistant_message) = if let Some(out) = review_output.clone() {
        let mut findings_str = String::new();
        let text = out.overall_explanation.trim();
        if !text.is_empty() {
            findings_str.push_str(text);
        }
        if !out.findings.is_empty() {
            let block = format_review_findings_block(&out.findings, /*selection*/ None);
            findings_str.push_str(&format!("\n{block}"));
        }
        let rendered = render_review_exit_success(&findings_str);
        let assistant_message = render_review_output_text(&out);
        (rendered, assistant_message)
    } else {
        let rendered = render_review_exit_interrupted();
        let assistant_message =
            "Review was interrupted. Please re-run /review and wait for it to complete."
                .to_string();
        (rendered, assistant_message)
    };

    session
        .record_conversation_items(
            &ctx,
            ctx.model_info(),
            &[ResponseItem::Message {
                id: Some(ResponseItemId::new("msg")),
                role: "user".to_string(),
                content: vec![ContentItem::InputText { text: user_message }],
                phase: None,
                internal_chat_message_metadata_passthrough: None,
            }],
        )
        .await;

    let item = TurnItem::ExitedReviewMode(ExitedReviewModeItem {
        id: uuid::Uuid::now_v7().to_string(),
        review_output,
    });
    session.emit_turn_item_started(ctx.as_ref(), &item).await;
    session.emit_turn_item_completed(ctx.as_ref(), item).await;
    session
        .record_response_item_and_emit_turn_item(
            ctx.as_ref(),
            ctx.model_info(),
            ResponseItem::Message {
                id: Some(ResponseItemId::new("msg")),
                role: "assistant".to_string(),
                content: vec![ContentItem::OutputText {
                    text: assistant_message,
                }],
                phase: None,
                internal_chat_message_metadata_passthrough: None,
            },
        )
        .await;

    // Review turns can run before any regular user turn, so explicitly
    // materialize rollout persistence. Do this after emitting review output so
    // file creation + git metadata collection cannot delay client-facing items.
    session
        .ensure_rollout_materialized(PersistContext::Standard)
        .await;
}

	use std::sync::Arc;

	use codex_prompts::render_review_exit_interrupted;
	use codex_prompts::render_review_exit_success;
	use codex_protocol::ResponseItemId;
	use codex_protocol::config_types::WebSearchMode;
	use codex_protocol::items::ExitedReviewModeItem;
	use codex_protocol::items::TurnItem;
	use codex_protocol::models::BaseInstructionsProvenance;
	use codex_protocol::models::ContentItem;
	use codex_protocol::models::ResponseItem;
	use codex_protocol::protocol::AgentMessageContentDeltaEvent;
	use codex_protocol::protocol::AskForApproval;
	use codex_protocol::protocol::Event;
	use codex_protocol::protocol::EventMsg;
	use codex_protocol::protocol::ItemCompletedEvent;
	use codex_protocol::protocol::ReviewOutputEvent;
	use codex_protocol::protocol::SubAgentSource;
	use codex_protocol::review_format::format_review_findings_block;
	use codex_protocol::review_format::render_review_output_text;
	use tokio_util::sync::CancellationToken;

	use crate::codex_delegate::run_codex_thread_one_shot;
	use crate::config::Constrained;
	use crate::session::TurnInput;
	use crate::session::session::Session;
	use crate::session::turn_context::TurnContext;
	use crate::state::TaskKind;
	use codex_features::Feature;
	use codex_protocol::user_input::UserInput;
	use codex_thread_store::PersistContext;

	use super::SessionTask;
	use super::SessionTaskResult;

	#[derive(Clone, Copy)]
	pub(crate) struct ReviewTask;

	impl ReviewTask {
	pub(crate) fn new() -> Self {
	Self
	}
	}

	impl SessionTask for ReviewTask {
	fn kind(&self) -> TaskKind {
	TaskKind::Review
	}

	fn span_name(&self) -> &'static str {
	"session_task.review"
	}

	async fn run(
	self: Arc<Self>,
	session: Arc<Session>,
	ctx: Arc<TurnContext>,
	input: Vec<TurnInput>,
	cancellation_token: CancellationToken,
	) -> SessionTaskResult {
	session
	.services
	.session_telemetry
	.counter("codex.task.review", /inc/ 1, &[]);

	let mut user_input = Vec::new();
	for item in input {
	match item {
	TurnInput::UserInput { mut content, .. } => user_input.append(&mut content),
	TurnInput::ResponseItem(_)
	\| TurnInput::FunctionCallOutput(_)
	\| TurnInput::InterAgentCommunication(_) => {}
	}
	}

	// Start sub-codex conversation and get the receiver for events.
	let output = match start_review_conversation(
	session.clone(),
	ctx.clone(),
	user_input,
	cancellation_token.clone(),
	)
	.await
	{
	Some(receiver) => process_review_events(session.clone(), ctx.clone(), receiver).await,
	None => None,
	};
	if !cancellation_token.is_cancelled() {
	exit_review_mode(Arc::clone(&session), output.clone(), ctx.clone()).await;
	}
	Ok(None)
	}

	async fn abort(&self, session: Arc<Session>, ctx: Arc<TurnContext>) {
	exit_review_mode(session, /review_output/ None, ctx).await;
	}
	}

	async fn start_review_conversation(
	session: Arc<Session>,
	ctx: Arc<TurnContext>,
	input: Vec<UserInput>,
	cancellation_token: CancellationToken,
	) -> Option<async_channel::Receiver<Event>> {
	let config = ctx.config.clone();
	let mut sub_agent_config = config.as_ref().clone();
	// Carry over review-only feature restrictions so the delegate cannot
	// re-enable blocked tools (web search, collab tools, view image).
	if let Err(err) = sub_agent_config
	.web_search_mode
	.set(WebSearchMode::Disabled)
	{
	panic!("by construction Constrained<WebSearchMode> must always support Disabled: {err}");
	}
	let _ = sub_agent_config.features.disable(Feature::Collab);
	let _ = sub_agent_config.features.disable(Feature::MultiAgentV2);

	// Set explicit review rubric for the sub-agent
	sub_agent_config.base_instructions = Some(crate::REVIEW_PROMPT.to_string());
	sub_agent_config.base_instructions_provenance = Some(BaseInstructionsProvenance::Custom);
	sub_agent_config.permissions.approval_policy = Constrained::allow_only(AskForApproval::Never);

	let model = config
	.review_model
	.clone()
	.unwrap_or_else(\|\| ctx.model_info().slug.clone());
	sub_agent_config.model = Some(model);
	(run_codex_thread_one_shot(
	sub_agent_config,
	Arc::clone(&session.services.auth_manager),
	Arc::clone(&session.services.models_manager),
	input,
	Arc::clone(&session),
	ctx.clone(),
	cancellation_token,
	SubAgentSource::Review,
	/final_output_json_schema/ None,
	/initial_history/ None,
	)
	.await)
	.ok()
	.map(\|(_session, io)\| io.rx_event)
	}

	async fn process_review_events(
	session: Arc<Session>,
	ctx: Arc<TurnContext>,
	receiver: async_channel::Receiver<Event>,
	) -> Option<ReviewOutputEvent> {
	let mut prev_agent_message: Option<Event> = None;
	while let Ok(event) = receiver.recv().await {
	match event.clone().msg {
	EventMsg::AgentMessage(_) => {
	if let Some(prev) = prev_agent_message.take() {
	session.send_event(ctx.as_ref(), prev.msg).await;
	}
	prev_agent_message = Some(event);
	}
	// Suppress ItemCompleted only for assistant messages: forwarding it
	// would trigger legacy AgentMessage via as_legacy_events(), which this
	// review flow intentionally hides in favor of structured output.
	EventMsg::ItemCompleted(ItemCompletedEvent {
	item: TurnItem::AgentMessage(_),
	..
	})
	\| EventMsg::AgentMessageContentDelta(AgentMessageContentDeltaEvent { .. }) => {}
	EventMsg::TurnComplete(task_complete) => {
	// Parse review output from the last agent message (if present).
	let out = task_complete
	.last_agent_message
	.as_deref()
	.map(parse_review_output_event);
	return out;
	}
	EventMsg::TurnAborted(_) => {
	// Cancellation or abort: consumer will finalize with None.
	return None;
	}
	other => {
	session.send_event(ctx.as_ref(), other).await;
	}
	}
	}
	// Channel closed without TurnComplete: treat as interrupted.
	None
	}

	/// Parse a ReviewOutputEvent from a text blob returned by the reviewer model.
	/// If the text is valid JSON matching ReviewOutputEvent, deserialize it.
	/// Otherwise, attempt to extract the first JSON object substring and parse it.
	/// If parsing still fails, return a structured fallback carrying the plain text
	/// in `overall_explanation`.
	fn parse_review_output_event(text: &str) -> ReviewOutputEvent {
	if let Ok(ev) = serde_json::from_str::<ReviewOutputEvent>(text) {
	return ev;
	}
	if let (Some(start), Some(end)) = (text.find('{'), text.rfind('}'))
	&& start < end
	&& let Some(slice) = text.get(start..=end)
	&& let Ok(ev) = serde_json::from_str::<ReviewOutputEvent>(slice)
	{
	return ev;
	}
	ReviewOutputEvent {
	overall_explanation: text.to_string(),
	..Default::default()
	}
	}

	/// Emits ExitedReviewMode item lifecycle with optional ReviewOutput,
	/// and records the review output back into conversation history.
	pub(crate) async fn exit_review_mode(
	session: Arc<Session>,
	review_output: Option<ReviewOutputEvent>,
	ctx: Arc<TurnContext>,
	) {
	let (user_message, assistant_message) = if let Some(out) = review_output.clone() {
	let mut findings_str = String::new();
	let text = out.overall_explanation.trim();
	if !text.is_empty() {
	findings_str.push_str(text);
	}
	if !out.findings.is_empty() {
	let block = format_review_findings_block(&out.findings, /selection/ None);
	findings_str.push_str(&format!("\n{block}"));
	}
	let rendered = render_review_exit_success(&findings_str);
	let assistant_message = render_review_output_text(&out);
	(rendered, assistant_message)
	} else {
	let rendered = render_review_exit_interrupted();
	let assistant_message =
	"Review was interrupted. Please re-run /review and wait for it to complete."
	.to_string();
	(rendered, assistant_message)
	};

	session
	.record_conversation_items(
	&ctx,
	ctx.model_info(),
	&[ResponseItem::Message {
	id: Some(ResponseItemId::new("msg")),
	role: "user".to_string(),
	content: vec![ContentItem::InputText { text: user_message }],
	phase: None,
	internal_chat_message_metadata_passthrough: None,
	}],
	)
	.await;

	let item = TurnItem::ExitedReviewMode(ExitedReviewModeItem {
	id: uuid::Uuid::now_v7().to_string(),
	review_output,
	});
	session.emit_turn_item_started(ctx.as_ref(), &item).await;
	session.emit_turn_item_completed(ctx.as_ref(), item).await;
	session
	.record_response_item_and_emit_turn_item(
	ctx.as_ref(),
	ctx.model_info(),
	ResponseItem::Message {
	id: Some(ResponseItemId::new("msg")),
	role: "assistant".to_string(),
	content: vec![ContentItem::OutputText {
	text: assistant_message,
	}],
	phase: None,
	internal_chat_message_metadata_passthrough: None,
	},
	)
	.await;

	// Review turns can run before any regular user turn, so explicitly
	// materialize rollout persistence. Do this after emitting review output so
	// file creation + git metadata collection cannot delay client-facing items.
	session
	.ensure_rollout_materialized(PersistContext::Standard)
	.await;
	}