From cb837b61f4c737f0051688c9bd25bee83ec2aaac Mon Sep 17 00:00:00 2001 From: justinchuby <223556219+Copilot@users.noreply.github.com> Date: Tue, 18 Aug 2026 07:41:43 -0700 Subject: [PATCH] refactor(engine): unify native + ORT session-state policy behind SessionStore Extract the duplicated session lookup, rewind bound-checks + error text, checkpoint arithmetic, and token truncation policy shared by the native and ORT decode backends into a single `session_state` module, mirroring the `KvPrefixStore` precedent. Backend-specific mechanism (ORT draft/target validation, scheduler completion, paged-KV; native persistent decoder rewind) stays in per-backend adapters implementing the shared `SessionStore` trait. Adds a two-part DRY guard: a compile-time `CheckedPosition` newtype and a test-time tripwire that fails if the rewind bound-check string is open-coded outside the shared policy module. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- crates/onnx-genai-engine/src/engine/mod.rs | 1 + .../onnx-genai-engine/src/engine/runtime.rs | 466 ++++++++++-------- .../src/engine/session_state.rs | 403 +++++++++++++++ .../onnx-genai-engine/tests/native_engine.rs | 50 +- 4 files changed, 706 insertions(+), 214 deletions(-) create mode 100644 crates/onnx-genai-engine/src/engine/session_state.rs diff --git a/crates/onnx-genai-engine/src/engine/mod.rs b/crates/onnx-genai-engine/src/engine/mod.rs index 1dcf96585e..e4dd0a4040 100644 --- a/crates/onnx-genai-engine/src/engine/mod.rs +++ b/crates/onnx-genai-engine/src/engine/mod.rs @@ -74,6 +74,7 @@ mod model; #[cfg(feature = "native-backend")] mod placement; mod runtime; +pub(crate) mod session_state; mod speculative_load; pub(crate) use decode_backend::*; diff --git a/crates/onnx-genai-engine/src/engine/runtime.rs b/crates/onnx-genai-engine/src/engine/runtime.rs index a29e1f94ce..43214dccfb 100644 --- a/crates/onnx-genai-engine/src/engine/runtime.rs +++ b/crates/onnx-genai-engine/src/engine/runtime.rs @@ -2,10 +2,251 @@ use super::*; +use super::session_state::{CheckedPosition, SessionLen, SessionStore}; + fn generate_uses_scheduler(_backend: EngineDecodeBackend) -> bool { true } +/// Logical token length of an ORT-backed session, the single source of the +/// "session not found" answer for the ORT arm of the shared session policy. +fn ort_logical_len(engine: &Engine, id: SessionId) -> Option { + engine.sessions.get(&id).map(|state| state.tokens.len()) +} + +/// Logical token length of a native-backed session. +#[cfg(feature = "native-backend")] +fn native_logical_len(engine: &Engine, id: SessionId) -> Option { + engine + .native_sessions + .get(&id) + .map(|state| state.tokens.len()) +} + +/// The ORT backend as a [`SessionStore`]: a session in `Engine::sessions`, its +/// paged KV sequence, its scheduler entry, and any aligned draft state. +struct OrtSessions<'a>(&'a mut Engine); + +/// Read-only view of the ORT backend for `&self` policy calls (checkpoint, +/// token count), which need only the logical length. +struct OrtSessionsRef<'a>(&'a Engine); + +impl SessionLen for OrtSessions<'_> { + fn logical_len(&self, id: SessionId) -> Option { + ort_logical_len(self.0, id) + } +} + +impl SessionLen for OrtSessionsRef<'_> { + fn logical_len(&self, id: SessionId) -> Option { + ort_logical_len(self.0, id) + } +} + +impl SessionStore for OrtSessions<'_> { + fn validate_rewind(&self, id: SessionId, target: CheckedPosition) -> anyhow::Result<()> { + let engine = &*self.0; + let position = target.get(); + // Existence is guaranteed by the shared policy; a vanished session is a + // no-op here rather than a second not-found path. + let Some(state) = engine.sessions.get(&id) else { + return Ok(()); + }; + validate_target_state_rewind_to_len( + engine.kv_model.as_ref(), + &engine.kv_cache, + id, + state, + RewindRequest::new(position, RewindRunnerPolicy::RejectRunnerRewind), + )?; + if let (Some(draft_model), Some(draft)) = (&engine.draft, &state.draft) { + let draft_target = position.min(draft.tokens.len()); + validate_draft_state_rewind_to_len( + draft_model, + draft, + RewindRequest::new(draft_target, RewindRunnerPolicy::RejectRunnerRewind), + )?; + } + Ok(()) + } + + fn rewind(&mut self, id: SessionId, target: CheckedPosition) -> anyhow::Result<()> { + let engine = &mut *self.0; + let position = target.get(); + engine.scheduler.complete(id); + let mut state = engine + .sessions + .remove(&id) + .with_context(|| format!("session {id} not found"))?; + let result = (|| { + let session = engine.session.as_deref().context(MISSING_ORT_SESSION)?; + rewind_target_state_to_len( + session, + engine.kv_model.as_ref(), + &mut engine.kv_cache, + id, + &mut state, + RewindRequest::new(position, RewindRunnerPolicy::RejectRunnerRewind), + )?; + if let (Some(draft_model), Some(draft)) = (&mut engine.draft, &mut state.draft) { + let draft_target = position.min(draft.tokens.len()); + rewind_draft_state_to_len( + draft_model, + draft, + RewindRequest::new(draft_target, RewindRunnerPolicy::RejectRunnerRewind), + )?; + } + Ok(()) + })(); + engine.sessions.insert(id, state); + result + } + + fn reset(&mut self, id: SessionId) -> anyhow::Result<()> { + let engine = &mut *self.0; + engine.scheduler.complete(id); + engine + .kv_cache + .remove(id) + .map_err(|e| anyhow::anyhow!("Failed to reset KV sequence {id}: {e}"))?; + engine.kv_cache.page_table.create_sequence(id); + let decode_state = engine.new_target_decode_state()?; + let state = engine + .sessions + .get_mut(&id) + .context("session disappeared during reset")?; + state.tokens.clear(); + state.kv_token_count = 0; + state.decode_state = decode_state; + if let (Some(draft_model), Some(draft)) = (&mut engine.draft, &mut state.draft) { + draft_model + .kv_cache + .remove(draft.seq) + .map_err(|e| anyhow::anyhow!("Failed to reset draft KV sequence: {e}"))?; + draft.seq = draft_model.kv_cache.create_sequence(); + draft.tokens.clear(); + draft.kv_token_count = 0; + draft.decode_state = DecodeState::new_for_path_with_io( + &draft_model.session, + &draft_model.decode_path, + draft_model.io.as_ref(), + )?; + } + Ok(()) + } + + fn close(&mut self, id: SessionId) -> anyhow::Result<()> { + let engine = &mut *self.0; + engine.scheduler.complete(id); + let state = engine + .sessions + .remove(&id) + .with_context(|| format!("session {id} not found"))?; + engine + .kv_cache + .remove(id) + .map_err(|e| anyhow::anyhow!("Failed to remove KV sequence {id}: {e}"))?; + if let (Some(draft_model), Some(draft)) = (&mut engine.draft, state.draft) { + draft_model + .kv_cache + .remove(draft.seq) + .map_err(|e| anyhow::anyhow!("Failed to remove draft KV sequence: {e}"))?; + } + Ok(()) + } +} + +/// The native backend as a [`SessionStore`]: a token history in +/// `Engine::native_sessions` and one in-process decoder that is rewound or reset +/// only while it holds this session's KV (`native_active_session`). +#[cfg(feature = "native-backend")] +struct NativeSessions<'a>(&'a mut Engine); + +/// Read-only view of the native backend for `&self` policy calls. +#[cfg(feature = "native-backend")] +struct NativeSessionsRef<'a>(&'a Engine); + +#[cfg(feature = "native-backend")] +impl SessionLen for NativeSessions<'_> { + fn logical_len(&self, id: SessionId) -> Option { + native_logical_len(self.0, id) + } +} + +#[cfg(feature = "native-backend")] +impl SessionLen for NativeSessionsRef<'_> { + fn logical_len(&self, id: SessionId) -> Option { + native_logical_len(self.0, id) + } +} + +#[cfg(feature = "native-backend")] +impl SessionStore for NativeSessions<'_> { + fn validate_rewind(&self, _id: SessionId, _target: CheckedPosition) -> anyhow::Result<()> { + // The native decoder always admits a rewind and clamps to its own + // materialized length in `rewind`; there is no runner/draft state to + // reject it the way ORT's does. + Ok(()) + } + + fn rewind(&mut self, id: SessionId, target: CheckedPosition) -> anyhow::Result<()> { + let engine = &mut *self.0; + let position = target.get(); + if let Some(state) = engine.native_sessions.get_mut(&id) { + state.tokens.truncate(position); + } + if engine.native_active_session == Some(id) { + let native = engine + .native_session + .as_mut() + .context("native decoder session is unavailable")?; + native.rewind(position.min(native.current_len()))?; + } + let last_access = engine.touch_native_session(); + if let Some(state) = engine.native_sessions.get_mut(&id) { + state.last_access = last_access; + } + Ok(()) + } + + fn reset(&mut self, id: SessionId) -> anyhow::Result<()> { + let engine = &mut *self.0; + let last_access = engine.touch_native_session(); + if let Some(state) = engine.native_sessions.get_mut(&id) { + state.tokens.clear(); + state.last_access = last_access; + } + if engine.native_active_session == Some(id) { + let native = engine + .native_session + .as_mut() + .context("native decoder session is unavailable")?; + native.reset()?; + engine.native_active_session = None; + } + Ok(()) + } + + fn close(&mut self, id: SessionId) -> anyhow::Result<()> { + let engine = &mut *self.0; + engine.native_sessions.remove(&id); + if engine.native_active_session == Some(id) { + let native = engine + .native_session + .as_mut() + .context("native decoder session is unavailable")?; + native.reset()?; + engine.native_active_session = None; + } + if engine.native_default_session == Some(id) { + engine.native_default_session = None; + } + Ok(()) + } +} + + + fn generation_budget_cap(cap: ScheduledBudgetCap) -> GenerationBudgetCap { GenerationBudgetCap { requested_max_new_tokens: cap.requested_max_tokens, @@ -992,24 +1233,10 @@ impl Engine { pub fn checkpoint_session(&self, session_id: SessionId) -> anyhow::Result { #[cfg(feature = "native-backend")] if self.decode_backend == EngineDecodeBackend::Native { - let state = self - .native_sessions - .get(&session_id) - .with_context(|| format!("session {session_id} not found"))?; - return Ok(SessionCheckpoint { - session_id, - position: SessionPosition::new(state.tokens.len()), - }); + return session_state::checkpoint(&NativeSessionsRef(self), session_id); } self.require_ort_backend("session checkpoints")?; - let state = self - .sessions - .get(&session_id) - .with_context(|| format!("session {session_id} not found"))?; - Ok(SessionCheckpoint { - session_id, - position: SessionPosition::new(state.tokens.len()), - }) + session_state::checkpoint(&OrtSessionsRef(self), session_id) } /// Restore a persistent session to a previously checkpointed token boundary. @@ -1030,36 +1257,10 @@ impl Engine { ) -> anyhow::Result { #[cfg(feature = "native-backend")] if self.decode_backend == EngineDecodeBackend::Native { - let tokens = tokens.get(); - let current = self - .native_sessions - .get(&session_id) - .with_context(|| format!("session {session_id} not found"))? - .tokens - .len(); - let target = current.checked_sub(tokens).with_context(|| { - format!( - "cannot rewind session {session_id} by {tokens} tokens from length {current}" - ) - })?; - let target = SessionPosition::new(target); - self.rewind_session_to(session_id, target)?; - return Ok(target); + return session_state::rewind_by(&mut NativeSessions(self), session_id, tokens.get()); } self.require_ort_backend("session rewind")?; - let tokens = tokens.get(); - let current = self - .sessions - .get(&session_id) - .with_context(|| format!("session {session_id} not found"))? - .tokens - .len(); - let target = current.checked_sub(tokens).with_context(|| { - format!("cannot rewind session {session_id} by {tokens} tokens from length {current}") - })?; - let target = SessionPosition::new(target); - self.rewind_session_to(session_id, target)?; - Ok(target) + session_state::rewind_by(&mut OrtSessions(self), session_id, tokens.get()) } /// Rewind a persistent session to an absolute logical token position. @@ -1074,91 +1275,10 @@ impl Engine { ) -> anyhow::Result<()> { #[cfg(feature = "native-backend")] if self.decode_backend == EngineDecodeBackend::Native { - let position = position.get(); - let current = self - .native_sessions - .get(&session_id) - .with_context(|| format!("session {session_id} not found"))? - .tokens - .len(); - if position > current { - anyhow::bail!( - "cannot rewind session {session_id} to token {position}; current length is {current}" - ); - } - let state = self - .native_sessions - .get_mut(&session_id) - .with_context(|| format!("session {session_id} not found"))?; - state.tokens.truncate(position); - if self.native_active_session == Some(session_id) { - let native = self - .native_session - .as_mut() - .context("native decoder session is unavailable")?; - native.rewind(position.min(native.current_len()))?; - } - let last_access = self.touch_native_session(); - if let Some(state) = self.native_sessions.get_mut(&session_id) { - state.last_access = last_access; - } - return Ok(()); + return session_state::rewind_to(&mut NativeSessions(self), session_id, position); } self.require_ort_backend("session rewind")?; - let position = position.get(); - let state = self - .sessions - .get(&session_id) - .with_context(|| format!("session {session_id} not found"))?; - let current = state.tokens.len(); - if position > current { - anyhow::bail!( - "cannot rewind session {session_id} to token {position}; current length is {current}" - ); - } - validate_target_state_rewind_to_len( - self.kv_model.as_ref(), - &self.kv_cache, - session_id, - state, - RewindRequest::new(position, RewindRunnerPolicy::RejectRunnerRewind), - )?; - if let (Some(draft_model), Some(draft)) = (&self.draft, &state.draft) { - let draft_target = position.min(draft.tokens.len()); - validate_draft_state_rewind_to_len( - draft_model, - draft, - RewindRequest::new(draft_target, RewindRunnerPolicy::RejectRunnerRewind), - )?; - } - - self.scheduler.complete(session_id); - let mut state = self - .sessions - .remove(&session_id) - .with_context(|| format!("session {session_id} not found"))?; - let result = (|| { - let session = self.session.as_deref().context(MISSING_ORT_SESSION)?; - rewind_target_state_to_len( - session, - self.kv_model.as_ref(), - &mut self.kv_cache, - session_id, - &mut state, - RewindRequest::new(position, RewindRunnerPolicy::RejectRunnerRewind), - )?; - if let (Some(draft_model), Some(draft)) = (&mut self.draft, &mut state.draft) { - let draft_target = position.min(draft.tokens.len()); - rewind_draft_state_to_len( - draft_model, - draft, - RewindRequest::new(draft_target, RewindRunnerPolicy::RejectRunnerRewind), - )?; - } - Ok(()) - })(); - self.sessions.insert(session_id, state); - result + session_state::rewind_to(&mut OrtSessions(self), session_id, position) } /// Capability for session fork, if the selected backend supports safe CoW @@ -1205,55 +1325,10 @@ impl Engine { pub fn reset_session(&mut self, session_id: SessionId) -> anyhow::Result<()> { #[cfg(feature = "native-backend")] if self.decode_backend == EngineDecodeBackend::Native { - let last_access = self.touch_native_session(); - let state = self - .native_sessions - .get_mut(&session_id) - .with_context(|| format!("session {session_id} not found"))?; - state.tokens.clear(); - state.last_access = last_access; - if self.native_active_session == Some(session_id) { - let native = self - .native_session - .as_mut() - .context("native decoder session is unavailable")?; - native.reset()?; - self.native_active_session = None; - } - return Ok(()); + return session_state::reset(&mut NativeSessions(self), session_id); } self.require_ort_backend("persistent sessions")?; - if !self.sessions.contains_key(&session_id) { - anyhow::bail!("session {session_id} not found"); - } - self.scheduler.complete(session_id); - self.kv_cache - .remove(session_id) - .map_err(|e| anyhow::anyhow!("Failed to reset KV sequence {session_id}: {e}"))?; - self.kv_cache.page_table.create_sequence(session_id); - let decode_state = self.new_target_decode_state()?; - let state = self - .sessions - .get_mut(&session_id) - .context("session disappeared during reset")?; - state.tokens.clear(); - state.kv_token_count = 0; - state.decode_state = decode_state; - if let (Some(draft_model), Some(draft)) = (&mut self.draft, &mut state.draft) { - draft_model - .kv_cache - .remove(draft.seq) - .map_err(|e| anyhow::anyhow!("Failed to reset draft KV sequence: {e}"))?; - draft.seq = draft_model.kv_cache.create_sequence(); - draft.tokens.clear(); - draft.kv_token_count = 0; - draft.decode_state = DecodeState::new_for_path_with_io( - &draft_model.session, - &draft_model.decode_path, - draft_model.io.as_ref(), - )?; - } - Ok(()) + session_state::reset(&mut OrtSessions(self), session_id) } fn new_target_decode_state(&self) -> anyhow::Result { @@ -1293,55 +1368,20 @@ impl Engine { pub fn close_session(&mut self, session_id: SessionId) -> anyhow::Result<()> { #[cfg(feature = "native-backend")] if self.decode_backend == EngineDecodeBackend::Native { - self.native_sessions - .remove(&session_id) - .with_context(|| format!("session {session_id} not found"))?; - if self.native_active_session == Some(session_id) { - let native = self - .native_session - .as_mut() - .context("native decoder session is unavailable")?; - native.reset()?; - self.native_active_session = None; - } - if self.native_default_session == Some(session_id) { - self.native_default_session = None; - } - return Ok(()); + return session_state::close(&mut NativeSessions(self), session_id); } self.require_ort_backend("persistent sessions")?; - self.scheduler.complete(session_id); - let state = self - .sessions - .remove(&session_id) - .with_context(|| format!("session {session_id} not found"))?; - self.kv_cache - .remove(session_id) - .map_err(|e| anyhow::anyhow!("Failed to remove KV sequence {session_id}: {e}"))?; - if let (Some(draft_model), Some(draft)) = (&mut self.draft, state.draft) { - draft_model - .kv_cache - .remove(draft.seq) - .map_err(|e| anyhow::anyhow!("Failed to remove draft KV sequence: {e}"))?; - } - Ok(()) + session_state::close(&mut OrtSessions(self), session_id) } /// Number of logical tokens retained in a persistent session. pub fn session_token_count(&self, session_id: SessionId) -> anyhow::Result { #[cfg(feature = "native-backend")] if self.decode_backend == EngineDecodeBackend::Native { - return self - .native_sessions - .get(&session_id) - .map(|state| state.tokens.len()) - .with_context(|| format!("session {session_id} not found")); + return session_state::token_count(&NativeSessionsRef(self), session_id); } self.require_ort_backend("persistent sessions")?; - self.sessions - .get(&session_id) - .map(|state| state.tokens.len()) - .with_context(|| format!("session {session_id} not found")) + session_state::token_count(&OrtSessionsRef(self), session_id) } /// Get the loaded metadata. diff --git a/crates/onnx-genai-engine/src/engine/session_state.rs b/crates/onnx-genai-engine/src/engine/session_state.rs new file mode 100644 index 0000000000..7e132f8b96 --- /dev/null +++ b/crates/onnx-genai-engine/src/engine/session_state.rs @@ -0,0 +1,403 @@ +//! The single persistent-session state policy, shared by every decode backend. +//! +//! Both backends answer the same questions about a persistent session: does it +//! exist, how many logical tokens does it hold, is a rewind target in bounds, +//! and what token boundary does a checkpoint record? Only the *mechanism* +//! differs. The ORT backend keeps its session in `Engine::sessions` alongside a +//! paged KV cache, a scheduler entry, and optional draft state; the native +//! backend keeps only a token history in `Engine::native_sessions` plus a single +//! in-process decoder it rewinds. That difference belongs in [`SessionStore`], +//! and nowhere else. +//! +//! Keeping the policy in one place is not tidiness. When it lived in two, the +//! copies drifted risk was real and concrete: the `checked_sub` bound check and +//! its error string were duplicated *verbatim* between the native and ORT arms +//! of `rewind_session_by`, the `position > current` bound check and its string +//! between the two arms of `rewind_session_to`, and the "session {id} not found" +//! lookup across every method. A silent divergence in any of these would not +//! crash — it would reject a valid rewind, accept an out-of-bounds one, or emit +//! a different error than a caller's test expects, on exactly one backend. +//! +//! This mirrors [`super::prefix_reuse`], which unified the KV *length* policy +//! the same way: one pure decision, one trait for the mechanism, one DRY guard. + +use crate::config::{SessionCheckpoint, SessionId, SessionPosition}; +use anyhow::Context; + +/// A rewind target that has already been bound-checked against the session's +/// logical length by [`rewind_to`]. +/// +/// This is the DRY mechanism, not decoration. The field is private and the only +/// constructor lives in this module, so a new backend **cannot produce a rewind +/// target at all** without going through the shared bound check — the exact +/// mistake that produced two verbatim copies of that bound check in the first +/// place is now a compile error rather than a review catch. +/// +/// If you find yourself wanting another constructor, that is the signal to stop: +/// you are about to add a second copy of the bound-check policy. Route the new +/// backend through [`rewind_to`] instead, so every backend gets the same check. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct CheckedPosition(usize); + +impl CheckedPosition { + /// The bound-checked absolute token boundary this rewind targets. + pub(crate) fn get(self) -> usize { + self.0 + } +} + +/// Read side of a persistent session: the logical token vector's length is the +/// single source of the "session not found" answer that every method needs. +pub(crate) trait SessionLen { + /// Logical tokens retained by `id`, or `None` when the session is unknown. + fn logical_len(&self, id: SessionId) -> Option; +} + +/// The backend-specific half of session state: everything that touches KV, +/// scheduler, or decoder mechanism. The shared policy owns the lookup, the +/// bound checks, the error text, and the checkpoint arithmetic; a `SessionStore` +/// owns only the mechanism, and its methods are handed inputs the policy has +/// already validated. +pub(crate) trait SessionStore: SessionLen { + /// Check whether a rewind to `target` is admissible **without mutating + /// anything**, so a rejected rewind leaves the session untouched. The ORT + /// backend validates that its decode runner and any draft state can honour + /// the truncation here; the native backend always admits and clamps in + /// [`SessionStore::rewind`]. + fn validate_rewind(&self, id: SessionId, target: CheckedPosition) -> anyhow::Result<()>; + + /// Truncate the session's logical tokens and backend KV to `target`. + /// Called only after [`SessionStore::validate_rewind`] has accepted it. + fn rewind(&mut self, id: SessionId, target: CheckedPosition) -> anyhow::Result<()>; + + /// Clear the session's logical tokens and reset its backend KV, keeping the + /// id usable. + fn reset(&mut self, id: SessionId) -> anyhow::Result<()>; + + /// Remove the session and free its backend state. + fn close(&mut self, id: SessionId) -> anyhow::Result<()>; +} + +/// The single source of the "session {id} not found" lookup and message. +/// +/// Every policy function starts here, so the observable not-found text lives in +/// exactly one place across both backends. +fn require_session(store: &S, id: SessionId) -> anyhow::Result { + store + .logical_len(id) + .with_context(|| format!("session {id} not found")) +} + +/// Record the current logical token boundary as a checkpoint. +pub(crate) fn checkpoint( + store: &S, + id: SessionId, +) -> anyhow::Result { + let len = require_session(store, id)?; + Ok(SessionCheckpoint { + session_id: id, + position: SessionPosition::new(len), + }) +} + +/// Number of logical tokens retained by a session. +pub(crate) fn token_count( + store: &S, + id: SessionId, +) -> anyhow::Result { + require_session(store, id) +} + +/// Rewind a session by `tokens` logical tokens, returning the new position. +pub(crate) fn rewind_by( + store: &mut S, + id: SessionId, + tokens: usize, +) -> anyhow::Result { + let current = require_session(store, id)?; + let target = current.checked_sub(tokens).with_context(|| { + format!("cannot rewind session {id} by {tokens} tokens from length {current}") + })?; + let position = SessionPosition::new(target); + rewind_to(store, id, position)?; + Ok(position) +} + +/// Rewind a session to an absolute logical token position. +/// +/// This is the *only* producer of a [`CheckedPosition`]: the bound check that +/// gates it cannot be skipped or reimplemented by a backend. +pub(crate) fn rewind_to( + store: &mut S, + id: SessionId, + position: SessionPosition, +) -> anyhow::Result<()> { + let requested = position.get(); + let current = require_session(store, id)?; + if requested > current { + anyhow::bail!( + "cannot rewind session {id} to token {requested}; current length is {current}" + ); + } + let target = CheckedPosition(requested); + store.validate_rewind(id, target)?; + store.rewind(id, target) +} + +/// Reset a session, clearing its logical tokens and backend KV. +pub(crate) fn reset( + store: &mut S, + id: SessionId, +) -> anyhow::Result<()> { + require_session(store, id)?; + store.reset(id) +} + +/// Close a session and free its backend state. +pub(crate) fn close( + store: &mut S, + id: SessionId, +) -> anyhow::Result<()> { + require_session(store, id)?; + store.close(id) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + + /// A backend that records what the shared policy asked it to do, so the + /// policy can be exercised without a model, KV cache, or device. + #[derive(Default)] + struct FakeStore { + sessions: HashMap, + rewind_rejects: bool, + rewound_to: Option, + was_reset: bool, + was_closed: bool, + } + + impl FakeStore { + fn with_session(id: SessionId, len: usize) -> Self { + let mut store = FakeStore::default(); + store.sessions.insert(id, len); + store + } + } + + impl SessionLen for FakeStore { + fn logical_len(&self, id: SessionId) -> Option { + self.sessions.get(&id).copied() + } + } + + impl SessionStore for FakeStore { + fn validate_rewind(&self, _id: SessionId, target: CheckedPosition) -> anyhow::Result<()> { + if self.rewind_rejects { + anyhow::bail!("backend rejected the rewind"); + } + // A shared-side error would mutate nothing; record only. + let _ = target; + Ok(()) + } + + fn rewind(&mut self, id: SessionId, target: CheckedPosition) -> anyhow::Result<()> { + self.rewound_to = Some(target.get()); + self.sessions.insert(id, target.get()); + Ok(()) + } + + fn reset(&mut self, id: SessionId) -> anyhow::Result<()> { + self.was_reset = true; + self.sessions.insert(id, 0); + Ok(()) + } + + fn close(&mut self, id: SessionId) -> anyhow::Result<()> { + self.was_closed = true; + self.sessions.remove(&id); + Ok(()) + } + } + + fn sid(id: u64) -> SessionId { + id + } + + #[test] + fn a_missing_session_is_reported_the_same_way_everywhere() { + let store = FakeStore::default(); + for message in [ + checkpoint(&store, sid(1)).unwrap_err().to_string(), + token_count(&store, sid(1)).unwrap_err().to_string(), + ] { + assert_eq!(message, "session 1 not found"); + } + + let mut store = FakeStore::default(); + assert_eq!( + reset(&mut store, sid(1)).unwrap_err().to_string(), + "session 1 not found" + ); + assert_eq!( + close(&mut store, sid(1)).unwrap_err().to_string(), + "session 1 not found" + ); + assert_eq!( + rewind_to(&mut store, sid(1), SessionPosition::new(0)) + .unwrap_err() + .to_string(), + "session 1 not found" + ); + assert_eq!( + rewind_by(&mut store, sid(1), 0).unwrap_err().to_string(), + "session 1 not found" + ); + } + + #[test] + fn a_checkpoint_records_the_logical_length() { + let store = FakeStore::with_session(sid(7), 5); + assert_eq!( + checkpoint(&store, sid(7)).unwrap().position, + SessionPosition::new(5) + ); + assert_eq!(token_count(&store, sid(7)).unwrap(), 5); + } + + #[test] + fn rewinding_past_the_length_is_rejected_with_the_shared_message() { + let mut store = FakeStore::with_session(sid(3), 4); + assert_eq!( + rewind_to(&mut store, sid(3), SessionPosition::new(5)) + .unwrap_err() + .to_string(), + "cannot rewind session 3 to token 5; current length is 4" + ); + assert_eq!(store.rewound_to, None, "a rejected rewind must not mutate"); + } + + #[test] + fn rewinding_by_more_than_the_length_is_rejected_with_the_shared_message() { + let mut store = FakeStore::with_session(sid(3), 4); + assert_eq!( + rewind_by(&mut store, sid(3), 5).unwrap_err().to_string(), + "cannot rewind session 3 by 5 tokens from length 4" + ); + assert_eq!(store.rewound_to, None); + } + + #[test] + fn an_in_bounds_rewind_reaches_the_backend() { + let mut store = FakeStore::with_session(sid(3), 8); + rewind_to(&mut store, sid(3), SessionPosition::new(5)).unwrap(); + assert_eq!(store.rewound_to, Some(5)); + } + + #[test] + fn rewind_by_converts_to_an_absolute_target() { + let mut store = FakeStore::with_session(sid(3), 8); + let position = rewind_by(&mut store, sid(3), 3).unwrap(); + assert_eq!(position, SessionPosition::new(5)); + assert_eq!(store.rewound_to, Some(5)); + } + + #[test] + fn a_backend_rejection_leaves_the_session_untouched() { + let mut store = FakeStore { + rewind_rejects: true, + ..FakeStore::with_session(sid(3), 8) + }; + assert!(rewind_to(&mut store, sid(3), SessionPosition::new(5)).is_err()); + assert_eq!(store.rewound_to, None, "a rejected rewind must not mutate"); + } + + #[test] + fn reset_and_close_require_the_session_then_reach_the_backend() { + let mut store = FakeStore::with_session(sid(3), 8); + reset(&mut store, sid(3)).unwrap(); + assert!(store.was_reset); + assert_eq!(store.logical_len(sid(3)), Some(0)); + + close(&mut store, sid(3)).unwrap(); + assert!(store.was_closed); + assert_eq!(store.logical_len(sid(3)), None); + } + + /// Files allowed to spell the session-state policy strings. + /// + /// This module holds the one copy of the bound-check policy. A new backend + /// that open-codes a session rewind would reintroduce this fragment + /// elsewhere, which is the drift this module exists to prevent. + const SESSION_POLICY_CALLERS: [&str; 1] = ["engine\\session_state.rs"]; + + /// The bound-check message fragment unique to the rewind policy. `rewind_by` + /// and `rewind_to` both format it; `fork_session` uses "cannot fork + /// session", which is deliberately distinct. + const SESSION_POLICY_FRAGMENT: &str = "cannot rewind session"; + + fn rust_sources(dir: &std::path::Path, out: &mut Vec) { + for entry in std::fs::read_dir(dir).expect("crate src is readable") { + let path = entry.expect("readable entry").path(); + if path.is_dir() { + rust_sources(&path, out); + } else if path.extension().is_some_and(|ext| ext == "rs") { + out.push(path); + } + } + } + + /// The DRY tripwire, mirroring `prefix_reuse`'s KV-rewind guard. + /// + /// The rewind bound check is the piece that was duplicated *verbatim* + /// between the two backends. Confining its message fragment to this module, + /// together with the [`CheckedPosition`] compile-time guard, means a third + /// backend cannot open-code a rewind: it can neither construct a validated + /// target nor spell the bound-check error without tripping a test. + /// + /// If this fails, route the new backend through [`rewind_to`] and add a + /// [`SessionStore`] adapter -- **not** add the offending file to + /// `SESSION_POLICY_CALLERS`. Widening the allowlist is how the duplication + /// this module exists to prevent gets back in. + #[test] + fn the_rewind_bound_check_lives_only_in_the_shared_policy() { + let src = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("src"); + let mut sources = Vec::new(); + rust_sources(&src, &mut sources); + assert!( + sources.len() > 10, + "the source scan found almost nothing, so it cannot be trusted to \ + have looked: {} files under {}", + sources.len(), + src.display() + ); + + let mut offenders = Vec::new(); + for path in sources { + let relative = path.strip_prefix(&src).expect("scanned under src"); + let display = relative.display().to_string().replace('/', "\\"); + if SESSION_POLICY_CALLERS + .iter() + .any(|allowed| display.ends_with(allowed)) + { + continue; + } + let text = std::fs::read_to_string(&path).expect("source is readable"); + for (line_number, line) in text.lines().enumerate() { + if line.contains(SESSION_POLICY_FRAGMENT) { + offenders.push(format!("{display}:{} {}", line_number + 1, line.trim())); + } + } + } + + assert!( + offenders.is_empty(), + "these open-code the session rewind bound check outside the shared \ + `session_state` policy, which is how it gets reimplemented per \ + backend. Route them through `rewind_to` instead of widening \ + `SESSION_POLICY_CALLERS`:\n{}", + offenders.join("\n") + ); + } +} diff --git a/crates/onnx-genai-engine/tests/native_engine.rs b/crates/onnx-genai-engine/tests/native_engine.rs index 3925095060..6809ace505 100644 --- a/crates/onnx-genai-engine/tests/native_engine.rs +++ b/crates/onnx-genai-engine/tests/native_engine.rs @@ -2,7 +2,8 @@ use onnx_genai_engine::{ Engine, EngineConfig, EngineDecodeBackend, GeneratePrompt, GenerateRequest, - NATIVE_SESSION_INCREMENTAL_PREFILL_TEST_HITS, NativeDecodeDevice, SpeculativeMode, + NATIVE_SESSION_INCREMENTAL_PREFILL_TEST_HITS, NativeDecodeDevice, RewindTokenCount, + SpeculativeMode, }; use onnx_genai_ort::{SessionOptions, ep_selection}; use std::path::Path; @@ -513,6 +514,53 @@ fn native_session_switching_matches_cold_start() -> anyhow::Result<()> { } /// Native sessions use the unified session API and allow multiple logical sessions. +/// A valid explicit rewind through the shared `session_state` policy truncates +/// the native session's logical length. This drives the same shared bound check +/// that the ORT `failed_rewind_of_*` tests reach, so inverting the check in +/// `session_state::rewind_to` turns *both* backends red — the proof that the two +/// arms genuinely share one policy rather than two look-alike copies. +#[test] +fn native_session_rewind_by_truncates_logical_length() -> anyhow::Result<()> { + let fixture = + Path::new(env!("CARGO_MANIFEST_DIR")).join("../../tests/fixtures/tiny-native-sub4-engine"); + let mut engine = Engine::from_dir( + &fixture, + EngineConfig { + decode_backend: EngineDecodeBackend::Native, + ..EngineConfig::default() + }, + )?; + + let session = engine.create_session()?; + let mut request = GenerateRequest::new(GeneratePrompt::TokenIds(vec![0u32, 0, 0])); + request.options.max_new_tokens = 2; + request.options.temperature = 0.0; + request.options.stop_on_eos = false; + engine.generate_in_session(session, request)?; + + let before = engine.session_token_count(session)?; + assert!( + before >= 2, + "expected at least two logical tokens before rewind, got {before}" + ); + + let new_position = engine.rewind_session_by(session, RewindTokenCount::new(2))?; + assert_eq!(new_position.get(), before - 2); + assert_eq!(engine.session_token_count(session)?, before - 2); + + // Rewinding past the start is rejected by the same shared bound check. + let error = engine + .rewind_session_by(session, RewindTokenCount::new(before)) + .expect_err("rewinding past the start must fail"); + assert!( + error.to_string().contains("cannot rewind session"), + "unexpected error: {error:#}" + ); + + engine.close_session(session)?; + Ok(()) +} + #[test] fn native_session_creation_uses_unified_multi_session_api() -> anyhow::Result<()> { let fixture =