chore: use tokio mutex and async function to prevent blocking a worker (#3850)

### Why Use `tokio::sync::Mutex`

`std::sync::Mutex` are not _async-aware_. As a result, they will block
the entire thread instead of just yielding the task. Furthermore they
can be poisoned which is not the case of `tokio` Mutex.
This allows the Tokio runtime to continue running other tasks while
waiting for the lock, preventing deadlocks and performance bottlenecks.

In general, this is preferred in async environment
This commit is contained in:
jif-oai
2025-09-18 18:21:52 +01:00
committed by GitHub
parent 992b531180
commit 277fc6254e
2 changed files with 96 additions and 84 deletions

View File

@@ -1,7 +1,6 @@
use std::sync::Arc;
use super::AgentTask;
use super::MutexExt;
use super::Session;
use super::TurnContext;
use super::get_last_assistant_message_from_turn;
@@ -37,7 +36,7 @@ struct HistoryBridgeTemplate<'a> {
summary_text: &'a str,
}
pub(super) fn spawn_compact_task(
pub(super) async fn spawn_compact_task(
sess: Arc<Session>,
turn_context: Arc<TurnContext>,
sub_id: String,
@@ -50,7 +49,7 @@ pub(super) fn spawn_compact_task(
input,
SUMMARIZATION_PROMPT.to_string(),
);
sess.set_task(task);
sess.set_task(task).await;
}
pub(super) async fn run_inline_auto_compact_task(
@@ -114,7 +113,9 @@ async fn run_compact_task_inner(
) {
let initial_input_for_turn: ResponseInputItem = ResponseInputItem::from(input);
let instructions_override = compact_instructions;
let turn_input = sess.turn_input_with_history(vec![initial_input_for_turn.clone().into()]);
let turn_input = sess
.turn_input_with_history(vec![initial_input_for_turn.clone().into()])
.await;
let prompt = Prompt {
input: turn_input,
@@ -173,10 +174,10 @@ async fn run_compact_task_inner(
}
if remove_task_on_completion {
sess.remove_task(&sub_id);
sess.remove_task(&sub_id).await;
}
let history_snapshot = {
let state = sess.state.lock_unchecked();
let state = sess.state.lock().await;
state.history.contents()
};
let summary_text = get_last_assistant_message_from_turn(&history_snapshot).unwrap_or_default();
@@ -184,7 +185,7 @@ async fn run_compact_task_inner(
let initial_context = sess.build_initial_context(turn_context.as_ref());
let new_history = build_compacted_history(initial_context, &user_messages, &summary_text);
{
let mut state = sess.state.lock_unchecked();
let mut state = sess.state.lock().await;
state.history.replace(new_history);
}
@@ -288,7 +289,7 @@ async fn drain_to_completed(
};
match event {
Ok(ResponseEvent::OutputItemDone(item)) => {
let mut state = sess.state.lock_unchecked();
let mut state = sess.state.lock().await;
state.history.record_items(std::slice::from_ref(&item));
}
Ok(ResponseEvent::Completed { .. }) => {