agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: move extension scope ids into ExtensionData (#22490)

Browse Source 
## Summary
- add a scoped level_id to ExtensionData and expose it through
level_id()
- remove thread_id/turn_id parameters from extension contributor inputs
where the scoped ExtensionData already carries that identity
- move turn-scoped extension data onto TurnContext so token usage and
lifecycle contributors can share the same turn store

## Testing
- cargo check -p codex-extension-api -p codex-core --tests
- cargo test -p codex-extension-api
- cargo test -p codex-guardian
- cargo test -p codex-core --lib
record_token_usage_info_notifies_extension_contributors
- cargo test -p codex-core --lib
submission_loop_channel_close_emits_thread_stop_lifecycle
- cargo test -p codex-core --lib
submission_loop_channel_close_aborts_active_turn_before_thread_stop_lifecycle
- just fix -p codex-extension-api
- just fix -p codex-guardian
- just fix -p codex-core
- just fmt

## Note
- Attempted cargo test -p codex-core; it aborted in
agent::control::tests::spawn_agent_fork_last_n_turns_keeps_only_recent_turns
with the existing stack overflow before the full suite completed.
This commit is contained in:
jif-oai
2026-05-13 16:13:16 +02:00
committed by GitHub
parent 99157f3797
commit 1dcc89f1d4
16 changed files with 77 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/tasks/lifecycle.rs
+2 -18
View File
@@ -2,18 +2,11 @@ use codex_extension_api::ExtensionData;
use codex_protocol::protocol::TurnAbortReason;
use crate::session::session::Session;
use crate::session::turn_context::TurnContext;
impl Session {
pub(super) fn emit_turn_start_lifecycle(
&self,
turn_context: &TurnContext,
turn_store: &ExtensionData,
) {
pub(super) fn emit_turn_start_lifecycle(&self, turn_store: &ExtensionData) {
for contributor in self.services.extensions.turn_lifecycle_contributors() {
contributor.on_turn_start(codex_extension_api::TurnStartInput {
thread_id: self.conversation_id,
turn_id: &turn_context.sub_id,
session_store: &self.services.session_extension_data,
thread_store: &self.services.thread_extension_data,
turn_store,
@@ -21,15 +14,9 @@ impl Session {
}
}
pub(super) fn emit_turn_stop_lifecycle(
&self,
turn_context: &TurnContext,
turn_store: &ExtensionData,
) {
pub(super) fn emit_turn_stop_lifecycle(&self, turn_store: &ExtensionData) {
for contributor in self.services.extensions.turn_lifecycle_contributors() {
contributor.on_turn_stop(codex_extension_api::TurnStopInput {
thread_id: self.conversation_id,
turn_id: &turn_context.sub_id,
session_store: &self.services.session_extension_data,
thread_store: &self.services.thread_extension_data,
turn_store,
@@ -39,14 +26,11 @@ impl Session {
pub(super) fn emit_turn_abort_lifecycle(
&self,
turn_context: &TurnContext,
reason: TurnAbortReason,
turn_store: &ExtensionData,
) {
for contributor in self.services.extensions.turn_lifecycle_contributors() {
contributor.on_turn_abort(codex_extension_api::TurnAbortInput {
thread_id: self.conversation_id,
turn_id: &turn_context.sub_id,
reason: reason.clone(),
session_store: &self.services.session_extension_data,
thread_store: &self.services.thread_extension_data,
codex-rs/core/src/tasks/mod.rs
+9 -26
View File
@@ -346,7 +346,7 @@ impl Session {
debug_assert!(turn.tasks.is_empty());
Arc::clone(&turn.turn_state)
};
let turn_extension_data = {
{
let mut turn_state = turn_state.lock().await;
turn_state.token_usage_at_turn_start = token_usage_at_turn_start;
for item in queued_response_items {
@@ -355,9 +355,8 @@ impl Session {
for item in mailbox_items {
turn_state.push_pending_input(item);
}
Arc::clone(&turn_state.extension_data)
};
self.emit_turn_start_lifecycle(turn_context.as_ref(), turn_extension_data.as_ref());
}
self.emit_turn_start_lifecycle(turn_context.extension_data.as_ref());
let mut active = self.active_turn.lock().await;
let turn = active.get_or_insert_with(ActiveTurn::default);
@@ -428,7 +427,6 @@ impl Session {
task,
cancellation_token,
turn_context: Arc::clone(&turn_context),
turn_extension_data,
_timer: timer,
};
turn.add_task(running_task);
@@ -480,14 +478,10 @@ impl Session {
let mut aborted_turn = false;
let mut active_turn_to_clear = None;
let mut turn_context = None;
let mut turn_extension_data = None;
if let Some(mut active_turn) = self.take_active_turn().await {
let tasks = active_turn.drain_tasks();
aborted_turn = !tasks.is_empty();
turn_context = tasks.first().map(|task| Arc::clone(&task.turn_context));
turn_extension_data = tasks
.first()
.map(|task| Arc::clone(&task.turn_extension_data));
for task in tasks {
self.handle_task_abort(task, reason.clone()).await;
}
@@ -496,10 +490,8 @@ impl Session {
}
}
if let Some(turn_context) = turn_context.as_deref()
&& let Some(turn_extension_data) = turn_extension_data.as_deref()
{
self.emit_turn_abort_lifecycle(turn_context, reason.clone(), turn_extension_data);
if let Some(turn_context) = turn_context.as_deref() {
self.emit_turn_abort_lifecycle(reason.clone(), turn_context.extension_data.as_ref());
}
if (aborted_turn || reason == TurnAbortReason::Interrupted)
&& let Err(err) = self
@@ -543,16 +535,11 @@ impl Session {
let tasks = active_turn.drain_tasks();
let turn_context = tasks.first().map(|task| Arc::clone(&task.turn_context));
let turn_extension_data = tasks
.first()
.map(|task| Arc::clone(&task.turn_extension_data));
for task in tasks {
self.handle_task_abort(task, reason.clone()).await;
}
if let Some(turn_context) = turn_context.as_deref()
&& let Some(turn_extension_data) = turn_extension_data.as_deref()
{
self.emit_turn_abort_lifecycle(turn_context, reason.clone(), turn_extension_data);
if let Some(turn_context) = turn_context.as_deref() {
self.emit_turn_abort_lifecycle(reason.clone(), turn_context.extension_data.as_ref());
}
if let Err(err) = self
.goal_runtime_apply(GoalRuntimeEvent::TaskAborted {
@@ -589,7 +576,6 @@ impl Session {
let mut turn_had_memory_citation = false;
let mut turn_tool_calls = 0_u64;
let mut records_turn_token_usage_on_span = false;
let mut turn_extension_data = None;
let turn_state = {
let mut active = self.active_turn.lock().await;
if let Some(at) = active.as_mut()
@@ -598,7 +584,6 @@ impl Session {
records_turn_token_usage_on_span = removed_task.records_turn_token_usage_on_span;
if removed_task.active_turn_is_empty {
should_clear_active_turn = true;
turn_extension_data = Some(removed_task.turn_extension_data);
let turn_state = Arc::clone(&at.turn_state);
Some(turn_state)
} else {
@@ -756,10 +741,8 @@ impl Session {
.turn_timing_state
.time_to_first_token_ms()
.await;
if should_clear_active_turn
&& let Some(turn_extension_data) = turn_extension_data.as_deref()
{
self.emit_turn_stop_lifecycle(turn_context.as_ref(), turn_extension_data);
if should_clear_active_turn {
self.emit_turn_stop_lifecycle(turn_context.extension_data.as_ref());
}
if let Err(err) = self
.goal_runtime_apply(GoalRuntimeEvent::TurnFinished {