el_engine_candle/lib.rs
1//! `el-engine-candle` — inference engine adapter over **Candle** (ADR-002),
2//! implementing [`el_runtime::InferenceEngine`] / `RuntimeAcl`.
3//!
4//! Consumers supply their own model file; see [`CandleEngine::from_path`] and
5//! [`CandleEngine::from_bytes`]. For tests that need a working engine without
6//! a model asset, use [`CandleEngine::toy`].
7//!
8//! Expected GGUF tensor names:
9//! - `token_embd.weight` — embedding table `[vocab, dim]`
10//! - `output.weight` or `lm_head.weight` — lm-head `[vocab, dim]` (standard Llama layout)
11//!
12//! Float logits are quantised to integer milli-logits at the ACL boundary, so
13//! Candle's `Tensor`/`Device` types never cross into the domain.
14
15#![forbid(unsafe_code)]
16
17use candle_core::{Device, Tensor};
18use el_core::{
19 ChatMessage, ChatRequest, ChatResponse, ChatRole, ChatToken, DomainEvent, EdgeError,
20 LlmProvider, Result, SafetyMode, SessionConfig, SessionId, StopReason, Token,
21};
22use el_provenance::LoadPermit;
23use el_runtime::{
24 AnchorGuard, ContrastiveSteerer, ExpertLogits, InferenceEngine, InferenceSession,
25 LightweightFilter, NoSafety, Ports, SafetyModeSelector, SafetySteerer,
26};
27
28/// Candle-backed inference engine.
29pub struct CandleEngine {
30 embed: Tensor,
31 w_out: Tensor,
32 vocab: usize,
33 eos: Token,
34}
35
36impl CandleEngine {
37 /// Build a deterministic toy model on the CPU — no model file required.
38 ///
39 /// Uses fixed synthetic weights so tests are deterministic.
40 pub fn toy(vocab: usize, dim: usize, eos: Token) -> Result<Self> {
41 let device = Device::Cpu;
42
43 let embed_data: Vec<f32> = (0..vocab * dim)
44 .map(|k| {
45 let (i, j) = (k / dim, k % dim);
46 (((i + j) % 7) as f32) * 0.1
47 })
48 .collect();
49 let wout_data: Vec<f32> = (0..dim * vocab)
50 .map(|k| {
51 let (a, b) = (k / vocab, k % vocab);
52 ((((a * 31 + b * 17) % 13) as f32) * 0.1) - 0.6
53 })
54 .collect();
55
56 let embed = Tensor::from_vec(embed_data, (vocab, dim), &device)
57 .map_err(|_| EdgeError::Engine("candle: embed tensor build failed"))?;
58 let w_out = Tensor::from_vec(wout_data, (dim, vocab), &device)
59 .map_err(|_| EdgeError::Engine("candle: w_out tensor build failed"))?;
60
61 Ok(Self {
62 embed,
63 w_out,
64 vocab,
65 eos,
66 })
67 }
68
69 /// Load `token_embd.weight` and `output.weight` from a consumer-supplied GGUF file.
70 ///
71 /// # Limitations
72 /// This engine's forward pass is `embed[last_token] · w_out` — a single linear
73 /// projection. Only these two tensors are used; transformer blocks, attention,
74 /// RoPE, and norms present in the GGUF are ignored. Logits will not match a
75 /// real Llama/Mistral/etc. forward. This is the ADR-002 engine-seam proof; for
76 /// a full transformer forward implement a separate [`InferenceEngine`] using
77 /// `candle-transformers`.
78 pub fn from_path(path: impl AsRef<std::path::Path>, eos: Token) -> Result<Self> {
79 let file = std::fs::File::open(path.as_ref())
80 .map_err(|_| EdgeError::Engine("model file not found or not readable"))?;
81 Self::load_gguf(&mut std::io::BufReader::new(file), eos)
82 }
83
84 /// Load from raw bytes (WASM / memory-mapped scenarios).
85 ///
86 /// Same limitations as [`Self::from_path`]: only `token_embd.weight` and
87 /// `output.weight` are used; the forward is `embed[last] · w_out`.
88 pub fn from_bytes(data: &[u8], eos: Token) -> Result<Self> {
89 Self::load_gguf(&mut std::io::Cursor::new(data), eos)
90 }
91
92 fn load_gguf<R: std::io::Read + std::io::Seek>(reader: &mut R, eos: Token) -> Result<Self> {
93 use candle_core::quantized::gguf_file;
94
95 let content = gguf_file::Content::read(reader)
96 .map_err(|_| EdgeError::Engine("GGUF: invalid or unrecognised file"))?;
97 let device = Device::Cpu;
98
99 let embed = content
100 .tensor(reader, "token_embd.weight", &device)
101 .map_err(|_| EdgeError::Engine("GGUF: missing 'token_embd.weight'"))?
102 .dequantize(&device)
103 .map_err(|_| EdgeError::Engine("GGUF: cannot dequantize embed tensor"))?;
104
105 let (vocab, dim) = match embed.shape().dims() {
106 [v, d] => (*v, *d),
107 _ => return Err(EdgeError::Engine("GGUF: 'token_embd.weight' must be 2-D")),
108 };
109
110 let raw_w_q = match content.tensor(reader, "output.weight", &device) {
111 Ok(t) => t,
112 Err(_) => content
113 .tensor(reader, "lm_head.weight", &device)
114 .map_err(|_| {
115 EdgeError::Engine("GGUF: missing 'output.weight' / 'lm_head.weight'")
116 })?,
117 };
118 let raw_w = raw_w_q
119 .dequantize(&device)
120 .map_err(|_| EdgeError::Engine("GGUF: cannot dequantize output weight"))?;
121
122 // Standard GGUF / Llama convention: output.weight is [vocab, dim].
123 // We need [dim, vocab] so that embed_row [1,dim] × w_out [dim,vocab] → logits [1,vocab].
124 let w_out = match raw_w.shape().dims() {
125 [v, _d] if *v == vocab => raw_w
126 .t()
127 .map_err(|_| EdgeError::Engine("GGUF: failed to transpose output weight"))?,
128 _ => raw_w,
129 };
130
131 // Validate that the output weight's inner dimension matches the embedding dimension.
132 // A mismatch would silently produce all-zero logits at inference time.
133 match w_out.shape().dims() {
134 [d, v] if *d == dim && *v == vocab => {}
135 _ => return Err(EdgeError::Engine(
136 "GGUF: output weight shape incompatible with embed dim — expected [dim, vocab] after transpose",
137 )),
138 }
139
140 Ok(Self {
141 embed,
142 w_out,
143 vocab,
144 eos,
145 })
146 }
147
148 /// One real Candle forward: `embed[last] · w_out` → length-`vocab` logits.
149 fn forward(&self, last: usize) -> candle_core::Result<Vec<f32>> {
150 let row = self.embed.narrow(0, last, 1)?; // [1, dim]
151 let logits = row.matmul(&self.w_out)?; // [1, vocab]
152 Ok(logits.to_vec2::<f32>()?.remove(0))
153 }
154}
155
156impl InferenceEngine for CandleEngine {
157 fn prefill(&mut self, tokens: &[Token]) -> Result<u32> {
158 Ok(tokens.len() as u32)
159 }
160
161 fn next_logits(&mut self, committed: &[Token]) -> Vec<i32> {
162 let last = committed
163 .last()
164 .copied()
165 .unwrap_or(0)
166 .min(self.vocab as u32 - 1) as usize;
167 match self.forward(last) {
168 Ok(logits) => logits.iter().map(|x| (x * 1000.0).round() as i32).collect(),
169 Err(_) => vec![0; self.vocab],
170 }
171 }
172
173 fn eos_token(&self) -> Token {
174 self.eos
175 }
176
177 /// Stateless: this engine's forward is `embed[committed.last()] · w_out`, so
178 /// it holds no KV cache to restore — a rollback is a no-op.
179 fn rollback(&mut self, _keep_committed: u32) -> Result<()> {
180 Ok(())
181 }
182
183 /// Stateless: no conversation cache to discard between turns (ADR-018).
184 fn reset_cache(&mut self) -> Result<()> {
185 Ok(())
186 }
187}
188
189// ── LlmProvider (text-level) wrapper (ADR-010) ───────────────────────────────
190
191/// Wraps a `CandleEngine` behind the `LlmProvider` trait using a byte-level
192/// tokenizer. A production build would swap in a HuggingFace tokenizer loaded
193/// from the model file.
194pub struct LocalLlmProvider {
195 session: std::sync::Mutex<InferenceSession<CandleEngine>>,
196 vocab: usize,
197}
198
199impl LocalLlmProvider {
200 /// Load from a consumer-supplied GGUF file.
201 pub fn from_path(
202 path: impl AsRef<std::path::Path>,
203 eos: Token,
204 permit: LoadPermit,
205 ) -> Result<Self> {
206 let engine = CandleEngine::from_path(path, eos)?;
207 let vocab = engine.vocab;
208 let session = InferenceSession::new(SessionId(1), SessionConfig::default(), engine, permit);
209 Ok(Self {
210 session: std::sync::Mutex::new(session),
211 vocab,
212 })
213 }
214
215 /// Build a toy provider for testing.
216 pub fn toy(vocab: usize, dim: usize, eos: Token, permit: LoadPermit) -> Result<Self> {
217 let engine = CandleEngine::toy(vocab, dim, eos)?;
218 let session = InferenceSession::new(SessionId(1), SessionConfig::default(), engine, permit);
219 Ok(Self {
220 session: std::sync::Mutex::new(session),
221 vocab,
222 })
223 }
224
225 fn encode(&self, text: &str) -> Vec<Token> {
226 text.bytes()
227 .map(|b| (b as Token) % self.vocab as Token)
228 .collect()
229 }
230
231 fn decode(tokens: &[Token]) -> String {
232 tokens
233 .iter()
234 .map(|&t| {
235 let b = (t & 0xFF) as u8;
236 if b.is_ascii_graphic() || b == b' ' {
237 b as char
238 } else {
239 '?'
240 }
241 })
242 .collect()
243 }
244
245 fn format_messages(messages: &[ChatMessage]) -> String {
246 messages
247 .iter()
248 .map(|m| {
249 let role = match m.role {
250 ChatRole::System => "system",
251 ChatRole::User => "user",
252 ChatRole::Assistant => "assistant",
253 };
254 format!("{role}: {}", m.content)
255 })
256 .collect::<Vec<_>>()
257 .join("\n")
258 }
259}
260
261impl LlmProvider for LocalLlmProvider {
262 fn chat(&self, req: &ChatRequest) -> Result<ChatResponse> {
263 let prompt = Self::format_messages(&req.messages);
264 let prompt_tokens = self.encode(&prompt);
265 let prompt_len = prompt_tokens.len() as u32;
266 let max = req.max_tokens.unwrap_or(64);
267
268 let mut session = self.session.lock().unwrap();
269 session.reset()?;
270 let _ = session.drain_events(); // bound buffered events across reused turns
271 let ports = Ports::permissive();
272 session.load_prompt(&ports, &prompt_tokens)?;
273 session.generate(&ports, max)?;
274
275 let output = session.output().to_vec();
276 let completion_len = output.len() as u32;
277
278 Ok(ChatResponse {
279 content: Self::decode(&output),
280 model: "local/candle".into(),
281 prompt_tokens: prompt_len,
282 completion_tokens: completion_len,
283 })
284 }
285
286 fn chat_stream(&self, req: &ChatRequest, on_token: &mut dyn FnMut(ChatToken)) -> Result<()> {
287 let resp = self.chat(req)?;
288 for ch in resp.content.chars() {
289 on_token(ChatToken {
290 text: ch.to_string(),
291 is_final: false,
292 });
293 }
294 on_token(ChatToken {
295 text: String::new(),
296 is_final: true,
297 });
298 Ok(())
299 }
300}
301
302// ── Real Qwen2 transformer engine + chat provider (ADR-002 + ADR-010) ────────
303//
304// Unlike `CandleEngine` (a single linear projection used as the engine-seam
305// proof) this runs a genuine Qwen2 transformer forward via `candle-transformers`
306// with a real HuggingFace tokenizer, so it produces coherent chat. It plugs into
307// the SAME `el_runtime::InferenceSession` decode loop as every other engine —
308// nothing in the SDK pipeline is bypassed.
309
310use candle_transformers::models::quantized_qwen2::ModelWeights as Qwen2Weights;
311use el_core::{ModelId, ModelVersion};
312use el_provenance::{ModelArtifact, SignatureVerifier};
313use tokenizers::Tokenizer;
314
315// ── Opt-in benchmark instrumentation (EL_BENCH=1) ────────────────────────────
316//
317// Zero-cost when `EL_BENCH` is unset: `enabled()` short-circuits and no timing
318// is taken. When set, `QwenChatProvider::chat` prints a per-phase breakdown and
319// per-forward attribution (model compute vs. seam quantisation vs. runtime loop)
320// to stderr. Diagnostics only — not part of the SDK's public behaviour.
321mod bench {
322 use std::cell::Cell;
323 use std::sync::OnceLock;
324 use std::time::Duration;
325
326 static ENABLED: OnceLock<bool> = OnceLock::new();
327
328 /// True iff the `EL_BENCH` environment variable is present (read once).
329 pub fn enabled() -> bool {
330 *ENABLED.get_or_init(|| std::env::var_os("EL_BENCH").is_some())
331 }
332
333 thread_local! {
334 static FWD_TOTAL: Cell<Duration> = const { Cell::new(Duration::ZERO) };
335 static FWD_MODEL: Cell<Duration> = const { Cell::new(Duration::ZERO) };
336 static FWD_CALLS: Cell<u64> = const { Cell::new(0) };
337 }
338
339 /// Accumulate one `forward_one` sample: `total` is the whole seam call,
340 /// `model` is just the candle transformer forward inside it.
341 pub fn record(total: Duration, model: Duration) {
342 FWD_TOTAL.with(|c| c.set(c.get() + total));
343 FWD_MODEL.with(|c| c.set(c.get() + model));
344 FWD_CALLS.with(|c| c.set(c.get() + 1));
345 }
346
347 /// Read and reset the forward accumulators: `(total, model, calls)`.
348 pub fn take() -> (Duration, Duration, u64) {
349 (
350 FWD_TOTAL.replace(Duration::ZERO),
351 FWD_MODEL.replace(Duration::ZERO),
352 FWD_CALLS.replace(0),
353 )
354 }
355}
356
357/// A real Qwen2 transformer `InferenceEngine`.
358///
359/// Holds candle's stateful KV cache. Within one generation it is fed
360/// incrementally (prefill, then one new token per `next_logits` call). The engine
361/// is **loaded once and reused across conversations** (ADR-018): candle exposes no
362/// public cache-clear, but its attention *replaces* the cache on a forward at
363/// `index_pos == 0`, so [`reset_cache`](InferenceEngine::reset_cache) evicts the
364/// previous conversation's KV with a single benign position-0 forward — no engine
365/// reconstruction and no reload from disk between turns.
366///
367/// Across turns of the *same* conversation it also reuses the unchanged prefix's
368/// KV: [`prefill_reuse`](InferenceEngine::prefill_reuse) feeds only the suffix the
369/// re-rendered conversation adds beyond `cached`, the live token sequence behind
370/// the cache (ADR-018 AC-3 cross-turn incremental prefill).
371///
372/// A *within-generation* safety backtrack (ADR-012) is supported via
373/// [`InferenceEngine::rollback`]: candle's attention discards its cache when a
374/// forward runs at `index_pos == 0`, so we retain the prompt and replay it from
375/// position 0 to rebuild the cache for the safe prefix (the session then
376/// re-feeds the retained committed tokens). Float logits are quantised to
377/// integer milli-logits at the seam, exactly like [`CandleEngine`], so the
378/// runtime stays float-free.
379pub struct QwenEngine {
380 model: Qwen2Weights,
381 device: Device,
382 /// Absolute KV position written so far (candle's `index_pos`).
383 index_pos: usize,
384 /// How many of the runtime-`committed` tokens have already been fed.
385 fed: usize,
386 /// The prefill prompt, retained so a rollback can replay it from position 0
387 /// to rebuild candle's KV cache (which has no public truncation).
388 prompt: Vec<Token>,
389 /// The exact token sequence currently represented by the KV cache —
390 /// `prompt` plus the committed tokens fed so far. Its length equals
391 /// `index_pos` (every `forward_one` advances both in lock-step). It is the
392 /// basis for the cross-turn longest-common-prefix reuse check (ADR-018 AC-3,
393 /// [`prefill_reuse`](InferenceEngine::prefill_reuse)).
394 cached: Vec<Token>,
395 /// Milli-logits produced after the most recent forward.
396 last_logits: Vec<i32>,
397 vocab: usize,
398 eos: Token,
399 /// Whether candle's per-layer KV cache may hold conversation-derived K/V that
400 /// still needs clearing (ADR-018). Set by every `forward_one` (before the
401 /// fallible model forward) and cleared **only** after a fully successful
402 /// eviction forward in `reset_cache`, so a partially-failed eviction is retried
403 /// rather than skipped — `index_pos` alone can't carry that signal.
404 cache_dirty: bool,
405}
406
407/// Whether a failed forward has already appended the token to Candle's KV.
408///
409/// `next_logits` cannot return an error to the runtime. It must therefore know
410/// whether to consume the committed token before returning neutral logits.
411enum ForwardOneError {
412 BeforeForward(EdgeError),
413 AfterForward(EdgeError),
414}
415
416impl ForwardOneError {
417 fn into_edge(self) -> EdgeError {
418 match self {
419 Self::BeforeForward(error) | Self::AfterForward(error) => error,
420 }
421 }
422}
423
424fn apply_committed_forward_result(
425 result: std::result::Result<Vec<i32>, ForwardOneError>,
426 token: Token,
427 cached: &mut Vec<Token>,
428 fed: &mut usize,
429 last_logits: &mut Vec<i32>,
430 vocab: usize,
431) -> Option<Vec<i32>> {
432 match result {
433 Ok(logits) => {
434 *last_logits = logits;
435 cached.push(token);
436 *fed += 1;
437 None
438 }
439 Err(ForwardOneError::AfterForward(_)) => {
440 // Candle has already appended this token. Consume it in the Rust
441 // bookkeeping too so the next decode step cannot feed it twice.
442 cached.push(token);
443 *fed += 1;
444 Some(vec![0; vocab.max(1)])
445 }
446 Err(ForwardOneError::BeforeForward(_)) => Some(vec![0; vocab.max(1)]),
447 }
448}
449
450impl QwenEngine {
451 /// Load Qwen2 weights from a consumer-supplied GGUF file.
452 pub fn from_path(path: impl AsRef<std::path::Path>, eos: Token) -> Result<Self> {
453 use candle_core::quantized::gguf_file;
454 let mut file = std::fs::File::open(path.as_ref())
455 .map_err(|_| EdgeError::Engine("model file not found or not readable"))?;
456 let content = gguf_file::Content::read(&mut file)
457 .map_err(|_| EdgeError::Engine("GGUF: invalid or unrecognised file"))?;
458 let device = Device::Cpu;
459 let model = Qwen2Weights::from_gguf(content, &mut file, &device)
460 .map_err(|_| EdgeError::Engine("GGUF: failed to load Qwen2 weights"))?;
461 Ok(Self {
462 model,
463 device,
464 index_pos: 0,
465 fed: 0,
466 prompt: Vec::new(),
467 cached: Vec::new(),
468 last_logits: Vec::new(),
469 vocab: 0,
470 eos,
471 cache_dirty: false,
472 })
473 }
474
475 /// One forward over a single token at the current position; advances the KV
476 /// cache and returns milli-logits for the next token.
477 fn forward_one(&mut self, token: Token) -> std::result::Result<Vec<i32>, ForwardOneError> {
478 // Any forward may write conversation K/V into candle's cache; mark dirty
479 // before the fallible call so a forward that fails part-way still leaves
480 // the cache flagged for clearing (ADR-018).
481 self.cache_dirty = true;
482 let t_total = bench::enabled().then(std::time::Instant::now);
483
484 let input = Tensor::from_vec(vec![token], (1, 1), &self.device).map_err(|_| {
485 ForwardOneError::BeforeForward(EdgeError::Engine("candle: input tensor build failed"))
486 })?;
487
488 let t_model = bench::enabled().then(std::time::Instant::now);
489 let logits = self.model.forward(&input, self.index_pos).map_err(|_| {
490 ForwardOneError::BeforeForward(EdgeError::Engine("candle: Qwen2 forward failed"))
491 })?;
492 // Candle appends to its KV cache inside `forward`, before logits are
493 // extracted below. Keep the logical position aligned if extraction fails.
494 self.index_pos += 1;
495 let model_dur = t_model.map(|t| t.elapsed()).unwrap_or_default();
496
497 let row = logits.squeeze(0).map_err(|_| {
498 ForwardOneError::AfterForward(EdgeError::Engine("candle: squeeze logits failed"))
499 })?;
500 let floats = row.to_vec1::<f32>().map_err(|_| {
501 ForwardOneError::AfterForward(EdgeError::Engine("candle: logits to vec failed"))
502 })?;
503 let out: Vec<i32> = floats.iter().map(|x| (x * 1000.0).round() as i32).collect();
504 if let Some(t) = t_total {
505 bench::record(t.elapsed(), model_dur);
506 }
507 Ok(out)
508 }
509}
510
511impl InferenceEngine for QwenEngine {
512 fn prefill(&mut self, tokens: &[Token]) -> Result<u32> {
513 self.index_pos = 0;
514 self.fed = 0;
515 self.prompt = tokens.to_vec(); // retained for rollback replay
516 self.cached = Vec::with_capacity(tokens.len());
517 for &t in tokens {
518 self.last_logits = self.forward_one(t).map_err(ForwardOneError::into_edge)?;
519 self.cached.push(t);
520 }
521 self.vocab = self.last_logits.len();
522 Ok(tokens.len() as u32)
523 }
524
525 fn next_logits(&mut self, committed: &[Token]) -> Vec<i32> {
526 // Feed any newly committed (generated) tokens beyond what we've seen.
527 // `committed` grows by exactly one per decode step, so this feeds the
528 // token the runtime just sampled and returns the next distribution.
529 while self.fed < committed.len() {
530 let t = committed[self.fed];
531 if let Some(fallback) = apply_committed_forward_result(
532 self.forward_one(t),
533 t,
534 &mut self.cached,
535 &mut self.fed,
536 &mut self.last_logits,
537 self.vocab,
538 ) {
539 return fallback;
540 }
541 }
542 self.last_logits.clone()
543 }
544
545 fn eos_token(&self) -> Token {
546 self.eos
547 }
548
549 fn rollback(&mut self, _keep_committed: u32) -> Result<()> {
550 // candle's KV cache is append-only with no public truncation, but its
551 // attention discards the cache on a forward at `index_pos == 0` (see
552 // quantized_qwen2). So rebuild deterministically: replay the prompt from
553 // position 0 — the first forward resets the cache, the rest re-append it —
554 // leaving the engine in its exact post-prefill state. We reset `fed` to 0
555 // so the session's next `next_logits` re-feeds the retained committed
556 // prefix (already truncated to `keep_committed`) on top. Cost is bounded
557 // by `max_rollbacks` (ADR-012).
558 self.index_pos = 0;
559 self.fed = 0;
560 self.cached = Vec::with_capacity(self.prompt.len());
561 for i in 0..self.prompt.len() {
562 let t = self.prompt[i];
563 self.last_logits = self.forward_one(t).map_err(ForwardOneError::into_edge)?;
564 // Keep cached in lock-step with index_pos so a mid-replay error
565 // leaves the invariant intact rather than holding the old value.
566 self.cached.push(t);
567 }
568 Ok(())
569 }
570
571 /// Release the current conversation's KV **while keeping the resident weights
572 /// loaded** (ADR-018) — the separation of conversation lifecycle from model
573 /// lifecycle, and the engine half of [`InferenceSession::close`] / `reset`.
574 ///
575 /// candle's `quantized_qwen2` owns its per-layer KV with no public clear API,
576 /// but its attention *ignores and replaces* the cache on a forward at
577 /// `index_pos == 0`. So one forward over a benign token (id 0) drops the prior
578 /// (user) K/V tensors — freeing that memory and clearing the user's data from
579 /// the cache (PRD line 131) — without touching the weights. What remains is a
580 /// single non-user token's KV, itself overwritten by the next prefill or freed
581 /// when the engine is dropped. Skipped when nothing has been cached yet
582 /// (`index_pos == 0`), so a pristine or already-cleared engine does no work.
583 ///
584 /// Distinct from `rollback`, which *replays* a retained prefix to rewind
585 /// within a single generation. Fallible (it runs a forward); on error the
586 /// caller (`reset`/`close`) leaves session state untouched and surfaces it.
587 fn reset_cache(&mut self) -> Result<()> {
588 if self.cache_dirty {
589 // Overwrite (and thereby drop) the user K/V by forwarding a benign
590 // token at position 0; the resulting logits are discarded. `cache_dirty`
591 // is cleared **only after** a fully successful forward — if candle
592 // fails after replacing some layers, it stays set so the next call
593 // re-clears (a partially-cleared cache is never reported as clean).
594 self.index_pos = 0;
595 self.forward_one(0).map_err(ForwardOneError::into_edge)?;
596 self.cache_dirty = false;
597 }
598 self.index_pos = 0;
599 self.fed = 0;
600 // Release (not just `clear`) the conversation-derived buffers so their
601 // bytes aren't retained in an owned allocation (P2): `Vec::new()` drops
602 // the old allocation; `clear()` would keep capacity and the stale ids.
603 self.prompt = Vec::new();
604 self.cached = Vec::new();
605 self.last_logits = Vec::new();
606 Ok(())
607 }
608
609 /// Cross-turn incremental prefill (ADR-018 AC-3): reuse the KV already cached
610 /// for the longest prefix `full_context` shares with the live cache, and feed
611 /// only the divergent suffix at the live position — no reload, no whole-history
612 /// re-prefill.
613 ///
614 /// `cached` is the exact token sequence behind the current KV (length ==
615 /// `index_pos`). The token-level longest-common-prefix against it is the
616 /// tokenizer-round-trip guard: a re-rendered+re-tokenized conversation that
617 /// drifts from what was generated simply matches a shorter prefix and the rest
618 /// is fed fresh. When `full_context` exactly extends the cache, only the new
619 /// tail is forwarded (the fast path); otherwise — divergence, or a context
620 /// shorter than the cache — candle cannot truncate its append-only cache, so we
621 /// rebuild from position 0 (a forward at `index_pos == 0` drops the old cache),
622 /// which is never worse than the pre-ADR-018 full re-prefill.
623 ///
624 /// Either branch leaves the engine in the **same** state a `reset_cache()` +
625 /// `prefill(full_context)` would: the suffix is fed by the identical
626 /// `forward_one` calls at the identical positions, so subsequent logits are
627 /// bit-identical to a from-scratch prefill (the soundness contract).
628 fn prefill_reuse(&mut self, full_context: &[Token]) -> Result<u32> {
629 let reuse = longest_common_prefix(&self.cached, full_context);
630 if reuse == self.cached.len() && reuse == self.index_pos {
631 // Fast path: the cache is an exact prefix of `full_context`. Feed only
632 // the new suffix at the live position; the existing KV is reused as-is.
633 for &t in &full_context[reuse..] {
634 self.last_logits = self.forward_one(t).map_err(ForwardOneError::into_edge)?;
635 self.cached.push(t);
636 }
637 } else {
638 // Divergence (or a shorter context): rebuild from scratch. Setting
639 // `index_pos = 0` makes the first `forward_one` discard candle's old
640 // cache, exactly as a fresh `prefill` would. Clear `last_logits` first
641 // so an *empty* `full_context` leaves no stale distribution behind —
642 // matching `reset_cache()` + `prefill(&[])`; a non-empty context
643 // overwrites it in the loop.
644 self.index_pos = 0;
645 self.cached = Vec::with_capacity(full_context.len());
646 self.last_logits = Vec::new();
647 for &t in full_context {
648 self.last_logits = self.forward_one(t).map_err(ForwardOneError::into_edge)?;
649 self.cached.push(t);
650 }
651 }
652 self.fed = 0;
653 // Replay base for this turn's rollback.
654 self.prompt = full_context.to_vec();
655 // Set `vocab` unconditionally — exactly as `prefill` does — so an empty
656 // rebuild leaves `vocab == 0`, matching `reset_cache()` + `prefill(&[])`;
657 // a non-empty context sets it to the real vocab via the fed logits.
658 self.vocab = self.last_logits.len();
659 Ok(self.index_pos as u32)
660 }
661}
662
663/// Length of the longest common prefix of two token slices. The cross-turn KV
664/// reuse cutoff (ADR-018 AC-3): how many leading tokens of a re-tokenized
665/// conversation still match what the engine already cached.
666fn longest_common_prefix(a: &[Token], b: &[Token]) -> usize {
667 a.iter().zip(b).take_while(|(x, y)| x == y).count()
668}
669
670// ── On-device safety wiring (ADR-005 tier + ADR-012 control loop) ────────────
671//
672// The runtime ships the *primitives* (steerer, chunk guard, checkpointed
673// rollback). They only engage when a session is given a real steerer + guard in
674// its `Ports` — `Ports::permissive()` wires neither, so a provider must opt in.
675// This adapter does: it owns the tokenizer, so it is the one place that can turn
676// a human-readable unsafe-word list into the token-id patterns the runtime's
677// float-free guard consumes. The resolved patterns/bans then drive the standard
678// `InferenceSession::generate` control loop — nothing in the SDK is bypassed.
679
680/// A small, conservative built-in `Lightweight` safety list (ADR-005). These are
681/// unambiguous weapons/mass-harm manufacture terms — content the decode-time
682/// guard should never let the model emit. It is intentionally narrow to avoid
683/// false positives in ordinary chat; production swaps in the active tier's real
684/// safety model (the LoRA adapter / classifier of ADR-012's model inventory).
685const DEFAULT_UNSAFE_WORDS: &[&str] = &[
686 "bomb",
687 "explosive",
688 "detonator",
689 "methamphetamine",
690 "ricin",
691 "anthrax",
692 "sarin",
693 "nerve agent",
694];
695
696/// Deterministic hard refusal emitted when the control loop fails closed —
697/// rollbacks exhausted, no safe checkpoint, or refused at ingress (ADR-012
698/// §"Bounded rollback, fail-closed"; ADR-013 ingress triage).
699const SAFETY_REFUSAL: &str = "I can't help with that request.";
700
701/// Contrastive steering is restricted to the top-K base-logit tokens (ADR-013 /
702/// SafeDecoding): it keeps the per-step adjustment small (so the runtime's
703/// `pick` stays linear in the vocab) and avoids amplifying long-tail noise.
704const CONTRASTIVE_TOP_K: usize = 64;
705
706/// Encode a word to the token-id sequence(s) the model may actually emit for
707/// it. A word tokenizes differently depending on what precedes it, so both the
708/// **leading-space** form (mid-sentence, after another token) and the **bare**
709/// form (start of a line/turn or after punctuation) are returned — each as a
710/// distinct anchor n-gram. Empty/duplicate encodings are dropped, so the caller
711/// gets only matchable patterns.
712fn word_to_patterns(tokenizer: &Tokenizer, word: &str) -> Vec<Vec<Token>> {
713 let mut out: Vec<Vec<Token>> = Vec::new();
714 for variant in [format!(" {word}"), word.to_string()] {
715 if let Ok(enc) = tokenizer.encode(variant, false) {
716 let seq = enc.get_ids().to_vec();
717 if !seq.is_empty() && !out.contains(&seq) {
718 out.push(seq);
719 }
720 }
721 }
722 out
723}
724
725/// Resolved safety wiring for the provider, derived from the tokenizer once at
726/// construction. Holds only token-id data, so rebuilding a turn's `Ports` is a
727/// cheap clone with no tokenizer access on the hot path.
728#[derive(Debug, Clone)]
729struct SafetyConfig {
730 /// The ADR-005 tier. `Off` runs the plain single-pass decode (legacy path).
731 mode: SafetyMode,
732 /// Hard-banned single tokens — the always-on per-step `LightweightFilter`
733 /// layer (only words that encode to exactly one token; banning a shared
734 /// subword would be too blunt).
735 banned: Vec<Token>,
736 /// Built-in unsafe token-id n-grams — drive **both** the ADR-012 output
737 /// chunk guard and the ADR-013 prompt ingress triage.
738 patterns: Vec<Vec<Token>>,
739 /// Caller-supplied `--guard-word` n-grams — a **guard-only** demo/test hook.
740 /// Deliberately excluded from ingress so the documented rollback demo
741 /// (`--guard-word banana --prompt "…banana…"`) fires the *trajectory* loop
742 /// instead of refusing the prompt before decoding.
743 extra_guard_patterns: Vec<Vec<Token>>,
744}
745
746impl SafetyConfig {
747 /// Resolve the built-in `Lightweight` list against `tokenizer`.
748 fn lightweight(tokenizer: &Tokenizer) -> Self {
749 let mut banned = Vec::new();
750 let mut patterns = Vec::new();
751 for &word in DEFAULT_UNSAFE_WORDS {
752 for seq in word_to_patterns(tokenizer, word) {
753 // A single-token form is safe to hard-ban per step; multi-token
754 // forms are caught by the guard (banning a shared subword could
755 // hurt benign text).
756 if seq.len() == 1 && !banned.contains(&seq[0]) {
757 banned.push(seq[0]);
758 }
759 if !patterns.contains(&seq) {
760 patterns.push(seq);
761 }
762 }
763 }
764 Self {
765 mode: SafetyMode::Lightweight,
766 banned,
767 patterns,
768 extra_guard_patterns: Vec::new(),
769 }
770 }
771
772 /// Build the per-turn safety `Ports` (steerer + chunk guard) for this tier.
773 /// `Off`, or an empty list, yields no steering/guarding.
774 fn ports(&self) -> Ports {
775 let mut ports = Ports::permissive();
776 if matches!(self.mode, SafetyMode::Off) {
777 return ports;
778 }
779 let steerer: Box<dyn SafetySteerer> = if self.banned.is_empty() {
780 Box::new(NoSafety)
781 } else {
782 Box::new(LightweightFilter::new(self.banned.clone()))
783 };
784 ports.safety = steerer;
785 // Output chunk guard (ADR-012): built-in unsafe patterns + caller's
786 // --guard-word extras.
787 let guard_patterns: Vec<Vec<Token>> = self
788 .patterns
789 .iter()
790 .chain(self.extra_guard_patterns.iter())
791 .cloned()
792 .collect();
793 if !guard_patterns.is_empty() {
794 ports.guard = Some(Box::new(AnchorGuard::hard(guard_patterns)));
795 }
796 // Prompt ingress triage (ADR-013): built-in patterns ONLY — the
797 // --guard-word extras are a trajectory demo hook, not ingress refusals.
798 if !self.patterns.is_empty() {
799 ports.ingress = Some(Box::new(AnchorGuard::hard(self.patterns.clone())));
800 }
801 ports
802 }
803}
804
805/// Interior state of a [`QwenExpert`], guarded by one mutex so the
806/// rollback-detect-then-feed step is atomic.
807struct ExpertState {
808 engine: QwenEngine,
809 /// How many committed tokens the expert has fed since its last prime — used
810 /// to detect a base rollback (committed shrinks below this) and re-sync.
811 fed: usize,
812}
813
814/// A safety **expert** logit source for contrastive steering (ADR-013): a second
815/// Qwen engine — in production base + a safety LoRA; here any same-tokenizer Qwen
816/// GGUF — loaded through the ADR-006 provenance gate and **primed with the turn's
817/// prompt** so its logits align with the base engine's. The session feeds it the
818/// committed tokens via [`ExpertLogits::logits`].
819///
820/// The weights are **loaded once and kept resident** like the base model
821/// (ADR-018 expert persistence): the provider holds it across turns and calls
822/// [`reprime`](Self::reprime) per turn (`reset_cache` + prefill, no disk reload).
823/// State lives behind a `Mutex` (not `RefCell`/`Cell`) so the resident expert is
824/// `Send + Sync` and can sit in the `Send + Sync` provider.
825///
826/// Steering is bounded to the early-token window (ADR-013), so the expert runs
827/// only for the first `steer_window` tokens. When the base engine rolls back
828/// (committed output shrinks), the expert **re-primes to the prompt** and
829/// re-feeds the retained prefix, so its contrastive context stays aligned with
830/// the base rather than serving logits from the abandoned branch. Pointing this
831/// at the chat model itself yields ~zero contrast (a no-op); a safety-tuned Qwen
832/// GGUF gives real steering.
833pub struct QwenExpert {
834 state: std::sync::Mutex<ExpertState>,
835 /// Last-known vocab size, updated after each successful `logits()` call.
836 /// Enables returning zeros of the correct length on mutex poison, where
837 /// `ExpertState` is inaccessible. Initialized from the primed engine.
838 vocab: std::sync::atomic::AtomicUsize,
839 /// Evidence the expert weights passed the ADR-006 load gate (R5). Held for
840 /// the engine's lifetime; never used after construction.
841 _permit: LoadPermit,
842}
843
844impl QwenExpert {
845 /// Load the expert GGUF, gate it (ADR-006 — `permit` is required, not
846 /// optional), and prime it with `prompt` so its KV state matches the base
847 /// engine's post-prefill state.
848 pub fn from_path_primed(
849 path: impl AsRef<std::path::Path>,
850 eos: Token,
851 prompt: &[Token],
852 permit: LoadPermit,
853 ) -> Result<Self> {
854 let mut engine = QwenEngine::from_path(path, eos)?;
855 engine.prefill(prompt)?;
856 let init_vocab = engine.vocab;
857 Ok(Self {
858 state: std::sync::Mutex::new(ExpertState { engine, fed: 0 }),
859 vocab: std::sync::atomic::AtomicUsize::new(init_vocab),
860 _permit: permit,
861 })
862 }
863
864 /// Re-prime the **resident** expert to a new turn's prompt without reloading
865 /// the GGUF (ADR-018 expert persistence): discard the prior turn's KV and
866 /// prefill the new prompt on the same loaded weights. The expensive part —
867 /// reading + parsing the GGUF — happens once in `from_path_primed`; this only
868 /// re-runs the (cheap, bounded) prompt prefill.
869 pub fn reprime(&self, prompt: &[Token]) -> Result<()> {
870 let mut st = self
871 .state
872 .lock()
873 .map_err(|_| EdgeError::Engine("expert mutex poisoned"))?;
874 st.engine.reset_cache()?;
875 // reset_cache succeeded: engine is blank. Set fed to 0 before prefill so
876 // a prefill failure leaves fed consistent with the blank engine state.
877 st.fed = 0;
878 st.engine.prefill(prompt)?;
879 Ok(())
880 }
881
882 /// Release the expert's conversation KV while keeping its weights resident
883 /// (ADR-018) — the expert half of [`QwenChatProvider::end_session`].
884 fn release(&self) -> Result<()> {
885 let mut st = self
886 .state
887 .lock()
888 .map_err(|_| EdgeError::Engine("expert mutex poisoned"))?;
889 st.engine.reset_cache()?;
890 st.fed = 0;
891 Ok(())
892 }
893}
894
895impl ExpertLogits for QwenExpert {
896 fn logits(&self, committed: &[Token]) -> Vec<i32> {
897 let mut st = match self.state.lock() {
898 Ok(st) => st,
899 Err(_) => {
900 // Mutex is poisoned: a prior call panicked while holding the lock.
901 // Return zeros of the last-known vocab size — a neutral expert
902 // signal — rather than an empty vec that would mismatch the base
903 // logit length in the steerer.
904 let v = self.vocab.load(std::sync::atomic::Ordering::Relaxed);
905 return vec![0; v.max(1)];
906 }
907 };
908 // Base rolled back? `committed` shrank below what we've fed. Re-prime the
909 // expert to the prompt (QwenEngine::rollback replays the prompt and
910 // resets its feed cursor) so it re-feeds the retained prefix from a clean
911 // state — keeping the contrastive context aligned with the base. Cost is
912 // bounded by `max_rollbacks` (ADR-012), same as the base engine.
913 if committed.len() < st.fed {
914 if st.engine.rollback(committed.len() as u32).is_err() {
915 return vec![0; st.engine.vocab.max(1)];
916 }
917 st.fed = 0;
918 }
919 let out = st.engine.next_logits(committed);
920 let vocab = out.len();
921 if vocab > 0 {
922 self.vocab
923 .store(vocab, std::sync::atomic::Ordering::Relaxed);
924 }
925 st.fed = committed.len();
926 out
927 }
928}
929
930/// A resident expert shared into a turn's steerer (ADR-018 expert persistence):
931/// the weights are loaded once and held by the provider; each turn clones this
932/// `Arc` into the turn's [`ContrastiveSteerer`]. After the turn the steerer (and
933/// this clone) drops, but the provider's `Arc` keeps the weights resident. The
934/// newtype sidesteps the orphan rule (`ExpertLogits` cannot be implemented for a
935/// bare `Arc<QwenExpert>` outside `el-safety`).
936struct SharedExpert(std::sync::Arc<QwenExpert>);
937
938impl ExpertLogits for SharedExpert {
939 fn logits(&self, committed: &[Token]) -> Vec<i32> {
940 self.0.logits(committed)
941 }
942}
943
944/// The resident model behind a [`QwenChatProvider`] (ADR-018).
945///
946/// The weights are loaded **once** (in [`QwenChatProvider::from_paths`]) into
947/// `Loaded`. The first `chat` promotes them into a reusable [`InferenceSession`]
948/// (`Active`) — done lazily so the builder-configured safety tier / expert is
949/// finalized first — and every later turn reuses that one session: a follow-up
950/// turn via `continue_prompt` (reuse the cached KV prefix, AC-3), a fresh turn via
951/// `load_prompt`, and a turn after a mid-flight failure via `reset()` +
952/// `load_prompt`. The model is never re-read from disk per turn.
953enum ChatSession {
954 /// Weights resident, no conversation session yet.
955 Loaded(QwenEngine),
956 /// Reusable session wrapping the resident engine.
957 Active(InferenceSession<QwenEngine>),
958 /// Transient placeholder held only during the `Loaded` → `Active` swap.
959 Swapping,
960}
961
962/// A real local chat backend: a Qwen2 GGUF model + its tokenizer, driven
963/// through [`el_runtime::InferenceSession`].
964///
965/// The model weights are loaded **once** at construction and kept resident
966/// (ADR-018): each `chat` renders the whole conversation to Qwen2.5 ChatML and
967/// reuses one persistent provenance-gated session — a follow-up turn reuses the
968/// cached KV prefix and prefills only the new suffix (`continue_prompt`, AC-3
969/// cross-turn incremental prefill), while a fresh turn uses `load_prompt` —
970/// then `generate` (grammar mask → safety steer → guard + checkpointed rollback
971/// → greedy commit), instead of rebuilding the engine and re-reading the GGUF
972/// every turn.
973/// On-device safety (ADR-005 `Lightweight` tier + the ADR-012 control loop) is
974/// **on by default**; see [`with_safety`](Self::with_safety). The resident model
975/// lives behind a `Mutex`, so the provider stays `Send + Sync` and concurrent
976/// `chat` calls serialize on the one conversation.
977pub struct QwenChatProvider {
978 tokenizer: Tokenizer,
979 permit: LoadPermit,
980 eos: Token,
981 default_max_tokens: u32,
982 model_label: String,
983 safety: SafetyConfig,
984 /// Optional safety **expert** GGUF for ADR-013 contrastive steering. `None`
985 /// runs the token-only `Lightweight` steerer. The weights are loaded once and
986 /// kept resident in `expert` (ADR-018 expert persistence).
987 expert_model: Option<std::path::PathBuf>,
988 /// Contrastive steering strength ×1000 (1000 = 1.0×).
989 steer_alpha_milli: i32,
990 /// Resident model + reusable session (ADR-018).
991 session: std::sync::Mutex<ChatSession>,
992 /// The resident safety expert (ADR-018 expert persistence): loaded lazily on
993 /// the first `SecDecoding` turn and reused across turns via `reprime`, instead
994 /// of re-reading the expert GGUF from disk every turn. `None` until first use,
995 /// or always-`None` when no `--expert-model` is configured.
996 expert: std::sync::Mutex<Option<std::sync::Arc<QwenExpert>>>,
997}
998
999impl QwenChatProvider {
1000 /// Load a Qwen2 GGUF model and its `tokenizer.json` from local paths.
1001 pub fn from_paths(
1002 model_path: impl AsRef<std::path::Path>,
1003 tokenizer_path: impl AsRef<std::path::Path>,
1004 ) -> Result<Self> {
1005 let model_path = model_path.as_ref().to_path_buf();
1006 if !model_path.exists() {
1007 return Err(EdgeError::Engine("model file not found"));
1008 }
1009 let tokenizer = Tokenizer::from_file(tokenizer_path.as_ref())
1010 .map_err(|_| EdgeError::Engine("failed to load tokenizer.json"))?;
1011
1012 // Stop token: Qwen2.5 ChatML turn terminator (fallback to its known id).
1013 let eos = tokenizer.token_to_id("<|im_end|>").unwrap_or(151_645);
1014
1015 let model_label = model_path
1016 .file_stem()
1017 .and_then(|s| s.to_str())
1018 .map(|s| format!("local/{s}"))
1019 .unwrap_or_else(|| "local/qwen2".to_string());
1020
1021 let safety = SafetyConfig::lightweight(&tokenizer);
1022
1023 let permit = local_load_permit(&model_path)?;
1024
1025 // ADR-018: load the weights ONCE here and keep them resident, instead of
1026 // re-reading the GGUF on every `chat`. The first `chat` promotes this into
1027 // a reusable session (see `ChatSession`).
1028 let engine = QwenEngine::from_path(&model_path, eos)?;
1029
1030 Ok(Self {
1031 tokenizer,
1032 permit,
1033 eos,
1034 default_max_tokens: 512,
1035 model_label,
1036 safety,
1037 expert_model: None,
1038 steer_alpha_milli: 1000,
1039 session: std::sync::Mutex::new(ChatSession::Loaded(engine)),
1040 expert: std::sync::Mutex::new(None),
1041 })
1042 }
1043
1044 /// Select the on-device safety tier (ADR-005). [`SafetyMode::Off`] disables
1045 /// the steerer and the ADR-012 control loop (the plain single-pass decode);
1046 /// [`SafetyMode::Lightweight`] (the default) runs the token-anchor guard +
1047 /// hard-ban steerer + checkpointed rollback. `SecDecoding`/`Csd` need model
1048 /// assets not shipped here and fall back to the `Lightweight` wiring.
1049 pub fn with_safety(mut self, mode: SafetyMode) -> Self {
1050 self.safety.mode = mode;
1051 self
1052 }
1053
1054 /// Add extra words to the chunk guard's unsafe patterns (resolved to token
1055 /// ids via this model's tokenizer). Primarily a **test/demo hook**: e.g.
1056 /// `--guard-word banana` lets you watch the ADR-012 rollback / fail-closed
1057 /// refusal fire on a benign word, without needing the model to emit genuinely
1058 /// harmful content. Guard-only — these are not added to the hard-ban list, so
1059 /// the trajectory loop (not silent suppression) is what engages.
1060 pub fn with_extra_guard_words<I, S>(mut self, words: I) -> Self
1061 where
1062 I: IntoIterator<Item = S>,
1063 S: AsRef<str>,
1064 {
1065 for word in words {
1066 for seq in word_to_patterns(&self.tokenizer, word.as_ref()) {
1067 if !self.safety.extra_guard_patterns.contains(&seq) {
1068 self.safety.extra_guard_patterns.push(seq);
1069 }
1070 }
1071 }
1072 self
1073 }
1074
1075 /// Enable model-backed **contrastive** steering (ADR-013) with a safety
1076 /// **expert** GGUF (same tokenizer/family as the chat model). Steering runs
1077 /// only inside the early-token window. Pointing this at the chat model itself
1078 /// gives ~zero contrast (a no-op); a safety-tuned Qwen GGUF gives real
1079 /// steering. No effect under `--safety off`.
1080 pub fn with_expert_model(mut self, path: impl AsRef<std::path::Path>) -> Self {
1081 self.expert_model = Some(path.as_ref().to_path_buf());
1082 self
1083 }
1084
1085 /// Contrastive steering strength ×1000 (`1000` = 1.0×). Only meaningful with
1086 /// [`with_expert_model`](Self::with_expert_model).
1087 pub fn with_steer_alpha(mut self, alpha_milli: i32) -> Self {
1088 self.steer_alpha_milli = alpha_milli;
1089 self
1090 }
1091
1092 /// End the current conversation, releasing its KV / output / prompt / buffered
1093 /// events **while keeping the model resident** (ADR-018 separation of
1094 /// conversation and model lifecycles; the AC-4 explicit release / PRD line 131
1095 /// "KV caches … cleared on session end"). The next `chat` starts a fresh
1096 /// conversation on the same loaded weights — no reload. A no-op if no
1097 /// conversation has started yet. To free the weights too, drop the provider
1098 /// (Rust ownership).
1099 pub fn end_session(&self) -> Result<()> {
1100 let mut cell = self
1101 .session
1102 .lock()
1103 .map_err(|_| EdgeError::Engine("chat session mutex poisoned"))?;
1104 if let ChatSession::Active(session) = &mut *cell {
1105 session.close()?;
1106 }
1107 // Release the resident safety expert's conversation KV too (keeps its
1108 // weights). Locked after the session — the same order `chat` uses — so
1109 // the two mutexes never deadlock.
1110 let slot = self
1111 .expert
1112 .lock()
1113 .map_err(|_| EdgeError::Engine("expert mutex poisoned"))?;
1114 if let Some(expert) = slot.as_ref() {
1115 expert.release()?;
1116 }
1117 Ok(())
1118 }
1119
1120 fn encode(&self, text: &str) -> Result<Vec<Token>> {
1121 let enc = self
1122 .tokenizer
1123 .encode(text, false)
1124 .map_err(|_| EdgeError::Engine("tokenizer encode failed"))?;
1125 Ok(enc.get_ids().to_vec())
1126 }
1127
1128 fn decode(&self, ids: &[Token]) -> Result<String> {
1129 self.tokenizer
1130 .decode(ids, true)
1131 .map_err(|_| EdgeError::Engine("tokenizer decode failed"))
1132 }
1133}
1134
1135impl LlmProvider for QwenChatProvider {
1136 fn chat(&self, req: &ChatRequest) -> Result<ChatResponse> {
1137 let prompt = render_chatml(&req.messages);
1138
1139 let t_encode = bench::enabled().then(std::time::Instant::now);
1140 let prompt_tokens = self.encode(&prompt)?;
1141 let d_encode = t_encode.map(|t| t.elapsed()).unwrap_or_default();
1142
1143 // Carry the active safety tier on the session config so the runtime
1144 // derives the tier-aware ADR-012 `RollbackPolicy` and records the true
1145 // mode. A supplied expert promotes the tier to `SecDecoding`, so the
1146 // runtime's `SafetyModeSelector` can gate it on device class instead of
1147 // it masquerading as `Lightweight`. Both are deterministic from the
1148 // builder-set config, so they are identical on every turn.
1149 let requested = requested_session_safety(self.safety.mode, self.expert_model.is_some());
1150 let cfg = SessionConfig {
1151 safety: requested,
1152 ..SessionConfig::default()
1153 };
1154 // Resolve the same effective tier the runtime will: only install the
1155 // contrastive steerer if `SecDecoding` survives device selection (it
1156 // downgrades to `Lightweight` on non-accelerator devices, where the
1157 // expert is dropped — honest tier-aware behaviour).
1158 let effective = SafetyModeSelector::resolve(requested, cfg.device);
1159
1160 // ADR-018: reuse the resident model. Lock the session cell; on first use
1161 // promote the loaded weights into a reusable session (created with the now
1162 // final builder config); every later turn reuses it — no disk reload.
1163 let mut cell = self
1164 .session
1165 .lock()
1166 .map_err(|_| EdgeError::Engine("chat session mutex poisoned"))?;
1167 let t_load = bench::enabled().then(std::time::Instant::now);
1168 if matches!(&*cell, ChatSession::Loaded(_)) {
1169 let engine = match std::mem::replace(&mut *cell, ChatSession::Swapping) {
1170 ChatSession::Loaded(e) => e,
1171 _ => unreachable!("guarded by the matches! above"),
1172 };
1173 *cell = ChatSession::Active(InferenceSession::new(
1174 SessionId(1),
1175 cfg,
1176 engine,
1177 self.permit,
1178 ));
1179 }
1180 let d_load = t_load.map(|t| t.elapsed()).unwrap_or_default();
1181 let session = match &mut *cell {
1182 ChatSession::Active(s) => s,
1183 // `Swapping` only persists if a prior promotion panicked — in which
1184 // case `lock()` above would already have failed on the poisoned mutex.
1185 _ => return Err(EdgeError::Engine("chat session not initialized")),
1186 };
1187
1188 // Provider-owned turn isolation: drop any events buffered by a prior turn
1189 // (e.g. one that errored before its end-of-turn drain) so this turn's
1190 // safety count below cannot include another turn's stale violations.
1191 // `continue_prompt`/`reset` deliberately preserve events (generic
1192 // semantics); bounding them per turn is the reusing provider's job.
1193 let _ = session.drain_events();
1194 let mut ports = self.safety.ports();
1195
1196 if matches!(effective, SafetyMode::SecDecoding) {
1197 if let Some(expert_path) = &self.expert_model {
1198 // ADR-018 expert persistence: load the expert weights ONCE and keep
1199 // them resident; every later turn re-primes (no disk reload). The
1200 // expert lock is always taken while holding the session lock — a
1201 // fixed order, so no deadlock with `end_session`.
1202 let expert = {
1203 let mut slot = self
1204 .expert
1205 .lock()
1206 .map_err(|_| EdgeError::Engine("expert mutex poisoned"))?;
1207 match slot.as_ref() {
1208 Some(e) => {
1209 e.reprime(&prompt_tokens)?;
1210 std::sync::Arc::clone(e)
1211 }
1212 None => {
1213 let e = std::sync::Arc::new(QwenExpert::from_path_primed(
1214 expert_path,
1215 self.eos,
1216 &prompt_tokens,
1217 local_load_permit(expert_path)?,
1218 )?);
1219 *slot = Some(std::sync::Arc::clone(&e));
1220 e
1221 }
1222 }
1223 };
1224 ports.safety = Box::new(ContrastiveSteerer::new(
1225 SharedExpert(expert),
1226 self.safety.banned.clone(),
1227 self.steer_alpha_milli,
1228 CONTRASTIVE_TOP_K,
1229 effective,
1230 ));
1231 }
1232 }
1233
1234 let _ = bench::take(); // clear forward accumulators before prefill
1235 let t_prefill = bench::enabled().then(std::time::Instant::now);
1236 // ADR-018 AC-3: on a follow-up turn (a finished prior turn left the
1237 // session `Completed`), reuse the cached KV prefix and prefill only the
1238 // new suffix; a fresh turn does a full prefill. The engine's
1239 // longest-common-prefix check is the backstop, so a tokenizer round-trip
1240 // drift in the reused branch falls back to a correct full re-prefill.
1241 match session.phase() {
1242 el_core::Phase::Completed => session.continue_prompt(&ports, &prompt_tokens)?,
1243 // First use, or a fresh start after `end_session` / error recovery.
1244 el_core::Phase::Initialized => session.load_prompt(&ports, &prompt_tokens)?,
1245 // Dirty: a prior turn's prefill failed mid-transition and left the
1246 // session in `Prefilling`/`Decoding`. Now that the unconditional
1247 // per-turn `reset()` is gone, a bare `load_prompt` here would hit
1248 // `InvalidPhase` and wedge the provider — so discard the partial
1249 // conversation (clearing the engine's possibly half-fed cache) and
1250 // start fresh instead.
1251 _ => {
1252 let dirty_phase = session.phase().as_str();
1253 session.reset()?;
1254 session.load_prompt(&ports, &prompt_tokens)?;
1255 eprintln!(
1256 "[session] partial state ({dirty_phase}) detected — context reset, this turn starts fresh"
1257 );
1258 }
1259 }
1260 let d_prefill = t_prefill.map(|t| t.elapsed()).unwrap_or_default();
1261 let (pf_total, pf_model, pf_calls) = bench::take();
1262
1263 let max = req.max_tokens.unwrap_or(self.default_max_tokens);
1264 let t_decode = bench::enabled().then(std::time::Instant::now);
1265 let stop = session.generate(&ports, max)?;
1266 let d_decode = t_decode.map(|t| t.elapsed()).unwrap_or_default();
1267 let (dc_total, dc_model, dc_calls) = bench::take();
1268
1269 let out = session.output().to_vec();
1270 let completion_tokens = out.len() as u32;
1271
1272 let t_detok = bench::enabled().then(std::time::Instant::now);
1273 let decoded = self.decode(&out)?.trim().to_string();
1274 let d_detok = t_detok.map(|t| t.elapsed()).unwrap_or_default();
1275
1276 // ADR-018: always drain — a persistent session would otherwise accumulate
1277 // events across turns. ADR-012: surface what the decode-time control loop
1278 // did. A fail-closed stop (rollbacks exhausted / no safe checkpoint)
1279 // returns the deterministic refusal rather than the truncated unsafe
1280 // prefix; any intervention is reported on stderr so the test client can
1281 // show the guard working without corrupting the reply on stdout.
1282 let events = session.drain_events();
1283 let safety_active = !matches!(self.safety.mode, SafetyMode::Off);
1284 let content = if safety_active {
1285 let violations = events
1286 .iter()
1287 .filter(|e| matches!(e.event, DomainEvent::SafetyViolationDetected { .. }))
1288 .count();
1289 let rollbacks = events
1290 .iter()
1291 .filter(|e| matches!(e.event, DomainEvent::ClaimBacktracked { .. }))
1292 .count();
1293 let refused = stop == StopReason::Stopped && violations > 0;
1294 if violations > 0 || rollbacks > 0 {
1295 eprintln!(
1296 "[safety] {violations} violation(s), {rollbacks} rollback(s){}",
1297 if refused {
1298 " → refused (fail-closed)"
1299 } else {
1300 " → recovered"
1301 }
1302 );
1303 }
1304 if refused {
1305 SAFETY_REFUSAL.to_string()
1306 } else {
1307 decoded
1308 }
1309 } else {
1310 decoded
1311 };
1312
1313 if bench::enabled() {
1314 report_breakdown(
1315 prompt_tokens.len() as u32,
1316 completion_tokens,
1317 d_load,
1318 d_encode,
1319 d_prefill,
1320 d_decode,
1321 d_detok,
1322 (pf_total, pf_model, pf_calls),
1323 (dc_total, dc_model, dc_calls),
1324 );
1325 }
1326
1327 Ok(ChatResponse {
1328 content,
1329 model: self.model_label.clone(),
1330 prompt_tokens: prompt_tokens.len() as u32,
1331 completion_tokens,
1332 })
1333 }
1334
1335 fn chat_stream(&self, req: &ChatRequest, on_token: &mut dyn FnMut(ChatToken)) -> Result<()> {
1336 // The runtime decode loop runs to completion internally (no per-token
1337 // hook), so — like the toy `LocalLlmProvider` — we stream the finished
1338 // reply out character by character.
1339 let resp = self.chat(req)?;
1340 for ch in resp.content.chars() {
1341 on_token(ChatToken {
1342 text: ch.to_string(),
1343 is_final: false,
1344 });
1345 }
1346 on_token(ChatToken {
1347 text: String::new(),
1348 is_final: true,
1349 });
1350 Ok(())
1351 }
1352}
1353
1354/// Print an `EL_BENCH` per-phase + per-forward breakdown for one `chat()` call.
1355#[allow(clippy::too_many_arguments)]
1356fn report_breakdown(
1357 prompt_tokens: u32,
1358 completion_tokens: u32,
1359 d_load: std::time::Duration,
1360 d_encode: std::time::Duration,
1361 d_prefill: std::time::Duration,
1362 d_decode: std::time::Duration,
1363 d_detok: std::time::Duration,
1364 prefill_fwd: (std::time::Duration, std::time::Duration, u64),
1365 decode_fwd: (std::time::Duration, std::time::Duration, u64),
1366) {
1367 let ms = |d: std::time::Duration| d.as_secs_f64() * 1000.0;
1368 let total = d_load + d_encode + d_prefill + d_decode + d_detok;
1369 let pct = |d: std::time::Duration| {
1370 if total.as_secs_f64() > 0.0 {
1371 d.as_secs_f64() / total.as_secs_f64() * 100.0
1372 } else {
1373 0.0
1374 }
1375 };
1376 let tps = |n: u32, d: std::time::Duration| {
1377 if d.as_secs_f64() > 0.0 {
1378 n as f64 / d.as_secs_f64()
1379 } else {
1380 0.0
1381 }
1382 };
1383
1384 let (pf_total, pf_model, pf_calls) = prefill_fwd;
1385 let (dc_total, dc_model, dc_calls) = decode_fwd;
1386 let dc_loop = d_decode.saturating_sub(dc_total);
1387 let dc_seam = dc_total.saturating_sub(dc_model);
1388 let per_tok = |d: std::time::Duration, n: u64| if n > 0 { ms(d) / n as f64 } else { 0.0 };
1389
1390 eprintln!("\n┌─ EL_BENCH chat() breakdown ───────────────────────────────");
1391 eprintln!("│ prompt_tokens={prompt_tokens} completion_tokens={completion_tokens}");
1392 eprintln!("│ phase wall(ms) %total throughput");
1393 eprintln!(
1394 "│ session setup {:>9.1} {:>6.1}% (weights loaded once at startup — ADR-018)",
1395 ms(d_load),
1396 pct(d_load)
1397 );
1398 eprintln!(
1399 "│ tokenize {:>9.2} {:>6.1}%",
1400 ms(d_encode),
1401 pct(d_encode)
1402 );
1403 eprintln!(
1404 "│ prefill {:>9.1} {:>6.1}% {:>7.1} tok/s",
1405 ms(d_prefill),
1406 pct(d_prefill),
1407 tps(prompt_tokens, d_prefill)
1408 );
1409 eprintln!(
1410 "│ decode {:>9.1} {:>6.1}% {:>7.1} tok/s",
1411 ms(d_decode),
1412 pct(d_decode),
1413 tps(completion_tokens, d_decode)
1414 );
1415 eprintln!(
1416 "│ detokenize {:>9.2} {:>6.1}%",
1417 ms(d_detok),
1418 pct(d_detok)
1419 );
1420 eprintln!("│ TOTAL {:>9.1}", ms(total));
1421 eprintln!("│ ─ forward attribution (where prefill+decode time goes) ─");
1422 eprintln!(
1423 "│ prefill: {} fwd calls, model {:.1}ms, seam {:.1}ms, loop {:.1}ms",
1424 pf_calls,
1425 ms(pf_model),
1426 ms(pf_total.saturating_sub(pf_model)),
1427 ms(d_prefill.saturating_sub(pf_total)),
1428 );
1429 eprintln!(
1430 "│ decode : {} fwd calls, model {:.1}ms, seam {:.1}ms, loop {:.1}ms",
1431 dc_calls,
1432 ms(dc_model),
1433 ms(dc_seam),
1434 ms(dc_loop),
1435 );
1436 eprintln!(
1437 "│ per decoded token: {:.2}ms total = model {:.2} + seam {:.2} + loop {:.2}",
1438 per_tok(d_decode, dc_calls),
1439 per_tok(dc_model, dc_calls),
1440 per_tok(dc_seam, dc_calls),
1441 per_tok(dc_loop, dc_calls),
1442 );
1443 eprintln!("└───────────────────────────────────────────────────────────");
1444}
1445
1446/// Render a conversation as Qwen2.5 ChatML and open an assistant turn.
1447fn render_chatml(messages: &[ChatMessage]) -> String {
1448 let mut s = String::new();
1449 for m in messages {
1450 let role = match m.role {
1451 ChatRole::System => "system",
1452 ChatRole::User => "user",
1453 ChatRole::Assistant => "assistant",
1454 };
1455 s.push_str("<|im_start|>");
1456 s.push_str(role);
1457 s.push('\n');
1458 s.push_str(&m.content);
1459 s.push_str("<|im_end|>\n");
1460 }
1461 s.push_str("<|im_start|>assistant\n");
1462 s
1463}
1464
1465fn requested_session_safety(configured: SafetyMode, has_expert: bool) -> SafetyMode {
1466 match (configured, has_expert) {
1467 (SafetyMode::Off, _) => SafetyMode::Off,
1468 // A supplied expert is the only backed SecDecoding implementation in
1469 // this adapter. Promote any non-Off configured tier to that concrete
1470 // model-backed path so the runtime selector can gate it by device.
1471 (_, true) => SafetyMode::SecDecoding,
1472 // These public enum variants are not backed here without an expert.
1473 // Keep telemetry/policy honest by reflecting the lightweight ports that
1474 // will actually be installed.
1475 (SafetyMode::SecDecoding | SafetyMode::Csd, false) => SafetyMode::Lightweight,
1476 (mode, false) => mode,
1477 }
1478}
1479
1480/// Obtain a [`LoadPermit`] through the ADR-006 gate for a user-supplied local
1481/// model. There is no detached signature to check for a file the user downloaded
1482/// themselves, so this uses a trust-the-local-file verifier. This is explicitly
1483/// **not** cryptographic integrity over the GGUF bytes; production signed assets
1484/// must use a separate verifier path that reads the whole artifact and verifies
1485/// its detached signature before issuing a permit.
1486fn local_load_permit(path: &std::path::Path) -> Result<LoadPermit> {
1487 struct LocalFileTrust;
1488 impl SignatureVerifier for LocalFileTrust {
1489 fn verify(&self, _bytes: &[u8], _sig: &[u8], _key: u32) -> bool {
1490 true
1491 }
1492 }
1493 // Keep the local-trust path cheap: it proves callers go through the permit
1494 // gate, while deliberately avoiding fake "verification" of path strings or
1495 // header fragments that could be mistaken for artifact integrity.
1496 let _ = path;
1497 let mut artifact = ModelArtifact::new(
1498 ModelId(1),
1499 ModelVersion::new(0, 1, 0),
1500 el_core::ModelFormat::Gguf,
1501 );
1502 artifact.verify(&LocalFileTrust, b"local-trust", b"", 0);
1503 artifact.ensure_loadable()
1504}
1505
1506#[cfg(test)]
1507mod tests {
1508 use super::*;
1509 use el_runtime::InferenceEngine;
1510
1511 // ── helpers ──────────────────────────────────────────────────────────────
1512
1513 fn ok_permit() -> LoadPermit {
1514 use el_core::{ModelFormat, ModelId, ModelVersion};
1515 use el_provenance::{ModelArtifact, SignatureVerifier};
1516 struct OkV;
1517 impl SignatureVerifier for OkV {
1518 fn verify(&self, _: &[u8], _: &[u8], _: u32) -> bool {
1519 true
1520 }
1521 }
1522 let mut a = ModelArtifact::new(ModelId(1), ModelVersion::new(0, 1, 0), ModelFormat::Gguf);
1523 a.verify(&OkV, b"w", b"s", 0);
1524 a.ensure_loadable().unwrap()
1525 }
1526
1527 /// Build a minimal but spec-compliant GGUF v3 file in memory.
1528 ///
1529 /// Layout: no KV metadata, two F32 tensors:
1530 /// `token_embd.weight` [vocab, dim] at offset 0
1531 /// `output.weight` [vocab, dim] at offset vocab*dim*4
1532 ///
1533 /// GGUF stores dimensions innermost-first; candle reverses them on read.
1534 fn make_minimal_gguf(vocab: usize, dim: usize) -> Vec<u8> {
1535 let mut w: Vec<u8> = Vec::new();
1536
1537 // Header
1538 w.extend_from_slice(b"GGUF");
1539 w.extend_from_slice(&3u32.to_le_bytes()); // version 3
1540 w.extend_from_slice(&2u64.to_le_bytes()); // n_tensors
1541 w.extend_from_slice(&0u64.to_le_bytes()); // n_kv (none)
1542
1543 let tensor_bytes = (vocab * dim * 4) as u64;
1544
1545 // token_embd.weight: [vocab, dim] → GGUF dims [dim, vocab]
1546 let name = b"token_embd.weight";
1547 w.extend_from_slice(&(name.len() as u64).to_le_bytes());
1548 w.extend_from_slice(name);
1549 w.extend_from_slice(&2u32.to_le_bytes());
1550 w.extend_from_slice(&(dim as u64).to_le_bytes()); // innermost
1551 w.extend_from_slice(&(vocab as u64).to_le_bytes()); // outermost
1552 w.extend_from_slice(&0u32.to_le_bytes()); // F32
1553 w.extend_from_slice(&0u64.to_le_bytes()); // offset 0
1554
1555 // output.weight: [vocab, dim] → GGUF dims [dim, vocab]; loader will transpose
1556 let name = b"output.weight";
1557 w.extend_from_slice(&(name.len() as u64).to_le_bytes());
1558 w.extend_from_slice(name);
1559 w.extend_from_slice(&2u32.to_le_bytes());
1560 w.extend_from_slice(&(dim as u64).to_le_bytes());
1561 w.extend_from_slice(&(vocab as u64).to_le_bytes());
1562 w.extend_from_slice(&0u32.to_le_bytes());
1563 w.extend_from_slice(&tensor_bytes.to_le_bytes()); // offset after embed
1564
1565 // Pad to 32-byte alignment
1566 let pad = (32usize.wrapping_sub(w.len() % 32)) % 32;
1567 w.resize(w.len() + pad, 0u8);
1568
1569 // Tensor data (both tensors, row-major f32)
1570 for i in 0..(vocab * dim * 2) {
1571 w.extend_from_slice(&(i as f32 * 0.1f32).to_le_bytes());
1572 }
1573
1574 w
1575 }
1576
1577 // ── toy-model tests (unchanged) ──────────────────────────────────────────
1578
1579 #[test]
1580 fn real_candle_forward_is_deterministic_and_right_shape() {
1581 let mut eng = CandleEngine::toy(8, 4, 7).unwrap();
1582 let a = eng.next_logits(&[2]);
1583 let b = eng.next_logits(&[2]);
1584 assert_eq!(a.len(), 8, "logits length == vocab");
1585 assert_eq!(a, b, "fixed weights → deterministic real-tensor forward");
1586 let c = eng.next_logits(&[5]);
1587 assert_ne!(a, c);
1588 }
1589
1590 #[test]
1591 fn drives_the_runtime_end_to_end() {
1592 use el_core::{ModelFormat, ModelId, ModelVersion, SessionConfig, SessionId, StopReason};
1593 use el_provenance::{ModelArtifact, SignatureVerifier};
1594
1595 struct OkVerifier;
1596 impl SignatureVerifier for OkVerifier {
1597 fn verify(&self, _: &[u8], _: &[u8], _: u32) -> bool {
1598 true
1599 }
1600 }
1601 let mut art = ModelArtifact::new(
1602 ModelId(1),
1603 ModelVersion::new(0, 1, 0),
1604 ModelFormat::Safetensors,
1605 );
1606 art.verify(&OkVerifier, b"w", b"s", 1);
1607 let permit = art.ensure_loadable().unwrap();
1608
1609 let eng = CandleEngine::toy(16, 8, 9999).unwrap();
1610 let mut session =
1611 InferenceSession::new(SessionId(1), SessionConfig::default(), eng, permit);
1612 let ports = Ports::permissive();
1613 session.load_prompt(&ports, &[1, 2, 3]).unwrap();
1614
1615 let stop = session.generate(&ports, 4).unwrap();
1616 assert_eq!(stop, StopReason::MaxTokens);
1617 assert_eq!(session.output().len(), 4);
1618 }
1619
1620 // ── GGUF loading tests ───────────────────────────────────────────────────
1621
1622 #[test]
1623 fn from_bytes_rejects_invalid_magic() {
1624 let r = CandleEngine::from_bytes(b"not a gguf file", 0);
1625 assert!(matches!(r, Err(EdgeError::Engine(_))));
1626 }
1627
1628 #[test]
1629 fn from_bytes_loads_minimal_gguf_and_forward_has_correct_vocab() {
1630 let vocab = 8;
1631 let dim = 4;
1632 let gguf = make_minimal_gguf(vocab, dim);
1633 let mut engine = CandleEngine::from_bytes(&gguf, 7).unwrap();
1634
1635 let logits = engine.next_logits(&[0]);
1636 assert_eq!(logits.len(), vocab, "logit vec width == vocab from GGUF");
1637 assert_eq!(engine.eos_token(), 7);
1638 }
1639
1640 #[test]
1641 fn from_bytes_gguf_forward_is_deterministic() {
1642 let gguf = make_minimal_gguf(8, 4);
1643 let mut eng = CandleEngine::from_bytes(&gguf, 0).unwrap();
1644 assert_eq!(eng.next_logits(&[3]), eng.next_logits(&[3]));
1645 }
1646
1647 /// Same as `make_minimal_gguf` but `output.weight` has `wrong_dim` instead of `dim`,
1648 /// so the embed / output dimensions are incompatible.
1649 fn make_mismatched_gguf(vocab: usize, embed_dim: usize, output_dim: usize) -> Vec<u8> {
1650 let mut w: Vec<u8> = Vec::new();
1651 w.extend_from_slice(b"GGUF");
1652 w.extend_from_slice(&3u32.to_le_bytes());
1653 w.extend_from_slice(&2u64.to_le_bytes());
1654 w.extend_from_slice(&0u64.to_le_bytes());
1655
1656 let embed_bytes = (vocab * embed_dim * 4) as u64;
1657
1658 let name = b"token_embd.weight";
1659 w.extend_from_slice(&(name.len() as u64).to_le_bytes());
1660 w.extend_from_slice(name);
1661 w.extend_from_slice(&2u32.to_le_bytes());
1662 w.extend_from_slice(&(embed_dim as u64).to_le_bytes());
1663 w.extend_from_slice(&(vocab as u64).to_le_bytes());
1664 w.extend_from_slice(&0u32.to_le_bytes());
1665 w.extend_from_slice(&0u64.to_le_bytes());
1666
1667 let name = b"output.weight";
1668 w.extend_from_slice(&(name.len() as u64).to_le_bytes());
1669 w.extend_from_slice(name);
1670 w.extend_from_slice(&2u32.to_le_bytes());
1671 w.extend_from_slice(&(output_dim as u64).to_le_bytes()); // wrong dim
1672 w.extend_from_slice(&(vocab as u64).to_le_bytes());
1673 w.extend_from_slice(&0u32.to_le_bytes());
1674 w.extend_from_slice(&embed_bytes.to_le_bytes());
1675
1676 let pad = (32usize.wrapping_sub(w.len() % 32)) % 32;
1677 w.resize(w.len() + pad, 0u8);
1678
1679 for i in 0..(vocab * embed_dim + vocab * output_dim) {
1680 w.extend_from_slice(&(i as f32 * 0.1f32).to_le_bytes());
1681 }
1682 w
1683 }
1684
1685 #[test]
1686 fn from_path_missing_file_returns_engine_error() {
1687 let r = CandleEngine::from_path(std::path::Path::new("/nonexistent/model.gguf"), 0);
1688 assert!(matches!(r, Err(EdgeError::Engine(_))));
1689 }
1690
1691 #[test]
1692 fn from_bytes_rejects_mismatched_output_dim_at_load_time() {
1693 // embed dim=4, output dim=7 — incompatible; must error at load, not silently at forward.
1694 let gguf = make_mismatched_gguf(8, 4, 7);
1695 let r = CandleEngine::from_bytes(&gguf, 0);
1696 assert!(
1697 matches!(r, Err(EdgeError::Engine(_))),
1698 "mismatched output weight dim must be rejected at load time"
1699 );
1700 }
1701
1702 // ── LocalLlmProvider tests (unchanged + new from_path error path) ────────
1703
1704 #[test]
1705 fn local_provider_chat_returns_response() {
1706 let p = LocalLlmProvider::toy(32, 8, 31, ok_permit()).unwrap();
1707 let req = el_core::ChatRequest::new("local", vec![el_core::ChatMessage::user("hello")])
1708 .with_max_tokens(4);
1709 let resp = p.chat(&req).unwrap();
1710 assert_eq!(resp.model, "local/candle");
1711 assert_eq!(resp.completion_tokens, 4);
1712 assert!(!resp.content.is_empty());
1713 }
1714
1715 #[test]
1716 fn local_provider_stream_ends_with_final_token() {
1717 let p = LocalLlmProvider::toy(32, 8, 31, ok_permit()).unwrap();
1718 let req = el_core::ChatRequest::new("local", vec![el_core::ChatMessage::user("hi")])
1719 .with_max_tokens(3);
1720 let mut tokens: Vec<el_core::ChatToken> = Vec::new();
1721 p.chat_stream(&req, &mut |t| tokens.push(t)).unwrap();
1722 assert!(tokens.last().unwrap().is_final);
1723 assert!(tokens.len() > 1);
1724 }
1725
1726 #[test]
1727 fn local_provider_session_resets_between_calls() {
1728 let p = LocalLlmProvider::toy(32, 8, 31, ok_permit()).unwrap();
1729 let req = el_core::ChatRequest::new("local", vec![el_core::ChatMessage::user("a")])
1730 .with_max_tokens(4);
1731 let r1 = p.chat(&req).unwrap();
1732 let r2 = p.chat(&req).unwrap();
1733 assert_eq!(r1.content, r2.content);
1734 }
1735
1736 #[test]
1737 fn local_provider_from_path_missing_file_returns_error() {
1738 let r = LocalLlmProvider::from_path(
1739 std::path::Path::new("/nonexistent/model.gguf"),
1740 0,
1741 ok_permit(),
1742 );
1743 assert!(matches!(r, Err(EdgeError::Engine(_))));
1744 }
1745
1746 // ── Qwen provider helpers ─────────────────────────────────────────────────
1747
1748 #[test]
1749 fn render_chatml_wraps_each_turn_and_opens_assistant() {
1750 let msgs = vec![
1751 ChatMessage::system("be nice"),
1752 ChatMessage::user("hi"),
1753 ChatMessage::assistant("hello"),
1754 ChatMessage::user("bye"),
1755 ];
1756 let got = render_chatml(&msgs);
1757 let want = "<|im_start|>system\nbe nice<|im_end|>\n\
1758 <|im_start|>user\nhi<|im_end|>\n\
1759 <|im_start|>assistant\nhello<|im_end|>\n\
1760 <|im_start|>user\nbye<|im_end|>\n\
1761 <|im_start|>assistant\n";
1762 assert_eq!(got, want);
1763 }
1764
1765 #[test]
1766 fn longest_common_prefix_cutoff() {
1767 // The cross-turn KV reuse boundary (ADR-018 AC-3).
1768 assert_eq!(longest_common_prefix(&[], &[1, 2]), 0);
1769 assert_eq!(longest_common_prefix(&[1, 2, 3], &[1, 2, 3, 4, 5]), 3); // extends
1770 assert_eq!(longest_common_prefix(&[1, 2, 3], &[1, 2, 3]), 3); // equal
1771 assert_eq!(longest_common_prefix(&[1, 9, 3], &[1, 2, 3]), 1); // diverges at idx 1
1772 assert_eq!(longest_common_prefix(&[1, 2, 3], &[1, 2]), 2); // shorter context
1773 assert_eq!(longest_common_prefix(&[5, 6], &[1, 2]), 0); // immediate divergence
1774 }
1775
1776 #[test]
1777 fn post_forward_decode_error_consumes_the_committed_token_once() {
1778 let mut cached = vec![10];
1779 let mut fed = 0;
1780 let mut last_logits = vec![7, 8, 9, 10];
1781
1782 let fallback = apply_committed_forward_result(
1783 Err(ForwardOneError::AfterForward(EdgeError::Engine(
1784 "logit extraction failed",
1785 ))),
1786 11,
1787 &mut cached,
1788 &mut fed,
1789 &mut last_logits,
1790 4,
1791 );
1792
1793 assert_eq!(fallback, Some(vec![0, 0, 0, 0]));
1794 assert_eq!(cached, vec![10, 11]);
1795 assert_eq!(fed, 1);
1796 assert_eq!(last_logits, vec![7, 8, 9, 10]);
1797 }
1798
1799 #[test]
1800 fn local_load_permit_passes_the_provenance_gate() {
1801 // The runtime requires a LoadPermit; the local-trust path must yield one
1802 // for a GGUF artifact (ADR-006 gate exercised, not bypassed).
1803 let permit = local_load_permit(std::path::Path::new("models/qwen.gguf"))
1804 .expect("local permit issued");
1805 assert_eq!(permit.format, el_core::ModelFormat::Gguf);
1806 }
1807
1808 #[test]
1809 fn requested_safety_matches_the_backed_steerer_surface() {
1810 assert_eq!(
1811 requested_session_safety(SafetyMode::Off, true),
1812 SafetyMode::Off,
1813 "Off must stay off even if an expert path is configured"
1814 );
1815 assert_eq!(
1816 requested_session_safety(SafetyMode::Lightweight, true),
1817 SafetyMode::SecDecoding,
1818 "an expert promotes the concrete model-backed path"
1819 );
1820 assert_eq!(
1821 requested_session_safety(SafetyMode::SecDecoding, false),
1822 SafetyMode::Lightweight,
1823 "unbacked SecDecoding must not be reported as active"
1824 );
1825 assert_eq!(
1826 requested_session_safety(SafetyMode::Csd, false),
1827 SafetyMode::Lightweight,
1828 "unbacked Csd must not be reported as active"
1829 );
1830 }
1831
1832 #[test]
1833 fn qwen_provider_from_paths_missing_model_errors() {
1834 let r = QwenChatProvider::from_paths(
1835 std::path::Path::new("/nonexistent/model.gguf"),
1836 std::path::Path::new("/nonexistent/tokenizer.json"),
1837 );
1838 assert!(matches!(r, Err(EdgeError::Engine(_))));
1839 }
1840
1841 // ── safety wiring (ADR-005 tier + ADR-012 control loop) ──────────────────
1842
1843 #[test]
1844 fn safety_off_wires_no_guard_or_steering() {
1845 // Off → the plain single-pass decode: `Ports::permissive()` semantics
1846 // regardless of any resolved bans/patterns.
1847 let cfg = SafetyConfig {
1848 mode: SafetyMode::Off,
1849 banned: vec![1],
1850 patterns: vec![vec![2]],
1851 extra_guard_patterns: vec![],
1852 };
1853 let ports = cfg.ports();
1854 assert!(ports.guard.is_none(), "Off must not wire the chunk guard");
1855 assert!(ports.ingress.is_none(), "Off must not wire ingress triage");
1856 assert_eq!(
1857 ports.safety.mode(),
1858 SafetyMode::Off,
1859 "Off must keep the no-op steerer"
1860 );
1861 }
1862
1863 #[test]
1864 fn lightweight_wires_guard_and_hard_ban_steerer() {
1865 let cfg = SafetyConfig {
1866 mode: SafetyMode::Lightweight,
1867 banned: vec![1],
1868 patterns: vec![vec![2, 3]],
1869 extra_guard_patterns: vec![],
1870 };
1871 let ports = cfg.ports();
1872 assert!(
1873 ports.guard.is_some(),
1874 "Lightweight must wire the chunk guard"
1875 );
1876 assert!(
1877 ports.ingress.is_some(),
1878 "Lightweight must wire prompt ingress triage (ADR-013)"
1879 );
1880 assert_eq!(
1881 ports.safety.mode(),
1882 SafetyMode::Lightweight,
1883 "a non-empty ban list selects the LightweightFilter steerer"
1884 );
1885 }
1886
1887 #[test]
1888 fn lightweight_without_patterns_has_no_guard_or_ingress() {
1889 // No resolvable unsafe patterns (e.g. all multi-token and tokenizer
1890 // produced nothing) → guard/ingress stay off; the per-step ban can still
1891 // apply.
1892 let cfg = SafetyConfig {
1893 mode: SafetyMode::Lightweight,
1894 banned: vec![7],
1895 patterns: vec![],
1896 extra_guard_patterns: vec![],
1897 };
1898 let ports = cfg.ports();
1899 assert!(ports.guard.is_none());
1900 assert!(ports.ingress.is_none());
1901 }
1902
1903 #[test]
1904 fn extra_guard_words_drive_guard_but_not_ingress() {
1905 // Regression (review P2): --guard-word extras must NOT trigger ingress
1906 // refusal, or the documented rollback demo would refuse before decoding.
1907 let cfg = SafetyConfig {
1908 mode: SafetyMode::Lightweight,
1909 banned: vec![],
1910 patterns: vec![], // no built-in unsafe terms
1911 extra_guard_patterns: vec![vec![42]], // a --guard-word trip token
1912 };
1913 let ports = cfg.ports();
1914 assert!(
1915 ports.guard.is_some(),
1916 "extra guard words must drive the output guard"
1917 );
1918 assert!(
1919 ports.ingress.is_none(),
1920 "extra guard words must NOT drive ingress (trajectory demo, not refusal)"
1921 );
1922 }
1923
1924 #[test]
1925 fn qwen_expert_missing_file_errors_and_is_permit_gated() {
1926 // R5: the expert load requires an ADR-006 permit (required arg) and a
1927 // missing file is rejected, not silently ignored.
1928 let r = QwenExpert::from_path_primed(
1929 std::path::Path::new("/nonexistent/expert.gguf"),
1930 0,
1931 &[1, 2],
1932 ok_permit(),
1933 );
1934 assert!(matches!(r, Err(EdgeError::Engine(_))));
1935 }
1936
1937 #[test]
1938 fn provider_and_expert_stay_send_and_sync() {
1939 // ADR-018 expert persistence: making the expert resident must NOT cost the
1940 // provider its thread-safety. The resident expert (`Mutex<…Arc<QwenExpert>>`)
1941 // keeps `QwenChatProvider: Send + Sync`, which is why `QwenExpert` uses a
1942 // `Mutex` rather than `RefCell`/`Cell`. Compile-time guard.
1943 fn assert_send_sync<T: Send + Sync>() {}
1944 assert_send_sync::<QwenExpert>();
1945 assert_send_sync::<QwenChatProvider>();
1946 }
1947}