mirror of
https://github.com/tinyhumansai/openhuman.git
synced 2026-07-27 21:08:00 +00:00
refactor(rust): split oversized modules (#2991)
This commit is contained in:
+2
-354
@@ -24,6 +24,7 @@ mod fake_camera;
|
||||
mod file_logging;
|
||||
mod gmessages_scanner;
|
||||
mod imessage_scanner;
|
||||
mod local_data_reset;
|
||||
mod loopback_oauth;
|
||||
#[cfg(target_os = "macos")]
|
||||
mod mascot_native_window;
|
||||
@@ -324,359 +325,6 @@ async fn start_core_process(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Reset the user's local OpenHuman data and bounce the embedded core.
|
||||
///
|
||||
/// Replaces the prior two-step UI flow that called the core JSON-RPC
|
||||
/// `openhuman.config_reset_local_data` (in-process removal) followed by
|
||||
/// `restart_core_process`. The in-process removal failed on Windows with
|
||||
/// `ERROR_SHARING_VIOLATION` (os error 32) because the running core held
|
||||
/// open handles to SQLite databases, log files, the Sentry session store,
|
||||
/// etc. inside the directory it was being asked to delete — see
|
||||
/// OPENHUMAN-TAURI-AF.
|
||||
///
|
||||
/// New order:
|
||||
///
|
||||
/// 1. Query the core for the **paths** it would remove (`config_get_data_paths`)
|
||||
/// while the core is still up — these are derived from the loaded config
|
||||
/// and the active workspace marker, so the core is authoritative.
|
||||
/// 2. Acquire the restart lock so a concurrent `restart_core_process` cannot
|
||||
/// interleave with the remove.
|
||||
/// 3. Shut down the embedded core. `CoreProcessHandle::shutdown` cancels
|
||||
/// the cancellation token and awaits the tokio task, which drops the
|
||||
/// SQLite pool, log writer, etc. — releasing every Windows file handle.
|
||||
/// 4. Remove the three paths (current data dir, default data dir, active
|
||||
/// workspace marker) from this process. Missing entries are non-fatal.
|
||||
/// 5. Restart the embedded core via `ensure_running`.
|
||||
///
|
||||
/// Returns `Ok(())` only when the core is back up and the directories are
|
||||
/// gone (or were already absent). Any step's `Err` short-circuits and
|
||||
/// surfaces to the UI, which already renders the message as a toast.
|
||||
#[tauri::command]
|
||||
async fn reset_local_data(
|
||||
state: tauri::State<'_, core_process::CoreProcessHandle>,
|
||||
) -> Result<(), String> {
|
||||
log::info!("[core] reset_local_data: command invoked from frontend");
|
||||
|
||||
// ── 1. Ask the core for the paths it would remove ────────────────────
|
||||
//
|
||||
// The core is authoritative for path resolution (it owns config
|
||||
// loading, the workspace marker, and the staging-vs-prod default-dir
|
||||
// suffix). Resolve while the core is still up so we don't duplicate
|
||||
// that logic here.
|
||||
let paths = fetch_data_paths().await?;
|
||||
log::info!(
|
||||
"[core] reset_local_data: paths resolved current={} default={} marker={}",
|
||||
paths.current_openhuman_dir.display(),
|
||||
paths.default_openhuman_dir.display(),
|
||||
paths.active_workspace_marker_path.display()
|
||||
);
|
||||
|
||||
// ── 2. Acquire the restart lock ─────────────────────────────────────
|
||||
//
|
||||
// Prevents a concurrent `restart_core_process` from re-spawning the
|
||||
// embedded server in the middle of the remove step.
|
||||
let _guard = state.inner().restart_lock().await;
|
||||
log::debug!("[core] reset_local_data: acquired restart lock");
|
||||
|
||||
// ── 3. Shut down the embedded core ──────────────────────────────────
|
||||
//
|
||||
// Drops the tokio task, which drops the SQLite pool, log writer, and
|
||||
// every other RAII owner of a file handle inside the data directory.
|
||||
// On Windows this is the load-bearing step for OPENHUMAN-TAURI-AF.
|
||||
state.inner().shutdown().await;
|
||||
log::info!("[core] reset_local_data: embedded core stopped");
|
||||
|
||||
// ── 3b. Release the host-process log file handle (issue #1615) ──────
|
||||
//
|
||||
// The daily-rotating log appender at `<data_dir>/logs/openhuman-*.log`
|
||||
// is owned by *this* Tauri host process, not by the embedded core
|
||||
// tokio task — so `shutdown()` above does not release it. On Windows
|
||||
// that lingering OS file handle causes `remove_dir_all(.openhuman)`
|
||||
// below to fail with `ERROR_SHARING_VIOLATION` (os error 32). Drop
|
||||
// the writer guard now so the background flushing thread exits and
|
||||
// the file handle is closed before the removal walks the tree.
|
||||
let log_guard_dropped = openhuman_core::core::logging::shutdown_file_guard();
|
||||
log::info!("[core] reset_local_data: shutdown_file_guard dropped guard = {log_guard_dropped}");
|
||||
|
||||
// ── 4. Remove the paths ─────────────────────────────────────────────
|
||||
//
|
||||
// Missing entries are non-fatal: the user may already have manually
|
||||
// cleared the dir, or the marker may not exist for fresh installs.
|
||||
//
|
||||
// Capture the first delete error (if any) instead of propagating with
|
||||
// `?` — we must still restart the embedded core in step 5 so the app
|
||||
// doesn't end up with the sidecar dead. The original delete error is
|
||||
// surfaced after the restart attempt.
|
||||
let delete_result: Result<(), String> = async {
|
||||
remove_path_if_exists(
|
||||
&paths.active_workspace_marker_path,
|
||||
"active workspace marker",
|
||||
)
|
||||
.await?;
|
||||
remove_dir_if_exists(&paths.current_openhuman_dir, "current openhuman dir").await?;
|
||||
if paths.default_openhuman_dir != paths.current_openhuman_dir {
|
||||
remove_dir_if_exists(&paths.default_openhuman_dir, "default openhuman dir").await?;
|
||||
} else {
|
||||
log::debug!(
|
||||
"[core] reset_local_data: default dir == current dir; already removed above"
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
.await;
|
||||
if let Err(ref e) = delete_result {
|
||||
log::warn!("[core] reset_local_data: delete step failed: {e}; will still restart core");
|
||||
}
|
||||
|
||||
// ── 5. Restart the embedded core ────────────────────────────────────
|
||||
//
|
||||
// Always attempt restart, even if delete failed — otherwise the user
|
||||
// is left with a dead sidecar. If restart itself fails, prefer the
|
||||
// original delete error (more actionable) over the restart error.
|
||||
let restart_result = state.inner().ensure_running().await;
|
||||
match (&delete_result, &restart_result) {
|
||||
(Ok(()), Ok(())) => log::info!("[core] reset_local_data: embedded core back up"),
|
||||
(Err(_), Ok(())) => log::warn!(
|
||||
"[core] reset_local_data: core restarted but delete step failed; surfacing delete error"
|
||||
),
|
||||
(Ok(()), Err(e)) => log::error!("[core] reset_local_data: core restart failed: {e}"),
|
||||
(Err(_), Err(e)) => log::error!(
|
||||
"[core] reset_local_data: both delete and restart failed; restart error: {e}"
|
||||
),
|
||||
}
|
||||
delete_result?;
|
||||
restart_result?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Resolved data paths returned by `config_get_data_paths`.
|
||||
struct ResolvedDataPaths {
|
||||
current_openhuman_dir: std::path::PathBuf,
|
||||
default_openhuman_dir: std::path::PathBuf,
|
||||
active_workspace_marker_path: std::path::PathBuf,
|
||||
}
|
||||
|
||||
fn is_windows_file_lock_raw_os_error(raw_os_error: Option<i32>) -> bool {
|
||||
matches!(raw_os_error, Some(32 | 33))
|
||||
}
|
||||
|
||||
fn is_windows_file_lock_error(error: &std::io::Error) -> bool {
|
||||
cfg!(windows) && is_windows_file_lock_raw_os_error(error.raw_os_error())
|
||||
}
|
||||
|
||||
/// Returns:
|
||||
/// * `Ok(())` — the underlying remove failure should be swallowed (e.g.
|
||||
/// the path disappeared between the failed `remove_*` call and the
|
||||
/// reboot-fallback walk, so there is nothing left to clean up).
|
||||
/// * `Err(msg)` — a user-facing failure message the caller should surface
|
||||
/// to the UI / propagate up the reset flow.
|
||||
fn reset_local_data_delete_error(
|
||||
label: &str,
|
||||
path: &std::path::Path,
|
||||
error: &std::io::Error,
|
||||
) -> Result<(), String> {
|
||||
if is_windows_file_lock_error(error) {
|
||||
log::warn!(
|
||||
"[core] reset_local_data: Windows file lock blocked removal of {label} at {}: {error}",
|
||||
path.display()
|
||||
);
|
||||
|
||||
// Fallback: queue the still-locked sub-tree for deletion on the
|
||||
// next Windows boot via MoveFileExW + MOVEFILE_DELAY_UNTIL_REBOOT.
|
||||
// By this point in `reset_local_data` we have already:
|
||||
// * shut down the embedded core (drops every SQLite/log handle
|
||||
// the core task held), and
|
||||
// * released the host-process log appender via
|
||||
// `shutdown_file_guard()` (drops the rolling log file handle).
|
||||
// So any remaining lock now comes from *outside* this process —
|
||||
// anti-virus / file indexer / sibling app / Explorer — and cannot
|
||||
// be released by closing more OpenHuman windows. See issue #1615.
|
||||
#[cfg(target_os = "windows")]
|
||||
{
|
||||
return schedule_reboot_delete_or_describe(label, path, error);
|
||||
}
|
||||
// `is_windows_file_lock_error` is gated on `cfg!(windows)`, so on
|
||||
// Linux/macOS this branch is unreachable at runtime — but cargo
|
||||
// still type-checks the file for those targets and needs a value
|
||||
// of type `String`.
|
||||
#[cfg(not(target_os = "windows"))]
|
||||
{
|
||||
return Err(format!(
|
||||
"Failed to remove {label} at {} because it is locked by another OpenHuman window or process. Close all OpenHuman windows and try again. ({error})",
|
||||
path.display()
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
Err(format!(
|
||||
"Failed to remove {label} at {}: {error}",
|
||||
path.display()
|
||||
))
|
||||
}
|
||||
|
||||
/// Windows-only: ask the session manager to delete `path` (and its
|
||||
/// children if it is a directory) on the next reboot, and return either a
|
||||
/// user-facing message describing the outcome or `Ok(())` when the
|
||||
/// underlying failure should be treated as already-cleaned-up.
|
||||
#[cfg(target_os = "windows")]
|
||||
fn schedule_reboot_delete_or_describe(
|
||||
label: &str,
|
||||
path: &std::path::Path,
|
||||
original_error: &std::io::Error,
|
||||
) -> Result<(), String> {
|
||||
match reset_reboot_schedule::schedule_path_for_reboot_deletion(path) {
|
||||
Ok(summary) => {
|
||||
log::info!(
|
||||
"[core] reset_local_data: scheduled {label} at {} for reboot deletion (files={}, dirs={})",
|
||||
path.display(),
|
||||
summary.files,
|
||||
summary.dirs
|
||||
);
|
||||
Err(format!(
|
||||
"Couldn't remove {label} at {} right now because another process is holding it open ({original_error}). {} files and {} folders have been queued for deletion the next time you restart Windows — restart soon to finish the reset.",
|
||||
path.display(),
|
||||
summary.files,
|
||||
summary.dirs,
|
||||
))
|
||||
}
|
||||
// Race condition: the still-locked path disappeared between the
|
||||
// `remove_*` call that failed with `ERROR_SHARING_VIOLATION` and
|
||||
// the metadata read inside the reboot-schedule walk. Whoever else
|
||||
// held the handle has already finished cleaning up, so the reset
|
||||
// goal is achieved — swallow the original lock error and treat
|
||||
// this as success. The empty partial schedule (no entries queued
|
||||
// yet) is what distinguishes "vanished cleanly" from "started
|
||||
// walking, then hit a real error."
|
||||
Err(failure)
|
||||
if failure.error.kind() == std::io::ErrorKind::NotFound
|
||||
&& failure.partial.total() == 0 =>
|
||||
{
|
||||
log::info!(
|
||||
"[core] reset_local_data: {label} at {} disappeared between lock failure and reboot fallback; treating as removed",
|
||||
path.display(),
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(failure) => {
|
||||
let partial_total = failure.partial.total();
|
||||
log::error!(
|
||||
"[core] reset_local_data: reboot delete fallback failed for {label} at {}: {} (partial schedule: files={}, dirs={})",
|
||||
path.display(),
|
||||
failure.error,
|
||||
failure.partial.files,
|
||||
failure.partial.dirs,
|
||||
);
|
||||
if partial_total == 0 {
|
||||
Err(format!(
|
||||
"Failed to remove {label} at {} because it is locked by another OpenHuman window or process, and scheduling deletion on next reboot also failed ({}). Close all OpenHuman windows and try again. ({original_error})",
|
||||
path.display(),
|
||||
failure.error,
|
||||
))
|
||||
} else {
|
||||
Err(format!(
|
||||
"Failed to remove {label} at {} because it is locked by another OpenHuman window or process. {} files and {} folders were queued for the next reboot before scheduling failed ({}); the rest still needs manual cleanup. Close all OpenHuman windows and try again. ({original_error})",
|
||||
path.display(),
|
||||
failure.partial.files,
|
||||
failure.partial.dirs,
|
||||
failure.error,
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Call the core's `config_get_data_paths` RPC and parse the response.
|
||||
async fn fetch_data_paths() -> Result<ResolvedDataPaths, String> {
|
||||
let url = crate::core_rpc::core_rpc_url_value();
|
||||
let body = serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "openhuman.config_get_data_paths",
|
||||
"params": {}
|
||||
});
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(10))
|
||||
.build()
|
||||
.map_err(|e| format!("config_get_data_paths client build failed: {e}"))?;
|
||||
let req = crate::core_rpc::apply_auth(client.post(&url))?;
|
||||
let res = req
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("config_get_data_paths request failed: {e}"))?;
|
||||
if !res.status().is_success() {
|
||||
return Err(format!("config_get_data_paths http {}", res.status()));
|
||||
}
|
||||
let envelope: serde_json::Value = res
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("config_get_data_paths decode failed: {e}"))?;
|
||||
// JSON-RPC envelope wraps the `RpcOutcome` result twice:
|
||||
// `{ "result": { "result": { ...paths... }, "logs": [...] } }`.
|
||||
let inner = envelope
|
||||
.pointer("/result/result")
|
||||
.ok_or_else(|| "config_get_data_paths missing /result/result".to_string())?;
|
||||
let current = inner
|
||||
.get("current_openhuman_dir")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| "config_get_data_paths missing current_openhuman_dir".to_string())?;
|
||||
let default = inner
|
||||
.get("default_openhuman_dir")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| "config_get_data_paths missing default_openhuman_dir".to_string())?;
|
||||
let marker = inner
|
||||
.get("active_workspace_marker_path")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| "config_get_data_paths missing active_workspace_marker_path".to_string())?;
|
||||
Ok(ResolvedDataPaths {
|
||||
current_openhuman_dir: std::path::PathBuf::from(current),
|
||||
default_openhuman_dir: std::path::PathBuf::from(default),
|
||||
active_workspace_marker_path: std::path::PathBuf::from(marker),
|
||||
})
|
||||
}
|
||||
|
||||
/// Remove a regular file if present. Missing → debug log + Ok.
|
||||
async fn remove_path_if_exists(path: &std::path::Path, label: &str) -> Result<(), String> {
|
||||
match tokio::fs::remove_file(path).await {
|
||||
Ok(()) => {
|
||||
log::info!(
|
||||
"[core] reset_local_data: removed {label} at {}",
|
||||
path.display()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
|
||||
log::debug!(
|
||||
"[core] reset_local_data: {label} already absent at {}",
|
||||
path.display()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => reset_local_data_delete_error(label, path, &e),
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove a directory tree if present. Missing → debug log + Ok.
|
||||
async fn remove_dir_if_exists(path: &std::path::Path, label: &str) -> Result<(), String> {
|
||||
match tokio::fs::remove_dir_all(path).await {
|
||||
Ok(()) => {
|
||||
log::info!(
|
||||
"[core] reset_local_data: removed {label} at {}",
|
||||
path.display()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
|
||||
log::debug!(
|
||||
"[core] reset_local_data: {label} already absent at {}",
|
||||
path.display()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => reset_local_data_delete_error(label, path, &e),
|
||||
}
|
||||
}
|
||||
|
||||
/// Cleanly exit the application.
|
||||
///
|
||||
/// Called by the BootCheckGate "Quit" button when the core is unreachable and
|
||||
@@ -3427,7 +3075,7 @@ pub fn run() {
|
||||
restart_core_process,
|
||||
recover_port_conflict,
|
||||
start_core_process,
|
||||
reset_local_data,
|
||||
local_data_reset::reset_local_data,
|
||||
app_quit,
|
||||
restart_app,
|
||||
get_active_user_id,
|
||||
|
||||
@@ -58,90 +58,6 @@ fn overlay_parent_rpc_url_handles_empty() {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reset_local_data_windows_file_lock_error_codes_are_recognized() {
|
||||
assert!(is_windows_file_lock_raw_os_error(Some(32)));
|
||||
assert!(is_windows_file_lock_raw_os_error(Some(33)));
|
||||
assert!(!is_windows_file_lock_raw_os_error(Some(5)));
|
||||
assert!(!is_windows_file_lock_raw_os_error(None));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reset_local_data_delete_error_keeps_generic_message_for_other_errors() {
|
||||
let err = std::io::Error::from(std::io::ErrorKind::PermissionDenied);
|
||||
let result = reset_local_data_delete_error(
|
||||
"current openhuman dir",
|
||||
std::path::Path::new("/tmp/openhuman"),
|
||||
&err,
|
||||
);
|
||||
|
||||
let msg = result.expect_err("non-lock errors must still surface to the UI");
|
||||
assert!(msg.starts_with("Failed to remove current openhuman dir at /tmp/openhuman:"));
|
||||
assert!(!msg.contains("Close all OpenHuman windows and try again"));
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
#[test]
|
||||
fn reset_local_data_delete_error_swallows_lock_failure_when_path_disappeared() {
|
||||
// Race condition the reboot fallback now handles: the locked path
|
||||
// was gone by the time `schedule_path_for_reboot_deletion` ran its
|
||||
// `symlink_metadata` probe, so the reset goal is already met. The
|
||||
// helper must return `Ok(())` rather than surfacing a confusing
|
||||
// "couldn't remove (it's not there)" toast.
|
||||
let dir = tempfile::tempdir().expect("tempdir for reset error test");
|
||||
let missing = dir.path().join("definitely-not-there");
|
||||
|
||||
let err = std::io::Error::from_raw_os_error(32);
|
||||
let result = reset_local_data_delete_error("current openhuman dir", &missing, &err);
|
||||
|
||||
assert!(
|
||||
result.is_ok(),
|
||||
"expected NotFound + empty partial schedule to be swallowed as success, got {result:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
#[test]
|
||||
fn reset_local_data_delete_error_reports_reboot_schedule_counts() {
|
||||
// When the lock fallback can walk a real directory tree, the user
|
||||
// message should report how much has been queued so the support
|
||||
// log preserves "what was actually scheduled". Scheduling itself
|
||||
// may still fail at the MoveFileExW step in unprivileged test
|
||||
// processes (the registry key write requires administrator); the
|
||||
// fallback then carries a partial schedule that the error path
|
||||
// surfaces, so both branches must keep mentioning the lock cause
|
||||
// *and* expose either the queued counts or the schedule failure.
|
||||
let dir = tempfile::tempdir().expect("tempdir for reset error test");
|
||||
let target = dir.path().join("reset-mock");
|
||||
std::fs::create_dir_all(target.join("nested")).expect("mkdir nested");
|
||||
std::fs::write(target.join("a.txt"), b"x").expect("write a.txt");
|
||||
std::fs::write(target.join("nested").join("b.txt"), b"y").expect("write b.txt");
|
||||
|
||||
let err = std::io::Error::from_raw_os_error(32);
|
||||
let result = reset_local_data_delete_error("current openhuman dir", &target, &err);
|
||||
|
||||
// Path exists on disk, so the fallback must surface the outcome —
|
||||
// either an "all-queued" success-but-needs-reboot message (admin)
|
||||
// or one of the failure flavours (non-admin).
|
||||
let msg =
|
||||
result.expect_err("path exists, fallback must report queued counts or scheduling failure");
|
||||
let admin_path = msg.contains("queued for deletion the next time you restart Windows")
|
||||
&& msg.contains("2 files and 2 folders");
|
||||
let user_full_fail = msg.contains("scheduling deletion on next reboot also failed");
|
||||
let user_partial = msg.contains("queued for the next reboot before scheduling failed");
|
||||
assert!(
|
||||
admin_path || user_full_fail || user_partial,
|
||||
"expected reboot-scheduled, fully-failed, or partial-fail message, got: {msg}"
|
||||
);
|
||||
// Whatever branch we land on, the user must still be told the lock
|
||||
// is what blocked the immediate removal.
|
||||
assert!(
|
||||
msg.contains("locked by another OpenHuman window or process")
|
||||
|| msg.contains("another process is holding it open"),
|
||||
"missing lock cause: {msg}"
|
||||
);
|
||||
}
|
||||
|
||||
/// Tests for setup_tray conditional compilation
|
||||
/// The PR adds two versions of setup_tray():
|
||||
/// 1. No-op for linux + cef: logs warning and returns Ok(())
|
||||
|
||||
@@ -0,0 +1,360 @@
|
||||
use crate::core_process;
|
||||
#[cfg(target_os = "windows")]
|
||||
use crate::reset_reboot_schedule;
|
||||
|
||||
/// Reset the user's local OpenHuman data and bounce the embedded core.
|
||||
///
|
||||
/// Replaces the prior two-step UI flow that called the core JSON-RPC
|
||||
/// `openhuman.config_reset_local_data` (in-process removal) followed by
|
||||
/// `restart_core_process`. The in-process removal failed on Windows with
|
||||
/// `ERROR_SHARING_VIOLATION` (os error 32) because the running core held
|
||||
/// open handles to SQLite databases, log files, the Sentry session store,
|
||||
/// etc. inside the directory it was being asked to delete — see
|
||||
/// OPENHUMAN-TAURI-AF.
|
||||
///
|
||||
/// New order:
|
||||
///
|
||||
/// 1. Query the core for the **paths** it would remove (`config_get_data_paths`)
|
||||
/// while the core is still up — these are derived from the loaded config
|
||||
/// and the active workspace marker, so the core is authoritative.
|
||||
/// 2. Acquire the restart lock so a concurrent `restart_core_process` cannot
|
||||
/// interleave with the remove.
|
||||
/// 3. Shut down the embedded core. `CoreProcessHandle::shutdown` cancels
|
||||
/// the cancellation token and awaits the tokio task, which drops the
|
||||
/// SQLite pool, log writer, etc. — releasing every Windows file handle.
|
||||
/// 4. Remove the three paths (current data dir, default data dir, active
|
||||
/// workspace marker) from this process. Missing entries are non-fatal.
|
||||
/// 5. Restart the embedded core via `ensure_running`.
|
||||
///
|
||||
/// Returns `Ok(())` only when the core is back up and the directories are
|
||||
/// gone (or were already absent). Any step's `Err` short-circuits and
|
||||
/// surfaces to the UI, which already renders the message as a toast.
|
||||
#[tauri::command]
|
||||
pub async fn reset_local_data(
|
||||
state: tauri::State<'_, core_process::CoreProcessHandle>,
|
||||
) -> Result<(), String> {
|
||||
log::info!("[core] reset_local_data: command invoked from frontend");
|
||||
|
||||
// ── 1. Ask the core for the paths it would remove ────────────────────
|
||||
//
|
||||
// The core is authoritative for path resolution (it owns config
|
||||
// loading, the workspace marker, and the staging-vs-prod default-dir
|
||||
// suffix). Resolve while the core is still up so we don't duplicate
|
||||
// that logic here.
|
||||
let paths = fetch_data_paths().await?;
|
||||
log::info!(
|
||||
"[core] reset_local_data: paths resolved current={} default={} marker={}",
|
||||
paths.current_openhuman_dir.display(),
|
||||
paths.default_openhuman_dir.display(),
|
||||
paths.active_workspace_marker_path.display()
|
||||
);
|
||||
|
||||
// ── 2. Acquire the restart lock ─────────────────────────────────────
|
||||
//
|
||||
// Prevents a concurrent `restart_core_process` from re-spawning the
|
||||
// embedded server in the middle of the remove step.
|
||||
let _guard = state.inner().restart_lock().await;
|
||||
log::debug!("[core] reset_local_data: acquired restart lock");
|
||||
|
||||
// ── 3. Shut down the embedded core ──────────────────────────────────
|
||||
//
|
||||
// Drops the tokio task, which drops the SQLite pool, log writer, and
|
||||
// every other RAII owner of a file handle inside the data directory.
|
||||
// On Windows this is the load-bearing step for OPENHUMAN-TAURI-AF.
|
||||
state.inner().shutdown().await;
|
||||
log::info!("[core] reset_local_data: embedded core stopped");
|
||||
|
||||
// ── 3b. Release the host-process log file handle (issue #1615) ──────
|
||||
//
|
||||
// The daily-rotating log appender at `<data_dir>/logs/openhuman-*.log`
|
||||
// is owned by *this* Tauri host process, not by the embedded core
|
||||
// tokio task — so `shutdown()` above does not release it. On Windows
|
||||
// that lingering OS file handle causes `remove_dir_all(.openhuman)`
|
||||
// below to fail with `ERROR_SHARING_VIOLATION` (os error 32). Drop
|
||||
// the writer guard now so the background flushing thread exits and
|
||||
// the file handle is closed before the removal walks the tree.
|
||||
let log_guard_dropped = openhuman_core::core::logging::shutdown_file_guard();
|
||||
log::info!("[core] reset_local_data: shutdown_file_guard dropped guard = {log_guard_dropped}");
|
||||
|
||||
// ── 4. Remove the paths ─────────────────────────────────────────────
|
||||
//
|
||||
// Missing entries are non-fatal: the user may already have manually
|
||||
// cleared the dir, or the marker may not exist for fresh installs.
|
||||
//
|
||||
// Capture the first delete error (if any) instead of propagating with
|
||||
// `?` — we must still restart the embedded core in step 5 so the app
|
||||
// doesn't end up with the sidecar dead. The original delete error is
|
||||
// surfaced after the restart attempt.
|
||||
let delete_result: Result<(), String> = async {
|
||||
remove_path_if_exists(
|
||||
&paths.active_workspace_marker_path,
|
||||
"active workspace marker",
|
||||
)
|
||||
.await?;
|
||||
remove_dir_if_exists(&paths.current_openhuman_dir, "current openhuman dir").await?;
|
||||
if paths.default_openhuman_dir != paths.current_openhuman_dir {
|
||||
remove_dir_if_exists(&paths.default_openhuman_dir, "default openhuman dir").await?;
|
||||
} else {
|
||||
log::debug!(
|
||||
"[core] reset_local_data: default dir == current dir; already removed above"
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
.await;
|
||||
if let Err(ref e) = delete_result {
|
||||
log::warn!("[core] reset_local_data: delete step failed: {e}; will still restart core");
|
||||
}
|
||||
|
||||
// ── 5. Restart the embedded core ────────────────────────────────────
|
||||
//
|
||||
// Always attempt restart, even if delete failed — otherwise the user
|
||||
// is left with a dead sidecar. If restart itself fails, prefer the
|
||||
// original delete error (more actionable) over the restart error.
|
||||
let restart_result = state.inner().ensure_running().await;
|
||||
match (&delete_result, &restart_result) {
|
||||
(Ok(()), Ok(())) => log::info!("[core] reset_local_data: embedded core back up"),
|
||||
(Err(_), Ok(())) => log::warn!(
|
||||
"[core] reset_local_data: core restarted but delete step failed; surfacing delete error"
|
||||
),
|
||||
(Ok(()), Err(e)) => log::error!("[core] reset_local_data: core restart failed: {e}"),
|
||||
(Err(_), Err(e)) => log::error!(
|
||||
"[core] reset_local_data: both delete and restart failed; restart error: {e}"
|
||||
),
|
||||
}
|
||||
delete_result?;
|
||||
restart_result?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Resolved data paths returned by `config_get_data_paths`.
|
||||
struct ResolvedDataPaths {
|
||||
current_openhuman_dir: std::path::PathBuf,
|
||||
default_openhuman_dir: std::path::PathBuf,
|
||||
active_workspace_marker_path: std::path::PathBuf,
|
||||
}
|
||||
|
||||
fn is_windows_file_lock_raw_os_error(raw_os_error: Option<i32>) -> bool {
|
||||
matches!(raw_os_error, Some(32 | 33))
|
||||
}
|
||||
|
||||
fn is_windows_file_lock_error(error: &std::io::Error) -> bool {
|
||||
cfg!(windows) && is_windows_file_lock_raw_os_error(error.raw_os_error())
|
||||
}
|
||||
|
||||
/// Returns:
|
||||
/// * `Ok(())` — the underlying remove failure should be swallowed (e.g.
|
||||
/// the path disappeared between the failed `remove_*` call and the
|
||||
/// reboot-fallback walk, so there is nothing left to clean up).
|
||||
/// * `Err(msg)` — a user-facing failure message the caller should surface
|
||||
/// to the UI / propagate up the reset flow.
|
||||
fn reset_local_data_delete_error(
|
||||
label: &str,
|
||||
path: &std::path::Path,
|
||||
error: &std::io::Error,
|
||||
) -> Result<(), String> {
|
||||
if is_windows_file_lock_error(error) {
|
||||
log::warn!(
|
||||
"[core] reset_local_data: Windows file lock blocked removal of {label} at {}: {error}",
|
||||
path.display()
|
||||
);
|
||||
|
||||
// Fallback: queue the still-locked sub-tree for deletion on the
|
||||
// next Windows boot via MoveFileExW + MOVEFILE_DELAY_UNTIL_REBOOT.
|
||||
// By this point in `reset_local_data` we have already:
|
||||
// * shut down the embedded core (drops every SQLite/log handle
|
||||
// the core task held), and
|
||||
// * released the host-process log appender via
|
||||
// `shutdown_file_guard()` (drops the rolling log file handle).
|
||||
// So any remaining lock now comes from *outside* this process —
|
||||
// anti-virus / file indexer / sibling app / Explorer — and cannot
|
||||
// be released by closing more OpenHuman windows. See issue #1615.
|
||||
#[cfg(target_os = "windows")]
|
||||
{
|
||||
return schedule_reboot_delete_or_describe(label, path, error);
|
||||
}
|
||||
// `is_windows_file_lock_error` is gated on `cfg!(windows)`, so on
|
||||
// Linux/macOS this branch is unreachable at runtime — but cargo
|
||||
// still type-checks the file for those targets and needs a value
|
||||
// of type `String`.
|
||||
#[cfg(not(target_os = "windows"))]
|
||||
{
|
||||
return Err(format!(
|
||||
"Failed to remove {label} at {} because it is locked by another OpenHuman window or process. Close all OpenHuman windows and try again. ({error})",
|
||||
path.display()
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
Err(format!(
|
||||
"Failed to remove {label} at {}: {error}",
|
||||
path.display()
|
||||
))
|
||||
}
|
||||
|
||||
/// Windows-only: ask the session manager to delete `path` (and its
|
||||
/// children if it is a directory) on the next reboot, and return either a
|
||||
/// user-facing message describing the outcome or `Ok(())` when the
|
||||
/// underlying failure should be treated as already-cleaned-up.
|
||||
#[cfg(target_os = "windows")]
|
||||
fn schedule_reboot_delete_or_describe(
|
||||
label: &str,
|
||||
path: &std::path::Path,
|
||||
original_error: &std::io::Error,
|
||||
) -> Result<(), String> {
|
||||
match reset_reboot_schedule::schedule_path_for_reboot_deletion(path) {
|
||||
Ok(summary) => {
|
||||
log::info!(
|
||||
"[core] reset_local_data: scheduled {label} at {} for reboot deletion (files={}, dirs={})",
|
||||
path.display(),
|
||||
summary.files,
|
||||
summary.dirs
|
||||
);
|
||||
Err(format!(
|
||||
"Couldn't remove {label} at {} right now because another process is holding it open ({original_error}). {} files and {} folders have been queued for deletion the next time you restart Windows — restart soon to finish the reset.",
|
||||
path.display(),
|
||||
summary.files,
|
||||
summary.dirs,
|
||||
))
|
||||
}
|
||||
// Race condition: the still-locked path disappeared between the
|
||||
// `remove_*` call that failed with `ERROR_SHARING_VIOLATION` and
|
||||
// the metadata read inside the reboot-schedule walk. Whoever else
|
||||
// held the handle has already finished cleaning up, so the reset
|
||||
// goal is achieved — swallow the original lock error and treat
|
||||
// this as success. The empty partial schedule (no entries queued
|
||||
// yet) is what distinguishes "vanished cleanly" from "started
|
||||
// walking, then hit a real error."
|
||||
Err(failure)
|
||||
if failure.error.kind() == std::io::ErrorKind::NotFound
|
||||
&& failure.partial.total() == 0 =>
|
||||
{
|
||||
log::info!(
|
||||
"[core] reset_local_data: {label} at {} disappeared between lock failure and reboot fallback; treating as removed",
|
||||
path.display(),
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(failure) => {
|
||||
let partial_total = failure.partial.total();
|
||||
log::error!(
|
||||
"[core] reset_local_data: reboot delete fallback failed for {label} at {}: {} (partial schedule: files={}, dirs={})",
|
||||
path.display(),
|
||||
failure.error,
|
||||
failure.partial.files,
|
||||
failure.partial.dirs,
|
||||
);
|
||||
if partial_total == 0 {
|
||||
Err(format!(
|
||||
"Failed to remove {label} at {} because it is locked by another OpenHuman window or process, and scheduling deletion on next reboot also failed ({}). Close all OpenHuman windows and try again. ({original_error})",
|
||||
path.display(),
|
||||
failure.error,
|
||||
))
|
||||
} else {
|
||||
Err(format!(
|
||||
"Failed to remove {label} at {} because it is locked by another OpenHuman window or process. {} files and {} folders were queued for the next reboot before scheduling failed ({}); the rest still needs manual cleanup. Close all OpenHuman windows and try again. ({original_error})",
|
||||
path.display(),
|
||||
failure.partial.files,
|
||||
failure.partial.dirs,
|
||||
failure.error,
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Call the core's `config_get_data_paths` RPC and parse the response.
|
||||
async fn fetch_data_paths() -> Result<ResolvedDataPaths, String> {
|
||||
let url = crate::core_rpc::core_rpc_url_value();
|
||||
let body = serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "openhuman.config_get_data_paths",
|
||||
"params": {}
|
||||
});
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(10))
|
||||
.build()
|
||||
.map_err(|e| format!("config_get_data_paths client build failed: {e}"))?;
|
||||
let req = crate::core_rpc::apply_auth(client.post(&url))?;
|
||||
let res = req
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("config_get_data_paths request failed: {e}"))?;
|
||||
if !res.status().is_success() {
|
||||
return Err(format!("config_get_data_paths http {}", res.status()));
|
||||
}
|
||||
let envelope: serde_json::Value = res
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("config_get_data_paths decode failed: {e}"))?;
|
||||
// JSON-RPC envelope wraps the `RpcOutcome` result twice:
|
||||
// `{ "result": { "result": { ...paths... }, "logs": [...] } }`.
|
||||
let inner = envelope
|
||||
.pointer("/result/result")
|
||||
.ok_or_else(|| "config_get_data_paths missing /result/result".to_string())?;
|
||||
let current = inner
|
||||
.get("current_openhuman_dir")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| "config_get_data_paths missing current_openhuman_dir".to_string())?;
|
||||
let default = inner
|
||||
.get("default_openhuman_dir")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| "config_get_data_paths missing default_openhuman_dir".to_string())?;
|
||||
let marker = inner
|
||||
.get("active_workspace_marker_path")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| "config_get_data_paths missing active_workspace_marker_path".to_string())?;
|
||||
Ok(ResolvedDataPaths {
|
||||
current_openhuman_dir: std::path::PathBuf::from(current),
|
||||
default_openhuman_dir: std::path::PathBuf::from(default),
|
||||
active_workspace_marker_path: std::path::PathBuf::from(marker),
|
||||
})
|
||||
}
|
||||
|
||||
/// Remove a regular file if present. Missing → debug log + Ok.
|
||||
async fn remove_path_if_exists(path: &std::path::Path, label: &str) -> Result<(), String> {
|
||||
match tokio::fs::remove_file(path).await {
|
||||
Ok(()) => {
|
||||
log::info!(
|
||||
"[core] reset_local_data: removed {label} at {}",
|
||||
path.display()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
|
||||
log::debug!(
|
||||
"[core] reset_local_data: {label} already absent at {}",
|
||||
path.display()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => reset_local_data_delete_error(label, path, &e),
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove a directory tree if present. Missing → debug log + Ok.
|
||||
async fn remove_dir_if_exists(path: &std::path::Path, label: &str) -> Result<(), String> {
|
||||
match tokio::fs::remove_dir_all(path).await {
|
||||
Ok(()) => {
|
||||
log::info!(
|
||||
"[core] reset_local_data: removed {label} at {}",
|
||||
path.display()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
|
||||
log::debug!(
|
||||
"[core] reset_local_data: {label} already absent at {}",
|
||||
path.display()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(e) => reset_local_data_delete_error(label, path, &e),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "local_data_reset_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,84 @@
|
||||
use super::*;
|
||||
#[test]
|
||||
fn reset_local_data_windows_file_lock_error_codes_are_recognized() {
|
||||
assert!(is_windows_file_lock_raw_os_error(Some(32)));
|
||||
assert!(is_windows_file_lock_raw_os_error(Some(33)));
|
||||
assert!(!is_windows_file_lock_raw_os_error(Some(5)));
|
||||
assert!(!is_windows_file_lock_raw_os_error(None));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reset_local_data_delete_error_keeps_generic_message_for_other_errors() {
|
||||
let err = std::io::Error::from(std::io::ErrorKind::PermissionDenied);
|
||||
let result = reset_local_data_delete_error(
|
||||
"current openhuman dir",
|
||||
std::path::Path::new("/tmp/openhuman"),
|
||||
&err,
|
||||
);
|
||||
|
||||
let msg = result.expect_err("non-lock errors must still surface to the UI");
|
||||
assert!(msg.starts_with("Failed to remove current openhuman dir at /tmp/openhuman:"));
|
||||
assert!(!msg.contains("Close all OpenHuman windows and try again"));
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
#[test]
|
||||
fn reset_local_data_delete_error_swallows_lock_failure_when_path_disappeared() {
|
||||
// Race condition the reboot fallback now handles: the locked path
|
||||
// was gone by the time `schedule_path_for_reboot_deletion` ran its
|
||||
// `symlink_metadata` probe, so the reset goal is already met. The
|
||||
// helper must return `Ok(())` rather than surfacing a confusing
|
||||
// "couldn't remove (it's not there)" toast.
|
||||
let dir = tempfile::tempdir().expect("tempdir for reset error test");
|
||||
let missing = dir.path().join("definitely-not-there");
|
||||
|
||||
let err = std::io::Error::from_raw_os_error(32);
|
||||
let result = reset_local_data_delete_error("current openhuman dir", &missing, &err);
|
||||
|
||||
assert!(
|
||||
result.is_ok(),
|
||||
"expected NotFound + empty partial schedule to be swallowed as success, got {result:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
#[test]
|
||||
fn reset_local_data_delete_error_reports_reboot_schedule_counts() {
|
||||
// When the lock fallback can walk a real directory tree, the user
|
||||
// message should report how much has been queued so the support
|
||||
// log preserves "what was actually scheduled". Scheduling itself
|
||||
// may still fail at the MoveFileExW step in unprivileged test
|
||||
// processes (the registry key write requires administrator); the
|
||||
// fallback then carries a partial schedule that the error path
|
||||
// surfaces, so both branches must keep mentioning the lock cause
|
||||
// *and* expose either the queued counts or the schedule failure.
|
||||
let dir = tempfile::tempdir().expect("tempdir for reset error test");
|
||||
let target = dir.path().join("reset-mock");
|
||||
std::fs::create_dir_all(target.join("nested")).expect("mkdir nested");
|
||||
std::fs::write(target.join("a.txt"), b"x").expect("write a.txt");
|
||||
std::fs::write(target.join("nested").join("b.txt"), b"y").expect("write b.txt");
|
||||
|
||||
let err = std::io::Error::from_raw_os_error(32);
|
||||
let result = reset_local_data_delete_error("current openhuman dir", &target, &err);
|
||||
|
||||
// Path exists on disk, so the fallback must surface the outcome —
|
||||
// either an "all-queued" success-but-needs-reboot message (admin)
|
||||
// or one of the failure flavours (non-admin).
|
||||
let msg =
|
||||
result.expect_err("path exists, fallback must report queued counts or scheduling failure");
|
||||
let admin_path = msg.contains("queued for deletion the next time you restart Windows")
|
||||
&& msg.contains("2 files and 2 folders");
|
||||
let user_full_fail = msg.contains("scheduling deletion on next reboot also failed");
|
||||
let user_partial = msg.contains("queued for the next reboot before scheduling failed");
|
||||
assert!(
|
||||
admin_path || user_full_fail || user_partial,
|
||||
"expected reboot-scheduled, fully-failed, or partial-fail message, got: {msg}"
|
||||
);
|
||||
// Whatever branch we land on, the user must still be told the lock
|
||||
// is what blocked the immediate removal.
|
||||
assert!(
|
||||
msg.contains("locked by another OpenHuman window or process")
|
||||
|| msg.contains("another process is holding it open"),
|
||||
"missing lock cause: {msg}"
|
||||
);
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -62,76 +62,12 @@ use std::sync::Arc;
|
||||
/// detect those at the `ChatMessage` boundary (where `bound_cached_transcript_messages`
|
||||
/// operates) we have to peek inside the JSON. See TAURI-RUST-7 for the
|
||||
/// failure mode this guards against.
|
||||
fn assistant_message_has_tool_calls(msg: &ChatMessage) -> bool {
|
||||
if msg.role != "assistant" {
|
||||
return false;
|
||||
}
|
||||
let Ok(value) = serde_json::from_str::<serde_json::Value>(&msg.content) else {
|
||||
return false;
|
||||
};
|
||||
// CodeRabbit follow-up: only treat this as the native tool_calls envelope
|
||||
// when the full expected shape is present:
|
||||
// - top-level JSON object
|
||||
// - `content` key present (the envelope `dispatcher.rs` emits — see
|
||||
// `to_provider_messages`)
|
||||
// - non-empty `tool_calls` array whose every element carries an `id`
|
||||
// string, a `name` string, and an `arguments` field
|
||||
// This stops a legitimate assistant text reply that happens to contain
|
||||
// the literal string `tool_calls` from being misclassified and dropped at
|
||||
// the bound-cached-transcript boundary.
|
||||
let Some(obj) = value.as_object() else {
|
||||
return false;
|
||||
};
|
||||
if !obj.contains_key("content") {
|
||||
return false;
|
||||
}
|
||||
let Some(tool_calls) = obj.get("tool_calls").and_then(|tc| tc.as_array()) else {
|
||||
return false;
|
||||
};
|
||||
if tool_calls.is_empty() {
|
||||
return false;
|
||||
}
|
||||
tool_calls.iter().all(|tc| {
|
||||
tc.get("id").and_then(|v| v.as_str()).is_some()
|
||||
&& tc.get("name").and_then(|v| v.as_str()).is_some()
|
||||
&& tc.get("arguments").is_some()
|
||||
})
|
||||
}
|
||||
|
||||
/// Instruction appended (as a synthetic user turn) to the provider
|
||||
/// messages when a turn hits the tool-call iteration cap. Asks the model
|
||||
/// to wrap up with a resumable checkpoint instead of letting the turn die.
|
||||
/// Native tools are disabled for this call so the model produces prose,
|
||||
/// not yet another tool call. See bug-report-2026-05-26 A1.
|
||||
const MAX_ITER_CHECKPOINT_INSTRUCTION: &str = "\
|
||||
You have reached the maximum number of tool calls allowed for this single turn, so you cannot call any more tools right now. \
|
||||
Do not attempt another tool call. Instead, write a short progress checkpoint for the user with two clearly labelled parts:\n\
|
||||
1. **Done so far** — what you have accomplished in this turn, grounded in the tool results above.\n\
|
||||
2. **Next steps** — exactly what you plan to do next.\n\
|
||||
Write it so you can pick up seamlessly where you left off when the user replies. Be concise.";
|
||||
|
||||
/// Build a deterministic checkpoint summary from this turn's tool-call
|
||||
/// records. Used only as a safety net when the model-written checkpoint
|
||||
/// call fails or returns empty, so a capped turn can never be left without
|
||||
/// a well-formed assistant message — which is what silently wedged the
|
||||
/// thread before (bug-report-2026-05-26 A1).
|
||||
fn build_deterministic_checkpoint(records: &[ToolCallRecord], max_iterations: usize) -> String {
|
||||
let mut out = format!(
|
||||
"I reached the tool-call limit for this turn ({max_iterations} steps), so I paused here.\n\n**Done so far:**\n"
|
||||
);
|
||||
if records.is_empty() {
|
||||
out.push_str("- (no tools completed yet)\n");
|
||||
} else {
|
||||
for r in records {
|
||||
let status = if r.success { "ok" } else { "failed" };
|
||||
out.push_str(&format!("- `{}` — {}\n", r.name, status));
|
||||
}
|
||||
}
|
||||
out.push_str(
|
||||
"\n**Next steps:** I'll continue from here — just reply (e.g. \"continue\") and I'll pick up where I left off.",
|
||||
);
|
||||
out
|
||||
}
|
||||
#[path = "turn_checkpoint.rs"]
|
||||
mod turn_checkpoint;
|
||||
use turn_checkpoint::{
|
||||
assistant_message_has_tool_calls, build_deterministic_checkpoint,
|
||||
MAX_ITER_CHECKPOINT_INSTRUCTION,
|
||||
};
|
||||
|
||||
impl Agent {
|
||||
/// Executes a single interaction "turn" with the agent.
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
use crate::openhuman::agent::hooks::ToolCallRecord;
|
||||
use crate::openhuman::inference::provider::ChatMessage;
|
||||
|
||||
pub(super) fn assistant_message_has_tool_calls(msg: &ChatMessage) -> bool {
|
||||
if msg.role != "assistant" {
|
||||
return false;
|
||||
}
|
||||
let Ok(value) = serde_json::from_str::<serde_json::Value>(&msg.content) else {
|
||||
return false;
|
||||
};
|
||||
// CodeRabbit follow-up: only treat this as the native tool_calls envelope
|
||||
// when the full expected shape is present:
|
||||
// - top-level JSON object
|
||||
// - `content` key present (the envelope `dispatcher.rs` emits — see
|
||||
// `to_provider_messages`)
|
||||
// - non-empty `tool_calls` array whose every element carries an `id`
|
||||
// string, a `name` string, and an `arguments` field
|
||||
// This stops a legitimate assistant text reply that happens to contain
|
||||
// the literal string `tool_calls` from being misclassified and dropped at
|
||||
// the bound-cached-transcript boundary.
|
||||
let Some(obj) = value.as_object() else {
|
||||
return false;
|
||||
};
|
||||
if !obj.contains_key("content") {
|
||||
return false;
|
||||
}
|
||||
let Some(tool_calls) = obj.get("tool_calls").and_then(|tc| tc.as_array()) else {
|
||||
return false;
|
||||
};
|
||||
if tool_calls.is_empty() {
|
||||
return false;
|
||||
}
|
||||
tool_calls.iter().all(|tc| {
|
||||
tc.get("id").and_then(|v| v.as_str()).is_some()
|
||||
&& tc.get("name").and_then(|v| v.as_str()).is_some()
|
||||
&& tc.get("arguments").is_some()
|
||||
})
|
||||
}
|
||||
|
||||
/// Instruction appended (as a synthetic user turn) to the provider
|
||||
/// messages when a turn hits the tool-call iteration cap. Asks the model
|
||||
/// to wrap up with a resumable checkpoint instead of letting the turn die.
|
||||
/// Native tools are disabled for this call so the model produces prose,
|
||||
/// not yet another tool call. See bug-report-2026-05-26 A1.
|
||||
pub(super) const MAX_ITER_CHECKPOINT_INSTRUCTION: &str = "\
|
||||
You have reached the maximum number of tool calls allowed for this single turn, so you cannot call any more tools right now. \
|
||||
Do not attempt another tool call. Instead, write a short progress checkpoint for the user with two clearly labelled parts:\n\
|
||||
1. **Done so far** — what you have accomplished in this turn, grounded in the tool results above.\n\
|
||||
2. **Next steps** — exactly what you plan to do next.\n\
|
||||
Write it so you can pick up seamlessly where you left off when the user replies. Be concise.";
|
||||
|
||||
/// Build a deterministic checkpoint summary from this turn's tool-call
|
||||
/// records. Used only as a safety net when the model-written checkpoint
|
||||
/// call fails or returns empty, so a capped turn can never be left without
|
||||
/// a well-formed assistant message — which is what silently wedged the
|
||||
/// thread before (bug-report-2026-05-26 A1).
|
||||
pub(super) fn build_deterministic_checkpoint(
|
||||
records: &[ToolCallRecord],
|
||||
max_iterations: usize,
|
||||
) -> String {
|
||||
let mut out = format!(
|
||||
"I reached the tool-call limit for this turn ({max_iterations} steps), so I paused here.\n\n**Done so far:**\n"
|
||||
);
|
||||
if records.is_empty() {
|
||||
out.push_str("- (no tools completed yet)\n");
|
||||
} else {
|
||||
for r in records {
|
||||
let status = if r.success { "ok" } else { "failed" };
|
||||
out.push_str(&format!("- `{}` — {}\n", r.name, status));
|
||||
}
|
||||
}
|
||||
out.push_str(
|
||||
"\n**Next steps:** I'll continue from here — just reply (e.g. \"continue\") and I'll pick up where I left off.",
|
||||
);
|
||||
out
|
||||
}
|
||||
@@ -1,6 +1,5 @@
|
||||
use async_trait::async_trait;
|
||||
use once_cell::sync::Lazy;
|
||||
use regex::Regex;
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Map, Value};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
@@ -211,18 +210,6 @@ static IN_FLIGHT: Lazy<Mutex<HashMap<String, InFlightEntry>>> =
|
||||
#[cfg(test)]
|
||||
static TEST_FORCED_RUN_CHAT_TASK_ERROR: Lazy<Mutex<Option<String>>> =
|
||||
Lazy::new(|| Mutex::new(None));
|
||||
static BUDGET_ERROR_NORMALIZE_RE: Lazy<Regex> =
|
||||
Lazy::new(|| Regex::new(r"[-_\s]+").expect("budget normalize regex"));
|
||||
static BUDGET_ERROR_PATTERNS: Lazy<Vec<Regex>> = Lazy::new(|| {
|
||||
vec![
|
||||
Regex::new(r"budget.*exceed").expect("budget exceeded regex"),
|
||||
Regex::new(r"top up").expect("top up regex"),
|
||||
Regex::new(r"add.*credits").expect("add credits regex"),
|
||||
Regex::new(r"out of credits").expect("out of credits regex"),
|
||||
Regex::new(r"no remaining credits").expect("no remaining credits regex"),
|
||||
]
|
||||
});
|
||||
|
||||
/// Key for the per-thread runtime maps (`THREAD_SESSIONS`, `IN_FLIGHT`).
|
||||
///
|
||||
/// Keyed by `thread_id` ALONE — the stable, persistent identity of a
|
||||
@@ -244,523 +231,19 @@ fn event_session_id_for(client_id: &str, thread_id: &str) -> String {
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn is_inference_budget_exceeded_error(message: &str) -> bool {
|
||||
let normalized = BUDGET_ERROR_NORMALIZE_RE
|
||||
.replace_all(&message.trim().to_ascii_lowercase(), " ")
|
||||
.into_owned();
|
||||
BUDGET_ERROR_PATTERNS
|
||||
.iter()
|
||||
.any(|pattern| pattern.is_match(&normalized))
|
||||
}
|
||||
|
||||
fn inference_budget_exceeded_user_message() -> &'static str {
|
||||
"I don't have any budget available right now. Please top up your credits or choose a plan to continue."
|
||||
}
|
||||
|
||||
fn generic_inference_error_user_message() -> &'static str {
|
||||
"Something went wrong. Please try again.\nThis error has been reported. You can also report it on Discord.\n<openhuman-link path=\"community/discord\">Report on Discord</openhuman-link>"
|
||||
}
|
||||
|
||||
/// Pull the structured provider error message out of a raw error string.
|
||||
///
|
||||
/// Provider error chains from OpenAI/Anthropic/OpenRouter/etc. arrive looking
|
||||
/// like `custom_openai API error (404 Not Found): {"error":{"message":"...","type":"..."}}`.
|
||||
/// We extract the `error.message` value so the UI can show the *real* reason
|
||||
/// — e.g. "Project ... does not have access to model `gpt-5.5`" — instead of
|
||||
/// a generic apology.
|
||||
///
|
||||
/// Returns `None` for transport-level failures (DNS, TLS, connect refused)
|
||||
/// where there is no provider body to quote — those have no actionable
|
||||
/// detail and the raw error text can leak internal infrastructure URLs,
|
||||
/// which the chat surface deliberately does not expose to end users.
|
||||
fn extract_provider_error_detail(err: &str) -> Option<String> {
|
||||
const MAX_DETAIL_CHARS: usize = 300;
|
||||
|
||||
// Find the first `"message"` JSON field anywhere in the error chain.
|
||||
let key = "\"message\"";
|
||||
let idx = err.find(key)?;
|
||||
let after_key = &err[idx + key.len()..];
|
||||
// Skip whitespace and the colon to the opening quote of the value.
|
||||
let after_colon = after_key.trim_start_matches(|c: char| c != '"');
|
||||
let stripped = after_colon.strip_prefix('"')?;
|
||||
|
||||
// Manual unescape — handle `\"` and `\\` only; everything else passes
|
||||
// through. Sufficient for OpenAI/Anthropic/etc. error bodies.
|
||||
let mut out = String::new();
|
||||
let mut chars = stripped.chars();
|
||||
while let Some(c) = chars.next() {
|
||||
match c {
|
||||
'"' => {
|
||||
let trimmed = out.trim();
|
||||
if trimmed.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let sanitized = crate::openhuman::inference::provider::sanitize_api_error(trimmed);
|
||||
return Some(crate::openhuman::util::truncate_with_ellipsis(
|
||||
&sanitized,
|
||||
MAX_DETAIL_CHARS,
|
||||
));
|
||||
}
|
||||
'\\' => {
|
||||
if let Some(esc) = chars.next() {
|
||||
match esc {
|
||||
'"' => out.push('"'),
|
||||
'\\' => out.push('\\'),
|
||||
'n' => out.push('\n'),
|
||||
't' => out.push('\t'),
|
||||
other => out.push(other),
|
||||
}
|
||||
}
|
||||
}
|
||||
other => out.push(other),
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
/// Append the upstream provider detail to a user-facing message, if a useful
|
||||
/// one can be extracted. Keeps the friendly summary first and the verbatim
|
||||
/// provider reason below as a quotable block.
|
||||
fn with_provider_detail(summary: &str, err: &str) -> String {
|
||||
match extract_provider_error_detail(err) {
|
||||
Some(detail) => format!("{summary}\n\n> {detail}"),
|
||||
None => summary.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Structured chat-error envelope produced by [`classify_inference_error`].
|
||||
///
|
||||
/// Carries the typed metadata the frontend needs to render a recovery UI
|
||||
/// (retry-after countdown, retry button, fallback CTA) without having to
|
||||
/// regex the human-readable `message`. Issue #2606.
|
||||
///
|
||||
/// `error_type` and `message` preserve the wire shape PR #2371 established
|
||||
/// — existing FE handlers that read those fields keep working. The new
|
||||
/// fields are additive and `Option`-typed where the value isn't always
|
||||
/// known at the classifier layer.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct ClassifiedError {
|
||||
/// Stable token: `rate_limited`, `action_budget_exceeded`,
|
||||
/// `max_iterations`, `timeout`, `auth_error`, `budget_exhausted`,
|
||||
/// `provider_error`, `context_overflow`, `model_unavailable`,
|
||||
/// `inference`.
|
||||
pub(crate) error_type: &'static str,
|
||||
/// User-facing copy (already includes provider detail block and the
|
||||
/// retry-after countdown sentence when available).
|
||||
pub(crate) message: String,
|
||||
/// Where the limit originated. One of:
|
||||
/// - `"provider"` — upstream LLM provider 429 / rate limit
|
||||
/// - `"openhuman_budget"` — local SecurityPolicy per-hour action cap
|
||||
/// - `"agent_loop"` — agent ran out of tool iterations
|
||||
/// - `"openhuman_billing"` — OpenHuman credit/quota exhaustion
|
||||
/// - `"transport"` — network / DNS / TLS / timeout
|
||||
/// - `"config"` — auth, model, context, generic
|
||||
pub(crate) source: &'static str,
|
||||
/// Can the user retry the same prompt in the same thread? `false` for
|
||||
/// non-retryable business 429s, auth failures, model_unavailable,
|
||||
/// context_overflow, and OpenHuman billing exhaustion.
|
||||
pub(crate) retryable: bool,
|
||||
/// Milliseconds the upstream asked us to wait. Surfaced verbatim from
|
||||
/// `Retry-After:` / `retry_after:` headers when present; `None` when
|
||||
/// the upstream didn't supply one OR the error class doesn't have a
|
||||
/// concept of retry-after (auth, config, etc.).
|
||||
pub(crate) retry_after_ms: Option<u64>,
|
||||
/// Provider name extracted from the leading
|
||||
/// `"<provider> API error (...)"` envelope emitted by
|
||||
/// `inference::provider::ops::api_error`. `None` for non-provider
|
||||
/// errors (OpenHuman budget cap, agent loop) and for transport
|
||||
/// failures that don't carry an identifiable provider prefix.
|
||||
pub(crate) provider: Option<String>,
|
||||
/// `Some(false)` once the reliable-provider chain has exhausted every
|
||||
/// configured `model_fallbacks` entry (the aggregate "All
|
||||
/// providers/models failed" branch). `None` means the classifier
|
||||
/// can't tell from the error string alone — the FE should treat it
|
||||
/// as "unknown, don't promise a fallback".
|
||||
pub(crate) fallback_available: Option<bool>,
|
||||
}
|
||||
|
||||
/// Best-effort extraction of the provider name from an error string.
|
||||
///
|
||||
/// `inference::provider::ops::api_error` formats upstream failures as
|
||||
/// `"<provider> API error (<status>): <body>"`, e.g.
|
||||
/// `"openrouter API error (429 Too Many Requests): ..."`. We pull the
|
||||
/// leading word and lowercase it so the wire value is stable across
|
||||
/// providers' own capitalisation.
|
||||
///
|
||||
/// Returns `None` when:
|
||||
/// - The error string doesn't carry the `" API error"` infix.
|
||||
/// - The candidate word contains characters that wouldn't appear in a
|
||||
/// provider name (slashes, colons, etc. — guards against transport
|
||||
/// error prefixes that happen to be followed by " API error").
|
||||
fn extract_provider_name(err: &str) -> Option<String> {
|
||||
const INFIX: &str = " API error";
|
||||
let idx = err.find(INFIX)?;
|
||||
let prefix = err[..idx].trim_end();
|
||||
let candidate = prefix
|
||||
.rsplit_once(char::is_whitespace)
|
||||
.map_or(prefix, |(_, last)| last);
|
||||
if candidate.is_empty()
|
||||
|| !candidate
|
||||
.chars()
|
||||
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
|
||||
{
|
||||
return None;
|
||||
}
|
||||
Some(candidate.to_ascii_lowercase())
|
||||
}
|
||||
|
||||
/// Detect the reliable-provider aggregate that fires once every
|
||||
/// configured `model_fallbacks` entry has been tried.
|
||||
///
|
||||
/// `reliable.rs::format_failure_aggregate` always opens with
|
||||
/// `"All providers/models failed. Attempts:"`. When that marker is
|
||||
/// present the FE should NOT offer a fallback retry — there is none
|
||||
/// left to try.
|
||||
fn is_fallback_chain_exhausted(err: &str) -> bool {
|
||||
err.contains("All providers/models failed")
|
||||
}
|
||||
|
||||
/// Extract a Retry-After / retry_after seconds hint from a free-form
|
||||
/// error string. Mirrors the typed [`crate::openhuman::inference::
|
||||
/// provider::reliable::parse_retry_after_ms`] helper but operates on
|
||||
/// the already-flattened `String` that reaches the channel-classifier
|
||||
/// layer.
|
||||
///
|
||||
/// Returns `Some(n)` when a non-negative integer or fractional value
|
||||
/// follows one of the canonical headers; fractional values are
|
||||
/// rounded up so the user is never told to retry sooner than the
|
||||
/// upstream actually allows.
|
||||
fn parse_retry_after_secs_from_str(err: &str) -> Option<u64> {
|
||||
// Normalise quoted JSON-key wrappers ("retry_after": 30) by
|
||||
// stripping double quotes before scanning for prefixes
|
||||
// (CodeRabbit review on #2371). A serialised provider body like
|
||||
// `{"retry_after": 30}` would otherwise miss every prefix and
|
||||
// the user would lose the retry hint the provider supplied.
|
||||
let normalized = err.to_ascii_lowercase().replace('"', "");
|
||||
for prefix in &[
|
||||
"retry-after:",
|
||||
"retry_after:",
|
||||
"retry-after ",
|
||||
"retry_after ",
|
||||
] {
|
||||
if let Some(pos) = normalized.find(prefix) {
|
||||
let after = &normalized[pos + prefix.len()..];
|
||||
let num_str: String = after
|
||||
.trim()
|
||||
.chars()
|
||||
.take_while(|c| c.is_ascii_digit() || *c == '.')
|
||||
.collect();
|
||||
if let Ok(secs) = num_str.parse::<f64>() {
|
||||
if secs.is_finite() && secs >= 0.0 {
|
||||
return Some(secs.ceil() as u64);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Format the retry-after hint as a short user-friendly suffix
|
||||
/// (`" Try again in 30 seconds."`). Returns an empty string when no
|
||||
/// hint is available so callers can `format!("{summary}{hint}")`
|
||||
/// without branching on `Option`.
|
||||
fn retry_after_hint(secs: Option<u64>) -> String {
|
||||
match secs {
|
||||
Some(0) => " You can retry immediately.".to_string(),
|
||||
Some(1) => " Try again in 1 second.".to_string(),
|
||||
Some(n) if n < 90 => format!(" Try again in {n} seconds."),
|
||||
Some(n) => {
|
||||
// Round UP — never tell the user to retry sooner than
|
||||
// the upstream actually allows. 90–119s used to render
|
||||
// as "about 1 minutes" both because of integer flooring
|
||||
// and missing singular/plural handling (CodeRabbit
|
||||
// review on #2371).
|
||||
let mins = (n / 60) + u64::from(n % 60 != 0);
|
||||
let unit = if mins == 1 { "minute" } else { "minutes" };
|
||||
format!(" Try again in about {mins} {unit}.")
|
||||
}
|
||||
None => String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Detect the SecurityPolicy global hourly action-budget signal
|
||||
/// emitted by the built-in tools (`web_fetch`, `curl`, `http_request`,
|
||||
/// `polymarket`, `composio`, etc.) — see `src/openhuman/security/
|
||||
/// policy.rs::SecurityPolicy::is_rate_limited`.
|
||||
///
|
||||
/// We match the canonical English strings those tools emit. This is
|
||||
/// load-bearing for issue #2364: before this check ran, any string
|
||||
/// containing "rate limit" was misclassified as a provider 429 and
|
||||
/// the user saw the generic "You're being rate-limited" copy, which
|
||||
/// hides that the cap is OpenHuman's own per-hour safety budget,
|
||||
/// not the upstream LLM provider.
|
||||
fn is_action_budget_exhausted(err_lower: &str) -> bool {
|
||||
err_lower.contains("rate limit exceeded: action budget exhausted")
|
||||
|| err_lower.contains("rate limit exceeded: too many actions in the last hour")
|
||||
|| err_lower.contains("action blocked: rate limit exceeded")
|
||||
}
|
||||
|
||||
fn classify_inference_error(err: &str) -> ClassifiedError {
|
||||
let lower = err.to_lowercase();
|
||||
let provider = extract_provider_name(err);
|
||||
let fallback_available = if is_fallback_chain_exhausted(err) {
|
||||
Some(false)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Order matters: the SecurityPolicy hourly cap and the
|
||||
// agent-loop max-iterations error both surface as strings that
|
||||
// contain "rate limit" / "iteration", so they MUST be checked
|
||||
// before the generic provider-429 branch — otherwise users see
|
||||
// a confusing "your AI provider is rate-limiting you" message
|
||||
// for limits OpenHuman itself enforced (issue #2364).
|
||||
if is_action_budget_exhausted(&lower) {
|
||||
ClassifiedError {
|
||||
error_type: "action_budget_exceeded",
|
||||
message: with_provider_detail(
|
||||
"You've hit OpenHuman's per-hour action budget — this is a local safety cap, \
|
||||
not your AI provider. The window decays gradually; you can keep chatting in \
|
||||
this thread and tool-heavy steps will resume as the budget refills.",
|
||||
err,
|
||||
),
|
||||
source: "openhuman_budget",
|
||||
// The window decays gradually so the same thread CAN recover
|
||||
// — we just can't predict the exact wait.
|
||||
retryable: true,
|
||||
retry_after_ms: None,
|
||||
// OpenHuman's own cap — provider name (if any was in the
|
||||
// surrounding error chain) is irrelevant; the limit isn't
|
||||
// from a provider.
|
||||
provider: None,
|
||||
fallback_available: None,
|
||||
}
|
||||
} else if crate::openhuman::agent::error::is_max_iterations_error(err) {
|
||||
ClassifiedError {
|
||||
error_type: "max_iterations",
|
||||
message: with_provider_detail(
|
||||
"The agent ran the maximum number of tool steps for one turn without \
|
||||
finishing. This usually means a tool kept failing (often a rate limit on a \
|
||||
web fetch). You can retry the same question in this thread once the \
|
||||
underlying limit clears.",
|
||||
err,
|
||||
),
|
||||
source: "agent_loop",
|
||||
retryable: true,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available: None,
|
||||
}
|
||||
} else if lower.contains("rate limit") || lower.contains("429") {
|
||||
let retry_secs = parse_retry_after_secs_from_str(err);
|
||||
// Non-retryable business 429s ("plan does not include", balance
|
||||
// exhausted, known provider business codes like Z.AI 1311/1113)
|
||||
// also surface here — mark them non-retryable so the FE can hide
|
||||
// the "Retry" button and route the user to settings/billing.
|
||||
let non_retryable = is_non_retryable_rate_limit_text(&lower);
|
||||
let summary = if non_retryable {
|
||||
"Your AI provider is rejecting requests for billing or plan reasons \
|
||||
(out of credits, plan limit, or unavailable model). Retrying won't \
|
||||
help — open Settings to top up, upgrade your plan, or pick a \
|
||||
different model."
|
||||
.to_string()
|
||||
} else {
|
||||
format!(
|
||||
"Your AI provider is rate-limiting requests. This is a transient upstream \
|
||||
limit, not a thread-level block — you can retry in this thread.{}",
|
||||
retry_after_hint(retry_secs)
|
||||
)
|
||||
};
|
||||
ClassifiedError {
|
||||
error_type: "rate_limited",
|
||||
message: with_provider_detail(summary.as_str(), err),
|
||||
source: "provider",
|
||||
retryable: !non_retryable,
|
||||
retry_after_ms: retry_secs.map(|s| s.saturating_mul(1000)),
|
||||
provider,
|
||||
fallback_available,
|
||||
}
|
||||
} else if lower.contains("timeout") || lower.contains("timed out") {
|
||||
ClassifiedError {
|
||||
error_type: "timeout",
|
||||
message: with_provider_detail(
|
||||
"The request timed out. Please check your connection and try again.",
|
||||
err,
|
||||
),
|
||||
source: "transport",
|
||||
retryable: true,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available,
|
||||
}
|
||||
} else if lower.contains("401") || lower.contains("unauthorized") || lower.contains("api key") {
|
||||
ClassifiedError {
|
||||
error_type: "auth_error",
|
||||
message: with_provider_detail(
|
||||
"There's an authentication issue with the AI provider. Please check your API key in settings.",
|
||||
err,
|
||||
),
|
||||
source: "config",
|
||||
retryable: false,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available: None,
|
||||
}
|
||||
} else if lower.contains("402")
|
||||
|| lower.contains("payment required")
|
||||
|| lower.contains("insufficient balance")
|
||||
{
|
||||
// `openhuman_billing` means OpenHuman's own credit/quota system —
|
||||
// a 402 carrying the "openhuman" envelope (or no envelope at all,
|
||||
// since OpenHuman's backend is the only origin without one in
|
||||
// practice). When the 402 comes from an upstream provider envelope
|
||||
// (`<provider> API error (402)`), the limit belongs to that
|
||||
// provider, not OpenHuman billing, so tag the source as `provider`.
|
||||
let source: &'static str = match provider.as_deref() {
|
||||
Some("openhuman") | None => "openhuman_billing",
|
||||
Some(_) => "provider",
|
||||
};
|
||||
ClassifiedError {
|
||||
error_type: "budget_exhausted",
|
||||
message: with_provider_detail("Insufficient credits. Please top up to continue.", err),
|
||||
source,
|
||||
retryable: false,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available: None,
|
||||
}
|
||||
} else if lower.contains("500")
|
||||
|| lower.contains("internal server")
|
||||
|| lower.contains("service unavailable")
|
||||
|| lower.contains("503")
|
||||
{
|
||||
ClassifiedError {
|
||||
error_type: "provider_error",
|
||||
message: with_provider_detail(
|
||||
"The AI provider is temporarily unavailable. Please try again later.",
|
||||
err,
|
||||
),
|
||||
source: "provider",
|
||||
retryable: true,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available,
|
||||
}
|
||||
} else if lower.contains("context")
|
||||
&& (lower.contains("length")
|
||||
|| lower.contains("limit")
|
||||
|| lower.contains("exceed")
|
||||
|| lower.contains("token"))
|
||||
{
|
||||
ClassifiedError {
|
||||
error_type: "context_overflow",
|
||||
message: with_provider_detail(
|
||||
"The conversation is too long. Please start a new chat.",
|
||||
err,
|
||||
),
|
||||
source: "config",
|
||||
retryable: false,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available: None,
|
||||
}
|
||||
} else if crate::openhuman::inference::provider::is_provider_config_rejection_message(err) {
|
||||
// #2079 / #2076 / #2202: an OpenHuman abstract tier alias leaked to
|
||||
// a custom provider, a stale model pin, or a model-specific
|
||||
// temperature constraint. Checked BEFORE the generic
|
||||
// model-unavailable arm so config-rejection bodies that also
|
||||
// contain "model"/"does not exist"/"does not have access" get the
|
||||
// specific "Settings → LLM" remediation instead of the generic
|
||||
// copy. Shared predicate keeps this in lockstep with the
|
||||
// Sentry-demotion classifier.
|
||||
ClassifiedError {
|
||||
error_type: "model_unavailable",
|
||||
message: with_provider_detail(
|
||||
"Your AI provider rejected the request's model or temperature setting. \
|
||||
Check your model and routing in Settings → LLM.",
|
||||
err,
|
||||
),
|
||||
source: "config",
|
||||
retryable: false,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available: None,
|
||||
}
|
||||
} else if lower.contains("model")
|
||||
&& (lower.contains("not found")
|
||||
|| lower.contains("unavailable")
|
||||
|| lower.contains("does not exist")
|
||||
|| lower.contains("does not have access"))
|
||||
{
|
||||
ClassifiedError {
|
||||
error_type: "model_unavailable",
|
||||
message: with_provider_detail(
|
||||
"The selected model isn't available on your provider. Check your model settings.",
|
||||
err,
|
||||
),
|
||||
source: "config",
|
||||
retryable: false,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available: None,
|
||||
}
|
||||
} else {
|
||||
ClassifiedError {
|
||||
error_type: "inference",
|
||||
message: with_provider_detail(generic_inference_error_user_message(), err),
|
||||
source: "provider",
|
||||
retryable: true,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// String-flat mirror of
|
||||
/// [`crate::openhuman::inference::provider::reliable::is_non_retryable_rate_limit`].
|
||||
///
|
||||
/// The reliable provider already classifies 429s into retryable vs
|
||||
/// non-retryable based on business-quota markers ("plan does not
|
||||
/// include", "insufficient balance", Z.AI codes 1311/1113, …) — but
|
||||
/// that typed `anyhow::Error` is collapsed to a `String` at the
|
||||
/// native-bus boundary before reaching this layer. We re-detect the
|
||||
/// same markers in the flattened string so the FE knows whether to
|
||||
/// offer a "Retry" button.
|
||||
///
|
||||
/// Caller passes the already-lowercased error string to avoid double
|
||||
/// allocation.
|
||||
fn is_non_retryable_rate_limit_text(lower: &str) -> bool {
|
||||
const BUSINESS_HINTS: &[&str] = &[
|
||||
"plan does not include",
|
||||
"doesn't include",
|
||||
"not include",
|
||||
"insufficient balance",
|
||||
"insufficient_balance",
|
||||
"insufficient quota",
|
||||
"insufficient_quota",
|
||||
"quota exhausted",
|
||||
"out of credits",
|
||||
"no available package",
|
||||
"package not active",
|
||||
"purchase package",
|
||||
"model not available for your plan",
|
||||
];
|
||||
if BUSINESS_HINTS.iter().any(|hint| lower.contains(hint)) {
|
||||
return true;
|
||||
}
|
||||
// Known provider business codes observed for 429 where retry is
|
||||
// futile (mirrors reliable.rs). Scan integer-like tokens.
|
||||
for token in lower.split(|c: char| !c.is_ascii_digit()) {
|
||||
if let Ok(code) = token.parse::<u16>() {
|
||||
if matches!(code, 1113 | 1311) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
#[path = "web_errors.rs"]
|
||||
mod web_errors;
|
||||
pub(crate) use web_errors::{
|
||||
classify_inference_error, inference_budget_exceeded_user_message,
|
||||
is_inference_budget_exceeded_error,
|
||||
};
|
||||
#[cfg(test)]
|
||||
#[allow(unused_imports)]
|
||||
pub(crate) use web_errors::{
|
||||
extract_provider_error_detail, extract_provider_name, generic_inference_error_user_message,
|
||||
is_action_budget_exhausted, is_fallback_chain_exhausted, is_non_retryable_rate_limit_text,
|
||||
parse_retry_after_secs_from_str, retry_after_hint, with_provider_detail, ClassifiedError,
|
||||
};
|
||||
|
||||
fn prompt_guard_user_message(action: PromptEnforcementAction) -> &'static str {
|
||||
match action {
|
||||
|
||||
@@ -0,0 +1,532 @@
|
||||
use once_cell::sync::Lazy;
|
||||
use regex::Regex;
|
||||
|
||||
static BUDGET_ERROR_NORMALIZE_RE: Lazy<Regex> =
|
||||
Lazy::new(|| Regex::new(r"[-_\s]+").expect("budget normalize regex"));
|
||||
static BUDGET_ERROR_PATTERNS: Lazy<Vec<Regex>> = Lazy::new(|| {
|
||||
vec![
|
||||
Regex::new(r"budget.*exceed").expect("budget exceeded regex"),
|
||||
Regex::new(r"top up").expect("top up regex"),
|
||||
Regex::new(r"add.*credits").expect("add credits regex"),
|
||||
Regex::new(r"out of credits").expect("out of credits regex"),
|
||||
Regex::new(r"no remaining credits").expect("no remaining credits regex"),
|
||||
]
|
||||
});
|
||||
|
||||
pub(crate) fn is_inference_budget_exceeded_error(message: &str) -> bool {
|
||||
let normalized = BUDGET_ERROR_NORMALIZE_RE
|
||||
.replace_all(&message.trim().to_ascii_lowercase(), " ")
|
||||
.into_owned();
|
||||
BUDGET_ERROR_PATTERNS
|
||||
.iter()
|
||||
.any(|pattern| pattern.is_match(&normalized))
|
||||
}
|
||||
|
||||
pub(crate) fn inference_budget_exceeded_user_message() -> &'static str {
|
||||
"I don't have any budget available right now. Please top up your credits or choose a plan to continue."
|
||||
}
|
||||
|
||||
pub(crate) fn generic_inference_error_user_message() -> &'static str {
|
||||
"Something went wrong. Please try again.\nThis error has been reported. You can also report it on Discord.\n<openhuman-link path=\"community/discord\">Report on Discord</openhuman-link>"
|
||||
}
|
||||
|
||||
/// Pull the structured provider error message out of a raw error string.
|
||||
///
|
||||
/// Provider error chains from OpenAI/Anthropic/OpenRouter/etc. arrive looking
|
||||
/// like `custom_openai API error (404 Not Found): {"error":{"message":"...","type":"..."}}`.
|
||||
/// We extract the `error.message` value so the UI can show the *real* reason
|
||||
/// — e.g. "Project ... does not have access to model `gpt-5.5`" — instead of
|
||||
/// a generic apology.
|
||||
///
|
||||
/// Returns `None` for transport-level failures (DNS, TLS, connect refused)
|
||||
/// where there is no provider body to quote — those have no actionable
|
||||
/// detail and the raw error text can leak internal infrastructure URLs,
|
||||
/// which the chat surface deliberately does not expose to end users.
|
||||
pub(crate) fn extract_provider_error_detail(err: &str) -> Option<String> {
|
||||
const MAX_DETAIL_CHARS: usize = 300;
|
||||
|
||||
// Find the first `"message"` JSON field anywhere in the error chain.
|
||||
let key = "\"message\"";
|
||||
let idx = err.find(key)?;
|
||||
let after_key = &err[idx + key.len()..];
|
||||
// Skip whitespace and the colon to the opening quote of the value.
|
||||
let after_colon = after_key.trim_start_matches(|c: char| c != '"');
|
||||
let stripped = after_colon.strip_prefix('"')?;
|
||||
|
||||
// Manual unescape — handle `\"` and `\\` only; everything else passes
|
||||
// through. Sufficient for OpenAI/Anthropic/etc. error bodies.
|
||||
let mut out = String::new();
|
||||
let mut chars = stripped.chars();
|
||||
while let Some(c) = chars.next() {
|
||||
match c {
|
||||
'"' => {
|
||||
let trimmed = out.trim();
|
||||
if trimmed.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let sanitized = crate::openhuman::inference::provider::sanitize_api_error(trimmed);
|
||||
return Some(crate::openhuman::util::truncate_with_ellipsis(
|
||||
&sanitized,
|
||||
MAX_DETAIL_CHARS,
|
||||
));
|
||||
}
|
||||
'\\' => {
|
||||
if let Some(esc) = chars.next() {
|
||||
match esc {
|
||||
'"' => out.push('"'),
|
||||
'\\' => out.push('\\'),
|
||||
'n' => out.push('\n'),
|
||||
't' => out.push('\t'),
|
||||
other => out.push(other),
|
||||
}
|
||||
}
|
||||
}
|
||||
other => out.push(other),
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
/// Append the upstream provider detail to a user-facing message, if a useful
|
||||
/// one can be extracted. Keeps the friendly summary first and the verbatim
|
||||
/// provider reason below as a quotable block.
|
||||
pub(crate) fn with_provider_detail(summary: &str, err: &str) -> String {
|
||||
match extract_provider_error_detail(err) {
|
||||
Some(detail) => format!("{summary}\n\n> {detail}"),
|
||||
None => summary.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Structured chat-error envelope produced by [`classify_inference_error`].
|
||||
///
|
||||
/// Carries the typed metadata the frontend needs to render a recovery UI
|
||||
/// (retry-after countdown, retry button, fallback CTA) without having to
|
||||
/// regex the human-readable `message`. Issue #2606.
|
||||
///
|
||||
/// `error_type` and `message` preserve the wire shape PR #2371 established
|
||||
/// — existing FE handlers that read those fields keep working. The new
|
||||
/// fields are additive and `Option`-typed where the value isn't always
|
||||
/// known at the classifier layer.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct ClassifiedError {
|
||||
/// Stable token: `rate_limited`, `action_budget_exceeded`,
|
||||
/// `max_iterations`, `timeout`, `auth_error`, `budget_exhausted`,
|
||||
/// `provider_error`, `context_overflow`, `model_unavailable`,
|
||||
/// `inference`.
|
||||
pub(crate) error_type: &'static str,
|
||||
/// User-facing copy (already includes provider detail block and the
|
||||
/// retry-after countdown sentence when available).
|
||||
pub(crate) message: String,
|
||||
/// Where the limit originated. One of:
|
||||
/// - `"provider"` — upstream LLM provider 429 / rate limit
|
||||
/// - `"openhuman_budget"` — local SecurityPolicy per-hour action cap
|
||||
/// - `"agent_loop"` — agent ran out of tool iterations
|
||||
/// - `"openhuman_billing"` — OpenHuman credit/quota exhaustion
|
||||
/// - `"transport"` — network / DNS / TLS / timeout
|
||||
/// - `"config"` — auth, model, context, generic
|
||||
pub(crate) source: &'static str,
|
||||
/// Can the user retry the same prompt in the same thread? `false` for
|
||||
/// non-retryable business 429s, auth failures, model_unavailable,
|
||||
/// context_overflow, and OpenHuman billing exhaustion.
|
||||
pub(crate) retryable: bool,
|
||||
/// Milliseconds the upstream asked us to wait. Surfaced verbatim from
|
||||
/// `Retry-After:` / `retry_after:` headers when present; `None` when
|
||||
/// the upstream didn't supply one OR the error class doesn't have a
|
||||
/// concept of retry-after (auth, config, etc.).
|
||||
pub(crate) retry_after_ms: Option<u64>,
|
||||
/// Provider name extracted from the leading
|
||||
/// `"<provider> API error (...)"` envelope emitted by
|
||||
/// `inference::provider::ops::api_error`. `None` for non-provider
|
||||
/// errors (OpenHuman budget cap, agent loop) and for transport
|
||||
/// failures that don't carry an identifiable provider prefix.
|
||||
pub(crate) provider: Option<String>,
|
||||
/// `Some(false)` once the reliable-provider chain has exhausted every
|
||||
/// configured `model_fallbacks` entry (the aggregate "All
|
||||
/// providers/models failed" branch). `None` means the classifier
|
||||
/// can't tell from the error string alone — the FE should treat it
|
||||
/// as "unknown, don't promise a fallback".
|
||||
pub(crate) fallback_available: Option<bool>,
|
||||
}
|
||||
|
||||
/// Best-effort extraction of the provider name from an error string.
|
||||
///
|
||||
/// `inference::provider::ops::api_error` formats upstream failures as
|
||||
/// `"<provider> API error (<status>): <body>"`, e.g.
|
||||
/// `"openrouter API error (429 Too Many Requests): ..."`. We pull the
|
||||
/// leading word and lowercase it so the wire value is stable across
|
||||
/// providers' own capitalisation.
|
||||
///
|
||||
/// Returns `None` when:
|
||||
/// - The error string doesn't carry the `" API error"` infix.
|
||||
/// - The candidate word contains characters that wouldn't appear in a
|
||||
/// provider name (slashes, colons, etc. — guards against transport
|
||||
/// error prefixes that happen to be followed by " API error").
|
||||
pub(crate) fn extract_provider_name(err: &str) -> Option<String> {
|
||||
const INFIX: &str = " API error";
|
||||
let idx = err.find(INFIX)?;
|
||||
let prefix = err[..idx].trim_end();
|
||||
let candidate = prefix
|
||||
.rsplit_once(char::is_whitespace)
|
||||
.map_or(prefix, |(_, last)| last);
|
||||
if candidate.is_empty()
|
||||
|| !candidate
|
||||
.chars()
|
||||
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
|
||||
{
|
||||
return None;
|
||||
}
|
||||
Some(candidate.to_ascii_lowercase())
|
||||
}
|
||||
|
||||
/// Detect the reliable-provider aggregate that fires once every
|
||||
/// configured `model_fallbacks` entry has been tried.
|
||||
///
|
||||
/// `reliable.rs::format_failure_aggregate` always opens with
|
||||
/// `"All providers/models failed. Attempts:"`. When that marker is
|
||||
/// present the FE should NOT offer a fallback retry — there is none
|
||||
/// left to try.
|
||||
pub(crate) fn is_fallback_chain_exhausted(err: &str) -> bool {
|
||||
err.contains("All providers/models failed")
|
||||
}
|
||||
|
||||
/// Extract a Retry-After / retry_after seconds hint from a free-form
|
||||
/// error string. Mirrors the typed [`crate::openhuman::inference::
|
||||
/// provider::reliable::parse_retry_after_ms`] helper but operates on
|
||||
/// the already-flattened `String` that reaches the channel-classifier
|
||||
/// layer.
|
||||
///
|
||||
/// Returns `Some(n)` when a non-negative integer or fractional value
|
||||
/// follows one of the canonical headers; fractional values are
|
||||
/// rounded up so the user is never told to retry sooner than the
|
||||
/// upstream actually allows.
|
||||
pub(crate) fn parse_retry_after_secs_from_str(err: &str) -> Option<u64> {
|
||||
// Normalise quoted JSON-key wrappers ("retry_after": 30) by
|
||||
// stripping double quotes before scanning for prefixes
|
||||
// (CodeRabbit review on #2371). A serialised provider body like
|
||||
// `{"retry_after": 30}` would otherwise miss every prefix and
|
||||
// the user would lose the retry hint the provider supplied.
|
||||
let normalized = err.to_ascii_lowercase().replace('"', "");
|
||||
for prefix in &[
|
||||
"retry-after:",
|
||||
"retry_after:",
|
||||
"retry-after ",
|
||||
"retry_after ",
|
||||
] {
|
||||
if let Some(pos) = normalized.find(prefix) {
|
||||
let after = &normalized[pos + prefix.len()..];
|
||||
let num_str: String = after
|
||||
.trim()
|
||||
.chars()
|
||||
.take_while(|c| c.is_ascii_digit() || *c == '.')
|
||||
.collect();
|
||||
if let Ok(secs) = num_str.parse::<f64>() {
|
||||
if secs.is_finite() && secs >= 0.0 {
|
||||
return Some(secs.ceil() as u64);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Format the retry-after hint as a short user-friendly suffix
|
||||
/// (`" Try again in 30 seconds."`). Returns an empty string when no
|
||||
/// hint is available so callers can `format!("{summary}{hint}")`
|
||||
/// without branching on `Option`.
|
||||
pub(crate) fn retry_after_hint(secs: Option<u64>) -> String {
|
||||
match secs {
|
||||
Some(0) => " You can retry immediately.".to_string(),
|
||||
Some(1) => " Try again in 1 second.".to_string(),
|
||||
Some(n) if n < 90 => format!(" Try again in {n} seconds."),
|
||||
Some(n) => {
|
||||
// Round UP — never tell the user to retry sooner than
|
||||
// the upstream actually allows. 90–119s used to render
|
||||
// as "about 1 minutes" both because of integer flooring
|
||||
// and missing singular/plural handling (CodeRabbit
|
||||
// review on #2371).
|
||||
let mins = (n / 60) + u64::from(n % 60 != 0);
|
||||
let unit = if mins == 1 { "minute" } else { "minutes" };
|
||||
format!(" Try again in about {mins} {unit}.")
|
||||
}
|
||||
None => String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Detect the SecurityPolicy global hourly action-budget signal
|
||||
/// emitted by the built-in tools (`web_fetch`, `curl`, `http_request`,
|
||||
/// `polymarket`, `composio`, etc.) — see `src/openhuman/security/
|
||||
/// policy.rs::SecurityPolicy::is_rate_limited`.
|
||||
///
|
||||
/// We match the canonical English strings those tools emit. This is
|
||||
/// load-bearing for issue #2364: before this check ran, any string
|
||||
/// containing "rate limit" was misclassified as a provider 429 and
|
||||
/// the user saw the generic "You're being rate-limited" copy, which
|
||||
/// hides that the cap is OpenHuman's own per-hour safety budget,
|
||||
/// not the upstream LLM provider.
|
||||
pub(crate) fn is_action_budget_exhausted(err_lower: &str) -> bool {
|
||||
err_lower.contains("rate limit exceeded: action budget exhausted")
|
||||
|| err_lower.contains("rate limit exceeded: too many actions in the last hour")
|
||||
|| err_lower.contains("action blocked: rate limit exceeded")
|
||||
}
|
||||
|
||||
pub(crate) fn classify_inference_error(err: &str) -> ClassifiedError {
|
||||
let lower = err.to_lowercase();
|
||||
let provider = extract_provider_name(err);
|
||||
let fallback_available = if is_fallback_chain_exhausted(err) {
|
||||
Some(false)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Order matters: the SecurityPolicy hourly cap and the
|
||||
// agent-loop max-iterations error both surface as strings that
|
||||
// contain "rate limit" / "iteration", so they MUST be checked
|
||||
// before the generic provider-429 branch — otherwise users see
|
||||
// a confusing "your AI provider is rate-limiting you" message
|
||||
// for limits OpenHuman itself enforced (issue #2364).
|
||||
if is_action_budget_exhausted(&lower) {
|
||||
ClassifiedError {
|
||||
error_type: "action_budget_exceeded",
|
||||
message: with_provider_detail(
|
||||
"You've hit OpenHuman's per-hour action budget — this is a local safety cap, \
|
||||
not your AI provider. The window decays gradually; you can keep chatting in \
|
||||
this thread and tool-heavy steps will resume as the budget refills.",
|
||||
err,
|
||||
),
|
||||
source: "openhuman_budget",
|
||||
// The window decays gradually so the same thread CAN recover
|
||||
// — we just can't predict the exact wait.
|
||||
retryable: true,
|
||||
retry_after_ms: None,
|
||||
// OpenHuman's own cap — provider name (if any was in the
|
||||
// surrounding error chain) is irrelevant; the limit isn't
|
||||
// from a provider.
|
||||
provider: None,
|
||||
fallback_available: None,
|
||||
}
|
||||
} else if crate::openhuman::agent::error::is_max_iterations_error(err) {
|
||||
ClassifiedError {
|
||||
error_type: "max_iterations",
|
||||
message: with_provider_detail(
|
||||
"The agent ran the maximum number of tool steps for one turn without \
|
||||
finishing. This usually means a tool kept failing (often a rate limit on a \
|
||||
web fetch). You can retry the same question in this thread once the \
|
||||
underlying limit clears.",
|
||||
err,
|
||||
),
|
||||
source: "agent_loop",
|
||||
retryable: true,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available: None,
|
||||
}
|
||||
} else if lower.contains("rate limit") || lower.contains("429") {
|
||||
let retry_secs = parse_retry_after_secs_from_str(err);
|
||||
// Non-retryable business 429s ("plan does not include", balance
|
||||
// exhausted, known provider business codes like Z.AI 1311/1113)
|
||||
// also surface here — mark them non-retryable so the FE can hide
|
||||
// the "Retry" button and route the user to settings/billing.
|
||||
let non_retryable = is_non_retryable_rate_limit_text(&lower);
|
||||
let summary = if non_retryable {
|
||||
"Your AI provider is rejecting requests for billing or plan reasons \
|
||||
(out of credits, plan limit, or unavailable model). Retrying won't \
|
||||
help — open Settings to top up, upgrade your plan, or pick a \
|
||||
different model."
|
||||
.to_string()
|
||||
} else {
|
||||
format!(
|
||||
"Your AI provider is rate-limiting requests. This is a transient upstream \
|
||||
limit, not a thread-level block — you can retry in this thread.{}",
|
||||
retry_after_hint(retry_secs)
|
||||
)
|
||||
};
|
||||
ClassifiedError {
|
||||
error_type: "rate_limited",
|
||||
message: with_provider_detail(summary.as_str(), err),
|
||||
source: "provider",
|
||||
retryable: !non_retryable,
|
||||
retry_after_ms: retry_secs.map(|s| s.saturating_mul(1000)),
|
||||
provider,
|
||||
fallback_available,
|
||||
}
|
||||
} else if lower.contains("timeout") || lower.contains("timed out") {
|
||||
ClassifiedError {
|
||||
error_type: "timeout",
|
||||
message: with_provider_detail(
|
||||
"The request timed out. Please check your connection and try again.",
|
||||
err,
|
||||
),
|
||||
source: "transport",
|
||||
retryable: true,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available,
|
||||
}
|
||||
} else if lower.contains("401") || lower.contains("unauthorized") || lower.contains("api key") {
|
||||
ClassifiedError {
|
||||
error_type: "auth_error",
|
||||
message: with_provider_detail(
|
||||
"There's an authentication issue with the AI provider. Please check your API key in settings.",
|
||||
err,
|
||||
),
|
||||
source: "config",
|
||||
retryable: false,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available: None,
|
||||
}
|
||||
} else if lower.contains("402")
|
||||
|| lower.contains("payment required")
|
||||
|| lower.contains("insufficient balance")
|
||||
{
|
||||
// `openhuman_billing` means OpenHuman's own credit/quota system —
|
||||
// a 402 carrying the "openhuman" envelope (or no envelope at all,
|
||||
// since OpenHuman's backend is the only origin without one in
|
||||
// practice). When the 402 comes from an upstream provider envelope
|
||||
// (`<provider> API error (402)`), the limit belongs to that
|
||||
// provider, not OpenHuman billing, so tag the source as `provider`.
|
||||
let source: &'static str = match provider.as_deref() {
|
||||
Some("openhuman") | None => "openhuman_billing",
|
||||
Some(_) => "provider",
|
||||
};
|
||||
ClassifiedError {
|
||||
error_type: "budget_exhausted",
|
||||
message: with_provider_detail("Insufficient credits. Please top up to continue.", err),
|
||||
source,
|
||||
retryable: false,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available: None,
|
||||
}
|
||||
} else if lower.contains("500")
|
||||
|| lower.contains("internal server")
|
||||
|| lower.contains("service unavailable")
|
||||
|| lower.contains("503")
|
||||
{
|
||||
ClassifiedError {
|
||||
error_type: "provider_error",
|
||||
message: with_provider_detail(
|
||||
"The AI provider is temporarily unavailable. Please try again later.",
|
||||
err,
|
||||
),
|
||||
source: "provider",
|
||||
retryable: true,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available,
|
||||
}
|
||||
} else if lower.contains("context")
|
||||
&& (lower.contains("length")
|
||||
|| lower.contains("limit")
|
||||
|| lower.contains("exceed")
|
||||
|| lower.contains("token"))
|
||||
{
|
||||
ClassifiedError {
|
||||
error_type: "context_overflow",
|
||||
message: with_provider_detail(
|
||||
"The conversation is too long. Please start a new chat.",
|
||||
err,
|
||||
),
|
||||
source: "config",
|
||||
retryable: false,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available: None,
|
||||
}
|
||||
} else if crate::openhuman::inference::provider::is_provider_config_rejection_message(err) {
|
||||
// #2079 / #2076 / #2202: an OpenHuman abstract tier alias leaked to
|
||||
// a custom provider, a stale model pin, or a model-specific
|
||||
// temperature constraint. Checked BEFORE the generic
|
||||
// model-unavailable arm so config-rejection bodies that also
|
||||
// contain "model"/"does not exist"/"does not have access" get the
|
||||
// specific "Settings → LLM" remediation instead of the generic
|
||||
// copy. Shared predicate keeps this in lockstep with the
|
||||
// Sentry-demotion classifier.
|
||||
ClassifiedError {
|
||||
error_type: "model_unavailable",
|
||||
message: with_provider_detail(
|
||||
"Your AI provider rejected the request's model or temperature setting. \
|
||||
Check your model and routing in Settings → LLM.",
|
||||
err,
|
||||
),
|
||||
source: "config",
|
||||
retryable: false,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available: None,
|
||||
}
|
||||
} else if lower.contains("model")
|
||||
&& (lower.contains("not found")
|
||||
|| lower.contains("unavailable")
|
||||
|| lower.contains("does not exist")
|
||||
|| lower.contains("does not have access"))
|
||||
{
|
||||
ClassifiedError {
|
||||
error_type: "model_unavailable",
|
||||
message: with_provider_detail(
|
||||
"The selected model isn't available on your provider. Check your model settings.",
|
||||
err,
|
||||
),
|
||||
source: "config",
|
||||
retryable: false,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available: None,
|
||||
}
|
||||
} else {
|
||||
ClassifiedError {
|
||||
error_type: "inference",
|
||||
message: with_provider_detail(generic_inference_error_user_message(), err),
|
||||
source: "provider",
|
||||
retryable: true,
|
||||
retry_after_ms: None,
|
||||
provider,
|
||||
fallback_available,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// String-flat mirror of
|
||||
/// [`crate::openhuman::inference::provider::reliable::is_non_retryable_rate_limit`].
|
||||
///
|
||||
/// The reliable provider already classifies 429s into retryable vs
|
||||
/// non-retryable based on business-quota markers ("plan does not
|
||||
/// include", "insufficient balance", Z.AI codes 1311/1113, …) — but
|
||||
/// that typed `anyhow::Error` is collapsed to a `String` at the
|
||||
/// native-bus boundary before reaching this layer. We re-detect the
|
||||
/// same markers in the flattened string so the FE knows whether to
|
||||
/// offer a "Retry" button.
|
||||
///
|
||||
/// Caller passes the already-lowercased error string to avoid double
|
||||
/// allocation.
|
||||
pub(crate) fn is_non_retryable_rate_limit_text(lower: &str) -> bool {
|
||||
const BUSINESS_HINTS: &[&str] = &[
|
||||
"plan does not include",
|
||||
"doesn't include",
|
||||
"not include",
|
||||
"insufficient balance",
|
||||
"insufficient_balance",
|
||||
"insufficient quota",
|
||||
"insufficient_quota",
|
||||
"quota exhausted",
|
||||
"out of credits",
|
||||
"no available package",
|
||||
"package not active",
|
||||
"purchase package",
|
||||
"model not available for your plan",
|
||||
];
|
||||
if BUSINESS_HINTS.iter().any(|hint| lower.contains(hint)) {
|
||||
return true;
|
||||
}
|
||||
// Known provider business codes observed for 429 where retry is
|
||||
// futile (mirrors reliable.rs). Scan integer-like tokens.
|
||||
for token in lower.split(|c: char| !c.is_ascii_digit()) {
|
||||
if let Ok(code) = token.parse::<u16>() {
|
||||
if matches!(code, 1113 | 1311) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
@@ -1189,165 +1189,5 @@ pub(crate) async fn run_message_dispatch_loop(
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn contains_any_hits_at_least_one_word() {
|
||||
assert!(contains_any("hello world", &["world"]));
|
||||
assert!(contains_any("hello world", &["not there", "world"]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn contains_any_returns_false_when_none_match() {
|
||||
assert!(!contains_any("hello world", &["nope"]));
|
||||
assert!(!contains_any("hello world", &[]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn starts_with_any_detects_leading_prefix() {
|
||||
assert!(starts_with_any("hello world", &["hello"]));
|
||||
assert!(starts_with_any("hey you", &["yo", "hey"]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn starts_with_any_returns_false_when_none_match() {
|
||||
assert!(!starts_with_any("bonjour", &["hello", "hey"]));
|
||||
assert!(!starts_with_any("x", &[]));
|
||||
}
|
||||
|
||||
// ── select_acknowledgment_reaction ────────────────────────────
|
||||
|
||||
fn is_in(emoji: &str, options: &[&str]) -> bool {
|
||||
options.contains(&emoji)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_gratitude_category() {
|
||||
for msg in ["thanks a lot", "Thank you", "THX friend", "I appreciate it"] {
|
||||
let r = select_acknowledgment_reaction(msg);
|
||||
assert!(is_in(r, &["❤️", "🙏"]), "`{msg}` → {r}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_celebration_category() {
|
||||
for msg in ["amazing job", "this is awesome", "incredible!!"] {
|
||||
let r = select_acknowledgment_reaction(msg);
|
||||
assert!(is_in(r, &["🔥", "🎉"]), "`{msg}` → {r}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_crypto_category() {
|
||||
for msg in ["BTC price today", "ETH pump", "gm on the defi timeline"] {
|
||||
let r = select_acknowledgment_reaction(msg);
|
||||
assert!(is_in(r, &["💯", "⚡"]), "`{msg}` → {r}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_technical_category() {
|
||||
for msg in ["deploy the api", "debug this code", "rust question"] {
|
||||
let r = select_acknowledgment_reaction(msg);
|
||||
assert!(is_in(r, &["👨💻", "🤓"]), "`{msg}` → {r}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_greeting_category() {
|
||||
for msg in ["hi there", "hello", "hey friend", "yo"] {
|
||||
let r = select_acknowledgment_reaction(msg);
|
||||
assert!(is_in(r, &["🤗", "😁"]), "`{msg}` → {r}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_question_category() {
|
||||
for msg in [
|
||||
"what is this?",
|
||||
"how does it work",
|
||||
"can you help",
|
||||
"is this correct",
|
||||
] {
|
||||
let r = select_acknowledgment_reaction(msg);
|
||||
assert!(is_in(r, &["🤔", "✍️"]), "`{msg}` → {r}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_default_category() {
|
||||
let r = select_acknowledgment_reaction("the task is running");
|
||||
assert!(is_in(r, &["👀", "✍️"]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_is_deterministic() {
|
||||
let a = select_acknowledgment_reaction("thanks");
|
||||
let b = select_acknowledgment_reaction("thanks");
|
||||
assert_eq!(a, b, "same input should always yield same reaction");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_handles_empty_input_without_panic() {
|
||||
// `content.chars().next()` is None on empty input — must not panic.
|
||||
let r = select_acknowledgment_reaction("");
|
||||
assert!(!r.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_handles_single_char() {
|
||||
let r = select_acknowledgment_reaction("?");
|
||||
// Single "?" falls into question category (contains '?').
|
||||
assert!(is_in(r, &["🤔", "✍️"]));
|
||||
}
|
||||
|
||||
// ── build_channel_context_block (#928) ───────────────────────
|
||||
|
||||
fn cm(channel: &str, reply_target: &str) -> traits::ChannelMessage {
|
||||
traits::ChannelMessage {
|
||||
channel: channel.into(),
|
||||
sender: "alice".into(),
|
||||
content: "hi".into(),
|
||||
id: "m1".into(),
|
||||
reply_target: reply_target.into(),
|
||||
thread_ts: None,
|
||||
timestamp: 0,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn channel_context_block_omitted_for_web_and_cli() {
|
||||
assert!(build_channel_context_block(&cm("web", "1")).is_empty());
|
||||
assert!(build_channel_context_block(&cm("cli", "1")).is_empty());
|
||||
assert!(build_channel_context_block(&cm("WEB", "1")).is_empty());
|
||||
assert!(build_channel_context_block(&cm("", "1")).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn channel_context_block_omitted_when_reply_target_missing() {
|
||||
assert!(build_channel_context_block(&cm("telegram", "")).is_empty());
|
||||
assert!(build_channel_context_block(&cm("telegram", " ")).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn channel_context_block_for_telegram_includes_routing_hint() {
|
||||
let block = build_channel_context_block(&cm("telegram", "123456"));
|
||||
assert!(block.contains("[Channel context]"));
|
||||
assert!(block.contains("\"telegram\""));
|
||||
assert!(block.contains("\"123456\""));
|
||||
// Hint must steer the model toward announce mode with the same channel/target.
|
||||
assert!(block.contains("announce"));
|
||||
assert!(block.contains("cron_add"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn channel_context_block_for_discord_and_slack_share_shape() {
|
||||
for ch in ["discord", "slack", "matrix"] {
|
||||
let block = build_channel_context_block(&cm(ch, "chan-42"));
|
||||
assert!(block.contains(ch), "missing channel name in `{ch}` block");
|
||||
assert!(block.contains("chan-42"));
|
||||
assert!(block.contains("announce"));
|
||||
}
|
||||
}
|
||||
}
|
||||
#[path = "dispatch_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,160 @@
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn contains_any_hits_at_least_one_word() {
|
||||
assert!(contains_any("hello world", &["world"]));
|
||||
assert!(contains_any("hello world", &["not there", "world"]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn contains_any_returns_false_when_none_match() {
|
||||
assert!(!contains_any("hello world", &["nope"]));
|
||||
assert!(!contains_any("hello world", &[]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn starts_with_any_detects_leading_prefix() {
|
||||
assert!(starts_with_any("hello world", &["hello"]));
|
||||
assert!(starts_with_any("hey you", &["yo", "hey"]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn starts_with_any_returns_false_when_none_match() {
|
||||
assert!(!starts_with_any("bonjour", &["hello", "hey"]));
|
||||
assert!(!starts_with_any("x", &[]));
|
||||
}
|
||||
|
||||
// ── select_acknowledgment_reaction ────────────────────────────
|
||||
|
||||
fn is_in(emoji: &str, options: &[&str]) -> bool {
|
||||
options.contains(&emoji)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_gratitude_category() {
|
||||
for msg in ["thanks a lot", "Thank you", "THX friend", "I appreciate it"] {
|
||||
let r = select_acknowledgment_reaction(msg);
|
||||
assert!(is_in(r, &["❤️", "🙏"]), "`{msg}` → {r}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_celebration_category() {
|
||||
for msg in ["amazing job", "this is awesome", "incredible!!"] {
|
||||
let r = select_acknowledgment_reaction(msg);
|
||||
assert!(is_in(r, &["🔥", "🎉"]), "`{msg}` → {r}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_crypto_category() {
|
||||
for msg in ["BTC price today", "ETH pump", "gm on the defi timeline"] {
|
||||
let r = select_acknowledgment_reaction(msg);
|
||||
assert!(is_in(r, &["💯", "⚡"]), "`{msg}` → {r}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_technical_category() {
|
||||
for msg in ["deploy the api", "debug this code", "rust question"] {
|
||||
let r = select_acknowledgment_reaction(msg);
|
||||
assert!(is_in(r, &["👨💻", "🤓"]), "`{msg}` → {r}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_greeting_category() {
|
||||
for msg in ["hi there", "hello", "hey friend", "yo"] {
|
||||
let r = select_acknowledgment_reaction(msg);
|
||||
assert!(is_in(r, &["🤗", "😁"]), "`{msg}` → {r}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_question_category() {
|
||||
for msg in [
|
||||
"what is this?",
|
||||
"how does it work",
|
||||
"can you help",
|
||||
"is this correct",
|
||||
] {
|
||||
let r = select_acknowledgment_reaction(msg);
|
||||
assert!(is_in(r, &["🤔", "✍️"]), "`{msg}` → {r}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_default_category() {
|
||||
let r = select_acknowledgment_reaction("the task is running");
|
||||
assert!(is_in(r, &["👀", "✍️"]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_is_deterministic() {
|
||||
let a = select_acknowledgment_reaction("thanks");
|
||||
let b = select_acknowledgment_reaction("thanks");
|
||||
assert_eq!(a, b, "same input should always yield same reaction");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_handles_empty_input_without_panic() {
|
||||
// `content.chars().next()` is None on empty input — must not panic.
|
||||
let r = select_acknowledgment_reaction("");
|
||||
assert!(!r.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ack_reaction_handles_single_char() {
|
||||
let r = select_acknowledgment_reaction("?");
|
||||
// Single "?" falls into question category (contains '?').
|
||||
assert!(is_in(r, &["🤔", "✍️"]));
|
||||
}
|
||||
|
||||
// ── build_channel_context_block (#928) ───────────────────────
|
||||
|
||||
fn cm(channel: &str, reply_target: &str) -> traits::ChannelMessage {
|
||||
traits::ChannelMessage {
|
||||
channel: channel.into(),
|
||||
sender: "alice".into(),
|
||||
content: "hi".into(),
|
||||
id: "m1".into(),
|
||||
reply_target: reply_target.into(),
|
||||
thread_ts: None,
|
||||
timestamp: 0,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn channel_context_block_omitted_for_web_and_cli() {
|
||||
assert!(build_channel_context_block(&cm("web", "1")).is_empty());
|
||||
assert!(build_channel_context_block(&cm("cli", "1")).is_empty());
|
||||
assert!(build_channel_context_block(&cm("WEB", "1")).is_empty());
|
||||
assert!(build_channel_context_block(&cm("", "1")).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn channel_context_block_omitted_when_reply_target_missing() {
|
||||
assert!(build_channel_context_block(&cm("telegram", "")).is_empty());
|
||||
assert!(build_channel_context_block(&cm("telegram", " ")).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn channel_context_block_for_telegram_includes_routing_hint() {
|
||||
let block = build_channel_context_block(&cm("telegram", "123456"));
|
||||
assert!(block.contains("[Channel context]"));
|
||||
assert!(block.contains("\"telegram\""));
|
||||
assert!(block.contains("\"123456\""));
|
||||
// Hint must steer the model toward announce mode with the same channel/target.
|
||||
assert!(block.contains("announce"));
|
||||
assert!(block.contains("cron_add"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn channel_context_block_for_discord_and_slack_share_shape() {
|
||||
for ch in ["discord", "slack", "matrix"] {
|
||||
let block = build_channel_context_block(&cm(ch, "chan-42"));
|
||||
assert!(block.contains(ch), "missing channel name in `{ch}` block");
|
||||
assert!(block.contains("chan-42"));
|
||||
assert!(block.contains("announce"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,834 @@
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::context::prompt::{
|
||||
ConnectedIntegration, ConnectedIntegrationTool, GatedIntegrationTool,
|
||||
};
|
||||
|
||||
use super::client::ComposioClient;
|
||||
use super::ops::should_forward_tags;
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::{LazyLock, RwLock};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
// ── Prompt integration discovery ────────────────────────────────────
|
||||
|
||||
/// Defensive TTL on the integrations cache.
|
||||
///
|
||||
/// Background: the primary invalidation path is the
|
||||
/// `ComposioConnectionCreated` → `wait_for_connection_active` bus flow
|
||||
/// (see [`super::bus::ComposioConnectionCreatedSubscriber`]), which
|
||||
/// polls the backend for up to 60 s after `composio_authorize` returns
|
||||
/// a `connectUrl`. On Windows the OAuth round-trip can exceed that
|
||||
/// window (Defender SmartScreen, slower browser launch, extra consent
|
||||
/// dialogs), so the invalidation call never fires and the chat
|
||||
/// runtime's cache stays frozen on the pre-connect snapshot even
|
||||
/// though the Settings UI polls `composio_list_connections` every 5 s
|
||||
/// and shows the user as "Connected".
|
||||
///
|
||||
/// The cross-platform defenses we layer on top:
|
||||
/// 1. [`composio_list_connections`] diff-invalidates the cache whenever
|
||||
/// the backend's active-toolkit set diverges from what's cached,
|
||||
/// so a running UI keeps the chat cache in sync within one poll
|
||||
/// interval.
|
||||
/// 2. This TTL caps worst-case staleness at 60 s regardless of
|
||||
/// whether the UI is open, the bus fires, or the user reconnected
|
||||
/// out-of-band.
|
||||
pub(crate) const CACHE_TTL: Duration = Duration::from_secs(60);
|
||||
|
||||
/// Cached entry: the integrations list plus the timestamp we wrote it.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct CachedIntegrations {
|
||||
pub(crate) entries: Vec<ConnectedIntegration>,
|
||||
pub(crate) cached_at: Instant,
|
||||
}
|
||||
|
||||
/// Process-wide cache for connected integrations, keyed by the config
|
||||
/// identity (the `config_path` string) so different user contexts don't
|
||||
/// collide. Each entry is populated on first fetch and returned on
|
||||
/// subsequent calls until explicitly invalidated or the TTL expires.
|
||||
pub(crate) static INTEGRATIONS_CACHE: LazyLock<RwLock<HashMap<String, CachedIntegrations>>> =
|
||||
LazyLock::new(|| RwLock::new(HashMap::new()));
|
||||
|
||||
/// Derive a stable cache key from a [`Config`]. We use the stringified
|
||||
/// `config_path` because it uniquely identifies a user context (it
|
||||
/// resolves to the per-user openhuman dir).
|
||||
pub(crate) fn cache_key(config: &Config) -> String {
|
||||
config.config_path.display().to_string()
|
||||
}
|
||||
|
||||
/// Clear cached connected integrations so the next call to
|
||||
/// [`fetch_connected_integrations`] hits the backend again.
|
||||
///
|
||||
/// Called by [`super::bus::ComposioConnectionCreatedSubscriber`] when a
|
||||
/// new OAuth connection completes, by [`composio_list_connections`]
|
||||
/// when it observes a divergence between the backend response and the
|
||||
/// cached snapshot, and from tests. Clears the entire map because the
|
||||
/// callers don't carry a config reference.
|
||||
pub fn invalidate_connected_integrations_cache() {
|
||||
if let Ok(mut guard) = INTEGRATIONS_CACHE.write() {
|
||||
let entries = guard.len();
|
||||
guard.clear();
|
||||
tracing::info!(
|
||||
cached_keys = entries,
|
||||
"[composio][integrations] cache invalidated"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Read-only snapshot of the currently cached connected integrations for
|
||||
/// the given config, or [`None`] when the cache is empty, expired, or
|
||||
/// the lock is held by a writer.
|
||||
///
|
||||
/// Designed for hot-path callers that want a cheap "what does the cache
|
||||
/// already say?" probe without triggering a backend fetch. The agent
|
||||
/// harness uses this on every turn to detect mid-session connection
|
||||
/// changes — it relies on the desktop UI's 5 s `composio_list_connections`
|
||||
/// poll (which calls into [`fetch_connected_integrations`] and
|
||||
/// repopulates this cache) plus the event-driven invalidation path to
|
||||
/// keep the cache current.
|
||||
///
|
||||
/// `try_read` (not `read`) so a writer in progress — e.g. the UI poll
|
||||
/// repopulating the cache — never blocks a turn. Worst case the agent
|
||||
/// sees `None` for one turn while the writer holds the lock; the next
|
||||
/// turn picks up the value naturally.
|
||||
///
|
||||
/// TTL is enforced defensively: entries older than [`CACHE_TTL`] are
|
||||
/// treated as missing even though they're still in the map (a stale
|
||||
/// entry would otherwise pin the agent to a frozen view if every
|
||||
/// invalidation path silently failed).
|
||||
pub fn cached_active_integrations(config: &Config) -> Option<Vec<ConnectedIntegration>> {
|
||||
let key = cache_key(config);
|
||||
let guard = match INTEGRATIONS_CACHE.try_read() {
|
||||
Ok(g) => g,
|
||||
Err(_) => {
|
||||
tracing::trace!(
|
||||
key = %key,
|
||||
"[composio][integrations_cache] cached_active_integrations:lock_contended"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let Some(cached) = guard.get(&key) else {
|
||||
tracing::trace!(
|
||||
key = %key,
|
||||
"[composio][integrations_cache] cached_active_integrations:miss"
|
||||
);
|
||||
return None;
|
||||
};
|
||||
let age = cached.cached_at.elapsed();
|
||||
if age > CACHE_TTL {
|
||||
tracing::trace!(
|
||||
key = %key,
|
||||
age_ms = age.as_millis() as u64,
|
||||
ttl_ms = CACHE_TTL.as_millis() as u64,
|
||||
"[composio][integrations_cache] cached_active_integrations:expired"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
tracing::trace!(
|
||||
key = %key,
|
||||
entries = cached.entries.len(),
|
||||
age_ms = age.as_millis() as u64,
|
||||
"[composio][integrations_cache] cached_active_integrations:hit"
|
||||
);
|
||||
Some(cached.entries.clone())
|
||||
}
|
||||
|
||||
/// Stable hash of the *routing-relevant* slice of a connected-integrations
|
||||
/// snapshot.
|
||||
///
|
||||
/// Two snapshots produce the same hash iff they would synthesise the
|
||||
/// same `delegate_<toolkit>` tool set in the orchestrator's
|
||||
/// function-calling schema. The hash is:
|
||||
///
|
||||
/// - **Order-independent** — callers don't need to sort the input.
|
||||
/// - **Description-insensitive** — Composio catalogue text edits don't
|
||||
/// trigger a refresh. The schema's tool-description field still
|
||||
/// picks up new text on the next *real* (membership-changing)
|
||||
/// refresh, so descriptions are never permanently stale.
|
||||
/// - **Process-local** — [`std::collections::hash_map::DefaultHasher`]
|
||||
/// is randomly seeded per process. Fine because we only compare
|
||||
/// hashes within one process lifetime.
|
||||
///
|
||||
/// Only `connected == true` entries contribute. Unconnected toolkits are
|
||||
/// stripped by [`super::super::tools::orchestrator_tools::collect_orchestrator_tools`]
|
||||
/// anyway, so churn among the unconnected set never changes the agent's
|
||||
/// surface and shouldn't trigger a refresh.
|
||||
pub fn connected_set_hash(integrations: &[ConnectedIntegration]) -> u64 {
|
||||
use std::collections::hash_map::DefaultHasher;
|
||||
use std::hash::{Hash, Hasher};
|
||||
|
||||
let mut slugs: Vec<&str> = integrations
|
||||
.iter()
|
||||
.filter(|i| i.connected)
|
||||
.map(|i| i.toolkit.as_str())
|
||||
.collect();
|
||||
slugs.sort();
|
||||
|
||||
let mut hasher = DefaultHasher::new();
|
||||
slugs.hash(&mut hasher);
|
||||
hasher.finish()
|
||||
}
|
||||
|
||||
/// Collect the set of toolkit slugs marked `connected` in a snapshot.
|
||||
///
|
||||
/// Exposed to [`sync_cache_with_connections`] so it can diff the live
|
||||
/// backend connection list against what the chat runtime currently
|
||||
/// believes is connected.
|
||||
fn connected_toolkit_set(integrations: &[ConnectedIntegration]) -> HashSet<String> {
|
||||
integrations
|
||||
.iter()
|
||||
.filter(|i| i.connected)
|
||||
.map(|i| i.toolkit.clone())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Reconcile the process-wide integrations cache with a fresh backend
|
||||
/// `list_connections` response.
|
||||
///
|
||||
/// Called from [`composio_list_connections`], which the desktop UI
|
||||
/// polls every 5 s (see `app/src/lib/composio/hooks.ts`). When the set
|
||||
/// of ACTIVE/CONNECTED toolkits in the response differs from what's in
|
||||
/// the cache, we invalidate so the chat runtime re-fetches on its next
|
||||
/// `fetch_connected_integrations` call. This keeps tool availability
|
||||
/// in chat in sync with the badge the user sees in Settings, even when
|
||||
/// the primary event-bus invalidation path misses (e.g. Windows OAuth
|
||||
/// flows that overrun the 60 s readiness poll).
|
||||
pub(crate) fn sync_cache_with_connections(connections: &[super::types::ComposioConnection]) {
|
||||
let live_active: HashSet<String> = connections
|
||||
.iter()
|
||||
.filter(|c| c.is_active())
|
||||
.map(|c| c.normalized_toolkit())
|
||||
.filter(|toolkit| !toolkit.is_empty())
|
||||
.collect();
|
||||
|
||||
// Read once to decide whether any cache entry is out of sync. We
|
||||
// clone out the keys + connected sets so we can release the read
|
||||
// lock before taking the write lock.
|
||||
let divergent_keys: Vec<(String, HashSet<String>, HashSet<String>)> = {
|
||||
let Ok(guard) = INTEGRATIONS_CACHE.read() else {
|
||||
return;
|
||||
};
|
||||
guard
|
||||
.iter()
|
||||
.filter_map(|(key, cached)| {
|
||||
let cached_set = connected_toolkit_set(&cached.entries);
|
||||
if cached_set != live_active {
|
||||
Some((key.clone(), cached_set, live_active.clone()))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
|
||||
if divergent_keys.is_empty() {
|
||||
tracing::debug!(
|
||||
live_connected = live_active.len(),
|
||||
"[composio][integrations] list_connections matches cache — no invalidation needed"
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
if let Ok(mut guard) = INTEGRATIONS_CACHE.write() {
|
||||
for (key, cached_set, live_set) in divergent_keys {
|
||||
// Diff logging — makes Windows-timing regressions easy to
|
||||
// catch in user-supplied debug dumps without leaking any
|
||||
// PII (toolkit slugs are public strings like "gmail").
|
||||
let added: Vec<&String> = live_set.difference(&cached_set).collect();
|
||||
let removed: Vec<&String> = cached_set.difference(&live_set).collect();
|
||||
tracing::info!(
|
||||
key = %key,
|
||||
?added,
|
||||
?removed,
|
||||
"[composio][integrations] cache diverges from backend — invalidating"
|
||||
);
|
||||
guard.remove(&key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Fetch the user's active Composio connections and their available
|
||||
/// tool actions, returning a prompt-ready summary.
|
||||
///
|
||||
/// This is the **single source of truth** for connected integration
|
||||
/// data injected into system prompts — both the agent turn loop and
|
||||
/// the debug dump CLI call this function.
|
||||
///
|
||||
/// Results are cached process-wide (keyed by config identity) and
|
||||
/// returned instantly on subsequent calls. The cache is invalidated
|
||||
/// when a new connection is created
|
||||
/// (via [`invalidate_connected_integrations_cache`]), when a UI
|
||||
/// `list_connections` poll observes a divergent live set, when
|
||||
/// [`CACHE_TTL`] expires, or on process restart.
|
||||
///
|
||||
/// Best-effort: returns an empty vec when the user isn't signed in,
|
||||
/// the backend is unreachable, or any step fails.
|
||||
pub async fn fetch_connected_integrations(config: &Config) -> Vec<ConnectedIntegration> {
|
||||
match fetch_connected_integrations_status(config).await {
|
||||
FetchConnectedIntegrationsStatus::Authoritative(v) => v,
|
||||
FetchConnectedIntegrationsStatus::Unavailable => Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Discriminated outcome from [`fetch_connected_integrations_status`].
|
||||
///
|
||||
/// Lets callers distinguish "the backend confirmed the user has zero
|
||||
/// active connections right now" from "we couldn't talk to the backend
|
||||
/// (no client, transient failure, …) and have no truth to report".
|
||||
///
|
||||
/// The legacy [`fetch_connected_integrations`] collapses both into an
|
||||
/// empty `Vec`, which is fine for prompt-building (they look the same)
|
||||
/// but dangerous for spawn-time allowlist gates — using empty as truth
|
||||
/// in the unavailable case would silently wipe the user's allowlist
|
||||
/// during a transient 5xx.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum FetchConnectedIntegrationsStatus {
|
||||
/// Backend was reachable. Vec may legitimately be empty (no
|
||||
/// allowlisted toolkits, or no active connections).
|
||||
Authoritative(Vec<ConnectedIntegration>),
|
||||
/// Backend wasn't reachable (no auth client, transient error). The
|
||||
/// caller should fall back to its prior snapshot rather than treat
|
||||
/// "no connections" as truth.
|
||||
Unavailable,
|
||||
}
|
||||
|
||||
/// Status-returning variant of [`fetch_connected_integrations`].
|
||||
///
|
||||
/// Same caching, same cache-invalidation semantics — only the return
|
||||
/// shape differs. Cache hits are by definition `Authoritative` because
|
||||
/// we only cache the `Some(...)` arm of `_uncached` (i.e. results the
|
||||
/// backend confirmed).
|
||||
pub async fn fetch_connected_integrations_status(
|
||||
config: &Config,
|
||||
) -> FetchConnectedIntegrationsStatus {
|
||||
let key = cache_key(config);
|
||||
|
||||
// Fast path: return cached result if fresh. Stale entries fall
|
||||
// through to the backend fetch below so the chat runtime can never
|
||||
// be more than `CACHE_TTL` behind a real-world change.
|
||||
if let Ok(guard) = INTEGRATIONS_CACHE.read() {
|
||||
if let Some(cached) = guard.get(&key) {
|
||||
let age = cached.cached_at.elapsed();
|
||||
if age < CACHE_TTL {
|
||||
tracing::debug!(
|
||||
count = cached.entries.len(),
|
||||
age_ms = age.as_millis() as u64,
|
||||
key = %key,
|
||||
"[composio][integrations] returning cached result"
|
||||
);
|
||||
return FetchConnectedIntegrationsStatus::Authoritative(cached.entries.clone());
|
||||
}
|
||||
tracing::info!(
|
||||
count = cached.entries.len(),
|
||||
age_ms = age.as_millis() as u64,
|
||||
ttl_ms = CACHE_TTL.as_millis() as u64,
|
||||
key = %key,
|
||||
"[composio][integrations] cache entry expired — refetching"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
match fetch_connected_integrations_uncached(config).await {
|
||||
Some(result) => {
|
||||
// Backend was reachable — cache the result (even if empty).
|
||||
if let Ok(mut guard) = INTEGRATIONS_CACHE.write() {
|
||||
guard.insert(
|
||||
key,
|
||||
CachedIntegrations {
|
||||
entries: result.clone(),
|
||||
cached_at: Instant::now(),
|
||||
},
|
||||
);
|
||||
}
|
||||
FetchConnectedIntegrationsStatus::Authoritative(result)
|
||||
}
|
||||
None => {
|
||||
// No auth / client unavailable — do NOT cache so a
|
||||
// subsequent call with a different config can retry.
|
||||
FetchConnectedIntegrationsStatus::Unavailable
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The actual backend fetch, called on cache miss.
|
||||
///
|
||||
/// Returns `Some(vec)` when the backend was reachable. The returned
|
||||
/// vector is the merged **integration overview** — every toolkit in
|
||||
/// the backend allowlist appears as one entry, with a `connected`
|
||||
/// flag indicating whether the user has an active OAuth connection.
|
||||
/// Connected entries also carry the per-action tool catalogue
|
||||
/// (fetched in a single batched call).
|
||||
///
|
||||
/// Returns `None` when we couldn't even build a client (no auth),
|
||||
/// signalling the caller should NOT cache this result.
|
||||
async fn fetch_connected_integrations_uncached(
|
||||
config: &Config,
|
||||
) -> Option<Vec<ConnectedIntegration>> {
|
||||
use super::client::{create_composio_client, direct_list_connections, ComposioClientKind};
|
||||
use super::providers::toolkit_description;
|
||||
|
||||
// Route via the mode-aware factory so the chat-agent's
|
||||
// "connected_integrations" view reflects the live tenant — backend
|
||||
// (tinyhumans) or direct (user's personal Composio). Prior to #1710
|
||||
// Wave 3 this path called `build_composio_client` directly, which
|
||||
// is backend-only — after a `composio.mode = "direct"` toggle the
|
||||
// cache kept replaying the tinyhumans-tenant connections back into
|
||||
// the integration overview (e.g. gmail / notion appearing as
|
||||
// connected in direct mode even when the user's direct tenant had
|
||||
// a different set of toolkits). Resolving per call closes the
|
||||
// loop: `ComposioConfigChangedSubscriber` invalidates the cache on
|
||||
// toggle and the next miss re-populates it from the live tenant.
|
||||
let kind = match create_composio_client(config) {
|
||||
Ok(kind) => kind,
|
||||
Err(e) => {
|
||||
tracing::debug!(
|
||||
error = %e,
|
||||
"[composio] fetch_connected_integrations: no client (not signed in?)"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
// Pull the allowlist + connections + tool catalogue. Backend mode
|
||||
// walks the tinyhumans tenant's curated allowlist via
|
||||
// `list_toolkits`; direct mode has no centralised allowlist (per
|
||||
// `ops::composio_list_toolkits`'s direct-mode branch) so the
|
||||
// user's set of active connections IS the universe of valid
|
||||
// toolkit arguments.
|
||||
//
|
||||
// On transient errors we return `None` instead of a degraded
|
||||
// `Some(Vec::new())` so `fetch_connected_integrations` does NOT
|
||||
// cache the failure. Caching an empty allowlist would hide every
|
||||
// integration from the orchestrator until the process restarts or
|
||||
// the cache is explicitly invalidated — a single 5xx during
|
||||
// startup would silently break delegation for the whole session.
|
||||
let (allowlisted_toolkits, connections, tools_by_toolkit): (
|
||||
Vec<String>,
|
||||
Vec<super::types::ComposioConnection>,
|
||||
Vec<super::types::ComposioToolSchema>,
|
||||
) = match &kind {
|
||||
ComposioClientKind::Backend(client) => {
|
||||
let allowlist: Vec<String> = match client.list_toolkits().await {
|
||||
Ok(resp) => resp
|
||||
.toolkits
|
||||
.into_iter()
|
||||
.map(|toolkit| toolkit.trim().to_ascii_lowercase())
|
||||
.filter(|toolkit| !toolkit.is_empty())
|
||||
.collect(),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"[composio] fetch_connected_integrations: list_toolkits (backend) failed: {e}"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
if allowlist.is_empty() {
|
||||
tracing::debug!(
|
||||
"[composio] fetch_connected_integrations: backend allowlist is empty"
|
||||
);
|
||||
return Some(Vec::new());
|
||||
}
|
||||
|
||||
let connections = match client.list_connections().await {
|
||||
Ok(resp) => resp.connections,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"[composio] fetch_connected_integrations: list_connections (backend) failed: {e}"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
// Tool catalogue scoped to the active subset only —
|
||||
// not-connected toolkits won't be invoked from a sub-agent.
|
||||
let connected_slugs_for_tools: Vec<String> = {
|
||||
let mut v: Vec<String> = connections
|
||||
.iter()
|
||||
.filter(|c| c.is_active())
|
||||
.map(|c| c.normalized_toolkit())
|
||||
.filter(|t| !t.is_empty())
|
||||
.collect();
|
||||
v.sort();
|
||||
v.dedup();
|
||||
v
|
||||
};
|
||||
let tools = if connected_slugs_for_tools.is_empty() {
|
||||
Vec::new()
|
||||
} else {
|
||||
match client
|
||||
.list_tools(Some(&connected_slugs_for_tools), None)
|
||||
.await
|
||||
{
|
||||
Ok(resp) => resp.tools,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"[composio] fetch_connected_integrations: list_tools (backend) failed: {e}"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
(allowlist, connections, tools)
|
||||
}
|
||||
ComposioClientKind::Direct(direct) => {
|
||||
// Direct mode: walk the user's personal Composio tenant
|
||||
// for *connection state* (active accounts on their key) —
|
||||
// there's no central allowlist in direct mode, so the
|
||||
// active set IS the allowlist.
|
||||
//
|
||||
// Tool *schemas* are tenant-agnostic — Composio's action
|
||||
// definitions (e.g. GMAIL_SEND_EMAIL parameter shape) are
|
||||
// identical regardless of which Composio tenant a user is
|
||||
// connected via. So we best-effort fetch schemas through
|
||||
// the backend client (curated list_tools) if a backend
|
||||
// session is available, even though connection routing
|
||||
// goes to the user's direct tenant. This preserves the
|
||||
// chat agent's "21 gmail actions available" view while
|
||||
// execution itself (via `ComposioActionTool` / Wave 1
|
||||
// factory) still routes to the user's tenant. Direct-only
|
||||
// users without a backend session get empty tools — that
|
||||
// matches `composio_list_tools`'s direct-mode policy and
|
||||
// the `subagent_runner` LazyToolkitResolver still resolves
|
||||
// tools lazily at delegation time.
|
||||
let connections = match direct_list_connections(direct).await {
|
||||
Ok(resp) => resp.connections,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"[composio] fetch_connected_integrations: list_connections (direct) failed: {e:#}"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let allowlist: Vec<String> = {
|
||||
let mut v: Vec<String> = connections
|
||||
.iter()
|
||||
.filter(|c| c.is_active())
|
||||
.map(|c| c.normalized_toolkit())
|
||||
.filter(|t| !t.is_empty())
|
||||
.collect();
|
||||
v.sort();
|
||||
v.dedup();
|
||||
v
|
||||
};
|
||||
if allowlist.is_empty() {
|
||||
tracing::info!(
|
||||
"[composio-direct] fetch_connected_integrations: direct tenant has no active connections; returning empty overview"
|
||||
);
|
||||
return Some(Vec::new());
|
||||
}
|
||||
tracing::debug!(
|
||||
connected = allowlist.len(),
|
||||
"[composio-direct] fetch_connected_integrations: using direct tenant's active set as allowlist (no central allowlist in direct mode)"
|
||||
);
|
||||
|
||||
// Best-effort: pull tool schemas via the backend client
|
||||
// (definitional source). Failure is non-fatal — we fall
|
||||
// back to empty tools and let lazy resolution handle it.
|
||||
let tools = match super::client::build_composio_client(config) {
|
||||
Some(backend_client) => {
|
||||
match backend_client.list_tools(Some(&allowlist), None).await {
|
||||
Ok(resp) => {
|
||||
tracing::debug!(
|
||||
count = resp.tools.len(),
|
||||
"[composio-direct] fetch_connected_integrations: pulled tool schemas from backend (tenant-agnostic definitional source)"
|
||||
);
|
||||
resp.tools
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::info!(
|
||||
"[composio-direct] fetch_connected_integrations: backend list_tools failed (will use lazy fallback at delegation time): {e:#}"
|
||||
);
|
||||
Vec::new()
|
||||
}
|
||||
}
|
||||
}
|
||||
None => {
|
||||
tracing::info!(
|
||||
"[composio-direct] fetch_connected_integrations: no backend session for schema fetch; lazy fallback at delegation time"
|
||||
);
|
||||
Vec::new()
|
||||
}
|
||||
};
|
||||
(allowlist, connections, tools)
|
||||
}
|
||||
};
|
||||
|
||||
// Active connection slugs (status filter mirrors the original logic).
|
||||
let connected_slugs: std::collections::HashSet<String> = connections
|
||||
.iter()
|
||||
.filter(|c| c.is_active())
|
||||
.map(|c| c.normalized_toolkit())
|
||||
.filter(|toolkit| !toolkit.is_empty())
|
||||
.collect();
|
||||
|
||||
// Most-informative *non-active* status per toolkit slug. Lets the
|
||||
// integrations_agent spawn-gate (#2365) emit a precise message
|
||||
// when a connection row exists but isn't usable yet (`INITIATED`
|
||||
// — OAuth still in progress) or any longer (`EXPIRED` / `FAILED`)
|
||||
// — instead of the legacy generic "available but not authorized".
|
||||
//
|
||||
// Status priority (UI-actionability):
|
||||
// 1. EXPIRED — reconnect path
|
||||
// 2. FAILED / ERROR — reconnect path
|
||||
// 3. INITIATED / INITIALIZING / PENDING — finish OAuth in browser
|
||||
// 4. anything else — passes through verbatim
|
||||
let non_active_status_by_slug: std::collections::HashMap<String, String> = {
|
||||
fn priority(status: &str) -> u8 {
|
||||
let s = status.trim().to_ascii_uppercase();
|
||||
match s.as_str() {
|
||||
"EXPIRED" => 4,
|
||||
"FAILED" | "ERROR" => 3,
|
||||
"INITIATED" | "INITIALIZING" | "PENDING" => 2,
|
||||
_ => 1,
|
||||
}
|
||||
}
|
||||
let mut map: std::collections::HashMap<String, (u8, String)> =
|
||||
std::collections::HashMap::new();
|
||||
for conn in connections.iter().filter(|c| !c.is_active()) {
|
||||
let slug = conn.normalized_toolkit();
|
||||
if slug.is_empty() {
|
||||
continue;
|
||||
}
|
||||
// Don't override an ACTIVE-slug — those carry no non-active
|
||||
// status from this map's perspective.
|
||||
if connected_slugs.contains(&slug) {
|
||||
continue;
|
||||
}
|
||||
let p = priority(&conn.status);
|
||||
map.entry(slug.clone())
|
||||
.and_modify(|cur| {
|
||||
if p > cur.0 {
|
||||
tracing::debug!(
|
||||
target: "composio",
|
||||
toolkit = %slug,
|
||||
previous_status = %cur.1,
|
||||
previous_priority = cur.0,
|
||||
new_status = %conn.status,
|
||||
new_priority = p,
|
||||
"[composio] non_active_status_by_slug: upgraded most-informative status"
|
||||
);
|
||||
*cur = (p, conn.status.clone());
|
||||
} else {
|
||||
tracing::trace!(
|
||||
target: "composio",
|
||||
toolkit = %slug,
|
||||
kept_status = %cur.1,
|
||||
kept_priority = cur.0,
|
||||
candidate_status = %conn.status,
|
||||
candidate_priority = p,
|
||||
"[composio] non_active_status_by_slug: kept higher-priority status"
|
||||
);
|
||||
}
|
||||
})
|
||||
.or_insert_with(|| {
|
||||
tracing::debug!(
|
||||
target: "composio",
|
||||
toolkit = %slug,
|
||||
status = %conn.status,
|
||||
priority = p,
|
||||
"[composio] non_active_status_by_slug: first non-active row"
|
||||
);
|
||||
(p, conn.status.clone())
|
||||
});
|
||||
}
|
||||
let final_map: std::collections::HashMap<String, String> =
|
||||
map.into_iter().map(|(k, (_, v))| (k, v)).collect();
|
||||
tracing::debug!(
|
||||
target: "composio",
|
||||
entries = final_map.len(),
|
||||
"[composio] non_active_status_by_slug: final map"
|
||||
);
|
||||
final_map
|
||||
};
|
||||
|
||||
// Deduplicate the allowlist so a backend that returns duplicates
|
||||
// doesn't produce dual entries downstream.
|
||||
let mut unique_toolkits: Vec<String> = allowlisted_toolkits.clone();
|
||||
unique_toolkits.sort();
|
||||
unique_toolkits.dedup();
|
||||
|
||||
// Build one entry per allowlisted toolkit. Connected entries
|
||||
// carry their action catalogue; not-connected entries carry an
|
||||
// empty `tools` vec.
|
||||
let mut integrations: Vec<ConnectedIntegration> = Vec::with_capacity(unique_toolkits.len());
|
||||
for slug in &unique_toolkits {
|
||||
let connected = connected_slugs.contains(slug);
|
||||
// Anchor the prefix with an underscore so slugs that share
|
||||
// a text prefix (e.g. `git` vs `github`) don't false-match
|
||||
// each other's actions. `GMAIL_SEND_EMAIL` matches `gmail_`,
|
||||
// not just `gmail`, so siblings stay in their own buckets.
|
||||
let action_prefix = format!("{}_", slug.to_uppercase());
|
||||
let (tools, gated_tools): (Vec<ConnectedIntegrationTool>, Vec<GatedIntegrationTool>) =
|
||||
if connected {
|
||||
// Apply the same curated-whitelist + user-scope filter the
|
||||
// meta-tool layer uses, so the integrations_agent prompt
|
||||
// only advertises actions the agent is actually allowed to
|
||||
// call. One pref load per toolkit (not per action).
|
||||
//
|
||||
// Actions that the catalog *does* know about but the user's
|
||||
// current scope pref denies are routed into `gated_tools` so
|
||||
// the agent can honestly answer "I have this capability but
|
||||
// it needs the {scope} toggle in Connections → {toolkit}".
|
||||
// The agent cannot flip the scope itself — that's a UI-only
|
||||
// action. Without this gated surface the LLM has no way to
|
||||
// know the gated action exists at all and will tell the user
|
||||
// "I don't support that" — technically correct about its
|
||||
// callable surface, but misleading about the toolkit.
|
||||
let pref = super::providers::load_user_scope_or_default(slug).await;
|
||||
let mut visible: Vec<ConnectedIntegrationTool> = Vec::new();
|
||||
let mut gated: Vec<GatedIntegrationTool> = Vec::new();
|
||||
for t in tools_by_toolkit
|
||||
.iter()
|
||||
.filter(|t| t.function.name.starts_with(&action_prefix))
|
||||
{
|
||||
if super::providers::is_action_visible_with_pref(&t.function.name, &pref) {
|
||||
visible.push(ConnectedIntegrationTool {
|
||||
name: t.function.name.clone(),
|
||||
description: t.function.description.clone().unwrap_or_default(),
|
||||
parameters: t.function.parameters.clone(),
|
||||
});
|
||||
} else if let Some(required_scope) =
|
||||
super::providers::curated_scope_for(&t.function.name)
|
||||
{
|
||||
// Only surface CURATED actions as `gated` — uncurated
|
||||
// tools (which fall through to `classify_unknown` and
|
||||
// happen to land outside the user's pref) are not
|
||||
// first-class user-facing capabilities, and listing
|
||||
// them would clutter the prompt with internal slugs.
|
||||
// Deliberately NO `parameters` field: the LLM should
|
||||
// not be able to construct a call envelope; it can
|
||||
// only describe + point at the unlock path.
|
||||
// Ship the unlock path as data — single path today
|
||||
// (the Connections UI toggle). The agent does NOT
|
||||
// have a tool to flip scopes; that capability was
|
||||
// removed because LLM-mediated scope elevation made
|
||||
// the safety contract depend on model behavior and
|
||||
// was a soft gate the model could route around. If
|
||||
// more unlock paths exist in future (per-action
|
||||
// approval modal, time-boxed elevation, etc.) they
|
||||
// land here and the prompt renderer picks them up.
|
||||
let scope_str = required_scope.as_str();
|
||||
let unlock_paths = vec![format!(
|
||||
"the user enables it themselves in \
|
||||
Connections → {slug} → {scope_str}"
|
||||
)];
|
||||
gated.push(GatedIntegrationTool {
|
||||
name: t.function.name.clone(),
|
||||
description: t.function.description.clone().unwrap_or_default(),
|
||||
required_scope: scope_str.to_string(),
|
||||
unlock_paths,
|
||||
});
|
||||
}
|
||||
}
|
||||
tracing::debug!(
|
||||
toolkit = %slug,
|
||||
visible = visible.len(),
|
||||
gated = gated.len(),
|
||||
"[composio][scopes] integrations prompt action set"
|
||||
);
|
||||
(visible, gated)
|
||||
} else {
|
||||
(Vec::new(), Vec::new())
|
||||
};
|
||||
|
||||
integrations.push(ConnectedIntegration {
|
||||
toolkit: slug.clone(),
|
||||
description: toolkit_description(slug).to_string(),
|
||||
tools,
|
||||
gated_tools,
|
||||
connected,
|
||||
non_active_status: if connected {
|
||||
None
|
||||
} else {
|
||||
non_active_status_by_slug.get(slug).cloned()
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
integrations.sort_by(|a, b| a.toolkit.cmp(&b.toolkit));
|
||||
|
||||
let connected_count = integrations.iter().filter(|i| i.connected).count();
|
||||
tracing::info!(
|
||||
total = integrations.len(),
|
||||
connected = connected_count,
|
||||
"[composio] fetch_connected_integrations: done"
|
||||
);
|
||||
for ci in &integrations {
|
||||
tracing::debug!(
|
||||
toolkit = %ci.toolkit,
|
||||
connected = ci.connected,
|
||||
non_active_status = ?ci.non_active_status,
|
||||
tool_count = ci.tools.len(),
|
||||
"[composio] integration overview"
|
||||
);
|
||||
}
|
||||
|
||||
Some(integrations)
|
||||
}
|
||||
|
||||
/// Just-in-time fetch of every available action for a single Composio
|
||||
/// toolkit, returned in the [`ConnectedIntegrationTool`] shape the
|
||||
/// `integrations_agent` spawn path expects.
|
||||
///
|
||||
/// Unlike [`fetch_connected_integrations`] (which bulk-fetches every
|
||||
/// connected toolkit's tools once per session and caches the result),
|
||||
/// this helper is uncached and scoped to a single toolkit — meant to
|
||||
/// be called at `integrations_agent` spawn time so the sub-agent's
|
||||
/// prompt always reflects the toolkit's current action catalogue.
|
||||
///
|
||||
/// The filter `starts_with("{TOOLKIT}_")` matches
|
||||
/// `fetch_connected_integrations_uncached`'s own namespacing rule so
|
||||
/// siblings like `github` / `git` don't leak into each other's buckets.
|
||||
///
|
||||
/// `tags` narrows the result by Composio action tag (OR semantics). Only
|
||||
/// honoured for the GitHub toolkit; passed through to `list_tools` so the
|
||||
/// backend can skip the repo-list force-include and return a focused set.
|
||||
///
|
||||
/// Returns an empty vec when the backend has no actions for the
|
||||
/// toolkit (valid steady state for a freshly-authorised integration
|
||||
/// whose catalogue hasn't been published yet). Returns `Err` only for
|
||||
/// transport / auth failures the caller should surface to the user.
|
||||
pub async fn fetch_toolkit_actions(
|
||||
client: &ComposioClient,
|
||||
toolkit: &str,
|
||||
tags: Option<&[String]>,
|
||||
) -> anyhow::Result<Vec<ConnectedIntegrationTool>> {
|
||||
let toolkit_slug = toolkit.trim();
|
||||
if toolkit_slug.is_empty() {
|
||||
anyhow::bail!("fetch_toolkit_actions: toolkit must not be empty");
|
||||
}
|
||||
let effective_tags = if should_forward_tags(Some(&[toolkit_slug.to_string()])) {
|
||||
tags
|
||||
} else {
|
||||
None
|
||||
};
|
||||
tracing::debug!(toolkit = %toolkit_slug, ?effective_tags, "[composio] fetch_toolkit_actions");
|
||||
let resp = client
|
||||
.list_tools(Some(&[toolkit_slug.to_string()]), effective_tags)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("list_tools failed for toolkit `{toolkit_slug}`: {e}"))?;
|
||||
let action_prefix = format!("{}_", toolkit_slug.to_uppercase());
|
||||
// Apply curated whitelist + user scope so spawn-time tool
|
||||
// discovery agrees with the bulk path and the meta-tool layer.
|
||||
let pref = super::providers::load_user_scope_or_default(toolkit_slug).await;
|
||||
let actions: Vec<ConnectedIntegrationTool> = resp
|
||||
.tools
|
||||
.into_iter()
|
||||
.filter(|t| t.function.name.starts_with(&action_prefix))
|
||||
.filter(|t| super::providers::is_action_visible_with_pref(&t.function.name, &pref))
|
||||
.map(|t| ConnectedIntegrationTool {
|
||||
name: t.function.name,
|
||||
description: t.function.description.unwrap_or_default(),
|
||||
parameters: t.function.parameters,
|
||||
})
|
||||
.collect();
|
||||
tracing::debug!(
|
||||
toolkit = %toolkit_slug,
|
||||
action_count = actions.len(),
|
||||
"[composio] fetch_toolkit_actions: done"
|
||||
);
|
||||
Ok(actions)
|
||||
}
|
||||
@@ -39,6 +39,7 @@ pub mod action_tool;
|
||||
pub mod auth_retry;
|
||||
pub mod bus;
|
||||
pub mod client;
|
||||
mod connected_integrations;
|
||||
pub mod error_mapping;
|
||||
pub mod execute_dispatch;
|
||||
pub mod execute_prepare;
|
||||
|
||||
+13
-828
@@ -1402,835 +1402,20 @@ fn parse_sync_reason(raw: Option<&str>) -> OpResult<SyncReason> {
|
||||
}
|
||||
}
|
||||
|
||||
// ── Prompt integration discovery ────────────────────────────────────
|
||||
|
||||
use crate::openhuman::context::prompt::{
|
||||
ConnectedIntegration, ConnectedIntegrationTool, GatedIntegrationTool,
|
||||
#[cfg(test)]
|
||||
pub(crate) use super::connected_integrations::cache_key;
|
||||
use super::connected_integrations::sync_cache_with_connections;
|
||||
pub use super::connected_integrations::{
|
||||
cached_active_integrations, connected_set_hash, fetch_connected_integrations,
|
||||
fetch_connected_integrations_status, fetch_toolkit_actions,
|
||||
invalidate_connected_integrations_cache, FetchConnectedIntegrationsStatus,
|
||||
};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::{LazyLock, RwLock};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
/// Defensive TTL on the integrations cache.
|
||||
///
|
||||
/// Background: the primary invalidation path is the
|
||||
/// `ComposioConnectionCreated` → `wait_for_connection_active` bus flow
|
||||
/// (see [`super::bus::ComposioConnectionCreatedSubscriber`]), which
|
||||
/// polls the backend for up to 60 s after `composio_authorize` returns
|
||||
/// a `connectUrl`. On Windows the OAuth round-trip can exceed that
|
||||
/// window (Defender SmartScreen, slower browser launch, extra consent
|
||||
/// dialogs), so the invalidation call never fires and the chat
|
||||
/// runtime's cache stays frozen on the pre-connect snapshot even
|
||||
/// though the Settings UI polls `composio_list_connections` every 5 s
|
||||
/// and shows the user as "Connected".
|
||||
///
|
||||
/// The cross-platform defenses we layer on top:
|
||||
/// 1. [`composio_list_connections`] diff-invalidates the cache whenever
|
||||
/// the backend's active-toolkit set diverges from what's cached,
|
||||
/// so a running UI keeps the chat cache in sync within one poll
|
||||
/// interval.
|
||||
/// 2. This TTL caps worst-case staleness at 60 s regardless of
|
||||
/// whether the UI is open, the bus fires, or the user reconnected
|
||||
/// out-of-band.
|
||||
const CACHE_TTL: Duration = Duration::from_secs(60);
|
||||
|
||||
/// Cached entry: the integrations list plus the timestamp we wrote it.
|
||||
#[derive(Clone)]
|
||||
struct CachedIntegrations {
|
||||
entries: Vec<ConnectedIntegration>,
|
||||
cached_at: Instant,
|
||||
}
|
||||
|
||||
/// Process-wide cache for connected integrations, keyed by the config
|
||||
/// identity (the `config_path` string) so different user contexts don't
|
||||
/// collide. Each entry is populated on first fetch and returned on
|
||||
/// subsequent calls until explicitly invalidated or the TTL expires.
|
||||
static INTEGRATIONS_CACHE: LazyLock<RwLock<HashMap<String, CachedIntegrations>>> =
|
||||
LazyLock::new(|| RwLock::new(HashMap::new()));
|
||||
|
||||
/// Derive a stable cache key from a [`Config`]. We use the stringified
|
||||
/// `config_path` because it uniquely identifies a user context (it
|
||||
/// resolves to the per-user openhuman dir).
|
||||
fn cache_key(config: &Config) -> String {
|
||||
config.config_path.display().to_string()
|
||||
}
|
||||
|
||||
/// Clear cached connected integrations so the next call to
|
||||
/// [`fetch_connected_integrations`] hits the backend again.
|
||||
///
|
||||
/// Called by [`super::bus::ComposioConnectionCreatedSubscriber`] when a
|
||||
/// new OAuth connection completes, by [`composio_list_connections`]
|
||||
/// when it observes a divergence between the backend response and the
|
||||
/// cached snapshot, and from tests. Clears the entire map because the
|
||||
/// callers don't carry a config reference.
|
||||
pub fn invalidate_connected_integrations_cache() {
|
||||
if let Ok(mut guard) = INTEGRATIONS_CACHE.write() {
|
||||
let entries = guard.len();
|
||||
guard.clear();
|
||||
tracing::info!(
|
||||
cached_keys = entries,
|
||||
"[composio][integrations] cache invalidated"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Read-only snapshot of the currently cached connected integrations for
|
||||
/// the given config, or [`None`] when the cache is empty, expired, or
|
||||
/// the lock is held by a writer.
|
||||
///
|
||||
/// Designed for hot-path callers that want a cheap "what does the cache
|
||||
/// already say?" probe without triggering a backend fetch. The agent
|
||||
/// harness uses this on every turn to detect mid-session connection
|
||||
/// changes — it relies on the desktop UI's 5 s `composio_list_connections`
|
||||
/// poll (which calls into [`fetch_connected_integrations`] and
|
||||
/// repopulates this cache) plus the event-driven invalidation path to
|
||||
/// keep the cache current.
|
||||
///
|
||||
/// `try_read` (not `read`) so a writer in progress — e.g. the UI poll
|
||||
/// repopulating the cache — never blocks a turn. Worst case the agent
|
||||
/// sees `None` for one turn while the writer holds the lock; the next
|
||||
/// turn picks up the value naturally.
|
||||
///
|
||||
/// TTL is enforced defensively: entries older than [`CACHE_TTL`] are
|
||||
/// treated as missing even though they're still in the map (a stale
|
||||
/// entry would otherwise pin the agent to a frozen view if every
|
||||
/// invalidation path silently failed).
|
||||
pub fn cached_active_integrations(config: &Config) -> Option<Vec<ConnectedIntegration>> {
|
||||
let key = cache_key(config);
|
||||
let guard = match INTEGRATIONS_CACHE.try_read() {
|
||||
Ok(g) => g,
|
||||
Err(_) => {
|
||||
tracing::trace!(
|
||||
key = %key,
|
||||
"[composio][integrations_cache] cached_active_integrations:lock_contended"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let Some(cached) = guard.get(&key) else {
|
||||
tracing::trace!(
|
||||
key = %key,
|
||||
"[composio][integrations_cache] cached_active_integrations:miss"
|
||||
);
|
||||
return None;
|
||||
};
|
||||
let age = cached.cached_at.elapsed();
|
||||
if age > CACHE_TTL {
|
||||
tracing::trace!(
|
||||
key = %key,
|
||||
age_ms = age.as_millis() as u64,
|
||||
ttl_ms = CACHE_TTL.as_millis() as u64,
|
||||
"[composio][integrations_cache] cached_active_integrations:expired"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
tracing::trace!(
|
||||
key = %key,
|
||||
entries = cached.entries.len(),
|
||||
age_ms = age.as_millis() as u64,
|
||||
"[composio][integrations_cache] cached_active_integrations:hit"
|
||||
);
|
||||
Some(cached.entries.clone())
|
||||
}
|
||||
|
||||
/// Stable hash of the *routing-relevant* slice of a connected-integrations
|
||||
/// snapshot.
|
||||
///
|
||||
/// Two snapshots produce the same hash iff they would synthesise the
|
||||
/// same `delegate_<toolkit>` tool set in the orchestrator's
|
||||
/// function-calling schema. The hash is:
|
||||
///
|
||||
/// - **Order-independent** — callers don't need to sort the input.
|
||||
/// - **Description-insensitive** — Composio catalogue text edits don't
|
||||
/// trigger a refresh. The schema's tool-description field still
|
||||
/// picks up new text on the next *real* (membership-changing)
|
||||
/// refresh, so descriptions are never permanently stale.
|
||||
/// - **Process-local** — [`std::collections::hash_map::DefaultHasher`]
|
||||
/// is randomly seeded per process. Fine because we only compare
|
||||
/// hashes within one process lifetime.
|
||||
///
|
||||
/// Only `connected == true` entries contribute. Unconnected toolkits are
|
||||
/// stripped by [`super::super::tools::orchestrator_tools::collect_orchestrator_tools`]
|
||||
/// anyway, so churn among the unconnected set never changes the agent's
|
||||
/// surface and shouldn't trigger a refresh.
|
||||
pub fn connected_set_hash(integrations: &[ConnectedIntegration]) -> u64 {
|
||||
use std::collections::hash_map::DefaultHasher;
|
||||
use std::hash::{Hash, Hasher};
|
||||
|
||||
let mut slugs: Vec<&str> = integrations
|
||||
.iter()
|
||||
.filter(|i| i.connected)
|
||||
.map(|i| i.toolkit.as_str())
|
||||
.collect();
|
||||
slugs.sort();
|
||||
|
||||
let mut hasher = DefaultHasher::new();
|
||||
slugs.hash(&mut hasher);
|
||||
hasher.finish()
|
||||
}
|
||||
|
||||
/// Collect the set of toolkit slugs marked `connected` in a snapshot.
|
||||
///
|
||||
/// Exposed to [`sync_cache_with_connections`] so it can diff the live
|
||||
/// backend connection list against what the chat runtime currently
|
||||
/// believes is connected.
|
||||
fn connected_toolkit_set(integrations: &[ConnectedIntegration]) -> HashSet<String> {
|
||||
integrations
|
||||
.iter()
|
||||
.filter(|i| i.connected)
|
||||
.map(|i| i.toolkit.clone())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Reconcile the process-wide integrations cache with a fresh backend
|
||||
/// `list_connections` response.
|
||||
///
|
||||
/// Called from [`composio_list_connections`], which the desktop UI
|
||||
/// polls every 5 s (see `app/src/lib/composio/hooks.ts`). When the set
|
||||
/// of ACTIVE/CONNECTED toolkits in the response differs from what's in
|
||||
/// the cache, we invalidate so the chat runtime re-fetches on its next
|
||||
/// `fetch_connected_integrations` call. This keeps tool availability
|
||||
/// in chat in sync with the badge the user sees in Settings, even when
|
||||
/// the primary event-bus invalidation path misses (e.g. Windows OAuth
|
||||
/// flows that overrun the 60 s readiness poll).
|
||||
fn sync_cache_with_connections(connections: &[super::types::ComposioConnection]) {
|
||||
let live_active: HashSet<String> = connections
|
||||
.iter()
|
||||
.filter(|c| c.is_active())
|
||||
.map(|c| c.normalized_toolkit())
|
||||
.filter(|toolkit| !toolkit.is_empty())
|
||||
.collect();
|
||||
|
||||
// Read once to decide whether any cache entry is out of sync. We
|
||||
// clone out the keys + connected sets so we can release the read
|
||||
// lock before taking the write lock.
|
||||
let divergent_keys: Vec<(String, HashSet<String>, HashSet<String>)> = {
|
||||
let Ok(guard) = INTEGRATIONS_CACHE.read() else {
|
||||
return;
|
||||
};
|
||||
guard
|
||||
.iter()
|
||||
.filter_map(|(key, cached)| {
|
||||
let cached_set = connected_toolkit_set(&cached.entries);
|
||||
if cached_set != live_active {
|
||||
Some((key.clone(), cached_set, live_active.clone()))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
|
||||
if divergent_keys.is_empty() {
|
||||
tracing::debug!(
|
||||
live_connected = live_active.len(),
|
||||
"[composio][integrations] list_connections matches cache — no invalidation needed"
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
if let Ok(mut guard) = INTEGRATIONS_CACHE.write() {
|
||||
for (key, cached_set, live_set) in divergent_keys {
|
||||
// Diff logging — makes Windows-timing regressions easy to
|
||||
// catch in user-supplied debug dumps without leaking any
|
||||
// PII (toolkit slugs are public strings like "gmail").
|
||||
let added: Vec<&String> = live_set.difference(&cached_set).collect();
|
||||
let removed: Vec<&String> = cached_set.difference(&live_set).collect();
|
||||
tracing::info!(
|
||||
key = %key,
|
||||
?added,
|
||||
?removed,
|
||||
"[composio][integrations] cache diverges from backend — invalidating"
|
||||
);
|
||||
guard.remove(&key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Fetch the user's active Composio connections and their available
|
||||
/// tool actions, returning a prompt-ready summary.
|
||||
///
|
||||
/// This is the **single source of truth** for connected integration
|
||||
/// data injected into system prompts — both the agent turn loop and
|
||||
/// the debug dump CLI call this function.
|
||||
///
|
||||
/// Results are cached process-wide (keyed by config identity) and
|
||||
/// returned instantly on subsequent calls. The cache is invalidated
|
||||
/// when a new connection is created
|
||||
/// (via [`invalidate_connected_integrations_cache`]), when a UI
|
||||
/// `list_connections` poll observes a divergent live set, when
|
||||
/// [`CACHE_TTL`] expires, or on process restart.
|
||||
///
|
||||
/// Best-effort: returns an empty vec when the user isn't signed in,
|
||||
/// the backend is unreachable, or any step fails.
|
||||
pub async fn fetch_connected_integrations(config: &Config) -> Vec<ConnectedIntegration> {
|
||||
match fetch_connected_integrations_status(config).await {
|
||||
FetchConnectedIntegrationsStatus::Authoritative(v) => v,
|
||||
FetchConnectedIntegrationsStatus::Unavailable => Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Discriminated outcome from [`fetch_connected_integrations_status`].
|
||||
///
|
||||
/// Lets callers distinguish "the backend confirmed the user has zero
|
||||
/// active connections right now" from "we couldn't talk to the backend
|
||||
/// (no client, transient failure, …) and have no truth to report".
|
||||
///
|
||||
/// The legacy [`fetch_connected_integrations`] collapses both into an
|
||||
/// empty `Vec`, which is fine for prompt-building (they look the same)
|
||||
/// but dangerous for spawn-time allowlist gates — using empty as truth
|
||||
/// in the unavailable case would silently wipe the user's allowlist
|
||||
/// during a transient 5xx.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum FetchConnectedIntegrationsStatus {
|
||||
/// Backend was reachable. Vec may legitimately be empty (no
|
||||
/// allowlisted toolkits, or no active connections).
|
||||
Authoritative(Vec<ConnectedIntegration>),
|
||||
/// Backend wasn't reachable (no auth client, transient error). The
|
||||
/// caller should fall back to its prior snapshot rather than treat
|
||||
/// "no connections" as truth.
|
||||
Unavailable,
|
||||
}
|
||||
|
||||
/// Status-returning variant of [`fetch_connected_integrations`].
|
||||
///
|
||||
/// Same caching, same cache-invalidation semantics — only the return
|
||||
/// shape differs. Cache hits are by definition `Authoritative` because
|
||||
/// we only cache the `Some(...)` arm of `_uncached` (i.e. results the
|
||||
/// backend confirmed).
|
||||
pub async fn fetch_connected_integrations_status(
|
||||
config: &Config,
|
||||
) -> FetchConnectedIntegrationsStatus {
|
||||
let key = cache_key(config);
|
||||
|
||||
// Fast path: return cached result if fresh. Stale entries fall
|
||||
// through to the backend fetch below so the chat runtime can never
|
||||
// be more than `CACHE_TTL` behind a real-world change.
|
||||
if let Ok(guard) = INTEGRATIONS_CACHE.read() {
|
||||
if let Some(cached) = guard.get(&key) {
|
||||
let age = cached.cached_at.elapsed();
|
||||
if age < CACHE_TTL {
|
||||
tracing::debug!(
|
||||
count = cached.entries.len(),
|
||||
age_ms = age.as_millis() as u64,
|
||||
key = %key,
|
||||
"[composio][integrations] returning cached result"
|
||||
);
|
||||
return FetchConnectedIntegrationsStatus::Authoritative(cached.entries.clone());
|
||||
}
|
||||
tracing::info!(
|
||||
count = cached.entries.len(),
|
||||
age_ms = age.as_millis() as u64,
|
||||
ttl_ms = CACHE_TTL.as_millis() as u64,
|
||||
key = %key,
|
||||
"[composio][integrations] cache entry expired — refetching"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
match fetch_connected_integrations_uncached(config).await {
|
||||
Some(result) => {
|
||||
// Backend was reachable — cache the result (even if empty).
|
||||
if let Ok(mut guard) = INTEGRATIONS_CACHE.write() {
|
||||
guard.insert(
|
||||
key,
|
||||
CachedIntegrations {
|
||||
entries: result.clone(),
|
||||
cached_at: Instant::now(),
|
||||
},
|
||||
);
|
||||
}
|
||||
FetchConnectedIntegrationsStatus::Authoritative(result)
|
||||
}
|
||||
None => {
|
||||
// No auth / client unavailable — do NOT cache so a
|
||||
// subsequent call with a different config can retry.
|
||||
FetchConnectedIntegrationsStatus::Unavailable
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The actual backend fetch, called on cache miss.
|
||||
///
|
||||
/// Returns `Some(vec)` when the backend was reachable. The returned
|
||||
/// vector is the merged **integration overview** — every toolkit in
|
||||
/// the backend allowlist appears as one entry, with a `connected`
|
||||
/// flag indicating whether the user has an active OAuth connection.
|
||||
/// Connected entries also carry the per-action tool catalogue
|
||||
/// (fetched in a single batched call).
|
||||
///
|
||||
/// Returns `None` when we couldn't even build a client (no auth),
|
||||
/// signalling the caller should NOT cache this result.
|
||||
async fn fetch_connected_integrations_uncached(
|
||||
config: &Config,
|
||||
) -> Option<Vec<ConnectedIntegration>> {
|
||||
use super::client::{create_composio_client, direct_list_connections, ComposioClientKind};
|
||||
use super::providers::toolkit_description;
|
||||
|
||||
// Route via the mode-aware factory so the chat-agent's
|
||||
// "connected_integrations" view reflects the live tenant — backend
|
||||
// (tinyhumans) or direct (user's personal Composio). Prior to #1710
|
||||
// Wave 3 this path called `build_composio_client` directly, which
|
||||
// is backend-only — after a `composio.mode = "direct"` toggle the
|
||||
// cache kept replaying the tinyhumans-tenant connections back into
|
||||
// the integration overview (e.g. gmail / notion appearing as
|
||||
// connected in direct mode even when the user's direct tenant had
|
||||
// a different set of toolkits). Resolving per call closes the
|
||||
// loop: `ComposioConfigChangedSubscriber` invalidates the cache on
|
||||
// toggle and the next miss re-populates it from the live tenant.
|
||||
let kind = match create_composio_client(config) {
|
||||
Ok(kind) => kind,
|
||||
Err(e) => {
|
||||
tracing::debug!(
|
||||
error = %e,
|
||||
"[composio] fetch_connected_integrations: no client (not signed in?)"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
// Pull the allowlist + connections + tool catalogue. Backend mode
|
||||
// walks the tinyhumans tenant's curated allowlist via
|
||||
// `list_toolkits`; direct mode has no centralised allowlist (per
|
||||
// `ops::composio_list_toolkits`'s direct-mode branch) so the
|
||||
// user's set of active connections IS the universe of valid
|
||||
// toolkit arguments.
|
||||
//
|
||||
// On transient errors we return `None` instead of a degraded
|
||||
// `Some(Vec::new())` so `fetch_connected_integrations` does NOT
|
||||
// cache the failure. Caching an empty allowlist would hide every
|
||||
// integration from the orchestrator until the process restarts or
|
||||
// the cache is explicitly invalidated — a single 5xx during
|
||||
// startup would silently break delegation for the whole session.
|
||||
let (allowlisted_toolkits, connections, tools_by_toolkit): (
|
||||
Vec<String>,
|
||||
Vec<super::types::ComposioConnection>,
|
||||
Vec<super::types::ComposioToolSchema>,
|
||||
) = match &kind {
|
||||
ComposioClientKind::Backend(client) => {
|
||||
let allowlist: Vec<String> = match client.list_toolkits().await {
|
||||
Ok(resp) => resp
|
||||
.toolkits
|
||||
.into_iter()
|
||||
.map(|toolkit| toolkit.trim().to_ascii_lowercase())
|
||||
.filter(|toolkit| !toolkit.is_empty())
|
||||
.collect(),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"[composio] fetch_connected_integrations: list_toolkits (backend) failed: {e}"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
if allowlist.is_empty() {
|
||||
tracing::debug!(
|
||||
"[composio] fetch_connected_integrations: backend allowlist is empty"
|
||||
);
|
||||
return Some(Vec::new());
|
||||
}
|
||||
|
||||
let connections = match client.list_connections().await {
|
||||
Ok(resp) => resp.connections,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"[composio] fetch_connected_integrations: list_connections (backend) failed: {e}"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
// Tool catalogue scoped to the active subset only —
|
||||
// not-connected toolkits won't be invoked from a sub-agent.
|
||||
let connected_slugs_for_tools: Vec<String> = {
|
||||
let mut v: Vec<String> = connections
|
||||
.iter()
|
||||
.filter(|c| c.is_active())
|
||||
.map(|c| c.normalized_toolkit())
|
||||
.filter(|t| !t.is_empty())
|
||||
.collect();
|
||||
v.sort();
|
||||
v.dedup();
|
||||
v
|
||||
};
|
||||
let tools = if connected_slugs_for_tools.is_empty() {
|
||||
Vec::new()
|
||||
} else {
|
||||
match client
|
||||
.list_tools(Some(&connected_slugs_for_tools), None)
|
||||
.await
|
||||
{
|
||||
Ok(resp) => resp.tools,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"[composio] fetch_connected_integrations: list_tools (backend) failed: {e}"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
(allowlist, connections, tools)
|
||||
}
|
||||
ComposioClientKind::Direct(direct) => {
|
||||
// Direct mode: walk the user's personal Composio tenant
|
||||
// for *connection state* (active accounts on their key) —
|
||||
// there's no central allowlist in direct mode, so the
|
||||
// active set IS the allowlist.
|
||||
//
|
||||
// Tool *schemas* are tenant-agnostic — Composio's action
|
||||
// definitions (e.g. GMAIL_SEND_EMAIL parameter shape) are
|
||||
// identical regardless of which Composio tenant a user is
|
||||
// connected via. So we best-effort fetch schemas through
|
||||
// the backend client (curated list_tools) if a backend
|
||||
// session is available, even though connection routing
|
||||
// goes to the user's direct tenant. This preserves the
|
||||
// chat agent's "21 gmail actions available" view while
|
||||
// execution itself (via `ComposioActionTool` / Wave 1
|
||||
// factory) still routes to the user's tenant. Direct-only
|
||||
// users without a backend session get empty tools — that
|
||||
// matches `composio_list_tools`'s direct-mode policy and
|
||||
// the `subagent_runner` LazyToolkitResolver still resolves
|
||||
// tools lazily at delegation time.
|
||||
let connections = match direct_list_connections(direct).await {
|
||||
Ok(resp) => resp.connections,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"[composio] fetch_connected_integrations: list_connections (direct) failed: {e:#}"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
let allowlist: Vec<String> = {
|
||||
let mut v: Vec<String> = connections
|
||||
.iter()
|
||||
.filter(|c| c.is_active())
|
||||
.map(|c| c.normalized_toolkit())
|
||||
.filter(|t| !t.is_empty())
|
||||
.collect();
|
||||
v.sort();
|
||||
v.dedup();
|
||||
v
|
||||
};
|
||||
if allowlist.is_empty() {
|
||||
tracing::info!(
|
||||
"[composio-direct] fetch_connected_integrations: direct tenant has no active connections; returning empty overview"
|
||||
);
|
||||
return Some(Vec::new());
|
||||
}
|
||||
tracing::debug!(
|
||||
connected = allowlist.len(),
|
||||
"[composio-direct] fetch_connected_integrations: using direct tenant's active set as allowlist (no central allowlist in direct mode)"
|
||||
);
|
||||
|
||||
// Best-effort: pull tool schemas via the backend client
|
||||
// (definitional source). Failure is non-fatal — we fall
|
||||
// back to empty tools and let lazy resolution handle it.
|
||||
let tools = match super::client::build_composio_client(config) {
|
||||
Some(backend_client) => {
|
||||
match backend_client.list_tools(Some(&allowlist), None).await {
|
||||
Ok(resp) => {
|
||||
tracing::debug!(
|
||||
count = resp.tools.len(),
|
||||
"[composio-direct] fetch_connected_integrations: pulled tool schemas from backend (tenant-agnostic definitional source)"
|
||||
);
|
||||
resp.tools
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::info!(
|
||||
"[composio-direct] fetch_connected_integrations: backend list_tools failed (will use lazy fallback at delegation time): {e:#}"
|
||||
);
|
||||
Vec::new()
|
||||
}
|
||||
}
|
||||
}
|
||||
None => {
|
||||
tracing::info!(
|
||||
"[composio-direct] fetch_connected_integrations: no backend session for schema fetch; lazy fallback at delegation time"
|
||||
);
|
||||
Vec::new()
|
||||
}
|
||||
};
|
||||
(allowlist, connections, tools)
|
||||
}
|
||||
};
|
||||
|
||||
// Active connection slugs (status filter mirrors the original logic).
|
||||
let connected_slugs: std::collections::HashSet<String> = connections
|
||||
.iter()
|
||||
.filter(|c| c.is_active())
|
||||
.map(|c| c.normalized_toolkit())
|
||||
.filter(|toolkit| !toolkit.is_empty())
|
||||
.collect();
|
||||
|
||||
// Most-informative *non-active* status per toolkit slug. Lets the
|
||||
// integrations_agent spawn-gate (#2365) emit a precise message
|
||||
// when a connection row exists but isn't usable yet (`INITIATED`
|
||||
// — OAuth still in progress) or any longer (`EXPIRED` / `FAILED`)
|
||||
// — instead of the legacy generic "available but not authorized".
|
||||
//
|
||||
// Status priority (UI-actionability):
|
||||
// 1. EXPIRED — reconnect path
|
||||
// 2. FAILED / ERROR — reconnect path
|
||||
// 3. INITIATED / INITIALIZING / PENDING — finish OAuth in browser
|
||||
// 4. anything else — passes through verbatim
|
||||
let non_active_status_by_slug: std::collections::HashMap<String, String> = {
|
||||
fn priority(status: &str) -> u8 {
|
||||
let s = status.trim().to_ascii_uppercase();
|
||||
match s.as_str() {
|
||||
"EXPIRED" => 4,
|
||||
"FAILED" | "ERROR" => 3,
|
||||
"INITIATED" | "INITIALIZING" | "PENDING" => 2,
|
||||
_ => 1,
|
||||
}
|
||||
}
|
||||
let mut map: std::collections::HashMap<String, (u8, String)> =
|
||||
std::collections::HashMap::new();
|
||||
for conn in connections.iter().filter(|c| !c.is_active()) {
|
||||
let slug = conn.normalized_toolkit();
|
||||
if slug.is_empty() {
|
||||
continue;
|
||||
}
|
||||
// Don't override an ACTIVE-slug — those carry no non-active
|
||||
// status from this map's perspective.
|
||||
if connected_slugs.contains(&slug) {
|
||||
continue;
|
||||
}
|
||||
let p = priority(&conn.status);
|
||||
map.entry(slug.clone())
|
||||
.and_modify(|cur| {
|
||||
if p > cur.0 {
|
||||
tracing::debug!(
|
||||
target: "composio",
|
||||
toolkit = %slug,
|
||||
previous_status = %cur.1,
|
||||
previous_priority = cur.0,
|
||||
new_status = %conn.status,
|
||||
new_priority = p,
|
||||
"[composio] non_active_status_by_slug: upgraded most-informative status"
|
||||
);
|
||||
*cur = (p, conn.status.clone());
|
||||
} else {
|
||||
tracing::trace!(
|
||||
target: "composio",
|
||||
toolkit = %slug,
|
||||
kept_status = %cur.1,
|
||||
kept_priority = cur.0,
|
||||
candidate_status = %conn.status,
|
||||
candidate_priority = p,
|
||||
"[composio] non_active_status_by_slug: kept higher-priority status"
|
||||
);
|
||||
}
|
||||
})
|
||||
.or_insert_with(|| {
|
||||
tracing::debug!(
|
||||
target: "composio",
|
||||
toolkit = %slug,
|
||||
status = %conn.status,
|
||||
priority = p,
|
||||
"[composio] non_active_status_by_slug: first non-active row"
|
||||
);
|
||||
(p, conn.status.clone())
|
||||
});
|
||||
}
|
||||
let final_map: std::collections::HashMap<String, String> =
|
||||
map.into_iter().map(|(k, (_, v))| (k, v)).collect();
|
||||
tracing::debug!(
|
||||
target: "composio",
|
||||
entries = final_map.len(),
|
||||
"[composio] non_active_status_by_slug: final map"
|
||||
);
|
||||
final_map
|
||||
};
|
||||
|
||||
// Deduplicate the allowlist so a backend that returns duplicates
|
||||
// doesn't produce dual entries downstream.
|
||||
let mut unique_toolkits: Vec<String> = allowlisted_toolkits.clone();
|
||||
unique_toolkits.sort();
|
||||
unique_toolkits.dedup();
|
||||
|
||||
// Build one entry per allowlisted toolkit. Connected entries
|
||||
// carry their action catalogue; not-connected entries carry an
|
||||
// empty `tools` vec.
|
||||
let mut integrations: Vec<ConnectedIntegration> = Vec::with_capacity(unique_toolkits.len());
|
||||
for slug in &unique_toolkits {
|
||||
let connected = connected_slugs.contains(slug);
|
||||
// Anchor the prefix with an underscore so slugs that share
|
||||
// a text prefix (e.g. `git` vs `github`) don't false-match
|
||||
// each other's actions. `GMAIL_SEND_EMAIL` matches `gmail_`,
|
||||
// not just `gmail`, so siblings stay in their own buckets.
|
||||
let action_prefix = format!("{}_", slug.to_uppercase());
|
||||
let (tools, gated_tools): (Vec<ConnectedIntegrationTool>, Vec<GatedIntegrationTool>) =
|
||||
if connected {
|
||||
// Apply the same curated-whitelist + user-scope filter the
|
||||
// meta-tool layer uses, so the integrations_agent prompt
|
||||
// only advertises actions the agent is actually allowed to
|
||||
// call. One pref load per toolkit (not per action).
|
||||
//
|
||||
// Actions that the catalog *does* know about but the user's
|
||||
// current scope pref denies are routed into `gated_tools` so
|
||||
// the agent can honestly answer "I have this capability but
|
||||
// it needs the {scope} toggle in Connections → {toolkit}".
|
||||
// The agent cannot flip the scope itself — that's a UI-only
|
||||
// action. Without this gated surface the LLM has no way to
|
||||
// know the gated action exists at all and will tell the user
|
||||
// "I don't support that" — technically correct about its
|
||||
// callable surface, but misleading about the toolkit.
|
||||
let pref = super::providers::load_user_scope_or_default(slug).await;
|
||||
let mut visible: Vec<ConnectedIntegrationTool> = Vec::new();
|
||||
let mut gated: Vec<GatedIntegrationTool> = Vec::new();
|
||||
for t in tools_by_toolkit
|
||||
.iter()
|
||||
.filter(|t| t.function.name.starts_with(&action_prefix))
|
||||
{
|
||||
if super::providers::is_action_visible_with_pref(&t.function.name, &pref) {
|
||||
visible.push(ConnectedIntegrationTool {
|
||||
name: t.function.name.clone(),
|
||||
description: t.function.description.clone().unwrap_or_default(),
|
||||
parameters: t.function.parameters.clone(),
|
||||
});
|
||||
} else if let Some(required_scope) =
|
||||
super::providers::curated_scope_for(&t.function.name)
|
||||
{
|
||||
// Only surface CURATED actions as `gated` — uncurated
|
||||
// tools (which fall through to `classify_unknown` and
|
||||
// happen to land outside the user's pref) are not
|
||||
// first-class user-facing capabilities, and listing
|
||||
// them would clutter the prompt with internal slugs.
|
||||
// Deliberately NO `parameters` field: the LLM should
|
||||
// not be able to construct a call envelope; it can
|
||||
// only describe + point at the unlock path.
|
||||
// Ship the unlock path as data — single path today
|
||||
// (the Connections UI toggle). The agent does NOT
|
||||
// have a tool to flip scopes; that capability was
|
||||
// removed because LLM-mediated scope elevation made
|
||||
// the safety contract depend on model behavior and
|
||||
// was a soft gate the model could route around. If
|
||||
// more unlock paths exist in future (per-action
|
||||
// approval modal, time-boxed elevation, etc.) they
|
||||
// land here and the prompt renderer picks them up.
|
||||
let scope_str = required_scope.as_str();
|
||||
let unlock_paths = vec![format!(
|
||||
"the user enables it themselves in \
|
||||
Connections → {slug} → {scope_str}"
|
||||
)];
|
||||
gated.push(GatedIntegrationTool {
|
||||
name: t.function.name.clone(),
|
||||
description: t.function.description.clone().unwrap_or_default(),
|
||||
required_scope: scope_str.to_string(),
|
||||
unlock_paths,
|
||||
});
|
||||
}
|
||||
}
|
||||
tracing::debug!(
|
||||
toolkit = %slug,
|
||||
visible = visible.len(),
|
||||
gated = gated.len(),
|
||||
"[composio][scopes] integrations prompt action set"
|
||||
);
|
||||
(visible, gated)
|
||||
} else {
|
||||
(Vec::new(), Vec::new())
|
||||
};
|
||||
|
||||
integrations.push(ConnectedIntegration {
|
||||
toolkit: slug.clone(),
|
||||
description: toolkit_description(slug).to_string(),
|
||||
tools,
|
||||
gated_tools,
|
||||
connected,
|
||||
non_active_status: if connected {
|
||||
None
|
||||
} else {
|
||||
non_active_status_by_slug.get(slug).cloned()
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
integrations.sort_by(|a, b| a.toolkit.cmp(&b.toolkit));
|
||||
|
||||
let connected_count = integrations.iter().filter(|i| i.connected).count();
|
||||
tracing::info!(
|
||||
total = integrations.len(),
|
||||
connected = connected_count,
|
||||
"[composio] fetch_connected_integrations: done"
|
||||
);
|
||||
for ci in &integrations {
|
||||
tracing::debug!(
|
||||
toolkit = %ci.toolkit,
|
||||
connected = ci.connected,
|
||||
non_active_status = ?ci.non_active_status,
|
||||
tool_count = ci.tools.len(),
|
||||
"[composio] integration overview"
|
||||
);
|
||||
}
|
||||
|
||||
Some(integrations)
|
||||
}
|
||||
|
||||
/// Just-in-time fetch of every available action for a single Composio
|
||||
/// toolkit, returned in the [`ConnectedIntegrationTool`] shape the
|
||||
/// `integrations_agent` spawn path expects.
|
||||
///
|
||||
/// Unlike [`fetch_connected_integrations`] (which bulk-fetches every
|
||||
/// connected toolkit's tools once per session and caches the result),
|
||||
/// this helper is uncached and scoped to a single toolkit — meant to
|
||||
/// be called at `integrations_agent` spawn time so the sub-agent's
|
||||
/// prompt always reflects the toolkit's current action catalogue.
|
||||
///
|
||||
/// The filter `starts_with("{TOOLKIT}_")` matches
|
||||
/// `fetch_connected_integrations_uncached`'s own namespacing rule so
|
||||
/// siblings like `github` / `git` don't leak into each other's buckets.
|
||||
///
|
||||
/// `tags` narrows the result by Composio action tag (OR semantics). Only
|
||||
/// honoured for the GitHub toolkit; passed through to `list_tools` so the
|
||||
/// backend can skip the repo-list force-include and return a focused set.
|
||||
///
|
||||
/// Returns an empty vec when the backend has no actions for the
|
||||
/// toolkit (valid steady state for a freshly-authorised integration
|
||||
/// whose catalogue hasn't been published yet). Returns `Err` only for
|
||||
/// transport / auth failures the caller should surface to the user.
|
||||
pub async fn fetch_toolkit_actions(
|
||||
client: &ComposioClient,
|
||||
toolkit: &str,
|
||||
tags: Option<&[String]>,
|
||||
) -> anyhow::Result<Vec<ConnectedIntegrationTool>> {
|
||||
let toolkit_slug = toolkit.trim();
|
||||
if toolkit_slug.is_empty() {
|
||||
anyhow::bail!("fetch_toolkit_actions: toolkit must not be empty");
|
||||
}
|
||||
let effective_tags = if should_forward_tags(Some(&[toolkit_slug.to_string()])) {
|
||||
tags
|
||||
} else {
|
||||
None
|
||||
};
|
||||
tracing::debug!(toolkit = %toolkit_slug, ?effective_tags, "[composio] fetch_toolkit_actions");
|
||||
let resp = client
|
||||
.list_tools(Some(&[toolkit_slug.to_string()]), effective_tags)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("list_tools failed for toolkit `{toolkit_slug}`: {e}"))?;
|
||||
let action_prefix = format!("{}_", toolkit_slug.to_uppercase());
|
||||
// Apply curated whitelist + user scope so spawn-time tool
|
||||
// discovery agrees with the bulk path and the meta-tool layer.
|
||||
let pref = super::providers::load_user_scope_or_default(toolkit_slug).await;
|
||||
let actions: Vec<ConnectedIntegrationTool> = resp
|
||||
.tools
|
||||
.into_iter()
|
||||
.filter(|t| t.function.name.starts_with(&action_prefix))
|
||||
.filter(|t| super::providers::is_action_visible_with_pref(&t.function.name, &pref))
|
||||
.map(|t| ConnectedIntegrationTool {
|
||||
name: t.function.name,
|
||||
description: t.function.description.unwrap_or_default(),
|
||||
parameters: t.function.parameters,
|
||||
})
|
||||
.collect();
|
||||
tracing::debug!(
|
||||
toolkit = %toolkit_slug,
|
||||
action_count = actions.len(),
|
||||
"[composio] fetch_toolkit_actions: done"
|
||||
);
|
||||
Ok(actions)
|
||||
}
|
||||
#[cfg(test)]
|
||||
pub(crate) use super::connected_integrations::{CachedIntegrations, CACHE_TTL, INTEGRATIONS_CACHE};
|
||||
#[cfg(test)]
|
||||
pub(crate) use crate::openhuman::context::prompt::ConnectedIntegration;
|
||||
#[cfg(test)]
|
||||
pub(crate) use std::time::{Duration, Instant};
|
||||
|
||||
// ── Direct mode (BYO API key) ───────────────────────────────────────
|
||||
|
||||
|
||||
@@ -647,101 +647,15 @@ fn encrypt_config_secrets(config: &mut Config) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
const ACTIVE_USER_STATE_FILE: &str = "active_user.toml";
|
||||
#[path = "load_user_state.rs"]
|
||||
mod load_user_state;
|
||||
#[cfg(test)]
|
||||
pub(crate) use load_user_state::ACTIVE_USER_STATE_FILE;
|
||||
pub use load_user_state::{
|
||||
clear_active_user, pre_login_user_dir, read_active_user_id, user_openhuman_dir,
|
||||
write_active_user_id, PRE_LOGIN_USER_ID,
|
||||
};
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
struct ActiveUserState {
|
||||
user_id: String,
|
||||
}
|
||||
|
||||
/// Reads the active user id from `{default_openhuman_dir}/active_user.toml`.
|
||||
/// Returns `None` when the file does not exist, is empty, or cannot be parsed.
|
||||
pub fn read_active_user_id(default_openhuman_dir: &Path) -> Option<String> {
|
||||
let path = default_openhuman_dir.join(ACTIVE_USER_STATE_FILE);
|
||||
let contents = std::fs::read_to_string(&path).ok()?;
|
||||
let state: ActiveUserState = toml::from_str(&contents).ok()?;
|
||||
let id = state.user_id.trim().to_string();
|
||||
if id.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(id)
|
||||
}
|
||||
}
|
||||
|
||||
/// Writes the active user id to `{default_openhuman_dir}/active_user.toml`.
|
||||
pub fn write_active_user_id(default_openhuman_dir: &Path, user_id: &str) -> Result<()> {
|
||||
let path = default_openhuman_dir.join(ACTIVE_USER_STATE_FILE);
|
||||
let state = ActiveUserState {
|
||||
user_id: user_id.to_string(),
|
||||
};
|
||||
let toml_str = toml::to_string_pretty(&state).context("serialize active_user.toml")?;
|
||||
std::fs::write(&path, toml_str)
|
||||
.with_context(|| format!("Failed to write active user state: {}", path.display()))?;
|
||||
tracing::debug!(user_id = %user_id, path = %path.display(), "active user written");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Removes the active user marker. After this, the next config load will
|
||||
/// use the default (unauthenticated) openhuman directory.
|
||||
pub fn clear_active_user(default_openhuman_dir: &Path) -> Result<()> {
|
||||
let path = default_openhuman_dir.join(ACTIVE_USER_STATE_FILE);
|
||||
if path.exists() {
|
||||
std::fs::remove_file(&path)
|
||||
.with_context(|| format!("Failed to remove active user state: {}", path.display()))?;
|
||||
tracing::debug!(path = %path.display(), "active user cleared");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns the user-scoped openhuman directory for the given user id:
|
||||
/// `{default_openhuman_dir}/users/{user_id}`.
|
||||
pub fn user_openhuman_dir(default_openhuman_dir: &Path, user_id: &str) -> PathBuf {
|
||||
default_openhuman_dir.join("users").join(user_id)
|
||||
}
|
||||
|
||||
/// Stable id used to scope the openhuman directory before any user has
|
||||
/// logged in. All memory, state, config, sessions and workspace files
|
||||
/// created on first init land under `{root}/users/{PRE_LOGIN_USER_ID}`
|
||||
/// so nothing is ever written directly at the root `.openhuman` path.
|
||||
///
|
||||
/// On first successful login, this directory is migrated into the real
|
||||
/// user-scoped directory (see `credentials::ops::store_session`).
|
||||
pub const PRE_LOGIN_USER_ID: &str = "local";
|
||||
|
||||
/// Returns the pre-login (unauthenticated) user directory:
|
||||
/// `{default_openhuman_dir}/users/local`.
|
||||
pub fn pre_login_user_dir(default_openhuman_dir: &Path) -> PathBuf {
|
||||
user_openhuman_dir(default_openhuman_dir, PRE_LOGIN_USER_ID)
|
||||
}
|
||||
|
||||
/// Try to parse config TOML. On failure, try `.bak`, then fall back to `Config::default()`.
|
||||
///
|
||||
/// Returns `(config, was_corrupted)` where `was_corrupted == true` means the
|
||||
/// primary failed to parse and the caller is responsible for archiving the
|
||||
/// corrupt primary and persisting the recovered/default config.
|
||||
///
|
||||
/// This is a standalone async function (not a method on Config) so it can be
|
||||
/// called from both `load_or_init` and `load_from_default_paths`.
|
||||
///
|
||||
/// **Why the parse runs via `spawn_blocking`:** `toml::from_str::<Config>`
|
||||
/// is a recursive-descent parser whose serde-monomorphised `Visitor`
|
||||
/// frames for our deeply-nested `Config` cost several KB each. When this
|
||||
/// function is called from the bottom of a deep async tower — e.g.
|
||||
/// `composio_list_tools` reloading the config per call (#1710 Wave 4),
|
||||
/// reached via `chat → orchestrator → delegate_to_integrations_agent →
|
||||
/// sub-agent → composio_*` — running the parse inline on the tokio
|
||||
/// worker thread blows the ~2 MB worker stack and aborts the in-process
|
||||
/// core with `SIGBUS / KERN_PROTECTION_FAILURE` (see `crahs.log`,
|
||||
/// 2026-05-17, and `tests/composio_list_tools_stack_overflow_regression.rs`).
|
||||
/// Moving the parse onto the blocking-pool gives it a *fresh* thread
|
||||
/// stack with no async tower above it, so the same parser frames easily
|
||||
/// fit. (An earlier draft of this fix also fronted
|
||||
/// `config::ops::load_config_with_timeout` with a per-process cache to
|
||||
/// skip the parse on repeat calls, but it was reverted — the in-process
|
||||
/// integration tests in `tests/json_rpc_e2e.rs` reuse workspace paths
|
||||
/// and load config mid-mutation, racing the cache. The spawn_blocking
|
||||
/// move is sufficient on its own once paired with the Tauri worker
|
||||
/// stack bump in `app/src-tauri/src/lib.rs`.)
|
||||
async fn parse_config_with_recovery(config_path: &Path, contents: &str) -> (Config, bool) {
|
||||
let parse_err = match parse_toml_off_worker(contents.to_string()).await {
|
||||
Ok(config) => {
|
||||
|
||||
@@ -0,0 +1,122 @@
|
||||
use anyhow::{Context, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::io::Write;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
pub(crate) const ACTIVE_USER_STATE_FILE: &str = "active_user.toml";
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
struct ActiveUserState {
|
||||
user_id: String,
|
||||
}
|
||||
|
||||
/// Reads the active user id from `{default_openhuman_dir}/active_user.toml`.
|
||||
/// Returns `None` when the file does not exist, is empty, or cannot be parsed.
|
||||
pub fn read_active_user_id(default_openhuman_dir: &Path) -> Option<String> {
|
||||
let path = default_openhuman_dir.join(ACTIVE_USER_STATE_FILE);
|
||||
let contents = std::fs::read_to_string(&path).ok()?;
|
||||
let state: ActiveUserState = toml::from_str(&contents).ok()?;
|
||||
let id = state.user_id.trim().to_string();
|
||||
if id.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(id)
|
||||
}
|
||||
}
|
||||
|
||||
/// Writes the active user id to `{default_openhuman_dir}/active_user.toml`.
|
||||
pub fn write_active_user_id(default_openhuman_dir: &Path, user_id: &str) -> Result<()> {
|
||||
std::fs::create_dir_all(default_openhuman_dir).with_context(|| {
|
||||
format!(
|
||||
"Failed to create active user state directory: {}",
|
||||
default_openhuman_dir.display()
|
||||
)
|
||||
})?;
|
||||
|
||||
let path = default_openhuman_dir.join(ACTIVE_USER_STATE_FILE);
|
||||
let state = ActiveUserState {
|
||||
user_id: user_id.to_string(),
|
||||
};
|
||||
let toml_str = toml::to_string_pretty(&state).context("serialize active_user.toml")?;
|
||||
let temp_path = default_openhuman_dir.join(format!(
|
||||
".{ACTIVE_USER_STATE_FILE}.tmp-{}",
|
||||
uuid::Uuid::new_v4()
|
||||
));
|
||||
|
||||
let mut temp_file = std::fs::OpenOptions::new()
|
||||
.create_new(true)
|
||||
.write(true)
|
||||
.open(&temp_path)
|
||||
.with_context(|| {
|
||||
format!(
|
||||
"Failed to create temporary active user state: {}",
|
||||
temp_path.display()
|
||||
)
|
||||
})?;
|
||||
temp_file
|
||||
.write_all(toml_str.as_bytes())
|
||||
.context("Failed to write temporary active user state")?;
|
||||
temp_file
|
||||
.sync_all()
|
||||
.context("Failed to fsync temporary active user state")?;
|
||||
drop(temp_file);
|
||||
|
||||
if let Err(error) = std::fs::rename(&temp_path, &path) {
|
||||
let _ = std::fs::remove_file(&temp_path);
|
||||
anyhow::bail!(
|
||||
"Failed to atomically persist active user state {}: {error}",
|
||||
path.display()
|
||||
);
|
||||
}
|
||||
|
||||
sync_directory(default_openhuman_dir)?;
|
||||
tracing::debug!(user_id = %user_id, path = %path.display(), "active user written");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Removes the active user marker. After this, the next config load will
|
||||
/// use the default (unauthenticated) openhuman directory.
|
||||
pub fn clear_active_user(default_openhuman_dir: &Path) -> Result<()> {
|
||||
let path = default_openhuman_dir.join(ACTIVE_USER_STATE_FILE);
|
||||
if path.exists() {
|
||||
std::fs::remove_file(&path)
|
||||
.with_context(|| format!("Failed to remove active user state: {}", path.display()))?;
|
||||
tracing::debug!(path = %path.display(), "active user cleared");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns the user-scoped openhuman directory for the given user id:
|
||||
/// `{default_openhuman_dir}/users/{user_id}`.
|
||||
pub fn user_openhuman_dir(default_openhuman_dir: &Path, user_id: &str) -> PathBuf {
|
||||
default_openhuman_dir.join("users").join(user_id)
|
||||
}
|
||||
|
||||
/// Stable id used to scope the openhuman directory before any user has
|
||||
/// logged in. All memory, state, config, sessions and workspace files
|
||||
/// created on first init land under `{root}/users/{PRE_LOGIN_USER_ID}`
|
||||
/// so nothing is ever written directly at the root `.openhuman` path.
|
||||
///
|
||||
/// On first successful login, this directory is migrated into the real
|
||||
/// user-scoped directory (see `credentials::ops::store_session`).
|
||||
pub const PRE_LOGIN_USER_ID: &str = "local";
|
||||
|
||||
/// Returns the pre-login (unauthenticated) user directory:
|
||||
/// `{default_openhuman_dir}/users/local`.
|
||||
pub fn pre_login_user_dir(default_openhuman_dir: &Path) -> PathBuf {
|
||||
user_openhuman_dir(default_openhuman_dir, PRE_LOGIN_USER_ID)
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn sync_directory(path: &Path) -> Result<()> {
|
||||
let dir = std::fs::File::open(path)
|
||||
.with_context(|| format!("Failed to open directory for fsync: {}", path.display()))?;
|
||||
dir.sync_all()
|
||||
.with_context(|| format!("Failed to fsync directory metadata: {}", path.display()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
fn sync_directory(_path: &Path) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
@@ -6,6 +6,8 @@
|
||||
mod compatible_dump;
|
||||
#[path = "compatible_parse.rs"]
|
||||
mod compatible_parse;
|
||||
#[path = "compatible_request.rs"]
|
||||
mod compatible_request;
|
||||
#[path = "compatible_stream.rs"]
|
||||
mod compatible_stream;
|
||||
#[path = "compatible_types.rs"]
|
||||
@@ -25,10 +27,6 @@ use crate::openhuman::inference::provider::traits::{
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use futures_util::{stream, StreamExt};
|
||||
use reqwest::{
|
||||
header::{HeaderMap, HeaderValue, USER_AGENT},
|
||||
Client,
|
||||
};
|
||||
|
||||
use compatible_dump::{dump_prompt_if_enabled, dump_response_if_enabled, reserve_dump_seq};
|
||||
use compatible_parse::{
|
||||
@@ -207,234 +205,6 @@ impl OpenAiCompatibleProvider {
|
||||
self
|
||||
}
|
||||
|
||||
/// Resolve the effective temperature for `model`. Returns `None` when the
|
||||
/// model matches a pattern in `temperature_unsupported_models` (causing the
|
||||
/// field to be omitted from the serialised request). Otherwise yields the
|
||||
/// per-workload override if one was configured, else the caller's value.
|
||||
fn effective_temperature(&self, model: &str, temperature: f64) -> Option<f64> {
|
||||
if self
|
||||
.temperature_unsupported_models
|
||||
.iter()
|
||||
.any(|pat| super::temperature::glob_match(pat, model))
|
||||
{
|
||||
tracing::debug!(
|
||||
"[provider:{}] model='{}' matched temperature_unsupported_models — omitting temperature",
|
||||
self.name,
|
||||
model
|
||||
);
|
||||
None
|
||||
} else {
|
||||
Some(self.temperature_override.unwrap_or(temperature))
|
||||
}
|
||||
}
|
||||
|
||||
/// Read the ambient `thread_id` only when this provider has been
|
||||
/// opted in via [`with_openhuman_thread_id`]. Returns `None` for
|
||||
/// every third-party provider so the field is omitted by
|
||||
/// `skip_serializing_if`.
|
||||
fn outbound_thread_id(&self) -> Option<String> {
|
||||
if self.emit_openhuman_thread_id {
|
||||
super::thread_context::current_thread_id()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Collect all `system` role messages, concatenate their content,
|
||||
/// and prepend to the first `user` message. Drop all system messages.
|
||||
/// Used for providers (e.g. MiniMax) that reject `role: system`.
|
||||
fn flatten_system_messages(messages: &[ChatMessage]) -> Vec<ChatMessage> {
|
||||
let system_content: String = messages
|
||||
.iter()
|
||||
.filter(|m| m.role == "system")
|
||||
.map(|m| m.content.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n");
|
||||
|
||||
if system_content.is_empty() {
|
||||
return messages.to_vec();
|
||||
}
|
||||
|
||||
let mut result: Vec<ChatMessage> = messages
|
||||
.iter()
|
||||
.filter(|m| m.role != "system")
|
||||
.cloned()
|
||||
.collect();
|
||||
|
||||
if let Some(first_user) = result.iter_mut().find(|m| m.role == "user") {
|
||||
first_user.content = format!("{system_content}\n\n{}", first_user.content);
|
||||
} else {
|
||||
// No user message found: insert a synthetic user message with system content
|
||||
result.insert(0, ChatMessage::user(&system_content));
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
fn http_client(&self) -> Client {
|
||||
if let Some(ua) = self.user_agent.as_deref() {
|
||||
let mut headers = HeaderMap::new();
|
||||
if let Ok(value) = HeaderValue::from_str(ua) {
|
||||
headers.insert(USER_AGENT, value);
|
||||
}
|
||||
|
||||
// Platform-appropriate TLS backend — see [`crate::openhuman::tls`].
|
||||
let builder = crate::openhuman::tls::tls_client_builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.connect_timeout(std::time::Duration::from_secs(10))
|
||||
.default_headers(headers);
|
||||
let builder = crate::openhuman::config::apply_runtime_proxy_to_builder(
|
||||
builder,
|
||||
"provider.compatible",
|
||||
);
|
||||
|
||||
return builder.build().unwrap_or_else(|error| {
|
||||
tracing::warn!("Failed to build proxied timeout client with user-agent: {error}");
|
||||
crate::openhuman::tls::tls_client_builder()
|
||||
.build()
|
||||
.unwrap_or_default()
|
||||
});
|
||||
}
|
||||
|
||||
// Platform-appropriate TLS backend — see [`crate::openhuman::tls`].
|
||||
let builder = crate::openhuman::tls::tls_client_builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.connect_timeout(std::time::Duration::from_secs(10));
|
||||
let builder = crate::openhuman::config::apply_runtime_proxy_to_builder(
|
||||
builder,
|
||||
"provider.compatible",
|
||||
);
|
||||
builder.build().unwrap_or_else(|error| {
|
||||
tracing::warn!("Failed to build proxied timeout client: {error}");
|
||||
crate::openhuman::tls::tls_client_builder()
|
||||
.build()
|
||||
.unwrap_or_default()
|
||||
})
|
||||
}
|
||||
|
||||
/// Build the full URL for chat completions, detecting if base_url already includes the path.
|
||||
/// This allows custom providers with non-standard endpoints (e.g., VolcEngine ARK uses
|
||||
/// `/api/coding/v3/chat/completions` instead of `/v1/chat/completions`).
|
||||
fn chat_completions_url(&self) -> String {
|
||||
let has_full_endpoint = reqwest::Url::parse(&self.base_url)
|
||||
.map(|url| {
|
||||
url.path()
|
||||
.trim_end_matches('/')
|
||||
.ends_with("/chat/completions")
|
||||
})
|
||||
.unwrap_or_else(|_| {
|
||||
self.base_url
|
||||
.trim_end_matches('/')
|
||||
.ends_with("/chat/completions")
|
||||
});
|
||||
|
||||
let url = if has_full_endpoint {
|
||||
self.base_url.clone()
|
||||
} else {
|
||||
format!("{}/chat/completions", self.base_url)
|
||||
};
|
||||
log::info!(
|
||||
"[provider:{}] outbound chat/completions -> {}",
|
||||
self.name,
|
||||
url
|
||||
);
|
||||
url
|
||||
}
|
||||
|
||||
fn path_ends_with(&self, suffix: &str) -> bool {
|
||||
if let Ok(url) = reqwest::Url::parse(&self.base_url) {
|
||||
return url.path().trim_end_matches('/').ends_with(suffix);
|
||||
}
|
||||
|
||||
self.base_url.trim_end_matches('/').ends_with(suffix)
|
||||
}
|
||||
|
||||
fn has_explicit_api_path(&self) -> bool {
|
||||
let Ok(url) = reqwest::Url::parse(&self.base_url) else {
|
||||
return false;
|
||||
};
|
||||
|
||||
let path = url.path().trim_end_matches('/');
|
||||
!path.is_empty() && path != "/"
|
||||
}
|
||||
|
||||
/// Build the full URL for responses API, detecting if base_url already includes the path.
|
||||
fn responses_url(&self) -> String {
|
||||
if self.path_ends_with("/responses") {
|
||||
return self.base_url.clone();
|
||||
}
|
||||
|
||||
let normalized_base = self.base_url.trim_end_matches('/');
|
||||
|
||||
// If chat endpoint is explicitly configured, derive sibling responses endpoint.
|
||||
if let Some(prefix) = normalized_base.strip_suffix("/chat/completions") {
|
||||
return format!("{prefix}/responses");
|
||||
}
|
||||
|
||||
// If an explicit API path already exists (e.g. /v1, /openai, /api/coding/v3),
|
||||
// append responses directly to avoid duplicate /v1 segments.
|
||||
if self.has_explicit_api_path() {
|
||||
format!("{normalized_base}/responses")
|
||||
} else {
|
||||
format!("{normalized_base}/v1/responses")
|
||||
}
|
||||
}
|
||||
|
||||
fn tool_specs_to_openai_format(
|
||||
tools: &[crate::openhuman::tools::ToolSpec],
|
||||
) -> Vec<serde_json::Value> {
|
||||
tools
|
||||
.iter()
|
||||
.map(|tool| {
|
||||
serde_json::json!({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tool.name,
|
||||
"description": tool.description,
|
||||
"parameters": tool.parameters
|
||||
}
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn credential_for_request(&self) -> anyhow::Result<Option<&str>> {
|
||||
if matches!(&self.auth_header, AuthStyle::None) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
self.credential
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|credential| !credential.is_empty())
|
||||
.map(Some)
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"{} API key not set. Configure via the web UI or set the appropriate env var.",
|
||||
self.name
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn apply_auth_header(
|
||||
&self,
|
||||
req: reqwest::RequestBuilder,
|
||||
credential: Option<&str>,
|
||||
) -> reqwest::RequestBuilder {
|
||||
match (&self.auth_header, credential) {
|
||||
(AuthStyle::None, _) => req,
|
||||
(_, None) => req,
|
||||
(AuthStyle::Bearer, Some(credential)) => {
|
||||
req.header("Authorization", format!("Bearer {credential}"))
|
||||
}
|
||||
(AuthStyle::XApiKey, Some(credential)) => req.header("x-api-key", credential),
|
||||
(AuthStyle::Anthropic, Some(credential)) => req
|
||||
.header("x-api-key", credential)
|
||||
.header("anthropic-version", "2023-06-01"),
|
||||
(AuthStyle::Custom(header), Some(credential)) => req.header(header, credential),
|
||||
}
|
||||
}
|
||||
|
||||
async fn chat_via_responses(
|
||||
&self,
|
||||
credential: Option<&str>,
|
||||
|
||||
@@ -0,0 +1,239 @@
|
||||
use crate::openhuman::inference::provider::traits::ChatMessage;
|
||||
use reqwest::{
|
||||
header::{HeaderMap, HeaderValue, USER_AGENT},
|
||||
Client,
|
||||
};
|
||||
|
||||
use crate::openhuman::inference::provider::{temperature, thread_context};
|
||||
|
||||
use super::{AuthStyle, OpenAiCompatibleProvider};
|
||||
|
||||
impl OpenAiCompatibleProvider {
|
||||
/// Resolve the effective temperature for `model`. Returns `None` when the
|
||||
/// model matches a pattern in `temperature_unsupported_models` (causing the
|
||||
/// field to be omitted from the serialised request). Otherwise yields the
|
||||
/// per-workload override if one was configured, else the caller's value.
|
||||
pub(super) fn effective_temperature(&self, model: &str, temperature: f64) -> Option<f64> {
|
||||
if self
|
||||
.temperature_unsupported_models
|
||||
.iter()
|
||||
.any(|pat| temperature::glob_match(pat, model))
|
||||
{
|
||||
tracing::debug!(
|
||||
"[provider:{}] model='{}' matched temperature_unsupported_models — omitting temperature",
|
||||
self.name,
|
||||
model
|
||||
);
|
||||
None
|
||||
} else {
|
||||
Some(self.temperature_override.unwrap_or(temperature))
|
||||
}
|
||||
}
|
||||
|
||||
/// Read the ambient `thread_id` only when this provider has been
|
||||
/// opted in via [`with_openhuman_thread_id`]. Returns `None` for
|
||||
/// every third-party provider so the field is omitted by
|
||||
/// `skip_serializing_if`.
|
||||
pub(super) fn outbound_thread_id(&self) -> Option<String> {
|
||||
if self.emit_openhuman_thread_id {
|
||||
thread_context::current_thread_id()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Collect all `system` role messages, concatenate their content,
|
||||
/// and prepend to the first `user` message. Drop all system messages.
|
||||
/// Used for providers (e.g. MiniMax) that reject `role: system`.
|
||||
pub(super) fn flatten_system_messages(messages: &[ChatMessage]) -> Vec<ChatMessage> {
|
||||
let system_content: String = messages
|
||||
.iter()
|
||||
.filter(|m| m.role == "system")
|
||||
.map(|m| m.content.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n");
|
||||
|
||||
if system_content.is_empty() {
|
||||
return messages.to_vec();
|
||||
}
|
||||
|
||||
let mut result: Vec<ChatMessage> = messages
|
||||
.iter()
|
||||
.filter(|m| m.role != "system")
|
||||
.cloned()
|
||||
.collect();
|
||||
|
||||
if let Some(first_user) = result.iter_mut().find(|m| m.role == "user") {
|
||||
first_user.content = format!("{system_content}\n\n{}", first_user.content);
|
||||
} else {
|
||||
// No user message found: insert a synthetic user message with system content
|
||||
result.insert(0, ChatMessage::user(&system_content));
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
pub(super) fn http_client(&self) -> Client {
|
||||
if let Some(ua) = self.user_agent.as_deref() {
|
||||
let mut headers = HeaderMap::new();
|
||||
if let Ok(value) = HeaderValue::from_str(ua) {
|
||||
headers.insert(USER_AGENT, value);
|
||||
}
|
||||
|
||||
// Platform-appropriate TLS backend — see [`crate::openhuman::tls`].
|
||||
let builder = crate::openhuman::tls::tls_client_builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.connect_timeout(std::time::Duration::from_secs(10))
|
||||
.default_headers(headers);
|
||||
let builder = crate::openhuman::config::apply_runtime_proxy_to_builder(
|
||||
builder,
|
||||
"provider.compatible",
|
||||
);
|
||||
|
||||
return builder.build().unwrap_or_else(|error| {
|
||||
tracing::warn!("Failed to build proxied timeout client with user-agent: {error}");
|
||||
crate::openhuman::tls::tls_client_builder()
|
||||
.build()
|
||||
.unwrap_or_default()
|
||||
});
|
||||
}
|
||||
|
||||
// Platform-appropriate TLS backend — see [`crate::openhuman::tls`].
|
||||
let builder = crate::openhuman::tls::tls_client_builder()
|
||||
.timeout(std::time::Duration::from_secs(120))
|
||||
.connect_timeout(std::time::Duration::from_secs(10));
|
||||
let builder = crate::openhuman::config::apply_runtime_proxy_to_builder(
|
||||
builder,
|
||||
"provider.compatible",
|
||||
);
|
||||
builder.build().unwrap_or_else(|error| {
|
||||
tracing::warn!("Failed to build proxied timeout client: {error}");
|
||||
crate::openhuman::tls::tls_client_builder()
|
||||
.build()
|
||||
.unwrap_or_default()
|
||||
})
|
||||
}
|
||||
|
||||
/// Build the full URL for chat completions, detecting if base_url already includes the path.
|
||||
/// This allows custom providers with non-standard endpoints (e.g., VolcEngine ARK uses
|
||||
/// `/api/coding/v3/chat/completions` instead of `/v1/chat/completions`).
|
||||
pub(super) fn chat_completions_url(&self) -> String {
|
||||
let has_full_endpoint = reqwest::Url::parse(&self.base_url)
|
||||
.map(|url| {
|
||||
url.path()
|
||||
.trim_end_matches('/')
|
||||
.ends_with("/chat/completions")
|
||||
})
|
||||
.unwrap_or_else(|_| {
|
||||
self.base_url
|
||||
.trim_end_matches('/')
|
||||
.ends_with("/chat/completions")
|
||||
});
|
||||
|
||||
let url = if has_full_endpoint {
|
||||
self.base_url.clone()
|
||||
} else {
|
||||
format!("{}/chat/completions", self.base_url)
|
||||
};
|
||||
log::info!(
|
||||
"[provider:{}] outbound chat/completions -> {}",
|
||||
self.name,
|
||||
url
|
||||
);
|
||||
url
|
||||
}
|
||||
|
||||
pub(super) fn path_ends_with(&self, suffix: &str) -> bool {
|
||||
if let Ok(url) = reqwest::Url::parse(&self.base_url) {
|
||||
return url.path().trim_end_matches('/').ends_with(suffix);
|
||||
}
|
||||
|
||||
self.base_url.trim_end_matches('/').ends_with(suffix)
|
||||
}
|
||||
|
||||
pub(super) fn has_explicit_api_path(&self) -> bool {
|
||||
let Ok(url) = reqwest::Url::parse(&self.base_url) else {
|
||||
return false;
|
||||
};
|
||||
|
||||
let path = url.path().trim_end_matches('/');
|
||||
!path.is_empty() && path != "/"
|
||||
}
|
||||
|
||||
/// Build the full URL for responses API, detecting if base_url already includes the path.
|
||||
pub(super) fn responses_url(&self) -> String {
|
||||
if self.path_ends_with("/responses") {
|
||||
return self.base_url.clone();
|
||||
}
|
||||
|
||||
let normalized_base = self.base_url.trim_end_matches('/');
|
||||
|
||||
// If chat endpoint is explicitly configured, derive sibling responses endpoint.
|
||||
if let Some(prefix) = normalized_base.strip_suffix("/chat/completions") {
|
||||
return format!("{prefix}/responses");
|
||||
}
|
||||
|
||||
// If an explicit API path already exists (e.g. /v1, /openai, /api/coding/v3),
|
||||
// append responses directly to avoid duplicate /v1 segments.
|
||||
if self.has_explicit_api_path() {
|
||||
format!("{normalized_base}/responses")
|
||||
} else {
|
||||
format!("{normalized_base}/v1/responses")
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn tool_specs_to_openai_format(
|
||||
tools: &[crate::openhuman::tools::ToolSpec],
|
||||
) -> Vec<serde_json::Value> {
|
||||
tools
|
||||
.iter()
|
||||
.map(|tool| {
|
||||
serde_json::json!({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tool.name,
|
||||
"description": tool.description,
|
||||
"parameters": tool.parameters
|
||||
}
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(super) fn credential_for_request(&self) -> anyhow::Result<Option<&str>> {
|
||||
if matches!(&self.auth_header, AuthStyle::None) {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
self.credential
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|credential| !credential.is_empty())
|
||||
.map(Some)
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"{} API key not set. Configure via the web UI or set the appropriate env var.",
|
||||
self.name
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn apply_auth_header(
|
||||
&self,
|
||||
req: reqwest::RequestBuilder,
|
||||
credential: Option<&str>,
|
||||
) -> reqwest::RequestBuilder {
|
||||
match (&self.auth_header, credential) {
|
||||
(AuthStyle::None, _) => req,
|
||||
(_, None) => req,
|
||||
(AuthStyle::Bearer, Some(credential)) => {
|
||||
req.header("Authorization", format!("Bearer {credential}"))
|
||||
}
|
||||
(AuthStyle::XApiKey, Some(credential)) => req.header("x-api-key", credential),
|
||||
(AuthStyle::Anthropic, Some(credential)) => req
|
||||
.header("x-api-key", credential)
|
||||
.header("anthropic-version", "2023-06-01"),
|
||||
(AuthStyle::Custom(header), Some(credential)) => req.header(header, credential),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -3,7 +3,9 @@ use crate::openhuman::skills::types::ToolResult;
|
||||
use anyhow::Context;
|
||||
use base64::Engine;
|
||||
use parking_lot::Mutex;
|
||||
use reqwest::header::{HeaderMap, HeaderName, HeaderValue, ACCEPT, AUTHORIZATION, CONTENT_TYPE};
|
||||
#[cfg(test)]
|
||||
use reqwest::header::HeaderMap;
|
||||
use reqwest::header::{HeaderName, HeaderValue, ACCEPT, AUTHORIZATION, CONTENT_TYPE};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
@@ -647,132 +649,12 @@ pub fn redact_endpoint(raw: &str) -> String {
|
||||
format!("{scheme}://{authority}")
|
||||
}
|
||||
|
||||
fn parse_sse_message(body: &str) -> anyhow::Result<Value> {
|
||||
let events = parse_sse_events(body)?;
|
||||
let event = events
|
||||
.into_iter()
|
||||
.find_map(|event| event.data)
|
||||
.ok_or_else(|| anyhow::anyhow!("No SSE data frame found in MCP response: {body}"))?;
|
||||
Ok(event)
|
||||
}
|
||||
|
||||
fn parse_sse_events(body: &str) -> anyhow::Result<Vec<McpSseEvent>> {
|
||||
let mut events = Vec::new();
|
||||
let mut event_type: Option<String> = None;
|
||||
let mut event_id: Option<String> = None;
|
||||
let mut data_lines: Vec<String> = Vec::new();
|
||||
|
||||
let flush = |events: &mut Vec<McpSseEvent>,
|
||||
event_type: &mut Option<String>,
|
||||
event_id: &mut Option<String>,
|
||||
data_lines: &mut Vec<String>|
|
||||
-> anyhow::Result<()> {
|
||||
if event_type.is_none() && event_id.is_none() && data_lines.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let data = if data_lines.is_empty() {
|
||||
None
|
||||
} else {
|
||||
let joined = data_lines.join("\n");
|
||||
Some(
|
||||
serde_json::from_str(&joined)
|
||||
.with_context(|| format!("Failed to parse SSE data frame JSON: {joined}"))?,
|
||||
)
|
||||
};
|
||||
events.push(McpSseEvent {
|
||||
event: event_type.take(),
|
||||
id: event_id.take(),
|
||||
data,
|
||||
});
|
||||
data_lines.clear();
|
||||
Ok(())
|
||||
};
|
||||
|
||||
for raw_line in body.lines() {
|
||||
let line = raw_line.trim_end_matches('\r');
|
||||
if line.is_empty() {
|
||||
flush(&mut events, &mut event_type, &mut event_id, &mut data_lines)?;
|
||||
continue;
|
||||
}
|
||||
if line.starts_with(':') {
|
||||
continue;
|
||||
}
|
||||
if let Some(value) = line.strip_prefix("event:") {
|
||||
event_type = Some(value.trim_start().to_string());
|
||||
} else if let Some(value) = line.strip_prefix("id:") {
|
||||
event_id = Some(value.trim_start().to_string());
|
||||
} else if let Some(value) = line.strip_prefix("data:") {
|
||||
data_lines.push(value.trim_start().to_string());
|
||||
}
|
||||
}
|
||||
flush(&mut events, &mut event_type, &mut event_id, &mut data_lines)?;
|
||||
Ok(events)
|
||||
}
|
||||
|
||||
fn parse_www_authenticate_challenge(headers: &HeaderMap) -> Option<McpAuthChallenge> {
|
||||
let raw = headers.get("WWW-Authenticate")?.to_str().ok()?.trim();
|
||||
let mut parts = raw.splitn(2, ' ');
|
||||
let scheme = parts.next()?.trim().to_string();
|
||||
let params = parts.next().unwrap_or("").trim();
|
||||
let attrs = parse_auth_attribute_list(params);
|
||||
Some(McpAuthChallenge {
|
||||
scheme,
|
||||
realm: attrs.get("realm").cloned(),
|
||||
resource_metadata: attrs.get("resource_metadata").cloned(),
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_auth_attribute_list(input: &str) -> HashMap<String, String> {
|
||||
let mut attrs = HashMap::new();
|
||||
for part in input.split(',') {
|
||||
let Some((key, value)) = part.split_once('=') else {
|
||||
continue;
|
||||
};
|
||||
let value = value.trim().trim_matches('"').to_string();
|
||||
attrs.insert(key.trim().to_string(), value);
|
||||
}
|
||||
attrs
|
||||
}
|
||||
|
||||
fn header_to_string(headers: &HeaderMap, name: &str) -> Option<String> {
|
||||
headers.get(name)?.to_str().ok().map(|s| s.to_string())
|
||||
}
|
||||
|
||||
fn x_mcp_headers_from_schema(
|
||||
tool: &McpRemoteTool,
|
||||
arguments: &Value,
|
||||
) -> anyhow::Result<Vec<(HeaderName, HeaderValue)>> {
|
||||
let mut headers = Vec::new();
|
||||
let Some(args) = arguments.as_object() else {
|
||||
return Ok(headers);
|
||||
};
|
||||
let properties = tool
|
||||
.input_schema
|
||||
.get("properties")
|
||||
.and_then(Value::as_object)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
|
||||
for (param_name, schema) in properties {
|
||||
let Some(header_suffix) = schema.get("x-mcp-header").and_then(Value::as_str) else {
|
||||
continue;
|
||||
};
|
||||
let Some(value) = args.get(¶m_name) else {
|
||||
continue;
|
||||
};
|
||||
let header_name =
|
||||
HeaderName::from_bytes(format!("Mcp-Param-{header_suffix}").as_bytes())
|
||||
.with_context(|| format!("invalid x-mcp-header name for `{param_name}`"))?;
|
||||
let header_value = match value {
|
||||
Value::String(s) => HeaderValue::from_str(s),
|
||||
other => HeaderValue::from_str(&other.to_string()),
|
||||
}
|
||||
.with_context(|| format!("invalid x-mcp-header value for `{param_name}`"))?;
|
||||
headers.push((header_name, header_value));
|
||||
}
|
||||
|
||||
Ok(headers)
|
||||
}
|
||||
#[path = "client_helpers.rs"]
|
||||
mod client_helpers;
|
||||
use client_helpers::{
|
||||
header_to_string, parse_sse_events, parse_sse_message, parse_www_authenticate_challenge,
|
||||
x_mcp_headers_from_schema,
|
||||
};
|
||||
|
||||
impl McpHttpClient {
|
||||
async fn read_response(&self, response: reqwest::Response) -> anyhow::Result<ResponseEnvelope> {
|
||||
@@ -828,479 +710,5 @@ impl McpHttpClient {
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use axum::{
|
||||
extract::State,
|
||||
http::{HeaderMap as AxumHeaderMap, Method, StatusCode},
|
||||
response::{IntoResponse, Response},
|
||||
routing::{get, post},
|
||||
Json, Router,
|
||||
};
|
||||
use serde_json::Value;
|
||||
use std::sync::{
|
||||
atomic::{AtomicUsize, Ordering as AtomicOrdering},
|
||||
Arc,
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
struct TestState {
|
||||
init_count: Arc<AtomicUsize>,
|
||||
call_count: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
fn has_streamable_http_accept(headers: &AxumHeaderMap) -> bool {
|
||||
headers
|
||||
.get(ACCEPT)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(|value| value.contains("application/json") && value.contains("text/event-stream"))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
async fn mcp_handler(
|
||||
State(state): State<TestState>,
|
||||
headers: AxumHeaderMap,
|
||||
method: Method,
|
||||
Json(body): Json<Value>,
|
||||
) -> Response {
|
||||
if method == Method::POST && !has_streamable_http_accept(&headers) {
|
||||
return (
|
||||
StatusCode::NOT_ACCEPTABLE,
|
||||
"missing MCP Accept header".to_string(),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
let rpc_method = body.get("method").and_then(Value::as_str).unwrap_or("");
|
||||
if method == Method::POST && rpc_method == "initialize" {
|
||||
state.init_count.fetch_add(1, AtomicOrdering::SeqCst);
|
||||
return (
|
||||
[(HEADER_SESSION_ID, "session-1")],
|
||||
Json(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"].clone(),
|
||||
"result": {
|
||||
"protocolVersion": LATEST_PROTOCOL_VERSION,
|
||||
"capabilities": { "tools": { "listChanged": true } },
|
||||
"serverInfo": { "name": "test-server", "version": "1.0.0" }
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
if headers.get(HEADER_SESSION_ID).and_then(|v| v.to_str().ok()) != Some("session-1") {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
"missing or invalid session".to_string(),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
if headers
|
||||
.get(HEADER_PROTOCOL_VERSION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
!= Some(LATEST_PROTOCOL_VERSION)
|
||||
{
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
"missing protocol version".to_string(),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
match rpc_method {
|
||||
"notifications/initialized" => StatusCode::NO_CONTENT.into_response(),
|
||||
"tools/list" => Json(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"].clone(),
|
||||
"result": {
|
||||
"tools": [{
|
||||
"name": "needs_header",
|
||||
"description": "needs x-mcp-header",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"tenant": {
|
||||
"type": "string",
|
||||
"x-mcp-header": "tenant"
|
||||
}
|
||||
}
|
||||
}
|
||||
}]
|
||||
}
|
||||
}))
|
||||
.into_response(),
|
||||
"tools/call" => {
|
||||
state.call_count.fetch_add(1, AtomicOrdering::SeqCst);
|
||||
if headers
|
||||
.get("Mcp-Param-tenant")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
!= Some("acme")
|
||||
{
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
"missing mirrored tenant header".to_string(),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
Json(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"].clone(),
|
||||
"result": {
|
||||
"content": [{
|
||||
"type": "text",
|
||||
"text": "remote result"
|
||||
}]
|
||||
}
|
||||
}))
|
||||
.into_response()
|
||||
}
|
||||
_ => (
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!("unexpected method {rpc_method}"),
|
||||
)
|
||||
.into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn events_handler(headers: AxumHeaderMap) -> Response {
|
||||
if headers
|
||||
.get(ACCEPT)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.filter(|value| value.contains("text/event-stream"))
|
||||
.is_none()
|
||||
{
|
||||
return (
|
||||
StatusCode::NOT_ACCEPTABLE,
|
||||
"missing SSE Accept header".to_string(),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
if headers.get(HEADER_SESSION_ID).is_none() {
|
||||
return (StatusCode::BAD_REQUEST, "no session".to_string()).into_response();
|
||||
}
|
||||
(
|
||||
[(CONTENT_TYPE.as_str(), "text/event-stream")],
|
||||
"id: 1\nevent: message\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/message\",\"params\":{\"ok\":true}}\n\n",
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
async fn delete_handler() -> Response {
|
||||
StatusCode::NO_CONTENT.into_response()
|
||||
}
|
||||
|
||||
async fn bearer_required_handler(headers: AxumHeaderMap, Json(body): Json<Value>) -> Response {
|
||||
if headers.get(AUTHORIZATION).and_then(|v| v.to_str().ok()) != Some("Bearer secret-token") {
|
||||
return (StatusCode::UNAUTHORIZED, "missing bearer".to_string()).into_response();
|
||||
}
|
||||
Json(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"].clone(),
|
||||
"result": {
|
||||
"protocolVersion": LATEST_PROTOCOL_VERSION,
|
||||
"capabilities": {},
|
||||
"serverInfo": { "name": "bearer-server", "version": "1.0.0" }
|
||||
}
|
||||
}))
|
||||
.into_response()
|
||||
}
|
||||
|
||||
async fn retrying_mcp_handler(
|
||||
State(state): State<TestState>,
|
||||
headers: AxumHeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> Response {
|
||||
if !has_streamable_http_accept(&headers) {
|
||||
return (
|
||||
StatusCode::NOT_ACCEPTABLE,
|
||||
"missing MCP Accept header".to_string(),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
let rpc_method = body.get("method").and_then(Value::as_str).unwrap_or("");
|
||||
if rpc_method == "initialize" {
|
||||
state.init_count.fetch_add(1, AtomicOrdering::SeqCst);
|
||||
return (
|
||||
[(HEADER_SESSION_ID, "session-retry")],
|
||||
Json(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"].clone(),
|
||||
"result": {
|
||||
"protocolVersion": LATEST_PROTOCOL_VERSION,
|
||||
"capabilities": { "tools": {} },
|
||||
"serverInfo": { "name": "retry-server", "version": "1.0.0" }
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
if rpc_method == "notifications/initialized" {
|
||||
return StatusCode::NO_CONTENT.into_response();
|
||||
}
|
||||
if rpc_method == "tools/list" {
|
||||
let call_number = state.call_count.fetch_add(1, AtomicOrdering::SeqCst);
|
||||
if call_number == 0
|
||||
&& headers.get(HEADER_SESSION_ID).and_then(|v| v.to_str().ok())
|
||||
== Some("session-retry")
|
||||
{
|
||||
return (StatusCode::NOT_FOUND, "expired".to_string()).into_response();
|
||||
}
|
||||
return Json(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"].clone(),
|
||||
"result": { "tools": [] }
|
||||
}))
|
||||
.into_response();
|
||||
}
|
||||
(StatusCode::BAD_REQUEST, "unexpected".to_string()).into_response()
|
||||
}
|
||||
|
||||
async fn spawn_test_server() -> (String, TestState) {
|
||||
let state = TestState {
|
||||
init_count: Arc::new(AtomicUsize::new(0)),
|
||||
call_count: Arc::new(AtomicUsize::new(0)),
|
||||
};
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/",
|
||||
post(mcp_handler).get(events_handler).delete(delete_handler),
|
||||
)
|
||||
.with_state(state.clone());
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
(format!("http://{addr}/"), state)
|
||||
}
|
||||
|
||||
async fn spawn_retry_server() -> (String, TestState) {
|
||||
let state = TestState {
|
||||
init_count: Arc::new(AtomicUsize::new(0)),
|
||||
call_count: Arc::new(AtomicUsize::new(0)),
|
||||
};
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let app = Router::new()
|
||||
.route("/", post(retrying_mcp_handler))
|
||||
.with_state(state.clone());
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
(format!("http://{addr}/"), state)
|
||||
}
|
||||
|
||||
async fn spawn_discovery_server() -> String {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let base = format!("http://{addr}");
|
||||
let auth_header = format!(
|
||||
"Bearer realm=\"mcp\", resource_metadata=\"{base}/.well-known/oauth-protected-resource\""
|
||||
);
|
||||
let prm_base = base.clone();
|
||||
let issuer_base = base.clone();
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/",
|
||||
post(move || {
|
||||
let auth_header = auth_header.clone();
|
||||
async move {
|
||||
(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
[("WWW-Authenticate", auth_header.as_str())],
|
||||
"",
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/.well-known/oauth-protected-resource",
|
||||
get(move || {
|
||||
let prm_base = prm_base.clone();
|
||||
async move {
|
||||
let resource = format!("{prm_base}/");
|
||||
Json(json!({
|
||||
"resource": resource,
|
||||
"authorization_servers": [prm_base],
|
||||
"scopes_supported": ["mcp:tools"]
|
||||
}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/.well-known/openid-configuration",
|
||||
get(move || {
|
||||
let issuer_base = issuer_base.clone();
|
||||
async move {
|
||||
let authorization_endpoint = format!("{}/authorize", issuer_base);
|
||||
let token_endpoint = format!("{}/token", issuer_base);
|
||||
Json(json!({
|
||||
"issuer": issuer_base,
|
||||
"authorization_endpoint": authorization_endpoint,
|
||||
"token_endpoint": token_endpoint,
|
||||
"grant_types_supported": ["authorization_code"],
|
||||
"code_challenge_methods_supported": ["S256"]
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
format!("http://{addr}/")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn initialize_and_list_tools_negotiate_session() {
|
||||
let (endpoint, state) = spawn_test_server().await;
|
||||
let client = McpHttpClient::new(endpoint, 5);
|
||||
let tools = client.list_tools().await.expect("list_tools");
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(state.init_count.load(AtomicOrdering::SeqCst), 1);
|
||||
let snapshot = client.initialize_snapshot().expect("snapshot");
|
||||
assert_eq!(snapshot.protocol_version, LATEST_PROTOCOL_VERSION);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_tool_mirrors_x_mcp_header_parameters() {
|
||||
let (endpoint, state) = spawn_test_server().await;
|
||||
let client = McpHttpClient::new(endpoint, 5);
|
||||
let result = client
|
||||
.call_tool("needs_header", json!({"tenant": "acme"}))
|
||||
.await
|
||||
.expect("call_tool");
|
||||
assert_eq!(result.rendered.output(), "remote result");
|
||||
assert_eq!(state.call_count.load(AtomicOrdering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn session_404_triggers_reinitialize_and_retry() {
|
||||
let (endpoint, state) = spawn_retry_server().await;
|
||||
let client = McpHttpClient::new(endpoint, 5);
|
||||
let tools = client.list_tools().await.expect("list_tools");
|
||||
assert!(tools.is_empty());
|
||||
assert_eq!(state.init_count.load(AtomicOrdering::SeqCst), 2);
|
||||
assert_eq!(state.call_count.load(AtomicOrdering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn drain_events_parses_sse_stream() {
|
||||
let (endpoint, _) = spawn_test_server().await;
|
||||
let client = McpHttpClient::new(endpoint, 5);
|
||||
let events = client.drain_events(None).await.expect("drain events");
|
||||
assert_eq!(events.len(), 1);
|
||||
assert_eq!(events[0].id.as_deref(), Some("1"));
|
||||
assert_eq!(events[0].event.as_deref(), Some("message"));
|
||||
assert_eq!(events[0].data.as_ref().unwrap()["params"]["ok"], true);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn close_session_sends_delete() {
|
||||
let (endpoint, _) = spawn_test_server().await;
|
||||
let client = McpHttpClient::new(endpoint, 5);
|
||||
client.initialize().await.expect("initialize");
|
||||
client.close_session().await.expect("close_session");
|
||||
assert!(client.initialize_snapshot().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_endpoint_hides_paths_and_credentials() {
|
||||
assert_eq!(
|
||||
redact_endpoint("https://example.com/path?x=1"),
|
||||
"https://example.com"
|
||||
);
|
||||
assert_eq!(
|
||||
redact_endpoint("https://user:pw@example.com/a"),
|
||||
"<redacted>"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_sse_events_handles_multiple_frames() {
|
||||
let body = "id: 1\nevent: message\ndata: {\"a\":1}\n\ndata: {\"b\":2}\n\n";
|
||||
let events = parse_sse_events(body).expect("events");
|
||||
assert_eq!(events.len(), 2);
|
||||
assert_eq!(events[0].id.as_deref(), Some("1"));
|
||||
assert_eq!(events[1].data.as_ref().unwrap()["b"], 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_www_authenticate_extracts_resource_metadata() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
"WWW-Authenticate",
|
||||
HeaderValue::from_static(
|
||||
"Bearer realm=\"mcp\", resource_metadata=\"https://example.com/.well-known/oauth-protected-resource\"",
|
||||
),
|
||||
);
|
||||
let challenge = parse_www_authenticate_challenge(&headers).expect("challenge");
|
||||
assert_eq!(challenge.scheme, "Bearer");
|
||||
assert_eq!(challenge.realm.as_deref(), Some("mcp"));
|
||||
assert_eq!(
|
||||
challenge.resource_metadata.as_deref(),
|
||||
Some("https://example.com/.well-known/oauth-protected-resource")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn discover_authorization_returns_none_when_not_401() {
|
||||
let (endpoint, _) = spawn_test_server().await;
|
||||
let client = McpHttpClient::new(endpoint, 5);
|
||||
assert!(client.discover_authorization().await.unwrap().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn discover_authorization_fetches_metadata() {
|
||||
let endpoint = spawn_discovery_server().await;
|
||||
let client = McpHttpClient::new(endpoint, 2);
|
||||
let ctx = client
|
||||
.discover_authorization()
|
||||
.await
|
||||
.expect("discover")
|
||||
.expect("some");
|
||||
assert_eq!(ctx.challenge.scheme, "Bearer");
|
||||
assert_eq!(
|
||||
ctx.protected_resource_metadata
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.scopes_supported,
|
||||
vec!["mcp:tools"]
|
||||
);
|
||||
assert_eq!(ctx.authorization_server_metadata.len(), 1);
|
||||
let expected_authorization_endpoint = format!(
|
||||
"{}/authorize",
|
||||
ctx.protected_resource_metadata
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.authorization_servers[0]
|
||||
);
|
||||
assert_eq!(
|
||||
ctx.authorization_server_metadata[0]
|
||||
.authorization_endpoint
|
||||
.as_deref(),
|
||||
Some(expected_authorization_endpoint.as_str())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bearer_auth_is_attached_to_initialize() {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let app = Router::new().route("/", post(bearer_required_handler));
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
let client = McpHttpClient::with_options(
|
||||
format!("http://{addr}/"),
|
||||
2,
|
||||
McpAuthConfig::BearerToken {
|
||||
token: "secret-token".into(),
|
||||
},
|
||||
McpClientIdentityConfig::default(),
|
||||
);
|
||||
let init = client.initialize().await.expect("initialize");
|
||||
assert_eq!(init.server_info["name"], "bearer-server");
|
||||
}
|
||||
}
|
||||
#[path = "client_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
use super::{McpAuthChallenge, McpRemoteTool, McpSseEvent};
|
||||
use anyhow::Context;
|
||||
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
|
||||
pub(super) fn parse_sse_message(body: &str) -> anyhow::Result<Value> {
|
||||
let events = parse_sse_events(body)?;
|
||||
let event = events
|
||||
.into_iter()
|
||||
.find_map(|event| event.data)
|
||||
.ok_or_else(|| anyhow::anyhow!("No SSE data frame found in MCP response: {body}"))?;
|
||||
Ok(event)
|
||||
}
|
||||
|
||||
pub(super) fn parse_sse_events(body: &str) -> anyhow::Result<Vec<McpSseEvent>> {
|
||||
let mut events = Vec::new();
|
||||
let mut event_type: Option<String> = None;
|
||||
let mut event_id: Option<String> = None;
|
||||
let mut data_lines: Vec<String> = Vec::new();
|
||||
|
||||
let flush = |events: &mut Vec<McpSseEvent>,
|
||||
event_type: &mut Option<String>,
|
||||
event_id: &mut Option<String>,
|
||||
data_lines: &mut Vec<String>|
|
||||
-> anyhow::Result<()> {
|
||||
if event_type.is_none() && event_id.is_none() && data_lines.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let data = if data_lines.is_empty() {
|
||||
None
|
||||
} else {
|
||||
let joined = data_lines.join("\n");
|
||||
Some(
|
||||
serde_json::from_str(&joined)
|
||||
.with_context(|| format!("Failed to parse SSE data frame JSON: {joined}"))?,
|
||||
)
|
||||
};
|
||||
events.push(McpSseEvent {
|
||||
event: event_type.take(),
|
||||
id: event_id.take(),
|
||||
data,
|
||||
});
|
||||
data_lines.clear();
|
||||
Ok(())
|
||||
};
|
||||
|
||||
for raw_line in body.lines() {
|
||||
let line = raw_line.trim_end_matches('\r');
|
||||
if line.is_empty() {
|
||||
flush(&mut events, &mut event_type, &mut event_id, &mut data_lines)?;
|
||||
continue;
|
||||
}
|
||||
if line.starts_with(':') {
|
||||
continue;
|
||||
}
|
||||
if let Some(value) = line.strip_prefix("event:") {
|
||||
event_type = Some(value.trim_start().to_string());
|
||||
} else if let Some(value) = line.strip_prefix("id:") {
|
||||
event_id = Some(value.trim_start().to_string());
|
||||
} else if let Some(value) = line.strip_prefix("data:") {
|
||||
data_lines.push(value.trim_start().to_string());
|
||||
}
|
||||
}
|
||||
flush(&mut events, &mut event_type, &mut event_id, &mut data_lines)?;
|
||||
Ok(events)
|
||||
}
|
||||
|
||||
pub(super) fn parse_www_authenticate_challenge(headers: &HeaderMap) -> Option<McpAuthChallenge> {
|
||||
let raw = headers.get("WWW-Authenticate")?.to_str().ok()?.trim();
|
||||
let mut parts = raw.splitn(2, ' ');
|
||||
let scheme = parts.next()?.trim().to_string();
|
||||
let params = parts.next().unwrap_or("").trim();
|
||||
let attrs = parse_auth_attribute_list(params);
|
||||
Some(McpAuthChallenge {
|
||||
scheme,
|
||||
realm: attrs.get("realm").cloned(),
|
||||
resource_metadata: attrs.get("resource_metadata").cloned(),
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn parse_auth_attribute_list(input: &str) -> HashMap<String, String> {
|
||||
let mut attrs = HashMap::new();
|
||||
for part in input.split(',') {
|
||||
let Some((key, value)) = part.split_once('=') else {
|
||||
continue;
|
||||
};
|
||||
let value = value.trim().trim_matches('"').to_string();
|
||||
attrs.insert(key.trim().to_string(), value);
|
||||
}
|
||||
attrs
|
||||
}
|
||||
|
||||
pub(super) fn header_to_string(headers: &HeaderMap, name: &str) -> Option<String> {
|
||||
headers.get(name)?.to_str().ok().map(|s| s.to_string())
|
||||
}
|
||||
|
||||
pub(super) fn x_mcp_headers_from_schema(
|
||||
tool: &McpRemoteTool,
|
||||
arguments: &Value,
|
||||
) -> anyhow::Result<Vec<(HeaderName, HeaderValue)>> {
|
||||
let mut headers = Vec::new();
|
||||
let Some(args) = arguments.as_object() else {
|
||||
return Ok(headers);
|
||||
};
|
||||
let properties = tool
|
||||
.input_schema
|
||||
.get("properties")
|
||||
.and_then(Value::as_object)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
|
||||
for (param_name, schema) in properties {
|
||||
let Some(header_suffix) = schema.get("x-mcp-header").and_then(Value::as_str) else {
|
||||
continue;
|
||||
};
|
||||
let Some(value) = args.get(¶m_name) else {
|
||||
continue;
|
||||
};
|
||||
let header_name =
|
||||
HeaderName::from_bytes(format!("Mcp-Param-{header_suffix}").as_bytes())
|
||||
.with_context(|| format!("invalid x-mcp-header name for `{param_name}`"))?;
|
||||
let header_value = match value {
|
||||
Value::String(s) => HeaderValue::from_str(s),
|
||||
other => HeaderValue::from_str(&other.to_string()),
|
||||
}
|
||||
.with_context(|| format!("invalid x-mcp-header value for `{param_name}`"))?;
|
||||
headers.push((header_name, header_value));
|
||||
}
|
||||
|
||||
Ok(headers)
|
||||
}
|
||||
@@ -0,0 +1,473 @@
|
||||
use super::*;
|
||||
use axum::{
|
||||
extract::State,
|
||||
http::{HeaderMap as AxumHeaderMap, Method, StatusCode},
|
||||
response::{IntoResponse, Response},
|
||||
routing::{get, post},
|
||||
Json, Router,
|
||||
};
|
||||
use serde_json::Value;
|
||||
use std::sync::{
|
||||
atomic::{AtomicUsize, Ordering as AtomicOrdering},
|
||||
Arc,
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
struct TestState {
|
||||
init_count: Arc<AtomicUsize>,
|
||||
call_count: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
fn has_streamable_http_accept(headers: &AxumHeaderMap) -> bool {
|
||||
headers
|
||||
.get(ACCEPT)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(|value| value.contains("application/json") && value.contains("text/event-stream"))
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
async fn mcp_handler(
|
||||
State(state): State<TestState>,
|
||||
headers: AxumHeaderMap,
|
||||
method: Method,
|
||||
Json(body): Json<Value>,
|
||||
) -> Response {
|
||||
if method == Method::POST && !has_streamable_http_accept(&headers) {
|
||||
return (
|
||||
StatusCode::NOT_ACCEPTABLE,
|
||||
"missing MCP Accept header".to_string(),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
let rpc_method = body.get("method").and_then(Value::as_str).unwrap_or("");
|
||||
if method == Method::POST && rpc_method == "initialize" {
|
||||
state.init_count.fetch_add(1, AtomicOrdering::SeqCst);
|
||||
return (
|
||||
[(HEADER_SESSION_ID, "session-1")],
|
||||
Json(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"].clone(),
|
||||
"result": {
|
||||
"protocolVersion": LATEST_PROTOCOL_VERSION,
|
||||
"capabilities": { "tools": { "listChanged": true } },
|
||||
"serverInfo": { "name": "test-server", "version": "1.0.0" }
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
if headers.get(HEADER_SESSION_ID).and_then(|v| v.to_str().ok()) != Some("session-1") {
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
"missing or invalid session".to_string(),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
if headers
|
||||
.get(HEADER_PROTOCOL_VERSION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
!= Some(LATEST_PROTOCOL_VERSION)
|
||||
{
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
"missing protocol version".to_string(),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
|
||||
match rpc_method {
|
||||
"notifications/initialized" => StatusCode::NO_CONTENT.into_response(),
|
||||
"tools/list" => Json(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"].clone(),
|
||||
"result": {
|
||||
"tools": [{
|
||||
"name": "needs_header",
|
||||
"description": "needs x-mcp-header",
|
||||
"inputSchema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"tenant": {
|
||||
"type": "string",
|
||||
"x-mcp-header": "tenant"
|
||||
}
|
||||
}
|
||||
}
|
||||
}]
|
||||
}
|
||||
}))
|
||||
.into_response(),
|
||||
"tools/call" => {
|
||||
state.call_count.fetch_add(1, AtomicOrdering::SeqCst);
|
||||
if headers
|
||||
.get("Mcp-Param-tenant")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
!= Some("acme")
|
||||
{
|
||||
return (
|
||||
StatusCode::BAD_REQUEST,
|
||||
"missing mirrored tenant header".to_string(),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
Json(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"].clone(),
|
||||
"result": {
|
||||
"content": [{
|
||||
"type": "text",
|
||||
"text": "remote result"
|
||||
}]
|
||||
}
|
||||
}))
|
||||
.into_response()
|
||||
}
|
||||
_ => (
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!("unexpected method {rpc_method}"),
|
||||
)
|
||||
.into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn events_handler(headers: AxumHeaderMap) -> Response {
|
||||
if headers
|
||||
.get(ACCEPT)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.filter(|value| value.contains("text/event-stream"))
|
||||
.is_none()
|
||||
{
|
||||
return (
|
||||
StatusCode::NOT_ACCEPTABLE,
|
||||
"missing SSE Accept header".to_string(),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
if headers.get(HEADER_SESSION_ID).is_none() {
|
||||
return (StatusCode::BAD_REQUEST, "no session".to_string()).into_response();
|
||||
}
|
||||
(
|
||||
[(CONTENT_TYPE.as_str(), "text/event-stream")],
|
||||
"id: 1\nevent: message\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"notifications/message\",\"params\":{\"ok\":true}}\n\n",
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
async fn delete_handler() -> Response {
|
||||
StatusCode::NO_CONTENT.into_response()
|
||||
}
|
||||
|
||||
async fn bearer_required_handler(headers: AxumHeaderMap, Json(body): Json<Value>) -> Response {
|
||||
if headers.get(AUTHORIZATION).and_then(|v| v.to_str().ok()) != Some("Bearer secret-token") {
|
||||
return (StatusCode::UNAUTHORIZED, "missing bearer".to_string()).into_response();
|
||||
}
|
||||
Json(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"].clone(),
|
||||
"result": {
|
||||
"protocolVersion": LATEST_PROTOCOL_VERSION,
|
||||
"capabilities": {},
|
||||
"serverInfo": { "name": "bearer-server", "version": "1.0.0" }
|
||||
}
|
||||
}))
|
||||
.into_response()
|
||||
}
|
||||
|
||||
async fn retrying_mcp_handler(
|
||||
State(state): State<TestState>,
|
||||
headers: AxumHeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> Response {
|
||||
if !has_streamable_http_accept(&headers) {
|
||||
return (
|
||||
StatusCode::NOT_ACCEPTABLE,
|
||||
"missing MCP Accept header".to_string(),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
let rpc_method = body.get("method").and_then(Value::as_str).unwrap_or("");
|
||||
if rpc_method == "initialize" {
|
||||
state.init_count.fetch_add(1, AtomicOrdering::SeqCst);
|
||||
return (
|
||||
[(HEADER_SESSION_ID, "session-retry")],
|
||||
Json(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"].clone(),
|
||||
"result": {
|
||||
"protocolVersion": LATEST_PROTOCOL_VERSION,
|
||||
"capabilities": { "tools": {} },
|
||||
"serverInfo": { "name": "retry-server", "version": "1.0.0" }
|
||||
}
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
if rpc_method == "notifications/initialized" {
|
||||
return StatusCode::NO_CONTENT.into_response();
|
||||
}
|
||||
if rpc_method == "tools/list" {
|
||||
let call_number = state.call_count.fetch_add(1, AtomicOrdering::SeqCst);
|
||||
if call_number == 0
|
||||
&& headers.get(HEADER_SESSION_ID).and_then(|v| v.to_str().ok()) == Some("session-retry")
|
||||
{
|
||||
return (StatusCode::NOT_FOUND, "expired".to_string()).into_response();
|
||||
}
|
||||
return Json(json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": body["id"].clone(),
|
||||
"result": { "tools": [] }
|
||||
}))
|
||||
.into_response();
|
||||
}
|
||||
(StatusCode::BAD_REQUEST, "unexpected".to_string()).into_response()
|
||||
}
|
||||
|
||||
async fn spawn_test_server() -> (String, TestState) {
|
||||
let state = TestState {
|
||||
init_count: Arc::new(AtomicUsize::new(0)),
|
||||
call_count: Arc::new(AtomicUsize::new(0)),
|
||||
};
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/",
|
||||
post(mcp_handler).get(events_handler).delete(delete_handler),
|
||||
)
|
||||
.with_state(state.clone());
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
(format!("http://{addr}/"), state)
|
||||
}
|
||||
|
||||
async fn spawn_retry_server() -> (String, TestState) {
|
||||
let state = TestState {
|
||||
init_count: Arc::new(AtomicUsize::new(0)),
|
||||
call_count: Arc::new(AtomicUsize::new(0)),
|
||||
};
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let app = Router::new()
|
||||
.route("/", post(retrying_mcp_handler))
|
||||
.with_state(state.clone());
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
(format!("http://{addr}/"), state)
|
||||
}
|
||||
|
||||
async fn spawn_discovery_server() -> String {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let base = format!("http://{addr}");
|
||||
let auth_header = format!(
|
||||
"Bearer realm=\"mcp\", resource_metadata=\"{base}/.well-known/oauth-protected-resource\""
|
||||
);
|
||||
let prm_base = base.clone();
|
||||
let issuer_base = base.clone();
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/",
|
||||
post(move || {
|
||||
let auth_header = auth_header.clone();
|
||||
async move {
|
||||
(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
[("WWW-Authenticate", auth_header.as_str())],
|
||||
"",
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/.well-known/oauth-protected-resource",
|
||||
get(move || {
|
||||
let prm_base = prm_base.clone();
|
||||
async move {
|
||||
let resource = format!("{prm_base}/");
|
||||
Json(json!({
|
||||
"resource": resource,
|
||||
"authorization_servers": [prm_base],
|
||||
"scopes_supported": ["mcp:tools"]
|
||||
}))
|
||||
}
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/.well-known/openid-configuration",
|
||||
get(move || {
|
||||
let issuer_base = issuer_base.clone();
|
||||
async move {
|
||||
let authorization_endpoint = format!("{}/authorize", issuer_base);
|
||||
let token_endpoint = format!("{}/token", issuer_base);
|
||||
Json(json!({
|
||||
"issuer": issuer_base,
|
||||
"authorization_endpoint": authorization_endpoint,
|
||||
"token_endpoint": token_endpoint,
|
||||
"grant_types_supported": ["authorization_code"],
|
||||
"code_challenge_methods_supported": ["S256"]
|
||||
}))
|
||||
}
|
||||
}),
|
||||
);
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
format!("http://{addr}/")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn initialize_and_list_tools_negotiate_session() {
|
||||
let (endpoint, state) = spawn_test_server().await;
|
||||
let client = McpHttpClient::new(endpoint, 5);
|
||||
let tools = client.list_tools().await.expect("list_tools");
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(state.init_count.load(AtomicOrdering::SeqCst), 1);
|
||||
let snapshot = client.initialize_snapshot().expect("snapshot");
|
||||
assert_eq!(snapshot.protocol_version, LATEST_PROTOCOL_VERSION);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn call_tool_mirrors_x_mcp_header_parameters() {
|
||||
let (endpoint, state) = spawn_test_server().await;
|
||||
let client = McpHttpClient::new(endpoint, 5);
|
||||
let result = client
|
||||
.call_tool("needs_header", json!({"tenant": "acme"}))
|
||||
.await
|
||||
.expect("call_tool");
|
||||
assert_eq!(result.rendered.output(), "remote result");
|
||||
assert_eq!(state.call_count.load(AtomicOrdering::SeqCst), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn session_404_triggers_reinitialize_and_retry() {
|
||||
let (endpoint, state) = spawn_retry_server().await;
|
||||
let client = McpHttpClient::new(endpoint, 5);
|
||||
let tools = client.list_tools().await.expect("list_tools");
|
||||
assert!(tools.is_empty());
|
||||
assert_eq!(state.init_count.load(AtomicOrdering::SeqCst), 2);
|
||||
assert_eq!(state.call_count.load(AtomicOrdering::SeqCst), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn drain_events_parses_sse_stream() {
|
||||
let (endpoint, _) = spawn_test_server().await;
|
||||
let client = McpHttpClient::new(endpoint, 5);
|
||||
let events = client.drain_events(None).await.expect("drain events");
|
||||
assert_eq!(events.len(), 1);
|
||||
assert_eq!(events[0].id.as_deref(), Some("1"));
|
||||
assert_eq!(events[0].event.as_deref(), Some("message"));
|
||||
assert_eq!(events[0].data.as_ref().unwrap()["params"]["ok"], true);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn close_session_sends_delete() {
|
||||
let (endpoint, _) = spawn_test_server().await;
|
||||
let client = McpHttpClient::new(endpoint, 5);
|
||||
client.initialize().await.expect("initialize");
|
||||
client.close_session().await.expect("close_session");
|
||||
assert!(client.initialize_snapshot().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redact_endpoint_hides_paths_and_credentials() {
|
||||
assert_eq!(
|
||||
redact_endpoint("https://example.com/path?x=1"),
|
||||
"https://example.com"
|
||||
);
|
||||
assert_eq!(
|
||||
redact_endpoint("https://user:pw@example.com/a"),
|
||||
"<redacted>"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_sse_events_handles_multiple_frames() {
|
||||
let body = "id: 1\nevent: message\ndata: {\"a\":1}\n\ndata: {\"b\":2}\n\n";
|
||||
let events = parse_sse_events(body).expect("events");
|
||||
assert_eq!(events.len(), 2);
|
||||
assert_eq!(events[0].id.as_deref(), Some("1"));
|
||||
assert_eq!(events[1].data.as_ref().unwrap()["b"], 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_www_authenticate_extracts_resource_metadata() {
|
||||
let mut headers = HeaderMap::new();
|
||||
headers.insert(
|
||||
"WWW-Authenticate",
|
||||
HeaderValue::from_static(
|
||||
"Bearer realm=\"mcp\", resource_metadata=\"https://example.com/.well-known/oauth-protected-resource\"",
|
||||
),
|
||||
);
|
||||
let challenge = parse_www_authenticate_challenge(&headers).expect("challenge");
|
||||
assert_eq!(challenge.scheme, "Bearer");
|
||||
assert_eq!(challenge.realm.as_deref(), Some("mcp"));
|
||||
assert_eq!(
|
||||
challenge.resource_metadata.as_deref(),
|
||||
Some("https://example.com/.well-known/oauth-protected-resource")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn discover_authorization_returns_none_when_not_401() {
|
||||
let (endpoint, _) = spawn_test_server().await;
|
||||
let client = McpHttpClient::new(endpoint, 5);
|
||||
assert!(client.discover_authorization().await.unwrap().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn discover_authorization_fetches_metadata() {
|
||||
let endpoint = spawn_discovery_server().await;
|
||||
let client = McpHttpClient::new(endpoint, 2);
|
||||
let ctx = client
|
||||
.discover_authorization()
|
||||
.await
|
||||
.expect("discover")
|
||||
.expect("some");
|
||||
assert_eq!(ctx.challenge.scheme, "Bearer");
|
||||
assert_eq!(
|
||||
ctx.protected_resource_metadata
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.scopes_supported,
|
||||
vec!["mcp:tools"]
|
||||
);
|
||||
assert_eq!(ctx.authorization_server_metadata.len(), 1);
|
||||
let expected_authorization_endpoint = format!(
|
||||
"{}/authorize",
|
||||
ctx.protected_resource_metadata
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.authorization_servers[0]
|
||||
);
|
||||
assert_eq!(
|
||||
ctx.authorization_server_metadata[0]
|
||||
.authorization_endpoint
|
||||
.as_deref(),
|
||||
Some(expected_authorization_endpoint.as_str())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bearer_auth_is_attached_to_initialize() {
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let app = Router::new().route("/", post(bearer_required_handler));
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
let client = McpHttpClient::with_options(
|
||||
format!("http://{addr}/"),
|
||||
2,
|
||||
McpAuthConfig::BearerToken {
|
||||
token: "secret-token".into(),
|
||||
},
|
||||
McpClientIdentityConfig::default(),
|
||||
);
|
||||
let init = client.initialize().await.expect("initialize");
|
||||
assert_eq!(init.server_info["name"], "bearer-server");
|
||||
}
|
||||
@@ -941,153 +941,5 @@ fn type_name(value: &Value) -> &'static str {
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
// ── schemas() coverage ────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn schemas_registry_search_has_no_required_inputs() {
|
||||
let s = schemas("registry_search");
|
||||
assert_eq!(s.namespace, "mcp_clients");
|
||||
assert!(s.inputs.iter().all(|f| !f.required));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn schemas_registry_get_requires_qualified_name() {
|
||||
let s = schemas("registry_get");
|
||||
let qn = s
|
||||
.inputs
|
||||
.iter()
|
||||
.find(|f| f.name == "qualified_name")
|
||||
.unwrap();
|
||||
assert!(qn.required);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn schemas_install_requires_qualified_name_and_env() {
|
||||
let s = schemas("install");
|
||||
let names: Vec<_> = s
|
||||
.inputs
|
||||
.iter()
|
||||
.filter(|f| f.required)
|
||||
.map(|f| f.name)
|
||||
.collect();
|
||||
assert!(names.contains(&"qualified_name"));
|
||||
assert!(names.contains(&"env"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn schemas_connect_requires_server_id() {
|
||||
let s = schemas("connect");
|
||||
let si = s.inputs.iter().find(|f| f.name == "server_id").unwrap();
|
||||
assert!(si.required);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn schemas_tool_call_requires_three_fields() {
|
||||
let s = schemas("tool_call");
|
||||
let required: Vec<_> = s.inputs.iter().filter(|f| f.required).collect();
|
||||
assert_eq!(required.len(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn schemas_config_assist_history_is_optional() {
|
||||
let s = schemas("config_assist");
|
||||
let history = s.inputs.iter().find(|f| f.name == "history").unwrap();
|
||||
assert!(!history.required);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn schemas_unknown_function_returns_placeholder() {
|
||||
let s = schemas("not-a-real-function");
|
||||
assert_eq!(s.function, "unknown");
|
||||
assert_eq!(s.outputs[0].name, "error");
|
||||
}
|
||||
|
||||
// ── all_controller_schemas / all_registered_controllers ────────────────────
|
||||
|
||||
#[test]
|
||||
fn all_controller_schemas_covers_expected_methods() {
|
||||
let schemas = all_controller_schemas();
|
||||
// 10 mcp_clients + 6 mcp_setup
|
||||
assert_eq!(schemas.len(), 16);
|
||||
let mcp_clients_count = schemas
|
||||
.iter()
|
||||
.filter(|s| s.namespace == "mcp_clients")
|
||||
.count();
|
||||
let mcp_setup_count = schemas
|
||||
.iter()
|
||||
.filter(|s| s.namespace == "mcp_setup")
|
||||
.count();
|
||||
assert_eq!(mcp_clients_count, 10);
|
||||
assert_eq!(mcp_setup_count, 6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn all_registered_controllers_has_handler_per_schema() {
|
||||
let controllers = all_registered_controllers();
|
||||
assert_eq!(controllers.len(), 16);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn all_registered_controllers_use_expected_namespaces() {
|
||||
for c in all_registered_controllers() {
|
||||
assert!(
|
||||
matches!(c.schema.namespace, "mcp_clients" | "mcp_setup"),
|
||||
"unexpected namespace {}",
|
||||
c.schema.namespace
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ── read_required ─────────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn read_required_returns_value_for_present_key() {
|
||||
let mut params = Map::new();
|
||||
params.insert("server_id".into(), json!("srv-1"));
|
||||
let got: String = read_required(¶ms, "server_id").unwrap();
|
||||
assert_eq!(got, "srv-1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_required_errors_on_missing_key() {
|
||||
let err = read_required::<String>(&Map::new(), "server_id").unwrap_err();
|
||||
assert!(err.contains("missing required param 'server_id'"));
|
||||
}
|
||||
|
||||
// ── read_optional_u32 ─────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn read_optional_u32_absent_is_none() {
|
||||
assert_eq!(read_optional_u32(&Map::new(), "page").unwrap(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_optional_u32_valid_number() {
|
||||
let mut p = Map::new();
|
||||
p.insert("page".into(), json!(2));
|
||||
assert_eq!(read_optional_u32(&p, "page").unwrap(), Some(2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_optional_u32_rejects_negative() {
|
||||
let mut p = Map::new();
|
||||
p.insert("page".into(), json!(-1));
|
||||
assert!(read_optional_u32(&p, "page").is_err());
|
||||
}
|
||||
|
||||
// ── type_name ─────────────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn type_name_covers_all_variants() {
|
||||
assert_eq!(type_name(&Value::Null), "null");
|
||||
assert_eq!(type_name(&json!(true)), "bool");
|
||||
assert_eq!(type_name(&json!(1)), "number");
|
||||
assert_eq!(type_name(&json!("s")), "string");
|
||||
assert_eq!(type_name(&json!([])), "array");
|
||||
assert_eq!(type_name(&json!({})), "object");
|
||||
}
|
||||
}
|
||||
#[path = "schemas_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
// ── schemas() coverage ────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn schemas_registry_search_has_no_required_inputs() {
|
||||
let s = schemas("registry_search");
|
||||
assert_eq!(s.namespace, "mcp_clients");
|
||||
assert!(s.inputs.iter().all(|f| !f.required));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn schemas_registry_get_requires_qualified_name() {
|
||||
let s = schemas("registry_get");
|
||||
let qn = s
|
||||
.inputs
|
||||
.iter()
|
||||
.find(|f| f.name == "qualified_name")
|
||||
.unwrap();
|
||||
assert!(qn.required);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn schemas_install_requires_qualified_name_and_env() {
|
||||
let s = schemas("install");
|
||||
let names: Vec<_> = s
|
||||
.inputs
|
||||
.iter()
|
||||
.filter(|f| f.required)
|
||||
.map(|f| f.name)
|
||||
.collect();
|
||||
assert!(names.contains(&"qualified_name"));
|
||||
assert!(names.contains(&"env"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn schemas_connect_requires_server_id() {
|
||||
let s = schemas("connect");
|
||||
let si = s.inputs.iter().find(|f| f.name == "server_id").unwrap();
|
||||
assert!(si.required);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn schemas_tool_call_requires_three_fields() {
|
||||
let s = schemas("tool_call");
|
||||
let required: Vec<_> = s.inputs.iter().filter(|f| f.required).collect();
|
||||
assert_eq!(required.len(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn schemas_config_assist_history_is_optional() {
|
||||
let s = schemas("config_assist");
|
||||
let history = s.inputs.iter().find(|f| f.name == "history").unwrap();
|
||||
assert!(!history.required);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn schemas_unknown_function_returns_placeholder() {
|
||||
let s = schemas("not-a-real-function");
|
||||
assert_eq!(s.function, "unknown");
|
||||
assert_eq!(s.outputs[0].name, "error");
|
||||
}
|
||||
|
||||
// ── all_controller_schemas / all_registered_controllers ────────────────────
|
||||
|
||||
#[test]
|
||||
fn all_controller_schemas_covers_expected_methods() {
|
||||
let schemas = all_controller_schemas();
|
||||
// 10 mcp_clients + 6 mcp_setup
|
||||
assert_eq!(schemas.len(), 16);
|
||||
let mcp_clients_count = schemas
|
||||
.iter()
|
||||
.filter(|s| s.namespace == "mcp_clients")
|
||||
.count();
|
||||
let mcp_setup_count = schemas
|
||||
.iter()
|
||||
.filter(|s| s.namespace == "mcp_setup")
|
||||
.count();
|
||||
assert_eq!(mcp_clients_count, 10);
|
||||
assert_eq!(mcp_setup_count, 6);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn all_registered_controllers_has_handler_per_schema() {
|
||||
let controllers = all_registered_controllers();
|
||||
assert_eq!(controllers.len(), 16);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn all_registered_controllers_use_expected_namespaces() {
|
||||
for c in all_registered_controllers() {
|
||||
assert!(
|
||||
matches!(c.schema.namespace, "mcp_clients" | "mcp_setup"),
|
||||
"unexpected namespace {}",
|
||||
c.schema.namespace
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ── read_required ─────────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn read_required_returns_value_for_present_key() {
|
||||
let mut params = Map::new();
|
||||
params.insert("server_id".into(), json!("srv-1"));
|
||||
let got: String = read_required(¶ms, "server_id").unwrap();
|
||||
assert_eq!(got, "srv-1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_required_errors_on_missing_key() {
|
||||
let err = read_required::<String>(&Map::new(), "server_id").unwrap_err();
|
||||
assert!(err.contains("missing required param 'server_id'"));
|
||||
}
|
||||
|
||||
// ── read_optional_u32 ─────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn read_optional_u32_absent_is_none() {
|
||||
assert_eq!(read_optional_u32(&Map::new(), "page").unwrap(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_optional_u32_valid_number() {
|
||||
let mut p = Map::new();
|
||||
p.insert("page".into(), json!(2));
|
||||
assert_eq!(read_optional_u32(&p, "page").unwrap(), Some(2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_optional_u32_rejects_negative() {
|
||||
let mut p = Map::new();
|
||||
p.insert("page".into(), json!(-1));
|
||||
assert!(read_optional_u32(&p, "page").is_err());
|
||||
}
|
||||
|
||||
// ── type_name ─────────────────────────────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn type_name_covers_all_variants() {
|
||||
assert_eq!(type_name(&Value::Null), "null");
|
||||
assert_eq!(type_name(&json!(true)), "bool");
|
||||
assert_eq!(type_name(&json!(1)), "number");
|
||||
assert_eq!(type_name(&json!("s")), "string");
|
||||
assert_eq!(type_name(&json!([])), "array");
|
||||
assert_eq!(type_name(&json!({})), "object");
|
||||
}
|
||||
@@ -1320,249 +1320,5 @@ async fn stub_tts(text: &str) -> Vec<i16> {
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::meet_agent::session::registry;
|
||||
|
||||
#[tokio::test]
|
||||
async fn run_turn_skips_short_buffers() {
|
||||
registry().start("brain-skip", 16_000).unwrap();
|
||||
registry()
|
||||
.with_session("brain-skip", |s| {
|
||||
s.push_inbound_pcm(&vec![0; 800]); // 50ms — under floor
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(run_turn("brain-skip").await.unwrap(), false);
|
||||
let _ = registry().stop("brain-skip");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn run_turn_falls_back_to_stub_without_backend() {
|
||||
// No backend session in test env → STT/LLM/TTS all fail and
|
||||
// each stage falls back to its stub. The turn still produces
|
||||
// a Heard event, a Spoke event, and synthesized PCM, so the
|
||||
// smoke-test contract holds.
|
||||
registry().start("brain-fallback", 16_000).unwrap();
|
||||
registry()
|
||||
.with_session("brain-fallback", |s| {
|
||||
s.push_inbound_pcm(&vec![1000; 16_000]); // 1s
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(run_turn("brain-fallback").await.unwrap(), true);
|
||||
registry()
|
||||
.with_session("brain-fallback", |s| {
|
||||
let kinds: Vec<_> = s.events().iter().map(|e| format!("{:?}", e.kind)).collect();
|
||||
assert!(kinds.contains(&"Heard".to_string()));
|
||||
assert!(kinds.contains(&"Spoke".to_string()));
|
||||
assert_eq!(s.turn_count, 1);
|
||||
assert!(s.spoken_seconds() > 0.0);
|
||||
})
|
||||
.unwrap();
|
||||
let _ = registry().stop("brain-fallback");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_chat_completion_text_pulls_first_choice() {
|
||||
let raw = json!({
|
||||
"choices": [
|
||||
{ "message": { "content": " hello world " } }
|
||||
]
|
||||
});
|
||||
assert_eq!(
|
||||
extract_chat_completion_text(&raw),
|
||||
Some("hello world".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_chat_completion_text_returns_none_on_malformed() {
|
||||
assert_eq!(extract_chat_completion_text(&json!({})), None);
|
||||
assert_eq!(
|
||||
extract_chat_completion_text(&json!({ "choices": [] })),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_dialog_history_maps_event_kinds_to_chat_roles() {
|
||||
let now = 0;
|
||||
let events = vec![
|
||||
SessionEvent {
|
||||
kind: SessionEventKind::Heard,
|
||||
text: "Alice: how's the build going".into(),
|
||||
timestamp_ms: now,
|
||||
},
|
||||
SessionEvent {
|
||||
kind: SessionEventKind::Note,
|
||||
text: "wake word".into(),
|
||||
timestamp_ms: now,
|
||||
},
|
||||
SessionEvent {
|
||||
kind: SessionEventKind::Spoke,
|
||||
text: "Build is green.".into(),
|
||||
timestamp_ms: now,
|
||||
},
|
||||
SessionEvent {
|
||||
kind: SessionEventKind::Heard,
|
||||
text: "Bob: ship it".into(),
|
||||
timestamp_ms: now,
|
||||
},
|
||||
];
|
||||
let history = recent_dialog_history(&events, 10);
|
||||
assert_eq!(history.len(), 3, "Note events are dropped");
|
||||
assert_eq!(history[0].role, "user");
|
||||
assert_eq!(history[1].role, "assistant");
|
||||
assert_eq!(history[2].role, "user");
|
||||
assert_eq!(history[2].content, "Bob: ship it");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_dialog_history_caps_at_window_keeping_most_recent() {
|
||||
let events: Vec<SessionEvent> = (0..30)
|
||||
.map(|i| SessionEvent {
|
||||
kind: SessionEventKind::Heard,
|
||||
text: format!("line {i}"),
|
||||
timestamp_ms: 0,
|
||||
})
|
||||
.collect();
|
||||
let history = recent_dialog_history(&events, 5);
|
||||
assert_eq!(history.len(), 5);
|
||||
assert_eq!(history[0].content, "line 25");
|
||||
assert_eq!(history[4].content, "line 29");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_for_speech_removes_markdown_punctuation_and_fences() {
|
||||
let raw = "**Got it.** Adding `that` to your follow-ups.";
|
||||
assert_eq!(
|
||||
strip_for_speech(raw),
|
||||
"Got it. Adding that to your follow-ups."
|
||||
);
|
||||
let fenced = "Sure:\n```\ncode\n```\nDone.";
|
||||
assert_eq!(strip_for_speech(fenced), "Sure: Done.");
|
||||
let bullets = "- one\n- two";
|
||||
assert_eq!(strip_for_speech(bullets), "one two");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_for_speech_preserves_empty_when_input_empty() {
|
||||
assert_eq!(strip_for_speech(""), "");
|
||||
assert_eq!(strip_for_speech(" \n "), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn soft_deny_message_names_both_owner_and_asker() {
|
||||
let line = soft_deny_message("Bob", "Alice");
|
||||
assert!(line.contains("Bob"), "must address the asker: {line}");
|
||||
assert!(line.contains("Alice"), "must name the owner: {line}");
|
||||
assert!(
|
||||
line.to_lowercase().contains("allow"),
|
||||
"must hint the magic word: {line}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn soft_deny_message_handles_missing_names_gracefully() {
|
||||
// No asker, no owner — should still be a polite English sentence,
|
||||
// not a templated stub with empty placeholders.
|
||||
let line = soft_deny_message("", "");
|
||||
assert!(!line.is_empty());
|
||||
assert!(
|
||||
!line.contains("{"),
|
||||
"must not leak format placeholders: {line}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn looks_like_grant_intent_accepts_canonical_phrases() {
|
||||
// Whole-prompt approvals.
|
||||
for phrase in ["allow", "yes", "ok", "okay", "go ahead", "permit"] {
|
||||
assert!(
|
||||
looks_like_grant_intent(phrase),
|
||||
"must accept bare approval phrase: {phrase}"
|
||||
);
|
||||
}
|
||||
// Common longer forms.
|
||||
for phrase in [
|
||||
"allow them",
|
||||
"allow Bob to ask",
|
||||
"let them in",
|
||||
"let them ask",
|
||||
"let her ask",
|
||||
"go ahead and answer them",
|
||||
"yes go ahead",
|
||||
"permit Bob",
|
||||
"you can tell Bob",
|
||||
] {
|
||||
assert!(looks_like_grant_intent(phrase), "should accept: {phrase}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classify_unauthorized_intent_treats_bare_wake_as_greeting() {
|
||||
// Empty tail after the wake phrase — the non-owner just
|
||||
// said "hey openhuman" with nothing else. Friendly hi-back
|
||||
// is the right call, not a refusal.
|
||||
assert_eq!(
|
||||
classify_unauthorized_intent("hey openhuman"),
|
||||
UnauthorizedIntent::Greeting
|
||||
);
|
||||
assert_eq!(
|
||||
classify_unauthorized_intent("Hi openhuman."),
|
||||
UnauthorizedIntent::Greeting
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classify_unauthorized_intent_treats_filler_as_greeting() {
|
||||
// Common pleasantries that contain greeting words only.
|
||||
for text in [
|
||||
"hello openhuman there",
|
||||
"hi openhuman everyone",
|
||||
"hey openhuman hi",
|
||||
"hey openhuman good morning",
|
||||
] {
|
||||
assert_eq!(
|
||||
classify_unauthorized_intent(text),
|
||||
UnauthorizedIntent::Greeting,
|
||||
"should be greeting: {text}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classify_unauthorized_intent_flags_task_asks() {
|
||||
// Substantive task asks — refuse + tell owner how to grant.
|
||||
for text in [
|
||||
"hey openhuman read my slack",
|
||||
"hi openhuman what's on alice's calendar",
|
||||
"openhuman send the report",
|
||||
"hello openhuman remember the launch",
|
||||
] {
|
||||
assert_eq!(
|
||||
classify_unauthorized_intent(text),
|
||||
UnauthorizedIntent::TaskAsk,
|
||||
"should be task: {text}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn looks_like_grant_intent_rejects_unrelated_prompts() {
|
||||
// Words that happen to contain "allow" / "yes" mid-prompt
|
||||
// shouldn't hijack a normal question — the matcher only
|
||||
// honors prompts that BEGIN with a permit verb.
|
||||
for phrase in [
|
||||
"what's on my calendar today",
|
||||
"did i allow that meeting earlier",
|
||||
"yesterday's notes please",
|
||||
"remind me to ok the budget",
|
||||
"permittivity of free space",
|
||||
] {
|
||||
assert!(
|
||||
!looks_like_grant_intent(phrase),
|
||||
"must not match unrelated prompt: {phrase}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
#[path = "brain_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,244 @@
|
||||
use super::*;
|
||||
use crate::openhuman::meet_agent::session::registry;
|
||||
|
||||
#[tokio::test]
|
||||
async fn run_turn_skips_short_buffers() {
|
||||
registry().start("brain-skip", 16_000).unwrap();
|
||||
registry()
|
||||
.with_session("brain-skip", |s| {
|
||||
s.push_inbound_pcm(&vec![0; 800]); // 50ms — under floor
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(run_turn("brain-skip").await.unwrap(), false);
|
||||
let _ = registry().stop("brain-skip");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn run_turn_falls_back_to_stub_without_backend() {
|
||||
// No backend session in test env → STT/LLM/TTS all fail and
|
||||
// each stage falls back to its stub. The turn still produces
|
||||
// a Heard event, a Spoke event, and synthesized PCM, so the
|
||||
// smoke-test contract holds.
|
||||
registry().start("brain-fallback", 16_000).unwrap();
|
||||
registry()
|
||||
.with_session("brain-fallback", |s| {
|
||||
s.push_inbound_pcm(&vec![1000; 16_000]); // 1s
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(run_turn("brain-fallback").await.unwrap(), true);
|
||||
registry()
|
||||
.with_session("brain-fallback", |s| {
|
||||
let kinds: Vec<_> = s.events().iter().map(|e| format!("{:?}", e.kind)).collect();
|
||||
assert!(kinds.contains(&"Heard".to_string()));
|
||||
assert!(kinds.contains(&"Spoke".to_string()));
|
||||
assert_eq!(s.turn_count, 1);
|
||||
assert!(s.spoken_seconds() > 0.0);
|
||||
})
|
||||
.unwrap();
|
||||
let _ = registry().stop("brain-fallback");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_chat_completion_text_pulls_first_choice() {
|
||||
let raw = json!({
|
||||
"choices": [
|
||||
{ "message": { "content": " hello world " } }
|
||||
]
|
||||
});
|
||||
assert_eq!(
|
||||
extract_chat_completion_text(&raw),
|
||||
Some("hello world".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_chat_completion_text_returns_none_on_malformed() {
|
||||
assert_eq!(extract_chat_completion_text(&json!({})), None);
|
||||
assert_eq!(
|
||||
extract_chat_completion_text(&json!({ "choices": [] })),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_dialog_history_maps_event_kinds_to_chat_roles() {
|
||||
let now = 0;
|
||||
let events = vec![
|
||||
SessionEvent {
|
||||
kind: SessionEventKind::Heard,
|
||||
text: "Alice: how's the build going".into(),
|
||||
timestamp_ms: now,
|
||||
},
|
||||
SessionEvent {
|
||||
kind: SessionEventKind::Note,
|
||||
text: "wake word".into(),
|
||||
timestamp_ms: now,
|
||||
},
|
||||
SessionEvent {
|
||||
kind: SessionEventKind::Spoke,
|
||||
text: "Build is green.".into(),
|
||||
timestamp_ms: now,
|
||||
},
|
||||
SessionEvent {
|
||||
kind: SessionEventKind::Heard,
|
||||
text: "Bob: ship it".into(),
|
||||
timestamp_ms: now,
|
||||
},
|
||||
];
|
||||
let history = recent_dialog_history(&events, 10);
|
||||
assert_eq!(history.len(), 3, "Note events are dropped");
|
||||
assert_eq!(history[0].role, "user");
|
||||
assert_eq!(history[1].role, "assistant");
|
||||
assert_eq!(history[2].role, "user");
|
||||
assert_eq!(history[2].content, "Bob: ship it");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recent_dialog_history_caps_at_window_keeping_most_recent() {
|
||||
let events: Vec<SessionEvent> = (0..30)
|
||||
.map(|i| SessionEvent {
|
||||
kind: SessionEventKind::Heard,
|
||||
text: format!("line {i}"),
|
||||
timestamp_ms: 0,
|
||||
})
|
||||
.collect();
|
||||
let history = recent_dialog_history(&events, 5);
|
||||
assert_eq!(history.len(), 5);
|
||||
assert_eq!(history[0].content, "line 25");
|
||||
assert_eq!(history[4].content, "line 29");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_for_speech_removes_markdown_punctuation_and_fences() {
|
||||
let raw = "**Got it.** Adding `that` to your follow-ups.";
|
||||
assert_eq!(
|
||||
strip_for_speech(raw),
|
||||
"Got it. Adding that to your follow-ups."
|
||||
);
|
||||
let fenced = "Sure:\n```\ncode\n```\nDone.";
|
||||
assert_eq!(strip_for_speech(fenced), "Sure: Done.");
|
||||
let bullets = "- one\n- two";
|
||||
assert_eq!(strip_for_speech(bullets), "one two");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn strip_for_speech_preserves_empty_when_input_empty() {
|
||||
assert_eq!(strip_for_speech(""), "");
|
||||
assert_eq!(strip_for_speech(" \n "), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn soft_deny_message_names_both_owner_and_asker() {
|
||||
let line = soft_deny_message("Bob", "Alice");
|
||||
assert!(line.contains("Bob"), "must address the asker: {line}");
|
||||
assert!(line.contains("Alice"), "must name the owner: {line}");
|
||||
assert!(
|
||||
line.to_lowercase().contains("allow"),
|
||||
"must hint the magic word: {line}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn soft_deny_message_handles_missing_names_gracefully() {
|
||||
// No asker, no owner — should still be a polite English sentence,
|
||||
// not a templated stub with empty placeholders.
|
||||
let line = soft_deny_message("", "");
|
||||
assert!(!line.is_empty());
|
||||
assert!(
|
||||
!line.contains("{"),
|
||||
"must not leak format placeholders: {line}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn looks_like_grant_intent_accepts_canonical_phrases() {
|
||||
// Whole-prompt approvals.
|
||||
for phrase in ["allow", "yes", "ok", "okay", "go ahead", "permit"] {
|
||||
assert!(
|
||||
looks_like_grant_intent(phrase),
|
||||
"must accept bare approval phrase: {phrase}"
|
||||
);
|
||||
}
|
||||
// Common longer forms.
|
||||
for phrase in [
|
||||
"allow them",
|
||||
"allow Bob to ask",
|
||||
"let them in",
|
||||
"let them ask",
|
||||
"let her ask",
|
||||
"go ahead and answer them",
|
||||
"yes go ahead",
|
||||
"permit Bob",
|
||||
"you can tell Bob",
|
||||
] {
|
||||
assert!(looks_like_grant_intent(phrase), "should accept: {phrase}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classify_unauthorized_intent_treats_bare_wake_as_greeting() {
|
||||
// Empty tail after the wake phrase — the non-owner just
|
||||
// said "hey openhuman" with nothing else. Friendly hi-back
|
||||
// is the right call, not a refusal.
|
||||
assert_eq!(
|
||||
classify_unauthorized_intent("hey openhuman"),
|
||||
UnauthorizedIntent::Greeting
|
||||
);
|
||||
assert_eq!(
|
||||
classify_unauthorized_intent("Hi openhuman."),
|
||||
UnauthorizedIntent::Greeting
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classify_unauthorized_intent_treats_filler_as_greeting() {
|
||||
// Common pleasantries that contain greeting words only.
|
||||
for text in [
|
||||
"hello openhuman there",
|
||||
"hi openhuman everyone",
|
||||
"hey openhuman hi",
|
||||
"hey openhuman good morning",
|
||||
] {
|
||||
assert_eq!(
|
||||
classify_unauthorized_intent(text),
|
||||
UnauthorizedIntent::Greeting,
|
||||
"should be greeting: {text}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classify_unauthorized_intent_flags_task_asks() {
|
||||
// Substantive task asks — refuse + tell owner how to grant.
|
||||
for text in [
|
||||
"hey openhuman read my slack",
|
||||
"hi openhuman what's on alice's calendar",
|
||||
"openhuman send the report",
|
||||
"hello openhuman remember the launch",
|
||||
] {
|
||||
assert_eq!(
|
||||
classify_unauthorized_intent(text),
|
||||
UnauthorizedIntent::TaskAsk,
|
||||
"should be task: {text}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn looks_like_grant_intent_rejects_unrelated_prompts() {
|
||||
// Words that happen to contain "allow" / "yes" mid-prompt
|
||||
// shouldn't hijack a normal question — the matcher only
|
||||
// honors prompts that BEGIN with a permit verb.
|
||||
for phrase in [
|
||||
"what's on my calendar today",
|
||||
"did i allow that meeting earlier",
|
||||
"yesterday's notes please",
|
||||
"remind me to ok the budget",
|
||||
"permittivity of free space",
|
||||
] {
|
||||
assert!(
|
||||
!looks_like_grant_intent(phrase),
|
||||
"must not match unrelated prompt: {phrase}"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -935,141 +935,5 @@ pub(super) async fn parse_document(
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory_store::types::NamespaceDocumentInput;
|
||||
use serde_json::json;
|
||||
|
||||
fn sample_input() -> NamespaceDocumentInput {
|
||||
NamespaceDocumentInput {
|
||||
namespace: "global".into(),
|
||||
key: "doc-1".into(),
|
||||
title: "OpenHuman roadmap".into(),
|
||||
content: "Alice owns roadmap".into(),
|
||||
source_type: "manual".into(),
|
||||
priority: "normal".into(),
|
||||
tags: vec!["existing".into()],
|
||||
metadata: json!({"seed": true}),
|
||||
category: "core".into(),
|
||||
session_id: Some("session-1".into()),
|
||||
document_id: Some("doc-id-1".into()),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_sentences_breaks_on_punctuation_and_merges_tiny_fragments() {
|
||||
let parts = split_sentences("Hello world. Ok.\nNext line?");
|
||||
assert_eq!(parts.len(), 2);
|
||||
assert_eq!(parts[0], "Hello world Ok");
|
||||
assert_eq!(parts[1], "Next line?");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_units_respects_extraction_mode() {
|
||||
let chunks = vec!["One. Two.".to_string(), "Three".to_string()];
|
||||
let sentence_units = build_units(&chunks, ExtractionMode::Sentence);
|
||||
let chunk_units = build_units(&chunks, ExtractionMode::Chunk);
|
||||
|
||||
assert_eq!(sentence_units.len(), 2);
|
||||
assert_eq!(sentence_units[0].chunk_index, 0);
|
||||
assert_eq!(sentence_units[0].text, "One Two");
|
||||
assert_eq!(sentence_units[1].chunk_index, 1);
|
||||
assert_eq!(sentence_units[1].text, "Three");
|
||||
|
||||
assert_eq!(chunk_units.len(), 2);
|
||||
assert_eq!(chunk_units[0].text, "One. Two");
|
||||
assert_eq!(chunk_units[1].text, "Three");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn find_chunk_index_prefers_hint_then_wraps() {
|
||||
let chunks = vec![
|
||||
"alpha content".to_string(),
|
||||
"beta needle".to_string(),
|
||||
"gamma trailing".to_string(),
|
||||
];
|
||||
assert_eq!(find_chunk_index(&chunks, "needle", 1), 1);
|
||||
assert_eq!(find_chunk_index(&chunks, "alpha", 2), 0);
|
||||
assert_eq!(find_chunk_index(&chunks, "missing", 2), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn alias_map_builds_reverse_lookup() {
|
||||
let mut entities = HashMap::new();
|
||||
entities.insert(
|
||||
"ALICE".into(),
|
||||
RawEntity {
|
||||
name: "ALICE".into(),
|
||||
entity_type: "PERSON".into(),
|
||||
confidence: 0.8,
|
||||
},
|
||||
);
|
||||
entities.insert(
|
||||
"ALICE SMITH".into(),
|
||||
RawEntity {
|
||||
name: "ALICE SMITH".into(),
|
||||
entity_type: "PERSON".into(),
|
||||
confidence: 0.9,
|
||||
},
|
||||
);
|
||||
let aliases = build_alias_map(&entities);
|
||||
assert_eq!(
|
||||
aliases.get("ALICE").map(String::as_str),
|
||||
Some("ALICE SMITH")
|
||||
);
|
||||
assert_eq!(resolve_alias("ALICE", &aliases), "ALICE SMITH");
|
||||
|
||||
let reverse = reverse_aliases(&aliases);
|
||||
assert_eq!(reverse.get("ALICE SMITH"), Some(&vec!["ALICE".to_string()]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn enrich_document_metadata_merges_tags_and_ingestion_details() {
|
||||
let input = sample_input();
|
||||
let parsed = ParsedIngestion {
|
||||
tags: vec!["decision".into(), "existing".into()],
|
||||
metadata: json!({"kind": "profile", "extra": 1}),
|
||||
entities: vec![],
|
||||
relations: vec![],
|
||||
chunk_count: 3,
|
||||
preference_count: 1,
|
||||
decision_count: 2,
|
||||
};
|
||||
let config = MemoryIngestionConfig::default();
|
||||
let (enriched, tags) = enrich_document_metadata(&input, &parsed, &config);
|
||||
|
||||
assert_eq!(tags, vec!["decision".to_string(), "existing".to_string()]);
|
||||
assert_eq!(enriched.tags, tags);
|
||||
assert_eq!(enriched.metadata["seed"], json!(true));
|
||||
assert_eq!(enriched.metadata["extra"], json!(1));
|
||||
assert_eq!(
|
||||
enriched.metadata["ingestion"]["model_name"],
|
||||
config.model_name
|
||||
);
|
||||
assert_eq!(enriched.metadata["ingestion"]["chunk_count"], json!(3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_people_from_header_collects_named_people() {
|
||||
let mut acc = ExtractionAccumulator::default();
|
||||
let people = extract_people_from_header(
|
||||
"Alice Smith <alice@example.com>, Bob Jones <bob@example.com>",
|
||||
&mut acc,
|
||||
);
|
||||
assert_eq!(
|
||||
people,
|
||||
vec!["ALICE SMITH".to_string(), "BOB JONES".to_string()]
|
||||
);
|
||||
assert!(acc.entities.contains_key("ALICE SMITH"));
|
||||
assert!(acc.entities.contains_key("BOB JONES"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detect_primary_subject_only_matches_openhuman() {
|
||||
assert_eq!(
|
||||
detect_primary_subject("OpenHuman desktop roadmap"),
|
||||
Some("OPENHUMAN".to_string())
|
||||
);
|
||||
assert_eq!(detect_primary_subject("General roadmap"), None);
|
||||
}
|
||||
}
|
||||
#[path = "parse_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,136 @@
|
||||
use super::*;
|
||||
use crate::openhuman::memory_store::types::NamespaceDocumentInput;
|
||||
use serde_json::json;
|
||||
|
||||
fn sample_input() -> NamespaceDocumentInput {
|
||||
NamespaceDocumentInput {
|
||||
namespace: "global".into(),
|
||||
key: "doc-1".into(),
|
||||
title: "OpenHuman roadmap".into(),
|
||||
content: "Alice owns roadmap".into(),
|
||||
source_type: "manual".into(),
|
||||
priority: "normal".into(),
|
||||
tags: vec!["existing".into()],
|
||||
metadata: json!({"seed": true}),
|
||||
category: "core".into(),
|
||||
session_id: Some("session-1".into()),
|
||||
document_id: Some("doc-id-1".into()),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_sentences_breaks_on_punctuation_and_merges_tiny_fragments() {
|
||||
let parts = split_sentences("Hello world. Ok.\nNext line?");
|
||||
assert_eq!(parts.len(), 2);
|
||||
assert_eq!(parts[0], "Hello world Ok");
|
||||
assert_eq!(parts[1], "Next line?");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_units_respects_extraction_mode() {
|
||||
let chunks = vec!["One. Two.".to_string(), "Three".to_string()];
|
||||
let sentence_units = build_units(&chunks, ExtractionMode::Sentence);
|
||||
let chunk_units = build_units(&chunks, ExtractionMode::Chunk);
|
||||
|
||||
assert_eq!(sentence_units.len(), 2);
|
||||
assert_eq!(sentence_units[0].chunk_index, 0);
|
||||
assert_eq!(sentence_units[0].text, "One Two");
|
||||
assert_eq!(sentence_units[1].chunk_index, 1);
|
||||
assert_eq!(sentence_units[1].text, "Three");
|
||||
|
||||
assert_eq!(chunk_units.len(), 2);
|
||||
assert_eq!(chunk_units[0].text, "One. Two");
|
||||
assert_eq!(chunk_units[1].text, "Three");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn find_chunk_index_prefers_hint_then_wraps() {
|
||||
let chunks = vec![
|
||||
"alpha content".to_string(),
|
||||
"beta needle".to_string(),
|
||||
"gamma trailing".to_string(),
|
||||
];
|
||||
assert_eq!(find_chunk_index(&chunks, "needle", 1), 1);
|
||||
assert_eq!(find_chunk_index(&chunks, "alpha", 2), 0);
|
||||
assert_eq!(find_chunk_index(&chunks, "missing", 2), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn alias_map_builds_reverse_lookup() {
|
||||
let mut entities = HashMap::new();
|
||||
entities.insert(
|
||||
"ALICE".into(),
|
||||
RawEntity {
|
||||
name: "ALICE".into(),
|
||||
entity_type: "PERSON".into(),
|
||||
confidence: 0.8,
|
||||
},
|
||||
);
|
||||
entities.insert(
|
||||
"ALICE SMITH".into(),
|
||||
RawEntity {
|
||||
name: "ALICE SMITH".into(),
|
||||
entity_type: "PERSON".into(),
|
||||
confidence: 0.9,
|
||||
},
|
||||
);
|
||||
let aliases = build_alias_map(&entities);
|
||||
assert_eq!(
|
||||
aliases.get("ALICE").map(String::as_str),
|
||||
Some("ALICE SMITH")
|
||||
);
|
||||
assert_eq!(resolve_alias("ALICE", &aliases), "ALICE SMITH");
|
||||
|
||||
let reverse = reverse_aliases(&aliases);
|
||||
assert_eq!(reverse.get("ALICE SMITH"), Some(&vec!["ALICE".to_string()]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn enrich_document_metadata_merges_tags_and_ingestion_details() {
|
||||
let input = sample_input();
|
||||
let parsed = ParsedIngestion {
|
||||
tags: vec!["decision".into(), "existing".into()],
|
||||
metadata: json!({"kind": "profile", "extra": 1}),
|
||||
entities: vec![],
|
||||
relations: vec![],
|
||||
chunk_count: 3,
|
||||
preference_count: 1,
|
||||
decision_count: 2,
|
||||
};
|
||||
let config = MemoryIngestionConfig::default();
|
||||
let (enriched, tags) = enrich_document_metadata(&input, &parsed, &config);
|
||||
|
||||
assert_eq!(tags, vec!["decision".to_string(), "existing".to_string()]);
|
||||
assert_eq!(enriched.tags, tags);
|
||||
assert_eq!(enriched.metadata["seed"], json!(true));
|
||||
assert_eq!(enriched.metadata["extra"], json!(1));
|
||||
assert_eq!(
|
||||
enriched.metadata["ingestion"]["model_name"],
|
||||
config.model_name
|
||||
);
|
||||
assert_eq!(enriched.metadata["ingestion"]["chunk_count"], json!(3));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_people_from_header_collects_named_people() {
|
||||
let mut acc = ExtractionAccumulator::default();
|
||||
let people = extract_people_from_header(
|
||||
"Alice Smith <alice@example.com>, Bob Jones <bob@example.com>",
|
||||
&mut acc,
|
||||
);
|
||||
assert_eq!(
|
||||
people,
|
||||
vec!["ALICE SMITH".to_string(), "BOB JONES".to_string()]
|
||||
);
|
||||
assert!(acc.entities.contains_key("ALICE SMITH"));
|
||||
assert!(acc.entities.contains_key("BOB JONES"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detect_primary_subject_only_matches_openhuman() {
|
||||
assert_eq!(
|
||||
detect_primary_subject("OpenHuman desktop roadmap"),
|
||||
Some("OPENHUMAN".to_string())
|
||||
);
|
||||
assert_eq!(detect_primary_subject("General roadmap"), None);
|
||||
}
|
||||
@@ -1064,39 +1064,5 @@ fn to_json<T: serde::Serialize>(outcome: RpcOutcome<T>) -> Result<Value, String>
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn all_controller_schemas_and_registered_controllers_stay_in_sync() {
|
||||
let schemas = all_controller_schemas();
|
||||
let controllers = all_registered_controllers();
|
||||
assert_eq!(schemas.len(), controllers.len());
|
||||
assert!(schemas.iter().all(|s| s.namespace == NAMESPACE));
|
||||
assert!(controllers.iter().all(|c| c.schema.namespace == NAMESPACE));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_function_schema_returns_error_output() {
|
||||
let schema = schemas("not_real");
|
||||
assert_eq!(schema.namespace, NAMESPACE);
|
||||
assert_eq!(schema.function, "unknown");
|
||||
assert_eq!(schema.outputs.len(), 1);
|
||||
assert_eq!(schema.outputs[0].name, "error");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ingest_schema_requires_source_kind_source_id_and_payload() {
|
||||
let schema = schemas("ingest");
|
||||
assert_eq!(schema.function, "ingest");
|
||||
let required: Vec<&str> = schema
|
||||
.inputs
|
||||
.iter()
|
||||
.filter(|f| f.required)
|
||||
.map(|f| f.name)
|
||||
.collect();
|
||||
assert!(required.contains(&"source_kind"));
|
||||
assert!(required.contains(&"source_id"));
|
||||
assert!(required.contains(&"payload"));
|
||||
}
|
||||
}
|
||||
#[path = "schema_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn all_controller_schemas_and_registered_controllers_stay_in_sync() {
|
||||
let schemas = all_controller_schemas();
|
||||
let controllers = all_registered_controllers();
|
||||
assert_eq!(schemas.len(), controllers.len());
|
||||
assert!(schemas.iter().all(|s| s.namespace == NAMESPACE));
|
||||
assert!(controllers.iter().all(|c| c.schema.namespace == NAMESPACE));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_function_schema_returns_error_output() {
|
||||
let schema = schemas("not_real");
|
||||
assert_eq!(schema.namespace, NAMESPACE);
|
||||
assert_eq!(schema.function, "unknown");
|
||||
assert_eq!(schema.outputs.len(), 1);
|
||||
assert_eq!(schema.outputs[0].name, "error");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ingest_schema_requires_source_kind_source_id_and_payload() {
|
||||
let schema = schemas("ingest");
|
||||
assert_eq!(schema.function, "ingest");
|
||||
let required: Vec<&str> = schema
|
||||
.inputs
|
||||
.iter()
|
||||
.filter(|f| f.required)
|
||||
.map(|f| f.name)
|
||||
.collect();
|
||||
assert!(required.contains(&"source_kind"));
|
||||
assert!(required.contains(&"source_id"));
|
||||
assert!(required.contains(&"payload"));
|
||||
}
|
||||
@@ -782,730 +782,5 @@ async fn handle_reembed_backfill(config: &Config, job: &Job) -> Result<JobOutcom
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree_source::registry::get_or_create_source_tree;
|
||||
use crate::openhuman::memory_queue::store::{count_by_status, count_total};
|
||||
use crate::openhuman::memory_queue::types::JobStatus;
|
||||
use crate::openhuman::memory_store::chunks::store::with_connection;
|
||||
use crate::openhuman::memory_store::content as content_store;
|
||||
use crate::openhuman::memory_tree::tree::bucket_seal::{append_leaf_deferred, LeafRef};
|
||||
use crate::openhuman::memory_tree::tree::store as src_store;
|
||||
use chrono::TimeZone;
|
||||
use rusqlite::params;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn test_config() -> (TempDir, Config) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let mut cfg = Config::default();
|
||||
cfg.workspace_dir = tmp.path().to_path_buf();
|
||||
cfg.memory_tree.embedding_endpoint = None;
|
||||
cfg.memory_tree.embedding_model = None;
|
||||
cfg.memory_tree.embedding_strict = false;
|
||||
(tmp, cfg)
|
||||
}
|
||||
|
||||
/// Build a minimal `Job` row for direct handler invocation. Mirrors
|
||||
/// what `claim_next` would produce for a freshly-claimed row.
|
||||
fn mk_running_job(kind: JobKind, payload_json: String) -> Job {
|
||||
let now_ms = chrono::Utc::now().timestamp_millis();
|
||||
Job {
|
||||
id: "test-job-id".into(),
|
||||
kind,
|
||||
payload_json,
|
||||
dedupe_key: None,
|
||||
status: JobStatus::Running,
|
||||
attempts: 1,
|
||||
max_attempts: 5,
|
||||
available_at_ms: now_ms,
|
||||
locked_until_ms: Some(now_ms + 60_000),
|
||||
last_error: None,
|
||||
created_at_ms: now_ms,
|
||||
started_at_ms: Some(now_ms),
|
||||
completed_at_ms: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Count rows in `mem_tree_jobs` matching a specific kind.
|
||||
fn count_jobs_of_kind(cfg: &Config, kind: &str) -> u64 {
|
||||
with_connection(cfg, |conn| {
|
||||
let n: i64 = conn.query_row(
|
||||
"SELECT COUNT(*) FROM mem_tree_jobs WHERE kind = ?1",
|
||||
params![kind],
|
||||
|r| r.get(0),
|
||||
)?;
|
||||
Ok(n.max(0) as u64)
|
||||
})
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
/// Seed a source tree and push enough labeled leaves into its L0 buffer
|
||||
/// to cross `INPUT_TOKEN_BUDGET`, returning the tree. The caller can then
|
||||
/// fire `handle_seal` and inspect the result.
|
||||
async fn seed_source_tree_ready_to_seal(
|
||||
cfg: &Config,
|
||||
) -> crate::openhuman::memory_store::trees::types::Tree {
|
||||
use crate::openhuman::memory_store::chunks::store::upsert_chunks;
|
||||
use crate::openhuman::memory_store::chunks::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
let tree = get_or_create_source_tree(cfg, "slack:#eng").unwrap();
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "handler-seed"),
|
||||
content: "alice@example.com leading the rollout".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
// Bust budget so the L0 buffer is "ready" for seal.
|
||||
token_count: 60_000,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(cfg, &[chunk.clone()]).unwrap();
|
||||
// Stage to disk so `hydrate_leaf_inputs` can read the full body via
|
||||
// `read_chunk_body` when `handle_seal` fires and calls `seal_one_level`.
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory_store::chunks::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
let leaf = LeafRef {
|
||||
chunk_id: chunk.id,
|
||||
token_count: 60_000,
|
||||
timestamp: ts,
|
||||
content: chunk.content,
|
||||
entities: vec![],
|
||||
topics: vec![],
|
||||
score: 0.5,
|
||||
};
|
||||
// append_leaf_deferred only buffers; doesn't seal. handle_seal will.
|
||||
let _ = append_leaf_deferred(cfg, &tree, &leaf).unwrap();
|
||||
tree
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn source_tree_seal_handler_enqueues_summary_topic_route() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let tree = seed_source_tree_ready_to_seal(&cfg).await;
|
||||
|
||||
let payload = SealPayload {
|
||||
tree_id: tree.id.clone(),
|
||||
level: 0,
|
||||
force_now_ms: None,
|
||||
};
|
||||
let job = mk_running_job(JobKind::Seal, serde_json::to_string(&payload).unwrap());
|
||||
|
||||
// Pre-condition: queue has no topic_route jobs.
|
||||
assert_eq!(count_jobs_of_kind(&cfg, "topic_route"), 0);
|
||||
|
||||
super::handle_seal(&cfg, &job).await.unwrap();
|
||||
|
||||
// Post-condition: source-tree seal must enqueue exactly one
|
||||
// topic_route job carrying NodeRef::Summary { summary_id: <new> }.
|
||||
assert_eq!(
|
||||
count_jobs_of_kind(&cfg, "topic_route"),
|
||||
1,
|
||||
"source-tree seal must enqueue summary-side topic_route"
|
||||
);
|
||||
assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 1);
|
||||
|
||||
// Inspect the enqueued payload to confirm it's a Summary variant.
|
||||
let payload_json: String = with_connection(&cfg, |conn| {
|
||||
let s: String = conn
|
||||
.query_row(
|
||||
"SELECT payload_json FROM mem_tree_jobs WHERE kind = 'topic_route'",
|
||||
[],
|
||||
|r| r.get(0),
|
||||
)
|
||||
.unwrap();
|
||||
Ok(s)
|
||||
})
|
||||
.unwrap();
|
||||
let p: TopicRoutePayload = serde_json::from_str(&payload_json).unwrap();
|
||||
match p.node {
|
||||
NodeRef::Summary { summary_id } => {
|
||||
// Format: `summary:<13-digit-ms>:L<level>-<8hex>` —
|
||||
// see `tree::registry::new_summary_id`.
|
||||
assert!(
|
||||
summary_id.starts_with("summary:") && summary_id.contains(":L1-"),
|
||||
"expected summary id with L1 segment, got {summary_id}"
|
||||
);
|
||||
}
|
||||
other => panic!("expected NodeRef::Summary, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn topic_tree_seal_handler_does_not_enqueue_topic_route() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
// Spawn a topic tree directly via the registry (skipping curator's
|
||||
// hotness gate — we just need a TreeKind::Topic with leaves).
|
||||
let topic_tree = crate::openhuman::memory_store::trees::registry::get_or_create_topic_tree(
|
||||
&cfg,
|
||||
"topic:phoenix-migration",
|
||||
)
|
||||
.unwrap();
|
||||
// Push a single 10k-token leaf so L0 is gate-ready.
|
||||
use crate::openhuman::memory_store::chunks::store::upsert_chunks;
|
||||
use crate::openhuman::memory_store::chunks::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "topic-seed"),
|
||||
content: "topic content".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
token_count: 60_000,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
|
||||
// Stage to disk so `hydrate_leaf_inputs` can read the full body
|
||||
// when `handle_seal` fires.
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory_store::chunks::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
let leaf = LeafRef {
|
||||
chunk_id: chunk.id,
|
||||
token_count: 60_000,
|
||||
timestamp: ts,
|
||||
content: chunk.content,
|
||||
entities: vec![],
|
||||
topics: vec![],
|
||||
score: 0.5,
|
||||
};
|
||||
append_leaf_deferred(&cfg, &topic_tree, &leaf).unwrap();
|
||||
|
||||
let payload = SealPayload {
|
||||
tree_id: topic_tree.id.clone(),
|
||||
level: 0,
|
||||
force_now_ms: None,
|
||||
};
|
||||
let job = mk_running_job(JobKind::Seal, serde_json::to_string(&payload).unwrap());
|
||||
|
||||
super::handle_seal(&cfg, &job).await.unwrap();
|
||||
|
||||
// Topic-tree seals are sinks: must not enqueue any topic_route.
|
||||
assert_eq!(
|
||||
count_jobs_of_kind(&cfg, "topic_route"),
|
||||
0,
|
||||
"topic-tree seal must NOT enqueue topic_route (trees are sinks)"
|
||||
);
|
||||
// The seal itself should still have produced a summary node.
|
||||
assert_eq!(src_store::count_summaries(&cfg, &topic_tree.id).unwrap(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn handle_append_buffer_with_summary_payload_pushes_into_topic_tree() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
|
||||
// 1. Create a target topic tree with a clean L0 buffer.
|
||||
let topic_tree = crate::openhuman::memory_store::trees::registry::get_or_create_topic_tree(
|
||||
&cfg,
|
||||
"email:alice@example.com",
|
||||
)
|
||||
.unwrap();
|
||||
let l0_before = src_store::get_buffer(&cfg, &topic_tree.id, 0).unwrap();
|
||||
assert!(l0_before.is_empty());
|
||||
|
||||
// 2. Manually insert a summary node we can route. The simplest way
|
||||
// is to create a separate source tree, push two 6k leaves into
|
||||
// it, and let the seal produce a summary we can address.
|
||||
let source_tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap();
|
||||
use crate::openhuman::memory_store::chunks::store::upsert_chunks;
|
||||
use crate::openhuman::memory_store::chunks::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
use crate::openhuman::memory_tree::tree::bucket_seal::seal_one_level;
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
for seq in 0..2 {
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", seq, "summary-seed"),
|
||||
content: format!("source content {seq}"),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
token_count: 30_000,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
|
||||
// Stage to disk so `hydrate_leaf_inputs` can read the full body
|
||||
// during `seal_one_level`.
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory_store::chunks::store::upsert_staged_chunks_tx(
|
||||
&tx, &staged,
|
||||
)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
let leaf = LeafRef {
|
||||
chunk_id: chunk.id,
|
||||
token_count: 30_000,
|
||||
timestamp: ts,
|
||||
content: chunk.content,
|
||||
entities: vec![],
|
||||
topics: vec![],
|
||||
score: 0.5,
|
||||
};
|
||||
let _ = append_leaf_deferred(&cfg, &source_tree, &leaf).unwrap();
|
||||
}
|
||||
// Force-seal the source tree's L0 to mint the summary.
|
||||
use crate::openhuman::memory::chat::{test_override, ChatProvider, StaticChatProvider};
|
||||
let buf = src_store::get_buffer(&cfg, &source_tree.id, 0).unwrap();
|
||||
let provider: std::sync::Arc<dyn ChatProvider> =
|
||||
std::sync::Arc::new(StaticChatProvider::new("test summary content"));
|
||||
let summary_id = test_override::with_provider(provider, async {
|
||||
seal_one_level(
|
||||
&cfg,
|
||||
&source_tree,
|
||||
&buf,
|
||||
&crate::openhuman::memory_tree::tree::bucket_seal::LabelStrategy::Empty,
|
||||
// No follow-up enqueues — the test scopes assertions to the
|
||||
// append_buffer handler, not seal-side fan-out.
|
||||
false,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
})
|
||||
.await;
|
||||
|
||||
// 3. Build an append_buffer payload routing the summary into the
|
||||
// topic tree.
|
||||
let payload = AppendBufferPayload {
|
||||
node: NodeRef::Summary {
|
||||
summary_id: summary_id.clone(),
|
||||
},
|
||||
target: AppendTarget::Topic {
|
||||
tree_id: topic_tree.id.clone(),
|
||||
},
|
||||
};
|
||||
let job = mk_running_job(
|
||||
JobKind::AppendBuffer,
|
||||
serde_json::to_string(&payload).unwrap(),
|
||||
);
|
||||
|
||||
// Clear out any pending append_buffer jobs minted upstream so the
|
||||
// post-condition assertion below is unambiguous.
|
||||
let pre = count_total(&cfg).unwrap();
|
||||
|
||||
super::handle_append_buffer(&cfg, &job).await.unwrap();
|
||||
|
||||
// 4. Topic tree's L0 buffer should now hold the summary id.
|
||||
let l0_after = src_store::get_buffer(&cfg, &topic_tree.id, 0).unwrap();
|
||||
assert_eq!(l0_after.item_ids, vec![summary_id]);
|
||||
assert!(l0_after.token_sum > 0);
|
||||
|
||||
// No new jobs should have been enqueued (buffer didn't cross gate).
|
||||
assert_eq!(count_total(&cfg).unwrap(), pre);
|
||||
}
|
||||
|
||||
/// #1574 §6: a chunk with content but no sidecar vector at the active
|
||||
/// signature (the post-switch / dim-mismatch state) is re-embedded by
|
||||
/// `handle_reembed_backfill`; the chain `Defer`s while work remains and
|
||||
/// returns `Done` once the space is covered; a stale-signature job
|
||||
/// finishes immediately without touching anything.
|
||||
///
|
||||
/// (The process-global `backfill_in_progress` flag is intentionally not
|
||||
/// asserted here — it is shared across parallel tests and set widely by
|
||||
/// the §7 trigger, so asserting it would be flaky. The handler's
|
||||
/// deterministic effects are what this test pins.)
|
||||
#[tokio::test]
|
||||
async fn reembed_backfill_repopulates_then_completes() {
|
||||
use crate::openhuman::memory_store::chunks::store::{
|
||||
get_chunk_embedding_for_signature, tree_active_signature, upsert_chunks,
|
||||
upsert_staged_chunks_tx,
|
||||
};
|
||||
use crate::openhuman::memory_store::chunks::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
|
||||
let (_tmp, cfg) = test_config();
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "reembed-seed"),
|
||||
content: "memory content about the phoenix migration project".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
token_count: 12,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
|
||||
// Stage the body to disk so `read_chunk_body` succeeds in the handler.
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let sig = tree_active_signature(&cfg);
|
||||
assert!(
|
||||
get_chunk_embedding_for_signature(&cfg, &chunk.id, &sig)
|
||||
.unwrap()
|
||||
.is_none(),
|
||||
"precondition: no sidecar vector at the active signature"
|
||||
);
|
||||
|
||||
// Work present → re-embed + write sidecar, Defer to revisit.
|
||||
let job = mk_running_job(
|
||||
JobKind::ReembedBackfill,
|
||||
serde_json::to_string(&ReembedBackfillPayload {
|
||||
signature: sig.clone(),
|
||||
})
|
||||
.unwrap(),
|
||||
);
|
||||
let out = handle_reembed_backfill(&cfg, &job).await.unwrap();
|
||||
assert!(
|
||||
matches!(out, JobOutcome::Defer { .. }),
|
||||
"work present must Defer (self-continue), got {out:?}"
|
||||
);
|
||||
assert!(
|
||||
get_chunk_embedding_for_signature(&cfg, &chunk.id, &sig)
|
||||
.unwrap()
|
||||
.is_some(),
|
||||
"chunk re-embedded into the sidecar at the active signature"
|
||||
);
|
||||
|
||||
// Nothing left → Done.
|
||||
let out2 = handle_reembed_backfill(&cfg, &job).await.unwrap();
|
||||
assert_eq!(out2, JobOutcome::Done, "covered space must complete");
|
||||
|
||||
// Stale signature (embedder changed since enqueue) → finishes
|
||||
// immediately, no work, no panic.
|
||||
let stale = mk_running_job(
|
||||
JobKind::ReembedBackfill,
|
||||
serde_json::to_string(&ReembedBackfillPayload {
|
||||
signature: "provider=other;model=x;dims=1".into(),
|
||||
})
|
||||
.unwrap(),
|
||||
);
|
||||
assert_eq!(
|
||||
handle_reembed_backfill(&cfg, &stale).await.unwrap(),
|
||||
JobOutcome::Done
|
||||
);
|
||||
}
|
||||
|
||||
/// #1574 §6 regression gate: a terminal-failure chunk (its body file is
|
||||
/// missing on disk, despite the metadata row staying staged) is
|
||||
/// persistently tombstoned by `mark_chunk_reembed_skipped` on the first
|
||||
/// pass, then excluded from the next batch's worklist so the chain
|
||||
/// terminates (`Done`) instead of looping forever. Without this guard
|
||||
/// the §6 runaway-loop fix would silently regress — the same 16 orphans
|
||||
/// → ~8k defers → ~128k warns symptom observed in the wild before the
|
||||
/// fix landed (see PR body and store.rs:1195).
|
||||
///
|
||||
/// What the test pins:
|
||||
/// 1. Tombstone row is written for the failing chunk (exactly one).
|
||||
/// 2. The next-batch worklist `NOT EXISTS … reembed_skipped` clause
|
||||
/// excludes the tombstoned row — the handler returns `Done`.
|
||||
/// 3. The `ensure_reembed_backfill` migration probe agrees the space
|
||||
/// is covered (or the chain would re-arm on every config save).
|
||||
#[tokio::test]
|
||||
async fn reembed_backfill_tombstones_orphan_and_terminates() {
|
||||
use crate::openhuman::memory_store::chunks::store::{
|
||||
get_chunk_content_path, get_chunk_embedding_for_signature, tree_active_signature,
|
||||
upsert_chunks, upsert_staged_chunks_tx,
|
||||
};
|
||||
use crate::openhuman::memory_store::chunks::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
|
||||
let (_tmp, cfg) = test_config();
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "orphan-seed"),
|
||||
content: "memory content about the orphaned phoenix project".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
token_count: 12,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
|
||||
|
||||
// Stage the body file + metadata, then DELETE the body file from
|
||||
// disk while leaving the staged DB rows intact. Reproduces the
|
||||
// in-wild failure mode: chunk row + path hash both present, but
|
||||
// the body content was lost (user moved workspace dirs, partial
|
||||
// backup restore, manual file cleanup). `stage_chunks` returns
|
||||
// paths relative to `content_root`; resolve absolute before unlink.
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
let staged_rel = get_chunk_content_path(&cfg, &chunk.id)
|
||||
.unwrap()
|
||||
.expect("staged body path");
|
||||
let body_abs = content_root.join(&staged_rel);
|
||||
std::fs::remove_file(&body_abs).unwrap();
|
||||
|
||||
let sig = tree_active_signature(&cfg);
|
||||
let job = mk_running_job(
|
||||
JobKind::ReembedBackfill,
|
||||
serde_json::to_string(&ReembedBackfillPayload {
|
||||
signature: sig.clone(),
|
||||
})
|
||||
.unwrap(),
|
||||
);
|
||||
|
||||
// Pass 1: worklist picks up the orphan, body read fails, tombstone
|
||||
// written, `Defer` to revisit (the handler doesn't distinguish
|
||||
// "all rows tombstoned" from "more rows pending" inside this batch).
|
||||
let out1 = handle_reembed_backfill(&cfg, &job).await.unwrap();
|
||||
assert!(
|
||||
matches!(out1, JobOutcome::Defer { .. }),
|
||||
"first pass should Defer after failing to read body, got {out1:?}"
|
||||
);
|
||||
assert!(
|
||||
get_chunk_embedding_for_signature(&cfg, &chunk.id, &sig)
|
||||
.unwrap()
|
||||
.is_none(),
|
||||
"orphan chunk must not have a sidecar vector after failure"
|
||||
);
|
||||
|
||||
// (1) Tombstone row exists for exactly this (chunk, sig).
|
||||
let tombstone_count: i64 = with_connection(&cfg, |conn| {
|
||||
Ok(conn.query_row(
|
||||
"SELECT COUNT(*) FROM mem_tree_chunk_reembed_skipped
|
||||
WHERE chunk_id = ?1 AND model_signature = ?2",
|
||||
params![chunk.id, sig],
|
||||
|r| r.get(0),
|
||||
)?)
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
tombstone_count, 1,
|
||||
"orphan chunk must be tombstoned exactly once"
|
||||
);
|
||||
|
||||
// (2) Pass 2: worklist NOT EXISTS clause excludes the tombstoned
|
||||
// row; both worklists empty; chain completes.
|
||||
let out2 = handle_reembed_backfill(&cfg, &job).await.unwrap();
|
||||
assert_eq!(
|
||||
out2,
|
||||
JobOutcome::Done,
|
||||
"tombstoned-only state must complete the chain"
|
||||
);
|
||||
|
||||
// (3) Migration probe in `ensure_reembed_backfill` must agree the
|
||||
// space is covered, otherwise the chain re-arms on every config
|
||||
// save and we're back to the original infinite-loop bug.
|
||||
let probe_uncovered = with_connection(&cfg, |conn| {
|
||||
Ok(chunk_store::has_uncovered_reembed_work(conn, &sig)?)
|
||||
})
|
||||
.unwrap();
|
||||
assert!(
|
||||
!probe_uncovered,
|
||||
"after tombstoning the only orphan, the ensure_reembed_backfill probe must report covered"
|
||||
);
|
||||
}
|
||||
|
||||
/// #2358: clearing a tombstone re-opens the row for the backfill worklist.
|
||||
#[tokio::test]
|
||||
async fn clear_chunk_reembed_skipped_reopens_worklist() {
|
||||
use crate::openhuman::memory_store::chunks::store::{
|
||||
clear_chunk_reembed_skipped, get_chunk_content_path, mark_chunk_reembed_skipped,
|
||||
tree_active_signature, upsert_chunks, upsert_staged_chunks_tx,
|
||||
};
|
||||
use crate::openhuman::memory_store::chunks::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
|
||||
let (_tmp, cfg) = test_config();
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "clear-tombstone-seed"),
|
||||
content: "memory content for clear tombstone test".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
token_count: 12,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
let staged_rel = get_chunk_content_path(&cfg, &chunk.id)
|
||||
.unwrap()
|
||||
.expect("staged body path");
|
||||
std::fs::remove_file(content_root.join(&staged_rel)).unwrap();
|
||||
|
||||
let sig = tree_active_signature(&cfg);
|
||||
mark_chunk_reembed_skipped(&cfg, &chunk.id, &sig, "orphan").unwrap();
|
||||
|
||||
let covered_before_clear = with_connection(&cfg, |conn| {
|
||||
Ok(!chunk_store::has_uncovered_reembed_work(conn, &sig)?)
|
||||
})
|
||||
.unwrap();
|
||||
assert!(
|
||||
covered_before_clear,
|
||||
"tombstone must hide orphan from uncovered probe"
|
||||
);
|
||||
|
||||
clear_chunk_reembed_skipped(&cfg, &chunk.id, &sig).unwrap();
|
||||
|
||||
let uncovered_after_clear = with_connection(&cfg, |conn| {
|
||||
Ok(chunk_store::has_uncovered_reembed_work(conn, &sig)?)
|
||||
})
|
||||
.unwrap();
|
||||
assert!(
|
||||
uncovered_after_clear,
|
||||
"clearing tombstone must re-include chunk in worklist probe"
|
||||
);
|
||||
}
|
||||
|
||||
/// #1574 §4: `ensure_reembed_backfill` (the switch-path trigger) enqueues
|
||||
/// exactly one chain when there is uncovered work, is idempotent on
|
||||
/// re-call (per-signature dedupe), and enqueues nothing for an
|
||||
/// empty/covered space.
|
||||
#[tokio::test]
|
||||
async fn ensure_reembed_backfill_enqueues_only_when_uncovered() {
|
||||
use crate::openhuman::memory_queue::ensure_reembed_backfill;
|
||||
use crate::openhuman::memory_store::chunks::store::{
|
||||
upsert_chunks, upsert_staged_chunks_tx,
|
||||
};
|
||||
use crate::openhuman::memory_store::chunks::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
|
||||
// Empty space → nothing to do → no job.
|
||||
let (_t0, empty_cfg) = test_config();
|
||||
ensure_reembed_backfill(&empty_cfg);
|
||||
assert_eq!(
|
||||
count_jobs_of_kind(&empty_cfg, "reembed_backfill"),
|
||||
0,
|
||||
"empty/covered space must not enqueue a backfill"
|
||||
);
|
||||
|
||||
// Chunk with content but no sidecar vector → exactly one chain.
|
||||
let (_t1, cfg) = test_config();
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "ensure-seed"),
|
||||
content: "memory content needing a re-embed".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
token_count: 12,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
ensure_reembed_backfill(&cfg);
|
||||
assert_eq!(
|
||||
count_jobs_of_kind(&cfg, "reembed_backfill"),
|
||||
1,
|
||||
"uncovered work must enqueue exactly one backfill chain"
|
||||
);
|
||||
// Idempotent — re-call must not create a second chain (dedupe by sig).
|
||||
ensure_reembed_backfill(&cfg);
|
||||
assert_eq!(
|
||||
count_jobs_of_kind(&cfg, "reembed_backfill"),
|
||||
1,
|
||||
"re-call must dedupe to a single chain per signature"
|
||||
);
|
||||
}
|
||||
}
|
||||
#[path = "mod_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,721 @@
|
||||
use super::*;
|
||||
use crate::openhuman::memory::tree_source::registry::get_or_create_source_tree;
|
||||
use crate::openhuman::memory_queue::store::{count_by_status, count_total};
|
||||
use crate::openhuman::memory_queue::types::JobStatus;
|
||||
use crate::openhuman::memory_store::chunks::store::with_connection;
|
||||
use crate::openhuman::memory_store::content as content_store;
|
||||
use crate::openhuman::memory_tree::tree::bucket_seal::{append_leaf_deferred, LeafRef};
|
||||
use crate::openhuman::memory_tree::tree::store as src_store;
|
||||
use chrono::TimeZone;
|
||||
use rusqlite::params;
|
||||
use tempfile::TempDir;
|
||||
|
||||
fn test_config() -> (TempDir, Config) {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let mut cfg = Config::default();
|
||||
cfg.workspace_dir = tmp.path().to_path_buf();
|
||||
cfg.memory_tree.embedding_endpoint = None;
|
||||
cfg.memory_tree.embedding_model = None;
|
||||
cfg.memory_tree.embedding_strict = false;
|
||||
(tmp, cfg)
|
||||
}
|
||||
|
||||
/// Build a minimal `Job` row for direct handler invocation. Mirrors
|
||||
/// what `claim_next` would produce for a freshly-claimed row.
|
||||
fn mk_running_job(kind: JobKind, payload_json: String) -> Job {
|
||||
let now_ms = chrono::Utc::now().timestamp_millis();
|
||||
Job {
|
||||
id: "test-job-id".into(),
|
||||
kind,
|
||||
payload_json,
|
||||
dedupe_key: None,
|
||||
status: JobStatus::Running,
|
||||
attempts: 1,
|
||||
max_attempts: 5,
|
||||
available_at_ms: now_ms,
|
||||
locked_until_ms: Some(now_ms + 60_000),
|
||||
last_error: None,
|
||||
created_at_ms: now_ms,
|
||||
started_at_ms: Some(now_ms),
|
||||
completed_at_ms: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Count rows in `mem_tree_jobs` matching a specific kind.
|
||||
fn count_jobs_of_kind(cfg: &Config, kind: &str) -> u64 {
|
||||
with_connection(cfg, |conn| {
|
||||
let n: i64 = conn.query_row(
|
||||
"SELECT COUNT(*) FROM mem_tree_jobs WHERE kind = ?1",
|
||||
params![kind],
|
||||
|r| r.get(0),
|
||||
)?;
|
||||
Ok(n.max(0) as u64)
|
||||
})
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
/// Seed a source tree and push enough labeled leaves into its L0 buffer
|
||||
/// to cross `INPUT_TOKEN_BUDGET`, returning the tree. The caller can then
|
||||
/// fire `handle_seal` and inspect the result.
|
||||
async fn seed_source_tree_ready_to_seal(
|
||||
cfg: &Config,
|
||||
) -> crate::openhuman::memory_store::trees::types::Tree {
|
||||
use crate::openhuman::memory_store::chunks::store::upsert_chunks;
|
||||
use crate::openhuman::memory_store::chunks::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
let tree = get_or_create_source_tree(cfg, "slack:#eng").unwrap();
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "handler-seed"),
|
||||
content: "alice@example.com leading the rollout".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
// Bust budget so the L0 buffer is "ready" for seal.
|
||||
token_count: 60_000,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(cfg, &[chunk.clone()]).unwrap();
|
||||
// Stage to disk so `hydrate_leaf_inputs` can read the full body via
|
||||
// `read_chunk_body` when `handle_seal` fires and calls `seal_one_level`.
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory_store::chunks::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
let leaf = LeafRef {
|
||||
chunk_id: chunk.id,
|
||||
token_count: 60_000,
|
||||
timestamp: ts,
|
||||
content: chunk.content,
|
||||
entities: vec![],
|
||||
topics: vec![],
|
||||
score: 0.5,
|
||||
};
|
||||
// append_leaf_deferred only buffers; doesn't seal. handle_seal will.
|
||||
let _ = append_leaf_deferred(cfg, &tree, &leaf).unwrap();
|
||||
tree
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn source_tree_seal_handler_enqueues_summary_topic_route() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
let tree = seed_source_tree_ready_to_seal(&cfg).await;
|
||||
|
||||
let payload = SealPayload {
|
||||
tree_id: tree.id.clone(),
|
||||
level: 0,
|
||||
force_now_ms: None,
|
||||
};
|
||||
let job = mk_running_job(JobKind::Seal, serde_json::to_string(&payload).unwrap());
|
||||
|
||||
// Pre-condition: queue has no topic_route jobs.
|
||||
assert_eq!(count_jobs_of_kind(&cfg, "topic_route"), 0);
|
||||
|
||||
super::handle_seal(&cfg, &job).await.unwrap();
|
||||
|
||||
// Post-condition: source-tree seal must enqueue exactly one
|
||||
// topic_route job carrying NodeRef::Summary { summary_id: <new> }.
|
||||
assert_eq!(
|
||||
count_jobs_of_kind(&cfg, "topic_route"),
|
||||
1,
|
||||
"source-tree seal must enqueue summary-side topic_route"
|
||||
);
|
||||
assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 1);
|
||||
|
||||
// Inspect the enqueued payload to confirm it's a Summary variant.
|
||||
let payload_json: String = with_connection(&cfg, |conn| {
|
||||
let s: String = conn
|
||||
.query_row(
|
||||
"SELECT payload_json FROM mem_tree_jobs WHERE kind = 'topic_route'",
|
||||
[],
|
||||
|r| r.get(0),
|
||||
)
|
||||
.unwrap();
|
||||
Ok(s)
|
||||
})
|
||||
.unwrap();
|
||||
let p: TopicRoutePayload = serde_json::from_str(&payload_json).unwrap();
|
||||
match p.node {
|
||||
NodeRef::Summary { summary_id } => {
|
||||
// Format: `summary:<13-digit-ms>:L<level>-<8hex>` —
|
||||
// see `tree::registry::new_summary_id`.
|
||||
assert!(
|
||||
summary_id.starts_with("summary:") && summary_id.contains(":L1-"),
|
||||
"expected summary id with L1 segment, got {summary_id}"
|
||||
);
|
||||
}
|
||||
other => panic!("expected NodeRef::Summary, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn topic_tree_seal_handler_does_not_enqueue_topic_route() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
// Spawn a topic tree directly via the registry (skipping curator's
|
||||
// hotness gate — we just need a TreeKind::Topic with leaves).
|
||||
let topic_tree = crate::openhuman::memory_store::trees::registry::get_or_create_topic_tree(
|
||||
&cfg,
|
||||
"topic:phoenix-migration",
|
||||
)
|
||||
.unwrap();
|
||||
// Push a single 10k-token leaf so L0 is gate-ready.
|
||||
use crate::openhuman::memory_store::chunks::store::upsert_chunks;
|
||||
use crate::openhuman::memory_store::chunks::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "topic-seed"),
|
||||
content: "topic content".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
token_count: 60_000,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
|
||||
// Stage to disk so `hydrate_leaf_inputs` can read the full body
|
||||
// when `handle_seal` fires.
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory_store::chunks::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
let leaf = LeafRef {
|
||||
chunk_id: chunk.id,
|
||||
token_count: 60_000,
|
||||
timestamp: ts,
|
||||
content: chunk.content,
|
||||
entities: vec![],
|
||||
topics: vec![],
|
||||
score: 0.5,
|
||||
};
|
||||
append_leaf_deferred(&cfg, &topic_tree, &leaf).unwrap();
|
||||
|
||||
let payload = SealPayload {
|
||||
tree_id: topic_tree.id.clone(),
|
||||
level: 0,
|
||||
force_now_ms: None,
|
||||
};
|
||||
let job = mk_running_job(JobKind::Seal, serde_json::to_string(&payload).unwrap());
|
||||
|
||||
super::handle_seal(&cfg, &job).await.unwrap();
|
||||
|
||||
// Topic-tree seals are sinks: must not enqueue any topic_route.
|
||||
assert_eq!(
|
||||
count_jobs_of_kind(&cfg, "topic_route"),
|
||||
0,
|
||||
"topic-tree seal must NOT enqueue topic_route (trees are sinks)"
|
||||
);
|
||||
// The seal itself should still have produced a summary node.
|
||||
assert_eq!(src_store::count_summaries(&cfg, &topic_tree.id).unwrap(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn handle_append_buffer_with_summary_payload_pushes_into_topic_tree() {
|
||||
let (_tmp, cfg) = test_config();
|
||||
|
||||
// 1. Create a target topic tree with a clean L0 buffer.
|
||||
let topic_tree = crate::openhuman::memory_store::trees::registry::get_or_create_topic_tree(
|
||||
&cfg,
|
||||
"email:alice@example.com",
|
||||
)
|
||||
.unwrap();
|
||||
let l0_before = src_store::get_buffer(&cfg, &topic_tree.id, 0).unwrap();
|
||||
assert!(l0_before.is_empty());
|
||||
|
||||
// 2. Manually insert a summary node we can route. The simplest way
|
||||
// is to create a separate source tree, push two 6k leaves into
|
||||
// it, and let the seal produce a summary we can address.
|
||||
let source_tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap();
|
||||
use crate::openhuman::memory_store::chunks::store::upsert_chunks;
|
||||
use crate::openhuman::memory_store::chunks::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
use crate::openhuman::memory_tree::tree::bucket_seal::seal_one_level;
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
for seq in 0..2 {
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", seq, "summary-seed"),
|
||||
content: format!("source content {seq}"),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
token_count: 30_000,
|
||||
seq_in_source: seq,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
|
||||
// Stage to disk so `hydrate_leaf_inputs` can read the full body
|
||||
// during `seal_one_level`.
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
crate::openhuman::memory_store::chunks::store::upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
let leaf = LeafRef {
|
||||
chunk_id: chunk.id,
|
||||
token_count: 30_000,
|
||||
timestamp: ts,
|
||||
content: chunk.content,
|
||||
entities: vec![],
|
||||
topics: vec![],
|
||||
score: 0.5,
|
||||
};
|
||||
let _ = append_leaf_deferred(&cfg, &source_tree, &leaf).unwrap();
|
||||
}
|
||||
// Force-seal the source tree's L0 to mint the summary.
|
||||
use crate::openhuman::memory::chat::{test_override, ChatProvider, StaticChatProvider};
|
||||
let buf = src_store::get_buffer(&cfg, &source_tree.id, 0).unwrap();
|
||||
let provider: std::sync::Arc<dyn ChatProvider> =
|
||||
std::sync::Arc::new(StaticChatProvider::new("test summary content"));
|
||||
let summary_id = test_override::with_provider(provider, async {
|
||||
seal_one_level(
|
||||
&cfg,
|
||||
&source_tree,
|
||||
&buf,
|
||||
&crate::openhuman::memory_tree::tree::bucket_seal::LabelStrategy::Empty,
|
||||
// No follow-up enqueues — the test scopes assertions to the
|
||||
// append_buffer handler, not seal-side fan-out.
|
||||
false,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
})
|
||||
.await;
|
||||
|
||||
// 3. Build an append_buffer payload routing the summary into the
|
||||
// topic tree.
|
||||
let payload = AppendBufferPayload {
|
||||
node: NodeRef::Summary {
|
||||
summary_id: summary_id.clone(),
|
||||
},
|
||||
target: AppendTarget::Topic {
|
||||
tree_id: topic_tree.id.clone(),
|
||||
},
|
||||
};
|
||||
let job = mk_running_job(
|
||||
JobKind::AppendBuffer,
|
||||
serde_json::to_string(&payload).unwrap(),
|
||||
);
|
||||
|
||||
// Clear out any pending append_buffer jobs minted upstream so the
|
||||
// post-condition assertion below is unambiguous.
|
||||
let pre = count_total(&cfg).unwrap();
|
||||
|
||||
super::handle_append_buffer(&cfg, &job).await.unwrap();
|
||||
|
||||
// 4. Topic tree's L0 buffer should now hold the summary id.
|
||||
let l0_after = src_store::get_buffer(&cfg, &topic_tree.id, 0).unwrap();
|
||||
assert_eq!(l0_after.item_ids, vec![summary_id]);
|
||||
assert!(l0_after.token_sum > 0);
|
||||
|
||||
// No new jobs should have been enqueued (buffer didn't cross gate).
|
||||
assert_eq!(count_total(&cfg).unwrap(), pre);
|
||||
}
|
||||
|
||||
/// #1574 §6: a chunk with content but no sidecar vector at the active
|
||||
/// signature (the post-switch / dim-mismatch state) is re-embedded by
|
||||
/// `handle_reembed_backfill`; the chain `Defer`s while work remains and
|
||||
/// returns `Done` once the space is covered; a stale-signature job
|
||||
/// finishes immediately without touching anything.
|
||||
///
|
||||
/// (The process-global `backfill_in_progress` flag is intentionally not
|
||||
/// asserted here — it is shared across parallel tests and set widely by
|
||||
/// the §7 trigger, so asserting it would be flaky. The handler's
|
||||
/// deterministic effects are what this test pins.)
|
||||
#[tokio::test]
|
||||
async fn reembed_backfill_repopulates_then_completes() {
|
||||
use crate::openhuman::memory_store::chunks::store::{
|
||||
get_chunk_embedding_for_signature, tree_active_signature, upsert_chunks,
|
||||
upsert_staged_chunks_tx,
|
||||
};
|
||||
use crate::openhuman::memory_store::chunks::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
|
||||
let (_tmp, cfg) = test_config();
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "reembed-seed"),
|
||||
content: "memory content about the phoenix migration project".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
token_count: 12,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
|
||||
// Stage the body to disk so `read_chunk_body` succeeds in the handler.
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
let sig = tree_active_signature(&cfg);
|
||||
assert!(
|
||||
get_chunk_embedding_for_signature(&cfg, &chunk.id, &sig)
|
||||
.unwrap()
|
||||
.is_none(),
|
||||
"precondition: no sidecar vector at the active signature"
|
||||
);
|
||||
|
||||
// Work present → re-embed + write sidecar, Defer to revisit.
|
||||
let job = mk_running_job(
|
||||
JobKind::ReembedBackfill,
|
||||
serde_json::to_string(&ReembedBackfillPayload {
|
||||
signature: sig.clone(),
|
||||
})
|
||||
.unwrap(),
|
||||
);
|
||||
let out = handle_reembed_backfill(&cfg, &job).await.unwrap();
|
||||
assert!(
|
||||
matches!(out, JobOutcome::Defer { .. }),
|
||||
"work present must Defer (self-continue), got {out:?}"
|
||||
);
|
||||
assert!(
|
||||
get_chunk_embedding_for_signature(&cfg, &chunk.id, &sig)
|
||||
.unwrap()
|
||||
.is_some(),
|
||||
"chunk re-embedded into the sidecar at the active signature"
|
||||
);
|
||||
|
||||
// Nothing left → Done.
|
||||
let out2 = handle_reembed_backfill(&cfg, &job).await.unwrap();
|
||||
assert_eq!(out2, JobOutcome::Done, "covered space must complete");
|
||||
|
||||
// Stale signature (embedder changed since enqueue) → finishes
|
||||
// immediately, no work, no panic.
|
||||
let stale = mk_running_job(
|
||||
JobKind::ReembedBackfill,
|
||||
serde_json::to_string(&ReembedBackfillPayload {
|
||||
signature: "provider=other;model=x;dims=1".into(),
|
||||
})
|
||||
.unwrap(),
|
||||
);
|
||||
assert_eq!(
|
||||
handle_reembed_backfill(&cfg, &stale).await.unwrap(),
|
||||
JobOutcome::Done
|
||||
);
|
||||
}
|
||||
|
||||
/// #1574 §6 regression gate: a terminal-failure chunk (its body file is
|
||||
/// missing on disk, despite the metadata row staying staged) is
|
||||
/// persistently tombstoned by `mark_chunk_reembed_skipped` on the first
|
||||
/// pass, then excluded from the next batch's worklist so the chain
|
||||
/// terminates (`Done`) instead of looping forever. Without this guard
|
||||
/// the §6 runaway-loop fix would silently regress — the same 16 orphans
|
||||
/// → ~8k defers → ~128k warns symptom observed in the wild before the
|
||||
/// fix landed (see PR body and store.rs:1195).
|
||||
///
|
||||
/// What the test pins:
|
||||
/// 1. Tombstone row is written for the failing chunk (exactly one).
|
||||
/// 2. The next-batch worklist `NOT EXISTS … reembed_skipped` clause
|
||||
/// excludes the tombstoned row — the handler returns `Done`.
|
||||
/// 3. The `ensure_reembed_backfill` migration probe agrees the space
|
||||
/// is covered (or the chain would re-arm on every config save).
|
||||
#[tokio::test]
|
||||
async fn reembed_backfill_tombstones_orphan_and_terminates() {
|
||||
use crate::openhuman::memory_store::chunks::store::{
|
||||
get_chunk_content_path, get_chunk_embedding_for_signature, tree_active_signature,
|
||||
upsert_chunks, upsert_staged_chunks_tx,
|
||||
};
|
||||
use crate::openhuman::memory_store::chunks::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
|
||||
let (_tmp, cfg) = test_config();
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "orphan-seed"),
|
||||
content: "memory content about the orphaned phoenix project".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
token_count: 12,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
|
||||
|
||||
// Stage the body file + metadata, then DELETE the body file from
|
||||
// disk while leaving the staged DB rows intact. Reproduces the
|
||||
// in-wild failure mode: chunk row + path hash both present, but
|
||||
// the body content was lost (user moved workspace dirs, partial
|
||||
// backup restore, manual file cleanup). `stage_chunks` returns
|
||||
// paths relative to `content_root`; resolve absolute before unlink.
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
let staged_rel = get_chunk_content_path(&cfg, &chunk.id)
|
||||
.unwrap()
|
||||
.expect("staged body path");
|
||||
let body_abs = content_root.join(&staged_rel);
|
||||
std::fs::remove_file(&body_abs).unwrap();
|
||||
|
||||
let sig = tree_active_signature(&cfg);
|
||||
let job = mk_running_job(
|
||||
JobKind::ReembedBackfill,
|
||||
serde_json::to_string(&ReembedBackfillPayload {
|
||||
signature: sig.clone(),
|
||||
})
|
||||
.unwrap(),
|
||||
);
|
||||
|
||||
// Pass 1: worklist picks up the orphan, body read fails, tombstone
|
||||
// written, `Defer` to revisit (the handler doesn't distinguish
|
||||
// "all rows tombstoned" from "more rows pending" inside this batch).
|
||||
let out1 = handle_reembed_backfill(&cfg, &job).await.unwrap();
|
||||
assert!(
|
||||
matches!(out1, JobOutcome::Defer { .. }),
|
||||
"first pass should Defer after failing to read body, got {out1:?}"
|
||||
);
|
||||
assert!(
|
||||
get_chunk_embedding_for_signature(&cfg, &chunk.id, &sig)
|
||||
.unwrap()
|
||||
.is_none(),
|
||||
"orphan chunk must not have a sidecar vector after failure"
|
||||
);
|
||||
|
||||
// (1) Tombstone row exists for exactly this (chunk, sig).
|
||||
let tombstone_count: i64 = with_connection(&cfg, |conn| {
|
||||
Ok(conn.query_row(
|
||||
"SELECT COUNT(*) FROM mem_tree_chunk_reembed_skipped
|
||||
WHERE chunk_id = ?1 AND model_signature = ?2",
|
||||
params![chunk.id, sig],
|
||||
|r| r.get(0),
|
||||
)?)
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
tombstone_count, 1,
|
||||
"orphan chunk must be tombstoned exactly once"
|
||||
);
|
||||
|
||||
// (2) Pass 2: worklist NOT EXISTS clause excludes the tombstoned
|
||||
// row; both worklists empty; chain completes.
|
||||
let out2 = handle_reembed_backfill(&cfg, &job).await.unwrap();
|
||||
assert_eq!(
|
||||
out2,
|
||||
JobOutcome::Done,
|
||||
"tombstoned-only state must complete the chain"
|
||||
);
|
||||
|
||||
// (3) Migration probe in `ensure_reembed_backfill` must agree the
|
||||
// space is covered, otherwise the chain re-arms on every config
|
||||
// save and we're back to the original infinite-loop bug.
|
||||
let probe_uncovered = with_connection(&cfg, |conn| {
|
||||
Ok(chunk_store::has_uncovered_reembed_work(conn, &sig)?)
|
||||
})
|
||||
.unwrap();
|
||||
assert!(
|
||||
!probe_uncovered,
|
||||
"after tombstoning the only orphan, the ensure_reembed_backfill probe must report covered"
|
||||
);
|
||||
}
|
||||
|
||||
/// #2358: clearing a tombstone re-opens the row for the backfill worklist.
|
||||
#[tokio::test]
|
||||
async fn clear_chunk_reembed_skipped_reopens_worklist() {
|
||||
use crate::openhuman::memory_store::chunks::store::{
|
||||
clear_chunk_reembed_skipped, get_chunk_content_path, mark_chunk_reembed_skipped,
|
||||
tree_active_signature, upsert_chunks, upsert_staged_chunks_tx,
|
||||
};
|
||||
use crate::openhuman::memory_store::chunks::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
|
||||
let (_tmp, cfg) = test_config();
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "clear-tombstone-seed"),
|
||||
content: "memory content for clear tombstone test".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
token_count: 12,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
let staged_rel = get_chunk_content_path(&cfg, &chunk.id)
|
||||
.unwrap()
|
||||
.expect("staged body path");
|
||||
std::fs::remove_file(content_root.join(&staged_rel)).unwrap();
|
||||
|
||||
let sig = tree_active_signature(&cfg);
|
||||
mark_chunk_reembed_skipped(&cfg, &chunk.id, &sig, "orphan").unwrap();
|
||||
|
||||
let covered_before_clear = with_connection(&cfg, |conn| {
|
||||
Ok(!chunk_store::has_uncovered_reembed_work(conn, &sig)?)
|
||||
})
|
||||
.unwrap();
|
||||
assert!(
|
||||
covered_before_clear,
|
||||
"tombstone must hide orphan from uncovered probe"
|
||||
);
|
||||
|
||||
clear_chunk_reembed_skipped(&cfg, &chunk.id, &sig).unwrap();
|
||||
|
||||
let uncovered_after_clear = with_connection(&cfg, |conn| {
|
||||
Ok(chunk_store::has_uncovered_reembed_work(conn, &sig)?)
|
||||
})
|
||||
.unwrap();
|
||||
assert!(
|
||||
uncovered_after_clear,
|
||||
"clearing tombstone must re-include chunk in worklist probe"
|
||||
);
|
||||
}
|
||||
|
||||
/// #1574 §4: `ensure_reembed_backfill` (the switch-path trigger) enqueues
|
||||
/// exactly one chain when there is uncovered work, is idempotent on
|
||||
/// re-call (per-signature dedupe), and enqueues nothing for an
|
||||
/// empty/covered space.
|
||||
#[tokio::test]
|
||||
async fn ensure_reembed_backfill_enqueues_only_when_uncovered() {
|
||||
use crate::openhuman::memory_queue::ensure_reembed_backfill;
|
||||
use crate::openhuman::memory_store::chunks::store::{upsert_chunks, upsert_staged_chunks_tx};
|
||||
use crate::openhuman::memory_store::chunks::types::{
|
||||
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
|
||||
};
|
||||
|
||||
// Empty space → nothing to do → no job.
|
||||
let (_t0, empty_cfg) = test_config();
|
||||
ensure_reembed_backfill(&empty_cfg);
|
||||
assert_eq!(
|
||||
count_jobs_of_kind(&empty_cfg, "reembed_backfill"),
|
||||
0,
|
||||
"empty/covered space must not enqueue a backfill"
|
||||
);
|
||||
|
||||
// Chunk with content but no sidecar vector → exactly one chain.
|
||||
let (_t1, cfg) = test_config();
|
||||
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
|
||||
let chunk = Chunk {
|
||||
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "ensure-seed"),
|
||||
content: "memory content needing a re-embed".into(),
|
||||
metadata: Metadata {
|
||||
source_kind: SourceKind::Chat,
|
||||
source_id: "slack:#eng".into(),
|
||||
owner: "alice".into(),
|
||||
timestamp: ts,
|
||||
time_range: (ts, ts),
|
||||
tags: vec![],
|
||||
source_ref: Some(SourceRef::new("slack://x")),
|
||||
},
|
||||
token_count: 12,
|
||||
seq_in_source: 0,
|
||||
created_at: ts,
|
||||
partial_message: false,
|
||||
};
|
||||
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
|
||||
let content_root = cfg.memory_tree_content_root();
|
||||
std::fs::create_dir_all(&content_root).unwrap();
|
||||
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
|
||||
with_connection(&cfg, |conn| {
|
||||
let tx = conn.unchecked_transaction()?;
|
||||
upsert_staged_chunks_tx(&tx, &staged)?;
|
||||
tx.commit()?;
|
||||
Ok(())
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
ensure_reembed_backfill(&cfg);
|
||||
assert_eq!(
|
||||
count_jobs_of_kind(&cfg, "reembed_backfill"),
|
||||
1,
|
||||
"uncovered work must enqueue exactly one backfill chain"
|
||||
);
|
||||
// Idempotent — re-call must not create a second chain (dedupe by sig).
|
||||
ensure_reembed_backfill(&cfg);
|
||||
assert_eq!(
|
||||
count_jobs_of_kind(&cfg, "reembed_backfill"),
|
||||
1,
|
||||
"re-call must dedupe to a single chain per signature"
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,526 @@
|
||||
use anyhow::{Context, Result};
|
||||
use parking_lot::Mutex as PMutex;
|
||||
use rusqlite::Connection;
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
|
||||
#[cfg(test)]
|
||||
use std::sync::Mutex;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
|
||||
use super::{
|
||||
add_column_if_missing, migrate_legacy_embeddings_to_sidecar, DB_DIR, DB_FILE, SCHEMA,
|
||||
SQLITE_BUSY_TIMEOUT,
|
||||
};
|
||||
|
||||
// ── Schema-apply instrumentation (test-only) ─────────────────────────────────
|
||||
//
|
||||
// Per-path counter of how many times `apply_schema` ran for each DB path,
|
||||
// gated behind `cfg(test)` so the production binary carries no overhead.
|
||||
// Used by the concurrent-init regression test to assert "exactly once per
|
||||
// path" across racing workers; it survives even when the connection cache
|
||||
// is cleared between tests because tests use distinct tempdirs.
|
||||
#[cfg(test)]
|
||||
static SCHEMA_APPLY_COUNTS: OnceLock<Mutex<HashMap<PathBuf, usize>>> = OnceLock::new();
|
||||
|
||||
fn record_schema_apply(_path: &Path) {
|
||||
#[cfg(test)]
|
||||
{
|
||||
let counts = SCHEMA_APPLY_COUNTS.get_or_init(|| Mutex::new(HashMap::new()));
|
||||
let mut guard = counts
|
||||
.lock()
|
||||
.expect("memory_tree schema apply count mutex poisoned");
|
||||
*guard.entry(_path.to_path_buf()).or_insert(0) += 1;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[doc(hidden)]
|
||||
pub(crate) fn schema_apply_count_for_path_for_tests(path: &Path) -> usize {
|
||||
SCHEMA_APPLY_COUNTS
|
||||
.get()
|
||||
.and_then(|m| {
|
||||
m.lock()
|
||||
.ok()
|
||||
.map(|guard| guard.get(path).copied().unwrap_or(0))
|
||||
})
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
// SQLite extended result codes that fire during cold-start WAL/SHM bootstrap
|
||||
// races. NOTE on values: extended codes are `SQLITE_IOERR (10) | (sub << 8)`.
|
||||
// 4874 is `IOERR_SHMSIZE` (sub 19), NOT `SHMMAP` — the real `SHMMAP` is 5386
|
||||
// (sub 21) and the "open a new shared-memory segment" failure is `SHMOPEN`
|
||||
// 4618 (sub 18), which is what surfaced on macOS. The whole `-shm` family is
|
||||
// listed so the classifiers don't miss any of them.
|
||||
/// `CANTOPEN` — racing the lockfile/WAL creation done by another connection.
|
||||
const SQLITE_CANTOPEN: i32 = 14;
|
||||
/// `IOERR_TRUNCATE` — the WAL/db is being truncated during bootstrap.
|
||||
const SQLITE_IOERR_TRUNCATE: i32 = 1546;
|
||||
/// `IOERR_SHMOPEN` — opening a new `-shm` shared-memory segment failed (the
|
||||
/// macOS cold-start failure, e.g. Sentry TAURI-RUST-X1).
|
||||
const SQLITE_IOERR_SHMOPEN: i32 = 4618;
|
||||
/// `IOERR_SHMSIZE` — the `-shm` file is being resized during bootstrap.
|
||||
const SQLITE_IOERR_SHMSIZE: i32 = 4874;
|
||||
/// `IOERR_SHMMAP` — mapping a page of the `-shm` wal-index failed.
|
||||
const SQLITE_IOERR_SHMMAP: i32 = 5386;
|
||||
/// `IOERR_IN_PAGE` — an mmap-page I/O fault, also seen under WAL cold-start.
|
||||
const SQLITE_IOERR_IN_PAGE: i32 = 8714;
|
||||
|
||||
/// True if `err` (or anything in its cause chain) is one of the SQLite codes
|
||||
/// that fire during cold-start WAL/SHM bootstrap races: `CANTOPEN`,
|
||||
/// `IOERR_TRUNCATE`, the `-shm` family (`SHMOPEN` / `SHMSIZE` / `SHMMAP`), and
|
||||
/// `IOERR_IN_PAGE`.
|
||||
pub(crate) fn is_transient_cold_start(err: &anyhow::Error) -> bool {
|
||||
fn is_transient_sqlite(e: &(dyn std::error::Error + 'static)) -> bool {
|
||||
if let Some(rusqlite::Error::SqliteFailure(ffi, _)) = e.downcast_ref::<rusqlite::Error>() {
|
||||
return matches!(
|
||||
ffi.extended_code,
|
||||
SQLITE_CANTOPEN
|
||||
| SQLITE_IOERR_TRUNCATE
|
||||
| SQLITE_IOERR_SHMOPEN
|
||||
| SQLITE_IOERR_SHMSIZE
|
||||
| SQLITE_IOERR_SHMMAP
|
||||
| SQLITE_IOERR_IN_PAGE
|
||||
);
|
||||
}
|
||||
false
|
||||
}
|
||||
if is_transient_sqlite(err.root_cause()) {
|
||||
return true;
|
||||
}
|
||||
let mut src: Option<&(dyn std::error::Error + 'static)> = Some(err.as_ref());
|
||||
while let Some(cur) = src {
|
||||
if is_transient_sqlite(cur) {
|
||||
return true;
|
||||
}
|
||||
src = cur.source();
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
// ── Connection cache (#2206) ─────────────────────────────────────────────────
|
||||
|
||||
/// How many consecutive init failures before the circuit breaker trips.
|
||||
pub(crate) const CB_THRESHOLD: u32 = 3;
|
||||
/// How long the circuit breaker holds the DB closed after tripping.
|
||||
pub(crate) const CB_COOLDOWN: Duration = Duration::from_secs(30);
|
||||
|
||||
/// Per-path circuit breaker: after [`CB_THRESHOLD`] consecutive init failures
|
||||
/// the breaker trips and `get_or_init_connection` returns an error immediately
|
||||
/// until [`CB_COOLDOWN`] elapses. On the first success it resets to zero.
|
||||
struct CircuitBreaker {
|
||||
consecutive_failures: AtomicU32,
|
||||
tripped: AtomicBool,
|
||||
last_trip: PMutex<Option<Instant>>,
|
||||
}
|
||||
|
||||
impl CircuitBreaker {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
consecutive_failures: AtomicU32::new(0),
|
||||
tripped: AtomicBool::new(false),
|
||||
last_trip: PMutex::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn record_success(&self) {
|
||||
self.consecutive_failures.store(0, Ordering::Relaxed);
|
||||
self.tripped.store(false, Ordering::Relaxed);
|
||||
*self.last_trip.lock() = None;
|
||||
}
|
||||
|
||||
/// Records one more failure. Returns `true` if this call just tripped the
|
||||
/// breaker (i.e. the threshold was crossed right now).
|
||||
fn record_failure(&self) -> bool {
|
||||
let prev = self.consecutive_failures.fetch_add(1, Ordering::Relaxed);
|
||||
let count = prev + 1;
|
||||
if count >= CB_THRESHOLD && !self.tripped.swap(true, Ordering::Relaxed) {
|
||||
*self.last_trip.lock() = Some(Instant::now());
|
||||
return true;
|
||||
}
|
||||
// Re-arm the cooldown on each subsequent failure while already tripped
|
||||
// so that a failed post-cooldown retry restarts the 30 s window instead
|
||||
// of leaving the stale timestamp in place (which would let `is_open`
|
||||
// return false immediately and allow unlimited retries).
|
||||
if self.tripped.load(Ordering::Relaxed) {
|
||||
*self.last_trip.lock() = Some(Instant::now());
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// Returns `true` when the breaker is open AND the cooldown has not yet
|
||||
/// elapsed. Returns `false` (allowing a retry) once the cooldown passes.
|
||||
fn is_open(&self) -> bool {
|
||||
if !self.tripped.load(Ordering::Relaxed) {
|
||||
return false;
|
||||
}
|
||||
let guard = self.last_trip.lock();
|
||||
match *guard {
|
||||
Some(t) if t.elapsed() < CB_COOLDOWN => true,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Process-level cache — two separate maps so that a failing init cannot
|
||||
/// accidentally serve a dummy connection on the fast path.
|
||||
///
|
||||
/// `connections`: only fully-initialised (schema + migrations run) entries.
|
||||
/// `breakers`: persists across failed init attempts so the circuit breaker
|
||||
/// survives even when `connections` has no entry for that path.
|
||||
struct ConnectionCache {
|
||||
connections: PMutex<HashMap<PathBuf, Arc<PMutex<Connection>>>>,
|
||||
breakers: PMutex<HashMap<PathBuf, Arc<CircuitBreaker>>>,
|
||||
/// Per-path mutex held across the slow-path init so concurrent
|
||||
/// workers racing into `with_connection` on a cold DB serialise on
|
||||
/// the WAL+SHM bootstrap. Without this, N threads see "no cached
|
||||
/// connection" simultaneously and all run `apply_schema`, which is
|
||||
/// idempotent but reopens the cold-start race window
|
||||
/// (OPENHUMAN-TAURI-HH / -ZM / -MB).
|
||||
init_locks: PMutex<HashMap<PathBuf, Arc<PMutex<()>>>>,
|
||||
}
|
||||
|
||||
static CONN_CACHE: OnceLock<ConnectionCache> = OnceLock::new();
|
||||
|
||||
fn conn_cache() -> &'static ConnectionCache {
|
||||
CONN_CACHE.get_or_init(|| ConnectionCache {
|
||||
connections: PMutex::new(HashMap::new()),
|
||||
breakers: PMutex::new(HashMap::new()),
|
||||
init_locks: PMutex::new(HashMap::new()),
|
||||
})
|
||||
}
|
||||
|
||||
/// Compute the canonical DB path from `config`.
|
||||
pub(crate) fn db_path_for(config: &Config) -> PathBuf {
|
||||
config.workspace_dir.join(DB_DIR).join(DB_FILE)
|
||||
}
|
||||
|
||||
/// Delete stale WAL/SHM side-files (`<db>-shm`, `<db>-wal`) that can block a
|
||||
/// clean DB open after a crash. Logs what was deleted and returns `true` if
|
||||
/// anything was removed.
|
||||
///
|
||||
/// SQLite names these files `<db_path>-shm` and `<db_path>-wal`.
|
||||
/// For `chunks.db` that is `chunks.db-shm` / `chunks.db-wal`.
|
||||
pub(crate) fn try_cleanup_stale_files(db_path: &std::path::Path) -> bool {
|
||||
let mut cleaned = false;
|
||||
for suffix in &["-shm", "-wal"] {
|
||||
// Build the side-file path: append suffix to the full db filename.
|
||||
let side = {
|
||||
let mut p = db_path.to_path_buf();
|
||||
let name = p
|
||||
.file_name()
|
||||
.unwrap_or_default()
|
||||
.to_string_lossy()
|
||||
.into_owned();
|
||||
p.set_file_name(format!("{name}{suffix}"));
|
||||
p
|
||||
};
|
||||
if side.exists() {
|
||||
match std::fs::remove_file(&side) {
|
||||
Ok(()) => {
|
||||
log::warn!("[memory_tree] removed stale side-file: {}", side.display());
|
||||
cleaned = true;
|
||||
}
|
||||
Err(e) => {
|
||||
log::warn!(
|
||||
"[memory_tree] failed to remove stale side-file {}: {e}",
|
||||
side.display()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
cleaned
|
||||
}
|
||||
|
||||
/// Run the full one-time DB initialisation (journal mode, schema, migrations)
|
||||
/// against an already-open `Connection`. Used by `get_or_init_connection`.
|
||||
fn init_db(conn: &Connection, config: &Config) -> Result<()> {
|
||||
conn.busy_timeout(SQLITE_BUSY_TIMEOUT)
|
||||
.context("Failed to configure memory_tree busy timeout")?;
|
||||
// SQLite resets `foreign_keys` to off on every new connection. The
|
||||
// ConnectionCache holds one cached `Connection` per DB path, so
|
||||
// setting it here (alongside the rest of init) is the per-connection
|
||||
// surface — fast-path callers reuse the cached conn with FKs already
|
||||
// on.
|
||||
conn.execute_batch("PRAGMA foreign_keys = ON;")
|
||||
.context("Failed to enable memory_tree foreign_keys pragma")?;
|
||||
// memory_tree runs the TRUNCATE rollback journal (see `apply_schema`), so
|
||||
// crash-safety requires synchronous=FULL — NORMAL is only corruption-safe
|
||||
// under WAL. Set explicitly so a future global default can't weaken it.
|
||||
conn.execute_batch("PRAGMA synchronous = FULL;")
|
||||
.context("Failed to set memory_tree synchronous=FULL")?;
|
||||
apply_schema(conn)?;
|
||||
// #1574 §7: one-shot, version-gated legacy→sidecar embedding migration.
|
||||
migrate_legacy_embeddings_to_sidecar(conn, config)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn apply_schema(conn: &Connection) -> Result<()> {
|
||||
// Note: `init_db` runs the `#1574 §7` legacy→sidecar embedding migration
|
||||
// after this returns, so the dim-equal copy step is not duplicated here.
|
||||
// memory_tree uses the TRUNCATE rollback journal, NOT WAL. WAL's `-shm`
|
||||
// shared-memory index + `-wal` checkpoint machinery are the root of the
|
||||
// cold-start IOERR_SHMMAP (macOS) / IOERR_TRUNCATE (Windows, AV-held
|
||||
// handles) failures (Sentry TAURI-RUST-EV / TAURI-RUST-X1). All tree
|
||||
// access serialises on the single cached `PMutex<Connection>` (see
|
||||
// `get_or_init_connection`), so WAL's only real benefit — concurrent
|
||||
// readers — is unused here, which makes WAL pure liability. The sibling
|
||||
// tree DBs (cron / vault / redirect_links) already run the default
|
||||
// rollback journal without issue.
|
||||
//
|
||||
// Requesting TRUNCATE on a database a prior release left in WAL mode
|
||||
// checkpoints the `-wal` back into the main file and removes the
|
||||
// `-wal`/`-shm` side-files, so this also migrates existing WAL databases
|
||||
// in place on upgrade.
|
||||
let journal_mode: String = conn
|
||||
.query_row("PRAGMA journal_mode=TRUNCATE", [], |row| row.get(0))
|
||||
.context("Failed to set memory_tree journal_mode=TRUNCATE")?;
|
||||
if !journal_mode.eq_ignore_ascii_case("truncate") {
|
||||
log::warn!(
|
||||
"[memory_tree] journal_mode is '{journal_mode}' after requesting TRUNCATE \
|
||||
— a prior WAL connection or a locked -wal may be blocking the switch"
|
||||
);
|
||||
}
|
||||
conn.execute_batch(SCHEMA)
|
||||
.context("Failed to initialize memory_tree schema")?;
|
||||
// Phase 2 migrations — additive, idempotent.
|
||||
add_column_if_missing(conn, "mem_tree_chunks", "embedding", "BLOB")?;
|
||||
// Phase 2 LLM-NER follow-up: per-chunk LLM importance signal +
|
||||
// human-readable reason. Both nullable; absence is treated as
|
||||
// "no LLM signal available" by readers.
|
||||
add_column_if_missing(conn, "mem_tree_score", "llm_importance", "REAL")?;
|
||||
add_column_if_missing(conn, "mem_tree_score", "llm_importance_reason", "TEXT")?;
|
||||
// Phase 3a (#709): parent-summary backlink on leaves.
|
||||
add_column_if_missing(conn, "mem_tree_chunks", "parent_summary_id", "TEXT")?;
|
||||
// Phase 4 (#710): sealed-summary embeddings for semantic rerank.
|
||||
add_column_if_missing(conn, "mem_tree_summaries", "embedding", "BLOB")?;
|
||||
// Async-pipeline lifecycle flag. Default 'admitted' so chunks ingested
|
||||
// before the queue migration stay queryable.
|
||||
add_column_if_missing(
|
||||
conn,
|
||||
"mem_tree_chunks",
|
||||
"lifecycle_status",
|
||||
"TEXT NOT NULL DEFAULT 'admitted'",
|
||||
)?;
|
||||
conn.execute_batch(
|
||||
"CREATE INDEX IF NOT EXISTS idx_mem_tree_chunks_lifecycle \
|
||||
ON mem_tree_chunks(lifecycle_status);",
|
||||
)
|
||||
.context("Failed to create mem_tree_chunks lifecycle index")?;
|
||||
// Phase MD-content (#TBD): pointer + integrity hash.
|
||||
add_column_if_missing(conn, "mem_tree_chunks", "content_path", "TEXT")?;
|
||||
add_column_if_missing(conn, "mem_tree_chunks", "content_sha256", "TEXT")?;
|
||||
// Phase MD-content (summaries).
|
||||
add_column_if_missing(conn, "mem_tree_summaries", "content_path", "TEXT")?;
|
||||
add_column_if_missing(conn, "mem_tree_summaries", "content_sha256", "TEXT")?;
|
||||
// Raw-archive pointer column.
|
||||
add_column_if_missing(conn, "mem_tree_chunks", "raw_refs_json", "TEXT")?;
|
||||
// #1365: is_user flag on indexed entity rows.
|
||||
add_column_if_missing(
|
||||
conn,
|
||||
"mem_tree_entity_index",
|
||||
"is_user",
|
||||
"INTEGER NOT NULL DEFAULT 0",
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Whether `err` looks like one of the I/O error codes that warrant a
|
||||
/// stale-file cleanup + single retry before giving up.
|
||||
pub(crate) fn is_io_open_error(err: &anyhow::Error) -> bool {
|
||||
if let Some(rusqlite::Error::SqliteFailure(f, _)) = err.downcast_ref::<rusqlite::Error>() {
|
||||
return matches!(
|
||||
f.extended_code,
|
||||
SQLITE_CANTOPEN
|
||||
| SQLITE_IOERR_TRUNCATE
|
||||
| SQLITE_IOERR_SHMOPEN
|
||||
| SQLITE_IOERR_SHMSIZE
|
||||
| SQLITE_IOERR_SHMMAP
|
||||
| SQLITE_IOERR_IN_PAGE
|
||||
) || f.code == rusqlite::ErrorCode::CannotOpen;
|
||||
}
|
||||
let msg = format!("{err:#}").to_ascii_lowercase();
|
||||
msg.contains("disk i/o error")
|
||||
|| msg.contains("unable to open database file")
|
||||
|| msg.contains("xshmmap")
|
||||
|| msg.contains("truncate file")
|
||||
}
|
||||
|
||||
/// Obtain (or lazily create) a cached connection for the workspace described
|
||||
/// by `config`. Returns `Err` immediately when the circuit breaker is open.
|
||||
pub(crate) fn get_or_init_connection(config: &Config) -> Result<Arc<PMutex<Connection>>> {
|
||||
let db_path = db_path_for(config);
|
||||
|
||||
// ── Fast path: reject immediately if the breaker is open ─────────────
|
||||
{
|
||||
let breakers = conn_cache().breakers.lock();
|
||||
if let Some(breaker) = breakers.get(&db_path) {
|
||||
if breaker.is_open() {
|
||||
anyhow::bail!(
|
||||
"[memory_tree] circuit breaker open for {}: too many consecutive init failures",
|
||||
db_path.display()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
// ── Fast path: return cached connection if already initialised ────────
|
||||
{
|
||||
let guard = conn_cache().connections.lock();
|
||||
if let Some(conn) = guard.get(&db_path) {
|
||||
return Ok(Arc::clone(conn));
|
||||
}
|
||||
}
|
||||
|
||||
// ── Slow path: serialise init per-path so concurrent workers don't
|
||||
// all race into `open_and_init` on a cold DB.
|
||||
let init_lock = {
|
||||
let mut guard = conn_cache().init_locks.lock();
|
||||
guard
|
||||
.entry(db_path.clone())
|
||||
.or_insert_with(|| Arc::new(PMutex::new(())))
|
||||
.clone()
|
||||
};
|
||||
let _init_guard = init_lock.lock();
|
||||
|
||||
// Re-check the cache once we hold the init lock — another thread
|
||||
// may have completed init while we were queued.
|
||||
{
|
||||
let guard = conn_cache().connections.lock();
|
||||
if let Some(conn) = guard.get(&db_path) {
|
||||
return Ok(Arc::clone(conn));
|
||||
}
|
||||
}
|
||||
|
||||
log::debug!(
|
||||
"[memory_tree] opening and initialising DB at {}",
|
||||
db_path.display()
|
||||
);
|
||||
|
||||
// Attempt to open + init the connection (dir creation is inside
|
||||
// `open_and_init` so every failure — including EEXIST on the dir —
|
||||
// reaches the circuit-breaker recording logic below). On certain I/O
|
||||
// errors (#2206) we clean up stale WAL/SHM side-files and retry once.
|
||||
let conn = open_and_init(&db_path, config).or_else(|first_err| {
|
||||
if is_io_open_error(&first_err) {
|
||||
log::warn!(
|
||||
"[memory_tree] I/O error on first open attempt ({}), cleaning stale files and retrying",
|
||||
first_err
|
||||
);
|
||||
try_cleanup_stale_files(&db_path);
|
||||
open_and_init(&db_path, config)
|
||||
} else {
|
||||
Err(first_err)
|
||||
}
|
||||
});
|
||||
|
||||
match conn {
|
||||
Ok(conn) => {
|
||||
let arc_conn = Arc::new(PMutex::new(conn));
|
||||
conn_cache()
|
||||
.connections
|
||||
.lock()
|
||||
.insert(db_path.clone(), Arc::clone(&arc_conn));
|
||||
// Reset any prior failure counter now that init succeeded.
|
||||
if let Some(breaker) = conn_cache().breakers.lock().get(&db_path) {
|
||||
breaker.record_success();
|
||||
}
|
||||
log::debug!("[memory_tree] DB connection cached and ready");
|
||||
Ok(arc_conn)
|
||||
}
|
||||
Err(err) => {
|
||||
// Persist the breaker so the failure count accumulates across
|
||||
// calls even though no connection entry exists yet.
|
||||
let breaker = {
|
||||
let mut guard = conn_cache().breakers.lock();
|
||||
guard
|
||||
.entry(db_path.clone())
|
||||
.or_insert_with(|| Arc::new(CircuitBreaker::new()))
|
||||
.clone()
|
||||
};
|
||||
let just_tripped = breaker.record_failure();
|
||||
if just_tripped {
|
||||
log::error!(
|
||||
"[memory_tree] circuit breaker tripped for {}: {} consecutive init failures",
|
||||
db_path.display(),
|
||||
CB_THRESHOLD
|
||||
);
|
||||
let _ = crate::core::event_bus::publish_global(
|
||||
crate::core::event_bus::DomainEvent::HealthChanged {
|
||||
component: "memory_tree_db".to_string(),
|
||||
healthy: false,
|
||||
message: Some(format!(
|
||||
"Schema init failed {CB_THRESHOLD} consecutive times"
|
||||
)),
|
||||
},
|
||||
);
|
||||
}
|
||||
Err(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Ensure the DB directory exists, open the SQLite file, and run the full
|
||||
/// schema init sequence. All errors (dir creation, file open, schema init)
|
||||
/// are returned as `Err` so callers can funnel them through the circuit
|
||||
/// breaker logic in a single place.
|
||||
fn open_and_init(db_path: &std::path::Path, config: &Config) -> Result<Connection> {
|
||||
let dir = db_path.parent().expect("db_path always has a parent");
|
||||
std::fs::create_dir_all(dir)
|
||||
.with_context(|| format!("Failed to create memory_tree dir: {}", dir.display()))?;
|
||||
let conn = Connection::open(db_path)
|
||||
.with_context(|| format!("Failed to open memory_tree DB: {}", db_path.display()))?;
|
||||
init_db(&conn, config)
|
||||
.with_context(|| format!("Failed to init memory_tree schema: {}", db_path.display()))?;
|
||||
record_schema_apply(db_path);
|
||||
Ok(conn)
|
||||
}
|
||||
|
||||
/// Remove the cached connection for `config`'s workspace (forces a fresh open
|
||||
/// on the next `with_connection` call). Also clears the breaker so the next
|
||||
/// open attempt is not immediately rejected. Does nothing if no entry exists.
|
||||
#[allow(dead_code)]
|
||||
pub(crate) fn invalidate_connection(config: &Config) {
|
||||
let db_path = db_path_for(config);
|
||||
conn_cache().connections.lock().remove(&db_path);
|
||||
conn_cache().breakers.lock().remove(&db_path);
|
||||
log::debug!(
|
||||
"[memory_tree] connection invalidated for {}",
|
||||
db_path.display()
|
||||
);
|
||||
}
|
||||
|
||||
/// Clear the entire connection cache. For test isolation only.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn clear_connection_cache() {
|
||||
conn_cache().connections.lock().clear();
|
||||
conn_cache().breakers.lock().clear();
|
||||
conn_cache().init_locks.lock().clear();
|
||||
}
|
||||
|
||||
/// Open the memory_tree SQLite DB and run a closure against it.
|
||||
///
|
||||
/// Visible to sibling modules (e.g. `score::store`) so Phase 2 can reuse
|
||||
/// the same connection setup / schema initialisation without duplication.
|
||||
///
|
||||
/// # Connection caching (#2206)
|
||||
///
|
||||
/// The underlying connection is initialised once per workspace path and then
|
||||
/// reused from a process-level cache. Schema migrations run exactly once on
|
||||
/// the first call for a given `config.workspace_dir`. Subsequent calls pay
|
||||
/// only the cost of a `parking_lot::Mutex` lock and the closure itself.
|
||||
///
|
||||
/// `#[doc(hidden)] pub` (not `pub(crate)`) because the
|
||||
/// `memory-tree-init-smoke` bin in `src/bin/` is a separate crate target
|
||||
/// and must reach this entry point. It is NOT a stable API surface —
|
||||
/// downstream crates should treat it as internal.
|
||||
#[doc(hidden)]
|
||||
pub fn with_connection<T>(config: &Config, f: impl FnOnce(&Connection) -> Result<T>) -> Result<T> {
|
||||
let conn_arc = get_or_init_connection(config)?;
|
||||
let guard = conn_arc.lock();
|
||||
f(&guard)
|
||||
}
|
||||
@@ -0,0 +1,289 @@
|
||||
use super::with_connection;
|
||||
use crate::openhuman::config::Config;
|
||||
use anyhow::{Context, Result};
|
||||
use chrono::Utc;
|
||||
use rusqlite::{Connection, OptionalExtension};
|
||||
|
||||
// ── Phase 2: embedding column accessors ─────────────────────────────────
|
||||
|
||||
/// Resolve the active embedding signature for the memory tree from the global
|
||||
/// [`Config`] — the canonical key every per-model sidecar read/write is scoped
|
||||
/// by (#1574). Reuses the established local-AI workload derivation
|
||||
/// ([`Config::workload_local_model`]) and the probe-stable
|
||||
/// `active_embedding_signature`; introduces no parallel resolution path.
|
||||
/// `pub(crate)` so the sibling `tree` summary store shares the exact
|
||||
/// same resolution.
|
||||
pub(crate) fn tree_active_signature(config: &Config) -> String {
|
||||
let local_model = config.workload_local_model("embeddings");
|
||||
crate::openhuman::memory_store::active_embedding_signature(
|
||||
&config.memory,
|
||||
local_model.as_deref(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Store a chunk's embedding under the active model signature.
|
||||
///
|
||||
/// #1574 cutover: this now writes the per-model `mem_tree_chunk_embeddings`
|
||||
/// sidecar (via [`set_chunk_embedding_for_signature`]) instead of the legacy
|
||||
/// `mem_tree_chunks.embedding` column. Call sites are unchanged — the signature
|
||||
/// is resolved internally from `config`. The legacy column is left intact for
|
||||
/// the §7 one-shot migration to read; it is dropped only in a later release.
|
||||
pub fn set_chunk_embedding(config: &Config, chunk_id: &str, embedding: &[f32]) -> Result<()> {
|
||||
let signature = tree_active_signature(config);
|
||||
log::debug!(
|
||||
"[memory::chunk_store] set_chunk_embedding: chunk_id={chunk_id} sig={signature} dims={}",
|
||||
embedding.len()
|
||||
);
|
||||
set_chunk_embedding_for_signature(config, chunk_id, &signature, embedding)
|
||||
}
|
||||
|
||||
/// Core upsert into `mem_tree_chunk_embeddings` over an arbitrary
|
||||
/// `&Connection`. Shared by the standalone ([`set_chunk_embedding_for_signature`])
|
||||
/// and in-transaction ([`set_chunk_embedding_for_signature_tx`]) write paths so
|
||||
/// the SQL exists exactly once. `rusqlite::Transaction` derefs to `Connection`,
|
||||
/// so an in-tx caller passes `&tx` and the sidecar row commits atomically with
|
||||
/// the surrounding work (#1574 write-side cutover).
|
||||
fn upsert_chunk_embedding_conn(
|
||||
conn: &rusqlite::Connection,
|
||||
chunk_id: &str,
|
||||
model_signature: &str,
|
||||
embedding: &[f32],
|
||||
) -> Result<()> {
|
||||
let bytes = embedding_to_blob(embedding);
|
||||
let dim = i64::try_from(embedding.len()).context("embedding dimension does not fit i64")?;
|
||||
let created_at = Utc::now().timestamp_millis() as f64 / 1000.0;
|
||||
conn.execute(
|
||||
"INSERT INTO mem_tree_chunk_embeddings
|
||||
(chunk_id, model_signature, vector, dim, created_at)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5)
|
||||
ON CONFLICT(chunk_id, model_signature) DO UPDATE SET
|
||||
vector = excluded.vector,
|
||||
dim = excluded.dim,
|
||||
created_at = excluded.created_at",
|
||||
rusqlite::params![chunk_id, model_signature, bytes, dim, created_at],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Store a chunk embedding for a specific provider/model/dimension signature.
|
||||
///
|
||||
/// Per-model table write path for #1574. The legacy
|
||||
/// `mem_tree_chunks.embedding` column is intentionally left untouched by this
|
||||
/// helper (read by the §7 migration; dropped only in a later release).
|
||||
pub fn set_chunk_embedding_for_signature(
|
||||
config: &Config,
|
||||
chunk_id: &str,
|
||||
model_signature: &str,
|
||||
embedding: &[f32],
|
||||
) -> Result<()> {
|
||||
with_connection(config, |conn| {
|
||||
upsert_chunk_embedding_conn(conn, chunk_id, model_signature, embedding)
|
||||
})
|
||||
}
|
||||
|
||||
/// `true` when at least one chunk or summary still needs an embedding at
|
||||
/// `model_signature` and is not tombstoned as terminally unembeddable.
|
||||
///
|
||||
/// Shared by `ensure_reembed_backfill`, the §7 migration enqueue probe, and
|
||||
/// tests so the worklist and coverage probes cannot drift (#2358).
|
||||
pub(crate) fn has_uncovered_reembed_work(
|
||||
conn: &Connection,
|
||||
model_signature: &str,
|
||||
) -> rusqlite::Result<bool> {
|
||||
conn.query_row(
|
||||
"SELECT EXISTS(
|
||||
SELECT 1 FROM mem_tree_chunks c
|
||||
WHERE NOT EXISTS (SELECT 1 FROM mem_tree_chunk_embeddings e
|
||||
WHERE e.chunk_id = c.id AND e.model_signature = ?1)
|
||||
AND NOT EXISTS (SELECT 1 FROM mem_tree_chunk_reembed_skipped sk
|
||||
WHERE sk.chunk_id = c.id AND sk.model_signature = ?1))
|
||||
OR EXISTS(
|
||||
SELECT 1 FROM mem_tree_summaries s
|
||||
WHERE s.deleted = 0
|
||||
AND NOT EXISTS (SELECT 1 FROM mem_tree_summary_embeddings e
|
||||
WHERE e.summary_id = s.id AND e.model_signature = ?1)
|
||||
AND NOT EXISTS (SELECT 1 FROM mem_tree_summary_reembed_skipped sk
|
||||
WHERE sk.summary_id = s.id AND sk.model_signature = ?1))",
|
||||
rusqlite::params![model_signature],
|
||||
|r| r.get(0),
|
||||
)
|
||||
}
|
||||
|
||||
/// Persistently record that `(chunk_id, signature)` cannot be re-embedded.
|
||||
///
|
||||
/// Called by `handle_reembed_backfill` when the per-chunk body file is
|
||||
/// missing on disk (orphan) or the embedder rejects the row terminally
|
||||
/// (wrong dim / unrecoverable embed error). Inserting a row here causes
|
||||
/// the next backfill batch's worklist query to exclude this chunk via the
|
||||
/// `NOT EXISTS … mem_tree_chunk_reembed_skipped …` predicate, so the
|
||||
/// runaway "skipping" loop terminates instead of revisiting the same row
|
||||
/// every 5 s forever (#1574 §6 fix).
|
||||
pub fn mark_chunk_reembed_skipped(
|
||||
config: &Config,
|
||||
chunk_id: &str,
|
||||
model_signature: &str,
|
||||
reason: &str,
|
||||
) -> Result<()> {
|
||||
let chunk_id = validate_reembed_skip_key("chunk_id", chunk_id)?;
|
||||
let model_signature = validate_reembed_skip_key("model_signature", model_signature)?;
|
||||
with_connection(config, |conn| {
|
||||
let now_ms = Utc::now().timestamp_millis();
|
||||
conn.execute(
|
||||
"INSERT INTO mem_tree_chunk_reembed_skipped
|
||||
(chunk_id, model_signature, reason, skipped_at_ms)
|
||||
VALUES (?1, ?2, ?3, ?4)
|
||||
ON CONFLICT(chunk_id, model_signature) DO UPDATE SET
|
||||
reason = excluded.reason,
|
||||
skipped_at_ms = excluded.skipped_at_ms",
|
||||
rusqlite::params![chunk_id, model_signature, reason, now_ms],
|
||||
)?;
|
||||
log::debug!(
|
||||
"[memory::chunk_store] mark_chunk_reembed_skipped chunk_id={chunk_id} sig={model_signature} reason={reason}"
|
||||
);
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
/// Remove a single chunk tombstone so re-embed backfill can retry the row.
|
||||
///
|
||||
/// Idempotent: deleting a missing `(chunk_id, model_signature)` pair is a
|
||||
/// no-op. Intended for operator recovery after environmental failures (moved
|
||||
/// workspace, restored body files, fixed embedder config) — see #2358.
|
||||
pub fn clear_chunk_reembed_skipped(
|
||||
config: &Config,
|
||||
chunk_id: &str,
|
||||
model_signature: &str,
|
||||
) -> Result<()> {
|
||||
let chunk_id = validate_reembed_skip_key("chunk_id", chunk_id)?;
|
||||
let model_signature = validate_reembed_skip_key("model_signature", model_signature)?;
|
||||
with_connection(config, |conn| {
|
||||
conn.execute(
|
||||
"DELETE FROM mem_tree_chunk_reembed_skipped
|
||||
WHERE chunk_id = ?1 AND model_signature = ?2",
|
||||
rusqlite::params![chunk_id, model_signature],
|
||||
)?;
|
||||
log::debug!(
|
||||
"[memory::chunk_store] clear_chunk_reembed_skipped chunk_id={chunk_id} sig={model_signature}"
|
||||
);
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
/// Clear all chunk and summary tombstones for a model signature.
|
||||
///
|
||||
/// Returns the total number of rows removed across both tombstone tables.
|
||||
/// Idempotent when no tombstones exist for the signature.
|
||||
pub fn clear_reembed_skipped_for_signature(
|
||||
config: &Config,
|
||||
model_signature: &str,
|
||||
) -> Result<usize> {
|
||||
let model_signature = validate_reembed_skip_key("model_signature", model_signature)?;
|
||||
with_connection(config, |conn| {
|
||||
let chunk_deleted = conn.execute(
|
||||
"DELETE FROM mem_tree_chunk_reembed_skipped WHERE model_signature = ?1",
|
||||
rusqlite::params![model_signature],
|
||||
)?;
|
||||
let summary_deleted = conn.execute(
|
||||
"DELETE FROM mem_tree_summary_reembed_skipped WHERE model_signature = ?1",
|
||||
rusqlite::params![model_signature],
|
||||
)?;
|
||||
log::debug!(
|
||||
"[memory::chunk_store] clear_reembed_skipped_for_signature sig={model_signature} chunk_rows={chunk_deleted} summary_rows={summary_deleted}"
|
||||
);
|
||||
Ok(chunk_deleted + summary_deleted)
|
||||
})
|
||||
}
|
||||
|
||||
/// Bounds attacker-controlled ids/signatures passed to reembed-skipped admin
|
||||
/// helpers without affecting legitimate rows (typical ids are well under 512
|
||||
/// chars). Rejects NUL bytes so SQLite bindings cannot be truncated.
|
||||
pub(crate) const REEMBED_SKIP_KEY_MAX_LEN: usize = 2048;
|
||||
|
||||
pub(crate) fn validate_reembed_skip_key<'a>(label: &str, value: &'a str) -> Result<&'a str> {
|
||||
let trimmed = value.trim();
|
||||
if trimmed.is_empty() {
|
||||
anyhow::bail!("{label} must be non-empty");
|
||||
}
|
||||
if trimmed.len() > REEMBED_SKIP_KEY_MAX_LEN {
|
||||
anyhow::bail!("{label} exceeds maximum length ({REEMBED_SKIP_KEY_MAX_LEN})");
|
||||
}
|
||||
if trimmed.as_bytes().contains(&0) {
|
||||
anyhow::bail!("{label} must not contain NUL bytes");
|
||||
}
|
||||
Ok(trimmed)
|
||||
}
|
||||
|
||||
/// Transaction-scoped variant of [`set_chunk_embedding_for_signature`].
|
||||
///
|
||||
/// For callers that already hold a `Transaction` (e.g. the chunk-admission
|
||||
/// handler, which commits the sidecar row in the SAME tx as the lifecycle
|
||||
/// + score + job-enqueue writes — #1574 write-side cutover). Opening a fresh
|
||||
/// connection there would break atomicity / deadlock on the busy DB.
|
||||
pub(crate) fn set_chunk_embedding_for_signature_tx(
|
||||
tx: &rusqlite::Transaction<'_>,
|
||||
chunk_id: &str,
|
||||
model_signature: &str,
|
||||
embedding: &[f32],
|
||||
) -> Result<()> {
|
||||
upsert_chunk_embedding_conn(tx, chunk_id, model_signature, embedding)
|
||||
}
|
||||
|
||||
/// Fetch a chunk embedding for exactly one provider/model/dimension signature.
|
||||
pub fn get_chunk_embedding_for_signature(
|
||||
config: &Config,
|
||||
chunk_id: &str,
|
||||
model_signature: &str,
|
||||
) -> Result<Option<Vec<f32>>> {
|
||||
with_connection(config, |conn| {
|
||||
let row: Option<(Vec<u8>, i64)> = conn
|
||||
.query_row(
|
||||
"SELECT vector, dim
|
||||
FROM mem_tree_chunk_embeddings
|
||||
WHERE chunk_id = ?1 AND model_signature = ?2",
|
||||
rusqlite::params![chunk_id, model_signature],
|
||||
|r| Ok((r.get(0)?, r.get(1)?)),
|
||||
)
|
||||
.optional()?;
|
||||
match row {
|
||||
None => Ok(None),
|
||||
Some((bytes, dim)) => embedding_from_blob(&bytes, dim, "chunk embedding"),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Fetch a chunk's embedding for the active model signature.
|
||||
///
|
||||
/// #1574 cutover: reads the per-model `mem_tree_chunk_embeddings` sidecar at
|
||||
/// the active signature (via [`get_chunk_embedding_for_signature`]) instead of
|
||||
/// the legacy `mem_tree_chunks.embedding` column. Returns `Ok(None)` if the
|
||||
/// chunk has no vector under the active signature — e.g. during the §7
|
||||
/// backfill window, where this degrades retrieval gracefully (the row is
|
||||
/// simply absent from vector results, never cross-space compared).
|
||||
pub fn get_chunk_embedding(config: &Config, chunk_id: &str) -> Result<Option<Vec<f32>>> {
|
||||
let signature = tree_active_signature(config);
|
||||
get_chunk_embedding_for_signature(config, chunk_id, &signature)
|
||||
}
|
||||
|
||||
pub(crate) fn embedding_to_blob(embedding: &[f32]) -> Vec<u8> {
|
||||
embedding.iter().flat_map(|f| f.to_le_bytes()).collect()
|
||||
}
|
||||
|
||||
fn embedding_from_blob(bytes: &[u8], dim: i64, label: &str) -> Result<Option<Vec<f32>>> {
|
||||
if dim < 0 {
|
||||
anyhow::bail!("{label} has negative dimension {dim}");
|
||||
}
|
||||
if !bytes.len().is_multiple_of(4) {
|
||||
anyhow::bail!("{label} blob length {} not a multiple of 4", bytes.len());
|
||||
}
|
||||
let floats: Vec<f32> = bytes
|
||||
.chunks_exact(4)
|
||||
.map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
|
||||
.collect();
|
||||
if floats.len() != dim as usize {
|
||||
anyhow::bail!(
|
||||
"{label} dimension mismatch: dim column says {dim}, blob contains {} floats",
|
||||
floats.len()
|
||||
);
|
||||
}
|
||||
Ok(Some(floats))
|
||||
}
|
||||
@@ -24,15 +24,11 @@
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use chrono::{DateTime, TimeZone, Utc};
|
||||
use parking_lot::Mutex as PMutex;
|
||||
use rusqlite::{params, Connection, OptionalExtension, Transaction};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
|
||||
use std::collections::HashSet;
|
||||
#[cfg(test)]
|
||||
use std::sync::Mutex;
|
||||
use std::sync::{Arc, OnceLock};
|
||||
use std::time::{Duration, Instant};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::openhuman::config::Config;
|
||||
use crate::openhuman::memory::util::redact::{self, redact as redact_value};
|
||||
@@ -1024,529 +1020,18 @@ fn ms_to_utc(ms: i64) -> rusqlite::Result<DateTime<Utc>> {
|
||||
})
|
||||
}
|
||||
|
||||
// ── Schema-apply instrumentation (test-only) ─────────────────────────────────
|
||||
//
|
||||
// Per-path counter of how many times `apply_schema` ran for each DB path,
|
||||
// gated behind `cfg(test)` so the production binary carries no overhead.
|
||||
// Used by the concurrent-init regression test to assert "exactly once per
|
||||
// path" across racing workers; it survives even when the connection cache
|
||||
// is cleared between tests because tests use distinct tempdirs.
|
||||
#[path = "connection.rs"]
|
||||
mod connection;
|
||||
pub use connection::with_connection;
|
||||
#[cfg(test)]
|
||||
static SCHEMA_APPLY_COUNTS: OnceLock<Mutex<HashMap<PathBuf, usize>>> = OnceLock::new();
|
||||
|
||||
fn record_schema_apply(_path: &Path) {
|
||||
#[cfg(test)]
|
||||
{
|
||||
let counts = SCHEMA_APPLY_COUNTS.get_or_init(|| Mutex::new(HashMap::new()));
|
||||
let mut guard = counts
|
||||
.lock()
|
||||
.expect("memory_tree schema apply count mutex poisoned");
|
||||
*guard.entry(_path.to_path_buf()).or_insert(0) += 1;
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(unused_imports)]
|
||||
pub(crate) use connection::{
|
||||
clear_connection_cache, db_path_for, get_or_init_connection, invalidate_connection,
|
||||
is_io_open_error, schema_apply_count_for_path_for_tests, CB_THRESHOLD,
|
||||
};
|
||||
#[cfg(test)]
|
||||
#[doc(hidden)]
|
||||
pub(crate) fn schema_apply_count_for_path_for_tests(path: &Path) -> usize {
|
||||
SCHEMA_APPLY_COUNTS
|
||||
.get()
|
||||
.and_then(|m| {
|
||||
m.lock()
|
||||
.ok()
|
||||
.map(|guard| guard.get(path).copied().unwrap_or(0))
|
||||
})
|
||||
.unwrap_or(0)
|
||||
}
|
||||
pub(crate) use connection::{is_transient_cold_start, try_cleanup_stale_files};
|
||||
|
||||
// SQLite extended result codes that fire during cold-start WAL/SHM bootstrap
|
||||
// races. NOTE on values: extended codes are `SQLITE_IOERR (10) | (sub << 8)`.
|
||||
// 4874 is `IOERR_SHMSIZE` (sub 19), NOT `SHMMAP` — the real `SHMMAP` is 5386
|
||||
// (sub 21) and the "open a new shared-memory segment" failure is `SHMOPEN`
|
||||
// 4618 (sub 18), which is what surfaced on macOS. The whole `-shm` family is
|
||||
// listed so the classifiers don't miss any of them.
|
||||
/// `CANTOPEN` — racing the lockfile/WAL creation done by another connection.
|
||||
const SQLITE_CANTOPEN: i32 = 14;
|
||||
/// `IOERR_TRUNCATE` — the WAL/db is being truncated during bootstrap.
|
||||
const SQLITE_IOERR_TRUNCATE: i32 = 1546;
|
||||
/// `IOERR_SHMOPEN` — opening a new `-shm` shared-memory segment failed (the
|
||||
/// macOS cold-start failure, e.g. Sentry TAURI-RUST-X1).
|
||||
const SQLITE_IOERR_SHMOPEN: i32 = 4618;
|
||||
/// `IOERR_SHMSIZE` — the `-shm` file is being resized during bootstrap.
|
||||
const SQLITE_IOERR_SHMSIZE: i32 = 4874;
|
||||
/// `IOERR_SHMMAP` — mapping a page of the `-shm` wal-index failed.
|
||||
const SQLITE_IOERR_SHMMAP: i32 = 5386;
|
||||
/// `IOERR_IN_PAGE` — an mmap-page I/O fault, also seen under WAL cold-start.
|
||||
const SQLITE_IOERR_IN_PAGE: i32 = 8714;
|
||||
|
||||
/// True if `err` (or anything in its cause chain) is one of the SQLite codes
|
||||
/// that fire during cold-start WAL/SHM bootstrap races: `CANTOPEN`,
|
||||
/// `IOERR_TRUNCATE`, the `-shm` family (`SHMOPEN` / `SHMSIZE` / `SHMMAP`), and
|
||||
/// `IOERR_IN_PAGE`.
|
||||
pub(crate) fn is_transient_cold_start(err: &anyhow::Error) -> bool {
|
||||
fn is_transient_sqlite(e: &(dyn std::error::Error + 'static)) -> bool {
|
||||
if let Some(rusqlite::Error::SqliteFailure(ffi, _)) = e.downcast_ref::<rusqlite::Error>() {
|
||||
return matches!(
|
||||
ffi.extended_code,
|
||||
SQLITE_CANTOPEN
|
||||
| SQLITE_IOERR_TRUNCATE
|
||||
| SQLITE_IOERR_SHMOPEN
|
||||
| SQLITE_IOERR_SHMSIZE
|
||||
| SQLITE_IOERR_SHMMAP
|
||||
| SQLITE_IOERR_IN_PAGE
|
||||
);
|
||||
}
|
||||
false
|
||||
}
|
||||
if is_transient_sqlite(err.root_cause()) {
|
||||
return true;
|
||||
}
|
||||
let mut src: Option<&(dyn std::error::Error + 'static)> = Some(err.as_ref());
|
||||
while let Some(cur) = src {
|
||||
if is_transient_sqlite(cur) {
|
||||
return true;
|
||||
}
|
||||
src = cur.source();
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
// ── Connection cache (#2206) ─────────────────────────────────────────────────
|
||||
|
||||
/// How many consecutive init failures before the circuit breaker trips.
|
||||
const CB_THRESHOLD: u32 = 3;
|
||||
/// How long the circuit breaker holds the DB closed after tripping.
|
||||
const CB_COOLDOWN: Duration = Duration::from_secs(30);
|
||||
|
||||
/// Per-path circuit breaker: after [`CB_THRESHOLD`] consecutive init failures
|
||||
/// the breaker trips and `get_or_init_connection` returns an error immediately
|
||||
/// until [`CB_COOLDOWN`] elapses. On the first success it resets to zero.
|
||||
struct CircuitBreaker {
|
||||
consecutive_failures: AtomicU32,
|
||||
tripped: AtomicBool,
|
||||
last_trip: PMutex<Option<Instant>>,
|
||||
}
|
||||
|
||||
impl CircuitBreaker {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
consecutive_failures: AtomicU32::new(0),
|
||||
tripped: AtomicBool::new(false),
|
||||
last_trip: PMutex::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
fn record_success(&self) {
|
||||
self.consecutive_failures.store(0, Ordering::Relaxed);
|
||||
self.tripped.store(false, Ordering::Relaxed);
|
||||
*self.last_trip.lock() = None;
|
||||
}
|
||||
|
||||
/// Records one more failure. Returns `true` if this call just tripped the
|
||||
/// breaker (i.e. the threshold was crossed right now).
|
||||
fn record_failure(&self) -> bool {
|
||||
let prev = self.consecutive_failures.fetch_add(1, Ordering::Relaxed);
|
||||
let count = prev + 1;
|
||||
if count >= CB_THRESHOLD && !self.tripped.swap(true, Ordering::Relaxed) {
|
||||
*self.last_trip.lock() = Some(Instant::now());
|
||||
return true;
|
||||
}
|
||||
// Re-arm the cooldown on each subsequent failure while already tripped
|
||||
// so that a failed post-cooldown retry restarts the 30 s window instead
|
||||
// of leaving the stale timestamp in place (which would let `is_open`
|
||||
// return false immediately and allow unlimited retries).
|
||||
if self.tripped.load(Ordering::Relaxed) {
|
||||
*self.last_trip.lock() = Some(Instant::now());
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// Returns `true` when the breaker is open AND the cooldown has not yet
|
||||
/// elapsed. Returns `false` (allowing a retry) once the cooldown passes.
|
||||
fn is_open(&self) -> bool {
|
||||
if !self.tripped.load(Ordering::Relaxed) {
|
||||
return false;
|
||||
}
|
||||
let guard = self.last_trip.lock();
|
||||
match *guard {
|
||||
Some(t) if t.elapsed() < CB_COOLDOWN => true,
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Process-level cache — two separate maps so that a failing init cannot
|
||||
/// accidentally serve a dummy connection on the fast path.
|
||||
///
|
||||
/// `connections`: only fully-initialised (schema + migrations run) entries.
|
||||
/// `breakers`: persists across failed init attempts so the circuit breaker
|
||||
/// survives even when `connections` has no entry for that path.
|
||||
struct ConnectionCache {
|
||||
connections: PMutex<HashMap<PathBuf, Arc<PMutex<Connection>>>>,
|
||||
breakers: PMutex<HashMap<PathBuf, Arc<CircuitBreaker>>>,
|
||||
/// Per-path mutex held across the slow-path init so concurrent
|
||||
/// workers racing into `with_connection` on a cold DB serialise on
|
||||
/// the WAL+SHM bootstrap. Without this, N threads see "no cached
|
||||
/// connection" simultaneously and all run `apply_schema`, which is
|
||||
/// idempotent but reopens the cold-start race window
|
||||
/// (OPENHUMAN-TAURI-HH / -ZM / -MB).
|
||||
init_locks: PMutex<HashMap<PathBuf, Arc<PMutex<()>>>>,
|
||||
}
|
||||
|
||||
static CONN_CACHE: OnceLock<ConnectionCache> = OnceLock::new();
|
||||
|
||||
fn conn_cache() -> &'static ConnectionCache {
|
||||
CONN_CACHE.get_or_init(|| ConnectionCache {
|
||||
connections: PMutex::new(HashMap::new()),
|
||||
breakers: PMutex::new(HashMap::new()),
|
||||
init_locks: PMutex::new(HashMap::new()),
|
||||
})
|
||||
}
|
||||
|
||||
/// Compute the canonical DB path from `config`.
|
||||
fn db_path_for(config: &Config) -> PathBuf {
|
||||
config.workspace_dir.join(DB_DIR).join(DB_FILE)
|
||||
}
|
||||
|
||||
/// Delete stale WAL/SHM side-files (`<db>-shm`, `<db>-wal`) that can block a
|
||||
/// clean DB open after a crash. Logs what was deleted and returns `true` if
|
||||
/// anything was removed.
|
||||
///
|
||||
/// SQLite names these files `<db_path>-shm` and `<db_path>-wal`.
|
||||
/// For `chunks.db` that is `chunks.db-shm` / `chunks.db-wal`.
|
||||
pub(crate) fn try_cleanup_stale_files(db_path: &std::path::Path) -> bool {
|
||||
let mut cleaned = false;
|
||||
for suffix in &["-shm", "-wal"] {
|
||||
// Build the side-file path: append suffix to the full db filename.
|
||||
let side = {
|
||||
let mut p = db_path.to_path_buf();
|
||||
let name = p
|
||||
.file_name()
|
||||
.unwrap_or_default()
|
||||
.to_string_lossy()
|
||||
.into_owned();
|
||||
p.set_file_name(format!("{name}{suffix}"));
|
||||
p
|
||||
};
|
||||
if side.exists() {
|
||||
match std::fs::remove_file(&side) {
|
||||
Ok(()) => {
|
||||
log::warn!("[memory_tree] removed stale side-file: {}", side.display());
|
||||
cleaned = true;
|
||||
}
|
||||
Err(e) => {
|
||||
log::warn!(
|
||||
"[memory_tree] failed to remove stale side-file {}: {e}",
|
||||
side.display()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
cleaned
|
||||
}
|
||||
|
||||
/// Run the full one-time DB initialisation (journal mode, schema, migrations)
|
||||
/// against an already-open `Connection`. Used by `get_or_init_connection`.
|
||||
fn init_db(conn: &Connection, config: &Config) -> Result<()> {
|
||||
conn.busy_timeout(SQLITE_BUSY_TIMEOUT)
|
||||
.context("Failed to configure memory_tree busy timeout")?;
|
||||
// SQLite resets `foreign_keys` to off on every new connection. The
|
||||
// ConnectionCache holds one cached `Connection` per DB path, so
|
||||
// setting it here (alongside the rest of init) is the per-connection
|
||||
// surface — fast-path callers reuse the cached conn with FKs already
|
||||
// on.
|
||||
conn.execute_batch("PRAGMA foreign_keys = ON;")
|
||||
.context("Failed to enable memory_tree foreign_keys pragma")?;
|
||||
// memory_tree runs the TRUNCATE rollback journal (see `apply_schema`), so
|
||||
// crash-safety requires synchronous=FULL — NORMAL is only corruption-safe
|
||||
// under WAL. Set explicitly so a future global default can't weaken it.
|
||||
conn.execute_batch("PRAGMA synchronous = FULL;")
|
||||
.context("Failed to set memory_tree synchronous=FULL")?;
|
||||
apply_schema(conn)?;
|
||||
// #1574 §7: one-shot, version-gated legacy→sidecar embedding migration.
|
||||
migrate_legacy_embeddings_to_sidecar(conn, config)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn apply_schema(conn: &Connection) -> Result<()> {
|
||||
// Note: `init_db` runs the `#1574 §7` legacy→sidecar embedding migration
|
||||
// after this returns, so the dim-equal copy step is not duplicated here.
|
||||
// memory_tree uses the TRUNCATE rollback journal, NOT WAL. WAL's `-shm`
|
||||
// shared-memory index + `-wal` checkpoint machinery are the root of the
|
||||
// cold-start IOERR_SHMMAP (macOS) / IOERR_TRUNCATE (Windows, AV-held
|
||||
// handles) failures (Sentry TAURI-RUST-EV / TAURI-RUST-X1). All tree
|
||||
// access serialises on the single cached `PMutex<Connection>` (see
|
||||
// `get_or_init_connection`), so WAL's only real benefit — concurrent
|
||||
// readers — is unused here, which makes WAL pure liability. The sibling
|
||||
// tree DBs (cron / vault / redirect_links) already run the default
|
||||
// rollback journal without issue.
|
||||
//
|
||||
// Requesting TRUNCATE on a database a prior release left in WAL mode
|
||||
// checkpoints the `-wal` back into the main file and removes the
|
||||
// `-wal`/`-shm` side-files, so this also migrates existing WAL databases
|
||||
// in place on upgrade.
|
||||
let journal_mode: String = conn
|
||||
.query_row("PRAGMA journal_mode=TRUNCATE", [], |row| row.get(0))
|
||||
.context("Failed to set memory_tree journal_mode=TRUNCATE")?;
|
||||
if !journal_mode.eq_ignore_ascii_case("truncate") {
|
||||
log::warn!(
|
||||
"[memory_tree] journal_mode is '{journal_mode}' after requesting TRUNCATE \
|
||||
— a prior WAL connection or a locked -wal may be blocking the switch"
|
||||
);
|
||||
}
|
||||
conn.execute_batch(SCHEMA)
|
||||
.context("Failed to initialize memory_tree schema")?;
|
||||
// Phase 2 migrations — additive, idempotent.
|
||||
add_column_if_missing(conn, "mem_tree_chunks", "embedding", "BLOB")?;
|
||||
// Phase 2 LLM-NER follow-up: per-chunk LLM importance signal +
|
||||
// human-readable reason. Both nullable; absence is treated as
|
||||
// "no LLM signal available" by readers.
|
||||
add_column_if_missing(conn, "mem_tree_score", "llm_importance", "REAL")?;
|
||||
add_column_if_missing(conn, "mem_tree_score", "llm_importance_reason", "TEXT")?;
|
||||
// Phase 3a (#709): parent-summary backlink on leaves.
|
||||
add_column_if_missing(conn, "mem_tree_chunks", "parent_summary_id", "TEXT")?;
|
||||
// Phase 4 (#710): sealed-summary embeddings for semantic rerank.
|
||||
add_column_if_missing(conn, "mem_tree_summaries", "embedding", "BLOB")?;
|
||||
// Async-pipeline lifecycle flag. Default 'admitted' so chunks ingested
|
||||
// before the queue migration stay queryable.
|
||||
add_column_if_missing(
|
||||
conn,
|
||||
"mem_tree_chunks",
|
||||
"lifecycle_status",
|
||||
"TEXT NOT NULL DEFAULT 'admitted'",
|
||||
)?;
|
||||
conn.execute_batch(
|
||||
"CREATE INDEX IF NOT EXISTS idx_mem_tree_chunks_lifecycle \
|
||||
ON mem_tree_chunks(lifecycle_status);",
|
||||
)
|
||||
.context("Failed to create mem_tree_chunks lifecycle index")?;
|
||||
// Phase MD-content (#TBD): pointer + integrity hash.
|
||||
add_column_if_missing(conn, "mem_tree_chunks", "content_path", "TEXT")?;
|
||||
add_column_if_missing(conn, "mem_tree_chunks", "content_sha256", "TEXT")?;
|
||||
// Phase MD-content (summaries).
|
||||
add_column_if_missing(conn, "mem_tree_summaries", "content_path", "TEXT")?;
|
||||
add_column_if_missing(conn, "mem_tree_summaries", "content_sha256", "TEXT")?;
|
||||
// Raw-archive pointer column.
|
||||
add_column_if_missing(conn, "mem_tree_chunks", "raw_refs_json", "TEXT")?;
|
||||
// #1365: is_user flag on indexed entity rows.
|
||||
add_column_if_missing(
|
||||
conn,
|
||||
"mem_tree_entity_index",
|
||||
"is_user",
|
||||
"INTEGER NOT NULL DEFAULT 0",
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Whether `err` looks like one of the I/O error codes that warrant a
|
||||
/// stale-file cleanup + single retry before giving up.
|
||||
fn is_io_open_error(err: &anyhow::Error) -> bool {
|
||||
if let Some(rusqlite::Error::SqliteFailure(f, _)) = err.downcast_ref::<rusqlite::Error>() {
|
||||
return matches!(
|
||||
f.extended_code,
|
||||
SQLITE_CANTOPEN
|
||||
| SQLITE_IOERR_TRUNCATE
|
||||
| SQLITE_IOERR_SHMOPEN
|
||||
| SQLITE_IOERR_SHMSIZE
|
||||
| SQLITE_IOERR_SHMMAP
|
||||
| SQLITE_IOERR_IN_PAGE
|
||||
) || f.code == rusqlite::ErrorCode::CannotOpen;
|
||||
}
|
||||
let msg = format!("{err:#}").to_ascii_lowercase();
|
||||
msg.contains("disk i/o error")
|
||||
|| msg.contains("unable to open database file")
|
||||
|| msg.contains("xshmmap")
|
||||
|| msg.contains("truncate file")
|
||||
}
|
||||
|
||||
/// Obtain (or lazily create) a cached connection for the workspace described
|
||||
/// by `config`. Returns `Err` immediately when the circuit breaker is open.
|
||||
fn get_or_init_connection(config: &Config) -> Result<Arc<PMutex<Connection>>> {
|
||||
let db_path = db_path_for(config);
|
||||
|
||||
// ── Fast path: reject immediately if the breaker is open ─────────────
|
||||
{
|
||||
let breakers = conn_cache().breakers.lock();
|
||||
if let Some(breaker) = breakers.get(&db_path) {
|
||||
if breaker.is_open() {
|
||||
anyhow::bail!(
|
||||
"[memory_tree] circuit breaker open for {}: too many consecutive init failures",
|
||||
db_path.display()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
// ── Fast path: return cached connection if already initialised ────────
|
||||
{
|
||||
let guard = conn_cache().connections.lock();
|
||||
if let Some(conn) = guard.get(&db_path) {
|
||||
return Ok(Arc::clone(conn));
|
||||
}
|
||||
}
|
||||
|
||||
// ── Slow path: serialise init per-path so concurrent workers don't
|
||||
// all race into `open_and_init` on a cold DB.
|
||||
let init_lock = {
|
||||
let mut guard = conn_cache().init_locks.lock();
|
||||
guard
|
||||
.entry(db_path.clone())
|
||||
.or_insert_with(|| Arc::new(PMutex::new(())))
|
||||
.clone()
|
||||
};
|
||||
let _init_guard = init_lock.lock();
|
||||
|
||||
// Re-check the cache once we hold the init lock — another thread
|
||||
// may have completed init while we were queued.
|
||||
{
|
||||
let guard = conn_cache().connections.lock();
|
||||
if let Some(conn) = guard.get(&db_path) {
|
||||
return Ok(Arc::clone(conn));
|
||||
}
|
||||
}
|
||||
|
||||
log::debug!(
|
||||
"[memory_tree] opening and initialising DB at {}",
|
||||
db_path.display()
|
||||
);
|
||||
|
||||
// Attempt to open + init the connection (dir creation is inside
|
||||
// `open_and_init` so every failure — including EEXIST on the dir —
|
||||
// reaches the circuit-breaker recording logic below). On certain I/O
|
||||
// errors (#2206) we clean up stale WAL/SHM side-files and retry once.
|
||||
let conn = open_and_init(&db_path, config).or_else(|first_err| {
|
||||
if is_io_open_error(&first_err) {
|
||||
log::warn!(
|
||||
"[memory_tree] I/O error on first open attempt ({}), cleaning stale files and retrying",
|
||||
first_err
|
||||
);
|
||||
try_cleanup_stale_files(&db_path);
|
||||
open_and_init(&db_path, config)
|
||||
} else {
|
||||
Err(first_err)
|
||||
}
|
||||
});
|
||||
|
||||
match conn {
|
||||
Ok(conn) => {
|
||||
let arc_conn = Arc::new(PMutex::new(conn));
|
||||
conn_cache()
|
||||
.connections
|
||||
.lock()
|
||||
.insert(db_path.clone(), Arc::clone(&arc_conn));
|
||||
// Reset any prior failure counter now that init succeeded.
|
||||
if let Some(breaker) = conn_cache().breakers.lock().get(&db_path) {
|
||||
breaker.record_success();
|
||||
}
|
||||
log::debug!("[memory_tree] DB connection cached and ready");
|
||||
Ok(arc_conn)
|
||||
}
|
||||
Err(err) => {
|
||||
// Persist the breaker so the failure count accumulates across
|
||||
// calls even though no connection entry exists yet.
|
||||
let breaker = {
|
||||
let mut guard = conn_cache().breakers.lock();
|
||||
guard
|
||||
.entry(db_path.clone())
|
||||
.or_insert_with(|| Arc::new(CircuitBreaker::new()))
|
||||
.clone()
|
||||
};
|
||||
let just_tripped = breaker.record_failure();
|
||||
if just_tripped {
|
||||
log::error!(
|
||||
"[memory_tree] circuit breaker tripped for {}: {} consecutive init failures",
|
||||
db_path.display(),
|
||||
CB_THRESHOLD
|
||||
);
|
||||
let _ = crate::core::event_bus::publish_global(
|
||||
crate::core::event_bus::DomainEvent::HealthChanged {
|
||||
component: "memory_tree_db".to_string(),
|
||||
healthy: false,
|
||||
message: Some(format!(
|
||||
"Schema init failed {CB_THRESHOLD} consecutive times"
|
||||
)),
|
||||
},
|
||||
);
|
||||
}
|
||||
Err(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Ensure the DB directory exists, open the SQLite file, and run the full
|
||||
/// schema init sequence. All errors (dir creation, file open, schema init)
|
||||
/// are returned as `Err` so callers can funnel them through the circuit
|
||||
/// breaker logic in a single place.
|
||||
fn open_and_init(db_path: &std::path::Path, config: &Config) -> Result<Connection> {
|
||||
let dir = db_path.parent().expect("db_path always has a parent");
|
||||
std::fs::create_dir_all(dir)
|
||||
.with_context(|| format!("Failed to create memory_tree dir: {}", dir.display()))?;
|
||||
let conn = Connection::open(db_path)
|
||||
.with_context(|| format!("Failed to open memory_tree DB: {}", db_path.display()))?;
|
||||
init_db(&conn, config)
|
||||
.with_context(|| format!("Failed to init memory_tree schema: {}", db_path.display()))?;
|
||||
record_schema_apply(db_path);
|
||||
Ok(conn)
|
||||
}
|
||||
|
||||
/// Remove the cached connection for `config`'s workspace (forces a fresh open
|
||||
/// on the next `with_connection` call). Also clears the breaker so the next
|
||||
/// open attempt is not immediately rejected. Does nothing if no entry exists.
|
||||
#[allow(dead_code)]
|
||||
pub(crate) fn invalidate_connection(config: &Config) {
|
||||
let db_path = db_path_for(config);
|
||||
conn_cache().connections.lock().remove(&db_path);
|
||||
conn_cache().breakers.lock().remove(&db_path);
|
||||
log::debug!(
|
||||
"[memory_tree] connection invalidated for {}",
|
||||
db_path.display()
|
||||
);
|
||||
}
|
||||
|
||||
/// Clear the entire connection cache. For test isolation only.
|
||||
#[cfg(test)]
|
||||
pub(crate) fn clear_connection_cache() {
|
||||
conn_cache().connections.lock().clear();
|
||||
conn_cache().breakers.lock().clear();
|
||||
conn_cache().init_locks.lock().clear();
|
||||
}
|
||||
|
||||
/// Open the memory_tree SQLite DB and run a closure against it.
|
||||
///
|
||||
/// Visible to sibling modules (e.g. `score::store`) so Phase 2 can reuse
|
||||
/// the same connection setup / schema initialisation without duplication.
|
||||
///
|
||||
/// # Connection caching (#2206)
|
||||
///
|
||||
/// The underlying connection is initialised once per workspace path and then
|
||||
/// reused from a process-level cache. Schema migrations run exactly once on
|
||||
/// the first call for a given `config.workspace_dir`. Subsequent calls pay
|
||||
/// only the cost of a `parking_lot::Mutex` lock and the closure itself.
|
||||
///
|
||||
/// `#[doc(hidden)] pub` (not `pub(crate)`) because the
|
||||
/// `memory-tree-init-smoke` bin in `src/bin/` is a separate crate target
|
||||
/// and must reach this entry point. It is NOT a stable API surface —
|
||||
/// downstream crates should treat it as internal.
|
||||
#[doc(hidden)]
|
||||
pub fn with_connection<T>(config: &Config, f: impl FnOnce(&Connection) -> Result<T>) -> Result<T> {
|
||||
let conn_arc = get_or_init_connection(config)?;
|
||||
let guard = conn_arc.lock();
|
||||
f(&guard)
|
||||
}
|
||||
|
||||
/// One-shot migration (#1574 §7, vN): copy legacy `mem_tree_chunks.embedding`
|
||||
/// / `mem_tree_summaries.embedding` blobs into the per-model sidecar tables
|
||||
/// under the **active** signature, when (and only when) the legacy vector's
|
||||
/// dimensionality matches the active embedder's.
|
||||
///
|
||||
/// Version-gated via `PRAGMA user_version`: returns immediately once
|
||||
/// `>= TREE_EMBEDDING_MIGRATION_VERSION`, so the per-open cost is a single
|
||||
/// pragma read. Dim-mismatched rows are left for the §6 re-embed backfill —
|
||||
/// the blob's signature is unrecoverable (see spec §7b), so a same-length
|
||||
/// copy under the active signature is the only provably-safe move and
|
||||
/// anything else must be re-embedded. The legacy columns are **kept** (read
|
||||
/// here, dropped only in a later release — spec §7c). Idempotent: re-running
|
||||
/// before the version bumps re-copies the same rows harmlessly (sidecar
|
||||
/// upsert is ON CONFLICT); after the bump it is skipped entirely.
|
||||
fn migrate_legacy_embeddings_to_sidecar(conn: &Connection, config: &Config) -> Result<()> {
|
||||
let version: i64 = conn
|
||||
.query_row("PRAGMA user_version", [], |r| r.get(0))
|
||||
@@ -1802,289 +1287,19 @@ fn add_column_if_missing(conn: &Connection, table: &str, name: &str, sql_type: &
|
||||
}
|
||||
}
|
||||
|
||||
// ── Phase 2: embedding column accessors ─────────────────────────────────
|
||||
|
||||
/// Resolve the active embedding signature for the memory tree from the global
|
||||
/// [`Config`] — the canonical key every per-model sidecar read/write is scoped
|
||||
/// by (#1574). Reuses the established local-AI workload derivation
|
||||
/// ([`Config::workload_local_model`]) and the probe-stable
|
||||
/// `active_embedding_signature`; introduces no parallel resolution path.
|
||||
/// `pub(crate)` so the sibling `tree` summary store shares the exact
|
||||
/// same resolution.
|
||||
pub(crate) fn tree_active_signature(config: &Config) -> String {
|
||||
let local_model = config.workload_local_model("embeddings");
|
||||
crate::openhuman::memory_store::active_embedding_signature(
|
||||
&config.memory,
|
||||
local_model.as_deref(),
|
||||
)
|
||||
}
|
||||
|
||||
/// Store a chunk's embedding under the active model signature.
|
||||
///
|
||||
/// #1574 cutover: this now writes the per-model `mem_tree_chunk_embeddings`
|
||||
/// sidecar (via [`set_chunk_embedding_for_signature`]) instead of the legacy
|
||||
/// `mem_tree_chunks.embedding` column. Call sites are unchanged — the signature
|
||||
/// is resolved internally from `config`. The legacy column is left intact for
|
||||
/// the §7 one-shot migration to read; it is dropped only in a later release.
|
||||
pub fn set_chunk_embedding(config: &Config, chunk_id: &str, embedding: &[f32]) -> Result<()> {
|
||||
let signature = tree_active_signature(config);
|
||||
log::debug!(
|
||||
"[memory::chunk_store] set_chunk_embedding: chunk_id={chunk_id} sig={signature} dims={}",
|
||||
embedding.len()
|
||||
);
|
||||
set_chunk_embedding_for_signature(config, chunk_id, &signature, embedding)
|
||||
}
|
||||
|
||||
/// Core upsert into `mem_tree_chunk_embeddings` over an arbitrary
|
||||
/// `&Connection`. Shared by the standalone ([`set_chunk_embedding_for_signature`])
|
||||
/// and in-transaction ([`set_chunk_embedding_for_signature_tx`]) write paths so
|
||||
/// the SQL exists exactly once. `rusqlite::Transaction` derefs to `Connection`,
|
||||
/// so an in-tx caller passes `&tx` and the sidecar row commits atomically with
|
||||
/// the surrounding work (#1574 write-side cutover).
|
||||
fn upsert_chunk_embedding_conn(
|
||||
conn: &rusqlite::Connection,
|
||||
chunk_id: &str,
|
||||
model_signature: &str,
|
||||
embedding: &[f32],
|
||||
) -> Result<()> {
|
||||
let bytes = embedding_to_blob(embedding);
|
||||
let dim = i64::try_from(embedding.len()).context("embedding dimension does not fit i64")?;
|
||||
let created_at = Utc::now().timestamp_millis() as f64 / 1000.0;
|
||||
conn.execute(
|
||||
"INSERT INTO mem_tree_chunk_embeddings
|
||||
(chunk_id, model_signature, vector, dim, created_at)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5)
|
||||
ON CONFLICT(chunk_id, model_signature) DO UPDATE SET
|
||||
vector = excluded.vector,
|
||||
dim = excluded.dim,
|
||||
created_at = excluded.created_at",
|
||||
rusqlite::params![chunk_id, model_signature, bytes, dim, created_at],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Store a chunk embedding for a specific provider/model/dimension signature.
|
||||
///
|
||||
/// Per-model table write path for #1574. The legacy
|
||||
/// `mem_tree_chunks.embedding` column is intentionally left untouched by this
|
||||
/// helper (read by the §7 migration; dropped only in a later release).
|
||||
pub fn set_chunk_embedding_for_signature(
|
||||
config: &Config,
|
||||
chunk_id: &str,
|
||||
model_signature: &str,
|
||||
embedding: &[f32],
|
||||
) -> Result<()> {
|
||||
with_connection(config, |conn| {
|
||||
upsert_chunk_embedding_conn(conn, chunk_id, model_signature, embedding)
|
||||
})
|
||||
}
|
||||
|
||||
/// `true` when at least one chunk or summary still needs an embedding at
|
||||
/// `model_signature` and is not tombstoned as terminally unembeddable.
|
||||
///
|
||||
/// Shared by `ensure_reembed_backfill`, the §7 migration enqueue probe, and
|
||||
/// tests so the worklist and coverage probes cannot drift (#2358).
|
||||
pub(crate) fn has_uncovered_reembed_work(
|
||||
conn: &Connection,
|
||||
model_signature: &str,
|
||||
) -> rusqlite::Result<bool> {
|
||||
conn.query_row(
|
||||
"SELECT EXISTS(
|
||||
SELECT 1 FROM mem_tree_chunks c
|
||||
WHERE NOT EXISTS (SELECT 1 FROM mem_tree_chunk_embeddings e
|
||||
WHERE e.chunk_id = c.id AND e.model_signature = ?1)
|
||||
AND NOT EXISTS (SELECT 1 FROM mem_tree_chunk_reembed_skipped sk
|
||||
WHERE sk.chunk_id = c.id AND sk.model_signature = ?1))
|
||||
OR EXISTS(
|
||||
SELECT 1 FROM mem_tree_summaries s
|
||||
WHERE s.deleted = 0
|
||||
AND NOT EXISTS (SELECT 1 FROM mem_tree_summary_embeddings e
|
||||
WHERE e.summary_id = s.id AND e.model_signature = ?1)
|
||||
AND NOT EXISTS (SELECT 1 FROM mem_tree_summary_reembed_skipped sk
|
||||
WHERE sk.summary_id = s.id AND sk.model_signature = ?1))",
|
||||
rusqlite::params![model_signature],
|
||||
|r| r.get(0),
|
||||
)
|
||||
}
|
||||
|
||||
/// Persistently record that `(chunk_id, signature)` cannot be re-embedded.
|
||||
///
|
||||
/// Called by `handle_reembed_backfill` when the per-chunk body file is
|
||||
/// missing on disk (orphan) or the embedder rejects the row terminally
|
||||
/// (wrong dim / unrecoverable embed error). Inserting a row here causes
|
||||
/// the next backfill batch's worklist query to exclude this chunk via the
|
||||
/// `NOT EXISTS … mem_tree_chunk_reembed_skipped …` predicate, so the
|
||||
/// runaway "skipping" loop terminates instead of revisiting the same row
|
||||
/// every 5 s forever (#1574 §6 fix).
|
||||
pub fn mark_chunk_reembed_skipped(
|
||||
config: &Config,
|
||||
chunk_id: &str,
|
||||
model_signature: &str,
|
||||
reason: &str,
|
||||
) -> Result<()> {
|
||||
let chunk_id = validate_reembed_skip_key("chunk_id", chunk_id)?;
|
||||
let model_signature = validate_reembed_skip_key("model_signature", model_signature)?;
|
||||
with_connection(config, |conn| {
|
||||
let now_ms = Utc::now().timestamp_millis();
|
||||
conn.execute(
|
||||
"INSERT INTO mem_tree_chunk_reembed_skipped
|
||||
(chunk_id, model_signature, reason, skipped_at_ms)
|
||||
VALUES (?1, ?2, ?3, ?4)
|
||||
ON CONFLICT(chunk_id, model_signature) DO UPDATE SET
|
||||
reason = excluded.reason,
|
||||
skipped_at_ms = excluded.skipped_at_ms",
|
||||
rusqlite::params![chunk_id, model_signature, reason, now_ms],
|
||||
)?;
|
||||
log::debug!(
|
||||
"[memory::chunk_store] mark_chunk_reembed_skipped chunk_id={chunk_id} sig={model_signature} reason={reason}"
|
||||
);
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
/// Remove a single chunk tombstone so re-embed backfill can retry the row.
|
||||
///
|
||||
/// Idempotent: deleting a missing `(chunk_id, model_signature)` pair is a
|
||||
/// no-op. Intended for operator recovery after environmental failures (moved
|
||||
/// workspace, restored body files, fixed embedder config) — see #2358.
|
||||
pub fn clear_chunk_reembed_skipped(
|
||||
config: &Config,
|
||||
chunk_id: &str,
|
||||
model_signature: &str,
|
||||
) -> Result<()> {
|
||||
let chunk_id = validate_reembed_skip_key("chunk_id", chunk_id)?;
|
||||
let model_signature = validate_reembed_skip_key("model_signature", model_signature)?;
|
||||
with_connection(config, |conn| {
|
||||
conn.execute(
|
||||
"DELETE FROM mem_tree_chunk_reembed_skipped
|
||||
WHERE chunk_id = ?1 AND model_signature = ?2",
|
||||
rusqlite::params![chunk_id, model_signature],
|
||||
)?;
|
||||
log::debug!(
|
||||
"[memory::chunk_store] clear_chunk_reembed_skipped chunk_id={chunk_id} sig={model_signature}"
|
||||
);
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
/// Clear all chunk and summary tombstones for a model signature.
|
||||
///
|
||||
/// Returns the total number of rows removed across both tombstone tables.
|
||||
/// Idempotent when no tombstones exist for the signature.
|
||||
pub fn clear_reembed_skipped_for_signature(
|
||||
config: &Config,
|
||||
model_signature: &str,
|
||||
) -> Result<usize> {
|
||||
let model_signature = validate_reembed_skip_key("model_signature", model_signature)?;
|
||||
with_connection(config, |conn| {
|
||||
let chunk_deleted = conn.execute(
|
||||
"DELETE FROM mem_tree_chunk_reembed_skipped WHERE model_signature = ?1",
|
||||
rusqlite::params![model_signature],
|
||||
)?;
|
||||
let summary_deleted = conn.execute(
|
||||
"DELETE FROM mem_tree_summary_reembed_skipped WHERE model_signature = ?1",
|
||||
rusqlite::params![model_signature],
|
||||
)?;
|
||||
log::debug!(
|
||||
"[memory::chunk_store] clear_reembed_skipped_for_signature sig={model_signature} chunk_rows={chunk_deleted} summary_rows={summary_deleted}"
|
||||
);
|
||||
Ok(chunk_deleted + summary_deleted)
|
||||
})
|
||||
}
|
||||
|
||||
/// Bounds attacker-controlled ids/signatures passed to reembed-skipped admin
|
||||
/// helpers without affecting legitimate rows (typical ids are well under 512
|
||||
/// chars). Rejects NUL bytes so SQLite bindings cannot be truncated.
|
||||
const REEMBED_SKIP_KEY_MAX_LEN: usize = 2048;
|
||||
|
||||
pub(crate) fn validate_reembed_skip_key<'a>(label: &str, value: &'a str) -> Result<&'a str> {
|
||||
let trimmed = value.trim();
|
||||
if trimmed.is_empty() {
|
||||
anyhow::bail!("{label} must be non-empty");
|
||||
}
|
||||
if trimmed.len() > REEMBED_SKIP_KEY_MAX_LEN {
|
||||
anyhow::bail!("{label} exceeds maximum length ({REEMBED_SKIP_KEY_MAX_LEN})");
|
||||
}
|
||||
if trimmed.as_bytes().contains(&0) {
|
||||
anyhow::bail!("{label} must not contain NUL bytes");
|
||||
}
|
||||
Ok(trimmed)
|
||||
}
|
||||
|
||||
/// Transaction-scoped variant of [`set_chunk_embedding_for_signature`].
|
||||
///
|
||||
/// For callers that already hold a `Transaction` (e.g. the chunk-admission
|
||||
/// handler, which commits the sidecar row in the SAME tx as the lifecycle
|
||||
/// + score + job-enqueue writes — #1574 write-side cutover). Opening a fresh
|
||||
/// connection there would break atomicity / deadlock on the busy DB.
|
||||
pub(crate) fn set_chunk_embedding_for_signature_tx(
|
||||
tx: &rusqlite::Transaction<'_>,
|
||||
chunk_id: &str,
|
||||
model_signature: &str,
|
||||
embedding: &[f32],
|
||||
) -> Result<()> {
|
||||
upsert_chunk_embedding_conn(tx, chunk_id, model_signature, embedding)
|
||||
}
|
||||
|
||||
/// Fetch a chunk embedding for exactly one provider/model/dimension signature.
|
||||
pub fn get_chunk_embedding_for_signature(
|
||||
config: &Config,
|
||||
chunk_id: &str,
|
||||
model_signature: &str,
|
||||
) -> Result<Option<Vec<f32>>> {
|
||||
with_connection(config, |conn| {
|
||||
let row: Option<(Vec<u8>, i64)> = conn
|
||||
.query_row(
|
||||
"SELECT vector, dim
|
||||
FROM mem_tree_chunk_embeddings
|
||||
WHERE chunk_id = ?1 AND model_signature = ?2",
|
||||
rusqlite::params![chunk_id, model_signature],
|
||||
|r| Ok((r.get(0)?, r.get(1)?)),
|
||||
)
|
||||
.optional()?;
|
||||
match row {
|
||||
None => Ok(None),
|
||||
Some((bytes, dim)) => embedding_from_blob(&bytes, dim, "chunk embedding"),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Fetch a chunk's embedding for the active model signature.
|
||||
///
|
||||
/// #1574 cutover: reads the per-model `mem_tree_chunk_embeddings` sidecar at
|
||||
/// the active signature (via [`get_chunk_embedding_for_signature`]) instead of
|
||||
/// the legacy `mem_tree_chunks.embedding` column. Returns `Ok(None)` if the
|
||||
/// chunk has no vector under the active signature — e.g. during the §7
|
||||
/// backfill window, where this degrades retrieval gracefully (the row is
|
||||
/// simply absent from vector results, never cross-space compared).
|
||||
pub fn get_chunk_embedding(config: &Config, chunk_id: &str) -> Result<Option<Vec<f32>>> {
|
||||
let signature = tree_active_signature(config);
|
||||
get_chunk_embedding_for_signature(config, chunk_id, &signature)
|
||||
}
|
||||
|
||||
fn embedding_to_blob(embedding: &[f32]) -> Vec<u8> {
|
||||
embedding.iter().flat_map(|f| f.to_le_bytes()).collect()
|
||||
}
|
||||
|
||||
fn embedding_from_blob(bytes: &[u8], dim: i64, label: &str) -> Result<Option<Vec<f32>>> {
|
||||
if dim < 0 {
|
||||
anyhow::bail!("{label} has negative dimension {dim}");
|
||||
}
|
||||
if !bytes.len().is_multiple_of(4) {
|
||||
anyhow::bail!("{label} blob length {} not a multiple of 4", bytes.len());
|
||||
}
|
||||
let floats: Vec<f32> = bytes
|
||||
.chunks_exact(4)
|
||||
.map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
|
||||
.collect();
|
||||
if floats.len() != dim as usize {
|
||||
anyhow::bail!(
|
||||
"{label} dimension mismatch: dim column says {dim}, blob contains {} floats",
|
||||
floats.len()
|
||||
);
|
||||
}
|
||||
Ok(Some(floats))
|
||||
}
|
||||
#[path = "embeddings.rs"]
|
||||
mod embeddings;
|
||||
pub use embeddings::{
|
||||
clear_chunk_reembed_skipped, clear_reembed_skipped_for_signature, get_chunk_embedding,
|
||||
get_chunk_embedding_for_signature, mark_chunk_reembed_skipped, set_chunk_embedding,
|
||||
set_chunk_embedding_for_signature,
|
||||
};
|
||||
#[cfg(test)]
|
||||
pub(crate) use embeddings::{embedding_to_blob, REEMBED_SKIP_KEY_MAX_LEN};
|
||||
pub(crate) use embeddings::{
|
||||
has_uncovered_reembed_work, set_chunk_embedding_for_signature_tx, tree_active_signature,
|
||||
validate_reembed_skip_key,
|
||||
};
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "store_tests.rs"]
|
||||
|
||||
@@ -76,7 +76,7 @@ Insert is atomic (`ON CONFLICT DO NOTHING`) so concurrent shortens of the same U
|
||||
|
||||
- **Ids are content-addressed, not random**: same URL → same id (deterministic, deduped). Removing a link and re-shortening the same URL yields the same id again.
|
||||
- **Length threshold guards token waste**: the placeholder is ~24 bytes, so URLs below `DEFAULT_MIN_URL_LEN` (80) are left untouched by inbound rewrite.
|
||||
- **Trailing punctuation handling**: inbound rewrite and public-link tagging strip trailing `. , ; : !` so prose like "see https://…/path." doesn't capture the period into the stored/tagged URL.
|
||||
- **Trailing punctuation handling**: inbound rewrite and public-link tagging strip trailing `. , ; : !` so prose like "see an HTTPS URL ending in a period." doesn't capture the period into the stored/tagged URL.
|
||||
- **`append_user_id_to_public_links` is anchored to `openhm.xyz`** specifically and rejects lookalikes (`evil-openhm.xyz`, `openhm.xyz.evil.com`); it splits off `#fragment` so `?u=` always lands in the query, and is idempotent against existing `?u=`/`&u=`.
|
||||
- **`id_from_short`** accepts both `openhuman://link/<id>` and a bare hex `<id>`, lowercasing the result; non-hex input returns `None`.
|
||||
- **No agent tools, no event-bus subscribers, no `bus.rs`/`tools.rs`** — this domain is store + ops + RPC only.
|
||||
|
||||
@@ -318,803 +318,13 @@ impl Default for SecurityPolicy {
|
||||
}
|
||||
}
|
||||
|
||||
/// Environment variable names that can trigger arbitrary command execution
|
||||
/// when supplied as a leading inline assignment on an otherwise-allowed
|
||||
/// command. Each name here is either a hook variable that a downstream tool
|
||||
/// will spawn as a subprocess (`GIT_PAGER`, `GIT_SSH_COMMAND`, `EDITOR`,
|
||||
/// `LESS`/`LESSOPEN`, `MANPAGER`, `BROWSER`, `BAT_PAGER`), a runtime
|
||||
/// configuration knob that affects how Python or the shell evaluate user
|
||||
/// input (`PYTHONSTARTUP`, `BASH_ENV`, `ENV`, `PROMPT_COMMAND`), or a loader
|
||||
/// override that lets an attacker inject a library into the next process
|
||||
/// (`LD_PRELOAD`, `LD_LIBRARY_PATH`, `LD_AUDIT`, `DYLD_INSERT_LIBRARIES`,
|
||||
/// `DYLD_LIBRARY_PATH`, `DYLD_FORCE_FLAT_NAMESPACE`).
|
||||
///
|
||||
/// `PATH` and `SHELL` are listed so an inline override cannot redirect
|
||||
/// resolution of any allowed binary to an attacker-controlled path. `IFS`
|
||||
/// is listed because the shell uses it for word splitting and a malicious
|
||||
/// value can hide command boundaries from later parsers.
|
||||
const DANGEROUS_ENV_PREFIXES: &[&str] = &[
|
||||
"BASH_ENV",
|
||||
"BAT_PAGER",
|
||||
"BROWSER",
|
||||
"DYLD_FORCE_FLAT_NAMESPACE",
|
||||
"DYLD_INSERT_LIBRARIES",
|
||||
"DYLD_LIBRARY_PATH",
|
||||
"EDITOR",
|
||||
"ENV",
|
||||
"GIT_EDITOR",
|
||||
"GIT_EXTERNAL_DIFF",
|
||||
"GIT_EXTERNAL_FILTER",
|
||||
"GIT_PAGER",
|
||||
"GIT_SSH",
|
||||
"GIT_SSH_COMMAND",
|
||||
"IFS",
|
||||
"LD_AUDIT",
|
||||
"LD_LIBRARY_PATH",
|
||||
"LD_PRELOAD",
|
||||
"LESS",
|
||||
"LESSCLOSE",
|
||||
"LESSOPEN",
|
||||
"MANOPT",
|
||||
"MANPAGER",
|
||||
"PAGER",
|
||||
"PATH",
|
||||
"PROMPT_COMMAND",
|
||||
"PS1",
|
||||
"PS2",
|
||||
"PS3",
|
||||
"PS4",
|
||||
"PYTHONPATH",
|
||||
"PYTHONSTARTUP",
|
||||
"SHELL",
|
||||
"VISUAL",
|
||||
];
|
||||
|
||||
/// Returns true if `s` starts with one or more inline env assignments and any
|
||||
/// of the assigned names are in [`DANGEROUS_ENV_PREFIXES`].
|
||||
///
|
||||
/// The allowlist validation in [`SecurityPolicy::is_command_allowed`] uses
|
||||
/// [`skip_env_assignments`] to look past the env prefix before matching the
|
||||
/// command name. That leaves a class of attacks where the bare command (e.g.
|
||||
/// `git log`) is allowlisted but the env prefix mutates how it executes (e.g.
|
||||
/// `GIT_PAGER=<cmd> git log` — `git` spawns `<cmd>` as its pager). Because
|
||||
/// the prefix is stripped before allowlisting and the shell evaluates the
|
||||
/// prefix at execution time, the bypass lands without ever touching a
|
||||
/// blocked command name.
|
||||
///
|
||||
/// Treating any dangerous prefix as a denial keeps the allowlist
|
||||
/// semantically meaningful without having to enumerate every shape of every
|
||||
/// downstream tool's hook surface.
|
||||
fn has_dangerous_env_prefix(s: &str) -> bool {
|
||||
let mut rest = s.trim_start();
|
||||
loop {
|
||||
let Some(word) = rest.split_whitespace().next() else {
|
||||
return false;
|
||||
};
|
||||
if !word.contains('=') {
|
||||
return false;
|
||||
}
|
||||
if !word
|
||||
.chars()
|
||||
.next()
|
||||
.is_some_and(|c| c.is_ascii_alphabetic() || c == '_')
|
||||
{
|
||||
return false;
|
||||
}
|
||||
let (name, _) = word.split_once('=').unwrap_or((word, ""));
|
||||
let upper = name.to_ascii_uppercase();
|
||||
if DANGEROUS_ENV_PREFIXES.iter().any(|d| *d == upper.as_str()) {
|
||||
return true;
|
||||
}
|
||||
rest = rest[word.len()..].trim_start();
|
||||
}
|
||||
}
|
||||
|
||||
/// Skip leading environment variable assignments (e.g. `FOO=bar cmd args`).
|
||||
/// Returns the remainder starting at the first non-assignment word.
|
||||
fn skip_env_assignments(s: &str) -> &str {
|
||||
let mut rest = s;
|
||||
loop {
|
||||
let Some(word) = rest.split_whitespace().next() else {
|
||||
return rest;
|
||||
};
|
||||
// Environment assignment: contains '=' and starts with a letter or underscore
|
||||
if word.contains('=')
|
||||
&& word
|
||||
.chars()
|
||||
.next()
|
||||
.is_some_and(|c| c.is_ascii_alphabetic() || c == '_')
|
||||
{
|
||||
// Advance past this word
|
||||
rest = rest[word.len()..].trim_start();
|
||||
} else {
|
||||
return rest;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn command_basename(command: &str) -> &str {
|
||||
command
|
||||
.split(|ch| ch == '/' || ch == '\\')
|
||||
.next_back()
|
||||
.unwrap_or(command)
|
||||
}
|
||||
|
||||
fn normalized_command_name(command: &str) -> String {
|
||||
let command = command_basename(command).to_ascii_lowercase();
|
||||
command
|
||||
.strip_suffix(".exe")
|
||||
.unwrap_or(command.as_str())
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn is_python_command(command: &str) -> bool {
|
||||
let command = normalized_command_name(command);
|
||||
command == "python"
|
||||
|| command == "pythonw"
|
||||
|| command
|
||||
.strip_prefix("pythonw")
|
||||
.and_then(|suffix| suffix.chars().next())
|
||||
.is_some_and(|ch| ch.is_ascii_digit())
|
||||
|| command
|
||||
.strip_prefix("python")
|
||||
.and_then(|suffix| suffix.chars().next())
|
||||
.is_some_and(|ch| ch.is_ascii_digit())
|
||||
}
|
||||
|
||||
fn is_command_executor(command: &str) -> bool {
|
||||
let command = normalized_command_name(command);
|
||||
is_python_command(command.as_str())
|
||||
|| matches!(
|
||||
command.as_str(),
|
||||
"xargs"
|
||||
| "awk"
|
||||
| "gawk"
|
||||
| "mawk"
|
||||
| "nawk"
|
||||
| "perl"
|
||||
| "ruby"
|
||||
| "bash"
|
||||
| "sh"
|
||||
| "dash"
|
||||
| "zsh"
|
||||
| "ksh"
|
||||
| "fish"
|
||||
| "env"
|
||||
// JS/TS runtimes (the `node_exec`/`npm_exec` shell equivalents)
|
||||
| "node"
|
||||
| "nodejs"
|
||||
| "deno"
|
||||
| "bun"
|
||||
// Windows / PowerShell arbitrary-code launchers + LOLBins
|
||||
| "iex"
|
||||
| "invoke-expression"
|
||||
| "cmd"
|
||||
| "pwsh"
|
||||
| "powershell"
|
||||
| "wscript"
|
||||
| "cscript"
|
||||
| "mshta"
|
||||
| "rundll32"
|
||||
| "start-process"
|
||||
)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum QuoteState {
|
||||
None,
|
||||
Single,
|
||||
Double,
|
||||
}
|
||||
|
||||
/// Split a shell command into sub-commands by unquoted separators.
|
||||
///
|
||||
/// Separators:
|
||||
/// - `;` and newline
|
||||
/// - `|`
|
||||
/// - `&&`, `||`
|
||||
///
|
||||
/// Characters inside single or double quotes are treated as literals, so
|
||||
/// `sqlite3 db "SELECT 1; SELECT 2;"` remains a single segment.
|
||||
fn split_unquoted_segments(command: &str) -> Vec<String> {
|
||||
let mut segments = Vec::new();
|
||||
let mut current = String::new();
|
||||
let mut quote = QuoteState::None;
|
||||
let mut escaped = false;
|
||||
let mut chars = command.chars().peekable();
|
||||
|
||||
let push_segment = |segments: &mut Vec<String>, current: &mut String| {
|
||||
let trimmed = current.trim();
|
||||
if !trimmed.is_empty() {
|
||||
segments.push(trimmed.to_string());
|
||||
}
|
||||
current.clear();
|
||||
};
|
||||
|
||||
while let Some(ch) = chars.next() {
|
||||
match quote {
|
||||
QuoteState::Single => {
|
||||
if ch == '\'' {
|
||||
quote = QuoteState::None;
|
||||
}
|
||||
current.push(ch);
|
||||
}
|
||||
QuoteState::Double => {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
current.push(ch);
|
||||
continue;
|
||||
}
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
current.push(ch);
|
||||
continue;
|
||||
}
|
||||
if ch == '"' {
|
||||
quote = QuoteState::None;
|
||||
}
|
||||
current.push(ch);
|
||||
}
|
||||
QuoteState::None => {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
current.push(ch);
|
||||
continue;
|
||||
}
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
current.push(ch);
|
||||
continue;
|
||||
}
|
||||
|
||||
match ch {
|
||||
'\'' => {
|
||||
quote = QuoteState::Single;
|
||||
current.push(ch);
|
||||
}
|
||||
'"' => {
|
||||
quote = QuoteState::Double;
|
||||
current.push(ch);
|
||||
}
|
||||
';' | '\n' => push_segment(&mut segments, &mut current),
|
||||
'|' => {
|
||||
if chars.next_if_eq(&'|').is_some() {
|
||||
// Consume full `||`; both characters are separators.
|
||||
}
|
||||
push_segment(&mut segments, &mut current);
|
||||
}
|
||||
'&' => {
|
||||
if chars.next_if_eq(&'&').is_some() {
|
||||
// `&&` is a separator; single `&` is handled separately.
|
||||
push_segment(&mut segments, &mut current);
|
||||
} else {
|
||||
current.push(ch);
|
||||
}
|
||||
}
|
||||
_ => current.push(ch),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let trimmed = current.trim();
|
||||
if !trimmed.is_empty() {
|
||||
segments.push(trimmed.to_string());
|
||||
}
|
||||
|
||||
segments
|
||||
}
|
||||
|
||||
/// Detect a single unquoted `&` operator (background/chain). `&&` is allowed.
|
||||
///
|
||||
/// We treat any standalone `&` as unsafe in policy validation because it can
|
||||
/// chain hidden sub-commands and escape foreground timeout expectations.
|
||||
fn contains_unquoted_single_ampersand(command: &str) -> bool {
|
||||
let mut quote = QuoteState::None;
|
||||
let mut escaped = false;
|
||||
let mut chars = command.chars().peekable();
|
||||
|
||||
while let Some(ch) = chars.next() {
|
||||
match quote {
|
||||
QuoteState::Single => {
|
||||
if ch == '\'' {
|
||||
quote = QuoteState::None;
|
||||
}
|
||||
}
|
||||
QuoteState::Double => {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
continue;
|
||||
}
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
continue;
|
||||
}
|
||||
if ch == '"' {
|
||||
quote = QuoteState::None;
|
||||
}
|
||||
}
|
||||
QuoteState::None => {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
continue;
|
||||
}
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
continue;
|
||||
}
|
||||
match ch {
|
||||
'\'' => quote = QuoteState::Single,
|
||||
'"' => quote = QuoteState::Double,
|
||||
'&' => {
|
||||
if chars.next_if_eq(&'&').is_none() {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
/// Like [`contains_unquoted_single_ampersand`] but ignores file-descriptor
|
||||
/// duplication redirects, where the `&` is part of a redirect operator rather
|
||||
/// than a background/separator: `2>&1`, `>&2` (prev char `>`), and `&>file`
|
||||
/// (next char `>`). Used by [`has_hidden_execution`] so a benign `… 2>&1` —
|
||||
/// which `classify_command` already accounts for as a `Write` redirect — is not
|
||||
/// mistaken for a backgrounded command and hard-blocked after the human
|
||||
/// approved it. A standalone `&` (e.g. `cmd &`, `a & b`) still returns true,
|
||||
/// since it can run a second command `classify_command` wouldn't see.
|
||||
fn contains_unquoted_background_ampersand(command: &str) -> bool {
|
||||
let mut quote = QuoteState::None;
|
||||
let mut escaped = false;
|
||||
let mut prev = '\0';
|
||||
let mut chars = command.chars().peekable();
|
||||
|
||||
while let Some(ch) = chars.next() {
|
||||
match quote {
|
||||
QuoteState::Single => {
|
||||
if ch == '\'' {
|
||||
quote = QuoteState::None;
|
||||
}
|
||||
}
|
||||
QuoteState::Double => {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
prev = ch;
|
||||
continue;
|
||||
}
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
prev = ch;
|
||||
continue;
|
||||
}
|
||||
if ch == '"' {
|
||||
quote = QuoteState::None;
|
||||
}
|
||||
}
|
||||
QuoteState::None => {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
prev = ch;
|
||||
continue;
|
||||
}
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
prev = ch;
|
||||
continue;
|
||||
}
|
||||
match ch {
|
||||
'\'' => quote = QuoteState::Single,
|
||||
'"' => quote = QuoteState::Double,
|
||||
'&' => {
|
||||
if chars.next_if_eq(&'&').is_some() {
|
||||
// `&&` logical AND — consume both, not background.
|
||||
} else {
|
||||
let next = chars.peek().copied().unwrap_or('\0');
|
||||
// Skip fd-dup redirects: `2>&1`/`>&2` (prev `>`) and
|
||||
// `&>file` (next `>`).
|
||||
if prev != '>' && next != '>' {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
prev = ch;
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
/// Detect an unquoted character in a shell command.
|
||||
fn contains_unquoted_char(command: &str, target: char) -> bool {
|
||||
let mut quote = QuoteState::None;
|
||||
let mut escaped = false;
|
||||
|
||||
for ch in command.chars() {
|
||||
match quote {
|
||||
QuoteState::Single => {
|
||||
if ch == '\'' {
|
||||
quote = QuoteState::None;
|
||||
}
|
||||
}
|
||||
QuoteState::Double => {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
continue;
|
||||
}
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
continue;
|
||||
}
|
||||
if ch == '"' {
|
||||
quote = QuoteState::None;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
QuoteState::None => {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
continue;
|
||||
}
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
continue;
|
||||
}
|
||||
match ch {
|
||||
'\'' => quote = QuoteState::Single,
|
||||
'"' => quote = QuoteState::Double,
|
||||
_ if ch == target => return true,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
/// Provably read-only command bases (cross-platform union). A base **not** in
|
||||
/// this set — and not a recognized network/destructive/executor command, nor a
|
||||
/// read-only verb of git/npm/cargo — falls through to [`CommandClass::Write`]
|
||||
/// (the classifier is fail-closed). Conservative on purpose: anything that can
|
||||
/// write a file under a common flag is intentionally omitted (`sort -o`, `tee`).
|
||||
const READ_ONLY_BASES: &[&str] = &[
|
||||
// POSIX inspection / read-only coreutils
|
||||
"ls",
|
||||
"cat",
|
||||
"pwd",
|
||||
"echo",
|
||||
"wc",
|
||||
"head",
|
||||
"tail",
|
||||
"date",
|
||||
"grep",
|
||||
"egrep",
|
||||
"fgrep",
|
||||
"rg",
|
||||
"which",
|
||||
"whoami",
|
||||
"id",
|
||||
"hostname",
|
||||
"uname",
|
||||
"printenv",
|
||||
"stat",
|
||||
"file",
|
||||
"du",
|
||||
"df",
|
||||
"tree",
|
||||
"realpath",
|
||||
"readlink",
|
||||
"dirname",
|
||||
"basename",
|
||||
"cmp",
|
||||
"true",
|
||||
"false",
|
||||
"sleep",
|
||||
"seq",
|
||||
"tty",
|
||||
"groups",
|
||||
"locale",
|
||||
"ps",
|
||||
"top",
|
||||
"free",
|
||||
"uptime",
|
||||
"lsblk",
|
||||
"lscpu",
|
||||
"cut",
|
||||
// Windows cmd / PowerShell read verbs + common aliases
|
||||
"dir",
|
||||
"type",
|
||||
"where",
|
||||
"whereis",
|
||||
"get-childitem",
|
||||
"gci",
|
||||
"get-content",
|
||||
"gc",
|
||||
"get-location",
|
||||
"gl",
|
||||
"select-string",
|
||||
"sls",
|
||||
"measure-object",
|
||||
"get-item",
|
||||
"gi",
|
||||
"test-path",
|
||||
"resolve-path",
|
||||
"get-command",
|
||||
"gcm",
|
||||
"get-process",
|
||||
];
|
||||
|
||||
/// Commands that reach the network. Always-ask in every acting tier.
|
||||
const NETWORK_BASES: &[&str] = &[
|
||||
"curl",
|
||||
"wget",
|
||||
"ssh",
|
||||
"scp",
|
||||
"sftp",
|
||||
"rsync",
|
||||
"nc",
|
||||
"ncat",
|
||||
"netcat",
|
||||
"telnet",
|
||||
"ftp",
|
||||
"tftp",
|
||||
"socat",
|
||||
// Windows / PowerShell
|
||||
"invoke-webrequest",
|
||||
"iwr",
|
||||
"invoke-restmethod",
|
||||
"irm",
|
||||
"start-bitstransfer",
|
||||
"bitsadmin",
|
||||
];
|
||||
|
||||
/// Catastrophic / irreversible / privilege / system-control bases. Always-ask
|
||||
/// in every acting tier (Full included). Coarse on the broad Windows verbs
|
||||
/// (`reg`/`net`/`sc`) — over-prompting there is the safe default.
|
||||
const DESTRUCTIVE_BASES: &[&str] = &[
|
||||
// POSIX privilege / disk / system-control
|
||||
"sudo",
|
||||
"su",
|
||||
"doas",
|
||||
"dd",
|
||||
"mkfs",
|
||||
"fdisk",
|
||||
"sfdisk",
|
||||
"parted",
|
||||
"wipefs",
|
||||
"shred",
|
||||
"shutdown",
|
||||
"reboot",
|
||||
"halt",
|
||||
"poweroff",
|
||||
"init",
|
||||
"telinit",
|
||||
"mount",
|
||||
"umount",
|
||||
"swapoff",
|
||||
"iptables",
|
||||
"ip6tables",
|
||||
"nft",
|
||||
"ufw",
|
||||
"firewall-cmd",
|
||||
"useradd",
|
||||
"userdel",
|
||||
"usermod",
|
||||
"groupadd",
|
||||
"groupdel",
|
||||
"passwd",
|
||||
"chpasswd",
|
||||
"visudo",
|
||||
"modprobe",
|
||||
"insmod",
|
||||
"rmmod",
|
||||
// Windows / PowerShell
|
||||
"format",
|
||||
"diskpart",
|
||||
"bcdedit",
|
||||
"takeown",
|
||||
"cipher",
|
||||
"vssadmin",
|
||||
"reg",
|
||||
"regedit",
|
||||
"runas",
|
||||
"sc",
|
||||
"net",
|
||||
"set-executionpolicy",
|
||||
"stop-computer",
|
||||
"restart-computer",
|
||||
"clear-disk",
|
||||
"format-volume",
|
||||
"remove-partition",
|
||||
"disable-computerrestore",
|
||||
];
|
||||
|
||||
/// Git subcommands that only read repository state. Anything else — including
|
||||
/// `commit`/`push`/`branch`/`config`/unknown/bare `git` — is fail-closed to
|
||||
/// `Write`.
|
||||
const GIT_READ_VERBS: &[&str] = &[
|
||||
"status",
|
||||
"log",
|
||||
"diff",
|
||||
"show",
|
||||
"remote",
|
||||
"describe",
|
||||
"blame",
|
||||
"ls-files",
|
||||
"ls-tree",
|
||||
"rev-parse",
|
||||
"cat-file",
|
||||
"shortlog",
|
||||
"reflog",
|
||||
"rev-list",
|
||||
"name-rev",
|
||||
"var",
|
||||
"check-ignore",
|
||||
"check-attr",
|
||||
"verify-commit",
|
||||
"count-objects",
|
||||
"fsck",
|
||||
"whatchanged",
|
||||
"grep",
|
||||
"version",
|
||||
"help",
|
||||
];
|
||||
|
||||
/// npm/pnpm/yarn read-only subcommands. `install`/`run`/`test`/`exec` (which
|
||||
/// run arbitrary scripts) and unknown verbs are fail-closed to `Write`.
|
||||
const NODE_PKG_READ_VERBS: &[&str] = &[
|
||||
"ls", "list", "view", "info", "outdated", "ping", "whoami", "help", "why", "audit", "doctor",
|
||||
];
|
||||
|
||||
/// cargo read-only subcommands. `build`/`run`/`test`/`check` compile and may
|
||||
/// run build scripts, so they are fail-closed to `Write`.
|
||||
const CARGO_READ_VERBS: &[&str] = &["tree", "metadata", "search", "info", "version", "help"];
|
||||
|
||||
/// Detect a pacman *install/upgrade* from its bundled operation flag.
|
||||
///
|
||||
/// pacman packs its operation and modifiers into a single flag (`-Syu`, `-Ss`),
|
||||
/// and `args` reach us already lowercased — so the `-S` (sync) operation is
|
||||
/// indistinguishable from a literal `-s` by case alone. We therefore key off
|
||||
/// the *modifier* letters instead of a blanket `starts_with("-s")`, which would
|
||||
/// over-match every read-only `-S` query: a `-S`-family flag mutates the host
|
||||
/// only when it carries none of pacman's read-only query modifiers — search
|
||||
/// (`s`), info (`i`), list (`l`), groups (`g`) or print (`p`). So `-S pkg`,
|
||||
/// `-Sy`, `-Syu` are installs while `-Ss`/`-Si`/`-Sl`/`-Sg`/`-Sp` are reads.
|
||||
fn is_pacman_install(args: &[String]) -> bool {
|
||||
args.iter().any(|a| {
|
||||
a.strip_prefix("-s")
|
||||
.is_some_and(|modifiers| !modifiers.contains(['s', 'i', 'l', 'g', 'p']))
|
||||
})
|
||||
}
|
||||
|
||||
/// Detect a package-manager *install* invocation. These mutate the host /
|
||||
/// global environment, so they are the always-ask `Install` bucket (even in
|
||||
/// Full) — the same gate the dedicated `install_tool` enforces, applied to the
|
||||
/// shell escape hatch. Project-local installs (`npm install` without `-g`,
|
||||
/// `cargo add`) are ordinary `Write`s and are deliberately NOT matched here.
|
||||
/// `args` are already lowercased by the caller.
|
||||
fn is_install_command(base: &str, args: &[String]) -> bool {
|
||||
let has = |needle: &str| args.iter().any(|a| a == needle);
|
||||
let first_is = |verb: &str| args.first().map(String::as_str) == Some(verb);
|
||||
match base {
|
||||
// System package managers.
|
||||
"apt" | "apt-get" | "dnf" | "yum" | "zypper" => has("install"),
|
||||
"pacman" => is_pacman_install(args),
|
||||
"apk" => has("add"),
|
||||
"brew" | "snap" | "flatpak" | "winget" | "choco" | "scoop" => has("install"),
|
||||
// Language package managers — host/global-modifying installs only.
|
||||
"pip" | "pip3" | "pipx" | "gem" | "go" | "cargo" => first_is("install"),
|
||||
"npm" | "pnpm" => {
|
||||
(has("install") || has("i") || has("add")) && (has("-g") || has("--global"))
|
||||
}
|
||||
"yarn" => has("global"),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Classify a single already-split shell segment. `base` is the normalized
|
||||
/// (lowercased, `.exe`-stripped, basename-only) program name; `args` are the
|
||||
/// lowercased remaining words; `joined` is the lowercased segment used for
|
||||
/// pattern matching. Fail-closed: an unrecognized base resolves to `Write`.
|
||||
fn classify_segment(base: &str, args: &[String], joined: &str) -> CommandClass {
|
||||
// Catastrophic patterns first — they win regardless of the base command.
|
||||
if joined.contains("rm -rf /") || joined.contains("rm -fr /") || joined.contains(":(){:|:&};:")
|
||||
{
|
||||
return CommandClass::Destructive;
|
||||
}
|
||||
if DESTRUCTIVE_BASES.contains(&base) {
|
||||
return CommandClass::Destructive;
|
||||
}
|
||||
if NETWORK_BASES.contains(&base) {
|
||||
return CommandClass::Network;
|
||||
}
|
||||
// Package installs mutate the host → always-ask Install bucket (closes the
|
||||
// shell escape hatch around `install_tool`).
|
||||
if is_install_command(base, args) {
|
||||
return CommandClass::Install;
|
||||
}
|
||||
// Interpreters / code executors run arbitrary code. Fail-closed to Write
|
||||
// (not Destructive) so Full can still run code while Supervised prompts.
|
||||
if is_command_executor(base) {
|
||||
return CommandClass::Write;
|
||||
}
|
||||
// `find` is read-only unless it executes commands or deletes files.
|
||||
if base == "find" {
|
||||
if args.iter().any(|a| {
|
||||
matches!(
|
||||
a.as_str(),
|
||||
"-exec" | "-execdir" | "-ok" | "-okdir" | "-delete"
|
||||
)
|
||||
}) {
|
||||
return CommandClass::Write;
|
||||
}
|
||||
return CommandClass::Read;
|
||||
}
|
||||
// Verb-sensitive VCS / package tools.
|
||||
if base == "git" {
|
||||
return verb_class(args, GIT_READ_VERBS);
|
||||
}
|
||||
if matches!(base, "npm" | "pnpm" | "yarn") {
|
||||
return verb_class(args, NODE_PKG_READ_VERBS);
|
||||
}
|
||||
if base == "cargo" {
|
||||
return verb_class(args, CARGO_READ_VERBS);
|
||||
}
|
||||
if READ_ONLY_BASES.contains(&base) {
|
||||
return CommandClass::Read;
|
||||
}
|
||||
// Fail closed: unknown or known-mutating base → Write.
|
||||
CommandClass::Write
|
||||
}
|
||||
|
||||
/// `Read` when the first subcommand word is in `read_verbs`, else fail-closed
|
||||
/// `Write`. Mirrors the `args.first()` verb check used by `command_risk_level`.
|
||||
fn verb_class(args: &[String], read_verbs: &[&str]) -> CommandClass {
|
||||
match args.first().map(String::as_str) {
|
||||
Some(verb) if read_verbs.contains(&verb) => CommandClass::Read,
|
||||
_ => CommandClass::Write,
|
||||
}
|
||||
}
|
||||
|
||||
/// Structural-safety guard for the harness-gated command flow (Option 2). Even
|
||||
/// after a human approves a command, a hidden subshell / command substitution /
|
||||
/// output redirect / `tee` / background `&` could smuggle a *different* command
|
||||
/// past the approval summary, so these are refused outside Full (which is
|
||||
/// trusted to use redirects and pipes). Mirrors the structural checks in
|
||||
/// [`SecurityPolicy::is_command_allowed`].
|
||||
/// Detect shell structure that can **hide execution** from `classify_command`,
|
||||
/// which only inspects the base command of each `;`/`&&`/`|` segment. Command
|
||||
/// and process substitution and backticks run an *inner* command classification
|
||||
/// can't see (`echo $(rm -rf ~)` classifies as `echo` = Read and would run
|
||||
/// unprompted), and a trailing `&` detaches a process past the gate — so these
|
||||
/// stay hard-blocked outside Full.
|
||||
///
|
||||
/// Deliberately NOT flagged here: plain redirects (`>`, `2>&1`, `2>/dev/null`),
|
||||
/// `tee`, and `${VAR}` expansion. `classify_command` already lifts a redirect /
|
||||
/// `tee` to `Write`, so the gate prompts and — once the human approves — the
|
||||
/// command MUST actually run. Re-blocking an approved `… 2>&1` here was the bug
|
||||
/// that made Supervised mode unusable: every command the agent wrote carried a
|
||||
/// `2>&1`, got approved, then silently failed this in-tool guard and never ran.
|
||||
fn has_hidden_execution(command: &str) -> bool {
|
||||
// The backtick check is deliberately NOT quote-aware: any backtick in the
|
||||
// command string is blocked, even inside a double-quoted literal. Over-
|
||||
// blocking is the safe direction here. (By contrast the `&` case below is
|
||||
// quote-aware via `contains_unquoted_background_ampersand`, because that one
|
||||
// must still allow benign fd-dup redirects like `2>&1`.)
|
||||
command.contains('`')
|
||||
|| command.contains("$(")
|
||||
|| command.contains("<(")
|
||||
|| command.contains(">(")
|
||||
|| contains_unquoted_background_ampersand(command)
|
||||
}
|
||||
#[path = "policy_command.rs"]
|
||||
mod policy_command;
|
||||
use policy_command::{
|
||||
classify_segment, command_basename, contains_unquoted_char, contains_unquoted_single_ampersand,
|
||||
has_dangerous_env_prefix, has_hidden_execution, is_command_executor, normalized_command_name,
|
||||
skip_env_assignments, split_unquoted_segments,
|
||||
};
|
||||
|
||||
impl SecurityPolicy {
|
||||
/// Classify command risk. Any high-risk segment marks the whole command high.
|
||||
|
||||
@@ -0,0 +1,799 @@
|
||||
use super::CommandClass;
|
||||
|
||||
/// Environment variable names that can trigger arbitrary command execution
|
||||
/// when supplied as a leading inline assignment on an otherwise-allowed
|
||||
/// command. Each name here is either a hook variable that a downstream tool
|
||||
/// will spawn as a subprocess (`GIT_PAGER`, `GIT_SSH_COMMAND`, `EDITOR`,
|
||||
/// `LESS`/`LESSOPEN`, `MANPAGER`, `BROWSER`, `BAT_PAGER`), a runtime
|
||||
/// configuration knob that affects how Python or the shell evaluate user
|
||||
/// input (`PYTHONSTARTUP`, `BASH_ENV`, `ENV`, `PROMPT_COMMAND`), or a loader
|
||||
/// override that lets an attacker inject a library into the next process
|
||||
/// (`LD_PRELOAD`, `LD_LIBRARY_PATH`, `LD_AUDIT`, `DYLD_INSERT_LIBRARIES`,
|
||||
/// `DYLD_LIBRARY_PATH`, `DYLD_FORCE_FLAT_NAMESPACE`).
|
||||
///
|
||||
/// `PATH` and `SHELL` are listed so an inline override cannot redirect
|
||||
/// resolution of any allowed binary to an attacker-controlled path. `IFS`
|
||||
/// is listed because the shell uses it for word splitting and a malicious
|
||||
/// value can hide command boundaries from later parsers.
|
||||
const DANGEROUS_ENV_PREFIXES: &[&str] = &[
|
||||
"BASH_ENV",
|
||||
"BAT_PAGER",
|
||||
"BROWSER",
|
||||
"DYLD_FORCE_FLAT_NAMESPACE",
|
||||
"DYLD_INSERT_LIBRARIES",
|
||||
"DYLD_LIBRARY_PATH",
|
||||
"EDITOR",
|
||||
"ENV",
|
||||
"GIT_EDITOR",
|
||||
"GIT_EXTERNAL_DIFF",
|
||||
"GIT_EXTERNAL_FILTER",
|
||||
"GIT_PAGER",
|
||||
"GIT_SSH",
|
||||
"GIT_SSH_COMMAND",
|
||||
"IFS",
|
||||
"LD_AUDIT",
|
||||
"LD_LIBRARY_PATH",
|
||||
"LD_PRELOAD",
|
||||
"LESS",
|
||||
"LESSCLOSE",
|
||||
"LESSOPEN",
|
||||
"MANOPT",
|
||||
"MANPAGER",
|
||||
"PAGER",
|
||||
"PATH",
|
||||
"PROMPT_COMMAND",
|
||||
"PS1",
|
||||
"PS2",
|
||||
"PS3",
|
||||
"PS4",
|
||||
"PYTHONPATH",
|
||||
"PYTHONSTARTUP",
|
||||
"SHELL",
|
||||
"VISUAL",
|
||||
];
|
||||
|
||||
/// Returns true if `s` starts with one or more inline env assignments and any
|
||||
/// of the assigned names are in [`DANGEROUS_ENV_PREFIXES`].
|
||||
///
|
||||
/// The allowlist validation in [`SecurityPolicy::is_command_allowed`] uses
|
||||
/// [`skip_env_assignments`] to look past the env prefix before matching the
|
||||
/// command name. That leaves a class of attacks where the bare command (e.g.
|
||||
/// `git log`) is allowlisted but the env prefix mutates how it executes (e.g.
|
||||
/// `GIT_PAGER=<cmd> git log` — `git` spawns `<cmd>` as its pager). Because
|
||||
/// the prefix is stripped before allowlisting and the shell evaluates the
|
||||
/// prefix at execution time, the bypass lands without ever touching a
|
||||
/// blocked command name.
|
||||
///
|
||||
/// Treating any dangerous prefix as a denial keeps the allowlist
|
||||
/// semantically meaningful without having to enumerate every shape of every
|
||||
/// downstream tool's hook surface.
|
||||
pub(super) fn has_dangerous_env_prefix(s: &str) -> bool {
|
||||
let mut rest = s.trim_start();
|
||||
loop {
|
||||
let Some(word) = rest.split_whitespace().next() else {
|
||||
return false;
|
||||
};
|
||||
if !word.contains('=') {
|
||||
return false;
|
||||
}
|
||||
if !word
|
||||
.chars()
|
||||
.next()
|
||||
.is_some_and(|c| c.is_ascii_alphabetic() || c == '_')
|
||||
{
|
||||
return false;
|
||||
}
|
||||
let (name, _) = word.split_once('=').unwrap_or((word, ""));
|
||||
let upper = name.to_ascii_uppercase();
|
||||
if DANGEROUS_ENV_PREFIXES.iter().any(|d| *d == upper.as_str()) {
|
||||
return true;
|
||||
}
|
||||
rest = rest[word.len()..].trim_start();
|
||||
}
|
||||
}
|
||||
|
||||
/// Skip leading environment variable assignments (e.g. `FOO=bar cmd args`).
|
||||
/// Returns the remainder starting at the first non-assignment word.
|
||||
pub(super) fn skip_env_assignments(s: &str) -> &str {
|
||||
let mut rest = s;
|
||||
loop {
|
||||
let Some(word) = rest.split_whitespace().next() else {
|
||||
return rest;
|
||||
};
|
||||
// Environment assignment: contains '=' and starts with a letter or underscore
|
||||
if word.contains('=')
|
||||
&& word
|
||||
.chars()
|
||||
.next()
|
||||
.is_some_and(|c| c.is_ascii_alphabetic() || c == '_')
|
||||
{
|
||||
// Advance past this word
|
||||
rest = rest[word.len()..].trim_start();
|
||||
} else {
|
||||
return rest;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn command_basename(command: &str) -> &str {
|
||||
command
|
||||
.split(|ch| ch == '/' || ch == '\\')
|
||||
.next_back()
|
||||
.unwrap_or(command)
|
||||
}
|
||||
|
||||
pub(super) fn normalized_command_name(command: &str) -> String {
|
||||
let command = command_basename(command).to_ascii_lowercase();
|
||||
command
|
||||
.strip_suffix(".exe")
|
||||
.unwrap_or(command.as_str())
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn is_python_command(command: &str) -> bool {
|
||||
let command = normalized_command_name(command);
|
||||
command == "python"
|
||||
|| command == "pythonw"
|
||||
|| command
|
||||
.strip_prefix("pythonw")
|
||||
.and_then(|suffix| suffix.chars().next())
|
||||
.is_some_and(|ch| ch.is_ascii_digit())
|
||||
|| command
|
||||
.strip_prefix("python")
|
||||
.and_then(|suffix| suffix.chars().next())
|
||||
.is_some_and(|ch| ch.is_ascii_digit())
|
||||
}
|
||||
|
||||
pub(super) fn is_command_executor(command: &str) -> bool {
|
||||
let command = normalized_command_name(command);
|
||||
is_python_command(command.as_str())
|
||||
|| matches!(
|
||||
command.as_str(),
|
||||
"xargs"
|
||||
| "awk"
|
||||
| "gawk"
|
||||
| "mawk"
|
||||
| "nawk"
|
||||
| "perl"
|
||||
| "ruby"
|
||||
| "bash"
|
||||
| "sh"
|
||||
| "dash"
|
||||
| "zsh"
|
||||
| "ksh"
|
||||
| "fish"
|
||||
| "env"
|
||||
// JS/TS runtimes (the `node_exec`/`npm_exec` shell equivalents)
|
||||
| "node"
|
||||
| "nodejs"
|
||||
| "deno"
|
||||
| "bun"
|
||||
// Windows / PowerShell arbitrary-code launchers + LOLBins
|
||||
| "iex"
|
||||
| "invoke-expression"
|
||||
| "cmd"
|
||||
| "pwsh"
|
||||
| "powershell"
|
||||
| "wscript"
|
||||
| "cscript"
|
||||
| "mshta"
|
||||
| "rundll32"
|
||||
| "start-process"
|
||||
)
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum QuoteState {
|
||||
None,
|
||||
Single,
|
||||
Double,
|
||||
}
|
||||
|
||||
/// Split a shell command into sub-commands by unquoted separators.
|
||||
///
|
||||
/// Separators:
|
||||
/// - `;` and newline
|
||||
/// - `|`
|
||||
/// - `&&`, `||`
|
||||
///
|
||||
/// Characters inside single or double quotes are treated as literals, so
|
||||
/// `sqlite3 db "SELECT 1; SELECT 2;"` remains a single segment.
|
||||
pub(super) fn split_unquoted_segments(command: &str) -> Vec<String> {
|
||||
let mut segments = Vec::new();
|
||||
let mut current = String::new();
|
||||
let mut quote = QuoteState::None;
|
||||
let mut escaped = false;
|
||||
let mut chars = command.chars().peekable();
|
||||
|
||||
let push_segment = |segments: &mut Vec<String>, current: &mut String| {
|
||||
let trimmed = current.trim();
|
||||
if !trimmed.is_empty() {
|
||||
segments.push(trimmed.to_string());
|
||||
}
|
||||
current.clear();
|
||||
};
|
||||
|
||||
while let Some(ch) = chars.next() {
|
||||
match quote {
|
||||
QuoteState::Single => {
|
||||
if ch == '\'' {
|
||||
quote = QuoteState::None;
|
||||
}
|
||||
current.push(ch);
|
||||
}
|
||||
QuoteState::Double => {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
current.push(ch);
|
||||
continue;
|
||||
}
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
current.push(ch);
|
||||
continue;
|
||||
}
|
||||
if ch == '"' {
|
||||
quote = QuoteState::None;
|
||||
}
|
||||
current.push(ch);
|
||||
}
|
||||
QuoteState::None => {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
current.push(ch);
|
||||
continue;
|
||||
}
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
current.push(ch);
|
||||
continue;
|
||||
}
|
||||
|
||||
match ch {
|
||||
'\'' => {
|
||||
quote = QuoteState::Single;
|
||||
current.push(ch);
|
||||
}
|
||||
'"' => {
|
||||
quote = QuoteState::Double;
|
||||
current.push(ch);
|
||||
}
|
||||
';' | '\n' => push_segment(&mut segments, &mut current),
|
||||
'|' => {
|
||||
if chars.next_if_eq(&'|').is_some() {
|
||||
// Consume full `||`; both characters are separators.
|
||||
}
|
||||
push_segment(&mut segments, &mut current);
|
||||
}
|
||||
'&' => {
|
||||
if chars.next_if_eq(&'&').is_some() {
|
||||
// `&&` is a separator; single `&` is handled separately.
|
||||
push_segment(&mut segments, &mut current);
|
||||
} else {
|
||||
current.push(ch);
|
||||
}
|
||||
}
|
||||
_ => current.push(ch),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let trimmed = current.trim();
|
||||
if !trimmed.is_empty() {
|
||||
segments.push(trimmed.to_string());
|
||||
}
|
||||
|
||||
segments
|
||||
}
|
||||
|
||||
/// Detect a single unquoted `&` operator (background/chain). `&&` is allowed.
|
||||
///
|
||||
/// We treat any standalone `&` as unsafe in policy validation because it can
|
||||
/// chain hidden sub-commands and escape foreground timeout expectations.
|
||||
pub(super) fn contains_unquoted_single_ampersand(command: &str) -> bool {
|
||||
let mut quote = QuoteState::None;
|
||||
let mut escaped = false;
|
||||
let mut chars = command.chars().peekable();
|
||||
|
||||
while let Some(ch) = chars.next() {
|
||||
match quote {
|
||||
QuoteState::Single => {
|
||||
if ch == '\'' {
|
||||
quote = QuoteState::None;
|
||||
}
|
||||
}
|
||||
QuoteState::Double => {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
continue;
|
||||
}
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
continue;
|
||||
}
|
||||
if ch == '"' {
|
||||
quote = QuoteState::None;
|
||||
}
|
||||
}
|
||||
QuoteState::None => {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
continue;
|
||||
}
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
continue;
|
||||
}
|
||||
match ch {
|
||||
'\'' => quote = QuoteState::Single,
|
||||
'"' => quote = QuoteState::Double,
|
||||
'&' => {
|
||||
if chars.next_if_eq(&'&').is_none() {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
/// Like [`contains_unquoted_single_ampersand`] but ignores file-descriptor
|
||||
/// duplication redirects, where the `&` is part of a redirect operator rather
|
||||
/// than a background/separator: `2>&1`, `>&2` (prev char `>`), and `&>file`
|
||||
/// (next char `>`). Used by [`has_hidden_execution`] so a benign `… 2>&1` —
|
||||
/// which `classify_command` already accounts for as a `Write` redirect — is not
|
||||
/// mistaken for a backgrounded command and hard-blocked after the human
|
||||
/// approved it. A standalone `&` (e.g. `cmd &`, `a & b`) still returns true,
|
||||
/// since it can run a second command `classify_command` wouldn't see.
|
||||
fn contains_unquoted_background_ampersand(command: &str) -> bool {
|
||||
let mut quote = QuoteState::None;
|
||||
let mut escaped = false;
|
||||
let mut prev = '\0';
|
||||
let mut chars = command.chars().peekable();
|
||||
|
||||
while let Some(ch) = chars.next() {
|
||||
match quote {
|
||||
QuoteState::Single => {
|
||||
if ch == '\'' {
|
||||
quote = QuoteState::None;
|
||||
}
|
||||
}
|
||||
QuoteState::Double => {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
prev = ch;
|
||||
continue;
|
||||
}
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
prev = ch;
|
||||
continue;
|
||||
}
|
||||
if ch == '"' {
|
||||
quote = QuoteState::None;
|
||||
}
|
||||
}
|
||||
QuoteState::None => {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
prev = ch;
|
||||
continue;
|
||||
}
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
prev = ch;
|
||||
continue;
|
||||
}
|
||||
match ch {
|
||||
'\'' => quote = QuoteState::Single,
|
||||
'"' => quote = QuoteState::Double,
|
||||
'&' => {
|
||||
if chars.next_if_eq(&'&').is_some() {
|
||||
// `&&` logical AND — consume both, not background.
|
||||
} else {
|
||||
let next = chars.peek().copied().unwrap_or('\0');
|
||||
// Skip fd-dup redirects: `2>&1`/`>&2` (prev `>`) and
|
||||
// `&>file` (next `>`).
|
||||
if prev != '>' && next != '>' {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
prev = ch;
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
/// Detect an unquoted character in a shell command.
|
||||
pub(super) fn contains_unquoted_char(command: &str, target: char) -> bool {
|
||||
let mut quote = QuoteState::None;
|
||||
let mut escaped = false;
|
||||
|
||||
for ch in command.chars() {
|
||||
match quote {
|
||||
QuoteState::Single => {
|
||||
if ch == '\'' {
|
||||
quote = QuoteState::None;
|
||||
}
|
||||
}
|
||||
QuoteState::Double => {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
continue;
|
||||
}
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
continue;
|
||||
}
|
||||
if ch == '"' {
|
||||
quote = QuoteState::None;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
QuoteState::None => {
|
||||
if escaped {
|
||||
escaped = false;
|
||||
continue;
|
||||
}
|
||||
if ch == '\\' {
|
||||
escaped = true;
|
||||
continue;
|
||||
}
|
||||
match ch {
|
||||
'\'' => quote = QuoteState::Single,
|
||||
'"' => quote = QuoteState::Double,
|
||||
_ if ch == target => return true,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
false
|
||||
}
|
||||
|
||||
/// Provably read-only command bases (cross-platform union). A base **not** in
|
||||
/// this set — and not a recognized network/destructive/executor command, nor a
|
||||
/// read-only verb of git/npm/cargo — falls through to [`CommandClass::Write`]
|
||||
/// (the classifier is fail-closed). Conservative on purpose: anything that can
|
||||
/// write a file under a common flag is intentionally omitted (`sort -o`, `tee`).
|
||||
const READ_ONLY_BASES: &[&str] = &[
|
||||
// POSIX inspection / read-only coreutils
|
||||
"ls",
|
||||
"cat",
|
||||
"pwd",
|
||||
"echo",
|
||||
"wc",
|
||||
"head",
|
||||
"tail",
|
||||
"date",
|
||||
"grep",
|
||||
"egrep",
|
||||
"fgrep",
|
||||
"rg",
|
||||
"which",
|
||||
"whoami",
|
||||
"id",
|
||||
"hostname",
|
||||
"uname",
|
||||
"printenv",
|
||||
"stat",
|
||||
"file",
|
||||
"du",
|
||||
"df",
|
||||
"tree",
|
||||
"realpath",
|
||||
"readlink",
|
||||
"dirname",
|
||||
"basename",
|
||||
"cmp",
|
||||
"true",
|
||||
"false",
|
||||
"sleep",
|
||||
"seq",
|
||||
"tty",
|
||||
"groups",
|
||||
"locale",
|
||||
"ps",
|
||||
"top",
|
||||
"free",
|
||||
"uptime",
|
||||
"lsblk",
|
||||
"lscpu",
|
||||
"cut",
|
||||
// Windows cmd / PowerShell read verbs + common aliases
|
||||
"dir",
|
||||
"type",
|
||||
"where",
|
||||
"whereis",
|
||||
"get-childitem",
|
||||
"gci",
|
||||
"get-content",
|
||||
"gc",
|
||||
"get-location",
|
||||
"gl",
|
||||
"select-string",
|
||||
"sls",
|
||||
"measure-object",
|
||||
"get-item",
|
||||
"gi",
|
||||
"test-path",
|
||||
"resolve-path",
|
||||
"get-command",
|
||||
"gcm",
|
||||
"get-process",
|
||||
];
|
||||
|
||||
/// Commands that reach the network. Always-ask in every acting tier.
|
||||
const NETWORK_BASES: &[&str] = &[
|
||||
"curl",
|
||||
"wget",
|
||||
"ssh",
|
||||
"scp",
|
||||
"sftp",
|
||||
"rsync",
|
||||
"nc",
|
||||
"ncat",
|
||||
"netcat",
|
||||
"telnet",
|
||||
"ftp",
|
||||
"tftp",
|
||||
"socat",
|
||||
// Windows / PowerShell
|
||||
"invoke-webrequest",
|
||||
"iwr",
|
||||
"invoke-restmethod",
|
||||
"irm",
|
||||
"start-bitstransfer",
|
||||
"bitsadmin",
|
||||
];
|
||||
|
||||
/// Catastrophic / irreversible / privilege / system-control bases. Always-ask
|
||||
/// in every acting tier (Full included). Coarse on the broad Windows verbs
|
||||
/// (`reg`/`net`/`sc`) — over-prompting there is the safe default.
|
||||
const DESTRUCTIVE_BASES: &[&str] = &[
|
||||
// POSIX privilege / disk / system-control
|
||||
"sudo",
|
||||
"su",
|
||||
"doas",
|
||||
"dd",
|
||||
"mkfs",
|
||||
"fdisk",
|
||||
"sfdisk",
|
||||
"parted",
|
||||
"wipefs",
|
||||
"shred",
|
||||
"shutdown",
|
||||
"reboot",
|
||||
"halt",
|
||||
"poweroff",
|
||||
"init",
|
||||
"telinit",
|
||||
"mount",
|
||||
"umount",
|
||||
"swapoff",
|
||||
"iptables",
|
||||
"ip6tables",
|
||||
"nft",
|
||||
"ufw",
|
||||
"firewall-cmd",
|
||||
"useradd",
|
||||
"userdel",
|
||||
"usermod",
|
||||
"groupadd",
|
||||
"groupdel",
|
||||
"passwd",
|
||||
"chpasswd",
|
||||
"visudo",
|
||||
"modprobe",
|
||||
"insmod",
|
||||
"rmmod",
|
||||
// Windows / PowerShell
|
||||
"format",
|
||||
"diskpart",
|
||||
"bcdedit",
|
||||
"takeown",
|
||||
"cipher",
|
||||
"vssadmin",
|
||||
"reg",
|
||||
"regedit",
|
||||
"runas",
|
||||
"sc",
|
||||
"net",
|
||||
"set-executionpolicy",
|
||||
"stop-computer",
|
||||
"restart-computer",
|
||||
"clear-disk",
|
||||
"format-volume",
|
||||
"remove-partition",
|
||||
"disable-computerrestore",
|
||||
];
|
||||
|
||||
/// Git subcommands that only read repository state. Anything else — including
|
||||
/// `commit`/`push`/`branch`/`config`/unknown/bare `git` — is fail-closed to
|
||||
/// `Write`.
|
||||
const GIT_READ_VERBS: &[&str] = &[
|
||||
"status",
|
||||
"log",
|
||||
"diff",
|
||||
"show",
|
||||
"remote",
|
||||
"describe",
|
||||
"blame",
|
||||
"ls-files",
|
||||
"ls-tree",
|
||||
"rev-parse",
|
||||
"cat-file",
|
||||
"shortlog",
|
||||
"reflog",
|
||||
"rev-list",
|
||||
"name-rev",
|
||||
"var",
|
||||
"check-ignore",
|
||||
"check-attr",
|
||||
"verify-commit",
|
||||
"count-objects",
|
||||
"fsck",
|
||||
"whatchanged",
|
||||
"grep",
|
||||
"version",
|
||||
"help",
|
||||
];
|
||||
|
||||
/// npm/pnpm/yarn read-only subcommands. `install`/`run`/`test`/`exec` (which
|
||||
/// run arbitrary scripts) and unknown verbs are fail-closed to `Write`.
|
||||
const NODE_PKG_READ_VERBS: &[&str] = &[
|
||||
"ls", "list", "view", "info", "outdated", "ping", "whoami", "help", "why", "audit", "doctor",
|
||||
];
|
||||
|
||||
/// cargo read-only subcommands. `build`/`run`/`test`/`check` compile and may
|
||||
/// run build scripts, so they are fail-closed to `Write`.
|
||||
const CARGO_READ_VERBS: &[&str] = &["tree", "metadata", "search", "info", "version", "help"];
|
||||
|
||||
/// Detect a pacman *install/upgrade* from its bundled operation flag.
|
||||
///
|
||||
/// pacman packs its operation and modifiers into a single flag (`-Syu`, `-Ss`),
|
||||
/// and `args` reach us already lowercased — so the `-S` (sync) operation is
|
||||
/// indistinguishable from a literal `-s` by case alone. We therefore key off
|
||||
/// the *modifier* letters instead of a blanket `starts_with("-s")`, which would
|
||||
/// over-match every read-only `-S` query: a `-S`-family flag mutates the host
|
||||
/// only when it carries none of pacman's read-only query modifiers — search
|
||||
/// (`s`), info (`i`), list (`l`), groups (`g`) or print (`p`). So `-S pkg`,
|
||||
/// `-Sy`, `-Syu` are installs while `-Ss`/`-Si`/`-Sl`/`-Sg`/`-Sp` are reads.
|
||||
fn is_pacman_install(args: &[String]) -> bool {
|
||||
args.iter().any(|a| {
|
||||
a.strip_prefix("-s")
|
||||
.is_some_and(|modifiers| !modifiers.contains(['s', 'i', 'l', 'g', 'p']))
|
||||
})
|
||||
}
|
||||
|
||||
/// Detect a package-manager *install* invocation. These mutate the host /
|
||||
/// global environment, so they are the always-ask `Install` bucket (even in
|
||||
/// Full) — the same gate the dedicated `install_tool` enforces, applied to the
|
||||
/// shell escape hatch. Project-local installs (`npm install` without `-g`,
|
||||
/// `cargo add`) are ordinary `Write`s and are deliberately NOT matched here.
|
||||
/// `args` are already lowercased by the caller.
|
||||
fn is_install_command(base: &str, args: &[String]) -> bool {
|
||||
let has = |needle: &str| args.iter().any(|a| a == needle);
|
||||
let first_is = |verb: &str| args.first().map(String::as_str) == Some(verb);
|
||||
match base {
|
||||
// System package managers.
|
||||
"apt" | "apt-get" | "dnf" | "yum" | "zypper" => has("install"),
|
||||
"pacman" => is_pacman_install(args),
|
||||
"apk" => has("add"),
|
||||
"brew" | "snap" | "flatpak" | "winget" | "choco" | "scoop" => has("install"),
|
||||
// Language package managers — host/global-modifying installs only.
|
||||
"pip" | "pip3" | "pipx" | "gem" | "go" | "cargo" => first_is("install"),
|
||||
"npm" | "pnpm" => {
|
||||
(has("install") || has("i") || has("add")) && (has("-g") || has("--global"))
|
||||
}
|
||||
"yarn" => has("global"),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Classify a single already-split shell segment. `base` is the normalized
|
||||
/// (lowercased, `.exe`-stripped, basename-only) program name; `args` are the
|
||||
/// lowercased remaining words; `joined` is the lowercased segment used for
|
||||
/// pattern matching. Fail-closed: an unrecognized base resolves to `Write`.
|
||||
pub(super) fn classify_segment(base: &str, args: &[String], joined: &str) -> CommandClass {
|
||||
// Catastrophic patterns first — they win regardless of the base command.
|
||||
if joined.contains("rm -rf /") || joined.contains("rm -fr /") || joined.contains(":(){:|:&};:")
|
||||
{
|
||||
return CommandClass::Destructive;
|
||||
}
|
||||
if DESTRUCTIVE_BASES.contains(&base) {
|
||||
return CommandClass::Destructive;
|
||||
}
|
||||
if NETWORK_BASES.contains(&base) {
|
||||
return CommandClass::Network;
|
||||
}
|
||||
// Package installs mutate the host → always-ask Install bucket (closes the
|
||||
// shell escape hatch around `install_tool`).
|
||||
if is_install_command(base, args) {
|
||||
return CommandClass::Install;
|
||||
}
|
||||
// Interpreters / code executors run arbitrary code. Fail-closed to Write
|
||||
// (not Destructive) so Full can still run code while Supervised prompts.
|
||||
if is_command_executor(base) {
|
||||
return CommandClass::Write;
|
||||
}
|
||||
// `find` is read-only unless it executes commands or deletes files.
|
||||
if base == "find" {
|
||||
if args.iter().any(|a| {
|
||||
matches!(
|
||||
a.as_str(),
|
||||
"-exec" | "-execdir" | "-ok" | "-okdir" | "-delete"
|
||||
)
|
||||
}) {
|
||||
return CommandClass::Write;
|
||||
}
|
||||
return CommandClass::Read;
|
||||
}
|
||||
// Verb-sensitive VCS / package tools.
|
||||
if base == "git" {
|
||||
return verb_class(args, GIT_READ_VERBS);
|
||||
}
|
||||
if matches!(base, "npm" | "pnpm" | "yarn") {
|
||||
return verb_class(args, NODE_PKG_READ_VERBS);
|
||||
}
|
||||
if base == "cargo" {
|
||||
return verb_class(args, CARGO_READ_VERBS);
|
||||
}
|
||||
if READ_ONLY_BASES.contains(&base) {
|
||||
return CommandClass::Read;
|
||||
}
|
||||
// Fail closed: unknown or known-mutating base → Write.
|
||||
CommandClass::Write
|
||||
}
|
||||
|
||||
/// `Read` when the first subcommand word is in `read_verbs`, else fail-closed
|
||||
/// `Write`. Mirrors the `args.first()` verb check used by `command_risk_level`.
|
||||
fn verb_class(args: &[String], read_verbs: &[&str]) -> CommandClass {
|
||||
match args.first().map(String::as_str) {
|
||||
Some(verb) if read_verbs.contains(&verb) => CommandClass::Read,
|
||||
_ => CommandClass::Write,
|
||||
}
|
||||
}
|
||||
|
||||
/// Structural-safety guard for the harness-gated command flow (Option 2). Even
|
||||
/// after a human approves a command, a hidden subshell / command substitution /
|
||||
/// output redirect / `tee` / background `&` could smuggle a *different* command
|
||||
/// past the approval summary, so these are refused outside Full (which is
|
||||
/// trusted to use redirects and pipes). Mirrors the structural checks in
|
||||
/// [`SecurityPolicy::is_command_allowed`].
|
||||
/// Detect shell structure that can **hide execution** from `classify_command`,
|
||||
/// which only inspects the base command of each `;`/`&&`/`|` segment. Command
|
||||
/// and process substitution and backticks run an *inner* command classification
|
||||
/// can't see (`echo $(rm -rf ~)` classifies as `echo` = Read and would run
|
||||
/// unprompted), and a trailing `&` detaches a process past the gate — so these
|
||||
/// stay hard-blocked outside Full.
|
||||
///
|
||||
/// Deliberately NOT flagged here: plain redirects (`>`, `2>&1`, `2>/dev/null`),
|
||||
/// `tee`, and `${VAR}` expansion. `classify_command` already lifts a redirect /
|
||||
/// `tee` to `Write`, so the gate prompts and — once the human approves — the
|
||||
/// command MUST actually run. Re-blocking an approved `… 2>&1` here was the bug
|
||||
/// that made Supervised mode unusable: every command the agent wrote carried a
|
||||
/// `2>&1`, got approved, then silently failed this in-tool guard and never ran.
|
||||
pub(super) fn has_hidden_execution(command: &str) -> bool {
|
||||
// The backtick check is deliberately NOT quote-aware: any backtick in the
|
||||
// command string is blocked, even inside a double-quoted literal. Over-
|
||||
// blocking is the safe direction here. (By contrast the `&` case below is
|
||||
// quote-aware via `contains_unquoted_background_ampersand`, because that one
|
||||
// must still allow benign fd-dup redirects like `2>&1`.)
|
||||
command.contains('`')
|
||||
|| command.contains("$(")
|
||||
|| command.contains("<(")
|
||||
|| command.contains(">(")
|
||||
|| contains_unquoted_background_ampersand(command)
|
||||
}
|
||||
@@ -891,662 +891,5 @@ pub async fn execute_prepared(
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{extract::State, routing::post, Json, Router};
|
||||
use once_cell::sync::Lazy;
|
||||
use serde_json::{json, Value};
|
||||
use tempfile::TempDir;
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
use super::*;
|
||||
use crate::openhuman::config::rpc as config_rpc;
|
||||
use crate::openhuman::wallet::ops::{setup, WalletSetupParams, WalletSetupSource};
|
||||
|
||||
pub(crate) static TEST_LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
|
||||
|
||||
#[derive(Clone)]
|
||||
struct MockRpcState {
|
||||
estimate_calls: Arc<Mutex<Vec<Value>>>,
|
||||
raw_txs: Arc<Mutex<Vec<String>>>,
|
||||
chain_id: String,
|
||||
}
|
||||
|
||||
fn sample_account(chain: WalletChain) -> super::WalletAccount {
|
||||
super::WalletAccount {
|
||||
chain,
|
||||
address: match chain {
|
||||
WalletChain::Evm => "0x9858EfFD232B4033E47d90003D41EC34EcaEda94".to_string(),
|
||||
WalletChain::Btc => "bc1qcr8te4kr609gcawutmrza0j4xv80jy8z306fyu".to_string(),
|
||||
WalletChain::Solana => "HAgk14JpMQLgt6rVgv7cBQFJWFto5Dqxi472uT3DKpqk".to_string(),
|
||||
WalletChain::Tron => "TUEZSdKsoDHQMeZwihtdoBiN46zxhGWYdH".to_string(),
|
||||
},
|
||||
derivation_path: match chain {
|
||||
WalletChain::Evm => "m/44'/60'/0'/0/0".to_string(),
|
||||
WalletChain::Btc => "m/84'/0'/0'/0/0".to_string(),
|
||||
WalletChain::Solana => "m/44'/501'/0'/0'".to_string(),
|
||||
WalletChain::Tron => "m/44'/195'/0'/0/0".to_string(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
async fn mock_rpc(
|
||||
State(state): State<MockRpcState>,
|
||||
Json(payload): Json<Value>,
|
||||
) -> Json<Value> {
|
||||
let method = payload
|
||||
.get("method")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let params = payload
|
||||
.get("params")
|
||||
.and_then(Value::as_array)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let result = match method {
|
||||
"eth_chainId" => Value::String(state.chain_id.clone()),
|
||||
"eth_getTransactionCount" => Value::String("0x7".to_string()),
|
||||
"eth_gasPrice" => Value::String("0x3b9aca00".to_string()),
|
||||
"eth_estimateGas" => {
|
||||
state
|
||||
.estimate_calls
|
||||
.lock()
|
||||
.push(params.first().cloned().unwrap_or(Value::Null));
|
||||
Value::String("0x5208".to_string())
|
||||
}
|
||||
"eth_sendRawTransaction" => {
|
||||
if let Some(raw) = params.first().and_then(Value::as_str) {
|
||||
state.raw_txs.lock().push(raw.to_string());
|
||||
}
|
||||
Value::String(
|
||||
"0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
|
||||
.to_string(),
|
||||
)
|
||||
}
|
||||
"eth_getBalance" => Value::String("0xde0b6b3a7640000".to_string()),
|
||||
_ => Value::Null,
|
||||
};
|
||||
Json(json!({"jsonrpc":"2.0","id":1,"result":result}))
|
||||
}
|
||||
|
||||
pub(crate) async fn start_mock_rpc_with_chain_id(
|
||||
chain_id_hex: &str,
|
||||
) -> Result<(SocketAddr, Arc<Mutex<Vec<Value>>>, Arc<Mutex<Vec<String>>>), String> {
|
||||
let estimate_calls = Arc::new(Mutex::new(Vec::new()));
|
||||
let raw_txs = Arc::new(Mutex::new(Vec::new()));
|
||||
let state = MockRpcState {
|
||||
estimate_calls: estimate_calls.clone(),
|
||||
raw_txs: raw_txs.clone(),
|
||||
chain_id: chain_id_hex.to_string(),
|
||||
};
|
||||
let app = Router::new().route("/", post(mock_rpc)).with_state(state);
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.map_err(|e| format!("failed to bind mock rpc: {e}"))?;
|
||||
let addr = listener
|
||||
.local_addr()
|
||||
.map_err(|e| format!("failed to read mock rpc addr: {e}"))?;
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
Ok((addr, estimate_calls, raw_txs))
|
||||
}
|
||||
|
||||
async fn start_mock_rpc(
|
||||
) -> Result<(SocketAddr, Arc<Mutex<Vec<Value>>>, Arc<Mutex<Vec<String>>>), String> {
|
||||
start_mock_rpc_with_chain_id("0x1").await
|
||||
}
|
||||
|
||||
pub(crate) async fn setup_wallet_in(temp: &TempDir) -> Result<(), String> {
|
||||
std::env::set_var("OPENHUMAN_WORKSPACE", temp.path());
|
||||
let config = config_rpc::load_config_with_timeout().await?;
|
||||
let encrypted = crate::openhuman::encryption::rpc::encrypt_secret(
|
||||
&config,
|
||||
"abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about",
|
||||
)
|
||||
.await?
|
||||
.value;
|
||||
setup(WalletSetupParams {
|
||||
consent_granted: true,
|
||||
source: WalletSetupSource::Imported,
|
||||
mnemonic_word_count: 12,
|
||||
encrypted_mnemonic: Some(encrypted),
|
||||
accounts: [
|
||||
WalletChain::Evm,
|
||||
WalletChain::Btc,
|
||||
WalletChain::Solana,
|
||||
WalletChain::Tron,
|
||||
]
|
||||
.into_iter()
|
||||
.map(sample_account)
|
||||
.collect(),
|
||||
})
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_amount_rejects_empty_and_non_numeric() {
|
||||
assert!(validate_amount("").is_err());
|
||||
assert!(validate_amount("abc").is_err());
|
||||
assert_eq!(validate_amount("42").unwrap(), 42);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_calldata_requires_hex() {
|
||||
assert!(validate_calldata("deadbeef").is_err());
|
||||
assert!(validate_calldata("0xZZ").is_err());
|
||||
assert!(validate_calldata("0xabc").is_err());
|
||||
assert_eq!(validate_calldata("0xdeadbeef").unwrap(), "0xdeadbeef");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn formats_amount_with_decimals() {
|
||||
assert_eq!(format_amount(0, 18), "0.000000000000000000");
|
||||
assert_eq!(format_amount(1, 8), "0.00000001");
|
||||
assert_eq!(format_amount(123_456_789, 8), "1.23456789");
|
||||
assert_eq!(format_amount(100, 0), "100");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn next_quote_id_is_unique_and_prefixed() {
|
||||
let a = next_quote_id();
|
||||
let b = next_quote_id();
|
||||
assert_ne!(a, b);
|
||||
assert!(a.starts_with("q_"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quote_store_round_trips_and_expires() {
|
||||
let _guard = TEST_LOCK.lock();
|
||||
reset_quote_store_for_tests();
|
||||
let now = now_ms();
|
||||
let mut q = PreparedTransaction {
|
||||
quote_id: "q_test_1".to_string(),
|
||||
kind: PreparedKind::NativeTransfer,
|
||||
chain: WalletChain::Evm,
|
||||
evm_network: Some(EvmNetwork::EthereumMainnet),
|
||||
from_address: "0xfrom".to_string(),
|
||||
to_address: "0xto".to_string(),
|
||||
asset_symbol: "ETH".to_string(),
|
||||
amount_raw: "1".to_string(),
|
||||
amount_formatted: "0.000000000000000001".to_string(),
|
||||
receive_symbol: None,
|
||||
min_receive_raw: None,
|
||||
calldata: None,
|
||||
token_address: None,
|
||||
estimated_fee_raw: "0".to_string(),
|
||||
status: PreparedStatus::AwaitingConfirmation,
|
||||
created_at_ms: now,
|
||||
expires_at_ms: now + 60_000,
|
||||
notes: vec![],
|
||||
owner: None,
|
||||
};
|
||||
store_quote(q.clone());
|
||||
let taken = take_quote_for("q_test_1", None).expect("quote round-trips");
|
||||
assert_eq!(taken.quote_id, "q_test_1");
|
||||
assert!(
|
||||
take_quote_for("q_test_1", None).is_err(),
|
||||
"second take must fail"
|
||||
);
|
||||
|
||||
q.quote_id = "q_test_2".to_string();
|
||||
q.expires_at_ms = now.saturating_sub(1);
|
||||
store_quote(q);
|
||||
let err = take_quote_for("q_test_2", None).unwrap_err();
|
||||
assert!(err.contains("expired"), "got: {err}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_requires_confirmed_flag() {
|
||||
let err = execute_prepared(ExecutePreparedParams {
|
||||
quote_id: "missing".to_string(),
|
||||
confirmed: false,
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(err.contains("confirmed: true"), "got: {err}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn supported_assets_lists_default_erc20s_and_l2() {
|
||||
let out = supported_assets().await.unwrap();
|
||||
assert!(out
|
||||
.value
|
||||
.iter()
|
||||
.any(|asset| asset.symbol == "USDC"
|
||||
&& asset.evm_network == Some(EvmNetwork::BaseMainnet)));
|
||||
assert!(out
|
||||
.value
|
||||
.iter()
|
||||
.any(|asset| asset.symbol == "ETH" && asset.native));
|
||||
assert!(out
|
||||
.value
|
||||
.iter()
|
||||
.any(|asset| asset.symbol == "USDT" && asset.chain == WalletChain::Tron));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prepare_transfer_rejects_unknown_asset_symbol() {
|
||||
let _guard = TEST_LOCK.lock();
|
||||
let _env_guard = crate::openhuman::config::TEST_ENV_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
reset_quote_store_for_tests();
|
||||
let temp = TempDir::new().unwrap();
|
||||
setup_wallet_in(&temp).await.unwrap();
|
||||
let err = prepare_transfer(PrepareTransferParams {
|
||||
chain: WalletChain::Evm,
|
||||
to_address: "0x1111111111111111111111111111111111111111".into(),
|
||||
amount_raw: "1".into(),
|
||||
asset_symbol: Some("NOPE".into()),
|
||||
evm_network: None,
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(err.contains("unsupported asset_symbol"), "got: {err}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prepare_contract_call_rejects_non_evm_chain() {
|
||||
let err = prepare_contract_call(PrepareContractCallParams {
|
||||
chain: WalletChain::Btc,
|
||||
contract_address: "addr".into(),
|
||||
calldata: "0x".into(),
|
||||
value_raw: "0".into(),
|
||||
evm_network: None,
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(err.contains("only for EVM"), "got: {err}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_broadcasts_native_evm_transaction() {
|
||||
let _guard = TEST_LOCK.lock();
|
||||
let _env_guard = crate::openhuman::config::TEST_ENV_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
reset_quote_store_for_tests();
|
||||
let temp = TempDir::new().unwrap();
|
||||
setup_wallet_in(&temp).await.unwrap();
|
||||
let (addr, estimate_calls, raw_txs) = start_mock_rpc().await.unwrap();
|
||||
std::env::set_var("OPENHUMAN_WALLET_RPC_EVM", format!("http://{addr}"));
|
||||
|
||||
let prepared = prepare_transfer(PrepareTransferParams {
|
||||
chain: WalletChain::Evm,
|
||||
to_address: "0x1111111111111111111111111111111111111111".into(),
|
||||
amount_raw: "1000".into(),
|
||||
asset_symbol: None,
|
||||
evm_network: None,
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.value;
|
||||
let executed = execute_prepared(ExecutePreparedParams {
|
||||
quote_id: prepared.quote_id.clone(),
|
||||
confirmed: true,
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.value;
|
||||
|
||||
assert_eq!(executed.status, PreparedStatus::Broadcasted);
|
||||
assert!(executed.transaction_hash.starts_with("0xaaaa"));
|
||||
assert_eq!(raw_txs.lock().len(), 1);
|
||||
let estimate = estimate_calls.lock()[0].clone();
|
||||
assert_eq!(
|
||||
estimate.get("to").and_then(Value::as_str),
|
||||
Some("0x1111111111111111111111111111111111111111")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_broadcasts_erc20_transfer_using_default_token_catalog() {
|
||||
let _guard = TEST_LOCK.lock();
|
||||
let _env_guard = crate::openhuman::config::TEST_ENV_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
reset_quote_store_for_tests();
|
||||
let temp = TempDir::new().unwrap();
|
||||
setup_wallet_in(&temp).await.unwrap();
|
||||
let (addr, estimate_calls, raw_txs) = start_mock_rpc().await.unwrap();
|
||||
std::env::set_var("OPENHUMAN_WALLET_RPC_EVM", format!("http://{addr}"));
|
||||
|
||||
let prepared = prepare_transfer(PrepareTransferParams {
|
||||
chain: WalletChain::Evm,
|
||||
to_address: "0x1111111111111111111111111111111111111111".into(),
|
||||
amount_raw: "5000000".into(),
|
||||
asset_symbol: Some("USDC".into()),
|
||||
evm_network: None,
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.value;
|
||||
let executed = execute_prepared(ExecutePreparedParams {
|
||||
quote_id: prepared.quote_id.clone(),
|
||||
confirmed: true,
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.value;
|
||||
|
||||
assert_eq!(executed.status, PreparedStatus::Broadcasted);
|
||||
assert_eq!(raw_txs.lock().len(), 1);
|
||||
let estimate = estimate_calls.lock()[0].clone();
|
||||
assert_eq!(
|
||||
estimate.get("to").and_then(Value::as_str),
|
||||
Some("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48")
|
||||
);
|
||||
let data = estimate
|
||||
.get("data")
|
||||
.and_then(Value::as_str)
|
||||
.expect("token transfer calldata");
|
||||
assert!(data.starts_with("0xa9059cbb"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_broadcasts_native_evm_on_base_with_chain_id_8453() {
|
||||
let _guard = TEST_LOCK.lock();
|
||||
let _env_guard = crate::openhuman::config::TEST_ENV_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
reset_quote_store_for_tests();
|
||||
let temp = TempDir::new().unwrap();
|
||||
setup_wallet_in(&temp).await.unwrap();
|
||||
// Base uses chain_id 8453 = 0x2105.
|
||||
let (addr, _estimate_calls, raw_txs) =
|
||||
start_mock_rpc_with_chain_id("0x2105").await.unwrap();
|
||||
std::env::set_var("OPENHUMAN_WALLET_RPC_BASE", format!("http://{addr}"));
|
||||
|
||||
let prepared = prepare_transfer(PrepareTransferParams {
|
||||
chain: WalletChain::Evm,
|
||||
to_address: "0x1111111111111111111111111111111111111111".into(),
|
||||
amount_raw: "1000".into(),
|
||||
asset_symbol: None,
|
||||
evm_network: Some(EvmNetwork::BaseMainnet),
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.value;
|
||||
let executed = execute_prepared(ExecutePreparedParams {
|
||||
quote_id: prepared.quote_id.clone(),
|
||||
confirmed: true,
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.value;
|
||||
assert_eq!(executed.status, PreparedStatus::Broadcasted);
|
||||
assert_eq!(executed.evm_network, Some(EvmNetwork::BaseMainnet));
|
||||
assert_eq!(raw_txs.lock().len(), 1);
|
||||
}
|
||||
|
||||
/// Build a Quote-store fixture pinned to a specific owner.
|
||||
/// Bypasses prepare_* (which would need full wallet setup + mock RPC) so
|
||||
/// the owner-gate behavior can be exercised in isolation.
|
||||
fn insert_owned_quote(quote_id: &str, owner: Option<QuoteOwner>) -> PreparedTransaction {
|
||||
let now = now_ms();
|
||||
let q = PreparedTransaction {
|
||||
quote_id: quote_id.to_string(),
|
||||
kind: PreparedKind::NativeTransfer,
|
||||
chain: WalletChain::Evm,
|
||||
evm_network: Some(EvmNetwork::EthereumMainnet),
|
||||
from_address: "0xfrom".to_string(),
|
||||
to_address: "0xto".to_string(),
|
||||
asset_symbol: "ETH".to_string(),
|
||||
amount_raw: "1".to_string(),
|
||||
amount_formatted: "0.000000000000000001".to_string(),
|
||||
receive_symbol: None,
|
||||
min_receive_raw: None,
|
||||
calldata: None,
|
||||
token_address: None,
|
||||
estimated_fee_raw: "0".to_string(),
|
||||
status: PreparedStatus::AwaitingConfirmation,
|
||||
created_at_ms: now,
|
||||
expires_at_ms: now + 60_000,
|
||||
notes: vec![],
|
||||
owner,
|
||||
};
|
||||
insert_quote_for_test(q)
|
||||
}
|
||||
|
||||
fn owner_a() -> QuoteOwner {
|
||||
QuoteOwner {
|
||||
thread_id: "thread-A".into(),
|
||||
client_id: "client-A".into(),
|
||||
}
|
||||
}
|
||||
|
||||
fn owner_b() -> QuoteOwner {
|
||||
QuoteOwner {
|
||||
thread_id: "thread-B".into(),
|
||||
client_id: "client-B".into(),
|
||||
}
|
||||
}
|
||||
|
||||
fn chat_ctx_from(owner: &QuoteOwner) -> crate::openhuman::approval::ApprovalChatContext {
|
||||
crate::openhuman::approval::ApprovalChatContext {
|
||||
thread_id: owner.thread_id.clone(),
|
||||
client_id: owner.client_id.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Cross-thread execute must fail. The quote must remain in the store
|
||||
/// (mismatched caller cannot poison it by consuming on Alice's behalf).
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_rejects_cross_owner_execution() {
|
||||
use crate::openhuman::approval::APPROVAL_CHAT_CONTEXT;
|
||||
let _guard = TEST_LOCK.lock();
|
||||
reset_quote_store_for_tests();
|
||||
let q = insert_owned_quote("q_xowner_1", Some(owner_a()));
|
||||
|
||||
// Bob (owner_b) tries to execute Alice's quote. The error string is
|
||||
// intentionally identical to a not-found lookup.
|
||||
let err = APPROVAL_CHAT_CONTEXT
|
||||
.scope(
|
||||
chat_ctx_from(&owner_b()),
|
||||
execute_prepared(ExecutePreparedParams {
|
||||
quote_id: q.quote_id.clone(),
|
||||
confirmed: true,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(err, format!("quote '{}' not found", q.quote_id));
|
||||
|
||||
// The quote must still be present and executable by Alice — Bob's
|
||||
// failed attempt didn't poison the store.
|
||||
let still_present = prepared_quotes_for_test();
|
||||
assert!(
|
||||
still_present.iter().any(|p| p.quote_id == q.quote_id),
|
||||
"quote must survive owner-mismatch attempts"
|
||||
);
|
||||
}
|
||||
|
||||
/// Same-thread execute must pass the owner gate (it will fail later in
|
||||
/// the chain code because there's no mock RPC set up, but the failure
|
||||
/// must not be the "not found" oracle — that proves we got past the gate).
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_allows_same_owner_execution() {
|
||||
use crate::openhuman::approval::APPROVAL_CHAT_CONTEXT;
|
||||
let _guard = TEST_LOCK.lock();
|
||||
reset_quote_store_for_tests();
|
||||
let q = insert_owned_quote("q_same_owner_1", Some(owner_a()));
|
||||
|
||||
let result = APPROVAL_CHAT_CONTEXT
|
||||
.scope(
|
||||
chat_ctx_from(&owner_a()),
|
||||
execute_prepared(ExecutePreparedParams {
|
||||
quote_id: q.quote_id.clone(),
|
||||
confirmed: true,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
// Past the owner gate. Chain code may error (no mock RPC, no wallet
|
||||
// setup) — but it must NOT be the "not found" shape we use for the
|
||||
// owner-mismatch oracle.
|
||||
if let Err(err) = &result {
|
||||
assert_ne!(
|
||||
err,
|
||||
&format!("quote '{}' not found", q.quote_id),
|
||||
"same-owner path must not return the owner-mismatch oracle"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// No-context prepare + no-context execute must work. Keeps CLI / direct
|
||||
/// JSON-RPC flows usable.
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_allows_no_context_flows() {
|
||||
let _guard = TEST_LOCK.lock();
|
||||
reset_quote_store_for_tests();
|
||||
let q = insert_owned_quote("q_no_ctx_1", None);
|
||||
|
||||
// No APPROVAL_CHAT_CONTEXT scope — current_owner() returns None on
|
||||
// both prepare and execute, so the gate passes.
|
||||
let result = execute_prepared(ExecutePreparedParams {
|
||||
quote_id: q.quote_id.clone(),
|
||||
confirmed: true,
|
||||
})
|
||||
.await;
|
||||
if let Err(err) = &result {
|
||||
assert_ne!(
|
||||
err,
|
||||
&format!("quote '{}' not found", q.quote_id),
|
||||
"no-context path must not return the owner-mismatch oracle"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// A quote prepared inside a chat context must NOT be executable from a
|
||||
/// caller with no context. Prevents privilege-drop into background /
|
||||
/// triage / cron flows that wouldn't surface UI confirmation.
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_rejects_chat_quote_from_no_context_caller() {
|
||||
let _guard = TEST_LOCK.lock();
|
||||
reset_quote_store_for_tests();
|
||||
let q = insert_owned_quote("q_chat_to_bg_1", Some(owner_a()));
|
||||
|
||||
// No scope around execute → caller_owner = None ≠ Some(owner_a).
|
||||
let err = execute_prepared(ExecutePreparedParams {
|
||||
quote_id: q.quote_id.clone(),
|
||||
confirmed: true,
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(err, format!("quote '{}' not found", q.quote_id));
|
||||
}
|
||||
|
||||
/// Lock the error-shape invariant: cross-owner reject string MUST be
|
||||
/// byte-equal to the not-found string. Regressions here would re-open
|
||||
/// the enumeration-oracle gap.
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_owner_mismatch_error_matches_not_found_shape() {
|
||||
use crate::openhuman::approval::APPROVAL_CHAT_CONTEXT;
|
||||
let _guard = TEST_LOCK.lock();
|
||||
reset_quote_store_for_tests();
|
||||
|
||||
// Real (owner-mismatched) quote.
|
||||
let real = insert_owned_quote("q_real_1", Some(owner_a()));
|
||||
// Reach the mismatched-owner branch.
|
||||
let mismatch_err = APPROVAL_CHAT_CONTEXT
|
||||
.scope(
|
||||
chat_ctx_from(&owner_b()),
|
||||
execute_prepared(ExecutePreparedParams {
|
||||
quote_id: real.quote_id.clone(),
|
||||
confirmed: true,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
// Reach the genuine not-found branch (no quote with this id in store).
|
||||
let missing_err = APPROVAL_CHAT_CONTEXT
|
||||
.scope(
|
||||
chat_ctx_from(&owner_b()),
|
||||
execute_prepared(ExecutePreparedParams {
|
||||
quote_id: "q_does_not_exist".into(),
|
||||
confirmed: true,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
// Both branches surface the exact same template, parameterised only
|
||||
// by quote_id. No enumeration oracle.
|
||||
assert_eq!(mismatch_err, format!("quote '{}' not found", real.quote_id));
|
||||
assert_eq!(missing_err, "quote 'q_does_not_exist' not found");
|
||||
}
|
||||
|
||||
/// Verify that `prepare_transfer` inside an `APPROVAL_CHAT_CONTEXT` scope
|
||||
/// actually stamps `owner` via the task-local — not just via test helpers.
|
||||
#[tokio::test]
|
||||
async fn prepare_stamps_owner_via_task_local() {
|
||||
use crate::openhuman::approval::APPROVAL_CHAT_CONTEXT;
|
||||
let _guard = TEST_LOCK.lock();
|
||||
let _env_guard = crate::openhuman::config::TEST_ENV_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
reset_quote_store_for_tests();
|
||||
let temp = TempDir::new().unwrap();
|
||||
setup_wallet_in(&temp).await.unwrap();
|
||||
|
||||
let expected = owner_a();
|
||||
let ctx = chat_ctx_from(&expected);
|
||||
|
||||
let prepared = APPROVAL_CHAT_CONTEXT
|
||||
.scope(
|
||||
ctx,
|
||||
prepare_transfer(PrepareTransferParams {
|
||||
chain: WalletChain::Evm,
|
||||
to_address: "0x1111111111111111111111111111111111111111".into(),
|
||||
amount_raw: "1000".into(),
|
||||
asset_symbol: None,
|
||||
evm_network: Some(EvmNetwork::EthereumMainnet),
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.value;
|
||||
|
||||
assert_eq!(
|
||||
prepared.owner,
|
||||
Some(expected),
|
||||
"prepare_transfer must stamp owner from APPROVAL_CHAT_CONTEXT"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_rejects_evm_chain_id_mismatch() {
|
||||
let _guard = TEST_LOCK.lock();
|
||||
let _env_guard = crate::openhuman::config::TEST_ENV_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
reset_quote_store_for_tests();
|
||||
let temp = TempDir::new().unwrap();
|
||||
setup_wallet_in(&temp).await.unwrap();
|
||||
// Quote says Base; mock reports Ethereum (0x1) — must fail.
|
||||
let (addr, _e, _r) = start_mock_rpc_with_chain_id("0x1").await.unwrap();
|
||||
std::env::set_var("OPENHUMAN_WALLET_RPC_BASE", format!("http://{addr}"));
|
||||
|
||||
let prepared = prepare_transfer(PrepareTransferParams {
|
||||
chain: WalletChain::Evm,
|
||||
to_address: "0x1111111111111111111111111111111111111111".into(),
|
||||
amount_raw: "1000".into(),
|
||||
asset_symbol: None,
|
||||
evm_network: Some(EvmNetwork::BaseMainnet),
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.value;
|
||||
let err = execute_prepared(ExecutePreparedParams {
|
||||
quote_id: prepared.quote_id.clone(),
|
||||
confirmed: true,
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(err.contains("chain_id mismatch"), "got: {err}");
|
||||
}
|
||||
}
|
||||
#[path = "execution_tests.rs"]
|
||||
mod tests;
|
||||
|
||||
@@ -0,0 +1,651 @@
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{extract::State, routing::post, Json, Router};
|
||||
use once_cell::sync::Lazy;
|
||||
use serde_json::{json, Value};
|
||||
use tempfile::TempDir;
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
use super::*;
|
||||
use crate::openhuman::config::rpc as config_rpc;
|
||||
use crate::openhuman::wallet::ops::{setup, WalletSetupParams, WalletSetupSource};
|
||||
|
||||
pub(crate) static TEST_LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
|
||||
|
||||
#[derive(Clone)]
|
||||
struct MockRpcState {
|
||||
estimate_calls: Arc<Mutex<Vec<Value>>>,
|
||||
raw_txs: Arc<Mutex<Vec<String>>>,
|
||||
chain_id: String,
|
||||
}
|
||||
|
||||
fn sample_account(chain: WalletChain) -> super::WalletAccount {
|
||||
super::WalletAccount {
|
||||
chain,
|
||||
address: match chain {
|
||||
WalletChain::Evm => "0x9858EfFD232B4033E47d90003D41EC34EcaEda94".to_string(),
|
||||
WalletChain::Btc => "bc1qcr8te4kr609gcawutmrza0j4xv80jy8z306fyu".to_string(),
|
||||
WalletChain::Solana => "HAgk14JpMQLgt6rVgv7cBQFJWFto5Dqxi472uT3DKpqk".to_string(),
|
||||
WalletChain::Tron => "TUEZSdKsoDHQMeZwihtdoBiN46zxhGWYdH".to_string(),
|
||||
},
|
||||
derivation_path: match chain {
|
||||
WalletChain::Evm => "m/44'/60'/0'/0/0".to_string(),
|
||||
WalletChain::Btc => "m/84'/0'/0'/0/0".to_string(),
|
||||
WalletChain::Solana => "m/44'/501'/0'/0'".to_string(),
|
||||
WalletChain::Tron => "m/44'/195'/0'/0/0".to_string(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
async fn mock_rpc(State(state): State<MockRpcState>, Json(payload): Json<Value>) -> Json<Value> {
|
||||
let method = payload
|
||||
.get("method")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let params = payload
|
||||
.get("params")
|
||||
.and_then(Value::as_array)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let result = match method {
|
||||
"eth_chainId" => Value::String(state.chain_id.clone()),
|
||||
"eth_getTransactionCount" => Value::String("0x7".to_string()),
|
||||
"eth_gasPrice" => Value::String("0x3b9aca00".to_string()),
|
||||
"eth_estimateGas" => {
|
||||
state
|
||||
.estimate_calls
|
||||
.lock()
|
||||
.push(params.first().cloned().unwrap_or(Value::Null));
|
||||
Value::String("0x5208".to_string())
|
||||
}
|
||||
"eth_sendRawTransaction" => {
|
||||
if let Some(raw) = params.first().and_then(Value::as_str) {
|
||||
state.raw_txs.lock().push(raw.to_string());
|
||||
}
|
||||
Value::String(
|
||||
"0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".to_string(),
|
||||
)
|
||||
}
|
||||
"eth_getBalance" => Value::String("0xde0b6b3a7640000".to_string()),
|
||||
_ => Value::Null,
|
||||
};
|
||||
Json(json!({"jsonrpc":"2.0","id":1,"result":result}))
|
||||
}
|
||||
|
||||
pub(crate) async fn start_mock_rpc_with_chain_id(
|
||||
chain_id_hex: &str,
|
||||
) -> Result<(SocketAddr, Arc<Mutex<Vec<Value>>>, Arc<Mutex<Vec<String>>>), String> {
|
||||
let estimate_calls = Arc::new(Mutex::new(Vec::new()));
|
||||
let raw_txs = Arc::new(Mutex::new(Vec::new()));
|
||||
let state = MockRpcState {
|
||||
estimate_calls: estimate_calls.clone(),
|
||||
raw_txs: raw_txs.clone(),
|
||||
chain_id: chain_id_hex.to_string(),
|
||||
};
|
||||
let app = Router::new().route("/", post(mock_rpc)).with_state(state);
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.map_err(|e| format!("failed to bind mock rpc: {e}"))?;
|
||||
let addr = listener
|
||||
.local_addr()
|
||||
.map_err(|e| format!("failed to read mock rpc addr: {e}"))?;
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
Ok((addr, estimate_calls, raw_txs))
|
||||
}
|
||||
|
||||
async fn start_mock_rpc(
|
||||
) -> Result<(SocketAddr, Arc<Mutex<Vec<Value>>>, Arc<Mutex<Vec<String>>>), String> {
|
||||
start_mock_rpc_with_chain_id("0x1").await
|
||||
}
|
||||
|
||||
pub(crate) async fn setup_wallet_in(temp: &TempDir) -> Result<(), String> {
|
||||
std::env::set_var("OPENHUMAN_WORKSPACE", temp.path());
|
||||
let config = config_rpc::load_config_with_timeout().await?;
|
||||
let encrypted = crate::openhuman::encryption::rpc::encrypt_secret(
|
||||
&config,
|
||||
"abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about",
|
||||
)
|
||||
.await?
|
||||
.value;
|
||||
setup(WalletSetupParams {
|
||||
consent_granted: true,
|
||||
source: WalletSetupSource::Imported,
|
||||
mnemonic_word_count: 12,
|
||||
encrypted_mnemonic: Some(encrypted),
|
||||
accounts: [
|
||||
WalletChain::Evm,
|
||||
WalletChain::Btc,
|
||||
WalletChain::Solana,
|
||||
WalletChain::Tron,
|
||||
]
|
||||
.into_iter()
|
||||
.map(sample_account)
|
||||
.collect(),
|
||||
})
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_amount_rejects_empty_and_non_numeric() {
|
||||
assert!(validate_amount("").is_err());
|
||||
assert!(validate_amount("abc").is_err());
|
||||
assert_eq!(validate_amount("42").unwrap(), 42);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_calldata_requires_hex() {
|
||||
assert!(validate_calldata("deadbeef").is_err());
|
||||
assert!(validate_calldata("0xZZ").is_err());
|
||||
assert!(validate_calldata("0xabc").is_err());
|
||||
assert_eq!(validate_calldata("0xdeadbeef").unwrap(), "0xdeadbeef");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn formats_amount_with_decimals() {
|
||||
assert_eq!(format_amount(0, 18), "0.000000000000000000");
|
||||
assert_eq!(format_amount(1, 8), "0.00000001");
|
||||
assert_eq!(format_amount(123_456_789, 8), "1.23456789");
|
||||
assert_eq!(format_amount(100, 0), "100");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn next_quote_id_is_unique_and_prefixed() {
|
||||
let a = next_quote_id();
|
||||
let b = next_quote_id();
|
||||
assert_ne!(a, b);
|
||||
assert!(a.starts_with("q_"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quote_store_round_trips_and_expires() {
|
||||
let _guard = TEST_LOCK.lock();
|
||||
reset_quote_store_for_tests();
|
||||
let now = now_ms();
|
||||
let mut q = PreparedTransaction {
|
||||
quote_id: "q_test_1".to_string(),
|
||||
kind: PreparedKind::NativeTransfer,
|
||||
chain: WalletChain::Evm,
|
||||
evm_network: Some(EvmNetwork::EthereumMainnet),
|
||||
from_address: "0xfrom".to_string(),
|
||||
to_address: "0xto".to_string(),
|
||||
asset_symbol: "ETH".to_string(),
|
||||
amount_raw: "1".to_string(),
|
||||
amount_formatted: "0.000000000000000001".to_string(),
|
||||
receive_symbol: None,
|
||||
min_receive_raw: None,
|
||||
calldata: None,
|
||||
token_address: None,
|
||||
estimated_fee_raw: "0".to_string(),
|
||||
status: PreparedStatus::AwaitingConfirmation,
|
||||
created_at_ms: now,
|
||||
expires_at_ms: now + 60_000,
|
||||
notes: vec![],
|
||||
owner: None,
|
||||
};
|
||||
store_quote(q.clone());
|
||||
let taken = take_quote_for("q_test_1", None).expect("quote round-trips");
|
||||
assert_eq!(taken.quote_id, "q_test_1");
|
||||
assert!(
|
||||
take_quote_for("q_test_1", None).is_err(),
|
||||
"second take must fail"
|
||||
);
|
||||
|
||||
q.quote_id = "q_test_2".to_string();
|
||||
q.expires_at_ms = now.saturating_sub(1);
|
||||
store_quote(q);
|
||||
let err = take_quote_for("q_test_2", None).unwrap_err();
|
||||
assert!(err.contains("expired"), "got: {err}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_requires_confirmed_flag() {
|
||||
let err = execute_prepared(ExecutePreparedParams {
|
||||
quote_id: "missing".to_string(),
|
||||
confirmed: false,
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(err.contains("confirmed: true"), "got: {err}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn supported_assets_lists_default_erc20s_and_l2() {
|
||||
let out = supported_assets().await.unwrap();
|
||||
assert!(out
|
||||
.value
|
||||
.iter()
|
||||
.any(|asset| asset.symbol == "USDC" && asset.evm_network == Some(EvmNetwork::BaseMainnet)));
|
||||
assert!(out
|
||||
.value
|
||||
.iter()
|
||||
.any(|asset| asset.symbol == "ETH" && asset.native));
|
||||
assert!(out
|
||||
.value
|
||||
.iter()
|
||||
.any(|asset| asset.symbol == "USDT" && asset.chain == WalletChain::Tron));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prepare_transfer_rejects_unknown_asset_symbol() {
|
||||
let _guard = TEST_LOCK.lock();
|
||||
let _env_guard = crate::openhuman::config::TEST_ENV_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
reset_quote_store_for_tests();
|
||||
let temp = TempDir::new().unwrap();
|
||||
setup_wallet_in(&temp).await.unwrap();
|
||||
let err = prepare_transfer(PrepareTransferParams {
|
||||
chain: WalletChain::Evm,
|
||||
to_address: "0x1111111111111111111111111111111111111111".into(),
|
||||
amount_raw: "1".into(),
|
||||
asset_symbol: Some("NOPE".into()),
|
||||
evm_network: None,
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(err.contains("unsupported asset_symbol"), "got: {err}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn prepare_contract_call_rejects_non_evm_chain() {
|
||||
let err = prepare_contract_call(PrepareContractCallParams {
|
||||
chain: WalletChain::Btc,
|
||||
contract_address: "addr".into(),
|
||||
calldata: "0x".into(),
|
||||
value_raw: "0".into(),
|
||||
evm_network: None,
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(err.contains("only for EVM"), "got: {err}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_broadcasts_native_evm_transaction() {
|
||||
let _guard = TEST_LOCK.lock();
|
||||
let _env_guard = crate::openhuman::config::TEST_ENV_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
reset_quote_store_for_tests();
|
||||
let temp = TempDir::new().unwrap();
|
||||
setup_wallet_in(&temp).await.unwrap();
|
||||
let (addr, estimate_calls, raw_txs) = start_mock_rpc().await.unwrap();
|
||||
std::env::set_var("OPENHUMAN_WALLET_RPC_EVM", format!("http://{addr}"));
|
||||
|
||||
let prepared = prepare_transfer(PrepareTransferParams {
|
||||
chain: WalletChain::Evm,
|
||||
to_address: "0x1111111111111111111111111111111111111111".into(),
|
||||
amount_raw: "1000".into(),
|
||||
asset_symbol: None,
|
||||
evm_network: None,
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.value;
|
||||
let executed = execute_prepared(ExecutePreparedParams {
|
||||
quote_id: prepared.quote_id.clone(),
|
||||
confirmed: true,
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.value;
|
||||
|
||||
assert_eq!(executed.status, PreparedStatus::Broadcasted);
|
||||
assert!(executed.transaction_hash.starts_with("0xaaaa"));
|
||||
assert_eq!(raw_txs.lock().len(), 1);
|
||||
let estimate = estimate_calls.lock()[0].clone();
|
||||
assert_eq!(
|
||||
estimate.get("to").and_then(Value::as_str),
|
||||
Some("0x1111111111111111111111111111111111111111")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_broadcasts_erc20_transfer_using_default_token_catalog() {
|
||||
let _guard = TEST_LOCK.lock();
|
||||
let _env_guard = crate::openhuman::config::TEST_ENV_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
reset_quote_store_for_tests();
|
||||
let temp = TempDir::new().unwrap();
|
||||
setup_wallet_in(&temp).await.unwrap();
|
||||
let (addr, estimate_calls, raw_txs) = start_mock_rpc().await.unwrap();
|
||||
std::env::set_var("OPENHUMAN_WALLET_RPC_EVM", format!("http://{addr}"));
|
||||
|
||||
let prepared = prepare_transfer(PrepareTransferParams {
|
||||
chain: WalletChain::Evm,
|
||||
to_address: "0x1111111111111111111111111111111111111111".into(),
|
||||
amount_raw: "5000000".into(),
|
||||
asset_symbol: Some("USDC".into()),
|
||||
evm_network: None,
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.value;
|
||||
let executed = execute_prepared(ExecutePreparedParams {
|
||||
quote_id: prepared.quote_id.clone(),
|
||||
confirmed: true,
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.value;
|
||||
|
||||
assert_eq!(executed.status, PreparedStatus::Broadcasted);
|
||||
assert_eq!(raw_txs.lock().len(), 1);
|
||||
let estimate = estimate_calls.lock()[0].clone();
|
||||
assert_eq!(
|
||||
estimate.get("to").and_then(Value::as_str),
|
||||
Some("0xa0b86991c6218b36c1d19d4a2e9eb0ce3606eb48")
|
||||
);
|
||||
let data = estimate
|
||||
.get("data")
|
||||
.and_then(Value::as_str)
|
||||
.expect("token transfer calldata");
|
||||
assert!(data.starts_with("0xa9059cbb"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_broadcasts_native_evm_on_base_with_chain_id_8453() {
|
||||
let _guard = TEST_LOCK.lock();
|
||||
let _env_guard = crate::openhuman::config::TEST_ENV_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
reset_quote_store_for_tests();
|
||||
let temp = TempDir::new().unwrap();
|
||||
setup_wallet_in(&temp).await.unwrap();
|
||||
// Base uses chain_id 8453 = 0x2105.
|
||||
let (addr, _estimate_calls, raw_txs) = start_mock_rpc_with_chain_id("0x2105").await.unwrap();
|
||||
std::env::set_var("OPENHUMAN_WALLET_RPC_BASE", format!("http://{addr}"));
|
||||
|
||||
let prepared = prepare_transfer(PrepareTransferParams {
|
||||
chain: WalletChain::Evm,
|
||||
to_address: "0x1111111111111111111111111111111111111111".into(),
|
||||
amount_raw: "1000".into(),
|
||||
asset_symbol: None,
|
||||
evm_network: Some(EvmNetwork::BaseMainnet),
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.value;
|
||||
let executed = execute_prepared(ExecutePreparedParams {
|
||||
quote_id: prepared.quote_id.clone(),
|
||||
confirmed: true,
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.value;
|
||||
assert_eq!(executed.status, PreparedStatus::Broadcasted);
|
||||
assert_eq!(executed.evm_network, Some(EvmNetwork::BaseMainnet));
|
||||
assert_eq!(raw_txs.lock().len(), 1);
|
||||
}
|
||||
|
||||
/// Build a Quote-store fixture pinned to a specific owner.
|
||||
/// Bypasses prepare_* (which would need full wallet setup + mock RPC) so
|
||||
/// the owner-gate behavior can be exercised in isolation.
|
||||
fn insert_owned_quote(quote_id: &str, owner: Option<QuoteOwner>) -> PreparedTransaction {
|
||||
let now = now_ms();
|
||||
let q = PreparedTransaction {
|
||||
quote_id: quote_id.to_string(),
|
||||
kind: PreparedKind::NativeTransfer,
|
||||
chain: WalletChain::Evm,
|
||||
evm_network: Some(EvmNetwork::EthereumMainnet),
|
||||
from_address: "0xfrom".to_string(),
|
||||
to_address: "0xto".to_string(),
|
||||
asset_symbol: "ETH".to_string(),
|
||||
amount_raw: "1".to_string(),
|
||||
amount_formatted: "0.000000000000000001".to_string(),
|
||||
receive_symbol: None,
|
||||
min_receive_raw: None,
|
||||
calldata: None,
|
||||
token_address: None,
|
||||
estimated_fee_raw: "0".to_string(),
|
||||
status: PreparedStatus::AwaitingConfirmation,
|
||||
created_at_ms: now,
|
||||
expires_at_ms: now + 60_000,
|
||||
notes: vec![],
|
||||
owner,
|
||||
};
|
||||
insert_quote_for_test(q)
|
||||
}
|
||||
|
||||
fn owner_a() -> QuoteOwner {
|
||||
QuoteOwner {
|
||||
thread_id: "thread-A".into(),
|
||||
client_id: "client-A".into(),
|
||||
}
|
||||
}
|
||||
|
||||
fn owner_b() -> QuoteOwner {
|
||||
QuoteOwner {
|
||||
thread_id: "thread-B".into(),
|
||||
client_id: "client-B".into(),
|
||||
}
|
||||
}
|
||||
|
||||
fn chat_ctx_from(owner: &QuoteOwner) -> crate::openhuman::approval::ApprovalChatContext {
|
||||
crate::openhuman::approval::ApprovalChatContext {
|
||||
thread_id: owner.thread_id.clone(),
|
||||
client_id: owner.client_id.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Cross-thread execute must fail. The quote must remain in the store
|
||||
/// (mismatched caller cannot poison it by consuming on Alice's behalf).
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_rejects_cross_owner_execution() {
|
||||
use crate::openhuman::approval::APPROVAL_CHAT_CONTEXT;
|
||||
let _guard = TEST_LOCK.lock();
|
||||
reset_quote_store_for_tests();
|
||||
let q = insert_owned_quote("q_xowner_1", Some(owner_a()));
|
||||
|
||||
// Bob (owner_b) tries to execute Alice's quote. The error string is
|
||||
// intentionally identical to a not-found lookup.
|
||||
let err = APPROVAL_CHAT_CONTEXT
|
||||
.scope(
|
||||
chat_ctx_from(&owner_b()),
|
||||
execute_prepared(ExecutePreparedParams {
|
||||
quote_id: q.quote_id.clone(),
|
||||
confirmed: true,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(err, format!("quote '{}' not found", q.quote_id));
|
||||
|
||||
// The quote must still be present and executable by Alice — Bob's
|
||||
// failed attempt didn't poison the store.
|
||||
let still_present = prepared_quotes_for_test();
|
||||
assert!(
|
||||
still_present.iter().any(|p| p.quote_id == q.quote_id),
|
||||
"quote must survive owner-mismatch attempts"
|
||||
);
|
||||
}
|
||||
|
||||
/// Same-thread execute must pass the owner gate (it will fail later in
|
||||
/// the chain code because there's no mock RPC set up, but the failure
|
||||
/// must not be the "not found" oracle — that proves we got past the gate).
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_allows_same_owner_execution() {
|
||||
use crate::openhuman::approval::APPROVAL_CHAT_CONTEXT;
|
||||
let _guard = TEST_LOCK.lock();
|
||||
reset_quote_store_for_tests();
|
||||
let q = insert_owned_quote("q_same_owner_1", Some(owner_a()));
|
||||
|
||||
let result = APPROVAL_CHAT_CONTEXT
|
||||
.scope(
|
||||
chat_ctx_from(&owner_a()),
|
||||
execute_prepared(ExecutePreparedParams {
|
||||
quote_id: q.quote_id.clone(),
|
||||
confirmed: true,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
// Past the owner gate. Chain code may error (no mock RPC, no wallet
|
||||
// setup) — but it must NOT be the "not found" shape we use for the
|
||||
// owner-mismatch oracle.
|
||||
if let Err(err) = &result {
|
||||
assert_ne!(
|
||||
err,
|
||||
&format!("quote '{}' not found", q.quote_id),
|
||||
"same-owner path must not return the owner-mismatch oracle"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// No-context prepare + no-context execute must work. Keeps CLI / direct
|
||||
/// JSON-RPC flows usable.
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_allows_no_context_flows() {
|
||||
let _guard = TEST_LOCK.lock();
|
||||
reset_quote_store_for_tests();
|
||||
let q = insert_owned_quote("q_no_ctx_1", None);
|
||||
|
||||
// No APPROVAL_CHAT_CONTEXT scope — current_owner() returns None on
|
||||
// both prepare and execute, so the gate passes.
|
||||
let result = execute_prepared(ExecutePreparedParams {
|
||||
quote_id: q.quote_id.clone(),
|
||||
confirmed: true,
|
||||
})
|
||||
.await;
|
||||
if let Err(err) = &result {
|
||||
assert_ne!(
|
||||
err,
|
||||
&format!("quote '{}' not found", q.quote_id),
|
||||
"no-context path must not return the owner-mismatch oracle"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// A quote prepared inside a chat context must NOT be executable from a
|
||||
/// caller with no context. Prevents privilege-drop into background /
|
||||
/// triage / cron flows that wouldn't surface UI confirmation.
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_rejects_chat_quote_from_no_context_caller() {
|
||||
let _guard = TEST_LOCK.lock();
|
||||
reset_quote_store_for_tests();
|
||||
let q = insert_owned_quote("q_chat_to_bg_1", Some(owner_a()));
|
||||
|
||||
// No scope around execute → caller_owner = None ≠ Some(owner_a).
|
||||
let err = execute_prepared(ExecutePreparedParams {
|
||||
quote_id: q.quote_id.clone(),
|
||||
confirmed: true,
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(err, format!("quote '{}' not found", q.quote_id));
|
||||
}
|
||||
|
||||
/// Lock the error-shape invariant: cross-owner reject string MUST be
|
||||
/// byte-equal to the not-found string. Regressions here would re-open
|
||||
/// the enumeration-oracle gap.
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_owner_mismatch_error_matches_not_found_shape() {
|
||||
use crate::openhuman::approval::APPROVAL_CHAT_CONTEXT;
|
||||
let _guard = TEST_LOCK.lock();
|
||||
reset_quote_store_for_tests();
|
||||
|
||||
// Real (owner-mismatched) quote.
|
||||
let real = insert_owned_quote("q_real_1", Some(owner_a()));
|
||||
// Reach the mismatched-owner branch.
|
||||
let mismatch_err = APPROVAL_CHAT_CONTEXT
|
||||
.scope(
|
||||
chat_ctx_from(&owner_b()),
|
||||
execute_prepared(ExecutePreparedParams {
|
||||
quote_id: real.quote_id.clone(),
|
||||
confirmed: true,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
// Reach the genuine not-found branch (no quote with this id in store).
|
||||
let missing_err = APPROVAL_CHAT_CONTEXT
|
||||
.scope(
|
||||
chat_ctx_from(&owner_b()),
|
||||
execute_prepared(ExecutePreparedParams {
|
||||
quote_id: "q_does_not_exist".into(),
|
||||
confirmed: true,
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
// Both branches surface the exact same template, parameterised only
|
||||
// by quote_id. No enumeration oracle.
|
||||
assert_eq!(mismatch_err, format!("quote '{}' not found", real.quote_id));
|
||||
assert_eq!(missing_err, "quote 'q_does_not_exist' not found");
|
||||
}
|
||||
|
||||
/// Verify that `prepare_transfer` inside an `APPROVAL_CHAT_CONTEXT` scope
|
||||
/// actually stamps `owner` via the task-local — not just via test helpers.
|
||||
#[tokio::test]
|
||||
async fn prepare_stamps_owner_via_task_local() {
|
||||
use crate::openhuman::approval::APPROVAL_CHAT_CONTEXT;
|
||||
let _guard = TEST_LOCK.lock();
|
||||
let _env_guard = crate::openhuman::config::TEST_ENV_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
reset_quote_store_for_tests();
|
||||
let temp = TempDir::new().unwrap();
|
||||
setup_wallet_in(&temp).await.unwrap();
|
||||
|
||||
let expected = owner_a();
|
||||
let ctx = chat_ctx_from(&expected);
|
||||
|
||||
let prepared = APPROVAL_CHAT_CONTEXT
|
||||
.scope(
|
||||
ctx,
|
||||
prepare_transfer(PrepareTransferParams {
|
||||
chain: WalletChain::Evm,
|
||||
to_address: "0x1111111111111111111111111111111111111111".into(),
|
||||
amount_raw: "1000".into(),
|
||||
asset_symbol: None,
|
||||
evm_network: Some(EvmNetwork::EthereumMainnet),
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.value;
|
||||
|
||||
assert_eq!(
|
||||
prepared.owner,
|
||||
Some(expected),
|
||||
"prepare_transfer must stamp owner from APPROVAL_CHAT_CONTEXT"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn execute_prepared_rejects_evm_chain_id_mismatch() {
|
||||
let _guard = TEST_LOCK.lock();
|
||||
let _env_guard = crate::openhuman::config::TEST_ENV_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
reset_quote_store_for_tests();
|
||||
let temp = TempDir::new().unwrap();
|
||||
setup_wallet_in(&temp).await.unwrap();
|
||||
// Quote says Base; mock reports Ethereum (0x1) — must fail.
|
||||
let (addr, _e, _r) = start_mock_rpc_with_chain_id("0x1").await.unwrap();
|
||||
std::env::set_var("OPENHUMAN_WALLET_RPC_BASE", format!("http://{addr}"));
|
||||
|
||||
let prepared = prepare_transfer(PrepareTransferParams {
|
||||
chain: WalletChain::Evm,
|
||||
to_address: "0x1111111111111111111111111111111111111111".into(),
|
||||
amount_raw: "1000".into(),
|
||||
asset_symbol: None,
|
||||
evm_network: Some(EvmNetwork::BaseMainnet),
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.value;
|
||||
let err = execute_prepared(ExecutePreparedParams {
|
||||
quote_id: prepared.quote_id.clone(),
|
||||
confirmed: true,
|
||||
})
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(err.contains("chain_id mismatch"), "got: {err}");
|
||||
}
|
||||
Reference in New Issue
Block a user