storkit/server/src/http/chat.rs

use crate::http::context::{AppContext, OpenApiResult, bad_request};
use crate::llm::chat;
use poem_openapi::{OpenApi, Tags, payload::Json};
use std::sync::Arc;

#[derive(Tags)]
enum ChatTags {
    Chat,
}

pub struct ChatApi {
    pub ctx: Arc<AppContext>,
}

#[OpenApi(tag = "ChatTags::Chat")]
impl ChatApi {
    /// Cancel the currently running chat stream, if any.
    ///
    /// Returns `true` once the cancellation signal is issued.
    #[oai(path = "/chat/cancel", method = "post")]
    async fn cancel_chat(&self) -> OpenApiResult<Json<bool>> {
        chat::cancel_chat(&self.ctx.state).map_err(bad_request)?;
        Ok(Json(true))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    fn test_api(dir: &TempDir) -> ChatApi {
        ChatApi {
            ctx: Arc::new(AppContext::new_test(dir.path().to_path_buf())),
        }
    }

    #[tokio::test]
    async fn cancel_chat_returns_true() {
        let dir = TempDir::new().unwrap();
        let api = test_api(&dir);
        let result = api.cancel_chat().await;
        assert!(result.is_ok());
        assert!(result.unwrap().0);
    }

    #[tokio::test]
    async fn cancel_chat_sends_cancel_signal() {
        let dir = TempDir::new().unwrap();
        let api = test_api(&dir);
        let mut cancel_rx = api.ctx.state.cancel_rx.clone();
        cancel_rx.borrow_and_update();

        api.cancel_chat().await.unwrap();

        assert!(*cancel_rx.borrow());
    }
}