mirror of
https://github.com/openai/codex.git
synced 2026-04-30 09:26:44 +00:00
app-server: improve thread resume rejoin flow (#11776)
thread/resume response includes latest turn with all items, in band so no events are stale or lost Testing - e2e tested using app-server-test-client using flow described in "Testing Thread Rejoin Behavior" in codex-rs/app-server-test-client/README.md - e2e tested in codex desktop by reconnecting to a running turn
This commit is contained in:
@@ -45,7 +45,6 @@ use codex_app_server_protocol::McpToolCallResult;
|
||||
use codex_app_server_protocol::McpToolCallStatus;
|
||||
use codex_app_server_protocol::ModelReroutedNotification;
|
||||
use codex_app_server_protocol::PatchApplyStatus;
|
||||
use codex_app_server_protocol::PatchChangeKind as V2PatchChangeKind;
|
||||
use codex_app_server_protocol::PlanDeltaNotification;
|
||||
use codex_app_server_protocol::RawResponseItemCompletedNotification;
|
||||
use codex_app_server_protocol::ReasoningSummaryPartAddedNotification;
|
||||
@@ -72,6 +71,7 @@ use codex_app_server_protocol::TurnPlanStep;
|
||||
use codex_app_server_protocol::TurnPlanUpdatedNotification;
|
||||
use codex_app_server_protocol::TurnStatus;
|
||||
use codex_app_server_protocol::build_turns_from_rollout_items;
|
||||
use codex_app_server_protocol::convert_patch_changes;
|
||||
use codex_core::CodexThread;
|
||||
use codex_core::ThreadManager;
|
||||
use codex_core::parse_command::shlex_join;
|
||||
@@ -81,7 +81,6 @@ use codex_core::protocol::Event;
|
||||
use codex_core::protocol::EventMsg;
|
||||
use codex_core::protocol::ExecApprovalRequestEvent;
|
||||
use codex_core::protocol::ExecCommandEndEvent;
|
||||
use codex_core::protocol::FileChange as CoreFileChange;
|
||||
use codex_core::protocol::McpToolCallBeginEvent;
|
||||
use codex_core::protocol::McpToolCallEndEvent;
|
||||
use codex_core::protocol::Op;
|
||||
@@ -1706,46 +1705,6 @@ fn render_review_output_text(output: &ReviewOutputEvent) -> String {
|
||||
}
|
||||
}
|
||||
|
||||
fn convert_patch_changes(changes: &HashMap<PathBuf, CoreFileChange>) -> Vec<FileUpdateChange> {
|
||||
let mut converted: Vec<FileUpdateChange> = changes
|
||||
.iter()
|
||||
.map(|(path, change)| FileUpdateChange {
|
||||
path: path.to_string_lossy().into_owned(),
|
||||
kind: map_patch_change_kind(change),
|
||||
diff: format_file_change_diff(change),
|
||||
})
|
||||
.collect();
|
||||
converted.sort_by(|a, b| a.path.cmp(&b.path));
|
||||
converted
|
||||
}
|
||||
|
||||
fn map_patch_change_kind(change: &CoreFileChange) -> V2PatchChangeKind {
|
||||
match change {
|
||||
CoreFileChange::Add { .. } => V2PatchChangeKind::Add,
|
||||
CoreFileChange::Delete { .. } => V2PatchChangeKind::Delete,
|
||||
CoreFileChange::Update { move_path, .. } => V2PatchChangeKind::Update {
|
||||
move_path: move_path.clone(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn format_file_change_diff(change: &CoreFileChange) -> String {
|
||||
match change {
|
||||
CoreFileChange::Add { content } => content.clone(),
|
||||
CoreFileChange::Delete { content } => content.clone(),
|
||||
CoreFileChange::Update {
|
||||
unified_diff,
|
||||
move_path,
|
||||
} => {
|
||||
if let Some(path) = move_path {
|
||||
format!("{unified_diff}\n\nMoved to: {}", path.display())
|
||||
} else {
|
||||
unified_diff.clone()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn map_file_change_approval_decision(
|
||||
decision: FileChangeApprovalDecision,
|
||||
) -> (ReviewDecision, Option<PatchApplyStatus>) {
|
||||
|
||||
@@ -336,6 +336,7 @@ pub(crate) struct CodexMessageProcessor {
|
||||
outgoing: Arc<OutgoingMessageSender>,
|
||||
codex_linux_sandbox_exe: Option<PathBuf>,
|
||||
config: Arc<Config>,
|
||||
single_client_mode: bool,
|
||||
cli_overrides: Vec<(String, TomlValue)>,
|
||||
cloud_requirements: Arc<RwLock<CloudRequirementsLoader>>,
|
||||
active_login: Arc<Mutex<Option<ActiveLogin>>>,
|
||||
@@ -361,6 +362,7 @@ pub(crate) struct CodexMessageProcessorArgs {
|
||||
pub(crate) config: Arc<Config>,
|
||||
pub(crate) cli_overrides: Vec<(String, TomlValue)>,
|
||||
pub(crate) cloud_requirements: Arc<RwLock<CloudRequirementsLoader>>,
|
||||
pub(crate) single_client_mode: bool,
|
||||
pub(crate) feedback: CodexFeedback,
|
||||
}
|
||||
|
||||
@@ -397,6 +399,7 @@ impl CodexMessageProcessor {
|
||||
config,
|
||||
cli_overrides,
|
||||
cloud_requirements,
|
||||
single_client_mode,
|
||||
feedback,
|
||||
} = args;
|
||||
Self {
|
||||
@@ -405,6 +408,7 @@ impl CodexMessageProcessor {
|
||||
outgoing: outgoing.clone(),
|
||||
codex_linux_sandbox_exe,
|
||||
config,
|
||||
single_client_mode,
|
||||
cli_overrides,
|
||||
cloud_requirements,
|
||||
active_login: Arc::new(Mutex::new(None)),
|
||||
@@ -3042,21 +3046,14 @@ impl CodexMessageProcessor {
|
||||
return true;
|
||||
}
|
||||
|
||||
if let Err(err) = self
|
||||
.ensure_conversation_listener(
|
||||
existing_thread_id,
|
||||
request_id.connection_id,
|
||||
false,
|
||||
ApiVersion::V2,
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
"failed to attach listener for thread {}: {}",
|
||||
existing_thread_id,
|
||||
err.message
|
||||
);
|
||||
}
|
||||
let thread_state = self.thread_state_manager.thread_state(existing_thread_id);
|
||||
self.ensure_listener_task_running(
|
||||
existing_thread_id,
|
||||
existing_thread.clone(),
|
||||
thread_state.clone(),
|
||||
ApiVersion::V2,
|
||||
)
|
||||
.await;
|
||||
|
||||
let config_snapshot = existing_thread.config_snapshot().await;
|
||||
let mismatch_details = collect_resume_override_mismatches(params, &config_snapshot);
|
||||
@@ -3068,41 +3065,39 @@ impl CodexMessageProcessor {
|
||||
);
|
||||
}
|
||||
|
||||
let Some(mut thread) = self
|
||||
.load_thread_from_rollout_or_send_internal(
|
||||
request_id.clone(),
|
||||
existing_thread_id,
|
||||
rollout_path.as_path(),
|
||||
config_snapshot.model_provider_id.as_str(),
|
||||
)
|
||||
.await
|
||||
else {
|
||||
let listener_command_tx = {
|
||||
let thread_state = thread_state.lock().await;
|
||||
thread_state.listener_command_tx()
|
||||
};
|
||||
let Some(listener_command_tx) = listener_command_tx else {
|
||||
let err = JSONRPCErrorError {
|
||||
code: INTERNAL_ERROR_CODE,
|
||||
message: format!(
|
||||
"failed to enqueue running thread resume for thread {existing_thread_id}: thread listener is not running"
|
||||
),
|
||||
data: None,
|
||||
};
|
||||
self.outgoing.send_error(request_id, err).await;
|
||||
return true;
|
||||
};
|
||||
|
||||
let ThreadConfigSnapshot {
|
||||
model,
|
||||
model_provider_id,
|
||||
approval_policy,
|
||||
sandbox_policy,
|
||||
cwd,
|
||||
reasoning_effort,
|
||||
..
|
||||
} = config_snapshot;
|
||||
thread.status = self
|
||||
.thread_watch_manager
|
||||
.loaded_status_for_thread(&thread.id)
|
||||
.await;
|
||||
let response = ThreadResumeResponse {
|
||||
thread,
|
||||
model,
|
||||
model_provider: model_provider_id,
|
||||
cwd,
|
||||
approval_policy: approval_policy.into(),
|
||||
sandbox: sandbox_policy.into(),
|
||||
reasoning_effort,
|
||||
};
|
||||
self.outgoing.send_response(request_id, response).await;
|
||||
let command = crate::thread_state::ThreadListenerCommand::SendThreadResumeResponse(
|
||||
crate::thread_state::PendingThreadResumeRequest {
|
||||
request_id: request_id.clone(),
|
||||
rollout_path,
|
||||
config_snapshot,
|
||||
},
|
||||
);
|
||||
if listener_command_tx.send(command).is_err() {
|
||||
let err = JSONRPCErrorError {
|
||||
code: INTERNAL_ERROR_CODE,
|
||||
message: format!(
|
||||
"failed to enqueue running thread resume for thread {existing_thread_id}: thread listener command channel is closed"
|
||||
),
|
||||
data: None,
|
||||
};
|
||||
self.outgoing.send_error(request_id, err).await;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
false
|
||||
@@ -5817,17 +5812,18 @@ impl CodexMessageProcessor {
|
||||
api_version: ApiVersion,
|
||||
) {
|
||||
let (cancel_tx, mut cancel_rx) = oneshot::channel();
|
||||
{
|
||||
let mut listener_command_rx = {
|
||||
let mut thread_state = thread_state.lock().await;
|
||||
if thread_state.listener_matches(&conversation) {
|
||||
return;
|
||||
}
|
||||
thread_state.set_listener(cancel_tx, &conversation);
|
||||
}
|
||||
thread_state.set_listener(cancel_tx, &conversation)
|
||||
};
|
||||
let outgoing_for_task = self.outgoing.clone();
|
||||
let thread_manager = self.thread_manager.clone();
|
||||
let thread_watch_manager = self.thread_watch_manager.clone();
|
||||
let fallback_model_provider = self.config.model_provider_id.clone();
|
||||
let single_client_mode = self.single_client_mode;
|
||||
tokio::spawn(async move {
|
||||
loop {
|
||||
tokio::select! {
|
||||
@@ -5869,7 +5865,10 @@ impl CodexMessageProcessor {
|
||||
conversation_id.to_string().into(),
|
||||
);
|
||||
let (subscribed_connection_ids, raw_events_enabled) = {
|
||||
let thread_state = thread_state.lock().await;
|
||||
let mut thread_state = thread_state.lock().await;
|
||||
if !single_client_mode {
|
||||
thread_state.track_current_turn_event(&event.msg);
|
||||
}
|
||||
(
|
||||
thread_state.subscribed_connection_ids(),
|
||||
thread_state.experimental_raw_events,
|
||||
@@ -5908,6 +5907,25 @@ impl CodexMessageProcessor {
|
||||
)
|
||||
.await;
|
||||
}
|
||||
listener_command = listener_command_rx.recv() => {
|
||||
let Some(listener_command) = listener_command else {
|
||||
break;
|
||||
};
|
||||
match listener_command {
|
||||
crate::thread_state::ThreadListenerCommand::SendThreadResumeResponse(
|
||||
resume_request,
|
||||
) => {
|
||||
handle_pending_thread_resume_request(
|
||||
conversation_id,
|
||||
&thread_state,
|
||||
&thread_watch_manager,
|
||||
&outgoing_for_task,
|
||||
resume_request,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
@@ -6206,6 +6224,106 @@ impl CodexMessageProcessor {
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_pending_thread_resume_request(
|
||||
conversation_id: ThreadId,
|
||||
thread_state: &Arc<Mutex<ThreadState>>,
|
||||
thread_watch_manager: &ThreadWatchManager,
|
||||
outgoing: &Arc<OutgoingMessageSender>,
|
||||
pending: crate::thread_state::PendingThreadResumeRequest,
|
||||
) {
|
||||
let active_turn = {
|
||||
let state = thread_state.lock().await;
|
||||
state.active_turn_snapshot()
|
||||
};
|
||||
|
||||
let request_id = pending.request_id;
|
||||
let connection_id = request_id.connection_id;
|
||||
let mut thread = match load_thread_for_running_resume_response(
|
||||
conversation_id,
|
||||
pending.rollout_path.as_path(),
|
||||
pending.config_snapshot.model_provider_id.as_str(),
|
||||
active_turn.as_ref(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(thread) => thread,
|
||||
Err(message) => {
|
||||
outgoing
|
||||
.send_error(
|
||||
request_id,
|
||||
JSONRPCErrorError {
|
||||
code: INTERNAL_ERROR_CODE,
|
||||
message,
|
||||
data: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
thread.status = thread_watch_manager
|
||||
.loaded_status_for_thread(&thread.id)
|
||||
.await;
|
||||
|
||||
let ThreadConfigSnapshot {
|
||||
model,
|
||||
model_provider_id,
|
||||
approval_policy,
|
||||
sandbox_policy,
|
||||
cwd,
|
||||
reasoning_effort,
|
||||
..
|
||||
} = pending.config_snapshot;
|
||||
let response = ThreadResumeResponse {
|
||||
thread,
|
||||
model,
|
||||
model_provider: model_provider_id,
|
||||
cwd,
|
||||
approval_policy: approval_policy.into(),
|
||||
sandbox: sandbox_policy.into(),
|
||||
reasoning_effort,
|
||||
};
|
||||
outgoing.send_response(request_id, response).await;
|
||||
thread_state.lock().await.add_connection(connection_id);
|
||||
}
|
||||
|
||||
async fn load_thread_for_running_resume_response(
|
||||
conversation_id: ThreadId,
|
||||
rollout_path: &Path,
|
||||
fallback_provider: &str,
|
||||
active_turn: Option<&Turn>,
|
||||
) -> std::result::Result<Thread, String> {
|
||||
let mut thread = read_summary_from_rollout(rollout_path, fallback_provider)
|
||||
.await
|
||||
.map(summary_to_thread)
|
||||
.map_err(|err| {
|
||||
format!(
|
||||
"failed to load rollout `{}` for thread {conversation_id}: {err}",
|
||||
rollout_path.display()
|
||||
)
|
||||
})?;
|
||||
|
||||
let mut turns = read_rollout_items_from_rollout(rollout_path)
|
||||
.await
|
||||
.map(|items| build_turns_from_rollout_items(&items))
|
||||
.map_err(|err| {
|
||||
format!(
|
||||
"failed to load rollout `{}` for thread {conversation_id}: {err}",
|
||||
rollout_path.display()
|
||||
)
|
||||
})?;
|
||||
if let Some(active_turn) = active_turn {
|
||||
merge_turn_history_with_active_turn(&mut turns, active_turn.clone());
|
||||
}
|
||||
thread.turns = turns;
|
||||
Ok(thread)
|
||||
}
|
||||
|
||||
fn merge_turn_history_with_active_turn(turns: &mut Vec<Turn>, active_turn: Turn) {
|
||||
turns.retain(|turn| turn.id != active_turn.id);
|
||||
turns.push(active_turn);
|
||||
}
|
||||
|
||||
fn collect_resume_override_mismatches(
|
||||
request: &ThreadResumeParams,
|
||||
config_snapshot: &ThreadConfigSnapshot,
|
||||
|
||||
@@ -238,7 +238,8 @@ pub async fn run_main_with_transport(
|
||||
Some(start_websocket_acceptor(bind_address, transport_event_tx.clone()).await?);
|
||||
}
|
||||
}
|
||||
let shutdown_when_no_connections = matches!(transport, AppServerTransport::Stdio);
|
||||
let single_client_mode = matches!(transport, AppServerTransport::Stdio);
|
||||
let shutdown_when_no_connections = single_client_mode;
|
||||
|
||||
// Parse CLI overrides once and derive the base Config eagerly so later
|
||||
// components do not need to work with raw TOML values.
|
||||
@@ -439,6 +440,7 @@ pub async fn run_main_with_transport(
|
||||
outgoing: outgoing_message_sender,
|
||||
codex_linux_sandbox_exe,
|
||||
config: Arc::new(config),
|
||||
single_client_mode,
|
||||
cli_overrides,
|
||||
loader_overrides,
|
||||
cloud_requirements: cloud_requirements.clone(),
|
||||
|
||||
@@ -139,6 +139,7 @@ pub(crate) struct MessageProcessorArgs {
|
||||
pub(crate) outgoing: Arc<OutgoingMessageSender>,
|
||||
pub(crate) codex_linux_sandbox_exe: Option<PathBuf>,
|
||||
pub(crate) config: Arc<Config>,
|
||||
pub(crate) single_client_mode: bool,
|
||||
pub(crate) cli_overrides: Vec<(String, TomlValue)>,
|
||||
pub(crate) loader_overrides: LoaderOverrides,
|
||||
pub(crate) cloud_requirements: CloudRequirementsLoader,
|
||||
@@ -154,6 +155,7 @@ impl MessageProcessor {
|
||||
outgoing,
|
||||
codex_linux_sandbox_exe,
|
||||
config,
|
||||
single_client_mode,
|
||||
cli_overrides,
|
||||
loader_overrides,
|
||||
cloud_requirements,
|
||||
@@ -184,6 +186,7 @@ impl MessageProcessor {
|
||||
config: Arc::clone(&config),
|
||||
cli_overrides: cli_overrides.clone(),
|
||||
cloud_requirements: cloud_requirements.clone(),
|
||||
single_client_mode,
|
||||
feedback,
|
||||
});
|
||||
let config_api = ConfigApi::new(
|
||||
|
||||
@@ -1,13 +1,19 @@
|
||||
use crate::outgoing_message::ConnectionId;
|
||||
use crate::outgoing_message::ConnectionRequestId;
|
||||
use codex_app_server_protocol::ThreadHistoryBuilder;
|
||||
use codex_app_server_protocol::Turn;
|
||||
use codex_app_server_protocol::TurnError;
|
||||
use codex_core::CodexThread;
|
||||
use codex_core::ThreadConfigSnapshot;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::protocol::EventMsg;
|
||||
use std::collections::HashMap;
|
||||
use std::collections::HashSet;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Weak;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use uuid::Uuid;
|
||||
|
||||
@@ -16,6 +22,16 @@ type PendingInterruptQueue = Vec<(
|
||||
crate::codex_message_processor::ApiVersion,
|
||||
)>;
|
||||
|
||||
pub(crate) struct PendingThreadResumeRequest {
|
||||
pub(crate) request_id: ConnectionRequestId,
|
||||
pub(crate) rollout_path: PathBuf,
|
||||
pub(crate) config_snapshot: ThreadConfigSnapshot,
|
||||
}
|
||||
|
||||
pub(crate) enum ThreadListenerCommand {
|
||||
SendThreadResumeResponse(PendingThreadResumeRequest),
|
||||
}
|
||||
|
||||
/// Per-conversation accumulation of the latest states e.g. error message while a turn runs.
|
||||
#[derive(Default, Clone)]
|
||||
pub(crate) struct TurnSummary {
|
||||
@@ -31,6 +47,8 @@ pub(crate) struct ThreadState {
|
||||
pub(crate) turn_summary: TurnSummary,
|
||||
pub(crate) cancel_tx: Option<oneshot::Sender<()>>,
|
||||
pub(crate) experimental_raw_events: bool,
|
||||
listener_command_tx: Option<mpsc::UnboundedSender<ThreadListenerCommand>>,
|
||||
current_turn_history: ThreadHistoryBuilder,
|
||||
listener_thread: Option<Weak<CodexThread>>,
|
||||
subscribed_connections: HashSet<ConnectionId>,
|
||||
}
|
||||
@@ -47,17 +65,22 @@ impl ThreadState {
|
||||
&mut self,
|
||||
cancel_tx: oneshot::Sender<()>,
|
||||
conversation: &Arc<CodexThread>,
|
||||
) {
|
||||
) -> mpsc::UnboundedReceiver<ThreadListenerCommand> {
|
||||
if let Some(previous) = self.cancel_tx.replace(cancel_tx) {
|
||||
let _ = previous.send(());
|
||||
}
|
||||
let (listener_command_tx, listener_command_rx) = mpsc::unbounded_channel();
|
||||
self.listener_command_tx = Some(listener_command_tx);
|
||||
self.listener_thread = Some(Arc::downgrade(conversation));
|
||||
listener_command_rx
|
||||
}
|
||||
|
||||
pub(crate) fn clear_listener(&mut self) {
|
||||
if let Some(cancel_tx) = self.cancel_tx.take() {
|
||||
let _ = cancel_tx.send(());
|
||||
}
|
||||
self.listener_command_tx = None;
|
||||
self.current_turn_history.reset();
|
||||
self.listener_thread = None;
|
||||
}
|
||||
|
||||
@@ -76,6 +99,23 @@ impl ThreadState {
|
||||
pub(crate) fn set_experimental_raw_events(&mut self, enabled: bool) {
|
||||
self.experimental_raw_events = enabled;
|
||||
}
|
||||
|
||||
pub(crate) fn listener_command_tx(
|
||||
&self,
|
||||
) -> Option<mpsc::UnboundedSender<ThreadListenerCommand>> {
|
||||
self.listener_command_tx.clone()
|
||||
}
|
||||
|
||||
pub(crate) fn active_turn_snapshot(&self) -> Option<Turn> {
|
||||
self.current_turn_history.active_turn_snapshot()
|
||||
}
|
||||
|
||||
pub(crate) fn track_current_turn_event(&mut self, event: &EventMsg) {
|
||||
self.current_turn_history.handle_event(event);
|
||||
if !self.current_turn_history.has_active_turn() {
|
||||
self.current_turn_history.reset();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
@@ -200,12 +240,24 @@ impl ThreadStateManager {
|
||||
}
|
||||
|
||||
pub(crate) async fn remove_connection(&mut self, connection_id: ConnectionId) {
|
||||
let Some(thread_ids) = self.thread_ids_by_connection.remove(&connection_id) else {
|
||||
return;
|
||||
};
|
||||
let thread_ids = self
|
||||
.thread_ids_by_connection
|
||||
.remove(&connection_id)
|
||||
.unwrap_or_default();
|
||||
self.subscription_state_by_id
|
||||
.retain(|_, state| state.connection_id != connection_id);
|
||||
|
||||
if thread_ids.is_empty() {
|
||||
for thread_state in self.thread_states.values() {
|
||||
let mut thread_state = thread_state.lock().await;
|
||||
thread_state.remove_connection(connection_id);
|
||||
if thread_state.subscribed_connection_ids().is_empty() {
|
||||
thread_state.clear_listener();
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
for thread_id in thread_ids {
|
||||
if let Some(thread_state) = self.thread_states.get(&thread_id) {
|
||||
let mut thread_state = thread_state.lock().await;
|
||||
|
||||
Reference in New Issue
Block a user