diff --git a/crates/mesh-llm-host-runtime/src/runtime/survey.rs b/crates/mesh-llm-host-runtime/src/runtime/survey.rs index c4b071ea8..19502e266 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/survey.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/survey.rs @@ -26,6 +26,9 @@ const OTLP_METRICS_ENDPOINT_ENV: &str = "OTEL_EXPORTER_OTLP_METRICS_ENDPOINT"; const TELEMETRY_ATTRIBUTE_ALLOWLIST: &[&str] = &[ "llama_stage.verify_window.direct_return_reverse_fallback", "llama_stage.verify_window.direct_return_upstream_opened", + "llama_stage.linear_proposal.source_callback_us", + "llama_stage.linear_proposal.source_outcome", + "llama_stage.linear_proposal.source_queue_wait_us", "mesh_llm.architecture", "mesh_llm.attempt_outcome", "mesh_llm.backend", @@ -1385,6 +1388,9 @@ mod tests { assert_eq!( keys, BTreeSet::from([ + "llama_stage.linear_proposal.source_callback_us", + "llama_stage.linear_proposal.source_outcome", + "llama_stage.linear_proposal.source_queue_wait_us", "llama_stage.verify_window.direct_return_reverse_fallback", "llama_stage.verify_window.direct_return_upstream_opened", "mesh_llm.architecture", diff --git a/crates/mesh-native-serving-plugin-host/src/lib.rs b/crates/mesh-native-serving-plugin-host/src/lib.rs index a84c2ad05..70864694a 100644 --- a/crates/mesh-native-serving-plugin-host/src/lib.rs +++ b/crates/mesh-native-serving-plugin-host/src/lib.rs @@ -1,15 +1,15 @@ //! Loads one native serving plugin and adapts its stable ABI to Skippy hooks. +mod plugin_dispatch; +#[cfg(test)] +mod test_support; + use std::{ ffi::{c_char, c_void}, path::{Path, PathBuf}, ptr::NonNull, - sync::{ - Arc, Mutex, OnceLock, - atomic::{AtomicU64, Ordering}, - mpsc::{Receiver, SyncSender, TrySendError, sync_channel}, - }, - thread::{self, JoinHandle}, + sync::{Arc, Mutex, OnceLock}, + thread, time::{Duration, Instant}, }; @@ -21,16 +21,16 @@ use skippy_server::frontend::{ GenerationReceipt, GenerationReceiptConfig, GenerationStart, LinearProposal, LinearProposalDiscardReason, LinearProposalDisposition, LinearProposalIngress, LinearProposalIngressConfig, LinearProposalQuery, LinearProposalReceipt, - OpaqueProposalDecisionId, + LinearProposalSourceResponse, OpaqueProposalDecisionId, }; use skippy_server::serving_hooks::{ModelServingHooks, ModelServingHooksFactory}; use skippy_server::tokenizer::TokenizerCapability; +use plugin_dispatch::{PluginCommand, PluginDriver}; + const ERROR_BUFFER_BYTES: usize = 2_048; -const PLUGIN_COMMAND_CAPACITY: usize = 1_024; const MAX_NATIVE_PLUGIN_PROPOSAL_TOKENS: usize = 4_096; const PROPOSAL_POLL_INTERVAL: Duration = Duration::from_micros(50); -const CLEAN_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); /// Mesh-owned factory for one independently built native serving plugin. #[derive(Clone)] @@ -124,7 +124,7 @@ impl ModelServingHooksFactory for NativeServingPluginFactory { let active = ActivePlugin { definition: Arc::clone(&self.definition), instance: Some(instance), - proposal_token_buffer: vec![0; MAX_NATIVE_PLUGIN_PROPOSAL_TOKENS], + proposal_token_buffer: Mutex::new(vec![0; MAX_NATIVE_PLUGIN_PROPOSAL_TOKENS]), }; let driver = Arc::new(PluginDriver::spawn(active)?); let lifecycle: Arc = Arc::new(NativeLifecycleIngress { @@ -269,7 +269,7 @@ fn validate_table(table: &abi::NativeServingPluginV1) -> Result { struct ActivePlugin { definition: Arc, instance: Option>, - proposal_token_buffer: Vec, + proposal_token_buffer: Mutex>, } // SAFETY: activation succeeds only for plugins implementing the ABI's @@ -354,7 +354,7 @@ impl ActivePlugin { self.call_status("finish generation", status) } - fn propose(&mut self, query: LinearProposalQuery) -> Result> { + fn propose(&self, query: LinearProposalQuery) -> Result> { let event = abi::ProposalQuery { struct_size: size_of::(), request_id: query.request_id, @@ -374,17 +374,17 @@ impl ActivePlugin { } fn poll_until_deadline( - &mut self, + &self, operation: abi::ProposalOperation, deadline: Instant, max_proposal_tokens: usize, ) -> Result> { - let mut token_buffer = std::mem::take(&mut self.proposal_token_buffer); + let mut token_buffer = self + .proposal_token_buffer + .lock() + .map_err(|_| anyhow!("native serving plugin proposal token buffer lock poisoned"))?; let token_capacity = max_proposal_tokens.min(token_buffer.len()); - let result = - self.poll_with_buffer(operation, deadline, &mut token_buffer[..token_capacity]); - self.proposal_token_buffer = token_buffer; - result + self.poll_with_buffer(operation, deadline, &mut token_buffer[..token_capacity]) } fn poll_with_buffer( @@ -524,8 +524,17 @@ struct NativeProposalIngress { } impl LinearProposalIngress for NativeProposalIngress { - fn propose(&self, query: LinearProposalQuery) -> Result> { - self.driver.propose(query) + fn propose(&self, query: LinearProposalQuery) -> Result { + let response = self.driver.propose(query)?; + // A failing plugin callback must not fail the generation, so a source + // error degrades to an abstention for decode. The worker logs the + // plugin's message and the `SourceError` outcome keeps the failure + // distinguishable from a deliberate abstention in telemetry. + let proposal = response.proposal.unwrap_or_default(); + Ok(LinearProposalSourceResponse::with_telemetry( + proposal, + response.telemetry, + )) } fn report(&self, receipt: &LinearProposalReceipt) -> Result<()> { @@ -544,170 +553,6 @@ impl LinearProposalIngress for NativeProposalIngress { } } -enum PluginCommand { - Begin(GenerationStart), - Committed(GenerationCommit), - Abort(GenerationAbort), - Finish(GenerationReceipt), - Proposal( - LinearProposalQuery, - SyncSender, String>>, - ), - Report(LinearProposalReceipt), - Discard(Vec, LinearProposalDiscardReason), - Shutdown(SyncSender>), -} - -struct PluginDriver { - sender: SyncSender, - fatal_error: Arc>>, - lifecycle_delivery_failures: Arc, - worker: Mutex>>, -} - -impl PluginDriver { - fn spawn(active: ActivePlugin) -> Result { - let (sender, receiver) = sync_channel(PLUGIN_COMMAND_CAPACITY); - let fatal_error = Arc::new(Mutex::new(None)); - let worker_fatal_error = Arc::clone(&fatal_error); - let lifecycle_delivery_failures = Arc::new(AtomicU64::new(0)); - let worker_lifecycle_delivery_failures = Arc::clone(&lifecycle_delivery_failures); - let worker = thread::Builder::new() - .name("mesh-native-serving-plugin".to_string()) - .spawn(move || { - plugin_worker( - active, - receiver, - worker_fatal_error, - worker_lifecycle_delivery_failures, - ); - }) - .context("spawn native serving plugin worker")?; - Ok(Self { - sender, - fatal_error, - lifecycle_delivery_failures, - worker: Mutex::new(Some(worker)), - }) - } - - fn ensure_healthy(&self) -> Result<()> { - let error = self - .fatal_error - .lock() - .map_err(|_| anyhow!("native serving plugin health lock poisoned"))?; - if let Some(error) = error.as_deref() { - bail!("native serving plugin worker failed: {error}"); - } - Ok(()) - } - - fn enqueue(&self, command: PluginCommand) -> Result<()> { - self.ensure_healthy()?; - self.enqueue_recovery(command) - } - - /// Deliver lifecycle cleanup even after an earlier callback failed. - /// - /// This bypasses only the health gate. The command still uses the same - /// bounded FIFO queue and fails if the worker has stopped or is full. - fn enqueue_recovery(&self, command: PluginCommand) -> Result<()> { - self.sender.try_send(command).map_err(|error| match error { - TrySendError::Full(_) => anyhow!("native serving plugin command queue is full"), - TrySendError::Disconnected(_) => anyhow!("native serving plugin worker stopped"), - }) - } - - fn lifecycle_delivery_failures(&self) -> u64 { - self.lifecycle_delivery_failures.load(Ordering::Relaxed) - } - - fn propose(&self, query: LinearProposalQuery) -> Result> { - self.ensure_healthy()?; - let deadline = query.deadline; - let (reply, response) = sync_channel(1); - match self.sender.try_send(PluginCommand::Proposal(query, reply)) { - Ok(()) => {} - Err(TrySendError::Full(_)) => return Ok(None), - Err(TrySendError::Disconnected(_)) => { - bail!("native serving plugin worker stopped") - } - } - let remaining = deadline.saturating_duration_since(Instant::now()); - if remaining.is_zero() { - return Ok(None); - } - match response.recv_timeout(remaining) { - Ok(result) => result.map_err(anyhow::Error::msg), - Err(std::sync::mpsc::RecvTimeoutError::Timeout) => Ok(None), - Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => { - bail!("native serving plugin worker stopped before replying") - } - } - } -} - -impl Drop for PluginDriver { - fn drop(&mut self) { - let (reply, response) = sync_channel(1); - let queued = self.sender.try_send(PluginCommand::Shutdown(reply)).is_ok(); - if queued - && response.recv_timeout(CLEAN_SHUTDOWN_TIMEOUT).is_ok() - && let Ok(worker) = self.worker.get_mut() - && let Some(worker) = worker.take() - { - let _ = worker.join(); - } - } -} - -fn plugin_worker( - mut active: ActivePlugin, - receiver: Receiver, - fatal_error: Arc>>, - lifecycle_delivery_failures: Arc, -) { - while let Ok(command) = receiver.recv() { - let (result, terminal, lifecycle) = match command { - PluginCommand::Begin(event) => (active.begin(&event), false, true), - PluginCommand::Committed(event) => (active.committed(&event), false, true), - PluginCommand::Abort(event) => (active.abort(&event), false, true), - PluginCommand::Finish(event) => (active.finish(&event), false, true), - PluginCommand::Proposal(query, reply) => { - let result = active.propose(query); - let _ = reply.send( - result - .as_ref() - .map(Clone::clone) - .map_err(ToString::to_string), - ); - (result.map(|_| ()), false, false) - } - PluginCommand::Report(event) => (active.report(&event), false, false), - PluginCommand::Discard(decision_id, reason) => { - (active.discard(&decision_id, reason), false, false) - } - PluginCommand::Shutdown(reply) => { - let result = active.shutdown(); - let _ = reply.send(result.as_ref().map(|_| ()).map_err(ToString::to_string)); - (result, true, false) - } - }; - if lifecycle && result.is_err() { - lifecycle_delivery_failures.fetch_add(1, Ordering::Relaxed); - } - if terminal - && let Err(error) = result - && let Ok(mut fatal) = fatal_error.lock() - { - *fatal = Some(format!("{error:#}")); - } - if terminal { - break; - } - } -} - fn proposal_from_output( decision_id: &[u8; abi::MAX_DECISION_ID_BYTES], token_ids: &[i32], @@ -818,10 +663,7 @@ unsafe fn read_utf8(slice: abi::ByteSlice, label: &str) -> Result { #[cfg(test)] mod tests { use super::*; - use std::sync::atomic::{AtomicUsize, Ordering}; - - static FAKE_NAME: &[u8] = b"test-serving-plugin"; - static CANCEL_COUNT: AtomicUsize = AtomicUsize::new(0); + use crate::test_support::fake_table; #[test] fn tokenizer_inventory_view_borrows_host_owned_bytes_for_activation() { @@ -849,188 +691,6 @@ mod tests { ); } - struct FakeState { - start_delay: Duration, - begin_fails: bool, - events: Arc>>, - abort_count: Arc, - } - - unsafe extern "C" fn fake_activate( - _context: *const abi::ActivationContext, - _activation: *mut abi::PluginActivation, - ) -> abi::PluginStatus { - abi::PluginStatus::INTERNAL_ERROR - } - - unsafe extern "C" fn fake_shutdown(instance: abi::PluginInstance) -> abi::PluginStatus { - if !instance.is_null() { - drop(unsafe { Box::from_raw(instance.cast::()) }); - } - abi::PluginStatus::OK - } - - unsafe extern "C" fn fake_begin( - instance: abi::PluginInstance, - _event: *const abi::GenerationStart, - ) -> abi::PluginStatus { - let state = unsafe { &*instance.cast::() }; - state.events.lock().unwrap().push("begin"); - if state.begin_fails { - abi::PluginStatus::INTERNAL_ERROR - } else { - abi::PluginStatus::OK - } - } - - unsafe extern "C" fn fake_commit( - instance: abi::PluginInstance, - _event: *const abi::GenerationCommit, - ) -> abi::PluginStatus { - let state = unsafe { &*instance.cast::() }; - state.events.lock().unwrap().push("commit"); - abi::PluginStatus::OK - } - - unsafe extern "C" fn fake_abort( - instance: abi::PluginInstance, - _event: *const abi::GenerationAbort, - ) -> abi::PluginStatus { - let state = unsafe { &*instance.cast::() }; - state.abort_count.fetch_add(1, Ordering::SeqCst); - abi::PluginStatus::OK - } - - unsafe extern "C" fn fake_finish( - _instance: abi::PluginInstance, - _event: *const abi::GenerationFinish, - ) -> abi::PluginStatus { - abi::PluginStatus::OK - } - - unsafe extern "C" fn fake_start_proposal( - instance: abi::PluginInstance, - _query: *const abi::ProposalQuery, - operation: *mut abi::ProposalOperation, - ) -> abi::PluginStatus { - let state = unsafe { &*instance.cast::() }; - state.events.lock().unwrap().push("proposal"); - thread::sleep(state.start_delay); - unsafe { *operation = 1 }; - abi::PluginStatus::OK - } - - unsafe extern "C" fn fake_poll_proposal( - _instance: abi::PluginInstance, - _operation: abi::ProposalOperation, - _output: *mut abi::ProposalOutput, - ) -> abi::ProposalPollStatus { - abi::ProposalPollStatus::ABSTAIN - } - - unsafe extern "C" fn fake_cancel_proposal( - _instance: abi::PluginInstance, - _operation: abi::ProposalOperation, - ) { - CANCEL_COUNT.fetch_add(1, Ordering::SeqCst); - } - - unsafe extern "C" fn fake_report_proposal( - _instance: abi::PluginInstance, - _event: *const abi::ProposalOutcome, - ) -> abi::PluginStatus { - abi::PluginStatus::OK - } - - unsafe extern "C" fn fake_discard_proposal( - _instance: abi::PluginInstance, - _event: *const abi::ProposalDiscard, - ) -> abi::PluginStatus { - abi::PluginStatus::OK - } - - unsafe extern "C" fn fake_last_error( - _instance: abi::PluginInstance, - _output: *mut c_char, - _capacity: usize, - ) -> usize { - 0 - } - - fn fake_table() -> abi::NativeServingPluginV1 { - abi::NativeServingPluginV1 { - abi_version: abi::NATIVE_SERVING_PLUGIN_ABI_V1, - struct_size: size_of::(), - plugin_name: abi::ByteSlice::from_bytes(FAKE_NAME), - activate: fake_activate, - shutdown: fake_shutdown, - begin_generation: fake_begin, - commit_generation: fake_commit, - abort_generation: fake_abort, - finish_generation: fake_finish, - start_proposal: fake_start_proposal, - poll_proposal: fake_poll_proposal, - cancel_proposal: fake_cancel_proposal, - report_proposal: fake_report_proposal, - discard_proposal: fake_discard_proposal, - last_error: fake_last_error, - } - } - - fn fake_active(start_delay: Duration) -> ActivePlugin { - fake_active_with_events(start_delay).0 - } - - fn fake_active_with_events( - start_delay: Duration, - ) -> (ActivePlugin, Arc>>) { - let (active, events, _) = fake_active_with_observations(start_delay); - (active, events) - } - - fn fake_active_with_observations( - start_delay: Duration, - ) -> ( - ActivePlugin, - Arc>>, - Arc, - ) { - fake_active_with_options(start_delay, false) - } - - fn fake_active_with_options( - start_delay: Duration, - begin_fails: bool, - ) -> ( - ActivePlugin, - Arc>>, - Arc, - ) { - let table = Box::leak(Box::new(fake_table())); - let definition = Arc::new(LoadedDefinition { - _library: None, - api: NonNull::from(table), - name: "test-serving-plugin".to_string(), - }); - let events = Arc::new(Mutex::new(Vec::new())); - let abort_count = Arc::new(AtomicUsize::new(0)); - let state = Box::new(FakeState { - start_delay, - begin_fails, - events: Arc::clone(&events), - abort_count: Arc::clone(&abort_count), - }); - ( - ActivePlugin { - definition, - instance: NonNull::new(Box::into_raw(state).cast::()), - proposal_token_buffer: vec![0; MAX_NATIVE_PLUGIN_PROPOSAL_TOKENS], - }, - events, - abort_count, - ) - } - #[test] fn output_validation_is_fail_closed() { let decision = [1_u8; abi::MAX_DECISION_ID_BYTES]; @@ -1082,124 +742,4 @@ mod tests { .contains("table size") ); } - - #[test] - fn blocking_plugin_cannot_extend_the_decode_deadline() { - CANCEL_COUNT.store(0, Ordering::SeqCst); - let driver = PluginDriver::spawn(fake_active(Duration::from_millis(250))).unwrap(); - let started = Instant::now(); - let result = driver - .propose(LinearProposalQuery::new( - 1, - 2, - 16, - 16, - 0, - 8_192, - started + Duration::from_millis(5), - )) - .unwrap(); - let elapsed = started.elapsed(); - - assert!(result.is_none()); - assert!( - elapsed < Duration::from_millis(150), - "decode waited {elapsed:?} for a blocking plugin" - ); - drop(driver); - assert_eq!(CANCEL_COUNT.load(Ordering::SeqCst), 1); - } - - #[test] - fn lifecycle_ingress_shares_plugin_queue_order_with_proposals() { - let (active, events) = fake_active_with_events(Duration::ZERO); - let driver = Arc::new(PluginDriver::spawn(active).unwrap()); - let ingress = NativeLifecycleIngress { - driver: Arc::clone(&driver), - }; - ingress - .try_submit(GenerationLifecycleObservation::Started(GenerationStart { - request_id: 1, - session_id: 2, - agent_session_id: None, - prompt_token_ids: Arc::from([3]), - })) - .unwrap(); - ingress - .try_submit(GenerationLifecycleObservation::Committed( - GenerationCommit { - request_id: 1, - session_id: 2, - generated_token_count: 1, - token_ids: vec![4].into_boxed_slice(), - }, - )) - .unwrap(); - driver - .propose(LinearProposalQuery::new( - 1, - 2, - 1, - 1, - 0, - 8, - Instant::now() + Duration::from_millis(100), - )) - .unwrap(); - - assert_eq!(*events.lock().unwrap(), ["begin", "commit", "proposal"]); - } - - #[test] - fn lifecycle_callback_failure_is_observed_without_poisoning_the_driver() { - let (active, _, _) = fake_active_with_options(Duration::ZERO, true); - let driver = Arc::new(PluginDriver::spawn(active).unwrap()); - let ingress = NativeLifecycleIngress { - driver: Arc::clone(&driver), - }; - ingress - .try_submit(GenerationLifecycleObservation::Started(GenerationStart { - request_id: 7, - session_id: 9, - agent_session_id: None, - prompt_token_ids: Arc::from([3]), - })) - .unwrap(); - - driver - .propose(LinearProposalQuery::new( - 7, - 9, - 1, - 1, - 0, - 8, - Instant::now() + Duration::from_millis(100), - )) - .unwrap(); - - assert_eq!(driver.lifecycle_delivery_failures(), 1); - assert!(driver.ensure_healthy().is_ok()); - } - - #[test] - fn generation_abort_bypasses_unhealthy_driver_gate() { - let (active, _, abort_count) = fake_active_with_observations(Duration::ZERO); - let driver = Arc::new(PluginDriver::spawn(active).unwrap()); - *driver.fatal_error.lock().unwrap() = Some("report proposal failed".to_string()); - let ingress = NativeLifecycleIngress { - driver: Arc::clone(&driver), - }; - - ingress - .try_submit(GenerationLifecycleObservation::Aborted(GenerationAbort { - request_id: 7, - session_id: 9, - })) - .unwrap(); - - drop(ingress); - drop(driver); - assert_eq!(abort_count.load(Ordering::SeqCst), 1); - } } diff --git a/crates/mesh-native-serving-plugin-host/src/plugin_dispatch.rs b/crates/mesh-native-serving-plugin-host/src/plugin_dispatch.rs new file mode 100644 index 000000000..8913ca655 --- /dev/null +++ b/crates/mesh-native-serving-plugin-host/src/plugin_dispatch.rs @@ -0,0 +1,560 @@ +//! Bounded dispatch for one native serving plugin's callback threads. +//! +//! Lifecycle and proposal callbacks share one ordered queue so every proposal +//! observes the committed state that preceded it. Passive terminal callbacks +//! (`report` / `discard`) run on their own worker so they can never consume a +//! proposal's deadline. + +use std::{ + collections::VecDeque, + sync::{ + Arc, Condvar, Mutex, + atomic::{AtomicU64, Ordering}, + mpsc::{SyncSender, sync_channel}, + }, + thread::{self, JoinHandle}, + time::{Duration, Instant}, +}; + +use anyhow::{Context, Result, anyhow, bail}; +use skippy_server::frontend::{ + GenerationAbort, GenerationCommit, GenerationReceipt, GenerationStart, LinearProposal, + LinearProposalDiscardReason, LinearProposalQuery, LinearProposalReceipt, + LinearProposalSourceOutcome, LinearProposalSourceTelemetry, +}; + +use crate::ActivePlugin; + +/// Commands a caller may enqueue for either worker. +const PLUGIN_COMMAND_CAPACITY: usize = 1_024; +/// Headroom above [`PLUGIN_COMMAND_CAPACITY`] kept exclusively for terminal +/// dispositions, so a deadline discard is never lost to a full queue. +const PLUGIN_TERMINAL_RESERVE: usize = 64; +/// Bound on how long a drop waits for a worker to observe its closed queue. +const CLEAN_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); + +pub(crate) enum PluginCommand { + Begin(GenerationStart), + Committed(GenerationCommit), + Abort(GenerationAbort), + Finish(GenerationReceipt), + Proposal(LinearProposalQuery, SyncSender), + Report(LinearProposalReceipt), + Discard(Vec, LinearProposalDiscardReason), +} + +pub(crate) struct ProposalResponse { + pub(crate) proposal: std::result::Result, String>, + pub(crate) telemetry: LinearProposalSourceTelemetry, +} + +struct QueuedPluginCommand { + enqueued_at: Instant, + command: PluginCommand, +} + +struct QueueState { + commands: VecDeque, + closed: bool, +} + +/// Bounded, closable command queue owned by exactly one worker thread. +pub(crate) struct PluginCommandQueue { + state: Mutex, + available: Condvar, +} + +#[derive(Debug)] +pub(crate) enum PluginCommandQueueError { + Full, + Stopped, + Poisoned, +} + +impl PluginCommandQueue { + pub(crate) fn new() -> Self { + Self { + state: Mutex::new(QueueState { + commands: VecDeque::with_capacity(PLUGIN_COMMAND_CAPACITY), + closed: false, + }), + available: Condvar::new(), + } + } + + fn enqueue(&self, command: PluginCommand) -> Result<()> { + self.try_enqueue(command).map_err(|error| match error { + PluginCommandQueueError::Full => anyhow!("native serving plugin command queue is full"), + PluginCommandQueueError::Stopped => anyhow!("native serving plugin worker stopped"), + PluginCommandQueueError::Poisoned => { + anyhow!("native serving plugin command queue lock poisoned") + } + }) + } + + pub(crate) fn try_enqueue( + &self, + command: PluginCommand, + ) -> std::result::Result<(), PluginCommandQueueError> { + self.enqueue_within(command, PLUGIN_COMMAND_CAPACITY) + } + + /// Enqueues a terminal disposition using the reserved headroom. + /// + /// A late candidate is withheld from decode, so its `discard` is the only + /// remaining way the plugin can resolve that decision ID. It must not be + /// dropped just because ordinary traffic filled the queue. + fn try_enqueue_terminal( + &self, + command: PluginCommand, + ) -> std::result::Result<(), PluginCommandQueueError> { + self.enqueue_within(command, PLUGIN_COMMAND_CAPACITY + PLUGIN_TERMINAL_RESERVE) + } + + fn enqueue_within( + &self, + command: PluginCommand, + capacity: usize, + ) -> std::result::Result<(), PluginCommandQueueError> { + let mut state = self + .state + .lock() + .map_err(|_| PluginCommandQueueError::Poisoned)?; + if state.closed { + return Err(PluginCommandQueueError::Stopped); + } + if state.commands.len() >= capacity { + return Err(PluginCommandQueueError::Full); + } + state.commands.push_back(QueuedPluginCommand { + enqueued_at: Instant::now(), + command, + }); + self.available.notify_one(); + Ok(()) + } + + /// Refuses new work and wakes the worker so it can drain and exit. + /// + /// Closing does not require queue capacity, so a full queue can still be + /// shut down cleanly. + pub(crate) fn close(&self) { + if let Ok(mut state) = self.state.lock() { + state.closed = true; + } + self.available.notify_all(); + } + + /// Returns the next command, or `None` once the queue is closed and drained. + fn next(&self) -> Option { + let mut state = self + .state + .lock() + .expect("native serving plugin command queue lock must not be poisoned"); + loop { + // Lifecycle and proposal callbacks share this FIFO so every + // proposal observes its earlier committed state. Passive callbacks + // use their own worker and cannot delay either class. + if let Some(command) = state.commands.pop_front() { + return Some(command); + } + if state.closed { + return None; + } + state = self + .available + .wait(state) + .expect("native serving plugin command queue lock must not be poisoned"); + } + } +} + +/// One-shot notification that a worker thread has left its loop. +struct WorkerExit { + exited: Mutex, + signal: Condvar, +} + +impl WorkerExit { + fn new() -> Self { + Self { + exited: Mutex::new(false), + signal: Condvar::new(), + } + } + + fn mark_exited(&self) { + if let Ok(mut exited) = self.exited.lock() { + *exited = true; + } + self.signal.notify_all(); + } + + fn wait_for_exit(&self, timeout: Duration) -> bool { + let Ok(exited) = self.exited.lock() else { + return false; + }; + self.signal + .wait_timeout_while(exited, timeout, |exited| !*exited) + .is_ok_and(|(exited, _)| *exited) + } +} + +/// Closes the queue and signals exit even when a callback unwinds. +struct WorkerStopGuard { + queue: Arc, + exit: Arc, +} + +impl Drop for WorkerStopGuard { + fn drop(&mut self) { + self.queue.close(); + self.exit.mark_exited(); + } +} + +pub(crate) struct PluginDriver { + pub(crate) queue: Arc, + passive_queue: Arc, + active: Arc, + pub(crate) fatal_error: Arc>>, + lifecycle_delivery_failures: Arc, + worker: WorkerHandle, + passive_worker: WorkerHandle, +} + +struct WorkerHandle { + exit: Arc, + handle: Mutex>>, +} + +impl PluginDriver { + pub(crate) fn spawn(active: ActivePlugin) -> Result { + let queue = Arc::new(PluginCommandQueue::new()); + let passive_queue = Arc::new(PluginCommandQueue::new()); + let active = Arc::new(active); + let fatal_error = Arc::new(Mutex::new(None)); + let lifecycle_delivery_failures = Arc::new(AtomicU64::new(0)); + let exit = Arc::new(WorkerExit::new()); + let passive_exit = Arc::new(WorkerExit::new()); + + let worker_queue = Arc::clone(&queue); + let worker_passive_queue = Arc::clone(&passive_queue); + let worker_active = Arc::clone(&active); + let worker_failures = Arc::clone(&lifecycle_delivery_failures); + let worker_exit = Arc::clone(&exit); + let handle = thread::Builder::new() + .name("mesh-native-serving-plugin".to_string()) + .spawn(move || { + plugin_worker( + worker_active, + worker_queue, + worker_passive_queue, + worker_failures, + worker_exit, + ); + }) + .context("spawn native serving plugin worker")?; + + let passive_worker_queue = Arc::clone(&passive_queue); + let passive_worker_active = Arc::clone(&active); + let passive_worker_exit = Arc::clone(&passive_exit); + let passive_handle = thread::Builder::new() + .name("mesh-native-serving-plugin-passive".to_string()) + .spawn(move || { + plugin_passive_worker( + passive_worker_active, + passive_worker_queue, + passive_worker_exit, + ); + }) + .context("spawn native serving plugin passive worker")?; + + Ok(Self { + queue, + passive_queue, + active, + fatal_error, + lifecycle_delivery_failures, + worker: WorkerHandle { + exit, + handle: Mutex::new(Some(handle)), + }, + passive_worker: WorkerHandle { + exit: passive_exit, + handle: Mutex::new(Some(passive_handle)), + }, + }) + } + + pub(crate) fn ensure_healthy(&self) -> Result<()> { + let error = self + .fatal_error + .lock() + .map_err(|_| anyhow!("native serving plugin health lock poisoned"))?; + if let Some(error) = error.as_deref() { + bail!("native serving plugin worker failed: {error}"); + } + Ok(()) + } + + pub(crate) fn enqueue(&self, command: PluginCommand) -> Result<()> { + self.ensure_healthy()?; + self.enqueue_recovery(command) + } + + /// Deliver lifecycle cleanup even after an earlier callback failed. + /// + /// This bypasses only the health gate. The command still uses its bounded + /// callback queue and fails if that worker has stopped or is full. + pub(crate) fn enqueue_recovery(&self, command: PluginCommand) -> Result<()> { + self.queue_for(&command).enqueue(command) + } + + fn queue_for(&self, command: &PluginCommand) -> &Arc { + if matches!( + command, + PluginCommand::Report(_) | PluginCommand::Discard(_, _) + ) { + &self.passive_queue + } else { + &self.queue + } + } + + pub(crate) fn lifecycle_delivery_failures(&self) -> u64 { + self.lifecycle_delivery_failures.load(Ordering::Relaxed) + } + + pub(crate) fn propose(&self, query: LinearProposalQuery) -> Result { + self.ensure_healthy()?; + let deadline = query.deadline; + let submitted_at = Instant::now(); + if submitted_at >= deadline { + return Ok(abstention( + 0, + LinearProposalSourceOutcome::HostDeadlineExceeded, + )); + } + let (reply, response) = sync_channel(1); + match self + .queue + .try_enqueue(PluginCommand::Proposal(query, reply)) + { + Ok(()) => {} + Err(PluginCommandQueueError::Full) => { + return Ok(abstention(0, LinearProposalSourceOutcome::QueueFull)); + } + Err(PluginCommandQueueError::Stopped) => { + bail!("native serving plugin worker stopped before accepting proposal") + } + Err(PluginCommandQueueError::Poisoned) => { + bail!("native serving plugin command queue lock poisoned") + } + } + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return Ok(abstention( + elapsed_us(submitted_at), + LinearProposalSourceOutcome::HostDeadlineExceeded, + )); + } + match response.recv_timeout(remaining) { + Ok(result) => Ok(result), + Err(std::sync::mpsc::RecvTimeoutError::Timeout) => Ok(abstention( + elapsed_us(submitted_at), + LinearProposalSourceOutcome::HostDeadlineExceeded, + )), + Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => { + bail!("native serving plugin worker stopped before replying") + } + } + } +} + +fn abstention(queue_wait_us: u64, outcome: LinearProposalSourceOutcome) -> ProposalResponse { + ProposalResponse { + proposal: Ok(None), + telemetry: LinearProposalSourceTelemetry { + queue_wait_us, + callback_elapsed_us: 0, + outcome, + }, + } +} + +fn elapsed_us(started: Instant) -> u64 { + u64::try_from(started.elapsed().as_micros()).unwrap_or(u64::MAX) +} + +impl Drop for PluginDriver { + fn drop(&mut self) { + stop_worker(&self.queue, &self.worker, "callback"); + stop_worker(&self.passive_queue, &self.passive_worker, "passive"); + if let Some(active) = Arc::get_mut(&mut self.active) + && let Err(error) = active.shutdown() + { + eprintln!("native serving plugin shutdown failed: {error:#}"); + } + } +} + +/// Closes a worker's queue and joins it once it has actually left its loop. +/// +/// Closing never needs queue capacity, and the exit signal is set by the +/// worker's stop guard, so neither a full queue nor the gap between the last +/// callback and thread teardown can skip the join. +fn stop_worker(queue: &Arc, worker: &WorkerHandle, label: &str) { + queue.close(); + if !worker.exit.wait_for_exit(CLEAN_SHUTDOWN_TIMEOUT) { + eprintln!( + "native serving plugin {label} worker did not stop within {CLEAN_SHUTDOWN_TIMEOUT:?}; \ + deferring plugin shutdown to that thread" + ); + return; + } + if let Ok(mut handle) = worker.handle.lock() + && let Some(handle) = handle.take() + { + let _ = handle.join(); + } +} + +fn plugin_worker( + active: Arc, + queue: Arc, + passive_queue: Arc, + lifecycle_delivery_failures: Arc, + exit: Arc, +) { + let _stop_guard = WorkerStopGuard { + queue: Arc::clone(&queue), + exit, + }; + while let Some(QueuedPluginCommand { + enqueued_at, + command, + }) = queue.next() + { + let (result, lifecycle) = match command { + PluginCommand::Begin(event) => (active.begin(&event), true), + PluginCommand::Committed(event) => (active.committed(&event), true), + PluginCommand::Abort(event) => (active.abort(&event), true), + PluginCommand::Finish(event) => (active.finish(&event), true), + PluginCommand::Proposal(query, reply) => { + run_proposal(&active, &passive_queue, enqueued_at, query, &reply); + continue; + } + PluginCommand::Report(_) | PluginCommand::Discard(_, _) => { + unreachable!("passive plugin callbacks must use the passive worker queue") + } + }; + if lifecycle && result.is_err() { + lifecycle_delivery_failures.fetch_add(1, Ordering::Relaxed); + } + } +} + +fn run_proposal( + active: &ActivePlugin, + passive_queue: &PluginCommandQueue, + enqueued_at: Instant, + query: LinearProposalQuery, + reply: &SyncSender, +) { + let deadline = query.deadline; + let queue_wait_us = elapsed_us(enqueued_at); + if Instant::now() >= deadline { + let _ = reply.send(abstention( + queue_wait_us, + LinearProposalSourceOutcome::DeadlineExceededBeforeDispatch, + )); + return; + } + + let callback_started = Instant::now(); + let result = active.propose(query); + // One timestamp classifies both the forwarding decision and the telemetry, + // so a candidate can never be forwarded while telemetry calls it late. + let callback_finished = Instant::now(); + let callback_elapsed_us = elapsed_us(callback_started); + let deadline_missed = callback_finished >= deadline; + + let (proposal, outcome) = match result { + Ok(Some(proposal)) if deadline_missed => { + discard_late_candidate(passive_queue, &proposal); + ( + Ok(None), + LinearProposalSourceOutcome::CandidateReturnedTooLate, + ) + } + Ok(Some(proposal)) => (Ok(Some(proposal)), LinearProposalSourceOutcome::Ready), + Ok(None) if deadline_missed => ( + Ok(None), + LinearProposalSourceOutcome::DeadlineExceededInPlugin, + ), + Ok(None) => (Ok(None), LinearProposalSourceOutcome::Abstained), + Err(error) => { + // Fail open, but surface the plugin's message instead of + // silently degrading a failure into an abstention. + eprintln!("native serving plugin proposal failed: {error:#}"); + ( + Err(format!("{error:#}")), + LinearProposalSourceOutcome::SourceError, + ) + } + }; + + let _ = reply.send(ProposalResponse { + proposal, + telemetry: LinearProposalSourceTelemetry { + queue_wait_us, + callback_elapsed_us, + outcome, + }, + }); +} + +/// Resolves a decision the host refuses to forward after its deadline. +fn discard_late_candidate(passive_queue: &PluginCommandQueue, proposal: &LinearProposal) { + if let Err(error) = passive_queue.try_enqueue_terminal(PluginCommand::Discard( + proposal.decision_id.as_bytes().to_vec(), + LinearProposalDiscardReason::DeadlineExceeded, + )) { + eprintln!( + "native serving plugin could not deliver the terminal discard for a late proposal: \ + {error:?}" + ); + } +} + +fn plugin_passive_worker( + active: Arc, + queue: Arc, + exit: Arc, +) { + let _stop_guard = WorkerStopGuard { + queue: Arc::clone(&queue), + exit, + }; + while let Some(queued) = queue.next() { + match queued.command { + PluginCommand::Report(event) => { + let _ = active.report(&event); + } + PluginCommand::Discard(decision_id, reason) => { + let _ = active.discard(&decision_id, reason); + } + PluginCommand::Begin(_) + | PluginCommand::Committed(_) + | PluginCommand::Abort(_) + | PluginCommand::Finish(_) + | PluginCommand::Proposal(_, _) => { + unreachable!("lifecycle and proposal callbacks must use the primary worker queue") + } + } + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/mesh-native-serving-plugin-host/src/plugin_dispatch/tests.rs b/crates/mesh-native-serving-plugin-host/src/plugin_dispatch/tests.rs new file mode 100644 index 000000000..a5c3de06e --- /dev/null +++ b/crates/mesh-native-serving-plugin-host/src/plugin_dispatch/tests.rs @@ -0,0 +1,423 @@ +//! Deadline, ordering, and shutdown behavior of the plugin dispatch workers. + +use std::{ + sync::{Arc, atomic::Ordering, mpsc::sync_channel}, + time::{Duration, Instant}, +}; + +use skippy_server::frontend::{ + GenerationAbort, GenerationCommit, GenerationLifecycleIngress, GenerationLifecycleObservation, + GenerationStart, LinearProposalDiscardReason, LinearProposalQuery, LinearProposalSourceOutcome, +}; + +use super::{PluginCommand, PluginCommandQueue, PluginCommandQueueError, PluginDriver}; +use crate::{ + NativeLifecycleIngress, + test_support::{ + fake_active, fake_active_with_events, fake_active_with_failing_proposal, + fake_active_with_late_candidate, fake_active_with_observations, fake_active_with_options, + fake_active_with_timing, fake_observations, proposal_query, wait_for_event, + }, +}; + +#[test] +fn blocking_plugin_cannot_extend_the_decode_deadline() { + let active = fake_active(Duration::from_millis(250)); + let observations = fake_observations(&active); + let driver = PluginDriver::spawn(active).unwrap(); + let started = Instant::now(); + let result = driver + .propose(LinearProposalQuery::new( + 1, + 2, + 16, + 16, + 0, + 8_192, + started + Duration::from_millis(5), + )) + .unwrap(); + let elapsed = started.elapsed(); + + assert!(result.proposal.unwrap().is_none()); + assert_eq!( + result.telemetry.outcome, + LinearProposalSourceOutcome::HostDeadlineExceeded + ); + assert!( + elapsed < Duration::from_millis(150), + "decode waited {elapsed:?} for a blocking plugin" + ); + drop(driver); + assert_eq!(observations.cancel_count.load(Ordering::SeqCst), 1); +} + +#[test] +fn slow_commit_abstains_before_plugin_dispatch_and_later_positions_recover() { + let (active, events, _) = fake_active_with_timing( + Duration::ZERO, + Duration::from_millis(40), + Duration::ZERO, + false, + ); + let driver = Arc::new(PluginDriver::spawn(active).unwrap()); + let ingress = NativeLifecycleIngress { + driver: Arc::clone(&driver), + }; + ingress + .try_submit(GenerationLifecycleObservation::Committed( + GenerationCommit { + request_id: 1, + session_id: 2, + generated_token_count: 1, + token_ids: vec![4].into_boxed_slice(), + }, + )) + .unwrap(); + wait_for_event(events.as_ref(), "commit"); + + let started = Instant::now(); + let missed = driver + .propose(proposal_query(started + Duration::from_millis(5))) + .unwrap(); + assert!(missed.proposal.unwrap().is_none()); + assert_eq!( + missed.telemetry.outcome, + LinearProposalSourceOutcome::HostDeadlineExceeded + ); + assert!( + started.elapsed() < Duration::from_millis(30), + "proposal wait exceeded its deadline: {:?}", + started.elapsed() + ); + + let recovered = driver + .propose(proposal_query(Instant::now() + Duration::from_millis(100))) + .unwrap(); + assert!(recovered.proposal.unwrap().is_none()); + assert_eq!( + recovered.telemetry.outcome, + LinearProposalSourceOutcome::Abstained + ); + assert_eq!(*events.lock().unwrap(), ["commit", "proposal"]); +} + +#[test] +fn running_passive_discard_cannot_delay_the_next_proposal() { + let (active, events, _) = fake_active_with_timing( + Duration::ZERO, + Duration::ZERO, + Duration::from_millis(100), + false, + ); + let driver = Arc::new(PluginDriver::spawn(active).unwrap()); + driver + .enqueue(PluginCommand::Discard( + vec![1], + LinearProposalDiscardReason::PositionMismatch, + )) + .unwrap(); + wait_for_event(events.as_ref(), "discard"); + + let started = Instant::now(); + let response = driver + .propose(proposal_query(started + Duration::from_millis(20))) + .unwrap(); + assert!(response.proposal.unwrap().is_none()); + assert_eq!( + response.telemetry.outcome, + LinearProposalSourceOutcome::Abstained + ); + assert!(started.elapsed() < Duration::from_millis(60)); + assert_eq!(*events.lock().unwrap(), ["discard", "proposal"]); +} + +#[test] +fn worker_reports_pre_dispatch_deadlines_without_running_the_callback() { + let (active, events, _) = fake_active_with_timing( + Duration::ZERO, + Duration::from_millis(40), + Duration::ZERO, + false, + ); + let driver = PluginDriver::spawn(active).unwrap(); + driver + .enqueue(PluginCommand::Committed(GenerationCommit { + request_id: 1, + session_id: 2, + generated_token_count: 1, + token_ids: vec![4].into_boxed_slice(), + })) + .unwrap(); + wait_for_event(events.as_ref(), "commit"); + let (reply, response) = sync_channel(1); + driver + .queue + .try_enqueue(PluginCommand::Proposal( + proposal_query(Instant::now() + Duration::from_millis(5)), + reply, + )) + .unwrap(); + + let response = response.recv_timeout(Duration::from_millis(100)).unwrap(); + assert!(response.proposal.unwrap().is_none()); + assert_eq!( + response.telemetry.outcome, + LinearProposalSourceOutcome::DeadlineExceededBeforeDispatch + ); + assert_eq!(*events.lock().unwrap(), ["commit"]); +} + +#[test] +fn late_candidate_is_reported_and_not_forwarded_to_the_decode() { + let driver = + PluginDriver::spawn(fake_active_with_late_candidate(Duration::from_millis(20))).unwrap(); + let (reply, response) = sync_channel(1); + driver + .queue + .try_enqueue(PluginCommand::Proposal( + proposal_query(Instant::now() + Duration::from_millis(5)), + reply, + )) + .unwrap(); + + let response = response.recv_timeout(Duration::from_millis(100)).unwrap(); + assert!(response.proposal.unwrap().is_none()); + assert_eq!( + response.telemetry.outcome, + LinearProposalSourceOutcome::CandidateReturnedTooLate + ); +} + +#[test] +fn stopped_worker_rejects_lifecycle_delivery() { + let queue = PluginCommandQueue::new(); + queue.close(); + + assert!(matches!( + queue.try_enqueue(PluginCommand::Abort(GenerationAbort { + request_id: 1, + session_id: 2, + })), + Err(PluginCommandQueueError::Stopped) + )); +} + +#[test] +fn lifecycle_ingress_shares_plugin_queue_order_with_proposals() { + let (active, events) = fake_active_with_events(Duration::ZERO); + let driver = Arc::new(PluginDriver::spawn(active).unwrap()); + let ingress = NativeLifecycleIngress { + driver: Arc::clone(&driver), + }; + ingress + .try_submit(GenerationLifecycleObservation::Started(GenerationStart { + request_id: 1, + session_id: 2, + agent_session_id: None, + prompt_token_ids: Arc::from([3]), + })) + .unwrap(); + ingress + .try_submit(GenerationLifecycleObservation::Committed( + GenerationCommit { + request_id: 1, + session_id: 2, + generated_token_count: 1, + token_ids: vec![4].into_boxed_slice(), + }, + )) + .unwrap(); + driver + .propose(LinearProposalQuery::new( + 1, + 2, + 1, + 1, + 0, + 8, + Instant::now() + Duration::from_millis(100), + )) + .unwrap(); + + assert_eq!(*events.lock().unwrap(), ["begin", "commit", "proposal"]); +} + +#[test] +fn lifecycle_callback_failure_is_observed_without_poisoning_the_driver() { + let (active, _, _) = fake_active_with_options(Duration::ZERO, true); + let driver = Arc::new(PluginDriver::spawn(active).unwrap()); + let ingress = NativeLifecycleIngress { + driver: Arc::clone(&driver), + }; + ingress + .try_submit(GenerationLifecycleObservation::Started(GenerationStart { + request_id: 7, + session_id: 9, + agent_session_id: None, + prompt_token_ids: Arc::from([3]), + })) + .unwrap(); + + driver + .propose(LinearProposalQuery::new( + 7, + 9, + 1, + 1, + 0, + 8, + Instant::now() + Duration::from_millis(100), + )) + .unwrap(); + + assert_eq!(driver.lifecycle_delivery_failures(), 1); + assert!(driver.ensure_healthy().is_ok()); +} + +#[test] +fn generation_abort_bypasses_unhealthy_driver_gate() { + let (active, _, abort_count) = fake_active_with_observations(Duration::ZERO); + let driver = Arc::new(PluginDriver::spawn(active).unwrap()); + *driver.fatal_error.lock().unwrap() = Some("report proposal failed".to_string()); + let ingress = NativeLifecycleIngress { + driver: Arc::clone(&driver), + }; + + ingress + .try_submit(GenerationLifecycleObservation::Aborted(GenerationAbort { + request_id: 7, + session_id: 9, + })) + .unwrap(); + + drop(ingress); + drop(driver); + assert_eq!(abort_count.load(Ordering::SeqCst), 1); +} + +#[test] +fn a_full_queue_can_still_be_closed_and_drained() { + // Closing takes no queue capacity, so a saturated queue can never strand + // its worker the way a queued shutdown command could. + let queue = PluginCommandQueue::new(); + let mut queued = 0_usize; + while queue + .try_enqueue(PluginCommand::Abort(GenerationAbort { + request_id: 1, + session_id: 2, + })) + .is_ok() + { + queued += 1; + assert!(queued < 10_000, "queue never reported full"); + } + + queue.close(); + for _ in 0..queued { + assert!(queue.next().is_some(), "close must not discard queued work"); + } + assert!( + queue.next().is_none(), + "a drained, closed queue must stop its worker" + ); +} + +#[test] +fn dropping_the_driver_drains_its_backlog_and_shuts_the_plugin_down() { + let (active, events, _) = fake_active_with_timing( + Duration::ZERO, + Duration::from_millis(2), + Duration::ZERO, + false, + ); + let observations = fake_observations(&active); + let driver = PluginDriver::spawn(active).unwrap(); + for _ in 0..20 { + driver + .enqueue(PluginCommand::Committed(GenerationCommit { + request_id: 1, + session_id: 2, + generated_token_count: 1, + token_ids: vec![4].into_boxed_slice(), + })) + .unwrap(); + } + + drop(driver); + assert_eq!( + events.lock().unwrap().len(), + 20, + "backlog must be delivered" + ); + assert_eq!(observations.shutdown_count.load(Ordering::SeqCst), 1); +} + +#[test] +fn a_full_passive_queue_still_accepts_the_terminal_discard() { + // Reserved headroom keeps a deadline discard deliverable, so the plugin + // always learns the fate of a decision the host withheld. + let queue = PluginCommandQueue::new(); + while queue + .try_enqueue(PluginCommand::Discard( + vec![1], + LinearProposalDiscardReason::PositionMismatch, + )) + .is_ok() + {} + + assert!(matches!( + queue.try_enqueue(PluginCommand::Discard( + vec![2], + LinearProposalDiscardReason::PositionMismatch, + )), + Err(PluginCommandQueueError::Full) + )); + queue + .try_enqueue_terminal(PluginCommand::Discard( + vec![3], + LinearProposalDiscardReason::DeadlineExceeded, + )) + .expect("terminal discard must use the reserved headroom"); +} + +#[test] +fn a_late_candidate_delivers_its_discard_to_the_plugin() { + // The host withholds a late candidate, so `discard` is the only way the + // plugin can learn that decision's fate. + let active = fake_active_with_late_candidate(Duration::from_millis(20)); + let observations = fake_observations(&active); + let driver = PluginDriver::spawn(active).unwrap(); + let (reply, response) = sync_channel(1); + driver + .queue + .try_enqueue(PluginCommand::Proposal( + proposal_query(Instant::now() + Duration::from_millis(5)), + reply, + )) + .unwrap(); + + let response = response.recv_timeout(Duration::from_millis(500)).unwrap(); + assert!(response.proposal.unwrap().is_none()); + assert_eq!( + response.telemetry.outcome, + LinearProposalSourceOutcome::CandidateReturnedTooLate + ); + wait_for_event(observations.events.as_ref(), "discard"); +} + +#[test] +fn a_failing_proposal_callback_reports_source_error_instead_of_abstaining() { + // Fail open for decode, but keep the plugin's failure distinguishable + // from a deliberate abstention. + let driver = PluginDriver::spawn(fake_active_with_failing_proposal()).unwrap(); + let response = driver + .propose(proposal_query(Instant::now() + Duration::from_millis(200))) + .unwrap(); + + assert!(response.proposal.is_err()); + assert_eq!( + response.telemetry.outcome, + LinearProposalSourceOutcome::SourceError + ); +} diff --git a/crates/mesh-native-serving-plugin-host/src/test_support.rs b/crates/mesh-native-serving-plugin-host/src/test_support.rs new file mode 100644 index 000000000..791a4831e --- /dev/null +++ b/crates/mesh-native-serving-plugin-host/src/test_support.rs @@ -0,0 +1,311 @@ +//! Shared fake native plugin used by host and dispatch unit tests. + +use std::{ + ffi::{c_char, c_void}, + ptr::NonNull, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + thread, + time::{Duration, Instant}, +}; + +use mesh_native_serving_plugin_api as abi; +use skippy_server::frontend::LinearProposalQuery; + +use crate::{ActivePlugin, LoadedDefinition, MAX_NATIVE_PLUGIN_PROPOSAL_TOKENS}; + +pub(crate) static FAKE_NAME: &[u8] = b"test-serving-plugin"; + +/// Per-instance observations, so tests running in parallel cannot see each +/// other's callback counts. +pub(crate) struct FakeObservations { + pub(crate) events: Arc>>, + pub(crate) cancel_count: Arc, + pub(crate) shutdown_count: Arc, +} + +pub(crate) struct FakeState { + start_delay: Duration, + poll_delay: Duration, + poll_returns_candidate: bool, + poll_fails: bool, + commit_delay: Duration, + report_delay: Duration, + begin_fails: bool, + events: Arc>>, + abort_count: Arc, + cancel_count: Arc, + shutdown_count: Arc, +} + +unsafe extern "C" fn fake_activate( + _context: *const abi::ActivationContext, + _activation: *mut abi::PluginActivation, +) -> abi::PluginStatus { + abi::PluginStatus::INTERNAL_ERROR +} + +unsafe extern "C" fn fake_shutdown(instance: abi::PluginInstance) -> abi::PluginStatus { + if !instance.is_null() { + let state = unsafe { Box::from_raw(instance.cast::()) }; + state.shutdown_count.fetch_add(1, Ordering::SeqCst); + } + abi::PluginStatus::OK +} + +unsafe extern "C" fn fake_begin( + instance: abi::PluginInstance, + _event: *const abi::GenerationStart, +) -> abi::PluginStatus { + let state = unsafe { &*instance.cast::() }; + state.events.lock().unwrap().push("begin"); + if state.begin_fails { + abi::PluginStatus::INTERNAL_ERROR + } else { + abi::PluginStatus::OK + } +} + +unsafe extern "C" fn fake_commit( + instance: abi::PluginInstance, + _event: *const abi::GenerationCommit, +) -> abi::PluginStatus { + let state = unsafe { &*instance.cast::() }; + state.events.lock().unwrap().push("commit"); + thread::sleep(state.commit_delay); + abi::PluginStatus::OK +} + +unsafe extern "C" fn fake_abort( + instance: abi::PluginInstance, + _event: *const abi::GenerationAbort, +) -> abi::PluginStatus { + let state = unsafe { &*instance.cast::() }; + state.abort_count.fetch_add(1, Ordering::SeqCst); + abi::PluginStatus::OK +} + +unsafe extern "C" fn fake_finish( + _instance: abi::PluginInstance, + _event: *const abi::GenerationFinish, +) -> abi::PluginStatus { + abi::PluginStatus::OK +} + +unsafe extern "C" fn fake_start_proposal( + instance: abi::PluginInstance, + _query: *const abi::ProposalQuery, + operation: *mut abi::ProposalOperation, +) -> abi::PluginStatus { + let state = unsafe { &*instance.cast::() }; + state.events.lock().unwrap().push("proposal"); + thread::sleep(state.start_delay); + unsafe { *operation = 1 }; + abi::PluginStatus::OK +} + +unsafe extern "C" fn fake_poll_proposal( + instance: abi::PluginInstance, + _operation: abi::ProposalOperation, + output: *mut abi::ProposalOutput, +) -> abi::ProposalPollStatus { + let state = unsafe { &*instance.cast::() }; + thread::sleep(state.poll_delay); + if state.poll_fails { + return abi::ProposalPollStatus::FAILED; + } + if !state.poll_returns_candidate { + return abi::ProposalPollStatus::ABSTAIN; + } + let output = unsafe { &mut *output }; + unsafe { + *output.decision_id = 7; + *output.token_ids = 42; + } + output.decision_id_length = 1; + output.token_length = 1; + abi::ProposalPollStatus::READY +} + +unsafe extern "C" fn fake_cancel_proposal( + instance: abi::PluginInstance, + _operation: abi::ProposalOperation, +) { + if !instance.is_null() { + let state = unsafe { &*instance.cast::() }; + state.cancel_count.fetch_add(1, Ordering::SeqCst); + } +} + +unsafe extern "C" fn fake_report_proposal( + instance: abi::PluginInstance, + _event: *const abi::ProposalOutcome, +) -> abi::PluginStatus { + let state = unsafe { &*instance.cast::() }; + state.events.lock().unwrap().push("report"); + thread::sleep(state.report_delay); + abi::PluginStatus::OK +} + +unsafe extern "C" fn fake_discard_proposal( + instance: abi::PluginInstance, + _event: *const abi::ProposalDiscard, +) -> abi::PluginStatus { + let state = unsafe { &*instance.cast::() }; + state.events.lock().unwrap().push("discard"); + thread::sleep(state.report_delay); + abi::PluginStatus::OK +} + +unsafe extern "C" fn fake_last_error( + _instance: abi::PluginInstance, + _output: *mut c_char, + _capacity: usize, +) -> usize { + 0 +} + +pub(crate) fn fake_table() -> abi::NativeServingPluginV1 { + abi::NativeServingPluginV1 { + abi_version: abi::NATIVE_SERVING_PLUGIN_ABI_V1, + struct_size: size_of::(), + plugin_name: abi::ByteSlice::from_bytes(FAKE_NAME), + activate: fake_activate, + shutdown: fake_shutdown, + begin_generation: fake_begin, + commit_generation: fake_commit, + abort_generation: fake_abort, + finish_generation: fake_finish, + start_proposal: fake_start_proposal, + poll_proposal: fake_poll_proposal, + cancel_proposal: fake_cancel_proposal, + report_proposal: fake_report_proposal, + discard_proposal: fake_discard_proposal, + last_error: fake_last_error, + } +} + +pub(crate) fn fake_active(start_delay: Duration) -> ActivePlugin { + fake_active_with_events(start_delay).0 +} + +pub(crate) fn fake_active_with_events( + start_delay: Duration, +) -> (ActivePlugin, Arc>>) { + let (active, events, _) = fake_active_with_observations(start_delay); + (active, events) +} + +pub(crate) fn fake_active_with_observations( + start_delay: Duration, +) -> ( + ActivePlugin, + Arc>>, + Arc, +) { + fake_active_with_options(start_delay, false) +} + +pub(crate) fn fake_active_with_options( + start_delay: Duration, + begin_fails: bool, +) -> ( + ActivePlugin, + Arc>>, + Arc, +) { + fake_active_with_timing(start_delay, Duration::ZERO, Duration::ZERO, begin_fails) +} + +pub(crate) fn fake_active_with_timing( + start_delay: Duration, + commit_delay: Duration, + report_delay: Duration, + begin_fails: bool, +) -> ( + ActivePlugin, + Arc>>, + Arc, +) { + let table = Box::leak(Box::new(fake_table())); + let definition = Arc::new(LoadedDefinition { + _library: None, + api: NonNull::from(table), + name: "test-serving-plugin".to_string(), + }); + let events = Arc::new(Mutex::new(Vec::new())); + let abort_count = Arc::new(AtomicUsize::new(0)); + let cancel_count = Arc::new(AtomicUsize::new(0)); + let shutdown_count = Arc::new(AtomicUsize::new(0)); + let state = Box::new(FakeState { + start_delay, + poll_delay: Duration::ZERO, + poll_returns_candidate: false, + poll_fails: false, + commit_delay, + report_delay, + begin_fails, + events: Arc::clone(&events), + abort_count: Arc::clone(&abort_count), + cancel_count: Arc::clone(&cancel_count), + shutdown_count: Arc::clone(&shutdown_count), + }); + ( + ActivePlugin { + definition, + instance: NonNull::new(Box::into_raw(state).cast::()), + proposal_token_buffer: Mutex::new(vec![0; MAX_NATIVE_PLUGIN_PROPOSAL_TOKENS]), + }, + events, + abort_count, + ) +} + +pub(crate) fn fake_active_with_late_candidate(poll_delay: Duration) -> ActivePlugin { + let (active, _, _) = + fake_active_with_timing(Duration::ZERO, Duration::ZERO, Duration::ZERO, false); + let instance = active.instance.unwrap().as_ptr().cast::(); + unsafe { + (*instance).poll_delay = poll_delay; + (*instance).poll_returns_candidate = true; + } + active +} + +/// Clones the per-instance observation handles out of a fake plugin. +/// +/// The handles outlive the plugin, so a test can still assert on shutdown +/// after the driver has dropped it. +pub(crate) fn fake_observations(active: &ActivePlugin) -> FakeObservations { + let state = unsafe { &*active.instance.unwrap().as_ptr().cast::() }; + FakeObservations { + events: Arc::clone(&state.events), + cancel_count: Arc::clone(&state.cancel_count), + shutdown_count: Arc::clone(&state.shutdown_count), + } +} + +/// Builds a plugin whose proposal callback fails, exercising the fail-open path. +pub(crate) fn fake_active_with_failing_proposal() -> ActivePlugin { + let (active, _, _) = + fake_active_with_timing(Duration::ZERO, Duration::ZERO, Duration::ZERO, false); + let instance = active.instance.unwrap().as_ptr().cast::(); + unsafe { + (*instance).poll_fails = true; + } + active +} + +pub(crate) fn wait_for_event(events: &Mutex>, event: &str) { + let deadline = Instant::now() + Duration::from_secs(1); + while !events.lock().unwrap().contains(&event) { + assert!(Instant::now() < deadline, "timed out waiting for {event}"); + thread::yield_now(); + } +} + +pub(crate) fn proposal_query(deadline: Instant) -> LinearProposalQuery { + LinearProposalQuery::new(1, 2, 1, 1, 0, 8, deadline) +} diff --git a/crates/skippy-server/src/frontend.rs b/crates/skippy-server/src/frontend.rs index 4a6c339aa..dc1ab1281 100644 --- a/crates/skippy-server/src/frontend.rs +++ b/crates/skippy-server/src/frontend.rs @@ -48,6 +48,7 @@ pub use self::guardrails::{ pub use self::linear_proposal::{ LinearProposal, LinearProposalDiscardReason, LinearProposalDisposition, LinearProposalIngress, LinearProposalIngressConfig, LinearProposalQuery, LinearProposalReceipt, + LinearProposalSourceOutcome, LinearProposalSourceResponse, LinearProposalSourceTelemetry, OpaqueProposalDecisionId, }; pub use self::speculative::{ diff --git a/crates/skippy-server/src/frontend/linear_proposal.rs b/crates/skippy-server/src/frontend/linear_proposal.rs index ee0702f8a..b202431f1 100644 --- a/crates/skippy-server/src/frontend/linear_proposal.rs +++ b/crates/skippy-server/src/frontend/linear_proposal.rs @@ -279,14 +279,56 @@ impl LinearProposalReceipt { } } +/// Per-request result from an in-process linear proposal source. +/// +/// The telemetry travels with the corresponding request result so a shared +/// source cannot accidentally attribute timings from one decode to another. +pub struct LinearProposalSourceResponse { + proposal: Option, + telemetry: Option, +} + +impl LinearProposalSourceResponse { + /// Creates a result without source-specific timing data. + #[must_use] + pub fn new(proposal: Option) -> Self { + Self { + proposal, + telemetry: None, + } + } + + /// Creates a result with telemetry for this exact proposal request. + #[must_use] + pub fn with_telemetry( + proposal: Option, + telemetry: LinearProposalSourceTelemetry, + ) -> Self { + Self { + proposal, + telemetry: Some(telemetry), + } + } + + fn into_parts( + self, + ) -> ( + Option, + Option, + ) { + (self.proposal, self.telemetry) + } +} + /// In-process, source-neutral width-one proposal boundary. /// /// Implementations must honor `query.deadline`. Skippy independently rejects a /// proposal that arrives after it and calls `discard` so the source can resolve /// any pending decision without treating it as verified. pub trait LinearProposalIngress: Send + Sync { - /// Returns an optional bounded proposal for the committed query state. - fn propose(&self, query: LinearProposalQuery) -> Result>; + /// Returns an optional bounded proposal and telemetry for this exact + /// committed query state. + fn propose(&self, query: LinearProposalQuery) -> Result; /// Receives the target-authoritative outcome for a verified proposal. fn report(&self, receipt: &LinearProposalReceipt) -> Result<()>; @@ -301,6 +343,71 @@ pub trait LinearProposalIngress: Send + Sync { } } +/// Bounded source-side outcome for one proposal query. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub enum LinearProposalSourceOutcome { + /// The source returned before the configured deadline. + Ready, + /// The source deliberately returned no proposal before the deadline. + Abstained, + /// The host stopped waiting at the configured wall-clock deadline. + HostDeadlineExceeded, + /// The host queue was at capacity, so the source was never submitted. + QueueFull, + /// The worker established that the deadline had elapsed before dispatch. + DeadlineExceededBeforeDispatch, + /// A plugin callback exhausted the deadline without producing a candidate. + DeadlineExceededInPlugin, + /// A plugin callback returned a candidate after the deadline. + CandidateReturnedTooLate, + /// A plugin callback failed and was treated as a fail-open abstention. + SourceError, +} + +impl LinearProposalSourceOutcome { + const fn as_str(self) -> &'static str { + match self { + Self::Ready => "ready", + Self::Abstained => "abstained", + Self::HostDeadlineExceeded => "host_deadline_exceeded", + Self::QueueFull => "queue_full", + Self::DeadlineExceededBeforeDispatch => "deadline_exceeded_before_dispatch", + Self::DeadlineExceededInPlugin => "deadline_exceeded_in_plugin", + Self::CandidateReturnedTooLate => "candidate_returned_too_late", + Self::SourceError => "source_error", + } + } +} + +/// Privacy-safe source timing for one proposal query. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct LinearProposalSourceTelemetry { + /// Time from source submission until the worker began dispatch. + pub queue_wait_us: u64, + /// Time spent crossing the source's callback boundary. + pub callback_elapsed_us: u64, + /// Bounded completion or abstention outcome. + pub outcome: LinearProposalSourceOutcome, +} + +impl LinearProposalSourceTelemetry { + pub(crate) fn insert_telemetry_attrs(self, attrs: &mut BTreeMap) { + attrs.insert( + "llama_stage.linear_proposal.source_queue_wait_us".to_string(), + json!(self.queue_wait_us), + ); + attrs.insert( + "llama_stage.linear_proposal.source_callback_us".to_string(), + json!(self.callback_elapsed_us), + ); + attrs.insert( + "llama_stage.linear_proposal.source_outcome".to_string(), + json!(self.outcome.as_str()), + ); + } +} + #[derive(Clone)] pub struct LinearProposalIngressConfig { source: Arc, @@ -353,11 +460,17 @@ pub(crate) struct QueriedLinearProposal { pub(crate) proposal: LinearProposal, pub(crate) proposal_elapsed_us: u64, pub(crate) operation_started: Instant, + pub(crate) source_telemetry: Option, } pub(crate) enum LinearProposalQueryOutcome { - NoProposal, - DeadlineExceeded { proposal_elapsed_us: u64 }, + NoProposal { + source_telemetry: Option, + }, + DeadlineExceeded { + proposal_elapsed_us: u64, + source_telemetry: Option, + }, Ready(QueriedLinearProposal), } @@ -390,14 +503,16 @@ pub(crate) fn query_linear_proposal( .min(params.runtime_max_proposal_tokens) .min(config.max_proposal_tokens()); if max_proposal_tokens == 0 { - return Ok(LinearProposalQueryOutcome::NoProposal); + return Ok(LinearProposalQueryOutcome::NoProposal { + source_telemetry: None, + }); } let operation_started = Instant::now(); let deadline = operation_started .checked_add(config.deadline()) .ok_or_else(|| OpenAiError::backend("linear proposal deadline overflow"))?; let proposal_started = Instant::now(); - let proposal = config + let response = config .source() .propose(LinearProposalQuery::new( params.request_id, @@ -410,8 +525,9 @@ pub(crate) fn query_linear_proposal( )) .map_err(openai_backend_error)?; let proposal_elapsed_us = elapsed_us(proposal_started); + let (proposal, source_telemetry) = response.into_parts(); let Some(proposal) = proposal else { - return Ok(LinearProposalQueryOutcome::NoProposal); + return Ok(LinearProposalQueryOutcome::NoProposal { source_telemetry }); }; if Instant::now() > deadline { config @@ -423,6 +539,7 @@ pub(crate) fn query_linear_proposal( .map_err(openai_backend_error)?; return Ok(LinearProposalQueryOutcome::DeadlineExceeded { proposal_elapsed_us, + source_telemetry, }); } if proposal.token_ids.is_empty() || proposal.token_ids.len() > max_proposal_tokens { @@ -433,7 +550,7 @@ pub(crate) fn query_linear_proposal( LinearProposalDiscardReason::InvalidTokenCount, ) .map_err(openai_backend_error)?; - return Ok(LinearProposalQueryOutcome::NoProposal); + return Ok(LinearProposalQueryOutcome::NoProposal { source_telemetry }); } if proposal.token_ids.iter().any(|token| *token < 0) { config @@ -443,12 +560,13 @@ pub(crate) fn query_linear_proposal( LinearProposalDiscardReason::InvalidTokenId, ) .map_err(openai_backend_error)?; - return Ok(LinearProposalQueryOutcome::NoProposal); + return Ok(LinearProposalQueryOutcome::NoProposal { source_telemetry }); } Ok(LinearProposalQueryOutcome::Ready(QueriedLinearProposal { proposal, proposal_elapsed_us, operation_started, + source_telemetry, })) } @@ -549,7 +667,7 @@ mod tests { } impl LinearProposalIngress for FakeIngress { - fn propose(&self, query: LinearProposalQuery) -> Result> { + fn propose(&self, query: LinearProposalQuery) -> Result { self.queries.lock().unwrap().push(RecordedQuery { request_id: query.request_id, session_id: query.session_id, @@ -559,7 +677,9 @@ mod tests { max_proposal_tokens: query.max_proposal_tokens, }); thread::sleep(*self.delay.lock().unwrap()); - Ok(self.proposal.lock().unwrap().take()) + Ok(LinearProposalSourceResponse::new( + self.proposal.lock().unwrap().take(), + )) } fn report(&self, receipt: &LinearProposalReceipt) -> Result<()> { @@ -804,7 +924,7 @@ mod tests { .unwrap(); assert!(matches!( query_linear_proposal(&invalid_config, query_params(1, 2, 1, 0, 1, 5, 4)).unwrap(), - LinearProposalQueryOutcome::NoProposal + LinearProposalQueryOutcome::NoProposal { .. } )); assert_eq!( invalid_source.discards.lock().unwrap().as_slice(), @@ -821,6 +941,7 @@ mod tests { .unwrap(); let LinearProposalQueryOutcome::DeadlineExceeded { proposal_elapsed_us, + .. } = query_linear_proposal(&late_config, query_params(1, 2, 1, 0, 1, 5, 4)).unwrap() else { panic!("late proposal should produce deadline telemetry"); @@ -844,7 +965,7 @@ mod tests { assert!(matches!( query_linear_proposal(&invalid_token_config, query_params(1, 2, 1, 0, 1, 5, 4),) .unwrap(), - LinearProposalQueryOutcome::NoProposal + LinearProposalQueryOutcome::NoProposal { .. } )); assert_eq!( invalid_token_source.discards.lock().unwrap().as_slice(), @@ -901,7 +1022,7 @@ mod tests { assert!(matches!( query_linear_proposal(&config, query_params(1, 2, 1, 0, 1, 64, 7)).unwrap(), - LinearProposalQueryOutcome::NoProposal + LinearProposalQueryOutcome::NoProposal { .. } )); assert_eq!(source.queries.lock().unwrap()[0].max_proposal_tokens, 7); } @@ -934,6 +1055,12 @@ mod tests { let mut attrs = BTreeMap::new(); receipt.insert_telemetry_attrs(&mut attrs); + LinearProposalSourceTelemetry { + queue_wait_us: 7, + callback_elapsed_us: 11, + outcome: LinearProposalSourceOutcome::DeadlineExceededBeforeDispatch, + } + .insert_telemetry_attrs(&mut attrs); let encoded = serde_json::to_string(&attrs).unwrap(); assert!(!encoded.contains(secret)); @@ -941,5 +1068,9 @@ mod tests { assert!(!encoded.contains("67890")); assert!(!attrs.keys().any(|key| key.contains("decision_id"))); assert!(!attrs.keys().any(|key| key.contains("error"))); + assert_eq!( + attrs["llama_stage.linear_proposal.source_outcome"], + json!("deadline_exceeded_before_dispatch") + ); } } diff --git a/crates/skippy-server/src/frontend/local_generation/linear_decode.rs b/crates/skippy-server/src/frontend/local_generation/linear_decode.rs index 8e21123d0..b1139b119 100644 --- a/crates/skippy-server/src/frontend/local_generation/linear_decode.rs +++ b/crates/skippy-server/src/frontend/local_generation/linear_decode.rs @@ -4,13 +4,13 @@ use std::time::Duration; use openai_frontend::{OpenAiError, OpenAiResult}; use serde_json::json; -use crate::frontend::LinearProposalDisposition; use crate::frontend::generation::{LocalGeneration, StageOpenAiBackend, TokenControl}; use crate::frontend::linear_proposal::{ LinearProposalDiscardReason, LinearProposalExecutionParams, LinearProposalQueryOutcome, LinearProposalQueryParams, execute_linear_proposal_with_terminal_discard, query_linear_proposal, report_linear_proposal_receipt, }; +use crate::frontend::{LinearProposalDisposition, LinearProposalSourceTelemetry}; use super::token_generation::{DecodeState, LinearProposalProgress}; @@ -65,10 +65,21 @@ impl StageOpenAiBackend { runtime_max_proposal_tokens: state.linear_proposal_max_tokens, }, )? { - LinearProposalQueryOutcome::NoProposal => None, + LinearProposalQueryOutcome::NoProposal { source_telemetry } => { + self.emit_linear_proposal_source_telemetry( + source_telemetry, + state.emit_token_debug, + ); + None + } LinearProposalQueryOutcome::DeadlineExceeded { proposal_elapsed_us, + source_telemetry, } => { + self.emit_linear_proposal_source_telemetry( + source_telemetry, + state.emit_token_debug, + ); let mut attrs = BTreeMap::new(); attrs.insert( "llama_stage.linear_proposal.discard_reason".to_string(), @@ -82,7 +93,13 @@ impl StageOpenAiBackend { .emit("stage.openai_linear_proposal_late", attrs); None } - LinearProposalQueryOutcome::Ready(queried) => Some(queried), + LinearProposalQueryOutcome::Ready(queried) => { + self.emit_linear_proposal_source_telemetry( + queried.source_telemetry, + state.emit_token_debug, + ); + Some(queried) + } }; if request .cancellation @@ -176,4 +193,30 @@ impl StageOpenAiBackend { Ok(LinearProposalProgress::Continue) } } + + fn emit_linear_proposal_source_telemetry( + &self, + source_telemetry: Option, + emit_token_debug: bool, + ) { + let Some(source_telemetry) = source_telemetry else { + return; + }; + let mut attrs = BTreeMap::new(); + source_telemetry.insert_telemetry_attrs(&mut attrs); + match source_telemetry.outcome { + crate::frontend::LinearProposalSourceOutcome::Ready + | crate::frontend::LinearProposalSourceOutcome::Abstained + if emit_token_debug => + { + self.telemetry + .emit_debug("stage.openai_linear_proposal_source", attrs); + } + crate::frontend::LinearProposalSourceOutcome::Ready + | crate::frontend::LinearProposalSourceOutcome::Abstained => {} + _ => self + .telemetry + .emit("stage.openai_linear_proposal_source", attrs), + } + } } diff --git a/crates/skippy-server/src/lib.rs b/crates/skippy-server/src/lib.rs index 8a941c3e6..3abf29424 100644 --- a/crates/skippy-server/src/lib.rs +++ b/crates/skippy-server/src/lib.rs @@ -31,6 +31,7 @@ pub use frontend::{ EmbeddedOpenAiBackend, EmbeddedOpenAiRequestDefaults, EmbeddedReasoningBudget, EmbeddedReasoningEnabled, EmbeddedReasoningFormat, LinearProposal, LinearProposalDiscardReason, LinearProposalDisposition, LinearProposalIngress, LinearProposalQuery, LinearProposalReceipt, + LinearProposalSourceOutcome, LinearProposalSourceResponse, LinearProposalSourceTelemetry, NativeMtpProposalConfig, NgramExtensionConfig, NgramProposalConfig, NgramProposerKind, OpaqueProposalDecisionId, OpenAiGuardrailsConfig, OpenAiGuardrailsStatus, OpenAiGuardrailsTarget, SpeculativeDecodeConfig, VerifyWindowConfig, embedded_openai_backend, diff --git a/crates/skippy-server/src/serving_hooks.rs b/crates/skippy-server/src/serving_hooks.rs index 23ae0ffd8..beafc4a53 100644 --- a/crates/skippy-server/src/serving_hooks.rs +++ b/crates/skippy-server/src/serving_hooks.rs @@ -84,8 +84,8 @@ mod tests { use crate::frontend::{ GenerationAbort, GenerationCommit, GenerationReceipt, GenerationReceiptSink, - GenerationStart, LinearProposal, LinearProposalIngress, LinearProposalQuery, - LinearProposalReceipt, + GenerationStart, LinearProposalIngress, LinearProposalQuery, LinearProposalReceipt, + LinearProposalSourceResponse, }; use super::*; @@ -110,8 +110,8 @@ mod tests { struct ProposalIngress; impl LinearProposalIngress for ProposalIngress { - fn propose(&self, _query: LinearProposalQuery) -> Result> { - Ok(None) + fn propose(&self, _query: LinearProposalQuery) -> Result { + Ok(LinearProposalSourceResponse::new(None)) } fn report(&self, _receipt: &LinearProposalReceipt) -> Result<()> { diff --git a/docs/plugins/telemetry.md b/docs/plugins/telemetry.md index c42dec6b5..6ce5aa9a1 100644 --- a/docs/plugins/telemetry.md +++ b/docs/plugins/telemetry.md @@ -149,6 +149,9 @@ to an OTLP record. | `mesh_llm.guardrail.attempt_bucket` | guardrail outcome | Bounded retry bucket: `1`, `2`, or `3_plus`. | | `llama_stage.verify_window.direct_return_upstream_opened` | Skippy decode summary | Boolean indicating that the preferred upstream-opened v10 prediction-return sink completed its handshake. | | `llama_stage.verify_window.direct_return_reverse_fallback` | Skippy decode summary | Boolean indicating that the final stage used the bounded reverse-open v10 prediction-return fallback after the preferred sink was unavailable. | +| `llama_stage.linear_proposal.source_queue_wait_us` | Skippy linear proposal source | Queue wait in microseconds; contains no request, session, token, or plugin identity. | +| `llama_stage.linear_proposal.source_callback_us` | Skippy linear proposal source | Plugin callback duration in microseconds; contains no request, session, token, or plugin identity. | +| `llama_stage.linear_proposal.source_outcome` | Skippy linear proposal source | Bounded outcome enum: `ready`, `abstained`, `host_deadline_exceeded`, `queue_full`, `deadline_exceeded_before_dispatch`, `deadline_exceeded_in_plugin`, `candidate_returned_too_late`, or `source_error`. Ready/abstained events use token-debug sampling; deadline, pressure, late-candidate, and source-error outcomes are emitted unconditionally. | ## Review Checklist