superfluid_daemon/
lib.rs

1//! superfluid — the daemon.
2
3pub mod anthropic;
4pub mod api;
5pub mod apiweb;
6pub mod atem;
7pub mod bus;
8pub mod capabilities;
9pub mod cli;
10pub mod compat;
11pub mod codec;
12mod exec;
13pub mod completion_bucket;
14pub mod files;
15pub mod harmony;
16pub mod keypolicy;
17pub mod launch;
18pub mod batches;
19pub mod response_store;
20pub mod ratelimit;
21pub mod registry;
22pub mod template_codec;
23pub mod tool_fragment;
24pub mod fleet;
25pub mod fleet_manager;
26pub mod fleet_join;
27pub mod fleet_mdns;
28pub mod fleet_token;
29#[cfg(feature = "tui")]
30pub mod tui;
31#[cfg(feature = "tui")]
32pub mod tui_top;
33pub mod nodeagent;
34mod inference;
35mod block_inference;
36mod generation_output;
37mod openai_output;
38mod http_state;
39mod http_generation;
40mod keepalive;
41mod responses;
42pub mod openai;
43pub mod ollama;
44pub mod fleet_openai;
45pub mod park;
46pub mod pins;
47pub mod runtime;
48pub mod runtime_options;
49pub mod runtime_pick;
50pub mod installs;
51pub mod pulls;
52pub mod runtimes;
53pub mod runtimes_report;
54pub mod versions;
55pub mod selfupdate;
56pub mod linked;
57pub mod dylib_tokenizer;
58pub mod export;
59pub mod otlp;
60pub mod media;
61pub mod metrics;
62pub mod pressure;
63pub mod scheduler;
64pub mod serve;
65pub mod speculate_resolve;
66use superfluid_proto::linkw as proto_linkw;
67pub mod telemetry;
68pub mod store;
69pub mod store_id;
70pub mod wal;
71pub mod workerd;
72
73pub mod qos {
74    pub const INTERACTIVE_CHAT: u8 = 0;
75    pub const INLINE_COMPLETION: u8 = 1;
76    pub const FOREGROUND_AGENT: u8 = 2;
77    pub const BACKGROUND_AGENT: u8 = 3;
78
79    pub const NAMES: [&str; 4] = ["interactive", "completion", "agent", "background"];
80
81    pub fn parse(s: &str) -> Option<u8> {
82        match s.trim().to_ascii_lowercase().as_str() {
83            "interactive" | "interactive-chat" | "interactive_chat" | "0" => Some(INTERACTIVE_CHAT),
84            "completion" | "inline-completion" | "inline_completion" | "1" => Some(INLINE_COMPLETION),
85            "agent" | "foreground" | "foreground-agent" | "foreground_agent" | "2" => Some(FOREGROUND_AGENT),
86            "background" | "background-agent" | "background_agent" | "3" => Some(BACKGROUND_AGENT),
87            _ => None,
88        }
89    }
90
91    pub fn parse_class_lanes(spec: &str, max_lanes: usize) -> Result<[usize; 4], String> {
92        let mut caps = [0usize; 4];
93        let mut seen = [false; 4];
94        for part in spec.split(',').map(str::trim).filter(|p| !p.is_empty()) {
95            let (name, n) = part
96                .split_once('=')
97                .ok_or_else(|| format!("expected class=N, got {part:?}"))?;
98            let class = parse(name).ok_or_else(|| {
99                format!("unknown class {:?} (expected one of {})", name.trim(), NAMES.join(", "))
100            })? as usize;
101            if seen[class] {
102                return Err(format!("class {} is capped twice", NAMES[class]));
103            }
104            seen[class] = true;
105            let n: usize = n
106                .trim()
107                .parse()
108                .map_err(|_| format!("cap for {} must be a whole number, got {:?}", NAMES[class], n.trim()))?;
109            if n == 0 {
110                return Err(format!(
111                    "cap for {} must be at least 1 (omit the class to leave it uncapped)",
112                    NAMES[class]
113                ));
114            }
115            if n > max_lanes {
116                return Err(format!(
117                    "cap {n} for {} exceeds --max-batch {max_lanes}; it could never bind",
118                    NAMES[class]
119                ));
120            }
121            caps[class] = n;
122        }
123        Ok(caps)
124    }
125
126    #[cfg(test)]
127    mod tests {
128        use super::*;
129
130        #[test]
131        fn names_round_trip_and_aliases_parse() {
132            for (c, name) in NAMES.iter().enumerate() {
133                assert_eq!(parse(name), Some(c as u8));
134                assert_eq!(parse(&c.to_string()), Some(c as u8));
135            }
136            assert_eq!(parse(" Interactive-Chat "), Some(INTERACTIVE_CHAT));
137            assert_eq!(parse("foreground_agent"), Some(FOREGROUND_AGENT));
138            assert_eq!(parse("urgent"), None);
139            assert_eq!(parse("4"), None);
140        }
141
142        #[test]
143        fn class_lanes_spec_parses_and_refuses_bad_caps() {
144            assert_eq!(parse_class_lanes("", 8), Ok([0; 4]));
145            assert_eq!(parse_class_lanes("background=2, agent=6", 8), Ok([0, 0, 6, 2]));
146            assert_eq!(parse_class_lanes("interactive=8", 8), Ok([8, 0, 0, 0]));
147            assert!(parse_class_lanes("background", 8).unwrap_err().contains("class=N"));
148            assert!(parse_class_lanes("urgent=1", 8).unwrap_err().contains("unknown class"));
149            assert!(parse_class_lanes("agent=-1", 8).unwrap_err().contains("whole number"));
150            assert!(parse_class_lanes("agent=9", 8).unwrap_err().contains("exceeds"));
151            assert!(parse_class_lanes("background=0", 8).unwrap_err().contains("at least 1"));
152            assert!(parse_class_lanes("agent=1,foreground=2", 8).unwrap_err().contains("twice"));
153        }
154    }
155}
156
157pub use codec::{MockCodec, NoCodec, TextCodec, Utf8Stream};
158pub use superfluid_engine::StageOutput;
159pub use runtime::{EngineHost, WorkerSpec};
160pub use store::{CommittedEvent, SessionStore};
161pub use wal::{EventBody, GenParams};
162
163pub use codec::BundleCodec;
164
165#[derive(Debug, thiserror::Error)]
166pub enum DaemonError {
167    #[error("io: {0}")]
168    Io(#[from] std::io::Error),
169    #[error("wal encode/decode: {0}")]
170    Codec(#[from] postcard::Error),
171    #[error("wal record failed checksum")]
172    WalCorrupt,
173    #[error("unknown session {0}")]
174    UnknownSession(u64),
175    #[error("session {0} has no content to generate from")]
176    EmptySession(u64),
177    #[error(
178        "this model's maximum context length is {max} tokens, but your request has {len} \
179         tokens; shorten the prompt or drop earlier turns"
180    )]
181    StreamTooLong { len: u64, max: u64 },
182    #[error("model load: {0}")]
183    EngineLoad(String),
184    #[error("engine link: {0}")]
185    Agent(#[from] superfluid_agent::AgentError),
186    #[error("fleet link (Link F): {0}")]
187    Fleet(#[from] superfluid_linkf::LinkFError),
188    #[error("fleet node {addr}: {why}")]
189    FleetNode { addr: String, why: String },
190    #[error("engine rejected the plan with status {0}")]
191    Rejected(i32),
192    #[error("generation failed: {0}")]
193    Generation(String),
194    #[error("session {0} is generating; append/generate after it finishes or cancel it")]
195    SessionBusy(u64),
196    #[error("tool call {0} is unknown or already closed")]
197    UnknownToolCall(u64),
198    #[error("unknown message role {0}")]
199    UnknownRole(u32),
200    #[error("tool {index} is not a function declaration: {why}")]
201    InvalidTool { index: usize, why: String },
202    #[error("session {0} already has a prompt; AppendSystem must come before any other turn")]
203    SystemNotFirst(u64),
204    #[error("fork point {fork_at} is outside session {session}'s committed events (1..={max})")]
205    ForkPoint { session: u64, fork_at: u64, max: u64 },
206    #[error("tool call {0} belongs to an ancestor's timeline; only that session can close it")]
207    InheritedToolCall(u64),
208    #[error("session {session} metadata version is {actual}, command expected {expected}")]
209    MetaConflict { session: u64, expected: u64, actual: u64 },
210    #[error("session {session} generation is {actual}, command expected {expected}")]
211    GenerationConflict { session: u64, expected: u64, actual: u64 },
212    #[error("session {0} was purged")]
213    Purged(u64),
214    #[error("files quota exceeded: {stored} bytes stored, {incoming} more would pass the {max} limit")]
215    FilesQuota { stored: u64, incoming: u64, max: u64 },
216    #[error("rebase edit spec invalid: {0}")]
217    RebaseSpec(&'static str),
218    #[error("session {0} did not quiesce for purge (lane still running)")]
219    PurgeTimeout(u64),
220    #[error("permission request {0} is unknown or already answered")]
221    UnknownPermission(u64),
222    #[error("session {0} has an open permission request; answer it before generating")]
223    PermissionPending(u64),
224    #[error("tool call {0}'s executor lease expired; the result was recorded as expired, not applied")]
225    ToolLeaseExpired(u64),
226    #[error("tool call {0} already reached a terminal outcome; the late result was recorded as a reconciliation")]
227    ToolLateEffect(u64),
228    #[error("unknown block kind {0}")]
229    UnknownBlockKind(u32),
230    #[error("media blob {0} is not in the pool (put it first)")]
231    UnknownMedia(String),
232    #[error("image expands to {0} placeholders, beyond one prefill chunk")]
233    MediaTooLarge(u32),
234    #[error("this codec has no FIM dialect")]
235    NoFimDialect,
236    #[error("completion rate limit exceeded; retry in {retry_after_ms} ms")]
237    CompletionThrottled { retry_after_ms: u64 },
238    #[error("protocol: {0}")]
239    Protocol(&'static str),
240    #[error("constraint unsatisfiable: {0}")]
241    Constraint(String),
242    #[error("the model's chat template refused this conversation: {0}")]
243    TemplateRefused(String),
244    #[error("{0}")]
245    Unsupported(Refusal),
246    #[error("configuration: {0}")]
247    Config(&'static str),
248}
249
250#[derive(Debug, Clone, PartialEq, Eq)]
251pub struct Refusal {
252    pub code: &'static str,
253    pub param: Option<String>,
254    pub message: String,
255}
256
257impl Refusal {
258    pub fn new(code: &'static str, param: Option<&str>, message: impl Into<String>) -> Refusal {
259        Refusal { code, param: param.map(str::to_string), message: message.into() }
260    }
261}
262
263impl std::fmt::Display for Refusal {
264    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
265        f.write_str(&self.message)
266    }
267}
268
269#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
270pub struct SessionSummary {
271    pub id: u64,
272    pub parent: Option<u64>,
273    pub fork_at: u64,
274    pub title: Option<String>,
275    pub archived: bool,
276    pub meta_version: u64,
277    pub generation: u64,
278    pub events: u64,
279    pub tokens: u64,
280    pub open_tool_calls: u64,
281}
282
283#[derive(Debug)]
284pub struct GenerateOutcome {
285    pub events: Vec<CommittedEvent>,
286    pub tokens_generated: u32,
287    pub finish: u32,
288    pub warm_prefix: u64,
289    pub logprobs: Vec<scheduler::TokenLogprob>,
290    pub spec: (u64, u64),
291    /// Tokens decoded after the generation's first tick and the seconds they
292    /// took, tick end to tick end, as the scheduler delivered them.
293    pub decode: (u64, f64),
294}
295
296#[derive(Debug, Clone, Default, PartialEq)]
297pub struct ModelSamplingDefaults {
298    pub temperature: Option<f32>,
299    pub top_p: Option<f32>,
300    pub top_k: Option<u32>,
301    pub min_p: Option<f32>,
302    pub do_sample: Option<bool>,
303    pub repetition_penalty: Option<f32>,
304}
305
306impl ModelSamplingDefaults {
307    pub fn from_json(v: &serde_json::Value) -> Self {
308        let f = |k: &str| v.get(k).and_then(|x| x.as_f64()).map(|x| x as f32);
309        Self {
310            temperature: f("temperature"),
311            top_p: f("top_p"),
312            top_k: v.get("top_k").and_then(|x| x.as_u64()).map(|x| x as u32),
313            min_p: f("min_p"),
314            do_sample: v.get("do_sample").and_then(|x| x.as_bool()),
315            repetition_penalty: f("repetition_penalty"),
316        }
317    }
318}
319
320thread_local! {
321    static SYNC_CLASS: std::cell::Cell<Option<u8>> = const { std::cell::Cell::new(None) };
322}
323
324pub fn with_sync_class<R>(class: Option<u8>, f: impl FnOnce() -> R) -> R {
325    let _scope = SyncClassScope::enter(class);
326    f()
327}
328
329pub struct SyncClassScope {
330    prev: Option<Option<u8>>,
331    _not_send: std::marker::PhantomData<*const ()>,
332}
333
334impl SyncClassScope {
335    pub fn enter(class: Option<u8>) -> SyncClassScope {
336        let prev = class.map(|c| SYNC_CLASS.with(|cell| cell.replace(Some(c))));
337        SyncClassScope { prev, _not_send: std::marker::PhantomData }
338    }
339}
340
341impl Drop for SyncClassScope {
342    fn drop(&mut self) {
343        if let Some(prev) = self.prev {
344            SYNC_CLASS.with(|c| c.set(prev));
345        }
346    }
347}
348
349fn ambient_sync_class() -> Option<u8> {
350    SYNC_CLASS.with(|c| c.get())
351}
352
353struct Serving {
354    capabilities: runtime::SharedCapabilities,
355    retired: std::sync::Arc<std::sync::Mutex<Option<String>>>,
356    model_sampling: std::sync::OnceLock<ModelSamplingDefaults>,
357    codec: std::sync::Arc<dyn TextCodec + Send + Sync>,
358    sched: scheduler::SchedulerHandle,
359    speculation: std::sync::Arc<scheduler::Speculation>,
360    staged_tokens: u64,
361}
362
363/// Sessions whose generation should stop. A cancel also asks the running
364/// tick to end at its next step boundary, so the lane stops within a step;
365/// a queued chat asks the same.
366#[derive(Default)]
367pub struct CancelSet {
368    set: std::sync::Mutex<std::collections::HashSet<u64>>,
369    wake: std::sync::RwLock<Option<runtime::YieldSlot>>,
370    /// When a latency-sensitive request last arrived (`monotonic_ms`; 0:
371    /// never).
372    latency_at_ms: std::sync::atomic::AtomicU64,
373}
374
375/// Milliseconds on a monotonic clock that starts at 1 with the process: a
376/// wall clock stepped back would hold latency mode on until it caught up.
377fn monotonic_ms() -> u64 {
378    static BASE: std::sync::OnceLock<std::time::Instant> = std::sync::OnceLock::new();
379    BASE.get_or_init(std::time::Instant::now).elapsed().as_millis() as u64 + 1
380}
381
382impl std::ops::Deref for CancelSet {
383    type Target = std::sync::Mutex<std::collections::HashSet<u64>>;
384
385    fn deref(&self) -> &Self::Target {
386        &self.set
387    }
388}
389
390impl CancelSet {
391    pub fn cancel(&self, session: u64) {
392        self.set.lock().expect("cancel registry").insert(session);
393        self.wake();
394    }
395
396    /// Asks the running tick to end at its next step boundary.
397    pub fn wake(&self) {
398        if let Some(slot) = self.wake.read().expect("wake slot").as_ref() {
399            runtime::raise_yield(slot);
400        }
401    }
402
403    pub(crate) fn wake_with(&self, slot: runtime::YieldSlot) {
404        *self.wake.write().expect("wake slot") = Some(slot);
405    }
406
407    /// A chat or an inline completion arrived: the worker runs short
408    /// prefill steps for a while, and the running tick ends at its next one.
409    pub fn latency_request(&self) {
410        self.latency_at_ms.store(monotonic_ms(), std::sync::atomic::Ordering::Release);
411        if let Some(slot) = self.wake.read().expect("wake slot").as_ref() {
412            runtime::set_latency(slot, true);
413        }
414        self.wake();
415    }
416
417    /// Whether a latency-sensitive request arrived within `hold`.
418    pub(crate) fn latency_recent(&self, hold: std::time::Duration) -> bool {
419        let at = self.latency_at_ms.load(std::sync::atomic::Ordering::Acquire);
420        at > 0 && monotonic_ms().saturating_sub(at) < hold.as_millis() as u64
421    }
422}
423
424struct Ground<'a> {
425    store: &'a std::sync::Arc<std::sync::Mutex<SessionStore>>,
426    cancels: &'a std::sync::Arc<CancelSet>,
427    active: &'a std::sync::Arc<std::sync::Mutex<std::collections::HashSet<u64>>>,
428    stats: &'a std::sync::Arc<scheduler::SchedStats>,
429    background_tick_divisor: &'a std::sync::Arc<std::sync::atomic::AtomicU32>,
430}
431
432impl Serving {
433    fn start(
434        host: EngineHost,
435        codec: std::sync::Arc<dyn TextCodec + Send + Sync>,
436        ground: Ground<'_>,
437        opts: DaemonOptions,
438    ) -> Result<Serving, DaemonError> {
439        let mut opts = opts;
440        let mut speculation = std::sync::Arc::new(scheduler::Speculation::default());
441        let mut host = host;
442        ground.cancels.wake_with(host.yield_slot());
443        ground
444            .stats
445            .kv_bits
446            .store(host.kv_bits() as u64, std::sync::atomic::Ordering::Relaxed);
447        if let Some(first) = opts.speculate.clone() {
448            let mut candidates = vec![first];
449            if opts.speculate_auto {
450                candidates.extend(opts.speculate_fallbacks.iter().cloned());
451            }
452            let picked = first_registrable(
453                &candidates,
454                opts.speculate_auto,
455                |strategy| {
456                    speculation = std::sync::Arc::new(scheduler::Speculation::default());
457                    scheduler::Speculation::register(&speculation, &mut host, strategy, &opts.speculation)
458                },
459                |strategy, e, next| {
460                    eprintln!(
461                        "superfluid: speculation: auto picked {strategy}, but the engine refused it ({e}) — {}",
462                        match next {
463                            Some(n) => format!("trying {n}"),
464                            None => "decoding plainly".to_string(),
465                        }
466                    );
467                },
468            )?;
469            if picked.is_none() {
470                speculation = std::sync::Arc::new(scheduler::Speculation::default());
471            } else if picked != opts.speculate {
472                eprintln!("superfluid: speculation: {}", picked.as_deref().unwrap_or_default());
473            }
474            opts.speculate = picked;
475        }
476        let staged_tokens = host.staged_tokens();
477        let capabilities = host.capabilities();
478        let retired = std::sync::Arc::new(std::sync::Mutex::new(None));
479        let sched = scheduler::Scheduler::spawn(
480            host,
481            std::sync::Arc::clone(ground.store),
482            std::sync::Arc::clone(&codec),
483            std::sync::Arc::clone(ground.cancels),
484            std::sync::Arc::clone(ground.active),
485            std::sync::Arc::clone(ground.stats),
486            std::sync::Arc::clone(ground.background_tick_divisor),
487            std::sync::Arc::clone(&speculation),
488            std::sync::Arc::clone(&retired),
489            opts,
490        );
491        {
492            let st = ground.store.lock().expect("store");
493            let now = wal::now_unix_ms();
494            for id in st.session_ids() {
495                let Ok(s) = st.session(id) else { continue };
496                if s.pin_deadline_unix_ms > now {
497                    let (tokens, want) = pins::session_target(&s.tokens);
498                    let _ = sched.tx.send(scheduler::SchedCmd::Pin(pins::PinReq {
499                        key: pins::PinKey::Session(id),
500                        tokens,
501                        want,
502                        deadline_unix_ms: s.pin_deadline_unix_ms,
503                        anchor: Some(id),
504                        reply: None,
505                    }));
506                }
507            }
508        }
509        Ok(Serving { capabilities, retired, model_sampling: std::sync::OnceLock::new(), codec, sched, speculation, staged_tokens })
510    }
511}
512
513pub struct Daemon {
514    serving: std::sync::RwLock<std::sync::Arc<Serving>>,
515    reloading: std::sync::Mutex<()>,
516    reload_hook: std::sync::OnceLock<Box<dyn Fn() + Send + Sync>>,
517    runtime_id: Option<&'static str>,
518    store: std::sync::Arc<std::sync::Mutex<SessionStore>>,
519    cancels: std::sync::Arc<CancelSet>,
520    active: std::sync::Arc<std::sync::Mutex<std::collections::HashSet<u64>>>,
521    purging: std::sync::Arc<std::sync::Mutex<std::collections::HashSet<u64>>>,
522    stats: std::sync::Arc<scheduler::SchedStats>,
523    trace_scope: otlp::TraceScope,
524    background_tick_divisor: std::sync::Arc<std::sync::atomic::AtomicU32>,
525    park_dir: Option<std::path::PathBuf>,
526    tool_lease_ms: u64,
527    media: media::MediaPool,
528    files: files::FileStore,
529    batches: std::sync::Arc<batches::BatchStore>,
530    responses: response_store::ResponseStore,
531    completion: std::sync::Arc<CompletionSatellite>,
532    fim_satellite: std::sync::OnceLock<FimSatellite>,
533    built_with: DaemonOptions,
534}
535
536impl CompletionSatellite {
537    fn cache_get(&self, key: u64) -> Option<(Vec<u32>, String, u32)> {
538        let mut c = self.cache.lock().expect("completion cache");
539        if let Some(v) = c.map.get(&key).cloned() {
540            c.lru.retain(|&k| k != key);
541            c.lru.push_back(key);
542            Some(v)
543        } else {
544            None
545        }
546    }
547    fn cache_put(&self, key: u64, tokens: Vec<u32>, text: String, finish: u32) {
548        let mut c = self.cache.lock().expect("completion cache");
549        if c.map.insert(key, (tokens, text, finish)).is_none() {
550            c.lru.push_back(key);
551        }
552        while c.lru.len() > self.cache_cap {
553            if let Some(old) = c.lru.pop_front() {
554                c.map.remove(&old);
555            }
556        }
557    }
558}
559
560struct FimSatellite {
561    name: String,
562    daemon: std::sync::Arc<Daemon>,
563}
564
565#[derive(Default)]
566struct CompletionCounters {
567    generated: std::sync::atomic::AtomicU64,
568    cached: std::sync::atomic::AtomicU64,
569    expired: std::sync::atomic::AtomicU64,
570    throttled: std::sync::atomic::AtomicU64,
571}
572
573struct CompletionSatellite {
574    store: std::sync::Mutex<SessionStore>,
575    cache: std::sync::Mutex<CompletionCache>,
576    deadline_ms: u64,
577    cache_cap: usize,
578    bucket: completion_bucket::BucketConfig,
579    counters: CompletionCounters,
580}
581
582#[derive(Default)]
583struct CompletionCache {
584    map: std::collections::HashMap<u64, (Vec<u32>, String, u32)>,
585    lru: std::collections::VecDeque<u64>,
586}
587
588#[derive(Debug, Clone, PartialEq)]
589pub struct Completion {
590    pub text: String,
591    pub tokens: Vec<u32>,
592    pub cached: bool,
593    pub expired: bool,
594    pub finish: u32,
595    pub prompt_tokens: u32,
596}
597
598#[derive(Debug, Clone)]
599pub struct FimRequest {
600    pub prefix: String,
601    pub suffix: String,
602    pub mode: u8,
603    pub max_tokens: u32,
604    pub params: GenParams,
605    pub extras: scheduler::GenExtras,
606}
607
608impl FimRequest {
609    pub fn greedy(prefix: &str, suffix: &str, mode: u8, max_tokens: u32) -> FimRequest {
610        FimRequest {
611            prefix: prefix.to_string(),
612            suffix: suffix.to_string(),
613            mode,
614            max_tokens,
615            params: GenParams::default(),
616            extras: scheduler::GenExtras::default(),
617        }
618    }
619}
620
621#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
622pub enum Tier {
623    Resident,
624    Parked { covered: u64, encoding: u8 },
625    Cold,
626}
627
628#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
629pub struct Inspection {
630    pub summary: SessionSummary,
631    pub params: GenParams,
632    pub epoch: u64,
633    pub last_finish: u32,
634    pub qos_class: u8,
635    pub batch_invariant: bool,
636    pub pin_deadline_unix_ms: u64,
637    pub behavior_fingerprint: Option<u64>,
638    pub open_tool_calls: Vec<(u64, String, String)>,
639    pub tier: Tier,
640    pub tokens: Vec<u32>,
641    pub events: Vec<api::EventMsg>,
642    pub wal_version: u8,
643    #[serde(default, deserialize_with = "de_trailing_opt")]
644    pub store_id: Option<store::StoreId>,
645}
646
647fn de_trailing_opt<'de, D, T>(d: D) -> Result<Option<T>, D::Error>
648where
649    D: serde::Deserializer<'de>,
650    T: serde::Deserialize<'de>,
651{
652    use std::cell::Cell;
653    struct V<'a, T>(&'a Cell<bool>, std::marker::PhantomData<T>);
654    impl<'de, T: serde::Deserialize<'de>> serde::de::Visitor<'de> for V<'_, T> {
655        type Value = Option<T>;
656        fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
657            f.write_str("an optional trailing field")
658        }
659        fn visit_none<E: serde::de::Error>(self) -> Result<Option<T>, E> {
660            self.0.set(true);
661            Ok(None)
662        }
663        fn visit_some<D2: serde::Deserializer<'de>>(self, d: D2) -> Result<Option<T>, D2::Error> {
664            self.0.set(true);
665            T::deserialize(d).map(Some)
666        }
667    }
668    let reached = Cell::new(false);
669    match d.deserialize_option(V(&reached, std::marker::PhantomData)) {
670        Ok(v) => Ok(v),
671        Err(e) if !reached.get() && e.to_string() == POSTCARD_UNEXPECTED_END => Ok(None),
672        Err(e) => Err(e),
673    }
674}
675
676const POSTCARD_UNEXPECTED_END: &str = "Hit the end of buffer, expected more data";
677
678#[derive(Debug, Clone, Default, PartialEq)]
679pub struct SpeculationOptions {
680    pub max_draft: Option<u32>,
681    pub adaptive: Option<bool>,
682    pub min_yield: Option<f64>,
683    pub yield_rounds: Option<u32>,
684    pub throughput_gate: Option<bool>,
685    pub gate_probe_tokens: Option<u32>,
686    pub min_speedup: Option<f64>,
687    pub gate_reprobe: Option<u32>,
688    pub gate_reprobe_max: Option<u32>,
689    pub dspark_confidence: Option<f32>,
690}
691
692#[derive(Debug, Clone)]
693pub struct DaemonOptions {
694    pub max_lanes: usize,
695    pub park_dir: Option<std::path::PathBuf>,
696    pub pressure_high_pct: u8,
697    pub pressure_low_pct: u8,
698    pub pin_budget_pct: u8,
699    pub park_lossy: bool,
700    pub park_budget_bytes: u64,
701    pub class_lanes: [usize; 4],
702    pub tick_decode_budget: u32,
703    pub tick_target_ms: u32,
704    pub runtime_id: Option<&'static str>,
705    pub prefill_budget: u32,
706    pub agent_starvation_ticks: u32,
707    pub pressure_source: pressure::PressureConfig,
708    pub completion_deadline_ms: u64,
709    pub completion_bucket: completion_bucket::BucketConfig,
710    pub media_dir: Option<std::path::PathBuf>,
711    pub files_max_bytes: Option<u64>,
712    pub files_expiry_secs: Option<u64>,
713    pub tool_lease_ms: u64,
714    pub speculate: Option<String>,
715    pub speculate_auto: bool,
716    pub speculate_fallbacks: Vec<String>,
717    pub speculation: SpeculationOptions,
718    pub trace_scope: otlp::TraceScope,
719}
720
721impl Default for DaemonOptions {
722    fn default() -> DaemonOptions {
723        DaemonOptions {
724            max_lanes: 8,
725            park_dir: None,
726            pressure_high_pct: 85,
727            pressure_low_pct: 70,
728            pin_budget_pct: 50,
729            park_lossy: false,
730            park_budget_bytes: 20 << 30,
731            class_lanes: [0; 4],
732            tick_decode_budget: 0,
733            tick_target_ms: 0,
734            runtime_id: None,
735            prefill_budget: scheduler::PREFILL_BUDGET,
736            agent_starvation_ticks: 8,
737            pressure_source: pressure::PressureConfig::None,
738            tool_lease_ms: 0,
739            speculate: None,
740            speculate_auto: false,
741            speculate_fallbacks: Vec::new(),
742            speculation: SpeculationOptions::default(),
743            files_max_bytes: None,
744            files_expiry_secs: None,
745            media_dir: None,
746            completion_deadline_ms: 0,
747            trace_scope: otlp::TraceScope::default(),
748            completion_bucket: completion_bucket::BucketConfig::default(),
749        }
750    }
751}
752
753struct PurgeFence {
754    set: std::sync::Arc<std::sync::Mutex<std::collections::HashSet<u64>>>,
755    ids: Vec<u64>,
756}
757
758impl Drop for PurgeFence {
759    fn drop(&mut self) {
760        let mut p = self.set.lock().expect("purging");
761        for id in &self.ids {
762            p.remove(id);
763        }
764    }
765}
766
767pub(crate) fn render_refusal(
768    codec: &dyn TextCodec,
769    messages: &[codec::ChatMessage],
770    tools: &[String],
771    kwargs: &serde_json::Map<String, serde_json::Value>,
772    fallback: &'static str,
773) -> DaemonError {
774    match codec.template_refusal(messages, tools, kwargs) {
775        Some(why) => DaemonError::TemplateRefused(why),
776        None => DaemonError::Protocol(fallback),
777    }
778}
779
780impl Daemon {
781    pub fn new(
782        store: SessionStore,
783        host: EngineHost,
784        codec: Box<dyn TextCodec + Send + Sync>,
785        max_lanes: usize,
786    ) -> Daemon {
787        Daemon::with_park(store, host, codec, max_lanes, None)
788    }
789
790    pub fn with_park(
791        store: SessionStore,
792        host: EngineHost,
793        codec: Box<dyn TextCodec + Send + Sync>,
794        max_lanes: usize,
795        park_dir: Option<std::path::PathBuf>,
796    ) -> Daemon {
797        Daemon::with_options(
798            store,
799            host,
800            codec,
801            DaemonOptions {
802                max_lanes,
803                park_dir,
804                ..DaemonOptions::default()
805            },
806        )
807        .expect("default watermarks are valid")
808    }
809
810    pub fn with_options(
811        store: SessionStore,
812        host: EngineHost,
813        codec: Box<dyn TextCodec + Send + Sync>,
814        opts: DaemonOptions,
815    ) -> Result<Daemon, DaemonError> {
816        if opts.pressure_high_pct == 0 || opts.pressure_high_pct > 100 {
817            return Err(DaemonError::Config(
818                "pressure_high_pct must be in 1..=100",
819            ));
820        }
821        if opts.pressure_low_pct >= opts.pressure_high_pct {
822            return Err(DaemonError::Config(
823                "pressure_low_pct must be below pressure_high_pct",
824            ));
825        }
826        if opts.max_lanes == 0 {
827            return Err(DaemonError::Config("max_lanes must be at least 1"));
828        }
829        let mut opts = opts;
830        let runtime_id = opts.runtime_id;
831        opts.pin_budget_pct = opts.pin_budget_pct.min(opts.pressure_low_pct);
832        opts.completion_bucket.validate().map_err(DaemonError::Config)?;
833        if matches!(opts.speculation.max_draft, Some(0 | 16..)) {
834            return Err(DaemonError::Config("speculation max_draft must be in 1..=15"));
835        }
836        if opts.speculation.min_yield.is_some_and(|v| !v.is_finite() || v < 0.0) {
837            return Err(DaemonError::Config("speculation min_yield must be finite and non-negative"));
838        }
839        if opts.speculation.yield_rounds == Some(0) {
840            return Err(DaemonError::Config("speculation yield_rounds must be at least 1"));
841        }
842        if opts.speculation.min_speedup.is_some_and(|v| !v.is_finite() || v < 1.0) {
843            return Err(DaemonError::Config(
844                "speculation min_speedup must be finite and at least 1.0",
845            ));
846        }
847        if opts.speculation.dspark_confidence.is_some_and(|v| !v.is_finite() || v <= 0.0) {
848            return Err(DaemonError::Config(
849                "speculation dspark_confidence must be finite and positive",
850            ));
851        }
852        opts.trace_scope.ns = export::store_trace_ns(&store.store_id());
853        let store = std::sync::Arc::new(std::sync::Mutex::new(store));
854        let built_with = opts.clone();
855        let cancels = std::sync::Arc::new(CancelSet::default());
856        let active = std::sync::Arc::new(std::sync::Mutex::new(std::collections::HashSet::new()));
857        let purging = std::sync::Arc::new(std::sync::Mutex::new(std::collections::HashSet::new()));
858        let stats = std::sync::Arc::new(scheduler::SchedStats::default());
859        let background_tick_divisor = std::sync::Arc::new(std::sync::atomic::AtomicU32::new(1));
860        let park_dir = opts.park_dir.clone();
861        let tool_lease_ms = opts.tool_lease_ms;
862        let files_max_bytes = opts.files_max_bytes;
863        let files_expiry_secs = opts.files_expiry_secs;
864        let completion_deadline_ms = opts.completion_deadline_ms;
865        let completion_bucket = opts.completion_bucket;
866        let media_dir = opts.media_dir.clone().unwrap_or_else(|| {
867            park_dir
868                .as_ref()
869                .and_then(|p| p.parent().map(|d| d.join("media")))
870                .unwrap_or_else(|| {
871                    let uniq = std::time::SystemTime::now()
872                        .duration_since(std::time::UNIX_EPOCH)
873                        .map(|d| d.as_nanos())
874                        .unwrap_or(0);
875                    std::env::temp_dir().join(format!("superfluid-{}-{}", std::process::id(), uniq)).join("media")
876                })
877        });
878        let trace_scope = opts.trace_scope.clone();
879        let serving = Serving::start(
880            host,
881            std::sync::Arc::from(codec),
882            Ground {
883                store: &store,
884                cancels: &cancels,
885                active: &active,
886                stats: &stats,
887                background_tick_divisor: &background_tick_divisor,
888            },
889            opts,
890        )?;
891        Ok(Daemon {
892            trace_scope,
893            serving: std::sync::RwLock::new(std::sync::Arc::new(serving)),
894            reloading: std::sync::Mutex::new(()),
895            reload_hook: std::sync::OnceLock::new(),
896            runtime_id,
897            store,
898            cancels,
899            active,
900            purging,
901            stats,
902            background_tick_divisor,
903            park_dir,
904            tool_lease_ms,
905            files: files::FileStore::with_limits(
906                files::files_dir_for(media_dir.parent().unwrap_or(media_dir.as_path())),
907                files_max_bytes,
908                files_expiry_secs,
909            ),
910            batches: std::sync::Arc::new(batches::BatchStore::new(batches::batches_dir_for(
911                media_dir.parent().unwrap_or(media_dir.as_path()),
912            ))),
913            responses: response_store::ResponseStore::new(response_store::responses_dir_for(
914                media_dir.parent().unwrap_or(media_dir.as_path()),
915            )),
916            media: media::MediaPool::new(media_dir),
917            completion: std::sync::Arc::new(CompletionSatellite {
918                store: std::sync::Mutex::new(SessionStore::ephemeral()),
919                cache: std::sync::Mutex::new(CompletionCache::default()),
920                deadline_ms: completion_deadline_ms,
921                cache_cap: 256,
922                bucket: completion_bucket,
923                counters: CompletionCounters::default(),
924            }),
925            fim_satellite: std::sync::OnceLock::new(),
926            built_with,
927        })
928    }
929
930    fn serving(&self) -> std::sync::Arc<Serving> {
931        std::sync::Arc::clone(&self.serving.read().expect("serving"))
932    }
933
934    fn codec(&self) -> std::sync::Arc<dyn TextCodec + Send + Sync> {
935        std::sync::Arc::clone(&self.serving().codec)
936    }
937
938    fn sched_tx(&self) -> std::sync::mpsc::Sender<scheduler::SchedCmd> {
939        self.serving().sched.tx.clone()
940    }
941
942    pub fn reload_with(
943        &self,
944        build: impl FnOnce() -> Result<(EngineHost, Box<dyn TextCodec + Send + Sync>), DaemonError>,
945    ) -> Result<(), DaemonError> {
946        let _one_at_a_time = self.reloading.lock().expect("reloading");
947        self.shutdown();
948        {
949            let mut cancels = self.cancels.lock().expect("cancel registry");
950            let active = self.active.lock().expect("active set");
951            cancels.retain(|session| active.contains(session));
952        }
953        let started = build().and_then(|(host, codec)| {
954            Serving::start(
955                host,
956                std::sync::Arc::from(codec),
957                Ground {
958                    store: &self.store,
959                    cancels: &self.cancels,
960                    active: &self.active,
961                    stats: &self.stats,
962                    background_tick_divisor: &self.background_tick_divisor,
963                },
964                self.built_with.clone(),
965            )
966        });
967        match started {
968            Ok(serving) => {
969                *self.serving.write().expect("serving") = std::sync::Arc::new(serving);
970                Ok(())
971            }
972            Err(e) => {
973                let stopped = self.serving();
974                let mut why = stopped.retired.lock().expect("retired");
975                if why.is_none() {
976                    *why = Some(format!("the model's engine was stopped to load it again, which failed ({e}); it loads again on its next request, so retry"));
977                }
978                Err(e)
979            }
980        }
981    }
982
983    pub fn on_retired(&self, reload: Box<dyn Fn() + Send + Sync>) {
984        let _ = self.reload_hook.set(reload);
985    }
986
987    fn ensure_serving(&self) -> Result<(), DaemonError> {
988        if self.needs_reload().is_none() {
989            return Ok(());
990        }
991        if let Some(reload) = self.reload_hook.get() {
992            reload();
993        }
994        match self.needs_reload() {
995            None => Ok(()),
996            Some(why) => Err(DaemonError::Generation(why)),
997        }
998    }
999
1000    pub fn max_stream_tokens(&self) -> u64 {
1001        let serving = self.serving();
1002        let record = serving.capabilities.read().expect("capability record");
1003        runtime::stream_ceiling(serving.staged_tokens, &record)
1004    }
1005
1006    fn check_stream_fits(&self, len: usize) -> Result<(), DaemonError> {
1007        let len = len as u64;
1008        let max = self.max_stream_tokens();
1009        if max > 0 && len > max {
1010            return Err(DaemonError::StreamTooLong { len, max });
1011        }
1012        Ok(())
1013    }
1014
1015    pub fn session_fits(&self, session: u64, thinking: Option<bool>) -> Result<(), DaemonError> {
1016        self.ensure_serving()?;
1017        let len = {
1018            let store = self.store.lock().expect("store");
1019            let s = store.session(session)?;
1020            s.tokens.len() + self.generation_opener(&s.events, thinking).map_or(0, |op| op.len())
1021        };
1022        self.check_stream_fits(len)
1023    }
1024
1025    pub fn fim_fits(&self, req: &FimRequest) -> Result<(), DaemonError> {
1026        let d = match self.fim_satellite.get() {
1027            Some(sat) => &sat.daemon,
1028            None => self,
1029        };
1030        d.ensure_serving()?;
1031        let fim = d
1032            .codec()
1033            .render_fim(&req.prefix, &req.suffix, req.mode)
1034            .ok_or(DaemonError::NoFimDialect)?;
1035        d.check_stream_fits(fim.len())
1036    }
1037
1038    pub fn complete(
1039        &self,
1040        prefix: &str,
1041        suffix: &str,
1042        mode: u8,
1043        max_tokens: u32,
1044    ) -> Result<Completion, DaemonError> {
1045        self.complete_streaming(&FimRequest::greedy(prefix, suffix, mode, max_tokens), |_, _| {
1046            std::ops::ControlFlow::Continue(())
1047        })
1048    }
1049
1050    pub fn complete_streaming(
1051        &self,
1052        req: &FimRequest,
1053        on_text: impl FnMut(&str, usize) -> std::ops::ControlFlow<()>,
1054    ) -> Result<Completion, DaemonError> {
1055        let result = match self.fim_satellite.get() {
1056            Some(sat) => sat.daemon.complete_here(req, on_text),
1057            None => self.complete_here(req, on_text),
1058        };
1059        if let Ok(c) = &result {
1060            use std::sync::atomic::Ordering::Relaxed;
1061            let n = &self.completion.counters;
1062            if c.expired {
1063                n.expired.fetch_add(1, Relaxed);
1064            } else if c.cached {
1065                n.cached.fetch_add(1, Relaxed);
1066            } else {
1067                n.generated.fetch_add(1, Relaxed);
1068            }
1069        }
1070        result
1071    }
1072
1073    pub fn attach_fim_satellite(
1074        &self,
1075        name: impl Into<String>,
1076        satellite: std::sync::Arc<Daemon>,
1077    ) -> Result<(), DaemonError> {
1078        if satellite.codec().render_fim("", "", codec::fim_mode::PSM).is_none() {
1079            return Err(DaemonError::NoFimDialect);
1080        }
1081        if satellite.fim_satellite.get().is_some() {
1082            return Err(DaemonError::Config("a FIM satellite cannot itself route to another"));
1083        }
1084        self.fim_satellite
1085            .set(FimSatellite {
1086                name: name.into(),
1087                daemon: satellite,
1088            })
1089            .map_err(|_| DaemonError::Config("a FIM satellite is already attached"))
1090    }
1091
1092    pub fn fim_satellite_name(&self) -> Option<String> {
1093        self.fim_satellite.get().map(|s| s.name.clone())
1094    }
1095
1096    pub fn fim_satellite(&self) -> Option<std::sync::Arc<Daemon>> {
1097        self.fim_satellite.get().map(|s| std::sync::Arc::clone(&s.daemon))
1098    }
1099
1100    pub fn fim_supported(&self) -> bool {
1101        self.fim_satellite.get().is_some()
1102            || self.codec().render_fim("", "", codec::fim_mode::PSM).is_some()
1103    }
1104
1105    pub fn completion_bucket(&self) -> completion_bucket::BucketConfig {
1106        self.completion.bucket
1107    }
1108
1109    pub fn note_completion_throttled(&self) {
1110        self.completion
1111            .counters
1112            .throttled
1113            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1114    }
1115
1116    fn complete_here(
1117        &self,
1118        req: &FimRequest,
1119        mut on_text: impl FnMut(&str, usize) -> std::ops::ControlFlow<()>,
1120    ) -> Result<Completion, DaemonError> {
1121        self.ensure_serving()?;
1122        let _span = self.raw_generate_span(req.max_tokens);
1123        let submitted = std::time::Instant::now();
1124        let fim = self
1125            .codec()
1126            .render_fim(&req.prefix, &req.suffix, req.mode)
1127            .ok_or(DaemonError::NoFimDialect)?;
1128        self.check_stream_fits(fim.len())?;
1129        let prompt_tokens = fim.len() as u32;
1130        let expired = || Completion {
1131            text: String::new(),
1132            tokens: Vec::new(),
1133            cached: false,
1134            expired: true,
1135            finish: superfluid_abi::finish::NONE,
1136            prompt_tokens,
1137        };
1138        let cacheable = req.params.temperature <= 0.0;
1139        let key = {
1140            let mut h = superfluid_fingerprint::RecordHasher::new();
1141            for t in &fim {
1142                h.field_u32(1, *t);
1143            }
1144            h.field_u32(2, req.mode as u32);
1145            h.field_u32(3, req.max_tokens);
1146            if let Some(fp) = self.codec().behavior_fingerprint() {
1147                h.field_u64(4, fp.digest());
1148            }
1149            let x = &req.extras;
1150            h.field_u32(5, x.presence_penalty.to_bits());
1151            h.field_u32(6, x.frequency_penalty.to_bits());
1152            h.field_u32(7, x.repeat_penalty.to_bits());
1153            h.field_u32(8, x.ignore_eos as u32);
1154            h.finish()
1155        };
1156        if cacheable {
1157            if let Some((tokens, text, finish)) = self.completion.cache_get(key) {
1158                let _ = on_text(&text, tokens.len());
1159                return Ok(Completion {
1160                    text,
1161                    tokens,
1162                    cached: true,
1163                    expired: false,
1164                    finish,
1165                    prompt_tokens,
1166                });
1167            }
1168        }
1169        if self.completion.deadline_ms > 0
1170            && submitted.elapsed().as_millis() as u64 > self.completion.deadline_ms
1171        {
1172            return Ok(expired());
1173        }
1174        let session = {
1175            let mut store = self.completion.store.lock().expect("completion store");
1176            let s = store.create(None, req.params)?;
1177            store.set_qos(s, qos::INLINE_COMPLETION, false)?;
1178            store.append(s, None, fim.clone())?;
1179            s
1180        };
1181        let (stream, params) = {
1182            let store = self.completion.store.lock().expect("completion store");
1183            let st = store.session(session)?;
1184            (st.tokens.clone(), st.params)
1185        };
1186        let mut produced: Vec<u32> = Vec::new();
1187        let mut emitted = 0usize;
1188        let result = scheduler::run_completion_streaming(
1189            &self.sched_tx(),
1190            stream,
1191            params,
1192            req.max_tokens,
1193            self.completion.deadline_ms,
1194            scheduler::GenExtras {
1195                grammar_handle: 0,
1196                logit_bias_handle: 0,
1197                want_logprobs: false,
1198                ..req.extras.clone()
1199            },
1200            |toks| {
1201                produced.extend_from_slice(toks);
1202                let text = self.codec().decode(&produced);
1203                let stable = text.trim_end_matches('\u{FFFD}');
1204                if stable.len() > emitted && stable.is_char_boundary(emitted) {
1205                    let piece = &stable[emitted..];
1206                    emitted = stable.len();
1207                    return on_text(piece, produced.len());
1208                }
1209                std::ops::ControlFlow::Continue(())
1210            },
1211        );
1212        {
1213            let mut store = self.completion.store.lock().expect("completion store");
1214            let _ = store.drop_session(session);
1215        }
1216        let (tokens, finish, was_expired) = result?;
1217        if was_expired {
1218            return Ok(expired());
1219        }
1220        let text = self.codec().decode(&tokens);
1221        if text.len() > emitted && text.is_char_boundary(emitted) {
1222            let _ = on_text(&text[emitted..], tokens.len());
1223        }
1224        if cacheable && finish != superfluid_abi::finish::CANCELLED {
1225            self.completion.cache_put(key, tokens.clone(), text.clone(), finish);
1226        }
1227        Ok(Completion {
1228            text,
1229            tokens,
1230            cached: false,
1231            expired: false,
1232            finish,
1233            prompt_tokens,
1234        })
1235    }
1236
1237    pub fn generate_tokens(
1238        &self,
1239        stream: Vec<u32>,
1240        params: GenParams,
1241        max_tokens: u32,
1242        deadline_ms: u64,
1243    ) -> Result<(Vec<u32>, bool), DaemonError> {
1244        self.ensure_serving()?;
1245        let _span = self.raw_generate_span(max_tokens);
1246        self.check_stream_fits(stream.len())?;
1247        let deadline = self.effective_deadline(deadline_ms);
1248        let (tokens, _finish, expired) =
1249            scheduler::run_completion(&self.sched_tx(), stream, params, max_tokens, deadline, 0)?;
1250        Ok((tokens, expired))
1251    }
1252
1253    pub fn generate_tokens_streaming(
1254        &self,
1255        stream: Vec<u32>,
1256        params: GenParams,
1257        max_tokens: u32,
1258        deadline_ms: u64,
1259        extras: scheduler::GenExtras,
1260        on_tokens: impl FnMut(&[u32]) -> std::ops::ControlFlow<()>,
1261    ) -> Result<(Vec<u32>, u32, bool), DaemonError> {
1262        self.ensure_serving()?;
1263        let _span = self.raw_generate_span(max_tokens);
1264        self.check_stream_fits(stream.len())?;
1265        let deadline = self.effective_deadline(deadline_ms);
1266        scheduler::run_completion_streaming(
1267            &self.sched_tx(),
1268            stream,
1269            params,
1270            max_tokens,
1271            deadline,
1272            extras,
1273            on_tokens,
1274        )
1275    }
1276
1277    pub fn forward_stage(
1278        &self,
1279        tokens: Vec<u32>,
1280        start_layer: u32,
1281        end_layer: u32,
1282        hidden_in: Option<Vec<u8>>,
1283    ) -> Result<superfluid_engine::StageOutput, DaemonError> {
1284        self.ensure_serving()?;
1285        match self.engine_sync(proto_linkw::StateSyncReqMsg::ForwardStage {
1286            tokens,
1287            start_layer,
1288            end_layer,
1289            hidden_in,
1290        })? {
1291            proto_linkw::StateSyncOkMsg::ForwardStage {
1292                hidden: Some(h),
1293                token: None,
1294            } => Ok(superfluid_engine::StageOutput::Hidden(h)),
1295            proto_linkw::StateSyncOkMsg::ForwardStage {
1296                hidden: None,
1297                token: Some(t),
1298            } => Ok(superfluid_engine::StageOutput::Token(t)),
1299            _ => Err(DaemonError::Protocol("mismatched ForwardStage reply")),
1300        }
1301    }
1302
1303    fn effective_deadline(&self, deadline_ms: u64) -> u64 {
1304        if deadline_ms == 0 {
1305            self.completion.deadline_ms
1306        } else if self.completion.deadline_ms == 0 {
1307            deadline_ms
1308        } else {
1309            deadline_ms.min(self.completion.deadline_ms)
1310        }
1311    }
1312
1313    pub fn generate_tokens_constrained(
1314        &self,
1315        stream: Vec<u32>,
1316        params: GenParams,
1317        max_tokens: u32,
1318        deadline_ms: u64,
1319        json_schema: &str,
1320    ) -> Result<(Vec<u32>, bool), DaemonError> {
1321        self.ensure_serving()?;
1322        let _span = self.raw_generate_span(max_tokens);
1323        let handle = match self.engine_sync(proto_linkw::StateSyncReqMsg::CreateGrammar {
1324            json_schema: json_schema.to_string(),
1325        })? {
1326            proto_linkw::StateSyncOkMsg::CreateGrammar { grammar_handle } => grammar_handle,
1327            _ => return Err(DaemonError::Protocol("mismatched CreateGrammar reply")),
1328        };
1329        let deadline = self.effective_deadline(deadline_ms);
1330        let result =
1331            scheduler::run_completion(&self.sched_tx(), stream, params, max_tokens, deadline, handle);
1332        let _ = self.engine_sync(proto_linkw::StateSyncReqMsg::FreeGrammar {
1333            grammar_handle: handle,
1334        });
1335        let (tokens, finish, expired) = result?;
1336        if expired {
1337            return Err(DaemonError::Generation(format!(
1338                "constrained generation hit the wall deadline with the grammar incomplete                  ({} tokens)",
1339                tokens.len()
1340            )));
1341        }
1342        if finish != superfluid_abi::finish::GRAMMAR {
1343            return Err(DaemonError::Generation(format!(
1344                "constrained generation did not complete the grammar (finish={finish},                  {} tokens); raise max_tokens",
1345                tokens.len()
1346            )));
1347        }
1348        Ok((tokens, false))
1349    }
1350
1351    pub fn put_media(&self, bytes: &[u8]) -> Result<String, DaemonError> {
1352        Ok(self.media.put(bytes)?)
1353    }
1354
1355    pub fn media_pool(&self) -> &media::MediaPool {
1356        &self.media
1357    }
1358
1359    fn engine_sync(&self, req: proto_linkw::StateSyncReqMsg) -> Result<proto_linkw::StateSyncOkMsg, DaemonError> {
1360        self.engine_sync_as(req, None)
1361    }
1362
1363    fn engine_sync_as(
1364        &self,
1365        req: proto_linkw::StateSyncReqMsg,
1366        class: Option<u8>,
1367    ) -> Result<proto_linkw::StateSyncOkMsg, DaemonError> {
1368        let class = class.or_else(ambient_sync_class);
1369        let (tx, rx) = std::sync::mpsc::channel();
1370        let cmd = scheduler::SchedCmd::Sync(req, tx);
1371        self.sched_tx()
1372            .send(match class {
1373                Some(c) => scheduler::SchedCmd::SyncAs(c, Box::new(cmd)),
1374                None => cmd,
1375            })
1376            .map_err(|_| DaemonError::Protocol("scheduler is gone"))?;
1377        rx.recv()
1378            .map_err(|_| DaemonError::Protocol("scheduler dropped the query"))?
1379            .map_err(DaemonError::Generation)
1380    }
1381
1382    pub fn append_image(
1383        &self,
1384        session: u64,
1385        role: u32,
1386        blob: &str,
1387        pre_text: &str,
1388        post_text: &str,
1389    ) -> Result<CommittedEvent, DaemonError> {
1390        if self.active.lock().expect("active set").contains(&session) {
1391            return Err(DaemonError::SessionBusy(session));
1392        }
1393        if !matches!(role, wal::role::SYSTEM | wal::role::USER | wal::role::ASSISTANT) {
1394            return Err(DaemonError::UnknownRole(role));
1395        }
1396        if !self.media.exists(blob) {
1397            return Err(DaemonError::UnknownMedia(blob.to_string()));
1398        }
1399        let path = self.media.path_of(blob);
1400        let info = match self.engine_sync(proto_linkw::StateSyncReqMsg::MediaProbe {
1401            image_path: path.to_string_lossy().into_owned(),
1402        })? {
1403            proto_linkw::StateSyncOkMsg::MediaProbe { info } => info,
1404            _ => return Err(DaemonError::Protocol("mismatched MediaProbe reply")),
1405        };
1406        if info.n_tokens == 0 || info.n_tokens as u64 > scheduler::PREFILL_BUDGET as u64 {
1407            return Err(DaemonError::MediaTooLarge(info.n_tokens));
1408        }
1409        let mut run = Vec::with_capacity(info.n_tokens as usize + 2);
1410        if info.boi_token_id != 0 {
1411            run.push(info.boi_token_id);
1412        }
1413        run.extend(std::iter::repeat_n(info.image_token_id, info.n_tokens as usize));
1414        if info.eoi_token_id != 0 {
1415            run.push(info.eoi_token_id);
1416        }
1417        let mut span = self.separator_for(session);
1418        let turn = self
1419            .codec()
1420            .render_turn_with_tokens(role, pre_text, &run, post_text)
1421            .ok_or(DaemonError::Protocol("codec has no chat dialect"))?;
1422        span.extend(turn);
1423        let first_in_span = span
1424            .iter()
1425            .position(|&t| t == info.image_token_id)
1426            .ok_or(DaemonError::Protocol("codec dropped the placeholder run"))?;
1427        let mut store = self.store.lock().expect("store");
1428        let stream_len = store.session(session)?.tokens.len() as u64;
1429        let payload = serde_json::json!({
1430            "blob": blob,
1431            "n_tokens": info.n_tokens,
1432            "image_token_id": info.image_token_id,
1433            "offset": stream_len + first_in_span as u64,
1434            "pre_text": pre_text,
1435            "post_text": post_text,
1436            "media_identity": {
1437                "preprocessing_fp": format!("{:016x}", info.preprocess_fp),
1438                "encoder": "bundle-weights-identity",
1439            },
1440        })
1441        .to_string();
1442        store.commit_block(
1443            session,
1444            role,
1445            wal::block_kind::IMAGE,
1446            payload,
1447            span,
1448            codec::MODEL_VISIBILITY_VERSION,
1449        )
1450    }
1451
1452    fn media_of(&self, s: &store::SessionState) -> Vec<(u64, std::path::PathBuf, u32)> {
1453        let mut v = Vec::new();
1454        for e in &s.events {
1455            if let EventBody::Block {
1456                kind, payload, ..
1457            } = &e.body
1458            {
1459                if *kind != wal::block_kind::IMAGE {
1460                    continue;
1461                }
1462                let p: serde_json::Value = serde_json::from_str(payload).unwrap_or_default();
1463                let blob = p.get("blob").and_then(|b| b.as_str()).unwrap_or("");
1464                let off = p.get("offset").and_then(|o| o.as_u64()).unwrap_or(0);
1465                let n = p.get("n_tokens").and_then(|o| o.as_u64()).unwrap_or(0) as u32;
1466                if !blob.is_empty() && n > 0 {
1467                    v.push((off, self.media.path_of(blob), n));
1468                }
1469            }
1470        }
1471        v
1472    }
1473
1474    pub fn gc_media(&self) -> Result<(usize, usize), DaemonError> {
1475        let store = self.store.lock().expect("store");
1476        let mut live = std::collections::HashSet::new();
1477        for id in store.session_ids() {
1478            let Ok(s) = store.session(id) else { continue };
1479            for e in &s.events {
1480                if let EventBody::Block { kind, payload, .. } = &e.body {
1481                    if *kind == wal::block_kind::IMAGE {
1482                        let p: serde_json::Value = serde_json::from_str(payload).unwrap_or_default();
1483                        if let Some(b) = p.get("blob").and_then(|b| b.as_str()) {
1484                            live.insert(b.to_string());
1485                        }
1486                    }
1487                }
1488            }
1489        }
1490        Ok(self.media.sweep(&live)?)
1491    }
1492
1493    pub fn speculation(&self) -> Option<scheduler::SpeculationGrant> {
1494        self.serving().speculation.grant()
1495    }
1496
1497    pub fn tool_lease_ms(&self) -> u64 {
1498        self.tool_lease_ms
1499    }
1500
1501    pub fn inspect(&self, session: u64) -> Result<Inspection, DaemonError> {
1502        // The active set before the store, in the order a generation takes them.
1503        let resident = self.active.lock().expect("active set").contains(&session);
1504        let store = self.store.lock().expect("store");
1505        let s = store.session(session)?;
1506        let tier = if resident {
1507            Tier::Resident
1508        } else {
1509            match self
1510                .park_dir
1511                .as_deref()
1512                .and_then(|d| park::read(d, session))
1513            {
1514                Some(p) => Tier::Parked {
1515                    covered: p.covered,
1516                    encoding: p.encoding,
1517                },
1518                None => Tier::Cold,
1519            }
1520        };
1521        Ok(Inspection {
1522            summary: SessionSummary {
1523                id: s.id,
1524                parent: if s.rerooted { None } else { s.parent },
1525                fork_at: s.base,
1526                title: s.title.clone(),
1527                archived: s.archived,
1528                meta_version: s.meta_version,
1529                generation: s.generation,
1530                events: s.events.len() as u64,
1531                tokens: s.tokens.len() as u64,
1532                open_tool_calls: s.open_tool_calls.len() as u64,
1533            },
1534            params: s.params,
1535            epoch: s.epoch,
1536            last_finish: s.last_finish,
1537            qos_class: s.qos_class,
1538            batch_invariant: s.batch_invariant,
1539            pin_deadline_unix_ms: s.pin_deadline_unix_ms,
1540            behavior_fingerprint: s.behavior_fp,
1541            open_tool_calls: s
1542                .open_tool_calls
1543                .iter()
1544                .map(|(id, (n, a))| (*id, n.clone(), a.clone()))
1545                .collect(),
1546            tier,
1547            tokens: s.tokens.clone(),
1548            events: s.events.iter().cloned().map(Into::into).collect(),
1549            wal_version: store.wal_version(),
1550            store_id: Some(store.store_id()),
1551        })
1552    }
1553
1554    pub fn files(&self) -> &files::FileStore {
1555        &self.files
1556    }
1557
1558    pub fn batches(&self) -> std::sync::Arc<batches::BatchStore> {
1559        std::sync::Arc::clone(&self.batches)
1560    }
1561
1562    pub fn responses(&self) -> &response_store::ResponseStore {
1563        &self.responses
1564    }
1565
1566    pub fn tokenize(&self, text: &str) -> Vec<u32> {
1567        self.codec().encode(text)
1568    }
1569
1570    pub fn grammar_create(&self, json_schema: &str) -> Result<u32, DaemonError> {
1571        match self.engine_sync(proto_linkw::StateSyncReqMsg::CreateGrammar {
1572            json_schema: json_schema.to_string(),
1573        })? {
1574            proto_linkw::StateSyncOkMsg::CreateGrammar { grammar_handle } => Ok(grammar_handle),
1575            _ => Err(DaemonError::Protocol("mismatched CreateGrammar reply")),
1576        }
1577    }
1578
1579    pub fn grammar_free(&self, grammar_handle: u32) {
1580        let _ = self.engine_sync(proto_linkw::StateSyncReqMsg::FreeGrammar { grammar_handle });
1581    }
1582
1583    pub fn logit_bias_create(&self, tokens: &[i32], values: &[f32]) -> Result<u32, DaemonError> {
1584        match self.engine_sync(proto_linkw::StateSyncReqMsg::CreateLogitBias {
1585            tokens: tokens.to_vec(),
1586            values: values.iter().map(|v| v.to_bits()).collect(),
1587        })? {
1588            proto_linkw::StateSyncOkMsg::CreateLogitBias { handle } => Ok(handle),
1589            _ => Err(DaemonError::Protocol("mismatched CreateLogitBias reply")),
1590        }
1591    }
1592    pub fn logit_bias_free(&self, handle: u32) {
1593        let _ = self.engine_sync(proto_linkw::StateSyncReqMsg::FreeLogitBias { handle });
1594    }
1595
1596    pub fn lora_load(&self, adapter_path: &str) -> Result<Option<String>, DaemonError> {
1597        self.lora_reply(proto_linkw::StateSyncReqMsg::LoraLoad {
1598            adapter_path: adapter_path.to_string(),
1599        })
1600    }
1601
1602    pub fn lora_unload(&self) -> Result<Option<String>, DaemonError> {
1603        self.lora_reply(proto_linkw::StateSyncReqMsg::LoraUnload)
1604    }
1605
1606    pub fn lora_id(&self) -> Result<Option<String>, DaemonError> {
1607        self.lora_reply(proto_linkw::StateSyncReqMsg::LoraId)
1608    }
1609
1610    fn lora_reply(&self, req: proto_linkw::StateSyncReqMsg) -> Result<Option<String>, DaemonError> {
1611        match self.engine_sync(req)? {
1612            proto_linkw::StateSyncOkMsg::LoraId { id } => Ok(id),
1613            _ => Err(DaemonError::Protocol("mismatched LoraId reply")),
1614        }
1615    }
1616
1617    pub fn try_model_sampling_defaults(&self) -> Option<ModelSamplingDefaults> {
1618        let serving = self.serving();
1619        if let Some(cached) = serving.model_sampling.get() {
1620            return Some(cached.clone());
1621        }
1622        let descriptor = self.capability_descriptor()?;
1623        let resolved = descriptor
1624            .get("sampling_defaults")
1625            .map(ModelSamplingDefaults::from_json)
1626            .unwrap_or_default();
1627        let _ = serving.model_sampling.set(resolved.clone());
1628        Some(resolved)
1629    }
1630
1631    pub fn model_sampling_defaults(&self) -> ModelSamplingDefaults {
1632        self.try_model_sampling_defaults().unwrap_or_default()
1633    }
1634
1635    pub fn capability_descriptor(&self) -> Option<serde_json::Value> {
1636        self.serving().capabilities.read().expect("capability record").with_runtime(self.runtime_id)
1637    }
1638
1639    pub fn runtime_id(&self) -> Option<&'static str> {
1640        self.runtime_id
1641    }
1642
1643    pub fn needs_reload(&self) -> Option<String> {
1644        self.serving().retired.lock().expect("retired").clone()
1645    }
1646
1647    pub fn capabilities(&self) -> capabilities::Capabilities {
1648        self.serving().capabilities.read().expect("capability record").clone()
1649    }
1650
1651    pub fn require(&self, section: &str, key: &str, code: &'static str, param: Option<&str>, what: &str) -> Result<(), DaemonError> {
1652        let why = match self.serving().capabilities.read().expect("capability record").get(section, key) {
1653            capabilities::Cap::No(why) => why.to_string(),
1654            capabilities::Cap::Yes | capabilities::Cap::Unknown => return Ok(()),
1655        };
1656        Err(DaemonError::Unsupported(Refusal::new(code, param, format!("{} {what}: {why}", self.described()))))
1657    }
1658
1659    fn described(&self) -> String {
1660        let model = &self.trace_scope.model;
1661        let runtime = self.capability_descriptor().and_then(|r| r["runtime"]["id"].as_str().map(str::to_string));
1662        match (model.is_empty(), runtime) {
1663            (false, Some(r)) => format!("model '{model}' on runtime {r}"),
1664            (false, None) => format!("model '{model}'"),
1665            (true, Some(r)) => format!("this model (runtime {r})"),
1666            (true, None) => "this model".to_string(),
1667        }
1668    }
1669
1670    pub fn chat_template(&self) -> Option<String> {
1671        self.codec().chat_template_source()
1672    }
1673
1674    pub fn transcribe(
1675        &self,
1676        audio_path: &str,
1677        params: &superfluid_engine::TranscribeParams,
1678    ) -> Result<superfluid_engine::Transcription, DaemonError> {
1679        self.transcribe_as(audio_path, params, None)
1680    }
1681
1682    pub fn transcribe_as(
1683        &self,
1684        audio_path: &str,
1685        params: &superfluid_engine::TranscribeParams,
1686        class: Option<u8>,
1687    ) -> Result<superfluid_engine::Transcription, DaemonError> {
1688        self.ensure_serving()?;
1689        self.require_transcription(params)?;
1690        Self::decode_transcription(self.engine_sync_as(Self::transcribe_req(audio_path, params, false), class)?)
1691    }
1692
1693    fn require_transcription(&self, params: &superfluid_engine::TranscribeParams) -> Result<(), DaemonError> {
1694        if params.translate {
1695            self.require("modalities", "whisper_translate", "unsupported_model", Some("model"), "does not translate speech")
1696        } else {
1697            self.require("modalities", "whisper_transcribe", "unsupported_model", Some("model"), "does not transcribe speech")
1698        }
1699    }
1700
1701    #[allow(clippy::type_complexity)]
1702    pub fn transcribe_streaming(
1703        &self,
1704        audio_path: &str,
1705        params: &superfluid_engine::TranscribeParams,
1706    ) -> Result<
1707        (
1708            std::sync::mpsc::Receiver<proto_linkw::TranscribeSegmentMsg>,
1709            std::sync::mpsc::Receiver<Result<superfluid_engine::Transcription, DaemonError>>,
1710        ),
1711        DaemonError,
1712    > {
1713        self.transcribe_streaming_as(audio_path, params, None)
1714    }
1715
1716    #[allow(clippy::type_complexity)]
1717    pub fn transcribe_streaming_as(
1718        &self,
1719        audio_path: &str,
1720        params: &superfluid_engine::TranscribeParams,
1721        class: Option<u8>,
1722    ) -> Result<
1723        (
1724            std::sync::mpsc::Receiver<proto_linkw::TranscribeSegmentMsg>,
1725            std::sync::mpsc::Receiver<Result<superfluid_engine::Transcription, DaemonError>>,
1726        ),
1727        DaemonError,
1728    > {
1729        self.ensure_serving()?;
1730        self.require_transcription(params)?;
1731        let (seg_tx, seg_rx) = std::sync::mpsc::channel();
1732        let (res_tx, res_rx) = std::sync::mpsc::channel();
1733        let (raw_tx, raw_rx) = std::sync::mpsc::channel();
1734        self.sched_tx()
1735            .send({
1736                let cmd = scheduler::SchedCmd::SyncStream(Self::transcribe_req(audio_path, params, true), seg_tx, raw_tx);
1737                match class.or_else(ambient_sync_class) {
1738                    Some(c) => scheduler::SchedCmd::SyncAs(c, Box::new(cmd)),
1739                    None => cmd,
1740                }
1741            })
1742            .map_err(|_| DaemonError::Protocol("scheduler is gone"))?;
1743        std::thread::spawn(move || {
1744            let out = match raw_rx.recv() {
1745                Ok(Ok(ok)) => Self::decode_transcription(ok),
1746                Ok(Err(e)) => Err(DaemonError::Generation(e)),
1747                Err(_) => Err(DaemonError::Protocol("scheduler dropped the query")),
1748            };
1749            let _ = res_tx.send(out);
1750        });
1751        Ok((seg_rx, res_rx))
1752    }
1753
1754    fn transcribe_req(
1755        audio_path: &str,
1756        params: &superfluid_engine::TranscribeParams,
1757        stream: bool,
1758    ) -> proto_linkw::StateSyncReqMsg {
1759        proto_linkw::StateSyncReqMsg::Transcribe {
1760            audio_path: audio_path.to_string(),
1761            language: params.language.clone(),
1762            translate: params.translate,
1763            timestamps: params.timestamps,
1764            prompt: params.prompt.clone(),
1765            stream,
1766        }
1767    }
1768
1769    fn decode_transcription(
1770        ok: proto_linkw::StateSyncOkMsg,
1771    ) -> Result<superfluid_engine::Transcription, DaemonError> {
1772        match ok {
1773            proto_linkw::StateSyncOkMsg::Transcribe { text, language, duration_ms, segments } => {
1774                Ok(superfluid_engine::Transcription {
1775                    text,
1776                    language,
1777                    duration_ms,
1778                    segments: segments
1779                        .into_iter()
1780                        .map(|s| superfluid_engine::TranscriptSegment {
1781                            start_ms: s.start_ms,
1782                            end_ms: s.end_ms,
1783                            text: s.text,
1784                            avg_logprob: f32::from_bits(s.avg_logprob_bits),
1785                            no_speech_prob: f32::from_bits(s.no_speech_prob_bits),
1786                            compression_ratio: f32::from_bits(s.compression_ratio_bits),
1787                            temperature: f32::from_bits(s.temperature_bits),
1788                        })
1789                        .collect(),
1790                })
1791            }
1792            _ => Err(DaemonError::Protocol("engine returned a non-Transcribe reply")),
1793        }
1794    }
1795
1796    pub fn embed(&self, text: &str) -> Result<Vec<f32>, DaemonError> {
1797        self.embed_as(text, None)
1798    }
1799
1800    pub fn embed_as(&self, text: &str, class: Option<u8>) -> Result<Vec<f32>, DaemonError> {
1801        self.ensure_serving()?;
1802        self.require("workload", "embedding", "unsupported_model", Some("model"), "serves no embeddings")?;
1803        self.embed_tokens_as(self.codec().encode(text), class)
1804    }
1805
1806    pub(crate) fn embed_tokens_as(&self, tokens: Vec<u32>, class: Option<u8>) -> Result<Vec<f32>, DaemonError> {
1807        self.ensure_serving()?;
1808        self.require("workload", "embedding", "unsupported_model", Some("model"), "serves no embeddings")?;
1809        match self.engine_sync_as(proto_linkw::StateSyncReqMsg::Embed { tokens }, class)? {
1810            proto_linkw::StateSyncOkMsg::Embed { embedding } => Ok(embedding
1811                .as_chunks::<4>()
1812                .0
1813                .iter()
1814                .map(|b| f32::from_le_bytes(*b))
1815                .collect()),
1816            _ => Err(DaemonError::Protocol("engine returned a non-Embed reply")),
1817        }
1818    }
1819
1820    pub fn detokenize(&self, tokens: &[u32]) -> String {
1821        self.codec().decode(tokens)
1822    }
1823
1824    fn raw_generate_span(&self, max_tokens: u32) -> tracing::span::EnteredSpan {
1825        tracing::info_span!(
1826            parent: None,
1827            "generate",
1828            max_tokens,
1829            raw = true,
1830            model = self.trace_scope.model.as_str(),
1831            store_ns = self.trace_scope.ns
1832        )
1833        .entered()
1834    }
1835
1836    pub fn metrics_text(&self) -> String {
1837        let mut out = metrics::render_prometheus(&self.stats);
1838        out.push_str(&format!(
1839            "# HELP superfluid_telemetry_dropped_total Telemetry records dropped by the bounded sinks (log lines and OTLP spans).\n# TYPE superfluid_telemetry_dropped_total counter\nsuperfluid_telemetry_dropped_total {}\n",
1840            telemetry::dropped_total()
1841        ));
1842        let g = |a: &std::sync::atomic::AtomicU64| a.load(std::sync::atomic::Ordering::Relaxed);
1843        out.push_str(&format!(
1844            "# HELP superfluid_otlp_spans_exported_total Operational spans accepted by the OTLP collector.\n# TYPE superfluid_otlp_spans_exported_total counter\nsuperfluid_otlp_spans_exported_total {}\n",
1845            g(&otlp::EXPORTED)
1846        ));
1847        out.push_str(&format!(
1848            "# HELP superfluid_otlp_spans_dropped_total Operational spans dropped instead of delaying a tick.\n# TYPE superfluid_otlp_spans_dropped_total counter\nsuperfluid_otlp_spans_dropped_total{{reason=\"queue_full\"}} {}\nsuperfluid_otlp_spans_dropped_total{{reason=\"export_failed\"}} {}\nsuperfluid_otlp_spans_dropped_total{{reason=\"rejected\"}} {}\n",
1849            g(&otlp::DROPPED_QUEUE_FULL),
1850            g(&otlp::DROPPED_EXPORT_FAILED),
1851            g(&otlp::DROPPED_REJECTED)
1852        ));
1853        out.push_str(&self.completion_metrics_text());
1854        out
1855    }
1856
1857    fn completion_metrics_text(&self) -> String {
1858        use std::sync::atomic::Ordering::Relaxed;
1859        let n = &self.completion.counters;
1860        let mut out = String::from(
1861            "# HELP superfluid_fim_completions_total FIM completions by outcome.\n\
1862             # TYPE superfluid_fim_completions_total counter\n",
1863        );
1864        for (label, v) in [
1865            ("generated", &n.generated),
1866            ("cached", &n.cached),
1867            ("expired", &n.expired),
1868            ("throttled", &n.throttled),
1869        ] {
1870            out.push_str(&format!(
1871                "superfluid_fim_completions_total{{result=\"{label}\"}} {}\n",
1872                v.load(Relaxed)
1873            ));
1874        }
1875        if let Some(sat) = self.fim_satellite.get() {
1876            let s = &sat.daemon.stats;
1877            let model = sat.name.replace('\\', "\\\\").replace('"', "\\\"").replace('\n', "\\n");
1878            out.push_str(&format!(
1879                "# HELP superfluid_fim_satellite_info The --fim-model satellite serving completions.\n\
1880                 # TYPE superfluid_fim_satellite_info gauge\n\
1881                 superfluid_fim_satellite_info{{model=\"{model}\"}} 1\n\
1882                 # HELP superfluid_fim_satellite_decode_tokens_total Tokens the satellite decoded.\n\
1883                 # TYPE superfluid_fim_satellite_decode_tokens_total counter\n\
1884                 superfluid_fim_satellite_decode_tokens_total {}\n\
1885                 # HELP superfluid_fim_satellite_lanes_active Satellite lanes resident after its last tick.\n\
1886                 # TYPE superfluid_fim_satellite_lanes_active gauge\n\
1887                 superfluid_fim_satellite_lanes_active {}\n\
1888                 # HELP superfluid_fim_satellite_pool_blocks_used Satellite KV paged blocks in use.\n\
1889                 # TYPE superfluid_fim_satellite_pool_blocks_used gauge\n\
1890                 superfluid_fim_satellite_pool_blocks_used {}\n\
1891                 # HELP superfluid_fim_satellite_pool_blocks_total Satellite KV paged blocks total.\n\
1892                 # TYPE superfluid_fim_satellite_pool_blocks_total gauge\n\
1893                 superfluid_fim_satellite_pool_blocks_total {}\n",
1894                s.decode_tokens.load(Relaxed),
1895                s.lanes_active.load(Relaxed),
1896                s.pool_blocks_used.load(Relaxed),
1897                s.pool_blocks_total.load(Relaxed),
1898            ));
1899        }
1900        out
1901    }
1902
1903    pub fn set_background_tick_divisor(&self, divisor: u32) {
1904        self.background_tick_divisor
1905            .store(divisor.max(1), std::sync::atomic::Ordering::Relaxed);
1906    }
1907
1908    pub fn set_qos(
1909        &self,
1910        session: u64,
1911        class: u8,
1912        batch_invariant: bool,
1913    ) -> Result<CommittedEvent, DaemonError> {
1914        self.store
1915            .lock()
1916            .expect("store")
1917            .set_qos(session, class, batch_invariant)
1918    }
1919
1920    pub fn sched_stats(&self) -> std::sync::Arc<scheduler::SchedStats> {
1921        std::sync::Arc::clone(&self.stats)
1922    }
1923
1924    pub fn cancel_registry(&self) -> std::sync::Arc<CancelSet> {
1925        std::sync::Arc::clone(&self.cancels)
1926    }
1927
1928    pub fn active_registry(
1929        &self,
1930    ) -> std::sync::Arc<std::sync::Mutex<std::collections::HashSet<u64>>> {
1931        std::sync::Arc::clone(&self.active)
1932    }
1933
1934    pub fn store(&self) -> std::sync::Arc<std::sync::Mutex<SessionStore>> {
1935        std::sync::Arc::clone(&self.store)
1936    }
1937
1938    pub fn bus(&self) -> std::sync::Arc<bus::SessionBus> {
1939        self.store.lock().expect("store").bus()
1940    }
1941
1942    pub fn shutdown(&self) {
1943        self.serving().sched.shutdown();
1944    }
1945
1946    pub fn signal_shutdown(&self) {
1947        self.serving().sched.signal_shutdown();
1948    }
1949
1950    pub fn create(&self, parent: Option<u64>, params: GenParams) -> Result<u64, DaemonError> {
1951        self.store.lock().expect("store").create(parent, params)
1952    }
1953
1954    pub fn fork(
1955        &self,
1956        parent: u64,
1957        at_event: u64,
1958        params: Option<GenParams>,
1959    ) -> Result<u64, DaemonError> {
1960        self.refuse_if_purging(parent)?;
1961        self.store.lock().expect("store").fork(parent, at_event, params)
1962    }
1963
1964    pub fn tree(&self) -> Vec<(u64, Option<u64>, u64)> {
1965        self.sessions(false)
1966            .into_iter()
1967            .map(|s| (s.id, s.parent, s.fork_at))
1968            .collect()
1969    }
1970
1971    pub fn sessions(&self, include_archived: bool) -> Vec<SessionSummary> {
1972        let store = self.store.lock().expect("store");
1973        store
1974            .session_ids()
1975            .into_iter()
1976            .filter_map(|id| store.session(id).ok())
1977            .filter(|s| include_archived || !s.archived)
1978            .map(|s| SessionSummary {
1979                id: s.id,
1980                parent: if s.rerooted { None } else { s.parent },
1981                fork_at: s.base,
1982                title: s.title.clone(),
1983                archived: s.archived,
1984                meta_version: s.meta_version,
1985                generation: s.generation,
1986                events: s.events.len() as u64,
1987                tokens: s.tokens.len() as u64,
1988                open_tool_calls: s.open_tool_calls.len() as u64,
1989            })
1990            .collect()
1991    }
1992
1993    pub fn set_meta(
1994        &self,
1995        session: u64,
1996        expected_version: u64,
1997        title: Option<String>,
1998        archived: Option<bool>,
1999    ) -> Result<u64, DaemonError> {
2000        self.store
2001            .lock()
2002            .expect("store")
2003            .set_meta(session, expected_version, title, archived)
2004    }
2005
2006    pub fn rebase(
2007        &self,
2008        parent: u64,
2009        edits: Vec<wal::RebaseEdit>,
2010        params: Option<GenParams>,
2011    ) -> Result<u64, DaemonError> {
2012        use wal::RebaseEdit;
2013        if edits.is_empty() {
2014            return Err(DaemonError::RebaseSpec("no edits"));
2015        }
2016        self.refuse_if_purging(parent)?;
2017        let range = |e: &RebaseEdit| match e {
2018            RebaseEdit::Drop { from, to } | RebaseEdit::Replace { from, to, .. } => (*from, *to),
2019        };
2020        let mut store = self.store.lock().expect("store");
2021        let (max, parent_events) = {
2022            let p = store.session(parent)?;
2023            (p.next_event_id(), p.events.clone())
2024        };
2025        let mut last_to = 0;
2026        for e in &edits {
2027            let (from, to) = range(e);
2028            if from == 0 {
2029                return Err(DaemonError::RebaseSpec("edits never include Created (id 0)"));
2030            }
2031            if from >= to {
2032                return Err(DaemonError::RebaseSpec("empty range"));
2033            }
2034            if to > max {
2035                return Err(DaemonError::RebaseSpec("range beyond the parent's committed events"));
2036            }
2037            if from < last_to {
2038                return Err(DaemonError::RebaseSpec("ranges must be sorted and non-overlapping"));
2039            }
2040            if let RebaseEdit::Replace { role, .. } = e {
2041                if !matches!(*role, wal::role::SYSTEM | wal::role::USER | wal::role::ASSISTANT) {
2042                    return Err(DaemonError::UnknownRole(*role));
2043                }
2044            }
2045            last_to = to;
2046        }
2047        let first = range(&edits[0]).0;
2048        let child = store.rebase_branch(parent, first, params, edits.clone())?;
2049        let mut tool_ids: std::collections::HashMap<u64, u64> = std::collections::HashMap::new();
2050        let mut i = first as usize;
2051        while i < parent_events.len() {
2052            let id = i as u64;
2053            if let Some(e) = edits.iter().find(|e| {
2054                let (from, to) = range(e);
2055                id >= from && id < to
2056            }) {
2057                let (from, to) = range(e);
2058                if let RebaseEdit::Replace { role, text, .. } = e {
2059                    if id == from {
2060                        let span = self.render_for(&store, child, *role, text)?;
2061                        store.append_message(child, *role, text.clone(), span)?;
2062                    }
2063                }
2064                i = to as usize;
2065                continue;
2066            }
2067            match &parent_events[i].body {
2068                EventBody::Appended { text, span } => {
2069                    store.append(child, text.clone(), span.clone())?;
2070                }
2071                EventBody::Message { role, text, span } => {
2072                    store.append_message(child, *role, text.clone(), span.clone())?;
2073                }
2074                EventBody::GenerationPrompt { span } => {
2075                    store.commit_generation_prompt(child, span.clone())?;
2076                }
2077                EventBody::Generated {
2078                    span,
2079                    text,
2080                    channel,
2081                    finish,
2082                } => {
2083                    store.commit_generated(child, span.clone(), text.clone(), *channel, *finish)?;
2084                }
2085                EventBody::ToolUse { name, arguments } => {
2086                    let e = store.commit_tool_use(child, name.clone(), arguments.clone())?;
2087                    tool_ids.insert(id, e.event_id);
2088                }
2089                EventBody::ToolResult {
2090                    call_id,
2091                    content,
2092                    span,
2093                } => match tool_ids.get(call_id) {
2094                    Some(new_id) => {
2095                        store.commit_tool_result(child, *new_id, content.clone(), span.clone())?;
2096                    }
2097                    None => {
2098                        store.append_message(child, wal::role::TOOL, content.clone(), span.clone())?;
2099                    }
2100                },
2101                EventBody::ToolParseFailure { raw } => {
2102                    store.commit_tool_parse_failure(child, raw.clone())?;
2103                }
2104                EventBody::Created { .. }
2105                | EventBody::EpochBump
2106                | EventBody::GenerationFingerprint { .. }
2107                | EventBody::Forked { .. }
2108                | EventBody::Rebased { .. }
2109                | EventBody::MetaUpdated { .. }
2110                | EventBody::ToolExpired { .. }
2111                | EventBody::Purged { .. }
2112                | EventBody::Rerooted { .. }
2113                | EventBody::QosSet { .. } => {}
2114                EventBody::Block {
2115                    role,
2116                    kind,
2117                    payload,
2118                    span,
2119                    visibility_version,
2120                } => {
2121                    store.commit_block(child, *role, *kind, payload.clone(), span.clone(), *visibility_version)?;
2122                }
2123                EventBody::PermissionRequest { .. }
2124                | EventBody::PermissionResponse { .. }
2125                | EventBody::ToolCancelRequested { .. }
2126                | EventBody::ToolOutcome { .. }
2127                | EventBody::ToolReconciliation { .. }
2128                | EventBody::ToolLease { .. }
2129                | EventBody::Pinned { .. } => {}
2130            }
2131            i += 1;
2132        }
2133        Ok(child)
2134    }
2135
2136    pub fn trim(&self, session: u64, to_event: u64) -> Result<u64, DaemonError> {
2137        let max = {
2138            let store = self.store.lock().expect("store");
2139            store.session(session)?.next_event_id()
2140        };
2141        if to_event == 0 || to_event > max {
2142            return Err(DaemonError::ForkPoint {
2143                session,
2144                fork_at: to_event,
2145                max,
2146            });
2147        }
2148        if to_event == max {
2149            return self.fork(session, to_event, None);
2150        }
2151        self.rebase(
2152            session,
2153            vec![wal::RebaseEdit::Drop {
2154                from: to_event,
2155                to: max,
2156            }],
2157            None,
2158        )
2159    }
2160
2161    pub fn pin(&self, session: u64, ttl_ms: u64) -> Result<u64, DaemonError> {
2162        let deadline = if ttl_ms == 0 {
2163            0
2164        } else {
2165            wal::now_unix_ms().saturating_add(ttl_ms)
2166        };
2167        let (tokens, want) = {
2168            let mut store = self.store.lock().expect("store");
2169            store.set_pin(session, deadline)?;
2170            pins::session_target(&store.session(session)?.tokens)
2171        };
2172        self.send_pin(
2173            pins::PinKey::Session(session),
2174            tokens,
2175            want,
2176            deadline,
2177            Some(session),
2178        )?;
2179        Ok(deadline)
2180    }
2181
2182    pub fn pin_prefix(
2183        &self,
2184        anchor: Option<u64>,
2185        stream: &[u32],
2186        upto: usize,
2187        ttl_ms: u64,
2188    ) -> Result<u64, DaemonError> {
2189        let upto = upto.min(stream.len());
2190        let key = pins::PinKey::Prefix(pins::prefix_digest(&stream[..upto]));
2191        let tokens = if upto < stream.len() {
2192            stream[..=upto].to_vec()
2193        } else {
2194            pins::session_target(stream).0
2195        };
2196        let deadline = wal::now_unix_ms().saturating_add(ttl_ms.max(1));
2197        self.send_pin(key, tokens, upto as u64, deadline, anchor)
2198    }
2199
2200    pub fn pinned_prefix_tokens(
2201        &self,
2202        stream: &[u32],
2203        uptos: &[usize],
2204    ) -> Result<u64, DaemonError> {
2205        let keys = uptos
2206            .iter()
2207            .map(|&u| {
2208                let upto = u.min(stream.len());
2209                pins::PinKey::Prefix(pins::prefix_digest(&stream[..upto]))
2210            })
2211            .collect();
2212        let (tx, rx) = std::sync::mpsc::channel();
2213        self.sched_tx()
2214            .send(scheduler::SchedCmd::PinHeld(keys, tx))
2215            .map_err(|_| DaemonError::Protocol("scheduler is gone"))?;
2216        rx.recv()
2217            .map_err(|_| DaemonError::Protocol("scheduler dropped the query"))
2218    }
2219
2220    pub fn unanchor_pins(&self, session: u64) {
2221        let _ = self.sched_tx().send(scheduler::SchedCmd::PinUnanchor(session));
2222    }
2223
2224    fn send_pin(
2225        &self,
2226        key: pins::PinKey,
2227        tokens: Vec<u32>,
2228        want: u64,
2229        deadline_unix_ms: u64,
2230        anchor: Option<u64>,
2231    ) -> Result<u64, DaemonError> {
2232        let (tx, rx) = std::sync::mpsc::channel();
2233        self.sched_tx()
2234            .send(scheduler::SchedCmd::Pin(pins::PinReq {
2235                key,
2236                tokens,
2237                want,
2238                deadline_unix_ms,
2239                anchor,
2240                reply: Some(tx),
2241            }))
2242            .map_err(|_| DaemonError::Protocol("scheduler is gone"))?;
2243        rx.recv()
2244            .map_err(|_| DaemonError::Protocol("scheduler dropped the pin"))
2245    }
2246
2247    fn render_for(
2248        &self,
2249        store: &SessionStore,
2250        session: u64,
2251        role: u32,
2252        text: &str,
2253    ) -> Result<Vec<u32>, DaemonError> {
2254        let s = store.session(session)?;
2255        let mut span = if s.tokens.is_empty() {
2256            self.codec().stream_prologue()
2257        } else {
2258            match Self::last_stream_record(&s.events) {
2259                Some(EventBody::Generated { finish, .. })
2260                    if *finish == superfluid_abi::finish::EOS =>
2261                {
2262                    self.codec().turn_separator().unwrap_or_default()
2263                }
2264                _ => Vec::new(),
2265            }
2266        };
2267        match self.codec().render_message(role, text) {
2268            Some(rendered) => span.extend(rendered),
2269            None => {
2270                let encoded = self.codec().encode(text);
2271                if encoded.is_empty() {
2272                    return Err(DaemonError::Protocol("codec produced no span for the replacement"));
2273                }
2274                span.extend(encoded);
2275            }
2276        }
2277        Ok(span)
2278    }
2279
2280    pub fn purge(
2281        &self,
2282        session: u64,
2283        expected_generation: u64,
2284        mode: wal::PurgeMode,
2285    ) -> Result<Vec<u64>, DaemonError> {
2286        let targets: Vec<u64> = {
2287            let store = self.store.lock().expect("store");
2288            let s = store.session(session)?;
2289            if s.generation != expected_generation {
2290                return Err(DaemonError::GenerationConflict {
2291                    session,
2292                    expected: expected_generation,
2293                    actual: s.generation,
2294                });
2295            }
2296            let mut t = vec![session];
2297            if mode == wal::PurgeMode::Cascade {
2298                t.extend(store.descendants_of(session));
2299            }
2300            t
2301        };
2302        {
2303            let _a = self.active.lock().expect("active set");
2304            let mut p = self.purging.lock().expect("purging");
2305            p.extend(targets.iter().copied());
2306        }
2307        let _fence = PurgeFence {
2308            set: std::sync::Arc::clone(&self.purging),
2309            ids: targets.clone(),
2310        };
2311        for &t in &targets {
2312            self.quiesce(t)?;
2313        }
2314        let purged = {
2315            let mut store = self.store.lock().expect("store");
2316            for &t in &targets {
2317                let open: Vec<u64> = store.session(t)?.open_tool_calls.keys().copied().collect();
2318                for call in open {
2319                    store.expire_tool_call(t, call, "purged")?;
2320                }
2321            }
2322            store.purge(session, expected_generation, mode)?
2323        };
2324        let bus = self.bus();
2325        for &p in &purged {
2326            bus.end_session(p, bus::EndReason::Purged);
2327        }
2328        for &p in &purged {
2329            let _ = self.send_pin(pins::PinKey::Session(p), Vec::new(), 0, 0, None);
2330        }
2331        Ok(purged)
2332    }
2333
2334    fn refuse_if_purging(&self, session: u64) -> Result<(), DaemonError> {
2335        let _a = self.active.lock().expect("active set");
2336        if self.purging.lock().expect("purging").contains(&session) {
2337            return Err(DaemonError::Purged(session));
2338        }
2339        Ok(())
2340    }
2341
2342    fn quiesce(&self, session: u64) -> Result<(), DaemonError> {
2343        {
2344            let mut c = self.cancels.lock().expect("cancel registry");
2345            if self.active.lock().expect("active set").contains(&session) {
2346                c.insert(session);
2347            } else {
2348                return Ok(());
2349            }
2350        }
2351        self.cancels.wake();
2352        for _ in 0..2000 {
2353            if !self.active.lock().expect("active set").contains(&session) {
2354                return Ok(());
2355            }
2356            std::thread::sleep(std::time::Duration::from_millis(5));
2357        }
2358        Err(DaemonError::PurgeTimeout(session))
2359    }
2360
2361    pub fn append(
2362        &self,
2363        session: u64,
2364        text: Option<String>,
2365        span: Vec<u32>,
2366    ) -> Result<CommittedEvent, DaemonError> {
2367        {
2368            let active = self.active.lock().expect("active set");
2369            if active.contains(&session) {
2370                return Err(DaemonError::SessionBusy(session));
2371            }
2372            if self.purging.lock().expect("purging").contains(&session) {
2373                return Err(DaemonError::Purged(session));
2374            }
2375        }
2376        let span = if span.is_empty() {
2377            match &text {
2378                Some(t) => {
2379                    let encoded = self.codec().encode(t);
2380                    if encoded.is_empty() {
2381                        return Err(DaemonError::Protocol(
2382                            "append had no span and the codec produced none",
2383                        ));
2384                    }
2385                    encoded
2386                }
2387                None => return Err(DaemonError::Protocol("append carried neither text nor span")),
2388            }
2389        } else {
2390            span
2391        };
2392        self.store.lock().expect("store").append(session, text, span)
2393    }
2394
2395    fn separator_for(&self, session: u64) -> Vec<u32> {
2396        let store = self.store.lock().expect("store");
2397        let Ok(s) = store.session(session) else {
2398            return Vec::new();
2399        };
2400        if s.tokens.is_empty() {
2401            return self.codec().stream_prologue();
2402        }
2403        match Self::last_stream_record(&s.events) {
2404            Some(EventBody::Generated { finish, .. })
2405                if *finish == superfluid_abi::finish::EOS =>
2406            {
2407                self.codec().turn_separator().unwrap_or_default()
2408            }
2409            _ => Vec::new(),
2410        }
2411    }
2412
2413    pub fn append_message(
2414        &self,
2415        session: u64,
2416        role: u32,
2417        text: String,
2418    ) -> Result<CommittedEvent, DaemonError> {
2419        if self.active.lock().expect("active set").contains(&session) {
2420            return Err(DaemonError::SessionBusy(session));
2421        }
2422        if !matches!(role, wal::role::SYSTEM | wal::role::USER | wal::role::ASSISTANT) {
2423            return Err(DaemonError::UnknownRole(role));
2424        }
2425        let mut span = self.separator_for(session);
2426        span.extend(
2427            self.codec()
2428                .render_message(role, &text)
2429                .ok_or(DaemonError::Protocol("codec has no chat dialect"))?,
2430        );
2431        self.store
2432            .lock()
2433            .expect("store")
2434            .append_message(session, role, text, span)
2435    }
2436
2437    pub fn renders_per_message(&self) -> bool {
2438        self.codec().renders_per_message()
2439    }
2440
2441    pub fn append_system_full(
2442        &self,
2443        session: u64,
2444        system: Option<String>,
2445        tool_jsons: Vec<String>,
2446        kwargs: &serde_json::Map<String, serde_json::Value>,
2447    ) -> Result<CommittedEvent, DaemonError> {
2448        if self.active.lock().expect("active set").contains(&session) {
2449            return Err(DaemonError::SessionBusy(session));
2450        }
2451        let span = self
2452            .codec()
2453            .render_system_full(system.as_deref(), &tool_jsons, kwargs)
2454            .ok_or(DaemonError::Protocol("codec has no chat dialect"))?;
2455        self.store.lock().expect("store").append_message(
2456            session,
2457            wal::role::SYSTEM,
2458            system.unwrap_or_default(),
2459            span,
2460        )
2461    }
2462
2463    // Chat templates declare tools only in the opening system turn, so one rendered later would
2464    // put the model's input in a shape it was not trained on; it is refused instead.
2465    pub fn append_system(
2466        &self,
2467        session: u64,
2468        system: Option<String>,
2469        tools: Vec<String>,
2470    ) -> Result<CommittedEvent, DaemonError> {
2471        if system.is_none() && tools.is_empty() {
2472            return Err(DaemonError::Protocol("AppendSystem carried neither text nor tools"));
2473        }
2474        let tool_jsons = tools
2475            .iter()
2476            .enumerate()
2477            .map(|(index, t)| {
2478                codec::function_tool(t).map_err(|why| DaemonError::InvalidTool { index, why })
2479            })
2480            .collect::<Result<Vec<_>, _>>()?;
2481        if self.active.lock().expect("active set").contains(&session) {
2482            return Err(DaemonError::SessionBusy(session));
2483        }
2484        let codec = self.codec();
2485        let mut span = codec.stream_prologue();
2486        span.extend(
2487            codec
2488                .render_system_full(system.as_deref(), &tool_jsons, &serde_json::Map::new())
2489                .ok_or(DaemonError::Protocol("codec has no chat dialect"))?,
2490        );
2491        let mut store = self.store.lock().expect("store");
2492        if !store.session(session)?.tokens.is_empty() {
2493            return Err(DaemonError::SystemNotFirst(session));
2494        }
2495        store.append_message(session, wal::role::SYSTEM, system.unwrap_or_default(), span)
2496    }
2497
2498    pub(crate) fn render_conversation(
2499        &self,
2500        messages: &[codec::ChatMessage],
2501        tools: &[String],
2502    ) -> Option<Vec<u32>> {
2503        self.codec()
2504            .render_prompt_structured_with(messages, tools, &serde_json::Map::new())
2505    }
2506
2507    pub fn supports_reasoning_effort(&self) -> bool {
2508        self.codec().supports_reasoning_effort()
2509    }
2510
2511    pub fn reasoning_effort_levels(&self) -> Vec<&'static str> {
2512        self.codec().reasoning_effort_levels()
2513    }
2514
2515    pub fn supports_enable_thinking(&self) -> bool {
2516        self.codec().supports_enable_thinking()
2517    }
2518
2519    pub fn validate_template_kwargs(
2520        &self,
2521        kwargs: &serde_json::Map<String, serde_json::Value>,
2522    ) -> Result<(), String> {
2523        self.codec().validate_template_kwargs(kwargs)
2524    }
2525
2526    fn commit_opener(&self, session: u64, opener: Vec<u32>) -> Result<(), DaemonError> {
2527        if opener.is_empty() {
2528            return Ok(());
2529        }
2530        self.store
2531            .lock()
2532            .expect("store")
2533            .commit_generation_prompt(session, opener)
2534            .map(|_| ())
2535    }
2536
2537    pub fn append_conversation(
2538        &self,
2539        session: u64,
2540        messages: &[crate::codec::ChatMessage],
2541        tools: &[String],
2542        kwargs: &serde_json::Map<String, serde_json::Value>,
2543    ) -> Result<CommittedEvent, DaemonError> {
2544        let codec = self.codec();
2545        let mut span = codec
2546            .render_prompt_structured_with(messages, tools, kwargs)
2547            .ok_or_else(|| render_refusal(&*codec, messages, tools, kwargs, "codec has no chat dialect"))?;
2548        let opener_len = codec.trailing_generation_prompt(messages, tools, kwargs, &span);
2549        let opener: Vec<u32> = if opener_len > 0 && opener_len < span.len() {
2550            span.split_off(span.len() - opener_len)
2551        } else {
2552            Vec::new()
2553        };
2554        let text = messages.last().map(|m| m.content.clone());
2555        let blobs: Vec<(u32, String)> = messages
2556            .iter()
2557            .flat_map(|m| {
2558                m.parts.iter().filter_map(move |p| match p {
2559                    crate::codec::ContentPart::Image { blob } => Some((m.role, blob.clone())),
2560                    crate::codec::ContentPart::Audio { blob } => Some((m.role, blob.clone())),
2561                    crate::codec::ContentPart::Text(_) => None,
2562                })
2563            })
2564            .collect();
2565        if blobs.is_empty() {
2566            let ev = self.append(session, text, span)?;
2567            self.commit_opener(session, opener)?;
2568            return Ok(ev);
2569        }
2570        let mut probes = Vec::with_capacity(blobs.len());
2571        for (_, blob) in &blobs {
2572            if !self.media.exists(blob) {
2573                return Err(DaemonError::UnknownMedia(blob.clone()));
2574            }
2575            let path = self.media.path_of(blob);
2576            let info = match self.engine_sync(proto_linkw::StateSyncReqMsg::MediaProbe {
2577                image_path: path.to_string_lossy().into_owned(),
2578            })? {
2579                proto_linkw::StateSyncOkMsg::MediaProbe { info } => info,
2580                _ => return Err(DaemonError::Protocol("mismatched MediaProbe reply")),
2581            };
2582            if info.n_tokens == 0 || info.n_tokens as u64 > scheduler::PREFILL_BUDGET as u64 {
2583                return Err(DaemonError::MediaTooLarge(info.n_tokens));
2584            }
2585            probes.push(info);
2586        }
2587        let mut out = Vec::with_capacity(span.len() + 256);
2588        let mut offsets: Vec<usize> = Vec::with_capacity(probes.len());
2589        let mut img = 0usize;
2590        for (idx, &t) in span.iter().enumerate() {
2591            if img == probes.len() || t != probes[img].image_token_id {
2592                if probes.iter().any(|p| p.image_token_id == t) {
2593                    return Err(DaemonError::Protocol(
2594                        "rendered prompt contains a media marker with no matching pending part",
2595                    ));
2596                }
2597                out.push(t);
2598                continue;
2599            }
2600            let info = &probes[img];
2601            if info.boi_token_id != 0 && out.last() != Some(&info.boi_token_id) {
2602                out.push(info.boi_token_id);
2603            }
2604            offsets.push(out.len());
2605            out.extend(std::iter::repeat_n(info.image_token_id, info.n_tokens as usize));
2606            if info.eoi_token_id != 0 && span.get(idx + 1) != Some(&info.eoi_token_id) {
2607                out.push(info.eoi_token_id);
2608            }
2609            img += 1;
2610        }
2611        if img != probes.len() {
2612            return Err(DaemonError::Protocol(
2613                "this model's chat template did not emit an image marker for every image part;                  it cannot place media",
2614            ));
2615        }
2616        let stream_len = {
2617            let store = self.store.lock().expect("store");
2618            store.session(session)?.tokens.len() as u64
2619        };
2620        let ev = self.append(session, text, out)?;
2621        for (((role, blob), info), off) in blobs.iter().zip(&probes).zip(&offsets) {
2622            let payload = serde_json::json!({
2623                "blob": blob,
2624                "n_tokens": info.n_tokens,
2625                "image_token_id": info.image_token_id,
2626                "offset": stream_len + *off as u64,
2627                "pre_text": "",
2628                "post_text": "",
2629                "media_identity": {
2630                    "preprocessing_fp": format!("{:016x}", info.preprocess_fp),
2631                    "encoder": "bundle-weights-identity",
2632                },
2633            })
2634            .to_string();
2635            let mut store = self.store.lock().expect("store");
2636            store.commit_block(
2637                session,
2638                *role,
2639                wal::block_kind::IMAGE,
2640                payload,
2641                Vec::new(),
2642                codec::MODEL_VISIBILITY_VERSION,
2643            )?;
2644        }
2645        self.commit_opener(session, opener)?;
2646        Ok(ev)
2647    }
2648
2649    pub fn response_format_tag(&self, json_schema: &str) -> Option<String> {
2650        self.codec().response_format_tag(json_schema)
2651    }
2652
2653    pub fn tool_structural_tag(&self, tool_jsons: &[String], at_least_one: bool) -> Option<String> {
2654        self.codec().structural_tag(tool_jsons, at_least_one)
2655    }
2656
2657    pub fn grammar_create_structural(&self, tag_json: &str) -> Result<u32, DaemonError> {
2658        match self.engine_sync(proto_linkw::StateSyncReqMsg::CreateGrammarStructural {
2659            tag_json: tag_json.to_string(),
2660        })? {
2661            proto_linkw::StateSyncOkMsg::CreateGrammar { grammar_handle } => Ok(grammar_handle),
2662            _ => Err(DaemonError::Protocol("mismatched CreateGrammar reply")),
2663        }
2664    }
2665
2666    pub fn parse_tool_call(&self, raw: &str) -> Option<(String, String)> {
2667        self.codec().parse_tool_call(raw)
2668    }
2669
2670    pub fn named_tool_call(&self, raw: &str) -> Option<(String, String)> {
2671        self.codec().named_tool_call(raw)
2672    }
2673
2674    pub fn parse_tool_call_with(
2675        &self,
2676        raw: &str,
2677        schemas: Option<&crate::codec::ToolSchemas>,
2678    ) -> Option<(String, String)> {
2679        self.codec().parse_tool_call_with(raw, schemas)
2680    }
2681
2682    pub fn tool_fragment_mode(&self) -> crate::tool_fragment::ToolFragmentMode {
2683        self.codec().tool_fragment_mode()
2684    }
2685
2686    pub fn tool_call_grammar(&self, tool_jsons: &[String]) -> Option<String> {
2687        self.codec().call_grammar(tool_jsons)
2688    }
2689
2690    pub fn append_system_with_tools(
2691        &self,
2692        session: u64,
2693        system: Option<String>,
2694        tool_jsons: Vec<String>,
2695    ) -> Result<CommittedEvent, DaemonError> {
2696        if self.active.lock().expect("active set").contains(&session) {
2697            return Err(DaemonError::SessionBusy(session));
2698        }
2699        let span = self
2700            .codec()
2701            .render_system_with_tools(system.as_deref(), &tool_jsons)
2702            .ok_or(DaemonError::Protocol("codec has no tool dialect"))?;
2703        self.store.lock().expect("store").append_message(
2704            session,
2705            wal::role::SYSTEM,
2706            system.unwrap_or_default(),
2707            span,
2708        )
2709    }
2710
2711    /// `after_query`: no user query follows this turn in the conversation
2712    /// (a step of a tool loop), which some templates render differently.
2713    pub fn append_assistant_with_tool_calls(
2714        &self,
2715        session: u64,
2716        content: String,
2717        calls: Vec<(String, String)>,
2718        after_query: bool,
2719    ) -> Result<CommittedEvent, DaemonError> {
2720        if self.active.lock().expect("active set").contains(&session) {
2721            return Err(DaemonError::SessionBusy(session));
2722        }
2723        let mut span = self.separator_for(session);
2724        let codec = self.codec();
2725        let turn = if after_query {
2726            codec.render_assistant_with_tool_calls_after_query(&content, &calls)
2727        } else {
2728            codec.render_assistant_with_tool_calls(&content, &calls)
2729        };
2730        span.extend(turn.ok_or(DaemonError::Protocol("codec has no tool dialect"))?);
2731        self.store.lock().expect("store").append_message(
2732            session,
2733            wal::role::ASSISTANT,
2734            content,
2735            span,
2736        )
2737    }
2738
2739    pub fn append_tool_result(
2740        &self,
2741        session: u64,
2742        call_id: u64,
2743        content: String,
2744    ) -> Result<CommittedEvent, DaemonError> {
2745        if self.active.lock().expect("active set").contains(&session) {
2746            return Err(DaemonError::SessionBusy(session));
2747        }
2748        let name = {
2749            let mut store = self.store.lock().expect("store");
2750            let s = store.session(session)?;
2751            if call_id < s.base {
2752                return Err(DaemonError::InheritedToolCall(call_id));
2753            }
2754            match s.ledger.get(&call_id) {
2755                None => return Err(DaemonError::UnknownToolCall(call_id)),
2756                Some(e)
2757                    if e.state == wal::ledger_state::OPEN
2758                        || e.state == wal::ledger_state::CANCEL_REQUESTED =>
2759                {
2760                    if let Some(deadline) = e.lease_deadline_unix_ms {
2761                        if wal::now_unix_ms() > deadline {
2762                            store.expire_tool_call(session, call_id, "lease expired")?;
2763                            return Err(DaemonError::ToolLeaseExpired(call_id));
2764                        }
2765                    }
2766                    e.name.clone()
2767                }
2768                Some(e) if e.state == wal::ledger_state::SUCCEEDED => {
2769                    return Err(DaemonError::UnknownToolCall(call_id));
2770                }
2771                Some(_) => {
2772                    store.commit_tool_reconciliation(
2773                        session,
2774                        call_id,
2775                        format!("late result: {} bytes", content.len()),
2776                    )?;
2777                    return Err(DaemonError::ToolLateEffect(call_id));
2778                }
2779            }
2780        };
2781        let mut span = self.separator_for(session);
2782        span.extend(
2783            self.codec()
2784                .render_tool_result(&name, &content)
2785                .ok_or(DaemonError::Protocol("codec has no tool dialect"))?,
2786        );
2787        self.store
2788            .lock()
2789            .expect("store")
2790            .commit_tool_result(session, call_id, content, span)
2791    }
2792
2793    pub fn append_block(
2794        &self,
2795        session: u64,
2796        role: u32,
2797        kind: u32,
2798        payload: String,
2799    ) -> Result<CommittedEvent, DaemonError> {
2800        if self.active.lock().expect("active set").contains(&session) {
2801            return Err(DaemonError::SessionBusy(session));
2802        }
2803        if !matches!(role, wal::role::SYSTEM | wal::role::USER | wal::role::ASSISTANT | wal::role::TOOL) {
2804            return Err(DaemonError::UnknownRole(role));
2805        }
2806        if codec::project_block_text(kind, &payload).is_none() {
2807            return Err(DaemonError::UnknownBlockKind(kind));
2808        }
2809        let projected = self
2810            .codec()
2811            .render_block(role, kind, &payload)
2812            .ok_or(DaemonError::Protocol("codec has no chat dialect"))?;
2813        let span = if projected.is_empty() {
2814            Vec::new()
2815        } else {
2816            let mut span = self.separator_for(session);
2817            span.extend(projected);
2818            span
2819        };
2820        self.store.lock().expect("store").commit_block(
2821            session,
2822            role,
2823            kind,
2824            payload,
2825            span,
2826            codec::MODEL_VISIBILITY_VERSION,
2827        )
2828    }
2829
2830    pub fn request_permission(
2831        &self,
2832        session: u64,
2833        call_id: u64,
2834        text: String,
2835    ) -> Result<CommittedEvent, DaemonError> {
2836        let active = self.active.lock().expect("active set");
2837        if active.contains(&session) {
2838            return Err(DaemonError::SessionBusy(session));
2839        }
2840        if self.purging.lock().expect("purging").contains(&session) {
2841            return Err(DaemonError::Purged(session));
2842        }
2843        self.store
2844            .lock()
2845            .expect("store")
2846            .commit_permission_request(session, call_id, text)
2847    }
2848
2849    pub fn respond_permission(
2850        &self,
2851        session: u64,
2852        request_id: u64,
2853        granted: bool,
2854    ) -> Result<CommittedEvent, DaemonError> {
2855        self.store
2856            .lock()
2857            .expect("store")
2858            .commit_permission_response(session, request_id, granted)
2859    }
2860
2861    pub fn request_tool_cancel(&self, session: u64, call_id: u64) -> Result<CommittedEvent, DaemonError> {
2862        self.store
2863            .lock()
2864            .expect("store")
2865            .commit_tool_cancel_request(session, call_id)
2866    }
2867
2868    pub fn append_tool_outcome(
2869        &self,
2870        session: u64,
2871        call_id: u64,
2872        outcome: u8,
2873        note: String,
2874    ) -> Result<CommittedEvent, DaemonError> {
2875        self.store
2876            .lock()
2877            .expect("store")
2878            .commit_tool_outcome(session, call_id, outcome, note)
2879    }
2880
2881    pub fn ledger(&self, session: u64) -> Result<Vec<(u64, store::LedgerEntry)>, DaemonError> {
2882        Ok(self
2883            .store
2884            .lock()
2885            .expect("store")
2886            .session(session)?
2887            .ledger
2888            .iter()
2889            .map(|(id, e)| (*id, e.clone()))
2890            .collect())
2891    }
2892
2893    pub fn append_tool_history(
2894        &self,
2895        session: u64,
2896        content: String,
2897    ) -> Result<CommittedEvent, DaemonError> {
2898        if self.active.lock().expect("active set").contains(&session) {
2899            return Err(DaemonError::SessionBusy(session));
2900        }
2901        let mut span = self.separator_for(session);
2902        span.extend(
2903            self.codec()
2904                .render_tool_result("", &content)
2905                .ok_or(DaemonError::Protocol("codec has no tool dialect"))?,
2906        );
2907        self.store.lock().expect("store").append_message(
2908            session,
2909            wal::role::TOOL,
2910            content,
2911            span,
2912        )
2913    }
2914
2915    pub fn open_tool_calls(
2916        &self,
2917        session: u64,
2918    ) -> Result<Vec<(u64, String, String)>, DaemonError> {
2919        Ok(self
2920            .store
2921            .lock()
2922            .expect("store")
2923            .session(session)?
2924            .open_tool_calls
2925            .iter()
2926            .map(|(id, (n, a))| (*id, n.clone(), a.clone()))
2927            .collect())
2928    }
2929
2930    pub fn generate(&self, session: u64, max_tokens: u32) -> Result<GenerateOutcome, DaemonError> {
2931        self.generate_streaming(session, max_tokens, |_| Ok(()))
2932    }
2933
2934    pub fn generate_streaming(
2935        &self,
2936        session: u64,
2937        max_tokens: u32,
2938        on_event: impl FnMut(&CommittedEvent) -> Result<(), DaemonError>,
2939    ) -> Result<GenerateOutcome, DaemonError> {
2940        self.generate_streaming_ex(session, max_tokens, scheduler::GenExtras::default(), on_event)
2941    }
2942
2943    pub fn generate_streaming_ex(
2944        &self,
2945        session: u64,
2946        max_tokens: u32,
2947        extras: scheduler::GenExtras,
2948        on_event: impl FnMut(&CommittedEvent) -> Result<(), DaemonError>,
2949    ) -> Result<GenerateOutcome, DaemonError> {
2950        self.generate_streaming_ex_lp(session, max_tokens, extras, on_event, |_| {})
2951    }
2952
2953    pub fn generate_streaming_ex_lp(
2954        &self,
2955        session: u64,
2956        max_tokens: u32,
2957        extras: scheduler::GenExtras,
2958        on_event: impl FnMut(&CommittedEvent) -> Result<(), DaemonError>,
2959        on_logprobs: impl FnMut(&[scheduler::TokenLogprob]),
2960    ) -> Result<GenerateOutcome, DaemonError> {
2961        self.generate_streaming_inner(session, max_tokens, extras, on_event, on_logprobs, |_, _, _| {})
2962    }
2963
2964    pub fn generate_streaming_delta(
2965        &self,
2966        session: u64,
2967        max_tokens: u32,
2968        mut extras: scheduler::GenExtras,
2969        on_event: impl FnMut(&CommittedEvent) -> Result<(), DaemonError>,
2970        on_logprobs: impl FnMut(&[scheduler::TokenLogprob]),
2971        on_delta: impl FnMut(u32, &str, u32),
2972    ) -> Result<GenerateOutcome, DaemonError> {
2973        extras.wants_deltas = true;
2974        self.generate_streaming_inner(session, max_tokens, extras, on_event, on_logprobs, on_delta)
2975    }
2976
2977    fn generation_opener(&self, events: &[store::CommittedEvent], thinking: Option<bool>) -> Option<Vec<u32>> {
2978        match Self::last_stream_record(events) {
2979            Some(EventBody::Message { .. }) | Some(EventBody::ToolResult { .. }) => self
2980                .codec()
2981                .generation_prefix_with(codec::TurnState::AfterInput, thinking),
2982            Some(EventBody::Generated { finish, .. }) if *finish == superfluid_abi::finish::EOS => self
2983                .codec()
2984                .generation_prefix_with(codec::TurnState::AfterFinishedTurn, thinking),
2985            _ => None,
2986        }
2987    }
2988
2989    /// Where a continuation of `session` forks (one past its last event),
2990    /// and whether its last turn ended: an append after an unended turn
2991    /// lands inside it.
2992    pub(crate) fn session_end(&self, session: u64) -> Option<(u64, bool)> {
2993        let store = self.store.lock().expect("store");
2994        let s = store.session(session).ok()?;
2995        let ended = matches!(
2996            Self::last_stream_record(&s.events),
2997            Some(EventBody::Generated { finish, .. }) if *finish == superfluid_abi::finish::EOS
2998        );
2999        Some((s.next_event_id(), ended))
3000    }
3001
3002    fn last_stream_record(events: &[store::CommittedEvent]) -> Option<&EventBody> {
3003        events
3004            .iter()
3005            .rev()
3006            .map(|e| &e.body)
3007            .find(|b| b.span().is_some_and(|s| !s.is_empty()))
3008    }
3009
3010    fn open_turn_tokens(events: &[store::CommittedEvent]) -> Vec<u32> {
3011        let mut spans: Vec<&[u32]> = Vec::new();
3012        for e in events.iter().rev() {
3013            match &e.body {
3014                EventBody::Generated { span, finish, .. } => {
3015                    if spans.is_empty() && *finish == superfluid_abi::finish::EOS {
3016                        return Vec::new();
3017                    }
3018                    spans.push(span);
3019                }
3020                EventBody::GenerationPrompt { span } => {
3021                    spans.push(span);
3022                    break;
3023                }
3024                b => match b.span() {
3025                    None | Some([]) => continue,
3026                    Some(_) => break,
3027                },
3028            }
3029        }
3030        spans.into_iter().rev().flatten().copied().collect()
3031    }
3032
3033    fn generate_streaming_inner(
3034        &self,
3035        session: u64,
3036        max_tokens: u32,
3037        extras: scheduler::GenExtras,
3038        on_event: impl FnMut(&CommittedEvent) -> Result<(), DaemonError>,
3039        on_logprobs: impl FnMut(&[scheduler::TokenLogprob]),
3040        on_delta: impl FnMut(u32, &str, u32),
3041    ) -> Result<GenerateOutcome, DaemonError> {
3042        self.ensure_serving()?;
3043        let _span = tracing::info_span!(
3044            parent: None,
3045            "generate",
3046            session,
3047            max_tokens,
3048            model = self.trace_scope.model.as_str(),
3049            store_ns = self.trace_scope.ns
3050        )
3051        .entered();
3052        let (stream, params, class, batch_invariant, media) = {
3053            let mut cancels = self.cancels.lock().expect("cancel registry");
3054            let mut active = self.active.lock().expect("active set");
3055            if self.purging.lock().expect("purging").contains(&session) {
3056                return Err(DaemonError::Purged(session));
3057            }
3058            if self.tool_lease_ms > 0 {
3059                let now = crate::wal::now_unix_ms();
3060                let mut store = self.store.lock().expect("store");
3061                let elapsed: Vec<u64> = {
3062                    let s = store.session(session)?;
3063                    s.open_tool_calls
3064                        .keys()
3065                        .filter(|id| {
3066                            s.ledger
3067                                .get(id)
3068                                .and_then(|e| e.lease_deadline_unix_ms)
3069                                .is_some_and(|d| d <= now)
3070                        })
3071                        .copied()
3072                        .collect()
3073                };
3074                for call_id in elapsed {
3075                    let _ = store.expire_tool_call(session, call_id, "lease elapsed");
3076                }
3077            }
3078            let snap = {
3079                let store = self.store.lock().expect("store");
3080                let s = store.session(session)?;
3081                if !s.open_permissions.is_empty() {
3082                    return Err(DaemonError::PermissionPending(session));
3083                }
3084                if !s.open_tool_calls.is_empty() {
3085                    tracing::warn!(
3086                        session,
3087                        open_tool_calls = s.open_tool_calls.len(),
3088                        "generating with unanswered tool calls: the model cannot see its own \
3089                         tool output and will re-issue the call. Return a result for each open \
3090                         call id (OpenAI: a `role: \"tool\"` message carrying its `tool_call_id`), \
3091                         or run with --tool-lease-ms so unanswered calls expire instead of \
3092                         accumulating."
3093                    );
3094                }
3095                (
3096                    s.tokens.clone(),
3097                    s.params,
3098                    s.qos_class,
3099                    s.batch_invariant,
3100                    self.media_of(s),
3101                )
3102            };
3103            if snap.0.is_empty() {
3104                return Err(DaemonError::EmptySession(session));
3105            }
3106            self.check_stream_fits(snap.0.len())?;
3107            if max_tokens == 0 {
3108                return Ok(GenerateOutcome {
3109                    events: Vec::new(),
3110                    tokens_generated: 0,
3111                    finish: superfluid_abi::finish::NONE,
3112                    warm_prefix: 0,
3113                    logprobs: Vec::new(),
3114                    spec: (0, 0),
3115                    decode: (0, 0.0),
3116                });
3117            }
3118            if !active.insert(session) {
3119                return Err(DaemonError::SessionBusy(session));
3120            }
3121            // A cancel aimed at the previous generation can land after its
3122            // lane stopped reading cancels (finishing, or already retired);
3123            // it must not stop this one.
3124            cancels.remove(&session);
3125            snap
3126        };
3127        if let Some(mut fp) = self.codec().behavior_fingerprint() {
3128            fp.speculation = self.serving().speculation.identity_for_fingerprint();
3129            let commit = self
3130                .store
3131                .lock()
3132                .expect("store")
3133                .commit_generation_fingerprint(session, fp.digest());
3134            if let Err(e) = commit {
3135                self.active.lock().expect("active set").remove(&session);
3136                return Err(e);
3137            }
3138        }
3139        let mut stream = stream;
3140        let (opener, open_turn) = {
3141            let store = self.store.lock().expect("store");
3142            let events = &store.session(session)?.events;
3143            let opener = self.generation_opener(events, extras.thinking);
3144            let open_turn: Vec<u32> = if opener.is_some() {
3145                Vec::new()
3146            } else {
3147                Self::open_turn_tokens(events)
3148            };
3149            (opener, open_turn)
3150        };
3151        let opener_len = opener.as_ref().map(Vec::len);
3152        let input_end = if opener.is_some() { stream.len() as u64 } else { 0 };
3153        if let Some(op) = opener {
3154            if let Err(e) = self.check_stream_fits(stream.len() + op.len()) {
3155                self.active.lock().expect("active set").remove(&session);
3156                return Err(e);
3157            }
3158            let commit = self
3159                .store
3160                .lock()
3161                .expect("store")
3162                .commit_generation_prompt(session, op.clone());
3163            if let Err(e) = commit {
3164                self.active.lock().expect("active set").remove(&session);
3165                return Err(e);
3166            }
3167            stream.extend_from_slice(&op);
3168        }
3169        let mut extras = extras;
3170        if extras.open_channel == 0 {
3171            let scope: &[u32] = match opener_len {
3172                Some(n) => &stream[stream.len() - n..],
3173                None => &open_turn,
3174            };
3175            let open = codec::trailing_open_channel(scope, &self.codec().channel_markers());
3176            if open != 0 {
3177                extras.open_channel = open;
3178            }
3179        }
3180        let result = scheduler::run_job(
3181            &self.sched_tx(),
3182            session,
3183            stream,
3184            input_end,
3185            params,
3186            max_tokens,
3187            class,
3188            batch_invariant,
3189            media,
3190            extras,
3191            &self.cancels,
3192            on_event,
3193            on_logprobs,
3194            on_delta,
3195        );
3196        match result {
3197            Ok((events, tokens_generated, finish, warm_prefix, logprobs, spec, decode)) => Ok(GenerateOutcome {
3198                events,
3199                tokens_generated,
3200                finish,
3201                warm_prefix,
3202                logprobs,
3203                spec,
3204                decode,
3205            }),
3206            Err(e) => {
3207                self.active.lock().expect("active set").remove(&session);
3208                Err(e)
3209            }
3210        }
3211    }
3212
3213    pub fn read(&self, session: u64, cursor: u64) -> Result<Vec<CommittedEvent>, DaemonError> {
3214        Ok(self.store.lock().expect("store").read(session, cursor)?.to_vec())
3215    }
3216
3217    pub fn session_ids(&self) -> Vec<u64> {
3218        self.store.lock().expect("store").session_ids()
3219    }
3220}
3221
3222#[cfg(test)]
3223mod open_turn_tests {
3224    use super::*;
3225    use superfluid_abi::finish;
3226
3227    fn ev(event_id: u64, body: EventBody) -> store::CommittedEvent {
3228        store::CommittedEvent { event_id, epoch: 0, ts_unix_ms: 0, body }
3229    }
3230
3231    fn generated(span: &[u32], finish: u32) -> EventBody {
3232        EventBody::Generated { span: span.to_vec(), text: String::new(), channel: 0, finish }
3233    }
3234
3235    #[test]
3236    fn open_turn_skips_every_token_less_record() {
3237        let events = vec![
3238            ev(0, EventBody::Message { role: 1, text: "hi".into(), span: vec![900, 901] }),
3239            ev(1, EventBody::GenerationPrompt { span: vec![1, 2] }),
3240            ev(2, generated(&[3, 4], finish::NONE)),
3241            ev(3, EventBody::ToolUse { name: "x".into(), arguments: "{".into() }),
3242            ev(4, EventBody::ToolLease { call_id: 3, deadline_unix_ms: 1 }),
3243            ev(5, EventBody::ToolReconciliation { call_id: 3, note: String::new() }),
3244            ev(6, EventBody::EpochBump),
3245            ev(7, EventBody::Pinned { deadline_unix_ms: 1 }),
3246            ev(8, EventBody::Block { role: 1, kind: 7, payload: "{}".into(), span: vec![], visibility_version: 1 }),
3247            ev(9, generated(&[5], finish::LENGTH)),
3248        ];
3249        assert_eq!(Daemon::open_turn_tokens(&events), vec![1, 2, 3, 4, 5]);
3250    }
3251
3252    #[test]
3253    fn token_less_records_do_not_reopen_a_closed_turn() {
3254        let closed = vec![
3255            ev(0, EventBody::GenerationPrompt { span: vec![1] }),
3256            ev(1, generated(&[2], finish::EOS)),
3257            ev(2, EventBody::ToolLease { call_id: 9, deadline_unix_ms: 1 }),
3258            ev(3, EventBody::QosSet { class: 0, batch_invariant: false }),
3259        ];
3260        assert!(Daemon::open_turn_tokens(&closed).is_empty());
3261        let after_result = vec![
3262            ev(0, EventBody::GenerationPrompt { span: vec![1] }),
3263            ev(1, generated(&[2], finish::NONE)),
3264            ev(2, EventBody::ToolUse { name: "x".into(), arguments: "{}".into() }),
3265            ev(3, EventBody::ToolResult { call_id: 2, content: "ok".into(), span: vec![7, 8] }),
3266            ev(4, EventBody::ToolLease { call_id: 2, deadline_unix_ms: 1 }),
3267        ];
3268        assert!(Daemon::open_turn_tokens(&after_result).is_empty());
3269    }
3270
3271    #[test]
3272    fn the_last_stream_record_is_the_last_one_with_tokens() {
3273        let events = vec![
3274            ev(0, EventBody::GenerationPrompt { span: vec![1] }),
3275            ev(1, generated(&[2], finish::EOS)),
3276            ev(2, EventBody::ToolUse { name: "x".into(), arguments: "{}".into() }),
3277            ev(3, EventBody::ToolLease { call_id: 2, deadline_unix_ms: 1 }),
3278            ev(4, EventBody::Block { role: 1, kind: 7, payload: "{}".into(), span: vec![], visibility_version: 1 }),
3279            ev(5, EventBody::GenerationFingerprint { digest: 9 }),
3280        ];
3281        assert!(matches!(
3282            Daemon::last_stream_record(&events),
3283            Some(EventBody::Generated { finish: finish::EOS, .. })
3284        ));
3285        assert!(Daemon::last_stream_record(&events[2..]).is_none());
3286    }
3287}
3288
3289fn first_registrable<E>(
3290    candidates: &[String],
3291    auto: bool,
3292    mut register: impl FnMut(&str) -> Result<(), E>,
3293    mut refused: impl FnMut(&str, &E, Option<&str>),
3294) -> Result<Option<String>, E> {
3295    for (i, c) in candidates.iter().enumerate() {
3296        match register(c) {
3297            Ok(()) => return Ok(Some(c.clone())),
3298            Err(e) if auto => refused(c, &e, candidates.get(i + 1).map(String::as_str)),
3299            Err(e) => return Err(e),
3300        }
3301    }
3302    Ok(None)
3303}
3304
3305#[cfg(test)]
3306mod speculate_fallback_tests {
3307    use super::first_registrable;
3308
3309    fn strs(v: &[&str]) -> Vec<String> {
3310        v.iter().map(|s| s.to_string()).collect()
3311    }
3312
3313    #[test]
3314    fn auto_tries_each_candidate_in_order_until_one_registers() {
3315        let mut tried = Vec::new();
3316        let mut reported = Vec::new();
3317        let picked = first_registrable(
3318            &strs(&["dspark:a", "dflash:b", "eagle3:c"]),
3319            true,
3320            |s| {
3321                tried.push(s.to_string());
3322                if s == "dflash:b" { Ok(()) } else { Err(format!("no {s}")) }
3323            },
3324            |s, e, next| reported.push((s.to_string(), e.clone(), next.map(str::to_string))),
3325        )
3326        .unwrap();
3327        assert_eq!(picked.as_deref(), Some("dflash:b"));
3328        assert_eq!(tried, ["dspark:a", "dflash:b"]);
3329        assert_eq!(reported, [("dspark:a".to_string(), "no dspark:a".to_string(), Some("dflash:b".to_string()))]);
3330    }
3331
3332    #[test]
3333    fn auto_with_every_candidate_refused_serves_plainly() {
3334        let mut last_next = Some("x".to_string());
3335        let picked = first_registrable(
3336            &strs(&["dspark:a", "dflash:b"]),
3337            true,
3338            |_| Err::<(), _>("refused"),
3339            |_, _, next| last_next = next.map(str::to_string),
3340        )
3341        .unwrap();
3342        assert_eq!(picked, None);
3343        assert_eq!(last_next, None, "the last refusal says it is decoding plainly");
3344    }
3345
3346    #[test]
3347    fn an_explicit_strategy_fails_on_the_first_refusal() {
3348        let mut reported = 0;
3349        let r = first_registrable(
3350            &strs(&["dspark:a", "dflash:b"]),
3351            false,
3352            |s| if s == "dspark:a" { Err("refused") } else { Ok(()) },
3353            |_, _, _| reported += 1,
3354        );
3355        assert_eq!(r, Err("refused"));
3356        assert_eq!(reported, 0);
3357    }
3358}
3359
3360#[cfg(test)]
3361mod latency_tests {
3362    use super::CancelSet;
3363    use std::time::Duration;
3364
3365    #[test]
3366    fn a_latency_request_raises_the_word_and_holds_it_for_a_while() {
3367        let ctl = superfluid_shm::ControlSegment::create().unwrap();
3368        let worker = superfluid_shm::ControlSegment::attach(
3369            superfluid_shm::ShmSegment::from_fd(
3370                {
3371                    use std::os::fd::FromRawFd;
3372                    // SAFETY: a fresh duplicate of the segment's fd, owned here.
3373                    unsafe { std::os::fd::OwnedFd::from_raw_fd(libc::dup(ctl.raw_fd())) }
3374                },
3375                superfluid_shm::CONTROL_BYTES,
3376            )
3377            .unwrap(),
3378        )
3379        .unwrap();
3380        let cancels = CancelSet::default();
3381        cancels.wake_with(std::sync::Arc::new(std::sync::RwLock::new(Some(ctl.signal().unwrap()))));
3382        assert!(!cancels.latency_recent(Duration::from_secs(300)), "nothing has come");
3383        cancels.latency_request();
3384        assert!(worker.latency_wanted() && worker.yield_requested(), "short steps, and the running tick ends at its next");
3385        assert!(cancels.latency_recent(Duration::from_secs(300)));
3386        assert!(!cancels.latency_recent(Duration::ZERO), "past the hold");
3387    }
3388}
superfluid_daemon/lib.rs

superfluid_daemon/
lib.rs
