1use std::collections::HashSet;
2
3use serde::Serialize;
4
5use super::probes::OpenCodeProbeResult;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
8#[serde(rename_all = "snake_case")]
9pub enum AvailabilityStatus {
10 Runnable,
11 Unavailable,
12 Unknown,
13}
14
15#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
16#[serde(rename_all = "snake_case")]
17pub enum AvailabilitySource {
18 HarnessInstalled,
19 #[serde(rename = "opencode_probe")]
20 OpenCodeProbe,
21 #[serde(rename = "opencode_probe_negative")]
22 OpenCodeProbeNegative,
23 #[serde(rename = "opencode_probe_unknown")]
24 OpenCodeProbeUnknown,
25 NoHarness,
26 Offline,
27}
28
29#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
31pub struct RunnablePath {
32 pub harness: String,
33 pub mars_provider: String,
34 pub harness_model_id: String,
35}
36
37#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
39pub struct ModelAvailability {
40 pub status: AvailabilityStatus,
41 pub source: AvailabilitySource,
42 pub runnable_paths: Vec<RunnablePath>,
43}
44
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub struct DecomposedSlug {
47 pub oc_provider: String,
48 pub upstream_provider: Option<String>,
49 pub model_part: String,
50 pub full_slug: String,
51}
52
53pub fn decompose_slug(slug: &str) -> Option<DecomposedSlug> {
54 let parts: Vec<&str> = slug.split('/').collect();
55 match parts.as_slice() {
56 [oc_provider, model_part] if !oc_provider.is_empty() && !model_part.is_empty() => {
57 Some(DecomposedSlug {
58 oc_provider: (*oc_provider).to_string(),
59 upstream_provider: None,
60 model_part: (*model_part).to_string(),
61 full_slug: slug.to_string(),
62 })
63 }
64 [oc_provider, upstream_provider, model_part]
65 if !oc_provider.is_empty()
66 && !upstream_provider.is_empty()
67 && !model_part.is_empty() =>
68 {
69 Some(DecomposedSlug {
70 oc_provider: (*oc_provider).to_string(),
71 upstream_provider: Some((*upstream_provider).to_string()),
72 model_part: (*model_part).to_string(),
73 full_slug: slug.to_string(),
74 })
75 }
76 _ => None,
77 }
78}
79
80pub fn normalize_model_id(id: &str) -> String {
81 id.to_lowercase().replace('.', "-")
82}
83
84pub fn model_id_matches(mars_id: &str, oc_model: &str) -> bool {
85 normalize_model_id(mars_id) == normalize_model_id(oc_model)
86}
87
88pub fn provider_matches(mars_provider: &str, oc_segment: &str) -> bool {
89 mars_provider.eq_ignore_ascii_case(oc_segment)
90}
91
92pub fn classify_for_harness(
94 harness: &str,
95 provider: &str,
96 model_id: &str,
97 installed: &HashSet<String>,
98 probe_result: Option<&OpenCodeProbeResult>,
99) -> Option<(AvailabilityStatus, AvailabilitySource, Option<RunnablePath>)> {
100 let harness = harness.to_ascii_lowercase();
101 if !installed.contains(&harness) {
102 return Some((
103 AvailabilityStatus::Unavailable,
104 AvailabilitySource::NoHarness,
105 None,
106 ));
107 }
108
109 let direct_match = match harness.as_str() {
110 "claude" => provider_matches(provider, "anthropic"),
111 "codex" => provider_matches(provider, "openai"),
112 "gemini" => provider_matches(provider, "google"),
113 "opencode" => return classify_opencode(provider, model_id, probe_result),
114 _ => false,
115 };
116
117 if direct_match {
118 Some((
119 AvailabilityStatus::Runnable,
120 AvailabilitySource::HarnessInstalled,
121 Some(RunnablePath {
122 harness,
123 mars_provider: provider.to_string(),
124 harness_model_id: model_id.to_string(),
125 }),
126 ))
127 } else {
128 Some((
129 AvailabilityStatus::Unavailable,
130 AvailabilitySource::NoHarness,
131 None,
132 ))
133 }
134}
135
136fn classify_opencode(
137 provider: &str,
138 model_id: &str,
139 probe_result: Option<&OpenCodeProbeResult>,
140) -> Option<(AvailabilityStatus, AvailabilitySource, Option<RunnablePath>)> {
141 let Some(probe) = probe_result else {
142 return Some((
143 AvailabilityStatus::Unknown,
144 AvailabilitySource::OpenCodeProbeUnknown,
145 None,
146 ));
147 };
148
149 if !probe.provider_probe_success {
150 return Some((
151 AvailabilityStatus::Unknown,
152 AvailabilitySource::OpenCodeProbeUnknown,
153 None,
154 ));
155 }
156
157 let provider_lower = provider.to_lowercase();
158 let has_provider = probe
159 .providers
160 .get(&provider_lower)
161 .copied()
162 .unwrap_or(false);
163 let has_openrouter = probe.providers.get("openrouter").copied().unwrap_or(false);
164 let has_via_openrouter = has_openrouter && openrouter_supports_provider(&provider_lower);
165
166 if !has_provider && !has_via_openrouter {
167 return Some((
168 AvailabilityStatus::Unavailable,
169 AvailabilitySource::OpenCodeProbeNegative,
170 None,
171 ));
172 }
173
174 let harness_model_id = if probe.model_probe_success {
175 find_matching_slug(model_id, provider, &probe.model_slugs)
176 } else {
177 None
178 }
179 .unwrap_or_else(|| {
180 if has_via_openrouter && !has_provider {
181 format!("openrouter/{provider_lower}/{model_id}")
182 } else {
183 format!("{provider_lower}/{model_id}")
184 }
185 });
186
187 Some((
188 AvailabilityStatus::Runnable,
189 AvailabilitySource::OpenCodeProbe,
190 Some(RunnablePath {
191 harness: "opencode".to_string(),
192 mars_provider: provider.to_string(),
193 harness_model_id,
194 }),
195 ))
196}
197
198fn openrouter_supports_provider(provider: &str) -> bool {
199 matches!(
200 provider,
201 "anthropic" | "meta" | "mistral" | "deepseek" | "cohere"
202 )
203}
204
205fn find_matching_slug(
206 mars_model_id: &str,
207 mars_provider: &str,
208 slugs: &[String],
209) -> Option<String> {
210 for slug in slugs {
211 let Some(decomposed) = decompose_slug(slug) else {
212 continue;
213 };
214 let effective_provider = decomposed
215 .upstream_provider
216 .as_deref()
217 .unwrap_or(&decomposed.oc_provider);
218
219 if provider_matches(mars_provider, effective_provider)
220 && model_id_matches(mars_model_id, &decomposed.model_part)
221 {
222 return Some(slug.clone());
223 }
224 }
225
226 None
227}
228
229pub fn classify_model(
230 model_id: &str,
231 provider: &str,
232 installed: &HashSet<String>,
233 probe_result: Option<&OpenCodeProbeResult>,
234 offline: bool,
235) -> ModelAvailability {
236 let mut statuses = Vec::new();
237 let mut runnable_paths = Vec::new();
238
239 for harness in ["claude", "codex", "gemini"] {
240 let Some((status, source, path)) =
241 classify_for_harness(harness, provider, model_id, installed, None)
242 else {
243 continue;
244 };
245 if let Some(path) = path {
246 runnable_paths.push(path);
247 }
248 statuses.push((status, source));
249 }
250
251 if installed.contains("opencode") {
252 if offline {
253 statuses.push((AvailabilityStatus::Unknown, AvailabilitySource::Offline));
254 } else if let Some(result) = probe_result {
255 if let Some((status, source, path)) =
256 classify_for_harness("opencode", provider, model_id, installed, Some(result))
257 {
258 if let Some(path) = path {
259 runnable_paths.push(path);
260 }
261 statuses.push((status, source));
262 }
263 } else {
264 statuses.push((
265 AvailabilityStatus::Unknown,
266 AvailabilitySource::OpenCodeProbeUnknown,
267 ));
268 }
269 }
270
271 aggregate_statuses(statuses, runnable_paths)
272}
273
274fn aggregate_statuses(
275 statuses: Vec<(AvailabilityStatus, AvailabilitySource)>,
276 runnable_paths: Vec<RunnablePath>,
277) -> ModelAvailability {
278 if statuses.is_empty() {
279 return ModelAvailability {
280 status: AvailabilityStatus::Unavailable,
281 source: AvailabilitySource::NoHarness,
282 runnable_paths: Vec::new(),
283 };
284 }
285
286 if statuses
287 .iter()
288 .any(|(status, _)| *status == AvailabilityStatus::Runnable)
289 {
290 return ModelAvailability {
291 status: AvailabilityStatus::Runnable,
292 source: statuses
293 .iter()
294 .find_map(|(status, source)| {
295 (*status == AvailabilityStatus::Runnable).then(|| source.clone())
296 })
297 .expect("runnable status exists"),
298 runnable_paths,
299 };
300 }
301
302 if statuses
303 .iter()
304 .any(|(status, _)| *status == AvailabilityStatus::Unknown)
305 {
306 return ModelAvailability {
307 status: AvailabilityStatus::Unknown,
308 source: statuses
309 .iter()
310 .find_map(|(status, source)| {
311 (*status == AvailabilityStatus::Unknown).then(|| source.clone())
312 })
313 .unwrap_or(AvailabilitySource::OpenCodeProbeUnknown),
314 runnable_paths: Vec::new(),
315 };
316 }
317
318 ModelAvailability {
319 status: AvailabilityStatus::Unavailable,
320 source: statuses
321 .iter()
322 .find_map(|(_, source)| {
323 (*source != AvailabilitySource::NoHarness).then(|| source.clone())
324 })
325 .or_else(|| statuses.first().map(|(_, source)| source.clone()))
326 .unwrap_or(AvailabilitySource::NoHarness),
327 runnable_paths: Vec::new(),
328 }
329}
330
331#[cfg(test)]
332mod tests {
333 use super::*;
334 use std::collections::HashMap;
335
336 fn installed(names: &[&str]) -> HashSet<String> {
337 names.iter().map(|name| (*name).to_string()).collect()
338 }
339
340 #[test]
341 fn test_decompose_slug_two_segments() {
342 let slug = decompose_slug("openai/gpt-5.4").unwrap();
343 assert_eq!(slug.oc_provider, "openai");
344 assert_eq!(slug.upstream_provider, None);
345 assert_eq!(slug.model_part, "gpt-5.4");
346 assert_eq!(slug.full_slug, "openai/gpt-5.4");
347 }
348
349 #[test]
350 fn test_decompose_slug_three_segments() {
351 let slug = decompose_slug("openrouter/anthropic/claude-opus-4.7").unwrap();
352 assert_eq!(slug.oc_provider, "openrouter");
353 assert_eq!(slug.upstream_provider.as_deref(), Some("anthropic"));
354 assert_eq!(slug.model_part, "claude-opus-4.7");
355 assert_eq!(slug.full_slug, "openrouter/anthropic/claude-opus-4.7");
356 }
357
358 #[test]
359 fn test_decompose_slug_invalid() {
360 assert!(decompose_slug("gpt-5").is_none());
361 assert!(decompose_slug("openai/").is_none());
362 assert!(decompose_slug("a/b/c/d").is_none());
363 }
364
365 #[test]
366 fn test_normalize_model_id() {
367 assert_eq!(normalize_model_id("Claude-Opus-4.7"), "claude-opus-4-7");
368 }
369
370 #[test]
371 fn test_model_id_matches() {
372 assert!(model_id_matches("claude-opus-4-7", "Claude-Opus-4.7"));
373 assert!(!model_id_matches("claude-opus-4-7", "claude-sonnet-4-7"));
374 }
375
376 #[test]
377 fn test_provider_matches() {
378 assert!(provider_matches("Anthropic", "anthropic"));
379 assert!(!provider_matches("Anthropic", "openai"));
380 }
381
382 #[test]
383 fn test_classify_claude_anthropic() {
384 let result = classify_for_harness(
385 "claude",
386 "Anthropic",
387 "claude-opus-4-7",
388 &installed(&["claude"]),
389 None,
390 )
391 .unwrap();
392 assert_eq!(result.0, AvailabilityStatus::Runnable);
393 assert_eq!(result.1, AvailabilitySource::HarnessInstalled);
394 assert_eq!(
395 result.2.unwrap().harness_model_id,
396 "claude-opus-4-7".to_string()
397 );
398 }
399
400 #[test]
401 fn test_classify_codex_openai() {
402 let result =
403 classify_for_harness("codex", "OpenAI", "gpt-5.4", &installed(&["codex"]), None)
404 .unwrap();
405 assert_eq!(result.0, AvailabilityStatus::Runnable);
406 assert_eq!(result.1, AvailabilitySource::HarnessInstalled);
407 }
408
409 #[test]
410 fn test_classify_gemini_google() {
411 let result = classify_for_harness(
412 "gemini",
413 "Google",
414 "gemini-2.5-pro",
415 &installed(&["gemini"]),
416 None,
417 )
418 .unwrap();
419 assert_eq!(result.0, AvailabilityStatus::Runnable);
420 assert_eq!(result.1, AvailabilitySource::HarnessInstalled);
421 }
422
423 #[test]
424 fn test_classify_no_harness() {
425 let result = classify_for_harness(
426 "claude",
427 "Anthropic",
428 "claude-opus-4-7",
429 &installed(&[]),
430 None,
431 )
432 .unwrap();
433 assert_eq!(result.0, AvailabilityStatus::Unavailable);
434 assert_eq!(result.1, AvailabilitySource::NoHarness);
435 assert!(result.2.is_none());
436 }
437
438 #[test]
439 fn test_classify_multi_harness_any_runnable() {
440 let result = classify_model(
441 "claude-opus-4-7",
442 "Anthropic",
443 &installed(&["claude", "codex"]),
444 None,
445 false,
446 );
447 assert_eq!(result.status, AvailabilityStatus::Runnable);
448 assert_eq!(result.source, AvailabilitySource::HarnessInstalled);
449 assert_eq!(result.runnable_paths.len(), 1);
450 assert_eq!(result.runnable_paths[0].harness, "claude");
451 }
452
453 #[test]
454 fn test_classify_multi_harness_all_unavailable() {
455 let result = classify_model("custom-model", "Unknown", &installed(&[]), None, false);
456 assert_eq!(result.status, AvailabilityStatus::Unavailable);
457 assert_eq!(result.source, AvailabilitySource::NoHarness);
458 assert!(result.runnable_paths.is_empty());
459 }
460
461 #[test]
462 fn test_classify_offline_mode() {
463 let result = classify_model("gpt-5.4", "OpenAI", &installed(&["codex"]), None, true);
464 assert_eq!(result.status, AvailabilityStatus::Runnable);
465 assert_eq!(result.source, AvailabilitySource::HarnessInstalled);
466 assert_eq!(result.runnable_paths.len(), 1);
467 assert_eq!(result.runnable_paths[0].harness, "codex");
468
469 let result = classify_model("gpt-5.4", "OpenAI", &installed(&["opencode"]), None, true);
470 assert_eq!(result.status, AvailabilityStatus::Unknown);
471 assert_eq!(result.source, AvailabilitySource::Offline);
472 assert!(result.runnable_paths.is_empty());
473 }
474
475 #[test]
476 fn test_classify_opencode_direct_slug() {
477 let probe = OpenCodeProbeResult {
478 providers: HashMap::from([("openai".to_string(), true)]),
479 model_slugs: vec!["openai/gpt-5.4".to_string()],
480 provider_probe_success: true,
481 model_probe_success: true,
482 error: None,
483 };
484
485 let result = classify_model(
486 "gpt-5.4",
487 "OpenAI",
488 &installed(&["opencode"]),
489 Some(&probe),
490 false,
491 );
492
493 assert_eq!(result.status, AvailabilityStatus::Runnable);
494 assert_eq!(result.source, AvailabilitySource::OpenCodeProbe);
495 assert_eq!(result.runnable_paths.len(), 1);
496 assert_eq!(result.runnable_paths[0].harness, "opencode");
497 assert_eq!(result.runnable_paths[0].harness_model_id, "openai/gpt-5.4");
498 }
499
500 #[test]
501 fn test_classify_opencode_openrouter_slug() {
502 let probe = OpenCodeProbeResult {
503 providers: HashMap::from([("openrouter".to_string(), true)]),
504 model_slugs: vec!["openrouter/anthropic/claude-opus-4.7".to_string()],
505 provider_probe_success: true,
506 model_probe_success: true,
507 error: None,
508 };
509
510 let result = classify_model(
511 "claude-opus-4-7",
512 "Anthropic",
513 &installed(&["opencode"]),
514 Some(&probe),
515 false,
516 );
517
518 assert_eq!(result.status, AvailabilityStatus::Runnable);
519 assert_eq!(result.source, AvailabilitySource::OpenCodeProbe);
520 assert_eq!(
521 result.runnable_paths[0].harness_model_id,
522 "openrouter/anthropic/claude-opus-4.7"
523 );
524 }
525
526 #[test]
527 fn test_classify_opencode_provider_negative() {
528 let probe = OpenCodeProbeResult {
529 providers: HashMap::from([("google".to_string(), true)]),
530 provider_probe_success: true,
531 ..OpenCodeProbeResult::default()
532 };
533
534 let result = classify_model(
535 "gpt-5.4",
536 "OpenAI",
537 &installed(&["opencode"]),
538 Some(&probe),
539 false,
540 );
541
542 assert_eq!(result.status, AvailabilityStatus::Unavailable);
543 assert_eq!(result.source, AvailabilitySource::OpenCodeProbeNegative);
544 assert!(result.runnable_paths.is_empty());
545 }
546
547 #[test]
548 fn test_classify_opencode_empty_providers() {
549 let probe = OpenCodeProbeResult {
550 providers: HashMap::new(),
551 model_slugs: Vec::new(),
552 provider_probe_success: true,
553 model_probe_success: true,
554 error: None,
555 };
556
557 let result = classify_model(
558 "claude-opus-4-7",
559 "Anthropic",
560 &installed(&["opencode"]),
561 Some(&probe),
562 false,
563 );
564
565 assert_eq!(result.status, AvailabilityStatus::Unavailable);
566 assert_eq!(result.source, AvailabilitySource::OpenCodeProbeNegative);
567 assert!(result.runnable_paths.is_empty());
568 }
569
570 #[test]
571 fn test_classify_opencode_no_matching_slug() {
572 let probe = OpenCodeProbeResult {
573 providers: HashMap::from([("anthropic".to_string(), true)]),
574 model_slugs: vec!["anthropic/claude-3-5-sonnet".to_string()],
575 provider_probe_success: true,
576 model_probe_success: true,
577 error: None,
578 };
579
580 let result = classify_model(
581 "claude-opus-4-7",
582 "Anthropic",
583 &installed(&["opencode"]),
584 Some(&probe),
585 false,
586 );
587
588 assert_eq!(result.status, AvailabilityStatus::Runnable);
589 assert_eq!(result.source, AvailabilitySource::OpenCodeProbe);
590 assert_eq!(result.runnable_paths.len(), 1);
591 assert_eq!(
592 result.runnable_paths[0].harness_model_id,
593 "anthropic/claude-opus-4-7"
594 );
595 }
596
597 #[test]
598 fn test_classify_opencode_synthesizes_slug_when_model_probe_fails() {
599 let probe = OpenCodeProbeResult {
600 providers: HashMap::from([("anthropic".to_string(), true)]),
601 provider_probe_success: true,
602 model_probe_success: false,
603 error: Some("model probe failed: timeout".to_string()),
604 ..OpenCodeProbeResult::default()
605 };
606
607 let result = classify_model(
608 "claude-opus-4-7",
609 "Anthropic",
610 &installed(&["opencode"]),
611 Some(&probe),
612 false,
613 );
614
615 assert_eq!(result.status, AvailabilityStatus::Runnable);
616 assert_eq!(result.source, AvailabilitySource::OpenCodeProbe);
617 assert_eq!(result.runnable_paths.len(), 1);
618 assert_eq!(
619 result.runnable_paths[0].harness_model_id,
620 "anthropic/claude-opus-4-7"
621 );
622 }
623
624 #[test]
625 fn test_classify_opencode_unknown_when_probe_fails() {
626 let probe = OpenCodeProbeResult {
627 error: Some("provider probe failed: timeout".to_string()),
628 ..OpenCodeProbeResult::default()
629 };
630
631 let result = classify_model(
632 "gpt-5.4",
633 "OpenAI",
634 &installed(&["opencode"]),
635 Some(&probe),
636 false,
637 );
638
639 assert_eq!(result.status, AvailabilityStatus::Unknown);
640 assert_eq!(result.source, AvailabilitySource::OpenCodeProbeUnknown);
641 assert!(result.runnable_paths.is_empty());
642 }
643}