From 575409f453afce6387d09676b660647655afb0ab Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Mon, 1 Jun 2026 17:32:55 -0700 Subject: [PATCH] keep going --- codex-rs/core-api/src/lib.rs | 6 ++ codex-rs/ext/extension-api/src/lib.rs | 1 + codex-rs/thread-manager-sample/README.md | 15 ++- codex-rs/thread-manager-sample/src/main.rs | 110 ++++++++++++++++++--- 4 files changed, 113 insertions(+), 19 deletions(-) diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index e87ee82f30..ddf86ab656 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -46,6 +46,10 @@ pub use codex_core::skills::SkillsManager; pub use codex_core::thread_store_from_config; pub use codex_exec_server::EnvironmentManager; pub use codex_exec_server::ExecServerRuntimePaths; +pub use codex_extension_api::ExtensionRegistryBuilder; +pub use codex_extension_api::ThreadIdleInput; +pub use codex_extension_api::ThreadLifecycleContributor; +pub use codex_extension_api::async_trait; pub use codex_extension_api::empty_extension_registry; pub use codex_features::Feature; pub use codex_features::Features; @@ -64,7 +68,9 @@ pub use codex_protocol::config_types::ShellEnvironmentPolicy; pub use codex_protocol::config_types::WebSearchMode; pub use codex_protocol::dynamic_tools::DynamicToolSpec; pub use codex_protocol::error::Result as CodexResult; +pub use codex_protocol::models::ContentItem; pub use codex_protocol::models::PermissionProfile; +pub use codex_protocol::models::ResponseItem; pub use codex_protocol::openai_models::ModelPreset; pub use codex_protocol::protocol::AskForApproval; pub use codex_protocol::protocol::EventMsg; diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 2f00a8dd98..dc6450f1eb 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -3,6 +3,7 @@ mod contributors; mod registry; mod state; +pub use async_trait::async_trait; pub use capabilities::AgentSpawnFuture; pub use capabilities::AgentSpawner; pub use capabilities::ExtensionEventSink; diff --git a/codex-rs/thread-manager-sample/README.md b/codex-rs/thread-manager-sample/README.md index 7b021b7eb0..1656fda6ee 100644 --- a/codex-rs/thread-manager-sample/README.md +++ b/codex-rs/thread-manager-sample/README.md @@ -1,8 +1,8 @@ # ThreadManager Sample -Small one-shot binary that starts a Codex thread with `ThreadManager` from -`codex-core-api`, submits a single user turn, and prints the final assistant -message. +Small binary that starts a Codex thread with `ThreadManager` from +`codex-core-api`, submits a user turn, and prints mapped notifications as +newline-delimited JSON. ```sh cargo run -p codex-thread-manager-sample -- "Say hello" @@ -19,3 +19,12 @@ The prompt can also be piped through stdin: ```sh printf 'Say hello\n' | cargo run -p codex-thread-manager-sample ``` + +To install the demo idle extension, pass `--keep-going-on-idle`. The extension +listens for the thread idle lifecycle event and starts another turn with a +`keep going` user message each time the thread becomes idle. The sample keeps +running until interrupted or until the thread reports an error: + +```sh +cargo run -p codex-thread-manager-sample -- --keep-going-on-idle "Say hello" +``` diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index 9305a6f668..eb26ce1f3a 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -4,6 +4,7 @@ use std::io::IsTerminal; use std::io::Read; use std::io::Write; use std::sync::Arc; +use std::sync::Weak; use anyhow::Context; use anyhow::bail; @@ -20,9 +21,11 @@ use codex_core_api::CodexThread; use codex_core_api::Config; use codex_core_api::ConfigLayerStack; use codex_core_api::Constrained; +use codex_core_api::ContentItem; use codex_core_api::EnvironmentManager; use codex_core_api::EventMsg; use codex_core_api::ExecServerRuntimePaths; +use codex_core_api::ExtensionRegistryBuilder; use codex_core_api::Features; use codex_core_api::GhostSnapshotConfig; use codex_core_api::History; @@ -40,9 +43,13 @@ use codex_core_api::Permissions; use codex_core_api::ProjectConfig; use codex_core_api::RealtimeAudioConfig; use codex_core_api::RealtimeConfig; +use codex_core_api::ResponseItem; use codex_core_api::SessionPickerViewMode; use codex_core_api::SessionSource; use codex_core_api::TerminalResizeReflowConfig; +use codex_core_api::ThreadId; +use codex_core_api::ThreadIdleInput; +use codex_core_api::ThreadLifecycleContributor; use codex_core_api::ThreadManager; use codex_core_api::ThreadStoreConfig; use codex_core_api::ToolSuggestConfig; @@ -53,6 +60,7 @@ use codex_core_api::UriBasedFileOpener; use codex_core_api::UserInput; use codex_core_api::WebSearchMode; use codex_core_api::arg0_dispatch_or_else; +use codex_core_api::async_trait; use codex_core_api::built_in_model_providers; use codex_core_api::empty_extension_registry; use codex_core_api::find_codex_home; @@ -65,18 +73,69 @@ use codex_core_api::thread_store_from_config; #[derive(Debug, Parser)] #[command( name = "codex-thread-manager-sample", - about = "Run one Codex turn through ThreadManager and print mapped notifications as newline-delimited JSON." + about = "Run Codex turns through ThreadManager and print mapped notifications as newline-delimited JSON." )] struct Args { /// Override the model for this run. #[arg(long, value_name = "MODEL")] model: Option, + /// Install a demo extension that sends one "keep going" message when the thread becomes idle. + #[arg(long)] + keep_going_on_idle: bool, + /// Prompt text. If omitted, the prompt is read from piped stdin. #[arg(value_name = "PROMPT", num_args = 0.., trailing_var_arg = true)] prompt: Vec, } +struct KeepGoingOnIdleExtension { + thread_manager: Weak, +} + +impl KeepGoingOnIdleExtension { + fn new(thread_manager: Weak) -> Self { + Self { thread_manager } + } + + async fn send_keep_going_if_idle(&self, input: ThreadIdleInput<'_>) { + let Ok(thread_id) = ThreadId::from_string(input.thread_store.level_id()) else { + tracing::warn!( + "thread idle extension saw invalid thread id {}", + input.thread_store.level_id() + ); + return; + }; + let Some(thread_manager) = self.thread_manager.upgrade() else { + tracing::debug!("skipping idle keep-going because thread manager was dropped"); + return; + }; + let Ok(thread) = thread_manager.get_thread(thread_id).await else { + tracing::debug!("skipping idle keep-going because live thread is unavailable"); + return; + }; + let item = ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: "keep going".to_string(), + }], + phase: None, + }; + + if thread.try_start_turn_if_idle(vec![item]).await.is_err() { + tracing::debug!("skipping idle keep-going because the thread is no longer idle"); + } + } +} + +#[async_trait] +impl ThreadLifecycleContributor for KeepGoingOnIdleExtension { + async fn on_thread_idle(&self, input: ThreadIdleInput<'_>) { + self.send_keep_going_if_idle(input).await; + } +} + fn main() -> anyhow::Result<()> { arg0_dispatch_or_else(run_main) } @@ -87,6 +146,7 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { } let args = Args::parse(); + let keep_going_on_idle = args.keep_going_on_idle; let prompt = if args.prompt.is_empty() { if std::io::stdin().is_terminal() { bail!("no prompt provided; pass a prompt argument or pipe one into stdin"); @@ -120,18 +180,29 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { .await?, ); let installation_id = resolve_installation_id(&config.codex_home).await?; - let thread_manager = ThreadManager::new( - &config, - auth_manager, - SessionSource::Exec, - environment_manager, - empty_extension_registry(), - /*analytics_events_client*/ None, - Arc::clone(&thread_store), - state_db, - installation_id, - /*attestation_provider*/ None, - ); + let thread_manager: Arc = Arc::new_cyclic(|thread_manager| { + let extensions = if keep_going_on_idle { + let mut builder = ExtensionRegistryBuilder::::new(); + builder.thread_lifecycle_contributor(Arc::new(KeepGoingOnIdleExtension::new( + thread_manager.clone(), + ))); + Arc::new(builder.build()) + } else { + empty_extension_registry() + }; + ThreadManager::new( + &config, + auth_manager, + SessionSource::Exec, + environment_manager, + extensions, + /*analytics_events_client*/ None, + Arc::clone(&thread_store), + state_db, + installation_id, + /*attestation_provider*/ None, + ) + }); let NewThread { thread_id, thread, .. @@ -141,7 +212,7 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { .context("start Codex thread")?; let thread_id_string = thread_id.to_string(); - let turn_output = run_turn(&thread, &thread_id_string, prompt).await; + let turn_output = run_turn(&thread, &thread_id_string, prompt, keep_going_on_idle).await; let shutdown_result = thread.shutdown_and_wait().await; let _ = thread_manager.remove_thread(&thread_id).await; @@ -288,7 +359,12 @@ fn new_config(model: Option, arg0_paths: Arg0DispatchPaths) -> anyhow::R Ok(config) } -async fn run_turn(thread: &CodexThread, thread_id: &str, prompt: String) -> anyhow::Result<()> { +async fn run_turn( + thread: &CodexThread, + thread_id: &str, + prompt: String, + keep_running_after_turn_complete: bool, +) -> anyhow::Result<()> { thread .submit(Op::UserInput { items: vec![UserInput::Text { @@ -358,7 +434,9 @@ async fn run_turn(thread: &CodexThread, thread_id: &str, prompt: String) -> anyh match event.msg { EventMsg::TurnComplete(_) => { - return Ok(()); + if !keep_running_after_turn_complete { + return Ok(()); + } } EventMsg::Error(event) => { bail!(event.message);