From d1a763f6cd5f0820a0d0afaddef6b581982123e5 Mon Sep 17 00:00:00 2001 From: Abhinav Vedmala Date: Tue, 19 May 2026 12:28:01 -0700 Subject: [PATCH] Preserve sticky SessionStart hook context --- .../src/context/hook_additional_context.rs | 21 +++++++++ codex-rs/core/src/context/mod.rs | 1 + codex-rs/core/src/event_mapping.rs | 1 + codex-rs/core/src/event_mapping_tests.rs | 17 +++++++ codex-rs/core/src/hook_runtime.rs | 47 ++++++++++++------- codex-rs/core/tests/suite/hooks.rs | 15 ++++-- 6 files changed, 82 insertions(+), 20 deletions(-) diff --git a/codex-rs/core/src/context/hook_additional_context.rs b/codex-rs/core/src/context/hook_additional_context.rs index 5ef6218aa0..6450f8628b 100644 --- a/codex-rs/core/src/context/hook_additional_context.rs +++ b/codex-rs/core/src/context/hook_additional_context.rs @@ -5,12 +5,23 @@ pub(crate) struct HookAdditionalContext { text: String, } +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct StickyHookAdditionalContext { + text: String, +} + impl HookAdditionalContext { pub(crate) fn new(text: impl Into) -> Self { Self { text: text.into() } } } +impl StickyHookAdditionalContext { + pub(crate) fn new(text: impl Into) -> Self { + Self { text: text.into() } + } +} + impl ContextualUserFragment for HookAdditionalContext { const ROLE: &'static str = "developer"; const START_MARKER: &'static str = ""; @@ -20,3 +31,13 @@ impl ContextualUserFragment for HookAdditionalContext { self.text.clone() } } + +impl ContextualUserFragment for StickyHookAdditionalContext { + const ROLE: &'static str = "developer"; + const START_MARKER: &'static str = ""; + const END_MARKER: &'static str = ""; + + fn body(&self) -> String { + self.text.clone() + } +} diff --git a/codex-rs/core/src/context/mod.rs b/codex-rs/core/src/context/mod.rs index bd89025368..976e37bf27 100644 --- a/codex-rs/core/src/context/mod.rs +++ b/codex-rs/core/src/context/mod.rs @@ -43,6 +43,7 @@ pub(crate) use fragment::FragmentRegistrationProxy; pub(crate) use goal_context::GoalContext; pub(crate) use guardian_followup_review_reminder::GuardianFollowupReviewReminder; pub(crate) use hook_additional_context::HookAdditionalContext; +pub(crate) use hook_additional_context::StickyHookAdditionalContext; pub(crate) use image_generation_instructions::ImageGenerationInstructions; pub(crate) use legacy_apply_patch_exec_command_warning::LegacyApplyPatchExecCommandWarning; pub(crate) use legacy_model_mismatch_warning::LegacyModelMismatchWarning; diff --git a/codex-rs/core/src/event_mapping.rs b/codex-rs/core/src/event_mapping.rs index e7c79e6dd2..fde94c54ce 100644 --- a/codex-rs/core/src/event_mapping.rs +++ b/codex-rs/core/src/event_mapping.rs @@ -27,6 +27,7 @@ use crate::web_search::web_search_action_detail; const CONTEXTUAL_DEVELOPER_PREFIXES: &[&str] = &[ "", "", + "", COLLABORATION_MODE_OPEN_TAG, REALTIME_CONVERSATION_OPEN_TAG, "", diff --git a/codex-rs/core/src/event_mapping_tests.rs b/codex-rs/core/src/event_mapping_tests.rs index a70b2a69b0..55f0c3705d 100644 --- a/codex-rs/core/src/event_mapping_tests.rs +++ b/codex-rs/core/src/event_mapping_tests.rs @@ -1,3 +1,5 @@ +use super::has_non_contextual_dev_message_content; +use super::is_contextual_dev_message_content; use super::parse_turn_item; use crate::context::ContextualUserFragment; use crate::context::GoalContext; @@ -304,6 +306,21 @@ fn parses_hook_prompt_and_hides_other_contextual_fragments() { } } +#[test] +fn hook_context_fragments_distinguish_thread_scoped_from_sticky_state() { + let trimmable_context = vec![ContentItem::InputText { + text: "thread scoped note".to_string(), + }]; + assert!(is_contextual_dev_message_content(&trimmable_context)); + assert!(!has_non_contextual_dev_message_content(&trimmable_context)); + + let sticky_context = vec![ContentItem::InputText { + text: "session scoped note".to_string(), + }]; + assert!(!is_contextual_dev_message_content(&sticky_context)); + assert!(has_non_contextual_dev_message_content(&sticky_context)); +} + #[test] fn goal_context_does_not_parse_as_visible_turn_item() { let item = ResponseItem::Message { diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index 1b698c2d97..26779506cd 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -33,6 +33,7 @@ use serde_json::Value; use crate::context::ContextualUserFragment; use crate::context::HookAdditionalContext; +use crate::context::StickyHookAdditionalContext; use crate::context_manager::updates::build_developer_update_item; use crate::event_mapping::parse_turn_item; use crate::session::session::Session; @@ -126,15 +127,16 @@ pub(crate) async fn run_pending_session_start_hooks( }; let hooks = sess.hooks(); let preview_runs = hooks.preview_session_start(&request); - run_context_injecting_hook( + let outcome = run_context_injecting_hook( sess, turn_context, preview_runs, hooks.run_session_start(request, Some(turn_context.sub_id.clone())), ) - .await - .record_additional_contexts(sess, turn_context) - .await + .await; + // SessionStart context is durable developer state; rollback should not trim it. + record_sticky_additional_contexts(sess, turn_context, outcome.additional_contexts).await; + outcome.should_stop } /// Runs matching `PreToolUse` hooks before a tool executes. @@ -433,18 +435,6 @@ where outcome.outcome } -impl HookRuntimeOutcome { - async fn record_additional_contexts( - self, - sess: &Arc, - turn_context: &Arc, - ) -> bool { - record_additional_contexts(sess, turn_context, self.additional_contexts).await; - - self.should_stop - } -} - pub(crate) async fn record_additional_contexts( sess: &Arc, turn_context: &Arc, @@ -459,6 +449,21 @@ pub(crate) async fn record_additional_contexts( .await; } +async fn record_sticky_additional_contexts( + sess: &Arc, + turn_context: &Arc, + additional_contexts: Vec, +) { + // Use this for hook context that should survive rollback as persistent developer state. + let developer_messages = sticky_additional_context_messages(additional_contexts); + if developer_messages.is_empty() { + return; + } + + sess.record_conversation_items(turn_context, developer_messages.as_slice()) + .await; +} + fn additional_context_messages(additional_contexts: Vec) -> Vec { let sections = additional_contexts .into_iter() @@ -469,6 +474,16 @@ fn additional_context_messages(additional_contexts: Vec) -> Vec) -> Vec { + let sections = additional_contexts + .into_iter() + .map(StickyHookAdditionalContext::new) + .map(|context| context.render()) + .collect(); + + build_developer_update_item(sections).into_iter().collect() +} + async fn emit_hook_started_events( sess: &Arc, turn_context: &Arc, diff --git a/codex-rs/core/tests/suite/hooks.rs b/codex-rs/core/tests/suite/hooks.rs index 5c64363aa3..9c61b00597 100644 --- a/codex-rs/core/tests/suite/hooks.rs +++ b/codex-rs/core/tests/suite/hooks.rs @@ -785,7 +785,11 @@ fn request_hook_prompt_texts( fn spilled_hook_output_path(text: &str) -> Option<&str> { text.lines() .find_map(|line| line.strip_prefix("Full hook output saved to: ")) - .map(|path| path.trim_end_matches("")) + .map(|path| { + path.strip_suffix("") + .or_else(|| path.strip_suffix("")) + .unwrap_or(path) + }) } fn read_stop_hook_inputs(home: &Path) -> Result> { @@ -1103,6 +1107,7 @@ async fn session_start_hook_spills_large_additional_context() -> Result<()> { .find(|message| spilled_hook_output_path(message).is_some()) .context("spilled developer hook message")?; assert!(developer_message.contains("tokens truncated")); + assert!(developer_message.contains("")); let path = spilled_hook_output_path(developer_message).context("spill path")?; assert_eq!(fs::read_to_string(path)?, additional_context); @@ -1149,9 +1154,11 @@ async fn parallel_session_start_additional_contexts_share_one_developer_message( .filter_map(|content| content.get("text").and_then(Value::as_str)) .collect::>() .join("\n"); - contexts - .iter() - .all(|context| joined.contains(&format!("{context}"))) + contexts.iter().all(|context| { + joined.contains(&format!( + "{context}" + )) + }) }) .count(); assert_eq!(merged_developer_messages, 1);