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

handshake_core/harness/models/
discovery.rs
