better wiring

This commit is contained in:
jif-oai
2026-01-20 12:33:30 +00:00
parent e323e905ff
commit 8df6d0d4ee
9 changed files with 115 additions and 132 deletions
@@ -1860,7 +1860,7 @@ impl CodexMessageProcessor {
match self
.thread_manager
.resume_thread_with_history(config, thread_history, self.auth_manager.clone())
.resume_thread_with_history(config, thread_history)
.await
{
Ok(NewThread {
@@ -2783,7 +2783,7 @@ impl CodexMessageProcessor {
match self
.thread_manager
.resume_thread_with_history(config, thread_history, self.auth_manager.clone())
.resume_thread_with_history(config, thread_history)
.await
{
Ok(NewThread {
-1
View File
@@ -1,7 +1,6 @@
use crate::agent::AgentStatus;
use crate::error::CodexErr;
use crate::error::Result as CodexResult;
use crate::protocol::Event;
use crate::thread_manager::ThreadManagerState;
use codex_protocol::ThreadId;
use codex_protocol::protocol::InitialHistory;
-2
View File
@@ -3209,8 +3209,6 @@ pub(super) fn get_last_assistant_message_from_turn(responses: &[ResponseItem]) -
pub(crate) use tests::make_session_and_context;
use crate::git_info::get_git_repo_root;
#[cfg(test)]
pub(crate) use tests::make_session_and_context_with_rx;
#[cfg(test)]
mod tests {
+45 -73
View File
@@ -1,25 +1,23 @@
use async_trait::async_trait;
use codex_protocol::config_types::WebSearchMode;
use codex_protocol::items::TurnItem;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::AgentMessageContentDeltaEvent;
use codex_protocol::protocol::AgentMessageDeltaEvent;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::ExitedReviewModeEvent;
use codex_protocol::protocol::ItemCompletedEvent;
use codex_protocol::protocol::ReviewOutputEvent;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
use crate::agent::AgentStatus;
use crate::agent::status::is_final;
use crate::codex::Session;
use crate::codex::TurnContext;
use crate::review_format::format_review_findings_block;
use crate::review_format::render_review_output_text;
use crate::state::TaskKind;
use async_trait::async_trait;
use codex_protocol::config_types::WebSearchMode;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::ExitedReviewModeEvent;
use codex_protocol::protocol::ReviewOutputEvent;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::user_input::UserInput;
use std::sync::Arc;
use tokio::sync::watch;
use tokio_util::sync::CancellationToken;
use super::SessionTask;
use super::SessionTaskContext;
@@ -52,16 +50,11 @@ impl SessionTask for ReviewTask {
.otel_manager
.counter("codex.task.review", 1, &[]);
// Start sub-codex conversation and get the receiver for events.
let output = match start_review_conversation(
session.clone(),
ctx.clone(),
input,
cancellation_token.clone(),
)
.await
{
Some(receiver) => process_review_events(session.clone(), ctx.clone(), receiver).await,
// Start sub-codex conversation and await its final status.
let output = match start_review_conversation(session.clone(), ctx.clone(), input).await {
Some(sub_agent) => {
wait_for_review_output(sub_agent.status_rx, cancellation_token.clone()).await
}
None => None,
};
if !cancellation_token.is_cancelled() {
@@ -79,8 +72,7 @@ async fn start_review_conversation(
session: Arc<SessionTaskContext>,
ctx: Arc<TurnContext>,
input: Vec<UserInput>,
cancellation_token: CancellationToken,
) -> Option<async_channel::Receiver<Event>> {
) -> Option<ReviewSubAgent> {
let config = ctx.client.config();
let mut sub_agent_config = config.as_ref().clone();
// Carry over review-only feature restrictions so the sub-agent cannot
@@ -108,58 +100,38 @@ async fn start_review_conversation(
agent_control.send_input(thread_id, input).await.ok()?;
// todo this part is not done yet and won't be easy
None
let status_rx = agent_control.subscribe_status(thread_id).await.ok()?;
Some(ReviewSubAgent { status_rx })
}
async fn process_review_events(
session: Arc<SessionTaskContext>,
ctx: Arc<TurnContext>,
receiver: async_channel::Receiver<Event>,
struct ReviewSubAgent {
status_rx: watch::Receiver<AgentStatus>,
}
async fn wait_for_review_output(
mut status_rx: watch::Receiver<AgentStatus>,
cancellation_token: CancellationToken,
) -> 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
.clone_session()
.send_event(ctx.as_ref(), prev.msg)
.await;
let mut status = status_rx.borrow().clone();
while !is_final(&status) {
tokio::select! {
_ = cancellation_token.cancelled() => return None,
changed = status_rx.changed() => {
if changed.is_err() {
break;
}
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::AgentMessageDelta(AgentMessageDeltaEvent { .. })
| 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
.clone_session()
.send_event(ctx.as_ref(), other)
.await;
status = status_rx.borrow().clone();
}
}
}
// Channel closed without TurnComplete: treat as interrupted.
None
review_output_from_status(status)
}
fn review_output_from_status(status: AgentStatus) -> Option<ReviewOutputEvent> {
match status {
AgentStatus::Completed(Some(message)) => Some(parse_review_output_event(&message)),
_ => None,
}
}
/// Parse a ReviewOutputEvent from a text blob returned by the reviewer model.
+3 -3
View File
@@ -58,9 +58,9 @@ async fn review_subagent_ignores_legacy_deltas() {
}
}
assert_eq!(reasoning_delta_count, 1, "expected one new reasoning delta");
assert_eq!(reasoning_delta_count, 0, "expected no reasoning deltas");
assert_eq!(
legacy_reasoning_delta_count, 1,
"expected one legacy reasoning delta"
legacy_reasoning_delta_count, 0,
"expected no legacy reasoning deltas"
);
}
+60 -38
View File
@@ -74,7 +74,8 @@ async fn review_op_emits_lifecycle_and_review_output() {
let sse_raw = sse_template.replace("__REVIEW__", &review_json_escaped);
let (server, _request_log) = start_responses_server_with_sse(&sse_raw, 1).await;
let codex_home = TempDir::new().unwrap();
let codex = new_conversation_for_server(&server, &codex_home, |_| {}).await;
let conversation = new_conversation_for_server(&server, &codex_home, |_| {}).await;
let codex = &conversation.codex;
// Submit review request.
codex
@@ -90,8 +91,8 @@ async fn review_op_emits_lifecycle_and_review_output() {
.unwrap();
// Verify lifecycle: Entered -> Exited(Some(review)) -> TurnComplete.
let _entered = wait_for_event(&codex, |ev| matches!(ev, EventMsg::EnteredReviewMode(_))).await;
let closed = wait_for_event(&codex, |ev| matches!(ev, EventMsg::ExitedReviewMode(_))).await;
let _entered = wait_for_event(codex, |ev| matches!(ev, EventMsg::EnteredReviewMode(_))).await;
let closed = wait_for_event(codex, |ev| matches!(ev, EventMsg::ExitedReviewMode(_))).await;
let review = match closed {
EventMsg::ExitedReviewMode(ev) => ev
.review_output
@@ -116,7 +117,7 @@ async fn review_op_emits_lifecycle_and_review_output() {
overall_confidence_score: 0.8,
};
assert_eq!(expected, review);
let _complete = wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
let _complete = wait_for_event(codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
// Also verify that a user message with the header and a formatted finding
// was recorded back in the parent session's rollout.
@@ -195,7 +196,8 @@ async fn review_op_with_plain_text_emits_review_fallback() {
]"#;
let (server, _request_log) = start_responses_server_with_sse(sse_raw, 1).await;
let codex_home = TempDir::new().unwrap();
let codex = new_conversation_for_server(&server, &codex_home, |_| {}).await;
let conversation = new_conversation_for_server(&server, &codex_home, |_| {}).await;
let codex = &conversation.codex;
codex
.submit(Op::Review {
@@ -209,8 +211,8 @@ async fn review_op_with_plain_text_emits_review_fallback() {
.await
.unwrap();
let _entered = wait_for_event(&codex, |ev| matches!(ev, EventMsg::EnteredReviewMode(_))).await;
let closed = wait_for_event(&codex, |ev| matches!(ev, EventMsg::ExitedReviewMode(_))).await;
let _entered = wait_for_event(codex, |ev| matches!(ev, EventMsg::EnteredReviewMode(_))).await;
let closed = wait_for_event(codex, |ev| matches!(ev, EventMsg::ExitedReviewMode(_))).await;
let review = match closed {
EventMsg::ExitedReviewMode(ev) => ev
.review_output
@@ -224,7 +226,7 @@ async fn review_op_with_plain_text_emits_review_fallback() {
..Default::default()
};
assert_eq!(expected, review);
let _complete = wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
let _complete = wait_for_event(codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
server.verify().await;
}
@@ -255,7 +257,8 @@ async fn review_filters_agent_message_related_events() {
]"#;
let (server, _request_log) = start_responses_server_with_sse(sse_raw, 1).await;
let codex_home = TempDir::new().unwrap();
let codex = new_conversation_for_server(&server, &codex_home, |_| {}).await;
let conversation = new_conversation_for_server(&server, &codex_home, |_| {}).await;
let codex = &conversation.codex;
codex
.submit(Op::Review {
@@ -273,7 +276,7 @@ async fn review_filters_agent_message_related_events() {
let mut saw_exited = false;
// Drain until TurnComplete; assert streaming-related events never surface.
wait_for_event(&codex, |event| match event {
wait_for_event(codex, |event| match event {
EventMsg::TurnComplete(_) => true,
EventMsg::EnteredReviewMode(_) => {
saw_entered = true;
@@ -336,7 +339,8 @@ async fn review_does_not_emit_agent_message_on_structured_output() {
let sse_raw = sse_template.replace("__REVIEW__", &review_json_escaped);
let (server, _request_log) = start_responses_server_with_sse(&sse_raw, 1).await;
let codex_home = TempDir::new().unwrap();
let codex = new_conversation_for_server(&server, &codex_home, |_| {}).await;
let conversation = new_conversation_for_server(&server, &codex_home, |_| {}).await;
let codex = &conversation.codex;
codex
.submit(Op::Review {
@@ -355,7 +359,7 @@ async fn review_does_not_emit_agent_message_on_structured_output() {
let mut saw_entered = false;
let mut saw_exited = false;
let mut agent_messages = 0;
wait_for_event(&codex, |event| match event {
wait_for_event(codex, |event| match event {
EventMsg::TurnComplete(_) => true,
EventMsg::AgentMessage(_) => {
agent_messages += 1;
@@ -391,11 +395,12 @@ async fn review_uses_custom_review_model_from_config() {
let (server, request_log) = start_responses_server_with_sse(sse_raw, 1).await;
let codex_home = TempDir::new().unwrap();
// Choose a review model different from the main model; ensure it is used.
let codex = new_conversation_for_server(&server, &codex_home, |cfg| {
let conversation = new_conversation_for_server(&server, &codex_home, |cfg| {
cfg.model = Some("gpt-4.1".to_string());
cfg.review_model = Some("gpt-5.1".to_string());
})
.await;
let codex = &conversation.codex;
codex
.submit(Op::Review {
@@ -410,8 +415,8 @@ async fn review_uses_custom_review_model_from_config() {
.unwrap();
// Wait for completion
let _entered = wait_for_event(&codex, |ev| matches!(ev, EventMsg::EnteredReviewMode(_))).await;
let _closed = wait_for_event(&codex, |ev| {
let _entered = wait_for_event(codex, |ev| matches!(ev, EventMsg::EnteredReviewMode(_))).await;
let _closed = wait_for_event(codex, |ev| {
matches!(
ev,
EventMsg::ExitedReviewMode(ExitedReviewModeEvent {
@@ -420,7 +425,7 @@ async fn review_uses_custom_review_model_from_config() {
)
})
.await;
let _complete = wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
let _complete = wait_for_event(codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
// Assert the request body model equals the configured review model
let request = request_log.single_request();
@@ -443,11 +448,12 @@ async fn review_uses_session_model_when_review_model_unset() {
]"#;
let (server, request_log) = start_responses_server_with_sse(sse_raw, 1).await;
let codex_home = TempDir::new().unwrap();
let codex = new_conversation_for_server(&server, &codex_home, |cfg| {
let conversation = new_conversation_for_server(&server, &codex_home, |cfg| {
cfg.model = Some("gpt-4.1".to_string());
cfg.review_model = None;
})
.await;
let codex = &conversation.codex;
codex
.submit(Op::Review {
@@ -461,8 +467,8 @@ async fn review_uses_session_model_when_review_model_unset() {
.await
.unwrap();
let _entered = wait_for_event(&codex, |ev| matches!(ev, EventMsg::EnteredReviewMode(_))).await;
let _closed = wait_for_event(&codex, |ev| {
let _entered = wait_for_event(codex, |ev| matches!(ev, EventMsg::EnteredReviewMode(_))).await;
let _closed = wait_for_event(codex, |ev| {
matches!(
ev,
EventMsg::ExitedReviewMode(ExitedReviewModeEvent {
@@ -471,7 +477,7 @@ async fn review_uses_session_model_when_review_model_unset() {
)
})
.await;
let _complete = wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
let _complete = wait_for_event(codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
let request = request_log.single_request();
assert_eq!(request.path(), "/v1/responses");
@@ -561,8 +567,9 @@ async fn review_input_isolated_from_parent_history() {
.await
.unwrap();
}
let codex =
let conversation =
resume_conversation_for_server(&server, &codex_home, session_file.clone(), |_| {}).await;
let codex = &conversation.codex;
// Submit review request; it must start fresh (no parent history in `input`).
let review_prompt = "Please review only this".to_string();
@@ -578,8 +585,8 @@ async fn review_input_isolated_from_parent_history() {
.await
.unwrap();
let _entered = wait_for_event(&codex, |ev| matches!(ev, EventMsg::EnteredReviewMode(_))).await;
let _closed = wait_for_event(&codex, |ev| {
let _entered = wait_for_event(codex, |ev| matches!(ev, EventMsg::EnteredReviewMode(_))).await;
let _closed = wait_for_event(codex, |ev| {
matches!(
ev,
EventMsg::ExitedReviewMode(ExitedReviewModeEvent {
@@ -588,7 +595,7 @@ async fn review_input_isolated_from_parent_history() {
)
})
.await;
let _complete = wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
let _complete = wait_for_event(codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
// Assert the request `input` contains the environment context followed by the user review prompt.
let request = request_log.single_request();
@@ -674,7 +681,8 @@ async fn review_history_surfaces_in_parent_session() {
]"#;
let (server, request_log) = start_responses_server_with_sse(sse_raw, 2).await;
let codex_home = TempDir::new().unwrap();
let codex = new_conversation_for_server(&server, &codex_home, |_| {}).await;
let conversation = new_conversation_for_server(&server, &codex_home, |_| {}).await;
let codex = &conversation.codex;
// 1) Run a review turn that produces an assistant message (isolated in child).
codex
@@ -688,8 +696,8 @@ async fn review_history_surfaces_in_parent_session() {
})
.await
.unwrap();
let _entered = wait_for_event(&codex, |ev| matches!(ev, EventMsg::EnteredReviewMode(_))).await;
let _closed = wait_for_event(&codex, |ev| {
let _entered = wait_for_event(codex, |ev| matches!(ev, EventMsg::EnteredReviewMode(_))).await;
let _closed = wait_for_event(codex, |ev| {
matches!(
ev,
EventMsg::ExitedReviewMode(ExitedReviewModeEvent {
@@ -698,7 +706,7 @@ async fn review_history_surfaces_in_parent_session() {
)
})
.await;
let _complete = wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
let _complete = wait_for_event(codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
// 2) Continue in the parent session; request input must not include any review items.
let followup = "back to parent".to_string();
@@ -712,7 +720,7 @@ async fn review_history_surfaces_in_parent_session() {
})
.await
.unwrap();
let _complete = wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
let _complete = wait_for_event(codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
// Inspect the second request (parent turn) input contents.
// Parent turns include session initial messages (user_instructions, environment_context).
@@ -806,10 +814,11 @@ async fn review_uses_overridden_cwd_for_base_branch_merge_base() {
.to_string();
let codex_home = TempDir::new().unwrap();
let codex = new_conversation_for_server(&server, &codex_home, |config| {
let conversation = new_conversation_for_server(&server, &codex_home, |config| {
config.cwd = initial_cwd.path().to_path_buf();
})
.await;
let codex = &conversation.codex;
codex
.submit(Op::OverrideTurnContext {
@@ -836,8 +845,8 @@ async fn review_uses_overridden_cwd_for_base_branch_merge_base() {
.await
.unwrap();
let _entered = wait_for_event(&codex, |ev| matches!(ev, EventMsg::EnteredReviewMode(_))).await;
let _complete = wait_for_event(&codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
let _entered = wait_for_event(codex, |ev| matches!(ev, EventMsg::EnteredReviewMode(_))).await;
let _complete = wait_for_event(codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
let requests = request_log.requests();
assert_eq!(requests.len(), 1);
@@ -871,13 +880,18 @@ async fn start_responses_server_with_sse(
(server, request_log)
}
struct ReviewConversation {
codex: Arc<CodexThread>,
_manager: ThreadManager,
}
/// Create a conversation configured to talk to the provided mock server.
#[expect(clippy::expect_used)]
async fn new_conversation_for_server<F>(
server: &MockServer,
codex_home: &TempDir,
mutator: F,
) -> Arc<CodexThread>
) -> ReviewConversation
where
F: FnOnce(&mut Config),
{
@@ -892,11 +906,15 @@ where
CodexAuth::from_api_key("Test API Key"),
config.model_provider.clone(),
);
thread_manager
let codex = thread_manager
.start_thread(config)
.await
.expect("create conversation")
.thread
.thread;
ReviewConversation {
codex,
_manager: thread_manager,
}
}
/// Create a conversation resuming from a rollout file, configured to talk to the provided mock server.
@@ -906,7 +924,7 @@ async fn resume_conversation_for_server<F>(
codex_home: &TempDir,
resume_path: std::path::PathBuf,
mutator: F,
) -> Arc<CodexThread>
) -> ReviewConversation
where
F: FnOnce(&mut Config),
{
@@ -921,9 +939,13 @@ where
CodexAuth::from_api_key("Test API Key"),
config.model_provider.clone(),
);
thread_manager
let codex = thread_manager
.resume_thread_from_rollout(config, resume_path)
.await
.expect("resume conversation")
.thread
.thread;
ReviewConversation {
codex,
_manager: thread_manager,
}
}
+1 -1
View File
@@ -326,7 +326,7 @@ pub async fn run_main(cli: Cli, codex_linux_sandbox_exe: Option<PathBuf>) -> any
if let Some(path) = resume_path {
thread_manager
.resume_thread_from_rollout(config.clone(), path, auth_manager.clone())
.resume_thread_from_rollout(config.clone(), path)
.await?
} else {
thread_manager.start_thread(config.clone()).await?
+2 -6
View File
@@ -441,7 +441,7 @@ impl App {
}
SessionSelection::Resume(path) => {
let resumed = thread_manager
.resume_thread_from_rollout(config.clone(), path.clone(), auth_manager.clone())
.resume_thread_from_rollout(config.clone(), path.clone())
.await
.wrap_err_with(|| {
let path_display = path.display();
@@ -716,11 +716,7 @@ impl App {
);
match self
.server
.resume_thread_from_rollout(
self.config.clone(),
path.clone(),
self.auth_manager.clone(),
)
.resume_thread_from_rollout(self.config.clone(), path.clone())
.await
{
Ok(resumed) => {
+2 -6
View File
@@ -499,7 +499,7 @@ impl App {
}
SessionSelection::Resume(path) => {
let resumed = thread_manager
.resume_thread_from_rollout(config.clone(), path.clone(), auth_manager.clone())
.resume_thread_from_rollout(config.clone(), path.clone())
.await
.wrap_err_with(|| {
let path_display = path.display();
@@ -1490,11 +1490,7 @@ impl App {
);
match self
.server
.resume_thread_from_rollout(
self.config.clone(),
path.clone(),
self.auth_manager.clone(),
)
.resume_thread_from_rollout(self.config.clone(), path.clone())
.await
{
Ok(resumed) => {