refactor(rust): split oversized modules (#2991)

This commit is contained in:
Steven Enamakel
2026-05-29 16:36:56 -07:00
committed by GitHub
parent 5e7a45671e
commit 123e8b0353
40 changed files with 8126 additions and 7977 deletions
+2 -354
View File
@@ -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,
-84
View File
@@ -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(())
+360
View File
@@ -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
+6 -70
View File
@@ -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
}
+13 -530
View File
@@ -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. 90119s 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. 90119s 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
}
+2 -162
View File
@@ -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)
}
+1
View File
@@ -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
View File
@@ -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) ───────────────────────────────────────
+8 -94
View File
@@ -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(())
}
+2 -232
View File
@@ -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),
}
}
}
+11 -603
View File
@@ -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(&param_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;
+132
View File
@@ -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(&param_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)
}
+473
View File
@@ -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");
}
+2 -150
View File
@@ -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(&params, "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;
+148
View File
@@ -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(&params, "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");
}
+2 -246
View File
@@ -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;
+244
View File
@@ -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}"
);
}
}
+2 -138
View File
@@ -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);
}
+2 -36
View File
@@ -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;
+34
View File
@@ -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"));
}
+2 -727
View File
@@ -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))
}
+25 -810
View File
@@ -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"]
+1 -1
View File
@@ -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.
+7 -797
View File
@@ -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.
+799
View File
@@ -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)
}
+2 -659
View File
@@ -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;
+651
View File
@@ -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}");
}