valknar/llmx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Pass TurnContext around instead of sub_id (#5421)

Browse Source 
Today `sub_id` is an ID of a single incoming Codex Op submition. We then
associate all events triggered by this operation using the same
`sub_id`.

At the same time we are also creating a TurnContext per submission and
we'd like to start associating some events (item added/item completed)
with an entire turn instead of just the operation that started it.

Using turn context when sending events give us flexibility to change
notification scheme.
This commit is contained in:
pakrym-oai
2025-10-21 08:04:16 -07:00
committed by GitHub
parent 42d5c35020
commit 789e65b9d2
28 changed files with 293 additions and 401 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
codex-rs/core/src/tasks/mod.rs
View File

@@ -15,7 +15,6 @@ use tracing::warn;
use crate::codex::Session;
use crate::codex::TurnContext;
use crate::protocol::Event;
use crate::protocol::EventMsg;
use crate::protocol::TaskCompleteEvent;
use crate::protocol::TurnAbortReason;
@@ -55,13 +54,12 @@ pub(crate) trait SessionTask: Send + Sync + 'static {
self: Arc<Self>,
session: Arc<SessionTaskContext>,
ctx: Arc<TurnContext>,
sub_id: String,
input: Vec<UserInput>,
cancellation_token: CancellationToken,
) -> Option<String>;
async fn abort(&self, session: Arc<SessionTaskContext>, sub_id: &str) {
let _ = (session, sub_id);
async fn abort(&self, session: Arc<SessionTaskContext>, ctx: Arc<TurnContext>) {
let _ = (session, ctx);
}
}
@@ -69,7 +67,6 @@ impl Session {
pub async fn spawn_task<T: SessionTask>(
self: &Arc<Self>,
turn_context: Arc<TurnContext>,
sub_id: String,
input: Vec<UserInput>,
task: T,
) {
@@ -86,14 +83,13 @@ impl Session {
let session_ctx = Arc::new(SessionTaskContext::new(Arc::clone(self)));
let ctx = Arc::clone(&turn_context);
let task_for_run = Arc::clone(&task);
let sub_clone = sub_id.clone();
let task_cancellation_token = cancellation_token.child_token();
tokio::spawn(async move {
let ctx_for_finish = Arc::clone(&ctx);
let last_agent_message = task_for_run
.run(
Arc::clone(&session_ctx),
ctx,
sub_clone.clone(),
input,
task_cancellation_token.child_token(),
)
@@ -102,7 +98,8 @@ impl Session {
if !task_cancellation_token.is_cancelled() {
// Emit completion uniformly from spawn site so all tasks share the same lifecycle.
let sess = session_ctx.clone_session();
sess.on_task_finished(sub_clone, last_agent_message).await;
sess.on_task_finished(ctx_for_finish, last_agent_message)
.await;
}
done_clone.notify_waiters();
})
@@ -114,60 +111,54 @@ impl Session {
kind: task_kind,
task,
cancellation_token,
turn_context: Arc::clone(&turn_context),
};
self.register_new_active_task(sub_id, running_task).await;
self.register_new_active_task(running_task).await;
}
pub async fn abort_all_tasks(self: &Arc<Self>, reason: TurnAbortReason) {
for (sub_id, task) in self.take_all_running_tasks().await {
self.handle_task_abort(sub_id, task, reason.clone()).await;
for task in self.take_all_running_tasks().await {
self.handle_task_abort(task, reason.clone()).await;
}
}
pub async fn on_task_finished(
self: &Arc<Self>,
sub_id: String,
turn_context: Arc<TurnContext>,
last_agent_message: Option<String>,
) {
let mut active = self.active_turn.lock().await;
if let Some(at) = active.as_mut()
&& at.remove_task(&sub_id)
&& at.remove_task(&turn_context.sub_id)
{
*active = None;
}
drop(active);
let event = Event {
id: sub_id,
msg: EventMsg::TaskComplete(TaskCompleteEvent { last_agent_message }),
};
self.send_event(event).await;
let event = EventMsg::TaskComplete(TaskCompleteEvent { last_agent_message });
self.send_event(turn_context.as_ref(), event).await;
}
async fn register_new_active_task(&self, sub_id: String, task: RunningTask) {
async fn register_new_active_task(&self, task: RunningTask) {
let mut active = self.active_turn.lock().await;
let mut turn = ActiveTurn::default();
turn.add_task(sub_id, task);
turn.add_task(task);
*active = Some(turn);
}
async fn take_all_running_tasks(&self) -> Vec<(String, RunningTask)> {
async fn take_all_running_tasks(&self) -> Vec<RunningTask> {
let mut active = self.active_turn.lock().await;
match active.take() {
Some(mut at) => {
at.clear_pending().await;
let tasks = at.drain_tasks();
tasks.into_iter().collect()
at.drain_tasks()
}
None => Vec::new(),
}
}
async fn handle_task_abort(
self: &Arc<Self>,
sub_id: String,
task: RunningTask,
reason: TurnAbortReason,
) {
async fn handle_task_abort(self: &Arc<Self>, task: RunningTask, reason: TurnAbortReason) {
let sub_id = task.turn_context.sub_id.clone();
if task.cancellation_token.is_cancelled() {
return;
}
@@ -187,13 +178,12 @@ impl Session {
task.handle.abort();
let session_ctx = Arc::new(SessionTaskContext::new(Arc::clone(self)));
session_task.abort(session_ctx, &sub_id).await;
session_task
.abort(session_ctx, Arc::clone(&task.turn_context))
.await;
let event = Event {
id: sub_id.clone(),
msg: EventMsg::TurnAborted(TurnAbortedEvent { reason }),
};
self.send_event(event).await;
let event = EventMsg::TurnAborted(TurnAbortedEvent { reason });
self.send_event(task.turn_context.as_ref(), event).await;
}
}