use super::*;
#[test]
fn request_serialises_correctly() {
let texts = vec!["hello".to_string(), "world".to_string()];
let req = RpcRequest {
jsonrpc: JSONRPC_VERSION,
method: METHOD_EMBED,
params: EmbedParams { texts: &texts },
id: 1,
};
let s = serde_json::to_string(&req).unwrap();
assert!(s.contains("\"jsonrpc\":\"2.0\""), "must have jsonrpc 2.0");
assert!(s.contains("\"method\":\"embed\""), "must have embed method");
assert!(
s.contains("\"texts\":[\"hello\",\"world\"]"),
"must include texts"
);
assert!(s.contains("\"id\":1"), "must have id");
}
#[test]
fn error_response_maps_to_model_error() {
let json = r#"{"jsonrpc":"2.0","error":{"code":-32603,"message":"ort failed"},"id":1}"#;
let result = decode_response(json, 1);
assert!(
matches!(result, Err(EmbedderError::ModelError(_))),
"got: {result:?}"
);
}
#[test]
fn success_response_decoded() {
let json = r#"{"jsonrpc":"2.0","result":{"embeddings":[[0.1,0.2],[0.3,0.4]]},"id":1}"#;
let result = decode_response(json, 2).unwrap();
assert_eq!(result.len(), 2);
assert_eq!(result[0][0], 0.1_f32);
}
#[test]
fn count_mismatch_returns_dimension_error() {
let json = r#"{"jsonrpc":"2.0","result":{"embeddings":[[0.1],[0.2]]},"id":1}"#;
let result = decode_response(json, 3);
assert!(
matches!(
result,
Err(EmbedderError::DimensionMismatch { sent: 3, got: 2 })
),
"got: {result:?}"
);
}
#[test]
fn extract_response_id_numeric() {
let json = r#"{"jsonrpc":"2.0","result":{"embeddings":[]},"id":42}"#;
assert_eq!(extract_response_id(json), Some(42));
}
#[test]
fn extract_response_id_null_returns_none() {
let json = r#"{"jsonrpc":"2.0","error":{"code":-32700,"message":"parse error"},"id":null}"#;
assert_eq!(extract_response_id(json), None);
}
#[test]
fn extract_response_id_string_returns_none() {
let json = r#"{"jsonrpc":"2.0","result":{"embeddings":[]},"id":"abc"}"#;
assert_eq!(extract_response_id(json), None);
}
#[test]
fn extract_response_device_present_returns_it() {
let json = r#"{"jsonrpc":"2.0","result":{"embeddings":[[0.1]],"device":"mps"},"id":1}"#;
assert_eq!(extract_response_device(json), Some("mps".to_string()));
}
#[test]
fn extract_response_device_absent_returns_none() {
let json = r#"{"jsonrpc":"2.0","result":{"embeddings":[[0.1]]},"id":1}"#;
assert_eq!(extract_response_device(json), None);
}
#[test]
fn extract_response_device_error_frame_returns_none() {
let json = r#"{"jsonrpc":"2.0","error":{"code":-32603,"message":"boom"},"id":1}"#;
assert_eq!(extract_response_device(json), None);
}
#[test]
fn extract_response_device_malformed_json_returns_none() {
assert_eq!(extract_response_device("{not valid json"), None);
}
#[tokio::test]
async fn reader_task_captures_wire_device() {
use tokio::io::{AsyncWriteExt, duplex};
use tokio::sync::oneshot;
let (mut writer, reader_end) = duplex(4096);
let reader = tokio::io::BufReader::new(reader_end);
let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
let (unhealthy_tx, _unhealthy_rx) = watch::channel(false);
let timeout_tracker = Arc::new(TimeoutTracker::new());
let last_device: Arc<std::sync::Mutex<Option<String>>> = Arc::new(std::sync::Mutex::new(None));
let device_capture_attempted = Arc::new(std::sync::atomic::AtomicBool::new(false));
let handle = tokio::spawn(reader_task(
reader,
Arc::clone(&pending),
Duration::from_secs(30),
timeout_tracker,
unhealthy_tx,
Arc::clone(&last_device),
Arc::clone(&device_capture_attempted),
));
let (tx, rx) = oneshot::channel();
pending
.lock()
.await
.insert(1, PendingRequest { sent: 1, reply: tx });
assert!(
last_device.lock().unwrap().is_none(),
"last_device must be None before any response arrives"
);
let frame =
b"{\"jsonrpc\":\"2.0\",\"result\":{\"embeddings\":[[0.1]],\"device\":\"mps\"},\"id\":1}\n";
writer.write_all(frame).await.unwrap();
writer.flush().await.unwrap();
let result = tokio::time::timeout(Duration::from_secs(2), rx)
.await
.expect("rx timed out")
.expect("channel closed unexpectedly");
assert!(result.is_ok(), "the embed call itself must still succeed");
assert_eq!(
last_device.lock().unwrap().clone(),
Some("mps".to_string()),
"reader_task must capture the wire-reported device"
);
drop(writer);
let _ = tokio::time::timeout(Duration::from_secs(1), handle).await;
}
#[tokio::test]
async fn reader_task_captures_device_once_and_never_reparses() {
use tokio::io::{AsyncWriteExt, duplex};
use tokio::sync::oneshot;
let (mut writer, reader_end) = duplex(4096);
let reader = tokio::io::BufReader::new(reader_end);
let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
let (unhealthy_tx, _unhealthy_rx) = watch::channel(false);
let timeout_tracker = Arc::new(TimeoutTracker::new());
let last_device: Arc<std::sync::Mutex<Option<String>>> = Arc::new(std::sync::Mutex::new(None));
let device_capture_attempted = Arc::new(std::sync::atomic::AtomicBool::new(false));
let handle = tokio::spawn(reader_task(
reader,
Arc::clone(&pending),
Duration::from_secs(30),
timeout_tracker,
unhealthy_tx,
Arc::clone(&last_device),
Arc::clone(&device_capture_attempted),
));
let (tx1, rx1) = oneshot::channel();
pending.lock().await.insert(
1,
PendingRequest {
sent: 1,
reply: tx1,
},
);
let frame1 =
b"{\"jsonrpc\":\"2.0\",\"result\":{\"embeddings\":[[0.1]],\"device\":\"mps\"},\"id\":1}\n";
writer.write_all(frame1).await.unwrap();
writer.flush().await.unwrap();
tokio::time::timeout(Duration::from_secs(2), rx1)
.await
.expect("rx1 timed out")
.expect("channel closed unexpectedly")
.expect("frame 1 must decode successfully");
assert_eq!(
last_device.lock().unwrap().clone(),
Some("mps".to_string()),
"must capture device from the first frame"
);
let (tx2, rx2) = oneshot::channel();
pending.lock().await.insert(
2,
PendingRequest {
sent: 1,
reply: tx2,
},
);
let frame2 =
b"{\"jsonrpc\":\"2.0\",\"result\":{\"embeddings\":[[0.2]],\"device\":\"cpu\"},\"id\":2}\n";
writer.write_all(frame2).await.unwrap();
writer.flush().await.unwrap();
tokio::time::timeout(Duration::from_secs(2), rx2)
.await
.expect("rx2 timed out")
.expect("channel closed unexpectedly")
.expect("frame 2 must decode successfully");
assert_eq!(
last_device.lock().unwrap().clone(),
Some("mps".to_string()),
"device must NOT be overwritten by a later frame — reader_task must \
only ever parse for the device field once, on the first successful \
frame"
);
drop(writer);
let _ = tokio::time::timeout(Duration::from_secs(1), handle).await;
}
#[tokio::test]
async fn embed_call_stalled_reader_times_out() {
use tokio::io::AsyncBufReadExt;
use tokio::io::duplex;
let (_tx, rx) = duplex(1024);
let mut buf = String::new();
let mut reader = tokio::io::BufReader::new(rx);
let result = tokio::time::timeout(Duration::from_secs(1), reader.read_line(&mut buf)).await;
assert!(
result.is_err(),
"a read_line on a never-writing reader must time out under a 1 s deadline; \
got: {result:?}"
);
}
#[tokio::test]
async fn reader_task_survives_timeout_and_serves_next_request() {
use tokio::io::{AsyncWriteExt, duplex};
use tokio::sync::oneshot;
let short_timeout = Duration::from_millis(50);
let (mut writer, reader_end) = duplex(4096);
let reader = tokio::io::BufReader::new(reader_end);
let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
let pending_clone = Arc::clone(&pending);
let (unhealthy_tx, _unhealthy_rx) = watch::channel(false);
let timeout_tracker = Arc::new(TimeoutTracker::new());
let last_device: Arc<std::sync::Mutex<Option<String>>> = Arc::new(std::sync::Mutex::new(None));
let device_capture_attempted = Arc::new(std::sync::atomic::AtomicBool::new(false));
let handle = tokio::spawn(reader_task(
reader,
pending_clone,
short_timeout,
timeout_tracker,
unhealthy_tx,
Arc::clone(&last_device),
Arc::clone(&device_capture_attempted),
));
let (tx_a, mut rx_a) = oneshot::channel();
pending.lock().await.insert(
1,
PendingRequest {
sent: 2,
reply: tx_a,
},
);
tokio::time::sleep(short_timeout * 3).await;
let result_a = rx_a.try_recv();
assert!(
matches!(result_a, Ok(Err(EmbedderError::Stdio(_)))),
"request A after timeout must receive Err(Stdio): got {result_a:?}"
);
let stale_a =
b"{\"jsonrpc\":\"2.0\",\"result\":{\"embeddings\":[[0.1,0.2],[0.3,0.4]]},\"id\":1}\n";
writer.write_all(stale_a).await.unwrap();
writer.flush().await.unwrap();
let (tx_b, rx_b) = oneshot::channel();
pending.lock().await.insert(
2,
PendingRequest {
sent: 2,
reply: tx_b,
},
);
let real_b =
b"{\"jsonrpc\":\"2.0\",\"result\":{\"embeddings\":[[0.5,0.6],[0.7,0.8]]},\"id\":2}\n";
writer.write_all(real_b).await.unwrap();
writer.flush().await.unwrap();
let result_b = tokio::time::timeout(Duration::from_secs(2), rx_b)
.await
.expect("rx_b timed out — reader task may have exited instead of continuing")
.expect("rx_b channel closed unexpectedly");
assert!(
result_b.is_ok(),
"request B must succeed after reader task survived timeout (#763): \
got {result_b:?}"
);
let embeddings_b = result_b.unwrap();
assert_eq!(
embeddings_b.len(),
2,
"request B must return 2 embedding vectors"
);
assert!(
(embeddings_b[0][0] - 0.5_f32).abs() < 1e-6,
"request B must receive its OWN embeddings (0.5…), not A's stale \
embeddings (0.1…) — misattribution bug would put 0.1 here. \
Got: {:?}",
embeddings_b[0]
);
assert!(
(embeddings_b[1][0] - 0.7_f32).abs() < 1e-6,
"request B second vector must be B's own data (0.7…), not A's (0.3…). \
Got: {:?}",
embeddings_b[1]
);
drop(writer);
let _ = tokio::time::timeout(Duration::from_secs(1), handle).await;
}
#[test]
fn timeout_stall_hint_is_provider_aware() {
use crate::embedder::ExecutionProvider;
let cuda = super::timeout_stall_hint(ExecutionProvider::Cuda);
assert!(
cuda.contains("CUDA"),
"CUDA provider must mention CUDA; got: {cuda:?}"
);
assert!(
!cuda.contains("CoreML"),
"CUDA hint must not mention CoreML; got: {cuda:?}"
);
let coreml = super::timeout_stall_hint(ExecutionProvider::CoreML);
assert!(
coreml.contains("CoreML"),
"CoreML provider must mention CoreML; got: {coreml:?}"
);
assert!(
!coreml.contains("CUDA"),
"CoreML hint must not mention CUDA; got: {coreml:?}"
);
let coreml_ane = super::timeout_stall_hint(ExecutionProvider::CoreMLAne);
assert!(
coreml_ane.contains("CoreML"),
"CoreMLAne provider must mention CoreML; got: {coreml_ane:?}"
);
assert!(
!coreml_ane.contains("CUDA"),
"CoreMLAne hint must not mention CUDA; got: {coreml_ane:?}"
);
let cpu = super::timeout_stall_hint(ExecutionProvider::Cpu);
assert!(
!cpu.contains("CUDA"),
"CPU hint must not mention CUDA; got: {cpu:?}"
);
assert!(
!cpu.contains("CoreML"),
"CPU hint must not mention CoreML; got: {cpu:?}"
);
assert!(!cpu.is_empty(), "CPU hint must not be empty");
}
struct PanicReader;
impl tokio::io::AsyncRead for PanicReader {
fn poll_read(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
_buf: &mut tokio::io::ReadBuf<'_>,
) -> std::task::Poll<std::io::Result<()>> {
panic!("simulated reader task panic (#1448 regression test)");
}
}
impl tokio::io::AsyncBufRead for PanicReader {
fn poll_fill_buf(
self: std::pin::Pin<&mut Self>,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<std::io::Result<&[u8]>> {
panic!("simulated reader task panic (#1448 regression test)");
}
fn consume(self: std::pin::Pin<&mut Self>, _amt: usize) {}
}
#[tokio::test]
async fn reader_panic_drains_pending_and_signals_unhealthy() {
use tokio::sync::oneshot;
let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
let (tx, rx) = oneshot::channel();
pending
.lock()
.await
.insert(1, PendingRequest { sent: 1, reply: tx });
let long_timeout = Duration::from_secs(30);
let (unhealthy_tx, mut unhealthy_rx) = spawn_reader_with_monitor(
PanicReader,
Arc::clone(&pending),
long_timeout,
Arc::new(std::sync::Mutex::new(None)),
Arc::new(std::sync::atomic::AtomicBool::new(false)),
);
let _keepalive_tx = unhealthy_tx;
let result = tokio::time::timeout(Duration::from_secs(2), rx)
.await
.expect("pending request must resolve promptly after reader panic, not hang")
.expect("oneshot must not be dropped without a value");
assert!(
matches!(result, Err(EmbedderError::Stdio(_))),
"pending request must fail with EmbedderError::Stdio after reader panic: got {result:?}"
);
tokio::time::timeout(Duration::from_secs(2), unhealthy_rx.changed())
.await
.expect("unhealthy_signal must fire promptly after reader panic")
.expect("unhealthy watch channel must not close without a value");
assert!(
*unhealthy_rx.borrow(),
"unhealthy_signal must be true after reader panic"
);
}
#[tokio::test]
async fn reader_eof_exit_also_signals_unhealthy() {
use tokio::io::duplex;
let (writer, reader_end) = duplex(64);
let reader = tokio::io::BufReader::new(reader_end);
let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
let (unhealthy_tx, mut unhealthy_rx) = spawn_reader_with_monitor(
reader,
Arc::clone(&pending),
Duration::from_secs(30),
Arc::new(std::sync::Mutex::new(None)),
Arc::new(std::sync::atomic::AtomicBool::new(false)),
);
let _keepalive_tx = unhealthy_tx;
drop(writer);
tokio::time::timeout(Duration::from_secs(2), unhealthy_rx.changed())
.await
.expect("unhealthy_signal must fire after a normal EOF reader exit")
.expect("unhealthy watch channel must not close without a value");
assert!(
*unhealthy_rx.borrow(),
"unhealthy_signal must be true after EOF reader exit"
);
}
#[test]
fn timeout_tracker_fires_exactly_at_threshold() {
let tracker = TimeoutTracker::new();
for i in 0..WEDGED_TIMEOUT_THRESHOLD - 1 {
assert!(
!tracker.record_timeout(),
"must not fire before the threshold (call {i})"
);
}
assert!(
tracker.record_timeout(),
"must fire exactly on the call that reaches WEDGED_TIMEOUT_THRESHOLD"
);
}
#[test]
fn timeout_tracker_record_success_resets_count() {
let tracker = TimeoutTracker::new();
for _ in 0..WEDGED_TIMEOUT_THRESHOLD - 1 {
assert!(!tracker.record_timeout());
}
tracker.record_success();
for i in 0..WEDGED_TIMEOUT_THRESHOLD - 1 {
assert!(
!tracker.record_timeout(),
"must not fire before the threshold post-reset (call {i})"
);
}
assert!(
tracker.record_timeout(),
"must fire after a fresh threshold run"
);
}
#[tokio::test]
async fn wedge_threshold_fires_after_consecutive_timeouts() {
use tokio::io::duplex;
use tokio::sync::oneshot;
let short_timeout = Duration::from_millis(30);
let (_writer, reader_end) = duplex(64);
let reader = tokio::io::BufReader::new(reader_end);
let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
{
let mut guard = pending.lock().await;
for i in 0..WEDGED_TIMEOUT_THRESHOLD {
let (tx, _rx) = oneshot::channel();
guard.insert(u64::from(i) + 1, PendingRequest { sent: 1, reply: tx });
}
}
let (unhealthy_tx, mut unhealthy_rx) = watch::channel(false);
let timeout_tracker = Arc::new(TimeoutTracker::new());
let last_device: Arc<std::sync::Mutex<Option<String>>> = Arc::new(std::sync::Mutex::new(None));
let device_capture_attempted = Arc::new(std::sync::atomic::AtomicBool::new(false));
let handle = tokio::spawn(reader_task(
reader,
Arc::clone(&pending),
short_timeout,
Arc::clone(&timeout_tracker),
unhealthy_tx,
Arc::clone(&last_device),
Arc::clone(&device_capture_attempted),
));
tokio::time::timeout(Duration::from_secs(5), unhealthy_rx.changed())
.await
.expect("unhealthy_signal must fire once the wedge threshold is crossed")
.expect("unhealthy watch channel must not close without a value");
assert!(
*unhealthy_rx.borrow(),
"unhealthy_signal must be true after WEDGED_TIMEOUT_THRESHOLD consecutive timeouts"
);
assert!(
!handle.is_finished(),
"reader task must stay alive on wedge detection (fix #763 invariant) — \
only the supervisor acts on unhealthy_signal"
);
handle.abort();
}
#[tokio::test]
async fn wedge_threshold_resets_on_success() {
use tokio::io::{AsyncWriteExt, duplex};
use tokio::sync::oneshot;
const _: () = assert!(WEDGED_TIMEOUT_THRESHOLD >= 2);
let below = WEDGED_TIMEOUT_THRESHOLD - 1;
let short_timeout = Duration::from_millis(30);
let (mut writer, reader_end) = duplex(4096);
let reader = tokio::io::BufReader::new(reader_end);
let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
{
let mut guard = pending.lock().await;
for i in 0..below {
let (tx, _rx) = oneshot::channel();
guard.insert(u64::from(i) + 1, PendingRequest { sent: 1, reply: tx });
}
}
let (unhealthy_tx, unhealthy_rx) = watch::channel(false);
let timeout_tracker = Arc::new(TimeoutTracker::new());
let last_device: Arc<std::sync::Mutex<Option<String>>> = Arc::new(std::sync::Mutex::new(None));
let device_capture_attempted = Arc::new(std::sync::atomic::AtomicBool::new(false));
let handle = tokio::spawn(reader_task(
reader,
Arc::clone(&pending),
short_timeout,
Arc::clone(&timeout_tracker),
unhealthy_tx,
Arc::clone(&last_device),
Arc::clone(&device_capture_attempted),
));
wait_until_pending_empty(&pending, short_timeout * 20).await;
assert!(
!*unhealthy_rx.borrow(),
"must not be unhealthy yet — only {below} of {WEDGED_TIMEOUT_THRESHOLD} \
consecutive timeouts have occurred"
);
let success_id = 1_000_u64;
let (tx_ok, rx_ok) = oneshot::channel();
pending.lock().await.insert(
success_id,
PendingRequest {
sent: 1,
reply: tx_ok,
},
);
let frame = format!(
"{{\"jsonrpc\":\"2.0\",\"result\":{{\"embeddings\":[[0.1]]}},\"id\":{success_id}}}\n"
);
writer.write_all(frame.as_bytes()).await.unwrap();
writer.flush().await.unwrap();
tokio::time::timeout(Duration::from_secs(2), rx_ok)
.await
.expect("success response must be delivered promptly")
.expect("oneshot must not close without a value")
.expect("synthetic success response must decode without error");
{
let mut guard = pending.lock().await;
for i in 0..below {
let (tx, _rx) = oneshot::channel();
guard.insert(2_000 + u64::from(i), PendingRequest { sent: 1, reply: tx });
}
}
wait_until_pending_empty(&pending, short_timeout * 20).await;
assert!(
!*unhealthy_rx.borrow(),
"unhealthy_signal must not fire: the intervening success reset the \
consecutive-timeout counter, so only {below} timeouts have \
accumulated since — one below WEDGED_TIMEOUT_THRESHOLD"
);
handle.abort();
}
async fn wait_until_pending_empty(pending: &PendingMap, budget: Duration) {
let deadline = tokio::time::Instant::now() + budget;
loop {
if pending.lock().await.is_empty() {
return;
}
if tokio::time::Instant::now() >= deadline {
panic!("pending map did not drain within {budget:?}");
}
tokio::time::sleep(Duration::from_millis(5)).await;
}
}
#[tokio::test]
async fn request_after_reader_death_errors_promptly() {
let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
let (unhealthy_tx, unhealthy_rx) = watch::channel(false);
let (reply_tx, reply_rx) = oneshot::channel();
let id = 42_u64;
pending.lock().await.insert(
id,
PendingRequest {
sent: 1,
reply: reply_tx,
},
);
let _ = unhealthy_tx.send(true);
let result = tokio::time::timeout(
Duration::from_secs(2),
await_reply_or_unhealthy(id, &pending, reply_rx, unhealthy_rx),
)
.await
.expect("must error promptly instead of hanging when the client is already unhealthy");
assert!(
matches!(result, Err(EmbedderError::Stdio(_))),
"got: {result:?}"
);
assert!(
!pending.lock().await.contains_key(&id),
"the stranded pending entry must be removed, not leaked"
);
}
#[tokio::test]
async fn unhealthy_signal_during_wait_errors_promptly() {
let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
let (unhealthy_tx, unhealthy_rx) = watch::channel(false);
let (reply_tx, reply_rx) = oneshot::channel();
let id = 7_u64;
pending.lock().await.insert(
id,
PendingRequest {
sent: 1,
reply: reply_tx,
},
);
let flipper = tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(50)).await;
let _ = unhealthy_tx.send(true);
});
let result = tokio::time::timeout(
Duration::from_secs(2),
await_reply_or_unhealthy(id, &pending, reply_rx, unhealthy_rx),
)
.await
.expect("must not hang once the unhealthy signal fires mid-wait");
assert!(
matches!(result, Err(EmbedderError::Stdio(_))),
"got: {result:?}"
);
flipper.await.expect("flipper task must not panic");
}
#[tokio::test]
async fn await_reply_or_unhealthy_returns_ok_on_real_reply() {
let pending: PendingMap = Arc::new(Mutex::new(HashMap::new()));
let (_unhealthy_tx, unhealthy_rx) = watch::channel(false);
let (reply_tx, reply_rx) = oneshot::channel();
let _ = reply_tx.send(Ok(vec![vec![0.1_f32, 0.2_f32]]));
let result = tokio::time::timeout(
Duration::from_secs(2),
await_reply_or_unhealthy(1, &pending, reply_rx, unhealthy_rx),
)
.await
.expect("must resolve promptly")
.expect("must succeed when a real reply is already available");
assert_eq!(result, vec![vec![0.1_f32, 0.2_f32]]);
}