twine_core/harness/models/
discovery.rs

1//! Owns model-discovery workers and their bounded helper processes.
2
3use std::ffi::OsString;
4use std::io::{Read, Seek};
5use std::path::Path;
6use std::process::{Command, Stdio};
7use std::sync::atomic::{AtomicBool, Ordering};
8use std::sync::{Arc, Mutex, PoisonError, Weak, mpsc};
9use std::thread;
10use std::time::{Duration, Instant};
11
12use tracing::warn;
13
14use super::{HarnessModels, ModelListError, help_levels};
15use crate::harness::{HarnessError, HarnessId, LocatedHarness};
16use crate::process::CommandChild;
17
18const LIST_TIMEOUT: Duration = Duration::from_secs(20);
19const OPENCODE_CATALOG_GRACE: Duration = Duration::from_secs(2);
20/// Codex's catalog carries each model's instructions, so it runs to hundreds of KiB.
21const MAX_LIST_OUTPUT: u64 = 4 * 1024 * 1024;
22
23impl HarnessId {
24    /// The arguments that make the harness print its models.
25    fn model_list_arguments(self) -> &'static [&'static str] {
26        match self {
27            Self::Codex => &["debug", "models"],
28            Self::ClaudeCode => &["--help"],
29            Self::Pi => &["--list-models"],
30            Self::Antigravity => &["models"],
31            Self::Omp => &["models", "--json"],
32            Self::Opencode => &["api", "model.list"],
33        }
34    }
35}
36
37/// Discovery workers owned by one application runtime, even when requests outlive that runtime.
38#[derive(Default)]
39pub(crate) struct ModelListWorkers {
40    workers: Mutex<Vec<Weak<ModelListWorker>>>,
41}
42
43impl ModelListWorkers {
44    pub(crate) fn request(
45        &self,
46        harness: HarnessId,
47        folder: Option<std::path::PathBuf>,
48        locate: impl FnOnce() -> Result<LocatedHarness, HarnessError> + Send + 'static,
49    ) -> ModelListRequest {
50        let request = ModelListRequest::spawn(harness, folder, locate);
51        let mut workers = self.workers.lock().unwrap_or_else(PoisonError::into_inner);
52        workers.retain(|worker| worker.strong_count() > 0);
53        workers.push(Arc::downgrade(&request.worker));
54        request
55    }
56
57    pub(crate) fn shutdown(&self) {
58        let workers =
59            std::mem::take(&mut *self.workers.lock().unwrap_or_else(PoisonError::into_inner))
60                .into_iter()
61                .filter_map(|worker| worker.upgrade())
62                .collect::<Vec<_>>();
63        // Cancel every listing before joining any, so independent helpers stop together.
64        for worker in &workers {
65            worker.cancelled.store(true, Ordering::Relaxed);
66        }
67        for worker in workers {
68            worker.join();
69        }
70    }
71}
72
73struct ModelListWorker {
74    cancelled: Arc<AtomicBool>,
75    thread: Mutex<Option<thread::JoinHandle<()>>>,
76}
77
78impl ModelListWorker {
79    fn join(&self) {
80        // Keep the lock until the join finishes: a concurrent shutdown must also wait for cleanup.
81        let mut thread = self.thread.lock().unwrap_or_else(PoisonError::into_inner);
82        if let Some(thread) = thread.take()
83            && thread.join().is_err()
84        {
85            warn!("model discovery worker panicked");
86        }
87    }
88}
89
90/// A model list running on its own thread, so a slow harness never holds up the caller. Dropping
91/// it stops the harness and waits for its worker to clean up.
92pub struct ModelListRequest {
93    receiver: mpsc::Receiver<Result<HarnessModels, ModelListError>>,
94    worker: Arc<ModelListWorker>,
95}
96
97impl ModelListRequest {
98    /// Finds the harness with `locate`, then lists its models, all on a new thread.
99    pub(crate) fn spawn(
100        harness: HarnessId,
101        folder: Option<std::path::PathBuf>,
102        locate: impl FnOnce() -> Result<LocatedHarness, HarnessError> + Send + 'static,
103    ) -> Self {
104        let (sender, receiver) = mpsc::sync_channel(1);
105        let cancelled = Arc::new(AtomicBool::new(false));
106        let flag = Arc::clone(&cancelled);
107        let spawned = crate::blocking_worker::spawn("harness-models".into(), move || {
108            let located = locate();
109            let result = if flag.load(Ordering::Relaxed) {
110                Err(ModelListError::Cancelled)
111            } else {
112                located.map_err(ModelListError::from).and_then(|located| {
113                    list_models(
114                        harness,
115                        &located.program,
116                        &located.path,
117                        folder.as_deref(),
118                        &flag,
119                    )
120                })
121            };
122            let _ = sender.send(result);
123        });
124        let thread = spawned
125            .map_err(|error| warn!(%error, "couldn't start listing harness models"))
126            .ok();
127        Self {
128            receiver,
129            worker: Arc::new(ModelListWorker {
130                cancelled,
131                thread: Mutex::new(thread),
132            }),
133        }
134    }
135
136    /// The models once the harness has answered, without waiting for it.
137    #[must_use]
138    pub fn poll(&self) -> Option<Result<HarnessModels, ModelListError>> {
139        match self.receiver.try_recv() {
140            Ok(result) => Some(result),
141            Err(mpsc::TryRecvError::Empty) => None,
142            Err(mpsc::TryRecvError::Disconnected) => Some(Err(ModelListError::Cancelled)),
143        }
144    }
145}
146
147impl Drop for ModelListRequest {
148    fn drop(&mut self) {
149        self.worker.cancelled.store(true, Ordering::Relaxed);
150        self.worker.join();
151    }
152}
153
154/// Runs `program` to list the harness's models, with `path` as its `PATH`.
155fn list_models(
156    harness: HarnessId,
157    program: &Path,
158    path: &OsString,
159    folder: Option<&Path>,
160    cancelled: &AtomicBool,
161) -> Result<HarnessModels, ModelListError> {
162    let name = harness.definition().name;
163    let directory = if harness == HarnessId::Opencode {
164        folder
165    } else {
166        None
167    };
168    let run = |arguments| {
169        run(
170            program,
171            arguments,
172            path,
173            directory,
174            name,
175            LIST_TIMEOUT,
176            cancelled,
177        )
178    };
179    // The shared server uses its own default location unless the nested query is explicit.
180    let location =
181        directory.map(|directory| format!("location[directory]={}", directory.display()));
182    let mut arguments = harness.model_list_arguments().to_vec();
183    if let Some(location) = &location {
184        arguments.extend(["--param", location.as_str()]);
185    }
186    let mut models = harness.parse_models(&run(&arguments)?)?;
187    // A new v2 location can return its snapshot before provider plugins settle.
188    let deadline = Instant::now() + OPENCODE_CATALOG_GRACE;
189    while harness == HarnessId::Opencode && models.models.is_empty() && Instant::now() < deadline {
190        if cancelled.load(Ordering::Relaxed) {
191            return Err(ModelListError::Cancelled);
192        }
193        thread::sleep(Duration::from_millis(100));
194        models = harness.parse_models(&run(&arguments)?)?;
195    }
196    // These catalogs omit default-model effort levels, which their help lists instead.
197    let effort_option = match harness {
198        HarnessId::Pi => Some("--thinking <level>"),
199        HarnessId::Antigravity => Some("--effort"),
200        HarnessId::Omp => Some("--thinking=<value>"),
201        HarnessId::Codex | HarnessId::ClaudeCode | HarnessId::Opencode => None,
202    };
203    if let Some(option) = effort_option {
204        models.efforts = help_levels(&run(&["--help"])?, option);
205    }
206    Ok(models)
207}
208
209/// Runs `program` with `arguments` and returns what it printed, within a time and size limit. A
210/// harness that exits unsuccessfully has failed, whatever it printed.
211fn run(
212    program: &Path,
213    arguments: &[&str],
214    path: &OsString,
215    directory: Option<&Path>,
216    name: &'static str,
217    timeout: Duration,
218    cancelled: &AtomicBool,
219) -> Result<String, ModelListError> {
220    if cancelled.load(Ordering::Relaxed) {
221        return Err(ModelListError::Cancelled);
222    }
223    let failed = |error: std::io::Error| {
224        warn!(harness = name, %error, "couldn't run a harness to list its models");
225        ModelListError::Failed(name)
226    };
227    // A file avoids blocking on a full pipe while the harness is supervised.
228    let mut output = tempfile::tempfile().map_err(failed)?;
229    let stdout = output.try_clone().map_err(failed)?;
230    let mut command = Command::new(program);
231    if let Some(directory) = directory {
232        command.current_dir(directory);
233    }
234    command
235        .args(arguments)
236        .env("PATH", path)
237        .env("NO_COLOR", "1")
238        .stdin(Stdio::null())
239        .stdout(stdout)
240        .stderr(Stdio::null());
241    let mut child = CommandChild::spawn(&mut command).map_err(|error| {
242        warn!(harness = name, %error, "couldn't start a harness to list its models");
243        ModelListError::Start(name)
244    })?;
245    let deadline = Instant::now() + timeout;
246    let status = loop {
247        if cancelled.load(Ordering::Relaxed) {
248            return Err(ModelListError::Cancelled);
249        }
250        if output.metadata().map_err(failed)?.len() > MAX_LIST_OUTPUT {
251            warn!(
252                harness = name,
253                "a harness printed too much while listing its models"
254            );
255            return Err(ModelListError::Unreadable(name));
256        }
257        if child.has_exited().map_err(failed)? {
258            break child.finish().map_err(failed)?;
259        }
260        if Instant::now() >= deadline {
261            warn!(harness = name, "a harness took too long to list its models");
262            return Err(ModelListError::Timeout(name));
263        }
264        thread::sleep(Duration::from_millis(10));
265    };
266    if !status.success() {
267        warn!(harness = name, %status, "a harness failed to list its models");
268        return Err(ModelListError::Failed(name));
269    }
270    output.rewind().map_err(failed)?;
271    let mut text = String::new();
272    output
273        .take(MAX_LIST_OUTPUT)
274        .read_to_string(&mut text)
275        .map_err(|_| ModelListError::Unreadable(name))?;
276    Ok(text)
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282    use crate::process::tests::FixtureProcess;
283
284    #[test]
285    fn concurrent_request_destruction_and_shutdown_both_wait_for_cleanup() {
286        let workers = Arc::new(ModelListWorkers::default());
287        let (started, ready) = mpsc::channel();
288        let (release, gate) = mpsc::channel();
289        let request = workers.request(HarnessId::Pi, None, move || {
290            started.send(()).unwrap();
291            gate.recv().unwrap();
292            Err(HarnessError::NotFound {
293                name: "fixture",
294                binary: "fixture",
295            })
296        });
297        ready.recv_timeout(Duration::from_secs(3)).unwrap();
298        let cancellation = Arc::clone(&request.worker.cancelled);
299        let (dropped, drop_done) = mpsc::channel();
300        let dropping = thread::spawn(move || {
301            drop(request);
302            dropped.send(()).unwrap();
303        });
304        let deadline = Instant::now() + Duration::from_secs(3);
305        while !cancellation.load(Ordering::Relaxed) {
306            assert!(Instant::now() < deadline);
307            thread::sleep(Duration::from_millis(5));
308        }
309        let registry = Arc::clone(&workers);
310        let (stopped, shutdown_done) = mpsc::channel();
311        let stopping = thread::spawn(move || {
312            registry.shutdown();
313            stopped.send(()).unwrap();
314        });
315        // Wait until shutdown has taken ownership of its registrations before checking it blocks.
316        while !workers.workers.lock().unwrap().is_empty() {
317            assert!(Instant::now() < deadline);
318            thread::sleep(Duration::from_millis(5));
319        }
320        assert!(matches!(
321            drop_done.recv_timeout(Duration::from_millis(50)),
322            Err(mpsc::RecvTimeoutError::Timeout)
323        ));
324        assert!(matches!(
325            shutdown_done.recv_timeout(Duration::from_millis(50)),
326            Err(mpsc::RecvTimeoutError::Timeout)
327        ));
328        release.send(()).unwrap();
329        dropping.join().unwrap();
330        stopping.join().unwrap();
331        drop_done.recv().unwrap();
332        shutdown_done.recv().unwrap();
333    }
334
335    #[test]
336    fn omp_discovers_models_and_thinking_levels_from_its_cli() {
337        let directory = tempfile::tempdir().unwrap();
338        let program = stub(
339            directory.path(),
340            "omp",
341            "case \"$1 $2\" in\n'models --json') printf '%s\\n' '{\"models\":[{\"selector\":\"local/model\",\"provider\":\"local\",\"name\":\"Model\",\"thinking\":[\"high\"]}]}' ;;\n\
342             '--help ') printf '      --thinking=<value>  Set thinking level: off, minimal, low, medium, high, xhigh, max, auto\\n      --service-tier=<value>  Service tier\\n' ;;\n\
343             *) exit 1 ;;\nesac",
344        );
345        let models = list_models(
346            HarnessId::Omp,
347            &program,
348            &OsString::from("/bin:/usr/bin"),
349            None,
350            &AtomicBool::new(false),
351        )
352        .unwrap();
353        assert_eq!(models.models[0].id, "local/model");
354        assert_eq!(
355            models.efforts,
356            [
357                "off", "minimal", "low", "medium", "high", "xhigh", "max", "auto"
358            ]
359        );
360    }
361
362    #[test]
363    fn opencode_discovers_model_variants_through_its_cli_api() {
364        let directory = tempfile::tempdir().unwrap();
365        let program = stub(
366            directory.path(),
367            "opencode",
368            r#"
369            [ "$#" -eq 2 ] && [ "$1 $2" = 'api model.list' ] || exit 1
370            printf '%s\n' '{"data":[{"providerID":"local","id":"model","name":"Model","variants":[{"id":"high"}]}]}'
371        "#,
372        );
373        let models = list_models(
374            HarnessId::Opencode,
375            &program,
376            &OsString::from("/bin:/usr/bin"),
377            None,
378            &AtomicBool::new(false),
379        )
380        .unwrap();
381        assert_eq!(models.models[0].id, "local/model");
382        assert_eq!(models.models[0].efforts, Some(vec!["high".into()]));
383        assert!(models.efforts.is_empty() && !models.supports_yolo);
384    }
385
386    #[test]
387    fn opencode_discovers_each_folders_catalog_after_provider_startup() {
388        let directory = tempfile::tempdir().unwrap();
389        let program = stub(
390            directory.path(),
391            "opencode",
392            r#"
393            [ "$1 $2 $3" = 'api model.list --param' ] || exit 1
394            catalog_location=${4#*=}
395            [ "${4%%=*}" = 'location[directory]' ] && [ "$catalog_location" -ef "$PWD" ] || exit 1
396            if [ ! -f ready ]; then
397                touch ready
398                printf '%s\n' '{"data":[]}'
399            else
400                cat catalog.json
401            fi
402        "#,
403        );
404        for (folder, variant) in [("first", "custom-name"), ("second", "Custom_Name")] {
405            let path = directory.path().join(folder);
406            std::fs::create_dir(&path).unwrap();
407            std::fs::write(path.join("catalog.json"), serde_json::to_vec(&serde_json::json!({
408                "data": [{"providerID": "local", "id": folder, "name": folder, "variants": [{"id": variant}]}]
409            })).unwrap()).unwrap();
410            let models = list_models(
411                HarnessId::Opencode,
412                &program,
413                &OsString::from("/bin:/usr/bin"),
414                Some(&path),
415                &AtomicBool::new(false),
416            )
417            .unwrap();
418            assert_eq!(models.models[0].id, format!("local/{folder}"));
419            assert_eq!(models.models[0].efforts, Some(vec![variant.into()]));
420        }
421    }
422
423    #[test]
424    fn antigravity_reads_models_and_effort_levels_from_its_cli() {
425        let directory = tempfile::tempdir().unwrap();
426        let program = stub(
427            directory.path(),
428            "agy",
429            "case \"$1\" in\nmodels) printf 'gemini-3.8-flash-high\\tGemini 3.8 Flash (High)\\n' ;;\n\
430             --help) printf '  --effort  Reasoning effort for the current CLI session (low|medium|high)\\n  --model  Model\\n' ;;\n\
431             *) exit 1 ;;\nesac",
432        );
433        let models = list_models(
434            HarnessId::Antigravity,
435            &program,
436            &OsString::from("/bin:/usr/bin"),
437            None,
438            &AtomicBool::new(false),
439        )
440        .unwrap();
441        assert_eq!(models.models[0].id, "gemini-3.8-flash-high");
442        assert_eq!(models.efforts, ["low", "medium", "high"]);
443        assert!(models.supports_yolo);
444    }
445
446    fn stub(directory: &Path, name: &str, script: &str) -> std::path::PathBuf {
447        let program = directory.join(name);
448        std::fs::write(&program, format!("#!/bin/sh\n{script}\n")).unwrap();
449        std::fs::set_permissions(
450            &program,
451            std::os::unix::fs::PermissionsExt::from_mode(0o755),
452        )
453        .unwrap();
454        program
455    }
456
457    fn is_running(pid: &str) -> bool {
458        Command::new("/bin/kill")
459            .args(["-0", pid])
460            .stderr(Stdio::null())
461            .status()
462            .unwrap()
463            .success()
464    }
465
466    #[test]
467    fn listing_runs_the_harness_and_reports_a_harness_that_fails_to_answer() {
468        let directory = tempfile::tempdir().unwrap();
469        let program = stub(
470            directory.path(),
471            "pi",
472            "case \"$1\" in\n--list-models) printf 'provider model\\nlocal m1 x\\n' ;;\n\
473             --help) printf '  --thinking <level>  Set thinking level: off, low, high\\n  -e x\\n' ;;\n\
474             *) exit 1 ;;\nesac",
475        );
476        let path = OsString::from("/bin:/usr/bin");
477        let running = AtomicBool::new(false);
478        let models = list_models(HarnessId::Pi, &program, &path, None, &running).unwrap();
479        assert_eq!(models.models[0].id, "local/m1");
480        assert_eq!(models.efforts, ["off", "low", "high"]);
481        assert!(list_models(HarnessId::Codex, &program, &path, None, &running).is_err());
482        assert!(
483            list_models(
484                HarnessId::Pi,
485                &directory.path().join("missing"),
486                &path,
487                None,
488                &running
489            )
490            .is_err()
491        );
492        // An empty list from a harness that failed is a failure, not "no models".
493        let broken = stub(directory.path(), "claude", "exit 1");
494        assert!(matches!(
495            list_models(HarnessId::ClaudeCode, &broken, &path, None, &running),
496            Err(ModelListError::Failed("Claude Code"))
497        ));
498    }
499
500    #[test]
501    fn a_slow_or_flooding_harness_is_stopped() {
502        let directory = tempfile::tempdir().unwrap();
503        let path = OsString::from("/bin:/usr/bin");
504        let running = AtomicBool::new(false);
505        let pid_file = directory.path().join("pid");
506        let slow = stub(
507            directory.path(),
508            "slow",
509            &format!("echo $$ > '{}'; exec sleep 60", pid_file.display()),
510        );
511        let started = Instant::now();
512        assert!(matches!(
513            run(
514                &slow,
515                &[],
516                &path,
517                None,
518                "pi",
519                Duration::from_millis(300),
520                &running
521            ),
522            Err(ModelListError::Timeout("pi"))
523        ));
524        assert!(started.elapsed() < Duration::from_secs(5));
525        let pid = std::fs::read_to_string(&pid_file).unwrap();
526        assert!(!is_running(pid.trim()), "the timed-out harness was killed");
527        let child_file = directory.path().join("child");
528        let wrapper = stub(
529            directory.path(),
530            "wrapper",
531            &format!("sleep 60 & echo $! > '{}'; wait", child_file.display()),
532        );
533        assert!(
534            run(
535                &wrapper,
536                &[],
537                &path,
538                None,
539                "pi",
540                Duration::from_millis(300),
541                &running
542            )
543            .is_err()
544        );
545        let child = std::fs::read_to_string(&child_file).unwrap();
546        // A killed process can linger briefly until it's reaped.
547        let deadline = Instant::now() + Duration::from_secs(5);
548        while is_running(child.trim()) {
549            assert!(
550                Instant::now() < deadline,
551                "what the harness started was left running"
552            );
553            thread::sleep(Duration::from_millis(10));
554        }
555        let flood = stub(directory.path(), "flood", "exec yes model");
556        assert!(matches!(
557            run(&flood, &[], &path, None, "pi", LIST_TIMEOUT, &running),
558            Err(ModelListError::Unreadable("pi"))
559        ));
560    }
561
562    #[test]
563    fn completed_helpers_leave_no_background_children() {
564        let directory = tempfile::tempdir().unwrap();
565        let pid_path = directory.path().join("child.pid");
566        for exit_code in [0, 7] {
567            let program = stub(
568                directory.path(),
569                "helper",
570                &format!(
571                    "sleep 60 & echo $! > '{}'; printf 'catalog'; exit {exit_code}",
572                    pid_path.display()
573                ),
574            );
575            let result = run(
576                &program,
577                &[],
578                &OsString::from("/bin:/usr/bin"),
579                None,
580                "fixture",
581                LIST_TIMEOUT,
582                &AtomicBool::new(false),
583            );
584            let child = FixtureProcess::read(&pid_path);
585            if exit_code == 0 {
586                assert_eq!(result.unwrap(), "catalog");
587            } else {
588                assert!(matches!(result, Err(ModelListError::Failed("fixture"))));
589            }
590            child.assert_stopped();
591        }
592    }
593
594    #[test]
595    fn dropping_a_request_stops_its_harness() {
596        let directory = tempfile::tempdir().unwrap();
597        let pid_file = directory.path().join("pid");
598        let slow = stub(
599            directory.path(),
600            "pi",
601            &format!("echo $$ > '{}'; exec sleep 60", pid_file.display()),
602        );
603        let request = ModelListRequest::spawn(HarnessId::Pi, None, move || {
604            Ok(LocatedHarness {
605                program: slow,
606                path: OsString::from("/bin:/usr/bin"),
607            })
608        });
609        let deadline = Instant::now() + Duration::from_secs(10);
610        while std::fs::read_to_string(&pid_file).map_or(true, |pid| pid.trim().is_empty()) {
611            assert!(Instant::now() < deadline, "the harness never started");
612            thread::sleep(Duration::from_millis(10));
613        }
614        assert!(request.poll().is_none(), "still listing");
615        drop(request);
616        let pid = std::fs::read_to_string(&pid_file).unwrap();
617        assert!(
618            !is_running(pid.trim()),
619            "request destruction waits for the harness to stop"
620        );
621    }
622}
twine_core/harness/models/discovery.rs

twine_core/harness/models/
discovery.rs
