Skip to main content

superfluid_daemon/
scheduler.rs

1//! The multi-lane scheduler.
2
3use std::collections::{HashMap, HashSet, VecDeque};
4use std::path::PathBuf;
5use std::sync::atomic::Ordering;
6use std::sync::mpsc::{Receiver, Sender, TryRecvError};
7use std::sync::{Arc, Mutex};
8
9use superfluid_abi::{encoding, finish, op_state};
10use superfluid_proto::linkw::{self, ALL_SPACES};
11
12use crate::bus::SessionBus;
13use crate::codec::{Channelizer, TextCodec, Utf8Stream};
14use crate::park;
15use crate::runtime::EngineHost;
16use crate::store::{CommittedEvent, SessionStore};
17use crate::wal::GenParams;
18use crate::DaemonError;
19
20const KV_SPACE: u32 = 1;
21const MAX_WORKER_RESPAWNS: u32 = 3;
22// An admission stages at most two token-in slots: its prompt and a resumed prefix.
23const STAGES_PER_ADMISSION: u32 = 2;
24
25fn lane_cap(configured: usize, max_seqs: Option<u64>, emit_slots: Option<u32>) -> usize {
26    let lanes = match max_seqs {
27        Some(n) if (n as usize) < configured => {
28            eprintln!("[superfluid] {configured} lanes asked for; the runtime holds {n} sequences, so {n} run");
29            (n as usize).max(1)
30        }
31        _ => configured,
32    };
33    // Every decoding lane writes one token-out slot a tick, read back after
34    // it: more lanes than slots would overwrite the first emits.
35    match emit_slots {
36        Some(n) if (n as usize) < lanes => {
37            eprintln!("[superfluid] {lanes} lanes asked for; the worker's token-out ring holds {n} emits a tick, so {n} run");
38            (n as usize).max(1)
39        }
40        _ => lanes,
41    }
42}
43
44fn ring_slots(host: &mut EngineHost, ring_id: u32) -> Option<u32> {
45    if !host.is_alive() {
46        return None;
47    }
48    host.client()
49        .hello
50        .limits
51        .ring_specs
52        .iter()
53        .find(|r| r.ring_id == ring_id)
54        .map(|r| r.slots)
55}
56
57fn worker_is_gone(e: &DaemonError) -> bool {
58    matches!(
59        e,
60        DaemonError::Agent(
61            superfluid_agent::AgentError::Io(_)
62                | superfluid_agent::AgentError::Closed
63                | superfluid_agent::AgentError::Envelope(_)
64        )
65    )
66}
67
68pub const PREFILL_BUDGET: u32 = 4096;
69pub const SOLO_PREFILL_BUDGET: u32 = 16384;
70const TICK_DECODE: u16 = 32;
71const TICK_DECODE_MAX: u16 = 256;
72const DECODE_SHARE_BESIDE_PREFILL: f64 = 0.25;
73/// The same while a chat or an inline completion decodes: its reply is
74/// short and watched, so the prefill beside it waits more, not the reply.
75const DECODE_SHARE_BESIDE_PREFILL_INTERACTIVE: f64 = 2.0;
76const TICK_DECODE_BESIDE_PREFILL: u16 = 8;
77const TARGET_TICK_TOTAL_SECS: f64 = 2.0;
78const OPEN_TICK_SECS: f64 = 0.25;
79const MAX_PREFILL_TICK_SECS: f64 = 8.0;
80const MAX_LONE_PREFILL_TICK_SECS: f64 = 16.0;
81/// A tick's prefill while a request could still join: a newcomer waits for
82/// the tick in flight, so this bounds its time to first token.
83const OPEN_PREFILL_TICK_SECS: f64 = 0.5;
84/// And in tokens, whatever the learned rate says: an engine that returns
85/// before its GPU work ends makes an early tick look all but free.
86const OPEN_PREFILL_TICK_MAX: u32 = 1024;
87
88/// The open tick's target on an engine that ends a tick at its next step when a
89/// preempting request arrives.
90const OPEN_PREFILL_TICK_SECS_YIELDING: f64 = 2.0;
91
92/// How long after the last chat or inline completion the engine keeps its
93/// prefill steps short.
94const LATENCY_HOLD: std::time::Duration = std::time::Duration::from_secs(300);
95
96/// A queued request waits for a lane prefilling the same prefix (at least this long) when
97/// that lane retires soon, and starts from its cached prefix.
98const SHARED_PREFIX_MIN_TOKENS: u64 = 1024;
99/// "Soon": at most this many tokens left to generate.
100const SHARED_PREFIX_TAIL_TOKENS: u32 = 64;
101/// Never hold a request longer than this.
102const SHARED_PREFIX_MAX_HOLD_TICKS: u64 = 64;
103/// On a runtime whose state cannot be cut back, a chat's input of at least
104/// this many tokens is checkpointed where it ends (the executor's floor for
105/// the checkpoint where a prompt ends): less costs little to run again.
106const TURN_CHECKPOINT_MIN_TOKENS: u64 = 128;
107/// The turn checkpoints remembered so that a later turn's supersedes them.
108/// One forgotten stays in the cache until evicted.
109const TURN_CHECKPOINTS_TRACKED: usize = 64;
110
111const PREFILL_RATE_MIN_TOKENS: u32 = 256;
112
113
114/// The daemon's cache classes (a bitset the engine maps its entries onto):
115/// a resident prefix, and one that has started two lanes or more.
116const SHARED_PREFIX_CLASS: u64 = 2;
117/// The prefix cache's entries held in host memory (llama.cpp's exports): they
118/// free nothing in the pool, and are named only when the machine is short of
119/// memory.
120const HOST_EXPORT_CLASS: u64 = 4;
121
122#[derive(Clone, Copy, Debug, PartialEq)]
123struct TickShape {
124    open_secs: f64,
125    full_secs: f64,
126    decode_floor: u16,
127    beside_prefill: u16,
128}
129
130const NATIVE_TICKS: TickShape =
131    TickShape { open_secs: OPEN_TICK_SECS, full_secs: f64::INFINITY, decode_floor: TICK_DECODE, beside_prefill: TICK_DECODE_BESIDE_PREFILL };
132const ROUND_TICKS: TickShape = TickShape { open_secs: 0.05, full_secs: 0.25, decode_floor: 4, beside_prefill: 1 };
133
134fn tick_target(shape: &TickShape, configured: f64, lanes: usize, max_lanes: usize, finishing: bool) -> f64 {
135    if lanes < max_lanes || finishing {
136        configured.min(shape.open_secs)
137    } else {
138        configured.min(shape.full_secs)
139    }
140}
141
142pub struct GenerateJob {
143    pub session: u64,
144    pub stream: Vec<u32>,
145    /// Where the conversation's input ends in `stream`, before the
146    /// generation prompt the daemon appended to it (0: none appended).
147    pub input_end: u64,
148    pub params: GenParams,
149    pub max_tokens: u32,
150    pub events: Sender<LaneEvent>,
151    pub class: u8,
152    pub batch_invariant: bool,
153    pub carry: Option<LaneCarry>,
154    pub queued_at: u64,
155    pub completion: bool,
156    pub deadline: Option<std::time::Instant>,
157    pub media: Vec<(u64, std::path::PathBuf, u32)>,
158    pub grammar_handle: u32,
159    pub extras: GenExtras,
160    pub trace: tracing::Span,
161}
162
163#[derive(Debug, Clone, Default)]
164pub struct GenExtras {
165    pub presence_penalty: f32,
166    pub frequency_penalty: f32,
167    pub repeat_penalty: f32,
168    pub grammar_handle: u32,
169    pub logit_bias_handle: u32,
170    pub want_logprobs: bool,
171    pub top_logprobs: u8,
172    pub wants_deltas: bool,
173    pub ephemeral: bool,
174    pub open_channel: u32,
175    pub thinking: Option<bool>,
176    pub ignore_eos: bool,
177    pub seed_drawn: bool,
178    pub tool_schemas: Option<std::sync::Arc<crate::codec::ToolSchemas>>,
179}
180
181pub struct LaneCarry {
182    seg: Box<dyn Channelizer>,
183    utf8: Utf8Stream,
184    disp_seg: Box<dyn Channelizer>,
185    disp_utf8: Utf8Stream,
186    tool_buf: String,
187    produced_before: u32,
188    spec_before: (u64, u64),
189    warm_prefix_first: Option<u64>,
190}
191
192#[derive(Debug, Clone)]
193pub struct TokenLogprob {
194    pub token: u32,
195    pub logprob: f32,
196    pub top: Vec<(u32, f32)>,
197}
198
199pub enum LaneEvent {
200    Committed(CommittedEvent),
201    Provisional { channel: u32, text: String, produced: u32 },
202    TokenLogprobs(Vec<TokenLogprob>),
203    CompletionTokens(Vec<u32>),
204    Expired,
205    Done {
206        tokens_generated: u32,
207        finish: u32,
208        warm_prefix: u64,
209        spec: (u64, u64),
210    },
211    Failed(String),
212}
213
214fn provisional_ingest(
215    codec: &Arc<dyn TextCodec + Send + Sync>,
216    bus: &SessionBus,
217    lane: &mut Lane,
218    tokens: &[u32],
219) {
220    if tokens.is_empty() {
221        return;
222    }
223    lane.disp_tokens += tokens.len() as u32;
224    let produced = lane.disp_tokens;
225    for run in lane.disp_seg.split(tokens) {
226        let mut text = String::new();
227        for &t in &run.text {
228            text.push_str(&lane.disp_utf8.push(&codec.token_bytes(t)));
229        }
230        if text.is_empty() {
231            continue;
232        }
233        let channel = run.channel;
234        if channel == crate::wal::channel::TEXT || channel == crate::wal::channel::REASONING {
235            bus.publish_provisional(lane.session, channel, &text);
236            let _ = lane.job.events.send(LaneEvent::Provisional { channel, text, produced });
237        }
238    }
239}
240
241fn provisional_flush(bus: &SessionBus, lane: &mut Lane) {
242    let text = lane.disp_utf8.flush();
243    if text.is_empty() {
244        return;
245    }
246    let channel = lane.disp_seg.current_channel();
247    if channel == crate::wal::channel::TEXT || channel == crate::wal::channel::REASONING {
248        bus.publish_provisional(lane.session, channel, &text);
249        let produced = lane.disp_tokens;
250        let _ = lane.job.events.send(LaneEvent::Provisional { channel, text, produced });
251    }
252}
253
254pub enum SchedCmd {
255    Submit(Box<GenerateJob>),
256    Sync(linkw::StateSyncReqMsg, Sender<Result<linkw::StateSyncOkMsg, String>>),
257    SyncAs(u8, Box<SchedCmd>),
258    SyncStream(
259        linkw::StateSyncReqMsg,
260        Sender<linkw::TranscribeSegmentMsg>,
261        Sender<Result<linkw::StateSyncOkMsg, String>>,
262    ),
263    Pin(crate::pins::PinReq),
264    PinHeld(Vec<crate::pins::PinKey>, Sender<u64>),
265    PinUnanchor(u64),
266    Shutdown,
267}
268
269pub struct SchedulerHandle {
270    pub tx: Sender<SchedCmd>,
271    join: std::sync::Mutex<Option<std::thread::JoinHandle<()>>>,
272}
273
274impl SchedulerHandle {
275    pub fn signal_shutdown(&self) {
276        let _ = self.tx.send(SchedCmd::Shutdown);
277    }
278
279    pub fn shutdown(&self) {
280        let mut guard = self.join.lock().expect("join");
281        if let Some(j) = guard.take() {
282            self.signal_shutdown();
283            let _ = j.join();
284        }
285    }
286}
287
288impl Drop for SchedulerHandle {
289    fn drop(&mut self) {
290        self.shutdown();
291    }
292}
293
294enum LanePhase {
295    Prefilling,
296    Decoding,
297    Finishing,
298}
299
300struct Lane {
301    tag: u64,
302    session: u64,
303    job: GenerateJob,
304    phase: LanePhase,
305    prefilled: u64,
306    produced: u32,
307    fin: u32,
308    warm_prefix: u64,
309    utf8: Utf8Stream,
310    seg: Box<dyn Channelizer>,
311    disp_seg: Box<dyn Channelizer>,
312    disp_utf8: Utf8Stream,
313    frame_fed: usize,
314    emit_dead: bool,
315    tool_buf: String,
316    admitted: bool,
317    seed_handle: u64,
318    span: linkw::TokenRefMsg,
319    all_tokens: Vec<u32>,
320    parked: bool,
321    parked_lossy: bool,
322    preempted: bool,
323    starved_ticks: u32,
324    prefill_starved_ticks: u32,
325    /// On a recurrent runtime: where this lane's prefill stops to keep a state for queued
326    /// requests sharing that prefix (0: nowhere), whether this chunk ends there, and the
327    /// longest state kept.
328    checkpoint_at: u64,
329    checkpoint_due: bool,
330    checkpointed: u64,
331    /// Whether the state where the job's input ends has been kept (or tried).
332    turn_kept: bool,
333    produced_before: u32,
334    spec: (u64, u64),
335    disp_tokens: u32,
336    warm_prefix_first: Option<u64>,
337    admitted_at: std::time::Instant,
338    ttft_observed: bool,
339    ttft_frame_ms: Option<u64>,
340    fault_code: Option<u32>,
341}
342
343#[derive(Default)]
344struct ClassQueue {
345    jobs: VecDeque<GenerateJob>,
346}
347
348pub fn kv_bits_label(bits: u64) -> Option<&'static str> {
349    match bits {
350        16 => Some("f16"),
351        8 => Some("q8_0"),
352        4 => Some("q4_0"),
353        84 => Some("K q8_0 / V q4_0"),
354        _ => None,
355    }
356}
357
358#[derive(Debug, Default)]
359pub struct SchedStats {
360    pub pressure_evictions: std::sync::atomic::AtomicU64,
361    pub pressure_bytes_evicted: std::sync::atomic::AtomicU64,
362    pub unservable_cold_admissions: std::sync::atomic::AtomicU64,
363    pub unservable_park_refusals: std::sync::atomic::AtomicU64,
364    pub preemptions: std::sync::atomic::AtomicU64,
365    pub admits_deferred: std::sync::atomic::AtomicU64,
366    pub starvation_grants: std::sync::atomic::AtomicU64,
367    pub os_pressure_events: std::sync::atomic::AtomicU64,
368    pub ticks: std::sync::atomic::AtomicU64,
369    pub lanes_active: std::sync::atomic::AtomicU64,
370    pub queue_depth: [std::sync::atomic::AtomicU64; 4],
371    pub prefill_tokens: std::sync::atomic::AtomicU64,
372    pub decode_tokens: std::sync::atomic::AtomicU64,
373    pub warm_prefix_tokens: std::sync::atomic::AtomicU64,
374    pub cold_admissions: std::sync::atomic::AtomicU64,
375    pub parks_lossless: std::sync::atomic::AtomicU64,
376    pub parks_lossy: std::sync::atomic::AtomicU64,
377    pub resumes: std::sync::atomic::AtomicU64,
378    pub worker_respawns: std::sync::atomic::AtomicU64,
379    pub spec_proposed: std::sync::atomic::AtomicU64,
380    pub spec_accepted: std::sync::atomic::AtomicU64,
381    pub ttft_buckets: [std::sync::atomic::AtomicU64; TTFT_BUCKETS_MS.len() + 1],
382    pub ttft_sum_ms: std::sync::atomic::AtomicU64,
383    pub ttft_count: std::sync::atomic::AtomicU64,
384    pub ttft_last_ms: std::sync::atomic::AtomicU64,
385    pub pool_blocks_used: std::sync::atomic::AtomicU64,
386    pub prefill_budget_live: std::sync::atomic::AtomicU64,
387    pub decode_grant_live: std::sync::atomic::AtomicU64,
388    pub tick_target_live_ms: std::sync::atomic::AtomicU64,
389    pub pool_blocks_total: std::sync::atomic::AtomicU64,
390    pub kv_bits: std::sync::atomic::AtomicU64,
391    pub pins_held: std::sync::atomic::AtomicU64,
392    pub pinned_bytes: std::sync::atomic::AtomicU64,
393    pub pins_yielded: std::sync::atomic::AtomicU64,
394    pub pins_expired: std::sync::atomic::AtomicU64,
395    pub tick_wall_ms_live: std::sync::atomic::AtomicU64,
396    pub prefill_rate_live: std::sync::atomic::AtomicU64,
397    pub decode_rate_live: std::sync::atomic::AtomicU64,
398    pub work_ticks: std::sync::atomic::AtomicU64,
399    tick_seq: std::sync::atomic::AtomicU64,
400    tick_since_ms: std::sync::atomic::AtomicU64,
401    tick_prefill_tokens: std::sync::atomic::AtomicU64,
402    tick_prefill_lanes: std::sync::atomic::AtomicU64,
403    tick_decode_lanes: std::sync::atomic::AtomicU64,
404    tick_admits: std::sync::atomic::AtomicU64,
405    tick_retires: std::sync::atomic::AtomicU64,
406    tick_target_ms: std::sync::atomic::AtomicU64,
407}
408
409#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
410pub struct TickInFlight {
411    pub since_ms: u64,
412    pub prefill_tokens: u64,
413    pub prefill_lanes: u64,
414    pub decode_lanes: u64,
415    pub admits: u64,
416    pub retires: u64,
417    pub target_ms: u64,
418}
419
420impl SchedStats {
421    pub fn set_tick_in_flight(&self, t: TickInFlight) {
422        use std::sync::atomic::Ordering::SeqCst;
423        self.tick_seq.fetch_add(1, SeqCst);
424        self.tick_since_ms.store(t.since_ms, SeqCst);
425        self.tick_prefill_tokens.store(t.prefill_tokens, SeqCst);
426        self.tick_prefill_lanes.store(t.prefill_lanes, SeqCst);
427        self.tick_decode_lanes.store(t.decode_lanes, SeqCst);
428        self.tick_admits.store(t.admits, SeqCst);
429        self.tick_retires.store(t.retires, SeqCst);
430        self.tick_target_ms.store(t.target_ms, SeqCst);
431        self.tick_seq.fetch_add(1, SeqCst);
432    }
433
434    pub fn clear_tick_in_flight(&self) {
435        self.set_tick_in_flight(TickInFlight::default());
436    }
437
438    pub fn tick_in_flight(&self) -> Option<TickInFlight> {
439        use std::sync::atomic::Ordering::SeqCst;
440        loop {
441            let s1 = self.tick_seq.load(SeqCst);
442            if s1 & 1 == 1 {
443                std::hint::spin_loop();
444                continue;
445            }
446            let t = TickInFlight {
447                since_ms: self.tick_since_ms.load(SeqCst),
448                prefill_tokens: self.tick_prefill_tokens.load(SeqCst),
449                prefill_lanes: self.tick_prefill_lanes.load(SeqCst),
450                decode_lanes: self.tick_decode_lanes.load(SeqCst),
451                admits: self.tick_admits.load(SeqCst),
452                retires: self.tick_retires.load(SeqCst),
453                target_ms: self.tick_target_ms.load(SeqCst),
454            };
455            if self.tick_seq.load(SeqCst) == s1 {
456                return (t.since_ms != 0).then_some(t);
457            }
458        }
459    }
460}
461
462pub fn uptime_ms() -> u64 {
463    static START: std::sync::OnceLock<std::time::Instant> = std::sync::OnceLock::new();
464    START.get_or_init(std::time::Instant::now).elapsed().as_millis() as u64
465}
466
467fn per_sec(tokens: u64, secs: f64) -> u64 {
468    if tokens == 0 || secs <= 0.0 {
469        0
470    } else {
471        (tokens as f64 / secs) as u64
472    }
473}
474
475pub const TTFT_BUCKETS_MS: [u64; 13] =
476    [5, 10, 25, 50, 100, 250, 500, 1000, 2500, 5000, 10_000, 20_000, 40_000];
477
478impl SchedStats {
479    pub fn observe_ttft(&self, ms: u64) {
480        let idx = TTFT_BUCKETS_MS
481            .iter()
482            .position(|&b| ms <= b)
483            .unwrap_or(TTFT_BUCKETS_MS.len());
484        self.ttft_buckets[idx].fetch_add(1, Ordering::Relaxed);
485        self.ttft_sum_ms.fetch_add(ms, Ordering::Relaxed);
486        self.ttft_count.fetch_add(1, Ordering::Relaxed);
487        self.ttft_last_ms.store(ms, Ordering::Relaxed);
488    }
489}
490
491#[derive(Debug, Clone, PartialEq, Eq)]
492pub struct SpeculationGrant {
493    pub strategy_id: String,
494    pub slot: u32,
495    pub class: u8,
496    pub cert_id: u32,
497    pub identity: u64,
498    pub modes: u8,
499}
500
501fn proposal_config(id: &str, opts: &crate::SpeculationOptions) -> String {
502    fn env_depth(key: &str) -> Option<u32> {
503        std::env::var(key)
504            .ok()
505            .and_then(|v| v.trim().parse::<u32>().ok())
506            .filter(|n| (1..=15).contains(n))
507    }
508    fn env_bool(key: &str, default: bool) -> bool {
509        std::env::var(key)
510            .ok()
511            .map(|v| v.trim().parse::<i32>().unwrap_or(i32::from(default)) != 0)
512            .unwrap_or(default)
513    }
514    let strategy_depth = || match id {
515        "mtp-head" => env_depth("SUPERFLUID_MTP_DRAFTS").unwrap_or(0),
516        "eagle3" => env_depth("SUPERFLUID_EAGLE_DRAFTS").unwrap_or(2),
517        "dflash" | "dspark" => env_depth("SUPERFLUID_DFLASH_DRAFTS").unwrap_or(0),
518        _ => 4,
519    };
520    let max_draft = opts.max_draft.unwrap_or_else(strategy_depth);
521    let adaptive = opts
522        .adaptive
523        .unwrap_or_else(|| env_bool("SUPERFLUID_SPEC_ADAPTIVE", true));
524    let min_yield = opts.min_yield.unwrap_or_else(|| {
525        std::env::var("SUPERFLUID_SPEC_MIN_YIELD")
526            .ok()
527            .and_then(|v| v.trim().parse::<f64>().ok())
528            .filter(|v| v.is_finite() && *v >= 0.0)
529            .unwrap_or(0.75)
530    });
531    let yield_rounds = opts.yield_rounds.unwrap_or_else(|| {
532        std::env::var("SUPERFLUID_SPEC_LANE_YIELD_ROUNDS")
533            .ok()
534            .and_then(|v| v.trim().parse::<u32>().ok())
535            .filter(|n| *n > 0)
536            .unwrap_or(24)
537    });
538    let throughput_gate = opts
539        .throughput_gate
540        .unwrap_or_else(|| env_bool("SUPERFLUID_SPEC_GATE", true));
541    let gate_probe_tokens = opts.gate_probe_tokens.unwrap_or_else(|| {
542        std::env::var("SUPERFLUID_SPEC_GATE_PROBE")
543            .ok()
544            .and_then(|v| v.trim().parse::<u32>().ok())
545            .unwrap_or(16)
546    });
547    let min_speedup = opts.min_speedup.unwrap_or_else(|| {
548        std::env::var("SUPERFLUID_SPEC_MIN_SPEEDUP")
549            .ok()
550            .and_then(|v| v.trim().parse::<f64>().ok())
551            .filter(|v| v.is_finite() && *v >= 1.0)
552            .unwrap_or(1.08)
553    });
554    let env_u32 = |key: &str, default: u32| {
555        std::env::var(key)
556            .ok()
557            .and_then(|v| v.trim().parse::<u32>().ok())
558            .unwrap_or(default)
559    };
560    let gate_reprobe = opts
561        .gate_reprobe
562        .unwrap_or_else(|| env_u32("SUPERFLUID_SPEC_GATE_REPROBE", 32));
563    let gate_reprobe_max = opts
564        .gate_reprobe_max
565        .unwrap_or_else(|| env_u32("SUPERFLUID_SPEC_GATE_REPROBE_MAX", 1024));
566    let bitexact = env_bool("BASERT_SPEC_BITEXACT", false);
567    let mut cfg = format!(
568        "max_draft={max_draft},adaptive={},min_yield={min_yield:.6},yield_rounds={yield_rounds},throughput_gate={},gate_probe_tokens={gate_probe_tokens},min_speedup={min_speedup:.6},gate_reprobe={gate_reprobe},gate_reprobe_max={gate_reprobe_max},bitexact={}",
569        u8::from(adaptive),
570        u8::from(throughput_gate),
571        u8::from(bitexact),
572    );
573    if id == "dspark" {
574        let confidence = opts.dspark_confidence.or_else(|| {
575            std::env::var("SUPERFLUID_DSPARK_CONFIDENCE")
576                .ok()
577                .and_then(|v| v.trim().parse::<f32>().ok())
578                .filter(|t| t.is_finite() && *t > 0.0)
579        });
580        if let Some(t) = confidence {
581            cfg.push_str(&format!(",confidence={t:.6}"));
582        }
583    }
584    cfg
585}
586
587fn prefill_tick_cap(secs_per_token: f64, lone: bool) -> u32 {
588    let secs = if lone { MAX_LONE_PREFILL_TICK_SECS } else { MAX_PREFILL_TICK_SECS };
589    (secs / secs_per_token).clamp(512.0, PREFILL_BUDGET as f64) as u32
590}
591
592/// Whole 512-token chunks; whole 128-token steps on a model where one chunk passes the
593/// half second.
594/// How long an open tick's prefill may run: longer on an engine that yields at its next
595/// step, unless a latency-sensitive request is around.
596fn open_prefill_tick_target(yields: bool, latency: bool) -> f64 {
597    if yields && !latency {
598        OPEN_PREFILL_TICK_SECS_YIELDING
599    } else {
600        OPEN_PREFILL_TICK_SECS
601    }
602}
603
604fn open_prefill_tick_cap(secs_per_token: f64, target_secs: f64) -> u32 {
605    let tokens = (target_secs / secs_per_token).clamp(128.0, OPEN_PREFILL_TICK_MAX as f64) as u32;
606    let step = if tokens < 512 { 128 } else { 512 };
607    tokens / step * step
608}
609
610/// Why a shared-prefix checkpoint was not kept: the engine forks no lane's
611/// sequence (none will be), or this one failed.
612enum CheckpointError {
613    Refused(String),
614    Failed(String),
615}
616
617impl CheckpointError {
618    /// Only an engine that says it cannot fork refuses for good; anything
619    /// else (a lane that faulted, a full sequence table) fails this one.
620    fn of(e: superfluid_agent::AgentError) -> CheckpointError {
621        if matches!(e, superfluid_agent::AgentError::Rejected(s) if s == superfluid_abi::Status::Unsupported as i32) {
622            CheckpointError::Refused(e.to_string())
623        } else {
624            CheckpointError::Failed(e.to_string())
625        }
626    }
627}
628
629struct PrefillBudget {
630    left: u64,
631    tail_capped: bool,
632    held_below: Option<u8>,
633}
634
635impl PrefillBudget {
636    const TAIL_FLOOR: u64 = 256;
637
638    fn new(budget: u64) -> Self {
639        PrefillBudget { left: budget, tail_capped: false, held_below: None }
640    }
641
642    fn offer(&self, want: u64) -> u64 {
643        want.min(self.left)
644    }
645
646    fn offer_to(&self, want: u64, class: u8) -> u64 {
647        if self.held_below.is_some_and(|c| class > c) {
648            return 0;
649        }
650        self.offer(want)
651    }
652
653    /// A chat or an inline completion finished its prompt this tick: nothing
654    /// of a lower class prefills beside it, so its first token is not held
655    /// behind someone else's prompt.
656    fn hold_below(&mut self, class: u8) {
657        self.held_below = Some(self.held_below.map_or(class, |c| c.min(class)));
658    }
659
660    fn take(&mut self, chunk: u64, completes: bool) {
661        self.left -= chunk;
662        if completes && !self.tail_capped {
663            self.tail_capped = true;
664            self.left = self.left.min(chunk.max(Self::TAIL_FLOOR));
665        }
666    }
667}
668
669/// `plan` with only its first `keep` admissions, and none of the work of the
670/// lanes it drops.
671fn with_first_admits(plan: &linkw::TickPlanMsg, keep: usize) -> linkw::TickPlanMsg {
672    let mut p = plan.clone();
673    let dropped: std::collections::HashSet<u64> = p.admits.split_off(keep.min(p.admits.len())).iter().map(|a| a.lane_tag).collect();
674    p.prefills.retain(|x| !dropped.contains(&x.lane_tag));
675    p.decodes.retain(|x| !dropped.contains(&x.lane_tag));
676    p.commits.retain(|x| !dropped.contains(&x.lane_tag));
677    p
678}
679
680/// A lane's starved-tick count after a tick: `waiting` (prefill left, none
681/// this tick) counts only when the youngest lane served was admitted after
682/// it (lane tags rise with admission).
683fn starved_after(prev: u32, waiting: bool, youngest_served: Option<u64>, tag: u64) -> u32 {
684    if waiting && youngest_served.is_some_and(|y| y > tag) {
685        prev.saturating_add(1)
686    } else {
687        prev
688    }
689}
690
691fn prefill_rank(class: u8, remaining: u64, tag: u64, starved_ticks: u32, bound: u32) -> (bool, u8, u64, u64) {
692    let starving = remaining > 0 && starved_ticks >= bound;
693    (!starving, class, remaining, tag)
694}
695
696#[cfg(test)]
697mod prefill_order_tests {
698    use super::{prefill_rank, prefill_tick_cap, starved_after, PrefillBudget, PREFILL_BUDGET};
699
700    #[test]
701    fn prompts_that_arrived_together_prefill_one_after_another() {
702        // Eight equal prompts, one tick's budget each time: the shortest
703        // left goes first, and nobody is passed over by a younger lane, so
704        // they finish in order instead of in turns.
705        const BOUND: u32 = 8;
706        let mut lanes: Vec<(u64, u64, u32)> = (0..8).map(|t| (t, 7680, 0)).collect();
707        let mut finished_at = Vec::new();
708        for tick in 0..400u32 {
709            let mut order: Vec<usize> = (0..lanes.len()).filter(|&i| lanes[i].1 > 0).collect();
710            if order.is_empty() {
711                break;
712            }
713            order.sort_by_key(|&i| prefill_rank(3, lanes[i].1, lanes[i].0, lanes[i].2, BOUND));
714            let served = order[0];
715            lanes[served].1 -= 512;
716            lanes[served].2 = 0;
717            if lanes[served].1 == 0 {
718                finished_at.push(tick + 1);
719            }
720            let youngest = Some(lanes[served].0);
721            for (i, l) in lanes.iter_mut().enumerate() {
722                l.2 = starved_after(l.2, i != served && l.1 > 0, youngest, l.0);
723            }
724        }
725        assert_eq!(finished_at, (1..=8).map(|k| 15 * k).collect::<Vec<_>>(), "fifteen ticks each, in turn");
726    }
727
728    #[test]
729    fn only_an_engine_that_cannot_fork_refuses_checkpoints_for_good() {
730        use super::CheckpointError;
731        use superfluid_abi::Status;
732        use superfluid_agent::AgentError;
733        let refused = |e| matches!(CheckpointError::of(e), CheckpointError::Refused(_));
734        assert!(refused(AgentError::Rejected(Status::Unsupported as i32)));
735        assert!(!refused(AgentError::Rejected(Status::UnknownHandle as i32)), "a lane that faulted in its tick");
736        assert!(!refused(AgentError::Rejected(Status::NeedsReplan as i32)), "a full sequence table");
737        assert!(!refused(AgentError::Closed), "a worker that went away");
738    }
739
740    #[test]
741    fn an_open_tick_prefills_whole_chunks_for_about_half_a_second() {
742        use super::open_prefill_tick_cap;
743        assert_eq!(open_prefill_tick_cap(1.0 / 1200.0, super::OPEN_PREFILL_TICK_SECS), 512, "600 tokens round down to one chunk");
744        assert_eq!(open_prefill_tick_cap(1.0 / 5000.0, super::OPEN_PREFILL_TICK_SECS), 1024, "2500 tokens, at most two chunks");
745        assert_eq!(open_prefill_tick_cap(1.0, super::OPEN_PREFILL_TICK_SECS), 128, "never below one step");
746        assert_eq!(open_prefill_tick_cap(1.0 / 330.0, super::OPEN_PREFILL_TICK_SECS), 128, "a 27B: 165 tokens round down to one step");
747        assert_eq!(open_prefill_tick_cap(1.0 / 700.0, super::OPEN_PREFILL_TICK_SECS), 256, "350 tokens, two steps");
748        assert_eq!(open_prefill_tick_cap(1e-6, super::OPEN_PREFILL_TICK_SECS), 1024, "never above two chunks, whatever an early tick suggests");
749        assert_eq!(
750            open_prefill_tick_cap(1.0 / 330.0, super::OPEN_PREFILL_TICK_SECS_YIELDING),
751            512,
752            "an engine that yields mid-tick: a 27B's open tick is a whole chunk"
753        );
754        use super::open_prefill_tick_target as target;
755        assert_eq!(target(true, false), super::OPEN_PREFILL_TICK_SECS_YIELDING, "agents alone, on an engine that yields");
756        assert_eq!(target(true, true), super::OPEN_PREFILL_TICK_SECS, "a chat around: its next tokens wait out each tick");
757        assert_eq!(target(false, false), super::OPEN_PREFILL_TICK_SECS, "an engine that runs a tick whole");
758    }
759
760    #[test]
761    fn lone_lane_plans_longer_prefill_ticks() {
762        let spt = 0.0078125;
763        assert_eq!(prefill_tick_cap(spt, false), 1024);
764        assert_eq!(prefill_tick_cap(spt, true), 2048);
765        assert_eq!(prefill_tick_cap(1.0, true), 512);
766        assert_eq!(prefill_tick_cap(1e-6, false), PREFILL_BUDGET);
767        assert_eq!(prefill_tick_cap(1e-6, true), PREFILL_BUDGET);
768    }
769
770    fn grants(budget: u64, remaining: &[u64]) -> Vec<u64> {
771        let mut b = PrefillBudget::new(budget);
772        remaining
773            .iter()
774            .map(|&r| {
775                let c = b.offer(r);
776                if c > 0 {
777                    b.take(c, c == r);
778                }
779                c
780            })
781            .collect()
782    }
783
784    #[test]
785    fn the_tail_after_the_first_completion_is_capped_in_aggregate() {
786        let g = grants(4096, &[500, 4000, 4000, 4000, 4000, 4000, 4000, 4000]);
787        assert_eq!(g, vec![500, 500, 0, 0, 0, 0, 0, 0]);
788        let g = grants(4096, &[40, 4000, 4000]);
789        assert_eq!(g, vec![40, 256, 0]);
790        let g = grants(4096, &[1000, 300, 200, 4000]);
791        assert_eq!(g, vec![1000, 300, 200, 500]);
792        assert_eq!(grants(4096, &[3000, 3000]), vec![3000, 1096]);
793    }
794
795    #[test]
796    fn nothing_of_a_lower_class_prefills_beside_a_chat_that_finishes_its_prompt() {
797        let mut b = PrefillBudget::new(4096);
798        assert_eq!(b.offer_to(22, 0), 22);
799        b.take(22, true);
800        b.hold_below(0);
801        assert_eq!(b.offer_to(4000, 3), 0, "an agent waits a tick");
802        assert_eq!(b.offer_to(100, 0), 100, "another chat does not");
803        let mut b = PrefillBudget::new(4096);
804        b.take(500, false);
805        assert_eq!(b.offer_to(4000, 3), 3596, "a chat still prefilling holds nothing back");
806    }
807
808    #[test]
809    fn a_long_prefill_progresses_under_a_stream_of_short_prompts() {
810        const BUDGET: u64 = 512;
811        const BOUND: u32 = 3;
812        let mut lanes: Vec<(u64, u64, u32)> = vec![(0, 8 * BUDGET, 0)];
813        let mut next_tag = 1;
814        let mut long_ticks = Vec::new();
815        for tick in 0..64u32 {
816            lanes.retain(|l| l.0 == 0 || l.1 > 0);
817            while lanes.len() < 3 {
818                lanes.push((next_tag, BUDGET, 0));
819                next_tag += 1;
820            }
821            let mut order: Vec<usize> = (0..lanes.len()).collect();
822            order.sort_by_key(|&i| prefill_rank(0, lanes[i].1, lanes[i].0, lanes[i].2, BOUND));
823            let mut budget = PrefillBudget::new(BUDGET);
824            let mut served = Vec::new();
825            for i in order {
826                let l = &mut lanes[i];
827                if l.1 == 0 {
828                    continue;
829                }
830                let chunk = budget.offer(l.1);
831                l.1 -= chunk;
832                if chunk > 0 {
833                    l.2 = 0;
834                    served.push(l.0);
835                    budget.take(chunk, l.1 == 0);
836                }
837                if l.0 == 0 && chunk > 0 {
838                    long_ticks.push(tick);
839                }
840            }
841            let youngest = served.iter().copied().max();
842            for l in lanes.iter_mut() {
843                l.2 = starved_after(l.2, l.1 > 0 && !served.contains(&l.0), youngest, l.0);
844            }
845            assert!(lanes[0].2 <= BOUND, "tick {tick}: the long lane waited past the bound");
846            if lanes[0].1 == 0 {
847                break;
848            }
849        }
850        assert_eq!(lanes[0], (0, 0, 0), "the long prompt finished prefilling: {long_ticks:?}");
851        assert!(long_ticks.windows(2).all(|w| w[1] - w[0] <= BOUND + 1), "bounded gaps: {long_ticks:?}");
852        assert!(prefill_rank(0, 8 * BUDGET, 0, 0, BOUND) > prefill_rank(0, BUDGET, 9, 0, BOUND));
853        assert!(prefill_rank(5, 8 * BUDGET, 0, BOUND, BOUND) < prefill_rank(0, 1, 9, 0, BOUND));
854        assert!(prefill_rank(5, 0, 0, BOUND, BOUND) > prefill_rank(0, 1, 9, 0, BOUND));
855    }
856}
857
858#[cfg(test)]
859mod tick_target_tests {
860    use super::{tick_target, NATIVE_TICKS, OPEN_TICK_SECS, ROUND_TICKS};
861
862    #[test]
863    fn a_tick_is_short_while_a_joiner_could_be_admitted_next() {
864        assert_eq!(tick_target(&NATIVE_TICKS, 2.0, 1, 2, false), OPEN_TICK_SECS, "a slot is free");
865        assert_eq!(tick_target(&NATIVE_TICKS, 2.0, 2, 2, false), 2.0, "full and staying full: nothing could admit a joiner");
866        assert_eq!(tick_target(&ROUND_TICKS, 2.0, 1, 2, false), 0.05, "a round-granular runtime's open tick is a few rounds");
867        assert_eq!(tick_target(&ROUND_TICKS, 2.0, 2, 2, false), 0.25, "and its full one is never long");
868        assert_eq!(
869            tick_target(&NATIVE_TICKS, 2.0, 2, 2, true),
870            OPEN_TICK_SECS,
871            "a finishing lane frees a slot, and its client waits on the retire"
872        );
873        assert_eq!(tick_target(&NATIVE_TICKS, 0.1, 2, 2, true), 0.1, "a configured target under the open one stands");
874    }
875}
876
877#[cfg(test)]
878mod speculation_config_tests {
879    use super::proposal_config;
880    use crate::SpeculationOptions;
881
882    #[test]
883    fn cli_policy_is_serialized_for_the_engine_and_fingerprint() {
884        let cfg = proposal_config(
885            "dspark",
886            &SpeculationOptions {
887                max_draft: Some(2),
888                adaptive: Some(false),
889                min_yield: Some(0.5),
890                yield_rounds: Some(12),
891                throughput_gate: Some(false),
892                gate_probe_tokens: Some(8),
893                min_speedup: Some(1.12),
894                gate_reprobe: Some(0),
895                gate_reprobe_max: Some(64),
896                dspark_confidence: Some(0.42),
897            },
898        );
899        assert!(cfg.starts_with(
900            "max_draft=2,adaptive=0,min_yield=0.500000,yield_rounds=12,throughput_gate=0,gate_probe_tokens=8,min_speedup=1.120000,gate_reprobe=0,gate_reprobe_max=64,"
901        ));
902        assert!(cfg.ends_with(",confidence=0.420000"));
903    }
904
905    #[test]
906    fn one_draft_override_applies_to_mtp() {
907        let cfg = proposal_config(
908            "mtp-head",
909            &SpeculationOptions {
910                max_draft: Some(1),
911                ..SpeculationOptions::default()
912            },
913        );
914        assert!(cfg.starts_with("max_draft=1,"));
915        assert!(!cfg.contains("confidence="));
916    }
917}
918
919#[derive(Default)]
920pub struct Speculation {
921    reg: Mutex<Option<linkw::StrategyRegisterMsg>>,
922    grant: Mutex<Option<SpeculationGrant>>,
923}
924
925impl Speculation {
926    pub fn register(
927        this: &Arc<Speculation>,
928        host: &mut EngineHost,
929        strategy: &str,
930        opts: &crate::SpeculationOptions,
931    ) -> Result<(), DaemonError> {
932        let (id, draft_path) = match strategy.split_once(':') {
933            Some(("draft-model", path)) => ("draft-model", Some(path.to_string())),
934            Some(("dflash", path)) => ("dflash", Some(path.to_string())),
935            Some(("dspark", path)) => ("dspark", Some(path.to_string())),
936            Some(("eagle3", path)) => ("eagle3", Some(path.to_string())),
937            _ if strategy == "prompt-lookup" => ("prompt-lookup", None),
938            _ if strategy == "mtp-head" => ("mtp-head", None),
939            _ => {
940                return Err(DaemonError::Config(
941                    "unknown speculation strategy (prompt-lookup | mtp-head | dflash / dspark / eagle3 / draft-model followed by ':' and a model id or .base path)",
942                ))
943            }
944        };
945        let draft_identity = match draft_path.as_ref() {
946            Some(p) => Some(draft_weights_identity(p)?),
947            None => None,
948        };
949        let impl_hash = {
950            let mut h = [0u8; 32];
951            let tag = if id == "dflash" {
952                "dflash/1"
953            } else if id == "dspark" {
954                "dspark/1"
955            } else if id == "eagle3" {
956                "eagle3/1"
957            } else if draft_path.is_some() {
958                "draft-model/1"
959            } else if id == "mtp-head" {
960                "mtp-head/1"
961            } else {
962                "prompt-lookup/1"
963            };
964            h[..8].copy_from_slice(&superfluid_fingerprint::fnv1a64(tag.as_bytes(), 0xCBF2_9CE4_8422_2325).to_le_bytes());
965            h
966        };
967        let proposal_cfg = proposal_config(id, opts);
968        let config_hash = {
969            let mut h = [0u8; 32];
970            let cfg = match (&draft_path, &draft_identity) {
971                (Some(p), Some((fnv, size))) => {
972                    format!("draft={p},weights={fnv:016x},bytes={size},{proposal_cfg}")
973                }
974                _ => format!("n_max=4,{proposal_cfg}"),
975            };
976            h[..8].copy_from_slice(&superfluid_fingerprint::fnv1a64(cfg.as_bytes(), 0xCBF2_9CE4_8422_2325).to_le_bytes());
977            h
978        };
979        let artifacts = match (&draft_path, &draft_identity) {
980            (Some(path), Some((fnv, size))) => {
981                let mut content_hash = [0u8; 32];
982                content_hash[..8].copy_from_slice(&fnv.to_le_bytes());
983                vec![linkw::ArtifactMsg {
984                    role: superfluid_abi::artifact_role::DRAFT_WEIGHTS,
985                    content_hash,
986                    byte_size: *size,
987                    load_path: path.clone(),
988                }]
989            }
990            _ => vec![],
991        };
992        let reg = linkw::StrategyRegisterMsg {
993            strategy_id: id.to_string(),
994            impl_version: "1.0.0".into(),
995            impl_hash,
996            config_hash,
997            artifacts,
998            taps: vec![],
999            capabilities: vec![linkw::CapabilityReqMsg {
1000                kind_id: superfluid_abi::strategy_cap::PROPOSAL_LINEAR,
1001                params: proposal_cfg.into_bytes(),
1002            }],
1003            target_archs: vec![],
1004            kernel_caps_required: 0,
1005            est_state_bytes: 0,
1006            claimed_exactness: superfluid_abi::exactness::SEED_PATH_INVARIANT,
1007            rng_contract_version: 1,
1008        };
1009        *this.reg.lock().expect("spec reg") = Some(reg);
1010        this.reregister(host)
1011    }
1012
1013    pub fn reregister(&self, host: &mut EngineHost) -> Result<(), DaemonError> {
1014        let Some(reg) = self.reg.lock().expect("spec reg").clone() else {
1015            return Ok(());
1016        };
1017        let ok = host.client().register_strategy(reg.clone()).map_err(|e| {
1018            DaemonError::Generation(format!("speculation strategy {} refused: {e}", reg.strategy_id))
1019        })?;
1020        let best = ok
1021            .certificates
1022            .iter()
1023            .filter(|c| c.sampling_modes & (superfluid_abi::cert_mode::GREEDY | superfluid_abi::cert_mode::GUMBEL) != 0)
1024            .max_by_key(|c| c.exactness);
1025        let identity = {
1026            let mut h = superfluid_fingerprint::RecordHasher::new();
1027            h.field(1, reg.strategy_id.as_bytes())
1028                .field(2, &reg.impl_hash)
1029                .field(3, &reg.config_hash);
1030            h.finish()
1031        };
1032        let modes = ok.certificates.iter().fold(0u8, |m, c| m | c.sampling_modes);
1033        *self.grant.lock().expect("spec grant") = Some(SpeculationGrant {
1034            strategy_id: reg.strategy_id.clone(),
1035            slot: ok.strategy_slot,
1036            class: best.map(|c| c.exactness).unwrap_or(superfluid_abi::exactness::APPROXIMATE),
1037            cert_id: best.map(|c| c.cert_id).unwrap_or(0),
1038            identity,
1039            modes,
1040        });
1041        Ok(())
1042    }
1043
1044    pub fn grant(&self) -> Option<SpeculationGrant> {
1045        self.grant.lock().expect("spec grant").clone()
1046    }
1047
1048    pub fn slot_for(&self, sampled: bool) -> u32 {
1049        let mode = if sampled {
1050            superfluid_abi::cert_mode::GUMBEL
1051        } else {
1052            superfluid_abi::cert_mode::GREEDY
1053        };
1054        self.grant
1055            .lock()
1056            .expect("spec grant")
1057            .as_ref()
1058            .filter(|g| g.modes & mode != 0)
1059            .map(|g| g.slot)
1060            .unwrap_or(0)
1061    }
1062
1063    pub fn identity_for_fingerprint(&self) -> Option<u64> {
1064        self.grant
1065            .lock()
1066            .expect("spec grant")
1067            .as_ref()
1068            .filter(|g| g.class < superfluid_abi::exactness::SEED_PATH_INVARIANT)
1069            .map(|g| g.identity)
1070    }
1071}
1072
1073#[derive(Debug, Clone, PartialEq, Eq)]
1074pub enum ParkRefused {
1075    SpaceUnservable,
1076    Failed(String),
1077}
1078
1079impl std::fmt::Display for ParkRefused {
1080    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1081        match self {
1082            ParkRefused::SpaceUnservable => {
1083                write!(f, "a state space has no ops (hybrid GDN before P2.5); park refused")
1084            }
1085            ParkRefused::Failed(m) => f.write_str(m),
1086        }
1087    }
1088}
1089
1090impl From<String> for ParkRefused {
1091    fn from(m: String) -> Self {
1092        ParkRefused::Failed(m)
1093    }
1094}
1095
1096impl From<&str> for ParkRefused {
1097    fn from(m: &str) -> Self {
1098        ParkRefused::Failed(m.to_string())
1099    }
1100}
1101
1102enum ParkTask {
1103    Write(Box<ParkJob>),
1104    Remove { dir: PathBuf, session: u64 },
1105}
1106
1107struct ParkJob {
1108    dir: PathBuf,
1109    session: u64,
1110    covered: u64,
1111    digest: u64,
1112    encoding: u8,
1113    sealed: Vec<u8>,
1114    spaces: Option<Vec<park::ParkedSpace>>,
1115}
1116
1117pub struct Scheduler {
1118    host: EngineHost,
1119    store: Arc<Mutex<SessionStore>>,
1120    codec: Arc<dyn TextCodec + Send + Sync>,
1121    cancels: Arc<crate::CancelSet>,
1122    active: Arc<Mutex<HashSet<u64>>>,
1123    stats: Arc<SchedStats>,
1124    rx: Receiver<SchedCmd>,
1125    max_lanes: usize,
1126    configured_lanes: usize,
1127    retired: Arc<Mutex<Option<String>>>,
1128    park_dir: Option<PathBuf>,
1129    pressure_high_pct: u8,
1130    pressure_low_pct: u8,
1131    park_lossy: bool,
1132    trace_scope: crate::otlp::TraceScope,
1133    park_tx: Option<std::sync::mpsc::SyncSender<ParkTask>>,
1134    ephemeral_parked: std::collections::HashSet<u64>,
1135    park_removals_pending: Vec<u64>,
1136    park_writer: Option<std::thread::JoinHandle<()>>,
1137    bus: Arc<SessionBus>,
1138
1139    plan_seq: u64,
1140    next_lane: u64,
1141    lanes: Vec<Lane>,
1142    queues: [ClassQueue; 4],
1143    class_lanes: [usize; 4],
1144    tick_decode_budget: u32,
1145    tick_target_secs: f64,
1146    tick_shape: TickShape,
1147    prefill_budget: u32,
1148    prefill_budget_fixed: bool,
1149    decode_round_secs: Option<f64>,
1150    prefill_secs_per_token: Option<f64>,
1151    just_retired: bool,
1152    prefill_secs_per_token_bulk: Option<f64>,
1153    prefill_secs_per_token_last: Option<f64>,
1154    prefill_rate_samples: u32,
1155    /// The engine refused to fork a lane's sequence: no checkpoint where a
1156    /// shared prefix ends is kept, and no request is held for one.
1157    prefix_checkpoints_off: bool,
1158    /// The turn checkpoints kept, oldest first, by the length and digest of the tokens each
1159    /// covers; a turn's checkpoint is dropped once its next turn keeps one.
1160    turn_checkpoints: std::collections::VecDeque<(u64, [u8; 32])>,
1161    tick_decode_budget_fixed: bool,
1162    starvation_ticks: u32,
1163    background_tick_divisor: Arc<std::sync::atomic::AtomicU32>,
1164    pressure_source: Option<Box<dyn crate::pressure::PressureSource + Send>>,
1165    tick_count: u64,
1166    speculation: Arc<Speculation>,
1167    tool_lease_ms: u64,
1168    respawn_streak: u32,
1169    staged_unread: u32,
1170    deferred_sync: VecDeque<(u8, u64, SchedCmd)>,
1171    admit_retry: bool,
1172    pins: crate::pins::Pins,
1173}
1174
1175impl Scheduler {
1176    fn respawn_worker(&mut self) -> Result<(), DaemonError> {
1177        self.host.respawn()?;
1178        self.stats.kv_bits.store(self.host.kv_bits() as u64, Ordering::Relaxed);
1179        Ok(())
1180    }
1181
1182    #[allow(clippy::too_many_arguments)]
1183    pub fn spawn(
1184        host: EngineHost,
1185        store: Arc<Mutex<SessionStore>>,
1186        codec: Arc<dyn TextCodec + Send + Sync>,
1187        cancels: Arc<crate::CancelSet>,
1188        active: Arc<Mutex<HashSet<u64>>>,
1189        stats: Arc<SchedStats>,
1190        background_tick_divisor: Arc<std::sync::atomic::AtomicU32>,
1191        speculation: Arc<Speculation>,
1192        retired: Arc<Mutex<Option<String>>>,
1193        mut opts: crate::DaemonOptions,
1194    ) -> SchedulerHandle {
1195        let (tx, rx) = std::sync::mpsc::channel();
1196        let join = std::thread::Builder::new()
1197            .name("superfluid-scheduler".into())
1198            .spawn(move || {
1199                let bus = store.lock().expect("store").bus();
1200                let host_max_stream = host.max_stream_tokens();
1201                let mut host = host;
1202                let max_lanes = lane_cap(
1203                    opts.max_lanes,
1204                    host.max_seqs(),
1205                    ring_slots(&mut host, superfluid_shm::TOKEN_RING_OUT),
1206                );
1207                let park_refused = host.park_refused();
1208                let lossy_allowed = host.lossy_park_allowed();
1209                let tick_shape = if host.round_granular() { ROUND_TICKS } else { NATIVE_TICKS };
1210                let (park_tx, park_writer) = if opts.park_dir.is_some() && park_refused.is_none() {
1211                    let (tx, rx) = std::sync::mpsc::sync_channel::<ParkTask>(4);
1212                    let park_budget = opts.park_budget_bytes;
1213                    let writer = std::thread::Builder::new()
1214                        .name("superfluid-park-writer".into())
1215                        .spawn(move || {
1216                            while let Ok(task) = rx.recv() {
1217                                let job = match task {
1218                                    ParkTask::Write(job) => job,
1219                                    ParkTask::Remove { dir, session } => {
1220                                        park::remove(&dir, session);
1221                                        continue;
1222                                    }
1223                                };
1224                                let res = match &job.spaces {
1225                                    Some(spaces) => park::write_v3(
1226                                        &job.dir,
1227                                        job.session,
1228                                        job.covered,
1229                                        job.digest,
1230                                        spaces,
1231                                    ),
1232                                    None => park::write(
1233                                        &job.dir,
1234                                        job.session,
1235                                        job.covered,
1236                                        job.digest,
1237                                        job.encoding,
1238                                        &job.sealed,
1239                                    ),
1240                                };
1241                                if let Err(e) = res {
1242                                    eprintln!(
1243                                        "[superfluid] park write failed (session {}): {e}",
1244                                        job.session
1245                                    );
1246                                }
1247                                if park_budget > 0 {
1248                                    let freed = park::sweep(&job.dir, park_budget);
1249                                    if freed > 0 {
1250                                        tracing::debug!(
1251                                            freed_bytes = freed,
1252                                            budget_bytes = park_budget,
1253                                            "park retention sweep"
1254                                        );
1255                                    }
1256                                }
1257                            }
1258                        })
1259                        .expect("spawn park writer");
1260                    (Some(tx), Some(writer))
1261                } else {
1262                    (None, None)
1263                };
1264                Scheduler {
1265                    host,
1266                    store,
1267                    codec,
1268                    cancels,
1269                    active,
1270                    stats,
1271                    rx,
1272                    max_lanes,
1273                    configured_lanes: opts.max_lanes,
1274                    retired,
1275                    park_dir: match park_refused {
1276                        Some(why) if opts.park_dir.is_some() => {
1277                            eprintln!("[superfluid] parking is off for this model: {why}");
1278                            None
1279                        }
1280                        _ => opts.park_dir,
1281                    },
1282                    pressure_high_pct: opts.pressure_high_pct,
1283                    pressure_low_pct: opts.pressure_low_pct,
1284                    park_lossy: opts.park_lossy && lossy_allowed,
1285                    trace_scope: opts.trace_scope.clone(),
1286                    park_tx,
1287                    park_writer,
1288                    ephemeral_parked: std::collections::HashSet::new(),
1289                    park_removals_pending: Vec::new(),
1290                    bus,
1291                    plan_seq: 0,
1292                    next_lane: 1,
1293                    lanes: Vec::new(),
1294                    queues: Default::default(),
1295                    class_lanes: opts.class_lanes,
1296                    tick_target_secs: if opts.tick_target_ms == 0 {
1297                        TARGET_TICK_TOTAL_SECS
1298                    } else {
1299                        opts.tick_target_ms as f64 / 1000.0
1300                    },
1301                    tick_shape,
1302                    tick_decode_budget: if opts.tick_decode_budget == 0 {
1303                        (max_lanes as u32).saturating_mul(TICK_DECODE as u32)
1304                    } else {
1305                        opts.tick_decode_budget
1306                    },
1307                    prefill_budget: if opts.prefill_budget == 0 {
1308                        1024
1309                    } else {
1310                        if opts.prefill_budget > PREFILL_BUDGET {
1311                            eprintln!(
1312                                "superfluid: --prefill-budget {} is above the {PREFILL_BUDGET}-token per-tick \
1313                                 ceiling (the media run-size limit); using {PREFILL_BUDGET}",
1314                                opts.prefill_budget
1315                            );
1316                        }
1317                        opts.prefill_budget.clamp(64, PREFILL_BUDGET)
1318                    },
1319                    prefill_budget_fixed: opts.prefill_budget != 0,
1320                    decode_round_secs: None,
1321                    prefill_secs_per_token: None,
1322                    just_retired: false,
1323                    prefill_secs_per_token_bulk: None,
1324                    prefill_secs_per_token_last: None,
1325                    prefill_rate_samples: 0,
1326                    prefix_checkpoints_off: false,
1327                    turn_checkpoints: std::collections::VecDeque::new(),
1328                    tick_decode_budget_fixed: opts.tick_decode_budget != 0,
1329                    starvation_ticks: opts.agent_starvation_ticks.max(1),
1330                    background_tick_divisor,
1331                    pressure_source: opts.pressure_source.take_source(),
1332                    tick_count: 0,
1333                    speculation,
1334                    tool_lease_ms: opts.tool_lease_ms,
1335                    respawn_streak: 0,
1336                    staged_unread: 0,
1337                    deferred_sync: VecDeque::new(),
1338                    admit_retry: false,
1339                    pins: crate::pins::Pins::new(
1340                        opts.pin_budget_pct,
1341                        host_max_stream * max_lanes as u64,
1342                    ),
1343                }
1344                .run();
1345            })
1346            .expect("spawn scheduler");
1347        SchedulerHandle {
1348            tx,
1349            join: std::sync::Mutex::new(Some(join)),
1350        }
1351    }
1352
1353    fn run(mut self) {
1354        self.run_loop();
1355        drop(self.park_tx.take());
1356        if let Some(w) = self.park_writer.take() {
1357            let _ = w.join();
1358        }
1359        if let Some(dir) = self.park_dir.clone() {
1360            let pending = std::mem::take(&mut self.park_removals_pending);
1361            let tracked = std::mem::take(&mut self.ephemeral_parked);
1362            for session in pending.into_iter().chain(tracked) {
1363                park::remove(&dir, session);
1364            }
1365        }
1366    }
1367
1368    fn run_loop(&mut self) {
1369        loop {
1370            if self.lanes.is_empty() && self.queued() == 0 && self.deferred_sync.is_empty() {
1371                self.retry_park_removals();
1372                let retry = (!self.park_removals_pending.is_empty())
1373                    .then(|| std::time::Duration::from_millis(50));
1374                let wait = match (self.pins.until_next_deadline(), retry) {
1375                    (Some(a), Some(b)) => Some(a.min(b)),
1376                    (a, b) => a.or(b),
1377                };
1378                let cmd = match wait {
1379                    None => self.rx.recv().map_err(|_| ()),
1380                    Some(wait) => match self.rx.recv_timeout(wait) {
1381                        Ok(cmd) => Ok(cmd),
1382                        Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {
1383                            self.pins.sweep(&mut self.host, &self.stats);
1384                            continue;
1385                        }
1386                        Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => Err(()),
1387                    },
1388                };
1389                match cmd {
1390                    Ok(SchedCmd::Shutdown) | Err(()) => return,
1391                    Ok(cmd) => self.dispatch(cmd),
1392                }
1393            }
1394            // Lowered before what is waiting is taken in: anything that
1395            // arrives after raises it for the tick about to run.
1396            self.host.clear_yield();
1397            // A request that arrives between the look and the store raised
1398            // the word itself: looked at again, it is raised again.
1399            let slot = self.host.yield_slot();
1400            let latency = self.latency_sensitive();
1401            crate::runtime::set_latency(&slot, latency);
1402            if !latency && self.cancels.latency_recent(LATENCY_HOLD) {
1403                crate::runtime::set_latency(&slot, true);
1404            }
1405            loop {
1406                match self.rx.try_recv() {
1407                    Ok(SchedCmd::Shutdown) => return,
1408                    Ok(cmd) => self.dispatch(cmd),
1409                    Err(TryRecvError::Empty) => break,
1410                    Err(TryRecvError::Disconnected) => return,
1411                }
1412            }
1413            if !self.host.is_alive() {
1414                if self.respawn_streak < MAX_WORKER_RESPAWNS {
1415                    self.respawn_streak += 1;
1416                    if self.respawn_worker().is_ok() {
1417                        let _ = self.speculation.reregister(&mut self.host);
1418                        self.after_respawn();
1419                    }
1420                }
1421                if !self.host.is_alive() {
1422                    if self.respawn_streak >= MAX_WORKER_RESPAWNS {
1423                        self.give_up_on_worker();
1424                    }
1425                    let why = self
1426                        .retired
1427                        .lock()
1428                        .expect("retired")
1429                        .clone()
1430                        .unwrap_or_else(|| "engine worker unavailable (respawn failed)".into());
1431                    let orphaned: Vec<GenerateJob> =
1432                        self.queues.iter_mut().flat_map(|q| q.jobs.drain(..)).collect();
1433                    for job in orphaned {
1434                        self.finish_failed(&job, &why);
1435                    }
1436                    for (_, _, cmd) in std::mem::take(&mut self.deferred_sync) {
1437                        self.dispatch(cmd);
1438                    }
1439                    continue;
1440                }
1441            }
1442            self.pins.sweep(&mut self.host, &self.stats);
1443            let tick_span = if self.lanes.is_empty() && self.queued() == 0 {
1444                tracing::Span::none()
1445            } else {
1446                tracing::debug_span!(
1447                    parent: None,
1448                    "tick",
1449                    tick_seq = tracing::field::Empty,
1450                    no_tick = tracing::field::Empty,
1451                    lanes = tracing::field::Empty,
1452                    queued = self.queued(),
1453                    model = self.trace_scope.model.as_str(),
1454                    store_ns = self.trace_scope.ns
1455                )
1456            };
1457            let _tick = tick_span.enter();
1458            self.poll_os_pressure();
1459            self.preempt_for_queued();
1460            self.admit_from_queue();
1461            self.run_deferred_syncs();
1462            let ticks_before = self.tick_count;
1463            let tick_result = self.tick_once(&tick_span);
1464            if !tick_span.is_disabled() {
1465                if self.tick_count != ticks_before {
1466                    tick_span.record("tick_seq", self.tick_count);
1467                } else {
1468                    tick_span.record("no_tick", true);
1469                }
1470            }
1471            if tick_result.is_ok() {
1472                self.respawn_streak = 0;
1473            }
1474            if let Err(e) = tick_result {
1475                let admitted: Vec<u64> = self
1476                    .lanes
1477                    .iter()
1478                    .filter(|l| l.admitted)
1479                    .map(|l| l.tag)
1480                    .collect();
1481                if !admitted.is_empty() {
1482                    let mut p = linkw::TickPlanMsg {
1483                        plan_seq: {
1484                            self.plan_seq += 1;
1485                            self.plan_seq
1486                        },
1487                        flags: 0,
1488                        prefill_token_budget: PREFILL_BUDGET,
1489                        max_decode_lanes: self.max_lanes as u32,
1490                        admits: vec![],
1491                        commits: vec![],
1492                        prefills: vec![],
1493                        decodes: vec![],
1494                        retires: vec![],
1495                        shed_policy: linkw::ShedPolicyMsg {
1496                            victim_lanes: vec![],
1497                            evictable_cache_classes: u64::MAX,
1498                            protected_quota_bytes: 0,
1499                            max_evict_bytes: u64::MAX,
1500                        },
1501                    };
1502                    for tag in admitted {
1503                        p.retires.push(linkw::LaneRetireMsg {
1504                            lane_tag: tag,
1505                            publish_to_cache: true,
1506                        });
1507                    }
1508                    self.stats.set_tick_in_flight(TickInFlight {
1509                        since_ms: uptime_ms().max(1),
1510                        retires: p.retires.len() as u64,
1511                        target_ms: self.stats.tick_target_live_ms.load(Ordering::Relaxed),
1512                        ..Default::default()
1513                    });
1514                    let _ = self.host.client().tick(p);
1515                    self.stats.clear_tick_in_flight();
1516                }
1517                let msg = format!("scheduler tick failed: {e}");
1518                let drained: Vec<u64> = self.lanes.iter().map(|l| l.session).collect();
1519                for lane in self.lanes.drain(..) {
1520                    self.active.lock().expect("active set").remove(&lane.session);
1521                    let _ = lane.job.events.send(LaneEvent::Failed(msg.clone()));
1522                }
1523                for session in drained {
1524                    self.drop_ephemeral_artifact(session);
1525                }
1526                self.publish_occupancy();
1527                if worker_is_gone(&e) && self.host.is_process() {
1528                    let mut fresh = false;
1529                    if self.respawn_streak < MAX_WORKER_RESPAWNS {
1530                        self.respawn_streak += 1;
1531                        self.stats.worker_respawns.fetch_add(1, Ordering::Relaxed);
1532                        match self.respawn_worker().and_then(|()| self.speculation.reregister(&mut self.host)) {
1533                            Ok(()) => {
1534                                eprintln!("[superfluid] engine worker died ({e}); respawned a fresh one");
1535                                self.after_respawn();
1536                                fresh = true;
1537                            }
1538                            Err(err) => eprintln!(
1539                                "[superfluid] engine worker died ({e}); respawn FAILED: {err}"
1540                            ),
1541                        }
1542                    }
1543                    if !fresh && self.respawn_streak >= MAX_WORKER_RESPAWNS {
1544                        self.give_up_on_worker();
1545                    }
1546                }
1547            }
1548            self.pins.sweep(&mut self.host, &self.stats);
1549        }
1550    }
1551
1552    fn dispatch(&mut self, cmd: SchedCmd) {
1553        match cmd {
1554            SchedCmd::Submit(job) => self.enqueue(*job),
1555            SchedCmd::Sync(req, reply) => self.run_sync(req, reply),
1556            SchedCmd::SyncStream(req, segs, reply) => self.run_sync_stream(req, segs, reply),
1557            SchedCmd::Pin(req) => self.pins.handle(&mut self.host, &self.stats, self.tick_count, req),
1558            SchedCmd::PinHeld(keys, reply) => {
1559                let _ = reply.send(self.pins.held_max(&self.host, &keys));
1560            }
1561            SchedCmd::PinUnanchor(session) => self.pins.unanchor(session),
1562            SchedCmd::SyncAs(class, inner) => {
1563                self.deferred_sync.push_back((class, self.tick_count, *inner));
1564            }
1565            SchedCmd::Shutdown => {}
1566        }
1567    }
1568
1569    fn sync_must_wait(&self, class: u8) -> bool {
1570        self.lanes.iter().any(|l| l.job.class < class)
1571            || self.queues.iter().take(class as usize).any(|q| !q.jobs.is_empty())
1572    }
1573
1574    fn run_deferred_syncs(&mut self) {
1575        let period = self.starvation_ticks as u64;
1576        let stuck = self.lanes.is_empty() && !self.admit_retry;
1577        let starved = |queued: u64| self.tick_count.saturating_sub(queued) >= period;
1578        let pick = self
1579            .deferred_sync
1580            .iter()
1581            .enumerate()
1582            .filter(|(_, (class, queued, _))| stuck || !self.sync_must_wait(*class) || starved(*queued))
1583            .min_by_key(|&(i, (class, queued, _))| {
1584                let aged = starved(*queued);
1585                (!aged, if aged { 0 } else { self.sync_rank(*class) }, i)
1586            })
1587            .map(|(i, _)| i);
1588        if let Some((_, _, cmd)) = pick.and_then(|i| self.deferred_sync.remove(i)) {
1589            self.dispatch(cmd);
1590        }
1591    }
1592
1593    fn sync_rank(&self, class: u8) -> usize {
1594        class as usize
1595    }
1596
1597    fn run_sync(&mut self, req: linkw::StateSyncReqMsg, reply: Sender<Result<linkw::StateSyncOkMsg, String>>) {
1598        let res = if self.host.is_alive() {
1599            self.host.client().state_sync(req).map_err(|e| e.to_string())
1600        } else {
1601            Err("engine worker unavailable".into())
1602        };
1603        let _ = reply.send(res);
1604    }
1605
1606    fn run_sync_stream(
1607        &mut self,
1608        req: linkw::StateSyncReqMsg,
1609        segs: Sender<linkw::TranscribeSegmentMsg>,
1610        reply: Sender<Result<linkw::StateSyncOkMsg, String>>,
1611    ) {
1612        let res = if self.host.is_alive() {
1613            self.host
1614                .client()
1615                .state_sync_streaming(req, &mut |seg| {
1616                    segs.send(seg).is_ok()
1617                })
1618                .map_err(|e| e.to_string())
1619        } else {
1620            Err("engine worker unavailable".into())
1621        };
1622        let _ = reply.send(res);
1623    }
1624
1625    fn queued(&self) -> usize {
1626        self.queues.iter().map(|q| q.jobs.len()).sum()
1627    }
1628
1629    fn enqueue(&mut self, mut job: GenerateJob) {
1630        let c = (job.class as usize).min(3);
1631        job.queued_at = self.tick_count;
1632        self.queues[c].jobs.push_back(job);
1633    }
1634
1635    fn engine_exclusive(&self, job: &GenerateJob) -> bool {
1636        job.batch_invariant || (!job.media.is_empty() && self.host.has_recurrent_space())
1637    }
1638
1639    fn strategy_slot_for(speculation: &Speculation, job: &GenerateJob) -> u32 {
1640        if job.grammar_handle != 0 || job.extras.logit_bias_handle != 0 || job.extras.want_logprobs {
1641            0
1642        } else {
1643            speculation.slot_for(job.params.temperature != 0.0)
1644        }
1645    }
1646
1647    fn decode_replay_for(strategy_slot: u32, produced_before: u32) -> u32 {
1648        if strategy_slot == 0 { produced_before } else { 0 }
1649    }
1650
1651    fn prefill_end_of(speculation: &Speculation, lane: &Lane) -> u64 {
1652        let replay = Self::decode_replay_for(Self::strategy_slot_for(speculation, &lane.job), lane.produced_before);
1653        (lane.job.stream.len() as u64)
1654            .saturating_sub(replay as u64)
1655            .max(lane.prefilled)
1656    }
1657
1658    fn pop_admissible(&mut self) -> Option<GenerateJob> {
1659        if self.lanes.iter().any(|l| self.engine_exclusive(&l.job)) {
1660            return None;
1661        }
1662        let mut per_class = [0usize; 4];
1663        for l in &self.lanes {
1664            per_class[(l.job.class as usize).min(3)] += 1;
1665        }
1666        let now = self.tick_count;
1667        let period = self.starvation_ticks as u64;
1668        let mut best: Option<(u8, usize)> = None;
1669        let mut excl_waiting: Option<u8> = None;
1670        for (c, q) in self.queues.iter().enumerate() {
1671            let Some(head) = q.jobs.front() else {
1672                continue;
1673            };
1674            let cap = if self.class_lanes[c] == 0 {
1675                self.max_lanes
1676            } else {
1677                self.class_lanes[c]
1678            };
1679            if per_class[c] >= cap {
1680                continue;
1681            }
1682            let waited = now.saturating_sub(head.queued_at);
1683            let aged = (waited / period).min(c as u64) as u8;
1684            let eff = c as u8 - aged;
1685            if self.engine_exclusive(head) && !self.lanes.is_empty() {
1686                if excl_waiting.is_none_or(|b| eff < b) {
1687                    excl_waiting = Some(eff);
1688                }
1689                continue;
1690            }
1691            if best.is_none_or(|(b, _)| eff < b) {
1692                best = Some((eff, c));
1693            }
1694        }
1695        if let (Some(e), Some((b, _))) = (excl_waiting, best) {
1696            if e <= b {
1697                return None;
1698            }
1699        }
1700        let (_, c) = best?;
1701        self.queues[c].jobs.pop_front()
1702    }
1703
1704    /// Whether `job` should wait for a lane that is prefilling a long prefix
1705    /// it shares and will publish that prefix to the cache shortly.
1706    fn awaits_a_shared_prefix(&self, job: &GenerateJob) -> bool {
1707        if job.carry.is_some()
1708            || job.class <= crate::qos::INLINE_COMPLETION
1709            || !job.media.is_empty()
1710            || self.tick_count.saturating_sub(job.queued_at) >= SHARED_PREFIX_MAX_HOLD_TICKS
1711            || self.host.has_unservable_space()
1712            || !self.host.cache_serves_all_spaces()
1713        {
1714            return false;
1715        }
1716        let Some(page) = self.host.page_size().filter(|&p| p > 0) else {
1717            return false;
1718        };
1719        if self.host.needs_prefix_checkpoints() {
1720            if self.prefix_checkpoints_off {
1721                return false;
1722            }
1723            // A retired lane's entry covers its reply and seeds nothing that
1724            // only shares its prompt: wait instead for the state kept where
1725            // the shared part ends, while a lane can still stop there.
1726            return self.lanes.iter().any(|l| {
1727                if l.preempted || l.job.completion || !l.job.media.is_empty() {
1728                    return false;
1729                }
1730                let shared = job.stream.iter().zip(&l.all_tokens).take_while(|(a, b)| a == b).count() as u64;
1731                let total = Self::prefill_end_of(&self.speculation, l);
1732                shared >= SHARED_PREFIX_MIN_TOKENS && shared + 1 < total && !Self::kept_for(l, shared) && l.prefilled < shared
1733            });
1734        }
1735        self.lanes.iter().any(|l| {
1736            if l.preempted
1737                || l.parked_lossy
1738                || l.job.completion
1739                || !l.job.media.is_empty()
1740                || l.job.max_tokens.saturating_sub(l.produced) > SHARED_PREFIX_TAIL_TOKENS
1741            {
1742                return false;
1743            }
1744            let shared = job.stream.iter().zip(&l.all_tokens).take_while(|(a, b)| a == b).count() as u64;
1745            let shared = shared / page * page;
1746            shared >= SHARED_PREFIX_MIN_TOKENS && shared > l.warm_prefix
1747        })
1748    }
1749
1750    /// On a runtime whose state cannot be cut back, a prefilling lane stops
1751    /// where the requests queued behind it stop sharing its prompt (the
1752    /// shortest such share of at least SHARED_PREFIX_MIN_TOKENS), so its
1753    /// state there can be kept and seed them all; and where a chat's input
1754    /// ends, before the generation prompt. The next turn renders that input
1755    /// again but not the generation prompt as it was (a template drops the
1756    /// reasoning opener from an assistant turn in history), so the state the
1757    /// executor keeps where the prompt ends is no prefix of it. Of the points
1758    /// past what is prefilled, the nearest comes first; the next is planned
1759    /// once the prefill has passed it.
1760    fn plan_prefix_checkpoints(&mut self) {
1761        if !self.host.needs_prefix_checkpoints() || self.prefix_checkpoints_off {
1762            for l in self.lanes.iter_mut() {
1763                l.checkpoint_at = 0;
1764            }
1765            return;
1766        }
1767        // A prefill that stopped where the input ends without planning to
1768        // (a shed chunk planned again from there, or a budget that ran out
1769        // there) keeps the state there before it goes on.
1770        let mut landed = false;
1771        for l in self.lanes.iter_mut() {
1772            if Self::turn_point(&self.speculation, l) == Some(l.prefilled) {
1773                (l.checkpoint_at, l.checkpoint_due, landed) = (l.prefilled, true, true);
1774            }
1775        }
1776        if landed {
1777            self.publish_prefix_checkpoints();
1778        }
1779        for i in 0..self.lanes.len() {
1780            let l = &self.lanes[i];
1781            let total = Self::prefill_end_of(&self.speculation, l);
1782            let mut at = 0u64;
1783            if Self::checkpoint_eligible(l, total) {
1784                for job in self.queues.iter().flat_map(|q| q.jobs.iter()) {
1785                    let shared = job.stream.iter().zip(&l.all_tokens).take_while(|(a, b)| a == b).count() as u64;
1786                    let useful = shared >= SHARED_PREFIX_MIN_TOKENS
1787                        && shared + 1 < total
1788                        && shared > l.prefilled
1789                        && !Self::kept_for(l, shared);
1790                    if useful && (at == 0 || shared < at) {
1791                        at = shared;
1792                    }
1793                }
1794                if let Some(end) = Self::turn_point(&self.speculation, l).filter(|&end| end > l.prefilled) {
1795                    if at == 0 || end < at {
1796                        at = end;
1797                    }
1798                }
1799            }
1800            self.lanes[i].checkpoint_at = at;
1801        }
1802    }
1803
1804    fn checkpoint_eligible(l: &Lane, total: u64) -> bool {
1805        !l.preempted
1806            && !l.job.completion
1807            && l.job.media.is_empty()
1808            && !matches!(l.phase, LanePhase::Finishing)
1809            && l.prefilled < total
1810    }
1811
1812    /// Where this lane's state is kept for the chat's next turn: the end of the input,
1813    /// unless the lane started there or already prefilled past it.
1814    fn turn_point(speculation: &Speculation, l: &Lane) -> Option<u64> {
1815        let end = l.job.input_end;
1816        let total = Self::prefill_end_of(speculation, l);
1817        (Self::checkpoint_eligible(l, total)
1818            && !l.turn_kept
1819            && end >= TURN_CHECKPOINT_MIN_TOKENS
1820            && end + 1 < total
1821            && end >= l.prefilled
1822            && end != l.warm_prefix)
1823            .then_some(end)
1824    }
1825
1826    /// Whether a state that seeds a request sharing `shared` tokens of this
1827    /// lane's prompt is in the cache already: the lane kept one, or started
1828    /// from one.
1829    fn kept_for(l: &Lane, shared: u64) -> bool {
1830        let serves = |at: u64| at >= SHARED_PREFIX_MIN_TOKENS && at <= shared;
1831        serves(l.checkpointed) || serves(l.warm_prefix)
1832    }
1833
1834    /// Keeps a copy of each lane's state where this tick's prefill stopped
1835    /// for a checkpoint, in the prefix cache under the tokens it covers.
1836    fn publish_prefix_checkpoints(&mut self) {
1837        for i in 0..self.lanes.len() {
1838            if !std::mem::take(&mut self.lanes[i].checkpoint_due) {
1839                continue;
1840            }
1841            let (tag, at) = (self.lanes[i].tag, self.lanes[i].checkpoint_at);
1842            let tokens = self.lanes[i].all_tokens[..at as usize].to_vec();
1843            let turn = at == self.lanes[i].job.input_end;
1844            self.lanes[i].turn_kept |= turn;
1845            // A turn's checkpoint serves the chat's next turn alone, and goes
1846            // before a prefix requests share when the cache is short of room.
1847            match self.fork_into_cache(tag, &tokens, turn) {
1848                Ok(()) => {
1849                    self.lanes[i].checkpointed = at;
1850                    if turn {
1851                        self.supersede_turn_checkpoints(&tokens);
1852                    }
1853                }
1854                Err(CheckpointError::Refused(e)) => {
1855                    tracing::info!(lane = tag, "the engine forks no lane's sequence ({e}): no shared-prefix checkpoints");
1856                    self.prefix_checkpoints_off = true;
1857                }
1858                Err(CheckpointError::Failed(e)) => tracing::debug!(lane = tag, at, "prefix checkpoint not kept: {e}"),
1859            }
1860        }
1861    }
1862
1863    /// Records the turn checkpoint just kept under `tokens`, and drops those
1864    /// it extends: earlier turns of the same chat.
1865    fn supersede_turn_checkpoints(&mut self, tokens: &[u32]) {
1866        let digest = superfluid_fingerprint::content_digest(tokens);
1867        let mut kept = std::collections::VecDeque::with_capacity(self.turn_checkpoints.len() + 1);
1868        for (len, d) in std::mem::take(&mut self.turn_checkpoints) {
1869            let n = len as usize;
1870            if n == tokens.len() && d == digest {
1871                continue;
1872            }
1873            if n < tokens.len() && superfluid_fingerprint::content_digest(&tokens[..n]) == d {
1874                self.unpublish(&tokens[..n]);
1875                continue;
1876            }
1877            kept.push_back((len, d));
1878        }
1879        kept.push_back((tokens.len() as u64, digest));
1880        while kept.len() > TURN_CHECKPOINTS_TRACKED {
1881            kept.pop_front();
1882        }
1883        self.turn_checkpoints = kept;
1884    }
1885
1886    fn unpublish(&mut self, tokens: &[u32]) {
1887        let res = match self.stage(tokens) {
1888            Ok(span) => self.host.client().state_sync(linkw::StateSyncReqMsg::Unpublish { span }).map_err(|e| e.to_string()),
1889            Err(e) => Err(e.to_string()),
1890        };
1891        if let Err(e) = res {
1892            tracing::debug!(len = tokens.len(), "superseded turn checkpoint not dropped: {e}");
1893        }
1894    }
1895
1896    fn fork_into_cache(&mut self, tag: u64, tokens: &[u32], private: bool) -> Result<(), CheckpointError> {
1897        use CheckpointError::Failed;
1898        let seq = match self.host.client().state_sync(linkw::StateSyncReqMsg::LaneSequence { lane_tag: tag }) {
1899            Ok(linkw::StateSyncOkMsg::LaneSequence { sequence }) => sequence,
1900            Ok(_) => return Err(Failed("mismatched LaneSequence reply".into())),
1901            Err(e) => return Err(CheckpointError::of(e)),
1902        };
1903        let child = match self.host.client().state_sync(linkw::StateSyncReqMsg::Fork {
1904            parent_sequence: seq,
1905            flags: if private { superfluid_abi::fork_flags::PRIVATE } else { 0 },
1906        }) {
1907            Ok(linkw::StateSyncOkMsg::Fork { child_sequence }) => child_sequence,
1908            Ok(_) => return Err(Failed("mismatched Fork reply".into())),
1909            Err(e) => return Err(CheckpointError::of(e)),
1910        };
1911        let published = match self.stage(tokens) {
1912            Ok(span) => match self.host.client().state_sync(linkw::StateSyncReqMsg::PublishSequence { sequence: child, span }) {
1913                Ok(linkw::StateSyncOkMsg::PublishSequence) => Ok(()),
1914                Ok(_) => Err("mismatched PublishSequence reply".into()),
1915                Err(e) => Err(e.to_string()),
1916            },
1917            Err(e) => Err(e.to_string()),
1918        };
1919        if published.is_err() {
1920            let _ = self.host.client().state_sync(linkw::StateSyncReqMsg::FreeSequence { sequence: child });
1921        }
1922        published.map_err(Failed)
1923    }
1924
1925    /// Whether a chat or inline completion is queued, being served, or arrived within
1926    /// LATENCY_HOLD: such a tick stays short.
1927    fn latency_sensitive(&self) -> bool {
1928        self.cancels.latency_recent(LATENCY_HOLD)
1929            || self.queues[..=crate::qos::INLINE_COMPLETION as usize].iter().any(|q| !q.jobs.is_empty())
1930            || self.lanes.iter().any(|l| l.job.class <= crate::qos::INLINE_COMPLETION && !l.preempted)
1931    }
1932
1933    fn joiner_waits_on_this_tick(&self) -> bool {
1934        if self.max_lanes <= 1 {
1935            return false;
1936        }
1937        self.live_lanes() < self.max_lanes
1938            || self.queues[..=crate::qos::INLINE_COMPLETION as usize]
1939                .iter()
1940                .any(|q| !q.jobs.is_empty())
1941            || self.lanes.iter().any(|l| {
1942                l.job.class > crate::qos::INLINE_COMPLETION
1943                    && !l.preempted
1944                    && l.job.media.is_empty()
1945                    && !self.engine_exclusive(&l.job)
1946            })
1947    }
1948
1949    /// Lanes holding a slot past this tick: a preempted lane retires in it,
1950    /// and the engine retires before it admits, so its slot is already free.
1951    fn live_lanes(&self) -> usize {
1952        self.lanes
1953            .iter()
1954            .filter(|l| !(l.preempted && matches!(l.phase, LanePhase::Finishing)))
1955            .count()
1956    }
1957
1958    fn preempt_for_queued(&mut self) {
1959        if self.live_lanes() < self.max_lanes {
1960            return;
1961        }
1962        let Some(want) = self
1963            .queues
1964            .iter()
1965            .enumerate()
1966            .find(|(_, q)| !q.jobs.is_empty())
1967            .map(|(c, _)| c as u8)
1968        else {
1969            return;
1970        };
1971        if want > crate::qos::INLINE_COMPLETION {
1972            return;
1973        }
1974        // A decoding lane first; else one still prefilling, whose prefilled
1975        // pages go to the prefix cache on retire and seed it when it resumes.
1976        let victim_in = |phase_decoding: bool| {
1977            self.lanes
1978                .iter()
1979                .enumerate()
1980                .filter(|(_, l)| {
1981                    l.admitted
1982                        && !l.preempted
1983                        && !self.engine_exclusive(&l.job)
1984                        && l.job.class > want
1985                        && l.job.media.is_empty()
1986                        && if phase_decoding {
1987                            matches!(l.phase, LanePhase::Decoding)
1988                        } else {
1989                            matches!(l.phase, LanePhase::Prefilling)
1990                        }
1991                })
1992                .max_by_key(|(_, l)| (l.job.class, l.tag))
1993                .map(|(i, _)| i)
1994        };
1995        let victim = victim_in(true).or_else(|| victim_in(false));
1996        if let Some(i) = victim {
1997            self.preempt_lane(i);
1998        }
1999    }
2000
2001    fn preempt_lane(&mut self, i: usize) {
2002        let lane = &mut self.lanes[i];
2003        lane.preempted = true;
2004        lane.phase = LanePhase::Finishing;
2005        lane.fin = finish::NONE;
2006        tracing::info!(
2007            session = lane.session,
2008            lane = lane.tag,
2009            class = lane.job.class,
2010            produced = lane.produced,
2011            "lane preempted"
2012        );
2013        self.stats.preemptions.fetch_add(1, Ordering::Relaxed);
2014    }
2015
2016    fn poll_os_pressure(&mut self) {
2017        let Some(src) = self.pressure_source.as_mut() else {
2018            return;
2019        };
2020        match src.level() {
2021            crate::pressure::PressureLevel::Normal => {}
2022            crate::pressure::PressureLevel::Warning => {
2023                self.stats.os_pressure_events.fetch_add(1, Ordering::Relaxed);
2024                self.relief_round_now();
2025            }
2026            crate::pressure::PressureLevel::Critical => {
2027                self.stats.os_pressure_events.fetch_add(1, Ordering::Relaxed);
2028                self.relief_round_now();
2029                let victims: Vec<usize> = self
2030                    .lanes
2031                    .iter()
2032                    .enumerate()
2033                    .filter(|(_, l)| {
2034                        l.admitted
2035                            && !l.preempted
2036                            && l.job.class == crate::qos::BACKGROUND_AGENT
2037                            && matches!(l.phase, LanePhase::Decoding)
2038                    })
2039                    .map(|(i, _)| i)
2040                    .collect();
2041                for i in victims {
2042                    self.preempt_lane(i);
2043                }
2044            }
2045        }
2046    }
2047
2048    fn relief_round_now(&mut self) {
2049        // Memory is short on the machine: what the cache holds in host memory
2050        // goes first, whatever the pool holds.
2051        self.evict_cache(0, HOST_EXPORT_CLASS);
2052        if let Ok(pong) = self.host.client().ping() {
2053            let mem = pong.mem;
2054            let low = self.pressure_low_pct as u64;
2055            let (used, total) = (mem.pool_blocks_used, mem.pool_blocks_total);
2056            if total == 0 || used * 100 <= total * low {
2057                return;
2058            }
2059            let saved = self.pressure_high_pct;
2060            self.pressure_high_pct = 1;
2061            self.relieve_pressure(&mem, true);
2062            self.pressure_high_pct = saved;
2063        }
2064    }
2065
2066    fn after_respawn(&mut self) {
2067        self.max_lanes = lane_cap(
2068            self.configured_lanes,
2069            self.host.max_seqs(),
2070            ring_slots(&mut self.host, superfluid_shm::TOKEN_RING_OUT),
2071        );
2072        self.pins.set_fallback_pool_tokens(self.host.max_stream_tokens() * self.max_lanes as u64);
2073        if let Some(why) = self.host.take_record_change() {
2074            self.retire(&why);
2075        }
2076    }
2077
2078    // Respawning stops after a streak; without a retire nothing would load
2079    // the model again and every request would fail until a restart.
2080    fn give_up_on_worker(&mut self) {
2081        if self.retired.lock().expect("retired").is_none() {
2082            self.retire(&format!(
2083                "the engine worker died and {MAX_WORKER_RESPAWNS} respawns in a row did not keep one up"
2084            ));
2085        }
2086    }
2087
2088    fn retire(&mut self, why: &str) {
2089        let msg = format!("{why}; the model loads again on its next request, so retry");
2090        eprintln!("[superfluid] {msg}");
2091        *self.retired.lock().expect("retired") = Some(msg.clone());
2092        let drained: Vec<u64> = self.lanes.iter().map(|l| l.session).collect();
2093        for lane in self.lanes.drain(..) {
2094            self.active.lock().expect("active set").remove(&lane.session);
2095            let _ = lane.job.events.send(LaneEvent::Failed(msg.clone()));
2096        }
2097        for session in drained {
2098            self.drop_ephemeral_artifact(session);
2099        }
2100        let queued: Vec<GenerateJob> = self.queues.iter_mut().flat_map(|q| q.jobs.drain(..)).collect();
2101        for job in queued {
2102            self.finish_failed(&job, &msg);
2103        }
2104    }
2105
2106    fn stage(&mut self, tokens: &[u32]) -> Result<linkw::TokenRefMsg, superfluid_agent::AgentError> {
2107        self.staged_unread += 1;
2108        self.host.client().stage_prompt(tokens)
2109    }
2110
2111    fn admit_from_queue(&mut self) {
2112        self.admit_retry = false;
2113        let retired = self.retired.lock().expect("retired").clone();
2114        if let Some(msg) = retired {
2115            let queued: Vec<GenerateJob> = self.queues.iter_mut().flat_map(|q| q.jobs.drain(..)).collect();
2116            for job in queued {
2117                self.finish_failed(&job, &msg);
2118            }
2119            return;
2120        }
2121        let mut held: Vec<GenerateJob> = Vec::new();
2122        if self.lanes.iter().all(|l| l.admitted) {
2123            self.staged_unread = 0;
2124        }
2125        let stage_slots = ring_slots(&mut self.host, superfluid_shm::TOKEN_RING_IN).unwrap_or(u32::MAX);
2126        while self.live_lanes() < self.max_lanes {
2127            // A staged prompt is read when its tick admits it; staging past
2128            // the ring's size would overwrite one still waiting, and fail the
2129            // tick. The rest wait a tick in their queue.
2130            if self.staged_unread > 0 && self.staged_unread.saturating_add(STAGES_PER_ADMISSION) > stage_slots {
2131                break;
2132            }
2133            let Some(job) = self.pop_admissible() else {
2134                break;
2135            };
2136            if job.stream.is_empty() {
2137                self.finish_failed(&job, "session has no content to generate from");
2138                continue;
2139            }
2140            if self.awaits_a_shared_prefix(&job) {
2141                held.push(job);
2142                continue;
2143            }
2144            if let Some(deadline) = job.deadline {
2145                if std::time::Instant::now() > deadline {
2146                    let _ = job.events.send(LaneEvent::Expired);
2147                    self.active.lock().expect("active set").remove(&job.session);
2148                    self.drop_ephemeral_artifact(job.session);
2149                    continue;
2150                }
2151            }
2152            let max_stream = self.host.max_stream_tokens();
2153            if job.stream.len() as u64 > max_stream {
2154                let msg = DaemonError::StreamTooLong {
2155                    len: job.stream.len() as u64,
2156                    max: max_stream,
2157                }
2158                .to_string();
2159                self.finish_failed(&job, &msg);
2160                continue;
2161            }
2162            let span = match self.stage(&job.stream) {
2163                Ok(s) => s,
2164                Err(e) => {
2165                    let transport = matches!(
2166                        e,
2167                        superfluid_agent::AgentError::Io(_)
2168                            | superfluid_agent::AgentError::Closed
2169                            | superfluid_agent::AgentError::Shm(_)
2170                    );
2171                    if transport && self.host.is_process() && self.respawn_streak < MAX_WORKER_RESPAWNS {
2172                        self.respawn_streak += 1;
2173                        eprintln!(
2174                            "[superfluid] engine worker gone at admission ({e}); respawning"
2175                        );
2176                        self.stats.worker_respawns.fetch_add(1, Ordering::Relaxed);
2177                        if self
2178                            .respawn_worker()
2179                            .and_then(|()| self.speculation.reregister(&mut self.host))
2180                            .is_ok()
2181                        {
2182                            let c = (job.class as usize).min(3);
2183                            self.queues[c].jobs.push_front(job);
2184                            self.after_respawn();
2185                            self.admit_retry = true;
2186                            break;
2187                        }
2188                    }
2189                    self.finish_failed(&job, &format!("stage failed: {e}"));
2190                    continue;
2191                }
2192            };
2193            let mut seed_handle = 0u64;
2194            let mut warm_prefix = 0u64;
2195            if self.host.has_unservable_space() {
2196                self.stats
2197                    .unservable_cold_admissions
2198                    .fetch_add(1, Ordering::Relaxed);
2199            } else if let Some(page) = self.host.page_size() {
2200                if self.host.cache_serves_all_spaces() && job.media.is_empty() {
2201                if let Ok(m) = self.host.client().match_prefix(ALL_SPACES, span) {
2202                    let align_down = |l: u64| l.checked_div(page).unwrap_or(0) * page;
2203                    let replay = Self::decode_replay_for(
2204                        Self::strategy_slot_for(&self.speculation, &job),
2205                        job.carry.as_ref().map(|c| c.produced_before).unwrap_or(0),
2206                    );
2207                    let stream_len = job.stream.len() as u64;
2208                    let aligned_cap =
2209                        align_down(stream_len.saturating_sub(1).min(stream_len.saturating_sub(replay as u64)));
2210                    let best = crate::pins::common_seed_lengths(&self.host, &m, aligned_cap)
2211                        .first()
2212                        .copied()
2213                        .unwrap_or(0);
2214                    if best >= page {
2215                        if let Ok(h) = self.host.client().seed_acquire(span, best, 0) {
2216                            seed_handle = h;
2217                            warm_prefix = best;
2218                        }
2219                    }
2220                }
2221                }
2222                (seed_handle, warm_prefix) =
2223                    self.resume_from_park(&job, span, page, seed_handle, warm_prefix);
2224            }
2225            if warm_prefix > 0 {
2226                self.stats
2227                    .warm_prefix_tokens
2228                    .fetch_add(warm_prefix, Ordering::Relaxed);
2229            } else {
2230                self.stats.cold_admissions.fetch_add(1, Ordering::Relaxed);
2231            }
2232            let tag = self.next_lane;
2233            self.next_lane += 1;
2234            let mut job = job;
2235            let mut media_ok = true;
2236            let mut bound: Vec<u64> = Vec::new();
2237            for (offset, path, _n) in job.media.iter().filter(|(o, _, _)| *o >= warm_prefix) {
2238                let path_s = path.to_string_lossy().into_owned();
2239                let enc = self.host.client().state_sync(linkw::StateSyncReqMsg::MediaEncode {
2240                    image_path: path_s.clone(),
2241                });
2242                let handle = match enc {
2243                    Ok(linkw::StateSyncOkMsg::MediaEncode { media_handle, .. }) => media_handle,
2244                    other => {
2245                        eprintln!("[superfluid] media encode failed for {path_s}: {other:?}");
2246                        media_ok = false;
2247                        break;
2248                    }
2249                };
2250                let bind = self.host.client().state_sync(linkw::StateSyncReqMsg::MediaBind {
2251                    lane_tag: tag,
2252                    media_handle: handle,
2253                    token_offset: *offset as u32,
2254                });
2255                if matches!(bind, Ok(linkw::StateSyncOkMsg::MediaBind)) {
2256                    bound.push(handle);
2257                } else {
2258                    let _ = self
2259                        .host
2260                        .client()
2261                        .state_sync(linkw::StateSyncReqMsg::MediaRelease { media_handle: handle });
2262                    media_ok = false;
2263                    break;
2264                }
2265            }
2266            if !media_ok {
2267                for handle in &bound {
2268                    let _ = self.host.client().state_sync(
2269                        linkw::StateSyncReqMsg::MediaRelease { media_handle: *handle },
2270                    );
2271                }
2272                if seed_handle != 0 {
2273                    let _ = self.host.client().seed_release(seed_handle);
2274                }
2275                self.finish_failed(&job, "media encode/bind failed (see daemon log)");
2276                continue;
2277            }
2278            let spec_before = job.carry.as_ref().map(|c| c.spec_before).unwrap_or_default();
2279            let (seg, utf8, disp_seg, disp_utf8, tool_buf, produced_before, warm_prefix_first) =
2280                match job.carry.take() {
2281                    Some(c) => (
2282                        c.seg,
2283                        c.utf8,
2284                        c.disp_seg,
2285                        c.disp_utf8,
2286                        c.tool_buf,
2287                        c.produced_before,
2288                        c.warm_prefix_first,
2289                    ),
2290                    None => {
2291                        let mk = || {
2292                            let mut seg = self.codec.channelizer();
2293                            if job.extras.open_channel != 0 {
2294                                seg.prime(job.extras.open_channel);
2295                            }
2296                            seg
2297                        };
2298                        (
2299                            mk(),
2300                            Utf8Stream::default(),
2301                            mk(),
2302                            Utf8Stream::default(),
2303                            String::new(),
2304                            0,
2305                            None,
2306                        )
2307                    }
2308                };
2309            let all_tokens = job.stream.clone();
2310            self.lanes.push(Lane {
2311                tag,
2312                session: job.session,
2313                job,
2314                phase: LanePhase::Prefilling,
2315                prefilled: warm_prefix,
2316                produced: 0,
2317                fin: finish::NONE,
2318                warm_prefix,
2319                utf8,
2320                seg,
2321                disp_seg,
2322                disp_utf8,
2323                frame_fed: 0,
2324                emit_dead: false,
2325                tool_buf,
2326                admitted: false,
2327                seed_handle,
2328                span,
2329                all_tokens,
2330                parked: false,
2331                parked_lossy: false,
2332                preempted: false,
2333                starved_ticks: 0,
2334                prefill_starved_ticks: 0,
2335                checkpoint_at: 0,
2336                checkpoint_due: false,
2337                checkpointed: 0,
2338                turn_kept: false,
2339                produced_before,
2340                disp_tokens: produced_before,
2341                spec: spec_before,
2342                warm_prefix_first,
2343                admitted_at: std::time::Instant::now(),
2344                ttft_observed: false,
2345                ttft_frame_ms: None,
2346                fault_code: None,
2347            });
2348        }
2349        for job in held.into_iter().rev() {
2350            let c = (job.class as usize).min(3);
2351            self.queues[c].jobs.push_front(job);
2352        }
2353    }
2354
2355    fn resume_from_park(
2356        &mut self,
2357        job: &GenerateJob,
2358        span: linkw::TokenRefMsg,
2359        page: u64,
2360        seed_handle: u64,
2361        warm_prefix: u64,
2362    ) -> (u64, u64) {
2363        let Some(dir) = self.park_dir.clone() else {
2364            return (seed_handle, warm_prefix);
2365        };
2366        let Some(p) = park::read(&dir, job.session) else {
2367            return (seed_handle, warm_prefix);
2368        };
2369        if p.multi_space {
2370            if !self.host.multi_space() && self.host.parkable_spaces().len() != p.spaces.len() {
2371                return (seed_handle, warm_prefix);
2372            }
2373            return match self.resume_adopt(job, &p) {
2374                Ok((h, covered)) => {
2375                    if seed_handle != 0 {
2376                        let _ = self.host.client().seed_release(seed_handle);
2377                    }
2378                    self.stats.resumes.fetch_add(1, Ordering::Relaxed);
2379                    (h, covered)
2380                }
2381                Err(msg) => {
2382                    eprintln!("[superfluid] park resume fell back cold (session {}): {msg}", job.session);
2383                    (seed_handle, warm_prefix)
2384                }
2385            };
2386        }
2387        if self.host.multi_space() || !job.media.is_empty() {
2388            return (seed_handle, warm_prefix);
2389        }
2390        let align_down = |l: u64| l.checked_div(page).unwrap_or(0) * page;
2391        let aligned_cap = align_down(job.stream.len() as u64 - 1);
2392        let usable = p.covered.min(aligned_cap);
2393        if p.covered % page != 0
2394            || usable <= warm_prefix
2395            || p.covered as usize > job.stream.len()
2396            || park::stream_digest(&job.stream[..p.covered as usize]) != p.stream_digest
2397        {
2398            return (seed_handle, warm_prefix);
2399        }
2400        let covered_span = match self.stage(&job.stream[..p.covered as usize]) {
2401            Ok(s) => s,
2402            Err(e) => {
2403                eprintln!(
2404                    "[superfluid] park resume fell back cold (session {}): stage: {e}",
2405                    job.session
2406                );
2407                return (seed_handle, warm_prefix);
2408            }
2409        };
2410        if let Err(msg) = self.restore_to_cache(job.session, &p, covered_span) {
2411            eprintln!(
2412                "[superfluid] park resume fell back cold (session {}): {msg}",
2413                job.session
2414            );
2415            return (seed_handle, warm_prefix);
2416        }
2417        match self.host.client().seed_acquire(span, usable, 0) {
2418            Ok(h) => {
2419                if seed_handle != 0 {
2420                    let _ = self.host.client().seed_release(seed_handle);
2421                }
2422                self.stats.resumes.fetch_add(1, Ordering::Relaxed);
2423                (h, usable)
2424            }
2425            Err(e) => {
2426                eprintln!(
2427                    "[superfluid] park resume restored but could not seed (session {}): {e}",
2428                    job.session
2429                );
2430                (seed_handle, warm_prefix)
2431            }
2432        }
2433    }
2434
2435    fn restore_to_cache(
2436        &mut self,
2437        session: u64,
2438        p: &park::Parked,
2439        span: linkw::TokenRefMsg,
2440    ) -> Result<(), String> {
2441        let _op_span = tracing::info_span!(
2442            "restore",
2443            session,
2444            model = self.trace_scope.model.as_str(),
2445            store_ns = self.trace_scope.ns,
2446            covered = p.covered,
2447            bytes = p.sealed.len(),
2448            encoding = p.encoding
2449        )
2450        .entered();
2451        let client = self.host.client();
2452        let seq = match client.state_sync(linkw::StateSyncReqMsg::CreateSequence) {
2453            Ok(linkw::StateSyncOkMsg::CreateSequence { sequence }) => sequence,
2454            Ok(_) => return Err("mismatched CreateSequence reply".into()),
2455            Err(e) => return Err(e.to_string()),
2456        };
2457        let free_seq = |client: &mut superfluid_agent::WorkerClient| {
2458            let _ = client.state_sync(linkw::StateSyncReqMsg::FreeSequence { sequence: seq });
2459        };
2460        let buf = match client.buf_create(p.sealed.len() as u64) {
2461            Ok(b) => b,
2462            Err(e) => {
2463                free_seq(client);
2464                return Err(e.to_string());
2465            }
2466        };
2467        let result = (|| {
2468            client
2469                .buf_write(buf, &p.sealed)
2470                .map_err(|e| e.to_string())?;
2471            let op = client
2472                .state_op(linkw::StateOpKind::Restore {
2473                    sequence: seq,
2474                    space_id: KV_SPACE,
2475                    buf_id: buf,
2476                })
2477                .map_err(|e| e.to_string())?;
2478            let done = client.wait_op_done(op).map_err(|e| e.to_string())?;
2479            if done.state != op_state::DONE || done.error != 0 {
2480                return Err(format!(
2481                    "restore terminal state {} error {}",
2482                    done.state, done.error
2483                ));
2484            }
2485            Ok(())
2486        })();
2487        let _ = client.buf_release(buf);
2488        if let Err(e) = result {
2489            free_seq(client);
2490            return Err(e);
2491        }
2492        match client.state_sync(linkw::StateSyncReqMsg::PublishSequence {
2493            sequence: seq,
2494            span,
2495        }) {
2496            Ok(linkw::StateSyncOkMsg::PublishSequence) => Ok(()),
2497            Ok(_) => {
2498                free_seq(client);
2499                Err("mismatched PublishSequence reply".into())
2500            }
2501            Err(e) => {
2502                free_seq(client);
2503                Err(e.to_string())
2504            }
2505        }
2506    }
2507
2508    fn relieve_pressure(&mut self, mem: &linkw::MemCountersMsg, host_short: bool) {
2509        let (high, low) = (self.pressure_high_pct as u64, self.pressure_low_pct as u64);
2510        if std::env::var_os("SUPERFLUID_DEBUG_PRESSURE").is_some() {
2511            eprintln!(
2512                "[pressure] used={} total={} evictable={}",
2513                mem.pool_blocks_used, mem.pool_blocks_total, mem.pool_bytes_evictable
2514            );
2515        }
2516        if mem.pool_blocks_total == 0 || high == 0 {
2517            return;
2518        }
2519        let Some(block_bytes) = self.host.kv_block_bytes() else {
2520            return;
2521        };
2522        let (used, total) = (mem.pool_blocks_used, mem.pool_blocks_total);
2523        if used * 100 < total * high {
2524            return;
2525        }
2526        let low_blocks = total * low / 100;
2527        let excess = used.saturating_sub(low_blocks);
2528        if excess == 0 {
2529            return;
2530        }
2531        let slack = total * block_bytes / 4;
2532        let mut evictable = mem.pool_bytes_evictable;
2533        let need = excess * block_bytes;
2534        if need > evictable && self.pins.any_held() {
2535            evictable += self
2536                .pins
2537                .yield_for(&mut self.host, &self.stats, need - evictable);
2538        }
2539        let bytes_target = (excess * block_bytes + slack).min(evictable.max(1));
2540        // The pool's pressure leaves what the cache holds in host memory
2541        // alone, and the entries that leave the pool may be kept there.
2542        let all = if host_short { u64::MAX } else { !HOST_EXPORT_CLASS };
2543        // The slack is taken from entries no two requests have shared; a
2544        // shared prefix goes only for what the pool needs.
2545        let mut freed = self.evict_cache(bytes_target, all & !SHARED_PREFIX_CLASS);
2546        if freed < need {
2547            freed += self.evict_cache(need - freed, all);
2548        }
2549        if freed < need && self.pins.any_held() {
2550            let yielded = self
2551                .pins
2552                .yield_for(&mut self.host, &self.stats, need - freed);
2553            if yielded > 0 {
2554                freed += self.evict_cache((need - freed + slack).min(yielded), all);
2555            }
2556        }
2557        if freed > 0 {
2558            self.stats.pressure_evictions.fetch_add(1, Ordering::Relaxed);
2559            self.stats
2560                .pressure_bytes_evicted
2561                .fetch_add(freed, Ordering::Relaxed);
2562        }
2563    }
2564
2565    fn evict_cache(&mut self, bytes_target: u64, classes: u64) -> u64 {
2566        let res = self
2567            .host
2568            .client()
2569            .state_sync(linkw::StateSyncReqMsg::CacheEvict {
2570                bytes_target,
2571                policy: linkw::ShedPolicyMsg {
2572                    victim_lanes: vec![],
2573                    evictable_cache_classes: classes,
2574                    protected_quota_bytes: 0,
2575                    max_evict_bytes: bytes_target,
2576                },
2577            });
2578        match res {
2579            Ok(linkw::StateSyncOkMsg::CacheEvict { bytes_freed }) => bytes_freed,
2580            _ => 0,
2581        }
2582    }
2583
2584    fn retry_park_removals(&mut self) {
2585        if self.park_removals_pending.is_empty() {
2586            return;
2587        }
2588        let (Some(dir), Some(tx)) = (self.park_dir.clone(), self.park_tx.as_ref()) else {
2589            self.park_removals_pending.clear();
2590            return;
2591        };
2592        self.park_removals_pending.retain(|&session| {
2593            matches!(
2594                tx.try_send(ParkTask::Remove {
2595                    dir: dir.clone(),
2596                    session,
2597                }),
2598                Err(std::sync::mpsc::TrySendError::Full(_))
2599            )
2600        });
2601    }
2602
2603    /// Lanes whose admission the engine could not take this tick go back to
2604    /// the head of their class queue, their seed leases released, to be
2605    /// admitted on a later tick.
2606    fn defer_admits(&mut self, tags: &[u64]) {
2607        let mut back: Vec<GenerateJob> = Vec::new();
2608        let mut seeds: Vec<u64> = Vec::new();
2609        let mut i = 0;
2610        while i < self.lanes.len() {
2611            if tags.contains(&self.lanes[i].tag) {
2612                let lane = self.lanes.remove(i);
2613                if lane.seed_handle != 0 {
2614                    seeds.push(lane.seed_handle);
2615                }
2616                back.push(lane.job);
2617            } else {
2618                i += 1;
2619            }
2620        }
2621        for seed in seeds {
2622            let _ = self.host.client().seed_release(seed);
2623        }
2624        for job in back.into_iter().rev() {
2625            let c = (job.class as usize).min(3);
2626            self.queues[c].jobs.push_front(job);
2627        }
2628    }
2629
2630    fn drop_refused(&mut self, refused: &HashMap<u64, u32>) {
2631        let mut gone: Vec<(u64, u64)> = Vec::new();
2632        let mut active = self.active.lock().expect("active set");
2633        self.lanes.retain(|lane| {
2634            let Some(&code) = refused.get(&lane.tag) else { return true };
2635            active.remove(&lane.session);
2636            let status = superfluid_abi::Status::from_raw(code as i32)
2637                .map(|s| format!("{s:?}"))
2638                .unwrap_or_else(|| format!("status {}", code as i32));
2639            tracing::warn!(session = lane.session, lane = lane.tag, status = %status, "engine refused an admit");
2640            let _ = lane.job.events.send(LaneEvent::Failed(format!("the engine refused to admit this request ({status})")));
2641            gone.push((lane.session, lane.seed_handle));
2642            false
2643        });
2644        drop(active);
2645        for (session, seed) in gone {
2646            if seed != 0 {
2647                let _ = self.host.client().seed_release(seed);
2648            }
2649            self.drop_ephemeral_artifact(session);
2650        }
2651    }
2652
2653    fn drop_ephemeral_artifact(&mut self, session: u64) {
2654        if !self.ephemeral_parked.remove(&session) {
2655            return;
2656        }
2657        let Some(dir) = self.park_dir.clone() else {
2658            return;
2659        };
2660        let Some(tx) = self.park_tx.as_ref() else {
2661            return;
2662        };
2663        if tx.try_send(ParkTask::Remove { dir, session }).is_err() {
2664            self.park_removals_pending.push(session);
2665        }
2666    }
2667
2668    fn park_lane(&mut self, i: usize) -> Result<(), ParkRefused> {
2669        let Some(dir) = self.park_dir.clone() else {
2670            return Ok(());
2671        };
2672        if self.lanes[i].job.completion {
2673            return Ok(());
2674        }
2675        if (self.host.multi_space() || !self.lanes[i].job.media.is_empty()) && !self.host.has_unservable_space() {
2676            return self.park_lane_multi(i, dir);
2677        }
2678        if self.host.has_unservable_space() {
2679            self.stats
2680                .unservable_park_refusals
2681                .fetch_add(1, Ordering::Relaxed);
2682            return Err(ParkRefused::SpaceUnservable);
2683        }
2684        let Some(page) = self.host.page_size() else {
2685            return Ok(());
2686        };
2687        let (tag, session) = (self.lanes[i].tag, self.lanes[i].session);
2688        let n = self.lanes[i].all_tokens.len() as u64;
2689        if n <= 1 {
2690            return Ok(());
2691        }
2692        let mut aligned = ((n - 1) / page) * page;
2693        if aligned < page {
2694            return Ok(());
2695        }
2696        let park_encoding = if self.park_lossy {
2697            encoding::Q8
2698        } else {
2699            encoding::LOSSLESS
2700        };
2701        let client = self.host.client();
2702        let seq = match client.state_sync(linkw::StateSyncReqMsg::LaneSequence { lane_tag: tag }) {
2703            Ok(linkw::StateSyncOkMsg::LaneSequence { sequence }) => sequence,
2704            Ok(_) => return Err("mismatched LaneSequence reply".into()),
2705            Err(e) => return Err(e.to_string().into()),
2706        };
2707        if let Ok(linkw::StateSyncOkMsg::SnapshotBoundary { boundary }) =
2708            client.state_sync(linkw::StateSyncReqMsg::SnapshotBoundary { sequence: seq, space_id: KV_SPACE, cap: n - 1 })
2709        {
2710            aligned = aligned.min(boundary);
2711            if aligned < page {
2712                return Ok(());
2713            }
2714        }
2715        let digest = park::stream_digest(&self.lanes[i].all_tokens[..aligned as usize]);
2716        let _span = tracing::info_span!(
2717            "park",
2718            session,
2719            model = self.trace_scope.model.as_str(),
2720            store_ns = self.trace_scope.ns,
2721            covered = aligned,
2722            lossy = self.park_lossy
2723        )
2724        .entered();
2725        let (bytes, sizing_gen) = match client.state_sync(linkw::StateSyncReqMsg::ExportSize {
2726            sequence: seq,
2727            space_id: KV_SPACE,
2728            range: linkw::TokenRangeMsg {
2729                start: 0,
2730                end: aligned,
2731            },
2732            encoding: park_encoding,
2733        }) {
2734            Ok(linkw::StateSyncOkMsg::ExportSize {
2735                required_bytes,
2736                sizing_gen,
2737            }) => (required_bytes, sizing_gen),
2738            Ok(_) => return Err("mismatched ExportSize reply".into()),
2739            Err(e) => return Err(e.to_string().into()),
2740        };
2741        let buf = client.buf_create(bytes).map_err(|e| ParkRefused::from(e.to_string()))?;
2742        let result = (|| {
2743            let op_kind = if self.park_lossy {
2744                linkw::StateOpKind::Demote {
2745                    sequence: seq,
2746                    space_id: KV_SPACE,
2747                    range: linkw::TokenRangeMsg {
2748                        start: 0,
2749                        end: aligned,
2750                    },
2751                    encoding: encoding::Q8,
2752                    buf_id: buf,
2753                    sizing_gen,
2754                }
2755            } else {
2756                linkw::StateOpKind::Snapshot {
2757                    sequence: seq,
2758                    space_id: KV_SPACE,
2759                    boundary_pos: aligned,
2760                    buf_id: buf,
2761                    sizing_gen,
2762                }
2763            };
2764            let op = client.state_op(op_kind).map_err(|e| e.to_string())?;
2765            let done = client.wait_op_done(op).map_err(|e| e.to_string())?;
2766            if done.state != op_state::DONE || done.error != 0 {
2767                return Err(format!(
2768                    "snapshot terminal state {} error {}",
2769                    done.state, done.error
2770                ));
2771            }
2772            let take = if done.bytes_moved > 0 {
2773                done.bytes_moved.min(bytes)
2774            } else {
2775                bytes
2776            };
2777            client
2778                .buf_read(buf, take as usize)
2779                .map_err(|e| e.to_string())
2780        })();
2781        let _ = client.buf_release(buf);
2782        let sealed = result.map_err(ParkRefused::from)?;
2783        let job = ParkJob {
2784            dir,
2785            session,
2786            covered: aligned,
2787            digest,
2788            encoding: park_encoding,
2789            sealed,
2790            spaces: None,
2791        };
2792        match self
2793            .park_tx
2794            .as_ref()
2795            .expect("writer exists with park_dir")
2796            .try_send(ParkTask::Write(Box::new(job)))
2797        {
2798            Ok(()) => {
2799                if self.lanes[i].job.extras.ephemeral {
2800                    self.ephemeral_parked.insert(session);
2801                }
2802                if self.park_lossy {
2803                    self.stats.parks_lossy.fetch_add(1, Ordering::Relaxed);
2804                } else {
2805                    self.stats.parks_lossless.fetch_add(1, Ordering::Relaxed);
2806                }
2807            }
2808            Err(std::sync::mpsc::TrySendError::Full(_)) => {
2809                return Err("park writer backlogged; artifact dropped".into());
2810            }
2811            Err(std::sync::mpsc::TrySendError::Disconnected(_)) => {
2812                return Err("park writer gone".into());
2813            }
2814        }
2815        self.lanes[i].parked_lossy = self.park_lossy;
2816        Ok(())
2817    }
2818
2819    fn park_lane_multi(&mut self, i: usize, dir: PathBuf) -> Result<(), ParkRefused> {
2820        let (tag, session) = (self.lanes[i].tag, self.lanes[i].session);
2821        let n = self.lanes[i].all_tokens.len() as u64;
2822        if n <= 1 {
2823            return Ok(());
2824        }
2825        let spaces = self.host.parkable_spaces();
2826        let lossy = self.park_lossy;
2827        let fallback_boundary = self.host.park_boundary(n - 1);
2828        let client = self.host.client();
2829        let seq = match client.state_sync(linkw::StateSyncReqMsg::LaneSequence { lane_tag: tag }) {
2830            Ok(linkw::StateSyncOkMsg::LaneSequence { sequence }) => sequence,
2831            Ok(_) => return Err("mismatched LaneSequence reply".into()),
2832            Err(e) => return Err(e.to_string().into()),
2833        };
2834        let cap = n - 1;
2835        let mut covered = cap;
2836        let mut queried = false;
2837        for &(space_id, _, _, _) in &spaces {
2838            match client.state_sync(linkw::StateSyncReqMsg::SnapshotBoundary { sequence: seq, space_id, cap }) {
2839                Ok(linkw::StateSyncOkMsg::SnapshotBoundary { boundary }) => {
2840                    queried = true;
2841                    covered = covered.min(boundary);
2842                }
2843                Ok(_) => return Err("mismatched SnapshotBoundary reply".into()),
2844                Err(_) => break,
2845            }
2846        }
2847        if !queried {
2848            covered = fallback_boundary;
2849        }
2850        if covered == 0 {
2851            return Ok(());
2852        }
2853        let digest = park::stream_digest(&self.lanes[i].all_tokens[..covered as usize]);
2854        let _span = tracing::info_span!(
2855            "park",
2856            session,
2857            model = self.trace_scope.model.as_str(),
2858            store_ns = self.trace_scope.ns,
2859            covered,
2860            lossy = self.park_lossy,
2861            multi = true
2862        )
2863        .entered();
2864        let mut parked: Vec<park::ParkedSpace> = Vec::with_capacity(spaces.len());
2865        for &(space_id, kind, _, _) in &spaces {
2866            let paged = kind == superfluid_abi::space_kind::PAGED_TOKEN_KV;
2867            let encoding = if paged && lossy { encoding::Q8 } else { encoding::LOSSLESS };
2868            let range = if paged {
2869                linkw::TokenRangeMsg { start: 0, end: covered }
2870            } else {
2871                linkw::TokenRangeMsg { start: covered, end: covered }
2872            };
2873            let (bytes, sizing_gen) = match client.state_sync(linkw::StateSyncReqMsg::ExportSize {
2874                sequence: seq,
2875                space_id,
2876                range,
2877                encoding,
2878            }) {
2879                Ok(linkw::StateSyncOkMsg::ExportSize { required_bytes, sizing_gen }) => {
2880                    (required_bytes, sizing_gen)
2881                }
2882                Ok(_) => return Err("mismatched ExportSize reply".into()),
2883                Err(e) => return Err(format!("export size (space {space_id}): {e}").into()),
2884            };
2885            let buf = client.buf_create(bytes).map_err(|e| ParkRefused::from(e.to_string()))?;
2886            let result = (|| {
2887                let op_kind = if paged && lossy {
2888                    linkw::StateOpKind::Demote {
2889                        sequence: seq,
2890                        space_id,
2891                        range,
2892                        encoding: encoding::Q8,
2893                        buf_id: buf,
2894                        sizing_gen,
2895                    }
2896                } else {
2897                    linkw::StateOpKind::Snapshot {
2898                        sequence: seq,
2899                        space_id,
2900                        boundary_pos: covered,
2901                        buf_id: buf,
2902                        sizing_gen,
2903                    }
2904                };
2905                let op = client.state_op(op_kind).map_err(|e| e.to_string())?;
2906                let done = client.wait_op_done(op).map_err(|e| e.to_string())?;
2907                if done.state != op_state::DONE || done.error != 0 {
2908                    return Err(format!(
2909                        "space {space_id} export terminal state {} error {}",
2910                        done.state, done.error
2911                    ));
2912                }
2913                let take = if done.bytes_moved > 0 { done.bytes_moved.min(bytes) } else { bytes };
2914                client.buf_read(buf, take as usize).map_err(|e| e.to_string())
2915            })();
2916            let _ = client.buf_release(buf);
2917            let sealed = result.map_err(ParkRefused::from)?;
2918            parked.push(park::ParkedSpace {
2919                space_id,
2920                kind,
2921                encoding,
2922                sealed,
2923            });
2924        }
2925        let job = ParkJob {
2926            dir,
2927            session,
2928            covered,
2929            digest,
2930            encoding: if lossy { encoding::Q8 } else { encoding::LOSSLESS },
2931            sealed: Vec::new(),
2932            spaces: Some(parked),
2933        };
2934        match self
2935            .park_tx
2936            .as_ref()
2937            .expect("writer exists with park_dir")
2938            .try_send(ParkTask::Write(Box::new(job)))
2939        {
2940            Ok(()) => {
2941                if self.lanes[i].job.extras.ephemeral {
2942                    self.ephemeral_parked.insert(session);
2943                }
2944                if lossy {
2945                    self.stats.parks_lossy.fetch_add(1, Ordering::Relaxed);
2946                } else {
2947                    self.stats.parks_lossless.fetch_add(1, Ordering::Relaxed);
2948                }
2949            }
2950            Err(std::sync::mpsc::TrySendError::Full(_)) => {
2951                return Err("park writer backlogged; artifact dropped".into());
2952            }
2953            Err(std::sync::mpsc::TrySendError::Disconnected(_)) => {
2954                return Err("park writer gone".into());
2955            }
2956        }
2957        self.lanes[i].parked_lossy = lossy;
2958        Ok(())
2959    }
2960
2961    fn resume_adopt(&mut self, job: &GenerateJob, p: &park::Parked) -> Result<(u64, u64), String> {
2962        if p.covered as usize >= job.stream.len() {
2963            return Err("artifact covers the whole stream (nothing to prefill)".into());
2964        }
2965        if park::stream_digest(&job.stream[..p.covered as usize]) != p.stream_digest {
2966            return Err("stream digest mismatch (rewritten history)".into());
2967        }
2968        let _op_span = tracing::info_span!(
2969            "restore",
2970            session = job.session,
2971            model = self.trace_scope.model.as_str(),
2972            store_ns = self.trace_scope.ns,
2973            covered = p.covered,
2974            spaces = p.spaces.len(),
2975            multi = true
2976        )
2977        .entered();
2978        let want = self.host.parkable_spaces();
2979        for &(id, _, _, _) in &want {
2980            if !p.spaces.iter().any(|s| s.space_id == id) {
2981                return Err(format!("artifact lacks space {id}"));
2982            }
2983        }
2984        let covered_span = self
2985            .stage(&job.stream[..p.covered as usize])
2986            .map_err(|e| e.to_string())?;
2987        let client = self.host.client();
2988        let seq = match client.state_sync(linkw::StateSyncReqMsg::CreateSequence) {
2989            Ok(linkw::StateSyncOkMsg::CreateSequence { sequence }) => sequence,
2990            Ok(_) => return Err("mismatched CreateSequence reply".into()),
2991            Err(e) => return Err(e.to_string()),
2992        };
2993        let free_seq = |client: &mut superfluid_agent::WorkerClient| {
2994            let _ = client.state_sync(linkw::StateSyncReqMsg::FreeSequence { sequence: seq });
2995        };
2996        let mut order: Vec<&park::ParkedSpace> = p.spaces.iter().collect();
2997        order.sort_by_key(|s| (s.kind != superfluid_abi::space_kind::PAGED_TOKEN_KV, s.space_id));
2998        for sp in order {
2999            let buf = match client.buf_create(sp.sealed.len() as u64) {
3000                Ok(b) => b,
3001                Err(e) => {
3002                    free_seq(client);
3003                    return Err(e.to_string());
3004                }
3005            };
3006            let result = (|| {
3007                client.buf_write(buf, &sp.sealed).map_err(|e| e.to_string())?;
3008                let op = client
3009                    .state_op(linkw::StateOpKind::Restore {
3010                        sequence: seq,
3011                        space_id: sp.space_id,
3012                        buf_id: buf,
3013                    })
3014                    .map_err(|e| e.to_string())?;
3015                let done = client.wait_op_done(op).map_err(|e| e.to_string())?;
3016                if done.state != op_state::DONE || done.error != 0 {
3017                    return Err(format!(
3018                        "restore space {} terminal state {} error {}",
3019                        sp.space_id, done.state, done.error
3020                    ));
3021                }
3022                Ok(())
3023            })();
3024            let _ = client.buf_release(buf);
3025            if let Err(e) = result {
3026                free_seq(client);
3027                return Err(e);
3028            }
3029        }
3030        match client.state_sync(linkw::StateSyncReqMsg::AdoptSequence {
3031            sequence: seq,
3032            span: covered_span,
3033        }) {
3034            Ok(linkw::StateSyncOkMsg::AdoptSequence { seed_handle }) => Ok((seed_handle, p.covered)),
3035            Ok(_) => {
3036                free_seq(client);
3037                Err("mismatched AdoptSequence reply".into())
3038            }
3039            Err(e) => {
3040                free_seq(client);
3041                Err(e.to_string())
3042            }
3043        }
3044    }
3045
3046    fn commit_tool_block(
3047        codec: &Arc<dyn TextCodec + Send + Sync>,
3048        store: &Arc<Mutex<SessionStore>>,
3049        tool_lease_ms: u64,
3050        session: u64,
3051        raw: String,
3052        schemas: Option<&crate::codec::ToolSchemas>,
3053    ) -> Result<Option<crate::CommittedEvent>, DaemonError> {
3054        match codec.parse_tool_call_with(&raw, schemas) {
3055            Some((name, args)) => {
3056                let mut store = store.lock().expect("store");
3057                let e = store.commit_tool_use(session, name, args)?;
3058                if tool_lease_ms > 0 {
3059                    let deadline = crate::wal::now_unix_ms() + tool_lease_ms;
3060                    let _ = store.commit_tool_lease(session, e.event_id, deadline);
3061                }
3062                Ok(Some(e))
3063            }
3064            None => Ok(Some(
3065                store.lock().expect("store").commit_tool_parse_failure(session, raw)?,
3066            )),
3067        }
3068    }
3069
3070    fn finish_failed(&mut self, job: &GenerateJob, msg: &str) {
3071        self.drop_ephemeral_artifact(job.session);
3072        self.active.lock().expect("active set").remove(&job.session);
3073        let _ = job.events.send(LaneEvent::Failed(msg.to_string()));
3074    }
3075
3076    fn publish_occupancy(&self) {
3077        self.stats.clear_tick_in_flight();
3078        self.stats
3079            .lanes_active
3080            .store(self.lanes.len() as u64, Ordering::Relaxed);
3081        for (c, q) in self.queues.iter().enumerate() {
3082            self.stats.queue_depth[c].store(q.jobs.len() as u64, Ordering::Relaxed);
3083        }
3084    }
3085
3086    fn tick_once(&mut self, tick_span: &tracing::Span) -> Result<(), DaemonError> {
3087        self.publish_occupancy();
3088        if self.lanes.is_empty() {
3089            return Ok(());
3090        }
3091
3092        let mut locally_done: Vec<u64> = Vec::new();
3093        for i in 0..self.lanes.len() {
3094            let session = self.lanes[i].session;
3095            let is_active = !matches!(self.lanes[i].phase, LanePhase::Finishing);
3096            if is_active && self.cancels.lock().expect("cancels").remove(&session) {
3097                let lane = &mut self.lanes[i];
3098                lane.fin = finish::CANCELLED;
3099                provisional_flush(&self.bus, lane);
3100                let text = lane.utf8.flush();
3101                let e = self.store.lock().expect("store").commit_generated(
3102                    session,
3103                    Vec::new(),
3104                    text,
3105                    lane.seg.current_channel(),
3106                    finish::CANCELLED,
3107                )?;
3108                let _ = lane.job.events.send(LaneEvent::Committed(e));
3109                if lane.admitted {
3110                    lane.phase = LanePhase::Finishing;
3111                } else {
3112                    if lane.seed_handle != 0 {
3113                        let _ = self.host.client().seed_release(lane.seed_handle);
3114                    }
3115                    locally_done.push(lane.tag);
3116                }
3117            }
3118        }
3119        let now_dl = std::time::Instant::now();
3120        for i in 0..self.lanes.len() {
3121            if matches!(self.lanes[i].phase, LanePhase::Finishing) {
3122                continue;
3123            }
3124            let Some(deadline) = self.lanes[i].job.deadline else {
3125                continue;
3126            };
3127            if now_dl <= deadline {
3128                continue;
3129            }
3130            let completion = self.lanes[i].job.completion;
3131            let session = self.lanes[i].session;
3132            if !completion {
3133                provisional_flush(&self.bus, &mut self.lanes[i]);
3134                let text = self.lanes[i].utf8.flush();
3135                let channel = self.lanes[i].seg.current_channel();
3136                let e = self.store.lock().expect("store").commit_generated(
3137                    session,
3138                    Vec::new(),
3139                    text,
3140                    channel,
3141                    finish::LENGTH,
3142                )?;
3143                let _ = self.lanes[i].job.events.send(LaneEvent::Committed(e));
3144                self.lanes[i].fin = finish::LENGTH;
3145            } else {
3146                let _ = self.lanes[i].job.events.send(LaneEvent::Expired);
3147            }
3148            let lane = &mut self.lanes[i];
3149            if lane.admitted {
3150                lane.phase = LanePhase::Finishing;
3151            } else {
3152                if lane.seed_handle != 0 {
3153                    let _ = self.host.client().seed_release(lane.seed_handle);
3154                }
3155                locally_done.push(lane.tag);
3156            }
3157        }
3158        if !locally_done.is_empty() {
3159            let mut locally_retired: Vec<u64> = Vec::new();
3160            let mut active = self.active.lock().expect("active set");
3161            self.lanes.retain(|lane| {
3162                if locally_done.contains(&lane.tag) {
3163                    active.remove(&lane.session);
3164                    let _ = lane.job.events.send(LaneEvent::Done {
3165                        tokens_generated: lane.produced,
3166                        finish: lane.fin,
3167                        warm_prefix: lane.warm_prefix,
3168                        spec: lane.spec,
3169                    });
3170                    locally_retired.push(lane.session);
3171                    false
3172                } else {
3173                    true
3174                }
3175            });
3176            drop(active);
3177            for session in locally_retired {
3178                self.drop_ephemeral_artifact(session);
3179            }
3180        }
3181        self.publish_occupancy();
3182        if self.lanes.is_empty() {
3183            return Ok(());
3184        }
3185
3186        self.tick_count += 1;
3187        self.stats.ticks.fetch_add(1, Ordering::Relaxed);
3188
3189        if self.park_lossy {
3190            for lane in &mut self.lanes {
3191                if lane.preempted
3192                    && !lane.parked
3193                    && lane.job.params.seed != 0
3194                    && !lane.job.extras.seed_drawn
3195                {
3196                    lane.parked = true;
3197                }
3198            }
3199        }
3200        if self.park_lossy && !self.pins.is_empty() && self.pins.enforceable(&self.host) {
3201            for lane in &mut self.lanes {
3202                if matches!(lane.phase, LanePhase::Finishing)
3203                    && !lane.parked
3204                    && lane.job.media.is_empty()
3205                    && !lane.job.completion
3206                    && self.pins.awaits_lease(&self.host, lane.session)
3207                {
3208                    lane.parked = true;
3209                }
3210            }
3211        }
3212
3213        if !self.park_removals_pending.is_empty() {
3214            if let (Some(dir), Some(tx)) = (self.park_dir.clone(), self.park_tx.as_ref()) {
3215                self.park_removals_pending.retain(|&session| {
3216                    matches!(
3217                        tx.try_send(ParkTask::Remove {
3218                            dir: dir.clone(),
3219                            session,
3220                        }),
3221                        Err(std::sync::mpsc::TrySendError::Full(_))
3222                    )
3223                });
3224            } else {
3225                self.park_removals_pending.clear();
3226            }
3227        }
3228
3229        if self.park_dir.is_some() {
3230            for i in 0..self.lanes.len() {
3231                if matches!(self.lanes[i].phase, LanePhase::Finishing)
3232                    && self.lanes[i].admitted
3233                    && !self.lanes[i].parked
3234                    && (!self.lanes[i].job.extras.ephemeral || self.lanes[i].preempted)
3235                {
3236                    self.lanes[i].parked = true;
3237                    if let Err(msg) = self.park_lane(i) {
3238                        eprintln!(
3239                            "[superfluid] park skipped (session {}): {msg}",
3240                            self.lanes[i].session
3241                        );
3242                    }
3243                }
3244            }
3245        }
3246
3247        let rate_cap = |ceiling: u32| match self.prefill_secs_per_token_bulk {
3248            Some(spt) if spt > 0.0 => Some((MAX_PREFILL_TICK_SECS / spt).clamp(512.0, ceiling as f64) as u32),
3249            _ => None,
3250        };
3251        let solo = !self.prefill_budget_fixed
3252            && self.lanes.len() == 1
3253            && self.queued() == 0
3254            && self
3255                .lanes
3256                .iter()
3257                .all(|l| matches!(l.phase, LanePhase::Prefilling));
3258        let solo_budget = if solo { rate_cap(SOLO_PREFILL_BUDGET) } else { None };
3259        let open = self.joiner_waits_on_this_tick();
3260        let tick_prefill_budget = if self
3261            .lanes
3262            .iter()
3263            .any(|l| !l.job.media.is_empty() && matches!(l.phase, LanePhase::Prefilling))
3264        {
3265            PREFILL_BUDGET
3266        } else if open {
3267            // One chunk until two ticks have taught the rate; then the slower of the
3268            // average and the last tick.
3269            let spt = match (self.prefill_secs_per_token_bulk, self.prefill_secs_per_token_last) {
3270                (Some(a), Some(b)) => Some(a.max(b)),
3271                (a, b) => a.or(b),
3272            };
3273            let cap = match spt.filter(|&t| t > 0.0 && self.prefill_rate_samples >= 2) {
3274                Some(spt) => {
3275                    let yields = matches!(
3276                        self.host.capabilities().read().expect("capability record").get("serving", "tick_yield"),
3277                        crate::capabilities::Cap::Yes
3278                    );
3279                    open_prefill_tick_cap(spt, open_prefill_tick_target(yields, self.latency_sensitive()))
3280                }
3281                None => 512,
3282            };
3283            self.prefill_budget.min(cap)
3284        } else if let Some(cap) = solo_budget.filter(|&cap| cap > self.prefill_budget) {
3285            cap
3286        } else {
3287            match self.prefill_secs_per_token_bulk {
3288                Some(spt) if spt > 0.0 => {
3289                    let lone = self.lanes.len() == 1
3290                        && self.queues.iter().all(|q| q.jobs.is_empty())
3291                        && self.lanes.iter().all(|l| {
3292                            Self::prefill_end_of(&self.speculation, l).saturating_sub(l.prefilled)
3293                                > prefill_tick_cap(spt, true) as u64
3294                        });
3295                    self.prefill_budget.min(prefill_tick_cap(spt, lone))
3296                }
3297                _ => self.prefill_budget,
3298            }
3299        };
3300        let mut plan = linkw::TickPlanMsg {
3301            plan_seq: {
3302                self.plan_seq += 1;
3303                self.plan_seq
3304            },
3305            flags: 0,
3306            prefill_token_budget: tick_prefill_budget,
3307            max_decode_lanes: self.max_lanes as u32,
3308            admits: vec![],
3309            commits: vec![],
3310            prefills: vec![],
3311            decodes: vec![],
3312            retires: vec![],
3313            shed_policy: linkw::ShedPolicyMsg {
3314                victim_lanes: vec![],
3315                evictable_cache_classes: u64::MAX,
3316                protected_quota_bytes: 0,
3317                max_evict_bytes: u64::MAX,
3318            },
3319        };
3320
3321        let prefill_demand_lanes = self
3322            .lanes
3323            .iter()
3324            .filter(|l| {
3325                !matches!(l.phase, LanePhase::Finishing)
3326                    && l.prefilled < Self::prefill_end_of(&self.speculation, l)
3327            })
3328            .count()
3329            + self.queues.iter().map(|q| q.jobs.len()).sum::<usize>();
3330        self.plan_prefix_checkpoints();
3331        let mut budget = PrefillBudget::new(tick_prefill_budget as u64);
3332        let mut retiring: Vec<u64> = Vec::new();
3333        let starvation = self.starvation_ticks;
3334        let mut order: Vec<usize> = (0..self.lanes.len()).collect();
3335        order.sort_by_key(|&i| {
3336            let l = &self.lanes[i];
3337            let remaining = Self::prefill_end_of(&self.speculation, l).saturating_sub(l.prefilled);
3338            prefill_rank(l.job.class, remaining, l.tag, l.prefill_starved_ticks, starvation)
3339        });
3340        let mut decode_candidates: Vec<usize> = Vec::new();
3341        let speculation = Arc::clone(&self.speculation);
3342        for &li in &order {
3343            let lane = &mut self.lanes[li];
3344            if matches!(lane.phase, LanePhase::Finishing) {
3345                plan.retires.push(linkw::LaneRetireMsg {
3346                    lane_tag: lane.tag,
3347                    publish_to_cache: !lane.parked_lossy
3348                        && lane.job.media.is_empty()
3349                        && !lane.job.completion,
3350                });
3351                retiring.push(lane.tag);
3352                continue;
3353            }
3354            if !lane.admitted {
3355                let strategy_slot = Self::strategy_slot_for(&speculation, &lane.job);
3356                plan.admits.push(linkw::LaneAdmitMsg {
3357                    lane_tag: lane.tag,
3358                    prompt: lane.span,
3359                    seed_handle: lane.seed_handle,
3360                    sampling: if lane.job.params.temperature == 0.0 { 0 } else { 1 },
3361                    params: linkw::SamplingParamsMsg {
3362                        temperature: lane.job.params.temperature,
3363                        top_p: lane.job.params.top_p,
3364                        min_p: lane.job.params.min_p,
3365                        top_k: lane.job.params.top_k,
3366                        freq_penalty: lane.job.extras.frequency_penalty,
3367                        presence_penalty: lane.job.extras.presence_penalty,
3368                        repeat_penalty: lane.job.extras.repeat_penalty,
3369                        flags: if lane.job.extras.ignore_eos {
3370                            superfluid_abi::sampling_flag::IGNORE_EOS
3371                        } else {
3372                            0
3373                        },
3374                    },
3375                    rng_counter_base: lane
3376                        .job
3377                        .params
3378                        .seed
3379                        .wrapping_add(lane.job.stream.len() as u64),
3380                    grammar_handle: lane.job.grammar_handle,
3381                    grammar_replay: if lane.job.grammar_handle != 0 { lane.produced_before } else { 0 },
3382                    decode_replay: Self::decode_replay_for(strategy_slot, lane.produced_before),
3383                    logit_bias_handle: lane.job.extras.logit_bias_handle,
3384                    want_logprobs: lane.job.extras.want_logprobs,
3385                    top_logprobs: lane.job.extras.top_logprobs,
3386                    strategy_slot,
3387                    determinism_class: 0,
3388                    minimum_exactness: 0,
3389                    allow_approximate: false,
3390                    required_cert_id: 0,
3391                    host_sampler_identity: [0; 32],
3392                });
3393            }
3394            let total = Self::prefill_end_of(&speculation, lane);
3395            if lane.prefilled < total {
3396                let mut chunk = budget.offer_to(total - lane.prefilled, lane.job.class) as u32;
3397                if lane.checkpoint_at > lane.prefilled {
3398                    chunk = chunk.min((lane.checkpoint_at - lane.prefilled) as u32);
3399                }
3400                for &(off, _, n) in &lane.job.media {
3401                    let (rs, re) = (off, off + n as u64);
3402                    let (cs, ce) = (lane.prefilled, lane.prefilled + chunk as u64);
3403                    if rs < ce && re > ce {
3404                        chunk = if rs > cs { (rs - cs) as u32 } else { 0 };
3405                    }
3406                }
3407                if chunk > 0 {
3408                    lane.prefill_starved_ticks = 0;
3409                }
3410                if chunk > 0 {
3411                    plan.prefills.push(linkw::LanePrefillMsg {
3412                        lane_tag: lane.tag,
3413                        token_offset: lane.prefilled as u32,
3414                        token_count: chunk,
3415                    });
3416                    lane.prefilled += chunk as u64;
3417                    lane.checkpoint_due = lane.checkpoint_at > 0 && lane.prefilled == lane.checkpoint_at;
3418                    budget.take(chunk as u64, lane.prefilled >= total);
3419                    if lane.prefilled >= total && lane.job.class <= crate::qos::INLINE_COMPLETION {
3420                        budget.hold_below(lane.job.class);
3421                    }
3422                }
3423            }
3424            if lane.prefilled >= total {
3425                lane.phase = LanePhase::Decoding;
3426                if lane.job.max_tokens > lane.produced {
3427                    decode_candidates.push(li);
3428                }
3429            }
3430        }
3431        // A lane counts a starved tick only when the prefill went to a younger lane.
3432        let youngest_served = plan.prefills.iter().map(|p| p.lane_tag).max();
3433        for lane in self.lanes.iter_mut() {
3434            let waiting = lane.prefilled < Self::prefill_end_of(&self.speculation, lane)
3435                && !matches!(lane.phase, LanePhase::Finishing)
3436                && !plan.prefills.iter().any(|p| p.lane_tag == lane.tag);
3437            lane.prefill_starved_ticks = starved_after(lane.prefill_starved_ticks, waiting, youngest_served, lane.tag);
3438        }
3439        let divisor = self.background_tick_divisor.load(Ordering::Relaxed).max(1);
3440        let tick_no = self.tick_count;
3441        decode_candidates.sort_by_key(|&i| {
3442            let l = &self.lanes[i];
3443            let starving = l.starved_ticks >= starvation;
3444            (!starving, l.job.class, l.tag)
3445        });
3446        let planned_prefill_tokens: f64 = plan.prefills.iter().map(|p| p.token_count as f64).sum();
3447        let finishing = self.lanes.iter().any(|l| matches!(l.phase, LanePhase::Finishing));
3448        let shape = self.tick_shape;
3449        let target_secs = tick_target(&shape, self.tick_target_secs, self.lanes.len(), self.max_lanes, finishing);
3450        self.stats.tick_target_live_ms.store((target_secs * 1000.0).round() as u64, Ordering::Relaxed);
3451        let grant_cap: u16 = match (plan.prefills.len(), self.decode_round_secs) {
3452            (p, Some(c)) if c > 0.0 && !decode_candidates.is_empty() => {
3453                if p > 0 {
3454                    let prefill_secs = self.prefill_secs_per_token.map_or(0.0, |s| s * planned_prefill_tokens);
3455                    let interactive =
3456                        decode_candidates.iter().any(|&i| self.lanes[i].job.class <= crate::qos::INLINE_COMPLETION);
3457                    let share =
3458                        if interactive { DECODE_SHARE_BESIDE_PREFILL_INTERACTIVE } else { DECODE_SHARE_BESIDE_PREFILL };
3459                    ((prefill_secs * share / c) as u32)
3460                        .clamp(shape.beside_prefill as u32, TICK_DECODE_MAX as u32) as u16
3461                } else {
3462                    let rounds = ((target_secs / c) as u32).clamp(shape.decode_floor as u32, TICK_DECODE_MAX as u32);
3463                    let first_out = decode_candidates
3464                        .iter()
3465                        .map(|&i| self.lanes[i].job.max_tokens.saturating_sub(self.lanes[i].produced))
3466                        .min()
3467                        .unwrap_or(rounds);
3468                    let rounds = rounds.min(first_out.max(shape.beside_prefill as u32));
3469                    if self.just_retired {
3470                        rounds.min(shape.beside_prefill as u32) as u16
3471                    } else {
3472                        rounds as u16
3473                    }
3474                }
3475            }
3476            _ => shape.decode_floor,
3477        };
3478        self.stats.decode_grant_live.store(grant_cap as u64, Ordering::Relaxed);
3479        let mut decode_left = if self.tick_decode_budget_fixed {
3480            self.tick_decode_budget
3481        } else {
3482            (self.max_lanes as u32).saturating_mul(grant_cap as u32)
3483        };
3484        let retire_only = plan.prefills.is_empty() && !plan.retires.is_empty();
3485        if retire_only {
3486            decode_candidates.clear();
3487        }
3488        self.just_retired = retire_only;
3489        for &li in &decode_candidates {
3490            let lane = &mut self.lanes[li];
3491            let background = lane.job.class >= crate::qos::FOREGROUND_AGENT;
3492            let throttled = background && divisor > 1 && !tick_no.is_multiple_of(divisor as u64);
3493            let want = if throttled {
3494                0
3495            } else {
3496                (lane.job.max_tokens - lane.produced)
3497                    .min(grant_cap as u32)
3498                    .min(decode_left) as u16
3499            };
3500            if want > 0 {
3501                if lane.starved_ticks >= starvation {
3502                    self.stats.starvation_grants.fetch_add(1, Ordering::Relaxed);
3503                }
3504                lane.starved_ticks = 0;
3505                decode_left -= want as u32;
3506                let overshoot = (lane.job.max_tokens - lane.produced - want as u32).min(16) as u16;
3507                plan.decodes.push(linkw::LaneDecodeMsg {
3508                    lane_tag: lane.tag,
3509                    max_new_tokens: want,
3510                    overshoot,
3511                });
3512            } else {
3513                lane.starved_ticks = lane.starved_ticks.saturating_add(1);
3514            }
3515        }
3516
3517        let wants_partial = self.lanes.iter().any(|l| {
3518            !matches!(l.phase, LanePhase::Finishing)
3519                && (l.job.extras.wants_deltas
3520                    || l.job.completion
3521                    || self.bus.has_provisional_subscriber(l.session))
3522        });
3523        if wants_partial {
3524            plan.flags |= superfluid_abi::tick_flags::PARTIAL_EMITS;
3525        }
3526        let mut plan_prefill_tokens: u32 = plan.prefills.iter().map(|p| p.token_count).sum();
3527        let prefill_by_tag: HashMap<u64, u32> =
3528            plan.prefills.iter().map(|p| (p.lane_tag, p.token_count)).collect();
3529        let plan_decode_rounds: u32 =
3530            plan.decodes.iter().map(|d| d.max_new_tokens as u32).max().unwrap_or(0);
3531        let plan_prefill_lanes_planned = plan.prefills.len();
3532        let plan_prefill_lanes = plan.prefills.len().max(prefill_demand_lanes);
3533        let plan_decode_lanes = plan.decodes.len();
3534        if !tick_span.is_disabled() {
3535            let planned: HashSet<u64> = plan
3536                .admits
3537                .iter()
3538                .map(|a| a.lane_tag)
3539                .chain(plan.prefills.iter().map(|p| p.lane_tag))
3540                .chain(plan.decodes.iter().map(|d| d.lane_tag))
3541                .chain(plan.retires.iter().map(|r| r.lane_tag))
3542                .collect();
3543            tick_span.record("lanes", planned.len());
3544            for lane in self.lanes.iter().filter(|l| planned.contains(&l.tag)) {
3545                tick_span.follows_from(&lane.job.trace);
3546            }
3547        }
3548        let tick_started = std::time::Instant::now();
3549        self.stats.set_tick_in_flight(TickInFlight {
3550            since_ms: uptime_ms().max(1),
3551            prefill_tokens: plan_prefill_tokens as u64,
3552            prefill_lanes: plan_prefill_lanes_planned as u64,
3553            decode_lanes: plan_decode_lanes as u64,
3554            admits: plan.admits.len() as u64,
3555            retires: plan.retires.len() as u64,
3556            target_ms: (target_secs * 1000.0).round() as u64,
3557        });
3558        let mut deferred: Vec<u64> = Vec::new();
3559        let ev = {
3560            let Scheduler {
3561                host,
3562                lanes,
3563                bus,
3564                codec,
3565                pins,
3566                stats,
3567                ..
3568            } = self;
3569            let by_tag: std::collections::HashMap<u64, usize> =
3570                lanes.iter().enumerate().map(|(i, l)| (l.tag, i)).collect();
3571            // The engine answers NeedsReplan when this tick's admissions do
3572            // not fit (sequence slots or cells it cannot free); it has run
3573            // nothing. Admit fewer rather than fail every lane.
3574            let fewer = (!plan.admits.is_empty()
3575                && plan.admits.iter().all(|a| {
3576                    by_tag.get(&a.lane_tag).is_some_and(|&i| lanes[i].job.media.is_empty())
3577                }))
3578            .then(|| plan.clone());
3579            let mut on_emit = |em: linkw::TickEmitMsg| {
3580                let Some(&i) = by_tag.get(&em.lane_tag) else { return };
3581                let lane = &mut lanes[i];
3582                lane.frame_fed += em.tokens.len();
3583                stats
3584                    .decode_tokens
3585                    .fetch_add(em.tokens.len() as u64, Ordering::Relaxed);
3586                if !lane.ttft_observed && lane.ttft_frame_ms.is_none() && !em.tokens.is_empty() {
3587                    lane.ttft_frame_ms = Some(lane.admitted_at.elapsed().as_millis() as u64);
3588                }
3589                if lane.job.completion {
3590                    if !lane.emit_dead
3591                        && lane.job.events.send(LaneEvent::CompletionTokens(em.tokens)).is_err()
3592                    {
3593                        lane.emit_dead = true;
3594                    }
3595                } else {
3596                    provisional_ingest(codec, bus, lane, &em.tokens);
3597                }
3598            };
3599            let retry = pins.any_held().then(|| plan.clone());
3600            let needs_replan = |r: &Result<linkw::TickEventsMsg, superfluid_agent::AgentError>| {
3601                matches!(r, Err(superfluid_agent::AgentError::Rejected(s)) if *s == superfluid_abi::Status::NeedsReplan as i32)
3602            };
3603            let mut res = match (host.client().tick_streaming(plan, &mut on_emit), retry) {
3604                (Err(superfluid_agent::AgentError::Rejected(status)), Some(plan))
3605                    if status == superfluid_abi::Status::NeedsReplan as i32
3606                        && pins.yield_all(host, stats) =>
3607                {
3608                    host.client().tick_streaming(plan, &mut on_emit)
3609                }
3610                (res, _) => res,
3611            };
3612            if let Some(full) = fewer.filter(|_| needs_replan(&res)) {
3613                let mut keep = full.admits.len();
3614                while keep > 0 && needs_replan(&res) {
3615                    keep /= 2;
3616                    res = host.client().tick_streaming(with_first_admits(&full, keep), &mut on_emit);
3617                }
3618                if res.is_ok() {
3619                    deferred = full.admits[keep..].iter().map(|a| a.lane_tag).collect();
3620                }
3621            }
3622            stats.clear_tick_in_flight();
3623            res?
3624        };
3625        if !deferred.is_empty() {
3626            plan_prefill_tokens = plan_prefill_tokens
3627                .saturating_sub(deferred.iter().filter_map(|t| prefill_by_tag.get(t)).sum());
3628            self.stats.admits_deferred.fetch_add(deferred.len() as u64, Ordering::Relaxed);
3629            self.defer_admits(&deferred);
3630        }
3631        let tick_secs = tick_started.elapsed().as_secs_f64().max(1e-3);
3632        // Prefill the engine shed (a tick that yielded between steps) is
3633        // planned again from where it stopped.
3634        for e in ev.shed.entries.iter().filter(|e| e.kind == superfluid_abi::shed_kind::PREFILL_CHUNK) {
3635            if let Some(lane) = self.lanes.iter_mut().find(|l| l.tag == e.lane_tag) {
3636                lane.prefilled = lane.prefilled.saturating_sub(e.tokens);
3637                if matches!(lane.phase, LanePhase::Decoding) {
3638                    lane.phase = LanePhase::Prefilling;
3639                }
3640                if lane.checkpoint_due && lane.prefilled != lane.checkpoint_at {
3641                    lane.checkpoint_due = false;
3642                }
3643            }
3644        }
3645        let prefill_tokens_ran: u32 = plan_prefill_tokens
3646            .saturating_sub(
3647                ev.admit_results
3648                    .iter()
3649                    .filter(|a| a.status != superfluid_abi::admit_status::ADMITTED)
3650                    .filter_map(|a| prefill_by_tag.get(&a.lane_tag))
3651                    .sum(),
3652            )
3653            .saturating_sub(
3654                ev.faults
3655                    .iter()
3656                    .filter(|f| f.code == superfluid_abi::fault::MEDIA)
3657                    .filter_map(|f| prefill_by_tag.get(&f.lane_tag))
3658                    .sum(),
3659            )
3660            .saturating_sub(
3661                ev.shed
3662                    .entries
3663                    .iter()
3664                    .filter(|e| e.kind == superfluid_abi::shed_kind::PREFILL_CHUNK)
3665                    .map(|e| e.tokens.min(u32::MAX as u64) as u32)
3666                    .sum(),
3667            );
3668        let tick_rounds_exec: u32 = ev
3669            .emits
3670            .iter()
3671            .filter(|e| e.n_tokens > 0)
3672            .map(|e| e.n_tokens.saturating_sub(e.spec.accepted).max(1))
3673            .max()
3674            .unwrap_or(0);
3675        self.stats
3676            .prefill_tokens
3677            .fetch_add(prefill_tokens_ran as u64, Ordering::Relaxed);
3678        let phase_split = plan_prefill_tokens == 0 || ev.timings.prefill_ns > 0;
3679        let (prefill_secs_meas, decode_secs_meas) = if phase_split {
3680            (ev.timings.prefill_ns as f64 / 1e9, ev.timings.decode_ns as f64 / 1e9)
3681        } else {
3682            let decode_est = self.decode_round_secs.unwrap_or(0.0) * tick_rounds_exec as f64;
3683            ((tick_secs - decode_est).max(1e-3), decode_est)
3684        };
3685        // A tick the daemon ended early ran fewer steps than planned: the
3686        // rates are of what ran.
3687        let decode_rounds_ran = if tick_rounds_exec > 0 { tick_rounds_exec.min(plan_decode_rounds) } else { plan_decode_rounds };
3688        if decode_rounds_ran > 0 && phase_split && decode_secs_meas > 0.0 {
3689            let sample = decode_secs_meas / decode_rounds_ran as f64;
3690            self.decode_round_secs = Some(match self.decode_round_secs {
3691                Some(c) => 0.7 * c + 0.3 * sample,
3692                None => sample,
3693            });
3694        } else if plan_prefill_tokens == 0 && decode_rounds_ran > 0 {
3695            let sample = tick_secs / decode_rounds_ran as f64;
3696            self.decode_round_secs = Some(match self.decode_round_secs {
3697                Some(c) => 0.7 * c + 0.3 * sample,
3698                None => sample,
3699            });
3700        }
3701        if prefill_tokens_ran >= PREFILL_RATE_MIN_TOKENS && phase_split {
3702            let sample = prefill_secs_meas.max(1e-3) / prefill_tokens_ran as f64;
3703            self.prefill_secs_per_token = Some(match self.prefill_secs_per_token {
3704                Some(c) => 0.7 * c + 0.3 * sample,
3705                None => sample,
3706            });
3707            if prefill_tokens_ran >= 512 {
3708                self.prefill_secs_per_token_last = Some(sample);
3709                self.prefill_rate_samples = self.prefill_rate_samples.saturating_add(1);
3710                self.prefill_secs_per_token_bulk = Some(match self.prefill_secs_per_token_bulk {
3711                    Some(c) => 0.7 * c + 0.3 * sample,
3712                    None => sample,
3713                });
3714            }
3715        }
3716        if !self.prefill_budget_fixed && prefill_tokens_ran > 0 {
3717            // What the whole planned prefill would have taken at the rate of
3718            // what ran.
3719            let prefill_secs = prefill_secs_meas.max(1e-3) * plan_prefill_tokens as f64 / prefill_tokens_ran as f64;
3720            let share = plan_prefill_lanes as f64
3721                / (plan_prefill_lanes as f64 + plan_decode_lanes as f64);
3722            let factor = (self.tick_target_secs * share / prefill_secs).clamp(0.5, 2.0);
3723            self.prefill_budget =
3724                ((self.prefill_budget as f64 * factor) as u32).clamp(256, PREFILL_BUDGET);
3725        }
3726        if std::env::var_os("SUPERFLUID_DEBUG_TICK").is_some() {
3727            eprintln!(
3728                "[tick] secs={:.3} pf_secs={:.3} dec_secs={:.3} prefill_tok={} rounds={} P_eff={} P_plan={} D={} round_secs={:?} pf_secs_per_tok={:?} budget={} grant={}",
3729                tick_secs, prefill_secs_meas, decode_secs_meas, plan_prefill_tokens, plan_decode_rounds, plan_prefill_lanes, plan_prefill_lanes_planned,
3730                plan_decode_lanes, self.decode_round_secs, self.prefill_secs_per_token, self.prefill_budget, grant_cap
3731            );
3732        }
3733        self.stats
3734            .prefill_budget_live
3735            .store(self.prefill_budget as u64, Ordering::Relaxed);
3736        if !self.pins.is_empty() {
3737            let retired: Vec<(u64, Vec<u32>)> = self
3738                .lanes
3739                .iter()
3740                .filter(|l| retiring.contains(&l.tag) && !l.job.completion)
3741                .map(|l| (l.session, l.all_tokens.clone()))
3742                .collect();
3743            self.pins
3744                .after_tick(&mut self.host, &self.stats, self.tick_count, &retired);
3745        }
3746        self.publish_prefix_checkpoints();
3747        self.relieve_pressure(&ev.mem, false);
3748        self.stats
3749            .pool_blocks_used
3750            .store(ev.mem.pool_blocks_used, Ordering::Relaxed);
3751        self.stats
3752            .pool_blocks_total
3753            .store(ev.mem.pool_blocks_total, Ordering::Relaxed);
3754
3755        let refused: HashMap<u64, u32> = ev
3756            .admit_results
3757            .iter()
3758            .filter(|r| r.status == superfluid_abi::admit_status::REJECTED)
3759            .map(|r| (r.lane_tag, r.reject_code))
3760            .collect();
3761        if !refused.is_empty() {
3762            self.drop_refused(&refused);
3763        }
3764
3765        for f in &ev.faults {
3766            if let Some(lane) = self.lanes.iter_mut().find(|l| l.tag == f.lane_tag) {
3767                if lane.fault_code.is_some() {
3768                    continue;
3769                }
3770                lane.fault_code = Some(f.code);
3771                lane.fin = finish::ERROR;
3772                provisional_flush(&self.bus, lane);
3773                let text = lane.utf8.flush();
3774                let e = self.store.lock().expect("store").commit_generated(
3775                    lane.session,
3776                    Vec::new(),
3777                    text,
3778                    lane.seg.current_channel(),
3779                    finish::ERROR,
3780                )?;
3781                let _ = lane.job.events.send(LaneEvent::Committed(e));
3782                lane.phase = LanePhase::Finishing;
3783                tracing::warn!(session = lane.session, lane = lane.tag, code = f.code, "engine lane fault");
3784            }
3785        }
3786
3787        #[allow(clippy::type_complexity)]
3788        let mut logprobs_by_lane: HashMap<u64, Vec<(f32, Vec<(u32, f32)>)>> = HashMap::new();
3789        for lp in &ev.logprobs {
3790            let top: Vec<(u32, f32)> = lp
3791                .top_ids
3792                .iter()
3793                .zip(&lp.top_logprob_bits)
3794                .map(|(id, b)| (*id, f32::from_bits(*b)))
3795                .collect();
3796            logprobs_by_lane
3797                .entry(lp.lane_tag)
3798                .or_default()
3799                .push((f32::from_bits(lp.logprob_bits), top));
3800        }
3801        let mut emitted: HashMap<u64, (Vec<u32>, u32)> = HashMap::new();
3802        for emit in &ev.emits {
3803            if emit.spec.proposed > 0 {
3804                self.stats
3805                    .spec_proposed
3806                    .fetch_add(emit.spec.proposed as u64, Ordering::Relaxed);
3807                self.stats
3808                    .spec_accepted
3809                    .fetch_add(emit.spec.accepted as u64, Ordering::Relaxed);
3810                if let Some(lane) = self.lanes.iter_mut().find(|l| l.tag == emit.lane_tag) {
3811                    lane.spec.0 += emit.spec.proposed as u64;
3812                    lane.spec.1 += emit.spec.accepted as u64;
3813                }
3814            }
3815            if emit.n_tokens > 0 {
3816                let tokens = self.host.client().read_tokens(&emit.token_ref)?;
3817                emitted.insert(emit.lane_tag, (tokens, emit.finish));
3818            } else if emit.finish != finish::NONE {
3819                emitted.insert(emit.lane_tag, (Vec::new(), emit.finish));
3820            }
3821        }
3822        let tick_decode_tokens: u64 = emitted.values().map(|(t, _)| t.len() as u64).sum();
3823        let ttft_tail_secs = if tick_rounds_exec > 1 {
3824            decode_secs_meas * (tick_rounds_exec as f64 - 1.0) / tick_rounds_exec as f64
3825        } else {
3826            0.0
3827        };
3828        if prefill_tokens_ran > 0 || tick_decode_tokens > 0 {
3829            self.stats
3830                .tick_wall_ms_live
3831                .store((tick_secs * 1e3) as u64, Ordering::Relaxed);
3832            self.stats.prefill_rate_live.store(
3833                per_sec(prefill_tokens_ran as u64, prefill_secs_meas),
3834                Ordering::Relaxed,
3835            );
3836            let decode_secs_for_rate =
3837                if decode_secs_meas > 0.0 { decode_secs_meas } else { tick_secs };
3838            self.stats.decode_rate_live.store(
3839                per_sec(tick_decode_tokens, decode_secs_for_rate),
3840                Ordering::Relaxed,
3841            );
3842            self.stats.work_ticks.fetch_add(1, Ordering::Relaxed);
3843        }
3844        let terminators = self.codec.turn_terminators();
3845        for lane in &mut self.lanes {
3846            lane.admitted = true;
3847            let Some((tokens, fin)) = emitted.remove(&lane.tag) else {
3848                continue;
3849            };
3850            if tokens.is_empty() {
3851                lane.fin = fin;
3852                if !lane.job.completion {
3853                    provisional_flush(&self.bus, lane);
3854                    let text = lane.utf8.flush();
3855                    let e = self.store.lock().expect("store").commit_generated(
3856                        lane.session,
3857                        Vec::new(),
3858                        text,
3859                        lane.seg.current_channel(),
3860                        fin,
3861                    )?;
3862                    let _ = lane.job.events.send(LaneEvent::Committed(e));
3863                }
3864                lane.phase = LanePhase::Finishing;
3865                continue;
3866            }
3867            lane.all_tokens.extend_from_slice(&tokens);
3868            lane.produced += tokens.len() as u32;
3869            lane.fin = fin;
3870            let frame_fed = std::mem::take(&mut lane.frame_fed).min(tokens.len());
3871            if !lane.job.completion {
3872                provisional_ingest(&self.codec, &self.bus, lane, &tokens[frame_fed..]);
3873            }
3874            if lane.fin == finish::NONE
3875                && !lane.job.extras.ignore_eos
3876                && !terminators.is_empty()
3877                && tokens.iter().any(|t| terminators.contains(t))
3878            {
3879                lane.fin = finish::EOS;
3880            }
3881            if lane.job.extras.want_logprobs {
3882                if let Some(recs) = logprobs_by_lane.remove(&lane.tag) {
3883                    if !recs.is_empty() {
3884                        let items: Vec<TokenLogprob> = tokens
3885                            .iter()
3886                            .zip(recs)
3887                            .map(|(t, (lp, top))| TokenLogprob { token: *t, logprob: lp, top })
3888                            .collect();
3889                        let _ = lane.job.events.send(LaneEvent::TokenLogprobs(items));
3890                    }
3891                }
3892            }
3893            self.stats
3894                .decode_tokens
3895                .fetch_add((tokens.len() - frame_fed) as u64, Ordering::Relaxed);
3896            if !lane.ttft_observed && !tokens.is_empty() {
3897                lane.ttft_observed = true;
3898                let ms = lane.ttft_frame_ms.take().unwrap_or_else(|| {
3899                    let elapsed = lane.admitted_at.elapsed().as_secs_f64();
3900                    ((elapsed - ttft_tail_secs).max(0.0) * 1e3) as u64
3901                });
3902                self.stats.observe_ttft(ms);
3903            }
3904            let terminal_now = lane.fin != finish::NONE || lane.produced >= lane.job.max_tokens;
3905            if lane.job.completion {
3906                let leftover = tokens[frame_fed..].to_vec();
3907                if lane.emit_dead
3908                    || (!leftover.is_empty()
3909                        && lane.job.events.send(LaneEvent::CompletionTokens(leftover)).is_err())
3910                {
3911                    lane.phase = LanePhase::Finishing;
3912                    continue;
3913                }
3914                if terminal_now {
3915                    lane.phase = LanePhase::Finishing;
3916                }
3917                continue;
3918            }
3919            if terminal_now {
3920                provisional_flush(&self.bus, lane);
3921            }
3922            let runs = lane.seg.split(&tokens);
3923            let n_runs = runs.len().max(1);
3924            let mut dead = false;
3925            for (i, run) in runs.into_iter().enumerate() {
3926                let last = i + 1 == n_runs;
3927                let mut text = String::new();
3928                for &t in &run.text {
3929                    text.push_str(&lane.utf8.push(&self.codec.token_bytes(t)));
3930                }
3931                if last && terminal_now {
3932                    text.push_str(&lane.utf8.flush());
3933                }
3934                let channel = run.channel;
3935                let tool_closed = run.closes && channel == crate::wal::channel::TOOL_CALL;
3936                if channel == crate::wal::channel::TOOL_CALL {
3937                    lane.tool_buf.push_str(&text);
3938                }
3939                let e = self.store.lock().expect("store").commit_generated(
3940                    lane.session,
3941                    run.span,
3942                    text,
3943                    channel,
3944                    if last { fin } else { finish::NONE },
3945                )?;
3946                let tool_event = if tool_closed {
3947                    let raw = std::mem::take(&mut lane.tool_buf);
3948                    Self::commit_tool_block(
3949                        &self.codec,
3950                        &self.store,
3951                        self.tool_lease_ms,
3952                        lane.session,
3953                        raw,
3954                        lane.job.extras.tool_schemas.as_deref(),
3955                    )?
3956                } else {
3957                    None
3958                };
3959                if lane.job.events.send(LaneEvent::Committed(e)).is_err() {
3960                    dead = true;
3961                    break;
3962                }
3963                if let Some(te) = tool_event {
3964                    if lane.job.events.send(LaneEvent::Committed(te)).is_err() {
3965                        dead = true;
3966                        break;
3967                    }
3968                }
3969            }
3970            if dead {
3971                lane.fin = finish::CANCELLED;
3972                provisional_flush(&self.bus, lane);
3973                let text = lane.utf8.flush();
3974                let _ = self.store.lock().expect("store").commit_generated(
3975                    lane.session,
3976                    Vec::new(),
3977                    text,
3978                    lane.seg.current_channel(),
3979                    finish::CANCELLED,
3980                )?;
3981                lane.phase = LanePhase::Finishing;
3982                continue;
3983            }
3984            if terminal_now {
3985                if !lane.tool_buf.is_empty() {
3986                    let raw = std::mem::take(&mut lane.tool_buf);
3987                    let te = Self::commit_tool_block(
3988                        &self.codec,
3989                        &self.store,
3990                        self.tool_lease_ms,
3991                        lane.session,
3992                        raw,
3993                        lane.job.extras.tool_schemas.as_deref(),
3994                    )?;
3995                    if let Some(te) = te {
3996                        let _ = lane.job.events.send(LaneEvent::Committed(te));
3997                    }
3998                }
3999                lane.phase = LanePhase::Finishing;
4000            }
4001        }
4002
4003        if !retiring.is_empty() {
4004            let mut i = 0;
4005            let mut requeue: Vec<GenerateJob> = Vec::new();
4006            while i < self.lanes.len() {
4007                if !retiring.contains(&self.lanes[i].tag) {
4008                    i += 1;
4009                    continue;
4010                }
4011                let lane = self.lanes.remove(i);
4012                if lane.preempted && lane.fin == finish::NONE {
4013                    let mut job = lane.job;
4014                    job.stream = lane.all_tokens;
4015                    job.max_tokens -= lane.produced;
4016                    job.carry = Some(LaneCarry {
4017                        seg: lane.seg,
4018                        utf8: lane.utf8,
4019                        disp_seg: lane.disp_seg,
4020                        disp_utf8: lane.disp_utf8,
4021                        tool_buf: lane.tool_buf,
4022                        produced_before: lane.produced_before + lane.produced,
4023                        spec_before: lane.spec,
4024                        warm_prefix_first: Some(lane.warm_prefix_first.unwrap_or(lane.warm_prefix)),
4025                    });
4026                    requeue.push(job);
4027                } else {
4028                    self.drop_ephemeral_artifact(lane.session);
4029                    self.active.lock().expect("active set").remove(&lane.session);
4030                    if let Some(code) = lane.fault_code {
4031                        let _ = lane
4032                            .job
4033                            .events
4034                            .send(LaneEvent::Failed(format!("engine lane fault (code {code})")));
4035                    } else {
4036                        let _ = lane.job.events.send(LaneEvent::Done {
4037                            tokens_generated: lane.produced_before + lane.produced,
4038                            finish: lane.fin,
4039                            warm_prefix: lane.warm_prefix_first.unwrap_or(lane.warm_prefix),
4040                            spec: lane.spec,
4041                        });
4042                    }
4043                }
4044            }
4045            for job in requeue {
4046                let c = (job.class as usize).min(3);
4047                self.queues[c].jobs.push_front(job);
4048            }
4049        }
4050        self.publish_occupancy();
4051        Ok(())
4052    }
4053}
4054
4055fn draft_weights_identity(path: &str) -> Result<(u64, u64), DaemonError> {
4056    use std::io::Read;
4057    const SEED: u64 = 0xCBF2_9CE4_8422_2325;
4058    let mut f = std::fs::File::open(path).map_err(|_| {
4059        DaemonError::Config(
4060            "draft-model weights are not readable by the daemon; an exact draft identity \
4061             requires the daemon to read the weights file",
4062        )
4063    })?;
4064    let size = f.metadata().map(|m| m.len()).unwrap_or(0);
4065    let mut h = SEED;
4066    let mut buf = vec![0u8; 1 << 20];
4067    loop {
4068        match f.read(&mut buf) {
4069            Ok(0) => break,
4070            Ok(n) => h = superfluid_fingerprint::fnv1a64(&buf[..n], h),
4071            Err(_) => {
4072                return Err(DaemonError::Config(
4073                    "failed reading draft-model weights for identity",
4074                ))
4075            }
4076        }
4077    }
4078    Ok((h, size))
4079}
4080
4081pub fn run_completion(
4082    tx: &Sender<SchedCmd>,
4083    stream: Vec<u32>,
4084    params: GenParams,
4085    max_tokens: u32,
4086    deadline_ms: u64,
4087    grammar_handle: u32,
4088) -> Result<(Vec<u32>, u32, bool), DaemonError> {
4089    run_completion_streaming(
4090        tx,
4091        stream,
4092        params,
4093        max_tokens,
4094        deadline_ms,
4095        GenExtras { grammar_handle, ..GenExtras::default() },
4096        |_| std::ops::ControlFlow::Continue(()),
4097    )
4098}
4099
4100pub fn run_completion_streaming(
4101    tx: &Sender<SchedCmd>,
4102    stream: Vec<u32>,
4103    params: GenParams,
4104    max_tokens: u32,
4105    deadline_ms: u64,
4106    extras: GenExtras,
4107    mut on_tokens: impl FnMut(&[u32]) -> std::ops::ControlFlow<()>,
4108) -> Result<(Vec<u32>, u32, bool), DaemonError> {
4109    let (etx, erx) = std::sync::mpsc::channel();
4110    tx.send(SchedCmd::Submit(Box::new(GenerateJob {
4111        session: 0,
4112        stream,
4113        input_end: 0,
4114        params,
4115        max_tokens,
4116        events: etx,
4117        class: crate::qos::INLINE_COMPLETION,
4118        batch_invariant: false,
4119        carry: None,
4120        queued_at: 0,
4121        media: Vec::new(),
4122        completion: true,
4123        deadline: (deadline_ms > 0).then(|| std::time::Instant::now() + std::time::Duration::from_millis(deadline_ms)),
4124        grammar_handle: extras.grammar_handle,
4125        extras,
4126        trace: tracing::Span::current(),
4127    })))
4128    .map_err(|_| DaemonError::Protocol("scheduler is gone"))?;
4129    let mut tokens = Vec::new();
4130    loop {
4131        match erx.recv() {
4132            Ok(LaneEvent::Provisional { .. }) => {}
4133            Ok(LaneEvent::CompletionTokens(t)) => {
4134                if on_tokens(&t).is_break() {
4135                    tokens.extend(t);
4136                    return Ok((tokens, finish::CANCELLED, false));
4137                }
4138                tokens.extend(t);
4139            }
4140            Ok(LaneEvent::Done { finish, .. }) => return Ok((tokens, finish, false)),
4141            Ok(LaneEvent::Expired) => return Ok((Vec::new(), finish::NONE, true)),
4142            Ok(LaneEvent::Committed(_)) => {}
4143            Ok(LaneEvent::TokenLogprobs(_)) => {}
4144            Ok(LaneEvent::Failed(msg)) => return Err(DaemonError::Generation(msg)),
4145            Err(_) => return Err(DaemonError::Protocol("scheduler dropped the lane")),
4146        }
4147    }
4148}
4149
4150pub type RunJobResult = (Vec<CommittedEvent>, u32, u32, u64, Vec<TokenLogprob>, (u64, u64));
4151
4152#[allow(clippy::too_many_arguments)]
4153pub fn run_job(
4154    tx: &Sender<SchedCmd>,
4155    job_session: u64,
4156    stream: Vec<u32>,
4157    input_end: u64,
4158    params: GenParams,
4159    max_tokens: u32,
4160    class: u8,
4161    batch_invariant: bool,
4162    media: Vec<(u64, std::path::PathBuf, u32)>,
4163    extras: GenExtras,
4164    cancels: &Arc<crate::CancelSet>,
4165    mut on_event: impl FnMut(&CommittedEvent) -> Result<(), DaemonError>,
4166    mut on_logprobs: impl FnMut(&[TokenLogprob]),
4167    mut on_delta: impl FnMut(u32, &str, u32),
4168) -> Result<RunJobResult, DaemonError> {
4169    let (etx, erx) = std::sync::mpsc::channel();
4170    let grammar_handle = extras.grammar_handle;
4171    tx.send(SchedCmd::Submit(Box::new(GenerateJob {
4172        session: job_session,
4173        stream,
4174        input_end,
4175        params,
4176        max_tokens,
4177        events: etx,
4178        class,
4179        batch_invariant,
4180        carry: None,
4181        queued_at: 0,
4182        media,
4183        completion: false,
4184        deadline: None,
4185        grammar_handle,
4186        extras,
4187        trace: tracing::Span::current(),
4188    })))
4189    .map_err(|_| DaemonError::Protocol("scheduler is gone"))?;
4190    if class <= crate::qos::INLINE_COMPLETION {
4191        // A chat or a completion that may preempt: the running tick ends at
4192        // its next step instead of making it wait, and steps stay short while
4193        // such requests keep coming.
4194        cancels.latency_request();
4195    }
4196    let mut events = Vec::new();
4197    let mut logprobs: Vec<TokenLogprob> = Vec::new();
4198    let mut delivery_err: Option<DaemonError> = None;
4199    loop {
4200        match erx.recv() {
4201            Ok(LaneEvent::Provisional { channel, text, produced }) => {
4202                if delivery_err.is_none() {
4203                    on_delta(channel, &text, produced);
4204                }
4205            }
4206            Ok(LaneEvent::TokenLogprobs(items)) => {
4207                if delivery_err.is_none() {
4208                    on_logprobs(&items);
4209                }
4210                logprobs.extend(items);
4211            }
4212            Ok(LaneEvent::Committed(e)) => {
4213                if delivery_err.is_none() {
4214                    if let Err(err) = on_event(&e) {
4215                        cancels.cancel(job_session);
4216                        delivery_err = Some(err);
4217                    } else {
4218                        events.push(e);
4219                    }
4220                }
4221            }
4222            Ok(LaneEvent::Done {
4223                tokens_generated,
4224                finish,
4225                warm_prefix,
4226                spec,
4227            }) => {
4228                return match delivery_err {
4229                    Some(err) => Err(err),
4230                    None => Ok((events, tokens_generated, finish, warm_prefix, logprobs, spec)),
4231                };
4232            }
4233            Ok(LaneEvent::Failed(msg)) => {
4234                return Err(delivery_err.unwrap_or(DaemonError::Generation(msg)));
4235            }
4236            Ok(LaneEvent::CompletionTokens(_)) | Ok(LaneEvent::Expired) => {
4237                return Err(DaemonError::Protocol("completion event on a generate lane"));
4238            }
4239            Err(_) => {
4240                return Err(delivery_err.unwrap_or(DaemonError::Protocol("scheduler dropped the lane")));
4241            }
4242        }
4243    }
4244}