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:
@@ -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 { .. }) => {
|
||||
|
||||
Reference in New Issue
Block a user