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

superfluid_daemon/
lib.rs
