studio_worker/
local.rs

1//! Local image generation — submit a prompt straight to the engine, no studio.
2//!
3//! Mirrors the studio job path (`ws::session::run_offered_job`): resolve the
4//! model's [`ModelSource`] (here from the local [`Catalog`] instead of a studio
5//! offer), build a [`Task::Image`], dispatch, and record the finished job — into
6//! the dedicated local-queue ring so it shows up in the app.
7
8use std::sync::atomic::{AtomicU64, Ordering};
9
10use chrono::Utc;
11
12use crate::catalog::Catalog;
13use crate::engine::Engine;
14use crate::job_run::JobRun;
15use crate::runtime::{truncate_prompt, CurrentJob, JobOutcome, JobSource, WorkerObservers};
16use crate::types::{ImageParams, Task, TaskKind, TaskResult};
17
18/// A local image-generation request. Optional fields fall back to the model's
19/// CLI defaults from the catalog.
20#[derive(Debug, Clone, Default)]
21pub struct LocalImageRequest {
22    pub prompt: String,
23    /// Model id; `None` uses the catalog's default image model.
24    pub model: Option<String>,
25    pub negative_prompt: Option<String>,
26    pub width: Option<u32>,
27    pub height: Option<u32>,
28    pub steps: Option<u32>,
29    pub seed: Option<u64>,
30    pub ext: Option<String>,
31}
32
33/// Why a local generation could not run.
34#[derive(Debug, thiserror::Error)]
35pub enum LocalError {
36    #[error("unknown model '{0}' (not in the local catalog)")]
37    UnknownModel(String),
38    #[error("no {0} model configured in the local catalog")]
39    NoModelForKind(TaskKind),
40    #[error("no image model configured in the local catalog")]
41    NoDefaultModel,
42    #[error("model '{0}' is not an image model")]
43    NotImageModel(String),
44    #[error("model '{id}' is a {got} model, not {want}")]
45    WrongKind {
46        id: String,
47        want: TaskKind,
48        got: TaskKind,
49    },
50    #[error("engine error: {0}")]
51    Engine(String),
52}
53
54pub(crate) fn next_job_id() -> String {
55    static COUNTER: AtomicU64 = AtomicU64::new(0);
56    let n = COUNTER.fetch_add(1, Ordering::Relaxed);
57    format!("local-{}-{n}", Utc::now().timestamp_millis())
58}
59
60/// Run one local image job: resolve the model from `catalog`, dispatch it on
61/// `engine`, record it in the local-queue ring, and return the image bytes.
62pub fn run_image(
63    engine: &dyn Engine,
64    catalog: &Catalog,
65    observers: &WorkerObservers,
66    req: &LocalImageRequest,
67) -> Result<TaskResult, LocalError> {
68    let model = match &req.model {
69        Some(id) => catalog
70            .get(id)
71            .ok_or_else(|| LocalError::UnknownModel(id.clone()))?,
72        None => catalog
73            .default_image_model()
74            .ok_or(LocalError::NoDefaultModel)?,
75    };
76    if model.kind != TaskKind::Image {
77        return Err(LocalError::NotImageModel(model.id.clone()));
78    }
79
80    let defaults = &model.source.cli_defaults;
81    let params = ImageParams {
82        prompt: req.prompt.clone(),
83        negative_prompt: req.negative_prompt.clone(),
84        width: req.width.unwrap_or(defaults.width).max(1),
85        height: req.height.unwrap_or(defaults.height).max(1),
86        steps: req.steps.unwrap_or(defaults.steps).max(1),
87        seed: req.seed,
88        cfg_scale: Some(defaults.cfg_scale),
89        sampling_method: defaults.sampling_method.clone(),
90        ext: req.ext.clone().unwrap_or_else(|| "webp".to_string()),
91        ..Default::default()
92    };
93
94    dispatch_and_record(engine, model, observers, &req.prompt, Task::Image(params))
95}
96
97/// Resolve a model of `kind` (an explicit id, else the catalog's
98/// default for that kind), dispatch `task`, record the local job, and
99/// return the result.  The generic core behind every non-image local
100/// endpoint (chat / tts / stt / video) so each stays a thin adapter.
101pub fn run_kind(
102    engine: &dyn Engine,
103    catalog: &Catalog,
104    observers: &WorkerObservers,
105    kind: TaskKind,
106    model_id: Option<&str>,
107    prompt_preview: &str,
108    task: Task,
109) -> Result<TaskResult, LocalError> {
110    let model = match model_id {
111        Some(id) => catalog
112            .get(id)
113            .ok_or_else(|| LocalError::UnknownModel(id.to_string()))?,
114        None => catalog
115            .default_model_for(kind)
116            .ok_or(LocalError::NoModelForKind(kind))?,
117    };
118    if model.kind != kind {
119        return Err(LocalError::WrongKind {
120            id: model.id.clone(),
121            want: kind,
122            got: model.kind,
123        });
124    }
125    dispatch_and_record(engine, model, observers, prompt_preview, task)
126}
127
128/// Serve a chat on the lane of a loaded model, recording it like any
129/// local job.  `None` when the resolved model is not loaded (or is not
130/// a chat model): the caller runs it as a transient job instead.
131pub fn chat_on_lane(
132    host: &crate::host::ModelHost,
133    catalog: &Catalog,
134    observers: &WorkerObservers,
135    model_id: Option<&str>,
136    prompt_preview: &str,
137    params: crate::types::LlmParams,
138) -> Option<Result<TaskResult, LocalError>> {
139    let model = match model_id {
140        Some(id) => catalog.get(id)?,
141        None => catalog.default_model_for(TaskKind::Llm)?,
142    };
143    if model.kind != TaskKind::Llm {
144        return None;
145    }
146    let served = host.with_lane(&model.id, |loaded, lane| {
147        let chat = loaded.as_chat()?;
148        let run = JobRun::begin(
149            observers,
150            CurrentJob {
151                job_id: next_job_id(),
152                kind: TaskKind::Llm,
153                model: model.id.clone(),
154                prompt: truncate_prompt(prompt_preview),
155                started_at: Utc::now(),
156                source: JobSource::Lane,
157            },
158        );
159        let result = run
160            .span()
161            .in_scope(|| chat.chat(params, &|| lane.cancelled(), &mut |_| {}));
162        Some((run, result))
163    });
164    let (run, result) = match served {
165        Ok(Some(served)) => served,
166        // Not loaded, or loaded but not a chat model: the transient path decides.
167        Ok(None) | Err(_) => return None,
168    };
169    let result = result.map(|json| TaskResult::Llm { json });
170    run.finish(outcome_of(&result));
171    Some(result.map_err(|err| LocalError::Engine(err.to_string())))
172}
173
174/// The chat model a request names, or the default one: it must exist and
175/// be an LLM.
176pub fn resolve_llm<'a>(
177    catalog: &'a Catalog,
178    model_id: Option<&str>,
179) -> Result<&'a crate::catalog::CatalogModel, LocalError> {
180    let model = match model_id {
181        Some(id) => catalog
182            .get(id)
183            .ok_or_else(|| LocalError::UnknownModel(id.to_string()))?,
184        None => catalog
185            .default_model_for(TaskKind::Llm)
186            .ok_or(LocalError::NoModelForKind(TaskKind::Llm))?,
187    };
188    if model.kind != TaskKind::Llm {
189        return Err(LocalError::WrongKind {
190            id: model.id.clone(),
191            want: TaskKind::Llm,
192            got: model.kind,
193        });
194    }
195    Ok(model)
196}
197
198/// Stream a chat on `model_id`'s lane as server-sent events through
199/// `send` (which answers `false` once the client has gone, ending the
200/// generation).  Recorded like any local job.  The caller has checked the
201/// model is loaded; if it unloads first, the stream carries the error.
202pub fn stream_on_lane(
203    host: &crate::host::ModelHost,
204    observers: &WorkerObservers,
205    model_id: &str,
206    prompt_preview: &str,
207    params: crate::types::LlmParams,
208    send: &mut dyn FnMut(Vec<u8>) -> bool,
209) {
210    use crate::engine::llm_core::{
211        chunk_frame, final_frame, sse_event, Finish, ThinkSplitter, SSE_DONE,
212    };
213    let gone = std::cell::Cell::new(false);
214    let mut emit = |bytes: Vec<u8>| {
215        if !gone.get() && !send(bytes) {
216            gone.set(true);
217        }
218    };
219    let served = host.with_lane(model_id, |loaded, lane| {
220        let Some(chat) = loaded.as_chat() else {
221            return Err(anyhow::anyhow!("model {model_id} is not a chat model"));
222        };
223        let run = JobRun::begin(
224            observers,
225            CurrentJob {
226                job_id: next_job_id(),
227                kind: TaskKind::Llm,
228                model: model_id.to_string(),
229                prompt: truncate_prompt(prompt_preview),
230                started_at: Utc::now(),
231                source: JobSource::Lane,
232            },
233        );
234        let mut splitter = ThinkSplitter::default();
235        let result = run.span().in_scope(|| {
236            chat.chat(params, &|| lane.cancelled() || gone.get(), &mut |piece| {
237                for delta in splitter.push(piece) {
238                    emit(sse_event(&chunk_frame(model_id, &delta)));
239                }
240            })
241        });
242        for delta in splitter.finish() {
243            emit(sse_event(&chunk_frame(model_id, &delta)));
244        }
245        let recorded: Result<TaskResult, anyhow::Error> = match &result {
246            Ok(json) => Ok(TaskResult::Llm { json: json.clone() }),
247            Err(e) => Err(anyhow::anyhow!("{e:#}")),
248        };
249        run.finish(outcome_of(&recorded));
250        result
251    });
252    let last = match served {
253        Ok(Ok(json)) => {
254            let finish = match json["choices"][0]["finish_reason"].as_str() {
255                Some("length") => Finish::Length,
256                _ => Finish::Stop,
257            };
258            let usage = &json["usage"];
259            final_frame(
260                model_id,
261                finish,
262                usage["prompt_tokens"].as_u64().unwrap_or(0) as usize,
263                usage["completion_tokens"].as_u64().unwrap_or(0) as u32,
264            )
265        }
266        Ok(Err(e)) => serde_json::json!({ "error": { "message": format!("{e:#}") } }),
267        Err(e) => serde_json::json!({ "error": { "message": e.to_string() } }),
268    };
269    emit(sse_event(&last));
270    emit(SSE_DONE.to_vec());
271}
272
273/// Token ids for `text` from a loaded chat model, on its lane.
274pub fn tokenize_on_lane(
275    host: &crate::host::ModelHost,
276    model_id: &str,
277    text: &str,
278    add_special: bool,
279) -> Result<anyhow::Result<Vec<i32>>, crate::host::HostError> {
280    host.with_lane(model_id, |loaded, _lane| match loaded.as_chat() {
281        Some(chat) => chat.tokenize(text, add_special),
282        None => Err(anyhow::anyhow!("model {model_id} is not a chat model")),
283    })
284}
285
286/// The outcome a finished job is recorded with.
287fn outcome_of<E: std::fmt::Display>(result: &Result<TaskResult, E>) -> JobOutcome {
288    match result {
289        Ok(_) => JobOutcome::Completed,
290        Err(err) => JobOutcome::Failed {
291            reason: err.to_string(),
292        },
293    }
294}
295
296/// Dispatch `task` on `model` and record the finished job in the
297/// local-queue ring.  Shared by [`run_image`] and [`run_kind`] so the
298/// dispatch + bookkeeping lives in one place.
299fn dispatch_and_record(
300    engine: &dyn Engine,
301    model: &crate::catalog::CatalogModel,
302    observers: &WorkerObservers,
303    prompt_preview: &str,
304    task: Task,
305) -> Result<TaskResult, LocalError> {
306    let run = JobRun::begin(
307        observers,
308        CurrentJob {
309            job_id: next_job_id(),
310            kind: task.kind(),
311            model: model.id.clone(),
312            prompt: truncate_prompt(prompt_preview),
313            started_at: Utc::now(),
314            source: JobSource::Local,
315        },
316    );
317    let result = run
318        .span()
319        .in_scope(|| engine.dispatch_with_source(&model.id, task, &model.source));
320    if let Ok(result) = &result {
321        run.keep_thumbnail(result);
322    }
323    run.finish(outcome_of(&result));
324
325    result.map_err(|err| LocalError::Engine(err.to_string()))
326}
327
328#[cfg(test)]
329mod tests {
330    use super::*;
331    use crate::catalog::CatalogModel;
332    use crate::engine::SyntheticEngine;
333    use crate::types::{ModelCliDefaults, ModelEngine, ModelSource};
334
335    fn synthetic_model(id: &str, kind: TaskKind) -> CatalogModel {
336        CatalogModel {
337            id: id.into(),
338            display_name: id.into(),
339            kind,
340            vram_gb_estimate: 0.0,
341            description: None,
342            source: ModelSource {
343                engine: ModelEngine::Synthetic,
344                files: vec![],
345                cli_defaults: ModelCliDefaults {
346                    cfg_scale: 1.0,
347                    steps: 4,
348                    width: 64,
349                    height: 64,
350                    ..Default::default()
351                },
352            },
353            enabled: true,
354            origin: "local".into(),
355            exclusive_group: None,
356        }
357    }
358
359    fn catalog_with(models: Vec<CatalogModel>) -> Catalog {
360        Catalog {
361            models,
362            ..Default::default()
363        }
364    }
365
366    #[test]
367    fn generates_image_and_records_local_job() {
368        let engine = SyntheticEngine::new();
369        let catalog = catalog_with(vec![synthetic_model("synthetic-img", TaskKind::Image)]);
370        let observers = WorkerObservers::default();
371        let req = LocalImageRequest {
372            prompt: "a red fox".into(),
373            ..Default::default()
374        };
375
376        let result = run_image(&engine, &catalog, &observers, &req).unwrap();
377        match result {
378            TaskResult::Image { bytes, ext } => {
379                assert!(!bytes.is_empty());
380                assert_eq!(ext, "webp");
381            }
382            other => panic!("expected image, got {other:?}"),
383        }
384
385        let ring = observers.local_jobs.lock();
386        assert_eq!(ring.len(), 1);
387        let job = &ring[0];
388        assert_eq!(job.model, "synthetic-img");
389        assert_eq!(job.outcome, JobOutcome::Completed);
390        assert_eq!(job.prompt, "a red fox");
391        // The studio ring stays empty — local jobs are their own queue.
392        assert!(observers.recent_jobs.lock().is_empty());
393    }
394
395    #[test]
396    fn an_image_job_keeps_a_thumbnail_and_a_job_log() {
397        crate::test_support::install_job_log_capture();
398        let engine = SyntheticEngine::new();
399        let catalog = catalog_with(vec![synthetic_model("img", TaskKind::Image)]);
400        let observers = WorkerObservers::default();
401        let req = LocalImageRequest {
402            prompt: "a lighthouse".into(),
403            ..Default::default()
404        };
405        run_image(&engine, &catalog, &observers, &req).unwrap();
406
407        let job = observers.local_jobs.lock()[0].clone();
408        assert_eq!(job.source, JobSource::Local);
409        assert!(observers.thumbnails.contains(&job.job_id));
410        assert!(observers.active_jobs.lock().is_empty());
411        let log = crate::job_log::global().get(&job.job_id).expect("job log");
412        assert!(log.lines[0].message.starts_with("job started"));
413        assert!(log
414            .lines
415            .last()
416            .is_some_and(|l| l.message.starts_with("job finished")));
417    }
418
419    #[test]
420    fn a_chat_on_a_loaded_model_is_recorded_as_a_lane_job() {
421        let mut model = synthetic_model("chat", TaskKind::Llm);
422        model.source.engine = ModelEngine::LlamaCpp;
423        let catalog = catalog_with(vec![model]);
424        let shared = std::sync::Arc::new(parking_lot::Mutex::new(catalog.clone()));
425        let host = crate::host::ModelHost::new(
426            shared,
427            std::sync::Arc::new(crate::test_support::InstantRuntime),
428            std::sync::Arc::new(crate::test_support::FixedProbe(20.0)),
429            crate::residency::Residency::load_for_serving(None),
430        );
431        host.load("chat").unwrap();
432        host.wait_for(
433            "chat",
434            crate::lifecycle::ModelState::serves,
435            std::time::Duration::from_secs(5),
436        )
437        .expect("loaded");
438        let observers = WorkerObservers::default();
439        let params = crate::types::LlmParams {
440            messages: vec![crate::types::ChatMessage {
441                role: "user".into(),
442                content: "hello".into(),
443            }],
444            ..Default::default()
445        };
446
447        let answer = chat_on_lane(&host, &catalog, &observers, None, "hello", params)
448            .expect("served on the lane")
449            .unwrap();
450
451        assert!(matches!(answer, TaskResult::Llm { .. }));
452        let job = observers.local_jobs.lock()[0].clone();
453        assert_eq!(job.source, JobSource::Lane);
454        assert_eq!(job.outcome, JobOutcome::Completed);
455        assert!(observers.active_jobs.lock().is_empty());
456    }
457
458    #[test]
459    fn defaults_to_the_only_image_model_when_unspecified() {
460        let engine = SyntheticEngine::new();
461        let catalog = catalog_with(vec![synthetic_model("only-img", TaskKind::Image)]);
462        let observers = WorkerObservers::default();
463        let req = LocalImageRequest {
464            prompt: "x".into(),
465            model: None,
466            ..Default::default()
467        };
468        let out = run_image(&engine, &catalog, &observers, &req).unwrap();
469        assert!(matches!(out, TaskResult::Image { .. }));
470        assert_eq!(observers.local_jobs.lock()[0].model, "only-img");
471    }
472
473    #[test]
474    fn unknown_model_is_rejected() {
475        let engine = SyntheticEngine::new();
476        let catalog = catalog_with(vec![synthetic_model("a", TaskKind::Image)]);
477        let observers = WorkerObservers::default();
478        let req = LocalImageRequest {
479            prompt: "x".into(),
480            model: Some("missing".into()),
481            ..Default::default()
482        };
483        let err = run_image(&engine, &catalog, &observers, &req).unwrap_err();
484        assert!(matches!(err, LocalError::UnknownModel(m) if m == "missing"));
485        assert!(observers.local_jobs.lock().is_empty());
486    }
487
488    #[test]
489    fn no_image_model_yields_no_default() {
490        let engine = SyntheticEngine::new();
491        let catalog = catalog_with(vec![]);
492        let observers = WorkerObservers::default();
493        let req = LocalImageRequest {
494            prompt: "x".into(),
495            ..Default::default()
496        };
497        let err = run_image(&engine, &catalog, &observers, &req).unwrap_err();
498        assert!(matches!(err, LocalError::NoDefaultModel));
499    }
500
501    #[test]
502    fn non_image_model_is_rejected() {
503        let engine = SyntheticEngine::new();
504        let catalog = catalog_with(vec![synthetic_model("chat", TaskKind::Llm)]);
505        let observers = WorkerObservers::default();
506        let req = LocalImageRequest {
507            prompt: "x".into(),
508            model: Some("chat".into()),
509            ..Default::default()
510        };
511        let err = run_image(&engine, &catalog, &observers, &req).unwrap_err();
512        assert!(matches!(err, LocalError::NotImageModel(m) if m == "chat"));
513    }
514
515    #[test]
516    fn run_kind_dispatches_llm_and_records_the_job() {
517        let engine = SyntheticEngine::new();
518        let catalog = catalog_with(vec![synthetic_model("chat", TaskKind::Llm)]);
519        let observers = WorkerObservers::default();
520        let task = Task::Llm(crate::types::LlmParams {
521            messages: vec![crate::types::ChatMessage {
522                role: "user".into(),
523                content: "hi".into(),
524            }],
525            ..Default::default()
526        });
527        let out = run_kind(
528            &engine,
529            &catalog,
530            &observers,
531            TaskKind::Llm,
532            None,
533            "hi",
534            task,
535        )
536        .unwrap();
537        assert!(matches!(out, TaskResult::Llm { .. }));
538        assert_eq!(observers.local_jobs.lock()[0].kind, TaskKind::Llm);
539    }
540
541    #[test]
542    fn run_kind_rejects_a_wrong_kind_model() {
543        let engine = SyntheticEngine::new();
544        // An image model explicitly requested for an LLM job.
545        let catalog = catalog_with(vec![synthetic_model("img", TaskKind::Image)]);
546        let observers = WorkerObservers::default();
547        let task = Task::Llm(crate::types::LlmParams::default());
548        let err = run_kind(
549            &engine,
550            &catalog,
551            &observers,
552            TaskKind::Llm,
553            Some("img"),
554            "",
555            task,
556        )
557        .unwrap_err();
558        assert!(
559            matches!(err, LocalError::WrongKind { ref id, want, got }
560                if id == "img" && want == TaskKind::Llm && got == TaskKind::Image),
561            "got {err:?}"
562        );
563        assert!(
564            observers.local_jobs.lock().is_empty(),
565            "no job recorded on reject"
566        );
567    }
568
569    #[test]
570    fn run_kind_reports_no_model_for_kind() {
571        let engine = SyntheticEngine::new();
572        let catalog = catalog_with(vec![]);
573        let observers = WorkerObservers::default();
574        let err = run_kind(
575            &engine,
576            &catalog,
577            &observers,
578            TaskKind::AudioTts,
579            None,
580            "",
581            Task::AudioTts(crate::types::AudioTtsParams::default()),
582        )
583        .unwrap_err();
584        assert!(matches!(
585            err,
586            LocalError::NoModelForKind(TaskKind::AudioTts)
587        ));
588    }
589}
studio_worker/local.rs

studio_worker/
local.rs
