headless_lms_chatbot/azure_chatbot/turn/
cancellation.rs

1//! Keeps a turn that the client abandoned from losing what it had already streamed.
2
3use std::pin::Pin;
4use std::sync::{
5    Arc,
6    atomic::{self, AtomicBool},
7};
8use std::task::{Context, Poll};
9
10use bytes::Bytes;
11use futures::Stream;
12use pin_project::pin_project;
13use sqlx::PgPool;
14use tokio::sync::Mutex;
15
16use crate::azure_chatbot::client_tool_calls::repair::answer_unfinished_tool_calls;
17use crate::chatbot_error::ChatbotResult;
18use crate::llm_utils::estimate_tokens;
19use crate::prelude::*;
20
21/// Ties a turn's cancellation guard to its response stream, so cleanup runs when the client drops
22/// the stream rather than whenever the turn's driver function returns.
23#[pin_project]
24pub(super) struct GuardedStream<S> {
25    guard: RequestCancelledGuard,
26    #[pin]
27    stream: S,
28}
29
30impl<S> GuardedStream<S> {
31    pub(super) fn new(guard: RequestCancelledGuard, stream: S) -> Self {
32        Self { guard, stream }
33    }
34}
35
36impl<S> Stream for GuardedStream<S>
37where
38    S: Stream<Item = ChatbotResult<Bytes>> + Send,
39{
40    type Item = S::Item;
41
42    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
43        let this = self.project();
44        let polled = this.stream.poll_next(cx);
45        // Log stream errors here in the clean format; actix's dispatcher otherwise only
46        // surfaces them as a terse Display line once the error is in the response body.
47        if let Poll::Ready(Some(Err(error))) = &polled {
48            error!("Chatbot response stream error:\n{error:?}");
49        }
50        polled
51    }
52}
53
54pub(super) struct RequestCancelledGuard {
55    pub(super) conversation_id: Uuid,
56    /// The responses this turn's rounds were given, which is what bounds the tool calls the
57    /// cleanup may answer to the ones this turn made.
58    pub(super) response_ids: Arc<Mutex<Vec<String>>>,
59    pub(super) response_message_id: Arc<Mutex<Option<Uuid>>>,
60    pub(super) full_response_text: Arc<Mutex<String>>,
61    pub(super) pool: PgPool,
62    pub(super) done: Arc<AtomicBool>,
63}
64
65impl Drop for RequestCancelledGuard {
66    fn drop(&mut self) {
67        if self.done.load(atomic::Ordering::Relaxed) {
68            return;
69        }
70        info!("Request ended before the turn completed. Cleaning up.");
71        // Nothing awaits this task, so failures are logged instead of panicked on: a panic here
72        // would be invisible apart from a stray tracing event.
73        tokio::spawn(clean_up_abandoned_turn(
74            self.pool.clone(),
75            self.conversation_id,
76            self.response_ids.clone(),
77            self.response_message_id.clone(),
78            self.full_response_text.clone(),
79        ));
80    }
81}
82
83/// Cleans up after a turn the client abandoned mid-stream: answers the tool calls it left without
84/// an output, then deletes the response message if it never received any text, or saves the text
85/// it did receive as its (incomplete) answer.
86///
87/// The tool calls are answered first and whether or not there is a message: a call with no output
88/// makes the LLM reject every later message of the conversation, while a turn that died in a tool
89/// round created no message at all.
90///
91/// Unlike [`save_partial_answer`], this decides whether there is anything to save.
92async fn clean_up_abandoned_turn(
93    pool: PgPool,
94    conversation_id: Uuid,
95    response_ids: Arc<Mutex<Vec<String>>>,
96    response_message_id: Arc<Mutex<Option<Uuid>>>,
97    full_response_text: Arc<Mutex<String>>,
98) {
99    let mut conn = match pool.acquire().await {
100        Ok(conn) => conn,
101        Err(err) => {
102            error!(
103                "Could not acquire a connection to clean up after a cancelled chatbot request: {err}"
104            );
105            return;
106        }
107    };
108    let response_ids = response_ids.lock().await.clone();
109    if let Err(err) = answer_unfinished_tool_calls(&mut conn, conversation_id, &response_ids).await
110    {
111        error!("Could not answer the tool calls an abandoned chatbot turn left unfinished: {err}");
112    }
113    info!("Verifying the received message has been handled");
114    let Some(id) = response_message_id.lock().await.to_owned() else {
115        info!("No response message was created for this request, nothing else to clean up.");
116        return;
117    };
118    let full_response_text = full_response_text.lock().await;
119    if full_response_text.is_empty() {
120        info!("No response received. Deleting the response message");
121        if let Err(err) = models::chatbot_conversation_messages::delete(&mut conn, id).await {
122            error!("Could not delete the empty chatbot response message {id}: {err}");
123        }
124        return;
125    }
126    info!("Response received but not completed. Saving the text received so far.");
127    let estimated_cost = estimate_tokens(&full_response_text);
128    // Below the default log level, same as the equivalent line in `parse_text_response`: the
129    // answer text is learner-facing content, which `summarize_input_for_log` exists to keep out
130    // of the request-side logs, and this is the response-side counterpart of that.
131    trace!(
132        "End of chatbot response stream. Estimated cost: {}. Response: {}",
133        estimated_cost, *full_response_text
134    );
135    if let Err(err) = save_partial_answer(&mut conn, id, &full_response_text, estimated_cost).await
136    {
137        error!("Could not save the partial chatbot response message {id}: {err}");
138    }
139}
140
141/// Saves `text` as the (incomplete) answer of the message `message_id` already names, billed for
142/// `used_tokens`.
143///
144/// Bumps the parent message row transactionally — see
145/// [`chatbot_conversation_messages::update`](models::chatbot_conversation_messages::update).
146/// Callers decide for themselves whether an empty `text` is worth saving at all; this always does.
147pub(super) async fn save_partial_answer(
148    conn: &mut PgConnection,
149    message_id: Uuid,
150    text: &str,
151    used_tokens: i32,
152) -> ChatbotResult<()> {
153    models::chatbot_conversation_messages::update(conn, message_id, text, true, used_tokens)
154        .await?;
155    Ok(())
156}
headless_lms_chatbot/azure_chatbot/turn/cancellation.rs

headless_lms_chatbot/azure_chatbot/turn/
cancellation.rs
