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

superfluid_daemon/
lib.rs
