diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index c754b0d13e..3935ad957e 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2970,7 +2970,6 @@ dependencies = [ "anyhow", "async-trait", "chrono", - "codex-core", "codex-extension-api", "codex-otel", "codex-protocol", diff --git a/codex-rs/app-server/src/extensions.rs b/codex-rs/app-server/src/extensions.rs index 6cdd4ac206..ca1defe147 100644 --- a/codex-rs/app-server/src/extensions.rs +++ b/codex-rs/app-server/src/extensions.rs @@ -13,10 +13,13 @@ use codex_extension_api::ExtensionEventFuture; use codex_extension_api::ExtensionEventSink; use codex_extension_api::ExtensionRegistry; use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::ResponseItemInjectionFuture; +use codex_extension_api::ResponseItemInjector; use codex_features::Feature; use codex_login::AuthManager; use codex_protocol::ThreadId; use codex_protocol::error::CodexErr; +use codex_protocol::models::ResponseInputItem; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_rollout::state_db::StateDbHandle; @@ -40,7 +43,7 @@ where &mut builder, state_db, codex_otel::global(), - thread_manager, + Arc::new(ThreadManagerResponseItemInjector { thread_manager }), |config: &Config| config.features.enabled(Feature::Goals), ); } @@ -49,6 +52,29 @@ where Arc::new(builder.build()) } +struct ThreadManagerResponseItemInjector { + thread_manager: Weak, +} + +impl ResponseItemInjector for ThreadManagerResponseItemInjector { + fn inject_response_items<'a>( + &'a self, + thread_id: ThreadId, + items: Vec, + ) -> ResponseItemInjectionFuture<'a> { + let thread_manager = self.thread_manager.clone(); + Box::pin(async move { + let Some(thread_manager) = thread_manager.upgrade() else { + return Err(items); + }; + let Ok(thread) = thread_manager.get_thread(thread_id).await else { + return Err(items); + }; + thread.inject_response_items_into_active_turn(items).await + }) + } +} + pub(crate) fn app_server_extension_event_sink( outgoing: Arc, ) -> Arc { diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 26771ab550..b1c9f3e3f8 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -176,13 +176,11 @@ impl CodexThread { .thread_lifecycle_contributors() { contributor - .on_thread_resume_with_settings( - codex_extension_api::ThreadResumeWithSettingsInput { - thread_settings: &thread_settings, - session_store: &self.codex.session.services.session_extension_data, - thread_store: &self.codex.session.services.thread_extension_data, - }, - ) + .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, + }) .await; } } diff --git a/codex-rs/core/src/context/contextual_user_message.rs b/codex-rs/core/src/context/contextual_user_message.rs index afc1de9b09..2a7ee4d674 100644 --- a/codex-rs/core/src/context/contextual_user_message.rs +++ b/codex-rs/core/src/context/contextual_user_message.rs @@ -4,9 +4,9 @@ use codex_protocol::models::ContentItem; use super::AdditionalContextUserFragment; use super::EnvironmentContext; -use super::ExtensionContext; use super::FragmentRegistration; use super::FragmentRegistrationProxy; +use super::GoalContext; use super::LegacyApplyPatchExecCommandWarning; use super::LegacyModelMismatchWarning; use super::LegacyUnifiedExecProcessLimitWarning; @@ -22,7 +22,7 @@ static ENVIRONMENT_CONTEXT_REGISTRATION: FragmentRegistrationProxy = FragmentRegistrationProxy::new(); -static EXTENSION_CONTEXT_REGISTRATION: FragmentRegistrationProxy = +static GOAL_CONTEXT_REGISTRATION: FragmentRegistrationProxy = FragmentRegistrationProxy::new(); static SKILL_INSTRUCTIONS_REGISTRATION: FragmentRegistrationProxy = FragmentRegistrationProxy::new(); @@ -46,7 +46,7 @@ static CONTEXTUAL_USER_FRAGMENTS: &[&dyn FragmentRegistration] = &[ &USER_INSTRUCTIONS_REGISTRATION, &ENVIRONMENT_CONTEXT_REGISTRATION, &ADDITIONAL_CONTEXT_REGISTRATION, - &EXTENSION_CONTEXT_REGISTRATION, + &GOAL_CONTEXT_REGISTRATION, &SKILL_INSTRUCTIONS_REGISTRATION, &USER_SHELL_COMMAND_REGISTRATION, &TURN_ABORTED_REGISTRATION, diff --git a/codex-rs/core/src/context/contextual_user_message_tests.rs b/codex-rs/core/src/context/contextual_user_message_tests.rs index a4ccbbd738..cfba3c2320 100644 --- a/codex-rs/core/src/context/contextual_user_message_tests.rs +++ b/codex-rs/core/src/context/contextual_user_message_tests.rs @@ -14,9 +14,16 @@ fn detects_environment_context_fragment() { } #[test] -fn detects_extension_context_fragment() { +fn detects_legacy_goal_context_fragment() { assert!(is_contextual_user_fragment(&ContentItem::InputText { - text: "\ndata\n".to_string(), + text: "\nlegacy hidden goal prompt\n".to_string(), + })); +} + +#[test] +fn does_not_treat_generic_extension_context_as_hidden() { + assert!(!is_contextual_user_fragment(&ContentItem::InputText { + text: "\nuser-authored text\n".to_string(), })); } diff --git a/codex-rs/core/src/context/extension_context.rs b/codex-rs/core/src/context/extension_context.rs deleted file mode 100644 index 15e72191a0..0000000000 --- a/codex-rs/core/src/context/extension_context.rs +++ /dev/null @@ -1,44 +0,0 @@ -//! Hidden user-context wrappers owned by extensions. -//! -//! Extensions can use this host-recognized tag when they need hidden -//! model-visible context to survive history filtering. - -use super::ContextualUserFragment; -use super::fragment::matches_marked_text; - -const START_MARKER: &str = ""; -const END_MARKER: &str = ""; - -/// Hidden user-context fragment for extension-owned steering prompts. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ExtensionContext { - body: String, -} - -impl ExtensionContext { - pub fn new(body: impl Into) -> Self { - Self { body: body.into() } - } -} - -impl ContextualUserFragment for ExtensionContext { - fn role() -> &'static str { - "user" - } - - fn markers(&self) -> (&'static str, &'static str) { - (START_MARKER, END_MARKER) - } - - fn type_markers() -> (&'static str, &'static str) { - (START_MARKER, END_MARKER) - } - - fn matches_text(text: &str) -> bool { - matches_marked_text(START_MARKER, END_MARKER, text) - } - - fn body(&self) -> String { - format!("\n{}\n", self.body) - } -} diff --git a/codex-rs/core/src/context/goal_context.rs b/codex-rs/core/src/context/goal_context.rs new file mode 100644 index 0000000000..fa84b0ed5a --- /dev/null +++ b/codex-rs/core/src/context/goal_context.rs @@ -0,0 +1,29 @@ +//! Legacy hidden user-context marker for goal steering prompts. +//! +//! The goal implementation now lives in `codex-goal-extension`, but existing +//! rollouts can still contain this marker. Core keeps recognizing it so old +//! hidden goal prompts do not resurface as normal user messages after resume. + +use super::ContextualUserFragment; + +/// Marker-only registration for hidden goal steering context. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct GoalContext; + +impl ContextualUserFragment for GoalContext { + fn role() -> &'static str { + "user" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn body(&self) -> String { + String::new() + } + + fn type_markers() -> (&'static str, &'static str) { + ("", "") + } +} diff --git a/codex-rs/core/src/context/mod.rs b/codex-rs/core/src/context/mod.rs index 5ad1a8711e..70b4dda424 100644 --- a/codex-rs/core/src/context/mod.rs +++ b/codex-rs/core/src/context/mod.rs @@ -7,9 +7,9 @@ mod available_skills_instructions; mod collaboration_mode_instructions; mod contextual_user_message; mod environment_context; -mod extension_context; mod fragment; mod fragments; +mod goal_context; mod guardian_followup_review_reminder; mod hook_additional_context; mod image_generation_instructions; @@ -38,12 +38,12 @@ pub(crate) use collaboration_mode_instructions::CollaborationModeInstructions; pub(crate) use contextual_user_message::is_contextual_user_fragment; pub(crate) use contextual_user_message::parse_visible_hook_prompt_message; pub(crate) use environment_context::EnvironmentContext; -pub use extension_context::ExtensionContext; pub use fragment::ContextualUserFragment; pub(crate) use fragment::FragmentRegistration; pub(crate) use fragment::FragmentRegistrationProxy; pub(crate) use fragments::AdditionalContextDeveloperFragment; pub(crate) use fragments::AdditionalContextUserFragment; +pub(crate) use goal_context::GoalContext; pub(crate) use guardian_followup_review_reminder::GuardianFollowupReviewReminder; pub(crate) use hook_additional_context::HookAdditionalContext; pub(crate) use image_generation_instructions::ImageGenerationInstructions; diff --git a/codex-rs/core/src/session/input_queue.rs b/codex-rs/core/src/session/input_queue.rs index 620c410eb8..3ccb9701f2 100644 --- a/codex-rs/core/src/session/input_queue.rs +++ b/codex-rs/core/src/session/input_queue.rs @@ -3,6 +3,7 @@ use crate::state::MailboxDeliveryPhase; use crate::state::TurnState; use codex_protocol::models::ResponseInputItem; use codex_protocol::protocol::InterAgentCommunication; +use codex_protocol::protocol::TokenUsage; use codex_protocol::user_input::UserInput; use std::collections::VecDeque; use std::sync::Arc; @@ -181,6 +182,42 @@ impl InputQueue { turn_state.lock().await.pending_input.items.split_off(0) } + #[expect( + clippy::await_holding_invalid_type, + reason = "startup input must move only while the claimed active turn is still current" + )] + pub(crate) async fn prepare_starting_turn_input( + &self, + active_turn: &mut ActiveTurn, + token_usage_at_turn_start: TokenUsage, + ) { + let queued_response_items = self.take_queued_response_items_for_next_turn().await; + let mut turn_state = active_turn.turn_state.lock().await; + turn_state.token_usage_at_turn_start = token_usage_at_turn_start; + let existing_pending_input = turn_state.pending_input.items.split_off(0); + let accepts_mailbox_delivery = turn_state.accepts_mailbox_delivery_for_current_turn(); + turn_state.pending_input.items.extend( + queued_response_items + .into_iter() + .map(TurnInput::ResponseInputItem), + ); + turn_state + .pending_input + .items + .extend(existing_pending_input); + drop(turn_state); + if accepts_mailbox_delivery { + let mailbox_items = self.drain_mailbox_input_items().await; + active_turn + .turn_state + .lock() + .await + .pending_input + .items + .extend(mailbox_items.into_iter().map(TurnInput::ResponseInputItem)); + } + } + #[expect( clippy::await_holding_invalid_type, reason = "active turn checks and turn state updates must remain atomic" diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 65a8ab0d66..0071cc1cf6 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -8271,8 +8271,10 @@ async fn start_task_does_not_overwrite_racing_active_task() { #[tokio::test] async fn start_task_does_not_drain_input_while_another_start_is_claiming_turn() { let (sess, tc, _rx) = make_session_and_context_with_rx().await; - let mut active_turn = ActiveTurn::default(); - active_turn.task_starting = true; + let active_turn = ActiveTurn { + task_starting: true, + ..Default::default() + }; *sess.active_turn.lock().await = Some(active_turn); let queued_item = ResponseInputItem::Message { @@ -8310,6 +8312,70 @@ async fn start_task_does_not_drain_input_while_another_start_is_claiming_turn() ); } +#[tokio::test] +async fn abort_during_turn_start_lifecycle_preserves_queued_input() { + struct BlockingTurnStart { + entered: Arc, + release: Arc, + } + + #[async_trait::async_trait] + impl codex_extension_api::TurnLifecycleContributor for BlockingTurnStart { + async fn on_turn_start(&self, _input: codex_extension_api::TurnStartInput<'_>) { + self.entered.notify_waiters(); + self.release.notified().await; + } + } + + let (mut session, turn_context) = make_session_and_context().await; + let entered = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let mut builder = codex_extension_api::ExtensionRegistryBuilder::::new(); + builder.turn_lifecycle_contributor(Arc::new(BlockingTurnStart { + entered: Arc::clone(&entered), + release: Arc::clone(&release), + })); + session.services.extensions = Arc::new(builder.build()); + let sess = Arc::new(session); + let queued_item = ResponseInputItem::Message { + role: "assistant".to_string(), + content: vec![ContentItem::InputText { + text: "queued before lifecycle".to_string(), + }], + phase: None, + }; + sess.input_queue + .queue_response_items_for_next_turn(vec![queued_item]) + .await; + + let start_task = tokio::spawn({ + let sess = Arc::clone(&sess); + async move { + sess.start_task( + Arc::new(turn_context), + Vec::new(), + NeverEndingTask { + kind: TaskKind::Regular, + listen_to_cancellation_token: true, + }, + ) + .await; + } + }); + + entered.notified().await; + sess.abort_all_tasks(TurnAbortReason::Interrupted).await; + release.notify_waiters(); + start_task.await.expect("start task should finish"); + + assert!( + sess.input_queue + .has_queued_response_items_for_next_turn() + .await + ); + assert!(sess.active_turn.lock().await.is_none()); +} + #[tokio::test] async fn idle_interrupt_does_not_wake_queued_next_turn_items() { let (sess, _tc, _rx) = make_session_and_context_with_rx().await; diff --git a/codex-rs/core/src/tasks/idle_extension.rs b/codex-rs/core/src/tasks/idle_extension.rs index eadc25ae5d..b6aeac059c 100644 --- a/codex-rs/core/src/tasks/idle_extension.rs +++ b/codex-rs/core/src/tasks/idle_extension.rs @@ -3,12 +3,8 @@ use std::sync::Arc; use codex_extension_api::ThreadIdleRequest; -use codex_protocol::models::ContentItem; -use codex_protocol::models::ResponseInputItem; use codex_protocol::protocol::ThreadSettingsSnapshot; -use crate::context::ContextualUserFragment; -use crate::context::ExtensionContext; use crate::session::TurnInput; use crate::session::session::Session; use crate::state::ActiveTurn; @@ -89,7 +85,7 @@ pub(crate) async fn maybe_start_turn(session: Arc) { return; } - let input = vec![idle_extension_input(candidate.request.prompt)]; + let input = vec![idle_extension_input(candidate.request)]; // Treat extension-provided prompts as the new turn's initial input rather than stashing them // in turn-state pending input; this keeps rollback logic out of the scheduler. @@ -108,7 +104,7 @@ 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_with_settings(codex_extension_api::ThreadIdleWithSettingsInput { + .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, @@ -144,13 +140,11 @@ async fn next_idle_turn_candidate(session: &Session) -> Option TurnInput { +fn idle_extension_input(request: ThreadIdleRequest) -> TurnInput { let prompt = codex_utils_string::truncate_middle_with_token_budget( - &prompt, + &request.prompt, MAX_IDLE_EXTENSION_PROMPT_TOKENS, ) .0; - if ExtensionContext::matches_text(&prompt) { - return TurnInput::ResponseInputItem(ResponseInputItem::Message { - role: ExtensionContext::role().to_string(), - content: vec![ContentItem::InputText { text: prompt }], - phase: None, - }); - } - - TurnInput::ResponseInputItem(ExtensionContext::new(prompt).into_response_input_item()) + TurnInput::ResponseInputItem(request.context_marker.response_input_item(prompt)) } #[cfg(test)] mod tests { + use codex_extension_api::HiddenContextMarker; + use codex_protocol::models::ContentItem; + use codex_protocol::models::ResponseInputItem; + use super::*; #[test] @@ -238,9 +228,12 @@ mod tests { ); let original_len = prompt.len(); - let prompt = ExtensionContext::new(prompt).render(); + let request = ThreadIdleRequest::new( + HiddenContextMarker::new("", ""), + prompt, + ); let TurnInput::ResponseInputItem(ResponseInputItem::Message { content, .. }) = - idle_extension_input(prompt) + idle_extension_input(request) else { panic!("expected message input"); }; @@ -248,7 +241,7 @@ mod tests { panic!("expected one text content item"); }; - assert!(text.starts_with("")); + assert!(text.starts_with("")); assert!(text.contains("start")); assert!(text.contains("end")); assert!(text.len() < original_len); diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 528070db02..1a44feef01 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -311,6 +311,10 @@ impl Session { self.start_task(turn_context, input, task).await; } + #[expect( + clippy::await_holding_invalid_type, + reason = "the active-turn guard must cover spawning until RunningTask is installed" + )] pub(crate) async fn start_task( self: &Arc, turn_context: Arc, @@ -348,32 +352,10 @@ impl Session { .await .clear_turn(&turn_context.sub_id); - let queued_response_items = self - .input_queue - .take_queued_response_items_for_next_turn() - .await; - let mailbox_items = self.input_queue.get_pending_input(&self.active_turn).await; - turn_state.lock().await.token_usage_at_turn_start = token_usage_at_turn_start.clone(); - let mut pending_items = queued_response_items - .into_iter() - .map(TurnInput::ResponseInputItem) - .collect::>(); - pending_items.extend(mailbox_items); - self.input_queue - .extend_pending_input_for_turn_state(turn_state.as_ref(), pending_items) - .await; self.emit_turn_start_lifecycle(turn_context.as_ref(), &token_usage_at_turn_start) .await; let turn_extension_data = Arc::clone(&turn_context.extension_data); - let mut active = self.active_turn.lock().await; - let Some(turn) = active.as_mut() else { - return; - }; - if !Arc::ptr_eq(&turn.turn_state, &turn_state) || !turn.task_starting || turn.task.is_some() - { - return; - } let done_clone = Arc::clone(&done); let session_ctx = Arc::new(SessionTaskContext::new( Arc::clone(self), @@ -400,43 +382,53 @@ impl Session { codex.turn.token_usage.reasoning_output_tokens = field::Empty, codex.turn.token_usage.total_tokens = field::Empty, ); - let handle = tokio::spawn( - async move { - let ctx_for_finish = Arc::clone(&ctx); - let last_agent_message = task_for_run - .run( - Arc::clone(&session_ctx), - ctx, - task_input, - task_cancellation_token.child_token(), - ) - .await; - let sess = session_ctx.clone_session(); - if let Err(err) = sess.flush_rollout().await { - warn!("failed to flush rollout before completing turn: {err}"); - sess.send_event( - ctx_for_finish.as_ref(), - EventMsg::Warning(WarningEvent { - message: format!( - "Failed to save the conversation transcript; Codex will continue retrying. Error: {err}" - ), - }), - ) - .await; - } - if !task_cancellation_token.is_cancelled() { - // Emit completion uniformly from spawn site so all tasks share the same lifecycle. - sess.on_task_finished(Arc::clone(&ctx_for_finish), last_agent_message) - .await; - } - done_clone.notify_waiters(); + let task_future = async move { + let ctx_for_finish = Arc::clone(&ctx); + let last_agent_message = task_for_run + .run( + Arc::clone(&session_ctx), + ctx, + task_input, + task_cancellation_token.child_token(), + ) + .await; + let sess = session_ctx.clone_session(); + if let Err(err) = sess.flush_rollout().await { + warn!("failed to flush rollout before completing turn: {err}"); + sess.send_event( + ctx_for_finish.as_ref(), + EventMsg::Warning(WarningEvent { + message: format!( + "Failed to save the conversation transcript; Codex will continue retrying. Error: {err}" + ), + }), + ) + .await; } - .instrument(task_span), - ); + if !task_cancellation_token.is_cancelled() { + // Emit completion uniformly from spawn site so all tasks share the same lifecycle. + sess.on_task_finished(Arc::clone(&ctx_for_finish), last_agent_message) + .await; + } + done_clone.notify_waiters(); + } + .instrument(task_span); let timer = turn_context .session_telemetry .start_timer(TURN_E2E_DURATION_METRIC, &[]) .ok(); + let mut active = self.active_turn.lock().await; + let Some(turn) = active.as_mut() else { + return; + }; + if !Arc::ptr_eq(&turn.turn_state, &turn_state) || !turn.task_starting || turn.task.is_some() + { + return; + } + self.input_queue + .prepare_starting_turn_input(turn, token_usage_at_turn_start) + .await; + let handle = tokio::spawn(task_future); let running_task = RunningTask { done, handle: AbortOnDropHandle::new(handle), diff --git a/codex-rs/ext/extension-api/src/capabilities/response_items.rs b/codex-rs/ext/extension-api/src/capabilities/response_items.rs index 6c300e2bf1..80cbc906a6 100644 --- a/codex-rs/ext/extension-api/src/capabilities/response_items.rs +++ b/codex-rs/ext/extension-api/src/capabilities/response_items.rs @@ -1,6 +1,7 @@ use std::future::Future; use std::pin::Pin; +use codex_protocol::ThreadId; use codex_protocol::models::ResponseInputItem; /// Future returned when an extension asks the host to inject model-visible input. @@ -15,6 +16,7 @@ pub type ResponseItemInjectionFuture<'a> = pub trait ResponseItemInjector: Send + Sync { fn inject_response_items<'a>( &'a self, + thread_id: ThreadId, items: Vec, ) -> ResponseItemInjectionFuture<'a>; } @@ -26,6 +28,7 @@ pub struct NoopResponseItemInjector; impl ResponseItemInjector for NoopResponseItemInjector { fn inject_response_items<'a>( &'a self, + _thread_id: ThreadId, items: Vec, ) -> ResponseItemInjectionFuture<'a> { Box::pin(std::future::ready(Err(items))) diff --git a/codex-rs/ext/extension-api/src/context.rs b/codex-rs/ext/extension-api/src/context.rs new file mode 100644 index 0000000000..6996ecf4e7 --- /dev/null +++ b/codex-rs/ext/extension-api/src/context.rs @@ -0,0 +1,37 @@ +//! Helpers for extension-owned hidden context messages. + +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseInputItem; + +/// Marker pair for extension-owned hidden model-visible context. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct HiddenContextMarker { + start_marker: &'static str, + end_marker: &'static str, +} + +impl HiddenContextMarker { + /// Creates a marker pair for one extension-owned hidden context shape. + pub const fn new(start_marker: &'static str, end_marker: &'static str) -> Self { + Self { + start_marker, + end_marker, + } + } + + /// Renders this marker pair as a user-role response input item. + pub fn response_input_item(&self, body: impl AsRef) -> ResponseInputItem { + ResponseInputItem::Message { + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: format!( + "{}\n{}\n{}", + self.start_marker, + body.as_ref(), + self.end_marker + ), + }], + phase: None, + } + } +} diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 081221d530..2eea9d6d4d 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -22,9 +22,7 @@ pub use prompt::PromptSlot; pub use thread_lifecycle::ThreadIdleInput; pub use thread_lifecycle::ThreadIdleRequest; pub use thread_lifecycle::ThreadIdleTurnStartInput; -pub use thread_lifecycle::ThreadIdleWithSettingsInput; pub use thread_lifecycle::ThreadResumeInput; -pub use thread_lifecycle::ThreadResumeWithSettingsInput; pub use thread_lifecycle::ThreadStartInput; pub use thread_lifecycle::ThreadStopInput; pub use tool_lifecycle::ToolCallOutcome; @@ -60,21 +58,13 @@ pub trait ThreadLifecycleContributor: Send + Sync { /// Called after the host constructs a runtime from persisted history. async fn on_thread_resume(&self, _input: ThreadResumeInput<'_>) {} - /// Called after the host constructs a runtime from persisted history, with - /// a snapshot of host-owned thread settings. - async fn on_thread_resume_with_settings(&self, input: ThreadResumeWithSettingsInput<'_>) { - self.on_thread_resume(input.without_settings()).await; - } - /// Called after the host has drained immediately pending thread work. + /// + /// Implementations may use host capabilities captured by the extension to + /// submit follow-up input. The host remains responsible for deciding + /// whether that input starts a turn, is queued, or is ignored. async fn on_thread_idle(&self, _input: ThreadIdleInput<'_>) {} - /// Called after the host has drained immediately pending thread work, with - /// a snapshot of host-owned thread settings. - async fn on_thread_idle_with_settings(&self, input: ThreadIdleWithSettingsInput<'_>) { - self.on_thread_idle(input.without_settings()).await; - } - /// Called before the host drops the thread runtime and thread-scoped store. async fn on_thread_stop(&self, _input: ThreadStopInput<'_>) {} } @@ -101,15 +91,6 @@ pub trait ThreadIdleTurnContributor: Send + Sync { Box::pin(std::future::ready(None)) } - /// Returns one hidden prompt body to start an idle turn, with a snapshot of - /// host-owned thread settings. - fn request_thread_idle_turn_with_settings<'a>( - &'a self, - input: ThreadIdleWithSettingsInput<'a>, - ) -> ThreadIdleTurnRequestFuture<'a> { - self.request_thread_idle_turn(input.without_settings()) - } - /// Confirms that the idle-turn request is still current immediately before /// the host starts the turn. fn should_start_thread_idle_turn<'a>( diff --git a/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs b/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs index 6d1322f4af..d0c6b9afaa 100644 --- a/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs +++ b/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs @@ -1,4 +1,5 @@ use crate::ExtensionData; +use crate::HiddenContextMarker; use codex_protocol::protocol::ThreadSettingsSnapshot; /// Input supplied when the host starts a runtime for a thread. @@ -13,14 +14,6 @@ pub struct ThreadStartInput<'a, C> { /// Input supplied when the host resumes an existing thread. pub struct ThreadResumeInput<'a> { - /// Store scoped to the host session runtime. - pub session_store: &'a ExtensionData, - /// Store scoped to this thread runtime. - pub thread_store: &'a ExtensionData, -} - -/// Settings-aware input supplied when the host resumes an existing thread. -pub struct ThreadResumeWithSettingsInput<'a> { /// Current host-owned thread settings at resume time. pub thread_settings: &'a ThreadSettingsSnapshot, /// Store scoped to the host session runtime. @@ -29,25 +22,8 @@ pub struct ThreadResumeWithSettingsInput<'a> { pub thread_store: &'a ExtensionData, } -impl<'a> ThreadResumeWithSettingsInput<'a> { - pub fn without_settings(&self) -> ThreadResumeInput<'a> { - ThreadResumeInput { - session_store: self.session_store, - thread_store: self.thread_store, - } - } -} - /// Input supplied when the host has no immediately pending thread work. pub struct ThreadIdleInput<'a> { - /// Store scoped to the host session runtime. - pub session_store: &'a ExtensionData, - /// Store scoped to this thread runtime. - pub thread_store: &'a ExtensionData, -} - -/// Settings-aware input supplied when the host has no immediately pending thread work. -pub struct ThreadIdleWithSettingsInput<'a> { /// Current host-owned thread settings for the idle thread. pub thread_settings: &'a ThreadSettingsSnapshot, /// Store scoped to the host session runtime. @@ -56,27 +32,21 @@ pub struct ThreadIdleWithSettingsInput<'a> { pub thread_store: &'a ExtensionData, } -impl<'a> ThreadIdleWithSettingsInput<'a> { - pub fn without_settings(&self) -> ThreadIdleInput<'a> { - ThreadIdleInput { - session_store: self.session_store, - thread_store: self.thread_store, - } - } -} - /// Extension request to start a new idle turn. pub struct ThreadIdleRequest { /// Hidden prompt body that the host wraps as extension-owned context. pub prompt: String, + /// Marker pair used to wrap the prompt as hidden extension-owned context. + pub context_marker: HiddenContextMarker, /// Opaque extension-owned key used to reject stale requests before start. pub validation_key: Option, } impl ThreadIdleRequest { - pub fn new(prompt: impl Into) -> Self { + pub fn new(context_marker: HiddenContextMarker, prompt: impl Into) -> Self { Self { prompt: prompt.into(), + context_marker, validation_key: None, } } diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index f4b398f11f..1abe4d8b04 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -1,4 +1,5 @@ mod capabilities; +mod context; mod contributors; mod registry; mod state; @@ -23,6 +24,7 @@ pub use codex_tools::ToolPayload; pub use codex_tools::ToolSpec; pub use codex_tools::parse_tool_input_schema; pub use codex_tools::parse_tool_input_schema_without_compaction; +pub use context::HiddenContextMarker; pub use contributors::ApprovalReviewContributor; pub use contributors::ConfigContributor; pub use contributors::ContextContributor; @@ -34,10 +36,8 @@ pub use contributors::ThreadIdleTurnContributor; pub use contributors::ThreadIdleTurnRequestFuture; pub use contributors::ThreadIdleTurnStartFuture; pub use contributors::ThreadIdleTurnStartInput; -pub use contributors::ThreadIdleWithSettingsInput; pub use contributors::ThreadLifecycleContributor; pub use contributors::ThreadResumeInput; -pub use contributors::ThreadResumeWithSettingsInput; pub use contributors::ThreadStartInput; pub use contributors::ThreadStopInput; pub use contributors::TokenUsageContributor; diff --git a/codex-rs/ext/goal/Cargo.toml b/codex-rs/ext/goal/Cargo.toml index 6116383243..0138d28ade 100644 --- a/codex-rs/ext/goal/Cargo.toml +++ b/codex-rs/ext/goal/Cargo.toml @@ -15,7 +15,6 @@ workspace = true [dependencies] async-trait = { workspace = true } -codex-core = { workspace = true } codex-extension-api = { workspace = true } codex-otel = { workspace = true } codex-protocol = { workspace = true } diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index f7c932f356..fa7a2f8465 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -1,19 +1,18 @@ use std::sync::Arc; -use std::sync::Weak; use async_trait::async_trait; -use codex_core::ThreadManager; use codex_extension_api::ConfigContributor; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionEventSink; use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::ResponseItemInjector; +use codex_extension_api::ThreadIdleInput; use codex_extension_api::ThreadIdleTurnContributor; use codex_extension_api::ThreadIdleTurnRequestFuture; use codex_extension_api::ThreadIdleTurnStartFuture; use codex_extension_api::ThreadIdleTurnStartInput; -use codex_extension_api::ThreadIdleWithSettingsInput; use codex_extension_api::ThreadLifecycleContributor; -use codex_extension_api::ThreadResumeWithSettingsInput; +use codex_extension_api::ThreadResumeInput; use codex_extension_api::ThreadStartInput; use codex_extension_api::TokenUsageContributor; use codex_extension_api::ToolCallOutcome; @@ -49,7 +48,7 @@ pub struct GoalExtension { state_dbs: Arc, event_emitter: GoalEventEmitter, metrics: GoalMetrics, - thread_manager: Weak, + response_item_injector: Arc, goals_enabled: Arc bool + Send + Sync>, } @@ -64,14 +63,14 @@ impl GoalExtension { state_dbs: Arc, event_sink: Arc, metrics_client: Option, - thread_manager: Weak, + response_item_injector: Arc, goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static, ) -> Self { Self { state_dbs, event_emitter: GoalEventEmitter::new(event_sink), metrics: GoalMetrics::new(metrics_client), - thread_manager, + response_item_injector, goals_enabled: Arc::new(goals_enabled), } } @@ -96,7 +95,7 @@ where Arc::clone(&self.state_dbs), self.event_emitter.clone(), self.metrics.clone(), - self.thread_manager.clone(), + Arc::clone(&self.response_item_injector), accounting_state, enabled, ) @@ -104,7 +103,7 @@ where runtime.set_enabled(enabled); } - async fn on_thread_resume_with_settings(&self, input: ThreadResumeWithSettingsInput<'_>) { + async fn on_thread_resume(&self, input: ThreadResumeInput<'_>) { let Some(runtime) = goal_runtime_handle(input.thread_store) else { return; }; @@ -122,9 +121,9 @@ impl ThreadIdleTurnContributor for GoalExtension where C: Send + Sync + 'static, { - fn request_thread_idle_turn_with_settings<'a>( + fn request_thread_idle_turn<'a>( &'a self, - input: ThreadIdleWithSettingsInput<'a>, + input: ThreadIdleInput<'a>, ) -> ThreadIdleTurnRequestFuture<'a> { Box::pin(async move { let runtime = goal_runtime_handle(input.thread_store)?; @@ -456,7 +455,7 @@ pub fn install_with_backend( registry: &mut ExtensionRegistryBuilder, state_dbs: Arc, metrics_client: Option, - thread_manager: Weak, + response_item_injector: Arc, goals_enabled: impl Fn(&C) -> bool + Send + Sync + 'static, ) where C: Send + Sync + 'static, @@ -465,7 +464,7 @@ pub fn install_with_backend( state_dbs, registry.event_sink(), metrics_client, - thread_manager, + response_item_injector, goals_enabled, )); registry.thread_lifecycle_contributor(extension.clone()); diff --git a/codex-rs/ext/goal/src/runtime.rs b/codex-rs/ext/goal/src/runtime.rs index 51716aa17e..79def44635 100644 --- a/codex-rs/ext/goal/src/runtime.rs +++ b/codex-rs/ext/goal/src/runtime.rs @@ -1,9 +1,8 @@ use std::sync::Arc; -use std::sync::Weak; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; -use codex_core::ThreadManager; +use codex_extension_api::ResponseItemInjector; use codex_protocol::ThreadId; use codex_protocol::config_types::ModeKind; use codex_protocol::models::ResponseInputItem; @@ -28,7 +27,7 @@ struct GoalRuntimeInner { state_dbs: Arc, event_emitter: GoalEventEmitter, metrics: GoalMetrics, - thread_manager: Weak, + response_item_injector: Arc, accounting_state: Arc, enabled: AtomicBool, } @@ -50,7 +49,7 @@ impl GoalRuntimeHandle { state_dbs: Arc, event_emitter: GoalEventEmitter, metrics: GoalMetrics, - thread_manager: Weak, + response_item_injector: Arc, accounting_state: Arc, enabled: bool, ) -> Self { @@ -60,7 +59,7 @@ impl GoalRuntimeHandle { state_dbs, event_emitter, metrics, - thread_manager, + response_item_injector, accounting_state, enabled: AtomicBool::new(enabled), }), @@ -325,16 +324,10 @@ impl GoalRuntimeHandle { } pub(crate) async fn inject_active_turn_steering(&self, item: ResponseInputItem) { - let Some(thread_manager) = self.inner.thread_manager.upgrade() else { - tracing::debug!("skipping goal steering because thread manager is unavailable"); - return; - }; - let Ok(thread) = thread_manager.get_thread(self.inner.thread_id).await else { - tracing::debug!("skipping goal steering because live thread is unavailable"); - return; - }; - if thread - .inject_response_items_into_active_turn(vec![item]) + if self + .inner + .response_item_injector + .inject_response_items(self.inner.thread_id, vec![item]) .await .is_err() { diff --git a/codex-rs/ext/goal/src/steering.rs b/codex-rs/ext/goal/src/steering.rs index 92a05afb80..1907f72fee 100644 --- a/codex-rs/ext/goal/src/steering.rs +++ b/codex-rs/ext/goal/src/steering.rs @@ -1,11 +1,13 @@ -use codex_core::context::ContextualUserFragment; -use codex_core::context::ExtensionContext; +use codex_extension_api::HiddenContextMarker; use codex_extension_api::ThreadIdleRequest; use codex_protocol::models::ResponseInputItem; use codex_protocol::protocol::ThreadGoal; use codex_utils_template::Template; use std::sync::LazyLock; +const GOAL_CONTEXT_MARKER: HiddenContextMarker = + HiddenContextMarker::new("", ""); + static CONTINUATION_TEMPLATE: LazyLock