fix(ext/ai): use SHA-256 for ONNX session cache keys

ext/ai/Cargo.toml

@@ -20,17 +20,16 @@ ctor.workspace = true
 faster-hex.workspace = true
 futures.workspace = true
 futures-util = { workspace = true, features = ["io"] }
-fxhash.workspace = true
 log.workspace = true
 once_cell.workspace = true
 reqwest.workspace = true
 scopeguard.workspace = true
 serde.workspace = true
+sha2.workspace = true
 tokio.workspace = true
 tokio-util = { workspace = true, features = ["compat"] }
 tracing.workspace = true
 tracing-subscriber.workspace = true
-xxhash-rust = { workspace = true, features = ["std", "xxh3"] }
 
 blas-src = { version = "0.8", features = ['openblas'] }
 convert_case = "0.6"

ext/ai/onnxruntime/mod.rs

@@ -50,7 +50,7 @@ pub async fn op_ai_ort_init_session(
     }
     Err(_) => {
       trace!(kind = "bytes", len = model_bytes.len());
-      Model::from_bytes(&model_bytes).await?
+      Model::from_bytes(model_bytes).await?
     }
   };
 

ext/ai/onnxruntime/model.rs

@@ -75,7 +75,7 @@ impl Model {
     load_session_from_url(model_url).await.map(Self::new)
   }
 
-  pub async fn from_bytes(model_bytes: &[u8]) -> Result<Self> {
+  pub async fn from_bytes(model_bytes: Box<[u8]>) -> Result<Self> {
     load_session_from_bytes(model_bytes).await.map(Self::new)
   }
 }

ext/ai/onnxruntime/session.rs

@@ -1,16 +1,14 @@
 use deno_core::error::AnyError;
-use futures::io::AllowStdIo;
 use once_cell::sync::Lazy;
 use reqwest::Url;
+use sha2::Digest;
+use sha2::Sha256;
 use std::collections::HashMap;
-use std::hash::Hasher;
 use std::sync::Arc;
 use tokio::sync::Mutex;
-use tokio_util::compat::FuturesAsyncWriteCompatExt;
 use tracing::debug;
 use tracing::instrument;
 use tracing::trace;
-use xxhash_rust::xxh3::Xxh3;
 
 use anyhow::anyhow;
 use anyhow::Context;
@@ -125,18 +123,16 @@ fn create_session(model_bytes: &[u8]) -> Result<Arc<Session>, Error> {
 
 #[instrument(level = "debug", skip_all, fields(model_bytes = model_bytes.len()), err)]
 pub(crate) async fn load_session_from_bytes(
-  model_bytes: &[u8],
+  model_bytes: Box<[u8]>,
 ) -> Result<SessionWithId, Error> {
-  let session_id = {
-    let mut model_bytes = tokio::io::BufReader::new(model_bytes);
-    let mut hasher = AllowStdIo::new(Xxh3::new()).compat_write();
-    let _ = tokio::io::copy(&mut model_bytes, &mut hasher)
-      .await
-      .context("failed to calculate checksum")?;
-
-    let hasher = hasher.into_inner().into_inner();
-    faster_hex::hex_string(&hasher.finish().to_be_bytes())
-  };
+  // Hash off the isolate thread: this runs on every create(bytes) call, even
+  // on a cache hit, and would otherwise count toward the worker's CPU budget.
+  let (session_id, model_bytes) = tokio::task::spawn_blocking(move || {
+    let digest = Sha256::digest(&model_bytes);
+    (faster_hex::hex_string(&digest), model_bytes)
+  })
+  .await
+  .context("failed to calculate checksum")?;
 
   let mut sessions = SESSIONS.lock().await;
 
@@ -145,7 +141,7 @@ pub(crate) async fn load_session_from_bytes(
   }
 
   trace!(session_id, "new session");
-  let session = create_session(model_bytes)?;
+  let session = create_session(&model_bytes)?;
 
   sessions.insert(session_id.clone(), session.clone());
 
@@ -156,7 +152,7 @@ pub(crate) async fn load_session_from_bytes(
 pub(crate) async fn load_session_from_url(
   model_url: Url,
 ) -> Result<SessionWithId, Error> {
-  let session_id = fxhash::hash(model_url.as_str()).to_string();
+  let session_id = faster_hex::hex_string(&Sha256::digest(model_url.as_str()));
 
   let mut sessions = SESSIONS.lock().await;
 

ext/ai/utils.rs

@@ -1,27 +1,28 @@
-use std::hash::Hasher;
 use std::path::PathBuf;
 
 use anyhow::Context;
 use deno_core::error::AnyError;
 use futures::io::AllowStdIo;
 use futures::StreamExt;
 use reqwest::Url;
+use sha2::Digest;
+use sha2::Sha256;
 use tokio::io::AsyncWriteExt;
 use tokio_util::compat::FuturesAsyncWriteCompatExt;
 use tracing::info;
 use tracing::info_span;
 use tracing::instrument;
 use tracing::trace;
 use tracing::Instrument;
-use xxhash_rust::xxh3::Xxh3;
 
 #[instrument(fields(%kind, url = %url))]
 pub async fn fetch_and_cache_from_url(
   kind: &'static str,
   url: Url,
   cache_id: Option<String>,
 ) -> Result<PathBuf, AnyError> {
-  let cache_id = cache_id.unwrap_or(fxhash::hash(url.as_str()).to_string());
+  let cache_id = cache_id
+    .unwrap_or_else(|| faster_hex::hex_string(&Sha256::digest(url.as_str())));
   let download_dir = std::env::var("EXT_AI_CACHE_DIR")
     .map_err(AnyError::from)
     .map(PathBuf::from)
@@ -63,10 +64,10 @@ pub async fn fetch_and_cache_from_url(
       let filepath = filepath.clone();
       let checksum = tokio::task::spawn_blocking(move || {
         let mut file = std::fs::File::open(filepath).ok()?;
-        let mut hasher = Xxh3::new();
+        let mut hasher = Sha256::new();
         let _ = std::io::copy(&mut file, &mut hasher).ok()?;
 
-        Some(hasher.finish().to_be_bytes())
+        Some(hasher.finalize())
       })
       .await;
 
@@ -106,7 +107,7 @@ pub async fn fetch_and_cache_from_url(
 
       let mut stream = resp.bytes_stream();
       let mut writer = tokio::io::BufWriter::new(file);
-      let mut hasher = AllowStdIo::new(Xxh3::new()).compat_write();
+      let mut hasher = AllowStdIo::new(Sha256::new()).compat_write();
       let mut written = 0;
 
       while let Some(chunk) = stream.next().await {
@@ -125,7 +126,7 @@ pub async fn fetch_and_cache_from_url(
 
       let checksum_str = {
         let hasher = hasher.into_inner().into_inner();
-        faster_hex::hex_string(&hasher.finish().to_be_bytes())
+        faster_hex::hex_string(&hasher.finalize())
       };
 
       info!({ bytes_written = written, checksum = &checksum_str }, "done");