Add step-scoped data to extension contributors (#34525)

## What changed

- Add an `ExtensionData` store to each `StepContext`.
- Pass the step store to context, world-state, turn-input, and tool contributors so extensions can use capabilities bound to the current sampling step.
- Preserve the captured step context when rebuilding initial context during compaction.

GitOrigin-RevId: 51f3105d3a793f60534f122f9be8247c131c2ca5

codex-rs/ext/memories/src/extension.rs

@@ -52,6 +52,7 @@ impl ContextContributor for MemoriesExtension {
         &'a self,
         _session_store: &'a ExtensionData,
         thread_store: &'a ExtensionData,
+        _step_store: &'a ExtensionData,
     ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Vec<PromptFragment>> + Send + 'a>> {
         Box::pin(async move {
             let Some(config) = thread_store.get::<MemoriesExtensionConfig>() else {
@@ -100,6 +101,7 @@ impl ToolContributor for MemoriesExtension {
         &self,
         _session_store: &ExtensionData,
         thread_store: &ExtensionData,
+        _step_store: &ExtensionData,
     ) -> Vec<Arc<dyn codex_extension_api::ToolExecutor<codex_extension_api::ToolCall>>> {
         let Some(config) = thread_store.get::<MemoriesExtensionConfig>() else {
             return Vec::new();

codex-rs/ext/memories/src/tests.rs

@@ -41,7 +41,8 @@ fn tools_are_not_contributed_without_thread_config() {
         extension
             .tools(
                 &ExtensionData::new("session"),
-                &ExtensionData::new("thread")
+                &ExtensionData::new("thread"),
+                &ExtensionData::new("step")
             )
             .is_empty()
     );
@@ -59,7 +60,11 @@ fn tools_are_not_contributed_when_disabled() {
 
     assert!(
         extension
-            .tools(&ExtensionData::new("session"), &thread_store)
+            .tools(
+                &ExtensionData::new("session"),
+                &thread_store,
+                &ExtensionData::new("step"),
+            )
             .is_empty()
     );
 }
@@ -76,7 +81,11 @@ fn tools_are_not_contributed_when_dedicated_tools_disabled() {
 
     assert!(
         extension
-            .tools(&ExtensionData::new("session"), &thread_store)
+            .tools(
+                &ExtensionData::new("session"),
+                &thread_store,
+                &ExtensionData::new("step"),
+            )
             .is_empty()
     );
 }
@@ -92,7 +101,11 @@ fn tools_are_contributed_when_enabled_with_dedicated_tools() {
     });
 
     let tool_names = extension
-        .tools(&ExtensionData::new("session"), &thread_store)
+        .tools(
+            &ExtensionData::new("session"),
+            &thread_store,
+            &ExtensionData::new("step"),
+        )
         .into_iter()
         .map(|tool| tool.tool_name())
         .collect::<Vec<_>>();
@@ -123,7 +136,13 @@ fn install_registers_dedicated_tool_contributor() {
     let tool_names = registry
         .tool_contributors()
         .iter()
-        .flat_map(|contributor| contributor.tools(&ExtensionData::new("session"), &thread_store))
+        .flat_map(|contributor| {
+            contributor.tools(
+                &ExtensionData::new("session"),
+                &thread_store,
+                &ExtensionData::new("step"),
+            )
+        })
         .map(|tool| tool.tool_name())
         .collect::<Vec<_>>();
 
@@ -181,7 +200,11 @@ async fn prompt_contribution_uses_memory_summary_when_enabled() {
     });
 
     let fragments = extension
-        .contribute_thread_context(&ExtensionData::new("session"), &thread_store)
+        .contribute_thread_context(
+            &ExtensionData::new("session"),
+            &thread_store,
+            &ExtensionData::new("step"),
+        )
         .await;
 
     assert_eq!(fragments.len(), 1);