mirror of
https://github.com/MindWorkAI/AI-Studio.git
synced 2026-10-11 19:33:48 +00:00
Kept the runtime API responsive during blocking work (#1043)
This commit is contained in:
1 parent
4616d1a4b2
commit
fb756dd767
11 files changed
+389
-77
No files matched your search
@@ -130,6 +130,25 @@ Key modules:
|
||||
- `pandoc.rs` - Integration with Pandoc for document conversion
|
||||
- `log.rs` - Logging infrastructure using `flexi_logger`
|
||||
|
||||
**Runtime API handlers never block.** All calls of the .NET app share one HTTP/2 connection, and the
|
||||
task driving it may wait on exactly the Tokio worker which a blocking handler occupies, so a single
|
||||
blocking handler can hold up the whole app. Therefore:
|
||||
|
||||
- File and disk access, the OS keyring, waiting for a `std::sync::Mutex`, and CPU-bound work such as
|
||||
loading a tokenizer or scanning text run in `tokio::task::spawn_blocking`. Follow
|
||||
`run_qdrant_edge_request` in `qdrant_edge_database.rs`, `prepare_image` and `prepare_image_sync` in
|
||||
`image.rs`, or `sanitize_batch` in `prompt_injection/api.rs`.
|
||||
- When an API offers a callback, as the file dialogs do, await the callback instead of calling the
|
||||
blocking variant; see `await_dialog` in `file_actions.rs`.
|
||||
- Short, bounded work, such as a single metadata lookup or writing a log line, may stay on the worker.
|
||||
- When the blocking task fails, answer with an error, never with an empty value the app could take
|
||||
for a valid answer.
|
||||
- Never keep a `std::sync::MutexGuard` alive across an `.await`, not even with an explicit `drop`
|
||||
before it: the future of the handler is then no longer `Send`, and Axum refuses it. Scope the guard
|
||||
in a block instead.
|
||||
- `runtime_api::test_support::assert_runtime_stays_free` tests a handler which waits for a lock. Run
|
||||
such a test with `#[tokio::test(flavor = "multi_thread", worker_threads = 1)]`.
|
||||
|
||||
### .NET App (`app/MindWork AI Studio/`)
|
||||
**Entry point:** `app/MindWork AI Studio/Program.cs`
|
||||
|
||||
|
||||
@@ -26,6 +26,7 @@
|
||||
- Fixed AI Studio sometimes rebuilding the index of a data source from scratch when the configuration of your organization was applied while the data source was being updated.
|
||||
- Fixed justified texts being hyphenated by the rules of English even when AI Studio shows another language, such as German. Screen readers now also know which language AI Studio uses.
|
||||
- Fixed the log file of AI Studio showing the wrong time for some entries when AI Studio was busy. Each entry now shows when it actually happened.
|
||||
- Fixed AI Studio pausing for a moment while it opened the local index of a data source, such as a mailbox. Other work, for example applying the configuration of your organization, had to wait until then.
|
||||
- Updated the code contributions on the supporters page, which now thank everyone who has contributed code to AI Studio so far.
|
||||
- Upgraded several libraries to improve security.
|
||||
- Upgraded to Rust v1.99.0
|
||||
@@ -77,6 +77,26 @@ fn release_clipboard<B>(clipboard: &mut Option<B>) -> bool {
|
||||
|
||||
/// Sets the clipboard text to the provided encrypted text.
|
||||
pub async fn set_clipboard(_token: APIToken, encrypted_text: String) -> Json<SetClipboardResponse> {
|
||||
//
|
||||
// The clipboard backend of the operating system may take a moment to answer, a failed write is
|
||||
// retried with a fresh backend, and the lock waits for a write still running. None of this
|
||||
// belongs on a runtime worker, because all calls of the app share one HTTP/2 connection.
|
||||
//
|
||||
match tokio::task::spawn_blocking(move || set_clipboard_sync(encrypted_text)).await {
|
||||
Ok(response) => response,
|
||||
Err(error) => {
|
||||
let issue = format!("The clipboard worker failed: {error}");
|
||||
error!(Source = "Clipboard"; "{issue}");
|
||||
Json(SetClipboardResponse {
|
||||
success: false,
|
||||
issue,
|
||||
})
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Performs the blocking part of [`set_clipboard`].
|
||||
fn set_clipboard_sync(encrypted_text: String) -> Json<SetClipboardResponse> {
|
||||
let encrypted_text = EncryptedText::new(encrypted_text);
|
||||
|
||||
// Decrypt this text first:
|
||||
|
||||
+25
-24
@@ -1,4 +1,5 @@
|
||||
use crate::api_token::APIToken;
|
||||
use axum::http::StatusCode;
|
||||
use axum::Json;
|
||||
use log::{debug, error, info, warn};
|
||||
use serde::Serialize;
|
||||
@@ -641,35 +642,35 @@ struct EnterpriseSourceData {
|
||||
encryption_secret: String,
|
||||
}
|
||||
|
||||
pub async fn read_enterprise_env_config_id(_token: APIToken) -> String {
|
||||
debug!("Trying to read the effective enterprise configuration ID.");
|
||||
resolve_effective_enterprise_config_source()
|
||||
.configs
|
||||
.into_iter()
|
||||
.next()
|
||||
.map(|config| config.id)
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
pub async fn read_enterprise_env_config_server_url(_token: APIToken) -> String {
|
||||
debug!("Trying to read the effective enterprise configuration server URL.");
|
||||
resolve_effective_enterprise_config_source()
|
||||
.configs
|
||||
.into_iter()
|
||||
.next()
|
||||
.map(|config| config.server_url)
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
pub async fn read_enterprise_env_config_encryption_secret(_token: APIToken) -> String {
|
||||
pub async fn read_enterprise_env_config_encryption_secret(_token: APIToken) -> Result<String, (StatusCode, String)> {
|
||||
debug!("Trying to read the effective enterprise configuration encryption secret.");
|
||||
resolve_effective_enterprise_secret_source().encryption_secret
|
||||
read_enterprise_sources(|| resolve_effective_enterprise_secret_source().encryption_secret).await
|
||||
}
|
||||
|
||||
/// Returns all enterprise configurations from the effective source.
|
||||
pub async fn read_enterprise_configs(_token: APIToken) -> Json<Vec<EnterpriseConfig>> {
|
||||
pub async fn read_enterprise_configs(_token: APIToken) -> Result<Json<Vec<EnterpriseConfig>>, (StatusCode, String)> {
|
||||
info!("Trying to read the effective enterprise configurations.");
|
||||
Json(resolve_effective_enterprise_config_source().configs)
|
||||
read_enterprise_sources(|| Json(resolve_effective_enterprise_config_source().configs)).await
|
||||
}
|
||||
|
||||
/// Reads the enterprise sources on the blocking pool.
|
||||
///
|
||||
/// Every call reads the policy files and, on Windows, the registry again. That is quick most of
|
||||
/// the time, but a slow disk must not hold up every other call of the app, because they all share
|
||||
/// one HTTP/2 connection.
|
||||
///
|
||||
/// A failure is answered as an error, never as an empty value: the app would read an empty value
|
||||
/// as "no configuration of your organization" and act on it.
|
||||
async fn read_enterprise_sources<T, F>(read: F) -> Result<T, (StatusCode, String)>
|
||||
where
|
||||
T: Send + 'static,
|
||||
F: FnOnce() -> T + Send + 'static,
|
||||
{
|
||||
tokio::task::spawn_blocking(read).await.map_err(|error| {
|
||||
let issue = format!("The enterprise configuration worker failed: {error}");
|
||||
error!("{issue}");
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, issue)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn resolve_external_http_custom_root_certificate_policy() -> ExternalHttpCustomRootCertificatePolicy {
|
||||
|
||||
+51
-17
@@ -123,26 +123,25 @@ pub async fn select_directory(
|
||||
Query(query): Query<SelectDirectoryQuery>,
|
||||
previous_directory: Option<Json<PreviousDirectory>>,
|
||||
) -> Json<DirectorySelectionResponse> {
|
||||
let main_window_lock = MAIN_WINDOW.lock().unwrap();
|
||||
let main_window = match main_window_lock.as_ref() {
|
||||
Some(window) => window,
|
||||
None => {
|
||||
error!(Source = "Tauri"; "Cannot open directory dialog: main window not available.");
|
||||
return Json(DirectorySelectionResponse {
|
||||
user_cancelled: true,
|
||||
selected_directory: String::from(""),
|
||||
});
|
||||
}
|
||||
let dialog = MAIN_WINDOW
|
||||
.lock()
|
||||
.unwrap()
|
||||
.as_ref()
|
||||
.map(|w| w.dialog().file().set_parent(w).set_title(&query.title));
|
||||
|
||||
let Some(mut dialog) = dialog else {
|
||||
error!(Source = "Tauri"; "Cannot open directory dialog: main window not available.");
|
||||
return Json(DirectorySelectionResponse {
|
||||
user_cancelled: true,
|
||||
selected_directory: String::from(""),
|
||||
});
|
||||
};
|
||||
|
||||
let mut dialog = main_window.dialog().file().set_parent(main_window).set_title(&query.title);
|
||||
if let Some(previous) = previous_directory {
|
||||
dialog = dialog.set_directory(previous.path.clone());
|
||||
}
|
||||
|
||||
drop(main_window_lock);
|
||||
|
||||
let folder_path = dialog.blocking_pick_folder();
|
||||
let folder_path = await_dialog(|answer| dialog.pick_folder(answer)).await;
|
||||
match folder_path {
|
||||
Some(path) => {
|
||||
match path.into_path() {
|
||||
@@ -203,7 +202,7 @@ pub async fn select_file(
|
||||
}
|
||||
|
||||
// Show the file dialog and get the selected file path:
|
||||
let file_path = file_dialog.blocking_pick_file();
|
||||
let file_path = await_dialog(|answer| file_dialog.pick_file(answer)).await;
|
||||
match file_path {
|
||||
Some(path) => match path.into_path() {
|
||||
Ok(pb) => {
|
||||
@@ -262,7 +261,7 @@ pub async fn select_files(
|
||||
}
|
||||
|
||||
// Show the file dialog and get the selected file path:
|
||||
let file_paths = file_dialog.blocking_pick_files();
|
||||
let file_paths = await_dialog(|answer| file_dialog.pick_files(answer)).await;
|
||||
match file_paths {
|
||||
Some(paths) => {
|
||||
let converted: Vec<String> = paths.into_iter().filter_map(|p| p.into_path().ok()).map(|pb| pb.to_string_lossy().to_string()).collect();
|
||||
@@ -314,7 +313,7 @@ pub async fn save_file(_token: APIToken, payload: Json<SaveFileOptions>) -> Json
|
||||
}
|
||||
|
||||
// Displays the file dialogue box and select the file:
|
||||
let file_path = file_dialog.blocking_save_file();
|
||||
let file_path = await_dialog(|answer| file_dialog.save_file(answer)).await;
|
||||
match file_path {
|
||||
Some(path) => match path.into_path() {
|
||||
Ok(pb) => {
|
||||
@@ -882,6 +881,26 @@ fn apply_filter<R: tauri::Runtime>(file_dialog: FileDialogBuilder<R>, filter: &O
|
||||
}
|
||||
}
|
||||
|
||||
/// Shows a file dialog and waits for the user's answer without holding up the runtime.
|
||||
///
|
||||
/// A dialog stays open until the user answers it, which can take minutes. The blocking variants of
|
||||
/// the dialog API would hold a runtime worker that long, and with it every other call of the app,
|
||||
/// because they all share one HTTP/2 connection. The plugin shows the dialog on the main thread and
|
||||
/// hands the answer to a callback, so this function awaits that callback instead.
|
||||
async fn await_dialog<T: Send + 'static>(show: impl FnOnce(Box<dyn FnOnce(Option<T>) + Send>)) -> Option<T> {
|
||||
let (sender, receiver) = tokio::sync::oneshot::channel();
|
||||
show(Box::new(move |answer| {
|
||||
// The receiver is only gone when the request was dropped, and then nobody waits for the answer:
|
||||
let _ = sender.send(answer);
|
||||
}));
|
||||
|
||||
receiver.await.unwrap_or_else(|_| {
|
||||
// The plugin drops the callback without calling it when it cannot reach the main thread:
|
||||
error!(Source = "Tauri"; "The file dialog ended without an answer.");
|
||||
None
|
||||
})
|
||||
}
|
||||
|
||||
fn split_save_file_path(file_path: &str) -> (Option<PathBuf>, Option<String>) {
|
||||
let path = Path::new(file_path);
|
||||
let directory = path
|
||||
@@ -1047,6 +1066,21 @@ mod tests {
|
||||
use super::*;
|
||||
use std::fs;
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_dialog_answer_arrives_through_its_callback() {
|
||||
let answer = await_dialog(|answer| {
|
||||
std::thread::spawn(move || answer(Some("Quarterly briefing.html")));
|
||||
}).await;
|
||||
|
||||
assert_eq!(answer, Some("Quarterly briefing.html"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_dialog_without_an_answer_counts_as_cancelled() {
|
||||
let answer = await_dialog::<String>(drop).await;
|
||||
assert_eq!(answer, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn save_file_options_accept_the_previous_file_contract() {
|
||||
let options: SaveFileOptions = serde_json::from_str(
|
||||
|
||||
@@ -524,7 +524,7 @@ pub async fn extract_data(
|
||||
// whole file, and a failure here means we cannot answer the request at all.
|
||||
//
|
||||
let tokenizer = if query.include_token_count {
|
||||
match crate::tokenizer::get_tokenizer(&query.tokenizer_path) {
|
||||
match crate::tokenizer::get_tokenizer_off_worker(query.tokenizer_path.clone()).await {
|
||||
Ok(tokenizer) => Some(tokenizer),
|
||||
Err(e) => {
|
||||
let error = ExtractionError::new(ExtractionErrorCode::InvalidRequest, format!("The tokenizer could not be loaded: {e}"));
|
||||
|
||||
@@ -45,13 +45,30 @@ pub struct SanitizeBatchResponse {
|
||||
pub results: Vec<SanitizeResponse>,
|
||||
}
|
||||
|
||||
pub async fn sanitize(_token: APIToken, Json(request): Json<SanitizeRequest>) -> Json<SanitizeResponse> {
|
||||
let (sanitized_text, report) = sanitize_text(&request.text);
|
||||
pub async fn sanitize(
|
||||
_token: APIToken,
|
||||
Json(request): Json<SanitizeRequest>,
|
||||
) -> Result<Json<SanitizeResponse>, (StatusCode, String)> {
|
||||
//
|
||||
// A single web page or retrieval context can still be large, and scanning it is CPU-bound. Like
|
||||
// a batch, it goes to the blocking pool, so it does not stall every other call the app makes
|
||||
// meanwhile.
|
||||
//
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let (sanitized_text, report) = sanitize_text(&request.text);
|
||||
|
||||
Json(SanitizeResponse {
|
||||
sanitized_text,
|
||||
findings: report.findings,
|
||||
redacted_count: report.redacted_count,
|
||||
Json(SanitizeResponse {
|
||||
sanitized_text,
|
||||
findings: report.findings,
|
||||
redacted_count: report.redacted_count,
|
||||
})
|
||||
})
|
||||
.await
|
||||
.map_err(|error| {
|
||||
(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("The prompt injection filter failed: {error}"),
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -473,15 +473,25 @@ fn qdrant_edge_base_path() -> QdrantEdgeResult<PathBuf> {
|
||||
}
|
||||
|
||||
pub async fn qdrant_edge_info(_token: APIToken) -> Json<QdrantEdgeServiceInfo> {
|
||||
let status = QDRANT_EDGE_STATUS.lock().unwrap();
|
||||
let current_status = status.status;
|
||||
let unavailable_reason = status.unavailable_reason.clone();
|
||||
drop(status);
|
||||
let (current_status, mut unavailable_reason) = {
|
||||
let status = QDRANT_EDGE_STATUS.lock().unwrap();
|
||||
(status.status, status.unavailable_reason.clone())
|
||||
};
|
||||
|
||||
let database_guard = QDRANT_EDGE_DATABASE.lock().unwrap();
|
||||
let database_info = database_guard
|
||||
.as_ref()
|
||||
.and_then(|database| database.info().ok());
|
||||
//
|
||||
// The info needs the lock of the database, which another request may hold for seconds while
|
||||
// it loads a store, and it reads the stores directory. Both belong on the blocking pool, see
|
||||
// run_qdrant_edge_request.
|
||||
//
|
||||
let database_info = match tokio::task::spawn_blocking(read_qdrant_edge_info).await {
|
||||
Ok(database_info) => database_info,
|
||||
Err(error) => {
|
||||
let reason = format!("The Qdrant Edge worker failed: {error}");
|
||||
error!(Source = "Qdrant Edge"; "Reading the Qdrant Edge info failed: {reason}");
|
||||
unavailable_reason = Some(reason);
|
||||
None
|
||||
},
|
||||
};
|
||||
|
||||
let is_available = current_status == QdrantEdgeStatus::Available && database_info.is_some();
|
||||
Json(QdrantEdgeServiceInfo {
|
||||
@@ -495,40 +505,48 @@ pub async fn qdrant_edge_info(_token: APIToken) -> Json<QdrantEdgeServiceInfo> {
|
||||
})
|
||||
}
|
||||
|
||||
fn read_qdrant_edge_info() -> Option<QdrantEdgeInfo> {
|
||||
QDRANT_EDGE_DATABASE
|
||||
.lock()
|
||||
.unwrap()
|
||||
.as_ref()
|
||||
.and_then(|database| database.info().ok())
|
||||
}
|
||||
|
||||
pub async fn ensure_qdrant_edge_store(_token: APIToken, Json(request): Json<EnsureQdrantEdgeStoreRequest>) -> Json<QdrantEdgeResponse<QdrantEdgeEnsureStoreResult>> {
|
||||
execute_qdrant_edge_request(|database| {
|
||||
run_qdrant_edge_request(move |database| {
|
||||
database.ensure_store_exists(&request.store_name, &request.data_source_name, request.vector_size)
|
||||
})
|
||||
}).await
|
||||
}
|
||||
|
||||
pub async fn insert_qdrant_edge_embedding(_token: APIToken, Json(request): Json<InsertQdrantEdgeEmbeddingRequest>) -> Json<QdrantEdgeResponse<()>> {
|
||||
execute_qdrant_edge_request(|database| {
|
||||
run_qdrant_edge_request(move |database| {
|
||||
database.insert_embedding(&request.store_name, request.points)
|
||||
})
|
||||
}).await
|
||||
}
|
||||
|
||||
pub async fn search_qdrant_edge_embeddings(_token: APIToken, Json(request): Json<SearchQdrantEdgeEmbeddingRequest>) -> Json<QdrantEdgeResponse<Vec<QdrantEdgeSearchResult>>> {
|
||||
execute_qdrant_edge_request(|database| {
|
||||
run_qdrant_edge_request(move |database| {
|
||||
database.search_embedding(&request.store_name, request.vector, request.max_matches, request.point_ids)
|
||||
})
|
||||
}).await
|
||||
}
|
||||
|
||||
pub async fn delete_qdrant_edge_embedding_by_file(_token: APIToken, Json(request): Json<DeleteQdrantEdgeEmbeddingByFileRequest>) -> Json<QdrantEdgeResponse<()>> {
|
||||
execute_qdrant_edge_request(|database| {
|
||||
run_qdrant_edge_request(move |database| {
|
||||
database.delete_embedding_by_file(&request.store_name, &request.file_path)
|
||||
})
|
||||
}).await
|
||||
}
|
||||
|
||||
pub async fn optimize_qdrant_edge_store(_token: APIToken, Json(request): Json<OptimizeQdrantEdgeStoreRequest>) -> Json<QdrantEdgeResponse<()>> {
|
||||
execute_qdrant_edge_request(|database| {
|
||||
run_qdrant_edge_request(move |database| {
|
||||
database.optimize_store(&request.store_name)
|
||||
})
|
||||
}).await
|
||||
}
|
||||
|
||||
pub async fn delete_qdrant_edge_store(_token: APIToken, Json(request): Json<DeleteQdrantEdgeStoreRequest>) -> Json<QdrantEdgeResponse<()>> {
|
||||
execute_qdrant_edge_request(|database| {
|
||||
run_qdrant_edge_request(move |database| {
|
||||
database.delete_store(&request.store_name)
|
||||
})
|
||||
}).await
|
||||
}
|
||||
|
||||
pub fn start_qdrant_edge_database<R: tauri::Runtime>(app_handle: tauri::AppHandle<R>) {
|
||||
@@ -569,6 +587,32 @@ pub fn stop_qdrant_edge_database() {
|
||||
set_qdrant_edge_unavailable("Qdrant Edge was stopped.".to_string());
|
||||
}
|
||||
|
||||
/// Runs one request against the database on the blocking pool.
|
||||
///
|
||||
/// A request waits for the lock of the database and may load a store from disk, which together
|
||||
/// can take seconds. On a runtime worker, that would hold up every other call of the app as well:
|
||||
/// they all share one HTTP/2 connection, and the task driving it may be waiting in the LIFO slot of
|
||||
/// exactly this worker, where no other worker can take it over.
|
||||
async fn run_qdrant_edge_request<T, F>(operation: F) -> Json<QdrantEdgeResponse<T>>
|
||||
where
|
||||
T: Serialize + Send + 'static,
|
||||
F: FnOnce(&mut QdrantEdgeDatabase) -> QdrantEdgeResult<T> + Send + 'static,
|
||||
{
|
||||
match tokio::task::spawn_blocking(move || execute_qdrant_edge_request(operation)).await {
|
||||
Ok(response) => response,
|
||||
Err(error) => {
|
||||
let issue = format!("The Qdrant Edge worker failed: {error}");
|
||||
error!(Source = "Qdrant Edge"; "Qdrant Edge request failed: {issue}");
|
||||
Json(QdrantEdgeResponse {
|
||||
success: false,
|
||||
issue,
|
||||
issue_code: "",
|
||||
data: None,
|
||||
})
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn execute_qdrant_edge_request<T, F>(operation: F) -> Json<QdrantEdgeResponse<T>>
|
||||
where
|
||||
T: Serialize,
|
||||
@@ -1033,6 +1077,7 @@ fn validate_store_name(store_name: &str) -> QdrantEdgeResult<()> {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::runtime_api::test_support::assert_runtime_stays_free;
|
||||
use qdrant_edge::WalOptions;
|
||||
|
||||
#[test]
|
||||
@@ -1254,6 +1299,27 @@ mod tests {
|
||||
fs::remove_dir_all(test_directory).unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
|
||||
async fn a_request_waiting_for_the_database_leaves_the_runtime_free() {
|
||||
let request = optimize_qdrant_edge_store(
|
||||
APIToken::from_hex_text("test"),
|
||||
Json(OptimizeQdrantEdgeStoreRequest { store_name: SEARCH_STORE.to_string() }),
|
||||
);
|
||||
|
||||
let response = assert_runtime_stays_free(&QDRANT_EDGE_DATABASE, request).await;
|
||||
|
||||
// No database runs in the tests, so the request ends as soon as it gets the lock.
|
||||
assert!(!response.success);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
|
||||
async fn the_info_waiting_for_the_database_leaves_the_runtime_free() {
|
||||
let request = qdrant_edge_info(APIToken::from_hex_text("test"));
|
||||
|
||||
let info = assert_runtime_stays_free(&QDRANT_EDGE_DATABASE, request).await;
|
||||
assert!(!info.is_available);
|
||||
}
|
||||
|
||||
const SEARCH_STORE: &str = "rag_6cc665a82b1e4d42bc748015b7b391ec";
|
||||
const POINT_X: &str = "0b5f1e8a-3c2d-4e6f-9a1b-7c8d9e0f1a2b";
|
||||
const POINT_Y: &str = "1c6a2f9b-4d3e-4f70-8b2c-8d9e0f1a2b3c";
|
||||
|
||||
@@ -64,8 +64,6 @@ pub fn start_runtime_api() {
|
||||
.route("/system/runtime/info", get(crate::environment::get_runtime_info))
|
||||
.route("/system/language", get(crate::environment::read_user_language))
|
||||
.route("/system/username", get(crate::environment::read_user_name))
|
||||
.route("/system/enterprise/config/id", get(crate::environment::read_enterprise_env_config_id))
|
||||
.route("/system/enterprise/config/server", get(crate::environment::read_enterprise_env_config_server_url))
|
||||
.route("/system/enterprise/config/encryption_secret", get(crate::environment::read_enterprise_env_config_encryption_secret))
|
||||
.route("/system/enterprise/configs", get(crate::environment::read_enterprise_configs))
|
||||
.route("/retrieval/fs/extract", get(crate::file_data::extract_data))
|
||||
@@ -113,3 +111,51 @@ fn install_rustls_crypto_provider() {
|
||||
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
|
||||
});
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) mod test_support {
|
||||
use std::sync::Mutex;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
/// How long the lock stays held while the request waits for it.
|
||||
const LOCK_DURATION: Duration = Duration::from_secs(1);
|
||||
|
||||
/// How long the probe may wait for the runtime. Far below the lock duration, so a request which
|
||||
/// blocks the worker while it waits fails the test every time.
|
||||
const PROBE_LIMIT: Duration = Duration::from_millis(200);
|
||||
|
||||
/// Starts a request while another thread holds the given lock, as if a long operation were
|
||||
/// running, and checks that a probe still gets through the runtime meanwhile. Returns the
|
||||
/// response once the lock is free again.
|
||||
///
|
||||
/// Call it from a test with a single worker, `#[tokio::test(flavor = "multi_thread",
|
||||
/// worker_threads = 1)]`. A request which blocks that worker while it waits then holds up the
|
||||
/// probe, just as in the app it holds up the one HTTP/2 connection which all calls share.
|
||||
pub(crate) async fn assert_runtime_stays_free<T, F>(lock: &'static Mutex<T>, request: F) -> F::Output
|
||||
where
|
||||
T: Send + 'static,
|
||||
F: Future + Send + 'static,
|
||||
F::Output: Send + 'static,
|
||||
{
|
||||
let (locked_sender, locked_receiver) = std::sync::mpsc::channel();
|
||||
let lock_holder = std::thread::spawn(move || {
|
||||
let guard = lock.lock().unwrap();
|
||||
locked_sender.send(()).unwrap();
|
||||
std::thread::sleep(LOCK_DURATION);
|
||||
drop(guard);
|
||||
});
|
||||
|
||||
locked_receiver.recv().unwrap();
|
||||
let request = tokio::spawn(request);
|
||||
|
||||
let probe_started = Instant::now();
|
||||
tokio::spawn(async {}).await.unwrap();
|
||||
let probe_duration = probe_started.elapsed();
|
||||
|
||||
lock_holder.join().unwrap();
|
||||
let response = request.await.unwrap();
|
||||
|
||||
assert!(probe_duration < PROBE_LIMIT, "The probe waited {probe_duration:?} for the runtime.");
|
||||
response
|
||||
}
|
||||
}
|
||||
@@ -78,8 +78,37 @@ pub fn init_secret_store() {
|
||||
}
|
||||
}
|
||||
|
||||
/// Runs one request against the credential store on the blocking pool.
|
||||
///
|
||||
/// The credential store of the operating system may take a moment to answer, and on macOS and Linux
|
||||
/// it may wait until the user unlocks it, which can take minutes. On a runtime worker, that would
|
||||
/// hold up every other call of the app, because they all share one HTTP/2 connection.
|
||||
async fn run_secret_store_request<T, F>(request: F) -> Result<T, String>
|
||||
where
|
||||
T: Send + 'static,
|
||||
F: FnOnce() -> T + Send + 'static,
|
||||
{
|
||||
tokio::task::spawn_blocking(request).await.map_err(|error| {
|
||||
let issue = format!("The secret store worker failed: {error}");
|
||||
error!(Source = "Secret Store"; "{issue}");
|
||||
issue
|
||||
})
|
||||
}
|
||||
|
||||
/// Stores a secret in the secret store using the operating system's keyring.
|
||||
pub async fn store_secret(_token: APIToken, request: Json<StoreSecret>) -> Json<StoreSecretResponse> {
|
||||
match run_secret_store_request(move || store_secret_sync(request)).await {
|
||||
Ok(response) => response,
|
||||
Err(issue) => Json(StoreSecretResponse {
|
||||
success: false,
|
||||
issue,
|
||||
issue_code: SecretStoreIssueCode::Unknown,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Performs the blocking part of [`store_secret`].
|
||||
fn store_secret_sync(request: Json<StoreSecret>) -> Json<StoreSecretResponse> {
|
||||
let user_name = request.user_name.as_str();
|
||||
let decrypted_text = match ENCRYPTION.decrypt(&request.secret) {
|
||||
Ok(text) => text,
|
||||
@@ -145,6 +174,19 @@ pub struct StoreSecretResponse {
|
||||
|
||||
/// Retrieves a secret from the secret store using the operating system's keyring.
|
||||
pub async fn get_secret(_token: APIToken, request: Json<RequestSecret>) -> Json<RequestedSecret> {
|
||||
match run_secret_store_request(move || get_secret_sync(request)).await {
|
||||
Ok(response) => response,
|
||||
Err(issue) => Json(RequestedSecret {
|
||||
success: false,
|
||||
secret: EncryptedText::new(String::from("")),
|
||||
issue,
|
||||
issue_code: SecretStoreIssueCode::Unknown,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Performs the blocking part of [`get_secret`].
|
||||
fn get_secret_sync(request: Json<RequestSecret>) -> Json<RequestedSecret> {
|
||||
let user_name = request.user_name.as_str();
|
||||
let service = format!("mindwork-ai-studio::{}", request.destination);
|
||||
let entry = match Entry::new(service.as_str(), user_name) {
|
||||
@@ -227,6 +269,19 @@ pub struct RequestedSecret {
|
||||
|
||||
/// Deletes a secret from the secret store using the operating system's keyring.
|
||||
pub async fn delete_secret(_token: APIToken, request: Json<RequestSecret>) -> Json<DeleteSecretResponse> {
|
||||
match run_secret_store_request(move || delete_secret_sync(request)).await {
|
||||
Ok(response) => response,
|
||||
Err(issue) => Json(DeleteSecretResponse {
|
||||
success: false,
|
||||
was_entry_found: false,
|
||||
issue,
|
||||
issue_code: SecretStoreIssueCode::Unknown,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Performs the blocking part of [`delete_secret`].
|
||||
fn delete_secret_sync(request: Json<RequestSecret>) -> Json<DeleteSecretResponse> {
|
||||
let user_name = request.user_name.as_str();
|
||||
let service = format!("mindwork-ai-studio::{}", request.destination);
|
||||
let entry = match Entry::new(service.as_str(), user_name) {
|
||||
|
||||
@@ -104,33 +104,58 @@ pub fn set_default_tokenizer_path(app_handle: tauri::AppHandle) {
|
||||
}
|
||||
|
||||
pub async fn token_count(_token: APIToken, req: Json<SetTokenText>) -> Json<TokenizerResponse> {
|
||||
match get_token_count(&req.tokenizer_path, &req.text) {
|
||||
match run_tokenizer_work(move || get_token_count(&req.tokenizer_path, &req.text)).await {
|
||||
Ok(count) => Json(TokenizerResponse::available(count)),
|
||||
Err(e) => Json(TokenizerResponse::unavailable(e)),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn validate_tokenizer(_token: APIToken, payload: Json<TokenizerPath>) -> Json<TokenizerResponse> {
|
||||
match handle_tokenizer_validate(&PathBuf::from(payload.file_path.clone())) {
|
||||
match run_tokenizer_work(move || handle_tokenizer_validate(&PathBuf::from(payload.file_path.clone()))).await {
|
||||
Ok(count) => Json(TokenizerResponse::available(count)),
|
||||
Err(e) => Json(TokenizerResponse::unavailable(e)),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn store_tokenizer(_token: APIToken, payload: Json<TokenizerStorage>) -> Json<TokenizerResponse> {
|
||||
match handle_tokenizer_store(&payload) {
|
||||
match run_tokenizer_work(move || handle_tokenizer_store(&payload).map_err(|e| e.to_string())).await {
|
||||
Ok(dest_path) => Json(TokenizerResponse::stored(dest_path)),
|
||||
Err(e) => Json(TokenizerResponse::unavailable(e.to_string())),
|
||||
Err(e) => Json(TokenizerResponse::unavailable(e)),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn delete_tokenizer(_token: APIToken, payload: Json<TokenizerDelete>) -> Json<TokenizerResponse> {
|
||||
match handle_tokenizer_delete(&payload) {
|
||||
match run_tokenizer_work(move || handle_tokenizer_delete(&payload).map_err(|e| e.to_string())).await {
|
||||
Ok(_) => Json(TokenizerResponse::stored(String::new())),
|
||||
Err(e) => Json(TokenizerResponse::unavailable(e.to_string())),
|
||||
Err(e) => Json(TokenizerResponse::unavailable(e)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Gets the tokenizer for the given path, loading it on the blocking pool when it is not cached yet.
|
||||
pub async fn get_tokenizer_off_worker(path: String) -> Result<Arc<Tokenizer>, String> {
|
||||
run_tokenizer_work(move || get_tokenizer(&path)).await
|
||||
}
|
||||
|
||||
/// Runs tokenizer work on the blocking pool.
|
||||
///
|
||||
/// Loading a tokenizer reads and parses its file, which takes most of a second for a large one, and
|
||||
/// it waits for the storage lock while another request copies or deletes a tokenizer. Encoding a
|
||||
/// long text is CPU-bound as well. On a runtime worker, any of this would hold up every other call
|
||||
/// of the app, because they all share one HTTP/2 connection.
|
||||
async fn run_tokenizer_work<T, F>(work: F) -> Result<T, String>
|
||||
where
|
||||
T: Send + 'static,
|
||||
F: FnOnce() -> Result<T, String> + Send + 'static,
|
||||
{
|
||||
tokio::task::spawn_blocking(work)
|
||||
.await
|
||||
.unwrap_or_else(|error| {
|
||||
let issue = format!("The tokenizer worker failed: {error}");
|
||||
error!(Source = "Tokenizer"; "{issue}");
|
||||
Err(issue)
|
||||
})
|
||||
}
|
||||
|
||||
fn handle_tokenizer_validate(path: &PathBuf) -> Result<usize, String> {
|
||||
validate_tokenizer_file(path)
|
||||
}
|
||||
@@ -240,7 +265,7 @@ fn tokenizer_cache() -> &'static RwLock<HashMap<PathBuf, Arc<Tokenizer>>> {
|
||||
TOKENIZERS.get_or_init(|| RwLock::new(HashMap::new()))
|
||||
}
|
||||
|
||||
pub fn get_tokenizer(path: &str) -> Result<Arc<Tokenizer>, String> {
|
||||
fn get_tokenizer(path: &str) -> Result<Arc<Tokenizer>, String> {
|
||||
let resolved_path = resolve_tokenizer_path(path)?;
|
||||
let tokenizer_path = fs::canonicalize(&resolved_path)
|
||||
.map_err(|e| format!("Could not resolve tokenizer file '{}': {e}", resolved_path.display()))?;
|
||||
@@ -303,3 +328,31 @@ fn load_tokenizer_from_file(path: &PathBuf) -> Result<Tokenizer, String> {
|
||||
Tokenizer::from_file(path)
|
||||
.map_err(|e| format!("Failed to load tokenizer from '{}': {e}", path.display()))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::runtime_api::test_support::assert_runtime_stays_free;
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
|
||||
async fn a_token_count_waiting_for_the_storage_leaves_the_runtime_free() {
|
||||
//
|
||||
// Any file will do: the request waits for the storage lock before it reads the tokenizer.
|
||||
// Since this file is no tokenizer, the request fails once it gets the lock.
|
||||
//
|
||||
let tokenizer_path = std::env::temp_dir().join(format!("ai-studio-tokenizer-{}.json", std::process::id()));
|
||||
fs::write(&tokenizer_path, "This is not a tokenizer.").unwrap();
|
||||
|
||||
let request = token_count(
|
||||
APIToken::from_hex_text("test"),
|
||||
Json(SetTokenText {
|
||||
text: "Hello, world!".to_string(),
|
||||
tokenizer_path: tokenizer_path.to_string_lossy().to_string(),
|
||||
}),
|
||||
);
|
||||
|
||||
let response = assert_runtime_stays_free(&TOKENIZER_STORAGE_LOCK, request).await;
|
||||
fs::remove_file(&tokenizer_path).unwrap();
|
||||
assert!(!response.success);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user