mirror of
https://github.com/openai/codex.git
synced 2026-11-05 20:10:33 +00:00
Persist loaded threads before managed daemon shutdown (#44283)
## Why Loaded threads, including idle threads whose rollout files are still deferred, need to survive a managed daemon restart. Shutdown must also remain forceable when rollout I/O is blocked. ## What changed - Add a hidden `--managed-daemon` option for app-server Unix socket transports. After active turns and admitted requests drain, save loaded persistent root threads before exiting. - Skip ephemeral threads, non-root agent threads, and threads pending unload. Log persistence failures and continue saving other threads. - Apply shutdown admission checks to thread and turn settings updates, thread deletion, and archiving. - Keep force signals and daemon shutdown requests responsive during persistence. Return `AppServerExit::Forced` so executables can exit without waiting for runtime teardown. ## Testing Add integration coverage for resuming active and idle threads after restart, forcing shutdown during active work, and forcing shutdown with a blocked rollout writer. Extend shutdown rejection coverage to deletion and settings updates. GitOrigin-RevId: 4344e97d39f9f80c5d84c41c17300fd6aa99b4cc
This commit is contained in:
Generated
+1
@@ -2226,6 +2226,7 @@ dependencies = [
|
||||
"codex-state",
|
||||
"codex-thread-store",
|
||||
"codex-tools",
|
||||
"codex-uds",
|
||||
"codex-user-verification",
|
||||
"codex-utils-absolute-path",
|
||||
"codex-utils-cargo-bin",
|
||||
|
||||
@@ -123,6 +123,7 @@ tikv-jemallocator = { workspace = true }
|
||||
codex-windows-sandbox = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
codex-uds = { workspace = true }
|
||||
app_test_support = { workspace = true }
|
||||
axum = { workspace = true, default-features = false, features = [
|
||||
"http1",
|
||||
|
||||
@@ -445,6 +445,15 @@ pub async fn run_main(
|
||||
AppServerRuntimeOptions::default(),
|
||||
)
|
||||
.await
|
||||
.map(|_| ())
|
||||
}
|
||||
|
||||
/// Reports whether shutdown finished or the managed daemon must exit without waiting for I/O.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AppServerExit {
|
||||
Graceful,
|
||||
/// Executables should exit before dropping their runtime; libraries must not exit their host.
|
||||
Forced,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
@@ -459,6 +468,7 @@ pub struct AppServerRuntimeOptions {
|
||||
pub plugin_startup_tasks: PluginStartupTasks,
|
||||
pub remote_control_startup_mode: RemoteControlStartupMode,
|
||||
pub install_shutdown_signal_handler: bool,
|
||||
pub managed_daemon: bool,
|
||||
}
|
||||
|
||||
impl Default for AppServerRuntimeOptions {
|
||||
@@ -468,6 +478,7 @@ impl Default for AppServerRuntimeOptions {
|
||||
plugin_startup_tasks: PluginStartupTasks::Start,
|
||||
remote_control_startup_mode: RemoteControlStartupMode::ResolvePersisted,
|
||||
install_shutdown_signal_handler: true,
|
||||
managed_daemon: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -483,7 +494,7 @@ pub async fn run_main_with_transport_options(
|
||||
session_source: SessionSource,
|
||||
auth: AppServerWebsocketAuthSettings,
|
||||
runtime_options: AppServerRuntimeOptions,
|
||||
) -> IoResult<()> {
|
||||
) -> IoResult<AppServerExit> {
|
||||
let loader_overrides = loader_overrides_with_test_user_config_file(
|
||||
loader_overrides,
|
||||
test_user_config_file_from_env(),
|
||||
@@ -741,6 +752,8 @@ pub async fn run_main_with_transport_options(
|
||||
let single_client_mode = matches!(&transport, AppServerTransport::Stdio);
|
||||
let graceful_signal_restart_enabled =
|
||||
runtime_options.install_shutdown_signal_handler && !single_client_mode;
|
||||
let managed_daemon = matches!(&transport, AppServerTransport::UnixSocket { .. })
|
||||
&& runtime_options.managed_daemon;
|
||||
let mut app_server_client_name_rx = None;
|
||||
|
||||
match &transport {
|
||||
@@ -970,24 +983,43 @@ pub async fn run_main_with_transport_options(
|
||||
let transport_shutdown_token = transport_shutdown_token.clone();
|
||||
async move {
|
||||
let mut listen_for_threads = true;
|
||||
// Keep force signals and daemon control events responsive while saving.
|
||||
let mut snapshot = Box::pin(processor.persist_daemon_threads());
|
||||
let mut snapshot_finished = !managed_daemon;
|
||||
let mut clients_disconnected = false;
|
||||
let mut shutdown_state = ShutdownState::default();
|
||||
let exit_reason = loop {
|
||||
// Sample submissions first: one can publish a running turn before
|
||||
// releasing its permit, and shutdown must observe that new turn.
|
||||
let active_admissions = *active_admissions_rx.borrow_and_update();
|
||||
let running_turn_count = *running_turn_count_rx.borrow_and_update();
|
||||
if matches!(
|
||||
let ready_to_exit = matches!(
|
||||
shutdown_state.update(running_turn_count, active_admissions, connections.len()),
|
||||
ShutdownAction::Finish
|
||||
) {
|
||||
transport_shutdown_token.cancel();
|
||||
let _ = outbound_control_tx
|
||||
.send(OutboundControlEvent::DisconnectAll)
|
||||
.await;
|
||||
break "shutdown_requested";
|
||||
);
|
||||
if ready_to_exit {
|
||||
let finished = snapshot_finished || shutdown_state.forced();
|
||||
if finished {
|
||||
transport_shutdown_token.cancel();
|
||||
}
|
||||
if managed_daemon && shutdown_state.forced() {
|
||||
break "forced_shutdown_requested";
|
||||
}
|
||||
if !clients_disconnected {
|
||||
let _ = outbound_control_tx
|
||||
.send(OutboundControlEvent::DisconnectAll)
|
||||
.await;
|
||||
clients_disconnected = true;
|
||||
}
|
||||
if finished {
|
||||
break "shutdown_requested";
|
||||
}
|
||||
}
|
||||
|
||||
tokio::select! {
|
||||
_ = &mut snapshot, if ready_to_exit && !snapshot_finished => {
|
||||
snapshot_finished = true;
|
||||
}
|
||||
shutdown_signal_result = shutdown_signal(), if graceful_signal_restart_enabled && !shutdown_state.forced() => {
|
||||
let signal = match shutdown_signal_result {
|
||||
Ok(signal) => signal,
|
||||
@@ -1013,6 +1045,12 @@ pub async fn run_main_with_transport_options(
|
||||
let Some(event) = event else {
|
||||
break "transport_channel_closed";
|
||||
};
|
||||
if ready_to_exit && !matches!(event, TransportEvent::DaemonShutdown) {
|
||||
if let TransportEvent::ConnectionOpened { disconnect_sender: Some(token), .. } = event {
|
||||
token.cancel();
|
||||
}
|
||||
continue;
|
||||
}
|
||||
match event {
|
||||
TransportEvent::DaemonShutdown => {
|
||||
shutdown_state.on_signal(ShutdownSignal::Forceable, connections.len(), *running_turn_count_rx.borrow(), &processor.turn_admission);
|
||||
@@ -1188,7 +1226,7 @@ pub async fn run_main_with_transport_options(
|
||||
.send_server_notification(notification)
|
||||
.await;
|
||||
}
|
||||
created = thread_created_rx.recv(), if listen_for_threads => {
|
||||
created = thread_created_rx.recv(), if listen_for_threads && !ready_to_exit => {
|
||||
match created {
|
||||
Ok(thread_id) => {
|
||||
let mut initialized_connection_ids = Vec::new();
|
||||
@@ -1219,6 +1257,7 @@ pub async fn run_main_with_transport_options(
|
||||
}
|
||||
};
|
||||
|
||||
drop(snapshot);
|
||||
if !shutdown_state.forced() {
|
||||
futures::future::join_all(connections.iter().map(
|
||||
|(&connection_id, connection_state)| {
|
||||
@@ -1238,12 +1277,19 @@ pub async fn run_main_with_transport_options(
|
||||
shutdown_forced = shutdown_state.forced(),
|
||||
"processor task exited"
|
||||
);
|
||||
if managed_daemon && shutdown_state.forced() {
|
||||
AppServerExit::Forced
|
||||
} else {
|
||||
AppServerExit::Graceful
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
drop(transport_event_tx);
|
||||
|
||||
let _ = processor_handle.await;
|
||||
if matches!(processor_handle.await, Ok(AppServerExit::Forced)) {
|
||||
return Ok(AppServerExit::Forced);
|
||||
}
|
||||
let _ = outbound_handle.await;
|
||||
|
||||
transport_shutdown_token.cancel();
|
||||
@@ -1252,7 +1298,7 @@ pub async fn run_main_with_transport_options(
|
||||
let _ = handle.await;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
Ok(AppServerExit::Graceful)
|
||||
}
|
||||
|
||||
struct SqliteRecoveryNotice {
|
||||
|
||||
@@ -70,6 +70,10 @@ struct AppServerArgs {
|
||||
/// Enable remote control for this app-server process without changing persistence.
|
||||
#[arg(long = "remote-control", hide = true)]
|
||||
remote_control: bool,
|
||||
|
||||
/// Save loaded threads during managed daemon shutdown.
|
||||
#[arg(long, hide = true)]
|
||||
managed_daemon: bool,
|
||||
}
|
||||
|
||||
fn main() -> anyhow::Result<()> {
|
||||
@@ -85,6 +89,7 @@ fn main() -> anyhow::Result<()> {
|
||||
#[cfg(debug_assertions)]
|
||||
disable_plugin_startup_tasks_for_tests,
|
||||
remote_control,
|
||||
managed_daemon,
|
||||
} = AppServerArgs::parse();
|
||||
let loader_overrides = if disable_managed_config_from_debug_env() {
|
||||
LoaderOverrides::without_managed_config_for_tests()
|
||||
@@ -97,6 +102,7 @@ fn main() -> anyhow::Result<()> {
|
||||
let auth = auth.try_into_settings()?;
|
||||
let mut runtime_options = AppServerRuntimeOptions {
|
||||
code_mode_host_transport: code_mode_host.into(),
|
||||
managed_daemon,
|
||||
..Default::default()
|
||||
};
|
||||
#[cfg(debug_assertions)]
|
||||
@@ -110,7 +116,7 @@ fn main() -> anyhow::Result<()> {
|
||||
(false, false) => codex_app_server::RemoteControlStartupMode::ResolvePersisted,
|
||||
};
|
||||
|
||||
run_main_with_transport_options(
|
||||
let exit = run_main_with_transport_options(
|
||||
arg0_paths,
|
||||
config_overrides,
|
||||
loader_overrides,
|
||||
@@ -122,6 +128,10 @@ fn main() -> anyhow::Result<()> {
|
||||
runtime_options,
|
||||
)
|
||||
.await?;
|
||||
if exit == codex_app_server::AppServerExit::Forced {
|
||||
// Runtime teardown can wait forever for blocked rollout I/O.
|
||||
std::process::exit(0);
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
@@ -759,6 +759,10 @@ impl MessageProcessor {
|
||||
self.thread_processor.thread_created_receiver()
|
||||
}
|
||||
|
||||
pub(crate) async fn persist_daemon_threads(&self) {
|
||||
self.thread_processor.persist_daemon_threads().await
|
||||
}
|
||||
|
||||
pub(crate) async fn send_initialize_notifications_to_connection(
|
||||
&self,
|
||||
connection_id: ConnectionId,
|
||||
@@ -951,7 +955,11 @@ impl MessageProcessor {
|
||||
| ClientRequest::ThreadFork { .. }
|
||||
| ClientRequest::ThreadResume { .. }
|
||||
| ClientRequest::ThreadRollback { .. }
|
||||
| ClientRequest::ThreadRevert { .. } => (Some(self.turn_admission.admit()?), false),
|
||||
| ClientRequest::ThreadRevert { .. }
|
||||
| ClientRequest::ThreadSettingsUpdate { .. }
|
||||
| ClientRequest::TurnSettingsUpdate { .. }
|
||||
| ClientRequest::ThreadDelete { .. }
|
||||
| ClientRequest::ThreadArchive { .. } => (Some(self.turn_admission.admit()?), false),
|
||||
ClientRequest::TurnStart { .. }
|
||||
| ClientRequest::TurnSteer { .. }
|
||||
| ClientRequest::ReviewStart { .. }
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
//! Persists loaded persistent root threads before managed daemon shutdown.
|
||||
//! Failures are logged so other threads can still be saved.
|
||||
|
||||
use super::ThreadRequestProcessor;
|
||||
use codex_thread_store::PersistContext;
|
||||
use tracing::warn;
|
||||
|
||||
impl ThreadRequestProcessor {
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "snapshot selection must be serialized against pending unloads"
|
||||
)]
|
||||
pub(crate) async fn persist_daemon_threads(&self) {
|
||||
for thread_id in self.thread_manager.list_thread_ids().await {
|
||||
let pending_unloads = self.pending_thread_unloads.lock().await;
|
||||
if pending_unloads.contains(&thread_id) {
|
||||
continue;
|
||||
}
|
||||
let Ok(thread) = self.thread_manager.get_thread(thread_id).await else {
|
||||
continue;
|
||||
};
|
||||
let config = thread.config_snapshot().await;
|
||||
if !config.ephemeral
|
||||
&& config.parent_thread_id.is_none()
|
||||
&& !config.session_source.is_non_root_agent()
|
||||
&& let Err(err) = self
|
||||
.thread_store
|
||||
.persist_thread(thread_id, PersistContext::Standard)
|
||||
.await
|
||||
{
|
||||
warn!(%thread_id, %err, "failed to persist thread during daemon shutdown");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,6 @@
|
||||
#[path = "daemon_snapshot.rs"]
|
||||
mod daemon_snapshot;
|
||||
|
||||
use super::persisted_resume_settings::PersistedResumeSettings;
|
||||
use super::persisted_resume_settings::latest_persisted_resume_settings;
|
||||
use super::thread_enrichment::enrich_loaded_threads;
|
||||
|
||||
@@ -88,6 +88,15 @@ async fn websocket_transport_ctrl_c_waits_for_running_turn_before_exit() -> Resu
|
||||
("thread/start", json!({})),
|
||||
("thread/fork", json!({"threadId": thread_id})),
|
||||
("thread/resume", json!({"threadId": thread_id})),
|
||||
("thread/delete", json!({"threadId": thread_id})),
|
||||
(
|
||||
"thread/settings/update",
|
||||
json!({"threadId": thread_id, "model": "other"}),
|
||||
),
|
||||
(
|
||||
"turn/settings/update",
|
||||
json!({"threadId": thread_id, "turnId": "unused", "model": "other"}),
|
||||
),
|
||||
(
|
||||
"thread/rollback",
|
||||
json!({"threadId": thread_id, "numTurns": 1}),
|
||||
|
||||
@@ -0,0 +1,324 @@
|
||||
//! Best-effort thread persistence during managed app-server shutdown.
|
||||
|
||||
use super::connection_handling_websocket::DEFAULT_READ_TIMEOUT;
|
||||
use super::connection_handling_websocket::create_config_toml;
|
||||
use anyhow::Context;
|
||||
use anyhow::Result;
|
||||
use app_test_support::DISABLE_PLUGIN_STARTUP_TASKS_ARG;
|
||||
use app_test_support::create_final_assistant_message_sse_response;
|
||||
use codex_app_server_protocol::InitializeCapabilities;
|
||||
use codex_app_server_protocol::JSONRPCMessage;
|
||||
use codex_app_server_protocol::RequestId;
|
||||
use codex_app_server_protocol::ThreadStartResponse;
|
||||
use codex_uds::UnixStream;
|
||||
use core_test_support::streaming_sse::StreamingSseChunk;
|
||||
use core_test_support::streaming_sse::start_streaming_sse_server;
|
||||
use futures::SinkExt;
|
||||
use futures::StreamExt;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
use std::path::Path;
|
||||
#[cfg(unix)]
|
||||
use std::process::Command as StdCommand;
|
||||
use std::process::Stdio;
|
||||
use tempfile::TempDir;
|
||||
use tokio::process::Child;
|
||||
use tokio::process::Command;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio::time::Duration;
|
||||
use tokio::time::sleep;
|
||||
use tokio::time::timeout;
|
||||
use tokio_tungstenite::WebSocketStream;
|
||||
use tokio_tungstenite::client_async;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
|
||||
#[tokio::test]
|
||||
async fn managed_shutdown_persists_loaded_threads_after_active_turn() -> Result<()> {
|
||||
let home = TempDir::new()?;
|
||||
let (release_response, response_gate) = oneshot::channel();
|
||||
let (mock, _completions) =
|
||||
start_streaming_sse_server(vec![vec![stream_chunk(Some(response_gate), "Done")?]]).await;
|
||||
create_config_toml(home.path(), mock.uri(), "never")?;
|
||||
let socket_path = home.path().join("control/server.sock");
|
||||
let mut server = spawn_server(home.path(), &socket_path)?;
|
||||
let mut client = connect_default_daemon_client(&socket_path).await?;
|
||||
let active = start_thread(&mut client, /*id*/ 2, json!({})).await?;
|
||||
let idle = start_thread(&mut client, /*id*/ 3, json!({})).await?;
|
||||
start_turn(&mut client, /*id*/ 4, &active.thread.id).await?;
|
||||
wait_for_requests(&mock, /*count*/ 1).await?;
|
||||
request_shutdown(&server, &socket_path).await?;
|
||||
assert_still_running(&mut server, "active turn must finish first").await;
|
||||
release_response.send(()).expect("response is waiting");
|
||||
wait_success(&mut server).await?;
|
||||
|
||||
let mut successor = spawn_server(home.path(), &socket_path)?;
|
||||
let mut reconnected = connect_default_daemon_client(&socket_path).await?;
|
||||
for (id, thread) in [(2, &active.thread), (3, &idle.thread)] {
|
||||
let resumed = request(
|
||||
&mut reconnected,
|
||||
id,
|
||||
"thread/resume",
|
||||
json!({"threadId":thread.id}),
|
||||
)
|
||||
.await?;
|
||||
assert_eq!(resumed["thread"]["id"], thread.id);
|
||||
if thread.id == active.thread.id {
|
||||
assert!(
|
||||
resumed["thread"]["turns"][0]["items"]
|
||||
.as_array()
|
||||
.context("recovered turn items")?
|
||||
.iter()
|
||||
.any(|item| item["type"] == "agentMessage" && item["text"] == "Done")
|
||||
);
|
||||
}
|
||||
}
|
||||
request_shutdown(&successor, &socket_path).await?;
|
||||
wait_success(&mut successor).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn managed_force_shutdown_exits_with_active_work() -> Result<()> {
|
||||
let home = TempDir::new()?;
|
||||
let (_release_response, response_gate) = oneshot::channel();
|
||||
let (mock, _completions) =
|
||||
start_streaming_sse_server(vec![vec![stream_chunk(Some(response_gate), "Done")?]]).await;
|
||||
create_config_toml(home.path(), mock.uri(), "never")?;
|
||||
let socket_path = home.path().join("control/server.sock");
|
||||
let mut server = spawn_server(home.path(), &socket_path)?;
|
||||
let mut client = connect_default_daemon_client(&socket_path).await?;
|
||||
let thread = start_thread(&mut client, /*id*/ 2, json!({})).await?;
|
||||
start_turn(&mut client, /*id*/ 3, &thread.thread.id).await?;
|
||||
wait_for_requests(&mock, /*count*/ 1).await?;
|
||||
request_shutdown(&server, &socket_path).await?;
|
||||
assert_still_running(&mut server, "graceful shutdown must wait").await;
|
||||
request_shutdown(&server, &socket_path).await?;
|
||||
wait_success(&mut server).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn managed_force_shutdown_exits_with_blocked_rollout_writer() -> Result<()> {
|
||||
let home = TempDir::new()?;
|
||||
let mock = wiremock::MockServer::start().await;
|
||||
create_config_toml(home.path(), &mock.uri(), "never")?;
|
||||
let socket_path = home.path().join("control/server.sock");
|
||||
let mut server = spawn_server(home.path(), &socket_path)?;
|
||||
let mut client = connect_default_daemon_client(&socket_path).await?;
|
||||
let thread = start_thread(&mut client, /*id*/ 2, json!({})).await?;
|
||||
let rollout = thread
|
||||
.thread
|
||||
.path
|
||||
.context("persistent thread rollout path")?;
|
||||
assert!(!rollout.exists(), "idle rollout should still be deferred");
|
||||
let compressed = rollout.with_extension("jsonl.zst");
|
||||
std::fs::create_dir_all(compressed.parent().context("rollout directory")?)?;
|
||||
assert!(
|
||||
StdCommand::new("mkfifo")
|
||||
.arg(&compressed)
|
||||
.status()?
|
||||
.success()
|
||||
);
|
||||
|
||||
request_shutdown(&server, &socket_path).await?;
|
||||
// Opening the write end succeeds only once the real rollout writer opens
|
||||
// the compressed input. Keep it empty and open to block that writer's read.
|
||||
let _blocked_input = timeout(Duration::from_secs(10), async {
|
||||
loop {
|
||||
if let Ok(sender) = tokio::net::unix::pipe::OpenOptions::new().open_sender(&compressed)
|
||||
{
|
||||
return sender;
|
||||
}
|
||||
sleep(Duration::from_millis(25)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.context("rollout writer did not open the blocked input")?;
|
||||
assert_still_running(&mut server, "graceful shutdown must wait for the writer").await;
|
||||
request_shutdown(&server, &socket_path).await?;
|
||||
wait_success(&mut server).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn wait_success(server: &mut Child) -> Result<()> {
|
||||
anyhow::ensure!(
|
||||
timeout(Duration::from_secs(10), server.wait())
|
||||
.await??
|
||||
.success(),
|
||||
"server did not exit successfully"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn assert_still_running(server: &mut Child, message: &str) {
|
||||
assert!(
|
||||
timeout(Duration::from_millis(300), server.wait())
|
||||
.await
|
||||
.is_err(),
|
||||
"{message}"
|
||||
);
|
||||
}
|
||||
|
||||
fn stream_chunk(gate: Option<oneshot::Receiver<()>>, body: &str) -> Result<StreamingSseChunk> {
|
||||
Ok(StreamingSseChunk {
|
||||
gate,
|
||||
body: create_final_assistant_message_sse_response(body)?,
|
||||
})
|
||||
}
|
||||
|
||||
async fn wait_for_requests(
|
||||
mock: &core_test_support::streaming_sse::StreamingSseServer,
|
||||
count: usize,
|
||||
) -> Result<()> {
|
||||
timeout(Duration::from_secs(5), mock.wait_for_request_count(count)).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn start_turn(
|
||||
websocket: &mut WebSocketStream<UnixStream>,
|
||||
id: i64,
|
||||
thread_id: impl serde::Serialize,
|
||||
) -> Result<()> {
|
||||
request(
|
||||
websocket,
|
||||
id,
|
||||
"turn/start",
|
||||
json!({"threadId":thread_id,"input":[{"type":"text","text":"Hello"}]}),
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn start_thread(
|
||||
websocket: &mut WebSocketStream<UnixStream>,
|
||||
id: i64,
|
||||
params: serde_json::Value,
|
||||
) -> Result<ThreadStartResponse> {
|
||||
Ok(serde_json::from_value(
|
||||
request(websocket, id, "thread/start", params).await?,
|
||||
)?)
|
||||
}
|
||||
|
||||
async fn connect_default_daemon_client(socket_path: &Path) -> Result<WebSocketStream<UnixStream>> {
|
||||
connect_daemon_client(socket_path, InitializeCapabilities::default()).await
|
||||
}
|
||||
|
||||
async fn connect_daemon_client(
|
||||
socket_path: &Path,
|
||||
capabilities: InitializeCapabilities,
|
||||
) -> Result<WebSocketStream<UnixStream>> {
|
||||
connect_initialized(socket_path, capabilities, "daemon_recovery_test", "0.1.0").await
|
||||
}
|
||||
|
||||
fn spawn_server(home: &Path, socket_path: &Path) -> Result<Child> {
|
||||
let binary = codex_utils_cargo_bin::cargo_bin("codex-app-server")?;
|
||||
Ok(Command::new(binary)
|
||||
.args(["--listen", &format!("unix://{}", socket_path.display())])
|
||||
.arg(DISABLE_PLUGIN_STARTUP_TASKS_ARG)
|
||||
.env("CODEX_HOME", home)
|
||||
.arg("--managed-daemon")
|
||||
.env(
|
||||
codex_app_server_transport::DAEMON_SHUTDOWN_SOCKET_ENV,
|
||||
socket_path,
|
||||
)
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.kill_on_drop(true)
|
||||
.spawn()?)
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
async fn request_shutdown(server: &Child, _socket_path: &Path) -> Result<()> {
|
||||
let pid = server.id().context("server pid")?;
|
||||
let status = StdCommand::new("kill")
|
||||
.arg("-TERM")
|
||||
.arg(pid.to_string())
|
||||
.status()?;
|
||||
anyhow::ensure!(status.success(), "failed to signal server {pid}: {status}");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
async fn request_shutdown(server: &Child, socket_path: &Path) -> Result<()> {
|
||||
let pid = server.id().context("server pid")?.to_string();
|
||||
timeout(DEFAULT_READ_TIMEOUT, async {
|
||||
let stream = UnixStream::connect(socket_path).await?;
|
||||
let (mut websocket, _) = client_async("ws://localhost/daemon/shutdown", stream).await?;
|
||||
websocket.send(Message::Text(pid.clone().into())).await?;
|
||||
let reply = websocket
|
||||
.next()
|
||||
.await
|
||||
.context("missing shutdown acknowledgment")??;
|
||||
assert_eq!(reply, Message::Text(pid.into()));
|
||||
websocket.close(None).await?;
|
||||
Ok(())
|
||||
})
|
||||
.await?
|
||||
}
|
||||
|
||||
async fn connect_initialized(
|
||||
socket_path: &Path,
|
||||
capabilities: InitializeCapabilities,
|
||||
client_name: &str,
|
||||
client_version: &str,
|
||||
) -> Result<WebSocketStream<UnixStream>> {
|
||||
let mut websocket = timeout(DEFAULT_READ_TIMEOUT, async {
|
||||
loop {
|
||||
if let Ok(stream) = UnixStream::connect(socket_path).await
|
||||
&& let Ok((websocket, _)) = client_async("ws://localhost/", stream).await
|
||||
{
|
||||
return websocket;
|
||||
}
|
||||
sleep(Duration::from_millis(25)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.context("server did not open its control socket")?;
|
||||
request(
|
||||
&mut websocket,
|
||||
/*id*/ 1,
|
||||
"initialize",
|
||||
json!({"clientInfo":{"name":client_name,"version":client_version},"capabilities":capabilities}),
|
||||
)
|
||||
.await?;
|
||||
websocket
|
||||
.send(Message::Text(
|
||||
json!({"method":"initialized"}).to_string().into(),
|
||||
))
|
||||
.await?;
|
||||
Ok(websocket)
|
||||
}
|
||||
|
||||
async fn request(
|
||||
websocket: &mut WebSocketStream<UnixStream>,
|
||||
id: i64,
|
||||
method: &str,
|
||||
params: serde_json::Value,
|
||||
) -> Result<serde_json::Value> {
|
||||
websocket
|
||||
.send(Message::Text(
|
||||
json!({"id":id,"method":method,"params":params})
|
||||
.to_string()
|
||||
.into(),
|
||||
))
|
||||
.await?;
|
||||
timeout(DEFAULT_READ_TIMEOUT, async {
|
||||
loop {
|
||||
let frame = websocket.next().await.context("socket closed")??;
|
||||
let Message::Text(text) = frame else { continue };
|
||||
match serde_json::from_str::<JSONRPCMessage>(&text)? {
|
||||
JSONRPCMessage::Response(response) if response.id == RequestId::Integer(id) => {
|
||||
return Ok(response.result);
|
||||
}
|
||||
JSONRPCMessage::Error(error) if error.id == RequestId::Integer(id) => {
|
||||
anyhow::bail!("{}: {}", method, error.error.message);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.context("timed out waiting for app-server response")?
|
||||
}
|
||||
@@ -25,6 +25,7 @@ mod connection_handling_websocket_unix;
|
||||
mod curated_mcp_sync;
|
||||
mod current_time;
|
||||
mod cyber_access_program;
|
||||
mod daemon_update_recovery;
|
||||
mod daybreak_access;
|
||||
mod dynamic_tools;
|
||||
mod environment_add;
|
||||
|
||||
@@ -577,6 +577,10 @@ struct AppServerCommand {
|
||||
#[arg(long = "remote-control", hide = true)]
|
||||
remote_control: bool,
|
||||
|
||||
/// Save loaded threads during managed daemon shutdown.
|
||||
#[arg(long, hide = true)]
|
||||
managed_daemon: bool,
|
||||
|
||||
/// Controls whether analytics are enabled by default.
|
||||
///
|
||||
/// Analytics are disabled by default for app-server. Users have to explicitly opt in
|
||||
@@ -1319,6 +1323,7 @@ async fn cli_main(
|
||||
listen,
|
||||
stdio,
|
||||
remote_control,
|
||||
managed_daemon,
|
||||
analytics_default_enabled,
|
||||
auth,
|
||||
} = app_server_cli;
|
||||
@@ -1339,6 +1344,7 @@ async fn cli_main(
|
||||
let auth = auth.try_into_settings()?;
|
||||
let runtime_options = codex_app_server::AppServerRuntimeOptions {
|
||||
code_mode_host_transport: code_mode_host.into(),
|
||||
managed_daemon,
|
||||
remote_control_startup_mode: match (remote_control, remote_control_disabled)
|
||||
{
|
||||
(true, _) => {
|
||||
@@ -1353,7 +1359,7 @@ async fn cli_main(
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
codex_app_server::run_main_with_transport_options(
|
||||
let exit = codex_app_server::run_main_with_transport_options(
|
||||
arg0_paths.clone(),
|
||||
root_config_overrides,
|
||||
LoaderOverrides::default(),
|
||||
@@ -1365,6 +1371,10 @@ async fn cli_main(
|
||||
runtime_options,
|
||||
)
|
||||
.await?;
|
||||
if exit == codex_app_server::AppServerExit::Forced {
|
||||
// Runtime teardown can wait forever for blocked rollout I/O.
|
||||
std::process::exit(0);
|
||||
}
|
||||
}
|
||||
Some(AppServerSubcommand::Daemon(daemon_cli)) => match daemon_cli.subcommand {
|
||||
AppServerDaemonSubcommand::Start => {
|
||||
|
||||
@@ -129,7 +129,7 @@ async fn run_foreground_remote_control(
|
||||
..Default::default()
|
||||
};
|
||||
let (stop_rx, stop_signal_task) = foreground_stop_signal();
|
||||
let mut app_server_task = tokio::spawn(codex_app_server::run_main_with_transport_options(
|
||||
let app_server = codex_app_server::run_main_with_transport_options(
|
||||
arg0_paths,
|
||||
root_config_overrides,
|
||||
LoaderOverrides::default(),
|
||||
@@ -139,7 +139,8 @@ async fn run_foreground_remote_control(
|
||||
SessionSource::VSCode,
|
||||
AppServerWebsocketAuthSettings::default(),
|
||||
runtime_options,
|
||||
));
|
||||
);
|
||||
let mut app_server_task = tokio::spawn(async move { app_server.await.map(|_| ()) });
|
||||
|
||||
let summary = match wait_for_foreground_remote_control_start(
|
||||
&mut app_server_task,
|
||||
|
||||
Reference in New Issue
Block a user