Add per-probe runtime deadline

Author: zergzorg

src-tauri/src/plugin_engine/host_api.rs

@@ -1,16 +1,17 @@
 use aes_gcm::{
-    AesGcm, Nonce,
-    aead::{Aead, KeyInit, OsRng, generic_array::typenum::U16, rand_core::RngCore},
+    aead::{generic_array::typenum::U16, rand_core::RngCore, Aead, KeyInit, OsRng},
     aes::Aes256,
+    AesGcm, Nonce,
 };
-use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64_STANDARD};
+use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine as _};
 use rquickjs::{Ctx, Exception, Function, Object};
 use sha2::{Digest, Sha256};
 use std::collections::{HashMap, HashSet};
 use std::ffi::{OsStr, OsString};
 use std::path::{Path, PathBuf};
 use std::process::Command;
 use std::sync::{Mutex, OnceLock};
+use std::time::{Duration, Instant};
 
 const WHITELISTED_ENV_VARS: [&str; 16] = [
     "CODEX_HOME",
@@ -31,6 +32,45 @@ const WHITELISTED_ENV_VARS: [&str; 16] = [
     "PI_CODING_AGENT_DIR",
 ];
 
+#[derive(Clone, Copy, Debug)]
+pub(crate) struct ProbeDeadline {
+    expires_at: Option<Instant>,
+}
+
+impl ProbeDeadline {
+    #[cfg(test)]
+    pub(crate) fn none() -> Self {
+        Self { expires_at: None }
+    }
+
+    pub(crate) fn at(expires_at: Instant) -> Self {
+        Self {
+            expires_at: Some(expires_at),
+        }
+    }
+
+    pub(crate) fn has_elapsed(self) -> bool {
+        self.expires_at
+            .map(|expires_at| Instant::now() >= expires_at)
+            .unwrap_or(false)
+    }
+
+    fn clamp_duration(self, requested: Duration) -> Duration {
+        let Some(expires_at) = self.expires_at else {
+            return requested;
+        };
+        let remaining = expires_at
+            .checked_duration_since(Instant::now())
+            .unwrap_or_default();
+        let clamped = requested.min(remaining);
+        if clamped.is_zero() {
+            Duration::from_millis(1)
+        } else {
+            clamped
+        }
+    }
+}
+
 fn last_non_empty_trimmed_line(text: &str) -> Option<String> {
     text.lines()
         .map(|line| line.trim())
@@ -511,11 +551,28 @@ fn encrypt_aes_256_gcm_envelope(plaintext: &str, key_b64: &str) -> Result<String
     ))
 }
 
-pub fn inject_host_api<'js>(
+#[cfg(test)]
+pub(crate) fn inject_host_api<'js>(
+    ctx: &Ctx<'js>,
+    plugin_id: &str,
+    app_data_dir: &PathBuf,
+    app_version: &str,
+) -> rquickjs::Result<()> {
+    inject_host_api_with_deadline(
+        ctx,
+        plugin_id,
+        app_data_dir,
+        app_version,
+        ProbeDeadline::none(),
+    )
+}
+
+pub(crate) fn inject_host_api_with_deadline<'js>(
     ctx: &Ctx<'js>,
     plugin_id: &str,
     app_data_dir: &PathBuf,
     app_version: &str,
+    deadline: ProbeDeadline,
 ) -> rquickjs::Result<()> {
     let globals = ctx.globals();
     let probe_ctx = Object::new(ctx.clone())?;
@@ -545,11 +602,11 @@ pub fn inject_host_api<'js>(
     inject_fs(ctx, &host)?;
     inject_crypto(ctx, &host)?;
     inject_env(ctx, &host, plugin_id)?;
-    inject_http(ctx, &host, plugin_id)?;
+    inject_http(ctx, &host, plugin_id, deadline)?;
     inject_keychain(ctx, &host, plugin_id)?;
     inject_sqlite(ctx, &host)?;
     inject_ls(ctx, &host, plugin_id)?;
-    inject_ccusage(ctx, &host, plugin_id)?;
+    inject_ccusage(ctx, &host, plugin_id, deadline)?;
 
     probe_ctx.set("host", host)?;
     globals.set("__openusage_ctx", probe_ctx)?;
@@ -720,7 +777,12 @@ fn inject_env<'js>(ctx: &Ctx<'js>, host: &Object<'js>, _plugin_id: &str) -> rqui
     Ok(())
 }
 
-fn inject_http<'js>(ctx: &Ctx<'js>, host: &Object<'js>, plugin_id: &str) -> rquickjs::Result<()> {
+fn inject_http<'js>(
+    ctx: &Ctx<'js>,
+    host: &Object<'js>,
+    plugin_id: &str,
+    deadline: ProbeDeadline,
+) -> rquickjs::Result<()> {
     let http_obj = Object::new(ctx.clone())?;
     let pid = plugin_id.to_string();
 
@@ -733,6 +795,10 @@ fn inject_http<'js>(ctx: &Ctx<'js>, host: &Object<'js>, plugin_id: &str) -> rqui
                     Exception::throw_message(&ctx_inner, &format!("invalid request: {}", e))
                 })?;
 
