1use loopsmith_core::{LoopConfig, ProviderKind, ProviderSpec, Tier};
20use serde::{Deserialize, Serialize};
21use std::collections::BTreeMap;
22use std::io::Write;
23use std::path::PathBuf;
24use std::process::{Command, Stdio};
25use std::time::{Duration, Instant};
26
27#[derive(Debug, thiserror::Error)]
28pub enum ProviderError {
29 #[error("no provider available for tier {tier:?}; tried: {tried}")]
30 NoneAvailable { tier: Tier, tried: String },
31 #[error("provider `{id}` failed to start: {source}")]
32 Spawn {
33 id: String,
34 #[source]
35 source: std::io::Error,
36 },
37 #[error("provider `{id}` timed out after {seconds}s")]
38 Timeout { id: String, seconds: u64 },
39 #[error("provider `{id}` exited {code}: {stderr}")]
40 Failed {
41 id: String,
42 code: i32,
43 stderr: String,
44 },
45}
46
47#[derive(Debug, Clone)]
48pub struct InvokeRequest {
49 pub node_id: String,
51 pub system: String,
52 pub prompt: String,
53 pub tier: Tier,
54 pub workdir: PathBuf,
55}
56
57#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct InvokeResponse {
59 pub provider_id: String,
60 pub output: String,
61 pub exit_code: i32,
62 pub duration_ms: u64,
63 #[serde(default)]
64 pub stderr_tail: Option<String>,
65 #[serde(default)]
68 pub tokens: Option<u64>,
69 #[serde(default)]
72 pub tokens_estimated: bool,
73 #[serde(default)]
74 pub cost_usd: Option<f64>,
75}
76
77pub fn estimate_tokens(prompt: &str, output: &str) -> u64 {
84 ((prompt.chars().count() + output.chars().count()) as u64).div_ceil(4)
85}
86
87pub fn parse_usage(spec: &ProviderSpec, text: &str) -> Option<u64> {
89 let pattern = spec.usage_regex.as_ref()?;
90 let re = regex::Regex::new(pattern).ok()?;
91 let caps = re.captures(text)?;
92 let raw = caps.get(1).or_else(|| caps.get(0))?.as_str();
94 raw.trim().replace([',', '_'], "").parse::<u64>().ok()
95}
96
97pub fn digest(s: &str) -> String {
101 let mut h: u64 = 0xcbf29ce484222325;
103 for b in s.as_bytes() {
104 h ^= *b as u64;
105 h = h.wrapping_mul(0x100000001b3);
106 }
107 format!("{h:016x}")
108}
109
110pub fn availability(spec: &ProviderSpec) -> Availability {
112 let missing_env: Vec<String> = spec
113 .requires_env
114 .iter()
115 .filter(|k| std::env::var_os(k).is_none())
116 .cloned()
117 .collect();
118 let on_path = which(&spec.command).is_some();
119 Availability {
120 on_path,
121 missing_env,
122 }
123}
124
125#[derive(Debug, Clone)]
126pub struct Availability {
127 pub on_path: bool,
128 pub missing_env: Vec<String>,
130}
131
132impl Availability {
133 pub fn ok(&self) -> bool {
134 self.on_path && self.missing_env.is_empty()
135 }
136 pub fn why_not(&self) -> String {
137 let mut parts = Vec::new();
138 if !self.on_path {
139 parts.push("command not found on PATH".to_string());
140 }
141 if !self.missing_env.is_empty() {
142 parts.push(format!("missing env: {}", self.missing_env.join(", ")));
143 }
144 parts.join("; ")
145 }
146}
147
148pub use loopsmith_util::which;
152
153pub fn render(template: &str, vars: &BTreeMap<&str, &str>) -> String {
155 let mut out = template.to_string();
156 for (k, v) in vars {
157 out = out.replace(&format!("{{{k}}}"), v);
158 }
159 out
160}
161
162fn tier_name(t: Tier) -> &'static str {
163 match t {
164 Tier::Cheap => "cheap",
165 Tier::Standard => "standard",
166 Tier::Strong => "strong",
167 }
168}
169
170pub fn invoke(spec: &ProviderSpec, req: &InvokeRequest) -> Result<InvokeResponse, ProviderError> {
172 let model = spec.model.clone().unwrap_or_default();
173 let tier = tier_name(req.tier);
174 let vars: BTreeMap<&str, &str> = [
175 ("prompt", req.prompt.as_str()),
176 ("system", req.system.as_str()),
177 ("model", model.as_str()),
178 ("tier", tier),
179 ("node", req.node_id.as_str()),
180 ]
181 .into_iter()
182 .collect();
183
184 let args: Vec<String> = spec.args.iter().map(|a| render(a, &vars)).collect();
185
186 let mut cmd = Command::new(&spec.command);
187 cmd.args(&args)
188 .current_dir(&req.workdir)
189 .stdin(if spec.prompt_on_stdin {
190 Stdio::piped()
191 } else {
192 Stdio::null()
193 })
194 .stdout(Stdio::piped())
195 .stderr(Stdio::piped());
196
197 let started = Instant::now();
198 let mut child = cmd.spawn().map_err(|source| ProviderError::Spawn {
199 id: spec.id.clone(),
200 source,
201 })?;
202
203 if spec.prompt_on_stdin {
204 if let Some(mut sin) = child.stdin.take() {
205 let _ = sin.write_all(req.prompt.as_bytes());
208 }
209 }
210
211 let timeout = spec.timeout_seconds.map(Duration::from_secs);
214 let poll = Duration::from_millis(50);
215 loop {
216 match child.try_wait() {
217 Ok(Some(_)) => break,
218 Ok(None) => {
219 if let Some(limit) = timeout {
220 if started.elapsed() >= limit {
221 let _ = child.kill();
222 let _ = child.wait();
223 return Err(ProviderError::Timeout {
224 id: spec.id.clone(),
225 seconds: limit.as_secs(),
226 });
227 }
228 }
229 std::thread::sleep(poll);
230 }
231 Err(source) => {
232 return Err(ProviderError::Spawn {
233 id: spec.id.clone(),
234 source,
235 })
236 }
237 }
238 }
239
240 let out = child
241 .wait_with_output()
242 .map_err(|source| ProviderError::Spawn {
243 id: spec.id.clone(),
244 source,
245 })?;
246 let code = out.status.code().unwrap_or(-1);
247 let stderr = String::from_utf8_lossy(&out.stderr).to_string();
248
249 if code != 0 {
250 return Err(ProviderError::Failed {
251 id: spec.id.clone(),
252 code,
253 stderr: stderr.lines().last().unwrap_or("").to_string(),
254 });
255 }
256
257 let stdout = String::from_utf8_lossy(&out.stdout).to_string();
258
259 let reported = parse_usage(spec, &stdout).or_else(|| parse_usage(spec, &stderr));
262 let (tokens, estimated) = match reported {
263 Some(t) => (Some(t), false),
264 None => (Some(estimate_tokens(&req.prompt, &stdout)), true),
265 };
266 let cost = match (tokens, spec.cost_per_1k_tokens) {
267 (Some(t), Some(rate)) => Some((t as f64 / 1000.0) * rate),
268 _ => None,
269 };
270
271 Ok(InvokeResponse {
272 provider_id: spec.id.clone(),
273 output: stdout,
274 exit_code: code,
275 duration_ms: started.elapsed().as_millis() as u64,
276 stderr_tail: stderr.lines().last().map(|s| s.to_string()),
277 tokens,
278 tokens_estimated: estimated,
279 cost_usd: cost,
280 })
281}
282
283pub fn dispatch(
287 cfg: &LoopConfig,
288 req: &InvokeRequest,
289 pinned: Option<&str>,
290) -> Result<(InvokeResponse, Vec<String>), ProviderError> {
291 let candidates: Vec<&ProviderSpec> = match pinned {
292 Some(id) => cfg.provider(id).into_iter().collect(),
293 None => cfg.cascade_for(req.tier),
294 };
295
296 let mut skipped = Vec::new();
297 for spec in &candidates {
298 let av = availability(spec);
299 if !av.ok() {
300 skipped.push(format!("{} ({})", spec.id, av.why_not()));
301 continue;
302 }
303 match invoke(spec, req) {
304 Ok(resp) => return Ok((resp, skipped)),
305 Err(e) => skipped.push(format!("{}: {e}", spec.id)),
306 }
307 }
308
309 Err(ProviderError::NoneAvailable {
310 tier: req.tier,
311 tried: if skipped.is_empty() {
312 "none declared".to_string()
313 } else {
314 skipped.join("; ")
315 },
316 })
317}
318
319pub fn starter_providers() -> Vec<ProviderSpec> {
324 vec![
325 ProviderSpec {
326 id: "claude".into(),
327 kind: ProviderKind::ClaudeCode,
328 tiers: vec![Tier::Standard, Tier::Strong],
329 command: "claude".into(),
330 args: vec!["-p".into(), "{prompt}".into()],
331 model: None,
332 requires_env: vec![],
333 timeout_seconds: Some(900),
334 prompt_on_stdin: false,
335 usage_regex: None,
336 cost_per_1k_tokens: None,
337 },
338 ProviderSpec {
339 id: "ollama".into(),
340 kind: ProviderKind::Ollama,
341 tiers: vec![Tier::Cheap],
342 command: "ollama".into(),
343 args: vec!["run".into(), "{model}".into()],
344 model: Some("llama3".into()),
345 requires_env: vec![],
346 timeout_seconds: Some(120),
355 prompt_on_stdin: true,
356 usage_regex: None,
357 cost_per_1k_tokens: None,
358 },
359 ProviderSpec {
360 id: "grok".into(),
361 kind: ProviderKind::GrokCli,
362 tiers: vec![Tier::Standard],
363 command: "grok".into(),
364 args: vec!["-p".into(), "{prompt}".into()],
365 model: None,
366 requires_env: vec!["XAI_API_KEY".into()],
367 timeout_seconds: Some(600),
368 prompt_on_stdin: false,
369 usage_regex: None,
370 cost_per_1k_tokens: None,
371 },
372 ProviderSpec {
373 id: "openai".into(),
374 kind: ProviderKind::OpenAi,
375 tiers: vec![Tier::Strong],
376 command: "curl".into(),
377 args: vec![
378 "-sS".into(),
379 "https://api.openai.com/v1/chat/completions".into(),
380 "-H".into(),
381 "Content-Type: application/json".into(),
382 "-H".into(),
383 "Authorization: Bearer $OPENAI_API_KEY".into(),
386 "-d".into(),
387 "@-".into(),
388 ],
389 model: Some("gpt-4o-mini".into()),
390 requires_env: vec!["OPENAI_API_KEY".into()],
391 timeout_seconds: Some(300),
392 prompt_on_stdin: true,
393 usage_regex: Some(r#""total_tokens"\s*:\s*(\d+)"#.into()),
396 cost_per_1k_tokens: Some(0.0006),
397 },
398 ProviderSpec {
399 id: "gemini".into(),
400 kind: ProviderKind::Gemini,
401 tiers: vec![Tier::Standard],
402 command: "gemini".into(),
403 args: vec!["-p".into(), "{prompt}".into()],
404 model: None,
405 requires_env: vec!["GEMINI_API_KEY".into()],
406 timeout_seconds: Some(600),
407 prompt_on_stdin: false,
408 usage_regex: None,
409 cost_per_1k_tokens: None,
410 },
411 ]
412}
413
414#[cfg(test)]
415mod tests {
416 use super::*;
417
418 fn spec(id: &str, command: &str, args: &[&str]) -> ProviderSpec {
419 ProviderSpec {
420 id: id.into(),
421 kind: ProviderKind::Byok,
422 tiers: vec![Tier::Cheap, Tier::Standard, Tier::Strong],
423 command: command.into(),
424 args: args.iter().map(|s| s.to_string()).collect(),
425 model: None,
426 requires_env: vec![],
427 timeout_seconds: Some(30),
428 prompt_on_stdin: false,
429 usage_regex: None,
430 cost_per_1k_tokens: None,
431 }
432 }
433
434 fn req() -> InvokeRequest {
435 InvokeRequest {
436 node_id: "n1".into(),
437 system: "be terse".into(),
438 prompt: "hello".into(),
439 tier: Tier::Standard,
440 workdir: std::env::temp_dir(),
441 }
442 }
443
444 #[test]
445 fn placeholders_are_substituted() {
446 let vars: BTreeMap<&str, &str> =
447 [("prompt", "hi"), ("model", "m1")].into_iter().collect();
448 assert_eq!(render("say {prompt} via {model}", &vars), "say hi via m1");
449 }
450
451 #[test]
452 fn unknown_placeholders_are_left_alone() {
453 let vars: BTreeMap<&str, &str> = [("prompt", "hi")].into_iter().collect();
454 assert_eq!(render("{prompt} {unknown}", &vars), "hi {unknown}");
455 }
456
457 #[test]
458 fn digest_is_stable_and_differentiating() {
459 assert_eq!(digest("abc"), digest("abc"));
460 assert_ne!(digest("abc"), digest("abd"));
461 }
462
463 #[test]
464 fn which_finds_a_real_binary_and_misses_a_fake_one() {
465 assert!(which("sh").is_some());
466 assert!(which("definitely-not-a-real-binary-xyz").is_none());
467 }
468
469 #[test]
470 fn a_provider_with_a_missing_binary_is_unavailable() {
471 let s = spec("ghost", "definitely-not-a-real-binary-xyz", &[]);
472 let av = availability(&s);
473 assert!(!av.ok());
474 assert!(av.why_not().contains("not found on PATH"));
475 }
476
477 #[test]
478 fn a_provider_with_missing_env_is_unavailable_and_names_only_the_key() {
479 let mut s = spec("needs-key", "sh", &[]);
480 s.requires_env = vec!["LOOPSMITH_TEST_ABSENT_KEY".into()];
481 let av = availability(&s);
482 assert!(!av.ok());
483 let why = av.why_not();
484 assert!(why.contains("LOOPSMITH_TEST_ABSENT_KEY"));
485 assert!(!why.contains('='));
487 }
488
489 #[test]
490 fn invoking_echo_returns_its_stdout() {
491 let s = spec("echoer", "echo", &["{prompt}"]);
492 let r = invoke(&s, &req()).expect("echo runs");
493 assert_eq!(r.output.trim(), "hello");
494 assert_eq!(r.exit_code, 0);
495 assert_eq!(r.provider_id, "echoer");
496 }
497
498 #[test]
499 fn stdin_mode_pipes_the_prompt() {
500 let s = {
501 let mut s = spec("catter", "cat", &[]);
502 s.prompt_on_stdin = true;
503 s
504 };
505 let r = invoke(&s, &req()).expect("cat runs");
506 assert_eq!(r.output.trim(), "hello");
507 }
508
509 #[test]
510 fn a_nonzero_exit_is_an_error_not_a_silent_pass() {
511 let s = spec("failer", "false", &[]);
512 let e = invoke(&s, &req()).unwrap_err();
513 assert!(matches!(e, ProviderError::Failed { .. }));
514 }
515
516 #[test]
517 fn a_hanging_provider_is_killed_at_the_timeout() {
518 let mut s = spec("sleeper", "sleep", &["30"]);
519 s.timeout_seconds = Some(1);
520 let started = Instant::now();
521 let e = invoke(&s, &req()).unwrap_err();
522 assert!(matches!(e, ProviderError::Timeout { .. }));
523 assert!(started.elapsed() < Duration::from_secs(10), "kill was not prompt");
524 }
525
526 fn cfg_with(providers: Vec<ProviderSpec>, cascade: &[(&str, Vec<&str>)]) -> LoopConfig {
527 let mut cfg = loopsmith_core::parse_str(
528 r#"
529name: t
530goals:
531 - name: g1
532 description: a sufficiently long goal description
533validations:
534 - target: g1
535 name: v
536 mode: objective
537 statement: s
538 detector: { type: script, command: "true" }
539"#,
540 "test",
541 )
542 .unwrap();
543 cfg.providers.providers = providers;
544 cfg.providers.cascade = cascade
545 .iter()
546 .map(|(k, v)| (k.to_string(), v.iter().map(|s| s.to_string()).collect()))
547 .collect();
548 cfg
549 }
550
551 #[test]
552 fn cascade_falls_past_an_unavailable_provider() {
553 let cfg = cfg_with(
554 vec![
555 spec("missing", "definitely-not-a-real-binary-xyz", &[]),
556 spec("works", "echo", &["{prompt}"]),
557 ],
558 &[("standard", vec!["missing", "works"])],
559 );
560 let (resp, skipped) = dispatch(&cfg, &req(), None).expect("falls through");
561 assert_eq!(resp.provider_id, "works");
562 assert_eq!(skipped.len(), 1);
563 assert!(skipped[0].contains("missing"));
564 }
565
566 #[test]
567 fn cascade_falls_past_a_provider_that_errors() {
568 let cfg = cfg_with(
569 vec![spec("boom", "false", &[]), spec("works", "echo", &["{prompt}"])],
570 &[("standard", vec!["boom", "works"])],
571 );
572 let (resp, skipped) = dispatch(&cfg, &req(), None).unwrap();
573 assert_eq!(resp.provider_id, "works");
574 assert!(skipped[0].contains("boom"));
575 }
576
577 #[test]
578 fn exhausting_the_cascade_reports_every_attempt() {
579 let cfg = cfg_with(
580 vec![spec("a", "false", &[]), spec("b", "false", &[])],
581 &[("standard", vec!["a", "b"])],
582 );
583 let e = dispatch(&cfg, &req(), None).unwrap_err();
584 let msg = e.to_string();
585 assert!(msg.contains('a') && msg.contains('b'), "{msg}");
586 }
587
588 #[test]
589 fn pinning_a_provider_bypasses_the_cascade() {
590 let cfg = cfg_with(
591 vec![spec("cheap", "false", &[]), spec("pinned", "echo", &["pinned-out"])],
592 &[("standard", vec!["cheap"])],
593 );
594 let (resp, _) = dispatch(&cfg, &req(), Some("pinned")).unwrap();
595 assert_eq!(resp.output.trim(), "pinned-out");
596 }
597
598 #[test]
599 fn tiers_select_different_cascades() {
600 let cfg = cfg_with(
601 vec![
602 spec("small", "echo", &["small"]),
603 spec("big", "echo", &["big"]),
604 ],
605 &[("cheap", vec!["small"]), ("strong", vec!["big"])],
606 );
607 let mut r = req();
608 r.tier = Tier::Cheap;
609 assert_eq!(dispatch(&cfg, &r, None).unwrap().0.output.trim(), "small");
610 r.tier = Tier::Strong;
611 assert_eq!(dispatch(&cfg, &r, None).unwrap().0.output.trim(), "big");
612 }
613
614 #[test]
615 fn usage_is_estimated_when_the_provider_reports_nothing() {
616 let s = spec("echoer", "echo", &["{prompt}"]);
617 let r = invoke(&s, &req()).unwrap();
618 assert!(r.tokens.unwrap() > 0);
619 assert!(r.tokens_estimated, "must be flagged as an estimate");
620 }
621
622 #[test]
623 fn a_usage_regex_extracts_the_real_count() {
624 let mut s = spec("reporter", "echo", &["usage total_tokens=1234 end"]);
630 s.usage_regex = Some(r"total_tokens\s*=\s*(\d+)".into());
631 let r = invoke(&s, &req()).unwrap();
632 assert_eq!(r.tokens, Some(1234));
633 assert!(!r.tokens_estimated, "reported usage is not an estimate");
634 }
635
636 #[test]
637 fn cost_follows_from_tokens_and_rate() {
638 let mut s = spec("priced", "echo", &["hello"]);
639 s.usage_regex = Some(r"(\d+)".into());
640 s.cost_per_1k_tokens = Some(2.0);
641 let mut sp = s.clone();
642 sp.args = vec!["2000".into()];
643 let r = invoke(&sp, &req()).unwrap();
644 assert_eq!(r.tokens, Some(2000));
645 assert!((r.cost_usd.unwrap() - 4.0).abs() < 1e-9);
646 }
647
648 #[test]
649 fn no_rate_means_no_cost_rather_than_a_guessed_one() {
650 let s = spec("free", "echo", &["x"]);
651 let r = invoke(&s, &req()).unwrap();
652 assert!(r.cost_usd.is_none());
653 }
654
655 #[test]
656 fn a_malformed_usage_regex_falls_back_to_estimating() {
657 let mut s = spec("broken", "echo", &["x"]);
658 s.usage_regex = Some("([unclosed".into());
659 let r = invoke(&s, &req()).unwrap();
660 assert!(r.tokens_estimated);
661 }
662
663 #[test]
664 fn starter_providers_cover_every_tier() {
665 let ps = starter_providers();
666 for tier in [Tier::Cheap, Tier::Standard, Tier::Strong] {
667 assert!(
668 ps.iter().any(|p| p.tiers.contains(&tier)),
669 "no starter provider for {tier:?}"
670 );
671 }
672 }
673}