use std::collections::HashMap;
use std::path::PathBuf;
use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
use crate::guard::{self, Readiness};
use crate::schemas::{BillingPolicy, Task, WorkersFile};
use crate::state::Workspace;
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct RoutingOverrides {
#[serde(default)]
pub kind_overrides: HashMap<String, String>,
}
pub fn overrides_path(ws: &Workspace) -> PathBuf {
ws.agents_dir().join("routing-overrides.yaml")
}
pub fn load_overrides(ws: &Workspace) -> RoutingOverrides {
std::fs::read_to_string(overrides_path(ws))
.ok()
.and_then(|t| crate::yaml::from_str(&t).ok())
.unwrap_or_default()
}
pub struct Resolved {
pub worker_id: String,
pub bin: PathBuf,
pub reason: String,
}
pub fn resolve_worker_for_task(
ws: &Workspace,
workers: &WorkersFile,
billing: &BillingPolicy,
override_w: Option<&str>,
task: &Task,
) -> Result<Resolved> {
let overrides = load_overrides(ws);
let required: Vec<String> = task
.required_capabilities
.iter()
.map(|c| norm_cap(c))
.filter(|c| !c.is_empty())
.collect();
let mut candidate = candidate_for_task(workers, &overrides, override_w, task);
if !required.is_empty() && !worker_declares(workers, &candidate.worker_id, &required) {
if candidate.reason == "run override" {
return Err(anyhow!(
"worker '{}' was explicitly selected but does not declare the required \
capability/capabilities {:?}. Add it to that worker in .agents/workers.yaml \
or drop the override.",
candidate.worker_id,
required
));
}
match first_capable(workers, &required) {
Some(id) => {
candidate = Candidate {
worker_id: id,
reason: "capability route",
}
}
None => {
return Err(anyhow!(
"no enabled worker declares the required capability/capabilities \
{required:?}. Add it to a worker in .agents/workers.yaml."
))
}
}
}
resolve_candidate(
workers,
billing,
candidate.worker_id,
candidate.reason,
&required,
)
}
fn resolve_candidate(
workers: &WorkersFile,
billing: &BillingPolicy,
candidate: String,
source: &'static str,
required: &[String],
) -> Result<Resolved> {
let mut order = vec![candidate.clone()];
for w in &workers.routing.fallback_order {
if !order.contains(w) {
order.push(w.clone());
}
}
if !required.is_empty() {
order.retain(|id| worker_declares(workers, id, required));
}
if order.is_empty() {
return Err(anyhow!(
"no enabled worker declares the required capability/capabilities {required:?}. \
Add it to a worker in .agents/workers.yaml."
));
}
let mut tried = Vec::new();
for id in &order {
let Some(profile) = workers.workers.iter().find(|w| &w.id == id) else {
continue;
};
if !profile.enabled {
continue;
}
tried.push(id.clone());
let status = guard::probe(profile, billing);
if status.readiness == Readiness::Ready {
if let Some(bin) = status.binary_path {
let reason = if id == &candidate {
source.to_string()
} else {
format!("fallback ({candidate} not invocable)")
};
return Ok(Resolved {
worker_id: id.clone(),
bin,
reason,
});
}
}
}
Err(anyhow!(
"no invocable worker among {tried:?}. Run `yardlet worker status` to diagnose. \
Yardlet did not call an AI API and did not ask for an API key."
))
}
struct Candidate {
worker_id: String,
reason: &'static str,
}
pub fn candidate_for(
workers: &WorkersFile,
overrides: &RoutingOverrides,
override_w: Option<&str>,
preferred: &str,
kind: &str,
) -> (String, &'static str) {
if let Some(o) = override_w.filter(|s| !s.is_empty()) {
(o.to_string(), "run override")
} else if let Some(k) = overrides.kind_overrides.get(kind).filter(|s| !s.is_empty()) {
(k.clone(), "learned kind rule")
} else if !preferred.is_empty() {
(preferred.to_string(), "planner preferred")
} else {
(workers.routing.default_worker.clone(), "default")
}
}
fn candidate_for_task(
workers: &WorkersFile,
overrides: &RoutingOverrides,
override_w: Option<&str>,
task: &Task,
) -> Candidate {
let (worker_id, reason) = candidate_for(
workers,
overrides,
override_w,
&task.preferred_worker,
&task.kind,
);
Candidate { worker_id, reason }
}
pub(crate) fn norm_cap(s: &str) -> String {
s.trim().to_lowercase().replace([' ', '-'], "_")
}
fn worker_declares(workers: &WorkersFile, worker_id: &str, required: &[String]) -> bool {
workers
.workers
.iter()
.find(|w| w.id == worker_id)
.filter(|w| w.enabled)
.map(|w| {
let have: Vec<String> = w.capabilities.iter().map(|c| norm_cap(c)).collect();
required.iter().all(|r| have.iter().any(|h| h == r))
})
.unwrap_or(false)
}
fn first_capable(workers: &WorkersFile, required: &[String]) -> Option<String> {
workers
.workers
.iter()
.find(|w| {
w.enabled && {
let have: Vec<String> = w.capabilities.iter().map(|c| norm_cap(c)).collect();
required.iter().all(|r| have.iter().any(|h| h == r))
}
})
.map(|w| w.id.clone())
}
#[cfg(test)]
mod tests {
use super::*;
fn workers() -> WorkersFile {
crate::yaml::from_str(
"schema_version: 1\nrouting:\n default_worker: codex\n fallback_order: [codex, claude-code]\nworkers:\n - id: codex\n capabilities: [image_generation]\n invocation: { command: codex }\n - id: claude-code\n invocation: { command: claude }\n",
)
.unwrap()
}
#[test]
fn candidate_precedence() {
let w = workers();
let mut ov = RoutingOverrides::default();
ov.kind_overrides
.insert("refactor".into(), "claude-code".into());
assert_eq!(
candidate_for(&w, &ov, Some("codex"), "claude-code", "refactor").0,
"codex"
);
assert_eq!(
candidate_for(&w, &ov, None, "codex", "refactor").0,
"claude-code"
);
assert_eq!(
candidate_for(&w, &ov, None, "claude-code", "implementation").0,
"claude-code"
);
assert_eq!(
candidate_for(&w, &ov, None, "", "implementation").0,
"codex"
);
}
#[test]
fn norm_cap_folds_case_and_separators() {
assert_eq!(norm_cap(" Image-Generation "), "image_generation");
assert_eq!(norm_cap("image generation"), "image_generation");
}
#[test]
fn capability_gate_matches_only_declaring_workers() {
let w = workers();
let need = vec!["image_generation".to_string()];
assert!(worker_declares(&w, "codex", &need));
assert!(!worker_declares(&w, "claude-code", &need));
assert_eq!(first_capable(&w, &need).as_deref(), Some("codex"));
assert!(first_capable(&w, &["sorcery".to_string()]).is_none());
}
#[test]
fn capability_gate_requires_all_listed_capabilities() {
let w = workers();
let need = vec!["image_generation".to_string(), "video".to_string()];
assert!(!worker_declares(&w, "codex", &need));
assert!(first_capable(&w, &need).is_none());
}
}