+                if deadline.has_elapsed() {
+                    return Err(Exception::throw_message(&ctx_inner, "probe timed out"));
+                }
+
                 let method_str = req.method.as_deref().unwrap_or("GET");
                 let redacted_url = redact_url(&req.url);
                 log::info!("[plugin:{}] HTTP {} {}", pid, method_str, redacted_url);
@@ -758,9 +824,10 @@ fn inject_http<'js>(ctx: &Ctx<'js>, host: &Object<'js>, plugin_id: &str) -> rqui
                 }
 
                 let timeout_ms = req.timeout_ms.unwrap_or(10_000);
+                let timeout = deadline.clamp_duration(Duration::from_millis(timeout_ms));
                 let mut builder = reqwest::blocking::Client::builder()
-                    .timeout(std::time::Duration::from_millis(timeout_ms))
-                    .connect_timeout(std::time::Duration::from_millis(timeout_ms))
+                    .timeout(timeout)
+                    .connect_timeout(timeout)
                     .redirect(reqwest::redirect::Policy::none());
 
                 // Apply pre-resolved proxy (localhost bypass already configured)
@@ -1972,31 +2039,56 @@ fn format_ccusage_timeout(timeout: std::time::Duration) -> String {
     format!("{:.3}s", timeout.as_secs_f64())
 }
 
+#[cfg(test)]
 fn run_ccusage_with_runner(
     kind: CcusageRunnerKind,
     program: &str,
     opts: &CcusageQueryOpts,
     provider: CcusageProvider,
     plugin_id: &str,
 ) -> CcusageRunnerResult {
+    run_ccusage_with_runner_deadline(
+        kind,
+        program,
+        opts,
+        provider,
+        plugin_id,
+        ProbeDeadline::none(),
+    )
+}
+
+fn run_ccusage_with_runner_deadline(
+    kind: CcusageRunnerKind,
+    program: &str,
+    opts: &CcusageQueryOpts,
+    provider: CcusageProvider,
+    plugin_id: &str,
+    deadline: ProbeDeadline,
+) -> CcusageRunnerResult {
+    if deadline.has_elapsed() {
+        log::warn!("[plugin:{}] ccusage skipped: probe timed out", plugin_id);
+        return CcusageRunnerResult::TimedOut;
+    }
+
     let current = run_ccusage_with_runner_timeout(
         kind,
         program,
         opts,
         provider,
         plugin_id,
         CcusageCommandFlavor::Current,
-        std::time::Duration::from_secs(CCUSAGE_TIMEOUT_SECS),
+        deadline.clamp_duration(Duration::from_secs(CCUSAGE_TIMEOUT_SECS)),
     );
     match current {
+        CcusageRunnerResult::Failed if deadline.has_elapsed() => CcusageRunnerResult::TimedOut,
         CcusageRunnerResult::Failed => run_ccusage_with_runner_timeout(
             kind,
             program,
             opts,
             provider,
             plugin_id,
             CcusageCommandFlavor::Legacy,
-            std::time::Duration::from_secs(CCUSAGE_TIMEOUT_SECS),
+            deadline.clamp_duration(Duration::from_secs(CCUSAGE_TIMEOUT_SECS)),
         ),
         other => other,
     }
@@ -2195,6 +2287,7 @@ fn inject_ccusage<'js>(
     ctx: &Ctx<'js>,
     host: &Object<'js>,
     plugin_id: &str,
+    deadline: ProbeDeadline,
 ) -> rquickjs::Result<()> {
     let ccusage_obj = Object::new(ctx.clone())?;
     let pid = plugin_id.to_string();
@@ -2222,7 +2315,11 @@ fn inject_ccusage<'js>(
                     &opts,
                     provider,
                     &pid,
-                    run_ccusage_with_runner,
+                    |kind, program, opts, provider, plugin_id| {
+                        run_ccusage_with_runner_deadline(
+                            kind, program, opts, provider, plugin_id, deadline,
+                        )
+                    },
                 ))
             },
         )?,
@@ -4087,6 +4184,21 @@ esac
         );
     }
 
+    #[test]
+    fn probe_deadline_clamps_host_timeout_to_remaining_budget() {
+        let deadline = ProbeDeadline::at(Instant::now() + Duration::from_millis(25));
+        let clamped = deadline.clamp_duration(Duration::from_secs(10));
+
+        assert!(
+            clamped <= Duration::from_millis(25),
+            "host timeout should not exceed remaining probe budget"
+        );
+        assert!(
+            clamped >= Duration::from_millis(1),
+            "host timeout should stay non-zero for blocking clients"
+        );
+    }
+
     #[cfg(unix)]
     #[test]
     fn ccusage_timeout_kills_descendant_and_closes_pipes() {