use std::path::{Path, PathBuf};
use anyhow::{Context, Result};
use bitrouter_mcp::actions::route::{
ContextTierRates, EstimatedCost, PolicySelection, ProviderHop, ResolvedVia, RouteInput,
RouteQuery, RouteReport,
};
use bitrouter_mcp::error::ToolError;
use bitrouter_sdk::HeaderMap;
use bitrouter_sdk::caller::CallerContext;
use bitrouter_sdk::config::{Config, ConfigRoutingTable};
use bitrouter_sdk::language_model::types::{
GenerationParams, Message, Prompt, ProviderMetadata, Role,
};
use bitrouter_sdk::language_model::{RoutingPrefs, RoutingTable};
use crate::daemon::{DaemonCommand, DaemonResponse, RouteHop};
use crate::metering::PricingTable;
use crate::paths::ConfigSource;
use crate::policy_table_router::{PolicyDecision, PolicyTableRouter};
pub struct RouteAction {
source: ConfigSource,
socket: Option<PathBuf>,
}
impl RouteAction {
pub fn new(source: ConfigSource, socket: Option<PathBuf>) -> Self {
Self { source, socket }
}
pub async fn report(&self, input: RouteInput) -> Result<RouteReport> {
if let Some(socket) = self.socket.as_ref()
&& let Some(report) = self.via_daemon(socket, &input.model).await?
{
return Ok(report);
}
self.via_config(input).await
}
async fn via_daemon(&self, socket: &Path, model: &str) -> Result<Option<RouteReport>> {
match crate::daemon::send_command(
socket,
&DaemonCommand::Route {
model: model.to_string(),
},
)
.await
{
Ok(DaemonResponse::Route { chain }) => Ok(Some(assemble(
model,
model,
None,
ResolvedVia::Live,
None,
&chain,
&self.pricing().await,
))),
Ok(DaemonResponse::Error { message }) => {
Err(anyhow::anyhow!("resolving model '{model}': {message}"))
}
Ok(other) => {
tracing::debug!(response = ?other, "unexpected daemon route response — resolving from config");
Ok(None)
}
Err(e) => {
tracing::debug!(error = %e, "daemon route failed — resolving from config");
Ok(None)
}
}
}
async fn via_config(&self, input: RouteInput) -> Result<RouteReport> {
let resolved = self.resolved_config().await?;
let pricing = crate::assemble::build_pricing_table(&resolved);
let policy = PolicyTableRouter::from_config(&resolved.policy_table);
let table = ConfigRoutingTable::from_config(resolved);
let prompt = probe_prompt(&input);
let decision = policy.map(|p| p.decision_for(&prompt, &HeaderMap::new()));
let effective_model = decision
.as_ref()
.and_then(|d| d.selected_model.clone())
.unwrap_or_else(|| input.model.clone());
let effective_effort = decision.as_ref().and_then(|d| d.selected_effort);
let chain: Vec<RouteHop> = table
.route_chain(
&effective_model,
&RoutingPrefs::default(),
&CallerContext::local(),
)
.await
.with_context(|| format!("resolving model '{effective_model}'"))?
.into_iter()
.map(|t| RouteHop {
provider: t.provider_name,
service_id: t.service_id,
api_protocol: format!("{:?}", t.api_protocol).to_lowercase(),
})
.collect();
Ok(assemble(
&input.model,
&effective_model,
effective_effort,
if self.source.is_default() {
ResolvedVia::ZeroConfig
} else {
ResolvedVia::Config
},
decision.as_ref(),
&chain,
&pricing,
))
}
async fn resolved_config(&self) -> Result<Config> {
Ok(crate::commands::resolve_static(
crate::paths::load_config(&self.source).await?,
))
}
async fn pricing(&self) -> PricingTable {
match self.resolved_config().await {
Ok(config) => crate::assemble::build_pricing_table(&config),
Err(e) => {
tracing::debug!(error = %e, "no config for pricing — reporting the chain unpriced");
PricingTable::default()
}
}
}
}
fn assemble(
requested_model: &str,
effective_model: &str,
effective_effort: Option<bitrouter_sdk::language_model::types::ReasoningEffort>,
resolved_via: ResolvedVia,
decision: Option<&PolicyDecision>,
chain: &[RouteHop],
pricing: &PricingTable,
) -> RouteReport {
let estimated_cost = chain
.first()
.and_then(|h| pricing.resolve(&h.provider, &h.service_id))
.filter(|p| !p.is_unconfigured())
.map(|p| estimated_cost(&p));
RouteReport {
requested_model: requested_model.to_string(),
effective_model: effective_model.to_string(),
effective_effort,
resolved_via,
policy_decision: decision.map(policy_selection),
provider_chain: chain
.iter()
.map(|h| ProviderHop {
provider: h.provider.clone(),
service_id: h.service_id.clone(),
api_protocol: h.api_protocol.clone(),
})
.collect(),
estimated_cost,
}
}
fn probe_prompt(input: &RouteInput) -> Prompt {
let messages = match &input.prompt {
Some(text) => vec![Message::text(Role::User, text.clone())],
None => Vec::new(),
};
Prompt {
model: input.model.clone(),
system: None,
system_provider_metadata: ProviderMetadata::new(),
messages,
tools: Vec::new(),
params: GenerationParams::default(),
response_format: None,
tool_choice: None,
stream: false,
}
}
fn estimated_cost(p: &crate::metering::pricing::ModelPricing) -> EstimatedCost {
EstimatedCost::new(
p.input_micro_usd_per_token,
p.output_micro_usd_per_token,
p.context_tiers
.iter()
.map(|t| ContextTierRates {
above_input_tokens: t.above_input_tokens,
input_micro_usd_per_token: t.input_micro_usd_per_token,
output_micro_usd_per_token: t.output_micro_usd_per_token,
})
.collect(),
)
}
fn policy_selection(d: &PolicyDecision) -> PolicySelection {
PolicySelection {
request_key: d.request_key.clone(),
reason: d.reason.to_string(),
static_tier: d.static_tier.clone(),
static_model: d.static_model.clone(),
static_effort: d.static_effort,
selected_tier: d.selected_tier.clone(),
selected_model: d.selected_model.clone(),
selected_effort: d.selected_effort,
pinned: d.pinned,
locked: d.locked,
trialed: d.trialed,
}
}
#[async_trait::async_trait]
impl RouteQuery for RouteAction {
async fn route(&self, input: RouteInput) -> Result<RouteReport, ToolError> {
self.report(input)
.await
.map_err(|e| ToolError::new(format!("{e:#}")))
}
}
#[cfg(test)]
mod tests {
use super::*;
fn config_source(dir: &Path, yaml: &str) -> ConfigSource {
let path = dir.join("bitrouter.yaml");
std::fs::write(&path, yaml).expect("write config");
ConfigSource::File(path)
}
const ONE_MODEL: &str = r#"
providers:
demo:
api_base: https://api.example.test
api_key: sk-test
active: true
models:
- id: demo-model
"#;
const POLICY_REDIRECTS: &str = r#"
providers:
demo:
api_base: https://api.example.test
api_key: sk-test
active: true
models:
- id: demo-model
- id: demo-model-big
policy_table:
tiers:
big:
model: demo-model-big
effort: high
default_tier: big
"#;
#[tokio::test]
async fn both_surfaces_produce_the_same_report() {
let dir = tempfile::tempdir().expect("tempdir");
let action = RouteAction::new(config_source(dir.path(), ONE_MODEL), None);
let input = RouteInput {
model: "demo-model".to_string(),
prompt: Some("hello".to_string()),
};
let cli = action.report(input.clone()).await.expect("cli surface");
let tool = RouteQuery::route(&action, input)
.await
.expect("mcp surface");
assert_eq!(
serde_json::to_value(&cli).expect("cli json"),
serde_json::to_value(&tool).expect("mcp json"),
"the two surfaces of one action must be the same bytes"
);
assert_eq!(cli.requested_model, "demo-model");
assert_eq!(cli.resolved_via, ResolvedVia::Config);
assert_eq!(cli.provider_chain.len(), 1);
assert_eq!(cli.provider_chain[0].provider, "demo");
assert_eq!(cli.provider_chain[0].service_id, "demo-model");
}
#[tokio::test]
async fn both_surfaces_apply_the_policy_table() {
let dir = tempfile::tempdir().expect("tempdir");
let action = RouteAction::new(config_source(dir.path(), POLICY_REDIRECTS), None);
let input = RouteInput {
model: "demo-model".to_string(),
prompt: Some("write me a function".to_string()),
};
let cli = action.report(input.clone()).await.expect("cli surface");
let tool = RouteQuery::route(&action, input)
.await
.expect("mcp surface");
for (surface, report) in [("cli", &cli), ("mcp", &tool)] {
assert_eq!(report.requested_model, "demo-model", "{surface}");
assert_eq!(
report.effective_model, "demo-model-big",
"{surface} surface ignored the policy table — it would name a model the \
daemon would never pick"
);
assert_eq!(report.provider_chain[0].service_id, "demo-model-big");
let decision = report
.policy_decision
.as_ref()
.unwrap_or_else(|| panic!("{surface} surface reported no policy decision"));
assert_eq!(decision.selected_model.as_deref(), Some("demo-model-big"));
}
assert_eq!(
serde_json::to_value(&cli).expect("cli json"),
serde_json::to_value(&tool).expect("mcp json"),
);
}
#[tokio::test]
async fn an_edited_config_is_visible_to_the_next_call() {
let dir = tempfile::tempdir().expect("tempdir");
let source = config_source(dir.path(), ONE_MODEL);
let served: std::sync::Arc<dyn RouteQuery> =
std::sync::Arc::new(RouteAction::new(source, None));
let ask = || RouteInput {
model: "demo-model".to_string(),
prompt: None,
};
let before = served.route(ask()).await.expect("resolves");
assert_eq!(before.effective_model, "demo-model");
assert!(before.policy_decision.is_none(), "no policy table yet");
config_source(dir.path(), POLICY_REDIRECTS);
let after = served
.route(ask())
.await
.expect("resolves against the edited config");
assert_eq!(
after.effective_model, "demo-model-big",
"the same long-lived server must see the edited config, not a \
snapshot taken at `mcp serve` start"
);
}
#[tokio::test]
async fn a_dead_daemon_socket_falls_back_to_config() {
let dir = tempfile::tempdir().expect("tempdir");
let dead = dir.path().join("nothing-listening.sock");
let action = RouteAction::new(config_source(dir.path(), ONE_MODEL), Some(dead));
let report = action
.report(RouteInput {
model: "demo-model".to_string(),
prompt: None,
})
.await
.expect("resolves via config fallback");
assert_eq!(report.resolved_via, ResolvedVia::Config);
assert_eq!(report.provider_chain[0].provider, "demo");
}
#[tokio::test]
async fn an_unroutable_model_is_an_error() {
let dir = tempfile::tempdir().expect("tempdir");
let action = RouteAction::new(config_source(dir.path(), ONE_MODEL), None);
let err = RouteQuery::route(
&action,
RouteInput {
model: "nonexistent-model".to_string(),
prompt: None,
},
)
.await
.expect_err("an unroutable model surfaces a ToolError");
assert!(
err.0.contains("nonexistent-model"),
"the error names the unresolved model: {}",
err.0
);
}
#[test]
fn the_rate_card_surfaces_context_tiers() {
use crate::metering::pricing::{ContextTier, ModelPricing};
let flat = estimated_cost(&ModelPricing::new(1.0, 2.0));
assert!(flat.context_tiers.is_empty());
assert_eq!(flat.input_micro_usd_per_token, Some(1.0));
assert!(!flat.note.contains("context_tiers"), "note: {}", flat.note);
let tiered = estimated_cost(&ModelPricing {
input_micro_usd_per_token: Some(1.0),
output_micro_usd_per_token: Some(2.0),
cache_read_micro_usd_per_token: None,
cache_write_micro_usd_per_token: None,
context_tiers: vec![ContextTier {
above_input_tokens: 200_000,
input_micro_usd_per_token: Some(2.0),
output_micro_usd_per_token: Some(4.0),
cache_read_micro_usd_per_token: None,
cache_write_micro_usd_per_token: None,
}],
});
assert_eq!(tiered.context_tiers.len(), 1);
assert_eq!(tiered.context_tiers[0].above_input_tokens, 200_000);
assert_eq!(tiered.context_tiers[0].input_micro_usd_per_token, Some(2.0));
assert!(
tiered.note.contains("context_tiers"),
"the note has to explain the brackets: {}",
tiered.note
);
}
}