openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-05 20:10:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Replace thread settings snapshot with idle policy

Browse Source
This commit is contained in:
Eric Traut
2026-05-27 20:19:27 -07:00
parent 48118fa6a0
commit f690ea0b6a
10 changed files with 100 additions and 139 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/codex_thread.rs
-2
View File
@@ -167,7 +167,6 @@ impl CodexThread {
}
pub async fn emit_thread_resume_lifecycle(&self) {
let thread_settings = ThreadSettingsSnapshot::from(self.config_snapshot().await);
for contributor in self
.codex
.session
@@ -177,7 +176,6 @@ impl CodexThread {
{
contributor
.on_thread_resume(codex_extension_api::ThreadResumeInput {
thread_settings: &thread_settings,
session_store: &self.codex.session.services.session_extension_data,
thread_store: &self.codex.session.services.thread_extension_data,
})
codex-rs/core/src/session/mod.rs
-5
View File
@@ -2667,11 +2667,6 @@ impl Session {
state.session_configuration.collaboration_mode.clone()
}
pub(crate) async fn thread_config_snapshot(&self) -> ThreadConfigSnapshot {
let state = self.state.lock().await;
state.session_configuration.thread_config_snapshot()
}
async fn send_raw_response_items(&self, turn_context: &TurnContext, items: &[ResponseItem]) {
for item in items {
self.send_event(
codex-rs/core/src/tasks/idle_extension.rs
+17 -11
View File
@@ -4,7 +4,8 @@ use std::sync::Arc;
use codex_extension_api::ResponseInjectionItem;
use codex_extension_api::ThreadIdleRequest;
use codex_protocol::protocol::ThreadSettingsSnapshot;
use codex_protocol::config_types::CollaborationMode;
use codex_protocol::config_types::ModeKind;
use crate::session::TurnInput;
use crate::session::session::Session;
@@ -102,11 +103,9 @@ pub(crate) async fn maybe_start_turn(session: Arc<Session>) {
}
async fn notify_thread_idle(session: &Session) {
let thread_settings = thread_settings_snapshot(session).await;
for contributor in session.services.extensions.thread_lifecycle_contributors() {
contributor
.on_thread_idle(codex_extension_api::ThreadIdleInput {
thread_settings: &thread_settings,
session_store: &session.services.session_extension_data,
thread_store: &session.services.thread_extension_data,
})
@@ -132,7 +131,7 @@ struct IdleTurnCandidate {
}
async fn next_idle_turn_candidate(session: &Session) -> Option<IdleTurnCandidate> {
let thread_settings = thread_settings_snapshot(session).await;
let collaboration_mode = session.collaboration_mode().await;
for (contributor_index, contributor) in session
.services
.extensions
@@ -140,9 +139,11 @@ async fn next_idle_turn_candidate(session: &Session) -> Option<IdleTurnCandidate
.iter()
.enumerate()
{
if !idle_turn_policy_allows_mode(contributor.idle_turn_policy(), &collaboration_mode) {
continue;
}
let Some(request) = contributor
.request_thread_idle_turn(codex_extension_api::ThreadIdleInput {
thread_settings: &thread_settings,
session_store: &session.services.session_extension_data,
thread_store: &session.services.thread_extension_data,
})
@@ -168,6 +169,13 @@ fn is_non_empty_idle_input(item: &ResponseInjectionItem) -> bool {
}
}
fn idle_turn_policy_allows_mode(
policy: codex_extension_api::IdleTurnPolicy,
collaboration_mode: &CollaborationMode,
) -> bool {
!matches!(collaboration_mode.mode, ModeKind::Plan) || policy.allows_plan_mode()
}
async fn should_start_idle_turn(session: &Session, candidate: &IdleTurnCandidate) -> bool {
let Some(contributor) = session
.services
@@ -177,10 +185,12 @@ async fn should_start_idle_turn(session: &Session, candidate: &IdleTurnCandidate
else {
return false;
};
let thread_settings = thread_settings_snapshot(session).await;
let collaboration_mode = session.collaboration_mode().await;
if !idle_turn_policy_allows_mode(contributor.idle_turn_policy(), &collaboration_mode) {
return false;
}
contributor
.should_start_thread_idle_turn(codex_extension_api::ThreadIdleTurnStartInput {
thread_settings: &thread_settings,
request: &candidate.request,
session_store: &session.services.session_extension_data,
thread_store: &session.services.thread_extension_data,
@@ -188,10 +198,6 @@ async fn should_start_idle_turn(session: &Session, candidate: &IdleTurnCandidate
.await
}
async fn thread_settings_snapshot(session: &Session) -> ThreadSettingsSnapshot {
ThreadSettingsSnapshot::from(session.thread_config_snapshot().await)
}
async fn reserve_idle_turn(session: &Session) -> bool {
let mut active_turn = session.active_turn.lock().await;
if active_turn.is_some() {
codex-rs/ext/extension-api/src/contributors.rs
+8 -1
View File
@@ -19,6 +19,7 @@ mod turn_lifecycle;
pub use prompt::PromptFragment;
pub use prompt::PromptSlot;
pub use thread_lifecycle::IdleTurnPolicy;
pub use thread_lifecycle::ThreadIdleInput;
pub use thread_lifecycle::ThreadIdleRequest;
pub use thread_lifecycle::ThreadIdleTurnStartInput;
@@ -80,9 +81,15 @@ pub type ThreadIdleTurnStartFuture<'a> = std::pin::Pin<Box<dyn Future<Output = b
///
/// Implementations should normally return hidden context. Raw items are
/// available for extensions that intentionally need unwrapped input. The host
/// owns hidden-context wrapping and may ask the contributor to confirm the
/// owns hidden-context wrapping, applies the declared idle-turn policy before
/// calling into the contributor, and may ask the contributor to confirm the
/// request again immediately before the turn starts.
pub trait ThreadIdleTurnContributor: Send + Sync {
/// Returns the host scheduling policy for this contributor's idle turns.
fn idle_turn_policy(&self) -> IdleTurnPolicy {
IdleTurnPolicy::default()
}
/// Returns input to start an idle turn, if the extension still has work
/// that should run without user input.
fn request_thread_idle_turn<'a>(
codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs
+20 -7
View File
@@ -3,7 +3,26 @@ use std::sync::Arc;
use crate::ExtensionData;
use crate::ResponseInjectionItem;
use crate::ResponseItemInjector;
use codex_protocol::protocol::ThreadSettingsSnapshot;
/// Idle-turn scheduling policy declared by an extension contributor.
///
/// The default policy allows idle turns only outside plan mode.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct IdleTurnPolicy {
allow_plan_mode: bool,
}
impl IdleTurnPolicy {
pub const fn allow_plan_mode() -> Self {
Self {
allow_plan_mode: true,
}
}
pub fn allows_plan_mode(self) -> bool {
self.allow_plan_mode
}
}
/// Input supplied when the host starts a runtime for a thread.
pub struct ThreadStartInput<'a, C> {
@@ -20,8 +39,6 @@ pub struct ThreadStartInput<'a, C> {
/// Input supplied when the host resumes an existing thread.
pub struct ThreadResumeInput<'a> {
/// Current host-owned thread settings at resume time.
pub thread_settings: &'a ThreadSettingsSnapshot,
/// Store scoped to the host session runtime.
pub session_store: &'a ExtensionData,
/// Store scoped to this thread runtime.
@@ -30,8 +47,6 @@ pub struct ThreadResumeInput<'a> {
/// Input supplied when the host has no immediately pending thread work.
pub struct ThreadIdleInput<'a> {
/// Current host-owned thread settings for the idle thread.
pub thread_settings: &'a ThreadSettingsSnapshot,
/// Store scoped to the host session runtime.
pub session_store: &'a ExtensionData,
/// Store scoped to this thread runtime.
@@ -62,8 +77,6 @@ impl ThreadIdleRequest {
/// Input supplied before the host starts an extension-requested idle turn.
pub struct ThreadIdleTurnStartInput<'a> {
/// Host-owned thread settings for the default turn being started.
pub thread_settings: &'a ThreadSettingsSnapshot,
/// Request returned by this contributor for the candidate idle turn.
pub request: &'a ThreadIdleRequest,
/// Store scoped to the host session runtime.
codex-rs/ext/extension-api/src/lib.rs
+1
View File
@@ -30,6 +30,7 @@ pub use context::HiddenContextMarker;
pub use contributors::ApprovalReviewContributor;
pub use contributors::ConfigContributor;
pub use contributors::ContextContributor;
pub use contributors::IdleTurnPolicy;
pub use contributors::PromptFragment;
pub use contributors::PromptSlot;
pub use contributors::ThreadIdleInput;
codex-rs/ext/goal/src/accounting.rs
+3 -1
View File
@@ -158,13 +158,15 @@ impl GoalAccountingState {
Some(turn_id)
}
pub(crate) fn mark_idle_goal_active(&self, goal_id: impl Into<String>) {
pub(crate) fn mark_idle_goal_active(&self, goal_id: impl Into<String>) -> bool {
let mut inner = self.inner();
let goal_id = goal_id.into();
let was_active = inner.wall_clock.active_goal_id.as_deref() == Some(goal_id.as_str());
if inner.budget_limit_reported_goal_id.as_deref() != Some(goal_id.as_str()) {
inner.budget_limit_reported_goal_id = None;
}
inner.wall_clock.mark_active_goal(goal_id);
!was_active
}
pub(crate) fn clear_current_turn_goal(&self) -> Option<String> {
codex-rs/ext/goal/src/extension.rs
+2 -22
View File
@@ -11,7 +11,6 @@ use codex_extension_api::ThreadIdleTurnRequestFuture;
use codex_extension_api::ThreadIdleTurnStartFuture;
use codex_extension_api::ThreadIdleTurnStartInput;
use codex_extension_api::ThreadLifecycleContributor;
use codex_extension_api::ThreadResumeInput;
use codex_extension_api::ThreadStartInput;
use codex_extension_api::TokenUsageContributor;
use codex_extension_api::ToolCallOutcome;
@@ -98,19 +97,6 @@ where
});
runtime.set_enabled(enabled);
}
async fn on_thread_resume(&self, input: ThreadResumeInput<'_>) {
let Some(runtime) = goal_runtime_handle(input.thread_store) else {
return;
};
if let Err(err) = runtime.restore_after_resume(input.thread_settings).await {
tracing::warn!(
"failed to restore goal runtime after thread resume for {}: {err}",
runtime.thread_id()
);
}
}
}
impl<C> ThreadIdleTurnContributor for GoalExtension<C>
@@ -123,10 +109,7 @@ where
) -> ThreadIdleTurnRequestFuture<'a> {
Box::pin(async move {
let runtime = goal_runtime_handle(input.thread_store)?;
match runtime
.idle_continuation_request(input.thread_settings.collaboration_mode.mode)
.await
{
match runtime.idle_continuation_request().await {
Ok(request) => request,
Err(err) => {
tracing::warn!("failed to request idle goal continuation: {err}");
@@ -145,10 +128,7 @@ where
return false;
};
match runtime
.idle_continuation_is_current(
input.thread_settings.collaboration_mode.mode,
input.request.validation_key.as_deref(),
)
.idle_continuation_is_current(input.request.validation_key.as_deref())
.await
{
Ok(should_start) => should_start,
codex-rs/ext/goal/src/runtime.rs
+7 -38
View File
@@ -5,9 +5,7 @@ use std::sync::atomic::Ordering;
use codex_extension_api::HiddenContext;
use codex_extension_api::ResponseItemInjector;
use codex_protocol::ThreadId;
use codex_protocol::config_types::ModeKind;
use codex_protocol::protocol::ThreadGoal;
use codex_protocol::protocol::ThreadSettingsSnapshot;
use crate::accounting::BudgetLimitedGoalDisposition;
use crate::accounting::GoalAccountingState;
@@ -227,42 +225,10 @@ impl GoalRuntimeHandle {
Ok(())
}
pub async fn restore_after_resume(
&self,
thread_settings: &ThreadSettingsSnapshot,
) -> Result<(), String> {
if !self.is_enabled() {
return Ok(());
}
if matches!(thread_settings.collaboration_mode.mode, ModeKind::Plan) {
self.inner.accounting_state.clear_active_goal();
return Ok(());
}
let goal = self
.inner
.state_dbs
.thread_goals()
.get_thread_goal(self.thread_id())
.await
.map_err(|err| err.to_string())?;
match goal {
Some(goal) if goal.status == codex_state::ThreadGoalStatus::Active => {
self.inner
.accounting_state
.mark_idle_goal_active(goal.goal_id);
self.inner.metrics.record_resumed();
}
Some(_) | None => self.inner.accounting_state.clear_active_goal(),
}
Ok(())
}
pub async fn idle_continuation_request(
&self,
collaboration_mode: ModeKind,
) -> Result<Option<codex_extension_api::ThreadIdleRequest>, String> {
if !self.is_enabled() || matches!(collaboration_mode, ModeKind::Plan) {
if !self.is_enabled() {
return Ok(None);
}
@@ -282,9 +248,13 @@ impl GoalRuntimeHandle {
return Ok(None);
}
self.inner
let resumed_idle_goal = self
.inner
.accounting_state
.mark_idle_goal_active(goal.goal_id.clone());
if resumed_idle_goal {
self.inner.metrics.record_resumed();
}
let goal_id = goal.goal_id.clone();
let goal = protocol_goal_from_state(goal);
let request = continuation_steering_request(&goal).with_validation_key(goal_id);
@@ -293,10 +263,9 @@ impl GoalRuntimeHandle {
pub async fn idle_continuation_is_current(
&self,
collaboration_mode: ModeKind,
expected_goal_id: Option<&str>,
) -> Result<bool, String> {
if !self.is_enabled() || matches!(collaboration_mode, ModeKind::Plan) {
if !self.is_enabled() {
return Ok(false);
}
let Some(expected_goal_id) = expected_goal_id else {
codex-rs/ext/goal/tests/goal_extension_backend.rs
+42 -52
View File
@@ -12,7 +12,6 @@ use codex_extension_api::ResponseInjectionItem;
use codex_extension_api::ThreadIdleInput;
use codex_extension_api::ThreadIdleRequest;
use codex_extension_api::ThreadIdleTurnStartInput;
use codex_extension_api::ThreadResumeInput;
use codex_extension_api::ThreadStartInput;
use codex_extension_api::ToolCall;
use codex_extension_api::ToolCallOutcome;
@@ -27,19 +26,15 @@ use codex_extension_api::TurnStopInput;
use codex_goal_extension::GoalRuntimeHandle;
use codex_goal_extension::install_with_backend;
use codex_protocol::ThreadId;
use codex_protocol::config_types::ApprovalsReviewer;
use codex_protocol::config_types::CollaborationMode;
use codex_protocol::config_types::ModeKind;
use codex_protocol::config_types::Settings;
use codex_protocol::models::PermissionProfile;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::CodexErrorInfo;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::protocol::ThreadGoalStatus;
use codex_protocol::protocol::ThreadSettingsSnapshot;
use codex_protocol::protocol::TokenUsage;
use codex_protocol::protocol::TokenUsageInfo;
use codex_protocol::protocol::TruncationPolicy;
@@ -522,6 +517,11 @@ async fn thread_idle_hook_requests_hidden_goal_continuation() -> anyhow::Result<
.should_start_idle_turn(ModeKind::Default, &request)
.await
);
assert!(
!harness
.should_start_idle_turn(ModeKind::Plan, &request)
.await
);
Ok(())
}
@@ -1028,7 +1028,7 @@ async fn external_goal_set_active_resets_baseline_without_live_thread() -> anyho
}
#[tokio::test]
async fn thread_resume_rehydrates_active_goal_idle_accounting() -> anyhow::Result<()> {
async fn idle_continuation_request_rehydrates_active_goal_idle_accounting() -> anyhow::Result<()> {
let runtime = test_runtime().await?;
let thread_id = test_thread_id()?;
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
@@ -1043,7 +1043,28 @@ async fn thread_resume_rehydrates_active_goal_idle_accounting() -> anyhow::Resul
.await?;
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
harness.resume_thread().await;
tokio::time::sleep(Duration::from_millis(1_100)).await;
harness
.runtime_handle()
.prepare_external_goal_mutation()
.await
.map_err(anyhow::Error::msg)?;
let goal = runtime
.thread_goals()
.get_thread_goal(thread_id)
.await?
.ok_or_else(|| anyhow::anyhow!("goal should exist"))?;
assert_eq!(ThreadGoalStatus::Active, protocol_status(goal.status));
assert_eq!(
0, goal.time_used_seconds,
"resume alone should not rehydrate idle accounting"
);
let _request = harness
.next_idle_turn(ModeKind::Default)
.await
.ok_or_else(|| anyhow::anyhow!("active goal should request idle continuation"))?;
tokio::time::sleep(Duration::from_millis(1_100)).await;
harness
.runtime_handle()
@@ -1059,7 +1080,7 @@ async fn thread_resume_rehydrates_active_goal_idle_accounting() -> anyhow::Resul
assert_eq!(ThreadGoalStatus::Active, protocol_status(goal.status));
assert!(
goal.time_used_seconds >= 1,
"resumed idle accounting should add elapsed wall-clock time"
"idle continuation request should add elapsed wall-clock time"
);
assert_eq!(
vec![CapturedGoalEvent {
@@ -1074,7 +1095,7 @@ async fn thread_resume_rehydrates_active_goal_idle_accounting() -> anyhow::Resul
}
#[tokio::test]
async fn thread_resume_in_plan_mode_does_not_rehydrate_idle_accounting() -> anyhow::Result<()> {
async fn plan_mode_idle_policy_does_not_rehydrate_idle_accounting() -> anyhow::Result<()> {
let runtime = test_runtime().await?;
let thread_id = test_thread_id()?;
seed_thread_metadata(runtime.as_ref(), thread_id).await?;
@@ -1089,7 +1110,7 @@ async fn thread_resume_in_plan_mode_does_not_rehydrate_idle_accounting() -> anyh
.await?;
let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?;
harness.resume_thread_with_mode(ModeKind::Plan).await;
assert!(harness.next_idle_turn(ModeKind::Plan).await.is_none());
tokio::time::sleep(Duration::from_millis(1_100)).await;
harness
.runtime_handle()
@@ -1266,23 +1287,6 @@ impl GoalExtensionHarness {
}
}
async fn resume_thread(&self) {
self.resume_thread_with_mode(ModeKind::Default).await;
}
async fn resume_thread_with_mode(&self, mode: ModeKind) {
let thread_settings = thread_settings(mode);
for contributor in self.registry.thread_lifecycle_contributors() {
contributor
.on_thread_resume(ThreadResumeInput {
thread_settings: &thread_settings,
session_store: &self.session_store,
thread_store: &self.thread_store,
})
.await;
}
}
async fn notify_tool_finish(&self, turn_id: &str, call_id: &str, tool_name: &str) {
let turn_store = ExtensionData::new(turn_id);
let tool_name = codex_extension_api::ToolName::plain(tool_name);
@@ -1318,11 +1322,12 @@ impl GoalExtensionHarness {
}
async fn next_idle_turn(&self, mode: ModeKind) -> Option<ThreadIdleRequest> {
let thread_settings = thread_settings(mode);
for contributor in self.registry.thread_idle_turn_contributors() {
if !idle_turn_policy_allows_mode(contributor.idle_turn_policy(), mode) {
continue;
}
if let Some(request) = contributor
.request_thread_idle_turn(ThreadIdleInput {
thread_settings: &thread_settings,
session_store: &self.session_store,
thread_store: &self.thread_store,
})
@@ -1335,11 +1340,12 @@ impl GoalExtensionHarness {
}
async fn should_start_idle_turn(&self, mode: ModeKind, request: &ThreadIdleRequest) -> bool {
let thread_settings = thread_settings(mode);
for contributor in self.registry.thread_idle_turn_contributors() {
if !idle_turn_policy_allows_mode(contributor.idle_turn_policy(), mode) {
continue;
}
if contributor
.should_start_thread_idle_turn(ThreadIdleTurnStartInput {
thread_settings: &thread_settings,
request,
session_store: &self.session_store,
thread_store: &self.thread_store,
@@ -1462,27 +1468,11 @@ fn default_collaboration_mode() -> CollaborationMode {
}
}
fn thread_settings(mode: ModeKind) -> ThreadSettingsSnapshot {
let mut collaboration_mode = default_collaboration_mode();
collaboration_mode.mode = mode;
let cwd = match std::env::current_dir().and_then(TryInto::try_into) {
Ok(cwd) => cwd,
Err(err) => panic!("test cwd should be absolute: {err}"),
};
ThreadSettingsSnapshot {
model: collaboration_mode.model().to_string(),
model_provider_id: "test-provider".to_string(),
service_tier: None,
approval_policy: AskForApproval::Never,
approvals_reviewer: ApprovalsReviewer::User,
permission_profile: PermissionProfile::workspace_write(),
active_permission_profile: None,
cwd,
reasoning_effort: None,
reasoning_summary: None,
personality: None,
collaboration_mode,
}
fn idle_turn_policy_allows_mode(
policy: codex_extension_api::IdleTurnPolicy,
mode: ModeKind,
) -> bool {
!matches!(mode, ModeKind::Plan) || policy.allows_plan_mode()
}
fn token_usage(